[zheyuf]

@coderabbitai summary

Description

Makes the TRTLLM lowering of the non-causal windowed DSpark draft block cheap. This matters most when acceptance is fixed and only the per-step cost shows.

Why: since efa53ad, the TRTLLM backend serves a non-causal windowed draft layer as the causal windowed trtllm-gen decode plus a log-sum-exp merge of the block keys after each query (dflash_merge_future_block_attention). The merge ran eagerly: dozens of small kernels per layer. On GB300 it cost 61-148 us per layer on top of a 7-94 us decode; the MiniMax-M3.1 drafter runs it in 5 layers every step. #20922 carries the same eager merge (7d172e6).

Changes (one commit; the code is in modeling_dflash.py):

1. Compiled merge. dflash_merge_future_block_attention computes the same math, unrolled over the (at most block-size) block keys. dflash_merge_future_block_attention_ runs it under maybe_compile, in place on the decode's base-2 LSE. Default (automatic) dynamic shapes are used, so only the batch dimension becomes symbolic. It compiles at the eager warm-up that precedes each CUDA-graph capture.
2. No merge without later keys. With one block token or no right window, no block key follows a query. The layer then skips the LSE and the merge, and the causal decode is the exact result (dflash_merges_future_block_keys).
3. Workspace for the LSE. When the decode returns the LSE, flashinfer 0.6.18 carves 8 B x heads x requests x round_up(block, 256) + 1 MiB of softmax stats off the front of the workspace (trtllm_fmha_kernel_launcher.cu:279-305). The kernel scratch gets the rest. The drafter sized its workspace without the stats; only the 32 MiB trtllm-gen reservation (kTrtllmGenWorkspaceSize) left room for them. At 64 query heads per rank with a 3-token block that ran out above ~400 requests, and flashinfer refused the launch. The size now includes the stats whenever a layer takes the merge (dflash_trtllm_gen_workspace_size). No out-of-bounds write was possible: the carve is bounds-checked, and a byte-pattern probe measured the multi-CTA scratch at <= 4.1 MiB, used only at small batch, where the stats are small.

How the formulation was chosen (merge only, us per call, CUDA graphs, GB300; 16 q / 1 kv heads = MiniMax-M3.1 TP4 rank, K = 7):

formulation	B=1	B=8	B=32	B=64
eager (efa53ad)	67-69	77-86	84-85	88-89
efa53ad's merge under torch.compile, automatic dynamic	16.3	17.4	21.9	26.5
einsum + sigmoid/lerp (D2's fixup), dynamic=True	12.1	13.6	18.0	23.7
einsum + sigmoid/lerp, automatic dynamic	12.1	13.6-14.7	15.8-16.2	18.8-19.3
unrolled (this PR), automatic dynamic	4.9	6.2	8.4	11.4
one hand-written Triton kernel (reference only)	4.6	4.6	4.7	4.7

At 64 q / 4 kv heads (attention DP), B=64: eager 139-156, einsum + dynamic=True 114, einsum automatic 30-35, unrolled 25.6 us. Every compiled form matches the eager merge to bf16 rounding (max relative difference 3.4e-3).

Results on GB300 (the base's release image + this branch's Python):

Per-layer drafter attention: the cost of the non-causal layer over the causal decode, base -> this PR (us per layer; window 4096, 64-token pages, contexts 2K/8K/16K, CUDA graphs, mean of 2 runs in opposite order):

heads (q/kv)	K	B=1	B=8	B=16	B=64
16/1 (TP4)	7	+69..78 -> +6.2..6.7	+85 -> +6.8..7.0		+86..98 -> +13.3..13.9
16/1	3	+76 -> +5.2..5.5	+80..81 -> +6.5..6.6		+77..84 -> +7.8..8.3
16/1	1	+66 -> +0.0	+69 -> +0.0		+61..69 -> +0.0
64/4 (ADP)	7	+73..74 -> +6.2	+83..86 -> +7.9..8.2	+91..94 -> +10.8..11.1	+145..148 -> +27.9..28.9
64/4	3	+71 -> +5.5..5.7	+80..81 -> +6.4..6.8	+79..84 -> +7.0..8.0	+110 -> +11.9..13.7

The whole layer is 2.0-6.9x faster (geomean 4.0x over the 42 K > 1 points). K = 1 output is bit-identical to before.

E2E: MiniMax-M3.1-Flash TP4 with its DSpark drafter, K = 7, acceptance forced (TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS=2.6, AL 3.6 on both sides). Device step time is the median over generation-only iterations with exactly B requests (print_iter_log). Base (efa53ad) and this PR ran interleaved on one node, 2 rounds each:

config	requests	base r1 / r2 (ms)	this PR r1 / r2 (ms)	change (means)
TP4 EP4, bs64, CUDA graphs	1	8.167 / 7.965	7.667 / 7.596	-0.43 ms (-5.4%)
	8	14.783 / 14.602	14.288 / 14.111	-0.49 ms (-3.4%)
	64	28.020 / 28.085	27.744 / 27.954	-0.20 ms (-0.7%)
production recipe: TP4 EP1, bs30, MSL 1M, block reuse + host cache, piecewise CUDA graphs + torch.compile, PDL	1	7.879 / 7.979	8.053 / 7.659	-0.07 ms (-0.9%; rounds +2.2% / -4.0%)
	8	14.364 / 14.141	13.744 / 13.684	-0.54 ms (-3.8%)
	30	22.057 / 22.051	21.718 / 21.816	-0.29 ms (-1.3%)

• Forced AL is identical on every server.
• The client-side step medians agree.
• Startup: the merge compiles during the CUDA-graph warm-up. This adds ~20 s in the production recipe and nothing visible in the bs64 config.

AL is unchanged (80 SPEED-Bench coding prompts, /v1/completions, T=1, top_p 0.95, 4096 tokens, c64; mean of an unseeded and a seeded pass, K = 7):

config	base	this PR
default (use_rejection_sampling: false), TP4	3.624 (3.643 / 3.605)	3.601 (3.602 / 3.600)
golden semantics (RS + block verification + temperature-only q), with #21085	4.030 (base + #21085 + #21086)	4.024 (base + #21085 + this PR)
default, attention DP (DEP4: 64 q / 4 kv drafter heads per rank)	3.573	3.639
default, production recipe, fresh pass only	3.632 (D2's lowering, the same math)	3.605

• The SA golden at K = 7 is 4.04.
• Differences are within the run-to-run noise: three servers of identical code measured 3.537-3.620 in the default row.
• 0 request errors. Every server logs the non-causal lowering on all 5 draft layers. No error lines.

Test Coverage

• tests/unittest/_torch/speculative/hw_agnostic/test_dspark_noncausal_window.py:
  • The oracle test now also runs 64 q / 4 kv heads.
  • _PagedCase sizes the workspace as the drafter does.
  • New: which layers take the merge (CPU).
  • New: the compiled in-place merge equals the eager merge at B 1 / 5 / 64, K 2 / 3 / 7, 16/1 and 64/4 heads (CUDA).
  • New: the windowed non-causal decode at 64 q / 4 kv heads, 512 requests and K = 3 fits its workspace and matches the oracle. Without the stats, flashinfer refuses this shape.
  • Result: 86 passed (62 before).
• Also run on this branch, GB300:
  • A drafter-level check with a MiniMax-shaped drafter: TRTLLM vs VANILLA vs an fp32 reference, 16/1 and 64/4 heads, K 1/3/7, rebased context rows, CUDA-graph replay. 8 passed.
  • The DFlash / DSpark worker, draft-KV, CUDA-graph, rollback, heads, LiLiCorr, arena-sizing and config tests, test_dflash2_semantics.py, test_minimax_dspark_draft.py, test_modeling_speculative.py: 423 passed, 8 skipped. One failure, test_mla_dspark_auto_backend_resolves_to_cutedsl, fails the same way in this image without this PR.
  • api_stability: 38 passed.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

🤖 Generated with Claude Code