[Intel][XPU][KVCanary] Enable KV Canary on Intel XPU (#33520)
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5
parent
046cd6f4ea
commit
35eb7cf8d6
@@ -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"
|
||||||
@@ -4,6 +4,7 @@ from typing import Optional
|
|||||||
|
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
|
from sglang.kernels.ops.kv_canary._dispatch import use_torch_reference
|
||||||
from sglang.kernels.ops.kv_canary.plan.entries_kernel import (
|
from sglang.kernels.ops.kv_canary.plan.entries_kernel import (
|
||||||
launch_plan_entries_kernel,
|
launch_plan_entries_kernel,
|
||||||
)
|
)
|
||||||
@@ -99,6 +100,7 @@ def launch_canary_plan_kernels(
|
|||||||
Calling contract:
|
Calling contract:
|
||||||
- Pure side-effect; no host work, no D2H.
|
- Pure side-effect; no host work, no D2H.
|
||||||
- Safe in cuda-graph capture; caller refills all input tensors in-place before replay.
|
- 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.
|
- The wrapper launches the plan sub-kernels needed to fill both plans end-to-end.
|
||||||
- Padding rows contribute zero entries.
|
- 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
|
: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.
|
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])
|
bs = int(req_pool_indices.shape[0])
|
||||||
if bs > _PLAN_BS_BLOCK_SIZE:
|
if bs > _PLAN_BS_BLOCK_SIZE:
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
f"kv-canary: launch_canary_plan_kernels supports at most bs={_PLAN_BS_BLOCK_SIZE} reqs per launch, "
|
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."
|
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
|
device = verify_plan_out.verify_slot_indices.device
|
||||||
verify_offsets_scratch = torch.empty(
|
verify_offsets_scratch = torch.empty(
|
||||||
_PLAN_BS_BLOCK_SIZE + 1, dtype=torch.int64, device=device
|
_PLAN_BS_BLOCK_SIZE + 1, dtype=torch.int64, device=device
|
||||||
|
|||||||
@@ -8,6 +8,7 @@ import torch
|
|||||||
|
|
||||||
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
|
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 import consts
|
||||||
|
from sglang.kernels.ops.kv_canary._dispatch import use_torch_reference
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from tvm_ffi.module import Module
|
from tvm_ffi.module import Module
|
||||||
@@ -120,6 +121,16 @@ class RealKvSource:
|
|||||||
f"got {row_stride_bytes} bytes (shape={tuple(self.tensor.shape)}, "
|
f"got {row_stride_bytes} bytes (shape={tuple(self.tensor.shape)}, "
|
||||||
f"dtype={self.tensor.dtype})"
|
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)
|
@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
|
- Pure side-effect; never raises. Host polls violation_write_index[0] > 0 for is_errored and
|
||||||
violation_ring[0] for the first violation.
|
violation_ring[0] for the first violation.
|
||||||
- kernel_run_counter is bumped every call (canary-ran health signal).
|
- 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
|
Pinned by torch reference
|
||||||
:func:`sglang.kernels.ops.kv_canary.verify_ref.launch_canary_verify_kernel_torch_reference`; CUDA must match
|
: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
|
canary_buf = context.canary_buf
|
||||||
real_kv_sources = context.real_kv_sources
|
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:
|
if len(real_kv_sources) > consts.MAX_REAL_KV_SOURCES:
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
f"kv-canary: at most {consts.MAX_REAL_KV_SOURCES} RealKvSource entries supported by the CUDA ABI, "
|
f"kv-canary: at most {consts.MAX_REAL_KV_SOURCES} RealKvSource entries supported by the CUDA ABI, "
|
||||||
f"got {len(real_kv_sources)}"
|
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(canary_buf, "canary_buf")
|
||||||
_assert_contiguous(plan.verify_slot_indices, "plan.verify_slot_indices")
|
_assert_contiguous(plan.verify_slot_indices, "plan.verify_slot_indices")
|
||||||
_assert_contiguous(plan.verify_expected_tokens, "plan.verify_expected_tokens")
|
_assert_contiguous(plan.verify_expected_tokens, "plan.verify_expected_tokens")
|
||||||
|
|||||||
@@ -1,5 +1,7 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import NamedTuple, Sequence
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
from sglang.kernels.ops.kv_canary import consts
|
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}"
|
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]] = []
|
violation_rows: list[list[int]] = []
|
||||||
|
|
||||||
for k in range(active):
|
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(
|
expected_real_kv_hash_u64 = _compute_real_kv_hash_scalar(
|
||||||
slot_idx=slot_idx,
|
slot_idx=slot_idx,
|
||||||
real_kv_sources=real_kv_sources,
|
host_sources=host_real_kv_sources,
|
||||||
real_kv_hash_mode=real_kv_hash_mode,
|
|
||||||
work_device=work_device,
|
|
||||||
)
|
)
|
||||||
expected_real_kv_hash = _to_signed_int64(expected_real_kv_hash_u64)
|
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)
|
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(
|
def _compute_real_kv_hash_scalar(
|
||||||
*,
|
*,
|
||||||
slot_idx: int,
|
slot_idx: int,
|
||||||
real_kv_sources: tuple[RealKvSource, ...],
|
host_sources: tuple[_MaterializedRealKvSource, ...],
|
||||||
real_kv_hash_mode: consts.RealKvHashMode,
|
|
||||||
work_device: torch.device,
|
|
||||||
) -> int:
|
) -> int:
|
||||||
mode = int(real_kv_hash_mode)
|
if len(host_sources) == 0:
|
||||||
if mode == int(consts.RealKvHashMode.NONE) or len(real_kv_sources) == 0:
|
|
||||||
return 0
|
return 0
|
||||||
|
|
||||||
acc: int = 0
|
acc: int = 0
|
||||||
|
|
||||||
for source in real_kv_sources:
|
for source in host_sources:
|
||||||
page_size = source.page_size
|
row = source.row_lookup[slot_idx // source.page_size]
|
||||||
num_bytes_per_token = source.num_bytes_per_token
|
col_within_page = slot_idx % source.page_size
|
||||||
read_bytes = source.read_bytes
|
col_start = col_within_page * source.num_bytes_per_token
|
||||||
tensor_u8 = (
|
|
||||||
source.tensor.detach().to(device=work_device).contiguous().view(torch.uint8)
|
|
||||||
)
|
|
||||||
|
|
||||||
row = slot_idx // page_size
|
raw_bytes = source.tensor_u8[
|
||||||
col_within_page = slot_idx % page_size
|
row, col_start : col_start + source.effective_read_bytes
|
||||||
col_start = col_within_page * num_bytes_per_token
|
].tolist()
|
||||||
|
|
||||||
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()))
|
|
||||||
|
|
||||||
source_hash = _splitmix64_fold_bytes_scalar(raw_bytes=raw_bytes)
|
source_hash = _splitmix64_fold_bytes_scalar(raw_bytes=raw_bytes)
|
||||||
|
|
||||||
|
|||||||
@@ -7,6 +7,7 @@ import torch
|
|||||||
|
|
||||||
from sglang.kernels.jit.utils import cache_once, load_jit
|
from sglang.kernels.jit.utils import cache_once, load_jit
|
||||||
from sglang.kernels.ops.kv_canary import consts
|
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 (
|
from sglang.kernels.ops.kv_canary.verify import (
|
||||||
VerifyOrWriteContext,
|
VerifyOrWriteContext,
|
||||||
_assert_contiguous,
|
_assert_contiguous,
|
||||||
@@ -183,7 +184,8 @@ def launch_canary_write_kernel(
|
|||||||
- Input-verification mismatch records violations but does NOT abort the chain.
|
- Input-verification mismatch records violations but does NOT abort the chain.
|
||||||
- kernel_run_counter is bumped every call.
|
- kernel_run_counter is bumped every call.
|
||||||
- Safe in cuda-graph capture; caller refills input_ids / positions / out_cache_loc / plan
|
- 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
|
Pinned by torch reference
|
||||||
:func:`sglang.kernels.ops.kv_canary.write_ref.launch_canary_write_kernel_torch_reference`; CUDA must match
|
: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
|
canary_buf = context.canary_buf
|
||||||
real_kv_sources = context.real_kv_sources
|
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:
|
if len(real_kv_sources) > consts.MAX_REAL_KV_SOURCES:
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
f"kv-canary: at most {consts.MAX_REAL_KV_SOURCES} RealKvSource entries supported by the CUDA ABI, "
|
f"kv-canary: at most {consts.MAX_REAL_KV_SOURCES} RealKvSource entries supported by the CUDA ABI, "
|
||||||
f"got {len(real_kv_sources)}"
|
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(canary_buf, "canary_buf")
|
||||||
_assert_contiguous(plan.write_offsets, "plan.write_offsets")
|
_assert_contiguous(plan.write_offsets, "plan.write_offsets")
|
||||||
_assert_contiguous(plan.write_seed_slot_indices, "plan.write_seed_slot_indices")
|
_assert_contiguous(plan.write_seed_slot_indices, "plan.write_seed_slot_indices")
|
||||||
|
|||||||
@@ -10,6 +10,7 @@ from sglang.kernels.ops.kv_canary.verify_ref import (
|
|||||||
_compute_real_kv_hash_scalar,
|
_compute_real_kv_hash_scalar,
|
||||||
_to_signed_int64,
|
_to_signed_int64,
|
||||||
compute_slot_hash,
|
compute_slot_hash,
|
||||||
|
materialize_real_kv_sources,
|
||||||
splitmix64_mix3,
|
splitmix64_mix3,
|
||||||
)
|
)
|
||||||
from sglang.kernels.ops.kv_canary.write import WritePlan
|
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_tokens_host = None
|
||||||
expected_input_positions_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]] = []
|
violation_rows: list[list[int]] = []
|
||||||
total_slots_written = 0
|
total_slots_written = 0
|
||||||
|
|
||||||
@@ -129,9 +142,7 @@ def launch_canary_write_kernel_torch_reference(
|
|||||||
|
|
||||||
real_kv_hash_u64 = _compute_real_kv_hash_scalar(
|
real_kv_hash_u64 = _compute_real_kv_hash_scalar(
|
||||||
slot_idx=slot,
|
slot_idx=slot,
|
||||||
real_kv_sources=real_kv_sources,
|
host_sources=host_real_kv_sources,
|
||||||
real_kv_hash_mode=real_kv_hash_mode,
|
|
||||||
work_device=work_device,
|
|
||||||
)
|
)
|
||||||
|
|
||||||
if enable_write_input_assert:
|
if enable_write_input_assert:
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any, Callable, Optional
|
|||||||
|
|
||||||
import torch
|
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.capacities import CanaryLaunchCapacities
|
||||||
from sglang.srt.kv_canary.config import CanaryConfig, CanaryMode
|
from sglang.srt.kv_canary.config import CanaryConfig, CanaryMode
|
||||||
from sglang.srt.kv_canary.perturb.config import PerturbConfig
|
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,
|
check_cuda_graph_backend,
|
||||||
)
|
)
|
||||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
||||||
|
from sglang.srt.platforms import current_platform
|
||||||
from sglang.srt.runtime_context import (
|
from sglang.srt.runtime_context import (
|
||||||
get_disagg,
|
get_disagg,
|
||||||
get_parallel,
|
get_parallel,
|
||||||
@@ -32,6 +34,21 @@ if TYPE_CHECKING:
|
|||||||
logger = logging.getLogger(__name__)
|
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(
|
def install_canary(
|
||||||
*,
|
*,
|
||||||
server_args: ServerArgs,
|
server_args: ServerArgs,
|
||||||
@@ -50,6 +67,12 @@ def install_canary(
|
|||||||
|
|
||||||
perturb_config = PerturbConfig.from_env()
|
perturb_config = PerturbConfig.from_env()
|
||||||
device = torch.device(model_runner.device)
|
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``;
|
# 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.
|
# 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
|
kv_token_id_vs_position_offset = 1 if model_runner.is_draft_worker else 0
|
||||||
|
|||||||
@@ -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.state import CanaryDeviceState
|
||||||
from sglang.srt.kv_canary.token_oracle.oracle_manager import TokenOracleManager
|
from sglang.srt.kv_canary.token_oracle.oracle_manager import TokenOracleManager
|
||||||
|
from sglang.srt.utils import create_device_stream
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
|
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 = (
|
swa_divergence_interval = (
|
||||||
envs.SGLANG_KV_CANARY_SWA_DIVERGENCE_STATS_INTERVAL.get()
|
envs.SGLANG_KV_CANARY_SWA_DIVERGENCE_STATS_INTERVAL.get()
|
||||||
|
|||||||
@@ -12,7 +12,7 @@ logger = logging.getLogger(__name__)
|
|||||||
|
|
||||||
class CanaryEnableWarner:
|
class CanaryEnableWarner:
|
||||||
def __init__(
|
def __init__(
|
||||||
self, *, verify_capacity: int, d2h_stream: Optional[torch.cuda.Stream]
|
self, *, verify_capacity: int, d2h_stream: Optional[torch.Stream]
|
||||||
) -> None:
|
) -> None:
|
||||||
self._verify_capacity = verify_capacity
|
self._verify_capacity = verify_capacity
|
||||||
self._overflow_count_total: int = 0
|
self._overflow_count_total: int = 0
|
||||||
|
|||||||
@@ -6,6 +6,8 @@ from typing import Any, Optional, Union
|
|||||||
|
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
|
from sglang.srt.utils import is_device_stream_capturing
|
||||||
|
|
||||||
_PayloadDict = dict[str, Any]
|
_PayloadDict = dict[str, Any]
|
||||||
_TensorOrDict = Union[torch.Tensor, _PayloadDict]
|
_TensorOrDict = Union[torch.Tensor, _PayloadDict]
|
||||||
|
|
||||||
@@ -15,20 +17,14 @@ _DUMMY_DICT_KEY = "__dummy_key__"
|
|||||||
@dataclass(slots=True, kw_only=True)
|
@dataclass(slots=True, kw_only=True)
|
||||||
class FutureTensors:
|
class FutureTensors:
|
||||||
_data: Optional[_PayloadDict]
|
_data: Optional[_PayloadDict]
|
||||||
_event: Optional[torch.cuda.Event]
|
_event: Optional[torch.Event]
|
||||||
# Device-source clones must outlive the async d2h copy.
|
# Device-source clones must outlive the async d2h copy.
|
||||||
_retained_device_clones: Optional[dict[str, torch.Tensor]] = None
|
_retained_device_clones: Optional[dict[str, torch.Tensor]] = None
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def device_to_host(
|
def device_to_host(
|
||||||
cls, xs_device: _TensorOrDict, *, d2h_stream: torch.cuda.Stream
|
cls, xs_device: _TensorOrDict, *, d2h_stream: torch.Stream
|
||||||
) -> FutureTensors:
|
) -> 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):
|
if not isinstance(xs_device, dict):
|
||||||
xs_device = {_DUMMY_DICT_KEY: xs_device}
|
xs_device = {_DUMMY_DICT_KEY: xs_device}
|
||||||
|
|
||||||
@@ -43,6 +39,13 @@ class FutureTensors:
|
|||||||
device = first_tensor.device
|
device = first_tensor.device
|
||||||
del first_tensor
|
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 = {
|
tensors_device = {
|
||||||
k: v for k, v in xs_device.items() if isinstance(v, torch.Tensor)
|
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()
|
for key, x in tensors_device.items()
|
||||||
}
|
}
|
||||||
|
|
||||||
d2h_stream.wait_stream(torch.cuda.current_stream(device))
|
device_module = torch.get_device_module(device)
|
||||||
with torch.cuda.stream(d2h_stream):
|
d2h_stream.wait_stream(device_module.current_stream(device))
|
||||||
|
with device_module.stream(d2h_stream):
|
||||||
for key in tensors_device_cloned:
|
for key in tensors_device_cloned:
|
||||||
tensors_host[key].copy_(tensors_device_cloned[key], non_blocking=True)
|
tensors_host[key].copy_(tensors_device_cloned[key], non_blocking=True)
|
||||||
event = torch.cuda.Event()
|
event = device_module.Event()
|
||||||
event.record()
|
event.record()
|
||||||
|
|
||||||
return cls(
|
return cls(
|
||||||
@@ -100,7 +104,7 @@ class FutureTensors:
|
|||||||
class DelayedDeviceHostHandler:
|
class DelayedDeviceHostHandler:
|
||||||
"""Stage device-side compute at step T, drain + postprocess host copy at step T+1."""
|
"""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)
|
_future: Optional[FutureTensors] = field(default=None)
|
||||||
|
|
||||||
def step(
|
def step(
|
||||||
|
|||||||
@@ -34,7 +34,7 @@ class KernelRunCounterHealthChecker:
|
|||||||
device_state: CanaryDeviceState,
|
device_state: CanaryDeviceState,
|
||||||
active_tags: tuple[CanaryLaunchTag, ...],
|
active_tags: tuple[CanaryLaunchTag, ...],
|
||||||
outer_step_counter_getter: Callable[[], int],
|
outer_step_counter_getter: Callable[[], int],
|
||||||
d2h_stream: torch.cuda.Stream,
|
d2h_stream: torch.Stream,
|
||||||
) -> None:
|
) -> None:
|
||||||
self._config = config
|
self._config = config
|
||||||
self._device_state = device_state
|
self._device_state = device_state
|
||||||
|
|||||||
@@ -24,7 +24,7 @@ class PeriodicCanaryStatsLogger:
|
|||||||
active_tags: tuple[CanaryLaunchTag, ...],
|
active_tags: tuple[CanaryLaunchTag, ...],
|
||||||
outer_step_counter_getter: Callable[[], int],
|
outer_step_counter_getter: Callable[[], int],
|
||||||
sweep_orchestrator: SweepOrchestrator,
|
sweep_orchestrator: SweepOrchestrator,
|
||||||
d2h_stream: torch.cuda.Stream,
|
d2h_stream: torch.Stream,
|
||||||
) -> None:
|
) -> None:
|
||||||
self._config = config
|
self._config = config
|
||||||
self._device_state = device_state
|
self._device_state = device_state
|
||||||
|
|||||||
@@ -30,7 +30,7 @@ class SwaDivergenceReporter:
|
|||||||
self,
|
self,
|
||||||
*,
|
*,
|
||||||
device: torch.device,
|
device: torch.device,
|
||||||
d2h_stream: torch.cuda.Stream,
|
d2h_stream: torch.Stream,
|
||||||
interval: int,
|
interval: int,
|
||||||
swa_allocator: Optional[SWATokenToKVPoolAllocator] = None,
|
swa_allocator: Optional[SWATokenToKVPoolAllocator] = None,
|
||||||
req_to_token_pool: Optional[ReqToTokenPool] = None,
|
req_to_token_pool: Optional[ReqToTokenPool] = None,
|
||||||
|
|||||||
@@ -16,7 +16,7 @@ class ViolationManager:
|
|||||||
*,
|
*,
|
||||||
config: CanaryConfig,
|
config: CanaryConfig,
|
||||||
device_state: CanaryDeviceState,
|
device_state: CanaryDeviceState,
|
||||||
d2h_stream: torch.cuda.Stream,
|
d2h_stream: torch.Stream,
|
||||||
outer_step_counter_getter: Callable[[], int],
|
outer_step_counter_getter: Callable[[], int],
|
||||||
) -> None:
|
) -> None:
|
||||||
self._device_state = device_state
|
self._device_state = device_state
|
||||||
|
|||||||
@@ -71,7 +71,7 @@ class SingleForwardManager:
|
|||||||
per_forward_verify_capacity: int,
|
per_forward_verify_capacity: int,
|
||||||
per_forward_write_req_capacity: int,
|
per_forward_write_req_capacity: int,
|
||||||
per_forward_write_entry_capacity: int,
|
per_forward_write_entry_capacity: int,
|
||||||
d2h_stream: torch.cuda.Stream,
|
d2h_stream: torch.Stream,
|
||||||
token_oracle_manager: Optional[TokenOracleManager],
|
token_oracle_manager: Optional[TokenOracleManager],
|
||||||
swa_divergence_report: Optional[SwaDivergenceReporter],
|
swa_divergence_report: Optional[SwaDivergenceReporter],
|
||||||
is_eagle_draft_decode: bool,
|
is_eagle_draft_decode: bool,
|
||||||
|
|||||||
@@ -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 Qwen2MLP as Qwen3MLP
|
||||||
from sglang.srt.models.qwen2 import Qwen2Model
|
from sglang.srt.models.qwen2 import Qwen2Model
|
||||||
from sglang.srt.models.utils import apply_qk_norm
|
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.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
|
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
|
del self.lm_head.weight
|
||||||
self.model.embed_tokens.weight = embed
|
self.model.embed_tokens.weight = embed
|
||||||
self.lm_head.weight = head
|
self.lm_head.weight = head
|
||||||
torch.cuda.empty_cache()
|
current_platform.empty_cache()
|
||||||
torch.cuda.synchronize()
|
current_platform.synchronize()
|
||||||
|
|
||||||
def load_kv_cache_scales(self, quantization_param_path: str) -> None:
|
def load_kv_cache_scales(self, quantization_param_path: str) -> None:
|
||||||
self.model.load_kv_cache_scales(quantization_param_path)
|
self.model.load_kv_cache_scales(quantization_param_path)
|
||||||
|
|||||||
@@ -608,6 +608,17 @@ def device_stream_context(stream):
|
|||||||
return torch.get_device_module(stream.device).stream(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():
|
def get_amdgpu_memory_capacity():
|
||||||
try:
|
try:
|
||||||
# Run rocm-smi and capture the output
|
# Run rocm-smi and capture the output
|
||||||
|
|||||||
@@ -7,6 +7,7 @@ import triton
|
|||||||
import triton.language as tl
|
import triton.language as tl
|
||||||
|
|
||||||
from sglang.srt.environ import envs
|
from sglang.srt.environ import envs
|
||||||
|
from sglang.srt.utils import is_device_stream_capturing
|
||||||
|
|
||||||
|
|
||||||
def _phase_repr(phase: int | IntEnum) -> str:
|
def _phase_repr(phase: int | IntEnum) -> str:
|
||||||
@@ -51,6 +52,7 @@ class SimplePhaseChecker:
|
|||||||
|
|
||||||
def __init__(self, *, initial_phase: int | IntEnum, device: torch.device) -> None:
|
def __init__(self, *, initial_phase: int | IntEnum, device: torch.device) -> None:
|
||||||
self._initial_phase = int(initial_phase)
|
self._initial_phase = int(initial_phase)
|
||||||
|
self._device = device
|
||||||
self._phase = torch.tensor(
|
self._phase = torch.tensor(
|
||||||
self._initial_phase, dtype=torch.int32, device=device
|
self._initial_phase, dtype=torch.int32, device=device
|
||||||
)
|
)
|
||||||
@@ -81,7 +83,7 @@ class SimplePhaseChecker:
|
|||||||
f"caller_tag={caller_tag} "
|
f"caller_tag={caller_tag} "
|
||||||
f"expect={_phase_repr(expect_phase)} "
|
f"expect={_phase_repr(expect_phase)} "
|
||||||
f"next={_phase_repr(next_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,)](
|
_phase_check_kernel[(1,)](
|
||||||
self._phase,
|
self._phase,
|
||||||
|
|||||||
@@ -73,6 +73,12 @@ class CanaryE2EBase(CapturedServerE2EBase):
|
|||||||
# test methods send N sequential batches so the SWA allocator's full→swa index mapping
|
# test methods send N sequential batches so the SWA allocator's full→swa index mapping
|
||||||
# diverges from identity. Default 1 keeps MHA tests fast.
|
# diverges from identity. Default 1 keeps MHA tests fast.
|
||||||
workload_n_batches: ClassVar[int] = 1
|
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
|
_cfg: ClassVar[Optional[_ModeConfig]] = None
|
||||||
|
|
||||||
@@ -121,14 +127,24 @@ class CanaryE2EBase(CapturedServerE2EBase):
|
|||||||
|
|
||||||
def send_parallel_requests(
|
def send_parallel_requests(
|
||||||
self,
|
self,
|
||||||
n: int = 8,
|
n: Optional[int] = None,
|
||||||
*,
|
*,
|
||||||
assert_all_success: bool = True,
|
assert_all_success: bool = True,
|
||||||
max_new_tokens: int = 2048,
|
max_new_tokens: Optional[int] = None,
|
||||||
timeout: float = 240.0,
|
timeout: Optional[float] = None,
|
||||||
ignore_eos: Optional[bool] = None,
|
ignore_eos: Optional[bool] = None,
|
||||||
) -> list[dict]:
|
) -> 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:
|
if ignore_eos is None:
|
||||||
ignore_eos = self.model_mode == "swa"
|
ignore_eos = self.model_mode == "swa"
|
||||||
results = post_parallel_generate(
|
results = post_parallel_generate(
|
||||||
|
|||||||
@@ -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.kv_canary.pool_patcher.api import register_pool_attacher
|
||||||
from sglang.srt.mem_cache.radix_cache import RadixCache, TreeNode
|
from sglang.srt.mem_cache.radix_cache import RadixCache, TreeNode
|
||||||
from sglang.srt.model_executor.forward_batch_info import ForwardMode
|
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
|
@dataclass
|
||||||
|
|||||||
@@ -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()
|
||||||
@@ -20,6 +20,7 @@ from sglang.kernels.ops.kv_canary.verify import (
|
|||||||
from sglang.kernels.ops.kv_canary.verify_ref import (
|
from sglang.kernels.ops.kv_canary.verify_ref import (
|
||||||
_compute_real_kv_hash_scalar,
|
_compute_real_kv_hash_scalar,
|
||||||
launch_canary_verify_kernel_torch_reference,
|
launch_canary_verify_kernel_torch_reference,
|
||||||
|
materialize_real_kv_sources,
|
||||||
)
|
)
|
||||||
from sglang.kernels.ops.kv_canary.write_ref import (
|
from sglang.kernels.ops.kv_canary.write_ref import (
|
||||||
launch_canary_write_kernel_torch_reference,
|
launch_canary_write_kernel_torch_reference,
|
||||||
@@ -917,14 +918,18 @@ class TestRealKvHash:
|
|||||||
positions = [0, 1, 2]
|
positions = [0, 1, 2]
|
||||||
|
|
||||||
running = splitmix64(consts.CANARY_CHAIN_ANCHOR)
|
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] = []
|
real_kv_hashes: list[int] = []
|
||||||
for slot_idx in slot_indices:
|
for slot_idx in slot_indices:
|
||||||
real_kv_hashes.append(
|
real_kv_hashes.append(
|
||||||
_compute_real_kv_hash_scalar(
|
_compute_real_kv_hash_scalar(
|
||||||
real_kv_sources=sources_cuda,
|
|
||||||
real_kv_hash_mode=consts.RealKvHashMode.ALL,
|
|
||||||
slot_idx=slot_idx,
|
slot_idx=slot_idx,
|
||||||
work_device=torch.device("cpu"),
|
host_sources=host_sources,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -1030,6 +1035,21 @@ class TestRealKvSource:
|
|||||||
read_bytes=0,
|
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:
|
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."""
|
"""Host wrapper pads to 4 slots when fewer sources are supplied; dummy slots are never dereferenced."""
|
||||||
buf_pair = _buf_pair()
|
buf_pair = _buf_pair()
|
||||||
@@ -1310,12 +1330,16 @@ class TestLayoutAndScheduling:
|
|||||||
# byte-by-byte loop, so the stamped real_kv_hash matches what the kernel /
|
# 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
|
# verify reference will recompute. A byte-by-byte fold was the previous bug
|
||||||
# here and triggered REAL_KV_HASH violations on otherwise clean chains.
|
# 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 = [
|
rkv_values = [
|
||||||
_compute_real_kv_hash_scalar(
|
_compute_real_kv_hash_scalar(
|
||||||
slot_idx=slot_idx,
|
slot_idx=slot_idx,
|
||||||
real_kv_sources=sources_cuda,
|
host_sources=host_sources,
|
||||||
real_kv_hash_mode=consts.RealKvHashMode.ALL,
|
|
||||||
work_device=_DEVICE,
|
|
||||||
)
|
)
|
||||||
for slot_idx in slot_indices
|
for slot_idx in slot_indices
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -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()
|
||||||
@@ -1,16 +1,22 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import unittest
|
import unittest
|
||||||
from typing import cast
|
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
from sglang.srt.kv_canary.runner.future_tensor import FutureTensors
|
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
|
from sglang.test.test_utils import CustomTestCase
|
||||||
|
|
||||||
register_cuda_ci(est_time=10, stage="extra-a", runner_config="1-gpu-small")
|
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_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:
|
class _FakeEvent:
|
||||||
@@ -22,32 +28,30 @@ class _FakeEvent:
|
|||||||
|
|
||||||
|
|
||||||
class TestFutureTensors(CustomTestCase):
|
class TestFutureTensors(CustomTestCase):
|
||||||
def test_cuda_stage_then_wait_returns_host_copy(self) -> None:
|
def test_device_stage_then_wait_returns_host_copy(self) -> None:
|
||||||
"""Verify staged CUDA tensors are copied back on wait."""
|
"""Verify staged device tensors are copied back on wait."""
|
||||||
device = torch.device("cuda")
|
alt_stream = create_device_stream(DEFAULT_DEVICE)
|
||||||
alt_stream = torch.cuda.Stream(device=device)
|
default_stream = get_current_device_stream_fast()
|
||||||
default_stream = torch.cuda.current_stream(device)
|
|
||||||
self.assertNotEqual(alt_stream.stream_id, default_stream.stream_id)
|
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(
|
future_first = FutureTensors.device_to_host(
|
||||||
xs_device=src_first, d2h_stream=alt_stream
|
xs_device=src_first, d2h_stream=alt_stream
|
||||||
)
|
)
|
||||||
result_first = future_first.wait()
|
result_first = future_first.wait()
|
||||||
self.assertEqual(int(result_first.item()), 41)
|
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(
|
future_second = FutureTensors.device_to_host(
|
||||||
xs_device=src_second, d2h_stream=alt_stream
|
xs_device=src_second, d2h_stream=alt_stream
|
||||||
)
|
)
|
||||||
result_second = future_second.wait()
|
result_second = future_second.wait()
|
||||||
self.assertEqual(int(result_second.item()), 97)
|
self.assertEqual(int(result_second.item()), 97)
|
||||||
|
|
||||||
def test_cuda_pinned_when_stream_is_provided(self) -> None:
|
def test_device_pinned_when_stream_is_provided(self) -> None:
|
||||||
"""Verify CUDA staging uses pinned host memory with a stream."""
|
"""Verify device staging uses pinned host memory with a stream."""
|
||||||
device = torch.device("cuda")
|
alt_stream = create_device_stream(DEFAULT_DEVICE)
|
||||||
alt_stream = torch.cuda.Stream(device=device)
|
src = torch.tensor([5], dtype=torch.int32, device=DEFAULT_DEVICE)
|
||||||
src = torch.tensor([5], dtype=torch.int32, device=device)
|
|
||||||
future = FutureTensors.device_to_host(xs_device=src, d2h_stream=alt_stream)
|
future = FutureTensors.device_to_host(xs_device=src, d2h_stream=alt_stream)
|
||||||
staged_tensors = [
|
staged_tensors = [
|
||||||
v for v in future._data.values() if isinstance(v, torch.Tensor)
|
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.assertTrue(all(t.is_pinned() for t in staged_tensors))
|
||||||
self.assertEqual(int(future.wait().item()), 5)
|
self.assertEqual(int(future.wait().item()), 5)
|
||||||
|
|
||||||
def test_cuda_each_call_allocates_fresh_host(self) -> None:
|
def test_device_each_call_allocates_fresh_host(self) -> None:
|
||||||
"""Verify each CUDA staging call owns a fresh host buffer."""
|
"""Verify each device staging call owns a fresh host buffer."""
|
||||||
device = torch.device("cuda")
|
alt_stream = create_device_stream(DEFAULT_DEVICE)
|
||||||
alt_stream = torch.cuda.Stream(device=device)
|
src_a = torch.tensor([13], dtype=torch.int32, device=DEFAULT_DEVICE)
|
||||||
src_a = torch.tensor([13], dtype=torch.int32, device=device)
|
src_b = torch.tensor([29], dtype=torch.int32, device=DEFAULT_DEVICE)
|
||||||
src_b = torch.tensor([29], dtype=torch.int32, device=device)
|
|
||||||
future_a = FutureTensors.device_to_host(xs_device=src_a, d2h_stream=alt_stream)
|
future_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)
|
future_b = FutureTensors.device_to_host(xs_device=src_b, d2h_stream=alt_stream)
|
||||||
ptrs_a = {
|
ptrs_a = {
|
||||||
@@ -77,11 +80,10 @@ class TestFutureTensors(CustomTestCase):
|
|||||||
|
|
||||||
def test_dict_of_all_tensors_roundtrip(self) -> None:
|
def test_dict_of_all_tensors_roundtrip(self) -> None:
|
||||||
"""Verify a dict of multiple tensors round-trips entry-by-entry."""
|
"""Verify a dict of multiple tensors round-trips entry-by-entry."""
|
||||||
device = torch.device("cuda")
|
stream = create_device_stream(DEFAULT_DEVICE)
|
||||||
stream = torch.cuda.Stream(device=device)
|
|
||||||
src = {
|
src = {
|
||||||
"x": torch.tensor([11, 22], dtype=torch.int64, device=device),
|
"x": torch.tensor([11, 22], dtype=torch.int64, device=DEFAULT_DEVICE),
|
||||||
"y": torch.tensor([99], dtype=torch.int32, device=device),
|
"y": torch.tensor([99], dtype=torch.int32, device=DEFAULT_DEVICE),
|
||||||
}
|
}
|
||||||
future = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream)
|
future = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream)
|
||||||
out = future.wait()
|
out = future.wait()
|
||||||
@@ -93,14 +95,13 @@ class TestFutureTensors(CustomTestCase):
|
|||||||
|
|
||||||
def test_dict_mixes_tensor_and_passthrough(self) -> None:
|
def test_dict_mixes_tensor_and_passthrough(self) -> None:
|
||||||
"""Verify non-tensor dict entries ride through verbatim alongside staging."""
|
"""Verify non-tensor dict entries ride through verbatim alongside staging."""
|
||||||
device = torch.device("cuda")
|
stream = create_device_stream(DEFAULT_DEVICE)
|
||||||
stream = torch.cuda.Stream(device=device)
|
|
||||||
sentinel_obj = {"nested": [1, 2, 3]}
|
sentinel_obj = {"nested": [1, 2, 3]}
|
||||||
src = {
|
src = {
|
||||||
"step": 42,
|
"step": 42,
|
||||||
"label": "decode",
|
"label": "decode",
|
||||||
"extra": sentinel_obj,
|
"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)
|
future = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream)
|
||||||
out = future.wait()
|
out = future.wait()
|
||||||
@@ -113,9 +114,8 @@ class TestFutureTensors(CustomTestCase):
|
|||||||
|
|
||||||
def test_dict_passthrough_preserves_tensor_value(self) -> None:
|
def test_dict_passthrough_preserves_tensor_value(self) -> None:
|
||||||
"""Verify tensors share device memory but non-tensor types are not staged."""
|
"""Verify tensors share device memory but non-tensor types are not staged."""
|
||||||
device = torch.device("cuda")
|
stream = create_device_stream(DEFAULT_DEVICE)
|
||||||
stream = torch.cuda.Stream(device=device)
|
src_tensor = torch.tensor([3], dtype=torch.int32, device=DEFAULT_DEVICE)
|
||||||
src_tensor = torch.tensor([3], dtype=torch.int32, device=device)
|
|
||||||
src = {"step": 100, "buf": src_tensor}
|
src = {"step": 100, "buf": src_tensor}
|
||||||
future = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream)
|
future = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream)
|
||||||
out = future.wait()
|
out = future.wait()
|
||||||
@@ -128,8 +128,7 @@ class TestFutureTensors(CustomTestCase):
|
|||||||
|
|
||||||
def test_dict_without_tensor_raises(self) -> None:
|
def test_dict_without_tensor_raises(self) -> None:
|
||||||
"""Verify a tensor-less dict raises (no device to anchor the d2h sync)."""
|
"""Verify a tensor-less dict raises (no device to anchor the d2h sync)."""
|
||||||
device = torch.device("cuda")
|
stream = create_device_stream(DEFAULT_DEVICE)
|
||||||
stream = torch.cuda.Stream(device=device)
|
|
||||||
with self.assertRaises(ValueError):
|
with self.assertRaises(ValueError):
|
||||||
FutureTensors.device_to_host(
|
FutureTensors.device_to_host(
|
||||||
xs_device={"step": 0, "label": "decode"}, d2h_stream=stream
|
xs_device={"step": 0, "label": "decode"}, d2h_stream=stream
|
||||||
@@ -137,9 +136,8 @@ class TestFutureTensors(CustomTestCase):
|
|||||||
|
|
||||||
def test_wait_called_twice_raises(self) -> None:
|
def test_wait_called_twice_raises(self) -> None:
|
||||||
"""Verify wait() after the first drain raises (state cleared)."""
|
"""Verify wait() after the first drain raises (state cleared)."""
|
||||||
device = torch.device("cuda")
|
stream = create_device_stream(DEFAULT_DEVICE)
|
||||||
stream = torch.cuda.Stream(device=device)
|
src = torch.tensor([3], dtype=torch.int32, device=DEFAULT_DEVICE)
|
||||||
src = torch.tensor([3], dtype=torch.int32, device=device)
|
|
||||||
future = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream)
|
future = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream)
|
||||||
self.assertEqual(int(future.wait().item()), 3)
|
self.assertEqual(int(future.wait().item()), 3)
|
||||||
with self.assertRaises(RuntimeError):
|
with self.assertRaises(RuntimeError):
|
||||||
@@ -149,9 +147,7 @@ class TestFutureTensors(CustomTestCase):
|
|||||||
"""Verify wait() syncs the event exactly once and clears internal state."""
|
"""Verify wait() syncs the event exactly once and clears internal state."""
|
||||||
tensor = torch.tensor([1, 2, 3])
|
tensor = torch.tensor([1, 2, 3])
|
||||||
event = _FakeEvent()
|
event = _FakeEvent()
|
||||||
future = FutureTensors(
|
future = FutureTensors(_data={"x": tensor}, _event=event)
|
||||||
_data={"x": tensor}, _event=cast(torch.cuda.Event, event)
|
|
||||||
)
|
|
||||||
|
|
||||||
result = future.wait()
|
result = future.wait()
|
||||||
self.assertIs(result["x"], tensor)
|
self.assertIs(result["x"], tensor)
|
||||||
@@ -166,11 +162,10 @@ class TestFutureTensors(CustomTestCase):
|
|||||||
|
|
||||||
def test_dict_anchor_picked_from_first_tensor(self) -> None:
|
def test_dict_anchor_picked_from_first_tensor(self) -> None:
|
||||||
"""Verify staging works when the first key is a non-tensor (anchor must scan)."""
|
"""Verify staging works when the first key is a non-tensor (anchor must scan)."""
|
||||||
device = torch.device("cuda")
|
stream = create_device_stream(DEFAULT_DEVICE)
|
||||||
stream = torch.cuda.Stream(device=device)
|
|
||||||
src = {
|
src = {
|
||||||
"step": 5,
|
"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()
|
out = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream).wait()
|
||||||
self.assertEqual(out["step"], 5)
|
self.assertEqual(out["step"], 5)
|
||||||
|
|||||||
@@ -6,15 +6,21 @@ from types import SimpleNamespace
|
|||||||
import torch
|
import torch
|
||||||
|
|
||||||
from sglang.srt.kv_canary.plan_input import PlanInput
|
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 (
|
from sglang.test.kv_canary.fixtures import (
|
||||||
DEFAULT_DEVICE,
|
DEFAULT_DEVICE,
|
||||||
|
DEFAULT_DEVICE_MODULE,
|
||||||
make_forward_batch,
|
make_forward_batch,
|
||||||
)
|
)
|
||||||
from sglang.test.test_utils import CustomTestCase
|
from sglang.test.test_utils import CustomTestCase
|
||||||
|
|
||||||
register_cuda_ci(est_time=9, stage="extra-a", runner_config="1-gpu-small")
|
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_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:
|
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)
|
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 = _make_static_plan_input(bs_capacity=4, device=self.device)
|
||||||
plan.fill_from_forward_batch(forward_batch=fb)
|
plan.fill_from_forward_batch(forward_batch=fb)
|
||||||
torch.cuda.synchronize()
|
DEFAULT_DEVICE_MODULE.synchronize()
|
||||||
self.assertEqual(
|
self.assertEqual(
|
||||||
plan.req_to_verify_expected_tokens_valid_lens[:2].tolist(), [7, 9]
|
plan.req_to_verify_expected_tokens_valid_lens[:2].tolist(), [7, 9]
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -10,12 +10,21 @@ from sglang.srt.kv_canary.req_to_expected_token_ids_manager import (
|
|||||||
compute_req_all_ids_info,
|
compute_req_all_ids_info,
|
||||||
populate_req_to_expected_token_ids,
|
populate_req_to_expected_token_ids,
|
||||||
)
|
)
|
||||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
from sglang.test.ci.ci_register import (
|
||||||
from sglang.test.kv_canary.fixtures import DEFAULT_DEVICE, make_forward_batch
|
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
|
from sglang.test.test_utils import CustomTestCase
|
||||||
|
|
||||||
register_cuda_ci(est_time=11, stage="extra-a", runner_config="1-gpu-small")
|
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_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:
|
def _make_req(*, origin: list[int], output: list[int]) -> SimpleNamespace:
|
||||||
@@ -92,7 +101,7 @@ class TestPopulateReqToExpectedTokenIds(CustomTestCase):
|
|||||||
populate_req_to_expected_token_ids(
|
populate_req_to_expected_token_ids(
|
||||||
forward_batch=fb, req_to_verify_expected_tokens=pool
|
forward_batch=fb, req_to_verify_expected_tokens=pool
|
||||||
)
|
)
|
||||||
torch.cuda.synchronize()
|
DEFAULT_DEVICE_MODULE.synchronize()
|
||||||
self.assertTrue(torch.equal(pool, original))
|
self.assertTrue(torch.equal(pool, original))
|
||||||
|
|
||||||
def test_no_op_when_pool_is_none(self) -> None:
|
def test_no_op_when_pool_is_none(self) -> None:
|
||||||
@@ -117,7 +126,7 @@ class TestPopulateReqToExpectedTokenIds(CustomTestCase):
|
|||||||
populate_req_to_expected_token_ids(
|
populate_req_to_expected_token_ids(
|
||||||
forward_batch=fb, req_to_verify_expected_tokens=pool
|
forward_batch=fb, req_to_verify_expected_tokens=pool
|
||||||
)
|
)
|
||||||
torch.cuda.synchronize()
|
DEFAULT_DEVICE_MODULE.synchronize()
|
||||||
self.assertTrue(torch.equal(pool, original))
|
self.assertTrue(torch.equal(pool, original))
|
||||||
|
|
||||||
def test_raises_when_lens_length_mismatches_batch_size(self) -> None:
|
def test_raises_when_lens_length_mismatches_batch_size(self) -> None:
|
||||||
@@ -154,7 +163,7 @@ class TestPopulateReqToExpectedTokenIds(CustomTestCase):
|
|||||||
populate_req_to_expected_token_ids(
|
populate_req_to_expected_token_ids(
|
||||||
forward_batch=fb, req_to_verify_expected_tokens=pool
|
forward_batch=fb, req_to_verify_expected_tokens=pool
|
||||||
)
|
)
|
||||||
torch.cuda.synchronize()
|
DEFAULT_DEVICE_MODULE.synchronize()
|
||||||
|
|
||||||
pool_cpu = pool.cpu()
|
pool_cpu = pool.cpu()
|
||||||
self.assertEqual(pool_cpu[1, :3].tolist(), [10, 20, 30])
|
self.assertEqual(pool_cpu[1, :3].tolist(), [10, 20, 30])
|
||||||
|
|||||||
@@ -15,24 +15,28 @@ from sglang.srt.kv_canary.runner.swa_divergence import (
|
|||||||
SwaDivergenceReporter,
|
SwaDivergenceReporter,
|
||||||
compute_swa_full_idx_divergence,
|
compute_swa_full_idx_divergence,
|
||||||
)
|
)
|
||||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
from sglang.srt.utils import create_device_stream
|
||||||
from sglang.test.kv_canary.fixtures import make_buffer_group
|
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.kv_canary.runner_test_base import CanaryManagerTestCase, make_manager
|
||||||
from sglang.test.test_utils import CustomTestCase
|
from sglang.test.test_utils import CustomTestCase
|
||||||
|
|
||||||
register_cuda_ci(est_time=11, stage="extra-a", runner_config="1-gpu-small")
|
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")
|
register_amd_ci(est_time=45, suite="extra-a-test-1-gpu-small-amd")
|
||||||
|
register_xpu_ci(est_time=60, suite="stage-b-test-1-gpu-xpu")
|
||||||
_DEVICE = torch.device("cuda")
|
|
||||||
|
|
||||||
_EMPTY_FORWARD_BATCH = SimpleNamespace(
|
_EMPTY_FORWARD_BATCH = SimpleNamespace(
|
||||||
req_pool_indices=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=_DEVICE),
|
seq_lens=torch.empty(0, dtype=torch.int64, device=DEFAULT_DEVICE),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
def _make_verify_plan(value: int) -> VerifyPlan:
|
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))
|
plan.verify_num_valid.copy_(torch.tensor([value], dtype=torch.int32))
|
||||||
return plan
|
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:
|
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:
|
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)
|
return base.view(num_reqs, max_seq_len)
|
||||||
|
|
||||||
|
|
||||||
@@ -83,9 +89,9 @@ def _run_compute(
|
|||||||
|
|
||||||
class TestSwaDivergenceReporter(CustomTestCase):
|
class TestSwaDivergenceReporter(CustomTestCase):
|
||||||
def test_swa_divergence_log_emitted(self) -> None:
|
def test_swa_divergence_log_emitted(self) -> None:
|
||||||
d2h_stream = torch.cuda.Stream(device=_DEVICE)
|
d2h_stream = create_device_stream(DEFAULT_DEVICE)
|
||||||
stats = SwaDivergenceReporter(
|
stats = SwaDivergenceReporter(
|
||||||
device=_DEVICE,
|
device=DEFAULT_DEVICE,
|
||||||
d2h_stream=d2h_stream,
|
d2h_stream=d2h_stream,
|
||||||
interval=10,
|
interval=10,
|
||||||
swa_allocator=None,
|
swa_allocator=None,
|
||||||
@@ -96,13 +102,13 @@ class TestSwaDivergenceReporter(CustomTestCase):
|
|||||||
for forward_idx in range(3):
|
for forward_idx in range(3):
|
||||||
stats.observe_after_invoke_plan(
|
stats.observe_after_invoke_plan(
|
||||||
group=make_buffer_group(
|
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),
|
verify_plan=_make_verify_plan(10),
|
||||||
)
|
)
|
||||||
stats.observe_after_invoke_plan(
|
stats.observe_after_invoke_plan(
|
||||||
group=make_buffer_group(
|
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),
|
verify_plan=_make_verify_plan(3),
|
||||||
)
|
)
|
||||||
@@ -115,13 +121,13 @@ class TestSwaDivergenceReporter(CustomTestCase):
|
|||||||
# the staged future hangs onto it. forward_ct is now 4.
|
# the staged future hangs onto it. forward_ct is now 4.
|
||||||
stats.observe_after_invoke_plan(
|
stats.observe_after_invoke_plan(
|
||||||
group=make_buffer_group(
|
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),
|
verify_plan=_make_verify_plan(10),
|
||||||
)
|
)
|
||||||
stats.observe_after_invoke_plan(
|
stats.observe_after_invoke_plan(
|
||||||
group=make_buffer_group(
|
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),
|
verify_plan=_make_verify_plan(3),
|
||||||
)
|
)
|
||||||
@@ -150,9 +156,9 @@ class TestSwaDivergenceReporter(CustomTestCase):
|
|||||||
self.assertEqual(fields.swa_full_idx_divergence, 0)
|
self.assertEqual(fields.swa_full_idx_divergence, 0)
|
||||||
|
|
||||||
def test_swa_divergence_counts_monotonic_increasing(self) -> None:
|
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(
|
stats = SwaDivergenceReporter(
|
||||||
device=_DEVICE,
|
device=DEFAULT_DEVICE,
|
||||||
d2h_stream=d2h_stream,
|
d2h_stream=d2h_stream,
|
||||||
interval=10,
|
interval=10,
|
||||||
swa_allocator=None,
|
swa_allocator=None,
|
||||||
@@ -188,13 +194,19 @@ class TestSwaDivergenceReporter(CustomTestCase):
|
|||||||
for _ in range(5):
|
for _ in range(5):
|
||||||
stats.observe_after_invoke_plan(
|
stats.observe_after_invoke_plan(
|
||||||
group=make_buffer_group(
|
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),
|
verify_plan=_make_verify_plan(7),
|
||||||
)
|
)
|
||||||
stats.observe_after_invoke_plan(
|
stats.observe_after_invoke_plan(
|
||||||
group=make_buffer_group(
|
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),
|
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)
|
req_to_token = _make_identity_req_to_token(num_reqs=4, max_seq_len=16)
|
||||||
|
|
||||||
forward_batch = _make_forward_batch(
|
forward_batch = _make_forward_batch(
|
||||||
req_pool_indices=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=_DEVICE),
|
seq_lens=torch.empty(0, dtype=torch.int64, device=DEFAULT_DEVICE),
|
||||||
)
|
)
|
||||||
|
|
||||||
self.assertEqual(
|
self.assertEqual(
|
||||||
@@ -234,8 +246,10 @@ class TestSwaFullIdxDivergenceCompute(CustomTestCase):
|
|||||||
req_to_token = _make_identity_req_to_token(num_reqs=4, max_seq_len=16)
|
req_to_token = _make_identity_req_to_token(num_reqs=4, max_seq_len=16)
|
||||||
|
|
||||||
forward_batch = _make_forward_batch(
|
forward_batch = _make_forward_batch(
|
||||||
req_pool_indices=torch.tensor([0, 2], dtype=torch.int64, device=_DEVICE),
|
req_pool_indices=torch.tensor(
|
||||||
seq_lens=torch.tensor([8, 5], dtype=torch.int64, device=_DEVICE),
|
[0, 2], dtype=torch.int64, device=DEFAULT_DEVICE
|
||||||
|
),
|
||||||
|
seq_lens=torch.tensor([8, 5], dtype=torch.int64, device=DEFAULT_DEVICE),
|
||||||
)
|
)
|
||||||
|
|
||||||
self.assertEqual(
|
self.assertEqual(
|
||||||
@@ -256,8 +270,10 @@ class TestSwaFullIdxDivergenceCompute(CustomTestCase):
|
|||||||
mapping[17] = 60
|
mapping[17] = 60
|
||||||
|
|
||||||
forward_batch = _make_forward_batch(
|
forward_batch = _make_forward_batch(
|
||||||
req_pool_indices=torch.tensor([0, 1], dtype=torch.int64, device=_DEVICE),
|
req_pool_indices=torch.tensor(
|
||||||
seq_lens=torch.tensor([8, 8], dtype=torch.int64, device=_DEVICE),
|
[0, 1], dtype=torch.int64, device=DEFAULT_DEVICE
|
||||||
|
),
|
||||||
|
seq_lens=torch.tensor([8, 8], dtype=torch.int64, device=DEFAULT_DEVICE),
|
||||||
)
|
)
|
||||||
|
|
||||||
self.assertEqual(
|
self.assertEqual(
|
||||||
@@ -281,8 +297,10 @@ class TestSwaFullIdxDivergenceCompute(CustomTestCase):
|
|||||||
mapping[7] = 42
|
mapping[7] = 42
|
||||||
|
|
||||||
forward_batch = _make_forward_batch(
|
forward_batch = _make_forward_batch(
|
||||||
req_pool_indices=torch.tensor([0], dtype=torch.int64, device=_DEVICE),
|
req_pool_indices=torch.tensor(
|
||||||
seq_lens=torch.tensor([8], dtype=torch.int64, device=_DEVICE),
|
[0], dtype=torch.int64, device=DEFAULT_DEVICE
|
||||||
|
),
|
||||||
|
seq_lens=torch.tensor([8], dtype=torch.int64, device=DEFAULT_DEVICE),
|
||||||
)
|
)
|
||||||
|
|
||||||
self.assertEqual(
|
self.assertEqual(
|
||||||
@@ -302,8 +320,10 @@ class TestSwaFullIdxDivergenceCompute(CustomTestCase):
|
|||||||
mapping[28] = 77
|
mapping[28] = 77
|
||||||
|
|
||||||
forward_batch = _make_forward_batch(
|
forward_batch = _make_forward_batch(
|
||||||
req_pool_indices=torch.tensor([0], dtype=torch.int64, device=_DEVICE),
|
req_pool_indices=torch.tensor(
|
||||||
seq_lens=torch.tensor([10], dtype=torch.int64, device=_DEVICE),
|
[0], dtype=torch.int64, device=DEFAULT_DEVICE
|
||||||
|
),
|
||||||
|
seq_lens=torch.tensor([10], dtype=torch.int64, device=DEFAULT_DEVICE),
|
||||||
)
|
)
|
||||||
|
|
||||||
self.assertEqual(
|
self.assertEqual(
|
||||||
@@ -325,12 +345,16 @@ class TestSwaFullIdxDivergenceCompute(CustomTestCase):
|
|||||||
mapping[33] = 100
|
mapping[33] = 100
|
||||||
|
|
||||||
fb_req0 = _make_forward_batch(
|
fb_req0 = _make_forward_batch(
|
||||||
req_pool_indices=torch.tensor([0], dtype=torch.int64, device=_DEVICE),
|
req_pool_indices=torch.tensor(
|
||||||
seq_lens=torch.tensor([4], dtype=torch.int64, device=_DEVICE),
|
[0], dtype=torch.int64, device=DEFAULT_DEVICE
|
||||||
|
),
|
||||||
|
seq_lens=torch.tensor([4], dtype=torch.int64, device=DEFAULT_DEVICE),
|
||||||
)
|
)
|
||||||
fb_req2 = _make_forward_batch(
|
fb_req2 = _make_forward_batch(
|
||||||
req_pool_indices=torch.tensor([2], dtype=torch.int64, device=_DEVICE),
|
req_pool_indices=torch.tensor(
|
||||||
seq_lens=torch.tensor([4], dtype=torch.int64, device=_DEVICE),
|
[2], dtype=torch.int64, device=DEFAULT_DEVICE
|
||||||
|
),
|
||||||
|
seq_lens=torch.tensor([4], dtype=torch.int64, device=DEFAULT_DEVICE),
|
||||||
)
|
)
|
||||||
|
|
||||||
self.assertEqual(
|
self.assertEqual(
|
||||||
@@ -363,15 +387,17 @@ class TestSwaDivergenceReporterWithCompute(CustomTestCase):
|
|||||||
mapping[2] = 52
|
mapping[2] = 52
|
||||||
|
|
||||||
forward_batch = _make_forward_batch(
|
forward_batch = _make_forward_batch(
|
||||||
req_pool_indices=torch.tensor([0], dtype=torch.int64, device=_DEVICE),
|
req_pool_indices=torch.tensor(
|
||||||
seq_lens=torch.tensor([8], dtype=torch.int64, device=_DEVICE),
|
[0], dtype=torch.int64, device=DEFAULT_DEVICE
|
||||||
|
),
|
||||||
|
seq_lens=torch.tensor([8], dtype=torch.int64, device=DEFAULT_DEVICE),
|
||||||
)
|
)
|
||||||
|
|
||||||
swa_allocator = _make_allocator_stub(mapping)
|
swa_allocator = _make_allocator_stub(mapping)
|
||||||
req_to_token_pool = _make_req_to_token_pool_stub(req_to_token)
|
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(
|
stats = SwaDivergenceReporter(
|
||||||
device=_DEVICE,
|
device=DEFAULT_DEVICE,
|
||||||
d2h_stream=d2h_stream,
|
d2h_stream=d2h_stream,
|
||||||
interval=10,
|
interval=10,
|
||||||
swa_allocator=swa_allocator,
|
swa_allocator=swa_allocator,
|
||||||
@@ -379,13 +405,13 @@ class TestSwaDivergenceReporterWithCompute(CustomTestCase):
|
|||||||
)
|
)
|
||||||
stats.observe_after_invoke_plan(
|
stats.observe_after_invoke_plan(
|
||||||
group=make_buffer_group(
|
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),
|
verify_plan=_make_verify_plan(11),
|
||||||
)
|
)
|
||||||
stats.observe_after_invoke_plan(
|
stats.observe_after_invoke_plan(
|
||||||
group=make_buffer_group(
|
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),
|
verify_plan=_make_verify_plan(3),
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -10,12 +10,23 @@ from enum import IntEnum
|
|||||||
|
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
|
from sglang.srt.utils import get_device
|
||||||
from sglang.srt.utils.phase_checker import SimplePhaseChecker
|
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
|
from sglang.test.test_utils import CustomTestCase
|
||||||
|
|
||||||
register_cuda_ci(est_time=17, stage="base-b", runner_config="1-gpu-small")
|
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")
|
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):
|
class _Phase(IntEnum):
|
||||||
@@ -40,7 +51,7 @@ class TestConstruction(CustomTestCase):
|
|||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def setUpClass(cls) -> None:
|
def setUpClass(cls) -> None:
|
||||||
cls.device = torch.device("cuda:0")
|
cls.device = _DEVICE
|
||||||
|
|
||||||
def test_init_stores_initial_phase_int(self) -> None:
|
def test_init_stores_initial_phase_int(self) -> None:
|
||||||
checker = SimplePhaseChecker(initial_phase=7, device=self.device)
|
checker = SimplePhaseChecker(initial_phase=7, device=self.device)
|
||||||
@@ -72,19 +83,19 @@ class TestUpdateAssertDisabled(CustomTestCase):
|
|||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def setUpClass(cls) -> None:
|
def setUpClass(cls) -> None:
|
||||||
cls.device = torch.device("cuda:0")
|
cls.device = _DEVICE
|
||||||
|
|
||||||
def test_update_advances_phase_on_match(self) -> None:
|
def test_update_advances_phase_on_match(self) -> None:
|
||||||
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||||
checker.update(expect_phase=_Phase.IDLE, next_phase=_Phase.A, caller_name="t")
|
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))
|
self.assertEqual(_phase_value(checker), int(_Phase.A))
|
||||||
|
|
||||||
def test_update_advances_phase_on_mismatch(self) -> None:
|
def test_update_advances_phase_on_mismatch(self) -> None:
|
||||||
"""assert OFF tolerates mismatches — store still happens unconditionally."""
|
"""assert OFF tolerates mismatches — store still happens unconditionally."""
|
||||||
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||||
checker.update(expect_phase=_Phase.C, next_phase=_Phase.B, caller_name="t")
|
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))
|
self.assertEqual(_phase_value(checker), int(_Phase.B))
|
||||||
|
|
||||||
def test_init_time_lifecycle_violations_tolerated(self) -> None:
|
def test_init_time_lifecycle_violations_tolerated(self) -> None:
|
||||||
@@ -98,7 +109,7 @@ class TestUpdateAssertDisabled(CustomTestCase):
|
|||||||
checker.update(
|
checker.update(
|
||||||
expect_phase=_Phase.B, next_phase=_Phase.IDLE, caller_name="warmup"
|
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))
|
self.assertEqual(_phase_value(checker), int(_Phase.IDLE))
|
||||||
|
|
||||||
|
|
||||||
@@ -107,13 +118,13 @@ class TestUpdateAssertEnabled(CustomTestCase):
|
|||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def setUpClass(cls) -> None:
|
def setUpClass(cls) -> None:
|
||||||
cls.device = torch.device("cuda:0")
|
cls.device = _DEVICE
|
||||||
|
|
||||||
def test_update_advances_phase_on_match(self) -> None:
|
def test_update_advances_phase_on_match(self) -> None:
|
||||||
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||||
checker.enable_assert()
|
checker.enable_assert()
|
||||||
checker.update(expect_phase=_Phase.IDLE, next_phase=_Phase.A, caller_name="t")
|
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))
|
self.assertEqual(_phase_value(checker), int(_Phase.A))
|
||||||
|
|
||||||
def test_full_4_state_lifecycle_round_trip(self) -> None:
|
def test_full_4_state_lifecycle_round_trip(self) -> None:
|
||||||
@@ -129,28 +140,30 @@ class TestUpdateAssertEnabled(CustomTestCase):
|
|||||||
checker.update(
|
checker.update(
|
||||||
expect_phase=_Phase.C, next_phase=_Phase.IDLE, caller_name="p4"
|
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))
|
self.assertEqual(_phase_value(checker), int(_Phase.IDLE))
|
||||||
|
|
||||||
def test_update_mismatch_after_enable_raises_in_subprocess(self) -> None:
|
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.
|
"""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("""
|
script = textwrap.dedent("""
|
||||||
import sys
|
import sys
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
|
from sglang.srt.utils import get_device
|
||||||
from sglang.srt.utils.phase_checker import SimplePhaseChecker
|
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 = SimplePhaseChecker(initial_phase=0, device=device)
|
||||||
checker.enable_assert()
|
checker.enable_assert()
|
||||||
# phase=0 but we claim expect=99 — kernel must fire device_assert.
|
# phase=0 but we claim expect=99 — kernel must fire device_assert.
|
||||||
checker.update(expect_phase=99, next_phase=1, caller_name="bad")
|
checker.update(expect_phase=99, next_phase=1, caller_name="bad")
|
||||||
try:
|
try:
|
||||||
torch.cuda.synchronize()
|
device_module.synchronize()
|
||||||
except RuntimeError as e:
|
except RuntimeError as e:
|
||||||
msg = str(e).lower()
|
msg = str(e).lower()
|
||||||
if "device-side assert" in msg or "phase mismatch" in msg:
|
if "device-side assert" in msg or "phase mismatch" in msg:
|
||||||
@@ -164,7 +177,10 @@ class TestUpdateAssertEnabled(CustomTestCase):
|
|||||||
[sys.executable, "-c", script],
|
[sys.executable, "-c", script],
|
||||||
capture_output=True,
|
capture_output=True,
|
||||||
text=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
|
# 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
|
# process then dies is not: the CUDA coredump handler may abort it, and sync
|
||||||
@@ -189,7 +205,7 @@ class TestEnableAssert(CustomTestCase):
|
|||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def setUpClass(cls) -> None:
|
def setUpClass(cls) -> None:
|
||||||
cls.device = torch.device("cuda:0")
|
cls.device = _DEVICE
|
||||||
|
|
||||||
def test_enable_assert_sets_flag_to_one(self) -> None:
|
def test_enable_assert_sets_flag_to_one(self) -> None:
|
||||||
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||||
@@ -208,7 +224,7 @@ class TestEnableAssert(CustomTestCase):
|
|||||||
checker.update(
|
checker.update(
|
||||||
expect_phase=_Phase.IDLE, next_phase=_Phase.C, caller_name="warmup"
|
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))
|
self.assertEqual(_phase_value(checker), int(_Phase.C))
|
||||||
|
|
||||||
checker.enable_assert()
|
checker.enable_assert()
|
||||||
@@ -218,7 +234,7 @@ class TestEnableAssert(CustomTestCase):
|
|||||||
"""Reset target tracks the original initial_phase, not 0."""
|
"""Reset target tracks the original initial_phase, not 0."""
|
||||||
checker = SimplePhaseChecker(initial_phase=42, device=self.device)
|
checker = SimplePhaseChecker(initial_phase=42, device=self.device)
|
||||||
checker.update(expect_phase=42, next_phase=7, caller_name="t")
|
checker.update(expect_phase=42, next_phase=7, caller_name="t")
|
||||||
torch.cuda.synchronize()
|
_DEVICE_MODULE.synchronize()
|
||||||
self.assertEqual(_phase_value(checker), 7)
|
self.assertEqual(_phase_value(checker), 7)
|
||||||
|
|
||||||
checker.enable_assert()
|
checker.enable_assert()
|
||||||
@@ -237,12 +253,12 @@ class TestResetToIdle(CustomTestCase):
|
|||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def setUpClass(cls) -> None:
|
def setUpClass(cls) -> None:
|
||||||
cls.device = torch.device("cuda:0")
|
cls.device = _DEVICE
|
||||||
|
|
||||||
def test_reset_after_update_restores_initial_phase(self) -> None:
|
def test_reset_after_update_restores_initial_phase(self) -> None:
|
||||||
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||||
checker.update(expect_phase=_Phase.IDLE, next_phase=_Phase.B, caller_name="t")
|
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))
|
self.assertEqual(_phase_value(checker), int(_Phase.B))
|
||||||
|
|
||||||
checker._reset_to_idle()
|
checker._reset_to_idle()
|
||||||
@@ -264,7 +280,7 @@ class TestResetToIdle(CustomTestCase):
|
|||||||
def test_reset_with_nonzero_initial_phase(self) -> None:
|
def test_reset_with_nonzero_initial_phase(self) -> None:
|
||||||
checker = SimplePhaseChecker(initial_phase=5, device=self.device)
|
checker = SimplePhaseChecker(initial_phase=5, device=self.device)
|
||||||
checker.update(expect_phase=5, next_phase=9, caller_name="t")
|
checker.update(expect_phase=5, next_phase=9, caller_name="t")
|
||||||
torch.cuda.synchronize()
|
_DEVICE_MODULE.synchronize()
|
||||||
checker._reset_to_idle()
|
checker._reset_to_idle()
|
||||||
self.assertEqual(_phase_value(checker), 5)
|
self.assertEqual(_phase_value(checker), 5)
|
||||||
|
|
||||||
@@ -274,7 +290,7 @@ class TestCallerTagRegistry(CustomTestCase):
|
|||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def setUpClass(cls) -> None:
|
def setUpClass(cls) -> None:
|
||||||
cls.device = torch.device("cuda:0")
|
cls.device = _DEVICE
|
||||||
|
|
||||||
def test_first_caller_gets_tag_one(self) -> None:
|
def test_first_caller_gets_tag_one(self) -> None:
|
||||||
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
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 = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||||
checker.enable_assert()
|
checker.enable_assert()
|
||||||
checker.update(expect_phase=_Phase.IDLE, next_phase=_Phase.A) # caller_name=""
|
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.assertIn("", checker._caller_tag_registry)
|
||||||
self.assertEqual(_phase_value(checker), int(_Phase.A))
|
self.assertEqual(_phase_value(checker), int(_Phase.A))
|
||||||
|
|
||||||
@@ -311,7 +327,7 @@ class TestCallerTagRegistry(CustomTestCase):
|
|||||||
checker.update(
|
checker.update(
|
||||||
expect_phase=_Phase.A, next_phase=_Phase.IDLE, caller_name="beta"
|
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})
|
self.assertEqual(checker._caller_tag_registry, {"alpha": 1, "beta": 2})
|
||||||
|
|
||||||
|
|
||||||
@@ -320,13 +336,13 @@ class TestMultipleInstances(CustomTestCase):
|
|||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def setUpClass(cls) -> None:
|
def setUpClass(cls) -> None:
|
||||||
cls.device = torch.device("cuda:0")
|
cls.device = _DEVICE
|
||||||
|
|
||||||
def test_phase_tensors_are_independent(self) -> None:
|
def test_phase_tensors_are_independent(self) -> None:
|
||||||
a = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
a = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||||
b = 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")
|
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(a), int(_Phase.B))
|
||||||
self.assertEqual(_phase_value(b), int(_Phase.IDLE))
|
self.assertEqual(_phase_value(b), int(_Phase.IDLE))
|
||||||
|
|
||||||
@@ -345,6 +361,10 @@ class TestMultipleInstances(CustomTestCase):
|
|||||||
self.assertEqual(b._resolve_caller_tag("shared_name"), 1)
|
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):
|
class TestCudaGraphCapture(CustomTestCase):
|
||||||
"""The kernel is launched unconditionally so it is capture-safe; the device flag
|
"""The kernel is launched unconditionally so it is capture-safe; the device flag
|
||||||
decides at replay time whether the assert fires.
|
decides at replay time whether the assert fires.
|
||||||
@@ -352,7 +372,7 @@ class TestCudaGraphCapture(CustomTestCase):
|
|||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def setUpClass(cls) -> None:
|
def setUpClass(cls) -> None:
|
||||||
cls.device = torch.device("cuda:0")
|
cls.device = _DEVICE
|
||||||
|
|
||||||
def _capture_one_update(
|
def _capture_one_update(
|
||||||
self,
|
self,
|
||||||
@@ -379,7 +399,7 @@ class TestCudaGraphCapture(CustomTestCase):
|
|||||||
caller_name=caller_name,
|
caller_name=caller_name,
|
||||||
)
|
)
|
||||||
torch.cuda.current_stream(self.device).wait_stream(stream)
|
torch.cuda.current_stream(self.device).wait_stream(stream)
|
||||||
torch.cuda.synchronize()
|
_DEVICE_MODULE.synchronize()
|
||||||
|
|
||||||
graph = torch.cuda.CUDAGraph()
|
graph = torch.cuda.CUDAGraph()
|
||||||
with torch.cuda.graph(graph, stream=stream):
|
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.
|
# Enable assert (resets phase -> IDLE) and replay — captured expect=IDLE matches.
|
||||||
checker.enable_assert()
|
checker.enable_assert()
|
||||||
graph.replay()
|
graph.replay()
|
||||||
torch.cuda.synchronize()
|
_DEVICE_MODULE.synchronize()
|
||||||
self.assertEqual(_phase_value(checker), int(_Phase.B))
|
self.assertEqual(_phase_value(checker), int(_Phase.B))
|
||||||
|
|
||||||
# Reset + replay again — same result, no raise.
|
# Reset + replay again — same result, no raise.
|
||||||
checker._reset_to_idle()
|
checker._reset_to_idle()
|
||||||
graph.replay()
|
graph.replay()
|
||||||
torch.cuda.synchronize()
|
_DEVICE_MODULE.synchronize()
|
||||||
self.assertEqual(_phase_value(checker), int(_Phase.B))
|
self.assertEqual(_phase_value(checker), int(_Phase.B))
|
||||||
|
|
||||||
def test_assert_flag_toggle_visible_to_replayed_graph(self) -> None:
|
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.
|
# Replay with assert OFF tolerates a deliberately diverged phase.
|
||||||
checker._phase.fill_(999)
|
checker._phase.fill_(999)
|
||||||
graph.replay()
|
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))
|
self.assertEqual(_phase_value(checker), int(_Phase.A))
|
||||||
|
|
||||||
# Now turn on asserts (also resets phase -> IDLE) and replay.
|
# 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))
|
self.assertEqual(_phase_value(checker), int(_Phase.IDLE))
|
||||||
|
|
||||||
graph.replay()
|
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))
|
self.assertEqual(_phase_value(checker), int(_Phase.A))
|
||||||
|
|
||||||
|
|
||||||
@@ -451,13 +471,13 @@ class TestPhaseReprNoCrash(CustomTestCase):
|
|||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def setUpClass(cls) -> None:
|
def setUpClass(cls) -> None:
|
||||||
cls.device = torch.device("cuda:0")
|
cls.device = _DEVICE
|
||||||
|
|
||||||
def test_update_with_int_phases_does_not_crash(self) -> None:
|
def test_update_with_int_phases_does_not_crash(self) -> None:
|
||||||
checker = SimplePhaseChecker(initial_phase=0, device=self.device)
|
checker = SimplePhaseChecker(initial_phase=0, device=self.device)
|
||||||
checker.enable_assert()
|
checker.enable_assert()
|
||||||
checker.update(expect_phase=0, next_phase=1, caller_name="ints")
|
checker.update(expect_phase=0, next_phase=1, caller_name="ints")
|
||||||
torch.cuda.synchronize()
|
_DEVICE_MODULE.synchronize()
|
||||||
self.assertEqual(_phase_value(checker), 1)
|
self.assertEqual(_phase_value(checker), 1)
|
||||||
|
|
||||||
def test_update_with_intenum_phases_does_not_crash(self) -> None:
|
def test_update_with_intenum_phases_does_not_crash(self) -> None:
|
||||||
@@ -466,7 +486,7 @@ class TestPhaseReprNoCrash(CustomTestCase):
|
|||||||
checker.update(
|
checker.update(
|
||||||
expect_phase=_Phase.IDLE, next_phase=_Phase.A, caller_name="enums"
|
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))
|
self.assertEqual(_phase_value(checker), int(_Phase.A))
|
||||||
|
|
||||||
def test_update_mixing_int_and_intenum_phases(self) -> None:
|
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 = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||||
checker.enable_assert()
|
checker.enable_assert()
|
||||||
checker.update(expect_phase=_Phase.IDLE, next_phase=5, caller_name="mix1")
|
checker.update(expect_phase=_Phase.IDLE, next_phase=5, caller_name="mix1")
|
||||||
torch.cuda.synchronize()
|
_DEVICE_MODULE.synchronize()
|
||||||
self.assertEqual(_phase_value(checker), 5)
|
self.assertEqual(_phase_value(checker), 5)
|
||||||
checker.update(expect_phase=5, next_phase=_Phase.IDLE, caller_name="mix2")
|
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))
|
self.assertEqual(_phase_value(checker), int(_Phase.IDLE))
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -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()
|
||||||
Reference in New Issue
Block a user