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 15, 21:00 UTC
[V1][Bugfix] Standardize quantized kv cache rejection for attention backends (#14221)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Mar 6, 2025 at 17:18 UTC
6832707e90d460bc1d1eec550e0035af72db7a27
11 files changed
+59
-20
vllm/attention/backends/abstract.py
+4
@@ -294,3 +294,7 @@ class MLAAttentionImpl(AttentionImpl[T], Generic[T]):
294
output: Optional[torch.Tensor] = None,
295
) -> torch.Tensor:
296
raise NotImplementedError
297
+
298
+
299
+def is_quantized_kv_cache(kv_cache_dtype: str) -> bool:
300
+ return kv_cache_dtype != "auto"
vllm/attention/backends/flash_attn.py
+8
-1
@@ -8,11 +8,15 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type
8
import torch
9
10
from vllm import _custom_ops as ops
11
+# yapf conflicts with isort for this block
12
+# yapf: disable
13
from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
14
AttentionLayer,
15
AttentionMetadata,
16
AttentionMetadataBuilder,
15
- AttentionType)
17
+ AttentionType,
18
+ is_quantized_kv_cache)
19
+# yapf: enable
20
from vllm.attention.backends.utils import (
21
PAD_SLOT_ID, CommonAttentionState, compute_slot_mapping,
22
compute_slot_mapping_start_idx, get_flash_attn_version,
@@ -626,6 +630,9 @@ class FlashAttentionImpl(AttentionImpl):
630
self.sliding_window = ((sliding_window - 1,
631
0) if sliding_window is not None else (-1, -1))
632
self.kv_cache_dtype = kv_cache_dtype
633
+ if is_quantized_kv_cache(self.kv_cache_dtype):
634
+ raise NotImplementedError(
635
+ "FlashAttention with FP8 KV cache not yet supported")
636
if logits_soft_cap is None:
637
# In flash-attn, setting logits_soft_cap as 0 means no soft cap.
638
logits_soft_cap = 0
vllm/attention/backends/flashmla.py
+6
-3
@@ -6,7 +6,8 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type
6
7
import torch
8
9
-from vllm.attention.backends.abstract import AttentionType
9
+from vllm.attention.backends.abstract import (AttentionType,
10
+ is_quantized_kv_cache)
11
from vllm.attention.backends.mla.common import (MLACommonBackend,
12
MLACommonImpl,
13
MLACommonMetadata,
@@ -207,6 +208,10 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
208
"are not implemented for "
209
"FlashMLAImpl")
210
211
+ if is_quantized_kv_cache(self.kv_cache_dtype):
212
+ raise NotImplementedError(
213
+ "FlashMLA with FP8 KV cache not yet supported")
214
+
215
def _forward_decode(
216
self,
217
q_nope: torch.Tensor,
@@ -215,8 +220,6 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
220
attn_metadata: FlashMLAMetadata,
221
) -> torch.Tensor:
222
assert kv_c_and_k_pe_cache.numel() > 0
218
- if self.kv_cache_dtype.startswith("fp8"):
219
- raise NotImplementedError("FP8 FlashMLA not yet supported")
223
224
decode_meta = attn_metadata.decode_metadata
225
assert decode_meta is not None
vllm/attention/backends/hpu_attn.py
+6
-1
@@ -15,7 +15,8 @@ from vllm_hpu_extension.utils import (Matmul, ModuleFusedSDPA, Softmax,
15
16
from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
17
AttentionLayer,
18
- AttentionMetadata, AttentionType)
18
+ AttentionMetadata, AttentionType,
19
+ is_quantized_kv_cache)
20
from vllm.attention.backends.utils import CommonAttentionState
21
from vllm.attention.ops.hpu_paged_attn import (HPUPagedAttention,
22
HPUPagedAttentionMetadata)
@@ -158,6 +159,10 @@ class HPUAttentionImpl(AttentionImpl, torch.nn.Module):
159
"are not implemented for "
160
"HPUAttentionImpl")
161
162
+ if is_quantized_kv_cache(self.kv_cache_dtype):
163
+ raise NotImplementedError(
164
+ "HPUAttention with FP8 KV cache not yet supported")
165
+
166
def forward(
167
self,
168
layer: AttentionLayer,
vllm/attention/backends/ipex_attn.py
+3
-2
@@ -9,7 +9,8 @@ import torch
9
from vllm._ipex_ops import ipex_ops
10
from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
11
AttentionLayer,
12
- AttentionMetadata, AttentionType)
12
+ AttentionMetadata, AttentionType,
13
+ is_quantized_kv_cache)
14
from vllm.attention.backends.utils import CommonAttentionState
15
from vllm.attention.ops.paged_attn import (PagedAttention,
16
PagedAttentionMetadata)
@@ -145,7 +146,7 @@ class IpexAttnBackendImpl(AttentionImpl[IpexAttnMetadata]):
146
raise ValueError(
147
f"Head size {head_size} is not supported by PagedAttention. "
148
f"Supported head sizes are: {supported_head_sizes}.")
148
- if kv_cache_dtype != "auto":
149
+ if is_quantized_kv_cache(kv_cache_dtype):
150
raise NotImplementedError(
151
"IPEX backend does not support FP8 KV cache. "
152
"Please use xFormers backend instead.")
vllm/attention/backends/pallas.py
+3
-2
@@ -8,7 +8,8 @@ import torch_xla.experimental.custom_kernel # Required to register custom ops.
8
9
from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
10
AttentionLayer,
11
- AttentionMetadata, AttentionType)
11
+ AttentionMetadata, AttentionType,
12
+ is_quantized_kv_cache)
13
from vllm.attention.backends.utils import CommonAttentionState
14
15
@@ -119,7 +120,7 @@ class PallasAttentionBackendImpl(AttentionImpl):
120
raise NotImplementedError("Alibi slopes is not supported.")
121
if sliding_window is not None:
122
raise NotImplementedError("Sliding window is not supported.")
122
- if kv_cache_dtype != "auto":
123
+ if is_quantized_kv_cache(kv_cache_dtype):
124
raise NotImplementedError("FP8 KV cache dtype is not supported.")
125
if blocksparse_params is not None:
126
raise NotImplementedError("Blocksparse is not supported.")
vllm/attention/backends/torch_sdpa.py
+6
-2
@@ -7,11 +7,15 @@ from typing import Any, Dict, List, Optional, Tuple, Type
7
import torch
8
from torch.nn.functional import scaled_dot_product_attention
9
10
+# yapf conflicts with isort for this block
11
+# yapf: disable
12
from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
13
AttentionLayer,
14
AttentionMetadata,
15
AttentionMetadataBuilder,
14
- AttentionType)
16
+ AttentionType,
17
+ is_quantized_kv_cache)
18
+# yapf: enable
19
from vllm.attention.backends.utils import CommonAttentionState
20
from vllm.attention.ops.ipex_attn import PagedAttention
21
from vllm.attention.ops.paged_attn import PagedAttentionMetadata
@@ -427,7 +431,7 @@ class TorchSDPABackendImpl(AttentionImpl[TorchSDPAMetadata]):
431
raise ValueError(
432
f"Head size {head_size} is not supported by PagedAttention. "
433
f"Supported head sizes are: {supported_head_sizes}.")
430
- if kv_cache_dtype != "auto":
434
+ if is_quantized_kv_cache(kv_cache_dtype):
435
raise NotImplementedError(
436
"Torch SDPA backend does not support FP8 KV cache. "
437
"Please use xFormers backend instead.")
vllm/attention/backends/triton_mla.py
+6
-3
@@ -4,7 +4,8 @@ from typing import Any, Dict, List, Optional, Type
4
5
import torch
6
7
-from vllm.attention.backends.abstract import AttentionType
7
+from vllm.attention.backends.abstract import (AttentionType,
8
+ is_quantized_kv_cache)
9
from vllm.attention.backends.mla.common import (MLACommonBackend,
10
MLACommonImpl,
11
MLACommonMetadata)
@@ -58,6 +59,10 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
59
"are not implemented for "
60
"TritonMLAImpl")
61
62
+ if is_quantized_kv_cache(self.kv_cache_dtype):
63
+ raise NotImplementedError(
64
+ "TritonMLA with FP8 KV cache not yet supported")
65
+
66
def _forward_decode(
67
self,
68
q_nope: torch.Tensor,
@@ -66,8 +71,6 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
71
attn_metadata: MLACommonMetadata,
72
) -> torch.Tensor:
73
assert kv_c_and_k_pe_cache.numel() > 0
69
- if self.kv_cache_dtype.startswith("fp8"):
70
- raise NotImplementedError("FP8 Triton MLA not yet supported")
74
75
decode_meta = attn_metadata.decode_metadata
76
assert decode_meta is not None
vllm/v1/attention/backends/flash_attn.py
+5
-1
@@ -7,7 +7,8 @@ import numpy as np
7
import torch
8
9
from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
10
- AttentionMetadata, AttentionType)
10
+ AttentionMetadata, AttentionType,
11
+ is_quantized_kv_cache)
12
from vllm.attention.backends.utils import get_flash_attn_version
13
from vllm.attention.ops.triton_merge_attn_states import merge_attn_states
14
from vllm.logger import init_logger
@@ -180,6 +181,9 @@ class FlashAttentionImpl(AttentionImpl):
181
else:
182
self.sliding_window = (sliding_window - 1, 0)
183
self.kv_cache_dtype = kv_cache_dtype
184
+ if is_quantized_kv_cache(self.kv_cache_dtype):
185
+ raise NotImplementedError(
186
+ "FlashAttention V1 with FP8 KV cache not yet supported")
187
if logits_soft_cap is None:
188
# In flash-attn, setting logits_soft_cap as 0 means no soft cap.
189
logits_soft_cap = 0
vllm/v1/attention/backends/mla/flashmla.py
+6
-4
@@ -5,7 +5,8 @@ from typing import Any, Optional
5
6
import torch
7
8
-from vllm.attention.backends.abstract import AttentionType
8
+from vllm.attention.backends.abstract import (AttentionType,
9
+ is_quantized_kv_cache)
10
from vllm.attention.ops.flashmla import (flash_mla_with_kvcache,
11
get_mla_metadata,
12
is_flashmla_supported)
@@ -115,6 +116,10 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
116
"are not implemented for "
117
"FlashMLAImpl")
118
119
+ if is_quantized_kv_cache(self.kv_cache_dtype):
120
+ raise NotImplementedError(
121
+ "FlashMLA V1 with FP8 KV cache not yet supported")
122
+
123
def _forward_decode(
124
self,
125
q_nope: torch.Tensor,
@@ -125,9 +130,6 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
130
assert kv_c_and_k_pe_cache.numel() > 0
131
assert attn_metadata.decode is not None
132
128
- if self.kv_cache_dtype.startswith("fp8"):
129
- raise NotImplementedError("FP8 FlashMLA not yet supported")
130
-
133
q = torch.cat([q_nope, q_pe], dim=-1)\
134
.unsqueeze(1) # Add seqlen dim of 1 (decode)
135
vllm/v1/attention/backends/mla/triton_mla.py
+6
-1
@@ -4,7 +4,8 @@ from typing import Any, Optional
4
5
import torch
6
7
-from vllm.attention.backends.abstract import AttentionType
7
+from vllm.attention.backends.abstract import (AttentionType,
8
+ is_quantized_kv_cache)
9
from vllm.attention.ops.triton_decode_attention import decode_attention_fwd
10
from vllm.logger import init_logger
11
from vllm.v1.attention.backends.mla.common import (MLACommonBackend,
@@ -61,6 +62,10 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
62
"are not implemented for "
63
"TritonMLAImpl")
64
65
+ if is_quantized_kv_cache(self.kv_cache_dtype):
66
+ raise NotImplementedError(
67
+ "TritonMLA V1 with FP8 KV cache not yet supported")
68
+
69
def _forward_decode(
70
self,
71
q_nope: torch.Tensor,