A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 03:30 UTC
[xpu]feat: support multi-lora on xpu (#20616)
Signed-off-by: yan <yan.ma@intel.com>
Yan Ma committed
Jul 8, 2025 at 22:07 UTC
a4c23314c04a0ce3e507cd199d6372fb83cb6732
5 files changed
+28
-4
vllm/lora/ops/triton_ops/lora_expand_op.py
+2
@@ -13,6 +13,7 @@ import triton.language as tl
13
14
from vllm.lora.ops.triton_ops.kernel_utils import do_expand_kernel
15
from vllm.lora.ops.triton_ops.utils import _get_lora_b_ptr
16
+from vllm.platforms import current_platform
17
from vllm.utils import direct_register_custom_op
18
19
@@ -283,6 +284,7 @@ try:
284
op_func=_lora_expand,
285
mutates_args=["output_tensor"],
286
fake_impl=_lora_expand_fake,
287
+ dispatch_key=current_platform.dispatch_key,
288
)
289
lora_expand = torch.ops.vllm.lora_expand
290
vllm/lora/ops/triton_ops/lora_shrink_op.py
+2
@@ -13,6 +13,7 @@ import triton.language as tl
13
14
from vllm.lora.ops.triton_ops.kernel_utils import do_shrink_kernel
15
from vllm.lora.ops.triton_ops.utils import _get_lora_a_ptr
16
+from vllm.platforms import current_platform
17
from vllm.utils import direct_register_custom_op
18
19
@@ -237,6 +238,7 @@ try:
238
op_func=_lora_shrink,
239
mutates_args=["output_tensor"],
240
fake_impl=_lora_shrink_fake,
241
+ dispatch_key=current_platform.dispatch_key,
242
)
243
lora_shrink = torch.ops.vllm.lora_shrink
244
vllm/lora/ops/triton_ops/utils.py
+9
-3
@@ -35,7 +35,9 @@ def _get_lora_a_ptr(lora_a_weights: list[torch.Tensor], device: torch.device):
35
lora_strides_d1.append(lora_a_weight.stride(1))
36
lora_strides_d2.append(lora_a_weight.stride(2))
37
if len(lora_a_weights) > 1:
38
- lora_ptr_tensor = torch.tensor(tensor_ptrs, device=device)
38
+ lora_ptr_tensor = torch.tensor(tensor_ptrs,
39
+ device=device,
40
+ dtype=torch.uint64)
41
else:
42
lora_ptr_tensor = lora_a_weights[0]
43
@@ -89,8 +91,12 @@ def _get_lora_b_ptr(lora_weights: list[torch.Tensor], offset_start: int,
91
92
if len(lora_weights) > 1:
93
# note these are device tensors
92
- lora_ptr_tensor = torch.tensor(tensor_ptrs, device=device)
93
- slice_start_tensor = torch.tensor(slice_offset_lst, device=device)
94
+ lora_ptr_tensor = torch.tensor(tensor_ptrs,
95
+ device=device,
96
+ dtype=torch.uint64)
97
+ slice_start_tensor = torch.tensor(slice_offset_lst,
98
+ device=device,
99
+ dtype=torch.uint64)
100
else:
101
slice_start_tensor = slice_offset_lst[0]
102
lora_ptr_tensor = lora_b_weight[0]
vllm/model_executor/model_loader/tensorizer.py
+4
-1
@@ -27,6 +27,7 @@ from vllm.config import (ModelConfig, ParallelConfig, VllmConfig,
27
from vllm.logger import init_logger
28
from vllm.model_executor.layers.vocab_parallel_embedding import (
29
VocabParallelEmbedding)
30
+from vllm.platforms import current_platform
31
from vllm.utils import FlexibleArgumentParser, PlaceholderModule
32
33
if TYPE_CHECKING:
@@ -513,7 +514,9 @@ def deserialize_tensorizer_model(model: nn.Module,
514
**tensorizer_args.stream_kwargs) as stream, TensorDeserializer(
515
stream,
516
dtype=tensorizer_config.dtype,
516
- device=torch.device("cuda", torch.cuda.current_device()),
517
+ device=f'xpu:{torch.xpu.current_device()}'
518
+ if current_platform.is_xpu() else
519
+ f'cuda:{torch.cuda.current_device()}',
520
**tensorizer_args.deserialization_kwargs) as deserializer:
521
deserializer.load_into_module(model)
522
end = time.perf_counter()
vllm/platforms/xpu.py
+11
@@ -58,6 +58,10 @@ class XPUPlatform(Platform):
58
def get_device_name(cls, device_id: int = 0) -> str:
59
return torch.xpu.get_device_name(device_id)
60
61
+ @classmethod
62
+ def get_punica_wrapper(cls) -> str:
63
+ return "vllm.lora.punica_wrapper.punica_gpu.PunicaWrapperGPU"
64
+
65
@classmethod
66
def get_device_total_memory(cls, device_id: int = 0) -> int:
67
device_props = torch.xpu.get_device_properties(device_id)
@@ -78,6 +82,13 @@ class XPUPlatform(Platform):
82
if cache_config and cache_config.block_size is None:
83
cache_config.block_size = 64
84
85
+ # FIXME: Temporarily forcing eager mode
86
+ # remove after t.compile support stabilizes.
87
+ if (envs.VLLM_USE_V1 and vllm_config.model_config is not None
88
+ and not vllm_config.model_config.enforce_eager):
89
+ from vllm.config import CompilationLevel
90
+ vllm_config.compilation_config.level = CompilationLevel.NO_COMPILATION # noqa: E501
91
+
92
# Instances created using VllmConfig() typically have model_config as
93
# None by default. The modification involves adding a check to prevent
94
# potential null exceptions check and update model config.