diff --git a/python/sglang/srt/layers/attention/xpu_backend.py b/python/sglang/srt/layers/attention/xpu_backend.py index 5573bc477..60a2d6901 100644 --- a/python/sglang/srt/layers/attention/xpu_backend.py +++ b/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( - q=q, + o = self._forward_encoder_decoder_attn( + 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 - # output buffer at a stable device address so the graph can - # record writes to the same storage on every replay. - # _attn_output is that fixed buffer; None falls back to a - # freshly allocated tensor (eager / cascade-attn path). - 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 - else None + # `out` is injected via out_kwargs only on the graph path (buffer + # present, non-cascade); the eager path omits it for flash_attn + # builds that lack the kwarg. Piecewise XPU graph for prefill + # pre-allocates this fixed-address output buffer so graph replay + # writes to the same storage; radix_attention sets it only on the + # graph path, None on the eager path. + **( + { + "out": forward_batch._attn_output.view( + -1, layer.tp_q_head_num, layer.v_head_dim + ) + } + if not use_cascade_attn and forward_batch._attn_output is not 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, - v_cache=value_cache, - 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=None, - max_seqlen_q=self.forward_metadata_spec_decode_expand.max_seq_len_q, - softmax_scale=layer.scaling, - causal=False, - window_size=window_size, + key_cache=key_cache, + value_cache=value_cache, + page_table=expand.page_table, + cache_seqlens=expand.cache_seqlens_int32, + cu_seqlens_q=expand.cu_seqlens_q, + max_seqlen_q=expand.max_seq_len_q, + scale=layer.scaling, softcap=layer.logit_cap, - k_descale=k_descale, - v_descale=v_descale, + window_size=window_size, + 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, - key_cache, - value_cache, - layer, - page_table, - cache_seqlens, - cu_seqlens_q, - max_seqlen_q, - causal, + q: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + page_table: torch.Tensor, + cache_seqlens: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor], + max_seqlen_q: int, + scale: float, + softcap: float, + 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 - token-slot) page table. sgl-kernel-xpu PR #454 detects a page_size==1 - k_cache + page_table and gathers + runs varlen internally, so the backend - calls it like the FA (CUDA) path. Eager-only. A request with - 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). + """Attention against a ``page_table`` that indexes individual token + slots rather than ``self.page_size``-sized blocks. Serves cascade + tree-branch expand (target-verify / draft-decode) and encoder-decoder + cross-/self-attention. + """ + # page_size=1 attention has no sink-logit support today; drop-on-the-floor + # here would silently unsink the branch half of the merge. + assert sinks is None, ( + "flat-page-table attention does not support attention sinks" + ) + + # key_cache/value_cache carry the backend's configured page_size, but + # page_table indexes individual tokens regardless of that page_size, so + # the cache must be re-viewed as page_size=1 to line up with it -- + # otherwise each row's block id is read as page_table[row, 0], aliasing + # entry 0's block instead of the actual per-row token slots. + k_cache_unpaged = key_cache.reshape( + -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. """ - q_rows = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim) - # If the maximum cache_seqlens is 0, there are no keys to attend to. if int(cache_seqlens.max().item()) == 0: - return q_rows.new_zeros( - (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. - k_cache = key_cache.reshape(-1, 1, layer.tp_k_head_num, layer.head_dim) - v_cache = value_cache.reshape(-1, 1, layer.tp_v_head_num, layer.head_dim) - out = flash_attn_with_kvcache( - q=q_rows, - k_cache=k_cache, - v_cache=v_cache, + 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( - q=q, + o = self._forward_encoder_decoder_attn( + 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 = ( - flash_attn_with_kvcache( - q=q_reshaped, - k_cache=key_cache, - v_cache=value_cache, - 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=None, - 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, - **kwargs, - ) + expand = self.forward_metadata_spec_decode_expand + o_expand, lse_expand, *_ = self._forward_attn_flat_page_table( + q=q_reshaped, + key_cache=key_cache, + value_cache=value_cache, + page_table=expand.page_table, + cache_seqlens=expand.cache_seqlens_int32, + cu_seqlens_q=expand.cu_seqlens_q, + max_seqlen_q=expand.max_seq_len_q, + scale=layer.scaling, + softcap=layer.logit_cap, + window_size=window_size, + sinks=sinks, + return_softmax_lse=True, ) 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, ( - "XPUAttentionBackend does not support speculative decoding in XPU graph" + is_verify = forward_mode.is_target_verify() + 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_( - torch.cumsum(seq_lens.to(torch.int32), dim=0) - ) + metadata.cu_seqlens_k[1 : bs + 1].copy_(torch.cumsum(kv_seqlens, 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) diff --git a/python/sglang/srt/speculative/draft_utils.py b/python/sglang/srt/speculative/draft_utils.py index 365c11b81..b7ff407fd 100644 --- a/python/sglang/srt/speculative/draft_utils.py +++ b/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 diff --git a/test/registered/e2e/xpu/test_spec_eagle_intel_xpu.py b/test/registered/e2e/xpu/test_spec_eagle_intel_xpu.py new file mode 100644 index 000000000..5766d2152 --- /dev/null +++ b/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() diff --git a/test/registered/unit/layers/attention/test_encoder_decoder_varlen_gather.py b/test/registered/unit/layers/attention/test_encoder_decoder_varlen_gather.py index afa15db5f..1c0647d34 100644 --- a/test/registered/unit/layers/attention/test_encoder_decoder_varlen_gather.py +++ b/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 - # _forward_attn_flat_page_table, which must forward them unchanged with a + # cache_seqlens_int32 + causal=True -- then hands them to + # _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))) diff --git a/test/registered/xpu/test_xpu_encoder_decoder_varlen.py b/test/registered/xpu/test_xpu_encoder_decoder_varlen.py index 403e49340..5dfad6f95 100644 --- a/test/registered/xpu/test_xpu_encoder_decoder_varlen.py +++ b/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( - 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, - ) + scale, softcap = 0.5, 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))