Experimental

generate_imatrix.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 """
-HExState Importance Matrix Generator — HPC-Enhanced iMatrix from GGUF
+HexState Importance Matrix Generator — HPC-Enhanced iMatrix from GGUF
 
 Runs transformer forward passes over calibration text to collect per-channel
 E[x²] activation statistics, then uses HPC triality BP to propagate importance
@@ -147,6 +147,7 @@ class GGUFModel:
             '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),
         }
 
     def get_tensor_f32(self, name):
@@ -342,9 +343,21 @@ def softmax(x, axis=-1):
     e = np.exp(x - x_max)
     return e / np.sum(e, axis=axis, keepdims=True)
 
+def silu(x):
+    """SiLU / Swish activation — used by LLaMA, Mistral, Qwen, DeepSeek."""
+    return x * (1.0 / (1.0 + np.exp(-np.clip(x, -88, 88))))
+
 def gelu(x):
+    """GELU activation — used by Gemma, GPT-2."""
     return 0.5 * x * (1.0 + np.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * x**3)))
 
+# Architecture → activation function mapping
+ACTIVATION_MAP = {
+    'llama': silu, 'mistral': silu, 'qwen2': silu, 'qwen2moe': silu,
+    'phi3': silu, 'falcon': silu, 'deepseek': silu, 'deepseek2': silu,
+    'gemma': gelu, 'gemma2': gelu, 'gpt2': gelu,
+}
+
 
 class TransformerRunner:
     """Minimal Gemma transformer for importance collection."""
@@ -354,6 +367,7 @@ class TransformerRunner:
         self.cfg = config
         self.verbose = verbose
         self.head_dim = config['n_embd'] // config['n_head']
+        self.act_fn = ACTIVATION_MAP.get(config['arch'], silu)
 
         # Importance accumulators: tensor_name → (sum_x2, count)
         self.importance = {}
@@ -438,8 +452,16 @@ class TransformerRunner:
         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
+        # Causal mask (with optional sliding window)
         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:
+            # Sliding window: mask out positions beyond the window
+            for i in range(seq_len):
+                for j in range(max(0, i - swa), i):
+                    pass  # already 0 (visible)
+                for j in range(0, max(0, i - swa)):
+                    mask[i, j] = -1e9
         attn = attn + mask[np.newaxis, :, :]
         attn = softmax(attn, axis=-1)
 
@@ -466,7 +488,7 @@ class TransformerRunner:
             self._record(f'{pfx}.ffn_gate.weight', normed_ff)
             self._record(f'{pfx}.ffn_up.weight', normed_ff)
 
-            gate_out = gelu(normed_ff @ gate_w.T)
+            gate_out = self.act_fn(normed_ff @ gate_w.T)
             up_out = normed_ff @ up_w.T
             ff_mid = gate_out * up_out
 
@@ -499,7 +521,7 @@ class TransformerRunner:
                     self._record(f'{pfx}.ffn_gate.{exp_id}.weight', exp_input)
                     self._record(f'{pfx}.ffn_up.{exp_id}.weight', exp_input)
 
-                    g = gelu(exp_input @ ew_gate.T)
+                    g = self.act_fn(exp_input @ ew_gate.T)
                     u = exp_input @ ew_up.T
                     mid = g * u
                     self._record(f'{pfx}.ffn_down.{exp_id}.weight', mid)
@@ -515,6 +537,81 @@ class TransformerRunner:
 
         return hidden
 
+    def forward_linear_attn_layer(self, hidden, layer_idx):
+        """Forward pass through a DeltaNet (gated linear attention) layer.
+
+        Used by Qwen 3.5/3.6 for ~75% of layers. Records importance stats
+        for all SSM projection weights.
+        """
+        pfx = self._layer_prefix(layer_idx)
+        cfg = self.cfg
+        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'])
+
+        # ── DeltaNet projections ──
+        qkv_w = self._get_weight(f'{pfx}.ssm_in_qkv.weight')
+        z_w = self._get_weight(f'{pfx}.ssm_in_z.weight')
+        a_w = self._get_weight(f'{pfx}.ssm_in_a.weight')
+        b_w = self._get_weight(f'{pfx}.ssm_in_b.weight')
+        out_w = self._get_weight(f'{pfx}.ssm_out.weight')
+
+        if qkv_w is None or out_w is None:
+            return hidden
+
+        # Record importance on input activations
+        self._record(f'{pfx}.ssm_in_qkv.weight', normed)
+        if z_w is not None:
+            self._record(f'{pfx}.ssm_in_z.weight', normed)
+        if a_w is not None:
+            self._record(f'{pfx}.ssm_in_a.weight', normed)
+        if b_w is not None:
+            self._record(f'{pfx}.ssm_in_b.weight', normed)
+
+        # Approximate forward: project through QKV and output
+        # (Full DeltaNet recurrence is complex; for importance collection
+        # we just need the activation magnitudes at each projection)
+        qkv = normed @ qkv_w.T
+
+        # For importance: record output projection input
+        # Use qkv as a proxy for the recurrent state output
+        n_out = out_w.shape[1] if out_w.ndim >= 2 else hidden.shape[-1]
+        if qkv.shape[-1] >= n_out:
+            out_input = qkv[:, :n_out]
+        else:
+            out_input = qkv
+        self._record(f'{pfx}.ssm_out.weight', out_input)
+
+        attn_out = out_input @ out_w.T
+        hidden = hidden + attn_out
+
+        # ── FFN (same as standard transformer) ──
+        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_w = 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_w 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_w.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
@@ -532,7 +629,15 @@ class TransformerRunner:
 
         # Process each layer
         for layer_idx in range(cfg['n_layers']):
-            hidden = self.forward_layer(hidden, layer_idx, cos_f, sin_f)
+            # Check if this layer has DeltaNet tensors (Qwen 3.6 hybrid)
+            pfx = f"blk.{layer_idx}"
+            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:
+                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')