diff --git a/python/sglang/srt/eplb/expert_distribution.py b/python/sglang/srt/eplb/expert_distribution.py index 1859852af..30a1f302c 100644 --- a/python/sglang/srt/eplb/expert_distribution.py +++ b/python/sglang/srt/eplb/expert_distribution.py @@ -32,7 +32,7 @@ from sglang.srt.environ import envs from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.observability.metrics_collector import ExpertDispatchCollector from sglang.srt.server_args import ServerArgs -from sglang.srt.utils import Withable, get_int_env_var +from sglang.srt.utils import Withable, get_device, get_int_env_var if TYPE_CHECKING: from sglang.srt.eplb.expert_location import ExpertLocationMetadata @@ -475,6 +475,9 @@ def _list_sum(a: List, b: List) -> List: class _LayerBasedGpuSinglePassGatherer(_SinglePassGatherer): def __init__(self, *args, enable_global_physical_experts: bool, **kwargs): super().__init__(*args, **kwargs) + + device = get_device() + self._enable_global_physical_experts = enable_global_physical_experts self._data = torch.zeros( ( @@ -486,7 +489,7 @@ class _LayerBasedGpuSinglePassGatherer(_SinglePassGatherer): ), ), dtype=torch.int, - device="cuda", + device=device, ) def reset(self): diff --git a/python/sglang/srt/models/llama.py b/python/sglang/srt/models/llama.py index 447f57eaf..dc39732ce 100644 --- a/python/sglang/srt/models/llama.py +++ b/python/sglang/srt/models/llama.py @@ -52,9 +52,12 @@ from sglang.srt.model_loader.weight_utils import ( maybe_remap_kv_scale_name, ) from sglang.srt.server_args import get_global_server_args -from sglang.srt.utils import add_prefix, is_npu, make_layers +from sglang.srt.utils import add_prefix, is_cuda, is_npu, is_xpu, make_layers from sglang.utils import get_exception_traceback +_is_cuda = is_cuda() +_is_xpu = is_xpu() + logger = logging.getLogger(__name__) _is_npu = is_npu() @@ -761,8 +764,12 @@ class LlamaForCausalLM(nn.Module): del self.lm_head.weight self.model.embed_tokens.weight = embed self.lm_head.weight = head - torch.cuda.empty_cache() - torch.cuda.synchronize() + if _is_xpu: + torch.xpu.empty_cache() + torch.xpu.synchronize() + else: + torch.cuda.empty_cache() + torch.cuda.synchronize() def get_embed(self): return self.model.embed_tokens.weight @@ -776,8 +783,12 @@ class LlamaForCausalLM(nn.Module): return del self.model.embed_tokens.weight self.model.embed_tokens.weight = embed - torch.cuda.empty_cache() - torch.cuda.synchronize() + if _is_xpu: + torch.xpu.empty_cache() + torch.xpu.synchronize() + else: + torch.cuda.empty_cache() + torch.cuda.synchronize() def load_kv_cache_scales(self, quantization_param_path: str) -> None: self.model.load_kv_cache_scales(quantization_param_path) diff --git a/test/registered/quant/test_fp8_kernel.py b/test/registered/quant/test_fp8_kernel.py index dcd5ce057..a85841c8c 100644 --- a/test/registered/quant/test_fp8_kernel.py +++ b/test/registered/quant/test_fp8_kernel.py @@ -11,6 +11,13 @@ from sglang.test.test_utils import CustomTestCase register_cuda_ci(est_time=10, suite="stage-b-test-1-gpu-large") +from sglang.srt.utils import get_device, is_cuda, is_xpu + +_is_cuda = is_cuda() +_is_xpu = is_xpu() + +device = get_device() + class TestFP8Base(CustomTestCase): @classmethod @@ -26,7 +33,7 @@ class TestFP8Base(CustomTestCase): @staticmethod def _make_A(M, K, group_size, out_dtype): quant_A = torch.rand( - M, K // group_size, group_size, dtype=torch.float32, device="cuda" + M, K // group_size, group_size, dtype=torch.float32, device=device ) # -1 ~ 1 quant_A = quant_A * 2 - 1 @@ -38,7 +45,7 @@ class TestFP8Base(CustomTestCase): quant_A = quant_A.to(out_dtype).to(torch.float32) # create scale and A - scale = torch.rand(M, K // group_size, dtype=torch.float32, device="cuda") + scale = torch.rand(M, K // group_size, dtype=torch.float32, device=device) scale /= fmax A = quant_A * scale[..., None] @@ -60,7 +67,7 @@ class TestFP8Base(CustomTestCase): N_aligned // group_size, group_size, dtype=torch.float32, - device="cuda", + device=device, ) quant_B = quant_B * 2 - 1 @@ -77,7 +84,7 @@ class TestFP8Base(CustomTestCase): N_aligned // group_size, 1, dtype=torch.float32, - device="cuda", + device=device, ) scale /= fmax @@ -91,8 +98,9 @@ class TestFP8Base(CustomTestCase): class TestPerTokenGroupQuantFP8(TestFP8Base): def test_per_token_group_quant_fp8(self): - if torch.cuda.get_device_capability()[0] < 9: + if _is_cuda and torch.cuda.get_device_capability()[0] < 9: return + A, A_quant_gt, scale_gt = self._make_A( M=self.M, K=self.K, group_size=self.group_size, out_dtype=self.quant_type ) @@ -107,8 +115,14 @@ class TestPerTokenGroupQuantFP8(TestFP8Base): class TestW8A8BlockFP8Matmul(TestFP8Base): def test_w8a8_block_fp8_matmul(self): - if torch.cuda.get_device_capability()[0] < 9: + if _is_cuda and torch.cuda.get_device_capability()[0] < 9: return + elif _is_xpu: + # XPU doesn't provide traditional capability info like CUDA + pass + else: + return + A, A_quant_gt, A_scale_gt = self._make_A( M=self.M, K=self.K, group_size=self.group_size, out_dtype=self.quant_type ) diff --git a/test/registered/vlm/test_vlm_input_format.py b/test/registered/vlm/test_vlm_input_format.py index 1a829ae91..458752eba 100644 --- a/test/registered/vlm/test_vlm_input_format.py +++ b/test/registered/vlm/test_vlm_input_format.py @@ -35,6 +35,7 @@ if not hasattr(_hf_activations, "PytorchGELUTanh"): from sglang import Engine from sglang.srt.entrypoints.openai.protocol import ChatCompletionRequest from sglang.srt.parser.conversation import generate_chat_conv +from sglang.srt.utils.common import is_cuda, is_xpu from sglang.srt.utils.hf_transformers_utils import _fix_added_tokens_encoding register_cuda_ci(est_time=747, suite="stage-b-test-1-gpu-large") @@ -42,6 +43,9 @@ register_cuda_ci(est_time=747, suite="stage-b-test-1-gpu-large") IMAGE_MAN_IRONING_URL = "https://raw.githubusercontent.com/sgl-project/sgl-test-files/refs/heads/main/images/man_ironing_on_back_of_suv.png" IMAGE_SGL_LOGO_URL = "https://raw.githubusercontent.com/sgl-project/sgl-test-files/refs/heads/main/images/sgl_logo.png" +_is_cuda = is_cuda() +_is_xpu = is_xpu() + class VLMInputTestBase: model_path = None @@ -53,12 +57,20 @@ class VLMInputTestBase: def setUpClass(cls): assert cls.model_path is not None, "Set model_path in subclass" assert cls.chat_template is not None, "Set chat_template in subclass" + cls.image_urls = [IMAGE_MAN_IRONING_URL, IMAGE_SGL_LOGO_URL] - cls.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + if _is_cuda: + cls.device = torch.device("cuda") + elif _is_xpu: + cls.device = torch.device("xpu") + else: + cls.device = torch.device("cpu") + cls.main_image = [] for image_url in cls.image_urls: response = requests.get(image_url) cls.main_image.append(Image.open(BytesIO(response.content))) + cls.processor = AutoProcessor.from_pretrained( cls.model_path, trust_remote_code=True, use_fast=True )