recipes/docker-images/glm52-nf3-hybrid/overlay/vllm/v1/attention/backends/mla/indexer.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import os
from dataclasses import dataclass

import numpy as np
import torch

import vllm.envs as envs
from vllm.config import VllmConfig
from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.triton_utils import tl, triton
from vllm.utils.math_utils import cdiv
from vllm.utils.platform_utils import num_compute_units
from vllm.v1.attention.backend import (
    AttentionBackend,
    AttentionCGSupport,
    AttentionMetadataBuilder,
    CommonAttentionMetadata,
    MultipleOf,
)
from vllm.v1.attention.backends.mla.compressor_utils import get_compressed_slot_mapping
from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
from vllm.v1.kv_cache_interface import AttentionSpec, MLAAttentionSpec
from vllm.v1.worker.cp_utils import get_total_cp_world_size

logger = init_logger(__name__)

_B12X_PAGED_INDEX_SUPERTILE_K_DEFAULT = 32768
_B12X_PAGED_INDEX_TILE_BLOCK_K = 512


def _get_b12x_paged_indexer_supertile_k() -> int:
    raw = os.environ.get("B12X_PAGED_INDEX_SUPERTILE_K")
    tokens = _B12X_PAGED_INDEX_SUPERTILE_K_DEFAULT if raw is None else int(raw)
    tokens = max(tokens, _B12X_PAGED_INDEX_TILE_BLOCK_K)
    return (
        (tokens + _B12X_PAGED_INDEX_TILE_BLOCK_K - 1)
        // _B12X_PAGED_INDEX_TILE_BLOCK_K
        * _B12X_PAGED_INDEX_TILE_BLOCK_K
    )


@triton.jit
def _prepare_uniform_decode_kernel(
    seq_lens_ptr,
    decode_seq_lens_ptr,
    block_table_ptr,
    block_table_stride,
    expanded_block_table_ptr,
    expanded_bt_stride,
    decode_lens_ptr,
    max_decode_len,
    BLOCK_SIZE: tl.constexpr,
):
    idx = tl.program_id(0)
    req_id = idx // max_decode_len
    local_idx = idx % max_decode_len

    # Compute number of KVs attended to by this token.
    seq_len = tl.load(seq_lens_ptr + req_id)
    per_token_seq_len = seq_len - max_decode_len + local_idx + 1
    tl.store(decode_seq_lens_ptr + idx, per_token_seq_len)

    # Copy block table row.
    src = block_table_ptr + req_id * block_table_stride
    dst = expanded_block_table_ptr + idx * expanded_bt_stride
    for i in tl.range(0, expanded_bt_stride, BLOCK_SIZE):
        off = i + tl.arange(0, BLOCK_SIZE)
        mask = off < expanded_bt_stride
        src_block = tl.load(src + off, mask=mask)
        tl.store(dst + off, src_block, mask=mask)

    # All reqs now have decode_len = 1.
    tl.store(decode_lens_ptr + idx, 1)


def split_indexer_prefill_chunks(
    seq_lens_cpu: torch.Tensor | np.ndarray,
    query_lens_cpu: torch.Tensor | np.ndarray,
    workspace_size: int,
    max_logits_bytes: int,
    request_offset: int = 0,
) -> list[tuple[slice, slice]]:
    """
    Split prefill requests into chunks for the sparse indexer, respecting:
    - N constraint: total_seq_lens <= workspace_size (existing O(N) workspace)
    - Logits constraint: M * N * 4 <= max_logits_bytes

    When a single request-level chunk still exceeds the logits budget,
    sub-chunks on the query dimension (M) to bound peak memory.

    Returns list of (req_slice, query_slice) tuples.
    """
    if isinstance(seq_lens_cpu, torch.Tensor):
        seq_lens_cpu = seq_lens_cpu.numpy()
    if isinstance(query_lens_cpu, torch.Tensor):
        query_lens_cpu = query_lens_cpu.numpy()

    chunks: list[tuple[slice, slice]] = []
    n = len(seq_lens_cpu)
    max_logits_elems = max_logits_bytes // 4
    end = 0

    while end < n:
        start, chunk_m, chunk_n = end, 0, 0

        while end < n:
            q, s = int(query_lens_cpu[end]), int(seq_lens_cpu[end])
            new_m, new_n = chunk_m + q, chunk_n + s
            if new_n <= workspace_size and new_m * new_n <= max_logits_elems:
                chunk_m, chunk_n = new_m, new_n
                end += 1
            else:
                break

        # A single request can exceed the budget, requiring sub-chunking
        # on the query dimension.
        if end == start:
            chunk_m, chunk_n = int(query_lens_cpu[end]), int(seq_lens_cpu[end])
            end += 1

        req_slice = slice(start + request_offset, end + request_offset)
        max_q = max(1, max_logits_elems // chunk_n) if chunk_n > 0 else chunk_m
        for q_off in range(0, chunk_m, max_q):
            sub_m = min(max_q, chunk_m - q_off)
            chunks.append((req_slice, slice(q_off, q_off + sub_m)))

    return chunks


class DeepseekV32IndexerBackend(AttentionBackend):
    @staticmethod
    def get_name() -> str:
        return "DEEPSEEK_V32_INDEXER"

    @staticmethod
    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
        return [1, 64] if current_platform.is_rocm() else [64]

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

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

    @staticmethod
    def get_kv_cache_shape(
        num_blocks: int,
        block_size: int,
        num_kv_heads: int,
        head_size: int,
        cache_dtype_str: str = "auto",
    ) -> tuple[int, ...]:
        assert num_kv_heads == 1
        return (num_blocks, block_size, head_size)

    @staticmethod
    def get_kv_cache_stride_order(
        include_num_layers_dimension: bool = False,
    ) -> tuple[int, ...]:
        if include_num_layers_dimension:
            # DeepseekV32Indexer kernels do not support cross-layer
            # KV cache layout. Identity permutation keeps num_layers
            # first, signaling incompatibility.
            return (0, 1, 2, 3)
        return (0, 1, 2)


class B12xNonCompressedIndexerBackend(DeepseekV32IndexerBackend):
    @staticmethod
    def get_name() -> str:
        return "B12X_NON_COMPRESSED_INDEXER"


class DeepseekV4IndexerBackend(DeepseekV32IndexerBackend):
    @staticmethod
    def get_name() -> str:
        return "DEEPSEEK_V4_INDEXER"

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


@dataclass
class DeepseekV32IndexerPrefillChunkMetadata:
    block_table: torch.Tensor
    cu_seqlen_ks: torch.Tensor
    cu_seqlen_ke: torch.Tensor
    cu_seq_lens: torch.Tensor
    token_to_seq: torch.Tensor
    total_seq_lens: int
    token_start: int
    token_end: int
    num_reqs: int
    skip_kv_gather: bool = False


@dataclass
class DeepseekV32IndexerPrefillMetadata:
    chunks: list[DeepseekV32IndexerPrefillChunkMetadata]


@dataclass
class DeepSeekV32IndexerDecodeMetadata:
    block_table: torch.Tensor
    # seq_lens: per-token effective context lengths.
    #   - flatten path / plain decode: 1D (batch_size,)
    #   - native MTP path: 2D (B, next_n) where [b,j] = L_b - next_n + j + 1
    # Both fp8_fp4_paged_mqa_logits and the topk kernels accept both shapes.
    seq_lens: torch.Tensor
    max_seq_len: int
    decode_lens: torch.Tensor
    requires_padding: bool
    schedule_metadata: torch.Tensor | None
    compress_ratio: int = 1
    # Live scorer window (max compressed context across the batch) in cache
    # tokens, computed host-side in build() — a metadata tensor read by the
    # captured indexer kernel, never an in-kernel reduction. None => b12x uses
    # the capacity cap.
    active_width: torch.Tensor | None = None


@dataclass
class DeepseekV32IndexerMetadata:
    # FIXME (zyongye)
    # hacky way to access the data now, need to be in chunked meta
    seq_lens: torch.Tensor
    max_seq_len: int
    slot_mapping: torch.Tensor

    # New for MLA (compared to FlashAttention)
    # For handling prefill decode split
    num_decodes: int
    num_decode_tokens: int
    num_prefills: int
    num_prefill_tokens: int

    decode: DeepSeekV32IndexerDecodeMetadata | None = None
    prefill: DeepseekV32IndexerPrefillMetadata | None = None
    dcp_world_size: int = 1
    dcp_rank: int = 0
    cp_interleave_size: int = 1


def get_max_prefill_buffer_size(vllm_config: VllmConfig):
    max_model_len = vllm_config.model_config.max_model_len
    # NOTE(Chen): 40 is a magic number for controlling the prefill buffer size.
    # Each entry is 128 fp8 bytes and 4 scale bytes for a total of 132 bytes.
    # The flashmla_sparse backend uses a workspace size of 5 * max_model_len.
    # The memory usage of the workspace there is 576 * 2 bytes; so we size this as
    # (576 * 2 // 132) * 5 = 40 to maximize this workspace size while still fitting
    # within the flashmla_sparse workspace.
    # For DeepSeek-V3.2, the max_model_len is 163840.
    #   40 * 163840 * 132 = 865075200 bytes = 825 MB
    return max_model_len * 40


class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
    reorder_batch_threshold: int = 1

    @classmethod
    def get_cudagraph_support(
        cls,
        vllm_config: VllmConfig,
        kv_cache_spec: AttentionSpec,
    ) -> AttentionCGSupport:
        return AttentionCGSupport.UNIFORM_BATCH

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.storage_block_size = int(self.kv_cache_spec.storage_block_size)
        scheduler_config = self.vllm_config.scheduler_config
        parallel_config = self.vllm_config.parallel_config
        self.dcp_world_size = parallel_config.decode_context_parallel_size
        self.pcp_world_size = parallel_config.prefill_context_parallel_size
        self.cp_kv_cache_interleave_size = (
            parallel_config.cp_kv_cache_interleave_size
        )
        self.dcp_rank = 0
        if self.dcp_world_size > 1:
            assert self.pcp_world_size == 1, (
                "DeepseekV32IndexerMetadataBuilder supports DCP but not PCP."
            )
            from vllm.distributed.parallel_state import get_dcp_group

            self.dcp_rank = get_dcp_group().rank_in_group
        # NOTE(Chen):an estimated max size of flattened_kv. Need to double check.
        self.max_prefill_buffer_size = get_max_prefill_buffer_size(self.vllm_config)
        self.num_speculative_tokens = (
            self.vllm_config.speculative_config.num_speculative_tokens
            if self.vllm_config.speculative_config
            else 0
        )
        self.use_fp4_indexer_cache = (
            self.vllm_config.attention_config.use_fp4_indexer_cache
        )

        assert (
            current_platform.is_device_capability_family(100)
            or not self.use_fp4_indexer_cache
        ), (
            "use_fp4_indexer_cache requires Blackwell datacenter GPUs "
            "(sm_10x, e.g. B200/GB200); sm_120 (consumer Blackwell) and "
            "earlier architectures are not supported."
        )

        next_n = self.num_speculative_tokens + 1
        self.reorder_batch_threshold += self.num_speculative_tokens
        # NOTE: SM100 datacenter GPUs support any next_n natively via the
        # multi-atom paged MQA logits kernels (FP8 and FP4 indexer
        # caches). Outside the SM100 family the FP8
        # paged MQA logits kernel only supports next_n in (1, 2)
        # (deepgemm smxx_fp8_fp4_paged_mqa_logits.hpp:233), so flatten there.
        self.use_flattening = not current_platform.is_device_capability_family(
            100
        ) and next_n not in (1, 2)
        logger.info_once(
            "DSA indexer decode path: use_flattening=%s "
            "(next_n=%d, use_fp4_indexer_cache=%s)",
            self.use_flattening,
            next_n,
            self.use_fp4_indexer_cache,
        )

        sm_count = num_compute_units(self.device.index)
        self.num_sms = sm_count

        self.offsets_buffer = torch.arange(
            next_n, device=self.device, dtype=torch.int32
        )
        self.decode_lens_buffer = torch.zeros(
            (scheduler_config.max_num_batched_tokens,),
            dtype=torch.int32,
            device=self.device,
        )
        # Shared workspace for decode seq_lens. Native MTP views this as
        # (B, max_decode_len) at runtime, keeping context_lens contiguous even
        # when max_decode_len is smaller than next_n.
        self.decode_seq_lens_buffer = torch.zeros(
            (scheduler_config.max_num_batched_tokens,),
            dtype=torch.int32,
            device=self.device,
        )
        self.arange_buffer = torch.arange(
            max(
                scheduler_config.max_num_seqs * next_n,
                scheduler_config.max_num_batched_tokens,
            ),
            dtype=torch.int32,
            device=self.device,
        )
        max_num_blocks_per_req = cdiv(
            self.vllm_config.model_config.max_model_len,
            self.kv_cache_spec.block_size * get_total_cp_world_size(),
        )
        self.expanded_block_table_buffer = torch.zeros(
            (
                scheduler_config.max_num_batched_tokens,
                max_num_blocks_per_req + 8,  # reaper fix: runner tables run wider (461 vs 462 off-by-one)
            ),
            dtype=torch.int32,
            device=self.device,
        )

        # See: DeepGMM/csrc/apis/attention.hpp
        self.scheduler_metadata_buffer = torch.empty(
            (self.num_sms + 1, 2), dtype=torch.int32, device=self.device
        )

        # Persistent live-active-width buffer for the b12x indexer decode
        # scorer window. Filled host-side each build() (outside cudagraph
        # capture) and read by the captured kernel at a stable address.
        self.b12x_active_width_buffer = torch.zeros(
            (1,), dtype=torch.int32, device=self.device
        )

        # KV compression. Default to 1 for no compression.
        self.compress_ratio = 1
        # Get compress_ratio for DeepseekV4 support
        if isinstance(self.kv_cache_spec, MLAAttentionSpec):
            self.compress_ratio = self.kv_cache_spec.compress_ratio

        _prewarm_prefill_chunk_metadata_kernel(
            device=self.device,
            dcp_world_size=self.dcp_world_size,
            dcp_rank=self.dcp_rank,
            cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
            compress_ratio=self.compress_ratio,
        )

        # DCP writes the indexer KV cache through rank-local pages even when
        # compress_ratio == 1 (GLM/Kimi). Keep the mapped slots graph-stable.
        if self.compress_ratio > 1 or self.dcp_world_size > 1:
            self.compressed_slot_mapping_buffer = torch.zeros(
                (scheduler_config.max_num_batched_tokens,),
                dtype=torch.int64,
                device=self.device,
            )

        # Pre-allocate buffers for CUDA graph compatibility when
        if self.compress_ratio > 1:
            # compress_ratio > 1 (DeepseekV4)
            # Buffer for compressed seq_lens in decode path
            self.expanded_seq_lens_buffer = torch.zeros(
                (scheduler_config.max_num_batched_tokens,),
                dtype=torch.int32,
                device=self.device,
            )

    def _maybe_build_b12x_schedule_metadata(
        self,
        seq_lens: torch.Tensor,
        block_table: torch.Tensor,
        num_decode_tokens: int,
        requires_padding: bool,
    ) -> torch.Tensor | None:
        if not envs.VLLM_USE_B12X_SPARSE_INDEXER or requires_padding:
            return None

        schedule_seq_lens = seq_lens
        if schedule_seq_lens.dim() == 2:
            batch_size, next_n = schedule_seq_lens.shape
            if num_decode_tokens != int(batch_size * next_n):
                return None
            schedule_seq_lens = schedule_seq_lens.reshape(-1)
        if schedule_seq_lens.dim() != 1:
            return None

        from b12x.attention.indexer import (
            build_paged_mqa_schedule_metadata,
            uses_paged_mqa_schedule,
        )

        if not uses_paged_mqa_schedule(
            q_rows=int(schedule_seq_lens.shape[0]),
            max_pages=int(block_table.shape[1]),
        ):
            return None

        return build_paged_mqa_schedule_metadata(
            schedule_seq_lens.contiguous(),
            self.storage_block_size,
            self.num_sms,
            out=self.scheduler_metadata_buffer,
        )

    def _maybe_build_deep_gemm_schedule_metadata(
        self,
        seq_lens: torch.Tensor,
    ) -> torch.Tensor:
        if current_platform.is_cuda():
            from vllm.utils.deep_gemm import (
                get_paged_mqa_logits_metadata,
                has_deep_gemm,
            )

            if has_deep_gemm():
                self.scheduler_metadata_buffer[:] = get_paged_mqa_logits_metadata(
                    seq_lens,
                    self.storage_block_size,
                    self.num_sms,
                )
        return self.scheduler_metadata_buffer

    def _decode_topk_max_seq_len_from_cpu(
        self,
        common_attn_metadata: CommonAttentionMetadata,
        num_decodes: int,
        decode_lens_np: np.ndarray,
        max_decode_len: int,
        use_native: bool,
        dcp_local_seq_lens_cpu: torch.Tensor | None,
    ) -> int | None:
        """Return the exact decode max seq-len without reading CUDA tensors.

        B12X needs a host scalar for metadata/scheduling.  Computing it with
        ``seq_lens.max().item()`` synchronizes every decode step.  The model
        runner already maintains CPU shadows for normal decode; use those
        directly and mirror only the shape transforms that can affect the max.
        If async-spec has made the CPU shadow non-authoritative, return None so
        the caller can fall back to a safe graph-stable bound.
        """
        cpu_seq_lens = (
            dcp_local_seq_lens_cpu
            if dcp_local_seq_lens_cpu is not None
            else common_attn_metadata._seq_lens_cpu
        )
        if cpu_seq_lens is None:
            return None
        if cpu_seq_lens.device.type != "cpu":
            return None

        seq_lens_np = cpu_seq_lens[:num_decodes].numpy()
        if seq_lens_np.size == 0:
            return 0

        # Flattened variable-length MTP writes only real decode tokens and
        # zero-fills the tail.  Ignore padded decode rows to match that max.
        if not use_native and max_decode_len > 1:
            valid_decode_rows = decode_lens_np[:num_decodes] > 0
            if not bool(np.any(valid_decode_rows)):
                return 0
            seq_lens_np = seq_lens_np[valid_decode_rows]

        max_seq_len = int(seq_lens_np.max())
        if self.compress_ratio > 1:
            max_seq_len //= self.compress_ratio
        return max_seq_len

    def _prepare_decode_tensors(
        self,
        seq_lens: torch.Tensor,
        block_table: torch.Tensor,
        decode_lens: torch.Tensor,
        decode_lens_np: np.ndarray,
        query_start_loc: torch.Tensor,
        num_decodes: int,
        num_decode_tokens: int,
        use_native: bool,
        next_n: int,
        max_decode_len: int,
        global_seq_lens: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, bool]:
        """Expand seq_lens/block_table/decode_lens for the decode kernels.

        Flatten path (not use_native, max_decode_len > 1):
          Each multi-token decode request is expanded into individual
          single-token entries so the kernel always sees next_n=1.

        Native path (use_native or max_decode_len == 1):
          Plain decode or spec-decode with 2D per-token context lengths.

        Returns (seq_lens, block_table, decode_lens, batch_size, requires_padding).
        seq_lens is 1D (batch_size,) for flatten/plain, 2D (B, max_decode_len)
        for native MTP.
        """
        min_decode_len = int(decode_lens_np.min())
        if not use_native and max_decode_len > 1:
            assert self.decode_seq_lens_buffer.dim() == 1
            if min_decode_len == max_decode_len:
                # Uniform decode lengths.
                num_decode_tokens = num_decodes * max_decode_len
                _prepare_uniform_decode_kernel[(num_decode_tokens,)](
                    seq_lens,
                    self.decode_seq_lens_buffer,
                    block_table,
                    block_table.stride(0),
                    self.expanded_block_table_buffer,
                    self.expanded_block_table_buffer.stride(0),
                    self.decode_lens_buffer,
                    max_decode_len,
                    BLOCK_SIZE=1024,
                )
                if global_seq_lens is not None and self.dcp_world_size > 1:
                    expanded_global = (
                        global_seq_lens[:num_decodes].unsqueeze(1)
                        - max_decode_len
                        + 1
                        + self.offsets_buffer[:max_decode_len]
                    ).reshape(-1)
                    self.decode_seq_lens_buffer[:num_decode_tokens].copy_(
                        get_dcp_local_seq_lens(
                            expanded_global,
                            self.dcp_world_size,
                            self.dcp_rank,
                            self.cp_kv_cache_interleave_size,
                        ),
                        non_blocking=True,
                    )
                self.decode_seq_lens_buffer[num_decode_tokens:] = 0
                seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens]
                block_table = self.expanded_block_table_buffer[:num_decode_tokens]
                decode_lens = self.decode_lens_buffer[:num_decode_tokens]
                return seq_lens, block_table, decode_lens, num_decode_tokens, False
            else:
                # Variable decode lengths.
                # Assume 4 requests with seq_lens [10, 7, 12, 0] (the final req is
                # padding) and decode_lens [3, 1, 4, 0] in the below example comments.
                # The context lengths are therefore
                # [10-3, 7-1, 12-4, 0-0] = [7, 6, 8, 0].

                # 3 + 1 + 4 + 0 = 8
                actual_expanded = int(decode_lens_np.sum())

                # Fuse expanded_base and expanded_starts into a single
                # repeat_interleave:
                # seq_len_i = (context_start[b] - query_start_loc[b]) + arange[i] + 1
                # where context_start[b] = seq_lens[b] - decode_lens[b].
                # Example: offsets = [7-0, 6-3, 8-4, 0-8] = [7, 3, 4, -8]
                # expanded_offsets  = [7, 7, 7, 3, 4, 4, 4, 4]
                # result            = [8, 9, 10, 7, 9, 10, 11, 12]
                seq_lens_for_expansion = (
                    global_seq_lens if global_seq_lens is not None else seq_lens
                )
                expanded_offsets = torch.repeat_interleave(
                    seq_lens_for_expansion - decode_lens - query_start_loc,
                    decode_lens,
                    output_size=actual_expanded,
                )

                # [8, 9, 10, 7, 9, 10, 11, 12, ...] where ... is unused buffer space
                expanded_seq_lens = (
                    expanded_offsets + self.arange_buffer[:actual_expanded] + 1
                )
                if global_seq_lens is not None and self.dcp_world_size > 1:
                    expanded_seq_lens = get_dcp_local_seq_lens(
                        expanded_seq_lens,
                        self.dcp_world_size,
                        self.dcp_rank,
                        self.cp_kv_cache_interleave_size,
                    )
                self.decode_seq_lens_buffer[:actual_expanded] = expanded_seq_lens
                self.decode_seq_lens_buffer[actual_expanded:] = 0
                seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens]

                # Give each of the flattened entries the same block table row as the
                # original request.
                if int(block_table.shape[1]) > int(
                        self.expanded_block_table_buffer.shape[1]):
                    # reaper fix: grow-on-demand if a batch arrives wider than alloc
                    self.expanded_block_table_buffer = torch.zeros(
                        (int(self.expanded_block_table_buffer.shape[0]),
                         int(block_table.shape[1])),
                        dtype=torch.int32,
                        device=self.expanded_block_table_buffer.device,
                    )
                bt_w = int(block_table.shape[1])
                self.expanded_block_table_buffer[:actual_expanded, :bt_w] = (
                    torch.repeat_interleave(
                        block_table, decode_lens, dim=0, output_size=actual_expanded
                    )
                )
                if actual_expanded < num_decode_tokens:
                    self.expanded_block_table_buffer[
                        actual_expanded:num_decode_tokens, 0
                    ] = 0
                block_table = self.expanded_block_table_buffer[:num_decode_tokens]

                # All reqs now have decode_len=1
                self.decode_lens_buffer[:num_decode_tokens] = 1
                decode_lens = self.decode_lens_buffer[:num_decode_tokens]
                return seq_lens, block_table, decode_lens, num_decode_tokens, False
        else:
            # Native path: plain decode (next_n==1) or spec decode
            # with 2D per-token context lengths (next_n > 1).
            #
            # When decode_lens are not truly uniform (e.g. some requests have
            # decode_len < next_n due to padding or short prefills), the simple
            # reshape in sparse_attn_indexer won't work. Use pack_seq_triton
            # (requires_padding) instead.
            requires_padding = min_decode_len != max_decode_len
            if use_native and next_n > 1:
                assert self.decode_seq_lens_buffer.dim() == 1
                # (B, max_decode_len): token j attends to
                # L - max_decode_len + j + 1 KV tokens.
                seq_lens_buffer = self.decode_seq_lens_buffer[
                    : num_decodes * max_decode_len
                ].view(num_decodes, max_decode_len)
                seq_lens_for_expansion = (
                    global_seq_lens if global_seq_lens is not None else seq_lens
                )
                expanded_seq_lens = (
                    seq_lens_for_expansion.unsqueeze(1)
                    - max_decode_len
                    + 1
                    + self.offsets_buffer[:max_decode_len]
                )
                if global_seq_lens is not None and self.dcp_world_size > 1:
                    expanded_seq_lens = get_dcp_local_seq_lens(
                        expanded_seq_lens.reshape(-1),
                        self.dcp_world_size,
                        self.dcp_rank,
                        self.cp_kv_cache_interleave_size,
                    ).view(num_decodes, max_decode_len)
                seq_lens_buffer[:] = expanded_seq_lens
                seq_lens = seq_lens_buffer
            return seq_lens, block_table, decode_lens, num_decodes, requires_padding

    def build(
        self,
        common_prefix_len: int,
        common_attn_metadata: CommonAttentionMetadata,
        fast_build: bool = False,
    ) -> DeepseekV32IndexerMetadata:
        num_reqs = common_attn_metadata.num_reqs
        num_tokens = common_attn_metadata.num_actual_tokens
        query_start_loc = common_attn_metadata.query_start_loc
        query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
        if common_attn_metadata.batch_topology is not None:
            query_start_loc_np = common_attn_metadata.batch_topology.query_start_loc_np[
                : num_reqs + 1
            ]
            query_lens_np = common_attn_metadata.batch_topology.query_lens_np
        else:
            query_start_loc_np = query_start_loc_cpu.numpy()
            query_lens_np = np.diff(query_start_loc_np)
        seq_lens = common_attn_metadata.seq_lens
        slot_mapping = common_attn_metadata.slot_mapping
        block_table = common_attn_metadata.block_table_tensor
        use_dcp_local_kv = (
            self.dcp_world_size > 1
            and common_attn_metadata.dcp_local_seq_lens is not None
        )
        effective_dcp_world_size = self.dcp_world_size if use_dcp_local_kv else 1
        effective_dcp_rank = self.dcp_rank if use_dcp_local_kv else 0

        num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
            common_attn_metadata.split_decodes_and_prefills(
                decode_threshold=self.reorder_batch_threshold,
                require_uniform=not self.use_flattening,
            )
        )

        if (num_decodes + num_prefills != num_reqs
                or num_decode_tokens + num_prefill_tokens != num_tokens):
            raise AssertionError(
                f"[idxfix-diag] split mismatch: decodes={num_decodes} "
                f"prefills={num_prefills} reqs={num_reqs} "
                f"decode_toks={num_decode_tokens} prefill_toks={num_prefill_tokens} "
                f"toks={num_tokens} threshold={self.reorder_batch_threshold} "
                f"uniform={not self.use_flattening} "
                f"qsl={common_attn_metadata.query_start_loc_cpu.tolist()[:12]} "
                f"seq={common_attn_metadata._seq_lens_cpu.tolist()[:12] if common_attn_metadata._seq_lens_cpu is not None else None}"
            )

        compressed_slot_mapping = slot_mapping
        logical_compressed_seq_lens = seq_lens
        compressed_seq_lens = seq_lens
        if self.compress_ratio > 1 or use_dcp_local_kv:
            compressed_slot_mapping = get_compressed_slot_mapping(
                num_tokens,
                query_start_loc,
                seq_lens,
                block_table,
                self.storage_block_size,
                self.compress_ratio,
                out=self.compressed_slot_mapping_buffer,
                dcp_world_size=effective_dcp_world_size,
                dcp_rank=effective_dcp_rank,
                cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
            )
            logical_compressed_seq_lens = seq_lens // self.compress_ratio
            compressed_seq_lens = logical_compressed_seq_lens
        if use_dcp_local_kv:
            compressed_seq_lens = get_dcp_local_seq_lens(
                logical_compressed_seq_lens,
                self.dcp_world_size,
                self.dcp_rank,
                self.cp_kv_cache_interleave_size,
            )

        prefill_metadata = None
        if num_prefills > 0:
            # This CPU value is an upper bound for async-spec extend rows.  It
            # is safe for chunking/allocation because CUDA metadata below is
            # built from exact device seq_lens and gather ignores the tail.
            assert common_attn_metadata.seq_lens_cpu_upper_bound is not None
            seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
            logical_compressed_seq_lens_cpu = (
                seq_lens_cpu // self.compress_ratio
                if self.compress_ratio > 1
                else seq_lens_cpu
            )
            compressed_seq_lens_cpu = logical_compressed_seq_lens_cpu
            if use_dcp_local_kv:
                compressed_seq_lens_cpu = get_dcp_local_seq_lens(
                    logical_compressed_seq_lens_cpu,
                    self.dcp_world_size,
                    self.dcp_rank,
                    self.cp_kv_cache_interleave_size,
                )
            compressed_seq_lens_cpu_np = compressed_seq_lens_cpu.numpy()
            prefill_query_lens_cpu = query_lens_np[
                num_decodes : num_decodes + num_prefills
            ]
            max_logits_bytes = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
            if envs.VLLM_USE_B12X_SPARSE_INDEXER:
                chunk_specs = []
                b12x_budget_seq_lens = np.array(
                    [_get_b12x_paged_indexer_supertile_k()],
                    dtype=compressed_seq_lens_cpu_np.dtype,
                )
                for prefill_idx in range(num_prefills):
                    req_idx = num_decodes + prefill_idx
                    chunk_specs.extend(
                        split_indexer_prefill_chunks(
                            b12x_budget_seq_lens,
                            prefill_query_lens_cpu[
                                prefill_idx : prefill_idx + 1
                            ],
                            self.max_prefill_buffer_size,
                            max_logits_bytes,
                            request_offset=req_idx,
                        )
                    )
            else:
                chunk_specs = split_indexer_prefill_chunks(
                    compressed_seq_lens_cpu_np[num_decodes:],
                    prefill_query_lens_cpu,
                    self.max_prefill_buffer_size,
                    max_logits_bytes,
                    request_offset=num_decodes,
                )

            chunks = []
            for req_slice, query_slice in chunk_specs:
                metadata = build_prefill_chunk_metadata(
                    req_slice.start,
                    req_slice.stop,
                    query_start_loc,
                    query_start_loc_cpu,
                    seq_lens,
                    compressed_seq_lens,
                    compressed_seq_lens_cpu,
                    common_attn_metadata.block_table_tensor,
                    self.compress_ratio,
                    query_slice=query_slice,
                    query_start_loc_np=query_start_loc_np,
                    compressed_seq_lens_np=compressed_seq_lens_cpu_np,
                    skip_kv_gather=query_slice.start > 0,
                    dcp_world_size=effective_dcp_world_size,
                    dcp_rank=effective_dcp_rank,
                    cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
                )
                # Skip when total_seq_lens is 0 (i.e., no compressed token).
                if metadata is not None:
                    chunks.append(metadata)
            prefill_metadata = DeepseekV32IndexerPrefillMetadata(chunks)

        decode_metadata = None
        if num_decodes > 0:
            torch.diff(
                common_attn_metadata.query_start_loc[: num_decodes + 1],
                out=self.decode_lens_buffer[:num_decodes],
            )
            decode_lens = self.decode_lens_buffer[:num_decodes]
            decode_lens_np = query_lens_np[:num_decodes]

            global_decode_seq_lens = common_attn_metadata.seq_lens[:num_decodes]
            dcp_local_seq_lens = (
                common_attn_metadata.dcp_local_seq_lens[:num_decodes]
                if self.compress_ratio == 1
                and use_dcp_local_kv
                and common_attn_metadata.dcp_local_seq_lens is not None
                else None
            )
            dcp_local_seq_lens_cpu = (
                common_attn_metadata.dcp_local_seq_lens_cpu[:num_decodes]
                if dcp_local_seq_lens is not None
                and common_attn_metadata.dcp_local_seq_lens_cpu is not None
                else None
            )
            seq_lens = (
                dcp_local_seq_lens
                if dcp_local_seq_lens is not None
                else global_decode_seq_lens
            )
            block_table = common_attn_metadata.block_table_tensor[:num_decodes, ...]

            max_decode_len = int(decode_lens_np.max())
            next_n = 1 + self.num_speculative_tokens
            use_native = not self.use_flattening and max_decode_len <= next_n

            seq_lens, block_table, decode_lens, batch_size, requires_padding = (
                self._prepare_decode_tensors(
                    seq_lens=seq_lens,
                    block_table=block_table,
                    decode_lens=decode_lens,
                    decode_lens_np=decode_lens_np,
                    query_start_loc=common_attn_metadata.query_start_loc[:num_decodes],
                    num_decodes=num_decodes,
                    num_decode_tokens=num_decode_tokens,
                    use_native=use_native,
                    next_n=next_n,
                    max_decode_len=max_decode_len,
                    global_seq_lens=global_decode_seq_lens
                    if dcp_local_seq_lens is not None
                    else None,
                )
            )

            # For DeepseekV4 (compress_ratio > 1), the indexer KV cache stores
            # compressed tokens. Convert uncompressed seq_lens to compressed.
            if self.compress_ratio > 1:
                # True iff seq_lens aliases decode_seq_lens_buffer (flatten or
                # native wrote it); False iff it aliases common_attn_metadata.
                seq_lens_is_local_view = (use_native and next_n > 1) or (
                    not use_native and max_decode_len > 1
                )
                if seq_lens_is_local_view:
                    seq_lens //= self.compress_ratio
                    if use_dcp_local_kv:
                        if seq_lens.dim() == 1:
                            dcp_seq_lens = get_dcp_local_seq_lens(
                                seq_lens,
                                self.dcp_world_size,
                                self.dcp_rank,
                                self.cp_kv_cache_interleave_size,
                            )
                        else:
                            dcp_seq_lens = get_dcp_local_seq_lens(
                                seq_lens.reshape(-1),
                                self.dcp_world_size,
                                self.dcp_rank,
                                self.cp_kv_cache_interleave_size,
                            ).view(seq_lens.shape)
                        seq_lens.copy_(dcp_seq_lens)
                else:
                    # Copy to avoid mutating shared state; keeps CG address stable.
                    compressed_decode_seq_lens = seq_lens // self.compress_ratio
                    if use_dcp_local_kv:
                        compressed_decode_seq_lens = get_dcp_local_seq_lens(
                            compressed_decode_seq_lens,
                            self.dcp_world_size,
                            self.dcp_rank,
                            self.cp_kv_cache_interleave_size,
                        )
                    self.expanded_seq_lens_buffer[:num_decodes] = (
                        compressed_decode_seq_lens
                    )
                    self.expanded_seq_lens_buffer[num_decodes:num_decode_tokens] = 0
                    seq_lens = self.expanded_seq_lens_buffer[:num_decode_tokens]
            elif use_dcp_local_kv and dcp_local_seq_lens is None:
                if seq_lens.dim() == 1:
                    seq_lens = get_dcp_local_seq_lens(
                        seq_lens,
                        self.dcp_world_size,
                        self.dcp_rank,
                        self.cp_kv_cache_interleave_size,
                    )
                else:
                    seq_lens_shape = seq_lens.shape
                    seq_lens = get_dcp_local_seq_lens(
                        seq_lens.reshape(-1),
                        self.dcp_world_size,
                        self.dcp_rank,
                        self.cp_kv_cache_interleave_size,
                    ).view(seq_lens_shape)

            # Non-MTP: deep_gemm paged MQA logits requires 2D context_lens
            # (csrc/apis/attention.hpp). Unsqueeze to (B, 1) so downstream
            # kernels see the same (B, next_n) layout as the MTP path.
            if seq_lens.dim() == 1:
                seq_lens = seq_lens.unsqueeze(-1)

            active_width = None
            active_width_tokens = None
            if self.compress_ratio > 1:
                active_width_tokens = -(
                    -int(common_attn_metadata.max_seq_len) //
                    self.compress_ratio
                )
            if envs.VLLM_USE_B12X_SPARSE_INDEXER:
                # Live scorer window in cache tokens. ceil(max_seq_len /
                # compress_ratio) is an upper bound on the max compressed
                # context across the batch, so windowing to it is top-k-identical
                # to the capacity cap and only skips wasted k-tiles. Computed on
                # the host here (metadata-prep, outside cudagraph capture) and
                # filled into the persistent buffer the captured kernel reads.
                if active_width_tokens is None:
                    active_width_tokens = int(common_attn_metadata.max_seq_len)
                self.b12x_active_width_buffer.fill_(active_width_tokens)
                active_width = self.b12x_active_width_buffer
                if self.compress_ratio > 1:
                    # Compressed-MLA models already bound the B12X scorer with
                    # active_width. Avoiding a host reduction here removes a
                    # per-decode sync without changing the scorer window.
                    decode_topk_max_seq_len = active_width_tokens
                else:
                    # GLM/Kimi-style uncompressed indexer rows need the exact
                    # scalar for best throughput; use the CPU shadow maintained
                    # by the runner instead of synchronizing on seq_lens.
                    decode_topk_max_seq_len = (
                        self._decode_topk_max_seq_len_from_cpu(
                            common_attn_metadata,
                            num_decodes,
                            decode_lens_np,
                            max_decode_len,
                            use_native,
                            dcp_local_seq_lens_cpu,
                        )
                    )
                    if decode_topk_max_seq_len is None:
                        # Async-spec can make GPU seq_lens authoritative.  In
                        # that mode there is no exact host scalar without a
                        # sync, so use the same graph-stable live-window bound
                        # already consumed by the B12X scorer.
                        decode_topk_max_seq_len = active_width_tokens
            elif active_width_tokens is not None:
                # DeepGEMM still receives exact per-row seq_lens on device.
                # The scalar max_seq_len is only a host scorer bound; using the
                # metadata upper bound avoids a per-token CUDA scalar sync.
                decode_topk_max_seq_len = active_width_tokens
            else:
                decode_topk_max_seq_len = int(seq_lens.max().item())

            if envs.VLLM_USE_B12X_SPARSE_INDEXER:
                schedule_metadata = self._maybe_build_b12x_schedule_metadata(
                    seq_lens,
                    block_table,
                    num_decode_tokens,
                    requires_padding,
                )
            else:
                # DeepGEMM is required for paged MQA logits on CUDA devices.
                schedule_metadata = self._maybe_build_deep_gemm_schedule_metadata(
                    seq_lens
                )

            decode_metadata = DeepSeekV32IndexerDecodeMetadata(
                block_table=block_table,
                seq_lens=seq_lens,
                max_seq_len=decode_topk_max_seq_len,
                decode_lens=decode_lens,
                requires_padding=requires_padding,
                schedule_metadata=schedule_metadata,
                compress_ratio=self.compress_ratio,
                active_width=active_width,
            )

        attn_metadata = DeepseekV32IndexerMetadata(
            seq_lens=common_attn_metadata.seq_lens,
            max_seq_len=common_attn_metadata.max_seq_len,
            slot_mapping=compressed_slot_mapping,
            num_decodes=num_decodes,
            num_decode_tokens=num_decode_tokens,
            num_prefills=num_prefills,
            num_prefill_tokens=num_prefill_tokens,
            prefill=prefill_metadata,
            decode=decode_metadata,
            dcp_world_size=effective_dcp_world_size,
            dcp_rank=effective_dcp_rank,
            cp_interleave_size=self.cp_kv_cache_interleave_size,
        )

        return attn_metadata


def build_prefill_chunk_metadata(
    start_idx: int,
    end_idx: int,
    query_start_loc: torch.Tensor,
    query_start_loc_cpu: torch.Tensor,
    uncompressed_seq_lens: torch.Tensor,
    compressed_seq_lens: torch.Tensor,
    compressed_seq_lens_cpu: torch.Tensor,
    block_table: torch.Tensor,
    compress_ratio: int,
    query_slice: slice | None = None,
    query_start_loc_np: np.ndarray | None = None,
    compressed_seq_lens_np: np.ndarray | None = None,
    skip_kv_gather: bool = False,
    dcp_world_size: int = 1,
    dcp_rank: int = 0,
    cp_kv_cache_interleave_size: int = 1,
) -> DeepseekV32IndexerPrefillChunkMetadata | None:
    if compressed_seq_lens_np is not None:
        total_seq_lens = int(compressed_seq_lens_np[start_idx:end_idx].sum())
    else:
        total_seq_lens = compressed_seq_lens_cpu[start_idx:end_idx].sum().item()
    if total_seq_lens == 0:
        return None

    num_reqs = end_idx - start_idx
    device = block_table.device
    token_to_seq = torch.empty(total_seq_lens, dtype=torch.int32, device=device)

    cu_seq_lens = torch.empty(num_reqs + 1, dtype=torch.int32, device=device)
    # Assigning to slice avoids cpu sync.
    cu_seq_lens[:1] = 0
    torch.cumsum(compressed_seq_lens[start_idx:end_idx], dim=0, out=cu_seq_lens[1:])

    query_start_loc = (
        query_start_loc[start_idx : end_idx + 1] - query_start_loc[start_idx]
    )

    if query_start_loc_np is not None:
        total_query_len = int(
            query_start_loc_np[end_idx] - query_start_loc_np[start_idx]
        )
    else:
        total_query_len = int(
            (query_start_loc_cpu[end_idx] - query_start_loc_cpu[start_idx]).item()
        )
    if query_slice is not None:
        qs_start = query_slice.start
        qs_stop = query_slice.stop
    else:
        qs_start = 0
        qs_stop = total_query_len
    output_query_len = qs_stop - qs_start

    cu_seq_len_ks = torch.empty(output_query_len, dtype=torch.int32, device=device)
    cu_seq_len_ke = torch.empty(output_query_len, dtype=torch.int32, device=device)

    _build_prefill_chunk_metadata_kernel[(num_reqs,)](
        query_start_loc,
        uncompressed_seq_lens[start_idx:end_idx],
        cu_seq_lens,
        token_to_seq,
        cu_seq_len_ks,
        cu_seq_len_ke,
        qs_start,
        qs_stop,
        DCP_WORLD_SIZE=dcp_world_size,
        DCP_RANK=dcp_rank,
        CP_KV_CACHE_INTERLEAVE_SIZE=cp_kv_cache_interleave_size,
        BLOCK_SIZE=1024,
        COMPRESS_RATIO=compress_ratio,
    )

    token_start = (
        int(query_start_loc_np[start_idx])
        if query_start_loc_np is not None
        else query_start_loc_cpu[start_idx].item()
    )
    if query_slice is not None:
        token_end = token_start + qs_stop
        token_start = token_start + qs_start
        skip_kv_gather = skip_kv_gather or qs_start > 0
    else:
        token_end = (
            int(query_start_loc_np[end_idx])
            if query_start_loc_np is not None
            else query_start_loc_cpu[end_idx].item()
        )

    return DeepseekV32IndexerPrefillChunkMetadata(
        cu_seqlen_ks=cu_seq_len_ks,
        cu_seqlen_ke=cu_seq_len_ke,
        cu_seq_lens=cu_seq_lens,
        token_to_seq=token_to_seq,
        total_seq_lens=total_seq_lens,
        block_table=block_table[start_idx:end_idx],
        token_start=token_start,
        token_end=token_end,
        num_reqs=num_reqs,
        skip_kv_gather=skip_kv_gather,
    )


@triton.jit(do_not_specialize=["query_slice_start", "query_slice_stop"])
def _build_prefill_chunk_metadata_kernel(
    # Inputs
    query_start_loc_ptr,
    uncompressed_seq_lens_ptr,
    cu_compressed_seq_lens_ptr,
    # Outputs
    token_to_seq_ptr,
    cu_compressed_seq_len_ks_ptr,
    cu_compressed_seq_len_ke_ptr,
    query_slice_start,
    query_slice_stop,
    DCP_WORLD_SIZE: tl.constexpr,
    DCP_RANK: tl.constexpr,
    CP_KV_CACHE_INTERLEAVE_SIZE: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
    COMPRESS_RATIO: tl.constexpr,
):
    batch_idx = tl.program_id(0)

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

    seq_start = tl.load(cu_compressed_seq_lens_ptr + batch_idx)
    seq_end = tl.load(cu_compressed_seq_lens_ptr + batch_idx + 1)
    compressed_seq_len = seq_end - seq_start

    uncompressed_seq_len = tl.load(uncompressed_seq_lens_ptr + batch_idx)
    start_pos = uncompressed_seq_len - query_len

    for i in range(0, query_len, BLOCK_SIZE):
        offset = i + tl.arange(0, BLOCK_SIZE)
        abs_pos = query_start + offset
        mask = (
            (offset < query_len)
            & (abs_pos >= query_slice_start)
            & (abs_pos < query_slice_stop)
        )
        out_pos = abs_pos - query_slice_start

        # Compute cu_seq_len_ks
        tl.store(cu_compressed_seq_len_ks_ptr + out_pos, seq_start, mask=mask)

        # Compute cu_seq_len_ke
        seq_len_per_token = (start_pos + 1 + offset) // COMPRESS_RATIO
        if DCP_WORLD_SIZE > 1:
            base = (
                seq_len_per_token
                // CP_KV_CACHE_INTERLEAVE_SIZE
                // DCP_WORLD_SIZE
                * CP_KV_CACHE_INTERLEAVE_SIZE
            )
            remainder = seq_len_per_token - base * DCP_WORLD_SIZE
            rank_remainder = tl.minimum(
                tl.maximum(
                    remainder - DCP_RANK * CP_KV_CACHE_INTERLEAVE_SIZE,
                    0,
                ),
                CP_KV_CACHE_INTERLEAVE_SIZE,
            )
            seq_len_per_token = base + rank_remainder
        tl.store(
            cu_compressed_seq_len_ke_ptr + out_pos,
            seq_start + seq_len_per_token,
            mask=mask,
        )

    # Compute token_to_seq
    for i in range(0, compressed_seq_len, BLOCK_SIZE):
        offset = i + tl.arange(0, BLOCK_SIZE)
        mask = offset < compressed_seq_len
        tl.store(token_to_seq_ptr + seq_start + offset, batch_idx, mask=mask)


def _prewarm_prefill_chunk_metadata_kernel(
    *,
    device: torch.device,
    dcp_world_size: int,
    dcp_rank: int,
    cp_kv_cache_interleave_size: int,
    compress_ratio: int,
) -> None:
    if not current_platform.is_cuda():
        return

    query_start_loc = torch.tensor([0, 1], dtype=torch.int32, device=device)
    uncompressed_seq_lens = torch.tensor([1], dtype=torch.int32, device=device)
    cu_seq_lens = torch.tensor([0, 1], dtype=torch.int32, device=device)
    token_to_seq = torch.empty((1,), dtype=torch.int32, device=device)
    cu_seq_len_ks = torch.empty((1,), dtype=torch.int32, device=device)
    cu_seq_len_ke = torch.empty((1,), dtype=torch.int32, device=device)

    _build_prefill_chunk_metadata_kernel[(1,)](
        query_start_loc,
        uncompressed_seq_lens,
        cu_seq_lens,
        token_to_seq,
        cu_seq_len_ks,
        cu_seq_len_ke,
        0,
        1,
        DCP_WORLD_SIZE=int(dcp_world_size),
        DCP_RANK=int(dcp_rank),
        CP_KV_CACHE_INTERLEAVE_SIZE=int(cp_kv_cache_interleave_size),
        BLOCK_SIZE=1024,
        COMPRESS_RATIO=int(compress_ratio),
    )