@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 559d6710bf
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 2, 09:45 UTC

[PERF]MiniMax-M2 gate kernel (#38445)

Signed-off-by: Jee Jee Li <pandaleefree@gmail.com> Signed-off-by: qianlihuang <91178480+qianlihuang@users.noreply.github.com> Co-authored-by: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com>

Jee Jee Li committed May 30, 2026 at 09:28 UTC 559d6710bf45e6fb3d48429855702b6ff24bdfe0
10 files changed +716 -23
CMakeLists.txt
+24 -10
@@ -683,6 +683,22 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
683 "in CUDA target architectures.")
684 endif()
685
686 + # FP32 router GEMM (H=3072, E=256, M<=32). Requires SM90+ and CUDA >= 12.0.
687 + cuda_archs_sm90plus(FP32_ROUTER_GEMM_ARCHS "${CUDA_ARCHS}")
688 + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND FP32_ROUTER_GEMM_ARCHS)
689 + set(SRCS
690 + "csrc/libtorch_stable/fp32_router_gemm_entry.cu"
691 + "csrc/libtorch_stable/fp32_router_gemm.cu")
692 + set_gencode_flags_for_srcs(
693 + SRCS "${SRCS}"
694 + CUDA_ARCHS "${FP32_ROUTER_GEMM_ARCHS}")
695 + list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
696 + message(STATUS "Building fp32_router_gemm for archs: ${FP32_ROUTER_GEMM_ARCHS}")
697 + else()
698 + message(STATUS "Not building fp32_router_gemm as no compatible archs found "
699 + "(requires SM90+ and CUDA >= 12.0).")
700 + endif()
701 +
702 # Only build AllSpark kernels if we are building for at least some compatible archs.
703 cuda_archs_loose_intersection(ALLSPARK_ARCHS "8.0;8.6;8.7;8.9" "${CUDA_ARCHS}")
704 if (ALLSPARK_ARCHS)
@@ -1240,24 +1256,22 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
1256 " in CUDA target architectures")
1257 endif()
1258
1243 - # DeepSeek V3 router GEMM kernel - requires SM90+
1244 - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
1245 - cuda_archs_loose_intersection(DSV3_ROUTER_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
1246 - else()
1247 - cuda_archs_loose_intersection(DSV3_ROUTER_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
1248 - endif()
1249 - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_ROUTER_GEMM_ARCHS)
1259 + # DeepSeek V3 router GEMM kernel requires SM90+ and CUDA >= 12.0.
1260 + # (fp32_router_gemm has been migrated to _C_stable_libtorch above.)
1261 + cuda_archs_sm90plus(SM90PLUS_ROUTER_GEMM_ARCHS "${CUDA_ARCHS}")
1262 + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND SM90PLUS_ROUTER_GEMM_ARCHS)
1263 set(DSV3_ROUTER_GEMM_SRC
1264 "csrc/moe/dsv3_router_gemm_entry.cu"
1265 "csrc/moe/dsv3_router_gemm_float_out.cu"
1266 "csrc/moe/dsv3_router_gemm_bf16_out.cu")
1267 set_gencode_flags_for_srcs(
1268 SRCS "${DSV3_ROUTER_GEMM_SRC}"
1256 - CUDA_ARCHS "${DSV3_ROUTER_GEMM_ARCHS}")
1269 + CUDA_ARCHS "${SM90PLUS_ROUTER_GEMM_ARCHS}")
1270 list(APPEND VLLM_MOE_EXT_SRC "${DSV3_ROUTER_GEMM_SRC}")
1258 - message(STATUS "Building DSV3 router GEMM kernel for archs: ${DSV3_ROUTER_GEMM_ARCHS}")
1271 +
1272 + message(STATUS "Building DSV3 router GEMM kernels for archs: ${SM90PLUS_ROUTER_GEMM_ARCHS}")
1273 else()
1260 - message(STATUS "Not building DSV3 router GEMM kernel as no compatible archs found"
1274 + message(STATUS "Not building DSV3 router GEMM kernels as no compatible archs found"
1275 " (requires SM90+ and CUDA >= 12.0)")
1276 endif()
1277 endif()
benchmarks/kernels/benchmark_router_gemm.py new
+154
@@ -0,0 +1,154 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +import torch.nn.functional as F
6 +
7 +from vllm import _custom_ops as ops
8 +from vllm.platforms import current_platform
9 +from vllm.transformers_utils.config import get_config
10 +from vllm.triton_utils import triton
11 +from vllm.utils.argparse_utils import FlexibleArgumentParser
12 +
13 +# Dimensions supported by the DSV3 specialized kernel
14 +DSV3_SUPPORTED_NUM_EXPERTS = [256, 384]
15 +DSV3_SUPPORTED_HIDDEN_SIZES = [7168]
16 +
17 +# Dimensions supported by the gpt-oss specialized kernel
18 +GPT_OSS_SUPPORTED_NUM_EXPERTS = [32, 128]
19 +GPT_OSS_SUPPORTED_HIDDEN_SIZES = [2880]
20 +
21 +# Dimensions supported by the fp32 specialized kernel (MiniMax-M2)
22 +FP32_SUPPORTED_NUM_EXPERTS = [256]
23 +FP32_SUPPORTED_HIDDEN_SIZES = [3072]
24 +FP32_MAX_TOKENS = 32
25 +
26 +
27 +def get_batch_size_range(max_batch_size):
28 + return [2**x for x in range(14) if 2**x <= max_batch_size]
29 +
30 +
31 +def get_model_params(config):
32 + if config.architectures[0] in (
33 + "DeepseekV2ForCausalLM",
34 + "DeepseekV3ForCausalLM",
35 + "DeepseekV32ForCausalLM",
36 + ):
37 + num_experts = config.n_routed_experts
38 + hidden_size = config.hidden_size
39 + elif config.architectures[0] in ("GptOssForCausalLM",) or config.architectures[
40 + 0
41 + ] in ("MiniMaxM2ForCausalLM",):
42 + num_experts = config.num_local_experts
43 + hidden_size = config.hidden_size
44 + else:
45 + raise ValueError(f"Unsupported architecture: {config.architectures}")
46 + return num_experts, hidden_size
47 +
48 +
49 +def get_benchmark(model, max_batch_size, trust_remote_code):
50 + @triton.testing.perf_report(
51 + triton.testing.Benchmark(
52 + x_names=["batch_size"],
53 + x_vals=get_batch_size_range(max_batch_size),
54 + x_log=False,
55 + line_arg="provider",
56 + line_vals=[
57 + "torch",
58 + "vllm",
59 + ],
60 + line_names=["PyTorch", "vLLM"],
61 + styles=([("blue", "-"), ("red", "-")]),
62 + ylabel="TFLOPs",
63 + plot_name=f"{model} router gemm throughput",
64 + args={},
65 + )
66 + )
67 + def benchmark(batch_size, provider):
68 + config = get_config(model=model, trust_remote_code=trust_remote_code)
69 + num_experts, hidden_size = get_model_params(config)
70 +
71 + is_hopper_or_blackwell = current_platform.is_device_capability(
72 + 90
73 + ) or current_platform.is_device_capability_family(100)
74 + allow_dsv3_router_gemm = (
75 + is_hopper_or_blackwell
76 + and num_experts in DSV3_SUPPORTED_NUM_EXPERTS
77 + and hidden_size in DSV3_SUPPORTED_HIDDEN_SIZES
78 + )
79 + allow_gpt_oss_router_gemm = (
80 + is_hopper_or_blackwell
81 + and num_experts in GPT_OSS_SUPPORTED_NUM_EXPERTS
82 + and hidden_size in GPT_OSS_SUPPORTED_HIDDEN_SIZES
83 + )
84 + is_fp32_router_model = (
85 + is_hopper_or_blackwell
86 + and num_experts in FP32_SUPPORTED_NUM_EXPERTS
87 + and hidden_size in FP32_SUPPORTED_HIDDEN_SIZES
88 + )
89 + allow_fp32_router_gemm = is_fp32_router_model and batch_size <= FP32_MAX_TOKENS
90 +
91 + # Weight dtype: fp32 kernel requires fp32 weights; others use bf16.
92 + weight_dtype = torch.float32 if is_fp32_router_model else torch.bfloat16
93 + mat_a = torch.randn(
94 + (batch_size, hidden_size), dtype=torch.bfloat16, device="cuda"
95 + ).contiguous()
96 + mat_b = torch.randn(
97 + (num_experts, hidden_size), dtype=weight_dtype, device="cuda"
98 + ).contiguous()
99 + bias = torch.randn(
100 + num_experts, dtype=torch.bfloat16, device="cuda"
101 + ).contiguous()
102 +
103 + has_bias = allow_gpt_oss_router_gemm
104 +
105 + quantiles = [0.5, 0.2, 0.8]
106 +
107 + if provider == "torch":
108 +
109 + def runner():
110 + if allow_fp32_router_gemm:
111 + F.linear(mat_a.float(), mat_b)
112 + elif has_bias:
113 + F.linear(mat_a, mat_b, bias)
114 + else:
115 + F.linear(mat_a, mat_b)
116 + elif provider == "vllm":
117 +
118 + def runner():
119 + if allow_dsv3_router_gemm:
120 + ops.dsv3_router_gemm(mat_a, mat_b, torch.bfloat16)
121 + elif allow_fp32_router_gemm:
122 + ops.fp32_router_gemm(mat_a, mat_b)
123 + elif allow_gpt_oss_router_gemm:
124 + ops.gpt_oss_router_gemm(mat_a, mat_b, bias)
125 + elif is_fp32_router_model:
126 + # batch_size > FP32_MAX_TOKENS: fall back to F.linear
127 + F.linear(mat_a.float(), mat_b)
128 + else:
129 + F.linear(mat_a, mat_b)
130 +
131 + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
132 + runner, quantiles=quantiles
133 + )
134 +
135 + def tflops(t_ms):
136 + flops = 2 * batch_size * hidden_size * num_experts
137 + return flops / (t_ms * 1e-3) / 1e12
138 +
139 + return tflops(ms), tflops(max_ms), tflops(min_ms)
140 +
141 + return benchmark
142 +
143 +
144 +if __name__ == "__main__":
145 + parser = FlexibleArgumentParser()
146 + parser.add_argument("--model", type=str, default="openai/gpt-oss-20b")
147 + parser.add_argument("--max-batch-size", default=16, type=int)
148 + parser.add_argument("--trust-remote-code", action="store_true")
149 + args = parser.parse_args()
150 +
151 + # Get the benchmark function
152 + benchmark = get_benchmark(args.model, args.max_batch_size, args.trust_remote_code)
153 + # Run performance benchmark
154 + benchmark.run(print_data=True)
cmake/utils.cmake
+10
@@ -476,6 +476,16 @@ function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_AR
476 set(${OUT_CUDA_ARCHS} ${_CUDA_ARCHS} PARENT_SCOPE)
477 endfunction()
478
479 +
480 +function(cuda_archs_sm90plus OUT_CUDA_ARCHS TGT_CUDA_ARCHS)
481 + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
482 + cuda_archs_loose_intersection(_archs "9.0a;10.0f;11.0f" "${TGT_CUDA_ARCHS}")
483 + else()
484 + cuda_archs_loose_intersection(_archs "9.0a;10.0a;10.1a;10.3a" "${TGT_CUDA_ARCHS}")
485 + endif()
486 + set(${OUT_CUDA_ARCHS} ${_archs} PARENT_SCOPE)
487 +endfunction()
488 +
489 #
490 # Override the GPU architectures detected by cmake/torch and filter them by
491 # `GPU_SUPPORTED_ARCHES`. Sets the final set of architectures in
csrc/libtorch_stable/fp32_router_gemm.cu new
+223
@@ -0,0 +1,223 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +//
4 +// Router GEMM: activation(T) x weight(fp32) -> fp32, H=3072, E=256, M<=32.
5 +// Supports bf16 or fp32 activation; weight is always fp32.
6 +// Adapted from dsv3_router_gemm_float_out.cu.
7 +
8 +#include <cuda_bf16.h>
9 +#include <cuda_runtime.h>
10 +
11 +// ---------------------------------------------------------------------------
12 +// Load helpers
13 +// ---------------------------------------------------------------------------
14 +
15 +// Load VPT fp32 values from the weight matrix (always fp32).
16 +// VPT=4 when activation is fp32 (one float4 load)
17 +// VPT=8 when activation is bf16 (two float4 loads)
18 +template <int VPT>
19 +__device__ __forceinline__ void load_weight(float const* ptr, float* dst);
20 +
21 +template <>
22 +__device__ __forceinline__ void load_weight<4>(float const* ptr, float* dst) {
23 + float4 v = *reinterpret_cast<float4 const*>(ptr);
24 + dst[0] = v.x;
25 + dst[1] = v.y;
26 + dst[2] = v.z;
27 + dst[3] = v.w;
28 +}
29 +
30 +template <>
31 +__device__ __forceinline__ void load_weight<8>(float const* ptr, float* dst) {
32 + float4 v0 = *reinterpret_cast<float4 const*>(ptr);
33 + float4 v1 = *reinterpret_cast<float4 const*>(ptr + 4);
34 + dst[0] = v0.x;
35 + dst[1] = v0.y;
36 + dst[2] = v0.z;
37 + dst[3] = v0.w;
38 + dst[4] = v1.x;
39 + dst[5] = v1.y;
40 + dst[6] = v1.z;
41 + dst[7] = v1.w;
42 +}
43 +
44 +// Load VPT activation values and convert to fp32.
45 +template <typename T, int VPT>
46 +__device__ __forceinline__ void load_activation(T const* ptr, float* dst);
47 +
48 +// fp32 activation: one float4 load, no conversion needed.
49 +template <>
50 +__device__ __forceinline__ void load_activation<float, 4>(float const* ptr,
51 + float* dst) {
52 + float4 v = *reinterpret_cast<float4 const*>(ptr);
53 + dst[0] = v.x;
54 + dst[1] = v.y;
55 + dst[2] = v.z;
56 + dst[3] = v.w;
57 +}
58 +
59 +// bf16 activation: one uint4 load (8 × bf16) + element-wise conversion.
60 +template <>
61 +__device__ __forceinline__ void load_activation<__nv_bfloat16, 8>(
62 + __nv_bfloat16 const* ptr, float* dst) {
63 + uint4 v = *reinterpret_cast<uint4 const*>(ptr);
64 + __nv_bfloat16 const* bf16_ptr = reinterpret_cast<__nv_bfloat16 const*>(&v);
65 +#pragma unroll
66 + for (int i = 0; i < 8; i++) dst[i] = __bfloat162float(bf16_ptr[i]);
67 +}
68 +
69 +// ---------------------------------------------------------------------------
70 +// Kernel
71 +// ---------------------------------------------------------------------------
72 +
73 +// InputT : type of activation (float or __nv_bfloat16)
74 +// Weight is always fp32; output is always fp32.
75 +// VPT = 16 / sizeof(InputT): 4 for fp32, 8 for bf16
76 +template <typename InputT, int kBlockSize, int kNumTokens, int kNumExperts,
77 + int kHiddenDim>
78 +__global__ __launch_bounds__(128, 1) void fp32_router_gemm_kernel(
79 + float* out, InputT const* mat_a, float const* mat_b) {
80 + constexpr int VPT = 16 / sizeof(InputT);
81 + constexpr int k_elems_per_k_iteration = VPT * kBlockSize;
82 + constexpr int k_iterations = kHiddenDim / k_elems_per_k_iteration;
83 + constexpr int kWarpSize = 32;
84 + constexpr int kNumWarps = kBlockSize / kWarpSize;
85 +
86 + int const n_idx = blockIdx.x;
87 + int const tid = threadIdx.x;
88 + int const warpId = tid / kWarpSize;
89 + int const laneId = tid % kWarpSize;
90 +
91 + float acc[kNumTokens] = {};
92 + __shared__ float sm_reduction[kNumTokens][kNumWarps];
93 +
94 + float const* b_col = mat_b + n_idx * kHiddenDim;
95 +
96 + int k_bases[k_iterations];
97 +#pragma unroll
98 + for (int ki = 0; ki < k_iterations; ki++) {
99 + k_bases[ki] = ki * k_elems_per_k_iteration + tid * VPT;
100 + }
101 +
102 +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
103 + asm volatile("griddepcontrol.wait;");
104 +#endif
105 +
106 + for (int ki = 0; ki < k_iterations; ki++) {
107 + int const k_base = k_bases[ki];
108 +
109 + float b_float[VPT];
110 + load_weight<VPT>(b_col + k_base, b_float);
111 +
112 +#pragma unroll
113 + for (int m_idx = 0; m_idx < kNumTokens; m_idx++) {
114 + float a_float[VPT];
115 + load_activation<InputT, VPT>(mat_a + m_idx * kHiddenDim + k_base,
116 + a_float);
117 +#pragma unroll
118 + for (int k = 0; k < VPT; k++) {
119 + acc[m_idx] += a_float[k] * b_float[k];
120 + }
121 + }
122 + }
123 +
124 + // Warp-level butterfly reduction
125 +#pragma unroll
126 + for (int m = 0; m < kNumTokens; m++) {
127 + float sum = acc[m];
128 + sum += __shfl_xor_sync(0xffffffff, sum, 16);
129 + sum += __shfl_xor_sync(0xffffffff, sum, 8);
130 + sum += __shfl_xor_sync(0xffffffff, sum, 4);
131 + sum += __shfl_xor_sync(0xffffffff, sum, 2);
132 + sum += __shfl_xor_sync(0xffffffff, sum, 1);
133 + if (laneId == 0) sm_reduction[m][warpId] = sum;
134 + }
135 +
136 + __syncthreads();
137 +
138 + if (tid == 0) {
139 +#pragma unroll
140 + for (int m = 0; m < kNumTokens; m++) {
141 + float final_sum = 0.0f;
142 +#pragma unroll
143 + for (int w = 0; w < kNumWarps; w++) final_sum += sm_reduction[m][w];
144 + out[m * kNumExperts + n_idx] = final_sum;
145 + }
146 + }
147 +
148 +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
149 + asm volatile("griddepcontrol.launch_dependents;");
150 +#endif
151 +}
152 +
153 +// ---------------------------------------------------------------------------
154 +// Launcher
155 +// ---------------------------------------------------------------------------
156 +
157 +template <typename InputT, int kNumTokens, int kNumExperts, int kHiddenDim>
158 +void invokeFp32RouterGemm(float* output, InputT const* mat_a,
159 + float const* mat_b, cudaStream_t stream) {
160 + constexpr int kBlockSize = 128;
161 + cudaLaunchConfig_t config;
162 + config.gridDim = kNumExperts;
163 + config.blockDim = kBlockSize;
164 + config.dynamicSmemBytes = 0;
165 + config.stream = stream;
166 + cudaLaunchAttribute attrs[1];
167 + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
168 + attrs[0].val.programmaticStreamSerializationAllowed = 1;
169 + config.numAttrs = 1;
170 + config.attrs = attrs;
171 + cudaLaunchKernelEx(&config,
172 + fp32_router_gemm_kernel<InputT, kBlockSize, kNumTokens,
173 + kNumExperts, kHiddenDim>,
174 + output, mat_a, mat_b);
175 +}
176 +
177 +// ---------------------------------------------------------------------------
178 +// Explicit instantiations: M=1..32, E=256, H=3072, for both input types
179 +// ---------------------------------------------------------------------------
180 +
181 +#define INSTANTIATE(T, M) \
182 + template void invokeFp32RouterGemm<T, M, 256, 3072>( \
183 + float*, T const*, float const*, cudaStream_t);
184 +
185 +#define INSTANTIATE_ALL(T) \
186 + INSTANTIATE(T, 1) \
187 + INSTANTIATE(T, 2) \
188 + INSTANTIATE(T, 3) \
189 + INSTANTIATE(T, 4) \
190 + INSTANTIATE(T, 5) \
191 + INSTANTIATE(T, 6) \
192 + INSTANTIATE(T, 7) \
193 + INSTANTIATE(T, 8) \
194 + INSTANTIATE(T, 9) \
195 + INSTANTIATE(T, 10) \
196 + INSTANTIATE(T, 11) \
197 + INSTANTIATE(T, 12) \
198 + INSTANTIATE(T, 13) \
199 + INSTANTIATE(T, 14) \
200 + INSTANTIATE(T, 15) \
201 + INSTANTIATE(T, 16) \
202 + INSTANTIATE(T, 17) \
203 + INSTANTIATE(T, 18) \
204 + INSTANTIATE(T, 19) \
205 + INSTANTIATE(T, 20) \
206 + INSTANTIATE(T, 21) \
207 + INSTANTIATE(T, 22) \
208 + INSTANTIATE(T, 23) \
209 + INSTANTIATE(T, 24) \
210 + INSTANTIATE(T, 25) \
211 + INSTANTIATE(T, 26) \
212 + INSTANTIATE(T, 27) \
213 + INSTANTIATE(T, 28) \
214 + INSTANTIATE(T, 29) \
215 + INSTANTIATE(T, 30) \
216 + INSTANTIATE(T, 31) \
217 + INSTANTIATE(T, 32)
218 +
219 +INSTANTIATE_ALL(float)
220 +INSTANTIATE_ALL(__nv_bfloat16)
221 +
222 +#undef INSTANTIATE_ALL
223 +#undef INSTANTIATE
csrc/libtorch_stable/fp32_router_gemm_entry.cu new
+127
@@ -0,0 +1,127 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +#include <torch/csrc/stable/library.h>
5 +#include <torch/csrc/stable/tensor.h>
6 +#include <torch/headeronly/core/ScalarType.h>
7 +
8 +#include "core/registration.h"
9 +#include "libtorch_stable/torch_utils.h"
10 +
11 +#include <cuda_bf16.h>
12 +#include <cuda_runtime.h>
13 +
14 +#include <stdexcept>
15 +
16 +namespace {
17 +
18 +inline int getSMVersion() {
19 + auto* props = get_device_prop();
20 + return props->major * 10 + props->minor;
21 +}
22 +
23 +} // namespace
24 +
25 +static constexpr int FP32_NUM_EXPERTS = 256;
26 +static constexpr int FP32_HIDDEN_DIM = 3072;
27 +static constexpr int FP32_MAX_TOKENS = 32;
28 +
29 +// Forward declarations — 4 template params must match fp32_router_gemm.cu
30 +template <typename InputT, int kNumTokens, int kNumExperts, int kHiddenDim>
31 +void invokeFp32RouterGemm(float* output, InputT const* mat_a,
32 + float const* mat_b, cudaStream_t stream);
33 +
34 +// LoopUnroller templated on InputT
35 +template <typename InputT, int kBegin, int kEnd>
36 +struct Fp32LoopUnroller {
37 + static void unroll(int num_tokens, float* output, InputT const* mat_a,
38 + float const* mat_b, cudaStream_t stream) {
39 + if (num_tokens == kBegin) {
40 + invokeFp32RouterGemm<InputT, kBegin, FP32_NUM_EXPERTS, FP32_HIDDEN_DIM>(
41 + output, mat_a, mat_b, stream);
42 + } else {
43 + Fp32LoopUnroller<InputT, kBegin + 1, kEnd>::unroll(num_tokens, output,
44 + mat_a, mat_b, stream);
45 + }
46 + }
47 +};
48 +
49 +template <typename InputT, int kEnd>
50 +struct Fp32LoopUnroller<InputT, kEnd, kEnd> {
51 + static void unroll(int num_tokens, float* output, InputT const* mat_a,
52 + float const* mat_b, cudaStream_t stream) {
53 + if (num_tokens == kEnd) {
54 + invokeFp32RouterGemm<InputT, kEnd, FP32_NUM_EXPERTS, FP32_HIDDEN_DIM>(
55 + output, mat_a, mat_b, stream);
56 + } else {
57 + throw std::invalid_argument(
58 + "fp32_router_gemm: num_tokens must be in [1, 32]");
59 + }
60 + }
61 +};
62 +
63 +void fp32_router_gemm(
64 + torch::stable::Tensor& output, // [num_tokens, num_experts]
65 + torch::stable::Tensor const& mat_a, // [num_tokens, hidden_dim]
66 + torch::stable::Tensor const& mat_b // [num_experts, hidden_dim]
67 +) {
68 + STD_TORCH_CHECK(output.dim() == 2 && mat_a.dim() == 2 && mat_b.dim() == 2);
69 + STD_TORCH_CHECK(output.is_cuda() && mat_a.is_cuda() && mat_b.is_cuda(),
70 + "fp32_router_gemm: all tensors must be CUDA tensors");
71 + STD_TORCH_CHECK(output.get_device_index() == mat_a.get_device_index() &&
72 + output.get_device_index() == mat_b.get_device_index(),
73 + "fp32_router_gemm: all tensors must be on the same device");
74 + STD_TORCH_CHECK(
75 + output.is_contiguous() && mat_a.is_contiguous() && mat_b.is_contiguous(),
76 + "fp32_router_gemm: all tensors must be contiguous");
77 +
78 + const int num_tokens = mat_a.size(0);
79 + const int num_experts = mat_b.size(0);
80 + const int hidden_dim = mat_a.size(1);
81 +
82 + STD_TORCH_CHECK(output.size(0) == num_tokens && output.size(1) == num_experts,
83 + "fp32_router_gemm: output must have shape [num_tokens, "
84 + "num_experts]");
85 + STD_TORCH_CHECK(
86 + mat_a.size(1) == mat_b.size(1),
87 + "fp32_router_gemm: mat_a and mat_b must have the same hidden_dim");
88 + STD_TORCH_CHECK(hidden_dim == FP32_HIDDEN_DIM,
89 + "fp32_router_gemm: expected hidden_dim=3072");
90 + STD_TORCH_CHECK(num_experts == FP32_NUM_EXPERTS,
91 + "fp32_router_gemm: expected num_experts=256");
92 + STD_TORCH_CHECK(num_tokens <= FP32_MAX_TOKENS,
93 + "fp32_router_gemm: num_tokens must be in [0, 32]");
94 + STD_TORCH_CHECK(
95 + mat_a.scalar_type() == torch::headeronly::ScalarType::Float ||
96 + mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16,
97 + "fp32_router_gemm: mat_a must be float32 or bfloat16");
98 + STD_TORCH_CHECK(mat_b.scalar_type() == torch::headeronly::ScalarType::Float,
99 + "fp32_router_gemm: mat_b (weight) must be float32");
100 + STD_TORCH_CHECK(output.scalar_type() == torch::headeronly::ScalarType::Float,
101 + "fp32_router_gemm: output must be float32");
102 +
103 + if (num_tokens == 0) {
104 + return;
105 + }
106 +
107 + STD_TORCH_CHECK(getSMVersion() >= 90, "fp32_router_gemm: requires SM90+");
108 +
109 + auto stream = get_current_cuda_stream(mat_a.get_device_index());
110 + float* out_ptr = reinterpret_cast<float*>(output.mutable_data_ptr());
111 + float const* mat_b_ptr = reinterpret_cast<float const*>(mat_b.data_ptr());
112 +
113 + if (mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
114 + auto const* mat_a_ptr =
115 + reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr());
116 + Fp32LoopUnroller<__nv_bfloat16, 1, FP32_MAX_TOKENS>::unroll(
117 + num_tokens, out_ptr, mat_a_ptr, mat_b_ptr, stream);
118 + } else {
119 + auto const* mat_a_ptr = reinterpret_cast<float const*>(mat_a.data_ptr());
120 + Fp32LoopUnroller<float, 1, FP32_MAX_TOKENS>::unroll(
121 + num_tokens, out_ptr, mat_a_ptr, mat_b_ptr, stream);
122 + }
123 +}
124 +
125 +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
126 + m.impl("fp32_router_gemm", TORCH_BOX(&fp32_router_gemm));
127 +}
csrc/libtorch_stable/torch_bindings.cpp
+4
@@ -247,6 +247,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
247 ops.def(
248 "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
249
250 + // BF16/FP32 x FP32 -> FP32 router GEMM for H=3072, E=256, M<=32 (SM90+).
251 + // conditionally compiled so impl registration is in source file
252 + ops.def("fp32_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
253 +
254 // reorder weight for AllSpark Ampere W8A16 Fused Gemm kernel
255 ops.def(
256 "rearrange_kn_weight_as_n32k16_order(Tensor b_qweight, Tensor b_scales, "
tests/kernels/test_fp32_router_gemm.py new
+78
@@ -0,0 +1,78 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for fp32_router_gemm kernel: activation×weight→fp32, H=3072, E=256.
4 +
5 +Correctness baseline: torch.matmul in float64.
6 +"""
7 +
8 +import pytest
9 +import torch
10 +
11 +from vllm._custom_ops import fp32_router_gemm
12 +
13 +NUM_EXPERTS = 256
14 +HIDDEN_DIM = 3072
15 +# Absolute tolerance for fp32 kernel vs float64 reference
16 +ATOL_FP32 = 2e-4
17 +ATOL_BF16 = 2e-2 # bf16 activation has lower precision
18 +
19 +
20 +def _requires_sm90():
21 + if not torch.cuda.is_available():
22 + pytest.skip("CUDA not available")
23 + major, minor = torch.cuda.get_device_capability()
24 + if major * 10 + minor < 90:
25 + pytest.skip(f"fp32_router_gemm requires SM90+, got SM{major}{minor}")
26 +
27 +
28 +def _ref(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor:
29 + """Reference: F.linear in float32 on GPU."""
30 + return torch.nn.functional.linear(mat_a.float(), mat_b.float())
31 +
32 +
33 +@pytest.mark.parametrize("num_tokens", [1, 2, 4, 8, 16, 32])
34 +def test_fp32_activation(num_tokens: int):
35 + """fp32 activation → fp32 output should match reference closely."""
36 + _requires_sm90()
37 + torch.manual_seed(42)
38 + device = torch.device("cuda")
39 + mat_a = torch.randn(num_tokens, HIDDEN_DIM, dtype=torch.float32, device=device)
40 + mat_b = torch.randn(NUM_EXPERTS, HIDDEN_DIM, dtype=torch.float32, device=device)
41 +
42 + out = fp32_router_gemm(mat_a, mat_b)
43 + ref = _ref(mat_a, mat_b)
44 +
45 + assert out.shape == (num_tokens, NUM_EXPERTS)
46 + assert out.dtype == torch.float32
47 + torch.testing.assert_close(out, ref, atol=ATOL_FP32, rtol=0)
48 +
49 +
50 +@pytest.mark.parametrize("num_tokens", [1, 2, 4, 8, 16, 32])
51 +def test_bf16_activation(num_tokens: int):
52 + """bf16 activation → fp32 output should match reference within bf16 error."""
53 + _requires_sm90()
54 + torch.manual_seed(42)
55 + device = torch.device("cuda")
56 + mat_a_bf16 = torch.randn(
57 + num_tokens, HIDDEN_DIM, dtype=torch.bfloat16, device=device
58 + )
59 + mat_b = torch.randn(NUM_EXPERTS, HIDDEN_DIM, dtype=torch.float32, device=device)
60 +
61 + out = fp32_router_gemm(mat_a_bf16, mat_b)
62 + ref = _ref(mat_a_bf16, mat_b).to(device)
63 +
64 + assert out.shape == (num_tokens, NUM_EXPERTS)
65 + assert out.dtype == torch.float32
66 + torch.testing.assert_close(out, ref, atol=ATOL_BF16, rtol=0)
67 +
68 +
69 +def test_output_shape_and_dtype():
70 + """Basic shape and dtype checks."""
71 + _requires_sm90()
72 + device = torch.device("cuda")
73 + mat_a = torch.randn(4, HIDDEN_DIM, dtype=torch.float32, device=device)
74 + mat_b = torch.randn(NUM_EXPERTS, HIDDEN_DIM, dtype=torch.float32, device=device)
75 + out = fp32_router_gemm(mat_a, mat_b)
76 + assert out.shape == (4, NUM_EXPERTS)
77 + assert out.dtype == torch.float32
78 + assert out.device.type == "cuda"
vllm/_custom_ops.py
+25
@@ -2412,6 +2412,31 @@ def dsv3_router_gemm(
2412 return output
2413
2414
2415 +def fp32_router_gemm(
2416 + hidden_states: torch.Tensor,
2417 + router_weight: torch.Tensor,
2418 +) -> torch.Tensor:
2419 + output = torch.empty(
2420 + hidden_states.shape[0],
2421 + router_weight.shape[0],
2422 + device=hidden_states.device,
2423 + dtype=torch.float32,
2424 + )
2425 + torch.ops._C.fp32_router_gemm(output, hidden_states, router_weight)
2426 + return output
2427 +
2428 +
2429 +if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "fp32_router_gemm"):
2430 +
2431 + @register_fake("_C::fp32_router_gemm")
2432 + def fp32_router_gemm_fake(
2433 + output: torch.Tensor,
2434 + mat_a: torch.Tensor,
2435 + mat_b: torch.Tensor,
2436 + ) -> None:
2437 + return
2438 +
2439 +
2440 def topk_softmax(
2441 topk_weights: torch.Tensor,
2442 topk_ids: torch.Tensor,
vllm/model_executor/layers/fused_moe/router/gate_linear.py
+67 -9
@@ -3,18 +3,22 @@
3 import torch
4 from torch.nn.parameter import Parameter
5
6 +import vllm._custom_ops as ops
7 from vllm.model_executor.custom_op import PluggableLayer
8 from vllm.model_executor.layers.linear import ReplicatedLinear
9 from vllm.platforms import current_platform
10 +from vllm.utils.torch_utils import direct_register_custom_op
11
12
13 @PluggableLayer.register("gate_linear")
14 class GateLinear(ReplicatedLinear):
13 - """MoE gate linear layer with three-tier GEMM dispatch:
15 + """MoE gate linear layer with multi-tier GEMM dispatch:
16
15 - 1. DSV3 specialized kernel (SM90+, batch<=16, supported dims)
16 - 2. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 + fp32 out_dtype)
17 - 3. F.linear via ReplicatedLinear (ultimate fallback)
17 + 1. DSV3 specialized kernel (SM90+, fp32 out, M<=16, H=7168, E=256/384)
18 + 2. fp32 specialized kernel (SM90+, bf16/fp32 in, fp32 out,
19 + M<=32, H=3072, E=256)
20 + 3. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype)
21 + 4. F.linear via ReplicatedLinear (ultimate fallback)
22
23 The ``out_dtype`` attribute is mutable and can be set after init
24 (e.g. when the required dtype depends on the expert quantization
@@ -25,6 +29,11 @@ class GateLinear(ReplicatedLinear):
29 DSV3_SUPPORTED_NUM_EXPERTS = [256, 384]
30 DSV3_SUPPORTED_HIDDEN_SIZES = [7168]
31
32 + # Dimensions supported by the fp32 specialized kernel
33 + FP32_SUPPORTED_NUM_EXPERTS = [256]
34 + FP32_SUPPORTED_HIDDEN_SIZES = [3072]
35 + FP32_MAX_TOKENS = 32
36 +
37 def __init__(
38 self,
39 input_size: int,
@@ -43,7 +52,7 @@ class GateLinear(ReplicatedLinear):
52 )
53
54 # If fp32 compute is required and no specialized kernel is available,
46 - # store weights in fp32 so Tier 3 computes in fp32 natively.
55 + # store weights in fp32 so the fallback linear path computes in fp32.
56 if force_fp32_compute and not can_use_specialized_kernels:
57 params_dtype = torch.float32
58
@@ -65,6 +74,16 @@ class GateLinear(ReplicatedLinear):
74 and input_size in self.DSV3_SUPPORTED_HIDDEN_SIZES
75 )
76
77 + # fp32 specialized kernel eligibility (SM90+, exact dims, fp32 weight)
78 + self.allow_fp32_router_gemm = (
79 + not bias
80 + and self.weight.dtype == torch.float32
81 + and current_platform.is_cuda()
82 + and is_hopper_or_blackwell
83 + and output_size in self.FP32_SUPPORTED_NUM_EXPERTS
84 + and input_size in self.FP32_SUPPORTED_HIDDEN_SIZES
85 + )
86 +
87 # cuBLAS bf16→fp32 eligibility
88 self.allow_cublas_router_gemm = (
89 self.allow_specialized_router_gemm
@@ -92,8 +111,6 @@ class GateLinear(ReplicatedLinear):
111 def forward(
112 self, x: torch.Tensor
113 ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
95 - import vllm._custom_ops as ops
96 -
114 # Tier 1: DSV3 specialized kernel
115 if self.allow_dsv3_router_gemm and x.shape[0] <= 16:
116 output = ops.dsv3_router_gemm(
@@ -103,15 +120,56 @@ class GateLinear(ReplicatedLinear):
120 )
121 return output, None
122
106 - # Tier 2: cuBLAS bf16→fp32
123 + # Tier 2: fp32 specialized kernel (H=3072, E=256, M<=32)
124 + # Dispatch is wrapped in a custom op so that torch.compile/CUDA-graph
125 + # capture does not freeze the runtime num_tokens branch.
126 + if self.allow_fp32_router_gemm and x.dtype in (
127 + torch.float32,
128 + torch.bfloat16,
129 + ):
130 + output = torch.ops.vllm.fp32_router_gemm_dispatch(x, self.weight)
131 + return output, None
132 +
133 + # Tier 3: cuBLAS bf16→fp32
134 if self.allow_cublas_router_gemm and x.dtype == torch.bfloat16:
135 output = torch.mm(x, self.weight.T, out_dtype=torch.float32)
136 return output, None
137
111 - # Tier 3: F.linear (ReplicatedLinear)
138 + # Tier 4: F.linear (ReplicatedLinear)
139 if self.out_dtype is not None and x.dtype != self.weight.dtype:
140 x = x.to(self.weight.dtype)
141 output, output_bias = super().forward(x)
142 if self.out_dtype is not None and output.dtype != self.out_dtype:
143 output = output.to(self.out_dtype)
144 return output, output_bias
145 +
146 +
147 +_FP32_ROUTER_GEMM_MAX_TOKENS = GateLinear.FP32_MAX_TOKENS
148 +
149 +
150 +def fp32_router_gemm_dispatch_impl(
151 + x: torch.Tensor, weight: torch.Tensor
152 +) -> torch.Tensor:
153 + """
154 + Dynamically run fp32 specialized gemm if num_tokens <= FP32_MAX_TOKENS,
155 + otherwise fall back to F.linear.
156 + This must be wrapped in a custom op because our torch.compile integration
157 + does not support runtime dispatching on num_tokens.
158 + """
159 + if x.shape[0] <= _FP32_ROUTER_GEMM_MAX_TOKENS:
160 + return ops.fp32_router_gemm(x, weight)
161 + else:
162 + return torch.nn.functional.linear(x.float(), weight)
163 +
164 +
165 +def fp32_router_gemm_dispatch_fake(
166 + x: torch.Tensor, weight: torch.Tensor
167 +) -> torch.Tensor:
168 + return x.new_empty((x.shape[0], weight.shape[0]), dtype=torch.float32)
169 +
170 +
171 +direct_register_custom_op(
172 + op_name="fp32_router_gemm_dispatch",
173 + op_func=fp32_router_gemm_dispatch_impl,
174 + fake_impl=fp32_router_gemm_dispatch_fake,
175 +)
vllm/model_executor/models/minimax_m2.py
+4 -4
@@ -43,10 +43,10 @@ from vllm.model_executor.layers.fused_moe import (
43 FusedMoE,
44 fused_moe_make_expert_params_mapping,
45 )
46 +from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
47 from vllm.model_executor.layers.layernorm import RMSNorm
48 from vllm.model_executor.layers.linear import (
49 QKVParallelLinear,
49 - ReplicatedLinear,
50 RowParallelLinear,
51 )
52 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -113,12 +113,12 @@ class MiniMaxM2MoE(nn.Module):
113 router_logits_dtype=torch.float32,
114 )
115
116 - self.gate = ReplicatedLinear(
116 + self.gate = GateLinear(
117 config.hidden_size,
118 config.num_local_experts,
119 bias=False,
120 params_dtype=torch.float32,
121 - quant_config=None,
121 + out_dtype=torch.float32,
122 prefix=f"{prefix}.gate",
123 )
124
@@ -132,7 +132,7 @@ class MiniMaxM2MoE(nn.Module):
132 hidden_states = hidden_states.view(-1, hidden_dim)
133
134 # router_logits: (num_tokens, n_experts)
135 - router_logits, _ = self.gate(hidden_states.to(torch.float32))
135 + router_logits, _ = self.gate(hidden_states)
136 final_hidden_states = self.experts(
137 hidden_states=hidden_states, router_logits=router_logits
138 )