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 3, 00:30 UTC
[Attention] Add head_dim=512 support for FlashInfer trtllm attention backend (#38822)
Duncan Moss committed
May 22, 2026 at 17:27 UTC
552bbe6f4e6da23def2adf9fbc39c3ad0250f1e1
2 files changed
+13
-12
docs/design/attention_backends.md
+2
-2
@@ -170,8 +170,8 @@ Priority is **1 = highest** (tried first).
170
| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
171
| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
172
| `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 |
173
-| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
174
-| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x |
173
+| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
174
+| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64 | 64, 128, 256, 512 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x |
175
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
176
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
177
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
vllm/v1/attention/backends/flashinfer.py
+11
-10
@@ -402,7 +402,7 @@ class FlashInferBackend(AttentionBackend):
402
@classmethod
403
def get_supported_head_sizes(cls) -> list[int]:
404
# https://github.com/flashinfer-ai/flashinfer/blob/3d55c71a62052c590c130897d3a3db49b14fcc34/include/flashinfer/utils.cuh#L157
405
- return [64, 128, 256]
405
+ return [64, 128, 256, 512]
406
407
@classmethod
408
def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
@@ -1451,15 +1451,16 @@ class FlashInferImpl(AttentionImpl):
1451
1452
num_actual_tokens = attn_metadata.num_actual_tokens
1453
1454
- # The FlashInfer api requires data to be in fp8_e4m3 or fp8_e5m2
1455
- # to process the cache when the kv_cache_dtype is fp8
1456
- if self.kv_sharing_target_layer_name is None and is_quantized_kv_cache(
1457
- self.kv_cache_dtype
1458
- ):
1459
- torch_dtype = FlashInferBackend.get_dtype_for_flashinfer(
1460
- self.kv_cache_dtype
1461
- )
1462
- kv_cache = kv_cache.view(torch_dtype)
1454
+ # FlashInfer treats uint8 KV cache as NVFP4. vLLM stores FP8 KV cache
1455
+ # as uint8 bytes, so pass FP8 caches with their logical dtype.
1456
+ if not self.is_kvcache_nvfp4 and kv_cache.dtype == torch.uint8:
1457
+ fp8_view_dtype = None
1458
+ if self.kv_cache_dtype in ("fp8", "fp8_e4m3", torch.float8_e4m3fn):
1459
+ fp8_view_dtype = torch.float8_e4m3fn
1460
+ elif self.kv_cache_dtype in ("fp8_e5m2", torch.float8_e5m2):
1461
+ fp8_view_dtype = torch.float8_e5m2
1462
+ if fp8_view_dtype is not None:
1463
+ kv_cache = kv_cache.view(fp8_view_dtype)
1464
1465
# Inputs and outputs may be padded for CUDA graphs
1466
query = query[:num_actual_tokens]