Add the KV-canary plan JIT kernels (#26807)
This commit is contained in:
fzyzcjy
@@ -0,0 +1,500 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, replace
|
||||
from typing import Any, Callable, Iterator, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.kv_canary.plan import launch_canary_plan_kernels
|
||||
from sglang.jit_kernel.kv_canary.plan_ref import (
|
||||
launch_canary_plan_kernels_torch_reference,
|
||||
)
|
||||
from sglang.jit_kernel.kv_canary.verify import (
|
||||
CanaryLaunchTag,
|
||||
VerifyOrWriteContext,
|
||||
VerifyPlan,
|
||||
launch_canary_verify_kernel,
|
||||
)
|
||||
from sglang.jit_kernel.kv_canary.verify_ref import (
|
||||
launch_canary_verify_kernel_torch_reference,
|
||||
)
|
||||
from sglang.jit_kernel.kv_canary.write import WritePlan, launch_canary_write_kernel
|
||||
from sglang.jit_kernel.kv_canary.write_ref import (
|
||||
launch_canary_write_kernel_torch_reference,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
|
||||
FakeViolationLog,
|
||||
assert_canary_buf_equal,
|
||||
assert_canary_state_equal,
|
||||
make_log_pair,
|
||||
)
|
||||
|
||||
_DEVICE = torch.device("cuda")
|
||||
|
||||
|
||||
def _run_both_plan(
|
||||
*,
|
||||
triton_verify: VerifyPlan,
|
||||
triton_write: WritePlan,
|
||||
ref_verify: VerifyPlan,
|
||||
ref_write: WritePlan,
|
||||
req_pool_indices: torch.Tensor,
|
||||
prefix_lens: torch.Tensor,
|
||||
extend_seq_lens: torch.Tensor,
|
||||
req_to_token: torch.Tensor,
|
||||
extras: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
|
||||
swa_window_size: int,
|
||||
full_to_swa_index_mapping: Optional[torch.Tensor],
|
||||
assert_equal: bool = True,
|
||||
active_verify_entries: Optional[int] = None,
|
||||
active_write_reqs: Optional[int] = None,
|
||||
req_to_verify_expected_tokens: Optional[torch.Tensor] = None,
|
||||
req_to_verify_expected_tokens_valid_lens: Optional[torch.Tensor] = None,
|
||||
kv_token_id_vs_position_offset: int = 0,
|
||||
) -> None:
|
||||
_ = extras
|
||||
verify_capacity = int(triton_verify.verify_slot_indices.shape[0])
|
||||
# Default lens to "no tighter bound than pool width" so existing kernel tests that
|
||||
# only care about gather wiring keep their old semantics without each call site
|
||||
# explicitly building a per-req lens tensor.
|
||||
if (
|
||||
req_to_verify_expected_tokens is not None
|
||||
and req_to_verify_expected_tokens_valid_lens is None
|
||||
):
|
||||
req_to_verify_expected_tokens_valid_lens = torch.full(
|
||||
(int(req_pool_indices.shape[0]),),
|
||||
int(req_to_verify_expected_tokens.shape[1]),
|
||||
dtype=torch.int64,
|
||||
device=req_pool_indices.device,
|
||||
)
|
||||
launch_canary_plan_kernels(
|
||||
verify_plan_out=triton_verify,
|
||||
write_plan_out=triton_write,
|
||||
req_pool_indices=req_pool_indices,
|
||||
prefix_lens=prefix_lens,
|
||||
extend_seq_lens=extend_seq_lens,
|
||||
req_to_token=req_to_token,
|
||||
swa_window_size=swa_window_size,
|
||||
full_to_swa_index_mapping=full_to_swa_index_mapping,
|
||||
verify_capacity=verify_capacity,
|
||||
req_to_verify_expected_tokens=req_to_verify_expected_tokens,
|
||||
req_to_verify_expected_tokens_valid_lens=req_to_verify_expected_tokens_valid_lens,
|
||||
kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
|
||||
)
|
||||
launch_canary_plan_kernels_torch_reference(
|
||||
verify_plan_out=ref_verify,
|
||||
write_plan_out=ref_write,
|
||||
req_pool_indices=req_pool_indices,
|
||||
prefix_lens=prefix_lens,
|
||||
extend_seq_lens=extend_seq_lens,
|
||||
req_to_token=req_to_token,
|
||||
swa_window_size=swa_window_size,
|
||||
full_to_swa_index_mapping=full_to_swa_index_mapping,
|
||||
verify_capacity=int(ref_verify.verify_slot_indices.shape[0]),
|
||||
req_to_verify_expected_tokens=req_to_verify_expected_tokens,
|
||||
req_to_verify_expected_tokens_valid_lens=req_to_verify_expected_tokens_valid_lens,
|
||||
kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
if assert_equal:
|
||||
_assert_plans_byte_equal(
|
||||
triton_verify=triton_verify,
|
||||
triton_write=triton_write,
|
||||
ref_verify=ref_verify,
|
||||
ref_write=ref_write,
|
||||
active_verify_entries=active_verify_entries,
|
||||
active_write_reqs=active_write_reqs,
|
||||
)
|
||||
|
||||
|
||||
def _assert_plans_byte_equal(
|
||||
*,
|
||||
triton_verify: VerifyPlan,
|
||||
triton_write: WritePlan,
|
||||
ref_verify: VerifyPlan,
|
||||
ref_write: WritePlan,
|
||||
active_verify_entries: Optional[int] = None,
|
||||
active_write_reqs: Optional[int] = None,
|
||||
) -> None:
|
||||
"""Byte-equal check on (Triton vs ref) plan outputs.
|
||||
|
||||
Optional ``active_verify_entries`` / ``active_write_reqs`` truncate the comparison to the meaningful
|
||||
prefix; tail entries past the active count are kernel-undefined and need not match byte-equal.
|
||||
"""
|
||||
n_verify = (
|
||||
active_verify_entries
|
||||
if active_verify_entries is not None
|
||||
else int(triton_verify.verify_num_valid[0].item())
|
||||
)
|
||||
n_verify_ref = int(ref_verify.verify_num_valid[0].item())
|
||||
assert (
|
||||
n_verify == n_verify_ref
|
||||
), f"verify_num_valid diverged: triton={n_verify} ref={n_verify_ref}"
|
||||
# When total_verify > VERIFY_CAPACITY the offsets kernel clears verify_enable and
|
||||
# plan_entries skips its scatter — leaving verify_slot_indices/positions/prev_slot_indices
|
||||
# as whatever the (torch.empty) allocation contained. Skip the byte-equal probe in that
|
||||
# case; verify_num_valid being clamped + verify_enable=0 is the contract here.
|
||||
triton_enable = int(triton_verify.enable[0].item())
|
||||
ref_enable = int(ref_verify.enable[0].item())
|
||||
assert (
|
||||
triton_enable == ref_enable
|
||||
), f"verify_enable diverged: triton={triton_enable} ref={ref_enable}"
|
||||
if n_verify > 0 and triton_enable != 0:
|
||||
assert torch.equal(
|
||||
triton_verify.verify_slot_indices[:n_verify],
|
||||
ref_verify.verify_slot_indices[:n_verify],
|
||||
)
|
||||
assert torch.equal(
|
||||
triton_verify.verify_expected_tokens[:n_verify],
|
||||
ref_verify.verify_expected_tokens[:n_verify],
|
||||
)
|
||||
assert torch.equal(
|
||||
triton_verify.verify_expected_positions[:n_verify],
|
||||
ref_verify.verify_expected_positions[:n_verify],
|
||||
)
|
||||
assert torch.equal(
|
||||
triton_verify.verify_prev_slot_indices[:n_verify],
|
||||
ref_verify.verify_prev_slot_indices[:n_verify],
|
||||
)
|
||||
|
||||
n_write = (
|
||||
active_write_reqs
|
||||
if active_write_reqs is not None
|
||||
else int(triton_write.write_num_valid_reqs[0].item())
|
||||
)
|
||||
n_write_ref = int(ref_write.write_num_valid_reqs[0].item())
|
||||
assert (
|
||||
n_write == n_write_ref
|
||||
), f"write_num_valid_reqs diverged: triton={n_write} ref={n_write_ref}"
|
||||
assert torch.equal(
|
||||
triton_write.write_offsets[: n_write + 1],
|
||||
ref_write.write_offsets[: n_write + 1],
|
||||
)
|
||||
if n_write > 0:
|
||||
assert torch.equal(
|
||||
triton_write.write_seed_slot_indices[:n_write],
|
||||
ref_write.write_seed_slot_indices[:n_write],
|
||||
)
|
||||
|
||||
|
||||
def _run_both_verify(
|
||||
*,
|
||||
cuda_canary_buf: torch.Tensor,
|
||||
ref_canary_buf: torch.Tensor,
|
||||
plan_cuda,
|
||||
plan_ref,
|
||||
cuda_log: FakeViolationLog,
|
||||
ref_log: FakeViolationLog,
|
||||
kernel_kind: CanaryLaunchTag = CanaryLaunchTag.HEAD_K_FULL,
|
||||
assert_equal: bool = True,
|
||||
check_verify_expected_token: bool = True,
|
||||
) -> None:
|
||||
launch_canary_verify_kernel(
|
||||
context=VerifyOrWriteContext(
|
||||
canary_buf=cuda_canary_buf,
|
||||
kernel_kind=kernel_kind,
|
||||
violation_ring=cuda_log.ring,
|
||||
violation_write_index=cuda_log.write_index,
|
||||
slot_run_counter=cuda_log.slot_run_counter,
|
||||
kernel_run_counter=cuda_log.kernel_run_counter,
|
||||
enable_chain_position_assert=cuda_log.enable_chain_position_assert,
|
||||
),
|
||||
plan=plan_cuda,
|
||||
check_verify_expected_token=check_verify_expected_token,
|
||||
)
|
||||
launch_canary_verify_kernel_torch_reference(
|
||||
context=VerifyOrWriteContext(
|
||||
canary_buf=ref_canary_buf,
|
||||
kernel_kind=kernel_kind,
|
||||
violation_ring=ref_log.ring,
|
||||
violation_write_index=ref_log.write_index,
|
||||
slot_run_counter=ref_log.slot_run_counter,
|
||||
kernel_run_counter=ref_log.kernel_run_counter,
|
||||
enable_chain_position_assert=ref_log.enable_chain_position_assert,
|
||||
),
|
||||
plan=plan_ref,
|
||||
check_verify_expected_token=check_verify_expected_token,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
if assert_equal:
|
||||
assert_canary_state_equal(log_a=cuda_log, log_b=ref_log)
|
||||
|
||||
|
||||
def _run_both_write(
|
||||
*,
|
||||
cuda_canary_buf: torch.Tensor,
|
||||
ref_canary_buf: torch.Tensor,
|
||||
plan_cuda,
|
||||
plan_ref,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
out_cache_loc: torch.Tensor,
|
||||
enable_write_verify_inputs: bool,
|
||||
expected_input_tokens: torch.Tensor,
|
||||
expected_input_positions: torch.Tensor,
|
||||
cuda_log: FakeViolationLog,
|
||||
ref_log: FakeViolationLog,
|
||||
kernel_kind: CanaryLaunchTag = CanaryLaunchTag.HEAD_K_FULL,
|
||||
assert_equal: bool = True,
|
||||
) -> None:
|
||||
expected_tokens_for_launch = (
|
||||
expected_input_tokens if enable_write_verify_inputs else None
|
||||
)
|
||||
expected_positions_for_launch = (
|
||||
expected_input_positions if enable_write_verify_inputs else None
|
||||
)
|
||||
launch_canary_write_kernel(
|
||||
context=VerifyOrWriteContext(
|
||||
canary_buf=cuda_canary_buf,
|
||||
kernel_kind=kernel_kind,
|
||||
violation_ring=cuda_log.ring,
|
||||
violation_write_index=cuda_log.write_index,
|
||||
slot_run_counter=cuda_log.slot_run_counter,
|
||||
kernel_run_counter=cuda_log.kernel_run_counter,
|
||||
enable_chain_position_assert=cuda_log.enable_chain_position_assert,
|
||||
),
|
||||
plan=plan_cuda,
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
enable_write_input_assert=enable_write_verify_inputs,
|
||||
expected_input_tokens=expected_tokens_for_launch,
|
||||
expected_input_positions=expected_positions_for_launch,
|
||||
)
|
||||
launch_canary_write_kernel_torch_reference(
|
||||
context=VerifyOrWriteContext(
|
||||
canary_buf=ref_canary_buf,
|
||||
kernel_kind=kernel_kind,
|
||||
violation_ring=ref_log.ring,
|
||||
violation_write_index=ref_log.write_index,
|
||||
slot_run_counter=ref_log.slot_run_counter,
|
||||
kernel_run_counter=ref_log.kernel_run_counter,
|
||||
enable_chain_position_assert=ref_log.enable_chain_position_assert,
|
||||
),
|
||||
plan=plan_ref,
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
enable_write_input_assert=enable_write_verify_inputs,
|
||||
expected_input_tokens=expected_tokens_for_launch,
|
||||
expected_input_positions=expected_positions_for_launch,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
if assert_equal:
|
||||
assert_canary_buf_equal(buf_a=cuda_canary_buf, buf_b=ref_canary_buf)
|
||||
assert_canary_state_equal(log_a=cuda_log, log_b=ref_log)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True, kw_only=True)
|
||||
class ShrinkResult:
|
||||
inputs: Any
|
||||
mutations_applied: list[str]
|
||||
|
||||
|
||||
def shrink_inputs(
|
||||
inputs: Any,
|
||||
*,
|
||||
check_fn: Callable[[Any], bool],
|
||||
max_iterations: int = 50,
|
||||
) -> ShrinkResult:
|
||||
"""Greedy 1-step minify for a fuzz inputs dataclass.
|
||||
|
||||
``check_fn(candidate)`` returns True when ``candidate`` still reproduces the failure. Each round
|
||||
yields candidate-simpler-than-current mutations through ``_yield_simpler``; the first accepted
|
||||
candidate becomes the new current. Iteration stops when no mutation is accepted or ``max_iterations``
|
||||
is reached.
|
||||
"""
|
||||
current = inputs
|
||||
applied: list[str] = []
|
||||
for _ in range(max_iterations):
|
||||
improved = False
|
||||
for label, candidate in _yield_simpler(current):
|
||||
try:
|
||||
still_fails = check_fn(candidate)
|
||||
except Exception:
|
||||
still_fails = False
|
||||
if still_fails:
|
||||
current = candidate
|
||||
applied.append(label)
|
||||
improved = True
|
||||
break
|
||||
if not improved:
|
||||
break
|
||||
return ShrinkResult(inputs=current, mutations_applied=applied)
|
||||
|
||||
|
||||
def _yield_simpler(inputs: Any) -> Iterator[tuple[str, Any]]:
|
||||
"""Yield (label, simpler_candidate) tuples for generic fuzz-input minifiers.
|
||||
|
||||
The candidates touch only well-known field names; an inputs dataclass that lacks a field will simply
|
||||
have that mutation skipped. No kernel-specific knowledge is encoded here so the same shrinker drives
|
||||
Plan / Verify / Write fuzz failures uniformly.
|
||||
"""
|
||||
fields = {
|
||||
f: getattr(inputs, f) for f in inputs.__dataclass_fields__ # type: ignore[attr-defined]
|
||||
}
|
||||
|
||||
def emit(label: str, **overrides: Any) -> Iterator[tuple[str, Any]]:
|
||||
candidate = replace(inputs, **overrides)
|
||||
yield label, candidate
|
||||
|
||||
bs_field = (
|
||||
"req_pool_indices"
|
||||
if "req_pool_indices" in fields
|
||||
else ("input_ids" if "input_ids" in fields else None)
|
||||
)
|
||||
if bs_field is not None and isinstance(fields[bs_field], torch.Tensor):
|
||||
tensor = fields[bs_field]
|
||||
if tensor.numel() > 1:
|
||||
new_len = tensor.numel() - 1
|
||||
related_tensors_overrides: dict[str, Any] = {}
|
||||
for name in (
|
||||
"req_pool_indices",
|
||||
"prefix_lens",
|
||||
"extend_seq_lens",
|
||||
"input_ids",
|
||||
"positions",
|
||||
"out_cache_loc",
|
||||
"expected_input_tokens",
|
||||
"expected_input_positions",
|
||||
):
|
||||
t = fields.get(name)
|
||||
if (
|
||||
isinstance(t, torch.Tensor)
|
||||
and t.numel() >= new_len
|
||||
and t.dim() == 1
|
||||
):
|
||||
related_tensors_overrides[name] = t[:new_len].contiguous()
|
||||
if related_tensors_overrides:
|
||||
yield from emit("drop_last_row", **related_tensors_overrides)
|
||||
|
||||
if "swa_window_size" in fields and isinstance(fields["swa_window_size"], int):
|
||||
if fields["swa_window_size"] != 0:
|
||||
yield from emit(
|
||||
"swa_off", swa_window_size=0, full_to_swa_index_mapping=None
|
||||
)
|
||||
|
||||
if "extras_count" in fields and isinstance(fields["extras_count"], int):
|
||||
if fields["extras_count"] > 0:
|
||||
yield from emit("extras_zero", extras_count=0)
|
||||
|
||||
if "enable_write_verify_inputs" in fields:
|
||||
cur = fields["enable_write_verify_inputs"]
|
||||
if hasattr(cur, "value") and int(cur) != 0:
|
||||
cls = cur.__class__
|
||||
yield from emit("pseudo_off", enable_write_verify_inputs=cls(0))
|
||||
|
||||
for name in ("verify_capacity", "write_req_capacity"):
|
||||
if name in fields and isinstance(fields[name], int):
|
||||
current_value = fields[name]
|
||||
if current_value > 8:
|
||||
yield from emit(f"shrink_{name}", **{name: max(8, current_value // 2)})
|
||||
|
||||
|
||||
def run_verify_diff(
|
||||
*,
|
||||
buf_pair: tuple[torch.Tensor, torch.Tensor],
|
||||
plan_pair: tuple[VerifyPlan, VerifyPlan],
|
||||
kernel_kind: CanaryLaunchTag = CanaryLaunchTag.HEAD_K_FULL,
|
||||
device: torch.device = _DEVICE,
|
||||
assert_equal: bool = True,
|
||||
check_verify_expected_token: bool = True,
|
||||
) -> tuple[FakeViolationLog, FakeViolationLog]:
|
||||
"""Thin wrapper around ``_run_both_verify`` that creates a fresh log pair and packs (cuda, ref)
|
||||
buf/plan arguments into 2-tuples to drop ~8 lines of boilerplate per call site.
|
||||
"""
|
||||
cuda_log, ref_log = make_log_pair(device=device)
|
||||
_run_both_verify(
|
||||
cuda_canary_buf=buf_pair[0],
|
||||
ref_canary_buf=buf_pair[1],
|
||||
plan_cuda=plan_pair[0],
|
||||
plan_ref=plan_pair[1],
|
||||
cuda_log=cuda_log,
|
||||
ref_log=ref_log,
|
||||
kernel_kind=kernel_kind,
|
||||
assert_equal=assert_equal,
|
||||
check_verify_expected_token=check_verify_expected_token,
|
||||
)
|
||||
return cuda_log, ref_log
|
||||
|
||||
|
||||
def run_write_diff(
|
||||
*,
|
||||
buf_pair: tuple[torch.Tensor, torch.Tensor],
|
||||
plan_pair: tuple[WritePlan, WritePlan],
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
out_cache_loc: torch.Tensor,
|
||||
expected_input_tokens: torch.Tensor,
|
||||
expected_input_positions: torch.Tensor,
|
||||
enable_write_verify_inputs: bool = False,
|
||||
kernel_kind: CanaryLaunchTag = CanaryLaunchTag.HEAD_K_FULL,
|
||||
device: torch.device = _DEVICE,
|
||||
assert_equal: bool = True,
|
||||
) -> tuple[FakeViolationLog, FakeViolationLog]:
|
||||
"""Thin wrapper around ``_run_both_write`` that creates a fresh log pair and packs (cuda, ref)
|
||||
buf/plan arguments into 2-tuples to drop ~10 lines of boilerplate per call site.
|
||||
"""
|
||||
cuda_log, ref_log = make_log_pair(device=device)
|
||||
_run_both_write(
|
||||
cuda_canary_buf=buf_pair[0],
|
||||
ref_canary_buf=buf_pair[1],
|
||||
plan_cuda=plan_pair[0],
|
||||
plan_ref=plan_pair[1],
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
enable_write_verify_inputs=enable_write_verify_inputs,
|
||||
expected_input_tokens=expected_input_tokens,
|
||||
expected_input_positions=expected_input_positions,
|
||||
cuda_log=cuda_log,
|
||||
ref_log=ref_log,
|
||||
kernel_kind=kernel_kind,
|
||||
assert_equal=assert_equal,
|
||||
)
|
||||
return cuda_log, ref_log
|
||||
|
||||
|
||||
def run_plan_diff(
|
||||
*,
|
||||
plan_pair: tuple[tuple[VerifyPlan, WritePlan], tuple[VerifyPlan, WritePlan]],
|
||||
req_pool_indices: torch.Tensor,
|
||||
prefix_lens: torch.Tensor,
|
||||
extend_seq_lens: torch.Tensor,
|
||||
req_to_token: torch.Tensor,
|
||||
extras: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
|
||||
swa_window_size: int = 0,
|
||||
full_to_swa_index_mapping: Optional[torch.Tensor] = None,
|
||||
assert_equal: bool = True,
|
||||
active_verify_entries: Optional[int] = None,
|
||||
active_write_reqs: Optional[int] = None,
|
||||
req_to_verify_expected_tokens: Optional[torch.Tensor] = None,
|
||||
req_to_verify_expected_tokens_valid_lens: Optional[torch.Tensor] = None,
|
||||
kv_token_id_vs_position_offset: int = 0,
|
||||
) -> None:
|
||||
"""Thin wrapper around ``_run_both_plan`` that unpacks ``((triton_v, triton_w), (ref_v, ref_w))``
|
||||
plan pairs to drop the per-call-site ``triton_verify=.../triton_write=.../ref_verify=...`` block.
|
||||
"""
|
||||
(triton_verify, triton_write), (ref_verify, ref_write) = plan_pair
|
||||
_run_both_plan(
|
||||
triton_verify=triton_verify,
|
||||
triton_write=triton_write,
|
||||
ref_verify=ref_verify,
|
||||
ref_write=ref_write,
|
||||
req_pool_indices=req_pool_indices,
|
||||
prefix_lens=prefix_lens,
|
||||
extend_seq_lens=extend_seq_lens,
|
||||
req_to_token=req_to_token,
|
||||
extras=extras,
|
||||
swa_window_size=swa_window_size,
|
||||
full_to_swa_index_mapping=full_to_swa_index_mapping,
|
||||
assert_equal=assert_equal,
|
||||
active_verify_entries=active_verify_entries,
|
||||
active_write_reqs=active_write_reqs,
|
||||
req_to_verify_expected_tokens=req_to_verify_expected_tokens,
|
||||
req_to_verify_expected_tokens_valid_lens=req_to_verify_expected_tokens_valid_lens,
|
||||
kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
|
||||
)
|
||||
@@ -0,0 +1,45 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
from typing import Any, Callable
|
||||
|
||||
from sglang.jit_kernel.tests.kv_canary._differential import (
|
||||
ShrinkResult,
|
||||
shrink_inputs,
|
||||
)
|
||||
|
||||
FUZZ_SEEDS_PR: tuple[int, ...] = (0,)
|
||||
|
||||
|
||||
def check_repro(inputs: Any, *, run_one_fn: Callable[[Any], Any]) -> bool:
|
||||
try:
|
||||
run_one_fn(inputs)
|
||||
except (AssertionError, RuntimeError, ValueError):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def run_fuzz_combo(
|
||||
seed: int,
|
||||
*,
|
||||
draw_fn: Callable[[random.Random], Any],
|
||||
run_one_fn: Callable[[Any], Any],
|
||||
summarize_fn: Callable[[Any], str],
|
||||
n_iter: int,
|
||||
) -> None:
|
||||
rng = random.Random(seed)
|
||||
for iteration in range(n_iter):
|
||||
inputs = draw_fn(rng)
|
||||
try:
|
||||
run_one_fn(inputs)
|
||||
except AssertionError as exc:
|
||||
shrunk: ShrinkResult = shrink_inputs(
|
||||
inputs,
|
||||
check_fn=lambda i: check_repro(i, run_one_fn=run_one_fn),
|
||||
)
|
||||
raise AssertionError(
|
||||
f"seed={seed} iter={iteration} failure: {exc}\n"
|
||||
f"original: {summarize_fn(inputs)}\n"
|
||||
f"shrunk: {summarize_fn(shrunk.inputs)}\n"
|
||||
f"mutations applied: {shrunk.mutations_applied}"
|
||||
) from exc
|
||||
@@ -0,0 +1,425 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.kv_canary.verify import (
|
||||
CanaryLaunchTag,
|
||||
VerifyOrWriteContext,
|
||||
VerifyPlan,
|
||||
launch_canary_verify_kernel,
|
||||
)
|
||||
from sglang.jit_kernel.kv_canary.write import WritePlan
|
||||
from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
|
||||
FakeViolationLog,
|
||||
assert_canary_buf_equal,
|
||||
assert_canary_state_equal,
|
||||
make_canary_buf,
|
||||
make_canary_buf_pair,
|
||||
make_log_pair,
|
||||
make_verify_plan,
|
||||
make_verify_plan_pair,
|
||||
make_write_plan_pair,
|
||||
stamp_clean_chain,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._differential import (
|
||||
_assert_plans_byte_equal,
|
||||
_run_both_plan,
|
||||
_run_both_verify,
|
||||
_run_both_write,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._fixtures import (
|
||||
dummy_pseudo_tensors,
|
||||
empty_extras,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=60, suite="base-b-kernel-unit-1-gpu-large")
|
||||
|
||||
_DEVICE = torch.device("cuda")
|
||||
|
||||
|
||||
def _build_verify_plan_5_entries(
|
||||
*, device: torch.device
|
||||
) -> tuple[VerifyPlan, VerifyPlan]:
|
||||
num_slots = 16
|
||||
cuda_buf, ref_buf = make_canary_buf_pair(
|
||||
num_slots=num_slots, slot_stride_bytes=32, device=_DEVICE
|
||||
)
|
||||
|
||||
plan_cuda, plan_ref = make_verify_plan_pair(
|
||||
slot_indices=[0, 1, 2, 3, 4],
|
||||
positions=[0, 1, 2, 3, 4],
|
||||
prev_slot_indices=[-1, 0, 1, 2, 3],
|
||||
capacity=8,
|
||||
device=device,
|
||||
)
|
||||
return plan_cuda, plan_ref
|
||||
|
||||
|
||||
def _build_write_fixtures(
|
||||
*, device: torch.device
|
||||
) -> tuple[WritePlan, WritePlan, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
num_tokens = 5
|
||||
plan_cuda, plan_ref = make_write_plan_pair(
|
||||
write_offsets=[0, num_tokens],
|
||||
seed_slot_indices=[-1],
|
||||
num_valid_reqs=1,
|
||||
req_capacity=4,
|
||||
device=device,
|
||||
)
|
||||
input_ids = torch.tensor([10, 20, 30, 40, 50], dtype=torch.int64, device=device)
|
||||
positions = torch.tensor([0, 1, 2, 3, 4], dtype=torch.int64, device=device)
|
||||
out_cache_loc = torch.tensor([0, 1, 2, 3, 4], dtype=torch.int64, device=device)
|
||||
return plan_cuda, plan_ref, input_ids, positions, out_cache_loc
|
||||
|
||||
|
||||
def _build_plan_fixtures(
|
||||
*, device: torch.device, int64_req_to_token: bool = False
|
||||
) -> tuple[
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
]:
|
||||
bs = 3
|
||||
max_reqs = 4
|
||||
max_seq_len = 16
|
||||
req_pool_indices = torch.tensor([1, 2, 3], dtype=torch.int64, device=device)
|
||||
prefix_lens = torch.tensor([0, 4, 8], dtype=torch.int64, device=device)
|
||||
extend_seq_lens = torch.tensor([5, 1, 1], dtype=torch.int64, device=device)
|
||||
rp_axis = torch.arange(max_reqs, device=device, dtype=torch.int32).unsqueeze(1)
|
||||
pos_axis = torch.arange(max_seq_len, device=device, dtype=torch.int32).unsqueeze(0)
|
||||
req_to_token_int32 = (rp_axis * max_seq_len + pos_axis).contiguous()
|
||||
if int64_req_to_token:
|
||||
req_to_token = req_to_token_int32.to(torch.int64)
|
||||
else:
|
||||
req_to_token = req_to_token_int32
|
||||
return req_pool_indices, prefix_lens, extend_seq_lens, req_to_token
|
||||
|
||||
|
||||
def test_verify_byte_equal_across_repeated_launches_10x() -> None:
|
||||
num_launches = 10
|
||||
plan_cuda, plan_ref = _build_verify_plan_5_entries(device=_DEVICE)
|
||||
|
||||
snapshot_rings: list[torch.Tensor] = []
|
||||
snapshot_write_indices: list[torch.Tensor] = []
|
||||
snapshot_bufs: list[torch.Tensor] = []
|
||||
|
||||
for _ in range(num_launches):
|
||||
cuda_buf, ref_buf = make_canary_buf_pair(
|
||||
num_slots=16, slot_stride_bytes=32, device=_DEVICE
|
||||
)
|
||||
cuda_log, ref_log = make_log_pair(capacity=64, device=_DEVICE)
|
||||
|
||||
_run_both_verify(
|
||||
cuda_canary_buf=cuda_buf,
|
||||
ref_canary_buf=ref_buf,
|
||||
plan_cuda=plan_cuda,
|
||||
plan_ref=plan_ref,
|
||||
cuda_log=cuda_log,
|
||||
ref_log=ref_log,
|
||||
kernel_kind=CanaryLaunchTag.HEAD_K_FULL,
|
||||
)
|
||||
|
||||
assert_canary_buf_equal(buf_a=cuda_buf, buf_b=ref_buf)
|
||||
assert_canary_state_equal(log_a=cuda_log, log_b=ref_log)
|
||||
|
||||
snapshot_rings.append(cuda_log.ring.clone())
|
||||
snapshot_write_indices.append(cuda_log.write_index.clone())
|
||||
snapshot_bufs.append(cuda_buf.clone())
|
||||
|
||||
for i in range(1, num_launches):
|
||||
assert torch.equal(
|
||||
snapshot_rings[0], snapshot_rings[i]
|
||||
), f"violation_ring differs between launch 0 and {i}"
|
||||
assert torch.equal(
|
||||
snapshot_write_indices[0], snapshot_write_indices[i]
|
||||
), f"violation_write_index differs between launch 0 and {i}"
|
||||
assert torch.equal(
|
||||
snapshot_bufs[0], snapshot_bufs[i]
|
||||
), f"canary_buf differs between launch 0 and {i}"
|
||||
|
||||
|
||||
def test_write_byte_equal_across_repeated_launches_10x() -> None:
|
||||
num_launches = 10
|
||||
plan_cuda, plan_ref, input_ids, positions, out_cache_loc = _build_write_fixtures(
|
||||
device=_DEVICE
|
||||
)
|
||||
|
||||
snapshot_bufs: list[torch.Tensor] = []
|
||||
snapshot_rings: list[torch.Tensor] = []
|
||||
snapshot_counters: list[torch.Tensor] = []
|
||||
|
||||
for _ in range(num_launches):
|
||||
cuda_buf, ref_buf = make_canary_buf_pair(
|
||||
num_slots=16, slot_stride_bytes=32, device=_DEVICE
|
||||
)
|
||||
cuda_log, ref_log = make_log_pair(capacity=64, device=_DEVICE)
|
||||
pseudo_tok, pseudo_pos = dummy_pseudo_tensors(input_ids.shape[0])
|
||||
|
||||
_run_both_write(
|
||||
cuda_canary_buf=cuda_buf,
|
||||
ref_canary_buf=ref_buf,
|
||||
plan_cuda=plan_cuda,
|
||||
plan_ref=plan_ref,
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
enable_write_verify_inputs=False,
|
||||
expected_input_tokens=pseudo_tok,
|
||||
expected_input_positions=pseudo_pos,
|
||||
cuda_log=cuda_log,
|
||||
ref_log=ref_log,
|
||||
kernel_kind=CanaryLaunchTag.HEAD_K_FULL,
|
||||
)
|
||||
|
||||
assert_canary_buf_equal(buf_a=cuda_buf, buf_b=ref_buf)
|
||||
assert_canary_state_equal(log_a=cuda_log, log_b=ref_log)
|
||||
|
||||
snapshot_bufs.append(cuda_buf.clone())
|
||||
snapshot_rings.append(cuda_log.ring.clone())
|
||||
snapshot_counters.append(cuda_log.slot_run_counter.clone())
|
||||
|
||||
for i in range(1, num_launches):
|
||||
assert torch.equal(
|
||||
snapshot_bufs[0], snapshot_bufs[i]
|
||||
), f"canary_buf differs between launch 0 and {i}"
|
||||
assert torch.equal(
|
||||
snapshot_rings[0], snapshot_rings[i]
|
||||
), f"violation_ring differs between launch 0 and {i}"
|
||||
assert torch.equal(
|
||||
snapshot_counters[0], snapshot_counters[i]
|
||||
), f"slot_run_counter differs between launch 0 and {i}"
|
||||
|
||||
|
||||
def test_plan_byte_equal_across_repeated_launches_10x() -> None:
|
||||
num_launches = 10
|
||||
req_pool_indices, prefix_lens, extend_seq_lens, req_to_token = _build_plan_fixtures(
|
||||
device=_DEVICE
|
||||
)
|
||||
|
||||
snapshot_slots: list[torch.Tensor] = []
|
||||
snapshot_positions: list[torch.Tensor] = []
|
||||
snapshot_prevs: list[torch.Tensor] = []
|
||||
snapshot_write_offsets: list[torch.Tensor] = []
|
||||
|
||||
for _ in range(num_launches):
|
||||
triton_v = VerifyPlan.allocate(
|
||||
verify_capacity=64, device=_DEVICE
|
||||
).zero_for_testing_()
|
||||
triton_w = WritePlan.allocate(
|
||||
write_req_capacity=8, device=_DEVICE
|
||||
).zero_for_testing_()
|
||||
ref_v = VerifyPlan.allocate(
|
||||
verify_capacity=64, device=_DEVICE
|
||||
).zero_for_testing_()
|
||||
ref_w = WritePlan.allocate(
|
||||
write_req_capacity=8, device=_DEVICE
|
||||
).zero_for_testing_()
|
||||
|
||||
_run_both_plan(
|
||||
triton_verify=triton_v,
|
||||
triton_write=triton_w,
|
||||
ref_verify=ref_v,
|
||||
ref_write=ref_w,
|
||||
req_pool_indices=req_pool_indices,
|
||||
prefix_lens=prefix_lens,
|
||||
extend_seq_lens=extend_seq_lens,
|
||||
req_to_token=req_to_token,
|
||||
extras=empty_extras(),
|
||||
swa_window_size=0,
|
||||
full_to_swa_index_mapping=None,
|
||||
)
|
||||
|
||||
_assert_plans_byte_equal(
|
||||
triton_verify=triton_v,
|
||||
triton_write=triton_w,
|
||||
ref_verify=ref_v,
|
||||
ref_write=ref_w,
|
||||
)
|
||||
|
||||
n_verify = int(triton_v.verify_num_valid[0].item())
|
||||
snapshot_slots.append(triton_v.verify_slot_indices[:n_verify].clone())
|
||||
snapshot_positions.append(triton_v.verify_expected_positions[:n_verify].clone())
|
||||
snapshot_prevs.append(triton_v.verify_prev_slot_indices[:n_verify].clone())
|
||||
snapshot_write_offsets.append(triton_w.write_offsets.clone())
|
||||
|
||||
for i in range(1, num_launches):
|
||||
assert torch.equal(
|
||||
snapshot_slots[0], snapshot_slots[i]
|
||||
), f"verify_slot_indices differs between launch 0 and {i}"
|
||||
assert torch.equal(
|
||||
snapshot_positions[0], snapshot_positions[i]
|
||||
), f"verify_expected_positions differs between launch 0 and {i}"
|
||||
assert torch.equal(
|
||||
snapshot_prevs[0], snapshot_prevs[i]
|
||||
), f"verify_prev_slot_indices differs between launch 0 and {i}"
|
||||
assert torch.equal(
|
||||
snapshot_write_offsets[0], snapshot_write_offsets[i]
|
||||
), f"write_offsets differs between launch 0 and {i}"
|
||||
|
||||
|
||||
def test_verify_multi_launch_100x_counter_linear() -> None:
|
||||
num_launches = 100
|
||||
|
||||
plan_cuda = make_verify_plan(
|
||||
slot_indices=[0],
|
||||
positions=[0],
|
||||
prev_slot_indices=[-1],
|
||||
capacity=4,
|
||||
device=_DEVICE,
|
||||
)
|
||||
|
||||
cuda_log = FakeViolationLog.allocate(capacity=64, device=_DEVICE)
|
||||
|
||||
for _ in range(num_launches):
|
||||
cuda_buf = make_canary_buf(num_slots=16, slot_stride_bytes=32, device=_DEVICE)
|
||||
launch_canary_verify_kernel(
|
||||
context=VerifyOrWriteContext(
|
||||
canary_buf=cuda_buf,
|
||||
kernel_kind=CanaryLaunchTag.HEAD_K_FULL,
|
||||
violation_ring=cuda_log.ring,
|
||||
violation_write_index=cuda_log.write_index,
|
||||
slot_run_counter=cuda_log.slot_run_counter,
|
||||
kernel_run_counter=cuda_log.kernel_run_counter,
|
||||
enable_chain_position_assert=cuda_log.enable_chain_position_assert,
|
||||
),
|
||||
plan=plan_cuda,
|
||||
check_verify_expected_token=True,
|
||||
)
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
assert (
|
||||
int(cuda_log.kernel_run_counter[0].item()) == num_launches
|
||||
), f"kernel_run_counter expected {num_launches}, got {cuda_log.kernel_run_counter[0].item()}"
|
||||
assert int(cuda_log.slot_run_counter[0].item()) == num_launches, (
|
||||
f"slot_run_counter expected {num_launches} (1 active entry x 100 launches), "
|
||||
f"got {cuda_log.slot_run_counter[0].item()}"
|
||||
)
|
||||
|
||||
|
||||
def test_verify_check_disabled_byte_equal() -> None:
|
||||
"""check_verify_expected_token True vs False produce equivalent violation logs on a clean plan."""
|
||||
plan_true_cuda, plan_true_ref = _build_verify_plan_5_entries(device=_DEVICE)
|
||||
plan_false_cuda, plan_false_ref = _build_verify_plan_5_entries(device=_DEVICE)
|
||||
|
||||
chain_slot_indices = [0, 1, 2, 3, 4]
|
||||
chain_tokens = [10, 20, 30, 40, 50]
|
||||
chain_positions = [0, 1, 2, 3, 4]
|
||||
|
||||
cuda_buf_true, ref_buf_true = make_canary_buf_pair(
|
||||
num_slots=16, slot_stride_bytes=32, device=_DEVICE
|
||||
)
|
||||
stamp_clean_chain(
|
||||
cuda_buf=cuda_buf_true,
|
||||
ref_buf=ref_buf_true,
|
||||
slot_indices=chain_slot_indices,
|
||||
tokens=chain_tokens,
|
||||
positions=chain_positions,
|
||||
)
|
||||
cuda_buf_false, ref_buf_false = make_canary_buf_pair(
|
||||
num_slots=16, slot_stride_bytes=32, device=_DEVICE
|
||||
)
|
||||
stamp_clean_chain(
|
||||
cuda_buf=cuda_buf_false,
|
||||
ref_buf=ref_buf_false,
|
||||
slot_indices=chain_slot_indices,
|
||||
tokens=chain_tokens,
|
||||
positions=chain_positions,
|
||||
)
|
||||
cuda_log_true, ref_log_true = make_log_pair(capacity=64, device=_DEVICE)
|
||||
cuda_log_false, ref_log_false = make_log_pair(capacity=64, device=_DEVICE)
|
||||
|
||||
_run_both_verify(
|
||||
cuda_canary_buf=cuda_buf_true,
|
||||
ref_canary_buf=ref_buf_true,
|
||||
plan_cuda=plan_true_cuda,
|
||||
plan_ref=plan_true_ref,
|
||||
cuda_log=cuda_log_true,
|
||||
ref_log=ref_log_true,
|
||||
kernel_kind=CanaryLaunchTag.HEAD_K_FULL,
|
||||
check_verify_expected_token=True,
|
||||
)
|
||||
_run_both_verify(
|
||||
cuda_canary_buf=cuda_buf_false,
|
||||
ref_canary_buf=ref_buf_false,
|
||||
plan_cuda=plan_false_cuda,
|
||||
plan_ref=plan_false_ref,
|
||||
cuda_log=cuda_log_false,
|
||||
ref_log=ref_log_false,
|
||||
kernel_kind=CanaryLaunchTag.HEAD_K_FULL,
|
||||
check_verify_expected_token=False,
|
||||
)
|
||||
|
||||
assert int(cuda_log_true.write_index[0].item()) == 0
|
||||
assert int(cuda_log_false.write_index[0].item()) == 0
|
||||
assert torch.equal(cuda_log_true.ring, cuda_log_false.ring)
|
||||
assert torch.equal(cuda_log_true.write_index, cuda_log_false.write_index)
|
||||
assert torch.equal(cuda_log_true.slot_run_counter, cuda_log_false.slot_run_counter)
|
||||
assert torch.equal(
|
||||
cuda_log_true.kernel_run_counter, cuda_log_false.kernel_run_counter
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("per_req_present", [False, True])
|
||||
def test_plan_per_req_present_or_absent(per_req_present: bool) -> None:
|
||||
max_reqs = 4
|
||||
max_seq_len = 16
|
||||
rp_axis = torch.arange(max_reqs, device=_DEVICE, dtype=torch.int32).unsqueeze(1)
|
||||
pos_axis = torch.arange(max_seq_len, device=_DEVICE, dtype=torch.int32).unsqueeze(0)
|
||||
req_to_token = (rp_axis * max_seq_len + pos_axis).contiguous()
|
||||
|
||||
if per_req_present:
|
||||
req_pool_indices = torch.tensor([1, 2], dtype=torch.int64, device=_DEVICE)
|
||||
prefix_lens = torch.tensor([3, 5], dtype=torch.int64, device=_DEVICE)
|
||||
extend_seq_lens = torch.tensor([1, 1], dtype=torch.int64, device=_DEVICE)
|
||||
else:
|
||||
req_pool_indices = torch.tensor([0], dtype=torch.int64, device=_DEVICE)
|
||||
prefix_lens = torch.tensor([0], dtype=torch.int64, device=_DEVICE)
|
||||
extend_seq_lens = torch.tensor([0], dtype=torch.int64, device=_DEVICE)
|
||||
|
||||
triton_v = VerifyPlan.allocate(
|
||||
verify_capacity=64, device=_DEVICE
|
||||
).zero_for_testing_()
|
||||
triton_w = WritePlan.allocate(
|
||||
write_req_capacity=8, device=_DEVICE
|
||||
).zero_for_testing_()
|
||||
ref_v = VerifyPlan.allocate(verify_capacity=64, device=_DEVICE).zero_for_testing_()
|
||||
ref_w = WritePlan.allocate(write_req_capacity=8, device=_DEVICE).zero_for_testing_()
|
||||
|
||||
_run_both_plan(
|
||||
triton_verify=triton_v,
|
||||
triton_write=triton_w,
|
||||
ref_verify=ref_v,
|
||||
ref_write=ref_w,
|
||||
req_pool_indices=req_pool_indices,
|
||||
prefix_lens=prefix_lens,
|
||||
extend_seq_lens=extend_seq_lens,
|
||||
req_to_token=req_to_token,
|
||||
extras=empty_extras(),
|
||||
swa_window_size=0,
|
||||
full_to_swa_index_mapping=None,
|
||||
)
|
||||
|
||||
_assert_plans_byte_equal(
|
||||
triton_verify=triton_v,
|
||||
triton_write=triton_w,
|
||||
ref_verify=ref_v,
|
||||
ref_write=ref_w,
|
||||
)
|
||||
|
||||
if not per_req_present:
|
||||
assert int(triton_v.verify_num_valid[0].item()) == 0
|
||||
bs = int(req_pool_indices.shape[0])
|
||||
assert int(triton_w.write_offsets[bs].item()) == 0
|
||||
|
||||
if per_req_present:
|
||||
assert int(triton_v.verify_num_valid[0].item()) == 8
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
@@ -0,0 +1,790 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Optional
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.kv_canary import consts
|
||||
from sglang.jit_kernel.kv_canary.plan import launch_canary_plan_kernels
|
||||
from sglang.jit_kernel.kv_canary.plan_ref import (
|
||||
launch_canary_plan_kernels_torch_reference,
|
||||
)
|
||||
from sglang.jit_kernel.kv_canary.verify import (
|
||||
CanaryLaunchTag,
|
||||
VerifyOrWriteContext,
|
||||
VerifyPlan,
|
||||
launch_canary_verify_kernel,
|
||||
)
|
||||
from sglang.jit_kernel.kv_canary.verify_ref import (
|
||||
launch_canary_verify_kernel_torch_reference,
|
||||
)
|
||||
from sglang.jit_kernel.kv_canary.write import WritePlan, launch_canary_write_kernel
|
||||
from sglang.jit_kernel.kv_canary.write_ref import (
|
||||
launch_canary_write_kernel_torch_reference,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
|
||||
FakeViolationLog,
|
||||
assert_canary_buf_equal,
|
||||
assert_canary_state_equal,
|
||||
make_canary_buf,
|
||||
stamp_clean_chain,
|
||||
write_slot_fields,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._fixtures import (
|
||||
empty_extras,
|
||||
make_req_to_token,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
|
||||
|
||||
|
||||
_DEVICE = torch.device("cuda")
|
||||
|
||||
|
||||
def _run_pipeline(
|
||||
*,
|
||||
real: bool,
|
||||
req_pool_indices: torch.Tensor,
|
||||
prefix_lens: torch.Tensor,
|
||||
extend_seq_lens: torch.Tensor,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
out_cache_loc: torch.Tensor,
|
||||
req_to_token: torch.Tensor,
|
||||
canary_buf: torch.Tensor,
|
||||
log: FakeViolationLog,
|
||||
extras: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
|
||||
swa_window_size: int,
|
||||
full_to_swa_index_mapping: Optional[torch.Tensor],
|
||||
kernel_kind: CanaryLaunchTag,
|
||||
enable_write_verify_inputs: bool,
|
||||
expected_input_tokens: torch.Tensor,
|
||||
expected_input_positions: torch.Tensor,
|
||||
verify_capacity: int,
|
||||
write_req_capacity: int,
|
||||
req_to_verify_expected_tokens: Optional[torch.Tensor] = None,
|
||||
req_to_verify_expected_tokens_valid_lens: Optional[torch.Tensor] = None,
|
||||
kv_token_id_vs_position_offset: int = 0,
|
||||
check_verify_expected_token: bool = True,
|
||||
) -> tuple[VerifyPlan, WritePlan]:
|
||||
_ = extras
|
||||
plan_v = VerifyPlan.allocate(verify_capacity=verify_capacity, device=_DEVICE)
|
||||
plan_w = WritePlan.allocate(write_req_capacity=write_req_capacity, device=_DEVICE)
|
||||
|
||||
# Existing pipeline tests that supply a pool but no per-req lens want "bound by full
|
||||
# row width" semantics. Synthesise that bound here so callers don't have to.
|
||||
if (
|
||||
req_to_verify_expected_tokens is not None
|
||||
and req_to_verify_expected_tokens_valid_lens is None
|
||||
):
|
||||
req_to_verify_expected_tokens_valid_lens = torch.full(
|
||||
(int(req_pool_indices.shape[0]),),
|
||||
int(req_to_verify_expected_tokens.shape[1]),
|
||||
dtype=torch.int64,
|
||||
device=req_pool_indices.device,
|
||||
)
|
||||
|
||||
plan_fn = (
|
||||
launch_canary_plan_kernels
|
||||
if real
|
||||
else launch_canary_plan_kernels_torch_reference
|
||||
)
|
||||
plan_fn(
|
||||
verify_plan_out=plan_v,
|
||||
write_plan_out=plan_w,
|
||||
req_pool_indices=req_pool_indices,
|
||||
prefix_lens=prefix_lens,
|
||||
extend_seq_lens=extend_seq_lens,
|
||||
req_to_token=req_to_token,
|
||||
swa_window_size=swa_window_size,
|
||||
full_to_swa_index_mapping=full_to_swa_index_mapping,
|
||||
verify_capacity=verify_capacity,
|
||||
req_to_verify_expected_tokens=req_to_verify_expected_tokens,
|
||||
req_to_verify_expected_tokens_valid_lens=req_to_verify_expected_tokens_valid_lens,
|
||||
kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
|
||||
)
|
||||
|
||||
if real:
|
||||
context = VerifyOrWriteContext(
|
||||
canary_buf=canary_buf,
|
||||
kernel_kind=kernel_kind,
|
||||
violation_ring=log.ring,
|
||||
violation_write_index=log.write_index,
|
||||
slot_run_counter=log.slot_run_counter,
|
||||
kernel_run_counter=log.kernel_run_counter,
|
||||
enable_chain_position_assert=log.enable_chain_position_assert,
|
||||
)
|
||||
launch_canary_write_kernel(
|
||||
context=context,
|
||||
plan=plan_w,
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
enable_write_input_assert=enable_write_verify_inputs,
|
||||
expected_input_tokens=expected_input_tokens,
|
||||
expected_input_positions=expected_input_positions,
|
||||
)
|
||||
launch_canary_verify_kernel(
|
||||
context=context,
|
||||
plan=plan_v,
|
||||
check_verify_expected_token=check_verify_expected_token,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
else:
|
||||
launch_canary_write_kernel_torch_reference(
|
||||
context=VerifyOrWriteContext(
|
||||
canary_buf=canary_buf,
|
||||
kernel_kind=kernel_kind,
|
||||
violation_ring=log.ring,
|
||||
violation_write_index=log.write_index,
|
||||
slot_run_counter=log.slot_run_counter,
|
||||
kernel_run_counter=log.kernel_run_counter,
|
||||
enable_chain_position_assert=log.enable_chain_position_assert,
|
||||
),
|
||||
plan=plan_w,
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
enable_write_input_assert=enable_write_verify_inputs,
|
||||
expected_input_tokens=expected_input_tokens,
|
||||
expected_input_positions=expected_input_positions,
|
||||
)
|
||||
launch_canary_verify_kernel_torch_reference(
|
||||
context=VerifyOrWriteContext(
|
||||
canary_buf=canary_buf,
|
||||
kernel_kind=kernel_kind,
|
||||
violation_ring=log.ring,
|
||||
violation_write_index=log.write_index,
|
||||
slot_run_counter=log.slot_run_counter,
|
||||
kernel_run_counter=log.kernel_run_counter,
|
||||
enable_chain_position_assert=log.enable_chain_position_assert,
|
||||
),
|
||||
plan=plan_v,
|
||||
check_verify_expected_token=check_verify_expected_token,
|
||||
)
|
||||
|
||||
return plan_v, plan_w
|
||||
|
||||
|
||||
def _run_both_and_assert_pipeline_equal(
|
||||
*,
|
||||
req_pool_indices: torch.Tensor,
|
||||
prefix_lens: torch.Tensor,
|
||||
extend_seq_lens: torch.Tensor,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
out_cache_loc: torch.Tensor,
|
||||
req_to_token: torch.Tensor,
|
||||
num_slots: int,
|
||||
extras: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
|
||||
swa_window_size: int = 0,
|
||||
full_to_swa_index_mapping: Optional[torch.Tensor] = None,
|
||||
kernel_kind: CanaryLaunchTag = CanaryLaunchTag.HEAD_K_FULL,
|
||||
enable_write_verify_inputs: bool = False,
|
||||
expected_input_tokens: Optional[torch.Tensor] = None,
|
||||
expected_input_positions: Optional[torch.Tensor] = None,
|
||||
ring_capacity: int = 64,
|
||||
verify_capacity: int = 256,
|
||||
write_req_capacity: int = 16,
|
||||
assert_ring_equal: bool = True,
|
||||
initial_canary_buf: Optional[torch.Tensor] = None,
|
||||
req_to_verify_expected_tokens: Optional[torch.Tensor] = None,
|
||||
kv_token_id_vs_position_offset: int = 0,
|
||||
check_verify_expected_token: bool = True,
|
||||
) -> tuple[
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
FakeViolationLog,
|
||||
FakeViolationLog,
|
||||
VerifyPlan,
|
||||
WritePlan,
|
||||
VerifyPlan,
|
||||
WritePlan,
|
||||
]:
|
||||
# The kernel rejects non-None expected_* tensors when enable_write_verify_inputs=False
|
||||
# (sanity check to catch caller bugs), so only synthesise zero placeholders in the
|
||||
# branch that will actually assert against them.
|
||||
if enable_write_verify_inputs:
|
||||
total_tokens = int(input_ids.shape[0])
|
||||
if expected_input_tokens is None:
|
||||
expected_input_tokens = torch.zeros(
|
||||
total_tokens, dtype=torch.int64, device=_DEVICE
|
||||
)
|
||||
if expected_input_positions is None:
|
||||
expected_input_positions = torch.zeros(
|
||||
total_tokens, dtype=torch.int64, device=_DEVICE
|
||||
)
|
||||
|
||||
if initial_canary_buf is None:
|
||||
buf_real = make_canary_buf(num_slots=num_slots, device=_DEVICE)
|
||||
else:
|
||||
buf_real = initial_canary_buf.clone()
|
||||
buf_ref = buf_real.clone()
|
||||
log_real = FakeViolationLog.allocate(capacity=ring_capacity, device=_DEVICE)
|
||||
log_ref = FakeViolationLog.allocate(capacity=ring_capacity, device=_DEVICE)
|
||||
|
||||
shared: dict[str, Any] = dict(
|
||||
req_pool_indices=req_pool_indices,
|
||||
prefix_lens=prefix_lens,
|
||||
extend_seq_lens=extend_seq_lens,
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
req_to_token=req_to_token,
|
||||
extras=extras,
|
||||
swa_window_size=swa_window_size,
|
||||
full_to_swa_index_mapping=full_to_swa_index_mapping,
|
||||
kernel_kind=kernel_kind,
|
||||
enable_write_verify_inputs=enable_write_verify_inputs,
|
||||
expected_input_tokens=expected_input_tokens,
|
||||
expected_input_positions=expected_input_positions,
|
||||
verify_capacity=verify_capacity,
|
||||
write_req_capacity=write_req_capacity,
|
||||
req_to_verify_expected_tokens=req_to_verify_expected_tokens,
|
||||
kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
|
||||
check_verify_expected_token=check_verify_expected_token,
|
||||
)
|
||||
|
||||
plan_v_real, plan_w_real = _run_pipeline(
|
||||
real=True,
|
||||
canary_buf=buf_real,
|
||||
log=log_real,
|
||||
**shared,
|
||||
)
|
||||
plan_v_ref, plan_w_ref = _run_pipeline(
|
||||
real=False,
|
||||
canary_buf=buf_ref,
|
||||
log=log_ref,
|
||||
**shared,
|
||||
)
|
||||
|
||||
assert_canary_buf_equal(buf_a=buf_real, buf_b=buf_ref)
|
||||
if assert_ring_equal:
|
||||
assert_canary_state_equal(log_a=log_real, log_b=log_ref)
|
||||
else:
|
||||
assert torch.equal(log_real.write_index, log_ref.write_index)
|
||||
assert torch.equal(log_real.slot_run_counter, log_ref.slot_run_counter)
|
||||
assert torch.equal(log_real.kernel_run_counter, log_ref.kernel_run_counter)
|
||||
|
||||
return (
|
||||
buf_real,
|
||||
buf_ref,
|
||||
log_real,
|
||||
log_ref,
|
||||
plan_v_real,
|
||||
plan_w_real,
|
||||
plan_v_ref,
|
||||
plan_w_ref,
|
||||
)
|
||||
|
||||
|
||||
def _t(values: list[int]) -> torch.Tensor:
|
||||
return torch.tensor(values, dtype=torch.int64, device=_DEVICE)
|
||||
|
||||
|
||||
def _linear_r2t(*, max_reqs: int = 4, max_seq_len: int = 16) -> torch.Tensor:
|
||||
return make_req_to_token(
|
||||
kind="linear", max_reqs=max_reqs, max_seq_len=max_seq_len, device=_DEVICE
|
||||
)
|
||||
|
||||
|
||||
def _zero_no_write_inputs() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""``(input_ids, positions, out_cache_loc)`` zero placeholders for extend_seq_lens=0 tests."""
|
||||
zeros = torch.zeros(1, dtype=torch.int64, device=_DEVICE)
|
||||
return zeros.clone(), zeros.clone(), zeros.clone()
|
||||
|
||||
|
||||
def _contiguous_out_cache_loc(
|
||||
*, req_pool_idx: int, start: int, count: int, max_seq_len: int = 16
|
||||
) -> torch.Tensor:
|
||||
return _t([req_pool_idx * max_seq_len + start + i for i in range(count)])
|
||||
|
||||
|
||||
def _stamp_linear_prefix(
|
||||
*,
|
||||
initial_buf: torch.Tensor,
|
||||
initial_ref: torch.Tensor,
|
||||
req_pool_idx: int,
|
||||
prefix_len: int,
|
||||
tokens: list[int],
|
||||
max_seq_len: int = 16,
|
||||
) -> None:
|
||||
"""Stamp clean chain for slots ``[rp*max_seq_len + 0 .. + prefix_len)`` at positions ``0..prefix_len``."""
|
||||
stamp_clean_chain(
|
||||
cuda_buf=initial_buf,
|
||||
ref_buf=initial_ref,
|
||||
slot_indices=[req_pool_idx * max_seq_len + pos for pos in range(prefix_len)],
|
||||
tokens=tokens,
|
||||
positions=list(range(prefix_len)),
|
||||
)
|
||||
|
||||
|
||||
def test_pipeline_basic_5_step_single_req() -> None:
|
||||
"""Single req, prefix_len=0, extend_seq_len=5: basic plan→write→verify byte-equal."""
|
||||
_run_both_and_assert_pipeline_equal(
|
||||
req_pool_indices=_t([1]),
|
||||
prefix_lens=_t([0]),
|
||||
extend_seq_lens=_t([5]),
|
||||
input_ids=_t([10, 20, 30, 40, 50]),
|
||||
positions=_t([0, 1, 2, 3, 4]),
|
||||
out_cache_loc=_contiguous_out_cache_loc(req_pool_idx=1, start=0, count=5),
|
||||
req_to_token=_linear_r2t(),
|
||||
num_slots=64,
|
||||
extras=empty_extras(),
|
||||
swa_window_size=0,
|
||||
full_to_swa_index_mapping=None,
|
||||
)
|
||||
|
||||
|
||||
def test_pipeline_multi_req_mixed_extend_decode() -> None:
|
||||
"""bs=3: pure extend req, decode req (prefix+1 extend), and padding sentinel row."""
|
||||
max_seq_len = 16
|
||||
_run_both_and_assert_pipeline_equal(
|
||||
req_pool_indices=_t([1, 2, 0]),
|
||||
prefix_lens=_t([0, 5, 0]),
|
||||
extend_seq_lens=_t([4, 1, 0]),
|
||||
input_ids=_t([11, 12, 13, 14, 21]),
|
||||
positions=_t([0, 1, 2, 3, 5]),
|
||||
out_cache_loc=_t(
|
||||
[
|
||||
1 * max_seq_len + 0,
|
||||
1 * max_seq_len + 1,
|
||||
1 * max_seq_len + 2,
|
||||
1 * max_seq_len + 3,
|
||||
2 * max_seq_len + 5,
|
||||
]
|
||||
),
|
||||
req_to_token=_linear_r2t(max_reqs=8, max_seq_len=max_seq_len),
|
||||
num_slots=128,
|
||||
extras=empty_extras(),
|
||||
swa_window_size=0,
|
||||
full_to_swa_index_mapping=None,
|
||||
write_req_capacity=4,
|
||||
)
|
||||
|
||||
|
||||
def test_pipeline_swa_window() -> None:
|
||||
"""SWA window=4, prefix_len=6: verify covers window [2,6), write covers extend tokens."""
|
||||
max_seq_len = 16
|
||||
max_reqs = 4
|
||||
full_to_swa_index_mapping = torch.arange(
|
||||
max_reqs * max_seq_len + 1, dtype=torch.int64, device=_DEVICE
|
||||
)
|
||||
|
||||
_run_both_and_assert_pipeline_equal(
|
||||
req_pool_indices=_t([1]),
|
||||
prefix_lens=_t([6]),
|
||||
extend_seq_lens=_t([2]),
|
||||
input_ids=_t([100, 101]),
|
||||
positions=_t([6, 7]),
|
||||
out_cache_loc=_t(
|
||||
[
|
||||
full_to_swa_index_mapping[1 * max_seq_len + 6].item(),
|
||||
full_to_swa_index_mapping[1 * max_seq_len + 7].item(),
|
||||
]
|
||||
),
|
||||
req_to_token=_linear_r2t(max_reqs=max_reqs, max_seq_len=max_seq_len),
|
||||
num_slots=64,
|
||||
extras=empty_extras(),
|
||||
swa_window_size=4,
|
||||
full_to_swa_index_mapping=full_to_swa_index_mapping,
|
||||
)
|
||||
|
||||
|
||||
def test_pipeline_sweep_no_write() -> None:
|
||||
"""All extend_seq_lens=0: write_step is no-op, verify sweeps prefix, buf unchanged."""
|
||||
prefix_len = 4
|
||||
input_ids, positions, out_cache_loc = _zero_no_write_inputs()
|
||||
|
||||
initial_buf = make_canary_buf(num_slots=64, device=_DEVICE)
|
||||
initial_ref = initial_buf.clone()
|
||||
_stamp_linear_prefix(
|
||||
initial_buf=initial_buf,
|
||||
initial_ref=initial_ref,
|
||||
req_pool_idx=1,
|
||||
prefix_len=prefix_len,
|
||||
tokens=[100 + pos for pos in range(prefix_len)],
|
||||
)
|
||||
|
||||
buf_real, buf_ref, log_real, log_ref, plan_v_real, plan_w_real, _, _ = (
|
||||
_run_both_and_assert_pipeline_equal(
|
||||
req_pool_indices=_t([1]),
|
||||
prefix_lens=_t([prefix_len]),
|
||||
extend_seq_lens=_t([0]),
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
req_to_token=_linear_r2t(),
|
||||
num_slots=64,
|
||||
extras=empty_extras(),
|
||||
swa_window_size=0,
|
||||
full_to_swa_index_mapping=None,
|
||||
initial_canary_buf=initial_buf,
|
||||
)
|
||||
)
|
||||
|
||||
assert int(plan_v_real.verify_num_valid[0].item()) == prefix_len
|
||||
assert int(plan_w_real.write_num_valid_reqs[0].item()) == 1
|
||||
assert int(plan_w_real.write_offsets[1].item()) == 0
|
||||
assert torch.equal(buf_real, initial_buf)
|
||||
assert torch.equal(buf_ref, initial_buf)
|
||||
assert int(log_real.write_index[0].item()) == 0
|
||||
assert int(log_ref.write_index[0].item()) == 0
|
||||
assert int(log_real.slot_run_counter[0].item()) == prefix_len
|
||||
assert int(log_ref.slot_run_counter[0].item()) == prefix_len
|
||||
|
||||
|
||||
def test_pipeline_pseudo_mode_on_match() -> None:
|
||||
"""enable_write_verify_inputs=ON, expected==actual: zero violations, buf byte-equal."""
|
||||
input_ids = _t([1, 2, 3, 4])
|
||||
positions = _t([0, 1, 2, 3])
|
||||
|
||||
_, _, log_real, log_ref, _, _, _, _ = _run_both_and_assert_pipeline_equal(
|
||||
req_pool_indices=_t([1]),
|
||||
prefix_lens=_t([0]),
|
||||
extend_seq_lens=_t([4]),
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=_contiguous_out_cache_loc(req_pool_idx=1, start=0, count=4),
|
||||
req_to_token=_linear_r2t(),
|
||||
num_slots=64,
|
||||
extras=empty_extras(),
|
||||
enable_write_verify_inputs=True,
|
||||
expected_input_tokens=input_ids.clone(),
|
||||
expected_input_positions=positions.clone(),
|
||||
)
|
||||
|
||||
assert int(log_real.write_index[0].item()) == 0
|
||||
assert int(log_ref.write_index[0].item()) == 0
|
||||
|
||||
|
||||
def test_pipeline_pseudo_mode_on_token_mismatch_then_verify_clean() -> None:
|
||||
"""enable_write_verify_inputs=ON, expected tokens all wrong: write records N violations."""
|
||||
n_tokens = 3
|
||||
positions = _t([0, 1, 2])
|
||||
|
||||
_, _, log_real, log_ref, _, _, _, _ = _run_both_and_assert_pipeline_equal(
|
||||
req_pool_indices=_t([1]),
|
||||
prefix_lens=_t([0]),
|
||||
extend_seq_lens=_t([3]),
|
||||
input_ids=_t([10, 20, 30]),
|
||||
positions=positions,
|
||||
out_cache_loc=_contiguous_out_cache_loc(req_pool_idx=1, start=0, count=3),
|
||||
req_to_token=_linear_r2t(),
|
||||
num_slots=64,
|
||||
extras=empty_extras(),
|
||||
enable_write_verify_inputs=True,
|
||||
expected_input_tokens=_t([99, 99, 99]),
|
||||
expected_input_positions=positions.clone(),
|
||||
ring_capacity=64,
|
||||
)
|
||||
|
||||
write_violations = int(log_real.write_index[0].item())
|
||||
assert (
|
||||
write_violations == n_tokens
|
||||
), f"expected {n_tokens} write violations, got {write_violations}"
|
||||
|
||||
|
||||
def test_pipeline_empty_batch() -> None:
|
||||
"""bs=1 with req_pool_idx=0 (padding): write and verify are no-op, kernel_run_counter == 2 (write+verify)."""
|
||||
input_ids, positions, out_cache_loc = _zero_no_write_inputs()
|
||||
|
||||
_, _, log_real, log_ref, _, _, _, _ = _run_both_and_assert_pipeline_equal(
|
||||
req_pool_indices=_t([0]),
|
||||
prefix_lens=_t([0]),
|
||||
extend_seq_lens=_t([0]),
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
req_to_token=_linear_r2t(),
|
||||
num_slots=64,
|
||||
extras=empty_extras(),
|
||||
)
|
||||
|
||||
assert int(log_real.kernel_run_counter[0].item()) == 2
|
||||
assert int(log_ref.kernel_run_counter[0].item()) == 2
|
||||
assert int(log_real.write_index[0].item()) == 0
|
||||
|
||||
|
||||
def test_pipeline_negative_slot_swa_out_of_window() -> None:
|
||||
"""SWA: some out_cache_loc entries map to -1 (out-of-window); write_step skips them, buf unchanged."""
|
||||
max_seq_len = 16
|
||||
max_reqs = 4
|
||||
|
||||
full_to_swa_index_mapping = torch.arange(
|
||||
max_reqs * max_seq_len + 1, dtype=torch.int64, device=_DEVICE
|
||||
)
|
||||
full_to_swa_index_mapping[1 * max_seq_len + 6] = -1
|
||||
full_to_swa_index_mapping[1 * max_seq_len + 7] = -1
|
||||
|
||||
_run_both_and_assert_pipeline_equal(
|
||||
req_pool_indices=_t([1]),
|
||||
prefix_lens=_t([6]),
|
||||
extend_seq_lens=_t([4]),
|
||||
input_ids=_t([100, 101, 102, 103]),
|
||||
positions=_t([6, 7, 8, 9]),
|
||||
out_cache_loc=_t([-1, -1, 1 * max_seq_len + 8, 1 * max_seq_len + 9]),
|
||||
req_to_token=_linear_r2t(max_reqs=max_reqs, max_seq_len=max_seq_len),
|
||||
num_slots=128,
|
||||
extras=empty_extras(),
|
||||
swa_window_size=4,
|
||||
full_to_swa_index_mapping=full_to_swa_index_mapping,
|
||||
)
|
||||
|
||||
|
||||
def test_pipeline_ring_overflow_via_real_plan() -> None:
|
||||
"""Verify detects >capacity violations when prev_hash is pre-corrupted; write_index byte-equal, ring relaxed."""
|
||||
max_seq_len = 16
|
||||
max_reqs = 4
|
||||
req_to_token = _linear_r2t(max_reqs=max_reqs, max_seq_len=max_seq_len)
|
||||
n_slots = 8
|
||||
|
||||
req_pool_indices = _t([1])
|
||||
prefix_lens = _t([n_slots])
|
||||
extend_seq_lens = _t([0])
|
||||
input_ids, positions, out_cache_loc = _zero_no_write_inputs()
|
||||
|
||||
num_slots = max_reqs * max_seq_len
|
||||
|
||||
# Step 1: pre-pollute canary_buf slots [0..n_slots) with wrong prev_hash so verify fires n_slots violations.
|
||||
buf_real = make_canary_buf(num_slots=num_slots, device=_DEVICE)
|
||||
buf_ref = make_canary_buf(num_slots=num_slots, device=_DEVICE)
|
||||
for slot_idx in range(n_slots):
|
||||
full_slot = 1 * max_seq_len + slot_idx
|
||||
for buf in (buf_real, buf_ref):
|
||||
write_slot_fields(
|
||||
canary_buf=buf,
|
||||
slot_idx=full_slot,
|
||||
token=slot_idx + 1,
|
||||
position=slot_idx,
|
||||
prev_hash=0x1234_DEAD_BEEF_0000 + slot_idx,
|
||||
)
|
||||
|
||||
# Step 2: run real pipeline (plan + no write + verify); overflow ring capacity=4 with all n_slots violations.
|
||||
ring_capacity = 4
|
||||
log_real = FakeViolationLog.allocate(capacity=ring_capacity, device=_DEVICE)
|
||||
log_ref = FakeViolationLog.allocate(capacity=ring_capacity, device=_DEVICE)
|
||||
plan_v_real = VerifyPlan.allocate(verify_capacity=256, device=_DEVICE)
|
||||
plan_w_real = WritePlan.allocate(write_req_capacity=4, device=_DEVICE)
|
||||
plan_v_ref = VerifyPlan.allocate(verify_capacity=256, device=_DEVICE)
|
||||
plan_w_ref = WritePlan.allocate(write_req_capacity=4, device=_DEVICE)
|
||||
|
||||
launch_canary_plan_kernels(
|
||||
verify_plan_out=plan_v_real,
|
||||
write_plan_out=plan_w_real,
|
||||
req_pool_indices=req_pool_indices,
|
||||
prefix_lens=prefix_lens,
|
||||
extend_seq_lens=extend_seq_lens,
|
||||
req_to_token=req_to_token,
|
||||
swa_window_size=0,
|
||||
full_to_swa_index_mapping=None,
|
||||
verify_capacity=int(plan_v_real.verify_slot_indices.shape[0]),
|
||||
req_to_verify_expected_tokens=None,
|
||||
req_to_verify_expected_tokens_valid_lens=None,
|
||||
kv_token_id_vs_position_offset=0,
|
||||
)
|
||||
launch_canary_plan_kernels_torch_reference(
|
||||
verify_plan_out=plan_v_ref,
|
||||
write_plan_out=plan_w_ref,
|
||||
req_pool_indices=req_pool_indices,
|
||||
prefix_lens=prefix_lens,
|
||||
extend_seq_lens=extend_seq_lens,
|
||||
req_to_token=req_to_token,
|
||||
swa_window_size=0,
|
||||
full_to_swa_index_mapping=None,
|
||||
verify_capacity=int(plan_v_ref.verify_slot_indices.shape[0]),
|
||||
req_to_verify_expected_tokens=None,
|
||||
req_to_verify_expected_tokens_valid_lens=None,
|
||||
kv_token_id_vs_position_offset=0,
|
||||
)
|
||||
|
||||
launch_canary_verify_kernel(
|
||||
context=VerifyOrWriteContext(
|
||||
canary_buf=buf_real,
|
||||
kernel_kind=CanaryLaunchTag.HEAD_K_FULL,
|
||||
violation_ring=log_real.ring,
|
||||
violation_write_index=log_real.write_index,
|
||||
slot_run_counter=log_real.slot_run_counter,
|
||||
kernel_run_counter=log_real.kernel_run_counter,
|
||||
enable_chain_position_assert=log_real.enable_chain_position_assert,
|
||||
),
|
||||
plan=plan_v_real,
|
||||
check_verify_expected_token=True,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
launch_canary_verify_kernel_torch_reference(
|
||||
context=VerifyOrWriteContext(
|
||||
canary_buf=buf_ref,
|
||||
kernel_kind=CanaryLaunchTag.HEAD_K_FULL,
|
||||
violation_ring=log_ref.ring,
|
||||
violation_write_index=log_ref.write_index,
|
||||
slot_run_counter=log_ref.slot_run_counter,
|
||||
kernel_run_counter=log_ref.kernel_run_counter,
|
||||
enable_chain_position_assert=log_ref.enable_chain_position_assert,
|
||||
),
|
||||
plan=plan_v_ref,
|
||||
check_verify_expected_token=True,
|
||||
)
|
||||
|
||||
# Step 3: write_index byte-equal; ring contents relaxed (atomic order not guaranteed under overflow).
|
||||
assert torch.equal(log_real.write_index, log_ref.write_index)
|
||||
assert int(log_real.write_index[0].item()) == n_slots
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"kernel_kind", [CanaryLaunchTag.HEAD_K_FULL, CanaryLaunchTag.TAIL_V_SWA]
|
||||
)
|
||||
def test_pipeline_kernel_kind_propagates(kernel_kind: CanaryLaunchTag) -> None:
|
||||
"""Different CanaryLaunchTag values: violation ring's kernel_kind field matches on both sides."""
|
||||
max_seq_len = 16
|
||||
input_ids, positions, out_cache_loc = _zero_no_write_inputs()
|
||||
|
||||
initial_buf = make_canary_buf(num_slots=64, device=_DEVICE)
|
||||
write_slot_fields(
|
||||
canary_buf=initial_buf,
|
||||
slot_idx=1 * max_seq_len,
|
||||
token=7,
|
||||
position=99,
|
||||
prev_hash=0,
|
||||
)
|
||||
|
||||
_, _, log_real, log_ref, _, _, _, _ = _run_both_and_assert_pipeline_equal(
|
||||
req_pool_indices=_t([1]),
|
||||
prefix_lens=_t([1]),
|
||||
extend_seq_lens=_t([0]),
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
req_to_token=_linear_r2t(max_seq_len=max_seq_len),
|
||||
num_slots=64,
|
||||
extras=empty_extras(),
|
||||
kernel_kind=kernel_kind,
|
||||
initial_canary_buf=initial_buf,
|
||||
)
|
||||
|
||||
assert int(log_real.write_index[0].item()) == 1
|
||||
assert int(log_ref.write_index[0].item()) == 1
|
||||
assert int(log_real.ring[0, consts.VIOLATION_FIELD_KERNEL_KIND].item()) == int(
|
||||
kernel_kind
|
||||
)
|
||||
assert int(log_ref.ring[0, consts.VIOLATION_FIELD_KERNEL_KIND].item()) == int(
|
||||
kernel_kind
|
||||
)
|
||||
|
||||
|
||||
def test_pipeline_token_mismatch_detected_via_pool() -> None:
|
||||
"""plan-pool gather + verify-token check: stamped wrong token id raises VERIFY_TOKEN_MISMATCH."""
|
||||
max_seq_len = 16
|
||||
max_reqs = 4
|
||||
prefix_len = 4
|
||||
input_ids, positions, out_cache_loc = _zero_no_write_inputs()
|
||||
|
||||
expected_tokens = [1000 + pos for pos in range(prefix_len)]
|
||||
pool = torch.full((max_reqs, max_seq_len), -999, dtype=torch.int32, device=_DEVICE)
|
||||
for pos, token in enumerate(expected_tokens):
|
||||
pool[1, pos] = token
|
||||
|
||||
stored_tokens = [token + 1 for token in expected_tokens]
|
||||
|
||||
initial_buf = make_canary_buf(num_slots=64, device=_DEVICE)
|
||||
initial_ref = initial_buf.clone()
|
||||
_stamp_linear_prefix(
|
||||
initial_buf=initial_buf,
|
||||
initial_ref=initial_ref,
|
||||
req_pool_idx=1,
|
||||
prefix_len=prefix_len,
|
||||
tokens=stored_tokens,
|
||||
max_seq_len=max_seq_len,
|
||||
)
|
||||
|
||||
_, _, log_real, log_ref, _, _, _, _ = _run_both_and_assert_pipeline_equal(
|
||||
req_pool_indices=_t([1]),
|
||||
prefix_lens=_t([prefix_len]),
|
||||
extend_seq_lens=_t([0]),
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
req_to_token=_linear_r2t(max_reqs=max_reqs, max_seq_len=max_seq_len),
|
||||
num_slots=64,
|
||||
extras=empty_extras(),
|
||||
swa_window_size=0,
|
||||
full_to_swa_index_mapping=None,
|
||||
initial_canary_buf=initial_buf,
|
||||
req_to_verify_expected_tokens=pool,
|
||||
kv_token_id_vs_position_offset=0,
|
||||
check_verify_expected_token=True,
|
||||
)
|
||||
|
||||
assert int(log_real.write_index[0].item()) == prefix_len
|
||||
assert int(log_ref.write_index[0].item()) == prefix_len
|
||||
# Ring rows may land in any order; collect stored/expected pairs and compare as sets.
|
||||
observed_pairs: set[tuple[int, int]] = set()
|
||||
for row_idx in range(prefix_len):
|
||||
fail_bits = int(
|
||||
log_real.ring[row_idx, consts.VIOLATION_FIELD_FAIL_REASON_BITS].item()
|
||||
)
|
||||
assert fail_bits & int(
|
||||
consts.FailReason.VERIFY_TOKEN_MISMATCH
|
||||
), f"row {row_idx}: VERIFY_TOKEN_MISMATCH bit missing in {fail_bits:#b}"
|
||||
stored = int(log_real.ring[row_idx, consts.VIOLATION_FIELD_STORED_TOKEN].item())
|
||||
expected = int(
|
||||
log_real.ring[row_idx, consts.VIOLATION_FIELD_EXPECTED_TOKEN].item()
|
||||
)
|
||||
observed_pairs.add((stored, expected))
|
||||
expected_pairs = {(stored_tokens[i], expected_tokens[i]) for i in range(prefix_len)}
|
||||
assert observed_pairs == expected_pairs
|
||||
|
||||
|
||||
def test_pipeline_eagle_offset_plus_1_byte_equal() -> None:
|
||||
"""plan-pool + offset=+1 full pipeline: stamped tokens match pool[rp, pos+1], no violations CUDA vs ref byte-equal."""
|
||||
max_seq_len = 16
|
||||
max_reqs = 4
|
||||
prefix_len = 4
|
||||
input_ids, positions, out_cache_loc = _zero_no_write_inputs()
|
||||
|
||||
stored_tokens = [2000 + pos for pos in range(prefix_len)]
|
||||
pool = torch.full((max_reqs, max_seq_len), -999, dtype=torch.int32, device=_DEVICE)
|
||||
for pos in range(prefix_len):
|
||||
# offset=+1 means kernel gathers from pool[rp, pos + 1], so place stored_tokens[pos] there.
|
||||
pool[1, pos + 1] = stored_tokens[pos]
|
||||
|
||||
initial_buf = make_canary_buf(num_slots=64, device=_DEVICE)
|
||||
initial_ref = initial_buf.clone()
|
||||
_stamp_linear_prefix(
|
||||
initial_buf=initial_buf,
|
||||
initial_ref=initial_ref,
|
||||
req_pool_idx=1,
|
||||
prefix_len=prefix_len,
|
||||
tokens=stored_tokens,
|
||||
max_seq_len=max_seq_len,
|
||||
)
|
||||
|
||||
_, _, log_real, log_ref, _, _, _, _ = _run_both_and_assert_pipeline_equal(
|
||||
req_pool_indices=_t([1]),
|
||||
prefix_lens=_t([prefix_len]),
|
||||
extend_seq_lens=_t([0]),
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
req_to_token=_linear_r2t(max_reqs=max_reqs, max_seq_len=max_seq_len),
|
||||
num_slots=64,
|
||||
extras=empty_extras(),
|
||||
swa_window_size=0,
|
||||
full_to_swa_index_mapping=None,
|
||||
initial_canary_buf=initial_buf,
|
||||
req_to_verify_expected_tokens=pool,
|
||||
kv_token_id_vs_position_offset=1,
|
||||
check_verify_expected_token=True,
|
||||
)
|
||||
|
||||
assert int(log_real.write_index[0].item()) == 0
|
||||
assert int(log_ref.write_index[0].item()) == 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
@@ -0,0 +1,222 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.tests.kv_canary._differential import _run_both_plan
|
||||
from sglang.jit_kernel.tests.kv_canary._fixtures import (
|
||||
allocate_plan_pair,
|
||||
derive_plan_capacity,
|
||||
make_lut,
|
||||
make_padding_mask,
|
||||
make_req_to_token,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._fuzz_driver import (
|
||||
FUZZ_SEEDS_PR,
|
||||
run_fuzz_combo,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._invariants import PlanInvariants
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
|
||||
|
||||
|
||||
_DEVICE = torch.device("cuda")
|
||||
|
||||
_FUZZ_ITER_PER_SEED = 50
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True, kw_only=True)
|
||||
class PlanFuzzInputs:
|
||||
req_pool_indices: torch.Tensor
|
||||
prefix_lens: torch.Tensor
|
||||
extend_seq_lens: torch.Tensor
|
||||
req_to_token: torch.Tensor
|
||||
swa_window_size: int
|
||||
full_to_swa_index_mapping: Optional[torch.Tensor]
|
||||
verify_capacity: int
|
||||
write_req_capacity: int
|
||||
req_to_verify_expected_tokens: Optional[torch.Tensor]
|
||||
kv_token_id_vs_position_offset: int
|
||||
|
||||
|
||||
def _draw_random_plan_inputs(rng: random.Random) -> PlanFuzzInputs:
|
||||
bs = rng.randint(1, 16)
|
||||
max_seq_len = rng.choice([8, 16, 64, 128, 256])
|
||||
swa_enabled = rng.random() < 0.5
|
||||
swa_window_size = (
|
||||
rng.choice([4, 16, 64, max_seq_len, max(2, max_seq_len // 3)])
|
||||
if swa_enabled
|
||||
else 0
|
||||
)
|
||||
swa_window_size = min(swa_window_size, max_seq_len)
|
||||
lut_kind = (
|
||||
rng.choice(["identity", "shift", "permutation", "with_oob"])
|
||||
if swa_enabled
|
||||
else None
|
||||
)
|
||||
rtt_kind = rng.choice(["linear", "sparse_permuted"])
|
||||
padding_kind = rng.choice(["none", "trailing", "interleaved"])
|
||||
capacity_kind = rng.choice(["loose", "tight_match", "under_by_one"])
|
||||
|
||||
max_reqs = max(bs + 2, 4)
|
||||
pool_size = max_reqs * max_seq_len
|
||||
rtt = make_req_to_token(
|
||||
kind=rtt_kind,
|
||||
max_reqs=max_reqs,
|
||||
max_seq_len=max_seq_len,
|
||||
device=_DEVICE,
|
||||
rng=rng,
|
||||
)
|
||||
padding_mask = make_padding_mask(bs=bs, kind=padding_kind, rng=rng)
|
||||
req_pool_indices_list: list[int] = []
|
||||
prefix_lens_list: list[int] = []
|
||||
extend_seq_lens_list: list[int] = []
|
||||
for r in range(bs):
|
||||
if padding_mask[r]:
|
||||
req_pool_indices_list.append(0)
|
||||
prefix_lens_list.append(0)
|
||||
extend_seq_lens_list.append(0)
|
||||
else:
|
||||
req_pool_indices_list.append(rng.randint(1, max_reqs - 1))
|
||||
prefix_lens_list.append(rng.randint(0, max_seq_len - 1))
|
||||
extend_seq_lens_list.append(rng.randint(1, max(1, max_seq_len // 4)))
|
||||
req_pool_indices = torch.tensor(
|
||||
req_pool_indices_list, dtype=torch.int64, device=_DEVICE
|
||||
)
|
||||
prefix_lens = torch.tensor(prefix_lens_list, dtype=torch.int64, device=_DEVICE)
|
||||
extend_seq_lens = torch.tensor(
|
||||
extend_seq_lens_list, dtype=torch.int64, device=_DEVICE
|
||||
)
|
||||
|
||||
total_verify = 0
|
||||
for rpi, pfx in zip(req_pool_indices_list, prefix_lens_list):
|
||||
if rpi == 0:
|
||||
continue
|
||||
if swa_window_size > 0:
|
||||
window_start = max(0, pfx - swa_window_size)
|
||||
total_verify += max(0, pfx - window_start)
|
||||
else:
|
||||
total_verify += pfx
|
||||
|
||||
verify_capacity, write_req_capacity = derive_plan_capacity(
|
||||
kind=capacity_kind,
|
||||
total_verify=total_verify,
|
||||
extras_count=0,
|
||||
bs=bs,
|
||||
)
|
||||
|
||||
full_to_swa: Optional[torch.Tensor]
|
||||
if swa_window_size > 0 and lut_kind is not None:
|
||||
full_to_swa = make_lut(
|
||||
kind=lut_kind, pool_size=pool_size, device=_DEVICE, rng=rng
|
||||
)
|
||||
else:
|
||||
full_to_swa = None
|
||||
|
||||
expected_pool_present = rng.random() < 0.5
|
||||
kv_token_id_vs_position_offset = rng.choice([0, 1])
|
||||
expected_pool: Optional[torch.Tensor]
|
||||
if expected_pool_present:
|
||||
pool_max_context_len = rng.choice(
|
||||
[
|
||||
max(1, max_seq_len // 4),
|
||||
max(1, max_seq_len // 2),
|
||||
max_seq_len,
|
||||
]
|
||||
)
|
||||
expected_pool = torch.randint(
|
||||
low=0,
|
||||
high=50000,
|
||||
size=(max_reqs, pool_max_context_len),
|
||||
dtype=torch.int32,
|
||||
device=_DEVICE,
|
||||
)
|
||||
else:
|
||||
expected_pool = None
|
||||
|
||||
return PlanFuzzInputs(
|
||||
req_pool_indices=req_pool_indices,
|
||||
prefix_lens=prefix_lens,
|
||||
extend_seq_lens=extend_seq_lens,
|
||||
req_to_token=rtt,
|
||||
swa_window_size=swa_window_size,
|
||||
full_to_swa_index_mapping=full_to_swa,
|
||||
verify_capacity=verify_capacity,
|
||||
write_req_capacity=write_req_capacity,
|
||||
req_to_verify_expected_tokens=expected_pool,
|
||||
kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
|
||||
)
|
||||
|
||||
|
||||
def _run_one(inputs: PlanFuzzInputs) -> tuple:
|
||||
triton_v, triton_w, ref_v, ref_w = allocate_plan_pair(
|
||||
verify_capacity=inputs.verify_capacity,
|
||||
write_req_capacity=inputs.write_req_capacity,
|
||||
)
|
||||
_run_both_plan(
|
||||
triton_verify=triton_v,
|
||||
triton_write=triton_w,
|
||||
ref_verify=ref_v,
|
||||
ref_write=ref_w,
|
||||
req_pool_indices=inputs.req_pool_indices,
|
||||
prefix_lens=inputs.prefix_lens,
|
||||
extend_seq_lens=inputs.extend_seq_lens,
|
||||
req_to_token=inputs.req_to_token,
|
||||
extras=(
|
||||
torch.empty(0, dtype=torch.int64, device=_DEVICE),
|
||||
torch.empty(0, dtype=torch.int64, device=_DEVICE),
|
||||
torch.empty(0, dtype=torch.int64, device=_DEVICE),
|
||||
torch.zeros(1, dtype=torch.int32, device=_DEVICE),
|
||||
),
|
||||
swa_window_size=inputs.swa_window_size,
|
||||
full_to_swa_index_mapping=inputs.full_to_swa_index_mapping,
|
||||
req_to_verify_expected_tokens=inputs.req_to_verify_expected_tokens,
|
||||
kv_token_id_vs_position_offset=inputs.kv_token_id_vs_position_offset,
|
||||
)
|
||||
PlanInvariants.assert_all(
|
||||
verify_plan=triton_v,
|
||||
write_plan=triton_w,
|
||||
req_pool_indices=inputs.req_pool_indices,
|
||||
prefix_lens=inputs.prefix_lens,
|
||||
extend_seq_lens=inputs.extend_seq_lens,
|
||||
swa_window_size=inputs.swa_window_size,
|
||||
extras_slot_indices=torch.empty(0, dtype=torch.int64, device=_DEVICE),
|
||||
extras_positions=torch.empty(0, dtype=torch.int64, device=_DEVICE),
|
||||
extras_prev_slot_indices=torch.empty(0, dtype=torch.int64, device=_DEVICE),
|
||||
extras_count=0,
|
||||
)
|
||||
return triton_v, triton_w
|
||||
|
||||
|
||||
def _summarize(inputs: PlanFuzzInputs) -> str:
|
||||
return (
|
||||
f"bs={int(inputs.req_pool_indices.shape[0])} "
|
||||
f"swa={inputs.swa_window_size} "
|
||||
f"verify_cap={inputs.verify_capacity} write_cap={inputs.write_req_capacity} "
|
||||
f"has_lut={inputs.full_to_swa_index_mapping is not None} "
|
||||
f"has_pool={inputs.req_to_verify_expected_tokens is not None} "
|
||||
f"offset={inputs.kv_token_id_vs_position_offset}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("seed", FUZZ_SEEDS_PR)
|
||||
def test_plan_fuzz_full_combo(seed: int) -> None:
|
||||
"""Multi-dim plan fuzzer: random LUT/rtt/padding/capacity/swa × N iters, byte-equal."""
|
||||
run_fuzz_combo(
|
||||
seed,
|
||||
draw_fn=_draw_random_plan_inputs,
|
||||
run_one_fn=_run_one,
|
||||
summarize_fn=_summarize,
|
||||
n_iter=_FUZZ_ITER_PER_SEED,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,178 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.kv_canary import consts
|
||||
from sglang.jit_kernel.kv_canary.verify import (
|
||||
CanaryLaunchTag,
|
||||
VerifyPlan,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
|
||||
FakeViolationLog,
|
||||
make_canary_buf,
|
||||
make_log_pair,
|
||||
make_verify_plan_pair,
|
||||
stamp_clean_chain,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._differential import _run_both_verify
|
||||
from sglang.jit_kernel.tests.kv_canary._fuzz_driver import (
|
||||
FUZZ_SEEDS_PR,
|
||||
run_fuzz_combo,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._invariants import VerifyInvariants
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
|
||||
|
||||
|
||||
_DEVICE = torch.device("cuda")
|
||||
|
||||
_FUZZ_ITER_PER_SEED = 30
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True, kw_only=True)
|
||||
class VerifyFuzzInputs:
|
||||
cuda_canary_buf: torch.Tensor
|
||||
ref_canary_buf: torch.Tensor
|
||||
plan_cuda: VerifyPlan
|
||||
plan_ref: VerifyPlan
|
||||
kernel_kind: CanaryLaunchTag
|
||||
ring_capacity: int
|
||||
check_verify_expected_token: bool
|
||||
|
||||
|
||||
def _draw_random_verify_inputs(rng: random.Random) -> VerifyFuzzInputs:
|
||||
kernel_kind = rng.choice(list(CanaryLaunchTag))
|
||||
plan_size = rng.randint(0, 32)
|
||||
num_slots = max(plan_size + 8, 16)
|
||||
ring_capacity = rng.choice([16, 64, 256])
|
||||
|
||||
cuda_buf = make_canary_buf(
|
||||
num_slots=num_slots, slot_stride_bytes=32, device=_DEVICE
|
||||
)
|
||||
ref_buf = cuda_buf.clone()
|
||||
|
||||
slot_universe = list(range(1, num_slots))
|
||||
rng.shuffle(slot_universe)
|
||||
slot_indices = slot_universe[:plan_size]
|
||||
tokens = [rng.randint(0, 0xFFFFFFFF) for _ in range(plan_size)]
|
||||
positions = list(range(plan_size))
|
||||
prev_slot_indices: list[int] = []
|
||||
for i in range(plan_size):
|
||||
if i == 0:
|
||||
prev_slot_indices.append(-1)
|
||||
else:
|
||||
prev_slot_indices.append(slot_indices[i - 1])
|
||||
|
||||
if plan_size > 0:
|
||||
stamp_clean_chain(
|
||||
cuda_buf=cuda_buf,
|
||||
ref_buf=ref_buf,
|
||||
slot_indices=slot_indices,
|
||||
tokens=tokens,
|
||||
positions=positions,
|
||||
)
|
||||
|
||||
# Inject prev_slot == TOKEN_TO_KV_SLOT_PADDING into ~15% of entries so the differential
|
||||
# harness exercises the chain-check-skip branch (added for SWA-evicted ancestor handling).
|
||||
# Done AFTER stamp_clean_chain so the stored prev_hash on those slots is still chain-clean;
|
||||
# the kernel must rely on prev_slot==padding (not on stored hash) to decide whether to skip.
|
||||
for i in range(plan_size):
|
||||
if rng.random() < 0.15:
|
||||
prev_slot_indices[i] = consts.TOKEN_TO_KV_SLOT_PADDING
|
||||
|
||||
check_verify_expected_token = rng.random() < 0.5
|
||||
expected_input_ids: list[int] = []
|
||||
for i in range(plan_size):
|
||||
# Always pick a value; with check=False the kernel must not deref this column.
|
||||
if rng.random() < 0.3:
|
||||
expected_input_ids.append(-1)
|
||||
elif rng.random() < 0.5:
|
||||
expected_input_ids.append(int(tokens[i]))
|
||||
else:
|
||||
mutated = (int(tokens[i]) ^ 0x1) & 0xFFFFFFFF
|
||||
expected_input_ids.append(mutated)
|
||||
|
||||
plan_cuda, plan_ref = make_verify_plan_pair(
|
||||
slot_indices=slot_indices,
|
||||
positions=positions,
|
||||
prev_slot_indices=prev_slot_indices,
|
||||
expected_input_ids=expected_input_ids if plan_size > 0 else None,
|
||||
capacity=max(plan_size, 1),
|
||||
device=_DEVICE,
|
||||
)
|
||||
|
||||
return VerifyFuzzInputs(
|
||||
cuda_canary_buf=cuda_buf,
|
||||
ref_canary_buf=ref_buf,
|
||||
plan_cuda=plan_cuda,
|
||||
plan_ref=plan_ref,
|
||||
kernel_kind=kernel_kind,
|
||||
ring_capacity=ring_capacity,
|
||||
check_verify_expected_token=check_verify_expected_token,
|
||||
)
|
||||
|
||||
|
||||
def _run_one(inputs: VerifyFuzzInputs) -> None:
|
||||
cuda_buf_before = inputs.cuda_canary_buf.clone()
|
||||
cuda_log, ref_log = make_log_pair(capacity=inputs.ring_capacity, device=_DEVICE)
|
||||
log_before = FakeViolationLog.allocate(
|
||||
capacity=inputs.ring_capacity, device=_DEVICE
|
||||
)
|
||||
_run_both_verify(
|
||||
cuda_canary_buf=inputs.cuda_canary_buf,
|
||||
ref_canary_buf=inputs.ref_canary_buf,
|
||||
plan_cuda=inputs.plan_cuda,
|
||||
plan_ref=inputs.plan_ref,
|
||||
cuda_log=cuda_log,
|
||||
ref_log=ref_log,
|
||||
kernel_kind=inputs.kernel_kind,
|
||||
assert_equal=False,
|
||||
check_verify_expected_token=inputs.check_verify_expected_token,
|
||||
)
|
||||
assert int(cuda_log.kernel_run_counter[0].item()) == int(
|
||||
ref_log.kernel_run_counter[0].item()
|
||||
)
|
||||
assert int(cuda_log.slot_run_counter[0].item()) == int(
|
||||
ref_log.slot_run_counter[0].item()
|
||||
)
|
||||
assert int(cuda_log.write_index[0].item()) == int(ref_log.write_index[0].item())
|
||||
VerifyInvariants.assert_all(
|
||||
canary_buf_before=cuda_buf_before,
|
||||
canary_buf_after=inputs.cuda_canary_buf,
|
||||
log_before=log_before,
|
||||
log_after=cuda_log,
|
||||
plan=inputs.plan_cuda,
|
||||
kernel_kind=inputs.kernel_kind,
|
||||
)
|
||||
|
||||
|
||||
def _summarize(inputs: VerifyFuzzInputs) -> str:
|
||||
n_active = int(inputs.plan_cuda.verify_num_valid[0].item())
|
||||
return (
|
||||
f"plan_size={n_active} kind={inputs.kernel_kind.name} "
|
||||
f"ring={inputs.ring_capacity} "
|
||||
f"check_token={inputs.check_verify_expected_token}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("seed", FUZZ_SEEDS_PR)
|
||||
def test_verify_fuzz_full_combo(seed: int) -> None:
|
||||
"""Multi-dim verify fuzzer: random kernel kind × plan size × ring capacity × N iters, byte-equal."""
|
||||
run_fuzz_combo(
|
||||
seed,
|
||||
draw_fn=_draw_random_verify_inputs,
|
||||
run_one_fn=_run_one,
|
||||
summarize_fn=_summarize,
|
||||
n_iter=_FUZZ_ITER_PER_SEED,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,218 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.kv_canary.verify import CanaryLaunchTag
|
||||
from sglang.jit_kernel.kv_canary.write import WritePlan
|
||||
from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
|
||||
FakeViolationLog,
|
||||
make_canary_buf,
|
||||
make_log_pair,
|
||||
make_write_plan_pair,
|
||||
stamp_pair,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._differential import _run_both_write
|
||||
from sglang.jit_kernel.tests.kv_canary._fuzz_driver import (
|
||||
FUZZ_SEEDS_PR,
|
||||
run_fuzz_combo,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._invariants import WriteInvariants
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
|
||||
|
||||
|
||||
_DEVICE = torch.device("cuda")
|
||||
|
||||
_FUZZ_ITER_PER_SEED = 30
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True, kw_only=True)
|
||||
class WriteFuzzInputs:
|
||||
cuda_canary_buf: torch.Tensor
|
||||
ref_canary_buf: torch.Tensor
|
||||
plan_cuda: WritePlan
|
||||
plan_ref: WritePlan
|
||||
input_ids: torch.Tensor
|
||||
positions: torch.Tensor
|
||||
out_cache_loc: torch.Tensor
|
||||
kernel_kind: CanaryLaunchTag
|
||||
enable_write_verify_inputs: bool
|
||||
expected_input_tokens: torch.Tensor
|
||||
expected_input_positions: torch.Tensor
|
||||
ring_capacity: int
|
||||
|
||||
|
||||
def _draw_random_write_inputs(rng: random.Random) -> WriteFuzzInputs:
|
||||
enable_write_verify_inputs = rng.choice([False, True])
|
||||
kernel_kind = rng.choice(list(CanaryLaunchTag))
|
||||
ring_capacity = rng.choice([16, 64, 256])
|
||||
|
||||
n_reqs = rng.randint(1, 4)
|
||||
per_req_tokens: list[int] = [rng.randint(1, 5) for _ in range(n_reqs)]
|
||||
total_tokens = sum(per_req_tokens)
|
||||
num_slots = max(total_tokens + 8, 16)
|
||||
|
||||
cuda_buf = make_canary_buf(
|
||||
num_slots=num_slots, slot_stride_bytes=32, device=_DEVICE
|
||||
)
|
||||
ref_buf = cuda_buf.clone()
|
||||
|
||||
write_offsets: list[int] = [0]
|
||||
running = 0
|
||||
for t in per_req_tokens:
|
||||
running += t
|
||||
write_offsets.append(running)
|
||||
slot_pool = list(range(1, num_slots))
|
||||
rng.shuffle(slot_pool)
|
||||
seed_slot_indices: list[int] = []
|
||||
for _ in range(n_reqs):
|
||||
if rng.random() < 0.4 or len(slot_pool) <= total_tokens:
|
||||
seed_slot_indices.append(-1)
|
||||
else:
|
||||
seed_slot_indices.append(slot_pool.pop())
|
||||
|
||||
out_cache_loc_list: list[int] = []
|
||||
for _ in range(total_tokens):
|
||||
if not slot_pool:
|
||||
out_cache_loc_list.append(-1)
|
||||
else:
|
||||
out_cache_loc_list.append(slot_pool.pop())
|
||||
|
||||
plan_cuda, plan_ref = make_write_plan_pair(
|
||||
write_offsets=write_offsets,
|
||||
seed_slot_indices=seed_slot_indices,
|
||||
num_valid_reqs=n_reqs,
|
||||
device=_DEVICE,
|
||||
)
|
||||
|
||||
input_ids = torch.tensor(
|
||||
[rng.randint(-(1 << 31), (1 << 31) - 1) for _ in range(total_tokens)],
|
||||
dtype=torch.int64,
|
||||
device=_DEVICE,
|
||||
)
|
||||
# Per-chain sequential positions so the write kernel's chain-step position assert holds.
|
||||
# For chains with a real seed slot, stamp the seed with (chain_start_position - 1) so the
|
||||
# first chain entry's position == seed.position + 1.
|
||||
chain_start_positions: list[int] = [rng.randint(0, 1024) for _ in range(n_reqs)]
|
||||
positions_list: list[int] = []
|
||||
for r in range(n_reqs):
|
||||
start = chain_start_positions[r]
|
||||
positions_list.extend(start + i for i in range(per_req_tokens[r]))
|
||||
seed_slot = seed_slot_indices[r]
|
||||
if seed_slot >= 0:
|
||||
stamp_pair(
|
||||
(cuda_buf, ref_buf),
|
||||
slot_idx=seed_slot,
|
||||
token=0,
|
||||
position=start - 1,
|
||||
prev_hash=0,
|
||||
)
|
||||
positions = torch.tensor(positions_list, dtype=torch.int64, device=_DEVICE)
|
||||
out_cache_loc = torch.tensor(out_cache_loc_list, dtype=torch.int64, device=_DEVICE)
|
||||
expected_input_tokens = input_ids.clone()
|
||||
expected_input_positions = positions.clone()
|
||||
if enable_write_verify_inputs:
|
||||
candidate_indices = [
|
||||
idx for idx, slot in enumerate(out_cache_loc_list) if slot >= 0
|
||||
]
|
||||
rng.shuffle(candidate_indices)
|
||||
mismatch_count = rng.randint(0, len(candidate_indices))
|
||||
for idx in candidate_indices[:mismatch_count]:
|
||||
if rng.choice([False, True]):
|
||||
expected_input_tokens[idx] = expected_input_tokens[idx] + 1
|
||||
else:
|
||||
expected_input_positions[idx] = expected_input_positions[idx] + 1
|
||||
|
||||
return WriteFuzzInputs(
|
||||
cuda_canary_buf=cuda_buf,
|
||||
ref_canary_buf=ref_buf,
|
||||
plan_cuda=plan_cuda,
|
||||
plan_ref=plan_ref,
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
kernel_kind=kernel_kind,
|
||||
enable_write_verify_inputs=enable_write_verify_inputs,
|
||||
expected_input_tokens=expected_input_tokens,
|
||||
expected_input_positions=expected_input_positions,
|
||||
ring_capacity=ring_capacity,
|
||||
)
|
||||
|
||||
|
||||
def _run_one(inputs: WriteFuzzInputs) -> None:
|
||||
cuda_buf_before = inputs.cuda_canary_buf.clone()
|
||||
cuda_log, ref_log = make_log_pair(capacity=inputs.ring_capacity, device=_DEVICE)
|
||||
log_before = FakeViolationLog.allocate(
|
||||
capacity=inputs.ring_capacity, device=_DEVICE
|
||||
)
|
||||
_run_both_write(
|
||||
cuda_canary_buf=inputs.cuda_canary_buf,
|
||||
ref_canary_buf=inputs.ref_canary_buf,
|
||||
plan_cuda=inputs.plan_cuda,
|
||||
plan_ref=inputs.plan_ref,
|
||||
input_ids=inputs.input_ids,
|
||||
positions=inputs.positions,
|
||||
out_cache_loc=inputs.out_cache_loc,
|
||||
enable_write_verify_inputs=inputs.enable_write_verify_inputs,
|
||||
expected_input_tokens=inputs.expected_input_tokens,
|
||||
expected_input_positions=inputs.expected_input_positions,
|
||||
cuda_log=cuda_log,
|
||||
ref_log=ref_log,
|
||||
kernel_kind=inputs.kernel_kind,
|
||||
assert_equal=False,
|
||||
)
|
||||
assert torch.equal(
|
||||
inputs.cuda_canary_buf, inputs.ref_canary_buf
|
||||
), "CUDA vs ref canary_buf diverged"
|
||||
assert int(cuda_log.write_index[0].item()) == int(ref_log.write_index[0].item())
|
||||
assert int(cuda_log.slot_run_counter[0].item()) == int(
|
||||
ref_log.slot_run_counter[0].item()
|
||||
)
|
||||
assert int(cuda_log.kernel_run_counter[0].item()) == int(
|
||||
ref_log.kernel_run_counter[0].item()
|
||||
)
|
||||
WriteInvariants.assert_all(
|
||||
canary_buf_before=cuda_buf_before,
|
||||
canary_buf_after=inputs.cuda_canary_buf,
|
||||
plan=inputs.plan_cuda,
|
||||
input_ids=inputs.input_ids,
|
||||
positions=inputs.positions,
|
||||
out_cache_loc=inputs.out_cache_loc,
|
||||
enable_write_verify_inputs=inputs.enable_write_verify_inputs,
|
||||
expected_input_tokens=inputs.expected_input_tokens,
|
||||
expected_input_positions=inputs.expected_input_positions,
|
||||
log_before=log_before,
|
||||
log_after=cuda_log,
|
||||
)
|
||||
|
||||
|
||||
def _summarize(inputs: WriteFuzzInputs) -> str:
|
||||
n_active = int(inputs.plan_cuda.write_num_valid_reqs[0].item())
|
||||
total = int(inputs.plan_cuda.write_offsets[n_active].item())
|
||||
return (
|
||||
f"n_reqs={n_active} total_tokens={total} kind={inputs.kernel_kind.name} "
|
||||
f"pseudo={inputs.enable_write_verify_inputs}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("seed", FUZZ_SEEDS_PR)
|
||||
def test_write_fuzz_full_combo(seed: int) -> None:
|
||||
"""Multi-dim write fuzzer: random pseudo/kernel × N iters, byte-equal."""
|
||||
run_fuzz_combo(
|
||||
seed,
|
||||
draw_fn=_draw_random_write_inputs,
|
||||
run_one_fn=_run_one,
|
||||
summarize_fn=_summarize,
|
||||
n_iter=_FUZZ_ITER_PER_SEED,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
@@ -0,0 +1,892 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.kv_canary import consts
|
||||
from sglang.jit_kernel.kv_canary import write as write_module
|
||||
from sglang.jit_kernel.kv_canary.consts import splitmix64, splitmix64_mix3
|
||||
from sglang.jit_kernel.kv_canary.verify import (
|
||||
CANARY_SLOT_BYTES,
|
||||
CanaryLaunchTag,
|
||||
VerifyOrWriteContext,
|
||||
launch_canary_verify_kernel,
|
||||
)
|
||||
from sglang.jit_kernel.kv_canary.write import (
|
||||
launch_canary_write_kernel,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
|
||||
FakeViolationLog,
|
||||
assert_canary_state_equal,
|
||||
assert_only_bits_set,
|
||||
chain_anchor_signed,
|
||||
make_canary_buf,
|
||||
make_canary_buf_pair,
|
||||
make_log_pair,
|
||||
make_verify_plan,
|
||||
make_write_plan,
|
||||
make_write_plan_pair,
|
||||
read_slot_fields,
|
||||
stamp_pair,
|
||||
to_signed_int64,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._differential import (
|
||||
_run_both_write,
|
||||
run_write_diff,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._fixtures import (
|
||||
dummy_pseudo_tensors,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
|
||||
|
||||
|
||||
_DEVICE = torch.device("cuda")
|
||||
|
||||
|
||||
def _int32_tensor(values: list[int]) -> torch.Tensor:
|
||||
return torch.tensor(values, dtype=torch.int64, device=_DEVICE)
|
||||
|
||||
|
||||
def _make_default_buf_pair(
|
||||
num_slots: int = 16, slot_stride_bytes: int = 32
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
return make_canary_buf_pair(
|
||||
num_slots=num_slots, slot_stride_bytes=slot_stride_bytes, device=_DEVICE
|
||||
)
|
||||
|
||||
|
||||
def _run_write(
|
||||
*,
|
||||
buf_pair: tuple[torch.Tensor, torch.Tensor],
|
||||
input_ids: list[int] | torch.Tensor,
|
||||
positions: list[int] | torch.Tensor,
|
||||
out_cache_loc: list[int] | torch.Tensor,
|
||||
write_offsets: list[int] | None = None,
|
||||
seed_slot_indices: list[int] = (-1,),
|
||||
num_valid_reqs: int = 1,
|
||||
req_capacity: int | None = None,
|
||||
enable_write_verify_inputs: bool = False,
|
||||
expected_input_tokens: torch.Tensor | None = None,
|
||||
expected_input_positions: torch.Tensor | None = None,
|
||||
assert_equal: bool = True,
|
||||
) -> tuple[FakeViolationLog, FakeViolationLog]:
|
||||
"""Shared scaffold: build write plan + pseudo tensors and call ``run_write_diff``.
|
||||
|
||||
``write_offsets`` defaults to ``[0, len(input_ids)]`` (one req covering all entries).
|
||||
``expected_input_*`` default to ``dummy_pseudo_tensors(len(input_ids))``.
|
||||
"""
|
||||
ids_t = (
|
||||
input_ids
|
||||
if isinstance(input_ids, torch.Tensor)
|
||||
else _int32_tensor(list(input_ids))
|
||||
)
|
||||
pos_t = (
|
||||
positions
|
||||
if isinstance(positions, torch.Tensor)
|
||||
else _int32_tensor(list(positions))
|
||||
)
|
||||
loc_t = (
|
||||
out_cache_loc
|
||||
if isinstance(out_cache_loc, torch.Tensor)
|
||||
else _int32_tensor(list(out_cache_loc))
|
||||
)
|
||||
n_tokens = int(ids_t.shape[0])
|
||||
|
||||
if write_offsets is None:
|
||||
write_offsets = [0, n_tokens]
|
||||
|
||||
plan_kwargs = dict(
|
||||
write_offsets=write_offsets,
|
||||
seed_slot_indices=list(seed_slot_indices),
|
||||
num_valid_reqs=num_valid_reqs,
|
||||
device=_DEVICE,
|
||||
)
|
||||
if req_capacity is not None:
|
||||
plan_kwargs["req_capacity"] = req_capacity
|
||||
plan_pair = make_write_plan_pair(**plan_kwargs)
|
||||
|
||||
if expected_input_tokens is None or expected_input_positions is None:
|
||||
pseudo_tokens, pseudo_positions = dummy_pseudo_tensors(n_tokens)
|
||||
if expected_input_tokens is None:
|
||||
expected_input_tokens = pseudo_tokens
|
||||
if expected_input_positions is None:
|
||||
expected_input_positions = pseudo_positions
|
||||
|
||||
return run_write_diff(
|
||||
buf_pair=buf_pair,
|
||||
plan_pair=plan_pair,
|
||||
input_ids=ids_t,
|
||||
positions=pos_t,
|
||||
out_cache_loc=loc_t,
|
||||
enable_write_verify_inputs=enable_write_verify_inputs,
|
||||
expected_input_tokens=expected_input_tokens,
|
||||
expected_input_positions=expected_input_positions,
|
||||
assert_equal=assert_equal,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True, kw_only=True)
|
||||
class _WriteSingleSlotInput:
|
||||
token: int = 42
|
||||
position: int = 0
|
||||
enable_write_verify_inputs: bool = False
|
||||
|
||||
|
||||
class _RecordingWriteModule:
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[tuple[object, ...]] = []
|
||||
|
||||
def canary_write_step_cuda(self, *args: object) -> None:
|
||||
self.calls.append(args)
|
||||
|
||||
|
||||
def _run_write_single_slot_byte_equal(case: _WriteSingleSlotInput) -> None:
|
||||
_run_write(
|
||||
buf_pair=_make_default_buf_pair(),
|
||||
input_ids=[case.token],
|
||||
positions=[case.position],
|
||||
out_cache_loc=[0],
|
||||
enable_write_verify_inputs=case.enable_write_verify_inputs,
|
||||
)
|
||||
|
||||
|
||||
class TestSeedSlot:
|
||||
def setup_method(self) -> None:
|
||||
self.buf_pair = _make_default_buf_pair()
|
||||
|
||||
def test_seed_slot_idx_negative_uses_anchor(self) -> None:
|
||||
"""``seed_slot_idx == -1`` → initial ``running_prev_hash`` is ``splitmix64(consts.CANARY_CHAIN_ANCHOR)``."""
|
||||
_run_write(
|
||||
buf_pair=self.buf_pair,
|
||||
input_ids=[42],
|
||||
positions=[0],
|
||||
out_cache_loc=[3],
|
||||
)
|
||||
|
||||
stored_token, stored_position, stored_prev_hash, _ = read_slot_fields(
|
||||
canary_buf=self.buf_pair[0], slot_idx=3
|
||||
)
|
||||
assert stored_token == 42
|
||||
assert stored_position == 0
|
||||
assert stored_prev_hash == chain_anchor_signed()
|
||||
|
||||
def test_seed_slot_idx_loads_predecessor(self) -> None:
|
||||
"""``seed_slot_idx >= 0`` → load 3 fields (token, position, prev_hash) from ``canary_buf[seed]`` and splitmix64_mix3-advance into prev_hash."""
|
||||
# Step: pre-stamp slot 7 with a known chain link.
|
||||
seed_token, seed_position = 100, 4
|
||||
seed_prev_signed = to_signed_int64(splitmix64(consts.CANARY_CHAIN_ANCHOR))
|
||||
stamp_pair(
|
||||
self.buf_pair,
|
||||
slot_idx=7,
|
||||
token=seed_token,
|
||||
position=seed_position,
|
||||
prev_hash=seed_prev_signed,
|
||||
)
|
||||
|
||||
_run_write(
|
||||
buf_pair=self.buf_pair,
|
||||
input_ids=[999],
|
||||
positions=[5],
|
||||
out_cache_loc=[2],
|
||||
seed_slot_indices=[7],
|
||||
)
|
||||
|
||||
expected_prev_hash = splitmix64_mix3(
|
||||
splitmix64(consts.CANARY_CHAIN_ANCHOR), seed_token, seed_position
|
||||
)
|
||||
_, _, stored_prev_hash, _ = read_slot_fields(
|
||||
canary_buf=self.buf_pair[0], slot_idx=2
|
||||
)
|
||||
assert stored_prev_hash == to_signed_int64(expected_prev_hash)
|
||||
|
||||
def test_seed_slot_chain_link_continuous(self) -> None:
|
||||
"""After write, ``slot[0].prev_hash`` is consistent with verify's chain reconstruction from seed."""
|
||||
# Step 1: write a chain from seed slot=7 → newly written slot=2. Then run verify with prev=7 and
|
||||
# assert no violation — i.e., slot[2].prev_hash is splitmix64_mix3(seed.prev_hash, seed.token, seed.position).
|
||||
cuda_buf = self.buf_pair[0]
|
||||
seed_token, seed_position = 11, 0
|
||||
seed_prev_signed = to_signed_int64(splitmix64(consts.CANARY_CHAIN_ANCHOR))
|
||||
stamp_pair(
|
||||
self.buf_pair,
|
||||
slot_idx=7,
|
||||
token=seed_token,
|
||||
position=seed_position,
|
||||
prev_hash=seed_prev_signed,
|
||||
)
|
||||
|
||||
_run_write(
|
||||
buf_pair=self.buf_pair,
|
||||
input_ids=[222],
|
||||
positions=[1],
|
||||
out_cache_loc=[2],
|
||||
seed_slot_indices=[7],
|
||||
assert_equal=False,
|
||||
)
|
||||
|
||||
# Step 2: verify slot[2] with prev=7 — expects no violation.
|
||||
verify_plan = make_verify_plan(
|
||||
slot_indices=[2], positions=[1], prev_slot_indices=[7], device=_DEVICE
|
||||
)
|
||||
verify_log = FakeViolationLog.allocate(device=_DEVICE)
|
||||
launch_canary_verify_kernel(
|
||||
context=VerifyOrWriteContext(
|
||||
canary_buf=cuda_buf,
|
||||
kernel_kind=CanaryLaunchTag.HEAD_K_FULL,
|
||||
violation_ring=verify_log.ring,
|
||||
violation_write_index=verify_log.write_index,
|
||||
slot_run_counter=verify_log.slot_run_counter,
|
||||
kernel_run_counter=verify_log.kernel_run_counter,
|
||||
enable_chain_position_assert=verify_log.enable_chain_position_assert,
|
||||
),
|
||||
plan=verify_plan,
|
||||
check_verify_expected_token=True,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
assert int(verify_log.write_index[0].item()) == 0
|
||||
|
||||
def test_seed_slot_resume_5_step_hardcoded(self) -> None:
|
||||
cuda_buf = make_canary_buf(num_slots=50, slot_stride_bytes=32, device=_DEVICE)
|
||||
ref_buf = cuda_buf.clone()
|
||||
buf_pair = (cuda_buf, ref_buf)
|
||||
seed_token = 7
|
||||
seed_position = 10
|
||||
seed_prev_hash_signed = to_signed_int64(splitmix64(consts.CANARY_CHAIN_ANCHOR))
|
||||
stamp_pair(
|
||||
buf_pair,
|
||||
slot_idx=42,
|
||||
token=seed_token,
|
||||
position=seed_position,
|
||||
prev_hash=seed_prev_hash_signed,
|
||||
)
|
||||
|
||||
predecessor_advance = splitmix64_mix3(
|
||||
splitmix64(consts.CANARY_CHAIN_ANCHOR), seed_token, seed_position
|
||||
)
|
||||
|
||||
tokens = [101, 202, 303, 404, 505]
|
||||
positions = [11, 12, 13, 14, 15]
|
||||
out_cache_loc = [0, 1, 2, 3, 4]
|
||||
|
||||
expected_prev_hashes: list[int] = []
|
||||
running = predecessor_advance
|
||||
for t, p in zip(tokens, positions):
|
||||
expected_prev_hashes.append(running)
|
||||
running = splitmix64_mix3(running, t, p)
|
||||
|
||||
cuda_log, _ = _run_write(
|
||||
buf_pair=buf_pair,
|
||||
input_ids=tokens,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
write_offsets=[0, 5],
|
||||
seed_slot_indices=[42],
|
||||
)
|
||||
|
||||
for slot_idx, expected_token, expected_position, expected_prev_u64 in zip(
|
||||
out_cache_loc, tokens, positions, expected_prev_hashes
|
||||
):
|
||||
stored_token, stored_position, stored_prev_hash, stored_real_kv_hash = (
|
||||
read_slot_fields(canary_buf=cuda_buf, slot_idx=slot_idx)
|
||||
)
|
||||
assert stored_token == expected_token
|
||||
assert stored_position == expected_position
|
||||
assert stored_prev_hash == to_signed_int64(expected_prev_u64)
|
||||
assert stored_real_kv_hash == 0
|
||||
|
||||
assert int(cuda_log.write_index[0].item()) == 0
|
||||
|
||||
def test_seed_continues_existing_chain(self) -> None:
|
||||
"""Pre-stamp seed slot; subsequent write should continue chain from splitmix64_mix3(seed.*)."""
|
||||
seed_slot = 3
|
||||
seed_token = 7
|
||||
seed_position = 1
|
||||
expected_seed_prev_hash = splitmix64(consts.CANARY_CHAIN_ANCHOR)
|
||||
stamp_pair(
|
||||
self.buf_pair,
|
||||
slot_idx=seed_slot,
|
||||
token=seed_token,
|
||||
position=seed_position,
|
||||
prev_hash=to_signed_int64(expected_seed_prev_hash),
|
||||
)
|
||||
|
||||
new_slot = 4
|
||||
new_token = 13
|
||||
new_position = 2
|
||||
expected_running = splitmix64_mix3(
|
||||
expected_seed_prev_hash, seed_token, seed_position
|
||||
)
|
||||
|
||||
_run_write(
|
||||
buf_pair=self.buf_pair,
|
||||
input_ids=[new_token],
|
||||
positions=[new_position],
|
||||
out_cache_loc=[new_slot],
|
||||
seed_slot_indices=[seed_slot],
|
||||
)
|
||||
|
||||
new_stored = read_slot_fields(canary_buf=self.buf_pair[0], slot_idx=new_slot)
|
||||
assert new_stored[0] == new_token
|
||||
assert new_stored[1] == new_position
|
||||
assert new_stored[2] == to_signed_int64(
|
||||
expected_running
|
||||
), f"new slot prev_hash {new_stored[2]} != expected {to_signed_int64(expected_running)}"
|
||||
|
||||
|
||||
class TestChain:
|
||||
def setup_method(self) -> None:
|
||||
self.buf_pair = _make_default_buf_pair()
|
||||
|
||||
def test_chain_link_byte_equal_5_step(self) -> None:
|
||||
"""5-step chain, buf / ring / counters byte-equal against ref."""
|
||||
_run_write(
|
||||
buf_pair=self.buf_pair,
|
||||
input_ids=[10, 20, 30, 40, 50],
|
||||
positions=[0, 1, 2, 3, 4],
|
||||
out_cache_loc=[0, 1, 2, 3, 4],
|
||||
)
|
||||
|
||||
def test_chain_link_byte_equal_5_step_hardcoded(self) -> None:
|
||||
"""5-step write chain with hand-computed splitmix64 expected fields per slot."""
|
||||
tokens = [101, 202, 303, 404, 505]
|
||||
positions = [0, 1, 2, 3, 4]
|
||||
out_cache_loc = [0, 1, 2, 3, 4]
|
||||
|
||||
# Step 1: compute the expected stored prev_hash sequence in pure Python via splitmix64.
|
||||
expected_prev_hashes_u64: list[int] = []
|
||||
running = splitmix64(consts.CANARY_CHAIN_ANCHOR)
|
||||
for token, position in zip(tokens, positions):
|
||||
expected_prev_hashes_u64.append(running)
|
||||
running = splitmix64_mix3(running, token, position)
|
||||
expected_prev_hashes_signed = [
|
||||
to_signed_int64(h) for h in expected_prev_hashes_u64
|
||||
]
|
||||
|
||||
_run_write(
|
||||
buf_pair=self.buf_pair,
|
||||
input_ids=tokens,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
)
|
||||
|
||||
# Step 2: verify every slot's stored 4 fields match the hardcoded expected sequence.
|
||||
for slot_idx, expected_token, expected_position, expected_prev_signed in zip(
|
||||
out_cache_loc, tokens, positions, expected_prev_hashes_signed
|
||||
):
|
||||
stored_token, stored_position, stored_prev_hash, stored_real_kv_hash = (
|
||||
read_slot_fields(canary_buf=self.buf_pair[0], slot_idx=slot_idx)
|
||||
)
|
||||
assert stored_token == expected_token
|
||||
assert stored_position == expected_position
|
||||
assert stored_prev_hash == expected_prev_signed
|
||||
assert stored_real_kv_hash == 0
|
||||
|
||||
|
||||
class TestMockMode:
|
||||
def setup_method(self) -> None:
|
||||
self.buf_pair = _make_default_buf_pair()
|
||||
|
||||
def test_mock_mode_off_ignores_expected(self) -> None:
|
||||
"""``enable_write_verify_inputs = OFF`` → expected tensors are ignored (we pass garbage to prove the kernel skips them)."""
|
||||
# Garbage expected tensors that, if the kernel mistakenly reads, would generate mismatches.
|
||||
cuda_log, _ = _run_write(
|
||||
buf_pair=self.buf_pair,
|
||||
input_ids=[1, 2, 3],
|
||||
positions=[0, 1, 2],
|
||||
out_cache_loc=[0, 1, 2],
|
||||
expected_input_tokens=_int32_tensor([999, 999, 999]),
|
||||
expected_input_positions=_int32_tensor([999, 999, 999]),
|
||||
)
|
||||
|
||||
assert int(cuda_log.write_index[0].item()) == 0
|
||||
|
||||
def test_mock_mode_on_match_no_violation(self) -> None:
|
||||
"""``enable_write_verify_inputs = ON`` and expected matches actual → no violation, chain advances."""
|
||||
input_ids = _int32_tensor([7, 8, 9])
|
||||
positions = _int32_tensor([0, 1, 2])
|
||||
cuda_log, _ = _run_write(
|
||||
buf_pair=self.buf_pair,
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=[0, 1, 2],
|
||||
enable_write_verify_inputs=True,
|
||||
expected_input_tokens=input_ids.clone(),
|
||||
expected_input_positions=positions.clone(),
|
||||
)
|
||||
|
||||
assert int(cuda_log.write_index[0].item()) == 0
|
||||
|
||||
def test_mock_mode_on_token_mismatch_records_violation(self) -> None:
|
||||
"""``enable_write_verify_inputs = ON`` token mismatch → violation recorded; chain advances on ACTUAL token."""
|
||||
cuda_log, _ = _run_write(
|
||||
buf_pair=self.buf_pair,
|
||||
input_ids=[42],
|
||||
positions=[0],
|
||||
out_cache_loc=[0],
|
||||
enable_write_verify_inputs=True,
|
||||
expected_input_tokens=_int32_tensor([99]),
|
||||
expected_input_positions=_int32_tensor([0]),
|
||||
)
|
||||
|
||||
fail_bits = int(
|
||||
cuda_log.ring[0, consts.VIOLATION_FIELD_FAIL_REASON_BITS].item()
|
||||
)
|
||||
assert_only_bits_set(fail_bits, consts.FailReason.WRITE_TOKEN_MISMATCH)
|
||||
# Chain advances on actual (42), not expected (99). Stored token should be 42.
|
||||
stored_token, _, _, _ = read_slot_fields(
|
||||
canary_buf=self.buf_pair[0], slot_idx=0
|
||||
)
|
||||
assert stored_token == 42
|
||||
|
||||
def test_mock_mode_on_position_mismatch_records_violation(self) -> None:
|
||||
"""``enable_write_verify_inputs = ON`` position mismatch → violation recorded; chain advances on ACTUAL position."""
|
||||
cuda_log, _ = _run_write(
|
||||
buf_pair=self.buf_pair,
|
||||
input_ids=[42],
|
||||
positions=[7],
|
||||
out_cache_loc=[0],
|
||||
enable_write_verify_inputs=True,
|
||||
expected_input_tokens=_int32_tensor([42]),
|
||||
expected_input_positions=_int32_tensor([0]),
|
||||
)
|
||||
|
||||
fail_bits = int(
|
||||
cuda_log.ring[0, consts.VIOLATION_FIELD_FAIL_REASON_BITS].item()
|
||||
)
|
||||
assert_only_bits_set(fail_bits, consts.FailReason.WRITE_POSITION_MISMATCH)
|
||||
_, stored_position, _, _ = read_slot_fields(
|
||||
canary_buf=self.buf_pair[0], slot_idx=0
|
||||
)
|
||||
assert stored_position == 7
|
||||
|
||||
def test_mock_mode_chain_advances_on_actual_not_expected(self) -> None:
|
||||
"""Expected differs from actual on every entry → downstream verify must NOT cascade chain errors."""
|
||||
cuda_buf = self.buf_pair[0]
|
||||
# Every actual differs from expected.
|
||||
cuda_log, _ = _run_write(
|
||||
buf_pair=self.buf_pair,
|
||||
input_ids=[10, 20, 30],
|
||||
positions=[0, 1, 2],
|
||||
out_cache_loc=[1, 2, 3],
|
||||
enable_write_verify_inputs=True,
|
||||
expected_input_tokens=_int32_tensor([999, 999, 999]),
|
||||
expected_input_positions=_int32_tensor([999, 999, 999]),
|
||||
)
|
||||
|
||||
# All 3 entries should fire a violation row.
|
||||
assert int(cuda_log.write_index[0].item()) == 3
|
||||
# Run a downstream verify — it must see no chain mismatch because chain advanced on actuals.
|
||||
verify_plan = make_verify_plan(
|
||||
slot_indices=[1, 2, 3],
|
||||
positions=[0, 1, 2],
|
||||
prev_slot_indices=[-1, 1, 2],
|
||||
device=_DEVICE,
|
||||
)
|
||||
verify_log = FakeViolationLog.allocate(device=_DEVICE)
|
||||
launch_canary_verify_kernel(
|
||||
context=VerifyOrWriteContext(
|
||||
canary_buf=cuda_buf,
|
||||
kernel_kind=CanaryLaunchTag.HEAD_K_FULL,
|
||||
violation_ring=verify_log.ring,
|
||||
violation_write_index=verify_log.write_index,
|
||||
slot_run_counter=verify_log.slot_run_counter,
|
||||
kernel_run_counter=verify_log.kernel_run_counter,
|
||||
enable_chain_position_assert=verify_log.enable_chain_position_assert,
|
||||
),
|
||||
plan=verify_plan,
|
||||
check_verify_expected_token=True,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
assert int(verify_log.write_index[0].item()) == 0
|
||||
|
||||
@pytest.mark.parametrize("bit_to_trigger", ["MOCK_TOKEN", "MOCK_POSITION"])
|
||||
@pytest.mark.parametrize("injection_position", ["head", "mid", "last"])
|
||||
def test_mock_violation_bit_injection_position_matrix(
|
||||
self,
|
||||
bit_to_trigger: str,
|
||||
injection_position: str,
|
||||
) -> None:
|
||||
"""Sweep injection_position x bit_to_trigger for write-kernel pseudo-mode fail-reason coverage."""
|
||||
slot_count = 5
|
||||
tokens = [10, 20, 30, 40, 50]
|
||||
positions = [0, 1, 2, 3, 4]
|
||||
out_cache_locs = [0, 1, 2, 3, 4]
|
||||
corruption_index = {"head": 0, "mid": 2, "last": 4}[injection_position]
|
||||
corrupt_slot = out_cache_locs[corruption_index]
|
||||
|
||||
expected_bit = {
|
||||
"MOCK_TOKEN": consts.FailReason.WRITE_TOKEN_MISMATCH,
|
||||
"MOCK_POSITION": consts.FailReason.WRITE_POSITION_MISMATCH,
|
||||
}[bit_to_trigger]
|
||||
|
||||
input_ids = _int32_tensor(tokens)
|
||||
positions_t = _int32_tensor(positions)
|
||||
out_cache_loc = _int32_tensor(out_cache_locs)
|
||||
|
||||
pseudo_tokens = input_ids.clone()
|
||||
pseudo_positions = positions_t.clone()
|
||||
|
||||
if bit_to_trigger == "MOCK_TOKEN":
|
||||
pseudo_tokens[corruption_index] = tokens[corruption_index] + 999
|
||||
else:
|
||||
pseudo_positions[corruption_index] = positions[corruption_index] + 99
|
||||
|
||||
cuda_log, ref_log = _run_write(
|
||||
buf_pair=self.buf_pair,
|
||||
input_ids=input_ids,
|
||||
positions=positions_t,
|
||||
out_cache_loc=out_cache_loc,
|
||||
write_offsets=[0, slot_count],
|
||||
enable_write_verify_inputs=True,
|
||||
expected_input_tokens=pseudo_tokens,
|
||||
expected_input_positions=pseudo_positions,
|
||||
assert_equal=False,
|
||||
)
|
||||
|
||||
found = False
|
||||
for row_idx in range(int(cuda_log.write_index[0].item())):
|
||||
fail_bits = int(
|
||||
cuda_log.ring[row_idx, consts.VIOLATION_FIELD_FAIL_REASON_BITS].item()
|
||||
)
|
||||
row_slot = int(cuda_log.ring[row_idx, 1].item())
|
||||
if (fail_bits & expected_bit) and row_slot == corrupt_slot:
|
||||
found = True
|
||||
break
|
||||
assert found, (
|
||||
f"expected bit {expected_bit:#x} at slot {corrupt_slot} not found in ring "
|
||||
f"(bit_to_trigger={bit_to_trigger} injection_position={injection_position})"
|
||||
)
|
||||
assert_canary_state_equal(log_a=cuda_log, log_b=ref_log)
|
||||
|
||||
|
||||
class TestSlotHandling:
|
||||
def setup_method(self) -> None:
|
||||
self.buf_pair = _make_default_buf_pair()
|
||||
|
||||
def test_negative_slot_skips_entry(self) -> None:
|
||||
"""``out_cache_loc[i] < 0`` → that entry is skipped: no buf write, no violation, no
|
||||
canary slot mutation, and no write slot_run_counter increment.
|
||||
Covers both SWA out-of-window (after caller-side LUT gather) and explicit padding intents.
|
||||
"""
|
||||
# Two entries: first writes to slot 4 normally; second has slot=-1 and must be skipped.
|
||||
cuda_log, _ = _run_write(
|
||||
buf_pair=self.buf_pair,
|
||||
input_ids=[42, 99],
|
||||
positions=[0, 1],
|
||||
out_cache_loc=[4, -1],
|
||||
)
|
||||
|
||||
stored_token, _, _, _ = read_slot_fields(
|
||||
canary_buf=self.buf_pair[0], slot_idx=4
|
||||
)
|
||||
assert stored_token == 42
|
||||
assert int(cuda_log.slot_run_counter.item()) == 1
|
||||
|
||||
def test_pre_translated_slot_writes_normally(self) -> None:
|
||||
"""``out_cache_loc[i] >= 0`` → the kernel writes to exactly that slot, with no LUT applied. This
|
||||
confirms the kernel is SWA-agnostic: SWA endpoints feed the same shape of input here after their
|
||||
host-side gather, so the contract is symmetric across FULL / SWA groups.
|
||||
"""
|
||||
# Slot 4 here could equally be a FULL-group raw out_cache_loc value, or the result of an SWA
|
||||
# endpoint's host gather. The kernel can't tell the difference and that's the point.
|
||||
_run_write(
|
||||
buf_pair=self.buf_pair,
|
||||
input_ids=[55],
|
||||
positions=[0],
|
||||
out_cache_loc=[4],
|
||||
)
|
||||
|
||||
stored_token, _, _, _ = read_slot_fields(
|
||||
canary_buf=self.buf_pair[0], slot_idx=4
|
||||
)
|
||||
assert stored_token == 55
|
||||
|
||||
def test_padding_block_skipped(self) -> None:
|
||||
"""``blockIdx.x >= write_num_valid_reqs[0]`` → block early-exits, no write to canary_buf."""
|
||||
# Allocate plan with req_capacity=4 but only declare 1 active req.
|
||||
_run_write(
|
||||
buf_pair=self.buf_pair,
|
||||
input_ids=[1],
|
||||
positions=[0],
|
||||
out_cache_loc=[0],
|
||||
req_capacity=4,
|
||||
)
|
||||
|
||||
# Only slot 0 should have been written; padding blocks 1..3 must not touch the buffer.
|
||||
stored_token, _, _, _ = read_slot_fields(
|
||||
canary_buf=self.buf_pair[0], slot_idx=0
|
||||
)
|
||||
assert stored_token == 1
|
||||
for slot_idx in (1, 2, 3):
|
||||
stored_token_other, _, _, _ = read_slot_fields(
|
||||
canary_buf=self.buf_pair[0], slot_idx=slot_idx
|
||||
)
|
||||
assert stored_token_other == 0
|
||||
|
||||
def test_write_skip_when_out_cache_loc_is_minus_one(self) -> None:
|
||||
"""out_cache_loc[i] = -1 → that entry's slot is untouched by write kernel."""
|
||||
cuda_buf = self.buf_pair[0]
|
||||
cuda_buf_before_slot_view = cuda_buf.view(torch.int64).clone()
|
||||
|
||||
_run_write(
|
||||
buf_pair=self.buf_pair,
|
||||
input_ids=[100, 200, 300],
|
||||
positions=[0, 1, 2],
|
||||
out_cache_loc=[5, -1, 7],
|
||||
)
|
||||
|
||||
after = cuda_buf.view(torch.int64)
|
||||
for slot in range(cuda_buf.shape[0]):
|
||||
if slot in (5, 7):
|
||||
continue
|
||||
assert torch.equal(
|
||||
after[slot], cuda_buf_before_slot_view[slot]
|
||||
), f"slot {slot} should not have been written"
|
||||
|
||||
def test_shrink_active_reqs_does_not_write_stale_slots(self) -> None:
|
||||
"""Run write with bs=3 plan after a bs=8 run on same buffer: stale slots from bs=8 stay intact."""
|
||||
cuda_buf = make_canary_buf(num_slots=32, slot_stride_bytes=32, device=_DEVICE)
|
||||
ref_buf = cuda_buf.clone()
|
||||
buf_pair = (cuda_buf, ref_buf)
|
||||
|
||||
big_slots = list(range(1, 9))
|
||||
_run_write(
|
||||
buf_pair=buf_pair,
|
||||
input_ids=list(range(100, 108)),
|
||||
positions=[0] * 8,
|
||||
out_cache_loc=big_slots,
|
||||
write_offsets=[0, 1, 2, 3, 4, 5, 6, 7, 8],
|
||||
seed_slot_indices=[-1] * 8,
|
||||
num_valid_reqs=8,
|
||||
assert_equal=False,
|
||||
)
|
||||
|
||||
untouched_snapshot = cuda_buf.view(torch.int64).clone()
|
||||
|
||||
small_slots = [20, 21, 22]
|
||||
_run_write(
|
||||
buf_pair=buf_pair,
|
||||
input_ids=[7, 8, 9],
|
||||
positions=[0, 0, 0],
|
||||
out_cache_loc=small_slots,
|
||||
write_offsets=[0, 1, 2, 3],
|
||||
seed_slot_indices=[-1, -1, -1],
|
||||
num_valid_reqs=3,
|
||||
assert_equal=False,
|
||||
)
|
||||
|
||||
after = cuda_buf.view(torch.int64)
|
||||
for slot in big_slots:
|
||||
assert torch.equal(
|
||||
after[slot], untouched_snapshot[slot]
|
||||
), f"slot {slot} from earlier bs=8 run was overwritten by bs=3 run"
|
||||
|
||||
|
||||
class TestRunCounter:
|
||||
def setup_method(self) -> None:
|
||||
self.buf_pair = _make_default_buf_pair()
|
||||
|
||||
def test_kernel_run_counter_per_call(self) -> None:
|
||||
"""``kernel_run_counter`` increments by 1 per call (even when ``write_num_valid_reqs == 0``)."""
|
||||
plan_pair = make_write_plan_pair(
|
||||
write_offsets=[0, 0],
|
||||
seed_slot_indices=[-1],
|
||||
num_valid_reqs=0,
|
||||
device=_DEVICE,
|
||||
)
|
||||
input_ids = _int32_tensor([0])
|
||||
positions = _int32_tensor([0])
|
||||
out_cache_loc = _int32_tensor([0])
|
||||
pseudo_tokens, pseudo_positions = dummy_pseudo_tensors(1)
|
||||
cuda_log, ref_log = make_log_pair(device=_DEVICE)
|
||||
|
||||
for _ in range(3):
|
||||
_run_both_write(
|
||||
cuda_canary_buf=self.buf_pair[0],
|
||||
ref_canary_buf=self.buf_pair[1],
|
||||
plan_cuda=plan_pair[0],
|
||||
plan_ref=plan_pair[1],
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
enable_write_verify_inputs=False,
|
||||
expected_input_tokens=pseudo_tokens,
|
||||
expected_input_positions=pseudo_positions,
|
||||
cuda_log=cuda_log,
|
||||
ref_log=ref_log,
|
||||
assert_equal=False,
|
||||
)
|
||||
|
||||
assert int(cuda_log.kernel_run_counter[0].item()) == 3
|
||||
assert_canary_state_equal(log_a=cuda_log, log_b=ref_log)
|
||||
|
||||
def test_slot_run_counter_sums_entries(self) -> None:
|
||||
"""``slot_run_counter`` += sum(entry_count) across all active reqs in this call."""
|
||||
cuda_log, _ = _run_write(
|
||||
buf_pair=self.buf_pair,
|
||||
input_ids=[1, 2, 3, 4, 5],
|
||||
positions=[0, 1, 0, 1, 2],
|
||||
out_cache_loc=[0, 1, 2, 3, 4],
|
||||
write_offsets=[0, 2, 5],
|
||||
seed_slot_indices=[-1, -1],
|
||||
num_valid_reqs=2,
|
||||
)
|
||||
|
||||
assert int(cuda_log.slot_run_counter[0].item()) == 5
|
||||
|
||||
|
||||
class TestMisc:
|
||||
def test_empty_plan_no_op(self) -> None:
|
||||
"""``write_num_valid_reqs = 0`` → no buf write, no slot_run_counter bump, only kernel_run_counter += 1."""
|
||||
cuda_log, _ = _run_write(
|
||||
buf_pair=_make_default_buf_pair(),
|
||||
input_ids=[0],
|
||||
positions=[0],
|
||||
out_cache_loc=[0],
|
||||
write_offsets=[0, 0],
|
||||
num_valid_reqs=0,
|
||||
req_capacity=4,
|
||||
)
|
||||
|
||||
assert int(cuda_log.write_index[0].item()) == 0
|
||||
assert int(cuda_log.slot_run_counter[0].item()) == 0
|
||||
assert int(cuda_log.kernel_run_counter[0].item()) == 1
|
||||
|
||||
def test_disabled_input_verify_does_not_deref_expected_inputs(self) -> None:
|
||||
"""``enable_write_verify_inputs=False`` must short-circuit before any expected_input_* dereference;
|
||||
poison those tensors with 0x7F7F7F7F and assert nothing records a violation."""
|
||||
cuda_buf = make_canary_buf(num_slots=8, slot_stride_bytes=32, device=_DEVICE)
|
||||
ref_buf = cuda_buf.clone()
|
||||
buf_pair = (cuda_buf, ref_buf)
|
||||
|
||||
garbage_expected_tokens = torch.full(
|
||||
(1,), 0x7F7F7F7F, dtype=torch.int64, device=_DEVICE
|
||||
)
|
||||
garbage_expected_positions = torch.full(
|
||||
(1,), 0x7F7F7F7F, dtype=torch.int64, device=_DEVICE
|
||||
)
|
||||
|
||||
cuda_log, ref_log = _run_write(
|
||||
buf_pair=buf_pair,
|
||||
input_ids=[42],
|
||||
positions=[0],
|
||||
out_cache_loc=[0],
|
||||
expected_input_tokens=garbage_expected_tokens,
|
||||
expected_input_positions=garbage_expected_positions,
|
||||
assert_equal=False,
|
||||
)
|
||||
|
||||
assert int(cuda_log.write_index[0].item()) == 0
|
||||
assert int(ref_log.write_index[0].item()) == 0
|
||||
|
||||
def test_disabled_assert_inputs_passes_none_to_cuda(self) -> None:
|
||||
"""Disabled input assertions pass None instead of dummy tensors."""
|
||||
canary_buf = torch.zeros(
|
||||
4, CANARY_SLOT_BYTES, dtype=torch.uint8, device=_DEVICE
|
||||
)
|
||||
plan = make_write_plan(
|
||||
write_offsets=[0, 0],
|
||||
seed_slot_indices=[-1],
|
||||
num_valid_reqs=0,
|
||||
device=_DEVICE,
|
||||
)
|
||||
input_ids = torch.zeros(1, dtype=torch.int64, device=_DEVICE)
|
||||
positions = torch.zeros(1, dtype=torch.int64, device=_DEVICE)
|
||||
out_cache_loc = torch.zeros(1, dtype=torch.int64, device=_DEVICE)
|
||||
log = FakeViolationLog.allocate(capacity=2, device=_DEVICE)
|
||||
context = VerifyOrWriteContext(
|
||||
canary_buf=canary_buf,
|
||||
kernel_kind=CanaryLaunchTag.HEAD_K_FULL,
|
||||
violation_ring=log.ring,
|
||||
violation_write_index=log.write_index,
|
||||
slot_run_counter=log.slot_run_counter,
|
||||
kernel_run_counter=log.kernel_run_counter,
|
||||
enable_chain_position_assert=log.enable_chain_position_assert,
|
||||
)
|
||||
module = _RecordingWriteModule()
|
||||
|
||||
with patch.object(write_module, "_jit_canary_write_module", lambda: module):
|
||||
launch_canary_write_kernel(
|
||||
context=context,
|
||||
plan=plan,
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
enable_write_input_assert=False,
|
||||
expected_input_tokens=None,
|
||||
expected_input_positions=None,
|
||||
)
|
||||
|
||||
assert len(module.calls) == 1
|
||||
call = module.calls[0]
|
||||
assert call[8] == 0
|
||||
assert call[9] is None
|
||||
assert call[10] is None
|
||||
|
||||
|
||||
class TestBoundarySweep:
|
||||
@pytest.mark.parametrize(
|
||||
"token_val",
|
||||
[0, 1, 0xFFFFFFFF, -1, 0x80000000, 0x7FFFFFFF],
|
||||
)
|
||||
def test_token_boundary_byte_equal_sweep(self, token_val: int) -> None:
|
||||
"""Sweep token boundary values; assert CUDA write vs ref buf + state byte-equal."""
|
||||
_run_write_single_slot_byte_equal(_WriteSingleSlotInput(token=token_val))
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"position_val",
|
||||
[0, 1, 127, 128, 129, 0x7FFFFFFF],
|
||||
)
|
||||
def test_position_boundary_byte_equal_sweep(self, position_val: int) -> None:
|
||||
"""Sweep position boundary values; assert CUDA write vs ref buf + state byte-equal."""
|
||||
_run_write_single_slot_byte_equal(_WriteSingleSlotInput(position=position_val))
|
||||
|
||||
|
||||
class TestPseudoMode:
|
||||
def setup_method(self) -> None:
|
||||
self.buf_pair = _make_default_buf_pair()
|
||||
|
||||
def test_pseudo_mode_on_catches_token_mismatch(self) -> None:
|
||||
"""enable_write_verify_inputs=ON + intentional token mismatch → WRITE_TOKEN_MISMATCH bit recorded."""
|
||||
input_ids = _int32_tensor([10, 20, 30, 40, 50])
|
||||
positions = _int32_tensor([0, 1, 2, 3, 4])
|
||||
out_cache_loc = _int32_tensor([1, 2, 3, 4, 5])
|
||||
pseudo_tokens = _int32_tensor([10, 20, 30, 999, 50])
|
||||
pseudo_positions = positions.clone()
|
||||
|
||||
cuda_log, _ = _run_write(
|
||||
buf_pair=self.buf_pair,
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
enable_write_verify_inputs=True,
|
||||
expected_input_tokens=pseudo_tokens,
|
||||
expected_input_positions=pseudo_positions,
|
||||
)
|
||||
assert int(cuda_log.write_index[0].item()) >= 1
|
||||
bits = int(cuda_log.ring[0, consts.VIOLATION_FIELD_FAIL_REASON_BITS].item())
|
||||
assert (
|
||||
bits & consts.FailReason.WRITE_TOKEN_MISMATCH
|
||||
), f"expected WRITE_TOKEN_MISMATCH bit, got {bits:#b}"
|
||||
|
||||
def test_pseudo_mode_off_skips_token_check(self) -> None:
|
||||
"""enable_write_verify_inputs=False makes the caller pass no expected-input tensors."""
|
||||
cuda_log, _ = _run_write(
|
||||
buf_pair=self.buf_pair,
|
||||
input_ids=[10, 20, 30],
|
||||
positions=[0, 1, 2],
|
||||
out_cache_loc=[1, 2, 3],
|
||||
expected_input_tokens=_int32_tensor([99, 99, 99]),
|
||||
expected_input_positions=_int32_tensor([99, 99, 99]),
|
||||
)
|
||||
assert int(cuda_log.write_index[0].item()) == 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
Reference in New Issue
Block a user