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

[Kernel] TD operand loads for batched MoE GEMM (moe_mmk) on XPU (#46340)

Signed-off-by: oonyshch <xonyshch@gmail.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>

Lena Onyshchenko committed Jul 26, 2026 at 02:50 UTC 0111002323e226844a18e7f3e5834e9d0e1b6dcd

9 files changed +331 -33

tests/kernels/moe/test_batched_moe.py

+184

@@ -351,3 +351,187 @@ def test_fused_moe_batched_experts(

351 torch.testing.assert_close(batched_output, baseline_output, atol=3e-2, rtol=2e-2)

352

353 torch.testing.assert_close(triton_output, batched_output, atol=2e-2, rtol=2e-2)

354 +

355 +

356 +# USE_TD in moe_mmk

357 +

358 +# K % 8 == 0 (bf16, 16-byte alignment)

359 +_TD_SHAPES = [

360 + (4, 1, 128, 256),

361 + (8, 64, 512, 512),

362 + (4, 256, 1024, 2048),

363 +]

364 +

365 +

366 +def _td_supported() -> bool:

367 + return current_platform.is_xpu() or (

368 + current_platform.is_cuda() and current_platform.has_device_capability(90)

369 + )

370 +

371 +

372 +def _run_td(A, B, num_expert_tokens, use_td: bool):

373 + out = torch.zeros(

374 + A.shape[0],

375 + A.shape[1],

376 + B.shape[1],

377 + dtype=torch.bfloat16,

378 + device=A.device,

379 + )

380 + import triton

381 +

382 + from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import (

383 + batched_triton_kernel,

384 + )

385 +

386 + if use_td:

387 + from vllm.triton_utils.allocation import set_triton_allocator

388 +

389 + set_triton_allocator(A.device)

390 +

391 + E = A.shape[0]

392 + max_tokens = A.shape[1]

393 + K = A.shape[2]

394 + N = B.shape[1]

395 + BM = BN = BK = 64

396 + grid = (E, triton.cdiv(max_tokens, BM) * triton.cdiv(N, BN))

397 + batched_triton_kernel[grid](

398 + A,

399 + B,

400 + out,

401 + num_expert_tokens,

402 + tl.bfloat16,

403 + max_tokens,

404 + K,

405 + N,

406 + None,

407 + None,

408 + None,

409 + A.stride(0),

410 + A.stride(1),

411 + A.stride(2),

412 + B.stride(0),

413 + B.stride(2),

414 + B.stride(1),

415 + out.stride(0),

416 + out.stride(1),

417 + out.stride(2),

418 + 0,

419 + 0,

420 + 0,

421 + 0,

422 + 0,

423 + 0,

424 + 0,

425 + 0,

426 + False,

427 + False,

428 + False,

429 + BLOCK_M=BM,

430 + BLOCK_N=BN,

431 + BLOCK_K=BK,

432 + USE_TD=use_td,

433 + )

434 + return out

435 +

436 +

437 +@pytest.mark.parametrize("num_experts,max_tokens_per_expert,K,N", _TD_SHAPES)

438 +def test_batched_mm_td_matches_plain(num_experts, max_tokens_per_expert, K, N):

439 + if not _td_supported():

440 + pytest.skip("TD requires XPU or CUDA sm_90+")

441 + set_random_seed(42)

442 + device = current_platform.device_type

443 + A = (

444 + torch.randn(

445 + num_experts, max_tokens_per_expert, K, device=device, dtype=torch.bfloat16

446 + )

447 + / 10

448 + )

449 + B = torch.randn(num_experts, N, K, device=device, dtype=torch.bfloat16)

450 + num_expert_tokens = torch.randint(

451 + 1,

452 + max_tokens_per_expert + 1,

453 + size=(num_experts,),

454 + device=device,

455 + dtype=torch.int32,

456 + )

457 +

458 + out_plain = _run_td(A, B, num_expert_tokens, False)

459 + out_td = _run_td(A, B, num_expert_tokens, True)

460 +

461 + torch.testing.assert_close(out_td, out_plain, atol=6e-2, rtol=6e-2)

462 +

463 +

464 +def test_batched_mm_td_zero_expert_tokens():

465 + if not _td_supported():

466 + pytest.skip("TD requires XPU or CUDA sm_90+")

467 + set_random_seed(42)

468 + device = current_platform.device_type

469 + E, M, K, N = 8, 32, 256, 256

470 + A = torch.randn(E, M, K, device=device, dtype=torch.bfloat16) / 10

471 + B = torch.randn(E, N, K, device=device, dtype=torch.bfloat16)

472 + num_expert_tokens = torch.zeros(E, device=device, dtype=torch.int32)

473 + num_expert_tokens[::2] = M

474 +

475 + out_plain = _run_td(A, B, num_expert_tokens, False)

476 + out_td = _run_td(A, B, num_expert_tokens, True)

477 +

478 + for e in range(E):

479 + if num_expert_tokens[e].item() == 0:

480 + assert out_plain[e].abs().max().item() == 0.0

481 + assert out_td[e].abs().max().item() == 0.0

482 +

483 + torch.testing.assert_close(out_td, out_plain, atol=6e-2, rtol=6e-2)

484 +

485 +

486 +# BatchedTritonExperts device enablement (XPU)

487 +def test_batched_triton_experts_supports_current_device():

488 + from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import (

489 + BatchedTritonExperts,

490 + )

491 +

492 + if current_platform.is_xpu() or current_platform.is_cuda_alike():

493 + assert BatchedTritonExperts._supports_current_device()

494 + else:

495 + pytest.skip("No GPU device available")

496 +

497 +

498 +@pytest.mark.parametrize("m,n,k,e,topk", [(32, 512, 512, 8, 2), (45, 1024, 128, 8, 1)])

499 +def test_batched_experts_end_to_end(m, n, k, e, topk):

500 + """End-to-end BatchedTritonExperts via the reference (no-comms)

501 + BatchedPrepareAndFinalize, validated against a torch reference. Exercises

502 + the device-enablement path."""

503 + if not (current_platform.is_xpu() or current_platform.is_cuda_alike()):

504 + pytest.skip("No GPU device available")

505 +

506 + from vllm.v1.worker.workspace import init_workspace_manager

507 +

508 + set_random_seed(7)

509 + device = current_platform.device_type

510 + init_workspace_manager(torch.device(f"{device}:0"))

511 +

512 + a = torch.randn((m, k), device=device, dtype=torch.bfloat16) / 10

513 + score = torch.randn((m, e), device=device, dtype=torch.bfloat16)

514 + # w1 is the fused gate+up projection [E, 2N, K]; w2 is [E, K, N].

515 + w1 = torch.randn((e, 2 * n, k), device=device, dtype=torch.bfloat16) / 15

516 + w2 = torch.randn((e, k, n), device=device, dtype=torch.bfloat16) / 15

517 +

518 + with set_current_vllm_config(vllm_config):

519 + topk_weight, topk_ids, _ = fused_topk(a, score, topk, False)

520 +

521 + baseline_output = torch_experts(a, w1, w2, topk_weight, topk_ids)

522 + triton_output = batched_moe(a, w1, w2, topk_weight, topk_ids)

523 +

524 + torch.testing.assert_close(triton_output, baseline_output, atol=3e-2, rtol=2e-2)

525 +

526 +

527 +def test_batched_triton_backend_mapping():

528 + from vllm.model_executor.layers.fused_moe.oracle.unquantized import (

529 + UnquantizedMoeBackend,

530 + map_unquantized_backend,

531 + )

532 +

533 + assert (

534 + map_unquantized_backend("batched_triton")

535 + == UnquantizedMoeBackend.BATCHED_TRITON

536 + )

537 + assert map_unquantized_backend("triton") == UnquantizedMoeBackend.TRITON

vllm/config/kernel.py

@@ -122,6 +122,7 @@ class IrOpPriorityConfig:

122 MoEBackend = Literal[

123 "auto",

124 "triton",

125 + "batched_triton",

126 "deep_gemm",

127 "deep_gemm_mega_moe",

128 "cutlass",

@@ -192,6 +193,8 @@ class KernelConfig:

193

194 - "auto": Automatically select the best backend based on model and hardware

195 - "triton": Use Triton-based fused MoE kernels

196 + - "batched_triton": Use batched Triton experts (moe_mmk) on the batched

197 + activation format ([E_local, max_num_tokens, K])

198 - "deep_gemm": Use DeepGEMM kernels (FP8 block-quantized only)

199 - "deep_gemm_mega_moe": Use DeepGEMM mega MoE kernels

200 - "cutlass": Use vLLM CUTLASS kernels

vllm/model_executor/layers/fused_moe/all2all_utils.py

+11

@@ -19,6 +19,7 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import (

19 FusedMoEPrepareAndFinalize,

20 )

21 from vllm.model_executor.layers.fused_moe.prepare_finalize import (

22 + BatchedPrepareAndFinalize,

23 make_moe_prepare_and_finalize_naive_dp_ep,

24 make_moe_prepare_and_finalize_no_dp_ep,

25 )

@@ -140,6 +141,16 @@ def maybe_make_prepare_finalize(

141 if not allow_new_interface:

142 return None

143

144 + # Opt-in XPU batched path: reorganize tokens into E x T x K locally

145 + # (no all-to-all) so BatchedTritonExperts (moe_mmk TD) can run.

146 + if current_platform.is_xpu() and moe.moe_backend == "batched_triton":

147 + return BatchedPrepareAndFinalize(

148 + max_num_tokens=moe.max_num_tokens,

149 + num_local_experts=moe.num_local_experts,

150 + num_dispatchers=1,

151 + rank=moe.moe_parallel_config.ep_rank,

152 + )

153 +

154 # For DP/TP case, fall back to naive P/F.

155 if moe.moe_parallel_config.dp_size > 1:

156 logger.info_once(

vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py

+69 -13

@@ -32,7 +32,15 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

32 kFp8StaticTensorSym,

33 )

34 from vllm.platforms import current_platform

35 -from vllm.triton_utils import tl, triton

35 +from vllm.triton_utils import tl, triton, use_tensor_descriptor

36 +from vllm.triton_utils.allocation import set_triton_allocator

37 +

38 +

39 +def _is_capturing_or_compiling() -> bool:

40 + # torch.cuda.is_current_stream_capturing() is unavailable on non-CUDA (XPU) torch.

41 + return torch.compiler.is_compiling() or (

42 + current_platform.is_cuda_alike() and torch.cuda.is_current_stream_capturing()

43 + )

46 @triton.jit

@@ -71,9 +79,32 @@ def moe_mmk(

79 use_w8a8: tl.constexpr,

80 use_w8a16: tl.constexpr,

81 per_act_token_quant: tl.constexpr,

82 + # TD: a_base_ptr/b_base_ptr are the expert/CTA-offset bases of A[M,K]/B[N,K].

83 + a_base_ptr=None,

84 + b_base_ptr=None,

85 + M=0,

86 + N=0,

87 + stride_am: tl.int64 = 0,

88 + stride_bn: tl.int64 = 0,

89 + USE_TD: tl.constexpr = False,

90 ):

91 offs_k = tl.arange(0, BLOCK_K)

93 + if USE_TD:

94 + # make_tensor_descriptor requires the last (K) stride to be a

95 + # compile-time 1; the launcher only enables USE_TD for K-contiguous A/B.

96 + a_desc = tl.make_tensor_descriptor(

97 + a_base_ptr,

98 + shape=[M, K],

99 + strides=[stride_am, 1],

100 + block_shape=[BLOCK_M, BLOCK_K],

101 + )

102 + b_desc = tl.make_tensor_descriptor(

103 + b_base_ptr,

104 + shape=[N, K],

105 + strides=[stride_bn, 1],

106 + block_shape=[BLOCK_N, BLOCK_K],

107 + )

108 if use_w8a16:

109 b_scale_ptrs = (

110 b_scale_ptr + expert_id * stride_bse + offs_n[None, :] * stride_bsn

@@ -110,12 +141,17 @@ def moe_mmk(

141 for k in range(0, tl.cdiv(K, BLOCK_K)):

142 # Load the next block of A and B, generate a mask by checking the

143 # K dimension.

113 - a = tl.load(

114 - a_ptrs,

115 - mask=mask_m[:, None] & (offs_k[None, :] < K - k * BLOCK_K),

116 - other=0.0,

117 - )

118 - b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_K, other=0.0)

144 + if USE_TD:

145 + # B is [N, K]; tile is [BLOCK_N, BLOCK_K], transposed for dot.

146 + a = a_desc.load([0, k * BLOCK_K])

147 + b = tl.trans(b_desc.load([0, k * BLOCK_K]))

148 + else:

149 + a = tl.load(

150 + a_ptrs,

151 + mask=mask_m[:, None] & (offs_k[None, :] < K - k * BLOCK_K),

152 + other=0.0,

153 + )

154 + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_K, other=0.0)

155 # We accumulate along the K dimension.

156 if use_w8a16:

157 accumulator = tl.dot(a, b.to(compute_type), acc=accumulator)

@@ -193,6 +229,7 @@ def expert_triton_kernel(

229 BLOCK_M: tl.constexpr,

230 BLOCK_N: tl.constexpr,

231 BLOCK_K: tl.constexpr,

232 + USE_TD: tl.constexpr = False,

233 ):

234 offs_m = tl.arange(0, BLOCK_M)

235 offs_n = tl.arange(0, BLOCK_N) % N

@@ -238,6 +275,13 @@ def expert_triton_kernel(

275 use_fp8_w8a8,

276 use_int8_w8a16,

277 per_act_token_quant,

278 + a_ptr,

279 + b_ptr,

280 + M,

281 + N,

282 + stride_am,

283 + stride_bn,

284 + USE_TD,

285 )

286

287 # store in C

@@ -292,6 +336,7 @@ def batched_triton_kernel(

336 BLOCK_M: tl.constexpr,

337 BLOCK_N: tl.constexpr,

338 BLOCK_K: tl.constexpr,

339 + USE_TD: tl.constexpr = False,

340 ):

341 expert_id = tl.program_id(axis=0)

342 e_num_tokens = tl.load(expert_num_tokens + expert_id)

@@ -372,6 +417,7 @@ def batched_triton_kernel(

417 BLOCK_M,

418 BLOCK_N,

419 BLOCK_K,

420 + USE_TD,

421 )

422

423

@@ -444,6 +490,18 @@ def invoke_moe_batched_triton_kernel(

490 stride_asm = 0

491 stride_ask = 0

492

493 + use_td = (

494 + use_tensor_descriptor()

495 + and A.stride(2) == 1

496 + and B.stride(2) == 1

497 + and (K * A.element_size()) % 16 == 0

498 + and (BLOCK_M & (BLOCK_M - 1)) == 0

499 + and (BLOCK_N & (BLOCK_N - 1)) == 0

500 + and (BLOCK_K & (BLOCK_K - 1)) == 0

501 + )

502 + if use_td:

503 + set_triton_allocator(A.device)

504 +

505 batched_triton_kernel[grid](

506 A,

507 B,

@@ -485,6 +543,7 @@ def invoke_moe_batched_triton_kernel(

543 BLOCK_M=BLOCK_M,

544 BLOCK_N=BLOCK_N,

545 BLOCK_K=BLOCK_K,

546 + USE_TD=use_td,

547 )

548

549

@@ -616,10 +675,7 @@ class NaiveBatchedExperts(mk.FusedMoEExpertsModular):

675

676 for expert in range(num_local_experts):

677 # Indexing expert_num_tokens doesn't work w/cudagraphs or inductor

619 - if (

620 - torch.compiler.is_compiling()

621 - or torch.cuda.is_current_stream_capturing()

622 - ):

678 + if _is_capturing_or_compiling():

679 num = hidden_states.shape[1]

680 else:

681 num = int(expert_num_tokens[expert].item())

@@ -659,7 +715,7 @@ def batched_moe_kernel_quantize_input(

715 per_act_token_quant: bool,

716 block_shape: list[int] | None = None,

717 ) -> tuple[torch.Tensor, torch.Tensor | None]:

662 - if torch.compiler.is_compiling() or torch.cuda.is_current_stream_capturing():

718 + if _is_capturing_or_compiling():

719 # Note: this does a bunch of extra work because expert_num_tokens is

720 # ignored but it does support torch.compile + cudagraphs.

721 hidden_dim = A.size(-1)

@@ -745,7 +801,7 @@ class BatchedTritonExperts(mk.FusedMoEExpertsModular):

801

802 @staticmethod

803 def _supports_current_device() -> bool:

748 - return current_platform.is_cuda_alike()

804 + return current_platform.is_cuda_alike() or current_platform.is_xpu()

805

806 @staticmethod

807 def _supports_no_act_and_mul() -> bool:

vllm/model_executor/layers/fused_moe/oracle/unquantized.py

@@ -151,6 +151,7 @@ def map_unquantized_backend(runner_backend: MoEBackend) -> UnquantizedMoeBackend

151 """Map user's MoEBackend to UnquantizedMoeBackend."""

152 mapping = {

153 "triton": UnquantizedMoeBackend.TRITON,

154 + "batched_triton": UnquantizedMoeBackend.BATCHED_TRITON,

155 "flashinfer_trtllm": UnquantizedMoeBackend.FLASHINFER_TRTLLM,

156 "flashinfer_cutlass": UnquantizedMoeBackend.FLASHINFER_CUTLASS,

157 "aiter": UnquantizedMoeBackend.AITER,

@@ -233,6 +234,7 @@ def select_unquantized_moe_backend(

234 activation_format = (

235 mk.FusedMoEActivationFormat.BatchedExperts

236 if moe_config.moe_parallel_config.use_batched_activation_format

237 + or moe_config.moe_backend == "batched_triton"

238 else mk.FusedMoEActivationFormat.Standard

239 )

240

vllm/model_executor/layers/fused_moe/prepare_finalize/batched.py

+22 -11

@@ -112,15 +112,28 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):

112

113 a1_scale = normalize_scales_shape(quant_config.a1_scale)

114

115 - for expert_id in range(first_expert, last_expert):

116 - topks = torch.any(topk_ids == expert_id, dim=1).flatten()

117 - rows = torch.count_nonzero(topks.flatten())

118 - if rows == 0:

119 - continue

120 - idx = expert_id - first_expert

121 - tokens_per_expert[idx] = rows

122 - rhs = a1[: topks.numel()][topks]

123 - if quant_config.quant_dtype is not None:

115 + if quant_config.quant_dtype is None:

116 + # Vectorized dispatch: compute the [E_local, T] hit mask and

117 + # per-expert slot offsets in one shot (single nonzero sync instead

118 + # of a launch-bound per-expert Python loop).

119 + local_ids = torch.arange(first_expert, last_expert, device=a1.device).view(

120 + -1, 1, 1

121 + )

122 + hits = (topk_ids.unsqueeze(0) == local_ids).any(dim=2) # [E_local, T]

123 + tokens_per_expert[:num_local_experts] = hits.sum(dim=1).to(torch.int32)

124 + # slot of each token within its expert batch, preserving token order

125 + slots = hits.to(torch.int32).cumsum(dim=1) - 1 # [E_local, T]

126 + e_idx, t_idx = hits.nonzero(as_tuple=True)

127 + b_a1[e_idx, slots[e_idx, t_idx]] = a1[t_idx].to(b_type)

128 + else:

129 + for expert_id in range(first_expert, last_expert):

130 + topks = torch.any(topk_ids == expert_id, dim=1).flatten()

131 + rows = torch.count_nonzero(topks.flatten())

132 + if rows == 0:

133 + continue

134 + idx = expert_id - first_expert

135 + tokens_per_expert[idx] = rows

136 + rhs = a1[: topks.numel()][topks]

137 if a1_scale is not None:

138 if quant_config.is_per_act_token:

139 rhs_a1_scale = a1_scale[: topks.numel()][topks]

@@ -140,8 +153,6 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):

153 b_a1_scale[idx, :rows] = b_s[:rows]

154 else:

155 b_a1_scale[idx, : b_s.shape[0]] = b_s

143 - else:

144 - b_a1[idx, :rows, :] = rhs

156

157 assert b_a1_scale is None or b_a1_scale.ndim == 3

158

vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py

+15 -8

@@ -164,13 +164,20 @@ class TopKWeightAndReduceNaiveBatched(mk.TopKWeightAndReduce):

164 first_expert = num_local_experts * self.rank

165 last_expert = first_expert + num_local_experts

166

167 - for expert_id in range(first_expert, last_expert):

168 - matching_tokens = topk_ids == expert_id

169 - topks = torch.any(matching_tokens, dim=1).flatten()

170 - rows = torch.count_nonzero(topks)

171 - rhs = fused_expert_output[expert_id - first_expert, :rows, :]

172 - if not apply_router_weight_on_input:

173 - rhs.mul_(topk_weights[matching_tokens].view(rhs.size(0), 1))

174 - output[topks] = output[topks] + rhs

167 + # Vectorized weighted scatter-add (single nonzero sync instead of a

168 + # per-expert loop; mirrors the dispatch path).

169 + local_ids = torch.arange(

170 + first_expert, last_expert, device=topk_ids.device

171 + ).view(-1, 1, 1)

172 + matching = topk_ids.unsqueeze(0) == local_ids # [E_local, T, topk]

173 + hits = matching.any(dim=2) # [E_local, T]

174 + slots = hits.to(torch.int32).cumsum(dim=1) - 1 # [E_local, T]

175 + e_idx, t_idx = hits.nonzero(as_tuple=True)

176 + gathered = fused_expert_output[e_idx, slots[e_idx, t_idx], :]

177 + if not apply_router_weight_on_input:

178 + # weight of token t at expert e: the topk weight on the matching slot

179 + weights = (matching * topk_weights.unsqueeze(0)).sum(dim=2) # [E_local, T]

180 + gathered = gathered * weights[e_idx, t_idx].unsqueeze(1).to(gathered.dtype)

181 + output.index_add_(0, t_idx, gathered.to(output.dtype))

182

183 return output

vllm/triton_utils/__init__.py

+11 -1

@@ -17,7 +17,17 @@ else:

17 tl = TritonLanguagePlaceholder()

18 tldevice = TritonLanguagePlaceholder()

20 +from vllm.triton_utils.tensor_descriptor import use_tensor_descriptor

21 +

22 LOG2E = 1.4426950408889634

23 LOGE2 = 0.6931471805599453

23 -__all__ = ["HAS_TRITON", "triton", "tl", "tldevice", "LOG2E", "LOGE2"]

25 +__all__ = [

26 + "HAS_TRITON",

27 + "triton",

28 + "tl",

29 + "tldevice",

30 + "LOG2E",

31 + "LOGE2",

32 + "use_tensor_descriptor",

33 +]

vllm/triton_utils/tensor_descriptor.py new

+14

@@ -0,0 +1,14 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +

5 +def use_tensor_descriptor(override: bool | None = None) -> bool:

6 + """Tri-state VLLM_TRITON_USE_TD: unset=auto (on for XPU), 1/0=force on/off."""

7 + from vllm import envs

8 + from vllm.platforms import current_platform

9 +

10 + if override is None:

11 + override = envs.VLLM_TRITON_USE_TD

12 + if override is not None:

13 + return override

14 + return current_platform.is_xpu()