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

Replace hardcoded CUDA device with get_device() for XPU support (#13599)

Browse Source 
Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
This commit is contained in:
Kalyan Kumar
2026-05-01 07:13:46 +08:00
committed by GitHub
co-authored by Ma Mingfei
parent c5f1339773
commit 8a9e424faa
4 changed files with 54 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/eplb/expert_distribution.py
+5 -2
View File
@@ -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):
python/sglang/srt/models/llama.py
+16 -5
View File
@@ -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)
test/registered/quant/test_fp8_kernel.py
+20 -6
View File
@@ -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
)
test/registered/vlm/test_vlm_input_format.py
+13 -1
View File
@@ -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
)