A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 01:15 UTC
[CI][Bugfix] Make test_gpt2_cache_hit observable across V1 EngineCore (#42037)
Signed-off-by: haosdent <haosdent@gmail.com>
haosdent committed
May 9, 2026 at 11:53 UTC
e934e459e6bfdff0de6bcf255013a499434c3606
1 file changed
+41
-51
tests/compile/test_aot_compile.py
+41
-51
@@ -1,9 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-import functools
4
import hashlib
6
-import multiprocessing
5
import os
6
import pickle
7
import tempfile
@@ -15,7 +13,6 @@ import pytest
13
import torch
14
15
import vllm.envs as envs
18
-import vllm.model_executor.layers.activation
16
from vllm.compilation.backends import VllmBackend
17
from vllm.compilation.caching import (
18
StandaloneCompiledArtifacts,
@@ -476,64 +473,57 @@ def test_standalone_compile_correctness():
473
@create_new_process_for_each_test("spawn")
474
def test_gpt2_cache_hit(monkeypatch: pytest.MonkeyPatch):
475
"""
479
- Test that compiling gpt2 twice results in a cache hit and
480
- capture torch dynamic symbol creations to ensure make_symbol
481
- not called on cache hit.
482
- """
476
+ Test that compiling gpt2 twice results in a cache hit.
477
484
- import torch.fx.experimental.symbolic_shapes as symbolic_shapes_module
485
- from torch.utils._sympy.symbol import make_symbol
478
+ Counter values are read from the EngineCore subprocess via
479
+ ``LLM.collective_rpc`` so the test works under default V1
480
+ multiprocessing (no shared memory between test and engine).
481
+ """
482
483
from vllm import LLM
484
489
- create_symbol_counter = multiprocessing.Value("i", 0)
490
- original_make_symbol = make_symbol
485
+ def _snap(self):
486
+ from vllm.compilation.counter import compilation_counter
487
492
- @functools.wraps(original_make_symbol)
493
- def counting_make_symbol(prefix, idx, **kwargs):
494
- with create_symbol_counter.get_lock():
495
- create_symbol_counter.value += 1
496
- return original_make_symbol(prefix, idx, **kwargs)
497
-
498
- symbolic_shapes_module.make_symbol = counting_make_symbol
499
- try:
500
- with monkeypatch.context() as m, tempfile.TemporaryDirectory() as tmpdirname:
501
- m.setenv("VLLM_CACHE_ROOT", tmpdirname)
502
- m.setenv("VLLM_USE_AOT_COMPILE", "1")
503
- # First compilation - initialize model and generate
504
- llm_model = LLM(
505
- model="gpt2",
506
- compilation_config=CompilationConfig(
507
- mode=CompilationMode.VLLM_COMPILE,
508
- ),
509
- max_model_len=256,
510
- )
488
+ return (
489
+ compilation_counter.num_aot_compiles,
490
+ compilation_counter.num_aot_artifacts_saved,
491
+ compilation_counter.num_aot_artifacts_loaded,
492
+ )
493
512
- llm_model.generate("Hello, my name is")
513
- assert create_symbol_counter.value == 2
514
- create_symbol_counter.value = 0
494
+ # collective_rpc(callable) requires pickle-based serialization.
495
+ monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
496
516
- # Clean up first model
517
- del llm_model
518
- disable_envs_cache()
519
- vllm.model_executor.layers.activation._ACTIVATION_REGISTRY._dict.clear()
497
+ with monkeypatch.context() as m, tempfile.TemporaryDirectory() as tmpdirname:
498
+ m.setenv("VLLM_CACHE_ROOT", tmpdirname)
499
+ m.setenv("VLLM_USE_AOT_COMPILE", "1")
500
+ # First compilation - initialize model and generate
501
+ llm_model = LLM(
502
+ model="gpt2",
503
+ compilation_config=CompilationConfig(
504
+ mode=CompilationMode.VLLM_COMPILE,
505
+ ),
506
+ max_model_len=256,
507
+ )
508
521
- # Second compilation - should hit cache
522
- m.setenv("VLLM_FORCE_AOT_LOAD", "1")
523
- llm_model = LLM(
524
- model="gpt2",
525
- compilation_config=CompilationConfig(
526
- mode=CompilationMode.VLLM_COMPILE,
527
- ),
528
- max_model_len=256,
529
- )
530
- llm_model.generate("Hello, my name is")
509
+ llm_model.generate("Hello, my name is")
510
+ assert llm_model.collective_rpc(_snap)[0] == (1, 1, 0)
511
532
- assert create_symbol_counter.value == 0
512
+ # Clean up first model
513
+ del llm_model
514
+ disable_envs_cache()
515
534
- finally:
535
- # Restore original method
536
- symbolic_shapes_module.make_symbol = original_make_symbol
516
+ # Second compilation - should hit cache
517
+ m.setenv("VLLM_FORCE_AOT_LOAD", "1")
518
+ llm_model = LLM(
519
+ model="gpt2",
520
+ compilation_config=CompilationConfig(
521
+ mode=CompilationMode.VLLM_COMPILE,
522
+ ),
523
+ max_model_len=256,
524
+ )
525
+ llm_model.generate("Hello, my name is")
526
+ assert llm_model.collective_rpc(_snap)[0] == (0, 0, 1)
527
528
529
@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")