@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 57bd0ed441

[5/N][KV-Cache Layout Refactor] Backend-published KV packing via customize_spec (#51704)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Claude <noreply@anthropic.com>

Lucas Wilkinson committed Aug 14, 2026 at 09:47 UTC 57bd0ed441095b5c546707cabe25d3fa08b7f161

18 files changed +204 -219

tests/quantization/test_turboquant.py

+15 -6

@@ -281,7 +281,10 @@ class TestTurboQuantKVCacheSpec:

281 @pytest.mark.parametrize("preset", ALL_PRESETS)

282 def test_kv_cache_spec_sets_kv_quant_mode(self, preset):

283 from vllm.model_executor.layers.attention.attention import Attention

284 - from vllm.v1.kv_cache_interface import KVQuantMode, TQFullAttentionSpec

284 + from vllm.v1.attention.backends.turboquant_attn import (

285 + TurboQuantAttentionBackend,

286 + )

287 + from vllm.v1.kv_cache_interface import FullAttentionSpec

288

289 layer = SimpleNamespace(

290 attn_type="decoder",

@@ -294,10 +297,16 @@ class TestTurboQuantKVCacheSpec:

297 )

298 vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=32))

299

300 + # The layer builds an unpacked spec; the worker's spec-collection

301 + # loop applies TQ slot packing via the backend's customize_spec hook.

302 spec = Attention.get_kv_cache_spec(layer, vllm_config)

303 + assert isinstance(spec, FullAttentionSpec)

304 + assert spec.kv_quant_mode.is_turboquant

305 + assert spec.state_content_bytes is None

306

299 - assert isinstance(spec, TQFullAttentionSpec)

300 - assert spec.kv_quant_mode == KVQuantMode.TURBOQUANT

307 + spec = TurboQuantAttentionBackend.customize_spec(spec)

308 + expected_slot = TurboQuantConfig.from_cache_dtype(preset, 128).slot_size_aligned

309 + assert spec.state_content_bytes == expected_slot

310

311

312 class TestTurboQuantWorkspaceReservation:

@@ -333,15 +342,15 @@ class TestTurboQuantWorkspaceReservation:

342

343 @staticmethod

344 def _fake_kv_cache_spec():

336 - from vllm.v1.kv_cache_interface import TQFullAttentionSpec

345 + from vllm.v1.kv_cache_interface import FullAttentionSpec

346

338 - return TQFullAttentionSpec(

347 + return FullAttentionSpec(

348 block_size=32,

349 num_kv_heads=4,

350 head_size=128,

351 head_size_v=128,

352 dtype=torch.uint8,

344 - tq_slot_size=102,

353 + state_content_bytes=102,

354 )

355

356 def test_metadata_builder_reserves_decode_and_continuation_prefill_workspace(

tests/v1/core/test_kv_cache_utils.py

+10 -11

@@ -2766,21 +2766,20 @@ def test_unify_kv_cache_page_size_padding_requires_backend_support():

2766 kv_cache_utils.unify_kv_cache_spec_page_size(specs)

2767

2768

2769 -def test_unpadded_page_size_without_quant_matches_real_page():

2770 - # Without quantization the offload transfer width is just the raw page.

2771 - spec = new_kv_cache_spec()

2772 - assert spec.unpadded_page_size_bytes == spec.real_page_size_bytes

2773 - assert spec.page_size_bytes == spec.unpadded_page_size_bytes

2774 -

2775 -

2769 def test_unpadded_page_size_includes_per_token_head_scales():

2770 # Per-token-head quant carries inline fp32 scales that are carved from the

2778 - # raw KV allocation, so they must be budgeted into the offload width.

2779 - spec = new_kv_cache_spec(

2780 - dtype=torch.uint8, kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD

2771 + # raw KV allocation, so they must be budgeted into the offload width. The

2772 + # packing is published by the owning backend's customize_spec hook.

2773 + from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend

2774 +

2775 + dense = new_kv_cache_spec(dtype=torch.uint8)

2776 + spec = TritonAttentionBackend.customize_spec(

2777 + new_kv_cache_spec(

2778 + dtype=torch.uint8, kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD

2779 + )

2780 )

2781 scales = 2 * spec.block_size * spec.num_kv_heads * 4

2783 - assert spec.unpadded_page_size_bytes == spec.real_page_size_bytes + scales

2782 + assert spec.unpadded_page_size_bytes == dense.unpadded_page_size_bytes + scales

2783 assert spec.page_size_bytes == spec.unpadded_page_size_bytes

2784

2785

tests/v1/test_kv_cache_spec_registry.py

-11

@@ -34,7 +34,6 @@ from vllm.v1.kv_cache_interface import (

34 SinkFullAttentionSpec,

35 SlidingWindowMLASpec,

36 SlidingWindowSpec,

37 - TQFullAttentionSpec,

37 UniformTypeKVCacheSpecs,

38 get_kv_cache_spec_kind,

39 )

@@ -85,7 +84,6 @@ class _TrulyUnregisteredSpec(KVCacheSpec):

85 spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] = {

86 FullAttentionSpec: FullAttentionManager,

88 - TQFullAttentionSpec: FullAttentionManager,

87 MLAAttentionSpec: FullAttentionManager,

88 HiddenStateCacheSpec: FullAttentionManager,

89 SlidingWindowSpec: SlidingWindowManager,

@@ -98,7 +96,6 @@ spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] = {

97 spec_uniform_base_map: dict[type[KVCacheSpec], type[KVCacheSpec]] = {

98 FullAttentionSpec: FullAttentionSpec,

101 - TQFullAttentionSpec: FullAttentionSpec,

99 MLAAttentionSpec: FullAttentionSpec,

100 HiddenStateCacheSpec: FullAttentionSpec,

101 SlidingWindowSpec: SlidingWindowSpec,

@@ -113,13 +110,6 @@ spec_args_map: dict[type[KVCacheSpec], dict[str, Any]] = {

110 FullAttentionSpec: dict(

111 block_size=64, num_kv_heads=8, head_size=128, dtype=torch.bfloat16

112 ),

116 - TQFullAttentionSpec: dict(

117 - block_size=64,

118 - num_kv_heads=8,

119 - head_size=128,

120 - dtype=torch.bfloat16,

121 - tq_slot_size=256,

122 - ),

113 MLAAttentionSpec: dict(

114 block_size=64, num_kv_heads=1, head_size=128, dtype=torch.bfloat16

115 ),

@@ -262,7 +252,6 @@ class TestKVCacheSpecRegistry:

252 def test_full_attention_family_specs_are_uniform(self):

253 specs = [

254 make_spec(FullAttentionSpec),

265 - make_spec(TQFullAttentionSpec),

255 make_spec(MLAAttentionSpec),

256 make_spec(HiddenStateCacheSpec),

257 make_spec(SinkFullAttentionSpec),

tests/v1/worker/test_dsv4_packed_zeroer_geometry.py

@@ -93,6 +93,7 @@ def test_dsv4_packed_zeroer_geometry():

93 cache_dtype_str="fp8_ds_mla",

94 alignment=576,

95 model_version="deepseek_v4",

96 + state_content_bytes=584, # >576 to allocate room for scales at back of page

97 )

98 for _ in layer_names

99 ]

vllm/model_executor/layers/attention/attention.py

+11 -26

@@ -623,14 +623,17 @@ class Attention(nn.Module, AttentionLayerBase):

623 # bytes per block. Otherwise (page_size_padded is None) the smallest

624 # block is fine — ``unify`` scales it up by an integer ratio.

625 shared_page = vllm_config.cache_config.skip_page_size_padded

626 - sw_per_token = SlidingWindowSpec(

627 - block_size=1,

628 - num_kv_heads=self.num_kv_heads,

629 - head_size=self.head_size,

630 - head_size_v=self.head_size_v,

631 - dtype=self.kv_cache_torch_dtype,

632 - kv_quant_mode=quant_mode,

633 - sliding_window=self.sliding_window,

626 + # The backend owns its packing

627 + sw_per_token = self.attn_backend.customize_spec(

628 + SlidingWindowSpec(

629 + block_size=1,

630 + num_kv_heads=self.num_kv_heads,

631 + head_size=self.head_size,

632 + head_size_v=self.head_size_v,

633 + dtype=self.kv_cache_torch_dtype,

634 + kv_quant_mode=quant_mode,

635 + sliding_window=self.sliding_window,

636 + )

637 ).real_page_size_bytes

638 sw_block_size = _largest_kernel_block_within(

639 self.attn_backend,

@@ -649,24 +652,6 @@ class Attention(nn.Module, AttentionLayerBase):

652 sliding_window=self.sliding_window,

653 page_size_padded=shared_page,

654 )

652 - elif self.kv_cache_dtype.startswith("turboquant_"):

653 - from vllm.model_executor.layers.quantization.turboquant.config import (

654 - TurboQuantConfig,

655 - )

656 - from vllm.v1.kv_cache_interface import TQFullAttentionSpec

657 -

658 - tq_config = TurboQuantConfig.from_cache_dtype(

659 - self.kv_cache_dtype, self.head_size

660 - )

661 - return TQFullAttentionSpec(

662 - block_size=block_size,

663 - num_kv_heads=self.num_kv_heads,

664 - head_size=self.head_size,

665 - head_size_v=self.head_size,

666 - dtype=self.kv_cache_torch_dtype,

667 - kv_quant_mode=quant_mode,

668 - tq_slot_size=tq_config.slot_size_aligned,

669 - )

655 else:

656 return FullAttentionSpec(

657 block_size=block_size,

vllm/model_executor/layers/attention/mla_attention.py

+20 -1

@@ -204,7 +204,7 @@ import itertools

204 import math

205 from abc import abstractmethod

206 from collections.abc import Callable

207 -from dataclasses import dataclass

207 +from dataclasses import dataclass, replace

208 from enum import Enum

209 from math import lcm

210 from typing import ClassVar, Generic, TypeVar, cast

@@ -270,6 +270,7 @@ from vllm.utils.torch_utils import (

270 _encode_layer_name,

271 _resolve_layer_name,

272 direct_register_custom_op,

273 + get_dtype_size,

274 is_quantized_kv_cache,

275 kv_cache_dtype_str_to_dtype,

276 np_to_pinned_tensor,

@@ -298,6 +299,7 @@ from vllm.v1.attention.selector import get_attn_backend

299 from vllm.v1.kv_cache_interface import (

300 AttentionSpec,

301 KVCacheSpec,

302 + KVQuantMode,

303 MLAAttentionSpec,

304 SlidingWindowMLASpec,

305 get_kv_quant_mode,

@@ -1152,6 +1154,9 @@ class MLAAttention(nn.Module, AttentionLayerBase):

1154 dtype=kv_cache_dtype,

1155 cache_dtype_str=self.kv_cache_dtype,

1156 kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),

1157 + # fp8_ds_mla: 656-byte custom layout (kv_lora_rank=512 +

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

1159 + state_content_bytes=656 if self.kv_cache_dtype == "fp8_ds_mla" else None,

1160 )

1161 if self.sliding_window is not None:

1162 return SlidingWindowMLASpec(

@@ -1370,6 +1375,20 @@ class _DecodeConcatQuantFP8(QuantFP8):

1375

1376

1377 class MLACommonBackend(AttentionBackend):

1378 + @classmethod

1379 + def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec":

1380 + """Per-token-head modes pack an inline fp32 scale pair after the

1381 + latent data (single-sided: ``head_size_v == 0`` for MLA)."""

1382 + mode = spec.kv_quant_mode

1383 + if spec.state_content_bytes is not None or not mode.is_per_token_head:

1384 + return spec

1385 + head_size = spec.head_size

1386 + if mode == KVQuantMode.INT4_PER_TOKEN_HEAD:

1387 + head_size //= 2

1388 + scale_bytes = get_dtype_size(torch.float32)

1389 + content = head_size * get_dtype_size(spec.dtype) + 2 * scale_bytes

1390 + return replace(spec, state_content_bytes=content)

1391 +

1392 @staticmethod

1393 def get_name() -> str:

1394 return "TRITON_MLA"

vllm/models/deepseek_v4/attention.py

@@ -670,6 +670,9 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):

670 alignment=576 if uses_fp8_ds_mla_layout else 512,

671 model_version="deepseek_v4",

672 kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),

673 + # DeepseekV4: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B per token;

674 + # head_size stays semantic (512).

675 + state_content_bytes=584 if uses_fp8_ds_mla_layout else None,

676 )

677

678

vllm/models/kimi_k3/nvidia/mla.py

@@ -400,6 +400,8 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):

400 dtype=kv_cache_dtype,

401 cache_dtype_str=self.kv_cache_dtype,

402 kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),

403 + # fp8_ds_mla: 656-byte custom layout; see flashmla_sparse.py.

404 + state_content_bytes=656 if self.kv_cache_dtype == "fp8_ds_mla" else None,

405 non_causal_multi_token_decode=self.non_causal_multi_token_decode,

406 )

407

vllm/platforms/interface.py

+13 -13

@@ -687,13 +687,15 @@ class Platform:

687

688 def per_token_page_bytes(dtype: "torch.dtype", cache_dtype: str) -> int:

689 """Bytes one token occupies in one layer, for the given dtype."""

690 - return FullAttentionSpec(

690 + spec = FullAttentionSpec(

691 block_size=1,

692 num_kv_heads=model_config.get_num_kv_heads(parallel_config),

693 head_size=model_config.get_head_size(),

694 dtype=dtype,

695 kv_quant_mode=get_kv_quant_mode(cache_dtype),

696 - ).page_size_bytes

696 + )

697 + # The backend owns its packing

698 + return backend_cls.customize_spec(spec).page_size_bytes

699

700 primary_dtype = (

701 STR_DTYPE_TO_TORCH_DTYPE[cache_config.cache_dtype]

@@ -809,23 +811,18 @@ class Platform:

811 # when all attention layers are TQ. With mixed skip+TQ the skip

812 # layers still use the standard layout — take max so mamba

813 # padding covers the largest actual page.

812 - from vllm.model_executor.layers.quantization.turboquant.config import (

813 - TurboQuantConfig,

814 + from vllm.v1.attention.backends.turboquant_attn import (

815 + TurboQuantAttentionBackend,

816 )

815 - from vllm.v1.kv_cache_interface import TQFullAttentionSpec

817

817 - tq_cfg = TurboQuantConfig.from_cache_dtype(

818 - cache_config.cache_dtype, model_config.get_head_size()

819 - )

820 - tq_page = TQFullAttentionSpec(

818 + tq_spec = FullAttentionSpec(

819 block_size=1,

820 num_kv_heads=model_config.get_num_kv_heads(parallel_config),

821 head_size=model_config.get_head_size(),

824 - head_size_v=model_config.get_head_size(),

822 dtype=kv_cache_dtype,

823 kv_quant_mode=kv_quant_mode,

827 - tq_slot_size=tq_cfg.slot_size_aligned,

828 - ).page_size_bytes

824 + )

825 + tq_page = TurboQuantAttentionBackend.customize_spec(tq_spec).page_size_bytes

826 if cache_config.kv_cache_dtype_skip_layers:

827 skip_page = FullAttentionSpec(

828 block_size=1,

@@ -840,12 +837,15 @@ class Platform:

837 else:

838 attn_page_size_1_token = tq_page

839 else:

843 - attn_page_size_1_token = FullAttentionSpec(

840 + attn_spec = FullAttentionSpec(

841 block_size=1,

842 num_kv_heads=model_config.get_num_kv_heads(parallel_config),

843 head_size=model_config.get_head_size(),

844 dtype=kv_cache_dtype,

845 kv_quant_mode=kv_quant_mode,

846 + )

847 + attn_page_size_1_token = backend_cls.customize_spec(

848 + attn_spec

849 ).page_size_bytes

850

851 # Compute mamba page size

vllm/v1/attention/backend.py

+15 -1

@@ -24,7 +24,7 @@ if TYPE_CHECKING:

24 from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey

25 from vllm.platforms.interface import DeviceCapability

26 from vllm.v1.attention.backends.utils import KVCacheLayoutType

27 - from vllm.v1.kv_cache_interface import KVCacheSpec, KVQuantMode

27 + from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec, KVQuantMode

29 from vllm.v1.kv_cache_interface import get_kv_quant_mode

@@ -197,6 +197,20 @@ class AttentionBackend(ABC):

197 return True

198 return False

199

200 + @classmethod

201 + def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec":

202 + """Adjust the layer's KV cache spec for this backend's kernels. Used when the

203 + kernels want KV packed in a specific way.

204 +

205 + NOTE: temporary compatibility API. Today the Attention layer builds the spec

206 + from the model config and the backend only gets to adjust it post-hoc; the end

207 + state is for the backend to build and return the spec directly, at which point

208 + this hook goes away.

209 +

210 + (see: https://github.com/vllm-project/vllm/issues/42449)

211 + """

212 + return spec

213 +

214 @classmethod

215 def get_preferred_block_size(cls, default_block_size: int) -> int:

216 supported_sizes = cls.get_supported_kernel_block_sizes()

vllm/v1/attention/backends/flashinfer.py

+17 -1

@@ -2,7 +2,7 @@

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

3 """Attention layer with FlashInfer."""

5 -from dataclasses import dataclass

5 +from dataclasses import dataclass, replace

6 from enum import Enum

7 from functools import partial

8 from typing import ClassVar

@@ -50,6 +50,7 @@ from vllm.utils.math_utils import cdiv

50 from vllm.utils.torch_utils import (

51 PIN_MEMORY,

52 canonicalize_singleton_dim_strides,

53 + get_dtype_size,

54 is_quantized_kv_cache,

55 is_strictly_contiguous,

56 nvfp4_kv_cache_full_dim,

@@ -387,6 +388,21 @@ class BatchDCPPrefillWrapper:

388

389

390 class FlashInferBackend(AttentionBackend):

391 + @classmethod

392 + def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec":

393 + """NVFP4 stores K and V as separate per-head slots of packed fp4 data

394 + plus fp8 block scales."""

395 + if spec.state_content_bytes is not None or not spec.kv_quant_mode.is_nvfp4:

396 + return spec

397 + hs_k = nvfp4_kv_cache_full_dim(spec.head_size)

398 + hs_v = nvfp4_kv_cache_full_dim(spec.head_size_v)

399 + assert hs_k == hs_v, "nvfp4 with asymmetric K/V head sizes not yet supported"

400 + return replace(

401 + spec,

402 + num_head_slots=2 * spec.num_kv_heads,

403 + state_content_bytes=hs_k * get_dtype_size(spec.dtype),

404 + )

405 +

406 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]

407 supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [

408 "auto",

vllm/v1/attention/backends/mla/sparse_swa.py

@@ -97,6 +97,8 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):

97 dtype=self.dtype,

98 sliding_window=self.window_size,

99 cache_dtype_str=self.cache_config.cache_dtype,

100 + # DeepseekV4 fp8_ds_mla: 584B per token (448B NoPE + 128B RoPE + 8B scales)

101 + state_content_bytes=584 if uses_fp8_ds_mla_layout else None,

102 # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577).

103 alignment=576 if uses_fp8_ds_mla_layout else 512,

104 model_version="deepseek_v4",

vllm/v1/attention/backends/triton_attn.py

+16 -2

@@ -2,7 +2,7 @@

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

3 """High-Performance Triton-only Attention layer."""

5 -from dataclasses import dataclass

5 +from dataclasses import dataclass, replace

6 from typing import ClassVar

8 import torch

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

19 from vllm.platforms import current_platform

20 from vllm.platforms.interface import DeviceCapability

21 from vllm.utils.math_utils import next_power_of_2

22 -from vllm.utils.torch_utils import is_quantized_kv_cache

22 +from vllm.utils.torch_utils import get_dtype_size, is_quantized_kv_cache

23 from vllm.v1.attention.backend import (

24 AttentionBackend,

25 AttentionCGSupport,

@@ -274,6 +274,20 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet

274

275

276 class TritonAttentionBackend(AttentionBackend):

277 + @classmethod

278 + def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec":

279 + """Per-token-head modes pack inline fp32 scales after each head's

280 + data, so the content is (data + one scale) per K/V side."""

281 + mode = spec.kv_quant_mode

282 + if spec.state_content_bytes is not None or not mode.is_per_token_head:

283 + return spec

284 + hs_k, hs_v = spec.head_size, spec.head_size_v

285 + if mode == KVQuantMode.INT4_PER_TOKEN_HEAD:

286 + hs_k, hs_v = hs_k // 2, hs_v // 2

287 + scale_bytes = get_dtype_size(torch.float32)

288 + content = (hs_k + hs_v) * get_dtype_size(spec.dtype) + 2 * scale_bytes

289 + return replace(spec, state_content_bytes=content)

290 +

291 supported_dtypes: ClassVar[list[torch.dtype]] = [

292 torch.float16,

293 torch.bfloat16,

vllm/v1/attention/backends/turboquant_attn.py

+16 -1

@@ -19,7 +19,7 @@ Per-head per-position slot layout:

19 import contextlib

20 import functools

21 import math

22 -from dataclasses import dataclass

22 +from dataclasses import dataclass, replace

23 from typing import Any, ClassVar

25 import torch

@@ -137,6 +137,21 @@ class TurboQuantAttentionBackend(AttentionBackend):

137 "turboquant_3bit_nc",

138 ]

139

140 + @classmethod

141 + def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec:

142 + """TurboQuant packs K+V into one slot per head."""

143 + if spec.state_content_bytes is not None or not spec.kv_quant_mode.is_turboquant:

144 + return spec

145 + from vllm.model_executor.layers.quantization.turboquant.config import (

146 + TurboQuantConfig,

147 + )

148 +

149 + # KVQuantMode member names mirror the preset strings.

150 + tq = TurboQuantConfig.from_cache_dtype(

151 + spec.kv_quant_mode.name.lower(), spec.head_size

152 + )

153 + return replace(spec, state_content_bytes=tq.slot_size_aligned)

154 +

155 @staticmethod

156 def get_name() -> str:

157 return "TURBOQUANT"

vllm/v1/core/single_type_kv_cache_manager.py

-6

@@ -28,7 +28,6 @@ from vllm.v1.kv_cache_interface import (

28 SinkFullAttentionSpec,

29 SlidingWindowMLASpec,

30 SlidingWindowSpec,

31 - TQFullAttentionSpec,

31 )

32 from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry

33 from vllm.v1.request import Request

@@ -1929,11 +1928,6 @@ def register_all_kvcache_specs(vllm_config):

1928 )

1929

1930 # FullAttentionSpec subclasses — grouped with FullAttentionSpec

1932 - KVCacheSpecRegistry.register(

1933 - TQFullAttentionSpec,

1934 - FullAttentionManager,

1935 - uniform_type_base_spec=FullAttentionSpec,

1936 - )

1931 KVCacheSpecRegistry.register(

1932 MLAAttentionSpec, FullAttentionManager, uniform_type_base_spec=FullAttentionSpec

1933 )

vllm/v1/kv_cache_interface.py

+61 -140

@@ -16,7 +16,7 @@ from typing_extensions import Self

17 from vllm.logger import init_logger

18 from vllm.utils.math_utils import cdiv, round_up

19 -from vllm.utils.torch_utils import get_dtype_size, nvfp4_kv_cache_full_dim

19 +from vllm.utils.torch_utils import get_dtype_size

20 from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum

21 from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry

@@ -46,7 +46,11 @@ class KVQuantMode(IntEnum):

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

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

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

49 - TURBOQUANT = 6 # Hadamard-rotated Lloyd-Max quant, packed K+V per slot

49 + # Hadamard-rotated Lloyd-Max quant, packed K+V per slot.

50 + TURBOQUANT_K8V4 = 6

51 + TURBOQUANT_4BIT_NC = 7

52 + TURBOQUANT_K3V4_NC = 8

53 + TURBOQUANT_3BIT_NC = 9

55 @property

56 def is_per_token_head(self) -> bool:

@@ -64,8 +68,13 @@ class KVQuantMode(IntEnum):

69 @property

70 def is_turboquant(self) -> bool:

67 - """True for turboquant quantization mode."""

68 - return self == KVQuantMode.TURBOQUANT

71 + """True for any turboquant quantization mode."""

72 + return self in (

73 + KVQuantMode.TURBOQUANT_K8V4,

74 + KVQuantMode.TURBOQUANT_4BIT_NC,

75 + KVQuantMode.TURBOQUANT_K3V4_NC,

76 + KVQuantMode.TURBOQUANT_3BIT_NC,

77 + )

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

@@ -79,7 +88,7 @@ def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode:

88 if kv_cache_dtype.startswith("nvfp4"):

89 return KVQuantMode.NVFP4

90 if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("turboquant_"):

82 - return KVQuantMode.TURBOQUANT

91 + return KVQuantMode[kv_cache_dtype.upper()]

92 if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("fp8"):

93 return KVQuantMode.FP8_PER_TENSOR

94 return KVQuantMode.NONE

@@ -210,21 +219,37 @@ class AttentionSpec(KVCacheSpec):

219 num_kv_heads: int

220 head_size: int

221 dtype: torch.dtype

222 + head_size_v: int = None # type: ignore[assignment]

223 kv_quant_mode: KVQuantMode = KVQuantMode.NONE

224 page_size_padded: int | None = None

225 indexes_kv_by_block_stride: bool = False

226 + num_head_slots: int | None = None

227 + """H of the logical ``[B, H, N, C]`` page when packing diverges from one

228 + slot per KV head. None means one slot per KV head. Published by the backend.

229 + """

230 + state_content_bytes: int | None = None

231 + """C in bytes when packed; None means dense K/V content."""

232 +

233 + def __post_init__(self):

234 + if self.head_size_v is None:

235 + object.__setattr__(self, "head_size_v", self.head_size)

236 +

237 + @property

238 + def num_heads(self) -> int:

239 + if self.num_head_slots is not None:

240 + return self.num_head_slots

241 + return self.num_kv_heads

242 +

243 + @property

244 + def state_content_size_bytes(self) -> int:

245 + """Bytes per (head slot, stored state) cell of the page."""

246 + if self.state_content_bytes is not None:

247 + return self.state_content_bytes

248 + return (self.head_size + self.head_size_v) * get_dtype_size(self.dtype)

249

250 @property

251 def unpadded_page_size_bytes(self) -> int:

219 - unpadded = self.real_page_size_bytes

220 - # Per-token-head scales are stored in separate tensors managed

221 - # by the attention backend, but the memory is carved from the

222 - # raw KV cache allocation so it must be budgeted here.

223 - if self.kv_quant_mode.is_per_token_head:

224 - unpadded += (

225 - 2 * self.block_size * self.num_kv_heads * get_dtype_size(torch.float32)

226 - )

227 - return unpadded

252 + return self.num_heads * self.storage_block_size * self.state_content_size_bytes

253

254 @property

255 def page_size_bytes(self) -> int:

@@ -235,20 +260,10 @@ class AttentionSpec(KVCacheSpec):

260

261 @property

262 def real_page_size_bytes(self) -> int:

238 - if self.kv_quant_mode.is_nvfp4:

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

240 - head_dim = nvfp4_kv_cache_full_dim(self.head_size)

241 - elif self.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD:

242 - head_dim = self.head_size // 2

243 - else:

244 - head_dim = self.head_size

245 - return (

246 - 2

247 - * self.block_size

248 - * self.num_kv_heads

249 - * head_dim

250 - * get_dtype_size(self.dtype)

251 - )

263 + """Alias of ``unpadded_page_size_bytes``

264 + TODO(lucas): follow up with TPU backend to see if we can remove this property.

265 + """

266 + return self.unpadded_page_size_bytes

267

268 def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:

269 parallel_config = vllm_config.parallel_config

@@ -267,8 +282,6 @@ class FullAttentionSpec(AttentionSpec):

282 In this case, we use FullAttentionSpec and record the sliding window size.

283 """

284

270 - head_size_v: int = None # type: ignore[assignment]

271 -

285 sliding_window: int | None = None

286 """

287 Default to None for not using sliding window attention.

@@ -284,10 +297,6 @@ class FullAttentionSpec(AttentionSpec):

297 cache layout itself.

298 """

299

287 - def __post_init__(self):

288 - if self.head_size_v is None:

289 - object.__setattr__(self, "head_size_v", self.head_size)

290 -

300 def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:

301 max_model_len = vllm_config.model_config.max_model_len

302 dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size

@@ -337,6 +346,8 @@ class FullAttentionSpec(AttentionSpec):

346 kv_quant_mode=specs[0].kv_quant_mode,

347 page_size_padded=specs[0].page_size_padded,

348 indexes_kv_by_block_stride=specs[0].indexes_kv_by_block_stride,

349 + num_head_slots=specs[0].num_head_slots,

350 + state_content_bytes=specs[0].state_content_bytes,

351 sliding_window=cls.merge_window_sizes(sliding_window),

352 attention_chunk_size=cls.merge_window_sizes(attention_chunk_size),

353 # If any layer in the group is non-causal, treat the group as

@@ -357,23 +368,6 @@ class FullAttentionSpec(AttentionSpec):

368 )

369 return merged_spec

370

360 - @property

361 - def real_page_size_bytes(self) -> int:

362 - if self.kv_quant_mode.is_nvfp4:

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

364 - # fp4 data: head_size//2 bytes (2 fp4 values per byte)

365 - # fp8 block scale: head_size//16 bytes (1 scale per 16 elements)

366 - last_dim = nvfp4_kv_cache_full_dim(

367 - self.head_size

368 - ) + nvfp4_kv_cache_full_dim(self.head_size_v)

369 - elif self.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD:

370 - last_dim = self.head_size // 2 + self.head_size_v // 2

371 - else:

372 - last_dim = self.head_size + self.head_size_v

373 - return (

374 - self.block_size * self.num_kv_heads * last_dim * get_dtype_size(self.dtype)

375 - )

376 -

371

372 def _apply_alignment_padding(spec: MLAAttentionSpec | SlidingWindowMLASpec):

373 if spec.alignment is None:

@@ -384,32 +378,6 @@ def _apply_alignment_padding(spec: MLAAttentionSpec | SlidingWindowMLASpec):

378 object.__setattr__(spec, "page_size_padded", padded_page_size)

379

380

387 -@dataclass(frozen=True, kw_only=True)

388 -class TQFullAttentionSpec(FullAttentionSpec):

389 - """FullAttentionSpec with TQ-aware page size.

390 -

391 - Python equivalent of the C++ TQ4FullAttentionSpec. Overrides

392 - real_page_size_bytes to use TQ slot bytes instead of the raw

393 - head_size * dtype formula.

394 - """

395 -

396 - tq_slot_size: int = 0

397 -

398 - @property

399 - def real_page_size_bytes(self) -> int:

400 - if self.tq_slot_size > 0:

401 - return self.block_size * self.num_kv_heads * self.tq_slot_size

402 - return super().real_page_size_bytes

403 -

404 - @classmethod

405 - def merge(cls, specs: list[Self]) -> Self:

406 - merged = super().merge(specs)

407 - assert all(s.tq_slot_size == specs[0].tq_slot_size for s in specs), (

408 - "All TQ layers in the same KV cache group must use the same tq_slot_size."

409 - )

410 - return replace(merged, tq_slot_size=specs[0].tq_slot_size)

411 -

412 -

381 @dataclass(frozen=True, kw_only=True)

382 class MLAAttentionSpec(FullAttentionSpec):

383 # TODO(Lucas/Chen): less hacky way to do this

@@ -420,6 +388,8 @@ class MLAAttentionSpec(FullAttentionSpec):

388 model_version: str | None = None

389 # Marks draft groups that flatten a non-causal query block into decode rows.

390 non_causal_multi_token_decode: bool = False

391 + # MLA stores a single latent vector per state; there is no separate V.

392 + head_size_v: int = 0

393

394 def __post_init__(self):

395 super().__post_init__()

@@ -429,27 +399,6 @@ class MLAAttentionSpec(FullAttentionSpec):

399 def storage_block_size(self) -> int:

400 return self.block_size // self.compress_ratio

401

432 - @property

433 - def real_page_size_bytes(self) -> int:

434 - if self.cache_dtype_str == "fp8_ds_mla":

435 - if self.model_version == "deepseek_v4":

436 - # DeepseekV4: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B per token.

437 - # head_size stays semantic (512); bytes are determined here.

438 - return self.storage_block_size * 584

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

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

441 - return self.block_size * 656

442 - if self.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD:

443 - head_dim = self.head_size // 2

444 - else:

445 - head_dim = self.head_size

446 - return (

447 - self.storage_block_size

448 - * self.num_kv_heads

449 - * head_dim

450 - * get_dtype_size(self.dtype)

451 - )

452 -

402 @classmethod

403 def merge(cls, specs: list[Self]) -> Self:

404 assert all(isinstance(spec, MLAAttentionSpec) for spec in specs), (

@@ -476,6 +425,8 @@ class MLAAttentionSpec(FullAttentionSpec):

425 dtype=specs[0].dtype,

426 kv_quant_mode=specs[0].kv_quant_mode,

427 page_size_padded=specs[0].page_size_padded,

428 + num_head_slots=specs[0].num_head_slots,

429 + state_content_bytes=specs[0].state_content_bytes,

430 indexes_kv_by_block_stride=block_stride_set.pop(),

431 cache_dtype_str=cache_dtype_str_set.pop(),

432 compress_ratio=compress_ratio_set.pop(),

@@ -533,6 +484,8 @@ class RSWASpec(FullAttentionSpec):

484 kv_quant_mode=base.kv_quant_mode,

485 page_size_padded=base.page_size_padded,

486 indexes_kv_by_block_stride=base.indexes_kv_by_block_stride,

487 + num_head_slots=base.num_head_slots,

488 + state_content_bytes=base.state_content_bytes,

489 sliding_window=base.sliding_window,

490 attention_chunk_size=base.attention_chunk_size,

491 non_causal=base.non_causal,

@@ -583,7 +536,6 @@ class ChunkedLocalAttentionSpec(AttentionSpec):

536 @dataclass(frozen=True, kw_only=True)

537 class SlidingWindowSpec(AttentionSpec):

538 sliding_window: int

586 - head_size_v: int = None # type: ignore[assignment]

539 # The trailing edge of the window is extended by ``extra_retained_tokens``

540 # so that those extra trailing tokens' blocks are retained (but not

541 # attended). This is needed for multi-module spec decoding which can

@@ -591,30 +543,6 @@ class SlidingWindowSpec(AttentionSpec):

543 # of the sequence, and thus needs to delay freeing/caching of blocks.

544 extra_retained_tokens: int = 0

545

594 - def __post_init__(self):

595 - if self.head_size_v is None:

596 - object.__setattr__(self, "head_size_v", self.head_size)

597 -

598 - @property

599 - def real_page_size_bytes(self) -> int:

600 - # Mirror ``FullAttentionSpec.real_page_size_bytes`` for NVFP4 KV cache.

601 - if self.kv_quant_mode.is_nvfp4:

602 - last_dim = nvfp4_kv_cache_full_dim(

603 - self.head_size

604 - ) + nvfp4_kv_cache_full_dim(self.head_size_v)

605 - return (

606 - self.block_size

607 - * self.num_kv_heads

608 - * last_dim

609 - * get_dtype_size(self.dtype)

610 - )

611 - return (

612 - self.block_size

613 - * self.num_kv_heads

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

615 - * get_dtype_size(self.dtype)

616 - )

617 -

546 def max_admission_blocks_per_request(

547 self, max_in_flight_tokens: int, max_model_len: int

548 ) -> int:

@@ -672,31 +600,20 @@ class SlidingWindowMLASpec(SlidingWindowSpec):

600 alignment: int | None = None # Default to None for no padding.

601 compress_ratio: int = 1

602 model_version: str | None = None

603 + # MLA stores a single latent vector per state; there is no separate V.

604 + head_size_v: int = 0

605

606 def __post_init__(self):

607 + assert self.model_version in (None, "deepseek_v4"), (

608 + f"Unsupported model version: {self.model_version}"

609 + )

610 + super().__post_init__()

611 _apply_alignment_padding(self)

612

613 @property

614 def storage_block_size(self) -> int:

615 return self.block_size // self.compress_ratio

616

683 - @property

684 - def real_page_size_bytes(self) -> int:

685 - if self.model_version == "deepseek_v4" and self.cache_dtype_str == "fp8_ds_mla":

686 - # DeepseekV4 FlashMLA: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B

687 - # per token. FlashInfer's contiguous bf16/fp8 cache falls through to

688 - # the element-size formula below.

689 - return self.storage_block_size * 584

690 - assert self.model_version in (None, "deepseek_v4"), (

691 - f"Unsupported model version: {self.model_version}"

692 - )

693 - return (

694 - self.storage_block_size

695 - * self.num_kv_heads

696 - * self.head_size

697 - * get_dtype_size(self.dtype)

698 - )

699 -

617 @classmethod

618 def merge(cls, specs: list[Self]) -> Self:

619 assert all(isinstance(spec, SlidingWindowMLASpec) for spec in specs), (

@@ -727,6 +644,8 @@ class SlidingWindowMLASpec(SlidingWindowSpec):

644 head_size=specs[0].head_size,

645 dtype=specs[0].dtype,

646 page_size_padded=specs[0].page_size_padded,

647 + num_head_slots=specs[0].num_head_slots,

648 + state_content_bytes=specs[0].state_content_bytes,

649 indexes_kv_by_block_stride=block_stride_set.pop(),

650 sliding_window=sliding_window_set.pop(),

651 extra_retained_tokens=extra_retained_set.pop(),

@@ -853,6 +772,8 @@ class SinkFullAttentionSpec(FullAttentionSpec):

772 kv_quant_mode=specs[0].kv_quant_mode,

773 page_size_padded=specs[0].page_size_padded,

774 indexes_kv_by_block_stride=specs[0].indexes_kv_by_block_stride,

775 + num_head_slots=specs[0].num_head_slots,

776 + state_content_bytes=specs[0].state_content_bytes,

777 sliding_window=cls.merge_window_sizes(sliding_window),

778 attention_chunk_size=cls.merge_window_sizes(attention_chunk_size),

779 non_causal=any(spec.non_causal for spec in specs),

vllm/v1/worker/gpu/attn_utils.py

@@ -75,6 +75,7 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]:

75 with set_current_vllm_config(vllm_config):

76 indexes = backend.indexes_kv_by_block_stride()

77 spec = replace(spec, indexes_kv_by_block_stride=indexes)

78 + spec = backend.customize_spec(spec)

79 kv_cache_spec[layer_name] = spec

80 return kv_cache_spec

vllm/v1/worker/gpu_model_runner.py

@@ -7925,6 +7925,7 @@ class GPUModelRunner(

7925 with set_current_vllm_config(self.vllm_config):

7926 indexes = backend.indexes_kv_by_block_stride()

7927 spec = replace(spec, indexes_kv_by_block_stride=indexes)

7928 + spec = backend.customize_spec(spec)

7929 kv_cache_spec[layer_name] = spec

7930

7931 return kv_cache_spec