Repository: vllm-project/vllm
File: tests/compile/test_fuse_act_padding.py (Python)

Code before:
from vllm.config import (
    CompilationConfig,
    CompilationMode,
    ModelConfig,
    PassConfig,
    VllmConfig,
)
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.utils import rocm_unquantized_gemm

from .backend import TestBackend


class TestModel(torch.nn.Module):
    def __init__(
        self, hidden_size: int, num_local_experts: int, x_pad_to_multiple: int
    ):
        super().__init__()
        self.norm = RMSNorm(hidden_size, eps=1e-5)
        self.router = torch.nn.Linear(hidden_size, num_local_experts)
        self.x_pad_to_multiple = x_pad_to_multiple
        self.pad_dim = x_pad_to_multiple - (hidden_size % x_pad_to_multiple)

    def forward(self, x, residual):
        x, residual = self.norm(x, residual)
        router_logits = rocm_unquantized_gemm(
            self, x, self.router.weight, self.router.bias
        )
        x = torch.nn.functional.pad(x, (0, self.pad_dim), mode="constant", value=0.0)
        return x, residual, router_logits

    def ops_in_model_before(self):
        return [
            rocm_aiter_ops.get_rmsnorm_fused_add_op(),
            torch.ops.aten.constant_pad_nd,
        ]

    def ops_in_model_after(self):
        return [rocm_aiter_ops.get_triton_add_rmsnorm_pad_op()]


@pytest.mark.parametrize("dtype", [torch.bfloat16])
@pytest.mark.parametrize("hidden_size", [2880])
@pytest.mark.parametrize("num_local_experts", [128])
@pytest.mark.parametrize("x_pad_to_multiple", [256])
@pytest.mark.skipif(
    not is_aiter_found_and_supported(),
    reason="Only test on ROCm with AITER installed and supported",
)
def test_fuse_act_padding(
    dtype: torch.dtype,

Review comment (bug):
Let's replicate a few layers here to simulate residual reuse, as that can sometimes cause matching bugs.

Code after:
    CompilationConfig,
    CompilationMode,
    ModelConfig,
    PassConfig,
    VllmConfig,
)
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.utils import rocm_unquantized_gemm

from .backend import TestBackend


class TestModel(torch.nn.Module):
    def __init__(
        self,
        num_layers: int,
        hidden_size: int,
        num_local_experts: int,
        x_pad_to_multiple: int,
    ):
        super().__init__()
        self.num_layers = num_layers
        self.hidden_size = hidden_size
        self.x_pad_to_multiple = x_pad_to_multiple
        self.pad_dim = x_pad_to_multiple - (hidden_size % x_pad_to_multiple)

        self.norm = [RMSNorm(hidden_size, eps=1e-5) for _ in range(num_layers)]
        self.router = [
            torch.nn.Linear(hidden_size, num_local_experts) for _ in range(4)
        ]

    def forward(self, x):
        # avoid having graph input be an arg to a pattern directly
        x = resid = torch.relu(x)
        all_router_logits = []
        for layer in range(self.num_layers):
            x = x[:, : self.hidden_size]
            x, resid = self.norm[layer](x, resid)
            router_logits = rocm_unquantized_gemm(
                self, x, self.router[layer].weight, self.router[layer].bias
            )
            x = torch.nn.functional.pad(
                x, (0, self.pad_dim), mode="constant", value=0.0
            )
            all_router_logits.append(router_logits)

        return x, resid, *all_router_logits

    def ops_in_model_before(self):
        return [
            rocm_aiter_ops.get_rmsnorm_fused_add_op(),