[lucifer1004]

Description

Fixes #20551. Stacked on #20550, so this branch also contains that PR's commit. Review only the last commit, "MLA DSpark: run each drafter layer's kv path beside its q path".

Note: #19275 moves cute_dsl_kernels/blackwell/dspark_rmsnorm_rope.py to attention/kernels/blackwell/. Whichever of the two lands second updates the one DSparkRMSNormRoPEKernel import in modeling_dspark.py.

In MLADSparkForCausalLM.dflash_forward, each drafter layer runs its q path and its kv path one after another on one stream, although they are independent until attention. With TLLM_DSPARK_MLA_SPLIT_QKV_STREAMS=1, read per call and applied on the kernel block-decode paths only:

• after input_layernorm, the layer forks to a lazily created aux stream;
• the aux stream runs kv_a_proj_with_mqa -> RMSNorm/RoPE -> block write into the context pool;
• the main stream runs q_a_proj -> q_a_layernorm -> q_b_proj -> rope_into -> k absorption;
• the main stream joins before _mla_paged_attention, which skips its own block write (block_written=True).

Everything the aux stream reads is produced before the fork, and everything it writes is read only after the join. The next layer's fork orders the aux stream after this layer's main-stream work.

On the CuTe DSL block-decode variant, the RMSNorm/RoPE kernel stores the block latents straight into the pool at the fixup's pages / offsets, without a staging buffer and scatter. That kernel (_RMSNormRoPEPoolWriteKernel) is DSparkRMSNormRoPEKernel with only the store index changed. Other variants keep the staged write.

Two seams (_split_qkv_use_aux_stream, _split_qkv_direct_pool_write) let the two halves be measured separately.

Performance, Rubin, TP8 across 2 nodes (MNNVL), Kimi K3 with an MLA DSpark drafter, CUDA graphs, single request, decode step time, on top of #20550. This is a cross-job comparison, so the differences are close to cross-node noise; on a later stack, turning the switch off cost about +0.06 ms per step.

	2K	32K	64K	96K	128K	393K
split q / kv streams	-0.053 ms	-0.038 ms	-0.025 ms	-0.037 ms	-0.023 ms	-0.018 ms

Accepted tokens per step, accuracy and needle checks were unchanged.

In a single-GPU microbenchmark (Rubin single GPU, 200-SM engineering sample, drafter forward under CUDA graph replay), it saves 56-68 µs per step at B = 1 / 4. Almost all of that comes from the two streams; the direct pool write alone saves 6-19 µs but adds about nothing on top of the streams, because it is hidden behind the q path.

Test Coverage

• tests/unittest/_torch/speculative/test_dspark_cute_dsl_rmsnorm_rope.py (SM100 / SM103, in l0_b200.yml): new test that the direct pool write is bitwise what _rmsnorm_rope_batched plus the staged scatter write (B = 1 / 4).
• Depends on [#20365][fix] Use the fused DSpark RMSNorm/RoPE kernels on SM107 #20366 for SM107. The direct pool write is gated by is_fused_dspark_rmsnorm_rope_supported, and on main dspark_rmsnorm_rope_custom_op._DSV4_DSPARK_ARCH_BY_SM only maps SM100 / SM103, so on SM107 (Rubin) the kernel path is not taken and the switch falls back to the staged write. [#20365][fix] Use the fused DSpark RMSNorm/RoPE kernels on SM107 #20366 adds SM107 to that map; validating the pool write on SM107 needs it.
• GPU results at fee0fb7 (same code as now, Rubin single GPU, 200-SM engineering sample, SM107, without [#20365][fix] Use the fused DSpark RMSNorm/RoPE kernels on SM107 #20366):
  • the new pool-write test (run with -p no:skipping) fails as expected there, because the fused kernel is unsupported on SM107 without [#20365][fix] Use the fused DSpark RMSNorm/RoPE kernels on SM107 #20366 and the function reports the fallback; the existing test_fused_dspark_rmsnorm_rope_* cases fail the same way on main;
  • test_dspark_cute_dsl_rmsnorm_rope.py + test_dflash_fused_prologue.py with the normal skips: 17 passed, 22 skipped (the 17 are [#20548][perf] DFlash / MLA DSpark: fuse the decode-step drafter bookkeeping into Triton kernels #20550's).
• Pending: the pool-write test on SM100 / SM103, and on SM107 together with [#20365][fix] Use the fused DSpark RMSNorm/RoPE kernels on SM107 #20366. On CPU, the module imports and the dflash / dspark CPU tests have the same result as main.
• Earlier validation of the same code on Rubin, with a real DSparkWorker and MLADSparkForCausalLM at Kimi K3 per-rank shapes:
  • all switch combinations bitwise equal over 5 scenarios, eager and CUDA graph, with and without the low-M GEMM path;
  • negative controls: dropping the fork or the join gives illegal addresses or a capture error, while delaying either stream with correct fork / join stays bitwise equal.
• Not covered: multi-GPU TP, attention DP.
• pre-commit on the changed files: clean.

PR Checklist

• PR description clearly explains what and why.
• PR follows the TRT-LLM coding guidelines to the best of my knowledge.
• Test cases are provided for new code paths.
• No API changes, no new dependencies, no CODEOWNERS changes.
• Please check this after reviewing the above items as appropriate for this PR.