@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 19069bcbd5

FP32 router GEMV optimization (#48335)

FP32 router GEMV optimization (#48335) Signed-off-by: peiyuanz <peiyuanz@inferact.ai> Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: peiyuanz <peiyuanz@inferact.ai> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: zhouzhou <zhouzhou@zhouzhoudeMacBook-Pro.local>

Jee Jee Li committed Jul 11, 2026 at 21:07 UTC 19069bcbd5be607a9c502e007951d8325d6a340e
4 files changed +284 -67
csrc/libtorch_stable/fp32_router_gemm.cu
+179 -56
@@ -1,13 +1,18 @@
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.
4 +// Router GEMM: activation(T) x weight(fp32) -> fp32, M<=32, for the
5 +// supported (E, H) pairs listed at the bottom of this file.
6 // Supports bf16 or fp32 activation; weight is always fp32.
7 // Adapted from dsv3_router_gemm_float_out.cu.
8 +// (E=256, H=6144) bf16 uses a B300-tuned wide-block geometry; see
9 +// invokeFp32RouterGemm.
10
11 #include <cuda_bf16.h>
12 #include <cuda_runtime.h>
13
14 +#include <type_traits>
15 +
16 // ---------------------------------------------------------------------------
17 // Load helpers
18 // ---------------------------------------------------------------------------
@@ -73,94 +78,113 @@ __device__ __forceinline__ void load_activation<__nv_bfloat16, 8>(
78 // InputT : type of activation (float or __nv_bfloat16)
79 // Weight is always fp32; output is always fp32.
80 // 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) {
81 +// Each block computes kEPB expert columns; wider blocks / kEPB > 1 are
82 +// selected per (shape, M) in invokeFp32RouterGemm (B300-tuned, see below).
83 +// kTGroups > 1 splits the tokens across groups of kBlockSize threads within
84 +// the block: all groups scan the same weight K-slices (group 0 misses to
85 +// DRAM, later groups hit L1) so weight traffic stays 1x, while per-thread
86 +// accumulator registers drop by kTGroups (at M=16 the 32 fp32 accumulators
87 +// push the kernel to 128 regs/thread and 1 block/SM).
88 +template <typename InputT, int kBlockSize, int kNumTokens, int kEPB,
89 + int kNumExperts, int kHiddenDim, int kTGroups = 1>
90 +__global__ __launch_bounds__(
91 + kBlockSize* kTGroups, 1) void fp32_router_gemm_kernel(float* out,
92 + InputT const* mat_a,
93 + float const* mat_b) {
94 constexpr int VPT = 16 / sizeof(InputT);
95 constexpr int k_elems_per_k_iteration = VPT * kBlockSize;
96 constexpr int k_iterations = kHiddenDim / k_elems_per_k_iteration;
97 + static_assert(kHiddenDim % k_elems_per_k_iteration == 0);
98 + static_assert(kNumTokens % kTGroups == 0);
99 constexpr int kWarpSize = 32;
84 - constexpr int kNumWarps = kBlockSize / kWarpSize;
100 + constexpr int kNumWarps = kBlockSize / kWarpSize; // per token group
101 + constexpr int kMG = kNumTokens / kTGroups; // tokens per group
102
86 - int const n_idx = blockIdx.x;
87 - int const tid = threadIdx.x;
103 + int const e_base = blockIdx.x * kEPB;
104 + int const tid = threadIdx.x % kBlockSize;
105 + int const m0 = (threadIdx.x / kBlockSize) * kMG;
106 int const warpId = tid / kWarpSize;
107 int const laneId = tid % kWarpSize;
108
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 - }
109 + float acc[kMG][kEPB] = {};
110 + __shared__ float sm_reduction[kNumTokens][kEPB][kNumWarps];
111
112 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
113 cudaGridDependencySynchronize();
114 + // Fire the PDL trigger right after our own wait instead of at kernel end:
115 + // a gridsync-ing consumer is unaffected (its wait always targets full grid
116 + // completion), while a consumer that reads none of our outputs (e.g. the
117 + // NVFP4 activation quant, which reads the same hidden_states) can launch
118 + // now and fully overlap this kernel's body.
119 + cudaTriggerProgrammaticLaunchCompletion();
120 #endif
121
122 +#pragma unroll
123 for (int ki = 0; ki < k_iterations; ki++) {
107 - int const k_base = k_bases[ki];
124 + int const k_base = ki * k_elems_per_k_iteration + tid * VPT;
125
109 - float b_float[VPT];
110 - load_weight<VPT>(b_col + k_base, b_float);
126 + float b_float[kEPB][VPT];
127 +#pragma unroll
128 + for (int e = 0; e < kEPB; e++) {
129 + load_weight<VPT>(mat_b + (e_base + e) * kHiddenDim + k_base, b_float[e]);
130 + }
131
132 #pragma unroll
113 - for (int m_idx = 0; m_idx < kNumTokens; m_idx++) {
133 + for (int m_idx = 0; m_idx < kMG; m_idx++) {
134 float a_float[VPT];
115 - load_activation<InputT, VPT>(mat_a + m_idx * kHiddenDim + k_base,
116 - a_float);
135 + load_activation<InputT, VPT>(
136 + mat_a + (size_t)(m0 + m_idx) * kHiddenDim + k_base, a_float);
137 +#pragma unroll
138 + for (int e = 0; e < kEPB; e++) {
139 #pragma unroll
118 - for (int k = 0; k < VPT; k++) {
119 - acc[m_idx] += a_float[k] * b_float[k];
140 + for (int k = 0; k < VPT; k++) {
141 + acc[m_idx][e] += a_float[k] * b_float[e][k];
142 + }
143 }
144 }
145 }
146
147 // Warp-level butterfly reduction
148 #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;
149 + for (int m = 0; m < kMG; m++) {
150 +#pragma unroll
151 + for (int e = 0; e < kEPB; e++) {
152 + float sum = acc[m][e];
153 + sum += __shfl_xor_sync(0xffffffff, sum, 16);
154 + sum += __shfl_xor_sync(0xffffffff, sum, 8);
155 + sum += __shfl_xor_sync(0xffffffff, sum, 4);
156 + sum += __shfl_xor_sync(0xffffffff, sum, 2);
157 + sum += __shfl_xor_sync(0xffffffff, sum, 1);
158 + if (laneId == 0) sm_reduction[m0 + m][e][warpId] = sum;
159 + }
160 }
161
162 __syncthreads();
163
138 - if (tid == 0) {
139 -#pragma unroll
140 - for (int m = 0; m < kNumTokens; m++) {
141 - float final_sum = 0.0f;
164 + // Parallel finalize: one thread per (m, e) output.
165 + for (int idx = threadIdx.x; idx < kNumTokens * kEPB;
166 + idx += kBlockSize * kTGroups) {
167 + int const m = idx / kEPB;
168 + int const e = idx % kEPB;
169 + float final_sum = 0.0f;
170 #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 - }
171 + for (int w = 0; w < kNumWarps; w++) final_sum += sm_reduction[m][e][w];
172 + out[m * kNumExperts + e_base + e] = final_sum;
173 }
147 -
148 -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
149 - cudaTriggerProgrammaticLaunchCompletion();
150 -#endif
174 }
175
176 // ---------------------------------------------------------------------------
177 // Launcher
178 // ---------------------------------------------------------------------------
179
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;
180 +template <typename InputT, int kBlockSize, int kEPB, int kNumTokens,
181 + int kNumExperts, int kHiddenDim, int kTGroups = 1>
182 +static void launchFp32RouterGemm(float* output, InputT const* mat_a,
183 + float const* mat_b, cudaStream_t stream) {
184 + static_assert(kNumExperts % kEPB == 0);
185 cudaLaunchConfig_t config;
162 - config.gridDim = kNumExperts;
163 - config.blockDim = kBlockSize;
186 + config.gridDim = kNumExperts / kEPB;
187 + config.blockDim = kBlockSize * kTGroups;
188 config.dynamicSmemBytes = 0;
189 config.stream = stream;
190 cudaLaunchAttribute attrs[1];
@@ -168,15 +192,112 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a,
192 attrs[0].val.programmaticStreamSerializationAllowed = 1;
193 config.numAttrs = 1;
194 config.attrs = attrs;
171 - cudaLaunchKernelEx(&config,
172 - fp32_router_gemm_kernel<InputT, kBlockSize, kNumTokens,
173 - kNumExperts, kHiddenDim>,
174 - output, mat_a, mat_b);
195 + cudaLaunchKernelEx(
196 + &config,
197 + fp32_router_gemm_kernel<InputT, kBlockSize, kNumTokens, kEPB, kNumExperts,
198 + kHiddenDim, kTGroups>,
199 + output, mat_a, mat_b);
200 +}
201 +
202 +static bool isBlackwellFamily() {
203 + static int sm = []() {
204 + int dev = 0, major = 0, minor = 0;
205 + cudaGetDevice(&dev);
206 + cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, dev);
207 + cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, dev);
208 + return major * 10 + minor;
209 + }();
210 + return sm >= 100;
211 +}
212 +
213 +template <typename InputT, int kNumTokens, int kNumExperts, int kHiddenDim>
214 +void invokeFp32RouterGemm(float* output, InputT const* mat_a,
215 + float const* mat_b, cudaStream_t stream) {
216 + // Geometry tuned on B300 per supported shape, bf16 activation, under a
217 + // production-fidelity harness (CUDA-graph replay, per-layer cold weights).
218 + // GLM-5.2 (E=256, H=6144):
219 + // M <= 4 : BS=768, EPB=1 (2.7us vs cast+cuBLAS 8.1us at M=1)
220 + // M in [5, 15]
221 + // or odd : BS=384, EPB=2 (crossover vs BS=768 measured in (4, 8))
222 + // M >= 16, even : BS=192, EPB=2, 2 token groups (M=16 4.79us vs 5.04,
223 + // M=24 5.71 vs 6.38, M=32 6.79 vs 7.72; M=12 loses at
224 + // 0.97x, so the boundary is 16).
225 + // Only enabled on the Blackwell family where it was validated; Hopper and
226 + // other shapes / fp32 activation keep the legacy geometry.
227 + if constexpr (std::is_same_v<InputT, __nv_bfloat16> && kNumExperts == 256 &&
228 + kHiddenDim == 6144) {
229 + if (!isBlackwellFamily()) {
230 + launchFp32RouterGemm<InputT, 128, 1, kNumTokens, kNumExperts, kHiddenDim>(
231 + output, mat_a, mat_b, stream);
232 + return;
233 + }
234 + if constexpr (kNumTokens <= 4) {
235 + launchFp32RouterGemm<InputT, 768, 1, kNumTokens, kNumExperts, kHiddenDim>(
236 + output, mat_a, mat_b, stream);
237 + } else if constexpr (kNumTokens >= 16 && kNumTokens % 2 == 0) {
238 + launchFp32RouterGemm<InputT, 192, 2, kNumTokens, kNumExperts, kHiddenDim,
239 + 2>(output, mat_a, mat_b, stream);
240 + } else {
241 + launchFp32RouterGemm<InputT, 384, 2, kNumTokens, kNumExperts, kHiddenDim>(
242 + output, mat_a, mat_b, stream);
243 + }
244 + } else if constexpr (std::is_same_v<InputT, __nv_bfloat16> &&
245 + kNumExperts == 128 && kHiddenDim == 6144) {
246 + // MiniMax-M3. Legacy 128/1 only fills 128 blocks and pays the same
247 + // accumulator register cliffs; B300 sweep:
248 + // even M in [6, 10] : BS=384, EPB=1, 2 token groups (1.26-1.43x)
249 + // even M >= 12 : BS=192, EPB=1, 2 token groups (1.59-1.66x at
250 + // M >= 18; re-measured on B300+B200: 192 also wins
251 + // M=12/14 by 5-11%% on both, ties 384 at 16)
252 + // M <= 5 / odd : BS=384, EPB=1 (1.03-1.19x)
253 + if (!isBlackwellFamily()) {
254 + launchFp32RouterGemm<InputT, 128, 1, kNumTokens, kNumExperts, kHiddenDim>(
255 + output, mat_a, mat_b, stream);
256 + return;
257 + }
258 + if constexpr (kNumTokens >= 12 && kNumTokens % 2 == 0) {
259 + launchFp32RouterGemm<InputT, 192, 1, kNumTokens, kNumExperts, kHiddenDim,
260 + 2>(output, mat_a, mat_b, stream);
261 + } else if constexpr (kNumTokens >= 6 && kNumTokens % 2 == 0) {
262 + launchFp32RouterGemm<InputT, 384, 1, kNumTokens, kNumExperts, kHiddenDim,
263 + 2>(output, mat_a, mat_b, stream);
264 + } else {
265 + launchFp32RouterGemm<InputT, 384, 1, kNumTokens, kNumExperts, kHiddenDim>(
266 + output, mat_a, mat_b, stream);
267 + }
268 + } else if constexpr (std::is_same_v<InputT, __nv_bfloat16> &&
269 + kNumExperts == 256 && kHiddenDim == 3072) {
270 + // MiniMax-M2/M2.5. The 3.1MB weight is latency-floor bound at small M
271 + // (legacy already optimal); token groups win only at even M >= 8
272 + // (1.05-1.17x). EPB crossover measured between 12 and 16.
273 + if (!isBlackwellFamily()) {
274 + launchFp32RouterGemm<InputT, 128, 1, kNumTokens, kNumExperts, kHiddenDim>(
275 + output, mat_a, mat_b, stream);
276 + return;
277 + }
278 + if constexpr (kNumTokens >= 14 && kNumTokens % 2 == 0) {
279 + // M=14 originally measured 0.91x and stayed on legacy; two fresh
280 + // sweeps (B300 dev1 + B200) both put 192/2/tg2 ahead by 3.5-4%%.
281 + launchFp32RouterGemm<InputT, 192, 2, kNumTokens, kNumExperts, kHiddenDim,
282 + 2>(output, mat_a, mat_b, stream);
283 + } else if constexpr (kNumTokens >= 8 && kNumTokens <= 12 &&
284 + kNumTokens % 2 == 0) {
285 + launchFp32RouterGemm<InputT, 192, 1, kNumTokens, kNumExperts, kHiddenDim,
286 + 2>(output, mat_a, mat_b, stream);
287 + } else {
288 + launchFp32RouterGemm<InputT, 128, 1, kNumTokens, kNumExperts, kHiddenDim>(
289 + output, mat_a, mat_b, stream);
290 + }
291 + } else {
292 + launchFp32RouterGemm<InputT, 128, 1, kNumTokens, kNumExperts, kHiddenDim>(
293 + output, mat_a, mat_b, stream);
294 + }
295 }
296
297 // ---------------------------------------------------------------------------
298 // Explicit instantiations: M=1..32, for both input types, for the supported
179 -// (E, H) pairs: (256, 3072) [MiniMax-M2/M2.5] and (128, 6144) [MiniMax-M3].
299 +// (E, H) pairs: (256, 3072) [MiniMax-M2/M2.5], (128, 6144) [MiniMax-M3]
300 +// and (256, 6144) [GLM-5.2].
301 // ---------------------------------------------------------------------------
302
303 #define INSTANTIATE(T, M, E, H) \
@@ -221,6 +342,8 @@ INSTANTIATE_ALL(float, 256, 3072)
342 INSTANTIATE_ALL(__nv_bfloat16, 256, 3072)
343 INSTANTIATE_ALL(float, 128, 6144)
344 INSTANTIATE_ALL(__nv_bfloat16, 128, 6144)
345 +INSTANTIATE_ALL(float, 256, 6144)
346 +INSTANTIATE_ALL(__nv_bfloat16, 256, 6144)
347
348 #undef INSTANTIATE_ALL
349 #undef INSTANTIATE
csrc/libtorch_stable/fp32_router_gemm_entry.cu
+8 -3
@@ -25,10 +25,12 @@ inline int getSMVersion() {
25 static constexpr int FP32_MAX_TOKENS = 32;
26
27 // Supported (hidden_dim, num_experts) pairs (must match the instantiations in
28 -// fp32_router_gemm.cu): (3072, 256) for MiniMax-M2/M2.5, (6144, 128) for M3.
28 +// fp32_router_gemm.cu): (3072, 256) for MiniMax-M2/M2.5, (6144, 128) for M3,
29 +// (6144, 256) for GLM-5.2.
30 static inline bool fp32_router_gemm_supported(int hidden_dim, int num_experts) {
31 return (hidden_dim == 3072 && num_experts == 256) ||
31 - (hidden_dim == 6144 && num_experts == 128);
32 + (hidden_dim == 6144 && num_experts == 128) ||
33 + (hidden_dim == 6144 && num_experts == 256);
34 }
35
36 // Forward declarations — 4 template params must match fp32_router_gemm.cu
@@ -77,6 +79,9 @@ void dispatchFp32RouterGemm(int num_experts, int hidden_dim, int num_tokens,
79 } else if (num_experts == 128 && hidden_dim == 6144) {
80 Fp32LoopUnroller<InputT, 128, 6144, 1, FP32_MAX_TOKENS>::unroll(
81 num_tokens, output, mat_a, mat_b, stream);
82 + } else if (num_experts == 256 && hidden_dim == 6144) {
83 + Fp32LoopUnroller<InputT, 256, 6144, 1, FP32_MAX_TOKENS>::unroll(
84 + num_tokens, output, mat_a, mat_b, stream);
85 } else {
86 throw std::invalid_argument(
87 "fp32_router_gemm: unsupported (hidden_dim, num_experts) pair");
@@ -111,7 +116,7 @@ void fp32_router_gemm(
116 STD_TORCH_CHECK(
117 fp32_router_gemm_supported(hidden_dim, num_experts),
118 "fp32_router_gemm: supported (hidden_dim, num_experts) pairs are "
114 - "(3072, 256) and (6144, 128)");
119 + "(3072, 256), (6144, 128) and (6144, 256)");
120 STD_TORCH_CHECK(num_tokens <= FP32_MAX_TOKENS,
121 "fp32_router_gemm: num_tokens must be in [0, 32]");
122 STD_TORCH_CHECK(
tests/kernels/test_fp32_router_gemm.py
+94 -5
@@ -3,9 +3,13 @@
3 """Tests for fp32_router_gemm kernel: activation×weight→fp32.
4
5 Supported (hidden_size, num_experts) pairs:
6 - (3072, 256) -> MiniMax-M2/M2.5, (6144, 128) -> MiniMax-M3
6 + (3072, 256) -> MiniMax-M2/M2.5, (6144, 128) -> MiniMax-M3,
7 + (6144, 256) -> GLM-5.2
8
8 -Correctness baseline: torch.matmul in float64.
9 +Correctness baseline: F.linear in float32. Every M in [1, 32] is covered so
10 +all tuned geometries (wide-block, experts-per-block, token-group; boundaries
11 +at M=4/5, odd/even, M=15/16) are exercised on Blackwell, and the legacy
12 +128/1 geometry everywhere else.
13 """
14
15 import pytest
@@ -14,7 +18,8 @@ import torch
18 from vllm._custom_ops import fp32_router_gemm
19
20 # (hidden_size, num_experts)
17 -SHAPES = [(3072, 256), (6144, 128)]
21 +SHAPES = [(3072, 256), (6144, 128), (6144, 256)]
22 +ALL_M = list(range(1, 33))
23 # Absolute tolerance for fp32 kernel vs float64 reference
24 ATOL_FP32 = 2e-4
25 ATOL_BF16 = 2e-2 # bf16 activation has lower precision
@@ -34,7 +39,7 @@ def _ref(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor:
39
40
41 @pytest.mark.parametrize("hidden_dim,num_experts", SHAPES)
37 -@pytest.mark.parametrize("num_tokens", [1, 2, 4, 8, 16, 32])
42 +@pytest.mark.parametrize("num_tokens", ALL_M)
43 def test_fp32_activation(num_tokens: int, hidden_dim: int, num_experts: int):
44 """fp32 activation → fp32 output should match reference closely."""
45 _requires_sm90()
@@ -52,7 +57,7 @@ def test_fp32_activation(num_tokens: int, hidden_dim: int, num_experts: int):
57
58
59 @pytest.mark.parametrize("hidden_dim,num_experts", SHAPES)
55 -@pytest.mark.parametrize("num_tokens", [1, 2, 4, 8, 16, 32])
60 +@pytest.mark.parametrize("num_tokens", ALL_M)
61 def test_bf16_activation(num_tokens: int, hidden_dim: int, num_experts: int):
62 """bf16 activation → fp32 output should match reference within bf16 error."""
63 _requires_sm90()
@@ -82,3 +87,87 @@ def test_output_shape_and_dtype(hidden_dim: int, num_experts: int):
87 assert out.shape == (4, num_experts)
88 assert out.dtype == torch.float32
89 assert out.device.type == "cuda"
90 +
91 +
92 +@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES)
93 +@pytest.mark.parametrize("num_tokens", [1, 2, 4, 8, 16, 24, 32])
94 +def test_topk_routing_consistency(num_tokens: int, hidden_dim: int, num_experts: int):
95 + """The gate feeds top-k expert selection: the kernel's top-8 must match
96 + an fp64 reference's top-8 per token (ties tolerated). This is the
97 + business-level correctness of the router — numeric error only matters
98 + if it flips the argsort."""
99 + _requires_sm90()
100 + top_k = 8
101 + device = torch.device("cuda")
102 + for seed in range(5):
103 + torch.manual_seed(1000 + seed)
104 + mat_a = torch.randn(num_tokens, hidden_dim, dtype=torch.bfloat16, device=device)
105 + mat_b = torch.randn(num_experts, hidden_dim, dtype=torch.float32, device=device)
106 + out = fp32_router_gemm(mat_a, mat_b)
107 + ref = mat_a.double() @ mat_b.double().t()
108 + kernel_idx = out.topk(top_k, dim=-1).indices
109 + ref_vals, ref_idx = ref.topk(top_k, dim=-1)
110 + for t in range(num_tokens):
111 + got = set(kernel_idx[t].tolist())
112 + want = set(ref_idx[t].tolist())
113 + if got == want:
114 + continue
115 + # Tolerate genuine near-ties around the k-th value only.
116 + kth = ref_vals[t, -1].item()
117 + for e in got.symmetric_difference(want):
118 + gap = abs(ref[t, e].item() - kth)
119 + assert gap < 1e-3, (
120 + f"top-{top_k} mismatch beyond tie tolerance: token {t}, "
121 + f"expert {e}, gap {gap:.3e}"
122 + )
123 +
124 +
125 +def test_zero_tokens_returns_empty():
126 + """M=0 is a graceful no-op returning an empty [0, E] tensor."""
127 + _requires_sm90()
128 + device = torch.device("cuda")
129 + hidden_dim, num_experts = SHAPES[0]
130 + mat_a = torch.empty(0, hidden_dim, dtype=torch.float32, device=device)
131 + mat_b = torch.randn(num_experts, hidden_dim, dtype=torch.float32, device=device)
132 + out = fp32_router_gemm(mat_a, mat_b)
133 + assert out.shape == (0, num_experts)
134 + assert out.dtype == torch.float32
135 +
136 +
137 +def test_rejects_invalid_inputs():
138 + """The entry must fail loudly, never compute silently wrong results."""
139 + _requires_sm90()
140 + device = torch.device("cuda")
141 + hidden_dim, num_experts = SHAPES[0]
142 + mat_b = torch.randn(num_experts, hidden_dim, dtype=torch.float32, device=device)
143 +
144 + # num_tokens > 32 (beyond the instantiated range)
145 + with pytest.raises(Exception, match="num_tokens"):
146 + fp32_router_gemm(
147 + torch.randn(33, hidden_dim, dtype=torch.float32, device=device), mat_b
148 + )
149 +
150 + # unsupported (hidden_dim, num_experts) pair
151 + with pytest.raises(Exception, match="supported"):
152 + fp32_router_gemm(
153 + torch.randn(4, 1024, dtype=torch.float32, device=device),
154 + torch.randn(64, 1024, dtype=torch.float32, device=device),
155 + )
156 +
157 + # non-contiguous activation (a column-slice view)
158 + wide = torch.randn(4, hidden_dim * 2, dtype=torch.float32, device=device)
159 + with pytest.raises(Exception, match="contiguous"):
160 + fp32_router_gemm(wide[:, :hidden_dim], mat_b)
161 +
162 + # wrong weight dtype (bf16 weight is not a supported layout)
163 + with pytest.raises(Exception, match="float32"):
164 + fp32_router_gemm(
165 + torch.randn(4, hidden_dim, dtype=torch.float32, device=device),
166 + mat_b.to(torch.bfloat16),
167 + )
168 +
169 + # fp16 activation (only fp32 / bf16 are accepted)
170 + with pytest.raises(Exception, match="float32 or bfloat16"):
171 + fp32_router_gemm(
172 + torch.randn(4, hidden_dim, dtype=torch.float16, device=device), mat_b
173 + )
vllm/model_executor/layers/fused_moe/router/gate_linear.py
+3 -3
@@ -15,8 +15,8 @@ class GateLinear(ReplicatedLinear):
15 """MoE gate linear layer with multi-tier GEMM dispatch:
16
17 1. DSV3 specialized kernel (SM90+, M<=16, H=7168 E=256/384, H=6144 E=256)
18 - 2. fp32 specialized kernel (SM90+, bf16/fp32 in, fp32 out,
19 - M<=32, H=3072, E=256)
18 + 2. fp32 specialized kernel (SM90+, bf16/fp32 in, fp32 out, M<=32,
19 + (H, E) in {(3072, 256), (6144, 128), (6144, 256)})
20 3. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype)
21 4. F.linear via ReplicatedLinear (ultimate fallback)
22
@@ -36,7 +36,7 @@ class GateLinear(ReplicatedLinear):
36
37 # (hidden_size, num_experts) pairs with an instantiated fp32 kernel:
38 # (3072, 256) -> MiniMax-M2/M2.5, (6144, 128) -> MiniMax-M3
39 - FP32_SUPPORTED_SHAPES = {(3072, 256), (6144, 128)}
39 + FP32_SUPPORTED_SHAPES = {(3072, 256), (6144, 128), (6144, 256)}
40 FP32_MAX_TOKENS = 32
41
42 def __init__(