From 2280f0fd163b03d4331d49e08460f66502e35854 Mon Sep 17 00:00:00 2001 From: Aman Gupta Date: Sun, 25 Oct 2026 21:23:22 +0800 Subject: [PATCH] qwen: merge tensors on convert (#31028) * qwen: merge tensors on convert * convert: use generator functions * qwen4exp: remove conts and pre-scale --- conversion/qwen4exp.py | 69 ++++++++++++++++++------ gguf-py/gguf/constants.py | 12 +++++ src/llama-arch.cpp | 8 +++ src/llama-arch.h | 4 ++ src/llama-model.h | 4 ++ src/llama-quant.cpp | 3 +- src/models/models.h | 1 + src/models/qwen4exp.cpp | 111 +++++++++++++++++++++++++++++--------- 8 files changed, 172 insertions(+), 40 deletions(-) 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 - return super().modify_tensors(data_torch, name, bid) + 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 + + 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_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_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 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);