From 736ad1f32a442e09cdcfb9ede12faebe94613b9e Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Sun, 31 May 2026 09:53:56 +0800 Subject: [PATCH] Add the KV-canary plan JIT kernels (#26807) --- .../benchmark/kv_canary/bench_plan.py | 345 +++++ .../csrc/kv_canary/canary_plan_entries.cuh | 308 ++++ .../jit_kernel/kv_canary/plan/__init__.py | 1 + .../sglang/jit_kernel/kv_canary/plan/api.py | 165 ++ .../kv_canary/plan/entries_kernel.py | 71 + .../kv_canary/plan/offsets_kernel.py | 441 ++++++ .../sglang/jit_kernel/kv_canary/plan/utils.py | 97 ++ .../sglang/jit_kernel/kv_canary/plan_ref.py | 317 ++++ .../tests/kv_canary/_differential.py | 500 ++++++ .../tests/kv_canary/_fuzz_driver.py | 45 + .../tests/kv_canary/test_kernel_config.py | 425 +++++ .../tests/kv_canary/test_pipeline_e2e.py | 790 ++++++++++ .../tests/kv_canary/test_plan_fuzz.py | 222 +++ .../tests/kv_canary/test_plan_hand.py | 1371 +++++++++++++++++ .../tests/kv_canary/test_verify_fuzz.py | 178 +++ .../tests/kv_canary/test_verify_hand.py | 1266 +++++++++++++++ .../tests/kv_canary/test_write_fuzz.py | 218 +++ .../tests/kv_canary/test_write_hand.py | 892 +++++++++++ 18 files changed, 7652 insertions(+) create mode 100644 python/sglang/jit_kernel/benchmark/kv_canary/bench_plan.py create mode 100644 python/sglang/jit_kernel/csrc/kv_canary/canary_plan_entries.cuh create mode 100644 python/sglang/jit_kernel/kv_canary/plan/__init__.py create mode 100644 python/sglang/jit_kernel/kv_canary/plan/api.py create mode 100644 python/sglang/jit_kernel/kv_canary/plan/entries_kernel.py create mode 100644 python/sglang/jit_kernel/kv_canary/plan/offsets_kernel.py create mode 100644 python/sglang/jit_kernel/kv_canary/plan/utils.py create mode 100644 python/sglang/jit_kernel/kv_canary/plan_ref.py create mode 100644 python/sglang/jit_kernel/tests/kv_canary/_differential.py create mode 100644 python/sglang/jit_kernel/tests/kv_canary/_fuzz_driver.py create mode 100644 python/sglang/jit_kernel/tests/kv_canary/test_kernel_config.py create mode 100644 python/sglang/jit_kernel/tests/kv_canary/test_pipeline_e2e.py create mode 100644 python/sglang/jit_kernel/tests/kv_canary/test_plan_fuzz.py create mode 100644 python/sglang/jit_kernel/tests/kv_canary/test_plan_hand.py create mode 100644 python/sglang/jit_kernel/tests/kv_canary/test_verify_fuzz.py create mode 100644 python/sglang/jit_kernel/tests/kv_canary/test_verify_hand.py create mode 100644 python/sglang/jit_kernel/tests/kv_canary/test_write_fuzz.py create mode 100644 python/sglang/jit_kernel/tests/kv_canary/test_write_hand.py diff --git a/python/sglang/jit_kernel/benchmark/kv_canary/bench_plan.py b/python/sglang/jit_kernel/benchmark/kv_canary/bench_plan.py new file mode 100644 index 000000000..c130f3799 --- /dev/null +++ b/python/sglang/jit_kernel/benchmark/kv_canary/bench_plan.py @@ -0,0 +1,345 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch +import triton +import triton.testing + +from sglang.jit_kernel.benchmark.kv_canary.utils import ( + POOL_AXIS, + SWA_WINDOW, + BenchCase, + build_fast_matrix_cases, + build_full_matrix_cases, + naive_cumsum_fn, +) +from sglang.jit_kernel.benchmark.utils import ( + DEFAULT_DEVICE, + get_benchmark_range, + run_benchmark, +) +from sglang.jit_kernel.kv_canary.plan import launch_canary_plan_kernels +from sglang.jit_kernel.kv_canary.verify import VerifyPlan +from sglang.jit_kernel.kv_canary.write import WritePlan +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=900, suite="nightly-kernel-1-gpu", nightly=True) + + +_TOTAL_TOKENS_AXIS: list[int] = [256, 4096, 65536, 262144] +_TOTAL_TOKENS_BS_AXIS: list[int] = [1, 32, 256] + + +@dataclass(frozen=True, slots=True, kw_only=True) +class _TotalTokensBenchCase: + bs: int + total_tokens: int + pool_kind: str + + +def _build_total_tokens_cases() -> list[_TotalTokensBenchCase]: + cases: list[_TotalTokensBenchCase] = [] + for bs in _TOTAL_TOKENS_BS_AXIS: + for total_tokens in _TOTAL_TOKENS_AXIS: + if total_tokens < bs: + continue + for pool_kind in POOL_AXIS: + cases.append( + _TotalTokensBenchCase( + bs=bs, total_tokens=total_tokens, pool_kind=pool_kind + ) + ) + return cases + + +_POOL_CAPACITY_VERIFY_CAP_AXIS: list[int] = [16384, 262144, 1398028] +_POOL_CAPACITY_BS_AXIS: list[int] = [1, 4, 32] +_POOL_CAPACITY_PREFIX_LEN: int = 512 +_POOL_CAPACITY_BS_PADDED_AXIS: list[Optional[int]] = [None, 4096] + + +@dataclass(frozen=True, slots=True, kw_only=True) +class _PoolCapacityBenchCase: + """One pool-capacity bench point. + + Attributes: + bs: Number of active (non-padding) requests in the launch. + bs_padded: Total request-axis size of the input tensors. ``None`` means + no padding (``bs_padded == bs``); a concrete value pads ``req_pool_indices`` + with ``REQ_POOL_IDX_PADDING`` sentinels in rows ``[bs, bs_padded)``. + prefix_len: Per-active-request prefix length. + verify_capacity: Plan tensor row capacity. + pool_kind: "full". + """ + + bs: int + bs_padded: Optional[int] + prefix_len: int + verify_capacity: int + pool_kind: str + + +def _build_pool_capacity_cases() -> list[_PoolCapacityBenchCase]: + cases: list[_PoolCapacityBenchCase] = [] + for bs in _POOL_CAPACITY_BS_AXIS: + for verify_capacity in _POOL_CAPACITY_VERIFY_CAP_AXIS: + for bs_padded in _POOL_CAPACITY_BS_PADDED_AXIS: + if bs_padded is not None and bs_padded < bs: + continue + cases.append( + _PoolCapacityBenchCase( + bs=bs, + bs_padded=bs_padded, + prefix_len=_POOL_CAPACITY_PREFIX_LEN, + verify_capacity=verify_capacity, + pool_kind="full", + ) + ) + return cases + + +_X_NAMES_MATRIX = ["scenario", "bs", "prefix_len", "mode", "extend_len", "pool_kind"] + + +def _cases_to_matrix_x_vals( + cases: list[BenchCase], +) -> list[tuple[str, int, int, str, int, str]]: + return [ + (c.scenario, c.bs, c.prefix_len, c.mode, c.extend_len, c.pool_kind) + for c in cases + ] + + +_X_VALS_MATRIX = _cases_to_matrix_x_vals( + get_benchmark_range( + full_range=build_full_matrix_cases(), + ci_range=build_fast_matrix_cases(), + ) +) + +_X_NAMES_TT = ["bs", "total_tokens", "pool_kind"] +_X_VALS_TT = [(c.bs, c.total_tokens, c.pool_kind) for c in _build_total_tokens_cases()] + +_X_NAMES_PC = ["bs", "bs_padded", "prefix_len", "verify_capacity", "pool_kind"] +_X_VALS_PC = [ + ( + c.bs, + c.bs_padded if c.bs_padded is not None else c.bs, + c.prefix_len, + c.verify_capacity, + c.pool_kind, + ) + for c in _build_pool_capacity_cases() +] + + +def _build_plan_inputs( + *, + bs: int, + prefix_len: int, + extend_len: int, + pool_kind: str, + device: torch.device, + verify_capacity_override: Optional[int] = None, + bs_padded: Optional[int] = None, +) -> dict: + swa_window_size = SWA_WINDOW if pool_kind == "swa_window_128" else 0 + verify_per_req = min(prefix_len, SWA_WINDOW) if swa_window_size > 0 else prefix_len + if verify_capacity_override is not None: + verify_capacity = max(1, verify_capacity_override) + else: + verify_capacity = max(1, bs * verify_per_req) + + effective_bs = bs_padded if bs_padded is not None else bs + if effective_bs < bs: + raise ValueError(f"kv-canary bench: bs_padded={bs_padded} must be >= bs={bs}") + write_req_capacity = max(1, effective_bs) + + verify_plan = VerifyPlan.allocate(verify_capacity=verify_capacity, device=device) + write_plan = WritePlan.allocate( + write_req_capacity=write_req_capacity, device=device + ) + + req_pool_indices = torch.zeros(effective_bs, dtype=torch.int64, device=device) + req_pool_indices[:bs] = torch.arange(1, bs + 1, dtype=torch.int64, device=device) + prefix_lens = torch.zeros(effective_bs, dtype=torch.int64, device=device) + prefix_lens[:bs] = prefix_len + extend_seq_lens = torch.zeros(effective_bs, dtype=torch.int64, device=device) + extend_seq_lens[:bs] = extend_len + + max_seq_len = max(prefix_len + extend_len, 1) + req_to_token_rows = effective_bs + 1 + req_to_token = torch.zeros( + req_to_token_rows, + max_seq_len, + dtype=torch.int32, + device=device, + ) + if bs > 0: + row_idx = torch.arange(1, bs + 1, dtype=torch.int32, device=device).unsqueeze(1) + col_idx = torch.arange(max_seq_len, dtype=torch.int32, device=device).unsqueeze( + 0 + ) + req_to_token[1 : bs + 1] = (row_idx - 1) * max_seq_len + col_idx + + if swa_window_size > 0: + full_pool_size = effective_bs * max_seq_len + 1 + full_to_swa: Optional[torch.Tensor] = torch.arange( + full_pool_size + 1, dtype=torch.int64, device=device + ) + full_to_swa[-1] = -1 + else: + full_to_swa = None + + return dict( + verify_plan_out=verify_plan, + write_plan_out=write_plan, + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + extend_seq_lens=extend_seq_lens, + req_to_token=req_to_token, + swa_window_size=swa_window_size, + full_to_swa_index_mapping=full_to_swa, + verify_capacity=int(verify_plan.verify_slot_indices.shape[0]), + ) + + +def _make_plan_callable(inputs: dict): + def fn() -> None: + launch_canary_plan_kernels( + verify_plan_out=inputs["verify_plan_out"], + write_plan_out=inputs["write_plan_out"], + req_pool_indices=inputs["req_pool_indices"], + prefix_lens=inputs["prefix_lens"], + extend_seq_lens=inputs["extend_seq_lens"], + req_to_token=inputs["req_to_token"], + swa_window_size=inputs["swa_window_size"], + full_to_swa_index_mapping=inputs["full_to_swa_index_mapping"], + verify_capacity=inputs["verify_capacity"], + req_to_verify_expected_tokens=None, + req_to_verify_expected_tokens_valid_lens=None, + kv_token_id_vs_position_offset=0, + ) + + return fn + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=_X_NAMES_MATRIX, + x_vals=_X_VALS_MATRIX, + line_arg="provider", + line_vals=["canary", "naive"], + line_names=["canary_plan_step", "naive torch.cumsum"], + styles=[("blue", "-"), ("red", "--")], + ylabel="us", + plot_name="kv-canary-plan-matrix-perf", + args={}, + ) +) +def benchmark_matrix( + scenario: str, + bs: int, + prefix_len: int, + mode: str, + extend_len: int, + pool_kind: str, + provider: str, +) -> Tuple[float, float, float]: + del scenario + del mode + + device = torch.device(DEFAULT_DEVICE) + if provider == "canary": + inputs = _build_plan_inputs( + bs=bs, + prefix_len=prefix_len, + extend_len=extend_len, + pool_kind=pool_kind, + device=device, + ) + fn = _make_plan_callable(inputs) + else: + fn = naive_cumsum_fn(bs=bs, device=device) + return run_benchmark(fn) + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=_X_NAMES_TT, + x_vals=_X_VALS_TT, + line_arg="provider", + line_vals=["canary", "naive"], + line_names=["canary_plan_step", "naive torch.cumsum"], + styles=[("blue", "-"), ("red", "--")], + ylabel="us", + plot_name="kv-canary-plan-total-tokens-perf", + args={}, + ) +) +def benchmark_total_tokens( + bs: int, + total_tokens: int, + pool_kind: str, + provider: str, +) -> Tuple[float, float, float]: + device = torch.device(DEFAULT_DEVICE) + per_req_prefix = max(1, total_tokens // max(bs, 1)) + if provider == "canary": + inputs = _build_plan_inputs( + bs=bs, + prefix_len=per_req_prefix, + extend_len=1, + pool_kind=pool_kind, + device=device, + ) + fn = _make_plan_callable(inputs) + else: + fn = naive_cumsum_fn(bs=bs, device=device) + return run_benchmark(fn) + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=_X_NAMES_PC, + x_vals=_X_VALS_PC, + line_arg="provider", + line_vals=["canary", "naive"], + line_names=["canary_plan_step", "naive torch.cumsum"], + styles=[("blue", "-"), ("red", "--")], + ylabel="us", + plot_name="kv-canary-plan-pool-capacity-perf", + args={}, + ) +) +def benchmark_pool_capacity( + bs: int, + bs_padded: int, + prefix_len: int, + verify_capacity: int, + pool_kind: str, + provider: str, +) -> Tuple[float, float, float]: + device = torch.device(DEFAULT_DEVICE) + if provider == "canary": + inputs = _build_plan_inputs( + bs=bs, + prefix_len=prefix_len, + extend_len=1, + pool_kind=pool_kind, + device=device, + verify_capacity_override=verify_capacity, + bs_padded=bs_padded, + ) + fn = _make_plan_callable(inputs) + else: + fn = naive_cumsum_fn(bs=bs, device=device) + return run_benchmark(fn) + + +if __name__ == "__main__": + benchmark_matrix.run(print_data=True) + benchmark_total_tokens.run(print_data=True) + benchmark_pool_capacity.run(print_data=True) diff --git a/python/sglang/jit_kernel/csrc/kv_canary/canary_plan_entries.cuh b/python/sglang/jit_kernel/csrc/kv_canary/canary_plan_entries.cuh new file mode 100644 index 000000000..7c6e1c7ef --- /dev/null +++ b/python/sglang/jit_kernel/csrc/kv_canary/canary_plan_entries.cuh @@ -0,0 +1,308 @@ +#pragma once + +#include // For TensorMatcher, SymbolicSize, SymbolicDevice +#include // For RuntimeCheck + +#include // For host::runtime::get_sm_count +#include // For LaunchKernel, SGL_DEVICE + +#include +#include + +#include +#include + +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 +__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(total_verify), + static_cast(params.verify_capacity)); + } + __trap(); + } + + const int64_t tid_start = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t stride = static_cast(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(slot_raw)); + } else { + out_slot = static_cast(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(prev_raw)); + } else { + out_prev_slot = static_cast(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(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(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 +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 full_to_swa_index_mapping, + const tvm::ffi::TensorView verify_offsets_scratch, + const tvm::ffi::TensorView verify_enable, + const tvm::ffi::Optional req_to_verify_expected_tokens, + const tvm::ffi::Optional 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(); + + TensorMatcher({Nbs}) // + .with_dtype() + .with_device(device_) + .verify(req_pool_indices) + .verify(prefix_lens); + TensorMatcher({Nscratch}) // + .with_dtype() + .with_device(device_) + .verify(verify_offsets_scratch); + TensorMatcher({1}) // + .with_dtype() + .with_device(device_) + .verify(verify_enable); + TensorMatcher({Ncap}) // + .with_dtype() + .with_device(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() + .with_device(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() + .with_device(device_) + .verify(full_to_swa_index_mapping.value()); + } + if constexpr (HAS_VERIFY_EXPECTED_TOKEN_POOL) { + TensorMatcher({Npool_rows, Npool_cols}) // + .with_dtype() + .with_device(device_) + .verify(req_to_verify_expected_tokens.value()); + TensorMatcher({Nbs}) // + .with_dtype() + .with_device(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(full_to_swa_index_mapping.value().data_ptr()); + lut_len = static_cast(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(req_to_verify_expected_tokens.value().data_ptr()); + expected_token_ids_stride0 = static_cast(Npool_cols.unwrap()); + req_to_verify_expected_tokens_valid_lens_ptr = + static_cast(req_to_verify_expected_tokens_valid_lens.value().data_ptr()); + } + + const PlanEntriesParams params = PlanEntriesParams{ + .req_pool_indices = static_cast(req_pool_indices.data_ptr()), + .prefix_lens = static_cast(prefix_lens.data_ptr()), + .req_to_token = static_cast(req_to_token.data_ptr()), + .full_to_swa_lut = lut_ptr, + .verify_offsets_scratch = static_cast(verify_offsets_scratch.data_ptr()), + .verify_enable = static_cast(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(out_verify_slot_indices.data_ptr()), + .out_verify_expected_tokens = static_cast(out_verify_expected_tokens.data_ptr()), + .out_verify_expected_positions = static_cast(out_verify_expected_positions.data_ptr()), + .out_verify_prev_slot_indices = static_cast(out_verify_prev_slot_indices.data_ptr()), + .bs_padded = static_cast(bs_padded), + .verify_capacity = static_cast(Ncap.unwrap()), + .req_to_token_stride0 = static_cast(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, params, lut_len); + } +}; + +} // namespace diff --git a/python/sglang/jit_kernel/kv_canary/plan/__init__.py b/python/sglang/jit_kernel/kv_canary/plan/__init__.py new file mode 100644 index 000000000..e7bcc5c1d --- /dev/null +++ b/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 diff --git a/python/sglang/jit_kernel/kv_canary/plan/api.py b/python/sglang/jit_kernel/kv_canary/plan/api.py new file mode 100644 index 000000000..ac38ed8bc --- /dev/null +++ b/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; 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), + ) diff --git a/python/sglang/jit_kernel/kv_canary/plan/entries_kernel.py b/python/sglang/jit_kernel/kv_canary/plan/entries_kernel.py new file mode 100644 index 000000000..1a62b8c5b --- /dev/null +++ b/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), + ) diff --git a/python/sglang/jit_kernel/kv_canary/plan/offsets_kernel.py b/python/sglang/jit_kernel/kv_canary/plan/offsets_kernel.py new file mode 100644 index 000000000..7bf9cc725 --- /dev/null +++ b/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)) diff --git a/python/sglang/jit_kernel/kv_canary/plan/utils.py b/python/sglang/jit_kernel/kv_canary/plan/utils.py new file mode 100644 index 000000000..e59fe88dc --- /dev/null +++ b/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) diff --git a/python/sglang/jit_kernel/kv_canary/plan_ref.py b/python/sglang/jit_kernel/kv_canary/plan_ref.py new file mode 100644 index 000000000..a37d2440b --- /dev/null +++ b/python/sglang/jit_kernel/kv_canary/plan_ref.py @@ -0,0 +1,317 @@ +from __future__ import annotations + +from typing import Optional + +import torch + +from sglang.jit_kernel.kv_canary.consts import REQ_POOL_IDX_PADDING +from sglang.jit_kernel.kv_canary.verify import VerifyPlan +from sglang.jit_kernel.kv_canary.write import WritePlan + + +def launch_canary_plan_kernels_torch_reference( + *, + 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: + """Python reference for :func:`launch_canary_plan_kernels`. Same signature & byte-equal semantics.""" + bs = int(req_pool_indices.shape[0]) + work_device = torch.device("cpu") + + 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_torch_reference verify_capacity={verify_capacity} does not " + f"match verify_plan_out.verify_slot_indices.shape[0]={plan_verify_capacity}" + ) + write_req_capacity = int(write_plan_out.write_seed_slot_indices.shape[0]) + + req_pool_indices_host = req_pool_indices.detach().to( + device=work_device, dtype=torch.int64 + ) + prefix_lens_host = prefix_lens.detach().to(device=work_device, dtype=torch.int64) + extend_seq_lens_host = extend_seq_lens.detach().to( + device=work_device, dtype=torch.int64 + ) + req_to_token_host = req_to_token.detach().to(device=work_device, dtype=torch.int64) + + lut: Optional[torch.Tensor] = None + if full_to_swa_index_mapping is not None: + lut = full_to_swa_index_mapping.detach().to(device=work_device) + + expected_token_pool_host: Optional[torch.Tensor] = None + req_to_verify_expected_tokens_valid_lens_host: Optional[torch.Tensor] = None + if req_to_verify_expected_tokens is not None: + expected_token_pool_host = req_to_verify_expected_tokens.detach().to( + device=work_device, dtype=torch.int64 + ) + if req_to_verify_expected_tokens_valid_lens is None: + raise ValueError( + "kv-canary: launch_canary_plan_kernels_torch_reference requires " + "req_to_verify_expected_tokens_valid_lens when req_to_verify_expected_tokens is set" + ) + req_to_verify_expected_tokens_valid_lens_host = ( + req_to_verify_expected_tokens_valid_lens.detach().to( + device=work_device, dtype=torch.int64 + ) + ) + + total_verify = _materialize_verify_entries( + verify_plan_out=verify_plan_out, + req_pool_indices_host=req_pool_indices_host, + prefix_lens_host=prefix_lens_host, + req_to_token_host=req_to_token_host, + swa_window_size=swa_window_size, + lut=lut, + verify_capacity=verify_capacity, + work_device=work_device, + bs=bs, + expected_token_pool_host=expected_token_pool_host, + req_to_verify_expected_tokens_valid_lens_host=req_to_verify_expected_tokens_valid_lens_host, + kv_token_id_vs_position_offset=int(kv_token_id_vs_position_offset), + ) + + _materialize_write_metadata( + write_plan_out=write_plan_out, + req_pool_indices_host=req_pool_indices_host, + prefix_lens_host=prefix_lens_host, + extend_seq_lens_host=extend_seq_lens_host, + req_to_token_host=req_to_token_host, + lut=lut, + write_req_capacity=write_req_capacity, + work_device=work_device, + bs=bs, + ) + + _write_num_valid_and_enable( + verify_plan_out=verify_plan_out, + requested=total_verify, + verify_capacity=verify_capacity, + ) + + +def _write_num_valid_and_enable( + *, + verify_plan_out: VerifyPlan, + requested: int, + verify_capacity: int, +) -> None: + overflow = requested > verify_capacity + clamped = verify_capacity if overflow else requested + enable = 0 if overflow else 1 + verify_plan_out.verify_num_valid.fill_(int(clamped)) + verify_plan_out.enable.fill_(int(enable)) + + +def _swa_translate_slot(*, slot: int, lut: torch.Tensor) -> int: + if slot < 0: + return slot + lut_len = int(lut.shape[0]) + if slot >= lut_len: + raise ValueError( + f"kv-canary: SWA slot {slot} is outside full_to_swa_index_mapping length {lut_len}" + ) + return int(lut[slot].item()) + + +def _materialize_verify_entries( + *, + verify_plan_out: VerifyPlan, + req_pool_indices_host: torch.Tensor, + prefix_lens_host: torch.Tensor, + req_to_token_host: torch.Tensor, + swa_window_size: int, + lut: Optional[torch.Tensor], + verify_capacity: int, + work_device: torch.device, + bs: int, + expected_token_pool_host: Optional[torch.Tensor], + req_to_verify_expected_tokens_valid_lens_host: Optional[torch.Tensor], + kv_token_id_vs_position_offset: int, +) -> int: + out_slots: list[int] = [] + out_positions: list[int] = [] + out_expected_input_ids: list[int] = [] + out_prev_slots: list[int] = [] + + for r in range(bs): + rpi = int(req_pool_indices_host[r].item()) + prefix_len = int(prefix_lens_host[r].item()) + + if rpi == REQ_POOL_IDX_PADDING: + continue + + if swa_window_size > 0: + window_start = max(0, prefix_len - swa_window_size) + else: + window_start = 0 + verify_len = max(0, prefix_len - window_start) + + valid_len_r = ( + int(req_to_verify_expected_tokens_valid_lens_host[r].item()) + if req_to_verify_expected_tokens_valid_lens_host is not None + else 0 + ) + + for j in range(verify_len): + position = window_start + j + slot_full = int(req_to_token_host[rpi, position].item()) + + if lut is not None: + slot = _swa_translate_slot(slot=slot_full, lut=lut) + else: + slot = slot_full + + prev_position = position - 1 + if prev_position < 0: + prev_slot = -1 + else: + prev_slot_full = int(req_to_token_host[rpi, prev_position].item()) + if lut is not None: + prev_slot = _swa_translate_slot(slot=prev_slot_full, lut=lut) + else: + prev_slot = prev_slot_full + + expected_input_id = -1 + if expected_token_pool_host is not None: + sot_pos = position + kv_token_id_vs_position_offset + if 0 <= sot_pos < valid_len_r: + expected_input_id = int( + expected_token_pool_host[rpi, sot_pos].item() + ) + + out_slots.append(slot) + out_positions.append(position) + out_expected_input_ids.append(expected_input_id) + out_prev_slots.append(prev_slot) + + total_verify = len(out_slots) + if total_verify == 0: + return 0 + + # On overflow CUDA plan_entries skips scatter (verify_enable=0); mirror that. + if total_verify > verify_capacity: + return total_verify + + slots_t = torch.tensor(out_slots, dtype=torch.int64, device=work_device) + positions_t = torch.tensor(out_positions, dtype=torch.int64, device=work_device) + expected_input_ids_t = torch.tensor( + out_expected_input_ids, dtype=torch.int64, device=work_device + ) + prev_slots_t = torch.tensor(out_prev_slots, dtype=torch.int64, device=work_device) + + verify_plan_out.verify_slot_indices[:total_verify].copy_( + slots_t.to(verify_plan_out.verify_slot_indices.dtype).to( + verify_plan_out.verify_slot_indices.device + ) + ) + verify_plan_out.verify_expected_tokens[:total_verify].copy_( + expected_input_ids_t.to(verify_plan_out.verify_expected_tokens.dtype).to( + verify_plan_out.verify_expected_tokens.device + ) + ) + verify_plan_out.verify_expected_positions[:total_verify].copy_( + positions_t.to(verify_plan_out.verify_expected_positions.dtype).to( + verify_plan_out.verify_expected_positions.device + ) + ) + verify_plan_out.verify_prev_slot_indices[:total_verify].copy_( + prev_slots_t.to(verify_plan_out.verify_prev_slot_indices.dtype).to( + verify_plan_out.verify_prev_slot_indices.device + ) + ) + + return total_verify + + +def _materialize_write_metadata( + *, + write_plan_out: WritePlan, + req_pool_indices_host: torch.Tensor, + prefix_lens_host: torch.Tensor, + extend_seq_lens_host: torch.Tensor, + req_to_token_host: torch.Tensor, + lut: Optional[torch.Tensor], + write_req_capacity: int, + work_device: torch.device, + bs: int, +) -> None: + out_write_offsets_len = int(write_plan_out.write_offsets.shape[0]) + max_seq_len = int(req_to_token_host.shape[1]) + + write_offsets_list: list[int] = [] + seed_slots_list: list[int] = [] + + running_offset = 0 + for r in range(bs): + write_offsets_list.append(running_offset) + + rpi = int(req_pool_indices_host[r].item()) + extend_len = int(extend_seq_lens_host[r].item()) + + if rpi == REQ_POOL_IDX_PADDING or extend_len <= 0: + write_len = 0 + else: + write_len = max(0, extend_len) + + running_offset += write_len + + write_offsets_list.append(running_offset) + + copy_len = min(bs + 1, out_write_offsets_len) + write_offsets_t = torch.tensor( + write_offsets_list[:copy_len], dtype=torch.int64, device=work_device + ) + write_plan_out.write_offsets[:copy_len].copy_( + write_offsets_t.to(write_plan_out.write_offsets.dtype).to( + write_plan_out.write_offsets.device + ) + ) + if copy_len < out_write_offsets_len: + write_plan_out.write_offsets[copy_len:].zero_() + + capped_reqs = min(bs, write_req_capacity) + for r in range(capped_reqs): + rpi = int(req_pool_indices_host[r].item()) + prefix_len = int(prefix_lens_host[r].item()) + extend_len = int(extend_seq_lens_host[r].item()) + + if rpi == REQ_POOL_IDX_PADDING or extend_len <= 0: + seed_slots_list.append(-1) + continue + + if prefix_len <= 0: + seed_slots_list.append(-1) + continue + + safe_seed_pos = min(prefix_len - 1, max(max_seq_len - 1, 0)) + seed_slot_full = int(req_to_token_host[rpi, safe_seed_pos].item()) + + if lut is not None: + seed_slot = _swa_translate_slot(slot=seed_slot_full, lut=lut) + else: + seed_slot = seed_slot_full + + seed_slots_list.append(seed_slot) + + if len(seed_slots_list) > 0: + seed_slots_t = torch.tensor( + seed_slots_list, dtype=torch.int64, device=work_device + ) + write_plan_out.write_seed_slot_indices[:capped_reqs].copy_( + seed_slots_t.to(write_plan_out.write_seed_slot_indices.dtype).to( + write_plan_out.write_seed_slot_indices.device + ) + ) + + write_plan_out.write_num_valid_reqs.fill_(int(bs)) diff --git a/python/sglang/jit_kernel/tests/kv_canary/_differential.py b/python/sglang/jit_kernel/tests/kv_canary/_differential.py new file mode 100644 index 000000000..b05619777 --- /dev/null +++ b/python/sglang/jit_kernel/tests/kv_canary/_differential.py @@ -0,0 +1,500 @@ +from __future__ import annotations + +from dataclasses import dataclass, replace +from typing import Any, Callable, Iterator, Optional + +import torch + +from sglang.jit_kernel.kv_canary.plan import launch_canary_plan_kernels +from sglang.jit_kernel.kv_canary.plan_ref import ( + launch_canary_plan_kernels_torch_reference, +) +from sglang.jit_kernel.kv_canary.verify import ( + CanaryLaunchTag, + VerifyOrWriteContext, + VerifyPlan, + launch_canary_verify_kernel, +) +from sglang.jit_kernel.kv_canary.verify_ref import ( + launch_canary_verify_kernel_torch_reference, +) +from sglang.jit_kernel.kv_canary.write import WritePlan, launch_canary_write_kernel +from sglang.jit_kernel.kv_canary.write_ref import ( + launch_canary_write_kernel_torch_reference, +) +from sglang.jit_kernel.tests.kv_canary._canary_helpers import ( + FakeViolationLog, + assert_canary_buf_equal, + assert_canary_state_equal, + make_log_pair, +) + +_DEVICE = torch.device("cuda") + + +def _run_both_plan( + *, + triton_verify: VerifyPlan, + triton_write: WritePlan, + ref_verify: VerifyPlan, + ref_write: WritePlan, + req_pool_indices: torch.Tensor, + prefix_lens: torch.Tensor, + extend_seq_lens: torch.Tensor, + req_to_token: torch.Tensor, + extras: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], + swa_window_size: int, + full_to_swa_index_mapping: Optional[torch.Tensor], + assert_equal: bool = True, + active_verify_entries: Optional[int] = None, + active_write_reqs: Optional[int] = None, + req_to_verify_expected_tokens: Optional[torch.Tensor] = None, + req_to_verify_expected_tokens_valid_lens: Optional[torch.Tensor] = None, + kv_token_id_vs_position_offset: int = 0, +) -> None: + _ = extras + verify_capacity = int(triton_verify.verify_slot_indices.shape[0]) + # Default lens to "no tighter bound than pool width" so existing kernel tests that + # only care about gather wiring keep their old semantics without each call site + # explicitly building a per-req lens tensor. + if ( + req_to_verify_expected_tokens is not None + and req_to_verify_expected_tokens_valid_lens is None + ): + req_to_verify_expected_tokens_valid_lens = torch.full( + (int(req_pool_indices.shape[0]),), + int(req_to_verify_expected_tokens.shape[1]), + dtype=torch.int64, + device=req_pool_indices.device, + ) + launch_canary_plan_kernels( + verify_plan_out=triton_verify, + write_plan_out=triton_write, + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + extend_seq_lens=extend_seq_lens, + req_to_token=req_to_token, + swa_window_size=swa_window_size, + full_to_swa_index_mapping=full_to_swa_index_mapping, + verify_capacity=verify_capacity, + req_to_verify_expected_tokens=req_to_verify_expected_tokens, + req_to_verify_expected_tokens_valid_lens=req_to_verify_expected_tokens_valid_lens, + kv_token_id_vs_position_offset=kv_token_id_vs_position_offset, + ) + launch_canary_plan_kernels_torch_reference( + verify_plan_out=ref_verify, + write_plan_out=ref_write, + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + extend_seq_lens=extend_seq_lens, + req_to_token=req_to_token, + swa_window_size=swa_window_size, + full_to_swa_index_mapping=full_to_swa_index_mapping, + verify_capacity=int(ref_verify.verify_slot_indices.shape[0]), + req_to_verify_expected_tokens=req_to_verify_expected_tokens, + req_to_verify_expected_tokens_valid_lens=req_to_verify_expected_tokens_valid_lens, + kv_token_id_vs_position_offset=kv_token_id_vs_position_offset, + ) + torch.cuda.synchronize() + + if assert_equal: + _assert_plans_byte_equal( + triton_verify=triton_verify, + triton_write=triton_write, + ref_verify=ref_verify, + ref_write=ref_write, + active_verify_entries=active_verify_entries, + active_write_reqs=active_write_reqs, + ) + + +def _assert_plans_byte_equal( + *, + triton_verify: VerifyPlan, + triton_write: WritePlan, + ref_verify: VerifyPlan, + ref_write: WritePlan, + active_verify_entries: Optional[int] = None, + active_write_reqs: Optional[int] = None, +) -> None: + """Byte-equal check on (Triton vs ref) plan outputs. + + Optional ``active_verify_entries`` / ``active_write_reqs`` truncate the comparison to the meaningful + prefix; tail entries past the active count are kernel-undefined and need not match byte-equal. + """ + n_verify = ( + active_verify_entries + if active_verify_entries is not None + else int(triton_verify.verify_num_valid[0].item()) + ) + n_verify_ref = int(ref_verify.verify_num_valid[0].item()) + assert ( + n_verify == n_verify_ref + ), f"verify_num_valid diverged: triton={n_verify} ref={n_verify_ref}" + # When total_verify > VERIFY_CAPACITY the offsets kernel clears verify_enable and + # plan_entries skips its scatter — leaving verify_slot_indices/positions/prev_slot_indices + # as whatever the (torch.empty) allocation contained. Skip the byte-equal probe in that + # case; verify_num_valid being clamped + verify_enable=0 is the contract here. + triton_enable = int(triton_verify.enable[0].item()) + ref_enable = int(ref_verify.enable[0].item()) + assert ( + triton_enable == ref_enable + ), f"verify_enable diverged: triton={triton_enable} ref={ref_enable}" + if n_verify > 0 and triton_enable != 0: + assert torch.equal( + triton_verify.verify_slot_indices[:n_verify], + ref_verify.verify_slot_indices[:n_verify], + ) + assert torch.equal( + triton_verify.verify_expected_tokens[:n_verify], + ref_verify.verify_expected_tokens[:n_verify], + ) + assert torch.equal( + triton_verify.verify_expected_positions[:n_verify], + ref_verify.verify_expected_positions[:n_verify], + ) + assert torch.equal( + triton_verify.verify_prev_slot_indices[:n_verify], + ref_verify.verify_prev_slot_indices[:n_verify], + ) + + n_write = ( + active_write_reqs + if active_write_reqs is not None + else int(triton_write.write_num_valid_reqs[0].item()) + ) + n_write_ref = int(ref_write.write_num_valid_reqs[0].item()) + assert ( + n_write == n_write_ref + ), f"write_num_valid_reqs diverged: triton={n_write} ref={n_write_ref}" + assert torch.equal( + triton_write.write_offsets[: n_write + 1], + ref_write.write_offsets[: n_write + 1], + ) + if n_write > 0: + assert torch.equal( + triton_write.write_seed_slot_indices[:n_write], + ref_write.write_seed_slot_indices[:n_write], + ) + + +def _run_both_verify( + *, + cuda_canary_buf: torch.Tensor, + ref_canary_buf: torch.Tensor, + plan_cuda, + plan_ref, + cuda_log: FakeViolationLog, + ref_log: FakeViolationLog, + kernel_kind: CanaryLaunchTag = CanaryLaunchTag.HEAD_K_FULL, + assert_equal: bool = True, + check_verify_expected_token: bool = True, +) -> None: + launch_canary_verify_kernel( + context=VerifyOrWriteContext( + canary_buf=cuda_canary_buf, + kernel_kind=kernel_kind, + violation_ring=cuda_log.ring, + violation_write_index=cuda_log.write_index, + slot_run_counter=cuda_log.slot_run_counter, + kernel_run_counter=cuda_log.kernel_run_counter, + enable_chain_position_assert=cuda_log.enable_chain_position_assert, + ), + plan=plan_cuda, + check_verify_expected_token=check_verify_expected_token, + ) + launch_canary_verify_kernel_torch_reference( + context=VerifyOrWriteContext( + canary_buf=ref_canary_buf, + kernel_kind=kernel_kind, + violation_ring=ref_log.ring, + violation_write_index=ref_log.write_index, + slot_run_counter=ref_log.slot_run_counter, + kernel_run_counter=ref_log.kernel_run_counter, + enable_chain_position_assert=ref_log.enable_chain_position_assert, + ), + plan=plan_ref, + check_verify_expected_token=check_verify_expected_token, + ) + torch.cuda.synchronize() + + if assert_equal: + assert_canary_state_equal(log_a=cuda_log, log_b=ref_log) + + +def _run_both_write( + *, + cuda_canary_buf: torch.Tensor, + ref_canary_buf: torch.Tensor, + plan_cuda, + plan_ref, + input_ids: torch.Tensor, + positions: torch.Tensor, + out_cache_loc: torch.Tensor, + enable_write_verify_inputs: bool, + expected_input_tokens: torch.Tensor, + expected_input_positions: torch.Tensor, + cuda_log: FakeViolationLog, + ref_log: FakeViolationLog, + kernel_kind: CanaryLaunchTag = CanaryLaunchTag.HEAD_K_FULL, + assert_equal: bool = True, +) -> None: + expected_tokens_for_launch = ( + expected_input_tokens if enable_write_verify_inputs else None + ) + expected_positions_for_launch = ( + expected_input_positions if enable_write_verify_inputs else None + ) + launch_canary_write_kernel( + context=VerifyOrWriteContext( + canary_buf=cuda_canary_buf, + kernel_kind=kernel_kind, + violation_ring=cuda_log.ring, + violation_write_index=cuda_log.write_index, + slot_run_counter=cuda_log.slot_run_counter, + kernel_run_counter=cuda_log.kernel_run_counter, + enable_chain_position_assert=cuda_log.enable_chain_position_assert, + ), + plan=plan_cuda, + input_ids=input_ids, + positions=positions, + out_cache_loc=out_cache_loc, + enable_write_input_assert=enable_write_verify_inputs, + expected_input_tokens=expected_tokens_for_launch, + expected_input_positions=expected_positions_for_launch, + ) + launch_canary_write_kernel_torch_reference( + context=VerifyOrWriteContext( + canary_buf=ref_canary_buf, + kernel_kind=kernel_kind, + violation_ring=ref_log.ring, + violation_write_index=ref_log.write_index, + slot_run_counter=ref_log.slot_run_counter, + kernel_run_counter=ref_log.kernel_run_counter, + enable_chain_position_assert=ref_log.enable_chain_position_assert, + ), + plan=plan_ref, + input_ids=input_ids, + positions=positions, + out_cache_loc=out_cache_loc, + enable_write_input_assert=enable_write_verify_inputs, + expected_input_tokens=expected_tokens_for_launch, + expected_input_positions=expected_positions_for_launch, + ) + torch.cuda.synchronize() + + if assert_equal: + assert_canary_buf_equal(buf_a=cuda_canary_buf, buf_b=ref_canary_buf) + assert_canary_state_equal(log_a=cuda_log, log_b=ref_log) + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ShrinkResult: + inputs: Any + mutations_applied: list[str] + + +def shrink_inputs( + inputs: Any, + *, + check_fn: Callable[[Any], bool], + max_iterations: int = 50, +) -> ShrinkResult: + """Greedy 1-step minify for a fuzz inputs dataclass. + + ``check_fn(candidate)`` returns True when ``candidate`` still reproduces the failure. Each round + yields candidate-simpler-than-current mutations through ``_yield_simpler``; the first accepted + candidate becomes the new current. Iteration stops when no mutation is accepted or ``max_iterations`` + is reached. + """ + current = inputs + applied: list[str] = [] + for _ in range(max_iterations): + improved = False + for label, candidate in _yield_simpler(current): + try: + still_fails = check_fn(candidate) + except Exception: + still_fails = False + if still_fails: + current = candidate + applied.append(label) + improved = True + break + if not improved: + break + return ShrinkResult(inputs=current, mutations_applied=applied) + + +def _yield_simpler(inputs: Any) -> Iterator[tuple[str, Any]]: + """Yield (label, simpler_candidate) tuples for generic fuzz-input minifiers. + + The candidates touch only well-known field names; an inputs dataclass that lacks a field will simply + have that mutation skipped. No kernel-specific knowledge is encoded here so the same shrinker drives + Plan / Verify / Write fuzz failures uniformly. + """ + fields = { + f: getattr(inputs, f) for f in inputs.__dataclass_fields__ # type: ignore[attr-defined] + } + + def emit(label: str, **overrides: Any) -> Iterator[tuple[str, Any]]: + candidate = replace(inputs, **overrides) + yield label, candidate + + bs_field = ( + "req_pool_indices" + if "req_pool_indices" in fields + else ("input_ids" if "input_ids" in fields else None) + ) + if bs_field is not None and isinstance(fields[bs_field], torch.Tensor): + tensor = fields[bs_field] + if tensor.numel() > 1: + new_len = tensor.numel() - 1 + related_tensors_overrides: dict[str, Any] = {} + for name in ( + "req_pool_indices", + "prefix_lens", + "extend_seq_lens", + "input_ids", + "positions", + "out_cache_loc", + "expected_input_tokens", + "expected_input_positions", + ): + t = fields.get(name) + if ( + isinstance(t, torch.Tensor) + and t.numel() >= new_len + and t.dim() == 1 + ): + related_tensors_overrides[name] = t[:new_len].contiguous() + if related_tensors_overrides: + yield from emit("drop_last_row", **related_tensors_overrides) + + if "swa_window_size" in fields and isinstance(fields["swa_window_size"], int): + if fields["swa_window_size"] != 0: + yield from emit( + "swa_off", swa_window_size=0, full_to_swa_index_mapping=None + ) + + if "extras_count" in fields and isinstance(fields["extras_count"], int): + if fields["extras_count"] > 0: + yield from emit("extras_zero", extras_count=0) + + if "enable_write_verify_inputs" in fields: + cur = fields["enable_write_verify_inputs"] + if hasattr(cur, "value") and int(cur) != 0: + cls = cur.__class__ + yield from emit("pseudo_off", enable_write_verify_inputs=cls(0)) + + for name in ("verify_capacity", "write_req_capacity"): + if name in fields and isinstance(fields[name], int): + current_value = fields[name] + if current_value > 8: + yield from emit(f"shrink_{name}", **{name: max(8, current_value // 2)}) + + +def run_verify_diff( + *, + buf_pair: tuple[torch.Tensor, torch.Tensor], + plan_pair: tuple[VerifyPlan, VerifyPlan], + kernel_kind: CanaryLaunchTag = CanaryLaunchTag.HEAD_K_FULL, + device: torch.device = _DEVICE, + assert_equal: bool = True, + check_verify_expected_token: bool = True, +) -> tuple[FakeViolationLog, FakeViolationLog]: + """Thin wrapper around ``_run_both_verify`` that creates a fresh log pair and packs (cuda, ref) + buf/plan arguments into 2-tuples to drop ~8 lines of boilerplate per call site. + """ + cuda_log, ref_log = make_log_pair(device=device) + _run_both_verify( + cuda_canary_buf=buf_pair[0], + ref_canary_buf=buf_pair[1], + plan_cuda=plan_pair[0], + plan_ref=plan_pair[1], + cuda_log=cuda_log, + ref_log=ref_log, + kernel_kind=kernel_kind, + assert_equal=assert_equal, + check_verify_expected_token=check_verify_expected_token, + ) + return cuda_log, ref_log + + +def run_write_diff( + *, + buf_pair: tuple[torch.Tensor, torch.Tensor], + plan_pair: tuple[WritePlan, WritePlan], + input_ids: torch.Tensor, + positions: torch.Tensor, + out_cache_loc: torch.Tensor, + expected_input_tokens: torch.Tensor, + expected_input_positions: torch.Tensor, + enable_write_verify_inputs: bool = False, + kernel_kind: CanaryLaunchTag = CanaryLaunchTag.HEAD_K_FULL, + device: torch.device = _DEVICE, + assert_equal: bool = True, +) -> tuple[FakeViolationLog, FakeViolationLog]: + """Thin wrapper around ``_run_both_write`` that creates a fresh log pair and packs (cuda, ref) + buf/plan arguments into 2-tuples to drop ~10 lines of boilerplate per call site. + """ + cuda_log, ref_log = make_log_pair(device=device) + _run_both_write( + cuda_canary_buf=buf_pair[0], + ref_canary_buf=buf_pair[1], + plan_cuda=plan_pair[0], + plan_ref=plan_pair[1], + input_ids=input_ids, + positions=positions, + out_cache_loc=out_cache_loc, + enable_write_verify_inputs=enable_write_verify_inputs, + expected_input_tokens=expected_input_tokens, + expected_input_positions=expected_input_positions, + cuda_log=cuda_log, + ref_log=ref_log, + kernel_kind=kernel_kind, + assert_equal=assert_equal, + ) + return cuda_log, ref_log + + +def run_plan_diff( + *, + plan_pair: tuple[tuple[VerifyPlan, WritePlan], tuple[VerifyPlan, WritePlan]], + req_pool_indices: torch.Tensor, + prefix_lens: torch.Tensor, + extend_seq_lens: torch.Tensor, + req_to_token: torch.Tensor, + extras: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], + swa_window_size: int = 0, + full_to_swa_index_mapping: Optional[torch.Tensor] = None, + assert_equal: bool = True, + active_verify_entries: Optional[int] = None, + active_write_reqs: Optional[int] = None, + req_to_verify_expected_tokens: Optional[torch.Tensor] = None, + req_to_verify_expected_tokens_valid_lens: Optional[torch.Tensor] = None, + kv_token_id_vs_position_offset: int = 0, +) -> None: + """Thin wrapper around ``_run_both_plan`` that unpacks ``((triton_v, triton_w), (ref_v, ref_w))`` + plan pairs to drop the per-call-site ``triton_verify=.../triton_write=.../ref_verify=...`` block. + """ + (triton_verify, triton_write), (ref_verify, ref_write) = plan_pair + _run_both_plan( + triton_verify=triton_verify, + triton_write=triton_write, + ref_verify=ref_verify, + ref_write=ref_write, + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + extend_seq_lens=extend_seq_lens, + req_to_token=req_to_token, + extras=extras, + swa_window_size=swa_window_size, + full_to_swa_index_mapping=full_to_swa_index_mapping, + assert_equal=assert_equal, + active_verify_entries=active_verify_entries, + active_write_reqs=active_write_reqs, + req_to_verify_expected_tokens=req_to_verify_expected_tokens, + req_to_verify_expected_tokens_valid_lens=req_to_verify_expected_tokens_valid_lens, + kv_token_id_vs_position_offset=kv_token_id_vs_position_offset, + ) diff --git a/python/sglang/jit_kernel/tests/kv_canary/_fuzz_driver.py b/python/sglang/jit_kernel/tests/kv_canary/_fuzz_driver.py new file mode 100644 index 000000000..6af7c305b --- /dev/null +++ b/python/sglang/jit_kernel/tests/kv_canary/_fuzz_driver.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +import random +from typing import Any, Callable + +from sglang.jit_kernel.tests.kv_canary._differential import ( + ShrinkResult, + shrink_inputs, +) + +FUZZ_SEEDS_PR: tuple[int, ...] = (0,) + + +def check_repro(inputs: Any, *, run_one_fn: Callable[[Any], Any]) -> bool: + try: + run_one_fn(inputs) + except (AssertionError, RuntimeError, ValueError): + return True + return False + + +def run_fuzz_combo( + seed: int, + *, + draw_fn: Callable[[random.Random], Any], + run_one_fn: Callable[[Any], Any], + summarize_fn: Callable[[Any], str], + n_iter: int, +) -> None: + rng = random.Random(seed) + for iteration in range(n_iter): + inputs = draw_fn(rng) + try: + run_one_fn(inputs) + except AssertionError as exc: + shrunk: ShrinkResult = shrink_inputs( + inputs, + check_fn=lambda i: check_repro(i, run_one_fn=run_one_fn), + ) + raise AssertionError( + f"seed={seed} iter={iteration} failure: {exc}\n" + f"original: {summarize_fn(inputs)}\n" + f"shrunk: {summarize_fn(shrunk.inputs)}\n" + f"mutations applied: {shrunk.mutations_applied}" + ) from exc diff --git a/python/sglang/jit_kernel/tests/kv_canary/test_kernel_config.py b/python/sglang/jit_kernel/tests/kv_canary/test_kernel_config.py new file mode 100644 index 000000000..df2c1021a --- /dev/null +++ b/python/sglang/jit_kernel/tests/kv_canary/test_kernel_config.py @@ -0,0 +1,425 @@ +from __future__ import annotations + +import pytest +import torch + +from sglang.jit_kernel.kv_canary.verify import ( + CanaryLaunchTag, + VerifyOrWriteContext, + VerifyPlan, + launch_canary_verify_kernel, +) +from sglang.jit_kernel.kv_canary.write import WritePlan +from sglang.jit_kernel.tests.kv_canary._canary_helpers import ( + FakeViolationLog, + assert_canary_buf_equal, + assert_canary_state_equal, + make_canary_buf, + make_canary_buf_pair, + make_log_pair, + make_verify_plan, + make_verify_plan_pair, + make_write_plan_pair, + stamp_clean_chain, +) +from sglang.jit_kernel.tests.kv_canary._differential import ( + _assert_plans_byte_equal, + _run_both_plan, + _run_both_verify, + _run_both_write, +) +from sglang.jit_kernel.tests.kv_canary._fixtures import ( + dummy_pseudo_tensors, + empty_extras, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=60, suite="base-b-kernel-unit-1-gpu-large") + +_DEVICE = torch.device("cuda") + + +def _build_verify_plan_5_entries( + *, device: torch.device +) -> tuple[VerifyPlan, VerifyPlan]: + num_slots = 16 + cuda_buf, ref_buf = make_canary_buf_pair( + num_slots=num_slots, slot_stride_bytes=32, device=_DEVICE + ) + + plan_cuda, plan_ref = make_verify_plan_pair( + slot_indices=[0, 1, 2, 3, 4], + positions=[0, 1, 2, 3, 4], + prev_slot_indices=[-1, 0, 1, 2, 3], + capacity=8, + device=device, + ) + return plan_cuda, plan_ref + + +def _build_write_fixtures( + *, device: torch.device +) -> tuple[WritePlan, WritePlan, torch.Tensor, torch.Tensor, torch.Tensor]: + num_tokens = 5 + plan_cuda, plan_ref = make_write_plan_pair( + write_offsets=[0, num_tokens], + seed_slot_indices=[-1], + num_valid_reqs=1, + req_capacity=4, + device=device, + ) + input_ids = torch.tensor([10, 20, 30, 40, 50], dtype=torch.int64, device=device) + positions = torch.tensor([0, 1, 2, 3, 4], dtype=torch.int64, device=device) + out_cache_loc = torch.tensor([0, 1, 2, 3, 4], dtype=torch.int64, device=device) + return plan_cuda, plan_ref, input_ids, positions, out_cache_loc + + +def _build_plan_fixtures( + *, device: torch.device, int64_req_to_token: bool = False +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + bs = 3 + max_reqs = 4 + max_seq_len = 16 + req_pool_indices = torch.tensor([1, 2, 3], dtype=torch.int64, device=device) + prefix_lens = torch.tensor([0, 4, 8], dtype=torch.int64, device=device) + extend_seq_lens = torch.tensor([5, 1, 1], dtype=torch.int64, device=device) + rp_axis = torch.arange(max_reqs, device=device, dtype=torch.int32).unsqueeze(1) + pos_axis = torch.arange(max_seq_len, device=device, dtype=torch.int32).unsqueeze(0) + req_to_token_int32 = (rp_axis * max_seq_len + pos_axis).contiguous() + if int64_req_to_token: + req_to_token = req_to_token_int32.to(torch.int64) + else: + req_to_token = req_to_token_int32 + return req_pool_indices, prefix_lens, extend_seq_lens, req_to_token + + +def test_verify_byte_equal_across_repeated_launches_10x() -> None: + num_launches = 10 + plan_cuda, plan_ref = _build_verify_plan_5_entries(device=_DEVICE) + + snapshot_rings: list[torch.Tensor] = [] + snapshot_write_indices: list[torch.Tensor] = [] + snapshot_bufs: list[torch.Tensor] = [] + + for _ in range(num_launches): + cuda_buf, ref_buf = make_canary_buf_pair( + num_slots=16, slot_stride_bytes=32, device=_DEVICE + ) + cuda_log, ref_log = make_log_pair(capacity=64, device=_DEVICE) + + _run_both_verify( + cuda_canary_buf=cuda_buf, + ref_canary_buf=ref_buf, + plan_cuda=plan_cuda, + plan_ref=plan_ref, + cuda_log=cuda_log, + ref_log=ref_log, + kernel_kind=CanaryLaunchTag.HEAD_K_FULL, + ) + + assert_canary_buf_equal(buf_a=cuda_buf, buf_b=ref_buf) + assert_canary_state_equal(log_a=cuda_log, log_b=ref_log) + + snapshot_rings.append(cuda_log.ring.clone()) + snapshot_write_indices.append(cuda_log.write_index.clone()) + snapshot_bufs.append(cuda_buf.clone()) + + for i in range(1, num_launches): + assert torch.equal( + snapshot_rings[0], snapshot_rings[i] + ), f"violation_ring differs between launch 0 and {i}" + assert torch.equal( + snapshot_write_indices[0], snapshot_write_indices[i] + ), f"violation_write_index differs between launch 0 and {i}" + assert torch.equal( + snapshot_bufs[0], snapshot_bufs[i] + ), f"canary_buf differs between launch 0 and {i}" + + +def test_write_byte_equal_across_repeated_launches_10x() -> None: + num_launches = 10 + plan_cuda, plan_ref, input_ids, positions, out_cache_loc = _build_write_fixtures( + device=_DEVICE + ) + + snapshot_bufs: list[torch.Tensor] = [] + snapshot_rings: list[torch.Tensor] = [] + snapshot_counters: list[torch.Tensor] = [] + + for _ in range(num_launches): + cuda_buf, ref_buf = make_canary_buf_pair( + num_slots=16, slot_stride_bytes=32, device=_DEVICE + ) + cuda_log, ref_log = make_log_pair(capacity=64, device=_DEVICE) + pseudo_tok, pseudo_pos = dummy_pseudo_tensors(input_ids.shape[0]) + + _run_both_write( + cuda_canary_buf=cuda_buf, + ref_canary_buf=ref_buf, + plan_cuda=plan_cuda, + plan_ref=plan_ref, + input_ids=input_ids, + positions=positions, + out_cache_loc=out_cache_loc, + enable_write_verify_inputs=False, + expected_input_tokens=pseudo_tok, + expected_input_positions=pseudo_pos, + cuda_log=cuda_log, + ref_log=ref_log, + kernel_kind=CanaryLaunchTag.HEAD_K_FULL, + ) + + assert_canary_buf_equal(buf_a=cuda_buf, buf_b=ref_buf) + assert_canary_state_equal(log_a=cuda_log, log_b=ref_log) + + snapshot_bufs.append(cuda_buf.clone()) + snapshot_rings.append(cuda_log.ring.clone()) + snapshot_counters.append(cuda_log.slot_run_counter.clone()) + + for i in range(1, num_launches): + assert torch.equal( + snapshot_bufs[0], snapshot_bufs[i] + ), f"canary_buf differs between launch 0 and {i}" + assert torch.equal( + snapshot_rings[0], snapshot_rings[i] + ), f"violation_ring differs between launch 0 and {i}" + assert torch.equal( + snapshot_counters[0], snapshot_counters[i] + ), f"slot_run_counter differs between launch 0 and {i}" + + +def test_plan_byte_equal_across_repeated_launches_10x() -> None: + num_launches = 10 + req_pool_indices, prefix_lens, extend_seq_lens, req_to_token = _build_plan_fixtures( + device=_DEVICE + ) + + snapshot_slots: list[torch.Tensor] = [] + snapshot_positions: list[torch.Tensor] = [] + snapshot_prevs: list[torch.Tensor] = [] + snapshot_write_offsets: list[torch.Tensor] = [] + + for _ in range(num_launches): + triton_v = VerifyPlan.allocate( + verify_capacity=64, device=_DEVICE + ).zero_for_testing_() + triton_w = WritePlan.allocate( + write_req_capacity=8, device=_DEVICE + ).zero_for_testing_() + ref_v = VerifyPlan.allocate( + verify_capacity=64, device=_DEVICE + ).zero_for_testing_() + ref_w = WritePlan.allocate( + write_req_capacity=8, device=_DEVICE + ).zero_for_testing_() + + _run_both_plan( + triton_verify=triton_v, + triton_write=triton_w, + ref_verify=ref_v, + ref_write=ref_w, + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + extend_seq_lens=extend_seq_lens, + req_to_token=req_to_token, + extras=empty_extras(), + swa_window_size=0, + full_to_swa_index_mapping=None, + ) + + _assert_plans_byte_equal( + triton_verify=triton_v, + triton_write=triton_w, + ref_verify=ref_v, + ref_write=ref_w, + ) + + n_verify = int(triton_v.verify_num_valid[0].item()) + snapshot_slots.append(triton_v.verify_slot_indices[:n_verify].clone()) + snapshot_positions.append(triton_v.verify_expected_positions[:n_verify].clone()) + snapshot_prevs.append(triton_v.verify_prev_slot_indices[:n_verify].clone()) + snapshot_write_offsets.append(triton_w.write_offsets.clone()) + + for i in range(1, num_launches): + assert torch.equal( + snapshot_slots[0], snapshot_slots[i] + ), f"verify_slot_indices differs between launch 0 and {i}" + assert torch.equal( + snapshot_positions[0], snapshot_positions[i] + ), f"verify_expected_positions differs between launch 0 and {i}" + assert torch.equal( + snapshot_prevs[0], snapshot_prevs[i] + ), f"verify_prev_slot_indices differs between launch 0 and {i}" + assert torch.equal( + snapshot_write_offsets[0], snapshot_write_offsets[i] + ), f"write_offsets differs between launch 0 and {i}" + + +def test_verify_multi_launch_100x_counter_linear() -> None: + num_launches = 100 + + plan_cuda = make_verify_plan( + slot_indices=[0], + positions=[0], + prev_slot_indices=[-1], + capacity=4, + device=_DEVICE, + ) + + cuda_log = FakeViolationLog.allocate(capacity=64, device=_DEVICE) + + for _ in range(num_launches): + cuda_buf = make_canary_buf(num_slots=16, slot_stride_bytes=32, device=_DEVICE) + launch_canary_verify_kernel( + context=VerifyOrWriteContext( + canary_buf=cuda_buf, + kernel_kind=CanaryLaunchTag.HEAD_K_FULL, + violation_ring=cuda_log.ring, + violation_write_index=cuda_log.write_index, + slot_run_counter=cuda_log.slot_run_counter, + kernel_run_counter=cuda_log.kernel_run_counter, + enable_chain_position_assert=cuda_log.enable_chain_position_assert, + ), + plan=plan_cuda, + check_verify_expected_token=True, + ) + + torch.cuda.synchronize() + + assert ( + int(cuda_log.kernel_run_counter[0].item()) == num_launches + ), f"kernel_run_counter expected {num_launches}, got {cuda_log.kernel_run_counter[0].item()}" + assert int(cuda_log.slot_run_counter[0].item()) == num_launches, ( + f"slot_run_counter expected {num_launches} (1 active entry x 100 launches), " + f"got {cuda_log.slot_run_counter[0].item()}" + ) + + +def test_verify_check_disabled_byte_equal() -> None: + """check_verify_expected_token True vs False produce equivalent violation logs on a clean plan.""" + plan_true_cuda, plan_true_ref = _build_verify_plan_5_entries(device=_DEVICE) + plan_false_cuda, plan_false_ref = _build_verify_plan_5_entries(device=_DEVICE) + + chain_slot_indices = [0, 1, 2, 3, 4] + chain_tokens = [10, 20, 30, 40, 50] + chain_positions = [0, 1, 2, 3, 4] + + cuda_buf_true, ref_buf_true = make_canary_buf_pair( + num_slots=16, slot_stride_bytes=32, device=_DEVICE + ) + stamp_clean_chain( + cuda_buf=cuda_buf_true, + ref_buf=ref_buf_true, + slot_indices=chain_slot_indices, + tokens=chain_tokens, + positions=chain_positions, + ) + cuda_buf_false, ref_buf_false = make_canary_buf_pair( + num_slots=16, slot_stride_bytes=32, device=_DEVICE + ) + stamp_clean_chain( + cuda_buf=cuda_buf_false, + ref_buf=ref_buf_false, + slot_indices=chain_slot_indices, + tokens=chain_tokens, + positions=chain_positions, + ) + cuda_log_true, ref_log_true = make_log_pair(capacity=64, device=_DEVICE) + cuda_log_false, ref_log_false = make_log_pair(capacity=64, device=_DEVICE) + + _run_both_verify( + cuda_canary_buf=cuda_buf_true, + ref_canary_buf=ref_buf_true, + plan_cuda=plan_true_cuda, + plan_ref=plan_true_ref, + cuda_log=cuda_log_true, + ref_log=ref_log_true, + kernel_kind=CanaryLaunchTag.HEAD_K_FULL, + check_verify_expected_token=True, + ) + _run_both_verify( + cuda_canary_buf=cuda_buf_false, + ref_canary_buf=ref_buf_false, + plan_cuda=plan_false_cuda, + plan_ref=plan_false_ref, + cuda_log=cuda_log_false, + ref_log=ref_log_false, + kernel_kind=CanaryLaunchTag.HEAD_K_FULL, + check_verify_expected_token=False, + ) + + assert int(cuda_log_true.write_index[0].item()) == 0 + assert int(cuda_log_false.write_index[0].item()) == 0 + assert torch.equal(cuda_log_true.ring, cuda_log_false.ring) + assert torch.equal(cuda_log_true.write_index, cuda_log_false.write_index) + assert torch.equal(cuda_log_true.slot_run_counter, cuda_log_false.slot_run_counter) + assert torch.equal( + cuda_log_true.kernel_run_counter, cuda_log_false.kernel_run_counter + ) + + +@pytest.mark.parametrize("per_req_present", [False, True]) +def test_plan_per_req_present_or_absent(per_req_present: bool) -> None: + max_reqs = 4 + max_seq_len = 16 + rp_axis = torch.arange(max_reqs, device=_DEVICE, dtype=torch.int32).unsqueeze(1) + pos_axis = torch.arange(max_seq_len, device=_DEVICE, dtype=torch.int32).unsqueeze(0) + req_to_token = (rp_axis * max_seq_len + pos_axis).contiguous() + + if per_req_present: + req_pool_indices = torch.tensor([1, 2], dtype=torch.int64, device=_DEVICE) + prefix_lens = torch.tensor([3, 5], dtype=torch.int64, device=_DEVICE) + extend_seq_lens = torch.tensor([1, 1], dtype=torch.int64, device=_DEVICE) + else: + req_pool_indices = torch.tensor([0], dtype=torch.int64, device=_DEVICE) + prefix_lens = torch.tensor([0], dtype=torch.int64, device=_DEVICE) + extend_seq_lens = torch.tensor([0], dtype=torch.int64, device=_DEVICE) + + triton_v = VerifyPlan.allocate( + verify_capacity=64, device=_DEVICE + ).zero_for_testing_() + triton_w = WritePlan.allocate( + write_req_capacity=8, device=_DEVICE + ).zero_for_testing_() + ref_v = VerifyPlan.allocate(verify_capacity=64, device=_DEVICE).zero_for_testing_() + ref_w = WritePlan.allocate(write_req_capacity=8, device=_DEVICE).zero_for_testing_() + + _run_both_plan( + triton_verify=triton_v, + triton_write=triton_w, + ref_verify=ref_v, + ref_write=ref_w, + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + extend_seq_lens=extend_seq_lens, + req_to_token=req_to_token, + extras=empty_extras(), + swa_window_size=0, + full_to_swa_index_mapping=None, + ) + + _assert_plans_byte_equal( + triton_verify=triton_v, + triton_write=triton_w, + ref_verify=ref_v, + ref_write=ref_w, + ) + + if not per_req_present: + assert int(triton_v.verify_num_valid[0].item()) == 0 + bs = int(req_pool_indices.shape[0]) + assert int(triton_w.write_offsets[bs].item()) == 0 + + if per_req_present: + assert int(triton_v.verify_num_valid[0].item()) == 8 + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) diff --git a/python/sglang/jit_kernel/tests/kv_canary/test_pipeline_e2e.py b/python/sglang/jit_kernel/tests/kv_canary/test_pipeline_e2e.py new file mode 100644 index 000000000..5590997c7 --- /dev/null +++ b/python/sglang/jit_kernel/tests/kv_canary/test_pipeline_e2e.py @@ -0,0 +1,790 @@ +from __future__ import annotations + +from typing import Any, Optional + +import pytest +import torch + +from sglang.jit_kernel.kv_canary import consts +from sglang.jit_kernel.kv_canary.plan import launch_canary_plan_kernels +from sglang.jit_kernel.kv_canary.plan_ref import ( + launch_canary_plan_kernels_torch_reference, +) +from sglang.jit_kernel.kv_canary.verify import ( + CanaryLaunchTag, + VerifyOrWriteContext, + VerifyPlan, + launch_canary_verify_kernel, +) +from sglang.jit_kernel.kv_canary.verify_ref import ( + launch_canary_verify_kernel_torch_reference, +) +from sglang.jit_kernel.kv_canary.write import WritePlan, launch_canary_write_kernel +from sglang.jit_kernel.kv_canary.write_ref import ( + launch_canary_write_kernel_torch_reference, +) +from sglang.jit_kernel.tests.kv_canary._canary_helpers import ( + FakeViolationLog, + assert_canary_buf_equal, + assert_canary_state_equal, + make_canary_buf, + stamp_clean_chain, + write_slot_fields, +) +from sglang.jit_kernel.tests.kv_canary._fixtures import ( + empty_extras, + make_req_to_token, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large") + + +_DEVICE = torch.device("cuda") + + +def _run_pipeline( + *, + real: bool, + req_pool_indices: torch.Tensor, + prefix_lens: torch.Tensor, + extend_seq_lens: torch.Tensor, + input_ids: torch.Tensor, + positions: torch.Tensor, + out_cache_loc: torch.Tensor, + req_to_token: torch.Tensor, + canary_buf: torch.Tensor, + log: FakeViolationLog, + extras: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], + swa_window_size: int, + full_to_swa_index_mapping: Optional[torch.Tensor], + kernel_kind: CanaryLaunchTag, + enable_write_verify_inputs: bool, + expected_input_tokens: torch.Tensor, + expected_input_positions: torch.Tensor, + verify_capacity: int, + write_req_capacity: int, + req_to_verify_expected_tokens: Optional[torch.Tensor] = None, + req_to_verify_expected_tokens_valid_lens: Optional[torch.Tensor] = None, + kv_token_id_vs_position_offset: int = 0, + check_verify_expected_token: bool = True, +) -> tuple[VerifyPlan, WritePlan]: + _ = extras + plan_v = VerifyPlan.allocate(verify_capacity=verify_capacity, device=_DEVICE) + plan_w = WritePlan.allocate(write_req_capacity=write_req_capacity, device=_DEVICE) + + # Existing pipeline tests that supply a pool but no per-req lens want "bound by full + # row width" semantics. Synthesise that bound here so callers don't have to. + if ( + req_to_verify_expected_tokens is not None + and req_to_verify_expected_tokens_valid_lens is None + ): + req_to_verify_expected_tokens_valid_lens = torch.full( + (int(req_pool_indices.shape[0]),), + int(req_to_verify_expected_tokens.shape[1]), + dtype=torch.int64, + device=req_pool_indices.device, + ) + + plan_fn = ( + launch_canary_plan_kernels + if real + else launch_canary_plan_kernels_torch_reference + ) + plan_fn( + verify_plan_out=plan_v, + write_plan_out=plan_w, + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + extend_seq_lens=extend_seq_lens, + req_to_token=req_to_token, + swa_window_size=swa_window_size, + full_to_swa_index_mapping=full_to_swa_index_mapping, + verify_capacity=verify_capacity, + req_to_verify_expected_tokens=req_to_verify_expected_tokens, + req_to_verify_expected_tokens_valid_lens=req_to_verify_expected_tokens_valid_lens, + kv_token_id_vs_position_offset=kv_token_id_vs_position_offset, + ) + + if real: + context = VerifyOrWriteContext( + canary_buf=canary_buf, + kernel_kind=kernel_kind, + violation_ring=log.ring, + violation_write_index=log.write_index, + slot_run_counter=log.slot_run_counter, + kernel_run_counter=log.kernel_run_counter, + enable_chain_position_assert=log.enable_chain_position_assert, + ) + launch_canary_write_kernel( + context=context, + plan=plan_w, + input_ids=input_ids, + positions=positions, + out_cache_loc=out_cache_loc, + enable_write_input_assert=enable_write_verify_inputs, + expected_input_tokens=expected_input_tokens, + expected_input_positions=expected_input_positions, + ) + launch_canary_verify_kernel( + context=context, + plan=plan_v, + check_verify_expected_token=check_verify_expected_token, + ) + torch.cuda.synchronize() + else: + launch_canary_write_kernel_torch_reference( + context=VerifyOrWriteContext( + canary_buf=canary_buf, + kernel_kind=kernel_kind, + violation_ring=log.ring, + violation_write_index=log.write_index, + slot_run_counter=log.slot_run_counter, + kernel_run_counter=log.kernel_run_counter, + enable_chain_position_assert=log.enable_chain_position_assert, + ), + plan=plan_w, + input_ids=input_ids, + positions=positions, + out_cache_loc=out_cache_loc, + enable_write_input_assert=enable_write_verify_inputs, + expected_input_tokens=expected_input_tokens, + expected_input_positions=expected_input_positions, + ) + launch_canary_verify_kernel_torch_reference( + context=VerifyOrWriteContext( + canary_buf=canary_buf, + kernel_kind=kernel_kind, + violation_ring=log.ring, + violation_write_index=log.write_index, + slot_run_counter=log.slot_run_counter, + kernel_run_counter=log.kernel_run_counter, + enable_chain_position_assert=log.enable_chain_position_assert, + ), + plan=plan_v, + check_verify_expected_token=check_verify_expected_token, + ) + + return plan_v, plan_w + + +def _run_both_and_assert_pipeline_equal( + *, + req_pool_indices: torch.Tensor, + prefix_lens: torch.Tensor, + extend_seq_lens: torch.Tensor, + input_ids: torch.Tensor, + positions: torch.Tensor, + out_cache_loc: torch.Tensor, + req_to_token: torch.Tensor, + num_slots: int, + extras: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], + swa_window_size: int = 0, + full_to_swa_index_mapping: Optional[torch.Tensor] = None, + kernel_kind: CanaryLaunchTag = CanaryLaunchTag.HEAD_K_FULL, + enable_write_verify_inputs: bool = False, + expected_input_tokens: Optional[torch.Tensor] = None, + expected_input_positions: Optional[torch.Tensor] = None, + ring_capacity: int = 64, + verify_capacity: int = 256, + write_req_capacity: int = 16, + assert_ring_equal: bool = True, + initial_canary_buf: Optional[torch.Tensor] = None, + req_to_verify_expected_tokens: Optional[torch.Tensor] = None, + kv_token_id_vs_position_offset: int = 0, + check_verify_expected_token: bool = True, +) -> tuple[ + torch.Tensor, + torch.Tensor, + FakeViolationLog, + FakeViolationLog, + VerifyPlan, + WritePlan, + VerifyPlan, + WritePlan, +]: + # The kernel rejects non-None expected_* tensors when enable_write_verify_inputs=False + # (sanity check to catch caller bugs), so only synthesise zero placeholders in the + # branch that will actually assert against them. + if enable_write_verify_inputs: + total_tokens = int(input_ids.shape[0]) + if expected_input_tokens is None: + expected_input_tokens = torch.zeros( + total_tokens, dtype=torch.int64, device=_DEVICE + ) + if expected_input_positions is None: + expected_input_positions = torch.zeros( + total_tokens, dtype=torch.int64, device=_DEVICE + ) + + if initial_canary_buf is None: + buf_real = make_canary_buf(num_slots=num_slots, device=_DEVICE) + else: + buf_real = initial_canary_buf.clone() + buf_ref = buf_real.clone() + log_real = FakeViolationLog.allocate(capacity=ring_capacity, device=_DEVICE) + log_ref = FakeViolationLog.allocate(capacity=ring_capacity, device=_DEVICE) + + shared: dict[str, Any] = dict( + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + extend_seq_lens=extend_seq_lens, + input_ids=input_ids, + positions=positions, + out_cache_loc=out_cache_loc, + req_to_token=req_to_token, + extras=extras, + swa_window_size=swa_window_size, + full_to_swa_index_mapping=full_to_swa_index_mapping, + kernel_kind=kernel_kind, + enable_write_verify_inputs=enable_write_verify_inputs, + expected_input_tokens=expected_input_tokens, + expected_input_positions=expected_input_positions, + verify_capacity=verify_capacity, + write_req_capacity=write_req_capacity, + req_to_verify_expected_tokens=req_to_verify_expected_tokens, + kv_token_id_vs_position_offset=kv_token_id_vs_position_offset, + check_verify_expected_token=check_verify_expected_token, + ) + + plan_v_real, plan_w_real = _run_pipeline( + real=True, + canary_buf=buf_real, + log=log_real, + **shared, + ) + plan_v_ref, plan_w_ref = _run_pipeline( + real=False, + canary_buf=buf_ref, + log=log_ref, + **shared, + ) + + assert_canary_buf_equal(buf_a=buf_real, buf_b=buf_ref) + if assert_ring_equal: + assert_canary_state_equal(log_a=log_real, log_b=log_ref) + else: + assert torch.equal(log_real.write_index, log_ref.write_index) + assert torch.equal(log_real.slot_run_counter, log_ref.slot_run_counter) + assert torch.equal(log_real.kernel_run_counter, log_ref.kernel_run_counter) + + return ( + buf_real, + buf_ref, + log_real, + log_ref, + plan_v_real, + plan_w_real, + plan_v_ref, + plan_w_ref, + ) + + +def _t(values: list[int]) -> torch.Tensor: + return torch.tensor(values, dtype=torch.int64, device=_DEVICE) + + +def _linear_r2t(*, max_reqs: int = 4, max_seq_len: int = 16) -> torch.Tensor: + return make_req_to_token( + kind="linear", max_reqs=max_reqs, max_seq_len=max_seq_len, device=_DEVICE + ) + + +def _zero_no_write_inputs() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """``(input_ids, positions, out_cache_loc)`` zero placeholders for extend_seq_lens=0 tests.""" + zeros = torch.zeros(1, dtype=torch.int64, device=_DEVICE) + return zeros.clone(), zeros.clone(), zeros.clone() + + +def _contiguous_out_cache_loc( + *, req_pool_idx: int, start: int, count: int, max_seq_len: int = 16 +) -> torch.Tensor: + return _t([req_pool_idx * max_seq_len + start + i for i in range(count)]) + + +def _stamp_linear_prefix( + *, + initial_buf: torch.Tensor, + initial_ref: torch.Tensor, + req_pool_idx: int, + prefix_len: int, + tokens: list[int], + max_seq_len: int = 16, +) -> None: + """Stamp clean chain for slots ``[rp*max_seq_len + 0 .. + prefix_len)`` at positions ``0..prefix_len``.""" + stamp_clean_chain( + cuda_buf=initial_buf, + ref_buf=initial_ref, + slot_indices=[req_pool_idx * max_seq_len + pos for pos in range(prefix_len)], + tokens=tokens, + positions=list(range(prefix_len)), + ) + + +def test_pipeline_basic_5_step_single_req() -> None: + """Single req, prefix_len=0, extend_seq_len=5: basic plan→write→verify byte-equal.""" + _run_both_and_assert_pipeline_equal( + req_pool_indices=_t([1]), + prefix_lens=_t([0]), + extend_seq_lens=_t([5]), + input_ids=_t([10, 20, 30, 40, 50]), + positions=_t([0, 1, 2, 3, 4]), + out_cache_loc=_contiguous_out_cache_loc(req_pool_idx=1, start=0, count=5), + req_to_token=_linear_r2t(), + num_slots=64, + extras=empty_extras(), + swa_window_size=0, + full_to_swa_index_mapping=None, + ) + + +def test_pipeline_multi_req_mixed_extend_decode() -> None: + """bs=3: pure extend req, decode req (prefix+1 extend), and padding sentinel row.""" + max_seq_len = 16 + _run_both_and_assert_pipeline_equal( + req_pool_indices=_t([1, 2, 0]), + prefix_lens=_t([0, 5, 0]), + extend_seq_lens=_t([4, 1, 0]), + input_ids=_t([11, 12, 13, 14, 21]), + positions=_t([0, 1, 2, 3, 5]), + out_cache_loc=_t( + [ + 1 * max_seq_len + 0, + 1 * max_seq_len + 1, + 1 * max_seq_len + 2, + 1 * max_seq_len + 3, + 2 * max_seq_len + 5, + ] + ), + req_to_token=_linear_r2t(max_reqs=8, max_seq_len=max_seq_len), + num_slots=128, + extras=empty_extras(), + swa_window_size=0, + full_to_swa_index_mapping=None, + write_req_capacity=4, + ) + + +def test_pipeline_swa_window() -> None: + """SWA window=4, prefix_len=6: verify covers window [2,6), write covers extend tokens.""" + max_seq_len = 16 + max_reqs = 4 + full_to_swa_index_mapping = torch.arange( + max_reqs * max_seq_len + 1, dtype=torch.int64, device=_DEVICE + ) + + _run_both_and_assert_pipeline_equal( + req_pool_indices=_t([1]), + prefix_lens=_t([6]), + extend_seq_lens=_t([2]), + input_ids=_t([100, 101]), + positions=_t([6, 7]), + out_cache_loc=_t( + [ + full_to_swa_index_mapping[1 * max_seq_len + 6].item(), + full_to_swa_index_mapping[1 * max_seq_len + 7].item(), + ] + ), + req_to_token=_linear_r2t(max_reqs=max_reqs, max_seq_len=max_seq_len), + num_slots=64, + extras=empty_extras(), + swa_window_size=4, + full_to_swa_index_mapping=full_to_swa_index_mapping, + ) + + +def test_pipeline_sweep_no_write() -> None: + """All extend_seq_lens=0: write_step is no-op, verify sweeps prefix, buf unchanged.""" + prefix_len = 4 + input_ids, positions, out_cache_loc = _zero_no_write_inputs() + + initial_buf = make_canary_buf(num_slots=64, device=_DEVICE) + initial_ref = initial_buf.clone() + _stamp_linear_prefix( + initial_buf=initial_buf, + initial_ref=initial_ref, + req_pool_idx=1, + prefix_len=prefix_len, + tokens=[100 + pos for pos in range(prefix_len)], + ) + + buf_real, buf_ref, log_real, log_ref, plan_v_real, plan_w_real, _, _ = ( + _run_both_and_assert_pipeline_equal( + req_pool_indices=_t([1]), + prefix_lens=_t([prefix_len]), + extend_seq_lens=_t([0]), + input_ids=input_ids, + positions=positions, + out_cache_loc=out_cache_loc, + req_to_token=_linear_r2t(), + num_slots=64, + extras=empty_extras(), + swa_window_size=0, + full_to_swa_index_mapping=None, + initial_canary_buf=initial_buf, + ) + ) + + assert int(plan_v_real.verify_num_valid[0].item()) == prefix_len + assert int(plan_w_real.write_num_valid_reqs[0].item()) == 1 + assert int(plan_w_real.write_offsets[1].item()) == 0 + assert torch.equal(buf_real, initial_buf) + assert torch.equal(buf_ref, initial_buf) + assert int(log_real.write_index[0].item()) == 0 + assert int(log_ref.write_index[0].item()) == 0 + assert int(log_real.slot_run_counter[0].item()) == prefix_len + assert int(log_ref.slot_run_counter[0].item()) == prefix_len + + +def test_pipeline_pseudo_mode_on_match() -> None: + """enable_write_verify_inputs=ON, expected==actual: zero violations, buf byte-equal.""" + input_ids = _t([1, 2, 3, 4]) + positions = _t([0, 1, 2, 3]) + + _, _, log_real, log_ref, _, _, _, _ = _run_both_and_assert_pipeline_equal( + req_pool_indices=_t([1]), + prefix_lens=_t([0]), + extend_seq_lens=_t([4]), + input_ids=input_ids, + positions=positions, + out_cache_loc=_contiguous_out_cache_loc(req_pool_idx=1, start=0, count=4), + req_to_token=_linear_r2t(), + num_slots=64, + extras=empty_extras(), + enable_write_verify_inputs=True, + expected_input_tokens=input_ids.clone(), + expected_input_positions=positions.clone(), + ) + + assert int(log_real.write_index[0].item()) == 0 + assert int(log_ref.write_index[0].item()) == 0 + + +def test_pipeline_pseudo_mode_on_token_mismatch_then_verify_clean() -> None: + """enable_write_verify_inputs=ON, expected tokens all wrong: write records N violations.""" + n_tokens = 3 + positions = _t([0, 1, 2]) + + _, _, log_real, log_ref, _, _, _, _ = _run_both_and_assert_pipeline_equal( + req_pool_indices=_t([1]), + prefix_lens=_t([0]), + extend_seq_lens=_t([3]), + input_ids=_t([10, 20, 30]), + positions=positions, + out_cache_loc=_contiguous_out_cache_loc(req_pool_idx=1, start=0, count=3), + req_to_token=_linear_r2t(), + num_slots=64, + extras=empty_extras(), + enable_write_verify_inputs=True, + expected_input_tokens=_t([99, 99, 99]), + expected_input_positions=positions.clone(), + ring_capacity=64, + ) + + write_violations = int(log_real.write_index[0].item()) + assert ( + write_violations == n_tokens + ), f"expected {n_tokens} write violations, got {write_violations}" + + +def test_pipeline_empty_batch() -> None: + """bs=1 with req_pool_idx=0 (padding): write and verify are no-op, kernel_run_counter == 2 (write+verify).""" + input_ids, positions, out_cache_loc = _zero_no_write_inputs() + + _, _, log_real, log_ref, _, _, _, _ = _run_both_and_assert_pipeline_equal( + req_pool_indices=_t([0]), + prefix_lens=_t([0]), + extend_seq_lens=_t([0]), + input_ids=input_ids, + positions=positions, + out_cache_loc=out_cache_loc, + req_to_token=_linear_r2t(), + num_slots=64, + extras=empty_extras(), + ) + + assert int(log_real.kernel_run_counter[0].item()) == 2 + assert int(log_ref.kernel_run_counter[0].item()) == 2 + assert int(log_real.write_index[0].item()) == 0 + + +def test_pipeline_negative_slot_swa_out_of_window() -> None: + """SWA: some out_cache_loc entries map to -1 (out-of-window); write_step skips them, buf unchanged.""" + max_seq_len = 16 + max_reqs = 4 + + full_to_swa_index_mapping = torch.arange( + max_reqs * max_seq_len + 1, dtype=torch.int64, device=_DEVICE + ) + full_to_swa_index_mapping[1 * max_seq_len + 6] = -1 + full_to_swa_index_mapping[1 * max_seq_len + 7] = -1 + + _run_both_and_assert_pipeline_equal( + req_pool_indices=_t([1]), + prefix_lens=_t([6]), + extend_seq_lens=_t([4]), + input_ids=_t([100, 101, 102, 103]), + positions=_t([6, 7, 8, 9]), + out_cache_loc=_t([-1, -1, 1 * max_seq_len + 8, 1 * max_seq_len + 9]), + req_to_token=_linear_r2t(max_reqs=max_reqs, max_seq_len=max_seq_len), + num_slots=128, + extras=empty_extras(), + swa_window_size=4, + full_to_swa_index_mapping=full_to_swa_index_mapping, + ) + + +def test_pipeline_ring_overflow_via_real_plan() -> None: + """Verify detects >capacity violations when prev_hash is pre-corrupted; write_index byte-equal, ring relaxed.""" + max_seq_len = 16 + max_reqs = 4 + req_to_token = _linear_r2t(max_reqs=max_reqs, max_seq_len=max_seq_len) + n_slots = 8 + + req_pool_indices = _t([1]) + prefix_lens = _t([n_slots]) + extend_seq_lens = _t([0]) + input_ids, positions, out_cache_loc = _zero_no_write_inputs() + + num_slots = max_reqs * max_seq_len + + # Step 1: pre-pollute canary_buf slots [0..n_slots) with wrong prev_hash so verify fires n_slots violations. + buf_real = make_canary_buf(num_slots=num_slots, device=_DEVICE) + buf_ref = make_canary_buf(num_slots=num_slots, device=_DEVICE) + for slot_idx in range(n_slots): + full_slot = 1 * max_seq_len + slot_idx + for buf in (buf_real, buf_ref): + write_slot_fields( + canary_buf=buf, + slot_idx=full_slot, + token=slot_idx + 1, + position=slot_idx, + prev_hash=0x1234_DEAD_BEEF_0000 + slot_idx, + ) + + # Step 2: run real pipeline (plan + no write + verify); overflow ring capacity=4 with all n_slots violations. + ring_capacity = 4 + log_real = FakeViolationLog.allocate(capacity=ring_capacity, device=_DEVICE) + log_ref = FakeViolationLog.allocate(capacity=ring_capacity, device=_DEVICE) + plan_v_real = VerifyPlan.allocate(verify_capacity=256, device=_DEVICE) + plan_w_real = WritePlan.allocate(write_req_capacity=4, device=_DEVICE) + plan_v_ref = VerifyPlan.allocate(verify_capacity=256, device=_DEVICE) + plan_w_ref = WritePlan.allocate(write_req_capacity=4, device=_DEVICE) + + launch_canary_plan_kernels( + verify_plan_out=plan_v_real, + write_plan_out=plan_w_real, + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + extend_seq_lens=extend_seq_lens, + req_to_token=req_to_token, + swa_window_size=0, + full_to_swa_index_mapping=None, + verify_capacity=int(plan_v_real.verify_slot_indices.shape[0]), + req_to_verify_expected_tokens=None, + req_to_verify_expected_tokens_valid_lens=None, + kv_token_id_vs_position_offset=0, + ) + launch_canary_plan_kernels_torch_reference( + verify_plan_out=plan_v_ref, + write_plan_out=plan_w_ref, + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + extend_seq_lens=extend_seq_lens, + req_to_token=req_to_token, + swa_window_size=0, + full_to_swa_index_mapping=None, + verify_capacity=int(plan_v_ref.verify_slot_indices.shape[0]), + req_to_verify_expected_tokens=None, + req_to_verify_expected_tokens_valid_lens=None, + kv_token_id_vs_position_offset=0, + ) + + launch_canary_verify_kernel( + context=VerifyOrWriteContext( + canary_buf=buf_real, + kernel_kind=CanaryLaunchTag.HEAD_K_FULL, + violation_ring=log_real.ring, + violation_write_index=log_real.write_index, + slot_run_counter=log_real.slot_run_counter, + kernel_run_counter=log_real.kernel_run_counter, + enable_chain_position_assert=log_real.enable_chain_position_assert, + ), + plan=plan_v_real, + check_verify_expected_token=True, + ) + torch.cuda.synchronize() + + launch_canary_verify_kernel_torch_reference( + context=VerifyOrWriteContext( + canary_buf=buf_ref, + kernel_kind=CanaryLaunchTag.HEAD_K_FULL, + violation_ring=log_ref.ring, + violation_write_index=log_ref.write_index, + slot_run_counter=log_ref.slot_run_counter, + kernel_run_counter=log_ref.kernel_run_counter, + enable_chain_position_assert=log_ref.enable_chain_position_assert, + ), + plan=plan_v_ref, + check_verify_expected_token=True, + ) + + # Step 3: write_index byte-equal; ring contents relaxed (atomic order not guaranteed under overflow). + assert torch.equal(log_real.write_index, log_ref.write_index) + assert int(log_real.write_index[0].item()) == n_slots + + +@pytest.mark.parametrize( + "kernel_kind", [CanaryLaunchTag.HEAD_K_FULL, CanaryLaunchTag.TAIL_V_SWA] +) +def test_pipeline_kernel_kind_propagates(kernel_kind: CanaryLaunchTag) -> None: + """Different CanaryLaunchTag values: violation ring's kernel_kind field matches on both sides.""" + max_seq_len = 16 + input_ids, positions, out_cache_loc = _zero_no_write_inputs() + + initial_buf = make_canary_buf(num_slots=64, device=_DEVICE) + write_slot_fields( + canary_buf=initial_buf, + slot_idx=1 * max_seq_len, + token=7, + position=99, + prev_hash=0, + ) + + _, _, log_real, log_ref, _, _, _, _ = _run_both_and_assert_pipeline_equal( + req_pool_indices=_t([1]), + prefix_lens=_t([1]), + extend_seq_lens=_t([0]), + input_ids=input_ids, + positions=positions, + out_cache_loc=out_cache_loc, + req_to_token=_linear_r2t(max_seq_len=max_seq_len), + num_slots=64, + extras=empty_extras(), + kernel_kind=kernel_kind, + initial_canary_buf=initial_buf, + ) + + assert int(log_real.write_index[0].item()) == 1 + assert int(log_ref.write_index[0].item()) == 1 + assert int(log_real.ring[0, consts.VIOLATION_FIELD_KERNEL_KIND].item()) == int( + kernel_kind + ) + assert int(log_ref.ring[0, consts.VIOLATION_FIELD_KERNEL_KIND].item()) == int( + kernel_kind + ) + + +def test_pipeline_token_mismatch_detected_via_pool() -> None: + """plan-pool gather + verify-token check: stamped wrong token id raises VERIFY_TOKEN_MISMATCH.""" + max_seq_len = 16 + max_reqs = 4 + prefix_len = 4 + input_ids, positions, out_cache_loc = _zero_no_write_inputs() + + expected_tokens = [1000 + pos for pos in range(prefix_len)] + pool = torch.full((max_reqs, max_seq_len), -999, dtype=torch.int32, device=_DEVICE) + for pos, token in enumerate(expected_tokens): + pool[1, pos] = token + + stored_tokens = [token + 1 for token in expected_tokens] + + initial_buf = make_canary_buf(num_slots=64, device=_DEVICE) + initial_ref = initial_buf.clone() + _stamp_linear_prefix( + initial_buf=initial_buf, + initial_ref=initial_ref, + req_pool_idx=1, + prefix_len=prefix_len, + tokens=stored_tokens, + max_seq_len=max_seq_len, + ) + + _, _, log_real, log_ref, _, _, _, _ = _run_both_and_assert_pipeline_equal( + req_pool_indices=_t([1]), + prefix_lens=_t([prefix_len]), + extend_seq_lens=_t([0]), + input_ids=input_ids, + positions=positions, + out_cache_loc=out_cache_loc, + req_to_token=_linear_r2t(max_reqs=max_reqs, max_seq_len=max_seq_len), + num_slots=64, + extras=empty_extras(), + swa_window_size=0, + full_to_swa_index_mapping=None, + initial_canary_buf=initial_buf, + req_to_verify_expected_tokens=pool, + kv_token_id_vs_position_offset=0, + check_verify_expected_token=True, + ) + + assert int(log_real.write_index[0].item()) == prefix_len + assert int(log_ref.write_index[0].item()) == prefix_len + # Ring rows may land in any order; collect stored/expected pairs and compare as sets. + observed_pairs: set[tuple[int, int]] = set() + for row_idx in range(prefix_len): + fail_bits = int( + log_real.ring[row_idx, consts.VIOLATION_FIELD_FAIL_REASON_BITS].item() + ) + assert fail_bits & int( + consts.FailReason.VERIFY_TOKEN_MISMATCH + ), f"row {row_idx}: VERIFY_TOKEN_MISMATCH bit missing in {fail_bits:#b}" + stored = int(log_real.ring[row_idx, consts.VIOLATION_FIELD_STORED_TOKEN].item()) + expected = int( + log_real.ring[row_idx, consts.VIOLATION_FIELD_EXPECTED_TOKEN].item() + ) + observed_pairs.add((stored, expected)) + expected_pairs = {(stored_tokens[i], expected_tokens[i]) for i in range(prefix_len)} + assert observed_pairs == expected_pairs + + +def test_pipeline_eagle_offset_plus_1_byte_equal() -> None: + """plan-pool + offset=+1 full pipeline: stamped tokens match pool[rp, pos+1], no violations CUDA vs ref byte-equal.""" + max_seq_len = 16 + max_reqs = 4 + prefix_len = 4 + input_ids, positions, out_cache_loc = _zero_no_write_inputs() + + stored_tokens = [2000 + pos for pos in range(prefix_len)] + pool = torch.full((max_reqs, max_seq_len), -999, dtype=torch.int32, device=_DEVICE) + for pos in range(prefix_len): + # offset=+1 means kernel gathers from pool[rp, pos + 1], so place stored_tokens[pos] there. + pool[1, pos + 1] = stored_tokens[pos] + + initial_buf = make_canary_buf(num_slots=64, device=_DEVICE) + initial_ref = initial_buf.clone() + _stamp_linear_prefix( + initial_buf=initial_buf, + initial_ref=initial_ref, + req_pool_idx=1, + prefix_len=prefix_len, + tokens=stored_tokens, + max_seq_len=max_seq_len, + ) + + _, _, log_real, log_ref, _, _, _, _ = _run_both_and_assert_pipeline_equal( + req_pool_indices=_t([1]), + prefix_lens=_t([prefix_len]), + extend_seq_lens=_t([0]), + input_ids=input_ids, + positions=positions, + out_cache_loc=out_cache_loc, + req_to_token=_linear_r2t(max_reqs=max_reqs, max_seq_len=max_seq_len), + num_slots=64, + extras=empty_extras(), + swa_window_size=0, + full_to_swa_index_mapping=None, + initial_canary_buf=initial_buf, + req_to_verify_expected_tokens=pool, + kv_token_id_vs_position_offset=1, + check_verify_expected_token=True, + ) + + assert int(log_real.write_index[0].item()) == 0 + assert int(log_ref.write_index[0].item()) == 0 + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) diff --git a/python/sglang/jit_kernel/tests/kv_canary/test_plan_fuzz.py b/python/sglang/jit_kernel/tests/kv_canary/test_plan_fuzz.py new file mode 100644 index 000000000..de7d6c125 --- /dev/null +++ b/python/sglang/jit_kernel/tests/kv_canary/test_plan_fuzz.py @@ -0,0 +1,222 @@ +from __future__ import annotations + +import random +from dataclasses import dataclass +from typing import Optional + +import pytest +import torch + +from sglang.jit_kernel.tests.kv_canary._differential import _run_both_plan +from sglang.jit_kernel.tests.kv_canary._fixtures import ( + allocate_plan_pair, + derive_plan_capacity, + make_lut, + make_padding_mask, + make_req_to_token, +) +from sglang.jit_kernel.tests.kv_canary._fuzz_driver import ( + FUZZ_SEEDS_PR, + run_fuzz_combo, +) +from sglang.jit_kernel.tests.kv_canary._invariants import PlanInvariants +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large") + + +_DEVICE = torch.device("cuda") + +_FUZZ_ITER_PER_SEED = 50 + + +@dataclass(frozen=True, slots=True, kw_only=True) +class PlanFuzzInputs: + 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 + write_req_capacity: int + req_to_verify_expected_tokens: Optional[torch.Tensor] + kv_token_id_vs_position_offset: int + + +def _draw_random_plan_inputs(rng: random.Random) -> PlanFuzzInputs: + bs = rng.randint(1, 16) + max_seq_len = rng.choice([8, 16, 64, 128, 256]) + swa_enabled = rng.random() < 0.5 + swa_window_size = ( + rng.choice([4, 16, 64, max_seq_len, max(2, max_seq_len // 3)]) + if swa_enabled + else 0 + ) + swa_window_size = min(swa_window_size, max_seq_len) + lut_kind = ( + rng.choice(["identity", "shift", "permutation", "with_oob"]) + if swa_enabled + else None + ) + rtt_kind = rng.choice(["linear", "sparse_permuted"]) + padding_kind = rng.choice(["none", "trailing", "interleaved"]) + capacity_kind = rng.choice(["loose", "tight_match", "under_by_one"]) + + max_reqs = max(bs + 2, 4) + pool_size = max_reqs * max_seq_len + rtt = make_req_to_token( + kind=rtt_kind, + max_reqs=max_reqs, + max_seq_len=max_seq_len, + device=_DEVICE, + rng=rng, + ) + padding_mask = make_padding_mask(bs=bs, kind=padding_kind, rng=rng) + req_pool_indices_list: list[int] = [] + prefix_lens_list: list[int] = [] + extend_seq_lens_list: list[int] = [] + for r in range(bs): + if padding_mask[r]: + req_pool_indices_list.append(0) + prefix_lens_list.append(0) + extend_seq_lens_list.append(0) + else: + req_pool_indices_list.append(rng.randint(1, max_reqs - 1)) + prefix_lens_list.append(rng.randint(0, max_seq_len - 1)) + extend_seq_lens_list.append(rng.randint(1, max(1, max_seq_len // 4))) + req_pool_indices = torch.tensor( + req_pool_indices_list, dtype=torch.int64, device=_DEVICE + ) + prefix_lens = torch.tensor(prefix_lens_list, dtype=torch.int64, device=_DEVICE) + extend_seq_lens = torch.tensor( + extend_seq_lens_list, dtype=torch.int64, device=_DEVICE + ) + + total_verify = 0 + for rpi, pfx in zip(req_pool_indices_list, prefix_lens_list): + if rpi == 0: + continue + if swa_window_size > 0: + window_start = max(0, pfx - swa_window_size) + total_verify += max(0, pfx - window_start) + else: + total_verify += pfx + + verify_capacity, write_req_capacity = derive_plan_capacity( + kind=capacity_kind, + total_verify=total_verify, + extras_count=0, + bs=bs, + ) + + full_to_swa: Optional[torch.Tensor] + if swa_window_size > 0 and lut_kind is not None: + full_to_swa = make_lut( + kind=lut_kind, pool_size=pool_size, device=_DEVICE, rng=rng + ) + else: + full_to_swa = None + + expected_pool_present = rng.random() < 0.5 + kv_token_id_vs_position_offset = rng.choice([0, 1]) + expected_pool: Optional[torch.Tensor] + if expected_pool_present: + pool_max_context_len = rng.choice( + [ + max(1, max_seq_len // 4), + max(1, max_seq_len // 2), + max_seq_len, + ] + ) + expected_pool = torch.randint( + low=0, + high=50000, + size=(max_reqs, pool_max_context_len), + dtype=torch.int32, + device=_DEVICE, + ) + else: + expected_pool = None + + return PlanFuzzInputs( + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + extend_seq_lens=extend_seq_lens, + req_to_token=rtt, + swa_window_size=swa_window_size, + full_to_swa_index_mapping=full_to_swa, + verify_capacity=verify_capacity, + write_req_capacity=write_req_capacity, + req_to_verify_expected_tokens=expected_pool, + kv_token_id_vs_position_offset=kv_token_id_vs_position_offset, + ) + + +def _run_one(inputs: PlanFuzzInputs) -> tuple: + triton_v, triton_w, ref_v, ref_w = allocate_plan_pair( + verify_capacity=inputs.verify_capacity, + write_req_capacity=inputs.write_req_capacity, + ) + _run_both_plan( + triton_verify=triton_v, + triton_write=triton_w, + ref_verify=ref_v, + ref_write=ref_w, + req_pool_indices=inputs.req_pool_indices, + prefix_lens=inputs.prefix_lens, + extend_seq_lens=inputs.extend_seq_lens, + req_to_token=inputs.req_to_token, + extras=( + torch.empty(0, dtype=torch.int64, device=_DEVICE), + torch.empty(0, dtype=torch.int64, device=_DEVICE), + torch.empty(0, dtype=torch.int64, device=_DEVICE), + torch.zeros(1, dtype=torch.int32, device=_DEVICE), + ), + swa_window_size=inputs.swa_window_size, + full_to_swa_index_mapping=inputs.full_to_swa_index_mapping, + req_to_verify_expected_tokens=inputs.req_to_verify_expected_tokens, + kv_token_id_vs_position_offset=inputs.kv_token_id_vs_position_offset, + ) + PlanInvariants.assert_all( + verify_plan=triton_v, + write_plan=triton_w, + req_pool_indices=inputs.req_pool_indices, + prefix_lens=inputs.prefix_lens, + extend_seq_lens=inputs.extend_seq_lens, + swa_window_size=inputs.swa_window_size, + extras_slot_indices=torch.empty(0, dtype=torch.int64, device=_DEVICE), + extras_positions=torch.empty(0, dtype=torch.int64, device=_DEVICE), + extras_prev_slot_indices=torch.empty(0, dtype=torch.int64, device=_DEVICE), + extras_count=0, + ) + return triton_v, triton_w + + +def _summarize(inputs: PlanFuzzInputs) -> str: + return ( + f"bs={int(inputs.req_pool_indices.shape[0])} " + f"swa={inputs.swa_window_size} " + f"verify_cap={inputs.verify_capacity} write_cap={inputs.write_req_capacity} " + f"has_lut={inputs.full_to_swa_index_mapping is not None} " + f"has_pool={inputs.req_to_verify_expected_tokens is not None} " + f"offset={inputs.kv_token_id_vs_position_offset}" + ) + + +@pytest.mark.parametrize("seed", FUZZ_SEEDS_PR) +def test_plan_fuzz_full_combo(seed: int) -> None: + """Multi-dim plan fuzzer: random LUT/rtt/padding/capacity/swa × N iters, byte-equal.""" + run_fuzz_combo( + seed, + draw_fn=_draw_random_plan_inputs, + run_one_fn=_run_one, + summarize_fn=_summarize, + n_iter=_FUZZ_ITER_PER_SEED, + ) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) diff --git a/python/sglang/jit_kernel/tests/kv_canary/test_plan_hand.py b/python/sglang/jit_kernel/tests/kv_canary/test_plan_hand.py new file mode 100644 index 000000000..c0f9b1c3f --- /dev/null +++ b/python/sglang/jit_kernel/tests/kv_canary/test_plan_hand.py @@ -0,0 +1,1371 @@ +from __future__ import annotations + +import random + +import pytest +import torch + +from sglang.jit_kernel.kv_canary.plan import launch_canary_plan_kernels +from sglang.jit_kernel.kv_canary.plan_ref import ( + launch_canary_plan_kernels_torch_reference, +) +from sglang.jit_kernel.kv_canary.verify import VerifyPlan +from sglang.jit_kernel.kv_canary.write import WritePlan +from sglang.jit_kernel.tests.kv_canary._differential import run_plan_diff +from sglang.jit_kernel.tests.kv_canary._fixtures import ( + allocate_plan_pair, + derive_plan_capacity, + empty_extras, + make_lut, + make_req_to_token, +) +from sglang.jit_kernel.tests.kv_canary._invariants import PlanInvariants +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large") + + +_DEVICE = torch.device("cuda") + + +def _tensor(values: list[int]) -> torch.Tensor: + return torch.tensor(values, dtype=torch.int64, device=_DEVICE) + + +def _plan_pair( + *, verify_capacity: int, write_req_capacity: int +) -> tuple[tuple[VerifyPlan, WritePlan], tuple[VerifyPlan, WritePlan]]: + triton_v, triton_w, ref_v, ref_w = allocate_plan_pair( + verify_capacity=verify_capacity, write_req_capacity=write_req_capacity + ) + return (triton_v, triton_w), (ref_v, ref_w) + + +def _alloc_for_inputs( + *, + req_pool_indices: torch.Tensor, + prefix_lens: torch.Tensor, + extend_seq_lens: torch.Tensor, + extras_count: int, + swa_window_size: int, +) -> tuple[int, int]: + bs = int(req_pool_indices.shape[0]) + rpi_cpu = req_pool_indices.detach().cpu().tolist() + pfx_cpu = prefix_lens.detach().cpu().tolist() + ext_cpu = extend_seq_lens.detach().cpu().tolist() + total_verify = 0 + for rpi, pfx in zip(rpi_cpu, pfx_cpu): + if rpi == 0: + continue + if swa_window_size > 0: + window_start = max(0, pfx - swa_window_size) + total_verify += max(0, pfx - window_start) + else: + total_verify += max(0, pfx) + return derive_plan_capacity( + kind="loose", total_verify=total_verify, extras_count=extras_count, bs=bs + ) + + +def _run_label( + *, + label: str, + req_pool_indices: torch.Tensor, + prefix_lens: torch.Tensor, + extend_seq_lens: torch.Tensor, + req_to_token: torch.Tensor, + extras: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], + swa_window_size: int, + full_to_swa_index_mapping: torch.Tensor | None, + verify_capacity: int, + write_req_capacity: int, +) -> tuple[VerifyPlan, WritePlan]: + _ = extras + verify_plan = VerifyPlan.allocate( + verify_capacity=verify_capacity, device=_DEVICE + ).zero_for_testing_() + write_plan = WritePlan.allocate( + write_req_capacity=write_req_capacity, device=_DEVICE + ).zero_for_testing_() + runner = ( + launch_canary_plan_kernels + if label == "real" + else launch_canary_plan_kernels_torch_reference + ) + runner( + verify_plan_out=verify_plan, + write_plan_out=write_plan, + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + extend_seq_lens=extend_seq_lens, + req_to_token=req_to_token, + swa_window_size=swa_window_size, + full_to_swa_index_mapping=full_to_swa_index_mapping, + verify_capacity=verify_capacity, + req_to_verify_expected_tokens=None, + req_to_verify_expected_tokens_valid_lens=None, + kv_token_id_vs_position_offset=0, + ) + torch.cuda.synchronize() + return verify_plan, write_plan + + +class TestBasicShape: + def test_single_req_extend_basic(self) -> None: + """bs=1, prefix=0, extend=5 → verify entries empty; write_offsets[0:2] = [0, 5]; seed = -1.""" + # Step 1: build a one-req batch with no prefix and 5 extend tokens. + req_to_token = make_req_to_token( + kind="linear", max_reqs=4, max_seq_len=16, device=_DEVICE + ) + plans = _plan_pair(verify_capacity=64, write_req_capacity=4) + run_plan_diff( + plan_pair=plans, + req_pool_indices=_tensor([3]), + prefix_lens=_tensor([0]), + extend_seq_lens=_tensor([5]), + req_to_token=req_to_token, + extras=empty_extras(), + ) + + triton_v, triton_w = plans[0] + # Step 2: prefix=0 → no verify entries; seed = -1 because prefix==0. + assert int(triton_v.verify_num_valid[0].item()) == 0 + assert int(triton_w.write_num_valid_reqs[0].item()) == 1 + assert int(triton_w.write_offsets[0].item()) == 0 + assert int(triton_w.write_offsets[1].item()) == 5 + assert int(triton_w.write_seed_slot_indices[0].item()) == -1 + + def test_single_req_decode(self) -> None: + """extend=1, prefix=K → write_seed_slot = req_to_token[rp, K-1]; verify covers all K prefix tokens.""" + max_seq_len = 16 + req_to_token = make_req_to_token( + kind="linear", max_reqs=4, max_seq_len=max_seq_len, device=_DEVICE + ) + plans = _plan_pair(verify_capacity=64, write_req_capacity=4) + run_plan_diff( + plan_pair=plans, + req_pool_indices=_tensor([2]), + prefix_lens=_tensor([7]), + extend_seq_lens=_tensor([1]), + req_to_token=req_to_token, + extras=empty_extras(), + ) + + triton_v, triton_w = plans[0] + # Step: verify covers positions [0..7); seed slot for req rp=2 is at position 6 = rp * max_seq_len + 6. + assert int(triton_v.verify_num_valid[0].item()) == 7 + assert int(triton_w.write_seed_slot_indices[0].item()) == 2 * max_seq_len + 6 + + def test_multi_req_mixed_extend_decode(self) -> None: + """bs=3 mixed extend/decode → write_offsets cumsum is byte-equal across Triton + ref.""" + req_to_token = make_req_to_token( + kind="linear", max_reqs=4, max_seq_len=16, device=_DEVICE + ) + plans = _plan_pair(verify_capacity=64, write_req_capacity=8) + # req0: prefill extend=8; req1: decode extend=1; req2: decode extend=1. + run_plan_diff( + plan_pair=plans, + req_pool_indices=_tensor([1, 2, 3]), + prefix_lens=_tensor([0, 4, 10]), + extend_seq_lens=_tensor([8, 1, 1]), + req_to_token=req_to_token, + extras=empty_extras(), + ) + + triton_v, triton_w = plans[0] + # Step: write_offsets exclusive cumsum on extend_seq_lens. + expected_write_offsets = [0, 8, 9, 10] + for i, value in enumerate(expected_write_offsets): + assert int(triton_w.write_offsets[i].item()) == value + # Verify count = 0 + 4 + 10 = 14. + assert int(triton_v.verify_num_valid[0].item()) == 14 + + +class TestSeedSlot: + def test_prefix_zero_seed_is_minus_one(self) -> None: + """prefix=0 → seed_slot_idx = -1 (no predecessor to anchor on).""" + req_to_token = make_req_to_token( + kind="linear", max_reqs=4, max_seq_len=16, device=_DEVICE + ) + plans = _plan_pair(verify_capacity=64, write_req_capacity=4) + run_plan_diff( + plan_pair=plans, + req_pool_indices=_tensor([1]), + prefix_lens=_tensor([0]), + extend_seq_lens=_tensor([3]), + req_to_token=req_to_token, + extras=empty_extras(), + ) + + assert int(plans[0][1].write_seed_slot_indices[0].item()) == -1 + + def test_prev_slot_minus_one_at_chain_head(self) -> None: + """pos=0 entry → verify_prev_slot_indices == -1 (chain head).""" + req_to_token = make_req_to_token( + kind="linear", max_reqs=4, max_seq_len=16, device=_DEVICE + ) + plans = _plan_pair(verify_capacity=64, write_req_capacity=4) + run_plan_diff( + plan_pair=plans, + req_pool_indices=_tensor([1]), + prefix_lens=_tensor([3]), + extend_seq_lens=_tensor([1]), + req_to_token=req_to_token, + extras=empty_extras(), + ) + + # First entry has pos=0 → prev_slot = -1. + assert int(plans[0][0].verify_prev_slot_indices[0].item()) == -1 + + def test_prev_slot_is_self_minus_one(self) -> None: + """pos>0 entry → prev = req_to_token[rp, pos-1] (SWA-translated when SWA enabled).""" + max_seq_len = 16 + req_to_token = make_req_to_token( + kind="linear", max_reqs=4, max_seq_len=max_seq_len, device=_DEVICE + ) + plans = _plan_pair(verify_capacity=64, write_req_capacity=4) + run_plan_diff( + plan_pair=plans, + req_pool_indices=_tensor([2]), + prefix_lens=_tensor([4]), + extend_seq_lens=_tensor([1]), + req_to_token=req_to_token, + extras=empty_extras(), + ) + + triton_v = plans[0][0] + # entry[1] is pos=1: prev_slot = req_to_token[2, 0] = 2 * max_seq_len + 0 = 32. + assert int(triton_v.verify_prev_slot_indices[1].item()) == 2 * max_seq_len + 0 + # entry[2] is pos=2: prev_slot = req_to_token[2, 1] = 2 * max_seq_len + 1 = 33. + assert int(triton_v.verify_prev_slot_indices[2].item()) == 2 * max_seq_len + 1 + + def test_seed_translated_through_permuted_lut(self) -> None: + """Permuted LUT: seed slot is the LUT-lookup of req_to_token[rp, prefix-1], NOT identity.""" + rng = random.Random(42) + max_seq_len = 16 + max_reqs = 4 + pool_size = max_reqs * max_seq_len + lut = make_lut(kind="permutation", pool_size=pool_size, device=_DEVICE, rng=rng) + rtt = make_req_to_token( + kind="linear", max_reqs=max_reqs, max_seq_len=max_seq_len, device=_DEVICE + ) + + rp = 2 + prefix = 5 + req_pool_indices = _tensor([rp]) + prefix_lens = _tensor([prefix]) + extend_seq_lens = _tensor([1]) + + full_seed_slot = rp * max_seq_len + (prefix - 1) + expected_seed = int(lut[full_seed_slot].item()) + + extras = empty_extras() + verify_capacity, write_req_capacity = _alloc_for_inputs( + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + extend_seq_lens=extend_seq_lens, + extras_count=0, + swa_window_size=max_seq_len, + ) + for label in ("real", "ref"): + _, w_plan = _run_label( + label=label, + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + extend_seq_lens=extend_seq_lens, + req_to_token=rtt, + extras=extras, + swa_window_size=max_seq_len, + full_to_swa_index_mapping=lut, + verify_capacity=verify_capacity, + write_req_capacity=write_req_capacity, + ) + actual_seed = int(w_plan.write_seed_slot_indices[0].item()) + assert ( + actual_seed == expected_seed + ), f"[{label}] permuted-LUT seed expected {expected_seed} got {actual_seed}" + + def test_swa_window_head_prev_slot_is_real_predecessor(self) -> None: + """SWA window with non-zero window_start: head entry's prev_slot != -1; it is the real predecessor.""" + rng = random.Random(13) + max_seq_len = 256 + max_reqs = 2 + pool_size = max_reqs * max_seq_len + swa_window_size = 128 + prefix = 200 + rp = 1 + lut = make_lut(kind="permutation", pool_size=pool_size, device=_DEVICE, rng=rng) + rtt = make_req_to_token( + kind="linear", max_reqs=max_reqs, max_seq_len=max_seq_len, device=_DEVICE + ) + + req_pool_indices = _tensor([rp]) + prefix_lens = _tensor([prefix]) + extend_seq_lens = _tensor([1]) + extras = empty_extras() + + window_start = prefix - swa_window_size + full_prev_slot = int(rtt[rp, window_start - 1].item()) + expected_prev = int(lut[full_prev_slot].item()) + + verify_capacity, write_req_capacity = derive_plan_capacity( + kind="loose", total_verify=swa_window_size, extras_count=0, bs=1 + ) + + for label in ("real", "ref"): + v_plan, _ = _run_label( + label=label, + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + extend_seq_lens=extend_seq_lens, + req_to_token=rtt, + extras=extras, + swa_window_size=swa_window_size, + full_to_swa_index_mapping=lut, + verify_capacity=verify_capacity, + write_req_capacity=write_req_capacity, + ) + actual_prev = int(v_plan.verify_prev_slot_indices[0].item()) + assert ( + actual_prev != -1 + ), f"[{label}] SWA window head must have real predecessor, got -1" + assert ( + actual_prev == expected_prev + ), f"[{label}] expected prev={expected_prev} got {actual_prev}" + + +class TestPadding: + def test_padding_rows_contribute_zero(self) -> None: + """``req_pool_indices[r] == 0`` rows → no verify entry, no write entry, seed = -1.""" + req_to_token = make_req_to_token( + kind="linear", max_reqs=4, max_seq_len=16, device=_DEVICE + ) + plans = _plan_pair(verify_capacity=64, write_req_capacity=4) + # Step: bs=3 with row 1 marked as padding (rpi=0). + run_plan_diff( + plan_pair=plans, + req_pool_indices=_tensor([1, 0, 2]), + prefix_lens=_tensor([5, 99, 3]), + extend_seq_lens=_tensor([1, 99, 1]), + req_to_token=req_to_token, + extras=empty_extras(), + ) + + triton_v, triton_w = plans[0] + # verify count = 5 (req0) + 0 (padding) + 3 (req2) = 8. + assert int(triton_v.verify_num_valid[0].item()) == 8 + # write_offsets cumsum: [0, 1, 1, 2] — padding row contributes 0. + expected_write_offsets = [0, 1, 1, 2] + for i, value in enumerate(expected_write_offsets): + assert int(triton_w.write_offsets[i].item()) == value + # Padding row's seed must be -1. + assert int(triton_w.write_seed_slot_indices[1].item()) == -1 + + def test_per_req_slot_when_req_to_token_is_sparse(self) -> None: + """sparse_permuted rtt: verify_slot_indices read directly from the constructed table.""" + rng = random.Random(7) + max_seq_len = 8 + max_reqs = 3 + rtt = make_req_to_token( + kind="sparse_permuted", + max_reqs=max_reqs, + max_seq_len=max_seq_len, + device=_DEVICE, + rng=rng, + ) + rp = 1 + prefix = 4 + req_pool_indices = _tensor([rp]) + prefix_lens = _tensor([prefix]) + extend_seq_lens = _tensor([1]) + extras = empty_extras() + verify_capacity, write_req_capacity = _alloc_for_inputs( + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + extend_seq_lens=extend_seq_lens, + extras_count=0, + swa_window_size=0, + ) + + expected_slots = [int(rtt[rp, pos].item()) for pos in range(prefix)] + + for label in ("real", "ref"): + v_plan, _ = _run_label( + label=label, + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + extend_seq_lens=extend_seq_lens, + req_to_token=rtt, + extras=extras, + swa_window_size=0, + full_to_swa_index_mapping=None, + verify_capacity=verify_capacity, + write_req_capacity=write_req_capacity, + ) + actual_slots = v_plan.verify_slot_indices[:prefix].detach().cpu().tolist() + assert ( + actual_slots == expected_slots + ), f"[{label}] sparse-rtt slots expected {expected_slots} got {actual_slots}" + + def test_padding_row_with_garbage_prefix_does_not_oob(self) -> None: + """rpi==0 padding row with absurd prefix_lens must not OOB-read req_to_token (row is skipped).""" + req_pool_indices = _tensor([1, 0, 2]) + prefix_lens = _tensor([5, 99999, 3]) + extend_seq_lens = _tensor([1, 99999, 1]) + rtt = make_req_to_token( + kind="linear", max_reqs=4, max_seq_len=16, device=_DEVICE + ) + extras = empty_extras() + verify_capacity, write_req_capacity = derive_plan_capacity( + kind="loose", total_verify=8, extras_count=0, bs=3 + ) + + for label in ("real", "ref"): + v_plan, w_plan = _run_label( + label=label, + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + extend_seq_lens=extend_seq_lens, + req_to_token=rtt, + extras=extras, + swa_window_size=0, + full_to_swa_index_mapping=None, + verify_capacity=verify_capacity, + write_req_capacity=write_req_capacity, + ) + assert int(v_plan.verify_num_valid[0].item()) == 8, label + assert ( + int(w_plan.write_seed_slot_indices[1].item()) == -1 + ), f"[{label}] padding row seed must be -1" + PlanInvariants.assert_all( + verify_plan=v_plan, + write_plan=w_plan, + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + extend_seq_lens=extend_seq_lens, + swa_window_size=0, + extras_slot_indices=extras[0], + extras_positions=extras[1], + extras_prev_slot_indices=extras[2], + extras_count=0, + ) + + +class TestSwa: + def test_swa_window_clip_prefix_less_than_window(self) -> None: + """SWA: prefix=3 < window=128 → window_start=0, verify covers 3 entries (no clip).""" + req_to_token = make_req_to_token( + kind="linear", max_reqs=4, max_seq_len=512, device=_DEVICE + ) + # Identity LUT keeps slot indices unchanged after SWA translation. + full_pool_size = 4 * 512 + lut = torch.arange(full_pool_size + 1, dtype=torch.int64, device=_DEVICE) + plans = _plan_pair(verify_capacity=256, write_req_capacity=4) + run_plan_diff( + plan_pair=plans, + req_pool_indices=_tensor([1]), + prefix_lens=_tensor([3]), + extend_seq_lens=_tensor([1]), + req_to_token=req_to_token, + extras=empty_extras(), + swa_window_size=128, + full_to_swa_index_mapping=lut, + ) + + assert int(plans[0][0].verify_num_valid[0].item()) == 3 + + def test_swa_window_clip_prefix_gt_window(self) -> None: + """SWA: prefix=200 > window=128 → window_start=72, verify covers 128 entries.""" + req_to_token = make_req_to_token( + kind="linear", max_reqs=4, max_seq_len=512, device=_DEVICE + ) + full_pool_size = 4 * 512 + lut = torch.arange(full_pool_size + 1, dtype=torch.int64, device=_DEVICE) + plans = _plan_pair(verify_capacity=512, write_req_capacity=4) + run_plan_diff( + plan_pair=plans, + req_pool_indices=_tensor([1]), + prefix_lens=_tensor([200]), + extend_seq_lens=_tensor([1]), + req_to_token=req_to_token, + extras=empty_extras(), + swa_window_size=128, + full_to_swa_index_mapping=lut, + ) + + triton_v = plans[0][0] + assert int(triton_v.verify_num_valid[0].item()) == 128 + # First verify entry should be at position 72. + assert int(triton_v.verify_expected_positions[0].item()) == 72 + + def test_swa_lut_translates_verify_slots(self) -> None: + """FULL slot → SWA slot translation is performed inside the plan kernel for verify_slot_indices.""" + max_seq_len = 16 + req_to_token = make_req_to_token( + kind="linear", max_reqs=4, max_seq_len=max_seq_len, device=_DEVICE + ) + full_pool_size = 4 * max_seq_len + # Build a LUT that maps FULL slot S → SWA slot (S + 100) for every S; chosen so we can distinguish a + # translated value from a raw full slot. + lut = ( + torch.arange(full_pool_size + 1, dtype=torch.int64, device=_DEVICE) + 100 + ).contiguous() + plans = _plan_pair(verify_capacity=64, write_req_capacity=4) + run_plan_diff( + plan_pair=plans, + req_pool_indices=_tensor([1]), + prefix_lens=_tensor([3]), + extend_seq_lens=_tensor([1]), + req_to_token=req_to_token, + extras=empty_extras(), + swa_window_size=128, + full_to_swa_index_mapping=lut, + ) + + triton_v = plans[0][0] + # FULL slot for (rp=1, pos=0) = 1 * max_seq_len + 0 = 16; expected SWA slot = 16 + 100 = 116. + assert int(triton_v.verify_slot_indices[0].item()) == 1 * max_seq_len + 0 + 100 + assert int(triton_v.verify_slot_indices[1].item()) == 1 * max_seq_len + 1 + 100 + assert int(triton_v.verify_slot_indices[2].item()) == 1 * max_seq_len + 2 + 100 + + def test_swa_lut_translates_seed_slot(self) -> None: + """write_seed_slot_indices is also SWA-translated inside the plan kernel.""" + max_seq_len = 16 + req_to_token = make_req_to_token( + kind="linear", max_reqs=4, max_seq_len=max_seq_len, device=_DEVICE + ) + full_pool_size = 4 * max_seq_len + lut = ( + torch.arange(full_pool_size + 1, dtype=torch.int64, device=_DEVICE) + 100 + ).contiguous() + plans = _plan_pair(verify_capacity=64, write_req_capacity=4) + run_plan_diff( + plan_pair=plans, + req_pool_indices=_tensor([1]), + prefix_lens=_tensor([3]), + extend_seq_lens=_tensor([1]), + req_to_token=req_to_token, + extras=empty_extras(), + swa_window_size=128, + full_to_swa_index_mapping=lut, + ) + + triton_w = plans[0][1] + # FULL slot at (rp=1, pos=2) = 1 * max_seq_len + 2 = 18; expected SWA seed = 18 + 100 = 118. + assert ( + int(triton_w.write_seed_slot_indices[0].item()) == 1 * max_seq_len + 2 + 100 + ) + + def test_verify_covers_all_tokens_in_swa_window(self) -> None: + """SWA group with window=128 + bs=4 → verify_num_valid == Σ min(prefix_lens[r], 128).""" + window = 128 + prefix_values = [50, 128, 200, 1024] + max_seq_len = 2048 + req_to_token = make_req_to_token( + kind="linear", max_reqs=4, max_seq_len=max_seq_len, device=_DEVICE + ) + full_pool_size = 4 * max_seq_len + lut = torch.arange(full_pool_size + 1, dtype=torch.int64, device=_DEVICE) + plans = _plan_pair(verify_capacity=1024, write_req_capacity=8) + run_plan_diff( + plan_pair=plans, + req_pool_indices=_tensor([1, 2, 3, 1]), + prefix_lens=_tensor(prefix_values), + extend_seq_lens=_tensor([1, 1, 1, 1]), + req_to_token=req_to_token, + extras=empty_extras(), + swa_window_size=window, + full_to_swa_index_mapping=lut, + ) + + expected_total = sum(min(p, window) for p in prefix_values) + assert int(plans[0][0].verify_num_valid[0].item()) == expected_total + + +class TestNoExtras: + def test_plan_num_valid_counts_only_per_req_entries(self) -> None: + """Sweep extras are written directly into VerifyPlan; plan kernel does not append them.""" + req_to_token = make_req_to_token( + kind="linear", max_reqs=4, max_seq_len=16, device=_DEVICE + ) + plans = _plan_pair(verify_capacity=64, write_req_capacity=4) + run_plan_diff( + plan_pair=plans, + req_pool_indices=_tensor([1]), + prefix_lens=_tensor([3]), + extend_seq_lens=_tensor([1]), + req_to_token=req_to_token, + extras=empty_extras(), + ) + + triton_v = plans[0][0] + assert int(triton_v.verify_num_valid[0].item()) == 3 + + def test_zero_prefix_has_no_verify_entries(self) -> None: + req_to_token = make_req_to_token( + kind="linear", max_reqs=4, max_seq_len=16, device=_DEVICE + ) + plans = _plan_pair(verify_capacity=64, write_req_capacity=4) + run_plan_diff( + plan_pair=plans, + req_pool_indices=_tensor([1]), + prefix_lens=_tensor([0]), + extend_seq_lens=_tensor([5]), + req_to_token=req_to_token, + extras=empty_extras(), + ) + + assert int(plans[0][0].verify_num_valid[0].item()) == 0 + + def test_verify_capacity_just_fits_per_req_entries(self) -> None: + rp = 1 + prefix = 4 + req_pool_indices = _tensor([rp]) + prefix_lens = _tensor([prefix]) + extend_seq_lens = _tensor([1]) + rtt = make_req_to_token( + kind="linear", max_reqs=4, max_seq_len=16, device=_DEVICE + ) + + total_verify = prefix + verify_capacity, write_req_capacity = derive_plan_capacity( + kind="tight_match", + total_verify=total_verify, + extras_count=0, + bs=1, + ) + + for label in ("real", "ref"): + v_plan, _ = _run_label( + label=label, + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + extend_seq_lens=extend_seq_lens, + req_to_token=rtt, + extras=empty_extras(), + swa_window_size=0, + full_to_swa_index_mapping=None, + verify_capacity=verify_capacity, + write_req_capacity=write_req_capacity, + ) + n = int(v_plan.verify_num_valid[0].item()) + assert n == total_verify, f"[{label}] num_valid {n}" + assert int(v_plan.enable[0].item()) == 1 + + def test_verify_capacity_undershoot_by_one(self) -> None: + rp = 1 + prefix = 3 + req_pool_indices = _tensor([rp]) + prefix_lens = _tensor([prefix]) + extend_seq_lens = _tensor([1]) + rtt = make_req_to_token( + kind="linear", max_reqs=4, max_seq_len=16, device=_DEVICE + ) + + total_verify = prefix + verify_capacity, write_req_capacity = derive_plan_capacity( + kind="under_by_one", + total_verify=total_verify, + extras_count=0, + bs=1, + ) + + real_v, _ = _run_label( + label="real", + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + extend_seq_lens=extend_seq_lens, + req_to_token=rtt, + extras=empty_extras(), + swa_window_size=0, + full_to_swa_index_mapping=None, + verify_capacity=verify_capacity, + write_req_capacity=write_req_capacity, + ) + ref_v, _ = _run_label( + label="ref", + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + extend_seq_lens=extend_seq_lens, + req_to_token=rtt, + extras=empty_extras(), + swa_window_size=0, + full_to_swa_index_mapping=None, + verify_capacity=verify_capacity, + write_req_capacity=write_req_capacity, + ) + n_real = int(real_v.verify_num_valid[0].item()) + n_ref = int(ref_v.verify_num_valid[0].item()) + assert n_real == n_ref, f"real {n_real} vs ref {n_ref} diverged under cap" + assert n_real == verify_capacity + assert int(real_v.enable[0].item()) == 0 + assert int(ref_v.enable[0].item()) == 0 + + +class TestMisc: + def test_zero_extend_writes_empty_write_plan(self) -> None: + """``extend_seq_lens`` all zero → write offsets stay zero; VerifyPlan still populated.""" + req_to_token = make_req_to_token( + kind="linear", max_reqs=4, max_seq_len=16, device=_DEVICE + ) + plans = _plan_pair(verify_capacity=64, write_req_capacity=4) + run_plan_diff( + plan_pair=plans, + req_pool_indices=_tensor([1, 2]), + prefix_lens=_tensor([4, 6]), + extend_seq_lens=_tensor([0, 0]), + req_to_token=req_to_token, + extras=empty_extras(), + ) + + triton_v, triton_w = plans[0] + # VerifyPlan covers 4+6 = 10 entries; write offsets are zero. + assert int(triton_v.verify_num_valid[0].item()) == 10 + # write_offsets cumsum of zeros stays zero across the active prefix. + assert int(triton_w.write_offsets[0].item()) == 0 + assert int(triton_w.write_offsets[1].item()) == 0 + assert int(triton_w.write_offsets[2].item()) == 0 + # Seeds for write-empty reqs must be -1 per plan semantics. + assert int(triton_w.write_seed_slot_indices[0].item()) == -1 + assert int(triton_w.write_seed_slot_indices[1].item()) == -1 + + def test_replay_same_inputs_yields_same_outputs(self) -> None: + """Two consecutive runs on identical inputs produce byte-equal plans (kernel is pure).""" + req_pool_indices = _tensor([1, 2, 3]) + prefix_lens = _tensor([4, 7, 2]) + extend_seq_lens = _tensor([2, 1, 3]) + rtt = make_req_to_token( + kind="linear", max_reqs=8, max_seq_len=16, device=_DEVICE + ) + extras = empty_extras() + verify_capacity, write_req_capacity = derive_plan_capacity( + kind="loose", total_verify=13, extras_count=0, bs=3 + ) + + for label in ("real", "ref"): + run1_v, run1_w = _run_label( + label=label, + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + extend_seq_lens=extend_seq_lens, + req_to_token=rtt, + extras=extras, + swa_window_size=0, + full_to_swa_index_mapping=None, + verify_capacity=verify_capacity, + write_req_capacity=write_req_capacity, + ) + run2_v, run2_w = _run_label( + label=label, + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + extend_seq_lens=extend_seq_lens, + req_to_token=rtt, + extras=extras, + swa_window_size=0, + full_to_swa_index_mapping=None, + verify_capacity=verify_capacity, + write_req_capacity=write_req_capacity, + ) + assert torch.equal( + run1_v.verify_slot_indices, run2_v.verify_slot_indices + ), label + assert torch.equal( + run1_v.verify_expected_positions, run2_v.verify_expected_positions + ), label + assert torch.equal( + run1_v.verify_prev_slot_indices, run2_v.verify_prev_slot_indices + ), label + assert torch.equal(run1_v.verify_num_valid, run2_v.verify_num_valid), label + assert torch.equal(run1_w.write_offsets, run2_w.write_offsets), label + assert torch.equal( + run1_w.write_seed_slot_indices, run2_w.write_seed_slot_indices + ), label + assert torch.equal( + run1_w.write_num_valid_reqs, run2_w.write_num_valid_reqs + ), label + + def test_shrink_bs_clears_stale_write_offsets(self) -> None: + """Reusing a WritePlan with smaller bs: write_offsets beyond new bs must be zeroed by the kernel.""" + rtt = make_req_to_token( + kind="linear", max_reqs=16, max_seq_len=16, device=_DEVICE + ) + verify_capacity, write_req_capacity = derive_plan_capacity( + kind="loose", total_verify=80, extras_count=0, bs=8 + ) + + for label in ("real", "ref"): + big_rpi = _tensor([1, 2, 3, 4, 5, 6, 7, 8]) + big_pfx = _tensor([10] * 8) + big_ext = _tensor([1] * 8) + small_rpi = _tensor([1, 2, 3]) + small_pfx = _tensor([5, 5, 5]) + small_ext = _tensor([1, 1, 1]) + + verify_plan = VerifyPlan.allocate( + verify_capacity=verify_capacity, device=_DEVICE + ).zero_for_testing_() + write_plan = WritePlan.allocate( + write_req_capacity=write_req_capacity, device=_DEVICE + ).zero_for_testing_() + runner = ( + launch_canary_plan_kernels + if label == "real" + else launch_canary_plan_kernels_torch_reference + ) + runner( + verify_plan_out=verify_plan, + write_plan_out=write_plan, + req_pool_indices=big_rpi, + prefix_lens=big_pfx, + extend_seq_lens=big_ext, + req_to_token=rtt, + swa_window_size=0, + full_to_swa_index_mapping=None, + verify_capacity=verify_capacity, + req_to_verify_expected_tokens=None, + req_to_verify_expected_tokens_valid_lens=None, + kv_token_id_vs_position_offset=0, + ) + torch.cuda.synchronize() + runner( + verify_plan_out=verify_plan, + write_plan_out=write_plan, + req_pool_indices=small_rpi, + prefix_lens=small_pfx, + extend_seq_lens=small_ext, + req_to_token=rtt, + swa_window_size=0, + full_to_swa_index_mapping=None, + verify_capacity=verify_capacity, + req_to_verify_expected_tokens=None, + req_to_verify_expected_tokens_valid_lens=None, + kv_token_id_vs_position_offset=0, + ) + torch.cuda.synchronize() + n_active = int(write_plan.write_num_valid_reqs[0].item()) + tail_offsets = ( + write_plan.write_offsets[n_active + 1 : 8].detach().cpu().tolist() + ) + assert all( + v == 0 for v in tail_offsets + ), f"[{label}] stale write_offsets tail not cleared: {tail_offsets}" + + +class TestVerifyContent: + def test_verify_num_valid_aggregate(self) -> None: + """``verify_num_valid == sum(per-req verify_count)``.""" + req_to_token = make_req_to_token( + kind="linear", max_reqs=4, max_seq_len=16, device=_DEVICE + ) + plans = _plan_pair(verify_capacity=64, write_req_capacity=4) + run_plan_diff( + plan_pair=plans, + req_pool_indices=_tensor([1, 2, 3]), + prefix_lens=_tensor([2, 5, 4]), + extend_seq_lens=_tensor([1, 1, 1]), + req_to_token=req_to_token, + extras=empty_extras(), + ) + + assert int(plans[0][0].verify_num_valid[0].item()) == 11 + + def test_verify_covers_all_tokens_no_skip(self) -> None: + """FULL group + bs=4 → verify_num_valid == Σ(prefix_lens) — every prefix token verified.""" + # Step: 4 reqs with mixed prefix and extend; FULL group means no SWA window clip. + prefix_values = [0, 3, 7, 12] + extend_values = [4, 1, 1, 1] + req_to_token = make_req_to_token( + kind="linear", max_reqs=4, max_seq_len=32, device=_DEVICE + ) + plans = _plan_pair(verify_capacity=128, write_req_capacity=8) + run_plan_diff( + plan_pair=plans, + req_pool_indices=_tensor([1, 2, 3, 1]), + prefix_lens=_tensor(prefix_values), + extend_seq_lens=_tensor(extend_values), + req_to_token=req_to_token, + extras=empty_extras(), + ) + + assert int(plans[0][0].verify_num_valid[0].item()) == sum(prefix_values) + + def test_plan_verify_expected_positions_strictly_increment_per_req(self) -> None: + """Per req, verify_expected_positions[verify_offsets[r]:verify_offsets[r+1]] == [window_start..prefix-1].""" + req_to_token = make_req_to_token( + kind="linear", max_reqs=4, max_seq_len=32, device=_DEVICE + ) + plans = _plan_pair(verify_capacity=64, write_req_capacity=4) + run_plan_diff( + plan_pair=plans, + req_pool_indices=_tensor([1, 2]), + prefix_lens=_tensor([5, 8]), + extend_seq_lens=_tensor([1, 1]), + req_to_token=req_to_token, + extras=empty_extras(), + ) + + triton_v = plans[0][0] + # Req 0: positions [0..5); Req 1: positions [0..8). + req0_positions = triton_v.verify_expected_positions[:5].cpu().tolist() + req1_positions = triton_v.verify_expected_positions[5:13].cpu().tolist() + assert req0_positions == [0, 1, 2, 3, 4] + assert req1_positions == [0, 1, 2, 3, 4, 5, 6, 7] + + def test_write_num_valid_reqs_excludes_padding(self) -> None: + """Padding rows (rpi == 0) contribute zero to write_offsets so the write kernel does no work for them, but they ARE included in write_num_valid_reqs, which equals bs (the full batch size including padding).""" + req_to_token = make_req_to_token( + kind="linear", max_reqs=4, max_seq_len=16, device=_DEVICE + ) + plans = _plan_pair(verify_capacity=64, write_req_capacity=8) + # bs=4, last two rows are padding. + run_plan_diff( + plan_pair=plans, + req_pool_indices=_tensor([1, 2, 0, 0]), + prefix_lens=_tensor([3, 5, 99, 99]), + extend_seq_lens=_tensor([1, 1, 99, 99]), + req_to_token=req_to_token, + extras=empty_extras(), + ) + + triton_w = plans[0][1] + # Padding rows must contribute 0 to write_offsets cumsum: [0, 1, 2, 2, 2]. + expected_write_offsets = [0, 1, 2, 2, 2] + for i, value in enumerate(expected_write_offsets): + assert int(triton_w.write_offsets[i].item()) == value + + +class TestByteEqual: + def test_byte_equal_python_reference(self) -> None: + """End-to-end Triton vs Python ref byte-equal across a representative bs=4 case (no SWA).""" + req_to_token = make_req_to_token( + kind="linear", max_reqs=4, max_seq_len=32, device=_DEVICE + ) + plans = _plan_pair(verify_capacity=128, write_req_capacity=8) + run_plan_diff( + plan_pair=plans, + req_pool_indices=_tensor([1, 2, 3, 1]), + prefix_lens=_tensor([0, 3, 8, 15]), + extend_seq_lens=_tensor([4, 1, 1, 1]), + req_to_token=req_to_token, + extras=empty_extras(), + ) + + def test_byte_equal_python_reference_hardcoded(self) -> None: + """bs=3, three prefix combinations → hand-computed verify_offsets / write_offsets / seed slots.""" + # Step 1: pin (prefix, extend) per req. + prefixes = [0, 4, 7] + extends = [3, 1, 1] + rps = [1, 2, 3] + max_seq_len = 16 + req_to_token = make_req_to_token( + kind="linear", max_reqs=4, max_seq_len=max_seq_len, device=_DEVICE + ) + plans = _plan_pair(verify_capacity=64, write_req_capacity=4) + run_plan_diff( + plan_pair=plans, + req_pool_indices=_tensor(rps), + prefix_lens=_tensor(prefixes), + extend_seq_lens=_tensor(extends), + req_to_token=req_to_token, + extras=empty_extras(), + ) + + triton_v, triton_w = plans[0] + # Step 2: hand-compute expected write_offsets (exclusive cumsum of extends) and verify_num_valid. + expected_write_offsets = [0, 3, 4, 5] + expected_verify_num_valid = sum(prefixes) + expected_seeds = [ + -1, # prefix=0 → no predecessor + rps[1] * max_seq_len + (prefixes[1] - 1), + rps[2] * max_seq_len + (prefixes[2] - 1), + ] + + for i, value in enumerate(expected_write_offsets): + assert ( + int(triton_w.write_offsets[i].item()) == value + ), f"write_offsets[{i}] expected {value} got {int(triton_w.write_offsets[i].item())}" + assert ( + int(triton_v.verify_num_valid[0].item()) == expected_verify_num_valid + ), f"verify_num_valid expected {expected_verify_num_valid}" + for i, expected_seed in enumerate(expected_seeds): + assert ( + int(triton_w.write_seed_slot_indices[i].item()) == expected_seed + ), f"write_seed_slot_indices[{i}] expected {expected_seed}" + + +class TestBoundarySweep: + @pytest.mark.parametrize("bs", [1, 31, 32, 33, 128]) + def test_bs_boundary_byte_equal_sweep(self, bs: int) -> None: + """Sweep bs boundary values around Triton block boundaries; assert Triton vs ref byte-equal.""" + req_pool_indices = list(range(1, bs + 1)) + prefix_lens = [10] * bs + extend_seq_lens = [1] * bs + max_seq_len = 32 + req_to_token = make_req_to_token( + kind="linear", max_reqs=bs + 1, max_seq_len=max_seq_len, device=_DEVICE + ) + + total_verify = sum(min(p, max_seq_len) for p in prefix_lens) + plans = _plan_pair( + verify_capacity=max(total_verify + 64, 256), write_req_capacity=bs + 4 + ) + run_plan_diff( + plan_pair=plans, + req_pool_indices=_tensor(req_pool_indices), + prefix_lens=_tensor(prefix_lens), + extend_seq_lens=_tensor(extend_seq_lens), + req_to_token=req_to_token, + extras=empty_extras(), + ) + + @pytest.mark.parametrize("prefix_val", [0, 1, 127, 128, 129, 4096]) + def test_prefix_lens_boundary_byte_equal_sweep(self, prefix_val: int) -> None: + """Sweep prefix_lens boundary values; assert Triton vs ref byte-equal.""" + max_seq_len = max(prefix_val + 4, 256) + req_to_token = make_req_to_token( + kind="linear", max_reqs=4, max_seq_len=max_seq_len, device=_DEVICE + ) + total_verify = prefix_val + 10 + plans = _plan_pair( + verify_capacity=max(total_verify + 64, 256), write_req_capacity=8 + ) + run_plan_diff( + plan_pair=plans, + req_pool_indices=_tensor([1, 2]), + prefix_lens=_tensor([prefix_val, 10]), + extend_seq_lens=_tensor([1, 1]), + req_to_token=req_to_token, + extras=empty_extras(), + ) + + @pytest.mark.parametrize("extend_val", [1, 128, 4096]) + def test_extend_seq_lens_boundary_byte_equal_sweep(self, extend_val: int) -> None: + """Sweep extend_seq_lens boundary values; assert Triton vs ref byte-equal.""" + max_seq_len = max(extend_val + 4, 64) + req_to_token = make_req_to_token( + kind="linear", max_reqs=4, max_seq_len=max_seq_len, device=_DEVICE + ) + plans = _plan_pair(verify_capacity=64, write_req_capacity=4) + run_plan_diff( + plan_pair=plans, + req_pool_indices=_tensor([1]), + prefix_lens=_tensor([0]), + extend_seq_lens=_tensor([extend_val]), + req_to_token=req_to_token, + extras=empty_extras(), + ) + + +class TestExpectedTokenPool: + """Cover the optional ``req_to_verify_expected_tokens`` pool input and + ``kv_token_id_vs_position_offset`` shift in the plan_entries kernel. + + Pool dtype is int32 with layout ``[max_reqs, pool_max_context_len]``. + For each verify entry the kernel gathers ``expected_input_id = + req_to_verify_expected_tokens[rp, position + offset]`` and writes ``-1`` as + a sentinel when the pool is absent or the gather index is out of range. + """ + + _MAX_SEQ_LEN = 16 + _MAX_REQS = 4 + _POOL_COLS = 16 + _DEFAULT = object() + + def setup_method(self) -> None: + self.req_to_token = make_req_to_token( + kind="linear", + max_reqs=self._MAX_REQS, + max_seq_len=self._MAX_SEQ_LEN, + device=_DEVICE, + ) + self.pool = self._make_pool(fill_fn=lambda rp, pos: rp * 1000 + pos) + + def _make_pool( + self, *, fill_fn, pool_max_context_len: int | None = None + ) -> torch.Tensor: + cols = self._POOL_COLS if pool_max_context_len is None else pool_max_context_len + pool = torch.full( + (self._MAX_REQS, cols), -999, dtype=torch.int32, device=_DEVICE + ) + for rp in range(self._MAX_REQS): + for pos in range(cols): + pool[rp, pos] = fill_fn(rp, pos) + return pool + + def _run( + self, + *, + req_pool_indices: torch.Tensor, + prefix_lens: torch.Tensor, + pool: object = _DEFAULT, + offset: int = 0, + swa_window_size: int = 0, + full_to_swa_index_mapping: torch.Tensor | None = None, + ) -> torch.Tensor: + """Run plan_diff with the per-class fixed inputs; return the expected_tokens slice up to verify_num_valid.""" + bs = int(req_pool_indices.shape[0]) + pool_arg = self.pool if pool is self._DEFAULT else pool + plans = _plan_pair(verify_capacity=64, write_req_capacity=4) + run_plan_diff( + plan_pair=plans, + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + extend_seq_lens=_tensor([1] * bs), + req_to_token=self.req_to_token, + extras=empty_extras(), + swa_window_size=swa_window_size, + full_to_swa_index_mapping=full_to_swa_index_mapping, + req_to_verify_expected_tokens=pool_arg, + kv_token_id_vs_position_offset=offset, + ) + triton_v = plans[0][0] + n_valid = int(triton_v.verify_num_valid[0].item()) + return triton_v.verify_expected_tokens[:n_valid] + + @staticmethod + def _expected(values: list[int]) -> torch.Tensor: + return torch.tensor(values, dtype=torch.int64, device=_DEVICE) + + def test_pool_disabled_writes_minus_one_sentinel(self) -> None: + """pool=None default path: every verify entry's expected_token slot is -1.""" + got = self._run( + req_pool_indices=_tensor([1, 2]), + prefix_lens=_tensor([3, 5]), + pool=None, + ) + assert got.shape[0] == 8 + assert torch.equal(got, self._expected([-1] * 8)) + + def test_pool_enabled_target_offset_0_byte_equal(self) -> None: + """offset=0 (target pool): expected_token[i] == pool[rp, position[i]].""" + got = self._run( + req_pool_indices=_tensor([1, 2]), + prefix_lens=_tensor([3, 5]), + ) + assert torch.equal( + got, + self._expected( + [ + rp * 1000 + pos + for rp, plen in [(1, 3), (2, 5)] + for pos in range(plen) + ] + ), + ) + + def test_pool_enabled_eagle_offset_plus_1_byte_equal(self) -> None: + """offset=+1 (EAGLE draft): expected_token[i] == pool[rp, position[i] + 1].""" + got = self._run( + req_pool_indices=_tensor([1, 2]), + prefix_lens=_tensor([3, 5]), + offset=1, + ) + assert torch.equal( + got, + self._expected( + [ + rp * 1000 + pos + 1 + for rp, plen in [(1, 3), (2, 5)] + for pos in range(plen) + ] + ), + ) + + def test_pool_oob_above_size0_writes_sentinel(self) -> None: + """positions whose ``position + offset`` exceed pool_max_context_len get -1; in-range slots stay correct.""" + pool_cols = 4 + small_pool = self._make_pool( + fill_fn=lambda rp, pos: rp * 1000 + pos, pool_max_context_len=pool_cols + ) + got = self._run( + req_pool_indices=_tensor([1]), + prefix_lens=_tensor([6]), + pool=small_pool, + ) + assert torch.equal( + got, + self._expected([1000 + pos if pos < pool_cols else -1 for pos in range(6)]), + ) + + def test_pool_oob_offset_plus_1_byte_equal_triggers_sentinel(self) -> None: + """offset=+1 path that pushes the last entry past pool cols still byte-equals the ref (sentinel scatter).""" + pool_cols = 4 + small_pool = self._make_pool( + fill_fn=lambda rp, pos: rp * 1000 + pos, pool_max_context_len=pool_cols + ) + got = self._run( + req_pool_indices=_tensor([1]), + prefix_lens=_tensor([4]), + pool=small_pool, + offset=1, + ) + assert torch.equal( + got, + self._expected( + [1000 + pos + 1 if pos + 1 < pool_cols else -1 for pos in range(4)] + ), + ) + + +class TestExpectedTokenPoolValidLens: + _MAX_SEQ_LEN = 16 + _MAX_REQS = 4 + _POOL_COLS = 16 + + def setup_method(self) -> None: + self.req_to_token = make_req_to_token( + kind="linear", + max_reqs=self._MAX_REQS, + max_seq_len=self._MAX_SEQ_LEN, + device=_DEVICE, + ) + self.pool = self._make_pool(fill_fn=lambda rp, pos: rp * 1000 + pos) + + def _make_pool(self, *, fill_fn) -> torch.Tensor: + pool = torch.full( + (self._MAX_REQS, self._POOL_COLS), -999, dtype=torch.int32, device=_DEVICE + ) + for rp in range(self._MAX_REQS): + for pos in range(self._POOL_COLS): + pool[rp, pos] = fill_fn(rp, pos) + return pool + + def _run( + self, + *, + req_pool_indices: torch.Tensor, + prefix_lens: torch.Tensor, + valid_lens: torch.Tensor, + pool: torch.Tensor | None = None, + offset: int = 0, + swa_window_size: int = 0, + full_to_swa_index_mapping: torch.Tensor | None = None, + ) -> torch.Tensor: + """Run plan_diff with the per-class fixed inputs; return the expected_tokens slice up to verify_num_valid.""" + bs = int(req_pool_indices.shape[0]) + plans = _plan_pair(verify_capacity=64, write_req_capacity=4) + run_plan_diff( + plan_pair=plans, + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + extend_seq_lens=_tensor([1] * bs), + req_to_token=self.req_to_token, + extras=empty_extras(), + swa_window_size=swa_window_size, + full_to_swa_index_mapping=full_to_swa_index_mapping, + req_to_verify_expected_tokens=self.pool if pool is None else pool, + req_to_verify_expected_tokens_valid_lens=valid_lens, + kv_token_id_vs_position_offset=offset, + ) + triton_v = plans[0][0] + n_valid = int(triton_v.verify_num_valid[0].item()) + return triton_v.verify_expected_tokens[:n_valid] + + @staticmethod + def _expected(values: list[int]) -> torch.Tensor: + return torch.tensor(values, dtype=torch.int64, device=_DEVICE) + + def test_valid_lens_boundary_emits_sentinel_at_limit(self) -> None: + """sot_pos == valid_lens[r] is OUT of range; the kernel must emit -1 even though the pool has a real value at that slot.""" + got = self._run( + req_pool_indices=_tensor([1]), + prefix_lens=_tensor([3]), + valid_lens=_tensor([2]), + ) + assert torch.equal(got, self._expected([1000, 1001, -1])) + + def test_valid_lens_within_limit_reads_pool(self) -> None: + """sot_pos == valid_lens[r] - 1 is IN range; the kernel must gather the pool value, not -1.""" + got = self._run( + req_pool_indices=_tensor([1]), + prefix_lens=_tensor([3]), + valid_lens=_tensor([3]), + ) + assert torch.equal(got, self._expected([1000, 1001, 1002])) + + def test_valid_lens_mixed_across_reqs(self) -> None: + """Per-req different valid_lens in one batch: each req's gather is bounded by its own lens, not the batch max.""" + got = self._run( + req_pool_indices=_tensor([1, 2]), + prefix_lens=_tensor([3, 3]), + valid_lens=_tensor([2, 4]), + ) + assert torch.equal(got, self._expected([1000, 1001, -1, 2000, 2001, 2002])) + + def test_valid_lens_zero_emits_all_sentinel_for_that_req(self) -> None: + """valid_lens[r] == 0 disables every gather for req r regardless of pool content.""" + got = self._run( + req_pool_indices=_tensor([1]), + prefix_lens=_tensor([3]), + valid_lens=_tensor([0]), + ) + assert torch.equal(got, self._expected([-1, -1, -1])) + + def test_valid_lens_masks_stale_pool_data_above_bound(self) -> None: + """Pool has realistic-looking values past valid_lens (the recycled-slot motivation): kernel still emits -1 for them.""" + # Positions 2..15 carry a longer previous owner's leftover token; the bound must hide them. + stale_pool = self._make_pool( + fill_fn=lambda rp, pos: 7777 if pos >= 2 else (rp * 1000 + pos), + ) + got = self._run( + req_pool_indices=_tensor([1]), + prefix_lens=_tensor([5]), + valid_lens=_tensor([2]), + pool=stale_pool, + ) + assert torch.equal(got, self._expected([1000, 1001, -1, -1, -1])) + + def test_valid_lens_with_offset_plus_1_bounds_after_shift(self) -> None: + """``sot_pos = position + offset`` is compared against valid_lens; the offset shifts before the bound check.""" + # sot_pos for positions 0,1,2 is 1,2,3. valid_lens=2 → only sot_pos=1 reads pool. + got = self._run( + req_pool_indices=_tensor([1]), + prefix_lens=_tensor([3]), + valid_lens=_tensor([2]), + offset=1, + ) + assert torch.equal(got, self._expected([1001, -1, -1])) + + def test_valid_lens_applies_under_swa_window(self) -> None: + """SWA-windowed verify entries are bounded by valid_lens the same way as the FULL pool path.""" + # prefix=5, swa_window=3 → entries cover positions 2,3,4. valid_lens=4 admits pos 2,3 only. + got = self._run( + req_pool_indices=_tensor([1]), + prefix_lens=_tensor([5]), + valid_lens=_tensor([4]), + swa_window_size=3, + full_to_swa_index_mapping=make_lut( + kind="identity", + pool_size=self._MAX_REQS * self._MAX_SEQ_LEN, + device=_DEVICE, + ), + ) + assert torch.equal(got, self._expected([1002, 1003, -1])) + + def test_pool_set_but_valid_lens_missing_raises(self) -> None: + """One-way contract: passing the pool without per-req valid_lens is rejected at the Python wrapper.""" + triton_v, triton_w = _plan_pair(verify_capacity=64, write_req_capacity=4)[0] + with pytest.raises( + ValueError, match="req_to_verify_expected_tokens_valid_lens" + ): + launch_canary_plan_kernels( + verify_plan_out=triton_v, + write_plan_out=triton_w, + req_pool_indices=_tensor([1]), + prefix_lens=_tensor([3]), + extend_seq_lens=_tensor([1]), + req_to_token=self.req_to_token, + swa_window_size=0, + full_to_swa_index_mapping=None, + verify_capacity=int(triton_v.verify_slot_indices.shape[0]), + req_to_verify_expected_tokens=self.pool, + req_to_verify_expected_tokens_valid_lens=None, + kv_token_id_vs_position_offset=0, + ) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) diff --git a/python/sglang/jit_kernel/tests/kv_canary/test_verify_fuzz.py b/python/sglang/jit_kernel/tests/kv_canary/test_verify_fuzz.py new file mode 100644 index 000000000..22ecebd65 --- /dev/null +++ b/python/sglang/jit_kernel/tests/kv_canary/test_verify_fuzz.py @@ -0,0 +1,178 @@ +from __future__ import annotations + +import random +from dataclasses import dataclass + +import pytest +import torch + +from sglang.jit_kernel.kv_canary import consts +from sglang.jit_kernel.kv_canary.verify import ( + CanaryLaunchTag, + VerifyPlan, +) +from sglang.jit_kernel.tests.kv_canary._canary_helpers import ( + FakeViolationLog, + make_canary_buf, + make_log_pair, + make_verify_plan_pair, + stamp_clean_chain, +) +from sglang.jit_kernel.tests.kv_canary._differential import _run_both_verify +from sglang.jit_kernel.tests.kv_canary._fuzz_driver import ( + FUZZ_SEEDS_PR, + run_fuzz_combo, +) +from sglang.jit_kernel.tests.kv_canary._invariants import VerifyInvariants +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large") + + +_DEVICE = torch.device("cuda") + +_FUZZ_ITER_PER_SEED = 30 + + +@dataclass(frozen=True, slots=True, kw_only=True) +class VerifyFuzzInputs: + cuda_canary_buf: torch.Tensor + ref_canary_buf: torch.Tensor + plan_cuda: VerifyPlan + plan_ref: VerifyPlan + kernel_kind: CanaryLaunchTag + ring_capacity: int + check_verify_expected_token: bool + + +def _draw_random_verify_inputs(rng: random.Random) -> VerifyFuzzInputs: + kernel_kind = rng.choice(list(CanaryLaunchTag)) + plan_size = rng.randint(0, 32) + num_slots = max(plan_size + 8, 16) + ring_capacity = rng.choice([16, 64, 256]) + + cuda_buf = make_canary_buf( + num_slots=num_slots, slot_stride_bytes=32, device=_DEVICE + ) + ref_buf = cuda_buf.clone() + + slot_universe = list(range(1, num_slots)) + rng.shuffle(slot_universe) + slot_indices = slot_universe[:plan_size] + tokens = [rng.randint(0, 0xFFFFFFFF) for _ in range(plan_size)] + positions = list(range(plan_size)) + prev_slot_indices: list[int] = [] + for i in range(plan_size): + if i == 0: + prev_slot_indices.append(-1) + else: + prev_slot_indices.append(slot_indices[i - 1]) + + if plan_size > 0: + stamp_clean_chain( + cuda_buf=cuda_buf, + ref_buf=ref_buf, + slot_indices=slot_indices, + tokens=tokens, + positions=positions, + ) + + # Inject prev_slot == TOKEN_TO_KV_SLOT_PADDING into ~15% of entries so the differential + # harness exercises the chain-check-skip branch (added for SWA-evicted ancestor handling). + # Done AFTER stamp_clean_chain so the stored prev_hash on those slots is still chain-clean; + # the kernel must rely on prev_slot==padding (not on stored hash) to decide whether to skip. + for i in range(plan_size): + if rng.random() < 0.15: + prev_slot_indices[i] = consts.TOKEN_TO_KV_SLOT_PADDING + + check_verify_expected_token = rng.random() < 0.5 + expected_input_ids: list[int] = [] + for i in range(plan_size): + # Always pick a value; with check=False the kernel must not deref this column. + if rng.random() < 0.3: + expected_input_ids.append(-1) + elif rng.random() < 0.5: + expected_input_ids.append(int(tokens[i])) + else: + mutated = (int(tokens[i]) ^ 0x1) & 0xFFFFFFFF + expected_input_ids.append(mutated) + + plan_cuda, plan_ref = make_verify_plan_pair( + slot_indices=slot_indices, + positions=positions, + prev_slot_indices=prev_slot_indices, + expected_input_ids=expected_input_ids if plan_size > 0 else None, + capacity=max(plan_size, 1), + device=_DEVICE, + ) + + return VerifyFuzzInputs( + cuda_canary_buf=cuda_buf, + ref_canary_buf=ref_buf, + plan_cuda=plan_cuda, + plan_ref=plan_ref, + kernel_kind=kernel_kind, + ring_capacity=ring_capacity, + check_verify_expected_token=check_verify_expected_token, + ) + + +def _run_one(inputs: VerifyFuzzInputs) -> None: + cuda_buf_before = inputs.cuda_canary_buf.clone() + cuda_log, ref_log = make_log_pair(capacity=inputs.ring_capacity, device=_DEVICE) + log_before = FakeViolationLog.allocate( + capacity=inputs.ring_capacity, device=_DEVICE + ) + _run_both_verify( + cuda_canary_buf=inputs.cuda_canary_buf, + ref_canary_buf=inputs.ref_canary_buf, + plan_cuda=inputs.plan_cuda, + plan_ref=inputs.plan_ref, + cuda_log=cuda_log, + ref_log=ref_log, + kernel_kind=inputs.kernel_kind, + assert_equal=False, + check_verify_expected_token=inputs.check_verify_expected_token, + ) + assert int(cuda_log.kernel_run_counter[0].item()) == int( + ref_log.kernel_run_counter[0].item() + ) + assert int(cuda_log.slot_run_counter[0].item()) == int( + ref_log.slot_run_counter[0].item() + ) + assert int(cuda_log.write_index[0].item()) == int(ref_log.write_index[0].item()) + VerifyInvariants.assert_all( + canary_buf_before=cuda_buf_before, + canary_buf_after=inputs.cuda_canary_buf, + log_before=log_before, + log_after=cuda_log, + plan=inputs.plan_cuda, + kernel_kind=inputs.kernel_kind, + ) + + +def _summarize(inputs: VerifyFuzzInputs) -> str: + n_active = int(inputs.plan_cuda.verify_num_valid[0].item()) + return ( + f"plan_size={n_active} kind={inputs.kernel_kind.name} " + f"ring={inputs.ring_capacity} " + f"check_token={inputs.check_verify_expected_token}" + ) + + +@pytest.mark.parametrize("seed", FUZZ_SEEDS_PR) +def test_verify_fuzz_full_combo(seed: int) -> None: + """Multi-dim verify fuzzer: random kernel kind × plan size × ring capacity × N iters, byte-equal.""" + run_fuzz_combo( + seed, + draw_fn=_draw_random_verify_inputs, + run_one_fn=_run_one, + summarize_fn=_summarize, + n_iter=_FUZZ_ITER_PER_SEED, + ) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) diff --git a/python/sglang/jit_kernel/tests/kv_canary/test_verify_hand.py b/python/sglang/jit_kernel/tests/kv_canary/test_verify_hand.py new file mode 100644 index 000000000..9db088701 --- /dev/null +++ b/python/sglang/jit_kernel/tests/kv_canary/test_verify_hand.py @@ -0,0 +1,1266 @@ +from __future__ import annotations + +import random +import struct +from dataclasses import dataclass +from typing import Callable + +import pytest +import torch + +from sglang.jit_kernel.kv_canary import consts +from sglang.jit_kernel.kv_canary.consts import splitmix64, splitmix64_mix3 +from sglang.jit_kernel.kv_canary.verify import ( + CanaryLaunchTag, + VerifyOrWriteContext, + VerifyPlan, + launch_canary_verify_kernel, +) +from sglang.jit_kernel.kv_canary.verify_ref import ( + launch_canary_verify_kernel_torch_reference, +) +from sglang.jit_kernel.tests.kv_canary._canary_helpers import ( + FakeViolationLog, + assert_only_bits_set, + chain_anchor_signed, + make_canary_buf, + make_canary_buf_pair, + make_log_pair, + make_verify_plan, + make_verify_plan_pair, + read_slot_fields, + stamp_clean_chain, + stamp_pair, + to_signed_int64, + write_slot_fields, +) +from sglang.jit_kernel.tests.kv_canary._differential import ( + _run_both_verify, + run_verify_diff, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large") + + +_DEVICE = torch.device("cuda") + + +# --------------------------------------------------------------------------- +# Shared per-test scaffolding helpers. +# --------------------------------------------------------------------------- + + +def _buf_pair( + num_slots: int = 16, slot_stride_bytes: int = 32 +) -> tuple[torch.Tensor, torch.Tensor]: + return make_canary_buf_pair( + num_slots=num_slots, slot_stride_bytes=slot_stride_bytes, device=_DEVICE + ) + + +def _stamp_head( + buf_pair: tuple[torch.Tensor, torch.Tensor], + *, + slot_idx: int, + token: int = 42, + position: int = 0, + prev_hash: int | None = None, +) -> None: + """``stamp_pair`` with ``prev_hash`` defaulting to ``chain_anchor_signed()`` (the chain-head value).""" + stamp_pair( + buf_pair, + slot_idx=slot_idx, + token=token, + position=position, + prev_hash=chain_anchor_signed() if prev_hash is None else prev_hash, + ) + + +def _plan_pair_single( + *, + slot_idx: int, + position: int, + prev_slot_idx: int = -1, + expected_input_id: int | None = None, + capacity: int | None = None, +) -> tuple[VerifyPlan, VerifyPlan]: + """Single-entry verify plan pair, with optional expected_input_id (None → default sentinel).""" + expected = None if expected_input_id is None else [expected_input_id] + return make_verify_plan_pair( + slot_indices=[slot_idx], + positions=[position], + prev_slot_indices=[prev_slot_idx], + expected_input_ids=expected, + capacity=capacity, + device=_DEVICE, + ) + + +def _n_violations(log: FakeViolationLog) -> int: + return int(log.write_index[0].item()) + + +def _fail_bits(log: FakeViolationLog, row: int = 0) -> int: + return int(log.ring[row, consts.VIOLATION_FIELD_FAIL_REASON_BITS].item()) + + +def _run_both_verify_no_rkv( + *, + buf_pair: tuple[torch.Tensor, torch.Tensor], + plan_pair: tuple[VerifyPlan, VerifyPlan], + cuda_log: FakeViolationLog, + ref_log: FakeViolationLog, + assert_equal: bool = True, + kernel_kind: CanaryLaunchTag = CanaryLaunchTag.HEAD_K_FULL, +) -> None: + """``_run_both_verify`` — the most common in-place verify run.""" + _run_both_verify( + cuda_canary_buf=buf_pair[0], + ref_canary_buf=buf_pair[1], + plan_cuda=plan_pair[0], + plan_ref=plan_pair[1], + cuda_log=cuda_log, + ref_log=ref_log, + kernel_kind=kernel_kind, + assert_equal=assert_equal, + ) + + +@dataclass(frozen=True, slots=True, kw_only=True) +class _VerifySingleSlotInput: + token: int = 42 + position: int = 0 + stored_prev_hash_signed: int + + +def _run_verify_single_slot_byte_equal(case: _VerifySingleSlotInput) -> None: + buf_pair = _buf_pair() + stamp_pair( + buf_pair, + slot_idx=1, + token=case.token, + position=case.position, + prev_hash=case.stored_prev_hash_signed, + ) + plan_pair = _plan_pair_single(slot_idx=1, position=case.position) + run_verify_diff( + buf_pair=buf_pair, + plan_pair=plan_pair, + ) + + +# --------------------------------------------------------------------------- +# Kernel-contract invariants. +# --------------------------------------------------------------------------- + + +class TestChain: + def test_chain_head_anchor(self) -> None: + """``prev_slot_idx == -1`` → kernel uses ``splitmix64(consts.CANARY_CHAIN_ANCHOR)`` as the expected prev_hash.""" + # Step 1: stamp slot 5 such that stored.prev_hash already equals splitmix64(consts.CANARY_CHAIN_ANCHOR). + buf_pair = _buf_pair() + _stamp_head(buf_pair, slot_idx=5) + + # Step 2: a single-entry plan with prev_slot_idx = -1 should record no violation. + plan_pair = _plan_pair_single(slot_idx=5, position=0) + cuda_log, _ = run_verify_diff(buf_pair=buf_pair, plan_pair=plan_pair) + assert _n_violations(cuda_log) == 0 + + def test_chain_link_byte_equal_5_step(self) -> None: + """5-step chain, CUDA vs ref byte-equal across ring / counters / canary_buf (read-only).""" + cuda_buf, ref_buf = _buf_pair() + slot_indices = [1, 2, 3, 4, 5] + tokens = [11, 22, 33, 44, 55] + positions = [0, 1, 2, 3, 4] + stamp_clean_chain( + cuda_buf=cuda_buf, + ref_buf=ref_buf, + tokens=tokens, + positions=positions, + slot_indices=slot_indices, + ) + plan_pair = make_verify_plan_pair( + slot_indices=slot_indices, + positions=positions, + prev_slot_indices=[-1, 1, 2, 3, 4], + device=_DEVICE, + ) + cuda_log, _ = run_verify_diff(buf_pair=(cuda_buf, ref_buf), plan_pair=plan_pair) + assert _n_violations(cuda_log) == 0 + + def test_chain_link_byte_equal_5_step_hardcoded(self) -> None: + """5-step chain with hand-computed splitmix64 expected sequence; defends against ref + CUDA co-drift.""" + tokens = [101, 202, 303, 404, 505] + positions = [0, 1, 2, 3, 4] + slot_indices = [1, 2, 3, 4, 5] + + # Step 1: compute the expected stored prev_hash sequence in pure Python via splitmix64. + expected_prev_hashes_u64: list[int] = [] + running = splitmix64(consts.CANARY_CHAIN_ANCHOR) + for token, position in zip(tokens, positions): + expected_prev_hashes_u64.append(running) + running = splitmix64_mix3(running, token, position) + expected_prev_hashes_signed = [ + to_signed_int64(h) for h in expected_prev_hashes_u64 + ] + + # Step 2: stamp each slot manually with the hardcoded expected prev_hash. + buf_pair = _buf_pair() + cuda_buf, _ = buf_pair + for slot_idx, token, position, prev_hash in zip( + slot_indices, tokens, positions, expected_prev_hashes_signed + ): + stamp_pair( + buf_pair, + slot_idx=slot_idx, + token=token, + position=position, + prev_hash=prev_hash, + ) + + # Step 3: verify the 5-step chain — no violation expected and the ref vs CUDA state byte-equal. + plan_pair = make_verify_plan_pair( + slot_indices=slot_indices, + positions=positions, + prev_slot_indices=[-1, 1, 2, 3, 4], + device=_DEVICE, + ) + cuda_log, _ = run_verify_diff(buf_pair=buf_pair, plan_pair=plan_pair) + + assert _n_violations(cuda_log) == 0 + + # Step 4: also independently confirm the *stored* prev_hash at each slot matches the hardcoded sequence. + for slot_idx, expected_signed in zip(slot_indices, expected_prev_hashes_signed): + _, _, stored_prev_hash, _ = read_slot_fields( + canary_buf=cuda_buf, slot_idx=slot_idx + ) + assert stored_prev_hash == expected_signed + + def test_chain_advance_formula_matches_spec(self) -> None: + """Ref impl agrees with the chained splitmix64 chain-step formula. + + The chain step folds each of the 3 inputs into the accumulator sequentially via + ``acc = splitmix64(acc ^ next)``, starting from ``splitmix64(prev_hash)``. ``splitmix64_mix3`` + must produce the same result as the explicit chain. ``real_kv_hash`` is intentionally NOT + part of the chain hash — see ``compute_slot_hash`` for the radix-folding rationale. + """ + cases = [ + (consts.CANARY_CHAIN_ANCHOR, 0, 0), + (0x1234567890ABCDEF, 100, 5), + (0, 0xFFFF, 0x7FFFFFFF), + (0x123, 1, 1), + (0xFFFFFFFFFFFFFFFF, 0xFFFF, 0xFFFF), + ] + for prev_hash, token, position in cases: + u64_mask = (1 << 64) - 1 + h = splitmix64(prev_hash & u64_mask) + h = splitmix64(h ^ (token & u64_mask)) + expected = splitmix64(h ^ (position & u64_mask)) + + actual = ( + splitmix64_mix3( + prev_hash & u64_mask, token & u64_mask, position & u64_mask + ) + & u64_mask + ) + assert actual == expected, ( + f"chain advance mismatch: prev={prev_hash:#x} token={token:#x} pos={position:#x} " + f"expected={expected:#x} actual={actual:#x}" + ) + + def test_chain_head_anchored_on_constant(self) -> None: + """prev_slot==-1 + stored prev_hash != splitmix64(ANCHOR) → CHAIN_HASH bit set.""" + buf_pair = _buf_pair() + slot_idx = 5 + _stamp_head(buf_pair, slot_idx=slot_idx, prev_hash=to_signed_int64(0xDEADBEEF)) + + plan_pair = _plan_pair_single(slot_idx=slot_idx, position=0) + cuda_log, _ = run_verify_diff(buf_pair=buf_pair, plan_pair=plan_pair) + assert _n_violations(cuda_log) == 1 + assert _fail_bits(cuda_log) & consts.FailReason.VERIFY_CHAIN_HASH_MISMATCH + + def test_chain_head_prev_hash_equals_splitmix64_anchor_random_50(self) -> None: + random.seed(0) + expected_prev_hash_signed = to_signed_int64( + splitmix64(consts.CANARY_CHAIN_ANCHOR) + ) + + for _ in range(50): + token = random.randint(0, 0x7FFFFFFF) + position = random.randint(0, 0x7FFFFFFF) + slot_idx = random.randint(0, 15) + + buf_pair = _buf_pair() + _stamp_head( + buf_pair, + slot_idx=slot_idx, + token=token, + position=position, + prev_hash=expected_prev_hash_signed, + ) + + plan_pair = _plan_pair_single(slot_idx=slot_idx, position=position) + cuda_log, _ = run_verify_diff( + buf_pair=buf_pair, plan_pair=plan_pair, assert_equal=False + ) + + assert ( + _n_violations(cuda_log) == 0 + ), f"unexpected violation at iteration token={token} position={position} slot={slot_idx}" + + def test_prev_slot_padding_skips_chain_check_arbitrary_stored_hash(self) -> None: + """prev_slot_idx == TOKEN_TO_KV_SLOT_PADDING → chain check is skipped, regardless of stored chain hash.""" + buf_pair = _buf_pair() + # Stamp slot 5 with an arbitrary (non-anchor, non-derivable) prev_hash. Without the skip, + # the kernel would compute expected_chain_hash from slot 0's canary (all zeros = 0) and + # flag VERIFY_CHAIN_HASH_MISMATCH on every such row. + stamp_pair( + buf_pair, + slot_idx=5, + token=42, + position=0, + prev_hash=to_signed_int64(0xDEADBEEFCAFEBABE), + ) + + plan_pair = _plan_pair_single( + slot_idx=5, position=0, prev_slot_idx=consts.TOKEN_TO_KV_SLOT_PADDING + ) + cuda_log, _ = run_verify_diff(buf_pair=buf_pair, plan_pair=plan_pair) + assert _n_violations(cuda_log) == 0 + + def test_prev_slot_padding_does_not_mask_position_check(self) -> None: + """prev_slot == padding skips ONLY the chain check; position mismatch still fires.""" + buf_pair = _buf_pair() + stamp_pair( + buf_pair, + slot_idx=7, + token=11, + position=0, + prev_hash=to_signed_int64(0x12345678), + ) + + # Plan claims position 99 — chain check is skipped (prev=padding) but position must still fire. + plan_pair = _plan_pair_single( + slot_idx=7, position=99, prev_slot_idx=consts.TOKEN_TO_KV_SLOT_PADDING + ) + cuda_log, _ = run_verify_diff(buf_pair=buf_pair, plan_pair=plan_pair) + + assert _n_violations(cuda_log) == 1 + assert_only_bits_set( + _fail_bits(cuda_log), consts.FailReason.VERIFY_POSITION_MISMATCH + ) + + +class TestViolationField: + def test_violation_token_mismatch(self) -> None: + """Stored token differs from a fresh write at the same slot → TOKEN-side accounting via the chain bit.""" + # Verify kernel doesn't have a TOKEN fail bit per se — token mismatch propagates into next-slot + # CHAIN_HASH mismatch. Inject token corruption at slot 2 and verify slot 3 sees CHAIN_HASH bit set. + cuda_buf, ref_buf = _buf_pair() + slot_indices = [1, 2, 3] + tokens = [100, 200, 300] + positions = [0, 1, 2] + stamp_clean_chain( + cuda_buf=cuda_buf, + ref_buf=ref_buf, + tokens=tokens, + positions=positions, + slot_indices=slot_indices, + ) + + # Step: corrupt the stored token at slot 2 in both buffers — chain hash propagates downstream. + stamp_pair( + (cuda_buf, ref_buf), + slot_idx=2, + token=999, + position=1, + prev_hash=0, + ) + + plan_pair = make_verify_plan_pair( + slot_indices=[3], positions=[2], prev_slot_indices=[2], device=_DEVICE + ) + cuda_log, _ = run_verify_diff(buf_pair=(cuda_buf, ref_buf), plan_pair=plan_pair) + + assert _n_violations(cuda_log) == 1 + assert_only_bits_set( + _fail_bits(cuda_log), consts.FailReason.VERIFY_CHAIN_HASH_MISMATCH + ) + + def test_violation_position_mismatch(self) -> None: + """Stored position differs from what the slot's chain reconstruction would yield → POSITION bit.""" + # Stamp slot 7 with a valid head chain but stored position = 0; ask verify to expect position 5. + buf_pair = _buf_pair() + _stamp_head(buf_pair, slot_idx=7) + + plan_pair = _plan_pair_single(slot_idx=7, position=5) + cuda_log, _ = run_verify_diff(buf_pair=buf_pair, plan_pair=plan_pair) + + assert _n_violations(cuda_log) == 1 + assert_only_bits_set( + _fail_bits(cuda_log), consts.FailReason.VERIFY_POSITION_MISMATCH + ) + + def test_violation_position_diverges_from_plan(self) -> None: + """Plan-supplied position contradicts stored position → POSITION bit (verify trusts plan, not +1).""" + # Step: a clean chain head with stored position 0; plan claims position 99 — kernel must flag POSITION. + buf_pair = _buf_pair() + _stamp_head(buf_pair, slot_idx=3, token=11) + + plan_pair = _plan_pair_single(slot_idx=3, position=99) + cuda_log, _ = run_verify_diff(buf_pair=buf_pair, plan_pair=plan_pair) + + assert_only_bits_set( + _fail_bits(cuda_log), consts.FailReason.VERIFY_POSITION_MISMATCH + ) + + def test_violation_prev_hash_mismatch(self) -> None: + """Stored prev_hash differs from predecessor-derived expectation → CHAIN_HASH bit.""" + cuda_buf, ref_buf = _buf_pair() + slot_indices = [1, 2] + tokens = [10, 20] + positions = [0, 1] + stamp_clean_chain( + cuda_buf=cuda_buf, + ref_buf=ref_buf, + tokens=tokens, + positions=positions, + slot_indices=slot_indices, + ) + + # Step: corrupt slot 2's stored prev_hash with a bogus signed int64. + stamp_pair( + (cuda_buf, ref_buf), + slot_idx=2, + token=20, + position=1, + prev_hash=0x1234567812345678, + ) + + plan_pair = make_verify_plan_pair( + slot_indices=[2], positions=[1], prev_slot_indices=[1], device=_DEVICE + ) + cuda_log, _ = run_verify_diff(buf_pair=(cuda_buf, ref_buf), plan_pair=plan_pair) + + assert_only_bits_set( + _fail_bits(cuda_log), consts.FailReason.VERIFY_CHAIN_HASH_MISMATCH + ) + + @pytest.mark.parametrize("bit_to_trigger", ["POSITION", "PREV_HASH"]) + @pytest.mark.parametrize("injection_position", ["head", "mid", "last"]) + @pytest.mark.parametrize("ring_state", ["open", "full"]) + def test_violation_bit_injection_position_ring_state_matrix( + self, + bit_to_trigger: str, + injection_position: str, + ring_state: str, + ) -> None: + """Sweep injection_position x bit_to_trigger x ring_state for verify-kernel fail-reason coverage.""" + _RING_CAPACITY = 4 + slot_indices = [1, 2, 3, 4, 5] + tokens = [11, 22, 33, 44, 55] + positions = [0, 1, 2, 3, 4] + + corruption_index = {"head": 0, "mid": 2, "last": 4}[injection_position] + corrupt_slot = slot_indices[corruption_index] + + expected_bit = { + "POSITION": consts.FailReason.VERIFY_POSITION_MISMATCH, + "PREV_HASH": consts.FailReason.VERIFY_CHAIN_HASH_MISMATCH, + }[bit_to_trigger] + + cuda_buf, ref_buf = _buf_pair() + buf_pair = (cuda_buf, ref_buf) + stamp_clean_chain( + cuda_buf=cuda_buf, + ref_buf=ref_buf, + tokens=tokens, + positions=positions, + slot_indices=slot_indices, + ) + + if bit_to_trigger == "POSITION": + stored_token, stored_pos, stored_prev, _ = read_slot_fields( + canary_buf=cuda_buf, slot_idx=corrupt_slot + ) + stamp_pair( + buf_pair, + slot_idx=corrupt_slot, + token=stored_token, + position=stored_pos + 99, + prev_hash=stored_prev, + ) + else: + stored_token, stored_pos, stored_prev, _ = read_slot_fields( + canary_buf=cuda_buf, slot_idx=corrupt_slot + ) + flipped_prev = stored_prev ^ 1 + stamp_pair( + buf_pair, + slot_idx=corrupt_slot, + token=stored_token, + position=stored_pos, + prev_hash=flipped_prev, + ) + + ring_capacity = _RING_CAPACITY + cuda_log, ref_log = make_log_pair(capacity=ring_capacity, device=_DEVICE) + if ring_state == "full": + prefill_slots = list(range(8, 8 + ring_capacity)) + for slot_idx in prefill_slots: + _stamp_head(buf_pair, slot_idx=slot_idx, token=1) + prefill_plan_pair = make_verify_plan_pair( + slot_indices=prefill_slots, + positions=[99] * ring_capacity, + prev_slot_indices=[-1] * ring_capacity, + device=_DEVICE, + ) + _run_both_verify_no_rkv( + buf_pair=buf_pair, + plan_pair=prefill_plan_pair, + cuda_log=cuda_log, + ref_log=ref_log, + assert_equal=False, + ) + assert _n_violations(cuda_log) == ring_capacity + + prev_slot_indices = [-1, 1, 2, 3, 4] + plan_cuda, plan_ref = make_verify_plan_pair( + slot_indices=slot_indices, + positions=positions, + prev_slot_indices=prev_slot_indices, + device=_DEVICE, + ) + + _run_both_verify( + cuda_canary_buf=buf_pair[0], + ref_canary_buf=buf_pair[1], + plan_cuda=plan_cuda, + plan_ref=plan_ref, + cuda_log=cuda_log, + ref_log=ref_log, + assert_equal=False, + ) + + if ring_state == "open": + write_index = _n_violations(cuda_log) + rows_stored = min(write_index, ring_capacity) + found = any( + _fail_bits(cuda_log, row_idx) & expected_bit + for row_idx in range(rows_stored) + ) + assert found, ( + f"expected bit {expected_bit:#x} not found in any ring row " + f"(bit_to_trigger={bit_to_trigger} injection_position={injection_position})" + ) + else: + assert ( + _n_violations(cuda_log) > ring_capacity + ), "write_index did not advance beyond ring_capacity after overflow" + + def test_position_mismatch_sets_position_bit_only(self) -> None: + """Plan.position != stored.position with chain hash correct → only POSITION bit set.""" + buf_pair = _buf_pair() + slot_idx = 5 + _stamp_head(buf_pair, slot_idx=slot_idx, position=10) + + plan_pair = _plan_pair_single(slot_idx=slot_idx, position=99) + cuda_log, _ = run_verify_diff(buf_pair=buf_pair, plan_pair=plan_pair) + assert _n_violations(cuda_log) == 1 + bits = _fail_bits(cuda_log) + assert ( + bits & consts.FailReason.VERIFY_POSITION_MISMATCH + ), f"expected POSITION bit, got {bits:#b}" + assert ( + bits & consts.FailReason.VERIFY_CHAIN_HASH_MISMATCH + ) == 0, f"chain hash bit unexpectedly set: {bits:#b}" + + +class TestLayoutAndScheduling: + def test_swa_translated_slot_indices(self) -> None: + """SWA-translated slots already passed in plan; verify kernel does no further translation.""" + # SWA verify plans carry pre-translated slot indices — the verify kernel never sees the FULL slot + # index again. We pre-stamp the SWA-side slot and feed it directly into the verify plan to assert no + # extra translation happens kernel-side. + buf_pair = _buf_pair() + _stamp_head(buf_pair, slot_idx=2, token=99) + + plan_pair = _plan_pair_single(slot_idx=2, position=0) + cuda_log, _ = run_verify_diff(buf_pair=buf_pair, plan_pair=plan_pair) + assert _n_violations(cuda_log) == 0 + + def test_empty_plan_no_op(self) -> None: + """``verify_num_valid = 0`` → no ring write, no slot_run_counter bump, only kernel_run_counter += 1.""" + buf_pair = _buf_pair() + plan_pair = make_verify_plan_pair( + slot_indices=[], + positions=[], + prev_slot_indices=[], + capacity=4, + device=_DEVICE, + ) + cuda_log, _ = run_verify_diff(buf_pair=buf_pair, plan_pair=plan_pair) + + assert _n_violations(cuda_log) == 0 + assert int(cuda_log.slot_run_counter[0].item()) == 0 + assert int(cuda_log.kernel_run_counter[0].item()) == 1 + + def test_slot_zero_plan_entry_is_skipped(self) -> None: + """slot 0 is reserved padding: verify skips loads/violations but still counts the submitted entry.""" + canary_buf = make_canary_buf(num_slots=16, slot_stride_bytes=32, device=_DEVICE) + write_slot_fields( + canary_buf=canary_buf, + slot_idx=0, + token=999, + position=123, + prev_hash=to_signed_int64(0xDEADBEEF), + ) + plan = make_verify_plan( + slot_indices=[0], + positions=[0], + prev_slot_indices=[-1], + device=_DEVICE, + ) + log = FakeViolationLog.allocate(capacity=8, device=_DEVICE) + + launch_canary_verify_kernel( + context=VerifyOrWriteContext( + canary_buf=canary_buf, + kernel_kind=CanaryLaunchTag.HEAD_K_FULL, + violation_ring=log.ring, + violation_write_index=log.write_index, + slot_run_counter=log.slot_run_counter, + kernel_run_counter=log.kernel_run_counter, + enable_chain_position_assert=log.enable_chain_position_assert, + ), + plan=plan, + check_verify_expected_token=True, + ) + torch.cuda.synchronize() + + assert _n_violations(log) == 0 + assert int(log.slot_run_counter[0].item()) == 1 + assert int(log.kernel_run_counter[0].item()) == 1 + + @pytest.mark.parametrize( + "runner", + [launch_canary_verify_kernel, launch_canary_verify_kernel_torch_reference], + ) + def test_disabled_plan_skips_slots_but_counts_kernel( + self, runner: Callable[..., None] + ) -> None: + """``VerifyPlan.enable = 0`` skips active entries while still marking the verify launch as run.""" + canary_buf = make_canary_buf(num_slots=16, slot_stride_bytes=32, device=_DEVICE) + slot_idx = 5 + write_slot_fields( + canary_buf=canary_buf, + slot_idx=slot_idx, + token=42, + position=1, + prev_hash=to_signed_int64(0x1234), + ) + plan = make_verify_plan( + slot_indices=[slot_idx], + positions=[0], + prev_slot_indices=[-1], + device=_DEVICE, + ) + plan.enable[0] = 0 + + log = FakeViolationLog.allocate(capacity=8, device=_DEVICE) + log.ring.fill_(-777) + log.write_index[0] = 3 + log.slot_run_counter[0] = 11 + log.kernel_run_counter[0] = 13 + ring_before = log.ring.clone() + write_index_before = log.write_index.clone() + slot_run_before = log.slot_run_counter.clone() + kernel_run_before = log.kernel_run_counter.clone() + + runner( + context=VerifyOrWriteContext( + canary_buf=canary_buf, + kernel_kind=CanaryLaunchTag.HEAD_K_FULL, + violation_ring=log.ring, + violation_write_index=log.write_index, + slot_run_counter=log.slot_run_counter, + kernel_run_counter=log.kernel_run_counter, + enable_chain_position_assert=log.enable_chain_position_assert, + ), + plan=plan, + check_verify_expected_token=True, + ) + if runner is launch_canary_verify_kernel: + torch.cuda.synchronize() + + assert torch.equal(log.ring, ring_before) + assert torch.equal(log.write_index, write_index_before) + assert torch.equal(log.slot_run_counter, slot_run_before) + assert ( + int(log.kernel_run_counter[0].item()) + == int(kernel_run_before[0].item()) + 1 + ) + + def test_empty_plan_keeps_slot_counter_unchanged(self) -> None: + buf_pair = _buf_pair(num_slots=8) + _stamp_head(buf_pair, slot_idx=1, token=7) + + nonempty_plan_pair = _plan_pair_single(slot_idx=1, position=0) + empty_plan_pair = make_verify_plan_pair( + slot_indices=[], + positions=[], + prev_slot_indices=[], + capacity=4, + device=_DEVICE, + ) + cuda_log, ref_log = make_log_pair(device=_DEVICE) + for _ in range(30): + slot_before = int(cuda_log.slot_run_counter[0].item()) + kernel_before = int(cuda_log.kernel_run_counter[0].item()) + + _run_both_verify_no_rkv( + buf_pair=buf_pair, + plan_pair=empty_plan_pair, + cuda_log=cuda_log, + ref_log=ref_log, + assert_equal=False, + ) + + assert int(cuda_log.slot_run_counter[0].item()) == slot_before + assert int(cuda_log.kernel_run_counter[0].item()) == kernel_before + 1 + + _run_both_verify_no_rkv( + buf_pair=buf_pair, + plan_pair=nonempty_plan_pair, + cuda_log=cuda_log, + ref_log=ref_log, + assert_equal=False, + ) + + +class TestRunCounter: + def test_kernel_run_counter_per_call(self) -> None: + """``kernel_run_counter`` increments by 1 per call, even when ``verify_num_valid == 0``.""" + buf_pair = _buf_pair() + plan_pair = make_verify_plan_pair( + slot_indices=[], + positions=[], + prev_slot_indices=[], + capacity=4, + device=_DEVICE, + ) + # Use the low-level wrapper with explicit logs so we can observe the cross-call counter accumulation. + cuda_log, ref_log = make_log_pair(device=_DEVICE) + for _ in range(3): + _run_both_verify_no_rkv( + buf_pair=buf_pair, + plan_pair=plan_pair, + cuda_log=cuda_log, + ref_log=ref_log, + ) + + assert int(cuda_log.kernel_run_counter[0].item()) == 3 + + def test_slot_run_counter_per_entry(self) -> None: + """``slot_run_counter`` accumulates ``verify_num_valid`` entries per call.""" + cuda_buf, ref_buf = _buf_pair() + slot_indices = [1, 2, 3, 4] + tokens = [10, 11, 12, 13] + positions = [0, 1, 2, 3] + stamp_clean_chain( + cuda_buf=cuda_buf, + ref_buf=ref_buf, + tokens=tokens, + positions=positions, + slot_indices=slot_indices, + ) + + plan_pair = make_verify_plan_pair( + slot_indices=slot_indices, + positions=positions, + prev_slot_indices=[-1, 1, 2, 3], + device=_DEVICE, + ) + cuda_log, _ = run_verify_diff(buf_pair=(cuda_buf, ref_buf), plan_pair=plan_pair) + + assert int(cuda_log.slot_run_counter[0].item()) == 4 + + def test_grid_stride_processes_entries_beyond_grid_size(self) -> None: + """``verify_num_valid`` exceeding the persistent grid thread count is fully processed via grid-stride.""" + n_active = 40000 + cuda_buf, ref_buf = make_canary_buf_pair( + num_slots=n_active + 1, slot_stride_bytes=32, device=_DEVICE + ) + slot_indices = list(range(1, n_active + 1)) + tokens = [i + 10 for i in range(n_active)] + positions = list(range(n_active)) + stamp_clean_chain( + cuda_buf=cuda_buf, + ref_buf=ref_buf, + tokens=tokens, + positions=positions, + slot_indices=slot_indices, + ) + + plan_pair = make_verify_plan_pair( + slot_indices=slot_indices, + positions=positions, + prev_slot_indices=[-1] + slot_indices[:-1], + device=_DEVICE, + ) + cuda_log, _ = run_verify_diff(buf_pair=(cuda_buf, ref_buf), plan_pair=plan_pair) + + assert int(cuda_log.slot_run_counter[0].item()) == n_active + assert int(cuda_log.kernel_run_counter[0].item()) == 1 + + def test_replay_does_not_double_count_run_counters(self) -> None: + """Two consecutive runs on same plan: slot_run_counter += 2N, kernel_run_counter += 2.""" + cuda_buf, ref_buf = _buf_pair() + slot_indices = [1, 2, 3] + tokens = [10, 20, 30] + positions = [0, 1, 2] + stamp_clean_chain( + cuda_buf=cuda_buf, + ref_buf=ref_buf, + tokens=tokens, + positions=positions, + slot_indices=slot_indices, + ) + plan_pair = make_verify_plan_pair( + slot_indices=slot_indices, + positions=positions, + prev_slot_indices=[-1, 1, 2], + device=_DEVICE, + ) + cuda_log, ref_log = make_log_pair(device=_DEVICE) + for _ in range(2): + _run_both_verify_no_rkv( + buf_pair=(cuda_buf, ref_buf), + plan_pair=plan_pair, + cuda_log=cuda_log, + ref_log=ref_log, + ) + + assert int(cuda_log.slot_run_counter[0].item()) == 2 * len(slot_indices) + assert int(cuda_log.kernel_run_counter[0].item()) == 2 + + def test_slot_run_counter_delta_equals_active_entries_across_random_plans( + self, + ) -> None: + random.seed(0) + num_slots = 32 + cuda_buf = make_canary_buf( + num_slots=num_slots, slot_stride_bytes=32, device=_DEVICE + ) + ref_buf = cuda_buf.clone() + buf_pair = (cuda_buf, ref_buf) + + for slot_idx in range(num_slots): + _stamp_head( + buf_pair, + slot_idx=slot_idx, + token=slot_idx + 10, + position=slot_idx, + ) + + cuda_log, ref_log = make_log_pair(device=_DEVICE) + for _ in range(50): + bs = random.randint(1, 16) + entries_per_req = random.randint(1, 8) + n_entries = bs * entries_per_req + if n_entries > num_slots: + n_entries = num_slots + slot_indices = random.sample(range(num_slots), n_entries) + positions = [random.randint(0, 99) for _ in range(n_entries)] + prev_slot_indices = [-1] * n_entries + + plan_pair = make_verify_plan_pair( + slot_indices=slot_indices, + positions=positions, + prev_slot_indices=prev_slot_indices, + device=_DEVICE, + ) + + before = int(cuda_log.slot_run_counter[0].item()) + _run_both_verify_no_rkv( + buf_pair=buf_pair, + plan_pair=plan_pair, + cuda_log=cuda_log, + ref_log=ref_log, + assert_equal=False, + ) + after = int(cuda_log.slot_run_counter[0].item()) + assert after - before == n_entries + + def test_kernel_run_counter_per_call_invariant_50_calls(self) -> None: + buf_pair = _buf_pair(num_slots=8) + _stamp_head(buf_pair, slot_idx=1) + + plan_pair = _plan_pair_single(slot_idx=1, position=0) + cuda_log, ref_log = make_log_pair(device=_DEVICE) + for n in range(1, 51): + _run_both_verify_no_rkv( + buf_pair=buf_pair, + plan_pair=plan_pair, + cuda_log=cuda_log, + ref_log=ref_log, + assert_equal=False, + ) + assert int(cuda_log.kernel_run_counter[0].item()) == n + + +class TestViolationRing: + def test_violation_ring_fill_once_first_row(self) -> None: + """First violation lands at ring[0]; subsequent violations advance ``violation_write_index``.""" + buf_pair = _buf_pair() + # 3 chain-head entries with stored values that all yield POSITION mismatch (positions all 99). + for slot_idx in (1, 2, 3): + _stamp_head(buf_pair, slot_idx=slot_idx, token=1) + + plan_pair = make_verify_plan_pair( + slot_indices=[1, 2, 3], + positions=[99, 99, 99], + prev_slot_indices=[-1, -1, -1], + device=_DEVICE, + ) + cuda_log, _ = run_verify_diff( + buf_pair=buf_pair, plan_pair=plan_pair, assert_equal=False + ) + + assert _n_violations(cuda_log) == 3 + # row 0 is filled (slot 1 → POSITION bit). + assert _fail_bits(cuda_log) & consts.FailReason.VERIFY_POSITION_MISMATCH + + def test_violation_ring_overflow_counter_still_increments(self) -> None: + """Ring capacity exceeded → rows beyond are dropped but ``write_index`` still grows.""" + buf_pair = _buf_pair() + n_violations = 10 + slot_indices = list(range(1, n_violations + 1)) + for slot_idx in slot_indices: + _stamp_head(buf_pair, slot_idx=slot_idx, token=1) + + plan_pair = make_verify_plan_pair( + slot_indices=slot_indices, + positions=[99] * n_violations, + prev_slot_indices=[-1] * n_violations, + device=_DEVICE, + ) + cuda_log, ref_log = make_log_pair(capacity=4, device=_DEVICE) + _run_both_verify_no_rkv( + buf_pair=buf_pair, + plan_pair=plan_pair, + cuda_log=cuda_log, + ref_log=ref_log, + assert_equal=False, + ) + + assert _n_violations(cuda_log) == n_violations + assert _n_violations(ref_log) == n_violations + # Atomic-order may permute ring contents under overflow; only the write_index counter is + # byte-equal — we relax the ring-contents check here. + assert torch.equal(cuda_log.write_index, ref_log.write_index) + + def test_kernel_kind_stamped_into_row(self) -> None: + """Different ``CanaryLaunchTag`` values → violation row.kernel_kind reflects each.""" + buf_pair = _buf_pair() + _stamp_head(buf_pair, slot_idx=1, token=1) + + for tag in (CanaryLaunchTag.HEAD_K_FULL, CanaryLaunchTag.TAIL_V_SWA): + plan_pair = _plan_pair_single(slot_idx=1, position=99) + cuda_log, _ = run_verify_diff( + buf_pair=buf_pair, plan_pair=plan_pair, kernel_kind=tag + ) + kk = int(cuda_log.ring[0, consts.VIOLATION_FIELD_KERNEL_KIND].item()) + assert kk == int(tag) + + def test_violation_ring_row_byte_layout_hardcoded(self) -> None: + buf_pair = _buf_pair() + anchor_hash_signed = chain_anchor_signed() + _stamp_head(buf_pair, slot_idx=5, token=33) + + plan_pair = _plan_pair_single(slot_idx=5, position=99) + cuda_log, ref_log = make_log_pair(capacity=4, device=_DEVICE) + _run_both_verify_no_rkv( + buf_pair=buf_pair, + plan_pair=plan_pair, + cuda_log=cuda_log, + ref_log=ref_log, + ) + + assert _n_violations(cuda_log) == 1 + + # splitmix64(consts.CANARY_CHAIN_ANCHOR) = 0xde7fae23a9a1b716; signed = -2414019407054260458. + # Slot 5 was stamped with position=0 (stored); plan claims position=99 → POSITION mismatch. + # stored_chain_hash == expected_chain_hash (both splitmix64(ANCHOR)) → no CHAIN_HASH bit. + # plan has no populated verify_expected_tokens → entries read -1 sentinel, the + # verify kernel skips the token check (no TOKEN_MISMATCH bit) but the row's + # expected_token field reflects the gathered sentinel value. + # expected_aux = expected_chain_hash = splitmix64(ANCHOR) signed (same value as stored_chain_hash). + kernel_kind_val = int(CanaryLaunchTag.HEAD_K_FULL) + slot_idx_val = 5 + position_val = 0 + stored_token_val = 33 + expected_token_val = -1 + stored_chain_hash_val = anchor_hash_signed + expected_aux_val = anchor_hash_signed + fail_reason_bits_val = consts.FailReason.VERIFY_POSITION_MISMATCH + + expected_bytes = struct.pack( + "<8q", + kernel_kind_val, + slot_idx_val, + position_val, + stored_token_val, + expected_token_val, + stored_chain_hash_val, + expected_aux_val, + fail_reason_bits_val, + ) + + actual_bytes = cuda_log.ring[0].cpu().numpy().tobytes() + + if actual_bytes != expected_bytes: + expected_fields = struct.unpack("<8q", expected_bytes) + actual_fields = struct.unpack("<8q", actual_bytes) + field_names = [ + "kernel_kind", + "slot_idx", + "position", + "stored_token", + "expected_token", + "stored_chain_hash", + "expected_aux", + "fail_reason_bits", + ] + mismatches = [ + f" [{i}] {name}: expected {e} ({e:#x}) got {a} ({a:#x})" + for i, (name, e, a) in enumerate( + zip(field_names, expected_fields, actual_fields) + ) + if e != a + ] + raise AssertionError( + "violation_ring row binary layout mismatch:\n" + "\n".join(mismatches) + ) + + def test_violation_ring_atomic_with_many_violations(self) -> None: + """50 simultaneously-corrupted entries → write_index == 50 (no atomicity loss).""" + n = 50 + cuda_buf = make_canary_buf( + num_slots=n + 4, slot_stride_bytes=32, device=_DEVICE + ) + ref_buf = cuda_buf.clone() + slot_indices = list(range(1, n + 1)) + positions = [0] * n + + plan_pair = make_verify_plan_pair( + slot_indices=slot_indices, + positions=positions, + prev_slot_indices=[-1] * n, + capacity=n, + device=_DEVICE, + ) + cuda_log, ref_log = make_log_pair(capacity=128, device=_DEVICE) + _run_both_verify_no_rkv( + buf_pair=(cuda_buf, ref_buf), + plan_pair=plan_pair, + cuda_log=cuda_log, + ref_log=ref_log, + assert_equal=False, + ) + + assert _n_violations(cuda_log) == n + assert _n_violations(ref_log) == n + + def test_violation_rows_have_valid_kernel_kind_and_slot(self) -> None: + """Each violation row's kernel_kind matches the launch tag; slot_idx is one of the plan slots.""" + buf_pair = _buf_pair() + slot_indices = [1, 2, 3, 4] + positions = [0, 1, 2, 3] + plan_pair = make_verify_plan_pair( + slot_indices=slot_indices, + positions=positions, + prev_slot_indices=[-1] * 4, + device=_DEVICE, + ) + launch_tag = CanaryLaunchTag.HEAD_V_SWA + cuda_log, _ = run_verify_diff( + buf_pair=buf_pair, plan_pair=plan_pair, kernel_kind=launch_tag + ) + n_violations = _n_violations(cuda_log) + plan_slot_set = set(slot_indices) + for row in range(n_violations): + kind = int(cuda_log.ring[row, consts.VIOLATION_FIELD_KERNEL_KIND].item()) + assert kind == int( + launch_tag + ), f"row {row} kind {kind} != {int(launch_tag)}" + slot = int(cuda_log.ring[row, 1].item()) + assert slot in plan_slot_set, f"row {row} slot {slot} not in plan" + + def test_clear_resets_ring_and_write_index_zero(self) -> None: + buf_pair = _buf_pair(num_slots=8) + for slot_idx in range(1, 4): + _stamp_head(buf_pair, slot_idx=slot_idx, token=1, position=99) + + plan_pair = make_verify_plan_pair( + slot_indices=[1, 2, 3], + positions=[0, 0, 0], + prev_slot_indices=[-1, -1, -1], + device=_DEVICE, + ) + cuda_log, _ = run_verify_diff( + buf_pair=buf_pair, plan_pair=plan_pair, assert_equal=False + ) + + assert _n_violations(cuda_log) > 0 + + cuda_log_fresh = FakeViolationLog.allocate(device=_DEVICE) + assert _n_violations(cuda_log_fresh) == 0 + assert torch.all(cuda_log_fresh.ring == 0).item() + + +class TestBoundarySweep: + @pytest.mark.parametrize( + "token_val", + [0, 1, 0xFFFFFFFF, -1, 0x80000000, 0x7FFFFFFF], + ) + def test_token_boundary_byte_equal_sweep(self, token_val: int) -> None: + """Sweep token boundary values; assert CUDA vs ref state byte-equal.""" + _run_verify_single_slot_byte_equal( + _VerifySingleSlotInput( + token=token_val, stored_prev_hash_signed=chain_anchor_signed() + ) + ) + + @pytest.mark.parametrize( + "position_val", + [0, 1, 127, 128, 129, 0x7FFFFFFF], + ) + def test_position_boundary_byte_equal_sweep(self, position_val: int) -> None: + """Sweep position boundary values; assert CUDA vs ref state byte-equal.""" + _run_verify_single_slot_byte_equal( + _VerifySingleSlotInput( + position=position_val, stored_prev_hash_signed=chain_anchor_signed() + ) + ) + + @pytest.mark.parametrize( + "prev_hash_val", + [ + pytest.param(0, id="zero"), + pytest.param(None, id="splitmix64_of_zero"), + pytest.param(0xFFFFFFFFFFFFFFFF, id="all_ones"), + pytest.param(0x8000000000000000, id="sign_bit"), + ], + ) + def test_prev_hash_boundary_byte_equal_sweep( + self, prev_hash_val: int | None + ) -> None: + """Sweep prev_hash boundary values; assert CUDA vs ref state byte-equal.""" + if prev_hash_val is None: + prev_hash_val = splitmix64(0) + _run_verify_single_slot_byte_equal( + _VerifySingleSlotInput( + stored_prev_hash_signed=to_signed_int64(prev_hash_val) + ) + ) + + +class TestVerifyExpectedInputIds: + """Cover the new verify-time token-id check via VerifyPlan.verify_expected_tokens.""" + + def _run( + self, + *, + stored_token: int, + expected_input_id: int, + check_verify_expected_token: bool = True, + ) -> FakeViolationLog: + buf_pair = _buf_pair() + _stamp_head(buf_pair, slot_idx=1, token=stored_token) + plan_pair = _plan_pair_single( + slot_idx=1, position=0, expected_input_id=expected_input_id + ) + cuda_log, _ = run_verify_diff( + buf_pair=buf_pair, + plan_pair=plan_pair, + check_verify_expected_token=check_verify_expected_token, + ) + return cuda_log + + def test_sentinel_skips_token_check(self) -> None: + """``expected_input_id == -1`` must not fire even if stored token differs.""" + log = self._run(stored_token=42, expected_input_id=-1) + assert _n_violations(log) == 0 + + def test_match_records_no_violation(self) -> None: + """Matching stored vs expected token: zero violations.""" + log = self._run(stored_token=42, expected_input_id=42) + assert _n_violations(log) == 0 + + def test_mismatch_fires_verify_token_bit(self) -> None: + """Mismatch: kVerifyTokenMismatch bit set; expected_token field populated.""" + log = self._run(stored_token=42, expected_input_id=99) + assert _n_violations(log) == 1 + row = log.ring[0].tolist() + assert_only_bits_set( + int(row[consts.VIOLATION_FIELD_FAIL_REASON_BITS]), + int(consts.FailReason.VERIFY_TOKEN_MISMATCH), + ) + assert int(row[consts.VIOLATION_FIELD_STORED_TOKEN]) == 42 + assert int(row[consts.VIOLATION_FIELD_EXPECTED_TOKEN]) == 99 + + def test_check_disabled_never_raises_token_mismatch(self) -> None: + """check_verify_expected_token=False: token mismatch is silently ignored byte-equal CUDA vs ref.""" + log = self._run( + stored_token=42, + expected_input_id=99, + check_verify_expected_token=False, + ) + assert _n_violations(log) == 0 + + def test_check_disabled_does_not_read_expected_tokens_tensor(self) -> None: + """check_verify_expected_token=False with garbage expected_input_id: kernel must not deref the tensor.""" + buf_pair = _buf_pair() + _stamp_head(buf_pair, slot_idx=1) + garbage_value = (1 << 63) - 1 + plan_pair = _plan_pair_single( + slot_idx=1, position=0, expected_input_id=garbage_value + ) + cuda_log, _ = run_verify_diff( + buf_pair=buf_pair, + plan_pair=plan_pair, + check_verify_expected_token=False, + ) + assert _n_violations(cuda_log) == 0 + + def test_token_and_position_both_mismatch_set_both_bits(self) -> None: + """token mismatch + position mismatch: both fail bits set; expected_token row carries the gathered id.""" + buf_pair = _buf_pair() + slot_idx = 5 + _stamp_head(buf_pair, slot_idx=slot_idx, position=10) + plan_pair = _plan_pair_single( + slot_idx=slot_idx, position=99, expected_input_id=123 + ) + cuda_log, _ = run_verify_diff( + buf_pair=buf_pair, + plan_pair=plan_pair, + check_verify_expected_token=True, + ) + assert _n_violations(cuda_log) == 1 + row = cuda_log.ring[0].tolist() + expected_bits = int(consts.FailReason.VERIFY_TOKEN_MISMATCH) | int( + consts.FailReason.VERIFY_POSITION_MISMATCH + ) + assert_only_bits_set( + int(row[consts.VIOLATION_FIELD_FAIL_REASON_BITS]), expected_bits + ) + assert int(row[consts.VIOLATION_FIELD_STORED_TOKEN]) == 42 + assert int(row[consts.VIOLATION_FIELD_EXPECTED_TOKEN]) == 123 + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) diff --git a/python/sglang/jit_kernel/tests/kv_canary/test_write_fuzz.py b/python/sglang/jit_kernel/tests/kv_canary/test_write_fuzz.py new file mode 100644 index 000000000..6010796a5 --- /dev/null +++ b/python/sglang/jit_kernel/tests/kv_canary/test_write_fuzz.py @@ -0,0 +1,218 @@ +from __future__ import annotations + +import random +from dataclasses import dataclass + +import pytest +import torch + +from sglang.jit_kernel.kv_canary.verify import CanaryLaunchTag +from sglang.jit_kernel.kv_canary.write import WritePlan +from sglang.jit_kernel.tests.kv_canary._canary_helpers import ( + FakeViolationLog, + make_canary_buf, + make_log_pair, + make_write_plan_pair, + stamp_pair, +) +from sglang.jit_kernel.tests.kv_canary._differential import _run_both_write +from sglang.jit_kernel.tests.kv_canary._fuzz_driver import ( + FUZZ_SEEDS_PR, + run_fuzz_combo, +) +from sglang.jit_kernel.tests.kv_canary._invariants import WriteInvariants +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large") + + +_DEVICE = torch.device("cuda") + +_FUZZ_ITER_PER_SEED = 30 + + +@dataclass(frozen=True, slots=True, kw_only=True) +class WriteFuzzInputs: + cuda_canary_buf: torch.Tensor + ref_canary_buf: torch.Tensor + plan_cuda: WritePlan + plan_ref: WritePlan + input_ids: torch.Tensor + positions: torch.Tensor + out_cache_loc: torch.Tensor + kernel_kind: CanaryLaunchTag + enable_write_verify_inputs: bool + expected_input_tokens: torch.Tensor + expected_input_positions: torch.Tensor + ring_capacity: int + + +def _draw_random_write_inputs(rng: random.Random) -> WriteFuzzInputs: + enable_write_verify_inputs = rng.choice([False, True]) + kernel_kind = rng.choice(list(CanaryLaunchTag)) + ring_capacity = rng.choice([16, 64, 256]) + + n_reqs = rng.randint(1, 4) + per_req_tokens: list[int] = [rng.randint(1, 5) for _ in range(n_reqs)] + total_tokens = sum(per_req_tokens) + num_slots = max(total_tokens + 8, 16) + + cuda_buf = make_canary_buf( + num_slots=num_slots, slot_stride_bytes=32, device=_DEVICE + ) + ref_buf = cuda_buf.clone() + + write_offsets: list[int] = [0] + running = 0 + for t in per_req_tokens: + running += t + write_offsets.append(running) + slot_pool = list(range(1, num_slots)) + rng.shuffle(slot_pool) + seed_slot_indices: list[int] = [] + for _ in range(n_reqs): + if rng.random() < 0.4 or len(slot_pool) <= total_tokens: + seed_slot_indices.append(-1) + else: + seed_slot_indices.append(slot_pool.pop()) + + out_cache_loc_list: list[int] = [] + for _ in range(total_tokens): + if not slot_pool: + out_cache_loc_list.append(-1) + else: + out_cache_loc_list.append(slot_pool.pop()) + + plan_cuda, plan_ref = make_write_plan_pair( + write_offsets=write_offsets, + seed_slot_indices=seed_slot_indices, + num_valid_reqs=n_reqs, + device=_DEVICE, + ) + + input_ids = torch.tensor( + [rng.randint(-(1 << 31), (1 << 31) - 1) for _ in range(total_tokens)], + dtype=torch.int64, + device=_DEVICE, + ) + # Per-chain sequential positions so the write kernel's chain-step position assert holds. + # For chains with a real seed slot, stamp the seed with (chain_start_position - 1) so the + # first chain entry's position == seed.position + 1. + chain_start_positions: list[int] = [rng.randint(0, 1024) for _ in range(n_reqs)] + positions_list: list[int] = [] + for r in range(n_reqs): + start = chain_start_positions[r] + positions_list.extend(start + i for i in range(per_req_tokens[r])) + seed_slot = seed_slot_indices[r] + if seed_slot >= 0: + stamp_pair( + (cuda_buf, ref_buf), + slot_idx=seed_slot, + token=0, + position=start - 1, + prev_hash=0, + ) + positions = torch.tensor(positions_list, dtype=torch.int64, device=_DEVICE) + out_cache_loc = torch.tensor(out_cache_loc_list, dtype=torch.int64, device=_DEVICE) + expected_input_tokens = input_ids.clone() + expected_input_positions = positions.clone() + if enable_write_verify_inputs: + candidate_indices = [ + idx for idx, slot in enumerate(out_cache_loc_list) if slot >= 0 + ] + rng.shuffle(candidate_indices) + mismatch_count = rng.randint(0, len(candidate_indices)) + for idx in candidate_indices[:mismatch_count]: + if rng.choice([False, True]): + expected_input_tokens[idx] = expected_input_tokens[idx] + 1 + else: + expected_input_positions[idx] = expected_input_positions[idx] + 1 + + return WriteFuzzInputs( + cuda_canary_buf=cuda_buf, + ref_canary_buf=ref_buf, + plan_cuda=plan_cuda, + plan_ref=plan_ref, + input_ids=input_ids, + positions=positions, + out_cache_loc=out_cache_loc, + kernel_kind=kernel_kind, + enable_write_verify_inputs=enable_write_verify_inputs, + expected_input_tokens=expected_input_tokens, + expected_input_positions=expected_input_positions, + ring_capacity=ring_capacity, + ) + + +def _run_one(inputs: WriteFuzzInputs) -> None: + cuda_buf_before = inputs.cuda_canary_buf.clone() + cuda_log, ref_log = make_log_pair(capacity=inputs.ring_capacity, device=_DEVICE) + log_before = FakeViolationLog.allocate( + capacity=inputs.ring_capacity, device=_DEVICE + ) + _run_both_write( + cuda_canary_buf=inputs.cuda_canary_buf, + ref_canary_buf=inputs.ref_canary_buf, + plan_cuda=inputs.plan_cuda, + plan_ref=inputs.plan_ref, + input_ids=inputs.input_ids, + positions=inputs.positions, + out_cache_loc=inputs.out_cache_loc, + enable_write_verify_inputs=inputs.enable_write_verify_inputs, + expected_input_tokens=inputs.expected_input_tokens, + expected_input_positions=inputs.expected_input_positions, + cuda_log=cuda_log, + ref_log=ref_log, + kernel_kind=inputs.kernel_kind, + assert_equal=False, + ) + assert torch.equal( + inputs.cuda_canary_buf, inputs.ref_canary_buf + ), "CUDA vs ref canary_buf diverged" + assert int(cuda_log.write_index[0].item()) == int(ref_log.write_index[0].item()) + assert int(cuda_log.slot_run_counter[0].item()) == int( + ref_log.slot_run_counter[0].item() + ) + assert int(cuda_log.kernel_run_counter[0].item()) == int( + ref_log.kernel_run_counter[0].item() + ) + WriteInvariants.assert_all( + canary_buf_before=cuda_buf_before, + canary_buf_after=inputs.cuda_canary_buf, + plan=inputs.plan_cuda, + input_ids=inputs.input_ids, + positions=inputs.positions, + out_cache_loc=inputs.out_cache_loc, + enable_write_verify_inputs=inputs.enable_write_verify_inputs, + expected_input_tokens=inputs.expected_input_tokens, + expected_input_positions=inputs.expected_input_positions, + log_before=log_before, + log_after=cuda_log, + ) + + +def _summarize(inputs: WriteFuzzInputs) -> str: + n_active = int(inputs.plan_cuda.write_num_valid_reqs[0].item()) + total = int(inputs.plan_cuda.write_offsets[n_active].item()) + return ( + f"n_reqs={n_active} total_tokens={total} kind={inputs.kernel_kind.name} " + f"pseudo={inputs.enable_write_verify_inputs}" + ) + + +@pytest.mark.parametrize("seed", FUZZ_SEEDS_PR) +def test_write_fuzz_full_combo(seed: int) -> None: + """Multi-dim write fuzzer: random pseudo/kernel × N iters, byte-equal.""" + run_fuzz_combo( + seed, + draw_fn=_draw_random_write_inputs, + run_one_fn=_run_one, + summarize_fn=_summarize, + n_iter=_FUZZ_ITER_PER_SEED, + ) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) diff --git a/python/sglang/jit_kernel/tests/kv_canary/test_write_hand.py b/python/sglang/jit_kernel/tests/kv_canary/test_write_hand.py new file mode 100644 index 000000000..a7c5689f2 --- /dev/null +++ b/python/sglang/jit_kernel/tests/kv_canary/test_write_hand.py @@ -0,0 +1,892 @@ +from __future__ import annotations + +from dataclasses import dataclass +from unittest.mock import patch + +import pytest +import torch + +from sglang.jit_kernel.kv_canary import consts +from sglang.jit_kernel.kv_canary import write as write_module +from sglang.jit_kernel.kv_canary.consts import splitmix64, splitmix64_mix3 +from sglang.jit_kernel.kv_canary.verify import ( + CANARY_SLOT_BYTES, + CanaryLaunchTag, + VerifyOrWriteContext, + launch_canary_verify_kernel, +) +from sglang.jit_kernel.kv_canary.write import ( + launch_canary_write_kernel, +) +from sglang.jit_kernel.tests.kv_canary._canary_helpers import ( + FakeViolationLog, + assert_canary_state_equal, + assert_only_bits_set, + chain_anchor_signed, + make_canary_buf, + make_canary_buf_pair, + make_log_pair, + make_verify_plan, + make_write_plan, + make_write_plan_pair, + read_slot_fields, + stamp_pair, + to_signed_int64, +) +from sglang.jit_kernel.tests.kv_canary._differential import ( + _run_both_write, + run_write_diff, +) +from sglang.jit_kernel.tests.kv_canary._fixtures import ( + dummy_pseudo_tensors, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large") + + +_DEVICE = torch.device("cuda") + + +def _int32_tensor(values: list[int]) -> torch.Tensor: + return torch.tensor(values, dtype=torch.int64, device=_DEVICE) + + +def _make_default_buf_pair( + num_slots: int = 16, slot_stride_bytes: int = 32 +) -> tuple[torch.Tensor, torch.Tensor]: + return make_canary_buf_pair( + num_slots=num_slots, slot_stride_bytes=slot_stride_bytes, device=_DEVICE + ) + + +def _run_write( + *, + buf_pair: tuple[torch.Tensor, torch.Tensor], + input_ids: list[int] | torch.Tensor, + positions: list[int] | torch.Tensor, + out_cache_loc: list[int] | torch.Tensor, + write_offsets: list[int] | None = None, + seed_slot_indices: list[int] = (-1,), + num_valid_reqs: int = 1, + req_capacity: int | None = None, + enable_write_verify_inputs: bool = False, + expected_input_tokens: torch.Tensor | None = None, + expected_input_positions: torch.Tensor | None = None, + assert_equal: bool = True, +) -> tuple[FakeViolationLog, FakeViolationLog]: + """Shared scaffold: build write plan + pseudo tensors and call ``run_write_diff``. + + ``write_offsets`` defaults to ``[0, len(input_ids)]`` (one req covering all entries). + ``expected_input_*`` default to ``dummy_pseudo_tensors(len(input_ids))``. + """ + ids_t = ( + input_ids + if isinstance(input_ids, torch.Tensor) + else _int32_tensor(list(input_ids)) + ) + pos_t = ( + positions + if isinstance(positions, torch.Tensor) + else _int32_tensor(list(positions)) + ) + loc_t = ( + out_cache_loc + if isinstance(out_cache_loc, torch.Tensor) + else _int32_tensor(list(out_cache_loc)) + ) + n_tokens = int(ids_t.shape[0]) + + if write_offsets is None: + write_offsets = [0, n_tokens] + + plan_kwargs = dict( + write_offsets=write_offsets, + seed_slot_indices=list(seed_slot_indices), + num_valid_reqs=num_valid_reqs, + device=_DEVICE, + ) + if req_capacity is not None: + plan_kwargs["req_capacity"] = req_capacity + plan_pair = make_write_plan_pair(**plan_kwargs) + + if expected_input_tokens is None or expected_input_positions is None: + pseudo_tokens, pseudo_positions = dummy_pseudo_tensors(n_tokens) + if expected_input_tokens is None: + expected_input_tokens = pseudo_tokens + if expected_input_positions is None: + expected_input_positions = pseudo_positions + + return run_write_diff( + buf_pair=buf_pair, + plan_pair=plan_pair, + input_ids=ids_t, + positions=pos_t, + out_cache_loc=loc_t, + enable_write_verify_inputs=enable_write_verify_inputs, + expected_input_tokens=expected_input_tokens, + expected_input_positions=expected_input_positions, + assert_equal=assert_equal, + ) + + +@dataclass(frozen=True, slots=True, kw_only=True) +class _WriteSingleSlotInput: + token: int = 42 + position: int = 0 + enable_write_verify_inputs: bool = False + + +class _RecordingWriteModule: + def __init__(self) -> None: + self.calls: list[tuple[object, ...]] = [] + + def canary_write_step_cuda(self, *args: object) -> None: + self.calls.append(args) + + +def _run_write_single_slot_byte_equal(case: _WriteSingleSlotInput) -> None: + _run_write( + buf_pair=_make_default_buf_pair(), + input_ids=[case.token], + positions=[case.position], + out_cache_loc=[0], + enable_write_verify_inputs=case.enable_write_verify_inputs, + ) + + +class TestSeedSlot: + def setup_method(self) -> None: + self.buf_pair = _make_default_buf_pair() + + def test_seed_slot_idx_negative_uses_anchor(self) -> None: + """``seed_slot_idx == -1`` → initial ``running_prev_hash`` is ``splitmix64(consts.CANARY_CHAIN_ANCHOR)``.""" + _run_write( + buf_pair=self.buf_pair, + input_ids=[42], + positions=[0], + out_cache_loc=[3], + ) + + stored_token, stored_position, stored_prev_hash, _ = read_slot_fields( + canary_buf=self.buf_pair[0], slot_idx=3 + ) + assert stored_token == 42 + assert stored_position == 0 + assert stored_prev_hash == chain_anchor_signed() + + def test_seed_slot_idx_loads_predecessor(self) -> None: + """``seed_slot_idx >= 0`` → load 3 fields (token, position, prev_hash) from ``canary_buf[seed]`` and splitmix64_mix3-advance into prev_hash.""" + # Step: pre-stamp slot 7 with a known chain link. + seed_token, seed_position = 100, 4 + seed_prev_signed = to_signed_int64(splitmix64(consts.CANARY_CHAIN_ANCHOR)) + stamp_pair( + self.buf_pair, + slot_idx=7, + token=seed_token, + position=seed_position, + prev_hash=seed_prev_signed, + ) + + _run_write( + buf_pair=self.buf_pair, + input_ids=[999], + positions=[5], + out_cache_loc=[2], + seed_slot_indices=[7], + ) + + expected_prev_hash = splitmix64_mix3( + splitmix64(consts.CANARY_CHAIN_ANCHOR), seed_token, seed_position + ) + _, _, stored_prev_hash, _ = read_slot_fields( + canary_buf=self.buf_pair[0], slot_idx=2 + ) + assert stored_prev_hash == to_signed_int64(expected_prev_hash) + + def test_seed_slot_chain_link_continuous(self) -> None: + """After write, ``slot[0].prev_hash`` is consistent with verify's chain reconstruction from seed.""" + # Step 1: write a chain from seed slot=7 → newly written slot=2. Then run verify with prev=7 and + # assert no violation — i.e., slot[2].prev_hash is splitmix64_mix3(seed.prev_hash, seed.token, seed.position). + cuda_buf = self.buf_pair[0] + seed_token, seed_position = 11, 0 + seed_prev_signed = to_signed_int64(splitmix64(consts.CANARY_CHAIN_ANCHOR)) + stamp_pair( + self.buf_pair, + slot_idx=7, + token=seed_token, + position=seed_position, + prev_hash=seed_prev_signed, + ) + + _run_write( + buf_pair=self.buf_pair, + input_ids=[222], + positions=[1], + out_cache_loc=[2], + seed_slot_indices=[7], + assert_equal=False, + ) + + # Step 2: verify slot[2] with prev=7 — expects no violation. + verify_plan = make_verify_plan( + slot_indices=[2], positions=[1], prev_slot_indices=[7], device=_DEVICE + ) + verify_log = FakeViolationLog.allocate(device=_DEVICE) + launch_canary_verify_kernel( + context=VerifyOrWriteContext( + canary_buf=cuda_buf, + kernel_kind=CanaryLaunchTag.HEAD_K_FULL, + violation_ring=verify_log.ring, + violation_write_index=verify_log.write_index, + slot_run_counter=verify_log.slot_run_counter, + kernel_run_counter=verify_log.kernel_run_counter, + enable_chain_position_assert=verify_log.enable_chain_position_assert, + ), + plan=verify_plan, + check_verify_expected_token=True, + ) + torch.cuda.synchronize() + assert int(verify_log.write_index[0].item()) == 0 + + def test_seed_slot_resume_5_step_hardcoded(self) -> None: + cuda_buf = make_canary_buf(num_slots=50, slot_stride_bytes=32, device=_DEVICE) + ref_buf = cuda_buf.clone() + buf_pair = (cuda_buf, ref_buf) + seed_token = 7 + seed_position = 10 + seed_prev_hash_signed = to_signed_int64(splitmix64(consts.CANARY_CHAIN_ANCHOR)) + stamp_pair( + buf_pair, + slot_idx=42, + token=seed_token, + position=seed_position, + prev_hash=seed_prev_hash_signed, + ) + + predecessor_advance = splitmix64_mix3( + splitmix64(consts.CANARY_CHAIN_ANCHOR), seed_token, seed_position + ) + + tokens = [101, 202, 303, 404, 505] + positions = [11, 12, 13, 14, 15] + out_cache_loc = [0, 1, 2, 3, 4] + + expected_prev_hashes: list[int] = [] + running = predecessor_advance + for t, p in zip(tokens, positions): + expected_prev_hashes.append(running) + running = splitmix64_mix3(running, t, p) + + cuda_log, _ = _run_write( + buf_pair=buf_pair, + input_ids=tokens, + positions=positions, + out_cache_loc=out_cache_loc, + write_offsets=[0, 5], + seed_slot_indices=[42], + ) + + for slot_idx, expected_token, expected_position, expected_prev_u64 in zip( + out_cache_loc, tokens, positions, expected_prev_hashes + ): + stored_token, stored_position, stored_prev_hash, stored_real_kv_hash = ( + read_slot_fields(canary_buf=cuda_buf, slot_idx=slot_idx) + ) + assert stored_token == expected_token + assert stored_position == expected_position + assert stored_prev_hash == to_signed_int64(expected_prev_u64) + assert stored_real_kv_hash == 0 + + assert int(cuda_log.write_index[0].item()) == 0 + + def test_seed_continues_existing_chain(self) -> None: + """Pre-stamp seed slot; subsequent write should continue chain from splitmix64_mix3(seed.*).""" + seed_slot = 3 + seed_token = 7 + seed_position = 1 + expected_seed_prev_hash = splitmix64(consts.CANARY_CHAIN_ANCHOR) + stamp_pair( + self.buf_pair, + slot_idx=seed_slot, + token=seed_token, + position=seed_position, + prev_hash=to_signed_int64(expected_seed_prev_hash), + ) + + new_slot = 4 + new_token = 13 + new_position = 2 + expected_running = splitmix64_mix3( + expected_seed_prev_hash, seed_token, seed_position + ) + + _run_write( + buf_pair=self.buf_pair, + input_ids=[new_token], + positions=[new_position], + out_cache_loc=[new_slot], + seed_slot_indices=[seed_slot], + ) + + new_stored = read_slot_fields(canary_buf=self.buf_pair[0], slot_idx=new_slot) + assert new_stored[0] == new_token + assert new_stored[1] == new_position + assert new_stored[2] == to_signed_int64( + expected_running + ), f"new slot prev_hash {new_stored[2]} != expected {to_signed_int64(expected_running)}" + + +class TestChain: + def setup_method(self) -> None: + self.buf_pair = _make_default_buf_pair() + + def test_chain_link_byte_equal_5_step(self) -> None: + """5-step chain, buf / ring / counters byte-equal against ref.""" + _run_write( + buf_pair=self.buf_pair, + input_ids=[10, 20, 30, 40, 50], + positions=[0, 1, 2, 3, 4], + out_cache_loc=[0, 1, 2, 3, 4], + ) + + def test_chain_link_byte_equal_5_step_hardcoded(self) -> None: + """5-step write chain with hand-computed splitmix64 expected fields per slot.""" + tokens = [101, 202, 303, 404, 505] + positions = [0, 1, 2, 3, 4] + out_cache_loc = [0, 1, 2, 3, 4] + + # Step 1: compute the expected stored prev_hash sequence in pure Python via splitmix64. + expected_prev_hashes_u64: list[int] = [] + running = splitmix64(consts.CANARY_CHAIN_ANCHOR) + for token, position in zip(tokens, positions): + expected_prev_hashes_u64.append(running) + running = splitmix64_mix3(running, token, position) + expected_prev_hashes_signed = [ + to_signed_int64(h) for h in expected_prev_hashes_u64 + ] + + _run_write( + buf_pair=self.buf_pair, + input_ids=tokens, + positions=positions, + out_cache_loc=out_cache_loc, + ) + + # Step 2: verify every slot's stored 4 fields match the hardcoded expected sequence. + for slot_idx, expected_token, expected_position, expected_prev_signed in zip( + out_cache_loc, tokens, positions, expected_prev_hashes_signed + ): + stored_token, stored_position, stored_prev_hash, stored_real_kv_hash = ( + read_slot_fields(canary_buf=self.buf_pair[0], slot_idx=slot_idx) + ) + assert stored_token == expected_token + assert stored_position == expected_position + assert stored_prev_hash == expected_prev_signed + assert stored_real_kv_hash == 0 + + +class TestMockMode: + def setup_method(self) -> None: + self.buf_pair = _make_default_buf_pair() + + def test_mock_mode_off_ignores_expected(self) -> None: + """``enable_write_verify_inputs = OFF`` → expected tensors are ignored (we pass garbage to prove the kernel skips them).""" + # Garbage expected tensors that, if the kernel mistakenly reads, would generate mismatches. + cuda_log, _ = _run_write( + buf_pair=self.buf_pair, + input_ids=[1, 2, 3], + positions=[0, 1, 2], + out_cache_loc=[0, 1, 2], + expected_input_tokens=_int32_tensor([999, 999, 999]), + expected_input_positions=_int32_tensor([999, 999, 999]), + ) + + assert int(cuda_log.write_index[0].item()) == 0 + + def test_mock_mode_on_match_no_violation(self) -> None: + """``enable_write_verify_inputs = ON`` and expected matches actual → no violation, chain advances.""" + input_ids = _int32_tensor([7, 8, 9]) + positions = _int32_tensor([0, 1, 2]) + cuda_log, _ = _run_write( + buf_pair=self.buf_pair, + input_ids=input_ids, + positions=positions, + out_cache_loc=[0, 1, 2], + enable_write_verify_inputs=True, + expected_input_tokens=input_ids.clone(), + expected_input_positions=positions.clone(), + ) + + assert int(cuda_log.write_index[0].item()) == 0 + + def test_mock_mode_on_token_mismatch_records_violation(self) -> None: + """``enable_write_verify_inputs = ON`` token mismatch → violation recorded; chain advances on ACTUAL token.""" + cuda_log, _ = _run_write( + buf_pair=self.buf_pair, + input_ids=[42], + positions=[0], + out_cache_loc=[0], + enable_write_verify_inputs=True, + expected_input_tokens=_int32_tensor([99]), + expected_input_positions=_int32_tensor([0]), + ) + + fail_bits = int( + cuda_log.ring[0, consts.VIOLATION_FIELD_FAIL_REASON_BITS].item() + ) + assert_only_bits_set(fail_bits, consts.FailReason.WRITE_TOKEN_MISMATCH) + # Chain advances on actual (42), not expected (99). Stored token should be 42. + stored_token, _, _, _ = read_slot_fields( + canary_buf=self.buf_pair[0], slot_idx=0 + ) + assert stored_token == 42 + + def test_mock_mode_on_position_mismatch_records_violation(self) -> None: + """``enable_write_verify_inputs = ON`` position mismatch → violation recorded; chain advances on ACTUAL position.""" + cuda_log, _ = _run_write( + buf_pair=self.buf_pair, + input_ids=[42], + positions=[7], + out_cache_loc=[0], + enable_write_verify_inputs=True, + expected_input_tokens=_int32_tensor([42]), + expected_input_positions=_int32_tensor([0]), + ) + + fail_bits = int( + cuda_log.ring[0, consts.VIOLATION_FIELD_FAIL_REASON_BITS].item() + ) + assert_only_bits_set(fail_bits, consts.FailReason.WRITE_POSITION_MISMATCH) + _, stored_position, _, _ = read_slot_fields( + canary_buf=self.buf_pair[0], slot_idx=0 + ) + assert stored_position == 7 + + def test_mock_mode_chain_advances_on_actual_not_expected(self) -> None: + """Expected differs from actual on every entry → downstream verify must NOT cascade chain errors.""" + cuda_buf = self.buf_pair[0] + # Every actual differs from expected. + cuda_log, _ = _run_write( + buf_pair=self.buf_pair, + input_ids=[10, 20, 30], + positions=[0, 1, 2], + out_cache_loc=[1, 2, 3], + enable_write_verify_inputs=True, + expected_input_tokens=_int32_tensor([999, 999, 999]), + expected_input_positions=_int32_tensor([999, 999, 999]), + ) + + # All 3 entries should fire a violation row. + assert int(cuda_log.write_index[0].item()) == 3 + # Run a downstream verify — it must see no chain mismatch because chain advanced on actuals. + verify_plan = make_verify_plan( + slot_indices=[1, 2, 3], + positions=[0, 1, 2], + prev_slot_indices=[-1, 1, 2], + device=_DEVICE, + ) + verify_log = FakeViolationLog.allocate(device=_DEVICE) + launch_canary_verify_kernel( + context=VerifyOrWriteContext( + canary_buf=cuda_buf, + kernel_kind=CanaryLaunchTag.HEAD_K_FULL, + violation_ring=verify_log.ring, + violation_write_index=verify_log.write_index, + slot_run_counter=verify_log.slot_run_counter, + kernel_run_counter=verify_log.kernel_run_counter, + enable_chain_position_assert=verify_log.enable_chain_position_assert, + ), + plan=verify_plan, + check_verify_expected_token=True, + ) + torch.cuda.synchronize() + assert int(verify_log.write_index[0].item()) == 0 + + @pytest.mark.parametrize("bit_to_trigger", ["MOCK_TOKEN", "MOCK_POSITION"]) + @pytest.mark.parametrize("injection_position", ["head", "mid", "last"]) + def test_mock_violation_bit_injection_position_matrix( + self, + bit_to_trigger: str, + injection_position: str, + ) -> None: + """Sweep injection_position x bit_to_trigger for write-kernel pseudo-mode fail-reason coverage.""" + slot_count = 5 + tokens = [10, 20, 30, 40, 50] + positions = [0, 1, 2, 3, 4] + out_cache_locs = [0, 1, 2, 3, 4] + corruption_index = {"head": 0, "mid": 2, "last": 4}[injection_position] + corrupt_slot = out_cache_locs[corruption_index] + + expected_bit = { + "MOCK_TOKEN": consts.FailReason.WRITE_TOKEN_MISMATCH, + "MOCK_POSITION": consts.FailReason.WRITE_POSITION_MISMATCH, + }[bit_to_trigger] + + input_ids = _int32_tensor(tokens) + positions_t = _int32_tensor(positions) + out_cache_loc = _int32_tensor(out_cache_locs) + + pseudo_tokens = input_ids.clone() + pseudo_positions = positions_t.clone() + + if bit_to_trigger == "MOCK_TOKEN": + pseudo_tokens[corruption_index] = tokens[corruption_index] + 999 + else: + pseudo_positions[corruption_index] = positions[corruption_index] + 99 + + cuda_log, ref_log = _run_write( + buf_pair=self.buf_pair, + input_ids=input_ids, + positions=positions_t, + out_cache_loc=out_cache_loc, + write_offsets=[0, slot_count], + enable_write_verify_inputs=True, + expected_input_tokens=pseudo_tokens, + expected_input_positions=pseudo_positions, + assert_equal=False, + ) + + found = False + for row_idx in range(int(cuda_log.write_index[0].item())): + fail_bits = int( + cuda_log.ring[row_idx, consts.VIOLATION_FIELD_FAIL_REASON_BITS].item() + ) + row_slot = int(cuda_log.ring[row_idx, 1].item()) + if (fail_bits & expected_bit) and row_slot == corrupt_slot: + found = True + break + assert found, ( + f"expected bit {expected_bit:#x} at slot {corrupt_slot} not found in ring " + f"(bit_to_trigger={bit_to_trigger} injection_position={injection_position})" + ) + assert_canary_state_equal(log_a=cuda_log, log_b=ref_log) + + +class TestSlotHandling: + def setup_method(self) -> None: + self.buf_pair = _make_default_buf_pair() + + def test_negative_slot_skips_entry(self) -> None: + """``out_cache_loc[i] < 0`` → that entry is skipped: no buf write, no violation, no + canary slot mutation, and no write slot_run_counter increment. + Covers both SWA out-of-window (after caller-side LUT gather) and explicit padding intents. + """ + # Two entries: first writes to slot 4 normally; second has slot=-1 and must be skipped. + cuda_log, _ = _run_write( + buf_pair=self.buf_pair, + input_ids=[42, 99], + positions=[0, 1], + out_cache_loc=[4, -1], + ) + + stored_token, _, _, _ = read_slot_fields( + canary_buf=self.buf_pair[0], slot_idx=4 + ) + assert stored_token == 42 + assert int(cuda_log.slot_run_counter.item()) == 1 + + def test_pre_translated_slot_writes_normally(self) -> None: + """``out_cache_loc[i] >= 0`` → the kernel writes to exactly that slot, with no LUT applied. This + confirms the kernel is SWA-agnostic: SWA endpoints feed the same shape of input here after their + host-side gather, so the contract is symmetric across FULL / SWA groups. + """ + # Slot 4 here could equally be a FULL-group raw out_cache_loc value, or the result of an SWA + # endpoint's host gather. The kernel can't tell the difference and that's the point. + _run_write( + buf_pair=self.buf_pair, + input_ids=[55], + positions=[0], + out_cache_loc=[4], + ) + + stored_token, _, _, _ = read_slot_fields( + canary_buf=self.buf_pair[0], slot_idx=4 + ) + assert stored_token == 55 + + def test_padding_block_skipped(self) -> None: + """``blockIdx.x >= write_num_valid_reqs[0]`` → block early-exits, no write to canary_buf.""" + # Allocate plan with req_capacity=4 but only declare 1 active req. + _run_write( + buf_pair=self.buf_pair, + input_ids=[1], + positions=[0], + out_cache_loc=[0], + req_capacity=4, + ) + + # Only slot 0 should have been written; padding blocks 1..3 must not touch the buffer. + stored_token, _, _, _ = read_slot_fields( + canary_buf=self.buf_pair[0], slot_idx=0 + ) + assert stored_token == 1 + for slot_idx in (1, 2, 3): + stored_token_other, _, _, _ = read_slot_fields( + canary_buf=self.buf_pair[0], slot_idx=slot_idx + ) + assert stored_token_other == 0 + + def test_write_skip_when_out_cache_loc_is_minus_one(self) -> None: + """out_cache_loc[i] = -1 → that entry's slot is untouched by write kernel.""" + cuda_buf = self.buf_pair[0] + cuda_buf_before_slot_view = cuda_buf.view(torch.int64).clone() + + _run_write( + buf_pair=self.buf_pair, + input_ids=[100, 200, 300], + positions=[0, 1, 2], + out_cache_loc=[5, -1, 7], + ) + + after = cuda_buf.view(torch.int64) + for slot in range(cuda_buf.shape[0]): + if slot in (5, 7): + continue + assert torch.equal( + after[slot], cuda_buf_before_slot_view[slot] + ), f"slot {slot} should not have been written" + + def test_shrink_active_reqs_does_not_write_stale_slots(self) -> None: + """Run write with bs=3 plan after a bs=8 run on same buffer: stale slots from bs=8 stay intact.""" + cuda_buf = make_canary_buf(num_slots=32, slot_stride_bytes=32, device=_DEVICE) + ref_buf = cuda_buf.clone() + buf_pair = (cuda_buf, ref_buf) + + big_slots = list(range(1, 9)) + _run_write( + buf_pair=buf_pair, + input_ids=list(range(100, 108)), + positions=[0] * 8, + out_cache_loc=big_slots, + write_offsets=[0, 1, 2, 3, 4, 5, 6, 7, 8], + seed_slot_indices=[-1] * 8, + num_valid_reqs=8, + assert_equal=False, + ) + + untouched_snapshot = cuda_buf.view(torch.int64).clone() + + small_slots = [20, 21, 22] + _run_write( + buf_pair=buf_pair, + input_ids=[7, 8, 9], + positions=[0, 0, 0], + out_cache_loc=small_slots, + write_offsets=[0, 1, 2, 3], + seed_slot_indices=[-1, -1, -1], + num_valid_reqs=3, + assert_equal=False, + ) + + after = cuda_buf.view(torch.int64) + for slot in big_slots: + assert torch.equal( + after[slot], untouched_snapshot[slot] + ), f"slot {slot} from earlier bs=8 run was overwritten by bs=3 run" + + +class TestRunCounter: + def setup_method(self) -> None: + self.buf_pair = _make_default_buf_pair() + + def test_kernel_run_counter_per_call(self) -> None: + """``kernel_run_counter`` increments by 1 per call (even when ``write_num_valid_reqs == 0``).""" + plan_pair = make_write_plan_pair( + write_offsets=[0, 0], + seed_slot_indices=[-1], + num_valid_reqs=0, + device=_DEVICE, + ) + input_ids = _int32_tensor([0]) + positions = _int32_tensor([0]) + out_cache_loc = _int32_tensor([0]) + pseudo_tokens, pseudo_positions = dummy_pseudo_tensors(1) + cuda_log, ref_log = make_log_pair(device=_DEVICE) + + for _ in range(3): + _run_both_write( + cuda_canary_buf=self.buf_pair[0], + ref_canary_buf=self.buf_pair[1], + plan_cuda=plan_pair[0], + plan_ref=plan_pair[1], + input_ids=input_ids, + positions=positions, + out_cache_loc=out_cache_loc, + enable_write_verify_inputs=False, + expected_input_tokens=pseudo_tokens, + expected_input_positions=pseudo_positions, + cuda_log=cuda_log, + ref_log=ref_log, + assert_equal=False, + ) + + assert int(cuda_log.kernel_run_counter[0].item()) == 3 + assert_canary_state_equal(log_a=cuda_log, log_b=ref_log) + + def test_slot_run_counter_sums_entries(self) -> None: + """``slot_run_counter`` += sum(entry_count) across all active reqs in this call.""" + cuda_log, _ = _run_write( + buf_pair=self.buf_pair, + input_ids=[1, 2, 3, 4, 5], + positions=[0, 1, 0, 1, 2], + out_cache_loc=[0, 1, 2, 3, 4], + write_offsets=[0, 2, 5], + seed_slot_indices=[-1, -1], + num_valid_reqs=2, + ) + + assert int(cuda_log.slot_run_counter[0].item()) == 5 + + +class TestMisc: + def test_empty_plan_no_op(self) -> None: + """``write_num_valid_reqs = 0`` → no buf write, no slot_run_counter bump, only kernel_run_counter += 1.""" + cuda_log, _ = _run_write( + buf_pair=_make_default_buf_pair(), + input_ids=[0], + positions=[0], + out_cache_loc=[0], + write_offsets=[0, 0], + num_valid_reqs=0, + req_capacity=4, + ) + + assert int(cuda_log.write_index[0].item()) == 0 + assert int(cuda_log.slot_run_counter[0].item()) == 0 + assert int(cuda_log.kernel_run_counter[0].item()) == 1 + + def test_disabled_input_verify_does_not_deref_expected_inputs(self) -> None: + """``enable_write_verify_inputs=False`` must short-circuit before any expected_input_* dereference; + poison those tensors with 0x7F7F7F7F and assert nothing records a violation.""" + cuda_buf = make_canary_buf(num_slots=8, slot_stride_bytes=32, device=_DEVICE) + ref_buf = cuda_buf.clone() + buf_pair = (cuda_buf, ref_buf) + + garbage_expected_tokens = torch.full( + (1,), 0x7F7F7F7F, dtype=torch.int64, device=_DEVICE + ) + garbage_expected_positions = torch.full( + (1,), 0x7F7F7F7F, dtype=torch.int64, device=_DEVICE + ) + + cuda_log, ref_log = _run_write( + buf_pair=buf_pair, + input_ids=[42], + positions=[0], + out_cache_loc=[0], + expected_input_tokens=garbage_expected_tokens, + expected_input_positions=garbage_expected_positions, + assert_equal=False, + ) + + assert int(cuda_log.write_index[0].item()) == 0 + assert int(ref_log.write_index[0].item()) == 0 + + def test_disabled_assert_inputs_passes_none_to_cuda(self) -> None: + """Disabled input assertions pass None instead of dummy tensors.""" + canary_buf = torch.zeros( + 4, CANARY_SLOT_BYTES, dtype=torch.uint8, device=_DEVICE + ) + plan = make_write_plan( + write_offsets=[0, 0], + seed_slot_indices=[-1], + num_valid_reqs=0, + device=_DEVICE, + ) + input_ids = torch.zeros(1, dtype=torch.int64, device=_DEVICE) + positions = torch.zeros(1, dtype=torch.int64, device=_DEVICE) + out_cache_loc = torch.zeros(1, dtype=torch.int64, device=_DEVICE) + log = FakeViolationLog.allocate(capacity=2, device=_DEVICE) + context = VerifyOrWriteContext( + canary_buf=canary_buf, + kernel_kind=CanaryLaunchTag.HEAD_K_FULL, + violation_ring=log.ring, + violation_write_index=log.write_index, + slot_run_counter=log.slot_run_counter, + kernel_run_counter=log.kernel_run_counter, + enable_chain_position_assert=log.enable_chain_position_assert, + ) + module = _RecordingWriteModule() + + with patch.object(write_module, "_jit_canary_write_module", lambda: module): + launch_canary_write_kernel( + context=context, + plan=plan, + input_ids=input_ids, + positions=positions, + out_cache_loc=out_cache_loc, + enable_write_input_assert=False, + expected_input_tokens=None, + expected_input_positions=None, + ) + + assert len(module.calls) == 1 + call = module.calls[0] + assert call[8] == 0 + assert call[9] is None + assert call[10] is None + + +class TestBoundarySweep: + @pytest.mark.parametrize( + "token_val", + [0, 1, 0xFFFFFFFF, -1, 0x80000000, 0x7FFFFFFF], + ) + def test_token_boundary_byte_equal_sweep(self, token_val: int) -> None: + """Sweep token boundary values; assert CUDA write vs ref buf + state byte-equal.""" + _run_write_single_slot_byte_equal(_WriteSingleSlotInput(token=token_val)) + + @pytest.mark.parametrize( + "position_val", + [0, 1, 127, 128, 129, 0x7FFFFFFF], + ) + def test_position_boundary_byte_equal_sweep(self, position_val: int) -> None: + """Sweep position boundary values; assert CUDA write vs ref buf + state byte-equal.""" + _run_write_single_slot_byte_equal(_WriteSingleSlotInput(position=position_val)) + + +class TestPseudoMode: + def setup_method(self) -> None: + self.buf_pair = _make_default_buf_pair() + + def test_pseudo_mode_on_catches_token_mismatch(self) -> None: + """enable_write_verify_inputs=ON + intentional token mismatch → WRITE_TOKEN_MISMATCH bit recorded.""" + input_ids = _int32_tensor([10, 20, 30, 40, 50]) + positions = _int32_tensor([0, 1, 2, 3, 4]) + out_cache_loc = _int32_tensor([1, 2, 3, 4, 5]) + pseudo_tokens = _int32_tensor([10, 20, 30, 999, 50]) + pseudo_positions = positions.clone() + + cuda_log, _ = _run_write( + buf_pair=self.buf_pair, + input_ids=input_ids, + positions=positions, + out_cache_loc=out_cache_loc, + enable_write_verify_inputs=True, + expected_input_tokens=pseudo_tokens, + expected_input_positions=pseudo_positions, + ) + assert int(cuda_log.write_index[0].item()) >= 1 + bits = int(cuda_log.ring[0, consts.VIOLATION_FIELD_FAIL_REASON_BITS].item()) + assert ( + bits & consts.FailReason.WRITE_TOKEN_MISMATCH + ), f"expected WRITE_TOKEN_MISMATCH bit, got {bits:#b}" + + def test_pseudo_mode_off_skips_token_check(self) -> None: + """enable_write_verify_inputs=False makes the caller pass no expected-input tensors.""" + cuda_log, _ = _run_write( + buf_pair=self.buf_pair, + input_ids=[10, 20, 30], + positions=[0, 1, 2], + out_cache_loc=[1, 2, 3], + expected_input_tokens=_int32_tensor([99, 99, 99]), + expected_input_positions=_int32_tensor([99, 99, 99]), + ) + assert int(cuda_log.write_index[0].item()) == 0 + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"]))