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

fix: transpose B and SFB on the Python side at weight-load time, and adjust the SFB remap kernel to read from column-major source layout

Browse Source
This commit is contained in:
biondizzle
2026-05-15 04:35:45 +00:00
parent c56cc34ae1
commit c7f6a1dc4d
5 changed files with 41 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py
View File

@@ -20,8 +20,8 @@ except ImportError:
def cutlass_nvfp4_blockscaled_gemm(
A_packed, # (M, K_half) int8 packed E2M1
SFA, # scale factors for A (float8_e4m3fn)
B_packed, # (N, K_half) int8 packed E2M1
SFB, # scale factors for B (float8_e4m3fn)
B_packed, # (K_half, N) int8 packed E2M1, column-major for CUTLASS
SFB, # scale factors for B (sf_k, N) float8_e4m3fn, column-major for CUTLASS
M, N, K, # Problem dimensions (K in FP4 elements)
alpha=1.0, # fp32 scalar applied in epilogue: D = alpha * A @ B + beta * C
):
@@ -34,8 +34,8 @@ def cutlass_nvfp4_blockscaled_gemm(
def cutlass_grouped_nvfp4_gemm(
x_fp4, # (num_tokens, K_half) int8 packed E2M1
x_sf, # (num_tokens, sf_k) float8_e4m3fn block scales
weights, # (E_per_rank, N, K_half) int8 packed E2M1
weight_sf, # (E_per_rank, N, sf_k) float8_e4m3fn block scales
weights, # (E_per_rank, K_half, N) int8 packed E2M1, column-major for CUTLASS
weight_sf, # (E_per_rank, sf_k, N) float8_e4m3fn, column-major for CUTLASS
topk_ids, # (num_tokens, NUM_TOPK) int32
topk_weights, # (num_tokens, NUM_TOPK) float32
alpha=1.0, # fp32 scalar: D = alpha * A @ B (from stage_activation global scale)
@@ -48,7 +48,8 @@ def cutlass_grouped_nvfp4_gemm(
num_tokens = x_fp4.shape[0]
K_half = x_fp4.shape[1]
K = K_half * 2 # Actual K dimension (2 FP4 per byte)
N = weights.shape[1] # Output dimension
# Weights are (E, K_half, N) column-major (transposed at load time for CUTLASS ColumnMajor B)
N = weights.shape[2] # Output dimension
num_experts = weights.shape[0]
num_topk = topk_ids.shape[1]
@@ -69,8 +70,8 @@ def cutlass_grouped_nvfp4_gemm(
# Gather tokens for this expert
expert_x = x_fp4[token_indices] # (num_expert_tokens, K_half)
expert_x_sf = x_sf[token_indices] # (num_expert_tokens, sf_k)
expert_w = weights[e] # (N, K_half)
expert_w_sf = weight_sf[e] # (N, sf_k) — THIS IS SCALES, NOT WEIGHTS
expert_w = weights[e] # (K_half, N) column-major for CUTLASS
expert_w_sf = weight_sf[e] # (sf_k, N) column-major for CUTLASS
M_expert = token_indices.shape[0]