Add the KV-canary write JIT kernel and reference implementation (#26806)
This commit is contained in:
fzyzcjy
@@ -0,0 +1,330 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.kv_canary import consts
|
||||
from sglang.jit_kernel.kv_canary.consts import splitmix64, splitmix64_mix3
|
||||
from sglang.jit_kernel.kv_canary.verify import VerifyPlan
|
||||
from sglang.jit_kernel.kv_canary.write import WritePlan
|
||||
from sglang.jit_kernel.tests.kv_canary._constants import (
|
||||
_I64_SIGN_BIT,
|
||||
_U64_MASK,
|
||||
DEFAULT_NUM_SLOTS,
|
||||
DEFAULT_RING_CAPACITY,
|
||||
DEFAULT_SLOT_STRIDE_BYTES,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"FakeViolationLog",
|
||||
"assert_canary_buf_equal",
|
||||
"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_verify_plan_pair",
|
||||
"make_write_plan",
|
||||
"make_write_plan_pair",
|
||||
"read_slot_fields",
|
||||
"stamp_clean_chain",
|
||||
"stamp_pair",
|
||||
"to_signed_int64",
|
||||
"write_slot_fields",
|
||||
]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True, kw_only=True)
|
||||
class FakeViolationLog:
|
||||
ring: torch.Tensor
|
||||
write_index: torch.Tensor
|
||||
slot_run_counter: torch.Tensor
|
||||
kernel_run_counter: torch.Tensor
|
||||
enable_chain_position_assert: torch.Tensor
|
||||
|
||||
@classmethod
|
||||
def allocate(
|
||||
cls, *, capacity: int = DEFAULT_RING_CAPACITY, device: torch.device
|
||||
) -> "FakeViolationLog":
|
||||
return cls(
|
||||
ring=torch.zeros(
|
||||
capacity, consts.VIOLATION_FIELDS, dtype=torch.int64, device=device
|
||||
),
|
||||
write_index=torch.zeros(1, dtype=torch.int32, device=device),
|
||||
slot_run_counter=torch.zeros(1, dtype=torch.int64, device=device),
|
||||
kernel_run_counter=torch.zeros(1, dtype=torch.int64, device=device),
|
||||
enable_chain_position_assert=torch.ones(
|
||||
1, dtype=torch.int32, device=device
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def make_canary_buf(
|
||||
*,
|
||||
num_slots: int = DEFAULT_NUM_SLOTS,
|
||||
slot_stride_bytes: int = DEFAULT_SLOT_STRIDE_BYTES,
|
||||
device: torch.device,
|
||||
) -> torch.Tensor:
|
||||
return torch.zeros(num_slots, slot_stride_bytes, dtype=torch.uint8, device=device)
|
||||
|
||||
|
||||
def make_canary_buf_pair(
|
||||
*,
|
||||
num_slots: int = DEFAULT_NUM_SLOTS,
|
||||
slot_stride_bytes: int = DEFAULT_SLOT_STRIDE_BYTES,
|
||||
device: torch.device,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
cuda_buf = make_canary_buf(
|
||||
num_slots=num_slots, slot_stride_bytes=slot_stride_bytes, device=device
|
||||
)
|
||||
return cuda_buf, cuda_buf.clone()
|
||||
|
||||
|
||||
def make_log_pair(
|
||||
*,
|
||||
capacity: int = DEFAULT_RING_CAPACITY,
|
||||
device: torch.device,
|
||||
) -> tuple[FakeViolationLog, FakeViolationLog]:
|
||||
return (
|
||||
FakeViolationLog.allocate(capacity=capacity, device=device),
|
||||
FakeViolationLog.allocate(capacity=capacity, device=device),
|
||||
)
|
||||
|
||||
|
||||
def make_verify_plan(
|
||||
*,
|
||||
slot_indices: list[int],
|
||||
positions: list[int],
|
||||
prev_slot_indices: list[int],
|
||||
expected_input_ids: Optional[list[int]] = None,
|
||||
capacity: Optional[int] = None,
|
||||
device: torch.device,
|
||||
) -> VerifyPlan:
|
||||
"""Build a VerifyPlan whose active prefix matches the three input lists.
|
||||
|
||||
Active prefix mirrors the input lists. Tail entries are left at the
|
||||
allocate-time defaults; ``verify_num_valid = len(slot_indices)``.
|
||||
|
||||
``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.
|
||||
"""
|
||||
n_active = len(slot_indices)
|
||||
if not (len(positions) == n_active and len(prev_slot_indices) == n_active):
|
||||
raise ValueError(
|
||||
"make_verify_plan: slot_indices, positions, and prev_slot_indices must all have the same length"
|
||||
)
|
||||
if expected_input_ids is None:
|
||||
expected_input_ids = [-1] * n_active
|
||||
if len(expected_input_ids) != n_active:
|
||||
raise ValueError(
|
||||
"make_verify_plan: expected_input_ids must match len(slot_indices)"
|
||||
)
|
||||
cap = capacity if capacity is not None else max(n_active, 1)
|
||||
plan = VerifyPlan.allocate(verify_capacity=cap, device=device)
|
||||
if n_active > 0:
|
||||
plan.verify_slot_indices[:n_active] = torch.tensor(
|
||||
slot_indices, dtype=torch.int64, device=device
|
||||
)
|
||||
plan.verify_expected_tokens[:n_active] = torch.tensor(
|
||||
expected_input_ids, dtype=torch.int64, device=device
|
||||
)
|
||||
plan.verify_expected_positions[:n_active] = torch.tensor(
|
||||
positions, dtype=torch.int64, device=device
|
||||
)
|
||||
plan.verify_prev_slot_indices[:n_active] = torch.tensor(
|
||||
prev_slot_indices, dtype=torch.int64, device=device
|
||||
)
|
||||
plan.verify_num_valid[0] = n_active
|
||||
return plan
|
||||
|
||||
|
||||
def make_verify_plan_pair(
|
||||
*,
|
||||
slot_indices: list[int],
|
||||
positions: list[int],
|
||||
prev_slot_indices: list[int],
|
||||
expected_input_ids: Optional[list[int]] = None,
|
||||
capacity: Optional[int] = None,
|
||||
device: torch.device,
|
||||
) -> tuple[VerifyPlan, VerifyPlan]:
|
||||
return (
|
||||
make_verify_plan(
|
||||
slot_indices=slot_indices,
|
||||
positions=positions,
|
||||
prev_slot_indices=prev_slot_indices,
|
||||
expected_input_ids=expected_input_ids,
|
||||
capacity=capacity,
|
||||
device=device,
|
||||
),
|
||||
make_verify_plan(
|
||||
slot_indices=slot_indices,
|
||||
positions=positions,
|
||||
prev_slot_indices=prev_slot_indices,
|
||||
expected_input_ids=expected_input_ids,
|
||||
capacity=capacity,
|
||||
device=device,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def make_write_plan(
|
||||
*,
|
||||
write_offsets: list[int],
|
||||
seed_slot_indices: list[int],
|
||||
num_valid_reqs: int,
|
||||
req_capacity: Optional[int] = None,
|
||||
device: torch.device,
|
||||
) -> WritePlan:
|
||||
"""Build a WritePlan from raw offsets and seed slot lists.
|
||||
|
||||
``write_offsets`` must have length ``len(seed_slot_indices) + 1`` (the trailing total entry count).
|
||||
"""
|
||||
n_active = len(seed_slot_indices)
|
||||
if len(write_offsets) != n_active + 1:
|
||||
raise ValueError(
|
||||
"make_write_plan: write_offsets must have length len(seed_slot_indices) + 1"
|
||||
)
|
||||
cap = req_capacity if req_capacity is not None else max(n_active, 1)
|
||||
plan = WritePlan.allocate(write_req_capacity=cap, device=device)
|
||||
if n_active > 0:
|
||||
plan.write_seed_slot_indices[:n_active] = torch.tensor(
|
||||
seed_slot_indices, dtype=torch.int64, device=device
|
||||
)
|
||||
plan.write_offsets[: n_active + 1] = torch.tensor(
|
||||
write_offsets, dtype=torch.int64, device=device
|
||||
)
|
||||
plan.write_num_valid_reqs[0] = num_valid_reqs
|
||||
return plan
|
||||
|
||||
|
||||
def make_write_plan_pair(
|
||||
*,
|
||||
write_offsets: list[int],
|
||||
seed_slot_indices: list[int],
|
||||
num_valid_reqs: int,
|
||||
req_capacity: Optional[int] = None,
|
||||
device: torch.device,
|
||||
) -> tuple[WritePlan, WritePlan]:
|
||||
return (
|
||||
make_write_plan(
|
||||
write_offsets=write_offsets,
|
||||
seed_slot_indices=seed_slot_indices,
|
||||
num_valid_reqs=num_valid_reqs,
|
||||
req_capacity=req_capacity,
|
||||
device=device,
|
||||
),
|
||||
make_write_plan(
|
||||
write_offsets=write_offsets,
|
||||
seed_slot_indices=seed_slot_indices,
|
||||
num_valid_reqs=num_valid_reqs,
|
||||
req_capacity=req_capacity,
|
||||
device=device,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def to_signed_int64(value: int) -> int:
|
||||
value &= _U64_MASK
|
||||
if value >= _I64_SIGN_BIT:
|
||||
value -= 1 << 64
|
||||
return value
|
||||
|
||||
|
||||
def chain_anchor_signed() -> int:
|
||||
return to_signed_int64(splitmix64(consts.CANARY_CHAIN_ANCHOR))
|
||||
|
||||
|
||||
def write_slot_fields(
|
||||
*,
|
||||
canary_buf: torch.Tensor,
|
||||
slot_idx: int,
|
||||
token: int,
|
||||
position: int,
|
||||
prev_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
|
||||
|
||||
|
||||
def stamp_pair(
|
||||
buf_pair: tuple[torch.Tensor, torch.Tensor],
|
||||
*,
|
||||
slot_idx: int,
|
||||
token: int,
|
||||
position: int,
|
||||
prev_hash: int,
|
||||
) -> None:
|
||||
"""Stamp the same slot fields into both (cuda, ref) canary buffers."""
|
||||
for buf in buf_pair:
|
||||
write_slot_fields(
|
||||
canary_buf=buf,
|
||||
slot_idx=slot_idx,
|
||||
token=token,
|
||||
position=position,
|
||||
prev_hash=prev_hash,
|
||||
)
|
||||
|
||||
|
||||
def read_slot_fields(
|
||||
*, canary_buf: torch.Tensor, slot_idx: int
|
||||
) -> tuple[int, int, int, int]:
|
||||
row = canary_buf.view(torch.int64)[slot_idx, :4].detach().cpu().tolist()
|
||||
return int(row[0]), int(row[1]), int(row[2]), int(row[3])
|
||||
|
||||
|
||||
def stamp_clean_chain(
|
||||
*,
|
||||
cuda_buf: torch.Tensor,
|
||||
ref_buf: torch.Tensor,
|
||||
slot_indices: list[int],
|
||||
tokens: list[int],
|
||||
positions: list[int],
|
||||
) -> list[int]:
|
||||
running_prev_hash = splitmix64(consts.CANARY_CHAIN_ANCHOR)
|
||||
stored_prev_hashes: list[int] = []
|
||||
for slot_idx, token, position in zip(slot_indices, tokens, positions):
|
||||
signed_prev = to_signed_int64(running_prev_hash)
|
||||
for buf in (cuda_buf, ref_buf):
|
||||
write_slot_fields(
|
||||
canary_buf=buf,
|
||||
slot_idx=slot_idx,
|
||||
token=token,
|
||||
position=position,
|
||||
prev_hash=signed_prev,
|
||||
)
|
||||
stored_prev_hashes.append(signed_prev)
|
||||
running_prev_hash = splitmix64_mix3(running_prev_hash, token, position)
|
||||
return stored_prev_hashes
|
||||
|
||||
|
||||
def assert_canary_state_equal(
|
||||
*, log_a: FakeViolationLog, log_b: FakeViolationLog
|
||||
) -> None:
|
||||
for name in ("ring", "write_index", "slot_run_counter", "kernel_run_counter"):
|
||||
assert torch.equal(
|
||||
getattr(log_a, name), getattr(log_b, name)
|
||||
), f"{name} diverged (CUDA vs ref)"
|
||||
|
||||
|
||||
def assert_canary_buf_equal(*, buf_a: torch.Tensor, buf_b: torch.Tensor) -> None:
|
||||
assert torch.equal(buf_a, buf_b), "canary_buf diverged (CUDA vs ref)"
|
||||
|
||||
|
||||
def assert_only_bits_set(fail_bits: int, expected_bits: int) -> None:
|
||||
assert (
|
||||
fail_bits & expected_bits
|
||||
) == expected_bits, (
|
||||
f"missing expected bits: expected {expected_bits:#b} got {fail_bits:#b}"
|
||||
)
|
||||
assert (
|
||||
fail_bits & ~expected_bits
|
||||
) == 0, f"unexpected extra bits: got {fail_bits:#b} extras {fail_bits & ~expected_bits:#b}"
|
||||
@@ -0,0 +1,150 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
from typing import Literal, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.kv_canary.verify import VerifyPlan
|
||||
from sglang.jit_kernel.kv_canary.write import WritePlan
|
||||
|
||||
_DEVICE = torch.device("cuda")
|
||||
|
||||
|
||||
LutKind = Literal["identity", "shift", "permutation", "with_oob"]
|
||||
|
||||
|
||||
def make_lut(
|
||||
*,
|
||||
kind: LutKind,
|
||||
pool_size: int,
|
||||
device: torch.device,
|
||||
rng: Optional[random.Random] = None,
|
||||
) -> torch.Tensor:
|
||||
base = torch.arange(pool_size + 1, dtype=torch.int64, device=device)
|
||||
if kind == "identity":
|
||||
return base.contiguous()
|
||||
if kind == "shift":
|
||||
return (base + 100).contiguous()
|
||||
if kind in ("permutation", "with_oob"):
|
||||
if rng is None:
|
||||
rng = random.Random(0)
|
||||
perm = list(range(pool_size + 1))
|
||||
rng.shuffle(perm)
|
||||
out = torch.tensor(perm, dtype=torch.int64, device=device)
|
||||
if kind == "with_oob":
|
||||
out[-1] = pool_size + 999
|
||||
return out.contiguous()
|
||||
raise ValueError(f"unknown LutKind: {kind}")
|
||||
|
||||
|
||||
ReqToTokenKind = Literal["linear", "sparse_permuted"]
|
||||
|
||||
|
||||
def make_req_to_token(
|
||||
*,
|
||||
kind: ReqToTokenKind,
|
||||
max_reqs: int,
|
||||
max_seq_len: int,
|
||||
device: torch.device,
|
||||
rng: Optional[random.Random] = None,
|
||||
) -> torch.Tensor:
|
||||
if kind == "linear":
|
||||
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)
|
||||
return (rp_axis * max_seq_len + pos_axis).contiguous()
|
||||
if rng is None:
|
||||
rng = random.Random(0)
|
||||
pool_size = max_reqs * max_seq_len
|
||||
# Slots index into a full_to_swa LUT sized [pool_size + 1], so values must stay
|
||||
# in [0, pool_size]. The universe spans [1, pool_size] (skipping 0 as reserved),
|
||||
# giving exactly max_reqs * max_seq_len unique slots — one per (rp, pos) cell.
|
||||
slot_universe = list(range(1, pool_size + 1))
|
||||
rng.shuffle(slot_universe)
|
||||
rtt = torch.zeros((max_reqs, max_seq_len), dtype=torch.int32, device=device)
|
||||
cursor = 0
|
||||
for rp in range(max_reqs):
|
||||
per_req = slot_universe[cursor : cursor + max_seq_len]
|
||||
cursor += max_seq_len
|
||||
rtt[rp, :] = torch.tensor(per_req, dtype=torch.int32, device=device)
|
||||
return rtt.contiguous()
|
||||
|
||||
|
||||
PaddingKind = Literal["none", "trailing", "interleaved"]
|
||||
|
||||
|
||||
def make_padding_mask(
|
||||
*,
|
||||
bs: int,
|
||||
kind: PaddingKind,
|
||||
rng: Optional[random.Random] = None,
|
||||
padding_fraction: float = 0.25,
|
||||
) -> list[bool]:
|
||||
if bs == 0:
|
||||
return []
|
||||
if kind == "none":
|
||||
return [False] * bs
|
||||
n_pad = max(1, int(bs * padding_fraction)) if bs > 0 else 0
|
||||
n_pad = min(n_pad, bs)
|
||||
if kind == "trailing":
|
||||
return [False] * (bs - n_pad) + [True] * n_pad
|
||||
if kind == "interleaved":
|
||||
if rng is None:
|
||||
rng = random.Random(0)
|
||||
mask = [False] * bs
|
||||
chosen = rng.sample(range(bs), k=n_pad)
|
||||
for idx in chosen:
|
||||
mask[idx] = True
|
||||
return mask
|
||||
raise ValueError(f"unknown PaddingKind: {kind}")
|
||||
|
||||
|
||||
CapacityKind = Literal["loose", "tight_match", "under_by_one"]
|
||||
|
||||
|
||||
def derive_plan_capacity(
|
||||
*,
|
||||
kind: CapacityKind,
|
||||
total_verify: int,
|
||||
extras_count: int,
|
||||
bs: int,
|
||||
) -> tuple[int, int]:
|
||||
needed = total_verify + extras_count
|
||||
if kind == "loose":
|
||||
return max(needed + 64, 128), max(bs + 4, 8)
|
||||
if kind == "tight_match":
|
||||
return max(needed, 1), max(bs + 4, 8)
|
||||
if kind == "under_by_one":
|
||||
return max(needed - 1, 1), max(bs + 4, 8)
|
||||
raise ValueError(f"unknown CapacityKind: {kind}")
|
||||
|
||||
|
||||
def allocate_plan_pair(
|
||||
*,
|
||||
verify_capacity: int,
|
||||
write_req_capacity: int,
|
||||
) -> tuple[VerifyPlan, WritePlan, VerifyPlan, WritePlan]:
|
||||
return (
|
||||
VerifyPlan.allocate(verify_capacity=verify_capacity, device=_DEVICE),
|
||||
WritePlan.allocate(write_req_capacity=write_req_capacity, device=_DEVICE),
|
||||
VerifyPlan.allocate(verify_capacity=verify_capacity, device=_DEVICE),
|
||||
WritePlan.allocate(write_req_capacity=write_req_capacity, device=_DEVICE),
|
||||
)
|
||||
|
||||
|
||||
def empty_extras() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
return (
|
||||
torch.zeros(1, dtype=torch.int64, device=_DEVICE),
|
||||
torch.zeros(1, dtype=torch.int64, device=_DEVICE),
|
||||
torch.zeros(1, dtype=torch.int64, device=_DEVICE),
|
||||
torch.zeros(1, dtype=torch.int32, device=_DEVICE),
|
||||
)
|
||||
|
||||
|
||||
def dummy_pseudo_tensors(num_tokens: int) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
return (
|
||||
torch.zeros(num_tokens, dtype=torch.int64, device=_DEVICE),
|
||||
torch.zeros(num_tokens, dtype=torch.int64, device=_DEVICE),
|
||||
)
|
||||
@@ -0,0 +1,520 @@
|
||||
"""Ref/real-independent invariant assertions for kv_canary kernel tests.
|
||||
|
||||
Each invariant only looks at the kernel's inputs and outputs (shape relationships, monotonicity, tail
|
||||
positions, etc.) — it must never re-implement the reference algorithm. Hand and fuzz tests both call
|
||||
into this module so a single contract violation surfaces consistently.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.kv_canary import consts
|
||||
from sglang.jit_kernel.kv_canary.verify import CanaryLaunchTag, VerifyPlan
|
||||
from sglang.jit_kernel.kv_canary.write import WritePlan
|
||||
from sglang.jit_kernel.tests.kv_canary._canary_helpers import FakeViolationLog
|
||||
|
||||
|
||||
class PlanInvariants:
|
||||
@staticmethod
|
||||
def assert_all(
|
||||
*,
|
||||
verify_plan: VerifyPlan,
|
||||
write_plan: WritePlan,
|
||||
req_pool_indices: torch.Tensor,
|
||||
prefix_lens: torch.Tensor,
|
||||
extend_seq_lens: torch.Tensor,
|
||||
swa_window_size: int,
|
||||
extras_slot_indices: torch.Tensor,
|
||||
extras_positions: torch.Tensor,
|
||||
extras_prev_slot_indices: torch.Tensor,
|
||||
extras_count: int,
|
||||
) -> None:
|
||||
PlanInvariants._assert_write_offsets_monotone(write_plan)
|
||||
PlanInvariants._assert_write_offsets_total_matches_active_extend_sum(
|
||||
write_plan=write_plan,
|
||||
extend_seq_lens=extend_seq_lens,
|
||||
req_pool_indices=req_pool_indices,
|
||||
)
|
||||
derived = PlanInvariants._assert_verify_num_valid_equals_derived_plus_extras(
|
||||
verify_plan=verify_plan,
|
||||
prefix_lens=prefix_lens,
|
||||
req_pool_indices=req_pool_indices,
|
||||
swa_window_size=swa_window_size,
|
||||
extras_count=extras_count,
|
||||
)
|
||||
PlanInvariants._assert_padding_row_seed_is_minus_one(
|
||||
write_plan=write_plan,
|
||||
req_pool_indices=req_pool_indices,
|
||||
)
|
||||
# In overflow (derived + extras > verify_capacity) the plan kernel disables
|
||||
# verify (enable=0) and the verify entries buffer is partially populated; the
|
||||
# downstream entry-shape invariants only hold when the kernel actually emitted
|
||||
# the full set, so guard them on enable=1.
|
||||
verify_enabled = int(verify_plan.enable[0].item()) == 1
|
||||
if verify_enabled:
|
||||
PlanInvariants._assert_extras_land_at_tail(
|
||||
verify_plan=verify_plan,
|
||||
derived_verify_count=derived,
|
||||
extras_slot_indices=extras_slot_indices,
|
||||
extras_positions=extras_positions,
|
||||
extras_prev_slot_indices=extras_prev_slot_indices,
|
||||
extras_count=extras_count,
|
||||
)
|
||||
PlanInvariants._assert_prev_slot_minus_one_iff_chain_head(
|
||||
verify_plan=verify_plan,
|
||||
swa_window_size=swa_window_size,
|
||||
derived_verify_count=derived,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _assert_write_offsets_monotone(write_plan: WritePlan) -> None:
|
||||
n_active = int(write_plan.write_num_valid_reqs[0].item())
|
||||
if n_active < 0:
|
||||
raise AssertionError(f"write_num_valid_reqs negative: {n_active}")
|
||||
offsets = write_plan.write_offsets[: n_active + 1].detach().cpu().tolist()
|
||||
for i in range(len(offsets) - 1):
|
||||
assert (
|
||||
offsets[i] <= offsets[i + 1]
|
||||
), f"write_offsets non-monotone at {i}: {offsets[i]} > {offsets[i + 1]}"
|
||||
|
||||
@staticmethod
|
||||
def _assert_write_offsets_total_matches_active_extend_sum(
|
||||
*,
|
||||
write_plan: WritePlan,
|
||||
extend_seq_lens: torch.Tensor,
|
||||
req_pool_indices: torch.Tensor,
|
||||
) -> None:
|
||||
n_active = int(write_plan.write_num_valid_reqs[0].item())
|
||||
total = int(write_plan.write_offsets[n_active].item())
|
||||
rpi_cpu = req_pool_indices.detach().cpu().tolist()
|
||||
ext_cpu = extend_seq_lens.detach().cpu().tolist()
|
||||
expected_total = sum(ext for rpi, ext in zip(rpi_cpu, ext_cpu) if rpi != 0)
|
||||
assert (
|
||||
total == expected_total
|
||||
), f"write_offsets total {total} != active extend sum {expected_total}"
|
||||
|
||||
@staticmethod
|
||||
def _assert_extras_land_at_tail(
|
||||
*,
|
||||
verify_plan: VerifyPlan,
|
||||
derived_verify_count: int,
|
||||
extras_slot_indices: torch.Tensor,
|
||||
extras_positions: torch.Tensor,
|
||||
extras_prev_slot_indices: torch.Tensor,
|
||||
extras_count: int,
|
||||
) -> None:
|
||||
if extras_count == 0:
|
||||
return
|
||||
tail_start = derived_verify_count
|
||||
tail_end = derived_verify_count + extras_count
|
||||
n_valid = int(verify_plan.verify_num_valid[0].item())
|
||||
assert (
|
||||
tail_end <= n_valid
|
||||
), f"extras tail {tail_end} exceeds verify_num_valid {n_valid}"
|
||||
plan_slots = verify_plan.verify_slot_indices[tail_start:tail_end]
|
||||
plan_positions = verify_plan.verify_expected_positions[tail_start:tail_end]
|
||||
plan_prevs = verify_plan.verify_prev_slot_indices[tail_start:tail_end]
|
||||
assert torch.equal(plan_slots, extras_slot_indices[:extras_count])
|
||||
assert torch.equal(plan_positions, extras_positions[:extras_count])
|
||||
assert torch.equal(plan_prevs, extras_prev_slot_indices[:extras_count])
|
||||
|
||||
@staticmethod
|
||||
def _assert_padding_row_seed_is_minus_one(
|
||||
*,
|
||||
write_plan: WritePlan,
|
||||
req_pool_indices: torch.Tensor,
|
||||
) -> None:
|
||||
n_active = int(write_plan.write_num_valid_reqs[0].item())
|
||||
if n_active == 0:
|
||||
return
|
||||
rpi_cpu = req_pool_indices.detach().cpu().tolist()
|
||||
seeds_cpu = (
|
||||
write_plan.write_seed_slot_indices[:n_active].detach().cpu().tolist()
|
||||
)
|
||||
for r in range(min(n_active, len(rpi_cpu))):
|
||||
if rpi_cpu[r] == 0:
|
||||
assert (
|
||||
seeds_cpu[r] == -1
|
||||
), f"padding row {r} has seed {seeds_cpu[r]} != -1"
|
||||
|
||||
@staticmethod
|
||||
def _assert_prev_slot_minus_one_iff_chain_head(
|
||||
*,
|
||||
verify_plan: VerifyPlan,
|
||||
swa_window_size: int,
|
||||
derived_verify_count: int,
|
||||
) -> None:
|
||||
if derived_verify_count == 0:
|
||||
return
|
||||
positions_cpu = (
|
||||
verify_plan.verify_expected_positions[:derived_verify_count]
|
||||
.detach()
|
||||
.cpu()
|
||||
.tolist()
|
||||
)
|
||||
prevs_cpu = (
|
||||
verify_plan.verify_prev_slot_indices[:derived_verify_count]
|
||||
.detach()
|
||||
.cpu()
|
||||
.tolist()
|
||||
)
|
||||
for i, (pos, prev) in enumerate(zip(positions_cpu, prevs_cpu)):
|
||||
if pos == 0:
|
||||
assert (
|
||||
prev == -1
|
||||
), f"entry {i} at position 0 must have prev=-1, got {prev}"
|
||||
else:
|
||||
if swa_window_size == 0:
|
||||
assert (
|
||||
prev != -1
|
||||
), f"FULL entry {i} at position {pos} must have prev != -1, got {prev}"
|
||||
|
||||
@staticmethod
|
||||
def _assert_verify_num_valid_equals_derived_plus_extras(
|
||||
*,
|
||||
verify_plan: VerifyPlan,
|
||||
prefix_lens: torch.Tensor,
|
||||
req_pool_indices: torch.Tensor,
|
||||
swa_window_size: int,
|
||||
extras_count: int,
|
||||
) -> int:
|
||||
rpi_cpu = req_pool_indices.detach().cpu().tolist()
|
||||
pfx_cpu = prefix_lens.detach().cpu().tolist()
|
||||
derived = 0
|
||||
for rpi, pfx in zip(rpi_cpu, pfx_cpu):
|
||||
if rpi == 0:
|
||||
continue
|
||||
if swa_window_size > 0:
|
||||
window_start = max(0, pfx - swa_window_size)
|
||||
derived += max(0, pfx - window_start)
|
||||
else:
|
||||
derived += max(0, pfx)
|
||||
# The plan kernel clamps verify_num_valid to verify_capacity and turns enable
|
||||
# off when (derived + extras) overflows the slot indices buffer. The invariant
|
||||
# must match that: on overflow the kernel records the capacity, on no overflow
|
||||
# it records the exact derived total (so the verify kernel scans every row).
|
||||
verify_capacity = int(verify_plan.verify_slot_indices.shape[0])
|
||||
expected_unclamped = derived + extras_count
|
||||
expected = min(expected_unclamped, verify_capacity)
|
||||
overflow = expected_unclamped > verify_capacity
|
||||
actual = int(verify_plan.verify_num_valid[0].item())
|
||||
assert actual == expected, (
|
||||
f"verify_num_valid {actual} != min(derived {derived} + extras {extras_count}, "
|
||||
f"verify_capacity {verify_capacity}) = {expected}"
|
||||
)
|
||||
enable = int(verify_plan.enable[0].item())
|
||||
expected_enable = 0 if overflow else 1
|
||||
assert enable == expected_enable, (
|
||||
f"verify_plan.enable {enable} != expected {expected_enable} "
|
||||
f"(overflow={overflow}; derived+extras={expected_unclamped}, "
|
||||
f"verify_capacity={verify_capacity})"
|
||||
)
|
||||
return derived
|
||||
|
||||
|
||||
class VerifyInvariants:
|
||||
@staticmethod
|
||||
def assert_all(
|
||||
*,
|
||||
canary_buf_before: torch.Tensor,
|
||||
canary_buf_after: torch.Tensor,
|
||||
log_before: FakeViolationLog,
|
||||
log_after: FakeViolationLog,
|
||||
plan: VerifyPlan,
|
||||
kernel_kind: CanaryLaunchTag,
|
||||
) -> None:
|
||||
VerifyInvariants._assert_canary_buf_unchanged(
|
||||
canary_buf_before=canary_buf_before, canary_buf_after=canary_buf_after
|
||||
)
|
||||
VerifyInvariants._assert_violation_count_le_active_entries(
|
||||
log_after=log_after, log_before=log_before, plan=plan
|
||||
)
|
||||
VerifyInvariants._assert_violation_rows_have_valid_slot_and_kernel_kind(
|
||||
log_after=log_after,
|
||||
log_before=log_before,
|
||||
plan=plan,
|
||||
kernel_kind=kernel_kind,
|
||||
)
|
||||
VerifyInvariants._assert_slot_run_counter_incremented_by_active_entries(
|
||||
log_before=log_before, log_after=log_after, plan=plan
|
||||
)
|
||||
VerifyInvariants._assert_kernel_run_counter_incremented_by_one(
|
||||
log_before=log_before, log_after=log_after
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _assert_canary_buf_unchanged(
|
||||
*,
|
||||
canary_buf_before: torch.Tensor,
|
||||
canary_buf_after: torch.Tensor,
|
||||
) -> None:
|
||||
assert torch.equal(
|
||||
canary_buf_before, canary_buf_after
|
||||
), "verify kernel mutated canary_buf (must be read-only)"
|
||||
|
||||
@staticmethod
|
||||
def _assert_violation_count_le_active_entries(
|
||||
*,
|
||||
log_after: FakeViolationLog,
|
||||
log_before: FakeViolationLog,
|
||||
plan: VerifyPlan,
|
||||
) -> None:
|
||||
delta = int(log_after.write_index[0].item()) - int(
|
||||
log_before.write_index[0].item()
|
||||
)
|
||||
n_active = int(plan.verify_num_valid[0].item())
|
||||
assert (
|
||||
0 <= delta <= n_active
|
||||
), f"violation_write_index delta {delta} out of [0, {n_active}]"
|
||||
|
||||
@staticmethod
|
||||
def _assert_violation_rows_have_valid_slot_and_kernel_kind(
|
||||
*,
|
||||
log_after: FakeViolationLog,
|
||||
log_before: FakeViolationLog,
|
||||
plan: VerifyPlan,
|
||||
kernel_kind: CanaryLaunchTag,
|
||||
) -> None:
|
||||
write_idx_after = int(log_after.write_index[0].item())
|
||||
write_idx_before = int(log_before.write_index[0].item())
|
||||
if write_idx_after == write_idx_before:
|
||||
return
|
||||
ring_capacity = log_after.ring.shape[0]
|
||||
visible_start = write_idx_before
|
||||
visible_end = min(write_idx_after, ring_capacity)
|
||||
if visible_end <= visible_start:
|
||||
return
|
||||
n_active = int(plan.verify_num_valid[0].item())
|
||||
plan_slots = set(plan.verify_slot_indices[:n_active].detach().cpu().tolist())
|
||||
rows = log_after.ring[visible_start:visible_end].detach().cpu()
|
||||
for i in range(rows.shape[0]):
|
||||
kind = int(rows[i, consts.VIOLATION_FIELD_KERNEL_KIND].item())
|
||||
assert kind == int(
|
||||
kernel_kind
|
||||
), f"row {visible_start + i} kernel_kind {kind} != expected {int(kernel_kind)}"
|
||||
slot = int(rows[i, consts.VIOLATION_FIELD_SLOT_IDX].item())
|
||||
assert (
|
||||
slot in plan_slots
|
||||
), f"row {visible_start + i} slot {slot} not in plan_slots"
|
||||
|
||||
@staticmethod
|
||||
def _assert_slot_run_counter_incremented_by_active_entries(
|
||||
*,
|
||||
log_before: FakeViolationLog,
|
||||
log_after: FakeViolationLog,
|
||||
plan: VerifyPlan,
|
||||
) -> None:
|
||||
n_active = int(plan.verify_num_valid[0].item())
|
||||
delta = int(log_after.slot_run_counter[0].item()) - int(
|
||||
log_before.slot_run_counter[0].item()
|
||||
)
|
||||
assert (
|
||||
delta == n_active
|
||||
), f"slot_run_counter delta {delta} != active entries {n_active}"
|
||||
|
||||
@staticmethod
|
||||
def _assert_kernel_run_counter_incremented_by_one(
|
||||
*,
|
||||
log_before: FakeViolationLog,
|
||||
log_after: FakeViolationLog,
|
||||
) -> None:
|
||||
delta = int(log_after.kernel_run_counter[0].item()) - int(
|
||||
log_before.kernel_run_counter[0].item()
|
||||
)
|
||||
assert delta == 1, f"kernel_run_counter delta {delta} != 1"
|
||||
|
||||
|
||||
class WriteInvariants:
|
||||
@staticmethod
|
||||
def assert_all(
|
||||
*,
|
||||
canary_buf_before: torch.Tensor,
|
||||
canary_buf_after: torch.Tensor,
|
||||
plan: WritePlan,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
out_cache_loc: torch.Tensor,
|
||||
enable_write_verify_inputs: bool,
|
||||
expected_input_tokens: Optional[torch.Tensor],
|
||||
expected_input_positions: Optional[torch.Tensor],
|
||||
log_before: FakeViolationLog,
|
||||
log_after: FakeViolationLog,
|
||||
) -> None:
|
||||
WriteInvariants._assert_written_slots_token_position_match_input(
|
||||
canary_buf_after=canary_buf_after,
|
||||
plan=plan,
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
)
|
||||
WriteInvariants._assert_slot_minus_one_skipped(
|
||||
canary_buf_before=canary_buf_before,
|
||||
canary_buf_after=canary_buf_after,
|
||||
plan=plan,
|
||||
out_cache_loc=out_cache_loc,
|
||||
)
|
||||
WriteInvariants._assert_pseudo_violation_only_on_mismatch(
|
||||
enable_write_verify_inputs=enable_write_verify_inputs,
|
||||
log_before=log_before,
|
||||
log_after=log_after,
|
||||
expected_input_tokens=expected_input_tokens,
|
||||
expected_input_positions=expected_input_positions,
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
plan=plan,
|
||||
)
|
||||
WriteInvariants._assert_write_slot_run_counter_incremented(
|
||||
log_before=log_before,
|
||||
log_after=log_after,
|
||||
plan=plan,
|
||||
out_cache_loc=out_cache_loc,
|
||||
)
|
||||
WriteInvariants._assert_write_kernel_run_counter_incremented_by_one(
|
||||
log_before=log_before, log_after=log_after
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _assert_written_slots_token_position_match_input(
|
||||
*,
|
||||
canary_buf_after: torch.Tensor,
|
||||
plan: WritePlan,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
out_cache_loc: torch.Tensor,
|
||||
) -> None:
|
||||
n_active = int(plan.write_num_valid_reqs[0].item())
|
||||
if n_active == 0:
|
||||
return
|
||||
offsets = plan.write_offsets[: n_active + 1].detach().cpu().tolist()
|
||||
total = offsets[n_active]
|
||||
slots_cpu = out_cache_loc[:total].detach().cpu().tolist()
|
||||
tokens_cpu = input_ids[:total].detach().cpu().tolist()
|
||||
pos_cpu = positions[:total].detach().cpu().tolist()
|
||||
view = canary_buf_after.view(torch.int64)
|
||||
for i in range(total):
|
||||
slot = slots_cpu[i]
|
||||
if slot < 0:
|
||||
continue
|
||||
stored_token = int(view[slot, 0].item())
|
||||
stored_position = int(view[slot, 1].item())
|
||||
assert (
|
||||
stored_token == tokens_cpu[i]
|
||||
), f"slot {slot}: stored token {stored_token} != input {tokens_cpu[i]}"
|
||||
assert (
|
||||
stored_position == pos_cpu[i]
|
||||
), f"slot {slot}: stored position {stored_position} != input {pos_cpu[i]}"
|
||||
|
||||
@staticmethod
|
||||
def _assert_slot_minus_one_skipped(
|
||||
*,
|
||||
canary_buf_before: torch.Tensor,
|
||||
canary_buf_after: torch.Tensor,
|
||||
plan: WritePlan,
|
||||
out_cache_loc: torch.Tensor,
|
||||
) -> None:
|
||||
n_active = int(plan.write_num_valid_reqs[0].item())
|
||||
if n_active == 0:
|
||||
return
|
||||
total = int(plan.write_offsets[n_active].item())
|
||||
slots_cpu = out_cache_loc[:total].detach().cpu().tolist()
|
||||
written_slots = {s for s in slots_cpu if s >= 0}
|
||||
view_before = canary_buf_before.view(torch.int64)
|
||||
view_after = canary_buf_after.view(torch.int64)
|
||||
num_slots = canary_buf_after.shape[0]
|
||||
for slot in range(num_slots):
|
||||
if slot in written_slots:
|
||||
continue
|
||||
assert torch.equal(
|
||||
view_before[slot], view_after[slot]
|
||||
), f"slot {slot} not in out_cache_loc but canary_buf changed"
|
||||
|
||||
@staticmethod
|
||||
def _assert_pseudo_violation_only_on_mismatch(
|
||||
*,
|
||||
enable_write_verify_inputs: bool,
|
||||
log_before: FakeViolationLog,
|
||||
log_after: FakeViolationLog,
|
||||
expected_input_tokens: Optional[torch.Tensor],
|
||||
expected_input_positions: Optional[torch.Tensor],
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
out_cache_loc: torch.Tensor,
|
||||
plan: WritePlan,
|
||||
) -> None:
|
||||
delta = int(log_after.write_index[0].item()) - int(
|
||||
log_before.write_index[0].item()
|
||||
)
|
||||
if not enable_write_verify_inputs:
|
||||
assert (
|
||||
delta == 0
|
||||
), f"enable_write_verify_inputs=OFF must produce no violations, got {delta}"
|
||||
return
|
||||
if expected_input_tokens is None or expected_input_positions is None:
|
||||
return
|
||||
n_active = int(plan.write_num_valid_reqs[0].item())
|
||||
if n_active == 0:
|
||||
assert delta == 0, f"empty plan produced {delta} violations"
|
||||
return
|
||||
total = int(plan.write_offsets[n_active].item())
|
||||
tok = input_ids[:total].detach().cpu().tolist()
|
||||
pos = positions[:total].detach().cpu().tolist()
|
||||
exp_tok = expected_input_tokens[:total].detach().cpu().tolist()
|
||||
exp_pos = expected_input_positions[:total].detach().cpu().tolist()
|
||||
slots_cpu = out_cache_loc[:total].detach().cpu().tolist()
|
||||
mismatch_entries = sum(
|
||||
1
|
||||
for i in range(total)
|
||||
if slots_cpu[i] >= 0 and (tok[i] != exp_tok[i] or pos[i] != exp_pos[i])
|
||||
)
|
||||
no_mismatch = mismatch_entries == 0
|
||||
if no_mismatch:
|
||||
assert (
|
||||
delta == 0
|
||||
), f"enable_write_verify_inputs=ON with no mismatch produced {delta} violations"
|
||||
else:
|
||||
assert (
|
||||
delta == mismatch_entries
|
||||
), f"write input mismatch count {mismatch_entries} produced {delta} violations"
|
||||
|
||||
@staticmethod
|
||||
def _assert_write_slot_run_counter_incremented(
|
||||
*,
|
||||
log_before: FakeViolationLog,
|
||||
log_after: FakeViolationLog,
|
||||
plan: WritePlan,
|
||||
out_cache_loc: torch.Tensor,
|
||||
) -> None:
|
||||
n_active = int(plan.write_num_valid_reqs[0].item())
|
||||
if n_active == 0:
|
||||
delta = int(log_after.slot_run_counter[0].item()) - int(
|
||||
log_before.slot_run_counter[0].item()
|
||||
)
|
||||
assert delta == 0, f"empty plan incremented slot_run_counter by {delta}"
|
||||
return
|
||||
total = int(plan.write_offsets[n_active].item())
|
||||
# The write kernel skips entries where out_cache_loc < 0 (the documented "mark
|
||||
# skip" path used by SWA-translated callers), so the slot_run_counter delta
|
||||
# tracks the count of writeable entries, not the planned total.
|
||||
writeable = int((out_cache_loc[:total] >= 0).sum().item())
|
||||
delta = int(log_after.slot_run_counter[0].item()) - int(
|
||||
log_before.slot_run_counter[0].item()
|
||||
)
|
||||
assert delta == writeable, (
|
||||
f"slot_run_counter delta {delta} != writeable entries {writeable} "
|
||||
f"(total={total}, skipped={total - writeable})"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _assert_write_kernel_run_counter_incremented_by_one(
|
||||
*,
|
||||
log_before: FakeViolationLog,
|
||||
log_after: FakeViolationLog,
|
||||
) -> None:
|
||||
delta = int(log_after.kernel_run_counter[0].item()) - int(
|
||||
log_before.kernel_run_counter[0].item()
|
||||
)
|
||||
assert delta == 1, f"kernel_run_counter delta {delta} != 1"
|
||||
Reference in New Issue
Block a user