@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 70749fdcca

[Feature] Triton INT4 per-token-head KV cache quantization (#40835)

Signed-off-by: JartX <sagformas@epdcenter.es> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

JartX committed Jun 24, 2026 at 12:21 UTC 70749fdcca7da31be7206101e6f6fc77de9e6839

10 files changed +1463 -112

docs/design/attention_backends.md

+1 -1

@@ -170,7 +170,7 @@ Priority is **1 = highest** (tried first).

170 | `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ✅ | ✅ | ❌ | ❌ | Decoder | N/A |

171 | `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |

172 | `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |

173 -| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ✅ | ❌ | All | Any |

173 +| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int4_per_token_head`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ✅ | ❌ | All | Any |

174 | `TRITON_ATTN_DIFFKV` | | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |

175 | `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |

176

tests/models/quantization/test_per_token_kv_cache.py

+2 -1

@@ -31,7 +31,8 @@ from ..utils import check_logprobs_close

31 ],

32 )

33 @pytest.mark.parametrize(

34 - "kv_cache_dtype", ["int8_per_token_head", "fp8_per_token_head"]

34 + "kv_cache_dtype",

35 + ["int4_per_token_head", "int8_per_token_head", "fp8_per_token_head"],

36 )

37 @pytest.mark.parametrize("max_tokens", [4])

38 @pytest.mark.parametrize("enforce_eager", [True])

tests/quantization/test_per_token_kv_cache.py

+185 -70

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -"""Tests for per-token-head KV cache quantization (INT8 and FP8).

3 +"""Tests for per-token-head KV cache quantization (INT4, INT8 and FP8).

5 Covers:

6 - Per-token-head Triton reshape-and-cache kernel

@@ -23,6 +23,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

23 )

24 from vllm.platforms import current_platform

25 from vllm.utils.torch_utils import set_random_seed

26 +from vllm.v1.attention.ops.int4_per_token_head import single_rht

27 from vllm.v1.kv_cache_interface import KVQuantMode, is_quantized_kv_cache

29 DEVICE_TYPE = current_platform.device_type

@@ -81,11 +82,19 @@ FP8_CONFIG = QuantConfig(

82 kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD,

83 rounds_before_store=False,

84 )

84 -

85 -QUANT_CONFIGS = [INT8_CONFIG, FP8_CONFIG]

85 +INT4_CONFIG = QuantConfig(

86 + cache_dtype=torch.uint8,

87 + kv_cache_dtype_str="int4_per_token_head",

88 + quant_max=7.0,

89 + quant_min=-8.0,

90 + kv_quant_mode=KVQuantMode.INT4_PER_TOKEN_HEAD,

91 + # Unused for int4 (handled by its own rint path); kept for the dataclass.

92 + rounds_before_store=False,

93 +)

94 +QUANT_CONFIGS = [INT4_CONFIG, INT8_CONFIG, FP8_CONFIG]

88 -@pytest.fixture(params=QUANT_CONFIGS, ids=["int8", "fp8"])

97 +@pytest.fixture(params=QUANT_CONFIGS, ids=["int4", "int8", "fp8"])

98 def qcfg(request) -> QuantConfig:

99 return request.param

100

@@ -120,6 +129,9 @@ class TestIsQuantizedKvCache:

129 assert is_quantized_kv_cache("fp8_e4m3")

130 assert is_quantized_kv_cache("fp8_e5m2")

131

132 + def test_int4_per_token_head(self):

133 + assert is_quantized_kv_cache("int4_per_token_head")

134 +

135 def test_int8_per_token_head(self):

136 assert is_quantized_kv_cache("int8_per_token_head")

137

@@ -132,6 +144,13 @@ class TestIsQuantizedKvCache:

144 def test_bfloat16(self):

145 assert not is_quantized_kv_cache("bfloat16")

146

147 + def test_kv_quant_mode_int4(self):

148 + from vllm.v1.kv_cache_interface import get_kv_quant_mode

149 +

150 + assert (

151 + get_kv_quant_mode("int4_per_token_head") == KVQuantMode.INT4_PER_TOKEN_HEAD

152 + )

153 +

154 def test_kv_quant_mode_int8(self):

155 from vllm.v1.kv_cache_interface import get_kv_quant_mode

156

@@ -171,15 +190,17 @@ def test_reshape_and_cache_per_token_head(

190 torch.set_default_device(DEVICE_TYPE)

191

192 num_blocks = (num_tokens + block_size - 1) // block_size + 4

193 + is_int4 = qcfg.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD

194 + cache_head_size = head_size // 2 if is_int4 else head_size

195

196 key = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16)

197 value = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16)

198

199 key_cache = torch.zeros(

179 - num_blocks, block_size, num_heads, head_size, dtype=qcfg.cache_dtype

200 + num_blocks, block_size, num_heads, cache_head_size, dtype=qcfg.cache_dtype

201 )

202 value_cache = torch.zeros(

182 - num_blocks, block_size, num_heads, head_size, dtype=qcfg.cache_dtype

203 + num_blocks, block_size, num_heads, cache_head_size, dtype=qcfg.cache_dtype

204 )

205 k_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32)

206 v_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32)

@@ -197,45 +218,68 @@ def test_reshape_and_cache_per_token_head(

218 k_scale_cache,

219 v_scale_cache,

220 slot_mapping,

221 + kv_quant_mode=qcfg.kv_quant_mode,

222 )

223

202 - # Reference

203 - ref_k_quant, ref_k_scales = _quantize_per_token_head_ref(key, qcfg)

204 - ref_v_quant, ref_v_scales = _quantize_per_token_head_ref(value, qcfg)

224 + # INT4 (RHT + asymmetric), INT8/FP8 have different dequant paths. Only

225 + # INT8/FP8 can be compared to a PyTorch reference.

226 + if not is_int4:

227 + ref_k_quant, ref_k_scales = _quantize_per_token_head_ref(key, qcfg)

228 + ref_v_quant, ref_v_scales = _quantize_per_token_head_ref(value, qcfg)

229

206 - # Compare dequantized values rather than raw quantized values.

207 - # Triton and PyTorch reductions can differ at FP8 rounding boundaries

208 - # (up to 32 in quantized domain for fp8_e4m3), but the dequantized

209 - # error is bounded by the scale.

230 for i, slot in enumerate(slot_mapping.tolist()):

231 blk = slot // block_size

232 off = slot % block_size

233

214 - actual_k_scale = k_scale_cache[blk, off] # [num_heads]

215 - k_deq = key_cache[blk, off].float() * actual_k_scale[:, None]

216 - k_ref_deq = key[i].float()

217 - torch.testing.assert_close(

218 - k_deq,

219 - k_ref_deq,

220 - atol=0.1,

221 - rtol=0.1,

222 - )

223 - actual_v_scale = v_scale_cache[blk, off] # [num_heads]

224 - v_deq = value_cache[blk, off].float() * actual_v_scale[:, None]

225 - v_ref_deq = value[i].float()

226 - torch.testing.assert_close(

227 - v_deq,

228 - v_ref_deq,

229 - atol=0.1,

230 - rtol=0.1,

231 - )

232 - # Per-head scales: [num_heads]

233 - torch.testing.assert_close(

234 - k_scale_cache[blk, off], ref_k_scales[i], atol=1e-4, rtol=1e-3

235 - )

236 - torch.testing.assert_close(

237 - v_scale_cache[blk, off], ref_v_scales[i], atol=1e-4, rtol=1e-3

238 - )

234 + if is_int4:

235 + # Coarser quantization → wider tolerance.

236 + deq_atol = deq_rtol = 0.5

237 + for label, data, cache, sc in [

238 + ("key", key, key_cache, k_scale_cache),

239 + ("val", value, value_cache, v_scale_cache),

240 + ]:

241 + packed_scale = sc[blk, off] # [num_heads] float32

242 + scale_bits = packed_scale.view(torch.int32)

243 + zp = (scale_bits & 0xF).to(torch.float32)

244 + clean_scale = (scale_bits & -16).view(torch.float32)

245 +

246 + packed = cache[blk, off]

247 + lo = (packed & 0xF).to(torch.float32)

248 + hi = ((packed >> 4) & 0xF).to(torch.float32)

249 + full = torch.zeros(num_heads, head_size, dtype=torch.float32)

250 + full[:, 0::2] = lo

251 + full[:, 1::2] = hi

252 + # Asymmetric dequant in RHT domain, then IRHT/d → original

253 + deq_rht = (full - zp[:, None]) * clean_scale[:, None]

254 + deq = single_rht(deq_rht, inverse=True) / head_size

255 + ref_deq = data[i].float()

256 + torch.testing.assert_close(deq, ref_deq, atol=deq_atol, rtol=deq_rtol)

257 + else:

258 + actual_k_scale = k_scale_cache[blk, off] # [num_heads]

259 + k_deq = key_cache[blk, off].float() * actual_k_scale[:, None]

260 + k_ref_deq = key[i].float()

261 + torch.testing.assert_close(

262 + k_deq,

263 + k_ref_deq,

264 + atol=0.1,

265 + rtol=0.1,

266 + )

267 + actual_v_scale = v_scale_cache[blk, off] # [num_heads]

268 + v_deq = value_cache[blk, off].float() * actual_v_scale[:, None]

269 + v_ref_deq = value[i].float()

270 + torch.testing.assert_close(

271 + v_deq,

272 + v_ref_deq,

273 + atol=0.1,

274 + rtol=0.1,

275 + )

276 + # Per-head scales: [num_heads]

277 + torch.testing.assert_close(

278 + k_scale_cache[blk, off], ref_k_scales[i], atol=1e-4, rtol=1e-3

279 + )

280 + torch.testing.assert_close(

281 + v_scale_cache[blk, off], ref_v_scales[i], atol=1e-4, rtol=1e-3

282 + )

283

284

285 # ===========================================================================

@@ -265,16 +309,18 @@ def test_per_token_head_round_trip_accuracy(

309 torch.set_default_device(DEVICE_TYPE)

310 set_random_seed(42)

311

312 + is_int4 = qcfg.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD

313 num_blocks = (num_tokens + block_size - 1) // block_size + 2

314 + cache_head_size = head_size // 2 if is_int4 else head_size

315

316 key = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16) * 0.5

317 value = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16) * 0.5

318

319 key_cache = torch.zeros(

274 - num_blocks, block_size, num_heads, head_size, dtype=qcfg.cache_dtype

320 + num_blocks, block_size, num_heads, cache_head_size, dtype=qcfg.cache_dtype

321 )

322 value_cache = torch.zeros(

277 - num_blocks, block_size, num_heads, head_size, dtype=qcfg.cache_dtype

323 + num_blocks, block_size, num_heads, cache_head_size, dtype=qcfg.cache_dtype

324 )

325 k_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32)

326 v_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32)

@@ -289,8 +335,11 @@ def test_per_token_head_round_trip_accuracy(

335 k_scale_cache,

336 v_scale_cache,

337 slot_mapping,

338 + kv_quant_mode=qcfg.kv_quant_mode,

339 )

340

341 + rt_atol = 0.5 if is_int4 else 0.1

342 +

343 for i in range(num_tokens):

344 blk = i // block_size

345 off = i % block_size

@@ -300,18 +349,30 @@ def test_per_token_head_round_trip_accuracy(

349 ("val", value, value_cache, v_scale_cache),

350 ]:

351 for h in range(num_heads):

303 - orig = data[i, h].float() # [head_size]

304 -

305 - actual_q = cache[blk, off, h]

352 + orig = data[i, h].float()

353 actual_sc = sc[blk, off, h]

307 - actual_deq = actual_q.float() * actual_sc

308 -

309 - # Round-trip: dequantized should be close to original

354 + if is_int4:

355 + sc_bits = actual_sc.view(torch.int32)

356 + zp = (sc_bits & 0xF).to(torch.float32)

357 + clean_sc = (sc_bits & -16).view(torch.float32)

358 + packed = cache[blk, off, h]

359 + lo = (packed & 0xF).to(torch.float32)

360 + hi = ((packed >> 4) & 0xF).to(torch.float32)

361 + full = torch.zeros(head_size)

362 + full[0::2] = lo

363 + full[1::2] = hi

364 + deq_rht = (full - zp) * clean_sc

365 + actual_deq = (

366 + single_rht(deq_rht.unsqueeze(0), inverse=True).squeeze(0)

367 + / head_size

368 + )

369 + else:

370 + actual_deq = cache[blk, off, h].float() * actual_sc

371 torch.testing.assert_close(

372 actual_deq,

373 orig,

313 - atol=0.1,

314 - rtol=0.1,

374 + atol=rt_atol,

375 + rtol=rt_atol,

376 )

377

378

@@ -347,6 +408,7 @@ def test_int8_per_token_head_raw_cache_matches_round_reference():

408 k_scale_cache,

409 v_scale_cache,

410 slot_mapping,

411 + kv_quant_mode=INT8_CONFIG.kv_quant_mode,

412 )

413

414 ref_k_quant, ref_k_scales = _quantize_per_token_head_ref(key, INT8_CONFIG)

@@ -377,15 +439,17 @@ def test_per_token_head_negative_slot_skipped(qcfg: QuantConfig):

439 head_size = 64

440 block_size = 16

441 num_blocks = 2

442 + is_int4 = qcfg.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD

443 + cache_head_size = head_size // 2 if is_int4 else head_size

444

445 key = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16)

446 value = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16)

447

448 key_cache = torch.zeros(

385 - num_blocks, block_size, num_heads, head_size, dtype=qcfg.cache_dtype

449 + num_blocks, block_size, num_heads, cache_head_size, dtype=qcfg.cache_dtype

450 )

451 value_cache = torch.zeros(

388 - num_blocks, block_size, num_heads, head_size, dtype=qcfg.cache_dtype

452 + num_blocks, block_size, num_heads, cache_head_size, dtype=qcfg.cache_dtype

453 )

454 k_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32)

455 v_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32)

@@ -403,6 +467,7 @@ def test_per_token_head_negative_slot_skipped(qcfg: QuantConfig):

467 k_scale_cache,

468 v_scale_cache,

469 slot_mapping,

470 + kv_quant_mode=qcfg.kv_quant_mode,

471 )

472

473 # Slots 0 and 1 should have been written (tokens 0 and 2)

@@ -420,7 +485,8 @@ def test_per_token_head_negative_slot_skipped(qcfg: QuantConfig):

485 # 5. process_weights_after_loading -- per-token-head early return

486 # ===========================================================================

487 @pytest.mark.parametrize(

423 - "kv_cache_dtype", ["int8_per_token_head", "fp8_per_token_head"]

488 + "kv_cache_dtype",

489 + ["int4_per_token_head", "int8_per_token_head", "fp8_per_token_head"],

490 )

491 def test_process_weights_sets_placeholder_scales(kv_cache_dtype: str):

492 """Per-token-head should set _k_scale=1.0, _v_scale=1.0

@@ -454,7 +520,7 @@ def test_process_weights_sets_placeholder_scales(kv_cache_dtype: str):

520

521

522 # ===========================================================================

457 -# 6. Triton unified_attention -- per-token-head scale cache (INT8 and FP8)

523 +# 6. Triton unified_attention -- per-token-head scale cache (INT4/INT8/FP8)

524 # ===========================================================================

525 @pytest.mark.parametrize(

526 "seq_lens",

@@ -481,6 +547,8 @@ def test_triton_unified_attention_per_token_head_scale(

547 torch.set_default_device(DEVICE_TYPE)

548 set_random_seed(0)

549

550 + is_int4 = qcfg.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD

551 +

552 num_seqs = len(seq_lens)

553 query_lens = [s[0] for s in seq_lens]

554 kv_lens = [s[1] for s in seq_lens]

@@ -499,22 +567,68 @@ def test_triton_unified_attention_per_token_head_scale(

567 )

568 value_cache_bf16 = torch.randn_like(key_cache_bf16)

569

502 - # Per-token-head quantization: one scale per (block, slot, head)

503 - k_absmax = key_cache_bf16.float().abs().amax(dim=-1) # [..., num_kv_heads]

504 - v_absmax = value_cache_bf16.float().abs().amax(dim=-1)

505 - k_scale_cache = (k_absmax / qcfg.quant_max).clamp(min=1e-6).to(torch.float32)

506 - v_scale_cache = (v_absmax / qcfg.quant_max).clamp(min=1e-6).to(torch.float32)

507 -

508 - scaled_k = key_cache_bf16.float() / k_scale_cache[:, :, :, None]

509 - scaled_v = value_cache_bf16.float() / v_scale_cache[:, :, :, None]

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)

570 + if is_int4:

571 + # Asymmetric quantization reference (matches the Triton kernel).

572 + kf = key_cache_bf16.float()

573 + vf = value_cache_bf16.float()

574 + k_min = kf.amin(dim=-1)

575 + k_max = kf.amax(dim=-1)

576 + v_min = vf.amin(dim=-1)

577 + v_max = vf.amax(dim=-1)

578 + k_scale_cache = ((k_max - k_min) / 15.0).clamp(min=1e-6).to(torch.float32)

579 + v_scale_cache = ((v_max - v_min) / 15.0).clamp(min=1e-6).to(torch.float32)

580 + k_zp = (-k_min / k_scale_cache).round().clamp(0, 15)

581 + v_zp = (-v_min / v_scale_cache).round().clamp(0, 15)

582 +

583 + key_cache_q_full = (

584 + (kf / k_scale_cache[..., None] + k_zp[..., None]).round().clamp(0, 15)

585 )

514 - value_cache_q = (

515 - scaled_v.round().clamp(qcfg.quant_min, qcfg.quant_max).to(qcfg.cache_dtype)

586 + value_cache_q_full = (

587 + (vf / v_scale_cache[..., None] + v_zp[..., None]).round().clamp(0, 15)

588 )

589 +

590 + # Dequantized reference: x_hat = (q - zp) * scale

591 + key_cache_deq = (key_cache_q_full - k_zp[..., None]) * k_scale_cache[..., None]

592 + value_cache_deq = (value_cache_q_full - v_zp[..., None]) * v_scale_cache[

593 + ..., None

594 + ]

595 +

596 + # Pack two uint4 values into one byte

597 + def _pack_int4(data_float):

598 + u = data_float.to(torch.uint8)

599 + lo = u[..., 0::2]

600 + hi = u[..., 1::2]

601 + return (lo & 0xF) | ((hi & 0xF) << 4)

602 +

603 + key_cache_q = _pack_int4(key_cache_q_full)

604 + value_cache_q = _pack_int4(value_cache_q_full)

605 +

606 + # Steganography: pack zp into low 4 bits of scale

607 + k_zp_int = k_zp.to(torch.int32)

608 + k_bits = k_scale_cache.view(torch.int32)

609 + k_scale_cache = ((k_bits & -16) | (k_zp_int & 0xF)).view(torch.float32)

610 + v_zp_int = v_zp.to(torch.int32)

611 + v_bits = v_scale_cache.view(torch.int32)

612 + v_scale_cache = ((v_bits & -16) | (v_zp_int & 0xF)).view(torch.float32)

613 else:

614 + # Symmetric quantization for int8/fp8.

615 + k_absmax = key_cache_bf16.float().abs().amax(dim=-1)

616 + v_absmax = value_cache_bf16.float().abs().amax(dim=-1)

617 + k_scale_cache = (k_absmax / qcfg.quant_max).clamp(min=1e-6).to(torch.float32)

618 + v_scale_cache = (v_absmax / qcfg.quant_max).clamp(min=1e-6).to(torch.float32)

619 + scaled_k = key_cache_bf16.float() / k_scale_cache[:, :, :, None]

620 + scaled_v = value_cache_bf16.float() / v_scale_cache[:, :, :, None]

621 +

622 + key_cache_q_full = scaled_k.round().clamp(qcfg.quant_min, qcfg.quant_max)

623 + value_cache_q_full = scaled_v.round().clamp(qcfg.quant_min, qcfg.quant_max)

624 +

625 + key_cache_deq = key_cache_q_full * k_scale_cache[:, :, :, None]

626 + value_cache_deq = value_cache_q_full * v_scale_cache[:, :, :, None]

627 +

628 + if not is_int4 and qcfg.rounds_before_store:

629 + key_cache_q = key_cache_q_full.to(qcfg.cache_dtype)

630 + value_cache_q = value_cache_q_full.to(qcfg.cache_dtype)

631 + elif not is_int4:

632 key_cache_q = scaled_k.clamp(qcfg.quant_min, qcfg.quant_max).to(

633 qcfg.cache_dtype

634 )

@@ -522,10 +636,6 @@ def test_triton_unified_attention_per_token_head_scale(

636 qcfg.cache_dtype

637 )

638

525 - # Dequantized reference

526 - key_cache_deq = key_cache_q.float() * k_scale_cache[:, :, :, None]

527 - value_cache_deq = value_cache_q.float() * v_scale_cache[:, :, :, None]

528 -

639 cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(

640 dim=0, dtype=torch.int32

641 )

@@ -605,4 +715,9 @@ def test_triton_unified_attention_per_token_head_scale(

715 softmax_segm_expsum=softmax_segm_expsum,

716 )

717

608 - torch.testing.assert_close(output_q, output_ref, atol=5e-2, rtol=5e-2)

718 + # Coarser quantization → wider tolerance.

719 + if is_int4:

720 + atol, rtol = 0.5, 0.5

721 + else:

722 + atol, rtol = 5e-2, 5e-2

723 + torch.testing.assert_close(output_q, output_ref, atol=atol, rtol=rtol)

vllm/config/cache.py

@@ -29,6 +29,7 @@ CacheDType = Literal[

29 "turboquant_4bit_nc",

30 "turboquant_k3v4_nc",

31 "turboquant_3bit_nc",

32 + "int4_per_token_head",

33 "int8_per_token_head",

34 "fp8_per_token_head",

35 "nvfp4",

vllm/utils/torch_utils.py

@@ -39,6 +39,7 @@ STR_DTYPE_TO_TORCH_DTYPE = {

39 "fp8_e4m3": torch.uint8,

40 "fp8_e5m2": torch.uint8,

41 "int8": torch.int8,

42 + "int4_per_token_head": torch.uint8,

43 "int8_per_token_head": torch.int8,

44 "fp8_per_token_head": torch.uint8,

45 "fp8_inc": torch.float8_e4m3fn,

vllm/v1/attention/backends/triton_attn.py

+33 -16

@@ -258,6 +258,7 @@ class TritonAttentionBackend(AttentionBackend):

258 "fp8",

259 "fp8_e4m3",

260 "fp8_e5m2",

261 + "int4_per_token_head",

262 "int8_per_token_head",

263 "fp8_per_token_head",

264 ]

@@ -301,9 +302,11 @@ class TritonAttentionBackend(AttentionBackend):

302 if block_size % 16 != 0:

303 raise ValueError("Block size must be a multiple of 16.")

304 if kv_cache_uses_per_token_head_scales(cache_dtype_str):

304 - # Pad head_size by sizeof(float32)/sizeof(cache_dtype) so

305 - # the per-head scale fits inline. The backend extracts

306 - # data[:head_size] and scale[head_size:] via typed views.

305 + # Pad the head dim by sizeof(float32)/sizeof(cache_dtype) so the

306 + # per-(token, head) scale fits inline after the quantized data;

307 + # the backend extracts data[:head_size] and scale[head_size:] via

308 + # typed views (see _ensure_scale_caches). INT4 packs two values

309 + # per byte, so the data occupies only head_size // 2 bytes.

310 from vllm.utils.torch_utils import (

311 STR_DTYPE_TO_TORCH_DTYPE,

312 get_dtype_size,

@@ -311,7 +314,11 @@ class TritonAttentionBackend(AttentionBackend):

314

315 cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_dtype_str]

316 scale_pad = get_dtype_size(torch.float32) // get_dtype_size(cache_dtype)

314 - return (num_blocks, 2, block_size, num_kv_heads, head_size + scale_pad)

317 + if get_kv_quant_mode(cache_dtype_str) == KVQuantMode.INT4_PER_TOKEN_HEAD:

318 + data_head_size = head_size // 2

319 + else:

320 + data_head_size = head_size

321 + return (num_blocks, 2, block_size, num_kv_heads, data_head_size + scale_pad)

322 return (num_blocks, 2, block_size, num_kv_heads, head_size)

323

324 @staticmethod

@@ -587,18 +594,14 @@ class TritonAttentionImpl(AttentionImpl):

594 layer,

595 )

596

590 - # Per-token-head quantized KV cache: use separate scale caches.

597 + # Per-token-head quantized KV cache: handled by the core unified

598 + # kernel, which dequantizes per-(token, head) inline via constexpr

599 + # branches (INT8 / FP8) and dispatches to the packed INT4 kernel.

600 if self._is_per_token_head_quant:

592 - self._ensure_scale_caches(kv_cache)

593 - key_cache, value_cache = kv_cache.unbind(1)

594 - if key_cache.dtype == torch.uint8:

595 - key_cache = key_cache.view(self.fp8_dtype)

596 - value_cache = value_cache.view(self.fp8_dtype)

597 - q_descale = None

598 - k_descale = None

599 - v_descale = None

601 + key_cache, value_cache = self._pth_key_value_caches(kv_cache)

602 k_scale_cache = self._k_scale_cache

603 v_scale_cache = self._v_scale_cache

604 + q_descale = k_descale = v_descale = None

605 # FP8 per-tensor / auto path (original flow).

606 else:

607 key_cache, value_cache = kv_cache.unbind(1)

@@ -675,6 +678,17 @@ class TritonAttentionImpl(AttentionImpl):

678

679 return output

680

681 + def _pth_key_value_caches(

682 + self, kv_cache: torch.Tensor

683 + ) -> tuple[torch.Tensor, torch.Tensor]:

684 + """Per-token-head K/V cache views (ensures scale caches; FP8 retyped)."""

685 + self._ensure_scale_caches(kv_cache)

686 + key_cache, value_cache = kv_cache.unbind(1)

687 + if self._kv_quant_mode == KVQuantMode.FP8_PER_TOKEN_HEAD:

688 + key_cache = key_cache.view(self.fp8_dtype)

689 + value_cache = value_cache.view(self.fp8_dtype)

690 + return key_cache, value_cache

691 +

692 def _forward_encoder_attention(

693 self,

694 query: torch.Tensor,

@@ -737,7 +751,9 @@ class TritonAttentionImpl(AttentionImpl):

751 if self._is_per_token_head_quant:

752 self._ensure_scale_caches(kv_cache)

753 key_cache, value_cache = kv_cache.unbind(1)

740 - if key_cache.dtype == torch.uint8:

754 + k_scale_cache = self._k_scale_cache

755 + v_scale_cache = self._v_scale_cache

756 + if self._kv_quant_mode == KVQuantMode.FP8_PER_TOKEN_HEAD:

757 key_cache = key_cache.view(self.fp8_dtype)

758 value_cache = value_cache.view(self.fp8_dtype)

759 triton_reshape_and_cache_flash_per_token_head_quant(

@@ -745,9 +761,10 @@ class TritonAttentionImpl(AttentionImpl):

761 value,

762 key_cache,

763 value_cache,

748 - self._k_scale_cache,

749 - self._v_scale_cache,

764 + k_scale_cache,

765 + v_scale_cache,

766 slot_mapping,

767 + kv_quant_mode=self._kv_quant_mode,

768 )

769 return

770 # For decoder and cross-attention, use KV cache as before.

vllm/v1/attention/ops/int4_per_token_head.py new

+1155

@@ -0,0 +1,1155 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Sub-byte packed (INT4) per-token-head KV cache mode.

4 +

5 +INT4 packs two 4-bit values per cache byte, pre-rotates with a single RHT,

6 +and hides a 4-bit zero-point in the scale's low mantissa bits — too

7 +different from the core kernel to share it. Owns the whole mode: nibble

8 +pack/unpack, the reshape (write) kernel, the split-dot attention (read)

9 +kernel, the RHT transform, and the public ``reshape_and_cache_int4`` /

10 +``unified_attention_int4`` entry points.

11 +"""

12 +

13 +from __future__ import annotations

14 +

15 +from typing import Any

16 +

17 +import torch

18 +

19 +from vllm.platforms import current_platform

20 +from vllm.triton_utils import tl, triton

21 +from vllm.v1.attention.ops.triton_attention_helpers import (

22 + apply_alibi_to_score,

23 + apply_softcap,

24 + cdiv_fn,

25 + compute_kv_seq_mask,

26 + compute_tile_loop_bounds,

27 + init_softmax_M,

28 + load_qq_bias_tile,

29 + resolve_seq_and_query_len,

30 + softmax_step,

31 + store_segm_reduce_scalars,

32 +)

33 +from vllm.v1.attention.ops.triton_unified_attention import reduce_segments

34 +

35 +float8_info = torch.finfo(current_platform.fp8_dtype())

36 +

37 +# 2 x int4 packed per storage byte.

38 +_INT4_PACKING_FACTOR = 2

39 +

40 +

41 +# ----------------------------------------------------------------------

42 +# Nibble pack / unpack (shared write+read format)

43 +# ----------------------------------------------------------------------

44 +

45 +

46 +@triton.jit

47 +def pack_int4_nibbles(lo, hi):

48 + """Pack two uint8 values (each in [0, 15]) into one byte."""

49 + return (lo & 0xF) | ((hi & 0xF) << 4)

50 +

51 +

52 +@triton.jit

53 +def unpack_int4_nibbles(packed):

54 + """Split one packed byte into the (low, high) nibble pair as uint8."""

55 + return packed & 0xF, (packed >> 4) & 0xF

56 +

57 +

58 +# ----------------------------------------------------------------------

59 +# Write path: RHT + pack + per-(token, head) scale

60 +# ----------------------------------------------------------------------

61 +

62 +

63 +@triton.jit

64 +def _reshape_cache_int4_kernel(

65 + key_ptr,

66 + value_ptr,

67 + key_cache_ptr,

68 + value_cache_ptr,

69 + k_scale_cache_ptr,

70 + v_scale_cache_ptr,

71 + slot_mapping_ptr,

72 + stride_key_tok: tl.int64,

73 + stride_key_head: tl.int64,

74 + stride_val_tok: tl.int64,

75 + stride_val_head: tl.int64,

76 + stride_kc_blk: tl.int64,

77 + stride_kc_slot: tl.int64,

78 + stride_kc_head: tl.int64,

79 + stride_vc_blk: tl.int64,

80 + stride_vc_slot: tl.int64,

81 + stride_vc_head: tl.int64,

82 + stride_ks_blk: tl.int64,

83 + stride_ks_slot: tl.int64,

84 + stride_ks_head: tl.int64,

85 + stride_vs_blk: tl.int64,

86 + stride_vs_slot: tl.int64,

87 + stride_vs_head: tl.int64,

88 + block_size: tl.constexpr,

89 + head_size: tl.constexpr,

90 + head_size_v: tl.constexpr,

91 + PACKED_HEAD_PADDED: tl.constexpr,

92 +):

93 + """INT4 asymmetric quantization with zero-point steganography."""

94 + tok = tl.program_id(0)

95 + head = tl.program_id(1)

96 +

97 + slot = tl.load(slot_mapping_ptr + tok).to(tl.int64)

98 + if slot < 0:

99 + return

100 +

101 + blk = slot // block_size

102 + slot_in_blk = slot % block_size

103 +

104 + half_offs = tl.arange(0, PACKED_HEAD_PADDED)

105 + even_offs = half_offs * 2

106 + odd_offs = half_offs * 2 + 1

107 +

108 + half_k = head_size // 2

109 + even_k_mask = even_offs < head_size

110 + odd_k_mask = odd_offs < head_size

111 + key_base = key_ptr + tok * stride_key_tok + head * stride_key_head

112 +

113 + k_even = tl.load(key_base + even_offs, mask=even_k_mask, other=0.0).to(tl.float32)

114 + k_odd = tl.load(key_base + odd_offs, mask=odd_k_mask, other=0.0).to(tl.float32)

115 +

116 + k_min = tl.minimum(

117 + tl.min(tl.where(even_k_mask, k_even, float("inf"))),

118 + tl.min(tl.where(odd_k_mask, k_odd, float("inf"))),

119 + )

120 + k_max = tl.maximum(

121 + tl.max(tl.where(even_k_mask, k_even, float("-inf"))),

122 + tl.max(tl.where(odd_k_mask, k_odd, float("-inf"))),

123 + )

124 + k_scale = tl.maximum((k_max - k_min) / 15.0, 1e-6)

125 + k_zp_f = tl.clamp(

126 + tl.where(

127 + -k_min / k_scale >= 0,

128 + (-k_min / k_scale + 0.5).to(tl.int32),

129 + (-k_min / k_scale - 0.5).to(tl.int32),

130 + ).to(tl.float32),

131 + 0.0,

132 + 15.0,

133 + )

134 +

135 + inv_k = 1.0 / k_scale

136 + k_even_s = k_even * inv_k + k_zp_f

137 + k_odd_s = k_odd * inv_k + k_zp_f

138 + k_even_q = tl.clamp(

139 + tl.where(

140 + k_even_s >= 0,

141 + (k_even_s + 0.5).to(tl.int32),

142 + (k_even_s - 0.5).to(tl.int32),

143 + ).to(tl.float32),

144 + 0.0,

145 + 15.0,

146 + )

147 + k_odd_q = tl.clamp(

148 + tl.where(

149 + k_odd_s >= 0,

150 + (k_odd_s + 0.5).to(tl.int32),

151 + (k_odd_s - 0.5).to(tl.int32),

152 + ).to(tl.float32),

153 + 0.0,

154 + 15.0,

155 + )

156 +

157 + k_zp_int = k_zp_f.to(tl.int32)

158 + k_scale_bits = k_scale.to(tl.int32, bitcast=True)

159 + k_scale_packed = ((k_scale_bits & -16) | (k_zp_int & 0xF)).to(

160 + tl.float32, bitcast=True

161 + )

162 +

163 + tl.store(

164 + k_scale_cache_ptr

165 + + blk * stride_ks_blk

166 + + slot_in_blk * stride_ks_slot

167 + + head * stride_ks_head,

168 + k_scale_packed,

169 + )

170 +

171 + k_packed = pack_int4_nibbles(k_even_q.to(tl.uint8), k_odd_q.to(tl.uint8))

172 + tl.store(

173 + key_cache_ptr

174 + + blk * stride_kc_blk

175 + + slot_in_blk * stride_kc_slot

176 + + head * stride_kc_head

177 + + half_offs,

178 + k_packed,

179 + mask=half_offs < half_k,

180 + )

181 +

182 + half_v = head_size_v // 2

183 + even_v_mask = even_offs < head_size_v

184 + odd_v_mask = odd_offs < head_size_v

185 + val_base = value_ptr + tok * stride_val_tok + head * stride_val_head

186 +

187 + v_even = tl.load(val_base + even_offs, mask=even_v_mask, other=0.0).to(tl.float32)

188 + v_odd = tl.load(val_base + odd_offs, mask=odd_v_mask, other=0.0).to(tl.float32)

189 +

190 + v_min = tl.minimum(

191 + tl.min(tl.where(even_v_mask, v_even, float("inf"))),

192 + tl.min(tl.where(odd_v_mask, v_odd, float("inf"))),

193 + )

194 + v_max = tl.maximum(

195 + tl.max(tl.where(even_v_mask, v_even, float("-inf"))),

196 + tl.max(tl.where(odd_v_mask, v_odd, float("-inf"))),

197 + )

198 + v_scale = tl.maximum((v_max - v_min) / 15.0, 1e-6)

199 + v_zp_f = tl.clamp(

200 + tl.where(

201 + -v_min / v_scale >= 0,

202 + (-v_min / v_scale + 0.5).to(tl.int32),

203 + (-v_min / v_scale - 0.5).to(tl.int32),

204 + ).to(tl.float32),

205 + 0.0,

206 + 15.0,

207 + )

208 +

209 + inv_v = 1.0 / v_scale

210 + v_even_s = v_even * inv_v + v_zp_f

211 + v_odd_s = v_odd * inv_v + v_zp_f

212 + v_even_q = tl.clamp(

213 + tl.where(

214 + v_even_s >= 0,

215 + (v_even_s + 0.5).to(tl.int32),

216 + (v_even_s - 0.5).to(tl.int32),

217 + ).to(tl.float32),

218 + 0.0,

219 + 15.0,

220 + )

221 + v_odd_q = tl.clamp(

222 + tl.where(

223 + v_odd_s >= 0,

224 + (v_odd_s + 0.5).to(tl.int32),

225 + (v_odd_s - 0.5).to(tl.int32),

226 + ).to(tl.float32),

227 + 0.0,

228 + 15.0,

229 + )

230 +

231 + v_zp_int = v_zp_f.to(tl.int32)

232 + v_scale_bits = v_scale.to(tl.int32, bitcast=True)

233 + v_scale_packed = ((v_scale_bits & -16) | (v_zp_int & 0xF)).to(

234 + tl.float32, bitcast=True

235 + )

236 +

237 + tl.store(

238 + v_scale_cache_ptr

239 + + blk * stride_vs_blk

240 + + slot_in_blk * stride_vs_slot

241 + + head * stride_vs_head,

242 + v_scale_packed,

243 + )

244 +

245 + v_packed = pack_int4_nibbles(v_even_q.to(tl.uint8), v_odd_q.to(tl.uint8))

246 + tl.store(

247 + value_cache_ptr

248 + + blk * stride_vc_blk

249 + + slot_in_blk * stride_vc_slot

250 + + head * stride_vc_head

251 + + half_offs,

252 + v_packed,

253 + mask=half_offs < half_v,

254 + )

255 +

256 +

257 +def _run_reshape_kernel(

258 + kernel,

259 + *,

260 + key: torch.Tensor,

261 + value: torch.Tensor,

262 + key_cache: torch.Tensor,

263 + value_cache: torch.Tensor,

264 + k_scale_cache: torch.Tensor,

265 + v_scale_cache: torch.Tensor,

266 + slot_mapping: torch.Tensor,

267 + packing_factor: int,

268 +) -> None:

269 + """Launch the packed INT4 reshape kernel."""

270 + num_tokens, num_kv_heads, head_size = key.shape

271 + head_size_v = value.shape[2]

272 + assert head_size % packing_factor == 0 and head_size_v % packing_factor == 0

273 + packed_padded = triton.next_power_of_2(

274 + max(head_size, head_size_v) // packing_factor

275 + )

276 + if current_platform.is_rocm() or current_platform.is_xpu():

277 + num_warps = 4

278 + else:

279 + num_warps = min(16, max(1, packed_padded // 32))

280 +

281 + kernel[(num_tokens, num_kv_heads)](

282 + key_ptr=key,

283 + value_ptr=value,

284 + key_cache_ptr=key_cache,

285 + value_cache_ptr=value_cache,

286 + k_scale_cache_ptr=k_scale_cache,

287 + v_scale_cache_ptr=v_scale_cache,

288 + slot_mapping_ptr=slot_mapping,

289 + stride_key_tok=key.stride(0),

290 + stride_key_head=key.stride(1),

291 + stride_val_tok=value.stride(0),

292 + stride_val_head=value.stride(1),

293 + stride_kc_blk=key_cache.stride(0),

294 + stride_kc_slot=key_cache.stride(1),

295 + stride_kc_head=key_cache.stride(2),

296 + stride_vc_blk=value_cache.stride(0),

297 + stride_vc_slot=value_cache.stride(1),

298 + stride_vc_head=value_cache.stride(2),

299 + stride_ks_blk=k_scale_cache.stride(0),

300 + stride_ks_slot=k_scale_cache.stride(1),

301 + stride_ks_head=k_scale_cache.stride(2),

302 + stride_vs_blk=v_scale_cache.stride(0),

303 + stride_vs_slot=v_scale_cache.stride(1),

304 + stride_vs_head=v_scale_cache.stride(2),

305 + block_size=key_cache.shape[1],

306 + head_size=head_size,

307 + head_size_v=head_size_v,

308 + PACKED_HEAD_PADDED=packed_padded,

309 + num_warps=num_warps,

310 + )

311 +

312 +

313 +# ----------------------------------------------------------------------

314 +# Read path: split-dot attention over the packed cache

315 +# ----------------------------------------------------------------------

316 +

317 +

318 +@triton.jit

319 +def _attn_packed(

320 + # Output destinations. In 2D mode the final result is written into

321 + # ``output_ptr``; in 3D mode per-segment partials go into the three

322 + # ``segm_*`` tensors and ``output_ptr`` is unused.

323 + output_ptr,

324 + segm_output_ptr,

325 + segm_max_ptr,

326 + segm_expsum_ptr,

327 + query_ptr,

328 + key_cache_ptr,

329 + value_cache_ptr,

330 + sink_ptr,

331 + block_tables_ptr,

332 + seq_lens_ptr,

333 + alibi_slopes_ptr,

334 + qq_bias_ptr,

335 + scale,

336 + out_scale,

337 + softcap,

338 + k_scale_cache_ptr,

339 + v_scale_cache_ptr,

340 + num_query_heads: tl.constexpr,

341 + num_queries_per_kv: tl.constexpr,

342 + block_table_stride: tl.int64,

343 + query_stride_0: tl.int64,

344 + query_stride_1: tl.int64,

345 + output_stride_0: tl.int64,

346 + output_stride_1: tl.int64,

347 + qq_bias_stride_0: tl.int64,

348 + BLOCK_SIZE: tl.constexpr,

349 + TILE_SIZE: tl.constexpr,

350 + HEAD_SIZE: tl.constexpr,

351 + HEAD_SIZE_PADDED: tl.constexpr,

352 + PACKED_HEAD_PADDED: tl.constexpr, # HEAD_SIZE / PACKING_FACTOR, rounded up

353 + USE_ALIBI_SLOPES: tl.constexpr,

354 + USE_ALIBI_SQRT: tl.constexpr,

355 + USE_QQ_BIAS: tl.constexpr,

356 + USE_SOFTCAP: tl.constexpr,

357 + USE_SINKS: tl.constexpr,

358 + SLIDING_WINDOW: tl.constexpr,

359 + USE_MM_PREFIX: tl.constexpr,

360 + MAX_MM_RANGES: tl.constexpr,

361 + mm_prefix_range_ptr,

362 + stride_k_cache_0: tl.int64,

363 + stride_k_cache_1: tl.int64,

364 + stride_k_cache_2: tl.int64,

365 + stride_k_cache_3: tl.constexpr,

366 + stride_v_cache_0: tl.int64,

367 + stride_v_cache_1: tl.int64,

368 + stride_v_cache_2: tl.int64,

369 + stride_v_cache_3: tl.constexpr,

370 + stride_ks_blk: tl.int64,

371 + stride_ks_slot: tl.int64,

372 + stride_ks_head: tl.int64,

373 + stride_vs_blk: tl.int64,

374 + stride_vs_slot: tl.int64,

375 + stride_vs_head: tl.int64,

376 + query_start_len_ptr,

377 + BLOCK_Q: tl.constexpr,

378 + num_seqs: tl.int32,

379 + BLOCK_M: tl.constexpr,

380 + NUM_SEGMENTS_PER_SEQ: tl.constexpr,

381 + USE_FP8: tl.constexpr,

382 + IS_3D: tl.constexpr,

383 + # 2 → INT4 nibble pair (asymmetric + zp). The packed KV cache stores

384 + # one byte per ``packed_offs``, holding PACKING_FACTOR values.

385 + PACKING_FACTOR: tl.constexpr,

386 + FP8_MIN: tl.constexpr = float8_info.min,

387 + FP8_MAX: tl.constexpr = float8_info.max,

388 +):

389 + # Shared prologue: sequence lookup, q-block bounds, early returns.

390 + q_block_global_idx = tl.program_id(0)

391 + kv_head_idx = tl.program_id(1)

392 + segm_idx = tl.program_id(2) if IS_3D else 0

393 +

394 + (

395 + seq_idx,

396 + q_block_local_idx,

397 + cur_batch_in_all_start_index,

398 + cur_batch_query_len,

399 + seq_len,

400 + ) = resolve_seq_and_query_len(

401 + query_start_len_ptr, seq_lens_ptr, q_block_global_idx, num_seqs, BLOCK_Q

402 + )

403 +

404 + if q_block_local_idx * BLOCK_Q >= cur_batch_query_len:

405 + return

406 +

407 + if IS_3D:

408 + tiles_per_segment = cdiv_fn(seq_len, NUM_SEGMENTS_PER_SEQ * TILE_SIZE)

409 + if segm_idx * tiles_per_segment * TILE_SIZE >= seq_len:

410 + return

411 + else:

412 + tiles_per_segment = 0

413 +

414 + offs_m = tl.arange(0, BLOCK_M)

415 + offs_t = tl.arange(0, TILE_SIZE)

416 + query_pos = q_block_local_idx * BLOCK_Q + offs_m // num_queries_per_kv

417 +

418 + query_offset_0 = cur_batch_in_all_start_index + query_pos

419 + query_offset_1 = kv_head_idx * num_queries_per_kv + offs_m % num_queries_per_kv

420 +

421 + query_mask_0 = tl.where(query_pos < cur_batch_query_len, 1, 0).to(tl.int1)

422 + query_mask_1 = tl.where(query_offset_1 < num_query_heads, 1, 0).to(tl.int1)

423 +

424 + # Split-Q prologue: PACKING_FACTOR interleaved streams of Q.

425 + # INT4 uses 2 streams (even / odd). The packed KV cache stores one

426 + # byte per ``packed_offs``, which holds PACKING_FACTOR values.

427 + packed_offs = tl.arange(0, PACKED_HEAD_PADDED)

428 + offs_s0 = packed_offs * PACKING_FACTOR

429 + offs_s1 = packed_offs * PACKING_FACTOR + 1

430 + mask_s0 = tl.where(offs_s0 < HEAD_SIZE, 1, 0).to(tl.int1)

431 + mask_s1 = tl.where(offs_s1 < HEAD_SIZE, 1, 0).to(tl.int1)

432 + packed_dim_mask = tl.where(packed_offs < HEAD_SIZE // PACKING_FACTOR, 1, 0).to(

433 + tl.int1

434 + )

435 + q_base = (

436 + query_offset_0[:, None] * query_stride_0

437 + + query_offset_1[:, None] * query_stride_1

438 + )

439 + q_mask = query_mask_0[:, None] & query_mask_1[:, None]

440 + Q_s0 = tl.load(

441 + query_ptr + q_base + offs_s0[None, :],

442 + mask=mask_s0[None, :] & q_mask,

443 + other=0.0,

444 + ).to(tl.float32)

445 + Q_s1 = tl.load(

446 + query_ptr + q_base + offs_s1[None, :],

447 + mask=mask_s1[None, :] & q_mask,

448 + other=0.0,

449 + ).to(tl.float32)

450 +

451 + # INT4 asymmetric correction needs sum(Q) per row.

452 + Q_sum = tl.sum(Q_s0, axis=1) + tl.sum(Q_s1, axis=1)

453 +

454 + block_table_offset = seq_idx * block_table_stride

455 +

456 + # Online-softmax state + optional feature loads.

457 + M = init_softmax_M(

458 + sink_ptr, query_offset_1, query_mask_1, segm_idx, BLOCK_M, USE_SINKS, IS_3D

459 + )

460 + L = tl.full([BLOCK_M], 1.0, dtype=tl.float32)

461 + acc_s0 = tl.zeros([BLOCK_M, PACKED_HEAD_PADDED], dtype=tl.float32)

462 + acc_s1 = tl.zeros([BLOCK_M, PACKED_HEAD_PADDED], dtype=tl.float32)

463 +

464 + context_len = seq_len - cur_batch_query_len

465 +

466 + if USE_ALIBI_SLOPES:

467 + alibi_slope = tl.load(

468 + alibi_slopes_ptr + query_offset_1, mask=query_mask_1, other=0.0

469 + )

470 +

471 + if USE_QQ_BIAS:

472 + qq_bias_row_ptrs = qq_bias_ptr + query_pos[:, None] * qq_bias_stride_0

473 +

474 + loop_lo, loop_hi, max_seq_prefix_len = compute_tile_loop_bounds(

475 + context_len,

476 + seq_len,

477 + cur_batch_query_len,

478 + q_block_local_idx,

479 + segm_idx,

480 + tiles_per_segment,

481 + TILE_SIZE,

482 + BLOCK_M,

483 + BLOCK_Q,

484 + num_queries_per_kv,

485 + SLIDING_WINDOW,

486 + USE_MM_PREFIX,

487 + IS_3D,

488 + )

489 +

490 + # Tile loop. Per-tile: load packed KV + scales, dequantize into

491 + # PACKING_FACTOR streams, compute the split dot, run the shared

492 + # softmax step, accumulate per stream.

493 + for j in range(loop_lo, loop_hi):

494 + seq_offset = j * TILE_SIZE + offs_t

495 + tile_mask = seq_offset < max_seq_prefix_len

496 +

497 + physical_block_idx = tl.load(

498 + block_tables_ptr + block_table_offset + seq_offset // BLOCK_SIZE

499 + ).to(tl.int64)

500 +

501 + slot_in_blk = seq_offset % BLOCK_SIZE

502 + k_off = (

503 + physical_block_idx[None, :] * stride_k_cache_0

504 + + kv_head_idx * stride_k_cache_2

505 + + packed_offs[:, None] * stride_k_cache_3

506 + + slot_in_blk[None, :] * stride_k_cache_1

507 + )

508 + K_packed = tl.load(

509 + key_cache_ptr + k_off,

510 + mask=packed_dim_mask[:, None] & tile_mask[None, :],

511 + other=0,

512 + )

513 + v_off = (

514 + physical_block_idx[:, None] * stride_v_cache_0

515 + + kv_head_idx * stride_v_cache_2

516 + + packed_offs[None, :] * stride_v_cache_3

517 + + slot_in_blk[:, None] * stride_v_cache_1

518 + )

519 + V_packed = tl.load(

520 + value_cache_ptr + v_off,

521 + mask=packed_dim_mask[None, :] & tile_mask[:, None],

522 + other=0,

523 + )

524 + # Dequantize KV. INT4 unpacks nibbles as plain uint [0..15];

525 + # the zero-point is applied on the score side.

526 + K_s0_u, K_s1_u = unpack_int4_nibbles(K_packed)

527 + K_s0 = K_s0_u.to(tl.float32)

528 + K_s1 = K_s1_u.to(tl.float32)

529 + V_s0_u, V_s1_u = unpack_int4_nibbles(V_packed)

530 + V_s0 = V_s0_u.to(tl.float32)

531 + V_s1 = V_s1_u.to(tl.float32)

532 +

533 + ks_idx = (

534 + physical_block_idx * stride_ks_blk

535 + + slot_in_blk * stride_ks_slot

536 + + kv_head_idx * stride_ks_head

537 + )

538 + ks_raw = tl.load(k_scale_cache_ptr + ks_idx, mask=tile_mask, other=0)

539 + vs_idx = (

540 + physical_block_idx * stride_vs_blk

541 + + slot_in_blk * stride_vs_slot

542 + + kv_head_idx * stride_vs_head

543 + )

544 + vs_raw = tl.load(v_scale_cache_ptr + vs_idx, mask=tile_mask, other=0)

545 +

546 + # INT4 steganographs the 4-bit zero-point in the low 4 bits of

547 + # the float32 scale's mantissa.

548 + ks_bits = ks_raw.to(tl.int32, bitcast=True)

549 + k_zp = (ks_bits & 0xF).to(tl.float32)

550 + k_token_head_scales = (ks_bits & -16).to(tl.float32, bitcast=True)

551 + vs_bits = vs_raw.to(tl.int32, bitcast=True)

552 + v_zp = (vs_bits & 0xF).to(tl.float32)

553 + v_token_head_scales = (vs_bits & -16).to(tl.float32, bitcast=True)

554 +

555 + query_abs_pos = context_len + query_pos[:, None]

556 + seq_mask = compute_kv_seq_mask(

557 + query_abs_pos,

558 + seq_offset,

559 + seq_idx,

560 + seq_len,

561 + mm_prefix_range_ptr,

562 + SLIDING_WINDOW,

563 + USE_MM_PREFIX,

564 + MAX_MM_RANGES,

565 + )

566 +

567 + # Score: split-dot across the 2 INT4 streams; fused

568 + # softmax_scale * per-(token, head) k_scale in one mul. INT4

569 + # subtracts the ``zp * sum(Q)`` correction term.

570 + S = tl.zeros(shape=(BLOCK_M, TILE_SIZE), dtype=tl.float32)

571 + raw_dot = tl.dot(Q_s0, K_s0) + tl.dot(Q_s1, K_s1)

572 + S += (raw_dot - Q_sum[:, None] * k_zp[None, :]) * (

573 + scale * k_token_head_scales[None, :]

574 + )

575 +

576 + if USE_SOFTCAP:

577 + S = apply_softcap(S, softcap)

578 +

579 + S = tl.where(

580 + query_mask_1[:, None] & query_mask_0[:, None] & seq_mask, S, float("-inf")

581 + )

582 +

583 + if USE_ALIBI_SLOPES:

584 + S = apply_alibi_to_score(

585 + S, alibi_slope, seq_offset, context_len, query_pos, USE_ALIBI_SQRT

586 + )

587 +

588 + if USE_QQ_BIAS:

589 + S += load_qq_bias_tile(

590 + qq_bias_row_ptrs, seq_offset, context_len, qq_bias_stride_0

591 + )

592 +

593 + M, L, P, alpha = softmax_step(S, M, L)

594 + acc_s0 = acc_s0 * alpha[:, None]

595 + acc_s1 = acc_s1 * alpha[:, None]

596 +

597 + if SLIDING_WINDOW:

598 + qpos_lo = q_block_local_idx * BLOCK_Q

599 + sw_mask = (context_len + qpos_lo - seq_offset) < SLIDING_WINDOW

600 + V_s0 = tl.where(sw_mask[:, None], V_s0, 0.0)

601 + V_s1 = tl.where(sw_mask[:, None], V_s1, 0.0)

602 +

603 + # Fuse v per-(token, head) scale into P. INT4 also subtracts

604 + # the v-zero-point contribution from each stream once.

605 + P_v = (P * v_token_head_scales[None, :]).to(tl.float32)

606 + Pv_zp_sum = tl.sum(P_v * v_zp[None, :], axis=1)

607 + acc_s0 += tl.dot(P_v, V_s0) - Pv_zp_sum[:, None]

608 + acc_s1 += tl.dot(P_v, V_s1) - Pv_zp_sum[:, None]

609 +

610 + # Epilogue. 2D writes the final output with optional FP8 clamp;

611 + # 3D writes the per-segment partials (output / max / expsum) for

612 + # ``reduce_segments`` to finalize. Each stream writes its own

613 + # stripe in the output layout.

614 + out_mask = query_mask_0[:, None] & query_mask_1[:, None]

615 + if IS_3D:

616 + segm_base = (

617 + query_offset_0[:, None].to(tl.int64)

618 + * (num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED)

619 + + query_offset_1[:, None] * (NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED)

620 + + segm_idx * HEAD_SIZE_PADDED

621 + )

622 + tl.store(

623 + segm_output_ptr + segm_base + offs_s0[None, :],

624 + acc_s0,

625 + mask=mask_s0[None, :] & out_mask,

626 + )

627 + tl.store(

628 + segm_output_ptr + segm_base + offs_s1[None, :],

629 + acc_s1,

630 + mask=mask_s1[None, :] & out_mask,

631 + )

632 + store_segm_reduce_scalars(

633 + segm_max_ptr,

634 + segm_expsum_ptr,

635 + query_offset_0,

636 + query_offset_1,

637 + segm_idx,

638 + M,

639 + L,

640 + query_mask_0,

641 + query_mask_1,

642 + num_query_heads,

643 + NUM_SEGMENTS_PER_SEQ,

644 + )

645 + else:

646 + acc_s0 = acc_s0 / L[:, None]

647 + acc_s1 = acc_s1 / L[:, None]

648 + if USE_FP8:

649 + out_s = tl.load(out_scale)

650 + acc_s0 = tl.clamp(acc_s0 * out_s, FP8_MIN, FP8_MAX)

651 + acc_s1 = tl.clamp(acc_s1 * out_s, FP8_MIN, FP8_MAX)

652 + out_base = (

653 + query_offset_0[:, None] * output_stride_0

654 + + query_offset_1[:, None] * output_stride_1

655 + )

656 + tl.store(

657 + output_ptr + out_base + offs_s0[None, :],

658 + acc_s0,

659 + mask=mask_s0[None, :] & out_mask,

660 + )

661 + tl.store(

662 + output_ptr + out_base + offs_s1[None, :],

663 + acc_s1,

664 + mask=mask_s1[None, :] & out_mask,

665 + )

666 +

667 +

668 +def _launch_packed_attn(

669 + *,

670 + q,

671 + k_cache,

672 + v_cache,

673 + out,

674 + cu_seqlens_q,

675 + max_seqlen_q,

676 + seqused_k,

677 + softmax_scale,

678 + window_size,

679 + block_table,

680 + softcap,

681 + sinks,

682 + alibi_slopes,

683 + use_alibi_sqrt,

684 + qq_bias,

685 + output_scale,

686 + mm_prefix_range,

687 + k_scale_cache,

688 + v_scale_cache,

689 + seq_threshold_3D,

690 + num_par_softmax_segments,

691 + softmax_segm_output,

692 + softmax_segm_max,

693 + softmax_segm_expsum,

694 + packing_factor: int,

695 +):

696 + """Launch ``_attn_packed`` for one of the sub-byte modes.

697 +

698 + Handles 2D-vs-3D dispatch, placeholder pointers for the unused side

699 + of that split, and the trailing ``reduce_segments`` pass. Writes

700 + into ``out`` (directly for 2D; via the segm buffers for 3D).

701 + """

702 + import vllm.envs as envs

703 + from vllm.v1.attention.ops.triton_unified_attention import _get_tile_size

704 +

705 + is_batch_invariant = envs.VLLM_BATCH_INVARIANT

706 +

707 + use_mm_prefix = False

708 + max_mm_ranges = 0

709 + if mm_prefix_range is not None:

710 + assert mm_prefix_range.ndim == 3, (

711 + f"Unsupported mm_prefix_range shape: {mm_prefix_range.shape}"

712 + )

713 + use_mm_prefix = True

714 + max_mm_ranges = mm_prefix_range.shape[1]

715 +

716 + block_size = v_cache.shape[1]

717 + num_seqs = len(seqused_k)

718 + num_query_heads = q.shape[1]

719 + num_kv_heads = k_cache.shape[2]

720 + num_queries_per_kv = num_query_heads // num_kv_heads

721 + head_size = q.shape[2]

722 +

723 + BLOCK_M = (

724 + 16 if num_queries_per_kv <= 16 else triton.next_power_of_2(num_queries_per_kv)

725 + )

726 + BLOCK_Q = BLOCK_M // num_queries_per_kv

727 + total_num_q_blocks = q.shape[0] // BLOCK_Q + num_seqs

728 + sliding_window_val = 1 + window_size[0] if window_size[0] >= 0 else 0

729 + TILE_SIZE_PREFILL = _get_tile_size(

730 + head_size, sliding_window_val, q.element_size(), is_prefill=True

731 + )

732 + TILE_SIZE_DECODE = _get_tile_size(

733 + head_size, sliding_window_val, q.element_size(), is_prefill=False

734 + )

735 +

736 + use_3d = not (

737 + seq_threshold_3D is None

738 + or num_par_softmax_segments is None

739 + or softmax_segm_output is None

740 + or softmax_segm_max is None

741 + or softmax_segm_expsum is None

742 + or max_seqlen_q > 1

743 + or num_seqs > seq_threshold_3D

744 + or is_batch_invariant

745 + )

746 +

747 + # 3D never reads ``output_ptr`` and 2D never reads the segm tensors,

748 + # but Triton needs a non-null pointer everywhere; reuse ``out`` as

749 + # the placeholder for the unused side.

750 + segm_output_ptr = softmax_segm_output if use_3d else out

751 + segm_max_ptr = softmax_segm_max if use_3d else out

752 + segm_expsum_ptr = softmax_segm_expsum if use_3d else out

753 + num_segments = num_par_softmax_segments if use_3d else 1

754 +

755 + grid: tuple[Any, ...]

756 + if use_3d:

757 + grid = (total_num_q_blocks, num_kv_heads, num_par_softmax_segments)

758 + tile_size = TILE_SIZE_DECODE

759 + else:

760 + grid = (total_num_q_blocks, num_kv_heads)

761 + tile_size = TILE_SIZE_PREFILL

762 +

763 + _attn_packed[grid](

764 + output_ptr=out,

765 + segm_output_ptr=segm_output_ptr,

766 + segm_max_ptr=segm_max_ptr,

767 + segm_expsum_ptr=segm_expsum_ptr,

768 + query_ptr=q,

769 + key_cache_ptr=k_cache,

770 + value_cache_ptr=v_cache,

771 + sink_ptr=sinks,

772 + block_tables_ptr=block_table,

773 + seq_lens_ptr=seqused_k,

774 + alibi_slopes_ptr=alibi_slopes,

775 + qq_bias_ptr=qq_bias,

776 + scale=softmax_scale,

777 + out_scale=1 / output_scale if output_scale is not None else 1.0,

778 + softcap=softcap,

779 + k_scale_cache_ptr=k_scale_cache,

780 + v_scale_cache_ptr=v_scale_cache,

781 + num_query_heads=num_query_heads,

782 + num_queries_per_kv=num_queries_per_kv,

783 + block_table_stride=block_table.stride(0),

784 + query_stride_0=q.stride(0),

785 + query_stride_1=q.stride(1),

786 + output_stride_0=out.stride(0),

787 + output_stride_1=out.stride(1),

788 + qq_bias_stride_0=qq_bias.stride(0) if qq_bias is not None else 0,

789 + BLOCK_SIZE=block_size,

790 + TILE_SIZE=tile_size,

791 + HEAD_SIZE=head_size,

792 + HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),

793 + PACKED_HEAD_PADDED=triton.next_power_of_2(head_size) // packing_factor,

794 + USE_ALIBI_SLOPES=alibi_slopes is not None,

795 + USE_ALIBI_SQRT=use_alibi_sqrt,

796 + USE_QQ_BIAS=qq_bias is not None,

797 + USE_SOFTCAP=(softcap > 0),

798 + USE_SINKS=(sinks is not None),

799 + SLIDING_WINDOW=(1 + window_size[0]),

800 + USE_MM_PREFIX=use_mm_prefix,

801 + MAX_MM_RANGES=max_mm_ranges,

802 + mm_prefix_range_ptr=mm_prefix_range,

803 + stride_k_cache_0=k_cache.stride(0),

804 + stride_k_cache_1=k_cache.stride(1),

805 + stride_k_cache_2=k_cache.stride(2),

806 + stride_k_cache_3=k_cache.stride(3),

807 + stride_v_cache_0=v_cache.stride(0),

808 + stride_v_cache_1=v_cache.stride(1),

809 + stride_v_cache_2=v_cache.stride(2),

810 + stride_v_cache_3=v_cache.stride(3),

811 + stride_ks_blk=k_scale_cache.stride(0),

812 + stride_ks_slot=k_scale_cache.stride(1),

813 + stride_ks_head=k_scale_cache.stride(2),

814 + stride_vs_blk=v_scale_cache.stride(0),

815 + stride_vs_slot=v_scale_cache.stride(1),

816 + stride_vs_head=v_scale_cache.stride(2),

817 + query_start_len_ptr=cu_seqlens_q,

818 + BLOCK_Q=BLOCK_Q,

819 + num_seqs=num_seqs,

820 + BLOCK_M=BLOCK_M,

821 + NUM_SEGMENTS_PER_SEQ=num_segments,

822 + USE_FP8=output_scale is not None,

823 + IS_3D=use_3d,

824 + PACKING_FACTOR=packing_factor,

825 + )

826 +

827 + if use_3d:

828 + reduce_segments[(q.shape[0], num_query_heads)](

829 + output_ptr=out,

830 + segm_output_ptr=softmax_segm_output,

831 + segm_max_ptr=softmax_segm_max,

832 + segm_expsum_ptr=softmax_segm_expsum,

833 + seq_lens_ptr=seqused_k,

834 + num_seqs=num_seqs,

835 + num_query_heads=num_query_heads,

836 + out_scale_inv=1 / output_scale if output_scale is not None else 1.0,

837 + output_stride_0=out.stride(0),

838 + output_stride_1=out.stride(1),

839 + block_table_stride=block_table.stride(0),

840 + TILE_SIZE=TILE_SIZE_DECODE,

841 + HEAD_SIZE=head_size,

842 + HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),

843 + query_start_len_ptr=cu_seqlens_q,

844 + BLOCK_Q=BLOCK_Q,

845 + NUM_SEGMENTS_PER_SEQ=num_par_softmax_segments,

846 + USE_FP8=output_scale is not None,

847 + )

848 +

849 +

850 +# ----------------------------------------------------------------------

851 +# Public entry points

852 +# ----------------------------------------------------------------------

853 +

854 +

855 +def reshape_and_cache_int4(

856 + key: torch.Tensor,

857 + value: torch.Tensor,

858 + key_cache: torch.Tensor,

859 + value_cache: torch.Tensor,

860 + slot_mapping: torch.Tensor,

861 + *,

862 + k_scale_cache: torch.Tensor,

863 + v_scale_cache: torch.Tensor,

864 +) -> None:

865 + """Pre-rotate (RHT), pack to INT4 and write into the paged cache."""

866 + key = single_rht(key.float()).to(key.dtype)

867 + value = single_rht(value.float()).to(value.dtype)

868 + _run_reshape_kernel(

869 + _reshape_cache_int4_kernel,

870 + key=key,

871 + value=value,

872 + key_cache=key_cache,

873 + value_cache=value_cache,

874 + k_scale_cache=k_scale_cache,

875 + v_scale_cache=v_scale_cache,

876 + slot_mapping=slot_mapping,

877 + packing_factor=_INT4_PACKING_FACTOR,

878 + )

879 +

880 +

881 +def unified_attention_int4(

882 + q: torch.Tensor,

883 + k_cache: torch.Tensor,

884 + v_cache: torch.Tensor,

885 + out: torch.Tensor,

886 + *,

887 + cu_seqlens_q: torch.Tensor,

888 + max_seqlen_q: int,

889 + seqused_k: torch.Tensor,

890 + max_seqlen_k: int,

891 + softmax_scale: float,

892 + window_size: tuple[int, int],

893 + block_table: torch.Tensor,

894 + softcap: float,

895 + sinks: torch.Tensor | None,

896 + alibi_slopes: torch.Tensor | None,

897 + use_alibi_sqrt: bool,

898 + qq_bias: torch.Tensor | None,

899 + output_scale: torch.Tensor | None,

900 + mm_prefix_range: torch.Tensor | None,

901 + k_scale_cache: torch.Tensor,

902 + v_scale_cache: torch.Tensor,

903 + seq_threshold_3D: int | None = None,

904 + num_par_softmax_segments: int | None = None,

905 + softmax_segm_output: torch.Tensor | None = None,

906 + softmax_segm_max: torch.Tensor | None = None,

907 + softmax_segm_expsum: torch.Tensor | None = None,

908 +) -> None:

909 + """Paged attention over the INT4 packed cache, writing into *out*.

910 +

911 + The forward RHT has norm ``sqrt(head_size)``, so ``softmax_scale`` is

912 + divided by ``head_size`` and the inverse RHT divides the output by

913 + ``head_size`` as well.

914 + """

915 + q_orig_dtype = q.dtype

916 + q = single_rht(q.float()).to(q_orig_dtype)

917 + head_size = q.shape[2]

918 + softmax_scale = softmax_scale / head_size

919 +

920 + _launch_packed_attn(

921 + q=q,

922 + k_cache=k_cache,

923 + v_cache=v_cache,

924 + out=out,

925 + cu_seqlens_q=cu_seqlens_q,

926 + max_seqlen_q=max_seqlen_q,

927 + seqused_k=seqused_k,

928 + softmax_scale=softmax_scale,

929 + window_size=window_size,

930 + block_table=block_table,

931 + softcap=softcap,

932 + sinks=sinks,

933 + alibi_slopes=alibi_slopes,

934 + use_alibi_sqrt=use_alibi_sqrt,

935 + qq_bias=qq_bias,

936 + output_scale=output_scale,

937 + mm_prefix_range=mm_prefix_range,

938 + k_scale_cache=k_scale_cache,

939 + v_scale_cache=v_scale_cache,

940 + seq_threshold_3D=seq_threshold_3D,

941 + num_par_softmax_segments=num_par_softmax_segments,

942 + softmax_segm_output=softmax_segm_output,

943 + softmax_segm_max=softmax_segm_max,

944 + softmax_segm_expsum=softmax_segm_expsum,

945 + packing_factor=_INT4_PACKING_FACTOR,

946 + )

947 +

948 + out_f = single_rht(out.float(), inverse=True) / head_size

949 + out.copy_(out_f.to(q_orig_dtype))

950 +

951 +

952 +# ----------------------------------------------------------------------

953 +# Randomized Hadamard Transform (RHT) — gaussianizes K/V before INT4

954 +# quantization; applied on write and to Q before the read kernels.

955 +# ----------------------------------------------------------------------

956 +

957 +# Hadacore (CUDA tensor core kernel) availability check

958 +# Hadacore's CUDA impl is only registered when built for sm_80+, but the

959 +# schema def is unconditional — on ROCm ``hasattr`` is True yet dispatch

960 +# would crash, so we also gate on ``is_cuda()`` and the sm_80 capability.

961 +_HADACORE_AVAILABLE: bool | None = None

962 +

963 +

964 +def _hadacore_available() -> bool:

965 + global _HADACORE_AVAILABLE

966 + if _HADACORE_AVAILABLE is None:

967 + _HADACORE_AVAILABLE = (

968 + current_platform.is_cuda()

969 + and current_platform.has_device_capability(80)

970 + and hasattr(torch.ops._C, "hadacore_transform")

971 + )

972 + return _HADACORE_AVAILABLE

973 +

974 +

975 +# Cached Hadamard matrices (one per (size, dtype, device) tuple)

976 +_HADAMARD_MATRIX_CACHE: dict[tuple[int, torch.dtype, str], torch.Tensor] = {}

977 +

978 +

979 +def _get_hadamard_matrix(

980 + d: int, dtype: torch.dtype, device: torch.device

981 +) -> torch.Tensor:

982 + key = (d, dtype, str(device))

983 + cached = _HADAMARD_MATRIX_CACHE.get(key)

984 + if cached is None:

985 + H = torch.ones(1, 1, dtype=torch.float32, device=device)

986 + while H.shape[0] < d:

987 + H = torch.cat(

988 + [

989 + torch.cat([H, H], dim=1),

990 + torch.cat([H, -H], dim=1),

991 + ],

992 + dim=0,

993 + )

994 + cached = H.to(dtype).contiguous()

995 + _HADAMARD_MATRIX_CACHE[key] = cached

996 + return cached

997 +

998 +

999 +# Triton MMA Hadamard kernel (Tier 2)

1000 +@triton.jit

1001 +def _hadamard_mma_kernel(

1002 + x_ptr,

1003 + h_ptr,

1004 + out_ptr,

1005 + n_rows,

1006 + stride_x_row: tl.int64,

1007 + stride_x_col: tl.int64,

1008 + stride_o_row: tl.int64,

1009 + stride_o_col: tl.int64,

1010 + BLOCK_M: tl.constexpr,

1011 + D: tl.constexpr,

1012 +):

1013 + pid = tl.program_id(0)

1014 + rows = pid * BLOCK_M + tl.arange(0, BLOCK_M)

1015 + cols = tl.arange(0, D)

1016 + row_mask = rows < n_rows

1017 +

1018 + x = tl.load(

1019 + x_ptr + rows[:, None] * stride_x_row + cols[None, :] * stride_x_col,

1020 + mask=row_mask[:, None],

1021 + other=0.0,

1022 + )

1023 + H = tl.load(h_ptr + cols[:, None] * D + cols[None, :])

1024 +

1025 + out = tl.dot(x, H, out_dtype=tl.float32).to(x.dtype)

1026 +

1027 + tl.store(

1028 + out_ptr + rows[:, None] * stride_o_row + cols[None, :] * stride_o_col,

1029 + out,

1030 + mask=row_mask[:, None],

1031 + )

1032 +

1033 +

1034 +# H is D×D bf16 = 2·D² bytes of LDS. AMD CDNA has 64 KiB LDS, so D ≤ 128

1035 +# (32 KiB) leaves room for input + accumulator. Larger D falls back.

1036 +_TRITON_HADAMARD_MIN_D = 16

1037 +_TRITON_HADAMARD_MAX_D = 128

1038 +

1039 +

1040 +def _triton_hadamard_transform(x: torch.Tensor) -> torch.Tensor:

1041 + d = x.shape[-1]

1042 + orig_shape = x.shape

1043 + orig_dtype = x.dtype

1044 +

1045 + work_dtype = torch.bfloat16 if orig_dtype == torch.float32 else orig_dtype

1046 + x2d = x.contiguous().to(work_dtype).reshape(-1, d)

1047 + out2d = torch.empty_like(x2d)

1048 + n_rows = x2d.shape[0]

1049 + H_mat = _get_hadamard_matrix(d, work_dtype, x.device)

1050 +

1051 + BLOCK_M = 16

1052 + grid = (triton.cdiv(n_rows, BLOCK_M),)

1053 + # num_stages=1: the kernel has no loop, so default 3-stage pipelining

1054 + # would triple-buffer H and blow the AMD LDS budget.

1055 + _hadamard_mma_kernel[grid](

1056 + x2d,

1057 + H_mat,

1058 + out2d,

1059 + n_rows,

1060 + x2d.stride(0),

1061 + x2d.stride(1),

1062 + out2d.stride(0),

1063 + out2d.stride(1),

1064 + BLOCK_M=BLOCK_M,

1065 + D=d,

1066 + num_stages=1,

1067 + num_warps=4,

1068 + )

1069 + return out2d.reshape(orig_shape).to(orig_dtype)

1070 +

1071 +

1072 +# Public API

1073 +def fast_hadamard_transform(x: torch.Tensor) -> torch.Tensor:

1074 + """Unnormalized Walsh-Hadamard Transform along the last dimension.

1075 +

1076 + H_d × x where H_d × H_d = d × I. Last dim must be a power of 2.

1077 +

1078 + Three-tier dispatch:

1079 + 1. Hadacore CUDA Tensor Core kernel (sm_80+).

1080 + 2. Triton MMA matmul kernel (CUDA fallback + ROCm MFMA/WMMA path).

1081 + 3. PyTorch butterfly (CPU and any GPU/dtype combo Triton can't take).

1082 + """

1083 + d = x.shape[-1]

1084 + assert d & (d - 1) == 0, f"Requires power-of-2 dim, got {d}"

1085 +

1086 + # Tier 1 — hadacore on CUDA.

1087 + if _hadacore_available() and 0 < d <= (1 << 15):

1088 + from vllm import _custom_ops as ops

1089 +

1090 + # hadacore returns x @ (H/√d); rescale to the unnormalized H × x

1091 + # convention the INT4 scale math is calibrated to.

1092 + rescale = d**0.5

1093 + if x.dtype in (torch.float16, torch.bfloat16):

1094 + y = ops.hadacore_transform(x.contiguous().clone(), inplace=True)

1095 + return y * rescale

1096 + # fp32 → bf16 round-trip; precision loss is irrelevant before

1097 + # INT4 quantization.

1098 + orig_dtype = x.dtype

1099 + x_bf16 = x.contiguous().to(torch.bfloat16)

1100 + y_bf16 = ops.hadacore_transform(x_bf16, inplace=True)

1101 + return y_bf16.to(orig_dtype) * rescale

1102 +

1103 + # Tier 2 — Triton MMA kernel (covers ROCm via MFMA/WMMA codegen, and

1104 + # also CUDA when hadacore is unavailable).

1105 + if (

1106 + x.is_cuda

1107 + and _TRITON_HADAMARD_MIN_D <= d <= _TRITON_HADAMARD_MAX_D

1108 + and x.dtype in (torch.float16, torch.bfloat16, torch.float32)

1109 + ):

1110 + return _triton_hadamard_transform(x)

1111 +

1112 + # Tier 3 — PyTorch butterfly (CPU / unsupported dtype / D < 16).

1113 + h = 1

1114 + while h < d:

1115 + xv = x.view(*x.shape[:-1], d // (2 * h), 2, h)

1116 + a = xv[..., 0, :]

1117 + b = xv[..., 1, :]

1118 + x = torch.stack([a + b, a - b], dim=-2).reshape(x.shape)

1119 + h <<= 1

1120 + return x

1121 +

1122 +

1123 +# Randomized Hadamard Transform (used by INT4)

1124 +# Deterministic ±1 signs for Randomized Hadamard Transform.

1125 +# RHT = H × D × x (sign flip + Hadamard). Breaks residual structure

1126 +# in KV vectors, improving quantization quality.

1127 +_RHT_SIGNS_CACHE: dict[tuple[int, int, str], torch.Tensor] = {}

1128 +

1129 +

1130 +def _get_rht_signs(d: int, round_idx: int, device: torch.device) -> torch.Tensor:

1131 + """Return a cached deterministic ±1 sign vector of length *d*."""

1132 + key = (d, round_idx, str(device))

1133 + if key not in _RHT_SIGNS_CACHE:

1134 + gen = torch.Generator(device="cpu")

1135 + gen.manual_seed(0x9E3779B9 + round_idx * 0x517CC1B7)

1136 + signs = (

1137 + 2.0 * torch.bernoulli(torch.full((d,), 0.5, device="cpu"), generator=gen)

1138 + - 1.0

1139 + )

1140 + _RHT_SIGNS_CACHE[key] = signs.to(device)

1141 + return _RHT_SIGNS_CACHE[key]

1142 +

1143 +

1144 +def single_rht(x: torch.Tensor, inverse: bool = False) -> torch.Tensor:

1145 + """Single Randomized Hadamard Transform: H × D₁ × x.

1146 +

1147 + Used by INT4 per-token-head quantization to gaussianize data

1148 + before asymmetric quantization.

1149 + """

1150 + d = x.shape[-1]

1151 + d1 = _get_rht_signs(d, 0, x.device)

1152 + if inverse:

1153 + return fast_hadamard_transform(x) * d1

1154 + else:

1155 + return fast_hadamard_transform(x * d1)

vllm/v1/attention/ops/triton_reshape_and_cache_flash.py

+21 -2

@@ -10,6 +10,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

10 from vllm.platforms import current_platform

11 from vllm.triton_utils import tl, triton

12 from vllm.utils.torch_utils import is_quantized_kv_cache

13 +from vllm.v1.kv_cache_interface import KVQuantMode

15 FP8_MIN, FP8_MAX = get_fp8_min_max()

@@ -276,6 +277,7 @@ def triton_reshape_and_cache_flash_per_token_head_quant(

277 k_scale_cache: torch.Tensor, # [num_blocks, block_size, num_kv_heads] float32

278 v_scale_cache: torch.Tensor, # [num_blocks, block_size, num_kv_heads] float32

279 slot_mapping: torch.Tensor, # [num_tokens]

280 + kv_quant_mode: KVQuantMode,

281 ):

282 """Quantize key/value per (token, head) and write to paged cache.

283

@@ -283,9 +285,26 @@ def triton_reshape_and_cache_flash_per_token_head_quant(

285 quantized data in key_cache/value_cache, and stores the float32

286 scale in k_scale_cache/v_scale_cache.

287

286 - The quantization range (QUANT_MAX, QUANT_MIN) is derived from the

287 - cache tensor dtype so the same code path works for int8 and fp8.

288 + INT4 needs sub-byte packing + a Hadamard rotation, so it is handled by

289 + its own kernel; INT8 / FP8 share this kernel, with the quantization

290 + range (QUANT_MAX, QUANT_MIN) derived from the cache tensor dtype.

291 """

292 + if kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD:

293 + from vllm.v1.attention.ops.int4_per_token_head import (

294 + reshape_and_cache_int4,

295 + )

296 +

297 + reshape_and_cache_int4(

298 + key,

299 + value,

300 + key_cache,

301 + value_cache,

302 + slot_mapping,

303 + k_scale_cache=k_scale_cache,

304 + v_scale_cache=v_scale_cache,

305 + )

306 + return

307 +

308 cache_dtype = key_cache.dtype

309 quant_params = _PER_TOKEN_HEAD_QUANT_PARAMS.get(cache_dtype)

310 if quant_params is None:

vllm/v1/attention/ops/triton_unified_attention.py

+43 -1

@@ -259,7 +259,8 @@ def kernel_unified_attention(

259 stride_vs_head: tl.int64 = None,

260 # KV cache quantization mode handled inside this kernel via constexpr

261 # branches: NONE (0), FP8_PER_TENSOR (1), INT8_PER_TOKEN_HEAD (2),

262 - # FP8_PER_TOKEN_HEAD (3).

262 + # FP8_PER_TOKEN_HEAD (3). Sub-byte INT4 (4) uses its own

263 + # int4_per_token_head kernel, not this one.

264 KV_QUANT_MODE: tl.constexpr = 0,

265 FP8_MIN: tl.constexpr = float8_info.min,

266 FP8_MAX: tl.constexpr = float8_info.max,

@@ -824,6 +825,47 @@ def unified_attention(

825 use_causal = bool(causal) if not use_per_seq_causal else True

826 per_seq_causal_ptr = causal if use_per_seq_causal else None

827

828 + # Sub-byte packed mode (INT4) needs a bespoke kernel (split-dot +

829 + # sub-byte unpack); everything else goes through the core kernel below.

830 + if kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD:

831 + assert use_causal and not use_per_seq_causal, (

832 + "INT4_PER_TOKEN_HEAD only supports causal attention"

833 + )

834 + from vllm.v1.attention.ops.int4_per_token_head import (

835 + unified_attention_int4,

836 + )

837 +

838 + if sinks is not None:

839 + assert sinks.shape[0] == q.shape[1], "Sinks must be num_query_heads size"

840 + unified_attention_int4(

841 + q=q,

842 + k_cache=k,

843 + v_cache=v,

844 + out=out,

845 + cu_seqlens_q=cu_seqlens_q,

846 + max_seqlen_q=max_seqlen_q,

847 + seqused_k=seqused_k,

848 + max_seqlen_k=max_seqlen_k,

849 + softmax_scale=softmax_scale,

850 + window_size=window_size,

851 + block_table=block_table,

852 + softcap=softcap,

853 + sinks=sinks,

854 + alibi_slopes=alibi_slopes,

855 + use_alibi_sqrt=use_alibi_sqrt,

856 + qq_bias=qq_bias,

857 + output_scale=output_scale,

858 + mm_prefix_range=mm_prefix_range,

859 + k_scale_cache=k_scale_cache,

860 + v_scale_cache=v_scale_cache,

861 + seq_threshold_3D=seq_threshold_3D,

862 + num_par_softmax_segments=num_par_softmax_segments,

863 + softmax_segm_output=softmax_segm_output,

864 + softmax_segm_max=softmax_segm_max,

865 + softmax_segm_expsum=softmax_segm_expsum,

866 + )

867 + return

868 +

869 if sinks is not None:

870 assert sinks.shape[0] == q.shape[1], "Sinks must be num_query_heads size"

871

vllm/v1/kv_cache_interface.py

+21 -21

@@ -41,7 +41,8 @@ class KVQuantMode(IntEnum):

41 FP8_PER_TENSOR = 1 # per-tensor scales (current fp8 path)

42 INT8_PER_TOKEN_HEAD = 2 # per-token-head dynamic scales for int8

43 FP8_PER_TOKEN_HEAD = 3 # per-token-head dynamic scales for fp8

44 - NVFP4 = 4 # packed fp4 data + fp8 block scales

44 + INT4_PER_TOKEN_HEAD = 4 # packed 2×int4/byte, RHT + asymmetric zp

45 + NVFP4 = 5 # packed fp4 data + fp8 block scales

47 @property

48 def is_per_token_head(self) -> bool:

@@ -49,6 +50,7 @@ class KVQuantMode(IntEnum):

50 return self in (

51 KVQuantMode.INT8_PER_TOKEN_HEAD,

52 KVQuantMode.FP8_PER_TOKEN_HEAD,

53 + KVQuantMode.INT4_PER_TOKEN_HEAD,

54 )

56 @property

@@ -59,6 +61,8 @@ class KVQuantMode(IntEnum):

62 def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode:

63 """Map a ``kv_cache_dtype`` string to a :class:`KVQuantMode`."""

64 + if kv_cache_dtype == "int4_per_token_head":

65 + return KVQuantMode.INT4_PER_TOKEN_HEAD

66 if kv_cache_dtype == "int8_per_token_head":

67 return KVQuantMode.INT8_PER_TOKEN_HEAD

68 if kv_cache_dtype == "fp8_per_token_head":

@@ -184,19 +188,16 @@ class AttentionSpec(KVCacheSpec):

188 def real_page_size_bytes(self) -> int:

189 if self.kv_quant_mode.is_nvfp4:

190 # Packed layout: fp4 data + fp8 block scales per head.

187 - full_dim = nvfp4_kv_cache_full_dim(self.head_size)

188 - return (

189 - 2

190 - * self.block_size

191 - * self.num_kv_heads

192 - * full_dim

193 - * get_dtype_size(self.dtype)

194 - )

191 + head_dim = nvfp4_kv_cache_full_dim(self.head_size)

192 + elif self.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD:

193 + head_dim = self.head_size // 2

194 + else:

195 + head_dim = self.head_size

196 return (

197 2

198 * self.block_size

199 * self.num_kv_heads

199 - * self.head_size

200 + * head_dim

201 * get_dtype_size(self.dtype)

202 )

203

@@ -314,17 +315,12 @@ class FullAttentionSpec(AttentionSpec):

315 last_dim = nvfp4_kv_cache_full_dim(

316 self.head_size

317 ) + nvfp4_kv_cache_full_dim(self.head_size_v)

317 - return (

318 - self.block_size

319 - * self.num_kv_heads

320 - * last_dim

321 - * get_dtype_size(self.dtype)

322 - )

318 + elif self.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD:

319 + last_dim = self.head_size // 2 + self.head_size_v // 2

320 + else:

321 + last_dim = self.head_size + self.head_size_v

322 return (

324 - self.block_size

325 - * self.num_kv_heads

326 - * (self.head_size + self.head_size_v)

327 - * get_dtype_size(self.dtype)

323 + self.block_size * self.num_kv_heads * last_dim * get_dtype_size(self.dtype)

324 )

325

326

@@ -390,10 +386,14 @@ class MLAAttentionSpec(FullAttentionSpec):

386 # V3.2 main MLA: 656-byte custom layout (kv_lora_rank=512 +

387 # qk_rope_head_dim=64, head_size=576). See flashmla_sparse.py.

388 return self.block_size * 656

389 + if self.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD:

390 + head_dim = self.head_size // 2

391 + else:

392 + head_dim = self.head_size

393 return (

394 self.storage_block_size

395 * self.num_kv_heads

396 - * self.head_size

396 + * head_dim

397 * get_dtype_size(self.dtype)

398 )

399