From da7f890788071fb849b7e894069d3147e8e9fdf4 Mon Sep 17 00:00:00 2001 From: Polisetty V R K Jyothendra Varma Date: Fri, 1 May 2026 04:51:28 +0530 Subject: [PATCH] [Intel GPU] Integrate flash_mla_decode in Intel XPU attention backend (#23557) Signed-off-by: P V R K Jyothendra Varma Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Kangyan-Zhou Co-authored-by: Ma Mingfei --- .../srt/layers/attention/xpu_backend.py | 93 +++++++------------ .../sglang/srt/model_executor/model_runner.py | 1 + .../attention_backend_handler.py | 5 + .../srt/models/deepseek_common/utils.py | 1 + python/sglang/srt/server_args.py | 19 +++- test/srt/xpu/test_intel_xpu_backend.py | 15 ++- 6 files changed, 68 insertions(+), 66 deletions(-) diff --git a/python/sglang/srt/layers/attention/xpu_backend.py b/python/sglang/srt/layers/attention/xpu_backend.py index 4a40d25ee..0e0742938 100644 --- a/python/sglang/srt/layers/attention/xpu_backend.py +++ b/python/sglang/srt/layers/attention/xpu_backend.py @@ -14,12 +14,13 @@ from sglang.srt.layers.attention.flashattention_backend import ( ) from sglang.srt.managers.schedule_batch import get_global_server_args from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.utils import get_device_core_count if TYPE_CHECKING: from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.model_executor.model_runner import ModelRunner -from sgl_kernel import merge_state_v2 +from sgl_kernel import flash_mla_decode, flash_mla_get_workspace_size, merge_state_v2 from sgl_kernel.flash_attn import flash_attn_varlen_func, flash_attn_with_kvcache @@ -30,7 +31,7 @@ class XPUAttentionBackend(AttentionBackend): - Prefill and Decode disaggregation, currently only chunked prefill is supported - Speculative Decoding support - XPU Graph support, see https://github.com/pytorch/pytorch/issues/162143 - - MLA support + - MLA Prefill support """ def __init__( @@ -60,9 +61,6 @@ class XPUAttentionBackend(AttentionBackend): self.kv_cache_dtype_str = model_runner.server_args.kv_cache_dtype self.page_size = model_runner.page_size self.use_mla = model_runner.model_config.attention_arch == AttentionArch.MLA - assert ( - self.use_mla is False - ), "XPUAttentionBackend doesn't support MLA yet, please use --attention-backend triton instead." self.skip_prefill = skip_prefill self.is_hybrid_swa = model_runner.is_hybrid_swa if self.is_hybrid_swa: @@ -366,6 +364,21 @@ class XPUAttentionBackend(AttentionBackend): ), ] + if self.use_mla: + workspace_size = flash_mla_get_workspace_size( + self.max_context_len, + batch_size, + sm_count=get_device_core_count(), + num_kv_splits=-1, + ) + if ( + not hasattr(self, "workspace") + or self.workspace.numel() < workspace_size + ): + self.workspace = torch.empty( + workspace_size, device=self.device, dtype=torch.uint8 + ) + # Convert the page table to a strided format which is needed by FA3 API if self.page_size > 1: self.strided_indices = torch.arange( @@ -695,11 +708,14 @@ class XPUAttentionBackend(AttentionBackend): layer, cache_loc, k, v, layer.k_scale, layer.v_scale ) else: + k_rope_val = ( + k_rope if k_rope is not None else k[:, :, layer.v_head_dim :] + ) forward_batch.token_to_kv_pool.set_mla_kv_buffer( layer, cache_loc, k, - k_rope, + k_rope_val, ) # Use precomputed metadata across all layers @@ -857,17 +873,7 @@ class XPUAttentionBackend(AttentionBackend): kv_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id).to( q.dtype ) - k_rope = kv_cache[:, :, layer.v_head_dim :] - c_kv = kv_cache[:, :, : layer.v_head_dim] - k_rope_cache = k_rope.view( - -1, - self.page_size, - layer.tp_k_head_num, - layer.head_dim - layer.v_head_dim, - ) - c_kv_cache = c_kv.view( - -1, self.page_size, layer.tp_v_head_num, layer.v_head_dim - ) + assert not use_cascade_attn, "Cascade attention is not supported with MLA" if q_rope is not None: q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) @@ -878,53 +884,16 @@ class XPUAttentionBackend(AttentionBackend): q_all = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim) q_nope = q_all[:, :, : layer.v_head_dim] q_rope = q_all[:, :, layer.v_head_dim :] - max_seqlen_q = metadata.max_seq_len_q - result = flash_attn_with_kvcache( - q=q_rope, - k_cache=k_rope_cache, - v_cache=c_kv_cache, - qv=q_nope, - page_table=metadata.page_table, - cache_seqlens=metadata.cache_seqlens_int32, - cu_seqlens_q=metadata.cu_seqlens_q, - cu_seqlens_k_new=metadata.cu_seqlens_k, - max_seqlen_q=max_seqlen_q, - softmax_scale=layer.scaling, - causal=False if use_cascade_attn else causal, - softcap=layer.logit_cap, - k_descale=k_descale, - v_descale=v_descale, - return_softmax_lse=use_cascade_attn, # softmax_lse is needed for merge states + o = flash_mla_decode( + q_nope, + q_rope, + kv_cache.view(-1, self.page_size, layer.head_dim), + metadata.cache_seqlens_int32, + metadata.page_table, + self.workspace, + layer.scaling, ) - if use_cascade_attn: - o, softmax_lse, *rest = result - o_expand, softmax_lse_expand, *rest_expand = flash_attn_with_kvcache( - q=q_rope, - k_cache=k_rope_cache, - v_cache=c_kv_cache, - qv=q_nope, - page_table=self.forward_metadata_spec_decode_expand.page_table, - cache_seqlens=self.forward_metadata_spec_decode_expand.cache_seqlens_int32, - cu_seqlens_q=self.forward_metadata_spec_decode_expand.cu_seqlens_q, - cu_seqlens_k_new=self.forward_metadata_spec_decode_expand.cu_seqlens_k, - max_seqlen_q=self.forward_metadata_spec_decode_expand.max_seq_len_q, - softmax_scale=layer.scaling, - causal=False, - window_size=window_size, - softcap=layer.logit_cap, - k_descale=k_descale, - v_descale=v_descale, - return_softmax_lse=True, - ) - o, _ = merge_state_v2( - o, - softmax_lse.T.contiguous(), - o_expand, - softmax_lse_expand.T.contiguous(), - ) - else: - o = result return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 57486df63..005754b71 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -233,6 +233,7 @@ MLA_ATTENTION_BACKENDS = [ "trtllm_mla", "ascend", "nsa", + "intel_xpu", ] CHUNKED_PREFIX_CACHE_SUPPORTED_ATTENTION_BACKENDS = [ diff --git a/python/sglang/srt/models/deepseek_common/attention_backend_handler.py b/python/sglang/srt/models/deepseek_common/attention_backend_handler.py index dfc1d4c97..c1cd0e32c 100644 --- a/python/sglang/srt/models/deepseek_common/attention_backend_handler.py +++ b/python/sglang/srt/models/deepseek_common/attention_backend_handler.py @@ -172,6 +172,10 @@ def handle_attention_triton(attn, forward_batch): return _dispatch_mla_subtype(attn, forward_batch) +def handle_attention_intel_xpu(attn, forward_batch): + return _handle_attention_backend(attn, forward_batch, "intel_xpu") + + AttentionBackendRegistry.register("ascend", handle_attention_ascend) AttentionBackendRegistry.register("flashinfer", handle_attention_flashinfer) AttentionBackendRegistry.register("fa3", handle_attention_fa3) @@ -182,3 +186,4 @@ AttentionBackendRegistry.register("trtllm_mla", handle_attention_trtllm_mla) AttentionBackendRegistry.register("aiter", handle_attention_aiter) AttentionBackendRegistry.register("nsa", handle_attention_nsa) AttentionBackendRegistry.register("triton", handle_attention_triton) +AttentionBackendRegistry.register("intel_xpu", handle_attention_intel_xpu) diff --git a/python/sglang/srt/models/deepseek_common/utils.py b/python/sglang/srt/models/deepseek_common/utils.py index 59f0d8970..c8ac58c60 100644 --- a/python/sglang/srt/models/deepseek_common/utils.py +++ b/python/sglang/srt/models/deepseek_common/utils.py @@ -62,6 +62,7 @@ FORWARD_ABSORB_CORE_ATTENTION_BACKENDS = [ "cutlass_mla", "trtllm_mla", "ascend", + "intel_xpu", ] diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index aec174ee1..29294522f 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2683,10 +2683,23 @@ class ServerArgs: ) self.attention_backend = "triton" - if self.attention_backend == "intel_xpu": - if self.page_size not in [32, 64, 128]: + prefill_backend, decode_backend = self.get_attention_backends() + if self.use_mla_backend() and prefill_backend == "intel_xpu": + raise ValueError( + "intel_xpu backend is only supported on decode for MLA models, please set --decode-attention-backend to intel_xpu and do not set --attention-backend or --prefill-attention-backend to intel_xpu for prefill instead use triton." + ) + + if decode_backend == "intel_xpu": + if self.use_mla_backend(): + supported_page_sizes = [16, 32, 64, 128] + msg = "Intel XPU attention backend for MLA Decode" + else: + supported_page_sizes = [64, 128] + msg = "Intel XPU attention backend" + + if self.page_size not in supported_page_sizes: logger.warning( - f"Intel XPU attention backend only supports page_size of 32, 64 or 128, changing page_size from {self.page_size} to 128." + f"{msg} only supports page_sizes of {supported_page_sizes}, changing page_size from {self.page_size} to 128." ) self.page_size = 128 diff --git a/test/srt/xpu/test_intel_xpu_backend.py b/test/srt/xpu/test_intel_xpu_backend.py index c752b34dd..568ecc47d 100644 --- a/test/srt/xpu/test_intel_xpu_backend.py +++ b/test/srt/xpu/test_intel_xpu_backend.py @@ -7,6 +7,7 @@ import unittest from functools import wraps from sglang.test.test_utils import ( + DEFAULT_MODEL_NAME_FOR_TEST_FP8_WITH_MOE, DEFAULT_SMALL_MODEL_NAME_FOR_TEST_BASE, DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN, CustomTestCase, @@ -22,7 +23,7 @@ def intel_xpu_benchmark( @wraps(test_func) def wrapper(self): common_args = [ - "--disable-radix", + "--disable-radix-cache", "--trust-remote-code", "--mem-fraction-static", str(mem_fraction_static), @@ -65,6 +66,18 @@ class TestIntelXPUBackend(CustomTestCase): def test_attention_backend(self): return DEFAULT_SMALL_MODEL_NAME_FOR_TEST_BASE + @intel_xpu_benchmark( + [ + "--json-model-override-args", + '{"num_hidden_layers": 4}', + "--decode-attention-backend", + "intel_xpu", + ], + min_throughput=32, + ) + def test_mla_decode_attention_backend(self): + return DEFAULT_MODEL_NAME_FOR_TEST_FP8_WITH_MOE + if __name__ == "__main__": unittest.main()