[ROCm][CI] Loosen block-FP8 fused MoE test tolerance for large-K shapes (#48847)
Signed-off-by: Stefan Koncarevic <Stefan.Koncarevic@amd.com> Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
stefankoncarevic committed
Aug 7, 2026 at 10:19 UTC
0de0362ea1c69b93f9ed36126a1b5c94f0ce2f22
2 files changed
+56
-11
tests/kernels/moe/test_block_fp8.py
+46
-10
@@ -11,10 +11,7 @@ from tests.kernels.moe.utils import (
11
make_test_weights,
12
modular_triton_fused_moe,
13
)
14
-from tests.kernels.quant_utils import (
15
- native_per_token_group_quant_fp8,
16
- native_w8a8_block_matmul,
17
-)
14
+from tests.kernels.quant_utils import native_w8a8_block_matmul
15
from vllm.config import VllmConfig, set_current_vllm_config
16
from vllm.model_executor.layers.activation import SiluAndMul
17
from vllm.model_executor.layers.fused_moe import (
@@ -34,6 +31,9 @@ from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import (
31
from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import (
32
TritonOrDeepGemmExperts,
33
)
34
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
35
+ per_token_group_quant_fp8,
36
+)
37
from vllm.platforms import current_platform
38
from vllm.utils.deep_gemm import (
39
get_mk_alignment_for_contiguous_layout,
@@ -105,8 +105,16 @@ TOP_KS = [1, 2, 6]
105
SEEDS = [0]
106
107
108
-def torch_w8a8_block_fp8_moe(a, w1, w2, w1_s, w2_s, topk_weight, topk_ids, block_shape):
109
- """Fused moe with block-wise quantization using native torch."""
108
+def torch_w8a8_block_fp8_moe(
109
+ a, w1, w2, w1_s, w2_s, topk_weight, topk_ids, block_shape, silu_fp32=False
110
+):
111
+ """Fused MoE with block-wise fp8 quantization using native torch.
112
+
113
+ silu_fp32=True computes the intermediate SiLU in fp32 and quantizes it
114
+ directly, matching the modular Helion silu_and_mul_per_block_quant kernel;
115
+ False rounds the SiLU output to bf16 first, matching fused_experts. Each
116
+ kernel is checked against the reference variant matching its precision.
117
+ """
118
B, D = a.shape
119
topk = topk_ids.size(1)
120
a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
@@ -116,7 +124,10 @@ def torch_w8a8_block_fp8_moe(a, w1, w2, w1_s, w2_s, topk_weight, topk_ids, block
124
topk_ids = topk_ids.view(-1)
125
126
_, block_k = block_shape[0], block_shape[1]
119
- a_q, a_s = native_per_token_group_quant_fp8(a, block_k)
127
+ # Quantize with the production per-token-group fp8 kernel (same HIP/CUDA op the
128
+ # kernels use) so the reference is bit-identical here; removes ~0.06-0.10% fp8
129
+ # boundary-flip divergence that otherwise accumulates over K. Matmul stays fp32.
130
+ a_q, a_s = per_token_group_quant_fp8(a, block_k, dtype=current_platform.fp8_dtype())
131
a_q = a_q.to(torch.float32)
132
for i in range(w1.shape[0]):
133
mask = topk_ids == i
@@ -124,8 +135,12 @@ def torch_w8a8_block_fp8_moe(a, w1, w2, w1_s, w2_s, topk_weight, topk_ids, block
135
inter_out = native_w8a8_block_matmul(
136
a_q[mask], w1[i], a_s[mask], w1_s[i], block_shape, output_dtype=a.dtype
137
)
127
- act_out = SiluAndMul().forward_native(inter_out)
128
- act_out_q, act_out_s = native_per_token_group_quant_fp8(act_out, block_k)
138
+ act_out = SiluAndMul().forward_native(
139
+ inter_out.float() if silu_fp32 else inter_out
140
+ )
141
+ act_out_q, act_out_s = per_token_group_quant_fp8(
142
+ act_out, block_k, dtype=current_platform.fp8_dtype()
143
+ )
144
out[mask] = native_w8a8_block_matmul(
145
act_out_q, w2[i], act_out_s, w2_s[i], block_shape, output_dtype=a.dtype
146
)
@@ -206,8 +221,29 @@ def test_w8a8_block_fp8_fused_moe(
221
222
# 0.039 only needed for M >= 8192
223
tol = 0.035 if M < 8192 else 0.039
224
+
225
+ # The modular path fuses SiLU+quant in fp32 (silu_and_mul_per_block_quant),
226
+ # while fused_experts/the reference round SiLU to bf16 first. On large K/N this
227
+ # ~1-ULP gap pushes m_out past the base tol, so validate m_out against an
228
+ # fp32-SiLU reference — keeping the tight base tolerance, no widened override.
229
+ if current_platform.is_rocm() and K >= 4096 and N >= 1024:
230
+ with set_current_vllm_config(vllm_config):
231
+ ref_out_m = torch_w8a8_block_fp8_moe(
232
+ a,
233
+ w1,
234
+ w2,
235
+ quant_config.w1_scale,
236
+ quant_config.w2_scale,
237
+ topk_weights,
238
+ topk_ids,
239
+ block_size,
240
+ silu_fp32=True,
241
+ )
242
+ else:
243
+ ref_out_m = ref_out
244
+
245
torch.testing.assert_close(out, ref_out, atol=tol, rtol=tol)
210
- torch.testing.assert_close(m_out, ref_out, atol=tol, rtol=tol)
246
+ torch.testing.assert_close(m_out, ref_out_m, atol=tol, rtol=tol)
247
248
249
@pytest.mark.parametrize(("M", "N", "K"), MNK_FACTORS_DG)
vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py
+10
-1
@@ -5,7 +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
8
+from vllm.forward_context import get_forward_context, is_forward_context_available
9
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
10
from vllm.model_executor.layers.fused_moe.config import (
11
FusedMoEConfig,
@@ -27,6 +27,9 @@ __all__ = [
27
]
28
29
30
+# TODO(akaratza): Remove
31
+# _get_padding_mask and patch_gating_output after
32
+# https://github.com/ROCm/aiter/pull/4530 is released and consumed by vLLM.
33
def _get_padding_mask() -> torch.Tensor | None:
34
"""
35
Retrieves a boolean mask with non-padding (0) and padding (1) tokens.
@@ -36,7 +39,13 @@ def _get_padding_mask() -> torch.Tensor | None:
39
slot_mapping[num_tokens_unpadded:num_tokens_padded].fill_(-1)
40
41
in gpu_model_runner.py.
42
+
43
+ Direct kernel callers do not have model-runner padding metadata, so no
44
+ mask is needed when there is no forward context.
45
"""
46
+ if not is_forward_context_available():
47
+ return None
48
+
49
forward_context = get_forward_context()
50
51
# model runner v2.