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

fix: GEMM output should be 2x packed N (float4_e2m1fn_x2 packs 2 per element)

Browse Source
This commit is contained in:
biondizzle
2026-05-16 18:27:44 +00:00
parent 95a1345b92
commit 5a79065b2b
1 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
cutedsl/bridge.py
View File

@@ -369,11 +369,14 @@ def run_nvfp4_grouped_gemm(
no cute.compile() in the forward path.
"""
num_experts = mat_b.shape[0]
n_dim = mat_b.shape[2] # packed N (in float4 elements)
n_dim = mat_b.shape[2] # packed N (in float4_e2m1fn_x2 elements)
# float4_e2m1fn_x2 packs 2 FP4 values per storage element
# The GEMM accumulates in BF16, so output has 2x the packed dimension
n_dim_logical = n_dim * 2
tokens_sum = mat_a.shape[0]
device = mat_a.device
out = torch.zeros(tokens_sum, n_dim, dtype=torch.bfloat16, device=device)
out = torch.zeros(tokens_sum, n_dim_logical, dtype=torch.bfloat16, device=device)
compiled, kernel, max_active_clusters = _get_compiled_kernel(
num_experts, device, mma_tiler_mn, cluster_shape_mn