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 5, 04:45 UTC
[Attention Backend] TurboQuant: 2-bit KV cache compression with 4x capacity (#38479)
Signed-off-by: vibhavagarwal5 <vibhavagarwal5@gmail.com> Signed-off-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Xinyu Chen <xinyu1.chen@intel.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
Vibhav Agarwal committed
Apr 15, 2026 at 08:27 UTC
f4b42df04847dcbd3247f7f4c56dff45e40bdf0d
27 files changed
+2963
-3
.buildkite/test_areas/lm_eval.yaml
+10
@@ -91,6 +91,16 @@ steps:
91
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
92
93
94
+- label: LM Eval TurboQuant KV Cache
95
+ timeout_in_minutes: 75
96
+ source_file_dependencies:
97
+ - vllm/model_executor/layers/quantization/turboquant/
98
+ - vllm/v1/attention/backends/turboquant_attn.py
99
+ - vllm/v1/attention/ops/triton_turboquant_decode.py
100
+ - vllm/v1/attention/ops/triton_turboquant_store.py
101
+ commands:
102
+ - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant.txt
103
+
104
- label: GPQA Eval (GPT-OSS) (H100)
105
timeout_in_minutes: 120
106
device: h100
docs/design/attention_backends.md
+1
@@ -178,6 +178,7 @@ Priority is **1 = highest** (tried first).
178
| `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 |
179
| `TREE_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any |
180
| `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 |
181
+| `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | Decoder | Any |
182
183
> **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`.
184
>
pyproject.toml
+3
@@ -170,6 +170,9 @@ eles = "eles"
170
datas = "datas"
171
ser = "ser"
172
ure = "ure"
173
+# Walsh-Hadamard Transform
174
+wht = "wht"
175
+WHT = "WHT"
176
177
[tool.uv]
178
no-build-isolation-package = ["torch"]
tests/evals/gsm8k/configs/Qwen3-4B-TQ-k3v4nc.yaml
new
+5
@@ -0,0 +1,5 @@
1
+model_name: "Qwen/Qwen3-4B"
2
+accuracy_threshold: 0.78
3
+num_questions: 1319
4
+num_fewshot: 5
5
+server_args: "--kv-cache-dtype turboquant_k3v4_nc --enforce-eager --max-model-len 4096"
tests/evals/gsm8k/configs/Qwen3-4B-TQ-k8v4.yaml
new
+5
@@ -0,0 +1,5 @@
1
+model_name: "Qwen/Qwen3-4B"
2
+accuracy_threshold: 0.80
3
+num_questions: 1319
4
+num_fewshot: 5
5
+server_args: "--kv-cache-dtype turboquant_k8v4 --enforce-eager --max-model-len 4096"
tests/evals/gsm8k/configs/Qwen3-4B-TQ-t3nc.yaml
new
+5
@@ -0,0 +1,5 @@
1
+model_name: "Qwen/Qwen3-4B"
2
+accuracy_threshold: 0.75
3
+num_questions: 1319
4
+num_fewshot: 5
5
+server_args: "--kv-cache-dtype turboquant_3bit_nc --enforce-eager --max-model-len 4096"
tests/evals/gsm8k/configs/Qwen3-4B-TQ-t4nc.yaml
new
+5
@@ -0,0 +1,5 @@
1
+model_name: "Qwen/Qwen3-4B"
2
+accuracy_threshold: 0.80
3
+num_questions: 1319
4
+num_fewshot: 5
5
+server_args: "--kv-cache-dtype turboquant_4bit_nc --enforce-eager --max-model-len 4096"
tests/evals/gsm8k/configs/models-turboquant.txt
new
+4
@@ -0,0 +1,4 @@
1
+Qwen3-4B-TQ-k8v4.yaml
2
+Qwen3-4B-TQ-t4nc.yaml
3
+Qwen3-4B-TQ-k3v4nc.yaml
4
+Qwen3-4B-TQ-t3nc.yaml
tests/quantization/test_turboquant.py
new
+570
@@ -0,0 +1,570 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for TurboQuant KV-cache quantization.
4
+
5
+Run: .venv/bin/python -m pytest tests/quantization/test_turboquant.py -v
6
+"""
7
+
8
+import math
9
+
10
+import pytest
11
+import torch
12
+
13
+from vllm.model_executor.layers.quantization.turboquant.centroids import (
14
+ get_centroids,
15
+ solve_lloyd_max,
16
+)
17
+from vllm.model_executor.layers.quantization.turboquant.config import (
18
+ TQ_PRESETS,
19
+ TurboQuantConfig,
20
+)
21
+from vllm.model_executor.layers.quantization.turboquant.quantizer import (
22
+ generate_wht_signs,
23
+)
24
+from vllm.utils.math_utils import next_power_of_2
25
+
26
+# ============================================================================
27
+# Helpers
28
+# ============================================================================
29
+
30
+ALL_PRESETS = list(TQ_PRESETS.keys())
31
+
32
+
33
+def _assert_strictly_sorted(seq, name="sequence"):
34
+ for i in range(len(seq) - 1):
35
+ assert seq[i] < seq[i + 1], f"{name} not sorted at index {i}"
36
+
37
+
38
+def _is_power_of_2(n: int) -> bool:
39
+ return n > 0 and next_power_of_2(n) == n
40
+
41
+
42
+# Expected concrete values for each preset at head_dim=128.
43
+# fmt: off
44
+PRESET_EXPECTED = {
45
+ "turboquant_k8v4": dict(
46
+ key_fp8=True, key_quant_bits=8,
47
+ key_mse_bits=0, value_quant_bits=4,
48
+ mse_bits=4, n_centroids=16, centroid_bits=4,
49
+ norm_correction=False,
50
+ key_packed_size=128, value_packed_size=68,
51
+ slot_size=196, slot_size_aligned=196,
52
+ ),
53
+ "turboquant_4bit_nc": dict(
54
+ key_fp8=False, key_quant_bits=4,
55
+ key_mse_bits=4, value_quant_bits=4,
56
+ mse_bits=4, n_centroids=16, centroid_bits=4,
57
+ norm_correction=True,
58
+ key_packed_size=66, value_packed_size=68,
59
+ slot_size=134, slot_size_aligned=134,
60
+ ),
61
+ "turboquant_k3v4_nc": dict(
62
+ key_fp8=False, key_quant_bits=3,
63
+ key_mse_bits=3, value_quant_bits=4,
64
+ mse_bits=3, n_centroids=8, centroid_bits=3,
65
+ norm_correction=True,
66
+ key_packed_size=50, value_packed_size=68,
67
+ slot_size=118, slot_size_aligned=118,
68
+ ),
69
+ "turboquant_3bit_nc": dict(
70
+ key_fp8=False, key_quant_bits=3,
71
+ key_mse_bits=3, value_quant_bits=3,
72
+ mse_bits=3, n_centroids=8, centroid_bits=3,
73
+ norm_correction=True,
74
+ key_packed_size=50, value_packed_size=52,
75
+ slot_size=102, slot_size_aligned=102,
76
+ ),
77
+}
78
+# fmt: on
79
+
80
+
81
+# ============================================================================
82
+# Config tests (CPU-only, no dependencies beyond config.py)
83
+# ============================================================================
84
+
85
+
86
+class TestTurboQuantConfig:
87
+ @pytest.mark.parametrize("preset", ALL_PRESETS)
88
+ def test_preset_parses(self, preset):
89
+ cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
90
+ assert isinstance(cfg, TurboQuantConfig)
91
+
92
+ def test_invalid_preset_raises(self):
93
+ with pytest.raises(ValueError, match="Unknown TurboQuant"):
94
+ TurboQuantConfig.from_cache_dtype("turboquant_invalid", head_dim=128)
95
+
96
+ # ---- Per-preset concrete value checks (table-driven) ----
97
+
98
+ @pytest.mark.parametrize("preset", ALL_PRESETS)
99
+ def test_key_mode(self, preset):
100
+ cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
101
+ exp = PRESET_EXPECTED[preset]
102
+ assert cfg.key_fp8 is exp["key_fp8"]
103
+ assert cfg.key_quant_bits == exp["key_quant_bits"]
104
+ assert cfg.key_mse_bits == exp["key_mse_bits"]
105
+
106
+ @pytest.mark.parametrize("preset", ALL_PRESETS)
107
+ def test_value_mode(self, preset):
108
+ cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
109
+ exp = PRESET_EXPECTED[preset]
110
+ assert cfg.value_quant_bits == exp["value_quant_bits"]
111
+
112
+ @pytest.mark.parametrize("preset", ALL_PRESETS)
113
+ def test_bits_and_centroids(self, preset):
114
+ cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
115
+ exp = PRESET_EXPECTED[preset]
116
+ assert cfg.mse_bits == exp["mse_bits"]
117
+ assert cfg.n_centroids == exp["n_centroids"]
118
+ assert cfg.centroid_bits == exp["centroid_bits"]
119
+
120
+ @pytest.mark.parametrize("preset", ALL_PRESETS)
121
+ def test_norm_correction(self, preset):
122
+ cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
123
+ assert cfg.norm_correction is PRESET_EXPECTED[preset]["norm_correction"]
124
+
125
+ @pytest.mark.parametrize("preset", ALL_PRESETS)
126
+ def test_packed_sizes(self, preset):
127
+ cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
128
+ exp = PRESET_EXPECTED[preset]
129
+ assert cfg.key_packed_size == exp["key_packed_size"]
130
+ assert cfg.value_packed_size == exp["value_packed_size"]
131
+ assert cfg.slot_size == exp["slot_size"]
132
+ assert cfg.slot_size_aligned == exp["slot_size_aligned"]
133
+
134
+ # ---- Cross-preset structural invariants ----
135
+
136
+ @pytest.mark.parametrize("preset", ALL_PRESETS)
137
+ def test_slot_equals_key_plus_value(self, preset):
138
+ cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
139
+ assert cfg.slot_size == cfg.key_packed_size + cfg.value_packed_size
140
+
141
+ @pytest.mark.parametrize("preset", ALL_PRESETS)
142
+ def test_padded_slot_is_even(self, preset):
143
+ cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
144
+ assert cfg.slot_size_aligned >= cfg.slot_size
145
+ assert cfg.slot_size_aligned % 2 == 0, (
146
+ f"slot_size_aligned={cfg.slot_size_aligned} is not even"
147
+ )
148
+
149
+ @pytest.mark.parametrize("preset", ALL_PRESETS)
150
+ def test_key_value_packed_sizes_positive(self, preset):
151
+ cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
152
+ assert cfg.key_packed_size > 0
153
+ assert cfg.value_packed_size > 0
154
+
155
+ @pytest.mark.parametrize("preset", ALL_PRESETS)
156
+ def test_n_centroids_is_2_to_mse_bits(self, preset):
157
+ cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
158
+ assert cfg.n_centroids == 2**cfg.mse_bits
159
+
160
+ @pytest.mark.parametrize("preset", ALL_PRESETS)
161
+ def test_centroid_bits_always_positive(self, preset):
162
+ cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
163
+ assert cfg.centroid_bits > 0
164
+
165
+ @pytest.mark.parametrize("preset", ALL_PRESETS)
166
+ def test_mse_key_or_fp8_exclusive(self, preset):
167
+ """Each preset is either FP8 keys or MSE keys, never both."""
168
+ cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
169
+ if cfg.key_fp8:
170
+ assert cfg.key_mse_bits == 0
171
+ assert cfg.key_quant_bits == 8
172
+ else:
173
+ assert cfg.key_mse_bits > 0
174
+ assert cfg.key_quant_bits in (3, 4)
175
+
176
+ @pytest.mark.parametrize("preset", ALL_PRESETS)
177
+ @pytest.mark.parametrize("head_dim", [64, 96, 128, 256])
178
+ def test_all_presets_all_head_dims(self, preset, head_dim):
179
+ cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=head_dim)
180
+ assert cfg.head_dim == head_dim
181
+ assert cfg.slot_size == cfg.key_packed_size + cfg.value_packed_size
182
+ assert cfg.slot_size_aligned >= cfg.slot_size
183
+ assert cfg.slot_size_aligned % 2 == 0
184
+
185
+ # ---- Boundary skip layers ----
186
+
187
+ def test_boundary_skip_layers_basic(self):
188
+ layers = TurboQuantConfig.get_boundary_skip_layers(32)
189
+ assert layers == ["0", "1", "30", "31"]
190
+
191
+ def test_boundary_skip_layers_zero(self):
192
+ assert TurboQuantConfig.get_boundary_skip_layers(32, 0) == []
193
+
194
+ def test_boundary_skip_layers_small_model(self):
195
+ layers = TurboQuantConfig.get_boundary_skip_layers(4)
196
+ assert layers == ["0", "1", "2", "3"]
197
+
198
+ def test_boundary_skip_layers_cap_at_half(self):
199
+ layers = TurboQuantConfig.get_boundary_skip_layers(8, 10)
200
+ assert len(layers) == 8
201
+
202
+
203
+# ============================================================================
204
+# Centroids tests (CPU-only)
205
+# ============================================================================
206
+
207
+
208
+class TestCentroids:
209
+ @pytest.mark.parametrize("bits,expected_n", [(2, 4), (3, 8), (4, 16)])
210
+ def test_centroids_shape(self, bits, expected_n):
211
+ c = get_centroids(128, bits)
212
+ assert c.shape == (expected_n,)
213
+
214
+ @pytest.mark.parametrize("bits", [2, 3, 4])
215
+ def test_centroids_sorted(self, bits):
216
+ _assert_strictly_sorted(get_centroids(128, bits), "centroids")
217
+
218
+ def test_centroids_cached(self):
219
+ c1 = get_centroids(128, 3)
220
+ c2 = get_centroids(128, 3)
221
+ assert c1 is c2, "get_centroids should return cached object"
222
+
223
+ def test_centroids_different_dims_not_identical(self):
224
+ c64 = get_centroids(64, 3)
225
+ c128 = get_centroids(128, 3)
226
+ assert not torch.equal(c64, c128)
227
+
228
+ @pytest.mark.parametrize("bits", [2, 3, 4])
229
+ def test_centroids_symmetric_around_zero(self, bits):
230
+ """N(0, 1/d) is symmetric, so centroids should be ~symmetric."""
231
+ c = get_centroids(128, bits)
232
+ assert abs(c.mean().item()) < 0.01, "Centroids not centered near 0"
233
+ assert abs(c[0].item() + c[-1].item()) < 0.01
234
+
235
+ @pytest.mark.parametrize("bits", [2, 3, 4])
236
+ def test_centroids_within_4sigma(self, bits):
237
+ """All centroids should be within ~4 sigma of N(0, 1/d)."""
238
+ sigma = math.sqrt(1.0 / 128)
239
+ c = get_centroids(128, bits)
240
+ for i, val in enumerate(c):
241
+ assert abs(val.item()) < 4 * sigma, (
242
+ f"Centroid {i}={val:.6f} outside 4*sigma={4 * sigma:.6f}"
243
+ )
244
+
245
+
246
+class TestLloydMax:
247
+ @pytest.mark.parametrize("bits,expected_n", [(2, 4), (3, 8), (4, 16)])
248
+ def test_solve_shapes(self, bits, expected_n):
249
+ centroids, boundaries = solve_lloyd_max(128, bits)
250
+ assert centroids.shape == (expected_n,)
251
+ assert boundaries.shape == (expected_n - 1,)
252
+
253
+ @pytest.mark.parametrize("bits", [2, 3, 4])
254
+ def test_centroids_sorted(self, bits):
255
+ centroids, _ = solve_lloyd_max(128, bits)
256
+ _assert_strictly_sorted(centroids, "centroids")
257
+
258
+ @pytest.mark.parametrize("bits", [2, 3, 4])
259
+ def test_boundaries_sorted(self, bits):
260
+ _, boundaries = solve_lloyd_max(128, bits)
261
+ _assert_strictly_sorted(boundaries, "boundaries")
262
+
263
+ @pytest.mark.parametrize("bits", [2, 3, 4])
264
+ def test_boundaries_between_centroids(self, bits):
265
+ """Each boundary must lie between its adjacent centroids."""
266
+ centroids, boundaries = solve_lloyd_max(128, bits)
267
+ for i in range(len(boundaries)):
268
+ assert centroids[i] < boundaries[i] < centroids[i + 1], (
269
+ f"Boundary {i}={boundaries[i]:.6f} not between "
270
+ f"c[{i}]={centroids[i]:.6f} and c[{i + 1}]={centroids[i + 1]:.6f}"
271
+ )
272
+
273
+ @pytest.mark.parametrize("bits", [2, 3, 4])
274
+ def test_boundaries_are_midpoints(self, bits):
275
+ """Lloyd-Max boundaries are midpoints of adjacent centroids."""
276
+ centroids, boundaries = solve_lloyd_max(128, bits)
277
+ for i in range(len(boundaries)):
278
+ expected = (centroids[i] + centroids[i + 1]) / 2.0
279
+ assert abs(boundaries[i].item() - expected.item()) < 1e-6
280
+
281
+ def test_solve_deterministic(self):
282
+ c1, b1 = solve_lloyd_max(128, 3)
283
+ c2, b2 = solve_lloyd_max(128, 3)
284
+ assert torch.equal(c1, c2)
285
+ assert torch.equal(b1, b2)
286
+
287
+ def test_solve_dtype_float32(self):
288
+ centroids, boundaries = solve_lloyd_max(128, 3)
289
+ assert centroids.dtype == torch.float32
290
+ assert boundaries.dtype == torch.float32
291
+
292
+ @pytest.mark.parametrize("bits", [3, 4])
293
+ def test_centroids_match_scipy_reference(self, bits):
294
+ """Verify _trapz(n=200) centroids match scipy.integrate.quad reference.
295
+
296
+ This ensures our scipy-free trapezoid integration doesn't silently
297
+ drift from the published Lloyd-Max quality.
298
+ """
299
+ pytest.importorskip("scipy")
300
+ from scipy.integrate import quad
301
+
302
+ d = 128
303
+ sigma2 = 1.0 / d
304
+ sigma = math.sqrt(sigma2)
305
+
306
+ def pdf(x):
307
+ return (1.0 / math.sqrt(2 * math.pi * sigma2)) * math.exp(
308
+ -x * x / (2 * sigma2)
309
+ )
310
+
311
+ n_levels = 2**bits
312
+ lo, hi = -3.5 * sigma, 3.5 * sigma
313
+ ref_centroids = [lo + (hi - lo) * (i + 0.5) / n_levels for i in range(n_levels)]
314
+ for _ in range(200):
315
+ boundaries = [
316
+ (ref_centroids[i] + ref_centroids[i + 1]) / 2.0
317
+ for i in range(n_levels - 1)
318
+ ]
319
+ edges = [lo * 3] + boundaries + [hi * 3]
320
+ new_centroids = []
321
+ for i in range(n_levels):
322
+ a, b = edges[i], edges[i + 1]
323
+ num, _ = quad(lambda x: x * pdf(x), a, b)
324
+ den, _ = quad(pdf, a, b)
325
+ new_centroids.append(num / den if den > 1e-15 else ref_centroids[i])
326
+ if (
327
+ max(abs(new_centroids[i] - ref_centroids[i]) for i in range(n_levels))
328
+ < 1e-10
329
+ ):
330
+ break
331
+ ref_centroids = new_centroids
332
+
333
+ # Compare our _trapz centroids against scipy reference
334
+ our_centroids, _ = solve_lloyd_max(d, bits)
335
+ ref_t = torch.tensor(ref_centroids, dtype=torch.float32)
336
+ max_err = (our_centroids - ref_t).abs().max().item()
337
+ # _trapz(n=200) has ~O(h^2) error vs adaptive quad; 1e-3 is tight
338
+ # enough to catch regression while allowing trapezoid approximation.
339
+ assert max_err < 1e-3, (
340
+ f"d={d}, bits={bits}: max centroid error vs scipy = {max_err:.2e}"
341
+ )
342
+
343
+
344
+# ============================================================================
345
+# Rotation matrix tests (GPU required)
346
+# ============================================================================
347
+
348
+CUDA_AVAILABLE = torch.cuda.is_available()
349
+
350
+
351
+def generate_rotation_matrix(d: int, seed: int, device: str = "cpu") -> torch.Tensor:
352
+ """Haar-distributed random orthogonal matrix via QR (test/benchmark only)."""
353
+ gen = torch.Generator(device="cpu")
354
+ gen.manual_seed(seed)
355
+ G = torch.randn(d, d, generator=gen, device="cpu", dtype=torch.float32)
356
+ Q, R = torch.linalg.qr(G)
357
+ diag_sign = torch.sign(torch.diag(R))
358
+ diag_sign[diag_sign == 0] = 1.0
359
+ Q = Q * diag_sign.unsqueeze(0)
360
+ return Q.to(device)
361
+
362
+
363
+@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
364
+class TestRotationMatrix:
365
+ """Tests for the QR-based rotation (standalone benchmarks only)."""
366
+
367
+ @pytest.mark.parametrize("dim", [64, 96, 128, 256])
368
+ def test_rotation_matrix_shape_and_orthogonal(self, dim):
369
+ Pi = generate_rotation_matrix(dim, seed=42, device="cuda")
370
+ assert Pi.shape == (dim, dim)
371
+ eye = Pi @ Pi.T
372
+ assert torch.allclose(eye, torch.eye(dim, device="cuda"), atol=1e-5), (
373
+ f"Pi not orthogonal for dim={dim}"
374
+ )
375
+
376
+ def test_rotation_matrix_deterministic(self):
377
+ Pi1 = generate_rotation_matrix(128, seed=42)
378
+ Pi2 = generate_rotation_matrix(128, seed=42)
379
+ assert torch.equal(Pi1, Pi2)
380
+
381
+ def test_rotation_matrix_different_seeds(self):
382
+ Pi1 = generate_rotation_matrix(128, seed=42)
383
+ Pi2 = generate_rotation_matrix(128, seed=99)
384
+ assert not torch.equal(Pi1, Pi2)
385
+
386
+ def test_rotation_matrix_det_is_pm1(self):
387
+ """Orthogonal matrix determinant must be +1 or -1."""
388
+ Pi = generate_rotation_matrix(128, seed=42, device="cuda")
389
+ det = torch.linalg.det(Pi)
390
+ assert abs(abs(det.item()) - 1.0) < 1e-4
391
+
392
+
393
+# ============================================================================
394
+# WHT rotation tests (serving path: generate_wht_signs + _build_hadamard)
395
+# ============================================================================
396
+
397
+
398
+def _build_hadamard(d: int, device: str = "cpu") -> torch.Tensor:
399
+ """Reproduce the serving-path Hadamard construction."""
400
+ H = torch.tensor([[1.0]])
401
+ while H.shape[0] < d:
402
+ H = torch.cat([torch.cat([H, H], 1), torch.cat([H, -H], 1)], 0)
403
+ return (H / math.sqrt(d)).to(torch.device(device))
404
+
405
+
406
+@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
407
+class TestWHTRotation:
408
+ """Tests for the WHT rotation actually used in serving."""
409
+
410
+ @pytest.mark.parametrize("dim", [64, 128, 256])
411
+ def test_wht_orthonormal(self, dim):
412
+ """signs * H must be orthonormal: (signs*H) @ (signs*H)^T = I."""
413
+ signs = generate_wht_signs(dim, seed=42, device="cuda")
414
+ H = _build_hadamard(dim, "cuda")
415
+ PiT = (signs.unsqueeze(1) * H).contiguous()
416
+ eye = PiT @ PiT.T
417
+ assert torch.allclose(eye, torch.eye(dim, device="cuda"), atol=1e-5), (
418
+ f"WHT rotation not orthonormal for dim={dim}"
419
+ )
420
+
421
+ @pytest.mark.parametrize("dim", [64, 128, 256])
422
+ def test_wht_self_inverse(self, dim):
423
+ """PiT should be self-inverse: PiT @ PiT = I (up to sign flip)."""
424
+ signs = generate_wht_signs(dim, seed=42, device="cuda")
425
+ H = _build_hadamard(dim, "cuda")
426
+ PiT = (signs.unsqueeze(1) * H).contiguous()
427
+ Pi = PiT.T.contiguous()
428
+ # Pi @ PiT should be identity (rotation then inverse)
429
+ result = Pi @ PiT
430
+ assert torch.allclose(result, torch.eye(dim, device="cuda"), atol=1e-5), (
431
+ f"WHT rotation not self-inverse for dim={dim}"
432
+ )
433
+
434
+ def test_wht_signs_deterministic(self):
435
+ """Same seed must produce identical signs."""
436
+ s1 = generate_wht_signs(128, seed=42)
437
+ s2 = generate_wht_signs(128, seed=42)
438
+ assert torch.equal(s1, s2)
439
+
440
+ def test_wht_signs_different_seeds(self):
441
+ """Different seeds must produce different signs."""
442
+ s1 = generate_wht_signs(128, seed=42)
443
+ s2 = generate_wht_signs(128, seed=99)
444
+ assert not torch.equal(s1, s2)
445
+
446
+ def test_wht_signs_are_pm1(self):
447
+ """All sign values must be exactly +1 or -1."""
448
+ signs = generate_wht_signs(128, seed=42)
449
+ assert torch.all(signs.abs() == 1.0)
450
+
451
+
452
+# ============================================================================
453
+# Store → Decode round-trip test (GPU + Triton required)
454
+# ============================================================================
455
+
456
+
457
+@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
458
+class TestStoreDecodeRoundTrip:
459
+ """End-to-end: store KV into TQ cache, decode, compare vs fp16 ref."""
460
+
461
+ @pytest.mark.parametrize(
462
+ "preset",
463
+ ["turboquant_k8v4", "turboquant_4bit_nc"],
464
+ )
465
+ def test_single_token_roundtrip(self, preset):
466
+ """Store 1 token, decode with query=key, check attention output.
467
+
468
+ For a single token with query=key, attention output should equal
469
+ the value (softmax over single key = 1.0). Quantization error
470
+ means we check cosine similarity rather than exact equality.
471
+ """
472
+ from vllm.model_executor.layers.quantization.turboquant.centroids import (
473
+ solve_lloyd_max,
474
+ )
475
+ from vllm.v1.attention.ops.triton_turboquant_decode import (
476
+ triton_turboquant_decode_attention,
477
+ )
478
+ from vllm.v1.attention.ops.triton_turboquant_store import (
479
+ triton_turboquant_store,
480
+ )
481
+
482
+ cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
483
+ D = 128
484
+ Hk = 4 # num_kv_heads
485
+ Hq = 4 # num_q_heads (no GQA for simplicity)
486
+ B = 1 # single token
487
+ block_size = 16
488
+ num_blocks = 1
489
+
490
+ device = torch.device("cuda")
491
+
492
+ # Generate rotation
493
+ signs = generate_wht_signs(D, seed=42, device=device)
494
+ H = _build_hadamard(D, "cuda")
495
+ PiT = (signs.unsqueeze(1) * H).contiguous().float()
496
+ Pi = PiT.T.contiguous()
497
+
498
+ # Generate centroids
499
+ centroids, _ = solve_lloyd_max(D, cfg.centroid_bits)
500
+ centroids = centroids.float().to(device)
501
+ c_sorted, _ = centroids.sort()
502
+ midpoints = ((c_sorted[:-1] + c_sorted[1:]) / 2).to(device)
503
+
504
+ # Random K, V
505
+ torch.manual_seed(123)
506
+ key = torch.randn(B, Hk, D, device=device, dtype=torch.float16)
507
+ value = torch.randn(B, Hk, D, device=device, dtype=torch.float16)
508
+
509
+ # Allocate KV cache
510
+ padded_slot = cfg.slot_size_aligned
511
+ kv_cache = torch.zeros(
512
+ num_blocks,
513
+ block_size,
514
+ Hk,
515
+ padded_slot,
516
+ device=device,
517
+ dtype=torch.uint8,
518
+ )
519
+ slot_mapping = torch.tensor([0], device=device, dtype=torch.int32)
520
+
521
+ # Store
522
+ triton_turboquant_store(
523
+ key,
524
+ value,
525
+ kv_cache,
526
+ slot_mapping,
527
+ PiT,
528
+ midpoints,
529
+ mse_bits=cfg.key_mse_bits,
530
+ key_packed_size=cfg.key_packed_size,
531
+ value_quant_bits=cfg.effective_value_quant_bits,
532
+ key_fp8=cfg.key_fp8,
533
+ )
534
+
535
+ # Decode: use key as query so attention = softmax([1]) * V = V
536
+ query = key.expand(B, Hq, D).contiguous().to(torch.float16)
537
+ block_table = torch.tensor([[0]], device=device, dtype=torch.int32)
538
+ seq_lens = torch.tensor([1], device=device, dtype=torch.int32)
539
+
540
+ output = triton_turboquant_decode_attention(
541
+ query=query,
542
+ kv_cache=kv_cache,
543
+ block_table=block_table,
544
+ seq_lens=seq_lens,
545
+ Pi=Pi,
546
+ centroids=centroids,
547
+ scale=1.0 / math.sqrt(D),
548
+ mse_bits=cfg.key_mse_bits,
549
+ key_packed_size=cfg.key_packed_size,
550
+ value_quant_bits=cfg.effective_value_quant_bits,
551
+ key_fp8=cfg.key_fp8,
552
+ norm_correction=cfg.norm_correction,
553
+ PiT=PiT,
554
+ max_num_kv_splits=4,
555
+ )
556
+
557
+ # With single KV, output should approximate the stored value.
558
+ # Check per-head cosine similarity > threshold.
559
+ out_fp32 = output.float()
560
+ val_fp32 = value.expand(B, Hq, D).float()
561
+ for h in range(Hq):
562
+ cos_sim = torch.nn.functional.cosine_similarity(
563
+ out_fp32[0, h].unsqueeze(0),
564
+ val_fp32[0, h].unsqueeze(0),
565
+ ).item()
566
+ # FP8 keys should be very accurate; MSE keys have more error
567
+ threshold = 0.95 if cfg.key_fp8 else 0.85
568
+ assert cos_sim > threshold, (
569
+ f"Preset {preset} head {h}: cosine_sim={cos_sim:.4f} < {threshold}"
570
+ )
vllm/config/attention.py
+5
@@ -27,6 +27,11 @@ class AttentionConfig:
27
flash_attn_max_num_splits_for_cuda_graph: int = 32
28
"""Flash Attention max number splits for cuda graph decode."""
29
30
+ tq_max_kv_splits_for_cuda_graph: int = 32
31
+ """TurboQuant max NUM_KV_SPLITS for cuda graph decode.
32
+ Fixes the split count so grid dimensions are constant across captures,
33
+ and buffers can be pre-allocated to avoid inflating the memory estimate."""
34
+
35
use_cudnn_prefill: bool = False
36
"""Whether to use cudnn prefill."""
37
vllm/config/cache.py
+4
@@ -24,6 +24,10 @@ CacheDType = Literal[
24
"fp8_e5m2",
25
"fp8_inc",
26
"fp8_ds_mla",
27
+ "turboquant_k8v4",
28
+ "turboquant_4bit_nc",
29
+ "turboquant_k3v4_nc",
30
+ "turboquant_3bit_nc",
31
"int8_per_token_head",
32
"fp8_per_token_head",
33
]
vllm/engine/arg_utils.py
+38
@@ -1642,6 +1642,31 @@ class EngineArgs:
1642
kv_offloading_backend=self.kv_offloading_backend,
1643
)
1644
1645
+ # TurboQuant: auto-skip first/last 2 layers (boundary protection).
1646
+ # These layers are most sensitive to quantization error.
1647
+ # Users can add extra layers via --kv-cache-dtype-skip-layers.
1648
+ if resolved_cache_dtype.startswith("turboquant_"):
1649
+ if model_config.is_hybrid:
1650
+ raise NotImplementedError(
1651
+ "TurboQuant KV cache is not supported for hybrid "
1652
+ "(attention + Mamba) models. Boundary layer protection "
1653
+ "requires uniform attention layers."
1654
+ )
1655
+ from vllm.model_executor.layers.quantization.turboquant.config import (
1656
+ TurboQuantConfig,
1657
+ )
1658
+
1659
+ num_layers = model_config.hf_text_config.num_hidden_layers
1660
+ boundary = TurboQuantConfig.get_boundary_skip_layers(num_layers)
1661
+ existing = set(cache_config.kv_cache_dtype_skip_layers)
1662
+ merged = sorted(existing | set(boundary), key=lambda x: int(x))
1663
+ cache_config.kv_cache_dtype_skip_layers = merged
1664
+ logger.info(
1665
+ "TQ: skipping layers %s for boundary protection (num_layers=%d)",
1666
+ merged,
1667
+ num_layers,
1668
+ )
1669
+
1670
ray_runtime_env = None
1671
if is_ray_initialized():
1672
# Ray Serve LLM calls `create_engine_config` in the context
@@ -1948,6 +1973,19 @@ class EngineArgs:
1973
self.attention_backend
1974
)
1975
1976
+ # TurboQuant requires FlashAttention 2 — FA3 boundary layers assert
1977
+ # FlashAttentionImpl which fails with TurboQuantAttentionImpl.
1978
+ if resolved_cache_dtype.startswith("turboquant_") and (
1979
+ attention_config.flash_attn_version is None
1980
+ or attention_config.flash_attn_version >= 3
1981
+ ):
1982
+ logger.warning(
1983
+ "TurboQuant is not yet compatible with FlashAttention >= 3. "
1984
+ "Overriding flash_attn_version to 2. To silence this "
1985
+ "warning, pass --attention-config.flash_attn_version=2"
1986
+ )
1987
+ attention_config.flash_attn_version = 2
1988
+
1989
# Mamba config overrides
1990
mamba_config = copy.deepcopy(self.mamba_config)
1991
# Convert string to enum if needed (CLI parsing returns a string)
vllm/model_executor/layers/attention/attention.py
+82
@@ -379,6 +379,10 @@ class Attention(nn.Module, AttentionLayerBase):
379
# Initialize KV cache quantization attributes
380
_init_kv_cache_quant(self, quant_config, prefix)
381
382
+ # Initialize TurboQuant buffers (Pi, S, centroids) if tq cache dtype
383
+ if kv_cache_dtype.startswith("turboquant_"):
384
+ self._init_turboquant_buffers(kv_cache_dtype, head_size, prefix)
385
+
386
# for attn backends supporting query quantization
387
self.query_quant = None
388
if (
@@ -397,6 +401,67 @@ class Attention(nn.Module, AttentionLayerBase):
401
else GroupShape.PER_TENSOR,
402
)
403
404
+ def _init_turboquant_buffers(
405
+ self, cache_dtype: str, head_size: int, prefix: str
406
+ ) -> None:
407
+ """Initialize TurboQuant rotation/projection matrices and centroids."""
408
+ from vllm.model_executor.layers.quantization.turboquant.centroids import (
409
+ get_centroids,
410
+ )
411
+ from vllm.model_executor.layers.quantization.turboquant.config import (
412
+ TurboQuantConfig,
413
+ )
414
+ from vllm.model_executor.layers.quantization.turboquant.quantizer import (
415
+ generate_wht_signs,
416
+ )
417
+
418
+ tq_config = TurboQuantConfig.from_cache_dtype(cache_dtype, head_size)
419
+
420
+ # Each layer needs a unique rotation matrix so quantization errors
421
+ # don't correlate across layers. Stride must exceed max head_dim to
422
+ # ensure non-overlapping RNG streams between adjacent layers.
423
+ _TQ_LAYER_SEED_STRIDE = 1337
424
+
425
+ from vllm.model_executor.models.utils import extract_layer_index
426
+
427
+ layer_idx = extract_layer_index(prefix)
428
+ seed = tq_config.seed + layer_idx * _TQ_LAYER_SEED_STRIDE
429
+
430
+ self.register_buffer(
431
+ "_tq_signs",
432
+ generate_wht_signs(head_size, seed=seed),
433
+ )
434
+ self.register_buffer(
435
+ "_tq_centroids",
436
+ get_centroids(head_size, tq_config.centroid_bits),
437
+ )
438
+ self._tq_config = tq_config
439
+
440
+ # Pre-allocate decode intermediate buffers so model.to(device) moves
441
+ # them to GPU *before* the memory profiler runs. Without this the
442
+ # profiler gives all free memory to KV cache blocks and the first
443
+ # decode OOMs when these buffers are lazily allocated.
444
+ _vllm_cfg = get_current_vllm_config()
445
+ B = _vllm_cfg.scheduler_config.max_num_seqs
446
+ Hq = self.num_heads
447
+ S = _vllm_cfg.attention_config.tq_max_kv_splits_for_cuda_graph
448
+ D = head_size
449
+ self.register_buffer(
450
+ "_tq_mid_o_buf",
451
+ torch.empty(B, Hq, S, D + 1, dtype=torch.float32),
452
+ persistent=False,
453
+ )
454
+ self.register_buffer(
455
+ "_tq_output_buf",
456
+ torch.empty(B, Hq, D, dtype=torch.float32),
457
+ persistent=False,
458
+ )
459
+ self.register_buffer(
460
+ "_tq_lse_buf",
461
+ torch.empty(B, Hq, dtype=torch.float32),
462
+ persistent=False,
463
+ )
464
+
465
def forward(
466
self,
467
query: torch.Tensor,
@@ -544,6 +609,23 @@ class Attention(nn.Module, AttentionLayerBase):
609
kv_quant_mode=quant_mode,
610
sliding_window=self.sliding_window,
611
)
612
+ elif self.kv_cache_dtype.startswith("turboquant_"):
613
+ from vllm.model_executor.layers.quantization.turboquant.config import (
614
+ TurboQuantConfig,
615
+ )
616
+ from vllm.v1.kv_cache_interface import TQFullAttentionSpec
617
+
618
+ tq_config = TurboQuantConfig.from_cache_dtype(
619
+ self.kv_cache_dtype, self.head_size
620
+ )
621
+ return TQFullAttentionSpec(
622
+ block_size=block_size,
623
+ num_kv_heads=self.num_kv_heads,
624
+ head_size=self.head_size,
625
+ head_size_v=self.head_size,
626
+ dtype=self.kv_cache_torch_dtype,
627
+ tq_slot_size=tq_config.slot_size_aligned,
628
+ )
629
else:
630
return FullAttentionSpec(
631
block_size=block_size,
vllm/model_executor/layers/quantization/turboquant/__init__.py
new
+14
@@ -0,0 +1,14 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""TurboQuant: Near-optimal KV-cache quantization for vLLM.
4
+
5
+PolarQuant compression: random rotation + per-coordinate Lloyd-Max
6
+scalar quantization for keys, uniform quantization for values.
7
+
8
+Reference: "TurboQuant: Online Vector Quantization with Near-optimal
9
+Distortion Rate" (ICLR 2026), Zandieh et al.
10
+"""
11
+
12
+from vllm.model_executor.layers.quantization.turboquant.config import TurboQuantConfig
13
+
14
+__all__ = ["TurboQuantConfig"]
vllm/model_executor/layers/quantization/turboquant/centroids.py
new
+86
@@ -0,0 +1,86 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Lloyd-Max optimal scalar quantizer for TurboQuant.
4
+
5
+After rotating a d-dimensional unit vector by a random orthogonal matrix,
6
+each coordinate approximately follows N(0, 1/d) for d >= 64.
7
+We solve the Lloyd-Max conditions to find optimal centroids.
8
+
9
+Based on: turboquant-pytorch/lloyd_max.py (Zandieh et al.)
10
+"""
11
+
12
+import math
13
+from functools import lru_cache
14
+
15
+import torch
16
+
17
+
18
+def _gaussian_pdf(x: float, sigma2: float) -> float:
19
+ return (1.0 / math.sqrt(2 * math.pi * sigma2)) * math.exp(-x * x / (2 * sigma2))
20
+
21
+
22
+def _trapz(f, a: float, b: float, n: int = 200) -> float:
23
+ """Trapezoidal numerical integration (replaces scipy.integrate.quad)."""
24
+ h = (b - a) / n
25
+ result = 0.5 * (f(a) + f(b))
26
+ for i in range(1, n):
27
+ result += f(a + i * h)
28
+ return result * h
29
+
30
+
31
+def solve_lloyd_max(
32
+ d: int,
33
+ bits: int,
34
+ max_iter: int = 200,
35
+ tol: float = 1e-10,
36
+) -> tuple[torch.Tensor, torch.Tensor]:
37
+ """Solve Lloyd-Max optimal quantizer for N(0, 1/d) distribution.
38
+
39
+ Args:
40
+ d: Vector dimension (determines variance = 1/d).
41
+ bits: Number of quantization bits.
42
+ max_iter: Maximum Lloyd-Max iterations.
43
+ tol: Convergence tolerance.
44
+
45
+ Returns:
46
+ centroids: Sorted tensor of 2^bits optimal centroids.
47
+ boundaries: Sorted tensor of 2^bits - 1 decision boundaries.
48
+ """
49
+ n_levels = 2**bits
50
+ sigma2 = 1.0 / d
51
+ sigma = math.sqrt(sigma2)
52
+
53
+ def pdf(x):
54
+ return _gaussian_pdf(x, sigma2)
55
+
56
+ lo, hi = -3.5 * sigma, 3.5 * sigma
57
+ centroids = [lo + (hi - lo) * (i + 0.5) / n_levels for i in range(n_levels)]
58
+
59
+ for _ in range(max_iter):
60
+ boundaries = [
61
+ (centroids[i] + centroids[i + 1]) / 2.0 for i in range(n_levels - 1)
62
+ ]
63
+ edges = [lo * 3] + boundaries + [hi * 3]
64
+ new_centroids = []
65
+ for i in range(n_levels):
66
+ a, b = edges[i], edges[i + 1]
67
+ num = _trapz(lambda x: x * pdf(x), a, b)
68
+ den = _trapz(pdf, a, b)
69
+ new_centroids.append(num / den if den > 1e-15 else centroids[i])
70
+
71
+ if max(abs(new_centroids[i] - centroids[i]) for i in range(n_levels)) < tol:
72
+ break
73
+ centroids = new_centroids
74
+
75
+ boundaries = [(centroids[i] + centroids[i + 1]) / 2.0 for i in range(n_levels - 1)]
76
+ return (
77
+ torch.tensor(centroids, dtype=torch.float32),
78
+ torch.tensor(boundaries, dtype=torch.float32),
79
+ )
80
+
81
+
82
+@lru_cache(maxsize=32)
83
+def get_centroids(d: int, bits: int) -> torch.Tensor:
84
+ """Get precomputed Lloyd-Max centroids (cached)."""
85
+ centroids, _ = solve_lloyd_max(d, bits)
86
+ return centroids
vllm/model_executor/layers/quantization/turboquant/config.py
new
+185
@@ -0,0 +1,185 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""TurboQuant configuration."""
4
+
5
+import math
6
+from dataclasses import dataclass
7
+
8
+# Named TQ presets: each maps to frozen config parameters.
9
+# key_quant_bits: 8 = FP8 keys, 3-4 = MSE (Lloyd-Max) quantized keys.
10
+# value_quant_bits: 3-4 = uniform quantized values.
11
+TQ_PRESETS: dict[str, dict] = {
12
+ "turboquant_k8v4": {
13
+ "key_quant_bits": 8,
14
+ "value_quant_bits": 4,
15
+ "norm_correction": False,
16
+ },
17
+ "turboquant_4bit_nc": {
18
+ "key_quant_bits": 4,
19
+ "value_quant_bits": 4,
20
+ "norm_correction": True,
21
+ },
22
+ "turboquant_k3v4_nc": {
23
+ "key_quant_bits": 3,
24
+ "value_quant_bits": 4,
25
+ "norm_correction": True,
26
+ },
27
+ "turboquant_3bit_nc": {
28
+ "key_quant_bits": 3,
29
+ "value_quant_bits": 3,
30
+ "norm_correction": True,
31
+ },
32
+}
33
+
34
+
35
+@dataclass
36
+class TurboQuantConfig:
37
+ """Configuration for TurboQuant KV-cache quantization.
38
+
39
+ Uses PolarQuant (WHT rotation + Lloyd-Max scalar quantization) for keys
40
+ and uniform quantization for values. QJL is intentionally omitted —
41
+ community consensus (5+ independent groups) found it hurts attention
42
+ quality by amplifying variance through softmax.
43
+
44
+ Named presets (use via --kv-cache-dtype):
45
+ turboquant_k8v4: FP8 keys + 4-bit values, 2.6x, +1.17% PPL
46
+ turboquant_4bit_nc: 4-bit MSE keys + 4-bit values + NC, 3.8x, +2.71%
47
+ turboquant_k3v4_nc: 3-bit MSE keys + 4-bit values + NC, ~3.5x, +10.63%
48
+ turboquant_3bit_nc: 3-bit MSE keys + 3-bit values + NC, 4.9x, +20.59%
49
+
50
+ Args:
51
+ head_dim: Attention head dimension (e.g. 64, 96, 128).
52
+ key_quant_bits: Bits for key quantization. 8 = FP8 keys (no
53
+ rotation/MSE). 3-4 = Lloyd-Max MSE quantized keys.
54
+ value_quant_bits: Bits per value dimension for uniform quantization.
55
+ 3 = 8 levels, 4 = 16 levels (default).
56
+ seed: Base seed for deterministic random matrix generation.
57
+ Actual seed per layer = seed + layer_idx * 1337.
58
+ norm_correction: Re-normalize centroid vectors to unit norm before
59
+ inverse rotation during dequant. Fixes quantization-induced norm
60
+ distortion, improving PPL by ~0.8% at 4-bit.
61
+ """
62
+
63
+ head_dim: int = 128
64
+ key_quant_bits: int = 3 # 3-4 = MSE keys, 8 = FP8 keys
65
+ value_quant_bits: int = 4 # 3-4 = uniform quantized values
66
+ seed: int = 42
67
+ norm_correction: bool = False
68
+
69
+ @property
70
+ def key_fp8(self) -> bool:
71
+ """Whether keys are stored as FP8 — no rotation/quantization needed."""
72
+ return self.key_quant_bits == 8
73
+
74
+ @property
75
+ def mse_bits(self) -> int:
76
+ """MSE quantizer bit-width (determines centroid count: 2^mse_bits).
77
+
78
+ For MSE key modes, equals key_quant_bits.
79
+ For FP8 key mode, falls back to value_quant_bits (centroids are still
80
+ needed for continuation-prefill dequant and decode kernel params).
81
+ """
82
+ if self.key_fp8:
83
+ return self.value_quant_bits
84
+ return self.key_quant_bits
85
+
86
+ @property
87
+ def key_mse_bits(self) -> int:
88
+ """MSE bits actually used for key quantization (0 if FP8 keys)."""
89
+ if self.key_fp8:
90
+ return 0
91
+ return self.key_quant_bits
92
+
93
+ @property
94
+ def centroid_bits(self) -> int:
95
+ """Bits for centroid generation — always non-zero."""
96
+ return self.mse_bits
97
+
98
+ @property
99
+ def n_centroids(self) -> int:
100
+ return 2**self.mse_bits
101
+
102
+ @property
103
+ def key_packed_size(self) -> int:
104
+ """Packed bytes for a single KEY vector.
105
+
106
+ FP8 mode (key_quant_bits=8):
107
+ head_dim bytes (1 byte per element, no overhead).
108
+
109
+ TQ mode:
110
+ - MSE indices: ceil(head_dim * key_mse_bits / 8) bytes
111
+ - vec_norm: 2 bytes (float16)
112
+ """
113
+ if self.key_fp8:
114
+ return self.head_dim # 1 byte per element
115
+ mse_bytes = math.ceil(self.head_dim * self.key_mse_bits / 8)
116
+ norm_bytes = 2 # vec_norm fp16
117
+ return mse_bytes + norm_bytes
118
+
119
+ @property
120
+ def effective_value_quant_bits(self) -> int:
121
+ """Actual bits used for value storage."""
122
+ return self.value_quant_bits
123
+
124
+ @property
125
+ def value_packed_size(self) -> int:
126
+ """Packed bytes for a single VALUE vector.
127
+
128
+ Uniform quantization: ceil(head_dim * bits / 8) + 4 bytes (scale + zero fp16).
129
+ """
130
+ data_bytes = math.ceil(self.head_dim * self.value_quant_bits / 8)
131
+ return data_bytes + 4 # +2 scale(fp16) +2 zero(fp16)
132
+
133
+ @property
134
+ def slot_size(self) -> int:
135
+ """Total packed bytes per head per position (key + value combined).
136
+
137
+ Layout: [key_packed | value_packed]
138
+ """
139
+ return self.key_packed_size + self.value_packed_size
140
+
141
+ @property
142
+ def slot_size_aligned(self) -> int:
143
+ """Slot size rounded up to next even number.
144
+
145
+ Even-number is required so effective_head_size = slot_size_aligned // 2
146
+ is integral.
147
+ """
148
+ s = self.slot_size
149
+ return s + (s % 2) # round up to even
150
+
151
+ @staticmethod
152
+ def get_boundary_skip_layers(num_layers: int, n: int = 2) -> list[str]:
153
+ """Get layer indices to skip TQ compression (boundary protection).
154
+
155
+ Returns first N and last N layer indices as strings, suitable for
156
+ kv_cache_dtype_skip_layers.
157
+ """
158
+ if n <= 0 or num_layers <= 0:
159
+ return []
160
+ n = min(n, num_layers // 2) # don't skip more than half
161
+ first = list(range(n))
162
+ last = list(range(num_layers - n, num_layers))
163
+ # Deduplicate (if num_layers <= 2*n)
164
+ indices = sorted(set(first + last))
165
+ return [str(i) for i in indices]
166
+
167
+ @staticmethod
168
+ def from_cache_dtype(cache_dtype: str, head_dim: int) -> "TurboQuantConfig":
169
+ """Create config from a named preset.
170
+
171
+ Valid presets: turboquant_k8v4, turboquant_4bit_nc, etc.
172
+ """
173
+ if cache_dtype not in TQ_PRESETS:
174
+ valid = ", ".join(TQ_PRESETS.keys())
175
+ raise ValueError(
176
+ f"Unknown TurboQuant cache dtype: {cache_dtype!r}. "
177
+ f"Valid presets: {valid}"
178
+ )
179
+ preset = TQ_PRESETS[cache_dtype]
180
+ return TurboQuantConfig(
181
+ head_dim=head_dim,
182
+ key_quant_bits=preset["key_quant_bits"],
183
+ value_quant_bits=preset["value_quant_bits"],
184
+ norm_correction=preset["norm_correction"],
185
+ )
vllm/model_executor/layers/quantization/turboquant/quantizer.py
new
+24
@@ -0,0 +1,24 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""TurboQuant quantizer utilities.
4
+
5
+Serving path uses generate_wht_signs() for WHT rotation sign buffers.
6
+Triton kernels handle all quantization, packing, and dequantization on GPU.
7
+"""
8
+
9
+import torch
10
+
11
+_CPU = torch.device("cpu")
12
+
13
+
14
+def generate_wht_signs(d: int, seed: int, device: torch.device = _CPU) -> torch.Tensor:
15
+ """Generate deterministic random ±1 signs for WHT rotation.
16
+
17
+ Used with Walsh-Hadamard Transform for per-layer rotation randomization.
18
+ Same seed derivation as QR (per-layer via seed + layer_idx * stride).
19
+ """
20
+ gen = torch.Generator(device="cpu")
21
+ gen.manual_seed(seed)
22
+ bits = torch.randint(0, 2, (d,), generator=gen, device="cpu")
23
+ signs = bits.float() * 2 - 1
24
+ return signs.to(device)
vllm/platforms/cuda.py
+5
@@ -255,6 +255,11 @@ class CudaPlatformBase(Platform):
255
valid_backends_priorities = []
256
invalid_reasons: dict[AttentionBackendEnum, tuple[int, list[str]]] = {}
257
258
+ # TurboQuant KV cache: route directly to TQ backend
259
+ kv_cache_dtype = attn_selector_config.kv_cache_dtype
260
+ if kv_cache_dtype is not None and kv_cache_dtype.startswith("turboquant_"):
261
+ return [(AttentionBackendEnum.TURBOQUANT, 0)], {}
262
+
263
backend_priorities = _get_backend_priorities(
264
attn_selector_config.use_mla,
265
device_capability,
vllm/platforms/xpu.py
+6
@@ -61,6 +61,12 @@ class XPUPlatform(Platform):
61
"only NHD layout is supported by XPU attention kernels."
62
)
63
64
+ # TurboQuant KV cache: route directly to TQ backend
65
+ kv_cache_dtype = attn_selector_config.kv_cache_dtype
66
+ if kv_cache_dtype is not None and kv_cache_dtype.startswith("turboquant_"):
67
+ logger.info_once("Using TurboQuant attention backend.")
68
+ return AttentionBackendEnum.TURBOQUANT.get_path()
69
+
70
dtype = attn_selector_config.dtype
71
if attn_selector_config.use_sparse:
72
logger.info_once("Using XPU MLA Sparse backend.")
vllm/utils/torch_utils.py
+4
@@ -42,6 +42,10 @@ STR_DTYPE_TO_TORCH_DTYPE = {
42
"fp8_per_token_head": torch.uint8,
43
"fp8_inc": torch.float8_e4m3fn,
44
"fp8_ds_mla": torch.uint8,
45
+ "turboquant_k8v4": torch.uint8,
46
+ "turboquant_4bit_nc": torch.uint8,
47
+ "turboquant_k3v4_nc": torch.uint8,
48
+ "turboquant_3bit_nc": torch.uint8,
49
}
50
51
TORCH_DTYPE_TO_NUMPY_DTYPE = {
vllm/v1/attention/backends/registry.py
+1
@@ -82,6 +82,7 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
82
"RocmAiterUnifiedAttentionBackend"
83
)
84
CPU_ATTN = "vllm.v1.attention.backends.cpu_attn.CPUAttentionBackend"
85
+ TURBOQUANT = "vllm.v1.attention.backends.turboquant_attn.TurboQuantAttentionBackend"
86
# Placeholder for third-party/custom backends - must be registered before use
87
# set to None to avoid alias with other backend, whose value is an empty string
88
CUSTOM = None
vllm/v1/attention/backends/turboquant_attn.py
new
+812
@@ -0,0 +1,812 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""TurboQuant attention backend for vLLM.
4
+
5
+Prefill: Standard scaled dot-product attention on uncompressed K/V,
6
+ then quantize K and store K+V into combined cache slot.
7
+Decode: Compute TQ attention scores from compressed cache,
8
+ unpack FP16 values, softmax + weighted sum.
9
+
10
+Cache layout (no leading 2 dimension):
11
+ (num_blocks, block_size, num_kv_heads, slot_size)
12
+ where slot_size = key_packed_size + value_fp16_size
13
+
14
+Per-head per-position slot layout:
15
+ [key_packed (kps bytes) | value_fp16 (D*2 bytes)]
16
+ For turboquant_k3v4_nc head_dim=256: [100 bytes key | 512 bytes value] = 612
17
+"""
18
+
19
+import functools
20
+import math
21
+from dataclasses import dataclass
22
+from typing import Any, ClassVar
23
+
24
+import torch
25
+import torch.nn.functional as F
26
+
27
+from vllm.config import get_current_vllm_config
28
+from vllm.config.cache import CacheDType
29
+from vllm.triton_utils import triton
30
+from vllm.v1.attention.backend import (
31
+ AttentionBackend,
32
+ AttentionCGSupport,
33
+ AttentionImpl,
34
+ AttentionLayer,
35
+ AttentionMetadata,
36
+ AttentionMetadataBuilder,
37
+ AttentionType,
38
+ CommonAttentionMetadata,
39
+ MultipleOf,
40
+)
41
+from vllm.v1.attention.backends.fa_utils import (
42
+ is_flash_attn_varlen_func_available,
43
+)
44
+from vllm.v1.attention.backends.utils import split_decodes_and_prefills
45
+from vllm.v1.attention.ops.triton_turboquant_decode import (
46
+ _tq_full_dequant_kv,
47
+ _use_fp8_e4b15,
48
+ triton_turboquant_decode_attention,
49
+)
50
+from vllm.v1.attention.ops.triton_turboquant_store import triton_turboquant_store
51
+
52
+_HAS_FLASH_ATTN = is_flash_attn_varlen_func_available()
53
+if _HAS_FLASH_ATTN:
54
+ from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func
55
+
56
+# Continuation prefill: for small continuation chunks (q_len ≤ threshold),
57
+# use the TQ decode kernel directly instead of full-dequant + flash_attn.
58
+# do_kv_cache_update already stored all tokens to TQ cache, so the decode
59
+# kernel can read them efficiently. This avoids O(cached_len) dequant work
60
+# per continuation, eliminating the O(N²/chunk_size) collapse at long context.
61
+_CONTINUATION_DECODE_THRESHOLD = 128
62
+
63
+
64
+def _build_hadamard(d: int, device_str: str) -> torch.Tensor:
65
+ """Orthonormal Hadamard matrix (Sylvester construction), cached per (d, device).
66
+
67
+ Precomputed D×D matrix enables matmul-based WHT — single cuBLAS GEMM
68
+ instead of log2(D) butterfly kernel launches. 64KB for D=128.
69
+ """
70
+ # Normalize device string so "cuda" and "cuda:0" hit the same cache entry.
71
+ return _build_hadamard_cached(d, str(torch.device(device_str)))
72
+
73
+
74
+@functools.cache
75
+def _build_hadamard_cached(d: int, device_str: str) -> torch.Tensor:
76
+ H = torch.tensor([[1.0]])
77
+ while H.shape[0] < d:
78
+ H = torch.cat([torch.cat([H, H], 1), torch.cat([H, -H], 1)], 0)
79
+ return (H / math.sqrt(d)).to(torch.device(device_str))
80
+
81
+
82
+class TurboQuantAttentionBackend(AttentionBackend):
83
+ """Attention backend using TurboQuant KV-cache compression."""
84
+
85
+ accept_output_buffer: bool = True
86
+ forward_includes_kv_cache_update: bool = False
87
+
88
+ supported_dtypes: ClassVar[list[torch.dtype]] = [
89
+ torch.float16,
90
+ torch.bfloat16,
91
+ ]
92
+ supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
93
+ "turboquant_k8v4",
94
+ "turboquant_4bit_nc",
95
+ "turboquant_k3v4_nc",
96
+ "turboquant_3bit_nc",
97
+ ]
98
+
99
+ @staticmethod
100
+ def get_name() -> str:
101
+ return "TURBOQUANT"
102
+
103
+ @staticmethod
104
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
105
+ return [16, 32, 64, 128]
106
+
107
+ @classmethod
108
+ def supports_attn_type(cls, attn_type: str) -> bool:
109
+ return attn_type == AttentionType.DECODER
110
+
111
+ @classmethod
112
+ def supports_per_head_quant_scales(cls) -> bool:
113
+ return False
114
+
115
+ @staticmethod
116
+ def get_impl_cls() -> type["TurboQuantAttentionImpl"]:
117
+ return TurboQuantAttentionImpl
118
+
119
+ @staticmethod
120
+ def get_builder_cls() -> type["TurboQuantMetadataBuilder"]:
121
+ return TurboQuantMetadataBuilder
122
+
123
+ @staticmethod
124
+ def get_kv_cache_shape(
125
+ num_blocks: int,
126
+ block_size: int,
127
+ num_kv_heads: int,
128
+ head_size: int,
129
+ cache_dtype_str: str = "turboquant_4bit_nc",
130
+ ) -> tuple[int, ...]:
131
+ """Combined K+V cache shape — no leading 2 dimension.
132
+
133
+ Standard attention backends use (2, num_blocks, block_size, num_kv_heads,
134
+ head_dim) with a leading 2 to separate K and V. TurboQuant packs K+V
135
+ into a single interleaved slot per head per position, so the cache is:
136
+
137
+ (num_blocks, block_size, num_kv_heads, slot_size_aligned)
138
+
139
+ Each slot = [key_packed | value_packed | padding].
140
+ This is safe because TQ has its own get_kv_cache_shape override and
141
+ never shares cache tensors with other backends. Layers that fall back
142
+ to native dtype via kv_cache_dtype_skip_layers get their own
143
+ standard-shaped cache allocation.
144
+
145
+ head_size is the model's real head_dim. slot_size_aligned is computed
146
+ from the TQ config to ensure correct cache allocation for all head dims.
147
+ """
148
+ from vllm.model_executor.layers.quantization.turboquant.config import (
149
+ TurboQuantConfig,
150
+ )
151
+
152
+ tq_config = TurboQuantConfig.from_cache_dtype(cache_dtype_str, head_size)
153
+ return (num_blocks, block_size, num_kv_heads, tq_config.slot_size_aligned)
154
+
155
+ @classmethod
156
+ def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool:
157
+ if kv_cache_dtype is None:
158
+ return False
159
+ return kv_cache_dtype.startswith("turboquant_")
160
+
161
+ @classmethod
162
+ def supports_head_size(cls, head_size: int) -> bool:
163
+ # head_size from spec is effective_head_size (padded_slot//2),
164
+ # not the model's actual head_dim. Accept any positive value.
165
+ return head_size > 0
166
+
167
+
168
+@dataclass
169
+class TurboQuantMetadata(AttentionMetadata):
170
+ """Metadata for TurboQuant attention."""
171
+
172
+ seq_lens: torch.Tensor # (num_reqs,) — total context length per request
173
+ slot_mapping: torch.Tensor # (num_tokens,) — cache slot for each token
174
+ block_table: torch.Tensor # (num_reqs, max_num_blocks)
175
+ query_start_loc: torch.Tensor # (num_reqs + 1,) — cu_seqlens for queries
176
+ num_actual_tokens: int = 0 # actual tokens (excluding padding)
177
+ max_query_len: int = 0 # longest query in batch
178
+ max_seq_len: int = 0 # longest context in batch
179
+ is_prefill: bool = False
180
+ num_decodes: int = 0 # number of decode requests (first in batch)
181
+ num_decode_tokens: int = 0 # tokens from decode requests
182
+
183
+
184
+class TurboQuantMetadataBuilder(AttentionMetadataBuilder[TurboQuantMetadata]):
185
+ """Builds TurboQuantMetadata from scheduler output."""
186
+
187
+ _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
188
+
189
+ def __init__(self, kv_cache_spec, layer_names, vllm_config, device):
190
+ super().__init__(kv_cache_spec, layer_names, vllm_config, device)
191
+ self._init_reorder_batch_threshold(1, supports_spec_as_decode=False)
192
+
193
+ def build_for_cudagraph_capture(
194
+ self, common_attn_metadata: CommonAttentionMetadata
195
+ ) -> TurboQuantMetadata:
196
+ attn_metadata = self.build(0, common_attn_metadata)
197
+ # Set seq_lens to 1 so CUDA graph capture is fast
198
+ # (real seq_lens are filled at replay time).
199
+ attn_metadata.seq_lens.fill_(1)
200
+ return attn_metadata
201
+
202
+ def build(self, common_prefix_len, common_attn_metadata, fast_build=False):
203
+ """Build TurboQuantMetadata from common attention metadata."""
204
+ cam = common_attn_metadata
205
+
206
+ # With reorder_batch_threshold=1, the model runner guarantees
207
+ # decodes come first in the batch. split_decodes_and_prefills
208
+ # finds the boundary (operates on CPU tensors — no GPU sync).
209
+ assert self.reorder_batch_threshold is not None
210
+ num_decodes, num_prefills, num_decode_tokens, _ = split_decodes_and_prefills(
211
+ cam, decode_threshold=self.reorder_batch_threshold
212
+ )
213
+
214
+ return TurboQuantMetadata(
215
+ seq_lens=cam.seq_lens,
216
+ slot_mapping=cam.slot_mapping,
217
+ block_table=cam.block_table_tensor,
218
+ query_start_loc=cam.query_start_loc,
219
+ num_actual_tokens=cam.num_actual_tokens,
220
+ max_query_len=cam.max_query_len,
221
+ max_seq_len=cam.max_seq_len,
222
+ is_prefill=(cam.max_query_len > 1),
223
+ num_decodes=num_decodes,
224
+ num_decode_tokens=num_decode_tokens,
225
+ )
226
+
227
+
228
+class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
229
+ """TurboQuant attention implementation.
230
+
231
+ Vectorized PyTorch: batch quantize/store, vectorized bit-unpack
232
+ decode with einsum scores and value gather.
233
+ """
234
+
235
+ supports_quant_query_input: bool = False
236
+
237
+ def __init__(
238
+ self,
239
+ num_heads: int,
240
+ head_size: int,
241
+ scale: float,
242
+ num_kv_heads: int | None = None,
243
+ alibi_slopes: list[float] | None = None,
244
+ sliding_window: int | None = None,
245
+ kv_cache_dtype: str = "auto",
246
+ logits_soft_cap: float | None = None,
247
+ attn_type: str = AttentionType.DECODER,
248
+ kv_sharing_target_layer_name: str | None = None,
249
+ **kwargs,
250
+ ):
251
+ self.num_heads = num_heads
252
+ self.head_size = head_size
253
+ self.scale = scale
254
+ self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
255
+ self.num_kv_groups = num_heads // self.num_kv_heads
256
+ self.kv_cache_dtype = kv_cache_dtype
257
+
258
+ from vllm.model_executor.layers.quantization.turboquant.config import (
259
+ TurboQuantConfig,
260
+ )
261
+
262
+ self.tq_config = TurboQuantConfig.from_cache_dtype(kv_cache_dtype, head_size)
263
+
264
+ # Pre-compute kernel constants from config (avoid repeated arithmetic)
265
+ cfg = self.tq_config
266
+ self._mse_bytes = (
267
+ math.ceil(head_size * cfg.key_mse_bits / 8)
268
+ if not cfg.key_fp8
269
+ else head_size
270
+ )
271
+ self._val_data_bytes = math.ceil(head_size * cfg.effective_value_quant_bits / 8)
272
+ self._n_centroids = cfg.n_centroids if not cfg.key_fp8 else 1
273
+
274
+ # Fixed NUM_KV_SPLITS (grid dims must be constant for cudagraph,
275
+ # and benchmarks show no regression vs dynamic in eager mode).
276
+ vllm_config = get_current_vllm_config()
277
+ self.max_num_kv_splits = (
278
+ vllm_config.attention_config.tq_max_kv_splits_for_cuda_graph
279
+ )
280
+
281
+ def _ensure_on_device(self, layer, device):
282
+ """One-time derivation of TQ buffers (rotation matrices, midpoints).
283
+
284
+ Registered buffers (_tq_signs, _tq_centroids) are already on the
285
+ correct device via register_buffer + model.to(device).
286
+ """
287
+ if not hasattr(layer, "_tq_cached"):
288
+ D = layer._tq_signs.shape[0]
289
+ signs = layer._tq_signs.to(device=device, dtype=torch.float32)
290
+
291
+ # WHT rotation: orthonormal + self-inverse, enabling future
292
+ # in-kernel butterfly fusion and trivial inverse for continuation.
293
+ H = _build_hadamard(D, str(device))
294
+ layer._tq_PiT = (signs.unsqueeze(1) * H).contiguous()
295
+ layer._tq_Pi = layer._tq_PiT.T.contiguous()
296
+
297
+ c = layer._tq_centroids.to(device=device, dtype=torch.float32)
298
+ # Precompute midpoints for threshold-based quantization
299
+ c_sorted, _ = c.sort()
300
+ layer._tq_midpoints = (c_sorted[:-1] + c_sorted[1:]) / 2
301
+ # Decode buffers (_tq_mid_o_buf, _tq_output_buf, _tq_lse_buf)
302
+ # are pre-allocated via register_buffer in Attention.__init__
303
+ # and moved to GPU by model.to(device) — no allocation needed
304
+ # here. The memory profiler sees them before KV cache sizing.
305
+ layer._tq_cached = True
306
+
307
+ def do_kv_cache_update(
308
+ self,
309
+ layer: torch.nn.Module,
310
+ key: torch.Tensor,
311
+ value: torch.Tensor,
312
+ kv_cache: torch.Tensor,
313
+ slot_mapping: torch.Tensor,
314
+ ) -> None:
315
+ """Store compressed K/V into the combined TQ cache.
316
+
317
+ Called as a separate custom op (unified_kv_cache_update) BEFORE
318
+ the attention forward, matching FlashAttention's split pattern.
319
+ slot_mapping is already sliced to num_actual_tokens by the caller.
320
+ """
321
+ N = slot_mapping.shape[0]
322
+ if N <= 0:
323
+ return
324
+
325
+ device = key.device
326
+ self._ensure_on_device(layer, device)
327
+
328
+ k = key[:N].view(N, self.num_kv_heads, self.head_size)
329
+ v = value[:N].view(N, self.num_kv_heads, self.head_size)
330
+ self._store_kv(k, v, kv_cache, slot_mapping, layer)
331
+
332
+ def forward(
333
+ self,
334
+ layer: AttentionLayer,
335
+ query: torch.Tensor,
336
+ key: torch.Tensor,
337
+ value: torch.Tensor,
338
+ kv_cache: torch.Tensor,
339
+ attn_metadata: "TurboQuantMetadata",
340
+ output: torch.Tensor | None = None,
341
+ output_scale: torch.Tensor | None = None,
342
+ output_block_scale: torch.Tensor | None = None,
343
+ ) -> torch.Tensor:
344
+ num_tokens = query.shape[0]
345
+
346
+ if output is None:
347
+ output = torch.zeros(
348
+ num_tokens,
349
+ self.num_heads * self.head_size,
350
+ dtype=query.dtype,
351
+ device=query.device,
352
+ )
353
+
354
+ if attn_metadata is None:
355
+ return output.fill_(0)
356
+
357
+ # Slice to actual tokens
358
+ N = attn_metadata.num_actual_tokens
359
+ if N <= 0:
360
+ return output.fill_(0)
361
+
362
+ q = query[:N].view(N, self.num_heads, self.head_size)
363
+
364
+ # Get TQ buffers, ensure on device (one-time migration).
365
+ # Use Any-typed alias for dynamic _tq_* attrs set by _ensure_on_device.
366
+ tq_layer: Any = layer
367
+ device = q.device
368
+ self._ensure_on_device(tq_layer, device)
369
+ Pi = tq_layer._tq_Pi
370
+ PiT = tq_layer._tq_PiT
371
+ centroids = tq_layer._tq_centroids
372
+
373
+ # Compute attention (KV cache was already updated by do_kv_cache_update)
374
+ # With reorder_batch_threshold=1, decodes come first in the batch.
375
+ # num_decodes/num_decode_tokens from metadata give the split point.
376
+ num_decodes = attn_metadata.num_decodes
377
+ num_decode_tokens = attn_metadata.num_decode_tokens
378
+
379
+ if not attn_metadata.is_prefill:
380
+ # Pure decode batch — fast path
381
+ attn_out = self._decode_attention(
382
+ q, kv_cache, attn_metadata, Pi, centroids, PiT, layer
383
+ )
384
+ elif num_decodes == 0:
385
+ # Pure prefill batch
386
+ k = key[:N].view(N, self.num_kv_heads, self.head_size)
387
+ v = value[:N].view(N, self.num_kv_heads, self.head_size)
388
+ attn_out = self._prefill_attention(
389
+ q,
390
+ k,
391
+ v,
392
+ kv_cache,
393
+ attn_metadata,
394
+ Pi,
395
+ centroids,
396
+ PiT,
397
+ layer=layer,
398
+ )
399
+ else:
400
+ # Mixed batch: decodes first (guaranteed by reorder_batch).
401
+ attn_out = torch.zeros(
402
+ N, self.num_heads, self.head_size, device=device, dtype=q.dtype
403
+ )
404
+
405
+ # --- Decode portion (first num_decodes requests) ---
406
+ # Use full-batch max_seq_len as safe upper bound (no GPU sync).
407
+ decode_meta = TurboQuantMetadata(
408
+ seq_lens=attn_metadata.seq_lens[:num_decodes],
409
+ slot_mapping=attn_metadata.slot_mapping[:num_decode_tokens],
410
+ block_table=attn_metadata.block_table[:num_decodes],
411
+ query_start_loc=attn_metadata.query_start_loc[: num_decodes + 1],
412
+ num_actual_tokens=num_decode_tokens,
413
+ max_query_len=1,
414
+ max_seq_len=attn_metadata.max_seq_len,
415
+ is_prefill=False,
416
+ )
417
+ attn_out[:num_decode_tokens] = self._decode_attention(
418
+ q[:num_decode_tokens], kv_cache, decode_meta, Pi, centroids, PiT, layer
419
+ )
420
+
421
+ # --- Prefill portion (remaining requests) ---
422
+ # CRITICAL: use prefill-specific max_seq_len so flash_attn's
423
+ # fast path (max_query_len == max_seq_len) triggers for
424
+ # first-chunk prefills. Using full-batch max_seq_len breaks
425
+ # this because decode requests inflate max_seq_len.
426
+ prefill_seq_lens = attn_metadata.seq_lens[num_decodes:]
427
+ # Use CPU-side max to avoid GPU→CPU sync from .item()
428
+ prefill_max_seq = max(attn_metadata.seq_lens[num_decodes:].tolist())
429
+ prefill_qsl = (
430
+ attn_metadata.query_start_loc[num_decodes:] - num_decode_tokens
431
+ )
432
+ prefill_meta = TurboQuantMetadata(
433
+ seq_lens=prefill_seq_lens,
434
+ slot_mapping=attn_metadata.slot_mapping[num_decode_tokens:N],
435
+ block_table=attn_metadata.block_table[num_decodes:],
436
+ query_start_loc=prefill_qsl,
437
+ num_actual_tokens=N - num_decode_tokens,
438
+ max_query_len=attn_metadata.max_query_len,
439
+ max_seq_len=prefill_max_seq,
440
+ is_prefill=True,
441
+ )
442
+ k = key[:N].view(N, self.num_kv_heads, self.head_size)
443
+ v = value[:N].view(N, self.num_kv_heads, self.head_size)
444
+ attn_out[num_decode_tokens:] = self._prefill_attention(
445
+ q[num_decode_tokens:],
446
+ k[num_decode_tokens:],
447
+ v[num_decode_tokens:],
448
+ kv_cache,
449
+ prefill_meta,
450
+ Pi,
451
+ centroids,
452
+ PiT,
453
+ layer=layer,
454
+ )
455
+
456
+ # Write into output buffer: attn_out is (N, Hq, D)
457
+ # output may be 2D (N, Hq*D) or 3D (N, Hq, D)
458
+ if output.ndim == 3:
459
+ output[:N] = attn_out.to(output.dtype)
460
+ else:
461
+ output[:N] = attn_out.reshape(N, -1).to(output.dtype)
462
+ return output
463
+
464
+ # ------------------------------------------------------------------ #
465
+ # Store K/V into combined cache (vectorized) #
466
+ # ------------------------------------------------------------------ #
467
+ def _store_kv(
468
+ self,
469
+ key: torch.Tensor, # (N, Hk, D)
470
+ value: torch.Tensor, # (N, Hk, D)
471
+ kv_cache: torch.Tensor, # (num_blocks, block_size, Hk, slot_size)
472
+ slot_mapping: torch.Tensor,
473
+ layer: Any,
474
+ ):
475
+ """Quantize + store via fused Triton kernel."""
476
+ triton_turboquant_store(
477
+ key,
478
+ value,
479
+ kv_cache,
480
+ slot_mapping,
481
+ layer._tq_PiT,
482
+ layer._tq_midpoints,
483
+ mse_bits=self.tq_config.key_mse_bits,
484
+ key_packed_size=self.tq_config.key_packed_size,
485
+ value_quant_bits=self.tq_config.effective_value_quant_bits,
486
+ key_fp8=self.tq_config.key_fp8,
487
+ )
488
+
489
+ # ------------------------------------------------------------------ #
490
+ # Prefill: SDPA on raw Q/K/V with causal mask #
491
+ # ------------------------------------------------------------------ #
492
+ def _prefill_attention(
493
+ self,
494
+ query: torch.Tensor, # (N, Hq, D)
495
+ key: torch.Tensor, # (N, Hk, D)
496
+ value: torch.Tensor, # (N, Hk, D)
497
+ kv_cache: torch.Tensor, # (num_blocks, block_size, Hk, slot_size)
498
+ attn_metadata: TurboQuantMetadata,
499
+ Pi: torch.Tensor,
500
+ centroids: torch.Tensor,
501
+ PiT: torch.Tensor | None = None,
502
+ layer: Any = None,
503
+ ) -> torch.Tensor:
504
+ N, Hq, D = query.shape
505
+
506
+ # Fast path: use flash_attn for first-chunk prefills (all K/V in batch).
507
+ # max_query_len == max_seq_len means no request has prior cached KV.
508
+ # Both are Python ints — no GPU sync.
509
+ if _HAS_FLASH_ATTN and attn_metadata.max_query_len == attn_metadata.max_seq_len:
510
+ output = torch.empty(N, Hq, D, device=query.device, dtype=query.dtype)
511
+ flash_attn_varlen_func(
512
+ q=query,
513
+ k=key,
514
+ v=value,
515
+ cu_seqlens_q=attn_metadata.query_start_loc,
516
+ cu_seqlens_k=attn_metadata.query_start_loc,
517
+ max_seqlen_q=attn_metadata.max_query_len,
518
+ max_seqlen_k=attn_metadata.max_query_len,
519
+ softmax_scale=self.scale,
520
+ causal=True,
521
+ out=output,
522
+ )
523
+ return output
524
+
525
+ # Continuation or no flash_attn: per-request attention.
526
+ # For continuation chunks (seq_len > q_len), we must attend to
527
+ # previously cached K/V from the TQ cache, not just the current
528
+ # chunk's raw K/V.
529
+ Hk = key.shape[1]
530
+ use_gqa = Hk < Hq
531
+ query_start_loc = attn_metadata.query_start_loc
532
+ num_reqs = query_start_loc.shape[0] - 1
533
+
534
+ output = torch.zeros(N, Hq, D, device=query.device, dtype=query.dtype)
535
+
536
+ # Convert to Python lists once (single CPU-GPU sync) instead of
537
+ # per-request .item() calls that each force a sync.
538
+ qsl = query_start_loc.tolist()
539
+ seq_lens_list = attn_metadata.seq_lens.tolist()
540
+
541
+ # Pre-allocate cu_seqlens for single-request flash_attn calls
542
+ # to avoid per-request host→device tensor creation.
543
+ _cu_2 = torch.zeros(2, device=query.device, dtype=torch.int32)
544
+
545
+ for i in range(num_reqs):
546
+ q_start = qsl[i]
547
+ q_end = qsl[i + 1]
548
+ q_len = q_end - q_start
549
+ if q_len <= 0:
550
+ continue
551
+
552
+ seq_len = seq_lens_list[i]
553
+ q_seq = query[q_start:q_end] # (q_len, Hq, D)
554
+ k_seq = key[q_start:q_end] # (q_len, Hk, D)
555
+ v_seq = value[q_start:q_end] # (q_len, Hk, D)
556
+
557
+ if q_len == seq_len:
558
+ # First-chunk prefill: all K/V are in the current batch.
559
+ if _HAS_FLASH_ATTN:
560
+ out = torch.empty_like(q_seq)
561
+ _cu_2[1] = q_len
562
+ cu = _cu_2
563
+ flash_attn_varlen_func(
564
+ q=q_seq,
565
+ k=k_seq,
566
+ v=v_seq,
567
+ cu_seqlens_q=cu,
568
+ cu_seqlens_k=cu,
569
+ max_seqlen_q=q_len,
570
+ max_seqlen_k=q_len,
571
+ softmax_scale=self.scale,
572
+ causal=True,
573
+ out=out,
574
+ )
575
+ else:
576
+ q_t = q_seq.transpose(0, 1).contiguous()
577
+ k_t = k_seq.transpose(0, 1).contiguous()
578
+ v_t = v_seq.transpose(0, 1).contiguous()
579
+ out = F.scaled_dot_product_attention(
580
+ q_t,
581
+ k_t,
582
+ v_t,
583
+ is_causal=True,
584
+ scale=self.scale,
585
+ enable_gqa=use_gqa,
586
+ ).transpose(0, 1)
587
+ output[q_start:q_end] = out.to(query.dtype)
588
+ else:
589
+ # Continuation chunk: tokens already stored to TQ cache
590
+ # by do_kv_cache_update. Use decode kernel directly to
591
+ # avoid O(cached_len) full-dequant per continuation.
592
+ # For large continuations, fall back to _continuation_prefill.
593
+ cached_len = seq_len - q_len
594
+ if q_len <= _CONTINUATION_DECODE_THRESHOLD:
595
+ # Fast path: treat each query as a decode request
596
+ # with incremental seq_lens for causal masking.
597
+ synth_seq_lens = torch.arange(
598
+ cached_len + 1,
599
+ seq_len + 1,
600
+ device=query.device,
601
+ dtype=attn_metadata.seq_lens.dtype,
602
+ )
603
+ synth_bt = attn_metadata.block_table[i : i + 1].expand(q_len, -1)
604
+ out = triton_turboquant_decode_attention(
605
+ query=q_seq,
606
+ kv_cache=kv_cache,
607
+ block_table=synth_bt,
608
+ seq_lens=synth_seq_lens,
609
+ Pi=Pi,
610
+ centroids=centroids,
611
+ scale=self.scale,
612
+ mse_bits=self.tq_config.key_mse_bits,
613
+ key_packed_size=self.tq_config.key_packed_size,
614
+ value_quant_bits=(self.tq_config.effective_value_quant_bits),
615
+ key_fp8=self.tq_config.key_fp8,
616
+ norm_correction=self.tq_config.norm_correction,
617
+ PiT=PiT,
618
+ )
619
+ else:
620
+ # Large continuation: dequant cached K/V and use
621
+ # flash_attn for better throughput.
622
+ out = self._continuation_prefill(
623
+ layer,
624
+ q_seq,
625
+ k_seq,
626
+ v_seq,
627
+ kv_cache,
628
+ attn_metadata.block_table[i : i + 1],
629
+ cached_len,
630
+ seq_len,
631
+ Pi,
632
+ centroids,
633
+ )
634
+ output[q_start:q_end] = out.to(query.dtype)
635
+
636
+ return output
637
+
638
+ def _continuation_prefill(
639
+ self,
640
+ layer: Any,
641
+ query: torch.Tensor, # (q_len, Hq, D)
642
+ key_chunk: torch.Tensor, # (q_len, Hk, D)
643
+ val_chunk: torch.Tensor, # (q_len, Hk, D)
644
+ kv_cache: torch.Tensor, # (num_blocks, block_size, Hk, slot_size)
645
+ block_table: torch.Tensor, # (1, max_num_blocks)
646
+ cached_len: int,
647
+ seq_len: int,
648
+ Pi: torch.Tensor,
649
+ centroids: torch.Tensor,
650
+ ) -> torch.Tensor:
651
+ """Handle continuation chunk by dequanting cached K/V from TQ cache.
652
+
653
+ Dequants previously cached K/V, concatenates with the current
654
+ chunk's raw K/V, then runs flash_attn with causal masking.
655
+ """
656
+ q_len, Hq, D = query.shape
657
+ Hk = key_chunk.shape[1]
658
+ device = query.device
659
+ block_size = kv_cache.shape[1]
660
+ BLOCK_D = triton.next_power_of_2(D)
661
+
662
+ mse_bytes = self._mse_bytes
663
+ val_data_bytes = self._val_data_bytes
664
+
665
+ # Dequant cached K/V from TQ cache
666
+ # Allocate slightly over to align to block_size for the grid.
667
+ # Reuse cached buffers to avoid per-call allocation (~16MB at 8K).
668
+ alloc_len = math.ceil(cached_len / block_size) * block_size
669
+ buf_shape = (1, Hk, alloc_len, D)
670
+ k_buf = getattr(layer, "_tq_k_dequant_buf", None)
671
+ if k_buf is None or k_buf.shape[2] < alloc_len:
672
+ k_buf = torch.empty(buf_shape, dtype=torch.float16, device=device)
673
+ v_buf = torch.empty(buf_shape, dtype=torch.float16, device=device)
674
+ layer._tq_k_dequant_buf = k_buf
675
+ layer._tq_v_dequant_buf = v_buf
676
+ else:
677
+ v_buf = layer._tq_v_dequant_buf
678
+ k_cached = k_buf[:, :, :alloc_len, :].zero_()
679
+ v_cached = v_buf[:, :, :alloc_len, :].zero_()
680
+
681
+ grid = (alloc_len, 1 * Hk)
682
+ _tq_full_dequant_kv[grid](
683
+ kv_cache,
684
+ block_table,
685
+ centroids,
686
+ k_cached,
687
+ v_cached,
688
+ k_cached.stride(0),
689
+ k_cached.stride(1),
690
+ k_cached.stride(2),
691
+ v_cached.stride(0),
692
+ v_cached.stride(1),
693
+ v_cached.stride(2),
694
+ kv_cache.stride(0),
695
+ kv_cache.stride(1),
696
+ kv_cache.stride(2),
697
+ block_table.stride(0),
698
+ HEAD_DIM=D,
699
+ BLOCK_SIZE=block_size,
700
+ NUM_KV_HEADS=Hk,
701
+ MSE_BYTES=mse_bytes,
702
+ KPS=self.tq_config.key_packed_size,
703
+ VQB=self.tq_config.effective_value_quant_bits,
704
+ VAL_DATA_BYTES=val_data_bytes,
705
+ MSE_BITS=self.tq_config.key_mse_bits,
706
+ KEY_FP8=1 if self.tq_config.key_fp8 else 0,
707
+ BLOCK_D=BLOCK_D,
708
+ NORM_CORRECTION=1 if self.tq_config.norm_correction else 0,
709
+ FP8_E4B15=_use_fp8_e4b15(device.index or 0),
710
+ num_warps=4,
711
+ )
712
+
713
+ # Inverse-rotate MSE keys back to original space
714
+ if not self.tq_config.key_fp8:
715
+ k_flat = k_cached[0, :, :cached_len, :].reshape(-1, D).float()
716
+ k_flat = k_flat @ Pi
717
+ k_cached_trim = (
718
+ k_flat.to(torch.float16).reshape(Hk, cached_len, D).transpose(0, 1)
719
+ ) # (cached_len, Hk, D)
720
+ else:
721
+ k_cached_trim = (
722
+ k_cached[0, :, :cached_len, :].transpose(0, 1).contiguous()
723
+ ) # (cached_len, Hk, D)
724
+
725
+ v_cached_trim = (
726
+ v_cached[0, :, :cached_len, :].transpose(0, 1).contiguous()
727
+ ) # (cached_len, Hk, D)
728
+
729
+ # Concatenate cached + current chunk K/V (match query dtype)
730
+ qdtype = query.dtype
731
+ k_full = torch.cat([k_cached_trim.to(qdtype), key_chunk], dim=0)
732
+ v_full = torch.cat([v_cached_trim.to(qdtype), val_chunk], dim=0)
733
+
734
+ # Attention: q_len queries attending to seq_len K/V with causal mask
735
+ if _HAS_FLASH_ATTN:
736
+ output = torch.empty(q_len, Hq, D, device=device, dtype=query.dtype)
737
+ cu_seqlens_q = torch.tensor([0, q_len], device=device, dtype=torch.int32)
738
+ cu_seqlens_k = torch.tensor([0, seq_len], device=device, dtype=torch.int32)
739
+ flash_attn_varlen_func(
740
+ q=query,
741
+ k=k_full,
742
+ v=v_full,
743
+ cu_seqlens_q=cu_seqlens_q,
744
+ cu_seqlens_k=cu_seqlens_k,
745
+ max_seqlen_q=q_len,
746
+ max_seqlen_k=seq_len,
747
+ softmax_scale=self.scale,
748
+ causal=True,
749
+ out=output,
750
+ )
751
+ return output
752
+ else:
753
+ # SDPA fallback: expand KV for GQA, build causal mask
754
+ q_t = query.transpose(0, 1).unsqueeze(0) # (1, Hq, q_len, D)
755
+ k_t = k_full.transpose(0, 1).unsqueeze(0) # (1, Hk, seq_len, D)
756
+ v_t = v_full.transpose(0, 1).unsqueeze(0) # (1, Hk, seq_len, D)
757
+ # Build causal mask: query position p can attend to K position j
758
+ # where j <= cached_len + p (p is 0-indexed within chunk)
759
+ q_pos = torch.arange(q_len, device=device).unsqueeze(1) + cached_len
760
+ k_pos = torch.arange(seq_len, device=device).unsqueeze(0)
761
+ mask = k_pos <= q_pos # (q_len, seq_len)
762
+ out = F.scaled_dot_product_attention(
763
+ q_t,
764
+ k_t,
765
+ v_t,
766
+ attn_mask=mask,
767
+ scale=self.scale,
768
+ enable_gqa=(Hk < Hq),
769
+ ) # (1, Hq, q_len, D)
770
+ return out[0].transpose(0, 1) # (q_len, Hq, D)
771
+
772
+ # ------------------------------------------------------------------ #
773
+ # Decode: Triton TQ decode attention #
774
+ # ------------------------------------------------------------------ #
775
+ def _decode_attention(
776
+ self,
777
+ query: torch.Tensor, # (B, Hq, D)
778
+ kv_cache: torch.Tensor, # (num_blocks, block_size, Hk, slot_size)
779
+ attn_metadata: TurboQuantMetadata,
780
+ Pi: torch.Tensor,
781
+ centroids: torch.Tensor,
782
+ PiT: torch.Tensor | None = None,
783
+ layer: torch.nn.Module | None = None,
784
+ ) -> torch.Tensor:
785
+ # Grab cached decode buffers from the layer (lazily allocated).
786
+ mid_o_buf = output_buf = lse_buf = None
787
+ if layer is not None:
788
+ mid_o_buf = getattr(layer, "_tq_mid_o_buf", None)
789
+ output_buf = getattr(layer, "_tq_output_buf", None)
790
+ lse_buf = getattr(layer, "_tq_lse_buf", None)
791
+
792
+ result = triton_turboquant_decode_attention(
793
+ query=query,
794
+ kv_cache=kv_cache,
795
+ block_table=attn_metadata.block_table,
796
+ seq_lens=attn_metadata.seq_lens,
797
+ Pi=Pi,
798
+ centroids=centroids,
799
+ scale=self.scale,
800
+ mse_bits=self.tq_config.key_mse_bits,
801
+ key_packed_size=self.tq_config.key_packed_size,
802
+ value_quant_bits=self.tq_config.effective_value_quant_bits,
803
+ key_fp8=self.tq_config.key_fp8,
804
+ norm_correction=self.tq_config.norm_correction,
805
+ PiT=PiT,
806
+ mid_o_buf=mid_o_buf,
807
+ output_buf=output_buf,
808
+ lse_buf=lse_buf,
809
+ buf_holder=layer,
810
+ max_num_kv_splits=self.max_num_kv_splits,
811
+ )
812
+ return result
vllm/v1/attention/ops/triton_turboquant_decode.py
new
+617
@@ -0,0 +1,617 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Triton fused TurboQuant decode attention.
4
+
5
+Decode path: Triton stage1 (split-KV tiled attention scoring + value
6
+accumulation) + stage2 (log-sum-exp reduction across splits).
7
+
8
+Supports FP8 (E4M3) keys, 3-bit and 4-bit uniform quantized values.
9
+"""
10
+
11
+import math
12
+from typing import Any
13
+
14
+import torch
15
+
16
+from vllm.triton_utils import tl, triton
17
+from vllm.v1.attention.ops.triton_decode_attention import (
18
+ _fwd_kernel_stage2,
19
+)
20
+
21
+_FP8_E4B15: dict[int, int] = {}
22
+
23
+
24
+def _use_fp8_e4b15(device: int = 0) -> int:
25
+ """Return 1 if device needs fp8e4b15 (Ampere/Ada, SM < 8.9), else 0."""
26
+ if device not in _FP8_E4B15:
27
+ cap = torch.cuda.get_device_capability(device)
28
+ _FP8_E4B15[device] = 1 if cap < (8, 9) else 0
29
+ return _FP8_E4B15[device]
30
+
31
+
32
+# ---------------------------------------------------------------------------
33
+# Stage 1: Fused TQ score + value accumulation (BLOCK_KV tiled)
34
+# ---------------------------------------------------------------------------
35
+
36
+
37
+@triton.jit
38
+def _tq_decode_stage1(
39
+ # Precomputed query projection
40
+ Q_rot_ptr, # [B, Hq, D] float32
41
+ # Compressed KV cache (combined K+V)
42
+ KV_cache_ptr, # [num_blocks, block_size, Hk, padded_slot] uint8
43
+ # Block table and sequence info
44
+ Block_table_ptr, # [B, max_num_blocks] int32
45
+ Seq_lens_ptr, # [B] int32
46
+ # TQ parameters
47
+ Centroids_ptr, # [n_centroids] float32
48
+ # Output (intermediate for stage2)
49
+ Mid_o_ptr, # [B, Hq, NUM_KV_SPLITS, D+1] float32
50
+ # Strides
51
+ stride_qb,
52
+ stride_qh, # Q strides: [B, Hq, D]
53
+ stride_cache_block,
54
+ stride_cache_pos,
55
+ stride_cache_head, # KV cache
56
+ stride_bt_b, # block_table stride per batch
57
+ stride_mid_b,
58
+ stride_mid_h,
59
+ stride_mid_s, # mid_o strides
60
+ # Constexpr dims
61
+ NUM_KV_HEADS: tl.constexpr,
62
+ HEAD_DIM: tl.constexpr,
63
+ BLOCK_SIZE: tl.constexpr, # KV cache block_size (pages)
64
+ NUM_KV_SPLITS: tl.constexpr,
65
+ KV_GROUP_SIZE: tl.constexpr, # Hq // Hk
66
+ # TQ layout constants
67
+ MSE_BITS: tl.constexpr, # 3 or 4
68
+ MSE_BYTES: tl.constexpr, # ceil(D * mse_bits / 8)
69
+ KPS: tl.constexpr, # key_packed_size
70
+ VQB: tl.constexpr, # value_quant_bits (4 or 8=FP8)
71
+ VAL_DATA_BYTES: tl.constexpr, # ceil(D * vqb / 8) or D for FP8
72
+ # Score constants
73
+ ATTN_SCALE: tl.constexpr, # 1/sqrt(D)
74
+ # Block tile sizes
75
+ BLOCK_D: tl.constexpr, # next_power_of_2(HEAD_DIM)
76
+ BLOCK_KV: tl.constexpr, # tokens per tile (16)
77
+ KEY_FP8: tl.constexpr, # 1 if K is stored as FP8
78
+ NORM_CORRECTION: tl.constexpr = 0, # 1 = re-normalize centroids
79
+ FP8_E4B15: tl.constexpr = 0, # 1 = use e4b15 (Ampere/Ada), 0 = e4nv (Hopper+)
80
+):
81
+ bid = tl.program_id(0) # batch index
82
+ hid = tl.program_id(1) # q_head index
83
+ sid = tl.program_id(2) # kv_split index
84
+
85
+ kv_head = hid // KV_GROUP_SIZE
86
+
87
+ # Sequence length for this batch
88
+ seq_len = tl.load(Seq_lens_ptr + bid)
89
+
90
+ # KV split range
91
+ split_len = tl.cdiv(seq_len, NUM_KV_SPLITS)
92
+ split_start = split_len * sid
93
+ split_end = tl.minimum(split_start + split_len, seq_len)
94
+
95
+ if split_start >= split_end:
96
+ return
97
+
98
+ # Dimension offsets
99
+ d_offs = tl.arange(0, BLOCK_D)
100
+ d_mask = d_offs < HEAD_DIM
101
+ kv_range = tl.arange(0, BLOCK_KV)
102
+
103
+ # Load query vector: q_rot — [BLOCK_D] float32
104
+ q_base = bid * stride_qb + hid * stride_qh
105
+ q_rot = tl.load(Q_rot_ptr + q_base + d_offs, mask=d_mask, other=0.0).to(tl.float32)
106
+
107
+ # Precompute byte/bit index vectors for MSE gather loads
108
+ if not KEY_FP8:
109
+ mse_bit_off = d_offs * MSE_BITS
110
+ mse_byte_idx = mse_bit_off // 8
111
+ mse_bit_shift = mse_bit_off % 8
112
+ mse_mask = (1 << MSE_BITS) - 1
113
+
114
+ # Precompute value bit/byte index vectors (loop-invariant)
115
+ if VQB == 3:
116
+ val_bit_off = d_offs * 3
117
+ val_byte_idx = val_bit_off // 8
118
+ val_bit_shift = val_bit_off % 8
119
+
120
+ # Online softmax accumulators
121
+ m_prev = -float("inf")
122
+ l_prev = 0.0
123
+ acc = tl.zeros([BLOCK_D], dtype=tl.float32)
124
+
125
+ bt_base = bid * stride_bt_b
126
+
127
+ # ================================================================
128
+ # TILED LOOP: process BLOCK_KV tokens per iteration
129
+ # ================================================================
130
+ for start_n in range(split_start, split_end, BLOCK_KV):
131
+ kv_offs = start_n + kv_range
132
+ kv_mask = kv_offs < split_end
133
+
134
+ page_idx = kv_offs // BLOCK_SIZE
135
+ page_off = kv_offs % BLOCK_SIZE
136
+ block_nums = tl.load(
137
+ Block_table_ptr + bt_base + page_idx,
138
+ mask=kv_mask,
139
+ other=0,
140
+ )
141
+
142
+ slot_bases = (
143
+ block_nums * stride_cache_block
144
+ + page_off * stride_cache_pos
145
+ + kv_head * stride_cache_head
146
+ )
147
+
148
+ # ============================================================
149
+ # COMPUTE ATTENTION SCORES: [BLOCK_KV]
150
+ # ============================================================
151
+ if KEY_FP8:
152
+ k_addrs = slot_bases[:, None] + d_offs[None, :]
153
+ k_raw = tl.load(
154
+ KV_cache_ptr + k_addrs,
155
+ mask=kv_mask[:, None] & d_mask[None, :],
156
+ other=0,
157
+ )
158
+ if FP8_E4B15:
159
+ k_float = k_raw.to(tl.float8e4b15, bitcast=True).to(tl.float32)
160
+ else:
161
+ k_float = k_raw.to(tl.float8e4nv, bitcast=True).to(tl.float32)
162
+ scores = (
163
+ tl.sum(
164
+ tl.where(d_mask[None, :], q_rot[None, :] * k_float, 0.0),
165
+ axis=1,
166
+ )
167
+ * ATTN_SCALE
168
+ )
169
+ scores = tl.where(kv_mask, scores, -float("inf"))
170
+ else:
171
+ # MSE unpack + norms
172
+ mse_addrs0 = slot_bases[:, None] + mse_byte_idx[None, :]
173
+ mse_raw0 = tl.load(
174
+ KV_cache_ptr + mse_addrs0,
175
+ mask=kv_mask[:, None] & d_mask[None, :],
176
+ other=0,
177
+ ).to(tl.int32)
178
+ mse_raw1 = tl.load(
179
+ KV_cache_ptr + mse_addrs0 + 1,
180
+ mask=kv_mask[:, None] & d_mask[None, :],
181
+ other=0,
182
+ ).to(tl.int32)
183
+ raw16 = mse_raw0 | (mse_raw1 << 8)
184
+ mse_idx = (raw16 >> mse_bit_shift[None, :]) & mse_mask
185
+
186
+ # Centroid gather + dot product
187
+ c_vals = tl.load(
188
+ Centroids_ptr + mse_idx,
189
+ mask=kv_mask[:, None] & d_mask[None, :],
190
+ other=0.0,
191
+ )
192
+
193
+ # Norm correction: re-normalize centroid vector to unit norm
194
+ if NORM_CORRECTION:
195
+ c_norm_sq = tl.sum(
196
+ tl.where(d_mask[None, :], c_vals * c_vals, 0.0),
197
+ axis=1,
198
+ )
199
+ c_inv_norm = 1.0 / tl.sqrt(c_norm_sq + 1e-16)
200
+ c_vals = c_vals * c_inv_norm[:, None]
201
+
202
+ term1 = tl.sum(
203
+ tl.where(d_mask[None, :], q_rot[None, :] * c_vals, 0.0),
204
+ axis=1,
205
+ )
206
+
207
+ # Load norms (fp16 -> fp32): norms are at MSE_BYTES offset
208
+ norm_bases = slot_bases + MSE_BYTES
209
+ n_lo = tl.load(KV_cache_ptr + norm_bases, mask=kv_mask, other=0).to(
210
+ tl.uint16
211
+ )
212
+ n_hi = tl.load(KV_cache_ptr + norm_bases + 1, mask=kv_mask, other=0).to(
213
+ tl.uint16
214
+ )
215
+ vec_norms = (n_lo | (n_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32)
216
+
217
+ scores = vec_norms * term1 * ATTN_SCALE
218
+ scores = tl.where(kv_mask, scores, -float("inf"))
219
+
220
+ # ============================================================
221
+ # ONLINE SOFTMAX UPDATE (block-level)
222
+ # ============================================================
223
+ n_e_max = tl.maximum(tl.max(scores, 0), m_prev)
224
+ re_scale = tl.exp(m_prev - n_e_max)
225
+ p = tl.exp(scores - n_e_max)
226
+
227
+ # ============================================================
228
+ # VALUE LOAD + DEQUANTIZE: [BLOCK_KV, BLOCK_D]
229
+ # ============================================================
230
+ val_bases = slot_bases + KPS
231
+
232
+ if VQB == 3:
233
+ val_addrs0 = val_bases[:, None] + val_byte_idx[None, :]
234
+ val_raw0 = tl.load(
235
+ KV_cache_ptr + val_addrs0,
236
+ mask=kv_mask[:, None] & d_mask[None, :],
237
+ other=0,
238
+ ).to(tl.int32)
239
+ val_raw1 = tl.load(
240
+ KV_cache_ptr + val_addrs0 + 1,
241
+ mask=kv_mask[:, None] & d_mask[None, :],
242
+ other=0,
243
+ ).to(tl.int32)
244
+ raw16 = val_raw0 | (val_raw1 << 8)
245
+ v_idx = ((raw16 >> val_bit_shift[None, :]) & 0x7).to(tl.float32)
246
+
247
+ sc_bases = val_bases + VAL_DATA_BYTES
248
+ sc_lo = tl.load(KV_cache_ptr + sc_bases, mask=kv_mask, other=0).to(
249
+ tl.uint16
250
+ )
251
+ sc_hi = tl.load(KV_cache_ptr + sc_bases + 1, mask=kv_mask, other=0).to(
252
+ tl.uint16
253
+ )
254
+ v_scales = (
255
+ (sc_lo | (sc_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32)
256
+ )
257
+ zr_lo = tl.load(KV_cache_ptr + sc_bases + 2, mask=kv_mask, other=0).to(
258
+ tl.uint16
259
+ )
260
+ zr_hi = tl.load(KV_cache_ptr + sc_bases + 3, mask=kv_mask, other=0).to(
261
+ tl.uint16
262
+ )
263
+ v_zeros = (zr_lo | (zr_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32)
264
+ values = v_idx * v_scales[:, None] + v_zeros[:, None]
265
+ else: # VQB == 4
266
+ vb_idx = d_offs // 2
267
+ vb_shift = (d_offs % 2) * 4
268
+ val_addrs = val_bases[:, None] + vb_idx[None, :]
269
+ val_raw = tl.load(
270
+ KV_cache_ptr + val_addrs,
271
+ mask=kv_mask[:, None] & d_mask[None, :],
272
+ other=0,
273
+ ).to(tl.int32)
274
+ v_idx = ((val_raw >> vb_shift[None, :]) & 0xF).to(tl.float32)
275
+
276
+ sc_bases = val_bases + VAL_DATA_BYTES
277
+ sc_lo = tl.load(KV_cache_ptr + sc_bases, mask=kv_mask, other=0).to(
278
+ tl.uint16
279
+ )
280
+ sc_hi = tl.load(KV_cache_ptr + sc_bases + 1, mask=kv_mask, other=0).to(
281
+ tl.uint16
282
+ )
283
+ v_scales = (
284
+ (sc_lo | (sc_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32)
285
+ )
286
+ zr_lo = tl.load(KV_cache_ptr + sc_bases + 2, mask=kv_mask, other=0).to(
287
+ tl.uint16
288
+ )
289
+ zr_hi = tl.load(KV_cache_ptr + sc_bases + 3, mask=kv_mask, other=0).to(
290
+ tl.uint16
291
+ )
292
+ v_zeros = (zr_lo | (zr_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32)
293
+ values = v_idx * v_scales[:, None] + v_zeros[:, None]
294
+
295
+ # ============================================================
296
+ # WEIGHTED VALUE ACCUMULATION
297
+ # ============================================================
298
+ acc = acc * re_scale + tl.sum(p[:, None] * values, 0)
299
+ l_prev = l_prev * re_scale + tl.sum(p, 0)
300
+ m_prev = n_e_max
301
+
302
+ # Store partial result
303
+ out_base = bid * stride_mid_b + hid * stride_mid_h + sid * stride_mid_s
304
+ safe_l = tl.where(l_prev > 0.0, l_prev, 1.0)
305
+ tl.store(Mid_o_ptr + out_base + d_offs, acc / safe_l, mask=d_mask)
306
+ lse = m_prev + tl.log(safe_l)
307
+ tl.store(Mid_o_ptr + out_base + HEAD_DIM, lse)
308
+
309
+
310
+# ---------------------------------------------------------------------------
311
+# Pre-dequant kernel: Bulk dequant K (MSE+norms) and V to fp16
312
+# ---------------------------------------------------------------------------
313
+
314
+
315
+@triton.jit
316
+def _tq_full_dequant_kv(
317
+ KV_cache_ptr,
318
+ Block_table_ptr,
319
+ Centroids_ptr,
320
+ K_out_ptr, # [B, Hk, max_seq, D] float16
321
+ V_out_ptr, # [B, Hk, max_seq, D] float16
322
+ stride_ko_b,
323
+ stride_ko_h,
324
+ stride_ko_s,
325
+ stride_vo_b,
326
+ stride_vo_h,
327
+ stride_vo_s,
328
+ stride_cache_block,
329
+ stride_cache_pos,
330
+ stride_cache_head,
331
+ stride_bt_b,
332
+ HEAD_DIM: tl.constexpr,
333
+ BLOCK_SIZE: tl.constexpr,
334
+ NUM_KV_HEADS: tl.constexpr,
335
+ MSE_BYTES: tl.constexpr,
336
+ KPS: tl.constexpr,
337
+ VQB: tl.constexpr,
338
+ VAL_DATA_BYTES: tl.constexpr,
339
+ MSE_BITS: tl.constexpr,
340
+ KEY_FP8: tl.constexpr,
341
+ BLOCK_D: tl.constexpr,
342
+ NORM_CORRECTION: tl.constexpr = 0,
343
+ FP8_E4B15: tl.constexpr = 0, # 1 = use e4b15 (Ampere/Ada), 0 = e4nv (Hopper+)
344
+):
345
+ """Full dequant: reconstruct K (MSE centroids * norm or FP8) and V to fp16."""
346
+ pos = tl.program_id(0)
347
+ bh = tl.program_id(1)
348
+ bid = bh // NUM_KV_HEADS
349
+ hid = bh % NUM_KV_HEADS
350
+
351
+ page_idx = pos // BLOCK_SIZE
352
+ page_off = pos % BLOCK_SIZE
353
+ block_num = tl.load(Block_table_ptr + bid * stride_bt_b + page_idx)
354
+ slot_base = (
355
+ block_num * stride_cache_block
356
+ + page_off * stride_cache_pos
357
+ + hid * stride_cache_head
358
+ )
359
+
360
+ d_offs = tl.arange(0, BLOCK_D)
361
+ d_mask = d_offs < HEAD_DIM
362
+
363
+ # === K dequant ===
364
+ ko_base = bid * stride_ko_b + hid * stride_ko_h + pos * stride_ko_s
365
+ if KEY_FP8:
366
+ k_raw = tl.load(KV_cache_ptr + slot_base + d_offs, mask=d_mask, other=0)
367
+ if FP8_E4B15:
368
+ k_recon = k_raw.to(tl.float8e4b15, bitcast=True).to(tl.float32)
369
+ else:
370
+ k_recon = k_raw.to(tl.float8e4nv, bitcast=True).to(tl.float32)
371
+ tl.store(K_out_ptr + ko_base + d_offs, k_recon.to(tl.float16), mask=d_mask)
372
+ else:
373
+ # MSE unpack (3-bit or 4-bit) + norms
374
+ mse_bit_off = d_offs * MSE_BITS
375
+ mse_byte_idx = mse_bit_off // 8
376
+ mse_bit_shift = mse_bit_off % 8
377
+ mse_umask = (1 << MSE_BITS) - 1
378
+
379
+ mse_raw0 = tl.load(
380
+ KV_cache_ptr + slot_base + mse_byte_idx, mask=d_mask, other=0
381
+ ).to(tl.int32)
382
+ mse_raw1 = tl.load(
383
+ KV_cache_ptr + slot_base + mse_byte_idx + 1, mask=d_mask, other=0
384
+ ).to(tl.int32)
385
+ raw16_key = mse_raw0 | (mse_raw1 << 8)
386
+ mse_idx = (raw16_key >> mse_bit_shift) & mse_umask
387
+
388
+ k_mse = tl.load(Centroids_ptr + mse_idx, mask=d_mask, other=0.0)
389
+
390
+ # Norm correction: re-normalize centroid vector to unit norm
391
+ if NORM_CORRECTION:
392
+ c_norm_sq = tl.sum(tl.where(d_mask, k_mse * k_mse, 0.0), axis=0)
393
+ c_inv_norm = 1.0 / tl.sqrt(c_norm_sq + 1e-16)
394
+ k_mse = k_mse * c_inv_norm
395
+
396
+ # Norms at MSE_BYTES offset (no QJL bytes)
397
+ norm_base = slot_base + MSE_BYTES
398
+ n_lo = tl.load(KV_cache_ptr + norm_base).to(tl.uint16)
399
+ n_hi = tl.load(KV_cache_ptr + norm_base + 1).to(tl.uint16)
400
+ vec_norm = (n_lo | (n_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32)
401
+
402
+ k_recon = vec_norm * k_mse
403
+ tl.store(K_out_ptr + ko_base + d_offs, k_recon.to(tl.float16), mask=d_mask)
404
+
405
+ # === V dequant ===
406
+ val_base = slot_base + KPS
407
+ if VQB == 4:
408
+ vb_idx = d_offs // 2
409
+ vb_shift = (d_offs % 2) * 4
410
+ val_raw = tl.load(KV_cache_ptr + val_base + vb_idx, mask=d_mask, other=0).to(
411
+ tl.int32
412
+ )
413
+ v_idx = ((val_raw >> vb_shift) & 0xF).to(tl.float32)
414
+
415
+ sc_base = val_base + VAL_DATA_BYTES
416
+ sc_lo = tl.load(KV_cache_ptr + sc_base).to(tl.uint16)
417
+ sc_hi = tl.load(KV_cache_ptr + sc_base + 1).to(tl.uint16)
418
+ v_scale = (sc_lo | (sc_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32)
419
+ zr_lo = tl.load(KV_cache_ptr + sc_base + 2).to(tl.uint16)
420
+ zr_hi = tl.load(KV_cache_ptr + sc_base + 3).to(tl.uint16)
421
+ v_zero = (zr_lo | (zr_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32)
422
+ v_vals = v_idx * v_scale + v_zero
423
+ elif VQB == 3:
424
+ # 3-bit value unpack: 8 values per 3 bytes
425
+ val_bit_off = d_offs * 3
426
+ val_byte_idx = val_bit_off // 8
427
+ val_bit_shift = val_bit_off % 8
428
+ val_raw0 = tl.load(
429
+ KV_cache_ptr + val_base + val_byte_idx, mask=d_mask, other=0
430
+ ).to(tl.int32)
431
+ val_raw1 = tl.load(
432
+ KV_cache_ptr + val_base + val_byte_idx + 1, mask=d_mask, other=0
433
+ ).to(tl.int32)
434
+ raw16_val = val_raw0 | (val_raw1 << 8)
435
+ v_idx = ((raw16_val >> val_bit_shift) & 0x7).to(tl.float32)
436
+
437
+ sc_base = val_base + VAL_DATA_BYTES
438
+ sc_lo = tl.load(KV_cache_ptr + sc_base).to(tl.uint16)
439
+ sc_hi = tl.load(KV_cache_ptr + sc_base + 1).to(tl.uint16)
440
+ v_scale = (sc_lo | (sc_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32)
441
+ zr_lo = tl.load(KV_cache_ptr + sc_base + 2).to(tl.uint16)
442
+ zr_hi = tl.load(KV_cache_ptr + sc_base + 3).to(tl.uint16)
443
+ v_zero = (zr_lo | (zr_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32)
444
+ v_vals = v_idx * v_scale + v_zero
445
+ else:
446
+ v_vals = tl.zeros([BLOCK_D], dtype=tl.float32)
447
+
448
+ vo_base = bid * stride_vo_b + hid * stride_vo_h + pos * stride_vo_s
449
+ tl.store(V_out_ptr + vo_base + d_offs, v_vals.to(tl.float16), mask=d_mask)
450
+
451
+
452
+# ---------------------------------------------------------------------------
453
+# Stage 2: Reuse from triton_decode_attention.py
454
+# ---------------------------------------------------------------------------
455
+
456
+# ---------------------------------------------------------------------------
457
+# Launcher — cached constants + fused GEMM
458
+# ---------------------------------------------------------------------------
459
+
460
+_layout_cache: dict = {}
461
+
462
+
463
+def _get_layout(D, mse_bits, value_quant_bits, key_packed_size):
464
+ """Get cached layout constants."""
465
+ key = (D, mse_bits, value_quant_bits, key_packed_size)
466
+ cfg = _layout_cache.get(key)
467
+ if cfg is None:
468
+ val_data_bytes = math.ceil(D * value_quant_bits / 8)
469
+ cfg = {
470
+ "mse_bytes": math.ceil(D * mse_bits / 8),
471
+ "val_data_bytes": val_data_bytes,
472
+ "mse_bits": mse_bits,
473
+ "n_centroids": 2**mse_bits,
474
+ "BLOCK_D": triton.next_power_of_2(D),
475
+ }
476
+ _layout_cache[key] = cfg
477
+ return cfg
478
+
479
+
480
+def triton_turboquant_decode_attention(
481
+ query: torch.Tensor, # [B, Hq, D] — original query
482
+ kv_cache: torch.Tensor, # [num_blocks, block_size, Hk, padded_slot] uint8
483
+ block_table: torch.Tensor, # [B, max_num_blocks] int32
484
+ seq_lens: torch.Tensor, # [B] int32
485
+ Pi: torch.Tensor, # [D, D] float32
486
+ centroids: torch.Tensor, # [n_centroids] float32
487
+ scale: float,
488
+ mse_bits: int,
489
+ key_packed_size: int,
490
+ value_quant_bits: int,
491
+ key_fp8: bool = False,
492
+ norm_correction: bool = False,
493
+ PiT: torch.Tensor | None = None, # [D, D] pre-computed Pi.T contiguous
494
+ # Pre-allocated buffers (optional, avoids per-call allocation)
495
+ mid_o_buf: torch.Tensor | None = None,
496
+ output_buf: torch.Tensor | None = None,
497
+ lse_buf: torch.Tensor | None = None,
498
+ buf_holder: Any = None,
499
+ max_num_kv_splits: int = 32, # fixed split count (must be constant for cudagraph)
500
+) -> torch.Tensor:
501
+ """Launch fused TQ decode attention (Triton stage1 + stage2).
502
+
503
+ Returns: output tensor [B, Hq, D] in query's dtype.
504
+ """
505
+ B, Hq, D = query.shape
506
+ Hk = kv_cache.shape[2]
507
+ block_size = kv_cache.shape[1]
508
+ kv_group_size = Hq // Hk
509
+ device = query.device
510
+
511
+ cfg = _get_layout(D, mse_bits, value_quant_bits, key_packed_size)
512
+
513
+ # Compute q_rot = q @ Pi.T (rotated query for MSE key scoring)
514
+ # FP8 path: pass query directly (float16); kernel casts inline.
515
+ # MSE path: still needs external GEMM (cuBLAS), so q_rot is float32.
516
+ if key_fp8:
517
+ q_rot = query.contiguous()
518
+ else:
519
+ q_float = query.float()
520
+ if PiT is None:
521
+ PiT = Pi.T.contiguous()
522
+ q_rot = (q_float @ PiT).contiguous()
523
+
524
+ NUM_KV_SPLITS = max_num_kv_splits
525
+
526
+ if (
527
+ mid_o_buf is not None
528
+ and mid_o_buf.shape[0] >= B
529
+ and mid_o_buf.shape[2] >= NUM_KV_SPLITS
530
+ ):
531
+ mid_o = mid_o_buf[:B, :Hq, :NUM_KV_SPLITS, :]
532
+ else:
533
+ mid_o = torch.empty(
534
+ B,
535
+ Hq,
536
+ NUM_KV_SPLITS,
537
+ D + 1,
538
+ dtype=torch.float32,
539
+ device=device,
540
+ )
541
+ if buf_holder is not None:
542
+ buf_holder._tq_mid_o_buf = mid_o
543
+
544
+ # Stage 1: split-KV tiled attention scoring + value accumulation
545
+ fp8_e4b15 = _use_fp8_e4b15(device.index or 0)
546
+ BLOCK_KV = 4
547
+ grid = (B, Hq, NUM_KV_SPLITS)
548
+ _tq_decode_stage1[grid](
549
+ q_rot,
550
+ kv_cache,
551
+ block_table,
552
+ seq_lens,
553
+ centroids,
554
+ mid_o,
555
+ q_rot.stride(0),
556
+ q_rot.stride(1),
557
+ kv_cache.stride(0),
558
+ kv_cache.stride(1),
559
+ kv_cache.stride(2),
560
+ block_table.stride(0),
561
+ mid_o.stride(0),
562
+ mid_o.stride(1),
563
+ mid_o.stride(2),
564
+ NUM_KV_HEADS=Hk,
565
+ HEAD_DIM=D,
566
+ BLOCK_SIZE=block_size,
567
+ NUM_KV_SPLITS=NUM_KV_SPLITS,
568
+ KV_GROUP_SIZE=kv_group_size,
569
+ MSE_BITS=mse_bits,
570
+ MSE_BYTES=cfg["mse_bytes"],
571
+ KPS=key_packed_size,
572
+ VQB=value_quant_bits,
573
+ VAL_DATA_BYTES=cfg["val_data_bytes"],
574
+ ATTN_SCALE=scale,
575
+ BLOCK_D=cfg["BLOCK_D"],
576
+ BLOCK_KV=BLOCK_KV,
577
+ KEY_FP8=1 if key_fp8 else 0,
578
+ NORM_CORRECTION=1 if norm_correction else 0,
579
+ FP8_E4B15=fp8_e4b15,
580
+ num_warps=1,
581
+ num_stages=1,
582
+ )
583
+
584
+ # Stage 2: Reduce across KV splits
585
+ if output_buf is not None and output_buf.shape[0] >= B:
586
+ output = output_buf[:B, :Hq, :D]
587
+ else:
588
+ output = torch.empty(B, Hq, D, dtype=torch.float32, device=device)
589
+ if buf_holder is not None:
590
+ buf_holder._tq_output_buf = output
591
+ if lse_buf is not None and lse_buf.shape[0] >= B:
592
+ lse = lse_buf[:B, :Hq]
593
+ else:
594
+ lse = torch.empty(B, Hq, dtype=torch.float32, device=device)
595
+ if buf_holder is not None:
596
+ buf_holder._tq_lse_buf = lse
597
+
598
+ grid2 = (B, Hq)
599
+ _fwd_kernel_stage2[grid2](
600
+ mid_o,
601
+ output,
602
+ lse,
603
+ seq_lens,
604
+ mid_o.stride(0),
605
+ mid_o.stride(1),
606
+ mid_o.stride(2),
607
+ output.stride(0),
608
+ output.stride(1),
609
+ lse.stride(0),
610
+ NUM_KV_SPLITS=NUM_KV_SPLITS,
611
+ BLOCK_DV=cfg["BLOCK_D"],
612
+ Lv=D,
613
+ num_warps=4,
614
+ num_stages=2,
615
+ )
616
+
617
+ return output.to(query.dtype)
vllm/v1/attention/ops/triton_turboquant_store.py
new
+441
@@ -0,0 +1,441 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Fused Triton kernels for TurboQuant KV store.
4
+
5
+Two kernels:
6
+1. _tq_fused_store_fp8: FP8 key scatter + value uniform quantization.
7
+2. _tq_fused_store_mse: Fused binary-search bucketize + MSE index
8
+ packing + value quantization.
9
+
10
+The launcher `triton_turboquant_store` selects the appropriate kernel.
11
+"""
12
+
13
+import math
14
+
15
+import torch
16
+
17
+from vllm.triton_utils import tl, triton
18
+from vllm.v1.attention.ops.triton_turboquant_decode import _use_fp8_e4b15
19
+
20
+# ═══════════════════════════════════════════════════════════════════════
21
+# Shared: value uniform quantization + pack + scale/zero store
22
+# ═══════════════════════════════════════════════════════════════════════
23
+
24
+
25
+@triton.jit
26
+def _store_quantized_value(
27
+ Value_ptr,
28
+ KV_cache_ptr,
29
+ base, # pid * D offset into Value_ptr
30
+ slot_base, # byte offset into KV_cache_ptr for this slot+head
31
+ d_offs, # tl.arange(0, BLOCK_D)
32
+ d_mask, # d_offs < D
33
+ D: tl.constexpr,
34
+ KPS: tl.constexpr,
35
+ VQB: tl.constexpr,
36
+ VAL_DATA_BYTES: tl.constexpr,
37
+ BLOCK_D: tl.constexpr,
38
+ BLOCK_VAL: tl.constexpr,
39
+ BLOCK_GRP: tl.constexpr,
40
+):
41
+ """Uniform quantization of values to VQB bits, pack, and store with scale/zero."""
42
+ val_cache_offset = KPS
43
+
44
+ if VQB == 3:
45
+ val_vec = tl.load(Value_ptr + base + d_offs, mask=d_mask, other=0.0).to(
46
+ tl.float32
47
+ )
48
+ val_min = tl.min(tl.where(d_mask, val_vec, float("inf")), axis=0)
49
+ val_max = tl.max(tl.where(d_mask, val_vec, -float("inf")), axis=0)
50
+ v_scale = (val_max - val_min) / 7.0
51
+ v_scale = tl.where(v_scale > 1e-8, v_scale, 1e-8)
52
+
53
+ q_vals = tl.minimum(
54
+ tl.maximum(((val_vec - val_min) / v_scale + 0.5).to(tl.int32), 0), 7
55
+ )
56
+
57
+ grp_offs = tl.arange(0, BLOCK_GRP)
58
+ grp_mask = grp_offs < (D // 8)
59
+ q_grp = tl.reshape(q_vals, [BLOCK_GRP, 8])
60
+ shifts_3bit = tl.arange(0, 8) * 3
61
+ packed_24 = tl.sum(q_grp << shifts_3bit[None, :], axis=1)
62
+ b0 = (packed_24 & 0xFF).to(tl.uint8)
63
+ b1 = ((packed_24 >> 8) & 0xFF).to(tl.uint8)
64
+ b2 = ((packed_24 >> 16) & 0xFF).to(tl.uint8)
65
+ tl.store(
66
+ KV_cache_ptr + slot_base + val_cache_offset + grp_offs * 3,
67
+ b0,
68
+ mask=grp_mask,
69
+ )
70
+ tl.store(
71
+ KV_cache_ptr + slot_base + val_cache_offset + grp_offs * 3 + 1,
72
+ b1,
73
+ mask=grp_mask,
74
+ )
75
+ tl.store(
76
+ KV_cache_ptr + slot_base + val_cache_offset + grp_offs * 3 + 2,
77
+ b2,
78
+ mask=grp_mask,
79
+ )
80
+
81
+ sc_offset = val_cache_offset + VAL_DATA_BYTES
82
+ sc_f16 = v_scale.to(tl.float16)
83
+ sc_u16 = sc_f16.to(tl.uint16, bitcast=True)
84
+ tl.store(KV_cache_ptr + slot_base + sc_offset, (sc_u16 & 0xFF).to(tl.uint8))
85
+ tl.store(
86
+ KV_cache_ptr + slot_base + sc_offset + 1,
87
+ ((sc_u16 >> 8) & 0xFF).to(tl.uint8),
88
+ )
89
+ zr_f16 = val_min.to(tl.float16)
90
+ zr_u16 = zr_f16.to(tl.uint16, bitcast=True)
91
+ tl.store(KV_cache_ptr + slot_base + sc_offset + 2, (zr_u16 & 0xFF).to(tl.uint8))
92
+ tl.store(
93
+ KV_cache_ptr + slot_base + sc_offset + 3,
94
+ ((zr_u16 >> 8) & 0xFF).to(tl.uint8),
95
+ )
96
+
97
+ else: # VQB == 4
98
+ val_vec = tl.load(Value_ptr + base + d_offs, mask=d_mask, other=0.0).to(
99
+ tl.float32
100
+ )
101
+ val_min = tl.min(tl.where(d_mask, val_vec, float("inf")), axis=0)
102
+ val_max = tl.max(tl.where(d_mask, val_vec, -float("inf")), axis=0)
103
+ v_scale = (val_max - val_min) / 15.0
104
+ v_scale = tl.where(v_scale > 1e-8, v_scale, 1e-8)
105
+
106
+ # Quantize all D elements from register (no re-load)
107
+ q_all = tl.minimum(
108
+ tl.maximum(((val_vec - val_min) / v_scale + 0.5).to(tl.int32), 0), 15
109
+ )
110
+ # Reshape to pairs and pack two 4-bit values per byte
111
+ q_pairs = tl.reshape(q_all, [BLOCK_D // 2, 2])
112
+ shifts_4 = tl.arange(0, 2) * 4
113
+ packed_val = tl.sum((q_pairs & 0xF) << shifts_4[None, :], axis=1).to(tl.uint8)
114
+ val_offs = tl.arange(0, BLOCK_D // 2)
115
+ val_mask = val_offs < VAL_DATA_BYTES
116
+ tl.store(
117
+ KV_cache_ptr + slot_base + val_cache_offset + val_offs,
118
+ packed_val,
119
+ mask=val_mask,
120
+ )
121
+
122
+ sc_offset = val_cache_offset + VAL_DATA_BYTES
123
+ sc_f16 = v_scale.to(tl.float16)
124
+ sc_u16 = sc_f16.to(tl.uint16, bitcast=True)
125
+ tl.store(KV_cache_ptr + slot_base + sc_offset, (sc_u16 & 0xFF).to(tl.uint8))
126
+ tl.store(
127
+ KV_cache_ptr + slot_base + sc_offset + 1,
128
+ ((sc_u16 >> 8) & 0xFF).to(tl.uint8),
129
+ )
130
+ zr_f16 = val_min.to(tl.float16)
131
+ zr_u16 = zr_f16.to(tl.uint16, bitcast=True)
132
+ tl.store(KV_cache_ptr + slot_base + sc_offset + 2, (zr_u16 & 0xFF).to(tl.uint8))
133
+ tl.store(
134
+ KV_cache_ptr + slot_base + sc_offset + 3,
135
+ ((zr_u16 >> 8) & 0xFF).to(tl.uint8),
136
+ )
137
+
138
+
139
+# ═══════════════════════════════════════════════════════════════════════
140
+# FP8 key store + value uniform quantization
141
+# ═══════════════════════════════════════════════════════════════════════
142
+
143
+
144
+@triton.jit
145
+def _tq_fused_store_fp8(
146
+ Key_ptr, # [NH, D] float16/bfloat16 — raw keys
147
+ Value_ptr, # [NH, D] float16/bfloat16 — raw values
148
+ KV_cache_ptr, # [total_bytes] uint8 (flattened view)
149
+ Slot_mapping_ptr, # [N] int32 — per-token slot indices
150
+ # Cache strides (for computing byte offsets)
151
+ stride_cache_block: tl.constexpr,
152
+ stride_cache_pos: tl.constexpr,
153
+ stride_cache_head: tl.constexpr,
154
+ # Dimensions
155
+ D: tl.constexpr,
156
+ H: tl.constexpr,
157
+ BLOCK_SIZE: tl.constexpr,
158
+ BLOCK_D: tl.constexpr,
159
+ # TQ layout
160
+ KPS: tl.constexpr,
161
+ # Value quantization
162
+ VQB: tl.constexpr,
163
+ VAL_DATA_BYTES: tl.constexpr,
164
+ # Packing block sizes
165
+ BLOCK_VAL: tl.constexpr,
166
+ BLOCK_GRP: tl.constexpr = 16,
167
+ FP8_E4B15: tl.constexpr = 0, # 1 = e4b15 (Ampere/Ada), 0 = e4nv (Hopper+)
168
+):
169
+ """FP8 key cast+scatter + value uniform quantization."""
170
+ pid = tl.program_id(0)
171
+ token_idx = pid // H
172
+ head_idx = pid % H
173
+
174
+ slot = tl.load(Slot_mapping_ptr + token_idx)
175
+ if slot < 0:
176
+ return
177
+ blk = slot // BLOCK_SIZE
178
+ off = slot % BLOCK_SIZE
179
+ slot_base = (
180
+ blk * stride_cache_block + off * stride_cache_pos + head_idx * stride_cache_head
181
+ )
182
+
183
+ base = pid * D
184
+
185
+ # ── FP8 KEY: cast to FP8 in-kernel and store ─────────────────
186
+ d_offs = tl.arange(0, BLOCK_D)
187
+ d_mask = d_offs < D
188
+ k_vals = tl.load(Key_ptr + base + d_offs, mask=d_mask, other=0.0)
189
+ k_fp8 = k_vals.to(tl.float8e4b15) if FP8_E4B15 else k_vals.to(tl.float8e4nv)
190
+ k_bytes = k_fp8.to(tl.uint8, bitcast=True)
191
+ tl.store(KV_cache_ptr + slot_base + d_offs, k_bytes, mask=d_mask)
192
+
193
+ # ── VALUE QUANTIZE + PACK ───────────────────────────────────────
194
+ _store_quantized_value(
195
+ Value_ptr,
196
+ KV_cache_ptr,
197
+ base,
198
+ slot_base,
199
+ d_offs,
200
+ d_mask,
201
+ D=D,
202
+ KPS=KPS,
203
+ VQB=VQB,
204
+ VAL_DATA_BYTES=VAL_DATA_BYTES,
205
+ BLOCK_D=BLOCK_D,
206
+ BLOCK_VAL=BLOCK_VAL,
207
+ BLOCK_GRP=BLOCK_GRP,
208
+ )
209
+
210
+
211
+# ═══════════════════════════════════════════════════════════════════════
212
+# Fused MSE store: bucketize + MSE index pack + norm store + value pack
213
+# (eliminates 4 PyTorch kernel launches per layer vs pack-only kernel)
214
+# ═══════════════════════════════════════════════════════════════════════
215
+
216
+
217
+@triton.jit
218
+def _tq_fused_store_mse(
219
+ # Post-rotation inputs
220
+ Y_ptr, # [NH, D] float32 — rotated normalized keys (x_hat @ PiT)
221
+ Norms_ptr, # [NH] float32 — key vector norms (||k||)
222
+ Value_ptr, # [NH, D] float32 — raw values
223
+ # Quantization tables
224
+ Midpoints_ptr, # [n_centroids-1] float32
225
+ # Cache and indexing
226
+ KV_cache_ptr, # [total_bytes] uint8 (flattened view)
227
+ Slot_mapping_ptr, # [N] int32 — per-token slot indices
228
+ # Cache strides
229
+ stride_cache_block: tl.constexpr,
230
+ stride_cache_pos: tl.constexpr,
231
+ stride_cache_head: tl.constexpr,
232
+ # Dimensions
233
+ D: tl.constexpr,
234
+ H: tl.constexpr,
235
+ BLOCK_SIZE: tl.constexpr,
236
+ BLOCK_D: tl.constexpr,
237
+ # TQ layout
238
+ MSE_BYTES: tl.constexpr,
239
+ KPS: tl.constexpr,
240
+ # Value quantization
241
+ VQB: tl.constexpr,
242
+ VAL_DATA_BYTES: tl.constexpr,
243
+ # Packing block sizes
244
+ BLOCK_VAL: tl.constexpr,
245
+ # MSE params
246
+ MSE_BITS: tl.constexpr,
247
+ N_CENTROIDS: tl.constexpr,
248
+ BLOCK_GRP: tl.constexpr = 16,
249
+):
250
+ """Fused MSE quantize + pack + store.
251
+
252
+ Performs binary-search bucketize, MSE index packing, norm storage,
253
+ and value quantization in one kernel.
254
+ """
255
+ pid = tl.program_id(0)
256
+ token_idx = pid // H
257
+ head_idx = pid % H
258
+
259
+ slot = tl.load(Slot_mapping_ptr + token_idx)
260
+ if slot < 0:
261
+ return
262
+ blk = slot // BLOCK_SIZE
263
+ off = slot % BLOCK_SIZE
264
+ slot_base = (
265
+ blk * stride_cache_block + off * stride_cache_pos + head_idx * stride_cache_head
266
+ )
267
+
268
+ base = pid * D
269
+ d_offs = tl.arange(0, BLOCK_D)
270
+ d_mask = d_offs < D
271
+
272
+ # ── 1. BINARY SEARCH BUCKETIZE ───────────────────────────────────
273
+ # Midpoints are sorted (N_CENTROIDS-1 values); binary search finds
274
+ # insertion point in MSE_BITS iterations vs N_CENTROIDS-1 for linear.
275
+ y_vec = tl.load(Y_ptr + base + d_offs, mask=d_mask, other=0.0)
276
+ lo = tl.zeros([BLOCK_D], dtype=tl.int32)
277
+ hi = tl.full([BLOCK_D], N_CENTROIDS - 1, dtype=tl.int32)
278
+ for _ in range(MSE_BITS):
279
+ mid = (lo + hi) >> 1
280
+ # Clamp to valid midpoint index [0, N_CENTROIDS-2] for load safety;
281
+ # the search result (lo) is still correct since converged lanes
282
+ # don't change.
283
+ safe_mid = tl.minimum(mid, N_CENTROIDS - 2)
284
+ mid_val = tl.load(Midpoints_ptr + safe_mid, mask=d_mask, other=0.0)
285
+ lo = tl.where(y_vec >= mid_val, mid + 1, lo)
286
+ hi = tl.where(y_vec >= mid_val, hi, mid)
287
+ idx = tl.minimum(lo, N_CENTROIDS - 1)
288
+
289
+ # ── 2. PACK MSE INDICES from register idx ─────────────────────────
290
+ if MSE_BITS == 4:
291
+ idx_pairs = tl.reshape(idx, [BLOCK_D // 2, 2])
292
+ shifts_4 = tl.arange(0, 2) * 4
293
+ packed = tl.sum((idx_pairs & 0xF) << shifts_4[None, :], axis=1).to(tl.uint8)
294
+ mse_offs = tl.arange(0, BLOCK_D // 2)
295
+ mse_mask = mse_offs < MSE_BYTES
296
+ tl.store(KV_cache_ptr + slot_base + mse_offs, packed, mask=mse_mask)
297
+
298
+ elif MSE_BITS == 3:
299
+ grp_offs = tl.arange(0, BLOCK_GRP)
300
+ grp_mask = grp_offs < (D // 8)
301
+ idx_grp = tl.reshape(idx, [BLOCK_GRP, 8])
302
+ shifts_3 = tl.arange(0, 8) * 3
303
+ packed_24 = tl.sum((idx_grp & 0x7) << shifts_3[None, :], axis=1)
304
+ b0 = (packed_24 & 0xFF).to(tl.uint8)
305
+ b1 = ((packed_24 >> 8) & 0xFF).to(tl.uint8)
306
+ b2 = ((packed_24 >> 16) & 0xFF).to(tl.uint8)
307
+ tl.store(KV_cache_ptr + slot_base + grp_offs * 3, b0, mask=grp_mask)
308
+ tl.store(KV_cache_ptr + slot_base + grp_offs * 3 + 1, b1, mask=grp_mask)
309
+ tl.store(KV_cache_ptr + slot_base + grp_offs * 3 + 2, b2, mask=grp_mask)
310
+
311
+ # ── 3. STORE vec_norm (fp16, 2 bytes) ─────────────────────────────
312
+ norm_offset = MSE_BYTES
313
+
314
+ vn_f16 = tl.load(Norms_ptr + pid).to(tl.float16)
315
+ vn_u16 = vn_f16.to(tl.uint16, bitcast=True)
316
+ tl.store(KV_cache_ptr + slot_base + norm_offset, (vn_u16 & 0xFF).to(tl.uint8))
317
+ tl.store(
318
+ KV_cache_ptr + slot_base + norm_offset + 1, ((vn_u16 >> 8) & 0xFF).to(tl.uint8)
319
+ )
320
+
321
+ # ── 4. VALUE QUANTIZE + PACK ──────────────────────────────────────
322
+ _store_quantized_value(
323
+ Value_ptr,
324
+ KV_cache_ptr,
325
+ base,
326
+ slot_base,
327
+ d_offs,
328
+ d_mask,
329
+ D=D,
330
+ KPS=KPS,
331
+ VQB=VQB,
332
+ VAL_DATA_BYTES=VAL_DATA_BYTES,
333
+ BLOCK_D=BLOCK_D,
334
+ BLOCK_VAL=BLOCK_VAL,
335
+ BLOCK_GRP=BLOCK_GRP,
336
+ )
337
+
338
+
339
+# ═══════════════════════════════════════════════════════════════════════
340
+# Launcher
341
+# ═══════════════════════════════════════════════════════════════════════
342
+
343
+
344
+def triton_turboquant_store(
345
+ key: torch.Tensor, # [N, H, D] — raw keys (post-RoPE)
346
+ value: torch.Tensor, # [N, H, D] — raw values
347
+ kv_cache: torch.Tensor, # [num_blocks, block_size, Hk, padded_slot] uint8
348
+ slot_mapping: torch.Tensor, # [N] int32
349
+ PiT: torch.Tensor, # [D, D] float32
350
+ midpoints: torch.Tensor, # [n_centroids-1] float32
351
+ mse_bits: int,
352
+ key_packed_size: int,
353
+ value_quant_bits: int,
354
+ key_fp8: bool = False,
355
+):
356
+ """Launch TQ store kernel (FP8 or MSE path)."""
357
+ N, H, D = key.shape
358
+ NH = N * H
359
+ block_size = kv_cache.shape[1]
360
+ BLOCK_D = triton.next_power_of_2(D)
361
+ mse_bytes = math.ceil(D * mse_bits / 8)
362
+ n_centroids = 2**mse_bits
363
+
364
+ val_data_bytes = math.ceil(D * value_quant_bits / 8)
365
+
366
+ BLOCK_VAL = triton.next_power_of_2(val_data_bytes)
367
+
368
+ # Cache strides (element_size=1 for uint8, so stride in bytes = stride())
369
+ stride_block = kv_cache.stride(0)
370
+ stride_pos = kv_cache.stride(1)
371
+ stride_head = kv_cache.stride(2)
372
+
373
+ block_grp = triton.next_power_of_2(D // 8) if D >= 8 else 1
374
+
375
+ # ── FP8 PATH: in-kernel FP8 cast + scatter via fp8 kernel ──
376
+ if key_fp8:
377
+ k_flat = key.reshape(NH, D).contiguous()
378
+ v_flat = value.reshape(NH, D).contiguous()
379
+
380
+ fp8_e4b15 = _use_fp8_e4b15(key.device.index or 0)
381
+
382
+ grid = (NH,)
383
+ _tq_fused_store_fp8[grid](
384
+ k_flat,
385
+ v_flat,
386
+ kv_cache.view(-1),
387
+ slot_mapping,
388
+ stride_cache_block=stride_block,
389
+ stride_cache_pos=stride_pos,
390
+ stride_cache_head=stride_head,
391
+ D=D,
392
+ H=H,
393
+ BLOCK_SIZE=block_size,
394
+ BLOCK_D=BLOCK_D,
395
+ KPS=key_packed_size,
396
+ VQB=value_quant_bits,
397
+ VAL_DATA_BYTES=val_data_bytes,
398
+ BLOCK_VAL=BLOCK_VAL,
399
+ BLOCK_GRP=block_grp,
400
+ FP8_E4B15=fp8_e4b15,
401
+ num_warps=4,
402
+ num_stages=1,
403
+ )
404
+ return
405
+
406
+ # ── MSE PATH: external GEMM + fused bucketize/pack kernel ──
407
+ # Normalize + rotation GEMM externally (cuBLAS is faster than in-kernel)
408
+ k_flat = key.float().reshape(NH, D)
409
+ norms = k_flat.norm(dim=1, keepdim=True)
410
+ x_hat = k_flat / (norms + 1e-8)
411
+ y = x_hat @ PiT
412
+
413
+ v_flat = value.float().reshape(NH, D)
414
+
415
+ # Fused kernel: bucketize + MSE index pack + norm store + value pack
416
+ grid = (NH,)
417
+ _tq_fused_store_mse[grid](
418
+ y,
419
+ norms.squeeze(1),
420
+ v_flat,
421
+ midpoints,
422
+ kv_cache.view(-1),
423
+ slot_mapping,
424
+ stride_cache_block=stride_block,
425
+ stride_cache_pos=stride_pos,
426
+ stride_cache_head=stride_head,
427
+ D=D,
428
+ H=H,
429
+ BLOCK_SIZE=block_size,
430
+ BLOCK_D=BLOCK_D,
431
+ MSE_BYTES=mse_bytes,
432
+ KPS=key_packed_size,
433
+ VQB=value_quant_bits,
434
+ VAL_DATA_BYTES=val_data_bytes,
435
+ BLOCK_VAL=BLOCK_VAL,
436
+ MSE_BITS=mse_bits,
437
+ N_CENTROIDS=n_centroids,
438
+ BLOCK_GRP=block_grp,
439
+ num_warps=4,
440
+ num_stages=1,
441
+ )
vllm/v1/core/single_type_kv_cache_manager.py
+4
-2
@@ -21,6 +21,7 @@ from vllm.v1.kv_cache_interface import (
21
MLAAttentionSpec,
22
SinkFullAttentionSpec,
23
SlidingWindowSpec,
24
+ TQFullAttentionSpec,
25
)
26
from vllm.v1.request import Request
27
@@ -209,7 +210,7 @@ class SingleTypeKVCacheManager(ABC):
210
cdiv(num_total_computed_tokens, self.block_size) - len(req_blocks)
211
)
212
req_blocks.extend(allocated_blocks)
212
- if type(self.kv_cache_spec) is FullAttentionSpec:
213
+ if type(self.kv_cache_spec) in (FullAttentionSpec, TQFullAttentionSpec):
214
self.new_block_ids.extend(b.block_id for b in allocated_blocks)
215
216
def allocate_new_blocks(
@@ -237,7 +238,7 @@ class SingleTypeKVCacheManager(ABC):
238
else:
239
new_blocks = self.block_pool.get_new_blocks(num_new_blocks)
240
req_blocks.extend(new_blocks)
240
- if type(self.kv_cache_spec) is FullAttentionSpec:
241
+ if type(self.kv_cache_spec) in (FullAttentionSpec, TQFullAttentionSpec):
242
self.new_block_ids.extend(b.block_id for b in new_blocks)
243
return new_blocks
244
@@ -1114,6 +1115,7 @@ class SinkFullAttentionManager(FullAttentionManager):
1115
1116
spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] = {
1117
FullAttentionSpec: FullAttentionManager,
1118
+ TQFullAttentionSpec: FullAttentionManager,
1119
MLAAttentionSpec: FullAttentionManager,
1120
SlidingWindowSpec: SlidingWindowManager,
1121
ChunkedLocalAttentionSpec: ChunkedLocalAttentionManager,
vllm/v1/kv_cache_interface.py
+26
@@ -245,6 +245,32 @@ class FullAttentionSpec(AttentionSpec):
245
)
246
247
248
+@dataclass(frozen=True, kw_only=True)
249
+class TQFullAttentionSpec(FullAttentionSpec):
250
+ """FullAttentionSpec with TQ-aware page size.
251
+
252
+ Python equivalent of the C++ TQ4FullAttentionSpec. Overrides
253
+ real_page_size_bytes to use TQ slot bytes instead of the raw
254
+ head_size * dtype formula.
255
+ """
256
+
257
+ tq_slot_size: int = 0
258
+
259
+ @property
260
+ def real_page_size_bytes(self) -> int:
261
+ if self.tq_slot_size > 0:
262
+ return self.block_size * self.num_kv_heads * self.tq_slot_size
263
+ return super().real_page_size_bytes
264
+
265
+ @classmethod
266
+ def merge(cls, specs: list[Self]) -> Self:
267
+ merged = super().merge(specs)
268
+ assert all(s.tq_slot_size == specs[0].tq_slot_size for s in specs), (
269
+ "All TQ layers in the same KV cache group must use the same tq_slot_size."
270
+ )
271
+ return replace(merged, tq_slot_size=specs[0].tq_slot_size)
272
+
273
+
274
@dataclass(frozen=True, kw_only=True)
275
class MLAAttentionSpec(FullAttentionSpec):
276
# TODO(Lucas/Chen): less hacky way to do this
vllm/v1/worker/utils.py
+1
-1
@@ -120,7 +120,7 @@ class KVBlockZeroer:
120
121
for group in attn_groups_iter:
122
spec = group.kv_cache_spec
123
- if type(spec) is not FullAttentionSpec:
123
+ if not isinstance(spec, FullAttentionSpec):
124
continue
125
if group.kv_cache_group_id >= len(kernel_block_sizes):
126
continue