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

12
src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py
View File

@@ -91,8 +91,8 @@ MEGA_MOE_DEBUG = int(os.environ.get("MEGA_MOE_DEBUG", "0"))
def nvfp4_mega_moe_l1(
x_fp4, # (num_tokens, K//2) int8 packed E2M1
x_sf, # (num_tokens, sf_k_groups) uint32 packed UE4M3
l1_weights, # (E_per_rank, 2*INTER, K//2) int8 K-major
l1_scales, # (E_per_rank, 2*INTER, sf_k_groups) uint32 packed UE4M3
l1_weights, # (E_per_rank, K//2, 2*INTER) int8, column-major for CUTLASS
l1_scales, # (E_per_rank, sf_k_groups, 2*INTER) float8_e4m3fn, column-major
topk_ids, # (num_tokens, NUM_TOPK) int32
topk_weights, # (num_tokens, NUM_TOPK) float32
num_experts_per_rank,
@@ -108,7 +108,7 @@ def nvfp4_mega_moe_l1(
num_tokens = x_fp4.shape[0]
K_half = x_fp4.shape[1]
K = K_half * 2 # HIDDEN = 7168
N = l1_weights.shape[1] # 2 * INTERMEDIATE = 6144
N = l1_weights.shape[2] # 2 * INTERMEDIATE = 6144 (column-major: shape is E, K_half, N)
if MEGA_MOE_DEBUG:
print(f"[nvfp4_moe_l1] tokens={num_tokens} K={K} N={N} "
@@ -130,8 +130,8 @@ def nvfp4_mega_moe_l1(
def nvfp4_mega_moe_l2(
x_fp4, # (num_tokens, INTER//2) int8 packed E2M1
x_sf, # (num_tokens, sf_k_groups) uint32 packed UE4M3
l2_weights, # (E_per_rank, HIDDEN, INTER//2) int8 K-major
l2_scales, # (E_per_rank, HIDDEN, sf_k_groups) uint32 packed UE4M3
l2_weights, # (E_per_rank, INTER//2, HIDDEN) int8, column-major for CUTLASS
l2_scales, # (E_per_rank, sf_k_groups, HIDDEN) float8_e4m3fn, column-major
topk_ids, # (num_tokens, NUM_TOPK) int32
topk_weights, # (num_tokens, NUM_TOPK) float32
num_experts_per_rank,
@@ -144,7 +144,7 @@ def nvfp4_mega_moe_l2(
num_tokens = x_fp4.shape[0]
K_half = x_fp4.shape[1]
K = K_half * 2 # INTERMEDIATE = 3072
N = l2_weights.shape[1] # HIDDEN = 7168
N = l2_weights.shape[2] # HIDDEN = 7168 (column-major: shape is E, K_half, N)
if MEGA_MOE_DEBUG:
print(f"[nvfp4_moe_l2] tokens={num_tokens} K={K} N={N} "