diff --git a/docs_new/docs/hardware-platforms/xpu.mdx b/docs_new/docs/hardware-platforms/xpu.mdx index 82264340d..36fa064ca 100644 --- a/docs_new/docs/hardware-platforms/xpu.mdx +++ b/docs_new/docs/hardware-platforms/xpu.mdx @@ -141,6 +141,124 @@ Additionally, the requests can be formed with [OpenAI Completions API](../basic_usage/openai_api_completions) and sent via the command line (e.g. using `curl`) or via your own script. +## XPU Graph [Experimental] + +SGLang enables XPU graph capture to reduce per-step kernel-launch overhead. + +| Phase | Backend | Mechanism | Default | +|---|---|---|---| +| Decode | `full` | One `torch.xpu.XPUGraph` per batch size, captured on startup | **On** | +| Prefill | `tc_piecewise` | `torch.compile` + XPU graph, one graph segment per token-length bucket | **Off** (opt-in) | + +### Enable Prefill Graph + +Prefill graph capture is **opt-in** on XPU and requires `torch.compile` +and must be enabled explicitly: + +```bash +python -m sglang.launch_server --model-path --device xpu \ + --cuda-graph-backend-prefill tc_piecewise +``` + +By default the prefill subgraphs are compiled with `eager` mode. Switch to +`inductor` for higher-quality generated code at the cost of longer startup: + +```bash +python -m sglang.launch_server --model-path --device xpu \ + --cuda-graph-backend-prefill tc_piecewise \ + --cuda-graph-tc-compiler inductor +``` + +You can also configure both phases together with a single `--cuda-graph-config` JSON argument: + +```bash +python -m sglang.launch_server --model-path --device xpu \ + --cuda-graph-config '{"decode":{"backend":"full"},"prefill":{"backend":"tc_piecewise","tc_compiler":"eager"}}' +``` + +### Enable torch.compile for Decode + +`--enable-torch-compile` adds a `torch.compile` pass on top of the decode +XPU graph: the model forward is compiled first, and the compiled forward is +then captured as an `XPUGraph`. This can reduce per-kernel overhead further +but increases startup time. + +```bash +python -m sglang.launch_server --model-path --device xpu \ + --enable-torch-compile +``` + +> **Note:** `--enable-torch-compile` is mutually exclusive with the prefill +> `tc_piecewise` graph (the compatibility rules auto-disable it). Use them +> separately or lock the prefill backend explicitly via `--cuda-graph-config` +> if you need both. + +### Disable XPU Graph + +To opt out of one or both phases: + +```bash +# Disable decode graph +python -m sglang.launch_server --model-path --device xpu \ + --cuda-graph-backend-decode=disabled + +# Disable prefill graph (already off by default; explicit form) +python -m sglang.launch_server --model-path --device xpu \ + --cuda-graph-backend-prefill=disabled + +# Disable both phases +python -m sglang.launch_server --model-path --device xpu \ + --cuda-graph-backend-decode=disabled \ + --cuda-graph-backend-prefill=disabled +``` + +### Customize Capture Buckets + +By default, prefill capture sizes are derived from `--chunked-prefill-size`. +To specify explicit token-length buckets: + +```bash +python -m sglang.launch_server \ + --model-path --device xpu \ + --cuda-graph-backend-prefill tc_piecewise \ + --cuda-graph-bs-prefill 64 128 256 512 +``` + +To specify explicit decode graph batch sizes: + +```bash +python -m sglang.launch_server \ + --model-path --device xpu \ + --cuda-graph-bs-decode 1 2 4 8 +``` + +### Server Args + +| Argument | XPU allowed values | Default | Description | +|---|---|---|---| +| `--cuda-graph-backend-decode` | `full`, `disabled` | `full` | Backend for the decode phase. Only `full` is supported on XPU. | +| `--cuda-graph-backend-prefill` | `tc_piecewise`, `disabled` | `disabled`* | Backend for the prefill phase. Must be set to `tc_piecewise` explicitly to enable. | +| `--cuda-graph-tc-compiler` | `eager`, `inductor` | `eager` | Compiler for `tc_piecewise` prefill subgraphs. `inductor` produces more optimized code but has longer startup. | +| `--cuda-graph-bs-prefill` | list of ints | auto | Explicit token-length buckets to capture for prefill. | +| `--cuda-graph-bs-decode` | list of ints | auto | Explicit batch sizes to capture for decode. | +| `--cuda-graph-config` | JSON string | — | One-shot JSON config for both phases, e.g. `'{"decode":{"backend":"full"},"prefill":{"backend":"tc_piecewise","tc_compiler":"eager"}}'`. Overrides all per-phase flags. | +| `--disable-decode-cuda-graph` | — | `False` | Shorthand for `--cuda-graph-backend-decode=disabled`. | +| `--disable-prefill-cuda-graph` | — | `False` | Shorthand for `--cuda-graph-backend-prefill=disabled`. | +| `--enable-torch-compile` | — | `False` | Apply `torch.compile` on top of the decode XPU graph for further kernel optimization. | +| `--torch-compile-max-bs` | int | `32` | Maximum batch size compiled by `torch.compile` when `--enable-torch-compile` is set. | + +\* Prefill graph is auto-disabled on XPU unless you lock the backend explicitly +via `--cuda-graph-backend-prefill` or `--cuda-graph-config`. + +### Limitations + +| Feature | Status | +|---|---| +| Memory saver (`--enable-memory-saver`) | Not yet supported | +| Two-batch overlap (`--enable-two-batch-overlap`) | Not yet supported | +| Breakable CUDA graph | Not yet supported | +| Speculative decoding | Not yet implemented | + ## Prefill-Decode (P/D) Disaggregation on Intel XPU [Experimental] SGLang supports prefill-decode disaggregation on Intel XPU using the [NIXL](https://github.com/ai-dynamo/nixl) KV-transfer backend. diff --git a/python/sglang/srt/compilation/backend.py b/python/sglang/srt/compilation/backend.py index c735856b2..9c926d4a4 100644 --- a/python/sglang/srt/compilation/backend.py +++ b/python/sglang/srt/compilation/backend.py @@ -23,9 +23,10 @@ from sglang.srt.compilation.compiler_interface import EagerAdapter, InductorAdap from sglang.srt.compilation.cuda_piecewise_backend import CUDAPiecewiseBackend from sglang.srt.compilation.npu_piecewise_backend import NPUPiecewiseBackend from sglang.srt.compilation.pass_manager import PostGradPassManager +from sglang.srt.compilation.xpu_piecewise_backend import XPUPiecewiseBackend from sglang.srt.environ import envs from sglang.srt.platforms import current_platform -from sglang.srt.utils.common import is_npu +from sglang.srt.utils.common import is_npu, is_xpu logger = logging.getLogger(__name__) @@ -53,6 +54,8 @@ def make_backend( if current_platform.is_out_of_tree(): backend_cls = current_platform.get_piecewise_backend_cls() + elif is_xpu(): + backend_cls = XPUPiecewiseBackend elif is_npu(): backend_cls = NPUPiecewiseBackend else: diff --git a/python/sglang/srt/compilation/weak_ref_tensor.py b/python/sglang/srt/compilation/weak_ref_tensor.py index 21d5f5d10..d5f988b1f 100644 --- a/python/sglang/srt/compilation/weak_ref_tensor.py +++ b/python/sglang/srt/compilation/weak_ref_tensor.py @@ -2,14 +2,16 @@ from typing import Any, Union import torch -from sglang.srt.utils.common import is_cuda, is_hip, is_musa, is_npu +from sglang.srt.utils.common import is_cuda, is_hip, is_musa, is_npu, is_xpu -if is_cuda() or is_hip() or is_musa(): +if is_cuda() or is_hip() or is_musa() or is_xpu(): from sgl_kernel import weak_ref_tensor elif is_npu(): from torch_npu._C import _weak_ref_tensor as weak_ref_tensor else: - raise NotImplementedError("weak_ref_tensor is implemented only for CUDA and NPU.") + raise NotImplementedError( + "weak_ref_tensor is implemented only for CUDA, XPU, and NPU." + ) def weak_ref_tensors( diff --git a/python/sglang/srt/compilation/xpu_piecewise_backend.py b/python/sglang/srt/compilation/xpu_piecewise_backend.py new file mode 100644 index 000000000..e85ccb713 --- /dev/null +++ b/python/sglang/srt/compilation/xpu_piecewise_backend.py @@ -0,0 +1,99 @@ +from contextlib import ExitStack +from typing import Any +from unittest.mock import patch + +import torch + +from sglang.srt.compilation.compilation_counter import compilation_counter +from sglang.srt.compilation.compile_phase import ( + get_pcg_capture_stream, + is_in_torch_compile_warmup, +) +from sglang.srt.compilation.cuda_piecewise_backend import ( + CUDAPiecewiseBackend, + weak_ref_tensors, +) + + +class XPUPiecewiseBackend(CUDAPiecewiseBackend): + def __call__(self, *args) -> Any: + if not self.first_run_finished: + self.first_run_finished = True + self.check_for_ending_compilation() + return self.compiled_graph_for_general_shape(*args) + + if len(self.sym_shape_indices) == 0: + return self.compiled_graph_for_general_shape(*args) + + runtime_shape = args[self.sym_shape_indices[0]] + if runtime_shape not in self.concrete_size_entries: + return self.compiled_graph_for_general_shape(*args) + + entry = self.concrete_size_entries[runtime_shape] + + if entry.runnable is None: + entry.runnable = self.compiled_graph_for_general_shape + + if entry.need_to_compile and not entry.compiled: + entry.compiled = True + self.to_be_compiled_sizes.remove(runtime_shape) + entry.runnable = self.sglang_backend.compiler_manager.compile( + self.graph, + args, + self.inductor_config, + graph_index=self.piecewise_compile_index, + num_graphs=self.total_piecewise_compiles, + runtime_shape=runtime_shape, + ) + + if self.is_last_graph and not self.to_be_compiled_sizes: + self.check_for_ending_compilation() + + if is_in_torch_compile_warmup(): + return entry.runnable(*args) + + if entry.cudagraph is None: + if entry.num_finished_warmup < 1: # noqa + entry.num_finished_warmup += 1 + return entry.runnable(*args) + + stream = get_pcg_capture_stream() + assert ( + stream is not None + ), "PCG capture stream is not set, please check if runtime recompilation happened" + + if self.compile_config.get_enable_debug_mode(): + entry.input_addresses = [ + x.data_ptr() for x in args if isinstance(x, torch.Tensor) + ] + + xpugraph = torch.xpu.XPUGraph() + + with ExitStack() as stack: + if not self.is_first_graph: + stack.enter_context(patch("gc.collect", lambda: None)) + stack.enter_context(patch("torch.xpu.empty_cache", lambda: None)) + + with torch.xpu.graph( + xpu_graph=xpugraph, pool=self.graph_pool, stream=stream + ): + output = entry.runnable(*args) + if self.is_last_graph: + output = weak_ref_tensors(output) + + entry.output = weak_ref_tensors(output) + entry.cudagraph = xpugraph + + compilation_counter.num_cudagraph_captured += 1 + return output + + if self.compile_config.get_enable_debug_mode(): + new_input_addresses = [ + x.data_ptr() for x in args if isinstance(x, torch.Tensor) + ] + assert new_input_addresses == entry.input_addresses, ( + "Input addresses for cudagraphs are different during replay." + f" Expected {entry.input_addresses}, got {new_input_addresses}" + ) + entry.cudagraph.replay() + return entry.output diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 327123dd9..58690c38d 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -517,7 +517,7 @@ class GroupCoordinator: def graph_capture( self, graph_capture_context: Optional[GraphCaptureContext] = None, - stream: Optional[torch.cuda.Stream] = None, + stream=None, ): if graph_capture_context is None: if stream is None: @@ -606,7 +606,15 @@ class GroupCoordinator: return self.hpu_communicator.all_reduce(input_) if self.xpu_communicator is not None and not self.xpu_communicator.disabled: - return self.xpu_communicator.all_reduce(input_) + # Route through inplace_all_reduce custom op so Dynamo treats this as + # an opaque call and does not decompose it into _c10d_functional primitives + # (which invoke sycl_event.wait() and break XPU graph capture). + # Keeps the operation in-place; the all-reduce is performed by + # _all_reduce_in_place, which for XPU falls through to + # torch.distributed.all_reduce on self.device_group (the same group + # used by xpu_communicator). + inplace_all_reduce(input_, group_name=self.unique_name) + return input_ if self.npu_communicator is not None and not self.npu_communicator.disabled: return self.npu_communicator.all_reduce(input_) @@ -1002,9 +1010,13 @@ class GroupCoordinator: return envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.get() def all_gather_into_tensor(self, output: torch.Tensor, input: torch.Tensor): - if _is_npu or _is_xpu: + if _is_npu: self._all_gather_into_tensor(output, input) else: + # XPU and CUDA both go through reg_all_gather_into_tensor (custom_op) to + # stay opaque to Dynamo. Calling torch.distributed.all_gather_into_tensor + # directly causes Dynamo to rewrite it as _c10d_functional.all_gather_into_tensor + # + wait_tensor, which invokes sycl_event.wait() and breaks XPU graph capture. reg_all_gather_into_tensor(output, input, group_name=self.unique_name) def cp_all_gather_into_tensor_async( @@ -1762,7 +1774,7 @@ def get_mooncake_transfer_engine(): @contextmanager -def graph_capture(stream: Optional[torch.cuda.Stream] = None): +def graph_capture(stream=None): """ `graph_capture` is a context manager which should surround the code that is capturing the CUDA graph. Its main purpose is to ensure that the diff --git a/python/sglang/srt/hardware_backend/xpu/graph_runner/__init__.py b/python/sglang/srt/hardware_backend/xpu/graph_runner/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/sglang/srt/hardware_backend/xpu/graph_runner/xpu_full_graph_backend.py b/python/sglang/srt/hardware_backend/xpu/graph_runner/xpu_full_graph_backend.py new file mode 100644 index 000000000..6399c9df3 --- /dev/null +++ b/python/sglang/srt/hardware_backend/xpu/graph_runner/xpu_full_graph_backend.py @@ -0,0 +1,100 @@ +"""FullXPUGraphBackend — Intel XPU full-graph capture (torch.xpu.XPUGraph). + +Mirrors FullCudaGraphBackend with XPU-specific differences: + - Captures via torch.xpu.graph(xpu_graph=...) into torch.xpu.XPUGraph. + - Shares the global graph memory pool with the prefill backend so that + decode + prefill graphs occupy max(decode, prefill) rather than their sum. + - No set_graph_pool_id: SymmetricMemoryContext is never triggered on XPU + (oneCCL has no ncclMemAlloc equivalent; enable_symm_mem defaults False). +""" + +from __future__ import annotations + +from contextlib import contextmanager +from typing import TYPE_CHECKING, Any, Callable, Dict, Optional + +import torch + +from sglang.srt.model_executor.runner.shape_key import ShapeKey +from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import ( + BaseCudaGraphBackend, +) +from sglang.srt.model_executor.runner_utils.pool import ( + get_or_create_global_graph_memory_pool, +) + +if TYPE_CHECKING: + from sglang.srt.model_executor.forward_batch_info import ForwardBatch + from sglang.srt.model_executor.runner.base_cuda_graph_runner import ( + BaseCudaGraphRunner, + ) + + +class FullXPUGraphBackend(BaseCudaGraphBackend): + """One torch.xpu.XPUGraph per shape for Intel XPU devices.""" + + def __init__( + self, + cuda_graph_runner: BaseCudaGraphRunner, + ) -> None: + self._graphs: Dict[Any, torch.xpu.XPUGraph] = {} + self._outputs: Dict[Any, Any] = {} + self._pool = None + self._device_module = cuda_graph_runner.device_module + self._tp_group = cuda_graph_runner.model_runner.tp_group + self._capture_stream: Optional[torch.xpu.Stream] = None + + @contextmanager + def capture_session(self, stream: torch.xpu.Stream): + if self._pool is None: + self._pool = get_or_create_global_graph_memory_pool(self._device_module) + self._capture_stream = stream + try: + yield + finally: + self._capture_stream = None + + def capture_one( + self, + shape_key: ShapeKey, + forward_fn: Callable[[], Any], + dummies: Optional[Any] = None, + post_warmup_hook: Optional[Callable[[], None]] = None, + ) -> None: + for _ in range(2): + self._device_module.synchronize() + self._tp_group.barrier() + forward_fn() + if post_warmup_hook is not None: + post_warmup_hook() + + graph = torch.xpu.XPUGraph() + + with self._device_module.graph( + xpu_graph=graph, pool=self._pool, stream=self._capture_stream + ): + out = forward_fn() + + self._graphs[shape_key] = graph + self._outputs[shape_key] = out + + def can_run(self, forward_batch: ForwardBatch, shape_key: ShapeKey) -> bool: + return shape_key in self._graphs + + @contextmanager + def replay_session(self): + yield + + def replay( + self, + shape_key: ShapeKey, + static_forward_batch: ForwardBatch, + **kwargs, + ) -> Any: + self._graphs[shape_key].replay() + return self._outputs[shape_key] + + def cleanup(self) -> None: + self._graphs.clear() + self._outputs.clear() + self._pool = None diff --git a/python/sglang/srt/hardware_backend/xpu/graph_runner/xpu_graph_runner.py b/python/sglang/srt/hardware_backend/xpu/graph_runner/xpu_graph_runner.py new file mode 100644 index 000000000..1b53564eb --- /dev/null +++ b/python/sglang/srt/hardware_backend/xpu/graph_runner/xpu_graph_runner.py @@ -0,0 +1,173 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Run the model with xpu graph and torch.compile.""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +import torch +from torch.profiler import ProfilerActivity, profile + +from sglang.srt.model_executor.runner import DecodeCudaGraphRunner +from sglang.srt.speculative.spec_info import SpeculativeAlgorithm +from sglang.srt.utils import register_xpu_device_properties_for_dynamo + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from sglang.srt.model_executor.model_runner import ModelRunner + + +_fake_ops_registered = False + + +def register_fake_ops(): + """Register fake/abstract implementations for XPU sgl_kernel ops so that + torch.compile (Dynamo) can trace through them using FakeTensors for shape + and dtype propagation, without executing the real GPU kernels. + """ + global _fake_ops_registered + if _fake_ops_registered: + return + _fake_ops_registered = True + + @torch.library.register_fake("sgl_kernel::fwd") + def _( + q, + k, + v, + q_v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx, + leftpad_k, + rotary_cos, + rotary_sin, + seqlens_rotary, + q_descale, + k_descale, + v_descale, + softmax_scale, + sinks, + is_causal, + window_size_left, + window_size_right, + softcap, + is_rotary_interleaved, + scheduler_metadata, + num_kv_splits, + pack_gqa, + sm_margin, + out=None, + ): + total_q = q.shape[0] + num_heads_q = q.shape[1] + head_size_v = v.shape[-1] + if out is None: + out = q.new_empty(total_q, num_heads_q, head_size_v) + softmax_lse = q.new_empty(num_heads_q, total_q, dtype=torch.float32) + # out_accum and softmax_lse_accum are intermediate split-kv buffers; + # they are only read when num_kv_splits > 1, which is determined at + # runtime. Return empty tensors with correct rank so downstream ops + # that index into the list do not fail shape propagation. + out_accum = q.new_empty(0) + softmax_lse_accum = q.new_empty(0, dtype=torch.float32) + return (out, softmax_lse, out_accum, softmax_lse_accum) + + @torch.library.register_fake("sgl_kernel::flash_mla_decode") + def _( + out, + q_nope, + q_pe, + kv_c_and_k_pe_cache, + seq_lens, + page_table, + workspace, + sm_scale, + num_kv_splits, + ): + return + + +class XPUGraphRunner(DecodeCudaGraphRunner): + """A XPUGraphRunner runs the forward pass of a model with xpu graph and torch.compile.""" + + @staticmethod + def _apply_xpu_compile_config() -> None: + """Apply XPU-specific torch.compile / dynamo settings. + + Called unconditionally before super().__init__() so that the settings + are in place regardless of whether --enable-torch-compile is passed. + The critical flag is suppress_errors: when the Intel IGC compiler + crashes with SIGFPE on certain reduction kernels (ocloc -device bmg + returns exit code 245), dynamo falls back to eager for that subgraph + instead of propagating the crash. + """ + import torch._dynamo.config + + torch._dynamo.config.suppress_errors = True + + def __init__(self, model_runner: ModelRunner): + assert ( + not model_runner.server_args.enable_memory_saver + ), "XPUGraphRunner does not support Torch Memory Saver yet." + register_fake_ops() + self._apply_xpu_compile_config() + register_xpu_device_properties_for_dynamo() + super().__init__(model_runner) + + assert ( + not self.enable_two_batch_overlap + ), "XPUGraphRunner does not support two batch overlap yet." + assert ( + not self.require_mlp_tp_gather + ), "XPUGraphRunner does not support MLP TP gather yet." + assert ( + not self.require_mlp_sync + ), "XPUGraphRunner does not support MLP sync yet." + assert ( + not self.require_gathered_buffer + ), "XPUGraphRunner does not support gathered buffer yet." + assert ( + model_runner.spec_algorithm == SpeculativeAlgorithm.NONE + ), "XPUGraphRunner does not support speculative inference yet." + + def _init_profile_context_and_memory_record(self): + profile_context = profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.XPU], + record_shapes=True, + ) + torch.xpu.memory._record_memory_history() + return profile_context + + def _post_process_after_profile(self, prof_context): + torch.xpu.memory._dump_snapshot("xpu_graph_runner_memory_usage.pickle") + torch.xpu.memory._record_memory_history(enabled=None) + log_message = ( + "Sorted by XPU Time:\n" + + prof_context.key_averages(group_by_input_shape=True).table( + sort_by="self_xpu_time_total" + ) + + "\n\nSorted by CPU Time:\n" + + prof_context.key_averages(group_by_input_shape=True).table( + sort_by="self_cpu_time_total" + ) + + "\n\nMemory Usage is saved to xpu_graph_runner_memory_usage.pickle\n" + ) + logger.info(log_message) diff --git a/python/sglang/srt/layers/attention/fla/layernorm_gated.py b/python/sglang/srt/layers/attention/fla/layernorm_gated.py index a076d161c..0fa7e032f 100644 --- a/python/sglang/srt/layers/attention/fla/layernorm_gated.py +++ b/python/sglang/srt/layers/attention/fla/layernorm_gated.py @@ -6,6 +6,7 @@ # The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. +from contextlib import nullcontext from functools import lru_cache import torch @@ -254,7 +255,21 @@ def _layer_norm_fwd( # Update grid to use rows_per_block grid = (cdiv(M, rows_per_block), ngroups) pdl_kwargs = {"USE_GDC": True, "launch_pdl": True} if is_arch_support_pdl() else {} - with device_context(x.device): + # Workaround for PyTorch <= 2.12: torch.xpu.device is not Dynamo-compatible + # in that release — it creates a DynamoConfigPatchProxy that + # SourcelessBuilder cannot wrap, causing a hard error under + # torch.compile(fullgraph=True). The device context is a functional no-op + # for Triton kernel launches (device is determined by the tensor, not the + # surrounding context), so we simply skip it when Dynamo is tracing. + # PyTorch main already has the proper fix (XPUDeviceVariable registered in + # torch/_dynamo/variables/ctx_manager.py analogous to CUDADeviceVariable). + # TODO: remove this branch once we upgrade from PyTorch 2.12. + device_ctx = ( + nullcontext() + if x.device.type == "xpu" and torch.compiler.is_compiling() + else device_context(x.device) + ) + with device_ctx: _layer_norm_fwd_1pass_kernel[grid]( x, out, diff --git a/python/sglang/srt/layers/attention/xpu_backend.py b/python/sglang/srt/layers/attention/xpu_backend.py index 6dea33696..984e4c6bf 100644 --- a/python/sglang/srt/layers/attention/xpu_backend.py +++ b/python/sglang/srt/layers/attention/xpu_backend.py @@ -104,6 +104,7 @@ class XPUAttentionBackend(AttentionBackend): self.has_swa = ( self.sliding_window_size is not None and self.sliding_window_size > -1 ) + self.is_encoder_decoder = model_runner.model_config.is_encoder_decoder def init_forward_metadata(self, forward_batch: ForwardBatch): """Initialize forward metadata hence all layers in the forward pass can reuse it.""" @@ -411,22 +412,6 @@ class XPUAttentionBackend(AttentionBackend): workspace_size, device=self.device, dtype=torch.uint8 ) - # Translate full-pool indices to SWA-pool indices for hybrid models - if self.use_sliding_window_kv_pool: - # flash_attn_with_kvcache requires int32 page tables; the SWA index - # mapping is int64, so cast (matches flashattention_backend.py). - metadata.swa_page_table = ( - self.token_to_kv_pool.translate_loc_from_full_to_swa( - metadata.page_table - ).to(torch.int32) - ) - if forward_batch.out_cache_loc is not None: - metadata.swa_out_cache_loc = ( - self.token_to_kv_pool.translate_loc_from_full_to_swa( - forward_batch.out_cache_loc - ) - ) - # Convert the page table to a strided format which is needed by FA3 API if self.page_size > 1: self.strided_indices = torch.arange( @@ -477,7 +462,10 @@ class XPUAttentionBackend(AttentionBackend): if not self.use_mla: self.token_to_kv_pool.set_kv_buffer( layer, - KVWriteLoc(cache_loc, self.forward_metadata.swa_out_cache_loc), + KVWriteLoc( + cache_loc, + self.forward_metadata.swa_out_cache_loc, + ), k, v, layer.k_scale, @@ -599,6 +587,19 @@ class XPUAttentionBackend(AttentionBackend): k_descale=k_descale, v_descale=v_descale, return_softmax_lse=use_cascade_attn, + # Piecewise XPU graph for prefill requires a pre-allocated + # output buffer at a stable device address so the graph can + # record writes to the same storage on every replay. + # _attn_output is that fixed buffer; None falls back to a + # freshly allocated tensor (eager / cascade-attn path). + out=( + forward_batch._attn_output.view( + -1, layer.tp_q_head_num, layer.v_head_dim + ) + if not use_cascade_attn + and getattr(forward_batch, "_attn_output", None) is not None + else None + ), **kwargs, ) @@ -783,7 +784,10 @@ class XPUAttentionBackend(AttentionBackend): if not self.use_mla: self.token_to_kv_pool.set_kv_buffer( layer, - KVWriteLoc(cache_loc, self.forward_metadata.swa_out_cache_loc), + KVWriteLoc( + cache_loc, + self.forward_metadata.swa_out_cache_loc, + ), k, v, layer.k_scale, @@ -997,6 +1001,208 @@ class XPUAttentionBackend(AttentionBackend): """Get the fill value for sequence length in CUDA graph.""" return 1 + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + """Pre-allocate fixed-size tensors reused across XPU graph captures.""" + max_num_pages = (self.max_context_len + self.page_size - 1) // self.page_size + self.decode_cuda_graph_metadata = { + "cache_seqlens": torch.zeros(max_bs, dtype=torch.int32, device=self.device), + "cu_seqlens_q": torch.arange( + 0, max_bs + 1, dtype=torch.int32, device=self.device + ), + "cu_seqlens_k": torch.zeros( + max_bs + 1, dtype=torch.int32, device=self.device + ), + "page_table": torch.zeros( + max_bs, max_num_pages, dtype=torch.int32, device=self.device + ), + "strided_indices": torch.arange( + 0, self.max_context_len, self.page_size, device=self.device + ), + } + if self.use_sliding_window_kv_pool: + self.decode_cuda_graph_metadata["swa_page_table"] = torch.zeros( + max_bs, max_num_pages, dtype=torch.int32, device=self.device + ) + self.decode_cuda_graph_metadata["swa_out_cache_loc"] = torch.zeros( + max_num_tokens, dtype=torch.int64, device=self.device + ) + if self.is_encoder_decoder: + self.encoder_metadata = { + "encoder_page_table": torch.zeros( + max_bs, self.max_context_len, dtype=torch.int32, device=self.device + ), + "encoder_lens_int32": torch.zeros( + max_bs, dtype=torch.int32, device=self.device + ), + "encoder_cu_seqlens_k": torch.zeros( + max_bs + 1, dtype=torch.int32, device=self.device + ), + } + else: + self.encoder_metadata = {} + + def init_forward_metadata_out_graph( + self, + forward_batch: ForwardBatch, + in_capture: bool = False, + ): + """New unified graph capture/replay entry point (replaces the legacy + init_forward_metadata_capture_cuda_graph / + init_forward_metadata_replay_cuda_graph pair). + + Called by DecodeCudaGraphRunner: + - capture: in_capture=True → bind static metadata buffer slices, then fill + - replay: in_capture=False → update pre-allocated buffers in-place + - eager: via init_forward_metadata() default wrapper + """ + bs = forward_batch.batch_size + req_pool_indices = forward_batch.req_pool_indices + seq_lens = forward_batch.seq_lens + seq_lens_cpu = getattr(forward_batch, "seq_lens_cpu", None) + forward_mode = forward_batch.forward_mode + spec_info = forward_batch.spec_info + + assert ( + spec_info is None + ), "XPUAttentionBackend does not support speculative decoding in XPU graph" + assert ( + forward_mode.is_decode_or_idle() + ), "XPUAttentionBackend XPU graph only supports decode mode" + + if in_capture: + # Bind static-shape slices of the pre-allocated buffers so the + # captured graph always reads from the same storage addresses. + metadata = FlashAttentionMetadata() + metadata.cache_seqlens_int32 = self.decode_cuda_graph_metadata[ + "cache_seqlens" + ][:bs] + metadata.cu_seqlens_q = self.decode_cuda_graph_metadata["cu_seqlens_q"][ + : bs + 1 + ] + metadata.cu_seqlens_k = self.decode_cuda_graph_metadata["cu_seqlens_k"][ + : bs + 1 + ] + metadata.page_table = self.decode_cuda_graph_metadata["page_table"][:bs, :] + if self.use_sliding_window_kv_pool: + # Bind SWA page table slice so the graph captures the right tensor. + metadata.swa_page_table = self.decode_cuda_graph_metadata[ + "swa_page_table" + ][:bs, :] + if self.is_encoder_decoder and forward_batch.encoder_lens is not None: + encoder_bs = forward_batch.encoder_lens.numel() + metadata.encoder_lens_int32 = self.encoder_metadata[ + "encoder_lens_int32" + ][:encoder_bs] + metadata.encoder_cu_seqlens_k = self.encoder_metadata[ + "encoder_cu_seqlens_k" + ][: encoder_bs + 1] + metadata.encoder_page_table = self.encoder_metadata[ + "encoder_page_table" + ][:bs, :] + self.decode_cuda_graph_metadata[bs] = metadata + + # Both capture and replay: fill data into the pre-allocated buffers. + seq_lens = seq_lens[:bs] + seq_lens_cpu = seq_lens_cpu[:bs] if seq_lens_cpu is not None else None + req_pool_indices = req_pool_indices[:bs] + + metadata = self.decode_cuda_graph_metadata[bs] + max_len = ( + seq_lens_cpu.max().item() + if seq_lens_cpu is not None + else seq_lens.max().item() + ) + metadata.max_seq_len_k = max_len + + metadata.cache_seqlens_int32.copy_(seq_lens.to(torch.int32)) + + metadata.cu_seqlens_k[0] = 0 + metadata.cu_seqlens_k[1 : bs + 1].copy_( + torch.cumsum(seq_lens.to(torch.int32), dim=0) + ) + + if self.is_encoder_decoder and forward_batch.encoder_lens is not None: + encoder_lens = forward_batch.encoder_lens[:bs].to(torch.int32) + metadata.encoder_max_seq_len_k = int(encoder_lens.max().item()) + metadata.encoder_lens_int32.copy_(encoder_lens) + metadata.encoder_cu_seqlens_k[0] = 0 + metadata.encoder_cu_seqlens_k[1 : bs + 1].copy_( + torch.cumsum(encoder_lens, dim=0, dtype=torch.int32) + ) + metadata.encoder_page_table[:bs, : metadata.encoder_max_seq_len_k].copy_( + self.req_to_token[ + req_pool_indices, : metadata.encoder_max_seq_len_k + ].to(torch.int32) + ) + # Self-attention (text) page_table: decoder tokens start after encoder tokens. + text_max = metadata.max_seq_len_k + arange_text = torch.arange(text_max, device=req_pool_indices.device) + text_col = encoder_lens[:bs].long().unsqueeze(1) + arange_text.unsqueeze(0) + text_row = req_pool_indices.unsqueeze(1).expand(-1, text_max) + metadata.page_table[:bs, :text_max].copy_( + self.req_to_token[text_row, text_col].to(torch.int32) + ) + metadata.page_table[:bs, text_max:].zero_() + else: + raw_page = self.req_to_token[ + req_pool_indices[:, None], + self.decode_cuda_graph_metadata["strided_indices"][ + : ((metadata.max_seq_len_k + self.page_size - 1) // self.page_size) + ][None, :], + ] + if self.page_size > 1: + raw_page = raw_page // self.page_size + metadata.page_table[:bs, : raw_page.shape[1]].copy_( + raw_page.to(torch.int32) + ) + metadata.page_table[:bs, raw_page.shape[1] :].zero_() + + if self.use_sliding_window_kv_pool: + if forward_batch.out_cache_loc is None: + raise ValueError( + f"out_cache_loc is None for hybrid SWA model in graph " + f"{'capture' if in_capture else 'replay'} " + f"(forward_mode={forward_batch.forward_mode}). This should not happen." + ) + swa_out_cache_loc = self.decode_cuda_graph_metadata["swa_out_cache_loc"] + n = forward_batch.out_cache_loc.shape[0] + swa_out_cache_loc[n:].zero_() + swa_out_cache_loc[:n].copy_( + self.token_to_kv_pool.translate_loc_from_full_to_swa( + forward_batch.out_cache_loc + ) + ) + metadata.swa_out_cache_loc = swa_out_cache_loc[:n] + + if not (self.is_encoder_decoder and forward_batch.encoder_lens is not None): + max_seq_pages = ( + metadata.max_seq_len_k + self.page_size - 1 + ) // self.page_size + swa_page_table = self.decode_cuda_graph_metadata["swa_page_table"] + swa_page_table[:bs, max_seq_pages:].zero_() + swa_page_table[:bs, :max_seq_pages].copy_( + ( + self.token_to_kv_pool.translate_loc_from_full_to_swa(raw_page) + if self.page_size == 1 + else self.token_to_kv_pool.translate_loc_from_full_to_swa( + self.req_to_token[ + req_pool_indices[:, None], + self.decode_cuda_graph_metadata["strided_indices"][ + :max_seq_pages + ][None, :], + ] + ) + // self.page_size + ).to(torch.int32) + ) + metadata.swa_page_table = swa_page_table[:bs, :] + + self.forward_metadata = metadata + + def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch): + """Graph-recordable ops for XPU graph (no-op: all metadata setup is + host-side and lives in init_forward_metadata_out_graph).""" + def _init_local_attn_metadata( self, forwardbatch: ForwardBatch, diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 86eaf61a3..4d09d5d2b 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -112,6 +112,7 @@ from sglang.srt.eplb.lplb_solver import ( set_global_lplb_solver, ) from sglang.srt.hardware_backend.npu.graph_runner.npu_graph_runner import NPUGraphRunner +from sglang.srt.hardware_backend.xpu.graph_runner.xpu_graph_runner import XPUGraphRunner from sglang.srt.kv_canary.api import install_canary from sglang.srt.kv_canary.runner.canary_manager import context_tuple from sglang.srt.kv_canary.token_oracle.install import install_token_oracle_from_env @@ -879,7 +880,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): if self.device == "cuda" or self.device == "musa": self.init_cublas() self.init_attention_backend() - elif self.device == "cpu": + elif self.device in ["cpu", "xpu"]: self.init_attention_backend() elif self.device == "npu": self.init_attention_backend() @@ -922,7 +923,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): self.graph_mem_usage = 0 if capture_decode_cuda_graph: - if self.device in ("cuda", "musa", "cpu", "npu"): + if self.device in ("cuda", "musa", "cpu", "npu", "xpu"): self.init_decode_cuda_graph() elif ( current_platform.is_out_of_tree() @@ -2601,6 +2602,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): "musa": "CUDA graph", "cpu": "CPU graph", "npu": "NPU graph", + "xpu": "XPU graph", }, ) role = "draft" if self.is_draft_worker else "target" @@ -2636,6 +2638,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): { "cpu": CPUGraphRunner, "npu": NPUGraphRunner, + "xpu": XPUGraphRunner, }, ) self.decode_cuda_graph_runner = graph_runners[self.device](self) diff --git a/python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py b/python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py index 0054f0f2a..0974768a4 100644 --- a/python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py +++ b/python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py @@ -148,12 +148,17 @@ class TcPiecewiseCudaGraphBackend(BaseCudaGraphBackend): enable_torch_compile_warmup to drive FX / inductor through every shape without capturing cuda graphs yet.""" language_model = self._language_model + # Some multimodal models (e.g. Gemma4) store the inner transformer + # directly as `language_model` rather than wrapping it in a + # ForCausalLM that has a `.model` child. Fall back to the module + # itself when `.model` is absent. + inner_model = getattr(language_model, "model", language_model) compiler = self._compile_config.compiler with enable_tc_piecewise_cuda_graph(): try: if compiler != "eager": _toggle_multi_platform_ops( - language_model.model, reverse=False, num_tokens=16 + inner_model, reverse=False, num_tokens=16 ) cuda_graph_runner._run_dummy_forward( @@ -167,7 +172,7 @@ class TcPiecewiseCudaGraphBackend(BaseCudaGraphBackend): set_graph_pool_id(self._pool) self.install_compile( - language_model.model, + inner_model, compile_config=self._compile_config, graph_pool=self._pool, ) @@ -196,9 +201,7 @@ class TcPiecewiseCudaGraphBackend(BaseCudaGraphBackend): ) cuda_graph_runner._run_dummy_forward(num_tokens=num_tokens) finally: - _toggle_multi_platform_ops( - language_model.model, reverse=True, num_tokens=16 - ) + _toggle_multi_platform_ops(inner_model, reverse=True, num_tokens=16) @contextmanager def capture_session(self, stream: torch.cuda.Stream): diff --git a/python/sglang/srt/model_executor/runner_backend/utils.py b/python/sglang/srt/model_executor/runner_backend/utils.py index f36031ad4..ac11a8670 100644 --- a/python/sglang/srt/model_executor/runner_backend/utils.py +++ b/python/sglang/srt/model_executor/runner_backend/utils.py @@ -78,6 +78,17 @@ def resolve_decode_backend( return XPUCudaGraphBackend(cuda_graph_runner) + if model_runner.device == "xpu": + if backend_name not in (Backend.FULL, Backend.DISABLED): + raise ValueError( + f"XPU only supports cuda_graph_config decode backend 'full', got '{backend_name}'" + ) + from sglang.srt.hardware_backend.xpu.graph_runner.xpu_full_graph_backend import ( + FullXPUGraphBackend, + ) + + return FullXPUGraphBackend(cuda_graph_runner) + if backend_name == Backend.BREAKABLE: return BreakableCudaGraphBackend( cuda_graph_runner, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 459a36636..dc0945a6b 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -3054,12 +3054,15 @@ class ServerArgs: def _handle_xpu_backends(self): if self.device == "xpu": - if self.cuda_graph_config.prefill.backend != Backend.DISABLED: + if self.cuda_graph_config.prefill.backend not in ( + Backend.DISABLED, + Backend.TC_PIECEWISE, + ): logger.warning( - "XPU platform does not support piecewise CUDA graph, " - "disabling prefill cuda graph." + "XPU platform currently only supports prefill tc_piecewise CUDA graph; " + "disabling unsupported prefill backend." ) - self.cuda_graph_config.prefill.backend = Backend.DISABLED + self.cuda_graph_config.prefill.backend = Backend.DISABLED # ------------------------------------------------------------------ # CUDA graph configuration resolution diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 858b57ca7..686c118e8 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -179,6 +179,17 @@ def is_xpu() -> bool: return hasattr(torch, "xpu") and torch.xpu.is_available() +def register_xpu_device_properties_for_dynamo() -> None: + if not is_xpu(): + return + + import torch._dynamo.utils as dynamo_utils + + xpu_props_type = getattr(torch.xpu, "_XpuDeviceProperties", None) + if xpu_props_type is not None: + dynamo_utils.common_constant_types.add(xpu_props_type) + + @lru_cache(maxsize=1) def is_npu() -> bool: if not hasattr(torch, "npu"): @@ -610,9 +621,9 @@ def get_available_gpu_memory( if empty_cache: empty_device_cache(torch.xpu) - used_memory = torch.xpu.memory_allocated() - total_gpu_memory = torch.xpu.get_device_properties(gpu_id).total_memory - free_gpu_memory = total_gpu_memory - used_memory + # Use mem_get_info() to reflect true OS-level free memory + # including graph pool reservations; avoids KV-cache over-allocation. + free_gpu_memory, total_gpu_memory = torch.xpu.mem_get_info(gpu_id) elif device == "hpu": num_gpus = torch.hpu.device_count() diff --git a/test/registered/xpu/test_deepseek_ocr.py b/test/registered/xpu/test_deepseek_ocr.py index 33e38c789..ff8f6cc1e 100644 --- a/test/registered/xpu/test_deepseek_ocr.py +++ b/test/registered/xpu/test_deepseek_ocr.py @@ -38,6 +38,7 @@ class TestDeepSeekOCR(CustomTestCase): "xpu", "--attention-backend", "intel_xpu", + "--disable-decode-cuda-graph", ] os.environ["SGLANG_USE_SGL_XPU"] = "1" cls.process = popen_launch_server( diff --git a/test/registered/xpu/test_deepseek_ocr_triton.py b/test/registered/xpu/test_deepseek_ocr_triton.py index b17c29710..c34062e41 100644 --- a/test/registered/xpu/test_deepseek_ocr_triton.py +++ b/test/registered/xpu/test_deepseek_ocr_triton.py @@ -41,6 +41,7 @@ class TestDeepSeekOCRTriton(TestDeepSeekOCR): "xpu", "--attention-backend", "intel_xpu", + "--disable-decode-cuda-graph", ] os.environ["SGLANG_USE_SGL_XPU"] = "0" cls.process = popen_launch_server( diff --git a/test/registered/xpu/test_encoder_attention_backend.py b/test/registered/xpu/test_encoder_attention_backend.py index 1f64c15df..9dfc69fda 100644 --- a/test/registered/xpu/test_encoder_attention_backend.py +++ b/test/registered/xpu/test_encoder_attention_backend.py @@ -39,6 +39,7 @@ class TestEncoderAttention(CustomTestCase): "xpu", "--mm-attention-backend", "xpu_attn", + "--disable-decode-cuda-graph", ] os.environ["SGLANG_USE_SGL_XPU"] = "1" cls.process = popen_launch_server( @@ -127,6 +128,7 @@ class TestEncoderAttention_Triton(TestEncoderAttention): "xpu", "--mm-attention-backend", "triton_attn", + "--disable-decode-cuda-graph", ] os.environ["SGLANG_USE_SGL_XPU"] = "0" cls.process = popen_launch_server( diff --git a/test/registered/xpu/test_gemma_4_e2b.py b/test/registered/xpu/test_gemma_4_e2b.py index ebe1e2dd2..45a85c1c9 100644 --- a/test/registered/xpu/test_gemma_4_e2b.py +++ b/test/registered/xpu/test_gemma_4_e2b.py @@ -53,6 +53,7 @@ XPU_SERVER_ARGS = [ "intel_xpu", "--model-impl", "sglang", + "--disable-decode-cuda-graph", ] # Standard sglang e2e Q&A prompt (see test_openai_server.py::run_chat_completion). diff --git a/test/registered/xpu/test_intel_xpu_backend.py b/test/registered/xpu/test_intel_xpu_backend.py index b38c57984..cf6ec0077 100644 --- a/test/registered/xpu/test_intel_xpu_backend.py +++ b/test/registered/xpu/test_intel_xpu_backend.py @@ -34,6 +34,7 @@ def intel_xpu_benchmark( "1", "--device", "xpu", + "--disable-decode-cuda-graph", ] ci_args = ["--input", "64", "--output", "4"] if is_in_ci() else [] full_args = common_args + ci_args + (extra_args or []) diff --git a/test/registered/xpu/test_triton_attention_backend.py b/test/registered/xpu/test_triton_attention_backend.py index 72e0a99e7..4288fd81c 100644 --- a/test/registered/xpu/test_triton_attention_backend.py +++ b/test/registered/xpu/test_triton_attention_backend.py @@ -31,6 +31,7 @@ def triton_attention_benchmark(extra_args=None, mem_fraction_static="0.84"): "2050", "--attention-backend", "triton", + "--disable-decode-cuda-graph", ] full_args = common_args + (extra_args or []) diff --git a/test/registered/xpu/test_xpu_basic.py b/test/registered/xpu/test_xpu_basic.py index c1e4c30ea..024fd7809 100644 --- a/test/registered/xpu/test_xpu_basic.py +++ b/test/registered/xpu/test_xpu_basic.py @@ -33,6 +33,7 @@ class TestXPUBasic(CustomTestCase): "0.6", "--batch-size", "1", + "--disable-decode-cuda-graph", ] if is_in_ci(): args += ["--input", "64", "--output", "4"] diff --git a/test/registered/xpu/test_xpu_embedding.py b/test/registered/xpu/test_xpu_embedding.py index fdb8c5895..c07d9d3e4 100644 --- a/test/registered/xpu/test_xpu_embedding.py +++ b/test/registered/xpu/test_xpu_embedding.py @@ -34,7 +34,12 @@ class TestXPUEmbedding(CustomTestCase): cls.model, cls.base_url, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=["--is-embedding", "--device", "xpu"], + other_args=[ + "--is-embedding", + "--device", + "xpu", + "--disable-decode-cuda-graph", + ], ) cls.openai_url = cls.base_url + "/v1" diff --git a/test/registered/xpu/test_xpu_graph.py b/test/registered/xpu/test_xpu_graph.py new file mode 100644 index 000000000..d5123d60a --- /dev/null +++ b/test/registered/xpu/test_xpu_graph.py @@ -0,0 +1,71 @@ +""" +XPU graph tests: verifies decode full-graph and prefill tc_piecewise graph +on Intel XPU produce valid outputs. + + - TestXPUGraph : decode full-graph and prefill tc_piecewise graph enabled + together in a single bench_one_batch invocation. + +Usage: + python3 -m unittest test_xpu_graph.TestXPUGraph +""" + +import unittest + +from sglang.test.ci.ci_register import register_xpu_ci +from sglang.test.test_utils import ( + DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN, + CustomTestCase, + is_in_ci, + run_bench_one_batch, +) + +register_xpu_ci(est_time=600, suite="stage-b-test-1-gpu-xpu") + +_COMMON_ARGS = [ + "--device", + "xpu", + "--attention-backend", + "triton", + "--disable-radix-cache", + "--mem-fraction-static", + "0.6", + "--batch-size", + "1", +] + +_CI_IO_ARGS = ["--input", "64", "--output", "4"] +_FULL_IO_ARGS = ["--input", "128", "--output", "16"] + + +class TestXPUGraph(CustomTestCase): + """Decode full-graph + prefill tc_piecewise together.""" + + def test_full_graph_runs(self): + args = [ + *_COMMON_ARGS, + "--cuda-graph-config", + '{"decode":{"backend":"full"},"prefill":{"backend":"tc_piecewise","tc_compiler":"eager"}}', + "--cuda-graph-bs-prefill", + "64", + "128", + ] + if is_in_ci(): + args += _CI_IO_ARGS + else: + args += _FULL_IO_ARGS + + prefill_latency, decode_throughput, _ = run_bench_one_batch( + DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN, args + ) + self.assertGreater( + prefill_latency, + 0, + "prefill latency must be > 0 with tc_piecewise XPU graph", + ) + self.assertGreater( + decode_throughput, 0, "decode throughput must be > 0 with full XPU graph" + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/xpu/test_xpu_serving_features.py b/test/registered/xpu/test_xpu_serving_features.py index 1ea21b6d5..2f4746a9a 100644 --- a/test/registered/xpu/test_xpu_serving_features.py +++ b/test/registered/xpu/test_xpu_serving_features.py @@ -43,7 +43,7 @@ class TestXPUServingFeatures(CustomTestCase): cls.model, cls.base_url, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=["--device", "xpu"], + other_args=["--device", "xpu", "--disable-decode-cuda-graph"], ) cls.openai_url = cls.base_url + "/v1"