sgl-project--sglang/python/sglang/jit_kernel/kv_canary/plan/offsets_kernel.py

from __future__ import annotations

from typing import Optional

import torch
import triton
import triton.language as tl

from sglang.jit_kernel.kv_canary.consts import (
    REQ_POOL_IDX_PADDING,
    TOKEN_TO_KV_SLOT_PADDING,
)
from sglang.jit_kernel.kv_canary.plan.utils import (
    _compute_window_start,
    _require_1d,
    _require_2d,
    _require_dtype,
    _require_len,
    _require_min_len,
    _require_same_device,
    _resolve_swa_lut,
    _swa_translate_tile,
)
from sglang.jit_kernel.kv_canary.verify import _assert_contiguous

# Upper bound on bs for _plan_offsets_kernel's block-level cumsum. Reqs larger than this exceed Triton's
# single-program tl.cumsum reach. Increase if real workloads ever push past it; the cap is intentionally
# generous so the wrapper never silently truncates.
_PLAN_BS_BLOCK_SIZE: int = 4096


def launch_plan_offsets_kernel(
    *,
    req_pool_indices: torch.Tensor,
    prefix_lens: torch.Tensor,
    extend_seq_lens: torch.Tensor,
    req_to_token: torch.Tensor,
    full_to_swa_index_mapping: Optional[torch.Tensor],
    out_verify_offsets_scratch: torch.Tensor,
    out_write_offsets: torch.Tensor,
    out_write_seed_slot_indices: torch.Tensor,
    out_verify_num_valid: torch.Tensor,
    out_verify_enable: torch.Tensor,
    out_write_num_valid_reqs: torch.Tensor,
    swa_window_size: int,
    verify_capacity: int,
) -> None:
    bs = int(req_pool_indices.shape[0])
    lut_tensor, lut_len, has_swa_lut = _resolve_swa_lut(
        full_to_swa_index_mapping, out_verify_offsets_scratch.device
    )
    req_to_token_stride0 = int(req_to_token.stride(0))
    write_offsets_len = int(out_write_offsets.shape[0])
    write_req_capacity = int(out_write_seed_slot_indices.shape[0])

    _validate_offsets_kernel_inputs(
        req_pool_indices=req_pool_indices,
        prefix_lens=prefix_lens,
        extend_seq_lens=extend_seq_lens,
        req_to_token=req_to_token,
        lut_tensor=lut_tensor,
        out_verify_offsets_scratch=out_verify_offsets_scratch,
        out_write_offsets=out_write_offsets,
        out_write_seed_slot_indices=out_write_seed_slot_indices,
        out_verify_num_valid=out_verify_num_valid,
        out_verify_enable=out_verify_enable,
        out_write_num_valid_reqs=out_write_num_valid_reqs,
        bs=bs,
        req_to_token_stride0=req_to_token_stride0,
        lut_len=lut_len,
        has_swa_lut=has_swa_lut,
        write_offsets_len=write_offsets_len,
        write_req_capacity=write_req_capacity,
        verify_capacity=verify_capacity,
    )

    _plan_offsets_kernel[(1,)](
        req_pool_indices,
        prefix_lens,
        extend_seq_lens,
        req_to_token,
        lut_tensor,
        out_verify_offsets_scratch,
        out_write_offsets,
        out_write_seed_slot_indices,
        out_verify_num_valid,
        out_verify_enable,
        out_write_num_valid_reqs,
        bs,
        req_to_token_stride0,
        lut_len,
        BS_BLOCK=_PLAN_BS_BLOCK_SIZE,
        SWA_WINDOW=int(swa_window_size),
        HAS_SWA_LUT=has_swa_lut,
        WRITE_OFFSETS_LEN=write_offsets_len,
        WRITE_REQ_CAPACITY=write_req_capacity,
        VERIFY_CAPACITY=verify_capacity,
        REQ_POOL_IDX_PADDING=REQ_POOL_IDX_PADDING,
        TOKEN_TO_KV_SLOT_PADDING=TOKEN_TO_KV_SLOT_PADDING,
    )


def _validate_offsets_kernel_inputs(
    *,
    req_pool_indices: torch.Tensor,
    prefix_lens: torch.Tensor,
    extend_seq_lens: torch.Tensor,
    req_to_token: torch.Tensor,
    lut_tensor: torch.Tensor,
    out_verify_offsets_scratch: torch.Tensor,
    out_write_offsets: torch.Tensor,
    out_write_seed_slot_indices: torch.Tensor,
    out_verify_num_valid: torch.Tensor,
    out_verify_enable: torch.Tensor,
    out_write_num_valid_reqs: torch.Tensor,
    bs: int,
    req_to_token_stride0: int,
    lut_len: int,
    has_swa_lut: bool,
    write_offsets_len: int,
    write_req_capacity: int,
    verify_capacity: int,
) -> None:
    _assert_contiguous(req_pool_indices, "req_pool_indices")
    _assert_contiguous(prefix_lens, "prefix_lens")
    _assert_contiguous(extend_seq_lens, "extend_seq_lens")
    _assert_contiguous(req_to_token, "req_to_token")
    _assert_contiguous(lut_tensor, "lut_tensor")
    _assert_contiguous(out_verify_offsets_scratch, "out_verify_offsets_scratch")
    _assert_contiguous(out_write_offsets, "out_write_offsets")
    _assert_contiguous(out_write_seed_slot_indices, "out_write_seed_slot_indices")
    _assert_contiguous(out_verify_num_valid, "out_verify_num_valid")
    _assert_contiguous(out_verify_enable, "out_verify_enable")
    _assert_contiguous(out_write_num_valid_reqs, "out_write_num_valid_reqs")

    _require_dtype(req_pool_indices, "req_pool_indices", torch.int64)
    _require_dtype(prefix_lens, "prefix_lens", torch.int64)
    _require_dtype(extend_seq_lens, "extend_seq_lens", torch.int64)
    _require_dtype(req_to_token, "req_to_token", torch.int32)
    _require_dtype(lut_tensor, "lut_tensor", torch.int64)
    _require_dtype(
        out_verify_offsets_scratch, "out_verify_offsets_scratch", torch.int64
    )
    _require_dtype(out_write_offsets, "out_write_offsets", torch.int64)
    _require_dtype(
        out_write_seed_slot_indices, "out_write_seed_slot_indices", torch.int64
    )
    _require_dtype(out_verify_num_valid, "out_verify_num_valid", torch.int32)
    _require_dtype(out_verify_enable, "out_verify_enable", torch.int32)
    _require_dtype(out_write_num_valid_reqs, "out_write_num_valid_reqs", torch.int32)

    if bs < 0 or bs > _PLAN_BS_BLOCK_SIZE:
        raise ValueError(
            f"kv-canary: offsets kernel bs must be in [0, {_PLAN_BS_BLOCK_SIZE}], got {bs}"
        )
    if write_offsets_len <= 0:
        raise ValueError(
            f"kv-canary: write_offsets_len must be positive, got {write_offsets_len}"
        )
    if write_req_capacity < 0:
        raise ValueError(
            f"kv-canary: write_req_capacity must be non-negative, got {write_req_capacity}"
        )
    if verify_capacity < 0:
        raise ValueError(
            f"kv-canary: verify_capacity must be non-negative, got {verify_capacity}"
        )
    if req_to_token_stride0 <= 0:
        raise ValueError(
            f"kv-canary: req_to_token_stride0 must be positive, got {req_to_token_stride0}"
        )
    if lut_len < 0:
        raise ValueError(f"kv-canary: lut_len must be non-negative, got {lut_len}")
    if not isinstance(has_swa_lut, bool):
        raise ValueError(
            f"kv-canary: has_swa_lut must be bool, got {type(has_swa_lut).__name__}"
        )
    if has_swa_lut and lut_len <= 0:
        raise ValueError("kv-canary: lut_len must be positive when has_swa_lut is True")
    if not has_swa_lut and lut_len != 0:
        raise ValueError("kv-canary: lut_len must be 0 when has_swa_lut is False")

    _require_len(req_pool_indices, "req_pool_indices", bs)
    _require_len(prefix_lens, "prefix_lens", bs)
    _require_len(extend_seq_lens, "extend_seq_lens", bs)
    _require_2d(req_to_token, "req_to_token")
    _require_min_len(lut_tensor, "lut_tensor", max(lut_len, 1))
    _require_min_len(
        out_verify_offsets_scratch,
        "out_verify_offsets_scratch",
        _PLAN_BS_BLOCK_SIZE + 1,
    )
    _require_len(out_write_offsets, "out_write_offsets", write_offsets_len)
    _require_len(
        out_write_seed_slot_indices,
        "out_write_seed_slot_indices",
        write_req_capacity,
    )
    _require_len(out_verify_num_valid, "out_verify_num_valid", 1)
    _require_len(out_verify_enable, "out_verify_enable", 1)
    _require_len(out_write_num_valid_reqs, "out_write_num_valid_reqs", 1)
    _require_1d(lut_tensor, "lut_tensor")

    if write_offsets_len != write_req_capacity + 1:
        raise ValueError(
            f"kv-canary: write_offsets_len must equal write_req_capacity + 1, got "
            f"{write_offsets_len} and {write_req_capacity}"
        )
    if bs > write_req_capacity:
        raise ValueError(
            f"kv-canary: bs={bs} exceeds write_req_capacity={write_req_capacity}"
        )
    if req_to_token_stride0 != int(req_to_token.stride(0)):
        raise ValueError(
            f"kv-canary: req_to_token_stride0={req_to_token_stride0} does not match "
            f"req_to_token.stride(0)={int(req_to_token.stride(0))}"
        )

    _require_same_device(
        out_verify_offsets_scratch,
        "out_verify_offsets_scratch",
        (
            (req_pool_indices, "req_pool_indices"),
            (prefix_lens, "prefix_lens"),
            (extend_seq_lens, "extend_seq_lens"),
            (req_to_token, "req_to_token"),
            (lut_tensor, "lut_tensor"),
            (out_write_offsets, "out_write_offsets"),
            (out_write_seed_slot_indices, "out_write_seed_slot_indices"),
            (out_verify_num_valid, "out_verify_num_valid"),
            (out_verify_enable, "out_verify_enable"),
            (out_write_num_valid_reqs, "out_write_num_valid_reqs"),
        ),
    )


@triton.jit
def _plan_offsets_kernel(
    # Input pointers.
    req_pool_indices_ptr,
    prefix_lens_ptr,
    extend_seq_lens_ptr,
    req_to_token_ptr,
    full_to_swa_lut_ptr,
    # Output pointers.
    out_verify_offsets_ptr,
    out_write_offsets_ptr,
    out_write_seed_slot_indices_ptr,
    out_verify_num_valid_ptr,
    out_verify_enable_ptr,
    out_write_num_valid_reqs_ptr,
    # Runtime sizes.
    bs,
    req_to_token_stride0,
    swa_lut_len,
    # Compile-time constants.
    BS_BLOCK: tl.constexpr,
    SWA_WINDOW: tl.constexpr,
    HAS_SWA_LUT: tl.constexpr,
    WRITE_OFFSETS_LEN: tl.constexpr,
    WRITE_REQ_CAPACITY: tl.constexpr,
    VERIFY_CAPACITY: tl.constexpr,
    REQ_POOL_IDX_PADDING: tl.constexpr,
    TOKEN_TO_KV_SLOT_PADDING: tl.constexpr,
):
    bs_offs = tl.arange(0, BS_BLOCK)  # [BS_BLOCK]
    bs_mask = bs_offs < bs  # [BS_BLOCK] bool

    # Per-req inputs (int64 for canary-owned metadata; req_to_token keeps its pool dtype).
    rpi = tl.load(
        req_pool_indices_ptr + bs_offs, mask=bs_mask, other=REQ_POOL_IDX_PADDING
    )  # [BS_BLOCK]
    prefix_lens = tl.load(
        prefix_lens_ptr + bs_offs, mask=bs_mask, other=0
    )  # [BS_BLOCK]
    extend_lens = tl.load(
        extend_seq_lens_ptr + bs_offs, mask=bs_mask, other=0
    )  # [BS_BLOCK]

    is_active = (rpi != REQ_POOL_IDX_PADDING) & bs_mask  # [BS_BLOCK] bool
    has_prefix = is_active & (prefix_lens > 0)  # [BS_BLOCK] bool

    window_starts = _compute_window_start(prefix_lens, SWA_WINDOW)  # [BS_BLOCK]

    verify_lens = prefix_lens - window_starts  # [BS_BLOCK]
    verify_lens = tl.where(verify_lens > 0, verify_lens, 0)
    verify_lens = tl.where(is_active, verify_lens, 0)
    verify_exclusive, total_verify = _exclusive_offsets_and_total(verify_lens)

    write_lens = tl.where(extend_lens > 0, extend_lens, 0)  # [BS_BLOCK]
    write_lens = tl.where(is_active, write_lens, 0)
    write_exclusive, total_write = _exclusive_offsets_and_total(write_lens)

    _plan_verify_offsets(
        verify_exclusive,
        total_verify,
        bs_offs,
        bs_mask,
        out_verify_offsets_ptr,
        out_verify_num_valid_ptr,
        out_verify_enable_ptr,
        bs,
        VERIFY_CAPACITY,
    )
    _plan_write_offsets(
        rpi,
        prefix_lens,
        write_lens,
        write_exclusive,
        total_write,
        has_prefix,
        bs_offs,
        bs_mask,
        req_to_token_ptr,
        full_to_swa_lut_ptr,
        out_write_offsets_ptr,
        out_write_seed_slot_indices_ptr,
        out_write_num_valid_reqs_ptr,
        bs,
        req_to_token_stride0,
        swa_lut_len,
        BS_BLOCK,
        HAS_SWA_LUT,
        WRITE_OFFSETS_LEN,
        WRITE_REQ_CAPACITY,
        TOKEN_TO_KV_SLOT_PADDING,
    )


@triton.jit
def _exclusive_offsets_and_total(lens):
    inclusive = tl.cumsum(lens, axis=0)
    return inclusive - lens, tl.sum(lens, axis=0)


@triton.jit
def _plan_verify_offsets(
    verify_exclusive,
    total_verify,
    bs_offs,
    bs_mask,
    out_verify_offsets_ptr,
    out_verify_num_valid_ptr,
    out_verify_enable_ptr,
    bs,
    VERIFY_CAPACITY: tl.constexpr,
):
    tl.store(
        out_verify_offsets_ptr + bs_offs,
        verify_exclusive.to(tl.int64),
        mask=bs_mask,
    )
    tl.store(out_verify_offsets_ptr + bs, total_verify.to(tl.int64))

    # Scalar writes: out_verify_num_valid is clamped to the verify_capacity tensor extent so the verify kernel
    # never indexes past the buffer; enable carries the overflow bit (0 when total_verify > capacity) so the
    # verify kernel skips the whole launch and the host can warn-log this step.
    overflow = total_verify > VERIFY_CAPACITY  # scalar bool
    enable = tl.where(overflow, 0, 1)  # scalar
    clamped = tl.where(overflow, VERIFY_CAPACITY, total_verify)  # scalar
    tl.store(out_verify_num_valid_ptr, clamped.to(tl.int32))
    tl.store(out_verify_enable_ptr, tl.full((), enable, tl.int32))


@triton.jit
def _plan_write_offsets(
    rpi,
    prefix_lens,
    write_lens,
    write_exclusive,
    total_write,
    has_prefix,
    bs_offs,
    bs_mask,
    req_to_token_ptr,
    full_to_swa_lut_ptr,
    out_write_offsets_ptr,
    out_write_seed_slot_indices_ptr,
    out_write_num_valid_reqs_ptr,
    bs,
    req_to_token_stride0,
    swa_lut_len,
    BS_BLOCK: tl.constexpr,
    HAS_SWA_LUT: tl.constexpr,
    WRITE_OFFSETS_LEN: tl.constexpr,
    WRITE_REQ_CAPACITY: tl.constexpr,
    TOKEN_TO_KV_SLOT_PADDING: tl.constexpr,
):
    has_write_contribution = has_prefix & (write_lens > 0)  # [BS_BLOCK] bool

    # Seed slot per req. prefix_lens == 0 means no prefix → -1 sentinel. Padding row → no write contribution
    # → -1 sentinel either way; we also mask write_lens onto seed below to match the ref's "no write → -1".
    safe_prefix_pos = tl.where(prefix_lens > 0, prefix_lens - 1, 0)  # [BS_BLOCK]
    stride_i64 = req_to_token_stride0  # scalar
    seed_full = tl.load(  # [BS_BLOCK]
        req_to_token_ptr + rpi.to(tl.int64) * stride_i64 + safe_prefix_pos.to(tl.int64),
        mask=has_prefix,
        other=TOKEN_TO_KV_SLOT_PADDING,
    )

    if HAS_SWA_LUT:
        seed_translated = _swa_translate_tile(  # [BS_BLOCK]
            seed_full,
            has_prefix,
            full_to_swa_lut_ptr,
            swa_lut_len,
        )
    else:
        seed_translated = seed_full

    # Reqs with no write contribution should expose seed = -1 (ref's _seed_slot is masked by write_lens > 0).
    minus_one = tl.full((BS_BLOCK,), -1, dtype=seed_translated.dtype)  # [BS_BLOCK]
    seed_slot = tl.where(
        has_write_contribution, seed_translated, minus_one
    )  # [BS_BLOCK]

    write_offsets_mask = bs_offs < WRITE_OFFSETS_LEN  # [BS_BLOCK] bool
    tl.store(
        out_write_offsets_ptr + bs_offs,
        write_exclusive.to(tl.int64),
        mask=write_offsets_mask & bs_mask,
    )

    # Store the [bs] slot of out_write_offsets (one element past the last per-req entry).
    # out_write_offsets has length WRITE_OFFSETS_LEN = write_req_capacity + 1; only store if in range.
    write_tail_in_range = bs < WRITE_OFFSETS_LEN  # scalar bool
    tl.store(
        out_write_offsets_ptr + bs,
        total_write.to(tl.int64),
        mask=write_tail_in_range,
    )

    # Scatter seed slots (capped to write_req_capacity).
    seed_mask = bs_mask & (bs_offs < WRITE_REQ_CAPACITY)  # [BS_BLOCK] bool
    tl.store(
        out_write_seed_slot_indices_ptr + bs_offs,
        seed_slot.to(tl.int64),
        mask=seed_mask,
    )

    tl.store(out_write_num_valid_reqs_ptr, tl.full((), bs, tl.int32))