diff --git a/cutedsl/shared_expert_pipeline.py b/cutedsl/shared_expert_pipeline.py
index 9c9d4ef0..f2e66fd6 100644
--- a/cutedsl/shared_expert_pipeline.py
+++ b/cutedsl/shared_expert_pipeline.py
@@ -151,28 +151,26 @@ class CuTeDSLSharedExpertRunner:
         if not self._buffers_allocated:
             self._allocate_buffers()
 
-    def _assemble_scales_single_group(self, x_sf, padded_x_sf_buf):
+    def _assemble_scales_single_group(self, x_sf, num_tokens, padded_x_sf_buf):
         """Assemble 2D-side activation scales for num_groups=1.
 
         For a single group, scale assembly is just:
-        1. Zero the padded buffer
-        2. Copy x_sf into the top rows
-        3. Apply pad_and_swizzle_single (pads to 128 rows + Blackwell swizzle)
-        4. Reshape back to 2D (kernel expects 2D scale_a)
+        1. Copy x_sf into a correctly-sized buffer (padded to 128 rows, 4 cols)
+        2. Apply pad_and_swizzle_single (Blackwell swizzle)
+        3. Reshape back to 2D (kernel expects 2D scale_a)
 
-        Same as assemble_raw_scales_2d3d_2d_side but for a single group
-        (no cat of multiple expert scales).
+        The padded buffer must be sized exactly for 128-aligned num_tokens,
+        NOT the max_num_tokens buffer (which would be way too large).
         """
-        padded_x_sf = padded_x_sf_buf
-        padded_x_sf.zero_()
         num_rows, num_cols = x_sf.shape
-        padded_x_sf[:num_rows, :num_cols] = x_sf
-        swizzled_flat = pad_and_swizzle_single(padded_x_sf)
-        # pad_and_swizzle_single returns 1D flattened; reshape to 2D
-        # Total rows = round_up(num_rows, 128), cols = round_up(num_cols, 4)
-        total_rows = cutedsl_ceil_div(num_rows, 128) * 128
-        total_cols = cutedsl_ceil_div(num_cols, 4) * 4
-        return swizzled_flat.reshape(total_rows, total_cols)
+        padded_rows = cutedsl_ceil_div(num_rows, 128) * 128
+        padded_cols = cutedsl_ceil_div(num_cols, 4) * 4
+
+        # Use a temp buffer sized for this exact token count
+        buf = torch.zeros(padded_rows, padded_cols, dtype=torch.float16, device=x_sf.device).to(torch.float8_e4m3fn)
+        buf[:num_rows, :num_cols] = x_sf
+        swizzled_flat = pad_and_swizzle_single(buf)
+        return swizzled_flat.reshape(padded_rows, padded_cols)
 
     def compute_activation_global_scales(self, hidden_states_sample):
         """Compute activation global scales from a warmup forward pass.
@@ -221,7 +219,7 @@ class CuTeDSLSharedExpertRunner:
         padded_x_fp4.view(torch.uint8)[:num_tokens] = x_fp4.view(torch.uint8)
 
         # Assemble A-side scales
-        scale_a = self._assemble_scales_single_group(x_sf, self._padded_x_sf_buf_l1)
+        scale_a = self._assemble_scales_single_group(x_sf, num_tokens, self._padded_x_sf_buf_l1)
 
         # Expert offsets: [padded_rows] for 1 group
         expert_offsets = self._expert_offsets_buf
@@ -260,7 +258,7 @@ class CuTeDSLSharedExpertRunner:
         padded_x_fp4.view(torch.uint8)[:num_tokens] = x_fp4.view(torch.uint8)
 
         # Assemble A-side scales
-        scale_a = self._assemble_scales_single_group(x_sf, self._padded_x_sf_buf_l2)
+        scale_a = self._assemble_scales_single_group(x_sf, num_tokens, self._padded_x_sf_buf_l2)
 
         # Expert offsets
         expert_offsets = self._expert_offsets_buf