obscura/vllm
2
0
Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2026-08-19 12:10:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Attention] Add head_dim=512 support for FlashInfer trtllm attention backend (#38822)

Browse Source
This commit is contained in:
Duncan Moss
2026-05-22 20:27:35 -04:00
committed by GitHub
parent 6d30655b13
commit 552bbe6f4e
2 changed files with 13 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/design/attention_backends.md
+2 -2
View File
@@ -170,8 +170,8 @@ Priority is **1 = highest** (tried first).
| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | ❌ | All | N/A |
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x |
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64 | 64, 128, 256, 512 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x |
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
vllm/v1/attention/backends/flashinfer.py
+11 -10
View File
@@ -402,7 +402,7 @@ class FlashInferBackend(AttentionBackend):
@classmethod
def get_supported_head_sizes(cls) -> list[int]:
# https://github.com/flashinfer-ai/flashinfer/blob/3d55c71a62052c590c130897d3a3db49b14fcc34/include/flashinfer/utils.cuh#L157
return [64, 128, 256]
return [64, 128, 256, 512]
@classmethod
def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
@@ -1451,15 +1451,16 @@ class FlashInferImpl(AttentionImpl):
num_actual_tokens = attn_metadata.num_actual_tokens
# The FlashInfer api requires data to be in fp8_e4m3 or fp8_e5m2
# to process the cache when the kv_cache_dtype is fp8
if self.kv_sharing_target_layer_name is None and is_quantized_kv_cache(
self.kv_cache_dtype
):
torch_dtype = FlashInferBackend.get_dtype_for_flashinfer(
self.kv_cache_dtype
)
kv_cache = kv_cache.view(torch_dtype)
# FlashInfer treats uint8 KV cache as NVFP4. vLLM stores FP8 KV cache
# as uint8 bytes, so pass FP8 caches with their logical dtype.
if not self.is_kvcache_nvfp4 and kv_cache.dtype == torch.uint8:
fp8_view_dtype = None
if self.kv_cache_dtype in ("fp8", "fp8_e4m3", torch.float8_e4m3fn):
fp8_view_dtype = torch.float8_e4m3fn
elif self.kv_cache_dtype in ("fp8_e5m2", torch.float8_e5m2):
fp8_view_dtype = torch.float8_e5m2
if fp8_view_dtype is not None:
kv_cache = kv_cache.view(fp8_view_dtype)
# Inputs and outputs may be padded for CUDA graphs
query = query[:num_actual_tokens]