biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix scale assembly: use correctly-sized temp buffer for swizzle

Browse Source
This commit is contained in:
biondizzle
2026-05-18 20:09:50 +00:00
parent 97bdd604e9
commit 70f50a1ec6
1 changed files with 16 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
cutedsl/shared_expert_pipeline.py
View File

@@ -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