[ROCm][AITER] Hotfix for `memory access fault` errors in AITER triton MOE routing (#50859)
Signed-off-by: Felix Marty <Felix.Marty@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
fxmarty-amd committed
Aug 4, 2026 at 04:07 UTC
385d4c084ef0c52fe0f7baf7dd2dc4bf86fbe548
2 files changed
+269
-1
tests/models/quantization/test_gpt_oss.py
+212
-1
@@ -19,18 +19,38 @@ from dataclasses import dataclass
19
import huggingface_hub
20
import lm_eval
21
import pytest
22
+import torch
23
from packaging import version
24
25
+from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
26
+from vllm.config import get_current_vllm_config
27
+from vllm.forward_context import set_forward_context
28
+from vllm.model_executor.layers.fused_moe.config import (
29
+ mxfp4_w4a8_moe_quant_config,
30
+ mxfp4_w4a16_moe_quant_config,
31
+)
32
+from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a8_moe import (
33
+ aiter_triton_kernel_w4a8_moe_forward,
34
+ aiter_triton_kernel_w4a16_moe_forward,
35
+)
36
+from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
37
+ Mxfp4MoeBackend,
38
+ convert_gpt_oss_weight_to_mxfp4_moe_kernel_format,
39
+)
40
from vllm.platforms import current_platform
41
from vllm.transformers_utils.repo_utils import hf_api
42
+from vllm.utils.torch_utils import set_random_seed
43
44
if current_platform.is_rocm():
28
- from vllm.platforms.rocm import on_gfx950
45
+ from vllm.platforms.rocm import on_gfx950, on_gfx1250
46
else:
47
48
def on_gfx950() -> bool:
49
return False
50
51
+ def on_gfx1250() -> bool:
52
+ return False
53
+
54
55
MODEL_ACCURACIES = {
56
# Full quantization: attention linears and MoE linears
@@ -128,3 +148,194 @@ def test_gpt_oss_attention_quantization(
148
f"Accuracy {measured_accuracy:.4f} is below threshold "
149
f"{expected_accuracy - rtol:.4f} (expected >= {expected_accuracy} - {rtol})"
150
)
151
+
152
+
153
+@pytest.mark.skipif(
154
+ not (on_gfx950() or on_gfx1250()),
155
+ reason="AITER MXFP4 MoE requires gfx950 or gfx1250",
156
+)
157
+@pytest.mark.skipif(not is_aiter_found_and_supported(), reason="aiter is not installed")
158
+@pytest.mark.skipif(torch.accelerator.device_count() == 0, reason="no gpu available")
159
+@pytest.mark.parametrize(
160
+ "mxfp4_backend",
161
+ [Mxfp4MoeBackend.AITER_MXFP4_FP8, Mxfp4MoeBackend.AITER_MXFP4_BF16],
162
+ ids=["w4a8", "w4a16"],
163
+)
164
+@pytest.mark.parametrize("num_experts", [32, 64, 128])
165
+def test_aiter_mxfp4_moe_ignores_padded_rows(
166
+ mxfp4_backend: Mxfp4MoeBackend,
167
+ num_experts: int,
168
+ monkeypatch: pytest.MonkeyPatch,
169
+ dist_init,
170
+) -> None:
171
+ """
172
+ Garbage in cudagraph padding rows must not reach the unpadded outputs.
173
+
174
+ A cudagraph replay of a size-`TOKENS_PADDED` graph driven by
175
+ `TOKENS_UNPADDED` real tokens leaves the trailing padding row containing
176
+ -inf/inf/nan.
177
+ """
178
+ TOKENS_PADDED = 8
179
+ TOKENS_UNPADDED = 7
180
+
181
+ # `rocm_aiter_ops` snapshots the environment at import time, so setting the
182
+ # env var alone is not enough.
183
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
184
+ rocm_aiter_ops.refresh_env_variables()
185
+
186
+ set_random_seed(0)
187
+ device = torch.device("cuda")
188
+ hidden_size = 2048
189
+ intermediate_size = 512
190
+
191
+ num_experts = 32
192
+ topk = 4
193
+ is_w4a8 = mxfp4_backend is Mxfp4MoeBackend.AITER_MXFP4_FP8
194
+
195
+ # TODO: clean this in oracle/mxfp4.
196
+ # `convert_gpt_oss_weight_to_mxfp4_moe_kernel_format` maps
197
+ # AITER_MXFP4_BF16 onto the CK layout (`shuffle_weight_a16w4`), which the
198
+ # triton monolithic cannot consume -- it wants the `_swizzle_mxfp4`
199
+ # PrecisionConfig that the TRITON branch of the same converter produces.
200
+ weight_backend = (
201
+ Mxfp4MoeBackend.AITER_MXFP4_FP8 if is_w4a8 else Mxfp4MoeBackend.TRITON
202
+ )
203
+
204
+ layer = torch.nn.Module()
205
+ layer.w13_weight = torch.randint(
206
+ 0,
207
+ 256,
208
+ (num_experts, 2 * intermediate_size, hidden_size // 2),
209
+ dtype=torch.uint8,
210
+ device=device,
211
+ )
212
+ layer.w2_weight = torch.randint(
213
+ 0,
214
+ 256,
215
+ (num_experts, hidden_size, intermediate_size // 2),
216
+ dtype=torch.uint8,
217
+ device=device,
218
+ )
219
+
220
+ # Keep the exponents near 127 (2**0) so the dequantized weights stay in a sane
221
+ # range and 255 (e8m0 NaN) is never hit.
222
+ layer.w13_weight_scale = torch.randint(
223
+ 124,
224
+ 131,
225
+ (num_experts, 2 * intermediate_size, hidden_size // 32),
226
+ dtype=torch.uint8,
227
+ device=device,
228
+ )
229
+ layer.w2_weight_scale = torch.randint(
230
+ 124,
231
+ 131,
232
+ (num_experts, hidden_size, intermediate_size // 32),
233
+ dtype=torch.uint8,
234
+ device=device,
235
+ )
236
+ layer.w13_bias = torch.randn(
237
+ (num_experts, 2 * intermediate_size), dtype=torch.float32, device=device
238
+ )
239
+ layer.w2_bias = torch.randn(
240
+ (num_experts, hidden_size), dtype=torch.float32, device=device
241
+ )
242
+ # Read off `layer`, not off the arguments, by the AITER_MXFP4_FP8 branch.
243
+ layer.w13_input_scale = torch.full(
244
+ (num_experts,), 0.1, dtype=torch.float32, device=device
245
+ )
246
+ layer.w2_input_scale = torch.full(
247
+ (num_experts,), 0.1, dtype=torch.float32, device=device
248
+ )
249
+
250
+ (
251
+ w13_weight,
252
+ w2_weight,
253
+ w13_scale,
254
+ w2_scale,
255
+ w13_bias,
256
+ w2_bias,
257
+ ) = convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
258
+ weight_backend,
259
+ layer,
260
+ layer.w13_weight,
261
+ layer.w2_weight,
262
+ layer.w13_weight_scale,
263
+ layer.w2_weight_scale,
264
+ w13_bias=layer.w13_bias,
265
+ w2_bias=layer.w2_bias,
266
+ )
267
+
268
+ quant_config_factory = (
269
+ mxfp4_w4a8_moe_quant_config if is_w4a8 else mxfp4_w4a16_moe_quant_config
270
+ )
271
+ quant_config = quant_config_factory(
272
+ w1_scale=w13_scale,
273
+ w2_scale=w2_scale,
274
+ w1_bias=w13_bias,
275
+ w2_bias=w2_bias,
276
+ )
277
+
278
+ hidden_states = torch.randn(
279
+ (TOKENS_PADDED, hidden_size), dtype=torch.bfloat16, device=device
280
+ )
281
+ gating_output = torch.randn(
282
+ (TOKENS_PADDED, num_experts), dtype=torch.bfloat16, device=device
283
+ )
284
+
285
+ hidden_states[TOKENS_UNPADDED:, 0] = float("inf")
286
+ hidden_states[TOKENS_UNPADDED:, 1] = float("-inf")
287
+ hidden_states[TOKENS_UNPADDED:, 2] = float("nan")
288
+
289
+ gating_output[TOKENS_UNPADDED:, :] = float("-inf")
290
+ gating_output[TOKENS_UNPADDED:, 0] = float("inf")
291
+ gating_output[TOKENS_UNPADDED:, 1] = float("nan")
292
+ gating_output[TOKENS_UNPADDED:, 2] = -float("nan")
293
+
294
+ is_padding = torch.zeros(TOKENS_PADDED, dtype=torch.bool, device=device)
295
+ is_padding[TOKENS_UNPADDED:] = True
296
+
297
+ # Poison memory: after `del blocks`, further `torch.empty` in routing with wrongful
298
+ # expert ids may read into these memory sections resulting in
299
+ # potential memory access fault.
300
+ blocks = []
301
+ for numel in (16, 32, 64, 128, 256, 512, 1024, 4096, 1 << 14, 1 << 16, 1 << 20):
302
+ blocks.append(torch.full((numel,), 100e7, dtype=torch.int32, device=device))
303
+ # return blocks
304
+ del blocks
305
+ torch.accelerator.synchronize()
306
+
307
+ with set_forward_context(None, get_current_vllm_config(), is_padding=is_padding):
308
+ if is_w4a8:
309
+ # AITER_MXFP4_FP8 routes through the triton monolithic expert, which
310
+ # does its own top-k inside `aiter.ops.triton.moe.moe_routing`.
311
+ output = aiter_triton_kernel_w4a8_moe_forward(
312
+ hidden_states=hidden_states,
313
+ w1=w13_weight,
314
+ w2=w2_weight,
315
+ gating_output=gating_output,
316
+ topk=topk,
317
+ renormalize=True,
318
+ quant_config=quant_config,
319
+ global_num_experts=num_experts,
320
+ unpadded_N_w1=2 * intermediate_size,
321
+ unpadded_K_w1=hidden_size,
322
+ unpadded_N_w2=hidden_size,
323
+ unpadded_K_w2=intermediate_size,
324
+ )
325
+ else:
326
+ output = aiter_triton_kernel_w4a16_moe_forward(
327
+ hidden_states=hidden_states,
328
+ w1=w13_weight,
329
+ w2=w2_weight,
330
+ gating_output=gating_output,
331
+ topk=topk,
332
+ renormalize=True,
333
+ quant_config=quant_config,
334
+ global_num_experts=num_experts,
335
+ unpadded_N_w1=2 * intermediate_size,
336
+ unpadded_K_w1=hidden_size,
337
+ unpadded_N_w2=hidden_size,
338
+ unpadded_K_w2=intermediate_size,
339
+ )
340
+
341
+ assert torch.isfinite(output[:TOKENS_UNPADDED]).all()
vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py
+57
@@ -5,6 +5,7 @@ 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.forward_context import get_forward_context
9
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
10
from vllm.model_executor.layers.fused_moe.config import (
11
FusedMoEConfig,
@@ -26,6 +27,52 @@ __all__ = [
27
]
28
29
30
+def _get_padding_mask() -> torch.Tensor | None:
31
+ """
32
+ Retrieves a boolean mask with non-padding (0) and padding (1) tokens.
33
+
34
+ `slot_mapping < 0` comes from:
35
+
36
+ slot_mapping[num_tokens_unpadded:num_tokens_padded].fill_(-1)
37
+
38
+ in gpu_model_runner.py.
39
+ """
40
+ forward_context = get_forward_context()
41
+
42
+ # model runner v2.
43
+ if forward_context.is_padding is not None:
44
+ return forward_context.is_padding
45
+
46
+ # model runner v1.
47
+ slot_mapping = forward_context.slot_mapping
48
+
49
+ if isinstance(slot_mapping, list):
50
+ slot_mapping_dict = slot_mapping[0]
51
+ else:
52
+ slot_mapping_dict = slot_mapping
53
+
54
+ if isinstance(slot_mapping_dict, dict):
55
+ slot_mapping_sample = next(iter(slot_mapping_dict.values()), None)
56
+
57
+ if isinstance(slot_mapping_sample, torch.Tensor):
58
+ return slot_mapping_sample < 0
59
+ else:
60
+ return None
61
+
62
+
63
+def patch_gating_output(
64
+ gating_output: torch.Tensor, global_num_experts: int
65
+) -> torch.Tensor:
66
+ if global_num_experts != 128:
67
+ is_padding = _get_padding_mask()
68
+
69
+ if is_padding is not None:
70
+ assert is_padding.shape[0] == gating_output.shape[0]
71
+ gating_output = gating_output.masked_fill(is_padding[:, None], 0.0)
72
+
73
+ return gating_output
74
+
75
+
76
def aiter_triton_kernel_w4a8_moe_forward(
77
hidden_states: torch.Tensor,
78
w1, # Tensor or triton_kernels.Tensor
@@ -59,6 +106,11 @@ def aiter_triton_kernel_w4a8_moe_forward(
106
_routing_mod.is_tdm_avail = lambda: False
107
aiter_routing = _routing_mod.routing
108
109
+ # See context in #50859.
110
+ # TODO: Remove once https://github.com/ROCm/aiter/pull/4530 is merged,
111
+ # AITER released, and AITER pin in vLLM increased from 0.1.19.
112
+ gating_output = patch_gating_output(gating_output, global_num_experts)
113
+
114
routing_data, gather_idx, scatter_idx = aiter_routing(
115
gating_output, topk, sm_first=not renormalize
116
)
@@ -438,6 +490,11 @@ def aiter_triton_kernel_w4a16_moe_forward(
490
_routing_mod.is_tdm_avail = lambda: False
491
aiter_routing = _routing_mod.routing
492
493
+ # See context in #50859.
494
+ # TODO: Remove once https://github.com/ROCm/aiter/pull/4530 is merged,
495
+ # AITER released, and AITER pin in vLLM increased from 0.1.19.
496
+ gating_output = patch_gating_output(gating_output, global_num_experts)
497
+
498
if score_mode is not None:
499
use_grouped_topk = num_expert_group is not None and num_expert_group > 1
500
routing_data, gather_idx, scatter_idx = aiter_routing(