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

[XPU] weekly simple model enablement 2026/09/14 (#39439)

Browse Source 
Co-authored-by: Juan Muneton <102537701+jmunetong@users.noreply.github.com>
Co-authored-by: YangKai0616 <kai.yang@intel.com>
Co-authored-by: devan-carlin <devan-carlin@users.noreply.github.com>
Co-authored-by: Ashwini Rathi <arathi@habana.ai>
Co-authored-by: Ranjan Debnath <ranjan.debnath@intel.com>
Co-authored-by: Juan Muneton <juan.muneton.gallego@intel.com>
Co-authored-by: Amrutha M <amrutha.m@intel.com>
This commit is contained in:
Meng, Hengyu
2026-09-17 10:43:04 +08:00
committed by GitHub
co-authored by Juan Muneton YangKai0616 devan-carlin Ashwini Rathi Ranjan Debnath Juan Muneton Amrutha M
parent 4793f56835
commit 84d7604b7e
11 changed files with 309 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/benchmark/one_batch.py
+10
View File
@@ -69,6 +69,7 @@ from sglang.srt.arg_groups.overrides import (
resolution_result,
resolving_view,
)
from sglang.srt.configs.hybrid_arch import mambaish_config
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.distributed.parallel_state import (
destroy_distributed_environment,
@@ -369,6 +370,15 @@ def load_model(server_args, port_args, gpu_id, tp_rank):
model_runner.start_startup_weight_load()
model_runner.alloc_memory_pool()
model_runner.init_attention_backends()
# bench_one_batch bypasses the Scheduler, so the Mamba SSU backend that
# Scheduler.init_mamba_backend() would set up is never initialized. Do it
# here (per tp_rank, i.e. per worker process) for mamba/linear-attn models.
if mambaish_config(model_runner.model_config) is not None:
from sglang.kernels.ops.mamba.triton_ops import (
initialize_mamba_selective_state_update_backend,
)
initialize_mamba_selective_state_update_backend(server_args)
model_runner.init_cuda_graphs()
if get_model().is_startup_weight_load_overlap:
model_runner.finalize_startup_weight_load()
python/sglang/kernels/ops/attention/minimax_sparse/decode/flash_with_topk_idx.py
+7 -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_xpu
from ..common.utils import (
_bitonic_merge,
@@ -1039,7 +1040,12 @@ def flash_decode_with_topk_idx(
# Equivalent output to the 2-stage path (set of block ids, front-packed,
# -1 padded); ~2-16x faster for long context. See
# sglang/kernels/ops/attention/minimax_decode_topk.py.
from sglang.kernels.ops.attention.minimax_decode_topk import minimax_decode_topk
if is_xpu():
from sgl_kernel import minimax_decode_topk
else:
from sglang.kernels.ops.attention.minimax_decode_topk import (
minimax_decode_topk,
)
minimax_decode_topk(score, seq_lens, block_size, topk, out=topk_idx)
else:
python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_flash_with_topk_idx.py
+6 -1
View File
@@ -7,8 +7,9 @@ from sglang.kernels.ops.attention.minimax_sparse.decode.flash_with_topk_idx impo
flash_decode_with_topk_idx,
)
from sglang.srt.environ import envs
from sglang.srt.utils import get_device, is_xpu
DEVICE = "cuda"
DEVICE = get_device()
RTOL_VS_REF = 5e-3
ATOL_VS_REF = 5e-3
@@ -416,6 +417,10 @@ def test_flash_decode_jit_topk_trivial_rows_skip_score_writes():
assert (topk_new[h, b, actual_k:] == -1).all()
@pytest.mark.skipif(
is_xpu(),
reason="XPU does not support trtllm_mha/fa3 dense backend",
)
def test_flash_decode_dense_page_table_trivial_rows_skip_score_writes():
torch.manual_seed(321)
bs, nqh, nkh, hd, blk, tk, page_size = 3, 4, 1, 128, 64, 32, 1
python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_sparse_gqa.py
+2 -1
View File
@@ -13,8 +13,9 @@ import torch
from sglang.kernels.ops.attention.minimax_sparse.decode.topk_sparse import (
flash_decode_with_gqa_share_sparse,
)
from sglang.srt.utils import get_device
DEVICE = "cuda"
DEVICE = get_device()
RTOL = 5e-3
ATOL = 5e-3
python/sglang/srt/layers/attention/xpu_backend.py
+3 -4
View File
@@ -949,14 +949,13 @@ class XPUAttentionBackend(AttentionBackend):
layer.v_scale,
)
else:
k_rope_val = (
k_rope if k_rope is not None else k[:, :, layer.v_head_dim :]
)
# Pass k_rope as-is like forward_extend: when rope is folded into
# k (k_rope is None), set_mla_kv_buffer stores the whole kv row.
self.token_to_kv_pool.set_mla_kv_buffer(
layer,
cache_loc,
k,
k_rope_val,
k_rope,
)
# Use precomputed metadata across all layers
python/sglang/srt/layers/quantization/fp8.py
+3 -1
View File
@@ -1407,7 +1407,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):
if is_fp4_expert:
fp4_block_k = 32
if fp4_scale_dtype is None:
fp4_scale_dtype = torch.float8_e8m0fnu if _use_aiter else torch.float32
fp4_scale_dtype = (
torch.float8_e8m0fnu if _use_aiter or is_xpu() else torch.float32
)
w13_weight_scale = torch.nn.Parameter(
torch.ones(
num_experts,
python/sglang/srt/layers/quantization/marlin_utils.py
+2 -2
View File
@@ -72,7 +72,7 @@ def query_marlin_supported_quant_types(
):
if device_capability is None:
major, minor = get_device_capability()
capability = major * 10 + minor
capability = major * 10 + minor if major is not None else None
device_capability = -1 if capability is None else capability
if device_capability < 80:
@@ -110,7 +110,7 @@ def _check_marlin_supported(
if device_capability is None:
major, minor = get_device_capability()
capability = major * 10 + minor
capability = major * 10 + minor if major is not None else None
device_capability = -1 if capability is None else capability
supported_types = query_marlin_supported_quant_types(
python/sglang/srt/layers/quantization/mxfp4.py
+11 -14
View File
@@ -1067,21 +1067,18 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
return
elif _is_xpu:
# sgl-kernel-xpu's W4A16 grouped GEMM consumes the checkpoint MXFP4
# layout: packed e2m1 [E, N, K/2] plus N-outer ue8m0 scales
# [E, N, K/32] uint8, with GPT-OSS's interleaved
# layout as-is: packed e2m1 [E, N, K/2] uint8 plus N-outer ue8m0
# scales [E, N, K/32] uint8, with GPT-OSS's interleaved
# [gate_0, up_0, gate_1, up_1, ...] w13 row order (which is exactly
# what the swiglu epilogue expects). Scales and biases are already in
# the expected dtypes (uint8 / bf16 -- the launcher promotes bias to
# fp32 since the kernel accumulates it in fp32), so the only step is
# reinterpreting the packed nibbles as int8, matching the dtype the
# kernel keys the 4-bit path on. That is a free view, and crucially
# there is no bf16 upcast -- the whole point of MXFP4 on XPU.
layer.w13_weight = Parameter(
layer.w13_weight.data.view(torch.int8), requires_grad=False
)
layer.w2_weight = Parameter(
layer.w2_weight.data.view(torch.int8), requires_grad=False
)
# what the swiglu epilogue expects). A packed byte holds two e2m1
# nibbles rather than an integer, so the torch dtype is only a
# container label: the op accepts int8 or uint8, always casts to
# uint8_t*, and decodes each nibble (sign bit included) as
# float_e2m1_t -- a path selected by the explicit
# use_mxfp4_w4a16=True that apply() passes, not by the weight dtype.
# Biases stay bf16 (the launcher promotes them to fp32, which is how
# the kernel accumulates them). Crucially there is no bf16 upcast of
# the weights -- the whole point of MXFP4 on XPU.
return
else:
from triton_kernels.numerics_details.mxfp import upcast_from_mxfp
python/sglang/srt/mem_cache/memory_pool.py
+13
View File
@@ -107,6 +107,7 @@ GB = 1024 * 1024 * 1024
_is_cuda = is_cuda()
_is_npu = is_npu()
_is_cpu = is_cpu()
_is_xpu = is_xpu()
_cpu_has_amx_support = cpu_has_amx_support()
_is_hip = is_hip()
_is_fp8_fnuz = is_fp8_fnuz()
@@ -116,6 +117,9 @@ _is_fp8_fnuz = is_fp8_fnuz()
# silently ignored and the legacy NHD layout is used.
_use_aiter = bool(envs.SGLANG_USE_AITER.get()) and _is_hip
if _is_xpu:
from sgl_kernel import store_cache_xpu
def conv_window_dedup_enabled(
is_npu: bool, is_cpu: bool, speculative_eagle_topk: Optional[int], is_kda: bool
@@ -169,6 +173,15 @@ def _set_kv_buffer_impl(
size_limit=size_limit,
)
if _is_xpu and v_row_dim == row_dim:
return store_cache_xpu(
k.view(-1, row_dim),
v.view(-1, row_dim),
k_cache.view(-1, row_dim),
v_cache.view(-1, row_dim),
indices,
)
# store_cache_cpu takes a single row_dim for both K and V, so it only serves
# equal-width rows; asymmetric KV falls through to the naive path below.
if _is_cpu and _cpu_has_amx_support and v_row_dim == row_dim:
test/registered/moe/test_hash_topk.py
+1
View File
@@ -24,6 +24,7 @@ def _set_dummy_server_args():
def test_hash_topk_remaps_per_rank_fused_shared_slots(monkeypatch):
monkeypatch.setattr(hash_topk_module, "_is_xpu", False)
monkeypatch.setattr(
hash_topk_module, "has_per_rank_fused_shared_slots", lambda *_args: True
)
test/registered/xpu/test_store_cache_xpu.py
+251
View File
@@ -0,0 +1,251 @@
"""
Verifies the fused ``store_cache_xpu`` KV-cache write path on Intel XPU.
This branch wires the fused SYCL ``store_cache_xpu`` kernel (from
sgl-kernel-xpu) into sglang's KV-cache writer
``sglang.srt.mem_cache.memory_pool._set_kv_buffer_impl``. On XPU that
dispatch replaces 2x ``index_put`` with a single kernel launch.
The tests exercise sglang's own dispatch (not the kernel in isolation), so
they fail if the wiring regresses to the ``index_put`` fallback:
- ``test_parity_*`` : fused write matches an ``index_put`` reference.
- ``test_dispatches_*`` : the fused kernel is actually the path taken.
- ``test_single_token`` : the common decode (1 token) case.
- ``test_parity_strided_*`` : non-contiguous K/V (a per-head slice of a
wider ``[tokens, heads, dim]`` tensor) writes
correctly. The fused kernel addresses source
rows by their real stride, so the SWA-layer
layout used by Gemma-style models is handled
without a host-side ``.contiguous()`` copy.
- ``test_dispatches_strided``: the strided write still takes the fused path.
Run from test/registered::
python3 -m unittest xpu.test_store_cache_xpu
Requires Intel XPU. ``store_cache_xpu`` is exported by the sgl-kernel-xpu
wheel pinned in ``python/pyproject_xpu.toml``, so it is not optional: a missing
op is a broken install and must fail loudly rather than skip.
"""
from __future__ import annotations
import unittest
import torch
from sglang.srt.utils import is_xpu
from sglang.test.ci.ci_register import register_xpu_ci
from sglang.test.test_utils import CustomTestCase
# Pure unit test (no server); fast and runs on the 1-GPU XPU runner.
register_xpu_ci(est_time=60, suite="stage-b-test-1-gpu-xpu")
def _reference_store(k, v, k_cache, v_cache, indices):
"""Naive index_put write — the path the fused kernel replaces."""
k_cache[indices] = k
v_cache[indices] = v
@unittest.skipUnless(is_xpu(), "Intel XPU not available")
class TestStoreCacheXPU(CustomTestCase):
"""store_cache_xpu, exercised through sglang's _set_kv_buffer_impl."""
def _store(self, k, v, k_cache, v_cache, indices):
"""Invoke sglang's KV-cache writer (the integration point)."""
from sglang.srt.mem_cache.memory_pool import _set_kv_buffer_impl
row_dim = k.shape[-1]
cache_size = k_cache.shape[0]
_set_kv_buffer_impl(
k,
v,
k_cache,
v_cache,
indices,
row_dim,
k.dtype,
torch.xpu,
size_limit=cache_size,
alt_stream=None,
)
torch.xpu.synchronize()
def _assert_parity(self, num_tokens, row_dim, dtype):
torch.manual_seed(42)
cache_size = 2048
k = torch.randn(num_tokens, row_dim, dtype=dtype, device="xpu")
v = torch.randn(num_tokens, row_dim, dtype=dtype, device="xpu")
indices = torch.randperm(cache_size, device="xpu")[:num_tokens].to(torch.int64)
k_ref = torch.zeros(cache_size, row_dim, dtype=dtype, device="xpu")
v_ref = torch.zeros_like(k_ref)
k_test = torch.zeros_like(k_ref)
v_test = torch.zeros_like(k_ref)
_reference_store(k, v, k_ref, v_ref, indices)
self._store(k, v, k_test, v_test, indices)
torch.testing.assert_close(k_test, k_ref)
torch.testing.assert_close(v_test, v_ref)
@staticmethod
def _strided_head_slice(num_tokens, num_heads, row_dim, head, dtype):
"""A non-contiguous per-head K/V slice of a wider tensor.
``[num_tokens, num_heads, row_dim][:, head, :]`` has shape
``(num_tokens, row_dim)`` but row stride ``num_heads * row_dim`` (not
``row_dim``) — the SWA-layer layout Gemma-style models hand to the
KV-cache writer. The fused kernel must address rows by this real
stride; a naive ``.view``/contiguous assumption would corrupt or copy.
"""
kw = torch.randn(num_tokens, num_heads, row_dim, dtype=dtype, device="xpu")
vw = torch.randn(num_tokens, num_heads, row_dim, dtype=dtype, device="xpu")
k = kw[:, head, :]
v = vw[:, head, :]
assert not k.is_contiguous()
assert k.stride() == (num_heads * row_dim, 1)
return k, v
def _assert_parity_strided(self, num_tokens, row_dim, num_heads, head, dtype):
torch.manual_seed(123)
cache_size = 2048
k, v = self._strided_head_slice(num_tokens, num_heads, row_dim, head, dtype)
indices = torch.randperm(cache_size, device="xpu")[:num_tokens].to(torch.int64)
k_ref = torch.zeros(cache_size, row_dim, dtype=dtype, device="xpu")
v_ref = torch.zeros_like(k_ref)
k_test = torch.zeros_like(k_ref)
v_test = torch.zeros_like(k_ref)
_reference_store(k, v, k_ref, v_ref, indices)
self._store(k, v, k_test, v_test, indices)
torch.testing.assert_close(k_test, k_ref)
torch.testing.assert_close(v_test, v_ref)
def test_parity_shapes(self):
"""Fused write matches index_put across token counts and row dims."""
for num_tokens in (1, 4, 32, 128):
for row_dim in (128, 256, 512, 1024):
with self.subTest(num_tokens=num_tokens, row_dim=row_dim):
self._assert_parity(num_tokens, row_dim, torch.bfloat16)
def test_parity_dtypes(self):
"""Both KV-cache dtypes write correctly (contiguous K/V)."""
for dtype in (torch.bfloat16, torch.float16):
with self.subTest(dtype=dtype):
self._assert_parity(32, 256, dtype)
def test_parity_strided_shapes(self):
"""Non-contiguous K/V (per-head slice) matches index_put across shapes.
Covers a few head counts / slice positions / token counts so the
kernel's row-stride addressing is exercised for both the odd
(non-vectorizable) and aligned (16-byte OWord) row-base cases.
"""
# num_tokens > 1: a single-row slice is trivially contiguous, so it
# would not exercise the inter-row stride this test targets.
for num_heads in (2, 10):
for head in (0, num_heads - 1):
for num_tokens in (2, 33, 271):
with self.subTest(
num_heads=num_heads, head=head, num_tokens=num_tokens
):
self._assert_parity_strided(
num_tokens, 256, num_heads, head, torch.bfloat16
)
def test_parity_strided_dtypes(self):
"""Both KV-cache dtypes write correctly for non-contiguous K/V."""
for dtype in (torch.bfloat16, torch.float16):
with self.subTest(dtype=dtype):
self._assert_parity_strided(271, 256, 10, 1, dtype)
def test_single_token(self):
"""Single-token decode (the most common runtime case)."""
torch.manual_seed(0)
row_dim, cache_size = 512, 4096
k = torch.randn(1, row_dim, dtype=torch.bfloat16, device="xpu")
v = torch.randn(1, row_dim, dtype=torch.bfloat16, device="xpu")
indices = torch.tensor([42], dtype=torch.int64, device="xpu")
k_cache = torch.zeros(cache_size, row_dim, dtype=torch.bfloat16, device="xpu")
v_cache = torch.zeros_like(k_cache)
self._store(k, v, k_cache, v_cache, indices)
torch.testing.assert_close(k_cache[42], k[0])
torch.testing.assert_close(v_cache[42], v[0])
def _count_fused_calls(self, k, v, indices, cache_size, row_dim):
"""Run a store through sglang and return how many times the fused
``store_cache_xpu`` kernel was actually invoked."""
from sglang.srt.mem_cache import memory_pool
calls = {"n": 0}
# memory_pool imports the symbol at module level, so patch the binding
# it actually calls, not sgl_kernel's attribute.
original = memory_pool.store_cache_xpu
def counting_store(*args, **kwargs):
calls["n"] += 1
return original(*args, **kwargs)
memory_pool.store_cache_xpu = counting_store
try:
k_cache = torch.zeros(cache_size, row_dim, dtype=k.dtype, device="xpu")
v_cache = torch.zeros_like(k_cache)
self._store(k, v, k_cache, v_cache, indices)
finally:
memory_pool.store_cache_xpu = original
return calls["n"]
def test_dispatches_to_fused_kernel(self):
"""sglang must take the fused path on XPU, not the index_put fallback.
Wrap the kernel and assert it is invoked exactly once. Guards against
the dispatch silently regressing (e.g. if can_use_store_cache starts
gating XPU again, which can't JIT-compile the CUDA kernel).
"""
torch.manual_seed(7)
row_dim, cache_size, num_tokens = 256, 1024, 8
k = torch.randn(num_tokens, row_dim, dtype=torch.bfloat16, device="xpu")
v = torch.randn(num_tokens, row_dim, dtype=torch.bfloat16, device="xpu")
indices = torch.randperm(cache_size, device="xpu")[:num_tokens].to(torch.int64)
n = self._count_fused_calls(k, v, indices, cache_size, row_dim)
self.assertEqual(
n,
1,
"expected _set_kv_buffer_impl to call the fused store_cache_xpu "
"exactly once on XPU; it likely fell back to index_put",
)
def test_dispatches_to_fused_kernel_strided(self):
"""The fused path must also be taken for non-contiguous (per-head
slice) K/V — sglang must not silently fall back to index_put just
because the source rows are strided."""
torch.manual_seed(8)
row_dim, cache_size, num_tokens, num_heads = 256, 1024, 8, 10
k, v = self._strided_head_slice(
num_tokens, num_heads, row_dim, 1, torch.bfloat16
)
indices = torch.randperm(cache_size, device="xpu")[:num_tokens].to(torch.int64)
n = self._count_fused_calls(k, v, indices, cache_size, row_dim)
self.assertEqual(
n,
1,
"expected _set_kv_buffer_impl to call the fused store_cache_xpu "
"exactly once for strided K/V; it likely fell back to index_put",
)
if __name__ == "__main__":
unittest.main()