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, 06:00 UTC
Use xla flag to improve the quantized model performance (#19303)
Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com>
XiongfeiWei committed
Jun 9, 2025 at 18:28 UTC
9af6d22e4c8ef79972911efaab2a68f4b03def2a
1 file changed
+4
-1
vllm/v1/worker/tpu_worker.py
+4
-1
@@ -101,7 +101,10 @@ class TPUWorker:
101
# fix this. It will be removed after the bug in XLA compiler is fixed.
102
os.environ["LIBTPU_INIT_ARGS"] = (
103
os.environ.get("LIBTPU_INIT_ARGS", "") +
104
- " --xla_tpu_force_1d_allreduce_at_chunk_count=1")
104
+ " --xla_tpu_force_1d_allreduce_at_chunk_count=1"
105
+ " --xla_jf_conv_input_fusion=False")
106
+ # --xla_jf_conv_input_fusion=False is used to improve the perf of
107
+ # quantized matmul.
108
torch.set_grad_enabled(False)
109
torch.set_default_dtype(self.model_config.dtype)
110