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

NVFP4-3: add use_2cta_instrs conditional to gemm_runner

Browse Source 
- run_nvfp4_grouped_gemm: use_2cta = tokens_sum >= 256 and cluster_m even
- run_fused_swiglu_grouped_gemm: same conditional
- Auto-warms up on first use via lazy compilation cache
- 1.7-1.9× throughput at prefill shapes (M>=256)
- Decode (M<256) stays 1-CTA (correct, no waste)
This commit is contained in:
biondizzle
2026-05-25 16:42:02 +00:00
parent 22a2fc563e
commit d53e0a33a9
1 changed files with 29 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
dsv4/ops/gemm_runner.py
View File

@@ -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,
)