From e671780008b12ce045f3a213615e79ccb90a5a85 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 1 Jun 2026 00:30:37 +0000 Subject: [PATCH] fix: transpose checkpoint weights before make_b_k_major in Nvfp4Linear/SharedExpert MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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). --- dsv4/layers/linear.py | 12 +++++++++--- dsv4/layers/shared_expert.py | 14 +++++++++----- single_shot_inference.py | 8 +++++++- 3 files changed, 25 insertions(+), 9 deletions(-) diff --git a/dsv4/layers/linear.py b/dsv4/layers/linear.py index 6ae4b955..7b54e606 100644 --- a/dsv4/layers/linear.py +++ b/dsv4/layers/linear.py @@ -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 diff --git a/dsv4/layers/shared_expert.py b/dsv4/layers/shared_expert.py index ffb8eb37..930b5674 100644 --- a/dsv4/layers/shared_expert.py +++ b/dsv4/layers/shared_expert.py @@ -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) diff --git a/single_shot_inference.py b/single_shot_inference.py index 308dda1b..da9ab3d8 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -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)