[Kernel] Add swap AB optimization to fused_moe_kernel (#36559)
Signed-off-by: Xin Yang <xyangx@amazon.com>
Xin Yang committed
Jun 24, 2026 at 18:44 UTC
1273a8f05a1a3fa90c0edfdf1c6c87e70054e680
2 files changed
+59
-14
vllm/model_executor/layers/fused_moe/fused_moe.py
+49
-14
@@ -26,6 +26,7 @@ from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
26
moe_align_block_size,
27
)
28
from vllm.model_executor.layers.fused_moe.utils import (
29
+ enable_swap_ab,
30
moe_kernel_quantize_input,
31
)
32
from vllm.platforms import current_platform
@@ -343,6 +344,7 @@ def fused_moe_kernel(
344
use_int8_w8a16: tl.constexpr,
345
per_channel_quant: tl.constexpr,
346
HAS_BIAS: tl.constexpr,
347
+ SWAP_AB: tl.constexpr,
348
):
349
"""
350
Implements the fused computation for a Mixture of Experts (MOE) using
@@ -432,15 +434,25 @@ def fused_moe_kernel(
434
435
offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N
436
offs_k = tl.arange(0, BLOCK_SIZE_K)
435
- a_ptrs = a_ptr + (
436
- offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak
437
- )
437
+ if SWAP_AB:
438
+ a_ptrs = a_ptr + (
439
+ offs_k[:, None] * stride_ak + offs_token[None, :] // top_k * stride_am
440
+ )
441
+ b_ptrs = (
442
+ b_ptr
443
+ + off_experts * stride_be
444
+ + (offs_bn[:, None] * stride_bn + offs_k[None, :] * stride_bk)
445
+ )
446
+ else:
447
+ a_ptrs = a_ptr + (
448
+ offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak
449
+ )
450
+ b_ptrs = (
451
+ b_ptr
452
+ + off_experts * stride_be
453
+ + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
454
+ )
455
439
- b_ptrs = (
440
- b_ptr
441
- + off_experts * stride_be
442
- + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
443
- )
456
if use_int8_w8a16:
457
b_scale_ptrs = (
458
b_scale_ptr + off_experts * stride_bse + offs_bn[None, :] * stride_bsn
@@ -477,16 +489,25 @@ def fused_moe_kernel(
489
# We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block
490
# of fp32 values for higher accuracy.
491
# `accumulator` will be converted back to fp16 after the loop.
480
- accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
492
+ if SWAP_AB:
493
+ accumulator = tl.zeros((BLOCK_SIZE_N, BLOCK_SIZE_M), dtype=tl.float32)
494
+ else:
495
+ accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
496
for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
497
# Load the next block of A and B, generate a mask by checking the
498
# K dimension.
499
+ if SWAP_AB:
500
+ a_mask = (offs_k[:, None] < K - k * BLOCK_SIZE_K) & token_mask[None, :]
501
+ b_mask = offs_k[None, :] < K - k * BLOCK_SIZE_K
502
+ else:
503
+ a_mask = token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K)
504
+ b_mask = offs_k[:, None] < K - k * BLOCK_SIZE_K
505
a = tl.load(
506
a_ptrs,
486
- mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K),
507
+ mask=a_mask,
508
other=0.0,
509
)
489
- b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
510
+ b = tl.load(b_ptrs, mask=b_mask, other=0.0)
511
# We accumulate along the K dimension.
512
if use_int8_w8a16:
513
accumulator = tl.dot(a, b.to(compute_type), acc=accumulator)
@@ -498,12 +519,17 @@ def fused_moe_kernel(
519
a_scale_ptrs + offs_ks * stride_ask, mask=token_mask, other=0.0
520
)
521
b_scale = tl.load(b_scale_ptrs + offs_ks * stride_bsk)
501
-
502
- accumulator += tl.dot(a, b) * a_scale[:, None] * b_scale[None, :]
522
+ if SWAP_AB:
523
+ accumulator += tl.dot(b, a) * b_scale[:, None] * a_scale[None, :]
524
+ else:
525
+ accumulator += tl.dot(a, b) * a_scale[:, None] * b_scale[None, :]
526
else:
527
if use_fp8_w8a8:
528
# acc used to enable fp8_fast_accum
506
- accumulator = tl.dot(a, b, acc=accumulator)
529
+ if SWAP_AB:
530
+ accumulator = tl.dot(b, a, acc=accumulator)
531
+ else:
532
+ accumulator = tl.dot(a, b, acc=accumulator)
533
else:
534
accumulator += tl.dot(a, b)
535
else:
@@ -512,6 +538,9 @@ def fused_moe_kernel(
538
a_ptrs += BLOCK_SIZE_K * stride_ak
539
b_ptrs += BLOCK_SIZE_K * stride_bk
540
541
+ if SWAP_AB:
542
+ accumulator = tl.trans(accumulator, (1, 0))
543
+
544
# Dequantization for supported quantization schemes:
545
# - int8_w8a16
546
# - fp8_w8a8
@@ -729,6 +758,11 @@ def invoke_fused_moe_triton_kernel(
758
assert topk_weights is None or topk_weights.stride(1) == 1
759
assert sorted_token_ids is None or sorted_token_ids.stride(0) == 1
760
761
+ if use_fp8_w8a8:
762
+ SWAP_AB = enable_swap_ab(config["BLOCK_SIZE_M"], config["BLOCK_SIZE_N"])
763
+ else:
764
+ SWAP_AB = False
765
+
766
if use_fp8_w8a8 or use_int8_w8a8:
767
assert B_scale is not None
768
assert block_shape is None or triton.cdiv(
@@ -810,6 +844,7 @@ def invoke_fused_moe_triton_kernel(
844
naive_block_assignment=(sorted_token_ids is None),
845
HAS_BIAS=HAS_BIAS,
846
BLOCK_SIZE_K=BLOCK_SIZE_K,
847
+ SWAP_AB=SWAP_AB,
848
**config,
849
)
850
vllm/model_executor/layers/fused_moe/utils.py
+10
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import functools
4
from math import prod
5
6
import torch
@@ -446,3 +447,12 @@ def swiglu_limit_func(
447
up = torch.clamp(up, min=-swiglu_limit, max=swiglu_limit)
448
449
output.copy_(F.silu(gate) * up)
450
+
451
+
452
+@functools.lru_cache
453
+def enable_swap_ab(BLOCK_SIZE_M: int, BLOCK_SIZE_N: int) -> bool:
454
+ return (
455
+ current_platform.is_device_capability(90)
456
+ and BLOCK_SIZE_M < 64
457
+ and BLOCK_SIZE_N >= 64
458
+ )