[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 = (
-                        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

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

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