From 5338e4448300c5bf131f77a9308b0d0757e8ce84 Mon Sep 17 00:00:00 2001
From: Polisetty V R K Jyothendra Varma
 <polisetty.v.r.k.jyothendra.varma@intel.com>
Date: Wed, 24 Jun 2026 06:36:13 +0530
Subject: [PATCH] [Intel GPU] fix triton-mla attention on XPU by limiting
 max_kv_splits to 8 which is default (#28646)

Signed-off-by: P V R K Jyothendra Varma <polisetty.v.r.k.jyothendra.varma@intel.com>
Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
---
 .../srt/layers/attention/triton_backend.py    |  4 +-
 .../xpu/test_triton_attention_backend.py      | 66 +++++++++++++++++++
 2 files changed, 69 insertions(+), 1 deletion(-)
 create mode 100644 test/registered/xpu/test_triton_attention_backend.py

diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py
index 6541cd68e..3f0bab194 100644
--- a/python/sglang/srt/layers/attention/triton_backend.py
+++ b/python/sglang/srt/layers/attention/triton_backend.py
@@ -40,11 +40,13 @@ from sglang.srt.utils import (
     is_cuda,
     is_gfx95_supported,
     is_gfx942_supported,
+    is_xpu,
     next_power_of_2,
 )
 
 _is_cuda = is_cuda()
 _is_gfx942 = is_gfx942_supported()
+_is_xpu = is_xpu()
 
 if _is_cuda:
     from sgl_kernel.utils import is_arch_support_pdl
@@ -202,7 +204,7 @@ class TritonAttnBackend(AttentionBackend):
             "SGLANG_TRITON_DECODE_ATTN_STATIC_KV_SPLITS", "false"
         )
         self.max_kv_splits = model_runner.server_args.triton_attention_num_kv_splits
-        if self.use_mla:
+        if self.use_mla and not _is_xpu:
             self.max_kv_splits = _mla_decode_kv_splits_cap(
                 self.max_kv_splits,
                 self.device_core_count,
diff --git a/test/registered/xpu/test_triton_attention_backend.py b/test/registered/xpu/test_triton_attention_backend.py
new file mode 100644
index 000000000..72e0a99e7
--- /dev/null
+++ b/test/registered/xpu/test_triton_attention_backend.py
@@ -0,0 +1,66 @@
+"""
+Usage:
+python3 -m unittest test_triton_attention_backend.TestTritonAttentionBackend.test_mla_triton_attention_backend
+"""
+
+import unittest
+from functools import wraps
+
+from sglang.test.ci.ci_register import register_xpu_ci
+from sglang.test.test_utils import (
+    DEFAULT_MODEL_NAME_FOR_TEST_FP8_WITH_MOE,
+    CustomTestCase,
+    run_bench_serving,
+)
+
+register_xpu_ci(est_time=600, suite="stage-b-test-1-gpu-xpu")
+
+
+def triton_attention_benchmark(extra_args=None, mem_fraction_static="0.84"):
+    def decorator(test_func):
+        @wraps(test_func)
+        def wrapper(self):
+            common_args = [
+                "--disable-radix-cache",
+                "--trust-remote-code",
+                "--tp-size",
+                "1",
+                "--mem-fraction-static",
+                str(mem_fraction_static),
+                "--context-length",
+                "2050",
+                "--attention-backend",
+                "triton",
+            ]
+            full_args = common_args + (extra_args or [])
+
+            model = test_func(self)
+            res = run_bench_serving(
+                model,
+                256,
+                float("inf"),
+                full_args,
+                random_input_len=1024,
+                random_output_len=1024,
+                need_warmup=False,
+            )
+
+        return wrapper
+
+    return decorator
+
+
+class TestTritonAttentionBackend(CustomTestCase):
+
+    @triton_attention_benchmark(
+        [
+            "--json-model-override-args",
+            '{"num_hidden_layers": 4}',
+        ],
+    )
+    def test_mla_triton_attention_backend(self):
+        return DEFAULT_MODEL_NAME_FOR_TEST_FP8_WITH_MOE
+
+
+if __name__ == "__main__":
+    unittest.main()