Qwen patches

generate_imatrix.py

@@ -137,17 +137,35 @@ class GGUFModel:
 
     def get_config(self):
         arch = self.get_arch()
+        n_embd = self.kv.get(f'{arch}.embedding_length', 0)
+        n_head = self.kv.get(f'{arch}.attention.head_count', 0)
+        n_head_kv = self.kv.get(f'{arch}.attention.head_count_kv', 0)
+
+        # Auto-detect head_dim: prefer fused QKV tensor dims over n_embd/n_head
+        # (Qwen 3.6 has n_embd=5120, n_head=24 → non-integer ratio;
+        #  real head_dim=320, computed from QKV: 10240 / (24+2*4) = 320)
+        head_dim = 0
+        qkv_name = 'blk.0.attn_qkv.weight'
+        if qkv_name in self.tensor_infos:
+            qkv_cols = self.tensor_infos[qkv_name]['dims'][1]  # output dim
+            n_head_total = n_head + 2 * n_head_kv
+            if n_head_total > 0:
+                head_dim = qkv_cols // n_head_total
+        if head_dim == 0 and n_head > 0:
+            head_dim = n_embd // n_head
+
         return {
             'arch': arch,
             'n_layers': self.kv.get(f'{arch}.block_count', 0),
-            'n_embd': self.kv.get(f'{arch}.embedding_length', 0),
-            'n_head': self.kv.get(f'{arch}.attention.head_count', 0),
-            'n_head_kv': self.kv.get(f'{arch}.attention.head_count_kv', 0),
+            'n_embd': n_embd,
+            'n_head': n_head,
+            'n_head_kv': n_head_kv,
             'n_ff': self.kv.get(f'{arch}.feed_forward_length', 0),
             'vocab_size': self.kv.get(f'{arch}.vocab_size', 0),
             'rms_eps': self.kv.get(f'{arch}.attention.layer_norm_rms_epsilon', 1e-6),
             'rope_base': self.kv.get(f'{arch}.rope.freq_base', 10000.0),
             'swa_window': self.kv.get(f'{arch}.attention.sliding_window', 0),
+            'head_dim': head_dim,
         }
 
     def get_tensor_f32(self, name):
@@ -458,7 +476,7 @@ class TransformerRunner:
         self.model = model
         self.cfg = config
         self.verbose = verbose
-        self.head_dim = config['n_embd'] // config['n_head']
+        self.head_dim = config.get('head_dim', config['n_embd'] // config['n_head'])
         self.act_fn = ACTIVATION_MAP.get(config['arch'], silu)
 
         # Importance accumulators: tensor_name → (sum_x2, count)
@@ -701,6 +719,131 @@ class TransformerRunner:
 
         return hidden
 
+    def forward_qwen35_layer(self, hidden, layer_idx, cos_f, sin_f):
+        """Forward pass through a Qwen 3.6 hybrid layer (attention + SSM).
+
+        Qwen 3.6 uses:
+          - Fused attn_qkv.weight (Q+K+V in one tensor)
+          - attn_gate.weight (gated attention output, not attn_output)
+          - SSM tensors: ssm_alpha, ssm_beta, ssm_conv1d, ssm_out
+          - post_attention_norm.weight (instead of ffn_norm)
+        """
+        pfx = self._layer_prefix(layer_idx)
+        cfg = self.cfg
+        n_head = cfg['n_head']
+        n_head_kv = cfg['n_head_kv']
+        head_dim = self.head_dim
+        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'])
+
+        # ── Fused QKV projection ──
+        qkv_w = self._get_weight(f'{pfx}.attn_qkv.weight')
+        gate_w = self._get_weight(f'{pfx}.attn_gate.weight')
+
+        attn_out_vec = np.zeros_like(hidden)
+        if qkv_w is not None:
+            self._record(f'{pfx}.attn_qkv.weight', normed)
+
+            qkv = normed @ qkv_w.T  # [seq, (n_head + 2*n_head_kv) * head_dim]
+
+            # Split into Q, K, V
+            q_dim = n_head * head_dim
+            kv_dim = n_head_kv * head_dim
+            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)
+
+            # RoPE
+            q = apply_rope(q, cos_f, sin_f)
+            k = apply_rope(k, cos_f, sin_f)
+
+            # GQA: repeat KV heads
+            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)
+
+            # Scaled dot-product attention
+            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)
+
+            scale = 1.0 / np.sqrt(head_dim)
+            attn = np.matmul(q_t, k_t.transpose(0, 2, 1)) * scale
+
+            # Causal mask (with optional SWA)
+            mask = np.triu(np.full((seq_len, seq_len), -1e9, dtype=np.float32), k=1)
+            swa = cfg.get('swa_window', 0)
+            if swa and swa > 0:
+                for i in range(seq_len):
+                    for j in range(0, max(0, i - swa)):
+                        mask[i, j] = -1e9
+            attn = attn + mask[np.newaxis, :, :]
+            attn = softmax(attn, axis=-1)
+
+            out = np.matmul(attn, v_t)
+            attn_result = out.transpose(1, 0, 2).reshape(seq_len, -1)  # [seq, n_head*head_dim]
+
+            # Gated attention output
+            if gate_w is not None:
+                self._record(f'{pfx}.attn_gate.weight', attn_result)
+                attn_out_vec = attn_result @ gate_w.T
+
+        # ── SSM path ──
+        ssm_alpha_w = self._get_weight(f'{pfx}.ssm_alpha.weight')
+        ssm_beta_w = self._get_weight(f'{pfx}.ssm_beta.weight')
+        ssm_conv_w = self._get_weight(f'{pfx}.ssm_conv1d.weight')
+        ssm_out_w = self._get_weight(f'{pfx}.ssm_out.weight')
+
+        ssm_out_vec = np.zeros_like(hidden)
+        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_conv_w is not None:
+            # ssm_conv1d input is the QKV projection (reuse from attention)
+            if qkv_w is not None:
+                qkv_for_ssm = normed @ qkv_w.T
+                self._record(f'{pfx}.ssm_conv1d.weight', qkv_for_ssm)
+        if ssm_out_w is not None:
+            # SSM output projection — use normed as proxy for recurrent output
+            ssm_proxy = normed[:, :ssm_out_w.shape[1]] if normed.shape[-1] >= ssm_out_w.shape[1] else normed
+            self._record(f'{pfx}.ssm_out.weight', ssm_proxy)
+            ssm_out_vec = ssm_proxy @ ssm_out_w.T
+
+        # Combine attention + SSM
+        hidden = hidden + attn_out_vec + ssm_out_vec
+
+        # ── FFN (uses post_attention_norm instead of ffn_norm) ──
+        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(self, token_ids):
         """Full forward pass, collecting importance statistics."""
         cfg = self.cfg
@@ -718,14 +861,19 @@ class TransformerRunner:
 
         # Process each layer
         for layer_idx in range(cfg['n_layers']):
-            # Check if this layer has DeltaNet tensors (Qwen 3.6 hybrid)
             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
-            has_full_attn = f'{pfx}.attn_q.weight' in self.model.tensor_infos
 
-            if has_linear_attn and not has_full_attn:
+            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)
             else:
+                # Standard transformer (LLaMA, Mistral, Gemma, etc.)
                 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')