minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
This commit is contained in:
gaopengff
2026-08-24 16:25:47 +08:00
committed by GitHub
co-authored by Ma Mingfei
parent 7de80e566c
commit 317da0964e
4 changed files with 519 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/xpu/test_xpu_classification.py
+106
View File
@@ -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
+98 -47
View File
@@ -1,69 +1,120 @@
"""
XPU embedding server test: validates the OpenAI-compatible /v1/embeddings
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.
XPU embedding parity test: compares HF and SRT embedding outputs on Intel XPU.
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 (
DEFAULT_SMALL_EMBEDDING_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
from sglang.test.runners import DEFAULT_PROMPTS, HFRunner, SRTRunner
from sglang.test.test_utils import CustomTestCase, get_similarities
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
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--is-embedding",
"--device",
"xpu",
],
mp.set_start_method("spawn", force=True)
def _truncate_prompts(self, prompts, model_path):
config = AutoConfig.from_pretrained(model_path)
max_length = config.to_dict().get("max_position_embeddings", 2048)
tokenizer = AutoTokenizer.from_pretrained(model_path)
truncated_prompts = []
for prompt in prompts:
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
)
truncated_prompts.append(truncated_text)
else:
truncated_prompts.append(prompt)
return truncated_prompts
def assert_close_prefill_logits(
self,
prompts,
model_path,
tp_size,
torch_dtype,
prefill_tolerance,
matryoshka_dim: Optional[int] = None,
) -> None:
truncated_prompts = self._truncate_prompts(prompts, model_path)
with HFRunner(
model_path,
torch_dtype=torch_dtype,
model_type="embedding",
matryoshka_dim=matryoshka_dim,
) as hf_runner:
hf_outputs = hf_runner.forward(truncated_prompts)
with SRTRunner(
model_path,
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,
)
for prompt, hf_output, srt_output in zip(
prompts,
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,
)
cls.openai_url = cls.base_url + "/v1"
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def _client(self) -> openai.Client:
# Server has no API key, but openai client still requires a non-empty string.
return openai.Client(api_key="EMPTY", base_url=self.openai_url)
def test_embedding_single(self):
response = self._client().embeddings.create(
model=self.model, input="Hello world"
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), 1)
self.assertGreater(len(response.data[0].embedding), 0)
def test_embedding_batch(self):
response = self._client().embeddings.create(
model=self.model, input=["Hello world", "Test text"]
)
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
+227
View File
@@ -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
+88
View File
@@ -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()