[ROCm][BugFix] Trim MoE GEMM input to valid token prefix

tests/kernels/moe/test_rocm_mori_moe_ep.py

@@ -0,0 +1,219 @@
+# 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 MoriPrepareAndFinalize.prepare()/.finalize() and AiterExperts.apply()
+directly, bypassing FusedMoEFactory (see test_moe_layer_multinode.py).
+
+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=<node0-ip> MASTER_PORT=29525 \\
+    torchrun --nnodes 4 --nproc-per-node 8 --node-rank <0..3> \\
+      --rdzv_backend=c10d --rdzv_endpoint=<node0-ip>:29525 \\
+      -m pytest -v -s kernels/moe/test_rocm_mori_moe_ep.py
+"""
+
+import importlib.util
+import os
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_rocm() or importlib.util.find_spec("mori") is None,
+    reason="requires ROCm and mori",
+)
+
+GPU_PER_NODE = 8
+HIDDEN = 4096
+INTER = 1024
+TOPK = 8
+EXPERTS_PER_RANK = 8
+MAX_TOK = 128
+NUM_BIG_SMALL_ROUNDS = 20
+
+
+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:
+        pytest.skip(
+            f"world_size={world_size} < {GPU_PER_NODE} -- need at least one "
+            "full 8-GPU node"
+        )
+
+    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
+    # 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,
+        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 -- 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)
+        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() may be missing"
+        )
+
+    dist.barrier()
+    dist.destroy_process_group()

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

@@ -12,6 +12,47 @@
 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."""
 
@@ -107,6 +148,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
         )