minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Add real-data KV verification to the KV-canary (#26817)

Browse Source
This commit is contained in:
fzyzcjy
2026-05-31 09:58:32 +08:00
committed by GitHub
parent cdee16e144
commit 0ca610a6df
45 changed files with 2695 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/tests/kv_canary/_canary_helpers.py
+18 -4
View File
@@ -16,6 +16,10 @@ from sglang.jit_kernel.tests.kv_canary._constants import (
DEFAULT_RING_CAPACITY,
DEFAULT_SLOT_STRIDE_BYTES,
)
from sglang.jit_kernel.tests.kv_canary._fixtures import (
make_real_kv_source,
make_real_kv_sources,
)
__all__ = [
"FakeViolationLog",
@@ -26,6 +30,8 @@ __all__ = [
"make_canary_buf",
"make_canary_buf_pair",
"make_log_pair",
"make_real_kv_source",
"make_real_kv_sources",
"make_verify_plan",
"make_verify_plan_pair",
"make_write_plan",
@@ -111,7 +117,7 @@ def make_verify_plan(
``expected_input_ids`` defaults to ``[-1] * n_active`` (the verify-kernel
"skip token check" sentinel) so existing tests that only exercise the
chain / position paths keep working unchanged.
chain / position / real-kv-hash paths keep working unchanged.
"""
n_active = len(slot_indices)
if not (len(positions) == n_active and len(prev_slot_indices) == n_active):
@@ -246,13 +252,13 @@ def write_slot_fields(
token: int,
position: int,
prev_hash: int,
real_kv_hash: int,
) -> None:
view = canary_buf.view(torch.int64)
view[slot_idx, 0] = token
view[slot_idx, 1] = position
view[slot_idx, 2] = prev_hash
# field[3] (real_kv_hash) is always 0 in the naive canary.
view[slot_idx, 3] = 0
view[slot_idx, 3] = real_kv_hash
def stamp_pair(
@@ -262,6 +268,7 @@ def stamp_pair(
token: int,
position: int,
prev_hash: int,
real_kv_hash: int = 0,
) -> None:
"""Stamp the same slot fields into both (cuda, ref) canary buffers."""
for buf in buf_pair:
@@ -271,6 +278,7 @@ def stamp_pair(
token=token,
position=position,
prev_hash=prev_hash,
real_kv_hash=real_kv_hash,
)
@@ -288,10 +296,15 @@ def stamp_clean_chain(
slot_indices: list[int],
tokens: list[int],
positions: list[int],
real_kv_hashes: Optional[list[int]] = None,
) -> list[int]:
n = len(tokens)
real_kv_hashes = real_kv_hashes if real_kv_hashes is not None else [0] * n
running_prev_hash = splitmix64(consts.CANARY_CHAIN_ANCHOR)
stored_prev_hashes: list[int] = []
for slot_idx, token, position in zip(slot_indices, tokens, positions):
for slot_idx, token, position, real_kv_hash in zip(
slot_indices, tokens, positions, real_kv_hashes
):
signed_prev = to_signed_int64(running_prev_hash)
for buf in (cuda_buf, ref_buf):
write_slot_fields(
@@ -300,6 +313,7 @@ def stamp_clean_chain(
token=token,
position=position,
prev_hash=signed_prev,
real_kv_hash=to_signed_int64(real_kv_hash),
)
stored_prev_hashes.append(signed_prev)
running_prev_hash = splitmix64_mix3(running_prev_hash, token, position)
python/sglang/jit_kernel/tests/kv_canary/_differential.py
+48 -2
View File
@@ -5,12 +5,14 @@ from typing import Any, Callable, Iterator, Optional
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,
RealKvSource,
VerifyOrWriteContext,
VerifyPlan,
launch_canary_verify_kernel,
@@ -186,6 +188,9 @@ def _run_both_verify(
plan_ref,
cuda_log: FakeViolationLog,
ref_log: FakeViolationLog,
real_kv_sources_cuda: tuple[RealKvSource, ...],
real_kv_sources_ref: tuple[RealKvSource, ...],
real_kv_hash_mode: consts.RealKvHashMode,
kernel_kind: CanaryLaunchTag = CanaryLaunchTag.HEAD_K_FULL,
assert_equal: bool = True,
check_verify_expected_token: bool = True,
@@ -199,6 +204,8 @@ def _run_both_verify(
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,
real_kv_sources=real_kv_sources_cuda,
real_kv_hash_mode=real_kv_hash_mode,
),
plan=plan_cuda,
check_verify_expected_token=check_verify_expected_token,
@@ -212,6 +219,8 @@ def _run_both_verify(
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,
real_kv_sources=real_kv_sources_ref,
real_kv_hash_mode=real_kv_hash_mode,
),
plan=plan_ref,
check_verify_expected_token=check_verify_expected_token,
@@ -236,6 +245,9 @@ def _run_both_write(
expected_input_positions: torch.Tensor,
cuda_log: FakeViolationLog,
ref_log: FakeViolationLog,
real_kv_sources_cuda: tuple[RealKvSource, ...],
real_kv_sources_ref: tuple[RealKvSource, ...],
real_kv_hash_mode: consts.RealKvHashMode,
kernel_kind: CanaryLaunchTag = CanaryLaunchTag.HEAD_K_FULL,
assert_equal: bool = True,
) -> None:
@@ -254,6 +266,8 @@ def _run_both_write(
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,
real_kv_sources=real_kv_sources_cuda,
real_kv_hash_mode=real_kv_hash_mode,
),
plan=plan_cuda,
input_ids=input_ids,
@@ -272,6 +286,8 @@ def _run_both_write(
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,
real_kv_sources=real_kv_sources_ref,
real_kv_hash_mode=real_kv_hash_mode,
),
plan=plan_ref,
input_ids=input_ids,
@@ -381,6 +397,20 @@ def _yield_simpler(inputs: Any) -> Iterator[tuple[str, Any]]:
if fields["extras_count"] > 0:
yield from emit("extras_zero", extras_count=0)
if "real_kv_hash_mode" in fields:
cur = fields["real_kv_hash_mode"]
if hasattr(cur, "value"):
cls = cur.__class__
if int(cur) == 2:
yield from emit("hash_mode_bit", real_kv_hash_mode=cls(1))
elif int(cur) == 1:
yield from emit("hash_mode_off", real_kv_hash_mode=cls(0))
if "real_kv_sources" in fields:
srcs = fields["real_kv_sources"]
if isinstance(srcs, tuple) and len(srcs) > 1:
yield from emit("sources_to_one", real_kv_sources=srcs[:1])
if "enable_write_verify_inputs" in fields:
cur = fields["enable_write_verify_inputs"]
if hasattr(cur, "value") and int(cur) != 0:
@@ -398,13 +428,18 @@ def run_verify_diff(
*,
buf_pair: tuple[torch.Tensor, torch.Tensor],
plan_pair: tuple[VerifyPlan, VerifyPlan],
real_kv_sources_pair: tuple[tuple[RealKvSource, ...], tuple[RealKvSource, ...]] = (
(),
(),
),
real_kv_hash_mode: consts.RealKvHashMode = consts.RealKvHashMode.NONE,
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.
buf/plan/source 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(
@@ -414,6 +449,9 @@ def run_verify_diff(
plan_ref=plan_pair[1],
cuda_log=cuda_log,
ref_log=ref_log,
real_kv_sources_cuda=real_kv_sources_pair[0],
real_kv_sources_ref=real_kv_sources_pair[1],
real_kv_hash_mode=real_kv_hash_mode,
kernel_kind=kernel_kind,
assert_equal=assert_equal,
check_verify_expected_token=check_verify_expected_token,
@@ -431,12 +469,17 @@ def run_write_diff(
expected_input_tokens: torch.Tensor,
expected_input_positions: torch.Tensor,
enable_write_verify_inputs: bool = False,
real_kv_sources_pair: tuple[tuple[RealKvSource, ...], tuple[RealKvSource, ...]] = (
(),
(),
),
real_kv_hash_mode: consts.RealKvHashMode = consts.RealKvHashMode.NONE,
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.
buf/plan/source 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(
@@ -452,6 +495,9 @@ def run_write_diff(
expected_input_positions=expected_input_positions,
cuda_log=cuda_log,
ref_log=ref_log,
real_kv_sources_cuda=real_kv_sources_pair[0],
real_kv_sources_ref=real_kv_sources_pair[1],
real_kv_hash_mode=real_kv_hash_mode,
kernel_kind=kernel_kind,
assert_equal=assert_equal,
)
python/sglang/jit_kernel/tests/kv_canary/_fixtures.py
+82 -1
View File
@@ -5,8 +5,12 @@ from typing import Literal, Optional
import torch
from sglang.jit_kernel.kv_canary.verify import VerifyPlan
from sglang.jit_kernel.kv_canary.verify import (
RealKvSource,
VerifyPlan,
)
from sglang.jit_kernel.kv_canary.write import WritePlan
from sglang.jit_kernel.tests.kv_canary._constants import DEFAULT_NUM_SLOTS
_DEVICE = torch.device("cuda")
@@ -72,6 +76,83 @@ def make_req_to_token(
return rtt.contiguous()
def make_real_kv_source(
*,
num_slots: int = DEFAULT_NUM_SLOTS,
num_bytes_per_token: int = 16,
page_size: int = 1,
read_bytes: Optional[int] = None,
pad_dim1: int = 0,
device: torch.device,
fill: int = 0,
) -> RealKvSource:
"""Allocate one RealKvSource with the canonical [num_rows, dim1_bytes] uint8 shape.
``pad_dim1`` adds trailing per-row bytes the canary should skip — used by the "holey dim 1" case to
confirm the kernel never reads past ``page_size * num_bytes_per_token``.
"""
num_rows = (num_slots + page_size - 1) // page_size
cols = page_size * num_bytes_per_token + pad_dim1
tensor = torch.full(
(num_rows, cols), fill_value=fill, dtype=torch.uint8, device=device
)
effective_read = read_bytes if read_bytes is not None else num_bytes_per_token
return RealKvSource(
tensor=tensor,
page_size=page_size,
num_bytes_per_token=num_bytes_per_token,
read_bytes=effective_read,
)
FillStrategy = Literal["constant_per_source", "random_bytes"]
def make_real_kv_sources(
*,
count: int,
num_bytes_per_token: int = 16,
page_size: int = 1,
num_slots: int = DEFAULT_NUM_SLOTS,
device: torch.device,
rng: Optional[random.Random] = None,
fill_strategy: FillStrategy = "constant_per_source",
) -> tuple[RealKvSource, ...]:
sources: list[RealKvSource] = []
for i in range(count):
read_bytes_eff = num_bytes_per_token
src = make_real_kv_source(
num_slots=num_slots,
num_bytes_per_token=num_bytes_per_token,
page_size=page_size,
read_bytes=read_bytes_eff,
device=device,
fill=(i + 1) * 17,
)
if fill_strategy == "random_bytes":
if rng is None:
rng = random.Random(0)
seed = rng.randint(0, 0xFFFFFFFF)
gen = torch.Generator(device=device).manual_seed(seed)
src.tensor.random_(generator=gen)
sources.append(src)
return tuple(sources)
def clone_real_kv_sources(
sources: tuple[RealKvSource, ...],
) -> tuple[RealKvSource, ...]:
return tuple(
RealKvSource(
tensor=src.tensor.clone(),
page_size=src.page_size,
num_bytes_per_token=src.num_bytes_per_token,
read_bytes=src.read_bytes,
)
for src in sources
)
PaddingKind = Literal["none", "trailing", "interleaved"]
python/sglang/jit_kernel/tests/kv_canary/_hand_oracle.py
+31
View File
@@ -0,0 +1,31 @@
"""Hand-computed Python re-implementation of the real-kv-source fold, kept independent from
``verify_ref._splitmix64_fold_bytes_scalar`` so a ref / kernel co-regression cannot silently fix the
diff comparison."""
from __future__ import annotations
from sglang.jit_kernel.kv_canary.consts import splitmix64
def _fold_words(padded: bytes) -> int:
"""Pack padded bytes little-endian into 8-byte words, fold each via splitmix64 from acc=0."""
num_words = len(padded) // 8
acc = 0
for w in range(num_words):
chunk = padded[w * 8 : (w + 1) * 8]
word = sum(b << (8 * k) for k, b in enumerate(chunk))
acc = splitmix64(acc ^ word)
return splitmix64(0 ^ acc)
def _hand_fold_partial(raw_bytes: bytes) -> int:
"""PARTIAL-mode fold: first min(16, len) bytes, little-endian word-pack + splitmix64, same as ALL."""
truncated = raw_bytes[: min(16, len(raw_bytes))]
pad = (8 - len(truncated) % 8) % 8
return _fold_words(bytes(truncated) + bytes(pad))
def _hand_fold_all(raw_bytes: bytes) -> int:
"""ALL-mode fold: pack bytes little-endian into 8-byte words, fold each via splitmix64, then mix into acc=0."""
pad = (8 - len(raw_bytes) % 8) % 8
return _fold_words(raw_bytes + bytes(pad))
python/sglang/jit_kernel/tests/kv_canary/test_const_sync.py
+1 -1
View File
@@ -52,7 +52,7 @@ def test_int_consts_sync() -> None:
def test_enums_sync() -> None:
cuh = _CONSTS_CUH.read_text(encoding="utf-8")
for enum_name in ("FailReason",):
for enum_name in ("RealKvHashMode", "FailReason"):
cpp_members = _parse_enum_class(cuh, enum_name)
py_enum = getattr(consts, enum_name)
cpp_normalized = {
python/sglang/jit_kernel/tests/kv_canary/test_kernel_config.py
+15
View File
@@ -3,6 +3,7 @@ from __future__ import annotations
import pytest
import torch
from sglang.jit_kernel.kv_canary import consts
from sglang.jit_kernel.kv_canary.verify import (
CanaryLaunchTag,
VerifyOrWriteContext,
@@ -119,6 +120,9 @@ def test_verify_byte_equal_across_repeated_launches_10x() -> None:
plan_ref=plan_ref,
cuda_log=cuda_log,
ref_log=ref_log,
real_kv_sources_cuda=(),
real_kv_sources_ref=(),
real_kv_hash_mode=consts.RealKvHashMode.NONE,
kernel_kind=CanaryLaunchTag.HEAD_K_FULL,
)
@@ -171,6 +175,9 @@ def test_write_byte_equal_across_repeated_launches_10x() -> None:
expected_input_positions=pseudo_pos,
cuda_log=cuda_log,
ref_log=ref_log,
real_kv_sources_cuda=(),
real_kv_sources_ref=(),
real_kv_hash_mode=consts.RealKvHashMode.NONE,
kernel_kind=CanaryLaunchTag.HEAD_K_FULL,
)
@@ -284,6 +291,8 @@ def test_verify_multi_launch_100x_counter_linear() -> None:
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,
real_kv_sources=(),
real_kv_hash_mode=consts.RealKvHashMode.NONE,
),
plan=plan_cuda,
check_verify_expected_token=True,
@@ -339,6 +348,9 @@ def test_verify_check_disabled_byte_equal() -> None:
plan_ref=plan_true_ref,
cuda_log=cuda_log_true,
ref_log=ref_log_true,
real_kv_sources_cuda=(),
real_kv_sources_ref=(),
real_kv_hash_mode=consts.RealKvHashMode.NONE,
kernel_kind=CanaryLaunchTag.HEAD_K_FULL,
check_verify_expected_token=True,
)
@@ -349,6 +361,9 @@ def test_verify_check_disabled_byte_equal() -> None:
plan_ref=plan_false_ref,
cuda_log=cuda_log_false,
ref_log=ref_log_false,
real_kv_sources_cuda=(),
real_kv_sources_ref=(),
real_kv_hash_mode=consts.RealKvHashMode.NONE,
kernel_kind=CanaryLaunchTag.HEAD_K_FULL,
check_verify_expected_token=False,
)
python/sglang/jit_kernel/tests/kv_canary/test_pipeline_e2e.py
+52
View File
@@ -12,6 +12,7 @@ from sglang.jit_kernel.kv_canary.plan_ref import (
)
from sglang.jit_kernel.kv_canary.verify import (
CanaryLaunchTag,
RealKvSource,
VerifyOrWriteContext,
VerifyPlan,
launch_canary_verify_kernel,
@@ -28,10 +29,12 @@ from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
assert_canary_buf_equal,
assert_canary_state_equal,
make_canary_buf,
make_real_kv_sources,
stamp_clean_chain,
write_slot_fields,
)
from sglang.jit_kernel.tests.kv_canary._fixtures import (
clone_real_kv_sources,
empty_extras,
make_req_to_token,
)
@@ -62,6 +65,8 @@ def _run_pipeline(
enable_write_verify_inputs: bool,
expected_input_tokens: torch.Tensor,
expected_input_positions: torch.Tensor,
real_kv_sources: tuple[RealKvSource, ...],
real_kv_hash_mode: consts.RealKvHashMode,
verify_capacity: int,
write_req_capacity: int,
req_to_verify_expected_tokens: Optional[torch.Tensor] = None,
@@ -115,6 +120,8 @@ def _run_pipeline(
slot_run_counter=log.slot_run_counter,
kernel_run_counter=log.kernel_run_counter,
enable_chain_position_assert=log.enable_chain_position_assert,
real_kv_sources=real_kv_sources,
real_kv_hash_mode=real_kv_hash_mode,
)
launch_canary_write_kernel(
context=context,
@@ -142,6 +149,8 @@ def _run_pipeline(
slot_run_counter=log.slot_run_counter,
kernel_run_counter=log.kernel_run_counter,
enable_chain_position_assert=log.enable_chain_position_assert,
real_kv_sources=real_kv_sources,
real_kv_hash_mode=real_kv_hash_mode,
),
plan=plan_w,
input_ids=input_ids,
@@ -160,6 +169,8 @@ def _run_pipeline(
slot_run_counter=log.slot_run_counter,
kernel_run_counter=log.kernel_run_counter,
enable_chain_position_assert=log.enable_chain_position_assert,
real_kv_sources=real_kv_sources,
real_kv_hash_mode=real_kv_hash_mode,
),
plan=plan_v,
check_verify_expected_token=check_verify_expected_token,
@@ -185,6 +196,9 @@ def _run_both_and_assert_pipeline_equal(
enable_write_verify_inputs: bool = False,
expected_input_tokens: Optional[torch.Tensor] = None,
expected_input_positions: Optional[torch.Tensor] = None,
real_kv_sources_real: tuple[RealKvSource, ...] = (),
real_kv_sources_ref: tuple[RealKvSource, ...] = (),
real_kv_hash_mode: consts.RealKvHashMode = consts.RealKvHashMode.NONE,
ring_capacity: int = 64,
verify_capacity: int = 256,
write_req_capacity: int = 16,
@@ -240,6 +254,7 @@ def _run_both_and_assert_pipeline_equal(
enable_write_verify_inputs=enable_write_verify_inputs,
expected_input_tokens=expected_input_tokens,
expected_input_positions=expected_input_positions,
real_kv_hash_mode=real_kv_hash_mode,
verify_capacity=verify_capacity,
write_req_capacity=write_req_capacity,
req_to_verify_expected_tokens=req_to_verify_expected_tokens,
@@ -251,12 +266,14 @@ def _run_both_and_assert_pipeline_equal(
real=True,
canary_buf=buf_real,
log=log_real,
real_kv_sources=real_kv_sources_real,
**shared,
)
plan_v_ref, plan_w_ref = _run_pipeline(
real=False,
canary_buf=buf_ref,
log=log_ref,
real_kv_sources=real_kv_sources_ref,
**shared,
)
@@ -436,6 +453,35 @@ def test_pipeline_sweep_no_write() -> None:
assert int(log_ref.slot_run_counter[0].item()) == prefix_len
@pytest.mark.parametrize(
"real_kv_hash_mode",
[
consts.RealKvHashMode.NONE,
consts.RealKvHashMode.PARTIAL,
consts.RealKvHashMode.ALL,
],
)
def test_pipeline_real_kv_mode(real_kv_hash_mode: consts.RealKvHashMode) -> None:
"""real_kv_hash_mode OFF/PARTIAL/ALL: real and ref use cloned sources to prevent ALL-mode hash aliasing."""
sources_real = make_real_kv_sources(count=2, num_slots=64, device=_DEVICE)
sources_ref = clone_real_kv_sources(sources_real)
_run_both_and_assert_pipeline_equal(
req_pool_indices=_t([1]),
prefix_lens=_t([0]),
extend_seq_lens=_t([3]),
input_ids=_t([5, 6, 7]),
positions=_t([0, 1, 2]),
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(),
real_kv_sources_real=sources_real,
real_kv_sources_ref=sources_ref,
real_kv_hash_mode=real_kv_hash_mode,
)
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])
@@ -560,6 +606,7 @@ def test_pipeline_ring_overflow_via_real_plan() -> None:
token=slot_idx + 1,
position=slot_idx,
prev_hash=0x1234_DEAD_BEEF_0000 + slot_idx,
real_kv_hash=0,
)
# Step 2: run real pipeline (plan + no write + verify); overflow ring capacity=4 with all n_slots violations.
@@ -609,6 +656,8 @@ def test_pipeline_ring_overflow_via_real_plan() -> None:
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,
real_kv_sources=(),
real_kv_hash_mode=consts.RealKvHashMode.NONE,
),
plan=plan_v_real,
check_verify_expected_token=True,
@@ -624,6 +673,8 @@ def test_pipeline_ring_overflow_via_real_plan() -> None:
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,
real_kv_sources=(),
real_kv_hash_mode=consts.RealKvHashMode.NONE,
),
plan=plan_v_ref,
check_verify_expected_token=True,
@@ -649,6 +700,7 @@ def test_pipeline_kernel_kind_propagates(kernel_kind: CanaryLaunchTag) -> None:
token=7,
position=99,
prev_hash=0,
real_kv_hash=0,
)
_, _, log_real, log_ref, _, _, _, _ = _run_both_and_assert_pipeline_equal(
python/sglang/jit_kernel/tests/kv_canary/test_utils.py
+2
View File
@@ -50,6 +50,8 @@ def test_cases_to_x_vals_includes_scenario_axis() -> None:
case.mode,
case.extend_len,
case.pool_kind,
case.real_kv_kind,
case.hash_mode,
)
]
python/sglang/jit_kernel/tests/kv_canary/test_verify_fuzz.py
+38 -2
View File
@@ -9,6 +9,7 @@ import torch
from sglang.jit_kernel.kv_canary import consts
from sglang.jit_kernel.kv_canary.verify import (
CanaryLaunchTag,
RealKvSource,
VerifyPlan,
)
from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
@@ -19,6 +20,10 @@ from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
stamp_clean_chain,
)
from sglang.jit_kernel.tests.kv_canary._differential import _run_both_verify
from sglang.jit_kernel.tests.kv_canary._fixtures import (
clone_real_kv_sources,
make_real_kv_sources,
)
from sglang.jit_kernel.tests.kv_canary._fuzz_driver import (
FUZZ_SEEDS_PR,
run_fuzz_combo,
@@ -41,16 +46,39 @@ class VerifyFuzzInputs:
plan_cuda: VerifyPlan
plan_ref: VerifyPlan
kernel_kind: CanaryLaunchTag
real_kv_sources_cuda: tuple[RealKvSource, ...]
real_kv_sources_ref: tuple[RealKvSource, ...]
real_kv_hash_mode: consts.RealKvHashMode
ring_capacity: int
check_verify_expected_token: bool
def _draw_random_verify_inputs(rng: random.Random) -> VerifyFuzzInputs:
hash_mode = rng.choice(
[
consts.RealKvHashMode.NONE,
consts.RealKvHashMode.PARTIAL,
consts.RealKvHashMode.ALL,
]
)
src_count = rng.choice([1, 2, 4])
page_size = rng.choice([1, 16])
bytes_per = rng.choice([16, 64, 128])
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])
sources_cuda = make_real_kv_sources(
count=src_count,
num_bytes_per_token=bytes_per,
page_size=page_size,
num_slots=num_slots,
device=_DEVICE,
rng=rng,
)
sources_ref = clone_real_kv_sources(sources_cuda)
cuda_buf = make_canary_buf(
num_slots=num_slots, slot_stride_bytes=32, device=_DEVICE
)
@@ -68,7 +96,7 @@ def _draw_random_verify_inputs(rng: random.Random) -> VerifyFuzzInputs:
else:
prev_slot_indices.append(slot_indices[i - 1])
if plan_size > 0:
if hash_mode == consts.RealKvHashMode.NONE and plan_size > 0:
stamp_clean_chain(
cuda_buf=cuda_buf,
ref_buf=ref_buf,
@@ -112,6 +140,9 @@ def _draw_random_verify_inputs(rng: random.Random) -> VerifyFuzzInputs:
plan_cuda=plan_cuda,
plan_ref=plan_ref,
kernel_kind=kernel_kind,
real_kv_sources_cuda=sources_cuda,
real_kv_sources_ref=sources_ref,
real_kv_hash_mode=hash_mode,
ring_capacity=ring_capacity,
check_verify_expected_token=check_verify_expected_token,
)
@@ -130,6 +161,9 @@ def _run_one(inputs: VerifyFuzzInputs) -> None:
plan_ref=inputs.plan_ref,
cuda_log=cuda_log,
ref_log=ref_log,
real_kv_sources_cuda=inputs.real_kv_sources_cuda,
real_kv_sources_ref=inputs.real_kv_sources_ref,
real_kv_hash_mode=inputs.real_kv_hash_mode,
kernel_kind=inputs.kernel_kind,
assert_equal=False,
check_verify_expected_token=inputs.check_verify_expected_token,
@@ -155,6 +189,8 @@ 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"hash_mode={inputs.real_kv_hash_mode.name} "
f"sources={len(inputs.real_kv_sources_cuda)} "
f"ring={inputs.ring_capacity} "
f"check_token={inputs.check_verify_expected_token}"
)
@@ -162,7 +198,7 @@ def _summarize(inputs: VerifyFuzzInputs) -> str:
@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."""
"""Multi-dim verify fuzzer: random hash mode × kernel kind × page × bytes × N iters, byte-equal."""
run_fuzz_combo(
seed,
draw_fn=_draw_random_verify_inputs,
python/sglang/jit_kernel/tests/kv_canary/test_verify_hand.py
+694 -33
View File
@@ -12,13 +12,18 @@ from sglang.jit_kernel.kv_canary import consts
from sglang.jit_kernel.kv_canary.consts import splitmix64, splitmix64_mix3
from sglang.jit_kernel.kv_canary.verify import (
CanaryLaunchTag,
RealKvSource,
VerifyOrWriteContext,
VerifyPlan,
launch_canary_verify_kernel,
)
from sglang.jit_kernel.kv_canary.verify_ref import (
_compute_real_kv_hash_scalar,
launch_canary_verify_kernel_torch_reference,
)
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_only_bits_set,
@@ -26,8 +31,11 @@ from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
make_canary_buf,
make_canary_buf_pair,
make_log_pair,
make_real_kv_source,
make_real_kv_sources,
make_verify_plan,
make_verify_plan_pair,
make_write_plan,
read_slot_fields,
stamp_clean_chain,
stamp_pair,
@@ -38,6 +46,11 @@ from sglang.jit_kernel.tests.kv_canary._differential import (
_run_both_verify,
run_verify_diff,
)
from sglang.jit_kernel.tests.kv_canary._fixtures import clone_real_kv_sources
from sglang.jit_kernel.tests.kv_canary._hand_oracle import (
_hand_fold_all,
_hand_fold_partial,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
@@ -66,6 +79,7 @@ def _stamp_head(
token: int = 42,
position: int = 0,
prev_hash: int | None = None,
real_kv_hash: int = 0,
) -> None:
"""``stamp_pair`` with ``prev_hash`` defaulting to ``chain_anchor_signed()`` (the chain-head value)."""
stamp_pair(
@@ -74,6 +88,7 @@ def _stamp_head(
token=token,
position=position,
prev_hash=chain_anchor_signed() if prev_hash is None else prev_hash,
real_kv_hash=real_kv_hash,
)
@@ -114,7 +129,7 @@ def _run_both_verify_no_rkv(
assert_equal: bool = True,
kernel_kind: CanaryLaunchTag = CanaryLaunchTag.HEAD_K_FULL,
) -> None:
"""``_run_both_verify`` — the most common in-place verify run."""
"""``_run_both_verify`` with empty real_kv sources / NONE mode — the most common in-place verify run."""
_run_both_verify(
cuda_canary_buf=buf_pair[0],
ref_canary_buf=buf_pair[1],
@@ -122,6 +137,9 @@ def _run_both_verify_no_rkv(
plan_ref=plan_pair[1],
cuda_log=cuda_log,
ref_log=ref_log,
real_kv_sources_cuda=(),
real_kv_sources_ref=(),
real_kv_hash_mode=consts.RealKvHashMode.NONE,
kernel_kind=kernel_kind,
assert_equal=assert_equal,
)
@@ -132,6 +150,9 @@ class _VerifySingleSlotInput:
token: int = 42
position: int = 0
stored_prev_hash_signed: int
stored_real_kv_hash_signed: int = 0
real_kv_sources: tuple[RealKvSource, ...] = ()
real_kv_hash_mode: consts.RealKvHashMode = consts.RealKvHashMode.NONE
def _run_verify_single_slot_byte_equal(case: _VerifySingleSlotInput) -> None:
@@ -142,14 +163,66 @@ def _run_verify_single_slot_byte_equal(case: _VerifySingleSlotInput) -> None:
token=case.token,
position=case.position,
prev_hash=case.stored_prev_hash_signed,
real_kv_hash=case.stored_real_kv_hash_signed,
)
sources_cuda = case.real_kv_sources
sources_ref = clone_real_kv_sources(sources_cuda)
plan_pair = _plan_pair_single(slot_idx=1, position=case.position)
run_verify_diff(
buf_pair=buf_pair,
plan_pair=plan_pair,
real_kv_sources_pair=(sources_cuda, sources_ref),
real_kv_hash_mode=case.real_kv_hash_mode,
)
def _stamp_clean_kv_chain(
*,
buf_pair: tuple[torch.Tensor, torch.Tensor],
sources_cuda: tuple[RealKvSource, ...],
input_ids: torch.Tensor,
positions: torch.Tensor,
out_cache_loc: torch.Tensor,
real_kv_hash_mode: consts.RealKvHashMode,
) -> None:
"""Use the Python write ref impl to populate the canary buf for a fresh chain.
Lets verify tests start from a known-good chain without re-implementing splitmix64 by hand.
"""
n = int(input_ids.shape[0])
cuda_buf, ref_buf = buf_pair
write_plan = make_write_plan(
write_offsets=[0, n],
seed_slot_indices=[-1],
num_valid_reqs=1,
device=_DEVICE,
)
log = FakeViolationLog.allocate(device=_DEVICE)
# enable_write_input_assert=False is hard-wired here, so the kernel API requires the
# expected_* tensors be None (otherwise it raises ValueError).
launch_canary_write_kernel_torch_reference(
context=VerifyOrWriteContext(
canary_buf=cuda_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,
real_kv_sources=sources_cuda,
real_kv_hash_mode=real_kv_hash_mode,
),
plan=write_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,
)
ref_buf.copy_(cuda_buf)
# ---------------------------------------------------------------------------
# Kernel-contract invariants.
# ---------------------------------------------------------------------------
@@ -194,11 +267,12 @@ class TestChain:
tokens = [101, 202, 303, 404, 505]
positions = [0, 1, 2, 3, 4]
slot_indices = [1, 2, 3, 4, 5]
real_kv_hashes = [0, 0, 0, 0, 0]
# 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):
for token, position, real_kv_hash in zip(tokens, positions, real_kv_hashes):
expected_prev_hashes_u64.append(running)
running = splitmix64_mix3(running, token, position)
expected_prev_hashes_signed = [
@@ -447,7 +521,40 @@ class TestViolationField:
_fail_bits(cuda_log), consts.FailReason.VERIFY_CHAIN_HASH_MISMATCH
)
@pytest.mark.parametrize("bit_to_trigger", ["POSITION", "PREV_HASH"])
def test_violation_real_kv_hash_mismatch(self) -> None:
"""Mutate one byte of a RealKvSource tensor after writing the chain → REAL_KV_HASH bit on verify."""
buf_pair = _buf_pair()
sources_cuda = make_real_kv_sources(count=1, device=_DEVICE)
# Step: write a chain with real_kv_hash mixin, then mutate one byte in the source tensors so the next
# verify reconstructs a hash that differs from the stored one.
_stamp_clean_kv_chain(
buf_pair=buf_pair,
sources_cuda=sources_cuda,
input_ids=torch.tensor([7, 8, 9], dtype=torch.int64, device=_DEVICE),
positions=torch.tensor([0, 1, 2], dtype=torch.int64, device=_DEVICE),
out_cache_loc=torch.tensor([1, 2, 3], dtype=torch.int64, device=_DEVICE),
real_kv_hash_mode=consts.RealKvHashMode.ALL,
)
# Mutate one byte in BOTH copies so the verify recomputed hash diverges from stored.
sources_ref = clone_real_kv_sources(sources_cuda)
sources_cuda[0].tensor[1, 0] ^= 0xFF
sources_ref[0].tensor.copy_(sources_cuda[0].tensor)
plan_pair = _plan_pair_single(slot_idx=1, position=0)
cuda_log, _ = run_verify_diff(
buf_pair=buf_pair,
plan_pair=plan_pair,
real_kv_sources_pair=(sources_cuda, sources_ref),
real_kv_hash_mode=consts.RealKvHashMode.ALL,
)
assert_only_bits_set(
_fail_bits(cuda_log), consts.FailReason.VERIFY_REAL_KV_HASH_MISMATCH
)
@pytest.mark.parametrize("bit_to_trigger", ["POSITION", "PREV_HASH", "REAL_KV"])
@pytest.mark.parametrize("injection_position", ["head", "mid", "last"])
@pytest.mark.parametrize("ring_state", ["open", "full"])
def test_violation_bit_injection_position_ring_state_matrix(
@@ -468,41 +575,67 @@ class TestViolationField:
expected_bit = {
"POSITION": consts.FailReason.VERIFY_POSITION_MISMATCH,
"PREV_HASH": consts.FailReason.VERIFY_CHAIN_HASH_MISMATCH,
"REAL_KV": consts.FailReason.VERIFY_REAL_KV_HASH_MISMATCH,
}[bit_to_trigger]
cuda_buf, ref_buf = _buf_pair()
buf_pair = (cuda_buf, ref_buf)
stamp_clean_chain(
cuda_buf=cuda_buf,
ref_buf=ref_buf,
tokens=tokens,
positions=positions,
slot_indices=slot_indices,
)
if bit_to_trigger == "POSITION":
stored_token, stored_pos, stored_prev, _ = read_slot_fields(
canary_buf=cuda_buf, slot_idx=corrupt_slot
)
stamp_pair(
buf_pair,
slot_idx=corrupt_slot,
token=stored_token,
position=stored_pos + 99,
prev_hash=stored_prev,
if bit_to_trigger == "REAL_KV":
buf_pair = _buf_pair()
sources_cuda = make_real_kv_sources(count=1, device=_DEVICE)
_stamp_clean_kv_chain(
buf_pair=buf_pair,
sources_cuda=sources_cuda,
input_ids=torch.tensor(tokens, dtype=torch.int64, device=_DEVICE),
positions=torch.tensor(positions, dtype=torch.int64, device=_DEVICE),
out_cache_loc=torch.tensor(
slot_indices, dtype=torch.int64, device=_DEVICE
),
real_kv_hash_mode=consts.RealKvHashMode.ALL,
)
sources_ref = clone_real_kv_sources(sources_cuda)
sources_cuda[0].tensor[corrupt_slot, 0] ^= 0xFF
sources_ref[0].tensor.copy_(sources_cuda[0].tensor)
real_kv_hash_mode = consts.RealKvHashMode.ALL
real_kv_sources_cuda = sources_cuda
real_kv_sources_ref = sources_ref
else:
stored_token, stored_pos, stored_prev, _ = read_slot_fields(
canary_buf=cuda_buf, slot_idx=corrupt_slot
)
flipped_prev = stored_prev ^ 1
stamp_pair(
buf_pair,
slot_idx=corrupt_slot,
token=stored_token,
position=stored_pos,
prev_hash=flipped_prev,
cuda_buf, ref_buf = _buf_pair()
buf_pair = (cuda_buf, ref_buf)
stamp_clean_chain(
cuda_buf=cuda_buf,
ref_buf=ref_buf,
tokens=tokens,
positions=positions,
slot_indices=slot_indices,
)
real_kv_hash_mode = consts.RealKvHashMode.NONE
real_kv_sources_cuda = ()
real_kv_sources_ref = ()
if bit_to_trigger == "POSITION":
stored_token, stored_pos, stored_prev, stored_rkv = read_slot_fields(
canary_buf=cuda_buf, slot_idx=corrupt_slot
)
stamp_pair(
buf_pair,
slot_idx=corrupt_slot,
token=stored_token,
position=stored_pos + 99,
prev_hash=stored_prev,
real_kv_hash=stored_rkv,
)
else:
stored_token, stored_pos, stored_prev, stored_rkv = read_slot_fields(
canary_buf=cuda_buf, slot_idx=corrupt_slot
)
flipped_prev = stored_prev ^ 1
stamp_pair(
buf_pair,
slot_idx=corrupt_slot,
token=stored_token,
position=stored_pos,
prev_hash=flipped_prev,
real_kv_hash=stored_rkv,
)
ring_capacity = _RING_CAPACITY
cuda_log, ref_log = make_log_pair(capacity=ring_capacity, device=_DEVICE)
@@ -540,6 +673,9 @@ class TestViolationField:
plan_ref=plan_ref,
cuda_log=cuda_log,
ref_log=ref_log,
real_kv_sources_cuda=real_kv_sources_cuda,
real_kv_sources_ref=real_kv_sources_ref,
real_kv_hash_mode=real_kv_hash_mode,
assert_equal=False,
)
@@ -577,7 +713,451 @@ class TestViolationField:
) == 0, f"chain hash bit unexpectedly set: {bits:#b}"
class TestRealKvHash:
def test_real_kv_mode_off_yields_zero(self) -> None:
"""OFF mode → stored real_kv_hash field stays zero post-write; verify with OFF agrees byte-equal."""
buf_pair = _buf_pair()
sources = make_real_kv_sources(count=2, device=_DEVICE)
plan_pair = _plan_pair_single(slot_idx=1, position=0)
_stamp_head(buf_pair, slot_idx=1, token=1)
cuda_log, _ = run_verify_diff(
buf_pair=buf_pair,
plan_pair=plan_pair,
real_kv_sources_pair=(sources, sources),
)
assert _n_violations(cuda_log) == 0
@pytest.mark.parametrize(
"mode",
[
pytest.param(consts.RealKvHashMode.PARTIAL, id="partial"),
pytest.param(consts.RealKvHashMode.ALL, id="all"),
],
)
def test_real_kv_mode_byte_equal(self, mode: consts.RealKvHashMode) -> None:
"""PARTIAL / ALL modes both produce CUDA-vs-ref byte-equal state on a clean 3-step chain."""
buf_pair = _buf_pair()
sources_cuda = make_real_kv_sources(count=2, device=_DEVICE)
sources_ref = clone_real_kv_sources(sources_cuda)
# Write a chain through the ref so both buffers are byte-equal post-write.
_stamp_clean_kv_chain(
buf_pair=buf_pair,
sources_cuda=sources_cuda,
input_ids=torch.tensor([10, 20, 30], dtype=torch.int64, device=_DEVICE),
positions=torch.tensor([0, 1, 2], dtype=torch.int64, device=_DEVICE),
out_cache_loc=torch.tensor([1, 2, 3], dtype=torch.int64, device=_DEVICE),
real_kv_hash_mode=mode,
)
plan_pair = make_verify_plan_pair(
slot_indices=[1, 2, 3],
positions=[0, 1, 2],
prev_slot_indices=[-1, 1, 2],
device=_DEVICE,
)
cuda_log, _ = run_verify_diff(
buf_pair=buf_pair,
plan_pair=plan_pair,
real_kv_sources_pair=(sources_cuda, sources_ref),
real_kv_hash_mode=mode,
)
assert _n_violations(cuda_log) == 0
@pytest.mark.parametrize("count", [1, 2, 3, 4])
def test_real_kv_sources_fold_1_to_4(self, count: int) -> None:
"""Fold ``count`` sources sequentially → CUDA matches ref for every count in {1..4}."""
buf_pair = _buf_pair()
sources_cuda = make_real_kv_sources(count=count, device=_DEVICE)
sources_ref = clone_real_kv_sources(sources_cuda)
_stamp_clean_kv_chain(
buf_pair=buf_pair,
sources_cuda=sources_cuda,
input_ids=torch.tensor([1, 2], dtype=torch.int64, device=_DEVICE),
positions=torch.tensor([0, 1], dtype=torch.int64, device=_DEVICE),
out_cache_loc=torch.tensor([1, 2], dtype=torch.int64, device=_DEVICE),
real_kv_hash_mode=consts.RealKvHashMode.ALL,
)
plan_pair = make_verify_plan_pair(
slot_indices=[1, 2],
positions=[0, 1],
prev_slot_indices=[-1, 1],
device=_DEVICE,
)
cuda_log, _ = run_verify_diff(
buf_pair=buf_pair,
plan_pair=plan_pair,
real_kv_sources_pair=(sources_cuda, sources_ref),
real_kv_hash_mode=consts.RealKvHashMode.ALL,
)
assert _n_violations(cuda_log) == 0
@pytest.mark.parametrize(
"mode,fold_fn,expected_hash",
[
pytest.param(
consts.RealKvHashMode.PARTIAL,
_hand_fold_partial,
0x6041580849E6407D,
id="partial",
),
pytest.param(
consts.RealKvHashMode.ALL,
_hand_fold_all,
0x6041580849E6407D,
id="all",
),
],
)
def test_real_kv_hash_fold_mode_hardcoded(
self,
mode: consts.RealKvHashMode,
fold_fn: Callable[[bytes], int],
expected_hash: int,
) -> None:
# Step 1: build one RealKvSource with read_bytes=16 and a fixed byte pattern at slot 1.
_PATTERN = bytes(
[
0x01,
0x02,
0x04,
0x08,
0x10,
0x20,
0x40,
0x80,
0x81,
0x82,
0x84,
0x88,
0x90,
0xA0,
0xC0,
0xFF,
]
)
buf_pair = _buf_pair()
source_cuda = make_real_kv_source(
num_slots=16,
num_bytes_per_token=16,
page_size=1,
read_bytes=16,
device=_DEVICE,
)
source_cuda.tensor[1, :16] = torch.tensor(list(_PATTERN), dtype=torch.uint8)
source_ref = RealKvSource(
tensor=source_cuda.tensor.clone(),
page_size=source_cuda.page_size,
num_bytes_per_token=source_cuda.num_bytes_per_token,
read_bytes=source_cuda.read_bytes,
)
# Step 2: verify hand-computed fold matches the hex literal.
assert fold_fn(_PATTERN) == expected_hash
# Step 3: stamp slot 1 with a chain-head entry whose real_kv_hash equals the expected value.
_stamp_head(
buf_pair,
slot_idx=1,
token=7,
real_kv_hash=to_signed_int64(expected_hash),
)
# Step 4: 1-entry verify plan; no violation because stored matches recomputed.
plan_pair = _plan_pair_single(slot_idx=1, position=0)
cuda_log, _ = run_verify_diff(
buf_pair=buf_pair,
plan_pair=plan_pair,
real_kv_sources_pair=((source_cuda,), (source_ref,)),
real_kv_hash_mode=mode,
assert_equal=False,
)
assert _n_violations(cuda_log) == 0
# Step 5: mutate one byte in the source so the recomputed hash diverges from stored.
source_cuda.tensor[1, 0] ^= 0xFF
source_ref.tensor.copy_(source_cuda.tensor)
plan_pair2 = _plan_pair_single(slot_idx=1, position=0)
cuda_log2, _ = run_verify_diff(
buf_pair=buf_pair,
plan_pair=plan_pair2,
real_kv_sources_pair=((source_cuda,), (source_ref,)),
real_kv_hash_mode=mode,
)
assert_only_bits_set(
_fail_bits(cuda_log2), consts.FailReason.VERIFY_REAL_KV_HASH_MISMATCH
)
def test_real_kv_hash_all_mode_with_multiple_sources(self) -> None:
"""ALL mode with count=2 page=16 bytes=128 sources: chain still verifies clean."""
buf_pair = _buf_pair(num_slots=32)
sources_cuda = make_real_kv_sources(
count=2,
num_bytes_per_token=128,
page_size=16,
num_slots=32,
device=_DEVICE,
)
sources_ref = clone_real_kv_sources(sources_cuda)
slot_indices = [1, 2, 3]
tokens = [100, 200, 300]
positions = [0, 1, 2]
running = splitmix64(consts.CANARY_CHAIN_ANCHOR)
real_kv_hashes: list[int] = []
for slot_idx in slot_indices:
real_kv_hashes.append(
_compute_real_kv_hash_scalar(
real_kv_sources=sources_cuda,
real_kv_hash_mode=consts.RealKvHashMode.ALL,
slot_idx=slot_idx,
work_device=torch.device("cpu"),
)
)
for slot_idx, token, position, rkv in zip(
slot_indices, tokens, positions, real_kv_hashes
):
signed_prev = to_signed_int64(running)
stamp_pair(
buf_pair,
slot_idx=slot_idx,
token=token,
position=position,
prev_hash=signed_prev,
real_kv_hash=to_signed_int64(rkv),
)
running = splitmix64_mix3(running, token, position)
plan_pair = make_verify_plan_pair(
slot_indices=slot_indices,
positions=positions,
prev_slot_indices=[-1, 1, 2],
device=_DEVICE,
)
cuda_log, _ = run_verify_diff(
buf_pair=buf_pair,
plan_pair=plan_pair,
real_kv_sources_pair=(sources_cuda, sources_ref),
real_kv_hash_mode=consts.RealKvHashMode.ALL,
)
assert _n_violations(cuda_log) == 0
def test_real_kv_hash_partial_mode_detects_single_bit_flip(self) -> None:
"""PARTIAL mode + 1-bit flip in source tensor → REAL_KV_HASH bit set in violation row."""
buf_pair = _buf_pair()
sources_cuda = make_real_kv_sources(
count=1, num_bytes_per_token=16, device=_DEVICE
)
slot_idx = 3
row_bytes = (
sources_cuda[0]
.tensor[slot_idx, : sources_cuda[0].read_bytes]
.detach()
.cpu()
.tolist()
)
rkv_clean = _hand_fold_partial(bytes(row_bytes))
_stamp_head(
buf_pair,
slot_idx=slot_idx,
real_kv_hash=to_signed_int64(rkv_clean),
)
sources_cuda[0].tensor[slot_idx, 0] ^= 1
sources_ref = clone_real_kv_sources(sources_cuda)
plan_pair = _plan_pair_single(slot_idx=slot_idx, position=0)
cuda_log, _ = run_verify_diff(
buf_pair=buf_pair,
plan_pair=plan_pair,
real_kv_sources_pair=(sources_cuda, sources_ref),
real_kv_hash_mode=consts.RealKvHashMode.PARTIAL,
)
assert _n_violations(cuda_log) >= 1
bits = _fail_bits(cuda_log)
assert (
bits & consts.FailReason.VERIFY_REAL_KV_HASH_MISMATCH
), f"expected REAL_KV_HASH bit, got {bits:#b}"
def test_real_kv_off_does_not_deref_real_kv_sources(self) -> None:
buf_pair = _buf_pair(num_slots=8)
_stamp_head(buf_pair, slot_idx=1, token=1)
garbage_source = make_real_kv_source(
num_slots=8,
num_bytes_per_token=16,
page_size=1,
read_bytes=16,
device=_DEVICE,
fill=0xDE,
)
plan_pair = _plan_pair_single(slot_idx=1, position=0)
cuda_log, ref_log = run_verify_diff(
buf_pair=buf_pair,
plan_pair=plan_pair,
real_kv_sources_pair=((garbage_source,), (garbage_source,)),
assert_equal=False,
)
assert _n_violations(cuda_log) == 0
assert _n_violations(ref_log) == 0
class TestRealKvSource:
def test_real_kv_source_rejects_zero_read_bytes(self) -> None:
"""RealKvSource has no \"skip me\" sentinel — read_bytes=0 must raise rather than silently pass."""
with pytest.raises(ValueError, match="read_bytes"):
RealKvSource(
tensor=torch.zeros((1, 16), dtype=torch.uint8, device=_DEVICE),
page_size=1,
num_bytes_per_token=16,
read_bytes=0,
)
def test_real_kv_source_padding_below_4(self) -> None:
"""Host wrapper pads to 4 slots when fewer sources are supplied; dummy slots are never dereferenced."""
buf_pair = _buf_pair()
sources = make_real_kv_sources(count=2, device=_DEVICE)
plan_pair = _plan_pair_single(slot_idx=1, position=0)
_stamp_head(buf_pair, slot_idx=1, token=1)
run_verify_diff(
buf_pair=buf_pair,
plan_pair=plan_pair,
real_kv_sources_pair=(sources, sources),
)
def test_real_kv_source_above_4_raises(self) -> None:
"""``len(real_kv_sources) > 4`` → host wrapper raises ValueError before launching."""
canary_buf = make_canary_buf(device=_DEVICE)
plan = make_verify_plan(
slot_indices=[1], positions=[0], prev_slot_indices=[-1], device=_DEVICE
)
log = FakeViolationLog.allocate(device=_DEVICE)
sources = make_real_kv_sources(count=4, device=_DEVICE)
extra = make_real_kv_source(device=_DEVICE)
too_many = sources + (extra,)
with pytest.raises(ValueError, match="at most 4 RealKvSource"):
launch_canary_verify_kernel(
context=VerifyOrWriteContext(
canary_buf=canary_buf,
kernel_kind=CanaryLaunchTag.HEAD_K_FULL,
violation_ring=log.ring,
violation_write_index=log.write_index,
slot_run_counter=log.slot_run_counter,
kernel_run_counter=log.kernel_run_counter,
enable_chain_position_assert=log.enable_chain_position_assert,
real_kv_sources=too_many,
real_kv_hash_mode=consts.RealKvHashMode.NONE,
),
plan=plan,
check_verify_expected_token=True,
)
def test_real_kv_source_holey_dim1(self) -> None:
"""``tensor.shape[1] > page_size * num_bytes_per_token`` → trailing bytes are skipped."""
buf_pair = _buf_pair()
holey_source = make_real_kv_source(
num_slots=16,
num_bytes_per_token=16,
page_size=1,
read_bytes=16,
pad_dim1=16, # 16 trailing pad bytes per row; must be skipped.
device=_DEVICE,
)
trailing_start = holey_source.page_size * holey_source.num_bytes_per_token
# Fill those skipped trailing bytes with garbage; CUDA must not read them.
holey_source.tensor[:, trailing_start:].fill_(0xAA)
sources = (holey_source,)
sources_ref = clone_real_kv_sources(sources)
_stamp_clean_kv_chain(
buf_pair=buf_pair,
sources_cuda=sources,
input_ids=torch.tensor([1, 2], dtype=torch.int64, device=_DEVICE),
positions=torch.tensor([0, 1], dtype=torch.int64, device=_DEVICE),
out_cache_loc=torch.tensor([1, 2], dtype=torch.int64, device=_DEVICE),
real_kv_hash_mode=consts.RealKvHashMode.ALL,
)
plan_pair = make_verify_plan_pair(
slot_indices=[1, 2],
positions=[0, 1],
prev_slot_indices=[-1, 1],
device=_DEVICE,
)
cuda_log, _ = run_verify_diff(
buf_pair=buf_pair,
plan_pair=plan_pair,
real_kv_sources_pair=(sources, sources_ref),
real_kv_hash_mode=consts.RealKvHashMode.ALL,
)
assert _n_violations(cuda_log) == 0
class TestLayoutAndScheduling:
def test_page_size_gt_1_access_pattern(self) -> None:
"""``page_size > 1`` → byte access follows ``(row=slot//page, col=(slot%page)*bpt:)``."""
buf_pair = _buf_pair(num_slots=8)
src = make_real_kv_source(
num_slots=8,
num_bytes_per_token=16,
page_size=4, # 2 rows × 4 slots/page × 16 bytes/slot.
read_bytes=16,
device=_DEVICE,
)
# Each slot's 16 bytes get a slot-specific signature so kernel mis-indexing would shift the hash.
flat = src.tensor.view(-1)
for slot_idx in range(8):
row = slot_idx // src.page_size
col = (slot_idx % src.page_size) * src.num_bytes_per_token
for k in range(src.num_bytes_per_token):
flat_index = row * (src.page_size * src.num_bytes_per_token) + col + k
flat[flat_index] = (slot_idx * 13 + k) & 0xFF
sources = (src,)
sources_ref = clone_real_kv_sources(sources)
_stamp_clean_kv_chain(
buf_pair=buf_pair,
sources_cuda=sources,
input_ids=torch.tensor([1, 2], dtype=torch.int64, device=_DEVICE),
positions=torch.tensor([0, 1], dtype=torch.int64, device=_DEVICE),
out_cache_loc=torch.tensor([1, 5], dtype=torch.int64, device=_DEVICE),
real_kv_hash_mode=consts.RealKvHashMode.ALL,
)
plan_pair = make_verify_plan_pair(
slot_indices=[1, 5],
positions=[0, 1],
prev_slot_indices=[-1, 1],
device=_DEVICE,
)
cuda_log, _ = run_verify_diff(
buf_pair=buf_pair,
plan_pair=plan_pair,
real_kv_sources_pair=(sources, sources_ref),
real_kv_hash_mode=consts.RealKvHashMode.ALL,
)
assert _n_violations(cuda_log) == 0
def test_swa_translated_slot_indices(self) -> None:
"""SWA-translated slots already passed in plan; verify kernel does no further translation."""
# SWA verify plans carry pre-translated slot indices — the verify kernel never sees the FULL slot
@@ -615,6 +1195,7 @@ class TestLayoutAndScheduling:
token=999,
position=123,
prev_hash=to_signed_int64(0xDEADBEEF),
real_kv_hash=0,
)
plan = make_verify_plan(
slot_indices=[0],
@@ -633,6 +1214,8 @@ class TestLayoutAndScheduling:
slot_run_counter=log.slot_run_counter,
kernel_run_counter=log.kernel_run_counter,
enable_chain_position_assert=log.enable_chain_position_assert,
real_kv_sources=(),
real_kv_hash_mode=consts.RealKvHashMode.NONE,
),
plan=plan,
check_verify_expected_token=True,
@@ -659,6 +1242,7 @@ class TestLayoutAndScheduling:
token=42,
position=1,
prev_hash=to_signed_int64(0x1234),
real_kv_hash=0,
)
plan = make_verify_plan(
slot_indices=[slot_idx],
@@ -687,6 +1271,8 @@ class TestLayoutAndScheduling:
slot_run_counter=log.slot_run_counter,
kernel_run_counter=log.kernel_run_counter,
enable_chain_position_assert=log.enable_chain_position_assert,
real_kv_sources=(),
real_kv_hash_mode=consts.RealKvHashMode.NONE,
),
plan=plan,
check_verify_expected_token=True,
@@ -702,6 +1288,65 @@ class TestLayoutAndScheduling:
== int(kernel_run_before[0].item()) + 1
)
def test_paged_layout_page_size_16(self) -> None:
"""page_size=16: slot→page mapping doesn't change verify chain semantics on a clean chain."""
buf_pair = _buf_pair(num_slots=64)
sources_cuda = make_real_kv_sources(
count=1,
num_bytes_per_token=16,
page_size=16,
num_slots=64,
device=_DEVICE,
)
sources_ref = clone_real_kv_sources(sources_cuda)
# Step: cross a page boundary by writing slots [15, 16] which straddle pages 0 and 1.
slot_indices = [15, 16]
tokens = [77, 88]
positions = [0, 1]
running = splitmix64(consts.CANARY_CHAIN_ANCHOR)
# Use the reference fold (8-byte little-endian word pack + splitmix64), not a
# byte-by-byte loop, so the stamped real_kv_hash matches what the kernel /
# verify reference will recompute. A byte-by-byte fold was the previous bug
# here and triggered REAL_KV_HASH violations on otherwise clean chains.
rkv_values = [
_compute_real_kv_hash_scalar(
slot_idx=slot_idx,
real_kv_sources=sources_cuda,
real_kv_hash_mode=consts.RealKvHashMode.ALL,
work_device=_DEVICE,
)
for slot_idx in slot_indices
]
for slot_idx, token, position, rkv in zip(
slot_indices, tokens, positions, rkv_values
):
signed_prev = to_signed_int64(running)
stamp_pair(
buf_pair,
slot_idx=slot_idx,
token=token,
position=position,
prev_hash=signed_prev,
real_kv_hash=to_signed_int64(rkv),
)
running = splitmix64_mix3(running, token, position)
plan_pair = make_verify_plan_pair(
slot_indices=slot_indices,
positions=positions,
prev_slot_indices=[-1, 15],
device=_DEVICE,
)
cuda_log, _ = run_verify_diff(
buf_pair=buf_pair,
plan_pair=plan_pair,
real_kv_sources_pair=(sources_cuda, sources_ref),
real_kv_hash_mode=consts.RealKvHashMode.ALL,
)
assert _n_violations(cuda_log) == 0
def test_empty_plan_keeps_slot_counter_unchanged(self) -> None:
buf_pair = _buf_pair(num_slots=8)
_stamp_head(buf_pair, slot_idx=1, token=7)
@@ -1166,6 +1811,22 @@ class TestBoundarySweep:
)
)
@pytest.mark.parametrize(
"stored_rkv_val",
[0, 1, 0xFFFFFFFFFFFFFFFF, 0x8000000000000000],
)
def test_real_kv_hash_boundary_byte_equal_sweep(self, stored_rkv_val: int) -> None:
"""Sweep real_kv_hash boundary values; assert CUDA vs ref state byte-equal."""
sources_cuda = make_real_kv_sources(count=1, device=_DEVICE)
_run_verify_single_slot_byte_equal(
_VerifySingleSlotInput(
stored_prev_hash_signed=chain_anchor_signed(),
stored_real_kv_hash_signed=to_signed_int64(stored_rkv_val),
real_kv_sources=sources_cuda,
real_kv_hash_mode=consts.RealKvHashMode.PARTIAL,
)
)
class TestVerifyExpectedInputIds:
"""Cover the new verify-time token-id check via VerifyPlan.verify_expected_tokens."""
python/sglang/jit_kernel/tests/kv_canary/test_write_fuzz.py
+41 -3
View File
@@ -6,7 +6,11 @@ from dataclasses import dataclass
import pytest
import torch
from sglang.jit_kernel.kv_canary.verify import CanaryLaunchTag
from sglang.jit_kernel.kv_canary import consts
from sglang.jit_kernel.kv_canary.verify import (
CanaryLaunchTag,
RealKvSource,
)
from sglang.jit_kernel.kv_canary.write import WritePlan
from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
FakeViolationLog,
@@ -16,6 +20,10 @@ from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
stamp_pair,
)
from sglang.jit_kernel.tests.kv_canary._differential import _run_both_write
from sglang.jit_kernel.tests.kv_canary._fixtures import (
clone_real_kv_sources,
make_real_kv_sources,
)
from sglang.jit_kernel.tests.kv_canary._fuzz_driver import (
FUZZ_SEEDS_PR,
run_fuzz_combo,
@@ -44,11 +52,24 @@ class WriteFuzzInputs:
enable_write_verify_inputs: bool
expected_input_tokens: torch.Tensor
expected_input_positions: torch.Tensor
real_kv_sources_cuda: tuple[RealKvSource, ...]
real_kv_sources_ref: tuple[RealKvSource, ...]
real_kv_hash_mode: consts.RealKvHashMode
ring_capacity: int
def _draw_random_write_inputs(rng: random.Random) -> WriteFuzzInputs:
enable_write_verify_inputs = rng.choice([False, True])
hash_mode = rng.choice(
[
consts.RealKvHashMode.NONE,
consts.RealKvHashMode.PARTIAL,
consts.RealKvHashMode.ALL,
]
)
src_count = rng.choice([1, 2, 4])
page_size = rng.choice([1, 16])
bytes_per = rng.choice([16, 64, 128])
kernel_kind = rng.choice(list(CanaryLaunchTag))
ring_capacity = rng.choice([16, 64, 256])
@@ -57,6 +78,16 @@ def _draw_random_write_inputs(rng: random.Random) -> WriteFuzzInputs:
total_tokens = sum(per_req_tokens)
num_slots = max(total_tokens + 8, 16)
sources_cuda = make_real_kv_sources(
count=src_count,
num_bytes_per_token=bytes_per,
page_size=page_size,
num_slots=num_slots,
device=_DEVICE,
rng=rng,
)
sources_ref = clone_real_kv_sources(sources_cuda)
cuda_buf = make_canary_buf(
num_slots=num_slots, slot_stride_bytes=32, device=_DEVICE
)
@@ -140,6 +171,9 @@ def _draw_random_write_inputs(rng: random.Random) -> WriteFuzzInputs:
enable_write_verify_inputs=enable_write_verify_inputs,
expected_input_tokens=expected_input_tokens,
expected_input_positions=expected_input_positions,
real_kv_sources_cuda=sources_cuda,
real_kv_sources_ref=sources_ref,
real_kv_hash_mode=hash_mode,
ring_capacity=ring_capacity,
)
@@ -163,6 +197,9 @@ def _run_one(inputs: WriteFuzzInputs) -> None:
expected_input_positions=inputs.expected_input_positions,
cuda_log=cuda_log,
ref_log=ref_log,
real_kv_sources_cuda=inputs.real_kv_sources_cuda,
real_kv_sources_ref=inputs.real_kv_sources_ref,
real_kv_hash_mode=inputs.real_kv_hash_mode,
kernel_kind=inputs.kernel_kind,
assert_equal=False,
)
@@ -196,13 +233,14 @@ def _summarize(inputs: WriteFuzzInputs) -> str:
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}"
f"pseudo={inputs.enable_write_verify_inputs} hash_mode={inputs.real_kv_hash_mode.name} "
f"sources={len(inputs.real_kv_sources_cuda)}"
)
@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."""
"""Multi-dim write fuzzer: random pseudo/hash/kernel/page/source × N iters, byte-equal."""
run_fuzz_combo(
seed,
draw_fn=_draw_random_write_inputs,
python/sglang/jit_kernel/tests/kv_canary/test_write_hand.py
+333 -2
View File
@@ -1,6 +1,7 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import Callable
from unittest.mock import patch
import pytest
@@ -12,6 +13,7 @@ from sglang.jit_kernel.kv_canary.consts import splitmix64, splitmix64_mix3
from sglang.jit_kernel.kv_canary.verify import (
CANARY_SLOT_BYTES,
CanaryLaunchTag,
RealKvSource,
VerifyOrWriteContext,
launch_canary_verify_kernel,
)
@@ -26,6 +28,8 @@ from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
make_canary_buf,
make_canary_buf_pair,
make_log_pair,
make_real_kv_source,
make_real_kv_sources,
make_verify_plan,
make_write_plan,
make_write_plan_pair,
@@ -38,8 +42,13 @@ from sglang.jit_kernel.tests.kv_canary._differential import (
run_write_diff,
)
from sglang.jit_kernel.tests.kv_canary._fixtures import (
clone_real_kv_sources,
dummy_pseudo_tensors,
)
from sglang.jit_kernel.tests.kv_canary._hand_oracle import (
_hand_fold_all,
_hand_fold_partial,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
@@ -73,6 +82,10 @@ def _run_write(
enable_write_verify_inputs: bool = False,
expected_input_tokens: torch.Tensor | None = None,
expected_input_positions: torch.Tensor | None = None,
real_kv_sources_pair: (
tuple[tuple[RealKvSource, ...], tuple[RealKvSource, ...]] | None
) = None,
real_kv_hash_mode: consts.RealKvHashMode = consts.RealKvHashMode.NONE,
assert_equal: bool = True,
) -> tuple[FakeViolationLog, FakeViolationLog]:
"""Shared scaffold: build write plan + pseudo tensors and call ``run_write_diff``.
@@ -117,6 +130,9 @@ def _run_write(
if expected_input_positions is None:
expected_input_positions = pseudo_positions
extra_kwargs: dict = {}
if real_kv_sources_pair is not None:
extra_kwargs["real_kv_sources_pair"] = real_kv_sources_pair
return run_write_diff(
buf_pair=buf_pair,
plan_pair=plan_pair,
@@ -126,7 +142,9 @@ def _run_write(
enable_write_verify_inputs=enable_write_verify_inputs,
expected_input_tokens=expected_input_tokens,
expected_input_positions=expected_input_positions,
real_kv_hash_mode=real_kv_hash_mode,
assert_equal=assert_equal,
**extra_kwargs,
)
@@ -135,6 +153,8 @@ class _WriteSingleSlotInput:
token: int = 42
position: int = 0
enable_write_verify_inputs: bool = False
real_kv_sources: tuple[RealKvSource, ...] = ()
real_kv_hash_mode: consts.RealKvHashMode = consts.RealKvHashMode.NONE
class _RecordingWriteModule:
@@ -146,12 +166,16 @@ class _RecordingWriteModule:
def _run_write_single_slot_byte_equal(case: _WriteSingleSlotInput) -> None:
sources_cuda = case.real_kv_sources
sources_ref = clone_real_kv_sources(sources_cuda)
_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,
real_kv_sources_pair=(sources_cuda, sources_ref),
real_kv_hash_mode=case.real_kv_hash_mode,
)
@@ -242,6 +266,8 @@ class TestSeedSlot:
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,
real_kv_sources=(),
real_kv_hash_mode=consts.RealKvHashMode.NONE,
),
plan=verify_plan,
check_verify_expected_token=True,
@@ -270,11 +296,12 @@ class TestSeedSlot:
tokens = [101, 202, 303, 404, 505]
positions = [11, 12, 13, 14, 15]
real_kv = [0, 0, 0, 0, 0]
out_cache_loc = [0, 1, 2, 3, 4]
expected_prev_hashes: list[int] = []
running = predecessor_advance
for t, p in zip(tokens, positions):
for t, p, r in zip(tokens, positions, real_kv):
expected_prev_hashes.append(running)
running = splitmix64_mix3(running, t, p)
@@ -305,6 +332,7 @@ class TestSeedSlot:
seed_slot = 3
seed_token = 7
seed_position = 1
seed_real_kv = 0
expected_seed_prev_hash = splitmix64(consts.CANARY_CHAIN_ANCHOR)
stamp_pair(
self.buf_pair,
@@ -312,6 +340,7 @@ class TestSeedSlot:
token=seed_token,
position=seed_position,
prev_hash=to_signed_int64(expected_seed_prev_hash),
real_kv_hash=to_signed_int64(seed_real_kv),
)
new_slot = 4
@@ -355,11 +384,12 @@ class TestChain:
tokens = [101, 202, 303, 404, 505]
positions = [0, 1, 2, 3, 4]
out_cache_loc = [0, 1, 2, 3, 4]
real_kv_hashes = [0, 0, 0, 0, 0]
# 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):
for token, position, real_kv_hash in zip(tokens, positions, real_kv_hashes):
expected_prev_hashes_u64.append(running)
running = splitmix64_mix3(running, token, position)
expected_prev_hashes_signed = [
@@ -385,6 +415,41 @@ class TestChain:
assert stored_prev_hash == expected_prev_signed
assert stored_real_kv_hash == 0
def test_chain_advances_with_real_kv_hash_all(self) -> None:
"""ALL mode + 2 sources + 5-step chain: stored prev_hash recoverable from seed."""
cuda_buf = self.buf_pair[0]
sources_cuda = make_real_kv_sources(
count=2,
num_bytes_per_token=16,
page_size=1,
num_slots=16,
device=_DEVICE,
)
sources_ref = clone_real_kv_sources(sources_cuda)
slot_indices = [1, 2, 3, 4, 5]
tokens = [11, 22, 33, 44, 55]
positions = [0, 1, 2, 3, 4]
_run_write(
buf_pair=self.buf_pair,
input_ids=tokens,
positions=positions,
out_cache_loc=slot_indices,
real_kv_sources_pair=(sources_cuda, sources_ref),
real_kv_hash_mode=consts.RealKvHashMode.ALL,
)
running = splitmix64(consts.CANARY_CHAIN_ANCHOR)
for slot_idx, token, position in zip(slot_indices, tokens, positions):
stored_prev_signed, stored_real_kv_hash = read_slot_fields(
canary_buf=cuda_buf, slot_idx=slot_idx
)[2:]
assert stored_prev_signed == to_signed_int64(
running
), f"slot {slot_idx}: stored prev_hash != recomputed chain step"
running = splitmix64_mix3(running, token, position)
class TestMockMode:
def setup_method(self) -> None:
@@ -496,6 +561,8 @@ class TestMockMode:
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,
real_kv_sources=(),
real_kv_hash_mode=consts.RealKvHashMode.NONE,
),
plan=verify_plan,
check_verify_expected_token=True,
@@ -686,6 +753,265 @@ class TestSlotHandling:
), f"slot {slot} from earlier bs=8 run was overwritten by bs=3 run"
class TestRealKvHash:
def setup_method(self) -> None:
self.buf_pair = _make_default_buf_pair()
def test_real_kv_mode_off_writes_zero(self) -> None:
"""``consts.RealKvHashMode.NONE`` → ``real_kv_hash`` field is written as 0 regardless of source presence."""
sources = make_real_kv_sources(count=2, device=_DEVICE)
_run_write(
buf_pair=self.buf_pair,
input_ids=[1, 2],
positions=[0, 1],
out_cache_loc=[0, 1],
real_kv_sources_pair=(sources, sources),
)
_, _, _, real_kv_0 = read_slot_fields(canary_buf=self.buf_pair[0], slot_idx=0)
_, _, _, real_kv_1 = read_slot_fields(canary_buf=self.buf_pair[0], slot_idx=1)
assert real_kv_0 == 0
assert real_kv_1 == 0
@pytest.mark.parametrize(
"mode",
[
pytest.param(consts.RealKvHashMode.PARTIAL, id="partial"),
pytest.param(consts.RealKvHashMode.ALL, id="all"),
],
)
def test_real_kv_mode_byte_equal(self, mode: consts.RealKvHashMode) -> None:
"""PARTIAL / ALL modes both produce CUDA-vs-ref byte-equal write state on a 3-entry chain."""
sources_cuda = make_real_kv_sources(count=2, device=_DEVICE)
sources_ref = clone_real_kv_sources(sources_cuda)
_run_write(
buf_pair=self.buf_pair,
input_ids=[10, 20, 30],
positions=[0, 1, 2],
out_cache_loc=[0, 1, 2],
real_kv_sources_pair=(sources_cuda, sources_ref),
real_kv_hash_mode=mode,
)
@pytest.mark.parametrize("count", [1, 2, 3, 4])
def test_real_kv_sources_fold_1_to_4(self, count: int) -> None:
"""Folding ``count`` sources sequentially → CUDA matches ref for every count in {1..4}."""
sources_cuda = make_real_kv_sources(count=count, device=_DEVICE)
sources_ref = clone_real_kv_sources(sources_cuda)
_run_write(
buf_pair=self.buf_pair,
input_ids=[1, 2],
positions=[0, 1],
out_cache_loc=[0, 1],
real_kv_sources_pair=(sources_cuda, sources_ref),
real_kv_hash_mode=consts.RealKvHashMode.ALL,
)
def test_real_kv_source_above_4_raises(self) -> None:
"""``len(real_kv_sources) > 4`` → host wrapper raises ValueError before launching."""
cuda_buf = make_canary_buf(device=_DEVICE)
plan = make_write_plan(
write_offsets=[0, 1],
seed_slot_indices=[-1],
num_valid_reqs=1,
device=_DEVICE,
)
input_ids = _int32_tensor([1])
positions = _int32_tensor([0])
out_cache_loc = _int32_tensor([0])
log = FakeViolationLog.allocate(device=_DEVICE)
sources = make_real_kv_sources(count=4, device=_DEVICE)
extra = make_real_kv_source(device=_DEVICE)
too_many = sources + (extra,)
with pytest.raises(ValueError, match="at most 4 RealKvSource"):
launch_canary_write_kernel(
context=VerifyOrWriteContext(
canary_buf=cuda_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,
real_kv_sources=too_many,
real_kv_hash_mode=consts.RealKvHashMode.NONE,
),
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,
)
@pytest.mark.parametrize(
"mode,fold_fn,expected_hash",
[
pytest.param(
consts.RealKvHashMode.PARTIAL,
_hand_fold_partial,
0x6041580849E6407D,
id="partial",
),
pytest.param(
consts.RealKvHashMode.ALL,
_hand_fold_all,
0x6041580849E6407D,
id="all",
),
],
)
def test_real_kv_hash_fold_mode_writes_expected_hash_hardcoded(
self,
mode: consts.RealKvHashMode,
fold_fn: Callable[[bytes], int],
expected_hash: int,
) -> None:
# Step 1: build one RealKvSource with read_bytes=16 and a fixed byte pattern at slot 0.
_PATTERN = bytes(
[
0x01,
0x02,
0x04,
0x08,
0x10,
0x20,
0x40,
0x80,
0x81,
0x82,
0x84,
0x88,
0x90,
0xA0,
0xC0,
0xFF,
]
)
# Step 2: verify hand-computed fold matches the hex literal.
assert fold_fn(_PATTERN) == expected_hash
source_cuda = make_real_kv_source(
num_slots=16,
num_bytes_per_token=16,
page_size=1,
read_bytes=16,
device=_DEVICE,
)
source_cuda.tensor[0, :16] = torch.tensor(list(_PATTERN), dtype=torch.uint8)
source_ref = RealKvSource(
tensor=source_cuda.tensor.clone(),
page_size=source_cuda.page_size,
num_bytes_per_token=source_cuda.num_bytes_per_token,
read_bytes=source_cuda.read_bytes,
)
# Step 3: run write kernel on slot 0 with the given mode.
_run_write(
buf_pair=self.buf_pair,
input_ids=[7],
positions=[0],
out_cache_loc=[0],
real_kv_sources_pair=((source_cuda,), (source_ref,)),
real_kv_hash_mode=mode,
)
# Step 4: assert stored real_kv_hash equals the hand-computed hex literal.
_, _, _, stored_real_kv_hash = read_slot_fields(
canary_buf=self.buf_pair[0], slot_idx=0
)
assert stored_real_kv_hash == to_signed_int64(
expected_hash
), f"stored_real_kv_hash={stored_real_kv_hash:#x} expected={to_signed_int64(expected_hash):#x}"
def test_paged_real_kv_hash_consistent_across_slots(self) -> None:
"""page=16: writing two slots inside same page yields independent real_kv_hash per slot."""
sources_cuda = make_real_kv_sources(
count=1,
num_bytes_per_token=16,
page_size=16,
num_slots=16,
device=_DEVICE,
)
pattern_slot3 = bytes(range(1, 17))
pattern_slot7 = bytes(range(101, 117))
sources_cuda[0].tensor[0, 3 * 16 : 4 * 16] = torch.tensor(
list(pattern_slot3), dtype=torch.uint8, device=_DEVICE
)
sources_cuda[0].tensor[0, 7 * 16 : 8 * 16] = torch.tensor(
list(pattern_slot7), dtype=torch.uint8, device=_DEVICE
)
sources_ref = clone_real_kv_sources(sources_cuda)
_run_write(
buf_pair=self.buf_pair,
input_ids=[42, 84],
positions=[0, 1],
out_cache_loc=[3, 7],
real_kv_sources_pair=(sources_cuda, sources_ref),
real_kv_hash_mode=consts.RealKvHashMode.ALL,
)
slot3 = read_slot_fields(canary_buf=self.buf_pair[0], slot_idx=3)
slot7 = read_slot_fields(canary_buf=self.buf_pair[0], slot_idx=7)
assert slot3[3] == to_signed_int64(_hand_fold_all(pattern_slot3))
assert slot7[3] == to_signed_int64(_hand_fold_all(pattern_slot7))
assert slot3[3] != slot7[3]
def test_multi_source_real_kv_fold_order_matters(self) -> None:
"""Two sources folded in reverse order yields a different real_kv_hash (fold is ordered)."""
sources_a = make_real_kv_sources(
count=2, num_bytes_per_token=16, num_slots=8, device=_DEVICE
)
sources_b = tuple(reversed(sources_a))
def _run_with(srcs: tuple[RealKvSource, ...]) -> tuple[int, int, int, int]:
buf = make_canary_buf(num_slots=16, slot_stride_bytes=32, device=_DEVICE)
plan = make_write_plan(
write_offsets=[0, 1],
seed_slot_indices=[-1],
num_valid_reqs=1,
device=_DEVICE,
)
log = FakeViolationLog.allocate(device=_DEVICE)
launch_canary_write_kernel(
context=VerifyOrWriteContext(
canary_buf=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,
real_kv_sources=srcs,
real_kv_hash_mode=consts.RealKvHashMode.ALL,
),
plan=plan,
input_ids=_int32_tensor([1]),
positions=_int32_tensor([0]),
out_cache_loc=_int32_tensor([2]),
enable_write_input_assert=False,
expected_input_tokens=None,
expected_input_positions=None,
)
torch.cuda.synchronize()
return read_slot_fields(canary_buf=buf, slot_idx=2)
fields_a = _run_with(sources_a)
fields_b = _run_with(sources_b)
assert fields_a[3] != 0
assert fields_b[3] != 0
assert (
fields_a[3] != fields_b[3]
), "reversing source order must change real_kv_hash (fold is ordered)"
class TestRunCounter:
def setup_method(self) -> None:
self.buf_pair = _make_default_buf_pair()
@@ -718,6 +1044,9 @@ class TestRunCounter:
expected_input_positions=pseudo_positions,
cuda_log=cuda_log,
ref_log=ref_log,
real_kv_sources_cuda=(),
real_kv_sources_ref=(),
real_kv_hash_mode=consts.RealKvHashMode.NONE,
assert_equal=False,
)
@@ -806,6 +1135,8 @@ class TestMisc:
slot_run_counter=log.slot_run_counter,
kernel_run_counter=log.kernel_run_counter,
enable_chain_position_assert=log.enable_chain_position_assert,
real_kv_sources=(),
real_kv_hash_mode=consts.RealKvHashMode.NONE,
)
module = _RecordingWriteModule()