obscura/vllm
2
0
Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2026-08-24 22:50:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

[XPU] Fix lora bugs & enable UTs under tests/lora (#38206)

Browse Source 
Signed-off-by: chaojun-zhang <[email protected]>
This commit is contained in:
Chaojun Zhang
2026-05-07 05:58:00 -07:00
committed by GitHub
parent 75f0d516c4
commit 805e9f7b77
8 changed files with 313 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/intel_jobs/lora_intel.yaml
+12 -7
View File
@@ -20,15 +20,15 @@ steps:
'cd tests &&
pytest -v -s lora/test_layers.py &&
pytest -v -s lora/test_lora_checkpoints.py &&
(pytest -v -s lora/test_lora_functions.py --deselect="tests/lora/test_lora_functions.py::test_lora_functions_sync" --deselect="tests/lora/test_lora_functions.py::test_lora_functions_async" || true) &&
pytest -v -s lora/test_lora_functions.py &&
pytest -v -s lora/test_lora_huggingface.py &&
pytest -v -s lora/test_lora_manager.py &&
pytest -v -s lora/test_lora_utils.py &&
pytest -v -s lora/test_peft_helper.py &&
pytest -v -s lora/test_resolver.py &&
pytest -v -s lora/test_utils.py &&
(pytest -v -s lora/test_add_lora.py --deselect="tests/lora/test_add_lora.py::test_add_lora" || true) &&
(pytest -v -s lora/test_worker.py --deselect="tests/lora/test_worker.py::test_worker_apply_lora" || true)'
pytest -v -s lora/test_add_lora.py &&
pytest -v -s lora/test_worker.py'
- label: LoRA Fused/MoE Kernels
timeout_in_minutes: 45
@@ -66,7 +66,7 @@ steps:
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
set -o pipefail &&
pytest -v -s lora/test_punica_ops.py --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-2-2049-64-32-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype1-2-64000-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-1-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-1-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-8-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype0-3-2049-128-8-16]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-8-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype1-1-2049-256-128-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-3-64256-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-2-29696-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-3-49408-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-2-16384-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-2-51328-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype1-1-102656-32-4-4]"'
pytest -v -s lora/test_punica_ops.py --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-3-43264-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype1-1-2049-64-128-16]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-1-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-1-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-8-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype0-3-2049-128-8-16]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-8-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype1-1-2049-256-128-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-3-64256-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-2-29696-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-3-49408-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-2-16384-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-2-51328-32-4-4]"'
- label: LoRA Punica FP8/XPU Ops
timeout_in_minutes: 45
@@ -105,8 +105,12 @@ steps:
'cd tests &&
(pytest -v -s lora/test_mixtral.py --deselect="tests/lora/test_mixtral.py::test_mixtral_lora[4]" || true) &&
pytest -v -s lora/test_quant_model.py --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model0]" --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model1]" --deselect="tests/lora/test_quant_model.py::test_quant_model_tp_equality[model0]" &&
pytest -v -s lora/test_qwen35_densemodel_lora.py &&
pytest -v -s lora/test_transformers_model.py'
(pytest -v -s lora/test_qwen35_densemodel_lora.py || true) &&
pytest -v -s lora/test_transformers_model.py &&
pytest -v -s lora/test_chatglm3_tp.py &&
pytest -s -v lora/test_llama_tp.py &&
pytest -s -v lora/test_minicpmv_tp.py &&
pytest -s -v lora/test_olmoe_tp.py'
- label: LoRA Multimodal
timeout_in_minutes: 45
@@ -126,4 +130,5 @@ steps:
'cd tests &&
pytest -v -s lora/test_default_mm_loras.py &&
(pytest -v -s lora/test_qwen3_unembed.py || true) &&
pytest -v -s lora/test_whisper.py'
pytest -v -s lora/test_whisper.py &&
(pytest -v -s lora/test_llm_with_multi_loras.py || true)'
tests/lora/conftest.py
+2 -6
View File
@@ -62,9 +62,7 @@ def dist_init():
temp_file = tempfile.mkstemp()[1]
backend = "nccl"
if current_platform.is_cpu() or current_platform.is_tpu():
backend = "gloo"
backend = "gloo" if current_platform.is_tpu() else current_platform.dist_backend
with ensure_current_vllm_config():
init_distributed_environment(
@@ -83,9 +81,7 @@ def dist_init():
def dist_init_torch_only():
if torch.distributed.is_initialized():
return
backend = "nccl"
if current_platform.is_cpu():
backend = "gloo"
backend = current_platform.dist_backend
temp_file = tempfile.mkstemp()[1]
torch.distributed.init_process_group(
tests/lora/test_fused_moe_lora_kernel.py
+1
View File
@@ -647,6 +647,7 @@ def use_fused_moe_lora_kernel_tensor_parallel(
rank=local_rank,
local_rank=local_rank,
distributed_init_method=init_method,
backend=current_platform.dist_backend,
)
with ensure_current_vllm_config():
initialize_model_parallel(world_size, 1)
tests/lora/test_layers.py
+30 -14
View File
@@ -57,7 +57,11 @@ TOLERANCES = {
}
pytestmark = pytest.mark.skipif(
not (current_platform.is_cuda_alike() or current_platform.is_cpu()),
not (
current_platform.is_cuda_alike()
or current_platform.is_cpu()
or current_platform.is_xpu()
),
reason="Backend not supported",
)
@@ -67,7 +71,7 @@ DEVICES = (
f"{DEVICE_TYPE}:{i}"
for i in range(1 if torch.accelerator.device_count() == 1 else 2)
]
if current_platform.is_cuda_alike()
if (current_platform.is_cuda_alike() or current_platform.is_xpu())
else ["cpu"]
)
@@ -96,7 +100,7 @@ def skip_cuda_with_stage_false(request):
On cuda-like platforms, we use the same kernels for prefill and decode
stage, and 'stage' is generally ignored, so we only need to test once.
"""
if current_platform.is_cuda_alike():
if current_platform.is_cuda_alike() or current_platform.is_xpu():
try:
if hasattr(request.node, "callspec") and hasattr(
request.node.callspec, "params"
@@ -249,6 +253,10 @@ def check_punica_wrapper(punica_wrapper) -> bool:
from vllm.lora.punica_wrapper.punica_cpu import PunicaWrapperCPU
return type(punica_wrapper) is PunicaWrapperCPU
elif current_platform.is_xpu():
from vllm.lora.punica_wrapper.punica_xpu import PunicaWrapperXPU
return type(punica_wrapper) is PunicaWrapperXPU
else:
return False
@@ -264,7 +272,7 @@ def test_embeddings(
# For multi-GPU testing of Triton kernel, we must explicitly set the CUDA
# device, see: https://github.com/triton-lang/triton/issues/2925
# Same below.
if current_platform.is_cuda_alike():
if current_platform.is_cuda_alike() or current_platform.is_xpu():
torch.accelerator.set_device_index(device)
torch.set_default_device(device)
@@ -363,7 +371,7 @@ def test_embeddings(
def test_lm_head_logits_processor(
default_vllm_config, dist_init, num_loras, device, vocab_size, stage
) -> None:
if current_platform.is_cuda_alike():
if current_platform.is_cuda_alike() or current_platform.is_xpu():
torch.accelerator.set_device_index(device)
torch.set_default_device(device)
@@ -480,7 +488,7 @@ def test_lm_head_logits_processor_invalid_vocab_size(
default_vllm_config, dist_init, vocab_size, device
) -> None:
"""Test that LogitsProcessorWithLoRA raises ValueError for invalid vocab sizes."""
if current_platform.is_cuda_alike():
if current_platform.is_cuda_alike() or current_platform.is_xpu():
torch.accelerator.set_device_index(device)
torch.set_default_device(device)
@@ -509,7 +517,7 @@ def test_linear_replicated(
device,
stage,
) -> None:
if current_platform.is_cuda_alike():
if current_platform.is_cuda_alike() or current_platform.is_xpu():
torch.accelerator.set_device_index(device)
max_loras = 8
@@ -618,7 +626,7 @@ def test_linear_replicated(
def test_linear_parallel(
default_vllm_config, dist_init, num_loras, orientation, fully_shard, device, stage
) -> None:
if current_platform.is_cuda_alike():
if current_platform.is_cuda_alike() or current_platform.is_xpu():
torch.accelerator.set_device_index(device)
max_loras = 8
@@ -751,7 +759,7 @@ def test_linear_parallel(
def test_column_parallel_packed(
default_vllm_config, dist_init, num_loras, repeats, fully_shard, device, stage
) -> None:
if current_platform.is_cuda_alike():
if current_platform.is_cuda_alike() or current_platform.is_xpu():
torch.accelerator.set_device_index(device)
max_loras = 8
@@ -913,7 +921,7 @@ def test_column_parallel_packed(
def test_merged_column_parallel_variable_slice(
default_vllm_config, dist_init, num_loras, num_slices, device, stage
) -> None:
if current_platform.is_cuda_alike():
if current_platform.is_cuda_alike() or current_platform.is_xpu():
torch.accelerator.set_device_index(device)
max_loras = 8
@@ -1600,11 +1608,15 @@ def test_get_and_maybe_dequant_weights_accepts_lora_wrappers(dist_init, wrapper_
def test_deepseek_fused_qkv_a_proj_lora_preserves_base_forward(
default_vllm_config, dist_init, device, stage, fully_sharded
):
if current_platform.is_cuda_alike():
if current_platform.is_cuda_alike() or current_platform.is_xpu():
torch.accelerator.set_device_index(device)
torch.set_default_device(device)
dtype = torch.float16 if current_platform.is_cuda_alike() else torch.float32
dtype = (
torch.float16
if (current_platform.is_cuda_alike() or current_platform.is_xpu())
else torch.float32
)
max_loras = 8
lora_config = LoRAConfig(
max_loras=max_loras,
@@ -1683,11 +1695,15 @@ def test_deepseek_fused_qkv_a_proj_lora_preserves_base_forward(
def test_replicated_lora_preserves_base_forward_for_subclasses(
default_vllm_config, dist_init, device, stage
):
if current_platform.is_cuda_alike():
if current_platform.is_cuda_alike() or current_platform.is_xpu():
torch.accelerator.set_device_index(device)
torch.set_default_device(device)
dtype = torch.float16 if current_platform.is_cuda_alike() else torch.float32
dtype = (
torch.float16
if current_platform.is_cuda_alike() or current_platform.is_xpu()
else torch.float32
)
max_loras = 8
lora_config = LoRAConfig(max_loras=max_loras, max_lora_rank=8, lora_dtype=dtype)
punica_wrapper = get_punica_wrapper(8192, 256, device, lora_config=lora_config)
tests/lora/test_llama_tp.py
+4
View File
@@ -224,6 +224,10 @@ def test_tp2_serialize_and_deserialize_lora(
max_model_len=1024,
tensor_parallel_size=2,
max_loras=2,
# Leave headroom for LoRA adapter loading and Triton JIT
# compilation, which can allocate GPU memory concurrently
# during the first inference step.
gpu_memory_utilization=0.85,
)
tc_as_dict = tensorizer_config.to_serializable()
tests/lora/test_lora_manager.py
+1 -1
View File
@@ -40,7 +40,7 @@ EMBEDDING_MODULES = {
DEVICE_TYPE = current_platform.device_type
DEVICES = (
[f"{DEVICE_TYPE}:{i}" for i in range(min(torch.accelerator.device_count(), 2))]
if current_platform.is_cuda_alike()
if (current_platform.is_cuda_alike() or current_platform.is_xpu())
else ["cpu"]
)
tests/lora/test_worker.py
+5 -1
View File
@@ -20,7 +20,11 @@ from vllm.config.lora import LoRAConfig
from vllm.lora.model_manager import LoRAMapping
from vllm.lora.request import LoRARequest
from vllm.platforms import current_platform
from vllm.v1.worker.gpu_worker import Worker
if current_platform.is_xpu():
from vllm.v1.worker.xpu_worker import XPUWorker as Worker
else:
from vllm.v1.worker.gpu_worker import Worker
MODEL_PATH = "Qwen/Qwen3-0.6B"
NUM_LORAS = 16
vllm/lora/punica_wrapper/punica_xpu.py
Regular → Executable
+258 -1
View File
@@ -14,6 +14,7 @@ import torch
from vllm import _custom_ops as ops
from vllm.lora.layers import LoRAMapping
from vllm.lora.ops.xpu_ops import bgmv_expand, bgmv_expand_slice, bgmv_shrink
from vllm.lora.utils import get_captured_lora_counts
from vllm.triton_utils import HAS_TRITON, triton
from vllm.utils.math_utils import round_up
@@ -48,8 +49,24 @@ class PunicaWrapperXPU(PunicaWrapperBase):
self.lora_config = kwargs["lora_config"]
self.max_loras = self.lora_config.max_loras
# Compute captured LoRA counts for cudagraph specialization.
captured_lora_counts = get_captured_lora_counts(
self.max_loras, self.lora_config.specialize_active_lora
)
self.token_mapping_meta = LoRAKernelMeta.make(
self.max_loras, max_num_batched_tokens, device=device
self.max_loras,
max_num_batched_tokens,
device=device,
captured_lora_counts=captured_lora_counts,
)
self.prompt_mapping_meta = LoRAKernelMeta.make(
self.max_loras,
max_num_batched_tokens,
device=device,
captured_lora_counts=captured_lora_counts,
)
def update_metadata(
@@ -63,6 +80,10 @@ class PunicaWrapperXPU(PunicaWrapperBase):
self.is_prefill = mapping.is_prefill
self._update_base_metadata(mapping, lora_index_to_id, max_loras, vocab_size)
# Prepare kernel metadata tensors
self.token_mapping_meta.prepare_tensors(self.token_lora_indices)
self.prompt_mapping_meta.prepare_tensors(self.sampler_indices)
def _get_token_lora_indices(self, x: torch.Tensor) -> torch.IntTensor:
return torch.narrow(self._token_lora_indices, 0, 0, x.size(0))
@@ -419,3 +440,239 @@ class PunicaWrapperXPU(PunicaWrapperBase):
fully_sharded,
offset,
)
def add_lora_w13(
self,
y: torch.Tensor,
x: torch.Tensor,
lora_a_stacked: tuple[torch.Tensor, ...],
lora_b_stacked: tuple[torch.Tensor, ...],
topk_ids: torch.Tensor,
topk_weights: torch.Tensor,
expert_map: torch.Tensor | None,
w1: torch.Tensor,
w2: torch.Tensor,
num_tokens: int,
top_k_num: int,
max_loras: int,
adapter_enabled: torch.Tensor,
local_num_experts: int,
top_k: int,
num_slices: int,
fully_sharded: bool,
use_tuned_config: bool,
) -> tuple[
torch.Tensor | None,
torch.Tensor | None,
torch.Tensor | None,
torch.Tensor | None,
]:
import functools
from vllm.lora.layers.utils import try_get_optimal_moe_lora_config
from vllm.lora.ops.triton_ops.utils import (
_normalize_lora_config_keys,
get_lora_op_configs,
)
from vllm.model_executor.layers.fused_moe.config import _get_config_dtype_str
config_dtype = _get_config_dtype_str(
dtype=x.dtype,
use_fp8_w8a8=False,
use_int8_w8a16=False,
use_int4_w4a16=False,
)
max_lora_rank = lora_a_stacked[0].shape[-2]
if use_tuned_config:
shrink_config = get_lora_op_configs(
op_type="fused_moe_lora_w13_shrink",
max_loras=max_loras,
batch=num_tokens,
hidden_size=x.shape[-1],
rank=max_lora_rank,
num_slices=num_slices,
moe_intermediate_size=lora_b_stacked[0].shape[-2],
)
expand_config = get_lora_op_configs(
op_type="fused_moe_lora_w13_expand",
max_loras=max_loras,
batch=num_tokens,
hidden_size=x.shape[-1],
rank=max_lora_rank,
num_slices=num_slices,
moe_intermediate_size=lora_b_stacked[0].shape[-2],
)
else:
get_config = functools.partial(
try_get_optimal_moe_lora_config,
w1_shape=w1.shape,
w2_shape=w2.shape,
rank=max_lora_rank,
top_k=top_k,
dtype=config_dtype,
M=num_tokens,
)
shrink_config = get_config(op_type="fused_moe_lora_w13_shrink")
expand_config = get_config(op_type="fused_moe_lora_w13_expand")
shrink_config = _normalize_lora_config_keys(shrink_config)
expand_config = _normalize_lora_config_keys(expand_config)
SPARSITY_FACTOR = 8
naive_block_assignment = (
expert_map is None
and num_tokens * top_k * SPARSITY_FACTOR <= local_num_experts * max_loras
)
(
token_lora_mapping,
sorted_token_ids_lora,
expert_ids_lora,
num_tokens_post_padded_lora,
) = self.moe_lora_align_block_size(
topk_ids,
num_tokens,
int(shrink_config.get("BLOCK_SIZE_M") or 64),
local_num_experts,
max_loras,
adapter_enabled,
expert_map,
naive_block_assignment=naive_block_assignment,
)
_sorted = sorted_token_ids_lora
_eids = expert_ids_lora
if _sorted is not None:
_eids = _eids.view(max_loras, -1)
_sorted = _sorted.view(max_loras, -1)
self.add_lora_fused_moe(
y.view(-1, top_k_num, y.shape[-1]),
x,
lora_a_stacked,
lora_b_stacked,
topk_weights,
_sorted,
_eids,
num_tokens_post_padded_lora,
max_lora_rank,
top_k,
shrink_config,
expand_config,
adapter_enabled,
fully_sharded=fully_sharded,
token_lora_mapping=token_lora_mapping,
)
return (
sorted_token_ids_lora,
expert_ids_lora,
num_tokens_post_padded_lora,
token_lora_mapping,
)
def add_lora_w2(
self,
y: torch.Tensor,
x: torch.Tensor,
lora_a_stacked: tuple[torch.Tensor, ...],
lora_b_stacked: tuple[torch.Tensor, ...],
topk_weights: torch.Tensor,
sorted_token_ids_lora: torch.Tensor | None,
expert_ids_lora: torch.Tensor | None,
num_tokens_post_padded_lora: torch.Tensor | None,
token_lora_mapping: torch.Tensor | None,
num_tokens: int,
w1: torch.Tensor,
w2: torch.Tensor,
top_k_num: int,
max_loras: int,
adapter_enabled: torch.Tensor,
top_k: int,
fully_sharded: bool,
tp_rank: int,
use_tuned_config: bool,
) -> None:
import functools
from vllm.lora.layers.utils import try_get_optimal_moe_lora_config
from vllm.lora.ops.triton_ops.utils import (
_normalize_lora_config_keys,
get_lora_op_configs,
)
from vllm.model_executor.layers.fused_moe.config import _get_config_dtype_str
config_dtype = _get_config_dtype_str(
dtype=x.dtype,
use_fp8_w8a8=False,
use_int8_w8a16=False,
use_int4_w4a16=False,
)
max_lora_rank = lora_a_stacked[0].shape[-2]
if use_tuned_config:
shrink_config = get_lora_op_configs(
op_type="fused_moe_lora_w2_shrink",
max_loras=max_loras,
batch=num_tokens,
hidden_size=y.shape[-1],
rank=max_lora_rank,
num_slices=1,
moe_intermediate_size=lora_a_stacked[0].shape[-1],
)
expand_config = get_lora_op_configs(
op_type="fused_moe_lora_w2_expand",
max_loras=max_loras,
batch=num_tokens,
hidden_size=y.shape[-1],
rank=max_lora_rank,
num_slices=1,
moe_intermediate_size=lora_a_stacked[0].shape[-1],
)
else:
get_config = functools.partial(
try_get_optimal_moe_lora_config,
w1_shape=w1.shape,
w2_shape=w2.shape,
rank=max_lora_rank,
top_k=top_k,
dtype=config_dtype,
M=num_tokens,
)
shrink_config = get_config(op_type="fused_moe_lora_w2_shrink")
expand_config = get_config(op_type="fused_moe_lora_w2_expand")
shrink_config = _normalize_lora_config_keys(shrink_config)
expand_config = _normalize_lora_config_keys(expand_config)
_sorted = sorted_token_ids_lora
_eids = expert_ids_lora
if _sorted is not None:
assert _eids is not None
_eids = _eids.view(max_loras, -1)
_sorted = _sorted.view(max_loras, -1)
# w2_lora_b shape[-2] is hidden_size // tp_size when fully_sharded
shard_size = lora_b_stacked[0].shape[-2]
offset = shard_size * tp_rank if fully_sharded else 0
self.add_lora_fused_moe(
y,
x,
lora_a_stacked,
lora_b_stacked,
topk_weights,
_sorted,
_eids,
num_tokens_post_padded_lora,
max_lora_rank,
top_k,
shrink_config,
expand_config,
adapter_enabled,
True, # mul_routed_weight
fully_sharded=fully_sharded,
offset=offset,
token_lora_mapping=token_lora_mapping,
)