/python/sglang/srt/kv_canary/state.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Optional

import torch

from sglang.jit_kernel.kv_canary.consts import VIOLATION_FIELDS
from sglang.srt.kv_canary.config import CanaryConfig


@dataclass(frozen=True, slots=True, kw_only=True)
class ViolationLog:
    """Global violation sink shared across all canary launches.

    One instance per canary runner — every launch (head / tail / sweep, K / V half, FULL / SWA group) writes
    into the same ring. The kernel_kind field stamped into each violation row identifies which launch fired
    (kernel_kind is a static IntEnum tag — :class:`CanaryLaunchTag` in
    ``sglang.jit_kernel.kv_canary.verify`` — with a unique value per (head|tail|sweep, K|V, FULL|SWA) tuple).

    Ring capacity is sized generously (≥ 1024) so overflow is a non-concern in practice — violations are
    cold-path and the host raises at the first one anyway (or just logs it in mode="log"). atomicAdd
    contention on a single counter is also negligible since violation events are rare.

    Derived state (host computes on read; not stored):
        is_errored = violation_write_index[0] > 0
        first_violation = violation_ring[0]   (valid iff is_errored)
        ring_valid_count = min(violation_write_index[0], ring_capacity)

    The ring is fill-once: writes beyond ring_capacity are dropped but the counter still increments. Whoever
    wins atomicAdd for idx == 0 permanently occupies row 0.

    Fields:
        violation_ring: Append-only violation sink, shape [ring_capacity, VIOLATION_FIELDS], int64. Row 0 is
            the first violation; rows 1..min(write_index, capacity) follow in atomic order. Fill-once.
        violation_write_index: Monotonic violation counter, shape [1], int32. Incremented on every violation
            regardless of ring capacity.
    """

    violation_ring: torch.Tensor
    violation_write_index: torch.Tensor

    @classmethod
    def allocate(cls, *, ring_capacity: int, device: torch.device) -> ViolationLog:
        if ring_capacity <= 0:
            raise ValueError(
                f"kv-canary: ViolationLog ring_capacity must be positive, got {ring_capacity}"
            )
        return cls(
            violation_ring=torch.zeros(
                ring_capacity, VIOLATION_FIELDS, dtype=torch.int64, device=device
            ),
            violation_write_index=torch.zeros(1, dtype=torch.int32, device=device),
        )


@dataclass(frozen=True, slots=True, kw_only=True)
class CanaryDeviceState:
    """Device-side state owned by one CanaryManager instance.

    One instance per ModelRunner. Held on the same device as the KV pool. All tensors are allocated up
    front (sizes fixed by CanaryConfig + cuda-graph capture capacity) and reused across forward steps —
    no per-step allocation.

    Fields:
        violation_log: The single ViolationLog shared by every launch (head / tail / sweep × K / V ×
            FULL / SWA). All kernels atomicAdd into violation_log.violation_write_index and stamp their
            CanaryLaunchTag into each violation row.
        kernel_run_counters: Per-CanaryLaunchTag int64 counter array, shape [num_tags], device. The
            kernel itself does NOT index this array; runner takes a 1-element view at tag's slot (via
            CanaryEndpoint.kernel_run_counter_view) and hands a shape [1] tensor to the kernel,
            which atomicAdds 1 regardless of whether the plan had any active entry. Health watchdog
            reads this array to confirm "canary path actually ran".
        slot_run_counters: Per-CanaryLaunchTag int64 counter array, shape [num_tags], device. Same
            view-handed-to-kernel pattern as kernel_run_counters; each launch adds its active entry
            count to its slot. Used for periodic stats ("protected N tokens").
        enable_chain_position_assert: int32 [1] device flag gating the write kernel's chain-step
            write_position assert. allocate() defaults to 1; CanaryManager zeros it during
            __init__ for the warmup window and mark_init_finished() flips it back to 1.
        req_to_verify_expected_tokens: Optional int32 device tensor shape
            ``[req_to_token_alloc_size, max_context_len]``. Mirrors ReqToTokenPool layout;
            ``pool[req_idx, p]`` = source-of-truth token at logical position ``p`` for the
            req in slot ``req_idx``. Allocated only when
            ``CanaryConfig.enable_verify_token_assert`` is True. The plan-side entries
            kernel gathers from this pool (via ``kv_token_id_vs_position_offset`` per buffer
            group) into ``VerifyPlan.verify_expected_tokens``; the verify kernel then
            compares against each canary slot's stored token.
    """

    violation_log: ViolationLog
    kernel_run_counters: torch.Tensor
    slot_run_counters: torch.Tensor
    enable_chain_position_assert: torch.Tensor
    req_to_verify_expected_tokens: Optional[torch.Tensor]

    @classmethod
    def allocate(
        cls,
        *,
        config: CanaryConfig,
        device: torch.device,
        num_tags: int,
        req_to_token_alloc_size: Optional[int] = None,
        max_context_len: Optional[int] = None,
    ) -> CanaryDeviceState:
        if num_tags <= 0:
            raise ValueError(
                f"kv-canary: CanaryDeviceState num_tags must be positive, got {num_tags}"
            )
        violation_log = ViolationLog.allocate(
            ring_capacity=config.ring_capacity, device=device
        )
        kernel_run_counters = torch.zeros(num_tags, dtype=torch.int64, device=device)
        slot_run_counters = torch.zeros(num_tags, dtype=torch.int64, device=device)
        enable_chain_position_assert = torch.ones(1, dtype=torch.int32, device=device)
        if config.enable_verify_token_assert:
            if req_to_token_alloc_size is None or max_context_len is None:
                raise ValueError(
                    "kv-canary: CanaryDeviceState.allocate requires req_to_token_alloc_size "
                    "and max_context_len when CanaryConfig.enable_verify_token_assert is on"
                )
            req_to_verify_expected_tokens = torch.empty(
                (req_to_token_alloc_size, max_context_len),
                dtype=torch.int32,
                device=device,
            )
        else:
            req_to_verify_expected_tokens = None
        return cls(
            violation_log=violation_log,
            kernel_run_counters=kernel_run_counters,
            slot_run_counters=slot_run_counters,
            enable_chain_position_assert=enable_chain_position_assert,
            req_to_verify_expected_tokens=req_to_verify_expected_tokens,
        )