Repository navigation
minimax_m3: a CUDA backend for MiniMax-M3.1-Flash at TP over GB10s, from the release's own formats - #1388
Open
BobClawblaw wants to merge 15 commits into
Open
minimax_m3: a CUDA backend for MiniMax-M3.1-Flash at TP over GB10s, from the release's own formats#1388BobClawblaw wants to merge 15 commits into
BobClawblaw wants to merge 15 commits into
Conversation
… bf16 logits on request tf_route_logits gave each lane four 32-wide slices of k, which covers a hidden width of exactly 4096 (GLM-5.3-Flash and Step-5); at 2048 it read past each row and past the packed router. Each lane now takes hidden / 1024 slices, the same four at 4096, so those routers' arithmetic is unchanged, and `source` refuses a width that is not a multiple of 1024. `round_logits` rounds each logit to bf16 before the select, as a bf16 router matmul (MLX's) gives it.
…niMax XML as DSML)
MiniMaxAI/MiniMax-M3.1-Flash (428B-A23B) as a native Metal family, from checkpoints converted to MLX affine widths (experts 3- or 4-bit per layer, the rest 8-bit; the recipe's tools/convert.py): - Block-sparse attention with the lightning indexer: the index scores and the top-16 block selection a row and key head, decode attention over the kept blocks in 32-key partials, prompt chunks on the tensor units over each tile's kept blocks. 128 routed experts (4 picked, fp32 router with score bias) and a shared expert: decode windows run expert by expert, prompt chunks through the tensor-unit gathers. - Keys and values cached as int8 with a scale a position and head (57% of bf16's bytes) by default; MINIMAX_KV_CACHE=bf16 keeps bf16. Prompt chunks scale scores by the keys' scales and probabilities by the values' before the tensor-unit products, so windows and one-row steps stay bit for bit equal. - Drafts: incoai's DFlash2 (the default) and the release's DSpark head (MINIMAX_DRAFTER=dspark), on the lane core's drafted windows; drafted greedy replies equal undrafted ones. - Images (--vision): the CLIP tower on the shared Qwen-VL kernels (its axial t/h/w RoPE by moved head dims), the merging projector, the rows written over the image tokens; the release processor's 672 x 672 cap. - The prompt cache's snapshots, several streams in one decode window, and a load budget under MINIMAX_LOAD_LIMIT_GB (default 75% of RAM). Measured on one M5 Ultra (256 GB): against transformers' math on 4 layers KL <= 0.005; against the release's NVFP4 weights over a 1,024-token chat, KL 0.250 and top-1 80.0% (mixed 4/3-bit, int8 cache); a needle found in a 251,713-token prompt (the 3-bit mix with 4-bit layers 22, 24, 59 at a 262,144 window); the image tower against transformers' at a relative error of 0.030 to fp32 (torch's own bf16: 0.039).
… (one call block, xgrammar's minimax_m3_xml arguments)
…d the recipes index)
…96-row prompt chunks (same bits) A prompt tile now holds the 16 query heads of one key head on two rows instead of one query head on 32 rows. The heads share their rows' block selection, so each key and value block is read once for all of them, and a tile walks the union of two rows' blocks rather than 32's. Prompt chunks grow from 2,048 to 4,096 rows, so each expert's weights are read once for twice the rows. Last-row logits after an 8,192-token prompt are byte-identical to the previous kernel's with the int8 and the bf16 cache.
…are memory (gate|up and act with the experts', their outputs and the block scores with the sorted inputs), so 4,096-row chunks keep the 262,144-token window under 226 GB
…adgroup atomics), not one thread an expert scanning every pick
…d serves 8 experts; every logit's sum in the same order)
…ut share the attention output's and normed q's memory (each dead before the other is written), so mix34 keeps its 65,536-token window under 226 GB
…nd the 246,842-token needle
…he and MINIMAX_KV_CACHE removed) The int8 cache (a scale a position and head) changed the model's arithmetic, so it goes: keys, values and index keys are cached as bf16, as the checkpoint computes them. The prep and decode kernels lose their int8 variants and the prompt attention its KV8 path; the cache layout loses the scale planes. What stays is the old bf16 path unchanged: on mix256k the 8k prompt's last-row logits, every 16-row window's logits over a 2k prompt and 64 greedy one-row steps are byte-identical to the previous head with MINIMAX_KV_CACHE=bf16. A token takes 135 KiB of cache across the 60 layers (int8: 77 KiB), so under MINIMAX_LOAD_LIMIT_GB=226 on a 256 GB M5 Ultra the caches fit 150,400 tokens on mix256k and 37,376 on mix34 (--parallel 2 shares them).
…rom pinned sources, and their hashes MiniMaxAI/MiniMax-M3.1-Flash at 54e4e69d (mxfp8 projections, nvfp4 routed experts) and incoai's DFlash2 at ac07a7c1 to the minimax_m3 checkpoint: projections MLX affine 8-bit g64 from their exact bf16, experts a least-squares affine refit of the nvfp4 values. `--variant mix34|mix256k` sets each published variant's 4-bit (g128) layers; `--check` compares a folder with minimax_m3_checkpoints.json, the sizes and sha256s of every file at the published revisions (config.json by content and the shard index by its weight map: the published ones were reformatted and carry stale total_size fields). A full mix256k rebuild (mlx 0.32.4, numpy 2.5.4) checks 74/74. mix34 was not rebuilt; its download checks 74/74.
…edle) and how to rebuild its checkpoints The windows the exact cache fits under MINIMAX_LOAD_LIMIT_GB=226 (mix256k 147,456 of 150,400 tokens, mix34 32,768 of 37,376), KL against the NVFP4 release on this head (mix34 0.246, mix256k 0.301), DFlash2 decode, cold prompts and a 145,203-token needle on mix256k, all measured on the bf16 build; the recipe page's commands for minimax_m3_convert.py.
…ative MXFP8/NVFP4, DSpark, images, kernel contract)
7 of 8 tasks
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.
Stacked on #1235 (the family, its config, tokenizer handling, chat template and tool-call markup on Metal): this adds
the family's CUDA backend, so the same
minimax_m3family serves MiniMax-M3.1-Flash on two DGX Sparks at TP=2, asSynthema asked on #1281 (one family, a CUDA and a Metal backend). Nothing outside the family changes behaviour for other
families: the shared expert kernels gain a SwiGLU-OAI epilogue (new template instances, a
Layer.oaiflag off bydefault) and two of
experts.zig's tests move toexperts_test.zigto keep it under 600 lines.What it reads. The release's MXFP8 projections and NVFP4 routed experts as published. Two GB10s don't hold the
release (216.4 GiB of experts), so the folder keeps NVFP4 on 15 MoE layers and a 3-bit affine refit of the same NVFP4
values on the other 42 (
tools/zig/minimax_m3_nvfp4_mix.py, from the pinned revision, with--checkhashes); thelayers are the 15 whose routed outputs lose most at 3 bits (
tools/zig/minimax_m3_ref.py --capture,tools/zig/minimax_m3_requant_study.py).What it serves. TP over hosts (
TF_MINIMAX_RANK/TF_MINIMAX_WORLD, NCCL), the lane core's shared rounds (2streams at 32k on two GB10s), the release's DSpark drafter, kept prompt states at 2,048-token chunk starts, the
keyed sampler, structured output and forced tool calls (grammar masks on whole rows), images (the tower on the last
rank). Bf16 key/value cache, exact.
How it was checked. (receipts: recipe harness, target minimax-m31-nvfp4mix-gb10x2)
tools/zig/minimax_m3_ref.py):within the spread between two valid references (the MoE sum's rounding moved) over 20 layers, and over 3,026 tokens
where the sparse selection is active.
"draft": false(greedy, sampled with a seed, two concurrent streams); concurrent == solo; resumed ==fresh after a restart; seeded sampling repeats.
max-tokens 3/3, thinking switch 10/10, drafted == plain 12/12, structured output 18/18, tool calls 5/5, seeded
sampling 9/9. A prompt cancelled mid-prefill leaves the next request right; an image beside a second stream reads the
same as the image alone.
(0.933 / 0.950), GSM8K 0.972 (0.976 / 0.968), MGSM 0.838 (0.821 / 0.846), MMLU 0.836 (0.819 / 0.827), IFEval 0.587
(0.593 / 0.580), JSON schema 1.000 (1.000).
tf-cuda-test minimax-contract(also inkernel-contract), 10 kernels pass gates 1-4 and the 4planted controls go red.
first token 0.47-0.51 s, 2.7 tokens a round; two streams 29.8 / 49.3 / 32.6 / 53.9 in total. Weight reads run at
about 84% of the GB10's measured read floor.
Not covered
SwiGLU-OAI expert instances are not registered yet. Gate 5 records cover both MiniMax images (nvcc 13.0 and 13.2).
thinking-off suites (tools, GSM8K, MGSM, MMLU, IFEval, JSON) are within the M5 Ultra recipe's run-to-run range.