vLLM/vllm/models/deepseek_v4/amd/rocm.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import functools
from dataclasses import dataclass
from typing import cast

import torch

from vllm import envs
from vllm.config import VllmConfig
from vllm.distributed import (
    get_tensor_model_parallel_world_size,
    tensor_model_parallel_all_reduce,
)
from vllm.forward_context import get_forward_context
from vllm.logger import init_logger
from vllm.models.deepseek_v4.attention import DeepseekV4Attention
from vllm.models.deepseek_v4.common.ops import dequantize_and_gather_k_cache
from vllm.models.deepseek_v4.sparse_mla import (
    DeepseekV4FlashMLAMetadata,
    DeepseekV4SparseMLABackend,
    DeepseekV4SparseMLAMetadataBuilder,
)
from vllm.platforms import current_platform
from vllm.platforms.rocm import _ON_GFX950, on_gfx1250
from vllm.triton_utils import tl, triton
from vllm.utils.multi_stream_utils import execute_in_parallel
from vllm.v1.attention.backend import (
    AttentionCGSupport,
    CommonAttentionMetadata,
)
from vllm.v1.attention.backends.mla.sparse_swa import (
    DeepseekSparseSWAMetadata,
    DeepseekSparseSWAMetadataBuilder,
)
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
    build_prefill_topk_ragged_indices,
    build_ragged_indices_from_dense,
    rocm_inv_rope_einsum,
    rocm_inverse_rope_rows_,
    rocm_sparse_attn_decode,
    rocm_sparse_attn_prefill,
)
from vllm.v1.kv_cache_interface import KVCacheSpec
from vllm.v1.worker.workspace import current_workspace_manager

logger = init_logger(__name__)


def _wo_a_block_scale_to_e8m0(scale: torch.Tensor) -> torch.Tensor | None:
    """Normalize checkpoint WO_A scales to raw OCP MX E8M0 bytes.

    E8M0 is an unsigned exponent-only scale format with bias 127. A finite
    encoded byte ``b`` in ``[0, 254]`` represents ``2 ** (b - 127)``;
    ``0xFF`` is reserved for NaN. This is the vendor-neutral OCP encoding used
    by AMD AITER/OPUS, not an NVIDIA-specific convention. Reference: OCP
    Microscaling Formats (MX) Specification, section 5.4.1:
    https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf

    Loaders may preserve the encoded byte as ``float8_e8m0fnu``/``uint8``,
    or decode it to a floating-point power of two. Floating-point inputs are
    accepted only when the original E8M0 byte can be recovered losslessly.
    This function does not quantize or round arbitrary scales.
    """
    if scale.dtype == torch.float8_e8m0fnu:
        # Reinterpret the native E8M0 storage. ``to(uint8)`` would perform a
        # numeric conversion instead of preserving the encoded exponent byte.
        return scale.view(torch.uint8).contiguous()
    if scale.dtype == torch.uint8:
        # The checkpoint loader already exposed the E8M0 wire representation.
        return scale.contiguous()
    if not scale.dtype.is_floating_point:
        return None

    scale_f32 = scale.detach().float()
    if not bool(torch.isfinite(scale_f32).all()) or bool((scale_f32 <= 0).any()):
        return None

    # With no mantissa, E8M0 can represent only exact powers of two. Rebuild
    # the value before encoding so this adapter never silently quantizes a
    # general floating-point checkpoint scale.
    exponent = torch.round(torch.log2(scale_f32))
    if not torch.equal(torch.exp2(exponent), scale_f32):
        return None

    encoded = exponent.to(torch.int32) + 127
    # 0xFF is NaN in OCP E8M0, not a finite exponent.
    if int(encoded.min()) < 0 or int(encoded.max()) > 254:
        return None
    return encoded.to(torch.uint8).contiguous()


def _trust_dsv4_extra_cache_nan_free(
    kv_cache_dtype: str,
    has_kv_transfer: bool,
    has_extra_cache: bool,
) -> bool:
    return (
        _ON_GFX950
        and kv_cache_dtype == "fp8_ds_mla"
        and not has_kv_transfer
        and has_extra_cache
    )


def _aiter_sparse_mla_enabled(has_kv_transfer: bool) -> bool:
    from vllm._aiter_ops import rocm_aiter_ops

    if not rocm_aiter_ops.is_triton_sparse_mla_enabled():
        return False
    # The aiter kernel does not scrub NaNs, so like the trust check above it
    # needs a cache only the canonical writer has touched.
    if has_kv_transfer:
        logger.warning_once(
            "VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA is ignored with a KV "
            "connector; DeepSeek V4 keeps the Triton sparse attention."
        )
        return False
    return True


def _build_indptr_from_lengths(lengths: torch.Tensor) -> torch.Tensor:
    lengths = lengths.to(dtype=torch.int32).contiguous()
    indptr = torch.zeros(lengths.shape[0] + 1, dtype=torch.int32, device=lengths.device)
    torch.cumsum(lengths, dim=0, out=indptr[1:])
    return indptr


def apply_pre_quantized_block_scaled_mm(
    linear: torch.nn.Module,
    x_fp8: torch.Tensor,
    x_scale: torch.Tensor,
) -> torch.Tensor:
    """Block-scaled fp8 GEMM on pre-quantized activations.

    The fused q/kv norm kernel writes fp8 qr + per-1x128 scales; this
    drives the linear's block-scaled GEMM directly with them, bypassing
    apply_weights which would re-quantize the fp8 input. Only valid for
    the wq_b-style column/replicated linears: their output is the local
    TP shard, so no all-reduce is needed.
    """
    from vllm.model_executor.kernels.linear.scaled_mm.BlockScaledMMLinearKernel import (
        FP8BlockParams,
    )

    params = FP8BlockParams.from_layer(linear)
    weight_scale = (
        params.weight_scale
        if params.weight_scale_inv is None
        else params.weight_scale_inv
    )
    kernel = linear.quant_method.fp8_linear
    out = kernel.apply_block_scaled_mm(
        A=x_fp8, B=params.weight, As=x_scale, Bs=weight_scale
    )
    return out.to(dtype=kernel.config.out_dtype)


# ROCm sparse prefill keeps this dense combine local so AMD-specific SWA changes
# do not touch the shared DeepSeek V4 cache utilities.
_SPARSE_PREFILL_TOPK_ALIGNMENT = 128


def _default_first_token_threshold() -> int:
    """Max tokens at which the CSA/HCA fork enqueues the default chain first.

    At small batches starting the long default chain first shortens the step
    tail; larger batches saturate the device and aux-first is faster. Only
    tuned on gfx950 (24 tokens = conc 8 with MTP=2); 0 disables elsewhere.
    """
    return 24 if _ON_GFX950 else 0


@triton.jit
def _combine_topk_swa_indices_kernel(
    combined_indices_ptr,
    combined_indices_stride,
    combined_lens_ptr,
    topk_indices_ptr,
    topk_indices_stride,
    query_start_loc_ptr,
    seq_lens_ptr,
    gather_lens_ptr,
    left_visible_ptr,
    right_visible_ptr,
    M,
    N,
    TOP_K: tl.constexpr,
    COMPRESS_RATIO: tl.constexpr,
    WINDOW_SIZE: tl.constexpr,
    SWA_WIDTH: tl.constexpr,
    TOPK_WIDTH: tl.constexpr,
    PADDED_TOP_K: tl.constexpr,
    PADDED_SWA_WIDTH: tl.constexpr,
    HAS_IMAGE: tl.constexpr,
):
    batch_idx = tl.program_id(0)
    worker_id = tl.program_id(1)
    num_workers = tl.num_programs(1)

    base = tl.load(query_start_loc_ptr)
    query_start = tl.load(query_start_loc_ptr + batch_idx) - base
    query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base
    query_len = query_end - query_start
    seq_len = tl.load(seq_lens_ptr + batch_idx)
    gather_len = tl.load(gather_lens_ptr + batch_idx)
    start_pos = seq_len - query_len
    gather_start = seq_len - gather_len

    for token_idx in range(query_start + worker_id, query_end, num_workers):
        token_idx_in_query = token_idx - query_start
        pos = start_pos + token_idx_in_query
        topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
        if HAS_IMAGE:
            left = tl.load(left_visible_ptr + token_idx)
            right = tl.load(right_visible_ptr + token_idx)
        else:
            left = 0
            right = 0
        left_add = tl.maximum(left - (WINDOW_SIZE - 1), 0)
        # Prefix caching can resume inside an image span. Do not generate
        # indices outside the SWA rows present in the gathered workspace.
        swa_start = tl.maximum(
            tl.maximum(pos - (WINDOW_SIZE - 1) - left_add, 0), gather_start
        )
        swa_end = tl.minimum(pos + right + 1, seq_len)
        swa_len = tl.maximum(swa_end - swa_start, 0)

        topk_offset = tl.arange(0, PADDED_TOP_K)
        topk_mask = topk_offset < topk_len
        safe_topk_offset = tl.where(topk_offset < TOPK_WIDTH, topk_offset, 0)
        topk_indices = tl.load(
            topk_indices_ptr + token_idx * topk_indices_stride + safe_topk_offset,
            mask=topk_mask,
            other=-1,
        )
        valid_topk = (topk_indices >= 0) & (topk_indices < N)
        topk_indices = tl.where(valid_topk, topk_indices + M * batch_idx, -1)
        tl.store(
            combined_indices_ptr + token_idx * combined_indices_stride + topk_offset,
            topk_indices,
            mask=topk_mask,
        )

        swa_offset = tl.arange(0, PADDED_SWA_WIDTH)
        tl.store(
            combined_indices_ptr
            + token_idx * combined_indices_stride
            + topk_len
            + swa_offset,
            M * batch_idx + N + swa_offset + swa_start - gather_start,
            mask=(swa_offset < swa_len) & (swa_offset < SWA_WIDTH),
        )

        tl.store(combined_lens_ptr + token_idx, topk_len + swa_len)


def combine_topk_swa_indices(
    topk_indices: torch.Tensor,
    query_start_loc: torch.Tensor,
    seq_lens: torch.Tensor,
    gather_lens: torch.Tensor,
    window_size: int,
    compress_ratio: int,
    topk: int,
    M: int,
    N: int,
    max_image_tokens: int = 0,
    left_visible: torch.Tensor | None = None,
    right_visible: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    if (left_visible is None) != (right_visible is None):
        raise ValueError("left_visible and right_visible must be provided together")
    topk_indices = topk_indices.reshape(topk_indices.shape[0], -1).contiguous()
    num_tokens = topk_indices.shape[0]
    num_reqs = seq_lens.shape[0]
    has_image = left_visible is not None
    # Keep the row shape fixed for a vision model even when a particular batch
    # has no image.
    swa_width = window_size + max_image_tokens
    combined_topk = (
        (topk + swa_width + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1)
        // _SPARSE_PREFILL_TOPK_ALIGNMENT
        * _SPARSE_PREFILL_TOPK_ALIGNMENT
    )
    combined_indices = torch.full(
        (num_tokens, combined_topk),
        fill_value=-1,
        dtype=torch.int32,
        device=topk_indices.device,
    )
    combined_lens = torch.empty(
        num_tokens, dtype=torch.int32, device=topk_indices.device
    )

    num_workers = 128
    _combine_topk_swa_indices_kernel[(num_reqs, num_workers)](
        combined_indices,
        combined_indices.stride(0),
        combined_lens,
        topk_indices,
        topk_indices.stride(0),
        query_start_loc,
        seq_lens,
        gather_lens,
        left_visible if left_visible is not None else topk_indices,
        right_visible if right_visible is not None else topk_indices,
        M,
        N,
        TOP_K=topk,
        COMPRESS_RATIO=compress_ratio,
        WINDOW_SIZE=window_size,
        SWA_WIDTH=swa_width,
        TOPK_WIDTH=topk_indices.shape[-1],
        PADDED_TOP_K=triton.next_power_of_2(topk_indices.shape[-1]),
        PADDED_SWA_WIDTH=triton.next_power_of_2(swa_width),
        HAS_IMAGE=has_image,
    )
    return combined_indices, combined_lens


@triton.jit
def _compute_topk_lens_kernel(
    topk_lens_ptr,
    topk_indices_ptr,
    topk_indices_stride,
    topk,
    is_valid_token_ptr,
    TRITON_BLOCK_SIZE: tl.constexpr,
):
    token_idx = tl.program_id(0)
    is_valid_token = tl.load(is_valid_token_ptr + token_idx)

    count = tl.zeros((), dtype=tl.int32)
    for i in range(0, topk, TRITON_BLOCK_SIZE):
        offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
        mask = offset < topk
        local_idx = tl.load(
            topk_indices_ptr + token_idx * topk_indices_stride + offset,
            mask=mask,
            other=-1,
        )
        count += tl.sum((local_idx >= 0).to(tl.int32), axis=0)

    tl.store(topk_lens_ptr + token_idx, tl.where(is_valid_token, count, 0))


@triton.jit
def _pack_global_topk_ragged_kernel(
    global_topk_ragged_ptr,
    topk_indptr_ptr,
    topk_indices_ptr,
    topk_indices_stride,
    token_to_req_indices_ptr,
    block_table_ptr,
    block_table_stride,
    block_size,
    topk,
    BLOCK_SIZE: tl.constexpr,
):
    token_idx = tl.program_id(0)
    block_idx = tl.program_id(1)
    offset = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)

    out_start = tl.load(topk_indptr_ptr + token_idx)
    out_end = tl.load(topk_indptr_ptr + token_idx + 1)
    out_len = out_end - out_start
    if block_idx * BLOCK_SIZE >= out_len:
        return

    req_idx = tl.load(token_to_req_indices_ptr + token_idx)
    mask = (offset < out_len) & (offset < topk)
    local_idx = tl.load(
        topk_indices_ptr + token_idx * topk_indices_stride + offset,
        mask=mask,
        other=-1,
    )
    valid = mask & (local_idx >= 0)
    block_indices = local_idx // block_size
    block_numbers = tl.load(
        block_table_ptr + req_idx * block_table_stride + block_indices,
        mask=valid,
        other=0,
    )
    block_offsets = local_idx % block_size
    slot_ids = tl.where(valid, block_numbers * block_size + block_offsets, -1)
    tl.store(global_topk_ragged_ptr + out_start + offset, slot_ids, mask=mask)


def compute_global_topk_ragged_indices_and_indptr(
    topk_indices: torch.Tensor,
    token_to_req_indices: torch.Tensor,
    block_table: torch.Tensor,
    block_size: int,
    is_valid_token: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    topk_indices = topk_indices.reshape(topk_indices.shape[0], -1).contiguous()
    num_tokens = topk_indices.shape[0]
    topk = topk_indices.shape[1]

    topk_lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device)
    _compute_topk_lens_kernel[(num_tokens,)](
        topk_lens,
        topk_indices,
        topk_indices.stride(0),
        topk,
        is_valid_token,
        TRITON_BLOCK_SIZE=1024,
    )

    topk_indptr = _build_indptr_from_lengths(topk_lens)
    global_topk_ragged = torch.empty(
        num_tokens * topk,
        dtype=torch.int32,
        device=topk_indices.device,
    )
    if global_topk_ragged.numel() > 0:
        block = 128
        _pack_global_topk_ragged_kernel[(num_tokens, triton.cdiv(topk, block))](
            global_topk_ragged,
            topk_indptr,
            topk_indices,
            topk_indices.stride(0),
            token_to_req_indices,
            block_table,
            block_table.stride(0),
            block_size,
            topk,
            BLOCK_SIZE=block,
        )
    return global_topk_ragged, topk_indptr, topk_lens


def _copy_ragged_to_graph_buffers(
    ragged_indices: torch.Tensor,
    ragged_indptr: torch.Tensor,
    ragged_indices_buffer: torch.Tensor,
    ragged_indptr_buffer: torch.Tensor,
    num_rows: int,
    max_entries_per_row: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Copy dynamic ragged metadata into persistent CUDA graph buffers.

    FULL decode graphs capture kernel argument addresses. Keep the returned
    tensors backed by stable storage, while indptr continues to bound reads.
    """
    indptr_out = ragged_indptr_buffer[: num_rows + 1]
    indptr_out.copy_(ragged_indptr, non_blocking=True)

    max_entries = max(num_rows * max_entries_per_row, 1)
    ragged_out = ragged_indices_buffer[:max_entries]
    source_entries = ragged_indices.numel()
    if source_entries > 0:
        ragged_out[:source_entries].copy_(ragged_indices, non_blocking=True)
    if _ON_GFX950:
        # Preserve the graph-stable base pointer while exposing source capacity
        # to the sync-free split selector; indptr still carries the true NNZ.
        ragged_out = ragged_out[: max(source_entries, 1)]
    return ragged_out, indptr_out


@dataclass
class DeepseekV4ROCMAiterMLASparseMetadata(DeepseekV4FlashMLAMetadata):
    """ROCm-specific DeepSeek V4 metadata carrying ragged decode topk."""

    c128a_decode_topk_ragged_indices: torch.Tensor | None = None
    c128a_decode_topk_ragged_indptr: torch.Tensor | None = None
    for_cudagraph_capture: bool = False


@dataclass
class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata):
    decode_swa_ragged_indices: torch.Tensor | None = None
    decode_swa_ragged_indptr: torch.Tensor | None = None
    # Built only for the aiter prefill, which reads the SWA cache in place.
    prefill_swa_ragged_indices: torch.Tensor | None = None
    prefill_swa_ragged_indptr: torch.Tensor | None = None


class DeepseekV4ROCMAiterMLASparseMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder):
    @classmethod
    def get_cudagraph_support(
        cls,
        vllm_config: VllmConfig,
        kv_cache_spec: KVCacheSpec,
    ) -> AttentionCGSupport:
        spec_config = vllm_config.speculative_config
        if spec_config is not None and spec_config.enable_adaptive_verification:
            # All per-token metadata is built from device query boundaries into
            # persistent buffers, so adaptive verification can replay varlen
            # FULL decode graphs after reallocating drafts across requests.
            return AttentionCGSupport.ALWAYS
        return super().get_cudagraph_support(vllm_config, kv_cache_spec)

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.c128a_decode_topk_ragged_indices_buffer: torch.Tensor | None = None
        self.c128a_decode_topk_ragged_indptr_buffer: torch.Tensor | None = None
        if self.compress_ratio == 128:
            max_tokens = self.vllm_config.scheduler_config.max_num_batched_tokens
            self.c128a_decode_topk_ragged_indices_buffer = torch.empty(
                max_tokens * self.c128a_max_compressed,
                dtype=torch.int32,
                device=self.device,
            )
            self.c128a_decode_topk_ragged_indptr_buffer = torch.empty(
                max_tokens + 1,
                dtype=torch.int32,
                device=self.device,
            )

    def build(
        self,
        common_prefix_len: int,
        common_attn_metadata: CommonAttentionMetadata,
        fast_build: bool = False,
    ) -> DeepseekV4ROCMAiterMLASparseMetadata:
        base = super().build(
            common_prefix_len=common_prefix_len,
            common_attn_metadata=common_attn_metadata,
            fast_build=fast_build,
        )

        ragged_indices = None
        ragged_indptr = None
        dense_decode = base.c128a_global_decode_topk_indices
        decode_lens = base.c128a_decode_topk_lens
        if dense_decode is not None and decode_lens is not None:
            ragged_indices, ragged_indptr = build_ragged_indices_from_dense(
                dense_decode.reshape(dense_decode.shape[0], -1),
                decode_lens,
            )
            assert self.c128a_decode_topk_ragged_indices_buffer is not None
            assert self.c128a_decode_topk_ragged_indptr_buffer is not None
            ragged_indices, ragged_indptr = _copy_ragged_to_graph_buffers(
                ragged_indices,
                ragged_indptr,
                self.c128a_decode_topk_ragged_indices_buffer,
                self.c128a_decode_topk_ragged_indptr_buffer,
                dense_decode.shape[0],
                self.c128a_max_compressed,
            )

        return DeepseekV4ROCMAiterMLASparseMetadata(
            **vars(base),
            c128a_decode_topk_ragged_indices=ragged_indices,
            c128a_decode_topk_ragged_indptr=ragged_indptr,
        )

    def build_for_cudagraph_capture(
        self, common_attn_metadata: CommonAttentionMetadata
    ) -> DeepseekV4ROCMAiterMLASparseMetadata:
        metadata = cast(
            DeepseekV4ROCMAiterMLASparseMetadata,
            super().build_for_cudagraph_capture(common_attn_metadata),
        )
        metadata.for_cudagraph_capture = _ON_GFX950
        return metadata


class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuilder):
    @classmethod
    def get_cudagraph_support(
        cls,
        vllm_config: VllmConfig,
        kv_cache_spec: KVCacheSpec,
    ) -> AttentionCGSupport:
        spec_config = vllm_config.speculative_config
        if spec_config is not None and spec_config.enable_adaptive_verification:
            # SWA indices, lengths, and token-to-request mappings are built from
            # device boundaries into persistent buffers, so adaptive verification
            # can replay varlen FULL decode graphs safely.
            return AttentionCGSupport.ALWAYS
        return super().get_cudagraph_support(vllm_config, kv_cache_spec)

    # Keep fused multi-step decode disabled until update_draft_decode_metadata()
    # also refreshes the ROCm-specific ragged SWA indices and indptrs.
    supports_draft_decode_metadata_update = False

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        max_tokens = self.vllm_config.scheduler_config.max_num_batched_tokens
        # The non-causal (DSpark draft) path widens each token's SWA index list
        # to ``noncausal_index_width`` (>= window_size), so size the persistent
        # ragged buffer to the wider bound to cover both causal and non-causal.
        swa_index_width = max(self.window_size, self.noncausal_index_width)
        self.decode_swa_ragged_indices_buffer = torch.empty(
            max_tokens * swa_index_width,
            dtype=torch.int32,
            device=self.device,
        )
        self.decode_swa_ragged_indptr_buffer = torch.empty(
            max_tokens + 1,
            dtype=torch.int32,
            device=self.device,
        )
        self._use_aiter_sparse_mla = _aiter_sparse_mla_enabled(
            self.vllm_config.kv_transfer_config is not None
        )

    def build(
        self,
        common_prefix_len: int,
        common_attn_metadata: CommonAttentionMetadata,
        fast_build: bool = False,
        replay_start: torch.Tensor | None = None,
    ) -> DeepseekV4ROCMAiterSparseSWAMetadata:
        base = super().build(
            common_prefix_len=common_prefix_len,
            common_attn_metadata=common_attn_metadata,
            fast_build=fast_build,
            replay_start=replay_start,
        )

        prefill_ragged_indices = None
        prefill_ragged_indptr = None
        if (
            self._use_aiter_sparse_mla
            and base.num_prefill_tokens > 0
            and base.prefill_swa_indices is not None
            and base.prefill_swa_lens is not None
        ):
            # Paged-coordinate rows (image visibility included), so every
            # layer's prefill shares one build.
            prefill_ragged_indices, prefill_ragged_indptr = (
                build_ragged_indices_from_dense(
                    base.prefill_swa_indices.reshape(base.num_prefill_tokens, -1),
                    base.prefill_swa_lens,
                )
            )

        ragged_indices = None
        ragged_indptr = None
        if (
            base.num_decode_tokens > 0
            and base.decode_swa_indices is not None
            and base.decode_swa_lens is not None
        ):
            ragged_indices, ragged_indptr = build_ragged_indices_from_dense(
                base.decode_swa_indices.reshape(
                    base.num_decode_tokens, base.decode_swa_width
                ),
                base.decode_swa_lens,
            )
            ragged_indices, ragged_indptr = _copy_ragged_to_graph_buffers(
                ragged_indices,
                ragged_indptr,
                self.decode_swa_ragged_indices_buffer,
                self.decode_swa_ragged_indptr_buffer,
                base.num_decode_tokens,
                base.decode_swa_width,
            )

        return DeepseekV4ROCMAiterSparseSWAMetadata(
            **vars(base),
            decode_swa_ragged_indices=ragged_indices,
            decode_swa_ragged_indptr=ragged_indptr,
            prefill_swa_ragged_indices=prefill_ragged_indices,
            prefill_swa_ragged_indptr=prefill_ragged_indptr,
        )


class DeepseekV4ROCMAiterMLASparseBackend(DeepseekV4SparseMLABackend):
    @staticmethod
    def get_name() -> str:
        return "ROCM_FLASHMLA_SPARSE_DSV4"

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


class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
    """ROCm sparse MLA attention layer for DeepSeek V4."""

    backend_cls = DeepseekV4ROCMAiterMLASparseBackend
    _use_aiter_sparse_mla = False

    def __init__(self, *args, **kwargs):
        vllm_config = args[0] if args else kwargs["vllm_config"]
        super().__init__(*args, **kwargs)
        self._has_kv_transfer = vllm_config.kv_transfer_config is not None
        # Block scale for the preshuffled weight; None = not preshuffled.
        self._wqa_wkv_scale: torch.Tensor | None = None
        self._wo_b_scale: torch.Tensor | None = None
        self._wo_a_fp8_weight: torch.Tensor | None = None
        self._wo_a_e8m0_scale: torch.Tensor | None = None
        self._wo_a_cos_cache: torch.Tensor | None = None
        self._wo_a_sin_cache: torch.Tensor | None = None
        self._fused_compressor_weight: torch.Tensor | None
        self.register_buffer("_fused_compressor_weight", None, persistent=False)
        self._fused_compressor_split_sizes: tuple[int, int] | None = None

        if self.indexer is None:
            # Dense layers have no compressor work to overlap; HCA layers
            # (compressor, no indexer) keep the streams for the dual-stream
            # fork below.
            if self.compressor is None:
                self.aux_stream_list = None
        else:
            # Disable indexer inner overlap.
            self.indexer.aux_stream = None
        self._use_aiter_sparse_mla = _aiter_sparse_mla_enabled(self._has_kv_transfer)

    def _run_sequential_pipeline(
        self,
        hidden_states: torch.Tensor,
        positions: torch.Tensor,
        o_padded: torch.Tensor,
    ) -> None:
        """Disable ROCm streams when the current execution region cannot overlap."""
        aux_streams = self.aux_stream_list
        self.aux_stream_list = None
        try:
            qr_kv, kv_score, indexer_kv_score, indexer_weights = (
                self._run_parallel_input_projections(hidden_states)
            )
            qr, qr_scale, kv = self._split_qkv_and_norm(qr_kv)
            self._prepare_and_attn_fn(
                hidden_states,
                qr,
                kv,
                qr_scale,
                kv_score,
                indexer_kv_score,
                indexer_weights,
                positions,
                o_padded,
            )
        finally:
            self.aux_stream_list = aux_streams

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        llama_4_scaling: torch.Tensor | None = None,
    ) -> torch.Tensor:
        # Pre-allocate attention output with FlashMLA-padded head count.
        # The op writes into `o_padded`; we slice to n_local_heads after.
        num_tokens = hidden_states.shape[0]
        o_padded = torch.empty(
            (num_tokens, self.padded_heads, self.head_dim),
            dtype=hidden_states.dtype,
            device=hidden_states.device,
        )

        if current_platform.enable_multi_stream_overlap(
            self.aux_stream_list, get_forward_context().attn_metadata
        ):
            # The ROCm override consumes these sentinels inside the capture
            # boundary, moving the stream fan-out ahead of the projections.
            self._prepare_and_attn_fn(
                hidden_states,
                None,
                None,
                None,
                None,
                None,
                None,
                positions,
                o_padded,
            )
        else:
            self._run_sequential_pipeline(hidden_states, positions, o_padded)

        o = o_padded[:, : self.n_local_heads, :]

        # Inverse-RoPE + wo_a + wo_b output projection (platform-specific).
        return self._o_proj(o, positions)

    def _prepare_and_attn(
        self,
        hidden_states: torch.Tensor,
        qr: torch.Tensor | None,
        kv: torch.Tensor | None,
        qr_scale: torch.Tensor | None,
        kv_score: torch.Tensor | None,
        indexer_kv_score: torch.Tensor | None,
        indexer_weights: torch.Tensor | None,
        positions: torch.Tensor,
        o_padded: torch.Tensor,
    ) -> None:
        """Run the ROCm fork/join (HCA or CSA) inside the capture boundary."""
        aux_streams = self.aux_stream_list
        # The sequential pipeline disables aux_stream_list before calling
        # back with real projection inputs; aux_streams is None ends that
        # recursion here.
        if aux_streams is None:
            saved_streams = self.aux_stream_list
            self.aux_stream_list = None
            try:
                super()._prepare_and_attn(
                    hidden_states,
                    cast(torch.Tensor, qr),
                    cast(torch.Tensor, kv),
                    qr_scale,
                    cast(torch.Tensor, kv_score),
                    cast(torch.Tensor, indexer_kv_score),
                    cast(torch.Tensor, indexer_weights),
                    positions,
                    o_padded,
                )
            finally:
                self.aux_stream_list = saved_streams
            return

        # Re-check: forward's gate ran inside a captured segment that
        # _prepare_and_attn_eager (MRV1) then broke, making this region eager.
        if not current_platform.enable_multi_stream_overlap(
            self.aux_stream_list, get_forward_context().attn_metadata
        ):
            self._run_sequential_pipeline(hidden_states, positions, o_padded)
            return

        indexer = self.indexer
        compressor = self.compressor
        assert compressor is not None
        # Static per captured graph: each capture size bakes its own order.
        default_first = hidden_states.shape[0] <= _default_first_token_threshold()

        def default_chain():
            qr_kv = self._fused_wqa_wkv_gemm(hidden_states)
            qr_out, qr_scale_out, kv_out = self._split_qkv_and_norm(qr_kv)
            q = self._wq_b_proj(qr_out, qr_scale_out).view(
                -1, self.n_local_heads, self.head_dim
            )
            attn_metadata = get_forward_context().attn_metadata
            q = self._fused_qnorm_rope_kv_insert(q, kv_out, positions, attn_metadata)
            return q, qr_out, qr_scale_out, kv_out

        def main_compressor_chain() -> None:
            score = torch.mm(
                hidden_states,
                compressor.fused_wkv_wgate.weight.T,
                out_dtype=torch.float32,
            )
            compressor(score, positions, self.rotary_emb)

        if indexer is None:
            # HCA dual-stream: the main compressor runs on aux stream 0 while
            # the default stream produces q and inserts KV into the SWA cache.
            # Both branches only read hidden_states, so the join merely has to
            # precede the sparse attention that consumes the compressed KV.
            (q, _qr_out, _qr_scale_out, kv_out), _ = execute_in_parallel(
                default_chain,
                [main_compressor_chain],
                self.ln_events[0],
                [self.ln_events[1]],
                aux_streams[:1],
                enable=True,
                default_first=default_first,
            )
            self._sparse_indexer_and_attn(
                hidden_states, None, None, None, q, kv_out, positions, o_padded
            )
            return

        def indexer_compressor_chain() -> None:
            score = torch.mm(
                hidden_states,
                indexer.compressor.fused_wkv_wgate.weight.T,
                out_dtype=torch.float32,
            )
            indexer.compressor(score, positions, self.indexer_rotary_emb)

        # CSA three-stream: the main and indexer compressors run on aux
        # streams 0 and 1 while the default stream produces q and inserts KV
        # into the SWA cache. Every branch only reads hidden_states plus its
        # own state, so the join merely has to precede the indexer op and the
        # sparse attention, which consume the compressed KV caches.
        (q, qr_out, qr_scale_out, kv_out), _ = execute_in_parallel(
            default_chain,
            [main_compressor_chain, indexer_compressor_chain],
            self.ln_events[0],
            self.ln_events[1:3],
            aux_streams[:2],
            enable=True,
            default_first=default_first,
        )

        indexer_weights_out, _ = indexer.weights_proj(hidden_states)
        # The indexer compressor already ran on aux stream 1; build queries only.
        index_q, index_q_scale, weights = indexer(
            hidden_states,
            qr_out,
            None,
            indexer_weights_out,
            positions,
            self.indexer_rotary_emb,
            qr_scale_out,
            skip_compressor=True,
        )
        self._sparse_indexer_and_attn(
            hidden_states,
            index_q,
            index_q_scale,
            weights,
            q,
            kv_out,
            positions,
            o_padded,
        )

    @classmethod
    def get_padded_num_q_heads(cls, num_heads: int) -> int:
        return num_heads

    def prepare_attn_preshuffle(self) -> None:
        from vllm._aiter_ops import rocm_aiter_ops

        if not rocm_aiter_ops.is_enabled():
            return
        # aiter's gemm_a8w8_blockscale_bpreshuffle currently disabled on gfx1250.
        # Leaving the block scales as None -> "not preshuffled"
        # _fused_wqa_wkv_gemm and _o_proj both fall back to the standard linear path
        if on_gfx1250():
            return
        from vllm.model_executor.layers.quantization.utils.fp8_utils import (
            _upcast_e8m0_to_fp32,
            get_fp8_block_weight_scale,
        )
        from vllm.model_executor.utils import replace_parameter

        def _prep(linear) -> torch.Tensor | None:
            w = getattr(linear, "weight", None)
            if w is None or w.dim() != 2:
                return None
            # K % 128 (group-128 quant) and N % 16 (shuffle_weight) must hold.
            if w.shape[-1] % 128 != 0 or w.shape[0] % 16 != 0:
                return None
            ws = get_fp8_block_weight_scale(linear)
            if ws is None:
                return None
            if ws.dtype == torch.float8_e8m0fnu:
                ws = _upcast_e8m0_to_fp32(ws).contiguous()
            # Shuffle the weight in place (single weight, no unshuffled copy),
            # unless the linear kernel already B-preshuffled it at load time.
            if not getattr(linear, "aiter_bpreshuffled", False):
                replace_parameter(
                    linear,
                    "weight",
                    rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)),
                )
            return ws

        self._wqa_wkv_scale = _prep(self.fused_wqa_wkv)
        self._wo_b_scale = _prep(self.wo_b)
        if _ON_GFX950 and envs.VLLM_ROCM_USE_AITER_FP8BMM:
            self._prepare_fp8_wo_a()

    def _prepare_fp8_wo_a(self) -> None:
        try:
            from aiter.ops.batched_gemm_op_a8w8 import (
                batched_gemm_a8w8_mxscale as mxscale_op,
            )
            from aiter.ops.inverse_rope_group_quant import (
                inverse_rope_group_quant as inverse_quant_op,
            )
        except ImportError:
            logger.warning_once(
                "The DeepSeek V4 FP8 WO_A path requires AITER >= 0.1.20; "
                "falling back to BF16 WO_A."
            )
            return
        del mxscale_op, inverse_quant_op

        from vllm.model_executor.layers.quantization.utils.fp8_utils import (
            get_fp8_block_weight_scale,
            is_fp8,
        )

        weight = getattr(self.wo_a, "weight", None)
        scale = get_fp8_block_weight_scale(self.wo_a)
        if scale is None:
            # ModelOpt MXFP8 stores the multiplicative E8M0 scale without the
            # historical ``_inv`` suffix.
            scale = getattr(self.wo_a, "weight_scale", None)
        if weight is None or scale is None:
            logger.warning_once(
                "DeepSeek V4 FP8 WO_A needs a block-scaled FP8 wo_a weight; "
                "the layer exposes no weight/weight scale. Falling back to "
                "BF16 WO_A."
            )
            return
        if weight.dim() != 2 or scale.dim() != 2 or not is_fp8(weight.dtype):
            logger.warning_once(
                "DeepSeek V4 FP8 WO_A needs a 2-D FP8 wo_a weight with a 2-D "
                "block scale, got weight %s%s and scale %s. Falling back to "
                "BF16 WO_A.",
                weight.dtype,
                tuple(weight.shape),
                tuple(scale.shape),
            )
            return

        groups = self.n_local_groups
        out_per_group = self.o_lora_rank
        out_features, in_features = weight.shape
        if (
            out_features != groups * out_per_group
            or out_per_group % 128 != 0
            or in_features % 128 != 0
            or scale.shape != (out_features // 128, in_features // 128)
        ):
            logger.warning_once(
                "DeepSeek V4 FP8 WO_A needs group-128 blocks for %d groups of "
                "%d outputs, got weight %s and scale %s. Falling back to BF16 "
                "WO_A.",
                groups,
                out_per_group,
                tuple(weight.shape),
                tuple(scale.shape),
            )
            return

        e8m0_scale = _wo_a_block_scale_to_e8m0(scale)
        if e8m0_scale is None:
            logger.warning_once(
                "DeepSeek V4 FP8 WO_A could not losslessly encode the %s wo_a "
                "block scale as OCP E8M0. Falling back to BF16 WO_A.",
                scale.dtype,
            )
            return

        self._wo_a_fp8_weight = weight.view(groups, out_per_group, in_features)
        self._wo_a_e8m0_scale = e8m0_scale.view(
            groups, out_per_group // 128, in_features // 128
        )
        cache = getattr(self.rotary_emb, "cos_sin_cache_bf16", None)
        if cache is None:
            cache = self.rotary_emb.cos_sin_cache.to(dtype=torch.bfloat16)
        cos_cache, sin_cache = cache.chunk(2, dim=-1)
        self._wo_a_cos_cache = cos_cache.contiguous()
        self._wo_a_sin_cache = sin_cache.contiguous()

    def prepare_compressor_gemm_fusion(self) -> bool:
        if self._fused_compressor_weight is not None:
            return False

        from vllm.model_executor.offloader import NoopOffloader, get_offloader

        if not isinstance(get_offloader(), NoopOffloader):
            logger.warning_once(
                "DeepSeek V4 compressor GEMM fusion is incompatible with "
                "weight offloading and will remain disabled."
            )
            return False

        compressor = self.compressor
        indexer = self.indexer
        if compressor is None or indexer is None:
            return False

        main_weight = compressor.fused_wkv_wgate.weight
        indexer_weight = indexer.compressor.fused_wkv_wgate.weight
        if main_weight.ndim != 2 or indexer_weight.ndim != 2:
            raise ValueError("DeepSeek V4 compressor weights must be matrices")
        if main_weight.shape[1] != indexer_weight.shape[1]:
            raise ValueError("DeepSeek V4 compressor weights must share K")
        if main_weight.dtype != indexer_weight.dtype:
            raise ValueError("DeepSeek V4 compressor weights must share dtype")
        if main_weight.device != indexer_weight.device:
            raise ValueError("DeepSeek V4 compressor weights must share device")

        main_size = main_weight.shape[0]
        indexer_size = indexer_weight.shape[0]
        fused_weight = torch.cat((main_weight, indexer_weight), dim=0)
        with torch.no_grad():
            main_weight.set_(fused_weight[:main_size])
            indexer_weight.set_(fused_weight[main_size:])

        self._fused_compressor_weight = fused_weight
        self._fused_compressor_split_sizes = (main_size, indexer_size)
        return True

    def _bpre_attn_gemm(
        self,
        weight: torch.Tensor,
        scale: torch.Tensor,
        x: torch.Tensor,
        reduce_tp: bool,
    ) -> torch.Tensor:
        from vllm._aiter_ops import rocm_aiter_ops

        x_fp8, x_scale = rocm_aiter_ops.group_fp8_quant(x, transpose_scale=True)
        out = rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle(
            x_fp8, weight, x_scale, scale, output_dtype=x.dtype
        )
        if reduce_tp and get_tensor_model_parallel_world_size() > 1:
            out = tensor_model_parallel_all_reduce(out)
        return out

    def _fused_wqa_wkv_gemm(self, hidden_states: torch.Tensor) -> torch.Tensor:
        if self._wqa_wkv_scale is not None and hidden_states.dim() == 2:
            return self._bpre_attn_gemm(
                self.fused_wqa_wkv.weight, self._wqa_wkv_scale, hidden_states, False
            )
        return super()._fused_wqa_wkv_gemm(hidden_states)

    def _run_parallel_input_projections(
        self, hidden_states: torch.Tensor
    ) -> tuple[
        torch.Tensor,
        torch.Tensor | None,
        torch.Tensor | None,
        torch.Tensor | None,
    ]:
        fused_weight = self._fused_compressor_weight
        split_sizes = self._fused_compressor_split_sizes
        if fused_weight is None or split_sizes is None:
            return super()._run_parallel_input_projections(hidden_states)

        indexer = self.indexer
        if indexer is None:
            raise RuntimeError("Fused compressor weight requires a C4 indexer")

        qr_kv = self._fused_wqa_wkv_gemm(hidden_states)
        fused_scores = torch.mm(
            hidden_states,
            fused_weight.T,
            out_dtype=torch.float32,
        )
        kv_score, indexer_kv_score = fused_scores.split(split_sizes, dim=-1)
        indexer_weights, _ = indexer.weights_proj(hidden_states)
        return qr_kv, kv_score, indexer_kv_score, indexer_weights

    @functools.cached_property
    def _wq_b_uses_aiter_block_scaled(self) -> bool:
        """True when both wq_b GEMMs run the aiter block-scaled fp8 kernel.

        Cached: the linear kernels and the aiter env gates are fixed once
        the model is built, so this is evaluated at the first forward
        only.

        The fused norm+quant path is only valid if the quant and GEMM it
        replaces are exactly the aiter ones; otherwise fall back to the
        shared path.
        """
        from vllm._aiter_ops import rocm_aiter_ops
        from vllm.model_executor.kernels.linear.scaled_mm import (
            Fp8BlockScaledMMLinearKernel,
        )

        if not rocm_aiter_ops.is_linear_fp8_enabled():
            return False

        linears = [self.wq_b]
        if self.indexer is not None:
            linears.append(self.indexer.wq_b)
        for linear in linears:
            kernel = getattr(getattr(linear, "quant_method", None), "fp8_linear", None)
            if not isinstance(kernel, Fp8BlockScaledMMLinearKernel):
                return False
        return True

    def _split_qkv_and_norm(
        self, qr_kv: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]:
        """Fuse q/kv RMSNorm + per-1x128 fp8 q quant into one aiter kernel.

        The shared path norms q and kv in one triton kernel and the wq_b
        linears then re-read the bf16 qr to quantize it. The aiter kernel
        computes both RMSNorms (fp32 accumulate) and the fp8 group quant
        in a single pass, writing fp8 qr + group scales directly; both
        wq_b GEMMs (attention and indexer) then consume that pair and
        skip their own input quant. kv stays bf16: the fused insert
        kernel RoPE/quantizes it itself. Falls back to the shared path
        when the aiter linear path is not active.
        """
        qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1)
        if not (
            qr.dim() == 2
            and qr.shape[0] > 0
            and self.q_lora_rank % 128 == 0
            and self._wq_b_uses_aiter_block_scaled
        ):
            return super()._split_qkv_and_norm(qr_kv)

        from vllm._aiter_ops import rocm_aiter_ops

        return rocm_aiter_ops.fused_qk_rmsnorm_group_quant(
            q=qr,
            q_weight=self.q_norm.weight.data,
            q_epsilon=self.eps,
            kv=kv,
            kv_weight=self.kv_norm.weight.data,
            kv_epsilon=self.eps,
            group_size=128,
            transpose_scale=False,
        )

    def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
        if self._wo_a_fp8_weight is not None:
            from aiter.ops.batched_gemm_op_a8w8 import (
                batched_gemm_a8w8_mxscale,
            )
            from aiter.ops.inverse_rope_group_quant import (
                inverse_rope_group_quant,
            )

            assert self._wo_a_cos_cache is not None
            assert self._wo_a_sin_cache is not None
            o_fp8, o_scale = inverse_rope_group_quant(
                o.view(o.shape[0], self.n_local_heads, self.head_dim),
                positions.to(torch.int64),
                self._wo_a_cos_cache,
                self._wo_a_sin_cache,
                num_groups=self.n_local_groups,
                quant_group_size=128,
            )
            assert self._wo_a_e8m0_scale is not None
            zf = batched_gemm_a8w8_mxscale(
                o_fp8,
                self._wo_a_fp8_weight,
                o_scale,
                self._wo_a_e8m0_scale,
                dtype=o.dtype,
            ).flatten(1)
        else:
            # ROCm BF16 reference wo_a path (inverse RoPE + einsum) + wo_b.
            z = rocm_inv_rope_einsum(
                self.rotary_emb,
                o,
                positions,
                self.rope_head_dim,
                self.n_local_groups,
                self.o_lora_rank,
                self.wo_a,
                inverse_rope=False,
            )
            zf = z.flatten(1)
        if self._wo_b_scale is not None and zf.dim() == 2:
            return self._bpre_attn_gemm(self.wo_b.weight, self._wo_b_scale, zf, True)
        return self.wo_b(zf)

    def forward_mqa(
        self,
        q: torch.Tensor,
        kv: torch.Tensor,
        positions: torch.Tensor,
        output: torch.Tensor,
    ) -> None:
        assert output.shape == q.shape, (
            f"output buffer shape {output.shape} must match q shape {q.shape}"
        )
        assert output.dtype == q.dtype, (
            f"output buffer dtype {output.dtype} must match q dtype {q.dtype}"
        )

        forward_context = get_forward_context()
        attn_metadata = forward_context.attn_metadata

        if attn_metadata is None:
            # Warmup dummy run: no real metadata. Reserve the same bf16
            # gather workspace _forward_prefill would; the dequantize / topk
            # / sparse_fwd kernels are skipped this step. The aiter prefill
            # reads the caches in place and never asks for it.
            if not self._use_aiter_sparse_mla:
                swa_only = self.compress_ratio <= 1
                N = (
                    0
                    if swa_only
                    else (self.max_model_len + self.compress_ratio - 1)
                    // self.compress_ratio
                )
                M = N + self.window_size + self.max_num_batched_tokens
                current_workspace_manager().get_simultaneous(
                    ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
                )
            output.zero_()
            return

        assert isinstance(attn_metadata, dict)
        rocm_metadata = cast(
            DeepseekV4ROCMAiterMLASparseMetadata | None,
            attn_metadata.get(self.prefix),
        )
        swa_metadata = cast(
            DeepseekV4ROCMAiterSparseSWAMetadata | None,
            attn_metadata.get(self.swa_cache_layer.prefix),
        )
        assert swa_metadata is not None

        swa_only = self.compress_ratio <= 1
        self_kv_cache = self.kv_cache if not swa_only else None
        swa_kv_cache = self.swa_cache_layer.kv_cache

        num_decodes = swa_metadata.num_decodes
        num_prefills = swa_metadata.num_prefills
        num_decode_tokens = swa_metadata.num_decode_tokens

        if num_prefills > 0:
            self._forward_prefill(
                q=q[num_decode_tokens:],
                positions=positions[num_decode_tokens:],
                compressed_k_cache=self_kv_cache,
                swa_k_cache=swa_kv_cache,
                output=output[num_decode_tokens:],
                attn_metadata=rocm_metadata,
                swa_metadata=swa_metadata,
            )
        # The fp8 wo_a path rotates inside inverse_rope_group_quant, so folding
        # the rotation into the decode reduce would apply it twice. Only the
        # BF16 einsum path hands its rotation off to the decode.
        fuse_inv_rope = self._wo_a_fp8_weight is None
        rotated = 0
        if num_decodes > 0:
            fuse_decode_inv_rope = fuse_inv_rope and not on_gfx1250()
            rotated = self._forward_decode(
                q=q[:num_decode_tokens],
                positions=(
                    positions[:num_decode_tokens] if fuse_decode_inv_rope else None
                ),
                kv_cache=self_kv_cache,
                swa_metadata=swa_metadata,
                attn_metadata=rocm_metadata,
                swa_only=swa_only,
                output=output[:num_decode_tokens],
                adaptive_splits=(
                    _ON_GFX950
                    and not swa_only
                    and self.compress_ratio == 128
                    and rocm_metadata is not None
                    and rocm_metadata.for_cudagraph_capture
                ),
            )
        if fuse_inv_rope:
            # Only the decode reduce rotates its own rows, and only the leading
            # `rotated` of them; prefill rows and any decode path that did not
            # fuse still owe the standalone pass. Settle that here rather than
            # in _o_proj: the split is batch-dependent and _o_proj runs
            # compiled, where such a value freezes at its trace-time value.
            rocm_inverse_rope_rows_(
                output[rotated:, : self.n_local_heads, :],
                positions[rotated:],
                self.rotary_emb.cos_sin_cache,
                self.rope_head_dim,
            )

    def _forward_decode(
        self,
        q: torch.Tensor,
        positions: torch.Tensor | None,
        kv_cache: torch.Tensor | None,
        swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata,
        attn_metadata: DeepseekV4ROCMAiterMLASparseMetadata | None,
        swa_only: bool,
        output: torch.Tensor,
        adaptive_splits: bool,
    ) -> int:
        """Returns how many leading rows the decode epilogue inverse-RoPE'd."""
        num_decodes = swa_metadata.num_decodes
        num_decode_tokens = swa_metadata.num_decode_tokens

        topk_indices = None
        topk_lens = None
        topk_ragged_indices = None
        topk_ragged_indptr = None
        if not swa_only:
            assert attn_metadata is not None
            assert swa_metadata.is_valid_token is not None
            block_size = attn_metadata.block_size // self.compress_ratio
            is_valid = swa_metadata.is_valid_token[:num_decode_tokens]
            if self.compress_ratio == 4:
                assert self.topk_indices_buffer is not None
                (
                    topk_ragged_indices,
                    topk_ragged_indptr,
                    topk_lens,
                ) = compute_global_topk_ragged_indices_and_indptr(
                    self.topk_indices_buffer[:num_decode_tokens],
                    swa_metadata.token_to_req_indices,
                    attn_metadata.block_table[:num_decodes],
                    block_size,
                    is_valid,
                )
            else:
                topk_indices = attn_metadata.c128a_global_decode_topk_indices
                topk_lens = attn_metadata.c128a_decode_topk_lens
                topk_ragged_indices = attn_metadata.c128a_decode_topk_ragged_indices
                topk_ragged_indptr = attn_metadata.c128a_decode_topk_ragged_indptr

        if self._use_aiter_sparse_mla:
            assert swa_metadata.decode_swa_ragged_indices is not None
            assert swa_metadata.decode_swa_ragged_indptr is not None
            self._aiter_sparse_mla(
                q,
                output,
                self.swa_cache_layer.kv_cache,
                swa_metadata.decode_swa_ragged_indices,
                swa_metadata.decode_swa_ragged_indptr,
                kv_cache,
                topk_ragged_indices,
                topk_ragged_indptr,
            )
            # The aiter kernel leaves the inverse RoPE to the caller.
            return 0

        return rocm_sparse_attn_decode(
            q=q,
            kv_cache=kv_cache,
            swa_k_cache=self.swa_cache_layer.kv_cache,
            swa_only=swa_only,
            topk_indices=topk_indices,
            topk_lens=topk_lens,
            swa_indices=swa_metadata.decode_swa_indices,
            swa_lens=swa_metadata.decode_swa_lens,
            swa_ragged_indices=swa_metadata.decode_swa_ragged_indices,
            swa_ragged_indptr=swa_metadata.decode_swa_ragged_indptr,
            topk_ragged_indices=topk_ragged_indices,
            topk_ragged_indptr=topk_ragged_indptr,
            attn_sink=self.attn_sink,
            scale=self.scale,
            head_dim=self.head_dim,
            nope_head_dim=self.nope_head_dim,
            rope_head_dim=self.rope_head_dim,
            output=output,
            adaptive_splits=adaptive_splits,
            inv_rope_positions=positions,
            inv_rope_cos_sin_cache=self.rotary_emb.cos_sin_cache,
            extra_cache_nan_free=_trust_dsv4_extra_cache_nan_free(
                self.kv_cache_dtype,
                self._has_kv_transfer,
                not swa_only and kv_cache is not None,
            ),
        )

    def _forward_prefill(
        self,
        q: torch.Tensor,
        positions: torch.Tensor,
        compressed_k_cache: torch.Tensor | None,
        swa_k_cache: torch.Tensor,
        output: torch.Tensor,
        attn_metadata: DeepseekV4ROCMAiterMLASparseMetadata | None,
        swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata,
    ) -> None:
        if self._use_aiter_sparse_mla:
            self._forward_prefill_aiter(
                q=q,
                compressed_k_cache=compressed_k_cache,
                swa_k_cache=swa_k_cache,
                output=output,
                attn_metadata=attn_metadata,
                swa_metadata=swa_metadata,
            )
            return

        swa_only = attn_metadata is None

        num_prefill_tokens = swa_metadata.num_prefill_tokens
        num_decodes = swa_metadata.num_decodes
        num_decode_tokens = swa_metadata.num_decode_tokens

        seq_lens = swa_metadata.prefill_seq_lens
        gather_lens = swa_metadata.prefill_gather_lens
        assert seq_lens is not None
        assert gather_lens is not None

        query_start_loc_cpu = swa_metadata.query_start_loc_cpu
        query_start_loc = swa_metadata.query_start_loc
        assert query_start_loc_cpu is not None
        assert query_start_loc is not None
        prefill_token_base = query_start_loc_cpu[num_decodes]
        left_visible = swa_metadata.prefill_left_visible
        right_visible = swa_metadata.prefill_right_visible
        if left_visible is not None:
            left_visible = left_visible[num_decode_tokens:]
            assert right_visible is not None
            right_visible = right_visible[num_decode_tokens:]

        if not swa_only:
            if self.compress_ratio == 4:
                assert self.topk_indices_buffer is not None
                topk_indices = self.topk_indices_buffer[num_decode_tokens:]
                topk_indices = topk_indices[:num_prefill_tokens]
            else:
                assert attn_metadata is not None
                topk_indices = attn_metadata.c128a_prefill_topk_indices
            assert topk_indices is not None
            top_k = topk_indices.shape[-1]
        else:
            assert self.topk_indices_buffer is not None
            topk_indices = self.topk_indices_buffer[num_decode_tokens:]
            top_k = 0

        chunk_plan = swa_metadata.get_prefill_chunk_plan(
            compress_ratio=self.compress_ratio,
            prefill_chunk_size=self.PREFILL_CHUNK_SIZE,
        )
        assert chunk_plan, "prefill chunk plan must be non-empty when num_prefills > 0"
        workspace_manager = current_workspace_manager()
        for chunk_start, chunk_end, chunk_N, chunk_M in chunk_plan:
            chunk_size = chunk_end - chunk_start
            kv = workspace_manager.get_simultaneous(
                ((chunk_size, chunk_M, q.shape[-1]), torch.bfloat16),
            )[0]
            if not swa_only:
                assert attn_metadata is not None
                assert compressed_k_cache is not None
                block_table = attn_metadata.block_table[num_decodes:]
                # compressed_k_cache is OCP on every platform (Triton encoder).
                dequantize_and_gather_k_cache(
                    kv[:chunk_size],
                    compressed_k_cache,
                    seq_lens=seq_lens[chunk_start:chunk_end] // self.compress_ratio,
                    gather_lens=None,
                    block_table=block_table[chunk_start:chunk_end],
                    block_size=attn_metadata.block_size // self.compress_ratio,
                    offset=0,
                    use_fnuz=False,
                )

            swa_block_table = swa_metadata.block_table[num_decodes:]
            dequantize_and_gather_k_cache(
                kv[:chunk_size],
                swa_k_cache,
                seq_lens=seq_lens[chunk_start:chunk_end],
                gather_lens=gather_lens[chunk_start:chunk_end],
                block_table=swa_block_table[chunk_start:chunk_end],
                block_size=swa_metadata.block_size,
                offset=chunk_N,
                use_fnuz=current_platform.is_fp8_fnuz(),
            )

            query_start = (
                query_start_loc_cpu[num_decodes + chunk_start] - prefill_token_base
            )
            query_end = (
                query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base
            )

            combined_indices, combined_lens = combine_topk_swa_indices(
                topk_indices[query_start:query_end],
                query_start_loc[
                    num_decodes + chunk_start : num_decodes + chunk_end + 1
                ],
                seq_lens[chunk_start:chunk_end],
                gather_lens[chunk_start:chunk_end],
                self.window_size,
                self.compress_ratio,
                top_k,
                chunk_M,
                chunk_N,
                max_image_tokens=self.max_image_tokens,
                left_visible=(
                    left_visible[query_start:query_end]
                    if left_visible is not None
                    else None
                ),
                right_visible=(
                    right_visible[query_start:query_end]
                    if right_visible is not None
                    else None
                ),
            )
            rocm_sparse_attn_prefill(
                q=q[query_start:query_end],
                kv=kv.view(-1, 1, q.shape[-1]),
                indices=combined_indices,
                topk_length=combined_lens,
                scale=self.scale,
                head_dim=self.head_dim,
                nope_head_dim=self.nope_head_dim,
                rope_head_dim=self.rope_head_dim,
                attn_sink=self.attn_sink,
                output=output[query_start:query_end],
            )

    def _forward_prefill_aiter(
        self,
        q: torch.Tensor,
        compressed_k_cache: torch.Tensor | None,
        swa_k_cache: torch.Tensor,
        output: torch.Tensor,
        attn_metadata: DeepseekV4ROCMAiterMLASparseMetadata | None,
        swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata,
    ) -> None:
        """Prefill straight off the paged caches, with no bf16 gather."""
        num_decode_tokens = swa_metadata.num_decode_tokens
        num_prefill_tokens = swa_metadata.num_prefill_tokens
        assert swa_metadata.prefill_swa_ragged_indices is not None
        assert swa_metadata.prefill_swa_ragged_indptr is not None

        topk_ragged_indices = None
        topk_ragged_indptr = None
        if attn_metadata is None:
            compressed_k_cache = None
        else:
            assert compressed_k_cache is not None
            assert swa_metadata.token_to_req_indices is not None
            assert swa_metadata.query_start_loc is not None
            assert swa_metadata.seq_lens is not None
            assert swa_metadata.is_valid_token is not None
            if self.compress_ratio == 4:
                assert self.topk_indices_buffer is not None
                topk_indices = self.topk_indices_buffer[
                    num_decode_tokens : num_decode_tokens + num_prefill_tokens
                ]
            else:
                topk_indices = attn_metadata.c128a_prefill_topk_indices
            assert topk_indices is not None
            topk_ragged_indices, topk_ragged_indptr = build_prefill_topk_ragged_indices(
                topk_indices,
                swa_metadata.token_to_req_indices,
                swa_metadata.query_start_loc,
                swa_metadata.seq_lens,
                swa_metadata.is_valid_token,
                attn_metadata.block_table,
                block_size=attn_metadata.block_size // self.compress_ratio,
                compress_ratio=self.compress_ratio,
                num_compressed=(self.max_model_len + self.compress_ratio - 1)
                // self.compress_ratio,
                token_offset=num_decode_tokens,
                num_rows=compressed_k_cache.shape[0] * compressed_k_cache.shape[1],
            )

        self._aiter_sparse_mla(
            q[:num_prefill_tokens],
            output[:num_prefill_tokens],
            swa_k_cache,
            swa_metadata.prefill_swa_ragged_indices,
            swa_metadata.prefill_swa_ragged_indptr,
            compressed_k_cache,
            topk_ragged_indices,
            topk_ragged_indptr,
        )

    def _aiter_sparse_mla(
        self,
        q: torch.Tensor,
        output: torch.Tensor,
        swa_k_cache: torch.Tensor,
        swa_indices: torch.Tensor,
        swa_indptr: torch.Tensor,
        compressed_k_cache: torch.Tensor | None,
        topk_indices: torch.Tensor | None,
        topk_indptr: torch.Tensor | None,
    ) -> None:
        from vllm._aiter_ops import rocm_aiter_ops

        # Both fp8 caches are read in place: the SWA window as the main segment,
        # the top-k compressed tokens as the extra one.
        rocm_aiter_ops.triton_sparse_mla_fwd(
            q,
            swa_k_cache,
            output,
            self.scale,
            swa_indptr,
            swa_indices,
            attn_sink=self.attn_sink[: q.shape[1]],
            extra_kv_buffer=compressed_k_cache,
            extra_kv_indptr=topk_indptr,
            extra_kv_indices=topk_indices,
        )