[Kernel][Bugfix] Fix INT8 per-token-head KV cache rounding in Triton reshape-and-cache (#45361)
Signed-off-by: ZedongLiu <113341356+Zedong-Liu@users.noreply.github.com>
ZedongLiu committed
Jun 22, 2026 at 03:59 UTC
9c450b102788bb271c1710354db22290ed57f543
2 files changed
+65
-9
tests/quantization/test_per_token_kv_cache.py
+53
-7
@@ -61,8 +61,8 @@ class QuantConfig:
61
quant_max: float
62
quant_min: float
63
kv_quant_mode: KVQuantMode
64
- # INT8 Triton stores truncate; FP8 hardware casts round.
65
- uses_trunc: bool
64
+ # INT8 rounds explicitly; FP8 relies on dtype cast rounding.
65
+ rounds_before_store: bool
66
67
68
INT8_CONFIG = QuantConfig(
@@ -71,7 +71,7 @@ INT8_CONFIG = QuantConfig(
71
quant_max=127.0,
72
quant_min=-128.0,
73
kv_quant_mode=KVQuantMode.INT8_PER_TOKEN_HEAD,
74
- uses_trunc=True,
74
+ rounds_before_store=True,
75
)
76
FP8_CONFIG = QuantConfig(
77
cache_dtype=FP8_DTYPE,
@@ -79,7 +79,7 @@ FP8_CONFIG = QuantConfig(
79
quant_max=FP8_MAX,
80
quant_min=FP8_MIN,
81
kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD,
82
- uses_trunc=False,
82
+ rounds_before_store=False,
83
)
84
85
QUANT_CONFIGS = [INT8_CONFIG, FP8_CONFIG]
@@ -104,7 +104,7 @@ def _quantize_per_token_head_ref(
104
absmax = data.float().abs().amax(dim=2) # [num_tokens, num_heads]
105
scales = (absmax / cfg.quant_max).clamp(min=1e-6)
106
scaled = data.float() * (1.0 / scales[:, :, None])
107
- if cfg.uses_trunc:
107
+ if cfg.rounds_before_store:
108
q = scaled.round().clamp(cfg.quant_min, cfg.quant_max).to(cfg.cache_dtype)
109
else:
110
q = scaled.clamp(cfg.quant_min, cfg.quant_max).to(cfg.cache_dtype)
@@ -255,7 +255,7 @@ def test_per_token_head_round_trip_accuracy(
255
):
256
"""Verify per-token-head round-trip: kernel dequant matches reference.
257
258
- INT8: Triton truncates on float->int8 store.
258
+ INT8: round-to-nearest before int8 store.
259
FP8: hardware cast (clamp then cast).
260
"""
261
from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
@@ -315,6 +315,52 @@ def test_per_token_head_round_trip_accuracy(
315
)
316
317
318
+@torch.inference_mode()
319
+def test_int8_per_token_head_raw_cache_matches_round_reference():
320
+ """INT8 cache writes should match round-to-nearest quantization exactly."""
321
+ from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
322
+ triton_reshape_and_cache_flash_per_token_head_quant,
323
+ )
324
+
325
+ torch.set_default_device(DEVICE_TYPE)
326
+
327
+ head_size = 8
328
+ block_size = 4
329
+
330
+ key = torch.tensor(
331
+ [[[-127.0, -2.6, -2.4, -1.6, -1.4, -0.6, -0.4, 127.0]]],
332
+ dtype=torch.bfloat16,
333
+ )
334
+ value = -key
335
+
336
+ key_cache = torch.zeros(1, block_size, 1, head_size, dtype=torch.int8)
337
+ value_cache = torch.zeros_like(key_cache)
338
+ k_scale_cache = torch.ones(1, block_size, 1, dtype=torch.float32)
339
+ v_scale_cache = torch.ones_like(k_scale_cache)
340
+ slot_mapping = torch.tensor([2], dtype=torch.long)
341
+
342
+ triton_reshape_and_cache_flash_per_token_head_quant(
343
+ key,
344
+ value,
345
+ key_cache,
346
+ value_cache,
347
+ k_scale_cache,
348
+ v_scale_cache,
349
+ slot_mapping,
350
+ )
351
+
352
+ ref_k_quant, ref_k_scales = _quantize_per_token_head_ref(key, INT8_CONFIG)
353
+ ref_v_quant, ref_v_scales = _quantize_per_token_head_ref(value, INT8_CONFIG)
354
+
355
+ slot = slot_mapping.item()
356
+ blk = slot // block_size
357
+ off = slot % block_size
358
+ assert torch.equal(key_cache[blk, off], ref_k_quant[0])
359
+ assert torch.equal(value_cache[blk, off], ref_v_quant[0])
360
+ torch.testing.assert_close(k_scale_cache[blk, off], ref_k_scales[0])
361
+ torch.testing.assert_close(v_scale_cache[blk, off], ref_v_scales[0])
362
+
363
+
364
# ===========================================================================
365
# 4. Negative slot mapping (padding tokens should be skipped)
366
# ===========================================================================
@@ -461,7 +507,7 @@ def test_triton_unified_attention_per_token_head_scale(
507
508
scaled_k = key_cache_bf16.float() / k_scale_cache[:, :, :, None]
509
scaled_v = value_cache_bf16.float() / v_scale_cache[:, :, :, None]
464
- if qcfg.uses_trunc:
510
+ if qcfg.rounds_before_store:
511
key_cache_q = (
512
scaled_k.round().clamp(qcfg.quant_min, qcfg.quant_max).to(qcfg.cache_dtype)
513
)
vllm/v1/attention/ops/triton_reshape_and_cache_flash.py
+12
-2
@@ -181,6 +181,7 @@ def _reshape_cache_per_token_head(
181
HEAD_SIZE_PADDED: tl.constexpr, # next_power_of_2(max(head_size, head_size_v))
182
QUANT_MAX: tl.constexpr = 127.0,
183
QUANT_MIN: tl.constexpr = -128.0,
184
+ IS_INT_QUANT: tl.constexpr = False,
185
):
186
tok = tl.program_id(0)
187
head = tl.program_id(1)
@@ -211,7 +212,11 @@ def _reshape_cache_per_token_head(
212
k_scale,
213
)
214
214
- k_q = tl.clamp(k_h * (1.0 / k_scale), QUANT_MIN, QUANT_MAX)
215
+ k_q = k_h * (1.0 / k_scale)
216
+ if IS_INT_QUANT:
217
+ # Round half away from zero before the int8 store truncates.
218
+ k_q = tl.where(k_q >= 0, k_q + 0.5, k_q - 0.5)
219
+ k_q = tl.clamp(k_q, QUANT_MIN, QUANT_MAX)
220
tl.store(
221
key_cache_ptr
222
+ blk * stride_kc_blk
@@ -239,7 +244,11 @@ def _reshape_cache_per_token_head(
244
v_scale,
245
)
246
242
- v_q = tl.clamp(v_h * (1.0 / v_scale), QUANT_MIN, QUANT_MAX)
247
+ v_q = v_h * (1.0 / v_scale)
248
+ if IS_INT_QUANT:
249
+ # Round half away from zero before the int8 store truncates.
250
+ v_q = tl.where(v_q >= 0, v_q + 0.5, v_q - 0.5)
251
+ v_q = tl.clamp(v_q, QUANT_MIN, QUANT_MAX)
252
tl.store(
253
value_cache_ptr
254
+ blk * stride_vc_blk
@@ -327,6 +336,7 @@ def triton_reshape_and_cache_flash_per_token_head_quant(
336
HEAD_SIZE_PADDED=head_size_padded,
337
QUANT_MAX=quant_max,
338
QUANT_MIN=quant_min,
339
+ IS_INT_QUANT=cache_dtype == torch.int8,
340
num_warps=num_warps,
341
)
342