From 54e259b2af519baa5e9289dcdae0933d55dc6406 Mon Sep 17 00:00:00 2001 From: KakaruHayate Date: Sun, 16 Aug 2026 21:24:41 +0800 Subject: [PATCH] perf: fuse attn Q+KV into one qkv linear (llama-style) Fuse the EBF attention's q_linear and kv_linear into a single [3*H*D, D] linear at conversion time (qkv.weight/bias), matching the estimator's joint attention which already used packed QKV. One matmul instead of two per attention block (encoder 4 + segmenter 8 per D3PM pass): fewer graph nodes, one read of x per projection, one fewer backend submit per block on GPU. Numerics are exact: rows are quant-block aligned on the in-dim, so packed vs separate Q8 quantization is bit-identical. Verified byte-identical note CSV vs the pre-fusion model on CPU n8 (33 notes) and Vulkan Q8 n8 (33 notes). - ops_attn: AttentionWeights.w_qkv/b_qkv + chunk_three (removed chunk_two) - binders (encoder/segmenter): read attn.attn.qkv.{weight,bias} - converter: fuse q_linear/kv_linear pairs (order-independent); audit maps fused names; GAME_ARCH_VERSION 1 -> 2 (GGUF tensor-name schema change) --- scripts/convert_pt_to_gguf.py | 60 +++++++++++++++++++++++++++++++---- src/model_encoder.cpp | 8 ++--- src/model_segmenter.cpp | 6 ++-- src/ops_attn.cpp | 40 +++++++++++------------ src/ops_attn.h | 12 +++---- 5 files changed, 83 insertions(+), 43 deletions(-) diff --git a/scripts/convert_pt_to_gguf.py b/scripts/convert_pt_to_gguf.py index 75ad3aa..c358d87 100644 --- a/scripts/convert_pt_to_gguf.py +++ b/scripts/convert_pt_to_gguf.py @@ -55,7 +55,7 @@ # backwards-incompatible way (new required metadata, different tensor layout). # ----------------------------------------------------------------------------- GAME_ARCH = "game-me" -GAME_ARCH_VERSION = 1 +GAME_ARCH_VERSION = 2 # v2: attn q_linear/kv_linear fused into qkv (llama-style) @dataclass @@ -184,10 +184,11 @@ def _ebf_backbone_keys(prefix: str, backbone_cfg: dict, *, return_latent: bool) f"{block}.ffn2.ln2.weight", f"{block}.ffn2.ln2.bias", f"{block}.norm2.weight", }) - # PAC (self-attn + CgMLP merged) + # PAC (self-attn + CgMLP merged). q_linear/kv_linear are fused at + # conversion time into a single qkv (llama-style): one [3*H*D, D] + # linear instead of [H*D, D] + [2*H*D, D]. keys.update({ - f"{block}.attn.attn.q_linear.weight", f"{block}.attn.attn.q_linear.bias", - f"{block}.attn.attn.kv_linear.weight", f"{block}.attn.attn.kv_linear.bias", + f"{block}.attn.attn.qkv.weight", f"{block}.attn.attn.qkv.bias", f"{block}.attn.attn.out_linear.weight", f"{block}.attn.attn.out_linear.bias", f"{block}.attn.c.pw1.weight", f"{block}.attn.c.pw1.bias", f"{block}.attn.c.norm.weight", @@ -556,8 +557,21 @@ def convert(model_dir: pathlib.Path, output_path: pathlib.Path, *, strict: bool, log.info("stripped prefix '%s' from %d tensors", prefix, len(sd)) # --- audit --- + # q_linear/kv_linear are fused into qkv at write time; reflect that in the + # expected-vs-present comparison so the audit sees the emitted names. + def _fused_key(name: str) -> str: + if name.endswith("attn.attn.q_linear.weight"): + return name.replace("q_linear.weight", "qkv.weight") + if name.endswith("attn.attn.q_linear.bias"): + return name.replace("q_linear.bias", "qkv.bias") + if name.endswith("attn.attn.kv_linear.weight"): + return name.replace("kv_linear.weight", "qkv.weight") + if name.endswith("attn.attn.kv_linear.bias"): + return name.replace("kv_linear.bias", "qkv.bias") + return name + expected = _expected_keys(cfg) - present = set(sd) + present = {_fused_key(k) for k in sd} missing = sorted(expected - present) extras = sorted(present - expected) if missing: @@ -575,8 +589,42 @@ def convert(model_dir: pathlib.Path, output_path: pathlib.Path, *, strict: bool, total_params = 0 quant_report: dict[str, str] = {} + + # --- QKV fusion (llama-style): emit q_linear⊗kv_linear as one qkv tensor --- + # Rows are quant-block aligned on the in-dim and identical whether packed + # or separate, so Q8/etc. numbers are bit-identical to the un-fused form. + # Sorted state-dict order visits kv_linear before q_linear; whichever of + # the pair comes first fuses both and marks the sibling consumed. + writes: list[tuple[str, np.ndarray]] = [] + consumed: set[str] = set() for name in sorted(sd): - arr = _tensor_to_np(sd[name]) + if name in consumed: + continue + if name.endswith("attn.attn.q_linear.weight"): + kv = name.replace("q_linear.weight", "kv_linear.weight") + consumed.add(kv) + arr = np.concatenate([_tensor_to_np(sd[name]), _tensor_to_np(sd[kv])], axis=0) + writes.append((name.replace("q_linear.weight", "qkv.weight"), arr)) + elif name.endswith("attn.attn.kv_linear.weight"): + q = name.replace("kv_linear.weight", "q_linear.weight") + consumed.add(q) + arr = np.concatenate([_tensor_to_np(sd[q]), _tensor_to_np(sd[name])], axis=0) + writes.append((name.replace("kv_linear.weight", "qkv.weight"), arr)) + elif name.endswith("attn.attn.q_linear.bias"): + kv = name.replace("q_linear.bias", "kv_linear.bias") + consumed.add(kv) + arr = np.concatenate([_tensor_to_np(sd[name]), _tensor_to_np(sd[kv])], axis=0) + writes.append((name.replace("q_linear.bias", "qkv.bias"), arr)) + elif name.endswith("attn.attn.kv_linear.bias"): + q = name.replace("kv_linear.bias", "q_linear.bias") + consumed.add(q) + arr = np.concatenate([_tensor_to_np(sd[q]), _tensor_to_np(sd[name])], axis=0) + writes.append((name.replace("kv_linear.bias", "qkv.bias"), arr)) + else: + writes.append((name, _tensor_to_np(sd[name]))) + del consumed + + for name, arr in writes: qtype = resolve_tensor_type(name, arr, quant_rules or []) payload, eff = maybe_quantize(arr, qtype) # For block-quantized payloads the writer expects the *byte* shape and diff --git a/src/model_encoder.cpp b/src/model_encoder.cpp index 18a5c55..7bf943a 100644 --- a/src/model_encoder.cpp +++ b/src/model_encoder.cpp @@ -50,11 +50,9 @@ static ops::EBFBlockWeights bind_ebf_layer( P.w_a_norm = W.get(p + "attn.a_norm.weight"); P.w_c_norm = W.get(p + "attn.c_norm.weight"); - // Attention sub-weights - P.attn.w_q = W.get(p + "attn.attn.q_linear.weight"); - P.attn.b_q = W.get(p + "attn.attn.q_linear.bias"); - P.attn.w_kv = W.get(p + "attn.attn.kv_linear.weight"); - P.attn.b_kv = W.get(p + "attn.attn.kv_linear.bias"); + // Attention sub-weights (q/kv fused into one qkv tensor by the converter) + P.attn.w_qkv = W.get(p + "attn.attn.qkv.weight"); + P.attn.b_qkv = W.get(p + "attn.attn.qkv.bias"); P.attn.w_out = W.get(p + "attn.attn.out_linear.weight"); P.attn.b_out = W.get(p + "attn.attn.out_linear.bias"); diff --git a/src/model_segmenter.cpp b/src/model_segmenter.cpp index 6e748c1..a70c5f3 100644 --- a/src/model_segmenter.cpp +++ b/src/model_segmenter.cpp @@ -46,10 +46,8 @@ static ops::EBFBlockWeights bind_seg_layer( auto & P = B.pac_w; P.w_a_norm = W.get(p + "attn.a_norm.weight"); P.w_c_norm = W.get(p + "attn.c_norm.weight"); - P.attn.w_q = W.get(p + "attn.attn.q_linear.weight"); - P.attn.b_q = W.get(p + "attn.attn.q_linear.bias"); - P.attn.w_kv = W.get(p + "attn.attn.kv_linear.weight"); - P.attn.b_kv = W.get(p + "attn.attn.kv_linear.bias"); + P.attn.w_qkv = W.get(p + "attn.attn.qkv.weight"); + P.attn.b_qkv = W.get(p + "attn.attn.qkv.bias"); P.attn.w_out = W.get(p + "attn.attn.out_linear.weight"); P.attn.b_out = W.get(p + "attn.attn.out_linear.bias"); P.w_cg_pw1 = W.get(p + "attn.c.pw1.weight"); diff --git a/src/ops_attn.cpp b/src/ops_attn.cpp index df13fd3..20c0e84 100644 --- a/src/ops_attn.cpp +++ b/src/ops_attn.cpp @@ -17,23 +17,20 @@ namespace game_ggml::internal::ops { namespace { -// Split a (2*D, T, B) tensor along ne[0] into two contiguous (D, T, B) views. -// Both parts share memory with the input. -struct ChunkPair { - ggml_tensor * first; - ggml_tensor * second; -}; - -ChunkPair chunk_two(ggml_context * ctx, ggml_tensor * x) { - const int64_t half = x->ne[0] / 2; - const size_t esize = ggml_element_size(x); - ggml_tensor * a = ggml_view_4d(ctx, x, - half, x->ne[1], x->ne[2], x->ne[3], - x->nb[1], x->nb[2], x->nb[3], /*offset=*/0); - ggml_tensor * b = ggml_view_4d(ctx, x, - half, x->ne[1], x->ne[2], x->ne[3], - x->nb[1], x->nb[2], x->nb[3], /*offset=*/half * esize); - return {a, b}; +// Split a (3*D, T, B) tensor along ne[0] into q/k/v views (fused QKV). +struct QKVTriple { ggml_tensor * q; ggml_tensor * k; ggml_tensor * v; }; + +QKVTriple chunk_three(ggml_context * ctx, ggml_tensor * qkv) { + const int64_t D = qkv->ne[0] / 3; + const size_t esize = ggml_element_size(qkv); + QKVTriple r; + r.q = ggml_view_4d(ctx, qkv, D, qkv->ne[1], qkv->ne[2], qkv->ne[3], + qkv->nb[1], qkv->nb[2], qkv->nb[3], /*offset=*/0); + r.k = ggml_view_4d(ctx, qkv, D, qkv->ne[1], qkv->ne[2], qkv->ne[3], + qkv->nb[1], qkv->nb[2], qkv->nb[3], /*offset=*/D * esize); + r.v = ggml_view_4d(ctx, qkv, D, qkv->ne[1], qkv->ne[2], qkv->ne[3], + qkv->nb[1], qkv->nb[2], qkv->nb[3], /*offset=*/2 * D * esize); + return r; } } // namespace @@ -50,13 +47,12 @@ ggml_tensor * attention_with_rope( const int64_t T = x->ne[1]; const int64_t B = x->ne[2]; - // Linear projections. - ggml_tensor * q = linear(ctx, x, W.w_q, W.b_q); // (H*D, T, B) - ggml_tensor * kv = linear(ctx, x, W.w_kv, W.b_kv); // (2*H*D, T, B) - auto [k_flat, v_flat] = chunk_two(ctx, kv); + // Fused QKV projection (single [3*H*D, D] linear, then split). + ggml_tensor * qkv = linear(ctx, x, W.w_qkv, W.b_qkv); // (3*H*D, T, B) + auto [q_flat, k_flat, v_flat] = chunk_three(ctx, qkv); // Reshape to (D, H, T, B) so RoPE can be applied (ne[2] == T). - ggml_tensor * qr = ggml_reshape_4d(ctx, ggml_cont(ctx, q), head_dim, num_heads, T, B); + ggml_tensor * qr = ggml_reshape_4d(ctx, ggml_cont(ctx, q_flat), head_dim, num_heads, T, B); ggml_tensor * kr = ggml_reshape_4d(ctx, ggml_cont(ctx, k_flat), head_dim, num_heads, T, B); ggml_tensor * vr = ggml_reshape_4d(ctx, ggml_cont(ctx, v_flat), head_dim, num_heads, T, B); diff --git a/src/ops_attn.h b/src/ops_attn.h index a4b31a7..f8f1739 100644 --- a/src/ops_attn.h +++ b/src/ops_attn.h @@ -14,16 +14,16 @@ namespace game_ggml::internal::ops { // -------------------------------------------------------------------------- // AttentionWithRoPE (modules.backbones.ebf_with_joint_attention.AttentionWithRoPE) // -// q = q_linear(x) # [B, T, H*D] -// k, v = kv_linear(x).chunk(2, -1) # each [B, T, H*D] +// q, k, v = qkv_linear(x).chunk(3, -1) # each [B, T, H*D] // → reshape to (B, H, T, D) → RoPE → scaled dot-product attention → merge // out = out_linear(out) # [B, T, D_embed] +// q_linear/kv_linear are fused by the converter into a single qkv tensor +// (llama-style); the weights are stored row-aligned so quantization is +// bit-identical to the un-fused form. // -------------------------------------------------------------------------- struct AttentionWeights { - ggml_tensor * w_q = nullptr; - ggml_tensor * b_q = nullptr; - ggml_tensor * w_kv = nullptr; - ggml_tensor * b_kv = nullptr; + ggml_tensor * w_qkv = nullptr; + ggml_tensor * b_qkv = nullptr; ggml_tensor * w_out = nullptr; ggml_tensor * b_out = nullptr; };