diff --git a/conversion/base.py b/conversion/base.py index daae28e92adc..aa0865a8c862 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -130,7 +130,8 @@ def __init__(self, dir_model: Path, ftype: gguf.LlamaFileType, fname_out: Path, sentence_transformers_dense_modules: bool = False, target_model_dir: Path | None = None, fuse_gate_up_exps: bool = False, - fp8_as_q8: bool = False): + fp8_as_q8: bool = False, + fuse_qkv: bool = False): if type(self) is ModelBase or \ type(self) is TextModel or \ type(self) is MmprojModel: @@ -153,6 +154,15 @@ def __init__(self, dir_model: Path, ftype: gguf.LlamaFileType, fname_out: Path, self.fuse_gate_up_exps = fuse_gate_up_exps self._gate_exp_buffer: dict[int, Tensor] = {} self._up_exp_buffer: dict[int, Tensor] = {} + self.fuse_qkv = fuse_qkv + self._q_buffer: dict[int, Tensor] = {} + self._k_buffer: dict[int, Tensor] = {} + self._v_buffer: dict[int, Tensor] = {} + self._q_bias_buffer: dict[int, Tensor] = {} + self._k_bias_buffer: dict[int, Tensor] = {} + self._v_bias_buffer: dict[int, Tensor] = {} + self._fusable_qkv_weight_layers: set[int] = set() + self._fusable_qkv_bias_layers: set[int] = set() self.hparams = ModelBase.load_hparams(self.dir_model, self.is_mistral_format) if hparams is None else hparams self.model_tensors = self.index_tensors(remote_hf_model_id=remote_hf_model_id) self.metadata_override = metadata_override @@ -617,6 +627,43 @@ def map_tensor_name(self, name: str, try_suffixes: Sequence[str] = (".weight", " raise ValueError(f"Can not map tensor {name!r}") return new_name + def prepare_qkv_fusion(self) -> None: + self._fusable_qkv_weight_layers.clear() + self._fusable_qkv_bias_layers.clear() + if not self.fuse_qkv or gguf.MODEL_TENSOR.ATTN_QKV not in gguf.MODEL_TENSORS[self.model_arch]: + return + + qkv_types = { + gguf.MODEL_TENSOR.ATTN_Q, + gguf.MODEL_TENSOR.ATTN_K, + gguf.MODEL_TENSOR.ATTN_V, + } + weights: dict[int, set[gguf.MODEL_TENSOR]] = {} + biases: dict[int, set[gguf.MODEL_TENSOR]] = {} + + for name in self.model_tensors: + mapped = self.tensor_map.get_type_and_name(name, try_suffixes=(".weight", ".bias")) + if mapped is None: + continue + tensor_type, new_name = mapped + if tensor_type not in qkv_types: + continue + + bid = next((int(part) for part in new_name.split(".") if part.isdecimal()), None) + if bid is None: + continue + if new_name.endswith(".weight"): + weights.setdefault(bid, set()).add(tensor_type) + elif new_name.endswith(".bias"): + biases.setdefault(bid, set()).add(tensor_type) + + for bid, weight_types in weights.items(): + bias_types = biases.get(bid, set()) + if weight_types == qkv_types and (not bias_types or bias_types == qkv_types): + self._fusable_qkv_weight_layers.add(bid) + if bias_types: + self._fusable_qkv_bias_layers.add(bid) + def set_gguf_parameters(self): raise NotImplementedError("set_gguf_parameters() must be implemented in subclasses") @@ -645,6 +692,40 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.FFN_UP_EXP, bid): return [] + # Handle Q/K/V tensor fusion if enabled + qkv_bid = next((int(part) for part in new_name.split(".") if part.isdecimal()), None) if self.fuse_qkv else None + if qkv_bid is not None: + is_bias = new_name.endswith('.bias') + suffix = '.bias' if is_bias else '.weight' + fusable_layers = self._fusable_qkv_bias_layers if is_bias else self._fusable_qkv_weight_layers + if qkv_bid not in fusable_layers: + return [(new_name, data_torch)] + + buf_q = self._q_bias_buffer if is_bias else self._q_buffer + buf_k = self._k_bias_buffer if is_bias else self._k_buffer + buf_v = self._v_bias_buffer if is_bias else self._v_buffer + + if self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.ATTN_Q, qkv_bid, suffix): + buf_q[qkv_bid] = data_torch + elif self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.ATTN_K, qkv_bid, suffix): + buf_k[qkv_bid] = data_torch + elif self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.ATTN_V, qkv_bid, suffix): + buf_v[qkv_bid] = data_torch + + if qkv_bid in buf_q and qkv_bid in buf_k and qkv_bid in buf_v: + q_data = buf_q.pop(qkv_bid) + k_data = buf_k.pop(qkv_bid) + v_data = buf_v.pop(qkv_bid) + fused_data = torch.cat([q_data, k_data, v_data], dim=0) + fused_name = self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_QKV, qkv_bid, suffix=suffix) + logger.info(f"Fused Q, K, V {suffix[1:]} into QKV for layer {qkv_bid}") + return [(fused_name, fused_data)] + + if self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.ATTN_Q, qkv_bid, suffix) or \ + self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.ATTN_K, qkv_bid, suffix) or \ + self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.ATTN_V, qkv_bid, suffix): + return [] + return [(new_name, data_torch)] def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool: @@ -899,6 +980,8 @@ def load(): self.dequant_model() + self.prepare_qkv_fusion() + # Handle empty tensor_map for models with block_count=0 (like MobileNetV5) if self.tensor_map.mapping: max_name_len = max(len(s) for _, s in self.tensor_map.mapping.values()) + len(".weight,") @@ -1027,6 +1110,13 @@ def load(): self.gguf_writer.add_tensor(new_name, data, raw_dtype=data_qtype) + qkv_buffers = ( + self._q_buffer, self._k_buffer, self._v_buffer, + self._q_bias_buffer, self._k_bias_buffer, self._v_bias_buffer, + ) + if any(qkv_buffers): + raise ValueError("QKV fusion did not consume all buffered tensors") + def set_type(self): self.gguf_writer.add_type(gguf.GGUFType.MODEL) diff --git a/convert_hf_to_gguf.py b/convert_hf_to_gguf.py index 78ad26c65630..e09616b190cf 100755 --- a/convert_hf_to_gguf.py +++ b/convert_hf_to_gguf.py @@ -157,6 +157,10 @@ def parse_args() -> argparse.Namespace: help="Store tensors dequantized from FP8 as Q8_0 instead of BF16/F16.", ) + parser.add_argument( + "--fuse-qkv", action="store_true", + help="Fuse separate Q, K, V weight tensors into a single QKV tensor.", + ) parser.add_argument( "--target-model-dir", type=str, default=None, help=( @@ -290,6 +294,7 @@ def main() -> None: target_model_dir=Path(args.target_model_dir) if args.target_model_dir else None, fuse_gate_up_exps=args.fuse_gate_up_exps, fp8_as_q8=args.fp8_as_q8, + fuse_qkv=args.fuse_qkv, ) if args.vocab_only: diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index b85f62a31145..166682e98f37 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -2290,6 +2290,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -2310,6 +2311,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -2333,6 +2335,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -2353,6 +2356,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -2398,6 +2402,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -2500,6 +2505,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.TOKEN_TYPES, MODEL_TENSOR.ATTN_NORM_2, MODEL_TENSOR.ATTN_OUT_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -2528,6 +2534,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -2557,6 +2564,7 @@ class MODEL_TENSOR(IntEnum): 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, @@ -2569,6 +2577,7 @@ class MODEL_TENSOR(IntEnum): 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, @@ -2596,6 +2605,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -2627,6 +2637,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -2642,6 +2653,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -2657,6 +2669,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -2671,6 +2684,7 @@ class MODEL_TENSOR(IntEnum): 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, @@ -2685,6 +2699,7 @@ class MODEL_TENSOR(IntEnum): 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, @@ -2705,6 +2720,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -2721,6 +2737,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -2776,6 +2793,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -2792,6 +2810,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -2932,6 +2951,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3065,6 +3085,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3080,6 +3101,7 @@ class MODEL_TENSOR(IntEnum): 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, @@ -3098,6 +3120,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.ROPE_FACTORS_LONG, MODEL_TENSOR.ROPE_FACTORS_SHORT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3135,6 +3158,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3147,6 +3171,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.GEMMA2: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3163,6 +3188,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -3181,6 +3207,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -3217,6 +3244,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -3272,6 +3300,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.DENSE_2_OUT, MODEL_TENSOR.DENSE_3_OUT, MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -3292,6 +3321,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3455,6 +3485,7 @@ class MODEL_TENSOR(IntEnum): 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, @@ -3484,6 +3515,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3498,6 +3530,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3512,6 +3545,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3563,6 +3597,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.OLMO: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3575,6 +3610,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3590,6 +3626,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.SEED_OSS: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3606,6 +3643,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3656,6 +3694,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3677,6 +3716,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3739,6 +3779,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_A, MODEL_TENSOR.ATTN_Q_B, @@ -3861,6 +3902,7 @@ class MODEL_TENSOR(IntEnum): 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, @@ -3937,6 +3979,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4082,6 +4125,7 @@ class MODEL_TENSOR(IntEnum): 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, @@ -4096,6 +4140,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4117,6 +4162,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.SSM_D, MODEL_TENSOR.SSM_NORM, MODEL_TENSOR.SSM_OUT, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4136,6 +4182,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.SSM_D, MODEL_TENSOR.SSM_NORM, MODEL_TENSOR.SSM_OUT, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4166,6 +4213,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4181,6 +4229,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -4206,6 +4255,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -4237,6 +4287,7 @@ class MODEL_TENSOR(IntEnum): 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, @@ -4251,6 +4302,7 @@ class MODEL_TENSOR(IntEnum): 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, @@ -4276,6 +4328,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.SSM_D, MODEL_TENSOR.SSM_NORM, MODEL_TENSOR.SSM_OUT, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4339,6 +4392,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -4378,6 +4432,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4469,6 +4524,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -4532,6 +4588,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4547,6 +4604,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4598,6 +4656,7 @@ class MODEL_TENSOR(IntEnum): 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, @@ -4612,6 +4671,7 @@ class MODEL_TENSOR(IntEnum): 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, @@ -4629,6 +4689,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.ATTN_NORM, # Attention components + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, # Query projection MODEL_TENSOR.ATTN_K, # Key projection MODEL_TENSOR.ATTN_V, # Value projection @@ -4661,6 +4722,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -4681,6 +4743,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -4697,6 +4760,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -4745,6 +4809,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4761,6 +4826,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4784,6 +4850,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.ATTN_NORM, # operator_norm MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4804,6 +4871,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.ATTN_NORM, # operator_norm MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4819,6 +4887,7 @@ class MODEL_TENSOR(IntEnum): 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, @@ -4838,6 +4907,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4855,6 +4925,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4872,6 +4943,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -4910,6 +4982,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -4973,6 +5046,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -4990,6 +5064,7 @@ class MODEL_TENSOR(IntEnum): 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, @@ -5005,6 +5080,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -5158,6 +5234,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -5191,6 +5268,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -5210,6 +5288,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -5226,6 +5305,7 @@ class MODEL_TENSOR(IntEnum): 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, diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 274a6264336f..7f8f9e507abc 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1620,8 +1620,26 @@ llm_graph_qkv llm_graph_context::build_qkv( int64_t n_head, int64_t n_head_kv, int il) const { - const int64_t n_embd_q = n_embd_head * n_head; - const int64_t n_embd_kv = n_embd_head * n_head_kv; + return build_qkv(layer, cur, + n_embd_head, n_head, + n_embd_head, n_head_kv, + n_embd_head, n_head_kv, + il); +} + +llm_graph_qkv llm_graph_context::build_qkv( + const llama_layer & layer, + ggml_tensor * cur, + int64_t n_embd_head_q, + int64_t n_head_q, + int64_t n_embd_head_k, + int64_t n_head_k, + int64_t n_embd_head_v, + int64_t n_head_v, + int il, + bool reshape) const { + const int64_t n_embd_q = n_embd_head_q * n_head_q; + const int64_t n_embd_k = n_embd_head_k * n_head_k; ggml_tensor * Qcur, * Kcur, * Vcur; @@ -1632,59 +1650,93 @@ llm_graph_qkv llm_graph_context::build_qkv( if (layer.wqkv_b) { qkv = ggml_add(ctx0, qkv, layer.wqkv_b); cb(qkv, "wqkv_b", il); + } else if (layer.wq_b && layer.wk_b && layer.wv_b) { + // Fused weights may coexist with separate Q/K/V biases in legacy or custom GGUFs. + ggml_tensor * qkv_b = ggml_concat(ctx0, ggml_concat(ctx0, layer.wq_b, layer.wk_b, 0), layer.wv_b, 0); + qkv = ggml_add(ctx0, qkv, qkv_b); + cb(qkv, "wqkv_b", il); } - if (hparams.f_clamp_kqv > 0.0f) { + if (reshape && hparams.f_clamp_kqv > 0.0f) { qkv = ggml_clamp(ctx0, qkv, -hparams.f_clamp_kqv, hparams.f_clamp_kqv); cb(qkv, "wqkv_clamped", il); } - Qcur = ggml_view_3d(ctx0, qkv, n_embd_head, n_head, n_tokens, - ggml_row_size(qkv->type, n_embd_head), qkv->nb[1], 0); - Kcur = ggml_view_3d(ctx0, qkv, n_embd_head, n_head_kv, n_tokens, - ggml_row_size(qkv->type, n_embd_head), qkv->nb[1], - ggml_row_size(qkv->type, n_embd_q)); - Vcur = ggml_view_3d(ctx0, qkv, n_embd_head, n_head_kv, n_tokens, - ggml_row_size(qkv->type, n_embd_head), qkv->nb[1], - ggml_row_size(qkv->type, n_embd_q + n_embd_kv)); + if (reshape) { + Qcur = ggml_view_3d(ctx0, qkv, n_embd_head_q, n_head_q, n_tokens, + ggml_row_size(qkv->type, n_embd_head_q), qkv->nb[1], 0); + Kcur = ggml_view_3d(ctx0, qkv, n_embd_head_k, n_head_k, n_tokens, + ggml_row_size(qkv->type, n_embd_head_k), qkv->nb[1], + ggml_row_size(qkv->type, n_embd_q)); + Vcur = ggml_view_3d(ctx0, qkv, n_embd_head_v, n_head_v, n_tokens, + ggml_row_size(qkv->type, n_embd_head_v), qkv->nb[1], + ggml_row_size(qkv->type, n_embd_q + n_embd_k)); + } else { + Qcur = ggml_view_2d(ctx0, qkv, n_embd_q, n_tokens, qkv->nb[1], 0); + Kcur = ggml_view_2d(ctx0, qkv, n_embd_k, n_tokens, qkv->nb[1], + ggml_row_size(qkv->type, n_embd_q)); + Vcur = ggml_view_2d(ctx0, qkv, n_embd_head_v * n_head_v, n_tokens, qkv->nb[1], + ggml_row_size(qkv->type, n_embd_q + n_embd_k)); + } + if (!reshape) { + Qcur = ggml_cont(ctx0, Qcur); + Kcur = ggml_cont(ctx0, Kcur); + Vcur = ggml_cont(ctx0, Vcur); + } } else { // separate Q/K/V path Qcur = build_lora_mm(layer.wq, cur, layer.wq_s); - cb(Qcur, "Qcur", il); + if (reshape) { + cb(Qcur, "Qcur", il); + } if (layer.wq_b) { Qcur = ggml_add(ctx0, Qcur, layer.wq_b); - cb(Qcur, "Qcur", il); + if (reshape) { + cb(Qcur, "Qcur", il); + } } - if (hparams.f_clamp_kqv > 0.0f) { + if (reshape && hparams.f_clamp_kqv > 0.0f) { Qcur = ggml_clamp(ctx0, Qcur, -hparams.f_clamp_kqv, hparams.f_clamp_kqv); cb(Qcur, "Qcur_clamped", il); } Kcur = build_lora_mm(layer.wk, cur, layer.wk_s); - cb(Kcur, "Kcur", il); + if (reshape) { + cb(Kcur, "Kcur", il); + } if (layer.wk_b) { Kcur = ggml_add(ctx0, Kcur, layer.wk_b); - cb(Kcur, "Kcur", il); + if (reshape) { + cb(Kcur, "Kcur", il); + } } - if (hparams.f_clamp_kqv > 0.0f) { + if (reshape && hparams.f_clamp_kqv > 0.0f) { Kcur = ggml_clamp(ctx0, Kcur, -hparams.f_clamp_kqv, hparams.f_clamp_kqv); cb(Kcur, "Kcur_clamped", il); } Vcur = build_lora_mm(layer.wv, cur, layer.wv_s); - cb(Vcur, "Vcur", il); + if (reshape) { + cb(Vcur, "Vcur", il); + } if (layer.wv_b) { Vcur = ggml_add(ctx0, Vcur, layer.wv_b); - cb(Vcur, "Vcur", il); + if (reshape) { + cb(Vcur, "Vcur", il); + } } - if (hparams.f_clamp_kqv > 0.0f) { + if (reshape && hparams.f_clamp_kqv > 0.0f) { Vcur = ggml_clamp(ctx0, Vcur, -hparams.f_clamp_kqv, hparams.f_clamp_kqv); cb(Vcur, "Vcur_clamped", il); } - Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens); - Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); - Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); + if (reshape) { + Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head_q, n_head_q, n_tokens); + Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head_k, n_head_k, n_tokens); + Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head_v, n_head_v, n_tokens); + } } - cb(Qcur, "Qcur", il); - cb(Kcur, "Kcur", il); - cb(Vcur, "Vcur", il); + if (reshape) { + cb(Qcur, "Qcur", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); + } return { Qcur, Kcur, Vcur }; } diff --git a/src/llama-graph.h b/src/llama-graph.h index dddfdac7b51e..b486578c1338 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -1079,6 +1079,19 @@ struct llm_graph_context { int64_t n_head_kv, int il) const; + // Set reshape to false to return contiguous projections before clamp/reshape. + llm_graph_qkv build_qkv( + const llama_layer & layer, + ggml_tensor * cur, + int64_t n_embd_head_q, + int64_t n_head_q, + int64_t n_embd_head_k, + int64_t n_head_k, + int64_t n_embd_head_v, + int64_t n_head_v, + int il, + bool reshape = true) const; + ggml_tensor * build_ffn( ggml_tensor * cur, ggml_tensor * up, diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 6344f2d8aee4..6d6da3865cf0 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -3189,6 +3189,12 @@ void llama_model_base::create_tensor_qkv(llama_layer & layer, int bid, layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", bid), {n_embd_, n_embd_qkv}, TENSOR_NOT_REQUIRED | TENSOR_SKIP_IF_VIRTUAL); if (layer.wqkv) { layer.wqkv_b = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "bias", bid), {n_embd_qkv}, TENSOR_NOT_REQUIRED | TENSOR_SKIP_IF_VIRTUAL); + // Fused weights may coexist with separate Q/K/V biases in legacy or custom GGUFs. + if (!layer.wqkv_b) { + layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q, "bias", bid), {n_embd_q_}, TENSOR_NOT_REQUIRED); + layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K, "bias", bid), {n_embd_k_}, TENSOR_NOT_REQUIRED); + layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V, "bias", bid), {n_embd_v_}, TENSOR_NOT_REQUIRED); + } } else { layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", bid), {n_embd_, n_embd_q_}, flags); layer.wk = create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", bid), {n_embd_, n_embd_k_}, flags); diff --git a/src/models/deepseek2.cpp b/src/models/deepseek2.cpp index 4628ff4daf3b..deca86527978 100644 --- a/src/models/deepseek2.cpp +++ b/src/models/deepseek2.cpp @@ -475,21 +475,12 @@ llama_model_deepseek2::graph::graph(const llama_model & model, const llm_graph_p const int ocr_rope_type = GGML_ROPE_TYPE_NEOX; GGML_ASSERT(n_embed_head == n_embd_head_k && n_embed_head == n_embd_head_v); - ggml_tensor * Qcur = NULL; - ggml_tensor * Kcur = NULL; - ggml_tensor * Vcur = NULL; - - Qcur = ggml_mul_mat(ctx0, model.layers[il].wq, cur); - Kcur = ggml_mul_mat(ctx0, model.layers[il].wk, cur); - Vcur = ggml_mul_mat(ctx0, model.layers[il].wv, cur); + auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embed_head, n_head, n_head, il); cb(Qcur, "q", il); cb(Kcur, "k", il); cb(Vcur, "v", il); - Qcur = ggml_reshape_3d(ctx0, Qcur, n_embed_head, n_head, n_tokens); - Kcur = ggml_reshape_3d(ctx0, Kcur, n_embed_head, n_head, n_tokens); - Vcur = ggml_reshape_3d(ctx0, Vcur, n_embed_head, n_head, n_tokens); - GGML_ASSERT(fabs(freq_base - 10000.0) < 1e-4); Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, n_embed_head, ocr_rope_type, 0, freq_base, 1, 0, 1, 0, 0); Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, n_embed_head, ocr_rope_type, 0, freq_base, 1, 0, 1, 0, 0); diff --git a/src/models/deepseek2ocr.cpp b/src/models/deepseek2ocr.cpp index 1c5c452e96d4..3d630699ef2d 100644 --- a/src/models/deepseek2ocr.cpp +++ b/src/models/deepseek2ocr.cpp @@ -40,9 +40,7 @@ void llama_model_deepseek2ocr::load_arch_tensors(llama_model_loader &) { for (int i = 0; i < n_layer; ++i) { auto & layer = layers[i]; - layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_embd}, 0); - layer.wk = create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, n_embd}, 0); - layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, n_embd}, 0); + create_tensor_qkv(layer, i, n_embd, n_embd, n_embd, n_embd, 0); layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd, n_embd}, 0); // norm diff --git a/src/models/gemma3n.cpp b/src/models/gemma3n.cpp index 83eb8250aa94..ea616db3ba3c 100644 --- a/src/models/gemma3n.cpp +++ b/src/models/gemma3n.cpp @@ -176,7 +176,14 @@ llama_model_gemma3n::graph::graph(const llama_model & model, const llm_graph_par hparams.f_attention_scale, il); } else { // reuse KV cache of earlier layers - ggml_tensor * Qcur = build_lora_mm(model.layers[il].wq, cur); + ggml_tensor * Qcur; + if (model.layers[il].wqkv) { + ggml_tensor * qkv = build_lora_mm(model.layers[il].wqkv, cur); + const int64_t q_dim = n_embd_head * n_head; + Qcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, q_dim, n_tokens, qkv->nb[1], 0)); + } else { + Qcur = build_lora_mm(model.layers[il].wq, cur); + } cb(Qcur, "Qcur", il); Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens); diff --git a/src/models/gemma4.cpp b/src/models/gemma4.cpp index c6dd7d1bf9c8..e68420098a03 100644 --- a/src/models/gemma4.cpp +++ b/src/models/gemma4.cpp @@ -70,9 +70,13 @@ void llama_model_gemma4::load_arch_tensors(llama_model_loader &) { layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); // note: use_alternative_attention (v_proj is optional, if it's not present, use k_proj) - layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_embd_head * n_head}, 0); - layer.wk = create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, n_embd_k}, kv_flags); - layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, n_embd_v}, TENSOR_NOT_REQUIRED); + layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", i), + {n_embd, n_embd_head * n_head + n_embd_k + n_embd_v}, TENSOR_NOT_REQUIRED | TENSOR_SKIP_IF_VIRTUAL); + if (!layer.wqkv) { + layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_embd_head * n_head}, 0); + layer.wk = create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, n_embd_k}, kv_flags); + layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, n_embd_v}, TENSOR_NOT_REQUIRED); + } layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head * n_head, n_embd}, 0); layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head}, 0); @@ -197,9 +201,17 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para // Q projection (shared for both non-KV and KV layers) // this is to mirror Gemma4Attention in pytorch code + ggml_tensor * qkv_fused = nullptr; ggml_tensor * Qcur; - { + if (model.layers[il].wqkv) { + qkv_fused = build_lora_mm(model.layers[il].wqkv, cur, model.layers[il].wqkv_s); + cb(qkv_fused, "wqkv", il); + const int64_t q_dim = n_embd_head * n_head; + Qcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv_fused, q_dim, n_tokens, qkv_fused->nb[1], 0)); + } else { Qcur = build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s); + } + { cb(Qcur, "Qcur", il); Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens); @@ -214,12 +226,22 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para // self-attention if (hparams.has_kv(il)) { - ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur, model.layers[il].wk_s); + ggml_tensor * Kcur; + ggml_tensor * Vcur; + if (qkv_fused) { + const int64_t q_dim = n_embd_head * n_head; + const int64_t k_dim = n_embd_head * n_head_kv; + const int64_t v_dim = n_embd_head * n_head_kv; + const size_t esize = ggml_element_size(qkv_fused); + Kcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv_fused, k_dim, n_tokens, qkv_fused->nb[1], q_dim * esize)); + Vcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv_fused, v_dim, n_tokens, qkv_fused->nb[1], (q_dim + k_dim) * esize)); + } else { + Kcur = build_lora_mm(model.layers[il].wk, cur, model.layers[il].wk_s); + Vcur = model.layers[il].wv + ? build_lora_mm(model.layers[il].wv, cur, model.layers[il].wv_s) + : Kcur; // if v_proj is not present, use Kcur as Vcur + } cb(Kcur, "Kcur", il); - - ggml_tensor * Vcur = model.layers[il].wv - ? build_lora_mm(model.layers[il].wv, cur, model.layers[il].wv_s) - : Kcur; // if v_proj is not present, use Kcur as Vcur cb(Vcur, "Vcur", il); Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); diff --git a/src/models/jais2.cpp b/src/models/jais2.cpp index 8610fcc9f82f..64813b7b6b23 100644 --- a/src/models/jais2.cpp +++ b/src/models/jais2.cpp @@ -29,15 +29,9 @@ void llama_model_jais2::load_arch_tensors(llama_model_loader &) { layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); layer.attn_norm_b = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "bias", i), {n_embd}, 0); - layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_embd_head_k * n_head}, 0); - layer.wk = create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, n_embd_k_gqa}, 0); - layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, n_embd_v_gqa}, 0); + 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_k * n_head, n_embd}, 0); - // attention biases - all have shape n_embd (output dimension of projections) - layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q, "bias", i), {n_embd}, 0); - layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K, "bias", i), {n_embd}, 0); - layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V, "bias", i), {n_embd}, 0); layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias", i), {n_embd}, 0); layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); diff --git a/src/models/kimi-linear.cpp b/src/models/kimi-linear.cpp index 601d1d9beb8d..fb56e584487a 100644 --- a/src/models/kimi-linear.cpp +++ b/src/models/kimi-linear.cpp @@ -195,7 +195,7 @@ static ggml_tensor * causal_conv1d(ggml_cgraph * gf, ggml_context * ctx0, ggml_t // Causal Conv1d function for Q,K,V // When qkv is 0, it is Q, 1 is K, 2 is V // Step 1: Q, K, V projections -> [d_inner, n_tokens] - ggml_tensor * x_proj = ggml_mul_mat(ctx0, proj_w, x); + ggml_tensor * x_proj = proj_w ? ggml_mul_mat(ctx0, proj_w, x) : x; // Reshape input: {d_inner, n_tokens} -> {d_inner, n_seq_tokens, n_seqs} ggml_tensor * x_3d = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs); @@ -295,9 +295,20 @@ llama_model_kimi_linear::graph::graph(const llama_model & model, const llm_graph ggml_tensor * conv_states_all = mctx_cur->get_r_l(il); cb(conv_states_all, "conv_states_all", il); ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs); - ggml_tensor * Qcur = causal_conv1d(gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, kv_head); - ggml_tensor * Kcur = causal_conv1d(gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, kv_head); - ggml_tensor * Vcur = causal_conv1d(gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, kv_head); + ggml_tensor * q_in = cur, * k_in = cur, * v_in = cur; + ggml_tensor * q_w = layer.wq, * k_w = layer.wk, * v_w = layer.wv; + if (layer.wqkv) { + ggml_tensor * qkv = ggml_mul_mat(ctx0, layer.wqkv, cur); + const int64_t d_inner = head_dim * n_head; + const size_t esize = ggml_element_size(qkv); + q_in = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, d_inner, n_tokens, qkv->nb[1], 0)); + k_in = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, d_inner, n_tokens, qkv->nb[1], d_inner * esize)); + v_in = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, d_inner, n_tokens, qkv->nb[1], 2 * d_inner * esize)); + q_w = nullptr; k_w = nullptr; v_w = nullptr; + } + ggml_tensor * Qcur = causal_conv1d(gf, ctx0, conv_states_all, conv_state_all, 0, q_in, q_w, layer.ssm_q_conv, d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, kv_head); + ggml_tensor * Kcur = causal_conv1d(gf, ctx0, conv_states_all, conv_state_all, 1, k_in, k_w, layer.ssm_k_conv, d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, kv_head); + ggml_tensor * Vcur = causal_conv1d(gf, ctx0, conv_states_all, conv_state_all, 2, v_in, v_w, layer.ssm_v_conv, d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, kv_head); // g1 = -exp(A_log) * softplus(f_b(f_a(x)) + dt_bias) ggml_tensor * f_a = ggml_mul_mat(ctx0, layer.ssm_f_a, cur); diff --git a/src/models/llada.cpp b/src/models/llada.cpp index 87d4259f9a74..ae3d6925c136 100644 --- a/src/models/llada.cpp +++ b/src/models/llada.cpp @@ -36,12 +36,7 @@ void llama_model_llada::load_arch_tensors(llama_model_loader &) { layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), { n_embd }, 0); - // Use separate Q, K, V projections without bias, matching LLaDALlamaBlock - layer.wq = - create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), { n_embd, n_embd_head_k * n_head }, 0); - layer.wk = create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", i), { n_embd, n_embd_k_gqa }, 0); - layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", i), { n_embd, n_embd_v_gqa }, 0); - // No bias for QKV projections as per config: include_bias=false, include_qkv_bias=false + 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_k * n_head, n_embd }, 0); layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias", i), { n_embd }, TENSOR_NOT_REQUIRED); diff --git a/src/models/minimax-m2.cpp b/src/models/minimax-m2.cpp index c2e69bfaa89f..7a22af036bc7 100644 --- a/src/models/minimax-m2.cpp +++ b/src/models/minimax-m2.cpp @@ -71,14 +71,13 @@ llama_model_minimax_m2::graph::graph(const llama_model & model, const llm_graph_ 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); + auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embd_head, n_head, + n_embd_head, n_head_kv, + n_embd_head, n_head_kv, + il, false); cb(Qcur, "Qcur", il); - - ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur); cb(Kcur, "Kcur", il); - - ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur); cb(Vcur, "Vcur", il); Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, diff --git a/src/models/olmo2.cpp b/src/models/olmo2.cpp index cb52cdef7204..05b9394b8fe4 100644 --- a/src/models/olmo2.cpp +++ b/src/models/olmo2.cpp @@ -93,14 +93,13 @@ llama_model_olmo2::graph::graph(const llama_model & model, const llm_graph // self_attention { - // compute Q and K and RoPE them - ggml_tensor * Qcur = build_lora_mm(model.layers[il].wq, cur); + auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embd_head, n_head, + n_embd_head, n_head_kv, + n_embd_head, n_head_kv, + il, false); cb(Qcur, "Qcur", il); - - ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur); cb(Kcur, "Kcur", il); - - ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur); cb(Vcur, "Vcur", il); Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, diff --git a/src/models/olmoe.cpp b/src/models/olmoe.cpp index 1e2baeb207ff..11c53f3f4c9c 100644 --- a/src/models/olmoe.cpp +++ b/src/models/olmoe.cpp @@ -79,14 +79,13 @@ llama_model_olmoe::graph::graph(const llama_model & model, const llm_graph_param // self_attention { - // compute Q and K and RoPE them - ggml_tensor * Qcur = build_lora_mm(model.layers[il].wq, cur); + auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embd_head, n_head, + n_embd_head, n_head_kv, + n_embd_head, n_head_kv, + il, false); cb(Qcur, "Qcur", il); - - ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur); cb(Kcur, "Kcur", il); - - ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur); cb(Vcur, "Vcur", il); Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index 0b9210981d3b..ff9ddc840382 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -263,8 +263,14 @@ ggml_tensor * llama_model_qwen35::graph::build_layer_attn( // Order: joint QG projection, QG split, Q norm, KV projection, K norm, RoPE, attention // Qwen3Next uses a single Q projection that outputs query + gate - ggml_tensor * Qcur_full = build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s); // [ (n_embd_head * 2) * n_head, n_tokens ] + auto [Qcur_full, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embd_head * 2, n_head, + n_embd_head, n_head_kv, + n_embd_head, n_head_kv, + il, false); cb(Qcur_full, "Qcur_full", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); ggml_tensor * Qcur = ggml_view_3d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens, ggml_element_size(Qcur_full) * n_embd_head * 2, @@ -275,12 +281,6 @@ ggml_tensor * llama_model_qwen35::graph::build_layer_attn( Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, nullptr, LLM_NORM_RMS, il); cb(Qcur, "Qcur_normed", il); - ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur, model.layers[il].wk_s); - cb(Kcur, "Kcur", il); - - ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur, model.layers[il].wv_s); - cb(Vcur, "Vcur", il); - // Apply K normalization Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, il); @@ -553,7 +553,11 @@ llama_model_qwen35::graph_mtp::graph_mtp(const llama_model & model, const llm_gr cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); cb(cur, "mtp_attn_norm", il); - ggml_tensor * Qcur_full = build_lora_mm(layer.wq, cur, layer.wq_s); + auto [Qcur_full, Kcur, Vcur] = build_qkv(layer, cur, + n_embd_head * 2, n_head, + n_embd_head, n_head_kv, + n_embd_head, n_head_kv, + il, false); cb(Qcur_full, "mtp_Qcur_full", il); ggml_tensor * Qcur = ggml_view_3d(ctx0, Qcur_full, @@ -572,12 +576,10 @@ llama_model_qwen35::graph_mtp::graph_mtp(const llama_model & model, const llm_gr gate = ggml_cont_2d(ctx0, gate, n_embd_head * n_head, n_tokens); cb(gate, "mtp_gate", il); - ggml_tensor * Kcur = build_lora_mm(layer.wk, cur, layer.wk_s); Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); Kcur = build_norm(Kcur, layer.attn_k_norm, nullptr, LLM_NORM_RMS, il); cb(Kcur, "mtp_Kcur_normed", il); - ggml_tensor * Vcur = build_lora_mm(layer.wv, cur, layer.wv_s); Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); cb(Vcur, "mtp_Vcur", il); diff --git a/src/models/qwen35moe.cpp b/src/models/qwen35moe.cpp index ed4083f12b72..7be8214d4168 100644 --- a/src/models/qwen35moe.cpp +++ b/src/models/qwen35moe.cpp @@ -287,8 +287,14 @@ ggml_tensor * llama_model_qwen35moe::graph::build_layer_attn( // Order: joint QG projection, QG split, Q norm, KV projection, K norm, RoPE, attention // Qwen3Next uses a single Q projection that outputs query + gate - ggml_tensor * Qcur_full = build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s); // [ (n_embd_head * 2) * n_head, n_tokens ] + auto [Qcur_full, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embd_head * 2, n_head, + n_embd_head, n_head_kv, + n_embd_head, n_head_kv, + il, false); cb(Qcur_full, "Qcur_full", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); ggml_tensor * Qcur = ggml_view_3d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens, ggml_element_size(Qcur_full) * n_embd_head * 2, @@ -299,12 +305,6 @@ ggml_tensor * llama_model_qwen35moe::graph::build_layer_attn( Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, nullptr, LLM_NORM_RMS, il); cb(Qcur, "Qcur_normed", il); - ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur, model.layers[il].wk_s); - cb(Kcur, "Kcur", il); - - ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur, model.layers[il].wv_s); - cb(Vcur, "Vcur", il); - // Apply K normalization Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, il); @@ -617,7 +617,11 @@ llama_model_qwen35moe::graph_mtp::graph_mtp(const llama_model & model, const llm cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); cb(cur, "mtp_attn_norm", il); - ggml_tensor * Qcur_full = build_lora_mm(layer.wq, cur, layer.wq_s); + auto [Qcur_full, Kcur, Vcur] = build_qkv(layer, cur, + n_embd_head * 2, n_head, + n_embd_head, n_head_kv, + n_embd_head, n_head_kv, + il, false); cb(Qcur_full, "mtp_Qcur_full", il); ggml_tensor * Qcur = ggml_view_3d(ctx0, Qcur_full, @@ -636,12 +640,10 @@ llama_model_qwen35moe::graph_mtp::graph_mtp(const llama_model & model, const llm gate = ggml_cont_2d(ctx0, gate, n_embd_head * n_head, n_tokens); cb(gate, "mtp_gate", il); - ggml_tensor * Kcur = build_lora_mm(layer.wk, cur, layer.wk_s); Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); Kcur = build_norm(Kcur, layer.attn_k_norm, nullptr, LLM_NORM_RMS, il); cb(Kcur, "mtp_Kcur_normed", il); - ggml_tensor * Vcur = build_lora_mm(layer.wv, cur, layer.wv_s); Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); cb(Vcur, "mtp_Vcur", il); diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp index eb823b8eadda..7d5806b9547d 100644 --- a/src/models/qwen3next.cpp +++ b/src/models/qwen3next.cpp @@ -244,8 +244,14 @@ ggml_tensor * llama_model_qwen3next::graph::build_layer_attn( // Order: joint QG projection, QG split, Q norm, KV projection, K norm, RoPE, attention // Qwen3Next uses a single Q projection that outputs query + gate - ggml_tensor * Qcur_full = build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s); + auto [Qcur_full, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embd_head * 2, n_head, + n_embd_head, n_head_kv, + n_embd_head, n_head_kv, + il, false); cb(Qcur_full, "Qcur_full", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); Qcur_full = ggml_reshape_4d(ctx0, Qcur_full, n_embd_head * 2, n_head, n_tokens, 1); @@ -260,12 +266,6 @@ ggml_tensor * llama_model_qwen3next::graph::build_layer_attn( Qcur_full->nb[1], Qcur_full->nb[2], Qcur_full->nb[3], n_embd_head * ggml_element_size(Qcur_full)); cb(gate, "gate", il); - ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur, model.layers[il].wk_s); - cb(Kcur, "Kcur", il); - - ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur, model.layers[il].wv_s); - cb(Vcur, "Vcur", il); - Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); @@ -691,7 +691,11 @@ llama_model_qwen3next::graph_mtp::graph_mtp(const llama_model & model, const llm cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); cb(cur, "mtp_attn_norm", il); - ggml_tensor * Qcur_full = build_lora_mm(layer.wq, cur, layer.wq_s); + auto [Qcur_full, Kcur, Vcur] = build_qkv(layer, cur, + n_embd_head * 2, n_head, + n_embd_head, n_head_kv, + n_embd_head, n_head_kv, + il, false); cb(Qcur_full, "mtp_Qcur_full", il); ggml_tensor * Qcur = ggml_view_3d(ctx0, Qcur_full, @@ -702,12 +706,10 @@ llama_model_qwen3next::graph_mtp::graph_mtp(const llama_model & model, const llm Qcur = build_norm(Qcur, layer.attn_q_norm, nullptr, LLM_NORM_RMS, il); cb(Qcur, "mtp_Qcur_normed", il); - ggml_tensor * Kcur = build_lora_mm(layer.wk, cur, layer.wk_s); Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); Kcur = build_norm(Kcur, layer.attn_k_norm, nullptr, LLM_NORM_RMS, il); cb(Kcur, "mtp_Kcur_normed", il); - ggml_tensor * Vcur = build_lora_mm(layer.wv, cur, layer.wv_s); Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, diff --git a/src/models/step35.cpp b/src/models/step35.cpp index 53f3179c6357..946a3696000f 100644 --- a/src/models/step35.cpp +++ b/src/models/step35.cpp @@ -216,9 +216,11 @@ llama_model_step35::graph::graph(const llama_model & model, const llm_graph_para { cur = build_norm(cur, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il); cb(cur, "attn_norm", il); - ggml_tensor * Qcur = build_lora_mm(model.layers[il].wq, cur); - ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur); - ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur); + auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embd_head_k, n_head_l, + n_embd_head_k, n_head_kv_l, + n_embd_head_v, n_head_kv_l, + il, false); cb(Qcur, "Qcur", il); cb(Kcur, "Kcur", il); @@ -425,9 +427,11 @@ llama_model_step35::graph_mtp::graph_mtp(const llama_model & model, const llm_gr cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); cb(cur, "mtp_attn_norm", il); - ggml_tensor * Qcur = build_lora_mm(layer.wq, cur, layer.wq_s); - ggml_tensor * Kcur = build_lora_mm(layer.wk, cur, layer.wk_s); - ggml_tensor * Vcur = build_lora_mm(layer.wv, cur, layer.wv_s); + auto [Qcur, Kcur, Vcur] = build_qkv(layer, cur, + n_embd_head_k, n_head_l, + n_embd_head_k, n_head_kv_l, + n_embd_head_v, n_head_kv_l, + il, false); cb(Qcur, "mtp_Qcur", il); cb(Kcur, "mtp_Kcur", il); cb(Vcur, "mtp_Vcur", il);