Repository navigation
Conversation
MXFP8 non-expert projections with NVFP4 routed experts on the existing model: model_config expands the packaging's expert recipe into the mixed per-layer config, the weight mapper renames the packed expert tensors and inverts their global scales, MXFP8 index projections and fused MXFP8 shards load their own block-scale rows, the sparse cache manager takes the layer schedule from the checkpoint and counts index-K bytes, the router GEMM runs in FP32 cuBLAS, and the VL input processor loads the HF processor only for image requests. Signed-off-by: Pranav Shrestha <pshrestha@nvidia.com>
MiniMaxM3Gate keeps a BF16 split of its FP32 router weight (three terms along K, built at load_weights, the multi-term emulation of CUTLASS's 3xTF32 GEMMs). BF16 inputs run one cuBLAS BF16 GEMM with FP32 accumulation over the replicated input instead of an FP32 cast plus a SIMT SGEMM; FP32 inputs keep the FP32 cuBLAS GEMM. Top-k parity against an FP64 reference is tested. Signed-off-by: Pranav Shrestha <pshrestha@nvidia.com>
KV pool and block-scale views are built once per layer and layout, the fused QKV/index producer caches its per-layer-invariant arguments, block-id rows are copied from the page-index buffers, the MSA metadata builders do their batch-sized index arithmetic in numpy, and the NVFP4 decode and CSR prefill wrappers resolve the device capability and the fmha_sm100 API once. Same kernels, same arguments; adds prefix-reuse and BF16-cache decode tests. Signed-off-by: Pranav Shrestha <pshrestha@nvidia.com>
…SM107 The MiniMax-M3 Triton NVFP4 sparse decode kernels run on SM107, so keep NVFP4 KV for that path instead of falling back to FP8 and include SM107 in the decode capability list and test guards. Signed-off-by: Pranav Shrestha <pshrestha@nvidia.com>
…deo tokens per patch trtllm-serve and the multimodal prompt builders mark their rendered TextPrompt with chat_template_applied; the input processor uses it verbatim instead of wrapping it in a second user turn, normalises a framed marker, and rejects a marker/media count mismatch before the HF processor runs. Media arrive as uint8 numpy, and video tokens count one start/end bracket per temporal patch. Signed-off-by: Pranav Shrestha <pshrestha@nvidia.com>
…king on multimodal requests Forward the pre-fusion token ids as orig_input_ids so the spec worker (DSpark/DFlash, MTP, Eagle) sees them after fuse_input_embeds drops input_ids, and slice the vision embeddings with find_input_mm_embeds so a cached prefix or a chunk boundary inside a placeholder span no longer trips the token-count check. Signed-off-by: Pranav Shrestha <pshrestha@nvidia.com>
The input processor implements the token-ID expansion hooks but never set the opt-in flag, so a prompt_token_ids + image request was detokenized and re-templated, doubling the image marker and failing the marker count check. Set the flag and test the dispatch end to end. Signed-off-by: Pranav Shrestha <pshrestha@nvidia.com>
The drafter's nested HF config is flattened, its checkpoint tensors renamed, MXFP8 linears dequantized and Gemma norm weights folded at load. Its KV pool uses its own page size, is funded for its attention window plus the chunk in flight, and the unpaired draft mirror is sized from the target request and keeps the chunk start as history; the block decode reads a rebased table. Signed-off-by: Pranav Shrestha <pshrestha@nvidia.com>
… generation-init requests A generation worker receives the prompt KV instead of prefilling it and starts the drafter on a fresh slot, so the draft mirror never holds K/V for the transferred prompt. Treat that prompt as history on gen-init so the drafter's window trims the mirror to its tail plus headroom, instead of asking for a prompt-sized chunk that exceeds the pool funding. Signed-off-by: Pranav Shrestha <pshrestha@nvidia.com>
…ligned reused prefix On an unpaired draft pool the block decode reads a view rebased by the rounded-up base of a reused prefix. While the stored tail is shorter than that base, entry 0 of the view lies past the request's allocation and resolves to page 0, so the block's own K/V landed in another request's page and the drafter stayed blind until the context caught up. The view is now shifted only to the page holding the context length on those steps; the shortened read length is unchanged, so nothing below the base is read. Adds a CPU test that walks every context length of that window. Signed-off-by: Pranav Shrestha <pshrestha@nvidia.com>
…sing in step Resolve the parser's mode from thinking_mode or the prompt-resolved relay instead of the bare thinking flag, so a prefilled <mm:think> block is read as reasoning. Report truncated reasoning as reasoning in the non-streaming parse and make the streaming parser emit the same split: unopened </mm:think> is a delimiter, text before an opening tag is content, partial tags are withheld at chunk boundaries. Signed-off-by: Pranav Shrestha <pshrestha@nvidia.com>
The deployment guide gains a linkable "Serving MiniMax-M3.1-Flash" section (bundled DSpark drafter, KV cache reuse, NVFP4 KV cache on SM107, thinking_mode and reasoning parsing) and states that both MoE backends load the NVFP4 expert packaging with the same static recipe. The supported-models matrices and the sparse-attention page now say block reuse is supported and mixed MXFP8/NVFP4 loads natively. Signed-off-by: Pranav Shrestha <pshrestha@nvidia.com>
…okenization cache `TransformersTokenizer` delegates `num_special_tokens_to_add` to the wrapped tokenizer (the base class raises), so the MiniMax-M3 input processor's tokenization-cache check runs. `resolve_hf_chat_template` falls through to the tokenizer's template when the processor's `chat_template` is None. Signed-off-by: Pranav Shrestha <pshrestha@nvidia.com>
….4.0a1 The side-branch release job derives its wheel and container version from tensorrt_llm/version.py, so the branch gets its own alpha base instead of sharing the main nightlies' 1.4.0rc1 line. Signed-off-by: Pranav Shrestha <pshrestha@nvidia.com>
Use the 1.4.0a0 base required by the side-branch release convention and keep the README badge and example dependency constraint aligned. Signed-off-by: Yanchao Lu <yanchaol@nvidia.com>
Signed-off-by: Tyler Burt <195370667+tburt-nv@users.noreply.github.com> (cherry picked from commit c962edd) Signed-off-by: Yanchao Lu <yanchaol@nvidia.com>
Add SM103 TP8 single-request decode through eight rows and TP4 decode through 32 rows, with shared native kernels and guarded DSpark runtime integration. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
…3.1 TP4 target op16.py loads m3_moe16_kernel.py by path at prebuild time, but the file was not part of the target, so every modeling_v2 TP4 server failed during post_load_weights with FileNotFoundError. Add the 9..16-row MoE kernel that op16.py's M3_CONFIG contract (m_max, num_ctas, nb2_cap, fused_router, shared, timing) was written against. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…nfig DecoderRunner.__init__ stopped setting guided_decoder and lora_model_config, so warmup of any model without a guided decoder (e.g. the built-in MiniMax-M3 path with modeling_v2 off) failed with AttributeError on self.guided_decoder. Restore both defaults. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…pping msa_backend._build_msa_fields passes device_tables=not h2d_cache.enabled(), but build_paged_kv_slot_mapping had no such parameter, so every built-in MiniMax-M3 MSA step failed with TypeError. Add the parameter: with device_tables=False the device req_to_token/slot_ids tables are skipped and out_cache_loc is returned pinned on the host for the caller to stage. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…prefix dimension (NVIDIA#19975) Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> (cherry picked from commit 0c9b463) Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…slower than FP8 (cherry-pick of NVIDIA#19974) Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> (cherry picked from commit d68ea55) Conflict resolution against the MiniMax-M3.1 branch: - cache_manager.py: keep the branch's _cached_pool_view getters (they already cache every K/V and scale view), so upstream's _kv_buffers/_block_scale_buffers dicts are dropped; the eager per-layer view resolution at the end of __init__ now just warms that cache. copy_batch_block_offsets returns early when there are no attention-op pools, before the branch's unchanged-copy guard. - modeling_minimaxm3.py: keep the branch's cached _fused_producer_plan; the explicit kv_scale_cache from the piecewise producer op is threaded into the plan (and its cache key), and the plan resolves it from the manager otherwise. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…erate managers without is_nvfp4_layer _cached_pool_view keyed views by the raw kv_layout, so get_buffers(i) and get_buffers(i, "HND") built two views of the same pool and missed the views warmed at construction. Key by the resolved layout instead. _resolve_fused_producer_plan again treats a missing manager or one without is_nvfp4_layer as non-NVFP4 instead of raising AttributeError. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…n-tree kernel ops/gqamix_sparse_decode.py must stay a line-for-line copy of the in-tree triton_sparse_decode.py apart from its specialization (enforced by test_modeling_v2_m31_gqamix.py). It predated the SM107 NVFP4 capability change and the per-device capability cache, so the copy test failed. Re-copy the in-tree module and re-apply the two specialization edits. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…ode sampler ops Every TP rank runs the speculative-decoding sampler (draft sampling and the rejection test) on its own and relies on bit-identical results across ranks. top_p_renorm_probs_op and compute_probs_from_logits_op called FlashInfer's top_p_renorm_probs in its default, non-deterministic mode, so the renormed probs could differ in the last bits per rank; near-threshold accept/sample decisions then occasionally differed, the ranks' token counts drifted apart and, once a request finished on some ranks only, the TP group deadlocked. Pass is_deterministic=True, as the other FlashInfer sampling calls in this file already do. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
… GQAMIX The GB300 MiniMax-M3.1 modeling_v2 targets (tp4, tp8) register their prefill tile select (TSEL) on MsaIndexer.selector_hook and their target-local sparse decode (GQAMIX) on MsaDecodeFmha.sparse_decode_hook, but neither attribute existed in the MSA backend, so both registrations were no-ops. Add both attributes (default None). MsaIndexer._select offers the KV-head-reduced scores to the selector hook and keeps the built-in select when the hook is unset or returns None; MsaDecodeFmha._run_sparse calls the decode hook in place of minimax_m3_sparse_attn_decode with identical arguments when set. With no hook installed the behaviour is unchanged. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Add sparse_decode_tp2_policy with the GB300-measured split-K, query pairing, producer warps/stages and merge warps for the NVFP4 sparse decode at two KV heads per rank and decode_query_len 4, keyed on the local batch, and consult it from minimax_m3_sparse_attn_decode. Unlisted batches, other KV-head counts, other decode_query_len values (e.g. DSpark K=7) and non-SM103 devices keep the default routing. M3_O8_SPARSE_TABLES=0 or M3AX_STACK_A=0 disables it; M3_O8_PAIR=0 selects the no-pairing table. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…se-decode variant at warmup Ports the GB300 stall-fix prewarm (m3ax_msa_prewarm) to the built-in MiniMax-M3 MSA path: - core: enumerates every fmha_sm100 indexer-proxy variant (unsplit q 1..512 and the KV-split segmented family) with the library's own dispatch functions and compiles the missing ones in parallel under the library's file locks, before the stock split-proxy prewarm loads them. The dense family is dropped (MiniMax-M3.1 has no dense attention layers). - sdwarm: records the Triton sparse-decode call shapes during warmup and, when warmup finishes, warms every local batch 1..max_batch_size, every pointer-alignment signature eager mixed steps produce, and every decode query length the runner can use (recorded, 1 + max draft len, dynamic draft lens; DSpark K3 -> 4, K7 -> 8). - Serve-time observers log the first use of each fmha_sm100 variant, any inline fmha_sm100 compile, and every Triton compile after warmup. - DecoderRunner.warmup and _prewarm_minimax_m3_msa_indexer drive it; non-MSA runners are untouched. Switches (default on, as in the original): M3AX_STACK_A, M3AX_MSA_PREWARM_VARIANTS, M3AX_SPARSE_DECODE_PREWARM. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
… for GB300 TP2 Adds torch.ops.trtllm.minimax_m3_cute_sparse_decode, a block-sparse GQA verify-decode kernel for the TP2 NVFP4 shape on SM103: one CTA per (top-k split, KV head, query row), clusters merging over DSMEM, a bulk-copy smem ring feeding m16n8k16 HMMAs with NVFP4 dequantized in registers. With M3_CSD=1, minimax_m3_sparse_attn_decode routes a call to it when the device is SM103, the rank holds two KV heads, the (decode_query_len, local batch) key is in the built-in table (dql 4 at b3-b45, dql 6 at b2-b25) and the operands are covered; every other call (including dql 8) keeps Triton. M3_CSD_TABLE=<json> replaces the table for benchmarking. Default off. Ported from the GB300 optimized stack, iteration 4 (lever 05). Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…undary all-reduce, router fast path, NVFP4 hoist, long select) Built-in MiniMax-M3 / M3.1 model path, decode on SM100/SM103. Off by default: the master switch is M3AX_L5=1; under it each piece has its own switch, on unless set to 0 (M3AX_L5_MC_BOUNDARY, M3AX_L5_ROUTER_TC, M3AX_L5_NVFP4_HOIST, M3AX_L5_SELECT). Row ranges: M3AX_L5_MC_ROWS (1-48), M3AX_L5_ROUTER_TC_ROWS (2-48), M3AX_L5_SELECT_ROWS (1-72), M3AX_L5_HOIST_MAX_ROWS (1024). - mc_boundary: at each post-feed-forward boundary of a decode CUDA-graph forward of <= 48 rows, the shared + routed add, the AllReduce(RESIDUAL_RMS_NORM) and the next layer's MXFP8 input quantize become one minimax_m3_mc_allreduce_norm call. The next layer's QKV runs only its GEMM on the published MXFP8 copy, through a new prequantized slot in MXFP8LinearMethod. The residual gets one BF16 rounding instead of two (ulp-level difference). The workspaces reuse the M3.1 modeling_v2 tp4 multicast helpers. - router_tc: decode MoE forwards of 2..48 rows run minimax_m3_router_tc for the router logits and the experts' NVFP4 input in one kernel. - nvfp4_hoist: otherwise the experts' NVFP4 quantize runs on the shared-expert aux stream ahead of the shared expert, off the router chain (bitwise). - select: top-16 selects of 1..72 rows over 129..9216 blocks use minimax_m3_select_long (bitwise except on > 1024-way ties at the 16th score). At TP2 with DSpark K=3 (4 rows per request) the 48-row parts cover batch <= 12 and select covers batch <= 18. The NVFP4 pieces apply only to the TRT-LLM-op TRTLLM-gen NVFP4 experts with in-kernel routing, one chunk, no DP, comm or EPLB. Eagle3/DSpark capture layers, the last layer, eager, prefill and breakable-graph forwards keep the stock path. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
… workspace failures across ranks router_tc computes TF32 router logits, less precise than this model's FP32-grade MiniMaxM3Gate, so it now needs an explicit M3AX_L5_ROUTER_TC=1. Workspace creation also records ValueError/OSError as a local failure so that every rank still reaches tp_agree_built instead of one rank leaving the collective. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…mixed-step host trims Context (prefill) rows of the built-in MSA path, sm_103: - Block select: context rows take torch.ops.trtllm.minimax_m31_select_prefill_tile (one CTA per KV head and 8 queries, each score sector read once) instead of minimax_m3_select_blocks. Same table bit for bit, including tie order and init / local forcing. Falls back to the stock select off sm_103, without the op, under capture / compile, for scores without unit query stride, at or below 128 blocks (where the stock bitonic select is faster), past 5856 blocks (about 737K tokens) and below TRTLLM_M3_PREFILL_SELECT_TILE_MIN_ROWS (256) context tokens. Off with TRTLLM_M3_PREFILL_SELECT_TILE=0. In a mixed step it writes the context rows straight into the joined head-major table. - Proxy max score: the context proxy pass writes into an uninitialized [heads, max_k_tiles, tokens] block instead of fmha_sm100's -inf filled one; the proxy writes every entry the select reads. - k2q CSR: the NVFP4 prefill builds its CSR worklist and schedule into buffers kept on the metadata and reused by every layer, with the schedule sizing done once, instead of build_k2q_csr allocating and sizing them per layer. - per_token_valid_blocks runs in NumPy. The three host trims are off with TRTLLM_M3_MSA_PREFILL_HOST_TRIMS=0. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…-input cat MiniMaxM3Gate builds FP32-grade logits from BF16 activations as cat([x] * 3, -1) @ split_weight.t(). At 16K prefill rows the cat of 3 x [16K, 6144] costs more than the GEMM. A Triton kernel computes the same product reading x per term and chaining the BF16 products into one FP32 accumulator in the concatenated K order, which matches the cuBLAS kernel bit for bit on SM103 for every row count in 8625..16384 (exhaustively checked; below that cuBLAS alternates between kernels with a different K order). GB300, real gate weights: 16384 rows 0.298 -> 0.140 ms per layer (~9 ms per 16K step over 57 MoE layers). Rows outside the range, other SMs and torch.compile keep the stock path. M3_ROUTER_TRITON_ROWS=off disables it; it also takes other lo-hi ranges. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
The drafter's context store read the per-row draft block counts back from the device (.tolist()), blocking the host until the GPU finished the step, so the drafter, lm_head and sampling tail launched late on every mixed step. With KVCacheManagerV2's shared-page-table copy every copied row holds no negative entry, so the count is the row width, known on the host; other writers keep the device read. TRTLLM_DFLASH_SYNC_CTX_ALLOC=1 restores it. DSv4 DSpark window seeding reads the chunk's first position from the host position ids instead of .item(). Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
A forward waits on the KV tier copies its scheduling issued: offloads that free GPU slots for its new pages and onboarding of reused prefixes. Both can be issued earlier, after the previous forward is launched, so they overlap its compute. All of it is opt-in: * TRTLLM_M3_ASYNC_OFFLOAD=1: an OffloadHeadroom per v2 pool with a lower tier holds TRTLLM_M3_ASYNC_OFFLOAD_TOKENS (default max_num_tokens) of GPU slots in a private, stats-excluded cache on a side stream between launch and the next schedule. Backing it evicts LRU pages; releasing it leaves free slots whose ready event follows those offloads, so the next step's allocations take them without evicting. * TRTLLM_KV_ONBOARD_PREFETCH=<n>: up to n not-yet-admitted first-chunk context requests per step have their reused prefix onboarded through a transient cache (KvCache.prefetch to the GPU level). Requests arriving during a step are activated right after its launch so their copies get the lead (TRTLLM_KV_ONBOARD_PREFETCH_EARLY_FETCH=0 disables; off with attention DP). * TRTLLM_KV_ONBOARD_PREFETCH_DEFER_TOKENS=<t>: an arrival whose onboarding brings at least t reused tokens to the GPU is staged before its first schedule and held out of it for one iteration. Page ordering stays with the core's per-page and per-slot ready events. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…attention TRTLLM_M3_MSA_V1_DECODE=1 routes sparse generation rows to MiniMax MSA's CUTLASS Q8KV4 decode kernel over the existing NVFP4 pool, page table and top-k pages; calls outside its contract keep the Triton/CuTe kernels. Needs an MSA build that reads the M3 4x4 token-quad V-scale layout. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
… NVFP4 index-K mirror TRTLLM_M3_MSA_V1_INDEXER=1 (with indexer_kv_dtype fp8) widens each sparse layer's INDEX_KEY page to hold an NVFP4 copy of the E4M3 index-K, refreshed from this step's slots in run_indexer, and scores pure decode/verify steps with MiniMax MSA's q8kv4 indexer. The table is sorted back to the ascending, -1-padded contract; other steps keep the CuTe DSL scorer and FP8 prefill. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…TRTLLM_M3_CHEAP_PAD) CUDA-graph padding rows get zeroed logits before rejection sampling, so their probs are uniform. A uniform row is the worst case of FlashInfer's radix top-p renormalization (the whole vocabulary falls in one bucket) and one such row sets the duration of every top-p launch in the step: on MiniMax-M3.1 TP2 DSpark K6 a padded decode step pays ~9-10 ms in top-p (6 radix passes at 1.2-2 ms instead of ~0.13 ms), independent of how many padding rows there are. With TRTLLM_M3_CHEAP_PAD=1, padding rows get a one-hot logits row instead, which the filters pass through cheaply, and the uniform probs that zeroed logits give are written back after the filters, so rejection still accepts any stale draft token on a padding row. Real rows are untouched. Default off. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…ding rows Padding rows now always get a one-hot logits row, with uniform probs restored after the filters, instead of only under TRTLLM_M3_CHEAP_PAD=1. Real rows are bit-identical, and padding rows end with the same probs as before, so the change is lossless; it removes 9-10 ms of radix top-p work from every padded decode step. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
- write_context_onehot_draft_probs: index_fill_ instead of a Python scalar assigned through an advanced index, which staged a pageable host tensor and synchronized the stream on every step with context rows. - NVFP4 MSA prefill attention: validate inputs once per step geometry and combine straight into the output (no per-layer validation or copy). - Cache the KV pool device and the MSA-v1 device-capability check. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> (cherry picked from commit 7e90b63)
…ide torch.compile The registered custom op launches the same compiled kernel behind a Python custom-op dispatch and an input-description string formatted on every call; _cutedsl_score runs once per sparse layer of every eager mixed step. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> (cherry picked from commit 2fa0a79)
…task scores GSM8KInferenceX sets no scores_filter, so the average ran over the 'alias'/'name' string fields and raised a numpy dtype error after the results were already saved. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> (cherry picked from commit 6f7b009)
… over an NVFP4 index-K mirror" This reverts commit 3860fdd.
…step ahead" This reverts commit da32b44.
…t; agree workspace failures across ranks" This reverts commit 3780c9b.
…icast boundary all-reduce, router fast path, NVFP4 hoist, long select)" This reverts commit 8bb33b4.
…e kernel for GB300 TP2" This reverts commit dee76d1.
… (o8)" This reverts commit c7a463c.
…-decode launch table Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
… prewarm, prefill select, router and DSpark paths Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Squash of the MiniMax-M3.1 Rubin (sm_107) decode work onto this branch: the TP2 / TP4 / TP8 low-latency decode targets, the shared sm_107 opt-ins they arm at routing (shipped CuTe DSL AOT objects, drafter KV reserve and guard page, host-tier KV prefetch, int32 context token array, BF16X / E4M3X sparse-prefill fork, proxy max-score reuse, context one-hot index_fill, DFlash host context allocation, GPU timestamps), the TP2 router half-row TF32 pin, the TP2 BF16 draft gather with its native Markov-epilogue consumer (kill switch TRTLLM_M31_DISABLE_DRAFT_AG_BF16), and the fixes and tests that came with them. Every opt-in keeps its kill switch. Where this branch changed the same shared code, both behaviours are kept: - The sm_107 routing also arms msa_indexer.arm_sm107_msa_paths (no split context proxy plan, no prefill host trims, no prefill tile select, no MSA warmup prewarm) and interface.arm_sm107_spec_paths (CUDA-graph padding rows stay zeroed, FlashInfer's default top-p renorm mode). Every other process keeps this branch's defaults. - DFlash context allocation: armed processes use the host allocation the sm_107 targets were validated with; others use _ctx_alloc_tokens. - Context one-hot rows: index_fill_ everywhere; an armed process's kill switch restores the advanced-index write. - MSA prefill / indexer / backend: the segmented proxy plan, the reused CSR workspace and the direct NVFP4 attention launch run where the holds above are off; the policy proxy splits, proxy oracle, NVFP4 index-K and the BF16X / E4M3X / Q8KV4 fork are unchanged. Source commits: - 9d97565 [None][fix] Restore DecoderRunner optional state initialization - 162e93e [None][feat] Add Rubin (sm_107) MiniMax-M3.1 low-latency decode targets - bc6d981 [None][fix] Disaggregated generation: one-hot draft probs on the first verify under rejection sampling - fd38234 [None][fix] Gate TOPP_V6 to sm_107 and keep the token-array copy to host arrays - 0f4ab31 [None][fix] Arm the shared MiniMax-M3.1 sm_107 paths only for the targets that opt in - 0c88f87 [None][test] Expect a device token array to be rejected by LlmRequest - 3f48ba4 [None][test] Cover the KV prefetch startup state when the model does not opt in - 6aaa96e [None][fix] Arm the shared MiniMax-M3.1 sm_107 opt-ins at routing - bb04179 [None][fix] Avoid batch-shape specialization in DSpark drafter inputs - b3f097d [None][test] Cover DSpark input strides and context-view specializations - 6355f11 [None][fix] Await the iteration-stats RPC instead of blocking the event loop - 75dfb89 [None][feat] Add a phase-armed Nsight window and a per-rank host-path probe - a3f86ac [None][test] Run the iteration-stats RPC test in the CPU test list - 6cb186f [None][perf] Run the MiniMax-M3.1 TP2 router as an exact half-row TF32 GEMM at qualified row counts - 41cd32f [None][test] Cover the MiniMax-M3.1 TP2 router half-row pin - 82c5e77 [None][perf] Time the executor with GPU timestamps on the MiniMax-M3.1 sm_107 targets - f0af3b6 [None][fix] Drop the Triton stand-in stamp from the executor's GPU timestamps - 8b1d98e [None][chore] Drop the unused separate-library loader from the TP2 router half-row pin - 1d02b9279c [None][chore] Clarify the MiniMax-M3.1 MXDRAFT projection counts in the sm_107 routing log - e28c8a3d0b Add guarded M3.1 TP2 BF16 draft gather and native consumer gate - dfd23153a8 Relocate draft BF16 gather and preparation outside AOT digests - f8c9e69d72 Prewrite CI-listed B1 arm and native byte oracle tests - 2433ec82fa [None][perf] Widen M3.1 BF16 draft gather to measured decode rows Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
What the Rubin (sm_107) context work adds on top of the decode targets: the TP2 context (tp2ctx) target and its NVFP4 context path (the BF16X / E4M3X / Q8KV4 sparse-prefill fork, vendored MSA sources kept source-owned in precompiled installs), the late-compile detector armed after warmup, the tp2ctx route arming the same shared sm_107 opt-ins (and so the same MSA and padding-row holds) as the decode targets, and the ConversationRouter late-binding fix. Source commits: - 33f8296 [None][feat] MiniMax-M3.1 NVFP4 context (prefill) path on Rubin (SM107) - ac47665 [None][fix] Keep the vendored MSA sources source-owned in precompiled installs - 7dec51d [None][fix] MiniMax-M3.1 TP2 context: arm the late-compile detector after warmup - 249513c [None][fix] ConversationRouter late binding: dispatch waiters on every server-list change - 4ab4fc2 [None][chore] Router: state what _after_servers_updated guarantees about preparation - 8008839 [None][fix] MiniMax-M3.1 TP2 context: arm the shared sm_107 opt-ins its role runs with - 89be8d7 [None][fix] MiniMax-M3.1 TP2 context: arm the shared sm_107 opt-ins at routing Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
Collaborator
Author
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.
Depends on #21038 (MiniMax-M3.1 Rubin SM107 decode), which depends on #20917. This branch contains both of their commits; review only the top commit, "MiniMax-M3.1 NVFP4 context (prefill) path on Rubin (SM107)".
Description
Adds the MiniMax-M3.1 NVFP4 context (prefill) path for Rubin (SM107): a TP2 context target for disaggregated context workers, plus the prefill items of the TP2 aggregated target. It also vendors the MSA sparse-attention kernels this path needs and builds them with the wheel.
Everything is default ON on SM107 and exact against the stock path (bitwise, or the same math in a fixed order). Each item has a
TRTLLM_M31_DISABLE_<ITEM>=1kill switch;TRTLLM_M31_DISABLE_TP2CTX=1drops the whole context target.TP2CTX_QKV_FUSED,TP2CTX_QKV_ROWS)TP2CTX_NORMQ,TP2CTX_NORMQ_F32,TP2CTX_GATE_CAST,CTXNORMQ)TP2CTX_SP,TP2CTX_SP_OVL,TP2CTX_SP_PMOVL,TP2CTX_SP_ADDNORM; agg:TP2AGG_SP,AGG_SP_PREMOE)TP2CTX_MXFP8_CUBLASLT,TP2CTX_SWIGLU_MXFP8,AGG_PREFILL_CUBLASLT)TP2CTX_COMBINE_MXFP8,TP2CTX_TSEL,TP2CTX_PROXY_Q128,TP2CTX_SPARSE_OUT,AGG_TSEL,AGG_PROXY_Q128)TP2CTX_ITEMS_IN_BCG,TP2CTX_A1_IN_BCG); prefill output clone for graph output reuse (PREFILL_OUT_CLONE)TP2CTX_SCHED_LPT)PTC_ARRAY_SPLICE,PTC_INT32_REPLY); ConversationRouter per-server in-flight cap, late binding and session TTL (config keys, default off)TRTLLM_M31_ENABLE_TP2CTX_IDX_SCORER=1)Kernels: the C++ ops (SP add-norm, CTXNORMQ, SwiGLU + MXFP8, MXFP8 cuBLASLt
out=, the context peer ops, NVFP4 index-K) are built in-tree, with fake implementations for torch.compile. The MSAfmha_sm100,fmha_sm100_q8kv8andprefill_bf16xkernels are vendored from vllm-project/MSA@25c1447 (MIT; one Apache-2.0 file; lock, license and notice files included) and built ahead of time bybuild_wheel.py. That replaces the externalfmha_sm100package and the native-library fallback loaders.Results (Rubin, MiniMax-M3.1-Flash NVFP4, AgentX SA client)
Client: the SA InferenceX AgentX harness 1.0.12 preset. Acceptance: synthetic (forced) K7 AL 4.04; natural AL not demonstrated. Y = output-export total token throughput per GPU; X = per-user generation tokens/s.
The agg rows ran an image one step older than the disaggregated rows; the difference is context-only (the SP pre-MoE fix, plus a decode-side
register_fake, and their tests). Since the sweep, this commit only adds the ConversationRouter late-binding fixes listed below, which the swept runs did not exercise. dis-X and dis-T are latency cells: the generation servers run batch 1 by configuration.Changes to shared #20917 code
selectPrefillTileSm103): a selected score that is a negative NaN is now emitted as-1(it was emitted as its block id), as the referenceminimax_m3_select_blocksdoes for-inf. Outputs change only when an indexer score is a negative NaN. GB path changed, not GB-hardware-tested.3rdparty/vendor_patches/) instead of the externalfmha_sm100package. GB path changed, not GB-hardware-tested.ctxResidualRmsNormMxfp8Sm103): an optional FP32 side output and the SM107 device check; existing callers get the same kernel instantiation. GB path unchanged.ConversationRouter late-binding fixes (added after the sweep)
With
late_bind, a queued turn could stall the whole queue: one that excluded every live server waited forever, a server removal could leave the queued head without a candidate, and a head waiting for a busy server held back turns that could take the server it excluded. The coordinator path also ignoredlate_bind. Now dispatch walks the queue in order, skipping (not blocking on) a turn whose candidates are all at the cap; a turn with no live candidate fails with the ordinary "No available servers after excluding ..." error; adding or removing a server re-runs dispatch; and the coordinator path rejectslate_bind.Follow-up commits on this PR
TRTLLM_USE_PRECOMPILED) keep the vendored MSA sources from the checkout; only the CUTLASS headers the wheel build stages come from the wheel.Pending updates
Known limitations and follow-ups
3rdparty/fmha_sm100symlink; a follow-up updates it.max_batch_size 1); TP4 / TP8 GEN SwiGLU above 8 rows uses the stock Triton kernel (the native ops are TP2-only).max_batch_sizeaccordingly.nvidia-cutlass-dsl4.8.0 as in upstream requirements;flashinfer-cubin0.6.18 is an optional dependency for Rubin..sobytes).Test Coverage
Run on Rubin (SM107) with a wheel built from this stack,
nvidia-cutlass-dsl4.8.0:tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_107__tp2ctx/: 274 passed, 10 skipped, including the SP pre-MoE boundary tests.tests/unittest/_torch/attention/sparse/msa/test_msa_vendored_imports.py: the vendored MSA packages import without the top-level names, in-process and in a fresh interpreter.register_fakecoverage: every PR op with a CUDA kernel has a fake implementation.FakeMetadatafailures inherited from [None][perf] Add GB300 MiniMax-M3.1 low-latency decode targets #20917.tests/unittest/disaggregated/test_router.py113 passed (all 13 late-binding tests),test_openai_disagg_server.py21 passed,tests/unittest/others/test_precompiled_link_mode.py12 passed, and the TP2 context late-compile test 2 passed.New tests sit next to the existing MiniMax-M3.1 tests for the same target; they are not yet added to a CI test list (see the CI placement note).
CI placement: the SM107 tests need Rubin GPUs, which CI does not run yet; CPU-only cases run on any GPU stage once listed.
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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.