Fix wo_a: scatter each group's data at correct offset in padded buffer

The grouped GEMM expects each group's tokens at their own offset range:
- Group 0: rows [0, padded_T)
- Group 1: rows [padded_T, 2*padded_T)
- etc.

Previously we wrote all groups' data contiguously starting at row 0,
so group 1+ would read zeros from the padding area. Now we scatter
each group's quantized activation at the correct offset.

Also:
- Size buffer for total_max_rows = padded_max * n_groups
- Use assemble_scales_2d_side for multi-group scale assembly
- Extract output per-group at correct offsets

cutedsl/wo_a_grouped_linear.py

@@ -143,10 +143,11 @@ class CuTeDSLNvfp4WoA:
 
     def _allocate_buffers(self):
         """Pre-allocate buffers at max size for cudagraph compatibility."""
-        max_rows = cutedsl_ceil_div(self.max_num_tokens, 128) * 128
+        max_rows_per_group = cutedsl_ceil_div(self.max_num_tokens, 128) * 128
+        total_max_rows = max_rows_per_group * self.n_local_groups
 
         self._padded_x_fp4_buf = torch.zeros(
-            max_rows, self.group_in_features // 2, dtype=torch.uint8, device=self.device
+            total_max_rows, self.group_in_features // 2, dtype=torch.uint8, device=self.device
         ).view(torch.float4_e2m1fn_x2)
 
         self._gsa_buf = torch.zeros(self.n_local_groups, dtype=torch.float32, device=self.device)
@@ -212,38 +213,53 @@ class CuTeDSLNvfp4WoA:
         We reshape to (tokens, n_local_groups, heads_per_group * head_dim),
         then treat each group's (tokens, group_in_features) as one "expert"
         in our grouped GEMM. All tokens go to all groups.
+
+        The grouped GEMM layout requires each group's tokens to be
+        contiguous at their correct offset:
+        - Group 0: rows [0, padded_T)
+        - Group 1: rows [padded_T, 2*padded_T)
+        - ...
+        - Group G: rows [(G-1)*padded_T, G*padded_T)
         """
         self._ensure_initialized()
 
         num_tokens = o.shape[0]
+        padded_rows_per_group = cutedsl_ceil_div(num_tokens, 128) * 128
+
         # 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)
 
-        # 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)
+        # Permute to groups-first: (G, T, D)
+        o_grouped = o_grouped.permute(1, 0, 2)
 
-        padded_rows_per_group = cutedsl_ceil_div(num_tokens, 128) * 128
-        total_padded = padded_rows_per_group * self.n_local_groups
-
-        # Quantize activation
-        x_fp4, x_sf = quantize_activation_nvfp4(
-            o_flat, self._activation_global_scale
-        )
-
-        # Scatter into padded buffer
+        # Quantize each group's activation and scatter into padded buffer
         padded_x_fp4 = self._padded_x_fp4_buf
         padded_x_fp4.view(torch.uint8).zero_()
-        padded_x_fp4.view(torch.uint8)[:num_tokens * self.n_local_groups] = x_fp4.view(torch.uint8)
 
-        # Assemble A-side scales
-        scale_a = self._assemble_scales_single_group(x_sf)
+        # We need to collect scales for ALL groups for the GEMM
+        all_x_sf = []
 
-        # Expert offsets: cumulative [padded_rows, 2*padded_rows, ..., n_groups*padded_rows]
+        for g in range(self.n_local_groups):
+            group_act = o_grouped[g]  # (T, group_in_features)
+
+            # Quantize this group's activation
+            x_fp4_g, x_sf_g = quantize_activation_nvfp4(
+                group_act, self._activation_global_scale
+            )
+
+            # Scatter into the padded buffer at the correct offset
+            offset = g * padded_rows_per_group
+            padded_x_fp4.view(torch.uint8)[offset:offset + num_tokens] = x_fp4_g.view(torch.uint8)
+
+            all_x_sf.append(x_sf_g)
+
+        # Assemble A-side scales for all groups
+        # The grouped GEMM expects scales for all groups assembled together
+        # For 2Dx3D scenario, scale_a is assembled from per-group scale tensors
+        from cutedsl.bridge import assemble_scales_2d_side
+        scale_a = assemble_scales_2d_side(all_x_sf)
+
+        # Expert offsets: cumulative [padded_T, 2*padded_T, ..., n_groups*padded_T]
         expert_offsets = self._expert_offsets_buf
         for g in range(self.n_local_groups):
             expert_offsets[g] = (g + 1) * padded_rows_per_group
@@ -263,11 +279,13 @@ class CuTeDSLNvfp4WoA:
         )
 
         # Extract real outputs and reshape
-        out = out[:num_tokens * self.n_local_groups]
-        # 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)
+        # GEMM output has the same layout as mat_a: groups-first with padding
+        z = torch.empty(num_tokens, self.n_local_groups, self.o_lora_rank,
+                        dtype=torch.bfloat16, device=o.device)
+        for g in range(self.n_local_groups):
+            offset = g * padded_rows_per_group
+            z[:, g, :] = out[offset:offset + num_tokens, :]
+
         return z
 
     def __call__(self, o: torch.Tensor) -> torch.Tensor: