[zheyuf]

Description

This PR moves aggregated serving of MiniMax-M3.1-Flash with a standalone DSpark drafter to the unified KV layout that disaggregated serving uses, so one mechanism serves both. It is stacked on the unified-KV disagg PR (zheyuf/brb-dspark-unified-disagg-519-v2); only its last commit is new. The section below describes the change, its results and its tests.

Aggregated serving on the unified KV layout

On GB300, disaggregated serving keeps a standalone DSpark drafter's KV in the target's KVCacheManagerV2, as explicit draft layer specs with a lifecycle of their own: BF16, 4096-token window, 128-token pages, and a window-sized share of the KV budget. Aggregated serving kept #22's separate draft pool, paired with the target's block reuse. It now uses the unified layout too:

• KvCacheCreator._should_unify_dflash_kv_cache no longer requires disaggregated serving for a manager that hosts draft layers (MiniMaxM3KVCacheManagerV2).
• What the paired pool did carries over. The drafter's window is reused only where it is complete (the disagg PR's reuse commit), unwritten drafter pages are never published, and the draft share is window-sized for the same admissible sequences (two per batch slot with the overlap scheduler).
• A model that ships the separate pool's window itself (F0a, the sm_107 targets) keeps that pool, in both modes. TRTLLM_DISAGG_SHARED_DRAFT_KV=0 stays a disagg-only opt-out.
• Tests: the unify-gate test and the window-sized share test now cover aggregated serving too.

Validating aggregated serving on the c44 regression config (TP2, max_seq_len 1M) found three costs of the unified layout, which hit disaggregated serving as well. Three commits of the disagg PR fix them:

• "copy only the unified drafter's block offsets per step". The hybrid NVFP4 M3 manager kept one attention-op pool per layer, so every step converted each layer's block-table rows at full width (65 layers x 177 rows x 8,193 blocks at 1M). That cost about 3.5 ms per decode step at batch 1 on TP2. The manager now keeps only the drafter's pool and fills it with the copy a separate draft pool uses.
• "size a unified MiniMax-M3.1 drafter's context stores on the host". DFlash read each context row's allocation back from the device unless it recognized the row writer, so every context chunk waited for the target's forward.
• "KVCacheManagerV2: reuse an OPTIONAL group at a partial endpoint". The drafter's lifecycle was complete only at a page boundary, so a prefix hit that ended inside a page recomputed up to a page (see "Partial pages" below).

Results

Setup:

• GB300, MiniMax-M3.1-Flash NVFP4 + DSpark K7. The wheel is built from 519d279, with this PR's Python on top, as in the disagg PR.
• "Before" is the same build with an experiment-only switch, not part of this PR, that restores the paired separate pool in aggregated serving. Before and this PR always run on the same node.
• The tables were measured on the stack before the partial-endpoint commit, which changes C++. "Partial pages" below re-measures the agentic replay and the c1 step pass on a wheel built at that commit.
• AL comes from the server's counters. Every run had 0 client errors and 0 bad server lines.
• Two configurations:
  • TP4/EP4: the disagg PR's AL config (max batch 64, max_seq_len 32K, or 131K with block reuse).
  • The c44 regression config tot-regression-c44-tp2-s6: TP2, per_conversation block reuse, host cache, max batch 88, max_seq_len 1M, ARIA env. Two deviations: indexer_kv_dtype: fp8 (our image lacks the MiniMax MSA package) and an 89 GB host cache per rank (two TP2 servers per node, GPU pairs swapped between rounds).
• Every performance server starts from one frozen autotuner snapshot. Step-time cells also pin TRTLLM_MXFP8_GEMM_BACKEND=trtllm; otherwise the decode graphs' MXFP8 backend is re-profiled at every server start, which moves the step by about ±0.35 ms per server.

AL, TP4/EP4. Golden = rejection sampling with block verification; default = rejection sampling off.

	before	this PR	PR / before [95% CI]
golden SA gate replay (2 passes)	4.031	4.082	1.013 [0.998, 1.028]
golden ladder	2.411	2.406	0.998 [0.972, 1.026]
golden code	4.422	4.445	1.005 [0.960, 1.054]
default SA gate replay (2 passes)	3.538	3.604	1.018 [0.986, 1.050]
default ladder	2.136	2.128	0.996 [0.979, 1.015]
default code	3.806	3.727	0.979 [0.934, 1.027]
golden SA gate replay, modeling_v2 target, TP4 (4 passes)	4.065	4.020	0.989 [0.976, 1.002]

• The disagg PR's aggregated reference is 4.010 / 2.418 / 4.400 golden and 3.571 default SA.
• On the c44 config, the unforced MLPerf agentic v6 replay (531 requests, c44) gives AL 3.732 before and 3.722 with this PR.

KV capacity per rank:

config	before: target / separate draft pool	this PR: target / draft group (peak use)
TP4/EP4 golden	83,319 blocks / 6.66 GiB	85,708 blocks (+2.9%) / 4,505 blocks = 1.375 GiB (2,013)
c44 config, TP2	38,172-38,183 blocks / 18.32 GiB	43,087-43,110 blocks (+12.9%) / 6,092 blocks = 3.72 GiB (1,176-1,517)
c44 config, host tier for the target	64.57 GiB (the pool takes 18.32 of 82.89 GiB)	82.89 GiB (+28%)

Decode step time, forced AL, device step (iteration-log median of steady generation steps), MXFP8 backend pinned, mean of two rounds:

batch	c44 config, TP2: before / this PR	TP4/EP4: before / this PR
1	10.05 / 10.09 ms (+0.4%)
4	19.52 / 19.59 ms (+0.4%)	14.44 / 14.35 ms (-0.6%)
8	23.53 / 23.58 ms (+0.2%)
16		20.22 / 20.48 ms (+1.3%: 20.74 in one round, 20.22 in the other)
44	40.13 / 40.10 ms (-0.1%)
64		31.03 / 31.05 ms (+0.1%)

Client TPOT p50 is at parity at every concurrency. These are fresh servers. After two agentic replays at natural AL, this PR's step is 0.1-0.2 ms longer at batch 1 and 0.2 ms (0.4%) longer at batch 64. Without the pinned backend, a single server's decode step moves by up to ±0.35 ms either way.

At natural AL (rejection sampling off), TP4/EP4 waves of 1024-token requests give output tok/s 0.998 at c16 and 1.003 at c64 (this PR / before, two rounds). TPOT p50 is 1.000 / 0.987 and TTFT p50 1.010 / 0.987.

Agentic replay, c44 config (MLPerf agentic v6, 531 requests, c44). Forced AL 3.04 except the last two rows; 12 replays before and 6 with this PR, with GPU pairs swapped between rounds:

	before	this PR	PR / before
output tok/s/GPU	798.4	799.5	1.001
total tok/s/GPU	155.7K	153.8K	0.988
tok/s/user p50 / p90 (decode)	53.6 / 101.2	54.1 / 97.7	1.009 / 0.966
tok/s/user p50 / p90 (with TTFT)	46.1 / 79.8	46.6 / 73.6	1.011 / 0.923
TTFT p50 / p90	191.8 / 645.2 ms	193.8 / 613.9 ms	1.010 / 0.952
prefilled tokens per replay	565.1K	593.3K	1.050
natural AL (2 / 4 unforced replays)	3.732	3.722	0.997
output tok/s/GPU at natural AL	631.1	625.4	0.991

• The p90 per-user rates move by about ±5% between replays. A second run with an unfrozen autotuner gave 0.982 (decode) and 1.002 (with TTFT); over both runs the ratios are 0.974 and 0.962.
• The whole-replay TTFT p90 falls between first turns (9.8% of requests, 1.2-1.6 s) and continuation turns. Per turn type, TTFT p90 is equal (this PR vs before): 304.9 vs 305.1 ms on continuation turns, 1,557 vs 1,547 ms on first turns.
• The p90 per-user deficit comes from the extra prefill (next paragraph). The same unified layout with the partial page claimed instead (an experiment) matches before: 1.009 and 1.003 over both runs.

Partial pages. The unified drafter's window is reused only where it is complete. Before the partial-endpoint commit, an OPTIONAL lifecycle was complete only at a page boundary. A prefix hit that ended inside a page therefore claimed up to the last full page and recomputed the rest, at most 127 tokens per hit; the paired pool copies the partial page instead. On the stack without that commit, every remaining difference above comes from this recompute:

• Agentic: +5.0% prefilled tokens on the c44 config (+57 tokens per continuation turn) and +8.3% / +4.8% on TP4 (table below).
• One request repeating a cached prompt prefilled 65-105 tokens instead of 1, so c1 TTFT on such hits was 61-62 ms with this PR vs 52-59 ms before. Without a prefix hit, c1 TTFT is equal: p50 45.8 vs 45.6 ms over 2 x 300 requests.
• At c44 on warm prompts, the larger context step could cross the 3,072-token prefill CUDA-graph cap.

Dropping the recompute without completing the window costs AL. Claiming the partial page, so that the drafter starts with no window after such a hit, gives an unforced agentic AL of 3.219 instead of 3.722. The partial-endpoint commit instead completes the window: an OPTIONAL lifecycle is checked at any endpoint with the REQUIRED coverage rule, and its partial page is copied as for REQUIRED lifecycles. With it (c44 config, MXFP8 backend pinned, two rounds plus unforced replays):

	before	this PR	PR / before
prefilled tokens per replay	565.2K	565.2K	1.000
output / total tok/s/GPU (forced AL 3.04)	746.4 / 143.3K	755.2 / 146.1K	1.012 / 1.020
tok/s/user p50 / p90 (decode)	51.9 / 91.0	53.2 / 95.1	1.025 / 1.044
tok/s/user p50 / p90 (with TTFT)	45.1 / 71.4	46.1 / 75.3	1.021 / 1.055
TTFT p50 / p90	192.2 / 613.9 ms	193.6 / 620.1 ms	1.007 / 1.010
natural AL (2 unforced replays each)	3.762	3.716	0.988
output tok/s/GPU at natural AL	606.7	602.3	0.993

Two differences remain in the step pass. Its c1 requests are full-prompt hits, each sent right after a 1024-token request.

• TTFT, with iteration stats on: these requests start 2.6 / 9.0 ms later with this PR (two rounds). After each long request, trtllm-serve drains its iteration-stat records through the worker under the GIL, and the unified layout's records are larger: draft pool group, partial reuse counts. With enable_iter_perf_stats: false and print_iter_log: false, c1 TTFT p50 is 42.9 vs 43.1 ms; the 3-4 ms between rounds follows the GPU pair. A fix belongs in the stats path.
• TPOT, after two agentic replays: forced-AL decode of these hit requests is about 0.5 ms per step slower (c1 TPOT 2.64 vs 2.50 ms), with stats on or off. Fresh servers are at parity, and so are requests without a prefix hit after the same replays.
  • At natural AL after the replays, hits are at parity too: device step within ±0.3 ms at every batch size. SA gate AL is 3.910 vs 3.961 on hits and 3.908 vs 3.943 on misses.
  • Forced acceptance makes the target verify whatever the drafter proposes, so the verify's MoE cost depends on the drafts. This difference is that effect, not a KV cost.

Prefix reuse, TP4/EP4 golden, block reuse on. Hit set (greedy, c1) and the agentic replay (531 requests, c48, unforced):

	before, per_conversation	this PR, per_conversation	before, all_reusable	this PR, all_reusable
repeat, short / long prompt: cached tokens per hit	308 / 8,268	264 / 8,192	308 / 8,268	264 / 8,192
two-turn, short / long	346 / 20,220	304 / 20,096	346 / 20,220	304 / 20,096
16K shared prefix + short tail	12,416	12,416	15,879	16,512
16K shared prefix + long tail	16,570	16,570	16,570	16,570
agentic: prefilled tokens	565.6K	612.6K	535.3K	561.2K
agentic: AL	3.817	3.868	3.829	3.797

• Cold requests draft identically with both layouts: the hit set's cold AL matches before to four digits.
• Without the partial-endpoint commit, a repeat or two-turn hit is shorter by less than one page (partial pages, above).

AgentX c44 (agentx-harness 1.0.12 on aiperf; the c44 recipe with the nvfp4 indexer; the partial-endpoint commit included). The warmup is 488 prefill-only requests with many subagent branch-offs. Prefilled tokens over the whole warmup:

	before (#22's paired pool)	519 without #22 (unpaired pool)	this PR
prefilled tokens, whole warmup	73.0-75.4M in 1,720-1,782 s	23.0M in 602 s	18.0-19.0M in 456-483 s

kvCacheStats.cacheHitRate is not comparable across these layouts. KvCacheManagerV2 sums reuse and misses over all lifecycles, so in the unified layout the drafter lifecycle's page allocations count as misses: missed blocks are 2.1x the prefilled pages, against 1.0-1.2x with a separate pool. Compare prefilled tokens instead.

Soak (TP4/EP4 golden, block reuse on): 2,611 requests in 13 minutes: SA, 3 x 320 mixed stress, the agentic replay, 3 x 320 stress, SA.

• 0 errors and 0 bad server lines.
• AL stays flat: 3.41-3.45 per stress pass, and SA 4.018 before vs 4.065 after.
• The draft group peaked at 1,433 of 4,505 blocks.

End to end: InferenceX AgentX on Balaram's recipes

agentx-harness 1.0.12 in full mode at synthetic AL 4.04, on Balaram's recipes verbatim. "good" is the branch without the four DSpark commits (it fails the SA AL gate at 0.74 of golden); "519" is the branch head. "this PR" includes #25. UF is this PR plus the DSpark agg perf PR (fused ancestral drafting + drafter slot lifetime). P90 tok/s/user is 1/P90(ITL).

point	arm (runs)	P90 tok/s/user	tok/s/GPU	TTFT mean / p90	vs good (P90, tok/s/GPU)
c44-tp2-s6*	good (6)	62.2	109,886	2.30 / 5.92 s
c44-tp2-s6*	519 (7)	46.1	93,750	3.21 / 8.08 s	0.741, 0.853
c44-tp2-s6*	this PR (2)	67.1	109,817	1.97 / 4.39 s	1.079, 0.999
c44-tp2-s6*	UF (2)	67.3	110,253	1.95 / 4.71 s	1.083, 1.003
c1-tp4-s3	519 (2)	694.5	15,907		0.915, 0.990
c1-tp4-s3	UF (2)	736.8	16,030		0.971, 0.998
c16-tp4-s3-r2	519 (2)	325.0	58,704		0.898, 0.998
c16-tp4-s3-r2	UF (2)	352.2	59,025		0.973, 1.003

• At c44, this PR vs 519 is 1.455x P90 and 1.171x tok/s/GPU. The unified layout keeps the target's prefix cache through AgentX's prefill-only warmup: 18-19M prefilled tokens in 456-483 s, vs 21-23M in 534-602 s for good and 73-75M in about 1,750 s for 519.
• At TP4 (c1/c16) prefix reuse is equal across arms. There, UF costs about 1.3% per decode step vs the same stack on the separate pool (5.31 vs 5.24 ms at c1).
• *TP2 points are approximate: two servers per node, 89 GB host cache per rank, no MSA-v1 decode, nvfp4 index-K. TP4 points match the recipe exactly.

Not covered

• The iteration-stats cost on c1 requests after a long request (above): the unified layout's records are larger. A fix belongs in the stats path.
• A per-lifecycle (or target-only) cache hit rate in the KV cache stats, so the unified layout's hit rate is comparable (AgentX above).
• indexer_kv_dtype: nvfp4 of the c44 config needs the MiniMax MSA package, which our image lacks; every run used fp8. The c44 config's host cache was halved per rank to fit two servers per node.
• sm_107 targets (F0a) are unchanged and were not run.
• Disaggregated serving is unchanged by this commit; the disagg PR covers it.

Test coverage

• tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py: the unify gate (test_m3_unifies_a_dspark_drafter) and the window-sized share (test_window_sized_draft_ratio) are parametrized over aggregated and disaggregated serving.

• Unit runs on GB300 on the same wheel: this PR vs ff32d5d (the disagg PR before its host-sizing commit), so both commits are covered. No failure is new: every failure on this PR also fails on the base, with the same test ID.

  group	base ff32d5d	this PR
  MiniMax-M3 MSA attention files	135 passed, 6 failed, 9 skipped	137 passed, 6 failed, 9 skipped
  test_dflash_host_ctx_alloc.py, test_minimax_m3_kv_transfer.py	70 passed	70 passed
  speculative/hw_agnostic	1,003 passed, 44 failed	1,003 passed, 44 failed
  executor	4,166 passed, 119 failed, 5 errors, 3 skipped	4,166 passed, 119 failed, 5 errors, 3 skipped

  Not run on either side: test_hang_detector_kill.py (it SIGKILLs its process group), test_world_ranks_converge_on_hostless_fallback (it waits for an MPI world), and unittest/disaggregated apart from the M3 KV-transfer file.

• Unit runs on a wheel built at the partial-endpoint commit: this PR vs that commit, with the same result. The MSA attention files give 137 vs 135 passed and 6 failed on both sides. The touched files plus test_unified_draft_reuse.py, test_kvcm2_integration.py and test_dflash_unified_kv_gate.py give 326 passed and 2 failed on both sides: the MPI-world test above, and one reuse-rule expectation that the partial-endpoint commit then updated.

PR Checklist

• Please check this after reviewing the above items as appropriate for this PR.

🤖 Generated with Claude Code