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

fix: transpose checkpoint weights before make_b_k_major in Nvfp4Linear/SharedExpert

Browse Source 
Critical bug: checkpoint weights are (N_packed, K_packed) N-major format,
but make_b_k_major expects (E, K_packed, N_packed) input. Without the
permute, the K and N dimensions are swapped, producing garbage output
with wrong dimensions (e.g., q_a output was 3584 instead of 1536).

Also fix scale assembly: checkpoint scales are (N, K_sf) which should
use assemble_raw_scales_2d3d_3d_side (no transpose), not
assemble_scales_3d_side (which incorrectly transposes K_sf↔N).
This commit is contained in:
biondizzle
2026-06-01 00:30:37 +00:00
parent e8a7a9256f
commit e671780008
3 changed files with 25 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
dsv4/layers/linear.py
View File

@@ -14,7 +14,6 @@ from dsv4.ops.quantize import (
)
from dsv4.ops.layouts import (
make_b_k_major,
assemble_scales_3d_side,
)
from dsv4.ops.gemm_runner import (
run_nvfp4_grouped_gemm,
@@ -72,8 +71,15 @@ class Nvfp4Linear:
"""Process weights for CuTeDSL GEMM."""
# Convert uint8 checkpoint weights to float4_e2m1fn_x2 view
fp4_view = [w.view(torch.float4_e2m1fn_x2) if w.dtype == torch.uint8 else w for w in self.fp4]
self._mat_b = make_b_k_major(torch.stack(fp4_view)) # (1, K_packed, N_packed)
self._scale_b = assemble_scales_3d_side(self.sf)
# Checkpoint weight is (out_features//2, in_features//2) = (N_packed, K_packed)
# make_b_k_major expects (E, K_packed, N_packed), so we need to permute
stacked = torch.stack(fp4_view).permute(0, 2, 1).contiguous() # (1, K_packed, N_packed)
self._mat_b = make_b_k_major(stacked)
# Checkpoint scale is (N_packed, K_sf) — already in the right row order for the
# kernel's swizzle. Use assemble_raw_scales_2d3d_3d_side (no transpose),
# NOT assemble_scales_3d_side (which transposes K_sf↔N).
from dsv4.ops.layouts import assemble_raw_scales_2d3d_3d_side
self._scale_b = assemble_raw_scales_2d3d_3d_side(self.sf)
self._gsb = torch.tensor(self.gs, dtype=torch.float32, device=self.device)
# Fold weight_scale_2 into global_scale_b

14
dsv4/layers/shared_expert.py
View File

@@ -26,7 +26,6 @@ from dsv4.ops.quantize import (
)
from dsv4.ops.layouts import (
make_b_k_major,
assemble_scales_3d_side,
)
from dsv4.ops.gemm_runner import (
run_nvfp4_grouped_gemm,
@@ -105,11 +104,16 @@ class Nvfp4SharedExpert:
# Convert uint8 checkpoint weights to float4_e2m1fn_x2 view
l1_view = [w.view(torch.float4_e2m1fn_x2) if w.dtype == torch.uint8 else w for w in self.l1_fp4]
l2_view = [w.view(torch.float4_e2m1fn_x2) if w.dtype == torch.uint8 else w for w in self.l2_fp4]
# Checkpoint weight is (N_packed, K_packed), make_b_k_major expects (E, K_packed, N_packed)
l1_stacked = torch.stack(l1_view).permute(0, 2, 1).contiguous()
l2_stacked = torch.stack(l2_view).permute(0, 2, 1).contiguous()
# Stack weights and convert to K-major
self._l1_mat_b = make_b_k_major(torch.stack(l1_view)) # (1, K_packed, N_packed)
self._l2_mat_b = make_b_k_major(torch.stack(l2_view))
self._l1_scale_b = assemble_scales_3d_side(self.l1_sf) # (1, N, K_sf_padded)
self._l2_scale_b = assemble_scales_3d_side(self.l2_sf)
self._l1_mat_b = make_b_k_major(l1_stacked) # (1, K_packed, N_packed)
self._l2_mat_b = make_b_k_major(l2_stacked)
# Checkpoint scale is (N_packed, K_sf) — use assemble_raw_scales_2d3d_3d_side
from dsv4.ops.layouts import assemble_raw_scales_2d3d_3d_side
self._l1_scale_b = assemble_raw_scales_2d3d_3d_side(self.l1_sf)
self._l2_scale_b = assemble_raw_scales_2d3d_3d_side(self.l2_sf)
self._l1_gsb = torch.tensor(self.l1_gs, dtype=torch.float32, device=self.device)
self._l2_gsb = torch.tensor(self.l2_gs, dtype=torch.float32, device=self.device)

8
single_shot_inference.py
View File

@@ -129,9 +129,15 @@ def do_nvfp4_linear_ref(x, w, pfx, proj_name):
def make_nvfp4_linear(in_features, out_features, device, all_w, pfx, proj_name):
from dsv4.layers.linear import Nvfp4Linear
d = device
lin = Nvfp4Linear(in_features, out_features, max_num_tokens=8192, device=d)
weight, ws, ws2, isc = get_nvfp4_weight(all_w, pfx, proj_name)
assert weight is not None, f"{pfx}.{proj_name}.weight not found"
# Infer actual dimensions from checkpoint weight shape
# weight is (O, I//2) uint8 where O = out_features//2
actual_out = weight.shape[0] * 2
actual_in = weight.shape[1] * 2
if actual_out != out_features or actual_in != in_features:
log.warning(f"{pfx}.{proj_name}: declared ({in_features},{out_features}) but weight is ({actual_in},{actual_out})")
lin = Nvfp4Linear(actual_in, actual_out, max_num_tokens=8192, device=d)
lin.fp4 = [weight.to(d)]; lin.sf = [ws.to(d)]
lin.ws2 = [ws2.to(d) if ws2 is not None else None] # weight_scale_2
gs = isc.float().item() if isc is not None else 1.0 / (6.0 * 448.0)