diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index 8a261c502e..2dff668a01 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -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 | diff --git a/tests/kernels/attention/test_attention_selector.py b/tests/kernels/attention/test_attention_selector.py index 447b502293..1e85f76b64 100644 --- a/tests/kernels/attention/test_attention_selector.py +++ b/tests/kernels/attention/test_attention_selector.py @@ -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 diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 666c32bca8..044fc8c79c 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -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