[None][feat] Page the DSpark drafter context through the draft KV cache manager by dc3671 · Pull Request #18343 · NVIDIA/TensorRT-LLM

dc3671 · 2026-08-28T03:12:20Z

Description

The standalone DSpark/DFlash drafter kept its context KV in a private arena: dense in max_seq_len, allocated after the KV pool was carved (so free_gpu_memory_fraction never bounded it), and sized from max_position_embeddings (1048576 for K3). This pages it through the draft KV cache manager's per-request block tables instead.

	before	after
drafter KV	private arena, dense in `max_seq_len`	paged through the draft KV manager
sized from	`max_position_embeddings`	runtime `max_seq_len`
K3 @ 997696, batch 8	21.4 GiB/rank, cannot start	inside the pool budget

Three gates kept an external drafter off that path, each written for a mode it does not describe:

gate	written for	carve-out
`_util.py` attention-DP bail	MTP, whose draft layers are target-shaped and appendable	DFLASH/DSPARK only; PARD and DRAFT_TARGET_ONE_MODEL unchanged
`py_executor_creator.py` nvbugs/5807902 WAR	Eagle3 disagg RMSNorm crash	blanket disable stays (default is `True`); standalone drafters exempt
`dflash.py` slot assignment	aggregated prefill	disagg generation workers get slots; `release_index_slot` mirrored to the draft manager

V1 get_buffers() returns a per-layer strided slice, V2 a dense view over interleaved pages, so the block-table divisor is derived from stride(0). Guessing wrong is silent: ~20% acceptance, no error.

Test Coverage

TestKimiK3DSpark::test_gsm8k_tep8, on the 2x4 GB300 post-merge stage from #18363 (l0_gb300_multi_nodes_node2_gpu8.yml). Its SingleNvlinkDomain name forces --segment=2; straddling two NVL domains dies in symm_mem.rendezvous.

Spec dec is lossless, so accuracy cannot detect a broken drafter — hence the AL assertion. That is how the slot collapse was found:

K3 DSpark disagg, GSM8K 1319, TEP16	AL	decode steps	output tokens
before	1.087	123860	134575
after	3.441	39072	134463

Measured on this branch, 8x GB300, 1319 samples, three runs. Target is the NVFP4 requant (nothing else covers K3 NVFP4 end to end), which also fixes the MoE backend: trtllm-gen ships SiTu cubins only for W4A8_MXFP4_MXFP8.

recipe (same weights, same drafter)	GSM8K	AL
MEGAMOE_CUTEDSL, NVFP4 (this test)	95.07 / 95.15 / 95.30	4.318 / 4.323 / 4.337
CUTLASS, NVFP4	96.209	4.349
TRTLLM, MXFP4	96.361	4.355

Not about this change, but this is the first real-model accuracy data on MEGAMOE_CUTEDSL: it lands ~1.0 pt below CUTLASS on identical weights and is not reproducible across greedy runs (0.23 pt spread). AL is flat at 4.32-4.36 across all three, which is what this PR guards. The gsm8k reference stays at 96.5 so the kernel gap is not baked in; the hypothesis-testing threshold (93.297) absorbs it.

min_al 4.0: a collapse scores 1.09 with every accuracy gate green. AL figures of 5.2-5.9 quoted elsewhere are --apply_chat_template runs (3.81 plain vs 5.52-5.88 chat, same drafter and samples).

Also green on the embedded DSpark path (DeepSeek-V4-Pro DEP4 cutedsl, AL 4.436) — different drafter shape, same shared kv_cache_manager_v2 / py_executor code.

The drafter checkpoint is RadixArk-Kimi-K3-DSpark on the CI model share (llm-models!596, merged; RadixArk and Inferact both publish a repo named Kimi-K3-DSpark, so both carry their vendor).

Not covered: the VANILLA drafter attention backend still uses contiguous flash_attn_with_kvcache.

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.

…he manager The standalone drafter kept a private context arena, dense in max_seq_len and allocated after the KV pool had been carved, so free_gpu_memory_fraction never bounded it: at max_seq_len 997696 with max_batch 8 it wants 21.4 GiB/rank and the worker cannot start. Read the already-funded draft pool through the manager's per-request block tables instead, so the footprint follows the sequences served, and size what remains from max_seq_len rather than the drafter's advertised max_position_embeddings (1048576 for K3, hundreds of GiB on its own). The pool index space differs between the V1 and V2 managers, so the block-table divisor is derived from stride(0) rather than assumed. Three gates kept an external drafter off that path, each written for a mode it does not describe: - attention DP: the bail suits MTP, whose draft layers are target-shaped and appendable to the target pool. An external drafter has its own architecture, so nothing is appended and it stayed on the arena -- which under attention DP is sized with KV heads unsharded (20480 vs 2560 B/token for K3). - disaggregation: nvbugs/5807902 reported an Eagle3 RMSNorm failure and was worked around by disabling the separate draft KV cache for every speculative mode. Keep the workaround where it was reported. - disagg slot allocation: _store_prefill_context was the only place that assigned a drafter slot, so a generation worker -- which receives prompt KV instead of prefilling -- collapsed every concurrent request onto the single dummy slot. A context-only worker also releases the target's IndexMapper slot after prefill; the draft mirror never got that call and saturated. Accuracy cannot detect a broken drafter: speculative decoding is lossless, so one producing garbage scores the same and only runs slower. The added test therefore asserts on acceptance length. That is not hypothetical -- the slot collapse showed up as AL 1.087 vs 3.441 (decode steps 123860 -> 39072 for the same output length) while every accuracy gate passed. Signed-off-by: Zhenhuan Chen <zhenhuanc@nvidia.com>

dc3671 · 2026-09-07T10:09:34Z

/bot run --stage-list "GB300-8_GPUs-2_Nodes-PyTorch-SingleNvlinkDomain-Post-Merge-1"

tensorrt-cicd · 2026-09-07T10:17:13Z

PR_Github #71928 [ run ] triggered by Bot. Commit: 43befee Link to invocation

tensorrt-cicd · 2026-09-07T11:51:51Z

PR_Github #71928 [ run ] completed with state SUCCESS. Commit: 43befee
/LLM/main/L0_MergeRequest_PR pipeline #58996 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

dc3671 · 2026-09-07T11:54:30Z

/bot run

tensorrt-cicd · 2026-09-07T12:00:19Z

PR_Github #71946 [ run ] triggered by Bot. Commit: 43befee Link to invocation

tensorrt-cicd · 2026-09-07T13:05:57Z

PR_Github #71946 [ run ] completed with state SUCCESS. Commit: 43befee
/LLM/main/L0_MergeRequest_PR pipeline #59011 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

dc3671 · 2026-09-07T15:40:08Z

/bot run

tensorrt-cicd · 2026-09-07T15:48:02Z

PR_Github #71982 [ run ] triggered by Bot. Commit: 43befee Link to invocation

tensorrt-cicd · 2026-09-07T17:28:31Z

PR_Github #71982 [ run ] completed with state SUCCESS. Commit: 43befee
/LLM/main/L0_MergeRequest_PR pipeline #59043 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

dc3671 · 2026-09-07T17:40:21Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-09-07T17:46:11Z

PR_Github #71993 [ run ] triggered by Bot. Commit: 43befee Link to invocation

tensorrt-cicd · 2026-09-07T20:04:43Z

PR_Github #71993 [ run ] completed with state SUCCESS. Commit: 43befee
/LLM/main/L0_MergeRequest_PR pipeline #59054 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

dc3671 · 2026-09-07T23:57:03Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-09-08T00:02:48Z

PR_Github #72008 [ run ] triggered by Bot. Commit: 43befee Link to invocation

tensorrt-cicd · 2026-09-08T03:38:48Z

PR_Github #72008 [ run ] completed with state FAILURE. Commit: 43befee
/LLM/main/L0_MergeRequest_PR pipeline #59069 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

dc3671 · 2026-09-08T08:59:35Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-09-08T09:06:16Z

PR_Github #72131 [ run ] triggered by Bot. Commit: 43befee Link to invocation

tensorrt-cicd · 2026-09-08T10:15:55Z

PR_Github #72131 [ run ] completed with state SUCCESS. Commit: 43befee
/LLM/main/L0_MergeRequest_PR pipeline #59176 completed with status: 'SUCCESS'

CI Report

Link to invocation

brnguyen2

Approving — the comments below are optional touch-ups, not blockers.

The [None] ticket slot: this is a nontrivial feature (new paging path through the draft KV manager, three gate carve-outs, a new post-merge test) — it should carry a TRTLLM JIRA so the K3 DSpark memory work is traceable.

Two coverage notes, neither blocking:

The _util.py attention-DP carve-out means attention DP + standalone drafter now builds a separate draft manager for the first time, but the new CI test runs enable_attention_dp=False and the DEP4 run cited in the description was manual. If a cheap config exists that exercises DP + DSpark with the draft manager bound, it would pin the ATTENTION_DP_DUMMY_REQUEST_ID handling in prepare() and _mirror_draft_kv_cache.
The divisor-from-stride heuristic in _init_ctx_block_tables is the piece most likely to silently regress when a manager changes its pool layout (your own description notes wrong-guess = ~20% acceptance, no error). The min_al gate in the new test is the only tripwire; worth a one-line comment in the V1/V2 get_buffers implementations pointing back at this consumer so a layout change reviews against it.

Description matches the diff closely, including the limitation notes.

…om the coordinator send path Follows NVIDIA#18343, whose change landed in the method this PR moved. Signed-off-by: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com>

Upstream's NVIDIA#18343 clamps to the block-table width and drops _ctx_block_counts, which the MLA port's read-length truncation consumes -- so the managed-pool generation path raised AttributeError. Restore the count and bound both writes. Signed-off-by: Zhenhuan Chen <zhenhuanc@nvidia.com>

github-actions Bot assigned dc3671 Aug 28, 2026

dc3671 force-pushed the user/zhenhuanc/dspark-kvmgr branch 3 times, most recently from ab3b1c0 to 6f5df55 Compare August 28, 2026 09:37

chungen04 mentioned this pull request Aug 29, 2026

[None][feat] Add Qwen3-based DSpark drafter (DeepSpec dense checkpoints) #16813

Closed

1 task

dc3671 force-pushed the user/zhenhuanc/dspark-kvmgr branch 6 times, most recently from bd30278 to 3c6afb9 Compare September 1, 2026 09:08

dc3671 marked this pull request as ready for review September 1, 2026 09:14

dc3671 requested review from a team as code owners September 1, 2026 09:14

dc3671 requested review from 2ez4bz, allisonlim-nv, brnguyen2, chienchunhung, eopXD, nvpohanh, reasonsolo, schetlur-nv, yechank-nvidia and zhaoyangwang-nvidia September 1, 2026 09:14

dc3671 force-pushed the user/zhenhuanc/dspark-kvmgr branch from 6a4e7b0 to 43befee Compare September 7, 2026 10:09

longlee0622 approved these changes Sep 8, 2026

View reviewed changes

brnguyen2 approved these changes Sep 8, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/speculative/dflash.py

Comment thread tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py

Comment thread tensorrt_llm/_torch/speculative/dflash.py

Comment thread tensorrt_llm/_torch/speculative/dflash.py

dc3671 merged commit 9dfa2e2 into NVIDIA:main Sep 8, 2026
8 checks passed

chungen04 mentioned this pull request Sep 10, 2026

[TRTLLM-16005][feat] Store DFlash draft cross-attention context paged in the target KV cache manager #16150

Closed

1 task

dc3671 mentioned this pull request Sep 11, 2026

[None][feat] MLA-backboned standalone DSpark drafter (Inferact/Kimi-K3-DSpark) #19040

Open

Uh oh!

Conversation

dc3671 commented Aug 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

dc3671 commented Sep 7, 2026

Uh oh!

tensorrt-cicd commented Sep 7, 2026

Uh oh!

tensorrt-cicd commented Sep 7, 2026

Uh oh!

dc3671 commented Sep 7, 2026

Uh oh!

tensorrt-cicd commented Sep 7, 2026

Uh oh!

tensorrt-cicd commented Sep 7, 2026

Uh oh!

dc3671 commented Sep 7, 2026

Uh oh!

tensorrt-cicd commented Sep 7, 2026

Uh oh!

tensorrt-cicd commented Sep 7, 2026

Uh oh!

dc3671 commented Sep 7, 2026

Uh oh!

tensorrt-cicd commented Sep 7, 2026

Uh oh!

tensorrt-cicd commented Sep 7, 2026

Uh oh!

dc3671 commented Sep 7, 2026

Uh oh!

tensorrt-cicd commented Sep 8, 2026

Uh oh!

tensorrt-cicd commented Sep 8, 2026

Uh oh!

dc3671 commented Sep 8, 2026

Uh oh!

tensorrt-cicd commented Sep 8, 2026

Uh oh!

tensorrt-cicd commented Sep 8, 2026

Uh oh!

brnguyen2 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

11 participants

dc3671 commented Aug 28, 2026 •

edited

Loading