lightseekorg--tokenspeed/python/tokenspeed/runtime/layers/attention/kv_cache/dsa.py

# Copyright (c) 2026 LightSeek Foundation
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from __future__ import annotations

import torch
from tokenspeed_kernel.ops.quantization import quantize_fp8_with_scale

from tokenspeed.runtime.layers.attention.configs.dsa import dsa_index_k_row_bytes
from tokenspeed.runtime.layers.attention.kv_cache.mla import (
    MLATokenToKVPool,
    _get_tensor_size_bytes,
)

_INDEX_K_FP8_GROUP_SIZE = 128
_INDEX_K_SCALE_BYTES = torch._utils._element_size(torch.float32)


class DSATokenToKVPool(MLATokenToKVPool):
    def __init__(
        self,
        *args,
        index_head_dim: int,
        **kwargs,
    ):
        self.index_head_dim = int(index_head_dim)
        self.index_k_row_bytes = dsa_index_k_row_bytes(self.index_head_dim)
        super().__init__(*args, **kwargs)

        with self.memory_saver_adapter.region():
            self.index_k_buffer = [
                torch.zeros(
                    (self.size + self.page_size, self.index_k_row_bytes),
                    dtype=torch.uint8,
                    device=self.device,
                )
                for _ in range(self.layer_num)
            ]

        self.index_k_data_ptrs = torch.tensor(
            [buf.data_ptr() for buf in self.index_k_buffer],
            dtype=torch.uint64,
            device=self.device,
        )

    def _get_page_size_bytes(self):
        index_size_bytes = self.index_k_row_bytes
        return (
            super()._get_page_size_bytes()
            + self.page_size * self.layer_num * index_size_bytes
        )

    def get_kv_size_bytes(self):
        return super().get_kv_size_bytes() + _get_tensor_size_bytes(self.index_k_buffer)

    def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor):
        super().move_kv_cache(tgt_loc, src_loc)
        if tgt_loc.numel() == 0:
            return
        tgt_loc_flat = tgt_loc.view(-1).long()
        src_loc_flat = src_loc.view(-1).long()
        for buf in self.index_k_buffer:
            # Packed FP8 index-K is block-split per page, so a single token's
            # bytes are NOT a contiguous row; move the FP8 values and FP32
            # scales through their block-split views instead.
            fp8_view, scale_view = self._index_k_block_views(buf)
            ps = self.page_size
            tgt_page = tgt_loc_flat // ps
            tgt_slot = tgt_loc_flat % ps
            src_page = src_loc_flat // ps
            src_slot = src_loc_flat % ps
            fp8_view[tgt_page, tgt_slot] = fp8_view[src_page, src_slot]
            scale_view[tgt_page, tgt_slot] = scale_view[src_page, src_slot]

    def has_index_k_buffer(self) -> bool:
        return True

    def get_index_k_buffer(self, layer_id: int) -> torch.Tensor:
        if self.layer_transfer_counter is not None:
            self.layer_transfer_counter.wait_until(layer_id)
        return self.index_k_buffer[layer_id]

    def set_index_k_buffer(
        self,
        layer_id: int,
        loc: torch.Tensor,
        index_k: torch.Tensor,
    ) -> None:
        if index_k.dtype != self.model_dtype:
            index_k = index_k.to(self.model_dtype)
        index_k = index_k.view(-1, self.index_head_dim)
        self._set_index_k_buffer(layer_id, loc, index_k)

    def _index_k_block_views(
        self, buf: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Return per-page block-split views into a packed FP8 index-K buffer.

        DeepGEMM's ``fp8_paged_mqa_logits`` expects each page of ``page_size``
        tokens to be laid out as ``[page_size * head_dim FP8 values]`` followed
        by ``[page_size * num_groups FP32 scales]`` (block-split), NOT a
        per-token ``[fp8 | scale]`` interleave. The two ``as_strided`` views
        below alias the same storage as ``buf`` so writes land in place.

        Args:
            buf: Packed FP8 index-K buffer of shape ``[num_slots, row_bytes]``
                and dtype ``uint8``, where ``row_bytes == head_dim +
                scale_bytes`` and ``num_slots`` is a multiple of
                ``page_size``.

        Returns:
            ``(fp8_view, scale_view)`` where ``fp8_view`` has shape
            ``[num_pages, page_size, head_dim]`` (FP8 e4m3) and ``scale_view``
            has shape ``[num_pages, page_size, num_groups]`` (float32), both
            indexed as ``view[page, slot_in_page]``.
        """
        ps = self.page_size
        hd = self.index_head_dim
        ng = hd // _INDEX_K_FP8_GROUP_SIZE
        row = hd + ng * _INDEX_K_SCALE_BYTES
        num_pages = buf.shape[0] // ps
        page_bytes = ps * row
        flat = buf.reshape(-1)
        fp8_view = torch.as_strided(
            flat.view(torch.float8_e4m3fn),
            (num_pages, ps, hd),
            (page_bytes, hd, 1),
        )
        scale_view = torch.as_strided(
            flat.view(torch.float32),
            (num_pages, ps, ng),
            (page_bytes // 4, ng, 1),
            (ps * hd) // 4,
        )
        return fp8_view, scale_view

    def gather_index_k(
        self, layer_id: int, slots: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Gather per-token FP8 index-K values and scales from the cache.

        The packed FP8 index-K buffer is stored block-split per page (see
        :meth:`_index_k_block_views`), so the non-paged prefill
        scoring kernel (``fp8_mqa_logits``), which consumes contiguous
        ``(k_fp8, k_scale)`` tensors, must gather token rows through the
        block-split views rather than indexing raw rows.

        Args:
            layer_id: Layer whose index-K cache to read.
            slots: 1D int tensor of global token slot indices to gather.

        Returns:
            ``(k_fp8, k_scale)`` where ``k_fp8`` has shape
            ``[num_slots, head_dim]`` (FP8 e4m3) and ``k_scale`` has shape
            ``[num_slots, num_groups]`` (float32).
        """
        buf = self.get_index_k_buffer(layer_id)
        fp8_view, scale_view = self._index_k_block_views(buf)
        slots = slots.to(torch.long)
        page = slots // self.page_size
        slot_in_page = slots % self.page_size
        k_fp8 = fp8_view[page, slot_in_page]
        k_scale = scale_view[page, slot_in_page]
        return k_fp8, k_scale

    def _set_index_k_buffer(
        self,
        layer_id: int,
        loc: torch.Tensor,
        index_k: torch.Tensor,
    ) -> None:
        buf = self.index_k_buffer[layer_id]
        index_k_fp8, index_k_scale = quantize_fp8_with_scale(
            index_k,
            granularity="token_group",
            group_size=_INDEX_K_FP8_GROUP_SIZE,
            scale_encoding="float32",
        )

        fp8_view, scale_view = self._index_k_block_views(buf)
        loc = loc.to(torch.long)
        page = loc // self.page_size
        slot_in_page = loc % self.page_size
        fp8_view[page, slot_in_page] = index_k_fp8.view(-1, self.index_head_dim)
        scale_view[page, slot_in_page] = index_k_scale.view(
            -1, self.index_head_dim // _INDEX_K_FP8_GROUP_SIZE
        )

    def get_contiguous_buf_infos(self):
        data_ptrs, data_lens, item_lens = super().get_contiguous_buf_infos()
        data_ptrs = list(data_ptrs)
        data_lens = list(data_lens)
        item_lens = list(item_lens)
        for buf in self.index_k_buffer:
            data_ptrs.append(buf.data_ptr())
            data_lens.append(buf.nbytes)
            item_lens.append(buf[0].nbytes * self.page_size)
        return data_ptrs, data_lens, item_lens

    def get_layerwise_buf_info_offsets(self, start_idx=0):
        offsets = super().get_layerwise_buf_info_offsets(start_idx)
        if self.quant_method == "per_token_head":
            base_count = 3 * self.layer_num
        else:
            base_count = self.layer_num
        return [
            layer_offsets + [start_idx + base_count + layer_id]
            for layer_id, layer_offsets in enumerate(offsets)
        ]

    def get_cpu_copy(self, token_indices: list[int]) -> torch.Tensor:
        del token_indices
        raise NotImplementedError(
            "DSA KV cache offload is not implemented; sparse/indexer cache "
            "buffers require page-aware layout handling."
        )

    def load_cpu_copy(
        self, kv_cache_cpu: torch.Tensor, token_indices: list[int]
    ) -> None:
        del kv_cache_cpu, token_indices
        raise NotImplementedError(
            "DSA KV cache reload is not implemented; sparse/indexer cache "
            "buffers require page-aware layout handling."
        )