[nikwen]

Overview

Fixes the rotation of attention tensors for layers that reuse another layer's KV cache.

For these layers, the attention tensors need to use the source layer's rotation so that Q and KV heads are on the same device.

Required for the tests to pass on #31401. Tests failed without.

A nice side effect is that we move Gemma-specific logic out of llama-model.cpp into Gemma-specific files.

Additional information

Existing discussion of this PR here: #31401 (comment) + #31401 (comment)

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES. Opus 5.5 helped investigate and fix the bug. I heavily modified its output. (Claude's initial approach was bad for long-term maintenance.)

[ggerganov]

Nice - would be great to finally create gemma4 dummy models. I haven't tracked what is the root-cause for disabling them in the first-place. Could you rebase to latest master as we just merged some changes to the rotation matrices?

[ServeurpersoCom]

Reproduced on current master (with #31073 merged), using the gemma4 test enabled in #31401 on top:

Without this PR test-llama-archs aborts on gemma4 with 3 and 4 CUDA devices, and a real Gemma-4-E4B also crashes with -sm tensor on 3 devices.

With this PR both pass on 1 to 4 devices, and the E4B greedy output on 3 devices matches the single device one. Nice cleanup moving the Gemma logic out of llama-model.cpp too.

@nikwen one small nit, nothing blocking: the source rotation could be applied to the attn_* tensors only, so the FFN tensors of the shared layers keep their own rotation and the rounding stays spread across devices.

[nikwen]

@ServeurpersoCom @ggerganov Thanks for taking a look! I've rebased the PR and applied the source rotation to the attn_* tensors only. Thanks for noticing!

[nikwen]

Thanks for the reviews, @ServeurpersoCom, @ngxson, and @CISC!

[diabluz]

Hi! Thanks for fixing the shared-KV rotation — just wanted to flag a noticeable VRAM side-effect observed on a hybrid (recurrent + full-attention) arch with split-mode=tensor, in case it's useful.

Setup: 3x RTX 2000 Ada (16 GB), Qwen3.6-35B-A3B Q8_0 (qwen35moe, 40 layers: 30 recurrent GDN + 10 full-attention), ctx 262k, -np 4, --kv-unified, draft-mtp, ubatch 1024.

Symptom: after updating past this commit, the server starts fine but dies on the first long prompt with "CUDA error: out of memory" in cuMemCreate during a mul_mat_q. Same config, same model, same flags — it used to run with ~800 MiB of headroom per GPU.

What actually changed is visible already at startup: the Meta compute buffer reservation grows from 656 MiB (parent commit) to 1136 MiB (this commit, 0138937). Worth noting the graph is identical — same node count (4209), same input tensors. So the +480 MiB isn't coming from graph shape, it's purely in how split configs are derived on the Meta device. And the number is suspiciously clean: 480 MiB = 30 layers x 16 MiB, i.e. exactly ssm_d_inner (4096) x ubatch (1024) x f32 for each recurrent layer. Confirmed via git bisect run, and re-checked by hand (parent good, this commit bad).

It was also checked that this isn't caused by the nearby KV-rows work: with LLAMA_KV_ROWS_DISABLE=1 the regression is still fully there (#29510 only accounts for +4 MiB on this model).

Seems that at the mechanism level (read from the diff, happy to be wrong): for layers where has_kv(il) is false — which is 3 of every 4 layers on qwen35moe — the new reuse callback now resolves to kv_source(il) where it previously returned -1, and the attn_* rotation also goes through kv_source(il) now. On a meta/split device that apparently changes how the GDN intermediates get split (or forces them to materialize whole instead of per-device), and 30 layers pile that up fast.

My workaround for now: staying on a pre-#31406 build, so no urgency from this side — but for anyone running hybrid-SSM models tensor-split across GPUs with a tight VRAM budget, this commit is a pretty sharp cliff. Happy to re-test any candidate fix against the same harness if it helps (it's just a scripted build + one sched_reserve line to parse).

[ServeurpersoCom]

Symptom: after updating past this commit, the server starts fine but dies on the first long prompt with "CUDA error: out of memory" in cuMemCreate during a mul_mat_q. Same config, same model, same flags — it used to run with ~800 MiB of headroom per GPU.

Thanks for the detailed report! I couldn't reproduce it: with the same flags (-sm tensor on 3 devices, -c 262144 -np 4 --kv-unified -ub 1024, draft-mtp) and Qwen3.6-35B-A3B UD-Q8_K_XL on a rootless container with an RTX PRO 6000 emulating 3 devices via GGML_CUDA_DEVICES=3, the Meta compute buffer is 656.04 MiB on both the parent commit and 0138937, and 660.04 MiB on current master.

This matches the diff: for qwen35moe kv_source(il) is the identity, so has_kv() returns true on every layer as before, the reuse callback is only set for Gemma 3n/4, and the attn_* rotation resolves to the layer's own index, so the weight splits are unchanged.

Could you share the exact hashes you tested, any local patches or -ts you use, and the -lv 4 startup logs of both builds? Happy to dig further with that.

[diabluz]

Hello ServeurpersoCom, oh yes you're right, and sorry, my report was based on a bad baseline. Turns out the 656.04 MiB number I attributed to the parent commit that came from a local build of mine that carries a custom patch (it fuses the dequant of quantized KV into flash-attention), not from actual parent. Went back and re-checked ddf48f8 (the real parent) reserves 1136.04, same as the current master. And checking again the local patch on the exact same binary flips it 1136.04 off / 656.04 on. So the whole 480 MiB difference was my own patch all along — the kv_source/rotation reasoning in your reply is spot on for qwen35moe, and this PR really does nothing to this. Thanks a lot for your quick reply and analysis.

[nikwen]

@diabluz Thanks for getting back to us!