| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for IR inplace functionalization pass integration. |
| 4 | |
| 5 | This test suite verifies that the inplace functionalization pass, lowering pass, |
| 6 | and clone cleanup pass work together correctly with donated buffer tracking. |
| 7 | """ |
| 8 | |
| 9 | from collections.abc import Callable |
| 10 | |
| 11 | import pytest |
| 12 | import torch |
| 13 | import torch._dynamo.exc |
| 14 | from torch import nn |
| 15 | |
| 16 | import vllm.kernels # noqa: F401 to register kernels |
| 17 | from vllm.compilation.passes.inductor_pass import InductorPass, get_pass_context |
| 18 | from vllm.compilation.passes.ir.clone_elimination import ( |
| 19 | UnsafeCloneEliminationPass, |
| 20 | ) |
| 21 | from vllm.compilation.passes.ir.inplace_functionalization import ( |
| 22 | VllmIRInplaceFunctionalizationPass, |
| 23 | ) |
| 24 | from vllm.compilation.passes.ir.lowering_pass import VllmIRLoweringPass |
| 25 | from vllm.config import VllmConfig |
| 26 | from vllm.ir import ops |
| 27 | from vllm.platforms import current_platform |
| 28 | from vllm.triton_utils import HAS_TRITON, tl, triton |
| 29 | |
| 30 | from ...backend import TestBackend |
| 31 | |
| 32 | |
| 33 | class StoreDonationInfoPass(InductorPass): |
| 34 | def __init__(self): |
| 35 | self.donated_input_ids_sets: list[set[int]] = [] |
| 36 | |
| 37 | def __call__(self, *args, **kwargs): |
| 38 | ctx = get_pass_context() |
| 39 | self.donated_input_ids_sets += [ctx.donated_input_ids] |
| 40 | |
| 41 | |
| 42 | class MaybeInplaceModel(nn.Module): |
| 43 | """Model using only maybe_inplace variants.""" |
| 44 | |
| 45 | def __init__(self, hidden_size=16): |
| 46 | super().__init__() |
| 47 | self.weight1 = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16)) |
| 48 | self.weight2 = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16)) |
| 49 | |
| 50 | def forward( |
| 51 | self, x: torch.Tensor, residual1: torch.Tensor, residual2: torch.Tensor |
| 52 | ): |
| 53 | # First maybe_inplace - x & residual1 are donated |
| 54 | x_normed1, residual_out1 = ops.fused_add_rms_norm.maybe_inplace( |
| 55 | x, residual1, self.weight1, 1e-5 |
| 56 | ) |
| 57 | # Second maybe_inplace - residual2 is donated |
| 58 | x_normed2, residual_out2 = ops.fused_add_rms_norm.maybe_inplace( |
| 59 | x_normed1, residual2, self.weight2, 1e-5 |
| 60 | ) |
| 61 | return x_normed2, residual_out1, residual_out2 |
| 62 | |
| 63 | |
| 64 | class FunctionalModel(nn.Module): |
| 65 | """Model using only functional (default) variants.""" |
| 66 | |
| 67 | def __init__(self, hidden_size=16): |
| 68 | super().__init__() |
| 69 | self.weight1 = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16)) |
| 70 | self.weight2 = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16)) |
| 71 | |
| 72 | def forward( |
| 73 | self, x: torch.Tensor, residual1: torch.Tensor, residual2: torch.Tensor |
| 74 | ): |
| 75 | # First functional - no donation |
| 76 | x_normed1, residual_out1 = ops.fused_add_rms_norm( |
| 77 | x, residual1, self.weight1, 1e-5 |
| 78 | ) |
| 79 | # Second functional - no donation |
| 80 | x_normed2, residual_out2 = ops.fused_add_rms_norm( |
| 81 | x_normed1, residual2, self.weight2, 1e-5 |
| 82 | ) |
| 83 | return x_normed2, residual_out1, residual_out2 |
| 84 | |
| 85 | |
| 86 | class MixedModel(nn.Module): |
| 87 | """Model mixing maybe_inplace and functional variants.""" |
| 88 | |
| 89 | def __init__(self, hidden_size=16): |
| 90 | super().__init__() |
| 91 | self.weight1 = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16)) |
| 92 | self.weight2 = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16)) |
| 93 | |
| 94 | def forward( |
| 95 | self, x: torch.Tensor, residual1: torch.Tensor, residual2: torch.Tensor |
| 96 | ): |
| 97 | # First maybe_inplace - x & residual1 are donated |
| 98 | x_normed1, residual_out1 = ops.fused_add_rms_norm.maybe_inplace( |
| 99 | x, residual1, self.weight1, 1e-5 |
| 100 | ) |
| 101 | # Second functional - no donation, x_normed1 must be preserved as it's returned |
| 102 | x_normed2, residual_out2 = ops.fused_add_rms_norm( |
| 103 | x_normed1, residual2, self.weight2, 1e-5 |
| 104 | ) |
| 105 | # Return both to prevent x_normed1 from being optimized away |
| 106 | return x_normed1, x_normed2, residual_out1, residual_out2 |
| 107 | |
| 108 | |
| 109 | class ModelWithTritonAfterMaybeInplace(nn.Module): |
| 110 | """Model using maybe_inplace followed by a Triton kernel. |
| 111 | Test clone elimination can handle Triton in the graph |
| 112 | """ |
| 113 | |
| 114 | def __init__(self, hidden_size=16): |
| 115 | super().__init__() |
| 116 | self.weight = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16)) |
| 117 | |
| 118 | @triton.jit |
| 119 | def _triton_add_kernel( |
| 120 | x_ptr, |
| 121 | y_ptr, |
| 122 | n_elements, |
| 123 | BLOCK_SIZE: tl.constexpr, |
| 124 | ): |
| 125 | pid = tl.program_id(axis=0) |
| 126 | block_start = pid * BLOCK_SIZE |
| 127 | offsets = block_start + tl.arange(0, BLOCK_SIZE) |
| 128 | mask = offsets < n_elements |
| 129 | x = tl.load(x_ptr + offsets, mask=mask) |
| 130 | y = x + 0.1 |
| 131 | tl.store(y_ptr + offsets, y, mask=mask) |
| 132 | |
| 133 | def triton_add(x: torch.Tensor) -> torch.Tensor: |
| 134 | """Simple Triton add kernel.""" |
| 135 | y = torch.empty_like(x) |
| 136 | n_elements = x.numel() |
| 137 | grid = (triton.cdiv(n_elements, 256),) |
| 138 | _triton_add_kernel[grid](x, y, n_elements, BLOCK_SIZE=256) |
| 139 | return y |
| 140 | |
| 141 | self.triton_add = triton_add |
| 142 | |
| 143 | def forward(self, x: torch.Tensor, residual: torch.Tensor, residual2: torch.Tensor): |
| 144 | x_normed, residual_out = ops.fused_add_rms_norm.maybe_inplace( |
| 145 | x, residual, self.weight, 1e-5 |
| 146 | ) |
| 147 | |
| 148 | x_processed = self.triton_add(x_normed) |
| 149 | |
| 150 | # x_processed does not need to be cloned, residual2 does |
| 151 | x_normed2, residual_out2 = ops.fused_add_rms_norm( |
| 152 | x_processed, residual2, self.weight, 1e-5 |
| 153 | ) |
| 154 | return x_normed2, residual_out2 |
| 155 | |
| 156 | |
| 157 | skipif_no_triton = pytest.mark.skipif(not HAS_TRITON, reason="Requires Triton") |
| 158 | |
| 159 | |
| 160 | @pytest.mark.skipif( |
| 161 | not current_platform.is_cuda_alike(), |
| 162 | reason="Only test on cuda and rocm platform", |
| 163 | ) |
| 164 | @pytest.mark.parametrize( |
| 165 | "model_class,expected_functionalized,expected_donated,expected_clones", |
| 166 | [ |
| 167 | # 2 inplace calls, all activations donated, all clones eliminated |
| 168 | (MaybeInplaceModel, 2, 3, 0), |
| 169 | # No inplace calls, no donations, 3 clones (one eliminated) |
| 170 | (FunctionalModel, 0, 0, 3), |
| 171 | # One inplace call, two donated activations, 2 clones |
| 172 | (MixedModel, 1, 2, 2), |
| 173 | # One inplace call, two donated, 1 clone remaining |
| 174 | pytest.param(ModelWithTritonAfterMaybeInplace, 1, 2, 1, marks=skipif_no_triton), |
| 175 | ], |
| 176 | ) |
| 177 | def test_inplace_functionalization( |
| 178 | default_vllm_config: VllmConfig, |
| 179 | model_class, |
| 180 | expected_functionalized: int, |
| 181 | expected_clones: int, |
| 182 | expected_donated: int, |
| 183 | ): |
| 184 | """Test inplace functionalization, lowering, and clone cleanup.""" |
| 185 | torch.set_default_device(current_platform.device_type) |
| 186 | |
| 187 | # Use vllm_c so inplace path is triggered |
| 188 | default_vllm_config.kernel_config.ir_op_priority.fused_add_rms_norm = [ |
| 189 | "vllm_c", |
| 190 | "native", |
| 191 | ] |
| 192 | |
| 193 | # Create passes in order they run during compilation |
| 194 | functionalization_pass = VllmIRInplaceFunctionalizationPass(default_vllm_config) |
| 195 | lowering_pass = VllmIRLoweringPass(default_vllm_config) |
| 196 | donated_info_pass = StoreDonationInfoPass() |
| 197 | cleanup_pass = UnsafeCloneEliminationPass(default_vllm_config) |
| 198 | |
| 199 | # Set up backend with pre-grad pass |
| 200 | backend = TestBackend(lowering_pass, donated_info_pass, cleanup_pass) |
| 201 | backend.inductor_config["pre_grad_custom_pass"] = functionalization_pass |
| 202 | |
| 203 | model = model_class() |
| 204 | x = torch.randn(8, 16, dtype=torch.bfloat16) |
| 205 | residual1 = torch.randn(8, 16, dtype=torch.bfloat16) |
| 206 | residual2 = torch.randn(8, 16, dtype=torch.bfloat16) |
| 207 | |
| 208 | with default_vllm_config.kernel_config.ir_op_priority.set_priority(): |
| 209 | # Reference output without optimization |
| 210 | ref_output = model(x.clone(), residual1.clone(), residual2.clone()) |
| 211 | |
| 212 | # Compile with inplace optimization |
| 213 | compiled_model = torch.compile(model, backend=backend, fullgraph=True) |
| 214 | output = compiled_model(x.clone(), residual1.clone(), residual2.clone()) |
| 215 | |
| 216 | # Verify correctness (relaxed tolerance for bfloat16) |
| 217 | for i in range(len(ref_output)): |
| 218 | torch.testing.assert_close(output[i], ref_output[i], rtol=1e-2, atol=1e-2) |
| 219 | |
| 220 | # Verify expected number of ops were functionalized |
| 221 | func_ops = functionalization_pass.functionalized_ops |
| 222 | assert len(func_ops) == int(bool(expected_functionalized)) |
| 223 | if expected_functionalized > 0: |
| 224 | assert "fused_add_rms_norm" in func_ops |
| 225 | assert func_ops["fused_add_rms_norm"] == expected_functionalized |
| 226 | |
| 227 | # Verify lowering happened (2 ops in all cases) |
| 228 | assert "fused_add_rms_norm" in lowering_pass.selected_impls |
| 229 | assert len(lowering_pass.selected_impls["fused_add_rms_norm"]) == 2 |
| 230 | assert all( |
| 231 | provider == "vllm_c" |
| 232 | for node, provider in lowering_pass.selected_impls["fused_add_rms_norm"].items() |
| 233 | ), lowering_pass.selected_impls |
| 234 | |
| 235 | # Verify correct number of donated IDs |
| 236 | assert len(donated_info_pass.donated_input_ids_sets) == 1 |
| 237 | assert len(donated_info_pass.donated_input_ids_sets[0]) == expected_donated |
| 238 | |
| 239 | # Verify expected number of clones after cleanup |
| 240 | actual_clones = backend.op_count(torch.ops.aten.clone.default, before=False) |
| 241 | assert actual_clones == expected_clones, ( |
| 242 | f"Expected {expected_clones} clones, got {actual_clones}:" |
| 243 | f"{backend.print_graphs()}" |
| 244 | ) |
| 245 | |
| 246 | |
| 247 | @pytest.mark.skipif( |
| 248 | not current_platform.is_cuda_alike(), |
| 249 | reason="Only test on cuda and rocm platform", |
| 250 | ) |
| 251 | def test_donated_buffer_context_propagation(default_vllm_config): |
| 252 | """Test that donated_input_ids propagates correctly through pass_context.""" |
| 253 | torch.set_default_device(current_platform.device_type) |
| 254 | |
| 255 | # Create a custom backend that inspects pass_context in cleanup pass |
| 256 | functionalization_pass = VllmIRInplaceFunctionalizationPass(default_vllm_config) |
| 257 | lowering_pass = VllmIRLoweringPass(default_vllm_config) |
| 258 | |
| 259 | donation_info_pass = StoreDonationInfoPass() |
| 260 | cleanup_pass = UnsafeCloneEliminationPass(default_vllm_config) |
| 261 | |
| 262 | backend = TestBackend(lowering_pass, donation_info_pass, cleanup_pass) |
| 263 | backend.inductor_config["pre_grad_custom_pass"] = functionalization_pass |
| 264 | |
| 265 | model = MaybeInplaceModel() |
| 266 | x = torch.randn(8, 16, dtype=torch.bfloat16) |
| 267 | residual1 = torch.randn(8, 16, dtype=torch.bfloat16) |
| 268 | residual2 = torch.randn(8, 16, dtype=torch.bfloat16) |
| 269 | |
| 270 | compiled_model = torch.compile(model, backend=backend, fullgraph=True) |
| 271 | compiled_model(x.clone(), residual1.clone(), residual2.clone()) |
| 272 | |
| 273 | donated_ids_seen = donation_info_pass.donated_input_ids_sets |
| 274 | # Verify donated_input_ids was set and propagated |
| 275 | assert len(donated_ids_seen) == 1 |
| 276 | # Should have donated inputs (exact indices depend on AOTAutograd) |
| 277 | assert len(donated_ids_seen[0]) == 3 |
| 278 | # All donated ids should be valid non-negative integers |
| 279 | for idx in donated_ids_seen[0]: |
| 280 | assert isinstance(idx, int) and idx >= 0, f"Invalid donated index: {idx}" |
| 281 | |
| 282 | |
| 283 | @pytest.mark.skipif( |
| 284 | not current_platform.is_cuda_alike(), |
| 285 | reason="Only test on cuda and rocm platform", |
| 286 | ) |
| 287 | def test_maybe_inplace_reuse_error(default_vllm_config): |
| 288 | """Test that reusing a donated activation input raises ValueError.""" |
| 289 | torch.set_default_device(current_platform.device_type) |
| 290 | |
| 291 | class ReuseModel(nn.Module): |
| 292 | """Model that incorrectly reuses a donated activation input.""" |
| 293 | |
| 294 | def __init__(self, hidden_size=16): |
| 295 | super().__init__() |
| 296 | self.weight = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16)) |
| 297 | |
| 298 | def forward(self, x: torch.Tensor, residual: torch.Tensor): |
| 299 | # x is donated to maybe_inplace |
| 300 | x_normed, residual_out = ops.fused_add_rms_norm.maybe_inplace( |
| 301 | x, residual, self.weight, 1e-5 |
| 302 | ) |
| 303 | # ERROR: x is used again after being donated |
| 304 | return x_normed + x # This should raise ValueError |
| 305 | |
| 306 | functionalization_pass = VllmIRInplaceFunctionalizationPass(default_vllm_config) |
| 307 | lowering_pass = VllmIRLoweringPass(default_vllm_config) |
| 308 | cleanup_pass = UnsafeCloneEliminationPass(default_vllm_config) |
| 309 | |
| 310 | backend = TestBackend(lowering_pass, cleanup_pass) |
| 311 | backend.inductor_config["pre_grad_custom_pass"] = functionalization_pass |
| 312 | |
| 313 | model = ReuseModel() |
| 314 | x = torch.randn(8, 16, dtype=torch.bfloat16) |
| 315 | residual = torch.randn(8, 16, dtype=torch.bfloat16) |
| 316 | |
| 317 | # Compilation should raise BackendCompilerFailed wrapping ValueError |
| 318 | with pytest.raises( |
| 319 | torch._dynamo.exc.BackendCompilerFailed, |
| 320 | match="is used again after the node", |
| 321 | ): |
| 322 | compiled_model = torch.compile(model, backend=backend, fullgraph=True) |
| 323 | compiled_model(x.clone(), residual.clone()) |
| 324 | |
| 325 | |
| 326 | # Piecewise compilation tests with graph splitting |
| 327 | |
| 328 | |
| 329 | @torch.library.custom_op("vllm::test_split_marker", mutates_args=()) |
| 330 | def test_split_marker(x: torch.Tensor) -> torch.Tensor: |
| 331 | """Identity op that marks a split point for piecewise compilation.""" |
| 332 | return x.clone() |
| 333 | |
| 334 | |
| 335 | @test_split_marker.register_fake |
| 336 | def _fake_split_marker(x: torch.Tensor) -> torch.Tensor: |
| 337 | return torch.empty_like(x) |
| 338 | |
| 339 | |
| 340 | class TransformerBlockWithSplits(nn.Module): |
| 341 | """Transformer block with explicit split points for piecewise compilation.""" |
| 342 | |
| 343 | def __init__(self, hidden_size=32, intermediate_size=128): |
| 344 | super().__init__() |
| 345 | self.hidden_size = hidden_size |
| 346 | self.intermediate_size = intermediate_size |
| 347 | |
| 348 | # Attention-like projection |
| 349 | self.attn_proj = nn.Linear( |
| 350 | hidden_size, hidden_size, bias=False, dtype=torch.bfloat16 |
| 351 | ) |
| 352 | |
| 353 | # Post-attention norm |
| 354 | self.post_attn_norm = nn.Parameter( |
| 355 | torch.ones(hidden_size, dtype=torch.bfloat16) |
| 356 | ) |
| 357 | |
| 358 | # MLP |
| 359 | self.gate_proj = nn.Linear( |
| 360 | hidden_size, intermediate_size, bias=False, dtype=torch.bfloat16 |
| 361 | ) |
| 362 | self.up_proj = nn.Linear( |
| 363 | hidden_size, intermediate_size, bias=False, dtype=torch.bfloat16 |
| 364 | ) |
| 365 | self.down_proj = nn.Linear( |
| 366 | intermediate_size, hidden_size, bias=False, dtype=torch.bfloat16 |
| 367 | ) |
| 368 | |
| 369 | # Post-MLP norm |
| 370 | self.post_mlp_norm = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16)) |
| 371 | |
| 372 | def forward(self, x: torch.Tensor): |
| 373 | # Attention block with residual |
| 374 | residual1 = x |
| 375 | attn_out = self.attn_proj(x) |
| 376 | |
| 377 | # Fused add + norm (maybe_inplace: residual1 is donated) |
| 378 | normed1, residual1 = ops.fused_add_rms_norm.maybe_inplace( |
| 379 | attn_out, residual1, self.post_attn_norm, 1e-5 |
| 380 | ) |
| 381 | |
| 382 | # Force a graph split here |
| 383 | normed1 = torch.ops.vllm.test_split_marker(normed1) |
| 384 | |
| 385 | # MLP block |
| 386 | gate = self.gate_proj(normed1) |
| 387 | up = self.up_proj(normed1) |
| 388 | mlp_out = self.down_proj(gate * torch.nn.functional.silu(up)) |
| 389 | |
| 390 | # Fused add + norm (maybe_inplace: residual1 is donated) |
| 391 | normed2, residual2 = ops.fused_add_rms_norm.maybe_inplace( |
| 392 | mlp_out, residual1, self.post_mlp_norm, 1e-5 |
| 393 | ) |
| 394 | |
| 395 | return normed2, residual2 |
| 396 | |
| 397 | |
| 398 | def with_dyn_arg(fn: Callable, arg_index: int, dim_index: int): |
| 399 | def inner(*args): |
| 400 | torch._dynamo.mark_dynamic(args[arg_index], dim_index) |
| 401 | return fn(*args) |
| 402 | |
| 403 | return inner |
| 404 | |
| 405 | |
| 406 | @pytest.mark.skipif( |
| 407 | not current_platform.is_cuda_alike(), |
| 408 | reason="Only test on cuda and rocm platform", |
| 409 | ) |
| 410 | def test_piecewise_compilation_with_donated_buffers(monkeypatch, fresh_vllm_cache): |
| 411 | """Test piecewise compilation with donated buffers across graph splits. |
| 412 | Utilizes a custom splitting op. Uses fresh cache to avoid compilation caching. |
| 413 | """ |
| 414 | torch.set_default_device(current_platform.device_type) |
| 415 | |
| 416 | # Disable compilation cache to avoid serialization issues |
| 417 | monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1") |
| 418 | |
| 419 | from vllm.compilation.backends import VllmBackend |
| 420 | from vllm.config import CompilationConfig, VllmConfig |
| 421 | |
| 422 | # Create config with custom splitting op |
| 423 | store_donation_info = StoreDonationInfoPass() |
| 424 | vllm_config = VllmConfig( |
| 425 | compilation_config=CompilationConfig( |
| 426 | custom_ops=["all"], |
| 427 | splitting_ops=["vllm::test_split_marker"], |
| 428 | inductor_compile_config={"post_grad_custom_post_pass": store_donation_info}, |
| 429 | ) |
| 430 | ) |
| 431 | |
| 432 | backend = VllmBackend(vllm_config) |
| 433 | |
| 434 | model = TransformerBlockWithSplits() |
| 435 | x = torch.randn(8, 32, dtype=torch.bfloat16) |
| 436 | |
| 437 | # Reference output |
| 438 | ref_output = with_dyn_arg(model, 0, 0)(x.clone()) |
| 439 | |
| 440 | # Compile with piecewise compilation (graph will split at split_marker) |
| 441 | compiled_model = torch.compile(model, backend=backend, fullgraph=False) |
| 442 | output = with_dyn_arg(compiled_model, 0, 0)(x.clone()) |
| 443 | |
| 444 | # Verify correctness (relaxed tolerance for bfloat16) |
| 445 | torch.testing.assert_close(output[0], ref_output[0], rtol=1e-2, atol=1e-2) |
| 446 | torch.testing.assert_close(output[1], ref_output[1], rtol=1e-2, atol=1e-2) |
| 447 | |
| 448 | # Verify the model was split into multiple submodules |
| 449 | assert hasattr(backend, "split_gm"), "Backend should have split graph module" |
| 450 | |
| 451 | # Should have at least 2 submodules (split by test_split_marker op) |
| 452 | submodules = list(backend.split_gm.named_children()) |
| 453 | num_submodules = len(submodules) |
| 454 | assert num_submodules >= 2, ( |
| 455 | f"Expected at least 2 submodules (split), got {num_submodules}" |
| 456 | ) |
| 457 | |
| 458 | # Check that donation info was propagated correctly |
| 459 | donated_inputs_sets = store_donation_info.donated_input_ids_sets |
| 460 | assert len(donated_inputs_sets) == 2 |
| 461 | assert len(donated_inputs_sets[0]) == 1 |
| 462 | assert len(donated_inputs_sets[1]) == 1 |