| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import contextlib |
| 4 | import os |
| 5 | import weakref |
| 6 | |
| 7 | import pytest |
| 8 | |
| 9 | from tests.utils import wait_for_gpu_memory_to_clear |
| 10 | from tests.v1.attention.utils import full_cg_backend_configs as backend_configs |
| 11 | from vllm import LLM, SamplingParams |
| 12 | from vllm.config import CompilationConfig, CUDAGraphMode |
| 13 | from vllm.outputs import RequestOutput |
| 14 | from vllm.platforms import current_platform |
| 15 | from vllm.utils.torch_utils import is_torch_equal_or_newer |
| 16 | from vllm.v1.attention.backends.registry import AttentionBackendEnum |
| 17 | |
| 18 | |
| 19 | @contextlib.contextmanager |
| 20 | def temporary_environ(env_vars): |
| 21 | """Temporarily set environment variables and restore them afterward. |
| 22 | We have to do this vs monkeypatch because monkeypatch doesn't work |
| 23 | with "module" scoped fixtures. |
| 24 | """ |
| 25 | original_env = {k: os.environ.get(k) for k in env_vars} |
| 26 | try: |
| 27 | os.environ.update(env_vars) |
| 28 | yield |
| 29 | finally: |
| 30 | for k, v in original_env.items(): |
| 31 | if v is None: |
| 32 | os.environ.pop(k, None) |
| 33 | else: |
| 34 | os.environ[k] = v |
| 35 | |
| 36 | |
| 37 | model_backends_full_cudagraph = [] |
| 38 | |
| 39 | # deepseek-ai/DeepSeek-V2-Lite with MLA |
| 40 | MLA_backends = ["FlashMLA", "FlashAttentionMLA", "CutlassMLA", "FlashInferMLA"] |
| 41 | for mla_backend in MLA_backends: |
| 42 | model_backends_full_cudagraph.append( |
| 43 | ("deepseek-ai/DeepSeek-V2-Lite", backend_configs[mla_backend]) |
| 44 | ) |
| 45 | |
| 46 | # Qwen/Qwen2-1.5B-Instruct with other backends |
| 47 | other_backend_configs = [ |
| 48 | backend_configs[c] for c in backend_configs if c not in MLA_backends |
| 49 | ] |
| 50 | for backend_config in other_backend_configs: |
| 51 | model_backends_full_cudagraph.append(("Qwen/Qwen2-1.5B-Instruct", backend_config)) |
| 52 | |
| 53 | |
| 54 | @pytest.fixture(scope="class") |
| 55 | def llm_pair(request): |
| 56 | model, backend_config, use_inductor_graph_partition = request.param |
| 57 | if use_inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"): |
| 58 | pytest.skip("Inductor graph partition only supported in torch>=2.9") |
| 59 | |
| 60 | backend = AttentionBackendEnum[backend_config.attention_config["backend"]] |
| 61 | if current_platform.is_rocm() and backend not in ( |
| 62 | AttentionBackendEnum.TRITON_ATTN, |
| 63 | AttentionBackendEnum.ROCM_ATTN, |
| 64 | ): |
| 65 | pytest.skip(f"{backend_config.name} is a CUDA-only attention backend") |
| 66 | if backend == AttentionBackendEnum.ROCM_ATTN and not current_platform.is_rocm(): |
| 67 | pytest.skip("ROCM_ATTN requires ROCm") |
| 68 | |
| 69 | # Dynamically skip test if GPU capability is not met |
| 70 | if backend_config.specific_gpu_arch and ( |
| 71 | not current_platform.is_cuda() |
| 72 | or backend_config.specific_gpu_arch != current_platform.get_device_capability() |
| 73 | ): |
| 74 | if backend_config.specific_gpu_arch == (9, 0): |
| 75 | pytest.skip("Only Hopper GPUs support FA3 and FlashMLA") |
| 76 | elif backend_config.specific_gpu_arch == (10, 0): |
| 77 | pytest.skip("Only Blackwell GPUs support Cutlass MLA") |
| 78 | |
| 79 | env_vars = { |
| 80 | # Force native sampler to avoid potential nondeterminism in FlashInfer |
| 81 | # when per-request generators are not used in V1. |
| 82 | "VLLM_USE_FLASHINFER_SAMPLER": "0", |
| 83 | } |
| 84 | with temporary_environ(env_vars): |
| 85 | full = LLM( |
| 86 | model=model, |
| 87 | gpu_memory_utilization=0.43, |
| 88 | trust_remote_code=True, |
| 89 | max_model_len=1024, |
| 90 | max_num_seqs=128, |
| 91 | attention_config=backend_config.attention_config, |
| 92 | compilation_config=CompilationConfig( |
| 93 | **backend_config.comp_config, |
| 94 | use_inductor_graph_partition=use_inductor_graph_partition, |
| 95 | ), |
| 96 | generation_config="vllm", |
| 97 | seed=42, |
| 98 | ) |
| 99 | piecewise = LLM( |
| 100 | model=model, |
| 101 | gpu_memory_utilization=0.43, |
| 102 | trust_remote_code=True, |
| 103 | max_model_len=1024, |
| 104 | max_num_seqs=128, |
| 105 | attention_config=backend_config.attention_config, |
| 106 | compilation_config=CompilationConfig( |
| 107 | cudagraph_mode=CUDAGraphMode.PIECEWISE, |
| 108 | use_inductor_graph_partition=use_inductor_graph_partition, |
| 109 | ), |
| 110 | generation_config="vllm", |
| 111 | seed=42, |
| 112 | ) |
| 113 | |
| 114 | # PyTest caches the fixture values so we use weakref.proxy to enable GC |
| 115 | yield weakref.proxy(full), weakref.proxy(piecewise) |
| 116 | del full |
| 117 | del piecewise |
| 118 | |
| 119 | wait_for_gpu_memory_to_clear( |
| 120 | devices=[0], |
| 121 | threshold_ratio=0.1, |
| 122 | ) |
| 123 | |
| 124 | |
| 125 | @pytest.mark.parametrize( |
| 126 | "llm_pair", |
| 127 | [ |
| 128 | pytest.param( |
| 129 | (model, backend_config, use_inductor_graph_partition), |
| 130 | id=f"{backend_config.name}-partition={use_inductor_graph_partition}", |
| 131 | ) |
| 132 | for model, backend_config in model_backends_full_cudagraph |
| 133 | for use_inductor_graph_partition in [True, False] |
| 134 | ], |
| 135 | indirect=True, |
| 136 | ) |
| 137 | class TestFullCUDAGraph: |
| 138 | """Use a class such that an llm pair is constructed once for all |
| 139 | batch_size/max_tokens combinations and released immediately after. |
| 140 | |
| 141 | Module-scope fixtures would stick around the whole time, |
| 142 | meaning there would be multiple LLM instances hogging memory simultaneously. |
| 143 | """ |
| 144 | |
| 145 | @pytest.mark.parametrize( |
| 146 | ("batch_size", "max_tokens"), |
| 147 | [ |
| 148 | (1, 10), |
| 149 | (7, 10), |
| 150 | (16, 10), |
| 151 | (25, 10), |
| 152 | (32, 10), |
| 153 | (45, 10), |
| 154 | (64, 10), |
| 155 | (123, 10), |
| 156 | (8, 5), |
| 157 | (8, 30), |
| 158 | ], |
| 159 | ) |
| 160 | def test_full_cudagraph(self, batch_size, max_tokens, llm_pair: tuple[LLM, LLM]): |
| 161 | """Test various batch sizes and max_tokens to ensure that the |
| 162 | full cudagraph compilation works for padded cases too. |
| 163 | """ |
| 164 | full_cudagraph_llm, piecewise_llm = llm_pair |
| 165 | |
| 166 | prompts = ["the quick brown fox"] * batch_size |
| 167 | # Use purely greedy decoding to avoid top-p truncation sensitivity |
| 168 | # that can amplify tiny numeric differences across runtimes. |
| 169 | sampling_params = SamplingParams( |
| 170 | temperature=0.0, max_tokens=max_tokens, top_p=1.0 |
| 171 | ) |
| 172 | |
| 173 | responses = [] |
| 174 | for llm in (piecewise_llm, full_cudagraph_llm): |
| 175 | # Queue the whole batch before scheduling so both engines exercise |
| 176 | # the requested batch size instead of timing-dependent sub-batches. |
| 177 | llm.sleep(level=0, mode="keep") |
| 178 | try: |
| 179 | llm.enqueue(prompts, sampling_params) |
| 180 | finally: |
| 181 | llm.wake_up(tags=["scheduling"]) |
| 182 | responses.append(llm.wait_for_completion(output_type=RequestOutput)) |
| 183 | piecewise_responses, full_responses = responses |
| 184 | |
| 185 | # Check that all responses are the same |
| 186 | assert len(piecewise_responses) == len(full_responses) == batch_size |
| 187 | for piecewise_res, full_res in zip(piecewise_responses, full_responses): |
| 188 | assert ( |
| 189 | piecewise_res.outputs[0].text.lower() |
| 190 | == full_res.outputs[0].text.lower() |
| 191 | ) |