Add the KV-canary core: data layer, MHA KV-pool patcher, and per-forward runner (#26808)
This commit is contained in:
fzyzcjy
@@ -0,0 +1,31 @@
|
||||
import unittest
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.kv_canary.e2e_base import (
|
||||
_LONG_PROMPT_BODY,
|
||||
_UNIQUE_PROMPT_FIRST_CHARS,
|
||||
_make_unique_prompts,
|
||||
)
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=1, suite="base-b-test-cpu")
|
||||
|
||||
|
||||
class TestCanaryE2EBase(CustomTestCase):
|
||||
def test_make_unique_prompts_have_distinct_first_characters(self) -> None:
|
||||
"""Verify generated prompts use distinct first characters and all end with the shared long body."""
|
||||
prompts = _make_unique_prompts(8)
|
||||
|
||||
self.assertEqual(len({prompt[0] for prompt in prompts}), len(prompts))
|
||||
self.assertTrue(all(prompt.endswith(_LONG_PROMPT_BODY) for prompt in prompts))
|
||||
|
||||
def test_make_unique_prompts_rejects_more_prompts_than_distinct_first_characters(
|
||||
self,
|
||||
) -> None:
|
||||
"""Verify prompt generation rejects requests beyond the unique prefix budget."""
|
||||
with self.assertRaisesRegex(ValueError, "unique prompt count"):
|
||||
_make_unique_prompts(len(_UNIQUE_PROMPT_FIRST_CHARS) + 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,27 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.kv_canary.verify import CANARY_SLOT_BYTES
|
||||
from sglang.srt.kv_canary.pool_patcher.buffer_alloc import alloc_canary_buf
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.kv_canary.fixtures import DEFAULT_DEVICE
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=10, stage="extra-a", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
class TestAllocCanaryBuf(CustomTestCase):
|
||||
def test_alloc_canary_buf_shape_and_dtype(self) -> None:
|
||||
"""Verify alloc_canary_buf returns a zeroed uint8 buffer of [num_slots, CANARY_SLOT_BYTES]."""
|
||||
buf = alloc_canary_buf(num_slots=8, device=DEFAULT_DEVICE)
|
||||
self.assertEqual(buf.shape, (8, CANARY_SLOT_BYTES))
|
||||
self.assertEqual(buf.dtype, torch.uint8)
|
||||
self.assertEqual(buf.device.type, DEFAULT_DEVICE.type)
|
||||
self.assertTrue(torch.equal(buf, torch.zeros_like(buf)))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,89 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.kv_canary.capacities import CanaryLaunchCapacities
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=45, stage="extra-a", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
class TestComputeLaunchCapacities(CustomTestCase):
|
||||
@staticmethod
|
||||
def _make_server_args(*, max_bs: int) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
cuda_graph_max_bs=max_bs,
|
||||
speculative_num_draft_tokens=0,
|
||||
chunked_prefill_size=None,
|
||||
max_prefill_tokens=128,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _from_args(
|
||||
*,
|
||||
max_bs: int,
|
||||
max_seq_len: int,
|
||||
max_total_num_tokens: int | None = None,
|
||||
) -> CanaryLaunchCapacities:
|
||||
if max_total_num_tokens is None:
|
||||
max_total_num_tokens = max_bs * max_seq_len
|
||||
return CanaryLaunchCapacities.from_args(
|
||||
server_args=TestComputeLaunchCapacities._make_server_args(max_bs=max_bs),
|
||||
req_to_token_pool_size=max_bs,
|
||||
max_seq_len_per_req=max_seq_len,
|
||||
pool_slot_count=max_total_num_tokens,
|
||||
)
|
||||
|
||||
def test_per_forward_verify_capacity_covers_multi_req_prefix_sum(self) -> None:
|
||||
"""Verify per-forward verify capacity equals max_total_num_tokens * 3."""
|
||||
max_bs = 8
|
||||
max_seq_len = 64
|
||||
max_total_num_tokens = 1024
|
||||
capacities = self._from_args(
|
||||
max_bs=max_bs,
|
||||
max_seq_len=max_seq_len,
|
||||
max_total_num_tokens=max_total_num_tokens,
|
||||
)
|
||||
self.assertEqual(
|
||||
capacities.per_forward_verify_capacity,
|
||||
int(max_total_num_tokens * 3),
|
||||
)
|
||||
|
||||
def test_from_args_treats_missing_speculative_draft_tokens_as_zero(self) -> None:
|
||||
"""per_forward_write_entry_capacity is floored by max_prefill_tokens when batch * tokens_per_bs is smaller."""
|
||||
server_args = self._make_server_args(max_bs=2)
|
||||
server_args.speculative_num_draft_tokens = None
|
||||
|
||||
capacities = CanaryLaunchCapacities.from_args(
|
||||
server_args=server_args,
|
||||
req_to_token_pool_size=2,
|
||||
max_seq_len_per_req=32,
|
||||
pool_slot_count=64,
|
||||
)
|
||||
|
||||
self.assertEqual(capacities.per_forward_write_entry_capacity, 128)
|
||||
|
||||
def test_manual_capacities_reject_non_positive_fields(self) -> None:
|
||||
"""Verify manual launch capacities fail instead of being clamped."""
|
||||
with self.assertRaisesRegex(ValueError, "per_forward_verify_capacity"):
|
||||
CanaryLaunchCapacities(
|
||||
per_forward_verify_capacity=0,
|
||||
per_forward_write_req_capacity=1,
|
||||
per_forward_write_entry_capacity=1,
|
||||
)
|
||||
|
||||
def test_from_args_rejects_empty_pool_capacity(self) -> None:
|
||||
"""Verify derived launch capacities reject invalid pool sizing."""
|
||||
with self.assertRaisesRegex(ValueError, "pool_slot_count"):
|
||||
CanaryLaunchCapacities.from_args(
|
||||
server_args=self._make_server_args(max_bs=1),
|
||||
req_to_token_pool_size=1,
|
||||
max_seq_len_per_req=1,
|
||||
pool_slot_count=0,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,241 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.kv_canary.verify import (
|
||||
CANARY_SLOT_BYTES,
|
||||
CanaryLaunchTag,
|
||||
VerifyPlan,
|
||||
)
|
||||
from sglang.jit_kernel.kv_canary.write import WritePlan
|
||||
from sglang.srt.kv_canary import endpoint as endpoint_module
|
||||
from sglang.srt.kv_canary.endpoint import (
|
||||
CanaryEndpoint,
|
||||
)
|
||||
from sglang.srt.kv_canary.expected_inputs import ExpectedInputs
|
||||
from sglang.srt.kv_canary.state import (
|
||||
ViolationLog,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.kv_canary.fixtures import DEFAULT_DEVICE
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=20, stage="extra-a", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
def _make_endpoint(*, device, kernel_kind=CanaryLaunchTag.HEAD_K_FULL, swa_lut=None):
|
||||
canary_buf = torch.zeros(4, CANARY_SLOT_BYTES, dtype=torch.uint8, device=device)
|
||||
slot_view = torch.zeros(1, dtype=torch.int64, device=device)
|
||||
kernel_view = torch.zeros(1, dtype=torch.int64, device=device)
|
||||
enable_chain_position_assert = torch.ones(1, dtype=torch.int32, device=device)
|
||||
return CanaryEndpoint(
|
||||
kernel_kind=kernel_kind,
|
||||
canary_buf=canary_buf,
|
||||
full_to_swa_index_mapping=swa_lut,
|
||||
slot_run_counter_view=slot_view,
|
||||
kernel_run_counter_view=kernel_view,
|
||||
enable_chain_position_assert=enable_chain_position_assert,
|
||||
)
|
||||
|
||||
|
||||
def _make_kernel_args(device):
|
||||
verify_plan = VerifyPlan.allocate(verify_capacity=1, device=device)
|
||||
write_plan = WritePlan.allocate(write_req_capacity=1, device=device)
|
||||
log = ViolationLog.allocate(ring_capacity=2, device=device)
|
||||
return SimpleNamespace(
|
||||
verify_plan=verify_plan,
|
||||
write_plan=write_plan,
|
||||
violation_log=log,
|
||||
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),
|
||||
enable_write_input_assert=False,
|
||||
enable_verify_token_assert=False,
|
||||
expected_inputs=ExpectedInputs.allocate(capacity=1, device=device),
|
||||
)
|
||||
|
||||
|
||||
class TestSelfUnitEndpoint(CustomTestCase):
|
||||
def setUp(self):
|
||||
self.device = DEFAULT_DEVICE
|
||||
|
||||
def test_launch_per_forward_passes_kernel_kind(self):
|
||||
"""Verify per-forward launch passes the endpoint kernel kind."""
|
||||
captured: list[tuple[str, CanaryLaunchTag]] = []
|
||||
with patch.object(
|
||||
endpoint_module,
|
||||
"launch_canary_verify_kernel",
|
||||
lambda **kwargs: captured.append(("verify", kwargs["context"].kernel_kind)),
|
||||
), patch.object(
|
||||
endpoint_module,
|
||||
"launch_canary_write_kernel",
|
||||
lambda **kwargs: captured.append(("write", kwargs["context"].kernel_kind)),
|
||||
):
|
||||
ep = _make_endpoint(
|
||||
device=self.device, kernel_kind=CanaryLaunchTag.TAIL_V_SWA
|
||||
)
|
||||
args = _make_kernel_args(self.device)
|
||||
ep.launch_per_forward(
|
||||
verify_plan=args.verify_plan,
|
||||
write_plan=args.write_plan,
|
||||
input_ids=args.input_ids,
|
||||
positions=args.positions,
|
||||
out_cache_loc=args.out_cache_loc,
|
||||
enable_write_input_assert=args.enable_write_input_assert,
|
||||
enable_verify_token_assert=args.enable_verify_token_assert,
|
||||
expected_inputs=args.expected_inputs,
|
||||
violation_log=args.violation_log,
|
||||
)
|
||||
self.assertIn(("verify", CanaryLaunchTag.TAIL_V_SWA), captured)
|
||||
self.assertIn(("write", CanaryLaunchTag.TAIL_V_SWA), captured)
|
||||
|
||||
def test_endpoint_shares_violation_log_across_launches(self):
|
||||
"""Verify endpoints can reuse the same violation log."""
|
||||
captured_rings: list[int] = []
|
||||
with patch.object(
|
||||
endpoint_module,
|
||||
"launch_canary_verify_kernel",
|
||||
lambda **kwargs: captured_rings.append(
|
||||
kwargs["context"].violation_ring.data_ptr()
|
||||
),
|
||||
), patch.object(
|
||||
endpoint_module,
|
||||
"launch_canary_write_kernel",
|
||||
lambda **kwargs: None,
|
||||
):
|
||||
shared_log = ViolationLog.allocate(ring_capacity=2, device=self.device)
|
||||
ep_a = _make_endpoint(
|
||||
device=self.device, kernel_kind=CanaryLaunchTag.HEAD_K_FULL
|
||||
)
|
||||
ep_b = _make_endpoint(
|
||||
device=self.device, kernel_kind=CanaryLaunchTag.HEAD_V_FULL
|
||||
)
|
||||
|
||||
args = _make_kernel_args(self.device)
|
||||
ep_a.launch_per_forward(
|
||||
verify_plan=args.verify_plan,
|
||||
write_plan=args.write_plan,
|
||||
input_ids=args.input_ids,
|
||||
positions=args.positions,
|
||||
out_cache_loc=args.out_cache_loc,
|
||||
enable_write_input_assert=args.enable_write_input_assert,
|
||||
enable_verify_token_assert=args.enable_verify_token_assert,
|
||||
expected_inputs=args.expected_inputs,
|
||||
violation_log=shared_log,
|
||||
)
|
||||
ep_b.launch_per_forward(
|
||||
verify_plan=args.verify_plan,
|
||||
write_plan=args.write_plan,
|
||||
input_ids=args.input_ids,
|
||||
positions=args.positions,
|
||||
out_cache_loc=args.out_cache_loc,
|
||||
enable_write_input_assert=args.enable_write_input_assert,
|
||||
enable_verify_token_assert=args.enable_verify_token_assert,
|
||||
expected_inputs=args.expected_inputs,
|
||||
violation_log=shared_log,
|
||||
)
|
||||
self.assertEqual(captured_rings[0], captured_rings[1])
|
||||
self.assertEqual(captured_rings[0], shared_log.violation_ring.data_ptr())
|
||||
|
||||
def test_swa_endpoint_pre_translates_out_cache_loc(self):
|
||||
"""Verify SWA endpoints translate cache locations before write launch."""
|
||||
captured: list[torch.Tensor] = []
|
||||
with patch.object(
|
||||
endpoint_module, "launch_canary_verify_kernel", lambda **kwargs: None
|
||||
), patch.object(
|
||||
endpoint_module,
|
||||
"launch_canary_write_kernel",
|
||||
lambda **kwargs: captured.append(kwargs["out_cache_loc"]),
|
||||
):
|
||||
# LUT maps full slot i → swa slot (i + 100) so we can verify the gather happened.
|
||||
lut = torch.arange(8, dtype=torch.int64, device=self.device) + 100
|
||||
swa_ep = _make_endpoint(
|
||||
device=self.device,
|
||||
kernel_kind=CanaryLaunchTag.HEAD_K_SWA,
|
||||
swa_lut=lut,
|
||||
)
|
||||
full_ep = _make_endpoint(
|
||||
device=self.device,
|
||||
kernel_kind=CanaryLaunchTag.HEAD_K_FULL,
|
||||
swa_lut=None,
|
||||
)
|
||||
args = _make_kernel_args(self.device)
|
||||
|
||||
swa_ep.launch_per_forward(
|
||||
verify_plan=args.verify_plan,
|
||||
write_plan=args.write_plan,
|
||||
input_ids=args.input_ids,
|
||||
positions=args.positions,
|
||||
out_cache_loc=args.out_cache_loc,
|
||||
enable_write_input_assert=args.enable_write_input_assert,
|
||||
enable_verify_token_assert=args.enable_verify_token_assert,
|
||||
expected_inputs=args.expected_inputs,
|
||||
violation_log=args.violation_log,
|
||||
)
|
||||
full_ep.launch_per_forward(
|
||||
verify_plan=args.verify_plan,
|
||||
write_plan=args.write_plan,
|
||||
input_ids=args.input_ids,
|
||||
positions=args.positions,
|
||||
out_cache_loc=args.out_cache_loc,
|
||||
enable_write_input_assert=args.enable_write_input_assert,
|
||||
enable_verify_token_assert=args.enable_verify_token_assert,
|
||||
expected_inputs=args.expected_inputs,
|
||||
violation_log=args.violation_log,
|
||||
)
|
||||
# SWA call: out_cache_loc was rewritten via lut gather (so identity-shifted by +100 here).
|
||||
expected_swa = lut[args.out_cache_loc]
|
||||
self.assertTrue(torch.equal(captured[0], expected_swa))
|
||||
# FULL call: out_cache_loc keeps the same values and dtype.
|
||||
self.assertIs(captured[1], args.out_cache_loc)
|
||||
|
||||
def test_swa_endpoint_trailing_sentinel_row_yields_skip(self):
|
||||
"""Verify SWA sentinel cache rows become write-skip markers."""
|
||||
captured: list[torch.Tensor] = []
|
||||
with patch.object(
|
||||
endpoint_module, "launch_canary_verify_kernel", lambda **kwargs: None
|
||||
), patch.object(
|
||||
endpoint_module,
|
||||
"launch_canary_write_kernel",
|
||||
lambda **kwargs: captured.append(kwargs["out_cache_loc"]),
|
||||
):
|
||||
# 8 in-window rows + 1 trailing sentinel row at index 8.
|
||||
lut = torch.arange(8, dtype=torch.int64, device=self.device)
|
||||
lut = torch.cat(
|
||||
[lut, torch.tensor([-1], dtype=torch.int64, device=self.device)]
|
||||
)
|
||||
swa_ep = _make_endpoint(
|
||||
device=self.device, kernel_kind=CanaryLaunchTag.HEAD_K_SWA, swa_lut=lut
|
||||
)
|
||||
args = _make_kernel_args(self.device)
|
||||
# Point out_cache_loc at the trailing-sentinel-row index — this is how sglang signals
|
||||
# "this token is out-of-window for the SWA group" pre-cleanup, and the new host gather must
|
||||
# produce -1 here.
|
||||
args.out_cache_loc.fill_(8)
|
||||
|
||||
swa_ep.launch_per_forward(
|
||||
verify_plan=args.verify_plan,
|
||||
write_plan=args.write_plan,
|
||||
input_ids=args.input_ids,
|
||||
positions=args.positions,
|
||||
out_cache_loc=args.out_cache_loc,
|
||||
enable_write_input_assert=args.enable_write_input_assert,
|
||||
enable_verify_token_assert=args.enable_verify_token_assert,
|
||||
expected_inputs=args.expected_inputs,
|
||||
violation_log=args.violation_log,
|
||||
)
|
||||
|
||||
self.assertTrue(
|
||||
torch.equal(
|
||||
captured[0],
|
||||
torch.tensor([-1], dtype=torch.int64, device=self.device),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,180 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
from typing import cast
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.kv_canary.runner.future_tensor import FutureTensors
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=20, stage="extra-a", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
class _FakeEvent:
|
||||
def __init__(self) -> None:
|
||||
self.synchronize_count = 0
|
||||
|
||||
def synchronize(self) -> None:
|
||||
self.synchronize_count += 1
|
||||
|
||||
|
||||
class TestFutureTensors(CustomTestCase):
|
||||
def test_cuda_stage_then_wait_returns_host_copy(self) -> None:
|
||||
"""Verify staged CUDA tensors are copied back on wait."""
|
||||
device = torch.device("cuda")
|
||||
alt_stream = torch.cuda.Stream(device=device)
|
||||
default_stream = torch.cuda.current_stream(device)
|
||||
self.assertNotEqual(alt_stream.stream_id, default_stream.stream_id)
|
||||
|
||||
src_first = torch.tensor([41], dtype=torch.int32, device=device)
|
||||
future_first = FutureTensors.device_to_host(
|
||||
xs_device=src_first, d2h_stream=alt_stream
|
||||
)
|
||||
result_first = future_first.wait()
|
||||
self.assertEqual(int(result_first.item()), 41)
|
||||
|
||||
src_second = torch.tensor([97], dtype=torch.int32, device=device)
|
||||
future_second = FutureTensors.device_to_host(
|
||||
xs_device=src_second, d2h_stream=alt_stream
|
||||
)
|
||||
result_second = future_second.wait()
|
||||
self.assertEqual(int(result_second.item()), 97)
|
||||
|
||||
def test_cuda_pinned_when_stream_is_provided(self) -> None:
|
||||
"""Verify CUDA staging uses pinned host memory with a stream."""
|
||||
device = torch.device("cuda")
|
||||
alt_stream = torch.cuda.Stream(device=device)
|
||||
src = torch.tensor([5], dtype=torch.int32, device=device)
|
||||
future = FutureTensors.device_to_host(xs_device=src, d2h_stream=alt_stream)
|
||||
staged_tensors = [
|
||||
v for v in future._data.values() if isinstance(v, torch.Tensor)
|
||||
]
|
||||
self.assertTrue(staged_tensors)
|
||||
self.assertTrue(all(t.is_pinned() for t in staged_tensors))
|
||||
self.assertEqual(int(future.wait().item()), 5)
|
||||
|
||||
def test_cuda_each_call_allocates_fresh_host(self) -> None:
|
||||
"""Verify each CUDA staging call owns a fresh host buffer."""
|
||||
device = torch.device("cuda")
|
||||
alt_stream = torch.cuda.Stream(device=device)
|
||||
src_a = torch.tensor([13], dtype=torch.int32, device=device)
|
||||
src_b = torch.tensor([29], dtype=torch.int32, device=device)
|
||||
future_a = FutureTensors.device_to_host(xs_device=src_a, d2h_stream=alt_stream)
|
||||
future_b = FutureTensors.device_to_host(xs_device=src_b, d2h_stream=alt_stream)
|
||||
ptrs_a = {
|
||||
v.data_ptr() for v in future_a._data.values() if isinstance(v, torch.Tensor)
|
||||
}
|
||||
ptrs_b = {
|
||||
v.data_ptr() for v in future_b._data.values() if isinstance(v, torch.Tensor)
|
||||
}
|
||||
self.assertTrue(ptrs_a and ptrs_b)
|
||||
self.assertFalse(ptrs_a & ptrs_b)
|
||||
self.assertEqual(int(future_a.wait().item()), 13)
|
||||
self.assertEqual(int(future_b.wait().item()), 29)
|
||||
|
||||
def test_dict_of_all_tensors_roundtrip(self) -> None:
|
||||
"""Verify a dict of multiple tensors round-trips entry-by-entry."""
|
||||
device = torch.device("cuda")
|
||||
stream = torch.cuda.Stream(device=device)
|
||||
src = {
|
||||
"x": torch.tensor([11, 22], dtype=torch.int64, device=device),
|
||||
"y": torch.tensor([99], dtype=torch.int32, device=device),
|
||||
}
|
||||
future = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream)
|
||||
out = future.wait()
|
||||
self.assertIsInstance(out, dict)
|
||||
self.assertEqual(out["x"].tolist(), [11, 22])
|
||||
self.assertEqual(int(out["y"].item()), 99)
|
||||
self.assertTrue(out["x"].is_pinned())
|
||||
self.assertTrue(out["y"].is_pinned())
|
||||
|
||||
def test_dict_mixes_tensor_and_passthrough(self) -> None:
|
||||
"""Verify non-tensor dict entries ride through verbatim alongside staging."""
|
||||
device = torch.device("cuda")
|
||||
stream = torch.cuda.Stream(device=device)
|
||||
sentinel_obj = {"nested": [1, 2, 3]}
|
||||
src = {
|
||||
"step": 42,
|
||||
"label": "decode",
|
||||
"extra": sentinel_obj,
|
||||
"counter": torch.tensor([7], dtype=torch.int32, device=device),
|
||||
}
|
||||
future = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream)
|
||||
out = future.wait()
|
||||
self.assertEqual(out["step"], 42)
|
||||
self.assertEqual(out["label"], "decode")
|
||||
# Identity (not deep-copy) — callers can rely on shared mutable references.
|
||||
self.assertIs(out["extra"], sentinel_obj)
|
||||
self.assertEqual(int(out["counter"].item()), 7)
|
||||
self.assertTrue(out["counter"].is_pinned())
|
||||
|
||||
def test_dict_passthrough_preserves_tensor_value(self) -> None:
|
||||
"""Verify tensors share device memory but non-tensor types are not staged."""
|
||||
device = torch.device("cuda")
|
||||
stream = torch.cuda.Stream(device=device)
|
||||
src_tensor = torch.tensor([3], dtype=torch.int32, device=device)
|
||||
src = {"step": 100, "buf": src_tensor}
|
||||
future = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream)
|
||||
out = future.wait()
|
||||
# Tensor is staged to a fresh pinned-host buffer (different storage from src).
|
||||
self.assertNotEqual(out["buf"].data_ptr(), src_tensor.data_ptr())
|
||||
self.assertTrue(out["buf"].is_pinned())
|
||||
# Non-tensor passes through with no copy.
|
||||
self.assertEqual(out["step"], 100)
|
||||
self.assertIsInstance(out["step"], int)
|
||||
|
||||
def test_dict_without_tensor_raises(self) -> None:
|
||||
"""Verify a tensor-less dict raises (no device to anchor the d2h sync)."""
|
||||
device = torch.device("cuda")
|
||||
stream = torch.cuda.Stream(device=device)
|
||||
with self.assertRaises(ValueError):
|
||||
FutureTensors.device_to_host(
|
||||
xs_device={"step": 0, "label": "decode"}, d2h_stream=stream
|
||||
)
|
||||
|
||||
def test_wait_called_twice_raises(self) -> None:
|
||||
"""Verify wait() after the first drain raises (state cleared)."""
|
||||
device = torch.device("cuda")
|
||||
stream = torch.cuda.Stream(device=device)
|
||||
src = torch.tensor([3], dtype=torch.int32, device=device)
|
||||
future = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream)
|
||||
self.assertEqual(int(future.wait().item()), 3)
|
||||
with self.assertRaises(RuntimeError):
|
||||
future.wait()
|
||||
|
||||
def test_wait_clears_fields_and_rejects_second_wait(self) -> None:
|
||||
"""Verify wait() syncs the event exactly once and clears internal state."""
|
||||
tensor = torch.tensor([1, 2, 3])
|
||||
event = _FakeEvent()
|
||||
future = FutureTensors(
|
||||
_data={"x": tensor}, _event=cast(torch.cuda.Event, event)
|
||||
)
|
||||
|
||||
result = future.wait()
|
||||
self.assertIs(result["x"], tensor)
|
||||
self.assertEqual(event.synchronize_count, 1)
|
||||
self.assertIsNone(future._data)
|
||||
self.assertIsNone(future._event)
|
||||
|
||||
with self.assertRaisesRegex(RuntimeError, "called more than once"):
|
||||
future.wait()
|
||||
# Failed wait must not re-trigger event.synchronize.
|
||||
self.assertEqual(event.synchronize_count, 1)
|
||||
|
||||
def test_dict_anchor_picked_from_first_tensor(self) -> None:
|
||||
"""Verify staging works when the first key is a non-tensor (anchor must scan)."""
|
||||
device = torch.device("cuda")
|
||||
stream = torch.cuda.Stream(device=device)
|
||||
src = {
|
||||
"step": 5,
|
||||
"buf": torch.tensor([17], dtype=torch.int32, device=device),
|
||||
}
|
||||
out = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream).wait()
|
||||
self.assertEqual(out["step"], 5)
|
||||
self.assertEqual(int(out["buf"].item()), 17)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,128 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.kv_canary.plan_input import PlanInput
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.kv_canary.fixtures import (
|
||||
DEFAULT_DEVICE,
|
||||
make_forward_batch,
|
||||
)
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=30, stage="extra-a", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
def _make_static_plan_input(*, bs_capacity: int, device) -> PlanInput:
|
||||
return PlanInput(
|
||||
req_pool_indices=torch.zeros(bs_capacity, dtype=torch.int64, device=device),
|
||||
prefix_lens=torch.zeros(bs_capacity, dtype=torch.int64, device=device),
|
||||
extend_seq_lens=torch.zeros(bs_capacity, dtype=torch.int64, device=device),
|
||||
req_to_verify_expected_tokens_valid_lens=torch.zeros(
|
||||
bs_capacity, dtype=torch.int64, device=device
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class TestSelfUnitPlanInput(CustomTestCase):
|
||||
def setUp(self):
|
||||
self.device = DEFAULT_DEVICE
|
||||
|
||||
def test_plan_input_fill_from_forward_batch_extend(self):
|
||||
"""Verify extend batches populate per-forward plan inputs."""
|
||||
fb = make_forward_batch(
|
||||
self.device,
|
||||
req_pool_indices=torch.tensor(
|
||||
[1, 2], dtype=torch.int64, device=self.device
|
||||
),
|
||||
seq_lens=torch.tensor([10, 12], dtype=torch.int32, device=self.device),
|
||||
extend_prefix_lens=torch.tensor(
|
||||
[3, 5], dtype=torch.int32, device=self.device
|
||||
),
|
||||
extend_seq_lens=torch.tensor([7, 7], dtype=torch.int32, device=self.device),
|
||||
is_extend=True,
|
||||
)
|
||||
plan = _make_static_plan_input(bs_capacity=4, device=self.device)
|
||||
plan.fill_from_forward_batch(forward_batch=fb)
|
||||
self.assertEqual(plan.req_pool_indices[:2].tolist(), [1, 2])
|
||||
self.assertEqual(plan.req_pool_indices[2:].tolist(), [0, 0])
|
||||
self.assertEqual(plan.prefix_lens[:2].tolist(), [3, 5])
|
||||
self.assertEqual(plan.extend_seq_lens[:2].tolist(), [7, 7])
|
||||
self.assertEqual(plan.prefix_lens.dtype, torch.int64)
|
||||
self.assertEqual(plan.extend_seq_lens.dtype, torch.int64)
|
||||
|
||||
def test_plan_input_fill_from_forward_batch_target_verify(self):
|
||||
"""Verify target-verify batches derive draft verification spans."""
|
||||
# TARGET_VERIFY writes spec_info.draft_token_num positions per req starting at the
|
||||
# current (un-bumped) seq_lens; extend_prefix_lens / extend_seq_lens are deliberately
|
||||
# NOT supplied because init_new does not populate them for this mode.
|
||||
spec_info = SimpleNamespace(draft_token_num=4)
|
||||
fb = make_forward_batch(
|
||||
self.device,
|
||||
req_pool_indices=torch.tensor(
|
||||
[6, 7], dtype=torch.int64, device=self.device
|
||||
),
|
||||
seq_lens=torch.tensor([10, 14], dtype=torch.int32, device=self.device),
|
||||
is_target_verify=True,
|
||||
spec_info=spec_info,
|
||||
)
|
||||
plan = _make_static_plan_input(bs_capacity=4, device=self.device)
|
||||
plan.fill_from_forward_batch(forward_batch=fb)
|
||||
self.assertEqual(plan.prefix_lens[:2].tolist(), [10, 14])
|
||||
self.assertEqual(plan.extend_seq_lens[:2].tolist(), [4, 4])
|
||||
|
||||
def test_plan_input_fill_from_forward_batch_draft_extend_v2(self):
|
||||
"""Verify draft-extend-v2 batches derive prefix lengths from sequence lengths."""
|
||||
# DRAFT_EXTEND_V2 has seq_lens already bumped by the per-req draft refill length;
|
||||
# extend_prefix_lens is intentionally absent (cuda-graph replay does not set it). The
|
||||
# builder must derive prefix as seq_lens - extend_seq_lens.
|
||||
fb = make_forward_batch(
|
||||
self.device,
|
||||
req_pool_indices=torch.tensor(
|
||||
[4, 5], dtype=torch.int64, device=self.device
|
||||
),
|
||||
seq_lens=torch.tensor([14, 18], dtype=torch.int32, device=self.device),
|
||||
extend_seq_lens=torch.tensor([4, 4], dtype=torch.int32, device=self.device),
|
||||
is_draft_extend_v2=True,
|
||||
)
|
||||
plan = _make_static_plan_input(bs_capacity=4, device=self.device)
|
||||
plan.fill_from_forward_batch(forward_batch=fb)
|
||||
self.assertEqual(plan.prefix_lens[:2].tolist(), [10, 14])
|
||||
self.assertEqual(plan.extend_seq_lens[:2].tolist(), [4, 4])
|
||||
|
||||
def test_plan_input_fill_from_forward_batch_decode(self):
|
||||
"""Verify decode batches populate one-token verification spans."""
|
||||
fb = make_forward_batch(
|
||||
self.device,
|
||||
req_pool_indices=torch.tensor(
|
||||
[1, 2, 3], dtype=torch.int64, device=self.device
|
||||
),
|
||||
seq_lens=torch.tensor([4, 7, 1], dtype=torch.int32, device=self.device),
|
||||
is_extend=False,
|
||||
)
|
||||
plan = _make_static_plan_input(bs_capacity=4, device=self.device)
|
||||
plan.fill_from_forward_batch(forward_batch=fb)
|
||||
self.assertEqual(plan.prefix_lens[:3].tolist(), [3, 6, 0])
|
||||
self.assertEqual(plan.extend_seq_lens[:3].tolist(), [1, 1, 1])
|
||||
|
||||
def test_plan_input_padding_dummy_sentinel(self):
|
||||
"""Verify padding sentinel rows remain valid plan input entries."""
|
||||
fb = make_forward_batch(
|
||||
self.device,
|
||||
req_pool_indices=torch.tensor(
|
||||
[0, 5, 0], dtype=torch.int64, device=self.device
|
||||
),
|
||||
seq_lens=torch.tensor([0, 3, 0], dtype=torch.int32, device=self.device),
|
||||
is_extend=False,
|
||||
)
|
||||
plan = _make_static_plan_input(bs_capacity=4, device=self.device)
|
||||
plan.fill_from_forward_batch(forward_batch=fb)
|
||||
self.assertEqual(plan.req_pool_indices[:3].tolist(), [0, 5, 0])
|
||||
self.assertEqual(plan.req_pool_indices.dtype, torch.int64)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,116 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
|
||||
from sglang.jit_kernel.kv_canary.verify import CANARY_SLOT_BYTES
|
||||
from sglang.srt.kv_canary.buffer_group import PoolKind
|
||||
from sglang.srt.kv_canary.pool_patcher.api import attach_canary_buffers
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.kv_canary.fixtures import (
|
||||
DEFAULT_DEVICE,
|
||||
make_base_config,
|
||||
make_mha_pool,
|
||||
)
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=45, stage="extra-a", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
class PoolPatcherHelper:
|
||||
def setUp(self):
|
||||
self.device = DEFAULT_DEVICE
|
||||
self.config = make_base_config()
|
||||
|
||||
|
||||
class TestAttachCanaryBuffers(PoolPatcherHelper, CustomTestCase):
|
||||
def test_canary_buffer_group_allocate_full_only(self):
|
||||
"""Verify MHA pools allocate only full canary buffers."""
|
||||
pool = make_mha_pool(self.device, num_slots=16, dim=8, layer_num=2)
|
||||
groups_tuple = attach_canary_buffers(
|
||||
pool=pool,
|
||||
config=self.config,
|
||||
device=self.device,
|
||||
kv_token_id_vs_position_offset=0,
|
||||
)
|
||||
groups = {g.kind: g for g in groups_tuple}
|
||||
self.assertEqual(set(groups.keys()), {PoolKind.FULL})
|
||||
group = groups[PoolKind.FULL]
|
||||
self.assertEqual(group.k_head.shape, (16, CANARY_SLOT_BYTES))
|
||||
self.assertEqual(group.k_tail.shape, (16, CANARY_SLOT_BYTES))
|
||||
self.assertIsNotNone(group.v_head)
|
||||
self.assertIsNotNone(group.v_tail)
|
||||
self.assertEqual(group.v_head.shape, (16, CANARY_SLOT_BYTES))
|
||||
|
||||
|
||||
class TestPoolPatcherBufferInfos(PoolPatcherHelper, CustomTestCase):
|
||||
def test_get_contiguous_buf_infos_inserts_canary_entries(self):
|
||||
"""Verify contiguous buffer metadata includes canary entries after patching."""
|
||||
for patched in (False, True):
|
||||
with self.subTest(patched=patched):
|
||||
pool = make_mha_pool(self.device, num_slots=16, dim=8, layer_num=2)
|
||||
ptrs_before, _, _ = pool.get_contiguous_buf_infos()
|
||||
n_before = len(ptrs_before)
|
||||
|
||||
if patched:
|
||||
attach_canary_buffers(
|
||||
pool=pool,
|
||||
config=self.config,
|
||||
device=self.device,
|
||||
kv_token_id_vs_position_offset=0,
|
||||
)
|
||||
ptrs_after, _, _ = pool.get_contiguous_buf_infos()
|
||||
self.assertEqual(len(ptrs_after), n_before + 4)
|
||||
else:
|
||||
ptrs_after, _, _ = pool.get_contiguous_buf_infos()
|
||||
self.assertEqual(ptrs_after, ptrs_before)
|
||||
|
||||
def test_pd_layout_canary_inserted_correctly(self):
|
||||
"""Verify PD (prefill-decode disaggregation) canary buffers are inserted in layout order."""
|
||||
pool = make_mha_pool(self.device, num_slots=16, dim=8, layer_num=2)
|
||||
k_ptrs_orig = [b.data_ptr() for b in pool.k_buffer]
|
||||
v_ptrs_orig = [b.data_ptr() for b in pool.v_buffer]
|
||||
|
||||
groups_tuple = attach_canary_buffers(
|
||||
pool=pool,
|
||||
config=self.config,
|
||||
device=self.device,
|
||||
kv_token_id_vs_position_offset=0,
|
||||
)
|
||||
group = {g.kind: g for g in groups_tuple}[PoolKind.FULL]
|
||||
ptrs_after, _, _ = pool.get_contiguous_buf_infos()
|
||||
|
||||
canary_k_ptrs = [group.k_head.data_ptr(), group.k_tail.data_ptr()]
|
||||
canary_v_ptrs = [group.v_head.data_ptr(), group.v_tail.data_ptr()]
|
||||
self.assertEqual(ptrs_after[0], canary_k_ptrs[0])
|
||||
self.assertEqual(ptrs_after[1 : 1 + len(k_ptrs_orig)], k_ptrs_orig)
|
||||
k_tail_idx = 1 + len(k_ptrs_orig)
|
||||
self.assertEqual(ptrs_after[k_tail_idx], canary_k_ptrs[1])
|
||||
self.assertEqual(ptrs_after[k_tail_idx + 1], canary_v_ptrs[0])
|
||||
v_start = k_tail_idx + 2
|
||||
self.assertEqual(ptrs_after[v_start : v_start + len(v_ptrs_orig)], v_ptrs_orig)
|
||||
self.assertEqual(ptrs_after[-1], canary_v_ptrs[1])
|
||||
|
||||
|
||||
class TestCanaryBufferBudget(PoolPatcherHelper, CustomTestCase):
|
||||
def test_canary_buf_per_token_bytes_within_budget(self):
|
||||
"""Verify canary per-token storage stays below the real KV budget."""
|
||||
pool = make_mha_pool(self.device, num_slots=16, dim=64, layer_num=2)
|
||||
groups_tuple = attach_canary_buffers(
|
||||
pool=pool,
|
||||
config=self.config,
|
||||
device=self.device,
|
||||
kv_token_id_vs_position_offset=0,
|
||||
)
|
||||
group = {g.kind: g for g in groups_tuple}[PoolKind.FULL]
|
||||
|
||||
slot_stride_bytes = group.k_head.stride(0) * group.k_head.element_size()
|
||||
self.assertLessEqual(slot_stride_bytes, CANARY_SLOT_BYTES)
|
||||
|
||||
real_kv_per_token_bytes = (
|
||||
pool.k_buffer[0].stride(0) * pool.k_buffer[0].element_size()
|
||||
)
|
||||
self.assertLess(slot_stride_bytes, real_kv_per_token_bytes)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,52 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
|
||||
from sglang.srt.kv_canary.pool_patcher.utils import wrap_method
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=10, stage="extra-a", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
class _FakeObj:
|
||||
def greet(self, name: str) -> str:
|
||||
return f"hello {name}"
|
||||
|
||||
|
||||
class TestPoolPatcherUtils(CustomTestCase):
|
||||
def test_wrap_method_delegates_to_wrapper(self) -> None:
|
||||
"""Verify wrapped methods delegate through the wrapper."""
|
||||
obj = _FakeObj()
|
||||
|
||||
def _with_validation(original, name: str) -> str:
|
||||
return original(name) + "!"
|
||||
|
||||
wrap_method(obj, "greet", wrapper=_with_validation)
|
||||
self.assertEqual(obj.greet("world"), "hello world!")
|
||||
|
||||
def test_wrap_method_missing_method_raises_attribute_error(self) -> None:
|
||||
"""Verify wrapping a missing method raises AttributeError."""
|
||||
obj = _FakeObj()
|
||||
with self.assertRaisesRegex(AttributeError, "missing required method"):
|
||||
wrap_method(
|
||||
obj, "nonexistent", wrapper=lambda orig, *a, **kw: orig(*a, **kw)
|
||||
)
|
||||
|
||||
def test_wrap_method_double_wrap_raises_runtime_error(self) -> None:
|
||||
"""Verify wrapping the same method twice raises RuntimeError."""
|
||||
obj = _FakeObj()
|
||||
wrap_method(obj, "greet", wrapper=lambda orig, *a, **kw: orig(*a, **kw))
|
||||
with self.assertRaisesRegex(RuntimeError, "already wrapped by kv-canary"):
|
||||
wrap_method(obj, "greet", wrapper=lambda orig, *a, **kw: orig(*a, **kw))
|
||||
|
||||
def test_wrap_method_preserves_functools_wraps_metadata(self) -> None:
|
||||
"""Verify wrapping preserves method metadata."""
|
||||
obj = _FakeObj()
|
||||
original_name = obj.greet.__name__
|
||||
wrap_method(obj, "greet", wrapper=lambda orig, *a, **kw: orig(*a, **kw))
|
||||
self.assertEqual(obj.greet.__name__, original_name)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,333 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.kv_canary import consts
|
||||
from sglang.jit_kernel.kv_canary.consts import FailReason
|
||||
from sglang.jit_kernel.kv_canary.verify import CanaryLaunchTag
|
||||
from sglang.srt.kv_canary.config import CanaryMode
|
||||
from sglang.srt.kv_canary.runner import violation_reporter as violation_reporter_module
|
||||
from sglang.srt.kv_canary.runner.violation_reporter import (
|
||||
ViolationReporter,
|
||||
_format_violation,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=5, stage="extra-a", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
def _make_row(
|
||||
*,
|
||||
kernel_kind: CanaryLaunchTag = CanaryLaunchTag.HEAD_K_FULL,
|
||||
slot_idx: int = 17,
|
||||
position: int = 42,
|
||||
stored_token: int = 111,
|
||||
expected_token: int = 0,
|
||||
stored_chain_hash: int = 0,
|
||||
expected_aux: int = 0,
|
||||
fail_reason_bits: int = 0,
|
||||
) -> list[int]:
|
||||
row = [0] * consts.VIOLATION_FIELDS
|
||||
row[consts.VIOLATION_FIELD_KERNEL_KIND] = int(kernel_kind)
|
||||
row[consts.VIOLATION_FIELD_SLOT_IDX] = slot_idx
|
||||
row[consts.VIOLATION_FIELD_POSITION] = position
|
||||
row[consts.VIOLATION_FIELD_STORED_TOKEN] = stored_token
|
||||
row[consts.VIOLATION_FIELD_EXPECTED_TOKEN] = expected_token
|
||||
row[consts.VIOLATION_FIELD_STORED_CHAIN_HASH] = stored_chain_hash
|
||||
row[consts.VIOLATION_FIELD_EXPECTED_AUX] = expected_aux
|
||||
row[consts.VIOLATION_FIELD_FAIL_REASON_BITS] = fail_reason_bits
|
||||
return row
|
||||
|
||||
|
||||
class TestViolationReporter(CustomTestCase):
|
||||
def test_format_violation_verify_path_labels_each_bit(self) -> None:
|
||||
"""Verify verify-path violations render each fail-reason bit."""
|
||||
row = _make_row(
|
||||
stored_chain_hash=0x1111111111111111,
|
||||
expected_aux=0x2222222222222222,
|
||||
fail_reason_bits=int(
|
||||
FailReason.VERIFY_CHAIN_HASH_MISMATCH
|
||||
| FailReason.VERIFY_POSITION_MISMATCH
|
||||
),
|
||||
)
|
||||
out = _format_violation(
|
||||
row=row, total=1, ring_overflow=False, step_when_pumped=7
|
||||
)
|
||||
self.assertEqual(
|
||||
out,
|
||||
"kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=verify_chain_hash+verify_position "
|
||||
"slot_idx=17 position=42 stored_token=111 expected_token=0 stored_chain_hash=0x1111111111111111 "
|
||||
"expected_aux=0x2222222222222222\n"
|
||||
"KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=17, position=42)\n"
|
||||
"canary_kind: per_forward_head_k_full\n"
|
||||
" fail_reasons: verify_chain_hash verify_position\n"
|
||||
" stored: token_id=111 position=42 prev_hash=0x1111111111111111\n"
|
||||
" expected: prev_hash=0x2222222222222222\n"
|
||||
" total_violations=1 ring_overflow=False step_when_pumped=7",
|
||||
)
|
||||
|
||||
def test_format_violation_write_token_mismatch_labels_and_position(self) -> None:
|
||||
"""Verify write-token violations render token and position details."""
|
||||
row = _make_row(
|
||||
position=42,
|
||||
stored_token=999,
|
||||
expected_token=888,
|
||||
stored_chain_hash=0xDEADBEEFCAFEBABE,
|
||||
expected_aux=43,
|
||||
fail_reason_bits=int(FailReason.WRITE_TOKEN_MISMATCH),
|
||||
)
|
||||
out = _format_violation(
|
||||
row=row, total=1, ring_overflow=False, step_when_pumped=0
|
||||
)
|
||||
self.assertEqual(
|
||||
out,
|
||||
"kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=write_token slot_idx=17 position=42 "
|
||||
"stored_token=999 expected_token=888 stored_chain_hash=0xdeadbeefcafebabe "
|
||||
"expected_aux=0x000000000000002b\n"
|
||||
"KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=17, position=42)\n"
|
||||
"canary_kind: per_forward_head_k_full\n"
|
||||
" fail_reasons: write_token\n"
|
||||
" actual: token_id=999 position=42 prev_hash=0xdeadbeefcafebabe\n"
|
||||
" expected: token_id=888 position=43\n"
|
||||
" total_violations=1 ring_overflow=False step_when_pumped=0",
|
||||
)
|
||||
|
||||
def test_format_violation_write_position_mismatch_uses_expected_aux_as_position(
|
||||
self,
|
||||
) -> None:
|
||||
"""Verify write-position violations render expected_aux as a position."""
|
||||
row = _make_row(
|
||||
position=42,
|
||||
stored_token=111,
|
||||
expected_token=111,
|
||||
expected_aux=99,
|
||||
fail_reason_bits=int(FailReason.WRITE_POSITION_MISMATCH),
|
||||
)
|
||||
out = _format_violation(
|
||||
row=row, total=1, ring_overflow=False, step_when_pumped=0
|
||||
)
|
||||
self.assertEqual(
|
||||
out,
|
||||
"kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=write_position slot_idx=17 position=42 "
|
||||
"stored_token=111 expected_token=111 stored_chain_hash=0x0000000000000000 "
|
||||
"expected_aux=0x0000000000000063\n"
|
||||
"KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=17, position=42)\n"
|
||||
"canary_kind: per_forward_head_k_full\n"
|
||||
" fail_reasons: write_position\n"
|
||||
" actual: token_id=111 position=42 prev_hash=0x0000000000000000\n"
|
||||
" expected: token_id=111 position=99\n"
|
||||
" total_violations=1 ring_overflow=False step_when_pumped=0",
|
||||
)
|
||||
|
||||
def test_format_violation_combined_write_bits_render_both_labels(self) -> None:
|
||||
"""Verify combined write violation bits render both labels."""
|
||||
row = _make_row(
|
||||
fail_reason_bits=int(
|
||||
FailReason.WRITE_TOKEN_MISMATCH | FailReason.WRITE_POSITION_MISMATCH
|
||||
),
|
||||
)
|
||||
out = _format_violation(
|
||||
row=row, total=1, ring_overflow=False, step_when_pumped=0
|
||||
)
|
||||
self.assertEqual(
|
||||
out,
|
||||
"kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=write_token+write_position "
|
||||
"slot_idx=17 position=42 stored_token=111 expected_token=0 stored_chain_hash=0x0000000000000000 "
|
||||
"expected_aux=0x0000000000000000\n"
|
||||
"KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=17, position=42)\n"
|
||||
"canary_kind: per_forward_head_k_full\n"
|
||||
" fail_reasons: write_token write_position\n"
|
||||
" actual: token_id=111 position=42 prev_hash=0x0000000000000000\n"
|
||||
" expected: token_id=0 position=0\n"
|
||||
" total_violations=1 ring_overflow=False step_when_pumped=0",
|
||||
)
|
||||
|
||||
def test_format_violation_unknown_kernel_kind_renders_unknown_label(self) -> None:
|
||||
"""Verify unknown kernel kinds render an unknown label."""
|
||||
row = _make_row(fail_reason_bits=int(FailReason.VERIFY_CHAIN_HASH_MISMATCH))
|
||||
row[consts.VIOLATION_FIELD_KERNEL_KIND] = 9999
|
||||
out = _format_violation(
|
||||
row=row, total=1, ring_overflow=False, step_when_pumped=0
|
||||
)
|
||||
self.assertEqual(
|
||||
out,
|
||||
"kv_canary violation: launch_tag=unknown(9999) fail_reason=verify_chain_hash slot_idx=17 position=42 "
|
||||
"stored_token=111 expected_token=0 stored_chain_hash=0x0000000000000000 "
|
||||
"expected_aux=0x0000000000000000\n"
|
||||
"KV cache canary violation detected (kernel_kind=unknown(9999), slot_idx=17, position=42)\n"
|
||||
"canary_kind: unknown(9999)\n"
|
||||
" fail_reasons: verify_chain_hash\n"
|
||||
" stored: token_id=111 position=42 prev_hash=0x0000000000000000\n"
|
||||
" expected: prev_hash=0x0000000000000000\n"
|
||||
" total_violations=1 ring_overflow=False step_when_pumped=0",
|
||||
)
|
||||
|
||||
|
||||
def _make_reporter(
|
||||
*,
|
||||
rows: list[list[int]],
|
||||
write_index: int,
|
||||
ring_capacity: int,
|
||||
mode: CanaryMode = CanaryMode.LOG,
|
||||
) -> ViolationReporter:
|
||||
ring = torch.zeros(ring_capacity, consts.VIOLATION_FIELDS, dtype=torch.int64)
|
||||
for i, row in enumerate(rows):
|
||||
ring[i] = torch.tensor(row, dtype=torch.int64)
|
||||
violation_log = SimpleNamespace(
|
||||
violation_ring=ring,
|
||||
violation_write_index=torch.tensor([write_index], dtype=torch.int32),
|
||||
)
|
||||
device_state = SimpleNamespace(violation_log=violation_log)
|
||||
config = SimpleNamespace(mode=mode)
|
||||
return ViolationReporter(config=config, device_state=device_state)
|
||||
|
||||
|
||||
class TestLogOrRaiseViolation(CustomTestCase):
|
||||
def test_log_or_raise_violation_empty_ring_is_noop(self) -> None:
|
||||
"""Empty ring (write_index=0) emits no warning and leaves reporter non-raised."""
|
||||
reporter = _make_reporter(
|
||||
rows=[], write_index=0, ring_capacity=4, mode=CanaryMode.LOG
|
||||
)
|
||||
with patch.object(violation_reporter_module.logger, "warning") as mock_warning:
|
||||
reporter.log_or_raise_violation(outer_step_counter=0)
|
||||
mock_warning.assert_not_called()
|
||||
self.assertFalse(reporter.is_raised)
|
||||
|
||||
def test_log_mode_emits_one_warning_per_violation(self) -> None:
|
||||
"""Log mode with 2 valid rows emits 2 warnings, each a full _format_violation snapshot for that row."""
|
||||
rows = [
|
||||
_make_row(
|
||||
slot_idx=11,
|
||||
position=101,
|
||||
fail_reason_bits=int(FailReason.VERIFY_CHAIN_HASH_MISMATCH),
|
||||
),
|
||||
_make_row(
|
||||
slot_idx=22,
|
||||
position=202,
|
||||
fail_reason_bits=int(FailReason.VERIFY_POSITION_MISMATCH),
|
||||
),
|
||||
]
|
||||
reporter = _make_reporter(
|
||||
rows=rows, write_index=2, ring_capacity=4, mode=CanaryMode.LOG
|
||||
)
|
||||
with patch.object(violation_reporter_module.logger, "warning") as mock_warning:
|
||||
reporter.log_or_raise_violation(outer_step_counter=7)
|
||||
|
||||
self.assertEqual(mock_warning.call_count, 2)
|
||||
messages: list[str] = [call.args[0] for call in mock_warning.call_args_list]
|
||||
self.assertEqual(
|
||||
messages[0],
|
||||
"kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=verify_chain_hash slot_idx=11 position=101 "
|
||||
"stored_token=111 expected_token=0 stored_chain_hash=0x0000000000000000 "
|
||||
"expected_aux=0x0000000000000000\n"
|
||||
"KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=11, position=101)\n"
|
||||
"canary_kind: per_forward_head_k_full\n"
|
||||
" fail_reasons: verify_chain_hash\n"
|
||||
" stored: token_id=111 position=101 prev_hash=0x0000000000000000\n"
|
||||
" expected: prev_hash=0x0000000000000000\n"
|
||||
" total_violations=2 ring_overflow=False step_when_pumped=7",
|
||||
)
|
||||
self.assertEqual(
|
||||
messages[1],
|
||||
"kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=verify_position slot_idx=22 position=202 "
|
||||
"stored_token=111 expected_token=0 stored_chain_hash=0x0000000000000000 "
|
||||
"expected_aux=0x0000000000000000\n"
|
||||
"KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=22, position=202)\n"
|
||||
"canary_kind: per_forward_head_k_full\n"
|
||||
" fail_reasons: verify_position\n"
|
||||
" stored: token_id=111 position=202 prev_hash=0x0000000000000000\n"
|
||||
" expected: prev_hash=0x0000000000000000\n"
|
||||
" total_violations=2 ring_overflow=False step_when_pumped=7",
|
||||
)
|
||||
self.assertFalse(reporter.is_raised)
|
||||
|
||||
def test_raise_mode_raises_one_error_containing_all_violations(self) -> None:
|
||||
"""Raise mode raises a single RuntimeError whose text is the 2 formatted rows joined with single newlines."""
|
||||
rows = [
|
||||
_make_row(
|
||||
slot_idx=11,
|
||||
position=101,
|
||||
fail_reason_bits=int(FailReason.VERIFY_CHAIN_HASH_MISMATCH),
|
||||
),
|
||||
_make_row(
|
||||
slot_idx=22,
|
||||
position=202,
|
||||
fail_reason_bits=int(FailReason.VERIFY_POSITION_MISMATCH),
|
||||
),
|
||||
]
|
||||
reporter = _make_reporter(
|
||||
rows=rows, write_index=2, ring_capacity=4, mode=CanaryMode.RAISE
|
||||
)
|
||||
with self.assertRaises(RuntimeError) as ctx:
|
||||
reporter.log_or_raise_violation(outer_step_counter=5)
|
||||
|
||||
self.assertEqual(
|
||||
str(ctx.exception),
|
||||
"kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=verify_chain_hash slot_idx=11 position=101 "
|
||||
"stored_token=111 expected_token=0 stored_chain_hash=0x0000000000000000 "
|
||||
"expected_aux=0x0000000000000000\n"
|
||||
"KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=11, position=101)\n"
|
||||
"canary_kind: per_forward_head_k_full\n"
|
||||
" fail_reasons: verify_chain_hash\n"
|
||||
" stored: token_id=111 position=101 prev_hash=0x0000000000000000\n"
|
||||
" expected: prev_hash=0x0000000000000000\n"
|
||||
" total_violations=2 ring_overflow=False step_when_pumped=5\n"
|
||||
"kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=verify_position slot_idx=22 position=202 "
|
||||
"stored_token=111 expected_token=0 stored_chain_hash=0x0000000000000000 "
|
||||
"expected_aux=0x0000000000000000\n"
|
||||
"KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=22, position=202)\n"
|
||||
"canary_kind: per_forward_head_k_full\n"
|
||||
" fail_reasons: verify_position\n"
|
||||
" stored: token_id=111 position=202 prev_hash=0x0000000000000000\n"
|
||||
" expected: prev_hash=0x0000000000000000\n"
|
||||
" total_violations=2 ring_overflow=False step_when_pumped=5",
|
||||
)
|
||||
self.assertTrue(reporter.is_raised)
|
||||
|
||||
def test_log_mode_ring_overflow_marks_overflow_in_each_row(self) -> None:
|
||||
"""Log mode with write_index=5 but ring_capacity=2 emits 2 warnings, each a full snapshot with overflow footer."""
|
||||
rows = [
|
||||
_make_row(slot_idx=11, position=101),
|
||||
_make_row(slot_idx=22, position=202),
|
||||
]
|
||||
reporter = _make_reporter(
|
||||
rows=rows, write_index=5, ring_capacity=2, mode=CanaryMode.LOG
|
||||
)
|
||||
with patch.object(violation_reporter_module.logger, "warning") as mock_warning:
|
||||
reporter.log_or_raise_violation(outer_step_counter=0)
|
||||
|
||||
self.assertEqual(mock_warning.call_count, 2)
|
||||
messages: list[str] = [call.args[0] for call in mock_warning.call_args_list]
|
||||
self.assertEqual(
|
||||
messages[0],
|
||||
"kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=none slot_idx=11 position=101 "
|
||||
"stored_token=111 expected_token=0 stored_chain_hash=0x0000000000000000 "
|
||||
"expected_aux=0x0000000000000000\n"
|
||||
"KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=11, position=101)\n"
|
||||
"canary_kind: per_forward_head_k_full\n"
|
||||
" fail_reasons: none\n"
|
||||
" stored: token_id=111 position=101 prev_hash=0x0000000000000000\n"
|
||||
" expected: prev_hash=0x0000000000000000\n"
|
||||
" total_violations=5 ring_overflow=True step_when_pumped=0",
|
||||
)
|
||||
self.assertEqual(
|
||||
messages[1],
|
||||
"kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=none slot_idx=22 position=202 "
|
||||
"stored_token=111 expected_token=0 stored_chain_hash=0x0000000000000000 "
|
||||
"expected_aux=0x0000000000000000\n"
|
||||
"KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=22, position=202)\n"
|
||||
"canary_kind: per_forward_head_k_full\n"
|
||||
" fail_reasons: none\n"
|
||||
" stored: token_id=111 position=202 prev_hash=0x0000000000000000\n"
|
||||
" expected: prev_hash=0x0000000000000000\n"
|
||||
" total_violations=5 ring_overflow=True step_when_pumped=0",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user