vLLM/tests/kernels/test_gate_linear_rocm_dispatch.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Platform-agnostic eligibility tests for GateLinear's ROCm router GEMMs.

These assert the dispatch flags directly, so they run device-free by mocking
the platform predicates. ``allow_cublas_router_gemm`` selects the
bf16xbf16->fp32 ``torch.mm`` epilogue, while ``allow_fp32_router_gemm`` selects
the gfx950 low-M kernel with fp32 weights and output.

The ROCm branch is guarded on ``not bias`` because ``torch.mm`` has no bias
term; a biased gate must fall back so the bias is not silently dropped.
"""

import inspect

import pytest
import torch

import vllm.model_executor.layers.fused_moe.router.gate_linear as gate_linear_mod
from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear


def _make_gate(
    monkeypatch,
    *,
    is_rocm: bool,
    is_cuda: bool = False,
    bias: bool = False,
    params_dtype: torch.dtype = torch.bfloat16,
    out_dtype: torch.dtype | None = torch.float32,
    input_size: int = 2048,
    output_size: int = 64,
    on_gfx950: bool = False,
    parallel_world_size: int = 1,
) -> GateLinear:
    """Build a GateLinear with platform predicates mocked, no GPU needed."""
    for target in (
        "vllm.model_executor.layers.linear",
        "vllm.model_executor.parameter",
    ):
        monkeypatch.setattr(
            f"{target}.get_tensor_model_parallel_rank",
            lambda: 0,
        )
        monkeypatch.setattr(
            f"{target}.get_tensor_model_parallel_world_size",
            lambda: parallel_world_size,
        )

    platform = gate_linear_mod.current_platform
    monkeypatch.setattr(platform, "is_cuda", lambda: is_cuda)
    monkeypatch.setattr(platform, "is_rocm", lambda: is_rocm)
    monkeypatch.setattr(platform, "is_device_capability", lambda *a, **k: False)
    monkeypatch.setattr(platform, "is_device_capability_family", lambda *a, **k: False)
    if is_rocm:
        import vllm.platforms.rocm as rocm_platform

        monkeypatch.setattr(rocm_platform, "on_gfx950", lambda: on_gfx950)

    return GateLinear(
        input_size=input_size,
        output_size=output_size,
        bias=bias,
        out_dtype=out_dtype,
        params_dtype=params_dtype,
    )


def test_rocm_no_bias_bf16_fp32_enables_fused_gemm(monkeypatch):
    gate = _make_gate(monkeypatch, is_rocm=True, bias=False)
    assert not gate.allow_specialized_router_gemm
    assert gate.allow_cublas_router_gemm


def test_rocm_bias_disables_fused_gemm(monkeypatch):
    # torch.mm cannot add a bias, so a biased gate must not take the fused path.
    gate = _make_gate(monkeypatch, is_rocm=True, bias=True)
    assert not gate.allow_cublas_router_gemm


def test_rocm_fp32_weight_disables_fused_gemm(monkeypatch):
    gate = _make_gate(monkeypatch, is_rocm=True, params_dtype=torch.float32)
    assert not gate.allow_cublas_router_gemm


def test_rocm_non_fp32_out_dtype_disables_fused_gemm(monkeypatch):
    gate = _make_gate(monkeypatch, is_rocm=True, out_dtype=torch.bfloat16)
    assert not gate.allow_cublas_router_gemm


def test_non_rocm_non_cuda_disables_fused_gemm(monkeypatch):
    # Neither the CUDA specialized path nor the ROCm branch applies.
    gate = _make_gate(monkeypatch, is_rocm=False, is_cuda=False)
    assert not gate.allow_cublas_router_gemm


def test_rocm_set_out_dtype_enables_fused_gemm(monkeypatch):
    gate = _make_gate(monkeypatch, is_rocm=True, bias=False, out_dtype=None)
    assert not gate.allow_cublas_router_gemm
    gate.set_out_dtype(torch.float32)
    assert gate.allow_cublas_router_gemm


def test_rocm_set_out_dtype_respects_bias_guard(monkeypatch):
    gate = _make_gate(monkeypatch, is_rocm=True, bias=True, out_dtype=None)
    gate.set_out_dtype(torch.float32)
    assert not gate.allow_cublas_router_gemm


@pytest.mark.parametrize(
    ("input_size", "output_size"),
    [(3072, 256), (4096, 8), (6144, 128), (2048, 64), (2560, 512), (7168, 384)],
)
def test_rocm_gfx950_enables_fp32_router_gemm(
    monkeypatch, input_size: int, output_size: int
) -> None:
    gate = _make_gate(
        monkeypatch,
        is_rocm=True,
        params_dtype=torch.float32,
        input_size=input_size,
        output_size=output_size,
        on_gfx950=True,
    )
    assert gate.allow_fp32_router_gemm
    assert not gate.allow_cublas_router_gemm


@pytest.mark.parametrize("ep_size", [2, 4, 8])
def test_rocm_fp32_router_gemm_is_replicated_across_ep_sizes(
    monkeypatch, ep_size: int
) -> None:
    gate = _make_gate(
        monkeypatch,
        is_rocm=True,
        params_dtype=torch.float32,
        input_size=6144,
        output_size=128,
        on_gfx950=True,
        parallel_world_size=ep_size,
    )
    assert gate.weight.shape == (128, 6144)
    assert gate.allow_fp32_router_gemm


@pytest.mark.parametrize(
    ("input_size", "output_size", "params_dtype", "bias", "on_gfx950"),
    [
        pytest.param(6144, 128, torch.float32, False, False, id="gfx942"),
        pytest.param(6144, 128, torch.bfloat16, False, True, id="bf16-weight"),
        pytest.param(6144, 128, torch.float32, True, True, id="bias"),
    ],
)
def test_rocm_fp32_router_gemm_rejects_unsupported_configs(
    monkeypatch,
    input_size: int,
    output_size: int,
    params_dtype: torch.dtype,
    bias: bool,
    on_gfx950: bool,
) -> None:
    gate = _make_gate(
        monkeypatch,
        is_rocm=True,
        params_dtype=params_dtype,
        bias=bias,
        input_size=input_size,
        output_size=output_size,
        on_gfx950=on_gfx950,
    )
    assert not gate.allow_fp32_router_gemm


# ---------------------------------------------------------------------------
# ROCm BF16x3 tier (fp32 router weights)
# ---------------------------------------------------------------------------
def _make_bf16x3_gate(monkeypatch, *, supported=True, **kwargs):
    """Build a gate with the ROCm bf16x3 platform check mocked."""
    from vllm.model_executor.layers.fused_moe.router import (
        bf16x3_router_gemm_rocm as rocm_bf16x3,
    )

    monkeypatch.setattr(rocm_bf16x3, "platform_supported", lambda: supported)
    kwargs.setdefault("params_dtype", torch.float32)
    return _make_gate(monkeypatch, is_rocm=True, **kwargs)


def test_rocm_bf16x3_enabled_for_fp32_weight(monkeypatch):
    gate = _make_bf16x3_gate(monkeypatch)
    assert gate.allow_rocm_bf16x3_router_gemm
    # Built by process_weights_after_loading, not lazily in forward().
    # Registered up front but still None, so named_buffers() skips it.
    assert gate._bf16x3_weight is None
    assert "_bf16x3_weight" in gate._buffers


def test_rocm_bf16x3_disabled_by_bias(monkeypatch):
    # torch.mm has no bias term.
    gate = _make_bf16x3_gate(monkeypatch, bias=True)
    assert not gate.allow_rocm_bf16x3_router_gemm


def test_rocm_bf16x3_disabled_for_bf16_weight(monkeypatch):
    # A bf16 weight belongs to the cuBLAS tier; there is nothing to split.
    gate = _make_bf16x3_gate(monkeypatch, params_dtype=torch.bfloat16)
    assert not gate.allow_rocm_bf16x3_router_gemm


def test_rocm_bf16x3_disabled_for_non_fp32_out_dtype(monkeypatch):
    gate = _make_bf16x3_gate(monkeypatch, out_dtype=torch.bfloat16)
    assert not gate.allow_rocm_bf16x3_router_gemm


def test_rocm_bf16x3_disabled_on_unsupported_arch(monkeypatch):
    # Only gfx950 is benchmarked; gfx942 has a different fp32:bf16 MFMA ratio.
    gate = _make_bf16x3_gate(monkeypatch, supported=False)
    assert not gate.allow_rocm_bf16x3_router_gemm
    assert not gate._rocm_bf16x3_weight_eligible


def test_rocm_bf16x3_no_split_built_when_tier_disabled(monkeypatch):
    """An ineligible out_dtype must not allocate a split that is never read;
    it is 1.5x the router weight on every layer."""
    gate = _make_bf16x3_gate(monkeypatch, out_dtype=torch.bfloat16)
    gate.quant_method.process_weights_after_loading(gate)
    assert gate._bf16x3_weight is None


def test_rocm_bf16x3_set_out_dtype_enables_tier(monkeypatch):
    gate = _make_bf16x3_gate(monkeypatch, out_dtype=None)
    assert not gate.allow_rocm_bf16x3_router_gemm
    gate.set_out_dtype(torch.float32)
    assert gate.allow_rocm_bf16x3_router_gemm


def test_rocm_bf16x3_dispatch_is_a_custom_op():
    """The num_tokens branch must stay opaque to Dynamo; see the op's impl."""
    assert hasattr(torch.ops.vllm, "rocm_bf16x3_router_gemm_dispatch")
    x = torch.zeros(8192, 6144, dtype=torch.bfloat16, device="meta")
    w = torch.zeros(128, 6144, dtype=torch.float32, device="meta")
    split = torch.zeros(3, 128, 6144, dtype=torch.bfloat16, device="meta")
    out = torch.ops.vllm.rocm_bf16x3_router_gemm_dispatch(x, w, split)
    assert out.shape == (8192, 128)
    assert out.dtype == torch.float32


def test_rocm_bf16x3_is_reachable_for_low_m_shapes(monkeypatch):
    """The gfx950 low-M fp32 tier claims these shapes too and returns
    unconditionally, so bf16x3 is only reachable if checked first."""
    gate = _make_bf16x3_gate(
        monkeypatch, input_size=6144, output_size=128, on_gfx950=True
    )
    assert gate.allow_fp32_router_gemm
    assert gate.allow_rocm_bf16x3_router_gemm

    source = inspect.getsource(type(gate).forward)
    assert source.index("allow_rocm_bf16x3_router_gemm") < source.index(
        "allow_fp32_router_gemm"
    )