@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 95f4925c3a

[ROCm][CI] Enable AITER FP8/unquantized cases in modular-kernel sweep + fix MoRI per-tensor FP8 dispatch (#57074)

Signed-off-by: Divakar Verma <divakar.verma@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>

Divakar Verma committed Sep 16, 2026 at 20:53 UTC 95f4925c3a03df8cfcaa21633ccc9dd5b426c7a4

6 files changed +107 -69

.buildkite/test-amd.yaml

+2 -1

@@ -1704,8 +1704,9 @@ steps:

1704 - vllm/_aiter_ops.py

1705 - vllm/platforms/rocm.py

1706 commands:

1707 - - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT

1707 + - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --ignore=kernels/moe/test_modular_kernel_combinations.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT

1708 - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT

1709 + - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s kernels/moe/test_modular_kernel_combinations.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT

1710

1711 - label: ":amd: (MI300) Quantization Kernels" # TBD

1712 timeout_in_minutes: 55

tests/kernels/moe/modular_kernel_tools/common.py

+32 -24

@@ -15,6 +15,7 @@ from tests.kernels.quantization.nvfp4_utils import (

15 dequantize_nvfp4_to_dtype,

16 )

17 from tests.kernels.utils import torch_experts

18 +from vllm._aiter_ops import rocm_aiter_ops

19 from vllm.config import VllmConfig

20 from vllm.distributed import (

21 get_dp_group,

@@ -33,7 +34,9 @@ from vllm.model_executor.layers.fused_moe.config import (

34 FusedMoEQuantConfig,

35 RoutingMethodType,

36 )

37 +from vllm.model_executor.layers.quantization.utils.fp8_utils import is_fp8

38 from vllm.model_executor.layers.quantization.utils.quant_utils import (

39 + QuantKey,

40 kFp8Dynamic128Sym,

41 kFp8DynamicTensorSym,

42 kFp8DynamicTokenSym,

@@ -165,31 +168,34 @@ class Config:

168

169 return vllm_config, env_dict

170

171 + def fp8_quant_key_pair(self) -> tuple[QuantKey, QuantKey]:

172 + """Derive the (weight_quant_key, activation_quant_key) pair an FP8

173 + quant config of this shape corresponds to (either OCP or FNUZ FP8,

174 + see ``current_platform.fp8_dtype()``)."""

175 + if self.quant_block_shape is not None:

176 + return kFp8Static128BlockSym, kFp8Dynamic128Sym

177 + if self.is_per_out_ch_quant:

178 + return (

179 + kFp8StaticChannelSym,

180 + kFp8DynamicTokenSym

181 + if self.is_per_act_token_quant

182 + else kFp8StaticTensorSym,

183 + )

184 + return (

185 + kFp8StaticTensorSym,

186 + kFp8DynamicTensorSym

187 + if self.is_per_act_token_quant

188 + else kFp8StaticTensorSym,

189 + )

190 +

191 def fe_supports_quant_scheme(self) -> bool:

192 """Check if the fused experts class supports this quant config.

193 See https://github.com/ROCm/aiter/issues/2419 for AITER gaps."""

194 if self.quant_config is None or self.quant_dtype is None:

195 return True

173 - if self.quant_dtype != torch.float8_e4m3fn:

196 + if not is_fp8(self.quant_dtype):

197 return True

175 - # Derive QuantKeys from test config

176 - if self.quant_block_shape is not None:

177 - w_key = kFp8Static128BlockSym

178 - a_key = kFp8Dynamic128Sym

179 - elif self.is_per_out_ch_quant:

180 - w_key = kFp8StaticChannelSym

181 - a_key = (

182 - kFp8DynamicTokenSym

183 - if self.is_per_act_token_quant

184 - else kFp8StaticTensorSym

185 - )

186 - else:

187 - w_key = kFp8StaticTensorSym

188 - a_key = (

189 - kFp8DynamicTensorSym

190 - if self.is_per_act_token_quant

191 - else kFp8StaticTensorSym

192 - )

198 + w_key, a_key = self.fp8_quant_key_pair()

199 fe_cls = self.fused_experts_type

200 if hasattr(fe_cls, "_supports_quant_scheme"):

201 try:

@@ -199,10 +205,7 @@ class Config:

205 return True

206

207 def is_fp8_block_quantized(self):

202 - return (

203 - self.quant_dtype == torch.float8_e4m3fn

204 - and self.quant_block_shape is not None

205 - )

208 + return is_fp8(self.quant_dtype) and self.quant_block_shape is not None

209

210 def is_batched_prepare_finalize(self):

211 info = prepare_finalize_info(self.prepare_finalize_type)

@@ -334,6 +337,11 @@ class Config:

337 return False, "Needs Aiter, but Aiter not available."

338 if self.needs_mori() and not has_mori(): # noqa: SIM103

339 return False, "Needs MoRI, but MoRI not available."

340 + if self.needs_mori() and not rocm_aiter_ops.is_fused_moe_enabled():

341 + return False, (

342 + "Mori requires AITER's fused-moe backend to be enabled "

343 + "(VLLM_ROCM_USE_AITER=1 and VLLM_ROCM_USE_AITER_MOE=1)."

344 + )

345

346 try:

347 if not self.fused_experts_type._supports_current_device():

@@ -370,7 +378,7 @@ class WeightTensors:

378 def is_quantized(self) -> bool:

379 # or w1_scale is not None?

380 return (

373 - self.w1.dtype == torch.float8_e4m3fn

381 + is_fp8(self.w1.dtype)

382 or self.w1.dtype == torch.uint8

383 or self.w1.dtype == torch.int8

384 )

tests/kernels/moe/modular_kernel_tools/mk_objects.py

+10 -8

@@ -86,14 +86,14 @@ MK_FUSED_EXPERT_TYPES: list[mk.FusedMoEExpertsModular] = []

86 standard_format = mk.FusedMoEActivationFormat.Standard

87 batched_format = mk.FusedMoEActivationFormat.BatchedExperts

88 common_float_types: list[torch.dtype | str] = [

89 - torch.float8_e4m3fn,

89 + current_platform.fp8_dtype(),

90 torch.bfloat16,

91 torch.float16,

92 torch.float32,

93 ]

94 common_float_and_int_types = common_float_types + [torch.int8]

95 nvfp4_types = ["nvfp4"]

96 -fp8_types = [torch.float8_e4m3fn]

96 +fp8_types = [current_platform.fp8_dtype()]

99 def register_prepare_and_finalize(

@@ -324,7 +324,9 @@ if has_aiter():

324 register_experts(

325 AiterExperts,

326 standard_format,

327 - fp8_types,

327 + # AiterExperts also supports the fully-unquantized (None, None)

328 + # scheme (see SUPPORTED_W_A in rocm_aiter_moe.py), not just fp8.

329 + common_float_types,

330 blocked_quantization_support=True,

331 needs_aiter=True,

332 )

@@ -397,35 +399,35 @@ MK_QUANT_CONFIGS: list[TestMoEQuantConfig | None] = [

399 None,

400 # per-channel / per-column weights and per-tensor activations

401 TestMoEQuantConfig(

400 - quant_dtype=torch.float8_e4m3fn,

402 + quant_dtype=current_platform.fp8_dtype(),

403 per_out_ch_quant=True,

404 per_act_token_quant=False,

405 block_shape=None,

406 ),

407 # per-channel / per-column weights and per-token activations

408 TestMoEQuantConfig(

407 - quant_dtype=torch.float8_e4m3fn,

409 + quant_dtype=current_platform.fp8_dtype(),

410 per_out_ch_quant=True,

411 per_act_token_quant=True,

412 block_shape=None,

413 ),

414 # per-tensor weights and per-tensor activations

415 TestMoEQuantConfig(

414 - quant_dtype=torch.float8_e4m3fn,

416 + quant_dtype=current_platform.fp8_dtype(),

417 per_out_ch_quant=False,

418 per_act_token_quant=False,

419 block_shape=None,

420 ),

421 # per-tensor weights and per-token activations

422 TestMoEQuantConfig(

421 - quant_dtype=torch.float8_e4m3fn,

423 + quant_dtype=current_platform.fp8_dtype(),

424 per_out_ch_quant=False,

425 per_act_token_quant=True,

426 block_shape=None,

427 ),

428 # block-quantized weights and 128 block per-token activations

429 TestMoEQuantConfig(

428 - quant_dtype=torch.float8_e4m3fn,

430 + quant_dtype=current_platform.fp8_dtype(),

431 per_out_ch_quant=False,

432 per_act_token_quant=False,

433 block_shape=[128, 128],

tests/kernels/moe/test_modular_kernel_combinations.py

+43 -32

@@ -37,6 +37,7 @@ from .modular_kernel_tools.parallel_utils import (

37 ProcessGroupInfo,

38 parallel_launch_with_config,

39 )

40 +from .utils import check_accuracy

42 has_any_multi_gpu_package = (

43 has_deep_ep() or has_deep_gemm() or has_flashinfer_cutlass_fused_moe()

@@ -47,12 +48,6 @@ meets_multi_gpu_requirements = pytest.mark.skipif(

48 reason="Requires deep_ep or deep_gemm or flashinfer packages",

49 )

50 -if current_platform.is_fp8_fnuz():

51 - pytest.skip(

52 - "Tests in this file require float8_e4m3fn and platform does not support",

53 - allow_module_level=True,

54 - )

55 -

52 def format_result(verbose, msg, ex=None):

53 if ex is not None:

@@ -70,6 +65,28 @@ def format_result(verbose, msg, ex=None):

65 print(".", end="")

68 +def assert_aiter_quant_scheme_case(config: Config) -> None:

69 + """Make the AITER (weight_quant_key, activation_quant_key) pair this

70 + config exercises explicit, instead of AiterExperts being reached only

71 + indirectly through the general quant-config sweep.

72 + See https://github.com/vllm-project/vllm/issues/54966."""

73 + fe_cls = config.fused_experts_type

74 + if getattr(fe_cls, "__name__", "") != "AiterExperts":

75 + return

76 +

77 + if config.quant_config is None:

78 + w_key, a_key = None, None

79 + else:

80 + w_key, a_key = config.fp8_quant_key_pair()

81 +

82 + assert fe_cls._supports_quant_scheme(w_key, a_key), (

83 + f"AITER case (weight_key={w_key}, activation_key={a_key}) reached "

84 + "the modular-kernel harness, but AiterExperts._supports_quant_scheme "

85 + "does not declare it supported."

86 + )

87 + print(f"[AITER case] weight_key={w_key}, activation_key={a_key}")

88 +

89 +

90 def rank_worker(

91 pgi: ProcessGroupInfo,

92 vllm_config: VllmConfig,

@@ -125,6 +142,25 @@ def rank_worker(

142 count -= 1

143 continue

144

145 + # Skip unsupported: AITER x DeepEP-HT/Mori dispatch crashes

146 + # with an illegal memory access at world_size>1 on gfx942.

147 + # https://github.com/vllm-project/vllm/issues/57029

148 + if (

149 + config.world_size > 1

150 + and getattr(config.fused_experts_type, "__name__", "") == "AiterExperts"

151 + and getattr(config.prepare_finalize_type, "__name__", "")

152 + in ("DeepEPHTPrepareAndFinalize", "MoriPrepareAndFinalize")

153 + ):

154 + print(

155 + f"Skipping[{pgi.rank}]: m={m}, topk={topk}"

156 + " (AITER x DeepEP-HT/Mori illegal memory access,"

157 + " https://github.com/vllm-project/vllm/issues/57029)"

158 + )

159 + count -= 1

160 + continue

161 +

162 + assert_aiter_quant_scheme_case(config)

163 +

164 # modular kernel out

165 mk_out = run_modular_kernel(pgi, vllm_config, config, weights, rank_tensors)

166

@@ -152,32 +188,7 @@ def rank_worker(

188 and config.quant_config is not None

189 )

190 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 - )

191 + check_accuracy(ref_out, mk_out, atol=atol, rtol=rtol, percent=0.9)

192 else:

193 torch.testing.assert_close(ref_out, mk_out, atol=atol, rtol=rtol)

194 format_result(verbose, config.describe())

vllm/model_executor/layers/fused_moe/all2all_utils.py

+11 -4

@@ -284,13 +284,20 @@ def maybe_make_prepare_finalize(

284 # Note: We may want to use FP8 dispatch just to reduce

285 # data movement.

286 use_fp8_dispatch = (

287 - quant_config.is_per_act_token or quant_config.is_block_quantized

287 + quant_config.is_per_act_token

288 + or quant_config.is_block_quantized

289 + or quant_config.is_per_tensor

290 )

291 if use_fp8_dispatch:

290 - # For PTPC (per token per channel) quant, scale dim is 1

291 - # For 1x128 quant, scale dim is hidden_dim // 128

292 + # For PTPC (per token per channel) or per-tensor quant,

293 + # scale dim is 1. For 1x128 quant, scale dim is

294 + # hidden_dim // 128

295 quant_dtype = quant_config.quant_dtype

293 - scale_dim = 1 if quant_config.is_per_act_token else moe.hidden_dim // 128

296 + scale_dim = (

297 + 1

298 + if (quant_config.is_per_act_token or quant_config.is_per_tensor)

299 + else moe.hidden_dim // 128

300 + )

301 else:

302 # Unquantized dispatch (e.g. AITER with defer_input_quant):

303 # dispatch raw BF16/FP16 data, no scales needed.

vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py

@@ -82,6 +82,15 @@ class MoriPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):

82 elif quant_config.is_per_act_token:

83 quant_func = get_hip_quant(QuantType.per_Token)

84 a1, scale = quant_func(a1, quant_dtype=current_platform.fp8_dtype())

85 + elif quant_config.is_per_tensor:

86 + quant_func = get_hip_quant(QuantType.per_Tensor)

87 + a1, scale = quant_func(

88 + a1,

89 + scale=quant_config.a1_scale,

90 + quant_dtype=current_platform.fp8_dtype(),

91 + )

92 + # mori expects one scale slot per token; broadcast.

93 + scale = scale.expand(a1.shape[0], 1).contiguous()

95 (

96 dispatch_a1,