launcher/dcp2/baked/models/deepseek_v4_1/nvidia/flashinfer_sparse.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""DeepSeek V4 FlashInfer sparse MLA backend."""

import os
from typing import TYPE_CHECKING, ClassVar, cast

import torch

from vllm.config import VllmConfig
from vllm.distributed import get_dcp_group
from vllm.logger import init_logger
from vllm.config.cache import CacheDType
from vllm.forward_context import get_forward_context
from vllm.models.deepseek_v4.nvidia.ops.o_proj import (
    compute_fp8_einsum_recipe,
    deep_gemm_fp8_o_proj,
)
from vllm.models.deepseek_v4_1.attention import DeepseekV4Attention
from vllm.models.deepseek_v4_1.common.ops import (
    build_flashinfer_mixed_sparse_indices,
    compute_global_topk_indices_and_lens,
)
from vllm.models.deepseek_v4_1.sparse_mla import (
    DeepseekV4FlashMLAMetadata,
    DeepseekV4SparseMLABackend,
    DeepseekV4SparseMLAMetadataBuilder,
    DeepseekV41SparseSWAMetadataBuilder,
)
from vllm.platforms import current_platform
from vllm.platforms.interface import DeviceCapability
from vllm.utils.flashinfer import flashinfer_trtllm_batch_decode_sparse_mla_dsv4
from vllm.v1.attention.backend import AttentionCGSupport, MultipleOf
from vllm.v1.attention.backends.mla.compressor_utils import (
    get_dspark_swa_index_width,
)
from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWABackend
from vllm.v1.attention.backends.mla.sparse_utils import (
    triton_filter_and_convert_dcp_index,
)
from vllm.v1.attention.ops.dcp import (
    CPTritonContext,
    cp_lse_ag_out_rs,
    dcp_a2a_lse_reduce,
)

if TYPE_CHECKING:
    from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata

logger = init_logger(__name__)

from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (  # noqa: E402
    _PREFILL_TIMING as _DSV41_PREFILL_TIMING,
    _PrefillTimer as _DSV41PrefillTimer,
    prefill_acc_add as _dsv41_prefill_acc_add,
)

# Stage 9: DCP prefill attention via compressed-KV gather (no q gather / combine).
_DCP_PREFILL_GATHER = os.environ.get("DSV41_DCP_PREFILL_GATHER", "0") == "1"

# FlashInfer's SM120 sparse-MLA kernels write this (not -inf) into the LSE of a
# row with no candidate and no sink; the DCP combine needs -inf there.
_SM120_EMPTY_LSE = -1e29


def _sm120_sparse_attention_with_lse(
    query: torch.Tensor,
    swa_kv_cache: torch.Tensor,
    workspace: torch.Tensor,
    swa_indices: torch.Tensor,
    swa_lens: torch.Tensor | None,
    compressed_kv_cache: torch.Tensor | None,
    extra_indices: torch.Tensor | None,
    extra_lens: torch.Tensor | None,
    out: torch.Tensor,
    sm_scale: float,
    sinks: torch.Tensor | None,
    lse: torch.Tensor,
) -> None:
    """FlashInfer's SM120 packed sparse-MLA attention, returning the LSE.

    Mirrors the SM120 branch of ``trtllm_batch_decode_sparse_mla_dsv4``
    (flashinfer 0.7.0rc1, ``flashinfer/mla/_core.py``), whose public entry
    hard-codes ``return_lse=False``. ``query``/``out`` are ``[T, H, 512]``
    bf16, ``lse`` is ``[T, H]`` fp32 and receives the LSE in base 2 of the
    ``sm_scale``-scaled logits with the sink folded in (``log2(2^lse +
    2^(sink*log2e))``); rows with no candidate and no sink hold ``-1e30``.
    """
    from flashinfer.mla._core import (
        _check_sm120_dsv4_kv_cache_layout,
        _SparseMLASegment,
        _trtllm_batch_decode_sparse_mla_sm120,
    )

    num_tokens = query.shape[0]
    swa_kv_cache = _check_sm120_dsv4_kv_cache_layout(swa_kv_cache, "NHD", "swa_kv_cache")
    segments = [
        _SparseMLASegment(indices=swa_indices.reshape(num_tokens, -1), lengths=swa_lens)
    ]
    if extra_indices is not None:
        assert compressed_kv_cache is not None
        compressed_kv_cache = _check_sm120_dsv4_kv_cache_layout(
            compressed_kv_cache, "NHD", "compressed_kv_cache"
        )
        segments.append(
            _SparseMLASegment(
                indices=extra_indices.reshape(num_tokens, -1),
                lengths=extra_lens,
                kv_cache=compressed_kv_cache,
            )
        )
    _trtllm_batch_decode_sparse_mla_sm120(
        query=query.unsqueeze(1),
        kv_cache=swa_kv_cache,
        workspace_buffer=workspace,
        sparse_mla_segments=segments,
        out=out.unsqueeze(1),
        sm_scale=float(sm_scale),
        sinks=sinks,
        lse=lse,
        return_lse=True,
        kv_scale_format="auto",
        kv_cache_format="fp8",
    )

_FLASHINFER_DSV4_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024
_flashinfer_dsv4_workspace_by_device: dict[torch.device, torch.Tensor] = {}


def _get_flashinfer_dsv4_workspace(device: torch.device) -> torch.Tensor:
    workspace = _flashinfer_dsv4_workspace_by_device.get(device)
    if workspace is None:
        workspace = torch.zeros(
            _FLASHINFER_DSV4_WORKSPACE_BUFFER_SIZE,
            dtype=torch.uint8,
            device=device,
        )
        _flashinfer_dsv4_workspace_by_device[device] = workspace
    return workspace


def _packed_block_span(pool: torch.Tensor) -> int:
    """Per-block stride of ``pool`` in tokens (``stride(0)//stride(-2)``): ==
    block_size for unpacked KV, larger when packed (#44577). Raises if not
    token-aligned."""
    block_stride = pool.stride(0)
    token_stride = pool.stride(-2)
    if block_stride % token_stride != 0:
        raise NotImplementedError(
            "FLASHINFER_MLA_SPARSE_DSV4 packed KV requires the per-block stride "
            f"({block_stride}) to be a multiple of the per-token stride "
            f"({token_stride}); this layout is not supported yet."
        )
    return block_stride // token_stride


# Sparse MLA h_q counts accepted natively (flashinfer>=0.6.14, #3545).
_SPARSE_MLA_SUPPORTED_Q_HEADS = (8, 16, 32, 64, 128)


def _pad_to_supported_q_heads(num_heads: int) -> int:
    for supported in _SPARSE_MLA_SUPPORTED_Q_HEADS:
        if num_heads <= supported:
            return supported
    raise ValueError(
        f"DeepseekV4 FlashInfer MLA Sparse does not support {num_heads} heads "
        "(sparse MLA kernel requires h_q in {8, 16, 32, 64, 128})."
    )


def _required_sm120_sparse_topk(vllm_config: VllmConfig, window_size: int) -> int:
    """Return the SM120 DSV4 SWA specialization needed by this model."""
    if not vllm_config.attention_config.use_non_causal:
        return window_size
    speculative_config = vllm_config.speculative_config
    if speculative_config is None:
        return window_size
    return get_dspark_swa_index_width(
        window_size,
        speculative_config.num_speculative_tokens,
    )


class DeepseekV4FlashInferMLASparseBackend(DeepseekV4SparseMLABackend):
    """FlashInfer backend using the DSv4 sparse metadata/cache layout.

    Inherits the base and backend reuses its``DeepseekV4SparseMLAMetadataBuilder``
    """

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

    @staticmethod
    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
        # Tech2Wild/Kai 2026-09-10 (SM12x): 64-state compressed pages, so a
        # ratio-1 kv-source layer needs 64-token pages (128 stays first so the
        # cache-config default remains 128 for ratio-2 layers).
        if current_platform.is_cuda() and current_platform.is_device_capability_family(120):
            return [128, 64]
        return [128]

    @classmethod
    def get_compressed_block_size(cls, compress_ratio: int) -> int | None:
        """Tech2Wild/Kai 2026-09-10: tokens per compressed-KV page for a kv-source
        layer on SM12x (FlashInfer SM120 sparse-MLA prefill/decode kernels are
        instantiated for 64-state pages); None = use cache_config.block_size."""
        if compress_ratio > 0 and current_platform.is_cuda() and current_platform.is_device_capability_family(120):
            return 64 * compress_ratio
        return None

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

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

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

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

    @classmethod
    def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
        return capability.major in [10, 12]

    @classmethod
    def supports_combination(
        cls,
        head_size: int,
        dtype: torch.dtype,
        kv_cache_dtype: CacheDType | None,
        block_size: int | None,
        use_mla: bool,
        has_sink: bool,
        use_sparse: bool,
        use_mm_prefix: bool,
        device_capability: DeviceCapability,
    ) -> str | None:
        if device_capability.major == 10:
            if kv_cache_dtype == "fp8_ds_mla":
                return (
                    "FLASHINFER_MLA_SPARSE_DSV4 SM10x uses the plain "
                    "per-tensor FP8 KV layout, not fp8_ds_mla"
                )
            if kv_cache_dtype not in (None, "auto", "bfloat16", "fp8", "fp8_e4m3"):
                return "kv_cache_dtype not supported"
            return None
        if device_capability.major == 12:
            if kv_cache_dtype not in ("fp8", "fp8_e4m3", "fp8_ds_mla"):
                return "kv_cache_dtype not supported"
            from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120

            if not has_flashinfer_sparse_mla_sm120():
                return (
                    "FLASHINFER_MLA_SPARSE_DSV4 SM120 requires FlashInfer's "
                    "sparse MLA decode API"
                )
            return None
        return "FLASHINFER_MLA_SPARSE_DSV4 requires SM10x or SM12x"

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


class DeepseekV4FlashInferSparseMLAMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder):
    """Varlen-capable metadata builder for the FlashInfer sparse MLA backend."""

    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # DCP port staging (DCP2-PLAN.md): the sharded-KV attention path
        # (per-rank sparse attention + LSE combine) lands in stage 3. Until
        # DSV41_DCP_STAGE >= 3, a DCP>1 boot must fail here -- after the KV
        # cache config is computed and logged, before any forward -- rather
        # than run the dcp=1 kernels over a sharded cache and serve garbage.
        import os

        dcp = self.vllm_config.parallel_config.decode_context_parallel_size
        if dcp > 1 and int(os.environ.get("DSV41_DCP_STAGE", "3")) < 3:
            raise NotImplementedError(
                f"DeepSeek-V4.1 sparse MLA under DCP{dcp}: the attention combine "
                "(port stage 3) is disabled by DSV41_DCP_STAGE; unset it or run "
                "with --decode-context-parallel-size 1."
            )


class DeepseekSparseSWAFlashInferMetadataBuilder(DeepseekV41SparseSWAMetadataBuilder):
    """SWA metadata for the FlashInfer sparse decode path (varlen decode)."""

    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS


class DeepseekSparseSWAFlashInferBackend(DeepseekSparseSWABackend):
    @staticmethod
    def get_builder_cls() -> type[DeepseekSparseSWAFlashInferMetadataBuilder]:
        return DeepseekSparseSWAFlashInferMetadataBuilder


class DeepseekSparseSWAFlashInferSM120Backend(DeepseekSparseSWAFlashInferBackend):
    """Tech2Wild/Kai 2026-09-10: SM12x. FlashInfer's SM120 sparse-MLA decode and
    prefill kernels are instantiated only for 64-token pages
    (_DECODE_DSV4_PAGE_BLOCK_SIZE = 64); the base backend's "any multiple of 32"
    made vLLM build the SWA cache at 32-token pages, which the SM120 path rejects
    ("no decode kernel for this shape ... page_block_size=32")."""

    @staticmethod
    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
        return [64]

    @classmethod
    def get_preferred_block_size(cls, default_block_size: int) -> int:
        return 64

    @classmethod
    def get_swa_block_size(cls) -> int:
        return 64


class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
    """FlashInfer TRTLLM-gen sparse MLA attention layer for SM100 DeepSeek V4."""

    backend_cls = DeepseekV4FlashInferMLASparseBackend
    swa_backend_cls = DeepseekSparseSWAFlashInferBackend
    use_fp8_ds_mla_layout: ClassVar[bool] = False

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

    def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
        return deep_gemm_fp8_o_proj(
            o,
            positions,
            self.rotary_emb.cos_sin_cache,
            self.wo_a,
            self.wo_b,
            n_groups=self.n_local_groups,
            heads_per_group=self.n_local_heads // self.n_local_groups,
            nope_dim=self.nope_head_dim,
            rope_dim=self.rope_head_dim,
            o_lora_rank=self.o_lora_rank,
            einsum_recipe=self._einsum_recipe,
            tma_aligned_scales=self._tma_aligned_scales,
        )

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe(
            self._o_proj_block_size
        )
        # Per-tensor FP8 scale buffers + precomputed scalar BMM scales. Only the
        # per-tensor FP8 cache path consumes these; bf16 reads ``self.scale``.
        if self.kv_cache_torch_dtype != torch.float8_e4m3fn:
            return
        fp8_q_scale = 1.0
        fp8_kv_scale = 1.0
        self.register_buffer(
            "_flashinfer_fp8_q_scale",
            torch.tensor([fp8_q_scale], dtype=torch.float32),
            persistent=False,
        )
        self.register_buffer(
            "_flashinfer_fp8_q_scale_inv",
            torch.tensor([1.0 / fp8_q_scale], dtype=torch.float32),
            persistent=False,
        )
        self.register_buffer(
            "_flashinfer_fp8_kv_scale",
            torch.tensor([fp8_kv_scale], dtype=torch.float32),
            persistent=False,
        )
        # TRTLLM-gen takes scalar scale args on a distinct C++ path vs
        # one-element tensors, so these are Python floats.
        self._flashinfer_fp8_bmm1_scale = self.scale * fp8_q_scale * fp8_kv_scale
        self._flashinfer_fp8_bmm2_scale = fp8_kv_scale

    def forward_mqa(
        self,
        q: torch.Tensor,
        kv: torch.Tensor,
        positions: torch.Tensor,
        output: torch.Tensor,
    ) -> None:
        # The TRTLLM-gen kernel requires h_q in {64, 128}, so the output buffer
        # is allocated at the padded head count while q arrives at the local
        # head count; _forward pads q to match before the launcher.
        assert output.shape[0] == q.shape[0] and output.shape[-1] == q.shape[-1], (
            f"output buffer shape {output.shape} incompatible with q shape {q.shape}"
        )
        assert output.shape[1] >= q.shape[1], (
            f"output heads {output.shape[1]} must be >= q heads {q.shape[1]}"
        )
        # Per-tensor FP8 q produces a bf16 attention output.
        expected_output_dtype = (
            torch.bfloat16 if q.dtype == torch.float8_e4m3fn else q.dtype
        )
        assert output.dtype == expected_output_dtype, (
            f"output dtype {output.dtype} must match expected {expected_output_dtype} "
            f"for q dtype {q.dtype}"
        )

        forward_context = get_forward_context()
        attn_metadata = forward_context.attn_metadata
        if attn_metadata is None:
            # Warmup dummy run: FlashInfer reads the cache directly and lazily
            # allocates its workspace, so nothing to reserve here.
            output.zero_()
            return

        assert isinstance(attn_metadata, dict)
        # Compressed-cache metadata lives on the kv-source layer's prefix;
        # consumers share that cache and its block table.
        flashmla_metadata = cast(
            DeepseekV4FlashMLAMetadata | None,
            attn_metadata.get(self.compressed_cache_prefix)
            if self.compressed_cache_prefix is not None
            else None,
        )
        swa_metadata = cast(
            "DeepseekSparseSWAMetadata | None",
            attn_metadata.get(self.swa_cache_layer.prefix),
        )
        assert swa_metadata is not None

        swa_only = self.compress_ratio == 0
        # SWA-only layers have no compressed KV cache; consumers read the kv
        # source's cache.
        self_kv_cache = None if swa_only else self._compressed_kv_cache()
        swa_kv_cache = self.swa_cache_layer.kv_cache

        self._forward(
            q=q,
            kv_cache=self_kv_cache,
            swa_k_cache=swa_kv_cache,
            swa_metadata=swa_metadata,
            attn_metadata=flashmla_metadata,
            swa_only=swa_only,
            output=output,
        )

    def _build_sparse_index_metadata(
        self,
        kv_cache: torch.Tensor | None,
        swa_k_cache: torch.Tensor,
        swa_metadata: "DeepseekSparseSWAMetadata",
        attn_metadata: DeepseekV4FlashMLAMetadata | None,
        swa_only: bool,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """Build the combined sparse-index tensors for the mixed batch.

        Returns ``(compressed_kv_cache, seq_lens, sparse_indices,
        sparse_topk_lens)``.
        """
        num_decodes = swa_metadata.num_decodes
        num_prefills = swa_metadata.num_prefills
        num_decode_tokens = swa_metadata.num_decode_tokens
        num_prefill_tokens = swa_metadata.num_prefill_tokens
        num_reqs = num_decodes + num_prefills
        num_tokens = num_decode_tokens + num_prefill_tokens

        assert swa_metadata.seq_lens is not None
        assert swa_metadata.query_start_loc is not None
        assert swa_metadata.token_to_req_indices is not None
        assert swa_metadata.decode_swa_indices is not None
        assert swa_metadata.block_table is not None

        decode_swa_indices = swa_metadata.decode_swa_indices.reshape(
            num_decode_tokens, swa_metadata.decode_swa_width
        )
        decode_compressed_topk_lens = None
        decode_compressed_indices_are_local = False
        decode_is_valid_token = None

        if swa_only:
            assert self.topk_indices_buffer is not None
            compressed_kv_cache = swa_k_cache
            decode_compressed_indices = None
            prefill_topk_indices = self.topk_indices_buffer[
                num_decode_tokens:num_tokens, :0
            ]
            compressed_block_table = None
            compressed_block_size = swa_metadata.block_size
            top_k = 0
        else:
            assert kv_cache is not None
            assert attn_metadata is not None
            assert self.topk_indices_buffer is not None
            assert swa_metadata.is_valid_token is not None
            compressed_kv_cache = kv_cache
            compressed_block_table = attn_metadata.block_table[:num_reqs]
            compressed_block_size = attn_metadata.block_size // self.compress_ratio

            # Local indices filled by the index-source layer's indexer.
            if num_prefill_tokens > 0:
                prefill_topk_indices = self.topk_indices_buffer[
                    num_decode_tokens:num_tokens
                ]
                top_k = prefill_topk_indices.shape[-1]
            else:
                prefill_topk_indices = self.topk_indices_buffer[:0, :0]
                top_k = 0

            decode_compressed_indices_are_local = True
            decode_is_valid_token = swa_metadata.is_valid_token[:num_decode_tokens]
            if num_decode_tokens > 0:
                decode_compressed_indices = self.topk_indices_buffer[:num_decode_tokens]
            else:
                # Keep the logical width aligned with the mixed-batch case so
                # pure-prefill steps reuse the same Triton specialization.
                decode_compressed_indices = prefill_topk_indices[:0]

        query_start_loc = swa_metadata.query_start_loc[: num_reqs + 1]
        seq_lens = swa_metadata.seq_lens[:num_reqs]
        assert seq_lens.dtype == torch.int32
        # SWA-only layers all build the same mixed sparse indices, so the first
        # one caches them for the step; indexer layers depend on their own topk
        # indices and stay uncached.
        cached_sparse = (
            swa_metadata.flashinfer_sparse_index_cache.get("swa_only")
            if swa_only
            else None
        )
        if cached_sparse is None:
            swa_block_span = _packed_block_span(swa_k_cache)
            compressed_block_span = _packed_block_span(compressed_kv_cache)
            sparse_indices, sparse_topk_lens = build_flashinfer_mixed_sparse_indices(
                decode_swa_indices,
                decode_compressed_indices,
                decode_compressed_topk_lens,
                prefill_topk_indices[:num_prefill_tokens],
                query_start_loc,
                seq_lens,
                swa_metadata.token_to_req_indices[:num_tokens],
                swa_metadata.block_table[:num_reqs],
                swa_metadata.block_size,
                compressed_block_table,
                compressed_block_size,
                self.window_size,
                self.compress_ratio,
                top_k,
                decode_compressed_indices_are_local=decode_compressed_indices_are_local,
                decode_is_valid_token=decode_is_valid_token,
                swa_block_span=swa_block_span,
                compressed_block_span=compressed_block_span,
                prefill_left_visible=swa_metadata.prefill_left_visible,
                prefill_right_visible=swa_metadata.prefill_right_visible,
                # getattr for tests that bypass __init__ via object.__new__.
                max_image_tokens=getattr(self, "max_image_tokens", 0),
            )
            if swa_only:
                swa_metadata.flashinfer_sparse_index_cache["swa_only"] = (
                    sparse_indices,
                    sparse_topk_lens,
                )
        else:
            sparse_indices, sparse_topk_lens = cached_sparse
        return compressed_kv_cache, seq_lens, sparse_indices, sparse_topk_lens

    def _forward(
        self,
        q: torch.Tensor,
        kv_cache: torch.Tensor | None,
        swa_k_cache: torch.Tensor,
        swa_metadata: "DeepseekSparseSWAMetadata",
        attn_metadata: DeepseekV4FlashMLAMetadata | None,
        swa_only: bool,
        output: torch.Tensor,
    ) -> None:
        assert self.kv_cache_torch_dtype in (torch.bfloat16, torch.float8_e4m3fn)
        num_decodes = swa_metadata.num_decodes
        num_prefills = swa_metadata.num_prefills
        num_decode_tokens = swa_metadata.num_decode_tokens
        num_prefill_tokens = swa_metadata.num_prefill_tokens
        num_reqs = num_decodes + num_prefills
        num_tokens = num_decode_tokens + num_prefill_tokens
        if num_tokens == 0:
            return

        (
            compressed_kv_cache,
            seq_lens,
            sparse_indices,
            sparse_topk_lens,
        ) = self._build_sparse_index_metadata(
            kv_cache=kv_cache,
            swa_k_cache=swa_k_cache,
            swa_metadata=swa_metadata,
            attn_metadata=attn_metadata,
            swa_only=swa_only,
        )

        # CUDA graph execution can pad q/output past the scheduled token count;
        # restrict to the real tokens (the launcher validates sparse indices).
        query = q[:num_tokens]
        output = output[:num_tokens]
        bmm1_scale: float | torch.Tensor = self.scale
        bmm2_scale: float | torch.Tensor = 1.0
        if self.kv_cache_torch_dtype == torch.float8_e4m3fn:
            assert query.dtype == torch.float8_e4m3fn
            bmm1_scale = self._flashinfer_fp8_bmm1_scale
            bmm2_scale = self._flashinfer_fp8_bmm2_scale
        else:
            assert query.dtype == torch.bfloat16
            query = query.contiguous()

        # The TRTLLM-gen sparse-MLA kernel requires h_q in {64, 128}; zero-pad
        # the query heads to the allocated output head count. Padded heads attend
        # to the shared KV and are sliced off downstream (output is padded too).
        padded_heads = output.shape[1]
        if query.shape[1] < padded_heads:
            padded_query = query.new_zeros(
                (query.shape[0], padded_heads, query.shape[2])
            )
            padded_query[:, : query.shape[1], :] = query
            query = padded_query

        workspace = _get_flashinfer_dsv4_workspace(q.device)
        query_start_loc = swa_metadata.query_start_loc
        query_start_loc_cpu = swa_metadata.query_start_loc_cpu
        assert query_start_loc is not None and query_start_loc_cpu is not None

        # Keep the TRTLLM-gen decode/prefill split: the launcher is tuned for
        # uniform-q batches, and this avoids flattening mixed batches into one call.
        if num_decode_tokens > 0:
            decode_cu = query_start_loc[: num_decodes + 1]
            flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
                query=query[:num_decode_tokens],
                swa_kv_cache=swa_k_cache,
                workspace_buffer=workspace,
                sparse_indices=sparse_indices[:num_decode_tokens],
                compressed_kv_cache=compressed_kv_cache,
                sparse_topk_lens=sparse_topk_lens[:num_decode_tokens],
                seq_lens=seq_lens[:num_decodes],
                out=output[:num_decode_tokens],
                bmm1_scale=bmm1_scale,
                bmm2_scale=bmm2_scale,
                sinks=self.attn_sink,
                cum_seq_lens_q=decode_cu,
                max_q_len=swa_metadata.max_decode_query_len,
            )

        if num_prefill_tokens > 0:
            # The prefill query view re-anchors at offset 0, so rebase the
            # cumulative query offsets to start at 0.
            prefill_cu = (
                query_start_loc[num_decodes : num_reqs + 1]
                - query_start_loc[num_decodes]
            )
            prefill_cu_cpu = query_start_loc_cpu[num_decodes : num_reqs + 1]
            prefill_lens_cpu = prefill_cu_cpu[1:] - prefill_cu_cpu[:-1]
            flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
                query=query[num_decode_tokens:num_tokens],
                swa_kv_cache=swa_k_cache,
                workspace_buffer=workspace,
                sparse_indices=sparse_indices[num_decode_tokens:num_tokens],
                compressed_kv_cache=compressed_kv_cache,
                sparse_topk_lens=sparse_topk_lens[num_decode_tokens:num_tokens],
                seq_lens=seq_lens[num_decodes:num_reqs],
                out=output[num_decode_tokens:num_tokens],
                bmm1_scale=bmm1_scale,
                bmm2_scale=bmm2_scale,
                sinks=self.attn_sink,
                cum_seq_lens_q=prefill_cu,
                max_q_len=int(prefill_lens_cpu.max().item()),
            )


class DeepseekV4FlashInferSM120Attention(DeepseekV4Attention):
    """DeepSeek V4 sparse MLA attention through FlashInfer's SM120 kernels."""

    backend_cls = DeepseekV4FlashInferMLASparseBackend
    swa_backend_cls = DeepseekSparseSWAFlashInferSM120Backend
    use_fp8_ds_mla_layout: ClassVar[bool] = True
    # DCP: per-rank partial attention with the LSE + cross-rank combine.
    can_return_lse_for_decode: ClassVar[bool] = True
    # Under DCP the per-token q gather / combine has no cross-token dependency:
    # bound the transient [T, 2*heads, 512] buffers inside a prefill chunk.
    DCP_PREFILL_TOKEN_CHUNK: ClassVar[int] = 2048

    @staticmethod
    def _get_workspace(device: torch.device) -> torch.Tensor:
        return _get_flashinfer_dsv4_workspace(device)

    @staticmethod
    def _as_sparse_cache(kv_cache: torch.Tensor) -> torch.Tensor:
        if kv_cache.dtype == torch.float8_e4m3fn:
            kv_cache = kv_cache.view(torch.uint8)
        if kv_cache.dim() == 4:
            return kv_cache
        return kv_cache.unsqueeze(-2)

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

    def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
        return deep_gemm_fp8_o_proj(
            o,
            positions,
            self.rotary_emb.cos_sin_cache,
            self.wo_a,
            self.wo_b,
            n_groups=self.n_local_groups,
            heads_per_group=self.n_local_heads // self.n_local_groups,
            nope_dim=self.nope_head_dim,
            rope_dim=self.rope_head_dim,
            o_lora_rank=self.o_lora_rank,
            einsum_recipe=self._einsum_recipe,
            tma_aligned_scales=self._tma_aligned_scales,
        )

    def __init__(self, vllm_config: VllmConfig, *args, **kwargs) -> None:
        super().__init__(vllm_config, *args, **kwargs)
        from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120_config

        required_topk = _required_sm120_sparse_topk(vllm_config, self.window_size)
        # DCP: the kernel sees the gathered heads [rank0 heads | rank1 heads |
        # ...] in TP order; the local slice must be unpadded so the head
        # reduce-scatter hands each rank exactly its own heads back.
        self.n_dcp_heads = self.n_local_heads * self.dcp_world_size
        if self.dcp_world_size > 1:
            if self.padded_heads != self.n_local_heads:
                raise NotImplementedError(
                    "DeepSeek-V4.1 DCP needs an unpadded local head count in "
                    f"{_SPARSE_MLA_SUPPORTED_Q_HEADS}; got {self.n_local_heads} "
                    f"(padded to {self.padded_heads})."
                )
            if _pad_to_supported_q_heads(self.n_dcp_heads) != self.n_dcp_heads:
                raise NotImplementedError(
                    f"DeepSeek-V4.1 DCP gathered head count {self.n_dcp_heads} is "
                    f"not a supported sparse-MLA h_q {_SPARSE_MLA_SUPPORTED_Q_HEADS}."
                )
        kernel_heads = self.n_dcp_heads if self.dcp_world_size > 1 else self.padded_heads
        if not has_flashinfer_sparse_mla_sm120_config(kernel_heads, required_topk):
            raise RuntimeError(
                "FLASHINFER_MLA_SPARSE_DSV4 on SM120 requires a FlashInfer "
                "DSV4 sparse MLA decode specialization for "
                f"(num_q_heads={kernel_heads}, top_k={required_topk}). "
                "Install a FlashInfer build containing "
                "flashinfer-ai/flashinfer#4380."
            )
        # DCP state: the group (consecutive TP ranks), the gathered sink
        # (built lazily on the first forward, after weight loading) and the
        # neutral SWA index rows used on ranks other than 0.
        self.dcp_group = get_dcp_group() if self.dcp_world_size > 1 else None
        self._dcp_sink: torch.Tensor | None = None
        self._dcp_empty_swa: dict[int, tuple[torch.Tensor, torch.Tensor]] = {}
        # One Triton launch context per layer for the LSE-combine kernel (a
        # fresh context per call re-specializes the launch on every step).
        self._cp_ctx = CPTritonContext() if self.dcp_world_size > 1 else None
        # Combine flavour: "ag_rs" (LSE all-gather + head reduce-scatter, two
        # collectives) or "a2a" (one all-to-all of packed output+LSE).
        self._dcp_combine_a2a = os.environ.get("DSV41_DCP_COMBINE", "ag_rs") == "a2a"
        # Optional per-phase timing (CUDA events, synchronises; diagnostics
        # only): DSV41_DCP_TIMING=1 logs gather / kernel / combine ms per call,
        # averaged over the last 100 calls, on DCP rank 0.
        self._dcp_timing = os.environ.get("DSV41_DCP_TIMING", "0") == "1"
        self._dcp_t_acc = [0.0, 0.0, 0.0, 0]
        if self.dcp_world_size > 1 and self.compress_ratio > 0:
            logger.info_once(
                "DeepSeek-V4.1 SM120 attention under DCP%d: %d local heads -> %d "
                "gathered, SWA window + sink on DCP rank 0, base-2 LSE combine",
                self.dcp_world_size,
                self.n_local_heads,
                self.n_dcp_heads,
            )
        self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe(
            self._o_proj_block_size
        )
        # Per-tensor FP8 cache path scales.
        if self.kv_cache_torch_dtype != torch.float8_e4m3fn:
            return
        fp8_q_scale = 1.0
        fp8_kv_scale = 1.0
        self.register_buffer(
            "_flashinfer_fp8_q_scale",
            torch.tensor([fp8_q_scale], dtype=torch.float32),
            persistent=False,
        )
        self.register_buffer(
            "_flashinfer_fp8_q_scale_inv",
            torch.tensor([1.0 / fp8_q_scale], dtype=torch.float32),
            persistent=False,
        )
        self.register_buffer(
            "_flashinfer_fp8_kv_scale",
            torch.tensor([fp8_kv_scale], dtype=torch.float32),
            persistent=False,
        )
        # FlashInfer expects scalar scale arguments for this path.
        self._flashinfer_fp8_bmm1_scale = self.scale * fp8_q_scale * fp8_kv_scale
        self._flashinfer_fp8_bmm2_scale = fp8_kv_scale

    def _reserve_empty_forward_workspace(self) -> None:
        self._get_workspace(
            torch.device("cuda", torch.accelerator.current_device_index())
        )

    def _forward_sparse_impl(
        self,
        q: torch.Tensor,
        output: torch.Tensor,
        flashmla_metadata: DeepseekV4FlashMLAMetadata | None,
        swa_metadata: "DeepseekSparseSWAMetadata",
        self_kv_cache: torch.Tensor | None,
        swa_kv_cache: torch.Tensor,
        swa_only: bool,
    ) -> None:
        num_decode_tokens = swa_metadata.num_decode_tokens
        if swa_metadata.num_prefills > 0:
            self._forward_prefill(
                q=q[num_decode_tokens:],
                compressed_k_cache=self_kv_cache,
                swa_k_cache=swa_kv_cache,
                output=output[num_decode_tokens:],
                attn_metadata=flashmla_metadata,
                swa_metadata=swa_metadata,
            )
        if swa_metadata.num_decodes > 0:
            self._forward_decode(
                q=q[:num_decode_tokens],
                kv_cache=self_kv_cache,
                swa_metadata=swa_metadata,
                attn_metadata=flashmla_metadata,
                swa_only=swa_only,
                output=output[:num_decode_tokens],
            )

    def forward_mqa(
        self,
        q: torch.Tensor,
        kv: torch.Tensor,
        positions: torch.Tensor,
        output: torch.Tensor,
    ) -> None:
        # Output may be padded to backend-supported head counts.
        assert output.shape[0] == q.shape[0] and output.shape[-1] == q.shape[-1], (
            f"output buffer shape {output.shape} incompatible with q shape {q.shape}"
        )
        assert output.shape[1] >= q.shape[1], (
            f"output heads {output.shape[1]} must be >= q heads {q.shape[1]}"
        )
        # Per-tensor FP8 q produces a bf16 attention output.
        expected_output_dtype = (
            torch.bfloat16 if q.dtype == torch.float8_e4m3fn else q.dtype
        )
        assert output.dtype == expected_output_dtype, (
            f"output dtype {output.dtype} must match expected {expected_output_dtype} "
            f"for q dtype {q.dtype}"
        )

        forward_context = get_forward_context()
        attn_metadata = forward_context.attn_metadata
        if attn_metadata is None:
            self._reserve_empty_forward_workspace()
            output.zero_()
            return

        assert isinstance(attn_metadata, dict)
        # Compressed-cache metadata lives on the kv-source layer's prefix;
        # consumers share that cache and its block table.
        flashmla_metadata = cast(
            DeepseekV4FlashMLAMetadata | None,
            attn_metadata.get(self.compressed_cache_prefix)
            if self.compressed_cache_prefix is not None
            else None,
        )
        swa_metadata = cast(
            "DeepseekSparseSWAMetadata | None",
            attn_metadata.get(self.swa_cache_layer.prefix),
        )
        assert swa_metadata is not None

        swa_only = self.compress_ratio == 0
        # SWA-only layers have no compressed KV cache; consumers read the kv
        # source's cache.
        self_kv_cache = None if swa_only else self._compressed_kv_cache()
        swa_kv_cache = self.swa_cache_layer.kv_cache

        self._forward_sparse_impl(
            q=q,
            output=output,
            flashmla_metadata=flashmla_metadata,
            swa_metadata=swa_metadata,
            self_kv_cache=self_kv_cache,
            swa_kv_cache=swa_kv_cache,
            swa_only=swa_only,
        )

    def _prepare_query(self, q: torch.Tensor, output: torch.Tensor) -> torch.Tensor:
        if self.kv_cache_torch_dtype == torch.float8_e4m3fn:
            assert q.dtype == torch.float8_e4m3fn
            q = q.to(torch.bfloat16)
        else:
            assert q.dtype == torch.bfloat16
        padded_heads = output.shape[1]
        if q.shape[1] < padded_heads:
            padded_query = q.new_zeros((q.shape[0], padded_heads, q.shape[2]))
            padded_query[:, : q.shape[1], :] = q
            q = padded_query
        return q.contiguous()

    # ---- DCP helpers -------------------------------------------------------

    def _compressed_topk_to_slots(
        self,
        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]:
        """Indexer top-k (compressed ids) -> physical slots + per-row count.

        DCP1: the request-local ids map through the block table as today.
        DCP>1: the ids are GLOBAL compressed ids (indexer merge / short-context
        bypass); de-interleave them in state coordinates, drop the states other
        ranks own and compact this rank's slots to a prefix (-1 tail, count).
        ``block_size`` is the number of states per page on this rank.
        """
        num_tokens = topk_indices.shape[0]
        if self.dcp_world_size == 1:
            return compute_global_topk_indices_and_lens(
                topk_indices,
                token_to_req_indices,
                block_table,
                block_size,
                is_valid_token,
            )
        slots, lens = triton_filter_and_convert_dcp_index(
            token_to_req_indices[:num_tokens].contiguous(),
            block_table,
            topk_indices,
            dcp_size=self.dcp_world_size,
            dcp_rank=self.dcp_rank,
            cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
            BLOCK_SIZE=block_size,
            BLOCK_STRIDE_ROWS=block_size,
            NUM_TOPK_TOKENS=topk_indices.shape[1],
            BLOCK_N=128,
            return_valid_counts=True,
        )
        # CUDA-graph padding rows attend nothing (parity with the DCP1 kernel).
        lens.masked_fill_(~is_valid_token[:num_tokens], 0)
        return slots, lens

    # Stage 9: one prefill step's gathered compressed KV per source cache and
    # remapped top-k per index buffer, shared by every layer of that step. The
    # metadata object is held so the cache resets exactly once per step.
    _dcp_gather_cache: ClassVar[dict] = {"meta": None, "kv": {}, "idx": {}, "geom": None}

    def _dcp_prefill_gather(
        self,
        compressed_k_cache: torch.Tensor,
        topk_indices: torch.Tensor,
        token_to_req: torch.Tensor,
        is_valid_token: torch.Tensor,
        attn_metadata: DeepseekV4FlashMLAMetadata,
        num_decodes: int,
        num_prefills: int,
        block_size: int,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """DCP prefill without q gather / LSE combine: gather the DCP ranks'
        compressed pages of every prefill request into one rank-major paged
        buffer and remap the GLOBAL top-k ids into it (see stage 9 notes).
        Returns (gathered cache [W*P_total, block_size, 1, B], slots, lens)."""
        cache = type(self)._dcp_gather_cache
        if cache["meta"] is not attn_metadata:
            cache["meta"] = attn_metadata
            cache["kv"] = {}
            cache["idx"] = {}
            cache["geom"] = None
        W = self.dcp_world_size
        K = self.cp_kv_cache_interleave_size
        bs = block_size
        dev = topk_indices.device
        if cache["geom"] is None:
            assert attn_metadata.seq_lens is not None
            seq_lens = attn_metadata.seq_lens[num_decodes : num_decodes + num_prefills].tolist()
            # states that exist after this step (state = pos // ratio) -> virtual
            # blocks of bs*W states = one page per rank; at least one page so
            # the collective never carries zero bytes.
            pages = [max(1, -(-(int(sl) // self.compress_ratio) // (bs * W))) for sl in seq_lens]
            base = [0]
            for p in pages:
                base.append(base[-1] + W * p)
            cache["geom"] = (
                pages,
                torch.tensor(pages, device=dev, dtype=torch.int64),
                torch.tensor(base[:-1], device=dev, dtype=torch.int64),
            )
        pages, pages_t, base_t = cache["geom"]
        key = compressed_k_cache.data_ptr()
        kv = cache["kv"].get(key)
        if kv is None:
            with _DSV41PrefillTimer("attn_gather", int(topk_indices.shape[0])):
                pool = self._as_sparse_cache(compressed_k_cache)
                parts = []
                for i, p in enumerate(pages):
                    page_ids = attn_metadata.block_table[num_decodes + i, :p].long()
                    local = pool.index_select(0, page_ids).contiguous()
                    parts.append(self.dcp_group.all_gather(local, dim=0))
                kv = parts[0] if len(parts) == 1 else torch.cat(parts, dim=0)
            cache["kv"][key] = kv
        ikey = (topk_indices.data_ptr(), tuple(topk_indices.shape))
        idx = cache["idx"].get(ikey)
        if idx is None:
            n_req = max(len(pages) - 1, 0)
            req = (token_to_req.long() - num_decodes).clamp_(0, n_req)
            ids = topk_indices.long()
            valid = ids >= 0
            st = ids.clamp_min(0)
            span = bs * W
            vblock = st // span
            off = st % span
            owner = (off // K) % W
            loff = (off // (K * W)) * K + off % K
            slot = (base_t[req][:, None] + owner * pages_t[req][:, None] + vblock) * bs + loff
            slots = torch.where(valid, slot, torch.full_like(slot, -1)).to(torch.int32)
            lens = valid.sum(dim=-1, dtype=torch.int32)
            lens.masked_fill_(~is_valid_token[: lens.shape[0]], 0)
            idx = (slots, lens)
            cache["idx"][ikey] = idx
        return kv, idx[0], idx[1]

    def _dcp_attn_sink(self) -> torch.Tensor:
        """The sink for the gathered heads, in TP head order (collective)."""
        if self._dcp_sink is None:
            assert self.dcp_group is not None
            local = self.attn_sink.data[: self.n_local_heads].contiguous()
            self._dcp_sink = self.dcp_group.all_gather(local, dim=0)
        return self._dcp_sink

    def _dcp_neutral_swa(
        self, num_tokens: int, width: int, device: torch.device
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """All -1 SWA indices and zero lengths (ranks other than 0)."""
        buf = self._dcp_empty_swa.get(width)
        if buf is None or buf[0].shape[0] < num_tokens:
            rows = max(num_tokens, self.max_num_batched_tokens)
            buf = (
                torch.full((rows, 1, width), -1, dtype=torch.int32, device=device),
                torch.zeros(rows, dtype=torch.int32, device=device),
            )
            self._dcp_empty_swa[width] = buf
        return buf[0][:num_tokens], buf[1][:num_tokens]

    def _run_sm120_dcp(
        self,
        q: torch.Tensor,
        swa_cache: torch.Tensor,
        swa_indices: torch.Tensor,
        swa_lens: torch.Tensor,
        extra_cache: torch.Tensor | None,
        extra_indices: torch.Tensor | None,
        extra_lens: torch.Tensor | None,
        output: torch.Tensor,
    ) -> None:
        """One DCP attention step for ``q``/``output`` of ``[T, n_local_heads, 512]``.

        Gather the local heads to ``[T, n_local_heads * W, 512]`` (TP order),
        run this rank's sparse attention over its compressed-state shard with
        the LSE, then combine across the DCP group and reduce-scatter the heads
        back. The replicated SWA window and the sink are attended on DCP rank 0
        only, so with ``L_r`` the per-rank base-2 LSEs (sink folded into
        ``L_0``) the combine ``L = log2(sum_r 2^L_r)`` and ``o = sum_r
        2^(L_r - L) o_r`` reproduce the DCP1 softmax exactly.
        """
        assert self.dcp_group is not None
        num_tokens = q.shape[0]
        # Timing events cannot be synchronised inside a CUDA-graph capture (the
        # eager break runs on the capture stream during warmup/capture).
        capturing = torch.cuda.is_current_stream_capturing()
        prefill_timing = _DSV41_PREFILL_TIMING and num_tokens >= 1024 and not capturing
        timing = (self._dcp_timing or prefill_timing) and not capturing
        if self._dcp_timing:
            logger.info_once(
                "DCP path active on layer %d (capturing=%s, T=%d)",
                self.layer_id,
                capturing,
                num_tokens,
            )
        if timing:
            ev = [torch.cuda.Event(enable_timing=True) for _ in range(4)]
            ev[0].record()
        # Collective: every rank builds the gathered sink (lazily, once); only
        # DCP rank 0 applies it, together with the replicated SWA window.
        sink = self._dcp_attn_sink()
        q_gathered = self.dcp_group.all_gather(q, dim=1)
        num_heads = q_gathered.shape[1]
        if timing:
            ev[1].record()
        out_gathered = torch.empty(
            (num_tokens, num_heads, self.head_dim), dtype=torch.bfloat16, device=q.device
        )
        lse = torch.empty((num_tokens, num_heads), dtype=torch.float32, device=q.device)
        if self.dcp_rank != 0:
            swa_indices, swa_lens = self._dcp_neutral_swa(
                num_tokens, swa_indices.shape[-1], q.device
            )
        _sm120_sparse_attention_with_lse(
            q_gathered,
            swa_cache,
            self._get_workspace(q.device),
            swa_indices,
            swa_lens,
            extra_cache,
            extra_indices,
            extra_lens,
            out_gathered,
            self.scale,
            sink if self.dcp_rank == 0 else None,
            lse,
        )
        if timing:
            ev[2].record()
        # Rows with no local candidate (and no sink) carry the kernel's -1e30
        # LSE and a zero output; the combine kernel's max-subtracted exp2 makes
        # that the merge identity without any masking (test_stage3_combine).
        if self._dcp_combine_a2a:
            combined = dcp_a2a_lse_reduce(
                out_gathered, lse, self.dcp_group, is_lse_base_on_e=False
            )
        else:
            combined = cp_lse_ag_out_rs(
                out_gathered,
                lse,
                self.dcp_group,
                ctx=self._cp_ctx,
                is_lse_base_on_e=False,
            )
        output.copy_(combined)
        if timing:
            ev[3].record()
            ev[3].synchronize()
            if prefill_timing:
                _dsv41_prefill_acc_add("attn_gather", ev[0].elapsed_time(ev[1]))
                _dsv41_prefill_acc_add("attn_kernel", ev[1].elapsed_time(ev[2]))
                _dsv41_prefill_acc_add("attn_combine", ev[2].elapsed_time(ev[3]))
            acc = self._dcp_t_acc
            acc[0] += ev[0].elapsed_time(ev[1])
            acc[1] += ev[1].elapsed_time(ev[2])
            acc[2] += ev[2].elapsed_time(ev[3])
            acc[3] += 1
            if self._dcp_timing and acc[3] % 100 == 0 and self.dcp_rank == 0:
                logger.info(
                    "DCP timing layer %d (T=%d): gather %.3f ms, kernel %.3f ms, "
                    "combine %.3f ms per call (avg of %d)",
                    self.layer_id,
                    num_tokens,
                    acc[0] / acc[3],
                    acc[1] / acc[3],
                    acc[2] / acc[3],
                    acc[3],
                )
                acc[0] = acc[1] = acc[2] = 0.0
                acc[3] = 0

    def _forward_decode(
        self,
        q: torch.Tensor,
        kv_cache: torch.Tensor | None,
        swa_metadata: "DeepseekSparseSWAMetadata",
        attn_metadata: DeepseekV4FlashMLAMetadata | None,
        swa_only: bool,
        output: torch.Tensor,
    ) -> None:
        num_decodes = swa_metadata.num_decodes
        num_decode_tokens = swa_metadata.num_decode_tokens

        extra_sparse_indices = None
        extra_sparse_lengths = None
        if not swa_only:
            if attn_metadata is None:
                raise RuntimeError(
                    "Sparse MLA metadata is required for compressed layers."
                )
            if swa_metadata.is_valid_token is None:
                raise RuntimeError(
                    "SWA validity metadata is required for compressed layers."
                )
            if self.topk_indices_buffer is None:
                raise RuntimeError(
                    "Compressed-layer decode requires top-k indices from the indexer."
                )
            # Indices filled by the index-source layer's indexer (request-local
            # at DCP1, global compressed ids under DCP).
            is_valid = swa_metadata.is_valid_token[:num_decode_tokens]
            block_size = attn_metadata.block_size // self.compress_ratio
            global_indices, extra_sparse_lengths = self._compressed_topk_to_slots(
                self.topk_indices_buffer[:num_decode_tokens],
                swa_metadata.token_to_req_indices,
                attn_metadata.block_table[:num_decodes],
                block_size,
                is_valid,
            )
            extra_sparse_indices = global_indices.view(num_decode_tokens, 1, -1)

        swa_indices = swa_metadata.decode_swa_indices
        swa_lens = swa_metadata.decode_swa_lens
        assert swa_indices is not None
        assert swa_lens is not None
        q = self._prepare_query(q, output)
        swa_cache = self._as_sparse_cache(self.swa_cache_layer.kv_cache)
        extra_cache = self._as_sparse_cache(kv_cache) if kv_cache is not None else None
        if extra_cache is not None and extra_sparse_indices is None:
            raise RuntimeError(
                "Compressed sparse MLA decode requires compressed sparse indices."
            )
        if self.dcp_world_size > 1 and not swa_only:
            # SWA-only layers keep the DCP1 path: their cache is replicated and
            # every rank runs its own heads over the whole window.
            self._run_sm120_dcp(
                q,
                swa_cache,
                swa_indices,
                swa_lens,
                extra_cache,
                extra_sparse_indices,
                extra_sparse_lengths,
                output,
            )
            return
        flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
            query=q,
            swa_kv_cache=swa_cache,
            workspace_buffer=self._get_workspace(q.device),
            sparse_indices=swa_indices,
            compressed_kv_cache=extra_cache,
            out=output,
            bmm1_scale=self.scale,
            sinks=self.attn_sink,
            kv_layout="NHD",
            swa_topk_lens=swa_lens,
            extra_sparse_indices=extra_sparse_indices,
            extra_sparse_topk_lens=extra_sparse_lengths,
        )

    def _forward_prefill(
        self,
        q: torch.Tensor,
        compressed_k_cache: torch.Tensor | None,
        swa_k_cache: torch.Tensor,
        output: torch.Tensor,
        attn_metadata: DeepseekV4FlashMLAMetadata | None,
        swa_metadata: "DeepseekSparseSWAMetadata",
    ) -> None:
        swa_only = self.compress_ratio == 0

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

        query_start_loc_cpu = swa_metadata.query_start_loc_cpu
        assert query_start_loc_cpu is not None
        prefill_token_base = query_start_loc_cpu[num_decodes]

        extra_sparse_indices: torch.Tensor | None = None
        extra_sparse_lengths: torch.Tensor | None = None
        gathered_kv: torch.Tensor | None = None
        if not swa_only:
            if self.topk_indices_buffer is None:
                raise RuntimeError(
                    "Compressed-layer prefill requires top-k indices from the indexer."
                )
            if attn_metadata is None:
                raise RuntimeError("Compressed-layer prefill metadata is missing.")
            if swa_metadata.token_to_req_indices is None:
                raise RuntimeError(
                    "Compressed-layer prefill request mapping is missing."
                )
            if swa_metadata.is_valid_token is None:
                raise RuntimeError(
                    "Compressed-layer prefill validity metadata is missing."
                )
            # Local indices filled by the index-source layer's indexer.
            local_topk_indices = self.topk_indices_buffer[
                num_decode_tokens : num_decode_tokens + num_prefill_tokens
            ]
            prefill_token_slice = slice(
                num_decode_tokens, num_decode_tokens + num_prefill_tokens
            )
            block_size = attn_metadata.block_size // self.compress_ratio
            if self.dcp_world_size > 1 and _DCP_PREFILL_GATHER:
                assert compressed_k_cache is not None
                gathered_kv, extra_sparse_indices, extra_sparse_lengths = (
                    self._dcp_prefill_gather(
                        compressed_k_cache,
                        local_topk_indices,
                        swa_metadata.token_to_req_indices[prefill_token_slice],
                        swa_metadata.is_valid_token[prefill_token_slice],
                        attn_metadata,
                        num_decodes,
                        num_prefills,
                        block_size,
                    )
                )
            else:
                extra_sparse_indices, extra_sparse_lengths = (
                    self._compressed_topk_to_slots(
                        local_topk_indices,
                        swa_metadata.token_to_req_indices[prefill_token_slice],
                        attn_metadata.block_table,
                        block_size,
                        swa_metadata.is_valid_token[prefill_token_slice],
                    )
                )

        assert swa_metadata.prefill_swa_indices is not None
        assert swa_metadata.prefill_swa_lens is not None

        q = self._prepare_query(q, output)
        swa_kv_paged = self._as_sparse_cache(swa_k_cache)
        if swa_only:
            extra_kv_paged = None
        else:
            if compressed_k_cache is None:
                raise RuntimeError(
                    "Compressed sparse MLA layers require their compressed KV cache."
                )
            extra_kv_paged = (
                gathered_kv
                if gathered_kv is not None
                else self._as_sparse_cache(compressed_k_cache)
            )

        num_chunks = (
            num_prefills + self.PREFILL_CHUNK_SIZE - 1
        ) // self.PREFILL_CHUNK_SIZE
        for chunk_idx in range(num_chunks):
            chunk_start = chunk_idx * self.PREFILL_CHUNK_SIZE
            chunk_end = min(chunk_start + self.PREFILL_CHUNK_SIZE, num_prefills)
            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
            )

            extra_sparse_indices_chunk = (
                extra_sparse_indices[query_start:query_end]
                if extra_sparse_indices is not None
                else None
            )
            extra_sparse_lengths_chunk = (
                extra_sparse_lengths[query_start:query_end]
                if extra_sparse_lengths is not None
                else None
            )

            q_chunk = q[query_start:query_end]
            swa_indices_chunk = swa_metadata.prefill_swa_indices[query_start:query_end]
            swa_lens_chunk = swa_metadata.prefill_swa_lens[query_start:query_end]
            if extra_kv_paged is not None and extra_sparse_indices_chunk is None:
                raise RuntimeError(
                    "Compressed sparse MLA prefill requires compressed sparse indices."
                )
            if self.dcp_world_size > 1 and not swa_only and gathered_kv is None:
                # Chunk bounds come from the CPU query_start_loc (identical on
                # every rank), so the collectives below line up across ranks.
                qs, qe = int(query_start), int(query_end)
                for sub_start in range(qs, qe, self.DCP_PREFILL_TOKEN_CHUNK):
                    sub_end = min(sub_start + self.DCP_PREFILL_TOKEN_CHUNK, qe)
                    assert extra_sparse_indices is not None
                    assert extra_sparse_lengths is not None
                    self._run_sm120_dcp(
                        q[sub_start:sub_end],
                        swa_kv_paged,
                        swa_metadata.prefill_swa_indices[sub_start:sub_end],
                        swa_metadata.prefill_swa_lens[sub_start:sub_end],
                        extra_kv_paged,
                        extra_sparse_indices[sub_start:sub_end],
                        extra_sparse_lengths[sub_start:sub_end],
                        output[sub_start:sub_end],
                    )
                continue
            with _DSV41PrefillTimer(
                "attn_kernel",
                int(query_end - query_start) if gathered_kv is not None else 0,
            ):
                flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
                    query=q_chunk,
                    swa_kv_cache=swa_kv_paged,
                    workspace_buffer=self._get_workspace(q.device),
                    sparse_indices=swa_indices_chunk,
                    compressed_kv_cache=extra_kv_paged,
                    out=output[query_start:query_end],
                    bmm1_scale=self.scale,
                    sinks=self.attn_sink,
                    kv_layout="NHD",
                    swa_topk_lens=swa_lens_chunk,
                    extra_sparse_indices=extra_sparse_indices_chunk,
                    extra_sparse_topk_lens=extra_sparse_lengths_chunk,
                )