iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[CI] Prune PyTorch Compilation Unit Tests (#51069)

Browse Source 
Signed-off-by: mgoin <mgoin64@gmail.com>
This commit is contained in:
Michael Goin
2026-08-05 09:53:50 +08:00
committed by GitHub
parent 77b519912b
commit d16e7f9da4
3 changed files with 124 additions and 172 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/test_areas/pytorch.yaml
+2 -3
View File
@@ -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
tests/compile/test_aot_compile.py
+21 -101
View File
@@ -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)
tests/compile/test_dynamic_shapes_compilation.py
+101 -68
View File
@@ -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,