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 9, 18:15 UTC
[Bugfix][Attention] Explicitly report support for kv_cache_dtype bfloat16 (#32795)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Matthew Bonanni committed
Jan 22, 2026 at 14:05 UTC
955b43a5a5142046bd8a1bd22fd219bf944a6265
13 files changed
+31
-11
vllm/model_executor/layers/quantization/kv_cache.py
+6
-2
@@ -9,6 +9,7 @@ from vllm.model_executor.layers.quantization.base_config import (
9
QuantizeMethodBase,
10
)
11
from vllm.platforms import current_platform
12
+from vllm.v1.attention.backend import is_quantized_kv_cache
13
14
logger = init_logger(__name__)
15
@@ -52,11 +53,14 @@ class BaseKVCacheMethod(QuantizeMethodBase):
53
assert not hasattr(layer, "prob_scale")
54
return
55
55
- # If the kv-cache dtype is auto, we enforce the k/v_scale to be 1.0
56
+ # If the kv-cache is not quantized, we enforce the k/v_scale to be 1.0
57
# regardless whether the kv-scale is available in the checkpoint.
58
# No need to process kv scales after loading if we are going to
59
# calculate them on the fly.
59
- if layer.kv_cache_dtype != "auto" and not layer.calculate_kv_scales:
60
+ if (
61
+ is_quantized_kv_cache(layer.kv_cache_dtype)
62
+ and not layer.calculate_kv_scales
63
+ ):
64
if layer.k_scale > 0.0 and layer.v_scale > 0.0:
65
# We prefer to use separate k_scale and v_scale if present
66
k_scale = layer.k_scale.to("cpu").tolist()
vllm/platforms/cpu.py
+3
-2
@@ -16,6 +16,7 @@ import torch
16
17
from vllm import envs
18
from vllm.logger import init_logger
19
+from vllm.v1.attention.backend import is_quantized_kv_cache
20
from vllm.v1.attention.backends.registry import AttentionBackendEnum
21
22
from .interface import CpuArchEnum, Platform, PlatformEnum
@@ -198,13 +199,13 @@ class CpuPlatform(Platform):
199
if (
200
scheduler_config.enable_chunked_prefill
201
or cache_config.enable_prefix_caching
201
- ) and cache_config.cache_dtype != "auto":
202
+ ) and is_quantized_kv_cache(cache_config.cache_dtype):
203
raise RuntimeError(
204
"Chunked-prefill and prefix-cache on the CPU "
205
"backend is not compatible with FP8 KV cache."
206
)
207
207
- if cache_config.cache_dtype != "auto":
208
+ if cache_config.cache_dtype.startswith("fp8"):
209
logger.warning(
210
"CPU backend doesn't support KV cache quantization fallback to auto."
211
)
vllm/v1/attention/backend.py
+2
-2
@@ -51,7 +51,7 @@ class AttentionBackend(ABC):
51
# makes sure the output tensor is allocated inside the cudagraph.
52
accept_output_buffer: bool = False
53
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
54
- supported_kv_cache_dtypes: ClassVar[list["CacheDType"]] = ["auto"]
54
+ supported_kv_cache_dtypes: ClassVar[list["CacheDType"]] = ["auto", "bfloat16"]
55
56
@staticmethod
57
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
@@ -747,7 +747,7 @@ class MLAAttentionImpl(AttentionImpl[T], Generic[T]):
747
748
749
def is_quantized_kv_cache(kv_cache_dtype: str) -> bool:
750
- return kv_cache_dtype != "auto"
750
+ return kv_cache_dtype.startswith("fp8")
751
752
753
def subclass_attention_backend(
vllm/v1/attention/backends/flash_attn.py
+1
-1
@@ -151,7 +151,7 @@ class FlashAttentionBackend(AttentionBackend):
151
return True
152
if kv_cache_dtype.startswith("fp8"):
153
return flash_attn_supports_fp8()
154
- return kv_cache_dtype in ["auto"]
154
+ return kv_cache_dtype in ["auto", "bfloat16"]
155
156
@classmethod
157
def supports_sink(cls) -> bool:
vllm/v1/attention/backends/flashinfer.py
+1
@@ -281,6 +281,7 @@ class FlashInferBackend(AttentionBackend):
281
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
282
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
283
"auto",
284
+ "bfloat16",
285
"fp8",
286
"fp8_e4m3",
287
"fp8_e5m2",
vllm/v1/attention/backends/flex_attention.py
+1
-1
@@ -80,7 +80,7 @@ class FlexAttentionBackend(AttentionBackend):
80
torch.bfloat16,
81
torch.float32,
82
]
83
- supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto"]
83
+ supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto", "bfloat16"]
84
85
@staticmethod
86
def get_name() -> str:
vllm/v1/attention/backends/mla/cutlass_mla.py
+1
@@ -38,6 +38,7 @@ class CutlassMLABackend(MLACommonBackend):
38
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
39
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
40
"auto",
41
+ "bfloat16",
42
"fp8",
43
"fp8_e4m3",
44
]
vllm/v1/attention/backends/mla/flashattn_mla.py
+4
-1
@@ -43,7 +43,10 @@ logger = init_logger(__name__)
43
44
class FlashAttnMLABackend(MLACommonBackend):
45
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
46
- supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto"]
46
+ supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
47
+ "auto",
48
+ "bfloat16",
49
+ ]
50
51
@staticmethod
52
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
vllm/v1/attention/backends/mla/flashinfer_mla.py
+1
@@ -38,6 +38,7 @@ class FlashInferMLABackend(MLACommonBackend):
38
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
39
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
40
"auto",
41
+ "bfloat16",
42
"fp8",
43
"fp8_e4m3",
44
]
vllm/v1/attention/backends/mla/flashmla.py
+1
@@ -48,6 +48,7 @@ class FlashMLABackend(MLACommonBackend):
48
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
49
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
50
"auto",
51
+ "bfloat16",
52
"fp8",
53
"fp8_e4m3",
54
]
vllm/v1/attention/backends/mla/flashmla_sparse.py
+5
-1
@@ -76,7 +76,11 @@ structured as:
76
class FlashMLASparseBackend(AttentionBackend):
77
accept_output_buffer: bool = True
78
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
79
- supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto", "fp8_ds_mla"]
79
+ supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
80
+ "auto",
81
+ "bfloat16",
82
+ "fp8_ds_mla",
83
+ ]
84
85
@staticmethod
86
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
vllm/v1/attention/backends/mla/triton_mla.py
+4
-1
@@ -28,7 +28,10 @@ logger = init_logger(__name__)
28
29
class TritonMLABackend(MLACommonBackend):
30
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
31
- supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto"]
31
+ supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
32
+ "auto",
33
+ "bfloat16",
34
+ ]
35
36
@staticmethod
37
def get_name() -> str:
vllm/v1/attention/backends/triton_attn.py
+1
@@ -259,6 +259,7 @@ class TritonAttentionBackend(AttentionBackend):
259
]
260
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
261
"auto",
262
+ "bfloat16",
263
"fp8",
264
"fp8_e4m3",
265
"fp8_e5m2",