[Fix][XPU/ROCm/NPU] Defer sgl_kernel.quantization import in expert_pack (#36529)

test/registered/xpu/test_xpu_classification.py

@@ -4,6 +4,7 @@ Usage:
 python3 -m unittest test_xpu_classification.TestXPUClassification
 """
 
+import gc
 import multiprocessing as mp
 import unittest
 
@@ -11,7 +12,7 @@ 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
+from sglang.test.test_utils import CustomTestCase, empty_gpu_cache
 
 register_xpu_ci(est_time=120, suite="stage-b-test-1-gpu-xpu")
 
@@ -60,6 +61,7 @@ class TestXPUClassification(CustomTestCase):
             model_type="embedding",
             attention_backend="intel_xpu",
             trust_remote_code=True,
+            mem_fraction_static=0.55,
         ) as srt_runner:
             srt_logits = srt_runner.forward(PROMPTS).embed_logits
 
@@ -72,6 +74,8 @@ class TestXPUClassification(CustomTestCase):
 
     def test_classification_logits(self):
         hf_probs = self._hf_probs()
+        gc.collect()
+        empty_gpu_cache()
         srt_probs = self._srt_probs()
 
         self.assertEqual(len(hf_probs), len(PROMPTS))

test/registered/xpu/test_xpu_embedding.py

@@ -5,6 +5,7 @@ Usage:
 python3 -m unittest test_xpu_embedding.TestXPUEmbedding
 """
 
+import gc
 import multiprocessing as mp
 import unittest
 from typing import Optional
@@ -14,7 +15,7 @@ from transformers import AutoConfig, AutoTokenizer
 
 from sglang.test.ci.ci_register import register_xpu_ci
 from sglang.test.runners import DEFAULT_PROMPTS, HFRunner, SRTRunner
-from sglang.test.test_utils import CustomTestCase, get_similarities
+from sglang.test.test_utils import CustomTestCase, empty_gpu_cache, get_similarities
 
 register_xpu_ci(est_time=180, suite="stage-b-test-1-gpu-xpu")
 
@@ -65,12 +66,16 @@ class TestXPUEmbedding(CustomTestCase):
         ) as hf_runner:
             hf_outputs = hf_runner.forward(truncated_prompts)
 
+        gc.collect()
+        empty_gpu_cache()
+
         with SRTRunner(
             model_path,
             tp_size=tp_size,
             torch_dtype=torch_dtype,
             model_type="embedding",
             attention_backend="intel_xpu",
+            mem_fraction_static=0.55,
             json_model_override_args=(
                 {"matryoshka_dimensions": [matryoshka_dim]}
                 if matryoshka_dim is not None

test/registered/xpu/test_xpu_rerank.py

@@ -21,7 +21,7 @@ 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
+from sglang.test.test_utils import CustomTestCase, empty_gpu_cache
 
 
 def _xpu_total_gib() -> float:
@@ -30,13 +30,6 @@ def _xpu_total_gib() -> float:
     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
@@ -214,8 +207,8 @@ 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()
+        gc.collect()
+        empty_gpu_cache()
 
         with SRTRunner(
             model_path,

test/registered/xpu/test_xpu_reward.py

@@ -5,6 +5,7 @@ Usage:
 python3 -m unittest test_xpu_reward.TestXPUReward
 """
 
+import gc
 import multiprocessing as mp
 import unittest
 
@@ -12,7 +13,7 @@ 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
+from sglang.test.test_utils import CustomTestCase, empty_gpu_cache
 
 register_xpu_ci(est_time=60, suite="stage-b-test-1-gpu-xpu")
 
@@ -53,12 +54,16 @@ class TestXPUReward(CustomTestCase):
         ) as hf_runner:
             hf_outputs = hf_runner.forward(convs)
 
+        gc.collect()
+        empty_gpu_cache()
+
         with SRTRunner(
             model_path,
             tp_size=tp_size,
             torch_dtype=torch_dtype,
             model_type="reward",
             attention_backend="intel_xpu",
+            mem_fraction_static=0.55,
         ) as srt_runner:
             prompts = srt_runner.tokenizer.apply_chat_template(
                 convs,