iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[SpecDecode] Support DFlash with FlashInfer (#43081)

Browse Source 
Signed-off-by: gss <2783977641@qq.com>
Co-authored-by: gss <2783977641@qq.com>
This commit is contained in:
gq112
2026-06-22 04:55:30 +00:00
committed by GitHub
co-authored by gss
parent b529bfd6c5
commit db32b53e30
3 changed files with 63 additions and 16 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, 128, 256, 512, 1024 | 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, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x |
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 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, 128, 256, 512, 1024 | 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 |
tests/kernels/attention/test_attention_selector.py
+1 -1
View File
@@ -435,7 +435,7 @@ def test_per_head_quant_scales_backend_selection(
]
+ (
[
("FLASHINFER", True, False), # FlashInfer does not support non-causal
("FLASHINFER", True, True), # FlashInfer supports non-causal
("FLASHINFER", False, True), # FlashInfer works with causal
]
if CudaPlatform is not None
vllm/v1/attention/backends/flashinfer.py
+60 -13
View File
@@ -359,6 +359,10 @@ class FlashInferBackend(AttentionBackend):
def get_name() -> str:
return "FLASHINFER"
@classmethod
def supports_non_causal(cls) -> bool:
return True
@staticmethod
def get_impl_cls() -> type["FlashInferImpl"]:
return FlashInferImpl
@@ -526,6 +530,7 @@ class FlashInferMetadata:
num_decode_tokens: int
num_prefills: int
num_prefill_tokens: int
causal: bool
prefill: FIPrefill | TRTLLMPrefill | None
"""
@@ -568,6 +573,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
self._prefill_wrapper: (
BatchPrefillWithPagedKVCacheWrapper | BatchDCPPrefillWrapper | None
) = None # Wrapper for prefill/append
self._noncausal_prefill_wrapper: BatchPrefillWithPagedKVCacheWrapper | None = (
None # Wrapper for non-causal prefill (DFlash)
)
self._decode_wrapper = None # Wrapper for decode (general shape)
if envs.VLLM_BATCH_INVARIANT:
@@ -784,7 +792,26 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
def _get_prefill_wrapper(
self,
causal: bool = True,
) -> BatchPrefillWithPagedKVCacheWrapper | BatchDCPPrefillWrapper:
if not causal:
if self.use_dcp:
raise NotImplementedError(
"FlashInfer non-causal prefill is not supported with DCP yet."
)
if self.is_kvcache_nvfp4:
raise NotImplementedError(
"FlashInfer non-causal attention is not supported with "
"NVFP4 KV cache."
)
if self._noncausal_prefill_wrapper is None:
self._noncausal_prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper(
self._get_workspace_buffer(),
get_kv_cache_layout(),
backend="auto",
)
return self._noncausal_prefill_wrapper
if self._prefill_wrapper is None:
if self.use_dcp:
self._prefill_wrapper = BatchDCPPrefillWrapper(
@@ -915,13 +942,22 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
) -> FlashInferMetadata:
num_reqs = common_attn_metadata.num_reqs
num_actual_tokens = common_attn_metadata.num_actual_tokens
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
split_decodes_and_prefills(
common_attn_metadata,
decode_threshold=self.reorder_batch_threshold,
require_uniform=True,
causal = common_attn_metadata.causal
if causal:
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
split_decodes_and_prefills(
common_attn_metadata,
decode_threshold=self.reorder_batch_threshold,
require_uniform=True,
)
)
)
else:
# FlashInfer decode/TRTLLM paths cannot express non-causal
# query-query attention, so DFlash runs as native prefill.
num_decodes = 0
num_prefills = num_reqs
num_decode_tokens = 0
num_prefill_tokens = num_actual_tokens
page_size = self.page_size
max_seq_len = common_attn_metadata.max_seq_len
@@ -940,7 +976,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
prefill_force_trtllm = (
True if page_size >= 128 else self.attention_config.use_trtllm_attention
)
prefill_use_trtllm = use_trtllm_attention(
prefill_use_trtllm = causal and use_trtllm_attention(
self.num_qo_heads,
self.num_kv_heads,
num_prefill_tokens,
@@ -954,11 +990,21 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
has_spec=uses_spec_reorder,
)
decode_use_trtllm = (
self.use_trtllm_decode_attention and self.dcp_world_size <= 1
causal and self.use_trtllm_decode_attention and self.dcp_world_size <= 1
)
all_uses_trtllm = (num_prefills == 0 or prefill_use_trtllm) and (
num_decodes == 0 or decode_use_trtllm
if not causal and self.use_dcp:
raise NotImplementedError(
"FlashInfer non-causal prefill is not supported with DCP yet."
)
if not causal and self.use_trtllm_decode_attention:
logger.warning_once(
"Using FlashInfer for draft model non-causal attention; TRTLLM "
"can still be used for target model causal attention."
)
all_uses_trtllm = causal and (
(num_prefills == 0 or prefill_use_trtllm)
and (num_decodes == 0 or decode_use_trtllm)
)
if not all_uses_trtllm:
@@ -997,6 +1043,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
num_decode_tokens=num_decode_tokens,
num_prefills=num_prefills,
num_prefill_tokens=num_prefill_tokens,
causal=causal,
use_cascade=use_cascade,
prefill=None,
decode=None,
@@ -1154,7 +1201,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
max_seq_len=max_seq_len,
)
else:
prefill_wrapper = self._get_prefill_wrapper()
prefill_wrapper = self._get_prefill_wrapper(causal=attn_metadata.causal)
# Slicing CPU buffers that are only needed for FI native prefills
paged_kv_last_page_len_prefill_cpu = self.paged_kv_last_page_len.cpu[
prefill_start:num_reqs
@@ -1204,7 +1251,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
num_kv_heads=self.num_kv_heads,
head_dim_qk=self.head_dim,
page_size=self.page_size,
causal=True,
causal=attn_metadata.causal,
sm_scale=self.sm_scale,
window_left=self.window_left,
logits_soft_cap=self.logits_soft_cap,
@@ -1579,7 +1626,7 @@ class FlashInferImpl(AttentionImpl):
self.logits_soft_cap or 0.0
)
assert prefill_wrapper._sm_scale == self.scale
assert prefill_wrapper._causal
assert prefill_wrapper._causal == attn_metadata.causal
if self.is_kvcache_nvfp4:
kv_cache_permute = nvfp4_kv_data