@tduncan / Nvidia-Optimization-Stack-VLLM / commits / cc4210f671

[Perf] Improve BF16x3 router GEMM accuracy and make it default on sm100 (#55899)

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>

Thien Tran committed Sep 9, 2026 at 22:33 UTC cc4210f6719e3285fdfe4da862ae04c8dbb77ab1

6 files changed +314 -88

tests/kernels/test_bf16x3_router_gemm_cutedsl.py

+14 -4

@@ -1,6 +1,6 @@

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."""

3 +"""Tests for the SM100 BF16x3 router GEMM."""

5 import pytest

6 import torch

@@ -20,7 +20,16 @@ def _requires_sm100_cutedsl():

21 @pytest.mark.parametrize(

22 ("num_tokens", "hidden_dim", "num_experts"),

23 - [(48, 6144, 128), (96, 3072, 256), (129, 3072, 17)],

23 + [

24 + (48, 6144, 128),

25 + (96, 3072, 256),

26 + (129, 3072, 17),

27 + # long-K cases exercise the multi-accumulation path (the split-K

28 + # heuristic leaves chains of 15 and 32 K-tiles here, above the

29 + # kernel's num_tmem_acc bound)

30 + (1024, 8192, 256),

31 + (2048, 8192, 256),

32 + ],

33 )

34 def test_bf16x3_router_gemm_matches_reference(

35 num_tokens: int, hidden_dim: int, num_experts: int

@@ -36,8 +45,9 @@ def test_bf16x3_router_gemm_matches_reference(

45 # Match the observed router weight scale

46 w *= 0.053

47 out = bf16x3_router_gemm(x, w)

39 - ref = torch.nn.functional.linear(x.float(), w)

48 + # FP64 reference: the FP32 reference itself drifts by ~5e-6 at N=2048

49 + ref = torch.nn.functional.linear(x.double(), w.double())

51 assert out.shape == (num_tokens, num_experts)

52 assert out.dtype == torch.float32

43 - assert torch.mean(torch.abs(out - ref)).item() < 5e-6

53 + assert torch.mean(torch.abs(out.double() - ref)).item() < 5e-6

vllm/config/kernel.py

-3

@@ -230,9 +230,6 @@ class KernelConfig:

230 enable_jit_warmup: bool = True

231 """If True, run JIT compile warmup during kernel warmup."""

232

233 - enable_bf16x3_router_gemm: bool = False

234 - """If True, use the experimental SM100 BF16x3 CuteDSL router GEMM."""

235 -

233 moe_backend: MoEBackend = "auto"

234 """Backend for MoE expert computation kernels. Available options:

235

vllm/engine/arg_utils.py

-7

@@ -695,7 +695,6 @@ class EngineArgs:

695 enable_flashinfer_autotune: bool = get_field(

696 KernelConfig, "enable_flashinfer_autotune"

697 )

698 - enable_bf16x3_router_gemm: bool | None = None

698 worker_cls: str = ParallelConfig.worker_cls

699 worker_extension_cls: str = ParallelConfig.worker_extension_cls

700

@@ -1655,10 +1654,6 @@ class EngineArgs:

1654 "--enable-flashinfer-autotune",

1655 **kernel_kwargs["enable_flashinfer_autotune"],

1656 )

1658 - kernel_group.add_argument(

1659 - "--enable-bf16x3-router-gemm",

1660 - **kernel_kwargs["enable_bf16x3_router_gemm"],

1661 - )

1657 moe_backend_kwargs = kernel_kwargs["moe_backend"]

1658 moe_backend_kwargs["type"] = lambda s: s.lower().replace("-", "_")

1659 kernel_group.add_argument("--moe-backend", **moe_backend_kwargs)

@@ -2519,8 +2514,6 @@ class EngineArgs:

2514 "are mutually exclusive"

2515 )

2516 kernel_config.enable_flashinfer_autotune = self.enable_flashinfer_autotune

2522 - if self.enable_bf16x3_router_gemm is not None:

2523 - kernel_config.enable_bf16x3_router_gemm = self.enable_bf16x3_router_gemm

2517 if self.moe_backend != "auto":

2518 kernel_config.moe_backend = self.moe_backend

2519 if self.linear_backend != "auto":

vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py

+249 -56

@@ -6,6 +6,10 @@ 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 +The TMEM accumulation chain is bounded to ``num_tmem_acc`` K-tiles: at chunk

11 +boundaries the epilogue warps drain the main-term accumulator into a register

12 +master accumulator and the MMA warp resets it.

13 """

15 from functools import cache

@@ -22,6 +26,7 @@ from quack.compile_utils import make_fake_tensor

26 from vllm.cute_utils import _tcgen05, simple_tma_copy

27 from vllm.triton_utils import tl, triton

28 from vllm.utils import math_utils

29 +from vllm.utils.torch_utils import direct_register_custom_op

31 __all__ = ["bf16x3_router_gemm"]

@@ -82,6 +87,7 @@ class Sm100BF16x3RouterGemm:

87 self.cta_tile = (BN, 128, 64)

88 self.num_stages = 2

89 self.num_warps = 10

90 + self.num_tmem_acc = 8

92 @cute.jit

93 def _make_tma(self, tensor: cute.Tensor, BM: int, BK: int):

@@ -129,6 +135,7 @@ class Sm100BF16x3RouterGemm:

135

136 BN, BM, BK = self.cta_tile

137 num_stages = self.num_stages

138 + num_tmem_acc = self.num_tmem_acc

139

140 N, K = X_tma.tma_tensor.shape

141 M, _ = W_tma.tma_tensor.shape

@@ -151,7 +158,8 @@ class Sm100BF16x3RouterGemm:

158 tma_full_mbar = smem.allocate_array(Int64, num_stages)

159 tma_empty_mbar = smem.allocate_array(Int64, num_stages)

160 w_full_mbar = smem.allocate_array(Int64, num_stages)

154 - mma_mbar = smem.allocate_array(Int64, 1)

161 + acc_full_mbar = smem.allocate_array(Int64, 1)

162 + acc_empty_mbar = smem.allocate_array(Int64, 1)

163 taddr = smem.allocate(Int32, 4)

164

165 BAR_TMEM_ALLOC = 1

@@ -171,7 +179,8 @@ class Sm100BF16x3RouterGemm:

179 cute.arch.mbarrier_init(tma_full_mbar + i, 1)

180 cute.arch.mbarrier_init(tma_empty_mbar + i, 1)

181 cute.arch.mbarrier_init(w_full_mbar + i, 128)

174 - cute.arch.mbarrier_init(mma_mbar, 1)

182 + cute.arch.mbarrier_init(acc_full_mbar, 1)

183 + cute.arch.mbarrier_init(acc_empty_mbar, 128)

184 cute.arch.mbarrier_init_fence()

185 elif warp_id == 1:

186 cpasync.prefetch_descriptor(X_tma.atom)

@@ -213,40 +222,52 @@ class Sm100BF16x3RouterGemm:

222

223 elif warp_id == 8:

224 # MMA warp

216 - stage_id = 0

217 - parity = 0

225 + tma_stage_id = 0

226 + tma_parity = 0

227 +

228 + tmem_acc_count = 0

229 + tmem_parity = 1

230

231 idesc = _tcgen05.make_bf16_idesc(BM, BN)

232 sdesc = _tcgen05.make_sdesc_128B_swizzle(0)

233

234 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)

235 + if tmem_acc_count == 0:

236 + cute.arch.mbarrier_wait(acc_empty_mbar, tmem_parity)

237 + cute.arch.mbarrier_wait(tma_full_mbar + tma_stage_id, tma_parity)

238 + cute.arch.mbarrier_wait(w_full_mbar + tma_stage_id, tma_parity)

239 _tcgen05.fence_after_thread_sync()

240

227 - w_tmem = w_tmem_base + stage_id * (BK // 2 * 3)

228 - x_desc = sdesc | (sX[None, None, stage_id].iterator.toint() >> 4)

241 + w_tmem = w_tmem_base + tma_stage_id * (BK // 2 * 3)

242 + x_desc = sdesc | (sX[None, None, tma_stage_id].iterator.toint() >> 4)

243

244 for k in cutlass.range_constexpr(BK // 16):

231 - enable_d = (tile_k > bid_k) or (k > 0)

245 + enable_d_main = (tmem_acc_count > 0) or (k > 0)

246 + enable_d_res = (tile_k > bid_k) or (k > 0)

247 _tcgen05.mma_ts_f16(

233 - acc_main, w_tmem + k * 8, x_desc, idesc, enable_d

248 + acc_main, w_tmem + k * 8, x_desc, idesc, enable_d_main

249 )

250 _tcgen05.mma_ts_f16(

236 - acc_res, w_tmem + 32 + k * 8, x_desc, idesc, enable_d

251 + acc_res, w_tmem + 32 + k * 8, x_desc, idesc, enable_d_res

252 )

253 _tcgen05.mma_ts_f16(

254 acc_res, w_tmem + 64 + k * 8, x_desc, idesc, True

255 )

256 x_desc += 32 >> 4

257

243 - _tcgen05.commit(tma_empty_mbar + stage_id)

258 + _tcgen05.commit(tma_empty_mbar + tma_stage_id)

259

245 - stage_id = (stage_id + 1) % self.num_stages

246 - if stage_id == 0:

247 - parity ^= 1

260 + tma_stage_id = (tma_stage_id + 1) % num_stages

261 + if tma_stage_id == 0:

262 + tma_parity ^= 1

263 +

264 + tmem_acc_count = (tmem_acc_count + 1) % num_tmem_acc

265 + if tmem_acc_count == 0:

266 + _tcgen05.commit(acc_full_mbar)

267 + tmem_parity ^= 1

268

249 - _tcgen05.commit(mma_mbar)

269 + if tmem_acc_count != 0:

270 + _tcgen05.commit(acc_full_mbar)

271

272 elif warp_id >= 4:

273 # prep warps: decompose FP32 W into 3xBF16

@@ -303,39 +324,93 @@ class Sm100BF16x3RouterGemm:

324 _tcgen05.alloc(taddr)

325 cute.arch.barrier(barrier_id=BAR_TMEM_ALLOC, number_of_threads=128)

326

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()

327 + tiles_local = cute.ceil_div(k_tiles - bid_k, split_k)

328 + num_chunks = cute.ceil_div(tiles_local, num_tmem_acc)

329

330 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]

331 + main_regs = cute.make_rmem_tensor(WIDTH, Float32)

332 + res_regs = cute.make_rmem_tensor(WIDTH, Float32)

333

334 + if num_chunks == 1:

335 + # single chunk

336 + cute.arch.mbarrier_wait(acc_full_mbar, 0)

337 + _tcgen05.fence_after_thread_sync()

338 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]

339 + for i in cutlass.range_constexpr(BN // WIDTH):

340 + tcol = i * WIDTH

341 + main_regs.store(_tcgen05.ld(warp_id * 32, tcol, "32x32b", WIDTH))

342 + res_regs.store(

343 + _tcgen05.ld(warp_id * 32, BN + tcol, "32x32b", WIDTH)

344 + )

345 + _tcgen05.wait_ld()

346 +

347 + # CuteDSL will codegen add.f32x2

348 + for j in cutlass.range(WIDTH, vectorize=True):

349 + main_regs[j] += res_regs[j]

350 +

351 + for j in cutlass.range_constexpr(WIDTH):

352 + x_row_idx = bid_n * BN + i * WIDTH + j

353 + if x_row_idx < N and w_row_idx < M:

354 + out[bid_k, x_row_idx, w_row_idx] = main_regs[j]

355 +

356 + else:

357 + # multiple chunks

358 + master_acc = cute.make_rmem_tensor(BN, Float32)

359 +

360 + # chunk 0: pure TMEM load, no accumulation.

361 + cute.arch.mbarrier_wait(acc_full_mbar, 0)

362 + _tcgen05.fence_after_thread_sync()

363 + master_acc.store(_tcgen05.ld(warp_id * 32, 0, "32x32b", BN))

364 + _tcgen05.wait_ld()

365 + _tcgen05.fence_before_thread_sync()

366 + cute.arch.mbarrier_arrive(acc_empty_mbar)

367 +

368 + # accumulate main tmem acc

369 + for chunk in cutlass.range(1, num_chunks, unroll=1):

370 + cute.arch.mbarrier_wait(acc_full_mbar, chunk & 1)

371 + _tcgen05.fence_after_thread_sync()

372 +

373 + for i in cutlass.range_constexpr(BN // WIDTH):

374 + tcol = i * WIDTH

375 + main_regs.store(

376 + _tcgen05.ld(warp_id * 32, tcol, "32x32b", WIDTH)

377 + )

378 + _tcgen05.wait_ld()

379 + # CuteDSL refuses to compile

380 + # master_acc[i * WIDTH + j] += main_regs[j]

381 + # when vectorize=True

382 + master_view = cute.local_tile(master_acc, (WIDTH,), (i,))

383 + for j in cutlass.range(WIDTH, vectorize=True):

384 + master_view[j] += main_regs[j]

385 +

386 + _tcgen05.fence_before_thread_sync()

387 + cute.arch.mbarrier_arrive(acc_empty_mbar)

388 +

389 + # fold in the residual accumulator

390 + for i in cutlass.range_constexpr(BN // WIDTH):

391 + tcol = i * WIDTH

392 + res_regs.store(

393 + _tcgen05.ld(warp_id * 32, BN + tcol, "32x32b", WIDTH)

394 + )

395 + _tcgen05.wait_ld()

396 + master_view = cute.local_tile(master_acc, (WIDTH,), (i,))

397 + for j in cutlass.range(WIDTH, vectorize=True):

398 + master_view[j] += res_regs[j]

399

400 + w_row_idx = bid_m * BM + tid

401 + for j in cutlass.range_constexpr(WIDTH):

402 + x_row_idx = bid_n * BN + i * WIDTH + j

403 + if x_row_idx < N and w_row_idx < M:

404 + out[bid_k, x_row_idx, w_row_idx] = master_acc[i * WIDTH + j]

405 +

406 + cute.arch.griddepcontrol_launch_dependents()

407 cute.arch.barrier(barrier_id=BAR_EPI, number_of_threads=128)

408 if warp_id == 0:

409 _tcgen05.dealloc()

410

411 @cache

412 @staticmethod

338 - def compile(BN: int = 128, K: int = 6144):

413 + def compile(K: int, BN: int = 128):

414 N = cute.sym_int()

415 M = cute.sym_int()

416 SPLIT_K = cute.sym_int()

@@ -392,16 +467,19 @@ def _splitk_reduce_kernel(

467 )

468

469

470 +def _splitk_reduce_config(split_k: int) -> tuple[int, int, int]:

471 + block_s = triton.next_power_of_2(split_k)

472 + if block_s >= 64:

473 + return 1, 32, block_s

474 + if block_s >= 8:

475 + return 1, 256, block_s

476 + return min(16, 32 // block_s), 32, block_s

477 +

478 +

479 def splitk_reduce_triton(partials: torch.Tensor, out: torch.Tensor):

480 split_k, N, M = partials.shape

397 - block_s = 1 << (split_k - 1).bit_length()

481 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

482 + BN, BM, block_s = _splitk_reduce_config(split_k)

483 grid = (triton.cdiv(N, BN), triton.cdiv(M, BM))

484 _splitk_reduce_kernel[grid](

485 partials,

@@ -419,27 +497,124 @@ def splitk_reduce_triton(partials: torch.Tensor, out: torch.Tensor):

497 )

498

499

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

500 +# Accuracy-constrained (BN, split_k) configs from GB300 sweeps.

501 +_TILE_CONFIG_OVERRIDES = {

502 + (2816, 256): { # Hunyuan-V4

503 + 64: (16, 16),

504 + 128: (32, 16),

505 + 256: (64, 16),

506 + 512: (128, 16),

507 + },

508 + (3072, 256): { # MiniMax-M2

509 + 64: (16, 16),

510 + 128: (32, 16),

511 + 256: (64, 16),

512 + 512: (128, 16),

513 + },

514 + (4096, 192): { # Hunyuan-V3

515 + 64: (32, 32),

516 + 128: (64, 32),

517 + 256: (128, 32),

518 + 512: (128, 19),

519 + },

520 + (6144, 128): { # MiniMax-M3

521 + 64: (32, 64),

522 + 128: (64, 64),

523 + 256: (64, 32),

524 + 512: (128, 32),

525 + },

526 +}

527 +

528 +

529 +def _pick_tile_config(N: int, K: int, M: int, num_sms: int) -> tuple[int, int]:

530 + """Return (BN, split_k): tuned table mid-range, generic rule otherwise."""

531 + k_tiles = math_utils.cdiv(K, 64)

532 +

533 + if 32 < N <= 512:

534 + token_bucket = triton.next_power_of_2(N)

535 + overrides = _TILE_CONFIG_OVERRIDES.get((K, M))

536 + if overrides is not None:

537 + BN, split_k = overrides[token_bucket]

538 + return BN, min(split_k, k_tiles)

539

540 # next power of 2 within 8 and 128

541 BN = triton.next_power_of_2(N)

542 BN = min(max(BN, 8), 128)

543

432 - BM = 128

433 - BK = 64

434 - k_tiles = math_utils.cdiv(K, BK)

435 - grid_m = math_utils.cdiv(M, BM)

544 + grid_m = math_utils.cdiv(M, 128)

545 grid_n = math_utils.cdiv(N, BN)

437 -

546 base_ctas = grid_m * grid_n

547 split_k = min(k_tiles, max(1, num_sms // base_ctas))

548 + return BN, split_k

549 +

550 +

551 +def warmup_bf16x3_router_gemm(

552 + K: int,

553 + M: int,

554 + min_num_tokens: int,

555 + max_num_tokens: int,

556 +) -> tuple[int, ...]:

557 + """Compile every BF16x3 router GEMM configuration reachable at runtime."""

558 + from vllm.model_executor.warmup.jit_warmup_triton_helper import (

559 + TritonWarmupTensor,

560 + triton_scalar_specialization_rep,

561 + )

562 +

563 + device = torch.accelerator.current_device_index()

564 + num_sms = torch.cuda.get_device_properties(device).multi_processor_count

565 + gemm_configs: set[int] = set()

566 + reduce_configs: set[tuple[int, int, int, int, int, int]] = set()

567 + for N in range(min_num_tokens, max_num_tokens + 1):

568 + BN, split_k = _pick_tile_config(N, K, M, num_sms)

569 + gemm_configs.add(BN)

570 + if split_k == 1:

571 + continue

572 + reduce_BN, reduce_BM, block_s = _splitk_reduce_config(split_k)

573 + reduce_configs.add(

574 + (

575 + triton_scalar_specialization_rep(N),

576 + triton_scalar_specialization_rep(N * M),

577 + triton_scalar_specialization_rep(split_k),

578 + reduce_BN,

579 + reduce_BM,

580 + block_s,

581 + )

582 + )

583 +

584 + configs = tuple(sorted(gemm_configs))

585 + for BN in configs:

586 + Sm100BF16x3RouterGemm.compile(K, BN)

587 + for N, split_stride, split_k, BN, BM, block_s in sorted(reduce_configs):

588 + partials = TritonWarmupTensor(torch.float32)

589 + out = TritonWarmupTensor(torch.float32)

590 + _splitk_reduce_kernel.warmup(

591 + partials,

592 + out,

593 + N,

594 + M,

595 + split_stride,

596 + split_k,

597 + BN=BN,

598 + BM=BM,

599 + BS=block_s,

600 + USE_PDL=True,

601 + num_warps=4,

602 + launch_pdl=True,

603 + grid=(1, 1),

604 + )

605 + return configs

606 +

607 +

608 +def _bf16x3_router_gemm(X: torch.Tensor, W: torch.Tensor) -> torch.Tensor:

609 + """Return ``X @ W.T`` using the SM100 BF16x3 router GEMM kernel."""

610 + N, K = X.shape

611 + M, _ = W.shape

612 + num_sms = torch.cuda.get_device_properties(X.device).multi_processor_count

613 +

614 + BN, split_k = _pick_tile_config(N, K, M, num_sms)

615

616 partials = X.new_empty(split_k, N, M, dtype=torch.float32)

442 - Sm100BF16x3RouterGemm.compile(BN, K)(X, W, partials, split_k)

617 + Sm100BF16x3RouterGemm.compile(K, BN)(X, W, partials, split_k)

618

619 if split_k == 1:

620 return partials.squeeze(0)

@@ -447,3 +622,21 @@ def bf16x3_router_gemm(X: torch.Tensor, W: torch.Tensor) -> torch.Tensor:

622 out = X.new_empty(N, M, dtype=torch.float32)

623 splitk_reduce_triton(partials, out)

624 return out

625 +

626 +

627 +def _bf16x3_router_gemm_fake(

628 + X: torch.Tensor,

629 + W: torch.Tensor,

630 +) -> torch.Tensor:

631 + return X.new_empty((X.shape[0], W.shape[0]), dtype=torch.float32)

632 +

633 +

634 +direct_register_custom_op(

635 + op_name="bf16x3_router_gemm",

636 + op_func=_bf16x3_router_gemm,

637 + fake_impl=_bf16x3_router_gemm_fake,

638 +)

639 +

640 +

641 +def bf16x3_router_gemm(X: torch.Tensor, W: torch.Tensor) -> torch.Tensor:

642 + return torch.ops.vllm.bf16x3_router_gemm(X, W)

vllm/model_executor/layers/fused_moe/router/gate_linear.py

+4 -16

@@ -4,15 +4,11 @@ import torch

4 from torch.nn.parameter import Parameter

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

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

14 -logger = init_logger(__name__)

15 -

13 @PluggableLayer.register("gate_linear")

14 class GateLinear(ReplicatedLinear):

@@ -22,7 +18,7 @@ class GateLinear(ReplicatedLinear):

18 K divisible by 8)

19 2. fp32 specialized kernel (SM90+ or gfx950, bf16/fp32 in, fp32 out,

20 M<=32, model-specific shapes)

25 - 3. experimental bf16x3 CuteDSL kernel (opt-in, SM100, bf16 in, fp32 weight)

21 + 3. bf16x3 CuteDSL kernel (SM100, bf16 in, fp32 weight)

22 4. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype)

23 5. F.linear via ReplicatedLinear (ultimate fallback)

@@ -85,11 +81,6 @@ class GateLinear(ReplicatedLinear):

81 self.allow_specialized_router_gemm = can_use_specialized_kernels

83 # fp32 specialized kernel eligibility (exact dims, fp32 weight)

88 - vllm_config = get_current_vllm_config_or_none()

89 - enable_bf16x3_router_gemm = (

90 - vllm_config is not None

91 - and vllm_config.kernel_config.enable_bf16x3_router_gemm

92 - )

84 self.allow_fp32_router_gemm = (

85 not bias

86 and self.weight.dtype == torch.float32

@@ -108,10 +99,7 @@ class GateLinear(ReplicatedLinear):

99 and current_platform.is_cuda()

100 and is_blackwell

101 and input_size % 8 == 0

111 - and enable_bf16x3_router_gemm

102 )

113 - if self.allow_bf16x3_router_gemm:

114 - logger.info_once("Enabled experimental SM100 BF16x3 router GEMM.")

103

104 # Fused bf16 x bf16 -> fp32 GEMM eligibility. torch.mm's out_dtype

105 # epilogue folds the fp32 cast into the GEMM, removing the standalone

@@ -198,7 +186,7 @@ class GateLinear(ReplicatedLinear):

186 )

187 return output, None

188

201 - # Tier 3: experimental bf16x3 CuteDSL kernel for fp32 router weights

189 + # Tier 3: bf16x3 CuteDSL kernel for fp32 router weights

190 if self.allow_bf16x3_router_gemm and x.dtype == torch.bfloat16:

191 from vllm.model_executor.layers.fused_moe.router.bf16x3_router_gemm_cutedsl import ( # noqa: E501

192 bf16x3_router_gemm,

@@ -251,10 +239,10 @@ def fp32_router_gemm_dispatch_impl(

239

240 if allow_bf16x3_router_gemm and x.dtype == torch.bfloat16:

241 from vllm.model_executor.layers.fused_moe.router.bf16x3_router_gemm_cutedsl import ( # noqa: E501

254 - bf16x3_router_gemm,

242 + _bf16x3_router_gemm,

243 )

244

257 - return bf16x3_router_gemm(x, weight)

245 + return _bf16x3_router_gemm(x, weight)

246

247 return torch.nn.functional.linear(x.float(), weight)

248

vllm/model_executor/warmup/kernel_warmup.py

+47 -2

@@ -101,6 +101,44 @@ def _warmup_ll_bf16_router_gemm(model: torch.nn.Module) -> None:

101 )

102

103

104 +def _warmup_bf16x3_router_gemm(

105 + model: torch.nn.Module,

106 + max_num_tokens: int,

107 +) -> None:

108 + from vllm.model_executor.layers.fused_moe.router.bf16x3_router_gemm_cutedsl import ( # noqa: E501

109 + warmup_bf16x3_router_gemm,

110 + )

111 + from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear

112 +

113 + gate = next(

114 + (

115 + module

116 + for module in model.modules()

117 + if isinstance(module, GateLinear) and module.allow_bf16x3_router_gemm

118 + ),

119 + None,

120 + )

121 + if gate is None:

122 + logger.debug_once(

123 + "Skipping BF16x3 router GEMM warmup: no eligible GateLinear found."

124 + )

125 + return

126 +

127 + min_num_tokens = gate.FP32_MAX_TOKENS + 1 if gate.allow_fp32_router_gemm else 1

128 + logger.info_once(

129 + "Warming up BF16x3 router GEMM for K=%d, M=%d.",

130 + gate.input_size,

131 + gate.output_size,

132 + )

133 + configs = warmup_bf16x3_router_gemm(

134 + gate.input_size,

135 + gate.output_size,

136 + min_num_tokens,

137 + max_num_tokens,

138 + )

139 + logger.info_once("Warmed up BF16x3 router GEMM configs: %s.", configs)

140 +

141 +

142 def _warmup_kimi_k3_gemm_rs_ar() -> None:

143 # Kimi-K3 model construction imports this module only when GEMM-RS/AR is

144 # enabled and initializes its singleton before kernel_warmup runs. Avoid

@@ -131,7 +169,8 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):

169 if zeroer is not None:

170 zeroer.warmup(worker.model_runner.kv_cache_config.num_blocks)

171

134 - if worker.vllm_config.kernel_config.enable_jit_warmup:

172 + enable_jit_warmup = worker.vllm_config.kernel_config.enable_jit_warmup

173 + if enable_jit_warmup:

174 logger.info("JIT kernel warmup starting.")

175 jit_warmup_start = time.perf_counter()

176 try:

@@ -164,11 +203,17 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):

203 )

204

205 # Run next so input-prep kernels JIT against pristine runner state.

167 - if worker.vllm_config.kernel_config.enable_jit_warmup:

206 + if enable_jit_warmup:

207 kimi_k3_triton_warmup(worker)

208 spec_decode_rejection_warmup(worker)

209 qwen4_exp_qsa_triton_warmup(worker)

210

211 + if enable_jit_warmup and current_platform.is_device_capability_family(100):

212 + _warmup_bf16x3_router_gemm(

213 + worker.get_model(),

214 + worker.scheduler_config.max_num_batched_tokens,

215 + )

216 +

217 if current_platform.has_device_capability(90):

218 _warmup_ll_bf16_router_gemm(worker.get_model())

219