From 84d7604b7e6b89f647b4ea6fa2bcabb3d5ffe44f Mon Sep 17 00:00:00 2001 From: "Meng, Hengyu" Date: Wed, 16 Sep 2026 19:43:04 -0700 Subject: [PATCH] [XPU] weekly simple model enablement 2026/09/14 (#39439) Co-authored-by: Juan Muneton <102537701+jmunetong@users.noreply.github.com> Co-authored-by: YangKai0616 Co-authored-by: devan-carlin Co-authored-by: Ashwini Rathi Co-authored-by: Ranjan Debnath Co-authored-by: Juan Muneton Co-authored-by: Amrutha M --- python/sglang/benchmark/one_batch.py | 10 + .../decode/flash_with_topk_idx.py | 8 +- .../tests/test_flash_with_topk_idx.py | 7 +- .../tests/test_sparse_gqa.py | 3 +- .../srt/layers/attention/xpu_backend.py | 7 +- python/sglang/srt/layers/quantization/fp8.py | 4 +- .../srt/layers/quantization/marlin_utils.py | 4 +- .../sglang/srt/layers/quantization/mxfp4.py | 25 +- python/sglang/srt/mem_cache/memory_pool.py | 13 + test/registered/moe/test_hash_topk.py | 1 + test/registered/xpu/test_store_cache_xpu.py | 251 ++++++++++++++++++ 11 files changed, 309 insertions(+), 24 deletions(-) create mode 100644 test/registered/xpu/test_store_cache_xpu.py diff --git a/python/sglang/benchmark/one_batch.py b/python/sglang/benchmark/one_batch.py index 7bc8da18c..a4d7d0de3 100644 --- a/python/sglang/benchmark/one_batch.py +++ b/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() diff --git a/python/sglang/kernels/ops/attention/minimax_sparse/decode/flash_with_topk_idx.py b/python/sglang/kernels/ops/attention/minimax_sparse/decode/flash_with_topk_idx.py index 4e9b5b2a7..091b5ee09 100644 --- a/python/sglang/kernels/ops/attention/minimax_sparse/decode/flash_with_topk_idx.py +++ b/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: diff --git a/python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_flash_with_topk_idx.py b/python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_flash_with_topk_idx.py index 8f3b88bb4..81991455c 100644 --- a/python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_flash_with_topk_idx.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_sparse_gqa.py b/python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_sparse_gqa.py index 39adece33..a0c9b3db5 100644 --- a/python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_sparse_gqa.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/xpu_backend.py b/python/sglang/srt/layers/attention/xpu_backend.py index a689114ed..464fc9025 100644 --- a/python/sglang/srt/layers/attention/xpu_backend.py +++ b/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 diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 32af025f5..687bd96bb 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/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, diff --git a/python/sglang/srt/layers/quantization/marlin_utils.py b/python/sglang/srt/layers/quantization/marlin_utils.py index 9b432028d..5e5c7c896 100644 --- a/python/sglang/srt/layers/quantization/marlin_utils.py +++ b/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( diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index b94894ea7..98787d599 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 450b22e6f..b18e8964e 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/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: diff --git a/test/registered/moe/test_hash_topk.py b/test/registered/moe/test_hash_topk.py index 1c5658313..f3d0e32ca 100644 --- a/test/registered/moe/test_hash_topk.py +++ b/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 ) diff --git a/test/registered/xpu/test_store_cache_xpu.py b/test/registered/xpu/test_store_cache_xpu.py new file mode 100644 index 000000000..69082d66c --- /dev/null +++ b/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()