chore: import upstream snapshot with attribution

python/sglang/jit_kernel/kv_canary/plan/__init__.py

@@ -0,0 +1 @@
+from sglang.jit_kernel.kv_canary.plan.api import launch_canary_plan_kernels

python/sglang/jit_kernel/kv_canary/plan/api.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from sglang.jit_kernel.kv_canary.plan.entries_kernel import (
+    launch_plan_entries_kernel,
+)
+from sglang.jit_kernel.kv_canary.plan.offsets_kernel import (
+    _PLAN_BS_BLOCK_SIZE,
+    launch_plan_offsets_kernel,
+)
+from sglang.jit_kernel.kv_canary.verify import VerifyPlan
+from sglang.jit_kernel.kv_canary.write import WritePlan
+
+
+def launch_canary_plan_kernels(
+    *,
+    verify_plan_out: VerifyPlan,
+    write_plan_out: WritePlan,
+    req_pool_indices: torch.Tensor,
+    prefix_lens: torch.Tensor,
+    extend_seq_lens: torch.Tensor,
+    req_to_token: torch.Tensor,
+    swa_window_size: int,
+    full_to_swa_index_mapping: Optional[torch.Tensor],
+    verify_capacity: int,
+    req_to_verify_expected_tokens: Optional[torch.Tensor],
+    req_to_verify_expected_tokens_valid_lens: Optional[torch.Tensor],
+    kv_token_id_vs_position_offset: int,
+) -> None:
+    """Fill verify_plan_out + write_plan_out from normalized canary plan inputs.
+
+    For each req r with req_pool_indices[r] != 0 (0 = padding sentinel):
+
+    - **Verify entries**: one per pos in [window_start, prefix_lens[r]), where window_start = max(0,
+      prefix_lens[r] - swa_window_size) if SWA else 0. slot_idx = req_to_token[req_pool_indices[r], pos]
+      (SWA-translated via full_to_swa_index_mapping if non-None); prev_slot_idx =
+      req_to_token[req_pool_indices[r], pos-1] for pos > 0, else -1. (SWA windows do NOT reset the chain —
+      the writer chains across the entire prefix; sweep verify within an SWA window dereferences the real
+      predecessor for chain-link reconstruction.) Expected-token gather: when
+      ``req_to_verify_expected_tokens`` is supplied, ``expected_input_id =
+      req_to_verify_expected_tokens[rp, pos + kv_token_id_vs_position_offset]`` when ``0 <= pos +
+      kv_token_id_vs_position_offset < req_to_verify_expected_tokens_valid_lens[r]``, else the ``-1``
+      sentinel (which the verify kernel treats as "skip token-id check").
+    - **Write metadata** (when extend_seq_lens[r] > 0): contribute extend_seq_lens[r] to the per-req
+      write count (for write_offsets cumsum). Per-req chain seed = req_to_token[req_pool_indices[r],
+      prefix_lens[r]-1] (SWA-translated), or -1 if prefix_lens[r] == 0. Per-token write data
+      (input_ids / positions / out_cache_loc) is NOT materialized here — launch_canary_write_kernel
+      reads it directly from ForwardBatch via write_offsets.
+
+    Args:
+        verify_plan_out: Pre-allocated VerifyPlan; filled in-place.
+        write_plan_out: Pre-allocated WritePlan; filled in-place.
+        req_pool_indices: Per-row ReqToTokenPool row index, shape [bs], int64. 0 is the padding sentinel.
+        prefix_lens: Per-req prefix length already written before this step, shape [bs], int64.
+        extend_seq_lens: Per-req tokens being written this step, shape [bs], int64.
+        req_to_token: ReqToTokenPool.req_to_token; full-pool slot index table, shape [max_reqs, max_seq_len],
+            int32.
+        swa_window_size: 0 for the FULL canary group; positive window length for the SWA group.
+        full_to_swa_index_mapping: SWA LUT, shape [full_pool_size + 1], int64, or None. Required (non-None) iff
+            swa_window_size > 0. Used to translate verify slot indices and chain-seed slot indices at plan time.
+            Loaded element-typed via Triton ``tl.load``; intermediate translated slot values are int64 inside the
+            kernel and stored in the int64 plan schema.
+        verify_capacity: Length of verify_plan_out.verify_*; on overflow the offsets kernel clears
+            verify_enable and plan_entries skips the scatter.
+        req_to_verify_expected_tokens: Optional source-of-truth token pool, shape [max_reqs, max_context_len],
+            int32. When supplied, the plan kernel gathers expected_input_id for each verify entry from
+            ``[rp, pos + kv_token_id_vs_position_offset]``; when None, every entry gets the ``-1`` sentinel.
+        req_to_verify_expected_tokens_valid_lens: Per-req snapshot length on ``req_to_verify_expected_tokens``,
+            shape [bs], int64. Required iff ``req_to_verify_expected_tokens`` is set. Reads past
+            ``valid_lens[r]`` skip the gather (emit ``-1``) — this is what makes the plan kernel correct in the
+            presence of EAGLE draft / verify positions written past the committed history, and across pool
+            rows recycled from a longer previous owner whose stale tail still lives at high indices.
+        kv_token_id_vs_position_offset: Per-buffer-group logical-position offset applied to ``pos`` before
+            indexing ``req_to_verify_expected_tokens``. 0 for target pools; +1 for EAGLE draft.
+
+    Implementation:
+        - Two sub-kernels launched in sequence:
+          1. Triton ``_plan_offsets_kernel`` (1-D grid ``(1,)``, single program over all ``bs`` reqs):
+             reads req_pool_indices[r], prefix_lens[r], extend_seq_lens[r] for each r; computes
+             verify_count = (prefix_lens - window_start) and write_count = extend_seq_lens (both 0 if rp == 0
+             padding); gathers seed_slot_full = req_to_token[rp, prefix_lens - 1] (or -1 if prefix_lens == 0),
+             SWA-translates seed_slot via full_to_swa_index_mapping[seed_slot_full] if non-None; runs
+             block-level cumsum (``tl.cumsum``) to produce verify_offsets[_PLAN_BS_BLOCK_SIZE + 1] and
+             write_plan_out.write_offsets[write_req_capacity + 1] in-place; scatters write_seed slots; writes
+             scalar totals ``verify_plan_out.verify_num_valid`` and ``write_plan_out.write_num_valid_reqs``.
+          2. CUDA ``plan_entries_persistent_kernel`` (1-D persistent grid sized to ``num_sms *
+             kBlocksPerSm`` blocks of ``kBlockSize`` threads), wrapped by Python
+             ``launch_plan_entries_kernel``: each thread grid-strides over ``tid ∈ [0, total_verify)``,
+             locates its owning req via ``find_req_id`` (binary search on verify_offsets), computes
+             out_position = window_start[req_id] + (tid - verify_offsets[req_id]), gathers slot =
+             req_to_token[rp, out_position] (SWA-translated when ``HAS_SWA_LUT``), prev_slot =
+             req_to_token[rp, out_position - 1] when out_position > 0 (also translated) else -1, and
+             scatters (slot, position, prev_slot) into verify_plan_out at flat index tid.
+        - All output tensors are addressed at addresses baked into the cuda-graph capture.
+
+    Calling contract:
+        - Pure side-effect; no host work, no D2H.
+        - Safe in cuda-graph capture; caller refills all input tensors in-place before replay.
+        - The wrapper launches the plan sub-kernels needed to fill both plans end-to-end.
+        - Padding rows contribute zero entries.
+
+    Pinned by Python reference
+    :func:`sglang.jit_kernel.kv_canary.plan_ref.launch_canary_plan_kernels_torch_reference`; both the Triton
+    offsets kernel and the CUDA JIT entries kernel must match byte-for-byte.
+    """
+    bs = int(req_pool_indices.shape[0])
+    if bs > _PLAN_BS_BLOCK_SIZE:
+        raise ValueError(
+            f"kv-canary: launch_canary_plan_kernels supports at most bs={_PLAN_BS_BLOCK_SIZE} reqs per launch, "
+            f"got bs={bs}. Bump _PLAN_BS_BLOCK_SIZE if real workloads need this."
+        )
+    if swa_window_size > 0 and full_to_swa_index_mapping is None:
+        raise ValueError(
+            "kv-canary: launch_canary_plan_kernels requires full_to_swa_index_mapping when swa_window_size > 0"
+        )
+
+    device = verify_plan_out.verify_slot_indices.device
+    verify_offsets_scratch = torch.empty(
+        _PLAN_BS_BLOCK_SIZE + 1, dtype=torch.int64, device=device
+    )
+
+    plan_verify_capacity = int(verify_plan_out.verify_slot_indices.shape[0])
+    if verify_capacity != plan_verify_capacity:
+        raise ValueError(
+            f"kv-canary: launch_canary_plan_kernels verify_capacity={verify_capacity} does not match "
+            f"verify_plan_out.verify_slot_indices.shape[0]={plan_verify_capacity}"
+        )
+
+    write_plan_out.write_offsets.zero_()
+
+    launch_plan_offsets_kernel(
+        req_pool_indices=req_pool_indices,
+        prefix_lens=prefix_lens,
+        extend_seq_lens=extend_seq_lens,
+        req_to_token=req_to_token,
+        full_to_swa_index_mapping=full_to_swa_index_mapping,
+        out_verify_offsets_scratch=verify_offsets_scratch,
+        out_write_offsets=write_plan_out.write_offsets,
+        out_write_seed_slot_indices=write_plan_out.write_seed_slot_indices,
+        out_verify_num_valid=verify_plan_out.verify_num_valid,
+        out_verify_enable=verify_plan_out.enable,
+        out_write_num_valid_reqs=write_plan_out.write_num_valid_reqs,
+        swa_window_size=int(swa_window_size),
+        verify_capacity=verify_capacity,
+    )
+
+    launch_plan_entries_kernel(
+        req_pool_indices=req_pool_indices,
+        prefix_lens=prefix_lens,
+        req_to_token=req_to_token,
+        full_to_swa_index_mapping=full_to_swa_index_mapping,
+        verify_offsets_scratch=verify_offsets_scratch,
+        verify_enable=verify_plan_out.enable,
+        req_to_verify_expected_tokens=req_to_verify_expected_tokens,
+        req_to_verify_expected_tokens_valid_lens=req_to_verify_expected_tokens_valid_lens,
+        out_verify_slot_indices=verify_plan_out.verify_slot_indices,
+        out_verify_expected_tokens=verify_plan_out.verify_expected_tokens,
+        out_verify_expected_positions=verify_plan_out.verify_expected_positions,
+        out_verify_prev_slot_indices=verify_plan_out.verify_prev_slot_indices,
+        kv_token_id_vs_position_offset=int(kv_token_id_vs_position_offset),
+        swa_window_size=int(swa_window_size),
+    )

python/sglang/jit_kernel/kv_canary/plan/entries_kernel.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_plan_entries_module(
+    has_swa_lut: bool, has_verify_expected_token_pool: bool
+) -> Module:
+    args = make_cpp_args(has_swa_lut, has_verify_expected_token_pool)
+    return load_jit(
+        "kv_canary_plan_entries",
+        *args,
+        cuda_files=["kv_canary/canary_plan_entries.cuh"],
+        cuda_wrappers=[
+            ("plan_entries", f"PlanEntriesKernel<{args}>::run"),
+        ],
+    )
+
+
+def launch_plan_entries_kernel(
+    *,
+    req_pool_indices: torch.Tensor,
+    prefix_lens: torch.Tensor,
+    req_to_token: torch.Tensor,
+    full_to_swa_index_mapping: Optional[torch.Tensor],
+    verify_offsets_scratch: torch.Tensor,
+    verify_enable: torch.Tensor,
+    req_to_verify_expected_tokens: Optional[torch.Tensor],
+    req_to_verify_expected_tokens_valid_lens: Optional[torch.Tensor],
+    out_verify_slot_indices: torch.Tensor,
+    out_verify_expected_tokens: torch.Tensor,
+    out_verify_expected_positions: torch.Tensor,
+    out_verify_prev_slot_indices: torch.Tensor,
+    kv_token_id_vs_position_offset: int,
+    swa_window_size: int,
+) -> None:
+    has_swa_lut = full_to_swa_index_mapping is not None
+    has_verify_expected_token_pool = req_to_verify_expected_tokens is not None
+    if (
+        has_verify_expected_token_pool
+        and req_to_verify_expected_tokens_valid_lens is None
+    ):
+        raise ValueError(
+            "kv-canary: launch_plan_entries_kernel requires "
+            "req_to_verify_expected_tokens_valid_lens when req_to_verify_expected_tokens is set"
+        )
+    module = _jit_plan_entries_module(has_swa_lut, has_verify_expected_token_pool)
+    module.plan_entries(
+        req_pool_indices,
+        prefix_lens,
+        req_to_token,
+        full_to_swa_index_mapping,
+        verify_offsets_scratch,
+        verify_enable,
+        req_to_verify_expected_tokens,
+        req_to_verify_expected_tokens_valid_lens,
+        out_verify_slot_indices,
+        out_verify_expected_tokens,
+        out_verify_expected_positions,
+        out_verify_prev_slot_indices,
+        int(kv_token_id_vs_position_offset),
+        int(swa_window_size),
+    )

python/sglang/jit_kernel/kv_canary/plan/offsets_kernel.py

@@ -0,0 +1,441 @@
+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))

python/sglang/jit_kernel/kv_canary/plan/utils.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+
+def _resolve_swa_lut(
+    lut: Optional[torch.Tensor], device: torch.device
+) -> tuple[torch.Tensor, int, bool]:
+    """Return the (tensor, length, has_lut) triple to launch the plan kernel with.
+
+    Triton requires a valid tensor pointer at every kernel-arg slot even when ``HAS_SWA_LUT`` is False, so
+    when the caller passes ``None`` we substitute a one-element sentinel tensor and set ``lut_len=0``;
+    the kernel's constexpr branch guarantees no dereference happens. Dtype matches the production LUT
+    (int64) so Triton ``tl.load`` element typing stays consistent.
+    """
+    if lut is not None:
+        return lut, int(lut.shape[0]), True
+    return torch.zeros(1, dtype=torch.int64, device=device), 0, False
+
+
+def _require_dtype(tensor: torch.Tensor, name: str, dtype: torch.dtype) -> None:
+    if tensor.dtype != dtype:
+        raise ValueError(
+            f"kv-canary: {name} must have dtype {dtype}, got {tensor.dtype}"
+        )
+
+
+def _require_1d(tensor: torch.Tensor, name: str) -> None:
+    if tensor.ndim != 1:
+        raise ValueError(
+            f"kv-canary: {name} must be 1-D, got shape {tuple(tensor.shape)}"
+        )
+
+
+def _require_2d(tensor: torch.Tensor, name: str) -> None:
+    if tensor.ndim != 2:
+        raise ValueError(
+            f"kv-canary: {name} must be 2-D, got shape {tuple(tensor.shape)}"
+        )
+
+
+def _require_len(tensor: torch.Tensor, name: str, expected: int) -> None:
+    _require_1d(tensor=tensor, name=name)
+    actual = int(tensor.shape[0])
+    if actual != expected:
+        raise ValueError(f"kv-canary: {name} length must be {expected}, got {actual}")
+
+
+def _require_min_len(tensor: torch.Tensor, name: str, minimum: int) -> None:
+    _require_1d(tensor=tensor, name=name)
+    actual = int(tensor.shape[0])
+    if actual < minimum:
+        raise ValueError(f"kv-canary: {name} length must be >= {minimum}, got {actual}")
+
+
+def _require_same_device(
+    reference: torch.Tensor,
+    reference_name: str,
+    tensors: tuple[tuple[torch.Tensor, str], ...],
+) -> None:
+    for tensor, name in tensors:
+        if tensor.device != reference.device:
+            raise ValueError(
+                f"kv-canary: {name} must be on {reference_name}'s device "
+                f"{reference.device}, got {tensor.device}"
+            )
+
+
+@triton.jit
+def _compute_window_start(prefix_lens, SWA_WINDOW: tl.constexpr):
+    """Per-req window start: max(prefix_lens - SWA_WINDOW, 0) when SWA, else 0.
+    Works for tile and scalar inputs (broadcasts via prefix_lens shape).
+    """
+    if SWA_WINDOW > 0:
+        clipped = prefix_lens - SWA_WINDOW
+        return tl.where(clipped > 0, clipped, 0)
+    else:
+        return prefix_lens - prefix_lens
+
+
+@triton.jit
+def _swa_translate_tile(raw, mask, lut_ptr, lut_len):
+    """SWA-translate a tile of slot indices. Sentinels (raw < 0) are passed through unchanged.
+
+    ``lut_len`` is the LUT's length (Python int from the host wrapper); when 0 the LUT is unused (the caller
+    will only enter this branch when HAS_SWA_LUT is True, so lut_len is always > 0 in practice).
+    """
+    sentinel = raw < 0
+    safe = tl.where(sentinel, 0, raw)
+    if lut_len > 0:
+        safe = tl.where(safe >= lut_len, lut_len - 1, safe)
+    xlat = tl.load(lut_ptr + safe, mask=mask & (~sentinel), other=0)
+    return tl.where(sentinel, raw, xlat)