[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.

tensorrt-cicd · 2026-09-03T02:55:21Z

PR_Github #71140 [ run ] triggered by Bot. Commit: eb07ced Link to invocation

dc3671 · 2026-09-03T02:59:49Z

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

tensorrt-cicd · 2026-09-03T03:06:23Z

PR_Github #71146 [ run ] triggered by Bot. Commit: eb07ced Link to invocation

tensorrt-cicd · 2026-09-03T03:11:01Z

PR_Github #71140 [ run ] completed with state ABORTED. Commit: eb07ced

Link to invocation

tensorrt-cicd · 2026-09-03T04:43:38Z

PR_Github #71146 [ run ] completed with state FAILURE. Commit: eb07ced
/LLM/main/L0_MergeRequest_PR pipeline #58288 (Partly Tested) 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

ZhanruiSunCh · 2026-09-03T09:53:01Z

The new GB300 test was selected in L0 #58288 but failed before executing GSM8K: /scratch.trt_llm_data/llm-models/RadixArk-Kimi-K3-DSpark was not recognized as a local model directory, so the loader treated it as a Hugging Face repo id and hit HFValidationError (then TypeError: HFValidationError object is not iterable). Please add RadixArk-Kimi-K3-DSpark to the CI llm-models provisioning for this lane (or use the correct downloadable model identifier), then rerun the GB300 2-node stage.

dc3671 · 2026-09-03T09:55:01Z

The new GB300 test was selected in L0 #58288 but failed before executing GSM8K: /scratch.trt_llm_data/llm-models/RadixArk-Kimi-K3-DSpark was not recognized as a local model directory, so the loader treated it as a Hugging Face repo id and hit HFValidationError (then TypeError: HFValidationError object is not iterable). Please add RadixArk-Kimi-K3-DSpark to the CI llm-models provisioning for this lane (or use the correct downloadable model identifier), then rerun the GB300 2-node stage.

@ZhanruiSunCh Yes, I've added it in llm-models two days ago. But it seems hasn't been synced yet.

nvpohanh · 2026-09-04T07:22:11Z

[by Codex] @yizhang-nv Could you review this PR? Thanks!

dc3671 · 2026-09-04T13:05:15Z

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

tensorrt-cicd · 2026-09-04T13:12:40Z

PR_Github #71537 [ run ] triggered by Bot. Commit: eb07ced Link to invocation

tensorrt-cicd · 2026-09-04T13:13:20Z

PR_Github #71538 [ run ] triggered by Bot. Commit: eb07ced Link to invocation

tensorrt-cicd · 2026-09-04T13:13:23Z

PR_Github #71537 [ run ] completed with state ABORTED. Commit: eb07ced

Link to invocation

tensorrt-cicd · 2026-09-04T13:59:56Z

PR_Github #71538 [ run ] completed with state SUCCESS. Commit: eb07ced
/LLM/main/L0_MergeRequest_PR pipeline #58626 (Partly Tested) 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-07T02:28:30Z

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

tensorrt-cicd · 2026-09-07T02:34:24Z

PR_Github #71828 [ run ] triggered by Bot. Commit: 6a4e7b0 Link to invocation

tensorrt-cicd · 2026-09-07T03:31:21Z

PR_Github #71828 [ run ] completed with state SUCCESS. Commit: 6a4e7b0
/LLM/main/L0_MergeRequest_PR pipeline #58907 (Partly Tested) 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

yizhang-nv

Approving to unblock main: the disagg-gen slot collapse (AL 1.09) needs to land now,
and the arena sizing fixes are strictly better than what we have.

One note for the record, not blocking: we plan to remove the separate draft KV cache
manager entirely, and this PR makes it the memory owner of the drafter ctx cache and
adds disagg / attention-DP as new consumers of it. When that removal happens the
drafter will fall back to the dense arena unless we move first.

The direction I'd like to take in a follow-up is to append the drafter's L layers to
the unified target manager, the way MTP/Eagle3 already do via
_build_per_layer_num_kv_heads (external drafters only lack this because
get_num_spec_layers returns 0 for them). get_buffers(layer_offset + i, "HND")
gives the same per-layer view you bind today, the block table is
kv_cache_block_offsets directly, and the transceiver would carry those layers to
the disagg gen worker, giving the drafter prompt context instead of an empty slot —
which should close the 3.44 vs 4.32 AL gap. Happy to sync on that separately.

…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

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

Open

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

brnguyen2 added ci: post-merge approved Approved by TRT-LLM CI approvers for broad post-merge CI requests and removed ci: full pre-merge approved labels Sep 3, 2026

longlee0622 added the ci: full pre-merge approved label Sep 3, 2026

nvpohanh requested a review from yizhang-nv September 4, 2026 07:08

dc3671 force-pushed the user/zhenhuanc/dspark-kvmgr branch from eb07ced to 6a4e7b0 Compare September 7, 2026 02:28

yizhang-nv approved these changes Sep 7, 2026

View reviewed changes

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

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!

tensorrt-cicd commented Sep 3, 2026

Uh oh!

dc3671 commented Sep 3, 2026

Uh oh!

tensorrt-cicd commented Sep 3, 2026

Uh oh!

tensorrt-cicd commented Sep 3, 2026

Uh oh!

tensorrt-cicd commented Sep 3, 2026

Uh oh!

ZhanruiSunCh commented Sep 3, 2026

Uh oh!

dc3671 commented Sep 3, 2026

Uh oh!

nvpohanh commented Sep 4, 2026

Uh oh!

dc3671 commented Sep 4, 2026

Uh oh!

tensorrt-cicd commented Sep 4, 2026

Uh oh!

tensorrt-cicd commented Sep 4, 2026

Uh oh!

tensorrt-cicd commented Sep 4, 2026

Uh oh!

tensorrt-cicd commented Sep 4, 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!

yizhang-nv left a comment

Choose a reason for hiding this comment

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

11 participants

dc3671 commented Aug 28, 2026 •

edited

Loading