From e3e67c3992414e48a888d0bf162be46428e15500 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 24 May 2026 22:57:49 +0000 Subject: [PATCH] NVFP4-3: enable 2-CTA UMMA when MMA tile M >= 256 (1.7-1.9x throughput) --- dsv4/ops/gemm_runner.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/dsv4/ops/gemm_runner.py b/dsv4/ops/gemm_runner.py index a29a357c..a6dbe50d 100644 --- a/dsv4/ops/gemm_runner.py +++ b/dsv4/ops/gemm_runner.py @@ -298,6 +298,11 @@ def warmup_fused_swiglu_compilation(num_experts, K_packed, N_packed, device, global_scale_a = torch.ones(1, dtype=torch.float32, device=device) global_scale_b = torch.ones(1, dtype=torch.float32, device=device) + # NVFP4-3: Enable 2-CTA UMMA when MMA tile M >= 256 and cluster M is even. + # 2-CTA gives 1.7-1.9× throughput at prefill shapes by doubling effective MMA tile M. + # Not beneficial at decode (M < 256) — would waste hardware. + use_2cta = mma_tiler_mn[0] >= 256 and cluster_shape_mn[0] % 2 == 0 + kernel = FusedSwiGLUScaledGroupedGemmKernel( scenario="2Dx3D", sf_vec_size=SF_VEC_SIZE, @@ -308,6 +313,7 @@ def warmup_fused_swiglu_compilation(num_experts, K_packed, N_packed, device, cluster_shape_mnk=(*cluster_shape_mn, 1), fused_swiglu=True, swiglu_limit=swiglu_limit, + use_2cta_instrs=use_2cta, ) def to_cute(t):