@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 85e296950c

BF16x3 router GEMM (#47973)

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>

Thien Tran committed Jul 16, 2026 at 17:04 UTC 85e296950cff695695ccf06155942c80c89ef86c
5 files changed +555 -10
tests/kernels/test_bf16x3_router_gemm_cutedsl.py new
+43
@@ -0,0 +1,43 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for the experimental SM100 BF16x3 router GEMM."""
4 +
5 +import pytest
6 +import torch
7 +
8 +from vllm.utils.import_utils import has_cutedsl
9 +
10 +
11 +def _requires_sm100_cutedsl():
12 + if not torch.cuda.is_available():
13 + pytest.skip("CUDA not available")
14 + major, _ = torch.cuda.get_device_capability()
15 + if major != 10:
16 + pytest.skip("bf16x3 router GEMM requires SM100-class GPU")
17 + if not has_cutedsl():
18 + pytest.skip("cutedsl (cutlass) not installed")
19 +
20 +
21 +@pytest.mark.parametrize(
22 + ("num_tokens", "hidden_dim", "num_experts"),
23 + [(48, 6144, 128), (96, 3072, 256), (129, 3072, 17)],
24 +)
25 +def test_bf16x3_router_gemm_matches_reference(
26 + num_tokens: int, hidden_dim: int, num_experts: int
27 +):
28 + _requires_sm100_cutedsl()
29 + from vllm.model_executor.layers.fused_moe.router.bf16x3_router_gemm_cutedsl import ( # noqa: E501
30 + bf16x3_router_gemm,
31 + )
32 +
33 + torch.manual_seed(42)
34 + x = torch.randn(num_tokens, hidden_dim, dtype=torch.bfloat16, device="cuda")
35 + w = torch.randn(num_experts, hidden_dim, dtype=torch.float32, device="cuda")
36 + # Match the observed router weight scale
37 + w *= 0.053
38 + out = bf16x3_router_gemm(x, w)
39 + ref = torch.nn.functional.linear(x.float(), w)
40 +
41 + assert out.shape == (num_tokens, num_experts)
42 + assert out.dtype == torch.float32
43 + assert torch.mean(torch.abs(out - ref)).item() < 5e-6
vllm/config/kernel.py
+3
@@ -178,6 +178,9 @@ class KernelConfig:
178 enable_cutedsl_warmup: bool = True
179 """If True, run CuTeDSL compile warmup during kernel warmup."""
180
181 + enable_bf16x3_router_gemm: bool = False
182 + """If True, use the experimental SM100 BF16x3 CuteDSL router GEMM."""
183 +
184 moe_backend: MoEBackend = "auto"
185 """Backend for MoE expert computation kernels. Available options:
186
vllm/engine/arg_utils.py
+7
@@ -665,6 +665,7 @@ class EngineArgs:
665 enable_flashinfer_autotune: bool = get_field(
666 KernelConfig, "enable_flashinfer_autotune"
667 )
668 + enable_bf16x3_router_gemm: bool | None = None
669 worker_cls: str = ParallelConfig.worker_cls
670 worker_extension_cls: str = ParallelConfig.worker_extension_cls
671
@@ -1508,6 +1509,10 @@ class EngineArgs:
1509 "--enable-flashinfer-autotune",
1510 **kernel_kwargs["enable_flashinfer_autotune"],
1511 )
1512 + kernel_group.add_argument(
1513 + "--enable-bf16x3-router-gemm",
1514 + **kernel_kwargs["enable_bf16x3_router_gemm"],
1515 + )
1516 moe_backend_kwargs = kernel_kwargs["moe_backend"]
1517 moe_backend_kwargs["type"] = lambda s: s.lower().replace("-", "_")
1518 kernel_group.add_argument("--moe-backend", **moe_backend_kwargs)
@@ -2286,6 +2291,8 @@ class EngineArgs:
2291 "are mutually exclusive"
2292 )
2293 kernel_config.enable_flashinfer_autotune = self.enable_flashinfer_autotune
2294 + if self.enable_bf16x3_router_gemm is not None:
2295 + kernel_config.enable_bf16x3_router_gemm = self.enable_bf16x3_router_gemm
2296 if self.moe_backend != "auto":
2297 kernel_config.moe_backend = self.moe_backend
2298 if self.linear_backend != "auto":
vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py new
+449
@@ -0,0 +1,449 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""CuteDSL BF16x3 router GEMM.
4 +
5 +Computes ``X @ W.T`` for BF16 ``X`` with shape ``[N, K]`` and FP32 router
6 +weights ``W`` with shape ``[M, K]`` by decomposing each FP32 weight value into
7 +three BF16 residual terms inside the kernel, then accumulating the three BF16
8 +MMA results into FP32 TMEM output.
9 +"""
10 +
11 +from functools import cache
12 +
13 +import cutlass
14 +import torch
15 +from cuda.bindings.driver import CUstream
16 +from cutlass import BFloat16, Float32, Int32, Int64, Uint32, cute
17 +from cutlass._mlir.dialects import llvm
18 +from cutlass.cute.nvgpu import cpasync
19 +from cutlass.cutlass_dsl import T, dsl_user_op
20 +from quack.compile_utils import make_fake_tensor
21 +
22 +from vllm.cute_utils import _tcgen05, simple_tma_copy
23 +from vllm.triton_utils import tl, triton
24 +from vllm.utils import math_utils
25 +
26 +__all__ = ["bf16x3_router_gemm"]
27 +
28 +
29 +@dsl_user_op
30 +def _decompose_fp32x2_to_3xbf16x2(
31 + w0: Float32,
32 + w1: Float32,
33 + *,
34 + loc=None,
35 + ip=None,
36 +) -> tuple[Uint32, Uint32, Uint32]:
37 + # this PTX snippets does the following
38 + # out0 = BF16(in); res = in - FP32(out0)
39 + # out1 = BF16(res); res = res - FP32(out1)
40 + # out2 = BF16(res)
41 + #
42 + # for normal FP32, this decomposition is exact
43 + # i.e. in = FP32(out0) + FP32(out1) + FP32(out2)
44 + #
45 + out = llvm.inline_asm(
46 + llvm.StructType.get_literal([T.i32(), T.i32(), T.i32()]),
47 + [w0.ir_value(loc=loc, ip=ip), w1.ir_value(loc=loc, ip=ip)],
48 + "{\n\t"
49 + ".reg .b32 r1_lo, r1_hi, r2_lo, r2_hi;\n\t"
50 + ".reg .b32 w0_lo, w0_hi, w1_lo, w1_hi;\n\t"
51 + ".reg .b64 a_pair, w0_pair, w1_pair, r1_pair, r2_pair;\n\t"
52 + "cvt.rn.bf16x2.f32 $0, $4, $3;\n\t"
53 + "shl.b32 w0_lo, $0, 16;\n\t"
54 + "and.b32 w0_hi, $0, 0xffff0000;\n\t"
55 + "mov.b64 a_pair, {$3, $4};\n\t"
56 + "mov.b64 w0_pair, {w0_lo, w0_hi};\n\t"
57 + "sub.rn.f32x2 r1_pair, a_pair, w0_pair;\n\t"
58 + "mov.b64 {r1_lo, r1_hi}, r1_pair;\n\t"
59 + "cvt.rn.bf16x2.f32 $1, r1_hi, r1_lo;\n\t"
60 + "shl.b32 w1_lo, $1, 16;\n\t"
61 + "and.b32 w1_hi, $1, 0xffff0000;\n\t"
62 + "mov.b64 w1_pair, {w1_lo, w1_hi};\n\t"
63 + "sub.rn.f32x2 r2_pair, r1_pair, w1_pair;\n\t"
64 + "mov.b64 {r2_lo, r2_hi}, r2_pair;\n\t"
65 + "cvt.rn.bf16x2.f32 $2, r2_hi, r2_lo;\n\t"
66 + "}\n",
67 + "=r,=r,=r,f,f",
68 + has_side_effects=False,
69 + is_align_stack=False,
70 + loc=loc,
71 + ip=ip,
72 + )
73 + return (
74 + Uint32(llvm.extractvalue(T.i32(), out, [0], loc=loc, ip=ip)),
75 + Uint32(llvm.extractvalue(T.i32(), out, [1], loc=loc, ip=ip)),
76 + Uint32(llvm.extractvalue(T.i32(), out, [2], loc=loc, ip=ip)),
77 + )
78 +
79 +
80 +class Sm100BF16x3RouterGemm:
81 + def __init__(self, BN: int = 128) -> None:
82 + self.cta_tile = (BN, 128, 64)
83 + self.num_stages = 2
84 + self.num_warps = 10
85 +
86 + @cute.jit
87 + def _make_tma(self, tensor: cute.Tensor, BM: int, BK: int):
88 + op = cpasync.CopyBulkTensorTileG2SOp()
89 + swizzle_128B = cute.make_swizzle(3, 4, 3)
90 + elems = 128 * 8 // tensor.element_type.width # 128B
91 + slayout = cute.make_layout(
92 + (BM, (elems, BK // elems), self.num_stages),
93 + stride=(elems, (1, BM * elems), BM * BK),
94 + )
95 + slayout = cute.make_composed_layout(swizzle_128B, 0, slayout)
96 + return cpasync.make_tiled_tma_atom(op, tensor, slayout, (BM, BK))
97 +
98 + @cute.jit
99 + def __call__(
100 + self,
101 + X: cute.Tensor,
102 + W: cute.Tensor,
103 + out: cute.Tensor,
104 + split_k: Int32,
105 + stream: CUstream,
106 + ):
107 + BN, BM, BK = self.cta_tile
108 + W_tma = self._make_tma(W, BM, BK)
109 + X_tma = self._make_tma(X, BN, BK)
110 +
111 + grid_m = cute.ceil_div(W.shape[0], BM)
112 + grid_n = cute.ceil_div(X.shape[0], BN)
113 +
114 + self.kernel(X_tma, W_tma, out).launch(
115 + grid=(grid_m, grid_n, split_k),
116 + block=(self.num_warps * 32, 1, 1),
117 + stream=stream,
118 + use_pdl=True,
119 + )
120 +
121 + @cute.kernel
122 + def kernel(self, X_tma: cpasync.TmaInfo, W_tma: cpasync.TmaInfo, out: cute.Tensor):
123 + tid, _, _ = cute.arch.thread_idx()
124 + bid_m, bid_n, bid_k = cute.arch.block_idx()
125 + _, _, split_k = cute.arch.grid_dim()
126 +
127 + warp_id = cute.arch.make_warp_uniform(cute.arch.warp_idx())
128 + lane_id = cute.arch.lane_idx()
129 +
130 + BN, BM, BK = self.cta_tile
131 + num_stages = self.num_stages
132 +
133 + N, K = X_tma.tma_tensor.shape
134 + M, _ = W_tma.tma_tensor.shape
135 + k_tiles = cute.ceil_div(K, BK)
136 +
137 + smem = cutlass.utils.SmemAllocator()
138 + sX = smem.allocate_tensor(
139 + BFloat16,
140 + X_tma.smem_layout.outer,
141 + byte_alignment=128,
142 + swizzle=X_tma.smem_layout.inner,
143 + )
144 + sW = smem.allocate_tensor(
145 + Float32,
146 + W_tma.smem_layout.outer,
147 + byte_alignment=128,
148 + swizzle=W_tma.smem_layout.inner,
149 + )
150 +
151 + tma_full_mbar = smem.allocate_array(Int64, num_stages)
152 + tma_empty_mbar = smem.allocate_array(Int64, num_stages)
153 + w_full_mbar = smem.allocate_array(Int64, num_stages)
154 + mma_mbar = smem.allocate_array(Int64, 1)
155 + taddr = smem.allocate(Int32, 4)
156 +
157 + BAR_TMEM_ALLOC = 1
158 + BAR_PREP = 2
159 + BAR_EPI = 3
160 +
161 + # tmem "allocation"
162 + # acc_main is for the 1st BF16 term. acc_res is for the 2nd and 3rd
163 + # BF16 terms, which are much smaller than the first term.
164 + acc_main = 0
165 + acc_res = BN
166 + w_tmem_base = BN * 2
167 +
168 + if warp_id == 0:
169 + with cute.arch.elect_one():
170 + for i in cutlass.range_constexpr(num_stages):
171 + cute.arch.mbarrier_init(tma_full_mbar + i, 1)
172 + cute.arch.mbarrier_init(tma_empty_mbar + i, 1)
173 + cute.arch.mbarrier_init(w_full_mbar + i, 128)
174 + cute.arch.mbarrier_init(mma_mbar, 1)
175 + cute.arch.mbarrier_init_fence()
176 + elif warp_id == 1:
177 + cpasync.prefetch_descriptor(X_tma.atom)
178 + cpasync.prefetch_descriptor(W_tma.atom)
179 + cute.arch.sync_threads()
180 + cute.arch.griddepcontrol_wait()
181 +
182 + if warp_id == 9:
183 + # TMA warp
184 + stage_id = 0
185 + parity = 1
186 +
187 + # (BM, BK, K/BK)
188 + gW_tiles = cute.local_tile(W_tma.tma_tensor, (BM, BK), (bid_m, None))
189 + gX_tiles = cute.local_tile(X_tma.tma_tensor, (BN, BK), (bid_n, None))
190 +
191 + for tile_k in cutlass.range(bid_k, k_tiles, split_k, unroll=1):
192 + cute.arch.mbarrier_wait(tma_empty_mbar + stage_id, parity)
193 + mbar = tma_full_mbar + stage_id
194 + with cute.arch.elect_one():
195 + stage_bytes = BN * BK * 2 + BM * BK * 4
196 + cute.arch.mbarrier_arrive_and_expect_tx(mbar, stage_bytes)
197 + simple_tma_copy(
198 + W_tma.atom,
199 + gW_tiles[None, None, tile_k],
200 + sW[None, None, stage_id],
201 + mbar,
202 + )
203 + simple_tma_copy(
204 + X_tma.atom,
205 + gX_tiles[None, None, tile_k],
206 + sX[None, None, stage_id],
207 + mbar,
208 + )
209 +
210 + stage_id = (stage_id + 1) % num_stages
211 + if stage_id == 0:
212 + parity ^= 1
213 +
214 + elif warp_id == 8:
215 + # MMA warp
216 + stage_id = 0
217 + parity = 0
218 +
219 + idesc = _tcgen05.make_bf16_idesc(BM, BN)
220 + sdesc = _tcgen05.make_sdesc_128B_swizzle(0)
221 +
222 + for tile_k in cutlass.range(bid_k, k_tiles, split_k, unroll=1):
223 + cute.arch.mbarrier_wait(tma_full_mbar + stage_id, parity)
224 + cute.arch.mbarrier_wait(w_full_mbar + stage_id, parity)
225 + _tcgen05.fence_after_thread_sync()
226 +
227 + w_tmem = w_tmem_base + stage_id * (BK // 2 * 3)
228 + x_desc = sdesc | (sX[None, None, stage_id].iterator.toint() >> 4)
229 +
230 + for k in cutlass.range_constexpr(BK // 16):
231 + enable_d = (tile_k > bid_k) or (k > 0)
232 + _tcgen05.mma_ts_f16(
233 + acc_main, w_tmem + k * 8, x_desc, idesc, enable_d
234 + )
235 + _tcgen05.mma_ts_f16(
236 + acc_res, w_tmem + 32 + k * 8, x_desc, idesc, enable_d
237 + )
238 + _tcgen05.mma_ts_f16(
239 + acc_res, w_tmem + 64 + k * 8, x_desc, idesc, True
240 + )
241 + x_desc += 32 >> 4
242 +
243 + _tcgen05.commit(tma_empty_mbar + stage_id)
244 +
245 + stage_id = (stage_id + 1) % self.num_stages
246 + if stage_id == 0:
247 + parity ^= 1
248 +
249 + _tcgen05.commit(mma_mbar)
250 +
251 + elif warp_id >= 4:
252 + # prep warps: decompose FP32 W into 3xBF16
253 + warp_id_ = warp_id % 4
254 +
255 + stage_id = 0
256 + parity = 0
257 +
258 + # ld.shared.v4.f32
259 + op = cute.nvgpu.CopyUniversalOp()
260 + cp_atom = cute.make_copy_atom(op, Float32, num_bits_per_copy=128)
261 +
262 + # sW_view: ((4, 1), (BK/4, num_stages))
263 + row = warp_id_ * 32 + lane_id
264 + sW_view = cute.zipped_divide(sW[row, None, None], (4, 1))
265 +
266 + for _ in cutlass.range(bid_k, k_tiles, split_k, unroll=1):
267 + if warp_id_ == 0:
268 + cute.arch.mbarrier_wait(tma_full_mbar + stage_id, parity)
269 + cute.arch.barrier(barrier_id=BAR_PREP, number_of_threads=128)
270 +
271 + row = warp_id_ * 32 + lane_id
272 + w_tmem = w_tmem_base + stage_id * (BK // 2 * 3)
273 + for kblock in cutlass.range_constexpr(BK // 4):
274 + w0 = cute.make_rmem_tensor(2, Uint32)
275 + w1 = cute.make_rmem_tensor(2, Uint32)
276 + w2 = cute.make_rmem_tensor(2, Uint32)
277 +
278 + w_tmp = cute.make_rmem_tensor(4, Float32)
279 + cute.copy(cp_atom, sW_view[None, (kblock, stage_id)], w_tmp)
280 + w0[0], w1[0], w2[0] = _decompose_fp32x2_to_3xbf16x2(
281 + w_tmp[0], w_tmp[1]
282 + )
283 + w0[1], w1[1], w2[1] = _decompose_fp32x2_to_3xbf16x2(
284 + w_tmp[2], w_tmp[3]
285 + )
286 +
287 + tcol = kblock * 2
288 + _tcgen05.st(warp_id_ * 32, w_tmem + 0 + tcol, "32x32b", 2, w0)
289 + _tcgen05.st(warp_id_ * 32, w_tmem + 32 + tcol, "32x32b", 2, w1)
290 + _tcgen05.st(warp_id_ * 32, w_tmem + 64 + tcol, "32x32b", 2, w2)
291 +
292 + _tcgen05.wait_st()
293 + _tcgen05.fence_before_thread_sync()
294 + cute.arch.mbarrier_arrive(w_full_mbar + stage_id)
295 +
296 + stage_id = (stage_id + 1) % self.num_stages
297 + if stage_id == 0:
298 + parity ^= 1
299 +
300 + else:
301 + # epilogue warps
302 + if warp_id == 0:
303 + _tcgen05.alloc(taddr)
304 + cute.arch.barrier(barrier_id=BAR_TMEM_ALLOC, number_of_threads=128)
305 +
306 + if warp_id == 0:
307 + cute.arch.mbarrier_wait(mma_mbar, 0)
308 + cute.arch.barrier(barrier_id=BAR_EPI, number_of_threads=128)
309 + _tcgen05.fence_after_thread_sync()
310 +
311 + cute.arch.griddepcontrol_launch_dependents()
312 +
313 + WIDTH = 8
314 + for i in cutlass.range_constexpr(BN // WIDTH):
315 + tcol = i * WIDTH
316 + main_regs = cute.make_rmem_tensor(WIDTH, Float32)
317 + res_regs = cute.make_rmem_tensor(WIDTH, Float32)
318 + main_regs.store(_tcgen05.ld(warp_id * 32, tcol, "32x32b", WIDTH))
319 + res_regs.store(_tcgen05.ld(warp_id * 32, BN + tcol, "32x32b", WIDTH))
320 + _tcgen05.wait_ld()
321 +
322 + # CuteDSL will codegen add.f32x2
323 + for j in cutlass.range(WIDTH, vectorize=True):
324 + main_regs[j] += res_regs[j]
325 +
326 + w_row_idx = bid_m * BM + tid
327 + for j in cutlass.range_constexpr(WIDTH):
328 + x_row_idx = bid_n * BN + i * WIDTH + j
329 + if x_row_idx < N and w_row_idx < M:
330 + out[bid_k, x_row_idx, w_row_idx] = main_regs[j]
331 +
332 + cute.arch.barrier(barrier_id=BAR_EPI, number_of_threads=128)
333 + if warp_id == 0:
334 + _tcgen05.dealloc()
335 +
336 + @cache
337 + @staticmethod
338 + def compile(BN: int = 128, K: int = 6144):
339 + N = cute.sym_int()
340 + M = cute.sym_int()
341 + SPLIT_K = cute.sym_int()
342 + X = make_fake_tensor(BFloat16, (N, K), divisibility=8)
343 + W = make_fake_tensor(Float32, (M, K), divisibility=4)
344 + out = make_fake_tensor(Float32, (SPLIT_K, N, M), divisibility=1)
345 + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
346 + kernel = Sm100BF16x3RouterGemm(BN)
347 + return cute.compile(
348 + kernel, X, W, out, Int32(1), stream, options="--enable-tvm-ffi"
349 + )
350 +
351 +
352 +@triton.jit
353 +def _splitk_reduce_kernel(
354 + partials,
355 + out,
356 + N,
357 + M: tl.constexpr,
358 + split_stride,
359 + k_splits,
360 + BN: tl.constexpr,
361 + BM: tl.constexpr,
362 + BS: tl.constexpr,
363 + USE_PDL: tl.constexpr,
364 +):
365 + pid_n = tl.program_id(0)
366 + pid_m = tl.program_id(1)
367 + offs_n = pid_n * BN + tl.arange(0, BN)
368 + offs_m = pid_m * BM + tl.arange(0, BM)
369 + offs_s = tl.arange(0, BS)
370 +
371 + if USE_PDL:
372 + tl.extra.cuda.gdc_wait()
373 + tl.extra.cuda.gdc_launch_dependents()
374 +
375 + vals = tl.load(
376 + partials
377 + + offs_s[:, None, None] * split_stride
378 + + offs_n[None, :, None] * M
379 + + offs_m[None, None, :],
380 + mask=(
381 + (offs_s[:, None, None] < k_splits)
382 + & (offs_n[None, :, None] < N)
383 + & (offs_m[None, None, :] < M)
384 + ),
385 + other=0.0,
386 + )
387 + acc = tl.sum(vals, axis=0)
388 + tl.store(
389 + out + offs_n[:, None] * M + offs_m[None, :],
390 + acc,
391 + mask=(offs_n[:, None] < N) & (offs_m[None, :] < M),
392 + )
393 +
394 +
395 +def splitk_reduce_triton(partials: torch.Tensor, out: torch.Tensor):
396 + split_k, N, M = partials.shape
397 + block_s = 1 << (split_k - 1).bit_length()
398 + split_stride = partials.stride(0)
399 + if block_s >= 64:
400 + BN, BM = 1, 32
401 + elif block_s >= 8:
402 + BN, BM = 1, 256
403 + else:
404 + BN, BM = min(16, 32 // block_s), 32
405 + grid = (triton.cdiv(N, BN), triton.cdiv(M, BM))
406 + _splitk_reduce_kernel[grid](
407 + partials,
408 + out,
409 + N,
410 + M,
411 + split_stride,
412 + split_k,
413 + BN=BN,
414 + BM=BM,
415 + BS=block_s,
416 + USE_PDL=True,
417 + num_warps=4,
418 + launch_pdl=True,
419 + )
420 +
421 +
422 +def bf16x3_router_gemm(X: torch.Tensor, W: torch.Tensor) -> torch.Tensor:
423 + """Return ``X @ W.T`` using the SM100 BF16x3 router GEMM kernel."""
424 + N, K = X.shape
425 + M, _ = W.shape
426 + num_sms = torch.cuda.get_device_properties(X.device).multi_processor_count
427 +
428 + # next power of 2 within 8 and 128
429 + BN = triton.next_power_of_2(N)
430 + BN = min(max(BN, 8), 128)
431 +
432 + BM = 128
433 + BK = 64
434 + k_tiles = math_utils.cdiv(K, BK)
435 + grid_m = math_utils.cdiv(M, BM)
436 + grid_n = math_utils.cdiv(N, BN)
437 +
438 + base_ctas = grid_m * grid_n
439 + split_k = min(k_tiles, max(1, num_sms // base_ctas))
440 +
441 + partials = X.new_empty(split_k, N, M, dtype=torch.float32)
442 + Sm100BF16x3RouterGemm.compile(BN, K)(X, W, partials, split_k)
443 +
444 + if split_k == 1:
445 + return partials.squeeze(0)
446 +
447 + out = X.new_empty(N, M, dtype=torch.float32)
448 + splitk_reduce_triton(partials, out)
449 + return out
vllm/model_executor/layers/fused_moe/router/gate_linear.py
+53 -10
@@ -4,11 +4,15 @@ import torch
4 from torch.nn.parameter import Parameter
5
6 import vllm._custom_ops as ops
7 +from vllm.config import get_current_vllm_config_or_none
8 +from vllm.logger import init_logger
9 from vllm.model_executor.custom_op import PluggableLayer
10 from vllm.model_executor.layers.linear import ReplicatedLinear
11 from vllm.platforms import current_platform
12 from vllm.utils.torch_utils import direct_register_custom_op
13
14 +logger = init_logger(__name__)
15 +
16
17 @PluggableLayer.register("gate_linear")
18 class GateLinear(ReplicatedLinear):
@@ -19,8 +23,9 @@ class GateLinear(ReplicatedLinear):
23 2. DSV3 specialized kernel (SM90+, M<=16, H=7168 E=256/384, H=6144 E=256)
24 3. fp32 specialized kernel (SM90+, bf16/fp32 in, fp32 out, M<=32,
25 (H, E) in {(3072, 256), (6144, 128), (6144, 256)})
22 - 4. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype)
23 - 5. F.linear via ReplicatedLinear (ultimate fallback)
26 + 4. experimental bf16x3 CuteDSL kernel (opt-in, SM100, bf16 in, fp32 weight)
27 + 5. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype)
28 + 6. F.linear via ReplicatedLinear (ultimate fallback)
29
30 The ``out_dtype`` attribute is mutable and can be set after init
31 (e.g. when the required dtype depends on the expert quantization
@@ -86,6 +91,11 @@ class GateLinear(ReplicatedLinear):
91 self._dsv3_max_batch = 16 if is_hopper else 8
92
93 # fp32 specialized kernel eligibility (SM90+, exact dims, fp32 weight)
94 + vllm_config = get_current_vllm_config_or_none()
95 + enable_bf16x3_router_gemm = (
96 + vllm_config is not None
97 + and vllm_config.kernel_config.enable_bf16x3_router_gemm
98 + )
99 self.allow_fp32_router_gemm = (
100 not bias
101 and self.weight.dtype == torch.float32
@@ -93,6 +103,16 @@ class GateLinear(ReplicatedLinear):
103 and (is_hopper or is_blackwell)
104 and (input_size, output_size) in self.FP32_SUPPORTED_SHAPES
105 )
106 + self.allow_bf16x3_router_gemm = (
107 + not bias
108 + and self.weight.dtype == torch.float32
109 + and current_platform.is_cuda()
110 + and is_blackwell
111 + and input_size % 8 == 0
112 + and enable_bf16x3_router_gemm
113 + )
114 + if self.allow_bf16x3_router_gemm:
115 + logger.info_once("Enabled experimental SM100 BF16x3 router GEMM.")
116
117 # cuBLAS bf16→fp32 eligibility
118 self.allow_cublas_router_gemm = (
@@ -173,15 +193,26 @@ class GateLinear(ReplicatedLinear):
193 torch.float32,
194 torch.bfloat16,
195 ):
176 - output = torch.ops.vllm.fp32_router_gemm_dispatch(x, self.weight)
196 + output = torch.ops.vllm.fp32_router_gemm_dispatch(
197 + x, self.weight, self.allow_bf16x3_router_gemm
198 + )
199 return output, None
200
179 - # Tier 4: cuBLAS bf16→fp32
201 + # Tier 4: experimental bf16x3 CuteDSL kernel for fp32 router weights
202 + if self.allow_bf16x3_router_gemm and x.dtype == torch.bfloat16:
203 + from vllm.model_executor.layers.fused_moe.router.bf16x3_router_gemm_cutedsl import ( # noqa: E501
204 + bf16x3_router_gemm,
205 + )
206 +
207 + output = bf16x3_router_gemm(x, self.weight)
208 + return output, None
209 +
210 + # Tier 5: cuBLAS bf16→fp32
211 if self.allow_cublas_router_gemm and x.dtype == torch.bfloat16:
212 output = torch.mm(x, self.weight.T, out_dtype=torch.float32)
213 return output, None
214
184 - # Tier 5: F.linear (ReplicatedLinear)
215 + # Tier 6: F.linear (ReplicatedLinear)
216 if self.out_dtype is not None and x.dtype != self.weight.dtype:
217 x = x.to(self.weight.dtype)
218 output, output_bias = super().forward(x)
@@ -194,22 +225,34 @@ _FP32_ROUTER_GEMM_MAX_TOKENS = GateLinear.FP32_MAX_TOKENS
225
226
227 def fp32_router_gemm_dispatch_impl(
197 - x: torch.Tensor, weight: torch.Tensor
228 + x: torch.Tensor,
229 + weight: torch.Tensor,
230 + allow_bf16x3_router_gemm: bool,
231 ) -> torch.Tensor:
232 """
233 Dynamically run fp32 specialized gemm if num_tokens <= FP32_MAX_TOKENS,
201 - otherwise fall back to F.linear.
234 + otherwise optionally run the experimental BF16x3 kernel for medium/large
235 + SM100 router batches, then fall back to F.linear.
236 This must be wrapped in a custom op because our torch.compile integration
237 does not support runtime dispatching on num_tokens.
238 """
239 if x.shape[0] <= _FP32_ROUTER_GEMM_MAX_TOKENS:
240 return ops.fp32_router_gemm(x, weight)
207 - else:
208 - return torch.nn.functional.linear(x.float(), weight)
241 +
242 + if allow_bf16x3_router_gemm and x.dtype == torch.bfloat16:
243 + from vllm.model_executor.layers.fused_moe.router.bf16x3_router_gemm_cutedsl import ( # noqa: E501
244 + bf16x3_router_gemm,
245 + )
246 +
247 + return bf16x3_router_gemm(x, weight)
248 +
249 + return torch.nn.functional.linear(x.float(), weight)
250
251
252 def fp32_router_gemm_dispatch_fake(
212 - x: torch.Tensor, weight: torch.Tensor
253 + x: torch.Tensor,
254 + weight: torch.Tensor,
255 + allow_bf16x3_router_gemm: bool,
256 ) -> torch.Tensor:
257 return x.new_empty((x.shape[0], weight.shape[0]), dtype=torch.float32)
258