From 387380d5c5db7e6d8c5bea063c4b1faec961b08d Mon Sep 17 00:00:00 2001 From: amd-weisun Date: Mon, 17 Aug 2026 15:52:31 +0000 Subject: [PATCH 01/10] [ROCm][AITER] Trim MoE GEMM input to valid token prefix before calling AITER AITER's fused_moe kernel silently corrupts its own output (returns all-zero rows) when the allocated input buffer (hidden_states.shape[0]) is much larger than the real number of valid tokens it is asked to process. This has been confirmed independently of MoRI or any other EP dispatcher, via a standalone microbenchmark that calls AITER's fused_moe directly and compares against a from-scratch reference GEMM. This condition is common with EP dispatchers (e.g. MoRI) that pre-allocate a fixed-size receive buffer sized for the worst case (max_num_batched_tokens * world_size), while the actual per-step valid token count is usually far smaller (a handful of tokens per rank in normal decode-heavy serving). AITER's fused_moe() call receives the full pre-allocated buffer shape even though num_local_tokens correctly identifies only a small valid prefix within it. Fix: in AiterExperts.apply(), when num_local_tokens is available and smaller than the input buffer, slice hidden_states/topk_weights/ topk_ids/a1q_scale down to the valid prefix before calling rocm_aiter_fused_experts(), then scatter the (smaller) result back into a zero-initialized output buffer of the original shape. Padding rows beyond the valid prefix are never read downstream (num_local_tokens already tells consumers where the real data ends), so zero-filling them is semantically a no-op. Applied unconditionally: this is a real correctness defect, not a behavior change some configurations might prefer to opt out of, so there is no feature flag to disable it. Co-Authored-By: Claude Sonnet 5 Signed-off-by: amd-weisun --- .../fused_moe/experts/rocm_aiter_moe.py | 102 +++++++++++++++++- .../layers/fused_moe/prepare_finalize/mori.py | 56 ++++++++++ 2 files changed, 154 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py index a433b18487dd..131758fa3dcc 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py @@ -7,6 +7,7 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm._aiter_ops import rocm_aiter_ops +from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import ( FUSED_MOE_UNQUANTIZED_CONFIG, @@ -28,6 +29,7 @@ kMxfp4Dynamic, kMxfp4Static, ) +from vllm.platforms import current_platform class QuantMethod(IntEnum): @@ -53,6 +55,35 @@ class ActivationMethod(IntEnum): GELU = 1 +def _is_uniform_full_graph_batch() -> bool: + # Under full-cudagraph capture/replay every rank is padded to the same + # static token count -- used below to pick a capture-safe trim bound + # instead of a live `.item()` count. + from vllm.config import CUDAGraphMode + + if not is_forward_context_available(): + return False + forward_context = get_forward_context() + batch_descriptor = forward_context.batch_descriptor + return ( + forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL + and batch_descriptor is not None + and batch_descriptor.uniform + ) + + +def _is_stream_capturing() -> bool: + # torch.cuda.is_current_stream_capturing() is unavailable on non-CUDA + # (XPU) torch; guard the same way fused_batched_moe.py does. + try: + return ( + current_platform.is_cuda_alike() + and torch.cuda.is_current_stream_capturing() + ) + except Exception: + return False + + aiter_topK_meta_data: tuple[torch.Tensor, torch.Tensor] | None = None @@ -543,22 +574,85 @@ def apply( else: num_local_tokens = None + # AITER's fused_moe silently corrupts output (all-zero rows) when its + # input buffer is much larger than the real valid-token count -- + # common with EP dispatchers like MoRI that pre-allocate worst-case + # buffers. Trim to the valid prefix before calling AITER, then + # scatter the result into a zero-filled full-size buffer (padding + # rows are never read downstream). + orig_num_tokens = hidden_states.shape[0] + if num_local_tokens is None: + valid_num_tokens = orig_num_tokens + elif _is_uniform_full_graph_batch(): + # `.item()` is unsafe mid-capture and a live-count shape breaks + # replay, so use a static worst-case bound instead: each of this + # rank's `graph_bs` tokens can select up to `topk` experts, and + # each selection can land on any of `dispatch_group_size` peer + # ranks, so a rank can receive up to + # `graph_bs * topk * dispatch_group_size` rows. A narrower bound + # that dropped the `topk` factor (assuming multiple topk + # selections landing on the same destination rank always get + # deduplicated into one send) is not safe to assume in general + # and risks under-counting real traffic. + batch_descriptor = get_forward_context().batch_descriptor + dispatch_group_size = ( + self.moe_config.ep_size + if self.moe_config.use_ep + else self.moe_config.dp_size + ) + valid_num_tokens = min( + batch_descriptor.num_tokens + * self.moe_config.experts_per_token + * dispatch_group_size, + orig_num_tokens, + ) + elif _is_stream_capturing(): + # Mid-capture but not the uniform-full-graph case handled above + # (e.g. PIECEWISE cudagraphs, which can still wrap this call in a + # captured region): `.item()` is illegal here too, and there is no + # safe static bound available, so skip trimming for this call. + valid_num_tokens = orig_num_tokens + else: + valid_num_tokens = int(num_local_tokens[0].item()) + needs_trim = valid_num_tokens < orig_num_tokens + if needs_trim: + hidden_states_in = hidden_states[:valid_num_tokens] + topk_weights_in = topk_weights[:valid_num_tokens] + topk_ids_in = topk_ids[:valid_num_tokens] + if ( + a1q_scale is not None + and a1q_scale.dim() > 0 + and a1q_scale.shape[0] == orig_num_tokens + ): + a1q_scale_in = a1q_scale[:valid_num_tokens] + else: + a1q_scale_in = a1q_scale + else: + hidden_states_in = hidden_states + topk_weights_in = topk_weights + topk_ids_in = topk_ids + a1q_scale_in = a1q_scale + result = rocm_aiter_fused_experts( - hidden_states=hidden_states, + hidden_states=hidden_states_in, w1=w1, w2=w2, - topk_weights=topk_weights, - topk_ids=topk_ids, + topk_weights=topk_weights_in, + topk_ids=topk_ids_in, activation=activation, apply_router_weight_on_input=apply_router_weight_on_input, expert_mask=expert_map, quant_config=self.quant_config, moe_config=self.moe_config, - a1q_scale=a1q_scale, + a1q_scale=a1q_scale_in, num_local_tokens=num_local_tokens, output_dtype=output.dtype, moe_sorting_dispatch_policy=rocm_aiter_ops.get_moe_dispatch_policy(), ) + if needs_trim: + output.zero_() + output[:valid_num_tokens] = result + return # avoid redundant copy when output is a view of the result if ( output.shape == result.shape diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py index fe3a53941806..57da33a30578 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py @@ -12,6 +12,49 @@ logger = init_logger(__name__) +def _is_stream_capturing() -> bool: + try: + return ( + current_platform.is_cuda_alike() + and torch.cuda.is_current_stream_capturing() + ) + except Exception: + return False + + +def _trim_dispatch_output( + dispatch_a1: torch.Tensor, + dispatch_scale: torch.Tensor | None, + dispatch_ids: torch.Tensor, + dispatch_weights: torch.Tensor, + dispatch_recv_token_num: torch.Tensor, + ep_size: int, + topk: int, + input_num_tokens: int, +) -> tuple[ + torch.Tensor, torch.Tensor | None, torch.Tensor, torch.Tensor +]: + orig_rows = dispatch_a1.shape[0] + + if _is_stream_capturing(): + valid_rows = input_num_tokens * topk * ep_size + elif dispatch_recv_token_num is not None and dispatch_recv_token_num.numel() > 0: + valid_rows = int(dispatch_recv_token_num.reshape(-1)[0].item()) + else: + return dispatch_a1, dispatch_scale, dispatch_ids, dispatch_weights + + valid_rows = max(0, min(valid_rows, orig_rows)) + if valid_rows <= 0 or valid_rows >= orig_rows: + return dispatch_a1, dispatch_scale, dispatch_ids, dispatch_weights + + dispatch_a1 = dispatch_a1[:valid_rows] + dispatch_ids = dispatch_ids[:valid_rows] + dispatch_weights = dispatch_weights[:valid_rows] + if dispatch_scale is not None: + dispatch_scale = dispatch_scale[:valid_rows] + return dispatch_a1, dispatch_scale, dispatch_ids, dispatch_weights + + class MoriPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): """ Prepare/Finalize using MoRI kernels. @@ -94,6 +137,19 @@ def prepare( dispatch_recv_token_num, ) = self.mori_op.dispatch(a1, topk_weights, scale, topk_ids) + dispatch_a1, dispatch_scale, dispatch_ids, dispatch_weights = ( + _trim_dispatch_output( + dispatch_a1, + dispatch_scale, + dispatch_ids, + dispatch_weights, + dispatch_recv_token_num, + ep_size=self.num_dispatchers_, + topk=topk_ids.shape[1], + input_num_tokens=a1.shape[0], + ) + ) + expert_tokens_meta = mk.ExpertTokensMetadata( expert_num_tokens=dispatch_recv_token_num, expert_num_tokens_cpu=None ) From 3fe04af6b0e53e83db669dac7106cb1875a3bcf1 Mon Sep 17 00:00:00 2001 From: amd-weisun Date: Wed, 9 Sep 2026 18:18:39 +0100 Subject: [PATCH 02/10] trim code comments Signed-off-by: amd-weisun --- .../fused_moe/experts/rocm_aiter_moe.py | 23 +++---------------- 1 file changed, 3 insertions(+), 20 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py index 131758fa3dcc..df4f95833a1b 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py @@ -574,26 +574,12 @@ def apply( else: num_local_tokens = None - # AITER's fused_moe silently corrupts output (all-zero rows) when its - # input buffer is much larger than the real valid-token count -- - # common with EP dispatchers like MoRI that pre-allocate worst-case - # buffers. Trim to the valid prefix before calling AITER, then - # scatter the result into a zero-filled full-size buffer (padding - # rows are never read downstream). + orig_num_tokens = hidden_states.shape[0] if num_local_tokens is None: valid_num_tokens = orig_num_tokens elif _is_uniform_full_graph_batch(): - # `.item()` is unsafe mid-capture and a live-count shape breaks - # replay, so use a static worst-case bound instead: each of this - # rank's `graph_bs` tokens can select up to `topk` experts, and - # each selection can land on any of `dispatch_group_size` peer - # ranks, so a rank can receive up to - # `graph_bs * topk * dispatch_group_size` rows. A narrower bound - # that dropped the `topk` factor (assuming multiple topk - # selections landing on the same destination rank always get - # deduplicated into one send) is not safe to assume in general - # and risks under-counting real traffic. + batch_descriptor = get_forward_context().batch_descriptor dispatch_group_size = ( self.moe_config.ep_size @@ -607,10 +593,7 @@ def apply( orig_num_tokens, ) elif _is_stream_capturing(): - # Mid-capture but not the uniform-full-graph case handled above - # (e.g. PIECEWISE cudagraphs, which can still wrap this call in a - # captured region): `.item()` is illegal here too, and there is no - # safe static bound available, so skip trimming for this call. + valid_num_tokens = orig_num_tokens else: valid_num_tokens = int(num_local_tokens[0].item()) From 54b7e025ee5664bf51001589f9ab8c1e12991b75 Mon Sep 17 00:00:00 2001 From: amd-weisun Date: Fri, 18 Sep 2026 11:02:50 +0000 Subject: [PATCH 03/10] Add regression test for AiterExperts oversized-buffer trim Drives AiterExperts.apply() directly with a synthetic post-dispatch-shaped buffer (most rows unwritten padding, a small valid prefix) at the row count where AITER's fused_moe kernel is known to corrupt its own output. Verified on MI300 hardware that reverting the trim makes the large-buffer case fail with non-finite output, confirming this catches the regression. Co-Authored-By: Claude Sonnet 5 Signed-off-by: amd-weisun --- tests/kernels/moe/test_rocm_aiter_moe.py | 151 +++++++++++++++++++++++ 1 file changed, 151 insertions(+) diff --git a/tests/kernels/moe/test_rocm_aiter_moe.py b/tests/kernels/moe/test_rocm_aiter_moe.py index 89244bed1b0f..06388ecc55a6 100644 --- a/tests/kernels/moe/test_rocm_aiter_moe.py +++ b/tests/kernels/moe/test_rocm_aiter_moe.py @@ -9,6 +9,9 @@ - gfx950-only AITER MXFP4 W4A16 MoE support, accuracy, and determinism - MoE-facing FP8 group-quant activation quality - deterministic routing and representative gfx942 / gfx950 coverage +- ``AiterExperts.apply()`` trimming an oversized EP-dispatch-shaped buffer + down to the valid token prefix before calling AITER (regression coverage + for the buffer-size correctness defect fixed by that trim) Generic fused-MoE backend selection and non-ROCm kernel coverage live in the generic MoE test files under ``tests/kernels/moe``. @@ -953,6 +956,154 @@ def test_aiter_fused_moe_end_to_end( ) +# Oversized dispatch-buffer regression tests ------------------------------- + + +@pytest.mark.parametrize( + "buffer_rows,valid_tokens", + [ + pytest.param(65536, 8, id="buffer_at_aiter_corruption_threshold"), + pytest.param(128, 8, id="small_buffer_control"), + ], +) +def test_aiter_experts_apply_trims_oversized_dispatch_buffer( + buffer_rows: int, + valid_tokens: int, +): + """AiterExperts.apply() must trim an oversized EP-dispatch-shaped buffer + down to the valid token prefix before calling AITER, which silently + corrupts its output once the row count crosses large thresholds + (e.g. 65536), regardless of how many rows hold real data. + """ + import vllm.model_executor.layers.fused_moe.modular_kernel as mk + from vllm.config import VllmConfig + from vllm.forward_context import set_forward_context + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + from vllm.model_executor.layers.fused_moe.config import ( + FUSED_MOE_UNQUANTIZED_CONFIG, + ) + from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( + AiterExperts, + ) + + from .utils import make_dummy_moe_config + + _assert_aiter_supported() + torch.set_default_device("cuda") + set_random_seed(0) + + hidden_dim = 128 + intermediate_dim = 256 + num_experts = 4 + topk = 2 + activation = MoEActivation.SILU + + moe_config = make_dummy_moe_config( + num_experts=num_experts, + experts_per_token=topk, + hidden_dim=hidden_dim, + intermediate_size=intermediate_dim, + in_dtype=torch.bfloat16, + max_num_tokens=buffer_rows, + ) + experts = AiterExperts( + moe_config=moe_config, + quant_config=FUSED_MOE_UNQUANTIZED_CONFIG, + ) + + w1 = torch.randn( + num_experts, intermediate_dim * 2, hidden_dim, dtype=torch.bfloat16 + ) / math.sqrt(hidden_dim) + w2 = torch.randn( + num_experts, hidden_dim, intermediate_dim, dtype=torch.bfloat16 + ) / math.sqrt(intermediate_dim) + w1_shuffled, w2_shuffled = _shuffle_moe_weights(w1, w2) + + # Only the first `valid_tokens` rows carry real data; the rest mimics + # the unreceived tail of a real EP dispatch buffer. + hidden_states = torch.zeros(buffer_rows, hidden_dim, dtype=torch.bfloat16) + topk_weights = torch.zeros(buffer_rows, topk, dtype=torch.float32) + topk_ids = torch.zeros(buffer_rows, topk, dtype=torch.int32) + + hidden_states[:valid_tokens] = torch.randn( + valid_tokens, hidden_dim, dtype=torch.bfloat16 + ) + valid_weights = torch.rand(valid_tokens, topk, dtype=torch.float32) + topk_weights[:valid_tokens] = valid_weights / valid_weights.sum( + dim=-1, keepdim=True + ) + topk_ids[:valid_tokens] = _make_topk_ids(valid_tokens, num_experts, topk) + + expert_tokens_meta = mk.ExpertTokensMetadata( + expert_num_tokens=torch.tensor([valid_tokens], dtype=torch.int32), + expert_num_tokens_cpu=None, + ) + + workspace13_shape, workspace2_shape, output_shape = experts.workspace_shapes( + M=buffer_rows, + N=intermediate_dim, + K=hidden_dim, + topk=topk, + global_num_experts=num_experts, + local_num_experts=num_experts, + expert_tokens_meta=expert_tokens_meta, + activation=activation, + ) + workspace13 = torch.empty(workspace13_shape, dtype=torch.bfloat16) + workspace2 = torch.empty(workspace2_shape, dtype=torch.bfloat16) + output = torch.full(output_shape, torch.nan, dtype=torch.bfloat16) + + vllm_config = VllmConfig() + with set_forward_context(None, vllm_config, num_tokens=buffer_rows): + experts.apply( + output=output, + hidden_states=hidden_states, + w1=w1_shuffled, + w2=w2_shuffled, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=activation, + global_num_experts=num_experts, + expert_map=None, + a1q_scale=None, + a2_scale=None, + workspace13=workspace13, + workspace2=workspace2, + expert_tokens_meta=expert_tokens_meta, + apply_router_weight_on_input=False, + ) + + assert torch.isfinite(output).all() + + valid_out = output[:valid_tokens] + assert valid_out.abs().sum() > 0, ( + "AiterExperts.apply() returned all-zero output for valid tokens -- " + "the oversized-buffer trim appears to be missing or broken" + ) + + ref_out = ref_moe_forward( + hidden_states[:valid_tokens], + w1, + w2, + topk_weights[:valid_tokens], + topk_ids[:valid_tokens], + activation="silu", + ) + _assert_close_budget( + valid_out.float(), + ref_out, + label=( + f"oversized_dispatch_buffer buffer_rows={buffer_rows} " + f"valid_tokens={valid_tokens}" + ), + atol=0.05, + rtol=0.0, + ) + + # Padding rows must never be touched/scattered into. + assert torch.all(output[valid_tokens:] == 0) + + # Arch-specific tests ----------------------------------------------------- From dd049890b21219fa811dc02c2b0b22e71e2572a0 Mon Sep 17 00:00:00 2001 From: amd-weisun Date: Fri, 18 Sep 2026 11:39:35 +0000 Subject: [PATCH 04/10] Fix pre-commit: ruff-format cleanup and mypy null-check mypy can't see across the _is_uniform_full_graph_batch() helper that batch_descriptor is non-None when it returns True, so assert it explicitly at the call site. Co-Authored-By: Claude Sonnet 5 Signed-off-by: amd-weisun --- .../model_executor/layers/fused_moe/experts/rocm_aiter_moe.py | 4 +--- vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py | 4 +--- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py index df4f95833a1b..fea882da5ff9 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py @@ -574,13 +574,12 @@ def apply( else: num_local_tokens = None - orig_num_tokens = hidden_states.shape[0] if num_local_tokens is None: valid_num_tokens = orig_num_tokens elif _is_uniform_full_graph_batch(): - batch_descriptor = get_forward_context().batch_descriptor + assert batch_descriptor is not None dispatch_group_size = ( self.moe_config.ep_size if self.moe_config.use_ep @@ -593,7 +592,6 @@ def apply( orig_num_tokens, ) elif _is_stream_capturing(): - valid_num_tokens = orig_num_tokens else: valid_num_tokens = int(num_local_tokens[0].item()) diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py index 57da33a30578..a68ad8b1cdc4 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py @@ -31,9 +31,7 @@ def _trim_dispatch_output( ep_size: int, topk: int, input_num_tokens: int, -) -> tuple[ - torch.Tensor, torch.Tensor | None, torch.Tensor, torch.Tensor -]: +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor, torch.Tensor]: orig_rows = dispatch_a1.shape[0] if _is_stream_capturing(): From 2bb98444b4a64ebcd7c681f01cd52f1b70d5f529 Mon Sep 17 00:00:00 2001 From: amd-weisun Date: Fri, 18 Sep 2026 12:48:15 +0000 Subject: [PATCH 05/10] Drop non-required padding-zero assertion from oversized-buffer test Padding content past the valid prefix is never read downstream (the fix's own commit message notes zero-filling it is "semantically a no-op"), so asserting it stays exactly zero checks an implementation detail, not a real contract. The whole-buffer isfinite() check already covers the actual correctness/safety concern for that region. Co-Authored-By: Claude Sonnet 5 Signed-off-by: amd-weisun --- tests/kernels/moe/test_rocm_aiter_moe.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/kernels/moe/test_rocm_aiter_moe.py b/tests/kernels/moe/test_rocm_aiter_moe.py index 06388ecc55a6..c17c877ed816 100644 --- a/tests/kernels/moe/test_rocm_aiter_moe.py +++ b/tests/kernels/moe/test_rocm_aiter_moe.py @@ -1100,9 +1100,6 @@ def test_aiter_experts_apply_trims_oversized_dispatch_buffer( rtol=0.0, ) - # Padding rows must never be touched/scattered into. - assert torch.all(output[valid_tokens:] == 0) - # Arch-specific tests ----------------------------------------------------- From 2eb5a3913fc57b4b42a1b6e717f90f2deb4f0423 Mon Sep 17 00:00:00 2001 From: amd-weisun Date: Mon, 21 Sep 2026 14:48:32 +0000 Subject: [PATCH 06/10] Add real multi-node MoRI/AITER regression test for buffer-trim fix Adds tests/kernels/moe/test_mori_aiter_ep_trim_regression.py, which drives real MoRI InterNodeV1 dispatch, real AiterExperts.apply() (AITER GEMM), and real MoRI combine directly (bypassing FusedMoEFactory, matching test_moe_layer_multinode.py's approach), repeating dispatch/combine rounds on the same reused mori_op buffers to reproduce the buffer-trim regression. Verified on real 4-node/EP32 InterNodeV1 hardware via torchrun (two independent node allocations): with the trim fix present, the test passes across repeated rounds; with it reverted, real tokens' combined output goes all-zero after buffer reuse and the test fails with an AssertionError, matching the corruption mechanism this branch fixes. Co-Authored-By: Claude Sonnet 5 Signed-off-by: amd-weisun --- .../moe/test_mori_aiter_ep_trim_regression.py | 223 ++++++++++++++++++ 1 file changed, 223 insertions(+) create mode 100644 tests/kernels/moe/test_mori_aiter_ep_trim_regression.py diff --git a/tests/kernels/moe/test_mori_aiter_ep_trim_regression.py b/tests/kernels/moe/test_mori_aiter_ep_trim_regression.py new file mode 100644 index 000000000000..c65f0c1beaad --- /dev/null +++ b/tests/kernels/moe/test_mori_aiter_ep_trim_regression.py @@ -0,0 +1,223 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Real multi-node MoRI dispatch/AITER GEMM/MoRI combine regression test for +the buffer-trim fix in rocm_aiter_moe.py / prepare_finalize/mori.py. + +Drives the three real production primitives directly (MoriPrepareAndFinalize +.prepare()/.finalize() and AiterExperts.apply()) instead of going through +FusedMoEFactory/the full FusedMoE layer -- this is intentional, see +test_moe_layer_multinode.py's docstring for why that heavier construction +path doesn't work here (older pinned vLLM install used for multi-node CI +lacks FusedMoEFactory). + +Reproduces on real 4-node/EP32 InterNodeV1 hardware: without the trim, +repeated dispatch/combine rounds on the same (reused) mori_op buffers +produce all-zero combined output for genuinely valid, real tokens (MoRI's +own reported valid-row count balloons toward the oversized buffer's full +capacity instead of the true tiny per-round count, and AITER's output for +those rows is discarded/corrupted downstream in combine()). A single round +right after construction is NOT enough to show this -- it takes at least +one prior large round to have populated/reused the buffer. With the trim +applied, results stay correct and nonzero across repeated rounds. + +Launch (mirrors test_moe_layer_multinode.py -- one torchrun per host, all +pointed at the same rendezvous, one host acting as rank 0): + + MASTER_ADDR= MASTER_PORT=29525 \\ + torchrun --nnodes 4 --nproc-per-node 8 --node-rank <0..3> \\ + --rdzv_backend=c10d --rdzv_endpoint=:29525 \\ + -m pytest -v -s kernels/moe/test_mori_aiter_ep_trim_regression.py +""" + +import importlib.util +import os + +import pytest +import torch + +pytestmark = pytest.mark.skipif( + importlib.util.find_spec("mori") is None, reason="mori not installed" +) + +GPU_PER_NODE = 8 +HIDDEN = 4096 +INTER = 1024 +TOPK = 8 +EXPERTS_PER_RANK = 8 +MAX_TOK = 128 +NUM_BIG_SMALL_ROUNDS = 3 + + +def _make_topk_ids(num_tokens, num_experts, topk): + logits = torch.randn(num_tokens, num_experts, device="cuda") + _, ids = torch.topk(torch.softmax(logits, dim=-1), k=topk, dim=-1) + return ids.to(torch.int32) + + +def test_mori_dispatch_aiter_combine_trim_regression(): + world_size = int(os.environ.get("WORLD_SIZE", "0")) + if world_size == 0: + pytest.skip( + "must be launched under torchrun (WORLD_SIZE unset) -- see " + "this file's module docstring" + ) + if world_size // GPU_PER_NODE < 2: + pytest.skip( + f"world_size={world_size} implies <2 nodes -- this test needs " + "genuine multi-node MoRI InterNodeV1 traffic (validated at " + "EP32/4 nodes) to reproduce the bug" + ) + + import mori + import torch.distributed as dist + + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + torch.accelerator.set_device_index(local_rank) + dist.init_process_group(backend="gloo", rank=rank, world_size=world_size) + cpu_group = dist.new_group(list(range(world_size)), backend="gloo") + torch._C._distributed_c10d._register_process_group("mori", cpu_group) + mori.shmem.shmem_torch_process_group_init("mori") + + from vllm.config import VllmConfig + from vllm.forward_context import set_forward_context + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + from vllm.model_executor.layers.fused_moe.config import ( + FUSED_MOE_UNQUANTIZED_CONFIG, + ) + from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( + AiterExperts, + ) + from vllm.model_executor.layers.fused_moe.prepare_finalize.mori import ( + MoriPrepareAndFinalize, + ) + + from .test_rocm_aiter_moe import _assert_aiter_supported, _shuffle_moe_weights + from .utils import make_dummy_moe_config + + _assert_aiter_supported() + + num_experts = EXPERTS_PER_RANK * world_size + mori_config = mori.ops.EpDispatchCombineConfig( + data_type=torch.bfloat16, + rank=rank, + world_size=world_size, + hidden_dim=HIDDEN, + scale_dim=0, + scale_type_size=0, + max_token_type_size=2, + max_num_inp_token_per_rank=MAX_TOK, + num_experts_per_rank=EXPERTS_PER_RANK, + num_experts_per_token=TOPK, + warp_num_per_block=16, + block_num=32, + rdma_block_num=16, + kernel_type=mori.ops.EpDispatchCombineKernelType.InterNodeV1, + gpu_per_node=GPU_PER_NODE, + ) + op = mori.ops.EpDispatchCombineOp(mori_config) + pf = MoriPrepareAndFinalize( + op, max_tokens_per_rank=MAX_TOK, num_dispatchers=world_size + ) + + torch.manual_seed(rank + 1) + w1 = torch.randn( + EXPERTS_PER_RANK, INTER * 2, HIDDEN, dtype=torch.bfloat16, device="cuda" + ) / (HIDDEN**0.5) + w2 = torch.randn( + EXPERTS_PER_RANK, HIDDEN, INTER, dtype=torch.bfloat16, device="cuda" + ) / (INTER**0.5) + w1s, w2s = _shuffle_moe_weights(w1, w2) + + moe_config = make_dummy_moe_config( + num_experts=num_experts, + num_local_experts=EXPERTS_PER_RANK, + experts_per_token=TOPK, + hidden_dim=HIDDEN, + intermediate_size=INTER, + in_dtype=torch.bfloat16, + max_num_tokens=MAX_TOK, + ) + experts = AiterExperts( + moe_config=moe_config, quant_config=FUSED_MOE_UNQUANTIZED_CONFIG + ) + vllm_config = VllmConfig() + + def round_trip(n_real, marker): + a1 = torch.zeros(MAX_TOK, HIDDEN, dtype=torch.bfloat16, device="cuda") + ids = torch.zeros(MAX_TOK, TOPK, dtype=torch.int32, device="cuda") + weights = torch.zeros(MAX_TOK, TOPK, dtype=torch.float32, device="cuda") + if n_real > 0: + a1[:n_real].fill_(marker) + ids[:n_real] = _make_topk_ids(n_real, num_experts, TOPK) + weights[:n_real] = 1.0 / TOPK + + ( + dispatch_a1, + dispatch_scale, + expert_tokens_meta, + dispatch_ids, + dispatch_weights, + ) = pf.prepare( + a1, weights, ids, num_experts, None, False, FUSED_MOE_UNQUANTIZED_CONFIG + ) + + m = dispatch_a1.shape[0] + ws13, ws2, out_shape = experts.workspace_shapes( + M=m, + N=INTER, + K=HIDDEN, + topk=TOPK, + global_num_experts=num_experts, + local_num_experts=EXPERTS_PER_RANK, + expert_tokens_meta=expert_tokens_meta, + activation=MoEActivation.SILU, + ) + workspace13 = torch.empty(ws13, dtype=torch.bfloat16, device="cuda") + workspace2 = torch.empty(ws2, dtype=torch.bfloat16, device="cuda") + aiter_out = torch.full( + out_shape, float("nan"), dtype=torch.bfloat16, device="cuda" + ) + + with set_forward_context(None, vllm_config, num_tokens=m): + experts.apply( + output=aiter_out, + hidden_states=dispatch_a1, + w1=w1s, + w2=w2s, + topk_weights=dispatch_weights, + topk_ids=dispatch_ids, + activation=MoEActivation.SILU, + global_num_experts=num_experts, + expert_map=None, + a1q_scale=dispatch_scale, + a2_scale=None, + workspace13=workspace13, + workspace2=workspace2, + expert_tokens_meta=expert_tokens_meta, + apply_router_weight_on_input=False, + ) + + final_out = torch.zeros(MAX_TOK, HIDDEN, dtype=torch.bfloat16, device="cuda") + pf.finalize(final_out, aiter_out, weights, ids, False, None) + torch.accelerator.synchronize() + dist.barrier() + return final_out[:n_real].float() if n_real > 0 else final_out[:0].float() + + # warmup, then repeated (big, small) cycles reusing the same mori_op / + # dispatch buffers -- validated on real EP32/4-node hardware to require + # >=1 prior big round before a small round's real tokens get corrupted + # (a single round right after construction is not enough). + round_trip(MAX_TOK, 5.0) + for _ in range(NUM_BIG_SMALL_ROUNDS): + round_trip(MAX_TOK, 99.0) + small_out = round_trip(4, 1.0) + assert torch.isfinite(small_out).all() + assert small_out.abs().sum() > 0, ( + "MoRI dispatch/combine returned all-zero output for real " + "tokens after buffer reuse -- the dispatch-output trim in " + "MoriPrepareAndFinalize.prepare() appears to be missing" + ) + + dist.barrier() + dist.destroy_process_group() From fc2028029ecb25f272d978f59b4290bb4afe0b68 Mon Sep 17 00:00:00 2001 From: amd-weisun Date: Wed, 23 Sep 2026 09:36:38 +0000 Subject: [PATCH 07/10] Trim explanatory comments from unit tests Move mechanism/rationale explanations out of code and into the PR description; keep only what's needed to run/understand each test. Co-Authored-By: Claude Sonnet 5 Signed-off-by: amd-weisun --- .../moe/test_mori_aiter_ep_trim_regression.py | 34 +++++-------------- tests/kernels/moe/test_rocm_aiter_moe.py | 8 ++--- 2 files changed, 11 insertions(+), 31 deletions(-) diff --git a/tests/kernels/moe/test_mori_aiter_ep_trim_regression.py b/tests/kernels/moe/test_mori_aiter_ep_trim_regression.py index c65f0c1beaad..6b60c8f6eed7 100644 --- a/tests/kernels/moe/test_mori_aiter_ep_trim_regression.py +++ b/tests/kernels/moe/test_mori_aiter_ep_trim_regression.py @@ -1,27 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Real multi-node MoRI dispatch/AITER GEMM/MoRI combine regression test for -the buffer-trim fix in rocm_aiter_moe.py / prepare_finalize/mori.py. - -Drives the three real production primitives directly (MoriPrepareAndFinalize -.prepare()/.finalize() and AiterExperts.apply()) instead of going through -FusedMoEFactory/the full FusedMoE layer -- this is intentional, see -test_moe_layer_multinode.py's docstring for why that heavier construction -path doesn't work here (older pinned vLLM install used for multi-node CI -lacks FusedMoEFactory). - -Reproduces on real 4-node/EP32 InterNodeV1 hardware: without the trim, -repeated dispatch/combine rounds on the same (reused) mori_op buffers -produce all-zero combined output for genuinely valid, real tokens (MoRI's -own reported valid-row count balloons toward the oversized buffer's full -capacity instead of the true tiny per-round count, and AITER's output for -those rows is discarded/corrupted downstream in combine()). A single round -right after construction is NOT enough to show this -- it takes at least -one prior large round to have populated/reused the buffer. With the trim -applied, results stay correct and nonzero across repeated rounds. - -Launch (mirrors test_moe_layer_multinode.py -- one torchrun per host, all -pointed at the same rendezvous, one host acting as rank 0): +"""Real multi-node MoRI dispatch + AITER GEMM + MoRI combine regression test +for the buffer-trim fix in rocm_aiter_moe.py / prepare_finalize/mori.py. + +Drives MoriPrepareAndFinalize.prepare()/.finalize() and AiterExperts.apply() +directly, bypassing FusedMoEFactory (see test_moe_layer_multinode.py). + +Launch (one torchrun per host, same rendezvous, one host as rank 0): MASTER_ADDR= MASTER_PORT=29525 \\ torchrun --nnodes 4 --nproc-per-node 8 --node-rank <0..3> \\ @@ -205,9 +190,8 @@ def round_trip(n_real, marker): return final_out[:n_real].float() if n_real > 0 else final_out[:0].float() # warmup, then repeated (big, small) cycles reusing the same mori_op / - # dispatch buffers -- validated on real EP32/4-node hardware to require - # >=1 prior big round before a small round's real tokens get corrupted - # (a single round right after construction is not enough). + # dispatch buffers -- a big round must precede a small round to trigger + # the bug. round_trip(MAX_TOK, 5.0) for _ in range(NUM_BIG_SMALL_ROUNDS): round_trip(MAX_TOK, 99.0) diff --git a/tests/kernels/moe/test_rocm_aiter_moe.py b/tests/kernels/moe/test_rocm_aiter_moe.py index c17c877ed816..c500dcd68520 100644 --- a/tests/kernels/moe/test_rocm_aiter_moe.py +++ b/tests/kernels/moe/test_rocm_aiter_moe.py @@ -10,8 +10,7 @@ - MoE-facing FP8 group-quant activation quality - deterministic routing and representative gfx942 / gfx950 coverage - ``AiterExperts.apply()`` trimming an oversized EP-dispatch-shaped buffer - down to the valid token prefix before calling AITER (regression coverage - for the buffer-size correctness defect fixed by that trim) + down to the valid token prefix before calling AITER Generic fused-MoE backend selection and non-ROCm kernel coverage live in the generic MoE test files under ``tests/kernels/moe``. @@ -971,10 +970,7 @@ def test_aiter_experts_apply_trims_oversized_dispatch_buffer( valid_tokens: int, ): """AiterExperts.apply() must trim an oversized EP-dispatch-shaped buffer - down to the valid token prefix before calling AITER, which silently - corrupts its output once the row count crosses large thresholds - (e.g. 65536), regardless of how many rows hold real data. - """ + down to the valid token prefix before calling AITER.""" import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm.config import VllmConfig from vllm.forward_context import set_forward_context From e8f8be0ccca148d9247d1254d3a179f0dadd5cf0 Mon Sep 17 00:00:00 2001 From: amd-weisun Date: Wed, 23 Sep 2026 09:59:02 +0000 Subject: [PATCH 08/10] Rename mori/aiter EP regression test, gate to ROCm Rename test_mori_aiter_ep_trim_regression.py to test_rocm_mori_moe_ep.py and skip on non-ROCm platforms (previously only gated on mori availability). Co-Authored-By: Claude Sonnet 5 Signed-off-by: amd-weisun --- ...iter_ep_trim_regression.py => test_rocm_mori_moe_ep.py} | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) rename tests/kernels/moe/{test_mori_aiter_ep_trim_regression.py => test_rocm_mori_moe_ep.py} (97%) diff --git a/tests/kernels/moe/test_mori_aiter_ep_trim_regression.py b/tests/kernels/moe/test_rocm_mori_moe_ep.py similarity index 97% rename from tests/kernels/moe/test_mori_aiter_ep_trim_regression.py rename to tests/kernels/moe/test_rocm_mori_moe_ep.py index 6b60c8f6eed7..e15c3a0b19ad 100644 --- a/tests/kernels/moe/test_mori_aiter_ep_trim_regression.py +++ b/tests/kernels/moe/test_rocm_mori_moe_ep.py @@ -11,7 +11,7 @@ MASTER_ADDR= MASTER_PORT=29525 \\ torchrun --nnodes 4 --nproc-per-node 8 --node-rank <0..3> \\ --rdzv_backend=c10d --rdzv_endpoint=:29525 \\ - -m pytest -v -s kernels/moe/test_mori_aiter_ep_trim_regression.py + -m pytest -v -s kernels/moe/test_rocm_mori_moe_ep.py """ import importlib.util @@ -20,8 +20,11 @@ import pytest import torch +from vllm.platforms import current_platform + pytestmark = pytest.mark.skipif( - importlib.util.find_spec("mori") is None, reason="mori not installed" + not current_platform.is_rocm() or importlib.util.find_spec("mori") is None, + reason="requires ROCm and mori", ) GPU_PER_NODE = 8 From 2644580780d7bb2b91501f07b63951a1e77fd2df Mon Sep 17 00:00:00 2001 From: amd-weisun Date: Thu, 24 Sep 2026 10:56:50 +0000 Subject: [PATCH 09/10] Support single-node in mori/aiter EP trim regression test Auto-select single-node (world_size == GPU_PER_NODE, split into 2 logical MoRI nodes) vs. genuine multi-node based on world_size, so the test no longer requires a real multi-host cluster. Bump rounds 20 to reliably catch the bug at single-node scale, where it's lower-probability per round than at real multi-node scale. Co-Authored-By: Claude Sonnet 5 Signed-off-by: amd-weisun --- tests/kernels/moe/test_rocm_mori_moe_ep.py | 23 +++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/tests/kernels/moe/test_rocm_mori_moe_ep.py b/tests/kernels/moe/test_rocm_mori_moe_ep.py index e15c3a0b19ad..86a02a82bde4 100644 --- a/tests/kernels/moe/test_rocm_mori_moe_ep.py +++ b/tests/kernels/moe/test_rocm_mori_moe_ep.py @@ -6,7 +6,13 @@ Drives MoriPrepareAndFinalize.prepare()/.finalize() and AiterExperts.apply() directly, bypassing FusedMoEFactory (see test_moe_layer_multinode.py). -Launch (one torchrun per host, same rendezvous, one host as rank 0): +Launch, single node (world_size == GPU_PER_NODE, no real multi-host needed): + + torchrun --nnodes 1 --nproc-per-node 8 \\ + -m pytest -v -s kernels/moe/test_rocm_mori_moe_ep.py + +Launch, genuine multi-node (one torchrun per host, same rendezvous, one host +as rank 0): MASTER_ADDR= MASTER_PORT=29525 \\ torchrun --nnodes 4 --nproc-per-node 8 --node-rank <0..3> \\ @@ -33,7 +39,7 @@ TOPK = 8 EXPERTS_PER_RANK = 8 MAX_TOK = 128 -NUM_BIG_SMALL_ROUNDS = 3 +NUM_BIG_SMALL_ROUNDS = 20 def _make_topk_ids(num_tokens, num_experts, topk): @@ -49,11 +55,10 @@ def test_mori_dispatch_aiter_combine_trim_regression(): "must be launched under torchrun (WORLD_SIZE unset) -- see " "this file's module docstring" ) - if world_size // GPU_PER_NODE < 2: + if world_size < GPU_PER_NODE: pytest.skip( - f"world_size={world_size} implies <2 nodes -- this test needs " - "genuine multi-node MoRI InterNodeV1 traffic (validated at " - "EP32/4 nodes) to reproduce the bug" + f"world_size={world_size} < {GPU_PER_NODE} -- need at least one " + "full 8-GPU node" ) import mori @@ -86,6 +91,10 @@ def test_mori_dispatch_aiter_combine_trim_regression(): _assert_aiter_supported() num_experts = EXPERTS_PER_RANK * world_size + # Single real node: split it into 2 logical MoRI "nodes" so InterNodeV1 + # traffic (and the bug it can trigger) still occurs; genuine multi-node + # already has real nodes, so use the real per-node count. + gpu_per_node = GPU_PER_NODE // 2 if world_size == GPU_PER_NODE else GPU_PER_NODE mori_config = mori.ops.EpDispatchCombineConfig( data_type=torch.bfloat16, rank=rank, @@ -101,7 +110,7 @@ def test_mori_dispatch_aiter_combine_trim_regression(): block_num=32, rdma_block_num=16, kernel_type=mori.ops.EpDispatchCombineKernelType.InterNodeV1, - gpu_per_node=GPU_PER_NODE, + gpu_per_node=gpu_per_node, ) op = mori.ops.EpDispatchCombineOp(mori_config) pf = MoriPrepareAndFinalize( From b09903509a16b3d106b667cee2e1b922b5bb27bd Mon Sep 17 00:00:00 2001 From: amd-weisun Date: Fri, 2 Oct 2026 12:33:17 +0000 Subject: [PATCH 10/10] Keep the MoRI prepare() trim only; drop the AiterExperts.apply() trim With the dispatch-output trim in MoriPrepareAndFinalize.prepare(), the trim in AiterExperts.apply() has nothing left to do for MoRI and only adds a second device-to-host sync and a second capture policy. Remove it and its helpers, drop the unit test that only exercised it, and soften the MoRI test's assertion message since either trim alone avoids the failure. Co-Authored-By: Claude Sonnet 5.5 Signed-off-by: amd-weisun --- tests/kernels/moe/test_rocm_aiter_moe.py | 144 ------------------ tests/kernels/moe/test_rocm_mori_moe_ep.py | 2 +- .../fused_moe/experts/rocm_aiter_moe.py | 83 +--------- 3 files changed, 5 insertions(+), 224 deletions(-) diff --git a/tests/kernels/moe/test_rocm_aiter_moe.py b/tests/kernels/moe/test_rocm_aiter_moe.py index ebf7ba8389a8..05692852e1d8 100644 --- a/tests/kernels/moe/test_rocm_aiter_moe.py +++ b/tests/kernels/moe/test_rocm_aiter_moe.py @@ -9,8 +9,6 @@ - gfx950-only AITER MXFP4 W4A16 MoE support, accuracy, and determinism - MoE-facing FP8 group-quant activation quality - deterministic routing and representative gfx942 / gfx950 coverage -- ``AiterExperts.apply()`` trimming an oversized EP-dispatch-shaped buffer - down to the valid token prefix before calling AITER Generic fused-MoE backend selection and non-ROCm kernel coverage live in the generic MoE test files under ``tests/kernels/moe``. @@ -1195,148 +1193,6 @@ def test_aiter_fused_moe_end_to_end( ) -# Oversized dispatch-buffer regression tests ------------------------------- - - -@pytest.mark.parametrize( - "buffer_rows,valid_tokens", - [ - pytest.param(65536, 8, id="buffer_at_aiter_corruption_threshold"), - pytest.param(128, 8, id="small_buffer_control"), - ], -) -def test_aiter_experts_apply_trims_oversized_dispatch_buffer( - buffer_rows: int, - valid_tokens: int, -): - """AiterExperts.apply() must trim an oversized EP-dispatch-shaped buffer - down to the valid token prefix before calling AITER.""" - import vllm.model_executor.layers.fused_moe.modular_kernel as mk - from vllm.config import VllmConfig - from vllm.forward_context import set_forward_context - from vllm.model_executor.layers.fused_moe.activation import MoEActivation - from vllm.model_executor.layers.fused_moe.config import ( - FUSED_MOE_UNQUANTIZED_CONFIG, - ) - from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( - AiterExperts, - ) - - from .utils import make_dummy_moe_config - - _assert_aiter_supported() - torch.set_default_device("cuda") - set_random_seed(0) - - hidden_dim = 128 - intermediate_dim = 256 - num_experts = 4 - topk = 2 - activation = MoEActivation.SILU - - moe_config = make_dummy_moe_config( - num_experts=num_experts, - experts_per_token=topk, - hidden_dim=hidden_dim, - intermediate_size=intermediate_dim, - in_dtype=torch.bfloat16, - max_num_tokens=buffer_rows, - ) - experts = AiterExperts( - moe_config=moe_config, - quant_config=FUSED_MOE_UNQUANTIZED_CONFIG, - ) - - w1 = torch.randn( - num_experts, intermediate_dim * 2, hidden_dim, dtype=torch.bfloat16 - ) / math.sqrt(hidden_dim) - w2 = torch.randn( - num_experts, hidden_dim, intermediate_dim, dtype=torch.bfloat16 - ) / math.sqrt(intermediate_dim) - w1_shuffled, w2_shuffled = _shuffle_moe_weights(w1, w2) - - # Only the first `valid_tokens` rows carry real data; the rest mimics - # the unreceived tail of a real EP dispatch buffer. - hidden_states = torch.zeros(buffer_rows, hidden_dim, dtype=torch.bfloat16) - topk_weights = torch.zeros(buffer_rows, topk, dtype=torch.float32) - topk_ids = torch.zeros(buffer_rows, topk, dtype=torch.int32) - - hidden_states[:valid_tokens] = torch.randn( - valid_tokens, hidden_dim, dtype=torch.bfloat16 - ) - valid_weights = torch.rand(valid_tokens, topk, dtype=torch.float32) - topk_weights[:valid_tokens] = valid_weights / valid_weights.sum( - dim=-1, keepdim=True - ) - topk_ids[:valid_tokens] = _make_topk_ids(valid_tokens, num_experts, topk) - - expert_tokens_meta = mk.ExpertTokensMetadata( - expert_num_tokens=torch.tensor([valid_tokens], dtype=torch.int32), - expert_num_tokens_cpu=None, - ) - - workspace13_shape, workspace2_shape, output_shape = experts.workspace_shapes( - M=buffer_rows, - N=intermediate_dim, - K=hidden_dim, - topk=topk, - global_num_experts=num_experts, - local_num_experts=num_experts, - expert_tokens_meta=expert_tokens_meta, - activation=activation, - ) - workspace13 = torch.empty(workspace13_shape, dtype=torch.bfloat16) - workspace2 = torch.empty(workspace2_shape, dtype=torch.bfloat16) - output = torch.full(output_shape, torch.nan, dtype=torch.bfloat16) - - vllm_config = VllmConfig() - with set_forward_context(None, vllm_config, num_tokens=buffer_rows): - experts.apply( - output=output, - hidden_states=hidden_states, - w1=w1_shuffled, - w2=w2_shuffled, - topk_weights=topk_weights, - topk_ids=topk_ids, - activation=activation, - global_num_experts=num_experts, - expert_map=None, - a1q_scale=None, - a2_scale=None, - workspace13=workspace13, - workspace2=workspace2, - expert_tokens_meta=expert_tokens_meta, - apply_router_weight_on_input=False, - ) - - assert torch.isfinite(output).all() - - valid_out = output[:valid_tokens] - assert valid_out.abs().sum() > 0, ( - "AiterExperts.apply() returned all-zero output for valid tokens -- " - "the oversized-buffer trim appears to be missing or broken" - ) - - ref_out = ref_moe_forward( - hidden_states[:valid_tokens], - w1, - w2, - topk_weights[:valid_tokens], - topk_ids[:valid_tokens], - activation="silu", - ) - _assert_close_budget( - valid_out.float(), - ref_out, - label=( - f"oversized_dispatch_buffer buffer_rows={buffer_rows} " - f"valid_tokens={valid_tokens}" - ), - atol=0.05, - rtol=0.0, - ) - - # Arch-specific tests ----------------------------------------------------- diff --git a/tests/kernels/moe/test_rocm_mori_moe_ep.py b/tests/kernels/moe/test_rocm_mori_moe_ep.py index 86a02a82bde4..4b26bf8c78fe 100644 --- a/tests/kernels/moe/test_rocm_mori_moe_ep.py +++ b/tests/kernels/moe/test_rocm_mori_moe_ep.py @@ -212,7 +212,7 @@ def round_trip(n_real, marker): assert small_out.abs().sum() > 0, ( "MoRI dispatch/combine returned all-zero output for real " "tokens after buffer reuse -- the dispatch-output trim in " - "MoriPrepareAndFinalize.prepare() appears to be missing" + "MoriPrepareAndFinalize.prepare() may be missing" ) dist.barrier() diff --git a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py index 9ae5939038e3..07eddbf4a092 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py @@ -8,7 +8,6 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm._aiter_ops import rocm_aiter_ops -from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import ( FUSED_MOE_UNQUANTIZED_CONFIG, @@ -30,7 +29,6 @@ kMxfp4Dynamic, kMxfp4Static, ) -from vllm.platforms import current_platform if TYPE_CHECKING: from aiter import ActivationType @@ -59,35 +57,6 @@ class ActivationMethod(IntEnum): GELU = 1 -def _is_uniform_full_graph_batch() -> bool: - # Under full-cudagraph capture/replay every rank is padded to the same - # static token count -- used below to pick a capture-safe trim bound - # instead of a live `.item()` count. - from vllm.config import CUDAGraphMode - - if not is_forward_context_available(): - return False - forward_context = get_forward_context() - batch_descriptor = forward_context.batch_descriptor - return ( - forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL - and batch_descriptor is not None - and batch_descriptor.uniform - ) - - -def _is_stream_capturing() -> bool: - # torch.cuda.is_current_stream_capturing() is unavailable on non-CUDA - # (XPU) torch; guard the same way fused_batched_moe.py does. - try: - return ( - current_platform.is_cuda_alike() - and torch.cuda.is_current_stream_capturing() - ) - except Exception: - return False - - aiter_topK_meta_data: tuple[torch.Tensor, torch.Tensor] | None = None @@ -600,66 +569,22 @@ def apply( else: num_local_tokens = None - orig_num_tokens = hidden_states.shape[0] - if num_local_tokens is None: - valid_num_tokens = orig_num_tokens - elif _is_uniform_full_graph_batch(): - batch_descriptor = get_forward_context().batch_descriptor - assert batch_descriptor is not None - dispatch_group_size = ( - self.moe_config.ep_size - if self.moe_config.use_ep - else self.moe_config.dp_size - ) - valid_num_tokens = min( - batch_descriptor.num_tokens - * self.moe_config.experts_per_token - * dispatch_group_size, - orig_num_tokens, - ) - elif _is_stream_capturing(): - valid_num_tokens = orig_num_tokens - else: - valid_num_tokens = int(num_local_tokens[0].item()) - needs_trim = valid_num_tokens < orig_num_tokens - if needs_trim: - hidden_states_in = hidden_states[:valid_num_tokens] - topk_weights_in = topk_weights[:valid_num_tokens] - topk_ids_in = topk_ids[:valid_num_tokens] - if ( - a1q_scale is not None - and a1q_scale.dim() > 0 - and a1q_scale.shape[0] == orig_num_tokens - ): - a1q_scale_in = a1q_scale[:valid_num_tokens] - else: - a1q_scale_in = a1q_scale - else: - hidden_states_in = hidden_states - topk_weights_in = topk_weights - topk_ids_in = topk_ids - a1q_scale_in = a1q_scale - result = rocm_aiter_fused_experts( - hidden_states=hidden_states_in, + hidden_states=hidden_states, w1=w1, w2=w2, - topk_weights=topk_weights_in, - topk_ids=topk_ids_in, + topk_weights=topk_weights, + topk_ids=topk_ids, activation=activation, apply_router_weight_on_input=apply_router_weight_on_input, expert_mask=expert_map, quant_config=self.quant_config, moe_config=self.moe_config, - a1q_scale=a1q_scale_in, + a1q_scale=a1q_scale, num_local_tokens=num_local_tokens, output_dtype=output.dtype, moe_sorting_dispatch_policy=rocm_aiter_ops.get_moe_dispatch_policy(), ) - if needs_trim: - output.zero_() - output[:valid_num_tokens] = result - return # avoid redundant copy when output is a view of the result if ( output.shape == result.shape