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 12, 13:30 UTC
[Kernel] Use pre-allocated output buffer for triton kernel fused_experts (#29219)
Signed-off-by: Xin Yang <xyangx@amazon.com>
Xin Yang committed
Nov 25, 2025 at 18:21 UTC
53d7f1f601a12b8fa58aa0bffb7fc27d63d1eb5e
1 file changed
+73
-11
vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py
+73
-11
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+
5
import torch
6
7
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
@@ -12,6 +13,7 @@ from vllm.model_executor.layers.fused_moe.config import (
13
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
14
TopKWeightAndReduceNoOP,
15
)
16
+from vllm.model_executor.layers.fused_moe.utils import _resize_cache
17
from vllm.triton_utils import tl, triton
18
from vllm.utils.import_utils import has_triton_kernels
19
@@ -88,14 +90,17 @@ def triton_kernel_moe_forward(
90
gating_output, topk, sm_first=not renormalize
91
)
92
93
+ output = torch.empty_like(hidden_states)
94
+
95
return triton_kernel_fused_experts(
92
- None,
96
+ output,
97
hidden_states,
98
w1,
99
w2,
100
routing_data,
101
gather_idx,
102
scatter_idx,
103
+ topk=topk,
104
activation=activation,
105
quant_config=quant_config,
106
apply_router_weight_on_input=apply_router_weight_on_input,
@@ -113,6 +118,7 @@ def triton_kernel_fused_experts(
118
routing_data, # RoutingData
119
gather_indx, # GatherIndx
120
scatter_indx, # ScatterIndx
121
+ topk: int,
122
activation: str = "silu",
123
quant_config: FusedMoEQuantConfig | None = None,
124
swiglu_alpha: float = 1.702,
@@ -120,6 +126,7 @@ def triton_kernel_fused_experts(
126
apply_router_weight_on_input: bool = False,
127
global_num_experts: int = -1,
128
expert_map: torch.Tensor | None = None,
129
+ intermediate_cache: torch.Tensor | None = None,
130
a1q_scale: torch.Tensor | None = None,
131
) -> torch.Tensor:
132
if quant_config is None:
@@ -131,14 +138,30 @@ def triton_kernel_fused_experts(
138
assert quant_config.w2_bias is None or quant_config.w2_bias.dtype == torch.float32
139
140
# Shape check, only check non-mxfp4
141
+ assert hidden_states.ndim == 2
142
assert hidden_states.shape[-1] == w1.shape[-2]
143
assert w2.shape[-1] == w1.shape[1]
144
145
+ batch_dim = 1
146
+ M, K = hidden_states.shape[-2:]
147
E, _, N = w1.shape
148
149
if global_num_experts == -1:
150
global_num_experts = E
151
152
+ if intermediate_cache is None:
153
+ intermediate_cache = torch.empty(
154
+ (batch_dim, M * topk, N // 2),
155
+ device=hidden_states.device,
156
+ dtype=hidden_states.dtype,
157
+ )
158
+
159
+ # Add batch_dim to output buffer because matmul_ogs expects 3D output
160
+ intermediate_cache = _resize_cache(
161
+ intermediate_cache, (batch_dim, M * topk, N // 2)
162
+ )
163
+ output_tensor = _resize_cache(output_tensor, (batch_dim, M, K))
164
+
165
act = FusedActivation(
166
FnSpecs("swiglu", triton_kernels.swiglu.swiglu_fn, ("alpha", "limit")),
167
(swiglu_alpha, swiglu_limit),
@@ -146,7 +169,7 @@ def triton_kernel_fused_experts(
169
)
170
gammas = routing_data.gate_scal if routing_data else None
171
149
- intermediate_cache1 = matmul_ogs(
172
+ matmul_ogs(
173
hidden_states,
174
w1,
175
quant_config.w1_bias,
@@ -155,10 +178,11 @@ def triton_kernel_fused_experts(
178
precision_config=quant_config.w1_precision,
179
gammas=gammas if apply_router_weight_on_input else None,
180
fused_activation=act,
181
+ y=intermediate_cache,
182
)
183
160
- intermediate_cache3 = matmul_ogs(
161
- intermediate_cache1,
184
+ matmul_ogs(
185
+ intermediate_cache.view(M * topk, N // 2),
186
w2,
187
quant_config.w2_bias,
188
routing_data,
@@ -167,7 +191,8 @@ def triton_kernel_fused_experts(
191
gammas=None if apply_router_weight_on_input else gammas,
192
y=output_tensor,
193
)
170
- return intermediate_cache3
194
+ output_tensor = output_tensor.view(M, K)
195
+ return output_tensor
196
197
198
def make_routing_data(
@@ -221,6 +246,42 @@ class BaseOAITritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
246
def supports_expert_map(self) -> bool:
247
return True
248
249
+ def moe_problem_size(
250
+ self,
251
+ a1: torch.Tensor,
252
+ w1: torch.Tensor,
253
+ w2: torch.Tensor,
254
+ topk_ids: torch.Tensor,
255
+ ) -> tuple[int, int, int, int, int]:
256
+ """
257
+ Extract the MoE problem size from the given tensor arguments:
258
+ - a: The hidden states, input to the MoE layer.
259
+ - w1: The first set of expert weights.
260
+ - w2: The second set of expert weights.
261
+ - topk_ids: The topk ids.
262
+ Note: extracting the problem shape from the weight and activation
263
+ tensors is not obvious. It needs to be done this way specifically
264
+ due to subtle issues with particular kernels, e.g. the int4 kernels
265
+ divide the trailing dimension by two, so it's not "correct" to
266
+ extract N or K from the trailing dimension of w1 or w2. Similarly,
267
+ some kernels transpose the weights, so this needs to be kept in mind.
268
+ Note: This implementation covers most cases. However, if experts
269
+ require a specialized implementation, like MarlinExperts, they are free
270
+ to override this function.
271
+ """
272
+ assert w1.dim() == 3 and w2.dim() == 3
273
+ E, _, N = w1.size()
274
+ K = a1.size(-1)
275
+
276
+ assert a1.dim() == 2
277
+ assert topk_ids.size(0) == a1.size(0), f"{topk_ids.size(0)} != {a1.size(0)}"
278
+ M = a1.size(0)
279
+
280
+ assert topk_ids.dim() == 2
281
+ topk = topk_ids.size(1)
282
+
283
+ return E, M, N, K, topk
284
+
285
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
286
# Weight application and reduction happens in the fused_experts kernel.
287
return TopKWeightAndReduceNoOP()
@@ -263,8 +324,8 @@ class OAITritonExperts(BaseOAITritonExperts):
324
expert_tokens_meta: mk.ExpertTokensMetadata | None,
325
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
326
# workspace are allocated inside the kernel
266
- workspace1 = (M, K)
267
- workspace2 = (0, 0)
327
+ workspace1 = (0, 0)
328
+ workspace2 = (M * topk, N // 2)
329
output = (M, K)
330
return (workspace1, workspace2, output)
331
@@ -297,20 +358,21 @@ class OAITritonExperts(BaseOAITritonExperts):
358
topk_ids, topk_weights, local_num_experts
359
)
360
300
- experts_output = triton_kernel_fused_experts(
301
- None,
361
+ topk = topk_ids.size(1)
362
+ triton_kernel_fused_experts(
363
+ output,
364
hidden_states,
365
w1,
366
w2,
367
routing_data,
368
gather_indx,
369
scatter_indx,
370
+ topk=topk,
371
activation=activation,
372
quant_config=self.quant_config,
373
apply_router_weight_on_input=False,
374
global_num_experts=local_num_experts,
375
expert_map=None, # applied already
376
+ intermediate_cache=workspace2,
377
a1q_scale=a1q_scale,
378
)
315
-
316
- output.copy_(experts_output, non_blocking=True)