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

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

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