vllm.v1.attention.backends.rocm_aiter_fa

Attention layer with AiterFlashAttention.

Classes:

• AiterFlashAttentionBackend –
• AiterFlashAttentionImpl –
• AiterFlashAttentionMetadataBuilder –

AiterFlashAttentionBackend

Bases: AttentionBackend

Methods:

• customize_spec –

  Validate the block size the ROCm gather kernels require, and under

• supports_attn_type –

  ENCODER_DECODER is not supported because the prefill path uses

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

class AiterFlashAttentionBackend(AttentionBackend):
    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]

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

    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
        "auto",
        "float16",
        "bfloat16",
        "fp8",
        "fp8_e4m3",
        "fp8_e5m2",
    ]

    @classmethod
    def supports_attn_type(cls, attn_type: str) -> bool:
        """ENCODER_DECODER is not supported because the prefill path uses
        flash_attn_varlen_func with cu_seqlens_k set to decoder
        query_start_loc (not encoder seq lens) and causal=True, both of
        which are incorrect for cross-attention layers.
        """
        return attn_type in (AttentionType.DECODER,)

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

    @classmethod
    def get_supported_head_sizes(cls) -> list[int]:
        return [64, 128, 256]

    forward_includes_kv_cache_update: bool = False

    @staticmethod
    def get_name() -> str:
        return "FLASH_ATTN"

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

    @staticmethod
    def get_impl_cls() -> type["AiterFlashAttentionImpl"]:
        return AiterFlashAttentionImpl

    @staticmethod
    def get_builder_cls() -> type["AiterFlashAttentionMetadataBuilder"]:
        return AiterFlashAttentionMetadataBuilder

    @classmethod
    def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec:
        """Validate the block size the ROCm gather kernels require, and under
        the shuffle layout publish K and V as two head slots.

        The shuffle read/write kernels rearrange a whole ``head_size x
        block_size`` tile per side, so each side needs its own contiguous run
        in the page. Packed into the content dim they interleave per token
        instead, which no layout can undo.
        """
        # block_size == 1 is the per-token page-size probe (see
        # Platform.get_page_size_bytes); real blocks must be gatherable in
        # 16-token units by the ROCm kernel.
        if spec.block_size != 1 and spec.block_size % 16 != 0:
            raise ValueError("Block size must be a multiple of 16.")
        if not rocm_aiter_ops.is_shuffle_kv_cache_enabled():
            return spec
        if spec.state_content_bytes is not None:
            return spec
        assert spec.head_size == spec.head_size_v, (
            "Separate K/V head slots require symmetric K/V head sizes."
        )
        return replace(
            spec,
            num_head_slots=2,
            state_content_bytes=spec.num_kv_heads
            * spec.head_size
            * get_dtype_size(spec.dtype),
        )

    @classmethod
    def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
        if rocm_aiter_ops.is_shuffle_kv_cache_enabled():
            # pa_fwd_asm strides between pages by a whole dense page, so the two
            # head slots customize_spec publishes have to be separate planes
            # spanning every block: H outermost. Under LBHNC the slots alternate
            # block by block, doubling the stride the kernel assumes. LBHNC
            # stays listed for models whose mixed HNC shapes need a
            # block-compact layout; those read K/V by stride instead.
            return (KVCacheLayout.LHBNC, KVCacheLayout.LBHNC)
        # K and V come out of the content dim as transposed views rather than
        # copies, so the head dim may sit on either side of the block dim, but
        # the layer must stay outermost.
        return (KVCacheLayout.LBHNC, KVCacheLayout.LHBNC)

    @classmethod
    def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
        from vllm.platforms.rocm import get_cdna_version

        # DeviceCapability is currently created using torch.cuda.get_device_capability()
        # which is known to be buggy on rocm systems. on CDNA uses amd-smi which is
        # more reliable.
        return get_cdna_version() > 2

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

customize_spec(spec) classmethod

Validate the block size the ROCm gather kernels require, and under the shuffle layout publish K and V as two head slots.

The shuffle read/write kernels rearrange a whole head_size x block_size tile per side, so each side needs its own contiguous run in the page. Packed into the content dim they interleave per token instead, which no layout can undo.

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

@classmethod
def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec:
    """Validate the block size the ROCm gather kernels require, and under
    the shuffle layout publish K and V as two head slots.

    The shuffle read/write kernels rearrange a whole ``head_size x
    block_size`` tile per side, so each side needs its own contiguous run
    in the page. Packed into the content dim they interleave per token
    instead, which no layout can undo.
    """
    # block_size == 1 is the per-token page-size probe (see
    # Platform.get_page_size_bytes); real blocks must be gatherable in
    # 16-token units by the ROCm kernel.
    if spec.block_size != 1 and spec.block_size % 16 != 0:
        raise ValueError("Block size must be a multiple of 16.")
    if not rocm_aiter_ops.is_shuffle_kv_cache_enabled():
        return spec
    if spec.state_content_bytes is not None:
        return spec
    assert spec.head_size == spec.head_size_v, (
        "Separate K/V head slots require symmetric K/V head sizes."
    )
    return replace(
        spec,
        num_head_slots=2,
        state_content_bytes=spec.num_kv_heads
        * spec.head_size
        * get_dtype_size(spec.dtype),
    )

supports_attn_type(attn_type) classmethod

ENCODER_DECODER is not supported because the prefill path uses flash_attn_varlen_func with cu_seqlens_k set to decoder query_start_loc (not encoder seq lens) and causal=True, both of which are incorrect for cross-attention layers.

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

@classmethod
def supports_attn_type(cls, attn_type: str) -> bool:
    """ENCODER_DECODER is not supported because the prefill path uses
    flash_attn_varlen_func with cu_seqlens_k set to decoder
    query_start_loc (not encoder seq lens) and causal=True, both of
    which are incorrect for cross-attention layers.
    """
    return attn_type in (AttentionType.DECODER,)

AiterFlashAttentionImpl

Bases: AttentionImpl

Methods:

• forward –

  Forward pass with AiterFlashAttention.

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

class AiterFlashAttentionImpl(AttentionImpl):
    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 = None,
        attn_type: AttentionType = AttentionType.DECODER,
        kv_sharing_target_layer_name: int | None = None,
        sinks: torch.Tensor | None = None,
    ) -> None:
        self.num_heads = num_heads
        self.head_size = head_size
        self.scale = float(scale)
        self.num_kv_heads = num_kv_heads
        if alibi_slopes is not None:
            alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)
        self.alibi_slopes = alibi_slopes
        if sliding_window is None:
            self.sliding_window = (-1, -1)
        else:
            self.sliding_window = (sliding_window - 1, 0)
        self.kv_cache_dtype = kv_cache_dtype
        if logits_soft_cap is None:
            # In flash-attn, setting logits_soft_cap as 0 means no soft cap.
            logits_soft_cap = 0.0
        self.logits_soft_cap = logits_soft_cap
        self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
        self.sinks = sinks

        assert self.num_heads % self.num_kv_heads == 0
        self.num_queries_per_kv = self.num_heads // self.num_kv_heads

        if attn_type != AttentionType.DECODER:
            raise NotImplementedError(
                "Only decoder self-attention is supported for "
                "AiterFlashAttentionImpl. ENCODER_DECODER is not supported "
                "because the prefill path uses cu_seqlens_k set to decoder "
                "query_start_loc with causal=True, which is incorrect for "
                "cross-attention."
            )

    def _get_kv_cache_descales(
        self,
        layer: AttentionLayer,
        num_decodes: int,
    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
        # AITER treats non-null descales as an instruction to dequantize.
        if not is_quantized_kv_cache(self.kv_cache_dtype):
            return None, None
        descale_shape = (num_decodes, self.num_kv_heads)
        return (
            layer._k_scale.expand(descale_shape),
            layer._v_scale.expand(descale_shape),
        )

    def extend_for_sliding_window(
        self,
        attn_metadata: AiterFlashAttentionMetadata,
        query: torch.Tensor,
        key_cache,
        value_cache,
        output: torch.Tensor,
        cu_seqlens_q: torch.Tensor,
        max_seqlen_q: int,
        block_table: torch.Tensor,
        k_scale: float,
        v_scale: float,
    ):
        assert attn_metadata.extend_metadata is not None
        assert attn_metadata.extend_metadata.chunk_context_metadata is not None
        chunked_metadata = attn_metadata.extend_metadata.chunk_context_metadata
        swa_metadata = chunked_metadata.swa_metadata
        assert swa_metadata is not None
        swa_cu_seqlens = swa_metadata.swa_cu_seqlens
        swa_seq_starts = swa_metadata.swa_seq_starts
        swa_token_to_batch = swa_metadata.swa_token_to_batch
        swa_max_seqlens = swa_metadata.swa_max_seqlens
        swa_total_tokens = swa_metadata.swa_total_tokens
        key_fetched, value_fetched = (
            swa_metadata.swa_workspace[0],
            swa_metadata.swa_workspace[1],
        )
        cp_mha_gather_cache(
            key_cache=key_cache,
            value_cache=value_cache,
            key=key_fetched,
            value=value_fetched,
            block_tables=block_table,
            k_scales=k_scale,
            v_scales=v_scale,
            cu_seqlens_kv=swa_cu_seqlens,
            token_to_batch=swa_token_to_batch,
            seq_starts=swa_seq_starts,
            dequant=is_quantized_kv_cache(self.kv_cache_dtype),
            kv_cache_layout="NHD",
            total_tokens=swa_total_tokens,
        )

        rocm_aiter_ops.flash_attn_varlen_func(
            q=query,
            k=key_fetched,
            v=value_fetched,
            cu_seqlens_q=cu_seqlens_q,
            cu_seqlens_k=swa_cu_seqlens,
            max_seqlen_q=max_seqlen_q,
            max_seqlen_k=swa_max_seqlens,
            min_seqlen_q=1,
            dropout_p=0.0,
            softmax_scale=self.scale,
            causal=True,
            window_size=self.sliding_window,
            alibi_slopes=self.alibi_slopes,
            return_lse=False,
            out=output,
            sink_ptr=self.sinks,
        )

    def extend_forward(
        self,
        attn_metadata: AiterFlashAttentionMetadata,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        key_cache: torch.Tensor,
        value_cache: torch.Tensor,
        output: torch.Tensor,
        cu_seqlens_q: torch.Tensor,
        max_seqlen_q: int,
        min_seqlen_q: int,
        max_seqlen_k: int,
        block_table: torch.Tensor,
        slot_mapping: torch.Tensor,
        k_scale: torch.Tensor,
        v_scale: torch.Tensor,
    ):
        if self.sliding_window[0] != -1:
            self.extend_for_sliding_window(
                attn_metadata,
                query,
                key_cache,
                value_cache,
                output,
                cu_seqlens_q,
                max_seqlen_q,
                block_table,
                k_scale,
                v_scale,
            )
            return
        out, lse = rocm_aiter_ops.flash_attn_varlen_func(
            q=query,
            k=key,
            v=value,
            cu_seqlens_q=cu_seqlens_q,
            cu_seqlens_k=cu_seqlens_q,
            max_seqlen_q=max_seqlen_q,
            max_seqlen_k=max_seqlen_q,
            min_seqlen_q=min_seqlen_q,
            dropout_p=0.0,
            softmax_scale=self.scale,
            causal=True,
            window_size=self.sliding_window,
            alibi_slopes=self.alibi_slopes,
            return_lse=True,
            sink_ptr=self.sinks,
        )
        assert attn_metadata.extend_metadata is not None
        chunk_context_metadata = attn_metadata.extend_metadata.chunk_context_metadata
        num_chunks = chunk_context_metadata.num_chunks
        workspace = chunk_context_metadata.workspace
        cu_seqlens_kv = chunk_context_metadata.cu_seq_lens_chunk
        max_seqlens = chunk_context_metadata.max_seq_lens
        chunk_starts = chunk_context_metadata.chunk_starts
        token_to_batch = chunk_context_metadata.token_to_batch
        total_token_per_batch = chunk_context_metadata.total_token_per_batch
        key_fetched, value_fetched = workspace[0], workspace[1]
        chunked_output = None
        chunked_lse = None
        for chunk_idx in range(num_chunks):
            cp_mha_gather_cache(
                key_cache=key_cache,
                value_cache=value_cache,
                key=key_fetched,
                value=value_fetched,
                block_tables=block_table,
                k_scales=k_scale,
                v_scales=v_scale,
                cu_seqlens_kv=cu_seqlens_kv[chunk_idx],
                token_to_batch=token_to_batch[chunk_idx],
                seq_starts=chunk_starts[chunk_idx],
                dequant=is_quantized_kv_cache(self.kv_cache_dtype),
                kv_cache_layout="SHUFFLE"
                if rocm_aiter_ops.is_shuffle_kv_cache_enabled()
                else "NHD",
                total_tokens=total_token_per_batch[chunk_idx],
            )

            suf_out, suf_lse = rocm_aiter_ops.flash_attn_varlen_func(
                q=query,
                k=key_fetched,
                v=value_fetched,
                cu_seqlens_q=cu_seqlens_q,
                cu_seqlens_k=cu_seqlens_kv[chunk_idx],
                max_seqlen_q=max_seqlen_q,
                max_seqlen_k=max_seqlens[chunk_idx],
                min_seqlen_q=min_seqlen_q,
                dropout_p=0.0,
                softmax_scale=self.scale,
                causal=False,
                window_size=self.sliding_window,
                alibi_slopes=self.alibi_slopes,
                return_lse=True,
                sink_ptr=self.sinks,
            )
            if chunked_output is None:
                chunked_output = suf_out
                chunked_lse = suf_lse
            else:
                tmp_output = torch.empty_like(out)
                tmp_lse = torch.empty_like(lse)
                merge_attn_states(
                    output=tmp_output,
                    output_lse=tmp_lse,
                    prefix_output=chunked_output,
                    prefix_lse=chunked_lse,
                    suffix_output=suf_out,
                    suffix_lse=suf_lse,
                )
                chunked_output = tmp_output
                chunked_lse = tmp_lse

        merge_attn_states(
            output=output,
            prefix_output=chunked_output,
            prefix_lse=chunked_lse,
            suffix_output=out,
            suffix_lse=lse,
        )

    def forward(
        self,
        layer: torch.nn.Module,
        query: torch.Tensor,
        key: torch.Tensor | None,
        value: torch.Tensor | None,
        kv_cache: torch.Tensor,
        attn_metadata: AiterFlashAttentionMetadata,
        output: torch.Tensor,
        output_scale: torch.Tensor | None = None,
        output_block_scale: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Forward pass with AiterFlashAttention.

        Args:
            layer: The attention layer, providing the q/k/v quantization scales.
            query: shape = [num_tokens, num_heads, head_size]
            key: shape = [num_tokens, num_kv_heads, head_size]
            value: shape = [num_tokens, num_kv_heads, head_size]
            kv_cache: shape =
                [num_blocks, 2, block_size, num_kv_heads, head_size]
            attn_metadata: Metadata for attention.
            output: Tensor that the attention result is written into.
            output_scale: Scale for fused output quantization; not supported
                by this backend.
            output_block_scale: Block scale for fused output quantization;
                not supported by this backend.

        Returns:
            shape = [num_tokens, num_heads * head_size]
        NOTE: FP8 quantization, flash-attn expect the size of
              {q,k,v}_descale to be (num_sequences, num_kv_heads).
              We use torch's .expand() to avoid duplicating values

        """
        if output_scale is not None or output_block_scale is not None:
            raise NotImplementedError(
                "fused output quantization is not yet supported "
                "for AiterFlashAttentionImpl"
            )

        if attn_metadata is None:
            # Profiling run.
            return output.fill_(0)

        # IMPORTANT!
        # NOTE(woosuk): With piece-wise CUDA graphs, this method is
        # executed in eager-mode PyTorch. Thus, we need to be careful
        # about any CPU overhead in this method. For example, `view`
        # and `slice` (or `[:n]`) operations are surprisingly slow even
        # in the case they do not invoke any GPU ops.
        # Minimize the PyTorch ops in this method as much as possible.
        # Whenever making a change in this method, please benchmark the
        # performance to make sure it does not introduce any overhead.
        num_actual_tokens = attn_metadata.num_actual_tokens

        key_cache, value_cache = self._split_kv_cache(kv_cache)
        if is_quantized_kv_cache(self.kv_cache_dtype):
            key_cache = key_cache.view(current_platform.fp8_dtype())
            value_cache = value_cache.view(current_platform.fp8_dtype())

        # decode:extend:prefill
        query = query[:num_actual_tokens]
        if key is not None:
            key = key[:num_actual_tokens]
        if value is not None:
            value = value[:num_actual_tokens]

        output_actual_tokens = output[:num_actual_tokens]

        num_decodes = attn_metadata.num_decodes
        num_prefills = attn_metadata.num_prefills
        num_extends = attn_metadata.num_extends

        num_decode_tokens = attn_metadata.num_decode_tokens
        num_extend_tokens = attn_metadata.num_extend_tokens
        if self.kv_sharing_target_layer_name is not None and (
            num_prefills > 0 or num_extends > 0
        ):
            # Shared layers project only Q. Fetch the new K/V tokens from the
            # target cache for the direct prefill and extend suffix kernels.
            shared = attn_metadata.kv_sharing_metadata
            assert shared is not None
            key, value = shared.workspace[:, :num_actual_tokens].unbind(0)
            cp_mha_gather_cache(
                key_cache=key_cache,
                value_cache=value_cache,
                key=key[num_decode_tokens:],
                value=value[num_decode_tokens:],
                block_tables=attn_metadata.block_table[num_decodes:],
                k_scales=layer._k_scale,
                v_scales=layer._v_scale,
                cu_seqlens_kv=shared.query_start_loc,
                token_to_batch=shared.token_to_batch,
                seq_starts=shared.seq_starts,
                dequant=is_quantized_kv_cache(self.kv_cache_dtype),
                kv_cache_layout="SHUFFLE"
                if rocm_aiter_ops.is_shuffle_kv_cache_enabled()
                else "NHD",
                total_tokens=num_actual_tokens - num_decode_tokens,
            )
        if not attn_metadata.use_cascade:
            # calculate for pure prefills
            if num_prefills > 0:
                assert attn_metadata.prefill_metadata is not None
                assert key is not None and value is not None

                prefill_query = query[num_decode_tokens + num_extend_tokens :]
                prefill_key = key[num_decode_tokens + num_extend_tokens :]
                prefill_value = value[num_decode_tokens + num_extend_tokens :]

                rocm_aiter_ops.flash_attn_varlen_func(
                    q=prefill_query,
                    k=prefill_key,
                    v=prefill_value,
                    cu_seqlens_q=attn_metadata.prefill_metadata.query_start_loc,
                    cu_seqlens_k=attn_metadata.prefill_metadata.query_start_loc,
                    max_seqlen_q=attn_metadata.prefill_metadata.max_query_len,
                    max_seqlen_k=attn_metadata.prefill_metadata.max_seq_len,
                    min_seqlen_q=1,
                    dropout_p=0.0,
                    softmax_scale=self.scale,
                    causal=attn_metadata.causal,
                    window_size=self.sliding_window,
                    alibi_slopes=self.alibi_slopes,
                    out=output_actual_tokens[num_decode_tokens + num_extend_tokens :],
                    sink_ptr=self.sinks,
                )

            # calculate for extends
            if num_extends > 0:
                assert attn_metadata.extend_metadata is not None
                assert key is not None and value is not None
                extend_tokens_slice = slice(
                    num_decode_tokens, num_decode_tokens + num_extend_tokens
                )
                extend_queries = query[extend_tokens_slice]
                extend_keys = key[extend_tokens_slice]
                extend_values = value[extend_tokens_slice]
                extend_outputs = output[extend_tokens_slice]
                k_scale = layer._k_scale
                v_scale = layer._v_scale
                if rocm_aiter_ops.is_shuffle_kv_cache_enabled():
                    k_scale = attn_metadata.k_scale
                    v_scale = attn_metadata.v_scale
                self.extend_forward(
                    attn_metadata=attn_metadata,
                    query=extend_queries,
                    key=extend_keys,
                    value=extend_values,
                    key_cache=key_cache,
                    value_cache=value_cache,
                    output=extend_outputs,
                    cu_seqlens_q=attn_metadata.extend_metadata.query_start_loc,
                    max_seqlen_q=attn_metadata.extend_metadata.max_query_len,
                    min_seqlen_q=1,
                    max_seqlen_k=attn_metadata.extend_metadata.max_seq_len,
                    block_table=attn_metadata.block_table[
                        num_decodes : num_decodes + num_extends
                    ],
                    slot_mapping=attn_metadata.slot_mapping[
                        num_decodes : num_decodes + num_extends
                    ],
                    k_scale=k_scale,
                    v_scale=v_scale,
                )

            # calculate for decodes
            if num_decodes > 0:
                assert attn_metadata.decode_metadata is not None
                decode_max_query_len = attn_metadata.decode_metadata.max_query_len

                # Use unified_attention for the decodes the paged kernels can't
                # take: sliding window, sinks, or a multi-token batch
                # (pa_fwd_asm and paged_attention_v1 don't support sinks).
                if (
                    self.sliding_window[0] != -1
                    or decode_max_query_len > 1
                    or self.sinks is not None
                ):
                    k_descale, v_descale = self._get_kv_cache_descales(
                        layer, num_decodes
                    )
                    assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), (
                        "Shuffle KV cache layout is not supported with sliding "
                        "window, sinks, or speculative decoding (multi-token decode)."
                    )
                    if not attn_metadata.causal:
                        from aiter.ops.triton.attention.mha_v3 import (
                            flash_attn_with_kvcache,
                        )

                        decode_query = query[:num_decode_tokens].reshape(
                            num_decodes,
                            decode_max_query_len,
                            query.shape[1],
                            query.shape[2],
                        )
                        decode_out = flash_attn_with_kvcache(
                            q=decode_query,
                            k_cache=key_cache,
                            v_cache=value_cache,
                            cache_seqlens=attn_metadata.seq_lens[:num_decodes],
                            softmax_scale=self.scale,
                            causal=attn_metadata.causal,
                            window_size=self.sliding_window,
                            softcap=self.logits_soft_cap,
                            q_descale=None,
                            k_descale=k_descale,
                            v_descale=v_descale,
                            page_table=attn_metadata.block_table[:num_decodes],
                        )
                        output[:num_decode_tokens].copy_(
                            decode_out.reshape(
                                num_decode_tokens,
                                query.shape[1],
                                query.shape[2],
                            )
                        )
                    else:
                        # Non-uniform query lengths can appear in real serving
                        # traffic (e.g. mixed datasets). Fall back to varlen
                        # unified_attention instead of asserting.
                        from aiter.ops.triton.unified_attention import (
                            unified_attention,
                        )

                        unified_attention(
                            q=query[:num_decode_tokens],
                            k=key_cache,
                            v=value_cache,
                            out=output[:num_decode_tokens],
                            cu_seqlens_q=attn_metadata.query_start_loc[
                                : num_decodes + 1
                            ],
                            max_seqlen_q=decode_max_query_len,
                            seqused_k=attn_metadata.seq_lens[:num_decodes],
                            max_seqlen_k=attn_metadata.max_seq_len,
                            softmax_scale=self.scale,
                            causal=True,
                            alibi_slopes=self.alibi_slopes,
                            window_size=self.sliding_window,
                            block_table=attn_metadata.block_table[:num_decodes],
                            softcap=self.logits_soft_cap,
                            q_descale=None,
                            k_descale=k_descale,
                            v_descale=v_descale,
                            sinks=self.sinks,
                        )
                    return

                # The ll4mi kernel in paged_attention_v1 requires
                # HEAD_SIZE >= 16 * NWARPS (= 64 on ROCm with NWARPS=4).
                # For smaller head sizes or sliding window attention,
                # fall back to the unified_attention triton kernel which
                # handles both correctly.
                _MIN_HEAD_SIZE_FOR_LL4MI = 64
                use_unified_attention = self.head_size < _MIN_HEAD_SIZE_FOR_LL4MI

                if use_unified_attention:
                    k_descale, v_descale = self._get_kv_cache_descales(
                        layer, num_decodes
                    )
                    assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), (
                        "unified_attention fallback with shuffle layout "
                        "is not supported yet."
                    )
                    from aiter.ops.triton.unified_attention import (
                        unified_attention,
                    )

                    decode_cu_seqlens_q = attn_metadata.query_start_loc[
                        : num_decodes + 1
                    ]
                    unified_attention(
                        q=query[:num_decode_tokens],
                        k=key_cache,
                        v=value_cache,
                        out=output[:num_decode_tokens],
                        cu_seqlens_q=decode_cu_seqlens_q,
                        max_seqlen_q=1,
                        seqused_k=attn_metadata.seq_lens[:num_decodes],
                        max_seqlen_k=attn_metadata.max_seq_len,
                        softmax_scale=self.scale,
                        causal=True,
                        alibi_slopes=self.alibi_slopes,
                        window_size=self.sliding_window,
                        block_table=attn_metadata.block_table[:num_decodes],
                        softcap=self.logits_soft_cap,
                        q_descale=None,
                        k_descale=k_descale,
                        v_descale=v_descale,
                    )
                elif rocm_aiter_ops.is_shuffle_kv_cache_enabled():
                    _, num_heads, head_size = query.shape
                    num_blocks, block_size, num_kv_heads, _ = key_cache.shape
                    x = 16 // key_cache.element_size()
                    new_key_cache = key_cache.reshape(
                        num_blocks, num_kv_heads, head_size // x, block_size, x
                    )
                    new_value_cache = value_cache.reshape(
                        num_blocks, num_kv_heads, block_size // x, head_size, x
                    )

                    # This kernel derives block addresses arithmetically, so it
                    # only reads the right bytes when each side is a dense
                    # plane of pages.
                    assert new_key_cache.stride(0) == (
                        num_kv_heads * head_size * block_size
                    ), (
                        "paged_attention_common needs the K/V sides as dense "
                        "planes; the resolved KV cache layout interleaves them "
                        "within a block."
                    )
                    num_seqs = attn_metadata.seq_lens.shape[0]
                    max_num_partitions = (
                        attn_metadata.max_seq_len + _PARTITION_SIZE_ROCM - 1
                    ) // _PARTITION_SIZE_ROCM
                    tmp_out = torch.empty(
                        (num_seqs, num_heads, max_num_partitions, head_size),
                        dtype=query.dtype,
                        device=query.device,
                    )
                    exp_sums = torch.empty(
                        (num_seqs, num_heads, max_num_partitions),
                        dtype=torch.float32,
                        device=query.device,
                    )
                    max_logits = torch.empty_like(exp_sums)
                    k_qscale = (
                        layer._k_scale
                        if attn_metadata.k_scale is None
                        else attn_metadata.k_scale
                    )
                    v_qscale = (
                        layer._v_scale
                        if attn_metadata.v_scale is None
                        else attn_metadata.v_scale
                    )
                    rocm_aiter_ops.paged_attention_common(
                        Q=query[:num_decode_tokens],
                        K=new_key_cache,
                        V=new_value_cache,
                        tmp_out=tmp_out,
                        max_logits=max_logits,
                        exp_sums=exp_sums,
                        max_seq_len=attn_metadata.max_seq_len,
                        block_tables=attn_metadata.block_table[:num_decodes],
                        context_lens=attn_metadata.seq_lens[:num_decodes],
                        block_tables_stride0=attn_metadata.block_table[
                            :num_decodes
                        ].stride(0),
                        scale=self.scale,
                        K_QScale_hip=k_qscale,
                        V_QScale_hip=v_qscale,
                        K_QScale_asm=k_qscale,
                        V_QScale_asm=v_qscale,
                        out_=output[:num_decode_tokens],
                        kv_cache_dtype=self.kv_cache_dtype,
                    )
                else:
                    _, num_heads, head_size = query.shape
                    nbytes_per_qo_elem = torch.finfo(query.dtype).bits // 8
                    num_seqs = attn_metadata.seq_lens.shape[0]
                    max_num_partitions = (
                        attn_metadata.max_seq_len + _PARTITION_SIZE_ROCM - 1
                    ) // _PARTITION_SIZE_ROCM

                    workspace_buffer = torch.empty(
                        (num_seqs * num_heads * max_num_partitions * head_size)
                        * nbytes_per_qo_elem
                        + 2 * (num_seqs * num_heads * max_num_partitions) * 4,
                        dtype=torch.uint8,
                        device=output.device,
                    )

                    # import so that aiter register the op to the namespace of
                    # torch.ops.aiter
                    import aiter  # noqa: F401

                    torch.ops.aiter.paged_attention_v1(
                        output[:num_decode_tokens],
                        workspace_buffer,
                        query[:num_decode_tokens],
                        key_cache,
                        value_cache,
                        self.scale,
                        attn_metadata.block_table[:num_decodes],
                        attn_metadata.query_start_loc[:num_decodes],
                        attn_metadata.seq_lens[:num_decodes],
                        attn_metadata.max_seq_len,
                        self.alibi_slopes,
                        self.kv_cache_dtype,
                        "NHD",
                        self.logits_soft_cap,
                        layer._k_scale,
                        layer._v_scale,
                        None,
                        _PARTITION_SIZE_ROCM,
                        1,
                        self.sliding_window[0] + 1,
                    )
        else:
            raise NotImplementedError(
                "Cascade attention is not implemented for ROCM AITER"
            )

        return output

    def _split_kv_cache(
        self, kv_cache: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if rocm_aiter_ops.is_shuffle_kv_cache_enabled():
            # (B, 2, N, H*hs) -> one (B, N, H, hs) per side, which is what the
            # shuffle read/write kernels reinterpret in place. The heads are
            # folded into the content dim by customize_spec.
            key_cache, value_cache = kv_cache.unflatten(
                -1, (self.num_kv_heads, self.head_size)
            ).unbind(1)
            return key_cache, value_cache
        # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs))
        return kv_cache.transpose(1, 2).split(self.head_size, dim=-1)

    def do_kv_cache_update(
        self,
        layer: AttentionLayer,
        key: torch.Tensor,
        value: torch.Tensor,
        kv_cache: torch.Tensor,
        slot_mapping: torch.Tensor,
    ):
        key_cache, value_cache = self._split_kv_cache(kv_cache)

        # key and value may be None in the case of cross attention. They are
        # calculated once based on the output from the encoder and then cached
        # in KV cache.
        if is_quantized_kv_cache(self.kv_cache_dtype):
            key_cache = key_cache.view(current_platform.fp8_dtype())
            value_cache = value_cache.view(current_platform.fp8_dtype())
        # Reshape the input keys and values and store them in the cache.
        # Skip this if sharing KV cache with an earlier attention layer.
        # NOTE(woosuk): Here, key and value are padded while slot_mapping
        # is not padded. However, we don't need to do
        # key[:num_actual_tokens] and value[:num_actual_tokens] because
        # the reshape_and_cache_flash op uses the slot_mapping's shape
        # to determine the number of actual tokens.
        if rocm_aiter_ops.is_shuffle_kv_cache_enabled():
            # We may calculate per token quant scale in
            # reshape_and_cache_shuffle_triton which might differ from
            # vllm's style when shuffle layout is used.
            k_scale = layer._k_scale
            v_scale = layer._v_scale
            assert k_scale is not None and v_scale is not None, (
                "k_scale and v_scale are required for shuffled update"
            )
            reshape_and_cache_shuffle_triton(
                key,
                value,
                key_cache,
                value_cache,
                slot_mapping,
                self.kv_cache_dtype,
                k_scale,
                v_scale,
            )
        else:
            torch.ops._C_cache_ops.reshape_and_cache_flash(
                key,
                value,
                key_cache,
                value_cache,
                slot_mapping,
                self.kv_cache_dtype,
                layer._k_scale,
                layer._v_scale,
            )

    def fused_rope_kvcache_supported(self):
        # Only support fusion when shuffle KV cache layout is not used;
        # shuffle layout uses a different cache update path.
        return (
            rocm_aiter_ops.is_enabled()
            and not rocm_aiter_ops.is_shuffle_kv_cache_enabled()
        )

    def fused_qk_norm_rope_kvcache_supported(self):
        # Only fuse when shuffle layout is off; the shuffle write path uses a
        # dedicated cache update, mirroring fused_rope_kvcache_supported.
        return (
            rocm_aiter_ops.is_enabled()
            and not rocm_aiter_ops.is_shuffle_kv_cache_enabled()
        )

    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,
    ):
        key_cache, value_cache = self._split_kv_cache(kv_cache)
        rocm_aiter_ops.do_qk_norm_rope_kvcache_update(
            qkv=qkv,
            q_weight=q_weight,
            k_weight=k_weight,
            cos_sin_cache=cos_sin_cache,
            positions=positions,
            num_heads_q=self.num_heads,
            num_heads_k=self.num_kv_heads,
            head_dim=self.head_size,
            is_neox=is_neox,
            rms_norm_eps=rms_norm_eps,
            q_out=q_out,
            k_out=k_out,
            key_cache=key_cache,
            value_cache=value_cache,
            slot_mapping=layer_slot_mapping,
            k_scale=layer._k_scale_cpu,
            v_scale=layer._v_scale_cpu,
            kv_cache_dtype=self.kv_cache_dtype,
            use_shuffle_layout=rocm_aiter_ops.is_shuffle_kv_cache_enabled(),
        )

    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,
    ):
        # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs))
        key_cache, value_cache = self._split_kv_cache(kv_cache)
        flash_layout = True

        is_fp8_kv_cache = is_quantized_kv_cache(self.kv_cache_dtype)
        if is_fp8_kv_cache:
            key_cache = key_cache.view(current_platform.fp8_dtype())
            value_cache = value_cache.view(current_platform.fp8_dtype())

        rocm_aiter_ops.triton_rope_and_cache(
            query,
            key,
            value,
            positions,
            cos_sin_cache,
            is_neox,
            key_cache,
            value_cache,
            layer_slot_mapping,
            layer._k_scale,
            layer._v_scale,
            flash_layout,
            is_fp8_kv_cache,
        )

forward(layer, query, key, value, kv_cache, attn_metadata, output, output_scale=None, output_block_scale=None)

Forward pass with AiterFlashAttention.

Parameters:

• layer

  (Module) –

  The attention layer, providing the q/k/v quantization scales.

• query

  (Tensor) –

  shape = [num_tokens, num_heads, head_size]

• key

  (Tensor | None) –

  shape = [num_tokens, num_kv_heads, head_size]

• value

  (Tensor | None) –

  shape = [num_tokens, num_kv_heads, head_size]

• kv_cache

  (Tensor) –

  shape = [num_blocks, 2, block_size, num_kv_heads, head_size]

• attn_metadata

  (AiterFlashAttentionMetadata) –

  Metadata for attention.

• output

  (Tensor) –

  Tensor that the attention result is written into.

• output_scale

  (Tensor | None, default: None ) –

  Scale for fused output quantization; not supported by this backend.

• output_block_scale

  (Tensor | None, default: None ) –

  Block scale for fused output quantization; not supported by this backend.

Returns:

• Tensor –

  shape = [num_tokens, num_heads * head_size]

NOTE: FP8 quantization, flash-attn expect the size of {q,k,v}_descale to be (num_sequences, num_kv_heads). We use torch's .expand() to avoid duplicating values

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

def forward(
    self,
    layer: torch.nn.Module,
    query: torch.Tensor,
    key: torch.Tensor | None,
    value: torch.Tensor | None,
    kv_cache: torch.Tensor,
    attn_metadata: AiterFlashAttentionMetadata,
    output: torch.Tensor,
    output_scale: torch.Tensor | None = None,
    output_block_scale: torch.Tensor | None = None,
) -> torch.Tensor:
    """Forward pass with AiterFlashAttention.

    Args:
        layer: The attention layer, providing the q/k/v quantization scales.
        query: shape = [num_tokens, num_heads, head_size]
        key: shape = [num_tokens, num_kv_heads, head_size]
        value: shape = [num_tokens, num_kv_heads, head_size]
        kv_cache: shape =
            [num_blocks, 2, block_size, num_kv_heads, head_size]
        attn_metadata: Metadata for attention.
        output: Tensor that the attention result is written into.
        output_scale: Scale for fused output quantization; not supported
            by this backend.
        output_block_scale: Block scale for fused output quantization;
            not supported by this backend.

    Returns:
        shape = [num_tokens, num_heads * head_size]
    NOTE: FP8 quantization, flash-attn expect the size of
          {q,k,v}_descale to be (num_sequences, num_kv_heads).
          We use torch's .expand() to avoid duplicating values

    """
    if output_scale is not None or output_block_scale is not None:
        raise NotImplementedError(
            "fused output quantization is not yet supported "
            "for AiterFlashAttentionImpl"
        )

    if attn_metadata is None:
        # Profiling run.
        return output.fill_(0)

    # IMPORTANT!
    # NOTE(woosuk): With piece-wise CUDA graphs, this method is
    # executed in eager-mode PyTorch. Thus, we need to be careful
    # about any CPU overhead in this method. For example, `view`
    # and `slice` (or `[:n]`) operations are surprisingly slow even
    # in the case they do not invoke any GPU ops.
    # Minimize the PyTorch ops in this method as much as possible.
    # Whenever making a change in this method, please benchmark the
    # performance to make sure it does not introduce any overhead.
    num_actual_tokens = attn_metadata.num_actual_tokens

    key_cache, value_cache = self._split_kv_cache(kv_cache)
    if is_quantized_kv_cache(self.kv_cache_dtype):
        key_cache = key_cache.view(current_platform.fp8_dtype())
        value_cache = value_cache.view(current_platform.fp8_dtype())

    # decode:extend:prefill
    query = query[:num_actual_tokens]
    if key is not None:
        key = key[:num_actual_tokens]
    if value is not None:
        value = value[:num_actual_tokens]

    output_actual_tokens = output[:num_actual_tokens]

    num_decodes = attn_metadata.num_decodes
    num_prefills = attn_metadata.num_prefills
    num_extends = attn_metadata.num_extends

    num_decode_tokens = attn_metadata.num_decode_tokens
    num_extend_tokens = attn_metadata.num_extend_tokens
    if self.kv_sharing_target_layer_name is not None and (
        num_prefills > 0 or num_extends > 0
    ):
        # Shared layers project only Q. Fetch the new K/V tokens from the
        # target cache for the direct prefill and extend suffix kernels.
        shared = attn_metadata.kv_sharing_metadata
        assert shared is not None
        key, value = shared.workspace[:, :num_actual_tokens].unbind(0)
        cp_mha_gather_cache(
            key_cache=key_cache,
            value_cache=value_cache,
            key=key[num_decode_tokens:],
            value=value[num_decode_tokens:],
            block_tables=attn_metadata.block_table[num_decodes:],
            k_scales=layer._k_scale,
            v_scales=layer._v_scale,
            cu_seqlens_kv=shared.query_start_loc,
            token_to_batch=shared.token_to_batch,
            seq_starts=shared.seq_starts,
            dequant=is_quantized_kv_cache(self.kv_cache_dtype),
            kv_cache_layout="SHUFFLE"
            if rocm_aiter_ops.is_shuffle_kv_cache_enabled()
            else "NHD",
            total_tokens=num_actual_tokens - num_decode_tokens,
        )
    if not attn_metadata.use_cascade:
        # calculate for pure prefills
        if num_prefills > 0:
            assert attn_metadata.prefill_metadata is not None
            assert key is not None and value is not None

            prefill_query = query[num_decode_tokens + num_extend_tokens :]
            prefill_key = key[num_decode_tokens + num_extend_tokens :]
            prefill_value = value[num_decode_tokens + num_extend_tokens :]

            rocm_aiter_ops.flash_attn_varlen_func(
                q=prefill_query,
                k=prefill_key,
                v=prefill_value,
                cu_seqlens_q=attn_metadata.prefill_metadata.query_start_loc,
                cu_seqlens_k=attn_metadata.prefill_metadata.query_start_loc,
                max_seqlen_q=attn_metadata.prefill_metadata.max_query_len,
                max_seqlen_k=attn_metadata.prefill_metadata.max_seq_len,
                min_seqlen_q=1,
                dropout_p=0.0,
                softmax_scale=self.scale,
                causal=attn_metadata.causal,
                window_size=self.sliding_window,
                alibi_slopes=self.alibi_slopes,
                out=output_actual_tokens[num_decode_tokens + num_extend_tokens :],
                sink_ptr=self.sinks,
            )

        # calculate for extends
        if num_extends > 0:
            assert attn_metadata.extend_metadata is not None
            assert key is not None and value is not None
            extend_tokens_slice = slice(
                num_decode_tokens, num_decode_tokens + num_extend_tokens
            )
            extend_queries = query[extend_tokens_slice]
            extend_keys = key[extend_tokens_slice]
            extend_values = value[extend_tokens_slice]
            extend_outputs = output[extend_tokens_slice]
            k_scale = layer._k_scale
            v_scale = layer._v_scale
            if rocm_aiter_ops.is_shuffle_kv_cache_enabled():
                k_scale = attn_metadata.k_scale
                v_scale = attn_metadata.v_scale
            self.extend_forward(
                attn_metadata=attn_metadata,
                query=extend_queries,
                key=extend_keys,
                value=extend_values,
                key_cache=key_cache,
                value_cache=value_cache,
                output=extend_outputs,
                cu_seqlens_q=attn_metadata.extend_metadata.query_start_loc,
                max_seqlen_q=attn_metadata.extend_metadata.max_query_len,
                min_seqlen_q=1,
                max_seqlen_k=attn_metadata.extend_metadata.max_seq_len,
                block_table=attn_metadata.block_table[
                    num_decodes : num_decodes + num_extends
                ],
                slot_mapping=attn_metadata.slot_mapping[
                    num_decodes : num_decodes + num_extends
                ],
                k_scale=k_scale,
                v_scale=v_scale,
            )

        # calculate for decodes
        if num_decodes > 0:
            assert attn_metadata.decode_metadata is not None
            decode_max_query_len = attn_metadata.decode_metadata.max_query_len

            # Use unified_attention for the decodes the paged kernels can't
            # take: sliding window, sinks, or a multi-token batch
            # (pa_fwd_asm and paged_attention_v1 don't support sinks).
            if (
                self.sliding_window[0] != -1
                or decode_max_query_len > 1
                or self.sinks is not None
            ):
                k_descale, v_descale = self._get_kv_cache_descales(
                    layer, num_decodes
                )
                assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), (
                    "Shuffle KV cache layout is not supported with sliding "
                    "window, sinks, or speculative decoding (multi-token decode)."
                )
                if not attn_metadata.causal:
                    from aiter.ops.triton.attention.mha_v3 import (
                        flash_attn_with_kvcache,
                    )

                    decode_query = query[:num_decode_tokens].reshape(
                        num_decodes,
                        decode_max_query_len,
                        query.shape[1],
                        query.shape[2],
                    )
                    decode_out = flash_attn_with_kvcache(
                        q=decode_query,
                        k_cache=key_cache,
                        v_cache=value_cache,
                        cache_seqlens=attn_metadata.seq_lens[:num_decodes],
                        softmax_scale=self.scale,
                        causal=attn_metadata.causal,
                        window_size=self.sliding_window,
                        softcap=self.logits_soft_cap,
                        q_descale=None,
                        k_descale=k_descale,
                        v_descale=v_descale,
                        page_table=attn_metadata.block_table[:num_decodes],
                    )
                    output[:num_decode_tokens].copy_(
                        decode_out.reshape(
                            num_decode_tokens,
                            query.shape[1],
                            query.shape[2],
                        )
                    )
                else:
                    # Non-uniform query lengths can appear in real serving
                    # traffic (e.g. mixed datasets). Fall back to varlen
                    # unified_attention instead of asserting.
                    from aiter.ops.triton.unified_attention import (
                        unified_attention,
                    )

                    unified_attention(
                        q=query[:num_decode_tokens],
                        k=key_cache,
                        v=value_cache,
                        out=output[:num_decode_tokens],
                        cu_seqlens_q=attn_metadata.query_start_loc[
                            : num_decodes + 1
                        ],
                        max_seqlen_q=decode_max_query_len,
                        seqused_k=attn_metadata.seq_lens[:num_decodes],
                        max_seqlen_k=attn_metadata.max_seq_len,
                        softmax_scale=self.scale,
                        causal=True,
                        alibi_slopes=self.alibi_slopes,
                        window_size=self.sliding_window,
                        block_table=attn_metadata.block_table[:num_decodes],
                        softcap=self.logits_soft_cap,
                        q_descale=None,
                        k_descale=k_descale,
                        v_descale=v_descale,
                        sinks=self.sinks,
                    )
                return

            # The ll4mi kernel in paged_attention_v1 requires
            # HEAD_SIZE >= 16 * NWARPS (= 64 on ROCm with NWARPS=4).
            # For smaller head sizes or sliding window attention,
            # fall back to the unified_attention triton kernel which
            # handles both correctly.
            _MIN_HEAD_SIZE_FOR_LL4MI = 64
            use_unified_attention = self.head_size < _MIN_HEAD_SIZE_FOR_LL4MI

            if use_unified_attention:
                k_descale, v_descale = self._get_kv_cache_descales(
                    layer, num_decodes
                )
                assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), (
                    "unified_attention fallback with shuffle layout "
                    "is not supported yet."
                )
                from aiter.ops.triton.unified_attention import (
                    unified_attention,
                )

                decode_cu_seqlens_q = attn_metadata.query_start_loc[
                    : num_decodes + 1
                ]
                unified_attention(
                    q=query[:num_decode_tokens],
                    k=key_cache,
                    v=value_cache,
                    out=output[:num_decode_tokens],
                    cu_seqlens_q=decode_cu_seqlens_q,
                    max_seqlen_q=1,
                    seqused_k=attn_metadata.seq_lens[:num_decodes],
                    max_seqlen_k=attn_metadata.max_seq_len,
                    softmax_scale=self.scale,
                    causal=True,
                    alibi_slopes=self.alibi_slopes,
                    window_size=self.sliding_window,
                    block_table=attn_metadata.block_table[:num_decodes],
                    softcap=self.logits_soft_cap,
                    q_descale=None,
                    k_descale=k_descale,
                    v_descale=v_descale,
                )
            elif rocm_aiter_ops.is_shuffle_kv_cache_enabled():
                _, num_heads, head_size = query.shape
                num_blocks, block_size, num_kv_heads, _ = key_cache.shape
                x = 16 // key_cache.element_size()
                new_key_cache = key_cache.reshape(
                    num_blocks, num_kv_heads, head_size // x, block_size, x
                )
                new_value_cache = value_cache.reshape(
                    num_blocks, num_kv_heads, block_size // x, head_size, x
                )

                # This kernel derives block addresses arithmetically, so it
                # only reads the right bytes when each side is a dense
                # plane of pages.
                assert new_key_cache.stride(0) == (
                    num_kv_heads * head_size * block_size
                ), (
                    "paged_attention_common needs the K/V sides as dense "
                    "planes; the resolved KV cache layout interleaves them "
                    "within a block."
                )
                num_seqs = attn_metadata.seq_lens.shape[0]
                max_num_partitions = (
                    attn_metadata.max_seq_len + _PARTITION_SIZE_ROCM - 1
                ) // _PARTITION_SIZE_ROCM
                tmp_out = torch.empty(
                    (num_seqs, num_heads, max_num_partitions, head_size),
                    dtype=query.dtype,
                    device=query.device,
                )
                exp_sums = torch.empty(
                    (num_seqs, num_heads, max_num_partitions),
                    dtype=torch.float32,
                    device=query.device,
                )
                max_logits = torch.empty_like(exp_sums)
                k_qscale = (
                    layer._k_scale
                    if attn_metadata.k_scale is None
                    else attn_metadata.k_scale
                )
                v_qscale = (
                    layer._v_scale
                    if attn_metadata.v_scale is None
                    else attn_metadata.v_scale
                )
                rocm_aiter_ops.paged_attention_common(
                    Q=query[:num_decode_tokens],
                    K=new_key_cache,
                    V=new_value_cache,
                    tmp_out=tmp_out,
                    max_logits=max_logits,
                    exp_sums=exp_sums,
                    max_seq_len=attn_metadata.max_seq_len,
                    block_tables=attn_metadata.block_table[:num_decodes],
                    context_lens=attn_metadata.seq_lens[:num_decodes],
                    block_tables_stride0=attn_metadata.block_table[
                        :num_decodes
                    ].stride(0),
                    scale=self.scale,
                    K_QScale_hip=k_qscale,
                    V_QScale_hip=v_qscale,
                    K_QScale_asm=k_qscale,
                    V_QScale_asm=v_qscale,
                    out_=output[:num_decode_tokens],
                    kv_cache_dtype=self.kv_cache_dtype,
                )
            else:
                _, num_heads, head_size = query.shape
                nbytes_per_qo_elem = torch.finfo(query.dtype).bits // 8
                num_seqs = attn_metadata.seq_lens.shape[0]
                max_num_partitions = (
                    attn_metadata.max_seq_len + _PARTITION_SIZE_ROCM - 1
                ) // _PARTITION_SIZE_ROCM

                workspace_buffer = torch.empty(
                    (num_seqs * num_heads * max_num_partitions * head_size)
                    * nbytes_per_qo_elem
                    + 2 * (num_seqs * num_heads * max_num_partitions) * 4,
                    dtype=torch.uint8,
                    device=output.device,
                )

                # import so that aiter register the op to the namespace of
                # torch.ops.aiter
                import aiter  # noqa: F401

                torch.ops.aiter.paged_attention_v1(
                    output[:num_decode_tokens],
                    workspace_buffer,
                    query[:num_decode_tokens],
                    key_cache,
                    value_cache,
                    self.scale,
                    attn_metadata.block_table[:num_decodes],
                    attn_metadata.query_start_loc[:num_decodes],
                    attn_metadata.seq_lens[:num_decodes],
                    attn_metadata.max_seq_len,
                    self.alibi_slopes,
                    self.kv_cache_dtype,
                    "NHD",
                    self.logits_soft_cap,
                    layer._k_scale,
                    layer._v_scale,
                    None,
                    _PARTITION_SIZE_ROCM,
                    1,
                    self.sliding_window[0] + 1,
                )
    else:
        raise NotImplementedError(
            "Cascade attention is not implemented for ROCM AITER"
        )

    return output

AiterFlashAttentionMetadataBuilder

Bases: AttentionMetadataBuilder[AiterFlashAttentionMetadata]

Methods:

• build_for_drafting –

  Build attention metadata for draft model without CPU-GPU sync.

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

class AiterFlashAttentionMetadataBuilder(
    AttentionMetadataBuilder[AiterFlashAttentionMetadata]
):
    _cudagraph_support = AttentionCGSupport.UNIFORM_BATCH

    def __init__(
        self,
        kv_cache_spec: AttentionSpec,
        layer_names: list[str],
        vllm_config: VllmConfig,
        device: torch.device,
    ):
        super().__init__(kv_cache_spec, layer_names, vllm_config, device)

        self.model_config = vllm_config.model_config
        self.parallel_config = vllm_config.parallel_config
        self.cache_config = vllm_config.cache_config

        self.num_heads_q = self.model_config.get_num_attention_heads(
            self.parallel_config
        )
        self.num_heads_kv = self.model_config.get_num_kv_heads(self.parallel_config)
        self.headdim = self.model_config.get_head_size()
        self.block_size = kv_cache_spec.block_size
        # Sliding window size to be used with the AOT scheduler will be
        # populated on first build() call.
        self.aot_sliding_window: tuple[int, int] | None = None
        self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)

        sliding_window_configs: set[tuple[int, int] | None] = set()
        kv_sharing_shape = None
        layers = get_layers_from_vllm_config(self.vllm_config, Attention)
        for name, layer in layers.items():
            if name not in layer_names:
                continue
            assert isinstance(layer.impl, AiterFlashAttentionImpl), (
                "Aiter Flash Attention Metadata Builder can only be used "
                "with Aiter Flash Attention Impl."
            )
            sliding_window_configs.add(layer.impl.sliding_window)
            if layer.kv_sharing_target_layer_name is not None:
                kv_sharing_shape = (layer.impl.num_kv_heads, layer.impl.head_size)

        while len(sliding_window_configs) > 0:
            sliding_window_config = sliding_window_configs.pop()
            if sliding_window_config is not None and sliding_window_config[0] != -1:
                assert self.aot_sliding_window is None, (
                    "Aiter Flash ATTENTION can only support one valid sliding window!"
                )
                self.aot_sliding_window = sliding_window_config

        self.extend_workspace = torch.empty(
            [2, _CP_TOKENS_PER_ITER_ROCM, self.num_heads_kv, self.headdim],
            dtype=self.model_config.dtype,
            device=device,
        )
        self.scale = torch.tensor([1.0], dtype=torch.float, device=self.device)
        self.kv_sharing_workspace = (
            torch.empty(
                (
                    2,
                    vllm_config.scheduler_config.max_num_batched_tokens,
                    *kv_sharing_shape,
                ),
                dtype=self.model_config.dtype,
                device=device,
            )
            if kv_sharing_shape is not None
            else None
        )

    def build_for_cudagraph_capture(
        self, common_attn_metadata: CommonAttentionMetadata
    ):
        return self.build(
            common_prefix_len=0, common_attn_metadata=common_attn_metadata
        )

    def build(
        self,
        common_prefix_len: int,
        common_attn_metadata: CommonAttentionMetadata,
        fast_build: bool = False,
    ) -> "AiterFlashAttentionMetadata":
        assert self.reorder_batch_threshold is not None
        split_ret = split_decodes_prefills_and_extends(
            common_attn_metadata,
            decode_threshold=self.reorder_batch_threshold,
        )
        # Allocate scales for fp8 shuffle kv cache with shuffle_kv_cache enabled
        if (
            rocm_aiter_ops.is_shuffle_kv_cache_enabled()
            and self.scale.numel() == 1
            and is_quantized_kv_cache(self.vllm_config.cache_config.cache_dtype)
        ):
            # Size the scales from a layer this builder owns. The draft model
            # runs its own builder over its own KV cache, so the first layer of
            # the whole config can carry an unrelated block count.
            kv_cache_shape = self.vllm_config.compilation_config.static_forward_context[
                self.layer_names[0]
            ].kv_cache.shape
            num_blocks = kv_cache_shape[0]
            self.scale = torch.ones(
                [num_blocks, self.num_heads_kv, self.block_size],
                dtype=torch.float32,
                device=self.device,
            )
        (
            num_decodes,
            num_extends,
            num_prefills,
            num_decode_tokens,
            num_extend_tokens,
            _,
        ) = split_ret

        query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu

        with gpu_sync_allowed():
            # Only copy seq_lens to CPU when prefill or extend is present to avoid a
            # blocking device→host transfer.
            seq_lens = (
                common_attn_metadata.seq_lens.cpu()
                if num_prefills > 0 or num_extends > 0
                else None
            )

        query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]

        kv_sharing_metadata = None
        if self.kv_sharing_workspace is not None and seq_lens is not None:
            query_lens = query_lens_cpu[num_decodes:]
            token_to_batch = torch.repeat_interleave(
                torch.arange(len(query_lens), dtype=torch.int32, device="cpu"),
                query_lens,
            )
            query_start_loc = common_attn_metadata.query_start_loc[num_decodes:]
            kv_sharing_metadata = AiterKVSharingMetadata(
                workspace=self.kv_sharing_workspace,
                query_start_loc=query_start_loc - query_start_loc[0],
                token_to_batch=async_tensor_h2d(token_to_batch, self.device),
                seq_starts=async_tensor_h2d(
                    seq_lens[num_decodes:] - query_lens, self.device
                ),
            )

        decode_metadata = None
        if num_decodes > 0:
            decode_metadata = AiterFlashAttentionDecodeMetadata(
                max_query_len=query_lens_cpu[:num_decodes].max().item(),
            )

        prefill_metadata = None
        if num_prefills > 0:
            assert seq_lens is not None
            query_lens_for_prefill = query_lens_cpu[num_decodes + num_extends :]
            query_start_loc_device = common_attn_metadata.query_start_loc[
                num_decodes + num_extends :
            ]
            prefill_metadata = AiterFlashAttentionPrefillMetadata(
                max_query_len=query_lens_for_prefill.max().item(),
                max_seq_len=seq_lens[num_decodes + num_extends :].max().item(),
                query_start_loc=query_start_loc_device - query_start_loc_device[0],
            )

        extend_metadata = None
        if num_extends > 0:
            assert seq_lens is not None
            num_extends_slice = slice(num_decodes, num_decodes + num_extends)
            query_lens_for_extend = query_lens_cpu[num_extends_slice]
            seq_lens_for_extend = seq_lens[num_extends_slice]
            computed_kv_lens = seq_lens_for_extend - query_lens_for_extend
            swa_metadata = None
            if self.aot_sliding_window is not None:
                swa_seqlen_for_extend = torch.minimum(
                    seq_lens_for_extend,
                    query_lens_for_extend + self.aot_sliding_window[0] + 1,
                )
                cu_seq_lens = torch.zeros(
                    num_extends + 1,
                    dtype=torch.int32,
                    device=seq_lens_for_extend.device,
                    pin_memory=True,
                )
                torch.cumsum(
                    swa_seqlen_for_extend,
                    dim=0,
                    dtype=cu_seq_lens.dtype,
                    out=cu_seq_lens[1:],
                )
                token_to_seq = torch.arange(
                    0,
                    num_extends,
                    dtype=torch.int32,
                    device=seq_lens_for_extend.device,
                )
                token_to_seq = torch.repeat_interleave(
                    token_to_seq, swa_seqlen_for_extend
                )
                fetched_shape = cu_seq_lens[-1].item()
                # TODO(ganyi): Maybe reuse these 2 buffer from extend_workspace
                swa_workspace = torch.empty(
                    (2, fetched_shape, self.num_heads_kv, self.headdim),
                    dtype=self.vllm_config.model_config.dtype,
                    device=self.device,
                )

                seq_starts = seq_lens_for_extend - swa_seqlen_for_extend
                max_seqlen_k = swa_seqlen_for_extend.max().item()
                total_tokens = cu_seq_lens[-1].item()

                swa_metadata = AiterChunkSlidingWindowMetadata(
                    swa_cu_seqlens=cu_seq_lens.to(self.device, non_blocking=True),
                    swa_seq_starts=async_tensor_h2d(seq_starts, self.device),
                    swa_token_to_batch=async_tensor_h2d(token_to_seq, self.device),
                    swa_max_seqlens=max_seqlen_k,
                    swa_total_tokens=total_tokens,
                    swa_workspace=swa_workspace,
                )

            # allocate the equal amount of workspace for
            # each chunk prefill request
            max_context_chunk = _CP_TOKENS_PER_ITER_ROCM // num_extends
            num_chunks = cdiv(computed_kv_lens.max().item(), max_context_chunk)

            chunk_starts = (
                torch.arange(num_chunks, dtype=torch.int32)
                .unsqueeze(1)
                .expand(-1, num_extends)
                * max_context_chunk
            )
            chunk_ends = torch.min(
                computed_kv_lens.unsqueeze(0), chunk_starts + max_context_chunk
            )
            chunk_seq_lens = chunk_ends - chunk_starts
            chunk_seq_lens.clamp_(min=0)  # [num_chunks, num_extends]
            cu_seq_lens_cpu = torch.zeros(
                [num_chunks, num_extends + 1], dtype=torch.int32, pin_memory=True
            )
            torch.cumsum(
                chunk_seq_lens, dim=1, out=cu_seq_lens_cpu[:, 1:], dtype=torch.int32
            )
            # Avoid .max() on an empty tensor when there is no context
            # (num_chunks == 0, e.g. Whisper encoder's first pass).
            max_cum_tokens = (
                cu_seq_lens_cpu[:, -1].max().item() if num_chunks > 0 else 0
            )

            range_idx = torch.arange(max_cum_tokens, dtype=torch.int32)[None, None, :]
            idx_to_batch_tensor = range_idx == cu_seq_lens_cpu[:, 1:][:, :, None]
            idx_to_batch_tensor = idx_to_batch_tensor.sum(
                dim=1
            )  # [num_chunks, max_cum_tokens]
            token_to_batch_tensor = torch.cumsum(idx_to_batch_tensor, dim=1)

            chunk_context_metadata = AiterChunkContextMetadata(
                workspace=self.extend_workspace,
                cu_seq_lens_chunk=cu_seq_lens_cpu.to(self.device, non_blocking=True),
                chunk_starts=async_tensor_h2d(chunk_starts, self.device),
                token_to_batch=async_tensor_h2d(token_to_batch_tensor, self.device),
                max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),
                num_chunks=num_chunks,
                total_token_per_batch=cu_seq_lens_cpu[:, -1].tolist(),
                swa_metadata=swa_metadata,
            )

            query_start_loc_device = common_attn_metadata.query_start_loc[
                num_decodes : num_decodes + num_extends + 1
            ]
            seq_lens_device = common_attn_metadata.seq_lens[num_extends_slice]
            cu_seq_lens = torch.zeros(
                num_extends + 1, dtype=torch.int32, device=seq_lens_device.device
            )
            torch.cumsum(
                seq_lens_device, dim=0, dtype=cu_seq_lens.dtype, out=cu_seq_lens[1:]
            )
            extend_metadata = AiterFlashAttentionChunkPrefillMetadata(
                max_query_len=query_lens_for_extend.max().item(),
                max_seq_len=seq_lens[num_extends_slice].max().item(),
                query_start_loc=query_start_loc_device - query_start_loc_device[0],
                chunk_context_metadata=chunk_context_metadata,
            )

        use_cascade = common_prefix_len > 0

        attn_metadata = AiterFlashAttentionMetadata(
            num_actual_tokens=common_attn_metadata.num_actual_tokens,
            query_start_loc=common_attn_metadata.query_start_loc,
            max_seq_len=common_attn_metadata.max_seq_len,
            seq_lens=common_attn_metadata.seq_lens,
            block_table=common_attn_metadata.block_table_tensor,
            causal=common_attn_metadata.causal,
            slot_mapping=common_attn_metadata.slot_mapping,
            num_decodes=num_decodes,
            num_decode_tokens=num_decode_tokens,
            num_prefills=num_prefills,
            num_extends=num_extends,
            num_extend_tokens=num_extend_tokens,
            decode_metadata=decode_metadata,
            prefill_metadata=prefill_metadata,
            extend_metadata=extend_metadata,
            use_cascade=use_cascade,
            k_scale=self.scale,
            v_scale=self.scale,
            kv_sharing_metadata=kv_sharing_metadata,
        )
        return attn_metadata

    def build_for_drafting(
        self,
        common_attn_metadata: CommonAttentionMetadata,
        draft_index: int,
    ) -> AiterFlashAttentionMetadata:
        """Build attention metadata for draft model without CPU-GPU sync.

        During EAGLE drafting all requests are uniform decodes, so we can
        skip split_decodes_prefills_and_extends() and avoid all .cpu() /
        .item() calls that would otherwise break CUDA graph capture.
        """
        num_reqs = common_attn_metadata.num_reqs
        num_tokens = common_attn_metadata.num_actual_tokens

        decode_metadata = AiterFlashAttentionDecodeMetadata(
            max_query_len=common_attn_metadata.max_query_len,
        )

        return AiterFlashAttentionMetadata(
            num_actual_tokens=num_tokens,
            query_start_loc=common_attn_metadata.query_start_loc,
            max_seq_len=common_attn_metadata.max_seq_len,
            seq_lens=common_attn_metadata.seq_lens,
            block_table=common_attn_metadata.block_table_tensor,
            causal=common_attn_metadata.causal,
            slot_mapping=common_attn_metadata.slot_mapping,
            num_decodes=num_reqs,
            num_decode_tokens=num_tokens,
            num_prefills=0,
            num_extends=0,
            num_extend_tokens=0,
            decode_metadata=decode_metadata,
            prefill_metadata=None,
            extend_metadata=None,
            use_cascade=False,
            k_scale=self.scale,
            v_scale=self.scale,
        )

    def use_cascade_attention(self, *args, **kwargs) -> bool:
        return False

build_for_drafting(common_attn_metadata, draft_index)

Build attention metadata for draft model without CPU-GPU sync.

During EAGLE drafting all requests are uniform decodes, so we can skip split_decodes_prefills_and_extends() and avoid all .cpu() / .item() calls that would otherwise break CUDA graph capture.

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

def build_for_drafting(
    self,
    common_attn_metadata: CommonAttentionMetadata,
    draft_index: int,
) -> AiterFlashAttentionMetadata:
    """Build attention metadata for draft model without CPU-GPU sync.

    During EAGLE drafting all requests are uniform decodes, so we can
    skip split_decodes_prefills_and_extends() and avoid all .cpu() /
    .item() calls that would otherwise break CUDA graph capture.
    """
    num_reqs = common_attn_metadata.num_reqs
    num_tokens = common_attn_metadata.num_actual_tokens

    decode_metadata = AiterFlashAttentionDecodeMetadata(
        max_query_len=common_attn_metadata.max_query_len,
    )

    return AiterFlashAttentionMetadata(
        num_actual_tokens=num_tokens,
        query_start_loc=common_attn_metadata.query_start_loc,
        max_seq_len=common_attn_metadata.max_seq_len,
        seq_lens=common_attn_metadata.seq_lens,
        block_table=common_attn_metadata.block_table_tensor,
        causal=common_attn_metadata.causal,
        slot_mapping=common_attn_metadata.slot_mapping,
        num_decodes=num_reqs,
        num_decode_tokens=num_tokens,
        num_prefills=0,
        num_extends=0,
        num_extend_tokens=0,
        decode_metadata=decode_metadata,
        prefill_metadata=None,
        extend_metadata=None,
        use_cascade=False,
        k_scale=self.scale,
        v_scale=self.scale,
    )