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 14, 05:30 UTC
[Bugfix][LoRA][FusedMoE] Select MxFP4 Backend based on LoRA Enablement (#27487)
Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com> Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
Varun Sundar Rabindranath committed
Oct 27, 2025 at 10:32 UTC
5d3be3ba4c109f3ce5f0ffb78245296f3b9aa5ac
3 files changed
+34
-6
vllm/model_executor/layers/fused_moe/config.py
+2
@@ -825,6 +825,8 @@ class FusedMoEConfig:
825
826
is_act_and_mul: bool = True
827
828
+ is_lora_enabled: bool = False
829
+
830
def __post_init__(self):
831
if self.dp_size > 1:
832
logger.debug_once(
vllm/model_executor/layers/fused_moe/layer.py
+12
-3
@@ -982,6 +982,7 @@ def maybe_roundup_hidden_size(
982
act_dtype: torch.dtype,
983
quant_config: QuantizationConfig | None,
984
moe_parallel_config: FusedMoEParallelConfig,
985
+ is_lora_enabled: bool,
986
) -> int:
987
"""
988
Given layer hidden size and MoE configurations, round up hidden_size
@@ -992,6 +993,9 @@ def maybe_roundup_hidden_size(
993
act_dtype: Data type of the layer activations.
994
quant_config: Fused MoE quantization configuration.
995
moe_parallel_config: Fused MoE parallelization strategy configuration.
996
+ is_lora_enabled: True if the engine is enabled with LoRA. This
997
+ is used in the case of mxfp4 quantization in selecting the
998
+ MxFP4Backend.
999
1000
Return:
1001
Rounded up hidden_size if rounding up is required based on the configs.
@@ -1015,7 +1019,7 @@ def maybe_roundup_hidden_size(
1019
get_mxfp4_backend,
1020
)
1021
1018
- current_mxfp4_backend = get_mxfp4_backend()
1022
+ current_mxfp4_backend = get_mxfp4_backend(is_lora_enabled)
1023
if (
1024
current_mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16
1025
or current_mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS
@@ -1139,7 +1143,11 @@ class FusedMoE(CustomOp):
1143
1144
# Round up hidden size if needed.
1145
hidden_size = maybe_roundup_hidden_size(
1142
- hidden_size, moe_in_dtype, quant_config, self.moe_parallel_config
1146
+ hidden_size,
1147
+ moe_in_dtype,
1148
+ quant_config,
1149
+ self.moe_parallel_config,
1150
+ is_lora_enabled=self.vllm_config.lora_config is not None,
1151
)
1152
1153
# For smuggling this layer into the fused moe custom op
@@ -1270,8 +1278,9 @@ class FusedMoE(CustomOp):
1278
max_num_tokens=envs.VLLM_MOE_DP_CHUNK_SIZE,
1279
has_bias=has_bias,
1280
is_act_and_mul=is_act_and_mul,
1281
+ is_lora_enabled=vllm_config.lora_config is not None,
1282
)
1274
- self.moe_config = moe
1283
+ self.moe_config: FusedMoEConfig = moe
1284
self.moe_quant_config: FusedMoEQuantConfig | None = None
1285
self.quant_config = quant_config
1286
vllm/model_executor/layers/quantization/mxfp4.py
+20
-3
@@ -73,8 +73,24 @@ class Mxfp4Backend(Enum):
73
TRITON = 6
74
75
76
-def get_mxfp4_backend():
76
+def get_mxfp4_backend_with_lora() -> Mxfp4Backend:
77
+ """
78
+ Not all MXFP4 backends support LoRA. Select backends that are known to
79
+ have LoRA support.
80
+ """
81
+ if not current_platform.is_cuda():
82
+ return Mxfp4Backend.NONE
83
+
84
+ logger.info_once("[get_mxfp4_backend_with_lora] Using Marlin backend")
85
+ return Mxfp4Backend.MARLIN
86
+
87
+
88
+def get_mxfp4_backend(with_lora_support: bool) -> Mxfp4Backend:
89
# Backend Selection
90
+
91
+ if with_lora_support:
92
+ return get_mxfp4_backend_with_lora()
93
+
94
if current_platform.is_cuda():
95
if (
96
current_platform.is_device_capability(90)
@@ -183,13 +199,14 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
199
super().__init__(moe)
200
self.topk_indices_dtype = None
201
self.moe = moe
186
- self.mxfp4_backend = get_mxfp4_backend()
202
+ self.mxfp4_backend = get_mxfp4_backend(moe.is_lora_enabled)
203
self.max_capture_size = (
204
get_current_vllm_config().compilation_config.max_cudagraph_capture_size
205
)
206
207
assert self.mxfp4_backend != Mxfp4Backend.NONE, (
192
- "No MXFP4 MoE backend (FlashInfer/Marlin/Triton) available."
208
+ f"get_mxfp4_backend(with_lora_support={moe.is_lora_enabled}) found"
209
+ "no compatible MXFP4 MoE backend (FlashInfer/Marlin/Triton)."
210
"Please check your environment and try again."
211
)
212
self._cache_permute_indices: dict[torch.Size, torch.Tensor] = {}