[ROCm][MoE] Split AITER CK and Triton MXFP4 W4A16 into separate backends (#50622)
Signed-off-by: afriedri <afriedri@amd.com> Signed-off-by: Andy Friedrich <afriedri@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
afriedri committed
Aug 31, 2026 at 23:14 UTC
dc9114b2011928842a87d64776e3b0ef48224e93
6 files changed
+464
-390
tests/kernels/moe/test_ocp_mx_moe.py
+25
-21
@@ -134,10 +134,17 @@ def test_mxfp4_loading_and_execution_moe(vllm_runner, model_case: ModelCase):
134
assert output
135
136
137
-def swiglu(x, alpha: float = 1.702, beta: float = 1.0, limit: float | None = None):
138
- # Note we add an extra bias of 1 to the linear layer
139
- # Uses chunked layout: first half is gate, second half is up
140
- x_glu, x_linear = torch.chunk(x, 2, dim=-1)
137
+def swiglu(
138
+ x,
139
+ alpha: float = 1.702,
140
+ beta: float = 1.0,
141
+ limit: float | None = None,
142
+ interleaved: bool = False,
143
+):
144
+ if interleaved:
145
+ x_glu, x_linear = x[..., ::2], x[..., 1::2]
146
+ else:
147
+ x_glu, x_linear = torch.chunk(x, 2, dim=-1)
148
if limit is not None:
149
x_glu = x_glu.clamp(max=limit)
150
x_linear = x_linear.clamp(min=-limit, max=limit)
@@ -145,16 +152,6 @@ def swiglu(x, alpha: float = 1.702, beta: float = 1.0, limit: float | None = Non
152
return out_glu * (x_linear + beta)
153
154
148
-def swigluoai(x, alpha: float = 1.702, limit: float = 7.0):
149
- # OAI swiglu uses interleaved layout: gate/up alternating
150
- # See SwigluOAIAndMul in vllm/model_executor/layers/activation.py
151
- gate, up = x[..., ::2], x[..., 1::2]
152
- gate = gate.clamp(max=limit)
153
- up = up.clamp(min=-limit, max=limit)
154
- glu = gate * torch.sigmoid(gate * alpha)
155
- return (up + 1) * glu
156
-
157
-
155
fp4_lookup_table = [0, 0.5, 1, 1.5, 2, 3, 4, 6, -0, -0.5, -1, -1.5, -2, -3, -4, -6]
156
157
@@ -242,12 +239,9 @@ def reference_moe(
239
240
# Apply activation
241
if activation in ("swiglu", "silu"):
245
- if use_interleaved_layout:
246
- # SWIGLUOAI: interleaved gate/up layout
247
- t = swigluoai(t, alpha=alpha, limit=limit)
248
- else:
249
- # Standard swiglu/silu: chunked layout
250
- t = swiglu(t, alpha=alpha, beta=beta, limit=limit)
242
+ t = swiglu(
243
+ t, alpha=alpha, beta=beta, limit=limit, interleaved=use_interleaved_layout
244
+ )
245
elif activation == "relu2":
246
# RELU2_NO_MUL: relu(x)^2
247
t = torch.relu(t)
@@ -1403,6 +1397,14 @@ ROCM_BACKEND_CONFIGS = {
1397
"requires_aiter": True,
1398
"requires_gfx950": True,
1399
},
1400
+ "AITER_TRITON_MXFP4_BF16": {
1401
+ "activation": "SILU",
1402
+ "rtol": 0.3,
1403
+ "percent": 0.95,
1404
+ "requires_aiter": True,
1405
+ "requires_gfx950": False,
1406
+ "interleaved_layout": True,
1407
+ },
1408
"AITER_MXFP4_FP8": {
1409
"activation": "SWIGLUOAI",
1410
"rtol": 0.5,
@@ -1670,7 +1672,9 @@ def test_rocm_mxfp4_moe_oracle(
1672
# Determine activation type and layout
1673
# SWIGLUOAI uses interleaved layout (gate/up alternating)
1674
# SILU uses chunked layout (first half gate, second half up)
1673
- use_interleaved = activation == MoEActivation.SWIGLUOAI
1675
+ use_interleaved = bool(
1676
+ config.get("interleaved_layout", activation == MoEActivation.SWIGLUOAI)
1677
+ )
1678
if activation in [MoEActivation.SWIGLUOAI, MoEActivation.SILU]:
1679
act_name = "swiglu"
1680
else:
tests/models/quantization/test_gpt_oss.py
+2
@@ -31,6 +31,8 @@ from vllm.model_executor.layers.fused_moe.config import (
31
)
32
from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a8_moe import (
33
aiter_triton_kernel_w4a8_moe_forward,
34
+)
35
+from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a16_moe import (
36
aiter_triton_kernel_w4a16_moe_forward,
37
)
38
from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
vllm/config/kernel.py
+3
@@ -137,6 +137,7 @@ MoEBackend = Literal[
137
"humming",
138
"triton_unfused",
139
"aiter",
140
+ "aiter_triton_mxfp4_bf16",
141
"flydsl",
142
"hpc",
143
"emulation",
@@ -260,6 +261,8 @@ class KernelConfig:
261
- "humming": Use Humming Mixed Precision kernels
262
- "triton_unfused": Use Triton unfused MoE kernels
263
- "aiter": Use AMD AITer kernels (ROCm only)
264
+ - "aiter_triton_mxfp4_bf16": Use the AITER Triton MXFP4 W4A16
265
+ (moe_gemm_a16w4) MoE kernel (ROCm gfx942/gfx950/gfx1250)
266
- "flydsl": Use AMD FlyDSL kernels (ROCm only)
267
- "hpc": Use HPC kernels (FP8 and Hopper only)
268
- "emulation": use BF16/FP16 GEMM, dequantizing weights and
vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py
new
+403
@@ -0,0 +1,403 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
7
+from vllm._aiter_ops import rocm_aiter_ops
8
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
9
+from vllm.model_executor.layers.fused_moe.config import (
10
+ FusedMoEConfig,
11
+ FusedMoEParallelConfig,
12
+ FusedMoEQuantConfig,
13
+ RoutingMethodType,
14
+)
15
+from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a8_moe import (
16
+ patch_gating_output,
17
+)
18
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
19
+ QuantKey,
20
+ kMxfp4Static,
21
+)
22
+
23
+__all__ = [
24
+ "AiterW4A16ExpertsMonolithic",
25
+ "aiter_triton_kernel_w4a16_moe_forward",
26
+]
27
+
28
+
29
+def _aiter_raw(t):
30
+ """Unwrap the `triton_kernels.tensor.Tensor` that `_swizzle_mxfp4` returns."""
31
+ if t is None or isinstance(t, torch.Tensor):
32
+ return t
33
+ assert hasattr(t, "storage"), (
34
+ f"expected a triton_kernels wrapped tensor with a .storage attribute, "
35
+ f"got {type(t)}"
36
+ )
37
+ return t.storage.data
38
+
39
+
40
+def _aiter_w4a16_silu_via_a8w4(
41
+ hidden_states: torch.Tensor,
42
+ w1_data,
43
+ w2_data,
44
+ w1_wscale,
45
+ w2_wscale,
46
+ w1_bias,
47
+ w2_bias,
48
+ routing_data,
49
+ gather_idx,
50
+ scatter_idx,
51
+ gammas,
52
+ apply_router_weight_on_input: bool,
53
+ swiglu_limit: float,
54
+ unpadded_N_w1,
55
+ unpadded_K_w1,
56
+ unpadded_N_w2,
57
+ unpadded_K_w2,
58
+) -> torch.Tensor:
59
+ """
60
+ MXFP4 w4a16 MoE with a SILU (concatenated ``[gate | up]``) activation.
61
+ """
62
+ from aiter.ops.triton.fusions.fused_clamp_act_mul import fused_clamp_act_mul
63
+ from aiter.ops.triton.quant import dynamic_mxfp8_quant
64
+
65
+ try:
66
+ from aiter.ops.triton.moe.moe_op_gemm_a8w4 import moe_gemm_a8w4
67
+ except ImportError:
68
+ from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4
69
+
70
+ from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
71
+ should_use_cdna4_mx_scale_swizzle,
72
+ )
73
+
74
+ swz = "CDNA4_SCALE" if should_use_cdna4_mx_scale_swizzle() else None
75
+ quant_dtype = torch.float8_e4m3fn
76
+
77
+ g1_gammas = gammas if apply_router_weight_on_input else None
78
+ g2_gammas = None if apply_router_weight_on_input else gammas
79
+
80
+ hidden_q, a1_scale = dynamic_mxfp8_quant(hidden_states, quant_dtype=quant_dtype)
81
+ raw_gate_up = moe_gemm_a8w4(
82
+ hidden_q,
83
+ w1_data,
84
+ a1_scale,
85
+ w1_wscale,
86
+ None,
87
+ None,
88
+ w1_bias,
89
+ routing_data,
90
+ gather_indx=gather_idx,
91
+ gammas=g1_gammas,
92
+ swizzle_mx_scale=swz,
93
+ out_dtype=torch.bfloat16,
94
+ apply_swiglu=False,
95
+ unpadded_N=unpadded_N_w1,
96
+ unpadded_K=unpadded_K_w1,
97
+ )
98
+ if unpadded_N_w1 is not None:
99
+ raw_gate_up = raw_gate_up[:, :unpadded_N_w1]
100
+
101
+ # convert interleaved -> chunked for kernel
102
+ gate, up = raw_gate_up[:, ::2], raw_gate_up[:, 1::2]
103
+ raw_gate_up = torch.cat((gate, up), dim=1).contiguous()
104
+
105
+ interim_fp8, a2_scale = fused_clamp_act_mul(
106
+ raw_gate_up,
107
+ swiglu_limit=swiglu_limit,
108
+ activation="silu",
109
+ dtype_quant=quant_dtype,
110
+ scale_dtype_fmt="ue8m0",
111
+ quant_block_size=32,
112
+ )
113
+
114
+ out = moe_gemm_a8w4(
115
+ interim_fp8,
116
+ w2_data,
117
+ a2_scale,
118
+ w2_wscale,
119
+ None,
120
+ None,
121
+ w2_bias,
122
+ routing_data,
123
+ scatter_indx=scatter_idx,
124
+ gammas=g2_gammas,
125
+ swizzle_mx_scale=swz,
126
+ unpadded_N=unpadded_N_w2,
127
+ unpadded_K=unpadded_K_w2,
128
+ )
129
+ return out
130
+
131
+
132
+def aiter_triton_kernel_w4a16_moe_forward(
133
+ hidden_states: torch.Tensor,
134
+ w1,
135
+ w2,
136
+ gating_output: torch.Tensor,
137
+ topk: int,
138
+ renormalize: bool,
139
+ activation: MoEActivation = MoEActivation.SWIGLUOAI,
140
+ quant_config: FusedMoEQuantConfig | None = None,
141
+ apply_router_weight_on_input: bool = False,
142
+ global_num_experts: int = -1,
143
+ expert_map: torch.Tensor | None = None,
144
+ unpadded_N_w1=None,
145
+ unpadded_K_w1=None,
146
+ unpadded_N_w2=None,
147
+ unpadded_K_w2=None,
148
+ num_expert_group: int | None = None,
149
+ topk_group: int | None = None,
150
+ e_score_correction_bias: torch.Tensor | None = None,
151
+ routed_scaling_factor: float | None = None,
152
+ score_mode: str | None = None,
153
+):
154
+ assert quant_config is not None and rocm_aiter_ops.is_enabled()
155
+ from vllm.platforms.rocm import on_gfx1250
156
+
157
+ try:
158
+ from aiter.ops.triton.moe.moe_op_gemm_a16w4 import moe_gemm_a16w4
159
+ except ImportError:
160
+ from aiter.ops.triton.moe_op_gemm_a16w4 import moe_gemm_a16w4
161
+ try:
162
+ from aiter.ops.triton.moe.moe_routing import routing as _routing_mod
163
+ except ImportError:
164
+ from aiter.ops.triton.moe_routing import routing as _routing_mod
165
+
166
+ if on_gfx1250():
167
+ _routing_mod.is_tdm_avail = lambda: False
168
+ aiter_routing = _routing_mod.routing
169
+
170
+ # See context in #50859.
171
+ # TODO: Remove once https://github.com/ROCm/aiter/pull/4530 is merged,
172
+ # AITER released, and AITER pin in vLLM increased from 0.1.19.
173
+ gating_output = patch_gating_output(gating_output, global_num_experts)
174
+
175
+ if score_mode is not None:
176
+ use_grouped_topk = num_expert_group is not None and num_expert_group > 1
177
+ routing_data, gather_idx, scatter_idx = aiter_routing(
178
+ gating_output,
179
+ topk,
180
+ score_mode=score_mode,
181
+ bias=e_score_correction_bias,
182
+ renorm=renormalize,
183
+ routed_scaling_factor=(
184
+ routed_scaling_factor if routed_scaling_factor is not None else 1.0
185
+ ),
186
+ use_grouped_topk=use_grouped_topk,
187
+ num_expert_group=num_expert_group,
188
+ topk_group=topk_group,
189
+ )
190
+ else:
191
+ routing_data, gather_idx, scatter_idx = aiter_routing(
192
+ gating_output, topk, sm_first=not renormalize
193
+ )
194
+
195
+ if on_gfx1250():
196
+ gather_src = gather_idx.to(torch.long) // topk
197
+ hidden_states = hidden_states[gather_src]
198
+ gather_idx = None
199
+
200
+ assert quant_config.w1_precision is not None
201
+ assert quant_config.w2_precision is not None
202
+
203
+ w1_data = _aiter_raw(w1)
204
+ w2_data = _aiter_raw(w2)
205
+ w1_wscale = _aiter_raw(quant_config.w1_precision.weight_scale)
206
+ w2_wscale = _aiter_raw(quant_config.w2_precision.weight_scale)
207
+
208
+ gammas = routing_data.gate_scal if routing_data else None
209
+
210
+ swiglu_alpha = (
211
+ quant_config.gemm1_alpha if quant_config.gemm1_alpha is not None else 1.0
212
+ )
213
+ # 7.0 is the SwigluOAI clamp, so only default to it for that activation;
214
+ # a plain-SILU model that configures no limit must not be clamped.
215
+ if quant_config.gemm1_clamp_limit is not None:
216
+ swiglu_limit = quant_config.gemm1_clamp_limit
217
+ elif activation == MoEActivation.SWIGLUOAI:
218
+ swiglu_limit = 7.0
219
+ else:
220
+ swiglu_limit = None
221
+
222
+ # SILU on gfx1250: use the verified a8w4 kernel (dynamic MXFP8); a16w4 faults.
223
+ if activation == MoEActivation.SILU and on_gfx1250():
224
+ return _aiter_w4a16_silu_via_a8w4(
225
+ hidden_states,
226
+ w1_data,
227
+ w2_data,
228
+ w1_wscale,
229
+ w2_wscale,
230
+ quant_config.w1_bias,
231
+ quant_config.w2_bias,
232
+ routing_data,
233
+ gather_idx,
234
+ scatter_idx,
235
+ gammas,
236
+ apply_router_weight_on_input,
237
+ # fused_clamp_act_mul encodes "no clamp" as 0, not None.
238
+ swiglu_limit if swiglu_limit is not None else 0.0,
239
+ unpadded_N_w1,
240
+ unpadded_K_w1,
241
+ unpadded_N_w2,
242
+ unpadded_K_w2,
243
+ )
244
+
245
+ # SILU: silu(gate) * up — same kernel, just no "+1" residual in swiglu.
246
+ swiglu_add_residual = activation != MoEActivation.SILU
247
+
248
+ intermediate = moe_gemm_a16w4(
249
+ hidden_states,
250
+ w1_data,
251
+ None,
252
+ w1_wscale,
253
+ None,
254
+ None,
255
+ quant_config.w1_bias,
256
+ routing_data,
257
+ gather_indx=gather_idx,
258
+ gammas=gammas if apply_router_weight_on_input else None,
259
+ swizzle_mx_scale=None,
260
+ apply_swiglu=True,
261
+ alpha=swiglu_alpha,
262
+ limit=swiglu_limit,
263
+ swiglu_add_residual=swiglu_add_residual,
264
+ unpadded_N=unpadded_N_w1,
265
+ unpadded_K=unpadded_K_w1,
266
+ )
267
+
268
+ out = moe_gemm_a16w4(
269
+ intermediate,
270
+ w2_data,
271
+ None,
272
+ w2_wscale,
273
+ None,
274
+ None,
275
+ quant_config.w2_bias,
276
+ routing_data,
277
+ scatter_indx=scatter_idx,
278
+ gammas=None if apply_router_weight_on_input else gammas,
279
+ swizzle_mx_scale=None,
280
+ unpadded_N=unpadded_N_w2,
281
+ unpadded_K=unpadded_K_w2,
282
+ )
283
+
284
+ return out
285
+
286
+
287
+class AiterW4A16ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
288
+ def __init__(
289
+ self,
290
+ moe_config: FusedMoEConfig,
291
+ quant_config: FusedMoEQuantConfig,
292
+ ):
293
+ super().__init__(moe_config, quant_config)
294
+ self.topk = moe_config.experts_per_token
295
+ self.renormalize = moe_config.routing_method in (
296
+ RoutingMethodType.Renormalize,
297
+ RoutingMethodType.RenormalizeNaive,
298
+ RoutingMethodType.DeepseekV4,
299
+ )
300
+
301
+ @staticmethod
302
+ def activation_format() -> mk.FusedMoEActivationFormat:
303
+ return mk.FusedMoEActivationFormat.Standard
304
+
305
+ @staticmethod
306
+ def _supports_current_device() -> bool:
307
+ if not rocm_aiter_ops.is_enabled():
308
+ return False
309
+ from vllm.platforms.rocm import on_gfx942, on_gfx950, on_gfx1250
310
+
311
+ return on_gfx942() or on_gfx950() or on_gfx1250()
312
+
313
+ @staticmethod
314
+ def _supports_no_act_and_mul() -> bool:
315
+ return False
316
+
317
+ @staticmethod
318
+ def _supports_quant_scheme(
319
+ weight_key: QuantKey | None,
320
+ activation_key: QuantKey | None,
321
+ ) -> bool:
322
+ return (weight_key, activation_key) == (kMxfp4Static, None)
323
+
324
+ @staticmethod
325
+ def _supports_activation(activation: MoEActivation) -> bool:
326
+ return activation in (MoEActivation.SWIGLUOAI, MoEActivation.SILU)
327
+
328
+ @staticmethod
329
+ def _supports_parallel_config(
330
+ moe_parallel_config: FusedMoEParallelConfig,
331
+ ) -> bool:
332
+ return (
333
+ not moe_parallel_config.use_all2all_kernels
334
+ and not moe_parallel_config.enable_eplb
335
+ and moe_parallel_config.dp_size <= 1
336
+ )
337
+
338
+ @staticmethod
339
+ def _supports_routing_method(
340
+ routing_method: RoutingMethodType,
341
+ weight_key: QuantKey | None,
342
+ activation_key: QuantKey | None,
343
+ ) -> bool:
344
+ return routing_method in [
345
+ RoutingMethodType.Renormalize,
346
+ RoutingMethodType.RenormalizeNaive,
347
+ RoutingMethodType.DeepseekV4,
348
+ ]
349
+
350
+ @staticmethod
351
+ def _supports_router_logits_dtype(
352
+ router_logits_dtype: torch.dtype | None,
353
+ routing_method: RoutingMethodType,
354
+ ) -> bool:
355
+ return True
356
+
357
+ @property
358
+ def expects_unquantized_inputs(self) -> bool:
359
+ return True
360
+
361
+ def apply(
362
+ self,
363
+ hidden_states: torch.Tensor,
364
+ w1: torch.Tensor,
365
+ w2: torch.Tensor,
366
+ router_logits: torch.Tensor,
367
+ activation: MoEActivation,
368
+ global_num_experts: int,
369
+ expert_map: torch.Tensor | None,
370
+ a1q_scale: torch.Tensor | None,
371
+ apply_router_weight_on_input: bool,
372
+ num_expert_group: int | None = None,
373
+ e_score_correction_bias: torch.Tensor | None = None,
374
+ routed_scaling_factor: float | None = None,
375
+ topk_group: int | None = None,
376
+ ) -> torch.Tensor:
377
+ score_mode = (
378
+ "sqrtsoftplus"
379
+ if self.moe_config.routing_method == RoutingMethodType.DeepseekV4
380
+ else None
381
+ )
382
+ return aiter_triton_kernel_w4a16_moe_forward(
383
+ hidden_states=hidden_states,
384
+ w1=w1,
385
+ w2=w2,
386
+ gating_output=router_logits,
387
+ topk=self.topk,
388
+ renormalize=self.renormalize,
389
+ activation=activation,
390
+ global_num_experts=global_num_experts,
391
+ expert_map=expert_map,
392
+ quant_config=self.quant_config,
393
+ apply_router_weight_on_input=apply_router_weight_on_input,
394
+ unpadded_N_w1=None,
395
+ unpadded_K_w1=None,
396
+ unpadded_N_w2=None,
397
+ unpadded_K_w2=None,
398
+ num_expert_group=num_expert_group,
399
+ topk_group=topk_group,
400
+ e_score_correction_bias=e_score_correction_bias,
401
+ routed_scaling_factor=routed_scaling_factor,
402
+ score_mode=score_mode,
403
+ )
vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py
-362
@@ -21,9 +21,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
21
22
__all__ = [
23
"AiterW4A8ExpertsMonolithic",
24
- "AiterW4A16ExpertsMonolithic",
24
"aiter_triton_kernel_w4a8_moe_forward",
26
- "aiter_triton_kernel_w4a16_moe_forward",
25
]
26
27
@@ -371,363 +369,3 @@ class AiterW4A8ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
369
unpadded_N_w2=self.moe_config.hidden_dim_unpadded,
370
unpadded_K_w2=self.moe_config.intermediate_size_per_partition_unpadded,
371
)
374
-
375
-
376
-def _aiter_raw(t):
377
- if t is None or isinstance(t, torch.Tensor):
378
- return t
379
- return t.storage.data if hasattr(t, "storage") else t
380
-
381
-
382
-def _aiter_w4a16_silu_via_a8w4(
383
- hidden_states: torch.Tensor,
384
- w1_data,
385
- w2_data,
386
- w1_wscale,
387
- w2_wscale,
388
- w1_bias,
389
- w2_bias,
390
- routing_data,
391
- gather_idx,
392
- scatter_idx,
393
- gammas,
394
- apply_router_weight_on_input: bool,
395
- swiglu_limit: float,
396
- unpadded_N_w1,
397
- unpadded_K_w1,
398
- unpadded_N_w2,
399
- unpadded_K_w2,
400
-) -> torch.Tensor:
401
- """
402
- MXFP4 w4a16 MoE with a SILU (concatenated ``[gate | up]``) activation.
403
- """
404
- from aiter.ops.triton.fusions.fused_clamp_act_mul import fused_clamp_act_mul
405
- from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4
406
- from aiter.ops.triton.quant import dynamic_mxfp8_quant
407
-
408
- from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
409
- should_use_cdna4_mx_scale_swizzle,
410
- )
411
-
412
- swz = "CDNA4_SCALE" if should_use_cdna4_mx_scale_swizzle() else None
413
- quant_dtype = torch.float8_e4m3fn
414
-
415
- g1_gammas = gammas if apply_router_weight_on_input else None
416
- g2_gammas = None if apply_router_weight_on_input else gammas
417
-
418
- hidden_q, a1_scale = dynamic_mxfp8_quant(hidden_states, quant_dtype=quant_dtype)
419
- raw_gate_up = moe_gemm_a8w4(
420
- hidden_q,
421
- w1_data,
422
- a1_scale,
423
- w1_wscale,
424
- None,
425
- None,
426
- w1_bias,
427
- routing_data,
428
- gather_indx=gather_idx,
429
- gammas=g1_gammas,
430
- swizzle_mx_scale=swz,
431
- out_dtype=torch.bfloat16,
432
- apply_swiglu=False,
433
- unpadded_N=unpadded_N_w1,
434
- unpadded_K=unpadded_K_w1,
435
- )
436
- if unpadded_N_w1 is not None:
437
- raw_gate_up = raw_gate_up[:, :unpadded_N_w1]
438
-
439
- interim_fp8, a2_scale = fused_clamp_act_mul(
440
- raw_gate_up,
441
- swiglu_limit=swiglu_limit,
442
- activation="silu",
443
- dtype_quant=quant_dtype,
444
- scale_dtype_fmt="ue8m0",
445
- quant_block_size=32,
446
- )
447
-
448
- out = moe_gemm_a8w4(
449
- interim_fp8,
450
- w2_data,
451
- a2_scale,
452
- w2_wscale,
453
- None,
454
- None,
455
- w2_bias,
456
- routing_data,
457
- scatter_indx=scatter_idx,
458
- gammas=g2_gammas,
459
- swizzle_mx_scale=swz,
460
- unpadded_N=unpadded_N_w2,
461
- unpadded_K=unpadded_K_w2,
462
- )
463
- return out
464
-
465
-
466
-def aiter_triton_kernel_w4a16_moe_forward(
467
- hidden_states: torch.Tensor,
468
- w1,
469
- w2,
470
- gating_output: torch.Tensor,
471
- topk: int,
472
- renormalize: bool,
473
- activation: MoEActivation = MoEActivation.SWIGLUOAI,
474
- quant_config: FusedMoEQuantConfig | None = None,
475
- apply_router_weight_on_input: bool = False,
476
- global_num_experts: int = -1,
477
- expert_map: torch.Tensor | None = None,
478
- unpadded_N_w1=None,
479
- unpadded_K_w1=None,
480
- unpadded_N_w2=None,
481
- unpadded_K_w2=None,
482
- num_expert_group: int | None = None,
483
- topk_group: int | None = None,
484
- e_score_correction_bias: torch.Tensor | None = None,
485
- routed_scaling_factor: float | None = None,
486
- score_mode: str | None = None,
487
-):
488
- assert quant_config is not None and rocm_aiter_ops.is_enabled()
489
- from vllm.platforms.rocm import on_gfx1250
490
-
491
- try:
492
- from aiter.ops.triton.moe.moe_op_gemm_a16w4 import moe_gemm_a16w4
493
- from aiter.ops.triton.moe.moe_routing import routing as _routing_mod
494
- except ImportError:
495
- from aiter.ops.triton.moe.moe_op_gemm_a16w4 import moe_gemm_a16w4
496
- from aiter.ops.triton.moe_routing import routing as _routing_mod
497
-
498
- if on_gfx1250():
499
- _routing_mod.is_tdm_avail = lambda: False
500
- aiter_routing = _routing_mod.routing
501
-
502
- # See context in #50859.
503
- # TODO: Remove once https://github.com/ROCm/aiter/pull/4530 is merged,
504
- # AITER released, and AITER pin in vLLM increased from 0.1.19.
505
- gating_output = patch_gating_output(gating_output, global_num_experts)
506
-
507
- if score_mode is not None:
508
- use_grouped_topk = num_expert_group is not None and num_expert_group > 1
509
- routing_data, gather_idx, scatter_idx = aiter_routing(
510
- gating_output,
511
- topk,
512
- score_mode=score_mode,
513
- bias=e_score_correction_bias,
514
- renorm=renormalize,
515
- routed_scaling_factor=(
516
- routed_scaling_factor if routed_scaling_factor is not None else 1.0
517
- ),
518
- use_grouped_topk=use_grouped_topk,
519
- num_expert_group=num_expert_group,
520
- topk_group=topk_group,
521
- )
522
- else:
523
- routing_data, gather_idx, scatter_idx = aiter_routing(
524
- gating_output, topk, sm_first=not renormalize
525
- )
526
-
527
- if on_gfx1250():
528
- gather_src = gather_idx.to(torch.long) // topk
529
- hidden_states = hidden_states[gather_src]
530
- gather_idx = None
531
-
532
- assert quant_config.w1_precision is not None
533
- assert quant_config.w2_precision is not None
534
-
535
- w1_data = _aiter_raw(w1)
536
- w2_data = _aiter_raw(w2)
537
- w1_wscale = _aiter_raw(quant_config.w1_precision.weight_scale)
538
- w2_wscale = _aiter_raw(quant_config.w2_precision.weight_scale)
539
-
540
- gammas = routing_data.gate_scal if routing_data else None
541
-
542
- swiglu_alpha = (
543
- quant_config.gemm1_alpha if quant_config.gemm1_alpha is not None else 1.0
544
- )
545
- swiglu_limit = (
546
- quant_config.gemm1_clamp_limit
547
- if quant_config.gemm1_clamp_limit is not None
548
- else 7.0
549
- )
550
-
551
- # SILU on gfx1250: use the verified a8w4 kernel (dynamic MXFP8); a16w4 faults.
552
- if activation == MoEActivation.SILU and on_gfx1250():
553
- return _aiter_w4a16_silu_via_a8w4(
554
- hidden_states,
555
- w1_data,
556
- w2_data,
557
- w1_wscale,
558
- w2_wscale,
559
- quant_config.w1_bias,
560
- quant_config.w2_bias,
561
- routing_data,
562
- gather_idx,
563
- scatter_idx,
564
- gammas,
565
- apply_router_weight_on_input,
566
- swiglu_limit,
567
- unpadded_N_w1,
568
- unpadded_K_w1,
569
- unpadded_N_w2,
570
- unpadded_K_w2,
571
- )
572
-
573
- # SILU: silu(gate) * up — same kernel, just no "+1" residual in swiglu.
574
- swiglu_add_residual = activation != MoEActivation.SILU
575
-
576
- intermediate = moe_gemm_a16w4(
577
- hidden_states,
578
- w1_data,
579
- None,
580
- w1_wscale,
581
- None,
582
- None,
583
- quant_config.w1_bias,
584
- routing_data,
585
- gather_indx=gather_idx,
586
- gammas=gammas if apply_router_weight_on_input else None,
587
- swizzle_mx_scale=None,
588
- apply_swiglu=True,
589
- alpha=swiglu_alpha,
590
- limit=swiglu_limit,
591
- swiglu_add_residual=swiglu_add_residual,
592
- unpadded_N=unpadded_N_w1,
593
- unpadded_K=unpadded_K_w1,
594
- )
595
-
596
- out = moe_gemm_a16w4(
597
- intermediate,
598
- w2_data,
599
- None,
600
- w2_wscale,
601
- None,
602
- None,
603
- quant_config.w2_bias,
604
- routing_data,
605
- scatter_indx=scatter_idx,
606
- gammas=None if apply_router_weight_on_input else gammas,
607
- swizzle_mx_scale=None,
608
- unpadded_N=unpadded_N_w2,
609
- unpadded_K=unpadded_K_w2,
610
- )
611
-
612
- return out
613
-
614
-
615
-class AiterW4A16ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
616
- def __init__(
617
- self,
618
- moe_config: FusedMoEConfig,
619
- quant_config: FusedMoEQuantConfig,
620
- ):
621
- super().__init__(moe_config, quant_config)
622
- self.topk = moe_config.experts_per_token
623
- self.renormalize = moe_config.routing_method in (
624
- RoutingMethodType.Renormalize,
625
- RoutingMethodType.RenormalizeNaive,
626
- RoutingMethodType.DeepseekV4,
627
- )
628
-
629
- @staticmethod
630
- def activation_format() -> mk.FusedMoEActivationFormat:
631
- return mk.FusedMoEActivationFormat.Standard
632
-
633
- @staticmethod
634
- def _supports_current_device() -> bool:
635
- if not rocm_aiter_ops.is_enabled():
636
- return False
637
- from vllm.platforms.rocm import on_gfx950, on_gfx1250
638
-
639
- return on_gfx950() or on_gfx1250()
640
-
641
- @staticmethod
642
- def _supports_no_act_and_mul() -> bool:
643
- return False
644
-
645
- @staticmethod
646
- def _supports_quant_scheme(
647
- weight_key: QuantKey | None,
648
- activation_key: QuantKey | None,
649
- ) -> bool:
650
- return (weight_key, activation_key) == (kMxfp4Static, None)
651
-
652
- @staticmethod
653
- def _supports_activation(activation: MoEActivation) -> bool:
654
- return activation in (MoEActivation.SWIGLUOAI, MoEActivation.SILU)
655
-
656
- @staticmethod
657
- def _supports_parallel_config(
658
- moe_parallel_config: FusedMoEParallelConfig,
659
- ) -> bool:
660
- return (
661
- not moe_parallel_config.use_all2all_kernels
662
- and not moe_parallel_config.enable_eplb
663
- and moe_parallel_config.dp_size <= 1
664
- )
665
-
666
- @staticmethod
667
- def _supports_routing_method(
668
- routing_method: RoutingMethodType,
669
- weight_key: QuantKey | None,
670
- activation_key: QuantKey | None,
671
- ) -> bool:
672
- return routing_method in [
673
- RoutingMethodType.Renormalize,
674
- RoutingMethodType.RenormalizeNaive,
675
- RoutingMethodType.DeepseekV4,
676
- ]
677
-
678
- @staticmethod
679
- def _supports_router_logits_dtype(
680
- router_logits_dtype: torch.dtype | None,
681
- routing_method: RoutingMethodType,
682
- ) -> bool:
683
- return True
684
-
685
- @property
686
- def expects_unquantized_inputs(self) -> bool:
687
- return True
688
-
689
- def apply(
690
- self,
691
- hidden_states: torch.Tensor,
692
- w1: torch.Tensor,
693
- w2: torch.Tensor,
694
- router_logits: torch.Tensor,
695
- activation: MoEActivation,
696
- global_num_experts: int,
697
- expert_map: torch.Tensor | None,
698
- a1q_scale: torch.Tensor | None,
699
- apply_router_weight_on_input: bool,
700
- num_expert_group: int | None = None,
701
- e_score_correction_bias: torch.Tensor | None = None,
702
- routed_scaling_factor: float | None = None,
703
- topk_group: int | None = None,
704
- ) -> torch.Tensor:
705
- assert self.moe_config.intermediate_size_per_partition_unpadded is not None
706
- assert self.moe_config.hidden_dim_unpadded is not None
707
- score_mode = (
708
- "sqrtsoftplus"
709
- if self.moe_config.routing_method == RoutingMethodType.DeepseekV4
710
- else None
711
- )
712
- return aiter_triton_kernel_w4a16_moe_forward(
713
- hidden_states=hidden_states,
714
- w1=w1,
715
- w2=w2,
716
- gating_output=router_logits,
717
- topk=self.topk,
718
- renormalize=self.renormalize,
719
- activation=activation,
720
- global_num_experts=global_num_experts,
721
- expert_map=expert_map,
722
- quant_config=self.quant_config,
723
- apply_router_weight_on_input=apply_router_weight_on_input,
724
- unpadded_N_w1=self.moe_config.intermediate_size_per_partition_unpadded * 2,
725
- unpadded_K_w1=self.moe_config.hidden_dim_unpadded,
726
- unpadded_N_w2=self.moe_config.hidden_dim_unpadded,
727
- unpadded_K_w2=self.moe_config.intermediate_size_per_partition_unpadded,
728
- num_expert_group=num_expert_group,
729
- topk_group=topk_group,
730
- e_score_correction_bias=e_score_correction_bias,
731
- routed_scaling_factor=routed_scaling_factor,
732
- score_mode=score_mode,
733
- )
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+31
-7
@@ -116,9 +116,12 @@ class Mxfp4MoeBackend(Enum):
116
BATCHED_MARLIN = "BATCHED_MARLIN"
117
MARLIN = "MARLIN"
118
# ROCm AITER backends
119
- AITER_MXFP4_BF16 = "AITER_MXFP4_BF16" # W4A16: CK kernel
120
- # Keep the legacy name as an alias while the ROCm split backend rename settles.
119
+ AITER_MXFP4_BF16 = "AITER_MXFP4_BF16" # W4A16: CK kernel (gfx950)
120
+ # Legacy alias, resolves to the CK kernel. New code should name either
121
+ # AITER_MXFP4_BF16 (CK) or AITER_TRITON_MXFP4_BF16 (Triton) explicitly.
122
AITER = "AITER_MXFP4_BF16"
123
+ # W4A16: aiter Triton moe_gemm_a16w4 kernel (gfx942/gfx950/gfx1250)
124
+ AITER_TRITON_MXFP4_BF16 = "AITER_TRITON_MXFP4_BF16"
125
AITER_MXFP4_FP8 = "AITER_MXFP4_FP8" # W4A8: triton kernel
126
AITER_MXFP4_MXFP4 = "AITER_MXFP4_MXFP4" # W4A4: CK kernel
127
# Triton
@@ -143,6 +146,8 @@ TRTLLM_BACKENDS = (
146
TRITON_BACKENDS = (
147
Mxfp4MoeBackend.TRITON,
148
Mxfp4MoeBackend.TRITON_UNFUSED,
149
+ # aiter Triton W4A16 shares the triton_kernels weight format
150
+ Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16,
151
)
152
153
B12X_BACKENDS = (
@@ -232,14 +237,18 @@ def backend_to_kernel_cls(
237
return [BatchedMarlinExperts]
238
239
elif backend == Mxfp4MoeBackend.AITER_MXFP4_BF16:
235
- from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a8_moe import (
236
- AiterW4A16ExpertsMonolithic,
237
- )
240
from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
241
AiterExperts,
242
)
243
242
- return [AiterExperts, AiterW4A16ExpertsMonolithic]
244
+ return [AiterExperts]
245
+
246
+ elif backend == Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16:
247
+ from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a16_moe import (
248
+ AiterW4A16ExpertsMonolithic,
249
+ )
250
+
251
+ return [AiterW4A16ExpertsMonolithic]
252
253
elif backend == Mxfp4MoeBackend.AITER_MXFP4_FP8:
254
from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a8_moe import (
@@ -301,9 +310,11 @@ def map_mxfp4_backend(runner_backend: MoEBackend) -> list[Mxfp4MoeBackend]:
310
"marlin": [Mxfp4MoeBackend.MARLIN],
311
"aiter": [
312
Mxfp4MoeBackend.AITER_MXFP4_BF16,
313
+ Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16,
314
Mxfp4MoeBackend.AITER_MXFP4_FP8,
315
Mxfp4MoeBackend.AITER_MXFP4_MXFP4,
316
],
317
+ "aiter_triton_mxfp4_bf16": [Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16],
318
"aiter_mxfp4_fp8": [Mxfp4MoeBackend.AITER_MXFP4_FP8],
319
"aiter_mxfp4_mxfp4": [Mxfp4MoeBackend.AITER_MXFP4_MXFP4],
320
"xpu": [Mxfp4MoeBackend.XPU],
@@ -325,6 +336,7 @@ def _get_priority_backends_for_gpt_oss() -> list[Mxfp4MoeBackend]:
336
Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
337
Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
338
Mxfp4MoeBackend.AITER_MXFP4_BF16,
339
+ Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16,
340
Mxfp4MoeBackend.AITER_MXFP4_FP8,
341
Mxfp4MoeBackend.AITER_MXFP4_MXFP4,
342
Mxfp4MoeBackend.TRITON,
@@ -1647,7 +1659,18 @@ def convert_weight_to_mxfp4_moe_kernel_format(
1659
):
1660
from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
1661
1650
- if mxfp4_backend == Mxfp4MoeBackend.TRITON:
1662
+ if mxfp4_backend == Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16:
1663
+ # AITER moe_gemm_a16w4 needs gate/up interleaved
1664
+ def interleave_gate_up(w: torch.Tensor) -> torch.Tensor:
1665
+ gate, up = w.chunk(2, dim=1)
1666
+ return torch.stack((gate, up), dim=2).reshape(w.shape)
1667
+
1668
+ w13_weight = interleave_gate_up(w13_weight)
1669
+ w13_weight_scale = interleave_gate_up(w13_weight_scale)
1670
+
1671
+ if w13_bias is not None:
1672
+ w13_bias = interleave_gate_up(w13_bias.to(torch.float32))
1673
+ elif mxfp4_backend == Mxfp4MoeBackend.TRITON:
1674
1675
def shuffle_weight(w: torch.Tensor) -> torch.Tensor:
1676
shape = w.shape
@@ -1891,6 +1914,7 @@ def make_mxfp4_moe_quant_config(
1914
Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
1915
Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
1916
Mxfp4MoeBackend.AITER_MXFP4_BF16,
1917
+ Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16,
1918
Mxfp4MoeBackend.CPU,
1919
):
1920
return mxfp4_w4a16_moe_quant_config(