god dam i just want the gemm in

2026-05-15 05:31:13 +00:00
parent 3b3c506af5
commit 6a4f52cedc
1 changed files with 1 additions and 1 deletions
--- a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py
+++ b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py
@@ -76,7 +76,7 @@ def cutlass_grouped_nvfp4_gemm(
        M_expert = token_indices.shape[0]
        
        # DEBUG: verify data going into GEMM
-        if e == 0:
+        if e < 3 and M_expert > 0:
            print(f"[GEMM-IN] expert={e} M={M_expert} N={N} K={K} "
                  f"w shape={expert_w.shape} w_sf shape={expert_w_sf.shape} "
                  f"w absmax={expert_w.view(torch.int8).abs().max().item()} "