From 596e53067548fab65de6fc3abdcd020f752e7ae7 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Fri, 16 Oct 2026 15:39:46 -0700 Subject: [PATCH] [Kernel] Add router GEMM kernels for (5120, 256) and BF16 output to ll_bf16_gemm (#62042) Signed-off-by: Woosuk Kwon Co-authored-by: Lain Co-authored-by: Codex Co-authored-by: Claude Opus 5.5 Co-authored-by: vllm-agent --- csrc/fp32_router_gemm.cu | 54 ++++++++++-- csrc/fp32_router_gemm_entry.cu | 19 +++-- tests/kernels/moe/test_routing.py | 14 +++ tests/kernels/test_fp32_router_gemm.py | 61 ++++++++++--- tests/kernels/test_ll_bf16_gemm.py | 26 +++++- vllm/_custom_ops.py | 5 ++ .../linear/cute_dsl/_ll_bf16_dotprod.py | 12 +-- .../linear/cute_dsl/_ll_bf16_splitk.py | 4 +- .../linear/cute_dsl/_ll_router_common.py | 12 ++- .../kernels/linear/cute_dsl/ll_bf16.py | 85 +++++++++++++++---- .../router/bf16x3_router_gemm_cutedsl.py | 5 +- .../layers/fused_moe/router/gate_linear.py | 25 ++++-- 12 files changed, 258 insertions(+), 64 deletions(-) diff --git a/csrc/fp32_router_gemm.cu b/csrc/fp32_router_gemm.cu index 9719adb513..832874f914 100644 --- a/csrc/fp32_router_gemm.cu +++ b/csrc/fp32_router_gemm.cu @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project // -// Router GEMM: activation(T) x weight(fp32) -> fp32, M<=32, for the +// Router GEMM: activation(T) x weight(fp32) -> fp32, M<=48, for the // supported (E, H) pairs listed at the bottom of this file. // Supports bf16 or fp32 activation; weight is always fp32. @@ -214,8 +214,26 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a, // production-fidelity harness (CUDA-graph replay, per-layer cold weights). // Only enabled on the Blackwell family where it was validated; Hopper and // other shapes / fp32 activation keep the legacy geometry. - if constexpr (std::is_same_v && kNumExperts == 128 && - kHiddenDim == 6144) { + if constexpr (std::is_same_v && kNumExperts == 256 && + kHiddenDim == 5120) { + if (!isBlackwellFamily()) { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + return; + } + if constexpr (kNumTokens > 16 && kNumTokens % 8 == 0) { + constexpr int block_size = kNumTokens == 24 ? 160 : 128; + launchFp32RouterGemm(output, mat_a, mat_b, stream); + } else { + constexpr int block_size = kNumTokens == 1 ? 64 : 128; + constexpr int token_groups = kNumTokens % 2 == 0 ? 2 : 1; + launchFp32RouterGemm(output, mat_a, mat_b, + stream); + } + } else if constexpr (std::is_same_v && + kNumExperts == 128 && kHiddenDim == 6144) { // MiniMax-M3. Legacy 128/1 only fills 128 blocks and pays the same // accumulator register cliffs; B300 sweep: // even M in [6, 10] : BS=384, EPB=1, 2 token groups (1.26-1.43x) @@ -268,8 +286,8 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a, } // --------------------------------------------------------------------------- -// Explicit instantiations: M=1..32, for both input types, for the supported -// (E, H) pairs: (256, 3072) [MiniMax-M2/M2.5] and (128, 6144) [MiniMax-M3]. +// Explicit instantiations for both input types. +// Supported (E, H) pairs: (256, 3072), (128, 6144), and (256, 5120). // --------------------------------------------------------------------------- #define INSTANTIATE(T, M, E, H) \ @@ -315,5 +333,31 @@ INSTANTIATE_ALL(__nv_bfloat16, 256, 3072) INSTANTIATE_ALL(float, 128, 6144) INSTANTIATE_ALL(__nv_bfloat16, 128, 6144) +INSTANTIATE_ALL(float, 256, 5120) +INSTANTIATE_ALL(__nv_bfloat16, 256, 5120) + +#define INSTANTIATE_LARGE(T, E, H) \ + INSTANTIATE(T, 33, E, H) \ + INSTANTIATE(T, 34, E, H) \ + INSTANTIATE(T, 35, E, H) \ + INSTANTIATE(T, 36, E, H) \ + INSTANTIATE(T, 37, E, H) \ + INSTANTIATE(T, 38, E, H) \ + INSTANTIATE(T, 39, E, H) \ + INSTANTIATE(T, 40, E, H) \ + INSTANTIATE(T, 41, E, H) \ + INSTANTIATE(T, 42, E, H) \ + INSTANTIATE(T, 43, E, H) \ + INSTANTIATE(T, 44, E, H) \ + INSTANTIATE(T, 45, E, H) \ + INSTANTIATE(T, 46, E, H) \ + INSTANTIATE(T, 47, E, H) \ + INSTANTIATE(T, 48, E, H) + +INSTANTIATE_LARGE(float, 256, 5120) +INSTANTIATE_LARGE(__nv_bfloat16, 256, 5120) + +#undef INSTANTIATE_LARGE + #undef INSTANTIATE_ALL #undef INSTANTIATE diff --git a/csrc/fp32_router_gemm_entry.cu b/csrc/fp32_router_gemm_entry.cu index e55ac7e721..a6efc32831 100644 --- a/csrc/fp32_router_gemm_entry.cu +++ b/csrc/fp32_router_gemm_entry.cu @@ -25,10 +25,11 @@ inline int getSMVersion() { static constexpr int FP32_MAX_TOKENS = 32; // Supported (hidden_dim, num_experts) pairs (must match the instantiations in -// fp32_router_gemm.cu): (3072, 256) for MiniMax-M2/M2.5, (6144, 128) for M3. +// fp32_router_gemm.cu). 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 == 6144 && num_experts == 128) || + (hidden_dim == 5120 && num_experts == 256); } // Forward declarations — 4 template params must match fp32_router_gemm.cu @@ -60,8 +61,7 @@ struct Fp32LoopUnroller { invokeFp32RouterGemm( output, mat_a, mat_b, stream); } else { - throw std::invalid_argument( - "fp32_router_gemm: num_tokens must be in [1, 32]"); + throw std::invalid_argument("fp32_router_gemm: unsupported num_tokens"); } } }; @@ -77,6 +77,9 @@ void dispatchFp32RouterGemm(int num_experts, int hidden_dim, int num_tokens, } else if (num_experts == 128 && hidden_dim == 6144) { Fp32LoopUnroller::unroll( num_tokens, output, mat_a, mat_b, stream); + } else if (num_experts == 256 && hidden_dim == 5120) { + Fp32LoopUnroller::unroll(num_tokens, output, + mat_a, mat_b, stream); } else { throw std::invalid_argument( "fp32_router_gemm: unsupported (hidden_dim, num_experts) pair"); @@ -111,9 +114,11 @@ 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) and (6144, 128)"); - STD_TORCH_CHECK(num_tokens <= FP32_MAX_TOKENS, - "fp32_router_gemm: num_tokens must be in [0, 32]"); + "(3072, 256), (6144, 128), and (5120, 256)"); + const int max_tokens = + hidden_dim == 5120 && num_experts == 256 ? 48 : FP32_MAX_TOKENS; + STD_TORCH_CHECK(num_tokens <= max_tokens, + "fp32_router_gemm: num_tokens exceeds supported range"); STD_TORCH_CHECK( mat_a.scalar_type() == torch::headeronly::ScalarType::Float || mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16, diff --git a/tests/kernels/moe/test_routing.py b/tests/kernels/moe/test_routing.py index b0c3e26db9..77eac550d9 100644 --- a/tests/kernels/moe/test_routing.py +++ b/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( diff --git a/tests/kernels/test_fp32_router_gemm.py b/tests/kernels/test_fp32_router_gemm.py index d330abba53..0f55d860a6 100644 --- a/tests/kernels/test_fp32_router_gemm.py +++ b/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) diff --git a/tests/kernels/test_ll_bf16_gemm.py b/tests/kernels/test_ll_bf16_gemm.py index 1c7a15d514..76ea03f60e 100644 --- a/tests/kernels/test_ll_bf16_gemm.py +++ b/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"]) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 70fed78573..d74bb67816 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2424,6 +2424,11 @@ def grouped_topk( raise NotImplementedError( "The fused grouped_topk kernel is only available on CUDA and XPU platforms" ) + if scores.shape[0] == 0: + return ( + scores.new_empty((0, topk), dtype=torch.float32), + scores.new_empty((0, topk), dtype=torch.int32), + ) return torch.ops._moe_C.grouped_topk( scores, num_expert_group, diff --git a/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py b/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py index f06de6bd0c..707e2cdf45 100644 --- a/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py +++ b/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py @@ -356,10 +356,12 @@ class LLBf16Dotprod: if tidx == 0: for m in cutlass.range_constexpr(M): partials = sm[m, None].load() - gC[m, n_idx] = partials.reduce( - cute.ReductionOp.ADD, - init_val=cutlass.Float32(0.0), - reduction_profile=0, - ) + gC[m, n_idx] = cutlass.Float32( + partials.reduce( + cute.ReductionOp.ADD, + init_val=cutlass.Float32(0.0), + reduction_profile=0, + ) + ).to(gC.element_type) if const_expr(self.use_pdl): cute.arch.griddepcontrol_launch_dependents() diff --git a/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py b/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py index b44a5f6e2c..2ae3d59922 100644 --- a/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py +++ b/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py @@ -47,7 +47,7 @@ class LLBf16SplitK: :param ab_dtype: Element type for A and B operands. :param acc_dtype: Accumulator type used by MMA and reductions. - :param out_dtype: Output element type. The public wrapper uses fp32. + :param out_dtype: Output element type (fp32 or bf16). :param tile_n: CTA tile size in N. M is fixed to 16 for router batches. :type tile_n: int :param tile_k: K tile size staged through shared memory. @@ -580,6 +580,6 @@ class LLBf16SplitK: reduction_profile=0, ) ) - gC[local_coord] = acc + gC[local_coord] = cutlass.Float32(acc).to(gC.element_type) cute.arch.sync_threads() diff --git a/vllm/model_executor/kernels/linear/cute_dsl/_ll_router_common.py b/vllm/model_executor/kernels/linear/cute_dsl/_ll_router_common.py index 23fcefdac9..7770c0f8de 100644 --- a/vllm/model_executor/kernels/linear/cute_dsl/_ll_router_common.py +++ b/vllm/model_executor/kernels/linear/cute_dsl/_ll_router_common.py @@ -76,8 +76,8 @@ def make_fake_gemm_tensors( a_dtype: torch.dtype, b_dtype: torch.dtype, divisibility: int, + output_dtype: torch.dtype = torch.float32, ): - from cutlass import Float32 from quack.compile_utils import make_fake_tensor hidden_states = make_fake_tensor( @@ -86,7 +86,7 @@ def make_fake_gemm_tensors( router_weight = make_fake_tensor( cutlass_dtype(b_dtype), (N, K), divisibility=divisibility ) - output = make_fake_tensor(Float32, (M, N), divisibility=1) + output = make_fake_tensor(cutlass_dtype(output_dtype), (M, N), divisibility=1) return hidden_states, router_weight, output @@ -97,6 +97,7 @@ def validate_common_gemm_inputs( *, op_name: str, k_multiple: int | None = None, + output_dtypes: tuple[torch.dtype, ...] = (torch.float32,), ) -> None: if hidden_states.dim() != 2 or router_weight.dim() != 2: raise ValueError("hidden_states and router_weight must be 2D tensors") @@ -106,8 +107,11 @@ def validate_common_gemm_inputs( raise ValueError( "hidden_states and router_weight must be on the same CUDA device" ) - if output_dtype != torch.float32: - raise ValueError(f"{op_name} only supports output_dtype=torch.float32") + if output_dtype not in output_dtypes: + raise ValueError( + f"{op_name} only supports output_dtype=" + + " or ".join(str(dtype) for dtype in output_dtypes) + ) if hidden_states.shape[1] != router_weight.shape[1]: raise ValueError( "hidden_states and router_weight must have matching K dimensions" diff --git a/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py b/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py index 95cb78c6a1..9d1573bb80 100644 --- a/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py +++ b/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py @@ -47,6 +47,8 @@ _SM100F_TUNED_SPLITK_CONFIGS: dict[tuple[int, int], dict[int, tuple[int, int]]] **{M: (8, 5) for M in (5, 8)}, 9: (8, 2), }, + (5120, 1792): {M: (2, 4) for M in range(5, 17)}, + (5120, 3584): {M: (2, 4) for M in range(5, 17)}, (7168, 256): {14: (8, 2)}, (6144, 256): {M: (8, 2) for M in (9, 12, 16)}, (7168, 384): {M: (7, 5) for M in (13, 16)}, @@ -101,39 +103,52 @@ class LLBf16Gemm: bs: int = 0 split_k: int = 0 num_stages: int = 0 + output_dtype: torch.dtype = torch.float32 def __init__(self, *, prefetch_pdl_weights: bool = False) -> None: self._prefetch_pdl_weights = prefetch_pdl_weights - # Dot-prod: keyed on (M, K, bs), because M and K are Constexpr. - self._compiled_cache: dict[tuple[int, int, int], Any] = {} - # Split-K: keyed on (split_k, num_stages), fully shape-dynamic. - self._splitk_cache: dict[tuple[int, int], Any] = {} + # M and K are compile-time constants for dot-product. + self._compiled_cache: dict[tuple[int, int, int, torch.dtype], Any] = {} + # Split-K is shape-dynamic; both caches distinguish output dtype. + self._splitk_cache: dict[tuple[int, int, torch.dtype], Any] = {} - def dispatch(self, *, M: int, K: int, N: int) -> CompileKey: + def dispatch( + self, *, M: int, K: int, N: int, output_dtype: torch.dtype = torch.float32 + ) -> CompileKey: tuned_bs, tuned_splitk = _arch_tuned_configs() if M <= _DEFAULT_DOTPROD_MAX_M or K < 2048: bs = tuned_bs.get((K, N), {}).get(M, _DEFAULT_DOTPROD_BS) - return self.CompileKey(backend="dotprod", M=M, K=K, bs=bs) + return self.CompileKey( + backend="dotprod", M=M, K=K, bs=bs, output_dtype=output_dtype + ) split_k, num_stages = tuned_splitk.get((K, N), {}).get( M, _DEFAULT_SPLITK_CONFIG ) - return self.CompileKey(backend="splitk", split_k=split_k, num_stages=num_stages) + return self.CompileKey( + backend="splitk", + split_k=split_k, + num_stages=num_stages, + output_dtype=output_dtype, + ) def get_warmup_keys( self, *, shapes: Iterable[tuple[int, int]], m_values: Iterable[int], + output_dtype: torch.dtype = torch.float32, ) -> list[CompileKey]: return list( dict.fromkeys( - self.dispatch(M=M, K=K, N=N) for K, N in shapes for M in m_values + self.dispatch(M=M, K=K, N=N, output_dtype=output_dtype) + for K, N in shapes + for M in m_values ) ) @staticmethod - def _fake_gemm_tensors(*, M, K, N, divisibility: int): + def _fake_gemm_tensors(*, M, K, N, divisibility: int, output_dtype: torch.dtype): return make_fake_gemm_tensors( M=M, K=K, @@ -141,10 +156,13 @@ class LLBf16Gemm: a_dtype=torch.bfloat16, b_dtype=torch.bfloat16, divisibility=divisibility, + output_dtype=output_dtype, ) def _compile_splitk(self, compile_key: CompileKey) -> None: cute, _ = cute_context() + from cutlass import BFloat16, Float32 + from ._ll_bf16_splitk import LLBf16SplitK hidden_states, router_weight, output = self._fake_gemm_tensors( @@ -152,8 +170,12 @@ class LLBf16Gemm: K=cute.sym_int(), N=cute.sym_int(), divisibility=8, + output_dtype=compile_key.output_dtype, ) gemm = LLBf16SplitK( + out_dtype=BFloat16 + if compile_key.output_dtype == torch.bfloat16 + else Float32, tile_n=16, tile_k=256, num_stages=compile_key.num_stages, @@ -169,7 +191,9 @@ class LLBf16Gemm: current_cuda_stream(), options="--enable-tvm-ffi", ) - self._splitk_cache[(compile_key.split_k, compile_key.num_stages)] = compiled + self._splitk_cache[ + (compile_key.split_k, compile_key.num_stages, compile_key.output_dtype) + ] = compiled logger.debug( "Compiled ll_bf16_splitk: sk=%d ns=%d", compile_key.split_k, @@ -187,6 +211,7 @@ class LLBf16Gemm: K=compile_key.K, N=N, divisibility=stride_divisibility, + output_dtype=compile_key.output_dtype, ) gemm = LLBf16Dotprod( k=compile_key.K, @@ -205,7 +230,9 @@ class LLBf16Gemm: current_cuda_stream(), options="--enable-tvm-ffi --ptxas-options -maxrregcount=64", ) - self._compiled_cache[(compile_key.M, compile_key.K, compile_key.bs)] = compiled + self._compiled_cache[ + (compile_key.M, compile_key.K, compile_key.bs, compile_key.output_dtype) + ] = compiled logger.debug( "Compiled ll_bf16_dotprod: M=%d, K=%d, bs=%d", compile_key.M, @@ -215,12 +242,21 @@ class LLBf16Gemm: def compile(self, compile_key: CompileKey) -> None: if compile_key.backend == "splitk": - splitk_cache_key = (compile_key.split_k, compile_key.num_stages) + splitk_cache_key = ( + compile_key.split_k, + compile_key.num_stages, + compile_key.output_dtype, + ) if splitk_cache_key not in self._splitk_cache: self._compile_splitk(compile_key) return - dotprod_cache_key = (compile_key.M, compile_key.K, compile_key.bs) + dotprod_cache_key = ( + compile_key.M, + compile_key.K, + compile_key.bs, + compile_key.output_dtype, + ) if dotprod_cache_key not in self._compiled_cache: self._compile_dotprod(compile_key) @@ -229,8 +265,11 @@ class LLBf16Gemm: *, shapes: Iterable[tuple[int, int]], m_values: Iterable[int], + output_dtype: torch.dtype = torch.float32, ) -> None: - for compile_key in self.get_warmup_keys(shapes=shapes, m_values=m_values): + for compile_key in self.get_warmup_keys( + shapes=shapes, m_values=m_values, output_dtype=output_dtype + ): self.compile(compile_key) @staticmethod @@ -251,6 +290,7 @@ class LLBf16Gemm: output_dtype, op_name="ll_bf16_gemm", k_multiple=8, + output_dtypes=(torch.float32, torch.bfloat16), ) def __call__( @@ -258,19 +298,28 @@ class LLBf16Gemm: hidden_states: torch.Tensor, # [M, K] bf16 router_weight: torch.Tensor, # [N, K] bf16 output_dtype: torch.dtype = torch.float32, - ) -> torch.Tensor: # [M, N] fp32 + ) -> torch.Tensor: # [M, N] self._validate_inputs(hidden_states, router_weight, output_dtype) M, K = hidden_states.shape N = router_weight.shape[0] - compile_key = self.dispatch(M=M, K=K, N=N) + compile_key = self.dispatch(M=M, K=K, N=N, output_dtype=output_dtype) if compile_key.backend == "splitk": - splitk_cache_key = (compile_key.split_k, compile_key.num_stages) + splitk_cache_key = ( + compile_key.split_k, + compile_key.num_stages, + compile_key.output_dtype, + ) if splitk_cache_key not in self._splitk_cache: self.compile(compile_key) kernel = self._splitk_cache[splitk_cache_key] else: - dotprod_cache_key = (compile_key.M, compile_key.K, compile_key.bs) + dotprod_cache_key = ( + compile_key.M, + compile_key.K, + compile_key.bs, + compile_key.output_dtype, + ) if dotprod_cache_key not in self._compiled_cache: self.compile(compile_key) kernel = self._compiled_cache[dotprod_cache_key] diff --git a/vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py b/vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py index 2190eea157..395d7d68ea 100644 --- a/vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py +++ b/vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py @@ -517,6 +517,7 @@ _TILE_CONFIG_OVERRIDES = { 256: (128, 32), 512: (128, 19), }, + (5120, 256): {64: (64, 64)}, # GB200 (6144, 128): { # MiniMax-M3 64: (32, 64), 128: (64, 64), @@ -533,7 +534,7 @@ def _pick_tile_config(N: int, K: int, M: int, num_sms: int) -> tuple[int, int]: if 32 < N <= 512: token_bucket = triton.next_power_of_2(N) overrides = _TILE_CONFIG_OVERRIDES.get((K, M)) - if overrides is not None: + if overrides is not None and token_bucket in overrides: BN, split_k = overrides[token_bucket] return BN, min(split_k, k_tiles) @@ -609,6 +610,8 @@ def _bf16x3_router_gemm(X: torch.Tensor, W: torch.Tensor) -> torch.Tensor: """Return ``X @ W.T`` using the SM100 BF16x3 router GEMM kernel.""" N, K = X.shape M, _ = W.shape + if N == 0: + return X.new_empty((N, M), dtype=torch.float32) num_sms = torch.cuda.get_device_properties(X.device).multi_processor_count BN, split_k = _pick_tile_config(N, K, M, num_sms) 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 c6b1da890b..7d14f105c9 100644 --- a/vllm/model_executor/layers/fused_moe/router/gate_linear.py +++ b/vllm/model_executor/layers/fused_moe/router/gate_linear.py @@ -24,8 +24,8 @@ class GateLinear(ReplicatedLinear): 1. cuteDSL ll_bf16_gemm (SM90+, M<=16, bf16 in, fp32 out, any shape w/ K divisible by 8) - 2. fp32 specialized kernel (bf16/fp32 in, fp32 out, M<=32; SM90+ - model-specific shapes, gfx950 any shape) + 2. fp32 specialized kernel (bf16/fp32 in, fp32 out; SM90+ + model-specific shapes and batch limits, gfx950 any shape with M<=32) 3. cuteDSL ll_fp32w_gemm (SM90+, M<=32, bf16/fp32 in, fp32 out, any shape) 4. bf16x3 CuteDSL kernel (SM100, bf16 in, fp32 weight) @@ -46,9 +46,8 @@ class GateLinear(ReplicatedLinear): split-K from the token count and skip the batch-invariant linear path. """ - # (hidden_size, num_experts) pairs with an instantiated fp32 kernel: - # (3072, 256) -> MiniMax-M2/M2.5, (6144, 128) -> MiniMax-M3 - FP32_SUPPORTED_SHAPES = {(3072, 256), (6144, 128)} + # (hidden_size, num_experts) pairs with an instantiated fp32 kernel. + FP32_SUPPORTED_SHAPES = {(3072, 256), (6144, 128), (5120, 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. @@ -454,9 +453,9 @@ def fp32_router_gemm_dispatch_impl( weight: torch.Tensor, allow_bf16x3_router_gemm: bool, ) -> torch.Tensor: - """Dynamically run fp32 specialized gemm if num_tokens <= FP32_MAX_TOKENS, - otherwise optionally run the experimental BF16x3 kernel for medium/large - SM100 router batches, then fall back to F.linear. + """Run the FP32 kernel for small batches, with shape-specific limits. + + Larger SM100 BF16 batches use BF16x3; other inputs fall back to F.linear. This must be wrapped in a custom op because our torch.compile integration does not support runtime dispatching on num_tokens. """ @@ -473,7 +472,15 @@ def fp32_router_gemm_dispatch_impl( if weight.dtype == torch.bfloat16: return torch.mm(x, weight.T, out_dtype=torch.float32) return torch.nn.functional.linear(x.float(), weight) - if x.shape[0] <= _FP32_ROUTER_GEMM_MAX_TOKENS: + max_tokens = _FP32_ROUTER_GEMM_MAX_TOKENS + if ( + allow_bf16x3_router_gemm + and x.dtype == torch.bfloat16 + and weight.shape == (256, 5120) + ): + # GB200: four token groups win through 48 for multiples of eight. + max_tokens = 48 if x.shape[0] % 8 == 0 else 16 + if x.shape[0] <= max_tokens: return ops.fp32_router_gemm(x, weight) if allow_bf16x3_router_gemm and x.dtype == torch.bfloat16: