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, 17:15 UTC
[TPU] support disabling xla compilation cache (#15567)
Signed-off-by: Chengji Yao <chengjiyao@google.com>
Chengji Yao committed
Mar 26, 2025 at 17:09 UTC
e74ff409e0f8f3cacb8a251a1cae8b478721cead
2 files changed
+20
-6
vllm/v1/worker/tpu_worker.py
+10
-3
@@ -113,9 +113,16 @@ class TPUWorker:
113
# can have slightly different XLA graphs.
114
world_size = self.parallel_config.world_size
115
rank = xr.global_ordinal()
116
- per_rank_path = os.path.join(envs.VLLM_XLA_CACHE_PATH,
117
- f"tp{world_size}_rank{rank}")
118
- xr.initialize_cache(per_rank_path, readonly=False)
116
+ # The PyTorch/XLA compilation cache uses the Torch IR to generate keys.
117
+ # Consequently, changes in optimization flags, which affect compilation
118
+ # results, don't change the cache key. This can result in the wrong
119
+ # compilation being used. To prevent this, disabling the XLA compilation
120
+ # cache during development is recommended.We can disable it by
121
+ # `export VLLM_XLA_CACHE_PATH=`
122
+ if envs.VLLM_XLA_CACHE_PATH:
123
+ per_rank_path = os.path.join(envs.VLLM_XLA_CACHE_PATH,
124
+ f"tp{world_size}_rank{rank}")
125
+ xr.initialize_cache(per_rank_path, readonly=False)
126
127
# Init ModelRunner here, so that we have access to self.device.
128
self.model_runner = TPUModelRunner(self.vllm_config, self.device)
vllm/worker/tpu_worker.py
+10
-3
@@ -93,9 +93,16 @@ class TPUWorker(LoRANotSupportedWorkerBase, LocalOrDistributedWorkerBase):
93
# can have slightly different XLA graphs.
94
world_size = self.parallel_config.world_size
95
rank = xr.global_ordinal()
96
- per_rank_path = os.path.join(envs.VLLM_XLA_CACHE_PATH,
97
- f"tp{world_size}_rank{rank}")
98
- xr.initialize_cache(per_rank_path, readonly=False)
96
+ # The PyTorch/XLA compilation cache uses the Torch IR to generate keys.
97
+ # Consequently, changes in optimization flags, which affect compilation
98
+ # results, don't change the cache key. This can result in the wrong
99
+ # compilation being used. To prevent this, disabling the XLA compilation
100
+ # cache during development is recommended.We can disable it by
101
+ # `export VLLM_XLA_CACHE_PATH=`
102
+ if envs.VLLM_XLA_CACHE_PATH:
103
+ per_rank_path = os.path.join(envs.VLLM_XLA_CACHE_PATH,
104
+ f"tp{world_size}_rank{rank}")
105
+ xr.initialize_cache(per_rank_path, readonly=False)
106
107
self.profiler = None
108
if envs.VLLM_TORCH_PROFILER_DIR and self.rank < 1: