launcher/dcp2/base-theirs/v1/attention/backends/mla/sparse_swa.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from dataclasses import dataclass, field
from typing import Any, ClassVar, cast

import torch

from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
from vllm.model_executor.warmup.jit_warmup import (
    WarmupIntRange,
    zip_inputs,
)
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
    LaunchSpec,
    TritonWarmupTensor,
    VllmTritonJitKernel,
    kernel_launcher,
)
from vllm.platforms import current_platform
from vllm.triton_utils import tl, triton
from vllm.utils.math_utils import cdiv, next_power_of_2
from vllm.v1.attention.backend import (
    AttentionBackend,
    AttentionCGSupport,
    AttentionMetadataBuilder,
    CommonAttentionMetadata,
    MultipleOf,
)
from vllm.v1.attention.backends.mla.compressor_utils import (
    get_dspark_swa_index_width,
)
from vllm.v1.attention.backends.utils import split_decodes_and_prefills
from vllm.v1.attention.ops.flashmla import FlashMLASchedMeta, get_mla_metadata
from vllm.v1.kv_cache_interface import (
    KVCacheSpec,
    MLAAttentionSpec,
    SlidingWindowMLASpec,
    get_kv_quant_mode,
)

# DeepseekV4 decode layer types, keyed by compress_ratio. Each type has a distinct
# (topk, extra_topk, extra_page_block_size) config, so they cannot share a
# FlashMLA tile-scheduler plan. Within a type, all ~60 DeepseekV4 layers share one
# plan per step because b / s_q / h_q / page_block_sizes / topks are identical.
_LAYER_TYPE_SWAONLY = "swaonly"
_LAYER_TYPE_C4A = "c4a"
_LAYER_TYPE_C128A = "c128a"


def _layer_type_for(compress_ratio: int) -> str:
    if compress_ratio <= 1:
        return _LAYER_TYPE_SWAONLY
    if compress_ratio == 4:
        return _LAYER_TYPE_C4A
    if compress_ratio == 128:
        return _LAYER_TYPE_C128A
    raise ValueError(
        f"Unsupported DeepseekV4 compress_ratio={compress_ratio}; "
        "expected 1, 4, or 128."
    )


class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):
    def __init__(
        self,
        head_dim: int,
        window_size: int,
        dtype: torch.dtype,
        prefix: str,
        cache_config: CacheConfig,
        backend_cls: "type[AttentionBackend] | None" = None,
    ):
        super().__init__()
        self.backend_cls = backend_cls or DeepseekSparseSWABackend
        self.kv_cache = torch.tensor([])
        self.head_dim = head_dim
        self.window_size = window_size
        self.prefix = prefix
        self.cache_config = cache_config
        self.dtype = dtype
        compilation_config = get_current_vllm_config().compilation_config
        if prefix in compilation_config.static_forward_context:
            raise ValueError(f"Duplicate layer name: {prefix}")
        compilation_config.static_forward_context[prefix] = self

        # Block size is constrained by tensor sharing between SWA and C4A KV blocks.
        # Since both block types share the same physical tensor, they must use the
        # same page size. The C4A KV block shape [256//4, head_dim] = [64, head_dim]
        # determines the SWA block size of 64 tokens per block.
        # TODO(yifan): make SWA block size automatically determined and configurable.
        self.block_size = 64
        # uint8: fp8_ds_mla UE8M0 paged layout. bfloat16 / float8_e4m3fn:
        # contiguous full-cache layout.
        assert self.dtype in (torch.uint8, torch.bfloat16, torch.float8_e4m3fn)

    def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
        # [B, H=1, N, C] -> [B, N, C]
        self.kv_cache = kv_cache.squeeze(1)

    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
        # fp8_ds_mla's UE8M0 paged layout needs 576B alignment; contiguous
        # bf16/fp8 cache uses the natural element-size page.
        uses_fp8_ds_mla_layout = self.cache_config.cache_dtype == "fp8_ds_mla"
        return SlidingWindowMLASpec(
            block_size=self.block_size,
            num_kv_heads=1,
            head_size=self.head_dim,
            dtype=self.dtype,
            sliding_window=self.window_size,
            cache_dtype_str=self.cache_config.cache_dtype,
            # DeepseekV4 fp8_ds_mla: 584B per token (448B NoPE + 128B RoPE + 8B scales)
            state_content_bytes=584 if uses_fp8_ds_mla_layout else None,
            # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577).
            alignment=576 if uses_fp8_ds_mla_layout else 512,
            model_version="deepseek_v4",
            kv_quant_mode=get_kv_quant_mode(self.cache_config.cache_dtype),
        )

    def forward(self): ...

    def get_attn_backend(self) -> type[AttentionBackend]:
        return self.backend_cls


class DeepseekSparseSWABackend(AttentionBackend):
    @staticmethod
    def get_name() -> str:
        return "DEEPSEEK_SPARSE_SWA"

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

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

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

    @staticmethod
    def get_builder_cls() -> type["DeepseekSparseSWAMetadataBuilder"]:
        if current_platform.is_rocm():
            from vllm.models.deepseek_v4.amd.rocm import (
                DeepseekV4ROCMAiterSparseSWAMetadataBuilder,
            )

            return DeepseekV4ROCMAiterSparseSWAMetadataBuilder
        return DeepseekSparseSWAMetadataBuilder


@dataclass
class DeepseekSparseSWAMetadata:
    block_table: torch.Tensor
    slot_mapping: torch.Tensor
    block_size: int
    seq_lens: torch.Tensor | None = None  # [num_seqs]
    query_start_loc: torch.Tensor | None = None  # [num_seqs + 1]
    query_start_loc_cpu: torch.Tensor | None = None  # [num_seqs + 1]

    is_valid_token: torch.Tensor | None = None  # [num_tokens]
    token_to_req_indices: torch.Tensor | None = None  # [num_tokens]
    decode_swa_indices: torch.Tensor | None = None  # [num_decode_tokens, width]
    decode_swa_lens: torch.Tensor | None = None  # [num_decode_tokens]
    # window_size (causal) or noncausal_index_width (DSpark non-causal).
    decode_swa_width: int = 0
    # Paged-coordinate prefill SWA indices/lens (FP8 paged-direct prefill).
    prefill_swa_indices: torch.Tensor | None = (
        None  # [num_prefill_tokens, 1, prefill_index_width]
    )
    prefill_swa_lens: torch.Tensor | None = None  # [num_prefill_tokens]
    # In-image bidirectional visibility (vision variant): per-token counts of
    # extra visible tokens to the left/right inside an image span. Indexed by
    # absolute (decode-first) token position; only prefill rows are written.
    # None when the model is text-only or the batch has no image spans.
    prefill_left_visible: torch.Tensor | None = None
    prefill_right_visible: torch.Tensor | None = None

    # Number of decode/prefill requests/tokens (batch is reordered: decodes first)
    num_decodes: int = 0
    num_prefills: int = 0
    num_decode_tokens: int = 0
    num_prefill_tokens: int = 0
    max_decode_query_len: int = 1

    # Pre-computed prefill metadata shared across all DeepseekV4 attention layers.
    prefill_seq_lens: torch.Tensor | None = None
    prefill_seq_lens_cpu: torch.Tensor | None = None
    prefill_gather_lens: torch.Tensor | None = None
    prefill_query_lens_cpu: torch.Tensor | None = None
    prefill_window_size: int = 0
    prefill_max_model_len: int = 0
    prefill_max_num_batched_tokens: int = 0

    # Per-layer-type FlashMLA tile-scheduler metadata. One FlashMLASchedMeta
    # per present DeepseekV4 layer type, shared across all ~60 layers of that type
    # within a decode step. The first forward call of a given type triggers
    # the in-kernel planner (which also allocates tile_scheduler_metadata and
    # num_splits via PyTorch's graph-aware allocator); subsequent same-type
    # calls skip planning and reuse the plan. Fresh instance per build(), so
    # have_initialized is always False at the start of a step and the plan
    # is re-derived from current seq_lens / topk_length on replay.
    # None for layer types the model does not use (or when num_decode_tokens
    # is zero).
    tile_sched_swaonly: "FlashMLASchedMeta | None" = None
    tile_sched_c4a: "FlashMLASchedMeta | None" = None
    tile_sched_c128a: "FlashMLASchedMeta | None" = None
    flashinfer_sparse_index_cache: dict[str, tuple[torch.Tensor, torch.Tensor]] = field(
        default_factory=dict
    )

    def get_prefill_chunk_plan(
        self, compress_ratio: int, prefill_chunk_size: int
    ) -> list[tuple[int, int, int, int]]:
        if self.num_prefills == 0:
            return []

        assert self.prefill_seq_lens_cpu is not None
        assert self.prefill_query_lens_cpu is not None

        # query_len <= max_num_batched_tokens and
        # gather_len = query_len + min(prefix_len, window_size - 1), so the
        # worst-case gathered width is bounded by
        # max_num_batched_tokens + window_size - 1. The compressed prefix pool
        # is bounded by ceil(max_model_len / compress_ratio).
        max_workspace_area = prefill_chunk_size * (
            (
                0
                if compress_ratio <= 1
                else cdiv(self.prefill_max_model_len, compress_ratio)
            )
            + self.prefill_window_size
            + self.prefill_max_num_batched_tokens
        )
        prefix_lens_cpu = self.prefill_seq_lens_cpu - self.prefill_query_lens_cpu
        gather_lens_cpu = self.prefill_query_lens_cpu + torch.clamp(
            prefix_lens_cpu, min=0, max=self.prefill_window_size - 1
        )
        compressed_lens_cpu = (
            torch.zeros_like(self.prefill_seq_lens_cpu)
            if compress_ratio <= 1
            else torch.div(
                self.prefill_seq_lens_cpu,
                compress_ratio,
                rounding_mode="floor",
            )
        )

        chunk_plan: list[tuple[int, int, int, int]] = []
        chunk_start = 0
        while chunk_start < self.num_prefills:
            chunk_max_compressed = int(compressed_lens_cpu[chunk_start].item())
            chunk_max_gather = int(gather_lens_cpu[chunk_start].item())
            chunk_end = chunk_start + 1

            while chunk_end < self.num_prefills:
                candidate_max_compressed = max(
                    chunk_max_compressed,
                    int(compressed_lens_cpu[chunk_end].item()),
                )
                candidate_max_gather = max(
                    chunk_max_gather,
                    int(gather_lens_cpu[chunk_end].item()),
                )
                candidate_width = candidate_max_compressed + candidate_max_gather
                candidate_area = (chunk_end - chunk_start + 1) * candidate_width
                if candidate_area > max_workspace_area:
                    break
                chunk_max_compressed = candidate_max_compressed
                chunk_max_gather = candidate_max_gather
                chunk_end += 1

            chunk_plan.append(
                (
                    chunk_start,
                    chunk_end,
                    chunk_max_compressed,
                    chunk_max_compressed + chunk_max_gather,
                )
            )
            chunk_start = chunk_end

        return chunk_plan


class ComputePrefillMetadataKernel(
    VllmTritonJitKernel["ComputePrefillMetadataKernel.CompileKey"]
):
    @dataclass(frozen=True)
    class CompileKey:
        block_size: int

    @staticmethod
    @triton.jit(do_not_specialize=["num_prefills", "num_decodes", "window_size"])
    def kernel(
        # Outputs
        prefill_gather_lens_ptr,
        # Inputs
        seq_lens_ptr,
        query_start_loc_ptr,
        num_prefills,
        num_decodes,
        window_size,
        BLOCK_SIZE: tl.constexpr,
    ):
        """Compute prefill gather_lens in a single pass."""
        offset = tl.arange(0, BLOCK_SIZE)
        mask = offset < num_prefills
        # SM12x + Triton 3.6 raises IMA on out-of-bounds address arithmetic for
        # masked-off lanes even though the load mask gates the actual read, so
        # clamp the offset. Caller guarantees num_prefills > 0.
        safe_offset = tl.minimum(offset, num_prefills - 1)

        seq_len = tl.load(seq_lens_ptr + num_decodes + safe_offset, mask=mask)
        qsl_start = tl.load(query_start_loc_ptr + num_decodes + safe_offset, mask=mask)
        qsl_end = tl.load(
            query_start_loc_ptr + num_decodes + safe_offset + 1, mask=mask
        )

        query_len = qsl_end - qsl_start
        prefix_len = seq_len - query_len
        gather_len = query_len + tl.minimum(prefix_len, window_size - 1)

        tl.store(prefill_gather_lens_ptr + offset, gather_len, mask=mask)

    def dispatch(  # type: ignore[override]
        self,
        *,
        num_prefills: int,
    ) -> CompileKey:
        return self.CompileKey(
            block_size=next_power_of_2(num_prefills),
        )

    def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]:
        scheduler_config = vllm_config.scheduler_config
        max_prefills = max(
            1,
            min(
                scheduler_config.max_num_seqs,
                scheduler_config.max_num_batched_tokens,
            ),
        )
        return self._trace_dispatch(self.dispatch)(
            num_prefills=WarmupIntRange(1, max_prefills + 1),
        )

    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
        int32_ptr = TritonWarmupTensor(torch.int32)
        return dict(
            prefill_gather_lens=int32_ptr,
            seq_lens=int32_ptr,
            query_start_loc=int32_ptr,
            num_prefills=compile_key.block_size,
            num_decodes=0,
            window_size=1,
        )

    @kernel_launcher
    def __call__(
        self,
        prefill_gather_lens: torch.Tensor,
        seq_lens: torch.Tensor,
        query_start_loc: torch.Tensor,
        num_prefills: int,
        num_decodes: int,
        window_size: int,
    ) -> LaunchSpec:
        compile_key = self.dispatch(num_prefills=num_prefills)
        return (1,), dict(
            BLOCK_SIZE=compile_key.block_size,
        )


class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
    """Builds metadata for DeepseekV4 SWA cache.

    Similar to the indexer, this handles mixed batches by:
    1. Using split_decodes_and_prefills() to determine the boundary
    2. Building separate metadata for decode and prefill portions

    Supports:
    - Mixed decode/prefill batches
    - MTP (Multi-Token Prediction) where decode has query_len > 1
    - Chunked prefill (aligns with the indexer's chunking)
    """

    reorder_batch_threshold: int | None = None
    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
    supports_draft_decode_metadata_update = True

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        assert isinstance(self.kv_cache_spec, SlidingWindowMLASpec | MLAAttentionSpec)
        mla_spec = cast(SlidingWindowMLASpec | MLAAttentionSpec, self.kv_cache_spec)
        self.head_size = mla_spec.head_size  # Already considered quantization.
        assert isinstance(mla_spec.tokens_per_state, int)
        self.compress_ratio = mla_spec.tokens_per_state
        self.block_size = mla_spec.block_size
        self.max_model_len = self.vllm_config.model_config.max_model_len
        self.max_num_batched_tokens = (
            self.vllm_config.scheduler_config.max_num_batched_tokens
        )

        # Handle MTP: adjust decode_threshold like the indexer does
        spec_config = self.vllm_config.speculative_config
        self.num_speculative_tokens = (
            spec_config.num_speculative_tokens if spec_config else 0
        )
        # Decode can have query_len up to
        #   1 + (2 if parallel drafting else 1) * num_speculative_tokens.
        # sparse_swa has no MQA-vs-dense-MHA routing, so multi-token queries take
        # the prefill path and the decode/prefill split stays at that width.
        spec_mult = (
            2 if (spec_config is not None and spec_config.parallel_drafting) else 1
        )
        self.decode_threshold = 1 + spec_mult * self.num_speculative_tokens
        self.reorder_batch_threshold = None

        hf_config = self.vllm_config.model_config.hf_config
        assert hasattr(hf_config, "sliding_window")
        self.window_size = hf_config.sliding_window

        # Vision variant: image spans (up to vision_max_n_token tokens) are
        # visible bidirectionally, so prefill index rows widen from
        # window_size to window_size + max_image_tokens. Text-only models keep
        # max_image_tokens == 0 and take the original code paths everywhere.
        self.max_image_tokens = (
            getattr(hf_config, "vision_max_n_token", 0)
            if getattr(hf_config, "vision_n_layers", 0) > 0
            else 0
        )
        self.prefill_index_width = self.window_size + self.max_image_tokens

        # Detect which DeepseekV4 layer types this model uses so we only build a
        # FlashMLA tile-scheduler plan for types that will actually be called.
        # Models without compress_ratios (pure SWA) fall back to swaonly.
        compress_ratios = getattr(hf_config, "compress_ratios", None) or [1]
        self._layer_types: set[str] = set()
        for ratio in compress_ratios:
            self._layer_types.add(_layer_type_for(int(ratio)))

        max_tokens = self.vllm_config.scheduler_config.max_num_batched_tokens
        self.token_to_req_indices = torch.zeros(
            max_tokens,
            dtype=torch.int32,
            device=self.device,
        )
        self.decode_swa_indices = torch.zeros(
            max_tokens,
            1,
            self.window_size,
            dtype=torch.int32,
            device=self.device,
        )
        self.decode_swa_lens = torch.zeros(
            max_tokens,
            dtype=torch.int32,
            device=self.device,
        )
        # Allocated unconditionally — consumer picks paged-direct vs dequant
        # at call time.
        self.prefill_swa_indices = torch.zeros(
            max_tokens,
            1,
            self.prefill_index_width,
            dtype=torch.int32,
            device=self.device,
        )
        self.prefill_swa_lens = torch.zeros(
            max_tokens,
            dtype=torch.int32,
            device=self.device,
        )
        # In-image visibility side buffers (vision variant only). The span
        # CSR is uploaded per step from mm_req_doc_ranges; span capacity is
        # bounded by max_tokens since each span holds at least one token.
        if self.max_image_tokens > 0:
            self.left_visible = torch.zeros(
                max_tokens, dtype=torch.int32, device=self.device
            )
            self.right_visible = torch.zeros(
                max_tokens, dtype=torch.int32, device=self.device
            )
            self.span_indptr = torch.zeros(
                self.vllm_config.scheduler_config.max_num_seqs + 1,
                dtype=torch.int32,
                device=self.device,
            )
            self.span_starts = torch.zeros(
                max_tokens, dtype=torch.int32, device=self.device
            )
            self.span_ends = torch.zeros(
                max_tokens, dtype=torch.int32, device=self.device
            )
        self.is_valid_token = torch.zeros(
            max_tokens,
            dtype=torch.bool,
            device=self.device,
        )

        # DSpark draft: the block is non-causal (every query attends to the
        # trailing window of context PLUS all query tokens, including future ones),
        # so its per-token index list is wider than `window_size`. The kernel pads
        # the q-head count to B_TOPK. Pad to a kernel-supported width; the logical
        # SWA window remains unchanged when the padded matrix is built.
        self.is_dspark = spec_config is not None and spec_config.use_dspark()
        self.noncausal_index_width = (
            get_dspark_swa_index_width(
                self.window_size,
                self.num_speculative_tokens,
            )
            if self.is_dspark
            else 0
        )
        self.decode_swa_indices_noncausal: torch.Tensor | None = None
        self._max_tokens = max_tokens

    def build(
        self,
        common_prefix_len: int,
        common_attn_metadata: CommonAttentionMetadata,
        fast_build: bool = False,
    ) -> DeepseekSparseSWAMetadata:
        """Build SWA metadata for mixed decode/prefill batches.

        The batch is assumed to be reordered with decodes first (by vLLM scheduler).
        We use split_decodes_and_prefills() to find the boundary, then build
        separate window_topk_idxs for each portion.

        For prefill, we use chunked prefill to align with the indexer's chunking.
        """
        seq_lens = common_attn_metadata.seq_lens
        seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
        query_start_loc = common_attn_metadata.query_start_loc
        query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
        block_table = common_attn_metadata.block_table_tensor
        slot_mapping = common_attn_metadata.slot_mapping

        # Split into decode and prefill portions using configurable threshold
        (num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens) = (
            split_decodes_and_prefills(
                common_attn_metadata, decode_threshold=self.decode_threshold
            )
        )

        # NOTE: Ensure all metadata tensors maintain fixed memory addresses
        # for CUDA graph compatibility.
        token_to_req_indices = common_attn_metadata.token_to_req_indices(
            self.token_to_req_indices
        )

        is_valid_token = self.is_valid_token[: slot_mapping.shape[0]]
        is_valid_token.copy_(slot_mapping >= 0)

        non_causal = not common_attn_metadata.causal
        decode_swa_width = (
            self.noncausal_index_width if non_causal else self.window_size
        )
        decode_swa_indices = self.decode_swa_indices
        if num_decode_tokens > 0:
            self.decode_swa_lens[num_decode_tokens:] = 0
            if non_causal:
                assert self.is_dspark, (
                    "Non-causal DeepseekV4 SWA is only supported for the DSpark "
                    "speculation mode, but causal=False was set without DSpark."
                )
                if self.decode_swa_indices_noncausal is None:
                    self.decode_swa_indices_noncausal = torch.zeros(
                        self._max_tokens,
                        1,
                        self.noncausal_index_width,
                        dtype=torch.int32,
                        device=self.device,
                    )
                decode_swa_indices = self.decode_swa_indices_noncausal
                _COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL(
                    decode_swa_indices,
                    self.decode_swa_lens,
                    self.window_size,
                    self.noncausal_index_width,
                    query_start_loc,
                    seq_lens,
                    token_to_req_indices,
                    is_valid_token,
                    block_table,
                    self.block_size,
                    num_tokens=num_decode_tokens,
                    token_offset=0,
                )
            else:
                _COMPUTE_SWA_INDICES_AND_LENS_KERNEL(
                    decode_swa_indices,
                    self.decode_swa_lens,
                    self.window_size,
                    decode_swa_indices.shape[-1],
                    self.decode_swa_lens,  # unused (HAS_IMAGE=False)
                    self.decode_swa_lens,  # unused (HAS_IMAGE=False)
                    query_start_loc,
                    seq_lens,
                    token_to_req_indices,
                    is_valid_token,
                    block_table,
                    self.block_size,
                    num_tokens=num_decode_tokens,
                    token_offset=0,
                )

        # Vision variant: per-token in-image visibility for prefill tokens.
        # Decode tokens are always past the image spans (spans are prefilled
        # atomically), so the decode path above never needs them.
        prefill_left_visible: torch.Tensor | None = None
        prefill_right_visible: torch.Tensor | None = None
        mm_ranges = common_attn_metadata.mm_req_doc_ranges
        if (
            self.max_image_tokens > 0
            and num_prefill_tokens > 0
            and mm_ranges
            and any(mm_ranges.values())
        ):
            prefill_left_visible, prefill_right_visible = self._build_image_visibility(
                common_attn_metadata.num_reqs,
                mm_ranges,
                num_decode_tokens,
                num_prefill_tokens,
                seq_lens,
                query_start_loc,
                token_to_req_indices,
            )

        # Prefill SWA indices live in paged coordinates. `token_offset` lets
        # the kernel read is_valid_token / token_to_req_indices at absolute
        # prefill positions while writing output starting at index 0.
        if num_prefill_tokens > 0:
            has_image = prefill_left_visible is not None
            prefill_swa_indices = self.prefill_swa_indices[:num_prefill_tokens]
            prefill_swa_lens = self.prefill_swa_lens[:num_prefill_tokens]
            _COMPUTE_SWA_INDICES_AND_LENS_KERNEL(
                prefill_swa_indices,
                prefill_swa_lens,
                self.window_size,
                self.prefill_index_width,
                prefill_left_visible if has_image else prefill_swa_lens,
                prefill_right_visible if has_image else prefill_swa_lens,
                query_start_loc,
                seq_lens,
                token_to_req_indices,
                is_valid_token,
                block_table,
                self.block_size,
                num_tokens=num_prefill_tokens,
                token_offset=num_decode_tokens,
                has_image=has_image,
            )

        # Pre-compute DeepseekV4 prefill metadata shared across all attention layers.
        deepseek_v4_fields = self._build_deepseek_v4_metadata(
            num_decodes,
            num_prefills,
            seq_lens,
            seq_lens_cpu,
            query_start_loc,
            query_start_loc_cpu,
        )

        # Per-layer-type tile-scheduler plan holders. Empty FlashMLASchedMeta
        # per present DeepseekV4 layer type; the first flash_mla_with_kvcache call of
        # each type triggers the planner and all same-type layers reuse the
        # resulting plan for the rest of the step.
        tile_sched = self.build_tile_scheduler(num_decode_tokens)

        return DeepseekSparseSWAMetadata(
            seq_lens=seq_lens,
            query_start_loc=query_start_loc,
            query_start_loc_cpu=query_start_loc_cpu,
            block_table=block_table,
            slot_mapping=slot_mapping,
            is_valid_token=is_valid_token,
            token_to_req_indices=token_to_req_indices,
            decode_swa_indices=decode_swa_indices[:num_decode_tokens],
            decode_swa_lens=self.decode_swa_lens[:num_decode_tokens],
            decode_swa_width=decode_swa_width,
            prefill_swa_indices=(
                self.prefill_swa_indices[:num_prefill_tokens]
                if num_prefill_tokens > 0
                else None
            ),
            prefill_swa_lens=(
                self.prefill_swa_lens[:num_prefill_tokens]
                if num_prefill_tokens > 0
                else None
            ),
            prefill_left_visible=prefill_left_visible,
            prefill_right_visible=prefill_right_visible,
            block_size=self.block_size,
            num_decodes=num_decodes,
            num_prefills=num_prefills,
            num_decode_tokens=num_decode_tokens,
            num_prefill_tokens=num_prefill_tokens,
            # Upper bound on decode-split rows for the kernel's max_q_len
            # hint. common max_query_len bounds every row (scheduled max under
            # adaptive verification), clamped to what the split can admit so a
            # mixed batch's prefill max does not inflate decode scheduling.
            max_decode_query_len=min(
                common_attn_metadata.max_query_len, self.decode_threshold
            ),
            tile_sched_swaonly=tile_sched[_LAYER_TYPE_SWAONLY],
            tile_sched_c4a=tile_sched[_LAYER_TYPE_C4A],
            tile_sched_c128a=tile_sched[_LAYER_TYPE_C128A],
            **deepseek_v4_fields,  # type: ignore[arg-type]
        )

    def _build_image_visibility(
        self,
        num_reqs: int,
        mm_ranges: dict[int, list[tuple[int, int]]],
        num_decode_tokens: int,
        num_prefill_tokens: int,
        seq_lens: torch.Tensor,
        query_start_loc: torch.Tensor,
        token_to_req_indices: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Upload image spans and compute per-token in-image visibility.

        Returns (left, right) int32 tensors indexed by absolute (decode-first)
        token position; only prefill rows are meaningful.
        """
        indptr = [0] * (num_reqs + 1)
        starts: list[int] = []
        ends: list[int] = []
        for req_idx in range(num_reqs):
            for span_start, span_end in mm_ranges.get(req_idx, ()):
                starts.append(span_start)
                ends.append(span_end)
            indptr[req_idx + 1] = len(starts)
        num_spans = len(starts)
        assert num_spans <= self.span_starts.shape[0]
        self.span_indptr[: num_reqs + 1].copy_(torch.tensor(indptr, dtype=torch.int32))
        if num_spans > 0:
            self.span_starts[:num_spans].copy_(torch.tensor(starts, dtype=torch.int32))
            self.span_ends[:num_spans].copy_(torch.tensor(ends, dtype=torch.int32))
        _compute_image_visibility_kernel[(num_prefill_tokens,)](
            self.left_visible,
            self.right_visible,
            self.span_indptr,
            self.span_starts,
            self.span_ends,
            query_start_loc,
            seq_lens,
            token_to_req_indices,
            self.max_image_tokens,
            token_offset=num_decode_tokens,
        )
        end = num_decode_tokens + num_prefill_tokens
        return self.left_visible[:end], self.right_visible[:end]

    def update_draft_decode_metadata(
        self,
        metadata: DeepseekSparseSWAMetadata,
    ) -> None:
        if metadata.num_decode_tokens == 0:
            return
        assert metadata.query_start_loc is not None
        assert metadata.seq_lens is not None
        assert metadata.token_to_req_indices is not None
        assert metadata.is_valid_token is not None
        assert metadata.decode_swa_indices is not None
        assert metadata.decode_swa_lens is not None

        _COMPUTE_SWA_INDICES_AND_LENS_KERNEL(
            metadata.decode_swa_indices,
            metadata.decode_swa_lens,
            metadata.decode_swa_indices.shape[-1],
            metadata.decode_swa_indices.shape[-1],
            metadata.decode_swa_lens,  # unused (HAS_IMAGE=False)
            metadata.decode_swa_lens,  # unused (HAS_IMAGE=False)
            metadata.query_start_loc,
            metadata.seq_lens,
            metadata.token_to_req_indices,
            metadata.is_valid_token,
            metadata.block_table,
            self.block_size,
            num_tokens=metadata.num_decode_tokens,
            token_offset=0,
        )
        tile_sched = self.build_tile_scheduler(metadata.num_decode_tokens)
        metadata.tile_sched_swaonly = tile_sched[_LAYER_TYPE_SWAONLY]
        metadata.tile_sched_c4a = tile_sched[_LAYER_TYPE_C4A]
        metadata.tile_sched_c128a = tile_sched[_LAYER_TYPE_C128A]
        metadata.flashinfer_sparse_index_cache.clear()

    def build_tile_scheduler(
        self, num_decode_tokens: int
    ) -> dict[str, FlashMLASchedMeta | None]:
        """Allocate one empty ``FlashMLASchedMeta`` per present DeepseekV4 layer type.

        Returned instances have ``tile_scheduler_metadata`` / ``num_splits``
        set to ``None``; the FlashMLA C++ decode path will allocate them and
        run the tile-scheduler planner on the first ``flash_mla_with_kvcache``
        call of each type. Subsequent same-type calls reuse the plan because
        the tensors (and ``have_initialized``) are populated on the struct.

        Returns all-``None`` when there are no decode tokens this step, so
        ``_forward_decode`` sees a clean sentinel.
        """
        out: dict[str, FlashMLASchedMeta | None] = {
            _LAYER_TYPE_SWAONLY: None,
            _LAYER_TYPE_C4A: None,
            _LAYER_TYPE_C128A: None,
        }
        if (
            num_decode_tokens == 0
            or current_platform.is_rocm()
            or current_platform.is_xpu()
            or current_platform.is_device_capability_family(120)
        ):
            return out
        for layer_type in self._layer_types:
            # get_mla_metadata() is the official FlashMLA entry point that
            # returns a fresh empty FlashMLASchedMeta; using it keeps this
            # call site aligned with the rest of the vLLM FlashMLA backends
            # that already go through the same stub.
            out[layer_type] = get_mla_metadata()[0]
        return out

    def _build_deepseek_v4_metadata(
        self,
        num_decodes: int,
        num_prefills: int,
        seq_lens: torch.Tensor,
        seq_lens_cpu: torch.Tensor | None,
        query_start_loc: torch.Tensor,
        query_start_loc_cpu: torch.Tensor,
    ) -> dict[str, torch.Tensor | int | None]:
        """Pre-compute DeepseekV4 prefill metadata during the metadata build phase.

        Returns a dict of keyword arguments to pass to the
        DeepseekSparseSWAMetadata constructor.

        Note: C128A sparse metadata is computed by the FlashMLASparse builder
        (which owns the C128A block_table), not here.
        """
        result: dict[str, torch.Tensor | int | None] = {}

        # --- Prefill query metadata (single Triton kernel + CPU slicing) ---
        if num_prefills > 0:
            assert seq_lens_cpu is not None
            pfx_gather_lens = torch.empty(
                num_prefills, dtype=torch.int32, device=seq_lens.device
            )
            _COMPUTE_PREFILL_METADATA_KERNEL(
                pfx_gather_lens,
                seq_lens,
                query_start_loc,
                num_prefills,
                num_decodes,
                self.window_size,
            )

            result["prefill_seq_lens"] = seq_lens[num_decodes:]
            result["prefill_seq_lens_cpu"] = seq_lens_cpu[num_decodes:]
            result["prefill_gather_lens"] = pfx_gather_lens
            result["prefill_query_lens_cpu"] = (
                query_start_loc_cpu[num_decodes + 1 : num_decodes + num_prefills + 1]
                - query_start_loc_cpu[num_decodes : num_decodes + num_prefills]
            ).to(dtype=torch.int32)
            result["prefill_window_size"] = self.window_size
            result["prefill_max_model_len"] = self.max_model_len
            result["prefill_max_num_batched_tokens"] = self.max_num_batched_tokens

        return result


@triton.jit(do_not_specialize=["token_offset"])
def _compute_image_visibility_kernel(
    left_visible_ptr,
    right_visible_ptr,
    span_indptr_ptr,
    span_starts_ptr,
    span_ends_ptr,
    query_start_loc_ptr,
    seq_lens_ptr,
    token_to_req_indices_ptr,
    max_image_tokens,
    token_offset,
):
    """Per-token in-image visible counts (port of `get_image_visible`).

    One program per prefill token. A token at position pos inside a span
    [span_start, span_end] sees min(pos - span_start, max_image_tokens - 1)
    extra tokens to its left and min(span_end - pos, max_image_tokens) to its
    right; tokens outside every span get 0/0 (plain causal window).
    """
    pid = tl.program_id(0)
    token_idx = pid + token_offset
    req_idx = tl.load(token_to_req_indices_ptr + token_idx)

    query_start = tl.load(query_start_loc_ptr + req_idx)
    query_end = tl.load(query_start_loc_ptr + req_idx + 1)
    seq_len = tl.load(seq_lens_ptr + req_idx)
    pos = seq_len - (query_end - query_start) + token_idx - query_start

    span_lo = tl.load(span_indptr_ptr + req_idx)
    span_hi = tl.load(span_indptr_ptr + req_idx + 1)
    left = tl.zeros((), dtype=tl.int32)
    right = tl.zeros((), dtype=tl.int32)
    for i in range(span_lo, span_hi):
        span_start = tl.load(span_starts_ptr + i)
        span_end = tl.load(span_ends_ptr + i)
        in_span = (pos >= span_start) & (pos <= span_end)
        left = tl.where(
            in_span, tl.minimum(pos - span_start, max_image_tokens - 1), left
        )
        right = tl.where(in_span, tl.minimum(span_end - pos, max_image_tokens), right)
    tl.store(left_visible_ptr + token_idx, left)
    tl.store(right_visible_ptr + token_idx, right)


# TODO(ben): unify this kernel to reduce duplication


@triton.jit(
    do_not_specialize=[
        "swa_indices_stride",
        "block_table_stride",
        "token_offset",
    ]
)
def _compute_swa_indices_and_lens_kernel(
    swa_indices_ptr,
    swa_indices_stride,
    swa_lens_ptr,
    window_size,
    index_width,
    left_visible_ptr,
    right_visible_ptr,
    query_start_loc_ptr,
    seq_lens_ptr,
    token_to_req_indices_ptr,
    is_valid_token_ptr,
    block_table_ptr,
    block_table_stride,
    block_size,
    token_offset,
    HAS_IMAGE: tl.constexpr,
    TRITON_BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(0)
    token_idx = pid + token_offset
    is_valid = tl.load(is_valid_token_ptr + token_idx)
    if not is_valid:
        tl.store(swa_lens_ptr + pid, 0)
        # Clear the row so a padded token cannot gather through stale indices.
        for i in range(0, index_width, TRITON_BLOCK_SIZE):
            offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
            tl.store(
                swa_indices_ptr + pid * swa_indices_stride + offset,
                -1,
                mask=offset < index_width,
            )
        return

    req_idx = tl.load(token_to_req_indices_ptr + token_idx)

    query_start = tl.load(query_start_loc_ptr + req_idx)
    query_end = tl.load(query_start_loc_ptr + req_idx + 1)
    query_len = query_end - query_start

    seq_len = tl.load(seq_lens_ptr + req_idx)
    prefix_len = seq_len - query_len

    pos = prefix_len + token_idx - query_start
    if HAS_IMAGE:
        # In-image bidirectional visibility widens the window: the window
        # starts up to max(left - (window - 1), 0) positions earlier and
        # extends `right` positions past the query token.
        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)
    start_pos = tl.maximum(pos - (window_size - 1) - left_add, 0)
    end_pos = pos + right + 1

    swa_len = end_pos - start_pos
    tl.store(swa_lens_ptr + pid, swa_len)

    for i in range(0, index_width, TRITON_BLOCK_SIZE):
        offset = i + tl.arange(0, TRITON_BLOCK_SIZE)

        pos_offset = start_pos + offset
        block_indices = pos_offset // block_size
        block_numbers = tl.load(
            block_table_ptr + req_idx * block_table_stride + block_indices,
            mask=pos_offset < end_pos,
        )
        block_offsets = pos_offset % block_size
        slot_ids = block_numbers * block_size + block_offsets

        slot_ids = tl.where(offset < swa_len, slot_ids, -1)
        tl.store(
            swa_indices_ptr + pid * swa_indices_stride + offset,
            slot_ids,
            mask=offset < index_width,
        )


class ComputeSWAIndicesAndLensKernel(
    VllmTritonJitKernel["ComputeSWAIndicesAndLensKernel.CompileKey"]
):
    @dataclass(frozen=True)
    class CompileKey:
        window_size: int
        index_width: int
        has_image: int
        block_size: int
        triton_block_size: int

    kernel = staticmethod(_compute_swa_indices_and_lens_kernel)

    def dispatch(  # type: ignore[override]
        self,
        **compile_key_fields: int,
    ) -> CompileKey:
        return self.CompileKey(
            **compile_key_fields,
            triton_block_size=1024,
        )

    def get_warmup_keys(
        self,
        *,
        window_size: int,
        block_size: int,
        max_image_tokens: int = 0,
    ) -> list[CompileKey]:
        # Decode rows always use the plain causal window. Vision models widen
        # prefill rows by max_image_tokens and additionally need the in-image
        # bidirectional variant for batches that contain image spans.
        width_has_image = [(window_size, 0)]
        if max_image_tokens > 0:
            width_has_image += [
                (window_size + max_image_tokens, 0),
                (window_size + max_image_tokens, 1),
            ]
        return self._trace_dispatch(self.dispatch)(
            zip_inputs(
                *(
                    dict(index_width=index_width, has_image=has_image)
                    for index_width, has_image in width_has_image
                )
            ),
            window_size=window_size,
            block_size=block_size,
        )

    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
        int32_ptr = TritonWarmupTensor(torch.int32)
        return dict(
            swa_indices=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)),
            swa_lens=int32_ptr,
            window_size=compile_key.window_size,
            index_width=compile_key.index_width,
            left_visible=int32_ptr,
            right_visible=int32_ptr,
            query_start_loc=int32_ptr,
            seq_lens=int32_ptr,
            token_to_req_indices=int32_ptr,
            is_valid_token=TritonWarmupTensor(torch.bool),
            block_table=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)),
            block_size=compile_key.block_size,
            num_tokens=1,
            token_offset=0,
            has_image=bool(compile_key.has_image),
        )

    @kernel_launcher
    def __call__(
        self,
        swa_indices: torch.Tensor,
        swa_lens: torch.Tensor,
        window_size: int,
        index_width: int,
        left_visible: torch.Tensor,
        right_visible: torch.Tensor,
        query_start_loc: torch.Tensor,
        seq_lens: torch.Tensor,
        token_to_req_indices: torch.Tensor,
        is_valid_token: torch.Tensor,
        block_table: torch.Tensor,
        block_size: int,
        *,
        num_tokens: int,
        token_offset: int,
        has_image: bool = False,
    ) -> LaunchSpec:
        return (num_tokens,), dict(
            swa_indices_stride=swa_indices.stride(0),
            block_table_stride=block_table.stride(0),
            HAS_IMAGE=has_image,
            TRITON_BLOCK_SIZE=1024,
        )


class ComputeDSparkNoncausalSWAIndicesKernel(
    VllmTritonJitKernel["ComputeDSparkNoncausalSWAIndicesKernel.CompileKey"]
):
    @dataclass(frozen=True)
    class CompileKey:
        window_size: int
        index_width: int
        block_size: int
        triton_block_size: int

    @staticmethod
    @triton.jit(
        do_not_specialize=[
            "swa_indices_stride",
            "block_table_stride",
            "token_offset",
        ]
    )
    def kernel(
        swa_indices_ptr,
        swa_indices_stride,
        swa_lens_ptr,
        window_size,
        index_width,
        query_start_loc_ptr,
        seq_lens_ptr,
        token_to_req_indices_ptr,
        is_valid_token_ptr,
        block_table_ptr,
        block_table_stride,
        block_size,
        token_offset,
        TRITON_BLOCK_SIZE: tl.constexpr,
    ):
        """Non-causal per-token indices for the DSpark draft block.

        Here, we populate the topk indices with the trailing window of context tokens,
        plus all query tokens (including future ones).
        """
        pid = tl.program_id(0)
        token_idx = pid + token_offset
        is_valid = tl.load(is_valid_token_ptr + token_idx)
        if not is_valid:
            tl.store(swa_lens_ptr + pid, 0)
            # Clear the row so a padded token cannot gather through stale indices.
            for i in range(0, index_width, TRITON_BLOCK_SIZE):
                offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
                tl.store(
                    swa_indices_ptr + pid * swa_indices_stride + offset,
                    -1,
                    mask=offset < index_width,
                )
            return

        req_idx = tl.load(token_to_req_indices_ptr + token_idx)

        query_start = tl.load(query_start_loc_ptr + req_idx)
        query_end = tl.load(query_start_loc_ptr + req_idx + 1)
        query_len = query_end - query_start

        seq_len = tl.load(seq_lens_ptr + req_idx)
        prefix_len = seq_len - query_len

        # Block-anchored window (shared by every token in the block) + full block.
        start_pos = tl.maximum(prefix_len - window_size, 0)
        end_pos = seq_len

        swa_len = end_pos - start_pos
        tl.store(swa_lens_ptr + pid, swa_len)

        for i in range(0, index_width, TRITON_BLOCK_SIZE):
            offset = i + tl.arange(0, TRITON_BLOCK_SIZE)

            pos_offset = start_pos + offset
            block_indices = pos_offset // block_size
            block_numbers = tl.load(
                block_table_ptr + req_idx * block_table_stride + block_indices,
                mask=pos_offset < end_pos,
            )
            block_offsets = pos_offset % block_size
            slot_ids = block_numbers * block_size + block_offsets

            slot_ids = tl.where(offset < swa_len, slot_ids, -1)
            tl.store(
                swa_indices_ptr + pid * swa_indices_stride + offset,
                slot_ids,
                mask=offset < index_width,
            )

    def dispatch(  # type: ignore[override]
        self,
        *,
        window_size: int,
        num_speculative_tokens: int,
        block_size: int,
    ) -> CompileKey:
        return self.CompileKey(
            window_size=window_size,
            index_width=get_dspark_swa_index_width(window_size, num_speculative_tokens),
            block_size=block_size,
            triton_block_size=1024,
        )

    def get_warmup_keys(
        self,
        *,
        window_size: int,
        num_speculative_tokens: int,
        block_size: int,
    ) -> list[CompileKey]:
        return self._trace_dispatch(self.dispatch)(
            window_size=window_size,
            num_speculative_tokens=num_speculative_tokens,
            block_size=block_size,
        )

    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
        int32_ptr = TritonWarmupTensor(torch.int32)
        return dict(
            swa_indices=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)),
            swa_lens=int32_ptr,
            window_size=compile_key.window_size,
            index_width=compile_key.index_width,
            query_start_loc=int32_ptr,
            seq_lens=int32_ptr,
            token_to_req_indices=int32_ptr,
            is_valid_token=TritonWarmupTensor(torch.bool),
            block_table=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)),
            block_size=compile_key.block_size,
            num_tokens=1,
            token_offset=0,
        )

    @kernel_launcher
    def __call__(
        self,
        swa_indices: torch.Tensor,
        swa_lens: torch.Tensor,
        window_size: int,
        index_width: int,
        query_start_loc: torch.Tensor,
        seq_lens: torch.Tensor,
        token_to_req_indices: torch.Tensor,
        is_valid_token: torch.Tensor,
        block_table: torch.Tensor,
        block_size: int,
        *,
        num_tokens: int,
        token_offset: int,
    ) -> LaunchSpec:
        return (num_tokens,), dict(
            swa_indices_stride=swa_indices.stride(0),
            block_table_stride=block_table.stride(0),
            TRITON_BLOCK_SIZE=1024,
        )


_COMPUTE_PREFILL_METADATA_KERNEL = ComputePrefillMetadataKernel()
_COMPUTE_SWA_INDICES_AND_LENS_KERNEL = ComputeSWAIndicesAndLensKernel()
_COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL = ComputeDSparkNoncausalSWAIndicesKernel()