[zheyuf]

Description

Problem

The DFlash drafter's VANILLA attention backend keeps the drafter's context K/V in a private arena outside the KV cache manager. Verification is lossless, so nothing fails, but acceptance length (AL) silently drops in three situations:

• Prefix-cache hit. The target skips the reused prefix and the arena starts the request empty, so the drafter drafts from the un-reused tail only. The worker warns that AL "drops as if unpaired". SPEED-Bench never hits this: each pass uses a unique cache_salt. Multi-turn AgentX traffic hits it constantly.
• Disaggregated serving. The arena is never transferred, so the generation worker drafts without the prompt. Muse 1P1D coding AL was 3.91 on VANILLA vs 4.13 aggregated ([None][feat] Run the Muse modeling_v2 composition with a TRTLLM DFlash drafter #20812). Qwen3-8B loses about 17%.
• Memory. The arena is outside the KV budget and dense in max_seq_len (245-309 GiB per GPU at 1M tokens on the Muse Vera Rubin deployments). A separate budget-check apparatus exists only because of it.

TRTLLM keeps the drafter KV in the KV cache manager. With KVCMv2 and TRTLLM target attention it lives in the unified pool; otherwise it uses a separate draft KVCM that is paired for joint reuse on V2. That makes the drafter's context reusable on a prefix hit, transferable to the generation worker, and counted in the KV budget.

The tracker item reads "remove VANILLA, add AL testing, make TRTLLM the default". This PR does the second and third parts. For the first, it scopes VANILLA down instead of deleting it, because deleting it is not safe yet:

• Other GPUs. VANILLA is the only DFlash backend outside the SM100 family, apart from FA4, which is opt-in and SM90-only. CI runs DFlash on VANILLA on H100.
• Unwindowed drafters. On TRTLLM they run the non-causal block on the context kernel, 2-9x slower than VANILLA ([None][feat] Migrate DFlash2 sliding-window drafters to the TRTLLM attention backend #20776). There is no end-to-end measurement. Every DFlash checkpoint in llm-models except DFlash2 is unwindowed.
• Muse on Vera Rubin. A TRTLLM drafter's decode step is +2.8% at c1 on VR (+1-2% on GB300). Five fast paths in [None][feat] Muse Spark Open modeling_v2 TEP target: a per-deployment lever resolver and the Vera Rubin high-concurrency, mid-concurrency and low-latency fast paths #21377 only work over the dense or windowed arena (drafter_fc_early_fork, drafter_kv_side_chain, drafter_preamble_from_commit, incremental_block_offsets, windowed_commit_parallel), but they are not where the gap is: turning the four that apply at c1 off on VANILLA costs +0.22%, so about 2.6 points are the paged drafter attention itself.
• Cold-page KV-cache compression. The Muse cold-page point (pareto9) cannot run a TRTLLM drafter at speed yet, so AUTO keeps VANILLA under compression (item 4 below).

What changed

1. AUTO resolves to TRTLLM wherever it can run the drafter.
   • Windowed drafters already resolved to TRTLLM, including under modeling_v2. That is unchanged.
   • New: a disaggregated server that will share its draft KV now picks TRTLLM for unwindowed drafters too. The rule (dflash_auto_shares_draft_kv) mirrors the config-time half of _dflash_unify_blocker: a transceiver is set, sharing is not disabled, the target uses TRTLLM attention, KVCMv2 is set explicitly, pp=1, cp=1, there is no KV compression and no draft_len_schedule.
   • The frontend, the rank-side pin (_pin_dflash_auto_attention_backend) and the HF-hub download path all resolve AUTO the same way.
   • AUTO now also checks the drafter's head geometry against the TRTLLM kernels (Q/KV ratio, context head_dim) and keeps VANILLA if the kernels cannot take it. Before, such a drafter failed at its first forward.
   • AUTO still picks VANILLA in four cases: TRTLLM-gen is unavailable, KV cache compression is enabled, an unwindowed drafter on an aggregated server, or TRTLLM_DISAGG_SHARED_DRAFT_KV=0.
2. Explicit VANILLA is deprecated where TRTLLM can run the drafter. If AUTO would pick TRTLLM for the deployment, setting attention_backend: VANILLA logs a deprecation warning that names the lost context.
3. No silent AL loss in disaggregated serving.
   • A TRTLLM drafter that ends up in an unshared pool on a disaggregated worker now logs the blocker and that the generation worker drafts without the prompt. Windowed drafters already raised an error in this case.
   • A generation worker whose context worker runs without the drafter fails peer registration with an explicit message, which was already the case for windowed drafters.
4. Cold-page KV-cache compression admits DFlash on every backend; AUTO keeps VANILLA under it.
   • validate_kv_cache_compression_compatibility now admits DFlash on every backend.
   • The DFlash drafter's separate draft KVCM is built without a cold-page codec (draft_kv_cold_page_codec_provider), so draft pages move to and from Host/Disk losslessly. VANILLA's arena never leaves the GPU, so this matches it. The codec only ever sees target pages.
   • Compression still blocks the unified pool, so a TRTLLM drafter uses the separate draft KVCM, which prefix reuse does not reach. That pool is sized from the drafter's declared windows (the DFlash2 config marks all 5 layers sliding_attention, sliding_window: 2048, so the split charges it 0 bytes/token and gives it 0.55 GiB per GPU), while the unpaired mirror reserves the full prompt. On AgentX traffic at the Muse cold-page point (pareto9, TEP4 c96, 64k-336k-token contexts) only 1-3 requests fit at a time: 44 active requests, 2.9 generating on average, prefix hit 22%, output throughput 123 vs 955-968 tok/s/GPU on VANILLA (-87%), TTFT p90 ~400 s. SPEED-Bench (80 short prompts) passes and does not show it.
   • So AUTO resolves to VANILLA whenever kv_cache_compression_config is set (resolve_from_checkpoint(kv_cache_compression=...), passed by every place that pins AUTO). An explicit TRTLLM still runs.
5. Docs for DFlash attention_backend and the draft KV cache, and the cold-page example's list of supported speculative modes.

Out of scope: the CTX memory transient when TRTLLM_SKIP_KV_CACHE_ESTIMATION=1 (Remaining blockers).

Interplay with #21377 (either order)

• _util.py and llm_args.py auto-merge. I checked this with git merge-tree against [None][feat] Muse Spark Open modeling_v2 TEP target: a per-deployment lever resolver and the Vera Rubin high-concurrency, mid-concurrency and low-latency fast paths #21377 at e05bfc4.
  • This PR's cold-page admission is an early return above the lines [None][feat] Muse Spark Open modeling_v2 TEP target: a per-deployment lever resolver and the Vera Rubin high-concurrency, mid-concurrency and low-latency fast paths #21377 edits.
  • After both land, [None][feat] Muse Spark Open modeling_v2 TEP target: a per-deployment lever resolver and the Vera Rubin high-concurrency, mid-concurrency and low-latency fast paths #21377's dflash_without_draft_kvcm term is dead code and can be dropped.
• One textual conflict: tests/unittest/_torch/kv_cache_compression/test_quantization_for_cold_page.py.
  • Both PRs edit test_speculative_admission_accepts_verified_one_model_modes identically, but [None][feat] Muse Spark Open modeling_v2 TEP target: a per-deployment lever resolver and the Vera Rubin high-concurrency, mid-concurrency and low-latency fast paths #21377 appends tests right after it.
  • Whichever PR lands second should keep [None][feat] Muse Spark Open modeling_v2 TEP target: a per-deployment lever resolver and the Vera Rubin high-concurrency, mid-concurrency and low-latency fast paths #21377's added tests and delete test_speculative_admission_rejects_dflash_with_draft_kv_cache. That test asserts the rejection this PR removes; test_cold_page_dflash.py covers TRTLLM admission and the codec-free draft manager instead.
• [None][feat] Muse Spark Open modeling_v2 TEP target: a per-deployment lever resolver and the Vera Rubin high-concurrency, mid-concurrency and low-latency fast paths #21377's global_scale / uncompressed_layer_types then apply only to target pages. That matches the deployment's uncompressed_layer_types: [sliding_window] for its all-SWA drafter.
• [None][feat] Muse Spark Open modeling_v2 TEP target: a per-deployment lever resolver and the Vera Rubin high-concurrency, mid-concurrency and low-latency fast paths #21377's VANILLA-only levers are not touched. With AUTO resolving to TRTLLM they are declined, which accounts for the measured +2.8%.

Test Coverage

New tests, and the CI lists that carry them:

• unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py, 5 new tests (CPU, l0_cpu_musespark):
  • the AUTO matrix: disaggregated unwindowed resolves to TRTLLM, each blocker keeps VANILLA, unsupported geometry keeps VANILLA, KV cache compression keeps VANILLA (frontend with and without a GPU, and the rank pin);
  • the rank-pin and frontend parity;
  • the VANILLA deprecation warning.
• unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py, 1 new test (CPU): the unshared TRTLLM pool on a disaggregated worker warns, and TRTLLM_DISAGG_SHARED_DRAFT_KV=0 silences it.
• unittest/_torch/kv_cache_compression/test_cold_page_dflash.py, new (CPU): cold-page compression admits DFlash on every backend, the DFlash draft KVCM gets no codec, and EAGLE3 keeps its codec.
• accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_dflash_acceptance_under_prefix_reuse[TRTLLM|VANILLA] (1 GPU SM100+, l0_b200, l0_gb300_musespark):
  • Qwen3-8B with the unwindowed Qwen3-8B-DFlash-b16 drafter, K=15, 64 GSM8K 5-shot prompts, greedy.
  • Each prompt has its own cache_salt. A cold pass is followed by a warm pass that must reuse blocks; warm AL must be at least 0.95 x cold AL.
  • The VANILLA case is a strict xfail: it documents the bug class and proves the test detects it.
• accuracy/test_disaggregated_serving.py::TestQwen3_8BDFlash::test_default_1p1d_acceptance_matches_aggregated (2 GPU, l0_dgx_b200, l0_gb300_musespark):
  • The config sets neither attention_backend nor the sharing env.
  • The test asserts the worker log line "target and draft KV layers share one KV cache manager".
  • 1P1D AL must be at least 0.95 x the aggregated TRTLLM AL on the same prompts.
• accuracy/modeling_v2/test_musespark_open.py::TestModelingV2MuseSparkOpenNvfp4Tep4 (4x GB300/VR, l0_gb300_multi_gpus; commented out in l0_gb300_musespark until the DFlash2 checkpoint is in the CI cache):
  • test_dflash2_speed_bench_coding_golden_acceptance: the SemiAnalysis golden protocol, with the drafter on AUTO (TRTLLM).
    • 80 SPEED-Bench coding prompts, chat template, T=1.0, top_p 0.95, max 4096 tokens, rejection sampling, K=7.
    • Two seeded passes, pooled.
    • Floor min_al 3.969, the golden 4.05 less 2%. ref_al is the golden itself, not a measurement.
    • The prompts are read from MUSESPARK_SPEED_BENCH_CODING or llm-models/datasets/nvidia/SPEED-Bench/qualitative_coding_turn0.jsonl, and the test skips if they are absent.
  • test_dflash2_acceptance_under_prefix_reuse: the Muse modeling_v2 target, drafter on AUTO.
• test_quantization_for_cold_page.py is updated: DFlash is no longer in the rejected list.

Validated

Run on one Vera Rubin node, inside the Muse delivery image, with this PR's six changed tensorrt_llm files bind-mounted over the image. Those files are identical to this branch except that _util.py keeps #21377's cold-page exemption.

• Unit tests: 246 passed. Files: test_dflash_draft_kv.py, test_dflash_unified_kv_gate.py, test_cold_page_dflash.py, test_quantization_for_cold_page.py, test_disagg_dflash_guard.py, test_dflash_worker.py, test_dflash_trtllm_window.py, test_dflash2_semantics.py.

• TestQwen3_8B::test_dflash_acceptance_under_prefix_reuse: 1 passed, 1 xfailed.

  • TRTLLM: AL cold 3.4317, warm 3.4179 (-0.40%). Worker log: unified pool, reuse_addressable=True.
  • VANILLA: AL cold 3.4179, warm 2.6548 (-22.33%), the expected strict xfail. Worker log: reuse_addressable=False.

• TestQwen3_8BDFlash::test_default_1p1d_acceptance_matches_aggregated: passed.

  • AL aggregated (TRTLLM) 3.3863, 1P1D with AUTO 3.4045 (+0.54%).
  • The workers logged the shared target and draft KV manager. The drafter is unwindowed, so before this PR AUTO resolved to VANILLA on these workers. That is inferred from the code; the pre-PR baseline was not run.

• Unwindowed drafter cost on TRTLLM, measured on the same node. Ad-hoc run, not a repo test. Qwen3-8B with Qwen3-8B-DFlash-b16, K=15, prompts of about 1.2K tokens, no reuse. Median GPU forward time per pure-decode step, TRTLLM vs VANILLA:

  • c1: 3.222 vs 3.192 ms (+0.9%);
  • c16: 4.668 vs 4.611 ms (+1.2%).
  • Against the 22% AL that VANILLA loses on a prefix hit, this is small. Long contexts, where the context kernel's cost grows, are unmeasured, so aggregated AUTO still keeps VANILLA for unwindowed drafters.

• Cold-page compression in the same session. The ad-hoc run left the backend on AUTO, so this unwindowed aggregated drafter resolved to VANILLA. It started with NVFP4 cold pages and a host tier, and AL fell 24.5% on the warm pass. The TRTLLM + cold-page combination (separate draft KVCM without a codec) is therefore only unit-tested so far.

Muse Spark Open on Vera Rubin (delivery image 32862f7e42 with this PR's files bind-mounted, MR !163 recipes with attention_backend unset):

• Unit tests after the compression fix: 248 passed (same files).
• test_musespark_open.py, 4x VR: it needs MUSESPARK_OPEN_SM107_FUSED_KERNELS=1 on sm107, now in the test's environment.
  • test_dflash2_speed_bench_coding_golden_acceptance: passed, AL 4.0815 (golden 4.05, floor 3.969).
  • test_dflash2_acceptance_under_prefix_reuse: passed, AL cold 3.6568, warm 3.6963 (+1.08%).
• SPEED-Bench golden protocol through trtllm-serve (80 coding prompts, 4 cold passes pooled):
  • pareto9 aggregated, cold-page compression, AUTO (then TRTLLM, before the compression fix): 4.1449 (+2.34%);
  • pareto3 disaggregated, AUTO on both servers (TRTLLM, one shared KV manager on each): 4.1313 (+2.01%).
• AgentX traffic, forced AL 3.05, rejection sampling, pareto9 with compression: before the compression fix AUTO picked TRTLLM, with the -87% above. With the fix AUTO resolves to VANILLA (AUTO keeps VANILLA: with KV cache compression ...): 965 tok/s/GPU, x p10 97.8 tok/s/user, TTFT p90 3.3 s, prefix hit 86.7%, against an explicit VANILLA at 955 and 968, 97.8 and 99.5, 3.2 s.
• AgentX natural AL (rejection sampling, prefix hit 86-91%, contexts up to 336k tokens): VANILLA and TRTLLM agree within noise on pareto1 (dense arena, -1.3%, paired 95% CI +-6%), pareto6 (+0.7%) and pareto8 (-0.2%). The Muse windowed and dense arenas do not lose AL on prefix hits the way the generic VANILLA arena does.

Not validated here:

• TRTLLM drafter + cold-page compression at speed (see item 4);
• Muse performance parity (see Remaining blockers).

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.

Remaining blockers to deleting VANILLA

1. Muse perf parity. On VR the TRTLLM drafter costs +2.8% per decode step at c1 (pareto1), about 2.6 points of it the paged drafter attention itself. In aggregated AgentX serving it also costs 10-12% per context iteration at pareto6/9: the paged arena's context store reads _ctx_block_counts[:num_contexts].tolist() from the GPU, which blocks the host for the whole context forward (nsys: 37 blocking 8-40-byte copies, 105 ms each, per 300 iterations at pareto9; none on VANILLA). In disaggregated serving at pareto7 the shared pool also costs the CTX server 18% of its KV blocks, and its prefix hit falls from 86% to about 50% when long agentic sessions resume (TTFT p99 17 s to 122-149 s, throughput -10%).
2. Cold-page compression. Compression keeps the target and draft KV out of one manager, and the codec rejects mixed-dtype layer specs (FP8 target + BF16 draft). In aggregated serving the separate draft pool is undersized for long contexts (item 4 of What changed); a windowed TRTLLM drafter on a disaggregated server with compression still errors and points to VANILLA. Either the shared manager has to take compression, or the separate draft pool has to keep only the drafter's window.
3. Unwindowed drafters on TRTLLM. At ~1.2K-token contexts the decode-step cost is about 1% (above). Long contexts still need measuring before aggregated AUTO can switch, which would also make prefix reuse lossless for those drafters. The other route is to generalize the windowed decode-plus-correction path.
4. A paged backend outside the SM100 family (FA4 beyond SM90, or FlashInfer).
5. CTX memory when KV-cache estimation is skipped. The pool is sized by fraction before any forward, so the drafter's prefill transient on the CTX worker must fit in 1 - free_gpu_memory_fraction (Muse CTX: 0.8 to 0.7). With estimation on, the transient is measured. Accounting for it in code needs a measured transient model; this is a follow-up.
6. Deployed recipes. The Muse Vera Rubin delivery recipes pin VANILLA on all 9 points and must drop it.

🤖 Generated with Claude Code