lightseekorg--tokenspeed/python/tokenspeed/runtime/layers/attention/chunk.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 dataclasses import dataclass
from functools import cached_property

import torch
import triton
import triton.language as tl

from tokenspeed.runtime.utils import get_colorful_logger
from tokenspeed.runtime.utils.env import global_server_args_dict

logger = get_colorful_logger(__name__)


@triton.jit
def create_chunked_cache_kv_indices_paged(
    req_to_page_ptr,  # (max_batch, max_pages)
    req_pool_indices_ptr,  # (batch_size,)
    chunk_start_idx_ptr,  # (batch_size,)
    chunk_seq_lens_ptr,  # (batch_size,)
    chunk_cum_seq_lens_ptr,  # (batch_size + 1,)
    chunk_kv_indices_ptr,  # (num_chunk_tokens,)
    req_to_page_ptr_stride: tl.constexpr,
    PAGE_SIZE: tl.constexpr,
):
    BLOCK_SIZE: tl.constexpr = 512
    pid = tl.program_id(axis=0)

    req_pool_index = tl.load(req_pool_indices_ptr + pid)
    chunk_kv_indices_offset = tl.load(chunk_cum_seq_lens_ptr + pid)

    chunk_start_pos = tl.load(chunk_start_idx_ptr + pid).to(tl.int32)
    chunk_seq_len = tl.load(chunk_seq_lens_ptr + pid).to(tl.int32)

    num_loop = tl.cdiv(chunk_seq_len, BLOCK_SIZE)
    for i in range(num_loop):
        offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
        mask = offset < chunk_seq_len
        token_pos = chunk_start_pos + offset
        page_idx = token_pos // PAGE_SIZE
        page_id = tl.load(
            req_to_page_ptr + req_pool_index * req_to_page_ptr_stride + page_idx,
            mask=mask,
        )
        kv_slot = page_id * PAGE_SIZE + token_pos % PAGE_SIZE
        tl.store(
            chunk_kv_indices_ptr + chunk_kv_indices_offset + offset,
            kv_slot,
            mask=mask,
        )


def get_max_chunk_capacity():
    return (
        global_server_args_dict["chunked_prefill_size"]
        * global_server_args_dict["mla_chunk_multiplier"]
    )


# Here we suppose the length of each chunk is equal
# For example, if we have 4 sequences with seq length [256, 512, 768, 1024], chunk_len = 256
# num_chunks = cdiv(1024, 256) = 4
# chunk_starts = [[0, 0, 0, 0], [256, 256, 256, 256], [512, 512, 512, 512], [768, 768, 768, 768]]
# chunk_ends = [[256, 256, 256, 256], [256, 512, 512, 512], [256, 512, 768, 768], [256, 512, 768, 1024]]
# chunk_seq_lens = [[256, 256, 256, 256], [0, 256, 256, 256], [0, 0, 256, 256], [0, 0, 0, 256]]
"""
        seq0 seq1 seq2 seq3
chunk0   --   --   --   --
chunk1   --   --   --   --
chunk2   --   --   --   --
chunk3   --   --   --   --
"""


# starts, ends, len_in_chunk, cum_seq_lens, all satisfy the above layout
@dataclass
class Chunks:
    starts: torch.Tensor
    ends: torch.Tensor
    len_in_chunk: torch.Tensor

    @cached_property
    def cum_seq_lens(self):
        num_chunks = self.starts.shape[0]
        bs = self.starts.shape[1]
        result = torch.zeros(
            num_chunks, bs + 1, device=self.starts.device, dtype=torch.int32
        )
        torch.cumsum(self.len_in_chunk, dim=1, out=result[:, 1:])
        return result


def chunking(prefix_lens: torch.Tensor, num_chunks, batch_size, chunk_len):
    starts = (
        torch.arange(num_chunks, device=prefix_lens.device, dtype=torch.int32)
        .unsqueeze(1)
        .expand(-1, batch_size)
        * chunk_len
    )
    ends = torch.min(prefix_lens.unsqueeze(0), starts + chunk_len).to(torch.int32)

    chunks = Chunks(
        starts=starts,
        ends=ends,
        len_in_chunk=(ends - starts).clamp(min=0).to(torch.int32),
    )
    return chunks


def get_chunks_paged(
    prefix_lens, prefix_lens_cpu, req_to_page, req_pool_indices, page_size
):
    """Page-table aware version of get_chunks."""
    device: torch.device = prefix_lens.device
    batch_size = len(prefix_lens_cpu)

    chunk_capacity = get_max_chunk_capacity()
    chunk_len = chunk_capacity // batch_size
    max_prefix = prefix_lens_cpu.max().item()
    num_chunks = (max_prefix + chunk_len - 1) // chunk_len

    chunks = chunking(prefix_lens, num_chunks, batch_size, chunk_len)
    chunks_cpu = chunking(prefix_lens_cpu, num_chunks, batch_size, chunk_len)

    num_tokens_per_forward = chunks_cpu.len_in_chunk.sum(dim=1).tolist()

    chunk_kv_indices_list = []
    for idx in range(num_chunks):
        chunk_kv_indices = torch.empty(
            num_tokens_per_forward[idx], dtype=torch.int32, device=device
        )
        create_chunked_cache_kv_indices_paged[(batch_size,)](
            req_to_page,
            req_pool_indices,
            chunks.starts[idx],
            chunks.len_in_chunk[idx],
            chunks.cum_seq_lens[idx],
            chunk_kv_indices,
            req_to_page.shape[1],
            page_size,
        )
        chunk_kv_indices_list.append(chunk_kv_indices)

    return chunks, chunk_kv_indices_list, chunks_cpu


def build_chunked_prefill_metadata_arrays(
    extend_prefix_lens,
    extend_prefix_lens_cpu,
    req_to_page,
    req_pool_indices,
    page_size,
):
    """Build the per-prefix-loop arrays for chunked-prefill MLA.

    Run once per chunked-prefill iteration in the backend's
    ``_init_prefill_metadata``. Returns:

    - ``chunked_loop_num``: number of prefix loop iterations
    - ``chunk_kv_indices_list``: List[Tensor], paged KV indices per loop_idx
    - ``chunked_seq_len``: (chunked_loop_num, num_extends) int32 GPU — per-seq
      KV length within each loop_idx (zero for seqs whose prefix doesn't
      reach this chunk).
    - ``cu_chunked_seq_len``: (chunked_loop_num, num_extends+1) int32 GPU —
      cumsum along the seq dim, fed to the chunker as ``cum_seq_lens_kv``.
    - ``max_chunk_len_per_loop``: List[int], CPU max-seq-len per loop_idx,
      fed to the chunker as ``max_kv_len``.

    The q-side cumsum / max do not appear here: callers alias them to the
    causal pass's ``cum_extend_seq_lens`` / ``max_extend_seq_len``, since
    every prefix-chunk forward sees the same ``q_lens == extend_seq_lens``.
    """
    chunks, chunk_kv_indices_list, chunks_cpu = get_chunks_paged(
        extend_prefix_lens,
        extend_prefix_lens_cpu,
        req_to_page,
        req_pool_indices,
        page_size,
    )
    chunked_loop_num = chunks.starts.shape[0]
    max_chunk_len_per_loop = [
        chunks_cpu.len_in_chunk[i].max().item() for i in range(chunked_loop_num)
    ]
    return (
        chunked_loop_num,
        chunk_kv_indices_list,
        chunks.len_in_chunk,
        chunks.cum_seq_lens,
        max_chunk_len_per_loop,
    )