From 35eb7cf8d62396e9bfd16900c1dc97d76e0cab02 Mon Sep 17 00:00:00 2001 From: Dayananda V Date: Tue, 22 Sep 2026 06:49:01 +0530 Subject: [PATCH] [Intel][XPU][KVCanary] Enable KV Canary on Intel XPU (#33520) Co-authored-by: Claude Opus 5 --- .../sglang/kernels/ops/kv_canary/_dispatch.py | 12 ++ .../sglang/kernels/ops/kv_canary/plan/api.py | 37 ++++++- python/sglang/kernels/ops/kv_canary/verify.py | 30 ++++- .../kernels/ops/kv_canary/verify_ref.py | 104 +++++++++++++----- python/sglang/kernels/ops/kv_canary/write.py | 25 ++++- .../sglang/kernels/ops/kv_canary/write_ref.py | 17 ++- python/sglang/srt/kv_canary/api.py | 23 ++++ .../srt/kv_canary/runner/canary_manager.py | 3 +- .../srt/kv_canary/runner/enable_warner.py | 2 +- .../srt/kv_canary/runner/future_tensor.py | 28 +++-- .../srt/kv_canary/runner/health_checker.py | 2 +- .../srt/kv_canary/runner/stats_logger.py | 2 +- .../srt/kv_canary/runner/swa_divergence.py | 2 +- .../srt/kv_canary/runner/violation_manager.py | 2 +- .../single_forward_manager/manager.py | 2 +- python/sglang/srt/models/qwen3.py | 5 +- python/sglang/srt/utils/common.py | 11 ++ python/sglang/srt/utils/phase_checker.py | 4 +- python/sglang/test/kv_canary/e2e_base.py | 24 +++- python/sglang/test/kv_canary/fixtures.py | 8 +- .../manual/kv_canary/test_xpu_kv_canary_pp.py | 36 ++++++ .../kernels/ops/kv_canary/test_verify_hand.py | 36 +++++- .../test_self_unit_capture_refusal.py | 77 +++++++++++++ .../kv_canary/test_self_unit_future_tensor.py | 77 ++++++------- .../kv_canary/test_self_unit_plan_input.py | 10 +- ..._unit_req_to_expected_token_ids_manager.py | 19 +++- .../test_self_unit_runner_swa_divergence.py | 104 +++++++++++------- test/registered/utils/test_phase_checker.py | 90 +++++++++------ test/registered/xpu/test_xpu_kv_canary.py | 88 +++++++++++++++ 29 files changed, 690 insertions(+), 190 deletions(-) create mode 100644 python/sglang/kernels/ops/kv_canary/_dispatch.py create mode 100644 test/manual/kv_canary/test_xpu_kv_canary_pp.py create mode 100644 test/registered/kv_canary/test_self_unit_capture_refusal.py create mode 100644 test/registered/xpu/test_xpu_kv_canary.py diff --git a/python/sglang/kernels/ops/kv_canary/_dispatch.py b/python/sglang/kernels/ops/kv_canary/_dispatch.py new file mode 100644 index 000000000..f5775ebb6 --- /dev/null +++ b/python/sglang/kernels/ops/kv_canary/_dispatch.py @@ -0,0 +1,12 @@ +from __future__ import annotations + +import torch + + +def use_torch_reference(device: torch.device) -> bool: + """Whether a canary launcher must fall back to its byte-equal torch reference. + + The write / verify / plan-entries kernels are CUDA-JIT only; HIP keeps them + since torch reports it as ``"cuda"``. XPU / CPU / anything else falls back. + """ + return device.type != "cuda" diff --git a/python/sglang/kernels/ops/kv_canary/plan/api.py b/python/sglang/kernels/ops/kv_canary/plan/api.py index 84d34fd52..19ab6ef61 100644 --- a/python/sglang/kernels/ops/kv_canary/plan/api.py +++ b/python/sglang/kernels/ops/kv_canary/plan/api.py @@ -4,6 +4,7 @@ from typing import Optional import torch +from sglang.kernels.ops.kv_canary._dispatch import use_torch_reference from sglang.kernels.ops.kv_canary.plan.entries_kernel import ( launch_plan_entries_kernel, ) @@ -99,6 +100,7 @@ def launch_canary_plan_kernels( Calling contract: - Pure side-effect; no host work, no D2H. - Safe in cuda-graph capture; caller refills all input tensors in-place before replay. + The reference path is not (host work, D2H) and must not be launched under capture. - The wrapper launches the plan sub-kernels needed to fill both plans end-to-end. - Padding rows contribute zero entries. @@ -106,17 +108,42 @@ def launch_canary_plan_kernels( :func:`sglang.kernels.ops.kv_canary.plan_ref.launch_canary_plan_kernels_torch_reference`; both the Triton offsets kernel and the CUDA JIT entries kernel must match byte-for-byte. """ + # SWA plans are meaningless without the full->swa LUT (entries would carry + # untranslated full-pool slots), so this is a cross-backend contract, not a + # CUDA-only guard. Enforce it before dispatching: the torch reference does not + # re-check, so leaving it below the early-return would silently skip it. + if swa_window_size > 0 and full_to_swa_index_mapping is None: + raise ValueError( + "kv-canary: launch_canary_plan_kernels requires full_to_swa_index_mapping when swa_window_size > 0" + ) + + if use_torch_reference(verify_plan_out.verify_slot_indices.device): + from sglang.kernels.ops.kv_canary.plan_ref import ( + launch_canary_plan_kernels_torch_reference, + ) + + launch_canary_plan_kernels_torch_reference( + verify_plan_out=verify_plan_out, + write_plan_out=write_plan_out, + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + extend_seq_lens=extend_seq_lens, + req_to_token=req_to_token, + swa_window_size=swa_window_size, + full_to_swa_index_mapping=full_to_swa_index_mapping, + verify_capacity=verify_capacity, + req_to_verify_expected_tokens=req_to_verify_expected_tokens, + req_to_verify_expected_tokens_valid_lens=req_to_verify_expected_tokens_valid_lens, + kv_token_id_vs_position_offset=kv_token_id_vs_position_offset, + ) + return + bs = int(req_pool_indices.shape[0]) if bs > _PLAN_BS_BLOCK_SIZE: raise ValueError( f"kv-canary: launch_canary_plan_kernels supports at most bs={_PLAN_BS_BLOCK_SIZE} reqs per launch, " f"got bs={bs}. Bump _PLAN_BS_BLOCK_SIZE if real workloads need this." ) - if swa_window_size > 0 and full_to_swa_index_mapping is None: - raise ValueError( - "kv-canary: launch_canary_plan_kernels requires full_to_swa_index_mapping when swa_window_size > 0" - ) - device = verify_plan_out.verify_slot_indices.device verify_offsets_scratch = torch.empty( _PLAN_BS_BLOCK_SIZE + 1, dtype=torch.int64, device=device diff --git a/python/sglang/kernels/ops/kv_canary/verify.py b/python/sglang/kernels/ops/kv_canary/verify.py index f44e438cf..38de70a1c 100644 --- a/python/sglang/kernels/ops/kv_canary/verify.py +++ b/python/sglang/kernels/ops/kv_canary/verify.py @@ -8,6 +8,7 @@ import torch from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args from sglang.kernels.ops.kv_canary import consts +from sglang.kernels.ops.kv_canary._dispatch import use_torch_reference if TYPE_CHECKING: from tvm_ffi.module import Module @@ -120,6 +121,16 @@ class RealKvSource: f"got {row_stride_bytes} bytes (shape={tuple(self.tensor.shape)}, " f"dtype={self.tensor.dtype})" ) + # A row is addressed as page_size slots of num_bytes_per_token, unchecked at fold time; + # a narrower row hashes fewer bytes than asked and still reports the chain clean. + min_row_bytes = self.page_size * self.num_bytes_per_token + if row_stride_bytes < min_row_bytes: + raise ValueError( + f"kv-canary: RealKvSource.tensor dim-1 is {row_stride_bytes} bytes but " + f"page_size={self.page_size} x num_bytes_per_token={self.num_bytes_per_token} " + f"needs {min_row_bytes} (shape={tuple(self.tensor.shape)}, " + f"dtype={self.tensor.dtype})" + ) @dataclass(frozen=True, slots=True, kw_only=True) @@ -301,7 +312,8 @@ def launch_canary_verify_kernel( - Pure side-effect; never raises. Host polls violation_write_index[0] > 0 for is_errored and violation_ring[0] for the first violation. - kernel_run_counter is bumped every call (canary-ran health signal). - - Safe in cuda-graph capture; caller refills plan in-place before replay. + - Safe in cuda-graph capture; caller refills plan in-place before replay. The reference + path is not (host work, D2H) and must not be launched under capture. Pinned by torch reference :func:`sglang.kernels.ops.kv_canary.verify_ref.launch_canary_verify_kernel_torch_reference`; CUDA must match @@ -309,12 +321,28 @@ def launch_canary_verify_kernel( """ canary_buf = context.canary_buf real_kv_sources = context.real_kv_sources + # Enforce the source-count cap before dispatching: the torch reference is + # pinned to match the CUDA ABI byte-for-byte, so the limit is a cross-backend + # contract, not a CUDA-only guard. Checking after the reference early-return + # (XPU / CPU path) would silently skip it. if len(real_kv_sources) > consts.MAX_REAL_KV_SOURCES: raise ValueError( f"kv-canary: at most {consts.MAX_REAL_KV_SOURCES} RealKvSource entries supported by the CUDA ABI, " f"got {len(real_kv_sources)}" ) + if use_torch_reference(canary_buf.device): + from sglang.kernels.ops.kv_canary.verify_ref import ( + launch_canary_verify_kernel_torch_reference, + ) + + launch_canary_verify_kernel_torch_reference( + context=context, + plan=plan, + check_verify_expected_token=check_verify_expected_token, + ) + return + _assert_contiguous(canary_buf, "canary_buf") _assert_contiguous(plan.verify_slot_indices, "plan.verify_slot_indices") _assert_contiguous(plan.verify_expected_tokens, "plan.verify_expected_tokens") diff --git a/python/sglang/kernels/ops/kv_canary/verify_ref.py b/python/sglang/kernels/ops/kv_canary/verify_ref.py index b759d2106..7dd9a5e41 100644 --- a/python/sglang/kernels/ops/kv_canary/verify_ref.py +++ b/python/sglang/kernels/ops/kv_canary/verify_ref.py @@ -1,5 +1,7 @@ from __future__ import annotations +from typing import NamedTuple, Sequence + import torch from sglang.kernels.ops.kv_canary import consts @@ -93,6 +95,13 @@ def launch_canary_verify_kernel_torch_reference( f"kv-canary: canary_buf slot stride must hold at least 4 int64 fields, got {slot_stride_i64}" ) + host_real_kv_sources = materialize_real_kv_sources( + real_kv_sources=real_kv_sources, + real_kv_hash_mode=real_kv_hash_mode, + slot_indices=slot_indices_list, + work_device=work_device, + ) + violation_rows: list[list[int]] = [] for k in range(active): @@ -118,9 +127,7 @@ def launch_canary_verify_kernel_torch_reference( expected_real_kv_hash_u64 = _compute_real_kv_hash_scalar( slot_idx=slot_idx, - real_kv_sources=real_kv_sources, - real_kv_hash_mode=real_kv_hash_mode, - work_device=work_device, + host_sources=host_real_kv_sources, ) expected_real_kv_hash = _to_signed_int64(expected_real_kv_hash_u64) @@ -190,37 +197,84 @@ def compute_slot_hash(buf_i64: torch.Tensor, source_slot_idx: int) -> int: return splitmix64_mix3(prev_hash, token, position) +class _MaterializedRealKvSource(NamedTuple): + """A ``RealKvSource`` narrowed to the rows one launch reads, on ``work_device``. + + ``row_lookup`` maps a source row (``slot_idx // page_size``) to its index in + ``tensor_u8``, which holds only the gathered rows. + """ + + tensor_u8: torch.Tensor + row_lookup: dict[int, int] + page_size: int + num_bytes_per_token: int + effective_read_bytes: int + + +def materialize_real_kv_sources( + *, + real_kv_sources: tuple[RealKvSource, ...], + real_kv_hash_mode: consts.RealKvHashMode, + slot_indices: Sequence[int], + work_device: torch.device, +) -> tuple[_MaterializedRealKvSource, ...]: + """Gather each source's read rows onto ``work_device`` once per launch. + + An empty tuple means nothing to hash; callers skip the per-slot fold.""" + mode = int(real_kv_hash_mode) + if ( + mode == int(consts.RealKvHashMode.NONE) + or len(real_kv_sources) == 0 + or len(slot_indices) == 0 + ): + return () + + materialized: list[_MaterializedRealKvSource] = [] + for source in real_kv_sources: + # Gather on device first: copying the whole source is a KV-layer-sized transfer + # (one row per token of the pool) on every launch. + rows = sorted({slot_idx // source.page_size for slot_idx in slot_indices}) + row_index = torch.tensor(rows, dtype=torch.int64, device=source.tensor.device) + tensor_u8 = ( + source.tensor.detach() + .index_select(0, row_index) + .to(device=work_device) + .contiguous() + .view(torch.uint8) + ) + effective_read_bytes = ( + 16 if mode == int(consts.RealKvHashMode.PARTIAL) else source.read_bytes + ) + materialized.append( + _MaterializedRealKvSource( + tensor_u8=tensor_u8, + row_lookup={row: i for i, row in enumerate(rows)}, + page_size=source.page_size, + num_bytes_per_token=source.num_bytes_per_token, + effective_read_bytes=effective_read_bytes, + ) + ) + return tuple(materialized) + + def _compute_real_kv_hash_scalar( *, slot_idx: int, - real_kv_sources: tuple[RealKvSource, ...], - real_kv_hash_mode: consts.RealKvHashMode, - work_device: torch.device, + host_sources: tuple[_MaterializedRealKvSource, ...], ) -> int: - mode = int(real_kv_hash_mode) - if mode == int(consts.RealKvHashMode.NONE) or len(real_kv_sources) == 0: + if len(host_sources) == 0: return 0 acc: int = 0 - for source in real_kv_sources: - page_size = source.page_size - num_bytes_per_token = source.num_bytes_per_token - read_bytes = source.read_bytes - tensor_u8 = ( - source.tensor.detach().to(device=work_device).contiguous().view(torch.uint8) - ) + for source in host_sources: + row = source.row_lookup[slot_idx // source.page_size] + col_within_page = slot_idx % source.page_size + col_start = col_within_page * source.num_bytes_per_token - row = slot_idx // page_size - col_within_page = slot_idx % page_size - col_start = col_within_page * num_bytes_per_token - - effective_read_bytes = ( - 16 if mode == int(consts.RealKvHashMode.PARTIAL) else read_bytes - ) - raw_bytes: list[int] = [] - for b in range(effective_read_bytes): - raw_bytes.append(int(tensor_u8[row, col_start + b].item())) + raw_bytes = source.tensor_u8[ + row, col_start : col_start + source.effective_read_bytes + ].tolist() source_hash = _splitmix64_fold_bytes_scalar(raw_bytes=raw_bytes) diff --git a/python/sglang/kernels/ops/kv_canary/write.py b/python/sglang/kernels/ops/kv_canary/write.py index 6b68fa5e2..40cdfd35c 100644 --- a/python/sglang/kernels/ops/kv_canary/write.py +++ b/python/sglang/kernels/ops/kv_canary/write.py @@ -7,6 +7,7 @@ import torch from sglang.kernels.jit.utils import cache_once, load_jit from sglang.kernels.ops.kv_canary import consts +from sglang.kernels.ops.kv_canary._dispatch import use_torch_reference from sglang.kernels.ops.kv_canary.verify import ( VerifyOrWriteContext, _assert_contiguous, @@ -183,7 +184,8 @@ def launch_canary_write_kernel( - Input-verification mismatch records violations but does NOT abort the chain. - kernel_run_counter is bumped every call. - Safe in cuda-graph capture; caller refills input_ids / positions / out_cache_loc / plan - in-place before replay. + in-place before replay. The reference path is not (host work, D2H) and must not be + launched under capture. Pinned by torch reference :func:`sglang.kernels.ops.kv_canary.write_ref.launch_canary_write_kernel_torch_reference`; CUDA must match @@ -191,12 +193,33 @@ def launch_canary_write_kernel( """ canary_buf = context.canary_buf real_kv_sources = context.real_kv_sources + # Enforce the source-count cap before dispatching: the torch reference is + # pinned to match the CUDA ABI byte-for-byte, so the limit is a cross-backend + # contract, not a CUDA-only guard. Checking after the reference early-return + # (XPU / CPU path) would silently skip it. if len(real_kv_sources) > consts.MAX_REAL_KV_SOURCES: raise ValueError( f"kv-canary: at most {consts.MAX_REAL_KV_SOURCES} RealKvSource entries supported by the CUDA ABI, " f"got {len(real_kv_sources)}" ) + if use_torch_reference(canary_buf.device): + from sglang.kernels.ops.kv_canary.write_ref import ( + launch_canary_write_kernel_torch_reference, + ) + + launch_canary_write_kernel_torch_reference( + context=context, + plan=plan, + input_ids=input_ids, + positions=positions, + out_cache_loc=out_cache_loc, + enable_write_input_assert=enable_write_input_assert, + expected_input_tokens=expected_input_tokens, + expected_input_positions=expected_input_positions, + ) + return + _assert_contiguous(canary_buf, "canary_buf") _assert_contiguous(plan.write_offsets, "plan.write_offsets") _assert_contiguous(plan.write_seed_slot_indices, "plan.write_seed_slot_indices") diff --git a/python/sglang/kernels/ops/kv_canary/write_ref.py b/python/sglang/kernels/ops/kv_canary/write_ref.py index 01b391dcc..f8b042581 100644 --- a/python/sglang/kernels/ops/kv_canary/write_ref.py +++ b/python/sglang/kernels/ops/kv_canary/write_ref.py @@ -10,6 +10,7 @@ from sglang.kernels.ops.kv_canary.verify_ref import ( _compute_real_kv_hash_scalar, _to_signed_int64, compute_slot_hash, + materialize_real_kv_sources, splitmix64_mix3, ) from sglang.kernels.ops.kv_canary.write import WritePlan @@ -96,6 +97,18 @@ def launch_canary_write_kernel_torch_reference( expected_input_tokens_host = None expected_input_positions_host = None + # A superset of the slots the loop below folds: the per-req entry ranges all lie + # inside [0, total_entries), and gathering a spare row is harmless. + write_slot_indices = [ + slot for slot in out_cache_loc_host[:total_entries].tolist() if slot >= 0 + ] + host_real_kv_sources = materialize_real_kv_sources( + real_kv_sources=real_kv_sources, + real_kv_hash_mode=real_kv_hash_mode, + slot_indices=write_slot_indices, + work_device=work_device, + ) + violation_rows: list[list[int]] = [] total_slots_written = 0 @@ -129,9 +142,7 @@ def launch_canary_write_kernel_torch_reference( real_kv_hash_u64 = _compute_real_kv_hash_scalar( slot_idx=slot, - real_kv_sources=real_kv_sources, - real_kv_hash_mode=real_kv_hash_mode, - work_device=work_device, + host_sources=host_real_kv_sources, ) if enable_write_input_assert: diff --git a/python/sglang/srt/kv_canary/api.py b/python/sglang/srt/kv_canary/api.py index 55d8102a1..da9032f7b 100644 --- a/python/sglang/srt/kv_canary/api.py +++ b/python/sglang/srt/kv_canary/api.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any, Callable, Optional import torch +from sglang.kernels.ops.kv_canary._dispatch import use_torch_reference from sglang.srt.kv_canary.capacities import CanaryLaunchCapacities from sglang.srt.kv_canary.config import CanaryConfig, CanaryMode from sglang.srt.kv_canary.perturb.config import PerturbConfig @@ -18,6 +19,7 @@ from sglang.srt.model_executor.cuda_graph_config import ( check_cuda_graph_backend, ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.platforms import current_platform from sglang.srt.runtime_context import ( get_disagg, get_parallel, @@ -32,6 +34,21 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +def torch_reference_conflicts_with_decode_graph(device: torch.device) -> bool: + """Whether ``device`` would capture a decode graph over the canary torch reference. + + install_canary runs before capture, and the reference does host work and D2H, so its + launches never land in the graph and every replayed decode verifies clean. + """ + # An unpublished cuda_graph_config reads as "not disabled" here, so this refuses rather + # than waves through: a startup error beats a canary that reports clean forever. + return ( + use_torch_reference(device) + and current_platform.support_cuda_graph() + and not check_cuda_graph_backend(Phase.DECODE, Backend.DISABLED) + ) + + def install_canary( *, server_args: ServerArgs, @@ -50,6 +67,12 @@ def install_canary( perturb_config = PerturbConfig.from_env() device = torch.device(model_runner.device) + if torch_reference_conflicts_with_decode_graph(device): + raise ValueError( + f"kv-canary: {device.type} has no canary CUDA kernels and its torch reference " + "cannot be graph-captured; pass --disable-cuda-graph (or " + "--cuda-graph-backend-decode=disabled) when canary is enabled" + ) # EAGLE draft worker pools rotate input_ids so slot ``p`` stores K/V for the token at position ``p+1``; # target pools have no such shift. Threaded into the plan-side expected-token gather kernel. kv_token_id_vs_position_offset = 1 if model_runner.is_draft_worker else 0 diff --git a/python/sglang/srt/kv_canary/runner/canary_manager.py b/python/sglang/srt/kv_canary/runner/canary_manager.py index 66bf8d0fc..79f4527e6 100644 --- a/python/sglang/srt/kv_canary/runner/canary_manager.py +++ b/python/sglang/srt/kv_canary/runner/canary_manager.py @@ -29,6 +29,7 @@ from sglang.srt.kv_canary.single_forward_manager.manager import ( ) from sglang.srt.kv_canary.state import CanaryDeviceState from sglang.srt.kv_canary.token_oracle.oracle_manager import TokenOracleManager +from sglang.srt.utils import create_device_stream if TYPE_CHECKING: from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator @@ -91,7 +92,7 @@ class CanaryManager: ) ) - self._d2h_stream: torch.cuda.Stream = torch.cuda.Stream(device=device) + self._d2h_stream: torch.Stream = create_device_stream(device) swa_divergence_interval = ( envs.SGLANG_KV_CANARY_SWA_DIVERGENCE_STATS_INTERVAL.get() diff --git a/python/sglang/srt/kv_canary/runner/enable_warner.py b/python/sglang/srt/kv_canary/runner/enable_warner.py index ca922859b..69e3af420 100644 --- a/python/sglang/srt/kv_canary/runner/enable_warner.py +++ b/python/sglang/srt/kv_canary/runner/enable_warner.py @@ -12,7 +12,7 @@ logger = logging.getLogger(__name__) class CanaryEnableWarner: def __init__( - self, *, verify_capacity: int, d2h_stream: Optional[torch.cuda.Stream] + self, *, verify_capacity: int, d2h_stream: Optional[torch.Stream] ) -> None: self._verify_capacity = verify_capacity self._overflow_count_total: int = 0 diff --git a/python/sglang/srt/kv_canary/runner/future_tensor.py b/python/sglang/srt/kv_canary/runner/future_tensor.py index eff94502c..32fb64852 100644 --- a/python/sglang/srt/kv_canary/runner/future_tensor.py +++ b/python/sglang/srt/kv_canary/runner/future_tensor.py @@ -6,6 +6,8 @@ from typing import Any, Optional, Union import torch +from sglang.srt.utils import is_device_stream_capturing + _PayloadDict = dict[str, Any] _TensorOrDict = Union[torch.Tensor, _PayloadDict] @@ -15,20 +17,14 @@ _DUMMY_DICT_KEY = "__dummy_key__" @dataclass(slots=True, kw_only=True) class FutureTensors: _data: Optional[_PayloadDict] - _event: Optional[torch.cuda.Event] + _event: Optional[torch.Event] # Device-source clones must outlive the async d2h copy. _retained_device_clones: Optional[dict[str, torch.Tensor]] = None @classmethod def device_to_host( - cls, xs_device: _TensorOrDict, *, d2h_stream: torch.cuda.Stream + cls, xs_device: _TensorOrDict, *, d2h_stream: torch.Stream ) -> FutureTensors: - assert not torch.cuda.is_current_stream_capturing(), ( - "FutureTensors.device_to_host must not be called during cuda-graph " - "capture: the d2h side-stream copy + pinned-host alloc cannot be " - "captured. Upper-layer callers are responsible for placing the d2h " - "staging OUTSIDE the cuda graph (not inside it)." - ) if not isinstance(xs_device, dict): xs_device = {_DUMMY_DICT_KEY: xs_device} @@ -43,6 +39,13 @@ class FutureTensors: device = first_tensor.device del first_tensor + assert not is_device_stream_capturing(device), ( + "FutureTensors.device_to_host must not be called during cuda-graph " + "capture: the d2h side-stream copy + pinned-host alloc cannot be " + "captured. Upper-layer callers are responsible for placing the d2h " + "staging OUTSIDE the cuda graph (not inside it)." + ) + tensors_device = { k: v for k, v in xs_device.items() if isinstance(v, torch.Tensor) } @@ -61,11 +64,12 @@ class FutureTensors: for key, x in tensors_device.items() } - d2h_stream.wait_stream(torch.cuda.current_stream(device)) - with torch.cuda.stream(d2h_stream): + device_module = torch.get_device_module(device) + d2h_stream.wait_stream(device_module.current_stream(device)) + with device_module.stream(d2h_stream): for key in tensors_device_cloned: tensors_host[key].copy_(tensors_device_cloned[key], non_blocking=True) - event = torch.cuda.Event() + event = device_module.Event() event.record() return cls( @@ -100,7 +104,7 @@ class FutureTensors: class DelayedDeviceHostHandler: """Stage device-side compute at step T, drain + postprocess host copy at step T+1.""" - d2h_stream: torch.cuda.Stream + d2h_stream: torch.Stream _future: Optional[FutureTensors] = field(default=None) def step( diff --git a/python/sglang/srt/kv_canary/runner/health_checker.py b/python/sglang/srt/kv_canary/runner/health_checker.py index 7d5535945..62a2b5e23 100644 --- a/python/sglang/srt/kv_canary/runner/health_checker.py +++ b/python/sglang/srt/kv_canary/runner/health_checker.py @@ -34,7 +34,7 @@ class KernelRunCounterHealthChecker: device_state: CanaryDeviceState, active_tags: tuple[CanaryLaunchTag, ...], outer_step_counter_getter: Callable[[], int], - d2h_stream: torch.cuda.Stream, + d2h_stream: torch.Stream, ) -> None: self._config = config self._device_state = device_state diff --git a/python/sglang/srt/kv_canary/runner/stats_logger.py b/python/sglang/srt/kv_canary/runner/stats_logger.py index 5551d7ec8..e38e67b5f 100644 --- a/python/sglang/srt/kv_canary/runner/stats_logger.py +++ b/python/sglang/srt/kv_canary/runner/stats_logger.py @@ -24,7 +24,7 @@ class PeriodicCanaryStatsLogger: active_tags: tuple[CanaryLaunchTag, ...], outer_step_counter_getter: Callable[[], int], sweep_orchestrator: SweepOrchestrator, - d2h_stream: torch.cuda.Stream, + d2h_stream: torch.Stream, ) -> None: self._config = config self._device_state = device_state diff --git a/python/sglang/srt/kv_canary/runner/swa_divergence.py b/python/sglang/srt/kv_canary/runner/swa_divergence.py index d47f12c5a..7dda13930 100644 --- a/python/sglang/srt/kv_canary/runner/swa_divergence.py +++ b/python/sglang/srt/kv_canary/runner/swa_divergence.py @@ -30,7 +30,7 @@ class SwaDivergenceReporter: self, *, device: torch.device, - d2h_stream: torch.cuda.Stream, + d2h_stream: torch.Stream, interval: int, swa_allocator: Optional[SWATokenToKVPoolAllocator] = None, req_to_token_pool: Optional[ReqToTokenPool] = None, diff --git a/python/sglang/srt/kv_canary/runner/violation_manager.py b/python/sglang/srt/kv_canary/runner/violation_manager.py index d533eac43..ecadc885d 100644 --- a/python/sglang/srt/kv_canary/runner/violation_manager.py +++ b/python/sglang/srt/kv_canary/runner/violation_manager.py @@ -16,7 +16,7 @@ class ViolationManager: *, config: CanaryConfig, device_state: CanaryDeviceState, - d2h_stream: torch.cuda.Stream, + d2h_stream: torch.Stream, outer_step_counter_getter: Callable[[], int], ) -> None: self._device_state = device_state diff --git a/python/sglang/srt/kv_canary/single_forward_manager/manager.py b/python/sglang/srt/kv_canary/single_forward_manager/manager.py index 7773eb232..b4b40844f 100644 --- a/python/sglang/srt/kv_canary/single_forward_manager/manager.py +++ b/python/sglang/srt/kv_canary/single_forward_manager/manager.py @@ -71,7 +71,7 @@ class SingleForwardManager: per_forward_verify_capacity: int, per_forward_write_req_capacity: int, per_forward_write_entry_capacity: int, - d2h_stream: torch.cuda.Stream, + d2h_stream: torch.Stream, token_oracle_manager: Optional[TokenOracleManager], swa_divergence_report: Optional[SwaDivergenceReporter], is_eagle_draft_decode: bool, diff --git a/python/sglang/srt/models/qwen3.py b/python/sglang/srt/models/qwen3.py index b4d88886c..dca81ae25 100644 --- a/python/sglang/srt/models/qwen3.py +++ b/python/sglang/srt/models/qwen3.py @@ -30,6 +30,7 @@ from sglang.srt.model_loader.weight_utils import ( from sglang.srt.models.qwen2 import Qwen2MLP as Qwen3MLP from sglang.srt.models.qwen2 import Qwen2Model from sglang.srt.models.utils import apply_qk_norm +from sglang.srt.platforms import current_platform from sglang.srt.runtime_context import get_exec, get_parallel, get_stream from sglang.srt.utils import add_prefix, get_bool_env_var, is_cuda, is_hip, is_npu @@ -676,8 +677,8 @@ class Qwen3ForCausalLM(nn.Module): del self.lm_head.weight self.model.embed_tokens.weight = embed self.lm_head.weight = head - torch.cuda.empty_cache() - torch.cuda.synchronize() + current_platform.empty_cache() + current_platform.synchronize() def load_kv_cache_scales(self, quantization_param_path: str) -> None: self.model.load_kv_cache_scales(quantization_param_path) diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index b6e1bafb6..f8d6ec80b 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -608,6 +608,17 @@ def device_stream_context(stream): return torch.get_device_module(stream.device).stream(stream) +def is_device_stream_capturing(device: torch.device) -> bool: + """Whether ``device``'s current stream is mid graph capture (False if unsupported).""" + # Every platform answering support_cuda_graph() already calls + # device_module.is_current_stream_capturing() during capture, so it cannot be missing. + if device.type != current_platform.device_type: + return False + if not current_platform.support_cuda_graph(): + return False + return torch.get_device_module(device).is_current_stream_capturing() + + def get_amdgpu_memory_capacity(): try: # Run rocm-smi and capture the output diff --git a/python/sglang/srt/utils/phase_checker.py b/python/sglang/srt/utils/phase_checker.py index 5546fa0fc..d5c487ceb 100644 --- a/python/sglang/srt/utils/phase_checker.py +++ b/python/sglang/srt/utils/phase_checker.py @@ -7,6 +7,7 @@ import triton import triton.language as tl from sglang.srt.environ import envs +from sglang.srt.utils import is_device_stream_capturing def _phase_repr(phase: int | IntEnum) -> str: @@ -51,6 +52,7 @@ class SimplePhaseChecker: def __init__(self, *, initial_phase: int | IntEnum, device: torch.device) -> None: self._initial_phase = int(initial_phase) + self._device = device self._phase = torch.tensor( self._initial_phase, dtype=torch.int32, device=device ) @@ -81,7 +83,7 @@ class SimplePhaseChecker: f"caller_tag={caller_tag} " f"expect={_phase_repr(expect_phase)} " f"next={_phase_repr(next_phase)} " - f"capturing={torch.cuda.is_current_stream_capturing()}" + f"capturing={is_device_stream_capturing(self._device)}" ) _phase_check_kernel[(1,)]( self._phase, diff --git a/python/sglang/test/kv_canary/e2e_base.py b/python/sglang/test/kv_canary/e2e_base.py index 788576766..69d94262d 100644 --- a/python/sglang/test/kv_canary/e2e_base.py +++ b/python/sglang/test/kv_canary/e2e_base.py @@ -73,6 +73,12 @@ class CanaryE2EBase(CapturedServerE2EBase): # test methods send N sequential batches so the SWA allocator's full→swa index mapping # diverges from identity. Default 1 keeps MHA tests fast. workload_n_batches: ClassVar[int] = 1 + # Default workload for send_parallel_requests, tuned for the CUDA-kernel canary + # path; a slower backend retunes its own subclass here instead of passing sizes at + # every call site. + default_parallel_n: ClassVar[int] = 8 + default_max_new_tokens: ClassVar[int] = 2048 + default_request_timeout: ClassVar[float] = 240.0 _cfg: ClassVar[Optional[_ModeConfig]] = None @@ -121,14 +127,24 @@ class CanaryE2EBase(CapturedServerE2EBase): def send_parallel_requests( self, - n: int = 8, + n: Optional[int] = None, *, assert_all_success: bool = True, - max_new_tokens: int = 2048, - timeout: float = 240.0, + max_new_tokens: Optional[int] = None, + timeout: Optional[float] = None, ignore_eos: Optional[bool] = None, ) -> list[dict]: - """Fan out n parallel /generate requests; return list of response dicts.""" + """Fan out n parallel /generate requests; return list of response dicts. + + Unset sizes fall back to the ``default_*`` class attributes, so a subclass can + retune the whole workload for its backend in one place. + """ + if n is None: + n = self.default_parallel_n + if max_new_tokens is None: + max_new_tokens = self.default_max_new_tokens + if timeout is None: + timeout = self.default_request_timeout if ignore_eos is None: ignore_eos = self.model_mode == "swa" results = post_parallel_generate( diff --git a/python/sglang/test/kv_canary/fixtures.py b/python/sglang/test/kv_canary/fixtures.py index 1125405c0..524a1ffe1 100644 --- a/python/sglang/test/kv_canary/fixtures.py +++ b/python/sglang/test/kv_canary/fixtures.py @@ -15,8 +15,14 @@ from sglang.srt.kv_canary.pool_patcher.adapters.swa import attach_swa from sglang.srt.kv_canary.pool_patcher.api import register_pool_attacher from sglang.srt.mem_cache.radix_cache import RadixCache, TreeNode from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.srt.utils import get_device -DEFAULT_DEVICE: torch.device = torch.device("cuda") +# Resolve the active accelerator (cuda/xpu/...) instead of hardcoding cuda: a torch +# build without CUDA cannot allocate cuda tensors or call torch.cuda.*. Tests that +# need the runtime API (synchronize, streams, ...) go through DEFAULT_DEVICE_MODULE +# rather than torch.cuda. +DEFAULT_DEVICE: torch.device = torch.device(get_device()) +DEFAULT_DEVICE_MODULE = torch.get_device_module(DEFAULT_DEVICE) @dataclass diff --git a/test/manual/kv_canary/test_xpu_kv_canary_pp.py b/test/manual/kv_canary/test_xpu_kv_canary_pp.py new file mode 100644 index 000000000..f6f7a77f7 --- /dev/null +++ b/test/manual/kv_canary/test_xpu_kv_canary_pp.py @@ -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() diff --git a/test/registered/kernels/ops/kv_canary/test_verify_hand.py b/test/registered/kernels/ops/kv_canary/test_verify_hand.py index a1bea876a..a71c53ac6 100644 --- a/test/registered/kernels/ops/kv_canary/test_verify_hand.py +++ b/test/registered/kernels/ops/kv_canary/test_verify_hand.py @@ -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 ] diff --git a/test/registered/kv_canary/test_self_unit_capture_refusal.py b/test/registered/kv_canary/test_self_unit_capture_refusal.py new file mode 100644 index 000000000..84a1009c2 --- /dev/null +++ b/test/registered/kv_canary/test_self_unit_capture_refusal.py @@ -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() diff --git a/test/registered/kv_canary/test_self_unit_future_tensor.py b/test/registered/kv_canary/test_self_unit_future_tensor.py index a887b6f8c..6b6cb6e66 100644 --- a/test/registered/kv_canary/test_self_unit_future_tensor.py +++ b/test/registered/kv_canary/test_self_unit_future_tensor.py @@ -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) diff --git a/test/registered/kv_canary/test_self_unit_plan_input.py b/test/registered/kv_canary/test_self_unit_plan_input.py index eb25077ba..98073e2ef 100644 --- a/test/registered/kv_canary/test_self_unit_plan_input.py +++ b/test/registered/kv_canary/test_self_unit_plan_input.py @@ -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] ) diff --git a/test/registered/kv_canary/test_self_unit_req_to_expected_token_ids_manager.py b/test/registered/kv_canary/test_self_unit_req_to_expected_token_ids_manager.py index a5f615cf5..575336628 100644 --- a/test/registered/kv_canary/test_self_unit_req_to_expected_token_ids_manager.py +++ b/test/registered/kv_canary/test_self_unit_req_to_expected_token_ids_manager.py @@ -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]) diff --git a/test/registered/kv_canary/test_self_unit_runner_swa_divergence.py b/test/registered/kv_canary/test_self_unit_runner_swa_divergence.py index 1fc1887b0..2c1a350db 100644 --- a/test/registered/kv_canary/test_self_unit_runner_swa_divergence.py +++ b/test/registered/kv_canary/test_self_unit_runner_swa_divergence.py @@ -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), ) diff --git a/test/registered/utils/test_phase_checker.py b/test/registered/utils/test_phase_checker.py index 9c9675a63..d7bf6a12d 100644 --- a/test/registered/utils/test_phase_checker.py +++ b/test/registered/utils/test_phase_checker.py @@ -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)) diff --git a/test/registered/xpu/test_xpu_kv_canary.py b/test/registered/xpu/test_xpu_kv_canary.py new file mode 100644 index 000000000..f316078b1 --- /dev/null +++ b/test/registered/xpu/test_xpu_kv_canary.py @@ -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()