vllm.v1.attention.backend

Classes:

• AttentionBackend –

  Abstract class for attention backends.

• AttentionCGSupport –

  Constants for the cudagraph support of the attention backend

• AttentionImpl –

  Standard attention implementation with forward method.

• AttentionImplBase –

  Base class for attention implementations.

• AttentionMetadataBuilder –
• AttentionType –

  Attention type.

• CommonAttentionMetadata –

  Per-batch attention metadata, shared across layers and backends.

• MLAAttentionImpl –

  MLA attention implementation with forward_mqa and forward_mha methods.

Functions:

• max_decode_query_len –

  Widest request a spec-as-decode builder treats as a decode.

• subclass_attention_backend –

  Return a new subclass where get_builder_cls returns builder_cls.

AttentionBackend

Bases: ABC

Abstract class for attention backends.

Methods:

• customize_spec –

  Adjust the layer's KV cache spec for this backend's kernels. Used when the

• supported_kv_cache_layouts –

  Layouts this backend's kernels can consume, most preferred first, or

• supports_attn_type –

  Check if backend supports a given attention type.

• supports_device_cpu_query_lens_mismatch –

  Whether this backend can run a batch whose device query_start_loc disagrees

• supports_non_causal –

  Check if backend supports non-causal (bidirectional) attention

Source code in vllm/v1/attention/backend.py

class AttentionBackend(ABC):
    """Abstract class for attention backends."""

    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
    supported_kv_cache_dtypes: ClassVar[list["CacheDType"]] = [
        "auto",
        "float16",
        "bfloat16",
    ]

    # Does attention's forward() include kv cache update?
    forward_includes_kv_cache_update: bool = True

    @staticmethod
    def get_supported_kernel_block_sizes(
        kv_cache_spec: "KVCacheSpec | None" = None,
    ) -> list[int | MultipleOf]:
        return [MultipleOf(1)]

    @staticmethod
    @abstractmethod
    def get_name() -> str:
        raise NotImplementedError

    @staticmethod
    @abstractmethod
    def get_impl_cls() -> type["AttentionImplBase"]:
        raise NotImplementedError

    @staticmethod
    @abstractmethod
    def get_builder_cls():  # -> Type["AttentionMetadataBuilder"]:
        raise NotImplementedError

    @classmethod
    def full_cls_name(cls) -> tuple[str, str]:
        return (cls.__module__, cls.__qualname__)

    @classmethod
    def get_supported_head_sizes(cls) -> list[int]:
        return []

    @classmethod
    def supports_head_size(cls, head_size: int) -> bool:
        supported_head_sizes = cls.get_supported_head_sizes()
        return (not supported_head_sizes) or head_size in supported_head_sizes

    @classmethod
    def supports_dtype(cls, dtype: torch.dtype) -> bool:
        return dtype in cls.supported_dtypes

    @classmethod
    def supports_kv_cache_dtype(cls, kv_cache_dtype: "CacheDType | None") -> bool:
        if kv_cache_dtype is None:
            return True
        return (not cls.supported_kv_cache_dtypes) or (
            kv_cache_dtype in cls.supported_kv_cache_dtypes
        )

    @classmethod
    def supports_block_size(cls, block_size: int | None) -> bool:
        if block_size is None:
            return True

        supported_kernel_block_sizes = cls.get_supported_kernel_block_sizes()
        if not supported_kernel_block_sizes:
            return True

        for supported_size in supported_kernel_block_sizes:
            if isinstance(supported_size, MultipleOf):
                supported_size = supported_size.base
            # With hybrid_blocks feature, the framework-level block size
            # only needs to be a multiple of the kernel's requirement,
            # even if the kernel requires a fixed block_size.
            if block_size % supported_size == 0:
                return True
        return False

    @classmethod
    def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec":
        """Adjust the layer's KV cache spec for this backend's kernels. Used when the
        kernels want KV packed in a specific way.

        NOTE: temporary compatibility API. Today the Attention layer builds the spec
        from the model config and the backend only gets to adjust it post-hoc; the end
        state is for the backend to build and return the spec directly, at which point
        this hook goes away.

        (see: https://github.com/vllm-project/vllm/issues/42449)
        """
        return spec

    @classmethod
    def get_preferred_block_size(cls, default_block_size: int) -> int:
        supported_sizes = cls.get_supported_kernel_block_sizes()
        if not supported_sizes:
            return default_block_size

        if cls.supports_block_size(default_block_size):
            return default_block_size

        return min(s.base if isinstance(s, MultipleOf) else s for s in supported_sizes)

    @classmethod
    def is_mla(cls) -> bool:
        return False

    @classmethod
    def supports_sink(cls) -> bool:
        return False

    @classmethod
    def supports_alibi_sqrt(cls) -> bool:
        return False

    @classmethod
    def supports_mm_prefix(cls) -> bool:
        return False

    @classmethod
    def is_sparse(cls) -> bool:
        return False

    @classmethod
    def supports_per_head_quant_scales(cls) -> bool:
        return False

    @classmethod
    def supports_sliding_window(cls) -> bool:
        return False

    @classmethod
    def supports_rswa(cls) -> bool:
        return False

    @classmethod
    def supports_non_causal(cls) -> bool:
        """Check if backend supports non-causal (bidirectional) attention
        for decoder models.

        Unlike ENCODER_ONLY attention type which implies a different
        execution model, this refers to non-causal attention within the
        standard paged-KV-cache decoder path.
        """
        return False

    @classmethod
    def supports_batch_invariance(cls) -> bool:
        return False

    @classmethod
    def supports_kv_connector(cls) -> bool:
        return True

    @classmethod
    def supports_device_cpu_query_lens_mismatch(cls) -> bool:
        """Whether this backend can run a batch whose device query_start_loc disagrees
        with the CPU one; backends that plan off the CPU query lengths must opt out.

        Currently only verification requests are affected: adaptive verification trims
        their drafts on device. On the CPU the draft budget is evenly distributed across
        requests, so the total draft budget, the decode/prefill split point and the CPU
        prefill query lengths all stay correct.

        SSM backends opt out: their recurrent-state planning is built from the CPU
        per-request boundaries, which the trimmed batch no longer matches.
        """
        return not cls.is_ssm()

    @classmethod
    def supports_pcp(cls) -> bool:
        try:
            return cls.get_impl_cls().supports_pcp
        except NotImplementedError:
            return False

    @classmethod
    def supports_dcp(cls) -> bool:
        try:
            return cls.get_impl_cls().supports_dcp
        except NotImplementedError:
            return False

    @classmethod
    def supports_non_causal_dcp(cls) -> bool:
        builder_cls = cls.get_builder_cls()
        return bool(getattr(builder_cls, "supports_non_causal_multi_token_dcp", False))

    @classmethod
    def supports_attn_type(cls, attn_type: str) -> bool:
        """Check if backend supports a given attention type.

        By default, only supports decoder attention.
        Backends should override this to support other attention types.
        """
        return attn_type == AttentionType.DECODER

    @classmethod
    def supports_compute_capability(cls, capability: "DeviceCapability") -> bool:
        return True

    @classmethod
    def supports_combination(
        cls,
        head_size: int,
        dtype: torch.dtype,
        kv_cache_dtype: "CacheDType | None",
        block_size: int | None,
        use_mla: bool,
        has_sink: bool,
        use_sparse: bool,
        use_mm_prefix: bool,
        device_capability: "DeviceCapability",
    ) -> str | None:
        return None

    @classmethod
    def validate_configuration(
        cls,
        head_size: int,
        dtype: torch.dtype,
        kv_cache_dtype: "CacheDType | None",
        block_size: int | None,
        use_mla: bool,
        has_sink: bool,
        use_sparse: bool,
        use_mm_prefix: bool,
        use_per_head_quant_scales: bool,
        device_capability: "DeviceCapability",
        attn_type: str,
        has_sliding_window: bool = False,
        use_non_causal: bool = False,
        use_batch_invariant: bool = False,
        use_kv_connector: bool = False,
        use_pcp: bool = False,
        use_adaptive_verification: bool = False,
        use_dcp: bool = False,
        use_rswa: bool = False,
    ) -> list[str]:
        invalid_reasons = []
        if not cls.supports_head_size(head_size):
            invalid_reasons.append("head_size not supported")
        if not cls.supports_dtype(dtype):
            invalid_reasons.append("dtype not supported")
        if not cls.supports_kv_cache_dtype(kv_cache_dtype):
            invalid_reasons.append("kv_cache_dtype not supported")
        if not cls.supports_block_size(block_size):
            invalid_reasons.append("block_size not supported")
        if use_mm_prefix and not cls.supports_mm_prefix():
            invalid_reasons.append(
                "partial multimodal token full attention not supported"
            )
        if use_mla != cls.is_mla():
            if use_mla:
                invalid_reasons.append("MLA not supported")
            else:
                invalid_reasons.append("non-MLA not supported")
        if has_sink and not cls.supports_sink():
            invalid_reasons.append("attention sinks not supported")
        if use_sparse != cls.is_sparse():
            if use_sparse:
                invalid_reasons.append("sparse not supported")
            else:
                invalid_reasons.append("non-sparse not supported")
        if use_per_head_quant_scales and not cls.supports_per_head_quant_scales():
            invalid_reasons.append("per-head quant scales not supported")
        if not cls.supports_compute_capability(device_capability):
            invalid_reasons.append("compute capability not supported")
        if not cls.supports_attn_type(attn_type):
            invalid_reasons.append(f"attention type {attn_type} not supported")
        if has_sliding_window and not cls.supports_sliding_window():
            invalid_reasons.append("sliding window not supported")
        if use_rswa and not cls.supports_rswa():
            invalid_reasons.append("R-SWA not supported")
        if use_non_causal and not cls.supports_non_causal():
            invalid_reasons.append("non-causal attention not supported")
        if use_mla and use_non_causal and use_dcp and not cls.supports_non_causal_dcp():
            invalid_reasons.append("non-causal MLA attention with DCP not supported")
        if use_batch_invariant and not cls.supports_batch_invariance():
            invalid_reasons.append("batch invariance not supported")
        if use_kv_connector and not cls.supports_kv_connector():
            invalid_reasons.append("KV connector not supported")
        if use_pcp and not cls.supports_pcp():
            invalid_reasons.append("PCP not supported")
        if use_dcp and not cls.supports_dcp():
            invalid_reasons.append("DCP not supported")
        if (
            use_adaptive_verification
            and not cls.supports_device_cpu_query_lens_mismatch()
        ):
            invalid_reasons.append(
                "device-cpu query lens mismatch not supported, "
                "this is needed for adaptive verification"
            )
        combination_reason = cls.supports_combination(
            head_size,
            dtype,
            kv_cache_dtype,
            block_size,
            use_mla,
            has_sink,
            use_sparse,
            use_mm_prefix,
            device_capability,
        )
        if combination_reason is not None:
            invalid_reasons.append(combination_reason)
        return invalid_reasons

    @classmethod
    def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...] | None:
        """Layouts this backend's kernels can consume, most preferred first, or
        None when the kernels consume any layout and express no preference."""
        return None

    @classmethod
    def is_ssm(cls) -> bool:
        return False

customize_spec(spec) classmethod

Adjust the layer's KV cache spec for this backend's kernels. Used when the kernels want KV packed in a specific way.

NOTE: temporary compatibility API. Today the Attention layer builds the spec from the model config and the backend only gets to adjust it post-hoc; the end state is for the backend to build and return the spec directly, at which point this hook goes away.

(see: https://github.com/vllm-project/vllm/issues/42449)

Source code in vllm/v1/attention/backend.py

@classmethod
def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec":
    """Adjust the layer's KV cache spec for this backend's kernels. Used when the
    kernels want KV packed in a specific way.

    NOTE: temporary compatibility API. Today the Attention layer builds the spec
    from the model config and the backend only gets to adjust it post-hoc; the end
    state is for the backend to build and return the spec directly, at which point
    this hook goes away.

    (see: https://github.com/vllm-project/vllm/issues/42449)
    """
    return spec

supported_kv_cache_layouts() classmethod

Layouts this backend's kernels can consume, most preferred first, or None when the kernels consume any layout and express no preference.

Source code in vllm/v1/attention/backend.py

@classmethod
def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...] | None:
    """Layouts this backend's kernels can consume, most preferred first, or
    None when the kernels consume any layout and express no preference."""
    return None

supports_attn_type(attn_type) classmethod

Check if backend supports a given attention type.

By default, only supports decoder attention. Backends should override this to support other attention types.

Source code in vllm/v1/attention/backend.py

@classmethod
def supports_attn_type(cls, attn_type: str) -> bool:
    """Check if backend supports a given attention type.

    By default, only supports decoder attention.
    Backends should override this to support other attention types.
    """
    return attn_type == AttentionType.DECODER

supports_device_cpu_query_lens_mismatch() classmethod

Whether this backend can run a batch whose device query_start_loc disagrees with the CPU one; backends that plan off the CPU query lengths must opt out.

Currently only verification requests are affected: adaptive verification trims their drafts on device. On the CPU the draft budget is evenly distributed across requests, so the total draft budget, the decode/prefill split point and the CPU prefill query lengths all stay correct.

SSM backends opt out: their recurrent-state planning is built from the CPU per-request boundaries, which the trimmed batch no longer matches.

Source code in vllm/v1/attention/backend.py

@classmethod
def supports_device_cpu_query_lens_mismatch(cls) -> bool:
    """Whether this backend can run a batch whose device query_start_loc disagrees
    with the CPU one; backends that plan off the CPU query lengths must opt out.

    Currently only verification requests are affected: adaptive verification trims
    their drafts on device. On the CPU the draft budget is evenly distributed across
    requests, so the total draft budget, the decode/prefill split point and the CPU
    prefill query lengths all stay correct.

    SSM backends opt out: their recurrent-state planning is built from the CPU
    per-request boundaries, which the trimmed batch no longer matches.
    """
    return not cls.is_ssm()

supports_non_causal() classmethod

Check if backend supports non-causal (bidirectional) attention for decoder models.

Unlike ENCODER_ONLY attention type which implies a different execution model, this refers to non-causal attention within the standard paged-KV-cache decoder path.

Source code in vllm/v1/attention/backend.py

@classmethod
def supports_non_causal(cls) -> bool:
    """Check if backend supports non-causal (bidirectional) attention
    for decoder models.

    Unlike ENCODER_ONLY attention type which implies a different
    execution model, this refers to non-causal attention within the
    standard paged-KV-cache decoder path.
    """
    return False

AttentionCGSupport

Bases: Enum

Constants for the cudagraph support of the attention backend Here we do not consider the cascade attention, as currently it is never cudagraph supported.

Attributes:

• ALWAYS –

  Cudagraph always supported; supports mixed-prefill-decode

• NEVER –

  NO cudagraph support

• UNIFORM_BATCH –

  Cudagraph supported for batches the only contain query lengths that are

• UNIFORM_SINGLE_TOKEN_DECODE –

  Cudagraph supported for batches the only contain query_len==1 decodes

Source code in vllm/v1/attention/backend.py

class AttentionCGSupport(Enum):
    """Constants for the cudagraph support of the attention backend
    Here we do not consider the cascade attention, as currently
    it is never cudagraph supported."""

    ALWAYS = 3
    """Cudagraph always supported; supports mixed-prefill-decode"""
    UNIFORM_BATCH = 2
    """Cudagraph supported for batches the only contain query lengths that are
    the same, this can be used for spec-decode
        i.e. "decodes" are 1 + num_speculative_tokens"""
    UNIFORM_SINGLE_TOKEN_DECODE = 1
    """Cudagraph supported for batches the only contain query_len==1 decodes"""
    NEVER = 0
    """NO cudagraph support"""

ALWAYS = 3 class-attribute instance-attribute

Cudagraph always supported; supports mixed-prefill-decode

NEVER = 0 class-attribute instance-attribute

NO cudagraph support

UNIFORM_BATCH = 2 class-attribute instance-attribute

Cudagraph supported for batches the only contain query lengths that are the same, this can be used for spec-decode i.e. "decodes" are 1 + num_speculative_tokens

UNIFORM_SINGLE_TOKEN_DECODE = 1 class-attribute instance-attribute

Cudagraph supported for batches the only contain query_len==1 decodes

AttentionImpl

Bases: AttentionImplBase[T], Generic[T]

Standard attention implementation with forward method.

Methods:

• do_qk_norm_mrope_kvcache_update –

  Apply QK-norm and MRoPE, then write K/V to the cache.

• do_qk_norm_rope_kvcache_update –

  If fused_qk_norm_rope_kvcache_supported returns True, this method

• do_rope_and_kv_cache_update –

  If fused_rope_kvcache_supported returns True, this method will be called

• fused_output_quant_supported –

  Does this attention implementation support fused output quantization.

• fused_qk_norm_mrope_kvcache_supported –

  Whether this implementation supports fused QKNorm+MRoPE+KVCache.

• fused_qk_norm_rope_kvcache_supported –

  Does this attention implementation support fused QKNorm+RoPE+KVCache fusion.

• fused_rope_kvcache_supported –

  Does this attention implementation support RoPE+KVCache fusion.

Attributes:

• kv_quant_mode (KVQuantMode) –

  Return the KV cache quantization mode for this layer.

Source code in vllm/v1/attention/backend.py

class AttentionImpl(AttentionImplBase[T], Generic[T]):
    """Standard attention implementation with forward method."""

    kv_cache_dtype: str

    @property
    def kv_quant_mode(self) -> "KVQuantMode":
        """Return the KV cache quantization mode for this layer."""
        return get_kv_quant_mode(self.kv_cache_dtype)

    @abstractmethod
    def __init__(
        self,
        num_heads: int,
        head_size: int,
        scale: float,
        num_kv_heads: int | None = None,
        alibi_slopes: list[float] | None = None,
        sliding_window: int | None = None,
        kv_cache_dtype: str = "auto",
        logits_soft_cap: float | None = None,
        attn_type: str = AttentionType.DECODER,
        kv_sharing_target_layer_name: str | None = None,
    ) -> None:
        raise NotImplementedError

    @abstractmethod
    def forward(
        self,
        layer: AttentionLayer,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        kv_cache: torch.Tensor,
        attn_metadata: T,
        output: torch.Tensor,
        output_scale: torch.Tensor | None = None,
        output_block_scale: torch.Tensor | None = None,
    ) -> torch.Tensor:
        raise NotImplementedError

    def fused_output_quant_supported(self, quant_key: "QuantKey") -> bool:
        """Does this attention implementation support fused output quantization.
        This is used by the AttnFusionPass to only fuse output quantization
        onto implementations that support it.

        Args:
            quant_key: QuantKey object that describes the quantization op

        Returns:
            is fusion supported for this type of quantization

        """
        return False

    def fused_qk_norm_rope_kvcache_supported(self):
        """Does this attention implementation support fused QKNorm+RoPE+KVCache fusion.
        This is used by the QkNormRopeKvCachePattern to only fuse the QKNorm ops
        with the RoPE ops and the KV cache update for implementations that support it.
        """
        return False

    def fused_qk_norm_mrope_kvcache_supported(self):
        """Whether this implementation supports fused QKNorm+MRoPE+KVCache."""
        return False

    def fused_rope_kvcache_supported(self):
        """Does this attention implementation support RoPE+KVCache fusion.
        This is used by the RopeKVCacheFusionPass to only fuse the RoPE ops
        with the KV cache update for implementations that support it.
        """
        return False

    def do_qk_norm_rope_kvcache_update(
        self,
        layer: AttentionLayer,
        qkv: torch.Tensor,
        q_out: torch.Tensor,
        k_out: torch.Tensor,
        positions: torch.Tensor,
        q_weight: torch.Tensor,
        k_weight: torch.Tensor,
        rms_norm_eps: float,
        cos_sin_cache: torch.Tensor,
        is_neox: bool,
        kv_cache: torch.Tensor,
        layer_slot_mapping: torch.Tensor,
    ):
        """If `fused_qk_norm_rope_kvcache_supported` returns True, this method
        will be called by the fused custom op. Applies QK-norm + RoPE and
        writes K/V to the KV cache. Results are written to the pre-allocated
        q_out and k_out tensors; V is split from QKV at the graph level.
        """
        raise NotImplementedError

    def do_qk_norm_mrope_kvcache_update(
        self,
        layer: AttentionLayer,
        qkv: torch.Tensor,
        q_out: torch.Tensor,
        positions: torch.Tensor,
        q_weight: torch.Tensor,
        k_weight: torch.Tensor,
        rms_norm_eps: float,
        cos_sin_cache: torch.Tensor,
        is_neox: bool,
        mrope_section: tuple[int, int, int],
        is_interleaved: bool,
        rotary_dim: int,
        kv_cache: torch.Tensor,
        layer_slot_mapping: torch.Tensor,
    ):
        """Apply QK-norm and MRoPE, then write K/V to the cache."""
        raise NotImplementedError

    def do_rope_and_kv_cache_update(
        self,
        layer: AttentionLayer,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        positions: torch.Tensor,
        cos_sin_cache: torch.Tensor,
        is_neox: bool,
        kv_cache: torch.Tensor,
        layer_slot_mapping: torch.Tensor,
    ):
        """If `fused_rope_kvcache_supported` returns True, this method will be called
        by torch.ops.vllm.fused_rope_and_unified_kv_cache_update
        to perform the inplace RoPE and KV cache update.
        """
        raise NotImplementedError

kv_quant_mode property

Return the KV cache quantization mode for this layer.

do_qk_norm_mrope_kvcache_update(layer, qkv, q_out, positions, q_weight, k_weight, rms_norm_eps, cos_sin_cache, is_neox, mrope_section, is_interleaved, rotary_dim, kv_cache, layer_slot_mapping)

Apply QK-norm and MRoPE, then write K/V to the cache.

Source code in vllm/v1/attention/backend.py

def do_qk_norm_mrope_kvcache_update(
    self,
    layer: AttentionLayer,
    qkv: torch.Tensor,
    q_out: torch.Tensor,
    positions: torch.Tensor,
    q_weight: torch.Tensor,
    k_weight: torch.Tensor,
    rms_norm_eps: float,
    cos_sin_cache: torch.Tensor,
    is_neox: bool,
    mrope_section: tuple[int, int, int],
    is_interleaved: bool,
    rotary_dim: int,
    kv_cache: torch.Tensor,
    layer_slot_mapping: torch.Tensor,
):
    """Apply QK-norm and MRoPE, then write K/V to the cache."""
    raise NotImplementedError

do_qk_norm_rope_kvcache_update(layer, qkv, q_out, k_out, positions, q_weight, k_weight, rms_norm_eps, cos_sin_cache, is_neox, kv_cache, layer_slot_mapping)

If fused_qk_norm_rope_kvcache_supported returns True, this method will be called by the fused custom op. Applies QK-norm + RoPE and writes K/V to the KV cache. Results are written to the pre-allocated q_out and k_out tensors; V is split from QKV at the graph level.

Source code in vllm/v1/attention/backend.py

def do_qk_norm_rope_kvcache_update(
    self,
    layer: AttentionLayer,
    qkv: torch.Tensor,
    q_out: torch.Tensor,
    k_out: torch.Tensor,
    positions: torch.Tensor,
    q_weight: torch.Tensor,
    k_weight: torch.Tensor,
    rms_norm_eps: float,
    cos_sin_cache: torch.Tensor,
    is_neox: bool,
    kv_cache: torch.Tensor,
    layer_slot_mapping: torch.Tensor,
):
    """If `fused_qk_norm_rope_kvcache_supported` returns True, this method
    will be called by the fused custom op. Applies QK-norm + RoPE and
    writes K/V to the KV cache. Results are written to the pre-allocated
    q_out and k_out tensors; V is split from QKV at the graph level.
    """
    raise NotImplementedError

do_rope_and_kv_cache_update(layer, query, key, value, positions, cos_sin_cache, is_neox, kv_cache, layer_slot_mapping)

If fused_rope_kvcache_supported returns True, this method will be called by torch.ops.vllm.fused_rope_and_unified_kv_cache_update to perform the inplace RoPE and KV cache update.

Source code in vllm/v1/attention/backend.py

def do_rope_and_kv_cache_update(
    self,
    layer: AttentionLayer,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    positions: torch.Tensor,
    cos_sin_cache: torch.Tensor,
    is_neox: bool,
    kv_cache: torch.Tensor,
    layer_slot_mapping: torch.Tensor,
):
    """If `fused_rope_kvcache_supported` returns True, this method will be called
    by torch.ops.vllm.fused_rope_and_unified_kv_cache_update
    to perform the inplace RoPE and KV cache update.
    """
    raise NotImplementedError

fused_output_quant_supported(quant_key)

Does this attention implementation support fused output quantization. This is used by the AttnFusionPass to only fuse output quantization onto implementations that support it.

Parameters:

• quant_key

  (QuantKey) –

  QuantKey object that describes the quantization op

Returns:

• bool –

  is fusion supported for this type of quantization

Source code in vllm/v1/attention/backend.py

def fused_output_quant_supported(self, quant_key: "QuantKey") -> bool:
    """Does this attention implementation support fused output quantization.
    This is used by the AttnFusionPass to only fuse output quantization
    onto implementations that support it.

    Args:
        quant_key: QuantKey object that describes the quantization op

    Returns:
        is fusion supported for this type of quantization

    """
    return False

fused_qk_norm_mrope_kvcache_supported()

Whether this implementation supports fused QKNorm+MRoPE+KVCache.

Source code in vllm/v1/attention/backend.py

def fused_qk_norm_mrope_kvcache_supported(self):
    """Whether this implementation supports fused QKNorm+MRoPE+KVCache."""
    return False

fused_qk_norm_rope_kvcache_supported()

Does this attention implementation support fused QKNorm+RoPE+KVCache fusion. This is used by the QkNormRopeKvCachePattern to only fuse the QKNorm ops with the RoPE ops and the KV cache update for implementations that support it.

Source code in vllm/v1/attention/backend.py

def fused_qk_norm_rope_kvcache_supported(self):
    """Does this attention implementation support fused QKNorm+RoPE+KVCache fusion.
    This is used by the QkNormRopeKvCachePattern to only fuse the QKNorm ops
    with the RoPE ops and the KV cache update for implementations that support it.
    """
    return False

fused_rope_kvcache_supported()

Does this attention implementation support RoPE+KVCache fusion. This is used by the RopeKVCacheFusionPass to only fuse the RoPE ops with the KV cache update for implementations that support it.

Source code in vllm/v1/attention/backend.py

def fused_rope_kvcache_supported(self):
    """Does this attention implementation support RoPE+KVCache fusion.
    This is used by the RopeKVCacheFusionPass to only fuse the RoPE ops
    with the KV cache update for implementations that support it.
    """
    return False

AttentionImplBase

Bases: ABC, Generic[T]

Base class for attention implementations.

Contains common attributes and initialization logic shared by both standard AttentionImpl and MLAAttentionImpl. Does not define a forward method - subclasses define their own forward interfaces.

Methods:

• prepare_for_batch –

  Prepare implementation-specific state for the current batch.

Source code in vllm/v1/attention/backend.py

class AttentionImplBase(ABC, Generic[T]):
    """Base class for attention implementations.

    Contains common attributes and initialization logic shared by both
    standard AttentionImpl and MLAAttentionImpl. Does not define a forward
    method - subclasses define their own forward interfaces.
    """

    # Whether this impl uses a sparse (top-k) attention path. Used by MLA to
    # route between the dense-MHA prefill and sparse-MQA paths.
    is_sparse: ClassVar[bool] = False

    # Whether this impl provides a dense-MHA prefill path (forward_mha). Sparse
    # impls without one run the top-k MQA path for all requests.
    supports_dense_mha_prefill: ClassVar[bool] = True

    # Required attributes that all impls should have
    num_heads: int
    head_size: int
    scale: float

    # Whether the attention impl can return the softmax lse for decode.
    # Some features like decode context parallelism require the softmax lse.
    can_return_lse_for_decode: bool = False

    # Base of the logarithm used by this backend when returning softmax lse.
    # True  => natural log (lse = ln(sum(exp(qk))))
    #          -- e.g. Triton MLA, FlashAttention, FlashMLA, Cutlass MLA
    # False => base 2      (lse = log2(sum(exp(qk))))
    #          -- e.g. FlashInfer trtllm-gen MLA
    # The DCP combine kernel (cp_lse_ag_out_rs / dcp_a2a_lse_reduce in
    # vllm/v1/attention/ops/dcp.py) branches on this via its IS_BASE_E
    # constexpr; getting it wrong silently corrupts the cross-shard
    # softmax denominator.
    lse_base_on_e: bool = True

    # Whether the attention impl supports Prefill Context Parallelism.
    supports_pcp: bool = False
    # Whether the attention impl supports Decode Context Parallelism.
    supports_dcp: bool = False
    # Whether the attention impl(or ops) supports MTP
    # when cp_kv_cache_interleave_size > 1
    supports_mtp_with_cp_non_trivial_interleave_size: bool = False

    # some attention backends might not always want to return lse
    # even if they can return lse (for efficiency reasons)
    need_to_return_lse_for_decode: bool = False

    # Whether this attention implementation supports pre-quantized query input.
    # When True, the attention layer will quantize queries before passing them
    # to this backend, allowing torch.compile to fuse the quantization with
    # previous operations. This is typically supported when using FP8 KV cache
    # with compatible attention kernels (e.g., TRT-LLM).
    # Subclasses should set this in __init__.
    # TODO add support to more backends:
    # https://github.com/vllm-project/vllm/issues/25584
    supports_quant_query_input: bool = False

    dcp_world_size: int
    dcp_rank: int

    pcp_world_size: int
    pcp_rank: int

    total_cp_world_size: int
    total_cp_rank: int

    def __new__(cls, *args, **kwargs):
        # use __new__ so that all subclasses will call this
        self = super().__new__(cls)
        from vllm.config import get_current_vllm_config_or_none
        from vllm.distributed.parallel_state import get_dcp_world_size_and_rank

        # Replicated drafts run at DCP=1 inside a DCP target's process group.
        config = get_current_vllm_config_or_none()
        self.dcp_world_size, self.dcp_rank = get_dcp_world_size_and_rank(
            config is None or config.parallel_config.decode_context_parallel_size > 1
        )
        try:
            from vllm.distributed.parallel_state import get_pcp_group

            self.pcp_world_size = get_pcp_group().world_size
            self.pcp_rank = get_pcp_group().rank_in_group
        except AssertionError:
            self.pcp_world_size = 1
            self.pcp_rank = 0
        self.total_cp_world_size = self.dcp_world_size
        self.total_cp_rank = self.dcp_rank

        self.need_to_return_lse_for_decode = (
            self.dcp_world_size > 1 and self.can_return_lse_for_decode
        )
        return self

    def process_weights_after_loading(self, act_dtype: torch.dtype):
        pass

    def get_impl_variants(self) -> tuple["AttentionImplBase", ...]:
        return (self,)

    def prepare_for_batch(self, attn_metadata: T | None) -> None:
        """Prepare implementation-specific state for the current batch."""

prepare_for_batch(attn_metadata)

Prepare implementation-specific state for the current batch.

Source code in vllm/v1/attention/backend.py

def prepare_for_batch(self, attn_metadata: T | None) -> None:
    """Prepare implementation-specific state for the current batch."""

AttentionMetadataBuilder

Bases: ABC, Generic[M]

Methods:

• build –

  Central method that builds attention metadata.

• build_for_cudagraph_capture –

  Build attention metadata for CUDA graph capture. Uses build by default.

• build_for_drafting –

  Build attention metadata for draft model. Uses build by default.

• get_cudagraph_support –

  Get the cudagraph support level of this builder class.

• get_varlen_cudagraph_max_query_len –

  Get the largest per-request query length L of the variable-length

• update_block_table –

  Update the block table for the attention metadata.

• update_draft_decode_metadata –

  Update step-dependent draft decode metadata in place.

Source code in vllm/v1/attention/backend.py

class AttentionMetadataBuilder(ABC, Generic[M]):
    # Does this backend/builder support CUDA Graphs for attention (default: no).
    # Do not access directly. Call get_cudagraph_support() instead.
    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.NEVER
    # Does this backend/builder reorder the batch?
    # If not, set this to None. Otherwise set it to the query
    # length that will be pulled into the front of the batch.
    reorder_batch_threshold: int | None = None
    # Does this backend/builder support updating the block table in existing
    # metadata
    supports_update_block_table: bool = False
    # Whether the builder constructor requires the block-table width.
    requires_block_table_width: ClassVar[bool] = False
    # Whether all step-dependent draft decode metadata can be updated in place,
    # allowing one metadata build to be reused across autoregressive draft steps.
    supports_draft_decode_metadata_update: bool = False

    @abstractmethod
    def __init__(
        self,
        kv_cache_spec: "KVCacheSpec",
        layer_names: list[str],
        vllm_config: "VllmConfig",
        device: torch.device,
    ):
        self.kv_cache_spec = kv_cache_spec
        self.layer_names = layer_names
        self.vllm_config = vllm_config
        self.device = device
        self.kernel_block_size: int | None = None

    def set_kernel_block_size(self, kernel_block_size: int) -> None:
        self.kernel_block_size = kernel_block_size

    @classmethod
    def get_cudagraph_support(
        cls: type["AttentionMetadataBuilder"],
        vllm_config: "VllmConfig",
        kv_cache_spec: "KVCacheSpec",
    ) -> AttentionCGSupport:
        """Get the cudagraph support level of this builder class."""
        return cls._cudagraph_support

    @classmethod
    def get_varlen_cudagraph_max_query_len(
        cls: type["AttentionMetadataBuilder"],
        vllm_config: "VllmConfig",
        kv_cache_spec: "KVCacheSpec",
    ) -> int | None:
        """Get the largest per-request query length L of the variable-length
        decode batches a FULL cudagraph of this builder class can replay.

        The graph must replay any decode batch in which every real request has
        between 1 and L query tokens and every padding request has 0. Lengths
        come from the device query_start_loc; host metadata carries only the
        token count and an upper bound on the per-request length. Batches with
        a prefill never replay these graphs. Whether host metadata may
        understate device query lengths at all is a separate backend question;
        see supports_device_cpu_query_lens_mismatch().

        Returns:
            L, or None for builders reporting ALWAYS, which replay any batch,
            and for builders that cannot replay variable-length batches.

        """
        return None

    def _init_reorder_batch_threshold(
        self,
        reorder_batch_threshold: int | None = 1,
        supports_spec_as_decode: bool = False,
        supports_dcp_with_varlen: bool = False,
    ) -> None:
        self.reorder_batch_threshold = reorder_batch_threshold
        if self.reorder_batch_threshold is not None and supports_spec_as_decode:
            # If the backend supports spec-as-decode kernels, then we can set
            # the reorder_batch_threshold based on the number of speculative
            # tokens from the config.
            speculative_config = self.vllm_config.speculative_config
            if (
                speculative_config is not None
                and speculative_config.num_speculative_tokens is not None
            ):
                self.reorder_batch_threshold = max(
                    self.reorder_batch_threshold,
                    max_decode_query_len(self.vllm_config),
                )

        if (
            self.vllm_config.parallel_config.decode_context_parallel_size > 1
            and not supports_dcp_with_varlen
        ):
            self.reorder_batch_threshold = 1

    @abstractmethod
    def build(
        self,
        common_prefix_len: int,
        common_attn_metadata: CommonAttentionMetadata,
        fast_build: bool = False,
    ) -> M:
        """Central method that builds attention metadata.
        Some builders (MLA) require reorder_batch to be called prior to build.

        Args:
            common_prefix_len: The length of the common prefix of the batch.
            common_attn_metadata: The common attention metadata.
            fast_build: The meta-data will prioritize speed of building over
                then speed at execution. Can be used for spec-decode where the
                result of a build call may only be used for few layers/iters.

        """
        raise NotImplementedError

    def update_block_table(
        self,
        metadata: M,
        blk_table: torch.Tensor,
        slot_mapping: torch.Tensor,
    ) -> M:
        """Update the block table for the attention metadata.
        Faster when theres multiple kv-cache groups that create virtually the
        same metadata but just with different block tables.

        Only needs to be implemented if supports_update_block_table is True.
        """
        raise NotImplementedError

    def build_for_cudagraph_capture(
        self, common_attn_metadata: CommonAttentionMetadata
    ) -> M:
        """Build attention metadata for CUDA graph capture. Uses build by default.
        Subclasses that override this method should call self.build or
        super().build_for_cudagraph_capture.
        """
        return self.build(
            common_prefix_len=0, common_attn_metadata=common_attn_metadata
        )

    def build_for_drafting(
        self,
        common_attn_metadata: CommonAttentionMetadata,
        draft_index: int,
    ) -> M:
        """Build attention metadata for draft model. Uses build by default.

        Args:
            common_attn_metadata: The common attention metadata.
            draft_index: The index of the current draft operation.
                When speculating a chain of tokens, this index refers to the
                draft attempt for the i-th token.
                For tree-based attention, this index instead refers to the
                draft attempt for the i-th level in the tree of tokens.

        """
        return self.build(
            common_prefix_len=0,
            common_attn_metadata=common_attn_metadata,
            fast_build=True,
        )

    def update_draft_decode_metadata(self, metadata: M) -> None:
        """Update step-dependent draft decode metadata in place.

        The fused draft loop may call this method during full CUDA graph
        capture. CUDA graph replay does not run this Python method, so
        implementations must emit capture-safe operations and keep replayed
        tensor state in persistent storage.
        """
        raise NotImplementedError

    def use_cascade_attention(
        self,
        common_prefix_len: int,
        query_lens: np.ndarray,
        num_query_heads: int,
        num_kv_heads: int,
        use_alibi: bool,
        use_sliding_window: bool,
        use_local_attention: bool,
        num_sms: int,
        dcp_world_size: int,
    ) -> bool:
        return False

build(common_prefix_len, common_attn_metadata, fast_build=False) abstractmethod

Central method that builds attention metadata. Some builders (MLA) require reorder_batch to be called prior to build.

Parameters:

• common_prefix_len

  (int) –

  The length of the common prefix of the batch.

• common_attn_metadata

  (CommonAttentionMetadata) –

  The common attention metadata.

• fast_build

  (bool, default: False ) –

  The meta-data will prioritize speed of building over then speed at execution. Can be used for spec-decode where the result of a build call may only be used for few layers/iters.

Source code in vllm/v1/attention/backend.py

@abstractmethod
def build(
    self,
    common_prefix_len: int,
    common_attn_metadata: CommonAttentionMetadata,
    fast_build: bool = False,
) -> M:
    """Central method that builds attention metadata.
    Some builders (MLA) require reorder_batch to be called prior to build.

    Args:
        common_prefix_len: The length of the common prefix of the batch.
        common_attn_metadata: The common attention metadata.
        fast_build: The meta-data will prioritize speed of building over
            then speed at execution. Can be used for spec-decode where the
            result of a build call may only be used for few layers/iters.

    """
    raise NotImplementedError

build_for_cudagraph_capture(common_attn_metadata)

Build attention metadata for CUDA graph capture. Uses build by default. Subclasses that override this method should call self.build or super().build_for_cudagraph_capture.

Source code in vllm/v1/attention/backend.py

def build_for_cudagraph_capture(
    self, common_attn_metadata: CommonAttentionMetadata
) -> M:
    """Build attention metadata for CUDA graph capture. Uses build by default.
    Subclasses that override this method should call self.build or
    super().build_for_cudagraph_capture.
    """
    return self.build(
        common_prefix_len=0, common_attn_metadata=common_attn_metadata
    )

build_for_drafting(common_attn_metadata, draft_index)

Build attention metadata for draft model. Uses build by default.

Parameters:

• common_attn_metadata

  (CommonAttentionMetadata) –

  The common attention metadata.

• draft_index

  (int) –

  The index of the current draft operation. When speculating a chain of tokens, this index refers to the draft attempt for the i-th token. For tree-based attention, this index instead refers to the draft attempt for the i-th level in the tree of tokens.

Source code in vllm/v1/attention/backend.py

def build_for_drafting(
    self,
    common_attn_metadata: CommonAttentionMetadata,
    draft_index: int,
) -> M:
    """Build attention metadata for draft model. Uses build by default.

    Args:
        common_attn_metadata: The common attention metadata.
        draft_index: The index of the current draft operation.
            When speculating a chain of tokens, this index refers to the
            draft attempt for the i-th token.
            For tree-based attention, this index instead refers to the
            draft attempt for the i-th level in the tree of tokens.

    """
    return self.build(
        common_prefix_len=0,
        common_attn_metadata=common_attn_metadata,
        fast_build=True,
    )

get_cudagraph_support(vllm_config, kv_cache_spec) classmethod

Get the cudagraph support level of this builder class.

Source code in vllm/v1/attention/backend.py

@classmethod
def get_cudagraph_support(
    cls: type["AttentionMetadataBuilder"],
    vllm_config: "VllmConfig",
    kv_cache_spec: "KVCacheSpec",
) -> AttentionCGSupport:
    """Get the cudagraph support level of this builder class."""
    return cls._cudagraph_support

get_varlen_cudagraph_max_query_len(vllm_config, kv_cache_spec) classmethod

Get the largest per-request query length L of the variable-length decode batches a FULL cudagraph of this builder class can replay.

The graph must replay any decode batch in which every real request has between 1 and L query tokens and every padding request has 0. Lengths come from the device query_start_loc; host metadata carries only the token count and an upper bound on the per-request length. Batches with a prefill never replay these graphs. Whether host metadata may understate device query lengths at all is a separate backend question; see supports_device_cpu_query_lens_mismatch().

Returns:

• int | None –

  L, or None for builders reporting ALWAYS, which replay any batch,

• int | None –

  and for builders that cannot replay variable-length batches.

Source code in vllm/v1/attention/backend.py

@classmethod
def get_varlen_cudagraph_max_query_len(
    cls: type["AttentionMetadataBuilder"],
    vllm_config: "VllmConfig",
    kv_cache_spec: "KVCacheSpec",
) -> int | None:
    """Get the largest per-request query length L of the variable-length
    decode batches a FULL cudagraph of this builder class can replay.

    The graph must replay any decode batch in which every real request has
    between 1 and L query tokens and every padding request has 0. Lengths
    come from the device query_start_loc; host metadata carries only the
    token count and an upper bound on the per-request length. Batches with
    a prefill never replay these graphs. Whether host metadata may
    understate device query lengths at all is a separate backend question;
    see supports_device_cpu_query_lens_mismatch().

    Returns:
        L, or None for builders reporting ALWAYS, which replay any batch,
        and for builders that cannot replay variable-length batches.

    """
    return None

update_block_table(metadata, blk_table, slot_mapping)

Update the block table for the attention metadata. Faster when theres multiple kv-cache groups that create virtually the same metadata but just with different block tables.

Only needs to be implemented if supports_update_block_table is True.

Source code in vllm/v1/attention/backend.py

def update_block_table(
    self,
    metadata: M,
    blk_table: torch.Tensor,
    slot_mapping: torch.Tensor,
) -> M:
    """Update the block table for the attention metadata.
    Faster when theres multiple kv-cache groups that create virtually the
    same metadata but just with different block tables.

    Only needs to be implemented if supports_update_block_table is True.
    """
    raise NotImplementedError

update_draft_decode_metadata(metadata)

Update step-dependent draft decode metadata in place.

The fused draft loop may call this method during full CUDA graph capture. CUDA graph replay does not run this Python method, so implementations must emit capture-safe operations and keep replayed tensor state in persistent storage.

Source code in vllm/v1/attention/backend.py

def update_draft_decode_metadata(self, metadata: M) -> None:
    """Update step-dependent draft decode metadata in place.

    The fused draft loop may call this method during full CUDA graph
    capture. CUDA graph replay does not run this Python method, so
    implementations must emit capture-safe operations and keep replayed
    tensor state in persistent storage.
    """
    raise NotImplementedError

AttentionType

Bases: str, Enum

Attention type. Use string to be compatible with torch.compile.

Attributes:

• DECODER –

  Decoder attention between previous layer Q/K/V.

• ENCODER –

  Encoder attention between previous layer Q/K/V for encoder-decoder.

• ENCODER_DECODER –

  Attention between dec. Q and enc. K/V for encoder-decoder.

• ENCODER_ONLY –

  Encoder attention between previous layer Q/K/V.

Source code in vllm/v1/attention/backend.py

class AttentionType(str, Enum):
    """Attention type.
    Use string to be compatible with `torch.compile`.
    """

    DECODER = "decoder"
    """Decoder attention between previous layer Q/K/V."""
    ENCODER = "encoder"
    """Encoder attention between previous layer Q/K/V for encoder-decoder."""
    ENCODER_ONLY = "encoder_only"
    """Encoder attention between previous layer Q/K/V."""
    ENCODER_DECODER = "encoder_decoder"
    """Attention between dec. Q and enc. K/V for encoder-decoder."""

DECODER = 'decoder' class-attribute instance-attribute

Decoder attention between previous layer Q/K/V.

ENCODER = 'encoder' class-attribute instance-attribute

Encoder attention between previous layer Q/K/V for encoder-decoder.

ENCODER_DECODER = 'encoder_decoder' class-attribute instance-attribute

Attention between dec. Q and enc. K/V for encoder-decoder.

ENCODER_ONLY = 'encoder_only' class-attribute instance-attribute

Encoder attention between previous layer Q/K/V.

CommonAttentionMetadata dataclass

Per-batch attention metadata, shared across layers and backends. AttentionMetadataBuilder instances use it to construct per-layer metadata.

For many of the tensors we keep both GPU and CPU versions.

Methods:

• compute_num_computed_tokens –

  Compute num_computed_tokens on device (seq_lens - query_lens).

• naive_query_lens –

  Naive because it assumes that query ends where the next query starts.

• token_to_req_indices –

  Build or reuse the per-token request index mapping.

Attributes:

• dcp_local_seq_lens (Tensor | None) –

  Sequence lengths of the local rank in decode context parallelism world

• dcp_local_seq_lens_cpu_upper_bound (Tensor | None) –

  (batch_size,) CPU upper bound on dcp_local_seq_lens. Under PCP+DCP it

• is_prefilling (Tensor | None) –

  (batch_size,) bool tensor: True if request is still in prefill phase

• max_query_len (int) –

  Longest query in batch

• max_seq_len (int) –

  Longest context length (may be an upper bound)

• mm_req_doc_ranges (dict[int, list[tuple[int, int]]] | None) –

  PrefixLM bidirectional ranges for multimodal tokens. Maps

• num_actual_tokens (int) –

  Total number of tokens in batch

• num_reqs (int) –

  Number of requests

• positions (Tensor | None) –

  (num_actual_tokens,) token positions. Optional; set when the caller

• query_start_loc_cpu (Tensor) –

  (batch_size + 1,), the start location of each request in query Tensor

• replayssm_decode_base_cpu (Tensor | None) –

  (batch_size,) CPU ring origin for Mamba2 ReplaySSM decode: num_computed

• req_idx (ndarray | None) –

  (batch_size,) index of each row's request in the runner's request

• rswa_prefix_lens (Tensor | None) –

  (batch_size,) per-request prefix length (prompt/image token count) for

• seq_lens (Tensor) –

  (batch_size,), the number of computed tokens for each request

• seq_lens_cpu_upper_bound (Tensor | None) –

  (batch_size,) CPU upper bound on seq_lens. Precise for prefill rows

Source code in vllm/v1/attention/backend.py

@dataclass
class CommonAttentionMetadata:
    """Per-batch attention metadata, shared across layers and backends.
    AttentionMetadataBuilder instances use it to construct per-layer metadata.

    For many of the tensors we keep both GPU and CPU versions.
    """

    query_start_loc: torch.Tensor
    query_start_loc_cpu: torch.Tensor
    """(batch_size + 1,), the start location of each request in query Tensor"""

    seq_lens: torch.Tensor
    """(batch_size,), the number of computed tokens for each request"""

    num_reqs: int
    """Number of requests"""
    # TODO(lucas): rename to num_tokens since it may be padded and this is misleading
    num_actual_tokens: int
    """Total number of tokens in batch"""
    max_query_len: int
    """Longest query in batch"""
    max_seq_len: int
    """Longest context length (may be an upper bound)"""

    block_table_tensor: torch.Tensor
    slot_mapping: torch.Tensor

    causal: bool | torch.Tensor = True

    # Needed by FastPrefillAttentionBuilder
    logits_indices_padded: torch.Tensor | None = None
    num_logits_indices: int | None = None
    max_logits_per_req: int | None = None

    # Needed by CrossAttentionBuilder
    encoder_seq_lens: torch.Tensor | None = None
    encoder_seq_lens_cpu: np.ndarray | None = None

    dcp_local_seq_lens: torch.Tensor | None = None
    """Sequence lengths of the local rank in decode context parallelism world"""

    dcp_local_seq_lens_cpu_upper_bound: torch.Tensor | None = None
    """(batch_size,) CPU upper bound on dcp_local_seq_lens. Under PCP+DCP it
    holds, on every row, the largest shard any DCP rank has of the row's
    whole request, identical on every PCP rank, so the sparse backends pad
    their KV gather to it."""

    positions: torch.Tensor | None = None
    """(num_actual_tokens,) token positions.  Optional; set when the caller
    has positions available so that builders can pre-compute position-dependent
    sparse metadata for DeepSeek V4 C128A layers."""

    is_prefilling: torch.Tensor | None = None
    """(batch_size,) bool tensor: True if request is still in prefill phase
    (num_computed_tokens < num_prompt_tokens). Used by some backends to
    distinguish actual decodes from short extends."""

    seq_lens_cpu_upper_bound: torch.Tensor | None = None
    """(batch_size,) CPU upper bound on seq_lens. Precise for prefill rows
    and for all rows outside async spec decode; optimistic for async-spec
    decode rows (assumes every draft was accepted). Not safe for kernels
    that need exact per-row context lengths on decode rows."""

    req_idx: np.ndarray | None = None
    """(batch_size,) index of each row's request in the runner's request
    table. Rows of one request are adjacent, so equal neighbours are PCP
    chunks sharing one KV context."""

    mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None
    """PrefixLM bidirectional ranges for multimodal tokens. Maps
    request index to list of (start, end) token position ranges
    where bidirectional attention should apply. None for text-only
    batches or non-PrefixLM models. A request's ranges must not overlap."""

    rswa_prefix_lens: torch.Tensor | None = None
    """(batch_size,) per-request prefix length (prompt/image token count) for
    Reference Sliding Window Attention (R-SWA). Tokens with logical index below
    this stay globally visible; later (generated) tokens additionally see a
    fixed sliding window. None disables R-SWA. The attention backend copies this
    into its own persistent buffer and reads ``rswa_window`` from model config."""

    replayssm_decode_base_cpu: torch.Tensor | None = None
    """(batch_size,) CPU ring origin for Mamba2 ReplaySSM decode: num_computed
    at the current decode run's last full-state write. write_pos counts from
    here, so a preemption-resumed request re-anchors past the prompt boundary."""

    _num_computed_tokens_cache: torch.Tensor | None = None
    _token_to_req_indices_cache: torch.Tensor | None = None

    def batch_size(self) -> int:
        return self.seq_lens.shape[0]

    def naive_query_lens(self) -> torch.Tensor:
        """Naive because it assumes that query ends where the next query starts."""
        return self.query_start_loc[1:] - self.query_start_loc[:-1]

    def replace(self, **kwargs) -> "CommonAttentionMetadata":
        return replace(self, **kwargs)

    def compute_num_computed_tokens(self) -> torch.Tensor:
        """Compute num_computed_tokens on device (seq_lens - query_lens)."""
        if self._num_computed_tokens_cache is None:
            query_lens = self.query_start_loc[1:] - self.query_start_loc[:-1]
            self._num_computed_tokens_cache = self.seq_lens - query_lens
        return self._num_computed_tokens_cache

    def token_to_req_indices(self, buffer: torch.Tensor) -> torch.Tensor:
        """Build or reuse the per-token request index mapping."""
        num_tokens = self.num_actual_tokens
        if self._token_to_req_indices_cache is not None:
            assert self._token_to_req_indices_cache.device == buffer.device
            assert self._token_to_req_indices_cache.dtype == torch.int32
            assert self._token_to_req_indices_cache.shape[0] >= num_tokens
            return self._token_to_req_indices_cache[:num_tokens]

        # Built from the device query_start_loc: adaptive verification decides the
        # per-request draft split on device, so the CPU copy carries the right total
        # but not the right per-request boundaries. Padding requests have a query
        # length of zero and are skipped by the device boundary search.
        num_mapped_tokens = int(self.query_start_loc_cpu[-1])
        from vllm.v1.attention.ops.metadata import _token_request_mapping_kernel

        num_output_tokens = max(num_mapped_tokens, num_tokens)
        assert buffer.shape[0] >= num_output_tokens
        _token_request_mapping_kernel[((num_output_tokens + 255) // 256,)](
            self.query_start_loc,
            buffer,
            self.query_start_loc.shape[0] - 1,
            num_mapped_tokens,
            num_output_tokens,
            num_warps=4,
        )
        self._token_to_req_indices_cache = buffer[:num_output_tokens]
        return self._token_to_req_indices_cache[:num_tokens]

    # TODO(lucas): remove once we have FULL-CG spec-decode support
    def unpadded(
        self, num_actual_tokens: int, num_actual_reqs: int
    ) -> "CommonAttentionMetadata":
        maybe_slice_reqs = lambda x: x[:num_actual_reqs] if x is not None else None
        return CommonAttentionMetadata(
            query_start_loc=self.query_start_loc[: num_actual_reqs + 1],
            query_start_loc_cpu=self.query_start_loc_cpu[: num_actual_reqs + 1],
            seq_lens=self.seq_lens[:num_actual_reqs],
            num_reqs=num_actual_reqs,
            num_actual_tokens=num_actual_tokens,
            max_query_len=self.max_query_len,
            max_seq_len=self.max_seq_len,
            block_table_tensor=self.block_table_tensor[:num_actual_reqs],
            slot_mapping=self.slot_mapping[:num_actual_tokens],
            causal=self.causal[:num_actual_reqs]
            if isinstance(self.causal, torch.Tensor)
            else self.causal,
            logits_indices_padded=self.logits_indices_padded,
            num_logits_indices=self.num_logits_indices,
            max_logits_per_req=self.max_logits_per_req,
            encoder_seq_lens=maybe_slice_reqs(self.encoder_seq_lens),
            encoder_seq_lens_cpu=maybe_slice_reqs(self.encoder_seq_lens_cpu),
            dcp_local_seq_lens=maybe_slice_reqs(self.dcp_local_seq_lens),
            dcp_local_seq_lens_cpu_upper_bound=maybe_slice_reqs(
                self.dcp_local_seq_lens_cpu_upper_bound
            ),
            seq_lens_cpu_upper_bound=maybe_slice_reqs(self.seq_lens_cpu_upper_bound),
            is_prefilling=maybe_slice_reqs(self.is_prefilling),
            req_idx=maybe_slice_reqs(self.req_idx),
            rswa_prefix_lens=maybe_slice_reqs(self.rswa_prefix_lens),
            replayssm_decode_base_cpu=maybe_slice_reqs(self.replayssm_decode_base_cpu),
        )

dcp_local_seq_lens = None class-attribute instance-attribute

Sequence lengths of the local rank in decode context parallelism world

dcp_local_seq_lens_cpu_upper_bound = None class-attribute instance-attribute

(batch_size,) CPU upper bound on dcp_local_seq_lens. Under PCP+DCP it holds, on every row, the largest shard any DCP rank has of the row's whole request, identical on every PCP rank, so the sparse backends pad their KV gather to it.

is_prefilling = None class-attribute instance-attribute

(batch_size,) bool tensor: True if request is still in prefill phase (num_computed_tokens < num_prompt_tokens). Used by some backends to distinguish actual decodes from short extends.

max_query_len instance-attribute

Longest query in batch

max_seq_len instance-attribute

Longest context length (may be an upper bound)

mm_req_doc_ranges = None class-attribute instance-attribute

PrefixLM bidirectional ranges for multimodal tokens. Maps request index to list of (start, end) token position ranges where bidirectional attention should apply. None for text-only batches or non-PrefixLM models. A request's ranges must not overlap.

num_actual_tokens instance-attribute

Total number of tokens in batch

num_reqs instance-attribute

Number of requests

positions = None class-attribute instance-attribute

(num_actual_tokens,) token positions. Optional; set when the caller has positions available so that builders can pre-compute position-dependent sparse metadata for DeepSeek V4 C128A layers.

query_start_loc_cpu instance-attribute

(batch_size + 1,), the start location of each request in query Tensor

replayssm_decode_base_cpu = None class-attribute instance-attribute

(batch_size,) CPU ring origin for Mamba2 ReplaySSM decode: num_computed at the current decode run's last full-state write. write_pos counts from here, so a preemption-resumed request re-anchors past the prompt boundary.

req_idx = None class-attribute instance-attribute

(batch_size,) index of each row's request in the runner's request table. Rows of one request are adjacent, so equal neighbours are PCP chunks sharing one KV context.

rswa_prefix_lens = None class-attribute instance-attribute

(batch_size,) per-request prefix length (prompt/image token count) for Reference Sliding Window Attention (R-SWA). Tokens with logical index below this stay globally visible; later (generated) tokens additionally see a fixed sliding window. None disables R-SWA. The attention backend copies this into its own persistent buffer and reads rswa_window from model config.

seq_lens instance-attribute

(batch_size,), the number of computed tokens for each request

seq_lens_cpu_upper_bound = None class-attribute instance-attribute

(batch_size,) CPU upper bound on seq_lens. Precise for prefill rows and for all rows outside async spec decode; optimistic for async-spec decode rows (assumes every draft was accepted). Not safe for kernels that need exact per-row context lengths on decode rows.

compute_num_computed_tokens()

Compute num_computed_tokens on device (seq_lens - query_lens).

Source code in vllm/v1/attention/backend.py

def compute_num_computed_tokens(self) -> torch.Tensor:
    """Compute num_computed_tokens on device (seq_lens - query_lens)."""
    if self._num_computed_tokens_cache is None:
        query_lens = self.query_start_loc[1:] - self.query_start_loc[:-1]
        self._num_computed_tokens_cache = self.seq_lens - query_lens
    return self._num_computed_tokens_cache

naive_query_lens()

Naive because it assumes that query ends where the next query starts.

Source code in vllm/v1/attention/backend.py

def naive_query_lens(self) -> torch.Tensor:
    """Naive because it assumes that query ends where the next query starts."""
    return self.query_start_loc[1:] - self.query_start_loc[:-1]

token_to_req_indices(buffer)

Build or reuse the per-token request index mapping.

Source code in vllm/v1/attention/backend.py

def token_to_req_indices(self, buffer: torch.Tensor) -> torch.Tensor:
    """Build or reuse the per-token request index mapping."""
    num_tokens = self.num_actual_tokens
    if self._token_to_req_indices_cache is not None:
        assert self._token_to_req_indices_cache.device == buffer.device
        assert self._token_to_req_indices_cache.dtype == torch.int32
        assert self._token_to_req_indices_cache.shape[0] >= num_tokens
        return self._token_to_req_indices_cache[:num_tokens]

    # Built from the device query_start_loc: adaptive verification decides the
    # per-request draft split on device, so the CPU copy carries the right total
    # but not the right per-request boundaries. Padding requests have a query
    # length of zero and are skipped by the device boundary search.
    num_mapped_tokens = int(self.query_start_loc_cpu[-1])
    from vllm.v1.attention.ops.metadata import _token_request_mapping_kernel

    num_output_tokens = max(num_mapped_tokens, num_tokens)
    assert buffer.shape[0] >= num_output_tokens
    _token_request_mapping_kernel[((num_output_tokens + 255) // 256,)](
        self.query_start_loc,
        buffer,
        self.query_start_loc.shape[0] - 1,
        num_mapped_tokens,
        num_output_tokens,
        num_warps=4,
    )
    self._token_to_req_indices_cache = buffer[:num_output_tokens]
    return self._token_to_req_indices_cache[:num_tokens]

MLAAttentionImpl

Bases: AttentionImplBase[T], Generic[T]

MLA attention implementation with forward_mqa and forward_mha methods.

Methods:

• forward_mha –

  MHA-style prefill forward pass.

• forward_mqa –

  MQA-style decode forward pass.

• fused_output_quant_supported –

  Does this attention implementation support fused output quantization.

Source code in vllm/v1/attention/backend.py

class MLAAttentionImpl(AttentionImplBase[T], Generic[T]):
    """MLA attention implementation with forward_mqa and forward_mha methods."""

    hisparse_cache: "HiSparseCacheHandle | None" = None
    supports_pcp: bool = True

    @abstractmethod
    def __init__(
        self,
        num_heads: int,
        head_size: int,
        scale: float,
        num_kv_heads: int,
        alibi_slopes: list[float] | None,
        sliding_window: int | None,
        kv_cache_dtype: str,
        logits_soft_cap: float | None,
        attn_type: str,
        kv_sharing_target_layer_name: str | None,
        # MLA Specific Arguments
        q_lora_rank: int | None,
        kv_lora_rank: int,
        qk_nope_head_dim: int,
        qk_rope_head_dim: int,
        qk_head_dim: int,
        v_head_dim: int,
        kv_b_proj: "ColumnParallelLinear",
        indexer: object | None = None,
        q_pad_num_heads: int | None = None,
    ) -> None:
        raise NotImplementedError

    def forward_mha(
        self,
        q: torch.Tensor,
        kv_c_normed: torch.Tensor,
        k_pe: torch.Tensor,
        kv_c_and_k_pe_cache: torch.Tensor,
        attn_metadata: T,
        k_scale: torch.Tensor,
        output: torch.Tensor,
        output_scale: torch.Tensor | None = None,
    ) -> None:
        """MHA-style prefill forward pass."""
        raise NotImplementedError

    @abstractmethod
    def forward_mqa(
        self,
        q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
        kv_c_and_k_pe_cache: torch.Tensor,
        attn_metadata: T,
        layer: AttentionLayer,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """MQA-style decode forward pass."""
        raise NotImplementedError

    def fused_output_quant_supported(self, quant_key: "QuantKey"):
        """Does this attention implementation support fused output quantization.
        Since MLA quantization is done manually in forward_impl (common code),
        all MLA backends support it by default.
        """
        return quant_key in (
            kFp8StaticTensorSym,
            kNvfp4Dynamic,
            kFp8Dynamic128Sym,
            kFp8Dynamic64Sym,
        )

    def do_kv_cache_update(
        self,
        kv_c_normed: torch.Tensor,
        k_pe: torch.Tensor,
        kv_cache: torch.Tensor,
        slot_mapping: torch.Tensor,
        kv_cache_dtype: str,
        k_scale: torch.Tensor,
    ) -> None:
        if kv_cache.numel() == 0:
            return
        from vllm import _custom_ops as ops

        ops.concat_and_cache_mla(
            kv_c_normed,
            k_pe.squeeze(1),
            kv_cache,
            slot_mapping.flatten(),
            kv_cache_dtype=kv_cache_dtype,
            scale=k_scale,
        )

forward_mha(q, kv_c_normed, k_pe, kv_c_and_k_pe_cache, attn_metadata, k_scale, output, output_scale=None)

MHA-style prefill forward pass.

Source code in vllm/v1/attention/backend.py

def forward_mha(
    self,
    q: torch.Tensor,
    kv_c_normed: torch.Tensor,
    k_pe: torch.Tensor,
    kv_c_and_k_pe_cache: torch.Tensor,
    attn_metadata: T,
    k_scale: torch.Tensor,
    output: torch.Tensor,
    output_scale: torch.Tensor | None = None,
) -> None:
    """MHA-style prefill forward pass."""
    raise NotImplementedError

forward_mqa(q, kv_c_and_k_pe_cache, attn_metadata, layer) abstractmethod

MQA-style decode forward pass.

Source code in vllm/v1/attention/backend.py

@abstractmethod
def forward_mqa(
    self,
    q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
    kv_c_and_k_pe_cache: torch.Tensor,
    attn_metadata: T,
    layer: AttentionLayer,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    """MQA-style decode forward pass."""
    raise NotImplementedError

fused_output_quant_supported(quant_key)

Does this attention implementation support fused output quantization. Since MLA quantization is done manually in forward_impl (common code), all MLA backends support it by default.

Source code in vllm/v1/attention/backend.py

def fused_output_quant_supported(self, quant_key: "QuantKey"):
    """Does this attention implementation support fused output quantization.
    Since MLA quantization is done manually in forward_impl (common code),
    all MLA backends support it by default.
    """
    return quant_key in (
        kFp8StaticTensorSym,
        kNvfp4Dynamic,
        kFp8Dynamic128Sym,
        kFp8Dynamic64Sym,
    )

max_decode_query_len(vllm_config)

Widest request a spec-as-decode builder treats as a decode.

On model runner V2 this is a verification request, 1 + num_speculative_tokens: no V2 speculator builds a wider query. The reorder threshold and the code that mirrors the builders' decode/prefill split all read this, so they cannot drift apart.

Source code in vllm/v1/attention/backend.py

def max_decode_query_len(vllm_config: "VllmConfig") -> int:
    """Widest request a spec-as-decode builder treats as a decode.

    On model runner V2 this is a verification request, 1 +
    num_speculative_tokens: no V2 speculator builds a wider query. The reorder
    threshold and the code that mirrors the builders' decode/prefill split all
    read this, so they cannot drift apart.
    """
    speculative_config = getattr(vllm_config, "speculative_config", None)
    if speculative_config is None or speculative_config.num_speculative_tokens is None:
        return 1
    num_speculative_tokens = speculative_config.num_speculative_tokens
    max_query_len = 1 + num_speculative_tokens
    if speculative_config.parallel_drafting and not vllm_config.use_v2_model_runner:
        # Model runner V1 only; remove with it. Its parallel drafter appends up
        # to num_speculative_tokens mask slots to each verified request.
        max_query_len += num_speculative_tokens
    return max_query_len

subclass_attention_backend(name_prefix, attention_backend_cls, builder_cls)

Return a new subclass where get_builder_cls returns builder_cls.

Source code in vllm/v1/attention/backend.py

def subclass_attention_backend(
    name_prefix: str,
    attention_backend_cls: type[AttentionBackend],
    builder_cls: type[AttentionMetadataBuilder[M]],
) -> type[AttentionBackend]:
    """Return a new subclass where `get_builder_cls` returns `builder_cls`."""
    name: str = name_prefix + attention_backend_cls.__name__  # type: ignore

    return type(
        name, (attention_backend_cls,), {"get_builder_cls": lambda: builder_cls}
    )