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 6, 00:45 UTC
[ROCm] Fix MoE kernel test failures on gfx950 (#37833)
Signed-off-by: Andreas Karatzas <akaratza@amd.com> Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com> Co-authored-by: Woosuk Kwon <woosuk.kwon@berkeley.edu> Co-authored-by: Matthew Wong <Matthew.Wong2@amd.com>
Andreas Karatzas committed
Mar 25, 2026 at 13:46 UTC
7d6917bef552d6aff70142ab9fb8af648081d4db
12 files changed
+480
-88
tests/kernels/moe/modular_kernel_tools/common.py
+50
@@ -32,6 +32,14 @@ from vllm.model_executor.layers.fused_moe.config import (
32
FusedMoEQuantConfig,
33
RoutingMethodType,
34
)
35
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
36
+ kFp8Dynamic128Sym,
37
+ kFp8DynamicTensorSym,
38
+ kFp8DynamicTokenSym,
39
+ kFp8Static128BlockSym,
40
+ kFp8StaticChannelSym,
41
+ kFp8StaticTensorSym,
42
+)
43
from vllm.utils.import_utils import (
44
has_aiter,
45
has_deep_ep,
@@ -152,6 +160,39 @@ class Config:
160
161
return vllm_config, env_dict
162
163
+ def fe_supports_quant_scheme(self) -> bool:
164
+ """Check if the fused experts class supports this quant config.
165
+ See https://github.com/ROCm/aiter/issues/2419 for AITER gaps."""
166
+ if self.quant_config is None or self.quant_dtype is None:
167
+ return True
168
+ if self.quant_dtype != torch.float8_e4m3fn:
169
+ return True
170
+ # Derive QuantKeys from test config
171
+ if self.quant_block_shape is not None:
172
+ w_key = kFp8Static128BlockSym
173
+ a_key = kFp8Dynamic128Sym
174
+ elif self.is_per_out_ch_quant:
175
+ w_key = kFp8StaticChannelSym
176
+ a_key = (
177
+ kFp8DynamicTokenSym
178
+ if self.is_per_act_token_quant
179
+ else kFp8StaticTensorSym
180
+ )
181
+ else:
182
+ w_key = kFp8StaticTensorSym
183
+ a_key = (
184
+ kFp8DynamicTensorSym
185
+ if self.is_per_act_token_quant
186
+ else kFp8StaticTensorSym
187
+ )
188
+ fe_cls = self.fused_experts_type
189
+ if hasattr(fe_cls, "_supports_quant_scheme"):
190
+ try:
191
+ return fe_cls._supports_quant_scheme(w_key, a_key)
192
+ except NotImplementedError:
193
+ pass
194
+ return True
195
+
196
def is_fp8_block_quantized(self):
197
return (
198
self.quant_dtype == torch.float8_e4m3fn
@@ -253,6 +294,15 @@ class Config:
294
f"{self.fe_supported_types()}."
295
)
296
297
+ # Check quant scheme compatibility with fused experts class
298
+ if not self.fe_supports_quant_scheme():
299
+ return False, (
300
+ f"FE {self.fused_experts_type.__name__} does not support "
301
+ f"quant scheme (per_out_ch={self.is_per_out_ch_quant}, "
302
+ f"per_act_token={self.is_per_act_token_quant}, "
303
+ f"block={self.quant_block_shape})"
304
+ )
305
+
306
# Check block quantization support
307
is_block_quantized = self.quant_block_shape is not None
308
if is_block_quantized and self.quant_dtype is None:
tests/kernels/moe/test_gpt_oss_triton_kernels.py
+15
-6
@@ -384,12 +384,21 @@ def test_legacy_routing(
384
logits = gating_output
385
if sm_first:
386
logits = torch.softmax(logits, dim=-1)
387
- sparse_logits = topk_fn(logits, topk, apply_softmax=not sm_first)
388
- topk_ids = sparse_logits.indx.to(torch.long)
389
- topk_weights = sparse_logits.vals
390
- routing_data_ref, gather_indx_ref, scatter_indx_ref = make_routing_data(
391
- topk_ids, topk_weights, num_experts
392
- )
387
+ topk_result = topk_fn(logits, topk, apply_softmax=not sm_first)
388
+ # topk_fn returns SparseMatrix on NVIDIA, plain tuple on ROCm.
389
+ if isinstance(topk_result, tuple):
390
+ topk_weights, topk_ids_raw, bitmatrix = topk_result
391
+ from triton_kernels.routing import routing_from_bitmatrix
392
+
393
+ routing_data_ref, gather_indx_ref, scatter_indx_ref = routing_from_bitmatrix(
394
+ bitmatrix, topk_weights, topk_ids_raw, num_experts, topk
395
+ )
396
+ else:
397
+ topk_ids = topk_result.indx.to(torch.long)
398
+ topk_weights = topk_result.vals
399
+ routing_data_ref, gather_indx_ref, scatter_indx_ref = make_routing_data(
400
+ topk_ids, topk_weights, num_experts
401
+ )
402
403
routing_data, gather_indx, scatter_indx = legacy_routing(
404
gating_output, topk, sm_first=sm_first
tests/kernels/moe/test_modular_kernel_combinations.py
+59
-1
@@ -108,6 +108,23 @@ def rank_worker(
108
# inputs for rank
109
rank_tensors = RankTensors.make(config, pgi)
110
111
+ # Skip unsupported: AITER block-scaled MoE does not
112
+ # support apply_router_weight_on_input (topk=1 path).
113
+ # https://github.com/ROCm/aiter/issues/2418
114
+ if (
115
+ topk == 1
116
+ and config.supports_apply_weight_on_input()
117
+ and getattr(config.fused_experts_type, "__name__", "") == "AiterExperts"
118
+ and config.quant_block_shape is not None
119
+ ):
120
+ print(
121
+ f"Skipping[{pgi.rank}]: m={m}, topk={topk}"
122
+ " (AITER block-scaled + weight-on-input,"
123
+ " https://github.com/ROCm/aiter/issues/2418)"
124
+ )
125
+ count -= 1
126
+ continue
127
+
128
# modular kernel out
129
mk_out = run_modular_kernel(pgi, vllm_config, config, weights, rank_tensors)
130
@@ -121,7 +138,48 @@ def rank_worker(
138
atol = 3e-2
139
rtol = 3e-2
140
124
- torch.testing.assert_close(ref_out, mk_out, atol=atol, rtol=rtol)
141
+ # On ROCm, AITER FP8 fused MoE uses hardware FP8
142
+ # dot-product which can produce slightly larger error
143
+ # than dequant+f32 matmul at FP8 representable-value
144
+ # boundaries. Allow a small percentage of elements to
145
+ # exceed the base tolerance by a bounded margin.
146
+ # https://github.com/ROCm/aiter/issues/2421
147
+ from vllm.platforms import current_platform as _cp
148
+
149
+ is_aiter_fp8 = (
150
+ _cp.is_rocm()
151
+ and getattr(config.fused_experts_type, "__name__", "") == "AiterExperts"
152
+ and config.quant_config is not None
153
+ )
154
+ if is_aiter_fp8:
155
+ diff = (ref_out - mk_out).abs()
156
+ n_total = diff.numel()
157
+ max_diff = diff.max().item()
158
+ n_exceed = int((diff > atol).sum().item())
159
+ pct_exceed = n_exceed / n_total * 100
160
+ # FP8 hw matmul vs f32 reference: up to ~4% of
161
+ # elements may exceed base tolerance, but max
162
+ # error should stay within 3x base tolerance.
163
+ max_pct_allowed = 5.0
164
+ relaxed_atol = atol * 4
165
+ print(
166
+ f"[AITER FP8 precision] "
167
+ f"max_diff={max_diff:.6f}, "
168
+ f"exceed_atol={n_exceed}/{n_total} "
169
+ f"({pct_exceed:.4f}%), "
170
+ f"max_pct_allowed={max_pct_allowed}%, "
171
+ f"relaxed_limit={relaxed_atol}"
172
+ )
173
+ assert pct_exceed <= max_pct_allowed, (
174
+ f"AITER FP8: {pct_exceed:.2f}% elements exceed "
175
+ f"atol={atol} (max allowed {max_pct_allowed}%)"
176
+ )
177
+ assert max_diff <= relaxed_atol, (
178
+ f"AITER FP8: max_diff={max_diff:.6f} exceeds "
179
+ f"relaxed limit {relaxed_atol}"
180
+ )
181
+ else:
182
+ torch.testing.assert_close(ref_out, mk_out, atol=atol, rtol=rtol)
183
format_result(verbose, config.describe())
184
except Exception as ex:
185
format_result(verbose, config.describe(), ex)
tests/kernels/moe/test_routing.py
+97
-6
@@ -1,15 +1,31 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections.abc import Callable
4
+from unittest.mock import patch
5
6
import pytest
7
import torch
8
9
+from vllm._aiter_ops import rocm_aiter_ops
10
from vllm.distributed.eplb.eplb_state import EplbLayerState
11
from vllm.model_executor.layers.fused_moe.router.router_factory import (
12
create_fused_moe_router,
13
)
14
from vllm.model_executor.models.llama4 import Llama4MoE
15
+from vllm.platforms import current_platform
16
+
17
+
18
+def _is_aiter_capable() -> bool:
19
+ """Check if the platform supports AITER (gfx942/gfx950)."""
20
+ if not current_platform.is_rocm():
21
+ return False
22
+ try:
23
+ from vllm.platforms.rocm import _ON_MI3XX
24
+
25
+ return _ON_MI3XX
26
+ except ImportError:
27
+ return False
28
+
29
30
# Test parameters
31
MK_S = [(32, 256), (64, 512)]
@@ -96,6 +112,60 @@ def assert_routing_results_close(
112
)
113
114
115
+def assert_aiter_routing_valid(
116
+ topk_weights: torch.Tensor,
117
+ topk_ids: torch.Tensor,
118
+ top_k: int,
119
+ num_experts: int,
120
+ renormalize: bool,
121
+ routed_scaling_factor: float = 1.0,
122
+):
123
+ """Validate AITER routing outputs are structurally correct.
124
+
125
+ AITER grouped_topk is a fundamentally different implementation from
126
+ the Python baseline (different group selection, scoring internals),
127
+ so numerical comparison is not meaningful. Instead we verify the
128
+ outputs satisfy the routing contract: correct shapes, valid expert
129
+ IDs, non-negative weights, and proper normalization."""
130
+ n_tokens = topk_weights.shape[0]
131
+
132
+ # Shape
133
+ assert topk_weights.shape == (n_tokens, top_k), (
134
+ f"weights shape {topk_weights.shape} != ({n_tokens}, {top_k})"
135
+ )
136
+ assert topk_ids.shape == (n_tokens, top_k), (
137
+ f"ids shape {topk_ids.shape} != ({n_tokens}, {top_k})"
138
+ )
139
+
140
+ # Expert IDs in valid range
141
+ assert (topk_ids >= 0).all() and (topk_ids < num_experts).all(), (
142
+ f"expert IDs out of range [0, {num_experts}): "
143
+ f"min={topk_ids.min().item()}, max={topk_ids.max().item()}"
144
+ )
145
+
146
+ # No duplicate expert IDs per token
147
+ for i in range(n_tokens):
148
+ ids = topk_ids[i]
149
+ assert ids.unique().numel() == top_k, (
150
+ f"token {i}: duplicate expert IDs {ids.tolist()}"
151
+ )
152
+
153
+ # Weights are non-negative
154
+ assert (topk_weights >= 0).all(), "negative routing weights"
155
+
156
+ # If renormalized, weights should sum to ~scaling_factor per token
157
+ # (renormalization to 1.0 happens before scaling)
158
+ if renormalize:
159
+ expected_sum = routed_scaling_factor
160
+ sums = topk_weights.sum(dim=-1)
161
+ torch.testing.assert_close(
162
+ sums,
163
+ torch.full_like(sums, expected_sum),
164
+ atol=1e-3,
165
+ rtol=1e-3,
166
+ )
167
+
168
+
169
def baseline_fused_topk(
170
router_logits: torch.Tensor, top_k: int, renormalize: bool
171
) -> tuple[torch.Tensor, torch.Tensor]:
@@ -400,10 +470,7 @@ def test_grouped_topk(
470
471
hidden_states, router_logits = make_test_data(m, k, global_num_experts)
472
403
- # Get router output
404
- topk_weights, topk_ids = router.select_experts(hidden_states, router_logits)
405
-
406
- # Compute baseline
473
+ # Compute baseline (pure Python implementation)
474
baseline_weights, baseline_ids = baseline_grouped_topk(
475
router_logits,
476
top_k,
@@ -415,8 +482,32 @@ def test_grouped_topk(
482
routed_scaling_factor,
483
)
484
418
- # Compare results
419
- assert_routing_results_close(topk_weights, topk_ids, baseline_weights, baseline_ids)
485
+ # Test 1: Python/Triton path against baseline (exact match)
486
+ with patch(
487
+ "vllm.model_executor.layers.fused_moe.router.grouped_topk_router.rocm_aiter_ops.is_fused_moe_enabled",
488
+ return_value=False,
489
+ ):
490
+ py_weights, py_ids = router.select_experts(hidden_states, router_logits)
491
+ assert_routing_results_close(py_weights, py_ids, baseline_weights, baseline_ids)
492
+
493
+ # Test 2: AITER path — verify outputs are structurally valid.
494
+ # AITER grouped_topk is a different implementation so we can't
495
+ # compare numerically against the Python baseline.
496
+ if _is_aiter_capable():
497
+ # Force-enable AITER for gfx942/gfx950 regardless of env var,
498
+ # so CI always exercises this path on capable hardware.
499
+ with patch.object(rocm_aiter_ops, "_AITER_ENABLED", True):
500
+ aiter_weights, aiter_ids = router.select_experts(
501
+ hidden_states, router_logits
502
+ )
503
+ assert_aiter_routing_valid(
504
+ aiter_weights,
505
+ aiter_ids,
506
+ top_k,
507
+ global_num_experts,
508
+ renormalize,
509
+ routed_scaling_factor,
510
+ )
511
512
513
@pytest.mark.parametrize("m,k", MK_S)
tests/kernels/moe/test_shared_fused_moe_routed_transform.py
+86
-7
@@ -14,6 +14,7 @@ import torch.nn as nn
14
from vllm.config import VllmConfig, set_current_vllm_config
15
from vllm.forward_context import set_forward_context
16
from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE
17
+from vllm.platforms import current_platform
18
from vllm.utils.torch_utils import is_torch_equal_or_newer
19
20
@@ -51,6 +52,60 @@ class SimpleSharedExperts(nn.Module):
52
return self.down(nn.functional.silu(gate) * up)
53
54
55
+def _assert_close(
56
+ actual: torch.Tensor,
57
+ expected: torch.Tensor,
58
+ atol: float,
59
+ rtol: float,
60
+ label: str,
61
+) -> None:
62
+ """assert_close that prints diff diagnostics on both success and failure."""
63
+ actual_nans = int(actual.isnan().sum().item())
64
+ expected_nans = int(expected.isnan().sum().item())
65
+ actual_zeros = int((actual == 0).sum().item())
66
+ expected_zeros = int((expected == 0).sum().item())
67
+ n_total = actual.numel()
68
+
69
+ diff = (actual - expected).abs()
70
+ max_diff = diff.max().item()
71
+ mean_diff = diff.mean().item()
72
+ n_exceed = int((diff > atol).sum().item())
73
+ pct_exceed = n_exceed / n_total * 100
74
+
75
+ print(
76
+ f"[{label}] "
77
+ f"shape={list(actual.shape)}, "
78
+ f"max_diff={max_diff:.6e}, "
79
+ f"mean_diff={mean_diff:.6e}, "
80
+ f"exceed_atol({atol})={n_exceed}/{n_total} ({pct_exceed:.2f}%), "
81
+ f"actual=[{actual.min().item():.4f}, {actual.max().item():.4f}], "
82
+ f"expected=[{expected.min().item():.4f}, {expected.max().item():.4f}], "
83
+ f"nan(actual/expected)={actual_nans}/{expected_nans}, "
84
+ f"zeros(actual/expected)={actual_zeros}/{expected_zeros}"
85
+ )
86
+
87
+ assert actual_nans == 0, (
88
+ f"{label}: actual has {actual_nans}/{n_total} NaN values "
89
+ f"(expected has {expected_nans}). "
90
+ f"This indicates a kernel bug, not a precision issue."
91
+ )
92
+ assert expected_nans == 0, (
93
+ f"{label}: expected has {expected_nans}/{n_total} NaN values. "
94
+ f"This indicates a kernel bug, not a precision issue."
95
+ )
96
+
97
+ torch.testing.assert_close(
98
+ actual,
99
+ expected,
100
+ atol=atol,
101
+ rtol=rtol,
102
+ msg=(
103
+ f"{label}: max_diff={max_diff:.6e}, mean_diff={mean_diff:.6e}, "
104
+ f"exceed_atol({atol})={n_exceed}/{n_total} ({pct_exceed:.2f}%)"
105
+ ),
106
+ )
107
+
108
+
109
@pytest.fixture(autouse=True)
110
def setup_cuda():
111
if not torch.cuda.is_available():
@@ -61,6 +116,9 @@ def setup_cuda():
116
@pytest.mark.parametrize("num_tokens", [1, 32])
117
@pytest.mark.parametrize("hidden_size,latent_size", [(256, 128), (128, 64)])
118
@pytest.mark.parametrize("dtype", [torch.bfloat16])
119
+@pytest.mark.parametrize(
120
+ "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
121
+)
122
@pytest.mark.skipif(
123
is_torch_equal_or_newer("2.10.0"),
124
reason="Test fails with PyTorch 2.10.0 see: https://github.com/vllm-project/vllm/issues/33995",
@@ -70,14 +128,24 @@ def test_routed_input_transform_inside_vs_outside(
128
hidden_size: int,
129
latent_size: int,
130
dtype: torch.dtype,
131
+ use_rocm_aiter: bool,
132
dist_init,
133
workspace_init,
134
+ monkeypatch,
135
):
136
"""Compare SharedFusedMoE with transform inside vs manually applying outside.
137
Method A (inside): SharedFusedMoE with routed_input_transform
138
Method B (outside): Manually transform, then SharedFusedMoE without transform
139
"""
140
+ if current_platform.is_rocm() and use_rocm_aiter:
141
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1" if use_rocm_aiter else "0")
142
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1" if use_rocm_aiter else "0")
143
+ from vllm._aiter_ops import rocm_aiter_ops
144
+
145
+ rocm_aiter_ops.refresh_env_variables()
146
+
147
torch.manual_seed(42)
148
+ torch.cuda.manual_seed(42)
149
150
num_experts = 8
151
top_k = 2
@@ -125,7 +193,13 @@ def test_routed_input_transform_inside_vs_outside(
193
prefix="moe_without_transform",
194
)
195
196
+ # Weights are created via torch.empty (uninitialized).
197
+ # Initialize with seeded random values for reproducibility.
198
with torch.no_grad():
199
+ moe_with_transform.w13_weight.normal_()
200
+ moe_with_transform.w13_weight.div_(10)
201
+ moe_with_transform.w2_weight.normal_()
202
+ moe_with_transform.w2_weight.div_(10)
203
moe_without_transform.w13_weight.copy_(moe_with_transform.w13_weight)
204
moe_without_transform.w2_weight.copy_(moe_with_transform.w2_weight)
205
@@ -139,9 +213,14 @@ def test_routed_input_transform_inside_vs_outside(
213
hidden_states = torch.randn(num_tokens, hidden_size, device="cuda", dtype=dtype)
214
router_logits = torch.randn(num_tokens, num_experts, device="cuda", dtype=dtype)
215
216
+ # Clone inputs so any in-place modification by Method A
217
+ # cannot affect Method B's computation.
218
+ hidden_states_A = hidden_states.clone()
219
+ router_logits_A = router_logits.clone()
220
+
221
with set_forward_context(None, vllm_config, num_tokens=num_tokens):
222
shared_out_A, routed_out_A = moe_with_transform(
144
- hidden_states, router_logits
223
+ hidden_states_A, router_logits_A
224
)
225
226
transformed_hidden = routed_transform(hidden_states)
@@ -149,19 +228,19 @@ def test_routed_input_transform_inside_vs_outside(
228
transformed_hidden, router_logits
229
)
230
152
- torch.testing.assert_close(
231
+ expected_shared_out = shared_experts(hidden_states)
232
+
233
+ _assert_close(
234
routed_out_A,
235
routed_out_B,
236
atol=1e-3,
237
rtol=1e-3,
157
- msg="Routed output should match: transform inside vs outside",
238
+ label="Routed output: transform inside vs outside",
239
)
159
-
160
- expected_shared_out = shared_experts(hidden_states)
161
-
162
- torch.testing.assert_close(
240
+ _assert_close(
241
shared_out_A,
242
expected_shared_out,
243
atol=1e-3,
244
rtol=1e-3,
245
+ label="Shared expert output",
246
)
tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py
+67
-42
@@ -10,18 +10,15 @@ import torch
10
from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
11
persistent_masked_m_silu_mul_quant,
12
)
13
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
14
+ get_fp8_min_max,
15
+)
16
from vllm.platforms import current_platform
17
from vllm.utils.deep_gemm import DeepGemmQuantScaleFMT, has_deep_gemm
18
from vllm.utils.math_utils import cdiv, round_up
19
from vllm.utils.torch_utils import set_random_seed
20
18
-if current_platform.is_fp8_fnuz():
19
- pytest.skip(
20
- "Tests in this file require float8_e4m3fn and platform does not support",
21
- allow_module_level=True,
22
- )
23
-
24
-fp8_dtype = torch.float8_e4m3fn
21
+fp8_dtype = current_platform.fp8_dtype()
22
23
CASES = [
24
(1, 1, 128, fp8_dtype),
@@ -58,22 +55,21 @@ def as_uint8(x) -> torch.Tensor:
55
56
57
def silu(x: torch.Tensor) -> torch.Tensor:
61
- one_f32 = torch.tensor([1.0], device=x.device, dtype=torch.float32)
58
x_f32 = x.to(torch.float32)
63
- act_f32 = x_f32 / (one_f32 + torch.exp(-x_f32))
64
- assert act_f32.dtype == torch.float32
65
- return act_f32.to(torch.bfloat16)
59
+ act_f32 = x_f32 / (1.0 + torch.exp(-x_f32))
60
+ if current_platform.is_cuda():
61
+ # C++ kernel returns bf16
62
+ return act_f32.to(torch.bfloat16)
63
+ # Triton fallback stays in f32
64
+ return act_f32
65
66
67
def do_quant(x: torch.Tensor, group_size: int, ceil_ue8m0: bool):
68
+ fp8_min_val, fp8_max_val = get_fp8_min_max()
69
eps_bf16 = torch.tensor([1e-10], device=x.device, dtype=torch.bfloat16)
70
one_bf16 = torch.tensor([1.0], device=x.device, dtype=torch.bfloat16)
71
- fp8_max_bf16 = torch.tensor(
72
- [torch.finfo(fp8_dtype).max], device=x.device, dtype=torch.bfloat16
73
- )
74
- fp8_min_bf16 = torch.tensor(
75
- [torch.finfo(fp8_dtype).min], device=x.device, dtype=torch.bfloat16
76
- )
71
+ fp8_max_bf16 = torch.tensor([fp8_max_val], device=x.device, dtype=torch.bfloat16)
72
+ fp8_min_bf16 = torch.tensor([fp8_min_val], device=x.device, dtype=torch.bfloat16)
73
fp8_max_inv = one_bf16 / fp8_max_bf16
74
assert fp8_max_inv.dtype == torch.bfloat16
75
@@ -81,22 +77,36 @@ def do_quant(x: torch.Tensor, group_size: int, ceil_ue8m0: bool):
77
num_groups = x.numel() // group_size
78
x_og_shape = x.shape
79
84
- x = x.to(torch.bfloat16)
85
- x = x.view((-1, group_size))
86
- amax = x.abs().amax(dim=1).clamp(min=eps_bf16)
87
- assert amax.dtype == torch.bfloat16
88
- s = amax * fp8_max_inv
89
-
90
- if ceil_ue8m0:
91
- s = torch.exp2(
92
- torch.ceil(torch.log2(s).to(torch.bfloat16)).to(torch.bfloat16)
93
- ).to(torch.bfloat16)
94
-
95
- inv_s = one_bf16 / s
96
- inv_s = inv_s.view((num_groups, 1))
97
- xq = torch.clamp(x * inv_s, min=fp8_min_bf16.item(), max=fp8_max_bf16.item()).to(
98
- fp8_dtype
99
- )
80
+ if current_platform.is_cuda():
81
+ # C++ kernel computes entirely in bf16
82
+ x = x.to(torch.bfloat16)
83
+ x = x.view((-1, group_size))
84
+ amax = x.abs().amax(dim=1).clamp(min=eps_bf16)
85
+ assert amax.dtype == torch.bfloat16
86
+ s = amax * fp8_max_inv
87
+
88
+ if ceil_ue8m0:
89
+ s = torch.exp2(
90
+ torch.ceil(torch.log2(s).to(torch.bfloat16)).to(torch.bfloat16)
91
+ ).to(torch.bfloat16)
92
+
93
+ inv_s = one_bf16 / s
94
+ inv_s = inv_s.view((num_groups, 1))
95
+ xq = torch.clamp(
96
+ x * inv_s, min=fp8_min_bf16.item(), max=fp8_max_bf16.item()
97
+ ).to(fp8_dtype)
98
+ else:
99
+ # Triton fallback computes in f32. Use multiply-by-reciprocal
100
+ # to match Triton's constexpr evaluation of 1.0/fp8_max.
101
+ fp8_min_f, fp8_max_f = get_fp8_min_max()
102
+
103
+ x = x.to(torch.float32).view((-1, group_size))
104
+ amax = x.abs().amax(dim=1).clamp(min=1e-10)
105
+ s = amax * (1.0 / fp8_max_f)
106
+ if ceil_ue8m0:
107
+ s = torch.exp2(torch.ceil(torch.log2(s)))
108
+ inv_s = (1.0 / s).view((num_groups, 1))
109
+ xq = torch.clamp(x * inv_s, min=fp8_min_f, max=fp8_max_f).to(fp8_dtype)
110
111
xq = xq.view(x_og_shape)
112
xs = s.view((-1, xq.size(-1) // group_size))
@@ -112,12 +122,10 @@ def silu_mul_quant(
122
assert gate.dtype == torch.bfloat16
123
assert up.dtype == torch.bfloat16
124
115
- act_bf16 = silu(gate)
116
- assert act_bf16.dtype == torch.bfloat16
125
+ act = silu(gate)
126
127
# act & mul
119
- a_m = act_bf16 * up
120
- assert a_m.dtype == torch.bfloat16
128
+ a_m = act * up
129
130
q, s = do_quant(a_m, group_size, ceil_ue8m0)
131
return q, s
@@ -221,8 +229,12 @@ def test_silu_mul_fp8_quant_deep_gemm(E: int, T: int, H: int, fp8_type: torch.dt
229
scale_fmts = [
230
DeepGemmQuantScaleFMT.FLOAT32,
231
DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0,
224
- DeepGemmQuantScaleFMT.UE8M0,
232
]
233
+ # UE8M0 (int32 packed) scales require the C++ kernel which is
234
+ # not available on ROCm (#ifndef USE_ROCM).
235
+ # https://github.com/ROCm/aiter/issues/2420
236
+ if current_platform.is_cuda():
237
+ scale_fmts.append(DeepGemmQuantScaleFMT.UE8M0)
238
239
# Run the SiLU V2 kernel
240
for scale_fmt in scale_fmts:
@@ -274,10 +286,23 @@ def test_silu_mul_fp8_quant_deep_gemm(E: int, T: int, H: int, fp8_type: torch.dt
286
for e in range(E):
287
nt = tokens_per_expert[e].item()
288
277
- torch.testing.assert_close(
278
- y_q[e, :nt].to(torch.float32),
279
- ref_y_q[e, :nt].to(torch.float32),
280
- )
289
+ if current_platform.is_rocm():
290
+ # On ROCm the Triton fallback kernel uses f32 math
291
+ # intrinsics (tl.exp) that may differ from PyTorch's
292
+ # torch.exp by 1 ULP. At FP8 quantization
293
+ # boundaries this can flip one representable value.
294
+ # Allow 1 FP8 quantum of tolerance.
295
+ torch.testing.assert_close(
296
+ y_q[e, :nt].to(torch.float32),
297
+ ref_y_q[e, :nt].to(torch.float32),
298
+ atol=32.0,
299
+ rtol=0.2,
300
+ )
301
+ else:
302
+ torch.testing.assert_close(
303
+ y_q[e, :nt].to(torch.float32),
304
+ ref_y_q[e, :nt].to(torch.float32),
305
+ )
306
307
if scale_fmt == DeepGemmQuantScaleFMT.UE8M0:
308
G = H // group_size
tests/kernels/moe/test_unquantized_backend_selection.py
+41
-2
@@ -16,7 +16,7 @@ from vllm.platforms import current_platform
16
"platform_method,expected_backend",
17
[
18
("is_cuda", UnquantizedMoeBackend.TRITON), # Default CUDA without FlashInfer
19
- ("is_rocm", UnquantizedMoeBackend.TRITON),
19
+ ("is_rocm", UnquantizedMoeBackend.TRITON), # ROCm without AITER
20
("is_cpu", UnquantizedMoeBackend.CPU),
21
("is_xpu", UnquantizedMoeBackend.XPU),
22
("is_tpu", UnquantizedMoeBackend.TPU),
@@ -27,13 +27,19 @@ from vllm.platforms import current_platform
27
"vllm.model_executor.layers.fused_moe.oracle.unquantized.has_flashinfer",
28
return_value=False,
29
)
30
+@patch(
31
+ "vllm.model_executor.layers.fused_moe.oracle.unquantized.rocm_aiter_ops.is_fused_moe_enabled",
32
+ return_value=False,
33
+)
34
def test_select_default_backend_by_platform(
35
+ mock_aiter_enabled,
36
mock_has_flashinfer,
37
monkeypatch,
38
platform_method,
39
expected_backend,
40
):
36
- """Test backend selection for different platforms."""
41
+ """Test default backend selection per platform with all optional
42
+ accelerators (FlashInfer, AITER) disabled."""
43
with patch(
44
"vllm.model_executor.layers.fused_moe.oracle.unquantized.current_platform"
45
) as mock_platform:
@@ -58,6 +64,39 @@ def test_select_default_backend_by_platform(
64
assert selected_backend == expected_backend
65
66
67
+@patch(
68
+ "vllm.model_executor.layers.fused_moe.oracle.unquantized.has_flashinfer",
69
+ return_value=False,
70
+)
71
+@patch(
72
+ "vllm.model_executor.layers.fused_moe.oracle.unquantized.rocm_aiter_ops.is_fused_moe_enabled",
73
+ return_value=True,
74
+)
75
+@pytest.mark.skipif(
76
+ not current_platform.is_rocm(), reason="ROCm-specific backend selection test"
77
+)
78
+def test_select_rocm_aiter_backend(mock_aiter_enabled, mock_has_flashinfer):
79
+ """Test ROCm backend selection when AITER is available."""
80
+ with patch(
81
+ "vllm.model_executor.layers.fused_moe.oracle.unquantized.current_platform"
82
+ ) as mock_platform:
83
+ mock_platform.is_cuda.return_value = False
84
+ mock_platform.is_rocm.return_value = True
85
+ mock_platform.is_cpu.return_value = False
86
+ mock_platform.is_xpu.return_value = False
87
+ mock_platform.is_tpu.return_value = False
88
+ mock_platform.is_out_of_tree.return_value = False
89
+
90
+ moe_config = make_dummy_moe_config()
91
+ selected_backend = select_unquantized_moe_backend(
92
+ moe_config=moe_config,
93
+ use_ep=False,
94
+ use_dp=False,
95
+ )
96
+
97
+ assert selected_backend == UnquantizedMoeBackend.AITER
98
+
99
+
100
@patch(
101
"vllm.model_executor.layers.fused_moe.oracle.unquantized.has_flashinfer",
102
return_value=True,
tests/kernels/utils.py
+1
-1
@@ -941,7 +941,7 @@ def torch_experts(
941
if b_bias1 is not None:
942
tmp1 = tmp1 + b_bias1[i].view(1, -1).to(out.dtype)
943
944
- tmp2 = SiluAndMul()(tmp1).to(out.dtype)
944
+ tmp2 = act()(tmp1).to(out.dtype)
945
946
tmp2, b_scale = moe_kernel_quantize_input(
947
tmp2, a2_scale, quant_dtype, per_act_token_quant, block_shape
vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py
+13
-8
@@ -19,6 +19,7 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
19
from vllm.model_executor.layers.fused_moe.utils import _resize_cache
20
from vllm.model_executor.layers.quantization.utils.quant_utils import (
21
QuantKey,
22
+ get_fp8_min_max,
23
kFp8Dynamic128Sym,
24
kFp8Static128BlockSym,
25
)
@@ -117,7 +118,10 @@ def _silu_mul_fp8_quant_deep_gemm(
118
gate = gate * (1.0 / (1.0 + tl.exp(-gate)))
119
y = gate * up
120
120
- y_s = tl.maximum(tl.max(tl.abs(y)), eps) / fp8_max
121
+ # Use multiply-by-reciprocal to match PyTorch's tensor/scalar
122
+ # division precision (Triton GPU fast-division for constexpr
123
+ # divisors can introduce 1-ULP error).
124
+ y_s = tl.maximum(tl.max(tl.abs(y)), eps) * (1.0 / fp8_max)
125
if ceil_ue8m0:
126
y_s = tl.exp2(tl.ceil(tl.log2(y_s)))
127
@@ -190,7 +194,7 @@ def persistent_masked_m_silu_mul_quant(
194
195
tokens_per_expert = tokens_per_expert.to(device=y.device, dtype=torch.int32)
196
193
- fp8_dtype = torch.float8_e4m3fn
197
+ fp8_dtype = current_platform.fp8_dtype()
198
y_q = torch.empty((E, T, H), dtype=fp8_dtype, device=y.device)
199
200
ys_shape, ys_strides, ys_dtype = scales_shape_stride_dtype(E, T, G, quant_scale_fmt)
@@ -210,11 +214,14 @@ def persistent_masked_m_silu_mul_quant(
214
device_id=y.device.index
215
).to_int()
216
213
- if cuda_arch >= 80:
217
+ if current_platform.is_cuda() and cuda_arch >= 80:
218
torch.ops._C.persistent_masked_m_silu_mul_quant(
219
y, tokens_per_expert, y_q, y_s, ceil_ue8m0
220
)
221
else:
222
+ # Triton fallback for ROCm -- the C++ kernel is guarded by
223
+ # #ifndef USE_ROCM in activation_kernels.cu.
224
+ # https://github.com/ROCm/aiter/issues/2420
225
stride_cnt_e = tokens_per_expert.stride()[0]
226
227
# Static grid over experts and H-groups.
@@ -224,13 +231,11 @@ def persistent_masked_m_silu_mul_quant(
231
stride_i_e, stride_i_t, stride_i_h = y.stride()
232
stride_yq_e, stride_yq_t, stride_yq_h = y_q.stride()
233
227
- f_info = torch.finfo(fp8_dtype)
228
- fp8_max = f_info.max
229
- fp8_min = f_info.min
234
+ fp8_min, fp8_max = get_fp8_min_max()
235
eps: float = 1e-10
236
assert y_s.dtype == torch.float32, (
232
- "_silu_mul_fp8_quant_deep_gemm does"
233
- "not support {y_s.dtype} scales. Only torch.float32 supported."
237
+ "_silu_mul_fp8_quant_deep_gemm Triton fallback does not "
238
+ f"support {y_s.dtype} scales. Only torch.float32 supported."
239
)
240
_silu_mul_fp8_quant_deep_gemm[grid](
241
y,
vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py
+23
-6
@@ -253,10 +253,16 @@ def triton_kernel_moe_forward(
253
logits = gating_output
254
if sm_first:
255
logits = torch.softmax(logits, dim=-1)
256
- sparse_logits = topk_fn(logits, topk, apply_softmax=not sm_first)
257
- # sparse_logits.indx contains global expert IDs – remap to local.
258
- topk_ids = expert_map[sparse_logits.indx.to(torch.long)]
259
- topk_weights = sparse_logits.vals
256
+ topk_result = topk_fn(logits, topk, apply_softmax=not sm_first)
257
+ # topk may return a tuple (vals, indx, bitmatrix) or a
258
+ # SparseMatrix depending on the triton_kernels version.
259
+ if isinstance(topk_result, tuple):
260
+ topk_weights, topk_ids_raw, _ = topk_result
261
+ else:
262
+ topk_weights = topk_result.vals
263
+ topk_ids_raw = topk_result.indx
264
+ # topk_ids_raw contains global expert IDs - remap to local.
265
+ topk_ids = expert_map[topk_ids_raw.to(torch.long)]
266
local_num_experts = w1.shape[0]
267
routing_data, gather_idx, scatter_idx = make_routing_data(
268
topk_ids, topk_weights, local_num_experts
@@ -422,8 +428,13 @@ def triton_kernel_fused_mxfp4_w4a8_experts(
428
assert quant_config.w1_bias is None or quant_config.w1_bias.dtype == torch.float32
429
assert quant_config.w2_bias is None or quant_config.w2_bias.dtype == torch.float32
430
425
- # Shape check, only check non-mxfp4
426
- assert hidden_states.shape[-1] == w1.shape[-2]
431
+ # Shape check: when weights are padded (e.g. hidden_size padded for
432
+ # GFX950 swizzle), unpadded_K_w1 carries the original dimension.
433
+ expected_K_w1 = unpadded_K_w1 if unpadded_K_w1 is not None else w1.shape[-2]
434
+ assert hidden_states.shape[-1] == expected_K_w1, (
435
+ f"hidden_states K={hidden_states.shape[-1]} != "
436
+ f"expected K={expected_K_w1} (w1 K={w1.shape[-2]})"
437
+ )
438
assert w2.shape[-1] == w1.shape[1]
439
440
E, _, N = w1.shape
@@ -483,6 +494,12 @@ def triton_kernel_fused_mxfp4_w4a8_experts(
494
unpadded_K=unpadded_K_w2,
495
)
496
497
+ # When hidden_size was padded for alignment (e.g. GFX950 swizzle),
498
+ # the kernel output has the padded dimension. Slice back to the
499
+ # original hidden_size so downstream layers see the expected shape.
500
+ if unpadded_N_w2 is not None and intermediate_cache3.shape[-1] != unpadded_N_w2:
501
+ intermediate_cache3 = intermediate_cache3[..., :unpadded_N_w2].contiguous()
502
+
503
return intermediate_cache3
504
505
vllm/model_executor/layers/quantization/quark/quark_moe.py
+15
@@ -741,11 +741,14 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
741
# TP=4 yields intermediate_size_per_partition=384), AITER raises:
742
# "device_gemm ... does not support this GEMM problem".
743
# Fall back to emulation in that case.
744
+ # For gpt_oss models, create_weights rounds up the dimensions
745
+ # internally, so the alignment check is skipped.
746
if (
747
not self.emulate
748
and self.use_rocm_aiter_moe
749
and self.ocp_mx_scheme is not None
750
and self.ocp_mx_scheme.startswith("w_mxfp4")
751
+ and self.model_type != "gpt_oss"
752
and moe.intermediate_size_per_partition % CK_MXFP4_MOE_DIM_ALIGNMENT != 0
753
):
754
logger.warning_once(
@@ -819,6 +822,18 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
822
"unpadded_hidden_size", hidden_size
823
)
824
825
+ # On GFX950, the GFX950MXScaleLayout swizzle requires
826
+ # hidden_size to be a multiple of 256 (SCALE_K = hidden_size / 32
827
+ # must be divisible by 8). Pad hidden_size for weight/scale
828
+ # allocation; the original value is preserved in unpadded_hidden_size.
829
+ # Only applies to the native (non-emulated) CK path on GFX950.
830
+ if (
831
+ self.model_type == "gpt_oss"
832
+ and current_platform.is_rocm()
833
+ and not self.emulate
834
+ ):
835
+ hidden_size = round_up(hidden_size, 256)
836
+
837
# WEIGHTS
838
w13_weight = torch.nn.Parameter(
839
torch.empty(
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+13
-9
@@ -615,8 +615,8 @@ def _per_token_group_quant_fp8(
615
# Avoid to divide zero
616
eps,
617
# Information for float8
618
- fp8_min,
619
- fp8_max,
618
+ fp8_min: tl.constexpr,
619
+ fp8_max: tl.constexpr,
620
use_ue8m0: tl.constexpr,
621
# Meta-parameters
622
BLOCK: tl.constexpr,
@@ -647,8 +647,12 @@ def _per_token_group_quant_fp8(
647
648
y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32)
649
# Quant
650
+ # Use multiply-by-reciprocal instead of division to match PyTorch's
651
+ # tensor/scalar division precision (GPU fast-division for constexpr
652
+ # divisors can introduce 1-ULP error that flips FP8 quantization at
653
+ # representable-value boundaries).
654
_absmax = tl.maximum(tl.max(tl.abs(y)), eps)
651
- scale_raw = _absmax / fp8_max
655
+ scale_raw = _absmax * (1.0 / fp8_max)
656
y_s = tl.math.exp2(tl.ceil(tl.log2(scale_raw))) if use_ue8m0 else scale_raw
657
y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty)
658
@@ -667,8 +671,8 @@ def _silu_mul_per_token_group_quant_fp8_colmajor(
671
y_s_col_stride: tl.int64,
672
# Information for float8
673
eps,
670
- fp8_min,
671
- fp8_max,
674
+ fp8_min: tl.constexpr,
675
+ fp8_max: tl.constexpr,
676
use_ue8m0: tl.constexpr,
677
# Meta-parameters
678
GROUP_SIZE: tl.constexpr,
@@ -709,7 +713,7 @@ def _silu_mul_per_token_group_quant_fp8_colmajor(
713
714
# quant
715
_absmax = tl.maximum(tl.max(tl.abs(y), axis=1), eps)
712
- scale_raw = _absmax / fp8_max
716
+ scale_raw = _absmax * (1.0 / fp8_max)
717
y_s = tl.math.exp2(tl.ceil(tl.log2(scale_raw))) if use_ue8m0 else scale_raw
718
y_s = tl.reshape(y_s, (BLOCK_M, 1))
719
y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty)
@@ -808,8 +812,8 @@ def _per_token_group_quant_fp8_colmajor(
812
# Avoid to divide zero
813
eps,
814
# Information for float8
811
- fp8_min,
812
- fp8_max,
815
+ fp8_min: tl.constexpr,
816
+ fp8_max: tl.constexpr,
817
use_ue8m0: tl.constexpr,
818
# Meta-parameters
819
BLOCK: tl.constexpr,
@@ -849,7 +853,7 @@ def _per_token_group_quant_fp8_colmajor(
853
y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32)
854
# Quant
855
_absmax = tl.maximum(tl.max(tl.abs(y)), eps)
852
- scale_raw = _absmax / fp8_max
856
+ scale_raw = _absmax * (1.0 / fp8_max)
857
y_s = tl.math.exp2(tl.ceil(tl.log2(scale_raw))) if use_ue8m0 else scale_raw
858
y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty)
859