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, 03:45 UTC
[XPU] Fix lora bugs & enable UTs under tests/lora (#38206)
Signed-off-by: chaojun-zhang <chaojun.zhang@intel.com>
Chaojun Zhang committed
May 7, 2026 at 20:58 UTC
805e9f7b777b6bf39e18674f3db7c950bcbd9db1
8 files changed
+313
-30
.buildkite/intel_jobs/lora_intel.yaml
+12
-7
@@ -20,15 +20,15 @@ steps:
20
'cd tests &&
21
pytest -v -s lora/test_layers.py &&
22
pytest -v -s lora/test_lora_checkpoints.py &&
23
- (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) &&
23
+ pytest -v -s lora/test_lora_functions.py &&
24
pytest -v -s lora/test_lora_huggingface.py &&
25
pytest -v -s lora/test_lora_manager.py &&
26
pytest -v -s lora/test_lora_utils.py &&
27
pytest -v -s lora/test_peft_helper.py &&
28
pytest -v -s lora/test_resolver.py &&
29
pytest -v -s lora/test_utils.py &&
30
- (pytest -v -s lora/test_add_lora.py --deselect="tests/lora/test_add_lora.py::test_add_lora" || true) &&
31
- (pytest -v -s lora/test_worker.py --deselect="tests/lora/test_worker.py::test_worker_apply_lora" || true)'
30
+ pytest -v -s lora/test_add_lora.py &&
31
+ pytest -v -s lora/test_worker.py'
32
33
- label: LoRA Fused/MoE Kernels
34
timeout_in_minutes: 45
@@ -66,7 +66,7 @@ steps:
66
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
67
'cd tests &&
68
set -o pipefail &&
69
- 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]"'
69
+ 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]"'
70
71
- label: LoRA Punica FP8/XPU Ops
72
timeout_in_minutes: 45
@@ -105,8 +105,12 @@ steps:
105
'cd tests &&
106
(pytest -v -s lora/test_mixtral.py --deselect="tests/lora/test_mixtral.py::test_mixtral_lora[4]" || true) &&
107
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]" &&
108
- pytest -v -s lora/test_qwen35_densemodel_lora.py &&
109
- pytest -v -s lora/test_transformers_model.py'
108
+ (pytest -v -s lora/test_qwen35_densemodel_lora.py || true) &&
109
+ pytest -v -s lora/test_transformers_model.py &&
110
+ pytest -v -s lora/test_chatglm3_tp.py &&
111
+ pytest -s -v lora/test_llama_tp.py &&
112
+ pytest -s -v lora/test_minicpmv_tp.py &&
113
+ pytest -s -v lora/test_olmoe_tp.py'
114
115
- label: LoRA Multimodal
116
timeout_in_minutes: 45
@@ -126,4 +130,5 @@ steps:
130
'cd tests &&
131
pytest -v -s lora/test_default_mm_loras.py &&
132
(pytest -v -s lora/test_qwen3_unembed.py || true) &&
129
- pytest -v -s lora/test_whisper.py'
133
+ pytest -v -s lora/test_whisper.py &&
134
+ (pytest -v -s lora/test_llm_with_multi_loras.py || true)'
tests/lora/conftest.py
+2
-6
@@ -62,9 +62,7 @@ def dist_init():
62
63
temp_file = tempfile.mkstemp()[1]
64
65
- backend = "nccl"
66
- if current_platform.is_cpu() or current_platform.is_tpu():
67
- backend = "gloo"
65
+ backend = "gloo" if current_platform.is_tpu() else current_platform.dist_backend
66
67
with ensure_current_vllm_config():
68
init_distributed_environment(
@@ -83,9 +81,7 @@ def dist_init():
81
def dist_init_torch_only():
82
if torch.distributed.is_initialized():
83
return
86
- backend = "nccl"
87
- if current_platform.is_cpu():
88
- backend = "gloo"
84
+ backend = current_platform.dist_backend
85
86
temp_file = tempfile.mkstemp()[1]
87
torch.distributed.init_process_group(
tests/lora/test_fused_moe_lora_kernel.py
+1
@@ -647,6 +647,7 @@ def use_fused_moe_lora_kernel_tensor_parallel(
647
rank=local_rank,
648
local_rank=local_rank,
649
distributed_init_method=init_method,
650
+ backend=current_platform.dist_backend,
651
)
652
with ensure_current_vllm_config():
653
initialize_model_parallel(world_size, 1)
tests/lora/test_layers.py
+30
-14
@@ -57,7 +57,11 @@ TOLERANCES = {
57
}
58
59
pytestmark = pytest.mark.skipif(
60
- not (current_platform.is_cuda_alike() or current_platform.is_cpu()),
60
+ not (
61
+ current_platform.is_cuda_alike()
62
+ or current_platform.is_cpu()
63
+ or current_platform.is_xpu()
64
+ ),
65
reason="Backend not supported",
66
)
67
@@ -67,7 +71,7 @@ DEVICES = (
71
f"{DEVICE_TYPE}:{i}"
72
for i in range(1 if torch.accelerator.device_count() == 1 else 2)
73
]
70
- if current_platform.is_cuda_alike()
74
+ if (current_platform.is_cuda_alike() or current_platform.is_xpu())
75
else ["cpu"]
76
)
77
@@ -96,7 +100,7 @@ def skip_cuda_with_stage_false(request):
100
On cuda-like platforms, we use the same kernels for prefill and decode
101
stage, and 'stage' is generally ignored, so we only need to test once.
102
"""
99
- if current_platform.is_cuda_alike():
103
+ if current_platform.is_cuda_alike() or current_platform.is_xpu():
104
try:
105
if hasattr(request.node, "callspec") and hasattr(
106
request.node.callspec, "params"
@@ -249,6 +253,10 @@ def check_punica_wrapper(punica_wrapper) -> bool:
253
from vllm.lora.punica_wrapper.punica_cpu import PunicaWrapperCPU
254
255
return type(punica_wrapper) is PunicaWrapperCPU
256
+ elif current_platform.is_xpu():
257
+ from vllm.lora.punica_wrapper.punica_xpu import PunicaWrapperXPU
258
+
259
+ return type(punica_wrapper) is PunicaWrapperXPU
260
else:
261
return False
262
@@ -264,7 +272,7 @@ def test_embeddings(
272
# For multi-GPU testing of Triton kernel, we must explicitly set the CUDA
273
# device, see: https://github.com/triton-lang/triton/issues/2925
274
# Same below.
267
- if current_platform.is_cuda_alike():
275
+ if current_platform.is_cuda_alike() or current_platform.is_xpu():
276
torch.accelerator.set_device_index(device)
277
278
torch.set_default_device(device)
@@ -363,7 +371,7 @@ def test_embeddings(
371
def test_lm_head_logits_processor(
372
default_vllm_config, dist_init, num_loras, device, vocab_size, stage
373
) -> None:
366
- if current_platform.is_cuda_alike():
374
+ if current_platform.is_cuda_alike() or current_platform.is_xpu():
375
torch.accelerator.set_device_index(device)
376
377
torch.set_default_device(device)
@@ -480,7 +488,7 @@ def test_lm_head_logits_processor_invalid_vocab_size(
488
default_vllm_config, dist_init, vocab_size, device
489
) -> None:
490
"""Test that LogitsProcessorWithLoRA raises ValueError for invalid vocab sizes."""
483
- if current_platform.is_cuda_alike():
491
+ if current_platform.is_cuda_alike() or current_platform.is_xpu():
492
torch.accelerator.set_device_index(device)
493
494
torch.set_default_device(device)
@@ -509,7 +517,7 @@ def test_linear_replicated(
517
device,
518
stage,
519
) -> None:
512
- if current_platform.is_cuda_alike():
520
+ if current_platform.is_cuda_alike() or current_platform.is_xpu():
521
torch.accelerator.set_device_index(device)
522
523
max_loras = 8
@@ -618,7 +626,7 @@ def test_linear_replicated(
626
def test_linear_parallel(
627
default_vllm_config, dist_init, num_loras, orientation, fully_shard, device, stage
628
) -> None:
621
- if current_platform.is_cuda_alike():
629
+ if current_platform.is_cuda_alike() or current_platform.is_xpu():
630
torch.accelerator.set_device_index(device)
631
632
max_loras = 8
@@ -751,7 +759,7 @@ def test_linear_parallel(
759
def test_column_parallel_packed(
760
default_vllm_config, dist_init, num_loras, repeats, fully_shard, device, stage
761
) -> None:
754
- if current_platform.is_cuda_alike():
762
+ if current_platform.is_cuda_alike() or current_platform.is_xpu():
763
torch.accelerator.set_device_index(device)
764
765
max_loras = 8
@@ -913,7 +921,7 @@ def test_column_parallel_packed(
921
def test_merged_column_parallel_variable_slice(
922
default_vllm_config, dist_init, num_loras, num_slices, device, stage
923
) -> None:
916
- if current_platform.is_cuda_alike():
924
+ if current_platform.is_cuda_alike() or current_platform.is_xpu():
925
torch.accelerator.set_device_index(device)
926
927
max_loras = 8
@@ -1600,11 +1608,15 @@ def test_get_and_maybe_dequant_weights_accepts_lora_wrappers(dist_init, wrapper_
1608
def test_deepseek_fused_qkv_a_proj_lora_preserves_base_forward(
1609
default_vllm_config, dist_init, device, stage, fully_sharded
1610
):
1603
- if current_platform.is_cuda_alike():
1611
+ if current_platform.is_cuda_alike() or current_platform.is_xpu():
1612
torch.accelerator.set_device_index(device)
1613
1614
torch.set_default_device(device)
1607
- dtype = torch.float16 if current_platform.is_cuda_alike() else torch.float32
1615
+ dtype = (
1616
+ torch.float16
1617
+ if (current_platform.is_cuda_alike() or current_platform.is_xpu())
1618
+ else torch.float32
1619
+ )
1620
max_loras = 8
1621
lora_config = LoRAConfig(
1622
max_loras=max_loras,
@@ -1683,11 +1695,15 @@ def test_deepseek_fused_qkv_a_proj_lora_preserves_base_forward(
1695
def test_replicated_lora_preserves_base_forward_for_subclasses(
1696
default_vllm_config, dist_init, device, stage
1697
):
1686
- if current_platform.is_cuda_alike():
1698
+ if current_platform.is_cuda_alike() or current_platform.is_xpu():
1699
torch.accelerator.set_device_index(device)
1700
1701
torch.set_default_device(device)
1690
- dtype = torch.float16 if current_platform.is_cuda_alike() else torch.float32
1702
+ dtype = (
1703
+ torch.float16
1704
+ if current_platform.is_cuda_alike() or current_platform.is_xpu()
1705
+ else torch.float32
1706
+ )
1707
max_loras = 8
1708
lora_config = LoRAConfig(max_loras=max_loras, max_lora_rank=8, lora_dtype=dtype)
1709
punica_wrapper = get_punica_wrapper(8192, 256, device, lora_config=lora_config)
tests/lora/test_llama_tp.py
+4
@@ -224,6 +224,10 @@ def test_tp2_serialize_and_deserialize_lora(
224
max_model_len=1024,
225
tensor_parallel_size=2,
226
max_loras=2,
227
+ # Leave headroom for LoRA adapter loading and Triton JIT
228
+ # compilation, which can allocate GPU memory concurrently
229
+ # during the first inference step.
230
+ gpu_memory_utilization=0.85,
231
)
232
233
tc_as_dict = tensorizer_config.to_serializable()
tests/lora/test_lora_manager.py
+1
-1
@@ -40,7 +40,7 @@ EMBEDDING_MODULES = {
40
DEVICE_TYPE = current_platform.device_type
41
DEVICES = (
42
[f"{DEVICE_TYPE}:{i}" for i in range(min(torch.accelerator.device_count(), 2))]
43
- if current_platform.is_cuda_alike()
43
+ if (current_platform.is_cuda_alike() or current_platform.is_xpu())
44
else ["cpu"]
45
)
46
tests/lora/test_worker.py
+5
-1
@@ -20,7 +20,11 @@ from vllm.config.lora import LoRAConfig
20
from vllm.lora.model_manager import LoRAMapping
21
from vllm.lora.request import LoRARequest
22
from vllm.platforms import current_platform
23
-from vllm.v1.worker.gpu_worker import Worker
23
+
24
+if current_platform.is_xpu():
25
+ from vllm.v1.worker.xpu_worker import XPUWorker as Worker
26
+else:
27
+ from vllm.v1.worker.gpu_worker import Worker
28
29
MODEL_PATH = "Qwen/Qwen3-0.6B"
30
NUM_LORAS = 16
vllm/lora/punica_wrapper/punica_xpu.py
+258
-1
@@ -14,6 +14,7 @@ import torch
14
from vllm import _custom_ops as ops
15
from vllm.lora.layers import LoRAMapping
16
from vllm.lora.ops.xpu_ops import bgmv_expand, bgmv_expand_slice, bgmv_shrink
17
+from vllm.lora.utils import get_captured_lora_counts
18
from vllm.triton_utils import HAS_TRITON, triton
19
from vllm.utils.math_utils import round_up
20
@@ -48,8 +49,24 @@ class PunicaWrapperXPU(PunicaWrapperBase):
49
50
self.lora_config = kwargs["lora_config"]
51
self.max_loras = self.lora_config.max_loras
52
+
53
+ # Compute captured LoRA counts for cudagraph specialization.
54
+ captured_lora_counts = get_captured_lora_counts(
55
+ self.max_loras, self.lora_config.specialize_active_lora
56
+ )
57
+
58
self.token_mapping_meta = LoRAKernelMeta.make(
52
- self.max_loras, max_num_batched_tokens, device=device
59
+ self.max_loras,
60
+ max_num_batched_tokens,
61
+ device=device,
62
+ captured_lora_counts=captured_lora_counts,
63
+ )
64
+
65
+ self.prompt_mapping_meta = LoRAKernelMeta.make(
66
+ self.max_loras,
67
+ max_num_batched_tokens,
68
+ device=device,
69
+ captured_lora_counts=captured_lora_counts,
70
)
71
72
def update_metadata(
@@ -63,6 +80,10 @@ class PunicaWrapperXPU(PunicaWrapperBase):
80
self.is_prefill = mapping.is_prefill
81
self._update_base_metadata(mapping, lora_index_to_id, max_loras, vocab_size)
82
83
+ # Prepare kernel metadata tensors
84
+ self.token_mapping_meta.prepare_tensors(self.token_lora_indices)
85
+ self.prompt_mapping_meta.prepare_tensors(self.sampler_indices)
86
+
87
def _get_token_lora_indices(self, x: torch.Tensor) -> torch.IntTensor:
88
return torch.narrow(self._token_lora_indices, 0, 0, x.size(0))
89
@@ -419,3 +440,239 @@ class PunicaWrapperXPU(PunicaWrapperBase):
440
fully_sharded,
441
offset,
442
)
443
+
444
+ def add_lora_w13(
445
+ self,
446
+ y: torch.Tensor,
447
+ x: torch.Tensor,
448
+ lora_a_stacked: tuple[torch.Tensor, ...],
449
+ lora_b_stacked: tuple[torch.Tensor, ...],
450
+ topk_ids: torch.Tensor,
451
+ topk_weights: torch.Tensor,
452
+ expert_map: torch.Tensor | None,
453
+ w1: torch.Tensor,
454
+ w2: torch.Tensor,
455
+ num_tokens: int,
456
+ top_k_num: int,
457
+ max_loras: int,
458
+ adapter_enabled: torch.Tensor,
459
+ local_num_experts: int,
460
+ top_k: int,
461
+ num_slices: int,
462
+ fully_sharded: bool,
463
+ use_tuned_config: bool,
464
+ ) -> tuple[
465
+ torch.Tensor | None,
466
+ torch.Tensor | None,
467
+ torch.Tensor | None,
468
+ torch.Tensor | None,
469
+ ]:
470
+ import functools
471
+
472
+ from vllm.lora.layers.utils import try_get_optimal_moe_lora_config
473
+ from vllm.lora.ops.triton_ops.utils import (
474
+ _normalize_lora_config_keys,
475
+ get_lora_op_configs,
476
+ )
477
+ from vllm.model_executor.layers.fused_moe.config import _get_config_dtype_str
478
+
479
+ config_dtype = _get_config_dtype_str(
480
+ dtype=x.dtype,
481
+ use_fp8_w8a8=False,
482
+ use_int8_w8a16=False,
483
+ use_int4_w4a16=False,
484
+ )
485
+ max_lora_rank = lora_a_stacked[0].shape[-2]
486
+
487
+ if use_tuned_config:
488
+ shrink_config = get_lora_op_configs(
489
+ op_type="fused_moe_lora_w13_shrink",
490
+ max_loras=max_loras,
491
+ batch=num_tokens,
492
+ hidden_size=x.shape[-1],
493
+ rank=max_lora_rank,
494
+ num_slices=num_slices,
495
+ moe_intermediate_size=lora_b_stacked[0].shape[-2],
496
+ )
497
+ expand_config = get_lora_op_configs(
498
+ op_type="fused_moe_lora_w13_expand",
499
+ max_loras=max_loras,
500
+ batch=num_tokens,
501
+ hidden_size=x.shape[-1],
502
+ rank=max_lora_rank,
503
+ num_slices=num_slices,
504
+ moe_intermediate_size=lora_b_stacked[0].shape[-2],
505
+ )
506
+ else:
507
+ get_config = functools.partial(
508
+ try_get_optimal_moe_lora_config,
509
+ w1_shape=w1.shape,
510
+ w2_shape=w2.shape,
511
+ rank=max_lora_rank,
512
+ top_k=top_k,
513
+ dtype=config_dtype,
514
+ M=num_tokens,
515
+ )
516
+ shrink_config = get_config(op_type="fused_moe_lora_w13_shrink")
517
+ expand_config = get_config(op_type="fused_moe_lora_w13_expand")
518
+
519
+ shrink_config = _normalize_lora_config_keys(shrink_config)
520
+ expand_config = _normalize_lora_config_keys(expand_config)
521
+
522
+ SPARSITY_FACTOR = 8
523
+ naive_block_assignment = (
524
+ expert_map is None
525
+ and num_tokens * top_k * SPARSITY_FACTOR <= local_num_experts * max_loras
526
+ )
527
+
528
+ (
529
+ token_lora_mapping,
530
+ sorted_token_ids_lora,
531
+ expert_ids_lora,
532
+ num_tokens_post_padded_lora,
533
+ ) = self.moe_lora_align_block_size(
534
+ topk_ids,
535
+ num_tokens,
536
+ int(shrink_config.get("BLOCK_SIZE_M") or 64),
537
+ local_num_experts,
538
+ max_loras,
539
+ adapter_enabled,
540
+ expert_map,
541
+ naive_block_assignment=naive_block_assignment,
542
+ )
543
+
544
+ _sorted = sorted_token_ids_lora
545
+ _eids = expert_ids_lora
546
+ if _sorted is not None:
547
+ _eids = _eids.view(max_loras, -1)
548
+ _sorted = _sorted.view(max_loras, -1)
549
+
550
+ self.add_lora_fused_moe(
551
+ y.view(-1, top_k_num, y.shape[-1]),
552
+ x,
553
+ lora_a_stacked,
554
+ lora_b_stacked,
555
+ topk_weights,
556
+ _sorted,
557
+ _eids,
558
+ num_tokens_post_padded_lora,
559
+ max_lora_rank,
560
+ top_k,
561
+ shrink_config,
562
+ expand_config,
563
+ adapter_enabled,
564
+ fully_sharded=fully_sharded,
565
+ token_lora_mapping=token_lora_mapping,
566
+ )
567
+
568
+ return (
569
+ sorted_token_ids_lora,
570
+ expert_ids_lora,
571
+ num_tokens_post_padded_lora,
572
+ token_lora_mapping,
573
+ )
574
+
575
+ def add_lora_w2(
576
+ self,
577
+ y: torch.Tensor,
578
+ x: torch.Tensor,
579
+ lora_a_stacked: tuple[torch.Tensor, ...],
580
+ lora_b_stacked: tuple[torch.Tensor, ...],
581
+ topk_weights: torch.Tensor,
582
+ sorted_token_ids_lora: torch.Tensor | None,
583
+ expert_ids_lora: torch.Tensor | None,
584
+ num_tokens_post_padded_lora: torch.Tensor | None,
585
+ token_lora_mapping: torch.Tensor | None,
586
+ num_tokens: int,
587
+ w1: torch.Tensor,
588
+ w2: torch.Tensor,
589
+ top_k_num: int,
590
+ max_loras: int,
591
+ adapter_enabled: torch.Tensor,
592
+ top_k: int,
593
+ fully_sharded: bool,
594
+ tp_rank: int,
595
+ use_tuned_config: bool,
596
+ ) -> None:
597
+ import functools
598
+
599
+ from vllm.lora.layers.utils import try_get_optimal_moe_lora_config
600
+ from vllm.lora.ops.triton_ops.utils import (
601
+ _normalize_lora_config_keys,
602
+ get_lora_op_configs,
603
+ )
604
+ from vllm.model_executor.layers.fused_moe.config import _get_config_dtype_str
605
+
606
+ config_dtype = _get_config_dtype_str(
607
+ dtype=x.dtype,
608
+ use_fp8_w8a8=False,
609
+ use_int8_w8a16=False,
610
+ use_int4_w4a16=False,
611
+ )
612
+ max_lora_rank = lora_a_stacked[0].shape[-2]
613
+
614
+ if use_tuned_config:
615
+ shrink_config = get_lora_op_configs(
616
+ op_type="fused_moe_lora_w2_shrink",
617
+ max_loras=max_loras,
618
+ batch=num_tokens,
619
+ hidden_size=y.shape[-1],
620
+ rank=max_lora_rank,
621
+ num_slices=1,
622
+ moe_intermediate_size=lora_a_stacked[0].shape[-1],
623
+ )
624
+ expand_config = get_lora_op_configs(
625
+ op_type="fused_moe_lora_w2_expand",
626
+ max_loras=max_loras,
627
+ batch=num_tokens,
628
+ hidden_size=y.shape[-1],
629
+ rank=max_lora_rank,
630
+ num_slices=1,
631
+ moe_intermediate_size=lora_a_stacked[0].shape[-1],
632
+ )
633
+ else:
634
+ get_config = functools.partial(
635
+ try_get_optimal_moe_lora_config,
636
+ w1_shape=w1.shape,
637
+ w2_shape=w2.shape,
638
+ rank=max_lora_rank,
639
+ top_k=top_k,
640
+ dtype=config_dtype,
641
+ M=num_tokens,
642
+ )
643
+ shrink_config = get_config(op_type="fused_moe_lora_w2_shrink")
644
+ expand_config = get_config(op_type="fused_moe_lora_w2_expand")
645
+
646
+ shrink_config = _normalize_lora_config_keys(shrink_config)
647
+ expand_config = _normalize_lora_config_keys(expand_config)
648
+
649
+ _sorted = sorted_token_ids_lora
650
+ _eids = expert_ids_lora
651
+ if _sorted is not None:
652
+ assert _eids is not None
653
+ _eids = _eids.view(max_loras, -1)
654
+ _sorted = _sorted.view(max_loras, -1)
655
+
656
+ # w2_lora_b shape[-2] is hidden_size // tp_size when fully_sharded
657
+ shard_size = lora_b_stacked[0].shape[-2]
658
+ offset = shard_size * tp_rank if fully_sharded else 0
659
+
660
+ self.add_lora_fused_moe(
661
+ y,
662
+ x,
663
+ lora_a_stacked,
664
+ lora_b_stacked,
665
+ topk_weights,
666
+ _sorted,
667
+ _eids,
668
+ num_tokens_post_padded_lora,
669
+ max_lora_rank,
670
+ top_k,
671
+ shrink_config,
672
+ expand_config,
673
+ adapter_enabled,
674
+ True, # mul_routed_weight
675
+ fully_sharded=fully_sharded,
676
+ offset=offset,
677
+ token_lora_mapping=token_lora_mapping,
678
+ )