| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from unittest.mock import patch |
| 4 | |
| 5 | import pytest |
| 6 | import regex as re |
| 7 | import torch |
| 8 | from torch import nn |
| 9 | |
| 10 | import tests.compile.silly_attention # noqa |
| 11 | from vllm.compilation.decorators import support_torch_compile |
| 12 | from vllm.config import VllmConfig, set_current_vllm_config |
| 13 | from vllm.config.compilation import ( |
| 14 | CompilationConfig, |
| 15 | CompilationMode, |
| 16 | CUDAGraphMode, |
| 17 | ) |
| 18 | from vllm.config.scheduler import SchedulerConfig |
| 19 | from vllm.forward_context import set_forward_context |
| 20 | from vllm.platforms import current_platform |
| 21 | |
| 22 | MLP_SIZE = 64 |
| 23 | DEVICE_TYPE = current_platform.device_type |
| 24 | |
| 25 | |
| 26 | @support_torch_compile |
| 27 | class SimpleModel(nn.Module): |
| 28 | """A simple model with a splitting op for piecewise compilation.""" |
| 29 | |
| 30 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = "", **kwargs): |
| 31 | super().__init__() |
| 32 | |
| 33 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 34 | x = x + x |
| 35 | attn_output = torch.empty_like(x) |
| 36 | torch.ops.silly.attention(x, x, x, attn_output) |
| 37 | x = attn_output * 2 |
| 38 | return x |
| 39 | |
| 40 | |
| 41 | class TraceStructuredCapture: |
| 42 | """Captures trace_structured calls for testing.""" |
| 43 | |
| 44 | def __init__(self): |
| 45 | self.calls: list[dict] = [] |
| 46 | |
| 47 | def __call__(self, event_type: str, metadata_fn=None, payload_fn=None, **kwargs): |
| 48 | """Capture a trace_structured call.""" |
| 49 | metadata = metadata_fn() if metadata_fn else {} |
| 50 | self.calls.append( |
| 51 | { |
| 52 | "event_type": event_type, |
| 53 | "metadata": metadata, |
| 54 | } |
| 55 | ) |
| 56 | |
| 57 | def get(self, event_type: str, name_pattern: str) -> list[dict]: |
| 58 | """Get all calls with the given event type and name matching pattern. |
| 59 | |
| 60 | Args: |
| 61 | event_type: The event type to filter by (e.g., "artifact", "graph_dump") |
| 62 | name_pattern: Regex pattern to match against the artifact name |
| 63 | |
| 64 | """ |
| 65 | regex = re.compile(name_pattern) |
| 66 | return [ |
| 67 | c |
| 68 | for c in self.calls |
| 69 | if c["event_type"] == event_type |
| 70 | and regex.fullmatch(c.get("metadata", {}).get("name", "")) |
| 71 | ] |
| 72 | |
| 73 | |
| 74 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") |
| 75 | def test_vllm_structured_logging_artifacts(use_fresh_inductor_cache): |
| 76 | """Test that all expected vLLM artifacts are logged during compilation.""" |
| 77 | torch.set_default_device(DEVICE_TYPE) |
| 78 | |
| 79 | capture = TraceStructuredCapture() |
| 80 | |
| 81 | vllm_config = VllmConfig( |
| 82 | compilation_config=CompilationConfig( |
| 83 | mode=CompilationMode.VLLM_COMPILE, |
| 84 | cudagraph_mode=CUDAGraphMode.PIECEWISE, |
| 85 | compile_sizes=[8], |
| 86 | splitting_ops=["silly::attention"], |
| 87 | ), |
| 88 | scheduler_config=SchedulerConfig( |
| 89 | max_num_seqs=8, |
| 90 | max_model_len=8192, |
| 91 | is_encoder_decoder=False, |
| 92 | ), |
| 93 | ) |
| 94 | |
| 95 | # Patch trace_structured to capture calls |
| 96 | with ( |
| 97 | patch("vllm.compilation.backends.trace_structured", capture), |
| 98 | patch("vllm.compilation.piecewise_backend.trace_structured", capture), |
| 99 | set_current_vllm_config(vllm_config), |
| 100 | ): |
| 101 | model = SimpleModel(vllm_config=vllm_config, prefix="test") |
| 102 | with set_forward_context({}, vllm_config=vllm_config): |
| 103 | model(torch.randn(8, MLP_SIZE)) |
| 104 | |
| 105 | config_artifacts = capture.get("artifact", "vllm_compilation_config") |
| 106 | assert len(config_artifacts) == 1, ( |
| 107 | f"Expected 1 vllm_compilation_config, got {len(config_artifacts)}" |
| 108 | ) |
| 109 | vllm_piecewise_split_graph = capture.get("graph_dump", "vllm_piecewise_split_graph") |
| 110 | assert len(vllm_piecewise_split_graph) == 1, ( |
| 111 | "Expected 1 toplevel piecewise split graph, " |
| 112 | f"got {len(vllm_piecewise_split_graph)}" |
| 113 | ) |
| 114 | compile_start_artifacts = capture.get("artifact", "vllm_piecewise_compile_start") |
| 115 | assert len(compile_start_artifacts) == 4, ( |
| 116 | "Expected 4 vllm_piecewise_compile_start " |
| 117 | "(2 subgraphs x 2 ranges each: dynamic + compile size), " |
| 118 | f"got {len(compile_start_artifacts)}" |
| 119 | ) |
| 120 | submod_dumps = capture.get("graph_dump", r"vllm_submod_.*") |
| 121 | assert len(submod_dumps) == 2, ( |
| 122 | "Expected 2 submods (one before attention, one after attention), " |
| 123 | f"got {len(submod_dumps)}" |
| 124 | ) |