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 10, 09:00 UTC
[LoRA][Perf] Improve FusedMoE LoRA performance for small rank (#32019)
Signed-off-by: Xin Yang <xyangx@amazon.com>
Xin Yang committed
Jan 10, 2026 at 11:04 UTC
543c23be789cef54bcb8facacaf7417f42fb95cf
3 files changed
+61
-11
vllm/lora/layers/fused_moe.py
+14
-9
@@ -24,7 +24,6 @@ from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
24
)
25
from vllm.model_executor.layers.fused_moe.fused_moe import (
26
TritonExperts,
27
- try_get_optimal_moe_config,
27
)
28
from vllm.model_executor.layers.fused_moe.fused_moe_modular_method import (
29
FusedMoEModularMethod,
@@ -39,7 +38,7 @@ from vllm.model_executor.layers.fused_moe.prepare_finalize import (
38
MoEPrepareAndFinalizeNoEP,
39
)
40
42
-from .utils import _get_lora_device
41
+from .utils import _get_lora_device, try_get_optimal_moe_lora_config
42
43
44
class FusedMoEWithLoRA(BaseLayerWithLoRA):
@@ -103,15 +102,21 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
102
)
103
else: # fall back to the default config
104
get_config_func = functools.partial(
106
- try_get_optimal_moe_config,
107
- layer.w13_weight.size(),
108
- layer.w2_weight.size(),
109
- top_k,
110
- config_dtype,
105
+ try_get_optimal_moe_lora_config,
106
+ w1_shape=layer.w13_weight.size(),
107
+ w2_shape=layer.w2_weight.size(),
108
+ rank=rank,
109
+ top_k=top_k,
110
+ dtype=config_dtype,
111
+ M=M,
112
block_shape=layer.quant_method.moe_quant_config.block_shape,
113
)
113
- shrink_config = get_config_func(M)
114
- expand_config = get_config_func(M)
114
+ shrink_config = get_config_func(
115
+ op_type=f"fused_moe_lora_{op_prefix}_shrink"
116
+ )
117
+ expand_config = get_config_func(
118
+ op_type=f"fused_moe_lora_{op_prefix}_expand"
119
+ )
120
shrink_config = self._normalize_keys(shrink_config)
121
expand_config = self._normalize_keys(expand_config)
122
return shrink_config, expand_config
vllm/lora/layers/utils.py
+33
@@ -7,6 +7,9 @@ from enum import Enum
7
import torch
8
import torch.nn as nn
9
10
+from vllm.model_executor.layers.fused_moe.fused_moe import try_get_optimal_moe_config
11
+from vllm.utils.math_utils import next_power_of_2
12
+
13
14
class LoRAMappingType(Enum):
15
LANGUAGE = 1
@@ -80,3 +83,33 @@ def _fully_sharded_can_replace(can_replace):
83
)
84
85
return dec
86
+
87
+
88
+def try_get_optimal_moe_lora_config(
89
+ op_type: str,
90
+ w1_shape: tuple[int, ...],
91
+ w2_shape: tuple[int, ...],
92
+ rank: int,
93
+ top_k: int,
94
+ dtype: str | None,
95
+ M: int,
96
+ block_shape: list[int] | None = None,
97
+) -> dict[str, int | None]:
98
+ config = try_get_optimal_moe_config(
99
+ w1_shape, w2_shape, top_k, dtype, M, block_shape
100
+ ).copy()
101
+ if op_type in [
102
+ "fused_moe_lora_w13_shrink",
103
+ "fused_moe_lora_w2_shrink",
104
+ ]:
105
+ config["BLOCK_SIZE_N"] = min(
106
+ config.get("BLOCK_SIZE_N", 64), next_power_of_2(rank)
107
+ )
108
+ elif op_type in [
109
+ "fused_moe_lora_w13_expand",
110
+ "fused_moe_lora_w2_expand",
111
+ ]:
112
+ config["BLOCK_SIZE_K"] = max(
113
+ 16, min(config.get("BLOCK_SIZE_K", 32), next_power_of_2(rank))
114
+ )
115
+ return config
vllm/lora/ops/triton_ops/utils.py
+14
-2
@@ -13,6 +13,7 @@ from vllm import envs
13
from vllm.logger import init_logger
14
from vllm.model_executor.layers.batch_invariant import vllm_is_batch_invariant
15
from vllm.platforms import current_platform
16
+from vllm.utils.math_utils import next_power_of_2
17
18
logger = init_logger(__name__)
19
is_batch_invariant = vllm_is_batch_invariant()
@@ -223,14 +224,25 @@ def get_lora_op_configs(
224
# The default config for fused_moe_lora ops
225
elif op_type in [
226
"fused_moe_lora_w13_shrink",
226
- "fused_moe_lora_w13_expand",
227
"fused_moe_lora_w2_shrink",
228
+ ]:
229
+ default = {
230
+ "block_m": 64,
231
+ "block_n": min(64, next_power_of_2(rank)),
232
+ "block_k": 32,
233
+ "num_warps": 4,
234
+ "num_stages": 3,
235
+ "group_size_m": 8,
236
+ "split_k": 1,
237
+ }
238
+ elif op_type in [
239
+ "fused_moe_lora_w13_expand",
240
"fused_moe_lora_w2_expand",
241
]:
242
default = {
243
"block_m": 64,
244
"block_n": 64,
233
- "block_k": 32,
245
+ "block_k": max(16, min(32, next_power_of_2(rank))),
246
"num_warps": 4,
247
"num_stages": 3,
248
"group_size_m": 8,