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 15, 18:30 UTC
[Hardware][TPU] Add check for no additional graph compilation during runtime (#14710)
Signed-off-by: Siyuan Liu <lsiyuan@google.com>
Siyuan Liu committed
Mar 20, 2025 at 20:05 UTC
b15fd2be2acae28b25253840588f3eb135a3f709
3 files changed
+32
-6
.buildkite/run-tpu-v1-test.sh
+8
-6
@@ -19,17 +19,19 @@ docker run --privileged --net host --shm-size=16G -it \
19
vllm-tpu /bin/bash -c "python3 -m pip install git+https://github.com/thuml/depyf.git \
20
&& python3 -m pip install pytest \
21
&& python3 -m pip install lm_eval[api]==0.4.4 \
22
+ && export VLLM_USE_V1=1 \
23
+ && export VLLM_XLA_CHECK_RECOMPILATION=1 \
24
&& echo TEST_1 \
23
- && VLLM_USE_V1=1 python3 /workspace/vllm/tests/tpu/test_compilation.py \
25
+ && python3 /workspace/vllm/tests/tpu/test_compilation.py \
26
&& echo TEST_2 \
25
- && VLLM_USE_V1=1 pytest -v -s /workspace/vllm/tests/v1/tpu/test_basic.py \
27
+ && pytest -v -s /workspace/vllm/tests/v1/tpu/test_basic.py \
28
&& echo TEST_3 \
27
- && VLLM_USE_V1=1 pytest -v -s /workspace/vllm/tests/entrypoints/llm/test_accuracy.py::test_lm_eval_accuracy_v1_engine \
29
+ && pytest -v -s /workspace/vllm/tests/entrypoints/llm/test_accuracy.py::test_lm_eval_accuracy_v1_engine \
30
&& echo TEST_4 \
29
- && VLLM_USE_V1=1 pytest -s -v /workspace/vllm/tests/tpu/test_quantization_accuracy.py \
31
+ && pytest -s -v /workspace/vllm/tests/tpu/test_quantization_accuracy.py \
32
&& echo TEST_5 \
31
- && VLLM_USE_V1=1 python3 /workspace/vllm/examples/offline_inference/tpu.py" \
32
-
33
+ && python3 /workspace/vllm/examples/offline_inference/tpu.py" \
34
+
35
36
# TODO: This test fails because it uses RANDOM_SEED sampling
37
# && VLLM_USE_V1=1 pytest -v -s /workspace/vllm/tests/tpu/test_custom_dispatcher.py \
vllm/envs.py
+5
@@ -45,6 +45,7 @@ if TYPE_CHECKING:
45
VLLM_OPENVINO_CPU_KV_CACHE_PRECISION: Optional[str] = None
46
VLLM_OPENVINO_ENABLE_QUANTIZED_WEIGHTS: bool = False
47
VLLM_XLA_CACHE_PATH: str = os.path.join(VLLM_CACHE_ROOT, "xla_cache")
48
+ VLLM_XLA_CHECK_RECOMPILATION: bool = False
49
VLLM_FUSED_MOE_CHUNK_SIZE: int = 64 * 1024
50
VLLM_USE_RAY_SPMD_WORKER: bool = False
51
VLLM_USE_RAY_COMPILED_DAG: bool = False
@@ -446,6 +447,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
447
"VLLM_XLA_CACHE_PATH",
448
os.path.join(get_default_cache_root(), "vllm", "xla_cache"),
449
)),
450
+
451
+ # If set, assert on XLA recompilation after each execution step.
452
+ "VLLM_XLA_CHECK_RECOMPILATION":
453
+ lambda: bool(int(os.getenv("VLLM_XLA_CHECK_RECOMPILATION", "0"))),
454
"VLLM_FUSED_MOE_CHUNK_SIZE":
455
lambda: int(os.getenv("VLLM_FUSED_MOE_CHUNK_SIZE", "32768")),
456
vllm/v1/worker/tpu_model_runner.py
+19
@@ -11,6 +11,7 @@ import torch.nn as nn
11
import torch_xla.core.xla_model as xm
12
import torch_xla.runtime as xr
13
14
+import vllm.envs as envs
15
from vllm.attention.backends.abstract import AttentionType
16
from vllm.attention.layer import Attention
17
from vllm.config import VllmConfig
@@ -73,6 +74,10 @@ class TPUModelRunner:
74
scheduler_config = self.scheduler_config
75
parallel_config = self.parallel_config
76
self.device = device
77
+ self.check_recompilation = envs.VLLM_XLA_CHECK_RECOMPILATION
78
+ if self.check_recompilation:
79
+ self.num_xla_graphs = xr.get_num_cached_compilation_graph()
80
+ self.enforce_eager = model_config.enforce_eager
81
self.pin_memory = is_pin_memory_available()
82
self.dtype = self.model_config.dtype
83
@@ -671,6 +676,12 @@ class TPUModelRunner:
676
logprobs=None,
677
prompt_logprobs_dict=prompt_logprobs_dict,
678
)
679
+ # Check there is no new graph compilation, all the graphs should be
680
+ # captured and compiled during warming up.
681
+ if self.check_recompilation and not self.enforce_eager:
682
+ curr_cached_graph = xr.get_num_cached_compilation_graph()
683
+ assert self.num_xla_graphs == curr_cached_graph, (
684
+ "Recompilation after warm up is detected.")
685
return model_runner_output
686
687
def load_model(self) -> None:
@@ -810,6 +821,14 @@ class TPUModelRunner:
821
xm.wait_device_ops()
822
end = time.perf_counter()
823
logger.info("Compilation finished in in %.2f [secs].", end - start)
824
+ # Record the number cached XLA graph after warming up, this will be
825
+ # used for checking there is no additional graph compilation during
826
+ # runtime execution.
827
+ if self.check_recompilation:
828
+ total_cached_graphs = xr.get_num_cached_compilation_graph()
829
+ num_compiled_graphs = total_cached_graphs - self.num_xla_graphs
830
+ logger.info("Compiled %d XLA graphs.", num_compiled_graphs)
831
+ self.num_xla_graphs += num_compiled_graphs
832
833
def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
834
"""