diff --git a/cutedsl/wo_a_grouped_linear.py b/cutedsl/wo_a_grouped_linear.py
index d8d0d49e..16006fea 100644
--- a/cutedsl/wo_a_grouped_linear.py
+++ b/cutedsl/wo_a_grouped_linear.py
@@ -219,7 +219,12 @@ class CuTeDSLNvfp4WoA:
         # Reshape: (tokens, n_local_heads, head_dim) → (tokens, n_local_groups, group_in_features)
         o_grouped = o.reshape(num_tokens, self.n_local_groups, self.group_in_features)
 
-        # Flatten for GEMM: (tokens * n_groups, group_in_features)
+        # CRITICAL: permute to (n_local_groups, tokens, group_in_features) before flattening.
+        # The grouped GEMM expects mat_a to be laid out contiguously per group:
+        # [all tokens for group0, all tokens for group1, ...]
+        # A simple reshape of (T, G, D) → (T*G, D) gives interleaved:
+        # [t0_g0, t0_g1, ..., t0_g7, t1_g0, ...] which is WRONG.
+        o_grouped = o_grouped.permute(1, 0, 2)  # (G, T, D)
         o_flat = o_grouped.reshape(num_tokens * self.n_local_groups, self.group_in_features)
 
         padded_rows_per_group = cutedsl_ceil_div(num_tokens, 128) * 128
@@ -259,7 +264,10 @@ class CuTeDSLNvfp4WoA:
 
         # Extract real outputs and reshape
         out = out[:num_tokens * self.n_local_groups]
-        z = out.reshape(num_tokens, self.n_local_groups, self.o_lora_rank)
+        # GEMM output is (G*T, o_lora_rank) in groups-first order
+        # Reshape to (G, T, o_lora_rank) then permute to (T, G, o_lora_rank)
+        z = out.reshape(self.n_local_groups, num_tokens, self.o_lora_rank)
+        z = z.permute(1, 0, 2)  # (T, G, o_lora_rank)
         return z
 
     def __call__(self, o: torch.Tensor) -> torch.Tensor: