vllm.v1.attention.backends.rocm_attn

Attention layer with PagedAttention and Triton prefix prefill.

Classes:

• RocmAttentionBackend –
• RocmAttentionImpl –

RocmAttentionBackend

Bases: AttentionBackend

Methods:

• customize_spec –

  K and V as two head groups so the native HIP kernels address each side

• supports_attn_type –

  ENCODER_DECODER is not supported because

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

class RocmAttentionBackend(AttentionBackend):
    supported_dtypes: ClassVar[list[torch.dtype]] = [
        torch.float16,
        torch.bfloat16,
        torch.float32,
    ]
    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
        "auto",
        "float16",
        "bfloat16",
        "fp8",
        "fp8_e4m3",
        "fp8_e5m2",
    ]

    @staticmethod
    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
        # ROCM paged attention native C++ kernel only supports block sizes 16 and 32
        # due to shared memory (LDS) constraints on AMD GPUs.
        # The (head_size, block_size) pairs that have a native kernel are enumerated
        # by csrc/rocm/attention_paged_attn_<head>_<dtype>_<kv>.cu, each of which
        # instantiates ATTN_LAUNCH from csrc/rocm/attention_kernels.cuh; the runtime
        # gate is use_rocm_custom_paged_attention() in vllm/platforms/rocm.py.
        # However, vLLM allows support for any multiple of 16 via the Triton path.
        # As addressed in PR: https://github.com/vllm-project/vllm/pull/31380,
        # non-standard models (like qwen3-next with block_size 544, or qwen3_5
        # with 784 and 1056) are dynamically routed to our optimized Triton kernel
        # in `do_kv_cache_update`.
        return [MultipleOf(16)]

    @classmethod
    def get_supported_head_sizes(cls) -> list[int]:
        # Head sizes with a native ROCm kernel are only 64 and 128 (see
        # csrc/rocm/attention_paged_attn_h{64,128}_*.cu); the remaining entries are
        # served by the Triton fallback selected in use_rocm_custom_paged_attention().
        return [32, 64, 80, 96, 128, 160, 192, 224, 256]

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

    @classmethod
    def supports_sink(cls) -> bool:
        # ROCM custom attention kernel does not support sinks.
        # Callink this backend with sinks will cause it to fall back to the Triton
        # kernel, which is less efficient than the proper triton backends.
        return False

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

    @classmethod
    def supports_kv_connector(cls) -> bool:
        # ROCM_ATTN uses (2, num_blocks, ...) KV cache layout which is
        # incompatible with KV connectors that require blocks-first layout.
        return False

    forward_includes_kv_cache_update: bool = False

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

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

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

    @classmethod
    def supports_attn_type(cls, attn_type: str) -> bool:
        """ENCODER_DECODER is not supported because
        chunked_prefill_paged_decode's prefill kernel (context_attention_fwd)
        assumes self-attention semantics: it treats passed K/V as new tokens
        to mix with cached K/V. For cross-attention layers the encoder K/V
        are already fully cached, so mixing them again produces incorrect
        results when max_query_len > 1 (e.g. beam search).
        """
        return attn_type in (
            AttentionType.DECODER,
            AttentionType.ENCODER,
            AttentionType.ENCODER_ONLY,
        )

    @classmethod
    def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec:
        """K and V as two head groups so the native HIP kernels address each side
        as one contiguous region (x-packed interior applied in split_kv_cache)."""
        if spec.state_content_bytes is not None:
            return spec
        assert spec.head_size == spec.head_size_v, (
            "Separate K/V head groups 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, ...]:
        # The native HIP kernels hardcode group-contiguous block addressing, so
        # they need the K and V groups to span all blocks (H outermost within the layer)
        # The Triton fallbacks are stride-aware, so block-interior LBHNC is also
        # supported, just without the native kernels.
        return (KVCacheLayout.LHBNC, KVCacheLayout.LBHNC)

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

customize_spec(spec) classmethod

K and V as two head groups so the native HIP kernels address each side as one contiguous region (x-packed interior applied in split_kv_cache).

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

@classmethod
def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec:
    """K and V as two head groups so the native HIP kernels address each side
    as one contiguous region (x-packed interior applied in split_kv_cache)."""
    if spec.state_content_bytes is not None:
        return spec
    assert spec.head_size == spec.head_size_v, (
        "Separate K/V head groups 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 chunked_prefill_paged_decode's prefill kernel (context_attention_fwd) assumes self-attention semantics: it treats passed K/V as new tokens to mix with cached K/V. For cross-attention layers the encoder K/V are already fully cached, so mixing them again produces incorrect results when max_query_len > 1 (e.g. beam search).

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

@classmethod
def supports_attn_type(cls, attn_type: str) -> bool:
    """ENCODER_DECODER is not supported because
    chunked_prefill_paged_decode's prefill kernel (context_attention_fwd)
    assumes self-attention semantics: it treats passed K/V as new tokens
    to mix with cached K/V. For cross-attention layers the encoder K/V
    are already fully cached, so mixing them again produces incorrect
    results when max_query_len > 1 (e.g. beam search).
    """
    return attn_type in (
        AttentionType.DECODER,
        AttentionType.ENCODER,
        AttentionType.ENCODER_ONLY,
    )

RocmAttentionImpl

Bases: AttentionImpl

Methods:

• forward –

  Forward pass with FlashAttention.

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

class RocmAttentionImpl(AttentionImpl):
    def fused_output_quant_supported(self, quant_key: QuantKey):
        return quant_key == kFp8StaticTensorSym

    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.attn_type = attn_type
        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)
        elif attn_type in (AttentionType.ENCODER, AttentionType.ENCODER_ONLY):
            self.sliding_window = (sliding_window - 1, sliding_window - 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
        self.logits_soft_cap = logits_soft_cap
        self.kv_sharing_target_layer_name = kv_sharing_target_layer_name

        self.num_queries_per_kv = self.num_heads // self.num_kv_heads

        self.fp8_dtype = current_platform.fp8_dtype()

        self.sinks = sinks
        if sinks is not None:
            assert sinks.shape[0] == num_heads, (
                "Sinks must have the same number of heads as the number of "
                f"heads in the layer. Sinks shape: {sinks.shape}, "
                f"num_heads: {num_heads}."
            )

        self._use_interleaved_v_cache = False

    def _forward_encoder_attention(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        output: torch.Tensor,
        attn_metadata: RocmAttentionMetadata,
        layer: torch.nn.Module,
    ) -> torch.Tensor:
        """Forward pass for encoder attention without KV cache.

        Args:
            query: shape = [num_encoder_tokens, num_heads, head_size]
            key: shape = [num_encoder_tokens, num_kv_heads, head_size]
            value: shape = [num_encoder_tokens, num_kv_heads, head_size]
            output: shape = [num_encoder_tokens, num_heads, head_size]
            attn_metadata: Encoder attention metadata
            layer: The attention layer

        """
        # For encoder attention, process FP8 quantization if needed
        if is_quantized_kv_cache(self.kv_cache_dtype):
            raise NotImplementedError(
                "quantization is not supported for encoder attention"
            )

        # Use encoder-specific metadata for sequence information
        query_start_loc = attn_metadata.query_start_loc
        seq_lens = attn_metadata.seq_lens
        max_query_len = attn_metadata.max_query_len

        # Call flash attention directly on Q, K, V tensors
        from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd

        context_attention_fwd(
            q=query,
            k=key,
            v=value,
            o=output,
            b_start_loc=query_start_loc,
            b_seq_len=seq_lens,
            max_input_len=max_query_len,
            is_causal=False,
            softmax_scale=self.scale,
            sliding_window_q=self.sliding_window[0],
            sliding_window_k=self.sliding_window[1],
            sinks=self.sinks,
        )
        return output

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

        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: logical [num_blocks, 2, block_size, num_kv_heads *
                head_size] under LHBNC (physically K/V-group-first)
            attn_metadata: Metadata for attention.
            output: Tensor that the attention result is written into.
            output_scale: Scale for fused output quantization.
            output_block_scale: Block scale for fused output quantization;
                not supported by this backend.

        Returns:
            shape = [num_tokens, num_heads * head_size]

        """
        if output_block_scale is not None:
            raise NotImplementedError(
                "fused block_scale output quantization is not yet supported"
                " for RocmAttentionImpl"
            )

        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

        if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
            return self._forward_encoder_attention(
                query[:num_actual_tokens],
                key[:num_actual_tokens],
                value[:num_actual_tokens],
                output[:num_actual_tokens],
                attn_metadata,
                layer,
            )

        # The bound view is logical [B, 2, N, H*hs]; split_kv_cache expects
        # the K/V groups first.
        key_cache, value_cache = PagedAttention.split_kv_cache(
            kv_cache.transpose(0, 1), self.num_kv_heads, self.head_size
        )

        if is_quantized_kv_cache(self.kv_cache_dtype):
            key_cache = key_cache.view(self.fp8_dtype)
            value_cache = value_cache.view(self.fp8_dtype)
            # q_scale only applies to an fp8 query; this path keeps the query
            # in full precision, so a non-1.0 q_scale is not applicable here.
            if query.dtype == self.fp8_dtype and layer._q_scale_float != 1.0:
                raise NotImplementedError(
                    "A non 1.0 q_scale with an fp8 query is not currently "
                    "supported by RocmAttentionImpl."
                )

        cu_seqlens_q = attn_metadata.query_start_loc
        seqused_k = attn_metadata.seq_lens
        max_seqlen_q = attn_metadata.max_query_len
        max_seqlen_k = attn_metadata.max_seq_len
        block_table = attn_metadata.block_table

        # Compute attention and update output up to `num_actual_tokens`.
        chunked_prefill_paged_decode(
            query=query[:num_actual_tokens],
            key=key[:num_actual_tokens] if key is not None else None,
            value=value[:num_actual_tokens] if value is not None else None,
            output=output[:num_actual_tokens],
            kv_cache_dtype=self.kv_cache_dtype,
            key_cache=key_cache,
            value_cache=value_cache,
            block_table=block_table,
            query_start_loc=cu_seqlens_q,
            seq_lens=seqused_k,
            max_seq_len=max_seqlen_k,
            max_query_len=max_seqlen_q,
            k_scale=layer._k_scale,
            v_scale=layer._v_scale,
            alibi_slopes=self.alibi_slopes,
            # self.sliding_window[0] is the FlashAttention-style left span (W - 1).
            # chunked_prefill_paged_decode expects the full window length W.
            sliding_window=1 + self.sliding_window[0],
            sm_scale=self.scale,
            output_scale=output_scale,
            sinks=self.sinks,
            use_interleaved_v_cache=self._use_interleaved_v_cache,
            causal=attn_metadata.causal,
        )

        return output

    def do_kv_cache_update(
        self,
        layer: AttentionLayer,
        key: torch.Tensor,
        value: torch.Tensor,
        kv_cache: torch.Tensor,
        slot_mapping: torch.Tensor,
    ):
        if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
            return
        key_cache, value_cache = PagedAttention.split_kv_cache(
            kv_cache.transpose(0, 1), self.num_kv_heads, self.head_size
        )

        # Reshape the input keys and values and store them in the cache.
        # Get the actual block_size from value_cache
        # value_cache shape: [num_blocks, num_heads, head_size, block_size]
        block_size = value_cache.shape[3]
        has_native_layout = has_native_kv_cache_layout(key_cache, value_cache)

        # Match the fused writer's interleaved V-cache when fusion is active.
        interleaved_v = self._use_interleaved_v_cache
        interleaved_v_pack_factor = (
            (16 // value_cache.element_size()) if interleaved_v else 0
        )

        if block_size in (16, 32) and has_native_layout:
            # Normal 16, 32 with contiguous blocks: use vLLM native HIP C++ logic.
            PagedAttention.write_to_paged_cache(
                key,
                value,
                key_cache,
                value_cache,
                slot_mapping,
                self.kv_cache_dtype,
                layer._k_scale,
                layer._v_scale,
                interleaved_v,
            )
        else:
            # Non-standard blocks and hybrid attention/Mamba layouts need the
            # stride-aware Triton writer. The native reshape_and_cache kernel
            # assumes contiguous block storage and writes to the wrong hybrid
            # cache blocks.
            triton_reshape_and_cache_flash(
                key,
                value,
                key_cache,
                value_cache,
                slot_mapping,
                self.kv_cache_dtype,
                layer._k_scale,
                layer._v_scale,
                interleaved_v_pack_factor,
            )

    def fused_qk_norm_rope_kvcache_supported(self):
        return rocm_aiter_ops.is_enabled()

    def set_interleaved_v_cache(self):
        # Decode/prefill kernels must read the interleaved V-cache the AITER
        # fused write emits when use_shuffle_layout=True.
        self._use_interleaved_v_cache = True

    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,
    ):
        # kv_cache is logical [num_blocks, 2, block_size, num_kv_heads * head_size]
        # (see forward); the AITER shuffle writer takes one (B, N, H, hs) per side.
        key_cache, value_cache = kv_cache.unflatten(
            -1, (self.num_kv_heads, self.head_size)
        ).unbind(1)
        use_shuffle_layout = self._use_interleaved_v_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=use_shuffle_layout,
            return_kv=True,
        )

    def fused_rope_kvcache_supported(self):
        return rocm_aiter_ops.is_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,
    ):
        if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
            return
        if self._use_interleaved_v_cache:
            # The Triton RoPE + cache writer only stores the regular V layout.
            # With layer_name hoisted, one RoPE + KV-cache pattern serves every
            # layer, so a layer switched to the interleaved V cache is handled
            # here at run time rather than left out of the fusion.
            ops.rotary_embedding(
                positions,
                query,
                key,
                layer.head_size,  # type: ignore[attr-defined]
                cos_sin_cache,
                is_neox,
            )
            self.do_kv_cache_update(layer, key, value, kv_cache, layer_slot_mapping)
            return
        key_cache, value_cache = PagedAttention.split_kv_cache(
            kv_cache.transpose(0, 1),
            layer.num_kv_heads,  # type: ignore[attr-defined]
            layer.head_size,  # type: ignore[attr-defined]
        )
        flash_layout = False

        is_fp8_kv_cache = is_quantized_kv_cache(self.kv_cache_dtype)
        if is_fp8_kv_cache:
            key_cache = key_cache.view(self.fp8_dtype)
            value_cache = value_cache.view(self.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_encoder_attention(query, key, value, output, attn_metadata, layer)

Forward pass for encoder attention without KV cache.

Parameters:

• query

  (Tensor) –

  shape = [num_encoder_tokens, num_heads, head_size]

• key

  (Tensor) –

  shape = [num_encoder_tokens, num_kv_heads, head_size]

• value

  (Tensor) –

  shape = [num_encoder_tokens, num_kv_heads, head_size]

• output

  (Tensor) –

  shape = [num_encoder_tokens, num_heads, head_size]

• attn_metadata

  (RocmAttentionMetadata) –

  Encoder attention metadata

• layer

  (Module) –

  The attention layer

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

def _forward_encoder_attention(
    self,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    output: torch.Tensor,
    attn_metadata: RocmAttentionMetadata,
    layer: torch.nn.Module,
) -> torch.Tensor:
    """Forward pass for encoder attention without KV cache.

    Args:
        query: shape = [num_encoder_tokens, num_heads, head_size]
        key: shape = [num_encoder_tokens, num_kv_heads, head_size]
        value: shape = [num_encoder_tokens, num_kv_heads, head_size]
        output: shape = [num_encoder_tokens, num_heads, head_size]
        attn_metadata: Encoder attention metadata
        layer: The attention layer

    """
    # For encoder attention, process FP8 quantization if needed
    if is_quantized_kv_cache(self.kv_cache_dtype):
        raise NotImplementedError(
            "quantization is not supported for encoder attention"
        )

    # Use encoder-specific metadata for sequence information
    query_start_loc = attn_metadata.query_start_loc
    seq_lens = attn_metadata.seq_lens
    max_query_len = attn_metadata.max_query_len

    # Call flash attention directly on Q, K, V tensors
    from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd

    context_attention_fwd(
        q=query,
        k=key,
        v=value,
        o=output,
        b_start_loc=query_start_loc,
        b_seq_len=seq_lens,
        max_input_len=max_query_len,
        is_causal=False,
        softmax_scale=self.scale,
        sliding_window_q=self.sliding_window[0],
        sliding_window_k=self.sliding_window[1],
        sinks=self.sinks,
    )
    return output

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

Forward pass with FlashAttention.

Parameters:

• layer

  (Module) –

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

• query

  (Tensor) –

  shape = [num_tokens, num_heads, head_size]

• key

  (Tensor) –

  shape = [num_tokens, num_kv_heads, head_size]

• value

  (Tensor) –

  shape = [num_tokens, num_kv_heads, head_size]

• kv_cache

  (Tensor) –

  logical [num_blocks, 2, block_size, num_kv_heads * head_size] under LHBNC (physically K/V-group-first)

• attn_metadata

  (RocmAttentionMetadata) –

  Metadata for attention.

• output

  (Tensor) –

  Tensor that the attention result is written into.

• output_scale

  (Tensor | None, default: None ) –

  Scale for fused output quantization.

• 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]

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

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

    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: logical [num_blocks, 2, block_size, num_kv_heads *
            head_size] under LHBNC (physically K/V-group-first)
        attn_metadata: Metadata for attention.
        output: Tensor that the attention result is written into.
        output_scale: Scale for fused output quantization.
        output_block_scale: Block scale for fused output quantization;
            not supported by this backend.

    Returns:
        shape = [num_tokens, num_heads * head_size]

    """
    if output_block_scale is not None:
        raise NotImplementedError(
            "fused block_scale output quantization is not yet supported"
            " for RocmAttentionImpl"
        )

    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

    if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
        return self._forward_encoder_attention(
            query[:num_actual_tokens],
            key[:num_actual_tokens],
            value[:num_actual_tokens],
            output[:num_actual_tokens],
            attn_metadata,
            layer,
        )

    # The bound view is logical [B, 2, N, H*hs]; split_kv_cache expects
    # the K/V groups first.
    key_cache, value_cache = PagedAttention.split_kv_cache(
        kv_cache.transpose(0, 1), self.num_kv_heads, self.head_size
    )

    if is_quantized_kv_cache(self.kv_cache_dtype):
        key_cache = key_cache.view(self.fp8_dtype)
        value_cache = value_cache.view(self.fp8_dtype)
        # q_scale only applies to an fp8 query; this path keeps the query
        # in full precision, so a non-1.0 q_scale is not applicable here.
        if query.dtype == self.fp8_dtype and layer._q_scale_float != 1.0:
            raise NotImplementedError(
                "A non 1.0 q_scale with an fp8 query is not currently "
                "supported by RocmAttentionImpl."
            )

    cu_seqlens_q = attn_metadata.query_start_loc
    seqused_k = attn_metadata.seq_lens
    max_seqlen_q = attn_metadata.max_query_len
    max_seqlen_k = attn_metadata.max_seq_len
    block_table = attn_metadata.block_table

    # Compute attention and update output up to `num_actual_tokens`.
    chunked_prefill_paged_decode(
        query=query[:num_actual_tokens],
        key=key[:num_actual_tokens] if key is not None else None,
        value=value[:num_actual_tokens] if value is not None else None,
        output=output[:num_actual_tokens],
        kv_cache_dtype=self.kv_cache_dtype,
        key_cache=key_cache,
        value_cache=value_cache,
        block_table=block_table,
        query_start_loc=cu_seqlens_q,
        seq_lens=seqused_k,
        max_seq_len=max_seqlen_k,
        max_query_len=max_seqlen_q,
        k_scale=layer._k_scale,
        v_scale=layer._v_scale,
        alibi_slopes=self.alibi_slopes,
        # self.sliding_window[0] is the FlashAttention-style left span (W - 1).
        # chunked_prefill_paged_decode expects the full window length W.
        sliding_window=1 + self.sliding_window[0],
        sm_scale=self.scale,
        output_scale=output_scale,
        sinks=self.sinks,
        use_interleaved_v_cache=self._use_interleaved_v_cache,
        causal=attn_metadata.causal,
    )

    return output