[lucifer1004]

Description

Fixes #20548.

Each decode step of DFlash / DSpark with an MLA drafter runs a long chain of small torch kernels around the drafter. This PR adds two opt-in switches. Both are read per call, so a CUDA graph keeps what was read at capture, and both are bit-exact with the code they replace.

TLLM_DFLASH_FUSED_PROLOGUE=1

• dflash.py: these groups now run as one Triton kernel each, instead of chains of small torch kernels:
  • _refresh_ctx_block_tables: page counts and decoded table;
  • the gen-request prologue in prepare_1st_drafter_inputs: bonus token, query and context positions, embedding ids;
  • the noise block;
  • the gen tokens' context-KV store: rejected-row mask, clamp to the request's allocation, page lookup, write to every layer;
  • the ctx_len advance with the clamp, dummy-slot reset and advertised context length;
  • the single-latent _store_context_kv_paged scatter used by prefill.
• modeling_dspark.py (MLADSparkForCausalLM):
  • the layer-invariant block fixup is built by one kernel;
  • each layer's block latents go into the pool with one scatter kernel instead of an index_put;
  • the decode-step context-KV precompute runs the same compiled RMSNorm/RoPE kernel into one per-layer buffer, with one transpose copy instead of a contiguous copy and an output copy per layer.
• Each fused path checks its layout preconditions (contiguity, dtype, single-latent pool geometry shared by all layers) and otherwise falls back to the existing code.
• LiLiCorr drafters and a trained mask row whose dtype differs from the embedding keep the existing path.

TLLM_DFLASH_OVERLAP_CONTEXT_PROJECTION=1

• On pure-generation steps, the gen tokens' context projection (fc + hidden_norm + latent GEMM + RMSNorm/RoPE) is forked onto an aux stream before sample / accept and joined before the accepted-mask write.
• It reads only the captured target hidden states and the pre-step ctx_len; nothing between the fork and the join writes either.
• If the fused store is not taken, the overlapped result is dropped after the join and the existing code recomputes it.

Performance, Rubin, TP8 across 2 nodes (MNNVL), Kimi K3 with an MLA DSpark drafter, CUDA graphs, single request, decode step time. Measured with the same change on our integration branch, against the same stack without it:

	2K	32K	64K	96K	128K	393K
both switches	-0.195 ms	-0.206 ms	-0.186 ms	-0.212 ms	-0.224 ms	-0.210 ms
fused prologue only, vs. both	+0.002 ms	+0.002 ms	+0.001 ms	-0.003 ms	+0.006 ms	+0.004 ms

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

End to end, the overlap switch is within noise on this configuration. In a single-GPU microbenchmark (Rubin single GPU, 200-SM engineering sample, whole worker step under CUDA graph replay), it saved 53-84 µs on top of the fusion, bounded by the length of the sample / accept window it hides behind. It is kept as a separate switch so it can be measured on configurations with a longer window. I am happy to drop it from this PR if reviewers prefer.

Test Coverage

• New tests/unittest/_torch/speculative/test_dflash_fused_prologue.py (any CUDA GPU; added to l0_b200.yml). Each fused kernel is checked bitwise against the torch code it replaces:
  • block-table refresh: switch on vs. off through _refresh_ctx_block_tables, on a strided view of the offsets tensor;
  • single-latent context scatter: on vs. off through _store_context_kv_paged;
  • gen prologue on a strided accepted-token tensor, with requests at the context ceiling;
  • noise block vs. dflash_noise_block_embedding;
  • gen context store and ctx_len bookkeeping vs. the torch store, _advance_ctx_len and dflash_allocated_ctx_limit;
  • MLA block fixup vs. _build_mla_block_fixup (all fields), and the block-latent pool write on vs. off.
• tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py::test_mla_dspark_block_decode_matches_unabsorbed_reference counts which fixup builder the paged block decode used. With TLLM_DFLASH_FUSED_PROLOGUE=1 the fused builder runs instead, so the check now accepts either builder (and requires the fused one when the env is on); the numeric comparison is unchanged.
• GPU results at 4b80185 (Rubin single GPU, 200-SM engineering sample, SM107):
  • test_dflash_fused_prologue.py: 17 passed (17 failed on main, as expected).
  • hw_agnostic -k "dflash or dspark": the same pass / fail set as main (the existing failures there need the flash-attn package).
  • With TLLM_DFLASH_FUSED_PROLOGUE=1 TLLM_DFLASH_OVERLAP_CONTEXT_PROJECTION=1: one extra failure, the builder count above (a path assertion before the numeric check, not a numeric mismatch). Fixed in the test by the latest commit; the rerun is pending.
  • test_dspark_cute_dsl_attention.py / test_dspark_cute_dsl_rmsnorm_rope.py skip on SM107 (their skipif allows SM100 / SM103 only); pending on SM100 / SM103.
• Locally with TRITON_INTERPRET=1 on CPU: the new test passes (17); mutating any of four kernels makes 11 of them fail.
• Earlier validation of the same kernels on Rubin drove a real DSparkWorker with MLADSparkForCausalLM at Kimi K3 per-rank shapes (5 layers, 8 heads, page 32, block 7):
  • eager and CUDA graph (B = 1 / 4 with a dummy row), both pool layouts, all four switch combinations: step outputs, the whole pool and all bookkeeping bitwise equal to the unmodified code;
  • negative controls (a missing fork or join) are caught.
• Not covered: GQA drafters with the fused store (the single-latent geometry check keeps them on the existing path), 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.