name, gen = item
# Skip multimodal tensors
- if name.startswith(("mlp", "vit.", "vpm.", "siglip2.", "conformer.", "merger.", "resampler.", "sound_encoder.", "sound_projection.")) \
+ if name.startswith(("mlp", "vit.", "vpm.", "siglip2.", "conformer.", "merger.", "resampler.", "sound_encoder.", "sound_projection.", "speech_embeddings.")) \
or "visual." in name or "vision." in name or "audio." in name or "talker." in name \
or "vision_" in name or "audio_" in name or "sam_model" in name \
or "token2wav." in name or "code2wav." in name \
yield from super().modify_tensors(data_torch, name, bid)
-@ModelBase.register("MiMoV2FlashForCausalLM")
+@ModelBase.register("MiMoV2FlashForCausalLM", "MiMoV2ForCausalLM")
class MimoV2Model(TextModel):
model_arch = gguf.MODEL_ARCH.MIMO2
+ # MiMo V2-Flash, V2.5 and V2.5-Pro all ship 3 trained MTP layers under model.mtp.layers.{0,1,2}.
+ # The HF config does not expose the count, so it's hardcoded to match the count found in the safetensors.
+ _n_nextn = 3
+
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+
+ self.block_count = self.hparams["num_hidden_layers"] + self._n_nextn
+ self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
+
+ @staticmethod
+ def _tp_aware_qkv_dequant(weight: Tensor, scale_inv: Tensor,
+ n_q: int, n_kv: int, hd: int, vhd: int,
+ bs: int = 128) -> Tensor:
+ # MiMo-V2.5 (TP=4) and V2.5-Pro (TP=8) ship qkv_proj sharded across TP
+ # ranks; per rank, rows are stacked as [Q_per | K_per | V_per].
+ # weight_scale_inv has ceil(rows_per_rank/bs) block-rows per rank (last
+ # may extend past rows_per_rank with phantom rows not in the weight).
+ # Naive repeat_interleave aligns rank 0 only and mis-applies scales to
+ # later ranks once rows_per_rank isn't a multiple of bs.
+ # Re-group the per-rank [Q_per|K_per|V_per] rows into a single fused
+ # [Q | K | V] tensor matching the un-sharded original layout.
+ q_size = n_q * hd
+ k_size = n_kv * hd
+ v_size = n_kv * vhd
+ total_rows = q_size + k_size + v_size
+ if weight.shape[0] != total_rows:
+ raise ValueError(f"qkv_proj weight rows {weight.shape[0]} != q+k+v {total_rows}")
+
+ # detect TP from scale_inv block count, descending order so larger matches first
+ tp = None
+ for cand in (8, 4):
+ if total_rows % cand != 0:
+ continue
+ rpr = total_rows // cand
+ bpr = (rpr + bs - 1) // bs
+ if scale_inv.shape[0] == cand * bpr:
+ tp = cand
+ break
+ if tp is None:
+ raise ValueError(
+ f"qkv_proj: cannot detect TP - scale_inv rows {scale_inv.shape[0]}, "
+ f"q+k+v {total_rows}")
+
+ q_per = q_size // tp
+ k_per = k_size // tp
+ v_per = v_size // tp
+ rows_per_rank = q_per + k_per + v_per
+ blocks_per_rank = (rows_per_rank + bs - 1) // bs
+
+ scale_inv = scale_inv.float()
+ # per-row scale-row index: rank * blocks_per_rank + (rr_in_rank // bs)
+ row_idx = torch.arange(total_rows)
+ rr = row_idx % rows_per_rank
+ rank = row_idx // rows_per_rank
+ scale_row_idx = rank * blocks_per_rank + (rr // bs)
+ # gather: (total_rows, n_col_blocks)
+ scale_per_row_block = scale_inv[scale_row_idx]
+ # expand col-blocks -> cols: each block-col covers `bs` weight cols
+ scale_full = scale_per_row_block.repeat_interleave(bs, dim=1)
+ # crop to weight col count (in case last col-block isn't full)
+ scale_full = scale_full[:, : weight.shape[1]]
+ dequant = weight.float() * scale_full
+
+ if tp == 1:
+ return dequant
+
+ # Re-group per-rank [Q_per|K_per|V_per] rows into unified [Q | K | V]
+ qs, ks, vs = [], [], []
+ for r in range(tp):
+ base = r * rows_per_rank
+ qs.append(dequant[base : base + q_per])
+ ks.append(dequant[base + q_per : base + q_per + k_per])
+ vs.append(dequant[base + q_per + k_per : base + rows_per_rank])
+ return torch.cat(qs + ks + vs, dim=0)
+
+ def dequant_model(self):
+ # Capture raw FP8 (weight, scale_inv) lambdas for qkv_proj BEFORE super
+ # rewrites them with the existing dequant. Replace super's lambda after
+ # it runs so scale_inv removal still happens via the standard path.
+ qkv_overrides: dict[str, tuple[Callable, Callable, int]] = {}
+ qc = self.hparams.get("quantization_config")
+ if isinstance(qc, dict) and qc.get("quant_method") == "fp8":
+ pat = re.compile(r"^model\.layers\.(\d+)\.self_attn\.qkv_proj\.weight_scale_inv$")
+ for name in list(self.model_tensors.keys()):
+ m = pat.match(name)
+ if not m:
+ continue
+ weight_name = name.removesuffix("_scale_inv")
+ if weight_name not in self.model_tensors:
+ continue
+ qkv_overrides[weight_name] = (
+ self.model_tensors[weight_name],
+ self.model_tensors[name],
+ int(m.group(1)),
+ )
+
+ super().dequant_model()
+
+ if not qkv_overrides:
+ return
+
+ n_q = self.hparams["num_attention_heads"]
+ hd = self.hparams["head_dim"]
+ vhd = self.hparams["v_head_dim"]
+ hybrid = self.hparams["hybrid_layer_pattern"]
+ n_layer_text = self.hparams["num_hidden_layers"]
+ for weight_name, (w_fn, s_fn, bid) in qkv_overrides.items():
+ # MTP layers (bid >= n_layer_text) use SWA-style attention dims
+ is_swa = True if bid >= n_layer_text else hybrid[bid] == 1
+ n_kv = self.hparams["swa_num_key_value_heads" if is_swa else "num_key_value_heads"]
+ self.model_tensors[weight_name] = (
+ lambda w_fn=w_fn, s_fn=s_fn, n_q=n_q, n_kv=n_kv, hd=hd, vhd=vhd:
+ MimoV2Model._tp_aware_qkv_dequant(w_fn(), s_fn(), n_q, n_kv, hd, vhd)
+ )
+
def set_gguf_parameters(self):
super().set_gguf_parameters()
n_head_kv = self.hparams["num_key_value_heads"]
n_head_kv_swa = self.hparams["swa_num_key_value_heads"]
- n_head_kv_arr = [n_head_kv_swa if use_swa == 1 else n_head_kv for use_swa in self.hparams["hybrid_layer_pattern"]]
+ # Extend the per-layer pattern with SWA entries for the MTP blocks so the
+ # runtime arrays (sized to extended block_count) are fully populated.
+ hybrid = list(self.hparams["hybrid_layer_pattern"]) + [1] * self._n_nextn
+ n_head_kv_arr = [n_head_kv_swa if use_swa == 1 else n_head_kv for use_swa in hybrid]
self.gguf_writer.add_head_count_kv(n_head_kv_arr)
self.gguf_writer.add_sliding_window(self.hparams["sliding_window"])
- self.gguf_writer.add_sliding_window_pattern(self.hparams["hybrid_layer_pattern"])
+ self.gguf_writer.add_sliding_window_pattern(hybrid)
self.gguf_writer.add_value_length(self.hparams["v_head_dim"])
self.gguf_writer.add_expert_count(self.hparams["n_routed_experts"])
self.gguf_writer.add_expert_feed_forward_length(self.hparams["moe_intermediate_size"])
self.gguf_writer.add_layer_norm_rms_eps(self.hparams.get("layernorm_epsilon", 1e-5))
+ v_scale = self.hparams.get("attention_value_scale")
+ if v_scale is not None:
+ self.gguf_writer.add_attn_value_scale(float(v_scale))
+
+ self.gguf_writer.add_nextn_predict_layers(self._n_nextn)
+
_experts: list[dict[str, Tensor]] | None = None
@classmethod
if "attention_sink" in name and not name.endswith(".weight"):
name += ".weight"
- # TODO: mimo v2 does not indicate the number of next-token-prediction layers, therefore we cannot do the same way as GLM4_MOE
- if "model.mtp." in name:
- return None
-
return super().filter_tensors((name, gen))
def modify_tensors(self, data_torch, name, bid):
+ # Remap MTP/NextN tensors to additional layer slots so the standard tensor map handles them.
+ # HF: model.mtp.layers.{i}.foo -> model.layers.{n_layer_text + i}.foo
+ m = re.match(r"^model\.mtp\.layers\.(\d+)\.(.*)$", name)
+ if m is not None:
+ mtp_idx = int(m.group(1))
+ assert mtp_idx < self._n_nextn, f"MTP layer index {mtp_idx} >= _n_nextn ({self._n_nextn})"
+ rest = m.group(2)
+ n_layer_text = self.hparams["num_hidden_layers"]
+ new_bid = n_layer_text + mtp_idx
+ name = f"model.layers.{new_bid}.{rest}"
+ bid = new_bid
+
# process the experts separately
if name.find("mlp.experts") != -1:
n_experts = self.hparams["n_routed_experts"]
SLIDING_WINDOW = "{arch}.attention.sliding_window"
SCALE = "{arch}.attention.scale"
OUTPUT_SCALE = "{arch}.attention.output_scale"
+ VALUE_SCALE = "{arch}.attention.value_scale"
TEMPERATURE_LENGTH = "{arch}.attention.temperature_length"
KEY_LENGTH_MLA = "{arch}.attention.key_length_mla"
VALUE_LENGTH_MLA = "{arch}.attention.value_length_mla"
MODEL_TENSOR.OUTPUT_NORM,
MODEL_TENSOR.OUTPUT,
MODEL_TENSOR.ATTN_NORM,
+ MODEL_TENSOR.ATTN_QKV,
MODEL_TENSOR.ATTN_Q,
MODEL_TENSOR.ATTN_K,
MODEL_TENSOR.ATTN_V,
MODEL_TENSOR.FFN_DOWN_EXP,
MODEL_TENSOR.FFN_UP_EXP,
MODEL_TENSOR.FFN_EXP_PROBS_B,
+ MODEL_TENSOR.LAYER_OUT_NORM,
+ MODEL_TENSOR.NEXTN_EH_PROJ,
+ MODEL_TENSOR.NEXTN_ENORM,
+ MODEL_TENSOR.NEXTN_HNORM,
],
MODEL_ARCH.STEP35: [
MODEL_TENSOR.TOKEN_EMBD,
def add_attn_output_scale(self, value: float) -> None:
self.add_float32(Keys.Attention.OUTPUT_SCALE.format(arch=self.arch), value)
+ def add_attn_value_scale(self, value: float) -> None:
+ self.add_float32(Keys.Attention.VALUE_SCALE.format(arch=self.arch), value)
+
def add_attn_temperature_length(self, value: int) -> None:
self.add_uint32(Keys.Attention.TEMPERATURE_LENGTH.format(arch=self.arch), value)
{ LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, "%s.attention.sliding_window_pattern" },
{ LLM_KV_ATTENTION_SCALE, "%s.attention.scale" },
{ LLM_KV_ATTENTION_OUTPUT_SCALE, "%s.attention.output_scale" },
+ { LLM_KV_ATTENTION_VALUE_SCALE, "%s.attention.value_scale" },
{ LLM_KV_ATTENTION_TEMPERATURE_LENGTH, "%s.attention.temperature_length" },
{ LLM_KV_ATTENTION_TEMPERATURE_SCALE, "%s.attention.temperature_scale" },
{ LLM_KV_ATTENTION_KEY_LENGTH_MLA, "%s.attention.key_length_mla" },
LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN,
LLM_KV_ATTENTION_SCALE,
LLM_KV_ATTENTION_OUTPUT_SCALE,
+ LLM_KV_ATTENTION_VALUE_SCALE,
LLM_KV_ATTENTION_TEMPERATURE_LENGTH,
LLM_KV_ATTENTION_TEMPERATURE_SCALE,
LLM_KV_ATTENTION_KEY_LENGTH_MLA,
float f_attn_out_scale = 0.0f;
uint32_t attn_temp_length = 0;
+ float f_attn_value_scale = 0.0f;
+
bool causal_attn = true;
bool use_alibi = false;
bool attn_soft_cap = false;
// add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, ???);
add_kv(LLM_KV_ATTENTION_SCALE, hparams.f_attention_scale);
add_kv(LLM_KV_ATTENTION_OUTPUT_SCALE, hparams.f_attn_out_scale);
+ add_kv(LLM_KV_ATTENTION_VALUE_SCALE, hparams.f_attn_value_scale);
add_kv(LLM_KV_ATTENTION_TEMPERATURE_LENGTH, hparams.attn_temp_length);
add_kv(LLM_KV_ATTENTION_TEMPERATURE_SCALE, hparams.f_attn_temp_scale);
add_kv(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl);
LLAMA_LOG_INFO("%s: f_max_alibi_bias = %.1e\n", __func__, hparams.f_max_alibi_bias);
LLAMA_LOG_INFO("%s: f_logit_scale = %.1e\n", __func__, hparams.f_logit_scale);
LLAMA_LOG_INFO("%s: f_attn_scale = %.1e\n", __func__, hparams.f_attention_scale);
+ LLAMA_LOG_INFO("%s: f_attn_value_scale = %.4f\n", __func__, hparams.f_attn_value_scale);
LLAMA_LOG_INFO("%s: n_ff = %s\n", __func__, print_f([&](uint32_t il) { return hparams.n_ff(il); }, hparams.n_layer).c_str());
LLAMA_LOG_INFO("%s: n_expert = %u\n", __func__, hparams.n_expert);
LLAMA_LOG_INFO("%s: n_expert_used = %u\n", __func__, hparams.n_expert_used);
ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false);
ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.swa_layers, hparams.n_layer);
- switch (hparams.n_layer) {
+ float value_scale = 0.0f;
+ if (ml.get_key(LLM_KV_ATTENTION_VALUE_SCALE, value_scale, false) && value_scale != 1.0f) {
+ hparams.f_attn_value_scale = value_scale;
+ }
+
+ ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.nextn_predict_layers, false);
+ GGML_ASSERT(hparams.nextn_predict_layers < hparams.n_layer && "nextn_predict_layers must be < n_layer");
+ hparams.n_layer_kv_from_start = hparams.n_layer - hparams.nextn_predict_layers;
+
+ switch (hparams.n_layer - hparams.nextn_predict_layers) {
case 48: type = LLM_TYPE_310B_A15B; break;
default: type = LLM_TYPE_UNKNOWN;
}
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0);
+ const uint32_t n_nextn = hparams.nextn_predict_layers;
+
for (int i = 0; i < n_layer; ++i) {
auto & layer = layers[i];
uint32_t n_embd_k_gqa = hparams.n_embd_k_gqa(i);
uint32_t n_embd_v_gqa = hparams.n_embd_v_gqa(i);
uint32_t n_head = hparams.n_head(i);
- create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, 0);
- layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_v * n_head, n_embd }, 0);
+ // NextN/MTP layers (the last n_nextn blocks) are preserved but disabled pending support
+ const bool is_nextn = (n_nextn > 0) && (static_cast<uint32_t>(i) >= n_layer - n_nextn);
+ const int skip = is_nextn ? TENSOR_SKIP : 0;
- layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
- layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, TENSOR_NOT_REQUIRED);
+ create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, skip);
+ layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_v * n_head, n_embd }, skip);
- layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
+ layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, skip);
+ layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, TENSOR_NOT_REQUIRED | skip);
+
+ layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, skip);
// non-MoE branch
- layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, TENSOR_NOT_REQUIRED);
- layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, TENSOR_NOT_REQUIRED);
- layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, TENSOR_NOT_REQUIRED);
+ layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, TENSOR_NOT_REQUIRED | skip);
+ layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, TENSOR_NOT_REQUIRED | skip);
+ layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, TENSOR_NOT_REQUIRED | skip);
// MoE branch
int64_t n_ff_exp = hparams.n_ff_exp;
- layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, TENSOR_NOT_REQUIRED);
- layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED);
- layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, TENSOR_NOT_REQUIRED);
- layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED);
- layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED);
+ layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, TENSOR_NOT_REQUIRED | skip);
+ layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED | skip);
+ layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, TENSOR_NOT_REQUIRED | skip);
+ layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED | skip);
+ layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED | skip);
+
+ if (is_nextn) {
+ layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2 * n_embd, n_embd}, skip);
+ layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, skip);
+ layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, skip);
+ layer.layer_out_norm = create_tensor(tn(LLM_TENSOR_LAYER_OUT_NORM, "weight", i), {n_embd}, skip);
+ }
}
}
auto * inp_attn = build_attn_inp_kv_iswa();
ggml_tensor * inp_out_ids = build_inp_out_ids();
- for (int il = 0; il < n_layer; ++il) {
+ const float v_scale = hparams.f_attn_value_scale;
+
+ // The last hparams.nextn_predict_layers blocks are MTP heads, currently inactive
+ const int n_transformer_layers = n_layer - hparams.nextn_predict_layers;
+
+ for (int il = 0; il < n_transformer_layers; ++il) {
ggml_tensor * inpSA = inpL;
uint32_t n_head_l = hparams.n_head(il);
cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
cb(cur, "attn_norm", il);
- // compute Q and K and RoPE them
- ggml_tensor * Qcur = build_lora_mm(model.layers[il].wq, cur);
- cb(Qcur, "Qcur", il);
+ ggml_tensor * Qcur;
+ ggml_tensor * Kcur;
+ ggml_tensor * Vcur;
- ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur);
- cb(Kcur, "Kcur", il);
+ if (model.layers[il].wqkv) {
+ // Fused qkv_proj - Q/K share head_dim_k, V uses head_dim_v
+ ggml_tensor * qkv = build_lora_mm(model.layers[il].wqkv, cur);
+ cb(qkv, "wqkv", il);
- ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur);
- cb(Vcur, "Vcur", il);
+ const size_t row_k = ggml_row_size(qkv->type, n_embd_head_k);
+ const size_t row_v = ggml_row_size(qkv->type, n_embd_head_v);
+ const size_t row_full = qkv->nb[1];
+ const size_t k_off = row_k * n_head_l;
+ const size_t v_off = k_off + row_k * n_head_kv_l;
+
+ Qcur = ggml_view_3d(ctx0, qkv, n_embd_head_k, n_head_l, n_tokens, row_k, row_full, 0);
+ Kcur = ggml_view_3d(ctx0, qkv, n_embd_head_k, n_head_kv_l, n_tokens, row_k, row_full, k_off);
+ Vcur = ggml_view_3d(ctx0, qkv, n_embd_head_v, n_head_kv_l, n_tokens, row_v, row_full, v_off);
+ } else {
+ // Split path
+ Qcur = build_lora_mm(model.layers[il].wq, cur);
+ cb(Qcur, "Qcur", il);
- Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head_k, n_head_l, n_tokens);
- Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head_k, n_head_kv_l, n_tokens);
- Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head_v, n_head_kv_l, n_tokens);
+ Kcur = build_lora_mm(model.layers[il].wk, cur);
+ cb(Kcur, "Kcur", il);
+
+ Vcur = build_lora_mm(model.layers[il].wv, cur);
+ cb(Vcur, "Vcur", il);
+
+ Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head_k, n_head_l, n_tokens);
+ Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head_k, n_head_kv_l, n_tokens);
+ Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head_v, n_head_kv_l, n_tokens);
+ }
Qcur = ggml_rope_ext(
ctx0, Qcur, inp_pos, nullptr,
cur = build_attn(inp_attn,
model.layers[il].wo, NULL, model.layers[il].wo_s,
Qcur, Kcur, Vcur, nullptr, sinks, nullptr, 1.0f/sqrtf(float(n_embd_head_k)), il);
+ cb(cur, "attn_out", il);
+
+ if (v_scale) {
+ cur = ggml_scale(ctx0, cur, v_scale);
+ cb(cur, "attn_out_scaled", il);
+ }
}
- if (il == n_layer - 1 && inp_out_ids) {
+ if (il == n_transformer_layers - 1 && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
overview / pkg / ggml / sources / llama.cpp / commitdiff