A high-throughput and memory-efficient inference and serving engine for LLMs
[DSV4][Minor] Fix supported KV cache dtypes (#44892)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Jun 15, 2026 at 17:14 UTC
f4359a70f9e04b0223ef9209db6f0d4d6a10f094
3 files changed
+8
-8
docs/design/attention_backends.md
+2
-2
@@ -252,6 +252,6 @@ default on NVIDIA is `FLASHMLA_SPARSE_DSV4`.
252
253
| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
254
| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ |
255
-| `FLASHINFER_MLA_SPARSE_DSV4` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
256
-| `FLASHMLA_SPARSE_DSV4` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla`, `fp8` | 256 | 512 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
255
+| `FLASHINFER_MLA_SPARSE_DSV4` | bf16 | `auto`, `bfloat16`, `fp8` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
256
+| `FLASHMLA_SPARSE_DSV4` | bf16 | `auto`, `fp8_ds_mla`, `fp8` | 256 | 512 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
257
| `ROCM_FLASHMLA_SPARSE_DSV4` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py
+6
-5
@@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, ClassVar, cast
13
14
import torch
15
16
+from vllm.config.cache import CacheDType
17
from vllm.forward_context import get_forward_context
18
from vllm.models.deepseek_v4.attention import DeepseekV4Attention
19
from vllm.models.deepseek_v4.common.ops import (
@@ -52,13 +53,13 @@ def _get_flashinfer_dsv4_workspace(device: torch.device) -> torch.Tensor:
53
class DeepseekV4FlashInferMLASparseBackend(DeepseekV4FlashMLABackend):
54
"""Shares the FlashMLA V4 metadata/cache pipeline; swaps the attention impl.
55
55
- Inheriting from the FlashMLA V4 backend reuses its
56
- ``DeepseekV4FlashMLAMetadata`` builder (which the V4 sparse-index
57
- pipeline needs — the V3.2 FlashInfer builder lacks the ``c128a_*`` fields),
58
- 256-token blocks, head_size 512, and the (num_blocks, block_size, 512) cache
59
- shape for non-``fp8_ds_mla`` dtypes.
56
+ Inheriting from the FlashMLA V4 backend reuses its ``DeepseekV4FlashMLAMetadata``
57
+ builder.
58
"""
59
60
+ supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
61
+ supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto", "bfloat16", "fp8"]
62
+
63
@staticmethod
64
def get_name() -> str:
65
return "FLASHINFER_MLA_SPARSE_DSV4"
vllm/models/deepseek_v4/sparse_mla.py
-1
@@ -46,7 +46,6 @@ class DeepseekV4FlashMLABackend(AttentionBackend):
46
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
47
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
48
"auto",
49
- "bfloat16",
49
"fp8_ds_mla",
50
"fp8", # alias for fp8_ds_mla
51
]