sgl-project--sglang/python/sglang/jit_kernel/kv_canary/write.py

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING

import torch

from sglang.jit_kernel.kv_canary import consts
from sglang.jit_kernel.kv_canary.verify import (
    VerifyOrWriteContext,
    _assert_contiguous,
    _build_real_kv_source_abi,
)
from sglang.jit_kernel.utils import cache_once, load_jit

if TYPE_CHECKING:
    from tvm_ffi.module import Module


@dataclass(frozen=True, slots=True, kw_only=True)
class WritePlan:
    """Write plan consumed by launch_canary_write_kernel: per-token slot indices + per-req metadata.

    Fully per-req — no per-token tile. launch_canary_write_kernel uses write_offsets to map each thread's
    (req, j) into a flat index i, then reads token-level data from input_ids / positions /
    out_cache_loc[i] directly.
    SWA translation of per-token slots is done **host-side by the caller** (typically the endpoint) before
    invoking launch_canary_write_kernel — the kernel is SWA-agnostic and only understands "slot ≥ 0 ⇒ write;
    slot < 0 ⇒ skip this entry". Only the chain-seed slot (a per-req gather from req_to_token at plan time)
    is SWA-translated by the plan kernel and lives in write_seed_slot_indices.

    Req r's write entries occupy flat indices [write_offsets[r], write_offsets[r+1]). seed_slot_idx == -1 means
    K_req_old == 0 (anchor on CANARY_CHAIN_ANCHOR).

    Fields:
        write_offsets: Exclusive prefix-sum offsets indexing into ForwardBatch's input_ids / positions /
            out_cache_loc, shape [write_req_capacity + 1], int64. write_offsets[0] == 0;
            write_offsets[write_num_valid_reqs[0]] == total_write_entries.
        write_seed_slot_indices: Chain-seed slot per write req, shape [write_req_capacity], int64. Already
            SWA-translated. -1 = no prefix (chain anchors on CANARY_CHAIN_ANCHOR).
        write_num_valid_reqs: Active write-req count, shape [1], int32. launch_canary_write_kernel skips blocks
            with block_id >= write_num_valid_reqs[0].
    """

    write_offsets: torch.Tensor
    write_seed_slot_indices: torch.Tensor
    write_num_valid_reqs: torch.Tensor

    @classmethod
    def allocate(
        cls,
        *,
        write_req_capacity: int,
        device: torch.device,
    ) -> WritePlan:
        if write_req_capacity <= 0:
            raise ValueError(
                f"kv-canary: WritePlan write_req_capacity must be positive, got {write_req_capacity}"
            )
        return cls(
            write_offsets=torch.empty(
                write_req_capacity + 1, dtype=torch.int64, device=device
            ),
            write_seed_slot_indices=torch.empty(
                write_req_capacity, dtype=torch.int64, device=device
            ),
            write_num_valid_reqs=torch.empty(1, dtype=torch.int32, device=device),
        )

    def zero_for_testing_(self) -> WritePlan:
        """WARN: ONLY use it when testing plan kernel. Do not use it when testing verify or
        write kernel to avoid hiding bugs."""
        self.write_offsets.zero_()
        self.write_seed_slot_indices.zero_()
        self.write_num_valid_reqs.zero_()
        return self


def launch_canary_write_kernel(
    *,
    context: VerifyOrWriteContext,
    plan: WritePlan,
    input_ids: torch.Tensor,
    positions: torch.Tensor,
    out_cache_loc: torch.Tensor,
    enable_write_input_assert: bool,
    expected_input_tokens: torch.Tensor | None,
    expected_input_positions: torch.Tensor | None,
) -> None:
    """Write canary fingerprints into one canary buffer per a WritePlan.

    Grid: one CUDA block per active write req, single thread per block (chain is intrinsically serial).
    Block r walks entries ``[plan.write_offsets[r], plan.write_offsets[r+1])``. Per chain step ``i``:

    - ``slot`` = ``out_cache_loc[i]`` (caller-pre-translated for SWA groups; entries set to -1 are skipped).
    - ``token / position`` = ``input_ids[i] / positions[i]``.
    - ``real_kv_hash`` = ``real_kv_fold_sources(real_kv_sources, slot)`` if ``real_kv_hash_mode != NONE`` else 0.
    - Store 4 int64s ``(token, position, running_prev_hash, real_kv_hash)`` into ``canary_buf[slot]``.
    - Advance ``running_prev_hash = splitmix64_mix3(prev, token, position)``, where
      splitmix64_mix3 folds each input via ``acc = splitmix64(acc ^ next)`` starting from ``splitmix64(prev)``.
      ``real_kv_hash`` is intentionally not folded into the chain — see ``compute_slot_hash`` in
      ``csrc/kv_canary/canary_common.cuh`` for the radix-folding rationale.

    Initial ``running_prev_hash`` when ``seed_slot_idx >= 0``: load (token, position, prev_hash) from
    ``canary_buf[plan.write_seed_slot_indices[r]]`` and set
    ``running_prev_hash = splitmix64_mix3(seed.prev_hash, seed.token, seed.position)``
    (i.e. apply the same advance step that produced ``seed``'s successor — this keeps slot[0]'s stored
    ``prev_hash`` consistent with the chain link). Else
    ``running_prev_hash = splitmix64(CANARY_CHAIN_ANCHOR)``. ``write_seed_slot_indices`` is already
    SWA-translated by the plan kernel; ``CANARY_CHAIN_ANCHOR`` is hardcoded module-level (no runtime seed).

    Write-time input verification (caller-driven, kernel is oracle-agnostic): when
    ``enable_write_input_assert`` is True the kernel additionally compares ``input_ids[i]`` against
    ``expected_input_tokens[i]`` and ``positions[i]`` against ``expected_input_positions[i]``; mismatch
    on either field records a violation. The chain still advances on the actual values (not the expected
    ones) so a downstream verify won't cascade. Whoever produced the expected tensors is responsible for
    filling them; the kernel runs no oracle internally.

    Write only writes canary_buf (reads only at seed slots). Block uses no shared memory.

    The ForwardBatch-derived arguments are passed through unchanged from the source ForwardBatch — canary does not transform
    them.

    Args:
        context: Shared verify/write launch context, including canary buffer, launch tag, violation sink,
            health counters, and real KV fingerprint sources.
        plan: Pre-allocated WritePlan.
        input_ids: ForwardBatch.input_ids; token ids being written, shape [num_tokens_padded], int64.
            Flattened across reqs in plan.write_offsets order; tail beyond
            plan.write_offsets[plan.write_num_valid_reqs[0]] is cuda-graph padding.
        positions: ForwardBatch.positions; sequence positions of input_ids, shape [num_tokens_padded], int64.
        out_cache_loc: Per-token canary slot index, shape [num_tokens_padded], int64. The caller is
            responsible for translating ForwardBatch.out_cache_loc into the canary's index space for SWA
            groups (typically a host-side LUT gather in the endpoint); FULL groups pass it through
            unchanged. A -1 entry signals skip-this-token (used for SWA out-of-window slots or padding).
            The kernel does not consult any LUT.
        enable_write_input_assert: bool toggle. False = expected_input_* tensors must be None. True = compare
            each chain step's actual (token, position) against the caller-supplied expected tensors below.
        expected_input_tokens: Expected token id per write entry, shape [num_tokens_padded], int64. Only read
            when enable_write_input_assert is True; must be None when enable_write_input_assert is False.
            Layout mirrors input_ids (flattened across reqs in plan.write_offsets order); padding tail
            is ignored. Filled by the caller from whichever oracle produces expected inputs — the kernel
            knows no oracle.
        expected_input_positions: Expected position per write entry, shape [num_tokens_padded], int64, or None.
            Same shape/layout/lifetime rules as expected_input_tokens.

    Implementation:
        - CUDA __global__ `canary_write_kernel`: 1-D grid `(write_req_capacity, 1, 1)` blocks × `(1, 1, 1)` thread
          per block. block_id r = blockIdx.x = one write req; chains are intrinsically serial so a single thread
          per block is optimal (warp-level parallelism would idle 31 lanes).
        - Per block, early-exit on r >= plan.write_num_valid_reqs[0]. Else load entry_start = plan.write_offsets[r],
          entry_count = plan.write_offsets[r+1] - entry_start, seed_slot_idx = plan.write_seed_slot_indices[r] into
          registers.
        - Initialize running_prev_hash: if seed_slot_idx >= 0, load (token, position, prev_hash) from
          canary_buf[seed_slot_idx] and set running_prev_hash = splitmix64_mix3(prev_hash, token, position);
          else running_prev_hash = splitmix64(kCanaryChainAnchor).
        - Serial chain loop `for j in range(entry_count)`:
              i = entry_start + j;
              slot = out_cache_loc[i];  // caller-pre-translated; the kernel never consults a LUT
              if (slot < 0) continue;       // -1 sentinel = skip (SWA out-of-window or padding)
              token = input_ids[i]; position = positions[i];
              real_kv_hash = (real_kv_hash_mode == NONE) ? 0 : real_kv_fold_sources(real_kv_sources, slot);
                  // applies RealKvSource access invariant
              if enable_write_input_assert:
                  if token != expected_input_tokens[i] or position != expected_input_positions[i]:
                      record_violation();  // chain still advances on the ACTUAL (token, position) below
              store (token, position, running_prev_hash, real_kv_hash) to canary_buf[slot] as 4 int64 fields;
              running_prev_hash = splitmix64_mix3(running_prev_hash, token, position);
        - All chain state lives in the block's single thread's registers. No shared memory, no cross-block
          coordination.
        - record_violation() identical to verify (atomicAdd + atomic-write).
        - Counters: thread of block 0 does atomicAdd(kernel_run_counter, 1); each block accumulates its
          entry_count and atomicAdds to slot_run_counter once at exit.

    Calling contract:
        - Pure side-effect; never raises.
        - Input-verification mismatch records violations but does NOT abort the chain.
        - kernel_run_counter is bumped every call.
        - Safe in cuda-graph capture; caller refills input_ids / positions / out_cache_loc / plan
          in-place before replay.

    Pinned by torch reference
    :func:`sglang.jit_kernel.kv_canary.write_ref.launch_canary_write_kernel_torch_reference`; CUDA must match
    byte-for-byte.
    """
    canary_buf = context.canary_buf
    real_kv_sources = context.real_kv_sources
    if len(real_kv_sources) > consts.MAX_REAL_KV_SOURCES:
        raise ValueError(
            f"kv-canary: at most {consts.MAX_REAL_KV_SOURCES} RealKvSource entries supported by the CUDA ABI, "
            f"got {len(real_kv_sources)}"
        )

    _assert_contiguous(canary_buf, "canary_buf")
    _assert_contiguous(plan.write_offsets, "plan.write_offsets")
    _assert_contiguous(plan.write_seed_slot_indices, "plan.write_seed_slot_indices")
    _assert_contiguous(plan.write_num_valid_reqs, "plan.write_num_valid_reqs")
    _assert_contiguous(input_ids, "input_ids")
    _assert_contiguous(positions, "positions")
    _assert_contiguous(out_cache_loc, "out_cache_loc")
    if enable_write_input_assert:
        if expected_input_tokens is None or expected_input_positions is None:
            raise ValueError(
                "kv-canary: expected input tensors are required when enable_write_input_assert=True"
            )
        _assert_contiguous(expected_input_tokens, "expected_input_tokens")
        _assert_contiguous(expected_input_positions, "expected_input_positions")
    else:
        if expected_input_tokens is not None or expected_input_positions is not None:
            raise ValueError(
                "kv-canary: expected input tensors must be None when enable_write_input_assert=False"
            )
    _assert_contiguous(context.violation_ring, "violation_ring")
    _assert_contiguous(context.violation_write_index, "violation_write_index")
    _assert_contiguous(context.slot_run_counter, "slot_run_counter")
    _assert_contiguous(context.kernel_run_counter, "kernel_run_counter")
    _assert_contiguous(
        context.enable_chain_position_assert, "enable_chain_position_assert"
    )

    padded_bufs, source_params = _build_real_kv_source_abi(
        real_kv_sources=real_kv_sources, device=canary_buf.device
    )

    module = _jit_canary_write_module()
    module.canary_write_step_cuda(
        canary_buf,
        plan.write_offsets,
        plan.write_seed_slot_indices,
        plan.write_num_valid_reqs,
        input_ids,
        positions,
        out_cache_loc,
        int(context.kernel_kind),
        int(enable_write_input_assert),
        expected_input_tokens,
        expected_input_positions,
        context.violation_ring,
        context.violation_write_index,
        context.slot_run_counter,
        context.kernel_run_counter,
        context.enable_chain_position_assert,
        padded_bufs[0],
        padded_bufs[1],
        padded_bufs[2],
        padded_bufs[3],
        source_params,
        len(real_kv_sources),
        int(context.real_kv_hash_mode),
    )


@cache_once
def _jit_canary_write_module() -> Module:
    return load_jit(
        "kv_canary_write",
        cuda_files=["kv_canary/canary_write.cuh"],
        cuda_wrappers=[
            ("canary_write_step_cuda", "canary::canary_write_step_cuda"),
        ],
    )