A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 12:45 UTC
feat(moe): Add is_act_and_mul=False support for Triton MoE kernels (#31645)
Signed-off-by: rabi <ramishra@redhat.com>
Rabi Mishra committed
Jan 8, 2026 at 07:57 UTC
25eef3dc2ec82ca6ae4a029b247791b72a75ce15
7 files changed
+191
-9
tests/kernels/moe/test_triton_moe_no_act_mul.py
new
+129
@@ -0,0 +1,129 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Test for is_act_and_mul=False MoE using Triton.
5
+
6
+This tests the code path used by models like Nemotron-H that use
7
+non-fused activations (e.g., relu2_no_mul) instead of SwiGLU-style
8
+fused activations.
9
+
10
+This feature is supported on both CUDA and ROCm (with AITER disabled).
11
+"""
12
+
13
+import pytest
14
+import torch
15
+
16
+from vllm.platforms import current_platform
17
+
18
+pytestmark = pytest.mark.skipif(
19
+ not current_platform.is_cuda_alike(),
20
+ reason="Tests for is_act_and_mul=False MoE require CUDA or ROCm",
21
+)
22
+
23
+
24
+@pytest.fixture
25
+def disable_aiter_on_rocm(monkeypatch):
26
+ """Fixture to disable AITER on ROCm to use Triton path."""
27
+ if current_platform.is_rocm():
28
+ from vllm._aiter_ops import rocm_aiter_ops
29
+
30
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER", "0")
31
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "0")
32
+ rocm_aiter_ops.refresh_env_variables()
33
+
34
+ yield
35
+
36
+ rocm_aiter_ops.refresh_env_variables()
37
+ else:
38
+ # On CUDA, no special setup needed
39
+ yield
40
+
41
+
42
+@pytest.fixture
43
+def init_workspace():
44
+ """Initialize workspace manager for MoE tests."""
45
+ from vllm.v1.worker.workspace import (
46
+ init_workspace_manager,
47
+ reset_workspace_manager,
48
+ )
49
+
50
+ torch.manual_seed(42)
51
+ init_workspace_manager(torch.cuda.current_device())
52
+
53
+ yield
54
+
55
+ reset_workspace_manager()
56
+
57
+
58
+@pytest.mark.parametrize("m", [1, 33, 64, 222])
59
+@pytest.mark.parametrize("n", [128, 256, 1024])
60
+@pytest.mark.parametrize("k", [128, 512])
61
+@pytest.mark.parametrize("e", [4, 8])
62
+@pytest.mark.parametrize("topk", [2])
63
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
64
+@pytest.mark.parametrize("activation", ["relu2_no_mul", "silu_no_mul", "gelu_no_mul"])
65
+@torch.inference_mode()
66
+def test_moe_no_act_mul(
67
+ disable_aiter_on_rocm,
68
+ init_workspace,
69
+ m: int,
70
+ n: int,
71
+ k: int,
72
+ e: int,
73
+ topk: int,
74
+ dtype: torch.dtype,
75
+ activation: str,
76
+):
77
+ """Test MoE with is_act_and_mul=False using Triton."""
78
+ from vllm.model_executor.layers.fused_moe import TritonExperts, fused_topk
79
+ from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
80
+ from vllm.model_executor.layers.fused_moe.modular_kernel import (
81
+ FusedMoEModularKernel,
82
+ )
83
+ from vllm.model_executor.layers.fused_moe.prepare_finalize import (
84
+ MoEPrepareAndFinalizeNoEP,
85
+ )
86
+
87
+ a = torch.randn((m, k), device="cuda", dtype=dtype)
88
+ w1 = torch.randn((e, n, k), device="cuda", dtype=dtype) / 10
89
+ w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
90
+
91
+ quant_config = FusedMoEQuantConfig.make(is_act_and_mul=False)
92
+
93
+ score = torch.randn((m, e), device="cuda", dtype=dtype)
94
+ topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=True)
95
+
96
+ fused_experts = FusedMoEModularKernel(
97
+ MoEPrepareAndFinalizeNoEP(),
98
+ TritonExperts(quant_config),
99
+ )
100
+
101
+ output = fused_experts(
102
+ hidden_states=a,
103
+ w1=w1,
104
+ w2=w2,
105
+ topk_weights=topk_weights,
106
+ topk_ids=topk_ids,
107
+ activation=activation,
108
+ )
109
+
110
+ assert output.shape == (m, k), f"Expected shape {(m, k)}, got {output.shape}"
111
+ assert not torch.isnan(output).any(), "Output contains NaN"
112
+ assert not torch.isinf(output).any(), "Output contains Inf"
113
+ assert output.abs().sum() > 0, "Output is all zeros"
114
+
115
+
116
+@torch.inference_mode()
117
+def test_moe_workspace_shapes_no_act_mul(disable_aiter_on_rocm):
118
+ """Test workspace_shapes returns correct sizes for is_act_and_mul=False."""
119
+ from vllm.model_executor.layers.fused_moe import TritonExperts
120
+ from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
121
+
122
+ M, N, K, topk = 64, 256, 128, 2
123
+
124
+ quant_config = FusedMoEQuantConfig.make(is_act_and_mul=False)
125
+ experts = TritonExperts(quant_config)
126
+ ws1, ws2, out = experts.workspace_shapes(M, N, K, topk, 8, 8, None)
127
+
128
+ assert ws1[2] == max(N, K)
129
+ assert out == (M, K)
vllm/model_executor/layers/fused_moe/config.py
+9
@@ -201,6 +201,11 @@ class FusedMoEQuantConfig:
201
_w1: FusedMoEQuantDesc
202
_w2: FusedMoEQuantDesc
203
204
+ # Whether activation is fused with gate multiplication (SwiGLU-style).
205
+ # When True: intermediate_size = N // 2 (gate and up are combined)
206
+ # When False: intermediate_size = N (no gate multiplication)
207
+ is_act_and_mul: bool = True
208
+
209
def __post_init__(self):
210
assert not self.per_act_token_quant or self.block_shape is None, (
211
"illegal quantization"
@@ -435,6 +440,7 @@ class FusedMoEQuantConfig:
440
w1_zp: torch.Tensor | None = None,
441
w2_zp: torch.Tensor | None = None,
442
weight_dtype: torch.dtype | str | None = None,
443
+ is_act_and_mul: bool = True,
444
) -> "FusedMoEQuantConfig":
445
"""
446
General builder function for a FusedMoEQuantConfig.
@@ -494,6 +500,7 @@ class FusedMoEQuantConfig:
500
_w2=FusedMoEQuantDesc(
501
weight_dtype, w_shape, w2_scale, g2_alphas, w2_zp, w2_bias
502
),
503
+ is_act_and_mul=is_act_and_mul,
504
)
505
assert quant_config.per_act_token_quant == per_act_token_quant
506
assert quant_config.per_out_ch_quant == per_out_ch_quant
@@ -806,6 +813,7 @@ def awq_marlin_moe_quant_config(
813
def biased_moe_quant_config(
814
w1_bias: torch.Tensor | None,
815
w2_bias: torch.Tensor | None,
816
+ is_act_and_mul: bool = True,
817
) -> FusedMoEQuantConfig:
818
"""
819
Construct a quant config for unquantized activations with biases.
@@ -815,6 +823,7 @@ def biased_moe_quant_config(
823
_a2=FusedMoEQuantDesc(),
824
_w1=FusedMoEQuantDesc(bias=w1_bias),
825
_w2=FusedMoEQuantDesc(bias=w2_bias),
826
+ is_act_and_mul=is_act_and_mul,
827
)
828
829
vllm/model_executor/layers/fused_moe/fused_batched_moe.py
+10
-3
@@ -871,8 +871,11 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
871
num_dp = self.num_dispatchers
872
num_experts = local_num_experts
873
max_num_tokens = self.max_num_tokens
874
+ # For fused activations (SwiGLU): N = 2 * intermediate, after act = N/2
875
+ # For non-fused activations: N = intermediate, after act = N
876
+ intermediate_size = N // 2 if self.quant_config.is_act_and_mul else N
877
workspace13 = (num_experts, max_num_tokens * num_dp, max(K, N))
875
- workspace2 = (num_experts, max_num_tokens * num_dp, (N // 2))
878
+ workspace2 = (num_experts, max_num_tokens * num_dp, intermediate_size)
879
output = (num_experts, max_num_tokens * num_dp, K)
880
return (workspace13, workspace2, output)
881
@@ -947,7 +950,11 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
950
# We can reuse the memory between these because by the time we need
951
# cache3, we're done with cache1
952
intermediate_cache1 = _resize_cache(workspace13, (E, max_num_tokens, N))
950
- intermediate_cache2 = _resize_cache(workspace2, (E, max_num_tokens, N // 2))
953
+ # For fused activations (SwiGLU): output is N/2, for non-fused: output is N
954
+ intermediate_size = N // 2 if self.quant_config.is_act_and_mul else N
955
+ intermediate_cache2 = _resize_cache(
956
+ workspace2, (E, max_num_tokens, intermediate_size)
957
+ )
958
959
# TODO(bnell): should this be done for any quantized type?
960
if self.quant_config.use_fp8_w8a8:
@@ -978,7 +985,7 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
985
# TODO (bnell): use triton utility from batched deep gemm.
986
self.activation(
987
activation,
981
- intermediate_cache2.view(-1, N // 2),
988
+ intermediate_cache2.view(-1, intermediate_size),
989
intermediate_cache1.view(-1, N),
990
)
991
vllm/model_executor/layers/fused_moe/fused_moe.py
+7
-2
@@ -2292,7 +2292,10 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
2292
local_num_experts: int,
2293
expert_tokens_meta: mk.ExpertTokensMetadata | None,
2294
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
2295
- workspace1 = (M, topk, max(N // 2, K))
2295
+ # For fused activations (SwiGLU): N = 2 * intermediate, after act = N/2
2296
+ # For non-fused activations: N = intermediate, after act = N
2297
+ intermediate_size = N // 2 if self.quant_config.is_act_and_mul else N
2298
+ workspace1 = (M, topk, max(intermediate_size, K))
2299
workspace2 = (M, topk, max(N, K))
2300
output = (M, K)
2301
return (workspace1, workspace2, output)
@@ -2367,8 +2370,10 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
2370
2371
# Note that the output tensor might be in workspace1
2372
intermediate_cache1 = _resize_cache(workspace2, (num_tokens, top_k_num, N))
2373
+ # For fused activations (SwiGLU): output is N/2, for non-fused: output is N
2374
+ intermediate_size = N // 2 if self.quant_config.is_act_and_mul else N
2375
intermediate_cache2 = _resize_cache(
2371
- workspace13, (num_tokens * top_k_num, N // 2)
2376
+ workspace13, (num_tokens * top_k_num, intermediate_size)
2377
)
2378
intermediate_cache3 = _resize_cache(workspace2, (num_tokens, top_k_num, K))
2379
vllm/model_executor/layers/fused_moe/layer.py
+8
-2
@@ -603,9 +603,15 @@ class FusedMoE(CustomOp):
603
"is_act_and_mul=False is supported only for unquantized "
604
", ModelOpt FP8, and ModelOpt NvFp4 checkpoints"
605
)
606
- if not current_platform.is_cuda():
606
+ # ROCm without AITER MoE uses Triton which supports
607
+ # is_act_and_mul=False via standard PyTorch ops (F.silu, F.gelu)
608
+ rocm_without_aiter_moe = (
609
+ current_platform.is_rocm() and not rocm_aiter_ops.is_fused_moe_enabled()
610
+ )
611
+ if not current_platform.is_cuda() and not rocm_without_aiter_moe:
612
raise NotImplementedError(
608
- "is_act_and_mul=False is supported only for CUDA for now"
613
+ "is_act_and_mul=False is supported only for CUDA, or ROCm "
614
+ "(when AITER MoE is disabled) for now"
615
)
616
617
if self.enable_eplb and not self.quant_method.supports_eplb:
vllm/model_executor/layers/fused_moe/modular_kernel.py
+23
-2
@@ -4,7 +4,7 @@ from abc import ABC, abstractmethod
4
from collections.abc import Callable
5
from dataclasses import dataclass
6
from enum import Enum
7
-from math import prod
7
+from math import prod, sqrt
8
from typing import final
9
10
import torch
@@ -575,14 +575,35 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
575
def activation(
576
self, activation: str, output: torch.Tensor, input: torch.Tensor
577
) -> None:
578
- assert output.size(-1) * 2 == input.size(-1)
578
+ # Fused activations (SwiGLU-style): output is half the size of input
579
if activation == "silu":
580
+ assert output.size(-1) * 2 == input.size(-1)
581
torch.ops._C.silu_and_mul(output, input)
582
elif activation == "gelu":
583
+ assert output.size(-1) * 2 == input.size(-1)
584
torch.ops._C.gelu_and_mul(output, input)
585
elif activation == "swigluoai":
586
# alpha = 1.702, limit = 7.0
587
+ assert output.size(-1) * 2 == input.size(-1)
588
torch.ops._C.swigluoai_and_mul(output, input)
589
+ # Non-fused activations (is_act_and_mul=False): output same size as input
590
+ elif activation == "silu_no_mul":
591
+ assert output.size(-1) == input.size(-1)
592
+ # Use out= argument to avoid intermediate tensor
593
+ torch.sigmoid(input, out=output)
594
+ output.mul_(input)
595
+ elif activation == "gelu_no_mul":
596
+ assert output.size(-1) == input.size(-1)
597
+ # GELU(x) = 0.5 * x * (1 + erf(x / sqrt(2)))
598
+ # Use out= and in-place ops to avoid intermediate tensors
599
+ output.copy_(input).div_(sqrt(2))
600
+ torch.erf(output, out=output)
601
+ output.add_(1).mul_(input).mul_(0.5)
602
+ elif activation == "relu2_no_mul":
603
+ assert output.size(-1) == input.size(-1)
604
+ # ReLU²: clamp has out=, then in-place square
605
+ torch.clamp(input, min=0, out=output)
606
+ output.square_()
607
else:
608
raise ValueError(f"Unsupported FusedMoe activation: {activation}")
609
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+5
@@ -299,7 +299,12 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
299
return biased_moe_quant_config(
300
layer.w13_bias,
301
layer.w2_bias,
302
+ is_act_and_mul=self.moe.is_act_and_mul,
303
)
304
+ elif not self.moe.is_act_and_mul:
305
+ # Create a config with is_act_and_mul=False since
306
+ # FUSED_MOE_UNQUANTIZED_CONFIG has is_act_and_mul=True
307
+ return FusedMoEQuantConfig.make(is_act_and_mul=False)
308
else:
309
return FUSED_MOE_UNQUANTIZED_CONFIG
310