[xpu][lora]: Align LoRA implementation with Punica GPU: fix _apply_expand rank mismatch, add_inputs hardcode, and MoE EP (#45368)
Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com>
Chaojun Zhang committed
Jul 1, 2026 at 22:14 UTC
13c49f9845d5f64c53a75d7b5c44c3997c4f6ee6
3 files changed
+182
-35
.buildkite/intel_jobs/lora_intel.yaml
+1
-1
@@ -81,7 +81,7 @@ steps:
81
'cd tests &&
82
export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
83
set -o pipefail &&
84
- 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]"'
84
+ 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]" --deselect="tests/lora/test_kernels_hidden_size[shrink-0-xpu:0-dtype0-3-32000-32-4-4]" --deselect="tests/lora/test_kernels_hidden_size[shrink-0-xpu:0-dtype0-3-32000-32-4-4]"'
85
86
- label: LoRA Punica FP8/XPU Ops
87
timeout_in_minutes: 45
vllm/lora/ops/xpu_ops/lora_ops.py
+105
-9
@@ -3,29 +3,27 @@
3
4
import torch
5
6
-from vllm.logger import init_logger
6
+from vllm.utils.torch_utils import direct_register_custom_op
7
8
-logger = init_logger(__name__)
8
10
-
11
-def bgmv_shrink(
9
+def _bgmv_shrink_impl(
10
inputs: torch.Tensor,
11
lora_a_weights: torch.Tensor,
12
output_tensor: torch.Tensor,
13
lora_indices_tensor: torch.Tensor,
16
- scaling: float = 1.0,
14
+ scaling: float,
15
) -> None:
16
torch.ops._xpu_C.bgmv_shrink(
17
output_tensor, inputs, lora_a_weights, lora_indices_tensor, scaling
18
)
19
20
23
-def bgmv_expand(
21
+def _bgmv_expand_impl(
22
inputs: torch.Tensor,
23
lora_b_weights: torch.Tensor,
24
output_tensor: torch.Tensor,
25
lora_indices_tensor: torch.Tensor,
28
- add_inputs: bool = True,
26
+ add_inputs: bool,
27
) -> None:
28
weight_out_dim = lora_b_weights.size(-2)
29
output_dim = output_tensor.size(1)
@@ -65,14 +63,14 @@ def bgmv_expand(
63
)
64
65
68
-def bgmv_expand_slice(
66
+def _bgmv_expand_slice_impl(
67
inputs: torch.Tensor,
68
lora_b_weights: torch.Tensor,
69
output_tensor: torch.Tensor,
70
lora_indices_tensor: torch.Tensor,
71
slice_offset: int,
72
slice_size: int,
75
- add_inputs: bool = True,
73
+ add_inputs: bool,
74
) -> None:
75
assert slice_size == lora_b_weights.size(-2)
76
assert slice_offset + slice_size <= output_tensor.size(1)
@@ -85,3 +83,101 @@ def bgmv_expand_slice(
83
slice_size,
84
add_inputs,
85
)
86
+
87
+
88
+def _bgmv_shrink_fake(
89
+ inputs: torch.Tensor,
90
+ lora_a_weights: torch.Tensor,
91
+ output_tensor: torch.Tensor,
92
+ lora_indices_tensor: torch.Tensor,
93
+ scaling: float,
94
+) -> None:
95
+ return None
96
+
97
+
98
+def _bgmv_expand_fake(
99
+ inputs: torch.Tensor,
100
+ lora_b_weights: torch.Tensor,
101
+ output_tensor: torch.Tensor,
102
+ lora_indices_tensor: torch.Tensor,
103
+ add_inputs: bool,
104
+) -> None:
105
+ return None
106
+
107
+
108
+def _bgmv_expand_slice_fake(
109
+ inputs: torch.Tensor,
110
+ lora_b_weights: torch.Tensor,
111
+ output_tensor: torch.Tensor,
112
+ lora_indices_tensor: torch.Tensor,
113
+ slice_offset: int,
114
+ slice_size: int,
115
+ add_inputs: bool,
116
+) -> None:
117
+ return None
118
+
119
+
120
+direct_register_custom_op(
121
+ op_name="xpu_bgmv_shrink",
122
+ op_func=_bgmv_shrink_impl,
123
+ mutates_args=["output_tensor"],
124
+ fake_impl=_bgmv_shrink_fake,
125
+)
126
+
127
+direct_register_custom_op(
128
+ op_name="xpu_bgmv_expand",
129
+ op_func=_bgmv_expand_impl,
130
+ mutates_args=["output_tensor"],
131
+ fake_impl=_bgmv_expand_fake,
132
+)
133
+
134
+direct_register_custom_op(
135
+ op_name="xpu_bgmv_expand_slice",
136
+ op_func=_bgmv_expand_slice_impl,
137
+ mutates_args=["output_tensor"],
138
+ fake_impl=_bgmv_expand_slice_fake,
139
+)
140
+
141
+
142
+def bgmv_shrink(
143
+ inputs: torch.Tensor,
144
+ lora_a_weights: torch.Tensor,
145
+ output_tensor: torch.Tensor,
146
+ lora_indices_tensor: torch.Tensor,
147
+ scaling: float = 1.0,
148
+) -> None:
149
+ torch.ops.vllm.xpu_bgmv_shrink(
150
+ inputs, lora_a_weights, output_tensor, lora_indices_tensor, scaling
151
+ )
152
+
153
+
154
+def bgmv_expand(
155
+ inputs: torch.Tensor,
156
+ lora_b_weights: torch.Tensor,
157
+ output_tensor: torch.Tensor,
158
+ lora_indices_tensor: torch.Tensor,
159
+ add_inputs: bool = True,
160
+) -> None:
161
+ torch.ops.vllm.xpu_bgmv_expand(
162
+ inputs, lora_b_weights, output_tensor, lora_indices_tensor, add_inputs
163
+ )
164
+
165
+
166
+def bgmv_expand_slice(
167
+ inputs: torch.Tensor,
168
+ lora_b_weights: torch.Tensor,
169
+ output_tensor: torch.Tensor,
170
+ lora_indices_tensor: torch.Tensor,
171
+ slice_offset: int,
172
+ slice_size: int,
173
+ add_inputs: bool = True,
174
+) -> None:
175
+ torch.ops.vllm.xpu_bgmv_expand_slice(
176
+ inputs,
177
+ lora_b_weights,
178
+ output_tensor,
179
+ lora_indices_tensor,
180
+ slice_offset,
181
+ slice_size,
182
+ add_inputs,
183
+ )
vllm/lora/punica_wrapper/punica_xpu.py
+76
-25
@@ -62,6 +62,10 @@ class PunicaWrapperXPU(PunicaWrapperBase):
62
captured_lora_counts=captured_lora_counts,
63
)
64
65
+ # When speculative decoding is enabled, max_num_samples is
66
+ # max_batches * (num_speculative_decoding_tokens + 1).
67
+ # This line can be optimized by replacing max_num_batched_tokens
68
+ # to max_batches * (num_speculative_decoding_tokens + 1).
69
self.prompt_mapping_meta = LoRAKernelMeta.make(
70
self.max_loras,
71
max_num_batched_tokens,
@@ -106,6 +110,14 @@ class PunicaWrapperXPU(PunicaWrapperBase):
110
add_inputs: bool,
111
):
112
token_lora_indices = self._get_token_lora_indices(x)
113
+ # After tensor-parallel all-gather (non-fully-sharded LoRA), x may
114
+ # have been gathered along the rank dim so x.size(1) == max_lora_rank
115
+ # * tp_size, while lora_b only uses max_lora_rank elements. The XPU
116
+ # C++ kernel requires inputs.size(1) == lora_b.size(-1), so truncate
117
+ # to the actual rank. x[:, :rank] is non-contiguous, hence the copy.
118
+ rank = w_t_all.size(-1)
119
+ if x.size(1) != rank:
120
+ x = x[:, :rank].contiguous()
121
bgmv_expand_slice(
122
x, w_t_all, y, token_lora_indices, y_offset, y_slice_size, add_inputs
123
)
@@ -179,7 +191,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):
191
add_inputs=add_inputs,
192
)
193
offset_start += output_slices[slice_idx]
182
- y.view_as(y_org)
194
+ y = y.view_as(y_org)
195
196
def add_lora_embedding(
197
self,
@@ -227,7 +239,6 @@ class PunicaWrapperXPU(PunicaWrapperBase):
239
@ lora_b_stacked[indices[i], layer_idx, :, :]
240
* scale
241
).squeeze(0)
230
-
242
Args:
243
y (torch.Tensor): Output tensor. Will be changed in-place.
244
x (torch.Tensor): Input tensor
@@ -240,13 +251,17 @@ class PunicaWrapperXPU(PunicaWrapperBase):
251
252
assert len(lora_a_stacked) == len(lora_b_stacked) == len(output_slices)
253
243
- if buffer is None:
244
- r = lora_b_stacked[0].size(-1)
245
- buffer = torch.zeros( # type: ignore
246
- (len(output_slices), x.size(0), r),
247
- dtype=x.dtype,
248
- device=x.device,
249
- )
254
+ assert buffer is None, (
255
+ "To minimize overhead, the buffer should be created by "
256
+ ".add_lora_linear() instead of being passed in."
257
+ )
258
+ r = lora_b_stacked[0].size(-1)
259
+ buffer = torch.zeros( # type: ignore
260
+ (len(output_slices), x.size(0), r),
261
+ dtype=x.dtype,
262
+ device=x.device,
263
+ )
264
+ add_inputs = kwargs.pop("add_inputs", True)
265
self.add_shrink(
266
buffer, # type: ignore
267
x,
@@ -259,7 +274,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):
274
buffer, # type: ignore
275
lora_b_stacked,
276
output_slices,
262
- add_inputs=True,
277
+ add_inputs=add_inputs,
278
**kwargs,
279
)
280
@@ -300,12 +315,16 @@ class PunicaWrapperXPU(PunicaWrapperBase):
315
y = y.view(-1, y.shape[-1])
316
x = x.view(-1, x.shape[-1])
317
r = lora_b_stacked.size(-1)
303
- if buffer is None:
304
- buffer = torch.zeros((x.size(0), r), dtype=x.dtype, device=x.device)
318
+
319
+ assert buffer is None, (
320
+ "To minimize overhead, the buffer should be created by "
321
+ ".add_lora_linear() instead of being passed in."
322
+ )
323
+ buffer = torch.zeros((x.size(0), r), dtype=x.dtype, device=x.device)
324
sampler_indices = torch.narrow(self._sampler_indices, 0, 0, x.size(0))
325
bgmv_shrink(x, lora_a_stacked, buffer, sampler_indices, scale)
326
bgmv_expand(buffer, lora_b_stacked, y, sampler_indices, add_inputs=True)
308
- return y.view_as(y_org)
327
+ y = y.view_as(y_org)
328
329
def moe_lora_align_block_size(
330
self,
@@ -318,33 +337,62 @@ class PunicaWrapperXPU(PunicaWrapperBase):
337
expert_map: torch.Tensor | None = None,
338
pad_sorted_ids: bool = False,
339
naive_block_assignment: bool = False,
340
+ token_lora_mapping: torch.Tensor | None = None,
341
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
342
"""
343
Aligns tokens and experts into block-sized chunks for LoRA-based
344
mixture-of-experts (MoE) execution.
345
+
346
+ When `token_lora_mapping` is provided, it overrides the global mapping
347
+ read from `self.token_mapping_meta`. This is how EP+LoRA injects the
348
+ per-rank-local token→LoRA map after all-to-all dispatch.
349
"""
326
- (token_lora_mapping, _, _, _, lora_ids, _, _) = (
327
- self.token_mapping_meta.meta_args(
328
- num_tokens, self.lora_config.specialize_active_lora
329
- )
350
+ (
351
+ token_lora_mapping_meta,
352
+ _,
353
+ _,
354
+ _,
355
+ lora_ids,
356
+ _,
357
+ _,
358
+ ) = self.token_mapping_meta.meta_args(
359
+ num_tokens, self.lora_config.specialize_active_lora
360
+ )
361
+ if token_lora_mapping is None:
362
+ token_lora_mapping = token_lora_mapping_meta
363
+ # Under EP the caller passes local_num_experts but topk_ids carries
364
+ # GLOBAL expert indices. The CUDA kernel uses num_experts to size
365
+ # its bucketing table; with EP we must size by global_num_experts
366
+ # so global topk_ids don't overflow. expert_map inside the kernel
367
+ # then translates global→local so the output expert_ids are local
368
+ # (mirrors the non-LoRA moe_align_block_size behavior).
369
+ kernel_num_experts = (
370
+ expert_map.numel() if expert_map is not None else num_experts
371
)
372
if naive_block_assignment:
373
expert_ids = topk_ids.reshape(-1)
374
sorted_ids = None
375
num_tokens_post_pad = None
376
else:
336
- max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
377
+ max_num_tokens_padded = topk_ids.numel() + kernel_num_experts * (
378
+ block_size - 1
379
+ )
380
if pad_sorted_ids:
381
max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
382
+ if topk_ids.numel() < kernel_num_experts:
383
+ max_num_tokens_padded = topk_ids.numel() * block_size
384
sorted_ids = torch.empty(
385
(max_loras * max_num_tokens_padded,),
386
dtype=torch.int32,
387
device=topk_ids.device,
388
)
389
max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size)
345
- # Expert ids must be set default to -1 to prevent a blank block
346
- expert_ids = torch.empty(
390
+ # Expert ids are initialized to -1 so unused (lora, expert)
391
+ # slots don't drive the LoRA Triton kernel into the wrong bucket.
392
+ # The kernel overwrites only active slots.
393
+ expert_ids = torch.full(
394
(max_loras * max_num_m_blocks,),
395
+ -1,
396
dtype=torch.int32,
397
device=topk_ids.device,
398
)
@@ -355,7 +403,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):
403
ops.moe_lora_align_block_size(
404
topk_ids,
405
token_lora_mapping,
358
- num_experts,
406
+ kernel_num_experts,
407
block_size,
408
max_loras,
409
max_num_tokens_padded,
@@ -365,11 +413,10 @@ class PunicaWrapperXPU(PunicaWrapperBase):
413
num_tokens_post_pad,
414
adapter_enabled,
415
lora_ids,
416
+ expert_map,
417
)
369
- if expert_map is not None:
370
- expert_ids = expert_map[expert_ids]
418
372
- return None, sorted_ids, expert_ids, num_tokens_post_pad
419
+ return token_lora_mapping, sorted_ids, expert_ids, num_tokens_post_pad
420
421
def add_lora_fused_moe(
422
self,
@@ -525,7 +572,8 @@ class PunicaWrapperXPU(PunicaWrapperBase):
572
573
SPARSITY_FACTOR = 8
574
naive_block_assignment = (
528
- expert_map is None
575
+ not fully_sharded
576
+ and expert_map is None
577
and num_tokens * top_k * SPARSITY_FACTOR <= local_num_experts * max_loras
578
)
579
@@ -543,6 +591,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):
591
adapter_enabled,
592
expert_map,
593
naive_block_assignment=naive_block_assignment,
594
+ token_lora_mapping=token_lora_mapping,
595
)
596
597
_sorted = sorted_token_ids_lora
@@ -567,6 +616,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):
616
adapter_enabled,
617
fully_sharded=fully_sharded,
618
token_lora_mapping=token_lora_mapping,
619
+ add_inputs=add_inputs,
620
)
621
622
return (
@@ -680,4 +730,5 @@ class PunicaWrapperXPU(PunicaWrapperBase):
730
fully_sharded=fully_sharded,
731
offset=offset,
732
token_lora_mapping=token_lora_mapping,
733
+ add_inputs=add_inputs,
734
)