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

[Intel][XPU][KVCanary] Enable KV Canary on Intel XPU (#33520)

Browse Source 
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Dayananda V
2026-09-22 09:19:01 +08:00
committed by GitHub
co-authored by Claude Opus 5
parent 046cd6f4ea
commit 35eb7cf8d6
29 changed files with 690 additions and 190 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/manual/kv_canary/test_xpu_kv_canary_pp.py
+36
View File
@@ -0,0 +1,36 @@
"""KV-canary end-to-end on Intel XPU with pipeline parallelism.
``--pp 2`` routes the run through ``Qwen3ForCausalLM.set_embed_and_head`` (mha mode
is Qwen/Qwen3-0.6B), the embedding/head handoff that syncs and releases the device
cache. Needs two XPU cards, so it is manual until a 2-card lane is confirmed.
"""
from __future__ import annotations
import unittest
from sglang.srt.kv_canary.config import CanaryMode
from sglang.test.kv_canary.e2e_base import CanaryE2EBase
class TestXPUCanaryPipelineParallel(CanaryE2EBase):
"""Clean canary run across a pipeline-parallel XPU pair."""
model_mode = "mha"
kv_canary_mode = CanaryMode.LOG
# --disable-cuda-graph is mandatory, not tuning: install_canary refuses a captured decode
# on a device that routes to the torch reference (host work and D2H, so replay checks nothing).
extra_server_args = ("--device", "xpu", "--disable-cuda-graph", "--pp", "2")
# The torch reference folds the chain slot-by-slot on the host, so the workload is much
# smaller than the CUDA-tuned defaults on the shared base.
default_parallel_n = 2
default_max_new_tokens = 32
default_request_timeout = 120.0
def test_no_violation(self) -> None:
self.send_parallel_requests()
self.assert_no_violation(wait_seconds=2.0)
if __name__ == "__main__":
unittest.main()
test/registered/kernels/ops/kv_canary/test_verify_hand.py
+30 -6
View File
@@ -20,6 +20,7 @@ from sglang.kernels.ops.kv_canary.verify import (
from sglang.kernels.ops.kv_canary.verify_ref import (
_compute_real_kv_hash_scalar,
launch_canary_verify_kernel_torch_reference,
materialize_real_kv_sources,
)
from sglang.kernels.ops.kv_canary.write_ref import (
launch_canary_write_kernel_torch_reference,
@@ -917,14 +918,18 @@ class TestRealKvHash:
positions = [0, 1, 2]
running = splitmix64(consts.CANARY_CHAIN_ANCHOR)
host_sources = materialize_real_kv_sources(
real_kv_sources=sources_cuda,
real_kv_hash_mode=consts.RealKvHashMode.ALL,
slot_indices=slot_indices,
work_device=torch.device("cpu"),
)
real_kv_hashes: list[int] = []
for slot_idx in slot_indices:
real_kv_hashes.append(
_compute_real_kv_hash_scalar(
real_kv_sources=sources_cuda,
real_kv_hash_mode=consts.RealKvHashMode.ALL,
slot_idx=slot_idx,
work_device=torch.device("cpu"),
host_sources=host_sources,
)
)
@@ -1030,6 +1035,21 @@ class TestRealKvSource:
read_bytes=0,
)
def test_real_kv_source_rejects_row_narrower_than_page(self) -> None:
"""A row too narrow for its page must raise: neither fold reports it.
The CUDA fold reads past the row and the torch fold's dim-1 slice clamps to
the row end, so the tail slots of the page hash 0 bytes and the chain still
verifies clean.
"""
with pytest.raises(ValueError, match="page_size"):
RealKvSource(
tensor=torch.zeros((1, 16), dtype=torch.uint8, device=_DEVICE),
page_size=2,
num_bytes_per_token=16,
read_bytes=16,
)
def test_real_kv_source_padding_below_4(self) -> None:
"""Host wrapper pads to 4 slots when fewer sources are supplied; dummy slots are never dereferenced."""
buf_pair = _buf_pair()
@@ -1310,12 +1330,16 @@ class TestLayoutAndScheduling:
# byte-by-byte loop, so the stamped real_kv_hash matches what the kernel /
# verify reference will recompute. A byte-by-byte fold was the previous bug
# here and triggered REAL_KV_HASH violations on otherwise clean chains.
host_sources = materialize_real_kv_sources(
real_kv_sources=sources_cuda,
real_kv_hash_mode=consts.RealKvHashMode.ALL,
slot_indices=slot_indices,
work_device=_DEVICE,
)
rkv_values = [
_compute_real_kv_hash_scalar(
slot_idx=slot_idx,
real_kv_sources=sources_cuda,
real_kv_hash_mode=consts.RealKvHashMode.ALL,
work_device=_DEVICE,
host_sources=host_sources,
)
for slot_idx in slot_indices
]
test/registered/kv_canary/test_self_unit_capture_refusal.py
+77
View File
@@ -0,0 +1,77 @@
from __future__ import annotations
import unittest
from unittest import mock
import torch
from sglang.srt.kv_canary import api
from sglang.srt.kv_canary.api import torch_reference_conflicts_with_decode_graph
from sglang.srt.model_executor.cuda_graph_config import (
Backend,
CudaGraphConfig,
PhaseConfig,
)
from sglang.srt.runtime_context import get_context
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
class TestTorchReferenceConflictsWithDecodeGraph(CustomTestCase):
"""The refusal that keeps a graph-captured torch reference from passing silently.
The reference path does host work and D2H, so its launches leave nothing in a
captured decode graph and every replay verifies clean. Each case below pins one
branch of the gate; the platform capability is patched rather than probed so the
CPU lane exercises all four.
"""
def _publish_decode_backend(self, backend: str) -> None:
override = get_context().override_server_args(
cuda_graph_config=CudaGraphConfig(decode=PhaseConfig(backend=backend))
)
override.install()
self.addCleanup(override.restore)
def _patch_graph_support(self, supported: bool) -> None:
patcher = mock.patch.object(
api.current_platform, "support_cuda_graph", return_value=supported
)
patcher.start()
self.addCleanup(patcher.stop)
def test_reference_device_with_captured_decode_conflicts(self) -> None:
self._patch_graph_support(True)
self._publish_decode_backend(Backend.FULL)
self.assertTrue(
torch_reference_conflicts_with_decode_graph(torch.device("xpu"))
)
def test_reference_device_with_decode_graph_disabled_is_allowed(self) -> None:
self._patch_graph_support(True)
self._publish_decode_backend(Backend.DISABLED)
self.assertFalse(
torch_reference_conflicts_with_decode_graph(torch.device("xpu"))
)
def test_platform_without_graph_capture_is_allowed(self) -> None:
"""A device that never captures (CPU) keeps canary on the reference path."""
self._patch_graph_support(False)
self._publish_decode_backend(Backend.FULL)
self.assertFalse(
torch_reference_conflicts_with_decode_graph(torch.device("cpu"))
)
def test_cuda_device_is_never_refused(self) -> None:
"""CUDA/HIP run the real kernels, so the gate must not fire on them."""
self._patch_graph_support(True)
self._publish_decode_backend(Backend.FULL)
self.assertFalse(
torch_reference_conflicts_with_decode_graph(torch.device("cuda"))
)
if __name__ == "__main__":
unittest.main()
test/registered/kv_canary/test_self_unit_future_tensor.py
+36 -41
View File
@@ -1,16 +1,22 @@
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_amd_ci, register_cuda_ci
from sglang.srt.utils import create_device_stream, get_current_device_stream_fast
from sglang.test.ci.ci_register import (
register_amd_ci,
register_cuda_ci,
register_xpu_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")
register_amd_ci(est_time=20, suite="extra-a-test-1-gpu-small-amd")
register_xpu_ci(est_time=20, suite="stage-b-test-1-gpu-xpu")
class _FakeEvent:
@@ -22,32 +28,30 @@ class _FakeEvent:
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)
def test_device_stage_then_wait_returns_host_copy(self) -> None:
"""Verify staged device tensors are copied back on wait."""
alt_stream = create_device_stream(DEFAULT_DEVICE)
default_stream = get_current_device_stream_fast()
self.assertNotEqual(alt_stream.stream_id, default_stream.stream_id)
src_first = torch.tensor([41], dtype=torch.int32, device=device)
src_first = torch.tensor([41], dtype=torch.int32, device=DEFAULT_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)
src_second = torch.tensor([97], dtype=torch.int32, device=DEFAULT_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)
def test_device_pinned_when_stream_is_provided(self) -> None:
"""Verify device staging uses pinned host memory with a stream."""
alt_stream = create_device_stream(DEFAULT_DEVICE)
src = torch.tensor([5], dtype=torch.int32, device=DEFAULT_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)
@@ -56,12 +60,11 @@ class TestFutureTensors(CustomTestCase):
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)
def test_device_each_call_allocates_fresh_host(self) -> None:
"""Verify each device staging call owns a fresh host buffer."""
alt_stream = create_device_stream(DEFAULT_DEVICE)
src_a = torch.tensor([13], dtype=torch.int32, device=DEFAULT_DEVICE)
src_b = torch.tensor([29], dtype=torch.int32, device=DEFAULT_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 = {
@@ -77,11 +80,10 @@ class TestFutureTensors(CustomTestCase):
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)
stream = create_device_stream(DEFAULT_DEVICE)
src = {
"x": torch.tensor([11, 22], dtype=torch.int64, device=device),
"y": torch.tensor([99], dtype=torch.int32, device=device),
"x": torch.tensor([11, 22], dtype=torch.int64, device=DEFAULT_DEVICE),
"y": torch.tensor([99], dtype=torch.int32, device=DEFAULT_DEVICE),
}
future = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream)
out = future.wait()
@@ -93,14 +95,13 @@ class TestFutureTensors(CustomTestCase):
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)
stream = create_device_stream(DEFAULT_DEVICE)
sentinel_obj = {"nested": [1, 2, 3]}
src = {
"step": 42,
"label": "decode",
"extra": sentinel_obj,
"counter": torch.tensor([7], dtype=torch.int32, device=device),
"counter": torch.tensor([7], dtype=torch.int32, device=DEFAULT_DEVICE),
}
future = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream)
out = future.wait()
@@ -113,9 +114,8 @@ class TestFutureTensors(CustomTestCase):
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)
stream = create_device_stream(DEFAULT_DEVICE)
src_tensor = torch.tensor([3], dtype=torch.int32, device=DEFAULT_DEVICE)
src = {"step": 100, "buf": src_tensor}
future = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream)
out = future.wait()
@@ -128,8 +128,7 @@ class TestFutureTensors(CustomTestCase):
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)
stream = create_device_stream(DEFAULT_DEVICE)
with self.assertRaises(ValueError):
FutureTensors.device_to_host(
xs_device={"step": 0, "label": "decode"}, d2h_stream=stream
@@ -137,9 +136,8 @@ class TestFutureTensors(CustomTestCase):
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)
stream = create_device_stream(DEFAULT_DEVICE)
src = torch.tensor([3], dtype=torch.int32, device=DEFAULT_DEVICE)
future = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream)
self.assertEqual(int(future.wait().item()), 3)
with self.assertRaises(RuntimeError):
@@ -149,9 +147,7 @@ class TestFutureTensors(CustomTestCase):
"""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)
)
future = FutureTensors(_data={"x": tensor}, _event=event)
result = future.wait()
self.assertIs(result["x"], tensor)
@@ -166,11 +162,10 @@ class TestFutureTensors(CustomTestCase):
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)
stream = create_device_stream(DEFAULT_DEVICE)
src = {
"step": 5,
"buf": torch.tensor([17], dtype=torch.int32, device=device),
"buf": torch.tensor([17], dtype=torch.int32, device=DEFAULT_DEVICE),
}
out = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream).wait()
self.assertEqual(out["step"], 5)
test/registered/kv_canary/test_self_unit_plan_input.py
+8 -2
View File
@@ -6,15 +6,21 @@ from types import SimpleNamespace
import torch
from sglang.srt.kv_canary.plan_input import PlanInput
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.ci.ci_register import (
register_amd_ci,
register_cuda_ci,
register_xpu_ci,
)
from sglang.test.kv_canary.fixtures import (
DEFAULT_DEVICE,
DEFAULT_DEVICE_MODULE,
make_forward_batch,
)
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=9, stage="extra-a", runner_config="1-gpu-small")
register_amd_ci(est_time=30, suite="extra-a-test-1-gpu-small-amd")
register_xpu_ci(est_time=30, suite="stage-b-test-1-gpu-xpu")
def _make_static_plan_input(*, bs_capacity: int, device) -> PlanInput:
@@ -122,7 +128,7 @@ class TestSelfUnitPlanInput(CustomTestCase):
fb.req_all_ids_lens = torch.tensor([7, 9], dtype=torch.int64, pin_memory=True)
plan = _make_static_plan_input(bs_capacity=4, device=self.device)
plan.fill_from_forward_batch(forward_batch=fb)
torch.cuda.synchronize()
DEFAULT_DEVICE_MODULE.synchronize()
self.assertEqual(
plan.req_to_verify_expected_tokens_valid_lens[:2].tolist(), [7, 9]
)
test/registered/kv_canary/test_self_unit_req_to_expected_token_ids_manager.py
+14 -5
View File
@@ -10,12 +10,21 @@ from sglang.srt.kv_canary.req_to_expected_token_ids_manager import (
compute_req_all_ids_info,
populate_req_to_expected_token_ids,
)
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.kv_canary.fixtures import DEFAULT_DEVICE, make_forward_batch
from sglang.test.ci.ci_register import (
register_amd_ci,
register_cuda_ci,
register_xpu_ci,
)
from sglang.test.kv_canary.fixtures import (
DEFAULT_DEVICE,
DEFAULT_DEVICE_MODULE,
make_forward_batch,
)
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=11, stage="extra-a", runner_config="1-gpu-small")
register_amd_ci(est_time=15, suite="extra-a-test-1-gpu-small-amd")
register_xpu_ci(est_time=30, suite="stage-b-test-1-gpu-xpu")
def _make_req(*, origin: list[int], output: list[int]) -> SimpleNamespace:
@@ -92,7 +101,7 @@ class TestPopulateReqToExpectedTokenIds(CustomTestCase):
populate_req_to_expected_token_ids(
forward_batch=fb, req_to_verify_expected_tokens=pool
)
torch.cuda.synchronize()
DEFAULT_DEVICE_MODULE.synchronize()
self.assertTrue(torch.equal(pool, original))
def test_no_op_when_pool_is_none(self) -> None:
@@ -117,7 +126,7 @@ class TestPopulateReqToExpectedTokenIds(CustomTestCase):
populate_req_to_expected_token_ids(
forward_batch=fb, req_to_verify_expected_tokens=pool
)
torch.cuda.synchronize()
DEFAULT_DEVICE_MODULE.synchronize()
self.assertTrue(torch.equal(pool, original))
def test_raises_when_lens_length_mismatches_batch_size(self) -> None:
@@ -154,7 +163,7 @@ class TestPopulateReqToExpectedTokenIds(CustomTestCase):
populate_req_to_expected_token_ids(
forward_batch=fb, req_to_verify_expected_tokens=pool
)
torch.cuda.synchronize()
DEFAULT_DEVICE_MODULE.synchronize()
pool_cpu = pool.cpu()
self.assertEqual(pool_cpu[1, :3].tolist(), [10, 20, 30])
test/registered/kv_canary/test_self_unit_runner_swa_divergence.py
+65 -39
View File
@@ -15,24 +15,28 @@ from sglang.srt.kv_canary.runner.swa_divergence import (
SwaDivergenceReporter,
compute_swa_full_idx_divergence,
)
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.kv_canary.fixtures import make_buffer_group
from sglang.srt.utils import create_device_stream
from sglang.test.ci.ci_register import (
register_amd_ci,
register_cuda_ci,
register_xpu_ci,
)
from sglang.test.kv_canary.fixtures import DEFAULT_DEVICE, make_buffer_group
from sglang.test.kv_canary.runner_test_base import CanaryManagerTestCase, make_manager
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=11, stage="extra-a", runner_config="1-gpu-small")
register_amd_ci(est_time=45, suite="extra-a-test-1-gpu-small-amd")
_DEVICE = torch.device("cuda")
register_xpu_ci(est_time=60, suite="stage-b-test-1-gpu-xpu")
_EMPTY_FORWARD_BATCH = SimpleNamespace(
req_pool_indices=torch.empty(0, dtype=torch.int64, device=_DEVICE),
seq_lens=torch.empty(0, dtype=torch.int64, device=_DEVICE),
req_pool_indices=torch.empty(0, dtype=torch.int64, device=DEFAULT_DEVICE),
seq_lens=torch.empty(0, dtype=torch.int64, device=DEFAULT_DEVICE),
)
def _make_verify_plan(value: int) -> VerifyPlan:
plan = VerifyPlan.allocate(verify_capacity=4, device=_DEVICE)
plan = VerifyPlan.allocate(verify_capacity=4, device=DEFAULT_DEVICE)
plan.verify_num_valid.copy_(torch.tensor([value], dtype=torch.int32))
return plan
@@ -46,11 +50,13 @@ def _make_req_to_token_pool_stub(req_to_token: torch.Tensor) -> SimpleNamespace:
def _make_identity_mapping(size: int) -> torch.Tensor:
return torch.arange(size, dtype=torch.int64, device=_DEVICE)
return torch.arange(size, dtype=torch.int64, device=DEFAULT_DEVICE)
def _make_identity_req_to_token(num_reqs: int, max_seq_len: int) -> torch.Tensor:
base = torch.arange(num_reqs * max_seq_len, dtype=torch.int64, device=_DEVICE)
base = torch.arange(
num_reqs * max_seq_len, dtype=torch.int64, device=DEFAULT_DEVICE
)
return base.view(num_reqs, max_seq_len)
@@ -83,9 +89,9 @@ def _run_compute(
class TestSwaDivergenceReporter(CustomTestCase):
def test_swa_divergence_log_emitted(self) -> None:
d2h_stream = torch.cuda.Stream(device=_DEVICE)
d2h_stream = create_device_stream(DEFAULT_DEVICE)
stats = SwaDivergenceReporter(
device=_DEVICE,
device=DEFAULT_DEVICE,
d2h_stream=d2h_stream,
interval=10,
swa_allocator=None,
@@ -96,13 +102,13 @@ class TestSwaDivergenceReporter(CustomTestCase):
for forward_idx in range(3):
stats.observe_after_invoke_plan(
group=make_buffer_group(
device=_DEVICE, kind=PoolKind.FULL, has_v=False, num_slots=1
device=DEFAULT_DEVICE, kind=PoolKind.FULL, has_v=False, num_slots=1
),
verify_plan=_make_verify_plan(10),
)
stats.observe_after_invoke_plan(
group=make_buffer_group(
device=_DEVICE, kind=PoolKind.SWA, has_v=False, num_slots=1
device=DEFAULT_DEVICE, kind=PoolKind.SWA, has_v=False, num_slots=1
),
verify_plan=_make_verify_plan(3),
)
@@ -115,13 +121,13 @@ class TestSwaDivergenceReporter(CustomTestCase):
# the staged future hangs onto it. forward_ct is now 4.
stats.observe_after_invoke_plan(
group=make_buffer_group(
device=_DEVICE, kind=PoolKind.FULL, has_v=False, num_slots=1
device=DEFAULT_DEVICE, kind=PoolKind.FULL, has_v=False, num_slots=1
),
verify_plan=_make_verify_plan(10),
)
stats.observe_after_invoke_plan(
group=make_buffer_group(
device=_DEVICE, kind=PoolKind.SWA, has_v=False, num_slots=1
device=DEFAULT_DEVICE, kind=PoolKind.SWA, has_v=False, num_slots=1
),
verify_plan=_make_verify_plan(3),
)
@@ -150,9 +156,9 @@ class TestSwaDivergenceReporter(CustomTestCase):
self.assertEqual(fields.swa_full_idx_divergence, 0)
def test_swa_divergence_counts_monotonic_increasing(self) -> None:
d2h_stream = torch.cuda.Stream(device=_DEVICE)
d2h_stream = create_device_stream(DEFAULT_DEVICE)
stats = SwaDivergenceReporter(
device=_DEVICE,
device=DEFAULT_DEVICE,
d2h_stream=d2h_stream,
interval=10,
swa_allocator=None,
@@ -188,13 +194,19 @@ class TestSwaDivergenceReporter(CustomTestCase):
for _ in range(5):
stats.observe_after_invoke_plan(
group=make_buffer_group(
device=_DEVICE, kind=PoolKind.FULL, has_v=False, num_slots=1
device=DEFAULT_DEVICE,
kind=PoolKind.FULL,
has_v=False,
num_slots=1,
),
verify_plan=_make_verify_plan(7),
)
stats.observe_after_invoke_plan(
group=make_buffer_group(
device=_DEVICE, kind=PoolKind.SWA, has_v=False, num_slots=1
device=DEFAULT_DEVICE,
kind=PoolKind.SWA,
has_v=False,
num_slots=1,
),
verify_plan=_make_verify_plan(2),
)
@@ -216,8 +228,8 @@ class TestSwaFullIdxDivergenceCompute(CustomTestCase):
req_to_token = _make_identity_req_to_token(num_reqs=4, max_seq_len=16)
forward_batch = _make_forward_batch(
req_pool_indices=torch.empty(0, dtype=torch.int64, device=_DEVICE),
seq_lens=torch.empty(0, dtype=torch.int64, device=_DEVICE),
req_pool_indices=torch.empty(0, dtype=torch.int64, device=DEFAULT_DEVICE),
seq_lens=torch.empty(0, dtype=torch.int64, device=DEFAULT_DEVICE),
)
self.assertEqual(
@@ -234,8 +246,10 @@ class TestSwaFullIdxDivergenceCompute(CustomTestCase):
req_to_token = _make_identity_req_to_token(num_reqs=4, max_seq_len=16)
forward_batch = _make_forward_batch(
req_pool_indices=torch.tensor([0, 2], dtype=torch.int64, device=_DEVICE),
seq_lens=torch.tensor([8, 5], dtype=torch.int64, device=_DEVICE),
req_pool_indices=torch.tensor(
[0, 2], dtype=torch.int64, device=DEFAULT_DEVICE
),
seq_lens=torch.tensor([8, 5], dtype=torch.int64, device=DEFAULT_DEVICE),
)
self.assertEqual(
@@ -256,8 +270,10 @@ class TestSwaFullIdxDivergenceCompute(CustomTestCase):
mapping[17] = 60
forward_batch = _make_forward_batch(
req_pool_indices=torch.tensor([0, 1], dtype=torch.int64, device=_DEVICE),
seq_lens=torch.tensor([8, 8], dtype=torch.int64, device=_DEVICE),
req_pool_indices=torch.tensor(
[0, 1], dtype=torch.int64, device=DEFAULT_DEVICE
),
seq_lens=torch.tensor([8, 8], dtype=torch.int64, device=DEFAULT_DEVICE),
)
self.assertEqual(
@@ -281,8 +297,10 @@ class TestSwaFullIdxDivergenceCompute(CustomTestCase):
mapping[7] = 42
forward_batch = _make_forward_batch(
req_pool_indices=torch.tensor([0], dtype=torch.int64, device=_DEVICE),
seq_lens=torch.tensor([8], dtype=torch.int64, device=_DEVICE),
req_pool_indices=torch.tensor(
[0], dtype=torch.int64, device=DEFAULT_DEVICE
),
seq_lens=torch.tensor([8], dtype=torch.int64, device=DEFAULT_DEVICE),
)
self.assertEqual(
@@ -302,8 +320,10 @@ class TestSwaFullIdxDivergenceCompute(CustomTestCase):
mapping[28] = 77
forward_batch = _make_forward_batch(
req_pool_indices=torch.tensor([0], dtype=torch.int64, device=_DEVICE),
seq_lens=torch.tensor([10], dtype=torch.int64, device=_DEVICE),
req_pool_indices=torch.tensor(
[0], dtype=torch.int64, device=DEFAULT_DEVICE
),
seq_lens=torch.tensor([10], dtype=torch.int64, device=DEFAULT_DEVICE),
)
self.assertEqual(
@@ -325,12 +345,16 @@ class TestSwaFullIdxDivergenceCompute(CustomTestCase):
mapping[33] = 100
fb_req0 = _make_forward_batch(
req_pool_indices=torch.tensor([0], dtype=torch.int64, device=_DEVICE),
seq_lens=torch.tensor([4], dtype=torch.int64, device=_DEVICE),
req_pool_indices=torch.tensor(
[0], dtype=torch.int64, device=DEFAULT_DEVICE
),
seq_lens=torch.tensor([4], dtype=torch.int64, device=DEFAULT_DEVICE),
)
fb_req2 = _make_forward_batch(
req_pool_indices=torch.tensor([2], dtype=torch.int64, device=_DEVICE),
seq_lens=torch.tensor([4], dtype=torch.int64, device=_DEVICE),
req_pool_indices=torch.tensor(
[2], dtype=torch.int64, device=DEFAULT_DEVICE
),
seq_lens=torch.tensor([4], dtype=torch.int64, device=DEFAULT_DEVICE),
)
self.assertEqual(
@@ -363,15 +387,17 @@ class TestSwaDivergenceReporterWithCompute(CustomTestCase):
mapping[2] = 52
forward_batch = _make_forward_batch(
req_pool_indices=torch.tensor([0], dtype=torch.int64, device=_DEVICE),
seq_lens=torch.tensor([8], dtype=torch.int64, device=_DEVICE),
req_pool_indices=torch.tensor(
[0], dtype=torch.int64, device=DEFAULT_DEVICE
),
seq_lens=torch.tensor([8], dtype=torch.int64, device=DEFAULT_DEVICE),
)
swa_allocator = _make_allocator_stub(mapping)
req_to_token_pool = _make_req_to_token_pool_stub(req_to_token)
d2h_stream = torch.cuda.Stream(device=_DEVICE)
d2h_stream = create_device_stream(DEFAULT_DEVICE)
stats = SwaDivergenceReporter(
device=_DEVICE,
device=DEFAULT_DEVICE,
d2h_stream=d2h_stream,
interval=10,
swa_allocator=swa_allocator,
@@ -379,13 +405,13 @@ class TestSwaDivergenceReporterWithCompute(CustomTestCase):
)
stats.observe_after_invoke_plan(
group=make_buffer_group(
device=_DEVICE, kind=PoolKind.FULL, has_v=False, num_slots=1
device=DEFAULT_DEVICE, kind=PoolKind.FULL, has_v=False, num_slots=1
),
verify_plan=_make_verify_plan(11),
)
stats.observe_after_invoke_plan(
group=make_buffer_group(
device=_DEVICE, kind=PoolKind.SWA, has_v=False, num_slots=1
device=DEFAULT_DEVICE, kind=PoolKind.SWA, has_v=False, num_slots=1
),
verify_plan=_make_verify_plan(3),
)
test/registered/utils/test_phase_checker.py
+55 -35
View File
@@ -10,12 +10,23 @@ from enum import IntEnum
import torch
from sglang.srt.utils import get_device
from sglang.srt.utils.phase_checker import SimplePhaseChecker
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.ci.ci_register import (
register_amd_ci,
register_cuda_ci,
register_xpu_ci,
)
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=17, stage="base-b", runner_config="1-gpu-small")
register_amd_ci(est_time=120, stage="stage-b", runner_config="1-gpu-small-amd")
# Nightly, not a blocking lane: one case spawns a subprocess that trips a device-side
# assert, so a wedge costs the whole subprocess timeout below.
register_xpu_ci(est_time=300, suite="nightly-xpu-1-gpu", nightly=True)
_DEVICE: torch.device = torch.device(get_device(device_id=0))
_DEVICE_MODULE = torch.get_device_module(_DEVICE)
class _Phase(IntEnum):
@@ -40,7 +51,7 @@ class TestConstruction(CustomTestCase):
@classmethod
def setUpClass(cls) -> None:
cls.device = torch.device("cuda:0")
cls.device = _DEVICE
def test_init_stores_initial_phase_int(self) -> None:
checker = SimplePhaseChecker(initial_phase=7, device=self.device)
@@ -72,19 +83,19 @@ class TestUpdateAssertDisabled(CustomTestCase):
@classmethod
def setUpClass(cls) -> None:
cls.device = torch.device("cuda:0")
cls.device = _DEVICE
def test_update_advances_phase_on_match(self) -> None:
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
checker.update(expect_phase=_Phase.IDLE, next_phase=_Phase.A, caller_name="t")
torch.cuda.synchronize()
_DEVICE_MODULE.synchronize()
self.assertEqual(_phase_value(checker), int(_Phase.A))
def test_update_advances_phase_on_mismatch(self) -> None:
"""assert OFF tolerates mismatches — store still happens unconditionally."""
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
checker.update(expect_phase=_Phase.C, next_phase=_Phase.B, caller_name="t")
torch.cuda.synchronize()
_DEVICE_MODULE.synchronize()
self.assertEqual(_phase_value(checker), int(_Phase.B))
def test_init_time_lifecycle_violations_tolerated(self) -> None:
@@ -98,7 +109,7 @@ class TestUpdateAssertDisabled(CustomTestCase):
checker.update(
expect_phase=_Phase.B, next_phase=_Phase.IDLE, caller_name="warmup"
)
torch.cuda.synchronize() # no raise
_DEVICE_MODULE.synchronize() # no raise
self.assertEqual(_phase_value(checker), int(_Phase.IDLE))
@@ -107,13 +118,13 @@ class TestUpdateAssertEnabled(CustomTestCase):
@classmethod
def setUpClass(cls) -> None:
cls.device = torch.device("cuda:0")
cls.device = _DEVICE
def test_update_advances_phase_on_match(self) -> None:
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
checker.enable_assert()
checker.update(expect_phase=_Phase.IDLE, next_phase=_Phase.A, caller_name="t")
torch.cuda.synchronize()
_DEVICE_MODULE.synchronize()
self.assertEqual(_phase_value(checker), int(_Phase.A))
def test_full_4_state_lifecycle_round_trip(self) -> None:
@@ -129,28 +140,30 @@ class TestUpdateAssertEnabled(CustomTestCase):
checker.update(
expect_phase=_Phase.C, next_phase=_Phase.IDLE, caller_name="p4"
)
torch.cuda.synchronize()
_DEVICE_MODULE.synchronize()
self.assertEqual(_phase_value(checker), int(_Phase.IDLE))
def test_update_mismatch_after_enable_raises_in_subprocess(self) -> None:
"""A mismatched update with assert ON must fire device_assert at the next sync.
Run in a subprocess because device-side asserts poison the CUDA context.
Run in a subprocess because device-side asserts poison the accelerator context.
"""
script = textwrap.dedent("""
import sys
import torch
from sglang.srt.utils import get_device
from sglang.srt.utils.phase_checker import SimplePhaseChecker
device = torch.device("cuda:0")
device = torch.device(get_device(device_id=0))
device_module = torch.get_device_module(device)
checker = SimplePhaseChecker(initial_phase=0, device=device)
checker.enable_assert()
# phase=0 but we claim expect=99 — kernel must fire device_assert.
checker.update(expect_phase=99, next_phase=1, caller_name="bad")
try:
torch.cuda.synchronize()
device_module.synchronize()
except RuntimeError as e:
msg = str(e).lower()
if "device-side assert" in msg or "phase mismatch" in msg:
@@ -164,7 +177,10 @@ class TestUpdateAssertEnabled(CustomTestCase):
[sys.executable, "-c", script],
capture_output=True,
text=True,
timeout=180,
# Cold-Triton-cache compile takes minutes on XPU, where this file runs
# nightly; too tight a timeout surfaces as a spurious returncode=-9, not
# a real assert regression.
timeout=180 if _DEVICE.type == "cuda" else 600,
)
# The FAIL line is the evidence that the kernel-side check fired. How the
# process then dies is not: the CUDA coredump handler may abort it, and sync
@@ -189,7 +205,7 @@ class TestEnableAssert(CustomTestCase):
@classmethod
def setUpClass(cls) -> None:
cls.device = torch.device("cuda:0")
cls.device = _DEVICE
def test_enable_assert_sets_flag_to_one(self) -> None:
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
@@ -208,7 +224,7 @@ class TestEnableAssert(CustomTestCase):
checker.update(
expect_phase=_Phase.IDLE, next_phase=_Phase.C, caller_name="warmup"
)
torch.cuda.synchronize()
_DEVICE_MODULE.synchronize()
self.assertEqual(_phase_value(checker), int(_Phase.C))
checker.enable_assert()
@@ -218,7 +234,7 @@ class TestEnableAssert(CustomTestCase):
"""Reset target tracks the original initial_phase, not 0."""
checker = SimplePhaseChecker(initial_phase=42, device=self.device)
checker.update(expect_phase=42, next_phase=7, caller_name="t")
torch.cuda.synchronize()
_DEVICE_MODULE.synchronize()
self.assertEqual(_phase_value(checker), 7)
checker.enable_assert()
@@ -237,12 +253,12 @@ class TestResetToIdle(CustomTestCase):
@classmethod
def setUpClass(cls) -> None:
cls.device = torch.device("cuda:0")
cls.device = _DEVICE
def test_reset_after_update_restores_initial_phase(self) -> None:
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
checker.update(expect_phase=_Phase.IDLE, next_phase=_Phase.B, caller_name="t")
torch.cuda.synchronize()
_DEVICE_MODULE.synchronize()
self.assertEqual(_phase_value(checker), int(_Phase.B))
checker._reset_to_idle()
@@ -264,7 +280,7 @@ class TestResetToIdle(CustomTestCase):
def test_reset_with_nonzero_initial_phase(self) -> None:
checker = SimplePhaseChecker(initial_phase=5, device=self.device)
checker.update(expect_phase=5, next_phase=9, caller_name="t")
torch.cuda.synchronize()
_DEVICE_MODULE.synchronize()
checker._reset_to_idle()
self.assertEqual(_phase_value(checker), 5)
@@ -274,7 +290,7 @@ class TestCallerTagRegistry(CustomTestCase):
@classmethod
def setUpClass(cls) -> None:
cls.device = torch.device("cuda:0")
cls.device = _DEVICE
def test_first_caller_gets_tag_one(self) -> None:
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
@@ -299,7 +315,7 @@ class TestCallerTagRegistry(CustomTestCase):
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
checker.enable_assert()
checker.update(expect_phase=_Phase.IDLE, next_phase=_Phase.A) # caller_name=""
torch.cuda.synchronize()
_DEVICE_MODULE.synchronize()
self.assertIn("", checker._caller_tag_registry)
self.assertEqual(_phase_value(checker), int(_Phase.A))
@@ -311,7 +327,7 @@ class TestCallerTagRegistry(CustomTestCase):
checker.update(
expect_phase=_Phase.A, next_phase=_Phase.IDLE, caller_name="beta"
)
torch.cuda.synchronize()
_DEVICE_MODULE.synchronize()
self.assertEqual(checker._caller_tag_registry, {"alpha": 1, "beta": 2})
@@ -320,13 +336,13 @@ class TestMultipleInstances(CustomTestCase):
@classmethod
def setUpClass(cls) -> None:
cls.device = torch.device("cuda:0")
cls.device = _DEVICE
def test_phase_tensors_are_independent(self) -> None:
a = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
b = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
a.update(expect_phase=_Phase.IDLE, next_phase=_Phase.B, caller_name="a")
torch.cuda.synchronize()
_DEVICE_MODULE.synchronize()
self.assertEqual(_phase_value(a), int(_Phase.B))
self.assertEqual(_phase_value(b), int(_Phase.IDLE))
@@ -345,6 +361,10 @@ class TestMultipleInstances(CustomTestCase):
self.assertEqual(b._resolve_caller_tag("shared_name"), 1)
@unittest.skipUnless(
_DEVICE.type == "cuda",
"capture-safety is a CUDA-only contract (torch.cuda.CUDAGraph has no portable equivalent)",
)
class TestCudaGraphCapture(CustomTestCase):
"""The kernel is launched unconditionally so it is capture-safe; the device flag
decides at replay time whether the assert fires.
@@ -352,7 +372,7 @@ class TestCudaGraphCapture(CustomTestCase):
@classmethod
def setUpClass(cls) -> None:
cls.device = torch.device("cuda:0")
cls.device = _DEVICE
def _capture_one_update(
self,
@@ -379,7 +399,7 @@ class TestCudaGraphCapture(CustomTestCase):
caller_name=caller_name,
)
torch.cuda.current_stream(self.device).wait_stream(stream)
torch.cuda.synchronize()
_DEVICE_MODULE.synchronize()
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph, stream=stream):
@@ -404,13 +424,13 @@ class TestCudaGraphCapture(CustomTestCase):
# Enable assert (resets phase -> IDLE) and replay — captured expect=IDLE matches.
checker.enable_assert()
graph.replay()
torch.cuda.synchronize()
_DEVICE_MODULE.synchronize()
self.assertEqual(_phase_value(checker), int(_Phase.B))
# Reset + replay again — same result, no raise.
checker._reset_to_idle()
graph.replay()
torch.cuda.synchronize()
_DEVICE_MODULE.synchronize()
self.assertEqual(_phase_value(checker), int(_Phase.B))
def test_assert_flag_toggle_visible_to_replayed_graph(self) -> None:
@@ -433,7 +453,7 @@ class TestCudaGraphCapture(CustomTestCase):
# Replay with assert OFF tolerates a deliberately diverged phase.
checker._phase.fill_(999)
graph.replay()
torch.cuda.synchronize() # no raise — flag is OFF
_DEVICE_MODULE.synchronize() # no raise -- flag is OFF
self.assertEqual(_phase_value(checker), int(_Phase.A))
# Now turn on asserts (also resets phase -> IDLE) and replay.
@@ -442,7 +462,7 @@ class TestCudaGraphCapture(CustomTestCase):
self.assertEqual(_phase_value(checker), int(_Phase.IDLE))
graph.replay()
torch.cuda.synchronize() # no raise — phase matched expect
_DEVICE_MODULE.synchronize() # no raise -- phase matched expect
self.assertEqual(_phase_value(checker), int(_Phase.A))
@@ -451,13 +471,13 @@ class TestPhaseReprNoCrash(CustomTestCase):
@classmethod
def setUpClass(cls) -> None:
cls.device = torch.device("cuda:0")
cls.device = _DEVICE
def test_update_with_int_phases_does_not_crash(self) -> None:
checker = SimplePhaseChecker(initial_phase=0, device=self.device)
checker.enable_assert()
checker.update(expect_phase=0, next_phase=1, caller_name="ints")
torch.cuda.synchronize()
_DEVICE_MODULE.synchronize()
self.assertEqual(_phase_value(checker), 1)
def test_update_with_intenum_phases_does_not_crash(self) -> None:
@@ -466,7 +486,7 @@ class TestPhaseReprNoCrash(CustomTestCase):
checker.update(
expect_phase=_Phase.IDLE, next_phase=_Phase.A, caller_name="enums"
)
torch.cuda.synchronize()
_DEVICE_MODULE.synchronize()
self.assertEqual(_phase_value(checker), int(_Phase.A))
def test_update_mixing_int_and_intenum_phases(self) -> None:
@@ -474,10 +494,10 @@ class TestPhaseReprNoCrash(CustomTestCase):
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
checker.enable_assert()
checker.update(expect_phase=_Phase.IDLE, next_phase=5, caller_name="mix1")
torch.cuda.synchronize()
_DEVICE_MODULE.synchronize()
self.assertEqual(_phase_value(checker), 5)
checker.update(expect_phase=5, next_phase=_Phase.IDLE, caller_name="mix2")
torch.cuda.synchronize()
_DEVICE_MODULE.synchronize()
self.assertEqual(_phase_value(checker), int(_Phase.IDLE))
test/registered/xpu/test_xpu_kv_canary.py
+88
View File
@@ -0,0 +1,88 @@
"""KV-canary end-to-end on Intel XPU.
Exercises ``--kv-canary`` on ``--device xpu``, where the write / verify /
plan-entries kernels are CUDA-JIT only, so they route to their torch references
via ``kv_canary._dispatch.use_torch_reference`` and the D2H stream/event
machinery runs through ``torch.xpu``.
Both directions are needed: a dispatch shim that silently no-oped would pass the
baseline too, so only an injected corruption going *undetected* separates a
working fallback from a dead one.
"""
from __future__ import annotations
import unittest
from sglang.srt.kv_canary.config import CanaryMode
from sglang.test.ci.ci_register import register_xpu_ci
from sglang.test.kv_canary.e2e_base import CanaryE2EBase
register_xpu_ci(est_time=600, suite="stage-b-test-1-gpu-xpu")
# --disable-cuda-graph is mandatory, not tuning: install_canary refuses a captured decode
# on a device that routes to the torch reference (host work and D2H, so replay checks nothing).
_XPU_SERVER_ARGS = ("--device", "xpu", "--disable-cuda-graph")
class _XPUCanaryE2EBase(CanaryE2EBase):
"""Shared XPU server config for the cases below.
The torch reference folds the chain slot-by-slot on the host, so it runs orders
of magnitude slower than the CUDA kernels; this subclass shrinks the workload
rather than the shared base, which stays on its CUDA-tuned defaults.
"""
model_mode = "mha"
kv_canary_mode = CanaryMode.LOG
extra_server_args = _XPU_SERVER_ARGS
# Enough decode steps for the chain to span several forwards; measured at roughly
# 3 tok/s on the reference path, so the timeout is generous rather than tight.
default_parallel_n = 2
default_max_new_tokens = 32
default_request_timeout = 120.0
class TestXPUCanaryBaseline(_XPUCanaryE2EBase):
"""Clean XPU canary run: no violations, all requests succeed."""
def test_no_violation(self) -> None:
self.send_parallel_requests()
self.assert_no_violation(wait_seconds=2.0)
class TestXPUCanaryRealKvBaseline(_XPUCanaryE2EBase):
"""Clean run with real-KV fingerprinting on, the reference's other fold path.
``--kv-canary-real-data partial`` is what makes verify/write read the KV pool
itself; without a case that sets it, the reference's real-KV gather stays
unexecuted on XPU no matter how many chain-only cases pass.
"""
extra_server_args = (*_XPU_SERVER_ARGS, "--kv-canary-real-data", "partial")
def test_no_violation(self) -> None:
self.send_parallel_requests()
self.assert_no_violation(wait_seconds=2.0)
class TestXPUCanaryPerturbDetected(_XPUCanaryE2EBase):
"""Injected req_to_token corruption must be detected on XPU."""
extra_env = {
# Every forward, so the short reference workload cannot end before it fires.
"SGLANG_KV_CANARY_PERTURB_REQ_TO_TOKEN_PROB": "1.0",
"SGLANG_KV_CANARY_PERTURB_WARMUP_STEPS": "0",
# Corrupting the slot mapping looks like a pool leak to the on-idle checker.
# Expected here, so strict mode stays off or the scheduler crashes before we
# can assert.
"SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE": "0",
}
def test_req_to_token_perturbation_reports_chain_hash_violation(self) -> None:
self.send_parallel_requests()
self.assert_per_forward_violation_reported(fail_reason="verify_chain_hash")
if __name__ == "__main__":
unittest.main()