From 5a79065b2b6abd2b68e2e434ae44109d4fda7131 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 16 May 2026 18:27:44 +0000 Subject: [PATCH] fix: GEMM output should be 2x packed N (float4_e2m1fn_x2 packs 2 per element) --- cutedsl/bridge.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/cutedsl/bridge.py b/cutedsl/bridge.py index 2822d8c5..b96211b7 100644 --- a/cutedsl/bridge.py +++ b/cutedsl/bridge.py @@ -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