sgl-project--sglang/python/sglang/jit_kernel/tests/deepseek_v4/common.py

from __future__ import annotations

from dataclasses import dataclass
from typing import List, Literal, Optional, Tuple

import torch

from sglang.jit_kernel.dsv4 import CompressorDecodePlan, CompressorPrefillPlan


@dataclass
class LegacyContext:
    """Per-request ring buffer (no req_to_token / full_to_swa).

    `req_pool_indices[i]` directly maps to the request's ring base slot.
    """

    bs: int
    head_dim: int
    compress_ratio: int
    req_pool_indices: torch.Tensor  # int64 [bs] on cuda
    pages_per_req: int

    @property
    def num_pages(self) -> int:
        # Reserve enough pages to hold all batched requests' rings.
        return int(self.req_pool_indices.max().item() + 1) * self.pages_per_req

    def state_loc(self, b: int, position: int) -> int:
        rid = int(self.req_pool_indices[b].item())
        if self.compress_ratio == 4:
            page = rid * 2 + (position // 4) % 2
        else:
            page = rid
        return page * self.compress_ratio + position % self.compress_ratio

    def make_prefill_plan(
        self,
        seq_lens_cpu: torch.Tensor,
        extend_lens_cpu: torch.Tensor,
        num_q_tokens: int,
    ) -> CompressorPrefillPlan:
        return CompressorPrefillPlan.generate_legacy(
            compress_ratio=self.compress_ratio,  # type: ignore
            req_pool_indices=self.req_pool_indices,
            seq_lens=seq_lens_cpu,
            extend_lens=extend_lens_cpu,
            num_q_tokens=num_q_tokens,
            device=torch.device("cuda"),
        )

    def make_decode_plan(self, seq_lens_gpu: torch.Tensor) -> CompressorDecodePlan:
        return CompressorDecodePlan.generate_legacy(
            compress_ratio=self.compress_ratio,  # type: ignore
            req_pool_indices=self.req_pool_indices,
            seq_lens=seq_lens_gpu,
        )


@dataclass
class PagedContext:
    """SWA paged layout with identity req_to_token + identity full_to_swa.

    Each request occupies `num_swa_pages_per_req` contiguous swa_pages, so
    `req_to_token[r, p] = r * (num_swa_pages_per_req * swa_page_size) + p`.
    """

    bs: int
    head_dim: int
    compress_ratio: int
    swa_page_size: int
    ring_size: int
    num_swa_pages_per_req: int
    req_pool_indices: torch.Tensor  # int64 [bs] on cuda
    req_to_token: torch.Tensor  # int64 [num_reqs_capacity, max_tokens_per_req] on cuda
    full_to_swa: torch.Tensor  # int64 [num_swa_slots] on cuda

    @property
    def num_pages(self) -> int:
        # Upper bound: every (request, position) state slot fits.
        max_state_loc = (
            self.bs * self.num_swa_pages_per_req * self.ring_size
            + self.swa_page_size  # slack for the largest tail
        )
        return max_state_loc // self.compress_ratio + 1

    def state_loc(self, b: int, position: int) -> int:
        rid = int(self.req_pool_indices[b].item())
        loc = int(self.req_to_token[rid, position].item())
        swa_loc = int(self.full_to_swa[loc].item())
        swa_page = swa_loc // self.swa_page_size
        return swa_page * self.ring_size + swa_loc % self.ring_size

    def make_prefill_plan(
        self,
        seq_lens_cpu: torch.Tensor,
        extend_lens_cpu: torch.Tensor,
        num_q_tokens: int,
    ) -> CompressorPrefillPlan:
        return CompressorPrefillPlan.generate(
            compress_ratio=self.compress_ratio,  # type: ignore
            req_pool_indices=self.req_pool_indices,
            seq_lens=seq_lens_cpu,
            extend_lens=extend_lens_cpu,
            req_to_token=self.req_to_token,
            full_to_state=self.full_to_swa,
            swa_page_size=self.swa_page_size,
            ring_size=self.ring_size,
            num_q_tokens=num_q_tokens,
        )

    def make_decode_plan(self, seq_lens_gpu: torch.Tensor) -> CompressorDecodePlan:
        return CompressorDecodePlan.generate(
            compress_ratio=self.compress_ratio,  # type: ignore
            req_pool_indices=self.req_pool_indices,
            req_to_token=self.req_to_token,
            full_to_state=self.full_to_swa,
            seq_lens=seq_lens_gpu,
            swa_page_size=self.swa_page_size,
            ring_size=self.ring_size,
        )


def make_legacy_context(
    bs: int,
    compress_ratio: Literal[4, 128],
    head_dim: int = 512,
) -> LegacyContext:
    pages_per_req = 2 if compress_ratio == 4 else 1
    req_pool_indices = torch.arange(bs, dtype=torch.int64, device="cuda")
    return LegacyContext(
        bs=bs,
        head_dim=head_dim,
        compress_ratio=compress_ratio,
        req_pool_indices=req_pool_indices,
        pages_per_req=pages_per_req,
    )


def make_paged_context(
    bs: int,
    compress_ratio: Literal[4, 128],
    head_dim: int = 512,
    swa_page_size: int = 256,
    ring_size: Optional[int] = None,
    num_swa_pages_per_req: int = 8,
    max_tokens_per_req: int = 8192,
    num_reqs_capacity: int = 16,
) -> PagedContext:
    if ring_size is None:
        ring_size = 8 if compress_ratio == 4 else 128
    assert swa_page_size % ring_size == 0
    assert ring_size % compress_ratio == 0
    assert num_swa_pages_per_req * swa_page_size <= max_tokens_per_req

    stride = num_swa_pages_per_req * swa_page_size
    req_to_token = torch.zeros(
        (num_reqs_capacity, max_tokens_per_req), dtype=torch.int32
    )
    for r in range(bs):
        req_to_token[r, :stride] = torch.arange(
            r * stride, (r + 1) * stride, dtype=torch.int32
        )
    total_swa_slots = num_reqs_capacity * stride
    full_to_swa = torch.arange(total_swa_slots, dtype=torch.int64)
    req_pool_indices = torch.arange(bs, dtype=torch.int64)
    return PagedContext(
        bs=bs,
        head_dim=head_dim,
        compress_ratio=compress_ratio,
        swa_page_size=swa_page_size,
        ring_size=ring_size,
        num_swa_pages_per_req=num_swa_pages_per_req,
        req_pool_indices=req_pool_indices.cuda(),
        req_to_token=req_to_token.cuda(),
        full_to_swa=full_to_swa.cuda(),
    )


def make_state_pool(num_pages: int, compress_ratio: int, head_dim: int) -> torch.Tensor:
    last_dim = head_dim * (4 if compress_ratio == 4 else 2)
    return torch.zeros(
        (num_pages, compress_ratio, last_dim),
        dtype=torch.float32,
        device="cuda",
    )


def to_seq_extend(
    seq_extend_pairs: List[Tuple[int, int]],
) -> Tuple[torch.Tensor, torch.Tensor, int]:
    seq_lens = torch.tensor([s for s, _ in seq_extend_pairs], dtype=torch.int64)
    extend_lens = torch.tensor([e for _, e in seq_extend_pairs], dtype=torch.int64)
    num_q = int(extend_lens.sum().item())
    return seq_lens, extend_lens, num_q