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, 11:30 UTC
Revert "feat(moe): Add is_act_and_mul=False support for Triton MoE kernels" (#31978)
Michael Goin committed
Jan 8, 2026 at 14:31 UTC
87e07a6b46fbd85725cd19bfa14b1387e90003cf
7 files changed
+9
-191
tests/kernels/moe/test_triton_moe_no_act_mul.py
deleted
-129
@@ -1,129 +0,0 @@
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,11 +201,6 @@ 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
-
204
def __post_init__(self):
205
assert not self.per_act_token_quant or self.block_shape is None, (
206
"illegal quantization"
@@ -444,7 +439,6 @@ class FusedMoEQuantConfig:
439
w1_zp: torch.Tensor | None = None,
440
w2_zp: torch.Tensor | None = None,
441
weight_dtype: torch.dtype | str | None = None,
447
- is_act_and_mul: bool = True,
442
) -> "FusedMoEQuantConfig":
443
"""
444
General builder function for a FusedMoEQuantConfig.
@@ -504,7 +498,6 @@ class FusedMoEQuantConfig:
498
_w2=FusedMoEQuantDesc(
499
weight_dtype, w_shape, w2_scale, g2_alphas, w2_zp, w2_bias
500
),
507
- is_act_and_mul=is_act_and_mul,
501
)
502
assert quant_config.per_act_token_quant == per_act_token_quant
503
assert quant_config.per_out_ch_quant == per_out_ch_quant
@@ -836,7 +829,6 @@ def awq_marlin_moe_quant_config(
829
def biased_moe_quant_config(
830
w1_bias: torch.Tensor | None,
831
w2_bias: torch.Tensor | None,
839
- is_act_and_mul: bool = True,
832
) -> FusedMoEQuantConfig:
833
"""
834
Construct a quant config for unquantized activations with biases.
@@ -846,7 +838,6 @@ def biased_moe_quant_config(
838
_a2=FusedMoEQuantDesc(),
839
_w1=FusedMoEQuantDesc(bias=w1_bias),
840
_w2=FusedMoEQuantDesc(bias=w2_bias),
849
- is_act_and_mul=is_act_and_mul,
841
)
842
843
vllm/model_executor/layers/fused_moe/fused_batched_moe.py
+3
-10
@@ -871,11 +871,8 @@ 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
874
workspace13 = (num_experts, max_num_tokens * num_dp, max(K, N))
878
- workspace2 = (num_experts, max_num_tokens * num_dp, intermediate_size)
875
+ workspace2 = (num_experts, max_num_tokens * num_dp, (N // 2))
876
output = (num_experts, max_num_tokens * num_dp, K)
877
return (workspace13, workspace2, output)
878
@@ -950,11 +947,7 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
947
# We can reuse the memory between these because by the time we need
948
# cache3, we're done with cache1
949
intermediate_cache1 = _resize_cache(workspace13, (E, max_num_tokens, N))
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
- )
950
+ intermediate_cache2 = _resize_cache(workspace2, (E, max_num_tokens, N // 2))
951
952
# TODO(bnell): should this be done for any quantized type?
953
if self.quant_config.use_fp8_w8a8:
@@ -985,7 +978,7 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
978
# TODO (bnell): use triton utility from batched deep gemm.
979
self.activation(
980
activation,
988
- intermediate_cache2.view(-1, intermediate_size),
981
+ intermediate_cache2.view(-1, N // 2),
982
intermediate_cache1.view(-1, N),
983
)
984
vllm/model_executor/layers/fused_moe/fused_moe.py
+2
-7
@@ -2296,10 +2296,7 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
2296
local_num_experts: int,
2297
expert_tokens_meta: mk.ExpertTokensMetadata | None,
2298
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
2299
- # For fused activations (SwiGLU): N = 2 * intermediate, after act = N/2
2300
- # For non-fused activations: N = intermediate, after act = N
2301
- intermediate_size = N // 2 if self.quant_config.is_act_and_mul else N
2302
- workspace1 = (M, topk, max(intermediate_size, K))
2299
+ workspace1 = (M, topk, max(N // 2, K))
2300
workspace2 = (M, topk, max(N, K))
2301
output = (M, K)
2302
return (workspace1, workspace2, output)
@@ -2374,10 +2371,8 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
2371
2372
# Note that the output tensor might be in workspace1
2373
intermediate_cache1 = _resize_cache(workspace2, (num_tokens, top_k_num, N))
2377
- # For fused activations (SwiGLU): output is N/2, for non-fused: output is N
2378
- intermediate_size = N // 2 if self.quant_config.is_act_and_mul else N
2374
intermediate_cache2 = _resize_cache(
2380
- workspace13, (num_tokens * top_k_num, intermediate_size)
2375
+ workspace13, (num_tokens * top_k_num, N // 2)
2376
)
2377
intermediate_cache3 = _resize_cache(workspace2, (num_tokens, top_k_num, K))
2378
vllm/model_executor/layers/fused_moe/layer.py
+2
-8
@@ -600,15 +600,9 @@ class FusedMoE(CustomOp):
600
"is_act_and_mul=False is supported only for unquantized "
601
", ModelOpt FP8, and ModelOpt NvFp4 checkpoints"
602
)
603
- # ROCm without AITER MoE uses Triton which supports
604
- # is_act_and_mul=False via standard PyTorch ops (F.silu, F.gelu)
605
- rocm_without_aiter_moe = (
606
- current_platform.is_rocm() and not rocm_aiter_ops.is_fused_moe_enabled()
607
- )
608
- if not current_platform.is_cuda() and not rocm_without_aiter_moe:
603
+ if not current_platform.is_cuda():
604
raise NotImplementedError(
610
- "is_act_and_mul=False is supported only for CUDA, or ROCm "
611
- "(when AITER MoE is disabled) for now"
605
+ "is_act_and_mul=False is supported only for CUDA for now"
606
)
607
608
if self.enable_eplb and not self.quant_method.supports_eplb:
vllm/model_executor/layers/fused_moe/modular_kernel.py
+2
-23
@@ -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, sqrt
7
+from math import prod
8
from typing import final
9
10
import torch
@@ -575,35 +575,14 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
575
def activation(
576
self, activation: str, output: torch.Tensor, input: torch.Tensor
577
) -> None:
578
- # Fused activations (SwiGLU-style): output is half the size of input
578
+ assert output.size(-1) * 2 == input.size(-1)
579
if activation == "silu":
580
- assert output.size(-1) * 2 == input.size(-1)
580
torch.ops._C.silu_and_mul(output, input)
581
elif activation == "gelu":
583
- assert output.size(-1) * 2 == input.size(-1)
582
torch.ops._C.gelu_and_mul(output, input)
583
elif activation == "swigluoai":
584
# alpha = 1.702, limit = 7.0
587
- assert output.size(-1) * 2 == input.size(-1)
585
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_()
586
else:
587
raise ValueError(f"Unsupported FusedMoe activation: {activation}")
588
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
-5
@@ -299,12 +299,7 @@ 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,
302
)
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)
303
else:
304
return FUSED_MOE_UNQUANTIZED_CONFIG
305