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 7, 20:15 UTC
[Attn,KV-cache] Use per-head scales in the attention selector (#34281)
Signed-off-by: Your Name <you@example.com> Signed-off-by: Eldar Kurtic <research@neuralmagic.com> Co-authored-by: Eldar Kurtic <research@neuralmagic.com> Co-authored-by: Your Name <you@example.com>
Eldar Kurtić committed
Feb 24, 2026 at 15:02 UTC
a87cc508599d2257f39420536dd5bb52aaa557c3
6 files changed
+80
-15
tests/kernels/attention/test_attention_selector.py
+54
@@ -291,3 +291,57 @@ def test_invalid_backend():
291
):
292
# Invalid backend name should raise ValueError when creating enum
293
AttentionConfig(backend=AttentionBackendEnum["INVALID"])
294
+
295
+
296
+@pytest.mark.parametrize(
297
+ "backend_name,flash_attn_version,should_succeed",
298
+ [
299
+ ("FLASH_ATTN", 3, True), # FA3 supports per-head quant scales
300
+ ("FLASH_ATTN", 2, False), # FA2 does not support per-head quant scales
301
+ ("FLASHINFER", None, False), # FlashInfer does not support
302
+ ("FLEX_ATTENTION", None, False), # Flex does not support
303
+ ],
304
+)
305
+def test_per_head_quant_scales_backend_selection(
306
+ backend_name: str, flash_attn_version: int | None, should_succeed: bool
307
+):
308
+ """Test backend selection when use_per_head_quant_scales=True."""
309
+ # Clear cache to ensure fresh backend selection
310
+ _cached_get_attn_backend.cache_clear()
311
+
312
+ attention_config = AttentionConfig(
313
+ backend=AttentionBackendEnum[backend_name],
314
+ flash_attn_version=flash_attn_version,
315
+ )
316
+ vllm_config = VllmConfig(attention_config=attention_config)
317
+
318
+ with (
319
+ set_current_vllm_config(vllm_config),
320
+ patch("vllm.platforms.current_platform", CudaPlatform()),
321
+ ):
322
+ if backend_name == "FLASH_ATTN" and flash_attn_version == 3:
323
+ if not torch.cuda.is_available():
324
+ pytest.skip("FA3 requires CUDA")
325
+ capability = torch.cuda.get_device_capability()
326
+ if capability[0] != 9:
327
+ pytest.skip("FA3 is only supported on Hopper (SM 9.x) GPUs")
328
+
329
+ if should_succeed:
330
+ backend = get_attn_backend(
331
+ head_size=128,
332
+ dtype=torch.float16,
333
+ kv_cache_dtype="fp8",
334
+ block_size=64,
335
+ use_per_head_quant_scales=True,
336
+ )
337
+ assert backend.get_name() == backend_name
338
+ else:
339
+ with pytest.raises(ValueError) as exc_info:
340
+ get_attn_backend(
341
+ head_size=128,
342
+ dtype=torch.float16,
343
+ kv_cache_dtype="fp8",
344
+ block_size=64,
345
+ use_per_head_quant_scales=True,
346
+ )
347
+ assert backend_name in str(exc_info.value)
vllm/model_executor/layers/attention/attention.py
+9
-1
@@ -229,13 +229,20 @@ class Attention(nn.Module, AttentionLayerBase):
229
calculate_kv_scales = False
230
231
# llm-compressor mdls need to set cache_dtype to "fp8" manually.
232
- if getattr(quant_config, "kv_cache_scheme", None) is not None:
232
+ kv_cache_scheme = getattr(quant_config, "kv_cache_scheme", None)
233
+ if kv_cache_scheme is not None:
234
kv_cache_dtype = "fp8"
235
calculate_kv_scales = False
236
if cache_config is not None:
237
cache_config.cache_dtype = "fp8"
238
cache_config.calculate_kv_scales = False
239
240
+ # Check if per-head quant scales are required based on kv_cache_scheme
241
+ use_per_head_quant_scales = (
242
+ kv_cache_scheme is not None
243
+ and kv_cache_scheme.get("strategy") == "attn_head"
244
+ )
245
+
246
self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype(
247
kv_cache_dtype, vllm_config.model_config
248
)
@@ -272,6 +279,7 @@ class Attention(nn.Module, AttentionLayerBase):
279
use_mla=False,
280
has_sink=self.has_sink,
281
use_mm_prefix=self.use_mm_prefix,
282
+ use_per_head_quant_scales=use_per_head_quant_scales,
283
attn_type=attn_type,
284
)
285
else:
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+1
-8
@@ -985,14 +985,7 @@ class CompressedTensorsKVCacheMethod(BaseKVCacheMethod):
985
self.quant_config.kv_cache_scheme["strategy"]
986
)
987
988
- if strategy == QuantizationStrategy.ATTN_HEAD:
989
- assert layer.impl.supports_per_head_quant_scales, (
990
- f"Layer {layer.__class__.__name__} with implementation "
991
- f"{layer.impl.__class__.__name__} does not support per-head scales."
992
- )
993
- n_scales = int(layer.num_kv_heads)
994
- else:
995
- n_scales = 1
988
+ n_scales = int(layer.num_kv_heads) if strategy == "attn_head" else 1
989
990
layer.k_scale = torch.nn.Parameter(
991
torch.ones(n_scales, requires_grad=False, dtype=torch.float32)
vllm/v1/attention/backend.py
+7
-1
@@ -187,6 +187,10 @@ class AttentionBackend(ABC):
187
def is_sparse(cls) -> bool:
188
return False
189
190
+ @classmethod
191
+ def supports_per_head_quant_scales(cls) -> bool:
192
+ return False
193
+
194
@classmethod
195
def supports_attn_type(cls, attn_type: str) -> bool:
196
"""Check if backend supports a given attention type.
@@ -225,6 +229,7 @@ class AttentionBackend(ABC):
229
has_sink: bool,
230
use_sparse: bool,
231
use_mm_prefix: bool,
232
+ use_per_head_quant_scales: bool,
233
device_capability: "DeviceCapability",
234
attn_type: str,
235
) -> list[str]:
@@ -253,6 +258,8 @@ class AttentionBackend(ABC):
258
invalid_reasons.append("sparse not supported")
259
else:
260
invalid_reasons.append("non-sparse not supported")
261
+ if use_per_head_quant_scales and not cls.supports_per_head_quant_scales():
262
+ invalid_reasons.append("per-head quant scales not supported")
263
if not cls.supports_compute_capability(device_capability):
264
invalid_reasons.append("compute capability not supported")
265
if not cls.supports_attn_type(attn_type):
@@ -635,7 +642,6 @@ class AttentionImplBase(ABC, Generic[T]):
642
# TODO add support to more backends:
643
# https://github.com/vllm-project/vllm/issues/25584
644
supports_quant_query_input: bool = False
638
- supports_per_head_quant_scales: bool = False
645
646
dcp_world_size: int
647
dcp_rank: int
vllm/v1/attention/backends/flash_attn.py
+5
-5
@@ -95,6 +95,11 @@ class FlashAttentionBackend(AttentionBackend):
95
AttentionType.ENCODER_DECODER,
96
)
97
98
+ @classmethod
99
+ def supports_per_head_quant_scales(cls) -> bool:
100
+ fa_version = get_flash_attn_version()
101
+ return fa_version is not None and fa_version >= 3
102
+
103
@staticmethod
104
def get_impl_cls() -> type["FlashAttentionImpl"]:
105
return FlashAttentionImpl
@@ -595,11 +600,6 @@ class FlashAttentionImpl(AttentionImpl):
600
)
601
602
self.supports_quant_query_input = True
598
- self.supports_per_head_quant_scales = (
599
- self.vllm_flash_attn_version >= 3
600
- if self.vllm_flash_attn_version is not None
601
- else False
602
- )
603
604
def forward(
605
self,
vllm/v1/attention/selector.py
+4
@@ -27,6 +27,7 @@ class AttentionSelectorConfig(NamedTuple):
27
has_sink: bool = False
28
use_sparse: bool = False
29
use_mm_prefix: bool = False
30
+ use_per_head_quant_scales: bool = False
31
attn_type: str = AttentionType.DECODER
32
33
def __repr__(self):
@@ -39,6 +40,7 @@ class AttentionSelectorConfig(NamedTuple):
40
f"has_sink={self.has_sink}, "
41
f"use_sparse={self.use_sparse}, "
42
f"use_mm_prefix={self.use_mm_prefix}, "
43
+ f"use_per_head_quant_scales={self.use_per_head_quant_scales}, "
44
f"attn_type={self.attn_type})"
45
)
46
@@ -52,6 +54,7 @@ def get_attn_backend(
54
has_sink: bool = False,
55
use_sparse: bool = False,
56
use_mm_prefix: bool = False,
57
+ use_per_head_quant_scales: bool = False,
58
attn_type: str | None = None,
59
num_heads: int | None = None,
60
) -> type[AttentionBackend]:
@@ -77,6 +80,7 @@ def get_attn_backend(
80
has_sink=has_sink,
81
use_sparse=use_sparse,
82
use_mm_prefix=use_mm_prefix,
83
+ use_per_head_quant_scales=use_per_head_quant_scales,
84
attn_type=attn_type or AttentionType.DECODER,
85
)
86