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 7, 16:30 UTC
[XPU][8/N] Fix kernel bugs in XPU LoRA and MOE LORA (#34115)
Signed-off-by: chzhang <chaojun.zhang@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Chaojun Zhang committed
Feb 26, 2026 at 15:46 UTC
9f9a675b23c8276237e98b396c3242c5135808fb
5 files changed
+462
-20
tests/lora/test_fused_moe_lora_kernel.py
+3
-1
@@ -18,6 +18,7 @@ from vllm.distributed.parallel_state import (
18
get_tensor_model_parallel_world_size,
19
)
20
from vllm.lora.ops.triton_ops import fused_moe_lora
21
+from vllm.platforms import current_platform
22
from vllm.utils.network_utils import get_open_port
23
from vllm.utils.torch_utils import set_random_seed
24
@@ -244,8 +245,9 @@ def use_torch(
245
return torch.stack(outputs, dim=0)
246
247
248
+DEVICE_TYPE = current_platform.device_type
249
DTYPES = [torch.float16, torch.bfloat16]
248
-DEVICES = [f"cuda:{0}"]
250
+DEVICES = [f"{DEVICE_TYPE}:{0}"]
251
SEED = [42]
252
253
tests/lora/test_punica_xpu_ops.py
new
+298
@@ -0,0 +1,298 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+from tests.lora.utils import (
8
+ PunicaTensors,
9
+ assert_close,
10
+ generate_data,
11
+ generate_data_for_expand_nslices,
12
+)
13
+from vllm.lora.ops.xpu_ops import bgmv_expand, bgmv_expand_slice, bgmv_shrink
14
+from vllm.platforms import current_platform
15
+
16
+
17
+def torch_bgmv_expand(
18
+ inputs: torch.Tensor,
19
+ lora_b_weights: torch.Tensor,
20
+ output_tensor: torch.Tensor,
21
+ lora_indices_tensor: torch.Tensor,
22
+ add_inputs: bool = True,
23
+):
24
+ selected_loras = lora_b_weights[lora_indices_tensor].to(dtype=output_tensor.dtype)
25
+ if len(selected_loras.shape) == 4:
26
+ selected_loras = selected_loras.squeeze(dim=1)
27
+ inputs = inputs.to(dtype=output_tensor.dtype)
28
+ outputs = torch.einsum("bi, boi -> bo", inputs, selected_loras)
29
+
30
+ limit = output_tensor.shape[0]
31
+ if outputs.shape[0] == 1 and output_tensor.shape[0] != 1:
32
+ limit = 1
33
+
34
+ # LoRA adapter and model may add different amounts of padding to output
35
+ common_len = min(outputs.shape[1], output_tensor.shape[1])
36
+
37
+ if add_inputs:
38
+ output_tensor[:, :common_len] += outputs[:limit, :common_len]
39
+ else:
40
+ output_tensor[:, :common_len] = outputs[:limit, :common_len]
41
+
42
+
43
+def torch_bgmv_shrink(
44
+ inputs: torch.Tensor,
45
+ lora_b_weights: torch.Tensor,
46
+ output_tensor: torch.Tensor,
47
+ lora_indices_tensor: torch.Tensor,
48
+ scaling: float = 1.0,
49
+):
50
+ selected_loras = lora_b_weights[lora_indices_tensor].to(dtype=output_tensor.dtype)
51
+ if len(selected_loras.shape) == 4:
52
+ selected_loras = selected_loras.squeeze(dim=1)
53
+ inputs = inputs.to(dtype=output_tensor.dtype)
54
+ outputs = torch.einsum("bi, boi -> bo", inputs, selected_loras)
55
+
56
+ output_tensor[:, : outputs.shape[1]] = scaling * outputs[:]
57
+
58
+
59
+def torch_bgmv_expand_slice(
60
+ inputs: torch.Tensor,
61
+ lora_b_weights: torch.Tensor,
62
+ output_tensor: torch.Tensor,
63
+ lora_indices_tensor: torch.Tensor,
64
+ slice_offset: int,
65
+ slice_size: int,
66
+ add_inputs: bool = True,
67
+):
68
+ selected_loras = lora_b_weights[lora_indices_tensor].to(dtype=output_tensor.dtype)
69
+ inputs = inputs.to(dtype=output_tensor.dtype)
70
+ if len(selected_loras.shape) == 4:
71
+ selected_loras = selected_loras.squeeze(dim=1)
72
+ outputs = torch.einsum("bi, boi -> bo", inputs, selected_loras)
73
+
74
+ if add_inputs:
75
+ output_tensor[:, slice_offset : slice_offset + slice_size] += outputs[:]
76
+ else:
77
+ output_tensor[:, slice_offset : slice_offset + slice_size] = outputs[:]
78
+
79
+
80
+def check_bgmv_shrink(
81
+ batches: int,
82
+ num_loras: int,
83
+ rank: int,
84
+ hidden_size: int,
85
+ dtype: torch.dtype,
86
+ device: str,
87
+ scaling: float,
88
+):
89
+ """
90
+ Compare vllm.bgmv_shrink against a reference implementation.
91
+ """
92
+ seq_length = 1
93
+ data: PunicaTensors = generate_data(
94
+ batches,
95
+ hidden_size,
96
+ num_loras,
97
+ rank,
98
+ seq_length,
99
+ dtype,
100
+ "shrink",
101
+ device,
102
+ )
103
+
104
+ bgmv_shrink(
105
+ data.inputs_tensor,
106
+ data.lora_weights,
107
+ data.our_out_tensor,
108
+ data.token_lora_mapping,
109
+ scaling,
110
+ )
111
+
112
+ torch_bgmv_shrink(
113
+ data.inputs_tensor,
114
+ data.lora_weights,
115
+ data.ref_out_tensor,
116
+ data.token_lora_mapping,
117
+ scaling,
118
+ )
119
+
120
+ data.ref_out_tensor = data.ref_out_tensor.to(torch.float32)
121
+ assert_close(data.our_out_tensor, data.ref_out_tensor)
122
+
123
+
124
+def check_bgmv_expand(
125
+ batches: int,
126
+ num_loras: int,
127
+ rank: int,
128
+ hidden_size: int,
129
+ dtype: torch.dtype,
130
+ device: str,
131
+ add_inputs: bool,
132
+):
133
+ """
134
+ Compare vllm.bgmv_expand against a reference implementation.
135
+ """
136
+ seq_length = 1
137
+ data: PunicaTensors = generate_data(
138
+ batches,
139
+ hidden_size,
140
+ num_loras,
141
+ rank,
142
+ seq_length,
143
+ dtype,
144
+ "expand",
145
+ device,
146
+ )
147
+
148
+ bgmv_expand(
149
+ data.inputs_tensor,
150
+ data.lora_weights,
151
+ data.our_out_tensor,
152
+ data.token_lora_mapping,
153
+ add_inputs=add_inputs,
154
+ )
155
+ torch_bgmv_expand(
156
+ data.inputs_tensor,
157
+ data.lora_weights,
158
+ data.ref_out_tensor,
159
+ data.token_lora_mapping,
160
+ add_inputs=add_inputs,
161
+ )
162
+ assert_close(data.ref_out_tensor, data.our_out_tensor)
163
+
164
+
165
+def check_bgmv_expand_slice(
166
+ batches: int,
167
+ num_loras: int,
168
+ rank: int,
169
+ hidden_size: int,
170
+ nslices: int,
171
+ dtype: torch.dtype,
172
+ device: str,
173
+ add_inputs: bool,
174
+):
175
+ """
176
+ Compare vllm.bgmv_expand_slice against a reference implementation.
177
+ """
178
+ seq_length = 1
179
+ data: PunicaTensors = generate_data_for_expand_nslices(
180
+ batches,
181
+ hidden_size,
182
+ num_loras,
183
+ rank,
184
+ seq_length,
185
+ dtype,
186
+ nslices,
187
+ device,
188
+ )
189
+
190
+ slice_offset = 0
191
+ for index in range(nslices):
192
+ bgmv_expand_slice(
193
+ data.inputs_tensor,
194
+ data.lora_weights[index],
195
+ data.our_out_tensor,
196
+ data.token_lora_mapping,
197
+ slice_offset,
198
+ slice_size=hidden_size,
199
+ add_inputs=add_inputs,
200
+ )
201
+ torch_bgmv_expand_slice(
202
+ data.inputs_tensor,
203
+ data.lora_weights[index],
204
+ data.ref_out_tensor,
205
+ data.token_lora_mapping,
206
+ slice_offset,
207
+ slice_size=hidden_size,
208
+ add_inputs=add_inputs,
209
+ )
210
+
211
+ slice_offset += hidden_size
212
+ assert_close(data.ref_out_tensor, data.our_out_tensor)
213
+
214
+
215
+# General tests params that tests for variations in all dimensions
216
+# except hidden_size.
217
+test_params = {
218
+ "hidden_sizes": [2049],
219
+ "batches": [4],
220
+ "num_loras": [4],
221
+ "max_ranks": [32],
222
+}
223
+
224
+DTYPES = [torch.float16, torch.bfloat16]
225
+DEVICES = [f"xpu:{0}"]
226
+SEED = [0]
227
+
228
+
229
+@pytest.mark.parametrize("batches", test_params["batches"])
230
+@pytest.mark.parametrize("num_loras", test_params["num_loras"])
231
+@pytest.mark.parametrize("rank", test_params["max_ranks"])
232
+@pytest.mark.parametrize("hidden_size", test_params["hidden_sizes"])
233
+@pytest.mark.parametrize("dtype", DTYPES)
234
+@pytest.mark.parametrize("device", DEVICES)
235
+@pytest.mark.parametrize("seed", SEED)
236
+@pytest.mark.parametrize("op_type", ["shrink", "expand"])
237
+@pytest.mark.skipif(not current_platform.is_xpu(), reason="skip for non xpu platform")
238
+def test_bgmv(
239
+ batches: int,
240
+ num_loras: int,
241
+ rank: int,
242
+ hidden_size: int,
243
+ dtype: torch.dtype,
244
+ device: str,
245
+ seed: int,
246
+ op_type: str,
247
+):
248
+ if op_type == "shrink":
249
+ check_bgmv_shrink(
250
+ batches=batches,
251
+ num_loras=num_loras,
252
+ rank=rank,
253
+ hidden_size=hidden_size,
254
+ dtype=dtype,
255
+ device=device,
256
+ scaling=0.5,
257
+ )
258
+ else:
259
+ check_bgmv_expand(
260
+ batches=batches,
261
+ num_loras=num_loras,
262
+ rank=rank,
263
+ hidden_size=hidden_size,
264
+ dtype=dtype,
265
+ device=device,
266
+ add_inputs=True,
267
+ )
268
+
269
+
270
+@pytest.mark.parametrize("batches", test_params["batches"])
271
+@pytest.mark.parametrize("num_loras", test_params["num_loras"])
272
+@pytest.mark.parametrize("rank", test_params["max_ranks"])
273
+@pytest.mark.parametrize("hidden_size", test_params["hidden_sizes"])
274
+@pytest.mark.parametrize("nslices", [2, 3])
275
+@pytest.mark.parametrize("dtype", DTYPES)
276
+@pytest.mark.parametrize("device", DEVICES)
277
+@pytest.mark.parametrize("seed", SEED)
278
+@pytest.mark.skipif(not current_platform.is_xpu(), reason="skip for non xpu platform")
279
+def test_bgmv_expand_nslices(
280
+ batches: int,
281
+ num_loras: int,
282
+ rank: int,
283
+ hidden_size: int,
284
+ nslices: int,
285
+ dtype: torch.dtype,
286
+ device: str,
287
+ seed: int,
288
+):
289
+ check_bgmv_expand_slice(
290
+ batches=batches,
291
+ num_loras=num_loras,
292
+ rank=rank,
293
+ hidden_size=hidden_size,
294
+ nslices=nslices,
295
+ dtype=dtype,
296
+ device=device,
297
+ add_inputs=True,
298
+ )
vllm/lora/ops/xpu_ops/__init__.py
renamed
+1
-1
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from vllm.lora.ops.ipex_ops.lora_ops import bgmv_expand, bgmv_expand_slice, bgmv_shrink
4
+from vllm.lora.ops.xpu_ops.lora_ops import bgmv_expand, bgmv_expand_slice, bgmv_shrink
5
6
__all__ = ["bgmv_expand", "bgmv_expand_slice", "bgmv_shrink"]
vllm/lora/ops/xpu_ops/lora_ops.py
renamed
+8
-11
@@ -7,11 +7,6 @@ from vllm.logger import init_logger
7
8
logger = init_logger(__name__)
9
10
-try:
11
- import intel_extension_for_pytorch as ipex
12
-except ImportError as e:
13
- raise e
14
-
10
11
def bgmv_shrink(
12
inputs: torch.Tensor,
@@ -20,8 +15,8 @@ def bgmv_shrink(
15
lora_indices_tensor: torch.Tensor,
16
scaling: float = 1.0,
17
) -> None:
23
- ipex.llm.functional.bgmv_shrink(
24
- inputs, lora_a_weights, output_tensor, lora_indices_tensor, scaling
18
+ torch.ops._xpu_C.bgmv_shrink(
19
+ output_tensor, inputs, lora_a_weights, lora_indices_tensor, scaling
20
)
21
22
@@ -32,8 +27,8 @@ def bgmv_expand(
27
lora_indices_tensor: torch.Tensor,
28
add_inputs: bool = True,
29
) -> None:
35
- ipex.llm.functional.bgmv_expand(
36
- inputs, lora_b_weights, output_tensor, lora_indices_tensor, add_inputs
30
+ torch.ops._xpu_C.bgmv_expand(
31
+ output_tensor, inputs, lora_b_weights, lora_indices_tensor, add_inputs
32
)
33
34
@@ -46,10 +41,12 @@ def bgmv_expand_slice(
41
slice_size: int,
42
add_inputs: bool = True,
43
) -> None:
49
- ipex.llm.functional.bgmv_expand_slice(
44
+ assert slice_size == lora_b_weights.size(-2)
45
+ assert slice_offset + slice_size <= output_tensor.size(1)
46
+ torch.ops._xpu_C.bgmv_expand_slice(
47
+ output_tensor,
48
inputs,
49
lora_b_weights,
52
- output_tensor,
50
lora_indices_tensor,
51
slice_offset,
52
slice_size,
vllm/lora/punica_wrapper/punica_xpu.py
+152
-7
@@ -11,8 +11,17 @@ from typing import final
11
12
import torch
13
14
+from vllm import _custom_ops as ops
15
from vllm.lora.layers import LoRAMapping
15
-from vllm.lora.ops.ipex_ops import bgmv_expand, bgmv_expand_slice, bgmv_shrink
16
+from vllm.lora.ops.xpu_ops import bgmv_expand, bgmv_expand_slice, bgmv_shrink
17
+from vllm.triton_utils import HAS_TRITON, triton
18
+from vllm.utils.math_utils import round_up
19
+
20
+if HAS_TRITON:
21
+ from vllm.lora.ops.triton_ops import (
22
+ LoRAKernelMeta,
23
+ fused_moe_lora,
24
+ )
25
26
from .punica_base import PunicaWrapperBase
27
@@ -37,6 +46,12 @@ class PunicaWrapperXPU(PunicaWrapperBase):
46
torch._dynamo.mark_dynamic(self._embeddings_indices, 1)
47
torch._dynamo.mark_dynamic(self._sampler_indices_padded, 0)
48
49
+ self.lora_config = kwargs["lora_config"]
50
+ self.max_loras = self.lora_config.max_loras
51
+ self.token_mapping_meta = LoRAKernelMeta.make(
52
+ self.max_loras, max_num_batched_tokens, device=device
53
+ )
54
+
55
def update_metadata(
56
self,
57
mapping: LoRAMapping,
@@ -206,11 +221,9 @@ class PunicaWrapperXPU(PunicaWrapperBase):
221
222
if buffer is None:
223
r = lora_b_stacked[0].size(-1)
209
- # We set the buffer to be float32 by default, refer to:
210
- # https://github.com/triton-lang/triton/issues/1387
224
buffer = torch.zeros( # type: ignore
225
(len(output_slices), x.size(0), r),
213
- dtype=torch.float32,
226
+ dtype=x.dtype,
227
device=x.device,
228
)
229
self.add_shrink(
@@ -267,10 +280,142 @@ class PunicaWrapperXPU(PunicaWrapperBase):
280
x = x.view(-1, x.shape[-1])
281
r = lora_b_stacked.size(-1)
282
if buffer is None:
270
- # We set the buffer to be float32 by default, refer to:
271
- # https://github.com/triton-lang/triton/issues/1387
272
- buffer = torch.zeros((x.size(0), r), dtype=torch.float32, device=x.device)
283
+ buffer = torch.zeros((x.size(0), r), dtype=x.dtype, device=x.device)
284
sampler_indices = torch.narrow(self._sampler_indices, 0, 0, x.size(0))
285
bgmv_shrink(x, lora_a_stacked, buffer, sampler_indices, scale)
286
bgmv_expand(buffer, lora_b_stacked, y, sampler_indices, add_inputs=True)
287
return y.view_as(y_org)
288
+
289
+ def moe_lora_align_block_size(
290
+ self,
291
+ topk_ids: torch.Tensor,
292
+ num_tokens: int,
293
+ block_size: int,
294
+ num_experts: int,
295
+ max_loras: int,
296
+ adapter_enabled: torch.Tensor,
297
+ expert_map: torch.Tensor | None = None,
298
+ pad_sorted_ids: bool = False,
299
+ naive_block_assignment: bool = False,
300
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
301
+ """
302
+ Aligns tokens and experts into block-sized chunks for LoRA-based
303
+ mixture-of-experts (MoE) execution.
304
+ """
305
+ (token_lora_mapping, _, _, _, lora_ids, _, _) = (
306
+ self.token_mapping_meta.meta_args(
307
+ num_tokens, self.lora_config.specialize_active_lora
308
+ )
309
+ )
310
+ if naive_block_assignment:
311
+ expert_ids = topk_ids.reshape(-1)
312
+ sorted_ids = None
313
+ num_tokens_post_pad = None
314
+ else:
315
+ max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
316
+ if pad_sorted_ids:
317
+ max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
318
+ sorted_ids = torch.empty(
319
+ (max_loras * max_num_tokens_padded,),
320
+ dtype=torch.int32,
321
+ device=topk_ids.device,
322
+ )
323
+ max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size)
324
+ # Expert ids must be set default to -1 to prevent a blank block
325
+ expert_ids = torch.empty(
326
+ (max_loras * max_num_m_blocks,),
327
+ dtype=torch.int32,
328
+ device=topk_ids.device,
329
+ )
330
+ num_tokens_post_pad = torch.empty(
331
+ (max_loras), dtype=torch.int32, device=topk_ids.device
332
+ )
333
+
334
+ ops.moe_lora_align_block_size(
335
+ topk_ids,
336
+ token_lora_mapping,
337
+ num_experts,
338
+ block_size,
339
+ max_loras,
340
+ max_num_tokens_padded,
341
+ max_num_m_blocks,
342
+ sorted_ids,
343
+ expert_ids,
344
+ num_tokens_post_pad,
345
+ adapter_enabled,
346
+ lora_ids,
347
+ )
348
+ if expert_map is not None:
349
+ expert_ids = expert_map[expert_ids]
350
+
351
+ return None, sorted_ids, expert_ids, num_tokens_post_pad
352
+
353
+ def add_lora_fused_moe(
354
+ self,
355
+ y: torch.Tensor,
356
+ x: torch.Tensor,
357
+ lora_a_stacked: tuple[torch.Tensor, ...],
358
+ lora_b_stacked: tuple[torch.Tensor, ...],
359
+ topk_weights: torch.Tensor,
360
+ sorted_token_ids: torch.Tensor | None,
361
+ expert_ids: torch.Tensor,
362
+ num_tokens_post_padded: torch.Tensor | None,
363
+ max_lora_rank: int,
364
+ top_k_num: int,
365
+ shrink_config,
366
+ expand_config,
367
+ adapter_enabled: torch.Tensor,
368
+ mul_routed_weight=False,
369
+ fully_sharded: bool = False,
370
+ offset: int = 0,
371
+ token_lora_mapping: torch.Tensor | None = None,
372
+ ):
373
+ """
374
+ Performs a fused forward computation for LoRA of Mixture-of-Experts (MoE) layer.
375
+ """
376
+ (
377
+ token_lora_mapping_meta,
378
+ _,
379
+ _,
380
+ _,
381
+ lora_ids,
382
+ _,
383
+ num_active_loras,
384
+ ) = self.token_mapping_meta.meta_args(
385
+ x.size(0), self.lora_config.specialize_active_lora
386
+ )
387
+ if token_lora_mapping is None:
388
+ token_lora_mapping = token_lora_mapping_meta
389
+ fused_moe_lora(
390
+ y,
391
+ x,
392
+ lora_a_stacked,
393
+ lora_b_stacked,
394
+ topk_weights,
395
+ sorted_token_ids,
396
+ expert_ids,
397
+ num_tokens_post_padded,
398
+ token_lora_mapping,
399
+ max_lora_rank,
400
+ top_k_num,
401
+ lora_ids,
402
+ num_active_loras,
403
+ adapter_enabled,
404
+ shrink_config.get("BLOCK_SIZE_M", 64),
405
+ shrink_config.get("BLOCK_SIZE_N", 64),
406
+ shrink_config.get("BLOCK_SIZE_K", 32),
407
+ shrink_config.get("GROUP_SIZE_M", 8),
408
+ shrink_config.get("NUM_WARPS", 4),
409
+ shrink_config.get("NUM_STAGES", 3),
410
+ shrink_config.get("SPLIT_K", 1),
411
+ expand_config.get("BLOCK_SIZE_M", 64),
412
+ expand_config.get("BLOCK_SIZE_N", 64),
413
+ expand_config.get("BLOCK_SIZE_K", 32),
414
+ expand_config.get("GROUP_SIZE_M", 8),
415
+ expand_config.get("NUM_WARPS", 4),
416
+ expand_config.get("NUM_STAGES", 3),
417
+ expand_config.get("SPLIT_K", 1),
418
+ mul_routed_weight,
419
+ fully_sharded,
420
+ offset,
421
+ )