Repository navigation
[None][feat] MiniMax-M3.1 DSpark: unified target and drafter KV cache under disaggregated serving - #25
Draft
zheyuf wants to merge 15 commits into
Conversation
Signed-off-by: Allison Lim <allim@nvidia.com> (cherry picked from commit fce28f0) Backport to this branch: the cpp/ part (BufferConfig.coalescing_group, the attentionOp 3-D pool-pointer table) is not included; _coalescing_group_kwargs passes the group only when it is nonzero, so layer specs without NVFP4 buffers run on the existing bindings, and the two mixed NVFP4/BF16 pool tests skip on bindings without coalescing_group. Hunks for main-only features (mixed-precision KV, pool rebalancing) are dropped, and layer_specs is declared on the class for managers built without __init__. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
Port of NVIDIA#20486 at 219e26e (net diff against the main it merges, df95eda). Signed-off-by: Allison Lim <allim@nvidia.com> Backport to this branch: conflicts resolved against the separate draft pool's disagg window transfer (RankInfo keeps both draft fields), the drafter slot parking and seeds in DFlashSpecMetadata.prepare (the position seeding applies to the unified cache only), and the fused block-table refresh (skipped when the layer has a block offset). Hunks for code this branch does not have (spec_tail_kernels, model-level layer specs) are dropped; Gemma4 hybrid geometry stands in for _has_variable_kv_cache_geometry. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
Port of NVIDIA#20709 at 300a671 (net diff against the main it merges, 0746a23). Signed-off-by: Allison Lim <allim@nvidia.com> Backport to this branch: import and registration-check conflicts resolved by keeping both sides. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
DFlash took the number of valid block-table entries as the pages a request holds. On a unified cache with per-layer page tables (e.g. an FP8 target with a BF16 drafter), a sliding-window draft layer's evicted pages stay BAD_PAGE_INDEX in front of its live tail, so the count stopped at about one window (66 blocks for window 2048, page 32) whatever the context length. The decode-time allocation check then trapped once a context passed ~2.1K tokens, and prefill skipped the drafter's context store for later chunks and for prefix hits past the window. Use the extent of the allocation instead, the last valid entry + 1. Reads below the window are never issued and writes land in the live tail; for a full-attention drafter, and on the cpp backend, the extent equals the count. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> (cherry picked from commit ff11364)
CUDA-graph warmup builds every batch from dummy requests whose ids repeat across batches (0..bs-1), so a reused id keeps its drafter slot and the previous batch's context length, which the replay after each capture advances for good. A later batch can give the same id far fewer pages, and on a unified KV cache the drafter's allocation check then trapped at server start: during the advanced-sampling capture pass the filler rows reached 31..79 positions against a 24-position allocation. Only warmup runs a CUDA-graph batch eagerly, so there re-seed each generation row's slot from this batch's positions, as prepare() already seeds a new request. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> (cherry picked from commit 2b23449)
…'s manager under disagg On GB300 a disaggregated generation worker gets the target's prompt KV but not the DSpark drafter's, which lives in a separate pool: the drafter starts from an empty window at block 0 and acceptance collapses on long prompts. Hosting the draft layers in the target's KVCacheManagerV2 lets the transceiver move them with the target's pages. - MiniMaxM3KVCacheManagerV2 hosts explicit draft layer specs: index-K buffers follow the spec-built ones, dense target layers stay FP8, size-only coalescing keeps the target pools, draft layers use the base view. - KvCacheCreator unifies a standalone DSpark drafter with such a manager under disaggregated serving when the drafter reads the target's page size. Agg keeps the separate pool paired with block reuse, and a model that ships that pool's window (F0a) keeps it in disagg. TRTLLM_DISAGG_SHARED_DRAFT_KV=0 opts out. - An NVFP4 target transfers with a BF16 draft held in pools of its own. - A windowed draft lifecycle gets the hot-tier bytes of every admissible sequence's window plus one context chunk, instead of the bytes-per-token split that handed it a quarter of the KV cache. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…w is complete With the drafter's layers in the target's KV cache manager, the draft lifecycle's REQUIRED policy let the drafter's window constrain every prefix match: a request branching off a shared prefix found no committed window there and re-prefilled its whole prompt. - A windowed drafter's layer specs are OPTIONAL. The target alone sets the match, and a new cache's first resume selects the drafter's layers where their window is complete (page-aligned matches only). - Where the window is missing, prepare_context caps the claim so the prefill writes everything the drafter reads (from L + 2 - W on): at L + 1 - W, or one window before a branch point within W of the prompt end, whose window the request then publishes. A partial match first retries at its last full page. - The commit gate for unstored drafter context, its unwritten-prefix marks and branch-point publication cover a manager that hosts the drafter's layers, not only a paired draft pool. A disagg generation request needs no cap: the KV transfer requests an unselected OPTIONAL group in full. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…ck offsets per step A hybrid NVFP4 MiniMax-M3 cache is pool-less: no target layer reads its attention-op block offsets. With explicit draft layers it kept every layer's pool, so each step converted the whole per-layer table (layers x sequence slots x max_seq_len / tokens_per_block) on the device: about 3 ms per decode step at max_seq_len 1M on TP2. Keep one attention-op pool, the first draft layer's, which is the only one the unified drafter reads, and fill it from the draft lifecycle's rows with the copy a separate draft pool uses. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…get's KV cache A unified DFlash/DSpark drafter appends its block after the request's context, through the request's own pages. The manager sized every request with num_extra_kv_tokens = K-1 beyond its history, but the DSpark shift_label block is K wide (DFlash K+1). When the allocation ended on a page boundary, at the prompt end or after a step that accepts all K+1 tokens, the block's last slot had no page: the drafter then truncated the advertised context and its scratch overwrote the last context position in the draft layers, silently. When DRAFT-owner layers live in the manager, size each request beyond its history by max(num_extra_kv_tokens, drafter block width) at every allocation site and in the generation headroom. The creator passes the checkpoint-resolved width (shift_label). Adapted from NVIDIA/TensorRT-LLM Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> NVIDIA#20664 (76c4cef); num_extra_kv_tokens keeps its position-side meaning.
…ot be honored An explicit TRTLLM_DISAGG_SHARED_DRAFT_KV=1 on a disaggregated worker whose target and drafter cannot share one KV cache manager silently kept the separate draft pool, whose KV is not transferred. Raise at KV cache creation instead, as NVIDIA/TensorRT-LLM NVIDIA#20664's _check_shared_draft_kv_request does. disaggregated_mpi_worker now writes "auto" rather than "1" when both sections pass its coarse checks, so its automatic choice can still fall back where a worker's finer checks refuse; the creator shares under "1" and "auto" alike. The unset default is unchanged. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
The launcher wrote TRTLLM_DISAGG_SHARED_DRAFT_KV=0 for any drafter other than DFlash, and for any page size other than 16, 32 or 64, so a MiniMax-M3.1 DSpark deployment started through it (128-token pages) kept the separate draft pool, whose KV is not transferred. Accept a standalone DSpark drafter with attention_backend unset, AUTO or TRTLLM, as NVIDIA/TensorRT-LLM NVIDIA#20664 does, and leave its page size to each worker's KvCacheCreator: the launcher cannot know whether the drafter reads the target's page. It writes "auto"; if the two sides still disagree, the peers' draft-layer check fails loudly. An explicit TRTLLM_DISAGG_SHARED_DRAFT_KV=0 is now kept. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
… slot A generation row with the attention-DP dummy id must never inherit a real slot, because the C++ offset copy maps placeholders to page 0: a slot that a context warmup stored under that id would let the dummy row write its drafter K/V past its own pages, onto another request's page. As in NVIDIA/TensorRT-LLM NVIDIA#20664, such rows now always take the dummy slot. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…ores on the host DFlash bounds each context store by the row's allocation, which it reads back from the device unless the drafter's block-table rows are known to be written in full. That was only recognized for a separate draft pool, so with the unified drafter every context chunk waited for the target's forward before the store was issued. Let the KV cache manager say whether it rewrites copied rows in full, and check the target manager's own table for a unified drafter. The MiniMax-M3 manager's drafter-only pool is copied in full. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…ndpoint An OPTIONAL lifecycle group was checked only at a page-aligned endpoint, so a partial prefix match always left it incomplete. A unified DSpark drafter then lost its window's partial page and the caller cut its claim back to the last full page, recomputing up to a page per request. Check an OPTIONAL group at any endpoint with the REQUIRED coverage rule: the last page needs only the matched prefix. The first resume copies a selected group's partial page into a private page, as for REQUIRED, and counts it as a partial reuse. Partial snapshots now publish OPTIONAL pages as well, so a continuation of a prompt that ended mid-page finds them. A match whose partial page is missing still falls back to the last full page. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
The stock DFlashWorker._refresh_ctx_block_tables reads _ctx_block_cols since the allocation-extent fix; the test's object.__new__ stub never set it. 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
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
On GB300, a disaggregated MiniMax-M3.1-Flash + DSpark generation server received the target's prompt KV but not the drafter's. The drafter keeps its KV in a separate pool that the transceiver does not move, so every generation request drafted from an empty window. Acceptance length (AL) collapsed on long prompts: in 1P1D golden mode, the 1K-12K ladder fell to 1.59 (2.42 aggregated) and 2-16K code to 1.43 (4.40). The drafter-window transfer on this branch (F0a,
disaggregation/draft_kv.py) is enabled only for the sm_107 models.This PR keeps the DSpark drafter's KV in the target's own
KVCacheManagerV2, as explicit draft layer specs with a lifecycle of their own (BF16, 4096-token window, 128-token pages). The existing transceiver then moves the drafter's prompt window with the target's pages, and the generation-side drafter starts at the prompt end. Each run uses one mechanism:This PR includes a C++ change (commit 14,
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp): rebuild the wheel. Updating only the Python package of an installed wheel is not enough.Nothing to enable: ctx and gen servers with the same DSpark config use the unified layout.
TRTLLM_DISAGG_SHARED_DRAFT_KV=0keeps the separate pool, and an explicitTRTLLM_DISAGG_SHARED_DRAFT_KV=1makes a worker that cannot share fail at startup.Commits
BufferConfig.coalescing_groupand attentionOp's 3-D pool-pointer table. M3.1 needs neither._coalescing_group_kwargspasses a coalescing group only when it is nonzero, so the existing bindings accept layer specs that coalesce by size.NotImplementedError, and their 2 tests skip on bindings withoutcoalescing_group.DFlashSpecMetadata.prepare(the position seeding applies to the unified cache only);MiniMaxM3KVCacheManagerV2hosts explicit draft layer specs:KvCacheCreatorunifies a standalone DSpark drafter with such a manager under disaggregated serving when the drafter uses the target's page size. Aggregated serving keeps the separate pool, and so do models that ship the drafter's window themselves (F0a).num_extra_kv_tokens= K-1 = 6 beyond their history, but the DSpark block is K = 7 wide.num_extra_kv_tokens, block width). The width is the checkpoint's (shift_label).TRTLLM_DISAGG_SHARED_DRAFT_KV=1cannot be honored ([None][feat] DSpark Unified KVCache standalone NVIDIA/TensorRT-LLM#20664's_check_shared_draft_kv_request).disaggregated_mpi_workernow writes "auto" for its automatic choice, so a launcher run whose finer per-worker checks refuse still falls back. The unset default is unchanged.disaggregated_mpi_workershare a DSpark drafter's KV. The launcher wrote "0" for any drafter other than DFlash and for pages other than 16/32/64, so an M3.1 DSpark deployment started through it (128-token pages) kept the separate pool.attention_backendunset, AUTO or TRTLLM, as [None][feat] DSpark Unified KVCache standalone NVIDIA/TensorRT-LLM#20664 does.KvCacheCreator. If ctx and gen still disagree, peer registration fails loudly on the draft-layer count.TRTLLM_DISAGG_SHARED_DRAFT_KV=0is kept.kvCache.cpp).Commits 9 and 12 only guard the drafter's K/V, so they protect acceptance length. Speculative decoding stays lossless with or without them: a wrong draft is rejected, so outputs do not change.
Results
Setup:
Serving performance against the branch base 519d279
Measured at commit 13. Production recipe: max_seq_len 1M, 128-token pages, block reuse on, CUDA graphs, overlap scheduler. Prompts are unique after the chat template, so block reuse covers only the template. Commit 14 changes only that reuse, on the context side (see "Context side with block reuse" below); the generation-side numbers here do not depend on it.
ignore_eos.Natural AL (default sampling):
Forced AL 4 (
TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS=3), generation device step on SA prompts:On 8K prompts the forced step is 1.01-1.14x, but the two sides do not compute the same thing. nsys over 20 steps at B16 (1P1D, kernel time per step, rank 0):
=0)Agentic multi-turn. 44 conversations x 6 turns of an agentic coding trace (264 requests, 14.4M prompt tokens), c44, 1P1D TP2, production recipe with block reuse, MXFP8 GEMM backend pinned and autotuner tactics frozen on both sides:
Context side with block reuse (commit 14). 1P1D TP2, SA prompts that share a 187-token template, MXFP8 GEMM backend pinned, frozen tactics. Context tokens and TTFT at the context server:
TTFT at c1 after long requests. 1P1D TP2, 6 sequential requests with unique prompts and 4096 output tokens at forced AL 4, MXFP8 GEMM backend pinned, frozen tactics, measured at commit 13 (commit 14 does not touch this path):
enable_iter_perf_stats: true(production recipe)enable_iter_perf_stats: falseCommits 8-13 against commit 7
Unless a table says otherwise, the results further below were measured at commit 7 with max_seq_len 32K. This part compares commits 8-13 with commit 7.
Production shape, max_seq_len 1M. Generation device step at a forced AL of 4 (
TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS=3); medianprev_device_step_timeof steady generation-only iterations fromprint_iter_log. "Unified" is the default; "separate" setsTRTLLM_DISAGG_SHARED_DRAFT_KV=0on the same build. Batches B1 and B44 run as waves of exactly 1 and 44 requests, withignore_eosand 1024 output tokens. 0 errors.Context side at max_seq_len 1M (commit 13). 1P1D TP2 context server, unified vs separate, same build, forced AL 4. TTFT and prefill (scheduled -> first token) come from the context server's perf metrics; the host step comes from its iteration log.
Two rounds with the nodes swapped, block reuse off, max_seq_len 1M. On the c1 rows the host step is not used, because sporadic iterations include idle time.
AL at 32K. Golden mode, paired per prompt, 95% CI, against commit 7 on the same nodes (last row: against commit 13):
TRTLLM_DISAGG_SHARED_DRAFT_KV=1Commits 1-7
1P1D, ctx TP2/EP2 -> gen TP2/EP2, NIXL + Python transceiver, at commit 7.
TRTLLM_DISAGG_SHARED_DRAFT_KV=0.Golden (rejection sampling + block verification):
Default (rejection sampling off):
A second sample at commit 6 gave 4.051 / 2.465 / 4.363 golden and 3.680 / 2.150 / 3.737 default. The 32-prompt code pass varies by about ±4%.
TEP4 -> TEP4 (two nodes), golden, at commit 7 on the same wheel:
A second, streamed ladder gives 2.475. Per output position, its tokens per step match aggregated serving, with no deficit in the first tokens: 2.87 vs 2.88-2.92 at outputs 1-31, and 2.35 vs 2.32-2.36 at outputs 1024-2047.
Soak at commit 7 (TEP4 -> TEP4 golden, block reuse on at ctx and gen): 4,640 requests in 21 minutes (SA, 14 x 320 mixed stress, SA).
Topology matrix. 17 arms, 0 errors and 0 bad server lines. "Commits 1-6" rows ran on the previous base 20a1a00.
The 0.979 golden ladder of commits 1-6 is sampling noise: commit 7 gives 0.993, and 1.011 when streamed.
Generation step time at a forced AL of 4: unified / separate pool is 1.003 and 1.005 at batch 16, and 1.013 and 0.968 at batch 48-64, over two rounds. That is within run-to-run noise.
KV capacity per rank (free-memory fraction 0.85 on gen and 0.80 on ctx servers, max batch 64, 128-token blocks):
Transfer. The drafter's window adds bytes per prompt token, and transfer time does not grow:
DEP4 -> TEP4 re-shards the KV. Its median transfer is 70-136 / 402-449 / 555-693 ms, against 100 / 539 / 839 ms before.
Block reuse (commit 7). 1P1D TP2, block reuse on at ctx and gen.
per_conversationper_conversationall_reusableall_reusableper_conversation, the short-tail branch recomputes one drafter window plus its tail and publishes the window at the branch point. Without commit 7 it re-prefills the whole prompt.Relation to NVIDIA#20664
NVIDIA#20664 (DSpark Unified KVCache standalone, Allison Lim, head 4cd7c93) is the upstream form of this work for plain
KVCacheManagerV2targets. It is stacked on NVIDIA#20486._check_shared_draft_kv_request, with the launcher writing "auto" as [None][feat] DSpark Unified KVCache standalone NVIDIA/TensorRT-LLM#20664's does;_dflash_unify_blockerrefuses specialized managers ("the target does not use KVCacheManagerV2"), and M3.1 usesMiniMaxM3KVCacheManagerV2.TRTLLM_DISAGG_SHARED_DRAFT_KV.tokens_per_block in (16, 32, 64)it would refuse M3.1's 128-token pages; this PR's condition is that the drafter's page equals the target's._check_block_within_allocationturns both bugs into a device assert: past one window with a sliding-window drafter on per-layer page tables, and at server start in CUDA-graph warmup.Not covered
NotImplementedError. The M3.1 manager coalesces by size and does not need them.disaggregated_mpi_workerstill shares only when both sections use the sametensor_parallel_sizeandenable_attention_dp([TRTLLM-16722][feat] Dflash Unified KVCache NVIDIA/TensorRT-LLM#20486's rule). Workers launched one per server have no such limit; the TP and attention-DP mismatches in the matrix above ran that way.Test Coverage
New tests:
tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py:tests/unittest/_torch/executor/kv_cache/test_unified_draft_reuse.py: the reuse rules on a real unified manager.test_minimax_m3_unified_draft_kv.pychecks the single drafter pool and its rows;test_kvcm2_integration.py::test_shared_draft_pages_cover_the_drafter_block, with the allocation ending on a page edge at the prompt end and after a full-acceptance step. It fails at commit 7;test_dflash_unified_kv_gate.py::test_explicit_shared_draft_kv_request_fails_fast;test_dflash_unified_kv_gate.py::test_launcher_lets_dspark_workers_decide_on_128_token_pages;test_dflash_dummy_slot.py::test_attention_dp_dummy_generation_row_takes_the_dummy_slot;test_dflash_host_ctx_alloc.pyandtest_minimax_m3_unified_draft_kv.py.ReplayCacheTest.OptionalWindowClaimsAPartialEndpointinkvCacheManagerV2StatsTest.cpp, andtest_unified_draft_reuse.py::test_a_partial_match_reuses_the_drafters_partial_page. Intest_a_branch_point_window_is_recomputed_once_and_then_shared, the third request now claims all 43 shared tokens, because the second publishes the drafter's partial page.test_disagg_receive_ordering.pystubs the new_kv_alloc_extra_tokenson itsMock(spec=KVCacheManagerV2)fixture (one line).unittest/_torch/attentioninl0_b200.yml, andunittest/_torch/executorinl0_gb300_multi_gpus.ymlandl0_dgx_b300.yml.Ported upstream tests, with main-only cases adapted:
test_dflash_unified_kv_gate.py,test_kvcm2_integration.py,test_kv_cache_estimation.py,test_peer.py,test_rank_info.py,test_extractor.py,test_bounce.py.Unit runs on GB300, base 519d279 vs this PR on the same wheel. They cover CI's
unittest/_torch/executor,unittest/disaggregatedandunittest/_torch/speculative/hw_agnostic, plus the touched attention / MSA files.Not run to completion on either side:
test_hang_detector_kill.pySIGKILLs its process group.test_py_cache_transceiver_mp.py's MPI cases end the job step on this cluster. Its first 4 cases fail on both sides.test_transfer_lifecycle_*files, the cases that ran fail identically on both sides, at several minutes each. The rest were not waited for.test_world_ranks_converge_on_hostless_fallbackwaits for an MPI world.Commits 8-13 against commit 7, on the same wheel. No failure ID is new in any of these runs:
test_disagg_receive_ordering.py;test_cancelled_admission_releases_eventfails at commit 7 and at the head.Commit 14 on its own wheel:
kvCacheManagerV2*andradixBlockTreepass (199 tests).unittest/_torch/executor/kv_cacheran before a test-only amendment of commit 14. Every failure also fails on the base, with two exceptions:test_world_ranks_converge_on_hostless_fallback, which needs an MPI world (see above), and the branch-point reuse test, whose expectation the amendment updates.test_unified_draft_reuse.pyandtest_paired_draft_reuse.pypass (17 tests).Release Checks, at commit 7 for the first three points:
release_check.py'sbandit -r tensorrt_llmreports only the base's 10 existing issues, none in a file this PR touches.test_disaggregated_serving.py.PR Checklist
🤖 Generated with Claude Code