Add MiniMax-M3 (MSA: MiniMax Sparse Attention) support#24908
Add MiniMax-M3 (MSA: MiniMax Sparse Attention) support#24908timkhronos wants to merge 53 commits into
Conversation
Text-only port that re-uses existing components: MiniMax-M2 style GQA with per-head QK-norm and partial rotary, DeepSeek-V3 style leading-dense and routed/shared experts, and swigluoai activation. Sparse attention is not yet supported (dense fallback); vision tower and MTP heads are dropped.
…ch per group block picking
|
Hi @timkhronos, thanks for your contribution! Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:
Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below. |
|
Thanks. This is full support for a new model architecture (MiniMax-M3), so the core: arch definition, graph, KV-cache handling, indexer and conversion, has to land together to be functional; those aren't independently mergeable. The vision tower / mmproj is the one separable piece, but it's a small, self-contained addition (a standard CLIP-style ViT, the only model-specific parts are the patch embed, 3D-RoPE, and patch-merge projector), so I'd prefer to keep it in for complete model support in one pass rather than split it out. Of course, if a maintainer would still rather see vision as a follow-up, I'll split it. I am happy to open a tracking issue / RFC for the architecture now if you'd like the design discussed before review, though the PR description already covers the design, correctness validation, and limitations in detail. I went straight to the PR since this began as a vision-tower change and grew into full MSA support once I traced the long-context degradation past ~6k to the sparse-attention path, but I'm glad to write up an issue retroactively if that's the preferred process. |
4-way paths. Full debug harness remains at <8136a9c68ed7a5eb009aa67bba3fda8062f4648f> for reproducing the selection-parity validation.
@CISC Sorry, I don't have this model downloaded and examined in detail yet, just skimmed over the paper so I'm not qualified. |
Large part of the code is clearly AI-written. Please be honest about AI usage. Proof: you manually modified AI-generated comments to make it human-like: 2e82759 (To be clear: I'm ok with using AI and I myself use it from time to time, however, dishonesty is not something I will tolerate) |
This comment was marked as off-topic.
This comment was marked as off-topic.
Explained in #24908 (comment) and OP. |
|
|
@smalinin Could you please share the crash logs, your system specs (gpus, cpu, ram amount), when did the freeze occur on the last configuration(during loading, prefill/prompt processing or decoding/generation), and what launch command you used? I was unable to reproduce the behavior on my system ( other than --n-cpu-moe 0 where I also OOM-d). Also note that |
Log indexer cache size on launch Disallow ctx shift Support prompt caching
|
Since the PR says Vulkan is untested, I ran it on RADV. Short answer: it works. The sparse path is clearly being exercised, and I get identical output across two RDNA4 cards and an RDNA3 one. One caveat up front: this is a tiny random-weight model, so nothing here speaks to output quality or speed, only that the backend computes the right thing. Build is PR head Rather than pull the real weights I fabricated a small random-init M3 that keeps the parts MSA actually depends on: head_dim 128, 4 KV heads (= 4 index groups), index_dim 128, block 128, topk 16, 0.5 partial rotary, gemma-style norm, sigmoid+bias routing, and the VL-style On the ops: I ran For the model itself, two regimes. Short prompt (1339 tokens, below topk×block = 2048, so every block gets selected and MSA should reduce to plain dense): the MSA path, the dense bypass ( A couple of smaller things I noticed. With To be clear on scope again: random weights, so this only shows the Vulkan path computes correctly and deterministically, nothing about real quality or performance. If it helps I'm glad to run specific configs on the R9700s or the W7900; that would also answer whoever's asking in #24523 about 2× R9700 under ROCm. |
|
So I made a 2-bit GGUF quant with this PR if anyone else with similar hardware to me wants to give it a try: After messing around with the MSA version for a bit, it definitely works much better than the dense attention version from Unsloth's PR. The UD_IQ2_M model without MSA exhibited a strange sort of 'looping' for me. It wasn't just repeating the same text, rather, it would get stuck in the reasoning block repeating the same reasoning 'action'. (IE: Refining or drafting a part of the response over and over again but in slightly different ways). This dose not seem to happen with the MSA enabled version. I'm curious if higher quants also exhibit this behavior or if this is just showing up for me due to the smaller quant I was using. I did some very unscientific perplexity testing based off the files in the Serpen HF repo and comparing the UD-IQ2_M dense and MSA versions yielded a mild difference: I also did some testing with the indexer projections and what quants to run them at. The PR header recommends uping them to F32 from BF16. I compared that to just quanting them down to F16 and got identical KL numbers: Further personal testing and use of the model has also yielded no discernable change between the two. (Though this might be less true at higher context sizes! I was using 64K as it is all my hardware can run. And llama-perplexity used 512. So take this metric with quite a few grains of salt...) As for performance, MSA uses more RAM on my system than dense, as expected, and I get less TG speed. The dense model got around 13-14 t/s average at 64K context, while the MSA version gets 10 t/s average. Both have comparable prompt processing speed however. |
Fully rewrote minimax-m3.cpp for speed and buffer size gains: Unified the 4-way + decode, 1 FA call per layer instead of 4, with the groups mapped onto ne[3] Custom CPU op now emits block-level mask, expanded on GPU, which causes CPU to GPU transfer to shrinks at prefill Decode: ~25 nodes/layer vs ~50, no per-group concats/conts Unified selection semantics, so both regimes rank bs + local bias (position-anchored local force), which means prefill/decode can no longer disagree on selection can_reuse on the MSA bias input. Graph reuse at decode restored (was rebuilding the full graph every token) In-place mask adds, shrinking compute buffer ~6.8 to ~4.2 GiB at ub2048/62k Multi-stream: MSA now runs with -np N when kv_unified=false. Decode stays batched across streams (still 1 FA call), prefill loops per stream. dense fallback only for --kv-unified + multi-seq Measured effect on expert offload bound setup: decode 6.2(4WAY)–7.15(MSA_decode) -> 7.7~7.8 t/s, flat from 5k to 60k+. prefill around 10% faster. buffer about 20% smaller, multi-user support.
|
@neuromaniacMD Thanks for testing! Both issues should be fixed by the latest patch, although it is a fairly substantial rewrite of the logic itself. I removed the 4-way/decode split, instead now each sparse layer runs a single grouped FA call with the 4 groups mapped onto the mask's ne[3] broadcast dimension, decode gathers the selected blocks via get_rows instead of masking, graph reuse also now works at decode, and multi-stream (--parallel) should be supported. I also stripped all debug env vars, and the fallback warnings fire on my system. I also updated the description, since it was a bit outdated at this point indeed. Since the ops changed a bit, a re-run on RADV would be very valuable if you're willing, specifically whether Vulkan's FA handles distinct masks across ne[3], with ne[3] = 4 and ne[3] = 4 × N under -np N, and batched argsort/get_rows. The same determinism checks you ran, plus a -np 2 run with different-length prompts, would cover everything new. @Sciguy429 Thanks for the KLD numbers and especially the looping report! I should have stressed this harder in the description(fixed now), but MSA isn't a speed feature, it's a part of the model, it's what the model was trained with. The dense attention is an out of distribution approximation, and runs a computation the model was never trained for. It might look ok, but there's no guarantee it behaves like M3, and I personally noticed fairly significant degradation above 8-12k ctx, and also noticed the same looping you described on a Q4XL quant with the Dense PR. Dense mode in this PR exists purely as a fallback for configs that can't run MSA, with a quality warning. One important caveat on the KLD table though, llama-perplexity at 512 ctx can't measure any of this. Below topk*block = 2048 tokens there are fewer blocks than top-k, so MSA selects everything, and thus dense and MSA differ only by kernel fp noise there, and the IF16 vs IF32 comparison is guaranteed identical (selection can't change below 2048 ctx regardless of indexer precision). So those numbers unfortunately don't tell us much about indexer quant sensitivity, a meaningful test needs long-context eval where selection actually discards blocks. That's also why I kept the F32 recommendation for the indexer projections. F16 would most likely convert cleanly from bf16, but ~800 MiB is fairly cheap insurance for the tensors that decide which KV gets read until someone measures it at 32k+. Again it's probably fine, but the size difference is small enough where I'd personally just go F32. For me after the latest commit, decode speed is around 10% faster than the previous implementation, and only around 6% slower than the dense attention. Prompt processing is around 10% faster than before. Indexer cache size shrunk by around 50%(~800Mib at 64k) and compute buffer size shrunk by around 20%(2Gib at 64k). You will likely see noticeably better performance after repulling the latest changes. |
|
@timkhronos yeah that's what I was thinking too. Given the context size llama-perplexity was working with, the indexer tensors never really get 'used' at all. If we want some proper testing with that, someone else will need to run some perplexity tests at higher context levels. My hardware simply isn't capable of regenerating the Q8 logits file for it. So far from my own testing though, the F16 tensors do seem to work okay. I am more so running into general low-bit quant instability at the moment. Might be solvable with a better quant mix, but I will leave that for the typical players who are far better equipped than myself. As for performance, I pulled down the latest changes and did see a slight bump. I am now getting 12-13 t/s token generation from my typical 9-10 t/s and gained around +25 t/s in processing speed as well (though I would take this with a grain of salt because I have to run at a low -ub to conserve VRAM and am deliberately trading PP speed for it). I haven't seen a massive change in RAM usage, at least nothing I can directly notice from htop on my server. Perhaps 2-3gb at most. Comparing the perplexity runs from the old version and the new one shows a slight drop of around 1gb though, so it did do something. (Is this expected to pay off more for large context sizes?) Pre-commit: Post-commit: |
|
@Sciguy429 Yeah, I think comparing at around 20-30k would probably be more representative. I personally can only fit a Q4KM, so I can't run that either. Hopefully someone with enough memory can jump in to settle it either way, but for now I'd personally stick to F32, but I'd expect F16 to be mostly okay too most likely. I initially tried a Q4_K_M GGUF with the indexer projections quantized to Q4, and the output quality was noticeably worse. That is why I am fairly hesitant about reducing their precision in general without a proper long-context comparison. Your speed improvement numbers mostly check out, and almost +30% decode performance is actually quite a bit better than I expected. But your vram allocation is a bit off. Your compute buffer is identical pre/post (6828 in both), which means that's probably not dominated by the MSA tensors at all at your settings, since the changes cut those substantially, so whatever fills those ~6.8 GiB is probably something else(most likely something related to llama-perplexity?) Could you share your launch command? Is compute buffer size smaller when you launch the server? For reference, with the server, mine at 62k ctx / -ub 2048: (Compute buffer was ~6.4GiB pre-commit at 2048 ub/64k ctx) On another note, yes, the savings scale with context, the indexer cache saving is ~800 MiB at 64k, and the compute-buffer reduction scales with ctx*ub. |
|
@timkhronos yeah sure! (I left the command at the very top of my full log files in the HF repo.) It's pretty much a direct copy of the one from Serpen's repo, just altered to do a KL test rather than a logit generation. My server launch is this: (No batch specified, so it should default to '-b 2048 -ub 512'. This is slow for PP, but saves on VRAM.) NOTE: Just realized I never specified. 'MIX1' is the IF16 quant with f16 indexer tensors. |
|
So over the weekend I did some testing via runpod on the IF16/32 issue. My HF repo now has a flat Q8 and Q6_K quant in it if anyone else wants to avoid generating them. I also generated a 32K context Q8 logit dump via the following command (which is also in the HF repo): (Sadly, runpod ate the command output by randomly crashing the SSH session, so I don't have the final PPL for the Q8 file as I lost the entire command output...) Regardless, I went on to run a perplexity test on my UD-IQ2_M-IF32/16 quants and did find a slight difference between the two: The difference is within the error margin, but it is there. So while I can't say with certainty the exact importance of the F32 conversion, I think it is safe to say that it dose have an effect. @timkhronos can you run a similar comparison with your Q4KM quant when you get the chance? I am curious if this difference scales with BPW or not. I'm willing to bet it becomes more dramatic as the size decreases. |
|
Hi — first, thanks for this port, it's been running well for us on 2x RTX PRO What we're running
The existing fallback (confirmed working as intended)We found the explicit guard in const bool streams_ok = cparams.n_seq_max == 1 || !cparams.kv_unified;
const bool msa_enabled = fa_on && streams_ok;
...
LLAMA_LOG_WARN("%s: unified KV cache with n_seq_max > 1; MSA needs per-sequence streams "
"-> running DENSE attention. Output may be degraded. Drop --kv-unified to enable MSA.\n", __func__);This is a deliberate, correctly-documented fallback — not a silent bug. Confirmed empirically both ways:
Root cause (traced through the block-selection path)The MSA block-selection / maxpool-gather in the ggml graph indexes KV blocks With What a fix would needA sequence-relative page/slot indirection: something that maps logical Not filing this as a complaintThe current fallback is safe and correctly warns rather than silently |
Overview
Adds support for MiniMax-M3, a 60-layer / 128-expert MoE (3 dense + 57 MoE layers) using MiniMax Sparse Attention (MSA): a GQA block-sparse attention where a lightweight per-GQA-group indexer scores the visible causal context, max-pools the scores into 128-token key blocks, and selects the top-16 blocks (the query's local block is always force-included). The main attention then runs only over the selected blocks (16*128 = 2048 KV per query), independent of context length, so decode attention cost stops scaling with context.
Implements the architecture from the MSA paper (arXiv:2606.13392v2) and follows the transformers reference (modular_minimax_m3_vl) for the indexer construction, selection semantics, and RoPE. MTP heads are not present in the released weights. Vision support is a separate PR (#25113).
Note that MSA is not an optional speed optimization, as the model is trained with sparse attention, so block selection is part of the model's semantics. Running dense attention is an out-of-distribution approximation that degrades output quality. The dense path exists only as a fallback for configurations that can't run MSA.
Additional information
Implementation
Batch and decode implement the same block-scoring and selection semantics, block scores bs = maxpool_128(idx_q*idx_k^T + causal mask) plus a host-provided local-force bias (large positive bias on each query's local block), ranked together for top-k. Selection is per GQA group (4 indexer heads, one per group, shared across the group's 16 query heads).
Within each stream, the four GQA groups are mapped onto the FA mask's ne[3] broadcast dimension and executed in one grouped flash-attention call.
Batch (n_tokens_per_stream > 1, ie. prefill): a small CPU op turns block scores into a block-level 0/-inf mask, then the token-level expansion and combination with the causal mask happens on the GPU, and masked FA runs over the full cache views.
Decode (1 token per stream): batched top_k on the GPU, index arithmetic expands block ids to flat (token, kv-head) row ids, one ggml_get_rows each gathers the selected K/V/mask, and FA runs over just the 2048 gathered tokens.
Multi-stream (--parallel N) is supported whenever the KV cache is per-sequence (so kv_unified is false). In that configuration each stream's slots are then position-aligned, which MSA's slot-anchored block selection requires.
The indexer also needs its own per-layer key cache. It is allocated F32 on the same backend buffer as the main K/V cache. Context-shift is not supported, since shifting would desynchronize the index cache.
Requirements and fallbacks
Flash attention required. Both batch and decode call ggml_flash_attn_ext directly and assume the non-transposed V layout that llama.cpp only provides with FA on. FA off uses dense build_attn (no sparsity) with a one-time warning that output quality may be degraded.
Unified KV cache with multiple sequences (--kv-unified + -np > 1) would silently break block anchoring so it also falls back to dense with a one-time warning.
Limitations
Block selection is anchored to absolute KV slots. Erasing a prefix or middle of a sequence's history (llama_memory_seq_rm of a partial range) breaks selection silently. Removing whole sequences, trimming the suffix and reusing server slots is fine.
Main KV cache: F16/BF16 are tested. Quantized KV-cache types are not yet validated.
Quantization
Quants should keep the indexer projections (index_{q,k}_proj) at high precision(F32 ideally) since they drive block selection, so quant error there changes which blocks are read (a discrete retrieval loss), unlike a value projection. The cost is ~1 GB and these are some of the highest-value tensors in the model.
Conversion
GGUFs must be converted with this branch and must keep the indexer tensors, weights with the indexers stripped will not work.
Performance
My tested config is expert-offload bound. MSA decode is essentially flat with context (7.7 tok/s at 5k -> 7.67 at 62k), prefill is roughly ~420-480 tok/s at 62k with -b and -ub 2048. Compute buffer at -ub 2048 / 62k ctx is ~4.6 GiB.
How to test
As mentioned GGUFs converted with this PR are necessary, other's will fail to load. You can either convert one yourself, or you can use this Q4KM one (graciously provided by @AesSedai).
CUDA build assumed (
-DGGML_CUDA=ON), but in theory other backends with flash_attn_ext support for head-dim 128 should work (dense fallback otherwise), ROCm/HIP, Vulkan, Metal untested, reports welcome.Feed a ~4k and a ~30k+ prompt and generate. Decode throughput should stay roughly flat with context (fixed 2048-KV budget).
Multi-stream: llama-server --parallel 2 (or 4) with concurrent requests of different prompt lengths.
Credits
The initial MiniMax-M3 scaffolding was based on work by @danielhanchen. The sparse-attention indexer, the unified batch/decode MSA graphs, the KV-cache index handling, FA gating, and multi-stream support in this PR are new work.
Requirements