@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 3f72639d36
Fork 0

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, 06:30 UTC

[FIX] Add NO_MUL activation support for modular kernel path (#31528)

Signed-off-by: dafrimi <dafrimi@nvidia.com> Signed-off-by: <> Co-authored-by: root <root@gpu-267.slurm-workers-slurm.slurm.svc.cluster.local> Co-authored-by: root <root@gpu-537.slurm-workers-slurm.slurm.svc.cluster.local> Co-authored-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: root <root@pool0-01777.cm.cluster>

danielafrimi committed Jan 12, 2026 at 18:55 UTC 3f72639d36a15b348c3dfee211daa3da8140736f
17 files changed +368 -71
tests/kernels/moe/test_triton_moe_no_act_mul.py new
+201
@@ -0,0 +1,201 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for MoE with non-gated activations (*_no_mul).
4 +
5 +These tests verify that MoE layers work correctly with activations like
6 +silu_no_mul, gelu_no_mul, relu2_no_mul where the activation output dimension
7 +equals N (not N // 2 like gated activations).
8 +"""
9 +
10 +import pytest
11 +import torch
12 +
13 +from vllm.model_executor.layers.fused_moe.config import (
14 + FUSED_MOE_UNQUANTIZED_CONFIG,
15 +)
16 +from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts
17 +from vllm.model_executor.layers.fused_moe.utils import (
18 + GELU_NO_MUL,
19 + RELU2_NO_MUL,
20 + SILU_NO_MUL,
21 +)
22 +from vllm.platforms import current_platform
23 +
24 +# Test parameters
25 +M_SIZES = [1, 16, 64]
26 +N_SIZES = [128, 256]
27 +K_SIZES = [64, 128]
28 +TOPK_VALUES = [1, 2]
29 +NUM_EXPERTS = 8
30 +NO_MUL_ACTIVATIONS = [SILU_NO_MUL, GELU_NO_MUL, RELU2_NO_MUL]
31 +
32 +
33 +def make_test_tensors(
34 + m: int,
35 + n: int,
36 + k: int,
37 + num_experts: int,
38 + topk: int,
39 + dtype: torch.dtype = torch.bfloat16,
40 + device: str = "cuda",
41 +):
42 + """Create test tensors for MoE with non-gated activation.
43 +
44 + For non-gated activations (*_no_mul):
45 + - w1: (E, N, K) - projects from K to N
46 + - w2: (E, K, N) - projects from N back to K (note: N, not N//2)
47 + """
48 + hidden_states = torch.randn(m, k, dtype=dtype, device=device)
49 +
50 + # For non-gated: w1 projects K -> N, w2 projects N -> K
51 + w1 = torch.randn(num_experts, n, k, dtype=dtype, device=device) * 0.1
52 + w2 = torch.randn(num_experts, k, n, dtype=dtype, device=device) * 0.1
53 +
54 + topk_weights = torch.ones(m, topk, dtype=torch.float32, device=device) / topk
55 + topk_ids = torch.randint(0, num_experts, (m, topk), device=device)
56 +
57 + return hidden_states, w1, w2, topk_weights, topk_ids
58 +
59 +
60 +@pytest.mark.skipif(
61 + not current_platform.has_device_capability(80),
62 + reason="Requires compute capability >= 8.0",
63 +)
64 +@pytest.mark.parametrize("m", M_SIZES)
65 +@pytest.mark.parametrize("n", N_SIZES)
66 +@pytest.mark.parametrize("k", K_SIZES)
67 +@pytest.mark.parametrize("topk", TOPK_VALUES)
68 +@pytest.mark.parametrize("activation", NO_MUL_ACTIVATIONS)
69 +@torch.inference_mode()
70 +def test_triton_experts_no_mul_activation(
71 + m: int,
72 + n: int,
73 + k: int,
74 + topk: int,
75 + activation: str,
76 +):
77 + hidden_states, w1, w2, topk_weights, topk_ids = make_test_tensors(
78 + m, n, k, NUM_EXPERTS, topk
79 + )
80 +
81 + experts = TritonExperts(FUSED_MOE_UNQUANTIZED_CONFIG)
82 +
83 + ws1_shape, ws2_shape, out_shape = experts.workspace_shapes(
84 + M=m,
85 + N=n,
86 + K=k,
87 + topk=topk,
88 + global_num_experts=NUM_EXPERTS,
89 + local_num_experts=NUM_EXPERTS,
90 + expert_tokens_meta=None,
91 + activation=activation,
92 + )
93 +
94 + # Verify workspace shapes are correct for no_mul activation
95 + # workspace1 should handle activation_out_dim = N (not N//2)
96 + assert ws1_shape == (m, topk, max(n, k)), (
97 + f"workspace1 shape mismatch: expected {(m, topk, max(n, k))}, got {ws1_shape}"
98 + )
99 + # workspace2 should handle max(N, K) for intermediate_cache1/cache3
100 + assert ws2_shape == (m, topk, max(n, k)), (
101 + f"workspace2 shape mismatch: expected {(m, topk, max(n, k))}, got {ws2_shape}"
102 + )
103 + assert out_shape == (m, k), (
104 + f"output shape mismatch: expected {(m, k)}, got {out_shape}"
105 + )
106 +
107 + workspace1 = torch.empty(
108 + ws1_shape[0] * ws1_shape[1] * ws1_shape[2],
109 + dtype=hidden_states.dtype,
110 + device=hidden_states.device,
111 + )
112 + workspace2 = torch.empty(
113 + ws2_shape[0] * ws2_shape[1] * ws2_shape[2],
114 + dtype=hidden_states.dtype,
115 + device=hidden_states.device,
116 + )
117 + output = torch.zeros(m, k, dtype=hidden_states.dtype, device=hidden_states.device)
118 +
119 + experts.apply(
120 + output=output,
121 + hidden_states=hidden_states,
122 + w1=w1,
123 + w2=w2,
124 + topk_weights=topk_weights,
125 + topk_ids=topk_ids,
126 + activation=activation,
127 + global_num_experts=NUM_EXPERTS,
128 + expert_map=None,
129 + a1q_scale=None,
130 + a2_scale=None,
131 + workspace13=workspace1,
132 + workspace2=workspace2,
133 + expert_tokens_meta=None,
134 + apply_router_weight_on_input=False,
135 + )
136 +
137 + assert output.shape == (m, k), f"Expected shape {(m, k)}, got {output.shape}"
138 + assert not torch.isnan(output).any(), "Output contains NaN"
139 + assert not torch.isinf(output).any(), "Output contains Inf"
140 + assert output.abs().sum() > 0, "Output is all zeros"
141 +
142 +
143 +@pytest.mark.skipif(
144 + not current_platform.has_device_capability(80),
145 + reason="Requires compute capability >= 8.0",
146 +)
147 +@torch.inference_mode()
148 +def test_workspace_shapes_no_mul_vs_gated():
149 + """Test that workspace shapes differ correctly between gated and non-gated."""
150 + from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts
151 +
152 + M, N, K, topk = 64, 256, 128, 2
153 +
154 + experts = TritonExperts(FUSED_MOE_UNQUANTIZED_CONFIG)
155 +
156 + ws1_no_mul, _, out_no_mul = experts.workspace_shapes(
157 + M, N, K, topk, 8, 8, None, SILU_NO_MUL
158 + )
159 +
160 + ws1_gated, _, out_gated = experts.workspace_shapes(
161 + M, N, K, topk, 8, 8, None, "silu"
162 + )
163 +
164 + # For no_mul: activation_out_dim = N
165 + # For gated: activation_out_dim = N // 2
166 + # workspace1 should use max(activation_out_dim, K)
167 + activation_out_dim_no_mul = N
168 + activation_out_dim_gated = N // 2
169 +
170 + assert ws1_no_mul[2] == max(activation_out_dim_no_mul, K), (
171 + f"no_mul workspace1 last dim should be max({activation_out_dim_no_mul}, {K})"
172 + )
173 + assert ws1_gated[2] == max(activation_out_dim_gated, K), (
174 + f"gated workspace1 last dim should be max({activation_out_dim_gated}, {K})"
175 + )
176 +
177 + # Output shapes should be the same
178 + assert out_no_mul == out_gated == (M, K)
179 +
180 +
181 +@pytest.mark.skipif(
182 + not current_platform.has_device_capability(80),
183 + reason="Requires compute capability >= 8.0",
184 +)
185 +@torch.inference_mode()
186 +def test_adjust_n_for_activation():
187 + """Test the adjust_N_for_activation method."""
188 + from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts
189 +
190 + experts = TritonExperts(FUSED_MOE_UNQUANTIZED_CONFIG)
191 +
192 + N = 256
193 +
194 + # Gated activations should return N // 2
195 + assert experts.adjust_N_for_activation(N, "silu") == N // 2
196 + assert experts.adjust_N_for_activation(N, "gelu") == N // 2
197 +
198 + # Non-gated activations should return N
199 + assert experts.adjust_N_for_activation(N, SILU_NO_MUL) == N
200 + assert experts.adjust_N_for_activation(N, GELU_NO_MUL) == N
201 + assert experts.adjust_N_for_activation(N, RELU2_NO_MUL) == N
vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py
+3 -1
@@ -305,6 +305,7 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
305 global_num_experts: int,
306 local_num_experts: int,
307 expert_tokens_meta: mk.ExpertTokensMetadata | None,
308 + activation: str,
309 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
310 # FIXME (varun): We should be able to dispatch only from the leader
311 # DP ranks in the case of TP > 1. At the moment, all the Ranks
@@ -312,8 +313,9 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
313 num_dispatchers = self.num_dispatchers
314 num_experts = local_num_experts
315 max_num_tokens = M if self.max_num_tokens is None else self.max_num_tokens
316 + activation_out_dim = self.adjust_N_for_activation(N, activation)
317 workspace13 = (num_experts, max_num_tokens * num_dispatchers, max(K, N))
316 - workspace2 = (num_experts, max_num_tokens * num_dispatchers, (N // 2))
318 + workspace2 = (num_experts, max_num_tokens * num_dispatchers, activation_out_dim)
319 output = (num_experts, max_num_tokens * num_dispatchers, K)
320 return (workspace13, workspace2, output)
321
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+16 -4
@@ -355,9 +355,11 @@ class CutlassExpertsFp8(CutlassExpertsFp8Base):
355 global_num_experts: int,
356 local_num_experts: int,
357 expert_tokens_meta: mk.ExpertTokensMetadata | None,
358 + activation: str,
359 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
360 + activation_out_dim = self.adjust_N_for_activation(N, activation)
361 workspace1 = (M * topk, max(N, K))
360 - workspace2 = (M * topk, max(N // 2, K))
362 + workspace2 = (M * topk, max(activation_out_dim, K))
363 output = (M, K)
364 return (workspace1, workspace2, output)
365
@@ -402,11 +404,17 @@ class CutlassBatchedExpertsFp8(CutlassExpertsFp8Base):
404 global_num_experts: int,
405 local_num_experts: int,
406 expert_tokens_meta: mk.ExpertTokensMetadata | None,
407 + activation: str,
408 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
409 num_dp = self.num_dispatchers
410 assert num_dp is not None
411 + activation_out_dim = self.adjust_N_for_activation(N, activation)
412 workspace1 = (self.max_experts_per_worker, M * num_dp, max(N, K))
409 - workspace2 = (self.max_experts_per_worker, M * num_dp, max(N // 2, K))
413 + workspace2 = (
414 + self.max_experts_per_worker,
415 + M * num_dp,
416 + max(activation_out_dim, K),
417 + )
418 output = (self.max_experts_per_worker, M, K)
419 return (workspace1, workspace2, output)
420
@@ -635,13 +643,15 @@ class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):
643 global_num_experts: int,
644 local_num_experts: int,
645 expert_tokens_meta: mk.ExpertTokensMetadata | None,
646 + activation: str,
647 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
648 + activation_out_dim = self.adjust_N_for_activation(N, activation)
649 workspace1: tuple[int, ...] = ()
650 workspace2: tuple[int, ...] = ()
651 output: tuple[int, ...] = ()
652 if self.use_batched_format:
653 workspace1 = (self.max_experts_per_worker, M, max(N, K))
644 - workspace2 = (self.max_experts_per_worker, M, (N // 2))
654 + workspace2 = (self.max_experts_per_worker, M, activation_out_dim)
655 output = (self.max_experts_per_worker, M, K)
656 else:
657 workspace1 = (M * topk, max(2 * N, K))
@@ -896,9 +906,11 @@ class CutlassExpertsW4A8Fp8(mk.FusedMoEPermuteExpertsUnpermute):
906 global_num_experts: int,
907 local_num_experts: int,
908 expert_tokens_meta: mk.ExpertTokensMetadata | None,
909 + activation: str,
910 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
911 + activation_out_dim = self.adjust_N_for_activation(N, activation)
912 workspace1 = (M * topk, max(N, K))
901 - workspace2 = (M * topk, max(N // 2, K))
913 + workspace2 = (M * topk, max(activation_out_dim, K))
914 output = (M, K)
915 return (workspace1, workspace2, output)
916
vllm/model_executor/layers/fused_moe/deep_gemm_moe.py
+12 -6
@@ -143,6 +143,7 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
143 global_num_experts: int,
144 local_num_experts: int,
145 expert_tokens_meta: mk.ExpertTokensMetadata | None,
146 + activation: str,
147 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
148 assert self.block_shape is not None
149 block_m = self.block_shape[0]
@@ -151,7 +152,8 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
152 )
153 assert M_sum % block_m == 0
154
154 - workspace1 = (M_sum, max(N // 2, K))
155 + activation_out_dim = self.adjust_N_for_activation(N, activation)
156 + workspace1 = (M_sum, max(activation_out_dim, K))
157 workspace2 = (M_sum, max(N, K))
158 output = (M, K)
159 return (workspace1, workspace2, output)
@@ -163,11 +165,13 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
165 block_k = self.block_shape[1]
166 scale_fmt = DeepGemmQuantScaleFMT.from_oracle()
167
168 + M_sum, N = input.size()
169 + activation_out_dim = self.adjust_N_for_activation(N, activation)
170 +
171 # 1. DeepGemm UE8M0: use packed per-token-group quant
172 if scale_fmt == DeepGemmQuantScaleFMT.UE8M0:
168 - M_sum, N = input.size()
173 act_out = torch.empty(
170 - (M_sum, N // 2), dtype=input.dtype, device=input.device
174 + (M_sum, activation_out_dim), dtype=input.dtype, device=input.device
175 )
176 self.activation(activation, act_out, input)
177 a2q, a2q_scale = per_token_group_quant_fp8_packed_for_deepgemm(
@@ -187,8 +191,9 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
191 )
192
193 # 3. fallback path for non-SiLU activations in non‑UE8M0 cases.
190 - M_sum, N = input.size()
191 - act_out = torch.empty((M_sum, N // 2), dtype=input.dtype, device=input.device)
194 + act_out = torch.empty(
195 + (M_sum, activation_out_dim), dtype=input.dtype, device=input.device
196 + )
197 self.activation(activation, act_out, input)
198 return per_token_group_quant_fp8(
199 act_out, block_k, column_major_scales=True, out_q=output
@@ -254,8 +259,9 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
259 (a1q, a1q_scale), (w1, self.w1_scale), mm1_out, expert_ids
260 )
261
262 + activation_out_dim = self.adjust_N_for_activation(N, activation)
263 quant_out = _resize_cache(
258 - workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, N // 2)
264 + workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, activation_out_dim)
265 )
266 a2q, a2q_scale = self._act_mul_quant(
267 input=mm1_out.view(-1, N), output=quant_out, activation=activation
vllm/model_executor/layers/fused_moe/fallback.py
+1
@@ -76,6 +76,7 @@ class FallbackExperts(mk.FusedMoEPermuteExpertsUnpermute, ABC):
76 global_num_experts: int,
77 local_num_experts: int,
78 expert_tokens_meta: mk.ExpertTokensMetadata | None,
79 + activation: str,
80 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
81 raise NotImplementedError
82
vllm/model_executor/layers/fused_moe/flashinfer_cutedsl_moe.py
+1
@@ -91,6 +91,7 @@ class FlashInferCuteDSLExperts(mk.FusedMoEPermuteExpertsUnpermute):
91 global_num_experts: int,
92 local_num_experts: int,
93 expert_tokens_meta: mk.ExpertTokensMetadata | None,
94 + activation: str,
95 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
96 # We use global_num_experts due to how moe_align_block_size handles
97 # expert_maps.
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py
+1
@@ -103,6 +103,7 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
103 global_num_experts: int,
104 local_num_experts: int,
105 expert_tokens_meta: mk.ExpertTokensMetadata | None,
106 + activation: str,
107 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
108 # We use global_num_experts due to how moe_align_block_size handles
109 # expert_maps.
vllm/model_executor/layers/fused_moe/fused_batched_moe.py
+9 -3
@@ -673,6 +673,7 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
673 global_num_experts: int,
674 local_num_experts: int,
675 expert_tokens_meta: mk.ExpertTokensMetadata | None,
676 + activation: str,
677 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
678 num_dp = self.num_dispatchers
679 num_experts = local_num_experts
@@ -867,12 +868,14 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
868 global_num_experts: int,
869 local_num_experts: int,
870 expert_tokens_meta: mk.ExpertTokensMetadata | None,
871 + activation: str,
872 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
873 num_dp = self.num_dispatchers
874 num_experts = local_num_experts
875 max_num_tokens = self.max_num_tokens
876 + activation_out_dim = self.adjust_N_for_activation(N, activation)
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, activation_out_dim)
879 output = (num_experts, max_num_tokens * num_dp, K)
880 return (workspace13, workspace2, output)
881
@@ -947,7 +950,10 @@ 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 + activation_out_dim = self.adjust_N_for_activation(N, activation)
954 + intermediate_cache2 = _resize_cache(
955 + workspace2, (E, max_num_tokens, activation_out_dim)
956 + )
957
958 # TODO(bnell): should this be done for any quantized type?
959 if self.quant_config.use_fp8_w8a8:
@@ -978,7 +984,7 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
984 # TODO (bnell): use triton utility from batched deep gemm.
985 self.activation(
986 activation,
981 - intermediate_cache2.view(-1, N // 2),
987 + intermediate_cache2.view(-1, activation_out_dim),
988 intermediate_cache1.view(-1, N),
989 )
990
vllm/model_executor/layers/fused_moe/fused_marlin_moe.py
+2
@@ -640,6 +640,7 @@ class MarlinExperts(MarlinExpertsBase):
640 global_num_experts: int,
641 local_num_experts: int,
642 expert_tokens_meta: mk.ExpertTokensMetadata | None,
643 + activation: str,
644 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
645 # Modular Kernel provisions output buffer from workspace1. However in
646 # the fused_marlin_moe() function, the final torch.sum(), is defined
@@ -768,6 +769,7 @@ class BatchedMarlinExperts(MarlinExpertsBase):
769 global_num_experts: int,
770 local_num_experts: int,
771 expert_tokens_meta: mk.ExpertTokensMetadata | None,
772 + activation: str,
773 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
774 num_dispatchers = self.num_dispatchers
775 num_experts = local_num_experts
vllm/model_executor/layers/fused_moe/fused_moe.py
+17 -34
@@ -9,7 +9,6 @@ from collections.abc import Callable
9 from typing import Any
10
11 import torch
12 -import torch.nn.functional as F
12
13 import vllm.envs as envs
14 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
@@ -43,7 +42,7 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
42 )
43 from vllm.model_executor.layers.fused_moe.utils import (
44 _resize_cache,
46 - activation_without_mul,
45 + apply_moe_activation,
46 disable_inplace,
47 moe_kernel_quantize_input,
48 )
@@ -1957,11 +1956,6 @@ def fused_experts(
1956 )
1957
1958
1960 -SILU_NO_MUL: str = activation_without_mul("silu")
1961 -GELU_NO_MUL: str = activation_without_mul("gelu")
1962 -RELU2_NO_MUL: str = activation_without_mul("relu2")
1963 -
1964 -
1959 def _get_config_quant_dtype(
1960 use_fp8_w8a8: bool,
1961 use_int8_w8a8: bool,
@@ -2094,8 +2088,13 @@ def fused_experts_impl(
2088 intermediate_cache3 = cache13[: M * top_k_num * K].view(M, top_k_num, K)
2089
2090 # This needs separate memory since it's used concurrently with cache1
2091 + activation_out_dim = mk.FusedMoEPermuteExpertsUnpermute.adjust_N_for_activation(
2092 + N, activation
2093 + )
2094 intermediate_cache2 = torch.empty(
2098 - (M * top_k_num, N // 2), device=hidden_states.device, dtype=hidden_states.dtype
2095 + (M * top_k_num, activation_out_dim),
2096 + device=hidden_states.device,
2097 + dtype=hidden_states.dtype,
2098 )
2099
2100 if hidden_states.dtype == torch.bfloat16:
@@ -2235,29 +2234,9 @@ def fused_experts_impl(
2234 B_bias=w1_bias,
2235 )
2236
2238 - # Activation function with multiplication
2239 - if activation == "silu":
2240 - torch.ops._C.silu_and_mul(
2241 - intermediate_cache2, intermediate_cache1.view(-1, N)
2242 - )
2243 - elif activation == "gelu":
2244 - torch.ops._C.gelu_and_mul(
2245 - intermediate_cache2, intermediate_cache1.view(-1, N)
2246 - )
2247 - elif activation == "swigluoai":
2248 - # alpha = 1.702, limit = 7.0
2249 - torch.ops._C.swigluoai_and_mul(
2250 - intermediate_cache2, intermediate_cache1.view(-1, N)
2251 - )
2252 - # Activation function without multiplication
2253 - elif activation == SILU_NO_MUL:
2254 - intermediate_cache2 = F.silu(intermediate_cache1.view(-1, N))
2255 - elif activation == GELU_NO_MUL:
2256 - intermediate_cache2 = F.gelu(intermediate_cache1.view(-1, N))
2257 - elif activation == RELU2_NO_MUL:
2258 - intermediate_cache2 = torch.square(F.relu(intermediate_cache1.view(-1, N)))
2259 - else:
2260 - raise ValueError(f"Unsupported FusedMoe activation: {activation}.")
2237 + apply_moe_activation(
2238 + activation, intermediate_cache2, intermediate_cache1.view(-1, N)
2239 + )
2240
2241 qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(
2242 A=intermediate_cache2,
@@ -2336,8 +2315,10 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
2315 global_num_experts: int,
2316 local_num_experts: int,
2317 expert_tokens_meta: mk.ExpertTokensMetadata | None,
2318 + activation: str,
2319 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
2340 - workspace1 = (M, topk, max(N // 2, K))
2320 + activation_out_dim = self.adjust_N_for_activation(N, activation)
2321 + workspace1 = (M, topk, max(activation_out_dim, K))
2322 workspace2 = (M, topk, max(N, K))
2323 output = (M, K)
2324 return (workspace1, workspace2, output)
@@ -2412,8 +2393,9 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
2393
2394 # Note that the output tensor might be in workspace1
2395 intermediate_cache1 = _resize_cache(workspace2, (num_tokens, top_k_num, N))
2396 + cache2_dim = self.adjust_N_for_activation(N, activation)
2397 intermediate_cache2 = _resize_cache(
2416 - workspace13, (num_tokens * top_k_num, N // 2)
2398 + workspace13, (num_tokens * top_k_num, cache2_dim)
2399 )
2400 intermediate_cache3 = _resize_cache(workspace2, (num_tokens, top_k_num, K))
2401
@@ -2565,8 +2547,9 @@ class TritonWNA16Experts(TritonExperts):
2547
2548 # Note that the output tensor might be in workspace1
2549 intermediate_cache1 = _resize_cache(workspace2, (num_tokens, top_k_num, N))
2550 + activation_out_dim = self.adjust_N_for_activation(N, activation)
2551 intermediate_cache2 = _resize_cache(
2569 - workspace13, (num_tokens * top_k_num, N // 2)
2552 + workspace13, (num_tokens * top_k_num, activation_out_dim)
2553 )
2554 intermediate_cache3 = _resize_cache(workspace2, (num_tokens, top_k_num, K))
2555
vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py
+18 -13
@@ -323,10 +323,12 @@ class OAITritonExperts(BaseOAITritonExperts):
323 global_num_experts: int,
324 local_num_experts: int,
325 expert_tokens_meta: mk.ExpertTokensMetadata | None,
326 + activation: str,
327 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
328 # workspace are allocated inside the kernel
329 + activation_out_dim = self.adjust_N_for_activation(N, activation)
330 workspace1 = (0, 0)
329 - workspace2 = (M * topk, N // 2)
331 + workspace2 = (M * topk, activation_out_dim)
332 output = (M, K)
333 return (workspace1, workspace2, output)
334
@@ -415,9 +417,11 @@ class UnfusedOAITritonExperts(BaseOAITritonExperts):
417 global_num_experts: int,
418 local_num_experts: int,
419 expert_tokens_meta: mk.ExpertTokensMetadata | None,
420 + activation: str,
421 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
422 # workspace are allocated inside the kernel
420 - workspace1 = (M * topk, N // 2)
423 + activation_out_dim = self.adjust_N_for_activation(N, activation)
424 + workspace1 = (M * topk, activation_out_dim)
425 workspace2 = (M * topk, max(N, K))
426 output = (M, K)
427 return (workspace1, workspace2, output)
@@ -443,8 +447,10 @@ class UnfusedOAITritonExperts(BaseOAITritonExperts):
447 expert_tokens_meta: mk.ExpertTokensMetadata | None,
448 apply_router_weight_on_input: bool,
449 ):
446 - if self.quant_config is None:
447 - self.quant_config = FUSED_MOE_UNQUANTIZED_CONFIG
450 + # Use local variable to help mypy narrow the type after None check
451 + quant_config = self.quant_config
452 + if quant_config is None:
453 + quant_config = FUSED_MOE_UNQUANTIZED_CONFIG
454
455 if expert_map is not None:
456 topk_ids = expert_map[topk_ids]
@@ -462,12 +468,10 @@ class UnfusedOAITritonExperts(BaseOAITritonExperts):
468 # type check, uint8 means mxfp4
469 assert hidden_states.dtype == torch.bfloat16
470 assert (
465 - self.quant_config.w1_bias is None
466 - or self.quant_config.w1_bias.dtype == torch.float32
471 + quant_config.w1_bias is None or quant_config.w1_bias.dtype == torch.float32
472 )
473 assert (
469 - self.quant_config.w2_bias is None
470 - or self.quant_config.w2_bias.dtype == torch.float32
474 + quant_config.w2_bias is None or quant_config.w2_bias.dtype == torch.float32
475 )
476
477 # Shape check, only check non-mxfp4
@@ -485,17 +489,18 @@ class UnfusedOAITritonExperts(BaseOAITritonExperts):
489 # Note that the output tensor might be in workspace13
490 intermediate_cache1 = _resize_cache(workspace2, (batch_dim, M * topk, N))
491 intermediate_cache3 = _resize_cache(workspace2, (batch_dim, M * topk, K))
488 - intermediate_cache2 = _resize_cache(workspace13, (M * topk, N // 2))
492 + activation_out_dim = self.adjust_N_for_activation(N, activation)
493 + intermediate_cache2 = _resize_cache(workspace13, (M * topk, activation_out_dim))
494
495 gammas = routing_data.gate_scal if routing_data else None
496
497 matmul_ogs(
498 hidden_states,
499 w1,
495 - self.quant_config.w1_bias,
500 + quant_config.w1_bias,
501 routing_data,
502 gather_indx=gather_indx,
498 - precision_config=self.quant_config.w1_precision,
503 + precision_config=quant_config.w1_precision,
504 gammas=gammas if apply_router_weight_on_input else None,
505 fused_activation=None,
506 y=intermediate_cache1,
@@ -515,10 +520,10 @@ class UnfusedOAITritonExperts(BaseOAITritonExperts):
520 matmul_ogs(
521 intermediate_cache2[gather_indx.src_indx],
522 w2,
518 - self.quant_config.w2_bias,
523 + quant_config.w2_bias,
524 routing_data,
525 scatter_indx=scatter_indx,
521 - precision_config=self.quant_config.w2_precision,
526 + precision_config=quant_config.w2_precision,
527 gammas=None if apply_router_weight_on_input else gammas,
528 y=intermediate_cache3,
529 )
vllm/model_executor/layers/fused_moe/modular_kernel.py
+29 -10
@@ -18,6 +18,7 @@ from vllm.model_executor.layers.fused_moe.config import (
18 )
19 from vllm.model_executor.layers.fused_moe.utils import (
20 _resize_cache,
21 + apply_moe_activation,
22 count_expert_num_tokens,
23 disable_inplace,
24 )
@@ -542,6 +543,7 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
543 global_num_experts: int,
544 local_num_experts: int,
545 expert_tokens_meta: ExpertTokensMetadata | None,
546 + activation: str,
547 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
548 """
549 Compute the shapes for the temporary and final outputs of the two gemms
@@ -572,19 +574,31 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
574 """
575 raise NotImplementedError
576
577 + @staticmethod
578 + def adjust_N_for_activation(N: int, activation: str) -> int:
579 + """
580 + Calculate the output dimension for the activation function.
581 +
582 + For *_no_mul activations (e.g. relu2_no_mul),
583 + there's no gate/up split, so output size equals input size (N).
584 +
585 + For regular gated activations (e.g., silu, gelu, swigluoai),
586 + output size is N // 2 due to gate × activation(up) multiplication.
587 +
588 + Args:
589 + N: The intermediate size (width of w1/w3 weights).
590 + activation: The activation function name.
591 +
592 + Returns:
593 + The output dimension after activation.
594 + """
595 + is_no_mul = activation.endswith("_no_mul")
596 + return N if is_no_mul else N // 2
597 +
598 def activation(
599 self, activation: str, output: torch.Tensor, input: torch.Tensor
600 ) -> None:
578 - assert output.size(-1) * 2 == input.size(-1)
579 - if activation == "silu":
580 - torch.ops._C.silu_and_mul(output, input)
581 - elif activation == "gelu":
582 - torch.ops._C.gelu_and_mul(output, input)
583 - elif activation == "swigluoai":
584 - # alpha = 1.702, limit = 7.0
585 - torch.ops._C.swigluoai_and_mul(output, input)
586 - else:
587 - raise ValueError(f"Unsupported FusedMoe activation: {activation}")
601 + apply_moe_activation(activation, output, input)
602
603 def enable_chunking(self):
604 return (
@@ -761,6 +775,7 @@ class FusedMoEModularKernel(torch.nn.Module):
775 global_num_experts: int,
776 local_num_experts: int,
777 expert_tokens_meta: ExpertTokensMetadata | None,
778 + activation: str,
779 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
780 """
781 Allocate temporary and output buffers for the fused experts op.
@@ -796,6 +811,7 @@ class FusedMoEModularKernel(torch.nn.Module):
811 # amount of workspace. Mark it None, so we allocate for
812 # the worst-case scenario.
813 expert_tokens_meta=None,
814 + activation=activation,
815 )
816 )
817
@@ -814,6 +830,7 @@ class FusedMoEModularKernel(torch.nn.Module):
830 global_num_experts,
831 local_num_experts,
832 expert_tokens_meta,
833 + activation,
834 )
835
836 # Get final output shape based on the full M size.
@@ -825,6 +842,7 @@ class FusedMoEModularKernel(torch.nn.Module):
842 global_num_experts,
843 local_num_experts,
844 expert_tokens_meta,
845 + activation,
846 )
847
848 # We can reuse the memory between cache1 and cache3 because by the
@@ -1043,6 +1061,7 @@ class FusedMoEModularKernel(torch.nn.Module):
1061 global_num_experts,
1062 local_num_experts,
1063 expert_tokens_meta,
1064 + activation,
1065 )
1066
1067 for chunk_idx in range(num_chunks):
vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
+1
@@ -299,6 +299,7 @@ class AiterExperts(mk.FusedMoEPermuteExpertsUnpermute):
299 global_num_experts: int,
300 local_num_experts: int,
301 expert_tokens_meta: mk.ExpertTokensMetadata | None,
302 + activation: str,
303 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
304 # Workspaces are managed internally by AITER.
305 workspace1 = (0,)
vllm/model_executor/layers/fused_moe/triton_cutlass_moe.py
+3
@@ -39,6 +39,7 @@ class TritonOrCutlassExperts(FallbackExperts):
39 global_num_experts: int,
40 local_num_experts: int,
41 expert_tokens_meta: mk.ExpertTokensMetadata | None,
42 + activation: str,
43 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
44 # Small batch fallback for sm100.
45 if self.is_sm100 and M <= 8:
@@ -50,6 +51,7 @@ class TritonOrCutlassExperts(FallbackExperts):
51 global_num_experts,
52 local_num_experts,
53 expert_tokens_meta,
54 + activation,
55 )
56 else:
57 return self.experts.workspace_shapes(
@@ -60,6 +62,7 @@ class TritonOrCutlassExperts(FallbackExperts):
62 global_num_experts,
63 local_num_experts,
64 expert_tokens_meta,
65 + activation,
66 )
67
68 def _select_experts_impl(
vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py
+3
@@ -35,6 +35,7 @@ class TritonOrDeepGemmExperts(FallbackExperts):
35 global_num_experts: int,
36 local_num_experts: int,
37 expert_tokens_meta: mk.ExpertTokensMetadata | None,
38 + activation: str,
39 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
40 # Note: the deep gemm workspaces are strictly larger than the triton
41 # workspaces so we can be pessimistic here and allocate for DeepGemm
@@ -48,6 +49,7 @@ class TritonOrDeepGemmExperts(FallbackExperts):
49 global_num_experts,
50 local_num_experts,
51 expert_tokens_meta,
52 + activation,
53 )
54 else:
55 return self.fallback_experts.workspace_shapes(
@@ -58,6 +60,7 @@ class TritonOrDeepGemmExperts(FallbackExperts):
60 global_num_experts,
61 local_num_experts,
62 expert_tokens_meta,
63 + activation,
64 )
65
66 def _select_experts_impl(
vllm/model_executor/layers/fused_moe/trtllm_moe.py
+1
@@ -57,6 +57,7 @@ class TrtLlmGenExperts(mk.FusedMoEPermuteExpertsUnpermute):
57 global_num_experts: int,
58 local_num_experts: int,
59 expert_tokens_meta: mk.ExpertTokensMetadata | None,
60 + activation: str,
61 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
62 # The workspaces for this implementation are managed by flashinfer.
63 workspace1 = (0,)
vllm/model_executor/layers/fused_moe/utils.py
+50
@@ -4,6 +4,7 @@ import functools
4 from math import prod
5
6 import torch
7 +import torch.nn.functional as F
8
9 from vllm import _custom_ops as ops
10 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
@@ -324,6 +325,55 @@ def activation_without_mul(activation: str) -> str:
325 return activation + "_no_mul"
326
327
328 +RELU2_NO_MUL: str = activation_without_mul("relu2")
329 +SILU_NO_MUL: str = activation_without_mul("silu")
330 +GELU_NO_MUL: str = activation_without_mul("gelu")
331 +
332 +
333 +def apply_moe_activation(
334 + activation: str,
335 + output: torch.Tensor,
336 + input: torch.Tensor,
337 +) -> torch.Tensor:
338 + """
339 + Apply MoE activation function.
340 +
341 + For *_and_mul activations (silu, gelu, swigluoai):
342 + - Expects output.size(-1) * 2 == input.size(-1)
343 +
344 + For *_no_mul activations (silu_no_mul, gelu_no_mul, relu2_no_mul):
345 + - Expects output.size(-1) == input.size(-1)
346 + """
347 + is_no_mul = activation.endswith("_no_mul")
348 + if is_no_mul:
349 + assert output.size(-1) == input.size(-1), (
350 + f"{activation} expects equal sizes: {output.size(-1)} vs {input.size(-1)}"
351 + )
352 + else:
353 + assert output.size(-1) * 2 == input.size(-1), (
354 + f"{activation} expects 2x ratio: {output.size(-1) * 2} vs {input.size(-1)}"
355 + )
356 +
357 + # Activations with gated multiplication (gate × activation(up))
358 + if activation == "silu":
359 + torch.ops._C.silu_and_mul(output, input)
360 + elif activation == "gelu":
361 + torch.ops._C.gelu_and_mul(output, input)
362 + elif activation == "swigluoai":
363 + torch.ops._C.swigluoai_and_mul(output, input)
364 + # Activations without gated multiplication
365 + elif activation == SILU_NO_MUL:
366 + output.copy_(F.silu(input))
367 + elif activation == GELU_NO_MUL:
368 + output.copy_(F.gelu(input))
369 + elif activation == RELU2_NO_MUL:
370 + torch.square(F.relu(input), out=output)
371 + else:
372 + raise ValueError(f"Unsupported FusedMoe activation: {activation}")
373 +
374 + return output
375 +
376 +
377 # Torch custom ops can't deal with outputs aliasing inputs so we need to
378 # disable inplace for torch >= 2.9.
379 # See https://github.com/vllm-project/vllm/issues/26378