[None][feat] MiniMax-M3 eagle+disagg pipe clean: share Eagle3 draft layers via a sub-page view - #17457
Merged
pcicotti merged 32 commits intoAug 12, 2026
Conversation
…r when draft layers are shared get_draft_kv_cache_manager (both the worker method and the utils helper) gated on use_separate_draft_kv_cache, a flag derived from the raw speculative config. Under attention DP the manager-level decision (_should_create_separate_draft_kv_cache) can choose SHARED draft layers while that flag still reads True, so the worker looked up a DRAFT_KV_CACHE_MANAGER resource that was never registered and silently ran the draft loop without a swap. Harmless for same-geometry drafters, but wrong whenever the shared manager exposes a draft-specific attention view. Trust the registered resource as ground truth instead: return the separate manager when one exists, otherwise fall back to the target manager's draft_subpage_view (None for managers that do not expose one, preserving today's behavior). Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…layers via a sub-page view MiniMax-M3 was the only mainstream model forced onto the separate draft KV cache manager: its slot-coalesced pools (K/V/index-K of every layer packed per slot) break the uniform-stride assumptions of the generic AttentionOp tables, and the Eagle3 kernels are only healthy at 32-token pages on SM103 while the MSA target requires 128. The separate manager duplicated the drafter's block lifecycle, and every lifecycle feature had to be re-implemented against it (disaggregated transfer, prefix reuse, conversation retention) or silently go missing. Let the drafter share the target manager instead: - supports_shared_draft_layers=True: the drafter's buffers differ in per-block size from every M3 buffer, so V2 storage gives them their own sub-pages inside the shared slots; block lifecycle (radix reuse, eviction, disaggregated transfer) covers them natively, since a reused or transferred logical block physically contains the drafter sub-pages. - MiniMaxM3DraftSubpageView presents that pool to the attention ops at 32-token kernel pages inside the 128-token logical blocks (tokens_per_block=32, block table entry s -> s*scale*subdiv + j for K and +subdiv for V), following the same slot-geometry pattern as the dense-layer trtllm-gen adapter. This sidesteps both tokens_per_block=128 Eagle kernel issues (missing SM103 context cubin -> unfused-MHA workspace blow-up; ForGen IMA) with no kernel changes. The view owns no blocks; it rides the existing draft-swap machinery in the attention metadata. - sparse_layer_ids no longer mislabels appended one-model draft layers as sparse (derived from the per-layer kv-heads tail). - The nvbugs/5807902 WAR exemption from NVIDIA#17341 is retired: with shared draft layers the WAR's shared-manager fallback is exactly the path M3 wants in disaggregated mode. Validation (M3 NVFP4 + Eagle3 draft3, GB300 disagg unless noted): GSM8K AL 3.44 disagg / 3.54 agg (separate-manager refs 3.47 / 3.47); gsm8k_inferencemax accuracy 95.45 (reference 95.46, gate 94.0); multi-turn copy-task turns 2-4 3.57-3.67 (stock separate 3.03-3.11); AgentX 256K real traces AL 3.005; SPEED-Bench protocol 3.25/3.23 vs golden 2.83/2.93. Step-resolved probes show early-decode acceptance matching the fixed separate path, confirming drafter KV transfers natively. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
… the sub-page table build Self-review follow-ups on the unified KV cache change: - The manager no longer infers appended one-model draft layers from the per-layer kv-heads list tail (fragile for models with heterogeneous target heads). The creation site knows exactly how many draft layers it appended and now passes num_one_model_draft_layers explicitly; the base KVCacheManagerV2 ignores it. - The two draft-manager lookup sites shared duplicated fallback logic; both now delegate to a single resolve_draft_kv_cache_manager helper. - MiniMaxM3DraftSubpageView.copy_batch_block_offsets builds the sub-page table with one fused expansion over a cached padded buffer instead of per-row tensor ops, keeping the per-step host cost flat in batch size. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…d target layer count The per-layer kv-heads list can reach _create_kv_cache_manager either built in place (scalar heads + _build_per_layer_num_kv_heads) or pre-extended by the caller, in which case num_hidden_layers already includes the draft tail and the length delta reads zero. Anchor the count on config.num_hidden_layers (the target layer count), which is unambiguous in both flows; masked/cross managers yield a non-positive delta and correctly report zero. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…ds list num_layers at the creation site carries the pretrained target count (60) while the per-layer num_kv_heads list is already draft-extended (61); deriving the draft tail as range(num_layers - num_draft, num_layers) marked target layer 59 as draft, dropped its index-K allocation, and crashed startup. Extract derive_shared_draft_layout() anchored on the heads list's length (the unambiguous total) and cover both creation-site flows with unit tests. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…vfp4_eagle3 CI coverage for Eagle3 + the unified draft KV cache across the disaggregated transceiver. The new disagg parameter arm launches context TP2 -> attention-DP generation TP2 on one 4-GPU node with the AgentX/InferenceMAX serving shape (NIXL PYTHON transceiver, context block reuse, FP8 KV and indexer, CUDA graphs + overlap on gen) and asserts GSM8K / gsm8k_inferencemax accuracy through the router. The aggregated arm is unchanged and keeps the acceptance-length checks. The default-eval disagg point joins l0_dgx_b200_m3 pre-merge and the QA core list. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…draft page size Log the shared draft layout at manager creation (the counterpart of the first-dispatch view log), assert the view's mega-slot preconditions (single pool, draft layer resident in pool 0), and make the draft kernel page size overridable via TRTLLM_M3_DRAFT_KV_TOKENS_PER_BLOCK so retiring the 32-token workaround after the Eagle tokens_per_block=128 kernel fixes is a config flip; the subdiv=1 degeneration and the pool sanity check are unit-tested. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
Pure-logic tests for resolve_draft_kv_cache_manager: a registered separate manager wins, a shared target manager falls back to its draft_subpage_view, and a plain shared manager resolves to None. Guards the regression where a stale worker-level flag disagreed with the manager-level share decision and the drafting loop lost its manager. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…spatch Completes the unified-KV triage trio: manager creation logs the shared layout, first resolution logs which draft-side manager won (separate / sub-page view / none), and view construction logs its geometry. A silent fallback to direct shared-manager attention is now visible in startup logs. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…ion test The vectorized block-table expansion clamps unallocated tail slots to slot 0, so pad entries tile that slot's sub-pages instead of staying zero; kernels never read past the row's real block count, so the pads are inert (the BAD_PAGE_INDEX test encodes the same property). Full suite now passes against real torch: 13/13. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
… breaks resolve_draft_kv_cache_manager fetched the view with a defaulted getattr, which also swallows AttributeError raised inside view construction — a broken view silently downgraded to 'no view' and the drafter attended the shared manager at the wrong page size. Probe the class descriptor first, then access the property directly so construction failures propagate. The resolved-mode log now records every mode transition instead of only the first resolution. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…raft sub-page view The view-active log statement inside the draft_subpage_view property still referenced _draft_pool_ids, a field the sub-page-table refactor removed from the view. The f-string raised AttributeError through the view's __getattr__ delegation, the resolver's defaulted getattr swallowed it, and every resolution since the refactor silently downgraded to 'no view' — the drafter attended the shared manager at tokens_per_block=128 instead of through 32-token sub-pages. Log the view's real fields, and cover the property itself (construction, the log statement, caching, the is_draft guard) with a unit test; direct view construction in the existing tests had bypassed the property and missed this. Verified against real V2 storage: view constructs with draft_layers=[60], subdiv=4. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
Replace the per-row torch.as_tensor fill (~5 us/row) with a numpy padded fill parsed at C speed, and write the fused expansion into a persistent pinned staging buffer so the per-call stack/reshape allocations disappear and the H2D copy is genuinely asynchronous. copy_batch_block_offsets: 84/341/1202 -> 57/123/342 us at batch 8/64/256 (measured in-container, H2D included). Semantics pinned by the existing unit tests (expansion, V offset, pad clamping, subdiv=1 degeneration): 16/16 pass against real torch. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
Single source of truth on draft_manager_tokens_per_block (two-step plan: env flip to 128, then delete the view class, property, and add_dummy_requests override — a P128 control run with fixed kernels and the view disabled showed acceptance parity); the three WAR sites carry a greppable pointer to it. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…plify resolution Review feedback: the resolver's class-descriptor probe existed only because draft_subpage_view was a property (attribute access executes it, so a defaulted getattr would swallow construction failures). get_draft_subpage_view() is fetched by getattr without executing, so failures propagate from the call naturally — the probe, its five-line comment, and one branch disappear, and the lazy construction + logging side effects now live in an honest method instead of a property. 16/16 unit tests pass against real torch. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
Reviewer-facing polish, no behavior change: fix the view docstring's two stale statements (the pool pointer is re-rooted at the drafter's K and the mapping row rewritten — not reused verbatim — and the block-table formula is the general s*scale*subdiv+j form, not the two-page special case); tighten the WAR/retirement comments and kwarg rationale; merge the duplicated mega-slot comment in the view constructor and rename gid to layer_id; split the resolver's wrapped log ternary into a plain if/else; guard the creation log against an empty sparse range. 16/16 unit tests pass in-container. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
Review feedback: the resolved-mode log earned its keep as a discriminator while view resolution could silently return None; with get_draft_subpage_view now a method whose None cases are exactly the complement of the creation log's gate, the outcome is implied by the two M3-scoped logs (creation layout + view construction) or surfaces as a loud failure. Dropping it also removes a module-level mutable and stops tagging every non-M3 speculative deployment's logs with [unified-kv]. 16/16 unit tests pass in-container. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…le field draft_layers duplicates the creation log three lines above and subdiv is 128/tokens_per_block; tokens_per_block stays because it is the env-controlled knob (TRTLLM_M3_DRAFT_KV_TOKENS_PER_BLOCK) whose effective value the line confirms. Also stops reaching into the view's private _subdiv from the manager. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…rameter Review feedback: the class already declares its own kwargs (sparse_layer_ids etc.) as named parameters; num_one_model_draft_layers was the lone stringly-typed pop. As an explicit keyword parameter the signature documents it, typos fail loudly, and the pop dance disappears. The remaining kwargs.get calls are annotated for what they are: peeks at base-owned arguments that super().__init__ must still receive. Verified by constructing the real manager in-container (kwarg binds to the parameter; view constructs; 16/16 unit tests). Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
… configs Mirror every path-gating knob of the production disaggregated sweep (fix15a production-candidate, GB300) in the disagg arm: TRTLLM MoE, MAX_UTILIZATION scheduling, autotuner, fp8 KV on V2 pages of 128, full NIXL PYTHON transceiver settings, overlap on both server roles, context-side conversation-affinity routing config, generation-side postprocess workers / stream interval / pinned CUDA-graph batch sizes, and perf-metrics collection; workload-tuned sizes (1M-context batch and token budgets, GB300 memory fractions) are CI-adjusted with the production values called out inline. The generation server now honors the attention_dp parameter, so both production flavors have ids: attention-DP (the AgentX submission shape, in the test lists) and TEP (the production-candidate shape, enablable by listing it). Test lists slim to one entry per path: aggregated ADP-off inferencemax (separate draft manager + acceptance asserts) and disaggregated ADP-on GSM8K (unified view + transfer). Executed once end to end: PASSED in 6 min 55 s with the engagement logs present. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…ared drafters The equal-heads shortcut in _build_per_layer_num_kv_heads returned a scalar when the drafter's KV head count matches the target's, so the appended draft layers were invisible to the shared manager: MiniMax-M3 routed its GQA Eagle head (4 KV heads, like the target) through the target's fp8 attention machinery, which has no spec-dec kernel for that shape on SM103, instead of through the draft sub-page view. The extended list is what carries the draft tail; return it whenever a draft config is present. Unit test covers the uniform-list layout. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…t disagg acceptance Supersedes NVIDIA#17406. The drafter becomes Inferact/MiniMax-M3-EAGLE3-GQA in both arms, with thresholds recalibrated from the first GQA run: disagg chat-GSM8K measures length 2.773 / rate 0.591 (consistent with the GQA card's MT-Bench 2.668; the MHA head measured 3.44-3.48 here), so the MHA-derived floors do not carry over. The disagg arm now also asserts acceptance through the generation worker's /metrics iteration stats (enable_iter_perf_stats; the fixture's DuckLLM gains a router_url field to locate the worker via /cluster_info) — closing the gap where a drafter-KV transfer regression was invisible to the accuracy assert. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…ollow-up The GQA head's acceptance is genuinely lower (first runs: disagg chat-GSM8K length 2.73-2.77 / rate 0.58-0.59, consistent with its card's MT-Bench 2.668, vs 3.44-3.48 for the MHA head), so switching heads and recalibrating thresholds deserves its own review alongside the serving-shape decision. This PR stays on the MHA head with its established calibration (disagg floor 3.38 between healthy 3.44-3.48 and the broken-transfer plateau 3.33; aggregated thresholds from the drafter card). The equal-heads fix and the disagg acceptance assertion stay: both are head-agnostic. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…l workload The /metrics assertion measures acceptance on the accuracy eval itself (few-shot completion-format GSM8K, 256-token budget), where the chat-trained drafter accepts less than on chat-formatted prompts: the MHA head measures 2.97/0.66 here vs 3.44-3.48 on the chat-formatted smoke workload the previous floor was calibrated on. Set the floors as a catastrophic-breakage guard (corrupted or missing drafter KV collapses acceptance toward 1.x) and note that transfer-regression sensitivity needs a chat-format probe as a follow-up. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
zheyuf
changed the title
…rotocol Symmetric with the aggregated entry: the disagg test-list point moves from default GSM8K to gsm8k_inferencemax, so both the accuracy assert and the /metrics acceptance measurement ride the chat-formatted InferenceMAX workload — the drafter's training distribution, where the MHA calibration holds (healthy 3.44-3.48, broken-transfer plateau ~3.33). The acceptance floor is now selected per eval mode: 3.38/0.76 under inferencemax (transfer-regression sensitive), 2.5/0.5 under default (catastrophic guard for the out-of-distribution completion-format workload, which measures ~2.97). Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…ts own protocol First measurement on the InferenceMAX protocol: acceptance 3.083 / rate 0.694 (accuracy 95.53 vs reference 95.46) — each prompt format has its own acceptance level for the chat-trained drafter (plain chat 3.44-3.48, InferenceMAX 5-shot multiturn 3.08, few-shot completion 2.97). The floor moves between the measured healthy value and a broken-transfer plateau estimated by transferring the chat-workload healthy-to-broken delta; refine when a broken baseline is measured on this protocol. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
The resolver refactor moved the resource lookups from the worker method into resolve_draft_kv_cache_manager, leaving this import dead (caught by the autoflake pre-commit hook). Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
zheyuf
removed request for
JunyiXu-nv,
SimengLiu-nv,
StanleySun639,
chuangz0,
fredricz-20070104,
pcastonguay,
pengbowang-nv,
ruodil,
xinhe-nv and
yizhang-nv
August 11, 2026 08:55
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #65251 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #65253 [ run ] triggered by Bot. Commit: |
BowenFu
reviewed
Aug 11, 2026
Collaborator
|
PR_Github #65251 [ run ] completed with state |
zheyuf
changed the title
Collaborator
|
PR_Github #65253 [ run ] completed with state |
brb-nv
approved these changes
Aug 11, 2026
brb-nv
left a comment
brb-nv
left a comment
Collaborator
There was a problem hiding this comment.
Approving with a comment. Please take care of it if it's a legit concern.
4 tasks
…nned source Review catch (brb-nv): the sub-page view staged its expansion in a persistent pinned buffer refilled in place, then used it as the source of an asynchronous H2D copy — the exact pattern nvbug 6293536 was filed for and that KVCacheManager._stage_block_offsets_for_copy documents. The source is read at copy execution time, so with the overlap scheduler (which the new CI point pins on the generation server) the next iteration's refill can clobber a still-pending copy and the drafter indexes another batch's blocks — silent acceptance loss or an IMA, intermittent and load-dependent. Allocate the staging buffer per call instead and let the caching host allocator hold it until the copy retires; the vectorized expansion moves into _host_block_table unchanged. The numpy scratch stays persistent (only ever read synchronously). Cost: 123 -> 136 us per call at batch 64, still under 0.1% of an iteration. A unit test pins the invariant: two consecutive calls get distinct buffers and the first keeps its own values. Also takes BowenFu's suggestion to read the generation worker's address from /cluster_info's current_workers rather than server_lists. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #65407 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #65407 [ run ] completed with state |
brb-nv
pushed a commit
to brb-nv/TensorRT-LLM
that referenced
this pull request
Aug 25, 2026
…ayers via a sub-page view (NVIDIA#17457) Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
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
5 participants
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.
Description
This PR mainly does two things:
Todo in future PR: test on GQA and change the test case to GQA tests.
PR Checklist
[JIRA/NVBUG/None][type]format/bot run)