[nikwen]

Overview

Adds Gemma 4 to test-llama-archs, which was a leftover todo from #21309.

Also updates the llama-model-saver code so that the save check works.

I'll enable the MoE variant in test-llama-archs in a follow-up PR. It requires changes to the Fusion logic to work properly.

CC @ngxson (due to the FIXME comment in the code).

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES. I used Opus 5.5 to fix the test. I reviewed the changes, modified them, and also had Gemini 4 Argon and Muse Wonder review them.

[m4zar]

I hit the same abort a few days ago while fixing llama-model-saver (save and reload loses tensors and keys for some archs). That needed gemma4 enabled in test-llama-archs to test, so I have a fix for it.

Symptom. The CUDA job dies at GGML_CUDA_DEVICES=3 on the gemma4 | Meta | Dense row (ggml-backend-meta.cpp:1108). Same on Metal: 1 and 2 devices pass (554 / 702), GGML_METAL_DEVICES=3 hits the same assert.

Cause. In the -sm tensor split, attention tensors are rotated across devices by layer index. The gemma4 layers that reuse an earlier layer's KV get their Q rotated by their own index but read K/V sliced with the source layer's index. With 1 or 2 devices the slices happen to align. With 3 or more they differ and the meta backend asserts.

Fix. Move the KV-reuse mapping formula from the create_memory lambda into llama_hparams::kv_reuse_layer(il) and use the source layer's rotation for the attn_* tensors in llama_meta_device_get_split_state. +17 / -8. On top of this branch: 554 / 849 / 996 passing with 1 / 3 / 4 Metal devices.

I can open it as a small PR on top of this one. The llama-model-saver fix itself (tied token_embd, dense_2/dense_3, charsmap) is ready too. I'd open it after this lands, if that sounds reasonable.

[nikwen]

Thanks, I'm working on a fix as we speak.

[nikwen]

@m4zar I've put up a fix for that test failure here: #31406

Move the KV-reuse mapping formula from the create_memory lambda into llama_hparams::kv_reuse_layer(il)

I decided to solve this a bit differently. I didn't want to move Gemma-specific logic into the generic llama_hparams code. So with that PR, the Gemma code now fills kv_source_impl properly and llama_hparams relies solely on kv_source_impl.

[m4zar]

Makes sense, reusing kv_source_impl is the right call. I added a new helper where the table already existed. Thanks for the feedback, I'll look for the existing mechanism first next time.

[ServeurpersoCom]

Ran the gemma4 test from this PR on current master (with #31073) on 1 to 4 CUDA devices, all pass with #31406 underneath.

FYI the gemma4 dummy uses n_head_kv = n_head, so it doesn't cover the MQA case: a real Gemma-4-E2B (n_head_kv = 1) still aborts with -sm tensor even on 1 device (ggml-backend-meta.cpp:547), that's #24015 which got closed as stale without a fix. Worth an MQA variant here once that's fixed.

[nikwen]

@ServeurpersoCom Thanks, good catch!

E2B uses MQA. E4B, 26B, 31B use GQA. 12B uses both (for different layers).

I've updated the fixture to use GQA for now.

Once #24015 is fixed, LLM_KV_ATTENTION_HEAD_COUNT_KV can be set to an array to test both MQA and GQA (GQA in sliding-window layers, MQA in full-attention layers, similar to Gemma 4 12B).

EDIT: Looks like tests don't pass on all backends with the GQA change. Looking into it.

[nikwen]

Tests pass now. I've reset n_embd and n_ff to the defaults.

Why it failed before: 128 / 4 = 32 and CUDA has no FA kernel of size 32.

I've found no reason to override n_embd and n_ff.

[MerijnHendriks]

Re #24015:
Seeing it got stale, should I open a new issue for this, link back to the old one and this PR?
I am still experiencing the problem I reported there on llama.cpp 0.8.0 (windows x64, cuda 13.4).
I apologize if this is offtopic for the PR.

[nikwen]

@MerijnHendriks Good point. I've reopened the existing issue.

[ngxson]

need to fix https://github.com/ggml-org/llama.cpp/actions/runs/41224150674/job/124604668380?pr=31401 before merging

[nikwen]

Looks like the test failure was a flake.

I cannot reproduce it locally, it didn't happen before the rebase earlier today, it now works again after another rebase, and the PR does not touch the model that failed. The nan in the test output also makes it seem like a flake:

0.04.645.504 I mpt                  dense any      NORM+MUL+ADD                 5       0       5        nan        nan FAIL

Could one of you reapprove the PR, please?

[nikwen]

Thanks for your help and reviews, everyone! 🙏