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 6, 15:45 UTC
[Misc] Add `float16` to `CacheDType` (#37199)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Matthew Bonanni committed
Mar 16, 2026 at 16:24 UTC
93f3c8e53157f55b45cb902bb12ba68bb69e062c
19 files changed
+55
-22
docs/design/attention_backends.md
+19
-19
@@ -164,18 +164,18 @@ Priority is **1 = highest** (tried first).
164
| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | MM Prefix | DCP | Attention Types | Compute Cap. |
165
| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | --------- | --- | --------------- | ------------ |
166
| `CPU_ATTN` | | fp16, bf16, fp32 | `auto` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | All | N/A |
167
-| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
168
-| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
169
-| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
170
-| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
171
-| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 |
167
+| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
168
+| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
169
+| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
170
+| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
171
+| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 |
172
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any |
173
-| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
174
-| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder, Enc-Dec | N/A |
173
+| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
174
+| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder, Enc-Dec | N/A |
175
| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ✅ | ❌ | All | N/A |
176
-| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ✅ | ✅ | ❌ | All | N/A |
177
-| `TREE_ATTN` | | fp16, bf16 | `auto` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any |
178
-| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | All | Any |
176
+| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ✅ | ✅ | ❌ | All | N/A |
177
+| `TREE_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any |
178
+| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | All | Any |
179
180
> **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`.
181
>
@@ -204,14 +204,14 @@ configuration.
204
205
| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
206
| ------- | ------ | --------- | ----------- | ---------- | ---- | ------ | --------- | --- | --------------- | ------------ |
207
-| `CUTLASS_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
208
-| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
209
-| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
210
-| `FLASHMLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
207
+| `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
208
+| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
209
+| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
210
+| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
211
| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
212
-| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
213
-| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
214
-| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
212
+| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
213
+| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
214
+| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
215
| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
216
-| `TRITON_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
217
-| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | Any | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | Any |
216
+| `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
217
+| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | Any |
vllm/config/cache.py
+1
@@ -13,6 +13,7 @@ logger = init_logger(__name__)
13
14
CacheDType = Literal[
15
"auto",
16
+ "float16",
17
"bfloat16",
18
"fp8",
19
"fp8_e4m3",
vllm/v1/attention/backend.py
+5
-1
@@ -51,7 +51,11 @@ 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", "bfloat16"]
54
+ supported_kv_cache_dtypes: ClassVar[list["CacheDType"]] = [
55
+ "auto",
56
+ "float16",
57
+ "bfloat16",
58
+ ]
59
60
# Does attention's forward() include kv cache update?
61
forward_includes_kv_cache_update: bool = True
vllm/v1/attention/backends/flash_attn.py
+6
-1
@@ -64,6 +64,11 @@ logger = init_logger(__name__)
64
class FlashAttentionBackend(AttentionBackend):
65
accept_output_buffer: bool = True
66
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
67
+ supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
68
+ "auto",
69
+ "float16",
70
+ "bfloat16",
71
+ ]
72
73
@staticmethod
74
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
@@ -164,7 +169,7 @@ class FlashAttentionBackend(AttentionBackend):
169
return True
170
if kv_cache_dtype.startswith("fp8"):
171
return flash_attn_supports_fp8()
167
- return kv_cache_dtype in ["auto", "bfloat16"]
172
+ return kv_cache_dtype in ["auto", "float16", "bfloat16"]
173
174
@classmethod
175
def supports_sink(cls) -> bool:
vllm/v1/attention/backends/flashinfer.py
+1
@@ -291,6 +291,7 @@ class FlashInferBackend(AttentionBackend):
291
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
292
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
293
"auto",
294
+ "float16",
295
"bfloat16",
296
"fp8",
297
"fp8_e4m3",
vllm/v1/attention/backends/flex_attention.py
+5
-1
@@ -80,7 +80,11 @@ class FlexAttentionBackend(AttentionBackend):
80
torch.bfloat16,
81
torch.float32,
82
]
83
- supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto", "bfloat16"]
83
+ supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
84
+ "auto",
85
+ "float16",
86
+ "bfloat16",
87
+ ]
88
89
forward_includes_kv_cache_update: bool = False
90
vllm/v1/attention/backends/mla/cutlass_mla.py
+1
@@ -39,6 +39,7 @@ class CutlassMLABackend(MLACommonBackend):
39
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
40
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
41
"auto",
42
+ "float16",
43
"bfloat16",
44
"fp8",
45
"fp8_e4m3",
vllm/v1/attention/backends/mla/flashattn_mla.py
+1
@@ -46,6 +46,7 @@ class FlashAttnMLABackend(MLACommonBackend):
46
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
47
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
48
"auto",
49
+ "float16",
50
"bfloat16",
51
]
52
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
+ "float16",
42
"bfloat16",
43
"fp8",
44
"fp8_e4m3",
vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
+1
@@ -62,6 +62,7 @@ class FlashInferMLASparseBackend(AttentionBackend):
62
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
63
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
64
"auto",
65
+ "float16",
66
"bfloat16",
67
"fp8",
68
"fp8_e4m3",
vllm/v1/attention/backends/mla/flashmla.py
+1
@@ -49,6 +49,7 @@ class FlashMLABackend(MLACommonBackend):
49
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
50
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
51
"auto",
52
+ "float16",
53
"bfloat16",
54
"fp8",
55
"fp8_e4m3",
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+1
@@ -26,6 +26,7 @@ class AiterMLABackend(MLACommonBackend):
26
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
27
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
28
"auto",
29
+ "float16",
30
"bfloat16",
31
"fp8",
32
"fp8_e4m3",
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+1
@@ -82,6 +82,7 @@ class ROCMAiterMLASparseBackend(AttentionBackend):
82
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
83
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
84
"auto",
85
+ "float16",
86
"bfloat16",
87
]
88
vllm/v1/attention/backends/mla/triton_mla.py
+1
@@ -31,6 +31,7 @@ class TritonMLABackend(MLACommonBackend):
31
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
32
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
33
"auto",
34
+ "float16",
35
"bfloat16",
36
"fp8",
37
"fp8_e4m3",
vllm/v1/attention/backends/mla/xpu_mla_sparse.py
+1
@@ -38,6 +38,7 @@ class XPUMLASparseBackend(AttentionBackend):
38
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
39
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
40
"auto",
41
+ "float16",
42
"bfloat16",
43
]
44
vllm/v1/attention/backends/rocm_aiter_fa.py
+1
@@ -736,6 +736,7 @@ class AiterFlashAttentionBackend(AttentionBackend):
736
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
737
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
738
"auto",
739
+ "float16",
740
"bfloat16",
741
"fp8",
742
"fp8_e4m3",
vllm/v1/attention/backends/rocm_attn.py
+1
@@ -166,6 +166,7 @@ class RocmAttentionBackend(AttentionBackend):
166
]
167
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
168
"auto",
169
+ "float16",
170
"bfloat16",
171
"fp8",
172
"fp8_e4m3",
vllm/v1/attention/backends/tree_attn.py
+6
@@ -10,6 +10,7 @@ import torch
10
11
from vllm import _custom_ops as ops
12
from vllm.config import VllmConfig
13
+from vllm.config.cache import CacheDType
14
from vllm.logger import init_logger
15
from vllm.v1.attention.backend import (
16
AttentionBackend,
@@ -31,6 +32,11 @@ logger = init_logger(__name__)
32
class TreeAttentionBackend(AttentionBackend):
33
accept_output_buffer: bool = True
34
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
35
+ supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
36
+ "auto",
37
+ "float16",
38
+ "bfloat16",
39
+ ]
40
forward_includes_kv_cache_update: bool = False
41
42
@staticmethod
vllm/v1/attention/backends/triton_attn.py
+1
@@ -263,6 +263,7 @@ class TritonAttentionBackend(AttentionBackend):
263
]
264
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
265
"auto",
266
+ "float16",
267
"bfloat16",
268
"fp8",
269
"fp8_e4m3",