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 4, 23:45 UTC
[BugFix][XPU] fix lora ops bgmv_expand size not match (#39989)
Signed-off-by: Ma, Liangliang <liangliang.ma@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Liangliang Ma committed
Apr 20, 2026 at 08:24 UTC
898beca5a8f173e458b4a1f0bac0b2fead0a35a7
1 file changed
+36
-3
vllm/lora/ops/xpu_ops/lora_ops.py
+36
-3
@@ -27,9 +27,42 @@ def bgmv_expand(
27
lora_indices_tensor: torch.Tensor,
28
add_inputs: bool = True,
29
) -> None:
30
- torch.ops._xpu_C.bgmv_expand(
31
- output_tensor, inputs, lora_b_weights, lora_indices_tensor, add_inputs
32
- )
30
+ weight_out_dim = lora_b_weights.size(-2)
31
+ output_dim = output_tensor.size(1)
32
+
33
+ if weight_out_dim == output_dim:
34
+ torch.ops._xpu_C.bgmv_expand(
35
+ output_tensor,
36
+ inputs,
37
+ lora_b_weights,
38
+ lora_indices_tensor,
39
+ add_inputs,
40
+ )
41
+ elif weight_out_dim < output_dim:
42
+ # LoRA weight output dim can be smaller than the output tensor
43
+ # (e.g. vocab_size vs padded logits). Use expand_slice to write
44
+ # only the matching portion, mirroring torch_ops common_len logic.
45
+ torch.ops._xpu_C.bgmv_expand_slice(
46
+ output_tensor,
47
+ inputs,
48
+ lora_b_weights,
49
+ lora_indices_tensor,
50
+ 0,
51
+ weight_out_dim,
52
+ add_inputs,
53
+ )
54
+ else:
55
+ # Weight output dim larger than output tensor: truncate weights.
56
+ lora_b_weights = lora_b_weights[..., :output_dim, :].contiguous()
57
+ torch.ops._xpu_C.bgmv_expand_slice(
58
+ output_tensor,
59
+ inputs,
60
+ lora_b_weights,
61
+ lora_indices_tensor,
62
+ 0,
63
+ output_dim,
64
+ add_inputs,
65
+ )
66
67
68
def bgmv_expand_slice(