chore: import upstream snapshot with attribution

python/sglang/jit_kernel/csrc/kv_canary/canary_common.cuh

@@ -0,0 +1,136 @@
+#pragma once
+
+#include <sgl_kernel/utils.cuh>  // For SGL_DEVICE
+
+#include "consts.cuh"
+#include <cstdint>
+
+namespace canary {
+
+// Device-side handle for one real-KV source.
+struct RealKvSourceHandle {
+  const uint8_t* tensor;     // raw uint8 byte pointer to the source tensor
+  int32_t row_stride_bytes;  // tensor.shape[1] in bytes (may exceed page_size * num_bytes_per_token)
+  int32_t page_size;
+  int32_t num_bytes_per_token;
+  int32_t read_bytes;
+};
+
+SGL_DEVICE uint64_t splitmix64(uint64_t x) {
+  x = (x ^ (x >> 30)) * 0xBF58476D1CE4E5B9ULL;
+  x = (x ^ (x >> 27)) * 0x94D049BB133111EBULL;
+  return x ^ (x >> 31);
+}
+
+SGL_DEVICE uint64_t splitmix64_mix3(uint64_t a, uint64_t b, uint64_t c) {
+  uint64_t h = splitmix64(a);
+  h = splitmix64(h ^ b);
+  h = splitmix64(h ^ c);
+  return h;
+}
+
+// Read 16 aligned bytes from a source as two uint64 little-endian words, following the RealKvSource access
+// invariant. The invariant (from kv_canary/verify.py docstring) is:
+//
+//     tensor[slot_idx // page_size,
+//            (slot_idx % page_size) * num_bytes_per_token + byte_offset]
+//
+// row_stride_bytes is the dim-1 size of the underlying tensor in bytes (which may exceed
+// page_size * num_bytes_per_token; trailing bytes are skipped).
+SGL_DEVICE void real_kv_load_uint4(
+    const RealKvSourceHandle& src, int64_t slot_idx, int64_t byte_offset, uint64_t& word_lo, uint64_t& word_hi) {
+  const int64_t row = slot_idx / src.page_size;
+  const int64_t col_within_page = slot_idx % src.page_size;
+  const int64_t col = col_within_page * src.num_bytes_per_token + byte_offset;
+  const int64_t flat_index = row * static_cast<int64_t>(src.row_stride_bytes) + col;
+  const uint4 vec = *reinterpret_cast<const uint4*>(src.tensor + flat_index);
+  word_lo = static_cast<uint64_t>(vec.x) | (static_cast<uint64_t>(vec.y) << 32);
+  word_hi = static_cast<uint64_t>(vec.z) | (static_cast<uint64_t>(vec.w) << 32);
+}
+
+SGL_DEVICE uint64_t real_kv_fold_one_source(const RealKvSourceHandle& src, int64_t slot_idx, RealKvHashMode mode) {
+  const int64_t effective_read_bytes = (mode == RealKvHashMode::kPartial) ? static_cast<int64_t>(16) : src.read_bytes;
+  uint64_t acc = 0ULL;
+  for (int64_t byte_offset = 0; byte_offset < effective_read_bytes; byte_offset += 16) {
+    uint64_t word_lo;
+    uint64_t word_hi;
+    real_kv_load_uint4(src, slot_idx, byte_offset, word_lo, word_hi);
+    acc = splitmix64(acc ^ word_lo);
+    acc = splitmix64(acc ^ word_hi);
+  }
+  return acc;
+}
+
+// Fold all configured real-KV sources for a given slot. Iterates sequentially and combines each source's
+// contribution via acc = splitmix64(acc XOR source_hash); matches _compute_real_kv_hash_scalar in
+// kv_canary/verify_ref.py. In OFF mode the function returns 0 unconditionally (the running real_kv_hash
+// field is always 0).
+SGL_DEVICE uint64_t
+real_kv_fold_sources(const RealKvSourceHandle* sources, int num_sources, int64_t slot_idx, RealKvHashMode mode) {
+  if (mode == RealKvHashMode::kNone || num_sources <= 0) {
+    return 0ULL;
+  }
+  uint64_t acc = 0ULL;
+  for (int s = 0; s < num_sources; ++s) {
+    const uint64_t source_hash = real_kv_fold_one_source(sources[s], slot_idx, mode);
+    acc = splitmix64(acc ^ source_hash);
+  }
+  return acc;
+}
+
+struct ViolationSink {
+  int64_t* __restrict__ ring;
+  int32_t* __restrict__ write_index;
+  int32_t ring_capacity;
+  int32_t kernel_kind;
+};
+
+struct ViolationRow {
+  int64_t slot_idx;
+  int64_t position;
+  int64_t stored_token;
+  int64_t expected_token;
+  int64_t stored_chain_hash;
+  int64_t expected_aux;
+  int64_t fail_reason_bits;
+};
+
+SGL_DEVICE void record_violation(const ViolationSink& sink, const ViolationRow& row) {
+  const int32_t seq = atomicAdd(sink.write_index, 1);
+  if (seq < sink.ring_capacity) {
+    int64_t* dst = sink.ring + static_cast<int64_t>(seq) * kViolationFields;
+    dst[kViolationFieldKernelKind] = static_cast<int64_t>(sink.kernel_kind);
+    dst[kViolationFieldSlotIdx] = row.slot_idx;
+    dst[kViolationFieldPosition] = row.position;
+    dst[kViolationFieldStoredToken] = row.stored_token;
+    dst[kViolationFieldExpectedToken] = row.expected_token;
+    dst[kViolationFieldStoredChainHash] = row.stored_chain_hash;
+    dst[kViolationFieldExpectedAux] = row.expected_aux;
+    dst[kViolationFieldFailReasonBits] = row.fail_reason_bits;
+    __threadfence_system();
+  }
+}
+
+SGL_DEVICE int64_t canary_load_field(const uint8_t* buf, int64_t slot_idx, int64_t slot_stride_bytes, int field) {
+  const int64_t* p = reinterpret_cast<const int64_t*>(buf + slot_idx * slot_stride_bytes);
+  return p[field];
+}
+
+SGL_DEVICE void
+canary_store_field(uint8_t* buf, int64_t slot_idx, int64_t slot_stride_bytes, int field, int64_t value) {
+  int64_t* p = reinterpret_cast<int64_t*>(buf + slot_idx * slot_stride_bytes);
+  p[field] = value;
+}
+
+SGL_DEVICE uint64_t compute_slot_hash(const uint8_t* canary_buf, int64_t slot_stride_bytes, int64_t source_slot_idx) {
+  if (source_slot_idx < 0) {
+    return splitmix64(kCanaryChainAnchor);
+  }
+  const int64_t token = canary_load_field(canary_buf, source_slot_idx, slot_stride_bytes, kCanaryFieldToken);
+  const int64_t position = canary_load_field(canary_buf, source_slot_idx, slot_stride_bytes, kCanaryFieldPosition);
+  const int64_t prev_hash = canary_load_field(canary_buf, source_slot_idx, slot_stride_bytes, kCanaryFieldPrevHash);
+  return splitmix64_mix3(
+      static_cast<uint64_t>(prev_hash), static_cast<uint64_t>(token), static_cast<uint64_t>(position));
+}
+
+}  // namespace canary

python/sglang/jit_kernel/csrc/kv_canary/canary_plan_entries.cuh

@@ -0,0 +1,311 @@
+#pragma once
+
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/utils.h>   // For RuntimeCheck
+
+#include <sgl_kernel/runtime.cuh>  // For host::runtime::get_sm_count
+#include <sgl_kernel/utils.cuh>    // For LaunchKernel, SGL_DEVICE
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+
+namespace {
+
+struct PlanEntriesParams {
+  // Inputs.
+  const int64_t* __restrict__ req_pool_indices;               // [bs_padded] int64
+  const int64_t* __restrict__ prefix_lens;                    // [bs_padded] int64
+  const int32_t* __restrict__ req_to_token;                   // [max_reqs, max_seq_len] int32
+  const int64_t* __restrict__ full_to_swa_lut;                // [lut_len] int64, may be nullptr when !HAS_SWA_LUT
+  const int64_t* __restrict__ verify_offsets_scratch;         // [bs_padded + 1] int64 (cumulative prefix sum)
+  const int32_t* __restrict__ verify_enable;                  // [1] int32 — 0 ⇒ skip scatter entirely
+  const int32_t* __restrict__ req_to_verify_expected_tokens;  // [max_reqs, max_context_len] int32, may be nullptr
+  const int64_t* __restrict__ req_to_verify_expected_tokens_valid_lens;  // [bs_padded] int64 per-req snapshot length;
+                                                                         // nullptr iff pool is null
+  // Outputs.
+  int64_t* __restrict__ out_verify_slot_indices;        // [verify_capacity] int64
+  int64_t* __restrict__ out_verify_expected_tokens;     // [verify_capacity] int64
+  int64_t* __restrict__ out_verify_expected_positions;  // [verify_capacity] int64
+  int64_t* __restrict__ out_verify_prev_slot_indices;   // [verify_capacity] int64
+  // Sizes / strides.
+  int32_t bs_padded;
+  int64_t verify_capacity;  // out_verify_*[verify_capacity]; scatter is clamped to this length.
+  int64_t req_to_token_stride0;
+  int64_t req_to_verify_expected_tokens_stride0;
+  int32_t kv_token_id_vs_position_offset;  // 0 for target pools; +1 for EAGLE draft.
+  int32_t swa_window_size;
+};
+
+// Binary search for the largest req_id such that verify_offsets[req_id] <= tid. Pre-condition: tid is
+// strictly less than verify_offsets[bs_padded] = total_verify; bs_padded >= 1; verify_offsets[0] = 0.
+SGL_DEVICE int32_t find_req_id(const int64_t* __restrict__ verify_offsets, int32_t bs_padded, int64_t tid) {
+  int32_t lo = 0;
+  int32_t hi = bs_padded;  // exclusive upper bound; verify_offsets[hi] > tid
+  while (hi - lo > 1) {
+    const int32_t mid = (lo + hi) >> 1;
+    if (verify_offsets[mid] <= tid) {
+      lo = mid;
+    } else {
+      hi = mid;
+    }
+  }
+  return lo;
+}
+
+// Translate raw slot value via the SWA LUT. Sentinel passthrough (-1 stays -1). Clamp slot to
+// ``lut_len - 1`` defensively; in practice the caller never produces out-of-range slots.
+SGL_DEVICE int64_t swa_translate(const int64_t* __restrict__ lut, int64_t lut_len, int64_t raw_slot) {
+  if (raw_slot < 0) {
+    return raw_slot;
+  }
+  int64_t safe = raw_slot;
+  if (lut_len > 0 && safe >= lut_len) {
+    safe = lut_len - 1;
+  }
+  return lut[safe];
+}
+
+// Persistent grid; one thread = one verify entry (with stride). Template parameter HAS_SWA_LUT switches
+// the SWA-translate path off entirely in the FULL pool variant.
+// HAS_VERIFY_EXPECTED_TOKEN_POOL toggles the source-of-truth token gather; when off, every active entry
+// writes the ``-1`` sentinel so the verify kernel skips the token-mismatch check.
+template <bool HAS_SWA_LUT, bool HAS_VERIFY_EXPECTED_TOKEN_POOL>
+__global__ void plan_entries_persistent_kernel(
+    const PlanEntriesParams __grid_constant__ params,
+    int64_t lut_len  // only meaningful when HAS_SWA_LUT
+) {
+  const int64_t total_verify = params.verify_offsets_scratch[params.bs_padded];
+  if (total_verify <= 0) {
+    return;
+  }
+
+  if (*params.verify_enable == 0) {
+    return;
+  }
+
+  // Contract: offsets kernel clears verify_enable iff total_verify > verify_capacity, so reaching
+  // here implies total_verify <= verify_capacity. Trap on contract break (cassert is NDEBUG-gated).
+  if (total_verify > params.verify_capacity) {
+    if (blockIdx.x == 0 && threadIdx.x == 0) {
+      printf(
+          "kv-canary plan_entries: total_verify=%lld exceeds verify_capacity=%lld with "
+          "verify_enable=1 (offsets/entries contract broken)\n",
+          static_cast<long long>(total_verify),
+          static_cast<long long>(params.verify_capacity));
+    }
+#ifndef USE_ROCM
+    __trap();
+#else
+    __builtin_trap();  // HIP/clang device-side trap; __trap() is CUDA-only.
+#endif
+  }
+
+  const int64_t tid_start = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  const int64_t stride = static_cast<int64_t>(gridDim.x) * blockDim.x;
+
+  const int32_t swa_window = params.swa_window_size;
+  const int64_t req_to_token_stride0 = params.req_to_token_stride0;
+
+  for (int64_t tid = tid_start; tid < total_verify; tid += stride) {
+    // 1) Find the owning req via binary search over verify_offsets.
+    const int32_t req_id = find_req_id(params.verify_offsets_scratch, params.bs_padded, tid);
+    const int64_t req_start = params.verify_offsets_scratch[req_id];
+    const int64_t entry_idx = tid - req_start;
+
+    // 2) Load per-req metadata. Padding rows have verify_count=0 in the offsets prefix-sum, so a tid in
+    // the live range can never land on a padding req; no explicit padding-row check needed.
+    const int64_t rp = params.req_pool_indices[req_id];
+    const int64_t prefix_len = params.prefix_lens[req_id];
+    const int64_t window_start = (swa_window > 0) ? (prefix_len - swa_window > 0 ? prefix_len - swa_window : 0) : 0;
+    const int64_t out_position = window_start + entry_idx;
+
+    // 3) Gather slot + prev_slot via req_to_token.
+    const int64_t row_base = rp * req_to_token_stride0;
+    const int32_t slot_raw = params.req_to_token[row_base + out_position];
+    int64_t out_slot;
+    if constexpr (HAS_SWA_LUT) {
+      out_slot = swa_translate(params.full_to_swa_lut, lut_len, static_cast<int64_t>(slot_raw));
+    } else {
+      out_slot = static_cast<int64_t>(slot_raw);
+    }
+
+    int64_t out_prev_slot;
+    if (out_position > 0) {
+      const int32_t prev_raw = params.req_to_token[row_base + out_position - 1];
+      if constexpr (HAS_SWA_LUT) {
+        out_prev_slot = swa_translate(params.full_to_swa_lut, lut_len, static_cast<int64_t>(prev_raw));
+      } else {
+        out_prev_slot = static_cast<int64_t>(prev_raw);
+      }
+    } else {
+      out_prev_slot = -1;
+    }
+
+    int64_t out_expected_input_id = -1;
+    if constexpr (HAS_VERIFY_EXPECTED_TOKEN_POOL) {
+      const int64_t sot_pos = out_position + static_cast<int64_t>(params.kv_token_id_vs_position_offset);
+      const int64_t valid_len = params.req_to_verify_expected_tokens_valid_lens[req_id];
+      if (sot_pos >= 0 && sot_pos < valid_len) {
+        const int32_t token =
+            params.req_to_verify_expected_tokens[rp * params.req_to_verify_expected_tokens_stride0 + sot_pos];
+        out_expected_input_id = static_cast<int64_t>(token);
+      }
+    }
+
+    // 4) Scatter. out_idx == tid since verify_offsets[req_id] + entry_idx == tid by construction.
+    params.out_verify_slot_indices[tid] = out_slot;
+    params.out_verify_expected_tokens[tid] = out_expected_input_id;
+    params.out_verify_expected_positions[tid] = out_position;
+    params.out_verify_prev_slot_indices[tid] = out_prev_slot;
+  }
+}
+
+template <bool HAS_SWA_LUT, bool HAS_VERIFY_EXPECTED_TOKEN_POOL>
+struct PlanEntriesKernel {
+  static constexpr int kBlockSize = 128;
+  static constexpr int kBlocksPerSm = 8;
+
+  static auto get_num_sms(DLDevice device) {
+    static const auto kNumSM = host::runtime::get_sm_count(device.device_id);
+    return kNumSM;
+  }
+
+  static void
+  run(const tvm::ffi::TensorView req_pool_indices,
+      const tvm::ffi::TensorView prefix_lens,
+      const tvm::ffi::TensorView req_to_token,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> full_to_swa_index_mapping,
+      const tvm::ffi::TensorView verify_offsets_scratch,
+      const tvm::ffi::TensorView verify_enable,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> req_to_verify_expected_tokens,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> req_to_verify_expected_tokens_valid_lens,
+      const tvm::ffi::TensorView out_verify_slot_indices,
+      const tvm::ffi::TensorView out_verify_expected_tokens,
+      const tvm::ffi::TensorView out_verify_expected_positions,
+      const tvm::ffi::TensorView out_verify_prev_slot_indices,
+      int32_t kv_token_id_vs_position_offset,
+      int32_t swa_window_size) {
+    using namespace host;
+
+    SymbolicSize Nbs = {"bs_padded"};
+    SymbolicSize Nscratch = {"verify_offsets_scratch_len"};
+    SymbolicSize Ncap = {"verify_capacity"};
+    SymbolicSize Nmax_reqs = {"max_reqs"};
+    SymbolicSize Nmax_seq_len = {"max_seq_len"};
+    SymbolicSize Nlut = {"lut_len"};
+    SymbolicSize Npool_rows = {"req_to_verify_expected_tokens_rows"};
+    SymbolicSize Npool_cols = {"req_to_verify_expected_tokens_cols"};
+    SymbolicDevice device_;
+    device_.set_options<kDLGPU>();
+
+    TensorMatcher({Nbs})  //
+        .with_dtype<int64_t>()
+        .with_device<kDLGPU>(device_)
+        .verify(req_pool_indices)
+        .verify(prefix_lens);
+    TensorMatcher({Nscratch})  //
+        .with_dtype<int64_t>()
+        .with_device<kDLGPU>(device_)
+        .verify(verify_offsets_scratch);
+    TensorMatcher({1})  //
+        .with_dtype<int32_t>()
+        .with_device<kDLGPU>(device_)
+        .verify(verify_enable);
+    TensorMatcher({Ncap})  //
+        .with_dtype<int64_t>()
+        .with_device<kDLGPU>(device_)
+        .verify(out_verify_slot_indices)
+        .verify(out_verify_expected_tokens)
+        .verify(out_verify_expected_positions)
+        .verify(out_verify_prev_slot_indices);
+    TensorMatcher({Nmax_reqs, Nmax_seq_len})  //
+        .with_dtype<int32_t>()
+        .with_device<kDLGPU>(device_)
+        .verify(req_to_token);
+    RuntimeCheck(
+        full_to_swa_index_mapping.has_value() == HAS_SWA_LUT,
+        "full_to_swa_index_mapping presence does not match HAS_SWA_LUT specialization");
+    RuntimeCheck(
+        req_to_verify_expected_tokens.has_value() == HAS_VERIFY_EXPECTED_TOKEN_POOL,
+        "req_to_verify_expected_tokens presence does not match HAS_VERIFY_EXPECTED_TOKEN_POOL specialization");
+    if constexpr (HAS_VERIFY_EXPECTED_TOKEN_POOL) {
+      RuntimeCheck(
+          req_to_verify_expected_tokens_valid_lens.has_value(),
+          "req_to_verify_expected_tokens_valid_lens must be set when req_to_verify_expected_tokens is set");
+    }
+    if constexpr (HAS_SWA_LUT) {
+      TensorMatcher({Nlut})  //
+          .with_dtype<int64_t>()
+          .with_device<kDLGPU>(device_)
+          .verify(full_to_swa_index_mapping.value());
+    }
+    if constexpr (HAS_VERIFY_EXPECTED_TOKEN_POOL) {
+      TensorMatcher({Npool_rows, Npool_cols})  //
+          .with_dtype<int32_t>()
+          .with_device<kDLGPU>(device_)
+          .verify(req_to_verify_expected_tokens.value());
+      TensorMatcher({Nbs})  //
+          .with_dtype<int64_t>()
+          .with_device<kDLGPU>(device_)
+          .verify(req_to_verify_expected_tokens_valid_lens.value());
+    }
+    RuntimeCheck(Nscratch.unwrap() >= Nbs.unwrap() + 1, "verify_offsets_scratch length must be >= bs_padded + 1");
+
+    const int64_t bs_padded = Nbs.unwrap();
+    if (bs_padded <= 0) {
+      return;
+    }
+
+    const int64_t* lut_ptr = nullptr;
+    int64_t lut_len = 0;
+    if constexpr (HAS_SWA_LUT) {
+      lut_ptr = static_cast<const int64_t*>(full_to_swa_index_mapping.value().data_ptr());
+      lut_len = static_cast<int64_t>(Nlut.unwrap());
+    }
+
+    const int32_t* expected_token_ids_ptr = nullptr;
+    int64_t expected_token_ids_stride0 = 0;
+    const int64_t* req_to_verify_expected_tokens_valid_lens_ptr = nullptr;
+    if constexpr (HAS_VERIFY_EXPECTED_TOKEN_POOL) {
+      expected_token_ids_ptr = static_cast<const int32_t*>(req_to_verify_expected_tokens.value().data_ptr());
+      expected_token_ids_stride0 = static_cast<int64_t>(Npool_cols.unwrap());
+      req_to_verify_expected_tokens_valid_lens_ptr =
+          static_cast<const int64_t*>(req_to_verify_expected_tokens_valid_lens.value().data_ptr());
+    }
+
+    const PlanEntriesParams params = PlanEntriesParams{
+        .req_pool_indices = static_cast<const int64_t*>(req_pool_indices.data_ptr()),
+        .prefix_lens = static_cast<const int64_t*>(prefix_lens.data_ptr()),
+        .req_to_token = static_cast<const int32_t*>(req_to_token.data_ptr()),
+        .full_to_swa_lut = lut_ptr,
+        .verify_offsets_scratch = static_cast<const int64_t*>(verify_offsets_scratch.data_ptr()),
+        .verify_enable = static_cast<const int32_t*>(verify_enable.data_ptr()),
+        .req_to_verify_expected_tokens = expected_token_ids_ptr,
+        .req_to_verify_expected_tokens_valid_lens = req_to_verify_expected_tokens_valid_lens_ptr,
+        .out_verify_slot_indices = static_cast<int64_t*>(out_verify_slot_indices.data_ptr()),
+        .out_verify_expected_tokens = static_cast<int64_t*>(out_verify_expected_tokens.data_ptr()),
+        .out_verify_expected_positions = static_cast<int64_t*>(out_verify_expected_positions.data_ptr()),
+        .out_verify_prev_slot_indices = static_cast<int64_t*>(out_verify_prev_slot_indices.data_ptr()),
+        .bs_padded = static_cast<int32_t>(bs_padded),
+        .verify_capacity = static_cast<int64_t>(Ncap.unwrap()),
+        .req_to_token_stride0 = static_cast<int64_t>(Nmax_seq_len.unwrap()),
+        .req_to_verify_expected_tokens_stride0 = expected_token_ids_stride0,
+        .kv_token_id_vs_position_offset = kv_token_id_vs_position_offset,
+        .swa_window_size = swa_window_size,
+    };
+
+    const DLDevice device = device_.unwrap();
+    const int num_sms = get_num_sms(device);
+    const int num_blocks = num_sms * kBlocksPerSm;
+
+    const dim3 grid(num_blocks);
+    const dim3 block(kBlockSize);
+
+    LaunchKernel(grid, block, device)(
+        plan_entries_persistent_kernel<HAS_SWA_LUT, HAS_VERIFY_EXPECTED_TOKEN_POOL>, params, lut_len);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/kv_canary/canary_verify.cuh

@@ -0,0 +1,289 @@
+#pragma once
+
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/utils.h>   // For div_ceil, RuntimeCheck
+
+#include <sgl_kernel/utils.cuh>  // For LaunchKernel, SGL_DEVICE
+#include <sgl_kernel/warp.cuh>   // For device::warp::reduce_sum
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include "canary_common.cuh"
+#include <cstdint>
+
+namespace canary {
+
+namespace {
+
+constexpr uint32_t kVerifyBlockSize = 512;
+constexpr uint32_t kPersistentBlocks = 64;
+
+struct VerifyKernelParams {
+  // Canary buffer this launch verifies. Read-only.
+  const uint8_t* canary_buf;
+  int64_t slot_stride_bytes;
+
+  // Plan tensors.
+  const int64_t* verify_slot_indices;
+  const int64_t* verify_expected_tokens;
+  const int64_t* verify_expected_positions;
+  const int64_t* verify_prev_slot_indices;
+  const int32_t* verify_num_valid;
+  const int32_t* verify_enable;
+  int32_t verify_capacity;
+
+  // Violation sink (ring + write_index + capacity + kernel_kind bundled in canary_common.cuh).
+  ViolationSink violation_sink;
+
+  // Health counters.
+  int64_t* slot_run_counter;
+  int64_t* kernel_run_counter;
+
+  // Real-KV sources (fixed-size ABI; padding slots have read_bytes = 0).
+  RealKvSourceHandle sources[kMaxRealKvSources];
+  int32_t num_sources;
+  RealKvHashMode real_kv_hash_mode;
+};
+
+template <bool CHECK_VERIFY_EXPECTED_TOKEN>
+__global__ void canary_verify_kernel(const VerifyKernelParams __grid_constant__ p) {
+  const uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const uint32_t stride = gridDim.x * blockDim.x;
+
+  if (tid == 0) {
+    atomicAdd(reinterpret_cast<unsigned long long*>(p.kernel_run_counter), 1ULL);
+  }
+
+  if (*p.verify_enable == 0) {
+    return;
+  }
+
+  const int32_t active = min(*p.verify_num_valid, p.verify_capacity);
+
+  uint32_t local_active_count = 0;
+  for (uint32_t entry_idx = tid; entry_idx < static_cast<uint32_t>(active); entry_idx += stride) {
+    ++local_active_count;
+
+    const int64_t slot_idx = p.verify_slot_indices[entry_idx];
+    const int64_t expected_position = p.verify_expected_positions[entry_idx];
+    const int64_t prev_slot_idx = p.verify_prev_slot_indices[entry_idx];
+    int64_t expected_input_id = -1;
+    if constexpr (CHECK_VERIFY_EXPECTED_TOKEN) {
+      expected_input_id = p.verify_expected_tokens[entry_idx];
+    }
+
+    if (slot_idx == kTokenToKvSlotPadding) {
+      continue;
+    }
+
+    const int64_t stored_token = canary_load_field(p.canary_buf, slot_idx, p.slot_stride_bytes, kCanaryFieldToken);
+    const int64_t stored_position =
+        canary_load_field(p.canary_buf, slot_idx, p.slot_stride_bytes, kCanaryFieldPosition);
+    const int64_t stored_chain_hash =
+        canary_load_field(p.canary_buf, slot_idx, p.slot_stride_bytes, kCanaryFieldPrevHash);
+    const int64_t stored_real_kv_hash =
+        canary_load_field(p.canary_buf, slot_idx, p.slot_stride_bytes, kCanaryFieldRealKvHash);
+
+    const bool prev_reachable = (prev_slot_idx != kTokenToKvSlotPadding);
+    const int64_t expected_chain_hash =
+        prev_reachable ? static_cast<int64_t>(compute_slot_hash(p.canary_buf, p.slot_stride_bytes, prev_slot_idx))
+                       : stored_chain_hash;
+
+    const uint64_t expected_real_kv_hash_u64 =
+        real_kv_fold_sources(p.sources, p.num_sources, slot_idx, p.real_kv_hash_mode);
+    const int64_t expected_real_kv_hash = static_cast<int64_t>(expected_real_kv_hash_u64);
+
+    FailReason fail_reason_bits{};
+    if (prev_reachable && stored_chain_hash != expected_chain_hash) {
+      fail_reason_bits |= FailReason::kVerifyChainHashMismatch;
+    }
+    if constexpr (CHECK_VERIFY_EXPECTED_TOKEN) {
+      if (expected_input_id != -1 && stored_token != expected_input_id) {
+        fail_reason_bits |= FailReason::kVerifyTokenMismatch;
+      }
+    }
+    if (stored_position != expected_position) {
+      fail_reason_bits |= FailReason::kVerifyPositionMismatch;
+    }
+    if (stored_real_kv_hash != expected_real_kv_hash) {
+      fail_reason_bits |= FailReason::kVerifyRealKvHashMismatch;
+    }
+
+    if (fail_reason_bits != FailReason{}) {
+      record_violation(
+          p.violation_sink,
+          ViolationRow{
+              /* slot_idx = */ slot_idx,
+              /* position = */ stored_position,
+              /* stored_token = */ stored_token,
+              /* expected_token = */ expected_input_id,
+              /* stored_chain_hash = */ stored_chain_hash,
+              /* expected_aux = */ expected_chain_hash,
+              /* fail_reason_bits = */ static_cast<int64_t>(fail_reason_bits),
+          });
+    }
+  }
+
+  // Warp-size-agnostic reduction: device::warp::reduce_sum performs a width-kWarpThreads (32) xor
+  // all-reduce, so each 32-lane sub-group gets its own total with its own lane-0 leader. Correct on
+  // both CUDA (warp=32) and ROCm wave64 (two 32-lane sub-groups per wavefront); a hand-rolled
+  // __shfl_down over the full wavefront would cross the 32-lane boundary on wave64.
+  const uint32_t warp_active_count = device::warp::reduce_sum<device::kWarpThreads>(local_active_count);
+  const bool is_subgroup_leader = (threadIdx.x % device::kWarpThreads) == 0u;
+  if (is_subgroup_leader && warp_active_count != 0u) {
+    atomicAdd(
+        reinterpret_cast<unsigned long long*>(p.slot_run_counter), static_cast<unsigned long long>(warp_active_count));
+  }
+}
+
+}  // namespace
+
+// API source of truth: docstring of canary_verify_step in python/sglang/jit_kernel/kv_canary/verify.py.
+//
+// ABI notes:
+// - real_kv_buf_0 .. real_kv_buf_3 are 4 fixed uint8 tensor slots. Unused slots are dummy 1-byte tensors;
+//   the host wrapper sets real_kv_source_params[s, 2] (read_bytes) = 0 for them.
+// - real_kv_source_params is a CPU int32 tensor of shape [kMaxRealKvSources, 3]: per-source (page_size,
+//   num_bytes_per_token, read_bytes). Lives on CPU because the host wrapper materializes it from a tuple
+//   of Python dataclasses.
+// - num_sources is passed as int64 (tvm-ffi scalar convention) but always in [0, kMaxRealKvSources].
+template <bool CHECK_VERIFY_EXPECTED_TOKEN>
+struct CanaryVerifyKernel {
+  static void
+  run(tvm::ffi::TensorView canary_buf,
+      tvm::ffi::TensorView verify_slot_indices,
+      tvm::ffi::TensorView verify_expected_tokens,
+      tvm::ffi::TensorView verify_expected_positions,
+      tvm::ffi::TensorView verify_prev_slot_indices,
+      tvm::ffi::TensorView verify_num_valid,
+      tvm::ffi::TensorView verify_enable,
+      int64_t kernel_kind,
+      tvm::ffi::TensorView violation_ring,
+      tvm::ffi::TensorView violation_write_index,
+      tvm::ffi::TensorView slot_run_counter,
+      tvm::ffi::TensorView kernel_run_counter,
+      tvm::ffi::TensorView real_kv_buf_0,
+      tvm::ffi::TensorView real_kv_buf_1,
+      tvm::ffi::TensorView real_kv_buf_2,
+      tvm::ffi::TensorView real_kv_buf_3,
+      tvm::ffi::TensorView real_kv_source_params,
+      int64_t num_sources,
+      int64_t real_kv_hash_mode) {
+    using namespace host;
+
+    SymbolicSize N_slots = {"num_canary_slots"};
+    SymbolicSize N_stride = {"slot_stride_bytes"};
+    SymbolicSize N_verify = {"verify_capacity"};
+    SymbolicDevice device_;
+    device_.set_options<kDLGPU>();
+
+    TensorMatcher({N_slots, N_stride}).with_dtype<uint8_t>().with_device<kDLGPU>(device_).verify(canary_buf);
+
+    TensorMatcher({N_verify})
+        .with_dtype<int64_t>()
+        .with_device<kDLGPU>(device_)
+        .verify(verify_slot_indices)
+        .verify(verify_expected_tokens)
+        .verify(verify_expected_positions)
+        .verify(verify_prev_slot_indices);
+    TensorMatcher({1}).with_dtype<int32_t>().with_device<kDLGPU>(device_).verify(verify_num_valid);
+    TensorMatcher({1}).with_dtype<int32_t>().with_device<kDLGPU>(device_).verify(verify_enable);
+
+    TensorMatcher({1}).with_dtype<int32_t>().with_device<kDLGPU>(device_).verify(violation_write_index);
+    SymbolicSize N_ring = {"ring_capacity"};
+    TensorMatcher({N_ring, static_cast<int64_t>(kViolationFields)})
+        .with_dtype<int64_t>()
+        .with_device<kDLGPU>(device_)
+        .verify(violation_ring);
+    TensorMatcher({1})
+        .with_dtype<int64_t>()
+        .with_device<kDLGPU>(device_)
+        .verify(slot_run_counter)
+        .verify(kernel_run_counter);
+
+    // Real-KV source buffers are 2-D uint8 (any shape); dim 1 carries the row-stride in bytes. They live on
+    // CUDA. real_kv_source_params is a small CPU int32 table of length kMaxRealKvSources * 3.
+    SymbolicSize N_real_kv_rows_0 = {"real_kv_rows_0"};
+    SymbolicSize N_real_kv_cols_0 = {"real_kv_cols_0"};
+    TensorMatcher({N_real_kv_rows_0, N_real_kv_cols_0})
+        .with_dtype<uint8_t>()
+        .with_device<kDLGPU>(device_)
+        .verify(real_kv_buf_0);
+    SymbolicSize N_real_kv_rows_1 = {"real_kv_rows_1"};
+    SymbolicSize N_real_kv_cols_1 = {"real_kv_cols_1"};
+    TensorMatcher({N_real_kv_rows_1, N_real_kv_cols_1})
+        .with_dtype<uint8_t>()
+        .with_device<kDLGPU>(device_)
+        .verify(real_kv_buf_1);
+    SymbolicSize N_real_kv_rows_2 = {"real_kv_rows_2"};
+    SymbolicSize N_real_kv_cols_2 = {"real_kv_cols_2"};
+    TensorMatcher({N_real_kv_rows_2, N_real_kv_cols_2})
+        .with_dtype<uint8_t>()
+        .with_device<kDLGPU>(device_)
+        .verify(real_kv_buf_2);
+    SymbolicSize N_real_kv_rows_3 = {"real_kv_rows_3"};
+    SymbolicSize N_real_kv_cols_3 = {"real_kv_cols_3"};
+    TensorMatcher({N_real_kv_rows_3, N_real_kv_cols_3})
+        .with_dtype<uint8_t>()
+        .with_device<kDLGPU>(device_)
+        .verify(real_kv_buf_3);
+    TensorMatcher({static_cast<int64_t>(kMaxRealKvSources), static_cast<int64_t>(kRealKvSourceFieldsPerEntry)})
+        .with_dtype<int32_t>()
+        .with_device<kDLCPU>()
+        .verify(real_kv_source_params);
+
+    const int64_t slot_stride_bytes = N_stride.unwrap();
+    const int32_t verify_capacity = static_cast<int32_t>(N_verify.unwrap());
+    const int32_t ring_capacity = static_cast<int32_t>(N_ring.unwrap());
+    const DLDevice device = device_.unwrap();
+
+    RuntimeCheck(
+        slot_stride_bytes >= static_cast<int64_t>(kCanaryFieldsPerSlot * sizeof(int64_t)),
+        "canary_verify: slot_stride_bytes must hold at least ",
+        static_cast<int64_t>(kCanaryFieldsPerSlot * sizeof(int64_t)),
+        " bytes per slot, got ",
+        slot_stride_bytes);
+    RuntimeCheck(
+        num_sources >= 0 && num_sources <= static_cast<int64_t>(kMaxRealKvSources),
+        "canary_verify: num_sources must be in [0, ",
+        static_cast<int64_t>(kMaxRealKvSources),
+        "], got ",
+        num_sources);
+
+    VerifyKernelParams p{};
+    p.canary_buf = static_cast<const uint8_t*>(canary_buf.data_ptr());
+    p.slot_stride_bytes = slot_stride_bytes;
+    p.verify_slot_indices = static_cast<const int64_t*>(verify_slot_indices.data_ptr());
+    p.verify_expected_tokens = static_cast<const int64_t*>(verify_expected_tokens.data_ptr());
+    p.verify_expected_positions = static_cast<const int64_t*>(verify_expected_positions.data_ptr());
+    p.verify_prev_slot_indices = static_cast<const int64_t*>(verify_prev_slot_indices.data_ptr());
+    p.verify_num_valid = static_cast<const int32_t*>(verify_num_valid.data_ptr());
+    p.verify_enable = static_cast<const int32_t*>(verify_enable.data_ptr());
+    p.verify_capacity = verify_capacity;
+    p.violation_sink.ring = static_cast<int64_t*>(violation_ring.data_ptr());
+    p.violation_sink.write_index = static_cast<int32_t*>(violation_write_index.data_ptr());
+    p.violation_sink.ring_capacity = ring_capacity;
+    p.violation_sink.kernel_kind = static_cast<int32_t>(kernel_kind);
+    p.slot_run_counter = static_cast<int64_t*>(slot_run_counter.data_ptr());
+    p.kernel_run_counter = static_cast<int64_t*>(kernel_run_counter.data_ptr());
+
+    // Materialize the source handle array on the host. The CPU param tensor carries the per-source ints.
+    const int32_t* params = static_cast<const int32_t*>(real_kv_source_params.data_ptr());
+    tvm::ffi::TensorView source_bufs[kMaxRealKvSources] = {real_kv_buf_0, real_kv_buf_1, real_kv_buf_2, real_kv_buf_3};
+    for (int s = 0; s < kMaxRealKvSources; ++s) {
+      p.sources[s].tensor = static_cast<const uint8_t*>(source_bufs[s].data_ptr());
+      p.sources[s].row_stride_bytes = static_cast<int32_t>(source_bufs[s].size(1));
+      p.sources[s].page_size = params[s * kRealKvSourceFieldsPerEntry + kRealKvSourceFieldPageSize];
+      p.sources[s].num_bytes_per_token = params[s * kRealKvSourceFieldsPerEntry + kRealKvSourceFieldNumBytesPerToken];
+      p.sources[s].read_bytes = params[s * kRealKvSourceFieldsPerEntry + kRealKvSourceFieldReadBytes];
+    }
+    p.num_sources = static_cast<int32_t>(num_sources);
+    p.real_kv_hash_mode = static_cast<RealKvHashMode>(real_kv_hash_mode);
+
+    const uint32_t grid = kPersistentBlocks;
+    LaunchKernel(grid, kVerifyBlockSize, device)(canary_verify_kernel<CHECK_VERIFY_EXPECTED_TOKEN>, p);
+  }
+};
+
+}  // namespace canary

python/sglang/jit_kernel/csrc/kv_canary/canary_write.cuh

@@ -0,0 +1,349 @@
+#pragma once
+
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/utils.h>   // For RuntimeCheck
+
+#include <sgl_kernel/utils.cuh>  // For LaunchKernel, SGL_DEVICE
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include "canary_common.cuh"
+#include <cstdint>
+
+namespace canary {
+
+namespace {
+
+// Single thread per block — chain advance is inherently serial.
+constexpr uint32_t kWriteBlockSize = 1;
+
+struct WriteKernelParams {
+  uint8_t* canary_buf;
+  int64_t slot_stride_bytes;
+
+  // Plan tensors.
+  const int64_t* write_offsets;
+  const int64_t* write_seed_slot_indices;
+  const int32_t* write_num_valid_reqs;
+  int32_t write_req_capacity;
+
+  // ForwardBatch passthroughs. out_cache_loc is caller-pre-translated for SWA groups; the kernel
+  // treats it opaquely and skips entries with slot < 0.
+  const int64_t* input_ids;
+  const int64_t* positions;
+  const int64_t* out_cache_loc;
+
+  // Pseudo-mode oracle inputs.
+  bool enable_write_input_assert;
+  const int64_t* expected_input_tokens;
+  const int64_t* expected_input_positions;
+
+  // Violation sink (ring + write_index + capacity + kernel_kind bundled in canary_common.cuh).
+  ViolationSink violation_sink;
+
+  // Health counters.
+  int64_t* slot_run_counter;
+  int64_t* kernel_run_counter;
+
+  // Gates the chain-step position assert below. Default-on (1); CanaryManager zeros during the
+  // warmup window and flips back in mark_init_finished().
+  const int32_t* enable_chain_position_assert;
+
+  // Real-KV sources.
+  RealKvSourceHandle sources[kMaxRealKvSources];
+  int32_t num_sources;
+  RealKvHashMode real_kv_hash_mode;
+};
+
+__global__ void canary_write_kernel(const WriteKernelParams __grid_constant__ p) {
+  const uint32_t r = blockIdx.x;
+
+  // Unconditional kernel_run_counter bump (block 0 is always present).
+  if (r == 0) {
+    atomicAdd(reinterpret_cast<unsigned long long*>(p.kernel_run_counter), 1ULL);
+  }
+
+  const int32_t active = *p.write_num_valid_reqs;
+  if (r >= static_cast<uint32_t>(active)) {
+    return;
+  }
+
+  const int64_t entry_start = p.write_offsets[r];
+  const int64_t entry_end = p.write_offsets[r + 1];
+  const int64_t entry_count = entry_end - entry_start;
+  if (entry_count <= 0) {
+    return;
+  }
+
+  const int64_t seed_slot_idx = p.write_seed_slot_indices[r];
+
+  // Initialize running_prev_hash by advancing the chain from the seed slot
+  uint64_t running_prev_hash =
+      compute_slot_hash(p.canary_buf, p.slot_stride_bytes, static_cast<int64_t>(seed_slot_idx));
+
+  // Assumes eagle topk=1 (linear chain). Under topk>1 target_verify would be a tree and
+  // sibling positions share parent.pos+1, breaking this invariant.
+  const bool do_chain_position_assert = (seed_slot_idx >= 0) && (*p.enable_chain_position_assert != 0);
+  int64_t running_prev_position = 0;
+  if (do_chain_position_assert) {
+    running_prev_position = canary_load_field(p.canary_buf, seed_slot_idx, p.slot_stride_bytes, kCanaryFieldPosition);
+  }
+
+  int64_t entries_written = 0;
+  for (int64_t entry_offset = 0; entry_offset < entry_count; ++entry_offset) {
+    const int64_t entry_idx = entry_start + entry_offset;
+    const int64_t slot = p.out_cache_loc[entry_idx];
+
+    if (slot < 0) {
+      continue;
+    }
+    ++entries_written;
+
+    const int64_t token = p.input_ids[entry_idx];
+    const int64_t position = p.positions[entry_idx];
+
+    const uint64_t real_kv_hash_u64 = real_kv_fold_sources(p.sources, p.num_sources, slot, p.real_kv_hash_mode);
+    const int64_t real_kv_hash = static_cast<int64_t>(real_kv_hash_u64);
+
+    if (p.enable_write_input_assert) {
+      const int64_t expected_token = p.expected_input_tokens[entry_idx];
+      const int64_t expected_position = p.expected_input_positions[entry_idx];
+      FailReason mismatch_bits{};
+      if (token != expected_token) {
+        mismatch_bits |= FailReason::kWriteTokenMismatch;
+      }
+      if (position != expected_position) {
+        mismatch_bits |= FailReason::kWritePositionMismatch;
+      }
+      if (mismatch_bits != FailReason{}) {
+        record_violation(
+            p.violation_sink,
+            ViolationRow{
+                /* slot_idx = */ slot,
+                /* position = */ position,
+                /* stored_token = */ token,
+                /* expected_token = */ expected_token,
+                /* stored_chain_hash = */ static_cast<int64_t>(running_prev_hash),
+                /* expected_aux = expected_position */ expected_position,
+                /* fail_reason_bits = */ static_cast<int64_t>(mismatch_bits),
+            });
+      }
+    }
+
+    if (do_chain_position_assert) {
+      const int64_t expected_position_chain = running_prev_position + 1;
+      if (position != expected_position_chain) {
+        record_violation(
+            p.violation_sink,
+            ViolationRow{
+                /* slot_idx = */ slot,
+                /* position = */ position,
+                /* stored_token = */ token,
+                /* expected_token = */ token,
+                /* stored_chain_hash = */ static_cast<int64_t>(running_prev_hash),
+                /* expected_aux = expected_position */ expected_position_chain,
+                /* fail_reason_bits = */ static_cast<int64_t>(FailReason::kWritePositionMismatch),
+            });
+      }
+      running_prev_position = position;
+    }
+
+    canary_store_field(p.canary_buf, slot, p.slot_stride_bytes, kCanaryFieldToken, token);
+    canary_store_field(p.canary_buf, slot, p.slot_stride_bytes, kCanaryFieldPosition, position);
+    canary_store_field(
+        p.canary_buf, slot, p.slot_stride_bytes, kCanaryFieldPrevHash, static_cast<int64_t>(running_prev_hash));
+    canary_store_field(p.canary_buf, slot, p.slot_stride_bytes, kCanaryFieldRealKvHash, real_kv_hash);
+
+    running_prev_hash =
+        splitmix64_mix3(running_prev_hash, static_cast<uint64_t>(token), static_cast<uint64_t>(position));
+  }
+
+  // Each block contributes its non-skipped entry count to slot_run_counter once at exit.
+  atomicAdd(
+      reinterpret_cast<unsigned long long*>(p.slot_run_counter), static_cast<unsigned long long>(entries_written));
+}
+
+}  // namespace
+
+// API source of truth: docstring of canary_write_step in python/sglang/jit_kernel/kv_canary/write.py.
+//
+// ABI notes (same as verify):
+// - real_kv_buf_0 .. real_kv_buf_3 are 4 fixed uint8 tensor slots.
+// - real_kv_source_params is a CPU int32 [kMaxRealKvSources, 3] table of (page_size, num_bytes_per_token,
+//   read_bytes) triplets.
+// - out_cache_loc is caller-pre-translated for SWA groups; -1 entries mark skip. The kernel does not
+//   consult any LUT.
+inline void canary_write_step_cuda(
+    tvm::ffi::TensorView canary_buf,
+    tvm::ffi::TensorView write_offsets,
+    tvm::ffi::TensorView write_seed_slot_indices,
+    tvm::ffi::TensorView write_num_valid_reqs,
+    tvm::ffi::TensorView input_ids,
+    tvm::ffi::TensorView positions,
+    tvm::ffi::TensorView out_cache_loc,
+    int64_t kernel_kind,
+    int64_t enable_write_input_assert,
+    const tvm::ffi::Optional<tvm::ffi::TensorView> expected_input_tokens,
+    const tvm::ffi::Optional<tvm::ffi::TensorView> expected_input_positions,
+    tvm::ffi::TensorView violation_ring,
+    tvm::ffi::TensorView violation_write_index,
+    tvm::ffi::TensorView slot_run_counter,
+    tvm::ffi::TensorView kernel_run_counter,
+    tvm::ffi::TensorView enable_chain_position_assert,
+    tvm::ffi::TensorView real_kv_buf_0,
+    tvm::ffi::TensorView real_kv_buf_1,
+    tvm::ffi::TensorView real_kv_buf_2,
+    tvm::ffi::TensorView real_kv_buf_3,
+    tvm::ffi::TensorView real_kv_source_params,
+    int64_t num_sources,
+    int64_t real_kv_hash_mode) {
+  using namespace host;
+
+  SymbolicSize N_slots = {"num_canary_slots"};
+  SymbolicSize N_stride = {"slot_stride_bytes"};
+  SymbolicSize N_write_reqs = {"write_req_capacity"};
+  SymbolicSize N_tokens = {"num_tokens_padded"};
+  SymbolicDevice device_;
+  device_.set_options<kDLGPU>();
+
+  TensorMatcher({N_slots, N_stride}).with_dtype<uint8_t>().with_device<kDLGPU>(device_).verify(canary_buf);
+
+  // write_offsets has shape [write_req_capacity + 1]; the length relationship is pinned by the
+  // RuntimeCheck below, this matcher pins dtype + device.
+  SymbolicSize N_write_offsets = {"write_offsets_len"};
+  TensorMatcher({N_write_offsets}).with_dtype<int64_t>().with_device<kDLGPU>(device_).verify(write_offsets);
+  TensorMatcher({N_write_reqs}).with_dtype<int64_t>().with_device<kDLGPU>(device_).verify(write_seed_slot_indices);
+  TensorMatcher({1}).with_dtype<int32_t>().with_device<kDLGPU>(device_).verify(write_num_valid_reqs);
+
+  TensorMatcher({N_tokens})
+      .with_dtype<int64_t>()
+      .with_device<kDLGPU>(device_)
+      .verify(input_ids)
+      .verify(positions)
+      .verify(out_cache_loc);
+  const bool enable_write_input_assert_bool = (enable_write_input_assert != 0);
+  RuntimeCheck(
+      enable_write_input_assert_bool == expected_input_tokens.has_value(),
+      "canary_write: expected_input_tokens presence must match enable_write_input_assert");
+  RuntimeCheck(
+      enable_write_input_assert_bool == expected_input_positions.has_value(),
+      "canary_write: expected_input_positions presence must match enable_write_input_assert");
+  if (enable_write_input_assert_bool) {
+    TensorMatcher({N_tokens})
+        .with_dtype<int64_t>()
+        .with_device<kDLGPU>(device_)
+        .verify(expected_input_tokens.value())
+        .verify(expected_input_positions.value());
+  }
+
+  TensorMatcher({1}).with_dtype<int32_t>().with_device<kDLGPU>(device_).verify(violation_write_index);
+  SymbolicSize N_ring = {"ring_capacity"};
+  TensorMatcher({N_ring, static_cast<int64_t>(kViolationFields)})
+      .with_dtype<int64_t>()
+      .with_device<kDLGPU>(device_)
+      .verify(violation_ring);
+  TensorMatcher({1})
+      .with_dtype<int64_t>()
+      .with_device<kDLGPU>(device_)
+      .verify(slot_run_counter)
+      .verify(kernel_run_counter);
+  TensorMatcher({1}).with_dtype<int32_t>().with_device<kDLGPU>(device_).verify(enable_chain_position_assert);
+
+  SymbolicSize N_real_kv_rows_0 = {"real_kv_rows_0"};
+  SymbolicSize N_real_kv_cols_0 = {"real_kv_cols_0"};
+  TensorMatcher({N_real_kv_rows_0, N_real_kv_cols_0})
+      .with_dtype<uint8_t>()
+      .with_device<kDLGPU>(device_)
+      .verify(real_kv_buf_0);
+  SymbolicSize N_real_kv_rows_1 = {"real_kv_rows_1"};
+  SymbolicSize N_real_kv_cols_1 = {"real_kv_cols_1"};
+  TensorMatcher({N_real_kv_rows_1, N_real_kv_cols_1})
+      .with_dtype<uint8_t>()
+      .with_device<kDLGPU>(device_)
+      .verify(real_kv_buf_1);
+  SymbolicSize N_real_kv_rows_2 = {"real_kv_rows_2"};
+  SymbolicSize N_real_kv_cols_2 = {"real_kv_cols_2"};
+  TensorMatcher({N_real_kv_rows_2, N_real_kv_cols_2})
+      .with_dtype<uint8_t>()
+      .with_device<kDLGPU>(device_)
+      .verify(real_kv_buf_2);
+  SymbolicSize N_real_kv_rows_3 = {"real_kv_rows_3"};
+  SymbolicSize N_real_kv_cols_3 = {"real_kv_cols_3"};
+  TensorMatcher({N_real_kv_rows_3, N_real_kv_cols_3})
+      .with_dtype<uint8_t>()
+      .with_device<kDLGPU>(device_)
+      .verify(real_kv_buf_3);
+  TensorMatcher({static_cast<int64_t>(kMaxRealKvSources), static_cast<int64_t>(kRealKvSourceFieldsPerEntry)})
+      .with_dtype<int32_t>()
+      .with_device<kDLCPU>()
+      .verify(real_kv_source_params);
+
+  const int64_t slot_stride_bytes = N_stride.unwrap();
+  const int32_t write_req_capacity = static_cast<int32_t>(N_write_reqs.unwrap());
+  const int32_t ring_capacity = static_cast<int32_t>(N_ring.unwrap());
+  const DLDevice device = device_.unwrap();
+
+  RuntimeCheck(
+      write_offsets.size(0) == static_cast<int64_t>(write_req_capacity) + 1,
+      "canary_write: write_offsets.size(0) must equal write_req_capacity + 1 (",
+      static_cast<int64_t>(write_req_capacity) + 1,
+      "), got ",
+      write_offsets.size(0));
+  RuntimeCheck(
+      slot_stride_bytes >= static_cast<int64_t>(kCanaryFieldsPerSlot * sizeof(int64_t)),
+      "canary_write: slot_stride_bytes must hold at least ",
+      static_cast<int64_t>(kCanaryFieldsPerSlot * sizeof(int64_t)),
+      " bytes per slot, got ",
+      slot_stride_bytes);
+  RuntimeCheck(
+      num_sources >= 0 && num_sources <= static_cast<int64_t>(kMaxRealKvSources),
+      "canary_write: num_sources must be in [0, ",
+      static_cast<int64_t>(kMaxRealKvSources),
+      "], got ",
+      num_sources);
+
+  WriteKernelParams p{};
+  p.canary_buf = static_cast<uint8_t*>(canary_buf.data_ptr());
+  p.slot_stride_bytes = slot_stride_bytes;
+  p.write_offsets = static_cast<const int64_t*>(write_offsets.data_ptr());
+  p.write_seed_slot_indices = static_cast<const int64_t*>(write_seed_slot_indices.data_ptr());
+  p.write_num_valid_reqs = static_cast<const int32_t*>(write_num_valid_reqs.data_ptr());
+  p.write_req_capacity = write_req_capacity;
+  p.input_ids = static_cast<const int64_t*>(input_ids.data_ptr());
+  p.positions = static_cast<const int64_t*>(positions.data_ptr());
+  p.out_cache_loc = static_cast<const int64_t*>(out_cache_loc.data_ptr());
+  p.enable_write_input_assert = enable_write_input_assert_bool;
+  p.expected_input_tokens =
+      enable_write_input_assert_bool ? static_cast<const int64_t*>(expected_input_tokens.value().data_ptr()) : nullptr;
+  p.expected_input_positions = enable_write_input_assert_bool
+                                   ? static_cast<const int64_t*>(expected_input_positions.value().data_ptr())
+                                   : nullptr;
+  p.violation_sink.ring = static_cast<int64_t*>(violation_ring.data_ptr());
+  p.violation_sink.write_index = static_cast<int32_t*>(violation_write_index.data_ptr());
+  p.violation_sink.ring_capacity = ring_capacity;
+  p.violation_sink.kernel_kind = static_cast<int32_t>(kernel_kind);
+  p.slot_run_counter = static_cast<int64_t*>(slot_run_counter.data_ptr());
+  p.kernel_run_counter = static_cast<int64_t*>(kernel_run_counter.data_ptr());
+  p.enable_chain_position_assert = static_cast<const int32_t*>(enable_chain_position_assert.data_ptr());
+
+  const int32_t* params = static_cast<const int32_t*>(real_kv_source_params.data_ptr());
+  tvm::ffi::TensorView source_bufs[kMaxRealKvSources] = {real_kv_buf_0, real_kv_buf_1, real_kv_buf_2, real_kv_buf_3};
+  for (int s = 0; s < kMaxRealKvSources; ++s) {
+    p.sources[s].tensor = static_cast<const uint8_t*>(source_bufs[s].data_ptr());
+    p.sources[s].row_stride_bytes = static_cast<int32_t>(source_bufs[s].size(1));
+    p.sources[s].page_size = params[s * kRealKvSourceFieldsPerEntry + kRealKvSourceFieldPageSize];
+    p.sources[s].num_bytes_per_token = params[s * kRealKvSourceFieldsPerEntry + kRealKvSourceFieldNumBytesPerToken];
+    p.sources[s].read_bytes = params[s * kRealKvSourceFieldsPerEntry + kRealKvSourceFieldReadBytes];
+  }
+  p.num_sources = static_cast<int32_t>(num_sources);
+  p.real_kv_hash_mode = static_cast<RealKvHashMode>(real_kv_hash_mode);
+
+  // Grid: one block per write req capacity slot; the kernel early-exits on r >= write_num_valid_reqs[0].
+  // Always launch at least one block so the unconditional kernel_run_counter bump runs even when capacity
+  // == 0.
+  const uint32_t grid = write_req_capacity == 0 ? 1u : static_cast<uint32_t>(write_req_capacity);
+  LaunchKernel(grid, kWriteBlockSize, device)(canary_write_kernel, p);
+}
+
+}  // namespace canary

python/sglang/jit_kernel/csrc/kv_canary/consts.cuh

@@ -0,0 +1,62 @@
+#pragma once
+
+#include <cstdint>
+
+namespace canary {
+
+constexpr uint64_t kCanaryChainAnchor = 0xC0FFEE1234567890ULL;
+
+// Mirrors SGLang's TokenToKVPoolAllocator contract: token-to-KV slot 0 is reserved for padded-token dummy
+// writes. Since req_to_token stores token-to-KV slot ids and is zero-initialized, canary slot 0 is skipped
+// instead of treating unfilled entries as real KV slots.
+constexpr int64_t kTokenToKvSlotPadding = 0;
+constexpr int64_t kReqPoolIdxPadding = 0;
+
+constexpr int kCanaryFieldsPerSlot = 4;
+constexpr int kCanaryFieldToken = 0;
+constexpr int kCanaryFieldPosition = 1;
+constexpr int kCanaryFieldPrevHash = 2;
+constexpr int kCanaryFieldRealKvHash = 3;
+
+constexpr int kViolationFields = 8;
+constexpr int kViolationFieldKernelKind = 0;
+constexpr int kViolationFieldSlotIdx = 1;
+constexpr int kViolationFieldPosition = 2;
+constexpr int kViolationFieldStoredToken = 3;
+constexpr int kViolationFieldExpectedToken = 4;
+constexpr int kViolationFieldStoredChainHash = 5;
+constexpr int kViolationFieldExpectedAux = 6;
+constexpr int kViolationFieldFailReasonBits = 7;
+
+enum class FailReason : int64_t {
+  kVerifyChainHashMismatch = 1LL << 0,
+  kVerifyPositionMismatch = 1LL << 1,
+  kVerifyRealKvHashMismatch = 1LL << 2,
+  kWriteTokenMismatch = 1LL << 3,
+  kWritePositionMismatch = 1LL << 4,
+  kVerifyTokenMismatch = 1LL << 5,
+};
+
+constexpr FailReason operator|(FailReason a, FailReason b) {
+  return static_cast<FailReason>(static_cast<int64_t>(a) | static_cast<int64_t>(b));
+}
+
+constexpr FailReason& operator|=(FailReason& a, FailReason b) {
+  a = a | b;
+  return a;
+}
+
+enum class RealKvHashMode : int32_t {
+  kNone = 0,
+  kPartial = 1,
+  kAll = 2,
+};
+
+constexpr int kMaxRealKvSources = 4;
+
+constexpr int kRealKvSourceFieldsPerEntry = 3;
+constexpr int kRealKvSourceFieldPageSize = 0;
+constexpr int kRealKvSourceFieldNumBytesPerToken = 1;
+constexpr int kRealKvSourceFieldReadBytes = 2;
+
+}  // namespace canary