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

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,

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