diff --git a/python/sglang/srt/hardware_backend/xpu/attention/__init__.py b/python/sglang/srt/hardware_backend/xpu/attention/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/sglang/srt/hardware_backend/xpu/attention/xpu_gdn_backend.py b/python/sglang/srt/hardware_backend/xpu/attention/xpu_gdn_backend.py new file mode 100644 index 000000000..84ffec9e3 --- /dev/null +++ b/python/sglang/srt/hardware_backend/xpu/attention/xpu_gdn_backend.py @@ -0,0 +1,138 @@ +import torch +from sgl_kernel import gdn_attention as sgl_kernel_gdn_attention + +from sglang.srt.layers.attention.linear.gdn_backend import GDNAttnBackend +from sglang.srt.layers.radix_linear_attention import RadixLinearAttention +from sglang.srt.model_executor.forward_batch_info import ForwardBatch + + +class XpuGDNAttnBackend(GDNAttnBackend): + """XPU specialization of ``GDNAttnBackend``. + + Adds an optional fused path that dispatches the whole conv1d + gating + + delta-rule pipeline to the vendored vLLM SYCL kernel exposed as + ``torch.ops.sgl_kernel.gdn_attention``. This is opt-in via + ``--linear-attn-backend intel_xpu`` (the default remains ``triton``, same + as other platforms). + """ + + def supports_fused_gdn(self, layer, forward_batch: ForwardBatch) -> bool: + """Conservative guard: only the plain decode / non-prefix-cached, + non-speculative extend cases are handled by the fused kernel.""" + mode = forward_batch.forward_mode + backends = self.linear_attn_backends + selected = ( + backends.verify + if mode.is_target_verify() + else (backends.decode if mode.is_decode_or_idle() else backends.prefill) + ) + if not selected.is_intel_xpu(): + return False + if not hasattr(torch.ops.sgl_kernel, "gdn_attention"): + # User explicitly asked for intel_xpu but the op isn't built. + raise RuntimeError( + "--linear-attn-backend intel_xpu requires the " + "torch.ops.sgl_kernel.gdn_attention op, but it is not " + "available. Rebuild sgl-kernel-xpu or use " + "--linear-attn-backend triton." + ) + if mode.is_target_verify() or mode.is_draft_extend_v2(): + return False + fm = self.forward_metadata + if getattr(fm, "has_mamba_track_mask", False): + # chunked prefix-cache intermediate-state tracking unsupported + return False + if getattr(fm, "query_start_loc", None) is None: + return False + # GDN (not KDA) shared weights must be plain tensors + if not isinstance(layer.conv_weights, torch.Tensor): + return False + if layer.bias is not None and not isinstance(layer.bias, torch.Tensor): + return False + return True + + def forward_fused_gdn( + self, + layer: RadixLinearAttention, + forward_batch: ForwardBatch, + projected_states_qkvz: torch.Tensor, + projected_states_ba: torch.Tensor, + ): + """Run the fused SYCL GDN op and return ``(core_attn_out, z)``. + + Caches stay in the SGLang pool layout and are updated in place. The conv + pool is ``[cache, dim, width-1]``; we pass a transposed view so the op sees + its logical ``[cache, width-1, dim]`` layout while the kernels index via + explicit width/dim strides (no gather/transpose/scatter copies). The ssm + pool already matches the op layout. ``mamba_cache_indices`` indexes the + full pool directly for both conv and ssm. + """ + fm = self.forward_metadata + layer_cache = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id) + conv_states = layer_cache.conv[0] # [cache, dim, width-1] + ssm_states = layer_cache.temporal # [cache, nv, hv, hk] + cache_indices = fm.mamba_cache_indices + query_start_loc = fm.query_start_loc + + device = projected_states_qkvz.device + dtype = projected_states_qkvz.dtype + bs = forward_batch.batch_size + num_actual_tokens = projected_states_qkvz.shape[0] + + if forward_batch.forward_mode.is_decode_or_idle(): + num_decodes, num_prefills = bs, 0 + has_initial_state = torch.ones(bs, dtype=torch.bool, device=device) + else: + num_decodes, num_prefills = 0, bs + has_initial_state = forward_batch.extend_prefix_lens > 0 + + # Full-pool, zero-copy: transposed view for conv + native ssm pool, indexed + # directly by the full-pool cache indices. + conv_view = conv_states.transpose(1, 2) # [cache, width-1, dim] view + state_idx = cache_indices.to(torch.int32).contiguous() + + core_attn_out = torch.empty( + num_actual_tokens, + layer.num_v_heads, + layer.head_v_dim, + dtype=dtype, + device=device, + ) + z = torch.empty_like(core_attn_out) + + sgl_kernel_gdn_attention( + core_attn_out=core_attn_out, + z=z, + projected_states_qkvz=projected_states_qkvz, + projected_states_ba=projected_states_ba, + num_k_heads=layer.num_k_heads, + num_v_heads=layer.num_v_heads, + head_k_dim=layer.head_k_dim, + head_v_dim=layer.head_v_dim, + conv_state=conv_view, + ssm_state=ssm_states, + conv_weights=layer.conv_weights, + conv_bias=layer.bias, + activation=layer.activation, + A_log=layer.A_log, + dt_bias=layer.dt_bias, + num_prefills=num_prefills, + num_decodes=num_decodes, + num_spec_decodes=0, + has_initial_state=has_initial_state, + non_spec_query_start_loc=query_start_loc, + non_spec_token_indx=None, + non_spec_state_indices_tensor=state_idx, + spec_query_start_loc=None, + spec_token_indx=None, + spec_state_indices_tensor=None, + num_accepted_tokens=None, + num_actual_tokens=num_actual_tokens, + # Heads/tensors are already per-rank sharded; kernel needs no + # further in-kernel sharding, so this is always 1, not --tp-size. + tp_size=1, + reorder_input=True, + ) + + # conv/ssm states were updated in place via the pool views. + return core_attn_out, z diff --git a/python/sglang/srt/layers/attention/attention_registry.py b/python/sglang/srt/layers/attention/attention_registry.py index 35e075a37..c560895f7 100644 --- a/python/sglang/srt/layers/attention/attention_registry.py +++ b/python/sglang/srt/layers/attention/attention_registry.py @@ -378,6 +378,7 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac is_blackwell, is_npu, is_sm120_supported, + is_xpu, ) if not is_npu(): @@ -389,6 +390,11 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac GDNAttnBackend, flashinfer_gdn_prefill_default, ) + + if is_xpu(): + from sglang.srt.hardware_backend.xpu.attention.xpu_gdn_backend import ( + XpuGDNAttnBackend as GDNAttnBackend, + ) else: from sglang.srt.hardware_backend.npu.attention.ascend_gdn_backend import ( AscendGDNAttnBackend as GDNAttnBackend, diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py index 304fe1596..7131bfffb 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py @@ -142,6 +142,13 @@ class GDNKernelDispatcher: cutedsl_kernel = None if decode_backend.is_triton(): self.decode_kernel = triton_kernel + elif decode_backend.is_intel_xpu(): + if not is_xpu(): + raise ValueError("--linear-attn-backend intel_xpu requires Intel XPU") + # The fused SYCL kernel is dispatched via XpuGDNAttnBackend.forward_fused_gdn, + # outside this dispatcher; Triton is the dispatcher-level kernel for requests + # that hook doesn't handle (e.g. verify). + self.decode_kernel = triton_kernel elif decode_backend.is_cutedsl(): if not is_cuda(): raise ValueError("GDN CuTe DSL backend requires CUDA") @@ -169,6 +176,12 @@ class GDNKernelDispatcher: if prefill_backend.is_triton(): self.extend_kernel = triton_kernel + elif prefill_backend.is_intel_xpu(): + if not is_xpu(): + raise ValueError("--linear-attn-backend intel_xpu requires Intel XPU") + # See the decode branch above: intel_xpu uses Triton as its + # dispatcher-level fallback kernel. + self.extend_kernel = triton_kernel elif prefill_backend.is_cutedsl(): if not is_cuda(): raise ValueError("GDN CuTe DSL backend requires CUDA") @@ -384,6 +397,7 @@ class GDNAttnBackend(MambaAttnBackendBase): ), f"{self.conv_states_shape[-1]=} should be less than {FLA_CHUNK_SIZE}" backends = model_runner.linear_attn_backends + self.linear_attn_backends = backends self.kernel_dispatcher = GDNKernelDispatcher( backends.decode, backends.prefill, backends.verify ) diff --git a/python/sglang/srt/layers/attention/linear/utils.py b/python/sglang/srt/layers/attention/linear/utils.py index 74323c216..4e63068ec 100644 --- a/python/sglang/srt/layers/attention/linear/utils.py +++ b/python/sglang/srt/layers/attention/linear/utils.py @@ -21,6 +21,7 @@ class LinearAttnKernelBackend(Enum): NVIDIA_KDA = "nvidia_kda" PTX_KDA = "ptx_kda" HELION = "helion" + INTEL_XPU = "intel_xpu" CUSTOM = "custom" @classmethod @@ -51,6 +52,9 @@ class LinearAttnKernelBackend(Enum): def is_helion(self): return self == LinearAttnKernelBackend.HELION + def is_intel_xpu(self): + return self == LinearAttnKernelBackend.INTEL_XPU + def is_custom(self): return self == LinearAttnKernelBackend.CUSTOM diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 3371ac0d3..f5771a957 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -136,6 +136,7 @@ _qknorm_use_alt_stream = _is_cuda or ( get_bool_env_var("SGLANG_QK_NORM_ALT_STREAM", "False") and _hip_use_alt_stream ) _is_amx_available = cpu_has_amx_support() +_is_xpu = is_xpu() # Head-group ratios (num_v_heads // num_k_heads) served by the fused # split/reshape/cat Triton kernel. On AMD/aiter the ratio-8 layout is also @@ -670,6 +671,40 @@ class Qwen3_5GatedDeltaNet(nn.Module): projected_states_ba, _ = self.in_proj_ba(hs_bf16) return projected_states_qkvz, projected_states_ba + def _forward_xpu( + self, + backend: object, + projected_states_qkvz: torch.Tensor, + projected_states_ba: torch.Tensor, + forward_batch: ForwardBatch, + ): + core_attn_out, z = backend.forward_fused_gdn( + self.attn, + forward_batch, + projected_states_qkvz, + projected_states_ba, + ) + + assert core_attn_out is not None, "XPU backend must support fused GDN" + + z_shape_og = z.shape + # reshape input data into 2D tensor + core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1]) + z = z.reshape(-1, z.shape[-1]) + + # Add padding for DP-Attn + if core_attn_out.shape != z.shape: + core_attn_out_pad = torch.zeros_like(z) + core_attn_out_pad[: core_attn_out.shape[0], :] = core_attn_out + core_attn_out = core_attn_out_pad + + core_attn_out = self.norm(core_attn_out, z) + core_attn_out = core_attn_out.reshape(z_shape_og) + core_attn_out = core_attn_out.reshape(*core_attn_out.shape[:-2], -1) + + output, _ = self.out_proj(core_attn_out) + return output + def forward( self, hidden_states: torch.Tensor, @@ -685,6 +720,16 @@ class Qwen3_5GatedDeltaNet(nn.Module): hidden_states ) + if _is_xpu and get_exec().mamba.linear_attn_backend == "intel_xpu": + from sglang.srt.model_executor.forward_context import get_attn_backend + + backend = get_attn_backend() + backend = getattr(backend, "linear_attn_backend", backend) + if backend.supports_fused_gdn(self.attn, forward_batch): + return self._forward_xpu( + backend, projected_states_qkvz, projected_states_ba, forward_batch + ) + if ( self.num_v_heads // self.num_k_heads in _GDN_FUSED_QKVZBA_RATIOS and not _is_npu diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index bd850b330..aeafa014e 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -404,6 +404,7 @@ LINEAR_ATTN_KERNEL_BACKEND_CHOICES = [ "nvidia_kda", "ptx_kda", "helion", + "intel_xpu", ] diff --git a/test/registered/xpu/llm_models/test_xpu_qwen3_5_9b.py b/test/registered/xpu/llm_models/test_xpu_qwen3_5_9b.py index aacbe320c..53ef54538 100644 --- a/test/registered/xpu/llm_models/test_xpu_qwen3_5_9b.py +++ b/test/registered/xpu/llm_models/test_xpu_qwen3_5_9b.py @@ -1,6 +1,8 @@ """Qwen3.5-9B GSM8K accuracy on Intel XPU (TP=4). -Scored by ``simple_eval_gsm8k.GSM8KEval``. +Scored by ``simple_eval_gsm8k.GSM8KEval``. Covers both the opt-in fused GDN +SYCL kernel path (``--linear-attn-backend intel_xpu``) and the default Triton +GDN path (``triton``, unchanged from other platforms). """ import unittest @@ -18,9 +20,9 @@ register_xpu_ci(est_time=2400, suite="nightly-xpu-4-gpu", nightly=True) torch.xpu.is_available(), "Intel XPU not available (torch.xpu.is_available() returned False)", ) -class TestQwen3_5_9BXPU(SimpleEvalGSM8KXPUMixin, CustomTestCase): +class Qwen3_5_9BXPUBase(SimpleEvalGSM8KXPUMixin, CustomTestCase): model = "Qwen/Qwen3.5-9B" - tp_size = 4 + tp_size = 1 accuracy = 0.90 # max_tokens=8192 lets the GSM8K CoT complete under num_threads=4. num_examples = 50 @@ -37,5 +39,23 @@ class TestQwen3_5_9BXPU(SimpleEvalGSM8KXPUMixin, CustomTestCase): ] +class TestQwen3_5_9BXPUDefault(Qwen3_5_9BXPUBase): + """Default path: Triton GDN kernels (unchanged from other platforms).""" + + +class TestQwen3_5_9BXPUFusedGDN(Qwen3_5_9BXPUBase): + """Opt-in fused SYCL GDN kernel path (``--linear-attn-backend intel_xpu``). + + Small ``num_examples`` since this is a smoke check of the fused-kernel + dispatch, not a full accuracy regression test (already covered by the + default Triton path above); accuracy threshold is left at 0 accordingly. + """ + + other_args = Qwen3_5_9BXPUBase.other_args + [ + "--linear-attn-backend", + "intel_xpu", + ] + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/xpu/test_intel_xpu_linear_attn_dispatch.py b/test/registered/xpu/test_intel_xpu_linear_attn_dispatch.py new file mode 100644 index 000000000..f9e6162b2 --- /dev/null +++ b/test/registered/xpu/test_intel_xpu_linear_attn_dispatch.py @@ -0,0 +1,76 @@ +"""Intel XPU linear-attn (GDN) backend dispatch: intel_xpu is opt-in only +(default stays triton, like every other platform) and fails fast rather than +silently degrading when misconfigured. Pure dispatch-logic tests -- no XPU +device required -- kept under test/registered/xpu to separate Intel-XPU-only +behavior from the platform-agnostic linear-attn dispatch tests. +""" + +import unittest +from unittest.mock import patch + +from sglang.srt.layers.attention.linear import gdn_backend +from sglang.srt.layers.attention.linear.gdn_backend import GDNKernelDispatcher +from sglang.srt.layers.attention.linear.kda_backend import KDAKernelDispatcher +from sglang.srt.layers.attention.linear.kernels.gdn_triton import TritonGDNKernel +from sglang.srt.layers.attention.linear.utils import LinearAttnKernelBackend +from sglang.srt.server_args import LINEAR_ATTN_KERNEL_BACKEND_CHOICES +from sglang.test.ci.ci_register import register_xpu_ci +from sglang.test.test_utils import CustomTestCase + +register_xpu_ci(est_time=5, suite="stage-b-test-1-gpu-xpu") + + +class TestIntelXpuGDNDispatch(CustomTestCase): + def test_intel_xpu_is_a_registered_backend_choice(self): + self.assertIn("intel_xpu", LINEAR_ATTN_KERNEL_BACKEND_CHOICES) + + def test_intel_xpu_requires_xpu_hardware(self): + with patch.object(gdn_backend, "is_xpu", return_value=False): + with self.assertRaisesRegex(ValueError, "requires Intel XPU"): + GDNKernelDispatcher( + LinearAttnKernelBackend.INTEL_XPU, + LinearAttnKernelBackend.TRITON, + ) + with self.assertRaisesRegex(ValueError, "requires Intel XPU"): + GDNKernelDispatcher( + LinearAttnKernelBackend.TRITON, + LinearAttnKernelBackend.INTEL_XPU, + ) + + def test_intel_xpu_uses_triton_as_the_dispatcher_fallback_kernel(self): + # The fused SYCL kernel is dispatched outside GDNKernelDispatcher (via + # XpuGDNAttnBackend.forward_fused_gdn); the dispatcher itself only + # needs a valid fallback kernel for requests that hook declines + # (e.g. verify), which is Triton. + with patch.object(gdn_backend, "is_xpu", return_value=True): + dispatcher = GDNKernelDispatcher( + LinearAttnKernelBackend.INTEL_XPU, + LinearAttnKernelBackend.INTEL_XPU, + ) + + self.assertIsInstance(dispatcher.decode_kernel, TritonGDNKernel) + self.assertIsInstance(dispatcher.extend_kernel, TritonGDNKernel) + self.assertIsInstance(dispatcher.verify_kernel, TritonGDNKernel) + + +class TestIntelXpuKDADispatch(CustomTestCase): + def test_intel_xpu_is_not_a_supported_kda_backend(self): + # Unlike GDN, KDA has no Intel XPU SYCL kernel: intel_xpu must not be + # silently treated as Triton, it should fail fast like any other + # backend KDA does not implement. + with self.assertRaisesRegex(ValueError, "Unsupported KDA decode backend"): + KDAKernelDispatcher( + decode_backend=LinearAttnKernelBackend.INTEL_XPU, + prefill_backend=LinearAttnKernelBackend.TRITON, + verify_backend=LinearAttnKernelBackend.TRITON, + ) + with self.assertRaisesRegex(ValueError, "Unsupported KDA prefill backend"): + KDAKernelDispatcher( + decode_backend=LinearAttnKernelBackend.TRITON, + prefill_backend=LinearAttnKernelBackend.INTEL_XPU, + verify_backend=LinearAttnKernelBackend.TRITON, + ) + + +if __name__ == "__main__": + unittest.main()