[TRTLLM-14093][feat] Eagle3 for M3 MSA path + CUDA graphs compatibility - #3
Merged
Merged
Conversation
…V2 targets A structurally-V2 target (sparse-attention managers like MiniMax-M3's, which don't set use_kv_cache_manager_v2) paired with a plain-transformer draft (Eagle3) resolved the draft manager to V1, requiring V1-family special cases in KVCacheManagerV2.add_dummy_requests and leaving a latent AttributeError: KVCacheV2Scheduler calls suspend_request() on the draft manager, which only exists on V2. Promote the draft class to KVCacheManagerV2 whenever the target is V2 (shared helper used by both the creation and the cache-cost estimation paths, which previously disagreed). add_dummy_requests reverts to its original V2-only shape - kv_cache_manager_v2.py returns byte-identical to its pre-enablement upstream state. The V1-draft branches were reachable by exactly one configuration - MiniMax-M3 + one-model Eagle3; every other V2-target spec config already resolves a V2 draft (flag or target-config fallback) or shares the target manager. The V2-draft combination was validated on 4xB200: MMLU 84.84 / GSM8K 90.07 (explicit flag), plus boot+acceptance probe on the new default path (AR 0.465/AL 2.40, matching the V1-draft band). Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com>
zheyuf
force-pushed
the
feat/m3-msa-eagle3-graphs
branch
from
July 13, 2026 23:49
c51e9d6 to
4cbe580
Compare
… verify Removes the MSA+speculative-decoding restriction: the MSA kernel path (sparse_use_msa=True) now verifies draft tokens through a multi-token decode driver instead of rejecting at creation. Driver (decode_wrapper): hybrid scheme, validated bit-exact against the eager fmha_sm100 api. The proxy (OnlyScore) pass runs natively multi-token: the kernel's inclusive causal bound with offset = kv_len - qo_len produces exactly the verify ladder (token t attends kv_len - qo_len + t + 1 positions). Top-k selection gets per-token ladder valid-page counts, so a draft token cannot select blocks past its own attend bound. The sparse GQA pass row-expands each token to a qo_len=1 pseudo-row keeping its request's full kv_len and page-table base, with the ladder in the per-row offset - the same transform the eager api applies internally; the kernel consumes kv_block_indexes per row, so native multi-token sparse is not expressible. All device ops, capture-safe. Routing/metadata: on MSA, pure-generation uniform multi-token batches stay DECODE-shaped (decode_qo_len = 1 + draft_len), keeping the captured/overlap-safe device-plan path; mixed context+gen batches take the eager extend path; the reference backend is unchanged. KV slot staging generalizes to the causal ladder, and dense layers 0-2 get a ladder SDPA mask (identical to the old math at qo_len=1). Overlap-correction fix: the MSA hook now re-stages the flat page table on mixed batches, not just the CPU length mirrors - a correction that shrinks a row across a page boundary changes the indptr layout the eager kernels rebuild from the corrected lens, misbasing every subsequent row's pages (symptom: MMLU passes, GSM8K collapses). Tests: driver bit-diff suite extended with qo_len=4 legs - proxy bit-diff vs the eager api, top-k vs a per-token ladder reference, and CUDA-graph capture/replay with mutated lens. test_nvfp4_eagle3 runs the MSA path (use_msa single-choice parametrize); QA-list rows updated to the new test IDs (the old rows no longer matched any collected test). Validation (4xB200, TP4/EP4, NVFP4): MSA adp=False MMLU 84.97 / GSM8K 90.18, AR 0.698 / AL 3.094; MSA adp=True 85.21 / 91.24, AR 0.718 / AL 3.153; reference regression control 85.04 / 90.33, AR 0.707 / AL 3.121. MSA suite wall-clock 6:10 vs reference 11:50 on the same gates. Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com>
zheyuf
force-pushed
the
feat/m3-msa-eagle3-graphs
branch
from
July 13, 2026 23:54
4cbe580 to
aa4c9af
Compare
zheyuf
marked this pull request as ready for review
July 14, 2026 00:06
zheyuf
force-pushed
the
feat/m3-msa-eagle3-graphs
branch
3 times, most recently
from
July 14, 2026 02:12
87371b7 to
4ed35a9
Compare
The MSA decode driver was built capture-safe (device-only replan, prepare-time buffer allocation, stable data_ptrs), so multi-token verify no longer routes through the eager extend path on the MSA backend. Allow cuda_graph_config with sparse_use_msa=True; the reference path keeps raising. - speculative/eagle3.py: save/restore kv_lens_cuda across CUDA-graph warmup iterations. The one-model worker saved _seq_lens and the spec-decoding tensors, but not kv_lens_cuda, which the draft loop mutates in place - so the runner's pre-capture warmup iterations ran with drifted kv lens. Same pattern dflash/pard use; active only during graph warmup, not capture. - py_executor_creator.py: the graphs+spec rejection now applies only to the reference path (its verify routes through the eager extend path). - Capture hardening: the dense layers 0-2 decode branch now expands GQA K/V per KV head instead of all heads at once (bitwise-identical math; the expansion is captured into the graph pool, and under attention DP - unsharded heads - the whole-tensor transient exceeds the pool budget at large graph buckets). They also baked the prepare-time host upper bound max_seqlen_k into the captured gather/mask width; replays whose kv_len outgrew the capture-time value would silently truncate attention. Under graphs, bake min(page-table capacity, engine max_seq_len) instead (raw capacity alone OOMs during KV estimation; the seq_lens mask already invalidates positions past each row's true length). resolve_decode_state now raises if it would BUILD a decode state (JIT + allocation) while the stream is capturing. - test_nvfp4_eagle3: cuda_graph single-choice parametrize; the gated variants run the endgame config (MSA + Eagle3 + overlap + graphs). Support matrix updated (EAGLE-3 Linear: Yes). Probe (4xB200, TP4/EP4, batch-1, overlap, greedy): 25.98 tok/s with graphs vs 13.70 eager (1.90x), AR/AL in band; accuracy gate under graphs (adp=False): MMLU 84.62 / GSM8K 90.86, AR 0.719 / AL 3.156 - identical to eager; spec-dec graph capture confirmed on all ranks (draft_len=3 buckets); negative control: reference path + graphs + spec still raises at creation. Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com>
zheyuf
force-pushed
the
feat/m3-msa-eagle3-graphs
branch
from
July 14, 2026 02:54
4ed35a9 to
c8b09eb
Compare
zheyuf
changed the title
zheyuf
changed the title
zheyuf
changed the title
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
1 participant
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.
This PR contains three commits:
Validation (4×B200, TP4/EP4, NVFP4, all at this PR's commits)
Non-spec references: MMLU 85.11 / GSM8K 91.32. Batch-1 greedy decode:
graphs 25.98 tok/s vs eager 13.70 (1.90×).