[Perf] Tune MoE weighted-sum kernel launch configuration (#59731)

Signed-off-by: Jinzhen Lin <jinzhen.ljz@antgroup.com>
Co-authored-by: Codex <noreply@openai.com>
Co-authored-by: Misha Goin <mgoin64@gmail.com>

tests/kernels/moe/test_moe_fused_mul_sum.py

@@ -16,10 +16,8 @@ from vllm.utils.torch_utils import set_random_seed
 NUM_TOKENS = [1, 17, 256]
 TOP_KS = [2, 8]
 HIDDEN_SIZE = 512
-# One BLOCK_K tile, a partial trailing tile, and Kimi-K3's real dims
-# (moe_intermediate_size 3072, hidden_size 7168) -- exercises the per-tile loop
-# and its tail mask (BLOCK_K caps at 512 for <=2-byte dtypes).
-HIDDEN_SIZES = [512, 1000, 3072, 7168]
+# Include partial tiles and hidden sizes above the CUDA 4096-element tile cap.
+HIDDEN_SIZES = [512, 1000, 3072, 4097, 7168]
 DTYPES = [torch.bfloat16, torch.float32]
 NUM_EXPERTS = 16
 
@@ -233,7 +231,7 @@ def test_nonlocal_and_padding_adjacent():
 @pytest.mark.parametrize("num_tokens", NUM_TOKENS)
 @pytest.mark.parametrize("top_k", TOP_KS)
 @pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
-@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("dtype", [*DTYPES, torch.float16])
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
 def test_weighted_sum_matches_reference(
     num_tokens: int, top_k: int, hidden_size: int, dtype: torch.dtype
@@ -268,7 +266,11 @@ def test_weighted_sum_matches_reference(
 
     assert not out.isnan().any(), "invalid slots leaked into the reduction"
     ref = _reference(inputs.nan_to_num(), topk_weights, topk_ids, expert_map=None)
-    atol, rtol = (1e-2, 1e-2) if dtype == torch.bfloat16 else (1e-4, 1e-4)
+    atol, rtol = {
+        torch.bfloat16: (1e-2, 1e-2),
+        torch.float16: (1e-3, 1e-3),
+        torch.float32: (1e-4, 1e-4),
+    }[dtype]
     torch.testing.assert_close(out.float(), ref, atol=atol, rtol=rtol)
 
 
@@ -360,7 +362,7 @@ def test_num_valid_tokens_bounds_stale_padding(dtype: torch.dtype):
     """
     set_random_seed(0)
     device = "cuda"
-    num_tokens, num_recv, top_k, hidden_size = 17, 10, 8, 512
+    num_tokens, num_recv, top_k, hidden_size = 17, 10, 8, 7168
 
     inputs = torch.randn(num_tokens, top_k, hidden_size, dtype=dtype, device=device)
     topk_weights = torch.rand(num_tokens, top_k, dtype=dtype, device=device)

vllm/model_executor/layers/fused_moe/moe_fused_mul_sum.py

@@ -3,6 +3,7 @@
 import torch
 from torch._subclasses.fake_tensor import FakeTensor
 
+from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 
 
@@ -75,11 +76,18 @@ def moe_fused_mul_sum_kernel(
 def _heuristic_config(
     hidden_size: int,
     element_size: int,
+    device_index: int = 0,
 ):
     is_fp32 = element_size > 2
     max_block_k = 256 if is_fp32 else 512
-    BLOCK_K = max(128, min(triton.next_power_of_2(hidden_size), max_block_k))
     num_warps = 4 if is_fp32 else 2
+    if not is_fp32 and hidden_size >= 1024 and current_platform.is_cuda():
+        capability = current_platform.get_device_capability(device_index)
+        if capability is not None and capability >= (7, 5):
+            # Larger tiles reduce the serial hidden-dimension loop.
+            max_block_k = 4096
+            num_warps = 8 if capability == (7, 5) else 16
+    BLOCK_K = max(128, min(triton.next_power_of_2(hidden_size), max_block_k))
     num_stages = 3
     return BLOCK_K, num_warps, num_stages
 
@@ -149,6 +157,7 @@ def moe_fused_mul_sum(
         BLOCK_K, num_warps, num_stages = _heuristic_config(
             hidden_size,
             inputs.element_size(),
+            inputs.device.index or 0,
         )
         grid = (num_tokens,)
         moe_fused_mul_sum_kernel[grid](