[TRTLLM-14019][feat] MiniMax-M3 unified KV cache: share Eagle3 draft layers via a sub-page view - #17457
Draft
zheyuf wants to merge 13 commits into
Draft
[TRTLLM-14019][feat] MiniMax-M3 unified KV cache: share Eagle3 draft layers via a sub-page view#17457zheyuf wants to merge 13 commits into
zheyuf wants to merge 13 commits into
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>
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.
Description
MiniMax-M3 + Eagle3 could not share a KV cache manager, so one-model speculative
decoding ran a separate draft manager — and every block-lifecycle feature had
to be duplicated for it or silently go missing: disaggregated transfer dropped the
drafter's prompt KV (−0.14 AL on chat-GSM8K, −0.5 and worse on long-prompt and
multi-turn regimes), prefix reuse never applied to draft blocks (−0.2 to −0.35 AL
on turns 2+ in both aggregated and disaggregated serving), and each gap was a
silent acceptance regression rather than an error.
This PR lets the drafter share
MiniMaxM3KVCacheManagerV2(unified KV cache):supports_shared_draft_layers=True— the drafter's buffers differ inper-block size from every M3 buffer, so V2 storage places them as their own
sub-pages inside the shared slots. Block lifecycle then covers the drafter
natively: a reused, evicted, or disagg-transferred logical block
physically contains the drafter sub-pages. No mirror creation, no transfer
merging, no reuse re-implementation.
MiniMaxM3DraftSubpageView— presents the drafter's pool to the attentionops at 32-token kernel pages inside the 128-token logical blocks
(
tokens_per_block=32; block-table entrys → s·scale·subdiv + jfor K,+subdivfor V). Same slot-geometry pattern as the dense-layer trtllm-genadapter (
trtllm_gen_dense_decode.py). This sidesteps bothtokens_per_block=128Eagle kernel issues on SM103 (missing context cubin →unfused-MHA workspace blow-up; ForGen IMA) with no kernel changes. The
view owns no blocks and rides the existing draft-swap machinery in
TrtllmAttentionMetadata. The draft kernel page size is overridable viaTRTLLM_M3_DRAFT_KV_TOKENS_PER_BLOCK; once the kernels are fixed,=128degenerates the view to the identity expansion (unit-tested), making the
workaround retirement a config flip.
get_draft_kv_cache_managergated on aflag derived from the raw spec config, which disagrees with the manager-level
attention-DP share decision; the draft loop then silently ran without a swap.
Trust the registered resource, fall back to the target manager's
draft_subpage_view(None for managers without one — today's behavior).class descriptor and then accesses the property directly, so a failure inside
view construction propagates instead of silently downgrading to "no view"
(a defaulted
getattrswallowed exactly that during development: a stalefield in a log f-string disabled the view with zero symptoms). Three
[unified-kv]INFO logs pin the mode end to end: shared-layout at managercreation, resolved mode at first dispatch (and on every transition), view
geometry at construction. Sanity asserts pin the view's mega-slot
preconditions (single pool; draft layer resident in pool 0).
nvbugs/5807902 WAR's shared-manager fallback is exactly the path M3 wants in
disaggregated mode.
sparse_layer_idsno longer mislabels appended one-model draft layers assparse layers;
derive_shared_draft_layoutanchors the draft tail on theper-layer heads list (the unambiguous total), correct under both creation
conventions (
num_layers= pretrained target count with spec extension, orpre-extended).
CI coverage (new in this PR)
TestMiniMaxM3::test_nvfp4_eagle3gains adisaggparameter arm: context TP2(block reuse on, chunked prefill, no overlap) → attention-DP generation TP2
(CUDA graphs, overlap, reuse off) on one 4-GPU node, NIXL PYTHON transceiver,
FP8 KV + indexer, Eagle3 one-model draft3 — the AgentX/InferenceMAX serving
shape, asserting GSM8K /
gsm8k_inferencemaxaccuracy through the router. Thedefault-eval disagg point joins
l0_dgx_b200_m3pre-merge and the QA corelist; the aggregated arm is byte-identical to before (ids gain
-disagg=False).Validation (MiniMax-M3 NVFP4 + Eagle3 draft3, GB300)
gsm8k_inferencemaxaccuracy (disagg)Both PR-tip re-runs (aggregated and disaggregated) show all three
[unified-kv]engagement logs on every rank (draft layers [60] share the target KV cache manager (sparse layers 3..59)→draft KV resolved: target manager's draft sub-page view→draft subpage view active: draft_layers=[60] subdiv=4 tokens_per_block=32).Perf receipt — the view's only steady-state work is one block-table
expansion per draft window (
copy_batch_block_offsets): 84 / 341 / 1202 µs atbatch 8 / 64 / 256 (measured in-container, includes the H2D copy) ≈ 0.4–1 % of
a 30–80 ms generation iteration at concurrency 64. Acceptance is unaffected:
a direct-attention control run (view disabled) scores 3.53 vs 3.50 with the
view — the view exists for correctness at 32-token pages, not for AL.
Unit tests (16, pass against real torch/trtllm in-container): sub-page
addressing math (expansion, V offset, BAD_PAGE_INDEX padding, subdiv=1
degeneration, pool-0 sanity), draft-tail layout under both creation
conventions, the manager property itself (construction + logging + caching +
is_draftguard), and dispatch resolution (separate-wins, view fallback,loud propagation of construction failures).
Notes for reviewers
(default
_allow_separate_draft_kv_cache=True); production shapes (attentionDP, disaggregated) take the unified path.
non-ADP one-model users; its transfer/reuse gaps are tracked separately
([TRTLLM-14019][feat] Transfer the separate draft KV cache in disaggregated serving #17371 stays a draft for that audience).
the 341 µs @ bs64; happy to vectorize via a packed numpy staging buffer in
this PR or a follow-up if desired.
PR Checklist
[JIRA/NVBUG/None][type]format/bot run)