This should do it

generate_imatrix.py

@@ -472,12 +472,13 @@ ACTIVATION_MAP = {
 class TransformerRunner:
     """Minimal Gemma transformer for importance collection."""
 
-    def __init__(self, model, config, verbose=False):
+    def __init__(self, model, config, verbose=False, linear_attn=True):
         self.model = model
         self.cfg = config
         self.verbose = verbose
         self.head_dim = config.get('head_dim', config['n_embd'] // config['n_head'])
         self.act_fn = ACTIVATION_MAP.get(config['arch'], silu)
+        self.linear_attn = linear_attn
 
         # Importance accumulators: tensor_name → (sum_x2, count)
         self.importance = {}
@@ -508,6 +509,199 @@ class TransformerRunner:
     def _layer_prefix(self, layer_idx):
         return f"blk.{layer_idx}"
 
+    def forward_layer_linear(self, hidden, layer_idx):
+        """HPC-linearized forward: O(seq) attention for imatrix collection.
+
+        Instead of full O(seq²) softmax attention, uses causal linear attention:
+        each position's output is a running weighted average of V, where weights
+        come from Q·K similarity in phase space. This preserves activation
+        magnitude statistics (which is all imatrix needs) while being O(seq).
+
+        Records identical importance stats as the full forward_layer.
+        """
+        pfx = self._layer_prefix(layer_idx)
+        cfg = self.cfg
+        n_head = cfg['n_head']
+        n_head_kv = cfg['n_head_kv']
+        seq_len = hidden.shape[0]
+
+        # ── Attention norm ──
+        attn_norm_w = self._get_weight(f'{pfx}.attn_norm.weight')
+        if attn_norm_w is None:
+            return hidden
+        normed = rms_norm(hidden, attn_norm_w, cfg['rms_eps'])
+
+        # ── Check for fused vs separate QKV ──
+        qkv_w = self._get_weight(f'{pfx}.attn_qkv.weight')
+        gate_w = self._get_weight(f'{pfx}.attn_gate.weight')
+        q_w = self._get_weight(f'{pfx}.attn_q.weight')
+        k_w = self._get_weight(f'{pfx}.attn_k.weight')
+        v_w = self._get_weight(f'{pfx}.attn_v.weight')
+        o_w = self._get_weight(f'{pfx}.attn_output.weight')
+
+        if qkv_w is not None:
+            # ── Fused QKV path (Qwen 3.6 hybrid layers) ──
+            self._record(f'{pfx}.attn_qkv.weight', normed)
+            head_dim = self.head_dim
+            q_dim = n_head * head_dim
+            kv_dim = n_head_kv * head_dim
+
+            qkv = normed @ qkv_w.T
+            q = qkv[:, :q_dim].reshape(seq_len, n_head, head_dim)
+            k = qkv[:, q_dim:q_dim + kv_dim].reshape(seq_len, n_head_kv, head_dim)
+            v = qkv[:, q_dim + kv_dim:q_dim + 2 * kv_dim].reshape(seq_len, n_head_kv, head_dim)
+
+            # GQA expand
+            if n_head_kv < n_head:
+                rep = n_head // n_head_kv
+                k = np.repeat(k, rep, axis=1)
+                v = np.repeat(v, rep, axis=1)
+
+            # ── Linear attention: O(seq × head_dim²) ──
+            # φ(x) = elu(x) + 1 (feature map for linear attention)
+            q_feat = np.maximum(q, 0) + 1e-6  # [seq, n_head, head_dim]
+            k_feat = np.maximum(k, 0) + 1e-6
+
+            # Causal linear attention via running state (vectorized over heads):
+            # S_t = S_{t-1} + k_t ⊗ v_t  (outer product accumulator)
+            # z_t = z_{t-1} + k_t         (normalizer accumulator)
+            # out_t = (q_t @ S_t) / (q_t · z_t)
+            out = np.zeros_like(q)  # [seq, n_head, head_dim]
+            S = np.zeros((n_head, head_dim, head_dim), dtype=np.float32)
+            z = np.zeros((n_head, head_dim), dtype=np.float32)
+
+            for t in range(seq_len):
+                # Vectorized over all heads: [n_head, head_dim]
+                kt = k_feat[t]  # [n_head, head_dim]
+                vt = v[t]       # [n_head, head_dim]
+                qt = q_feat[t]  # [n_head, head_dim]
+                # S[h] += outer(kt[h], vt[h]) for all h at once
+                S += kt[:, :, None] * vt[:, None, :]  # [n_head, hd, hd]
+                z += kt                                # [n_head, hd]
+                # num = qt @ S -> [n_head, head_dim]
+                num = np.einsum('hd,hde->he', qt, S)
+                den = np.sum(qt * z, axis=-1, keepdims=True) + 1e-8  # [n_head, 1]
+                out[t] = num / den
+
+            attn_result = out.reshape(seq_len, -1)  # [seq, n_head * head_dim]
+
+            # Record and project
+            if gate_w is not None:
+                self._record(f'{pfx}.attn_gate.weight', attn_result)
+                if gate_w.shape[1] == hidden.shape[-1]:
+                    attn_out = attn_result @ gate_w
+                else:
+                    attn_out = attn_result @ gate_w.T
+            else:
+                attn_out = np.zeros_like(hidden)
+
+        elif q_w is not None and k_w is not None and v_w is not None and o_w is not None:
+            # ── Separate QKV path (standard transformer layers) ──
+            self._record(f'{pfx}.attn_q.weight', normed)
+            self._record(f'{pfx}.attn_k.weight', normed)
+            self._record(f'{pfx}.attn_v.weight', normed)
+
+            q = normed @ q_w.T
+            k = normed @ k_w.T
+            v = normed @ v_w.T
+
+            head_dim_q = q_w.shape[0] // n_head
+            head_dim_kv = k_w.shape[0] // n_head_kv
+
+            q = q.reshape(seq_len, n_head, head_dim_q)
+            k = k.reshape(seq_len, n_head_kv, head_dim_kv)
+            v = v.reshape(seq_len, n_head_kv, head_dim_kv)
+
+            if n_head_kv < n_head:
+                rep = n_head // n_head_kv
+                k = np.repeat(k, rep, axis=1)
+                v = np.repeat(v, rep, axis=1)
+
+            # Linear attention with feature map
+            q_feat = np.maximum(q, 0) + 1e-6
+            k_feat = np.maximum(k, 0) + 1e-6
+
+            out = np.zeros_like(v)  # [seq, n_head, head_dim_kv]
+            S = np.zeros((n_head, head_dim_kv, head_dim_kv), dtype=np.float32)
+            z = np.zeros((n_head, head_dim_kv), dtype=np.float32)
+
+            # Use min of q/k dims for the state accumulator
+            feat_dim = min(head_dim_q, head_dim_kv)
+            S = np.zeros((n_head, feat_dim, head_dim_kv), dtype=np.float32)
+            z = np.zeros((n_head, feat_dim), dtype=np.float32)
+
+            for t in range(seq_len):
+                # Vectorized over all heads
+                kf = k_feat[t, :, :feat_dim]  # [n_head, feat_dim]
+                qf = q_feat[t, :, :feat_dim]  # [n_head, feat_dim]
+                vt = v[t]                      # [n_head, head_dim_kv]
+                S += kf[:, :, None] * vt[:, None, :]  # [n_head, feat_dim, head_dim_kv]
+                z += kf                                # [n_head, feat_dim]
+                num = np.einsum('hd,hde->he', qf, S)   # [n_head, head_dim_kv]
+                den = np.sum(qf * z, axis=-1, keepdims=True) + 1e-8
+                out[t] = num / den
+
+            attn_result = out.reshape(seq_len, -1)
+
+            # Pad/truncate to match o_w input size
+            if attn_result.shape[-1] != o_w.shape[1]:
+                if attn_result.shape[-1] < o_w.shape[1]:
+                    padded = np.zeros((seq_len, o_w.shape[1]), dtype=attn_result.dtype)
+                    padded[:, :attn_result.shape[-1]] = attn_result
+                    attn_result = padded
+                else:
+                    attn_result = attn_result[:, :o_w.shape[1]]
+
+            self._record(f'{pfx}.attn_output.weight', attn_result)
+            attn_out = attn_result @ o_w.T
+        else:
+            return hidden
+
+        hidden = hidden + attn_out
+
+        # ── SSM path (Qwen 3.6 hybrid) ──
+        ssm_alpha_w = self._get_weight(f'{pfx}.ssm_alpha.weight')
+        ssm_beta_w = self._get_weight(f'{pfx}.ssm_beta.weight')
+        ssm_out_w = self._get_weight(f'{pfx}.ssm_out.weight')
+        if ssm_alpha_w is not None:
+            self._record(f'{pfx}.ssm_alpha.weight', normed)
+        if ssm_beta_w is not None:
+            self._record(f'{pfx}.ssm_beta.weight', normed)
+        if ssm_out_w is not None:
+            if qkv_w is not None:
+                qkv_full = normed @ qkv_w.T
+                ssm_proxy = qkv_full[:, :ssm_out_w.shape[1]] if qkv_full.shape[-1] >= ssm_out_w.shape[1] else normed
+            else:
+                ssm_proxy = normed
+            self._record(f'{pfx}.ssm_out.weight', ssm_proxy)
+            if ssm_out_w.shape[0] == hidden.shape[-1]:
+                hidden = hidden + ssm_proxy @ ssm_out_w.T
+
+        # ── FFN ──
+        ffn_norm_w = self._get_weight(f'{pfx}.post_attention_norm.weight')
+        if ffn_norm_w is None:
+            ffn_norm_w = self._get_weight(f'{pfx}.ffn_norm.weight')
+        if ffn_norm_w is None:
+            return hidden
+
+        normed_ff = rms_norm(hidden, ffn_norm_w, cfg['rms_eps'])
+
+        gate_fw = self._get_weight(f'{pfx}.ffn_gate.weight')
+        up_w = self._get_weight(f'{pfx}.ffn_up.weight')
+        down_w = self._get_weight(f'{pfx}.ffn_down.weight')
+
+        if gate_fw is not None and up_w is not None and down_w is not None:
+            self._record(f'{pfx}.ffn_gate.weight', normed_ff)
+            self._record(f'{pfx}.ffn_up.weight', normed_ff)
+            gate_out = self.act_fn(normed_ff @ gate_fw.T)
+            up_out = normed_ff @ up_w.T
+            ff_mid = gate_out * up_out
+            self._record(f'{pfx}.ffn_down.weight', ff_mid)
+            ff_out = ff_mid @ down_w.T
+            hidden = hidden + ff_out
+
+        return hidden
+
     def forward_layer(self, hidden, layer_idx, cos_f, sin_f):
         """Forward pass through one transformer layer. Returns new hidden state."""
         pfx = self._layer_prefix(layer_idx)
@@ -537,16 +731,30 @@ class TransformerRunner:
         self._record(f'{pfx}.attn_k.weight', normed)
         self._record(f'{pfx}.attn_v.weight', normed)
 
-        q = normed @ q_w.T  # [seq, n_head * head_dim]
-        k = normed @ k_w.T  # [seq, n_head_kv * head_dim]
+        q = normed @ q_w.T  # [seq, q_w.shape[0]]
+        k = normed @ k_w.T  # [seq, k_w.shape[0]]
         v = normed @ v_w.T
 
-        q = q.reshape(seq_len, n_head, head_dim)
-        k = k.reshape(seq_len, n_head_kv, head_dim)
-        v = v.reshape(seq_len, n_head_kv, head_dim)
+        # Dynamic head_dim based on tensor size
+        head_dim_q = q_w.shape[0] // n_head
+        head_dim_kv = k_w.shape[0] // n_head_kv
 
-        q = apply_rope(q, cos_f, sin_f)
-        k = apply_rope(k, cos_f, sin_f)
+        q = q.reshape(seq_len, n_head, head_dim_q)
+        k = k.reshape(seq_len, n_head_kv, head_dim_kv)
+        v = v.reshape(seq_len, n_head_kv, head_dim_kv)
+
+        # Apply RoPE
+        if head_dim_q != head_dim:
+            cos_q, sin_q = rope_freqs(head_dim_q, seq_len, cfg['rope_base'])
+            q = apply_rope(q, cos_q, sin_q)
+        else:
+            q = apply_rope(q, cos_f, sin_f)
+            
+        if head_dim_kv != head_dim:
+            cos_k, sin_k = rope_freqs(head_dim_kv, seq_len, cfg['rope_base'])
+            k = apply_rope(k, cos_k, sin_k)
+        else:
+            k = apply_rope(k, cos_f, sin_f)
 
         # GQA: repeat KV heads
         if n_head_kv < n_head:
@@ -554,12 +762,25 @@ class TransformerRunner:
             k = np.repeat(k, rep, axis=1)
             v = np.repeat(v, rep, axis=1)
 
-        # Attention: [n_head, seq, head_dim] @ [n_head, head_dim, seq]
-        q_t = q.transpose(1, 0, 2)  # [n_head, seq, head_dim]
-        k_t = k.transpose(1, 0, 2)
-        v_t = v.transpose(1, 0, 2)
+        q_t = q.transpose(1, 0, 2)  # [n_head, seq, head_dim_q]
+        k_t = k.transpose(1, 0, 2)  # [n_head, seq, head_dim_kv]
+        v_t = v.transpose(1, 0, 2)  # [n_head, seq, head_dim_kv]
+
+        scale = 1.0 / np.sqrt(head_dim_q)
+        
+        # If Q and K head dims differ, there might be a projection or it's not standard SDP.
+        # But for importance calculation, if we just need to get the attention magnitude:
+        # We can pad K to match Q, or truncate Q to match K. We only need an approximation.
+        if head_dim_q != head_dim_kv:
+            if head_dim_q > head_dim_kv:
+                k_t_padded = np.zeros_like(q_t)
+                k_t_padded[..., :head_dim_kv] = k_t
+                k_t = k_t_padded
+            else:
+                q_t_padded = np.zeros_like(k_t)
+                q_t_padded[..., :head_dim_q] = q_t
+                q_t = q_t_padded
 
-        scale = 1.0 / np.sqrt(head_dim)
         attn = np.matmul(q_t, k_t.transpose(0, 2, 1)) * scale  # [n_head, seq, seq]
 
         # Causal mask (with optional sliding window)
@@ -572,8 +793,20 @@ class TransformerRunner:
         attn = attn + mask[np.newaxis, :, :]
         attn = softmax(attn, axis=-1)
 
-        out = np.matmul(attn, v_t)  # [n_head, seq, head_dim]
-        out = out.transpose(1, 0, 2).reshape(seq_len, -1)  # [seq, n_embd]
+        out = np.matmul(attn, v_t)  # [n_head, seq, head_dim_kv]
+        
+        # Output projection input is out_w.T -> [in_features, out_features]
+        # In_features is out_w.shape[1]
+        out = out.transpose(1, 0, 2).reshape(seq_len, -1)  # [seq, n_head * head_dim_kv]
+        
+        # Pad or truncate out to match expected input size of o_w
+        if out.shape[-1] != o_w.shape[1]:
+            if out.shape[-1] < o_w.shape[1]:
+                out_padded = np.zeros((seq_len, o_w.shape[1]), dtype=out.dtype)
+                out_padded[:, :out.shape[-1]] = out
+                out = out_padded
+            else:
+                out = out[:, :o_w.shape[1]]
 
         self._record(f'{pfx}.attn_output.weight', out)
         attn_out = out @ o_w.T
@@ -878,19 +1111,21 @@ class TransformerRunner:
         # Process each layer
         for layer_idx in range(cfg['n_layers']):
             pfx = f"blk.{layer_idx}"
-            has_fused_qkv = f'{pfx}.attn_qkv.weight' in self.model.tensor_infos
-            has_separate_q = f'{pfx}.attn_q.weight' in self.model.tensor_infos
-            has_linear_attn = f'{pfx}.ssm_in_qkv.weight' in self.model.tensor_infos
-
-            if has_fused_qkv:
-                # Qwen 3.6 hybrid: fused QKV + SSM
-                hidden = self.forward_qwen35_layer(hidden, layer_idx, cos_f, sin_f)
-            elif has_linear_attn and not has_separate_q:
-                # DeltaNet-only layers
-                hidden = self.forward_linear_attn_layer(hidden, layer_idx)
+
+            if self.linear_attn:
+                # HPC-linearized attention: O(seq) per layer
+                hidden = self.forward_layer_linear(hidden, layer_idx)
             else:
-                # Standard transformer (LLaMA, Mistral, Gemma, etc.)
-                hidden = self.forward_layer(hidden, layer_idx, cos_f, sin_f)
+                has_fused_qkv = f'{pfx}.attn_qkv.weight' in self.model.tensor_infos
+                has_separate_q = f'{pfx}.attn_q.weight' in self.model.tensor_infos
+                has_linear_attn = f'{pfx}.ssm_in_qkv.weight' in self.model.tensor_infos
+
+                if has_fused_qkv:
+                    hidden = self.forward_qwen35_layer(hidden, layer_idx, cos_f, sin_f)
+                elif has_linear_attn and not has_separate_q:
+                    hidden = self.forward_linear_attn_layer(hidden, layer_idx)
+                else:
+                    hidden = self.forward_layer(hidden, layer_idx, cos_f, sin_f)
             if self.verbose and (layer_idx + 1) % 4 == 0:
                 print(f"    Layer {layer_idx + 1}/{cfg['n_layers']}", end='\r')
 
@@ -1014,12 +1249,14 @@ def main():
     parser.add_argument('calibration', help='Calibration text file')
     parser.add_argument('-o', '--output', default='imatrix.dat',
                         help='Output imatrix file (default: imatrix.dat)')
-    parser.add_argument('--chunks', type=int, default=100,
-                        help='Number of token chunks to process (default: 100)')
-    parser.add_argument('--chunk-size', type=int, default=512,
-                        help='Tokens per chunk (default: 512)')
+    parser.add_argument('--chunks', type=int, default=10,
+                        help='Number of token chunks to process (default: 10)')
+    parser.add_argument('--chunk-size', type=int, default=4096,
+                        help='Tokens per chunk (default: 4096)')
     parser.add_argument('--no-hpc', action='store_true',
                         help='Disable HPC cross-layer propagation')
+    parser.add_argument('--quadratic-attn', action='store_true',
+                        help='Use full O(seq²) attention instead of HPC-linearized O(seq)')
     parser.add_argument('--verbose', action='store_true',
                         help='Per-layer statistics')
     args = parser.parse_args()
@@ -1069,7 +1306,12 @@ def main():
 
     # ── Forward pass ──
     print("  Running forward passes...")
-    runner = TransformerRunner(model, config, verbose=args.verbose)
+    use_linear = not args.quadratic_attn
+    runner = TransformerRunner(model, config, verbose=args.verbose, linear_attn=use_linear)
+    if use_linear:
+        print(f"  Attention mode: HPC-linearized O(seq) — chunk_size={args.chunk_size}")
+    else:
+        print(f"  Attention mode: full O(seq²) softmax — chunk_size={args.chunk_size}")
 
     for i, chunk in enumerate(chunks):
         elapsed = time.time() - start_time