@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / eef364723c
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 15, 13:30 UTC

[FEAT] [ROCm]: AITER Fused MOE V1 Support (#16752)

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com> Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com>

vllmellm committed Apr 25, 2025 at 11:06 UTC eef364723c1850f62942d4d832c8f8a2eb31988c
3 files changed +306 -134
tests/model_executor/test_enabled_custom_ops.py
+4 -3
@@ -11,6 +11,8 @@ from vllm.model_executor.layers.fused_moe.fused_moe import (
11 dispatch_fused_experts_func, dispatch_topk_func,
12 torch_vllm_inplace_fused_experts, torch_vllm_outplace_fused_experts,
13 vllm_topk_softmax)
14 +from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
15 + is_rocm_aiter_moe_enabled)
16 from vllm.model_executor.layers.layernorm import (
17 RMSNorm, dispatch_cuda_rmsnorm_func, fused_add_rms_norm, rms_norm,
18 rocm_aiter_fused_add_rms_norm, rocm_aiter_rms_norm)
@@ -100,11 +102,10 @@ def test_enabled_ops_invalid(env: str):
102 def test_topk_dispatch(use_rocm_aiter: str, monkeypatch):
103 monkeypatch.setenv("VLLM_ROCM_USE_AITER", use_rocm_aiter)
104 topk_func = dispatch_topk_func()
103 -
105 + is_rocm_aiter_moe_enabled.cache_clear()
106 if current_platform.is_rocm() and int(use_rocm_aiter):
107 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
108 rocm_aiter_topk_softmax)
107 -
109 assert topk_func == rocm_aiter_topk_softmax
110 else:
111 assert topk_func == vllm_topk_softmax
@@ -116,11 +117,11 @@ def test_fused_experts_dispatch(use_rocm_aiter: str, inplace: bool,
117 monkeypatch):
118
119 monkeypatch.setenv("VLLM_ROCM_USE_AITER", use_rocm_aiter)
120 + is_rocm_aiter_moe_enabled.cache_clear()
121 fused_experts_func = dispatch_fused_experts_func(inplace)
122 if current_platform.is_rocm() and int(use_rocm_aiter):
123 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
124 rocm_aiter_fused_experts)
123 -
125 assert fused_experts_func == rocm_aiter_fused_experts
126 elif inplace:
127 assert fused_experts_func == torch_vllm_inplace_fused_experts
vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
+299 -128
@@ -1,31 +1,35 @@
1 # SPDX-License-Identifier: Apache-2.0
2 -from typing import List, Optional
2 +from functools import cache
3 +from typing import List, Optional, Tuple
4
5 import torch
6
6 -import vllm.envs as envs
7 +from vllm import envs
8 from vllm.platforms import current_platform
9 +from vllm.utils import direct_register_custom_op
10
11
12 +@cache
13 def is_rocm_aiter_moe_enabled() -> bool:
14 return current_platform.is_rocm() \
15 and envs.VLLM_ROCM_USE_AITER_MOE \
16 and envs.VLLM_ROCM_USE_AITER
17
18
16 -def rocm_aiter_asm_moe_tkw1(hidden_states,
17 - w1,
18 - w2,
19 - topk_weight,
20 - topk_ids,
21 - fc1_scale=None,
22 - fc2_scale=None,
23 - fc1_smooth_scale=None,
24 - fc2_smooth_scale=None,
25 - a16=False,
26 - per_tensor_quant_scale=None,
27 - expert_mask=None,
28 - activation_str: str = "silu") -> None:
19 +def rocm_aiter_asm_moe_tkw1_impl(
20 + hidden_states: torch.Tensor,
21 + w1: torch.Tensor,
22 + w2: torch.Tensor,
23 + topk_weight: torch.Tensor,
24 + topk_ids: torch.Tensor,
25 + fc1_scale: Optional[torch.Tensor] = None,
26 + fc2_scale: Optional[torch.Tensor] = None,
27 + fc1_smooth_scale: Optional[torch.Tensor] = None,
28 + fc2_smooth_scale: Optional[torch.Tensor] = None,
29 + a16: bool = False,
30 + per_tensor_quant_scale: Optional[torch.Tensor] = None,
31 + expert_mask: Optional[torch.Tensor] = None,
32 + activation_str: str = "silu") -> torch.Tensor:
33
34 from aiter import ActivationType
35 from aiter.fused_moe_bf16_asm import asm_moe_tkw1
@@ -48,34 +52,236 @@ def rocm_aiter_asm_moe_tkw1(hidden_states,
52 activation=activation)
53
54
51 -def rocm_aiter_fused_experts(
52 - hidden_states: torch.Tensor,
53 - w1: torch.Tensor,
54 - w2: torch.Tensor,
55 - topk_weights: torch.Tensor,
56 - topk_ids: torch.Tensor,
57 - inplace: bool = False,
58 - activation: str = "silu",
59 - apply_router_weight_on_input: bool = False,
60 - use_fp8_w8a8: bool = False,
61 - use_int8_w8a8: bool = False,
62 - use_int8_w8a16: bool = False,
63 - use_int4_w4a16: bool = False,
64 - per_channel_quant: bool = False,
65 - global_num_experts: int = -1,
66 - expert_map: Optional[torch.Tensor] = None,
67 - w1_scale: Optional[torch.Tensor] = None,
68 - w2_scale: Optional[torch.Tensor] = None,
69 - w1_zp: Optional[torch.Tensor] = None,
70 - w2_zp: Optional[torch.Tensor] = None,
71 - a1_scale: Optional[torch.Tensor] = None,
72 - a2_scale: Optional[torch.Tensor] = None,
73 - block_shape: Optional[List[int]] = None,
74 - allow_deep_gemm: bool = False,
75 -) -> torch.Tensor:
76 -
77 - import aiter as rocm_aiter
55 +def rocm_aiter_asm_moe_tkw1_fake(
56 + hidden_states: torch.Tensor,
57 + w1: torch.Tensor,
58 + w2: torch.Tensor,
59 + topk_weight: torch.Tensor,
60 + topk_ids: torch.Tensor,
61 + fc1_scale: Optional[torch.Tensor] = None,
62 + fc2_scale: Optional[torch.Tensor] = None,
63 + fc1_smooth_scale: Optional[torch.Tensor] = None,
64 + fc2_smooth_scale: Optional[torch.Tensor] = None,
65 + a16: bool = False,
66 + per_tensor_quant_scale: Optional[torch.Tensor] = None,
67 + expert_mask: Optional[torch.Tensor] = None,
68 + activation_str: str = "silu") -> torch.Tensor:
69 + return torch.empty_like(hidden_states)
70 +
71 +
72 +def rocm_aiter_ck_moe_impl(hidden_states: torch.Tensor, w1: torch.Tensor,
73 + w2: torch.Tensor, topk_weights: torch.Tensor,
74 + topk_ids: torch.Tensor) -> torch.Tensor:
75 + from aiter import ck_moe
76 + return ck_moe(hidden_states=hidden_states,
77 + w1=w1,
78 + w2=w2,
79 + topk_weights=topk_weights,
80 + topk_ids=topk_ids)
81 +
82 +
83 +def rocm_aiter_ck_moe_fake(hidden_states: torch.Tensor, w1: torch.Tensor,
84 + w2: torch.Tensor, topk_weights: torch.Tensor,
85 + topk_ids: torch.Tensor) -> torch.Tensor:
86 + return torch.empty_like(hidden_states)
87 +
88 +
89 +def rocm_aiter_fmoe_fp8_blockscale_g1u1_impl(
90 + topk_ids: torch.Tensor,
91 + topk_weights: torch.Tensor,
92 + hidden_states_dtype: torch.dtype,
93 + expert_mask: torch.Tensor,
94 + a1: torch.Tensor,
95 + w1: torch.Tensor,
96 + w2: torch.Tensor,
97 + w1_scale: torch.Tensor,
98 + w2_scale: torch.Tensor,
99 + a1_scale: torch.Tensor,
100 + block_shape: List[int],
101 + smooth_scale: Optional[torch.Tensor] = None) -> torch.Tensor:
102 + from aiter import fmoe_fp8_blockscale_g1u1
103 + from aiter.fused_moe_bf16_asm import moe_sorting_ck
104 +
105 + topk = topk_ids.shape[1]
106 + model_dim = w1.shape[-1]
107 + local_E = E = w1.shape[0]
108 + if expert_mask is not None:
109 + E = expert_mask.numel()
110 +
111 + (
112 + sorted_token_ids,
113 + sorted_weight_buf,
114 + sorted_expert_ids,
115 + num_valid_ids,
116 + out_asm,
117 + ) = moe_sorting_ck(topk_ids,
118 + topk_weights,
119 + E,
120 + model_dim,
121 + hidden_states_dtype,
122 + expert_mask=expert_mask)
123 +
124 + fmoe_fp8_blockscale_g1u1(out_asm, a1, w1, w2, sorted_token_ids,
125 + sorted_weight_buf, sorted_expert_ids,
126 + num_valid_ids, topk, w1_scale.view(local_E, -1),
127 + w2_scale.view(local_E, -1),
128 + a1_scale.t().contiguous(), *block_shape,
129 + smooth_scale)
130 +
131 + return out_asm
132 +
133 +
134 +def rocm_aiter_fmoe_fp8_blockscale_g1u1_fake(
135 + topk_ids: torch.Tensor,
136 + topk_weights: torch.Tensor,
137 + hidden_states_dtype: torch.dtype,
138 + expert_mask: torch.Tensor,
139 + a1: torch.Tensor,
140 + w1: torch.Tensor,
141 + w2: torch.Tensor,
142 + w1_scale: torch.Tensor,
143 + w2_scale: torch.Tensor,
144 + a1_scale: torch.Tensor,
145 + block_shape: List[int],
146 + smooth_scale: Optional[torch.Tensor] = None) -> torch.Tensor:
147 +
148 + return torch.empty_like(a1, dtype=torch.bf16)
149 +
150 +
151 +def rocm_aiter_asm_moe_impl(hidden_states: torch.Tensor,
152 + w1: torch.Tensor,
153 + w2: torch.Tensor,
154 + topk_weight: torch.Tensor,
155 + topk_ids: torch.Tensor,
156 + fc1_scale: Optional[torch.Tensor] = None,
157 + fc2_scale: Optional[torch.Tensor] = None,
158 + fc1_smooth_scale: Optional[torch.Tensor] = None,
159 + fc2_smooth_scale: Optional[torch.Tensor] = None,
160 + a16: bool = False,
161 + activation: str = "silu") -> torch.Tensor:
162 import aiter.fused_moe_bf16_asm as rocm_aiter_asm_fmoe
163 + from aiter import ActivationType
164 +
165 + assert activation in ["silu", "gelu"], "The given activation:" \
166 + f" {activation}" \
167 + " is not supported in" \
168 + " AITER."
169 + if activation == "silu":
170 + aiter_activation = ActivationType.Silu
171 + else:
172 + aiter_activation = ActivationType.Gelu
173 +
174 + return rocm_aiter_asm_fmoe.asm_moe(hidden_states=hidden_states,
175 + w1=w1,
176 + w2=w2,
177 + topk_weight=topk_weight,
178 + topk_ids=topk_ids,
179 + fc1_scale=fc1_scale,
180 + fc2_scale=fc2_scale,
181 + fc1_smooth_scale=fc1_smooth_scale,
182 + fc2_smooth_scale=fc2_smooth_scale,
183 + a16=a16,
184 + activation=aiter_activation)
185 +
186 +
187 +def rocm_aiter_asm_moe_fake(hidden_states: torch.Tensor,
188 + w1: torch.Tensor,
189 + w2: torch.Tensor,
190 + topk_weight: torch.Tensor,
191 + topk_ids: torch.Tensor,
192 + fc1_scale: Optional[torch.Tensor] = None,
193 + fc2_scale: Optional[torch.Tensor] = None,
194 + fc1_smooth_scale: Optional[torch.Tensor] = None,
195 + fc2_smooth_scale: Optional[torch.Tensor] = None,
196 + a16: bool = False,
197 + activation: str = "silu") -> torch.Tensor:
198 + return torch.empty_like(hidden_states)
199 +
200 +
201 +def rocm_aiter_topk_softmax_impl(topk_weights: torch.Tensor,
202 + topk_indices: torch.Tensor,
203 + token_expert_indices: torch.Tensor,
204 + gating_output: torch.Tensor,
205 + renormalize: bool) -> None:
206 + from aiter import topk_softmax
207 + topk_softmax(topk_weights, topk_indices, token_expert_indices,
208 + gating_output, renormalize)
209 +
210 +
211 +def rocm_aiter_topk_softmax_fake(topk_weights: torch.Tensor,
212 + topk_indices: torch.Tensor,
213 + token_expert_indices: torch.Tensor,
214 + gating_output: torch.Tensor,
215 + renormalize: bool) -> None:
216 + pass
217 +
218 +
219 +if current_platform.is_rocm():
220 +
221 + direct_register_custom_op(
222 + op_name="rocm_aiter_asm_moe_tkw1",
223 + op_func=rocm_aiter_asm_moe_tkw1_impl,
224 + mutates_args=[],
225 + fake_impl=rocm_aiter_asm_moe_tkw1_fake,
226 + dispatch_key=current_platform.dispatch_key,
227 + )
228 +
229 + direct_register_custom_op(
230 + op_name="rocm_aiter_ck_moe",
231 + op_func=rocm_aiter_ck_moe_impl,
232 + mutates_args=[],
233 + fake_impl=rocm_aiter_ck_moe_fake,
234 + dispatch_key=current_platform.dispatch_key,
235 + )
236 +
237 + direct_register_custom_op(
238 + op_name="rocm_aiter_fmoe_fp8_blockscale_g1u1",
239 + op_func=rocm_aiter_fmoe_fp8_blockscale_g1u1_impl,
240 + mutates_args=[],
241 + fake_impl=rocm_aiter_fmoe_fp8_blockscale_g1u1_fake,
242 + dispatch_key=current_platform.dispatch_key,
243 + )
244 +
245 + direct_register_custom_op(
246 + op_name="rocm_aiter_asm_moe",
247 + op_func=rocm_aiter_asm_moe_impl,
248 + mutates_args=[],
249 + fake_impl=rocm_aiter_asm_moe_fake,
250 + dispatch_key=current_platform.dispatch_key,
251 + )
252 +
253 + direct_register_custom_op(
254 + op_name="rocm_aiter_topk_softmax",
255 + op_func=rocm_aiter_topk_softmax_impl,
256 + mutates_args=["topk_weights", "topk_indices", "token_expert_indices"],
257 + fake_impl=rocm_aiter_topk_softmax_fake,
258 + dispatch_key=current_platform.dispatch_key,
259 + )
260 +
261 +
262 +def rocm_aiter_fused_experts(hidden_states: torch.Tensor,
263 + w1: torch.Tensor,
264 + w2: torch.Tensor,
265 + topk_weights: torch.Tensor,
266 + topk_ids: torch.Tensor,
267 + inplace: bool = False,
268 + activation: str = "silu",
269 + apply_router_weight_on_input: bool = False,
270 + use_fp8_w8a8: bool = False,
271 + use_int8_w8a8: bool = False,
272 + use_int8_w8a16: bool = False,
273 + use_int4_w4a16: bool = False,
274 + per_channel_quant: bool = False,
275 + global_num_experts: int = -1,
276 + expert_map: Optional[torch.Tensor] = None,
277 + w1_scale: Optional[torch.Tensor] = None,
278 + w2_scale: Optional[torch.Tensor] = None,
279 + w1_zp: Optional[torch.Tensor] = None,
280 + w2_zp: Optional[torch.Tensor] = None,
281 + a1_scale: Optional[torch.Tensor] = None,
282 + a2_scale: Optional[torch.Tensor] = None,
283 + block_shape: Optional[List[int]] = None,
284 + allow_deep_gemm: bool = False) -> torch.Tensor:
285
286 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
287 per_token_group_quant_fp8)
@@ -84,60 +290,24 @@ def rocm_aiter_fused_experts(
290 topk_weights = topk_weights.to(torch.float32)
291 topk_ids = topk_ids.to(torch.int32)
292
87 - if (block_shape is not None) and use_fp8_w8a8:
293 + # w8a8 block-scaled
294 + if block_shape is not None and use_fp8_w8a8:
295 assert not apply_router_weight_on_input, (
296 "apply_router_weight_on_input is not supported for block scaled moe"
297 )
91 -
298 assert w1_scale is not None
299 assert w2_scale is not None
300
95 - local_E = E = w1.shape[0]
96 - if expert_map is not None:
97 - E = expert_map.numel()
98 -
99 - topk = topk_ids.shape[1]
100 - model_dim = w1.shape[-1]
101 - dtype = hidden_states.dtype
301 # The default block sizes are 128 in AITER.
103 - if block_shape is None:
104 - block_shape = [128, 128]
105 -
106 - scale_blk_k = block_shape[1]
107 -
108 - (
109 - sorted_token_ids,
110 - sorted_weight_buf,
111 - sorted_expert_ids,
112 - num_valid_ids,
113 - out_asm,
114 - ) = rocm_aiter_asm_fmoe.moe_sorting_ck(topk_ids,
115 - topk_weights,
116 - E,
117 - model_dim,
118 - dtype,
119 - expert_mask=expert_map)
120 -
121 - a1, a1_scale = per_token_group_quant_fp8(hidden_states, scale_blk_k)
122 - rocm_aiter.fmoe_fp8_blockscale_g1u1(
123 - out_asm,
124 - a1,
125 - w1,
126 - w2,
127 - sorted_token_ids,
128 - sorted_weight_buf,
129 - sorted_expert_ids,
130 - num_valid_ids,
131 - topk,
132 - w1_scale.view(local_E, -1),
133 - w2_scale.view(local_E, -1),
134 - a1_scale.t().contiguous(),
135 - block_shape[0],
136 - block_shape[1],
137 - None,
138 - )
139 - return out_asm
302 + block_shape = [128, 128] if block_shape is None else block_shape
303 +
304 + a1, a1_scale = per_token_group_quant_fp8(hidden_states, block_shape[1])
305
306 + return torch.ops.vllm.rocm_aiter_fmoe_fp8_blockscale_g1u1(
307 + topk_ids, topk_weights, hidden_states.dtype, expert_map, a1, w1,
308 + w2, w1_scale, w2_scale, a1_scale, block_shape, None)
309 +
310 + # w8a8 per-channel quantization
311 elif per_channel_quant and apply_router_weight_on_input and use_fp8_w8a8:
312 # AITER tkw1 kernel for FP8 models with `apply_router_weight_on_input`
313 # This applies topk_weights on the GEMM output of the first FC layer
@@ -148,34 +318,36 @@ def rocm_aiter_fused_experts(
318 "Only support topk=1 when"
319 " `apply_router_weight_on_input` is True")
320
151 - return rocm_aiter_asm_moe_tkw1(hidden_states,
152 - w1,
153 - w2,
154 - topk_weights,
155 - topk_ids,
156 - fc1_scale=w1_scale,
157 - fc2_scale=w2_scale,
158 - fc1_smooth_scale=None,
159 - fc2_smooth_scale=None,
160 - a16=False,
161 - per_tensor_quant_scale=None,
162 - expert_mask=expert_map,
163 - activation_str=activation)
164 -
321 + return torch.ops.vllm.rocm_aiter_asm_moe_tkw1(
322 + hidden_states,
323 + w1,
324 + w2,
325 + topk_weights,
326 + topk_ids,
327 + fc1_scale=w1_scale,
328 + fc2_scale=w2_scale,
329 + fc1_smooth_scale=None,
330 + fc2_smooth_scale=None,
331 + a16=False,
332 + per_tensor_quant_scale=None,
333 + expert_mask=expert_map,
334 + activation_str=activation)
335 +
336 + # w8a8 per-tensor activation per-tensor weight
337 elif use_fp8_w8a8:
338 assert not apply_router_weight_on_input, (
339 "apply_router_weight_on_input is not supported for fp8_w8a8")
168 - return rocm_aiter_asm_fmoe.asm_moe(hidden_states=hidden_states,
169 - w1=w1,
170 - w2=w2,
171 - topk_weight=topk_weights,
172 - topk_ids=topk_ids,
173 - fc1_scale=w1_scale,
174 - fc2_scale=w2_scale,
175 - fc1_smooth_scale=None,
176 - fc2_smooth_scale=None,
177 - a16=False)
178 -
340 + return torch.ops.vllm.rocm_aiter_asm_moe(hidden_states=hidden_states,
341 + w1=w1,
342 + w2=w2,
343 + topk_weight=topk_weights,
344 + topk_ids=topk_ids,
345 + fc1_scale=w1_scale,
346 + fc2_scale=w2_scale,
347 + fc1_smooth_scale=None,
348 + fc2_smooth_scale=None,
349 + a16=False,
350 + activation=activation)
351 if apply_router_weight_on_input:
352 assert (topk_weights.dim() == 2
353 ), "`topk_weights` should be in shape (num_tokens, topk)"
@@ -188,26 +360,26 @@ def rocm_aiter_fused_experts(
360 topk_ids = topk_ids.to(torch.int32)
361 topk_weights = torch.ones_like(topk_weights, dtype=torch.float32)
362
191 - return rocm_aiter.ck_moe(hidden_states=hidden_states,
192 - w1=w1,
193 - w2=w2,
194 - topk_weights=topk_weights,
195 - topk_ids=topk_ids)
363 + # w16a16 fallback to rocm_aiter_ck_moe w16a16
364 + return torch.ops.vllm.rocm_aiter_ck_moe(hidden_states=hidden_states,
365 + w1=w1,
366 + w2=w2,
367 + topk_weights=topk_weights,
368 + topk_ids=topk_ids)
369
370
371 def rocm_aiter_topk_softmax(topk_weights: torch.Tensor,
372 topk_indices: torch.Tensor,
373 token_expert_indices: torch.Tensor,
374 gating_output: torch.Tensor,
202 - renormalize: bool) -> tuple[torch.Tensor, ...]:
203 - import aiter as rocm_aiter
204 - rocm_aiter.topk_softmax(topk_weights, topk_indices, token_expert_indices,
205 - gating_output, renormalize)
206 -
375 + renormalize: bool) -> Tuple[torch.Tensor, ...]:
376 + torch.ops.vllm.rocm_aiter_topk_softmax(topk_weights, topk_indices,
377 + token_expert_indices, gating_output,
378 + renormalize)
379 return topk_weights, topk_indices
380
381
210 -def shuffle_weights(*tensors: torch.Tensor) -> tuple[torch.Tensor, ...]:
382 +def shuffle_weights(*tensors: torch.Tensor) -> Tuple[torch.Tensor, ...]:
383 """
384 Applies shuffle_weight function from AITER to each
385 input tensor and returns them.
@@ -216,15 +388,14 @@ def shuffle_weights(*tensors: torch.Tensor) -> tuple[torch.Tensor, ...]:
388 *tensors: Variable number of torch.Tensor objects.
389
390 Returns:
219 - A tuple of shuffled tensors.
391 + A Tuple of shuffled tensors.
392 """
393 from aiter.ops.shuffle import shuffle_weight
222 -
394 return tuple(shuffle_weight(tensor) for tensor in tensors)
395
396
397 def expand_weights(*tensors: torch.Tensor,
227 - expansion_dims: list[int]) -> tuple[torch.Tensor, ...]:
398 + expansion_dims: list[int]) -> Tuple[torch.Tensor, ...]:
399 """
400 Expands the dimensions of input tensors.
401
@@ -234,7 +405,7 @@ def expand_weights(*tensors: torch.Tensor,
405 corresponding to each tensor.
406
407 Returns:
237 - A tuple of tensors with expanded dimensions.
408 + A Tuple of tensors with expanded dimensions.
409 """
410
411 assert len(tensors) == len(expansion_dims), \
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+3 -3
@@ -304,9 +304,9 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
304 e_score_correction_bias=e_score_correction_bias)
305
306 return self.fused_experts_func(
307 - x,
308 - layer.w13_weight,
309 - layer.w2_weight,
307 + hidden_states=x,
308 + w1=layer.w13_weight,
309 + w2=layer.w2_weight,
310 topk_weights=topk_weights,
311 topk_ids=topk_ids,
312 inplace=True,