vllm.v1.attention.backends.composite

Reusable attention backend composition and routing policies.

Classes:

• CompositeAttentionRouting –

  Select a child backend and declare routing constraints for composition.

• MMPrefixAttentionRouting –

  Route image masks to variant zero and causal queries to variant one.

Functions:

• create_composite_attention_backend –

  Compose two backends with shared storage and an explicit routing policy.

CompositeAttentionRouting

Bases: Protocol

Select a child backend and declare routing constraints for composition.

Methods:

• get_cudagraph_support –

  Limit graph support to batches whose routing stays invariant.

• select –

  Return 0 for the first child or 1 for the second child.

• variant_uses_mm_prefix –

  Whether this child must support the model's image-mask requirement.

Attributes:

• capture_variant (int) –

  Child used during graph capture; replay must select the same child.

Source code in vllm/v1/attention/backends/composite.py

class CompositeAttentionRouting(Protocol):
    """Select a child backend and declare routing constraints for composition."""

    capture_variant: ClassVar[int]
    """Child used during graph capture; replay must select the same child."""

    def __init__(self, layer_names: list[str], vllm_config: VllmConfig) -> None: ...

    def select(self, metadata: CommonAttentionMetadata) -> int:
        """Return 0 for the first child or 1 for the second child."""
        ...

    @staticmethod
    def get_cudagraph_support(
        vllm_config: VllmConfig, kv_cache_spec: KVCacheSpec
    ) -> AttentionCGSupport:
        """Limit graph support to batches whose routing stays invariant."""
        ...

    @staticmethod
    def variant_uses_mm_prefix(variant: int) -> bool:
        """Whether this child must support the model's image-mask requirement."""
        ...

capture_variant class-attribute

Child used during graph capture; replay must select the same child.

get_cudagraph_support(vllm_config, kv_cache_spec) staticmethod

Limit graph support to batches whose routing stays invariant.

Source code in vllm/v1/attention/backends/composite.py

@staticmethod
def get_cudagraph_support(
    vllm_config: VllmConfig, kv_cache_spec: KVCacheSpec
) -> AttentionCGSupport:
    """Limit graph support to batches whose routing stays invariant."""
    ...

select(metadata)

Return 0 for the first child or 1 for the second child.

Source code in vllm/v1/attention/backends/composite.py

def select(self, metadata: CommonAttentionMetadata) -> int:
    """Return 0 for the first child or 1 for the second child."""
    ...

variant_uses_mm_prefix(variant) staticmethod

Whether this child must support the model's image-mask requirement.

Source code in vllm/v1/attention/backends/composite.py

@staticmethod
def variant_uses_mm_prefix(variant: int) -> bool:
    """Whether this child must support the model's image-mask requirement."""
    ...

MMPrefixAttentionRouting

Bases: CompositeAttentionRouting

Route image masks to variant zero and causal queries to variant one.

Source code in vllm/v1/attention/backends/composite.py

class MMPrefixAttentionRouting(CompositeAttentionRouting):
    """Route image masks to variant zero and causal queries to variant one."""

    def __init__(self, layer_names: list[str], vllm_config: VllmConfig) -> None:
        layers = vllm_config.compilation_config.static_forward_context
        self.unclamped_window = _has_unclamped_window(
            layers[name] for name in layer_names
        )

    def select(self, metadata: CommonAttentionMetadata) -> int:
        return int(
            not requires_mm_prefix(metadata, unclamped_window=self.unclamped_window)
        )

    capture_variant: ClassVar[int] = 1

    @staticmethod
    def get_cudagraph_support(
        vllm_config: VllmConfig, kv_cache_spec: KVCacheSpec
    ) -> AttentionCGSupport:
        if _has_unclamped_window(
            vllm_config.compilation_config.static_forward_context.values()
        ):
            return AttentionCGSupport.NEVER
        return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE

    @staticmethod
    def variant_uses_mm_prefix(variant: int) -> bool:
        return variant == 0

create_composite_attention_backend(general_backend, causal_backend, *, name, backend_name, module, routing_policy, head_sizes=(), kernel_block_sizes=(), device_major=None) cached

Compose two backends with shared storage and an explicit routing policy.

Source code in vllm/v1/attention/backends/composite.py

@functools.cache
def create_composite_attention_backend(
    general_backend: type[AttentionBackend],
    causal_backend: type[AttentionBackend],
    *,
    name: str,
    backend_name: str,
    module: str,
    routing_policy: type[CompositeAttentionRouting],
    head_sizes: tuple[int, ...] = (),
    kernel_block_sizes: tuple[int, ...] = (),
    device_major: int | None = None,
) -> type[AttentionBackend]:
    """Compose two backends with shared storage and an explicit routing policy."""
    if general_backend.full_cls_name() == causal_backend.full_cls_name():
        return general_backend
    general_impl_cls = general_backend.get_impl_cls()
    causal_impl_cls = causal_backend.get_impl_cls()
    general_builder_cls = general_backend.get_builder_cls()
    causal_builder_cls = causal_backend.get_builder_cls()
    if (
        general_backend.forward_includes_kv_cache_update
        or causal_backend.forward_includes_kv_cache_update
    ):
        raise ValueError("Composite backends require separate KV cache updates")

    class CompositeAttentionImpl(AttentionImpl):
        def __init__(self, *args, **kwargs) -> None:
            self.general_impl = cast(AttentionImpl, general_impl_cls(*args, **kwargs))
            try:
                self.causal_impl = cast(AttentionImpl, causal_impl_cls(*args, **kwargs))
            except TypeError as e:
                raise ValueError(
                    f"Causal attention backend {causal_backend.get_name()} "
                    f"does not accept the attention arguments required by "
                    f"the general backend {general_backend.get_name()}: {e}"
                ) from e
            for name in (
                "num_heads",
                "num_kv_heads",
                "head_size",
                "scale",
                "kv_cache_dtype",
            ):
                if hasattr(self.general_impl, name):
                    setattr(self, name, getattr(self.general_impl, name))
            self.supports_quant_query_input = (
                self.general_impl.supports_quant_query_input
                and self.causal_impl.supports_quant_query_input
            )
            self.supports_pcp = self.general_impl.supports_pcp
            self.supports_dcp = self.causal_impl.supports_dcp
            self.can_return_lse_for_decode = self.causal_impl.can_return_lse_for_decode
            self.lse_base_on_e = self.causal_impl.lse_base_on_e
            self.need_to_return_lse_for_decode = (
                self.causal_impl.need_to_return_lse_for_decode
            )
            self.supports_mtp_with_cp_non_trivial_interleave_size = (
                self.causal_impl.supports_mtp_with_cp_non_trivial_interleave_size
            )

        def get_impl_for_metadata(self, attn_metadata):
            return (
                self.causal_impl
                if getattr(attn_metadata, "_attention_backend_variant", 0)
                else self.general_impl
            )

        def get_impl_variants(self):
            return self.general_impl, self.causal_impl

        def forward(
            self,
            layer,
            query,
            key,
            value,
            kv_cache,
            attn_metadata,
            output,
            output_scale=None,
            output_block_scale=None,
        ):
            impl = self.get_impl_for_metadata(attn_metadata)
            return impl.forward(
                layer,
                query,
                key,
                value,
                kv_cache,
                attn_metadata,
                output,
                output_scale,
                output_block_scale,
            )

        def process_weights_after_loading(self, act_dtype: torch.dtype):
            for impl in self.get_impl_variants():
                impl.process_weights_after_loading(act_dtype)

        def fused_output_quant_supported(self, quant_key):
            return all(
                impl.fused_output_quant_supported(quant_key)
                for impl in self.get_impl_variants()
            )

        def fused_qk_norm_rope_kvcache_supported(self):
            return self.general_impl.fused_qk_norm_rope_kvcache_supported()

        def do_qk_norm_rope_kvcache_update(self, *args, **kwargs):
            method = self.general_impl.do_qk_norm_rope_kvcache_update
            return method(*args, **kwargs)

        def fused_rope_kvcache_supported(self):
            return self.general_impl.fused_rope_kvcache_supported()

        def do_rope_and_kv_cache_update(self, *args, **kwargs):
            method = self.general_impl.do_rope_and_kv_cache_update
            return method(*args, **kwargs)

        def do_kv_cache_update(self, *args, **kwargs):
            method = self.general_impl.do_kv_cache_update  # type: ignore[attr-defined]
            return method(*args, **kwargs)

    class CompositeAttentionMetadataBuilder(AttentionMetadataBuilder):
        requires_block_table_width = (
            general_builder_cls.requires_block_table_width
            or causal_builder_cls.requires_block_table_width
        )

        def __init__(self, kv_cache_spec, layer_names, vllm_config, device, **kwargs):
            super().__init__(kv_cache_spec, layer_names, vllm_config, device)
            args = (kv_cache_spec, layer_names, vllm_config, device)
            # Only the variant that asked for the extra ctor kwargs receives them.
            self.general_builder = general_builder_cls(
                *args,
                **(kwargs if general_builder_cls.requires_block_table_width else {}),
            )
            self.causal_builder = causal_builder_cls(
                *args,
                **(kwargs if causal_builder_cls.requires_block_table_width else {}),
            )
            self.reorder_batch_threshold = min(
                (
                    builder.reorder_batch_threshold
                    for builder in self._builders
                    if builder.reorder_batch_threshold is not None
                ),
                default=None,
            )
            self.routing = routing_policy(layer_names, vllm_config)

        def set_kernel_block_size(self, kernel_block_size):
            super().set_kernel_block_size(kernel_block_size)
            for builder in self._builders:
                builder.set_kernel_block_size(kernel_block_size)

        @classmethod
        def get_cudagraph_support(cls, vllm_config, kv_cache_spec):
            return min(
                general_builder_cls.get_cudagraph_support(vllm_config, kv_cache_spec),
                causal_builder_cls.get_cudagraph_support(vllm_config, kv_cache_spec),
                routing_policy.get_cudagraph_support(vllm_config, kv_cache_spec),
                key=lambda support: support.value,
            )

        def _builder(self, common_attn_metadata):
            backend_variant = self.routing.select(common_attn_metadata)
            builder = self.causal_builder if backend_variant else self.general_builder
            return builder, backend_variant

        def _build(self, method, common_attn_metadata, **kwargs):
            builder, backend_variant = self._builder(common_attn_metadata)
            metadata = getattr(builder, method)(
                common_attn_metadata=common_attn_metadata, **kwargs
            )
            metadata._attention_backend_variant = backend_variant
            return metadata

        def build(self, common_prefix_len, common_attn_metadata, fast_build=False):
            return self._build(
                "build",
                common_attn_metadata,
                common_prefix_len=common_prefix_len,
                fast_build=fast_build,
            )

        def build_for_cudagraph_capture(self, common_attn_metadata):
            variant = routing_policy.capture_variant
            metadata = self._builders[variant].build_for_cudagraph_capture(
                common_attn_metadata
            )
            metadata._attention_backend_variant = variant
            return metadata

        @property
        def _builders(self):
            return self.general_builder, self.causal_builder

        # Forward the optional workspace protocol so that a wrapped builder
        # which allocates one (FlashInfer) still joins the runner's cross-group
        # sharing instead of allocating a second buffer behind the composite.
        def _get_workspace_buffer(self):
            for builder in self._builders:
                if hasattr(builder, "_get_workspace_buffer"):
                    workspace = builder._get_workspace_buffer()
                    # The runner only calls the getter for the group that
                    # provides the buffer, so hand it to the sibling here.
                    self.set_workspace_buffer(workspace)
                    return workspace
            return None

        def set_workspace_buffer(self, workspace_buffer):
            for builder in self._builders:
                if hasattr(builder, "set_workspace_buffer"):
                    builder.set_workspace_buffer(workspace_buffer)

    class CompositeAttentionBackend(AttentionBackend):
        forward_includes_kv_cache_update = False
        general_backend_cls = general_backend
        causal_backend_cls = causal_backend

        @staticmethod
        def get_name():
            return backend_name

        @staticmethod
        def get_impl_cls():
            return CompositeAttentionImpl

        @staticmethod
        def get_builder_cls():
            return CompositeAttentionMetadataBuilder

        @staticmethod
        def get_supported_kernel_block_sizes(
            kv_cache_spec: KVCacheSpec | None = None,
        ):
            if kernel_block_sizes:
                return [
                    size
                    for size in kernel_block_sizes
                    if all(
                        backend.supports_block_size(size)
                        for backend in (general_backend, causal_backend)
                    )
                ]
            return _intersect_kernel_block_sizes(
                general_backend, causal_backend, kv_cache_spec
            )

        @classmethod
        def supports_block_size(cls, block_size):
            return bool(
                cls.get_supported_kernel_block_sizes()
            ) and super().supports_block_size(block_size)

        @classmethod
        def get_supported_head_sizes(cls):
            return list(head_sizes)

        @classmethod
        def supports_head_size(cls, head_size):
            return (not head_sizes or head_size in head_sizes) and all(
                backend.supports_head_size(head_size)
                for backend in (general_backend, causal_backend)
            )

        @classmethod
        def supports_compute_capability(cls, capability):
            return (device_major is None or capability.major == device_major) and all(
                backend.supports_compute_capability(capability)
                for backend in (general_backend, causal_backend)
            )

        @classmethod
        def supports_dtype(cls, dtype):
            return all(
                backend.supports_dtype(dtype)
                for backend in (general_backend, causal_backend)
            )

        @classmethod
        def supports_kv_cache_dtype(cls, dtype):
            return all(
                backend.supports_kv_cache_dtype(dtype)
                for backend in (general_backend, causal_backend)
            )

        @classmethod
        def supports_mm_prefix(cls):
            return all(
                not routing_policy.variant_uses_mm_prefix(variant)
                or backend.supports_mm_prefix()
                for variant, backend in enumerate((general_backend, causal_backend))
            )

        @classmethod
        def supports_sliding_window(cls):
            return all(
                backend.supports_sliding_window()
                for backend in (general_backend, causal_backend)
            )

        @classmethod
        def supports_rswa(cls):
            return all(
                backend.supports_rswa() for backend in (general_backend, causal_backend)
            )

        @classmethod
        def supports_pcp(cls):
            return False

        @classmethod
        def supports_dcp(cls):
            return False

        @classmethod
        def supports_device_cpu_query_lens_mismatch(cls):
            return False

        @classmethod
        def supports_combination(
            cls,
            head_size,
            dtype,
            kv_cache_dtype,
            block_size,
            use_mla,
            has_sink,
            use_sparse,
            use_mm_prefix,
            device_capability,
        ):
            for variant, backend in enumerate((general_backend, causal_backend)):
                reason = backend.supports_combination(
                    head_size,
                    dtype,
                    kv_cache_dtype,
                    block_size,
                    use_mla,
                    has_sink,
                    use_sparse,
                    use_mm_prefix and routing_policy.variant_uses_mm_prefix(variant),
                    device_capability,
                )
                if reason:
                    return f"{backend.get_name()}: {reason}"
            return None

        @classmethod
        def customize_spec(cls, spec):
            general_spec = general_backend.customize_spec(spec)
            if general_spec != causal_backend.customize_spec(spec):
                raise ValueError("Composite backends require the same KV cache spec")
            return general_spec

        @classmethod
        def supported_kv_cache_layouts(cls):
            return tuple(
                get_supported_kv_cache_layouts((general_backend, causal_backend))
            )

    CompositeAttentionBackend.__name__ = name
    CompositeAttentionBackend.__qualname__ = name
    CompositeAttentionBackend.__module__ = module
    return CompositeAttentionBackend