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][V1][Bugfix] Fix w8a8 recompiilation with GSM8K (#15714)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Mar 29, 2025 at 05:13 UTC
da461f3cbf8be4094a6f14a1eaf89b5931f3625f
4 files changed
+16
-15
.buildkite/run-tpu-v1-test.sh
+4
-6
@@ -28,16 +28,14 @@ docker run --privileged --net host --shm-size=16G -it \
28
&& echo TEST_3 \
29
&& pytest -v -s /workspace/vllm/tests/entrypoints/llm/test_accuracy.py::test_lm_eval_accuracy_v1_engine \
30
&& echo TEST_4 \
31
- && python3 /workspace/vllm/examples/offline_inference/tpu.py \
31
+ && pytest -s -v /workspace/vllm/tests/tpu/test_quantization_accuracy.py \
32
&& echo TEST_5 \
33
- && pytest -s -v /workspace/vllm/tests/v1/tpu/worker/test_tpu_model_runner.py \
33
+ && python3 /workspace/vllm/examples/offline_inference/tpu.py \
34
&& echo TEST_6 \
35
+ && pytest -s -v /workspace/vllm/tests/v1/tpu/worker/test_tpu_model_runner.py \
36
+ && echo TEST_7 \
37
&& pytest -s -v /workspace/vllm/tests/v1/tpu/test_sampler.py" \
38
39
40
# TODO: This test fails because it uses RANDOM_SEED sampling
41
# && VLLM_USE_V1=1 pytest -v -s /workspace/vllm/tests/tpu/test_custom_dispatcher.py \
40
-
41
-# TODO: Re-enable this after fixing recompilation in quantization.
42
-# && echo TEST_4 \
43
-# && pytest -s -v /workspace/vllm/tests/tpu/test_quantization_accuracy.py \
vllm/model_executor/layers/quantization/kernels/scaled_mm/xla.py
+2
-1
@@ -97,7 +97,8 @@ class XLAScaledMMLinearKernel(ScaledMMLinearKernel):
97
block_size=-1,
98
int4_weight=False,
99
quantize_activation=True)
100
-
100
+ # `quantized_matmul` output is fp32, cast it down to bf16 for perf
101
+ out = out.to(x.dtype)
102
# Explicitly capture control flow to make dynamo happy.
103
# https://pytorch.org/docs/main/generated/exportdb/index.html#cond-branch-class-method # noqa: E501
104
return cond(bias is None, self.no_add_bias, self.add_bias, [out, bias])
vllm/v1/worker/tpu_model_runner.py
+8
-6
@@ -80,6 +80,7 @@ class TPUModelRunner:
80
self.enforce_eager = model_config.enforce_eager
81
self.pin_memory = is_pin_memory_available()
82
self.dtype = self.model_config.dtype
83
+ self._hidden_states_dtype = self.dtype
84
85
self.is_multimodal_model = model_config.is_multimodal_model
86
self.sliding_window = model_config.get_sliding_window()
@@ -771,10 +772,11 @@ class TPUModelRunner:
772
torch._dynamo.mark_dynamic(attn_metadata.slot_mapping, 0)
773
774
with set_forward_context(attn_metadata, self.vllm_config, 0):
774
- self.model(input_ids=input_ids,
775
- positions=position_ids,
776
- kv_caches=kv_caches,
777
- inputs_embeds=inputs_embeds)
775
+ out = self.model(input_ids=input_ids,
776
+ positions=position_ids,
777
+ kv_caches=kv_caches,
778
+ inputs_embeds=inputs_embeds)
779
+ self._hidden_states_dtype = out.dtype
780
781
def capture_model(self) -> None:
782
"""Compile the model."""
@@ -800,7 +802,7 @@ class TPUModelRunner:
802
num_reqs_to_sample = MIN_NUM_SEQS
803
dummy_hidden = torch.randn((num_tokens, hsize),
804
device=device,
803
- dtype=torch.bfloat16)
805
+ dtype=self._hidden_states_dtype)
806
# Compile for [8, 16, .., 128,.., `self.max_num_reqs`]
807
while True:
808
indices = torch.zeros(
@@ -823,7 +825,7 @@ class TPUModelRunner:
825
num_reqs_to_sample + 1, self.max_num_reqs)
826
xm.wait_device_ops()
827
end = time.perf_counter()
826
- logger.info("Compilation finished in in %.2f [secs].", end - start)
828
+ logger.info("Compilation finished in %.2f [secs].", end - start)
829
# Record the number cached XLA graph after warming up, this will be
830
# used for checking there is no additional graph compilation during
831
# runtime execution.
vllm/v1/worker/tpu_worker.py
+2
-2
@@ -105,8 +105,8 @@ class TPUWorker:
105
106
# Increase the cache size limit, which is the maximum number of
107
# dynamo graphs that can be compiled.
108
- # NOTE(woosuk): Usually, we compile 10-15 graphs for prefill and
109
- # 30-40 graphs for decode. 128 is an arbitrary safe number.
108
+ # TODO (NickLucche) On gsm we compile 80+ graphs.
109
+ # Re-evaluate limit, with MM we may get close to this limit.
110
torch._dynamo.config.cache_size_limit = 128
111
# Use persistent cache to avoid XLA recompilation.
112
# NOTE(woosuk): Set per-rank cache path since different ranks