From 39983a5e7a5a575a7829f5d5ada50a33d9337ec1 Mon Sep 17 00:00:00 2001 From: Jee Jee Li Date: Sat, 17 Oct 2026 13:17:55 +0800 Subject: [PATCH] [Kernel] Add FP32 router GEMM support for Step-5 Preview and Muse Spark Open (#62126) Signed-off-by: Jee Jee Li --- csrc/fp32_router_gemm.cu | 29 ++++++++++++++++++- csrc/fp32_router_gemm_entry.cu | 12 ++++++-- tests/kernels/test_fp32_router_gemm.py | 23 +++++++++------ .../layers/fused_moe/router/gate_linear.py | 8 ++++- 4 files changed, 59 insertions(+), 13 deletions(-) diff --git a/csrc/fp32_router_gemm.cu b/csrc/fp32_router_gemm.cu index 832874f914..d5784ff0e6 100644 --- a/csrc/fp32_router_gemm.cu +++ b/csrc/fp32_router_gemm.cu @@ -279,6 +279,27 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a, launchFp32RouterGemm( output, mat_a, mat_b, stream); } + } else if constexpr (std::is_same_v && + kNumExperts == 256 && kHiddenDim == 8192) { + // MuseSpark1x: the 8MB weight makes this shape FFMA-bound past M~10, so + // the wider block wins at every M and token groups clear the M >= 28 + // accumulator cliff. Thresholds are from a B300 sweep. BS=192/384 are not + // instantiable: VPT=8 requires kBlockSize to divide kHiddenDim / 8 = 1024. + if (!isBlackwellFamily()) { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + return; + } + if constexpr (kNumTokens >= 10 && kNumTokens % 2 == 0) { + launchFp32RouterGemm(output, mat_a, mat_b, stream); + } else if constexpr (kNumTokens >= 13) { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + } else { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + } } else { launchFp32RouterGemm( output, mat_a, mat_b, stream); @@ -287,7 +308,8 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a, // --------------------------------------------------------------------------- // Explicit instantiations for both input types. -// Supported (E, H) pairs: (256, 3072), (128, 6144), and (256, 5120). +// Supported (E, H) pairs: (256, 3072), (128, 6144), (256, 5120), +// (352, 4096) [Step-5], and (256, 8192) [Muse Spark]. // --------------------------------------------------------------------------- #define INSTANTIATE(T, M, E, H) \ @@ -336,6 +358,11 @@ INSTANTIATE_ALL(__nv_bfloat16, 128, 6144) INSTANTIATE_ALL(float, 256, 5120) INSTANTIATE_ALL(__nv_bfloat16, 256, 5120) +INSTANTIATE_ALL(float, 352, 4096) +INSTANTIATE_ALL(__nv_bfloat16, 352, 4096) +INSTANTIATE_ALL(float, 256, 8192) +INSTANTIATE_ALL(__nv_bfloat16, 256, 8192) + #define INSTANTIATE_LARGE(T, E, H) \ INSTANTIATE(T, 33, E, H) \ INSTANTIATE(T, 34, E, H) \ diff --git a/csrc/fp32_router_gemm_entry.cu b/csrc/fp32_router_gemm_entry.cu index a6efc32831..4a83dbd73a 100644 --- a/csrc/fp32_router_gemm_entry.cu +++ b/csrc/fp32_router_gemm_entry.cu @@ -29,7 +29,9 @@ static constexpr int FP32_MAX_TOKENS = 32; static inline bool fp32_router_gemm_supported(int hidden_dim, int num_experts) { return (hidden_dim == 3072 && num_experts == 256) || (hidden_dim == 6144 && num_experts == 128) || - (hidden_dim == 5120 && num_experts == 256); + (hidden_dim == 5120 && num_experts == 256) || + (hidden_dim == 4096 && num_experts == 352) || + (hidden_dim == 8192 && num_experts == 256); } // Forward declarations — 4 template params must match fp32_router_gemm.cu @@ -80,6 +82,12 @@ void dispatchFp32RouterGemm(int num_experts, int hidden_dim, int num_tokens, } else if (num_experts == 256 && hidden_dim == 5120) { Fp32LoopUnroller::unroll(num_tokens, output, mat_a, mat_b, stream); + } else if (num_experts == 352 && hidden_dim == 4096) { + Fp32LoopUnroller::unroll( + num_tokens, output, mat_a, mat_b, stream); + } else if (num_experts == 256 && hidden_dim == 8192) { + Fp32LoopUnroller::unroll( + num_tokens, output, mat_a, mat_b, stream); } else { throw std::invalid_argument( "fp32_router_gemm: unsupported (hidden_dim, num_experts) pair"); @@ -114,7 +122,7 @@ void fp32_router_gemm( STD_TORCH_CHECK( fp32_router_gemm_supported(hidden_dim, num_experts), "fp32_router_gemm: supported (hidden_dim, num_experts) pairs are " - "(3072, 256), (6144, 128), and (5120, 256)"); + "(3072, 256), (6144, 128), (5120, 256), (4096, 352), and (8192, 256)"); const int max_tokens = hidden_dim == 5120 && num_experts == 256 ? 48 : FP32_MAX_TOKENS; STD_TORCH_CHECK(num_tokens <= max_tokens, diff --git a/tests/kernels/test_fp32_router_gemm.py b/tests/kernels/test_fp32_router_gemm.py index 0f55d860a6..54331dc3d3 100644 --- a/tests/kernels/test_fp32_router_gemm.py +++ b/tests/kernels/test_fp32_router_gemm.py @@ -3,11 +3,11 @@ """Tests for fp32_router_gemm kernel: activation×weight→fp32. Supported (hidden_size, num_experts) pairs: - (3072, 256), (6144, 128), (5120, 256) + (3072, 256), (6144, 128), (5120, 256), (4096, 352), (8192, 256) 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 +at M=4/5, M=9/10, M=12/13, odd/even, M=15/16) are exercised on Blackwell, and the legacy 128/1 geometry everywhere else. """ @@ -18,7 +18,7 @@ from vllm._custom_ops import fp32_router_gemm from vllm.platforms import current_platform # (hidden_size, num_experts) -SHAPES = [(3072, 256), (6144, 128), (5120, 256)] +SHAPES = [(3072, 256), (6144, 128), (5120, 256), (4096, 352), (8192, 256)] CASES = [ (m, hidden, experts) for hidden, experts in SHAPES @@ -126,11 +126,11 @@ def test_topk_routing_consistency(num_tokens: int, hidden_dim: int, num_experts: ) -def test_zero_tokens_returns_empty(): +@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES) +def test_zero_tokens_returns_empty(hidden_dim, num_experts): """M=0 is a graceful no-op returning an empty [0, E] tensor.""" _requires_sm90() device = torch.device("cuda") - hidden_dim, num_experts = SHAPES[0] mat_a = torch.empty(0, hidden_dim, dtype=torch.float32, device=device) mat_b = torch.randn(num_experts, hidden_dim, dtype=torch.float32, device=device) out = fp32_router_gemm(mat_a, mat_b) @@ -183,8 +183,13 @@ def test_rejects_invalid_inputs(hidden_dim, num_experts): @pytest.mark.parametrize( "tokens", [*range(1, 17), 17, 24, 32, 33, 40, 48, 56, 64, 65, 128] ) +@pytest.mark.parametrize( + "hidden_dim,num_experts", [(5120, 256), (4096, 352), (8192, 256)] +) @torch.inference_mode() -def test_router_retains_fp32_weights_and_logits(default_vllm_config, dist_init, tokens): +def test_router_retains_fp32_weights_and_logits( + default_vllm_config, dist_init, tokens, hidden_dim, num_experts +): """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 @@ -193,16 +198,16 @@ def test_router_retains_fp32_weights_and_logits(default_vllm_config, dist_init, pytest.skip("SM100 required") torch.manual_seed(42) router = GateLinear( - 5120, 256, params_dtype=torch.float32, out_dtype=torch.float32 + hidden_dim, num_experts, 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) + x = torch.randn(tokens, hidden_dim, 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): + if tokens in (1, 16, 32, 33, 48, 56): graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): captured, _ = router(x) diff --git a/vllm/model_executor/layers/fused_moe/router/gate_linear.py b/vllm/model_executor/layers/fused_moe/router/gate_linear.py index 7d14f105c9..3237783ece 100644 --- a/vllm/model_executor/layers/fused_moe/router/gate_linear.py +++ b/vllm/model_executor/layers/fused_moe/router/gate_linear.py @@ -47,7 +47,13 @@ class GateLinear(ReplicatedLinear): """ # (hidden_size, num_experts) pairs with an instantiated fp32 kernel. - FP32_SUPPORTED_SHAPES = {(3072, 256), (6144, 128), (5120, 256)} + FP32_SUPPORTED_SHAPES = { + (3072, 256), + (6144, 128), + (5120, 256), + (4096, 352), + (8192, 256), + } FP32_MAX_TOKENS = 32 # Largest batch for which tier 1 beats the cuBLAS epilogue. Subclasses # whose router shape stays ahead further up the range may raise it.