launcher/dcp2/baked/v1/attention/backends/mla/compressor_utils.py

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

import torch

from vllm.model_executor.warmup.jit_warmup import zip_inputs
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
    LaunchSpec,
    TritonWarmupTensor,
    VllmTritonJitKernel,
    kernel_launcher,
    triton_scalar_specialization_rep,
)
from vllm.triton_utils import tl, triton
from vllm.utils.math_utils import cdiv

_DSPARK_SWA_INDEX_ALIGNMENT = 64


def get_dspark_swa_index_width(
    window_size: int,
    num_speculative_tokens: int,
) -> int:
    """Return the padded width of non-causal DSpark SWA indices."""
    width = max(int(window_size), 0) + max(int(num_speculative_tokens), 0)
    return cdiv(width, _DSPARK_SWA_INDEX_ALIGNMENT) * _DSPARK_SWA_INDEX_ALIGNMENT


class CompressedSlotMappingKernel(
    VllmTritonJitKernel["CompressedSlotMappingKernel.CompileKey"]
):
    TRITON_BLOCK_SIZE = 1024

    @dataclass(frozen=True)
    class CompileKey:
        compress_ratio: int
        triton_block_size: int
        block_size: int
        dcp_world: int = 1
        dcp_rank: int = 0
        dcp_interleave: int = 1

    @staticmethod
    @triton.jit(do_not_specialize=["block_table_stride"])
    def kernel(
        # [num_tokens]
        slot_mapping_ptr,
        # [num_reqs + 1]
        query_start_loc_ptr,
        # [num_reqs]
        seq_lens_ptr,
        # [num_reqs, max_num_blocks]
        block_table_ptr,
        block_table_stride,
        # states per page on THIS rank (kv_cache_spec.num_states)
        block_size,
        COMPRESS_RATIO: tl.constexpr,
        PAD_ID: tl.constexpr,
        TRITON_BLOCK_SIZE: tl.constexpr,
        DCP_WORLD: tl.constexpr,
        DCP_RANK: tl.constexpr,
        DCP_INTERLEAVE: 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_len = tl.load(seq_lens_ptr + batch_idx)
        start_pos = seq_len - query_len

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

            pos = start_pos + i + tl.arange(0, TRITON_BLOCK_SIZE)
            is_valid = (pos + 1) % COMPRESS_RATIO == 0
            # Global compressed state index of the group this token closes.
            state = pos // COMPRESS_RATIO
            if DCP_WORLD > 1:
                # DCP shards the compressed cache over STATES with the runner's
                # token-space formula (worker/gpu/block_table.py) applied in
                # state space: one block-table entry spans block_size * DCP_WORLD
                # states, this rank owns the states whose interleave group is
                # congruent to DCP_RANK, and stores them contiguously per entry.
                virtual_block_size = block_size * DCP_WORLD
                virtual_block_idx = state // virtual_block_size
                virtual_off = state % virtual_block_size
                is_local = (virtual_off // DCP_INTERLEAVE) % DCP_WORLD == DCP_RANK
                local_off = (virtual_off // (DCP_INTERLEAVE * DCP_WORLD)) * DCP_INTERLEAVE + (
                    virtual_off % DCP_INTERLEAVE
                )
                local_state = virtual_block_idx * block_size + local_off
                is_valid = is_valid & is_local
            else:
                local_state = state

            block_ids = local_state // block_size
            block_numbers = tl.load(
                block_table_ptr + batch_idx * block_table_stride + block_ids,
                mask=mask & is_valid,
            )
            slot_ids = block_numbers * block_size + local_state % block_size

            # NOTE
            slot_ids = tl.where(is_valid, slot_ids, PAD_ID)
            tl.store(slot_mapping_ptr + query_start + offset, slot_ids, mask=mask)

    def dispatch(  # type: ignore[override]
        self,
        *,
        compress_ratio: int,
        block_size: int,
        dcp_world_size: int = 1,
        dcp_rank: int = 0,
        cp_kv_cache_interleave_size: int = 1,
    ) -> CompileKey:
        return self.CompileKey(
            compress_ratio=compress_ratio,
            triton_block_size=self.TRITON_BLOCK_SIZE,
            block_size=triton_scalar_specialization_rep(block_size),
            dcp_world=dcp_world_size,
            dcp_rank=dcp_rank,
            dcp_interleave=cp_kv_cache_interleave_size,
        )

    def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
        hf_config = vllm_config.model_config.hf_config
        parallel_config = vllm_config.parallel_config
        dcp_world_size = parallel_config.decode_context_parallel_size
        cp_interleave = parallel_config.cp_kv_cache_interleave_size
        dcp_rank = 0
        if dcp_world_size > 1:
            from vllm.distributed import get_dcp_group

            dcp_rank = get_dcp_group().rank_in_group
        configured_ratios = (
            *(getattr(hf_config, "compress_ratios", None) or ()),
            getattr(hf_config, "index_kpool", 1) or 1,
        )
        compress_ratios = tuple(
            dict.fromkeys(int(ratio) for ratio in configured_ratios if int(ratio) > 1)
        )
        if not compress_ratios:
            return []
        return self._trace_dispatch(self.dispatch)(
            zip_inputs(
                *(
                    dict(
                        compress_ratio=ratio,
                        block_size=vllm_config.cache_config.block_size // ratio,
                        dcp_world_size=dcp_world_size,
                        dcp_rank=dcp_rank,
                        cp_kv_cache_interleave_size=cp_interleave,
                    )
                    for ratio in compress_ratios
                )
            )
        )

    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
        int32_ptr = TritonWarmupTensor(torch.int32)
        return dict(
            slot_mapping=TritonWarmupTensor(torch.int64),
            query_start_loc=int32_ptr,
            seq_lens=int32_ptr,
            block_table=int32_ptr,
            block_size=compile_key.block_size,
            compress_ratio=compile_key.compress_ratio,
            dcp_world_size=compile_key.dcp_world,
            dcp_rank=compile_key.dcp_rank,
            cp_kv_cache_interleave_size=compile_key.dcp_interleave,
        )

    @kernel_launcher
    def __call__(
        self,
        slot_mapping: torch.Tensor,
        query_start_loc: torch.Tensor,
        seq_lens: torch.Tensor,
        block_table: torch.Tensor,
        block_size: int,
        compress_ratio: int,
        dcp_world_size: int = 1,
        dcp_rank: int = 0,
        cp_kv_cache_interleave_size: int = 1,
    ) -> LaunchSpec:
        assert block_size % cp_kv_cache_interleave_size == 0, (
            f"states per page ({block_size}) must be a multiple of "
            f"cp_kv_cache_interleave_size ({cp_kv_cache_interleave_size})"
        )
        return (block_table.shape[0],), dict(
            block_table_stride=block_table.stride(0),
            COMPRESS_RATIO=compress_ratio,
            PAD_ID=-1,
            TRITON_BLOCK_SIZE=self.TRITON_BLOCK_SIZE,
            DCP_WORLD=dcp_world_size,
            DCP_RANK=dcp_rank,
            DCP_INTERLEAVE=cp_kv_cache_interleave_size,
        )


def get_compressed_slot_mapping(
    num_tokens: int,
    query_start_loc: torch.Tensor,
    seq_lens: torch.Tensor,
    block_table: torch.Tensor,
    block_size: int,
    compress_ratio: int,
    out: torch.Tensor | None = None,
    dcp_world_size: int = 1,
    dcp_rank: int = 0,
    cp_kv_cache_interleave_size: int = 1,
) -> torch.Tensor:
    """Slot of the compressed state each token closes (PAD otherwise).

    ``block_size`` is the number of states per page on this rank. Under DCP
    (``dcp_world_size > 1``) states are sharded across ranks and only this
    rank's states get a slot; the block table is the shared (manager) table
    whose entries each span ``block_size * dcp_world_size`` states.
    """
    if out is not None:
        # Guard: for padded / invalid sequences.
        # Negative positions produce bogus block indices that lead to illegal memory
        # accesses inside the block_table load.
        # NOTE: Fill -1 to the whole tensor, not just the first `num_tokens`.
        out.fill_(-1)
        slot_mapping = out[:num_tokens]
    else:
        slot_mapping = torch.full(
            (num_tokens,), -1, dtype=torch.int64, device=query_start_loc.device
        )

    _COMPRESSED_SLOT_MAPPING_KERNEL(
        slot_mapping,
        query_start_loc,
        seq_lens,
        block_table,
        block_size,
        compress_ratio,
        dcp_world_size,
        dcp_rank,
        cp_kv_cache_interleave_size,
    )
    return slot_mapping


_COMPRESSED_SLOT_MAPPING_KERNEL = CompressedSlotMappingKernel()