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

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>

python/sglang/benchmark/one_batch.py

@@ -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,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

@@ -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

@@ -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

@@ -949,14 +949,13 @@ class XPUAttentionBackend(AttentionBackend):
                         layer.v_scale,
                     )
                 else:
-                    k_rope_val = (
+                    # Pass k_rope as-is like forward_extend: when rope is folded into
-                        k_rope if k_rope is not None else k[:, :, layer.v_head_dim :]
+                    # 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

@@ -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

@@ -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

@@ -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
+            # layout as-is: packed e2m1 [E, N, K/2] uint8 plus N-outer ue8m0
-            # [E, N, K/32] uint8, with GPT-OSS's interleaved
+            # 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
+            # what the swiglu epilogue expects). A packed byte holds two e2m1
-            # the expected dtypes (uint8 / bf16 -- the launcher promotes bias to
+            # nibbles rather than an integer, so the torch dtype is only a
-            # fp32 since the kernel accumulates it in fp32), so the only step is
+            # container label: the op accepts int8 or uint8, always casts to
-            # reinterpreting the packed nibbles as int8, matching the dtype the
+            # uint8_t*, and decodes each nibble (sign bit included) as
-            # kernel keys the 4-bit path on. That is a free view, and crucially
+            # float_e2m1_t -- a path selected by the explicit
-            # there is no bf16 upcast -- the whole point of MXFP4 on XPU.
+            # use_mxfp4_w4a16=True that apply() passes, not by the weight dtype.
-            layer.w13_weight = Parameter(
+            # Biases stay bf16 (the launcher promotes them to fp32, which is how
-                layer.w13_weight.data.view(torch.int8), requires_grad=False
+            # the kernel accumulates them). Crucially there is no bf16 upcast of
-            )
+            # the weights -- the whole point of MXFP4 on XPU.
-            layer.w2_weight = Parameter(
-                layer.w2_weight.data.view(torch.int8), requires_grad=False
-            )
             return
         else:
             from triton_kernels.numerics_details.mxfp import upcast_from_mxfp

python/sglang/srt/mem_cache/memory_pool.py

@@ -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

@@ -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

@@ -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()