TP: fix split state and granularity for fused QKV gemma4, qwen35

[dfriehs]

Overview

#22780 adds the option to fuse QKV for models which previously had either no QKV tensors (gemma4) or for which QKV tensors only existed for specific layers (qwen35, qwen35moe, ...).

Gemma 4 31B: n_embd is 5376 but Q is 8192 (raising GGML_ASSERT(tensor->ne[axis] == n_embd + 2*n_embd_gqa);), so calculate from n_head * n_embd_head_k instead.

Qwen 27B/35B-A3B: added a copy of the default qkv handler for full attention layers, with n_embd doubled for Q gate tensors, and add the doubling in get_split_granularity as well.

Additional information

I tested Gemma 4 31B, Qwen 3.8 27B and Qwen 3.6 35B-A3B on 4 P100 cards.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES. I used Qwen 3.8 27B to analyze the errors I was getting with the fused models and then to reason through the changes required.

[ggerganov]

What is the effect before and after this PR?

[ggerganov]

We should deduplicate this condition if it has to match the one earlier:

Suggested change

	// fused full attention layers need Q gate tensors handled like above:
	if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE ||
	ud->model->arch == LLM_ARCH_QWEN4EXP) {
	// fused full attention layers need Q gate tensors handled like above:
	// TODO: deduplicate with the condition in `get_split_segments` [TAG_SPLIT_FUSED_QKV_QWEN]
	if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE ||
	ud->model->arch == LLM_ARCH_QWEN4EXP) {

[ggerganov]

Suggested change

	if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE ||
	ud->model->arch == LLM_ARCH_QWEN4EXP) {
	// TODO: clarify why this is necessary specifically for these models
	// TODO: deduplicate [TAG_SPLIT_FUSED_QKV_QWEN]
	if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE ||
	ud->model->arch == LLM_ARCH_QWEN4EXP) {

[dfriehs]

What is the effect before and after this PR?

On master both gemma4 and qwen35/qwen35moe crash for models exported with --fuse-qkv when -sm tensor is used.

I don't really have benchmarks, but I seem to get +2% tg on Qwen 3.6 35B-A3B with fused QKV on 4 cards (Q8_0: 70.86 -> 72.4), although both numbers use additional patches (force enabling cudagraphs on P100 which is only a net positive for MoE's on 4 cards or more).

TODO: clarify why this is necessary specifically for these models

My understanding is that attn_q contains both q and gate tensors, which is why it has double the elements and needs double the granularity.

[dfriehs]

@ggerganov I added a tag like your suggestions although I named it [TAG_SPLIT_QGATE_QWEN], hope this is alright. The doubled granularity is not specific for fused QKV and was already present for attn_q previously, I only extended it to fused QKV. I also added the tag to the attn_q handler that already existed previously.

The change to segmentation for qwen was necessary as the current code assumes the layout for recurrent layers for QKV tensors and no segmentation was necessary for other layers until now.