From 6f41a3a607edef63195c7a4f25b3df50bb26d855 Mon Sep 17 00:00:00 2001 From: Divakar Verma Date: Tue, 29 Sep 2026 23:33:32 +0000 Subject: [PATCH 1/7] [ROCm] Test AiterExperts token padding: inf/nan garbage rows Covers step 4c of the "Test padding" ask in vllm-project/vllm#54966: simulates the HIP-graph padding-token contract (grouped_topk_router.py masking padding rows to topk_ids=-1 / topk_weights=0) at the tensor level, without needing real HIP-graph capture. Fills a subset of rows' hidden_states with inf/-inf/nan and masks their routing accordingly, then verifies through AiterExperts that: (1) no NaN/Inf leaks into the real (non-garbage) rows' output, and (2) the real rows' output matches reference_moe_impl computed only over those rows. Parametrized across 3 garbage-row patterns (single inf, single nan, multiple mixed) x 5 quant schemes AiterExperts supports (unquantized + 4 fp8 variants). Adds a skipped MXFP4 stub pending MI350/gfx950 hardware. Builds on the Config.hidden_dim_unpadded / intermediate_size_per_partition_unpadded plumbing and _PADDING_* constants added in the 4b PR (stacked on that branch since this reuses those constants directly). Gated via require_aiter_moe / has_aiter(); inert when AITER isn't available or off ROCm. Signed-off-by: Divakar Verma --- .../moe/test_modular_kernel_combinations.py | 150 ++++++++++++++++++ 1 file changed, 150 insertions(+) diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py index 4b769931ea04..6f9963d881f6 100644 --- a/tests/kernels/moe/test_modular_kernel_combinations.py +++ b/tests/kernels/moe/test_modular_kernel_combinations.py @@ -988,6 +988,156 @@ def test_aiter_moe_padding_matrix_mxfp4(): padding matrix once MI350/gfx950 hardware is available in CI.""" +# --- 4c: HIP-graph token padding -- inf/nan garbage rows (topk_ids=-1) ------ + +# Simulates the runtime's padding-token contract for garbage/padding rows +# (grouped_topk_router.py: `topk_ids.masked_fill(is_padding, -1)`, +# `topk_weights.masked_fill(is_padding, 0)`), without needing real HIP-graph +# capture. Empirically confirmed (see +# https://github.com/vllm-project/vllm/issues/54966 discussion) that +# AiterExperts safely isolates these rows under both the default ("opus") +# and FlyDSL AITER sorting backends: garbage rows' inf/nan hidden_states +# never leak NaN/Inf into real rows' output. reference_moe_impl's dispatch +# (`mask = topk_ids == i`, tests/kernels/utils.py) is likewise row-local, so +# comparing real rows against it is a valid check. AiterExperts's per-tensor +# a1 quant scale (when applicable) is fixed once in run_modular_kernel()'s +# FusedMoEQuantConfig construction, from hidden_states_scale computed inside +# RankTensors.make() *before* garbage rows are injected, so it cannot be +# inflated by a garbage row's magnitude either. + +_TOKEN_PADDING_M = 16 +_TOKEN_PADDING_K = 1024 +_TOKEN_PADDING_N = 512 + +# {mode name -> {row index: fill value}}. Covers multiple garbage-row counts +# and positions, and both inf and nan. +_TOKEN_PADDING_GARBAGE_MODES: dict[str, dict[int, float]] = { + "single_inf": {3: float("inf")}, + "single_nan": {5: float("nan")}, + "multiple_mixed": { + 1: float("inf"), + 4: float("-inf"), + 7: float("nan"), + 12: float("inf"), + }, +} + + +def _apply_token_padding_garbage_rows( + rank_tensors: RankTensors, garbage_rows: dict[int, float] +) -> None: + """Simulate the HIP-graph padding-token contract in-place: garbage rows + get inf/nan hidden_states and their routing masked to topk_ids=-1 / + topk_weights=0, mirroring grouped_topk_router.py's `is_padding` masking. + """ + for row, fill in garbage_rows.items(): + rank_tensors.hidden_states[row].fill_(fill) + rank_tensors.topk_ids[row].fill_(-1) + rank_tensors.topk_weights[row].fill_(0.0) + + +def _aiter_token_padding_worker( + pgi: ProcessGroupInfo, + vllm_config: VllmConfig, + cpu_group, + config: Config, + weights: WeightTensors, + verbose: bool, + garbage_rows: dict[int, float], +): + device = torch.device(f"cuda:{pgi.local_rank}") + init_workspace_manager(device) + set_random_seed(pgi.rank) + + weights = copy.deepcopy(weights) + weights.to_current_device() + + rank_tensors = RankTensors.make(config, pgi) + _apply_token_padding_garbage_rows(rank_tensors, garbage_rows) + + with set_current_vllm_config(vllm_config): + mk_out = run_modular_kernel(pgi, vllm_config, config, weights, rank_tensors) + ref_out = reference_moe_impl(config, weights, rank_tensors) + + real_rows = [i for i in range(config.Ms) if i not in garbage_rows] + + assert torch.isfinite(mk_out[real_rows]).all(), ( + "Garbage (padding) token rows leaked NaN/Inf into real rows' output " + "through AiterExperts." + ) + + if config.quant_config is not None: + check_accuracy( + ref_out[real_rows], mk_out[real_rows], atol=3e-2, rtol=3e-2, percent=0.9 + ) + else: + torch.testing.assert_close( + ref_out[real_rows], mk_out[real_rows], atol=3e-2, rtol=3e-2 + ) + + +@require_aiter_moe +@pytest.mark.parametrize("quant_config", _PADDING_QUANT_CONFIGS, ids=_PADDING_QUANT_IDS) +@pytest.mark.parametrize("garbage_mode", list(_TOKEN_PADDING_GARBAGE_MODES)) +def test_aiter_moe_token_padding_garbage_rows( + garbage_mode: str, quant_config: TestMoEQuantConfig | None +): + """See https://github.com/vllm-project/vllm/issues/54966 ("Test padding"). + + Simulates HIP-graph token padding: a subset of rows get inf/nan + hidden_states with topk_ids=-1 / topk_weights=0, mirroring the router's + padding-token masking contract (grouped_topk_router.py). Verifies the + real (non-garbage) rows' output through AiterExperts remains finite and + matches reference_moe_impl, across multiple garbage-row counts/positions + and both inf and nan fill values, for unquantized + the 4 fp8 quant + schemes AiterExperts supports. + """ + from vllm.model_executor.layers.fused_moe.prepare_finalize import ( + MoEPrepareAndFinalizeNoDPEPModular, + ) + + garbage_rows = _TOKEN_PADDING_GARBAGE_MODES[garbage_mode] + + config = Config( + Ms=_TOKEN_PADDING_M, + K=_TOKEN_PADDING_K, + N=_TOKEN_PADDING_N, + E=_PADDING_E, + topks=_PADDING_TOPK, + dtype=torch.bfloat16, + quant_config=quant_config, + prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular, + fused_experts_type=AiterExperts, + world_size=1, + ) + assert config.is_valid()[0] + assert config.fe_supports_quant_scheme(), ( + f"AiterExperts does not support quant scheme {quant_config}." + ) + + weights = WeightTensors.make(config) + vllm_config, env_dict = config.make_env_data() + + parallel_launch_with_config( + config.world_size, + _aiter_token_padding_worker, + vllm_config, + env_dict, + config, + weights, + False, + garbage_rows, + ) + + +@pytest.mark.skip( + reason="MXFP4 AiterExperts token padding requires MI350/gfx950 hardware. " + 'See https://github.com/vllm-project/vllm/issues/54966 ("Test padding").' +) +def test_aiter_moe_token_padding_garbage_rows_mxfp4(): + pytest.skip("MXFP4 AiterExperts token padding requires MI350/gfx950 hardware.") + + if __name__ == "__main__": # Ability to test individual PrepareAndFinalize and FusedExperts combination from .modular_kernel_tools.cli_args import make_config, make_config_arg_parser From 3d30267db41b70b9862bdfa510bdc565ad7e5d39 Mon Sep 17 00:00:00 2001 From: Divakar Verma Date: Wed, 7 Oct 2026 01:28:56 +0000 Subject: [PATCH 2/7] [nit] Drop dead inner pytest.skip() in MXFP4 token-padding stub The outer @pytest.mark.skip already prevents this test from ever running; the inner pytest.skip() call is unreachable dead code. Signed-off-by: Divakar Verma --- tests/kernels/moe/test_modular_kernel_combinations.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py index 6f9963d881f6..343e3cfd60fa 100644 --- a/tests/kernels/moe/test_modular_kernel_combinations.py +++ b/tests/kernels/moe/test_modular_kernel_combinations.py @@ -1135,7 +1135,7 @@ def test_aiter_moe_token_padding_garbage_rows( 'See https://github.com/vllm-project/vllm/issues/54966 ("Test padding").' ) def test_aiter_moe_token_padding_garbage_rows_mxfp4(): - pytest.skip("MXFP4 AiterExperts token padding requires MI350/gfx950 hardware.") + pass if __name__ == "__main__": From 5ebcf4cd8d4be942cbf706bcabcb0612339e57cf Mon Sep 17 00:00:00 2001 From: Divakar Verma Date: Wed, 7 Oct 2026 01:59:22 +0000 Subject: [PATCH 3/7] [ROCm] Guard against degenerate (zeroed) output in token-padding test ref_out's typical magnitude here (~1e-3-1e-2) sits below the 3e-2 atol used by both check_accuracy and torch.testing.assert_close, so a fully zeroed-out mk_out -- e.g. a garbage row's inf/nan corrupting the real rows via some leak the isfinite check doesn't catch (zero isn't NaN/Inf) -- would still spuriously pass. Mirrors the same guard added to the stacked 4b padding-matrix test (#59333) for the same reason. Signed-off-by: Divakar Verma --- tests/kernels/moe/test_modular_kernel_combinations.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py index 343e3cfd60fa..ac791d13f755 100644 --- a/tests/kernels/moe/test_modular_kernel_combinations.py +++ b/tests/kernels/moe/test_modular_kernel_combinations.py @@ -1066,6 +1066,15 @@ def _aiter_token_padding_worker( "through AiterExperts." ) + # ref_out's magnitude here (~1e-3-1e-2) is below atol=3e-2, so a zeroed-out + # mk_out would still pass the checks below -- guard against that. + ref_scale = ref_out[real_rows].abs().mean() + mk_scale = mk_out[real_rows].abs().mean() + assert mk_scale > 0.5 * ref_scale, ( + f"AiterExperts output magnitude (mean |mk_out|={mk_scale:.6f}) looks " + f"degenerate/zeroed vs. reference (mean |ref_out|={ref_scale:.6f})." + ) + if config.quant_config is not None: check_accuracy( ref_out[real_rows], mk_out[real_rows], atol=3e-2, rtol=3e-2, percent=0.9 From ed20edd99d0cdb3b16e8804992686a38fe196696 Mon Sep 17 00:00:00 2001 From: Divakar Verma Date: Wed, 7 Oct 2026 02:00:22 +0000 Subject: [PATCH 4/7] [nit] Reuse 4b's padded K/N/M constants instead of re-declaring them _TOKEN_PADDING_M/K/N duplicated _PADDING_M/_PADDING_K_PADDED/ _PADDING_N_PADDED's values under new names. Reusing the originals avoids the two sets silently drifting out of sync. Signed-off-by: Divakar Verma --- .../moe/test_modular_kernel_combinations.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py index ac791d13f755..485e8e25d625 100644 --- a/tests/kernels/moe/test_modular_kernel_combinations.py +++ b/tests/kernels/moe/test_modular_kernel_combinations.py @@ -1005,10 +1005,6 @@ def test_aiter_moe_padding_matrix_mxfp4(): # RankTensors.make() *before* garbage rows are injected, so it cannot be # inflated by a garbage row's magnitude either. -_TOKEN_PADDING_M = 16 -_TOKEN_PADDING_K = 1024 -_TOKEN_PADDING_N = 512 - # {mode name -> {row index: fill value}}. Covers multiple garbage-row counts # and positions, and both inf and nan. _TOKEN_PADDING_GARBAGE_MODES: dict[str, dict[int, float]] = { @@ -1107,10 +1103,15 @@ def test_aiter_moe_token_padding_garbage_rows( garbage_rows = _TOKEN_PADDING_GARBAGE_MODES[garbage_mode] + # Reuses 4b's already-confirmed 128-aligned ("padded") K/N -- these are + # known to work across all 5 quant schemes here, including AITER's + # block-quant kernel (which 4b's _PADDING_K_UNPADDED/_PADDING_N_UNPADDED + # sizes do not support). Token-padding safety doesn't depend on + # hidden_pad/intermediate_pad, so there's no need for a separate shape. config = Config( - Ms=_TOKEN_PADDING_M, - K=_TOKEN_PADDING_K, - N=_TOKEN_PADDING_N, + Ms=_PADDING_M, + K=_PADDING_K_PADDED, + N=_PADDING_N_PADDED, E=_PADDING_E, topks=_PADDING_TOPK, dtype=torch.bfloat16, From cf05cb473c7e43601afa0e99f7e847f1716f345b Mon Sep 17 00:00:00 2001 From: Divakar Verma Date: Thu, 8 Oct 2026 17:16:37 +0000 Subject: [PATCH 5/7] [ROCm] Skip fp8_tensor_token inf-garbage-row padding cases (AITER bug) AiterExperts falls back to AITER's dynamic per-tensor activation quant for this scheme (rocm_aiter_moe.py ignores per_act_token_quant when weights are per-tensor), whose scale is a max-abs reduction over the whole batch. A real inf garbage row legitimately poisons that shared scale, zeroing every real row's output -- not something AiterExperts can be expected to isolate against today. Filed upstream as https://github.com/ROCm/aiter/issues/6275. Co-Authored-By: Claude Sonnet 5 Signed-off-by: Divakar Verma --- .../moe/test_modular_kernel_combinations.py | 24 ++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py index 485e8e25d625..6364b0eeba92 100644 --- a/tests/kernels/moe/test_modular_kernel_combinations.py +++ b/tests/kernels/moe/test_modular_kernel_combinations.py @@ -1003,7 +1003,9 @@ def test_aiter_moe_padding_matrix_mxfp4(): # a1 quant scale (when applicable) is fixed once in run_modular_kernel()'s # FusedMoEQuantConfig construction, from hidden_states_scale computed inside # RankTensors.make() *before* garbage rows are injected, so it cannot be -# inflated by a garbage row's magnitude either. +# inflated by a garbage row's magnitude either -- EXCEPT "fp8_tensor_token", +# whose a1_scale AITER computes live from the (garbage-injected) +# hidden_states. See the skip below for why it's excluded from inf modes. # {mode name -> {row index: fill value}}. Covers multiple garbage-row counts # and positions, and both inf and nan. @@ -1103,6 +1105,26 @@ def test_aiter_moe_token_padding_garbage_rows( garbage_rows = _TOKEN_PADDING_GARBAGE_MODES[garbage_mode] + # fp8_tensor_token falls back to AITER's dynamic per-tensor activation + # quant (rocm_aiter_moe.py ignores per_act_token_quant when weights are + # per-tensor), whose scale is a max-abs reduction over the whole batch -- + # a real inf garbage row poisons every real row's scale too. Upstream + # bug: https://github.com/ROCm/aiter/issues/6275. + is_fp8_tensor_token = ( + quant_config is not None + and not quant_config.per_out_ch_quant + and quant_config.per_act_token_quant + and quant_config.block_shape is None + ) + if is_fp8_tensor_token and any( + v in (float("inf"), float("-inf")) for v in garbage_rows.values() + ): + pytest.skip( + "Batch-wide dynamic per-tensor quant scale is poisoned by a " + "real inf garbage row. See " + "https://github.com/ROCm/aiter/issues/6275." + ) + # Reuses 4b's already-confirmed 128-aligned ("padded") K/N -- these are # known to work across all 5 quant schemes here, including AITER's # block-quant kernel (which 4b's _PADDING_K_UNPADDED/_PADDING_N_UNPADDED From 7bb2242a8ea52f252586bb0e67b286bf294919dd Mon Sep 17 00:00:00 2001 From: Divakar Verma Date: Thu, 8 Oct 2026 17:55:29 +0000 Subject: [PATCH 6/7] [nit] Reuse Config.fp8_quant_key_pair() in token-padding skip condition Avoids re-deriving (per_out_ch_quant, per_act_token_quant, block_shape) by hand -- matches the existing assert_aiter_quant_scheme_case() pattern. Co-Authored-By: Claude Sonnet 5 Signed-off-by: Divakar Verma --- .../moe/test_modular_kernel_combinations.py | 39 +++++++++---------- 1 file changed, 19 insertions(+), 20 deletions(-) diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py index 6364b0eeba92..5f7df709ccd7 100644 --- a/tests/kernels/moe/test_modular_kernel_combinations.py +++ b/tests/kernels/moe/test_modular_kernel_combinations.py @@ -17,6 +17,10 @@ from vllm.config import VllmConfig, set_current_vllm_config from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import AiterExperts +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8DynamicTensorSym, + kFp8StaticTensorSym, +) from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe from vllm.utils.import_utils import has_aiter, has_deep_ep, has_deep_gemm @@ -1105,26 +1109,6 @@ def test_aiter_moe_token_padding_garbage_rows( garbage_rows = _TOKEN_PADDING_GARBAGE_MODES[garbage_mode] - # fp8_tensor_token falls back to AITER's dynamic per-tensor activation - # quant (rocm_aiter_moe.py ignores per_act_token_quant when weights are - # per-tensor), whose scale is a max-abs reduction over the whole batch -- - # a real inf garbage row poisons every real row's scale too. Upstream - # bug: https://github.com/ROCm/aiter/issues/6275. - is_fp8_tensor_token = ( - quant_config is not None - and not quant_config.per_out_ch_quant - and quant_config.per_act_token_quant - and quant_config.block_shape is None - ) - if is_fp8_tensor_token and any( - v in (float("inf"), float("-inf")) for v in garbage_rows.values() - ): - pytest.skip( - "Batch-wide dynamic per-tensor quant scale is poisoned by a " - "real inf garbage row. See " - "https://github.com/ROCm/aiter/issues/6275." - ) - # Reuses 4b's already-confirmed 128-aligned ("padded") K/N -- these are # known to work across all 5 quant schemes here, including AITER's # block-quant kernel (which 4b's _PADDING_K_UNPADDED/_PADDING_N_UNPADDED @@ -1147,6 +1131,21 @@ def test_aiter_moe_token_padding_garbage_rows( f"AiterExperts does not support quant scheme {quant_config}." ) + # This (weight, activation) pair is AiterExperts's real per-tensor + # dynamic-activation-quant scheme -- its batch-wide max-abs scale is + # poisoned by a real inf garbage row. Upstream bug: + # https://github.com/ROCm/aiter/issues/6275. + if ( + quant_config is not None + and config.fp8_quant_key_pair() == (kFp8StaticTensorSym, kFp8DynamicTensorSym) + and any(v in (float("inf"), float("-inf")) for v in garbage_rows.values()) + ): + pytest.skip( + "Batch-wide dynamic per-tensor quant scale is poisoned by a " + "real inf garbage row. See " + "https://github.com/ROCm/aiter/issues/6275." + ) + weights = WeightTensors.make(config) vllm_config, env_dict = config.make_env_data() From 8dfe8636d018363398d53bda7c5ba5e21b250291 Mon Sep 17 00:00:00 2001 From: Divakar Verma Date: Thu, 8 Oct 2026 18:02:51 +0000 Subject: [PATCH 7/7] [nit] Trim module comment in token-padding test The dropped lines duplicated what the test's own docstring already says (what the HIP-graph padding contract simulates); keep only the non-duplicated "why this is a valid comparison" rationale. Co-Authored-By: Claude Sonnet 5 Signed-off-by: Divakar Verma --- .../moe/test_modular_kernel_combinations.py | 22 +++++-------------- 1 file changed, 6 insertions(+), 16 deletions(-) diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py index 5f7df709ccd7..f29c22085240 100644 --- a/tests/kernels/moe/test_modular_kernel_combinations.py +++ b/tests/kernels/moe/test_modular_kernel_combinations.py @@ -994,22 +994,12 @@ def test_aiter_moe_padding_matrix_mxfp4(): # --- 4c: HIP-graph token padding -- inf/nan garbage rows (topk_ids=-1) ------ -# Simulates the runtime's padding-token contract for garbage/padding rows -# (grouped_topk_router.py: `topk_ids.masked_fill(is_padding, -1)`, -# `topk_weights.masked_fill(is_padding, 0)`), without needing real HIP-graph -# capture. Empirically confirmed (see -# https://github.com/vllm-project/vllm/issues/54966 discussion) that -# AiterExperts safely isolates these rows under both the default ("opus") -# and FlyDSL AITER sorting backends: garbage rows' inf/nan hidden_states -# never leak NaN/Inf into real rows' output. reference_moe_impl's dispatch -# (`mask = topk_ids == i`, tests/kernels/utils.py) is likewise row-local, so -# comparing real rows against it is a valid check. AiterExperts's per-tensor -# a1 quant scale (when applicable) is fixed once in run_modular_kernel()'s -# FusedMoEQuantConfig construction, from hidden_states_scale computed inside -# RankTensors.make() *before* garbage rows are injected, so it cannot be -# inflated by a garbage row's magnitude either -- EXCEPT "fp8_tensor_token", -# whose a1_scale AITER computes live from the (garbage-injected) -# hidden_states. See the skip below for why it's excluded from inf modes. +# AiterExperts isolates garbage rows from real rows under both AITER sorting +# backends, and reference_moe_impl's dispatch (`mask = topk_ids == i`) is +# likewise row-local, so comparing real rows against it is valid. Exception: +# AiterExperts's a1 quant scale is normally fixed before garbage injection +# (computed in RankTensors.make()), but "fp8_tensor_token" computes it live +# from the (already-injected) hidden_states -- see the skip below. # {mode name -> {row index: fill value}}. Covers multiple garbage-row counts # and positions, and both inf and nan.