Support nvfp4 kv with kv-cache-dtype-skip-layers sliding_window (#42890)
Signed-off-by: Shiyang Chen <shiychen@nvidia.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
sychen52 committed
Jul 3, 2026 at 19:29 UTC
67ff0ae30fe6b1ab1a912e10977d99ddb169c4b2
6 files changed
+227
-4
docs/features/quantization/quantized_kvcache.md
+26
@@ -49,6 +49,32 @@ You can configure how the quantization scales are computed in vLLM using three d
49
- `kv_cache_dtype="fp8_e4m3"`: Supported on CUDA 11.8+ and ROCm (AMD GPUs)
50
- `kv_cache_dtype="fp8_e5m2"`: Supported on CUDA 11.8+
51
52
+### Skipping Specific Layers from KV-Cache Quantization
53
+
54
+Some attention layer types (e.g. sliding-window) are more sensitive to KV-cache quantization. The `--kv-cache-dtype-skip-layers` flag leaves the specified layers at the model's native dtype while keeping the rest of the layers under the chosen quantized dtype. The flag accepts either layer indices or layer-type names:
55
+
56
+```bash
57
+# Skip every sliding-window attention layer.
58
+vllm serve <model> \
59
+ --kv-cache-dtype fp8 \
60
+ --kv-cache-dtype-skip-layers sliding_window
61
+
62
+# Skip specific layer indices.
63
+vllm serve <model> \
64
+ --kv-cache-dtype fp8 \
65
+ --kv-cache-dtype-skip-layers 0 1 23
66
+```
67
+
68
+Programmatic usage:
69
+
70
+```python
71
+llm = LLM(
72
+ model="meta-llama/Llama-3.1-8B-Instruct",
73
+ kv_cache_dtype="fp8",
74
+ kv_cache_dtype_skip_layers=["sliding_window"],
75
+)
76
+```
77
+
78
---
79
80
## Examples
vllm/config/cache.py
+6
@@ -119,6 +119,11 @@ class CacheConfig:
119
mamba_page_size_padded: int | None = None
120
""" Optional override for mamba page size; used by hybrid mamba/attention
121
models to ensure exact alignment with attention page size."""
122
+ skip_page_size_padded: int | None = None
123
+ """Optional override for the page size of layers skipped from KV cache
124
+ quantization (``--kv-cache-dtype-skip-layers``); set during block-size
125
+ alignment so unquantized skip layers pad up to the quantized primary's
126
+ page."""
127
mamba_block_size: int | None = Field(default=None, gt=0)
128
"""Size of a contiguous cache block in number of tokens for mamba cache.
129
Can be set only when prefix caching is enabled.
@@ -207,6 +212,7 @@ class CacheConfig:
212
# Prefix-caching implementation detail (doesn't affect compiled graph).
213
"hash_block_size",
214
"mamba_page_size_padded",
215
+ "skip_page_size_padded",
216
"user_specified_block_size",
217
"user_specified_mamba_block_size",
218
"_block_size_resolved",
vllm/model_executor/layers/attention/attention.py
+51
-1
@@ -92,6 +92,35 @@ def should_load_quant_weights(quant_method: QuantizeMethodBase | None) -> bool:
92
)
93
94
95
+def _largest_kernel_block_within(
96
+ attn_backend: "type[AttentionBackend]",
97
+ per_token_bytes: int,
98
+ page_budget: int | None,
99
+ fallback: int,
100
+) -> int:
101
+ """Largest supported kernel block size whose page fits in ``page_budget``.
102
+
103
+ A padded spec (e.g. skip-quant layer) that pads its page up to a large shared page
104
+ wastes ``page_budget - block*per_token`` bytes per block. Picking the largest kernel
105
+ block whose natural page still fits under ``page_budget`` minimizes that waste.
106
+ Falls back to the smallest supported block when ``page_budget`` is None (no padding
107
+ — the block is handled by ``unify``'s integer scaling instead) or nothing fits.
108
+ """
109
+ from vllm.v1.attention.backend import MultipleOf
110
+
111
+ sizes = attn_backend.get_supported_kernel_block_sizes()
112
+ candidates = [s for s in sizes if isinstance(s, int)]
113
+ if not candidates:
114
+ candidates = [s.base for s in sizes if isinstance(s, MultipleOf)]
115
+ if not candidates:
116
+ return fallback
117
+ smallest = min(candidates)
118
+ if not page_budget or per_token_bytes <= 0:
119
+ return smallest
120
+ fitting = [b for b in candidates if b * per_token_bytes <= page_budget]
121
+ return max(fitting) if fitting else smallest
122
+
123
+
124
def set_default_quant_scales(layer: nn.Module, register_buffer: bool = False) -> None:
125
"""Sets default quantization scales for the layer."""
126
if register_buffer:
@@ -601,14 +630,35 @@ class Attention(nn.Module, AttentionLayerBase):
630
assert not vllm_config.model_config.use_mla, (
631
"MLA is not supported for slidingwindow"
632
)
633
+ # SW chooses its own block_size, decoupled from the user's
634
+ # ``--block-size`` (which only constrains primary attention).
635
+ # When this SW layer is a padded spec (skip-quant: its page is
636
+ # padded up to ``skip_page_size_padded``), pick the largest kernel
637
+ # block that still fits the shared page so we waste fewer padding
638
+ # bytes per block. Otherwise (page_size_padded is None) the smallest
639
+ # block is fine — ``unify`` scales it up by an integer ratio.
640
+ shared_page = vllm_config.cache_config.skip_page_size_padded
641
+ sw_per_token = SlidingWindowSpec(
642
+ block_size=1,
643
+ num_kv_heads=self.num_kv_heads,
644
+ head_size=self.head_size,
645
+ head_size_v=self.head_size_v,
646
+ dtype=self.kv_cache_torch_dtype,
647
+ kv_quant_mode=quant_mode,
648
+ sliding_window=self.sliding_window,
649
+ ).real_page_size_bytes
650
+ sw_block_size = _largest_kernel_block_within(
651
+ self.attn_backend, sw_per_token, shared_page, block_size
652
+ )
653
return SlidingWindowSpec(
605
- block_size=block_size,
654
+ block_size=sw_block_size,
655
num_kv_heads=self.num_kv_heads,
656
head_size=self.head_size,
657
head_size_v=self.head_size_v,
658
dtype=self.kv_cache_torch_dtype,
659
kv_quant_mode=quant_mode,
660
sliding_window=self.sliding_window,
661
+ page_size_padded=shared_page,
662
)
663
elif self.kv_cache_dtype.startswith("turboquant_"):
664
from vllm.model_executor.layers.quantization.turboquant.config import (
vllm/platforms/interface.py
+117
@@ -629,6 +629,123 @@ class Platform:
629
if model_config.is_hybrid:
630
cls._align_hybrid_block_size(vllm_config, backend_cls)
631
632
+ # Phase 3: Align block/page sizes when multiple KV dtypes share the
633
+ # block pool (e.g. nvfp4 primary + unquantized skip layers).
634
+ # May override the user's --block-size.
635
+ if cache_config.kv_cache_dtype_skip_layers:
636
+ cls._align_heterogeneous_kv_block_size(vllm_config, backend_cls)
637
+
638
+ @classmethod
639
+ def _align_heterogeneous_kv_block_size(
640
+ cls,
641
+ vllm_config: "VllmConfig",
642
+ backend_cls: "type[AttentionBackend]",
643
+ ) -> None:
644
+ """Align block size when several KV dtypes share one block pool.
645
+
646
+ A quantized primary (e.g. nvfp4) shares the block pool with one or more
647
+ higher-precision "padded specs" (skip layers today; the first/last-N
648
+ sibling in the future). A padded spec's per-token page is larger than
649
+ the primary's *and not an integer multiple of it*, so the trivial
650
+ ``unify_kv_cache_spec_page_size`` cannot reconcile them. We do it here
651
+ instead, before the specs are built:
652
+
653
+ 1. Bump the primary ``block_size`` (kernel-aligned) until the primary
654
+ page is large enough to cover the largest padded-spec page.
655
+ 2. Record that shared page in each padded spec's ``*_page_size_padded``
656
+ hint, so it pads up to the shared page.
657
+
658
+ ``unify`` then sees equal pages and stays trivial.
659
+
660
+ To add a padded-spec type: append its per-token page to ``padded_pages``
661
+ and set its ``*_page_size_padded`` hint below.
662
+ """
663
+ from vllm.config.vllm import set_current_vllm_config
664
+ from vllm.utils.math_utils import cdiv
665
+ from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
666
+ from vllm.v1.attention.backend import MultipleOf
667
+ from vllm.v1.kv_cache_interface import FullAttentionSpec, get_kv_quant_mode
668
+
669
+ cache_config = vllm_config.cache_config
670
+ model_config = vllm_config.model_config
671
+ parallel_config = vllm_config.parallel_config
672
+ if not model_config:
673
+ return
674
+
675
+ def per_token_page_bytes(dtype: "torch.dtype", cache_dtype: str) -> int:
676
+ """Bytes one token occupies in one layer, for the given dtype."""
677
+ return FullAttentionSpec(
678
+ block_size=1,
679
+ num_kv_heads=model_config.get_num_kv_heads(parallel_config),
680
+ head_size=model_config.get_head_size(),
681
+ dtype=dtype,
682
+ kv_quant_mode=get_kv_quant_mode(cache_dtype),
683
+ ).page_size_bytes
684
+
685
+ primary_dtype = (
686
+ STR_DTYPE_TO_TORCH_DTYPE[cache_config.cache_dtype]
687
+ if cache_config.cache_dtype != "auto"
688
+ else model_config.dtype
689
+ )
690
+ primary_page = per_token_page_bytes(primary_dtype, cache_config.cache_dtype)
691
+
692
+ # Per-token page of every higher-precision padded spec sharing the pool.
693
+ padded_pages: list[int] = []
694
+ if cache_config.kv_cache_dtype_skip_layers:
695
+ padded_pages.append(per_token_page_bytes(model_config.dtype, "auto"))
696
+ # To add the first/last-N sibling:
697
+ # padded_pages.append(per_token_page_bytes(<sibling_dtype>, "auto"))
698
+ if not padded_pages:
699
+ return
700
+
701
+ largest_padded_page = max(padded_pages)
702
+ assert largest_padded_page >= primary_page, (
703
+ f"padded-spec per-token page ({largest_padded_page}B) < primary "
704
+ f"({primary_page}B); a higher-precision padded spec must not be "
705
+ "smaller than the quantized primary."
706
+ )
707
+ if largest_padded_page == primary_page:
708
+ # Pages already match per token; ``unify`` reconciles the differing
709
+ # block sizes by integer scaling, so no bump or padding is needed.
710
+ return
711
+
712
+ # Smallest block the kernel supports, and the granularity the primary
713
+ # block is rounded up to (never below the already-chosen block_size).
714
+ with set_current_vllm_config(vllm_config):
715
+ supported = backend_cls.get_supported_kernel_block_sizes()
716
+ smallest_kernel_block = min(
717
+ s.base if isinstance(s, MultipleOf) else s for s in supported
718
+ )
719
+ block_alignment = max(smallest_kernel_block, cache_config.block_size)
720
+
721
+ # Bytes one padded-spec page spans at its own smallest kernel block;
722
+ # also cover any mamba page a hybrid model already padded.
723
+ required_page = max(
724
+ largest_padded_page * smallest_kernel_block,
725
+ cache_config.mamba_page_size_padded or 0,
726
+ )
727
+
728
+ # Smallest kernel-aligned primary block whose page covers required_page.
729
+ primary_block_size = block_alignment * cdiv(
730
+ required_page, block_alignment * primary_page
731
+ )
732
+ if cache_config.block_size < primary_block_size:
733
+ cache_config.block_size = primary_block_size
734
+ logger.info(
735
+ "Setting attention block size to %d tokens so the quantized "
736
+ "primary KV page covers the higher-precision padded-spec page.",
737
+ primary_block_size,
738
+ )
739
+
740
+ # The shared page that every padded spec (and mamba) pads up to.
741
+ shared_page = cache_config.block_size * primary_page
742
+ if cache_config.kv_cache_dtype_skip_layers:
743
+ cache_config.skip_page_size_padded = shared_page
744
+ # To add the first/last-N sibling:
745
+ # cache_config.sibling_page_size_padded = shared_page
746
+ if cache_config.mamba_page_size_padded is not None:
747
+ cache_config.mamba_page_size_padded = shared_page
748
+
749
@classmethod
750
def _align_hybrid_block_size(
751
cls,
vllm/v1/worker/gpu/attn_utils.py
+18
-2
@@ -23,6 +23,7 @@ from vllm.v1.kv_cache_interface import (
23
AttentionSpec,
24
KVCacheConfig,
25
KVCacheSpec,
26
+ KVQuantMode,
27
MambaSpec,
28
UniformTypeKVCacheSpecs,
29
)
@@ -300,12 +301,20 @@ def _reshape_kv_cache(
301
kv_cache_spec.storage_block_size // kernel_block_size
302
)
303
kernel_num_blocks = num_blocks * num_blocks_per_kv_block
304
+ # Skipped layers (--kv-cache-dtype-skip-layers) keep the
305
+ # unquantized shape; only the quantized primary uses the
306
+ # quantized cache dtype's (possibly packed) layout.
307
+ layer_cache_dtype = (
308
+ "auto"
309
+ if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE
310
+ else cache_dtype
311
+ )
312
kv_cache_shape = group.backend.get_kv_cache_shape(
313
kernel_num_blocks,
314
kernel_block_size,
315
kv_cache_spec.num_kv_heads,
316
kv_cache_spec.head_size,
308
- cache_dtype_str=cache_dtype,
317
+ cache_dtype_str=layer_cache_dtype,
318
)
319
320
# FIXME(woosuk): Add kv_cache_stride_order to all attention backends.
@@ -377,11 +386,18 @@ def _update_hybrid_attention_layout(
386
kv_cache_spec = group.kv_cache_spec
387
if not isinstance(kv_cache_spec, AttentionSpec):
388
continue
389
+ # Mirror the per-layer dtype selection used when building the shape
390
+ # above. The block-dim index is dtype-independent for current backends
391
+ # (quantization only changes the last dim), so this is a no-op today,
392
+ # but it keeps both call sites consistent for skip layers.
393
+ layer_cache_dtype = (
394
+ "auto" if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE else cache_dtype
395
+ )
396
block_dim = group.backend.get_kv_cache_block_dim(
397
kernel_block_sizes[group.kv_cache_group_id],
398
kv_cache_spec.num_kv_heads,
399
kv_cache_spec.head_size,
384
- cache_dtype_str=cache_dtype,
400
+ cache_dtype_str=layer_cache_dtype,
401
)
402
# if the first dim of the kvcache's layout is already num_blocks, continue
403
if block_dim == 0:
vllm/v1/worker/gpu_model_runner.py
+9
-1
@@ -150,6 +150,7 @@ from vllm.v1.kv_cache_interface import (
150
KVCacheConfig,
151
KVCacheGroupSpec,
152
KVCacheSpec,
153
+ KVQuantMode,
154
MambaSpec,
155
SlidingWindowSpec,
156
UniformTypeKVCacheSpecs,
@@ -7147,12 +7148,19 @@ class GPUModelRunner(
7148
else:
7149
shape_block_size = kernel_block_size
7150
7151
+ # Skipped layers (--kv-cache-dtype-skip-layers) need
7152
+ # the unquantized shape.
7153
+ layer_cache_dtype_str = (
7154
+ "auto"
7155
+ if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE
7156
+ else self.cache_config.cache_dtype
7157
+ )
7158
kv_cache_shape = attn_backend.get_kv_cache_shape(
7159
kernel_num_blocks,
7160
shape_block_size,
7161
kv_cache_spec.num_kv_heads,
7162
kv_cache_spec.head_size,
7155
- cache_dtype_str=self.cache_config.cache_dtype,
7163
+ cache_dtype_str=layer_cache_dtype_str,
7164
)
7165
try:
7166
kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()