diff --git a/.buildkite/test_areas/pytorch.yaml b/.buildkite/test_areas/pytorch.yaml index bf8014f6dd..b2bf2b0523 100644 --- a/.buildkite/test_areas/pytorch.yaml +++ b/.buildkite/test_areas/pytorch.yaml @@ -5,7 +5,7 @@ steps: - label: PyTorch Compilation Unit Tests device: h200_35gb key: pytorch-compilation-unit-tests - timeout_in_minutes: 150 + timeout_in_minutes: 60 source_file_dependencies: - vllm/__init__.py - vllm/_aiter_ops.py @@ -45,8 +45,7 @@ steps: - label: PyTorch Compilation Unit Tests (H100) key: pytorch-compilation-unit-tests-h100 timeout_in_minutes: 30 - device: h100 - num_devices: 1 + device: h200_18gb source_file_dependencies: - vllm/__init__.py - vllm/_aiter_ops.py diff --git a/tests/compile/test_aot_compile.py b/tests/compile/test_aot_compile.py index a7f32483a7..b3a71b4c30 100644 --- a/tests/compile/test_aot_compile.py +++ b/tests/compile/test_aot_compile.py @@ -43,7 +43,7 @@ def vllm_tmp_cache(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: def reference_fn(x: torch.Tensor): assert x.shape[0] <= 42 assert x.shape[0] % 2 == 0 - for _ in range(3000): + for _ in range(30): x = x + x.shape[0] return x @@ -52,7 +52,7 @@ def reference_fn_tuple(x: torch.Tensor): """Reference function that returns a tuple of tensors.""" assert x.shape[0] <= 42 assert x.shape[0] % 2 == 0 - for _ in range(3000): + for _ in range(30): x = x + x.shape[0] return x, x * 2 @@ -140,18 +140,35 @@ def test_save_and_load(monkeypatch: pytest.MonkeyPatch): m.setenv("VLLM_USE_AOT_COMPILE", "1") m.setenv("VLLM_USE_MEGA_AOT_ARTIFACT", "1") m.setenv("VLLM_USE_STANDALONE_COMPILE", "1") + disable_envs_cache() vllm_config = make_vllm_config() - with use_vllm_config(vllm_config): + with ( + use_vllm_config(vllm_config), + compilation_counter.expect( + num_aot_compiles=1, + num_aot_artifacts_saved=1, + num_aot_artifacts_loaded=0, + ), + ): compiled_mod = CompiledMod(vllm_config=vllm_config) expected = compiled_mod(*args) + assert isinstance(expected, torch.Tensor) disable_envs_cache() m.setenv("VLLM_FORCE_AOT_LOAD", "1") vllm_config = make_vllm_config() - with use_vllm_config(vllm_config): + with ( + use_vllm_config(vllm_config), + compilation_counter.expect( + num_aot_compiles=0, + num_aot_artifacts_saved=0, + num_aot_artifacts_loaded=1, + ), + ): cached_mod = CompiledMod(vllm_config=vllm_config) ret = cached_mod(*args) + assert isinstance(ret, torch.Tensor) assert cached_mod.was_aot_compile_fn_loaded_from_disk, ( "Expected was_aot_compile_fn_loaded_from_disk to be True" ) @@ -182,65 +199,6 @@ def test_save_and_load_slice(monkeypatch: pytest.MonkeyPatch): assert gm.code == loaded_gm.code -@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10") -def test_cache_load_returns_tuple_consistency(monkeypatch: pytest.MonkeyPatch): - """ - Test that cache loading correctly handles the returns_tuple logic. - - This verifies that when a model returns a single tensor (not a tuple), - the output type is consistent between fresh compilation and cache load. - Without the fix, cached artifacts would return [tensor] instead of tensor. - """ - with monkeypatch.context() as m: - args = (torch.randn(10, 10),) - - with tempfile.TemporaryDirectory() as tmpdirname: - m.setenv("VLLM_CACHE_ROOT", tmpdirname) - m.setenv("VLLM_USE_AOT_COMPILE", "1") - m.setenv("VLLM_USE_MEGA_AOT_ARTIFACT", "1") - m.setenv("VLLM_USE_STANDALONE_COMPILE", "1") - vllm_config = make_vllm_config() - - # Fresh compilation - with use_vllm_config(vllm_config): - compiled_mod = CompiledMod(vllm_config=vllm_config) - fresh_result = compiled_mod(*args) - fresh_result_type = type(fresh_result) - - # Verify fresh result is a tensor, not a tuple/list - assert isinstance(fresh_result, torch.Tensor), ( - f"Fresh compile should return tensor, got {fresh_result_type}" - ) - - disable_envs_cache() - - # Load from cache - m.setenv("VLLM_FORCE_AOT_LOAD", "1") - vllm_config = make_vllm_config() - with use_vllm_config(vllm_config): - cached_mod = CompiledMod(vllm_config=vllm_config) - cached_result = cached_mod(*args) - cached_result_type = type(cached_result) - - # Verify cache was actually loaded - assert cached_mod.was_aot_compile_fn_loaded_from_disk, ( - "Expected was_aot_compile_fn_loaded_from_disk to be True after " - "loading from cache" - ) - - # Verify cached result has same type as fresh result - assert isinstance(cached_result, torch.Tensor), ( - f"Cache load should return tensor, got {cached_result_type}. " - "This indicates the returns_tuple logic is not being applied " - "correctly when loading from cache." - ) - - # Verify values match - assert torch.allclose(cached_result, fresh_result), ( - "Cached result values should match fresh compilation" - ) - - @pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10") def test_cache_load_returns_tuple_consistency_tuple_output( monkeypatch: pytest.MonkeyPatch, @@ -872,41 +830,3 @@ def test_disable_compile_cache_skips_aot_load( mod(*args) assert not mod.was_aot_compile_fn_loaded_from_disk - - -@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10") -def test_aot_counters_on_save_and_load( - monkeypatch: pytest.MonkeyPatch, fresh_vllm_cache: str -): - """Verify AOT counters are incremented correctly on save and load.""" - monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "1") - disable_envs_cache() - - args = (torch.randn(10, 10),) - - # Phase 1: fresh compile + save - vllm_config = make_vllm_config() - with ( - use_vllm_config(vllm_config), - compilation_counter.expect( - num_aot_compiles=1, - num_aot_artifacts_saved=1, - num_aot_artifacts_loaded=0, - ), - ): - CompiledMod(vllm_config=vllm_config)(*args) - - # Phase 2: load from cache - monkeypatch.setenv("VLLM_FORCE_AOT_LOAD", "1") - disable_envs_cache() - - vllm_config = make_vllm_config() - with ( - use_vllm_config(vllm_config), - compilation_counter.expect( - num_aot_compiles=0, - num_aot_artifacts_saved=0, - num_aot_artifacts_loaded=1, - ), - ): - CompiledMod(vllm_config=vllm_config)(*args) diff --git a/tests/compile/test_dynamic_shapes_compilation.py b/tests/compile/test_dynamic_shapes_compilation.py index 3260d5aecc..71f66fce40 100644 --- a/tests/compile/test_dynamic_shapes_compilation.py +++ b/tests/compile/test_dynamic_shapes_compilation.py @@ -3,6 +3,7 @@ import tempfile from contextlib import contextmanager +from dataclasses import dataclass import pytest import torch @@ -23,40 +24,101 @@ from vllm.utils.torch_utils import is_torch_equal_or_newer def get_test_models(): """Get list of models to test based on PyTorch version""" models = [ + "Qwen/Qwen3-0.6B", "openai-community/gpt2", - "Qwen/Qwen2-7B-Instruct", - "meta-llama/Llama-3.1-8B", ] - if is_torch_equal_or_newer("2.12.0.dev"): - models.append("Qwen/Qwen3-4B-Instruct-2507") return models -@pytest.mark.parametrize("model_name", get_test_models()) -@pytest.mark.parametrize( - "shapes_type", - [ - DynamicShapesType.BACKED, - DynamicShapesType.UNBACKED, - DynamicShapesType.BACKED_SIZE_OBLIVIOUS, - ], -) -@pytest.mark.parametrize("use_aot_compile", ["0", "1"]) -@pytest.mark.parametrize("use_bytecode_hook", [True, False]) -@pytest.mark.parametrize("evaluate_guards", [False, True]) +@dataclass(frozen=True) +class DynamicShapesTestCase: + model_name: str + shapes_type: DynamicShapesType + use_aot_compile: bool + use_bytecode_hook: bool + + def __str__(self) -> str: + model_id = self.model_name.rsplit("/", 1)[-1] + return ( + f"{model_id}-{self.shapes_type.value}-" + f"aot{int(self.use_aot_compile)}-hook{int(self.use_bytecode_hook)}" + ) + + +def get_dynamic_shapes_test_cases(): + """Pairwise-cover compilation options and smoke-test model classes.""" + reference_model = "Qwen/Qwen3-0.6B" + cases = [ + DynamicShapesTestCase(reference_model, DynamicShapesType.BACKED, False, True), + DynamicShapesTestCase(reference_model, DynamicShapesType.BACKED, True, False), + DynamicShapesTestCase( + reference_model, DynamicShapesType.UNBACKED, False, False + ), + DynamicShapesTestCase(reference_model, DynamicShapesType.UNBACKED, True, True), + DynamicShapesTestCase( + reference_model, + DynamicShapesType.BACKED_SIZE_OBLIVIOUS, + False, + True, + ), + DynamicShapesTestCase( + reference_model, + DynamicShapesType.BACKED_SIZE_OBLIVIOUS, + True, + False, + ), + ] + cases.extend( + DynamicShapesTestCase(model_name, DynamicShapesType.BACKED, False, True) + for model_name in get_test_models() + if model_name != reference_model + ) + return cases + + +TEST_PROMPTS = ["Hello, my name is", "The capital of France is"] + + +def generate_outputs(vllm_model): + sampling_params = SamplingParams(max_tokens=5, temperature=0, logprobs=10) + outputs = [] + for prompt in TEST_PROMPTS: + output = vllm_model.llm.generate(prompt, sampling_params)[0].outputs[0] + assert len(output.text.strip()) > 0, "Model produced empty output" + outputs.append((output.token_ids, output.text, output.logprobs)) + return outputs + + +@pytest.fixture(scope="module") +def get_eager_outputs(vllm_runner): + cache = {} + + def get(model_name): + if model_name not in cache: + with vllm_runner( + model_name, + enforce_eager=True, + max_model_len=1024, + enable_chunked_prefill=None, + ) as vllm_model: + cache[model_name] = generate_outputs(vllm_model) + return cache[model_name] + + return get + + +@pytest.mark.parametrize("test_case", get_dynamic_shapes_test_cases(), ids=str) @pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10") def test_dynamic_shapes_compilation( monkeypatch, vllm_runner, - model_name, - shapes_type, - use_aot_compile, - use_bytecode_hook, - evaluate_guards, + get_eager_outputs, + test_case, ): - """Test that all dynamic shapes types compile successfully""" - if shapes_type == DynamicShapesType.UNBACKED and not is_torch_equal_or_newer( - "2.11.0" + """Test representative dynamic-shapes configurations end to end.""" + if ( + test_case.shapes_type == DynamicShapesType.UNBACKED + and not is_torch_equal_or_newer("2.11.0") ): # NOTE[ROCm]: shape_id (used by Qwen2/Llama to relate input dims) only # landed in torch 2.11, but the ROCm CI still runs torch 2.10.x. On @@ -64,64 +126,38 @@ def test_dynamic_shapes_compilation( # data-dependent and compilation blows up -- nothing to test. pytest.skip("unbacked dynamic shapes with shape_id require torch>=2.11") - if evaluate_guards and shapes_type == DynamicShapesType.UNBACKED: - pytest.skip("unbacked dynamic shapes do not add guards") + monkeypatch.setenv( + "VLLM_USE_AOT_COMPILE", "1" if test_case.use_aot_compile else "0" + ) + monkeypatch.setenv( + "VLLM_USE_BYTECODE_HOOK", "1" if test_case.use_bytecode_hook else "0" + ) - # TODO is this still a requirement? - if evaluate_guards and use_aot_compile: - pytest.skip("evaluate_guards requires use_aot_compile=0") + print(f"Testing {test_case.shapes_type.name} dynamic shapes...") - monkeypatch.setenv("VLLM_USE_AOT_COMPILE", use_aot_compile) - monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "1" if use_bytecode_hook else "0") - - prompt = "Hello, my name is" - - print(f"Testing {shapes_type.name} dynamic shapes...") - - sampling_params = SamplingParams(max_tokens=5, temperature=0, logprobs=10) - test_prompts = [prompt, "The capital of France is"] - - # VllmRunner shuts down the engine core on exit, so the eager model - # below never races a lingering compiled engine for GPU memory. + # The compiled engine shuts down before an uncached eager baseline starts. with vllm_runner( - model_name, + test_case.model_name, compilation_config={ "mode": CompilationMode.VLLM_COMPILE, "dynamic_shapes_config": { - "type": shapes_type.value, - "evaluate_guards": evaluate_guards, + "type": test_case.shapes_type.value, }, }, max_model_len=1024, enable_chunked_prefill=None, ) as vllm_model: - compiled_outputs = [] - for p in test_prompts: - output = vllm_model.llm.generate(p, sampling_params)[0].outputs[0] - assert len(output.text.strip()) > 0, "Compiled model produced empty output" - compiled_outputs.append((output.token_ids, output.text, output.logprobs)) - - with vllm_runner( - model_name, - enforce_eager=True, - max_model_len=1024, - enable_chunked_prefill=None, - ) as vllm_model: - eager_outputs = [] - for p in test_prompts: - output = vllm_model.llm.generate(p, sampling_params)[0].outputs[0] - assert len(output.text.strip()) > 0, "Eager model produced empty output" - eager_outputs.append((output.token_ids, output.text, output.logprobs)) + compiled_outputs = generate_outputs(vllm_model) check_logprobs_close( - outputs_0_lst=eager_outputs, + outputs_0_lst=get_eager_outputs(test_case.model_name), outputs_1_lst=compiled_outputs, name_0="eager", name_1="compiled", ) -@pytest.mark.parametrize("use_aot_compile", ["0", "1"]) +@pytest.mark.parametrize("use_aot_compile", [False, True]) @pytest.mark.parametrize( "dynamic_shapes_type", [ @@ -138,7 +174,7 @@ def test_model_specialization_with_evaluate_guards( """ if ( - use_aot_compile == "1" + use_aot_compile and dynamic_shapes_type == DynamicShapesType.BACKED and evaluate_guards ): @@ -177,12 +213,9 @@ def test_model_specialization_with_evaluate_guards( yield monkeypatch.setenv("TOKENIZERS_PARALLELISM", "true") - monkeypatch.setenv("VLLM_USE_AOT_COMPILE", use_aot_compile) + monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "1" if use_aot_compile else "0") monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "0") - # Create vllm config with the desired settings - from vllm.config import CompilationMode - vllm_config = VllmConfig( compilation_config=CompilationConfig( mode=CompilationMode.VLLM_COMPILE,