Repository navigation
[None][feat] Run the Muse modeling_v2 composition with a TRTLLM DFlash drafter - #20812
Merged
zheyuf merged 6 commits intoOct 28, 2026
Conversation
Merged
1 task done
zheyuf
force-pushed
the
musespark-open/mv2-trtllm-drafter
branch
3 times, most recently
from
October 28, 2026 03:58
6009ea9 to
935a207
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #81069 [ run ] triggered by Bot. Commit: |
… paged context The drafter's window-attention ops, dflash_window_attention and its multi-request form, read the DFlash context from a VANILLA drafter's dense per-slot arena. A TRTLLM drafter keeps that context in the KV cache manager's pages instead, which is what lets disaggregated serving transfer it with the target's KV. Given a page table, both ops now read and append the context through it: the caches are pools of pages and each request's batch index names its table row. The kernels see a pool as one space of rows, so a paged call differs from a dense one only in which row a column maps to, and on the same values the two are bitwise equal. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
… of its prepare prepare_1st_drafter_inputs stored the step's context K/V, committed it (lengths, dummy slot, window compaction), read the lengths the drafter attends with the allocation clamp, and built the drafter's inputs, all inline. A fast path that schedules part of that elsewhere -- the Muse modeling_v2 composition stores the context on a side stream, for one -- had to copy those rules, and every worker change then had to be made twice. Name each stage as a worker method: _store_step_context, _commit_context (_advance_ctx_len, then _settle_context: the window compaction and _drafter_context_lengths), _empty_dummy_slot and _drafter_inputs. The built-in prepare is those stages in order: the same operations on the same operands, so its results are unchanged. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…Flash worker The verify fold and the context-K/V fork carry their own copies of the DFlash worker's context commit, and the copies have drifted from the worker: - After every length update the worker empties the dummy slot, which CUDA-graph padding rows share: they attend no context, and their next context K/V starts inside the padding request's own pages. The fork's commit kernel and the fold's store did not, so the dummy slot's length grew by the padding rows' accepted counts every step. Over the KV cache manager's pages a padding row then stores past the padding request's pages, onto page 0 of another request. - On a unified target/draft KV cache the worker checks the drafter's block against the request's allocation before clamping the lengths it advertises; the fold clamped without the check. Rather than patch the copies again, both now run the worker's stages: the fold stores and commits with _store_step_context and _commit_context, the fork's side chain stores with _store_step_context, and the fork's commit kernel keeps only the part it fuses, the zeroing of the masked rows, before the worker's _commit_context. Over a windowed arena that drafter_window_compact compacts, the side stream advances the lengths and the prepare settles the context after the join, where the compaction ran before; the sm107 windowed commit still advances and compacts in its one launch, then applies the worker's dummy-slot rule and reads its lengths. The drafter's inputs are the worker's _drafter_inputs. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
Collaborator
|
PR_Github #81069 [ run ] completed with state |
…drafter context The musespark_open TEP composition required the VANILLA DFlash attention backend, so a TRTLLM drafter, whose context lives in the KV cache manager's pages, dropped every fast path, the target's included, and `modeling_v2: require` refused to start. Disaggregated serving needs that drafter: only then is the drafter's context transferred with the target's KV. The plan now admits either layout. A paged context is never windowed, so over pages the plan keeps every path but the window compaction (47 of the 48 sm103 TEP4 paths) and says so in its log line. Every path that touches the drafter's context reads it where the worker keeps it: - the block decode hands each layer's view of the manager's pool and the worker's one page table to the window ops; - the context-K/V fork's side chain decodes this step's page table at the fork point, which the worker then skips, and stores with the worker's own store; its commit is the worker's alone, since the drafter's block append overwrites the rows past the accepted count before anything reads them; - the verify fold takes the fork's paged steps. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
The Muse Spark Open modeling_v2 composition now runs over a paged drafter context, so AUTO no longer pins VANILLA under modeling_v2: it follows the checkpoint rule as on any other deployment, and a disaggregated deployment then gets the drafter whose context is transferred with the target's KV. Only a disaggregated worker with TRTLLM_DISAGG_SHARED_DRAFT_KV=0 still gets VANILLA. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…seSpark pre-merge stage Pre-merge CI on feat/musespark-open-agentx runs only the MuseSpark stages, so the tests of the composition over a TRTLLM drafter's paged context did not run on this branch. They need a GPU and none is cpu_only, so they go to l0_gb300_musespark. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
zheyuf
force-pushed
the
musespark-open/mv2-trtllm-drafter
branch
from
October 28, 2026 08:08
935a207 to
8af51a7
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
zheyuf
marked this pull request as ready for review
October 28, 2026 08:09
zheyuf
requested review from
a team,
Wanli-Jiang,
WeiHaocheng and
litaotju
as code owners
October 28, 2026 08:10
zheyuf
requested review from
JunyiXu-nv,
QiJune,
hanjingtian,
kaiyux,
laikhtewari and
weiminwang-nv
October 28, 2026 08:10
zheyuf
removed request for
a team,
JunyiXu-nv,
QiJune,
Wanli-Jiang,
WeiHaocheng,
hanjingtian,
kaiyux,
laikhtewari,
litaotju and
weiminwang-nv
October 28, 2026 08:10
zheyuf
added
the
api-compatible
Collaborator
|
PR_Github #81143 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #81143 [ run ] completed with state |
zheyuf
merged commit Oct 28, 2026
ff0697b
into
NVIDIA:feat/musespark-open-agentx
15 of 19 checks passed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
@coderabbitai summary
Description
In disaggregated serving the generation worker's drafter needs the prompt's draft KV. Since #20806 it is transferred with the target's KV, but only for a TRTLLM drafter, whose context lives in the KV cache. The Muse Spark Open modeling_v2 composition (#20770, #20844) accepted only the VANILLA drafter, whose private arena is not transferred, so 1P1D lost acceptance (coding AL 4.13 aggregated vs 3.91).
This PR:
context_kv_forkandverify_foldcall them instead of copying the worker's rules. This removes two copies that had drifted:drafter_window_compact.modeling_v2.Usage: give both workers the same
speculative_config; nothing else.attention_backend: VANILLAkeeps the previous behavior.Vera Rubin: AUTO now also runs Xiao's VR paths over pages. This is untested on VR hardware. GB300 ran the R200-geometry kernels, and the sm107 windowed commit is covered by a unit test only.
Results on GB300 (Muse Spark Open NVFP4, TEP4, RadixArk DFlash 2, K=7, T=1 rejection sampling), with the composition on both workers and no env:
attention_backend: VANILLA, step time is unchanged.Test Coverage
test_modeling_v2_drafter_context.py, onl0_gb300_musespark. Its 11 cases cover:KVCacheManagerV2unified cache vs the VANILLA arena, bitwise;test_dflash_draft_kv.py,test_dflash_dummy_slot.py.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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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