[Intel XPU] Fix cross-encoder rerank hang on B580 runners (#36360)

test/registered/xpu/test_xpu_rerank.py

@@ -10,6 +10,7 @@ python3 -m unittest test_xpu_rerank.TestXPUDecoderRerank
 python3 -m unittest test_xpu_rerank.TestXpuCrossEncoderReank
 """
 
+import gc
 import math
 import multiprocessing as mp
 import unittest
@@ -22,6 +23,24 @@ 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
 
+
+def _xpu_total_gib() -> float:
+    if not torch.xpu.is_available():
+        return 0.0
+    return torch.xpu.get_device_properties(0).total_memory / (1024**3)
+
+
+def _xpu_free_cache() -> None:
+    gc.collect()
+    if torch.xpu.is_available():
+        torch.xpu.empty_cache()
+        torch.xpu.synchronize()
+
+
+# fp32+Triton fits on B60 (22GiB) but hangs on B580 (~12GiB).
+_LARGE_XPU_VRAM_GIB = 20.0
+_HAS_LARGE_XPU = _xpu_total_gib() >= _LARGE_XPU_VRAM_GIB
+
 register_xpu_ci(est_time=180, suite="stage-b-test-1-gpu-xpu")
 
 MODEL_PATH = "Qwen/Qwen3-Reranker-0.6B"
@@ -156,16 +175,23 @@ class TestXPUDecoderRerank(CustomTestCase):
             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.
+# Ported from test/manual/prefill_only/test_cross_encoder_models.py.
+# fp32+triton fits on B60 (22GiB) but OOMs on B580 (~12GiB); bf16+intel_xpu on
+# this encoder model does not match HF (tracked separately), so on <20GiB XPU
+# the class is skipped rather than shipping a knowingly-wrong config.
 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
+CROSS_ENCODER_ATTENTION_BACKEND = "triton"
+CROSS_ENCODER_MEM_FRACTION_STATIC = 0.65
 
 
+@unittest.skipUnless(
+    _HAS_LARGE_XPU,
+    "bge-reranker-v2-m3 fp32+triton OOMs on <20GiB XPU (B580); "
+    "bf16+intel_xpu on this encoder produces wrong scores.",
+)
 class TestXPUCrossEncoderRerank(CustomTestCase):
     @classmethod
     def setUpClass(cls):
@@ -179,6 +205,7 @@ class TestXPUCrossEncoderRerank(CustomTestCase):
         torch_dtype,
         score_tolerance,
         attention_backend,
+        mem_fraction_static,
     ) -> None:
         with HFRunner(
             model_path,
@@ -187,6 +214,9 @@ class TestXPUCrossEncoderRerank(CustomTestCase):
         ) as hf_runner:
             hf_scores = hf_runner.forward(prompts).scores
 
+        # HFRunner leaks a ZMQ context on shutdown; free VRAM before SRT starts.
+        _xpu_free_cache()
+
         with SRTRunner(
             model_path,
             tp_size=tp_size,
@@ -195,6 +225,7 @@ class TestXPUCrossEncoderRerank(CustomTestCase):
             attention_backend=attention_backend,
             chunked_prefill_size=-1,
             disable_radix_cache=True,
+            mem_fraction_static=mem_fraction_static,
         ) as srt_runner:
             srt_scores = srt_runner.forward(prompts).scores
 
@@ -220,6 +251,7 @@ class TestXPUCrossEncoderRerank(CustomTestCase):
                 CROSS_ENCODER_TORCH_DTYPE,
                 CROSS_ENCODER_SCORE_TOLERANCE,
                 CROSS_ENCODER_ATTENTION_BACKEND,
+                CROSS_ENCODER_MEM_FRACTION_STATIC,
             )