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