MiniMax-M3 support — GQA + MSA block-sparse attention, converter (follow-up to #418) - #601
MiniMax-M3 support — GQA + MSA block-sparse attention, converter (follow-up to #418)#601steve-m wants to merge 22 commits into
Conversation
|
Really promising direction (GQA + MSA block-sparse + o200k), and thanks for pushing it. Current status so it's clear where it stands:
To move it forward:
Flag it ready-for-review (un-WIP) once it's rebased and split, and we'll go through it. No rush — better landed cleanly than fast. |
|
Rebase request — and an apology for how many of these you have had.
Before asking, three things changed on our side so that this is the last one of these you should need for a while. 1. The main source of these conflicts is closed. Gates are now derived from the build rules (#733). Adding a test means adding your 2. We resolved what we could ourselves instead of asking. Eight PRs were unblocked by a maintainer pushing the merge to the contributor's branch rather than requesting a rebase — including three where both sides carried a real change and had to be merged rather than picked. No commits were rewritten. Yours is here because its conflict is in engine code ( 3. Merge order is oldest-clean-first from now on. A PR that is green and unconflicted merges ahead of anything opened after it. The reason some of you rebased many times is that newer, smaller PRs kept jumping the queue and resetting you — a starvation loop we built, not bad luck on your side. What we need: one rebase onto current If you would rather not, say so and we will close it with thanks and the branch stays yours to reopen. No pressure either way, and no hard feelings — several of these have been open a while through no fault of the author. |
quant: AVX2 kernel for int3-g64 matmul (cherry-pick from #601)
|
Heads up — I implemented three of the WIP gaps listed here as a PR into your
Verified: MSA oracle still 24/24 + 20/20, kill-switch flips to full attention, tool-call render→parse round-trips (nested objects/arrays + schema type coercion), and |
steve-m
changed the title
|
Rebased onto current dev (v1.6.1) and un-WIPed. The split you asked for resolved itself in the meantime: the Vulkan backend landed via #418/#705, so this PR no longer carries one — it's now 16 commits, +1,335/−76 over 12 files (was +5,050): the MiniMax-M3 engine/converter/oracle as one piece, plus small M3 enablement of the in-tree VK backend (GQA attention core shader, swigluoai in the fused gate+up shader, a Dropped during the rebase as already upstream: everything that merged with #418/#705, the mirror/pilot commits, the space-separated-BPE-merges tokenizer fix (dev's parser is stricter and covers o200k), and the AVX2 int3-g64 kernel — which you already took as #926 (thanks @rofl0r for the ping; it landed with authorship preserved). Core-file blast radius, called out explicitly:
Also folds in @corzogac's contribution (steve-m#1, authorship preserved): the Validated on v1.6.1 (Ryzen 3900X + RX 9070/RADV): full C test suite; the in-tree tiny-model oracle chain ( |
… rendering Three of the WIP gaps called out in JustVugg#601: 1. COLI_MSA=0 A/B kill-switch (colibri.c). Force full causal attention even when the MSA indexer weights are present, by zeroing c->msa + idx_type the same two lines the indexer-absent auto-detect fallback uses — so it rides the already-tested full-attention path for measuring MSA's long-context win. 2. Banner/PROF cosmetics (colibri.c). The startup banner and the PROF attention-bound verdict now name MSA/M3 instead of hardcoding GLM/DSA (banner peeks config.json model_type before model_init; PROF branches on m->c.arch). 3. Tool calls for the MiniMax-M3 template (openai_server.py). Render the ]<]minimax[>[<invoke name=…> XML dialect — the developer <tools> declaration, assistant tool_calls, and ]~b]tool <response>… tool results — and parse it back, so render_chat no longer returns 400 for `tools` and parse_tool_calls dispatches to an M3 parser. Mirrors chat_template.jinja (to_xml recursion, nested objects/arrays, schema type coercion via _tool_param_types). Verified: MSA oracle still 24/24 + 20/20; COLI_MSA=0 flips to full attention; tool-call render/parse round-trips (nested objects/arrays + type coercion); full `make test` suite 167/167 OK (13 skipped).
|
I've reviewed this end-to-end and, importantly, ran it rather than reading it. Starting with what works, because it's substantial:
The port is correct and GLM does not regress. Your The gaps are all in the plumbing around it. Four of them block, and one affects people who have nothing to do with MiniMax. B1 —
|
… rendering Three of the WIP gaps called out in JustVugg#601: 1. COLI_MSA=0 A/B kill-switch (colibri.c). Force full causal attention even when the MSA indexer weights are present, by zeroing c->msa + idx_type the same two lines the indexer-absent auto-detect fallback uses — so it rides the already-tested full-attention path for measuring MSA's long-context win. 2. Banner/PROF cosmetics (colibri.c). The startup banner and the PROF attention-bound verdict now name MSA/M3 instead of hardcoding GLM/DSA (banner peeks config.json model_type before model_init; PROF branches on m->c.arch). 3. Tool calls for the MiniMax-M3 template (openai_server.py). Render the ]<]minimax[>[<invoke name=…> XML dialect — the developer <tools> declaration, assistant tool_calls, and ]~b]tool <response>… tool results — and parse it back, so render_chat no longer returns 400 for `tools` and parse_tool_calls dispatches to an M3 parser. Mirrors chat_template.jinja (to_xml recursion, nested objects/arrays, schema type coercion via _tool_param_types). Verified: MSA oracle still 24/24 + 20/20; COLI_MSA=0 flips to full attention; tool-call render/parse round-trips (nested objects/arrays + type coercion); full `make test` suite 167/167 OK (13 skipped).
steve-m
changed the title
|
Thank you for a review that came with measurements attached — B1's stop-filter chain and B2's B1 — one EOS resolver for every entry point (
|
| Check | Result |
|---|---|
make colibri / make colibri VK=1 (no suppression flag) |
clean, 0 warnings |
make m3-tiny-check (the new gate, end-to-end incl. new converter) |
OK — prefill 24/24, decode 20/20 |
TF_LOGITS vs oracle, IDOT=0 |
max|d| = 2.62e-06, 0 argmax flips |
VK expert tier (int4 fixture, COLI_PREFILL_CHUNK=2 so the S≤4 tier actually serves) |
VK-BLOCK nvk 2.00 / ncpu 0.00, CPU-vs-GPU logits bit-identical |
make test-c |
all pass |
make test-python |
530 tests, OK (32 skipped) |
One methodological note from that table: a plain TF prefill run never exercises the VK expert tier (it declines S>4 batches), so the earlier tier numbers and these were both taken with small chunks forced — worth knowing for anyone re-verifying.
|
Status, and a heads-up that will change the shape of your next rebase — mostly in your favour. Where this stood: your 2026-08-16 round closed all four blockers with measurements attached, re-rebased onto What changed underneath you, specifically, because it makes some of your diff obsolete rather than conflicted:
On v1.7.0: this is not in it, and should not be. A 1,652-line seventh-model-family PR going in on release evening is how a release gets a bad name — and the release was already cut when your last round landed. It is first in the queue after the tag, not last. What we can do to make this not be your fourth rebase: say the word and we will carry the mechanical part ourselves — a branch with your commits and authorship intact, conflicts resolved, each decision stated for you to review rather than perform. We did exactly this for #1024 → #1082 and #790 → #1113 this week. You would review the registry migration (where your judgement is actually needed) instead of re-fighting Either way, the MSA block-sparse work is wanted. It has just had the bad luck of landing in the busiest week this repo has had. |
|
@steve-m, please rebase |
… rendering Three of the WIP gaps called out in JustVugg#601: 1. COLI_MSA=0 A/B kill-switch (colibri.c). Force full causal attention even when the MSA indexer weights are present, by zeroing c->msa + idx_type the same two lines the indexer-absent auto-detect fallback uses — so it rides the already-tested full-attention path for measuring MSA's long-context win. 2. Banner/PROF cosmetics (colibri.c). The startup banner and the PROF attention-bound verdict now name MSA/M3 instead of hardcoding GLM/DSA (banner peeks config.json model_type before model_init; PROF branches on m->c.arch). 3. Tool calls for the MiniMax-M3 template (openai_server.py). Render the ]<]minimax[>[<invoke name=…> XML dialect — the developer <tools> declaration, assistant tool_calls, and ]~b]tool <response>… tool results — and parse it back, so render_chat no longer returns 400 for `tools` and parse_tool_calls dispatches to an M3 parser. Mirrors chat_template.jinja (to_xml recursion, nested objects/arrays, schema type coercion via _tool_param_types). Verified: MSA oracle still 24/24 + 20/20; COLI_MSA=0 flips to full attention; tool-call render/parse round-trips (nested objects/arrays + type coercion); full `make test` suite 167/167 OK (13 skipped).
|
Thank you for the status — and for the offer, which I'm declining for a reason that is entirely to your credit: the mechanical part turned out to be small, and the part that needed judgement was the registry migration, which is mine either way. A trial merge against The registry migration (#1063/#1068/#1103)You were right that most of the diff should delete. Finding
The one decision I'd like you to look at. Several launcher gates were written def colibri_core(family):
return family.engine_group == family_by_id("glm").engine_group
Two gates deliberately keep the narrow Your segment-conformance gate caught the other consequence: Serve codec (your point 2)Our diff does not touch serve framing — the only serve-side change is the B1 EOS resolver. Worth noting while I was in there:
|
| Check | Result |
|---|---|
make colibri / make colibri VK=1, no suppression flag |
clean, 0 warnings each |
make m3-tiny-check (the new gate, end to end incl. the converter) |
OK — prefill 24/24, decode 20/20 |
TF_LOGITS vs oracle, IDOT=0 |
max abs diff 2.62e-06, 0 argmax flips |
| Same gate under the VK binary | OK — 24/24, 20/20 |
| VK expert tier, CPU vs GPU (int4 fixture) | tier served all 20 decode steps (VK-BLOCK: 20 blocks, avg nvk 2.00 / ncpu 0.00); per-position decode output identical to the CPU arm |
make test-python |
695 tests, OK (35 skipped) |
make test-c |
ALL PASS, 0 failures |
A correction to my last round's methodological note. I wrote that the tier needs
COLI_PREFILL_CHUNK<=4 to serve during the gate. That was wrong, and if you tried it
you would have measured nothing: forward_all() calls layers_forward() directly, so
chunking in step() never reaches the teacher-forcing path — with COLI_PREFILL_CHUNK=2
and the decode arm switched off, VK-BLOCK prints zero blocks. What actually exercises
the tier is TF_DECODE=1: decode runs at S=1, inside the tier's S<=4 serving window.
The counters I reported last time were real, but they came from the decode arm, not the
chunk setting. The prefill arm of the gate is CPU-only by construction.
Worth stating plainly: an exact-zero CPU-vs-GPU diff on the prefill logits is not
evidence of agreement here — it is what "the tier never ran" looks like. The claim above
rests on the VK-BLOCK counters plus the per-position decode comparison.
One request, given what this round taught us both: if the shape is acceptable, a merge before dev moves another 80 commits would save a fifth rebase. If you would rather I split anything out — the colibri_core change is the obvious candidate for a separate GLM-side PR — say the word and I will.
|
This is the best rebase report I have read on this repo, and the work under it holds up. I reviewed it properly rather than trusting the summary, and I want to be specific about what I checked, because the thing that made this reviewable at all is that you told me where to look. What I verified, and why it was the right thing to checkM3 riding the GLM binary means every shared path is a risk to our flagship model, so I traced each one you touched:
The registry migration is the part I would have expected to go wrong and did not. I checked the claim that Two more things worth naming, because both are choices rather than accidents:
And declining the rebase offer for the reason you gave was the right call. One thing to fix before this mergesThe header comment on
But the body computes the indexer, picks the top-k blocks, and I followed It is not a correctness problem, and I would still be asking for it: in this tree comments carry as much weight as code, and the next person to read that function - to port it to a GPU backend, say - will trust it. Worth a minute. One thing neither of us could have seen alone
It is not a live bug: no shipping binary registers those adapters, and your PR predates both. But when M3 gets its Segment adapter, those globals need to hang off the model rather than the process. A comment at the declaration saying so costs nothing now and saves someone a genuinely nasty afternoon later. Smaller notes
Validation24/24 teacher-forced prefill and 20/20 decode token-exact against an independent numpy oracle, TF_LOGITS within 2.62e-06 with zero argmax flips, the VK tier serving all 20 decode steps with per-position output identical to the CPU arm, zero warnings on both builds, and the full suite green. That is the standard, and it is met. Thank you for carrying this for five weeks and for coming back to it with the diff smaller than it started. Fix the comment, add the globals note, rebase, and I will merge it. |
Maps the minimax_m3_vl VL checkpoint onto the GLM container scheme so the engine's MoE loader stays single-naming: strips the language_model. prefix, block_sparse_moe -> mlp (router gate.weight + e_score_correction_bias line up with the existing loader), Mixtral w1/w3/w2 -> gate/up/down_proj, drops the vision tower + the MSA index branch (a later pass, like DSA), keeps q/k/v_proj under the attn bits class, q_norm/k_norm as f32. Writes a flattened text_config as the container config.json and carries chat_template.jinja. Validated on a synthetic tiny checkpoint: full name flow + bit-exact int4-g64 vs quant_int4_grouped.
…igluoai, Gemma norms
model_type minimax_* in config.json switches the engine to the M3 family:
- load_cfg: M3 key set (num_local_experts, head_dim, num_key_value_heads,
rotary_dim, dense/shared intermediate, swiglu params, moe_layer_freq ->
first_dense); the GQA KV rows ride the MLA cache aliases (Lc=K, Rc=V,
kv_lora/qk_rope = n_kv_heads*head_dim) so kv_alloc/bind/persist/mux work
unchanged; arch-split config validation.
- loader: self_attn.{q,k,v,o}_proj + per-head q_norm/k_norm for ARCH_M3;
MoE/router/shared names already line up via the converter mapping.
- attention_gqa(): per-head Gemma QK-norm BEFORE partial split-half (NEOX)
RoPE on the first rotary_dim dims, standard per-token KV rows, grouped
scores/softmax/values (repeat_kv h -> h/(H/NK)), o_proj; honors the ragged
kvs/positions contract; full causal attention (MSA block selection exact
<=2048 ctx, deferred beyond).
- act_glu(): all 11 silu(g)*u sites now dispatch silu vs swigluoai
(clamps, alpha, and the (up+1) factor); rmsnorm gains the Gemma (1+w)
variant behind g_gemma_norm. Both default off -> GLM math untouched.
K3's parallel combine moves from expert_ffn's body into act_glu, where it
covers both activations and every caller; elementwise, so bit-identical.
- registry: a FamilyDescriptor entry, not per-file branches. minimax_m3 (and
the raw minimax_m3_vl wrapper) resolve to engine_artifact "colibri" with
internal_arch "minimax_m3" -- the shared binary self-dispatches on
config.json, and the descriptor keeps the two families distinct where it
matters (template, planner, limits, model id). _minimax_geometry() mirrors
kv_alloc's GQA rows + the MTP row and attention_gqa()'s scratch set; the
registry's own minimax double is superseded by the shipped descriptor, so
its test now pins the production geometry. Registering a family also
obligates a segment-conformance row (release_policy
"all_registered_families"), so the manifest, the C fixture matrix and its
roster gain minimax_m3: GQA KV + the MSA indexer, which reuses the DSA
indexer's own cache slots in the engine. The REAL segment adapter is a
separate subsystem and stays at six; the manifest records that M3 needs one.
- VK: expert tier + fused shared expert gated off for M3 (shaders hardcode
silu); the per-matmul dense chain stays available.
- chat: MiniMax template (]~!b[ / ]~b]role / [e~[) in the serve dialog loop,
coli run, and openai_server (text subset; tools 400 for now); [e~[ EOS
fallback.
- tools: make_m3tiny.py (tiny random M3 checkpoint) + oracle_m3.py (numpy
reference on the DEQUANTIZED container -> the engine's REF/TF gate).
TF_LOGITS=1 prints every teacher-forcing logit row (oracle bisection); TF_DECODE=1 re-runs the same REF continuation through the incremental S=1 decode path (fresh kv_alloc, prefill np then per-token steps) and compares argmax against tf_pred — validates KV append across calls, which the batch check alone cannot see. M3 tiny results: prefill 24/24 (max logit diff 3.2e-6 vs the numpy oracle with IDOT=0 exact kernels; the idot default's activation quant explains the earlier 23/24), decode 20/20.
… M3 expert tier
The fused gate+up shader hardcoded silu(gate)*up, so the VK expert tier was gated
off for MiniMax-M3 (swigluoai) and its experts ran on the CPU. Add a push-const-
selected activation: act=1 computes (up+1)*gate*sigmoid(alpha*gate) with gate/up
clamped to +/-limit (alpha 1.702, limit 7.0), matching act_glu and the numpy oracle;
act=0 keeps GLM's silu bit-identical. The activation is model-global, set once via
coli_vk_set_activation from vk_registry_fill (g_act_swigluoai -> the mode), read into
every gate_up dispatch's push constants (both the G and dev2 G2 paths).
Un-gates vk_registry_fill for M3. Validated on the real 225 GB int4-g64 container,
RX 9070: expert tier fills (384 resident, 10.9 GB, vk 40% of lookups), output stays
correct ('Paris' greedy + coherent prose), decode 0.64 -> 0.91 tok/s (+42%) with the
tier active. Attention (GQA) and the shared expert still run on the CPU (P4).
New attention_gqa.comp: one workgroup per (query row, head) computes
scores/softmax/weighted-V for GQA (query head h reads KV group h/(H/NK)),
reading the persistent on-device K/V mirror — K in the L buffer, V in the R
buffer, both NK*hd wide (the mirror already stores two arbitrary-width rows;
the absorb shader's R<=64 cap is its own, not the mirror's). coli_vk_gqa_attn
mirrors the absorb dispatch (5 bindings, one submit+fence+readback per layer).
Wired into attention_gqa behind COLI_VK_ATTN, single-sequence decode only,
CPU fallback on any failure.
Validated on the real 225 GB container, RX 9070: output correct ('Paris'
greedy + coherent prose) with the core active. HONEST PERF NOTE: it is
throughput-NEUTRAL on this box and stays opt-in/default-off, because the core
is only ~0.7s of an ~8.6s decode-attention cost — the projections dominate
(q/k/v proj+norm+rope ~4.8s, o-proj ~2.9s) and are still on the CPU. The per-
layer submit/readback slightly exceeds the tiny core it replaces. Offloading
the projections (fused, GLM-qprep style, to amortize submits) is the actual
lever and the next step. -Wno-missing-field-initializers quiets the shared
struct PC's zero-filled tail.
…evice)
coli_vk_gqa_attn_project runs the GQA core and the o-projection in one command
buffer: the core writes ctx into the att_ctx device scratch (never read back), a
compute barrier, then the o-proj matmul reads it and only [S,hidden] returns to
the host — the absorb_project pattern, for GQA. attention_gqa tries it first
(fills out directly, skips the CPU o-proj), then the core-only path, then CPU.
Validated correct ('Paris' greedy). Profile confirms the offload: decode
output-projection drops to 0.000s on the CPU, and the core+o-proj VK dispatch
(~2.3s/64tok) replaces CPU core+o-proj (~3.9s) — a ~1.6s attention saving.
HONEST e2e NOTE: still throughput-neutral on this box (fused 1.68/1.71 vs CPU
1.70/1.73). M3 decode here is expert-matmul + disk bound (expert-matmul ~16-17s,
disk-wait ~6s, attention ~9s of a ~34s decode; the Zen2 CPU is AVX2-maxed with
no VNNI, and the 16 GB VRAM caps the expert tier at 512). The saving is real but
masked by expert/disk variance. Kept behind COLI_VK_ATTN (default off for M3);
it pays off on a compute-bound or faster-disk box and sets up full projection
fusion.
M3's native attention is block-sparse beyond 2048 tokens: a small 4-head indexer
scores every key, max-pools into 128-token blocks, and the main GQA attention only
reads the top-16 blocks (+ the local block) per KV group. Our port ran full causal
attention (exact <=2048 ctx, divergent + O(ctx) beyond). This adds the real MSA.
Converter: keep + quantize the 228 indexer tensors (index_{q,k}_proj int8,
index_{q,k}_norm f32) instead of dropping them; --indexer mode gains M3 support for a
supplemental add-on pass (no full re-convert). The sparse config already rides the
flattened container config.
Engine: reuses the DSA index-key cache (Ic) + index_hd/index_nh/idx_type slots (M3's
sparse layers == MoE layers). load_cfg reads sparse_attention_config (index dim 128,
4 heads, block 128, top-16, local 1); guards the GLM DSA index read from clobbering
M3's fields. attention_gqa: the indexer (Gemma-norm + same partial-ROT rope as the
main attention, idx_k mirrored to Ic) runs per sparse layer, block-scores in f64, and
picks top-k blocks (local forced in; greedy ties -> lowest index). The main core then
attends only the selected blocks. VK attention is skipped on sparse layers.
Validated on a tiny indexer-bearing checkpoint (block_size 4 so 24 tokens span 6
blocks): engine vs numpy oracle bit-exact — PREFILL 24/24 + incremental DECODE 20/20
(IDOT=0). The 2 near-tie mismatches during bring-up were a config bug (JSON-bool
use_sparse_attention + the unguarded GLM index read zeroing index_hd), now fixed.
…ights absent A container converted before the indexer was kept still announces sparse_attention_ config, so the loader would try (and fail) to load index_q_proj on every sparse layer. Probe the first sparse layer's indexer weight; if absent, disable MSA (full causal attention, exact <=2048 ctx) and point the user at the --indexer pass. Same auto-detect discipline as the DSA indexer / MTP.
Three fixes so a container converted before the indexer was kept can get it added without a full 34-min re-convert: - drop the '--arch m3 --indexer not implemented' guard (only --mtp is refused now); - the shard pre-filter matched GLM's 'indexer' name only — also match M3's 'self_attn.index_' so the right shards are read; - layer_idx() only parsed 'model.layers.N...'; find 'layers.N' anywhere so it also handles M3's 'language_model.model.layers.N...' VL prefix. Extracts the 228 indexer weights (+scales) to out-idx-*.safetensors; drop them beside the main container (st.h scans any *.safetensors). Applied to the box container + mirror: MSA now active on the real model (no [MSA] fallback, output correct).
kv_persist.h gated the Ic rows on has_dsa (GLM DSA), which is false for
M3 (index_topk=0), so .coli_kv persisted K/V but not the MSA index keys:
a resumed conversation carried uninitialized Ic for every restored
position, silently corrupting block selection once the conversation
passed 2048 tokens (topk_blocks * block_size). Gate on (has_dsa || msa)
— the same condition kv_alloc uses for Ic — in the header, record size,
append and load. Old M3 .coli_kv files now mismatch the header
(h[3] 0 -> index_hd) and take the existing 'different model or version'
reject path (start over, no misparse).
Validated on the real 225 GB container:
- record size 245764 -> 274948 B/token (= +57 sparse layers x 128 f32)
- old-format file: clean reject + start-over
- end-to-end: a 2363-token needle document (needle in block 15 of 19)
saved by one process; a FRESH process resumed 2365 tokens (0.1 s, no
re-prefill) and answered the needle question exactly. Without the fix
the zero-key tie-break selects blocks {0..14, local} and drops block
15, so this recall only works through the restored index keys.
M3's GQA has no MLA low-rank bottleneck: q_proj and o_proj are 50 MB of int8 each across 60 layers, ~6 GB of weight reads EVERY token that rode the same ~28 GB/s DDR4 the routed experts stream over — the reason a 428B/23B-active model barely beat the 744B GLM (whose attention/dense lives in VRAM via CUDA_DENSE). matmul_qt_ex had Metal and CUDA branches but no Vulkan one; COLI_VK_DENSE only wired GLM's MLA call sites. Add a generic VK branch to matmul_qt_ex behind COLI_VK_DENSE=1, gated on a per-tensor vk_gemm flag set at load for M3's resident weights (q/k/v/o, indexer, dense MLP; shared experts already had wired sites) — routed expert QTs are slab-transient and never marked. At decode only tensors >= COLI_VK_GEMM_MB (default 8) are worth the per-call submit; a batched forward (S>=8) offloads every marked tensor. Weights upload lazily at the dense priority class (0.75 > tier 0.4), so the launcher reserve grows to 8 GB and the expert tier self-sizes into the rest. Measured (RX 9070, 512->~170-200-expert tier, 2-drive mirror, ngen 64): decode topp 0.7: 2.12-2.20 -> 2.6-2.84 tok/s (+25-30%) decode topp 0.5: 2.49 -> 3.0-3.10 tok/s prefill 516 tok: 152.7 -> 102.3 s (-33%); attention projections -67% Output verified (greedy Paris + coherent lighthouse prose). Launchers updated to COLI_VK_DENSE=1 RESERVE=8 and the stale MSA caveat rewritten.
Four levers on top of the VK dense matmul (95a565d):
- lm_head vk_gemm: the 1.2 GB int8 logit matmul rides COLI_VK_DENSE.
- q+k pair: k_p's 3 MB joins q_p's submit via vk_matmul_pair_qt (a lone
k submit costs more than its CPU read; paired it is free).
- fused shared expert un-gated for M3: the gate_up shader takes the
activation as a push constant since 8dd7fae, and fmt=4/g64 rides the
same grouped path the expert tier uses — the !g_act_swigluoai and
fmt gate predated both. 3 submits/layer -> 1 across 57 layers. Also
set the activation BEFORE vk_registry_fill's early returns, so a
tier-less run (COLI_VK_EXPERTS=0) can not dispatch silu on an
swigluoai model.
- MSA selection: t-outer loop (each cached key row read once for all 4
index heads), AVX2 4-lane f64 FMA scoring dot (the 64K-context decode
hot spot), OMP over prefill rows (was single-threaded on 12 cores).
Oracle gate after the selection rewrite: PREFILL 24/24 (exercises the
OMP path), incremental DECODE 20/20 (IDOT=0, tiny checkpoint).
Measured (RX 9070, R8 tier 168 + dense, 2-drive mirror, ngen 64):
decode topp 0.5: 3.02-3.10 -> 3.49-3.52 tok/s
decode topp 0.7: 2.78-2.84 (unchanged — 154 experts x 28.3 MB =
4.4 GB/token sits ON the dual-SSD read floor; compute offloads only
pay below it, which is the standing argument for topp<=0.55 or int3)
prefill 516 tok: 102.3 -> 96.8 s (152.7 before VK dense)
RAM_GB=58 probe: no gain (hit 84.8->85.9, wash) — RAM budget is spent.
Quality: Paris + coherent lighthouse prose at 0.5 and 0.7.
…ain pass Since the indexer weights are kept in the main pass (MSA), they fell into the generic resident class and silently followed --ebits. Block selection is discrete: quantization noise flips top-k choices rather than blurring them, and the validated real-model MSA config used int8 (supplemental pass). New 'idx' class + --idx-bits (default 8) pins them regardless of ebits/xbits; the tensors are tiny (~215 MB) so the cost is nil. Flag-path smoke on the tiny checkpoint: --ebits 4 --xbits 3 --group-size 64 converts (routed experts fmt=5 int3-g64, [MIXED] idx=8bit) and the engine loads and runs the result.
…xact compare The docstring predates the rebase: dev's engine takes the model dir as SNAP=<dir> (COLI_MODEL is not an engine env upstream). Also spell out that the T/T expectation holds under IDOT=0 -- the exact f32 kernels reproduce this oracle to ~1e-6 max logit delta, while the default int8-activation IDOT path can flip the argmax at a borderline position of the random tiny fixture (measured: one 0.058-margin position at 23/24), which is activation quantization, not a port defect.
… rendering Three of the WIP gaps called out in JustVugg#601: 1. COLI_MSA=0 A/B kill-switch (colibri.c). Force full causal attention even when the MSA indexer weights are present, by zeroing c->msa + idx_type the same two lines the indexer-absent auto-detect fallback uses — so it rides the already-tested full-attention path for measuring MSA's long-context win. 2. Banner/PROF cosmetics (colibri.c). The startup banner and the PROF attention-bound verdict now name MSA/M3 instead of hardcoding GLM/DSA (banner peeks config.json model_type before model_init; PROF branches on m->c.arch). 3. Tool calls for the MiniMax-M3 template (openai_server.py). Render the ]<]minimax[>[<invoke name=…> XML dialect — the developer <tools> declaration, assistant tool_calls, and ]~b]tool <response>… tool results — and parse it back, so render_chat no longer returns 400 for `tools` and parse_tool_calls dispatches to an M3 parser. Mirrors chat_template.jinja (to_xml recursion, nested objects/arrays, schema type coercion via _tool_param_types). Verified: MSA oracle still 24/24 + 20/20; COLI_MSA=0 flips to full attention; tool-call render/parse round-trips (nested objects/arrays + type coercion); full `make test` suite 167/167 OK (13 skipped).
…arker Three adjustments on top of the tool-call support from fork PR #1, all measured against the official chat_template.jinja rendered via jinja2: - invoke/to_xml rendering: the template's for-loops are whitespace-trimmed on both sides, so argument pairs and nested elements concatenate with NO newlines; dropped the extra \n after <invoke name="..."> and between argument pairs (the parser was already layout-agnostic). - developer-block example: the template's full two-invoke example including the nested param-2 item block, byte-identical. - _tool_stream_markers: minimax holds the visible stream at the ]<]minimax[>[<tool_call> opener. The arch-dispatched streaming scan landed on dev after this PR's base; without a marker a streamed reply would leak the raw tool-call block to the client token by token. Validated against the jinja reference render: tools declaration, assistant tool_call block, and tool-response run all byte-identical; parsing the template-rendered block round-trips nested objects/arrays/bools with schema type coercion.
The [e~[ fallback lived in run_serve only, but openai_server launches the engine with SERVE_BATCH=1 -> run_serve_mux, and coli run goes through run_text: both kept eos=-1 on an M3 tokenizer. With -1 the JustVugg#401 stop filter in sample.h is skipped and every tokenizer special token is armed as a hard stop, so the first tool-call marker ends the turn. One tok_eos_resolve() used by all three entry points. (Review B1.)
build_plan sized KV from the MLA keys only (kv_lora_rank/qk_rope_head_dim/ qk_nope_head_dim/v_head_dim), so an M3 container planned with kv_bytes=0, kv_buffer=0 and configured_experts=0: plan/doctor/--auto-tier under- reserved ~1.2 GB at 60L x 4096 ctx and the expert-cap clamp never fired. Since JustVugg#1103 those formulas belong to the family descriptor, so this lands in _minimax_geometry rather than as a branch in build_plan -- the planner itself is untouched. The geometry mirrors the engine: Lc/Rc hold full K and V rows (num_key_value_heads*head_dim each) over n_layers+1 rows, the MSA index-key cache exists only on the sparse layers and without the MTP row (kv_alloc allocates Ic over n_layers and skips !idx_type), and the sparse set is derived the way load_cfg derives it -- from moe_layer_freq, because M3's sparse attention layers ARE its MoE layers. Workspace is attention_gqa's scratch set plus the indexer's own projections. Two test layers: the registry test pins the arithmetic (and now asserts against the shipped descriptor instead of the double the registry landed with), and test_resource_plan runs a real M3 container end to end, so it also fails if build_plan ever stops routing M3 through the geometry. MLA arithmetic is unchanged. (Review B2.)
Two silent behavior changes for people not using MiniMax leave with this: - params manifest: bits_map["idx"] was always recorded (default 8) and "arch" always written, so check_or_record_params refused to resume any conversion started with the pre-M3 converter. The idx int8 default now lives at the consumer in convert_shard and both keys are recorded only when non-default: a default GLM/DeepSeek run builds a byte-identical params dict, while a REAL flag change (or --arch m3 on a GLM outdir) still aborts. "idx" classification itself only ever fires on M3's self_attn.index_* names (GLM's DSA indexer is "indexer", skipped or converted via --indexer), so no other arch's output bytes change. - layer_idx() matched "layers.<N>" anywhere; back to anchored "model.layers.N" plus the one vetted raw form, MiniMax-VL's "language_model.model.layers.N". (Review B3.)
…nnot infer
M3 rides the colibri binary, so a GLM-side refactor of the shared hot
paths (rmsnorm/act_glu/expert_ffn/MSA) could break it with every GLM
test green. Now something notices (Review B4):
- make m3-tiny-generate / m3-tiny-check: tiny random checkpoint ->
--arch m3 int8 conversion -> numpy oracle -> the engine's REF/TF
gate; tests/test_m3_tiny.py enforces token-exact counts (the engine
itself always exits 0). New ci.yml job runs it on Linux with pinned
deps (tools/requirements-m3-tiny.txt).
Everything the earlier revision of this commit added to coli and
openai_server -- banner roster, model_arch, engine dict, build target,
tune prompt, model ids, --arch choices, the KV-slot guard, the
cap sentinel -- is now supplied by the FamilyDescriptor and is gone from
this diff. What the descriptor could NOT express was the handful of
launcher gates written as `arch == "glm"`, which is a claim about the
engine binary, not about the family. They now ask the registry the
question they actually mean, via colibri_core():
family.engine_group == family_by_id("glm").engine_group
engine_group is the registry's own statement that two families share an
engine implementation -- and therefore its env contract, its CAP channel
and its byte-protocol REPL. Four sites move: env_for_engine (M3 was
about to get the sister-engine OMP treatment for a binary that sizes its
own team), operator_cap (CAP= was silently ignored), cmd_chat (M3 would
have spawned a gateway server instead of the byte-protocol REPL), and
engine_for's COLI_DOCKER_GLM_ONLY guard, which told users a GLM-only
image cannot run M3 -- it can; it is the same binary.
Two gates deliberately keep the narrow `arch == "glm"` test:
`coli tune`'s replay protocol (M3 tunes through the persistent-rotation
path; giving it GLM's fixed replay means touching six more branches in
autotune.py, unvalidated here) and need_worker_model (cluster expert
workers are not validated for M3, and refusing says so).
test_launcher_dispatch needs no M3 entry now: the registry asserts each
family resolves to its own declared engine_artifact, and the intentional
share is declared by the descriptor rather than by a set in the test.
Review follow-ups: - The Makefile suppression existed only because the VK GQA commit widened struct PC with act/alpha/limit push constants and left the existing positional initializers short. Name every field at all 16 sites (the gate_up dispatches overwrite act/alpha/limit right after, so semantics are unchanged -- the fields were implicitly zero before) and remove the flag from all five platform CFLAGS lines. gcc 15 builds warning-free. - kv_alloc's index-cache calloc: (unsigned) cast gives value-range propagation a provable bound so -Walloc-size-larger-than stays quiet on inlined paths (the alloc is dev's own line; our wider "|| c->msa" guard changed inlining enough to surface it on the reviewer's gcc). - docs/minimax-m3.md: family doc covering the shared-binary design, GQA/ MSA/swigluoai/router notes, the CTX=4096 default and the MSA 2048-token exactness bound, tool-calling dialect, and the m3-tiny CI gate.
…egment Review feedback from JustVugg#601. - attention_gqa's header claimed "MSA block selection is NOT implemented: full causal attention". It has been implemented since the Lightning Indexer landed: `sel` is built per (row, index head) and gates both the score loop and the value accumulation, so only the selected blocks are read. The comment now says what the body does, keeps the exactness note (selection is a no-op while the window fits in topk*block = 2048 tokens, because every causal block is selected then) and names the three ways out of the sparse path: dense layers, no indexer weights, COLI_MSA=0. - g_gemma_norm / g_act_swigluoai / g_swiglu_{alpha,limit} are process globals, which held while a process opened one model. JustVugg#1227 and JustVugg#1245 landed the Segment and Edge runtimes, so a GLM segment can now sit beside an M3 segment and the two would fight over these: last model_init wins. Latent -- no shipping binary registers those adapters -- but M3's segment adapter cannot land until they hang off the Model. Recorded at the declaration, where the next person to add an adapter will read it. - colibri_core() resolved family_by_id("glm") on every call to reach a string constant, and would have raised if that id were ever renamed. The group name moves into family_registry as COLIBRI_CORE_GROUP and both descriptors reference it, so the shared engine is single-source and the predicate no longer depends on another family's id.
|
All three done, rebased onto The
|
| Check | Result |
|---|---|
make colibri / make colibri VK=1 |
clean, 0 warnings each |
make segment-edge-library (new path via #1245) |
links; 0 new warnings |
make m3-tiny-check, CPU and VK binaries |
OK — prefill 24/24, decode 20/20 each |
TF_LOGITS vs oracle, IDOT=0 |
max abs diff 2.62e-06, 0 argmax flips |
| VK expert tier (int4 fixture) | VK-BLOCK: 20 blocks, avg nvk 2.00 / ncpu 0.00; decode output identical to the CPU arm at all 20 positions |
make test-c |
ALL PASS, incl. six real Segment and six real Edge adapters |
make test-python |
695 tests, OK (35 skipped) |
A footnote to last round's correction, since I walked into the near-miss again: TF_LOGITS dumps from the prefill loop only, so a CPU-vs-GPU compare on those rows reads 0.000e+00 no matter what the tier did. The decode comparison above is per-position engine output from the two arms, which is the arm the tier actually serves.
Rebase is clean — zero conflicts against the 8 new commits, and it picks up #1252, so the two red Windows checks should go green on this push.
There was a problem hiding this comment.
@steve-m, I am curious about the reason behind not making a dedicated backend for MiniMax M3.
5 participants
What this is
Early, testable support for MiniMax-M3 (428B total / 23B active MoE, the text backbone of the VL checkpoint) — GQA attention, the MSA "Lightning Indexer" block-sparse attention, the o200k-family tokenizer, converter support, and Vulkan offloads. Opening as WIP so people who want to play with M3 can start early; API/knobs may still shift.
What works
ARCH_M3, auto-detected frommodel_type): GQA 64Q/4KV attention with per-head Gemma QK-norm, partial split-half NEOX RoPE (rotary 64, θ=5e6), swigluoai activation, sigmoid+bias router with renormalized top-4 ×2.0, first-3-dense pattern, 64K context.convert --indexer.tok.hnow accepts classic space-separated BPE merges (o200k family) alongside the array form; MiniMax chat template (]~b]role … [e~[) in the serve loop andopenai_server.py.tools/convert_fp8_to_int4.py --arch m3, auto-detected): maps the VL checkpoint onto the container scheme (drops vision), int4-g64 or int3-g64 (--xbits 3) routed experts, indexer kept at int8 (--idx-bits), per-shard resume.COLI_VK_DENSE=1), fused shared expert, lm_head offload, opt-in GQA attention core (COLI_VK_ATTN=1, throughput-neutral on the reference box — see commit messages).matmul_i3was scalar on x86): 3.8 → 24 GB/s effective on a Zen2 12-core, which is what makes int3 containers viable for decode.Measured (reference box: Ryzen 9 3900X, 64 GB DDR4, RX 9070 16 GB, 2× NVMe striped mirror)
Prefill 516 tokens: 152.7 → 96.8 s over the same span.
--topphere is the expert-routing top-p (~228 → ~155 expert loads/token at 0.7).How to test
Without the weights (validates the engine math end-to-end, minutes):
With the real model (BF16 download is ~796 GB; the int4 container is 225 GB, int3 is 177 GB):
A second copy of the container on another drive via
COLI_MODEL_MIRROR=<dir>roughly doubles streaming bandwidth. CPU-only works (drop theCOLI_V*vars) but is ~5× slower.Note the model's custom license (minimax-community): convert your own weights; don't redistribute containers.
Known gaps / WIP
COLI_MSA=0kill switch for A/B against full attentiondevonce Vulkan backend: expert tier + dense + MLA attention on any Vulkan 1.2 GPU (successor to #84) #418 merges (drops the first ~38 commits)