from __future__ import annotations import random from typing import Literal, Optional import torch from sglang.jit_kernel.kv_canary.verify import ( RealKvSource, VerifyPlan, ) from sglang.jit_kernel.kv_canary.write import WritePlan from sglang.jit_kernel.tests.kv_canary._constants import DEFAULT_NUM_SLOTS _DEVICE = torch.device("cuda") LutKind = Literal["identity", "shift", "permutation", "with_oob"] def make_lut( *, kind: LutKind, pool_size: int, device: torch.device, rng: Optional[random.Random] = None, ) -> torch.Tensor: base = torch.arange(pool_size + 1, dtype=torch.int64, device=device) if kind == "identity": return base.contiguous() if kind == "shift": return (base + 100).contiguous() if kind in ("permutation", "with_oob"): if rng is None: rng = random.Random(0) perm = list(range(pool_size + 1)) rng.shuffle(perm) out = torch.tensor(perm, dtype=torch.int64, device=device) if kind == "with_oob": out[-1] = pool_size + 999 return out.contiguous() raise ValueError(f"unknown LutKind: {kind}") ReqToTokenKind = Literal["linear", "sparse_permuted"] def make_req_to_token( *, kind: ReqToTokenKind, max_reqs: int, max_seq_len: int, device: torch.device, rng: Optional[random.Random] = None, ) -> torch.Tensor: if kind == "linear": rp_axis = torch.arange(max_reqs, device=device, dtype=torch.int32).unsqueeze(1) pos_axis = torch.arange( max_seq_len, device=device, dtype=torch.int32 ).unsqueeze(0) return (rp_axis * max_seq_len + pos_axis).contiguous() if rng is None: rng = random.Random(0) pool_size = max_reqs * max_seq_len # Slots index into a full_to_swa LUT sized [pool_size + 1], so values must stay # in [0, pool_size]. The universe spans [1, pool_size] (skipping 0 as reserved), # giving exactly max_reqs * max_seq_len unique slots — one per (rp, pos) cell. slot_universe = list(range(1, pool_size + 1)) rng.shuffle(slot_universe) rtt = torch.zeros((max_reqs, max_seq_len), dtype=torch.int32, device=device) cursor = 0 for rp in range(max_reqs): per_req = slot_universe[cursor : cursor + max_seq_len] cursor += max_seq_len rtt[rp, :] = torch.tensor(per_req, dtype=torch.int32, device=device) return rtt.contiguous() def make_real_kv_source( *, num_slots: int = DEFAULT_NUM_SLOTS, num_bytes_per_token: int = 16, page_size: int = 1, read_bytes: Optional[int] = None, pad_dim1: int = 0, device: torch.device, fill: int = 0, ) -> RealKvSource: """Allocate one RealKvSource with the canonical [num_rows, dim1_bytes] uint8 shape. ``pad_dim1`` adds trailing per-row bytes the canary should skip — used by the "holey dim 1" case to confirm the kernel never reads past ``page_size * num_bytes_per_token``. """ num_rows = (num_slots + page_size - 1) // page_size cols = page_size * num_bytes_per_token + pad_dim1 tensor = torch.full( (num_rows, cols), fill_value=fill, dtype=torch.uint8, device=device ) effective_read = read_bytes if read_bytes is not None else num_bytes_per_token return RealKvSource( tensor=tensor, page_size=page_size, num_bytes_per_token=num_bytes_per_token, read_bytes=effective_read, ) FillStrategy = Literal["constant_per_source", "random_bytes"] def make_real_kv_sources( *, count: int, num_bytes_per_token: int = 16, page_size: int = 1, num_slots: int = DEFAULT_NUM_SLOTS, device: torch.device, rng: Optional[random.Random] = None, fill_strategy: FillStrategy = "constant_per_source", ) -> tuple[RealKvSource, ...]: sources: list[RealKvSource] = [] for i in range(count): read_bytes_eff = num_bytes_per_token src = make_real_kv_source( num_slots=num_slots, num_bytes_per_token=num_bytes_per_token, page_size=page_size, read_bytes=read_bytes_eff, device=device, fill=(i + 1) * 17, ) if fill_strategy == "random_bytes": if rng is None: rng = random.Random(0) seed = rng.randint(0, 0xFFFFFFFF) gen = torch.Generator(device=device).manual_seed(seed) src.tensor.random_(generator=gen) sources.append(src) return tuple(sources) def clone_real_kv_sources( sources: tuple[RealKvSource, ...], ) -> tuple[RealKvSource, ...]: return tuple( RealKvSource( tensor=src.tensor.clone(), page_size=src.page_size, num_bytes_per_token=src.num_bytes_per_token, read_bytes=src.read_bytes, ) for src in sources ) PaddingKind = Literal["none", "trailing", "interleaved"] def make_padding_mask( *, bs: int, kind: PaddingKind, rng: Optional[random.Random] = None, padding_fraction: float = 0.25, ) -> list[bool]: if bs == 0: return [] if kind == "none": return [False] * bs n_pad = max(1, int(bs * padding_fraction)) if bs > 0 else 0 n_pad = min(n_pad, bs) if kind == "trailing": return [False] * (bs - n_pad) + [True] * n_pad if kind == "interleaved": if rng is None: rng = random.Random(0) mask = [False] * bs chosen = rng.sample(range(bs), k=n_pad) for idx in chosen: mask[idx] = True return mask raise ValueError(f"unknown PaddingKind: {kind}") CapacityKind = Literal["loose", "tight_match", "under_by_one"] def derive_plan_capacity( *, kind: CapacityKind, total_verify: int, extras_count: int, bs: int, ) -> tuple[int, int]: needed = total_verify + extras_count if kind == "loose": return max(needed + 64, 128), max(bs + 4, 8) if kind == "tight_match": return max(needed, 1), max(bs + 4, 8) if kind == "under_by_one": return max(needed - 1, 1), max(bs + 4, 8) raise ValueError(f"unknown CapacityKind: {kind}") def allocate_plan_pair( *, verify_capacity: int, write_req_capacity: int, ) -> tuple[VerifyPlan, WritePlan, VerifyPlan, WritePlan]: return ( VerifyPlan.allocate(verify_capacity=verify_capacity, device=_DEVICE), WritePlan.allocate(write_req_capacity=write_req_capacity, device=_DEVICE), VerifyPlan.allocate(verify_capacity=verify_capacity, device=_DEVICE), WritePlan.allocate(write_req_capacity=write_req_capacity, device=_DEVICE), ) def empty_extras() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: return ( torch.zeros(1, dtype=torch.int64, device=_DEVICE), torch.zeros(1, dtype=torch.int64, device=_DEVICE), torch.zeros(1, dtype=torch.int64, device=_DEVICE), torch.zeros(1, dtype=torch.int32, device=_DEVICE), ) def dummy_pseudo_tensors(num_tokens: int) -> tuple[torch.Tensor, torch.Tensor]: return ( torch.zeros(num_tokens, dtype=torch.int64, device=_DEVICE), torch.zeros(num_tokens, dtype=torch.int64, device=_DEVICE), )