@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 82a85dc1d2

[Attention] Add Triton/FlashInfer composite for multimodal prefix attention (#56305)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: OpenAI Codex <codex@openai.com>

Matthew Bonanni committed Sep 13, 2026 at 10:29 UTC 82a85dc1d2d5b3ad4453eaa1bd7596de888a66b3
15 files changed +1095 -3
docs/design/attention_backends.md
+49
@@ -56,6 +56,55 @@ llm = LLM(
56
57 ## Backend Selection Behavior
58
59 +### Triton/FlashAttention Composite
60 +
61 +On Hopper, `TRITON_FLASH_ATTN` is preferred for compatible multimodal-prefix
62 +configurations. It uses Triton when the current queries need bidirectional
63 +image attention and FlashAttention for causal text prefills and decode, with a
64 +shared KV cache. The causal child must resolve to FA4. This happens
65 +automatically for FA4-only shapes such as head size 512 and models whose
66 +version policy promotes all layers to FA4, including Gemma 4. A standalone
67 +head-size-256 configuration whose version policy selects FA3 falls back to
68 +Triton for the whole layer.
69 +
70 +No attention override is needed for supported multimodal models:
71 +
72 +```bash
73 +vllm serve google/gemma-4-31B-it
74 +```
75 +
76 +### Triton/FlashInfer Composite
77 +
78 +On Blackwell, `TRITON_FLASHINFER` is preferred for compatible
79 +multimodal-prefix configurations, including Gemma 4 with BF16 or FP8 KV cache. It uses
80 +Triton when a batch's current queries require bidirectional image attention,
81 +and FlashInfer for causal text prefills and decode. Historical image tokens
82 +alone do not select Triton.
83 +
84 +This instantiates `create_composite_attention_backend` from
85 +`vllm/v1/attention/backends/composite.py` with Triton,
86 +FlashInfer, and `MMPrefixAttentionRouting`. The reusable factory owns child
87 +implementations, metadata dispatch, compatible cache requirements, and workspace
88 +sharing. The routing policy selects the child and defines graph-capture safety.
89 +Other combinations can reuse the same machinery.
90 +
91 +The backend supports head dimensions 256/512, FP16/BF16 and FP8 KV cache,
92 +and 64-token kernel pages with a head-major cache layout. TRTLLM handles causal
93 +attention at both head dimensions; its hdim512 kernels do not support 128-token
94 +pages.
95 +For Gemma 4, full CUDA graphs cover single-token batches; multi-token batches
96 +use the non-full-graph execution path. Models whose image masks extend beyond
97 +the sliding window cannot use full attention graphs with this composite.
98 +Context parallelism, R-SWA, attention sinks,
99 +and adaptive verification are not supported by this composite.
100 +
101 +It can also be selected explicitly:
102 +
103 +```bash
104 +vllm serve google/gemma-4-31B-it \
105 + --attention-backend TRITON_FLASHINFER
106 +```
107 +
108 ### Manual Selection
109
110 When you explicitly set a backend via `--attention-backend` or `AttentionConfig`:
tests/kernels/attention/test_attention_selector.py
+139
@@ -808,3 +808,142 @@ def test_fa4_hd256_impl_selection(attn_type, sliding_window, expected):
808 )
809 assert impl.vllm_flash_attn_version == expected
810 assert impl.fa4_hd256 == (expected == 4)
811 +
812 +
813 +@pytest.fixture
814 +def blackwell_selection():
815 + with (
816 + patch.object(
817 + type(current_platform),
818 + "get_device_capability",
819 + return_value=DeviceCapability(10, 0),
820 + ),
821 + patch(
822 + "vllm.v1.attention.backends.fa_utils.is_fa_version_supported",
823 + return_value=True,
824 + ),
825 + ):
826 + yield
827 +
828 +
829 +@pytest.fixture
830 +def hopper_selection():
831 + with (
832 + patch.object(
833 + type(current_platform),
834 + "get_device_capability",
835 + return_value=DeviceCapability(9, 0),
836 + ),
837 + patch(
838 + "vllm.v1.attention.backends.fa_utils.is_fa_version_supported",
839 + return_value=True,
840 + ),
841 + ):
842 + yield
843 +
844 +
845 +@blackwell_only
846 +@pytest.mark.parametrize("use_mm_prefix", [False, True])
847 +@pytest.mark.parametrize("flash_attn_version", [None, 3, 4])
848 +def test_hopper_mm_prefix_selects_triton_flash_attn(
849 + use_mm_prefix, flash_attn_version, hopper_selection
850 +):
851 + """Hopper selects the composite only when its causal route resolves FA4."""
852 + from vllm.engine.arg_utils import EngineArgs
853 +
854 + config = EngineArgs(
855 + model="google/gemma-4-31B-it",
856 + dtype="bfloat16",
857 + attention_config=AttentionConfig(flash_attn_version=flash_attn_version),
858 + ).create_engine_config()
859 + with set_current_vllm_config(config):
860 + backend = get_attn_backend(
861 + 256, torch.bfloat16, None, use_mm_prefix=use_mm_prefix
862 + )
863 + fa4_resolved = config.attention_config.flash_attn_version == 4
864 + if use_mm_prefix:
865 + expected = "TRITON_FLASH_ATTN" if fa4_resolved else "TRITON_ATTN"
866 + else:
867 + expected = "FLASH_ATTN"
868 + assert backend.get_name() == expected
869 +
870 +
871 +@blackwell_only
872 +@pytest.mark.parametrize("use_mm_prefix", [False, True])
873 +@pytest.mark.parametrize("kv_cache_dtype", [None, "fp8_e4m3"])
874 +def test_mm_prefix_selects_composite_without_changing_causal_default(
875 + use_mm_prefix, kv_cache_dtype, blackwell_selection
876 +):
877 + """The image-mask requirement must reach CUDA's automatic backend priority."""
878 + from vllm.engine.arg_utils import EngineArgs
879 +
880 + config = EngineArgs(
881 + model="google/gemma-4-31B-it", dtype="bfloat16"
882 + ).create_engine_config()
883 + with set_current_vllm_config(config):
884 + backend = get_attn_backend(
885 + 256, torch.bfloat16, kv_cache_dtype, use_mm_prefix=use_mm_prefix
886 + )
887 + expected = "TRITON_FLASHINFER" if use_mm_prefix else "FLASHINFER"
888 + assert backend.get_name() == expected
889 +
890 +
891 +@blackwell_only
892 +@pytest.mark.parametrize(
893 + "feature_kwargs,reason",
894 + [
895 + ({"use_rswa": True}, "R-SWA not supported"),
896 + ({"kv_cache_dtype": "int8_per_token_head"}, "kv_cache_dtype not supported"),
897 + ({"block_size": 16}, "block_size not supported"),
898 + ({"use_dcp": True}, "DCP not supported"),
899 + ({"use_pcp": True}, "PCP not supported"),
900 + ({"use_adaptive_verification": True}, "device-cpu query lens mismatch"),
901 + ],
902 +)
903 +def test_composite_rejects_features_not_shared_by_both_routes(feature_kwargs, reason):
904 + """Automatic selection must not admit features only one child can execute."""
905 + from vllm.engine.arg_utils import EngineArgs
906 + from vllm.v1.attention.backends.triton_flashinfer import TritonFlashInferBackend
907 +
908 + config = EngineArgs(
909 + model="google/gemma-4-31B-it", dtype="bfloat16"
910 + ).create_engine_config()
911 + args = dict(
912 + head_size=256,
913 + dtype=torch.bfloat16,
914 + kv_cache_dtype="auto",
915 + block_size=128,
916 + use_mla=False,
917 + has_sink=False,
918 + use_sparse=False,
919 + use_mm_prefix=True,
920 + use_per_head_quant_scales=False,
921 + device_capability=DeviceCapability(10, 0),
922 + attn_type=AttentionType.DECODER,
923 + )
924 + args.update(feature_kwargs)
925 + with set_current_vllm_config(config):
926 + reasons = TritonFlashInferBackend.validate_configuration(**args)
927 + assert any(reason in message for message in reasons), reasons
928 +
929 +
930 +@blackwell_only
931 +def test_rswa_selection_does_not_reuse_causal_result(blackwell_selection):
932 + """An R-SWA requirement must invalidate a previously cached causal selection."""
933 + from vllm.engine.arg_utils import EngineArgs
934 +
935 + config = EngineArgs(
936 + model="google/gemma-4-31B-it",
937 + dtype="bfloat16",
938 + attention_config={"backend": "TRITON_FLASHINFER"},
939 + ).create_engine_config()
940 + with set_current_vllm_config(config):
941 + assert (
942 + get_attn_backend(256, torch.bfloat16, None).get_name()
943 + == "TRITON_FLASHINFER"
944 + )
945 + config.model_config.model_arch_config.rswa_window = 128
946 + with pytest.raises(ValueError, match="R-SWA"):
947 + get_attn_backend(256, torch.bfloat16, None)
948 + config.attention_config.backend = AttentionBackendEnum.TRITON_ATTN
949 + assert get_attn_backend(256, torch.bfloat16, None).get_name() == "TRITON_ATTN"
tests/v1/attention/test_mm_prefix.py
+246
@@ -638,3 +638,249 @@ def test_mm_prefix_kv_cache_path(head_size: int):
638 assert not torch.allclose(expected, causal_only, atol=2e-2, rtol=2e-2), (
639 "test batch does not actually exercise the mm_prefix branch"
640 )
641 +
642 +
643 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA backend imports")
644 +@pytest.mark.parametrize(
645 + "query_len,seq_len,spans,expected,unclamped_expected",
646 + [
647 + (16, 32, [(20, 28)], True, True),
648 + (16, 32, [(0, 16)], False, True),
649 + (16, 32, [(0, 17)], True, True),
650 + (16, 32, [(32, 40)], False, False),
651 + (16, 32, [(31, 40)], False, True),
652 + (1, 24, [(20, 28)], False, True),
653 + (0, 24, [(20, 28)], False, False),
654 + (16, 32, [(20, 20)], False, False),
655 + (16, 64, [(20, 28)], False, False),
656 + ],
657 +)
658 +def test_composite_routes_queries_that_need_image_masking(
659 + query_len, seq_len, spans, expected, unclamped_expected
660 +):
661 + """Historical images and single-query steps must use the causal graph path."""
662 + from types import SimpleNamespace
663 +
664 + from vllm.v1.attention.backends.composite import requires_mm_prefix
665 +
666 + common = SimpleNamespace(
667 + max_query_len=query_len,
668 + mm_req_doc_ranges={0: spans},
669 + query_start_loc_cpu=torch.tensor([0, query_len]),
670 + seq_lens_cpu_upper_bound=torch.tensor([seq_len]),
671 + )
672 + assert requires_mm_prefix(common) is expected
673 + assert requires_mm_prefix(common, unclamped_window=True) is unclamped_expected
674 +
675 +
676 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA attention kernels")
677 +@pytest.mark.parametrize("head_size", [256, 512])
678 +@pytest.mark.parametrize(
679 + "reverse_children,cache_dtype",
680 + [(False, "auto"), (True, "auto"), (False, "fp8_e4m3")],
681 +)
682 +@pytest.mark.parametrize("causal_backend_name", ["FLASHINFER", "FLASH_ATTN"])
683 +def test_composite_shared_cache_across_image_and_causal_steps(
684 + head_size, reverse_children, cache_dtype, causal_backend_name
685 +):
686 + """Changing routes must preserve KV writes, image masking and graph replay."""
687 + if causal_backend_name == "FLASHINFER":
688 + if not current_platform.is_device_capability_family(100):
689 + pytest.skip("The Triton/FlashInfer composite supports Blackwell")
690 + else:
691 + if not current_platform.is_device_capability(90):
692 + pytest.skip("The Triton/FlashAttention composite supports Hopper")
693 + if reverse_children or cache_dtype != "auto":
694 + pytest.skip("Reverse-child and FP8 coverage is backend-independent")
695 +
696 + from vllm.config import set_current_vllm_config
697 + from vllm.engine.arg_utils import EngineArgs
698 + from vllm.model_executor.layers.attention import Attention
699 + from vllm.utils.torch_utils import set_default_torch_dtype
700 + from vllm.v1.attention.backend import AttentionCGSupport, CommonAttentionMetadata
701 + from vllm.v1.attention.backends.composite import (
702 + MMPrefixAttentionRouting,
703 + create_composite_attention_backend,
704 + )
705 + from vllm.v1.attention.backends.flash_attn import (
706 + FlashAttentionMetadata,
707 + )
708 + from vllm.v1.attention.backends.flashinfer import (
709 + FlashInferBackend,
710 + FlashInferMetadata,
711 + FlashInferTrtllmAPIDecode,
712 + TRTLLMPrefill,
713 + )
714 + from vllm.v1.attention.backends.triton_attn import (
715 + TritonAttentionBackend,
716 + TritonAttentionMetadata,
717 + )
718 + from vllm.v1.attention.backends.triton_flash_attn import (
719 + TritonFlashAttentionBackend,
720 + )
721 + from vllm.v1.attention.backends.triton_flashinfer import (
722 + TritonFlashInferBackend,
723 + )
724 + from vllm.v1.kv_cache_interface import FullAttentionSpec, SlidingWindowSpec
725 +
726 + if causal_backend_name == "FLASHINFER":
727 + backend = TritonFlashInferBackend
728 + causal_metadata_type = FlashInferMetadata
729 + block_size = backend.get_supported_kernel_block_sizes()[0]
730 + else:
731 + backend = TritonFlashAttentionBackend
732 + causal_metadata_type = FlashAttentionMetadata
733 + block_size = backend.get_preferred_block_size(16)
734 + num_blocks = (209 + block_size - 1) // block_size
735 + torch.manual_seed(42)
736 + attention_config = {"backend": backend.get_name()}
737 + if causal_backend_name == "FLASH_ATTN":
738 + attention_config["flash_attn_version"] = 4
739 + cfg = EngineArgs(
740 + model="google/gemma-4-31B-it",
741 + dtype="bfloat16",
742 + max_model_len=512,
743 + max_num_seqs=4,
744 + max_num_batched_tokens=512,
745 + block_size=block_size,
746 + kv_cache_dtype=cache_dtype,
747 + enforce_eager=True,
748 + attention_config=attention_config,
749 + ).create_engine_config()
750 + cfg.cache_config.kv_cache_layout = "LBHNC"
751 + window = 128 if head_size == 256 else None
752 + storage_dtype = torch.uint8 if cache_dtype == "fp8_e4m3" else DTYPE
753 + spec_kwargs = dict(
754 + block_size=block_size,
755 + num_kv_heads=16,
756 + head_size=head_size,
757 + dtype=storage_dtype,
758 + kv_quant_mode=get_kv_quant_mode(cache_dtype),
759 + )
760 + spec = (
761 + SlidingWindowSpec(**spec_kwargs, sliding_window=window)
762 + if window
763 + else FullAttentionSpec(**spec_kwargs)
764 + )
765 + if reverse_children:
766 + # Swapping the children and policy must preserve outputs and KV writes;
767 + # dispatch cannot depend on a particular child's class or position.
768 + class ReversedRouting(MMPrefixAttentionRouting):
769 + capture_variant = 0
770 +
771 + def select(self, metadata):
772 + return 1 - super().select(metadata)
773 +
774 + @staticmethod
775 + def variant_uses_mm_prefix(variant):
776 + return variant == 1
777 +
778 + backend = create_composite_attention_backend(
779 + FlashInferBackend,
780 + TritonAttentionBackend,
781 + name="ReversedCompositeBackend",
782 + backend_name="CUSTOM",
783 + module=__name__,
784 + routing_policy=ReversedRouting,
785 + )
786 + with set_current_vllm_config(cfg), set_default_torch_dtype(DTYPE):
787 + layer = Attention(
788 + 32,
789 + head_size,
790 + head_size**-0.5,
791 + num_kv_heads=16,
792 + cache_config=cfg.cache_config,
793 + prefix="composite_test",
794 + per_layer_sliding_window=window,
795 + attn_backend=backend,
796 + ).to(DEVICE)
797 + layer.mm_prefix_clamp_sliding_window = window is not None
798 + assert backend.get_builder_cls().get_cudagraph_support(cfg, spec) == (
799 + AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
800 + )
801 + if window:
802 + layer.mm_prefix_clamp_sliding_window = False
803 + assert backend.get_builder_cls().get_cudagraph_support(cfg, spec) == (
804 + AttentionCGSupport.NEVER
805 + )
806 + layer.mm_prefix_clamp_sliding_window = True
807 + builder = backend.get_builder_cls()(spec, ["composite_test"], cfg, DEVICE)
808 + builder.set_kernel_block_size(block_size)
809 + # Match production, where unwritten page tails are zero-initialized.
810 + cache = torch.zeros(
811 + num_blocks,
812 + 16,
813 + block_size,
814 + 2 * head_size,
815 + dtype=storage_dtype,
816 + device=DEVICE,
817 + )
818 + keys = torch.randn(209, 16, head_size, dtype=DTYPE, device=DEVICE)
819 + values = torch.randn_like(keys)
820 + if cache_dtype == "fp8_e4m3":
821 + keys = keys.to(torch.float8_e4m3fn).to(DTYPE)
822 + values = values.to(torch.float8_e4m3fn).to(DTYPE)
823 + table = torch.arange(num_blocks, dtype=torch.int32, device=DEVICE).unsqueeze(0)
824 + previous = 0
825 + for end, spans, metadata_type in [
826 + (128, [(32, 111)], TritonAttentionMetadata),
827 + (129, [(32, 111)], causal_metadata_type),
828 + (145, [(32, 111)], causal_metadata_type),
829 + (209, [(32, 111), (160, 207)], TritonAttentionMetadata),
830 + ]:
831 + qlen = end - previous
832 + starts = torch.tensor([0, qlen], dtype=torch.int32)
833 + lengths = torch.tensor([end], dtype=torch.int32)
834 + common = CommonAttentionMetadata(
835 + query_start_loc=starts.to(DEVICE),
836 + query_start_loc_cpu=starts,
837 + seq_lens=lengths.to(DEVICE),
838 + seq_lens_cpu_upper_bound=lengths,
839 + num_reqs=1,
840 + num_actual_tokens=qlen,
841 + max_query_len=qlen,
842 + max_seq_len=end,
843 + block_table_tensor=table,
844 + slot_mapping=torch.arange(previous, end, device=DEVICE),
845 + mm_req_doc_ranges={0: spans},
846 + )
847 + metadata = builder.build(0, common)
848 + assert isinstance(metadata, metadata_type)
849 + if isinstance(metadata, FlashInferMetadata):
850 + if qlen == 1:
851 + assert isinstance(metadata.decode, FlashInferTrtllmAPIDecode)
852 + else:
853 + assert isinstance(metadata.prefill, TRTLLMPrefill)
854 + query = torch.randn(qlen, 32, head_size, dtype=DTYPE, device=DEVICE)
855 + if cache_dtype == "fp8_e4m3":
856 + query = query.to(torch.float8_e4m3fn).to(DTYPE)
857 + key, value = keys[previous:end], values[previous:end]
858 + output = torch.empty_like(query)
859 + layer.impl.do_kv_cache_update(layer, key, value, cache, common.slot_mapping)
860 + layer.impl.forward(layer, query, key, value, cache, metadata, output)
861 + expected = _dense_reference(
862 + query,
863 + [keys[:end]],
864 + [values[:end]],
865 + [qlen],
866 + [end],
867 + {0: spans},
868 + window,
869 + head_size**-0.5,
870 + mm_clamp_sw=window or 0,
871 + )
872 + torch.testing.assert_close(output.float(), expected, atol=2e-2, rtol=2e-2)
873 + if qlen == 1:
874 + metadata = builder.build_for_cudagraph_capture(common)
875 + assert isinstance(metadata, metadata_type)
876 + graph = torch.cuda.CUDAGraph()
877 + with torch.cuda.graph(graph):
878 + layer.impl.forward(
879 + layer, query, key, value, cache, metadata, output
880 + )
881 + output.zero_()
882 + graph.replay()
883 + torch.testing.assert_close(
884 + output.float(), expected, atol=2e-2, rtol=2e-2
885 + )
886 + previous = end
tests/v1/spec_decode/test_dflash2.py
+1
@@ -159,6 +159,7 @@ def test_dflash2_model_decoder_layer_cls(monkeypatch):
159 model_config=SimpleNamespace(
160 dtype=torch.float32,
161 is_mm_prefix_lm=False,
162 + rswa_window=None,
163 ),
164 kernel_config=SimpleNamespace(
165 linear_backend="auto",
vllm/platforms/cuda.py
+8
@@ -87,6 +87,7 @@ def _get_backend_priorities(
87 kv_cache_dtype: CacheDType | None = None,
88 use_non_causal: bool = False,
89 head_size: int | None = None,
90 + use_mm_prefix: bool = False,
91 ) -> list[AttentionBackendEnum]:
92 """Get backend priorities with lazy import to avoid circular dependency."""
93 from vllm.utils.torch_utils import is_quantized_kv_cache
@@ -156,6 +157,7 @@ def _get_backend_priorities(
157 # So prefer FlashAttention when non-causal on SM100f.
158 if device_capability.major == 10 and not use_non_causal:
159 return [
160 + *([AttentionBackendEnum.TRITON_FLASHINFER] if use_mm_prefix else []),
161 AttentionBackendEnum.FLASHINFER,
162 AttentionBackendEnum.FLASH_ATTN,
163 AttentionBackendEnum.TRITON_ATTN,
@@ -164,6 +166,11 @@ def _get_backend_priorities(
166 ]
167 else:
168 return [
169 + *(
170 + [AttentionBackendEnum.TRITON_FLASH_ATTN]
171 + if device_capability.major == 9 and use_mm_prefix
172 + else []
173 + ),
174 AttentionBackendEnum.FLASH_ATTN,
175 AttentionBackendEnum.FLASHINFER,
176 AttentionBackendEnum.TRITON_ATTN,
@@ -394,6 +401,7 @@ class CudaPlatformBase(Platform):
401 kv_cache_dtype=attn_selector_config.kv_cache_dtype,
402 use_non_causal=attn_selector_config.use_non_causal,
403 head_size=attn_selector_config.head_size,
404 + use_mm_prefix=attn_selector_config.use_mm_prefix,
405 )
406 for priority, backend in enumerate(backend_priorities):
407 try:
vllm/v1/attention/backend.py
+10
@@ -185,6 +185,10 @@ class AttentionBackend(ABC):
185 def supports_sliding_window(cls) -> bool:
186 return False
187
188 + @classmethod
189 + def supports_rswa(cls) -> bool:
190 + return False
191 +
192 @classmethod
193 def supports_non_causal(cls) -> bool:
194 """Check if backend supports non-causal (bidirectional) attention
@@ -287,6 +291,7 @@ class AttentionBackend(ABC):
291 use_pcp: bool = False,
292 use_adaptive_verification: bool = False,
293 use_dcp: bool = False,
294 + use_rswa: bool = False,
295 ) -> list[str]:
296 invalid_reasons = []
297 if not cls.supports_head_size(head_size):
@@ -321,6 +326,8 @@ class AttentionBackend(ABC):
326 invalid_reasons.append(f"attention type {attn_type} not supported")
327 if has_sliding_window and not cls.supports_sliding_window():
328 invalid_reasons.append("sliding window not supported")
329 + if use_rswa and not cls.supports_rswa():
330 + invalid_reasons.append("R-SWA not supported")
331 if use_non_causal and not cls.supports_non_causal():
332 invalid_reasons.append("non-causal attention not supported")
333 if use_mla and use_non_causal and use_dcp and not cls.supports_non_causal_dcp():
@@ -854,6 +861,9 @@ class AttentionImplBase(ABC, Generic[T]):
861 def process_weights_after_loading(self, act_dtype: torch.dtype):
862 pass
863
864 + def get_impl_variants(self) -> tuple["AttentionImplBase", ...]:
865 + return (self,)
866 +
867 def prepare_for_batch(self, attn_metadata: T | None) -> None:
868 """Prepare implementation-specific state for the current batch."""
869
vllm/v1/attention/backends/composite.py new
+509
@@ -0,0 +1,509 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Reusable attention backend composition and routing policies."""
4 +
5 +import functools
6 +import math
7 +from typing import ClassVar, Protocol, cast
8 +
9 +import torch
10 +
11 +from vllm.config import VllmConfig
12 +from vllm.v1.attention.backend import (
13 + AttentionBackend,
14 + AttentionCGSupport,
15 + AttentionImpl,
16 + AttentionMetadataBuilder,
17 + CommonAttentionMetadata,
18 +)
19 +from vllm.v1.attention.backends.utils import get_supported_kv_cache_layouts
20 +from vllm.v1.kv_cache_interface import KVCacheSpec
21 +
22 +
23 +def requires_mm_prefix(
24 + metadata: CommonAttentionMetadata, *, unclamped_window: bool = False
25 +) -> bool:
26 + ranges = metadata.mm_req_doc_ranges
27 + # With one query per request there are no future keys to unmask. This
28 + # also keeps single-token prefills on the captured FlashInfer path.
29 + if not ranges or (metadata.max_query_len <= 1 and not unclamped_window):
30 + return False
31 + seq_lens = metadata.seq_lens_cpu_upper_bound
32 + assert seq_lens is not None
33 + starts = metadata.query_start_loc_cpu
34 + for req_idx, spans in ranges.items():
35 + query_len = int(starts[req_idx + 1] - starts[req_idx])
36 + if query_len == 0 or (query_len == 1 and not unclamped_window):
37 + continue
38 + end = int(seq_lens[req_idx])
39 + begin = end - query_len
40 + # An unclamped image can also expose keys behind the sliding window.
41 + if unclamped_window and any(
42 + start < stop and start < end and stop >= begin for start, stop in spans
43 + ):
44 + return True
45 + # Prefill lengths are exact. Historical image ranges must not divert
46 + # later text/decode queries to the image backend.
47 + if any(
48 + start < stop and start < end - 1 and stop > begin for start, stop in spans
49 + ):
50 + return True
51 + return False
52 +
53 +
54 +def _has_unclamped_window(layers) -> bool:
55 + return any(
56 + getattr(layer, "sliding_window", None) is not None
57 + and not getattr(layer, "mm_prefix_clamp_sliding_window", False)
58 + for layer in layers
59 + )
60 +
61 +
62 +class CompositeAttentionRouting(Protocol):
63 + """Select a child backend and declare routing constraints for composition."""
64 +
65 + capture_variant: ClassVar[int]
66 + """Child used during graph capture; replay must select the same child."""
67 +
68 + def __init__(self, layer_names: list[str], vllm_config: VllmConfig) -> None: ...
69 +
70 + def select(self, metadata: CommonAttentionMetadata) -> int:
71 + """Return 0 for the first child or 1 for the second child."""
72 + ...
73 +
74 + @staticmethod
75 + def get_cudagraph_support(
76 + vllm_config: VllmConfig, kv_cache_spec: KVCacheSpec
77 + ) -> AttentionCGSupport:
78 + """Limit graph support to batches whose routing stays invariant."""
79 + ...
80 +
81 + @staticmethod
82 + def variant_uses_mm_prefix(variant: int) -> bool:
83 + """Whether this child must support the model's image-mask requirement."""
84 + ...
85 +
86 +
87 +class MMPrefixAttentionRouting(CompositeAttentionRouting):
88 + """Route image masks to variant zero and causal queries to variant one."""
89 +
90 + def __init__(self, layer_names: list[str], vllm_config: VllmConfig) -> None:
91 + layers = vllm_config.compilation_config.static_forward_context
92 + self.unclamped_window = _has_unclamped_window(
93 + layers[name] for name in layer_names
94 + )
95 +
96 + def select(self, metadata: CommonAttentionMetadata) -> int:
97 + return int(
98 + not requires_mm_prefix(metadata, unclamped_window=self.unclamped_window)
99 + )
100 +
101 + capture_variant: ClassVar[int] = 1
102 +
103 + @staticmethod
104 + def get_cudagraph_support(
105 + vllm_config: VllmConfig, kv_cache_spec: KVCacheSpec
106 + ) -> AttentionCGSupport:
107 + if _has_unclamped_window(
108 + vllm_config.compilation_config.static_forward_context.values()
109 + ):
110 + return AttentionCGSupport.NEVER
111 + return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
112 +
113 + @staticmethod
114 + def variant_uses_mm_prefix(variant: int) -> bool:
115 + return variant == 0
116 +
117 +
118 +def _intersect_kernel_block_sizes(
119 + first: type[AttentionBackend], second: type[AttentionBackend]
120 +):
121 + from vllm.v1.attention.backend import MultipleOf
122 +
123 + result: list[int | MultipleOf] = []
124 + seen: set[tuple[type, int]] = set()
125 + for lhs in first.get_supported_kernel_block_sizes():
126 + for rhs in second.get_supported_kernel_block_sizes():
127 + candidate: int | MultipleOf | None = None
128 + if isinstance(lhs, MultipleOf) and isinstance(rhs, MultipleOf):
129 + candidate = MultipleOf(math.lcm(lhs.base, rhs.base))
130 + elif isinstance(lhs, MultipleOf) and isinstance(rhs, int):
131 + candidate = rhs if rhs % lhs.base == 0 else None
132 + elif isinstance(lhs, int) and isinstance(rhs, MultipleOf):
133 + candidate = lhs if lhs % rhs.base == 0 else None
134 + elif lhs == rhs:
135 + candidate = lhs
136 + if candidate is None:
137 + continue
138 + value = candidate.base if isinstance(candidate, MultipleOf) else candidate
139 + key = (type(candidate), value)
140 + if key not in seen:
141 + result.append(candidate)
142 + seen.add(key)
143 + return result
144 +
145 +
146 +@functools.cache
147 +def create_composite_attention_backend(
148 + general_backend: type[AttentionBackend],
149 + causal_backend: type[AttentionBackend],
150 + *,
151 + name: str,
152 + backend_name: str,
153 + module: str,
154 + routing_policy: type[CompositeAttentionRouting],
155 + head_sizes: tuple[int, ...] = (),
156 + kernel_block_sizes: tuple[int, ...] = (),
157 + device_major: int | None = None,
158 +) -> type[AttentionBackend]:
159 + """Compose two backends with shared storage and an explicit routing policy."""
160 + if general_backend.full_cls_name() == causal_backend.full_cls_name():
161 + return general_backend
162 + general_impl_cls = general_backend.get_impl_cls()
163 + causal_impl_cls = causal_backend.get_impl_cls()
164 + general_builder_cls = general_backend.get_builder_cls()
165 + causal_builder_cls = causal_backend.get_builder_cls()
166 + if (
167 + general_backend.forward_includes_kv_cache_update
168 + or causal_backend.forward_includes_kv_cache_update
169 + ):
170 + raise ValueError("Composite backends require separate KV cache updates")
171 +
172 + class CompositeAttentionImpl(AttentionImpl):
173 + def __init__(self, *args, **kwargs) -> None:
174 + self.general_impl = cast(AttentionImpl, general_impl_cls(*args, **kwargs))
175 + try:
176 + self.causal_impl = cast(AttentionImpl, causal_impl_cls(*args, **kwargs))
177 + except TypeError as e:
178 + raise ValueError(
179 + f"Causal attention backend {causal_backend.get_name()} "
180 + f"does not accept the attention arguments required by "
181 + f"the general backend {general_backend.get_name()}: {e}"
182 + ) from e
183 + for name in (
184 + "num_heads",
185 + "num_kv_heads",
186 + "head_size",
187 + "scale",
188 + "kv_cache_dtype",
189 + ):
190 + if hasattr(self.general_impl, name):
191 + setattr(self, name, getattr(self.general_impl, name))
192 + self.supports_quant_query_input = (
193 + self.general_impl.supports_quant_query_input
194 + and self.causal_impl.supports_quant_query_input
195 + )
196 + self.supports_pcp = self.general_impl.supports_pcp
197 + self.supports_dcp = self.causal_impl.supports_dcp
198 + self.can_return_lse_for_decode = self.causal_impl.can_return_lse_for_decode
199 + self.lse_base_on_e = self.causal_impl.lse_base_on_e
200 + self.need_to_return_lse_for_decode = (
201 + self.causal_impl.need_to_return_lse_for_decode
202 + )
203 + self.supports_mtp_with_cp_non_trivial_interleave_size = (
204 + self.causal_impl.supports_mtp_with_cp_non_trivial_interleave_size
205 + )
206 +
207 + def get_impl_for_metadata(self, attn_metadata):
208 + return (
209 + self.causal_impl
210 + if getattr(attn_metadata, "_attention_backend_variant", 0)
211 + else self.general_impl
212 + )
213 +
214 + def get_impl_variants(self):
215 + return self.general_impl, self.causal_impl
216 +
217 + def forward(
218 + self,
219 + layer,
220 + query,
221 + key,
222 + value,
223 + kv_cache,
224 + attn_metadata,
225 + output,
226 + output_scale=None,
227 + output_block_scale=None,
228 + ):
229 + impl = self.get_impl_for_metadata(attn_metadata)
230 + return impl.forward(
231 + layer,
232 + query,
233 + key,
234 + value,
235 + kv_cache,
236 + attn_metadata,
237 + output,
238 + output_scale,
239 + output_block_scale,
240 + )
241 +
242 + def process_weights_after_loading(self, act_dtype: torch.dtype):
243 + for impl in self.get_impl_variants():
244 + impl.process_weights_after_loading(act_dtype)
245 +
246 + def fused_output_quant_supported(self, quant_key):
247 + return all(
248 + impl.fused_output_quant_supported(quant_key)
249 + for impl in self.get_impl_variants()
250 + )
251 +
252 + def fused_qk_norm_rope_kvcache_supported(self):
253 + return self.general_impl.fused_qk_norm_rope_kvcache_supported()
254 +
255 + def do_qk_norm_rope_kvcache_update(self, *args, **kwargs):
256 + method = self.general_impl.do_qk_norm_rope_kvcache_update
257 + return method(*args, **kwargs)
258 +
259 + def fused_rope_kvcache_supported(self):
260 + return self.general_impl.fused_rope_kvcache_supported()
261 +
262 + def do_rope_and_kv_cache_update(self, *args, **kwargs):
263 + method = self.general_impl.do_rope_and_kv_cache_update
264 + return method(*args, **kwargs)
265 +
266 + def do_kv_cache_update(self, *args, **kwargs):
267 + method = self.general_impl.do_kv_cache_update # type: ignore[attr-defined]
268 + return method(*args, **kwargs)
269 +
270 + class CompositeAttentionMetadataBuilder(AttentionMetadataBuilder):
271 + requires_block_table_width = (
272 + general_builder_cls.requires_block_table_width
273 + or causal_builder_cls.requires_block_table_width
274 + )
275 +
276 + def __init__(self, kv_cache_spec, layer_names, vllm_config, device, **kwargs):
277 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)
278 + args = (kv_cache_spec, layer_names, vllm_config, device)
279 + # Only the variant that asked for the extra ctor kwargs receives them.
280 + self.general_builder = general_builder_cls(
281 + *args,
282 + **(kwargs if general_builder_cls.requires_block_table_width else {}),
283 + )
284 + self.causal_builder = causal_builder_cls(
285 + *args,
286 + **(kwargs if causal_builder_cls.requires_block_table_width else {}),
287 + )
288 + self.reorder_batch_threshold = min(
289 + (
290 + builder.reorder_batch_threshold
291 + for builder in self._builders
292 + if builder.reorder_batch_threshold is not None
293 + ),
294 + default=None,
295 + )
296 + self.routing = routing_policy(layer_names, vllm_config)
297 +
298 + def set_kernel_block_size(self, kernel_block_size):
299 + super().set_kernel_block_size(kernel_block_size)
300 + for builder in self._builders:
301 + builder.set_kernel_block_size(kernel_block_size)
302 +
303 + @classmethod
304 + def get_cudagraph_support(cls, vllm_config, kv_cache_spec):
305 + return min(
306 + general_builder_cls.get_cudagraph_support(vllm_config, kv_cache_spec),
307 + causal_builder_cls.get_cudagraph_support(vllm_config, kv_cache_spec),
308 + routing_policy.get_cudagraph_support(vllm_config, kv_cache_spec),
309 + key=lambda support: support.value,
310 + )
311 +
312 + def _builder(self, common_attn_metadata):
313 + backend_variant = self.routing.select(common_attn_metadata)
314 + builder = self.causal_builder if backend_variant else self.general_builder
315 + return builder, backend_variant
316 +
317 + def _build(self, method, common_attn_metadata, **kwargs):
318 + builder, backend_variant = self._builder(common_attn_metadata)
319 + metadata = getattr(builder, method)(
320 + common_attn_metadata=common_attn_metadata, **kwargs
321 + )
322 + metadata._attention_backend_variant = backend_variant
323 + return metadata
324 +
325 + def build(self, common_prefix_len, common_attn_metadata, fast_build=False):
326 + return self._build(
327 + "build",
328 + common_attn_metadata,
329 + common_prefix_len=common_prefix_len,
330 + fast_build=fast_build,
331 + )
332 +
333 + def build_for_cudagraph_capture(self, common_attn_metadata):
334 + variant = routing_policy.capture_variant
335 + metadata = self._builders[variant].build_for_cudagraph_capture(
336 + common_attn_metadata
337 + )
338 + metadata._attention_backend_variant = variant
339 + return metadata
340 +
341 + @property
342 + def _builders(self):
343 + return self.general_builder, self.causal_builder
344 +
345 + # Forward the optional workspace protocol so that a wrapped builder
346 + # which allocates one (FlashInfer) still joins the runner's cross-group
347 + # sharing instead of allocating a second buffer behind the composite.
348 + def _get_workspace_buffer(self):
349 + for builder in self._builders:
350 + if hasattr(builder, "_get_workspace_buffer"):
351 + workspace = builder._get_workspace_buffer()
352 + # The runner only calls the getter for the group that
353 + # provides the buffer, so hand it to the sibling here.
354 + self.set_workspace_buffer(workspace)
355 + return workspace
356 + return None
357 +
358 + def set_workspace_buffer(self, workspace_buffer):
359 + for builder in self._builders:
360 + if hasattr(builder, "set_workspace_buffer"):
361 + builder.set_workspace_buffer(workspace_buffer)
362 +
363 + class CompositeAttentionBackend(AttentionBackend):
364 + forward_includes_kv_cache_update = False
365 + general_backend_cls = general_backend
366 + causal_backend_cls = causal_backend
367 +
368 + @staticmethod
369 + def get_name():
370 + return backend_name
371 +
372 + @staticmethod
373 + def get_impl_cls():
374 + return CompositeAttentionImpl
375 +
376 + @staticmethod
377 + def get_builder_cls():
378 + return CompositeAttentionMetadataBuilder
379 +
380 + @staticmethod
381 + def get_supported_kernel_block_sizes():
382 + if kernel_block_sizes:
383 + return [
384 + size
385 + for size in kernel_block_sizes
386 + if all(
387 + backend.supports_block_size(size)
388 + for backend in (general_backend, causal_backend)
389 + )
390 + ]
391 + return _intersect_kernel_block_sizes(general_backend, causal_backend)
392 +
393 + @classmethod
394 + def supports_block_size(cls, block_size):
395 + return bool(
396 + cls.get_supported_kernel_block_sizes()
397 + ) and super().supports_block_size(block_size)
398 +
399 + @classmethod
400 + def get_supported_head_sizes(cls):
401 + return list(head_sizes)
402 +
403 + @classmethod
404 + def supports_head_size(cls, head_size):
405 + return (not head_sizes or head_size in head_sizes) and all(
406 + backend.supports_head_size(head_size)
407 + for backend in (general_backend, causal_backend)
408 + )
409 +
410 + @classmethod
411 + def supports_compute_capability(cls, capability):
412 + return (device_major is None or capability.major == device_major) and all(
413 + backend.supports_compute_capability(capability)
414 + for backend in (general_backend, causal_backend)
415 + )
416 +
417 + @classmethod
418 + def supports_dtype(cls, dtype):
419 + return all(
420 + backend.supports_dtype(dtype)
421 + for backend in (general_backend, causal_backend)
422 + )
423 +
424 + @classmethod
425 + def supports_kv_cache_dtype(cls, dtype):
426 + return all(
427 + backend.supports_kv_cache_dtype(dtype)
428 + for backend in (general_backend, causal_backend)
429 + )
430 +
431 + @classmethod
432 + def supports_mm_prefix(cls):
433 + return all(
434 + not routing_policy.variant_uses_mm_prefix(variant)
435 + or backend.supports_mm_prefix()
436 + for variant, backend in enumerate((general_backend, causal_backend))
437 + )
438 +
439 + @classmethod
440 + def supports_sliding_window(cls):
441 + return all(
442 + backend.supports_sliding_window()
443 + for backend in (general_backend, causal_backend)
444 + )
445 +
446 + @classmethod
447 + def supports_rswa(cls):
448 + return all(
449 + backend.supports_rswa() for backend in (general_backend, causal_backend)
450 + )
451 +
452 + @classmethod
453 + def supports_pcp(cls):
454 + return False
455 +
456 + @classmethod
457 + def supports_dcp(cls):
458 + return False
459 +
460 + @classmethod
461 + def supports_device_cpu_query_lens_mismatch(cls):
462 + return False
463 +
464 + @classmethod
465 + def supports_combination(
466 + cls,
467 + head_size,
468 + dtype,
469 + kv_cache_dtype,
470 + block_size,
471 + use_mla,
472 + has_sink,
473 + use_sparse,
474 + use_mm_prefix,
475 + device_capability,
476 + ):
477 + for variant, backend in enumerate((general_backend, causal_backend)):
478 + reason = backend.supports_combination(
479 + head_size,
480 + dtype,
481 + kv_cache_dtype,
482 + block_size,
483 + use_mla,
484 + has_sink,
485 + use_sparse,
486 + use_mm_prefix and routing_policy.variant_uses_mm_prefix(variant),
487 + device_capability,
488 + )
489 + if reason:
490 + return f"{backend.get_name()}: {reason}"
491 + return None
492 +
493 + @classmethod
494 + def customize_spec(cls, spec):
495 + general_spec = general_backend.customize_spec(spec)
496 + if general_spec != causal_backend.customize_spec(spec):
497 + raise ValueError("Composite backends require the same KV cache spec")
498 + return general_spec
499 +
500 + @classmethod
501 + def supported_kv_cache_layouts(cls):
502 + return tuple(
503 + get_supported_kv_cache_layouts((general_backend, causal_backend))
504 + )
505 +
506 + CompositeAttentionBackend.__name__ = name
507 + CompositeAttentionBackend.__qualname__ = name
508 + CompositeAttentionBackend.__module__ = module
509 + return CompositeAttentionBackend
vllm/v1/attention/backends/flash_attn.py
+4
@@ -129,6 +129,10 @@ class FlashAttentionBackend(AttentionBackend):
129 def get_name() -> str:
130 return "FLASH_ATTN"
131
132 + @classmethod
133 + def supports_rswa(cls) -> bool:
134 + return True
135 +
136 @classmethod
137 def supports_sliding_window(cls) -> bool:
138 return True
vllm/v1/attention/backends/flex_attention.py
+4
@@ -105,6 +105,10 @@ class FlexAttentionBackend(AttentionBackend):
105 def get_name() -> str:
106 return "FLEX_ATTENTION"
107
108 + @classmethod
109 + def supports_rswa(cls) -> bool:
110 + return True
111 +
112 @classmethod
113 def supports_sliding_window(cls) -> bool:
114 return True
vllm/v1/attention/backends/registry.py
+6
@@ -49,6 +49,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
49 TRITON_ATTN_DIFFKV = (
50 "vllm.v1.attention.backends.triton_attn_diffkv.TritonAttentionDiffKVBackend"
51 )
52 + TRITON_FLASH_ATTN = (
53 + "vllm.v1.attention.backends.triton_flash_attn.TritonFlashAttentionBackend"
54 + )
55 ROCM_ATTN = "vllm.v1.attention.backends.rocm_attn.RocmAttentionBackend"
56 ROCM_AITER_MLA = "vllm.v1.attention.backends.mla.rocm_aiter_mla.AiterMLABackend"
57 ROCM_AITER_TRITON_MLA = (
@@ -63,6 +66,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
66 XPU_MLA_SPARSE = "vllm.v1.attention.backends.mla.xpu_mla_sparse.XPUMLASparseBackend"
67 TORCH_SDPA = "" # this tag is only used for ViT
68 FLASHINFER = "vllm.v1.attention.backends.flashinfer.FlashInferBackend"
69 + TRITON_FLASHINFER = (
70 + "vllm.v1.attention.backends.triton_flashinfer.TritonFlashInferBackend"
71 + )
72 FLASHINFER_MLA = (
73 "vllm.v1.attention.backends.mla.flashinfer_mla.FlashInferMLABackend"
74 )
vllm/v1/attention/backends/triton_attn.py
+4
@@ -323,6 +323,10 @@ class TritonAttentionBackend(AttentionBackend):
323 def get_name() -> str:
324 return "TRITON_ATTN"
325
326 + @classmethod
327 + def supports_rswa(cls) -> bool:
328 + return True
329 +
330 @classmethod
331 def supports_sliding_window(cls) -> bool:
332 return True
vllm/v1/attention/backends/triton_flash_attn.py new
+67
@@ -0,0 +1,67 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Triton image-mask and FlashAttention causal attention composite."""
4 +
5 +import torch
6 +
7 +from vllm.config.cache import CacheDType
8 +from vllm.platforms.interface import DeviceCapability
9 +from vllm.v1.attention.backends.composite import (
10 + MMPrefixAttentionRouting,
11 + create_composite_attention_backend,
12 +)
13 +from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
14 +from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
15 +from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend
16 +
17 +
18 +class _FA4Backend(FlashAttentionBackend):
19 + @classmethod
20 + def supports_combination(
21 + cls,
22 + head_size: int,
23 + dtype: torch.dtype,
24 + kv_cache_dtype: CacheDType | None,
25 + block_size: int | None,
26 + use_mla: bool,
27 + has_sink: bool,
28 + use_sparse: bool,
29 + use_mm_prefix: bool,
30 + device_capability: DeviceCapability,
31 + ) -> str | None:
32 + reason = super().supports_combination(
33 + head_size,
34 + dtype,
35 + kv_cache_dtype,
36 + block_size,
37 + use_mla,
38 + has_sink,
39 + use_sparse,
40 + use_mm_prefix,
41 + device_capability,
42 + )
43 + if reason is not None:
44 + return reason
45 + if (
46 + get_flash_attn_version(
47 + head_size=head_size,
48 + has_sinks=has_sink,
49 + kv_cache_block_size=block_size,
50 + supports_fa4_hd256=True,
51 + )
52 + != 4
53 + ):
54 + return "causal route requires FlashAttention v4"
55 + return None
56 +
57 +
58 +TritonFlashAttentionBackend = create_composite_attention_backend(
59 + TritonAttentionBackend,
60 + _FA4Backend,
61 + name="TritonFlashAttentionBackend",
62 + backend_name="TRITON_FLASH_ATTN",
63 + module=__name__,
64 + routing_policy=MMPrefixAttentionRouting,
65 + head_sizes=(256, 512),
66 + device_major=9,
67 +)
vllm/v1/attention/backends/triton_flashinfer.py new
+22
@@ -0,0 +1,22 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Triton image-mask and FlashInfer causal attention composite."""
4 +
5 +from vllm.v1.attention.backends.composite import (
6 + MMPrefixAttentionRouting,
7 + create_composite_attention_backend,
8 +)
9 +from vllm.v1.attention.backends.flashinfer import FlashInferBackend
10 +from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend
11 +
12 +TritonFlashInferBackend = create_composite_attention_backend(
13 + TritonAttentionBackend,
14 + FlashInferBackend,
15 + name="TritonFlashInferBackend",
16 + backend_name="TRITON_FLASHINFER",
17 + module=__name__,
18 + routing_policy=MMPrefixAttentionRouting,
19 + head_sizes=(256, 512),
20 + kernel_block_sizes=(64,),
21 + device_major=10,
22 +)
vllm/v1/attention/backends/utils.py
+19 -2
@@ -8,6 +8,7 @@ from typing import (
8 TYPE_CHECKING,
9 Any,
10 Protocol,
11 + TypeVar,
12 cast,
13 )
14
@@ -45,6 +46,19 @@ logger = init_logger(__name__)
46 PAD_SLOT_ID = -1
47 NULL_BLOCK_ID = 0
48
49 +
50 +ImplT = TypeVar("ImplT", bound=AttentionImpl)
51 +
52 +
53 +def find_attention_impl_variant(
54 + impl: AttentionImpl, impl_cls: type[ImplT]
55 +) -> ImplT | None:
56 + for variant in impl.get_impl_variants():
57 + if isinstance(variant, impl_cls):
58 + return cast(ImplT, variant)
59 + return None
60 +
61 +
62 _LN_2 = math.log(2.0)
63
64
@@ -342,8 +356,8 @@ def get_per_layer_parameters(
356 per_layer_params: dict[str, PerLayerParameters] = {}
357
358 for key, layer in layers.items():
345 - impl = layer.impl
346 - assert isinstance(impl, cls_)
359 + impl = find_attention_impl_variant(layer.impl, cls_)
360 + assert impl is not None
361
362 # Infer hyperparameters from the attention layer
363 window_size = getattr(impl, "sliding_window", None)
@@ -1013,6 +1027,9 @@ def create_fast_prefill_custom_backend(
1027 common_attn_metadata.logits_indices_padded
1028 )
1029 self.num_logits_indices = common_attn_metadata.num_logits_indices
1030 + self._attention_backend_variant = getattr(
1031 + metadata, "_attention_backend_variant", 0
1032 + )
1033
1034 return KVSharingFastPrefillAttentionMetadata(metadata, common_attn_metadata)
1035
vllm/v1/attention/selector.py
+7 -1
@@ -36,6 +36,7 @@ class AttentionSelectorConfig(NamedTuple):
36 use_pcp: bool = False
37 use_adaptive_verification: bool = False
38 use_dcp: bool = False
39 + use_rswa: bool = False
40
41 def __repr__(self):
42 return (
@@ -55,7 +56,8 @@ class AttentionSelectorConfig(NamedTuple):
56 f"use_kv_connector={self.use_kv_connector}, "
57 f"use_adaptive_verification={self.use_adaptive_verification}, "
58 f"use_pcp={self.use_pcp}, "
58 - f"use_dcp={self.use_dcp})"
59 + f"use_dcp={self.use_dcp}, "
60 + f"use_rswa={self.use_rswa})"
61 )
62
63
@@ -168,6 +170,10 @@ def get_attn_backend(
170 use_pcp=vllm_config.parallel_config.prefill_context_parallel_size > 1,
171 use_adaptive_verification=use_adaptive_verification,
172 use_dcp=vllm_config.parallel_config.decode_context_parallel_size > 1,
173 + use_rswa=(
174 + vllm_config.model_config is not None
175 + and vllm_config.model_config.rswa_window is not None
176 + ),
177 )
178
179 # A per-KV-group override (keyed by KVCacheSpecKind) takes precedence over