[SpecDecode] Support DFlash with FlashInfer (#43081)
Signed-off-by: gss <2783977641@qq.com> Co-authored-by: gss <2783977641@qq.com>
gq112 committed
Jun 22, 2026 at 12:55 UTC
db32b53e302e840ff9eda291eee8699836ca61c1
3 files changed
+63
-16
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, 128, 256, 512, 1024 | 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, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x |
173
+| `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 |
174
+| `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 |
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 |
tests/kernels/attention/test_attention_selector.py
+1
-1
@@ -435,7 +435,7 @@ def test_per_head_quant_scales_backend_selection(
435
]
436
+ (
437
[
438
- ("FLASHINFER", True, False), # FlashInfer does not support non-causal
438
+ ("FLASHINFER", True, True), # FlashInfer supports non-causal
439
("FLASHINFER", False, True), # FlashInfer works with causal
440
]
441
if CudaPlatform is not None
vllm/v1/attention/backends/flashinfer.py
+60
-13
@@ -359,6 +359,10 @@ class FlashInferBackend(AttentionBackend):
359
def get_name() -> str:
360
return "FLASHINFER"
361
362
+ @classmethod
363
+ def supports_non_causal(cls) -> bool:
364
+ return True
365
+
366
@staticmethod
367
def get_impl_cls() -> type["FlashInferImpl"]:
368
return FlashInferImpl
@@ -526,6 +530,7 @@ class FlashInferMetadata:
530
num_decode_tokens: int
531
num_prefills: int
532
num_prefill_tokens: int
533
+ causal: bool
534
535
prefill: FIPrefill | TRTLLMPrefill | None
536
"""
@@ -568,6 +573,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
573
self._prefill_wrapper: (
574
BatchPrefillWithPagedKVCacheWrapper | BatchDCPPrefillWrapper | None
575
) = None # Wrapper for prefill/append
576
+ self._noncausal_prefill_wrapper: BatchPrefillWithPagedKVCacheWrapper | None = (
577
+ None # Wrapper for non-causal prefill (DFlash)
578
+ )
579
self._decode_wrapper = None # Wrapper for decode (general shape)
580
581
if envs.VLLM_BATCH_INVARIANT:
@@ -784,7 +792,26 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
792
793
def _get_prefill_wrapper(
794
self,
795
+ causal: bool = True,
796
) -> BatchPrefillWithPagedKVCacheWrapper | BatchDCPPrefillWrapper:
797
+ if not causal:
798
+ if self.use_dcp:
799
+ raise NotImplementedError(
800
+ "FlashInfer non-causal prefill is not supported with DCP yet."
801
+ )
802
+ if self.is_kvcache_nvfp4:
803
+ raise NotImplementedError(
804
+ "FlashInfer non-causal attention is not supported with "
805
+ "NVFP4 KV cache."
806
+ )
807
+ if self._noncausal_prefill_wrapper is None:
808
+ self._noncausal_prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper(
809
+ self._get_workspace_buffer(),
810
+ get_kv_cache_layout(),
811
+ backend="auto",
812
+ )
813
+ return self._noncausal_prefill_wrapper
814
+
815
if self._prefill_wrapper is None:
816
if self.use_dcp:
817
self._prefill_wrapper = BatchDCPPrefillWrapper(
@@ -915,13 +942,22 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
942
) -> FlashInferMetadata:
943
num_reqs = common_attn_metadata.num_reqs
944
num_actual_tokens = common_attn_metadata.num_actual_tokens
918
- num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
919
- split_decodes_and_prefills(
920
- common_attn_metadata,
921
- decode_threshold=self.reorder_batch_threshold,
922
- require_uniform=True,
945
+ causal = common_attn_metadata.causal
946
+ if causal:
947
+ num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
948
+ split_decodes_and_prefills(
949
+ common_attn_metadata,
950
+ decode_threshold=self.reorder_batch_threshold,
951
+ require_uniform=True,
952
+ )
953
)
924
- )
954
+ else:
955
+ # FlashInfer decode/TRTLLM paths cannot express non-causal
956
+ # query-query attention, so DFlash runs as native prefill.
957
+ num_decodes = 0
958
+ num_prefills = num_reqs
959
+ num_decode_tokens = 0
960
+ num_prefill_tokens = num_actual_tokens
961
962
page_size = self.page_size
963
max_seq_len = common_attn_metadata.max_seq_len
@@ -940,7 +976,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
976
prefill_force_trtllm = (
977
True if page_size >= 128 else self.attention_config.use_trtllm_attention
978
)
943
- prefill_use_trtllm = use_trtllm_attention(
979
+ prefill_use_trtllm = causal and use_trtllm_attention(
980
self.num_qo_heads,
981
self.num_kv_heads,
982
num_prefill_tokens,
@@ -954,11 +990,21 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
990
has_spec=uses_spec_reorder,
991
)
992
decode_use_trtllm = (
957
- self.use_trtllm_decode_attention and self.dcp_world_size <= 1
993
+ causal and self.use_trtllm_decode_attention and self.dcp_world_size <= 1
994
)
995
960
- all_uses_trtllm = (num_prefills == 0 or prefill_use_trtllm) and (
961
- num_decodes == 0 or decode_use_trtllm
996
+ if not causal and self.use_dcp:
997
+ raise NotImplementedError(
998
+ "FlashInfer non-causal prefill is not supported with DCP yet."
999
+ )
1000
+ if not causal and self.use_trtllm_decode_attention:
1001
+ logger.warning_once(
1002
+ "Using FlashInfer for draft model non-causal attention; TRTLLM "
1003
+ "can still be used for target model causal attention."
1004
+ )
1005
+ all_uses_trtllm = causal and (
1006
+ (num_prefills == 0 or prefill_use_trtllm)
1007
+ and (num_decodes == 0 or decode_use_trtllm)
1008
)
1009
1010
if not all_uses_trtllm:
@@ -997,6 +1043,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1043
num_decode_tokens=num_decode_tokens,
1044
num_prefills=num_prefills,
1045
num_prefill_tokens=num_prefill_tokens,
1046
+ causal=causal,
1047
use_cascade=use_cascade,
1048
prefill=None,
1049
decode=None,
@@ -1154,7 +1201,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1201
max_seq_len=max_seq_len,
1202
)
1203
else:
1157
- prefill_wrapper = self._get_prefill_wrapper()
1204
+ prefill_wrapper = self._get_prefill_wrapper(causal=attn_metadata.causal)
1205
# Slicing CPU buffers that are only needed for FI native prefills
1206
paged_kv_last_page_len_prefill_cpu = self.paged_kv_last_page_len.cpu[
1207
prefill_start:num_reqs
@@ -1204,7 +1251,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1251
num_kv_heads=self.num_kv_heads,
1252
head_dim_qk=self.head_dim,
1253
page_size=self.page_size,
1207
- causal=True,
1254
+ causal=attn_metadata.causal,
1255
sm_scale=self.sm_scale,
1256
window_left=self.window_left,
1257
logits_soft_cap=self.logits_soft_cap,
@@ -1579,7 +1626,7 @@ class FlashInferImpl(AttentionImpl):
1626
self.logits_soft_cap or 0.0
1627
)
1628
assert prefill_wrapper._sm_scale == self.scale
1582
- assert prefill_wrapper._causal
1629
+ assert prefill_wrapper._causal == attn_metadata.causal
1630
1631
if self.is_kvcache_nvfp4:
1632
kv_cache_permute = nvfp4_kv_data