[peihu-nv]

@coderabbitai summary

Description

Depends on #20917 (GB300 MiniMax-M3.1 decode targets). This PR builds on its shared DSpark and runtime-hook integration; until #20917 merges, the diff against feat/minimax-m3.1-flash also shows its commit. The context (prefill) target for Rubin is a separate PR stacked on this one.

MiniMax-M3.1 NVFP4 decode on Rubin (sm_107) currently falls back to the built-in model. This PR adds three modeling_v2 targets, routed on sm_107 by TP width:

Target	Topology	What it serves
minimax_m31_nvfp4__sm_107__tp8	TP8, experts TP8	low-latency K7 decode
minimax_m31_nvfp4__sm_107__tp4	TP4, experts TP4	low-latency K7 decode
minimax_m31_nvfp4__sm_107__tp2	TP2, experts TP2	decode baseline (aggregated serving and disaggregated generation)

Each target follows the layout of the GB300 targets in #20917:

• per-TP staircase targets;
• runtime hooks that stay inert unless a target arms them;
• every item default ON, with a TRTLLM_M31_DISABLE_<ITEM>=1 kill switch (m31_switches); items that are still opt-in use TRTLLM_M31_ENABLE_<ITEM>=1;
• compiled ops under cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp{2,4,8} and thop, built only for sm_10x, with fake (meta) registrations;
• bundled tactic pins (m31_tuning/) and ahead-of-time CuTe DSL objects (cute_aot_data/) shipped as package data, each guarded by the native build / toolchain it was produced with.

It also restores two DecoderRunner.__init__ defaults that #20917 dropped (guided_decoder and lora_model_config, both None); without them every server start fails in warmup with an AttributeError.

Shared changes armed by the sm_107 targets

These are inert unless a target arms them through RuntimeHooks (or an explicit config option):

• runtime hooks: step-input pack and prologue op, draft block-offset patching, fused spec offsets, first-content emission, one-event iteration timing, a late-compile detector, the MXFP8 pinned-tactic table, log-only work counters;
• DFlash / DSpark drafter APIs (fc TP split, fused all-reduce + RMSNorm, windowed-decode cap, full-slot context arena) and a parallel DSpark context advance;
• MSA backend forward_staging, MSA KV-cache fixes (pool zero-init, exact KV-size estimate), standalone-drafter KV pool fixes, saturating FP8 casts;
• disaggregated serving: the DSpark drafter's context window travels with the target KV (DisaggregatedParams.draft_ctx_state);
• indexer_kv_dtype='nvfp4' (prototype): NVFP4 index-K pages for the MSA implementation;
• native ops: TOPP_DET v3 list-path / wide-histogram variants (schema arguments default off), a slot-indexed chain speculative sampling variant, REJ v3 selected only on sm_107, a cluster-16 router GEMV, an FC-shard GEMM;
• requirements.txt pins nvidia-cutlass-dsl[cu13]==4.8.0 (the release, same as upstream main) instead of 4.8.0.dev0.

Shared paths taken only by the sm_107 targets

These shared paths are off by default for every model and GPU. Routing arms them, process-wide, when it selects an sm_107 MiniMax-M3.1 target, before the model is built (arm_sm107_shared_opt_ins). Each keeps its kill switch, and the same paths are available to other models as RuntimeHooks opt-ins.

Path	Kill switch
KV GPU prefix prefetch (py_executor): waiting context requests have their host-tier prefix staged to the GPU and are admitted after a short rank-agreed hold (KV cache manager V2 + block reuse).	TRTLLM_M31_DISABLE_KV_GPU_PREFETCH=1
DFlash / DSpark draft KV reserve + guard page (kv_cache_manager_v2): the draft KV manager reserves the drafter's block width per step and keeps page 0 as a guard, so drafter writes through unallocated block-table entries cannot clobber another request's page.	TRTLLM_M31_DISABLE_DRAFT_KV_RESERVE=1, TRTLLM_M31_DISABLE_DRAFT_KV_GUARD=1
Context-only prompt tokens as an int32 array (llm_request; host arrays only).	TRTLLM_M31_DISABLE_L1_TOKEN_ARRAY=1
MSA NVFP4 sparse prefill BF16X / E4M3X fork (msa_prefill), when the fmha_sm100 fork is installed and the geometry has 16 query heads per KV head.	TRTLLM_M31_DISABLE_TP2CTX_BF16X=1
MSA indexer proxy max-score buffer reuse across sparse layers (msa_indexer).	TRTLLM_M31_DISABLE_PROXY_MS_REUSE=1
Context one-hot draft probabilities via index_fill_ (speculative/interface; same elements, no host sync).	TRTLLM_M31_DISABLE_ONEHOT_INDEX_FILL=1
DFlash / DSpark host-side context allocation when the draft block table counts full rows (dflash).	TRTLLM_DFLASH_DISABLE_HOST_CTX_ALLOC=1
Shipped ahead-of-time CuTe DSL objects (cute_aot): an unarmed process compiles every kernel in-process with no lookup or export.	TRTLLM_M31_CUTE_AOT_DISABLE=1

toppRenormDetV3Sm103's TOPP_V6 chain-prefetch summation runs only on sm_107 devices; GB300 callers of the #20917 targets keep the staged loop (TRTLLM_M31_DISABLE_TOPP_V6=1 turns it off on sm_107 too).

Other shared behaviour changes (no kill switch; correctness fixes and validation):

• DFlash / DSpark context one-hot draft probabilities clear the full draft-probability width, not the draft logits' width, so a vocab-sharded drafter cannot leave stale columns for new requests.
• The one-model speculative sampler rejects requests with a logits_processor (it was silently never called), and the disaggregated peer check rejects replicated roles whose per-layer cache sizes differ (e.g. different indexer_kv_dtype).
• The scheduler reads a request's prompt length from num_input_tokens when present (same value, no token-list build).

Items and numerics

Unless noted, an item is exact: bitwise against its stock path in eager and graph mode, gated by kernel oracles. These items are not bitwise: same inputs and rounding points, different fp32 reduction order:

Item	What changes
VR_GEMV48 / LINEAR_CUBLASLT	33..64-row MXFP8 projections on the same quantized activation
S9B	drafter all-reduce fused with add + RMSNorm
MXDRAFT	drafter MXFP8 GEMMs; changes only the draft, never the accepted-token rule (ON at TP8 and TP2, OFF at TP4)
LC_ONEDIV (TP4)	up to 16 query rows, the long-context sparse-decode split-K merge divides once per output (fixed order, no atomics; within the float64 reference's error band)

S2 (drafter fc row split), S5 (window-capped drafter decode), MCAG, S9B_MC, REJ_SLOTS and SLOT_IDS_CHANGED are compiled in but default OFF.

Measured results (VR200, end to end)

AgentX-style multi-turn agentic workload (SemiAnalysis InferenceX client, agentx-harness 1.0.12 preset), K7 DSpark speculation, sampled (T 1, top-p 0.95, rejection sampling), 3600 s per cell. Acceptance: synthetic (forced) K7 AL 4.04; natural AL not demonstrated. (The acceptance counters read 4.039 on every generation endpoint, which reflects the forcing.) Y = output-inclusive total token throughput per GPU; X = per-user generation tokens/s.

Cell	Topology	Y per GPU	X median	X P90	TTFT p50
agg c20	1x TP2 (2 GPUs)	76.87K	535.8	479.3	0.17 s
agg c32	1x TP2	128.38K	410.4	282.0	0.25 s
agg c56	1x TP2	212.97K	262.2	200.9	0.23 s
agg c68	1x TP2	243.10K	209.4	163.6	0.34 s
disagg c16	ctx 1x TP2 + gen 1x TP8 (batch 1)	11.86K	1310.6	1227.2	0.60 s
disagg c24	ctx 1x TP2 + gen 2x TP4 (batch 1)	19.64K	1024.7	980.6	0.36 s
disagg c64	ctx 1x TP2 + gen 4x TP2 (batch 2)	50.72K	658.3	611.5	1.34 s

• The aggregated cells ran a build whose decode code differs from this commit only by the fake (meta) registrations of the native SwiGLU-OAI / MoE-finalize ops and their unit tests. The disaggregated cells ran this decode code exactly, with the context role served by the stacked prefill PR's target.
• Two files in this commit are formatting-only changes (identical AST) relative to the measured build: catalog/spec_decode/dspark_ctx_advance.py and tests/unittest/_torch/thop/parallel/test_minimax_m31_dspark_glue_ops.py.
• The disaggregated TP8 / TP4 generation cells run one request per generation server by configuration (latency cells).
• Native builds are not byte-reproducible; the measured images are the builds the tactic pins were recorded on (the pins carry a native-build guard).

Disclosures and known gaps

• The disaggregated c24 cell had 3 failed requests out of 4,049. Two come from the generation server's first verify after KV receive using stale draft probabilities (fixed in the follow-up commit "Disaggregated generation: one-hot draft probs on the first verify under rejection sampling" on this PR; the measured cells ran without it); each was followed by repeated stock post-processing errors on the same request. The third is a client transport error.
• One DSpark drafter-input kernel is loaded from disk on first use rather than at warmup (a fix is parked).
• The measured runs pre-seeded Triton / FlashInfer / Inductor JIT caches from a named seed; without it these compile at startup.
• TP4 / TP8 generation: the dense SwiGLU above 8 rows runs the Triton kernel (the native SwiGLU-OAI / MoE-finalize ops are TP2-only); the TP4 / TP8 cells above ran batch 1.
• The context-phase CTXNORMQ boundary is not taken on sm_107 (the op from [None][perf] Add GB300 MiniMax-M3.1 low-latency decode targets #20917 is built for SM103 only).
• With max_seq_len 1,048,576, set max_batch_size explicitly: at the default (2048) the MSA backend's worst-case plan overflows its int32 KV offsets at startup. A fix (int64 offsets or a clamp) follows.
• The sm_107 context target loads no bundled tactic pins yet, so it autotunes at startup.
• Drafter determinism above 8 concurrent requests relies on the draft KV reserve + guard above.
• MSA vendoring and the optional flashinfer-cubin dependency for Rubin are in the stacked prefill PR.

Updates on this PR after the first commit

Landed:

• Disaggregated generation first-verify fix: a request received from the context server enters with dummy draft tokens and a one-hot draft-probability row under rejection sampling, so its first verify cannot read a stale or uninitialized row (CPU unit test test_disagg_gen_first_verify_onehot.py).
• TOPP_V6 limited to sm_107 devices, and the LlmRequest token-array fast path limited to host arrays (a device array takes the sequence cast, which rejects it, as before the fast path).
• The shared paths above armed only for the sm_107 targets, from routing, so the target directories (the shipped AOT objects' source digests) are unchanged.
• DSpark drafter-inputs kernel: no Triton specialization on the batch-dependent accepted-token stride and view alignment, so a batch shape warmup never ran no longer compiles after the server is ready. No kill switch: it is a compile-specialization hint with no alternate path (bitwise unchanged).
• The iteration-stats RPC is awaited instead of blocking the event loop (TRTLLM_M31_DISABLE_A4_ASYNC_STATS=1 restores the blocking call).
• Two diagnostics, off unless their environment variable is set: a phase-armed Nsight capture window (TRTLLM_M31_PHASE_NSYS_ARM) and a per-rank host-path probe (TRTLLM_M31_HOST_PATH_PROBE).
• TP2 sequence-parallel router: each rank computes its owned half of the rows with a native TF32 GEMM (trtllm::m31_router_half_pin) at the full-step row counts qualified bitwise against the full-row reference; other layouts keep the full-row path (TRTLLM_M31_DISABLE_ROUTER_HALF_PIN=1). The measured effect is small (about 0.02-0.03% of mixed-step time; no decode-step change).
• The executor's iteration timing uses native GPU timestamps on the sm_107 targets instead of timed CUDA events (TRTLLM_M31_DISABLE_GPU_TIMESTAMPS=1).
• The GPU byte oracles for the router pin and the drafter-inputs kernel need an sm_107 GPU and run locally; their CPU counterparts are in the CPU test list.
• Validated on VR200 on a rebuilt image: native, speculative, top-p (1,002 sampling oracles, bitwise), opt-in, context-target and disaggregated unit-test suites, and a serve smoke of the aggregated and context roles with engagement, tactic pins and ahead-of-time objects unchanged from the measured build.

Pending:

• Bundled tactic pins for the TP4 / TP8 / TP2 disaggregated generation roles.
• The MSA int32 KV-offset fix above.

Test Coverage

• tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_107__tp{2,4,8}/: kernel oracles (bitwise or reduction-order), dispatch, kill switches, routing, prebuild and work-counter tests.
• tests/unittest/_torch/thop/parallel/ native-op tests (incl. fake registrations), tests/unittest/_torch/multi_gpu/test_minimax_m31_{f3_ar_prepushed,s2_moe_tail}.py (bitwise to the stock all-reduce chains).
• CPU tests added to l0_cpu.yml: test_m31_pins.py, test_cute_aot.py, test_flashinfer_autotune_cache.py; the two TP2 multi-GPU tests are listed in l0_dgx_b200.yml.
• tests/unittest/disaggregated/: draft-KV transfer, rank-info, hook-off and switch tests.
• On VR200 at the measured build: the native-op and speculative-decoding unit-test suites passed, and a serve smoke test passed for the aggregated, context and generation roles (pins loaded on every rank, no late compiles after ready).

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

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

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.