Qwen changes

generate_imatrix.py

@@ -166,6 +166,8 @@ class GGUFModel:
             '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,
+            'expert_count': self.kv.get(f'{arch}.expert_count', 0),
+            'expert_used_count': self.kv.get(f'{arch}.expert_used_count', 0),
         }
 
     def get_tensor_f32(self, name):
@@ -643,7 +645,7 @@ class TransformerRunner:
         imp_fu_p, cnt_fu_p = make_imp(f'{pfx}.ffn_up.weight', n_embd if ffn_up_w is not None else 0)
         imp_fd_p, cnt_fd_p = make_imp(f'{pfx}.ffn_down.weight', ffn_d if ffn_down_w is not None else 0)
 
-        # Call C function — entire layer in one call
+        # Call C function — entire layer in one call (FFN part will be NULL if MoE)
         lib.hexstate_forward_layer(
             h_ptr,
             norm_p,
@@ -669,16 +671,97 @@ class TransformerRunner:
             ctypes.c_int64(n_head), ctypes.c_int64(n_head_kv),
             ctypes.c_int64(head_dim), ctypes.c_float(eps))
 
-        # Read back importance
+        # Read back importance for the tensors that WERE processed in C
         for name, arr, cnt in imp_refs:
             self.importance[name] = (arr.astype(np.float64), cnt.value)
 
+        # Handle MoE FFN if C code skipped it
+        if ffn_gate_w is None:
+            # Re-normalize for FFN
+            normed_ff = self._hpc_rms_norm(hidden, ffn_norm_w, eps)
+            hidden = self._forward_moe_ffn(hidden, normed_ff, pfx)
+
         # Force-free per-layer weight buffers (~1.4 GB) before next layer
         del refs, imp_refs
         import gc; gc.collect()
 
         return hidden
 
+    def _forward_moe_ffn(self, hidden, normed_ff, pfx):
+        """Python-side MoE FFN handling (supports packed and shared experts)."""
+        gate_inp_w = self._get_weight(f'{pfx}.ffn_gate_inp.weight')
+        if gate_inp_w is None:
+            return hidden
+
+        self._record(f'{pfx}.ffn_gate_inp.weight', normed_ff)
+        router_logits = normed_ff @ gate_inp_w.T
+        n_experts = router_logits.shape[-1]
+        probs = softmax(router_logits, axis=-1)
+        topk = self.cfg.get('expert_used_count', 2)
+        top_k_indices = np.argsort(probs, axis=-1)[:, -topk:]
+
+        ff_out = np.zeros_like(normed_ff)
+
+        # Check for packed experts (Qwen style)
+        p_gate = self._get_weight(f'{pfx}.ffn_gate_exps.weight')
+        p_up   = self._get_weight(f'{pfx}.ffn_up_exps.weight')
+        p_down = self._get_weight(f'{pfx}.ffn_down_exps.weight')
+
+        for exp_id in range(n_experts):
+            if p_gate is not None:
+                ew_gate = p_gate[exp_id]
+                ew_up   = p_up[exp_id]
+                ew_down = p_down[exp_id]
+            else:
+                ew_gate = self._get_weight(f'{pfx}.ffn_gate.{exp_id}.weight')
+                ew_up   = self._get_weight(f'{pfx}.ffn_up.{exp_id}.weight')
+                ew_down = self._get_weight(f'{pfx}.ffn_down.{exp_id}.weight')
+
+            if ew_gate is None: continue
+
+            mask_exp = np.any(top_k_indices == exp_id, axis=-1)
+            if not np.any(mask_exp): continue
+
+            exp_input = normed_ff[mask_exp]
+            
+            # Record importance
+            if p_gate is not None:
+                self._record(f'{pfx}.ffn_gate_exps.weight', exp_input)
+                self._record(f'{pfx}.ffn_up_exps.weight', exp_input)
+            else:
+                self._record(f'{pfx}.ffn_gate.{exp_id}.weight', exp_input)
+                self._record(f'{pfx}.ffn_up.{exp_id}.weight', exp_input)
+
+            g = self.act_fn(exp_input @ ew_gate.T)
+            u = exp_input @ ew_up.T
+            mid = g * u
+            
+            if p_gate is not None:
+                self._record(f'{pfx}.ffn_down_exps.weight', mid)
+            else:
+                self._record(f'{pfx}.ffn_down.{exp_id}.weight', mid)
+
+            exp_out = mid @ ew_down.T
+            indices = np.where(mask_exp)[0]
+            for i, tidx in enumerate(indices):
+                w = probs[tidx, exp_id]
+                ff_out[tidx] += w * exp_out[i]
+
+        # Shared experts (Qwen style)
+        sh_gate = self._get_weight(f'{pfx}.ffn_gate_shexp.weight')
+        if sh_gate is not None:
+            sh_up = self._get_weight(f'{pfx}.ffn_up_shexp.weight')
+            sh_down = self._get_weight(f'{pfx}.ffn_down_shexp.weight')
+            self._record(f'{pfx}.ffn_gate_shexp.weight', normed_ff)
+            self._record(f'{pfx}.ffn_up_shexp.weight', normed_ff)
+            g = self.act_fn(normed_ff @ sh_gate.T)
+            u = normed_ff @ sh_up.T
+            mid = g * u
+            self._record(f'{pfx}.ffn_down_shexp.weight', mid)
+            ff_out += mid @ sh_down.T
+
+        return hidden + ff_out
+
     def _hpc_rms_norm(self, x, weight, eps):
         """RMS norm via HPC C library, falling back to numpy."""
         if self._hpc_lib and x.flags['C_CONTIGUOUS']:
@@ -1055,43 +1138,7 @@ class TransformerRunner:
             hidden = hidden + ff_out
         else:
             # MoE path
-            gate_inp_w = self._get_weight(f'{pfx}.ffn_gate_inp.weight')
-            if gate_inp_w is not None:
-                self._record(f'{pfx}.ffn_gate_inp.weight', normed_ff)
-                router_logits = normed_ff @ gate_inp_w.T
-                n_experts = router_logits.shape[-1]
-                probs = softmax(router_logits, axis=-1)
-                top2 = np.argsort(probs, axis=-1)[:, -2:]
-
-                ff_out = np.zeros_like(normed_ff)
-                for exp_id in range(n_experts):
-                    ew_gate = self._get_weight(f'{pfx}.ffn_gate.{exp_id}.weight')
-                    ew_up = self._get_weight(f'{pfx}.ffn_up.{exp_id}.weight')
-                    ew_down = self._get_weight(f'{pfx}.ffn_down.{exp_id}.weight')
-                    if ew_gate is None:
-                        continue
-
-                    mask_exp = np.any(top2 == exp_id, axis=-1)  # [seq]
-                    if not np.any(mask_exp):
-                        continue
-
-                    exp_input = normed_ff[mask_exp]
-                    self._record(f'{pfx}.ffn_gate.{exp_id}.weight', exp_input)
-                    self._record(f'{pfx}.ffn_up.{exp_id}.weight', exp_input)
-
-                    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)
-
-                    exp_out = mid @ ew_down.T
-                    # Weight by routing probability
-                    for token_idx in np.where(mask_exp)[0]:
-                        w = probs[token_idx, exp_id]
-                        local_idx = np.sum(mask_exp[:token_idx])
-                        ff_out[token_idx] += w * exp_out[local_idx]
-
-                hidden = hidden + ff_out
+            hidden = self._forward_moe_ffn(hidden, normed_ff, pfx)
 
         return hidden
 
@@ -1167,6 +1214,8 @@ class TransformerRunner:
             self._record(f'{pfx}.ffn_down.weight', ff_mid)
             ff_out = ff_mid @ down_w.T
             hidden = hidden + ff_out
+        else:
+            hidden = self._forward_moe_ffn(hidden, normed_ff, pfx)
 
         return hidden
 
@@ -1308,6 +1357,8 @@ class TransformerRunner:
             self._record(f'{pfx}.ffn_down.weight', ff_mid)
             ff_out = ff_mid @ down_w.T
             hidden = hidden + ff_out
+        else:
+            hidden = self._forward_moe_ffn(hidden, normed_ff, pfx)
 
         return hidden
 

hexstate_quantize.c

@@ -4252,16 +4252,16 @@ static void hpc_read_importance(HPCGraph *g, const float *x, int64_t dim,
  */
 static void hpc_matmul_graph(const float *x, const float *weight, float *out,
                              float *importance, int64_t *count,
-                             int64_t M, int64_t K, int64_t N)
+                             int64_t M, int64_t K, int64_t N, int trans_w)
 {
     /* Build HPCGraph over input columns for importance */
     int64_t stride = (K > 512) ? K / 512 : 1;
     int64_t n_sites = (K + stride - 1) / stride;
     HPCGraph *g = hpc_create(n_sites);
+    float *col_energy = (float *)calloc(K, sizeof(float));
 
-    if (g && importance) {
+    if (g && col_energy) {
         /* Compute per-column energies */
-        float *col_energy = (float *)calloc(K, sizeof(float));
         #pragma omp parallel for schedule(static)
         for (int64_t j = 0; j < K; j++) {
             float s = 0.0f;
@@ -4307,28 +4307,37 @@ static void hpc_matmul_graph(const float *x, const float *weight, float *out,
                 importance[j] += col_energy[j] * (float)boost;
         }
         if (count) *count += M;
-
-        free(col_energy);
-        hpc_destroy(g);
     }
 
-    /* Matmul: out = x @ W.T — OpenMP over rows */
+    /* Matmul: out = x @ W.T (trans_w=0) or x @ W (trans_w=1) */
     #pragma omp parallel for schedule(static)
     for (int64_t i = 0; i < M; i++) {
         const float *xi = x + i * K;
         float *oi = out + i * N;
-        for (int64_t n = 0; n < N; n++) {
-            const float *wn = weight + n * K;
-            float dot = 0.0f;
-            for (int64_t k = 0; k < K; k++)
-                dot += xi[k] * wn[k];
-            oi[n] = dot;
+        if (trans_w) {
+            for (int64_t n = 0; n < N; n++) {
+                float dot = 0.0f;
+                for (int64_t k = 0; k < K; k++)
+                    dot += xi[k] * weight[k * N + n];
+                oi[n] = dot;
+            }
+        } else {
+            for (int64_t n = 0; n < N; n++) {
+                const float *wn = weight + n * K;
+                float dot = 0.0f;
+                for (int64_t k = 0; k < K; k++)
+                    dot += xi[k] * wn[k];
+                oi[n] = dot;
+            }
         }
     }
+
+    if (col_energy) free(col_energy);
+    if (g) hpc_destroy(g);
 }
 
 /* ── Helper: RMS norm (OpenMP) ── */
-static void hpc_rms_norm(const float *x, const float *w, float *out,
+void hexstate_rms_norm(const float *x, const float *w, float *out,
                          int64_t seq, int64_t dim, float eps)
 {
     #pragma omp parallel for schedule(static)
@@ -4408,7 +4417,7 @@ void hexstate_forward_layer(
     if (!normed) return;
 
     /* ══════════════ Phase 1: Attention Norm ══════════════ */
-    hpc_rms_norm(hidden, norm_w, normed, seq_len, n_embd, eps);
+    hexstate_rms_norm(hidden, norm_w, normed, seq_len, n_embd, eps);
 
     /* ══════════════ Phase 2: QKV Projection via HPC Graph ══════════════ */
     float *attn_out = (float *)calloc(seq_len * n_embd, sizeof(float));
@@ -4421,7 +4430,7 @@ void hexstate_forward_layer(
 
         /* Graph-based matmul: importance via HPCGraph marginals */
         hpc_matmul_graph(normed, qkv_w, qkv, imp_qkv, cnt_qkv,
-                         seq_len, n_embd, qkv_dim);
+                         seq_len, n_embd, qkv_dim, 0);
 
         /* Split Q, K, V */
         int64_t q_total = n_head * head_dim;
@@ -4444,8 +4453,10 @@ void hexstate_forward_layer(
         HPCGraph *attn_graph = hpc_create(n_head);
         float *S = (float *)calloc(n_head * head_dim * head_dim, sizeof(float));
         float *z_acc = (float *)calloc(n_head * head_dim, sizeof(float));
+        int64_t inner_dim = n_head * head_dim;
+        float *attn_inner = (float *)calloc(seq_len * inner_dim, sizeof(float));
 
-        if (attn_graph && S && z_acc) {
+        if (attn_graph && S && z_acc && attn_inner) {
             for (int64_t t = 0; t < seq_len; t++) {
                 /* Extract Q/K/V for this timestep (handle strided layout) */
                 float *qt_base = qkv + t * qkv_dim;
@@ -4519,15 +4530,14 @@ void hexstate_forward_layer(
                         den += qf[d] * zh[d];
                     float inv_den = 1.0f / den;
 
-                    /* Write to attn_out at position [t, h*head_dim ... ] */
-                    float *ao = attn_out + t * n_embd;
+                    /* Write to attn_inner at position [t, h*head_dim ... ] */
+                    float *ao = attn_inner + t * inner_dim;
                     for (int64_t d2 = 0; d2 < head_dim; d2++) {
                         float num = 0.0f;
                         for (int64_t d1 = 0; d1 < head_dim; d1++)
                             num += qf[d1] * Sh[d1 * head_dim + d2];
-                        /* Accumulate into attn_out (multiple heads write here) */
-                        if (h * head_dim + d2 < n_embd)
-                            ao[h * head_dim + d2] = num * inv_den;
+                        /* Accumulate into attn_inner (multiple heads write here) */
+                        ao[h * head_dim + d2] = num * inv_den;
                     }
                 }
 
@@ -4539,25 +4549,25 @@ void hexstate_forward_layer(
 
         /* Gate projection if present */
         if (gate_w && gate_rows > 0) {
-            /* attn_out is [seq, n_head*head_dim], gate_w is [n_embd, n_head*head_dim] */
-            /* gate_rows may be larger than n_embd, allocate exactly what we need */
-            float *gated = (float *)malloc(seq_len * gate_rows * sizeof(float));
+            int trans_w = (gate_rows == inner_dim) ? 1 : 0;
+            int64_t N_out = trans_w ? n_embd : gate_rows;
+            float *gated = (float *)malloc(seq_len * N_out * sizeof(float));
             if (gated) {
-                hpc_matmul_graph(attn_out, gate_w, gated, imp_gate, cnt_gate,
-                                seq_len, n_head * head_dim, gate_rows);
-                
-                int64_t copy_dim = (gate_rows > n_embd) ? n_embd : gate_rows;
-                /* Must copy row-by-row since strides may differ */
-                #pragma omp parallel for schedule(static)
-                for (int64_t i = 0; i < seq_len; i++) {
-                    memcpy(attn_out + i * n_embd, gated + i * gate_rows, copy_dim * sizeof(float));
-                    if (copy_dim < n_embd) {
-                        memset(attn_out + i * n_embd + copy_dim, 0, (n_embd - copy_dim) * sizeof(float));
-                    }
+                hpc_matmul_graph(attn_inner, gate_w, gated, imp_gate, cnt_gate,
+                                seq_len, inner_dim, N_out, trans_w);
+                for (int64_t t = 0; t < seq_len; t++) {
+                    int64_t copy_dim = N_out < n_embd ? N_out : n_embd;
+                    memcpy(attn_out + t * n_embd, gated + t * N_out, copy_dim * sizeof(float));
                 }
                 free(gated);
             }
+        } else {
+            for (int64_t t = 0; t < seq_len; t++) {
+                int64_t copy_dim = inner_dim < n_embd ? inner_dim : n_embd;
+                memcpy(attn_out + t * n_embd, attn_inner + t * inner_dim, copy_dim * sizeof(float));
+            }
         }
+        if (attn_inner) free(attn_inner);
 
         if (attn_graph) hpc_destroy(attn_graph);
         free(S); free(z_acc); free(qkv);
@@ -4568,23 +4578,25 @@ void hexstate_forward_layer(
         float *K_buf = (float *)malloc(seq_len * k_dim * sizeof(float));
         float *V_buf = (float *)malloc(seq_len * v_dim * sizeof(float));
         if (!Q || !K_buf || !V_buf) {
-            free(Q); free(K_buf); free(V_buf);
+            if(Q) free(Q); if(K_buf) free(K_buf); if(V_buf) free(V_buf);
             free(normed); free(attn_out);
             return;
         }
 
-        hpc_matmul_graph(normed, q_w, Q, imp_q, cnt_q, seq_len, n_embd, q_dim);
-        hpc_matmul_graph(normed, k_w, K_buf, imp_k, cnt_k, seq_len, n_embd, k_dim);
-        hpc_matmul_graph(normed, v_w, V_buf, imp_v, cnt_v, seq_len, n_embd, v_dim);
+        hpc_matmul_graph(normed, q_w, Q, imp_q, cnt_q, seq_len, n_embd, q_dim, 0);
+        hpc_matmul_graph(normed, k_w, K_buf, imp_k, cnt_k, seq_len, n_embd, k_dim, 0);
+        hpc_matmul_graph(normed, v_w, V_buf, imp_v, cnt_v, seq_len, n_embd, v_dim, 0);
 
         /* Same HPC attention as above but with separate Q/K/V buffers */
         int64_t hd_q = q_dim / n_head;
         int64_t hd_kv = k_dim / n_head_kv;
+        int64_t inner_dim = n_head * hd_kv;
         HPCGraph *attn_graph = hpc_create(n_head);
         float *S = (float *)calloc(n_head * hd_kv * hd_kv, sizeof(float));
         float *z_acc = (float *)calloc(n_head * hd_kv, sizeof(float));
+        float *attn_inner = (float *)calloc(seq_len * inner_dim, sizeof(float));
 
-        if (attn_graph && S && z_acc) {
+        if (attn_graph && S && z_acc && attn_inner) {
             for (int64_t t = 0; t < seq_len; t++) {
                 /* Encode heads into graph */
                 for (int64_t h = 0; h < n_head; h++) {
@@ -4635,15 +4647,14 @@ void hexstate_forward_layer(
                         den += qf * zh[d];
                     }
                     float inv_den = 1.0f / den;
-                    float *ao = attn_out + t * n_embd;
+                    float *ao = attn_inner + t * inner_dim;
                     for (int64_t d2 = 0; d2 < hd_kv; d2++) {
                         float num = 0.0f;
                         for (int64_t d1 = 0; d1 < feat; d1++) {
                             float qf = (qh[d1] > 0 ? qh[d1] : 0) + 1e-6f;
                             num += qf * Sh[d1*hd_kv+d2];
                         }
-                        if (h*hd_kv+d2 < n_embd)
-                            ao[h*hd_kv+d2] = num * inv_den;
+                        ao[h*hd_kv+d2] = num * inv_den;
                     }
                 }
                 if (t > 0 && t % 64 == 0)
@@ -4653,16 +4664,36 @@ void hexstate_forward_layer(
 
         /* Output projection */
         if (o_w && o_cols > 0) {
-            int64_t attn_dim = n_head * hd_kv;
-            if (attn_dim > n_embd) attn_dim = n_embd;
+            float *proj_in = attn_inner;
+            int free_proj_in = 0;
+            if (inner_dim != o_cols) {
+                proj_in = (float *)calloc(seq_len * o_cols, sizeof(float));
+                if (proj_in) {
+                    for (int64_t t = 0; t < seq_len; t++) {
+                        int64_t copy_dim = inner_dim < o_cols ? inner_dim : o_cols;
+                        memcpy(proj_in + t * o_cols, attn_inner + t * inner_dim, copy_dim * sizeof(float));
+                    }
+                    free_proj_in = 1;
+                } else {
+                    proj_in = attn_inner;
+                }
+            }
+
             float *projected = (float *)calloc(seq_len * n_embd, sizeof(float));
             if (projected) {
-                hpc_matmul_graph(attn_out, o_w, projected, imp_o, cnt_o,
-                                seq_len, o_cols, n_embd);
+                hpc_matmul_graph(proj_in, o_w, projected, imp_o, cnt_o,
+                                seq_len, o_cols, n_embd, 0);
                 memcpy(attn_out, projected, seq_len * n_embd * sizeof(float));
                 free(projected);
             }
+            if (free_proj_in && proj_in != attn_inner) free(proj_in);
+        } else {
+            for (int64_t t = 0; t < seq_len; t++) {
+                int64_t copy_dim = inner_dim < n_embd ? inner_dim : n_embd;
+                memcpy(attn_out + t * n_embd, attn_inner + t * inner_dim, copy_dim * sizeof(float));
+            }
         }
+        if (attn_inner) free(attn_inner);
 
         if (attn_graph) hpc_destroy(attn_graph);
         free(S); free(z_acc);
@@ -4682,13 +4713,13 @@ void hexstate_forward_layer(
         float *up_out = (float *)malloc(seq_len * ffn_dim * sizeof(float));
 
         if (normed_ff && gate_out && up_out) {
-            hpc_rms_norm(hidden, ffn_norm_w, normed_ff, seq_len, n_embd, eps);
+            hexstate_rms_norm(hidden, ffn_norm_w, normed_ff, seq_len, n_embd, eps);
 
             /* Graph-based matmul for FFN with importance */
             hpc_matmul_graph(normed_ff, ffn_gate_w, gate_out,
-                            imp_ffn_gate, cnt_ffn_gate, seq_len, n_embd, ffn_dim);
+                            imp_ffn_gate, cnt_ffn_gate, seq_len, n_embd, ffn_dim, 0);
             hpc_matmul_graph(normed_ff, ffn_up_w, up_out,
-                            imp_ffn_up, cnt_ffn_up, seq_len, n_embd, ffn_dim);
+                            imp_ffn_up, cnt_ffn_up, seq_len, n_embd, ffn_dim, 0);
 
             /* SiLU(gate) * up */
             hpc_silu(gate_out, seq_len * ffn_dim);
@@ -4697,16 +4728,16 @@ void hexstate_forward_layer(
                 gate_out[i] *= up_out[i];
 
             /* Down projection: graph-based importance recording */
-            float *ff_out = (float *)malloc(seq_len * n_embd * sizeof(float));
-            if (ff_out) {
-                hpc_matmul_graph(gate_out, ffn_down_w, ff_out,
+            float *ff_out_buf = (float *)malloc(seq_len * n_embd * sizeof(float));
+            if (ff_out_buf) {
+                hpc_matmul_graph(gate_out, ffn_down_w, ff_out_buf,
                                 imp_ffn_down, cnt_ffn_down,
-                                seq_len, ffn_dim, n_embd);
+                                seq_len, ffn_dim, n_embd, 0);
                 /* Residual add */
                 #pragma omp parallel for schedule(static)
                 for (int64_t i = 0; i < total; i++)
-                    hidden[i] += ff_out[i];
-                free(ff_out);
+                    hidden[i] += ff_out_buf[i];
+                free(ff_out_buf);
             }
         }
 
@@ -4717,6 +4748,7 @@ void hexstate_forward_layer(
     free(attn_out);
 }
 
+
 #ifndef HEXSTATE_LIBRARY
 /* ═══════════════════════════════════════════════════════════════════════════
  * MAIN

hpc_forward_merged.c

@@ -75,16 +75,16 @@ static void hpc_read_importance(HPCGraph *g, const float *x, int64_t dim,
  */
 static void hpc_matmul_graph(const float *x, const float *weight, float *out,
                              float *importance, int64_t *count,
-                             int64_t M, int64_t K, int64_t N)
+                             int64_t M, int64_t K, int64_t N, int trans_w)
 {
     /* Build HPCGraph over input columns for importance */
     int64_t stride = (K > 512) ? K / 512 : 1;
     int64_t n_sites = (K + stride - 1) / stride;
     HPCGraph *g = hpc_create(n_sites);
+    float *col_energy = (float *)calloc(K, sizeof(float));
 
-    if (g && importance) {
+    if (g && col_energy) {
         /* Compute per-column energies */
-        float *col_energy = (float *)calloc(K, sizeof(float));
         #pragma omp parallel for schedule(static)
         for (int64_t j = 0; j < K; j++) {
             float s = 0.0f;
@@ -130,24 +130,33 @@ static void hpc_matmul_graph(const float *x, const float *weight, float *out,
                 importance[j] += col_energy[j] * (float)boost;
         }
         if (count) *count += M;
-
-        free(col_energy);
-        hpc_destroy(g);
     }
 
-    /* Matmul: out = x @ W.T — OpenMP over rows */
+    /* Matmul: out = x @ W.T (trans_w=0) or x @ W (trans_w=1) */
     #pragma omp parallel for schedule(static)
     for (int64_t i = 0; i < M; i++) {
         const float *xi = x + i * K;
         float *oi = out + i * N;
-        for (int64_t n = 0; n < N; n++) {
-            const float *wn = weight + n * K;
-            float dot = 0.0f;
-            for (int64_t k = 0; k < K; k++)
-                dot += xi[k] * wn[k];
-            oi[n] = dot;
+        if (trans_w) {
+            for (int64_t n = 0; n < N; n++) {
+                float dot = 0.0f;
+                for (int64_t k = 0; k < K; k++)
+                    dot += xi[k] * weight[k * N + n];
+                oi[n] = dot;
+            }
+        } else {
+            for (int64_t n = 0; n < N; n++) {
+                const float *wn = weight + n * K;
+                float dot = 0.0f;
+                for (int64_t k = 0; k < K; k++)
+                    dot += xi[k] * wn[k];
+                oi[n] = dot;
+            }
         }
     }
+
+    if (col_energy) free(col_energy);
+    if (g) hpc_destroy(g);
 }
 
 /* ── Helper: RMS norm (OpenMP) ── */
@@ -243,8 +252,8 @@ void hexstate_forward_layer(
         if (!qkv) { free(normed); free(attn_out); return; }
 
         /* Graph-based matmul: importance via HPCGraph marginals */
-        hpc_matmul_graph(normed, qkv_w, qkv, imp_qkv, cnt_qkv,
-                         seq_len, n_embd, qkv_dim);
+        printf("matmul qkv M=%ld K=%ld N=%ld\n", (long)seq_len, (long)n_embd, (long)qkv_dim); fflush(stdout); hpc_matmul_graph(normed, qkv_w, qkv, imp_qkv, cnt_qkv,
+                         seq_len, n_embd, qkv_dim, 0);
 
         /* Split Q, K, V */
         int64_t q_total = n_head * head_dim;
@@ -363,14 +372,15 @@ void hexstate_forward_layer(
 
         /* Gate projection if present */
         if (gate_w && gate_rows > 0) {
-            /* attn_inner is [seq, n_head*head_dim], gate_w is [gate_rows, n_head*head_dim] */
-            float *gated = (float *)malloc(seq_len * n_embd * sizeof(float));
+            int trans_w = (gate_rows == inner_dim) ? 1 : 0;
+            int64_t N_out = trans_w ? n_embd : gate_rows;
+            float *gated = (float *)malloc(seq_len * N_out * sizeof(float));
             if (gated) {
-                hpc_matmul_graph(attn_inner, gate_w, gated, imp_gate, cnt_gate,
-                                seq_len, inner_dim, gate_rows);
+                printf("matmul gate M=%ld K=%ld N=%ld trans=%d\n", (long)seq_len, (long)inner_dim, (long)N_out, trans_w); fflush(stdout); hpc_matmul_graph(attn_inner, gate_w, gated, imp_gate, cnt_gate,
+                                seq_len, inner_dim, N_out, trans_w);
                 for (int64_t t = 0; t < seq_len; t++) {
-                    int64_t copy_dim = gate_rows < n_embd ? gate_rows : n_embd;
-                    memcpy(attn_out + t * n_embd, gated + t * gate_rows, copy_dim * sizeof(float));
+                    int64_t copy_dim = N_out < n_embd ? N_out : n_embd;
+                    memcpy(attn_out + t * n_embd, gated + t * N_out, copy_dim * sizeof(float));
                 }
                 free(gated);
             }
@@ -391,14 +401,14 @@ void hexstate_forward_layer(
         float *K_buf = (float *)malloc(seq_len * k_dim * sizeof(float));
         float *V_buf = (float *)malloc(seq_len * v_dim * sizeof(float));
         if (!Q || !K_buf || !V_buf) {
-            free(Q); free(K_buf); free(V_buf);
+            if(Q) free(Q); if(K_buf) free(K_buf); if(V_buf) free(V_buf);
             free(normed); free(attn_out);
             return;
         }
 
-        hpc_matmul_graph(normed, q_w, Q, imp_q, cnt_q, seq_len, n_embd, q_dim);
-        hpc_matmul_graph(normed, k_w, K_buf, imp_k, cnt_k, seq_len, n_embd, k_dim);
-        hpc_matmul_graph(normed, v_w, V_buf, imp_v, cnt_v, seq_len, n_embd, v_dim);
+        hpc_matmul_graph(normed, q_w, Q, imp_q, cnt_q, seq_len, n_embd, q_dim, 0);
+        hpc_matmul_graph(normed, k_w, K_buf, imp_k, cnt_k, seq_len, n_embd, k_dim, 0);
+        hpc_matmul_graph(normed, v_w, V_buf, imp_v, cnt_v, seq_len, n_embd, v_dim, 0);
 
         /* Same HPC attention as above but with separate Q/K/V buffers */
         int64_t hd_q = q_dim / n_head;
@@ -495,7 +505,7 @@ void hexstate_forward_layer(
             float *projected = (float *)calloc(seq_len * n_embd, sizeof(float));
             if (projected) {
                 hpc_matmul_graph(proj_in, o_w, projected, imp_o, cnt_o,
-                                seq_len, o_cols, n_embd);
+                                seq_len, o_cols, n_embd, 0);
                 memcpy(attn_out, projected, seq_len * n_embd * sizeof(float));
                 free(projected);
             }
@@ -530,9 +540,9 @@ void hexstate_forward_layer(
 
             /* Graph-based matmul for FFN with importance */
             hpc_matmul_graph(normed_ff, ffn_gate_w, gate_out,
-                            imp_ffn_gate, cnt_ffn_gate, seq_len, n_embd, ffn_dim);
+                            imp_ffn_gate, cnt_ffn_gate, seq_len, n_embd, ffn_dim, 0);
             hpc_matmul_graph(normed_ff, ffn_up_w, up_out,
-                            imp_ffn_up, cnt_ffn_up, seq_len, n_embd, ffn_dim);
+                            imp_ffn_up, cnt_ffn_up, seq_len, n_embd, ffn_dim, 0);
 
             /* SiLU(gate) * up */
             hpc_silu(gate_out, seq_len * ffn_dim);
@@ -541,16 +551,16 @@ void hexstate_forward_layer(
                 gate_out[i] *= up_out[i];
 
             /* Down projection: graph-based importance recording */
-            float *ff_out = (float *)malloc(seq_len * n_embd * sizeof(float));
-            if (ff_out) {
-                hpc_matmul_graph(gate_out, ffn_down_w, ff_out,
+            float *ff_out_buf = (float *)malloc(seq_len * n_embd * sizeof(float));
+            if (ff_out_buf) {
+                hpc_matmul_graph(gate_out, ffn_down_w, ff_out_buf,
                                 imp_ffn_down, cnt_ffn_down,
-                                seq_len, ffn_dim, n_embd);
+                                seq_len, ffn_dim, n_embd, 0);
                 /* Residual add */
                 #pragma omp parallel for schedule(static)
                 for (int64_t i = 0; i < total; i++)
-                    hidden[i] += ff_out[i];
-                free(ff_out);
+                    hidden[i] += ff_out_buf[i];
+                free(ff_out_buf);
             }
         }