[Kernel] Add router GEMM kernels for (5120, 256) and BF16 output to ll_bf16_gemm (#62042)

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Co-authored-by: Lain <fusiyuan2000@gmail.com>
Co-authored-by: Codex <noreply@openai.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: vllm-agent <claw@inferact.ai>

tests/kernels/moe/test_routing.py

@@ -582,6 +582,20 @@ def test_fused_topk_bias(
     assert_routing_results_close(topk_weights, topk_ids, baseline_weights, baseline_ids)
 
 
+@pytest.mark.parametrize("scoring_func", [0, 1])
+def test_fused_grouped_topk_accepts_idle_rank(scoring_func):
+    """Zero-token routing must leave CUDA ready for the next batch."""
+    if not current_platform.is_cuda():
+        pytest.skip("CUDA required")
+    from vllm import _custom_ops as ops
+
+    scores = torch.empty((0, 8), device="cuda", dtype=torch.float32)
+    bias = torch.zeros(8, device="cuda")
+    weights, ids = ops.grouped_topk(scores, 2, 1, 2, True, 1.0, bias, scoring_func)
+    torch.testing.assert_close(weights, scores.new_empty((0, 2)))
+    torch.testing.assert_close(ids, scores.new_empty((0, 2), dtype=torch.int32))
+
+
 @pytest.mark.parametrize("m,k", MK_S)
 @pytest.mark.parametrize("top_k", TOP_KS)
 @pytest.mark.parametrize(

tests/kernels/test_fp32_router_gemm.py

@@ -3,9 +3,9 @@
 """Tests for fp32_router_gemm kernel: activation×weight→fp32.
 
 Supported (hidden_size, num_experts) pairs:
-  (3072, 256) -> MiniMax-M2/M2.5,  (6144, 128) -> MiniMax-M3
+  (3072, 256), (6144, 128), (5120, 256)
 
-Correctness baseline: F.linear in float32. Every M in [1, 32] is covered so
+Correctness baseline: F.linear in float32. Every supported M is covered so
 all tuned geometries (wide-block, experts-per-block, token-group; boundaries
 at M=4/5, odd/even, M=15/16) are exercised on Blackwell, and the legacy
 128/1 geometry everywhere else.
@@ -18,8 +18,12 @@ from vllm._custom_ops import fp32_router_gemm
 from vllm.platforms import current_platform
 
 # (hidden_size, num_experts)
-SHAPES = [(3072, 256), (6144, 128)]
-ALL_M = list(range(1, 33))
+SHAPES = [(3072, 256), (6144, 128), (5120, 256)]
+CASES = [
+    (m, hidden, experts)
+    for hidden, experts in SHAPES
+    for m in range(1, 49 if hidden == 5120 else 33)
+]
 # Absolute tolerance for fp32 kernel vs float64 reference
 ATOL_FP32 = 2e-4
 ATOL_BF16 = 2e-2  # bf16 activation has lower precision
@@ -40,8 +44,7 @@ def _ref(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor:
     return torch.nn.functional.linear(mat_a.float(), mat_b.float())
 
 
-@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES)
-@pytest.mark.parametrize("num_tokens", ALL_M)
+@pytest.mark.parametrize("num_tokens,hidden_dim,num_experts", CASES)
 def test_fp32_activation(num_tokens: int, hidden_dim: int, num_experts: int):
     """fp32 activation → fp32 output should match reference closely."""
     _requires_sm90()
@@ -58,8 +61,7 @@ def test_fp32_activation(num_tokens: int, hidden_dim: int, num_experts: int):
     torch.testing.assert_close(out, ref, atol=ATOL_FP32, rtol=0)
 
 
-@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES)
-@pytest.mark.parametrize("num_tokens", ALL_M)
+@pytest.mark.parametrize("num_tokens,hidden_dim,num_experts", CASES)
 def test_bf16_activation(num_tokens: int, hidden_dim: int, num_experts: int):
     """bf16 activation → fp32 output should match reference within bf16 error."""
     _requires_sm90()
@@ -136,17 +138,19 @@ def test_zero_tokens_returns_empty():
     assert out.dtype == torch.float32
 
 
-def test_rejects_invalid_inputs():
+@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES)
+def test_rejects_invalid_inputs(hidden_dim, num_experts):
     """The entry must fail loudly, never compute silently wrong results."""
     _requires_sm90()
     device = torch.device("cuda")
-    hidden_dim, num_experts = SHAPES[0]
     mat_b = torch.randn(num_experts, hidden_dim, dtype=torch.float32, device=device)
 
-    # num_tokens > 32 (beyond the instantiated range)
+    # Beyond the instantiated range for this shape.
+    invalid_tokens = 49 if hidden_dim == 5120 else 33
     with pytest.raises(Exception, match="num_tokens"):
         fp32_router_gemm(
-            torch.randn(33, hidden_dim, dtype=torch.float32, device=device), mat_b
+            torch.randn(invalid_tokens, hidden_dim, dtype=torch.float32, device=device),
+            mat_b,
         )
 
     # unsupported (hidden_dim, num_experts) pair
@@ -173,3 +177,36 @@ def test_rejects_invalid_inputs():
         fp32_router_gemm(
             torch.randn(4, hidden_dim, dtype=torch.float16, device=device), mat_b
         )
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize(
+    "tokens", [*range(1, 17), 17, 24, 32, 33, 40, 48, 56, 64, 65, 128]
+)
+@torch.inference_mode()
+def test_router_retains_fp32_weights_and_logits(default_vllm_config, dist_init, tokens):
+    """Low-latency routing must retain FP32 weight precision and logits."""
+    from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
+    from vllm.platforms import current_platform
+
+    if not current_platform.is_device_capability_family(100):
+        pytest.skip("SM100 required")
+    torch.manual_seed(42)
+    router = GateLinear(
+        5120, 256, params_dtype=torch.float32, out_dtype=torch.float32
+    ).cuda()
+    router.weight.normal_(std=0.01)
+    x = torch.randn(tokens, 5120, device="cuda", dtype=torch.bfloat16)
+    expected = (x.double() @ router.weight.double().T).float()
+    actual, bias = router(x)
+    assert bias is None and actual.dtype == torch.float32
+    torch.testing.assert_close(actual, expected, rtol=2e-5, atol=2e-5)
+    torch.testing.assert_close(actual.topk(8).indices, expected.topk(8).indices)
+    if tokens in (1, 16, 48, 56):
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            captured, _ = router(x)
+        x.normal_()
+        graph.replay()
+        expected = (x.double() @ router.weight.double().T).float()
+        torch.testing.assert_close(captured, expected, rtol=2e-5, atol=2e-5)

tests/kernels/test_ll_bf16_gemm.py

@@ -566,7 +566,7 @@ def test_invalid_output_dtype():
     with pytest.raises(ValueError, match="output_dtype=torch.float32"):
         from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ll_bf16_gemm
 
-        ll_bf16_gemm(a, b, output_dtype=torch.bfloat16)
+        ll_bf16_gemm(a, b, output_dtype=torch.float16)
 
 
 def test_cache_miss_compiles_dotprod():
@@ -581,5 +581,29 @@ def test_cache_miss_compiles_dotprod():
     _assert_close(out, _ref(a, b), context="cache miss dotprod")
 
 
+@pytest.mark.parametrize("M", [1, 2, 8, 16])
+@pytest.mark.parametrize("N", [1792, 3584])
+@torch.inference_mode()
+def test_bf16_output_and_fp32_cache_coexist(M, N):
+    """BF16 epilogues retain FP32 accumulation and can replay with new inputs."""
+    from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ll_bf16_gemm
+
+    torch.manual_seed(42)
+    x = torch.randn(M, 5120, device="cuda", dtype=torch.bfloat16)
+    w = torch.randn(N, 5120, device="cuda", dtype=torch.bfloat16) * 0.01
+    for dtype in (torch.float32, torch.bfloat16):
+        actual = ll_bf16_gemm(x, w, output_dtype=dtype)
+        assert actual.dtype == dtype
+        reference = (x.double() @ w.double().T).to(dtype)
+        torch.testing.assert_close(actual, reference, rtol=0.008, atol=1e-5)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        actual = ll_bf16_gemm(x, w, output_dtype=torch.bfloat16)
+    x.normal_()
+    graph.replay()
+    reference = (x.double() @ w.double().T).bfloat16()
+    torch.testing.assert_close(actual, reference, rtol=0.008, atol=1e-5)
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])