diff --git a/dsv4/ops/gemm_runner.py b/dsv4/ops/gemm_runner.py index a6dbe50d..e9ec40e9 100644 --- a/dsv4/ops/gemm_runner.py +++ b/dsv4/ops/gemm_runner.py @@ -175,8 +175,17 @@ def run_nvfp4_grouped_gemm( out = torch.zeros(tokens_sum, n_dim, dtype=torch.bfloat16, device=mat_a.device) - cache_key = (num_experts, str(mat_a.device), mma_tiler_mn, cluster_shape_mn, - mat_a.shape[1], mat_b.shape[2]) + # NVFP4-3: use 2-CTA UMMA for M>=256 (1.7-1.9× throughput at prefill) + use_2cta = tokens_sum >= 256 and cluster_shape_mn[0] % 2 == 0 + effective_mma_tiler_mn = ( + (mma_tiler_mn[0] * 2, mma_tiler_mn[1]) if use_2cta else mma_tiler_mn + ) + effective_cluster_shape_mn = ( + (cluster_shape_mn[0] * 2, cluster_shape_mn[1]) if use_2cta else cluster_shape_mn + ) + + cache_key = (num_experts, str(mat_a.device), effective_mma_tiler_mn, effective_cluster_shape_mn, + mat_a.shape[1], mat_b.shape[2], use_2cta) if cache_key not in _compiled_kernel_cache: # Lazy compilation — safety net if warmup_compilation wasn't called. @@ -187,8 +196,9 @@ def run_nvfp4_grouped_gemm( accumulate_on_output=False, separate_tensormap_init=True, consistent_token_padding=False, - mma_tiler_mnk=(*mma_tiler_mn, 256), - cluster_shape_mnk=(*cluster_shape_mn, 1), + mma_tiler_mnk=(*effective_mma_tiler_mn, 256), + cluster_shape_mnk=(*effective_cluster_shape_mn, 1), + use_2cta_instrs=use_2cta, ) def to_cute(t): @@ -385,8 +395,18 @@ def run_fused_swiglu_grouped_gemm( out = torch.zeros(tokens_sum, n_dim, dtype=torch.bfloat16, device=mat_a.device) - cache_key = ('fused', num_experts, str(mat_a.device), mma_tiler_mn, cluster_shape_mn, - mat_a.shape[1], mat_b.shape[2], swiglu_limit) + # NVFP4-3: use 2-CTA UMMA for M>=256 (1.7-1.9× throughput at prefill) + # At decode (M<256), 1-CTA is correct (2-CTA wastes hardware) + use_2cta = tokens_sum >= 256 and cluster_shape_mn[0] % 2 == 0 + effective_mma_tiler_mn = ( + (mma_tiler_mn[0] * 2, mma_tiler_mn[1]) if use_2cta else mma_tiler_mn + ) + effective_cluster_shape_mn = ( + (cluster_shape_mn[0] * 2, cluster_shape_mn[1]) if use_2cta else cluster_shape_mn + ) + + cache_key = ('fused', num_experts, str(mat_a.device), effective_mma_tiler_mn, effective_cluster_shape_mn, + mat_a.shape[1], mat_b.shape[2], swiglu_limit, use_2cta) if cache_key not in _fused_kernel_cache: # Lazy compilation @@ -396,8 +416,9 @@ def run_fused_swiglu_grouped_gemm( accumulate_on_output=False, separate_tensormap_init=True, consistent_token_padding=False, - mma_tiler_mnk=(*mma_tiler_mn, 256), - cluster_shape_mnk=(*cluster_shape_mn, 1), + mma_tiler_mnk=(*effective_mma_tiler_mn, 256), + cluster_shape_mnk=(*effective_cluster_shape_mn, 1), + use_2cta_instrs=use_2cta, fused_swiglu=True, swiglu_limit=swiglu_limit, )