[Intel XPU] support prefill only models for xpu (#35072)

Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>

test/registered/xpu/test_xpu_classification.py

@@ -0,0 +1,106 @@
+"""XPU classification parity test for Qwen2.5-1.5B-apeach.
+
+Usage:
+python3 -m unittest test_xpu_classification.TestXPUClassification
+"""
+
+import multiprocessing as mp
+import unittest
+
+import torch
+
+from sglang.test.ci.ci_register import register_xpu_ci
+from sglang.test.runners import HFRunner, SRTRunner
+from sglang.test.test_utils import CustomTestCase
+
+register_xpu_ci(est_time=120, suite="stage-b-test-1-gpu-xpu")
+
+MODEL_PATH = "jason9693/Qwen2.5-1.5B-apeach"
+TP_SIZE = 1
+TORCH_DTYPE = torch.bfloat16
+# Softmax probabilities are far less sensitive to bf16 rounding than raw
+# logits, so a modest probability tolerance is sufficient here.
+PROB_TOLERANCE = 5e-2
+
+PROMPTS = [
+    "This movie has a tight plot and keeps me engaged.",
+    "Shipping was late and the packaging arrived damaged.",
+    "The features are fine, but the price feels too high.",
+]
+
+
+class TestXPUClassification(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        mp.set_start_method("spawn", force=True)
+
+    def _hf_probs(self):
+        """Reference probabilities from HuggingFace sequence classification."""
+        with HFRunner(
+            MODEL_PATH,
+            torch_dtype=TORCH_DTYPE,
+            model_type="cross_encoder",
+        ) as hf_runner:
+            hf_scores = hf_runner.forward(PROMPTS).scores
+
+        probs = []
+        for row in hf_scores:
+            tensor = row if torch.is_tensor(row) else torch.tensor(row)
+            tensor = tensor.float().flatten()
+            probs.append(torch.softmax(tensor, dim=-1))
+        return probs
+
+    def _srt_probs(self):
+        """Reference probabilities from SRT sequence classification path."""
+        with SRTRunner(
+            MODEL_PATH,
+            tp_size=TP_SIZE,
+            torch_dtype=TORCH_DTYPE,
+            # SRT classify uses embedding-mode encode outputs (class-logit vectors);
+            model_type="embedding",
+            attention_backend="intel_xpu",
+            trust_remote_code=True,
+        ) as srt_runner:
+            srt_logits = srt_runner.forward(PROMPTS).embed_logits
+
+        probs = []
+        for row in srt_logits:
+            tensor = row if torch.is_tensor(row) else torch.tensor(row)
+            tensor = tensor.float().flatten()
+            probs.append(torch.softmax(tensor, dim=-1))
+        return probs
+
+    def test_classification_logits(self):
+        hf_probs = self._hf_probs()
+        srt_probs = self._srt_probs()
+
+        self.assertEqual(len(hf_probs), len(PROMPTS))
+        self.assertEqual(len(srt_probs), len(PROMPTS))
+
+        for index, (hf_row, srt_row) in enumerate(zip(hf_probs, srt_probs)):
+            self.assertEqual(
+                srt_row.shape,
+                hf_row.shape,
+                f"probability shape mismatch at sample {index}",
+            )
+
+            # Probabilities should be close (bf16-tolerant).
+            max_abs_diff = torch.max(torch.abs(hf_row - srt_row)).item()
+            self.assertLess(
+                max_abs_diff,
+                PROB_TOLERANCE,
+                f"classification probs diverged at sample {index}: {max_abs_diff}",
+            )
+
+            # Top-class agreement is the primary correctness signal.
+            hf_pred = int(torch.argmax(hf_row).item())
+            srt_pred = int(torch.argmax(srt_row).item())
+            self.assertEqual(
+                hf_pred,
+                srt_pred,
+                f"top class mismatch at sample {index}",
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/xpu/test_xpu_embedding.py

@@ -1,69 +1,120 @@
 """
-XPU embedding server test: validates the OpenAI-compatible /v1/embeddings
+XPU embedding parity test: compares HF and SRT embedding outputs on Intel XPU.
-endpoint on Intel XPU using a small embedding model. Lives in its own file
-because embedding models load with --is-embedding and use a different model
-than the chat fixtures in test_xpu_serving_features.py.
 
 Usage:
 python3 -m unittest test_xpu_embedding.TestXPUEmbedding
 """
 
+import multiprocessing as mp
 import unittest
+from typing import Optional
 
-import openai
+import torch
+from transformers import AutoConfig, AutoTokenizer
 
-from sglang.srt.utils import kill_process_tree
 from sglang.test.ci.ci_register import register_xpu_ci
-from sglang.test.test_utils import (
+from sglang.test.runners import DEFAULT_PROMPTS, HFRunner, SRTRunner
-    DEFAULT_SMALL_EMBEDDING_MODEL_NAME_FOR_TEST,
+from sglang.test.test_utils import CustomTestCase, get_similarities
-    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-    DEFAULT_URL_FOR_TEST,
-    CustomTestCase,
-    popen_launch_server,
-)
 
-register_xpu_ci(est_time=120, suite="stage-b-test-1-gpu-xpu")
+register_xpu_ci(est_time=180, suite="stage-b-test-1-gpu-xpu")
+
+MODEL_PATH = "Alibaba-NLP/gte-Qwen2-1.5B-instruct"
+TP_SIZE = 1
+PREFILL_TOLERANCE = 1e-3
+TORCH_DTYPE = torch.bfloat16
 
 
 class TestXPUEmbedding(CustomTestCase):
     @classmethod
     def setUpClass(cls):
-        cls.model = DEFAULT_SMALL_EMBEDDING_MODEL_NAME_FOR_TEST
+        mp.set_start_method("spawn", force=True)
-        cls.base_url = DEFAULT_URL_FOR_TEST
+
-        cls.process = popen_launch_server(
+    def _truncate_prompts(self, prompts, model_path):
-            cls.model,
+        config = AutoConfig.from_pretrained(model_path)
-            cls.base_url,
+        max_length = config.to_dict().get("max_position_embeddings", 2048)
-            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+        tokenizer = AutoTokenizer.from_pretrained(model_path)
-            other_args=[
+
-                "--is-embedding",
+        truncated_prompts = []
-                "--device",
+        for prompt in prompts:
-                "xpu",
+            tokens = tokenizer(prompt, return_tensors="pt", truncation=False)
-            ],
+            if len(tokens.input_ids[0]) > max_length:
+                truncated_text = tokenizer.decode(
+                    tokens.input_ids[0][: max_length - 1], skip_special_tokens=True
                 )
-        cls.openai_url = cls.base_url + "/v1"
+                truncated_prompts.append(truncated_text)
+            else:
+                truncated_prompts.append(prompt)
+        return truncated_prompts
 
-    @classmethod
+    def assert_close_prefill_logits(
-    def tearDownClass(cls):
+        self,
-        kill_process_tree(cls.process.pid)
+        prompts,
+        model_path,
+        tp_size,
+        torch_dtype,
+        prefill_tolerance,
+        matryoshka_dim: Optional[int] = None,
+    ) -> None:
+        truncated_prompts = self._truncate_prompts(prompts, model_path)
 
-    def _client(self) -> openai.Client:
+        with HFRunner(
-        # Server has no API key, but openai client still requires a non-empty string.
+            model_path,
-        return openai.Client(api_key="EMPTY", base_url=self.openai_url)
+            torch_dtype=torch_dtype,
+            model_type="embedding",
+            matryoshka_dim=matryoshka_dim,
+        ) as hf_runner:
+            hf_outputs = hf_runner.forward(truncated_prompts)
 
-    def test_embedding_single(self):
+        with SRTRunner(
-        response = self._client().embeddings.create(
+            model_path,
-            model=self.model, input="Hello world"
+            tp_size=tp_size,
+            torch_dtype=torch_dtype,
+            model_type="embedding",
+            attention_backend="intel_xpu",
+            json_model_override_args=(
+                {"matryoshka_dimensions": [matryoshka_dim]}
+                if matryoshka_dim is not None
+                else None
+            ),
+        ) as srt_runner:
+            srt_outputs = srt_runner.forward(
+                truncated_prompts,
+                dimensions=matryoshka_dim,
             )
-        self.assertEqual(len(response.data), 1)
-        self.assertGreater(len(response.data[0].embedding), 0)
 
-    def test_embedding_batch(self):
+        for prompt, hf_output, srt_output in zip(
-        response = self._client().embeddings.create(
+            prompts,
-            model=self.model, input=["Hello world", "Test text"]
+            hf_outputs.embed_logits,
+            srt_outputs.embed_logits,
+        ):
+            hf_logits = torch.Tensor(hf_output)
+            srt_logits = torch.Tensor(srt_output)
+
+            similarity = torch.tensor(get_similarities(hf_logits, srt_logits))
+            if len(prompt) <= 1000:
+                self.assertTrue(
+                    torch.all(torch.abs(similarity - 1) < prefill_tolerance),
+                    "embeddings are not all close",
+                )
+
+    def test_prefill_logits(self):
+        self.assert_close_prefill_logits(
+            DEFAULT_PROMPTS,
+            MODEL_PATH,
+            TP_SIZE,
+            TORCH_DTYPE,
+            PREFILL_TOLERANCE,
+        )
+
+    def test_matryoshka_embedding(self):
+        self.assert_close_prefill_logits(
+            DEFAULT_PROMPTS,
+            MODEL_PATH,
+            TP_SIZE,
+            TORCH_DTYPE,
+            PREFILL_TOLERANCE,
+            matryoshka_dim=128,
         )
-        self.assertEqual(len(response.data), 2)
-        self.assertGreater(len(response.data[0].embedding), 0)
-        self.assertGreater(len(response.data[1].embedding), 0)
 
 
 if __name__ == "__main__":

test/registered/xpu/test_xpu_rerank.py

@@ -0,0 +1,227 @@
+"""XPU rerank test suite.
+
+This file validates score parity between HuggingFace and SRT for two rerank
+serving styles:
+- Decoder-only reranker scoring (Qwen3-Reranker style).
+- Cross-encoder scoring (BAAI/bge-reranker-v2-m3).
+
+Usage:
+python3 -m unittest test_xpu_rerank.TestXPUDecoderRerank
+python3 -m unittest test_xpu_rerank.TestXpuCrossEncoderReank
+"""
+
+import math
+import multiprocessing as mp
+import unittest
+
+import torch
+from jinja2.sandbox import ImmutableSandboxedEnvironment
+
+from sglang.srt.utils.hf_transformers_utils import get_tokenizer
+from sglang.test.ci.ci_register import register_xpu_ci
+from sglang.test.runners import TEST_RERANK_QUERY_DOCS, HFRunner, SRTRunner
+from sglang.test.test_utils import CustomTestCase
+
+register_xpu_ci(est_time=180, suite="stage-b-test-1-gpu-xpu")
+
+MODEL_PATH = "Qwen/Qwen3-Reranker-0.6B"
+TP_SIZE = 1
+SCORE_TOLERANCE = 1e-2
+ATTENTION_BACKEND = "intel_xpu"
+TORCH_DTYPE = torch.bfloat16
+# Prompt template mirrored from examples/chat_template/qwen3_reranker.jinja.
+QWEN3_RERANKER_TEMPLATE = r"""<|im_start|>system
+Judge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be "yes" or "no".<|im_end|>
+<|im_start|>user
+<Instruct>: {{ instruct | default("Given a web search query, retrieve relevant passages that answer the query.") }}
+<Query>: {{ messages[0]["content"] }}
+<Document>: {{ messages[1]["content"] }}<|im_end|>
+<|im_start|>assistant{{ '\\n' }}
+"""
+
+JINJA_ENV = ImmutableSandboxedEnvironment(autoescape=False)
+QWEN3_RERANKER_JINJA = JINJA_ENV.from_string(QWEN3_RERANKER_TEMPLATE)
+
+# Small decoder-reranker dataset (from the Qwen3-Reranker cookbook style).
+# The documents intentionally include clear relevant/irrelevant contrast.
+
+RERANK_QUERY_DOCS = [
+    {
+        "query": "法国首都是哪里？",
+        "instruct": "Given a web search query, retrieve relevant passages that answer the query.",
+        "documents": [
+            "法国的首都是巴黎。",
+            "德国的首都是柏林。",
+            "香蕉是黄色的水果。",
+        ],
+    },
+]
+
+
+def format_prompt(query: str, document: str, instruct: str) -> str:
+    """Render the canonical Qwen3 reranker Jinja template used by serving."""
+    render_kwargs = {
+        "messages": [
+            {"role": "user", "content": query},
+            {"role": "user", "content": document},
+        ]
+    }
+    if instruct:
+        render_kwargs["instruct"] = instruct
+    return QWEN3_RERANKER_JINJA.render(**render_kwargs)
+
+
+def yes_no_token_ids(tokenizer) -> tuple[int, int]:
+    yes = tokenizer.encode("yes", add_special_tokens=False)
+    no = tokenizer.encode("no", add_special_tokens=False)
+    assert len(yes) == 1 and len(no) == 1, "yes/no must be single tokens"
+    return yes[0], no[0]
+
+
+def score_from_token_logprobs(logprob_yes: float, logprob_no: float) -> float:
+    """score = P(yes) / (P(yes) + P(no))."""
+    p_yes = math.exp(logprob_yes)
+    p_no = math.exp(logprob_no)
+    denom = p_yes + p_no
+    return p_yes / denom if denom > 0.0 else 0.0
+
+
+class TestXPUDecoderRerank(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        mp.set_start_method("spawn", force=True)
+        cls.tokenizer = get_tokenizer(MODEL_PATH)
+        cls.yes_id, cls.no_id = yes_no_token_ids(cls.tokenizer)
+
+    def _extract_scores(self, token_ids_output_logprobs) -> list[float]:
+        """token_ids_output_logprobs shape: [num_prompts][num_gen_tokens][num_token_ids].
+
+        We only generate 1 token and request exactly [yes_id, no_id], so we read
+        index [0] (first/only generated token) -> [yes_lp, no_lp].
+        """
+        scores = []
+        for per_prompt in token_ids_output_logprobs:
+            first_token_lps = per_prompt[0]  # logprobs for [yes_id, no_id]
+            yes_lp, no_lp = first_token_lps[0], first_token_lps[1]
+            scores.append(score_from_token_logprobs(yes_lp, no_lp))
+        return scores
+
+    def _assert_close_scores(self, prompts) -> None:
+        token_ids_logprob = [self.yes_id, self.no_id]
+
+        # --- HuggingFace reference (generation) ---
+        with HFRunner(
+            MODEL_PATH,
+            torch_dtype=TORCH_DTYPE,
+            model_type="generation",
+            output_str_only=False,
+        ) as hf_runner:
+            hf_out = hf_runner.forward(
+                prompts,
+                max_new_tokens=1,
+                token_ids_logprob=token_ids_logprob,
+            )
+        hf_scores = self._extract_scores(hf_out.token_ids_output_logprobs)
+
+        with SRTRunner(
+            MODEL_PATH,
+            tp_size=TP_SIZE,
+            torch_dtype=TORCH_DTYPE,
+            model_type="generation",
+            attention_backend=ATTENTION_BACKEND,
+        ) as srt_runner:
+            srt_out = srt_runner.forward(
+                prompts,
+                max_new_tokens=1,
+                token_ids_logprob=token_ids_logprob,
+            )
+        srt_scores = self._extract_scores(srt_out.token_ids_output_logprobs)
+
+        self.assertEqual(len(hf_scores), len(srt_scores))
+        for hf_score, srt_score in zip(hf_scores, srt_scores):
+            self.assertLess(
+                abs(hf_score - srt_score),
+                SCORE_TOLERANCE,
+                "decoder rerank scores are not all close",
+            )
+
+    def _preprocess_prompts(self, query_doc) -> list[str]:
+        query = query_doc["query"]
+        instruct = query_doc["instruct"]
+        return [format_prompt(query, doc, instruct) for doc in query_doc["documents"]]
+
+    def test_prefill_logits(self):
+        for query_doc in RERANK_QUERY_DOCS:
+            prompts = self._preprocess_prompts(query_doc)
+            self._assert_close_scores(prompts)
+
+
+# This cross-encoder test is ported from `test/manual/prefill_only/test_cross_encoder_models.py`,
+# which uses float32 with the triton backend. The `intel_xpu` attention backend currently only
+# supports the bfloat16 dtype, so we keep the triton backend here to preserve float32 parity.
+CROSS_ENCODER_MODEL_PATH = "BAAI/bge-reranker-v2-m3"
+CROSS_ENCODER_TP_SIZE = 1
+CROSS_ENCODER_SCORE_TOLERANCE = 1e-2
+CROSS_ENCODER_ATTENTION_BACKEND = "triton"
+CROSS_ENCODER_TORCH_DTYPE = torch.float32
+
+
+class TestXPUCrossEncoderRerank(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        mp.set_start_method("spawn", force=True)
+
+    def _assert_close_scores(
+        self,
+        prompts,
+        model_path,
+        tp_size,
+        torch_dtype,
+        score_tolerance,
+        attention_backend,
+    ) -> None:
+        with HFRunner(
+            model_path,
+            torch_dtype=torch_dtype,
+            model_type="cross_encoder",
+        ) as hf_runner:
+            hf_scores = hf_runner.forward(prompts).scores
+
+        with SRTRunner(
+            model_path,
+            tp_size=tp_size,
+            torch_dtype=torch_dtype,
+            model_type="cross_encoder",
+            attention_backend=attention_backend,
+            chunked_prefill_size=-1,
+            disable_radix_cache=True,
+        ) as srt_runner:
+            srt_scores = srt_runner.forward(prompts).scores
+
+        self.assertEqual(len(hf_scores), len(srt_scores))
+        for hf_score, srt_score in zip(hf_scores, srt_scores):
+            self.assertLess(
+                abs(hf_score - srt_score),
+                score_tolerance,
+                "cross encoder scores are not all close",
+            )
+
+    def _preprocess_prompts(self, query_doc):
+        query = query_doc["query"]
+        return [[query, document] for document in query_doc["documents"]]
+
+    def test_prefill_logits(self):
+        for query_doc in TEST_RERANK_QUERY_DOCS:
+            prompts = self._preprocess_prompts(query_doc)
+            self._assert_close_scores(
+                prompts,
+                CROSS_ENCODER_MODEL_PATH,
+                CROSS_ENCODER_TP_SIZE,
+                CROSS_ENCODER_TORCH_DTYPE,
+                CROSS_ENCODER_SCORE_TOLERANCE,
+                CROSS_ENCODER_ATTENTION_BACKEND,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/xpu/test_xpu_reward.py

@@ -0,0 +1,88 @@
+"""
+XPU reward parity test: compares HF and SRT reward scores on Intel XPU.
+
+Usage:
+python3 -m unittest test_xpu_reward.TestXPUReward
+"""
+
+import multiprocessing as mp
+import unittest
+
+import torch
+
+from sglang.test.ci.ci_register import register_xpu_ci
+from sglang.test.runners import HFRunner, SRTRunner
+from sglang.test.test_utils import CustomTestCase
+
+register_xpu_ci(est_time=60, suite="stage-b-test-1-gpu-xpu")
+
+MODEL_PATH = "Skywork/Skywork-Reward-V2-Qwen3-0.6B"
+TP_SIZE = 1
+TOLERANCE = 1.5e-1
+TORCH_DTYPE = torch.bfloat16
+
+PROMPT = (
+    "What is the range of the numeric output of a sigmoid node in a neural network?"
+)
+RESPONSE1 = "The output of a sigmoid node is bounded between -1 and 1."
+RESPONSE2 = "The output of a sigmoid node is bounded between 0 and 1."
+
+CONVS = [
+    [{"role": "user", "content": PROMPT}, {"role": "assistant", "content": RESPONSE1}],
+    [{"role": "user", "content": PROMPT}, {"role": "assistant", "content": RESPONSE2}],
+]
+
+
+class TestXPUReward(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        mp.set_start_method("spawn", force=True)
+
+    def assert_close_reward_scores(
+        self,
+        convs,
+        model_path,
+        tp_size,
+        torch_dtype,
+        tolerance,
+    ) -> None:
+        with HFRunner(
+            model_path,
+            torch_dtype=torch_dtype,
+            model_type="reward",
+        ) as hf_runner:
+            hf_outputs = hf_runner.forward(convs)
+
+        with SRTRunner(
+            model_path,
+            tp_size=tp_size,
+            torch_dtype=torch_dtype,
+            model_type="reward",
+            attention_backend="intel_xpu",
+        ) as srt_runner:
+            prompts = srt_runner.tokenizer.apply_chat_template(
+                convs,
+                tokenize=False,
+                return_dict=False,
+            )
+            srt_outputs = srt_runner.forward(prompts)
+
+        hf_scores = torch.tensor(hf_outputs.scores)
+        srt_scores = torch.tensor(srt_outputs.scores)
+        self.assertTrue(
+            torch.all(torch.abs(hf_scores - srt_scores) < tolerance),
+            "reward scores are not all close",
+        )
+
+    def test_reward_scores(self):
+        self.assert_close_reward_scores(
+            CONVS,
+            MODEL_PATH,
+            TP_SIZE,
+            TORCH_DTYPE,
+            TOLERANCE,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()