[XPU] Enable XPU graph support (decode full-graph + prefill tc_piecewise) (#29053)
This commit is contained in:
Cao E
@@ -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:
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user