Speculative Decoding support for intel_xpu attention backend on XPU target (#30548)

python/sglang/srt/layers/attention/xpu_backend.py

@@ -141,9 +141,6 @@ class XPUAttentionBackend(AttentionBackend):
         if forward_batch.forward_mode.is_decode_or_idle():
             # Draft Decode
             if forward_batch.spec_info is not None:
-                assert False, (
-                    "XPUAttentionBackend doesn't support speculative decoding yet, please use --attention-backend triton instead."
-                )
                 if self.topk <= 1:
                     metadata.cache_seqlens_int32 = (
                         seqlens_in_batch + (self.speculative_step_id + 1)
@@ -355,7 +352,9 @@ class XPUAttentionBackend(AttentionBackend):
                         metadata, metadata_expand
                     )
 
-        elif forward_batch.forward_mode.is_extend_or_draft_extend_or_mixed():
+        elif forward_batch.forward_mode.is_extend_or_draft_extend_or_mixed(
+            include_draft_extend_v2=True
+        ):
             metadata.cache_seqlens_int32 = seqlens_in_batch.to(torch.int32)
             metadata.max_seq_len_k = forward_batch.seq_lens_cpu.max().item()
             metadata.cu_seqlens_k = torch.nn.functional.pad(
@@ -365,7 +364,16 @@ class XPUAttentionBackend(AttentionBackend):
                 forward_batch.req_pool_indices, : metadata.max_seq_len_k
             ]
 
-            if any(forward_batch.extend_prefix_lens_cpu):
+            # Detect draft extend: either explicit DRAFT_EXTEND_V2 mode or EXTEND
+            # mode with an EagleDraftInput spec_info.
+            is_draft_extend = (
+                forward_batch.forward_mode == ForwardMode.DRAFT_EXTEND_V2
+                or (
+                    forward_batch.spec_info is not None
+                    and forward_batch.spec_info.is_draft_input()
+                )
+            )
+            if any(forward_batch.extend_prefix_lens_cpu) or is_draft_extend:
                 extend_seq_lens = forward_batch.extend_seq_lens
                 metadata.max_seq_len_q = max(forward_batch.extend_seq_lens_cpu)
                 metadata.cu_seqlens_q = torch.nn.functional.pad(
@@ -609,21 +617,22 @@ class XPUAttentionBackend(AttentionBackend):
                 -1, self.page_size, layer.tp_k_head_num, layer.head_dim
             )
             value_cache = value_cache.view(
-                -1, self.page_size, layer.tp_v_head_num, layer.head_dim
+                -1, self.page_size, layer.tp_v_head_num, layer.v_head_dim
             )
             if self.is_encoder_decoder and forward_batch.encoder_lens is not None:
                 page_table, cache_seqlens, causal = self._encoder_decoder_page_table(
                     layer, metadata
                 )
-                o = self._forward_attn_flat_page_table(
+                o = self._forward_encoder_decoder_attn(
-                    q=q,
+                    q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
                     key_cache=key_cache,
                     value_cache=value_cache,
-                    layer=layer,
                     page_table=page_table,
                     cache_seqlens=cache_seqlens,
                     cu_seqlens_q=metadata.cu_seqlens_q,
                     max_seqlen_q=metadata.max_seq_len_q,
+                    scale=layer.scaling,
+                    softcap=layer.logit_cap,
                     causal=causal,
                 )
                 return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
@@ -649,47 +658,47 @@ class XPUAttentionBackend(AttentionBackend):
                 k_descale=k_descale,
                 v_descale=v_descale,
                 return_softmax_lse=use_cascade_attn,
-                # Piecewise XPU graph for prefill requires a pre-allocated
+                # `out` is injected via out_kwargs only on the graph path (buffer
-                # output buffer at a stable device address so the graph can
+                # present, non-cascade); the eager path omits it for flash_attn
-                # record writes to the same storage on every replay.
+                # builds that lack the kwarg. Piecewise XPU graph for prefill
-                # _attn_output is that fixed buffer; None falls back to a
+                # pre-allocates this fixed-address output buffer so graph replay
-                # freshly allocated tensor (eager / cascade-attn path).
+                # writes to the same storage; radix_attention sets it only on the
-                out=(
+                # graph path, None on the eager path.
-                    forward_batch._attn_output.view(
+                **(
+                    {
+                        "out": forward_batch._attn_output.view(
                             -1, layer.tp_q_head_num, layer.v_head_dim
                         )
-                    if not use_cascade_attn
+                    }
-                    and getattr(forward_batch, "_attn_output", None) is not None
+                    if not use_cascade_attn and forward_batch._attn_output is not None
-                    else None
+                    else {}
                 ),
                 **kwargs,
             )
 
             if use_cascade_attn:
                 o, softmax_lse, *rest = result
-                o_expand, softmax_lse_expand, *rest_expand = flash_attn_with_kvcache(
+                expand = self.forward_metadata_spec_decode_expand
+                o_expand, lse_expand, *_ = self._forward_attn_flat_page_table(
                     q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
-                    k_cache=key_cache,
+                    key_cache=key_cache,
-                    v_cache=value_cache,
+                    value_cache=value_cache,
-                    page_table=self.forward_metadata_spec_decode_expand.page_table,
+                    page_table=expand.page_table,
-                    cache_seqlens=self.forward_metadata_spec_decode_expand.cache_seqlens_int32,
+                    cache_seqlens=expand.cache_seqlens_int32,
-                    cu_seqlens_q=self.forward_metadata_spec_decode_expand.cu_seqlens_q,
+                    cu_seqlens_q=expand.cu_seqlens_q,
-                    cu_seqlens_k_new=None,
+                    max_seqlen_q=expand.max_seq_len_q,
-                    max_seqlen_q=self.forward_metadata_spec_decode_expand.max_seq_len_q,
+                    scale=layer.scaling,
-                    softmax_scale=layer.scaling,
-                    causal=False,
-                    window_size=window_size,
                     softcap=layer.logit_cap,
-                    k_descale=k_descale,
+                    window_size=window_size,
-                    v_descale=v_descale,
+                    sinks=sinks,
                     return_softmax_lse=True,
-                    **kwargs,
                 )
                 o, _ = merge_state_v2_wrapper(
                     o,
                     softmax_lse.T.contiguous(),
                     o_expand,
-                    softmax_lse_expand.T.contiguous(),
+                    # lse_expand comes back as (Hq, total_q); merge_state wants (total_q, Hq).
+                    lse_expand.T.contiguous(),
                 )
             else:
                 o = result
@@ -798,48 +807,98 @@ class XPUAttentionBackend(AttentionBackend):
 
     def _forward_attn_flat_page_table(
         self,
-        *,
+        q: torch.Tensor,
-        q,
+        key_cache: torch.Tensor,
-        key_cache,
+        value_cache: torch.Tensor,
-        value_cache,
+        page_table: torch.Tensor,
-        layer,
+        cache_seqlens: torch.Tensor,
-        page_table,
+        cu_seqlens_q: Optional[torch.Tensor],
-        cache_seqlens,
+        max_seqlen_q: int,
-        cu_seqlens_q,
+        scale: float,
-        max_seqlen_q,
+        softcap: float,
-        causal,
+        causal: bool = False,
+        window_size: tuple[int, int] = (-1, -1),
+        sinks: Optional[torch.Tensor] = None,
+        return_softmax_lse: bool = False,
     ):
-        """MHA on XPU via flash_attn_with_kvcache with a page_size=1 (flat
+        """Attention against a ``page_table`` that indexes individual token
-        token-slot) page table. sgl-kernel-xpu PR #454 detects a page_size==1
+        slots rather than ``self.page_size``-sized blocks. Serves cascade
-        k_cache + page_table and gathers + runs varlen internally, so the backend
+        tree-branch expand (target-verify / draft-decode) and encoder-decoder
-        calls it like the FA (CUDA) path. Eager-only. A request with
+        cross-/self-attention.
-        cache_seqlens==0 attends to no keys and the kernel returns NaN for its
-        rows, so those rows are zeroed (an all-empty batch skips the launch).
         """
-        q_rows = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim)
+        # page_size=1 attention has no sink-logit support today; drop-on-the-floor
-        # If the maximum cache_seqlens is 0, there are no keys to attend to.
+        # here would silently unsink the branch half of the merge.
-        if int(cache_seqlens.max().item()) == 0:
+        assert sinks is None, (
-            return q_rows.new_zeros(
+            "flat-page-table attention does not support attention sinks"
-                (q_rows.shape[0], q_rows.shape[1], value_cache.shape[-1])
         )
-        # page_size=1 view of the KV pool: PR #454 detects if page size dim is 1
+
-        # i.e. k_cache.shape[1] == 1 and routes flash_attn_with_kvcache to varlen gather.
+        # key_cache/value_cache carry the backend's configured page_size, but
-        k_cache = key_cache.reshape(-1, 1, layer.tp_k_head_num, layer.head_dim)
+        # page_table indexes individual tokens regardless of that page_size, so
-        v_cache = value_cache.reshape(-1, 1, layer.tp_v_head_num, layer.head_dim)
+        # the cache must be re-viewed as page_size=1 to line up with it --
-        out = flash_attn_with_kvcache(
+        # otherwise each row's block id is read as page_table[row, 0], aliasing
-            q=q_rows,
+        # entry 0's block instead of the actual per-row token slots.
-            k_cache=k_cache,
+        k_cache_unpaged = key_cache.reshape(
-            v_cache=v_cache,
+            -1, 1, key_cache.shape[-2], key_cache.shape[-1]
+        )
+        v_cache_unpaged = value_cache.reshape(
+            -1, 1, value_cache.shape[-2], value_cache.shape[-1]
+        )
+
+        return flash_attn_with_kvcache(
+            q=q,
+            k_cache=k_cache_unpaged,
+            v_cache=v_cache_unpaged,
+            # Reads page_table[row, 0:cache_seqlens[row]] per row, so page_table
+            # must pack its valid entries first per row.
+            page_table=page_table,
+            cache_seqlens=cache_seqlens,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k_new=None,
+            max_seqlen_q=max_seqlen_q,
+            softmax_scale=scale,
+            causal=causal,
+            window_size=window_size,
+            softcap=softcap,
+            return_softmax_lse=return_softmax_lse,
+        )
+
+    def _forward_encoder_decoder_attn(
+        self,
+        q: torch.Tensor,
+        key_cache: torch.Tensor,
+        value_cache: torch.Tensor,
+        page_table: torch.Tensor,
+        cache_seqlens: torch.Tensor,
+        cu_seqlens_q: torch.Tensor,
+        max_seqlen_q: int,
+        scale: float,
+        softcap: float,
+        causal: bool,
+    ):
+        """Encoder-decoder cross-/self-attention via the flat page-table path.
+
+        A request with cache_seqlens==0 (empty encoder region, or an idle
+        decoder row) attends to no keys and the kernel returns NaN for it
+        rather than zeros, so those rows are zeroed here (an all-empty batch
+        skips the launch entirely). Cascade's expand rows are never empty, so
+        this check -- and its two device-to-host syncs -- stays out of
+        _forward_attn_flat_page_table and off that path.
+        """
+        if int(cache_seqlens.max().item()) == 0:
+            return q.new_zeros((q.shape[0], q.shape[1], value_cache.shape[-1]))
+
+        out = self._forward_attn_flat_page_table(
+            q=q,
+            key_cache=key_cache,
+            value_cache=value_cache,
             page_table=page_table,
             cache_seqlens=cache_seqlens,
             cu_seqlens_q=cu_seqlens_q,
             max_seqlen_q=max_seqlen_q,
-            softmax_scale=layer.scaling,
+            scale=scale,
+            softcap=softcap,
             causal=causal,
-            softcap=layer.logit_cap,
         )
-        # Mixed batch: requests with cache_seqlens==0 attend to no keys and come
-        # back as NaN, so zero their query rows (mapped via cu_seqlens_q).
         if int(cache_seqlens.min().item()) == 0:
             seg = cu_seqlens_q[1:] - cu_seqlens_q[:-1]
             out[(cache_seqlens == 0).repeat_interleave(seg)] = 0
@@ -947,22 +1006,23 @@ class XPUAttentionBackend(AttentionBackend):
                 -1, self.page_size, layer.tp_k_head_num, layer.head_dim
             )
             value_cache = value_cache.view(
-                -1, self.page_size, layer.tp_v_head_num, layer.head_dim
+                -1, self.page_size, layer.tp_v_head_num, layer.v_head_dim
             )
 
             if self.is_encoder_decoder and forward_batch.encoder_lens is not None:
                 page_table, cache_seqlens, causal = self._encoder_decoder_page_table(
                     layer, metadata
                 )
-                o = self._forward_attn_flat_page_table(
+                o = self._forward_encoder_decoder_attn(
-                    q=q,
+                    q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
                     key_cache=key_cache,
                     value_cache=value_cache,
-                    layer=layer,
                     page_table=page_table,
                     cache_seqlens=cache_seqlens,
                     cu_seqlens_q=metadata.cu_seqlens_q,
                     max_seqlen_q=1,
+                    scale=layer.scaling,
+                    softcap=layer.logit_cap,
                     causal=causal,
                 )
                 return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
@@ -1052,31 +1112,27 @@ class XPUAttentionBackend(AttentionBackend):
                 )
                 if use_cascade_attn:
                     o, softmax_lse, *rest = result
-                    o_expand, softmax_lse_expand, *rest_expand = (
+                    expand = self.forward_metadata_spec_decode_expand
-                        flash_attn_with_kvcache(
+                    o_expand, lse_expand, *_ = self._forward_attn_flat_page_table(
                         q=q_reshaped,
-                            k_cache=key_cache,
+                        key_cache=key_cache,
-                            v_cache=value_cache,
+                        value_cache=value_cache,
-                            page_table=self.forward_metadata_spec_decode_expand.page_table,
+                        page_table=expand.page_table,
-                            cache_seqlens=self.forward_metadata_spec_decode_expand.cache_seqlens_int32,
+                        cache_seqlens=expand.cache_seqlens_int32,
-                            cu_seqlens_q=self.forward_metadata_spec_decode_expand.cu_seqlens_q,
+                        cu_seqlens_q=expand.cu_seqlens_q,
-                            cu_seqlens_k_new=None,
+                        max_seqlen_q=expand.max_seq_len_q,
-                            max_seqlen_q=self.forward_metadata_spec_decode_expand.max_seq_len_q,
+                        scale=layer.scaling,
-                            softmax_scale=layer.scaling,
-                            causal=False,
-                            window_size=window_size,
                         softcap=layer.logit_cap,
-                            k_descale=k_descale,
+                        window_size=window_size,
-                            v_descale=v_descale,
+                        sinks=sinks,
                         return_softmax_lse=True,
-                            **kwargs,
-                        )
                     )
                     o, _ = merge_state_v2(
                         o,
                         softmax_lse.T.contiguous(),
                         o_expand,
-                        softmax_lse_expand.T.contiguous(),
+                        # lse_expand comes back as (Hq, total_q); merge_state wants (total_q, Hq).
+                        lse_expand.T.contiguous(),
                     )
                 else:
                     o = result
@@ -1158,6 +1214,28 @@ class XPUAttentionBackend(AttentionBackend):
         else:
             self.encoder_metadata = {}
 
+    def _spec_query_kv_offsets(
+        self,
+        is_verify: bool,
+        is_draft_extend: bool,
+        is_draft_decode: bool,
+        spec_info,
+    ) -> tuple[int, int]:
+        """Per-request query-row count and how far the KV length must extend
+        past ``seq_lens`` for each spec mode.
+        """
+        if is_verify:
+            # Packs speculative_num_draft_tokens rows, attending past seq_lens
+            # over those draft positions.
+            return self.speculative_num_draft_tokens, self.speculative_num_draft_tokens
+        if is_draft_extend:
+            # seq_lens already includes the extend tokens (standard extend
+            # convention), so no KV offset is needed.
+            return spec_info.num_tokens_per_req, 0
+        if is_draft_decode:
+            return 1, self.speculative_step_id + 1
+        return 1, 0
+
     def init_forward_metadata_out_graph(
         self,
         forward_batch: ForwardBatch,
@@ -1179,11 +1257,15 @@ class XPUAttentionBackend(AttentionBackend):
         forward_mode = forward_batch.forward_mode
         spec_info = forward_batch.spec_info
 
-        assert spec_info is None, (
+        is_verify = forward_mode.is_target_verify()
-            "XPUAttentionBackend does not support speculative decoding in XPU graph"
+        is_draft_decode = forward_mode.is_decode_or_idle() and spec_info is not None
+        is_draft_extend = forward_mode.is_draft_extend_v2()
+        assert forward_mode.is_decode_or_idle() or is_verify or is_draft_extend, (
+            "XPUAttentionBackend XPU graph only supports decode / target-verify / draft-extend modes"
         )
-        assert forward_mode.is_decode_or_idle(), (
+
-            "XPUAttentionBackend XPU graph only supports decode mode"
+        q_len_per_req, kv_len_offset = self._spec_query_kv_offsets(
+            is_verify, is_draft_extend, is_draft_decode, spec_info
         )
 
         if in_capture:
@@ -1229,14 +1311,30 @@ class XPUAttentionBackend(AttentionBackend):
             if seq_lens_cpu is not None
             else seq_lens.max().item()
         )
-        metadata.max_seq_len_k = max_len
+        metadata.max_seq_len_k = max_len + kv_len_offset
 
-        metadata.cache_seqlens_int32.copy_(seq_lens.to(torch.int32))
+        kv_seqlens = (seq_lens + kv_len_offset).to(torch.int32)
+        metadata.cache_seqlens_int32.copy_(kv_seqlens)
 
         metadata.cu_seqlens_k[0] = 0
-        metadata.cu_seqlens_k[1 : bs + 1].copy_(
+        metadata.cu_seqlens_k[1 : bs + 1].copy_(torch.cumsum(kv_seqlens, dim=0))
-            torch.cumsum(seq_lens.to(torch.int32), dim=0)
+
+        # target-verify and draft-extend pack multiple query rows per request;
+        # rebuild cu_seqlens_q as a strided ramp (0, q, 2q, ...). Plain/draft
+        # decode keep the identity ramp already stored in the pre-allocated buffer.
+        if q_len_per_req > 1:
+            metadata.max_seq_len_q = q_len_per_req
+            metadata.cu_seqlens_q[: bs + 1].copy_(
+                torch.arange(
+                    0,
+                    (bs + 1) * q_len_per_req,
+                    q_len_per_req,
+                    dtype=torch.int32,
+                    device=metadata.cu_seqlens_q.device,
                 )
+            )
+        else:
+            metadata.max_seq_len_q = 1
 
         if self.is_encoder_decoder and forward_batch.encoder_lens is not None:
             encoder_lens = forward_batch.encoder_lens[:bs].to(torch.int32)
@@ -1428,3 +1526,64 @@ class XPUAttentionBackend(AttentionBackend):
             metadata_swa.cu_seqlens_k.copy_(cu_seqlens_k)
 
         metadata.swa_spec_metadata = metadata_swa
+
+
+class XPUMultiStepDraftBackend:
+    """Wrap multiple XPU attention backends for consecutive draft decode steps."""
+
+    needs_cpu_seq_lens: bool = False
+
+    def __init__(
+        self,
+        model_runner: ModelRunner,
+        topk: int,
+        speculative_num_steps: int,
+    ):
+        self.topk = topk
+        self.speculative_num_steps = speculative_num_steps
+        self.attn_backends = [
+            XPUAttentionBackend(
+                model_runner,
+                skip_prefill=True,
+                speculative_step_id=i,
+                topk=topk,
+                speculative_num_steps=speculative_num_steps,
+            )
+            for i in range(speculative_num_steps - 1)
+        ]
+        self.max_context_len = self.attn_backends[0].max_context_len
+        self.device = model_runner.device
+        self.token_to_kv_pool = model_runner.token_to_kv_pool
+
+    def init_forward_metadata(self, forward_batch: ForwardBatch):
+        for attn_backend in self.attn_backends:
+            attn_backend.init_forward_metadata(forward_batch)
+
+    def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int):
+        for attn_backend in self.attn_backends:
+            attn_backend.init_cuda_graph_state(max_bs, max_num_tokens)
+
+    def init_forward_metadata_out_graph(
+        self,
+        forward_batch: ForwardBatch,
+        in_capture: bool = False,
+    ):
+        from sglang.srt.model_executor.forward_batch_info import build_inner_fb_view
+
+        assert forward_batch.spec_info is not None
+        assert forward_batch.spec_info.is_draft_input()
+
+        inner_fb = build_inner_fb_view(
+            forward_batch,
+            bs=forward_batch.batch_size,
+            forward_mode=ForwardMode.DECODE,
+            encoder_lens=forward_batch.encoder_lens,
+        )
+        for attn_backend in self.attn_backends:
+            attn_backend.init_forward_metadata_out_graph(
+                inner_fb, in_capture=in_capture
+            )
+
+    def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch) -> None:
+        for attn_backend in self.attn_backends:
+            attn_backend.init_forward_metadata_in_graph(forward_batch)

python/sglang/srt/speculative/draft_utils.py

@@ -97,6 +97,7 @@ class DraftBackendFactory:
             "flashinfer": self._create_flashinfer_decode_backend,
             "triton": self._create_triton_decode_backend,
             "intel_amx": self._create_intel_amx_decode_backend,
+            "intel_xpu": self._create_intel_xpu_decode_backend,
             "aiter": self._create_aiter_decode_backend,
             "fa3": self._create_fa3_decode_backend,
             "hybrid_linear_attn": self._create_hybrid_linear_attn_decode_backend,
@@ -127,6 +128,7 @@ class DraftBackendFactory:
             "flashinfer": self._create_flashinfer_prefill_backend,
             "triton": self._create_triton_prefill_backend,
             "intel_amx": self._create_intel_amx_prefill_backend,
+            "intel_xpu": self._create_intel_xpu_prefill_backend,
             "aiter": self._create_aiter_prefill_backend,
             "fa3": self._create_fa3_prefill_backend,
             "hybrid_linear_attn": self._create_hybrid_linear_attn_prefill_backend,
@@ -289,6 +291,16 @@ class DraftBackendFactory:
             return self._create_triton_prefill_backend()
         return self._create_fa3_prefill_backend()
 
+    def _create_intel_xpu_decode_backend(self):
+        from sglang.srt.layers.attention.xpu_backend import XPUMultiStepDraftBackend
+
+        return (
+            "intel_xpu",
+            XPUMultiStepDraftBackend(
+                self.draft_model_runner, self.topk, self.speculative_num_steps
+            ),
+        )
+
     def _create_aiter_decode_backend(self):
         from sglang.srt.layers.attention.aiter_backend import AiterMultiStepDraftBackend
 
@@ -477,6 +489,14 @@ class DraftBackendFactory:
 
         return ("intel_amx", IntelAMXAttnBackend(self.draft_model_runner))
 
+    def _create_intel_xpu_prefill_backend(self):
+        from sglang.srt.layers.attention.xpu_backend import XPUAttentionBackend
+
+        return (
+            "intel_xpu",
+            XPUAttentionBackend(self.draft_model_runner, skip_prefill=False),
+        )
+
     def _create_aiter_prefill_backend(self):
         from sglang.srt.layers.attention.aiter_backend import AiterAttnBackend
 

test/registered/e2e/xpu/test_spec_eagle_intel_xpu.py

@@ -0,0 +1,57 @@
+"""intel_xpu attention backend (EAGLE3 topk=1 chain + EAGLE/Llama-2 spec)."""
+
+import unittest
+
+from sglang.srt.environ import envs
+from sglang.test.ci.ci_register import register_xpu_ci
+from sglang.test.kits.matched_stop_kit import MatchedStopMixin
+from sglang.test.kits.spec_server_kits import (
+    SpecAccuracyKit,
+    SpecFeatureKit,
+    SpecHiddenStatesKit,
+    SpecLogprobKit,
+    SpecPenaltyKit,
+)
+from sglang.test.server_fixtures.spec_eagle_fixture import Eagle3Base, EagleLlama2Base
+
+register_xpu_ci(est_time=1800, suite="nightly-xpu-1-gpu", nightly=True)
+
+
+class TestEagle3IntelXPU(
+    Eagle3Base,
+    MatchedStopMixin,
+    SpecAccuracyKit,
+    SpecLogprobKit,
+    SpecPenaltyKit,
+    SpecFeatureKit,
+):
+    """EAGLE3 spec v2 on the intel_xpu attention backend (kits listed in bases)."""
+
+    attention_backend = "intel_xpu"
+    max_running_requests = 24
+    gsm8k_num_examples = 300
+    gsm8k_check_accept_len = True
+    mem_fraction_static = 0.95
+    extra_args = ("--max-total-tokens", "16384", "--disable-decode-cuda-graph")
+    env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)
+
+
+class TestEagleLlama2IntelXPU(
+    EagleLlama2Base, SpecAccuracyKit, SpecFeatureKit, SpecHiddenStatesKit
+):
+    """EAGLE/Llama-2 on intel_xpu using the supported topk = 1 paged config."""
+
+    attention_backend = "intel_xpu"
+    spec_topk = 1
+    page_size = 64
+    gsm8k_check_accept_len = True
+    gsm8k_num_examples = 300
+    enable_return_hidden_states = True
+    mem_fraction_static = 0.95
+    max_running_requests = 6
+    chunked_prefill_size = 512
+    extra_args = ("--max-total-tokens", "16384", "--disable-decode-cuda-graph")
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/layers/attention/test_encoder_decoder_varlen_gather.py

@@ -57,8 +57,8 @@ class TestEncoderDecoderForward(unittest.TestCase):
         # The caller picks (page_table, cache_seqlens, causal) via
         # _encoder_decoder_page_table -- cross-attn -> encoder_page_table +
         # encoder_lens_int32 + causal=False; self-attn -> page_table +
-        # cache_seqlens_int32 + causal=True -- then hands them to the generic
+        # cache_seqlens_int32 + causal=True -- then hands them to
-        # _forward_attn_flat_page_table, which must forward them unchanged with a
+        # _forward_encoder_decoder_attn, which must forward them unchanged with a
         # page_size=1 k_cache (shape[1]==1) so PR #454 routes to the varlen gather.
         enc_pt = torch.arange(5, dtype=torch.int32).unsqueeze(0)
         dec_pt = (torch.arange(4, dtype=torch.int32) + 10).unsqueeze(0)
@@ -70,7 +70,7 @@ class TestEncoderDecoderForward(unittest.TestCase):
         )
         key_cache = self.k_flat.view(-1, 1, self.HK, self.D)
         value_cache = self.v_flat.view(-1, 1, self.HK, self.D)
-        q = torch.randn(1, self.HQ * self.D)
+        q = torch.randn(1, self.HQ, self.D)
         cu_seqlens_q = torch.tensor([0, 1], dtype=torch.int32)
 
         for is_cross, exp_pt, exp_seqlens, exp_causal in (
@@ -94,15 +94,16 @@ class TestEncoderDecoderForward(unittest.TestCase):
                 )
 
             with patch.object(xpu_backend, "flash_attn_with_kvcache", fake_kvcache):
-                self.backend._forward_attn_flat_page_table(
+                self.backend._forward_encoder_decoder_attn(
                     q=q,
                     key_cache=key_cache,
                     value_cache=value_cache,
-                    layer=layer,
                     page_table=page_table,
                     cache_seqlens=cache_seqlens,
                     cu_seqlens_q=cu_seqlens_q,
                     max_seqlen_q=1,
+                    scale=layer.scaling,
+                    softcap=layer.logit_cap,
                     causal=causal,
                 )
             self.assertTrue(torch.equal(captured["page_table"], exp_pt))
@@ -117,18 +118,20 @@ class TestEncoderDecoderForward(unittest.TestCase):
         # zeros and never launch the kernel.
         key_cache = self.k_flat.view(-1, 1, self.HK, self.D)
         value_cache = self.v_flat.view(-1, 1, self.HK, self.D)
-        q = torch.randn(1, self.HQ * self.D)
+        q = torch.randn(1, self.HQ, self.D)
         sentinel = MagicMock(side_effect=AssertionError("kernel must not run"))
+        layer = self._layer(True)
         with patch.object(xpu_backend, "flash_attn_with_kvcache", sentinel):
-            out = self.backend._forward_attn_flat_page_table(
+            out = self.backend._forward_encoder_decoder_attn(
                 q=q,
                 key_cache=key_cache,
                 value_cache=value_cache,
-                layer=self._layer(True),
                 page_table=torch.zeros(1, 0, dtype=torch.int32),
                 cache_seqlens=torch.zeros(1, dtype=torch.int32),
                 cu_seqlens_q=torch.tensor([0, 1], dtype=torch.int32),
                 max_seqlen_q=1,
+                scale=layer.scaling,
+                softcap=layer.logit_cap,
                 causal=False,
             )
         sentinel.assert_not_called()
@@ -141,7 +144,7 @@ class TestEncoderDecoderForward(unittest.TestCase):
         # counts (2 and 3) exercise the cu_seqlens_q -> per-request row mapping.
         key_cache = self.k_flat.view(-1, 1, self.HK, self.D)
         value_cache = self.v_flat.view(-1, 1, self.HK, self.D)
-        q = torch.randn(5, self.HQ * self.D)
+        q = torch.randn(5, self.HQ, self.D)
 
         def fake_kvcache(*_, **kw):
             # All-ones (never-NaN) sentinel so zeroed rows are distinguishable.
@@ -149,16 +152,18 @@ class TestEncoderDecoderForward(unittest.TestCase):
                 (kw["q"].shape[0], kw["q"].shape[1], kw["v_cache"].shape[-1])
             )
 
+        layer = self._layer(True)
         with patch.object(xpu_backend, "flash_attn_with_kvcache", fake_kvcache):
-            out = self.backend._forward_attn_flat_page_table(
+            out = self.backend._forward_encoder_decoder_attn(
                 q=q,
                 key_cache=key_cache,
                 value_cache=value_cache,
-                layer=self._layer(True),
                 page_table=torch.zeros(2, 4, dtype=torch.int32),
                 cache_seqlens=torch.tensor([0, 4], dtype=torch.int32),
                 cu_seqlens_q=torch.tensor([0, 2, 5], dtype=torch.int32),
                 max_seqlen_q=3,
+                scale=layer.scaling,
+                softcap=layer.logit_cap,
                 causal=False,
             )
         self.assertTrue(torch.equal(out[:2], torch.zeros(2, self.HQ, self.D)))

test/registered/xpu/test_xpu_encoder_decoder_varlen.py

@@ -2,13 +2,12 @@
 
 The backend calls flash_attn_with_kvcache with a page_size=1 view; sgl-kernel-xpu
 PR #454 detects that and gathers + runs varlen inside the kernel. This runs on an
-actual XPU and guards what a mocked CPU test cannot: _forward_attn_flat_page_table
+actual XPU and guards what a mocked CPU test cannot: _forward_encoder_decoder_attn
 plus the real kernel produce correct attention for a scattered (non-page-aligned)
 token-slot layout, for both cross-attn (non-causal) and decoder self-attn (causal).
 """
 
 import unittest
-from types import SimpleNamespace
 
 import torch
 
@@ -85,30 +84,23 @@ class TestXPUEncoderDecoderVarlen(CustomTestCase):
             .to(self.dev)
         )
         q = torch.randn(num_rows, self.H, self.D, dtype=torch.bfloat16, device=self.dev)
-        layer = SimpleNamespace(
+        scale, softcap = 0.5, 0.0
-            is_cross_attention=not causal,
-            tp_q_head_num=self.H,
-            tp_k_head_num=self.H,
-            tp_v_head_num=self.H,
-            head_dim=self.D,
-            scaling=0.5,
-            logit_cap=0.0,
-        )
         key_cache = self.k_flat.view(-1, 1, self.H, self.D)
         value_cache = self.v_flat.view(-1, 1, self.H, self.D)
 
         # causal=True mirrors decoder self-attn, causal=False cross-attn; the
         # generic helper takes the (page_table, cache_seqlens, causal) that the
         # caller's _encoder_decoder_page_table dispatch would have selected.
-        got = self.backend._forward_attn_flat_page_table(
+        got = self.backend._forward_encoder_decoder_attn(
             q=q,
             key_cache=key_cache,
             value_cache=value_cache,
-            layer=layer,
             page_table=page_table,
             cache_seqlens=cache_seqlens,
             cu_seqlens_q=cu_seqlens_q,
             max_seqlen_q=1,
+            scale=scale,
+            softcap=softcap,
             causal=causal,
         )
         torch.xpu.synchronize()
@@ -119,7 +111,7 @@ class TestXPUEncoderDecoderVarlen(CustomTestCase):
             page_table=page_table,
             cache_seqlens=cache_seqlens,
             cu_seqlens_q=cu_seqlens_q,
-            scale=layer.scaling,
+            scale=scale,
             causal=causal,
         )
         self.assertEqual(tuple(got.shape), (num_rows, self.H, self.D))