llama : initial Mamba-2 support

[compilade]

Follow-up from #8519 (comment). This should fix #7727 and fix #8519.

I've implemented the fully recurrent mode of Mamba-2, because it's very similar to Mamba-1, and also because it seems like the most appropriate mode for text generation.

This does not implement the sequentially semistructured matrix mode, because I'm not yet sure how the block decomposition would fit within the batch and ubatch framework of llama.cpp, and how the chunk size should be chosen. If the recurrent mode is faster at single-user auto-regressive text generation, then I'm not sure how to keep the graph node structure constant when using the most appropriate technique for the batch size.

If the sequentially semistructured matrix mode is eventually implemented, it should help with prompt processing speed for large prompts.

What to expect

(mostly taken from #8519 (comment))

The state in Mamba-2 is bigger than I thought; Mamba-Codestral-7B-v0.1 takes 263.5 MiB (in F32) per sequence (e.g. with -np 1), compared to 38 MiB (also in F32) for Falcon-Mamba-7B (which is based on Mamba-1). But that remains constant whatever the context size. Mamba-2 is easier to implement efficiently, so the bigger state does not really impede inference speed.

However, a big downside right now with recurrent models in llama.cpp is the lack of state rollback (which is implemented through state checkpoints in #7531, but needs to be re-adapted to #8526), so the prompt will be reprocessed a lot if using llama-server. I think using llama-cli in conversation mode does not have this problem, however (or maybe only the bare interactive mode with --in-prefix and --in-suffix, not sure).

This initial implementation is CPU-only, but uses SIMD for the SSM scan, so even though the state is bigger than for Mamba-1 models, in my tests, the speed of Mamba2-130M is similar or better than Mamba-130M (but still not that fast compared to transformer-based models with an empty context), when both are run on CPU.

The speed of Mamba-2 models seems comparable to Transformer-based models when the latter have 2k to 4k tokens in their context.

Summary of changes

• Add support for Mamba2ForCausalLM (including the official Mamba-2 models, and Mamba-Codestral-7B-v0.1)
  • Note that config.json needs to contain "architectures": ["Mamba2ForCausalLM"], for the convert script to properly detect the architecture.
• View Mamba-1 as having d_inner (aka 2 * n_embd) heads of size 1.
  • This simplifies the handling of shapes in ggml_ssm_scan
• ggml
  • Implement Mamba-2's selective state update in ggml_ssm_scan.
    • Re-using the same operator as Mamba-1, because it's pretty much the same operation. (except for how ssm_a is broadcast)
  • Fuse the operation with ssm_d into ggml_ssm_scan
    • Otherwise it would need to be transposed, because the dot-products are done head-wise.
  • Implement Mamba-2's SSM scan with GGML_SIMD.
    • This is possible because there is no element-wise expf in the state update unlike with Mamba-1.
  • Avoid state copies for the SSM state (both for Mamba-1 and Mamba-2) by passing state ids to ggml_ssm_scan.
    • Mamba-2 states are huge. Otherwise masking and copying took close to 10% of the CPU time according to perf.

Other

Here's my favorite quote from Section 3.3 of https://arxiv.org/abs/2405.21060:

Furthermore—by a twist of fate—structured state space models and sequentially semiseparable matrices have the same acronyms, underscoring their equivalence! Conveniently we can use any of these acronyms SSM (state space model or semiseparable matrix), SSS (structured state space or sequentially semiseparable), or SS (state space or semiseparable) interchangeably to unambiguously refer to either concept.

TODO

• Rebase onto master after merging llama : simplify Mamba with advanced batch splits #8526.
• Avoid unnecessary moves of the state
• Adapt the Metal kernels and the tests from ggml : add SSM Metal kernels #8546 to the updated ggml_ssm_scan
• Remove the new GGML_MUL fast broadcast path because it's not used anymore to mask the states.
• Maybe use a new metadata key instead of {arch}.ssm.time_step_rank for the number of heads of Mamba-2, because it's not really the rank of the time step (well, maybe kind of).
  • The meaning of the number of heads and the time-step rank is overlapping enough in Mamba-2 that I think this is fine.
• Maybe not fuse the multiplication with ssm_d in ggml_ssm_scan?
• Maybe split ggml_ssm_scan to separate the implementations for Mamba-1 and Mamba-2, although they do have a lot in common.
  • Seems like they can be distinguished easily enough at the time of kernel dispatch.

---

• I have read the contributing guidelines
• Self-reported review complexity:
  • Medium

[ngxson]

Hey @compilade , thanks for implementing this!

I tried converting https://huggingface.co/mistralai/Mamba-Codestral-7B-v0.1 using convert_hf_to_gguf.py, but it gives error:

    with open(dir_model / "config.json", "r", encoding="utf-8") as f:
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
FileNotFoundError: [Errno 2] No such file or directory: 'config.json'

Nevertheless, I successfully converted a Mamba-Codestral transformers-compatible model: https://huggingface.co/Molbap/code2 (Need to comment out the line raise NotImplementedError("BPE pre-tokenizer was not recognized - update get_vocab_base_pre()") in convert_hf_to_gguf.py)

Run it output model (remember to select the correct chat template, since the model does not come with one):

make llama-cli -j && ./llama-cli -m ../models/mcode-7.3B-Q8_0.gguf -cnv -p "You are a helpful assistant" --chat-template mistral -ngl 0

The result looks promising, but I have no idea why there are [UNK_BYTE_0x29681...]. It seems like the there is a problem with space character:

<<SYS>>Youareahelpfulassistant<</SYS>>
> hi
[UNK_BYTE_0xe29681▁Hello]Hello![UNK_BYTE_0xe29681▁How]How[UNK_BYTE_0xe29681▁can]can[UNK_BYTE_0xe29681▁I]I[UNK_BYTE_0xe29681▁assist]assist[UNK_BYTE_0xe29681▁you]you[UNK_BYTE_0xe29681▁today]today?

Link to download GGUF: https://huggingface.co/ngxson/codestral-mamba-llamacpp-test/tree/main

[compilade]

Hey @compilade , thanks for implementing this!

I tried converting https://huggingface.co/mistralai/Mamba-Codestral-7B-v0.1 using convert_hf_to_gguf.py, but it gives error:

    with open(dir_model / "config.json", "r", encoding="utf-8") as f:
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
FileNotFoundError: [Errno 2] No such file or directory: 'config.json'

@ngxson

The steps I took to convert Mamba-Codestral-7B-v0.1 are the following:

1. Rename consolidated.safetensors to model.safetensors
2. Rename params.json to config.json
3. Add the line "architectures": ["Mamba2ForCausalLM"], in config.json
4. Rename tokenizer.model.v3 to tokenizer.model
5. Use convert_hf_to_gguf.py as usual.

I did not have tokenization problems in my tests. Maybe because I was using the original SentencePiece tokenizer instead of a BPE tokenizer.

That tokenizer.json in the transformers-compatible version seems to have problematic spaces. It uses the SentencePiece space escaping instead of the BPE one. Its normalizer seems to revert the escaping, but that's not handled in llama.cpp.

There are probably still problems with the SentencePiece tokenizer too, like the lack of special tokens (control tokens seem to be identified correctly, the only difference seems to be with the 20 [REFERENCE_DOC_{n}] tokens (where n is 0 to 19), which tokenzier.json identifies as non-special added tokens (maps to USER_DEFINED for llama.cpp), while tokenizer.model identifies them as NORMAL tokens).

I think the SentencePiece tokenizer should be preferred for this model; it should be easier to handle without workarounds. I should change that in convert_hf_to_gguf.py. Meanwhile either not include tokenizer.json or rename it to something else.

[ngxson]

Thanks for the guide! I've successfully converted the original repository the gguf by following your steps.

For the transformers-compatible, I will try to contact the one who made it. Hopefully it will be fixed soon.

I'm wondering if convert_hf_to_gguf.py can automatically handle the renaming of params.json, consolidated.safetensors and tokenizer.model.v3? For now, my fear is that someone who use automated tools like gguf-my-repo will be stuck due to this issue.

(Also cc @Vaibhavs10 since he's the maintainer of gguf-my-repo.)

[Vaibhavs10]

Hey @compilade/ @ngxson - JFYI - the transformers weights are now merged in the main repo: https://huggingface.co/mistralai/Mamba-Codestral-7B-v0.1

If you face any issues with the conversion with this could you open an issue on the repo for us to track! 🤗

[1ns0mni4c]

Any updates on when Codestral Mamba should be supported?

[learning-chip]

Nice work! Just a note on the ssm_scan kernel performance: a better fused implementation by the flash-linear-attention project can give the equivalent functionality as Mamba2's original kernel: fla-org/flash-linear-attention#49 , and runs 2x faster: fla-org/flash-linear-attention#50

[molbap]

Hi @compilade ! I worked on repo conversion for the transformers-compatible mamba2 version, let us know if you need anything from us to move forward with this PR :)

[HanClinto]

I'm wondering if convert_hf_to_gguf.py can automatically handle the renaming of params.json, consolidated.safetensors and tokenizer.model.v3? For now, my fear is that someone who use automated tools like gguf-my-repo will be stuck due to this issue.

(Also cc @Vaibhavs10 since he's the maintainer of gguf-my-repo.)

It sounds like having a simple fallback of expected filenames would be a reasonable thing to include here? I don't know that we want to maintain a ton of different ones, but adding a second layer of fallbacks for alternate filenames doesn't feel arduous.

[compilade]

It sounds like having a simple fallback of expected filenames would be a reasonable thing to include here? I don't know that we want to maintain a ton of different ones, but adding a second layer of fallbacks for alternate filenames doesn't feel arduous.

@HanClinto

That's not really a problem anymore (at least for Mamba-Codestral) since the official repo was updated in https://huggingface.co/mistralai/Mamba-Codestral-7B-v0.1/commit/88085f9cdfa832c3aca8a0315a4520cf7558c947 to use more standard names.

What is currently blocking this is that the Metal and CUDA kernels for ggml_ssm_scan need to be updated BUT before that, I want to refactor the operator to completely avoid copying Mamba-2 states (because otherwise the unnecessary copies use a non-negligible fraction of the memory bandwidth (10% of total text generation inference time on my laptop), since Mamba-2 states are big).

[hg0428]

Any updates on this?

[gabe-l-hart]

I took a whack at syncing this with the latest merges to master (actually right before #13834 was merged), but something broke pretty badly. My branch is here in case it's helpful: https://github.com/gabe-l-hart/llama.cpp/tree/broken/mamba2-sync

The broken symptom is that when I run a mamba (not 2) model, it trips this assertion from this call. When I debug, it shows that dimensions for x appear to be in the wrong order, but when I shuffle them, they then conflict with the dimensions of dt.

I also rebased Granite 4 on top of this broken point and the model did generate tokens, so the mamba2 portion didn't raise dimension conflicts, but the tokens were garbage, so something was also quite broken there.

[compilade]

The broken symptom is that when I run a mamba (not 2) model, it trips this assertion from this call. When I debug, it shows that dimensions for x appear to be in the wrong order, but when I shuffle them, they then conflict with the dimensions of dt.

@gabe-l-hart
In your branch, in the Mamba graph in build_mamba_layer, the reshape for the ssm right after build_recurrent_state from ssm_states_all is different than the one I have here, while it should be the same as in the Mamba2 graph. And Mamba also uses avoid_copies for the ssm (since ggml_ssm_scan was made indirect with the extra ids parameter).

I did not have the same problem when merging master here, (although I still have to fix Mamba2's SSM scan with ARM SVE, because apparently that is not directly compatible with GGML_SIMD).

[gabe-l-hart]

Awesome, thanks for finding that. I must have fat fingerd a merge resolution. I'll give it another go and see if I can get the Granite 4 branch working again.

[gabe-l-hart]

I've now got a temporary branch to sync this with the Hybrid Recurrent Cache (#13979): https://github.com/gabe-l-hart/llama.cpp/tree/mamba2-sync

This seems to work correctly for mamba and mamba2 models as before, so my fat-fingering is fixed.

(cc @younesbelkada since I think you were trying to get to the same merge point of hybrid + mamba2)

[gabe-l-hart]

@compilade I did also rebase the Granite 4 branch once I verified that my sync-point branch is working for non-hybrid models. Something isn't working correctly and is showing the same symptom I saw earlier where the model runs, but produces garbage. The last point in history where this wasn't happening was e94f39, so I'm going to look into the changes to mamba2 since then (I think there have been a bunch) and see if I can identify something that might lead to this.

My broken Granite 4 branch is here in case you get a moment to glance at it: https://github.com/gabe-l-hart/llama.cpp/tree/broken/GraniteFour

[gabe-l-hart]

@compilade I did also rebase the Granite 4 branch once I verified that my sync-point branch is working for non-hybrid models. Something isn't working correctly and is showing the same symptom I saw earlier where the model runs, but produces garbage. The last point in history where this wasn't happening was e94f39, so I'm going to look into the changes to mamba2 since then (I think there have been a bunch) and see if I can identify something that might lead to this.

My broken Granite 4 branch is here in case you get a moment to glance at it: https://github.com/gabe-l-hart/llama.cpp/tree/broken/GraniteFour

This turned out to be a bug in the hybrid recurrent cache, so it's resolved now.

@compilade with your latest changes to avoid copies, I now see state-spaces/mamba2-370m and ibm-granite/granite-4.0-tiny-preview working on both CPU and metal, but neither work on CUDA. I think I'm reading ssm-scan.cu correctly that it hasn't been updated for mamba2 yet, right? I think you mentioned that you were working on this in a previous comment, but I wanted to double check in case it was something I could look to help with (by finding someone at IBM who knows kernel development well).

Also, I'm seeing some interesting performance characteristics between metal and CPU on my M3 Max 64GB. If I run a simple inference with llama-cli, the generation timings output for CPU for Granite 4 are consistently around 35 t/s and on metal it's around 43 t/s. If I run with llama-batched-bench, the CPU numbers look very different but the metal numbers look almost exactly the same. I see the same pattern for state-spaces/mamba2-370m as well.

# simple run metal
./bin/llama-cli -m ~/models/mamba2-370m-hf/ggml-model-Q4_K_M.gguf --temp 0 -p "Tell me a story about a developer and their dog" -no-cnv -n 100
# llama_perf_sampler_print:    sampling time =       2.41 ms /   110 runs   (    0.02 ms per token, 45681.06 tokens per second)
# llama_perf_context_print:        load time =     141.28 ms
# llama_perf_context_print: prompt eval time =      61.99 ms /    10 tokens (    6.20 ms per token,   161.31 tokens per second)
# llama_perf_context_print:        eval time =    1748.05 ms /    99 runs   (   17.66 ms per token,    56.63 tokens per second)
# llama_perf_context_print:       total time =    1815.04 ms /   109 tokens

# simple run CPU
./bin/llama-cli -m ~/models/mamba2-370m-hf/ggml-model-Q4_K_M.gguf --temp 0 -p "Tell me a story about a developer and their dog" -no-cnv -n 100 -ngl 0
# llama_perf_sampler_print:    sampling time =       2.84 ms /   110 runs   (    0.03 ms per token, 38773.35 tokens per second)
# llama_perf_context_print:        load time =      99.32 ms
# llama_perf_context_print: prompt eval time =      33.33 ms /    10 tokens (    3.33 ms per token,   300.01 tokens per second)
# llama_perf_context_print:        eval time =    1009.74 ms /    99 runs   (   10.20 ms per token,    98.04 tokens per second)
# llama_perf_context_print:       total time =    1049.48 ms /   109 tokens

# batched-bench run metal
./bin/llama-batched-bench -m ~/models/mamba2-370m-hf/ggml-model-Q4_K_M.gguf -c 2048 -b 2048 -ub 512 -npp 10 -ntg 128 -npl 1
# llama_perf_context_print:        load time =     167.10 ms
# llama_perf_context_print: prompt eval time =     125.31 ms /    26 tokens (    4.82 ms per token,   207.48 tokens per second)
# llama_perf_context_print:        eval time =    2233.30 ms /   128 runs   (   17.45 ms per token,    57.31 tokens per second)
# llama_perf_context_print:       total time =    2444.29 ms /   154 tokens

# batched-bench run CPU
./bin/llama-batched-bench -m ~/models/mamba2-370m-hf/ggml-model-Q4_K_M.gguf -c 2048 -b 2048 -ub 512 -npp 10 -ntg 128 -npl 1 -ngl 0
# llama_perf_context_print:        load time =     108.84 ms
# llama_perf_context_print: prompt eval time =      55.39 ms /    26 tokens (    2.13 ms per token,   469.37 tokens per second)
# llama_perf_context_print:        eval time =    1006.83 ms /   128 runs   (    7.87 ms per token,   127.13 tokens per second)
# llama_perf_context_print:       total time =    1139.26 ms /   154 tokens

My only thought is that somehow the CPU version requires some kind of "warm up" beyond the basic warm up call where it's slower and that llama-batched-bench is accounting for this somehow?

[compilade]

I think I'm reading ssm-scan.cu correctly that it hasn't been updated for mamba2 yet, right? I think you mentioned that you were working on this in a previous comment, but I wanted to double check in case it was something I could look to help with

@gabe-l-hart That's correct, it's not yet updated. I did not really work on it enough since I wrote that.

But I'll be working on making/finishing a working version of the CUDA kernel for Mamba2 in the next days. It probably won't be the fastest it can be, so performance optimizations will be very welcome. If I don't give some progress update in the next 24 hours, please ping me.

---

My only thought is that somehow the CPU version requires some kind of "warm up" beyond the basic warm up call where it's slower and that llama-batched-bench is accounting for this somehow?

Yes, it could be related to the CPU clock boosting after an increase in workload, and the run in llama-cli was too short to trigger that boost (maybe). The frequency scaling for the GPU is likely different.

[gabe-l-hart]

But I'll be working on making/finishing a working version of the CUDA kernel for Mamba2 in the next days. It probably won't be the fastest it can be, so performance optimizations will be very welcome. If I don't give some progress update in the next 24 hours, please ping me.

Awesome, I super appreciate it. I'll ask around internally and look for folks that might be interested in helping once the initial update is done.

Yes, it could be related to the CPU clock boosting after an increase in workload, and the run in llama-cli was too short to trigger that boost (maybe). The frequency scaling for the GPU is likely different.

Fascinating, that would make a lot of sense. IIRC, I also saw something like this with llama-bench where when running in CPU-only mode, it would have pretty garbage performance for the first test in the matrix, then subsequent ones would be much more performant.

[gabe-l-hart]

@compilade @younesbelkada I've now updated mamba2-sync for the latest renames on #13979: https://github.com/gabe-l-hart/llama.cpp/tree/mamba2-sync

[compilade]

@gabe-l-hart I've implemented a working CUDA kernel for the Mamba-2 ssm scan in f8c7cae, and I've also fixed the Mamba-1 CUDA kernel (which I had previously broken by changing the shapes and reducing the copies with the extra ids arg).

I think this is mostly complete. (unless there's a remaining problem)

I will probably extract 0b6f6be in its own PR, though, since it's not specific to Mamba-2.

[gabe-l-hart]

Amazing, thank you @compilade! My CUDA box managed to become unreachable yesterday and the sysadmin is out for the week, so I'll test on my end as soon as I can get access to a machine.

[younesbelkada]

Hi @compilade @gabe-l-hart - thank you again for this work, I wanted to follow up about the status of this PR and see what needs to be done before merging it and if I can do anything to help merging it 🙏 thank you again !

[compilade]

what needs to be done before merging

• Solve the remaining conflicts
• split 0b6f6be in its own PR
• test conversion of Mamba-2 models
  • Mamba codestral https://huggingface.co/mistralai/Mamba-Codestral-7B-v0.1
  • Mamba-2 (from https://huggingface.co/state-spaces/models)
  • done by @younesbelkada in llama : initial Mamba-2 support #9126 (comment)
• test inference of Mamba-2 (and Mamba-1) models
  • I'm relatively confident that it works, because of test-backend-ops in the CI, and because I did test each of these at some point in the past.
  • On CPU
    • on ARM with SVE
    • on anything else with SIMD
    • without SIMD
  • With CUDA
  • With Metal
• Decide if it's okay to make breaking changes to SSM_SCAN to reduce copies
  • Should it be added to some changelog somewhere?

I think that is all? (I'll add more if missing)

[younesbelkada]

thank you @compilade !
On my side happy to take:

test conversion of Mamba-2 models

• Mamba codestral https://huggingface.co/mistralai/Mamba-Codestral-7B-v0.1
• Mamba-2 (from https://huggingface.co/state-spaces/models)

[younesbelkada]

I was able to successfully convert codestral 7B and mamba2-2.7b locally and run inference while getting coherent
generation:

Commands used for conversion (after git cloning locally the repos):

mamba2-2.7b

python convert_hf_to_gguf.py mamba2-2.7b --outfile models/mamba2-2b-bf16.gguf

codestral mamba-7b

python convert_hf_to_gguf.py Mamba-Codestral-7B-v0.1 --outfile models/Codestral-bf16.gguf

Output results:

Command used is below:

./build/bin/llama-cli -m models/mamba2-2b-bf16.gguf --prompt "def hello_world():" --temp 0.7

Mamba2-2.7B

def hello_world():
    print("Hello, world!")

hello_world()
 [end of text]

Codestral-Mamba7b:

 def hello_world():
    print("Hello World!")

if __name__ == "__main__":
    hello_world()

--> codestral 7b tends to repeat a lot of tokens when it's done generating, I found some community issues complaining about the same while not using llama.cpp: https://huggingface.co/mistralai/Mamba-Codestral-7B-v0.1/discussions/23 so it might be some tokenizer mis-configuration issue which is not related to the pure mamba2 modeling code.

Since the official mamba2 seems to give some positive signal with much less repetition, we should probably be good on conversion side. I can do more extensive tests if you suspect some commits might created some divergence/

Note I am testing everything on Apple Metal, so I can also take:

test inference of Mamba-2 (and Mamba-1) models

• Metal

How were you intended to test things? I can take that one and do some tests on Apple Metal

[ggerganov]

Decide if it's okay to make breaking changes to SSM_SCAN to reduce copies

It's OK to make breaking changes to SSM_SCAN. I'm pretty sure these have very small adoption (if any) at this point.

[gabe-l-hart]

Thanks for all the testing @younesbelkada! I'm still working on reviving my CUDA box and can test there once it's live.

@compilade when you test new model support, do you have a system for comparing activations against the corresponding transformers model? I've done this with granite models using llama-eval-callback, but it tends to be pretty manual, so just curious if you have any recommendations.