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, 01:00 UTC
[TPU] support fp8 kv cache quantization (#19292)
Signed-off-by: Chengji Yao <chengjiyao@google.com>
Chengji Yao committed
Jul 19, 2025 at 20:01 UTC
3a1d8940aea57999411b7ea47287d3ad5cb71676
6 files changed
+95
-28
tests/entrypoints/llm/test_accuracy.py
+30
-10
@@ -15,15 +15,18 @@ import pytest
15
from vllm.platforms import current_platform
16
17
MODEL_NAMES = [
18
- "Qwen/Qwen2-1.5B-Instruct",
18
+ "Qwen/Qwen3-1.7B",
19
"google/gemma-3-1b-it",
20
]
21
+FP8_KV_MODEL_NAMES = [
22
+ "Qwen/Qwen3-1.7B",
23
+]
24
NUM_CONCURRENT = 500
25
TASK = "gsm8k"
26
FILTER = "exact_match,strict-match"
27
RTOL = 0.03
28
EXPECTED_VALUES = {
26
- "Qwen/Qwen2-1.5B-Instruct": 0.58,
29
+ "Qwen/Qwen3-1.7B": 0.68,
30
"google/gemma-3-1b-it": 0.25,
31
}
32
@@ -70,10 +73,9 @@ def test_lm_eval_accuracy_v1_engine(model, monkeypatch: pytest.MonkeyPatch):
73
if current_platform.is_tpu():
74
# Limit compilation time for TPU V1
75
73
- if model == "google/gemma-3-1b-it":
74
- # TPU + google/gemma-3-1b-it + xet doesn't work well.
75
- m.setenv("HF_HUB_DISABLE_XET", "1")
76
-
76
+ # xet doesn't work well for both Qwen/Qwen3-1.7B and
77
+ # google/gemma-3-1b-it
78
+ m.setenv("HF_HUB_DISABLE_XET", "1")
79
more_args = "max_model_len=2048,max_num_seqs=64"
80
81
# Add TP test (if provided)
@@ -83,9 +85,27 @@ def test_lm_eval_accuracy_v1_engine(model, monkeypatch: pytest.MonkeyPatch):
85
run_test(model, more_args)
86
87
86
-def test_lm_eval_accuracy_v0_engine(monkeypatch: pytest.MonkeyPatch):
87
- """Run with the V0 Engine."""
88
+@pytest.mark.skipif(not current_platform.is_cuda()
89
+ and not current_platform.is_tpu(),
90
+ reason="V1 is currently only supported on CUDA and TPU")
91
+@pytest.mark.parametrize("model", FP8_KV_MODEL_NAMES)
92
+def test_lm_eval_accuracy_v1_engine_fp8_kv_cache(
93
+ model, monkeypatch: pytest.MonkeyPatch):
94
+ """Run with the V1 Engine."""
95
96
with monkeypatch.context() as m:
90
- m.setenv("VLLM_USE_V1", "0")
91
- run_test("Qwen/Qwen2-1.5B-Instruct")
97
+ m.setenv("VLLM_USE_V1", "1")
98
+
99
+ more_args = None
100
+ if current_platform.is_tpu():
101
+ # Limit compilation time for TPU V1
102
+
103
+ # xet doesn't work well for Qwen/Qwen3-1.7B
104
+ m.setenv("HF_HUB_DISABLE_XET", "1")
105
+ more_args = "max_model_len=2048,max_num_seqs=128,kv_cache_dtype=fp8"
106
+
107
+ # Add TP test (if provided)
108
+ if TPU_TP_TEST_STR:
109
+ more_args += ",{}".format(TPU_TP_TEST_STR)
110
+
111
+ run_test(model, more_args)
tests/v1/tpu/test_pallas.py
+2
@@ -95,4 +95,6 @@ def test_ragged_paged_attention():
95
sm_scale=scale,
96
sliding_window=sliding_window,
97
soft_cap=logits_soft_cap,
98
+ k_scale=1.0,
99
+ v_scale=1.0,
100
)
vllm/engine/arg_utils.py
+4
-4
@@ -1358,10 +1358,10 @@ class EngineArgs:
1358
and not envs.is_set("VLLM_ATTENTION_BACKEND")
1359
) or envs.VLLM_ATTENTION_BACKEND == "FLASH_ATTN_VLLM_V1"
1360
supported = False
1361
- if current_platform.is_rocm() or (
1362
- current_platform.is_cuda()
1363
- and current_platform.is_device_capability(100)
1364
- ): # handle hpu also for OOT platform
1361
+ if (current_platform.is_rocm()
1362
+ or (current_platform.is_cuda()
1363
+ and current_platform.is_device_capability(100))
1364
+ or current_platform.is_tpu()):
1365
supported = True
1366
elif fp8_attention and will_use_fa:
1367
from vllm.attention.utils.fa_utils import (
vllm/platforms/tpu.py
+3
-1
@@ -35,7 +35,9 @@ class TpuPlatform(Platform):
35
device_control_env_var: str = "TPU_VISIBLE_CHIPS"
36
simple_compile_backend: str = "openxla"
37
38
- supported_quantization: list[str] = ["tpu_int8", "compressed-tensors"]
38
+ supported_quantization: list[str] = [
39
+ "fp8", "tpu_int8", "compressed-tensors"
40
+ ]
41
42
additional_env_vars: list[str] = [
43
"TPU_CHIPS_PER_HOST_BOUNDS", "TPU_HOST_BOUNDS"
vllm/v1/attention/backends/pallas.py
+50
-8
@@ -24,6 +24,19 @@ logger = init_logger(__name__)
24
# TPU requires the head size to be a multiple of 128.
25
TPU_HEAD_SIZE_ALIGNMENT = 128
26
27
+# Note: TPU can fp8 as storage dtype but doesn't support converting from uint8
28
+# from to fp32 directly. That's why it has a dtype mapping different from GPU
29
+TPU_STR_DTYPE_TO_TORCH_DTYPE = {
30
+ "half": torch.half,
31
+ "bfloat16": torch.bfloat16,
32
+ "float": torch.float,
33
+ "fp8": torch.float8_e4m3fn,
34
+ "fp8_e4m3": torch.float8_e4m3fn,
35
+ "fp8_e5m2": torch.float8_e5m2,
36
+ "int8": torch.int8,
37
+ "uint8": torch.uint8,
38
+}
39
+
40
41
class PallasAttentionBackend(AttentionBackend):
42
@@ -152,8 +165,6 @@ class PallasAttentionBackendImpl(AttentionImpl):
165
self.num_queries_per_kv = self.num_heads // self.num_kv_heads
166
if alibi_slopes is not None:
167
raise NotImplementedError("Alibi slopes is not supported.")
155
- if kv_cache_dtype != "auto":
156
- raise NotImplementedError("FP8 KV cache dtype is not supported.")
168
169
if attn_type != AttentionType.DECODER:
170
raise NotImplementedError("Encoder self-attention and "
@@ -161,6 +172,11 @@ class PallasAttentionBackendImpl(AttentionImpl):
172
"are not implemented for "
173
"PallasAttentionBackendImpl")
174
175
+ self.kv_cache_quantized_dtype = None
176
+ if kv_cache_dtype != "auto":
177
+ self.kv_cache_quantized_dtype = TPU_STR_DTYPE_TO_TORCH_DTYPE.get(
178
+ kv_cache_dtype.lower().strip())
179
+
180
def forward(
181
self,
182
layer: AttentionLayer,
@@ -194,7 +210,6 @@ class PallasAttentionBackendImpl(AttentionImpl):
210
output = torch.ones_like(query)
211
return output
212
197
- assert layer._k_scale_float == 1.0 and layer._v_scale_float == 1.0
213
num_tokens, hidden_size = query.shape
214
query = query.view(num_tokens, self.num_heads, self.head_size)
215
key = key.view(-1, self.num_kv_heads, self.head_size)
@@ -215,10 +230,21 @@ class PallasAttentionBackendImpl(AttentionImpl):
230
# Skip this if sharing KV cache with an earlier attention layer.
231
slot_mapping = attn_metadata.slot_mapping
232
write_to_kv_cache(
218
- key, value, kv_cache, slot_mapping,
233
+ key,
234
+ value,
235
+ kv_cache,
236
+ slot_mapping,
237
attn_metadata.num_slices_per_kv_cache_update_block,
220
- attn_metadata.num_kv_update_slices)
221
-
238
+ attn_metadata.num_kv_update_slices,
239
+ self.kv_cache_quantized_dtype,
240
+ layer._k_scale_float,
241
+ layer._v_scale_float,
242
+ )
243
+
244
+ if self.kv_cache_quantized_dtype is not None and (
245
+ layer._k_scale_float == 0.0 or layer._v_scale_float == 0.0):
246
+ raise ValueError(
247
+ "k_scale_float and v_scale_float must be non-zero")
248
output = torch.ops.xla.ragged_paged_attention(
249
query,
250
kv_cache,
@@ -236,6 +262,8 @@ class PallasAttentionBackendImpl(AttentionImpl):
262
sm_scale=self.scale,
263
sliding_window=self.sliding_window,
264
soft_cap=self.logits_soft_cap,
265
+ k_scale=layer._k_scale_float,
266
+ v_scale=layer._v_scale_float,
267
)
268
269
if self.head_size % TPU_HEAD_SIZE_ALIGNMENT != 0:
@@ -251,18 +279,32 @@ def write_to_kv_cache(
279
slot_mapping: torch.Tensor,
280
num_slices_per_kv_cache_update_block: int,
281
num_kv_update_slices: torch.Tensor,
282
+ kv_cache_quantized_dtype: Optional[torch.dtype] = None,
283
+ k_scale: float = 1.0,
284
+ v_scale: float = 1.0,
285
) -> None:
286
""" Write the key and values to the KV cache.
287
288
Args:
258
- key: shape = [num_tokens, num_kv_heads * head_size]
259
- value: shape = [num_tokens, num_kv_heads * head_size]
289
+ key: shape = [num_tokens, num_kv_heads, head_size]
290
+ value: shape = [num_tokens, num_kv_heads, head_size]
291
kv_cache = [num_blocks, block_size, num_kv_heads * 2, head_size]
292
num_slices_per_kv_cache_update_block: int
293
"""
294
_, page_size, num_combined_kv_heads, head_size = kv_cache.shape
295
head_size = cdiv(head_size,
296
TPU_HEAD_SIZE_ALIGNMENT) * TPU_HEAD_SIZE_ALIGNMENT
297
+
298
+ if kv_cache_quantized_dtype is not None:
299
+ dtype_info = torch.finfo(kv_cache_quantized_dtype)
300
+ key = key.to(torch.float32) / k_scale
301
+ # NOTE: clamp is added here to avoid out of range of quantized dtype
302
+ key = torch.clamp(key, dtype_info.min, dtype_info.max)
303
+ key = key.to(kv_cache_quantized_dtype)
304
+ value = value.to(torch.float32) / v_scale
305
+ value = torch.clamp(value, dtype_info.min, dtype_info.max)
306
+ value = value.to(kv_cache_quantized_dtype)
307
+
308
kv = torch.cat([key, value], axis=-1).reshape(-1, num_combined_kv_heads,
309
head_size)
310
vllm/v1/worker/tpu_model_runner.py
+6
-5
@@ -32,9 +32,10 @@ from vllm.multimodal.inputs import (BatchedTensorInputs, MultiModalKwargs,
32
from vllm.multimodal.utils import group_mm_inputs_by_modality
33
from vllm.pooling_params import PoolingTask
34
from vllm.sequence import IntermediateTensors
35
-from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, LayerBlockType, cdiv,
36
- is_pin_memory_available, prev_power_of_2)
37
-from vllm.v1.attention.backends.pallas import (PallasAttentionBackend,
35
+from vllm.utils import (LayerBlockType, cdiv, is_pin_memory_available,
36
+ prev_power_of_2)
37
+from vllm.v1.attention.backends.pallas import (TPU_STR_DTYPE_TO_TORCH_DTYPE,
38
+ PallasAttentionBackend,
39
PallasMetadata,
40
get_page_size_bytes)
41
from vllm.v1.core.encoder_cache_manager import compute_encoder_budget
@@ -142,11 +143,11 @@ class TPUModelRunner(LoRAModelRunnerMixin):
143
if cache_config.cache_dtype == "auto":
144
model_dtype = self.dtype
145
if isinstance(model_dtype, str):
145
- self.kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[model_dtype]
146
+ self.kv_cache_dtype = TPU_STR_DTYPE_TO_TORCH_DTYPE[model_dtype]
147
else:
148
self.kv_cache_dtype = model_dtype
149
else:
149
- self.kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[
150
+ self.kv_cache_dtype = TPU_STR_DTYPE_TO_TORCH_DTYPE[
151
cache_config.cache_dtype]
152
self._hidden_states_dtype = self.dtype
153