| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | |
| 8 | import vllm.config |
| 9 | from tests.compile.backend import TestBackend |
| 10 | from vllm._aiter_ops import rocm_aiter_ops |
| 11 | from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass |
| 12 | from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass |
| 13 | from vllm.config import ( |
| 14 | CompilationConfig, |
| 15 | CompilationMode, |
| 16 | ModelConfig, |
| 17 | PassConfig, |
| 18 | VllmConfig, |
| 19 | ) |
| 20 | from vllm.model_executor.layers.layernorm import RMSNorm |
| 21 | from vllm.model_executor.layers.utils import rocm_unquantized_gemm |
| 22 | |
| 23 | |
| 24 | class TestModel(torch.nn.Module): |
| 25 | def __init__( |
| 26 | self, |
| 27 | num_layers: int, |
| 28 | hidden_size: int, |
| 29 | num_local_experts: int, |
| 30 | x_pad_to_multiple: int, |
| 31 | ): |
| 32 | super().__init__() |
| 33 | self.num_layers = num_layers |
| 34 | self.hidden_size = hidden_size |
| 35 | self.x_pad_to_multiple = x_pad_to_multiple |
| 36 | self.pad_dim = x_pad_to_multiple - (hidden_size % x_pad_to_multiple) |
| 37 | |
| 38 | self.norm = [RMSNorm(hidden_size, eps=1e-5) for _ in range(num_layers)] |
| 39 | self.router = [ |
| 40 | torch.nn.Linear(hidden_size, num_local_experts) for _ in range(4) |
| 41 | ] |
| 42 | |
| 43 | def forward(self, x): |
| 44 | # avoid having graph input be an arg to a pattern directly |
| 45 | x = resid = torch.relu(x) |
| 46 | all_router_logits = [] |
| 47 | for layer in range(self.num_layers): |
| 48 | x = x[:, : self.hidden_size] |
| 49 | x, resid = self.norm[layer](x, resid) |
| 50 | router_logits = rocm_unquantized_gemm( |
| 51 | self, x, self.router[layer].weight, self.router[layer].bias |
| 52 | ) |
| 53 | x = torch.nn.functional.pad( |
| 54 | x, (0, self.pad_dim), mode="constant", value=0.0 |
| 55 | ) |
| 56 | all_router_logits.append(router_logits) |
| 57 | |
| 58 | return x, resid, *all_router_logits |
| 59 | |
| 60 | def ops_in_model_before(self): |
| 61 | return [ |
| 62 | torch.ops.vllm_ir.fused_add_rms_norm, |
| 63 | torch.ops.aten.constant_pad_nd, |
| 64 | ] |
| 65 | |
| 66 | def ops_in_model_after(self): |
| 67 | return [rocm_aiter_ops.get_triton_add_rmsnorm_pad_op()] |
| 68 | |
| 69 | |
| 70 | @pytest.mark.parametrize("dtype", [torch.bfloat16]) |
| 71 | @pytest.mark.parametrize("num_layers", [3]) |
| 72 | @pytest.mark.parametrize("hidden_size", [2880]) |
| 73 | @pytest.mark.parametrize("num_local_experts", [128]) |
| 74 | @pytest.mark.parametrize("x_pad_to_multiple", [256]) |
| 75 | @pytest.mark.skip( |
| 76 | reason="Skipping for now because of the accuracy issue. See: https://github.com/ROCm/aiter/issues/2614" |
| 77 | ) |
| 78 | def test_fuse_act_padding( |
| 79 | dtype: torch.dtype, |
| 80 | num_layers: int, |
| 81 | hidden_size: int, |
| 82 | num_local_experts: int, |
| 83 | x_pad_to_multiple: int, |
| 84 | monkeypatch: pytest.MonkeyPatch, |
| 85 | ): |
| 86 | vllm_config = VllmConfig( |
| 87 | model_config=ModelConfig(dtype=dtype), |
| 88 | compilation_config=CompilationConfig( |
| 89 | mode=CompilationMode.VLLM_COMPILE, |
| 90 | custom_ops=["+rms_norm"], |
| 91 | pass_config=PassConfig(fuse_act_padding=True, eliminate_noops=True), |
| 92 | ), |
| 93 | ) |
| 94 | |
| 95 | with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m: |
| 96 | from vllm.compilation.passes.fusion.rocm_aiter_fusion import ( |
| 97 | RocmAiterTritonAddRMSNormPadFusionPass, |
| 98 | ) |
| 99 | |
| 100 | torch.set_default_device("cuda") |
| 101 | torch.set_default_dtype(dtype) |
| 102 | torch.manual_seed(1) |
| 103 | |
| 104 | m.setenv("VLLM_ROCM_USE_AITER", "1") |
| 105 | rocm_aiter_ops.refresh_env_variables() |
| 106 | |
| 107 | fusion_pass = RocmAiterTritonAddRMSNormPadFusionPass(vllm_config) |
| 108 | passes = [ |
| 109 | NoOpEliminationPass(vllm_config), |
| 110 | fusion_pass, |
| 111 | PostCleanupPass(vllm_config), |
| 112 | ] |
| 113 | backend = TestBackend(*passes) |
| 114 | model = TestModel(num_layers, hidden_size, num_local_experts, x_pad_to_multiple) |
| 115 | |
| 116 | x = torch.rand(1, hidden_size) |
| 117 | torch._dynamo.mark_dynamic(x, 0) |
| 118 | outputs_unfused = model(x) |
| 119 | |
| 120 | model_fused = torch.compile(model, backend=backend) |
| 121 | outputs_fused = model_fused(x) |
| 122 | |
| 123 | torch.testing.assert_close(outputs_unfused, outputs_fused) |
| 124 | |
| 125 | assert fusion_pass.matched_count == num_layers |
| 126 | |
| 127 | backend.check_before_ops(model.ops_in_model_before()) |
| 128 | backend.check_after_ops(model.ops_in_model_after()) |