Commit 2280f0fd1 for llama.cpp
commit 2280f0fd163b03d4331d49e08460f66502e35854
Author: Aman Gupta <amangupta052@gmail.com>
Date: Sun Oct 25 21:23:22 2026 +0800
qwen: merge tensors on convert (#31028)
* qwen: merge tensors on convert
* convert: use generator functions
* qwen4exp: remove conts and pre-scale
diff --git a/conversion/qwen4exp.py b/conversion/qwen4exp.py
index 7ca26c2a3..8a55d678d 100644
--- a/conversion/qwen4exp.py
+++ b/conversion/qwen4exp.py
@@ -41,6 +41,8 @@ class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
self._ple_shards: dict[int, str] = {}
self._ple_row_dim: int | None = None
self._mtp_fc: dict[str, Tensor] = {}
+ # halves of the merged weights, keyed by (bid, merged tensor), see _merge_shared_input
+ self._merge_parts: dict[tuple[int, gguf.MODEL_TENSOR], dict[int, Tensor]] = {}
@classmethod
def filter_tensors(cls, item):
@@ -131,44 +133,79 @@ class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
# int64 hash constants must stay exact; 1-D tensors force F32, so use KV
if name.endswith("ple_embedding.layer_multipliers"):
self._ple_multipliers = [int(x) for x in data_torch.tolist()]
- return []
+ return
if name.endswith("ple_embedding.ngram_heads_offsets"):
self._ple_head_offsets = [int(x) for x in data_torch.tolist()]
- return []
+ return
if name.endswith("ple_embedding.ngram_heads_vocab_sizes"):
self._ple_head_vocab_sizes = [int(x) for x in data_torch.tolist()]
- return []
+ return
if ".ngram_embedding.shard_" in name:
- return self._place_ple_shard(data_torch, name)
+ yield from self._place_ple_shard(data_torch, name)
+ return
# eh_proj([e ; h_s]) = fc_embedding(e) + fc_hidden(h_s) for every hc stream s
if name.endswith((".nextn_fc_embedding.weight", ".nextn_fc_hidden.weight")):
self._mtp_fc[name.rsplit(".", 2)[1]] = data_torch
if len(self._mtp_fc) < 2:
- return []
+ return
eh = torch.cat([self._mtp_fc.pop("nextn_fc_embedding"), self._mtp_fc.pop("nextn_fc_hidden")], dim=1)
- return [(self.format_tensor_name(gguf.MODEL_TENSOR.NEXTN_EH_PROJ, bid, ".weight"), eh)]
+ yield (self.format_tensor_name(gguf.MODEL_TENSOR.NEXTN_EH_PROJ, bid, ".weight"), eh)
+ return
# one projection feeds indexer q and k; split it, as minimax-m3 does
if ".indexer.index_qk_proj.weight" in name:
n_q = self.hparams["indexer_n_heads"] * self.hparams["indexer_head_dim"]
- q = data_torch[:n_q]
- k = data_torch[n_q:]
- return [
- (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_Q_PROJ, bid, ".weight"), q),
- (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_K_PROJ, bid, ".weight"), k),
- ]
+ yield (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_Q_PROJ, bid, ".weight"), data_torch[:n_q])
+ yield (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_K_PROJ, bid, ".weight"), data_torch[n_q:])
+ return
# Gemma zero-centred gammas the inherited norm.weight rule misses
if name.endswith((".ple.norm_key.weight", ".ple.norm_query.weight", ".ple.norm_conv.weight",
".indexer.q_layernorm.weight", ".indexer.k_layernorm.weight")):
- return [(self.map_tensor_name(name), data_torch + 1)]
+ yield (self.map_tensor_name(name), data_torch + 1)
+ return
if name.endswith(".ple.conv1d.weight"):
- return [(self.map_tensor_name(name), data_torch.squeeze())]
+ yield (self.map_tensor_name(name), data_torch.squeeze())
+ return
+
+ for new_name, data in super().modify_tensors(data_torch, name, bid):
+ yield from self._merge_shared_input(new_name, data, bid)
+
+ # matmuls that read the same input, written as one weight [a | b] so the runtime does one matmul:
+ # (merged tensor, (a, b), linear-attention layers only)
+ _SHARED_INPUT_MERGES = (
+ (gguf.MODEL_TENSOR.HC_ATTN_DOWN_INJECT, (gguf.MODEL_TENSOR.HC_ATTN_DOWN, gguf.MODEL_TENSOR.HC_ATTN_INJECT), False),
+ (gguf.MODEL_TENSOR.HC_FFN_DOWN_INJECT, (gguf.MODEL_TENSOR.HC_FFN_DOWN, gguf.MODEL_TENSOR.HC_FFN_INJECT), False),
+ (gguf.MODEL_TENSOR.ATTN_QKV_GATE, (gguf.MODEL_TENSOR.ATTN_QKV, gguf.MODEL_TENSOR.ATTN_GATE), True),
+ (gguf.MODEL_TENSOR.SSM_BETA_ALPHA_CAT, (gguf.MODEL_TENSOR.SSM_BETA, gguf.MODEL_TENSOR.SSM_ALPHA), True),
+ )
+
+ # caches one half of a merged weight until the other half arrives, then yields [a | b]
+ def _merge_shared_input(self, new_name: str, data: Tensor, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+ # a pre-quantized checkpoint carries per-tensor scales, which one merged weight cannot hold
+ if bid is None or self.hparams.get("quantization_config"):
+ yield (new_name, data)
+ return
+
+ layer_types = self.hparams.get("layer_types") or []
+ is_linear = bid < len(layer_types) and layer_types[bid] == "linear_attention"
+
+ match = next(((merged, parts.index(part)) for merged, parts, linear_only in self._SHARED_INPUT_MERGES
+ for part in parts
+ if (is_linear or not linear_only) and self.match_model_tensor_name(new_name, part, bid)), None)
+ if match is None:
+ yield (new_name, data)
+ return
- return super().modify_tensors(data_torch, name, bid)
+ merged, idx = match
+ halves = self._merge_parts.setdefault((bid, merged), {})
+ halves[idx] = data
+ if len(halves) == 2:
+ del self._merge_parts[(bid, merged)]
+ yield (self.format_tensor_name(merged, bid), torch.cat([halves[0], halves[1]], dim=0))
# the shards concatenate into a tensor of well over 100 GB
# use LazyChunkedTensor here, a single shard resident at a time
@@ -214,6 +251,8 @@ class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
super().prepare_tensors()
if self._mtp_fc:
raise ValueError(f"MTP projection missing its other half: {sorted(self._mtp_fc)}")
+ if self._merge_parts:
+ raise ValueError(f"merged weights missing a half: {sorted((bid, t.name) for bid, t in self._merge_parts)}")
n_parts = self.hparams.get("split_ngram_parts", 0)
if self._ple_shards and len(self._ple_shards) != n_parts:
raise ValueError(
diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py
index 706cad3ea..babc8bba4 100644
--- a/gguf-py/gguf/constants.py
+++ b/gguf-py/gguf/constants.py
@@ -920,6 +920,10 @@ class MODEL_TENSOR(IntEnum):
HC_FFN_DOWN = auto() # qwen4exp
HC_FFN_UP = auto() # qwen4exp
HC_FFN_INJECT = auto() # qwen4exp
+ HC_ATTN_DOWN_INJECT = auto() # qwen4exp, [hc_attn_down | hc_attn_inject]
+ HC_FFN_DOWN_INJECT = auto() # qwen4exp, [hc_ffn_down | hc_ffn_inject]
+ ATTN_QKV_GATE = auto() # qwen4exp, [attn_qkv | attn_gate]
+ SSM_BETA_ALPHA_CAT = auto() # qwen4exp, [ssm_beta | ssm_alpha], unlike the grouped qwen3next ssm_ba
PLE_KEY = auto() # qwen4exp
PLE_VALUE = auto() # qwen4exp
PLE_NORM_KEY = auto() # qwen4exp
@@ -1746,6 +1750,10 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
MODEL_TENSOR.HC_FFN_DOWN: "blk.{bid}.hc_ffn_down", # qwen4exp
MODEL_TENSOR.HC_FFN_UP: "blk.{bid}.hc_ffn_up", # qwen4exp
MODEL_TENSOR.HC_FFN_INJECT: "blk.{bid}.hc_ffn_inject", # qwen4exp
+ MODEL_TENSOR.HC_ATTN_DOWN_INJECT: "blk.{bid}.hc_attn_down_inject", # qwen4exp
+ MODEL_TENSOR.HC_FFN_DOWN_INJECT: "blk.{bid}.hc_ffn_down_inject", # qwen4exp
+ MODEL_TENSOR.ATTN_QKV_GATE: "blk.{bid}.attn_qkv_gate", # qwen4exp
+ MODEL_TENSOR.SSM_BETA_ALPHA_CAT: "blk.{bid}.ssm_beta_alpha", # qwen4exp
MODEL_TENSOR.PLE_KEY: "blk.{bid}.ple_key", # qwen4exp
MODEL_TENSOR.PLE_VALUE: "blk.{bid}.ple_value", # qwen4exp
MODEL_TENSOR.PLE_NORM_KEY: "blk.{bid}.ple_norm_key", # qwen4exp
@@ -3309,6 +3317,10 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
MODEL_TENSOR.HC_FFN_DOWN,
MODEL_TENSOR.HC_FFN_UP,
MODEL_TENSOR.HC_FFN_INJECT,
+ MODEL_TENSOR.HC_ATTN_DOWN_INJECT,
+ MODEL_TENSOR.HC_FFN_DOWN_INJECT,
+ MODEL_TENSOR.ATTN_QKV_GATE,
+ MODEL_TENSOR.SSM_BETA_ALPHA_CAT,
# full attention layers: ATTN_Q holds [q|gate] interleaved per head
MODEL_TENSOR.ATTN_Q,
MODEL_TENSOR.ATTN_Q_NORM,
diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp
index 39e29d09a..ea13e9c3f 100644
--- a/src/llama-arch.cpp
+++ b/src/llama-arch.cpp
@@ -635,6 +635,10 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
{ LLM_TENSOR_HC_FFN_DOWN, "blk.%d.hc_ffn_down" },
{ LLM_TENSOR_HC_FFN_UP, "blk.%d.hc_ffn_up" },
{ LLM_TENSOR_HC_FFN_INJECT, "blk.%d.hc_ffn_inject" },
+ { LLM_TENSOR_HC_ATTN_DOWN_INJECT, "blk.%d.hc_attn_down_inject" },
+ { LLM_TENSOR_HC_FFN_DOWN_INJECT, "blk.%d.hc_ffn_down_inject" },
+ { LLM_TENSOR_ATTN_QKV_GATE, "blk.%d.attn_qkv_gate" },
+ { LLM_TENSOR_SSM_BETA_ALPHA_CAT, "blk.%d.ssm_beta_alpha" },
{ LLM_TENSOR_PLE_KEY, "blk.%d.ple_key" },
{ LLM_TENSOR_PLE_VALUE, "blk.%d.ple_value" },
{ LLM_TENSOR_PLE_NORM_KEY, "blk.%d.ple_norm_key" },
@@ -907,6 +911,10 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
{LLM_TENSOR_HC_FFN_DOWN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_HC_FFN_UP, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_HC_FFN_INJECT, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+ {LLM_TENSOR_HC_ATTN_DOWN_INJECT, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+ {LLM_TENSOR_HC_FFN_DOWN_INJECT, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+ {LLM_TENSOR_ATTN_QKV_GATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+ {LLM_TENSOR_SSM_BETA_ALPHA_CAT, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_PLE_KEY, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_PLE_VALUE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_PLE_NORM_KEY, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
diff --git a/src/llama-arch.h b/src/llama-arch.h
index 3d0032c0c..7996bdc30 100644
--- a/src/llama-arch.h
+++ b/src/llama-arch.h
@@ -679,6 +679,10 @@ enum llm_tensor {
LLM_TENSOR_HC_FFN_DOWN, // qwen4exp
LLM_TENSOR_HC_FFN_UP, // qwen4exp
LLM_TENSOR_HC_FFN_INJECT, // qwen4exp
+ LLM_TENSOR_HC_ATTN_DOWN_INJECT, // qwen4exp
+ LLM_TENSOR_HC_FFN_DOWN_INJECT, // qwen4exp
+ LLM_TENSOR_ATTN_QKV_GATE, // qwen4exp
+ LLM_TENSOR_SSM_BETA_ALPHA_CAT, // qwen4exp
LLM_TENSOR_PLE_KEY, // qwen4exp
LLM_TENSOR_PLE_VALUE, // qwen4exp
LLM_TENSOR_PLE_NORM_KEY, // qwen4exp
diff --git a/src/llama-model.h b/src/llama-model.h
index 93164f433..01fbda319 100644
--- a/src/llama-model.h
+++ b/src/llama-model.h
@@ -623,6 +623,10 @@ struct llama_layer {
struct ggml_tensor * hc_ffn_down = nullptr;
struct ggml_tensor * hc_ffn_up = nullptr;
struct ggml_tensor * hc_ffn_inject = nullptr;
+ struct ggml_tensor * hc_attn_down_inject = nullptr; // [hc_attn_down | hc_attn_inject] in one weight
+ struct ggml_tensor * hc_ffn_down_inject = nullptr; // [hc_ffn_down | hc_ffn_inject] in one weight
+ struct ggml_tensor * wqkv_z = nullptr; // [wqkv | wqkv_gate] in one weight
+ struct ggml_tensor * ssm_beta_alpha_cat = nullptr; // [ssm_beta | ssm_alpha] in one weight, unlike the grouped ssm_beta_alpha
struct ggml_tensor * ple_key = nullptr;
struct ggml_tensor * ple_value = nullptr;
diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp
index 70feba14c..bdc4f653f 100644
--- a/src/llama-quant.cpp
+++ b/src/llama-quant.cpp
@@ -125,7 +125,8 @@ static tensor_category tensor_get_category(const std::string & tensor_name) {
if (tensor_name_match_token_embd(tensor_name.c_str())) {
return tensor_category::TOKEN_EMBD;
}
- if (tensor_name.find("attn_qkv.weight") != std::string::npos) {
+ // attn_qkv_gate is [attn_qkv | attn_gate] in one weight, the qkv rows decide its type
+ if (tensor_name.find("attn_qkv.weight") != std::string::npos || tensor_name.find("attn_qkv_gate.weight") != std::string::npos) {
return tensor_category::ATTENTION_QKV;
}
if (tensor_name.find("attn_kv_b.weight") != std::string::npos) {
diff --git a/src/models/models.h b/src/models/models.h
index 5fa552026..59a4e5703 100644
--- a/src/models/models.h
+++ b/src/models/models.h
@@ -2741,6 +2741,7 @@ struct llama_model_qwen4exp : public llama_model_base {
ggml_tensor * w_down,
ggml_tensor * w_up,
ggml_tensor * w_inject,
+ ggml_tensor * w_down_inject, // [w_down | w_inject] in one weight, can be null
ggml_tensor ** inject,
int il);
diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp
index 371db0035..1f3b1cdef 100644
--- a/src/models/qwen4exp.cpp
+++ b/src/models/qwen4exp.cpp
@@ -231,15 +231,32 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) {
const int64_t value_dim = head_v_dim * n_v_heads;
const int64_t conv_dim = key_dim * 2 + value_dim;
+ // matmuls that read the same input: try the merged [a | b] weight the converter writes, else the separate ones
+ auto create_merged = [&](llm_tensor t_merged, llm_tensor t_a, llm_tensor t_b, int64_t n_in, int64_t n_a, int64_t n_b,
+ ggml_tensor *& merged, ggml_tensor *& a, ggml_tensor *& b) {
+ if (flags & TENSOR_SKIP) {
+ const int skip = TENSOR_NOT_REQUIRED | TENSOR_SKIP;
+ create_tensor(tn(t_merged, "weight", il), { n_in, n_a + n_b }, skip | TENSOR_SKIP_IF_VIRTUAL);
+ create_tensor(tn(t_a, "weight", il), { n_in, n_a }, skip);
+ create_tensor(tn(t_b, "weight", il), { n_in, n_b }, skip);
+ return;
+ }
+ merged = create_tensor(tn(t_merged, "weight", il), { n_in, n_a + n_b }, flags | TENSOR_NOT_REQUIRED | TENSOR_SKIP_IF_VIRTUAL);
+ if (!merged) {
+ a = create_tensor(tn(t_a, "weight", il), { n_in, n_a }, flags);
+ b = create_tensor(tn(t_b, "weight", il), { n_in, n_b }, flags);
+ }
+ };
+
// two HC modules per layer: before the token mixer, before the MoE
layer.hc_attn_norm = create_tensor(tn(LLM_TENSOR_HC_ATTN_NORM, "weight", il), { n_embd, hc }, flags | TENSOR_ALLOW_RESHAPE);
- layer.hc_attn_down = create_tensor(tn(LLM_TENSOR_HC_ATTN_DOWN, "weight", il), { hc_dim, hc_lr }, flags);
layer.hc_attn_up = create_tensor(tn(LLM_TENSOR_HC_ATTN_UP, "weight", il), { hc_lr, hc_dim }, flags);
- layer.hc_attn_inject = create_tensor(tn(LLM_TENSOR_HC_ATTN_INJECT, "weight", il), { hc_dim, hc }, flags);
+ create_merged(LLM_TENSOR_HC_ATTN_DOWN_INJECT, LLM_TENSOR_HC_ATTN_DOWN, LLM_TENSOR_HC_ATTN_INJECT, hc_dim, hc_lr, hc,
+ layer.hc_attn_down_inject, layer.hc_attn_down, layer.hc_attn_inject);
layer.hc_ffn_norm = create_tensor(tn(LLM_TENSOR_HC_FFN_NORM, "weight", il), { n_embd, hc }, flags | TENSOR_ALLOW_RESHAPE);
- layer.hc_ffn_down = create_tensor(tn(LLM_TENSOR_HC_FFN_DOWN, "weight", il), { hc_dim, hc_lr }, flags);
layer.hc_ffn_up = create_tensor(tn(LLM_TENSOR_HC_FFN_UP, "weight", il), { hc_lr, hc_dim }, flags);
- layer.hc_ffn_inject = create_tensor(tn(LLM_TENSOR_HC_FFN_INJECT, "weight", il), { hc_dim, hc }, flags);
+ create_merged(LLM_TENSOR_HC_FFN_DOWN_INJECT, LLM_TENSOR_HC_FFN_DOWN, LLM_TENSOR_HC_FFN_INJECT, hc_dim, hc_lr, hc,
+ layer.hc_ffn_down_inject, layer.hc_ffn_down, layer.hc_ffn_inject);
if (!hparams.is_recr(il)) {
// full attention: wq holds [q|gate] interleaved per head
@@ -255,13 +272,13 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) {
layer.index_q_norm = create_tensor(tn(LLM_TENSOR_INDEXER_Q_NORM, "weight", il), { idx_dim }, flags);
layer.index_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", il), { idx_dim }, flags);
} else {
- layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", il), { n_embd, key_dim * 2 + value_dim }, flags);
- layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", il), { n_embd, value_dim }, flags);
+ create_merged(LLM_TENSOR_ATTN_QKV_GATE, LLM_TENSOR_ATTN_QKV, LLM_TENSOR_ATTN_GATE, n_embd, key_dim * 2 + value_dim, value_dim,
+ layer.wqkv_z, layer.wqkv, layer.wqkv_gate);
layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", il), { hparams.ssm_d_conv, conv_dim }, flags);
layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", il), { hparams.ssm_dt_rank }, flags);
layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, il), { hparams.ssm_dt_rank }, flags);
- layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", il), { n_embd, n_v_heads }, flags);
- layer.ssm_alpha = create_tensor(tn(LLM_TENSOR_SSM_ALPHA, "weight", il), { n_embd, n_v_heads }, flags);
+ create_merged(LLM_TENSOR_SSM_BETA_ALPHA_CAT, LLM_TENSOR_SSM_BETA, LLM_TENSOR_SSM_ALPHA, n_embd, n_v_heads, n_v_heads,
+ layer.ssm_beta_alpha_cat, layer.ssm_beta, layer.ssm_alpha);
layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", il), { head_v_dim }, flags);
layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", il), { value_dim, n_embd }, flags);
}
@@ -333,6 +350,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_hc_mix(
ggml_tensor * w_down,
ggml_tensor * w_up,
ggml_tensor * w_inject,
+ ggml_tensor * w_down_inject,
ggml_tensor ** inject,
int il) {
const int64_t hc = hparams.dsv4_hc_mult;
@@ -345,8 +363,20 @@ ggml_tensor * llama_model_qwen4exp::graph::build_hc_mix(
xn = ggml_reshape_2d(ctx0, xn, hc_dim, nt);
cb(xn, "hc_norm", il);
- ggml_tensor * lo = build_lora_mm(w_down, xn);
- lo = ggml_silu(ctx0, ggml_scale(ctx0, lo, 1.0f / (float) hc));
+ ggml_tensor * lo = nullptr;
+ if (inject && w_down_inject) {
+ // one matmul for both: the hc_down rows, then the hc_inject rows
+ const int64_t n_lo = w_down_inject->ne[1] - hc;
+ ggml_tensor * lo_inject = build_lora_mm(w_down_inject, xn);
+ // scale the combined rows before the split; both lo and inject use 1/hc
+ lo_inject = ggml_scale(ctx0, lo_inject, 1.0f / (float) hc);
+ lo = ggml_view_2d(ctx0, lo_inject, n_lo, nt, lo_inject->nb[1], 0);
+ *inject = ggml_view_2d(ctx0, lo_inject, hc, nt, lo_inject->nb[1], ggml_row_size(lo_inject->type, n_lo));
+ cb(*inject, "hc_inject", il);
+ } else {
+ lo = ggml_scale(ctx0, build_lora_mm(w_down, xn), 1.0f / (float) hc);
+ }
+ lo = ggml_silu(ctx0, lo);
ggml_tensor * gate = build_lora_mm(w_up, lo);
cb(gate, "hc_gate", il);
@@ -375,8 +405,8 @@ ggml_tensor * llama_model_qwen4exp::graph::build_hc_mix(
}
cb(mixed, "hc_mixed", il);
- if (inject) {
- *inject = build_lora_mm(w_inject, xn);
+ if (inject && !w_down_inject) {
+ *inject = ggml_scale(ctx0, build_lora_mm(w_inject, xn), 1.0f / (float) hc);
cb(*inject, "hc_inject", il);
}
@@ -392,7 +422,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_hc_combine(
const int64_t nt = residual->ne[2];
// 2*sigmoid centres the scatter weights on 1, so a zero injection is a plain residual add
- ggml_tensor * w = ggml_sigmoid(ctx0, ggml_scale(ctx0, inject, 1.0f / (float) hc));
+ ggml_tensor * w = ggml_sigmoid(ctx0, inject);
w = ggml_scale(ctx0, w, 2.0f);
ggml_tensor * cur = nullptr;
@@ -480,6 +510,7 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa
model.layers[il].hc_attn_down,
model.layers[il].hc_attn_up,
model.layers[il].hc_attn_inject,
+ model.layers[il].hc_attn_down_inject,
&inject, il);
ggml_build_forward_expand(gf, cur);
@@ -507,6 +538,7 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa
model.layers[il].hc_ffn_down,
model.layers[il].hc_ffn_up,
model.layers[il].hc_ffn_inject,
+ model.layers[il].hc_ffn_down_inject,
&inject, il);
cur = build_layer_ffn(cur, il);
@@ -532,7 +564,7 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa
// the final mixer is the output norm: there is no separate one
ggml_tensor * cur = build_hc_mix(res_hc,
model.hc_head_norm, model.hc_head_down, model.hc_head_up,
- nullptr, nullptr, -1);
+ nullptr, nullptr, nullptr, -1);
cb(cur, "result_norm", -1);
res->t_embd = cur;
@@ -610,11 +642,11 @@ llama_model_qwen4exp::graph_mtp::graph_mtp(const llama_model & model, const llm_
cb(res_hc, "mtp_eh_proj", il);
ggml_tensor * inject = nullptr;
- ggml_tensor * cur = build_hc_mix(res_hc, layer.hc_attn_norm, layer.hc_attn_down, layer.hc_attn_up, layer.hc_attn_inject, &inject, il);
+ ggml_tensor * cur = build_hc_mix(res_hc, layer.hc_attn_norm, layer.hc_attn_down, layer.hc_attn_up, layer.hc_attn_inject, layer.hc_attn_down_inject, &inject, il);
cur = build_layer_attn(inp_hyb->get_attn(), mctx_hyb, inp_kpool, cur, inp_pos, sections, il);
res_hc = build_hc_combine(res_hc, cur, inject, il);
- cur = build_hc_mix(res_hc, layer.hc_ffn_norm, layer.hc_ffn_down, layer.hc_ffn_up, layer.hc_ffn_inject, &inject, il);
+ cur = build_hc_mix(res_hc, layer.hc_ffn_norm, layer.hc_ffn_down, layer.hc_ffn_up, layer.hc_ffn_inject, layer.hc_ffn_down_inject, &inject, il);
cur = build_layer_ffn(cur, il);
res_hc = build_hc_combine(res_hc, cur, inject, il);
@@ -627,7 +659,7 @@ llama_model_qwen4exp::graph_mtp::graph_mtp(const llama_model & model, const llm_
cur = build_hc_mix(ggml_reshape_3d(ctx0, flat_out, n_embd, hc, flat_out->ne[1]),
layer.nextn.hc_head_norm, layer.nextn.hc_head_down, layer.nextn.hc_head_up,
- nullptr, nullptr, il);
+ nullptr, nullptr, nullptr, il);
cb(cur, "result_norm", -1);
res->t_embd = cur;
@@ -642,6 +674,23 @@ std::pair<ggml_tensor *, ggml_tensor *> llama_model_qwen4exp::graph::build_qkvz(
const int64_t n_seqs = ubatch.n_seqs;
const int64_t n_seq_tokens = ubatch.n_seq_tokens;
+ if (model.layers[il].wqkv_z) {
+ // qkv and z read the same input: one matmul, then the qkv rows and the z rows
+ ggml_tensor * qkvz = build_lora_mm(model.layers[il].wqkv_z, input);
+ cb(qkvz, "linear_attn_qkvz", il);
+
+ const int64_t n_z = hparams.ssm_d_state * hparams.ssm_dt_rank;
+ const int64_t n_qkv = qkvz->ne[0] - n_z;
+
+ ggml_tensor * qkv_mixed = ggml_view_3d(ctx0, qkvz, n_qkv, n_seq_tokens, n_seqs, qkvz->nb[1], qkvz->nb[1] * n_seq_tokens, 0);
+ cb(qkv_mixed, "linear_attn_qkv_mixed", il);
+
+ ggml_tensor * z = ggml_view_2d(ctx0, qkvz, n_z, n_seq_tokens * n_seqs, qkvz->nb[1], ggml_row_size(qkvz->type, n_qkv));
+ cb(z, "z", il);
+
+ return { qkv_mixed, z };
+ }
+
ggml_tensor * qkv_mixed = build_lora_mm(model.layers[il].wqkv, input, model.layers[il].wqkv_s);
qkv_mixed = ggml_reshape_3d(ctx0, qkv_mixed, qkv_mixed->ne[0], n_seq_tokens, n_seqs);
cb(qkv_mixed, "linear_attn_qkv_mixed", il);
@@ -1072,17 +1121,30 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn_linear(
ggml_tensor * qkv_mixed = qkvz.first;
ggml_tensor * z = qkvz.second;
- ggml_tensor * beta = build_lora_mm(model.layers[il].ssm_beta, cur, model.layers[il].ssm_beta_s);
- beta = ggml_reshape_4d(ctx0, beta, 1, num_v_heads, n_seq_tokens, n_seqs);
+ ggml_tensor * beta = nullptr;
+ ggml_tensor * alpha = nullptr;
+ if (model.layers[il].ssm_beta_alpha_cat) {
+ // beta and alpha read the same input: one matmul, then the beta rows and the alpha rows
+ ggml_tensor * ba = build_lora_mm(model.layers[il].ssm_beta_alpha_cat, cur);
+ cb(ba, "beta_alpha", il);
+
+ beta = ggml_view_4d(ctx0, ba, 1, num_v_heads, n_seq_tokens, n_seqs,
+ ggml_element_size(ba), ba->nb[1], ba->nb[1] * n_seq_tokens, 0);
+ alpha = ggml_view_3d(ctx0, ba, num_v_heads, n_seq_tokens, n_seqs,
+ ba->nb[1], ba->nb[1] * n_seq_tokens, ggml_row_size(ba->type, num_v_heads));
+ } else {
+ beta = build_lora_mm(model.layers[il].ssm_beta, cur, model.layers[il].ssm_beta_s);
+ beta = ggml_reshape_4d(ctx0, beta, 1, num_v_heads, n_seq_tokens, n_seqs);
+
+ alpha = build_lora_mm(model.layers[il].ssm_alpha, cur, model.layers[il].ssm_alpha_s);
+ alpha = ggml_reshape_3d(ctx0, alpha, num_v_heads, n_seq_tokens, n_seqs);
+ }
cb(beta, "beta", il);
+ cb(alpha, "alpha", il);
beta = ggml_sigmoid(ctx0, beta);
cb(beta, "beta_sigmoid", il);
- ggml_tensor * alpha = build_lora_mm(model.layers[il].ssm_alpha, cur, model.layers[il].ssm_alpha_s);
- alpha = ggml_reshape_3d(ctx0, alpha, num_v_heads, n_seq_tokens, n_seqs);
- cb(alpha, "alpha", il);
-
ggml_tensor * alpha_biased = ggml_add(ctx0, alpha, model.layers[il].ssm_dt);
ggml_tensor * alpha_softplus = ggml_softplus(ctx0, alpha_biased);
cb(alpha_softplus, "a_softplus", il);
@@ -1160,7 +1222,8 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn_linear(
ggml_tensor * output = build_recurrent_attn(inp, ssm_states_all, q_conv, k_conv, v_conv, gate, beta, state, il);
- ggml_tensor * z_2d = ggml_reshape_4d(ctx0, z, head_v_dim, num_v_heads, n_seq_tokens, n_seqs);
+ ggml_tensor * z_2d = ggml_view_4d(ctx0, z, head_v_dim, num_v_heads, n_seq_tokens, n_seqs,
+ ggml_row_size(z->type, head_v_dim), z->nb[1], z->nb[1] * n_seq_tokens, 0);
// gated normalization, as self.norm(core_attn_out, z) in the reference
ggml_tensor * attn_out_norm = build_norm_gated(output, model.layers[il].ssm_norm, z_2d, il);