| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import copy |
| 4 | from contextlib import nullcontext |
| 5 | from types import SimpleNamespace |
| 6 | from unittest.mock import MagicMock, patch |
| 7 | |
| 8 | import pytest |
| 9 | import torch |
| 10 | from pydantic import ValidationError |
| 11 | |
| 12 | from vllm.compilation.counter import compilation_counter |
| 13 | from vllm.compilation.passes.utility.fix_functionalization import ( |
| 14 | FixFunctionalizationPass, |
| 15 | ) |
| 16 | from vllm.config import ( |
| 17 | CompilationConfig, |
| 18 | CUDAGraphMode, |
| 19 | ModelConfig, |
| 20 | ParallelConfig, |
| 21 | SchedulerConfig, |
| 22 | SpeculativeConfig, |
| 23 | VllmConfig, |
| 24 | ) |
| 25 | from vllm.config.compilation import CompilationMode, PassConfig |
| 26 | from vllm.engine.arg_utils import EngineArgs |
| 27 | from vllm.platforms import current_platform |
| 28 | from vllm.utils.math_utils import cdiv |
| 29 | from vllm.utils.torch_utils import ( |
| 30 | _is_torch_equal_or_newer, |
| 31 | is_torch_equal, |
| 32 | ) |
| 33 | from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher |
| 34 | |
| 35 | # This import automatically registers `torch.ops.silly.attention` |
| 36 | from . import silly_attention # noqa: F401 |
| 37 | |
| 38 | DEVICE_TYPE = current_platform.device_type |
| 39 | |
| 40 | |
| 41 | def test_version(): |
| 42 | # Test the version comparison logic using the private function |
| 43 | assert _is_torch_equal_or_newer("2.8.0.dev20250624+cu128", "2.8.0.dev") |
| 44 | assert _is_torch_equal_or_newer("2.8.0a0+gitc82a174", "2.8.0.dev") |
| 45 | assert _is_torch_equal_or_newer("2.8.0", "2.8.0.dev") |
| 46 | assert _is_torch_equal_or_newer("2.8.1", "2.8.0.dev") |
| 47 | assert not _is_torch_equal_or_newer("2.7.1", "2.8.0.dev") |
| 48 | |
| 49 | |
| 50 | def test_get_raw_stream_patch(): |
| 51 | """Test that get_raw_stream patch is applied only for torch 2.9.0 or 2.9.1.""" |
| 52 | import builtins |
| 53 | |
| 54 | # Check if get_raw_stream exists in builtins |
| 55 | has_patch = hasattr(builtins, "get_raw_stream") |
| 56 | |
| 57 | # Import torch to get actual version |
| 58 | |
| 59 | is_torch_2_9 = is_torch_equal("2.9.0") or is_torch_equal("2.9.1") |
| 60 | |
| 61 | if is_torch_2_9: |
| 62 | # For torch 2.9.x, the patch should be applied |
| 63 | assert has_patch, "get_raw_stream should be patched for torch 2.9.x" |
| 64 | # Verify it's callable (it should be the _cuda_getCurrentRawStream function) |
| 65 | get_raw_stream = builtins.get_raw_stream # type: ignore[attr-defined] |
| 66 | assert callable(get_raw_stream) |
| 67 | # Verify it's the correct function from torch._C |
| 68 | from torch._C import _cuda_getCurrentRawStream |
| 69 | |
| 70 | assert get_raw_stream is _cuda_getCurrentRawStream |
| 71 | |
| 72 | |
| 73 | # Inductor can initialize CUDA even for CPU inputs. |
| 74 | @pytest.mark.forked |
| 75 | @pytest.mark.parametrize("use_v2", [False, True]) |
| 76 | def test_e8m0_custom_op_fullgraph(use_v2, use_fresh_inductor_cache): |
| 77 | """E8M0 inputs must not prevent mutable custom-op decomposition.""" |
| 78 | with torch.library._scoped_library("test_e8m0", "FRAGMENT") as lib: |
| 79 | lib.define("scale_(Tensor(a!) out, Tensor scale) -> ()") |
| 80 | |
| 81 | @torch.library.impl(lib, "scale_", "CPU") |
| 82 | def scale_impl(out, scale): |
| 83 | out.mul_(scale.float()) |
| 84 | |
| 85 | @torch.library.register_fake("test_e8m0::scale_", lib=lib) |
| 86 | def scale_fake(out, scale): |
| 87 | return None |
| 88 | |
| 89 | def forward(x, scale): |
| 90 | out = x.clone() |
| 91 | torch.ops.test_e8m0.scale_(out, scale) |
| 92 | return out |
| 93 | |
| 94 | x = torch.randn(4, 4) |
| 95 | scale = torch.full((4,), 128, dtype=torch.uint8).view(torch.float8_e8m0fnu) |
| 96 | compiled = torch.compile( |
| 97 | forward, |
| 98 | fullgraph=True, |
| 99 | options={"enable_auto_functionalized_v2": use_v2}, |
| 100 | ) |
| 101 | torch.testing.assert_close(compiled(x, scale), x * 2) |
| 102 | |
| 103 | |
| 104 | def test_copy_pass(): |
| 105 | vllm_config = VllmConfig() |
| 106 | inductor_pass = FixFunctionalizationPass(vllm_config) |
| 107 | copied_inductor_pass = copy.deepcopy(inductor_pass) |
| 108 | assert ( |
| 109 | copied_inductor_pass.compilation_config.use_inductor_graph_partition |
| 110 | == vllm_config.compilation_config.use_inductor_graph_partition |
| 111 | ) |
| 112 | assert ( |
| 113 | copied_inductor_pass.compilation_config.splitting_ops |
| 114 | == vllm_config.compilation_config.splitting_ops |
| 115 | ) |
| 116 | |
| 117 | |
| 118 | def test_custom_op(): |
| 119 | # proper syntax |
| 120 | _ = CompilationConfig(custom_ops=["+quant_fp8", "-silu_and_mul"]) |
| 121 | _ = CompilationConfig(custom_ops=["none", "+rms_norm"]) |
| 122 | _ = CompilationConfig(custom_ops=["+rms_norm", "+rms_norm"]) |
| 123 | |
| 124 | for custom_ops in (["quant_fp8"], ["+"], ["-"]): |
| 125 | with pytest.raises(ValueError, match="Invalid syntax '"): |
| 126 | CompilationConfig(custom_ops=custom_ops) |
| 127 | |
| 128 | |
| 129 | @pytest.mark.parametrize( |
| 130 | ("custom_ops", "config_kwargs", "match"), |
| 131 | [ |
| 132 | (["all", "none"], {}, "can contain only one base mode"), |
| 133 | ( |
| 134 | ["none", "+rms_norm", "-rms_norm"], |
| 135 | {}, |
| 136 | "cannot both enable and disable.*rms_norm", |
| 137 | ), |
| 138 | ( |
| 139 | ["-rotary_embedding"], |
| 140 | {"pass_config": PassConfig(enable_qk_norm_rope_fusion=True)}, |
| 141 | "cannot both enable and disable.*rotary_embedding", |
| 142 | ), |
| 143 | ], |
| 144 | ) |
| 145 | def test_reject_contradictory_custom_ops(custom_ops, config_kwargs, match): |
| 146 | with pytest.raises(ValueError, match=match): |
| 147 | CompilationConfig(custom_ops=custom_ops, **config_kwargs) |
| 148 | |
| 149 | |
| 150 | # forked needed to workaround https://github.com/vllm-project/vllm/issues/21073 |
| 151 | @pytest.mark.forked |
| 152 | # NB: We don't test VLLM_DISABLE_COMPILE_CACHE=0 because that depends |
| 153 | # on the state of the cache directory on the current machine, which |
| 154 | # may be influenced by other tests. |
| 155 | @pytest.mark.parametrize("val", ["1"]) |
| 156 | def test_VLLM_DISABLE_COMPILE_CACHE(vllm_runner, monkeypatch, val): |
| 157 | # Disable multiprocessing so that the counter is in the same process |
| 158 | monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") |
| 159 | monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", val) |
| 160 | |
| 161 | compilation_config = { |
| 162 | "cudagraph_mode": CUDAGraphMode.NONE, # speed things up a bit |
| 163 | } |
| 164 | with ( |
| 165 | compilation_counter.expect( |
| 166 | num_cache_entries_updated=0, num_compiled_artifacts_saved=0 |
| 167 | ), |
| 168 | # loading the model causes compilation (if enabled) to happen |
| 169 | vllm_runner( |
| 170 | "facebook/opt-125m", |
| 171 | compilation_config=compilation_config, |
| 172 | gpu_memory_utilization=0.4, |
| 173 | ) as _, |
| 174 | ): |
| 175 | pass |
| 176 | |
| 177 | |
| 178 | # forked needed to workaround https://github.com/vllm-project/vllm/issues/21073 |
| 179 | @pytest.mark.forked |
| 180 | @pytest.mark.parametrize( |
| 181 | "cudagraph_mode,num_cudagraph_captured", |
| 182 | [ |
| 183 | (CUDAGraphMode.NONE, 0), |
| 184 | (CUDAGraphMode.FULL_DECODE_ONLY, 1), |
| 185 | (CUDAGraphMode.PIECEWISE, 13), |
| 186 | (CUDAGraphMode.FULL_AND_PIECEWISE, 14), |
| 187 | ], |
| 188 | ) |
| 189 | def test_use_cudagraphs( |
| 190 | vllm_runner, monkeypatch, cudagraph_mode, num_cudagraph_captured |
| 191 | ): |
| 192 | # Disable multiprocessing so that the counter is in the same process |
| 193 | monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") |
| 194 | |
| 195 | compilation_config = { |
| 196 | "cudagraph_capture_sizes": [100], |
| 197 | "cudagraph_mode": cudagraph_mode, |
| 198 | } |
| 199 | num_gpu_runner_capture_triggers = 1 if cudagraph_mode != CUDAGraphMode.NONE else 0 |
| 200 | with ( |
| 201 | compilation_counter.expect( |
| 202 | num_graphs_seen=1, |
| 203 | num_gpu_runner_capture_triggers=num_gpu_runner_capture_triggers, |
| 204 | num_cudagraph_captured=num_cudagraph_captured, |
| 205 | ), |
| 206 | # loading the model causes compilation (if enabled) to happen |
| 207 | vllm_runner( |
| 208 | "facebook/opt-125m", |
| 209 | compilation_config=compilation_config, |
| 210 | gpu_memory_utilization=0.4, |
| 211 | ) as _, |
| 212 | ): |
| 213 | pass |
| 214 | |
| 215 | |
| 216 | # forked needed to workaround https://github.com/vllm-project/vllm/issues/21073 |
| 217 | @pytest.mark.forked |
| 218 | @pytest.mark.parametrize("use_v2_model_runner", [False, True]) |
| 219 | def test_stock_torch_compile(vllm_runner, monkeypatch, use_v2_model_runner): |
| 220 | monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", str(int(use_v2_model_runner))) |
| 221 | # Disable multiprocessing so that the counter is in the same process |
| 222 | monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") |
| 223 | |
| 224 | with ( |
| 225 | compilation_counter.expect(stock_torch_compile_count=1), |
| 226 | # loading the model causes compilation (if enabled) to happen |
| 227 | vllm_runner( |
| 228 | "facebook/opt-125m", |
| 229 | compilation_config={"mode": CompilationMode.STOCK_TORCH_COMPILE}, |
| 230 | gpu_memory_utilization=0.4, |
| 231 | ) as runner, |
| 232 | ): |
| 233 | outputs = runner.generate_greedy(["Hello, my name is"], max_tokens=5) |
| 234 | assert outputs[0][0] |
| 235 | |
| 236 | |
| 237 | # forked needed to workaround https://github.com/vllm-project/vllm/issues/21073 |
| 238 | @pytest.mark.forked |
| 239 | def test_no_compilation(vllm_runner, monkeypatch): |
| 240 | # Disable multiprocessing so that the counter is in the same process |
| 241 | monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") |
| 242 | with ( |
| 243 | compilation_counter.expect(num_graphs_seen=0, stock_torch_compile_count=0), |
| 244 | # loading the model causes compilation (if enabled) to happen |
| 245 | vllm_runner( |
| 246 | "facebook/opt-125m", |
| 247 | compilation_config={"mode": CompilationMode.NONE}, |
| 248 | gpu_memory_utilization=0.4, |
| 249 | ) as _, |
| 250 | ): |
| 251 | pass |
| 252 | |
| 253 | |
| 254 | # forked needed to workaround https://github.com/vllm-project/vllm/issues/21073 |
| 255 | @pytest.mark.forked |
| 256 | def test_enforce_eager(vllm_runner, monkeypatch): |
| 257 | # Disable multiprocessing so that the counter is in the same process |
| 258 | monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") |
| 259 | |
| 260 | with ( |
| 261 | compilation_counter.expect(num_graphs_seen=0, stock_torch_compile_count=0), |
| 262 | # loading the model causes compilation (if enabled) to happen |
| 263 | vllm_runner( |
| 264 | "facebook/opt-125m", enforce_eager=True, gpu_memory_utilization=0.4 |
| 265 | ) as _, |
| 266 | ): |
| 267 | pass |
| 268 | |
| 269 | |
| 270 | @pytest.mark.parametrize("enable_fault_tolerance", [False, True]) |
| 271 | def test_enforce_eager_jit_warmup(enable_fault_tolerance): |
| 272 | """Enforce-eager disables JIT warmup unless fault tolerance is on. |
| 273 | |
| 274 | FT fault detection runs against deadlines that in-inference Triton |
| 275 | compilation latency spikes can blow past, so warmup stays enabled. |
| 276 | """ |
| 277 | config = VllmConfig( |
| 278 | model_config=ModelConfig(model="facebook/opt-125m", enforce_eager=True), |
| 279 | parallel_config=ParallelConfig(enable_fault_tolerance=enable_fault_tolerance), |
| 280 | ) |
| 281 | assert config.compilation_config.mode == CompilationMode.NONE |
| 282 | assert config.kernel_config.enable_jit_warmup == enable_fault_tolerance |
| 283 | |
| 284 | |
| 285 | @pytest.mark.forked |
| 286 | def test_torch_compile_disable(vllm_runner, monkeypatch): |
| 287 | monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") |
| 288 | monkeypatch.setenv("TORCH_COMPILE_DISABLE", "1") |
| 289 | monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1") |
| 290 | |
| 291 | with ( |
| 292 | compilation_counter.expect(num_graphs_seen=0, stock_torch_compile_count=0), |
| 293 | vllm_runner( |
| 294 | "facebook/opt-125m", |
| 295 | gpu_memory_utilization=0.4, |
| 296 | ) as _, |
| 297 | ): |
| 298 | pass |
| 299 | |
| 300 | |
| 301 | def test_splitting_ops_dynamic(): |
| 302 | # Default config |
| 303 | config = VllmConfig() |
| 304 | # Default V1 config leaves cudagraph mode unset; splitting ops are only |
| 305 | # populated when the engine decides to use piecewise compilation. |
| 306 | assert config.compilation_config.cudagraph_mode == CUDAGraphMode.FULL_AND_PIECEWISE |
| 307 | assert config.compilation_config.splitting_ops_contain_attention() |
| 308 | splitting_ops = config.compilation_config.splitting_ops |
| 309 | assert splitting_ops is not None |
| 310 | assert { |
| 311 | "vllm::qwen_gdn_attention_core_fused_norm_packed", |
| 312 | } <= set(splitting_ops) |
| 313 | |
| 314 | # When use_inductor_graph_partition=True |
| 315 | config = VllmConfig( |
| 316 | compilation_config=CompilationConfig( |
| 317 | mode=CompilationMode.VLLM_COMPILE, |
| 318 | use_inductor_graph_partition=True, |
| 319 | splitting_ops=["vllm::unified_attention_with_output"], |
| 320 | ) |
| 321 | ) |
| 322 | # with inductor partition we use splitting_ops directly for |
| 323 | # partition rules |
| 324 | assert config.compilation_config.splitting_ops == [ |
| 325 | "vllm::unified_attention_with_output" |
| 326 | ] |
| 327 | |
| 328 | # When attn_fusion pass enabled. |
| 329 | config = VllmConfig( |
| 330 | compilation_config=CompilationConfig( |
| 331 | mode=CompilationMode.VLLM_COMPILE, |
| 332 | pass_config=PassConfig(fuse_attn_quant=True, eliminate_noops=True), |
| 333 | custom_ops=["+quant_fp8"], |
| 334 | cudagraph_mode=CUDAGraphMode.PIECEWISE, |
| 335 | ) |
| 336 | ) |
| 337 | assert config.compilation_config.splitting_ops == [] |
| 338 | # cudagraph mode also fall back to FULL |
| 339 | assert config.compilation_config.cudagraph_mode == CUDAGraphMode.FULL |
| 340 | |
| 341 | # splitting_ops can not contain attention ops when attn_fusion |
| 342 | # pass enabled. |
| 343 | with pytest.raises(ValidationError): |
| 344 | config = VllmConfig( |
| 345 | compilation_config=CompilationConfig( |
| 346 | mode=CompilationMode.VLLM_COMPILE, |
| 347 | pass_config=PassConfig(fuse_attn_quant=True, eliminate_noops=True), |
| 348 | custom_ops=["+quant_fp8"], |
| 349 | cudagraph_mode=CUDAGraphMode.PIECEWISE, |
| 350 | # work around for accessing all attntion ops |
| 351 | splitting_ops=CompilationConfig()._attention_ops, |
| 352 | ) |
| 353 | ) |
| 354 | |
| 355 | # When both use_inductor_graph_partition and attn_fusion pass enabled. |
| 356 | config = VllmConfig( |
| 357 | compilation_config=CompilationConfig( |
| 358 | mode=CompilationMode.VLLM_COMPILE, |
| 359 | use_inductor_graph_partition=True, |
| 360 | pass_config=PassConfig(fuse_attn_quant=True, eliminate_noops=True), |
| 361 | custom_ops=["+quant_fp8"], |
| 362 | cudagraph_mode=CUDAGraphMode.PIECEWISE, |
| 363 | ) |
| 364 | ) |
| 365 | # With inductor graph partition, attn_fusion and splitting_ops |
| 366 | # work together. Default splitting_ops include attention ops. |
| 367 | assert config.compilation_config.splitting_ops_contain_attention() |
| 368 | # fuse_attn_quant is directly supported under |
| 369 | # use_inductor_graph_partition=True, and cudagraph_mode |
| 370 | # is unchanged. |
| 371 | assert config.compilation_config.cudagraph_mode == CUDAGraphMode.PIECEWISE |
| 372 | |
| 373 | |
| 374 | def test_moe_splitting_ops_deepep_ht_inductor_partition(): |
| 375 | # Inductor partition case: user-provided splitting_ops should be |
| 376 | # preserved and MoE ops should be appended for DeepEP HT with dp>1. |
| 377 | config = VllmConfig( |
| 378 | parallel_config=ParallelConfig( |
| 379 | all2all_backend="deepep_high_throughput", |
| 380 | data_parallel_size=8, |
| 381 | ), |
| 382 | compilation_config=CompilationConfig( |
| 383 | mode=CompilationMode.VLLM_COMPILE, |
| 384 | use_inductor_graph_partition=True, |
| 385 | splitting_ops=[ |
| 386 | "vllm::unified_attention_with_output", |
| 387 | "vllm::moe_forward", |
| 388 | "vllm::moe_forward_shared", |
| 389 | ], |
| 390 | ), |
| 391 | ) |
| 392 | splitting_ops = config.compilation_config.splitting_ops |
| 393 | assert splitting_ops == [ |
| 394 | "vllm::unified_attention_with_output", |
| 395 | "vllm::moe_forward", |
| 396 | "vllm::moe_forward_shared", |
| 397 | ] |
| 398 | |
| 399 | |
| 400 | def test_should_split(): |
| 401 | import torch |
| 402 | |
| 403 | from vllm.compilation.partition_rules import should_split |
| 404 | |
| 405 | graph = torch.fx.Graph() |
| 406 | node = torch.fx.Node( |
| 407 | graph=graph, |
| 408 | name="dummy_node", |
| 409 | op="call_function", |
| 410 | target=torch.ops.aten.add.default, |
| 411 | args=(), |
| 412 | kwargs={}, |
| 413 | ) |
| 414 | |
| 415 | # supports OpOverloadPacket |
| 416 | splitting_ops = ["aten::add"] |
| 417 | assert should_split(node, splitting_ops) |
| 418 | |
| 419 | # supports OpOverload |
| 420 | splitting_ops = ["aten::add.default"] |
| 421 | assert should_split(node, splitting_ops) |
| 422 | |
| 423 | # supports OpOverload |
| 424 | splitting_ops = ["aten::add.Tensor"] |
| 425 | assert not should_split(node, splitting_ops) |
| 426 | |
| 427 | q, k, v, out = [torch.randn(1)] * 4 |
| 428 | |
| 429 | # supports custom ops as OpOverloadPacket |
| 430 | node = torch.fx.Node( |
| 431 | graph=graph, |
| 432 | name="dummy_node", |
| 433 | op="call_function", |
| 434 | target=torch.ops.silly.attention, |
| 435 | args=(q, k, v, out), |
| 436 | kwargs={}, |
| 437 | ) |
| 438 | |
| 439 | splitting_ops = ["silly::attention"] |
| 440 | assert should_split(node, splitting_ops) |
| 441 | |
| 442 | # supports custom ops as OpOverload |
| 443 | node = torch.fx.Node( |
| 444 | graph=graph, |
| 445 | name="dummy_node", |
| 446 | op="call_function", |
| 447 | target=torch.ops.silly.attention.default, |
| 448 | args=(q, k, v, out), |
| 449 | kwargs={}, |
| 450 | ) |
| 451 | |
| 452 | splitting_ops = ["silly::attention"] |
| 453 | assert should_split(node, splitting_ops) |
| 454 | |
| 455 | splitting_ops = ["silly::attention.default"] |
| 456 | assert should_split(node, splitting_ops) |
| 457 | |
| 458 | |
| 459 | @pytest.mark.skipif( |
| 460 | not current_platform.support_static_graph_mode(), |
| 461 | reason="Skip if not cudagraph mode supported", |
| 462 | ) |
| 463 | @pytest.mark.parametrize( |
| 464 | ( |
| 465 | "cudagraph_capture_sizes", |
| 466 | "max_cudagraph_capture_size", |
| 467 | "tp_size", |
| 468 | "enable_sp", |
| 469 | "max_num_batched_tokens", |
| 470 | "cudagraph_mode", |
| 471 | "expected_max_size", |
| 472 | ), |
| 473 | [ |
| 474 | (None, None, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256), |
| 475 | ([1, 2, 4], 4, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 4), |
| 476 | ( |
| 477 | [1, 2, 4], |
| 478 | 8, |
| 479 | 1, |
| 480 | False, |
| 481 | 2048, |
| 482 | CUDAGraphMode.FULL_AND_PIECEWISE, |
| 483 | ValidationError, |
| 484 | ), |
| 485 | ([1, 256], None, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256), |
| 486 | ([], None, 1, False, 2048, CUDAGraphMode.NONE, 0), |
| 487 | (None, 0, 1, False, 2048, CUDAGraphMode.NONE, 0), |
| 488 | # truncated to nearest multiple of 8 or 16 |
| 489 | (None, 257, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256), |
| 490 | # max_num_batched_tokens <= max_cudagraph_capture_size should always be |
| 491 | # captured even if not landing on a 16-stride step |
| 492 | (None, 2048, 1, False, 257, CUDAGraphMode.FULL_AND_PIECEWISE, 257), |
| 493 | # max from list |
| 494 | ([1, 2, 4, 15], None, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 15), |
| 495 | # SP forces full-graph compilation, sizes are filtered by TP |
| 496 | ([1, 2, 4, 15], None, 2, True, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 4), |
| 497 | # limited by the max_tokens |
| 498 | ([1, 2, 4, 15], None, 1, False, 8, CUDAGraphMode.FULL_AND_PIECEWISE, 4), |
| 499 | # the list should contain at least 1 element when use cudagraph |
| 500 | ([], None, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, ValidationError), |
| 501 | # the max capturing size should be >= 1 when use cudagraph |
| 502 | (None, 0, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, ValidationError), |
| 503 | ], |
| 504 | ) |
| 505 | def test_cudagraph_sizes_post_init( |
| 506 | cudagraph_capture_sizes, |
| 507 | max_cudagraph_capture_size, |
| 508 | tp_size, |
| 509 | enable_sp, |
| 510 | max_num_batched_tokens, |
| 511 | cudagraph_mode, |
| 512 | expected_max_size, |
| 513 | ): |
| 514 | ctx = nullcontext() |
| 515 | if expected_max_size == ValidationError: |
| 516 | ctx = pytest.raises(expected_max_size) |
| 517 | |
| 518 | with ( |
| 519 | ctx, |
| 520 | patch.object(current_platform, "device_count", return_value=tp_size), |
| 521 | ): |
| 522 | kwargs = {} |
| 523 | if cudagraph_capture_sizes is not None: |
| 524 | kwargs["cudagraph_capture_sizes"] = cudagraph_capture_sizes |
| 525 | if max_cudagraph_capture_size is not None: |
| 526 | kwargs["max_cudagraph_capture_size"] = max_cudagraph_capture_size |
| 527 | compilation_config = CompilationConfig( |
| 528 | pass_config=PassConfig( |
| 529 | enable_sp=enable_sp, |
| 530 | fuse_norm_quant=True, |
| 531 | fuse_act_quant=True, |
| 532 | eliminate_noops=True, |
| 533 | sp_min_token_num=512 if enable_sp else None, |
| 534 | ), |
| 535 | cudagraph_mode=cudagraph_mode, |
| 536 | **kwargs, |
| 537 | ) |
| 538 | engine_args = EngineArgs( |
| 539 | model="facebook/opt-125m", |
| 540 | tensor_parallel_size=tp_size, |
| 541 | max_num_seqs=min(max_num_batched_tokens, 128), |
| 542 | max_num_batched_tokens=max_num_batched_tokens, |
| 543 | compilation_config=compilation_config, |
| 544 | ) |
| 545 | vllm_config = engine_args.create_engine_config() |
| 546 | |
| 547 | assert ( |
| 548 | vllm_config.compilation_config.max_cudagraph_capture_size |
| 549 | == expected_max_size |
| 550 | ) |
| 551 | |
| 552 | |
| 553 | def _mock_config_for_cudagraph_sizes( |
| 554 | max_num_seqs: int, |
| 555 | num_speculative_tokens: int, |
| 556 | max_num_batched_tokens: int, |
| 557 | compilation_config: CompilationConfig, |
| 558 | num_speculative_tokens_per_batch_size: list[tuple[int, int, int]] | None = None, |
| 559 | ) -> MagicMock: |
| 560 | """Mock VllmConfig wired up enough to run `_set_cudagraph_sizes`. |
| 561 | |
| 562 | `num_speculative_tokens` and `uniform_decode_query_len` are filled in by |
| 563 | calling the real property functions, so the tests below cover the |
| 564 | derivation from `speculative_config` and not just the arithmetic |
| 565 | downstream of it. |
| 566 | """ |
| 567 | config = MagicMock(spec=VllmConfig) |
| 568 | config.compilation_config = compilation_config |
| 569 | config.scheduler_config = SchedulerConfig.default_factory( |
| 570 | max_num_seqs=max_num_seqs, |
| 571 | max_num_batched_tokens=max_num_batched_tokens, |
| 572 | ) |
| 573 | config.parallel_config = ParallelConfig() |
| 574 | config.model_config = MagicMock() |
| 575 | config.model_config.enforce_eager = False |
| 576 | config.performance_mode = None |
| 577 | config.diffusion_config = None |
| 578 | schedule = num_speculative_tokens_per_batch_size |
| 579 | config.speculative_config = ( |
| 580 | SimpleNamespace( |
| 581 | num_speculative_tokens=num_speculative_tokens, |
| 582 | num_speculative_tokens_per_batch_size=schedule, |
| 583 | uses_dynamic_speculative_decoding=lambda: schedule is not None, |
| 584 | ) |
| 585 | if num_speculative_tokens |
| 586 | else None |
| 587 | ) |
| 588 | config.num_speculative_tokens = VllmConfig.num_speculative_tokens.fget(config) |
| 589 | config.uniform_decode_query_len = VllmConfig.uniform_decode_query_len.fget(config) |
| 590 | return config |
| 591 | |
| 592 | |
| 593 | @pytest.mark.parametrize("max_num_seqs", [100, 101]) |
| 594 | def test_default_cudagraph_capture_sizes_cover_off_stride_max_num_seqs( |
| 595 | max_num_seqs: int, |
| 596 | ) -> None: |
| 597 | compilation_config = CompilationConfig( |
| 598 | cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE |
| 599 | ) |
| 600 | config = _mock_config_for_cudagraph_sizes( |
| 601 | max_num_seqs=max_num_seqs, |
| 602 | num_speculative_tokens=0, |
| 603 | max_num_batched_tokens=32768, |
| 604 | compilation_config=compilation_config, |
| 605 | ) |
| 606 | |
| 607 | VllmConfig._set_cudagraph_sizes(config) |
| 608 | |
| 609 | assert max_num_seqs in compilation_config.cudagraph_capture_sizes |
| 610 | |
| 611 | |
| 612 | @pytest.mark.parametrize( |
| 613 | ("max_num_seqs", "num_speculative_tokens", "widest_is_captured"), |
| 614 | [ |
| 615 | # No speculation: the 2x headroom under the platform ceiling, unchanged. |
| 616 | (8, 0, True), |
| 617 | (32, 0, True), |
| 618 | # Speculating, but the widest decode batch still fits under the ceiling. |
| 619 | (64, 7, True), |
| 620 | (256, 1, True), |
| 621 | # Wider decode batches must not raise the memory-safety ceiling. |
| 622 | (32, 16, False), |
| 623 | (64, 16, False), |
| 624 | # Exact H200 regressions: DFlash (2176) and suffix decoding (6400). |
| 625 | (128, 16, False), |
| 626 | (256, 24, False), |
| 627 | # Widest decode batch off the capture stride, above the ceiling. |
| 628 | (33, 16, False), |
| 629 | # ... and off the stride while below it, where the ceiling stands but |
| 630 | # the generated sizes would otherwise stop at 400. |
| 631 | (24, 16, True), |
| 632 | ], |
| 633 | ) |
| 634 | def test_default_cudagraph_capture_size_respects_platform_ceiling( |
| 635 | max_num_seqs, num_speculative_tokens, widest_is_captured |
| 636 | ): |
| 637 | """Uniform decode coverage must not raise the platform default ceiling. |
| 638 | |
| 639 | A decode step presents up to `max_num_seqs * (1 + num_speculative_tokens)` |
| 640 | tokens. Off-stride sizes within the platform ceiling are captured exactly. |
| 641 | Larger batches fall back to eager execution instead of expanding the |
| 642 | default capture range and risking OOM during engine initialization. |
| 643 | """ |
| 644 | compilation_config = CompilationConfig( |
| 645 | cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE |
| 646 | ) |
| 647 | config = _mock_config_for_cudagraph_sizes( |
| 648 | max_num_seqs=max_num_seqs, |
| 649 | num_speculative_tokens=num_speculative_tokens, |
| 650 | max_num_batched_tokens=32768, |
| 651 | compilation_config=compilation_config, |
| 652 | ) |
| 653 | |
| 654 | VllmConfig._set_cudagraph_sizes(config) |
| 655 | |
| 656 | decode_query_len = 1 + num_speculative_tokens |
| 657 | default_max_graph_size = ( |
| 658 | 1024 if current_platform.is_device_capability_family(100) else 512 |
| 659 | ) |
| 660 | token_grid_max = min( |
| 661 | max_num_seqs * decode_query_len * 2, |
| 662 | default_max_graph_size, |
| 663 | ) |
| 664 | widest_uniform_decode = max_num_seqs * decode_query_len |
| 665 | assert compilation_config.max_cudagraph_capture_size == token_grid_max |
| 666 | assert ( |
| 667 | compilation_config.max_cudagraph_capture_size |
| 668 | == compilation_config.cudagraph_capture_sizes[-1] |
| 669 | ) |
| 670 | assert ( |
| 671 | widest_uniform_decode in compilation_config.cudagraph_capture_sizes |
| 672 | ) is widest_is_captured |
| 673 | |
| 674 | |
| 675 | def test_default_cudagraph_capture_sizes_keep_all_sizes_bounded(): |
| 676 | """Keep both token and uniform-decode grids under the platform default. |
| 677 | |
| 678 | This guards both the measured 581 versus 100 capture-count regression and |
| 679 | the graph-memory regression from capturing shapes up to 8704 tokens. |
| 680 | """ |
| 681 | max_num_seqs = 512 |
| 682 | decode_query_len = 17 |
| 683 | compilation_config = CompilationConfig( |
| 684 | cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE |
| 685 | ) |
| 686 | config = _mock_config_for_cudagraph_sizes( |
| 687 | max_num_seqs=max_num_seqs, |
| 688 | num_speculative_tokens=16, |
| 689 | max_num_batched_tokens=32768, |
| 690 | compilation_config=compilation_config, |
| 691 | ) |
| 692 | |
| 693 | with patch.object( |
| 694 | current_platform, |
| 695 | "is_device_capability_family", |
| 696 | return_value=False, |
| 697 | ): |
| 698 | default_max_graph_size = ( |
| 699 | 1024 if current_platform.is_device_capability_family(100) else 512 |
| 700 | ) |
| 701 | VllmConfig._set_cudagraph_sizes(config) |
| 702 | |
| 703 | token_grid_max = min(max_num_seqs * decode_query_len * 2, default_max_graph_size) |
| 704 | token_grid = [size for size in [1, 2, 4] if size <= token_grid_max] |
| 705 | token_grid += list(range(8, min(token_grid_max + 1, 256), 8)) |
| 706 | token_grid += list(range(256, token_grid_max + 1, 16)) |
| 707 | |
| 708 | max_request_count = min(max_num_seqs, default_max_graph_size) |
| 709 | request_counts = [count for count in [1, 2, 4] if count <= max_request_count] |
| 710 | request_counts += list(range(8, min(max_request_count + 1, 256), 8)) |
| 711 | request_counts += list(range(256, max_request_count + 1, 16)) |
| 712 | request_counts.append(max_request_count) |
| 713 | expected_sizes = sorted( |
| 714 | set( |
| 715 | token_grid |
| 716 | + [ |
| 717 | count * decode_query_len |
| 718 | for count in request_counts |
| 719 | if count * decode_query_len <= default_max_graph_size |
| 720 | ] |
| 721 | ) |
| 722 | ) |
| 723 | |
| 724 | assert compilation_config.cudagraph_capture_sizes == expected_sizes |
| 725 | assert all( |
| 726 | size <= default_max_graph_size |
| 727 | for size in compilation_config.cudagraph_capture_sizes |
| 728 | ) |
| 729 | |
| 730 | |
| 731 | def test_cudagraph_capture_sizes_respect_sequence_parallelism(): |
| 732 | """Sequence-parallel capture sizes stay divisible by tensor parallel size.""" |
| 733 | compilation_config = CompilationConfig( |
| 734 | cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE |
| 735 | ) |
| 736 | compilation_config.pass_config.enable_sp = True |
| 737 | config = _mock_config_for_cudagraph_sizes( |
| 738 | max_num_seqs=32, |
| 739 | num_speculative_tokens=16, |
| 740 | max_num_batched_tokens=32768, |
| 741 | compilation_config=compilation_config, |
| 742 | ) |
| 743 | config.parallel_config = SimpleNamespace(tensor_parallel_size=2) |
| 744 | config.update_sizes_for_sequence_parallelism = lambda sizes: ( |
| 745 | VllmConfig.update_sizes_for_sequence_parallelism(config, sizes) |
| 746 | ) |
| 747 | |
| 748 | with patch.object( |
| 749 | current_platform, |
| 750 | "is_device_capability_family", |
| 751 | return_value=False, |
| 752 | ): |
| 753 | VllmConfig._set_cudagraph_sizes(config) |
| 754 | |
| 755 | assert all(size % 2 == 0 for size in compilation_config.cudagraph_capture_sizes) |
| 756 | assert 544 not in compilation_config.cudagraph_capture_sizes |
| 757 | assert compilation_config.max_cudagraph_capture_size == 512 |
| 758 | |
| 759 | |
| 760 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA") |
| 761 | def test_real_vllm_config_caps_widest_ngram_decode_batch(): |
| 762 | """Real config post-init preserves the default capture ceiling.""" |
| 763 | model_config = ModelConfig(model="facebook/opt-125m") |
| 764 | speculative_config = SpeculativeConfig( |
| 765 | prompt_lookup_min=1, |
| 766 | prompt_lookup_max=1, |
| 767 | num_speculative_tokens=16, |
| 768 | method="ngram", |
| 769 | ) |
| 770 | scheduler_config = SchedulerConfig.default_factory( |
| 771 | max_num_seqs=32, |
| 772 | max_num_batched_tokens=32768, |
| 773 | ) |
| 774 | compilation_config = CompilationConfig( |
| 775 | cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE |
| 776 | ) |
| 777 | |
| 778 | with patch.object( |
| 779 | current_platform, |
| 780 | "is_device_capability_family", |
| 781 | return_value=False, |
| 782 | ): |
| 783 | config = VllmConfig( |
| 784 | model_config=model_config, |
| 785 | speculative_config=speculative_config, |
| 786 | scheduler_config=scheduler_config, |
| 787 | compilation_config=compilation_config, |
| 788 | ) |
| 789 | |
| 790 | assert 544 not in config.compilation_config.cudagraph_capture_sizes |
| 791 | assert ( |
| 792 | config.compilation_config.max_cudagraph_capture_size |
| 793 | == config.compilation_config.cudagraph_capture_sizes[-1] |
| 794 | == 512 |
| 795 | ) |
| 796 | |
| 797 | |
| 798 | @pytest.mark.parametrize("max_num_seqs", [8, 32, 256, 300, 512, 600, 1024, 2048]) |
| 799 | def test_default_cudagraph_capture_size_unchanged_without_speculation(max_num_seqs): |
| 800 | """Without speculation the default must reproduce the historical formula. |
| 801 | |
| 802 | The platform ceiling bounds a request count, and without speculation a |
| 803 | request is one token, so `min(max_num_seqs, ceiling) * 1` can never lift it. |
| 804 | The result must match `min(max_num_seqs * 2, ceiling)` bit for bit. |
| 805 | """ |
| 806 | compilation_config = CompilationConfig( |
| 807 | cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE |
| 808 | ) |
| 809 | config = _mock_config_for_cudagraph_sizes( |
| 810 | max_num_seqs=max_num_seqs, |
| 811 | num_speculative_tokens=0, |
| 812 | max_num_batched_tokens=1_000_000, |
| 813 | compilation_config=compilation_config, |
| 814 | ) |
| 815 | |
| 816 | with patch.object( |
| 817 | current_platform, |
| 818 | "is_device_capability_family", |
| 819 | return_value=False, |
| 820 | ): |
| 821 | VllmConfig._set_cudagraph_sizes(config) |
| 822 | |
| 823 | default_max_graph_size = ( |
| 824 | 1024 if current_platform.is_device_capability_family(100) else 512 |
| 825 | ) |
| 826 | assert compilation_config.max_cudagraph_capture_size == min( |
| 827 | max_num_seqs * 2, |
| 828 | default_max_graph_size, |
| 829 | ) |
| 830 | |
| 831 | |
| 832 | def test_single_speculative_token_does_not_raise_default_capture_size(): |
| 833 | """One speculative token must not raise the default capture ceiling. |
| 834 | |
| 835 | MTP at depth 1 gives a query length of 2, so 300 requests is a 600-token |
| 836 | decode batch against the platform's token ceiling. The widest batch falls |
| 837 | back to eager execution rather than expanding the capture range. |
| 838 | """ |
| 839 | compilation_config = CompilationConfig( |
| 840 | cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE |
| 841 | ) |
| 842 | config = _mock_config_for_cudagraph_sizes( |
| 843 | max_num_seqs=300, |
| 844 | num_speculative_tokens=1, |
| 845 | max_num_batched_tokens=1_000_000, |
| 846 | compilation_config=compilation_config, |
| 847 | ) |
| 848 | |
| 849 | VllmConfig._set_cudagraph_sizes(config) |
| 850 | |
| 851 | default_max_graph_size = ( |
| 852 | 1024 if current_platform.is_device_capability_family(100) else 512 |
| 853 | ) |
| 854 | token_grid_max = min(300 * 2 * 2, default_max_graph_size) |
| 855 | assert compilation_config.max_cudagraph_capture_size == token_grid_max |
| 856 | assert ( |
| 857 | compilation_config.max_cudagraph_capture_size |
| 858 | == compilation_config.cudagraph_capture_sizes[-1] |
| 859 | ) |
| 860 | assert 600 not in compilation_config.cudagraph_capture_sizes |
| 861 | |
| 862 | |
| 863 | def test_default_cudagraph_capture_size_caps_tokens(): |
| 864 | """The ceiling remains a platform-bounded token count. |
| 865 | |
| 866 | `max_num_seqs` and speculative width must not push the default capture range |
| 867 | above the platform's memory-safety guard. |
| 868 | """ |
| 869 | compilation_config = CompilationConfig( |
| 870 | cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE |
| 871 | ) |
| 872 | config = _mock_config_for_cudagraph_sizes( |
| 873 | max_num_seqs=1024, |
| 874 | num_speculative_tokens=16, |
| 875 | max_num_batched_tokens=1_000_000, |
| 876 | compilation_config=compilation_config, |
| 877 | ) |
| 878 | |
| 879 | VllmConfig._set_cudagraph_sizes(config) |
| 880 | |
| 881 | default_max_graph_size = ( |
| 882 | 1024 if current_platform.is_device_capability_family(100) else 512 |
| 883 | ) |
| 884 | assert compilation_config.max_cudagraph_capture_size == default_max_graph_size |
| 885 | assert ( |
| 886 | compilation_config.max_cudagraph_capture_size |
| 887 | == compilation_config.cudagraph_capture_sizes[-1] |
| 888 | ) |
| 889 | assert all( |
| 890 | size <= default_max_graph_size |
| 891 | for size in compilation_config.cudagraph_capture_sizes |
| 892 | ) |
| 893 | |
| 894 | |
| 895 | def _widest_covered_request_count(capture_sizes, query_len, max_num_seqs): |
| 896 | """Widest request count `CudaGraphManager` can build a decode graph for. |
| 897 | |
| 898 | Mirrors `_init_candidates`: a capture size is rounded up to a multiple of |
| 899 | the tier's query length, then dropped once the implied request count runs |
| 900 | past `max_num_seqs`. |
| 901 | """ |
| 902 | reachable = [ |
| 903 | cdiv(size, query_len) |
| 904 | for size in capture_sizes |
| 905 | if cdiv(size, query_len) <= max_num_seqs |
| 906 | ] |
| 907 | return max(reachable, default=0) |
| 908 | |
| 909 | |
| 910 | def test_default_cudagraph_capture_sizes_cover_every_dynamic_decode_width(): |
| 911 | """Each scheduled draft width needs sizes over the range it applies to. |
| 912 | |
| 913 | Dynamic speculative decoding picks the width from the batch size, so |
| 914 | scaling by the widest one alone leaves the narrower tiers short: sizes |
| 915 | built from query length 17 round up to multiples of 3 that imply more |
| 916 | requests than the scheduler can run, and coverage stops at 227 of 256. |
| 917 | """ |
| 918 | compilation_config = CompilationConfig( |
| 919 | cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE |
| 920 | ) |
| 921 | config = _mock_config_for_cudagraph_sizes( |
| 922 | max_num_seqs=128, |
| 923 | num_speculative_tokens=16, |
| 924 | max_num_batched_tokens=32768, |
| 925 | compilation_config=compilation_config, |
| 926 | # 16 draft tokens up to a batch of 16, then 2 out to 128. Both |
| 927 | # tier maxima fit under the 512-token default capture ceiling. |
| 928 | num_speculative_tokens_per_batch_size=[(1, 16, 16), (17, 128, 2)], |
| 929 | ) |
| 930 | |
| 931 | VllmConfig._set_cudagraph_sizes(config) |
| 932 | |
| 933 | sizes = compilation_config.cudagraph_capture_sizes |
| 934 | # The wide tier only ever runs to a batch of 16, the narrow one to 128. |
| 935 | assert _widest_covered_request_count(sizes, 17, 128) >= 16 |
| 936 | assert _widest_covered_request_count(sizes, 3, 128) == 128 |
| 937 | |
| 938 | |
| 939 | def test_dynamic_decode_capture_clamps_configured_width(): |
| 940 | compilation_config = CompilationConfig( |
| 941 | cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE |
| 942 | ) |
| 943 | config = _mock_config_for_cudagraph_sizes( |
| 944 | max_num_seqs=16, |
| 945 | num_speculative_tokens=3, |
| 946 | max_num_batched_tokens=32768, |
| 947 | compilation_config=compilation_config, |
| 948 | num_speculative_tokens_per_batch_size=[(1, 16, 5)], |
| 949 | ) |
| 950 | |
| 951 | VllmConfig._set_cudagraph_sizes(config) |
| 952 | |
| 953 | sizes = compilation_config.cudagraph_capture_sizes |
| 954 | assert _widest_covered_request_count(sizes, 4, 16) == 16 |
| 955 | assert 6 not in sizes |
| 956 | |
| 957 | |
| 958 | def test_dynamic_decode_capture_covers_schedule_gap_and_tail(): |
| 959 | compilation_config = CompilationConfig( |
| 960 | cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE |
| 961 | ) |
| 962 | config = _mock_config_for_cudagraph_sizes( |
| 963 | max_num_seqs=17, |
| 964 | num_speculative_tokens=4, |
| 965 | max_num_batched_tokens=32768, |
| 966 | compilation_config=compilation_config, |
| 967 | num_speculative_tokens_per_batch_size=[(1, 2, 4), (5, 5, 1)], |
| 968 | ) |
| 969 | |
| 970 | VllmConfig._set_cudagraph_sizes(config) |
| 971 | |
| 972 | sizes = compilation_config.cudagraph_capture_sizes |
| 973 | assert 20 in sizes |
| 974 | assert 34 in sizes |
| 975 | |
| 976 | |
| 977 | def test_default_cudagraph_capture_size_still_clamped_by_token_budget(): |
| 978 | """Decode coverage does not override the `max_num_batched_tokens` clamp. |
| 979 | |
| 980 | A batch wider than the token budget cannot be scheduled in the first place, |
| 981 | so there is no decode step of that size to capture a graph for. |
| 982 | """ |
| 983 | compilation_config = CompilationConfig( |
| 984 | cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE |
| 985 | ) |
| 986 | config = _mock_config_for_cudagraph_sizes( |
| 987 | max_num_seqs=32, |
| 988 | num_speculative_tokens=16, |
| 989 | max_num_batched_tokens=512, |
| 990 | compilation_config=compilation_config, |
| 991 | ) |
| 992 | |
| 993 | VllmConfig._set_cudagraph_sizes(config) |
| 994 | |
| 995 | # The token budget still clamps the final capture size. |
| 996 | default_max_graph_size = ( |
| 997 | 1024 if current_platform.is_device_capability_family(100) else 512 |
| 998 | ) |
| 999 | decode_query_len = 17 |
| 1000 | token_grid_max = min( |
| 1001 | 32 * decode_query_len * 2, |
| 1002 | default_max_graph_size, |
| 1003 | ) |
| 1004 | widest_capturable_decode = min( |
| 1005 | 32 * decode_query_len, |
| 1006 | 512, |
| 1007 | ) |
| 1008 | expected_max_size = max(token_grid_max, widest_capturable_decode) |
| 1009 | assert compilation_config.max_cudagraph_capture_size == expected_max_size |
| 1010 | assert ( |
| 1011 | compilation_config.max_cudagraph_capture_size |
| 1012 | == compilation_config.cudagraph_capture_sizes[-1] |
| 1013 | ) |
| 1014 | assert 544 not in compilation_config.cudagraph_capture_sizes |
| 1015 | |
| 1016 | |
| 1017 | @pytest.mark.skipif( |
| 1018 | not current_platform.support_static_graph_mode(), |
| 1019 | reason="Skip if static graph mode is not supported", |
| 1020 | ) |
| 1021 | @pytest.mark.parametrize( |
| 1022 | ("is_blackwell", "expected_max_size"), [(False, 512), (True, 1024)] |
| 1023 | ) |
| 1024 | def test_blackwell_cudagraph_default(is_blackwell, expected_max_size): |
| 1025 | vllm_config = VllmConfig() |
| 1026 | vllm_config.model_config = MagicMock(enforce_eager=False) |
| 1027 | vllm_config.scheduler_config = SchedulerConfig( |
| 1028 | max_num_seqs=512, |
| 1029 | max_num_batched_tokens=2048, |
| 1030 | max_model_len=2048, |
| 1031 | is_encoder_decoder=False, |
| 1032 | ) |
| 1033 | vllm_config.compilation_config = CompilationConfig( |
| 1034 | cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE, |
| 1035 | ) |
| 1036 | |
| 1037 | with patch.object( |
| 1038 | current_platform, |
| 1039 | "is_device_capability_family", |
| 1040 | return_value=is_blackwell, |
| 1041 | ): |
| 1042 | vllm_config._set_cudagraph_sizes() |
| 1043 | |
| 1044 | assert ( |
| 1045 | vllm_config.compilation_config.max_cudagraph_capture_size == expected_max_size |
| 1046 | ) |
| 1047 | |
| 1048 | |
| 1049 | @pytest.mark.skipif( |
| 1050 | not current_platform.support_static_graph_mode(), |
| 1051 | reason="Skip if not cudagraph mode supported", |
| 1052 | ) |
| 1053 | @pytest.mark.parametrize( |
| 1054 | ( |
| 1055 | "cudagraph_mode", |
| 1056 | "use_inductor_graph_partition", |
| 1057 | "expected_enable_sp", |
| 1058 | "expected_cudagraph_mode", |
| 1059 | "expected_piecewise_compile", |
| 1060 | "expected_capture_sizes", |
| 1061 | "expected_max_size", |
| 1062 | ), |
| 1063 | [ |
| 1064 | (CUDAGraphMode.PIECEWISE, False, True, CUDAGraphMode.FULL, False, [2, 4], 4), |
| 1065 | ( |
| 1066 | CUDAGraphMode.FULL_DECODE_ONLY, |
| 1067 | False, |
| 1068 | True, |
| 1069 | CUDAGraphMode.FULL_DECODE_ONLY, |
| 1070 | False, |
| 1071 | [2, 4], |
| 1072 | 4, |
| 1073 | ), |
| 1074 | ( |
| 1075 | CUDAGraphMode.FULL_AND_PIECEWISE, |
| 1076 | False, |
| 1077 | True, |
| 1078 | CUDAGraphMode.FULL, |
| 1079 | False, |
| 1080 | [2, 4], |
| 1081 | 4, |
| 1082 | ), |
| 1083 | ( |
| 1084 | CUDAGraphMode.FULL_AND_PIECEWISE, |
| 1085 | True, |
| 1086 | True, |
| 1087 | CUDAGraphMode.FULL_AND_PIECEWISE, |
| 1088 | True, |
| 1089 | [2, 4], |
| 1090 | 4, |
| 1091 | ), |
| 1092 | ], |
| 1093 | ) |
| 1094 | def test_sequence_parallelism_requires_full_graph_compilation( |
| 1095 | cudagraph_mode: CUDAGraphMode, |
| 1096 | use_inductor_graph_partition: bool, |
| 1097 | expected_enable_sp: bool, |
| 1098 | expected_cudagraph_mode: CUDAGraphMode, |
| 1099 | expected_piecewise_compile: bool, |
| 1100 | expected_capture_sizes: list[int], |
| 1101 | expected_max_size: int, |
| 1102 | ): |
| 1103 | with patch.object(current_platform, "device_count", return_value=2): |
| 1104 | vllm_config = VllmConfig( |
| 1105 | parallel_config=ParallelConfig(tensor_parallel_size=2), |
| 1106 | scheduler_config=SchedulerConfig( |
| 1107 | max_num_seqs=128, |
| 1108 | max_num_batched_tokens=2048, |
| 1109 | max_model_len=2048, |
| 1110 | is_encoder_decoder=False, |
| 1111 | ), |
| 1112 | ) |
| 1113 | vllm_config.model_config = MagicMock( |
| 1114 | dtype=torch.float16, |
| 1115 | enforce_eager=False, |
| 1116 | is_moe=False, |
| 1117 | disable_cascade_attn=False, |
| 1118 | get_hidden_size=MagicMock(return_value=4096), |
| 1119 | ) |
| 1120 | vllm_config.compilation_config = CompilationConfig( |
| 1121 | mode=CompilationMode.VLLM_COMPILE, |
| 1122 | cudagraph_capture_sizes=[1, 2, 4, 15], |
| 1123 | max_cudagraph_capture_size=None, |
| 1124 | compile_sizes=["cudagraph_capture_sizes"], |
| 1125 | use_inductor_graph_partition=use_inductor_graph_partition, |
| 1126 | pass_config=PassConfig( |
| 1127 | enable_sp=True, |
| 1128 | fuse_gemm_comms=True, |
| 1129 | fuse_norm_quant=True, |
| 1130 | fuse_act_quant=True, |
| 1131 | eliminate_noops=True, |
| 1132 | sp_min_token_num=512, |
| 1133 | ), |
| 1134 | cudagraph_mode=cudagraph_mode, |
| 1135 | ) |
| 1136 | vllm_config.compilation_config.set_splitting_ops_for_v1( |
| 1137 | all2all_backend=vllm_config.parallel_config.all2all_backend, |
| 1138 | data_parallel_size=1, |
| 1139 | ) |
| 1140 | vllm_config._set_compile_ranges() |
| 1141 | vllm_config._set_cudagraph_sizes() |
| 1142 | |
| 1143 | assert ( |
| 1144 | vllm_config.compilation_config.use_inductor_graph_partition |
| 1145 | == use_inductor_graph_partition |
| 1146 | ) |
| 1147 | assert ( |
| 1148 | bool(vllm_config.compilation_config.splitting_ops) == expected_piecewise_compile |
| 1149 | ) |
| 1150 | assert vllm_config.compilation_config.pass_config.enable_sp == expected_enable_sp |
| 1151 | assert ( |
| 1152 | vllm_config.compilation_config.pass_config.fuse_gemm_comms == expected_enable_sp |
| 1153 | ) |
| 1154 | assert vllm_config.compilation_config.cudagraph_mode == expected_cudagraph_mode |
| 1155 | assert ( |
| 1156 | vllm_config.compilation_config.cudagraph_capture_sizes == expected_capture_sizes |
| 1157 | ) |
| 1158 | assert ( |
| 1159 | vllm_config.compilation_config.max_cudagraph_capture_size == expected_max_size |
| 1160 | ) |
| 1161 | assert ( |
| 1162 | 511 in vllm_config.compilation_config.compile_ranges_endpoints |
| 1163 | ) == expected_enable_sp |
| 1164 | |
| 1165 | |
| 1166 | def test_cached_compilation_config(default_vllm_config): |
| 1167 | import torch |
| 1168 | from torch._inductor.utils import run_and_get_code |
| 1169 | |
| 1170 | from vllm.config import get_cached_compilation_config, set_current_vllm_config |
| 1171 | from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 |
| 1172 | from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape |
| 1173 | |
| 1174 | dtype = torch.bfloat16 |
| 1175 | device = torch.device(f"{DEVICE_TYPE}:0") |
| 1176 | batch_size, num_qo_heads, head_size = 8, 16, 128 |
| 1177 | |
| 1178 | # access and cache default compilation config |
| 1179 | # default compilation config does not contain +quant_fp8 custom op. If this is |
| 1180 | # used, the generated code would use inductor-generated triton kernel instead |
| 1181 | # of the custom op `torch.ops._C.static_scaled_fp8_quant`. |
| 1182 | get_cached_compilation_config() |
| 1183 | |
| 1184 | vllm_config = VllmConfig( |
| 1185 | compilation_config=CompilationConfig( |
| 1186 | mode=CompilationMode.VLLM_COMPILE, |
| 1187 | custom_ops=["+quant_fp8"], |
| 1188 | ) |
| 1189 | ) |
| 1190 | |
| 1191 | # set_current_vllm_config should clear cached compilation config and |
| 1192 | # use the new compilation_config in vllm_config |
| 1193 | with set_current_vllm_config(vllm_config): |
| 1194 | query_quant = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR) |
| 1195 | query_quant = torch.compile(query_quant) |
| 1196 | |
| 1197 | _q_scale = torch.tensor(1.0, dtype=torch.float32, device=DEVICE_TYPE) |
| 1198 | query = torch.randn( |
| 1199 | batch_size, num_qo_heads * head_size, dtype=dtype, device=device |
| 1200 | ) |
| 1201 | |
| 1202 | _, code = run_and_get_code(query_quant, query, _q_scale) |
| 1203 | |
| 1204 | code = " ".join(code) |
| 1205 | assert "torch.ops._C.static_scaled_fp8_quant.default(" in code |
| 1206 | |
| 1207 | |
| 1208 | def _create_vllm_config_for_validation( |
| 1209 | compilation_config: CompilationConfig, |
| 1210 | ) -> MagicMock: |
| 1211 | """Helper to create a mock VllmConfig for padding validation testing.""" |
| 1212 | mock_config = MagicMock(spec=VllmConfig) |
| 1213 | mock_config.compilation_config = compilation_config |
| 1214 | mock_config.scheduler_config = SchedulerConfig.default_factory(max_num_seqs=8) |
| 1215 | mock_config.parallel_config = ParallelConfig() |
| 1216 | mock_config.speculative_config = None |
| 1217 | mock_config.lora_config = None |
| 1218 | return mock_config |
| 1219 | |
| 1220 | |
| 1221 | def test_compile_sizes_padding_validation(): |
| 1222 | """Test that compile_sizes with values that would be padded raises an error.""" |
| 1223 | # cudagraph_capture_sizes=[1, 2, 4, 8] means: |
| 1224 | # - size 1 -> padded to 1 |
| 1225 | # - size 2 -> padded to 2 |
| 1226 | # - size 3 -> padded to 4 |
| 1227 | # - size 4 -> padded to 4 |
| 1228 | # - size 5 -> padded to 8 |
| 1229 | # etc. |
| 1230 | # So compile_sizes=[3] should fail because 3 would be padded to 4 |
| 1231 | |
| 1232 | with pytest.raises(ValueError, match="would be padded to"): |
| 1233 | config = CompilationConfig( |
| 1234 | cudagraph_capture_sizes=[1, 2, 4, 8], |
| 1235 | max_cudagraph_capture_size=8, |
| 1236 | compile_sizes=[3], |
| 1237 | cudagraph_mode=CUDAGraphMode.FULL, |
| 1238 | ) |
| 1239 | config.post_init_cudagraph_sizes() |
| 1240 | dispatcher = CudagraphDispatcher(_create_vllm_config_for_validation(config)) |
| 1241 | dispatcher.initialize_cudagraph_keys(CUDAGraphMode.FULL) |
| 1242 | |
| 1243 | with pytest.raises(ValueError, match="would be padded to"): |
| 1244 | config = CompilationConfig( |
| 1245 | cudagraph_capture_sizes=[1, 2, 4, 8], |
| 1246 | max_cudagraph_capture_size=8, |
| 1247 | compile_sizes=[5], |
| 1248 | cudagraph_mode=CUDAGraphMode.FULL, |
| 1249 | ) |
| 1250 | config.post_init_cudagraph_sizes() |
| 1251 | dispatcher = CudagraphDispatcher(_create_vllm_config_for_validation(config)) |
| 1252 | dispatcher.initialize_cudagraph_keys(CUDAGraphMode.FULL) |
| 1253 | |
| 1254 | config = CompilationConfig( |
| 1255 | cudagraph_capture_sizes=[1, 2, 4, 8], |
| 1256 | max_cudagraph_capture_size=8, |
| 1257 | compile_sizes=[1, 2, 4, 8], |
| 1258 | cudagraph_mode=CUDAGraphMode.FULL, |
| 1259 | ) |
| 1260 | config.post_init_cudagraph_sizes() |
| 1261 | assert sorted(config.compile_sizes) == [1, 2, 4, 8] |
| 1262 | dispatcher = CudagraphDispatcher(_create_vllm_config_for_validation(config)) |
| 1263 | dispatcher.initialize_cudagraph_keys(CUDAGraphMode.FULL) # Should not raise |
| 1264 | |
| 1265 | config = CompilationConfig( |
| 1266 | cudagraph_capture_sizes=[1, 2, 4, 8], |
| 1267 | max_cudagraph_capture_size=8, |
| 1268 | compile_sizes=["cudagraph_capture_sizes"], |
| 1269 | cudagraph_mode=CUDAGraphMode.FULL, |
| 1270 | ) |
| 1271 | config.post_init_cudagraph_sizes() |
| 1272 | assert sorted(config.compile_sizes) == [1, 2, 4, 8] |
| 1273 | |
| 1274 | # When cudagraphs are disabled (max_cudagraph_capture_size=0), |
| 1275 | # padding validation should be skipped |
| 1276 | config = CompilationConfig( |
| 1277 | cudagraph_capture_sizes=[], |
| 1278 | max_cudagraph_capture_size=0, |
| 1279 | compile_sizes=[3, 5, 7], # would be invalid with cudagraphs |
| 1280 | ) |
| 1281 | config.post_init_cudagraph_sizes() |
| 1282 | assert sorted(config.compile_sizes) == [3, 5, 7] |
| 1283 | |
| 1284 | # When cudagraph_mode is NONE but capture_sizes is non-empty, |
| 1285 | # padding validation should still be skipped |
| 1286 | config = CompilationConfig( |
| 1287 | cudagraph_capture_sizes=[1, 2, 4, 8], |
| 1288 | max_cudagraph_capture_size=8, |
| 1289 | cudagraph_mode=CUDAGraphMode.NONE, |
| 1290 | compile_sizes=[3, 5, 7], # would be invalid if cudagraphs were enabled |
| 1291 | ) |
| 1292 | config.post_init_cudagraph_sizes() |
| 1293 | assert sorted(config.compile_sizes) == [3, 5, 7] |
| 1294 | dispatcher = CudagraphDispatcher(_create_vllm_config_for_validation(config)) |
| 1295 | dispatcher.initialize_cudagraph_keys(CUDAGraphMode.NONE) # Should not raise |
| 1296 | |
| 1297 | |
| 1298 | def test_inductor_asserts_default_disabled(monkeypatch): |
| 1299 | """Test that inductor runtime asserts are disabled by default |
| 1300 | (INFO logging level) on torch < 2.12.""" |
| 1301 | monkeypatch.setenv("VLLM_LOGGING_LEVEL", "INFO") |
| 1302 | |
| 1303 | import importlib |
| 1304 | |
| 1305 | import vllm.envs |
| 1306 | |
| 1307 | importlib.reload(vllm.envs) |
| 1308 | |
| 1309 | config = CompilationConfig() |
| 1310 | if not _is_torch_equal_or_newer(torch.__version__, "2.12.0.dev"): |
| 1311 | assert config.inductor_compile_config.get("size_asserts") is False |
| 1312 | assert config.inductor_compile_config.get("alignment_asserts") is False |
| 1313 | assert config.inductor_compile_config.get("scalar_asserts") is False |
| 1314 | |
| 1315 | |
| 1316 | def test_inductor_asserts_enabled_in_debug(monkeypatch): |
| 1317 | """Test that VLLM_LOGGING_LEVEL=DEBUG enables inductor runtime asserts |
| 1318 | on torch < 2.12.""" |
| 1319 | monkeypatch.setenv("VLLM_LOGGING_LEVEL", "DEBUG") |
| 1320 | |
| 1321 | import importlib |
| 1322 | |
| 1323 | import vllm.envs |
| 1324 | |
| 1325 | importlib.reload(vllm.envs) |
| 1326 | |
| 1327 | config = CompilationConfig() |
| 1328 | if not _is_torch_equal_or_newer(torch.__version__, "2.12.0.dev"): |
| 1329 | assert config.inductor_compile_config.get("size_asserts") is True |
| 1330 | assert config.inductor_compile_config.get("alignment_asserts") is True |
| 1331 | assert config.inductor_compile_config.get("scalar_asserts") is True |
| 1332 | |
| 1333 | |
| 1334 | def test_get_inductor_factors_includes_configs(): |
| 1335 | """Changing inductor or functorch config must change the cache key factors.""" |
| 1336 | from torch._functorch import config as functorch_config |
| 1337 | from torch._inductor import config as inductor_config |
| 1338 | |
| 1339 | from vllm.compilation.compiler_interface import get_inductor_factors |
| 1340 | |
| 1341 | baseline = get_inductor_factors() |
| 1342 | |
| 1343 | with inductor_config.patch("max_autotune", not inductor_config.max_autotune): |
| 1344 | patched = get_inductor_factors() |
| 1345 | assert baseline != patched, "inductor config change was not reflected" |
| 1346 | |
| 1347 | with functorch_config.patch("donated_buffer", not functorch_config.donated_buffer): |
| 1348 | patched = get_inductor_factors() |
| 1349 | assert baseline != patched, "functorch config change was not reflected" |
| 1350 | |
| 1351 | |
| 1352 | def test_inductor_asserts_user_override(monkeypatch): |
| 1353 | """Test that explicit inductor_compile_config overrides the |
| 1354 | debug-logging default.""" |
| 1355 | monkeypatch.setenv("VLLM_LOGGING_LEVEL", "INFO") |
| 1356 | |
| 1357 | import importlib |
| 1358 | |
| 1359 | import vllm.envs |
| 1360 | |
| 1361 | importlib.reload(vllm.envs) |
| 1362 | |
| 1363 | config = CompilationConfig( |
| 1364 | inductor_compile_config={"size_asserts": True}, |
| 1365 | ) |
| 1366 | assert config.inductor_compile_config.get("size_asserts") is True |
| 1367 | if not _is_torch_equal_or_newer(torch.__version__, "2.12.0.dev"): |
| 1368 | assert config.inductor_compile_config.get("alignment_asserts") is False |