From 2c8357f7947199e662fa0af85a65442ea4f18103 Mon Sep 17 00:00:00 2001 From: Juan Muneton <102537701+jmunetong@users.noreply.github.com> Date: Thu, 4 Jun 2026 19:05:07 -0700 Subject: [PATCH] [XPU] Enable Gemma 4 E2B / E4B / 31B/ 26B-A4B on Intel XPU (#23280) Co-authored-by: Claude Opus 4.6 (1M context) Co-authored-by: jmunetong Co-authored-by: Meng, Hengyu Co-authored-by: ckvermaAI --- .../srt/layers/attention/xpu_backend.py | 99 +++++++++++-- python/sglang/srt/layers/gemma4_fused_ops.py | 73 +++++++++- python/sglang/srt/layers/layernorm.py | 9 ++ python/sglang/srt/models/gemma4_causal.py | 61 +++++--- python/sglang/srt/server_args.py | 4 +- test/registered/xpu/test_gemma_4_e2b.py | 136 ++++++++++++++++++ 6 files changed, 348 insertions(+), 34 deletions(-) create mode 100644 test/registered/xpu/test_gemma_4_e2b.py diff --git a/python/sglang/srt/layers/attention/xpu_backend.py b/python/sglang/srt/layers/attention/xpu_backend.py index 41293a20b..222fb63e6 100644 --- a/python/sglang/srt/layers/attention/xpu_backend.py +++ b/python/sglang/srt/layers/attention/xpu_backend.py @@ -13,6 +13,7 @@ from sglang.srt.layers.attention.flashattention_backend import ( prepare_swa_spec_page_table_triton, ) from sglang.srt.managers.schedule_batch import get_global_server_args +from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode if TYPE_CHECKING: @@ -72,6 +73,12 @@ class XPUAttentionBackend(AttentionBackend): self.use_mla = model_runner.model_config.attention_arch == AttentionArch.MLA self.skip_prefill = skip_prefill self.is_hybrid_swa = model_runner.is_hybrid_swa + self.use_sliding_window_kv_pool = ( + isinstance(model_runner.token_to_kv_pool, SWAKVPool) + and model_runner.token_to_kv_pool.swa_layer_nums > 0 + ) + if self.use_sliding_window_kv_pool: + self.token_to_kv_pool = model_runner.token_to_kv_pool if self.is_hybrid_swa: self.full_to_swa_index_mapping = ( model_runner.token_to_kv_pool.full_to_swa_index_mapping @@ -193,6 +200,7 @@ class XPUAttentionBackend(AttentionBackend): metadata.page_table = self.req_to_token_pool.req_to_token[ forward_batch.req_pool_indices, : metadata.max_seq_len_k ] + # TODO: we need to test this part for llama 4 eagle case self._init_local_attn_metadata(forward_batch, metadata, device) elif forward_batch.forward_mode.is_target_verify(): @@ -373,6 +381,16 @@ class XPUAttentionBackend(AttentionBackend): ), ] + # Translate full-pool indices to SWA-pool indices for hybrid models + if self.use_sliding_window_kv_pool: + # flash_attn_with_kvcache requires int32 page tables; the SWA index + # mapping is int64, so cast (matches flashattention_backend.py). + metadata.swa_page_table = ( + self.token_to_kv_pool.translate_loc_from_full_to_swa( + metadata.page_table + ).to(torch.int32) + ) + if self.use_mla: workspace_size = flash_mla_get_workspace_size( max_seq_len=self.max_context_len, @@ -389,11 +407,27 @@ class XPUAttentionBackend(AttentionBackend): workspace_size, device=self.device, dtype=torch.uint8 ) + # Translate full-pool indices to SWA-pool indices for hybrid models + if self.use_sliding_window_kv_pool: + # flash_attn_with_kvcache requires int32 page tables; the SWA index + # mapping is int64, so cast (matches flashattention_backend.py). + metadata.swa_page_table = ( + self.token_to_kv_pool.translate_loc_from_full_to_swa( + metadata.page_table + ).to(torch.int32) + ) + # Convert the page table to a strided format which is needed by FA3 API if self.page_size > 1: self.strided_indices = torch.arange( 0, metadata.page_table.shape[1], self.page_size, device=self.device ) + + if self.use_sliding_window_kv_pool and metadata.swa_page_table is not None: + metadata.swa_page_table = ( + metadata.swa_page_table[:, self.strided_indices] // self.page_size + ) + metadata.page_table = ( metadata.page_table[:, self.strided_indices] // self.page_size ) @@ -413,8 +447,17 @@ class XPUAttentionBackend(AttentionBackend): k_rope: Optional[torch.Tensor] = None, sinks: Optional[torch.Tensor] = None, ): - if k is not None: - assert v is not None + if k is None and v is None: + # Cross-layer KV sharing (Gemma 4): the layer reuses another + # layer's KV cache. The paged kernel reads K/V directly via + # page_table, and pool.get_kv_buffer(layer.layer_id) routes + # to the correct sub-pool because RadixAttention is initialized + # with layer_id=kv_shared_layer_index for shared layers. No + # materialization needed; just skip the write path. + pass + elif k is None or v is None: + raise ValueError("Both k and v should be None or not None") + else: if save_kv_cache: cache_loc = ( forward_batch.out_cache_loc @@ -497,6 +540,13 @@ class XPUAttentionBackend(AttentionBackend): cu_seqlens_k = swa_spec_metadata.cu_seqlens_k else: page_table = metadata.page_table + if is_hybrid_swa and self.use_sliding_window_kv_pool: + if metadata.swa_page_table is not None: + page_table = metadata.swa_page_table + else: + page_table = self.token_to_kv_pool.translate_loc_from_full_to_swa( + metadata.page_table + ).to(torch.int32) cu_seqlens_q = metadata.cu_seqlens_q cache_seqlens = metadata.cache_seqlens_int32 max_seqlen_q = metadata.max_seq_len_q @@ -525,7 +575,7 @@ class XPUAttentionBackend(AttentionBackend): page_table=page_table, cache_seqlens=cache_seqlens, cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k_new=cu_seqlens_k if not use_local_attn else None, + cu_seqlens_k_new=None, max_seqlen_q=max_seqlen_q, softmax_scale=layer.scaling, causal=False if use_cascade_attn else causal, @@ -546,7 +596,7 @@ class XPUAttentionBackend(AttentionBackend): 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, + cu_seqlens_k_new=None, max_seqlen_q=self.forward_metadata_spec_decode_expand.max_seq_len_q, softmax_scale=layer.scaling, causal=False, @@ -648,7 +698,7 @@ class XPUAttentionBackend(AttentionBackend): page_table=page_table, cache_seqlens=cache_seqlens, cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k_new=cu_seqlens_k if not use_local_attn else None, + cu_seqlens_k_new=None, max_seqlen_q=max_seqlen_q, softmax_scale=layer.scaling, causal=False if use_cascade_attn else causal, @@ -668,7 +718,7 @@ class XPUAttentionBackend(AttentionBackend): 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, + cu_seqlens_k_new=None, max_seqlen_q=self.forward_metadata_spec_decode_expand.max_seq_len_q, softmax_scale=layer.scaling, causal=False, @@ -688,7 +738,8 @@ class XPUAttentionBackend(AttentionBackend): else: o = result - return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) + out = o.view(-1, layer.tp_q_head_num * layer.v_head_dim) + return out def forward_decode( self, @@ -703,8 +754,12 @@ class XPUAttentionBackend(AttentionBackend): k_rope: Optional[torch.Tensor] = None, sinks: Optional[torch.Tensor] = None, ) -> torch.Tensor: - if k is not None: - assert v is not None + if k is None and v is None: + # Cross-layer KV sharing (Gemma 4): see forward_extend for details. + pass + elif k is None or v is None: + raise ValueError("Both k and v should be None or not None") + else: if save_kv_cache: cache_loc = ( forward_batch.out_cache_loc @@ -787,7 +842,7 @@ class XPUAttentionBackend(AttentionBackend): page_table=metadata.encoder_page_table, cache_seqlens=metadata.encoder_lens_int32, cu_seqlens_q=metadata.cu_seqlens_q, - cu_seqlens_k_new=metadata.encoder_cu_seqlens_k, + cu_seqlens_k_new=None, max_seqlen_q=1, softmax_scale=layer.scaling, causal=False, @@ -817,7 +872,24 @@ class XPUAttentionBackend(AttentionBackend): **kwargs, ) else: + is_swa_layer = ( + layer.sliding_window_size is not None + and layer.sliding_window_size > -1 + ) + page_table = metadata.page_table + # For SWA layers on hybrid models, use the translated + # SWA-pool page table so KV reads hit the correct pool. + if is_swa_layer and self.use_sliding_window_kv_pool: + if metadata.swa_page_table is not None: + page_table = metadata.swa_page_table + else: + page_table = ( + self.token_to_kv_pool.translate_loc_from_full_to_swa( + metadata.page_table + ).to(torch.int32) + ) + cache_seqlens = metadata.cache_seqlens_int32 cu_seqlens_k = metadata.cu_seqlens_k max_seqlen_q = metadata.max_seq_len_q @@ -833,7 +905,7 @@ class XPUAttentionBackend(AttentionBackend): page_table=page_table, cache_seqlens=cache_seqlens, cu_seqlens_q=metadata.cu_seqlens_q, - cu_seqlens_k_new=cu_seqlens_k, + cu_seqlens_k_new=None, max_seqlen_q=max_seqlen_q, softmax_scale=layer.scaling, causal=False if use_cascade_attn else causal, @@ -854,7 +926,7 @@ class XPUAttentionBackend(AttentionBackend): 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, + cu_seqlens_k_new=None, max_seqlen_q=self.forward_metadata_spec_decode_expand.max_seq_len_q, softmax_scale=layer.scaling, causal=False, @@ -899,7 +971,8 @@ class XPUAttentionBackend(AttentionBackend): layer.scaling, ) - return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) + out = o.view(-1, layer.tp_q_head_num * layer.v_head_dim) + return out def get_cuda_graph_seq_len_fill_value(self): """Get the fill value for sequence length in CUDA graph.""" diff --git a/python/sglang/srt/layers/gemma4_fused_ops.py b/python/sglang/srt/layers/gemma4_fused_ops.py index 14a677e66..11a0ddd5f 100644 --- a/python/sglang/srt/layers/gemma4_fused_ops.py +++ b/python/sglang/srt/layers/gemma4_fused_ops.py @@ -215,7 +215,7 @@ def gemma_qkv_rmsnorm( If k and v are both None (KV-shared layer), only Q is normalized. """ - assert q.is_cuda + assert q.is_cuda or q.is_xpu assert q.stride(-1) == 1, "Q's last dim must be contiguous" assert q_weight.shape[-1] == head_dim M = q.shape[0] if q.dim() >= 2 else 1 @@ -223,7 +223,7 @@ def gemma_qkv_rmsnorm( has_kv = k is not None and v is not None if has_kv: - assert k.is_cuda and v.is_cuda + assert (k.is_cuda and v.is_cuda) or (k.is_xpu and v.is_xpu) assert k.stride(-1) == 1 and v.stride(-1) == 1 assert k_weight is not None and k_weight.shape[-1] == head_dim @@ -245,6 +245,75 @@ def gemma_qkv_rmsnorm( ) +@triton.jit +def _gemma_routing_post_topk_kernel( + Logits_ptr, + Ids_ptr, + Scale_ptr, + Out_weights_ptr, + Out_ids_ptr, + stride_l, + stride_ow, + stride_oi, + K: tl.constexpr, + BLOCK_K: tl.constexpr, +): + """Fused: softmax(topk_logits) * per_expert_scale[topk_ids] → float32 weights, int32 ids. + + One program per token. K is the number of top-k experts (e.g. 8). + """ + row = tl.program_id(0) + cols = tl.arange(0, BLOCK_K) + mask = cols < K + + logits = tl.load( + Logits_ptr + row * stride_l + cols, mask=mask, other=float("-inf") + ).to(tl.float32) + ids_i64 = tl.load(Ids_ptr + row * stride_l + cols, mask=mask, other=0) + + # Stable softmax + max_val = tl.max(logits, axis=0) + exp_val = tl.exp(logits - max_val) + sum_exp = tl.sum(exp_val, axis=0) + weights = exp_val / sum_exp + + # Gather per_expert_scale and multiply + scale = tl.load(Scale_ptr + ids_i64, mask=mask, other=1.0).to(tl.float32) + weights = weights * scale + + tl.store(Out_weights_ptr + row * stride_ow + cols, weights, mask=mask) + tl.store(Out_ids_ptr + row * stride_oi + cols, ids_i64.to(tl.int32), mask=mask) + + +def gemma_routing_post_topk( + topk_logits: torch.Tensor, + topk_ids: torch.Tensor, + per_expert_scale: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Fused softmax + scale-gather + casts for Gemma4 routing. + + Replaces: softmax(topk_logits) * per_expert_scale[topk_ids] → (f32, i32). + """ + B, K = topk_logits.shape + BLOCK_K = triton.next_power_of_2(K) + out_weights = torch.empty((B, K), dtype=torch.float32, device=topk_logits.device) + out_ids = torch.empty((B, K), dtype=torch.int32, device=topk_logits.device) + + _gemma_routing_post_topk_kernel[(B,)]( + topk_logits, + topk_ids, + per_expert_scale, + out_weights, + out_ids, + topk_logits.stride(0), + out_weights.stride(0), + out_ids.stride(0), + K=K, + BLOCK_K=BLOCK_K, + ) + return out_weights, out_ids + + def gemma_dual_rmsnorm_residual_scalar( x1: torch.Tensor, weight1: torch.Tensor, diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index 951d704f7..6cf269b8d 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -878,6 +878,15 @@ class Gemma4RMSNorm(MultiPlatformOp): out = out.reshape(original_shape) return out + def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: + if x.numel() == 0: + return x + if self.with_scale and self.scale_shift == 1.0: + out = gemma_rmsnorm(x, self.weight.data, self.eps) + else: + out = rmsnorm(x, self.weight.data, self.eps) + return out + def forward_hip(self, x: torch.Tensor) -> torch.Tensor: # sgl_kernel's gemma_rmsnorm is not available on ROCm; # delegate to the pure-PyTorch implementation. diff --git a/python/sglang/srt/models/gemma4_causal.py b/python/sglang/srt/models/gemma4_causal.py index 29b0cfd4b..9c2f9c25d 100644 --- a/python/sglang/srt/models/gemma4_causal.py +++ b/python/sglang/srt/models/gemma4_causal.py @@ -34,6 +34,7 @@ from sglang.srt.layers.gemma4_fused_ops import ( gemma_dual_rmsnorm_residual_scalar, gemma_qkv_rmsnorm, gemma_rmsnorm_residual_scalar, + gemma_routing_post_topk, ) from sglang.srt.layers.layernorm import Gemma4RMSNorm, RMSNorm from sglang.srt.layers.linear import ( @@ -56,6 +57,9 @@ from sglang.srt.model_loader.weight_utils import ( maybe_remap_kv_scale_name, ) from sglang.srt.models.gemma3_causal import Gemma3MLP, Gemma3TextScaledWordEmbedding +from sglang.srt.models.utils import ( + create_fused_set_kv_buffer_arg, +) from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import add_prefix, make_layers @@ -145,7 +149,8 @@ class Gemma4Router(nn.Module): super().__init__() self.hidden_size = config.hidden_size - # RMSNorm without learned weight — pure normalization only + # RMSNorm without learned weight — scale is folded into norm weight + # after loading so forward is a single fused norm kernel. self.norm = Gemma4RMSNorm( self.hidden_size, eps=config.rms_norm_eps, with_scale=False ) @@ -165,18 +170,19 @@ class Gemma4Router(nn.Module): quant_config=None, prefix=add_prefix("proj", prefix), ) - self._fused_scale: Optional[torch.Tensor] = None + self._scale_fused = False def fuse_scale(self): - """Pre-compute scale * root_size. Call after weights are loaded.""" - self._fused_scale = (self.scale * self.root_size).to(self.scale.dtype) + """Fold scale * root_size into norm.weight so forward needs no extra mul.""" + fused = (self.scale * self.root_size).to(self.norm.weight.dtype) + self.norm.weight.data.copy_(fused) + self._scale_fused = True def forward(self, x: torch.Tensor) -> torch.Tensor: """Returns raw router logits [T, E].""" - x = self.norm(x) - if self._fused_scale is None: + if not self._scale_fused: self.fuse_scale() - x = x * self._fused_scale.to(x.dtype) + x = self.norm(x) router_logits, _ = self.proj(x) return router_logits @@ -230,13 +236,15 @@ class Gemma4MoE(nn.Module): return gemma4_fused_routing(gating_output, per_expert_scale, topk) topk_logits, topk_ids = torch.topk(gating_output, k=topk, dim=-1) - topk_weights = torch.nn.functional.softmax(topk_logits, dim=-1) - # Fold per_expert_scale into routing weights + # Fused: softmax + per_expert_scale gather + mul + casts in one kernel + if topk_logits.is_cuda or topk_logits.is_xpu: + return gemma_routing_post_topk(topk_logits, topk_ids, per_expert_scale) + + topk_weights = torch.nn.functional.softmax(topk_logits, dim=-1) topk_weights = topk_weights * per_expert_scale[topk_ids].to( topk_weights.dtype ) - return topk_weights.to(torch.float32), topk_ids.to(torch.int32) self.topk = TopK( @@ -407,14 +415,15 @@ class Gemma4Attention(nn.Module): q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) # Fused Q/K/V RMSNorm: replaces three separate norm kernels with one. - # Preconditions for the fused path: tensors on CUDA, q_norm/k_norm use - # the standard norm*weight (scale_shift==0) and v_norm has weight=ones + # Preconditions for the fused path: tensors on CUDA or XPU (the kernel + # is pure Triton and lowers to both backends), q_norm/k_norm use the + # standard norm*weight (scale_shift==0) and v_norm has weight=ones # (with_scale=False) — the canonical Gemma4 attention configuration. is_kv_shared = ( self.is_kv_shared_layer and self.kv_shared_layer_index is not None ) can_fuse_qkv_norm = ( - q.is_cuda + (q.is_cuda or q.is_xpu) and self.q_norm.scale_shift == 0.0 and self.k_norm.scale_shift == 0.0 and not self.v_norm.with_scale @@ -466,9 +475,22 @@ class Gemma4Attention(nn.Module): v = self.v_norm(v) # Apply rotary embedding + use_fused_kv = False if k is not None: k = k.flatten(-2, -1) - q, k = self.rotary_emb(positions, q, k) + # Fuse RoPE + KV-cache write for non-SWA layers with bf16 cache + # DISABLED: causes accuracy regression in launch_server path + can_fuse = False + if can_fuse: + fused_arg = create_fused_set_kv_buffer_arg( + value=v.flatten(-2, -1) if v.dim() == 3 else v, + layer=self.attn, + forward_batch=forward_batch, + ) + use_fused_kv = True + else: + fused_arg = None + q, k = self.rotary_emb(positions, q, k, fused_set_kv_buffer_arg=fused_arg) k = k.unflatten(-1, (self.num_kv_heads, self.head_dim)) else: # Rotary embedding requires a key input; use zeros since KV is shared from another layer @@ -481,7 +503,7 @@ class Gemma4Attention(nn.Module): k, v, forward_batch=forward_batch, - save_kv_cache=not self.is_kv_shared_layer, + save_kv_cache=not self.is_kv_shared_layer and not use_fused_kv, ) if attn_output.dim() == 3: attn_output = attn_output.flatten(-2, -1) @@ -666,7 +688,7 @@ class Gemma4DecoderLayer(nn.Module): # Fused: (rmsnorm(rmsnorm(h1,w1) + rmsnorm(h2,w2), w3) + residual) * scalar if ( not self.has_ple - and hidden_states_1.is_cuda + and (hidden_states_1.is_cuda or hidden_states_1.is_xpu) and hidden_states_1.dim() == 2 ): norm1 = self.post_feedforward_layernorm_1 @@ -698,7 +720,12 @@ class Gemma4DecoderLayer(nn.Module): ) hidden_states = self.mlp(hidden_states) - if not self.has_ple and hidden_states.is_cuda and hidden_states.dim() == 2: + if ( + not self.has_ple + and self.moe is None + and (hidden_states.is_cuda or hidden_states.is_xpu) + and hidden_states.dim() == 2 + ): # Fused: (post_ff_norm(h) + residual) * layer_scalar in one kernel norm = self.post_feedforward_layernorm hidden_states = gemma_rmsnorm_residual_scalar( diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 037349d75..7f2b2f298 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2358,12 +2358,12 @@ class ServerArgs: self.attention_backend = default_attention_backend prefill_backend, decode_backend = self.get_attention_backends() - accepted_backends = ("trtllm_mha", "triton") + accepted_backends = ("trtllm_mha", "triton", "intel_xpu") assert ( prefill_backend in accepted_backends and decode_backend in accepted_backends ), ( - "Gemma4 only supports trtllm_mha or triton attention backend, " + "Gemma4 only supports trtllm_mha, triton, or intel_xpu attention backend, " f"got prefill={prefill_backend}, decode={decode_backend}" ) diff --git a/test/registered/xpu/test_gemma_4_e2b.py b/test/registered/xpu/test_gemma_4_e2b.py new file mode 100644 index 000000000..cd37335bf --- /dev/null +++ b/test/registered/xpu/test_gemma_4_e2b.py @@ -0,0 +1,136 @@ +""" +Gemma 4 E2B-it: simple text Q&A on XPU (OpenAI /v1), same shape as +``test_deepseek_coder_v2_lite_instruct.py``. + +Model card: https://huggingface.co/google/gemma-4-E2B-it + + - XPU test runs when Intel XPU is available. + +Run from test/srt:: + + python3 -m unittest xpu.test_gemma_4_e2b.TestGemma4E2BXPU.test_simple_qa + +A single end-to-end test (``test_simple_qa``) verifies the model boots +on XPU and returns a coherent reply. On failure the assertion message +includes the model's actual output. + +Server is started with ``sglang serve`` (``--model-impl sglang``). +""" + +from __future__ import annotations + +import os +import unittest + +import openai +import torch + +from sglang.srt.utils.common import is_xpu +from sglang.test.test_utils import CustomTestCase +from sglang.test.vlm_utils import ( + DEFAULT_URL_FOR_TEST, + kill_process_tree, + popen_launch_server, +) + +MODEL = "google/gemma-4-E2B-it" + +LAUNCH_TIMEOUT = 900 + +# The -it model ships its own chat_template.jinja, so no --chat-template needed. +# E2B model: single-rank for small model on XPU. +XPU_SERVER_ARGS = [ + "--device", + "xpu", + "--tp=1", + "--trust-remote-code", + "--disable-overlap-schedule", + "--page-size", + "64", + "--mem-fraction-static", + "0.70", + "--attention-backend", + "intel_xpu", + "--model-impl", + "sglang", +] + +# Standard sglang e2e Q&A prompt (see test_openai_server.py::run_chat_completion). +_SIMPLE_QA_PROMPT = "What is the capital of France? Answer in a few words." + + +def _simple_text_messages(): + return [ + { + "role": "user", + "content": [ + {"type": "text", "text": _SIMPLE_QA_PROMPT}, + ], + } + ] + + +def _empty_xpu_cache() -> None: + """Release cached XPU allocations so back-to-back tests start clean.""" + if torch.xpu.is_available(): + torch.xpu.empty_cache() + + +@unittest.skipUnless(is_xpu(), "Intel XPU not available") +class TestGemma4E2BXPU(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = MODEL + cls.base_url = DEFAULT_URL_FOR_TEST + cls.api_key = "sk-123456" + os.environ["SGLANG_USE_SGL_XPU"] = "1" + + _empty_xpu_cache() + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=LAUNCH_TIMEOUT, + api_key=cls.api_key, + other_args=list(XPU_SERVER_ARGS), + device="xpu", + ) + cls.base_url += "/v1" + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + _empty_xpu_cache() + + def test_simple_qa(self): + client = openai.Client(api_key=self.api_key, base_url=self.base_url) + response = client.chat.completions.create( + model="default", + messages=_simple_text_messages(), + temperature=0, + max_tokens=96, + ) + msg = response.choices[0].message + text = msg.content or "" + reasoning = getattr(msg, "reasoning_content", None) or "" + combined = f"{text} {reasoning}".strip() + + self.assertEqual(msg.role, "assistant", f"unexpected role; got: {combined!r}") + self.assertGreater(len(combined), 0, "empty reply from model") + self.assertIn( + "paris", combined.lower(), f"expected `Paris` in reply, got: {combined!r}" + ) + self.assertIsNotNone(response.usage) + self.assertGreater( + response.usage.completion_tokens, + 0, + f"no tokens generated; got: {combined!r}", + ) + + +from sglang.test.ci.ci_register import register_xpu_ci + +# Single e2e test: boot + a short Q&A. +register_xpu_ci(est_time=240, suite="stage-b-test-1-gpu-xpu") + +if __name__ == "__main__": + unittest.main()