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 4, 18:45 UTC
FlexAttention non-causal support (#40394)
Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Fynn Schmitt-Ulms committed
Apr 22, 2026 at 16:22 UTC
29f64c5f5e635e5071fe23ef6ed36edb520b9ae8
3 files changed
+122
-17
tests/v1/attention/test_attention_backends.py
+73
-4
@@ -315,6 +315,7 @@ def _test_backend_correctness(
315
backend_to_test: list[AttentionBackendEnum | str],
316
mask_mod,
317
*,
318
+ causal: bool = True,
319
attn_type: AttentionType = AttentionType.DECODER,
320
block_size: int = 16,
321
atol: float = 1e-2,
@@ -370,7 +371,7 @@ def _test_backend_correctness(
371
)
372
device = torch.device(f"{DEVICE_TYPE}:0")
373
373
- kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
374
+ kv_cache_spec = create_standard_kv_cache_spec(vllm_config, attn_type)
375
376
# 1. Setup
377
batch_size = batch_spec.batch_size
@@ -453,9 +454,7 @@ def _test_backend_correctness(
454
common_attn_metadata = create_common_attn_metadata(
455
batch_spec, vllm_config.cache_config.block_size, device
456
)
456
- if attn_type == AttentionType.ENCODER_ONLY:
457
- # For encoder-only, all tokens are prefill tokens
458
- common_attn_metadata.causal = False
457
+ common_attn_metadata.causal = causal
458
459
# 3. Simulate Paged KV Cache and a realistic slot_mapping
460
kv_cache = create_and_prepopulate_kv_cache(
@@ -736,6 +735,76 @@ def test_sliding_window_encoder_backend_correctness(
735
model,
736
SLIDING_WINDOW_BACKENDS_TO_TEST,
737
sliding_window_mask_mod_fn,
738
+ causal=False,
739
attn_type=AttentionType.ENCODER_ONLY,
740
tensor_parallel_size=tensor_parallel_size,
741
)
742
+
743
+
744
+NON_CAUSAL_BACKENDS_TO_TEST = [
745
+ AttentionBackendEnum.FLASH_ATTN,
746
+ AttentionBackendEnum.FLEX_ATTENTION,
747
+ "FLEX_ATTENTION_SLOW",
748
+]
749
+
750
+if current_platform.is_rocm():
751
+ NON_CAUSAL_BACKENDS_TO_TEST = [
752
+ x
753
+ for x in NON_CAUSAL_BACKENDS_TO_TEST
754
+ if x is not AttentionBackendEnum.FLASH_ATTN
755
+ ]
756
+
757
+
758
+@pytest.mark.parametrize(
759
+ "batch_spec_name",
760
+ [
761
+ "small_decode",
762
+ "small_prefill",
763
+ "mixed_small",
764
+ ],
765
+)
766
+@pytest.mark.parametrize("model", ["meta-llama/Meta-Llama-3-8B"])
767
+def test_non_causal_backend_correctness(
768
+ default_vllm_config, batch_spec_name: str, model: str
769
+):
770
+ """Test backend's correctness with non-causal (bidirectional) decoder
771
+ attention, as used by DFlash speculative decoding."""
772
+
773
+ def bidirectional_mask_mod(
774
+ b: torch.Tensor,
775
+ h: torch.Tensor,
776
+ q_idx: torch.Tensor,
777
+ kv_idx: torch.Tensor,
778
+ *,
779
+ context_len: int,
780
+ ):
781
+ return q_idx >= 0 # Always True
782
+
783
+ batch_spec = BATCH_SPECS[batch_spec_name]
784
+ LARGE_BLOCK_BACKENDS = (
785
+ [AttentionBackendEnum.FLEX_ATTENTION]
786
+ if is_torch_equal_or_newer("2.9.0.dev0")
787
+ else []
788
+ )
789
+
790
+ SMALL_BLOCK_BACKENDS = [
791
+ x for x in NON_CAUSAL_BACKENDS_TO_TEST if x not in LARGE_BLOCK_BACKENDS
792
+ ]
793
+
794
+ _test_backend_correctness(
795
+ batch_spec,
796
+ model,
797
+ SMALL_BLOCK_BACKENDS,
798
+ bidirectional_mask_mod,
799
+ causal=False,
800
+ )
801
+
802
+ if LARGE_BLOCK_BACKENDS:
803
+ _test_backend_correctness(
804
+ batch_spec,
805
+ model,
806
+ LARGE_BLOCK_BACKENDS,
807
+ bidirectional_mask_mod,
808
+ causal=False,
809
+ block_size=128,
810
+ )
tests/v1/attention/utils.py
+20
-3
@@ -21,10 +21,11 @@ from vllm.config.model import ModelDType
21
from vllm.v1.attention.backend import (
22
AttentionImpl,
23
AttentionMetadataBuilder,
24
+ AttentionType,
25
CommonAttentionMetadata,
26
)
27
from vllm.v1.attention.backends.registry import AttentionBackendEnum
27
-from vllm.v1.kv_cache_interface import FullAttentionSpec
28
+from vllm.v1.kv_cache_interface import EncoderOnlyAttentionSpec, FullAttentionSpec
29
30
31
@dataclass
@@ -142,8 +143,24 @@ def try_backend_includes_kv_cache_update(
143
raise AssertionError("unreachable") from None
144
145
145
-def create_standard_kv_cache_spec(vllm_config: VllmConfig) -> FullAttentionSpec:
146
- """Create a FullAttentionSpec from ModelParams only."""
146
+def create_standard_kv_cache_spec(
147
+ vllm_config: VllmConfig,
148
+ attn_type: AttentionType = AttentionType.DECODER,
149
+) -> FullAttentionSpec | EncoderOnlyAttentionSpec:
150
+ """Create an AttentionSpec from VllmConfig.
151
+
152
+ Returns an EncoderOnlyAttentionSpec for encoder-only attention (no KV
153
+ cache), and a FullAttentionSpec otherwise.
154
+ """
155
+ if attn_type == AttentionType.ENCODER_ONLY:
156
+ return EncoderOnlyAttentionSpec(
157
+ block_size=vllm_config.cache_config.block_size,
158
+ num_kv_heads=vllm_config.model_config.get_num_kv_heads(
159
+ vllm_config.parallel_config
160
+ ),
161
+ head_size=vllm_config.model_config.get_head_size(),
162
+ dtype=vllm_config.model_config.dtype,
163
+ )
164
return FullAttentionSpec(
165
block_size=vllm_config.cache_config.block_size,
166
num_kv_heads=vllm_config.model_config.get_num_kv_heads(
vllm/v1/attention/backends/flex_attention.py
+29
-10
@@ -36,7 +36,7 @@ from vllm.v1.attention.backend import (
36
AttentionType,
37
CommonAttentionMetadata,
38
)
39
-from vllm.v1.kv_cache_interface import AttentionSpec
39
+from vllm.v1.kv_cache_interface import AttentionSpec, EncoderOnlyAttentionSpec
40
41
logger = init_logger(__name__)
42
@@ -90,6 +90,10 @@ class FlexAttentionBackend(AttentionBackend):
90
def get_name() -> str:
91
return "FLEX_ATTENTION"
92
93
+ @classmethod
94
+ def supports_non_causal(cls) -> bool:
95
+ return True
96
+
97
@classmethod
98
def supports_attn_type(cls, attn_type: str) -> bool:
99
"""FlexAttention supports both decoder and encoder-only attention."""
@@ -294,6 +298,12 @@ def causal_mask_mod(
298
return q_idx >= kv_idx
299
300
301
+def bidirectional_mask_mod(
302
+ b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor
303
+):
304
+ return q_idx >= 0
305
+
306
+
307
# Type alias for the block sparsity hint callable signature.
308
_block_sparsity_hint_signature = Callable[
309
[torch.Tensor, torch.Tensor, int], torch.Tensor
@@ -364,6 +374,7 @@ class FlexAttentionMetadata:
374
block_mask: BlockMask | None = None
375
score_mod: _score_mod_signature | None = None
376
logical_mask_mod: _mask_mod_signature = causal_mask_mod
377
+ uses_paged_kv: bool = True
378
doc_ids: torch.Tensor | None = None
379
direct_build: bool = True
380
q_block_size: int = 16
@@ -497,7 +508,7 @@ class FlexAttentionMetadata:
508
False,
509
)
510
500
- return final_mask_mod if self.causal else sliding_window_mask_mod
511
+ return final_mask_mod if self.uses_paged_kv else sliding_window_mask_mod
512
513
def get_prefix_lm_mask_mod(self) -> _mask_mod_signature:
514
"""Creates the prefix LM mask_mod function for FlexAttention."""
@@ -541,8 +552,7 @@ class FlexAttentionMetadata:
552
def get_mask_mod(self):
553
# Stage-1: initialize the base mask_mod
554
# (causal mask for decoder or bidirectional mask for encoder)
544
- has_custom_mask = self.logical_mask_mod is not causal_mask_mod
545
- if self.causal or has_custom_mask:
555
+ if self.uses_paged_kv:
556
mask_mod = self.get_paged_mask_mod()
557
else:
558
mask_mod = self.get_bidirectional_mask_mod()
@@ -595,7 +605,7 @@ class FlexAttentionMetadata:
605
return transformed_score_mod
606
607
def _build_block_mask_direct(self) -> BlockMask:
598
- """Direct block mask construction for standard causal attention.
608
+ """Direct block mask construction for paged KV cache attention.
609
610
This method constructs the block mask directly using
611
BlockMask.from_kv_blocks which is much more efficient than the
@@ -693,7 +703,9 @@ class FlexAttentionMetadata:
703
704
def build_block_mask(self) -> BlockMask:
705
mask_mod = self.get_mask_mod()
696
- kv_len = self.total_cache_tokens if self.causal else self.num_actual_tokens
706
+ kv_len = (
707
+ self.total_cache_tokens if self.uses_paged_kv else self.num_actual_tokens
708
+ )
709
return create_block_mask_compiled(
710
mask_mod,
711
None,
@@ -842,8 +854,16 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
854
offset_tensor = common_attn_metadata.compute_num_computed_tokens()
855
offset_tensor = copy_to_persistent(self.persistent_offset_tensor, offset_tensor)
856
857
+ uses_paged_kv = not isinstance(self.kv_cache_spec, EncoderOnlyAttentionSpec)
858
+ logical_mask_mod = (
859
+ bidirectional_mask_mod
860
+ if uses_paged_kv and not common_attn_metadata.causal
861
+ else causal_mask_mod
862
+ )
863
+
864
out = FlexAttentionMetadata(
865
causal=common_attn_metadata.causal,
866
+ logical_mask_mod=logical_mask_mod,
867
num_actual_tokens=num_actual_tokens,
868
max_query_len=max_query_len,
869
query_start_loc=query_start_loc,
@@ -863,10 +883,11 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
883
total_cache_tokens=total_cache_tokens,
884
decode_offset=offset_tensor,
885
num_blocks_per_seq=num_blocks_per_seq,
886
+ uses_paged_kv=uses_paged_kv,
887
# FIXME(Isotr0py): direct build has issue to build bidirectional
888
# attention block mask for encoder-only models, disable it temporarily.
889
# see: https://github.com/vllm-project/vllm/pull/27329#issuecomment-3431484053
869
- direct_build=(self.direct_build and common_attn_metadata.causal),
890
+ direct_build=self.direct_build and uses_paged_kv,
891
q_block_size=self.q_block_size,
892
kv_block_size=self.kv_block_size,
893
persistent_kv_indices=self.persistent_kv_indices,
@@ -1055,9 +1076,7 @@ class FlexAttentionImpl(AttentionImpl):
1076
else:
1077
attn_metadata.block_mask = attn_metadata.build_block_mask()
1078
1058
- if not attn_metadata.causal:
1059
- assert self.attn_type == AttentionType.ENCODER_ONLY
1060
-
1079
+ if self.attn_type == AttentionType.ENCODER_ONLY:
1080
query, key_tensor, value_tensor = map(
1081
lambda x: self.view_as_4d(x).permute(0, 2, 1, 3),
1082
(query, key, value),