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, 23:45 UTC
Fix MoE backend selection for LoRA (unquantized MoE) (#40273)
Signed-off-by: Daniel Serebrenik <daserebrenik@nvidia.com>
danisereb committed
Apr 19, 2026 at 20:18 UTC
d1135a508705cc899946eae2d0cb63b7d9e94bfe
2 files changed
+90
tests/kernels/moe/test_unquantized_backend_selection.py
+85
@@ -11,6 +11,11 @@ from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
11
)
12
from vllm.platforms import current_platform
13
14
+skipif_not_cuda_rocm = pytest.mark.skipif(
15
+ not (current_platform.is_cuda() or current_platform.is_rocm()),
16
+ reason="Only supported on CUDA/ROCm platforms.",
17
+)
18
+
19
20
@pytest.mark.parametrize(
21
"platform_method,expected_backend",
@@ -190,3 +195,83 @@ def test_select_cuda_flashinfer_cutlass_backend(
195
196
assert selected_backend == UnquantizedMoeBackend.FLASHINFER_CUTLASS
197
assert experts_cls is not None
198
+
199
+
200
+@skipif_not_cuda_rocm
201
+def test_select_lora_backend_prefers_triton():
202
+ """LoRA-enabled unquantized MoE should select Triton backend."""
203
+ moe_config = make_dummy_moe_config()
204
+ moe_config.is_lora_enabled = True
205
+ selected_backend, experts_cls = select_unquantized_moe_backend(
206
+ moe_config=moe_config
207
+ )
208
+
209
+ assert selected_backend == UnquantizedMoeBackend.TRITON
210
+ assert experts_cls is not None
211
+
212
+
213
+@skipif_not_cuda_rocm
214
+def test_select_lora_explicit_non_triton_backend():
215
+ """LoRA should override explicit non-Triton backend to Triton."""
216
+ moe_config = make_dummy_moe_config()
217
+ moe_config.is_lora_enabled = True
218
+
219
+ # Use string from mapping in function map_unquantized_backend()
220
+ moe_config.moe_backend = "flashinfer_cutlass"
221
+
222
+ selected_backend, experts_cls = select_unquantized_moe_backend(
223
+ moe_config=moe_config
224
+ )
225
+
226
+ assert selected_backend == UnquantizedMoeBackend.TRITON
227
+ assert experts_cls is not None
228
+
229
+
230
+@skipif_not_cuda_rocm
231
+@pytest.mark.parametrize("is_lora_enabled", [False, True])
232
+def test_select_explicit_triton_backend(is_lora_enabled):
233
+ """Explicit triton backend selection should return Triton."""
234
+ moe_config = make_dummy_moe_config()
235
+ moe_config.is_lora_enabled = is_lora_enabled
236
+ moe_config.moe_backend = "triton"
237
+
238
+ selected_backend, experts_cls = select_unquantized_moe_backend(
239
+ moe_config=moe_config
240
+ )
241
+
242
+ assert selected_backend == UnquantizedMoeBackend.TRITON
243
+ assert experts_cls is not None
244
+
245
+
246
+@skipif_not_cuda_rocm
247
+def test_select_explicit_triton_ignores_flashinfer_env(monkeypatch):
248
+ """Explicit triton backend should override FlashInfer env selection."""
249
+ monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP16", "1")
250
+ monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "throughput")
251
+
252
+ moe_config = make_dummy_moe_config()
253
+ moe_config.is_lora_enabled = False
254
+ moe_config.moe_backend = "triton"
255
+
256
+ selected_backend, experts_cls = select_unquantized_moe_backend(
257
+ moe_config=moe_config
258
+ )
259
+
260
+ assert selected_backend == UnquantizedMoeBackend.TRITON
261
+ assert experts_cls is not None
262
+
263
+
264
+@skipif_not_cuda_rocm
265
+def test_select_lora_ignores_flashinfer_env(monkeypatch):
266
+ """LoRA path should still choose Triton even if FlashInfer env is on."""
267
+ monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP16", "1")
268
+ monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "throughput")
269
+
270
+ moe_config = make_dummy_moe_config()
271
+ moe_config.is_lora_enabled = True
272
+ selected_backend, experts_cls = select_unquantized_moe_backend(
273
+ moe_config=moe_config
274
+ )
275
+
276
+ assert selected_backend == UnquantizedMoeBackend.TRITON
277
+ assert experts_cls is not None
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
+5
@@ -163,6 +163,11 @@ def select_unquantized_moe_backend(
163
if current_platform.is_out_of_tree():
164
return UnquantizedMoeBackend.OOT, None
165
166
+ if moe_config.is_lora_enabled:
167
+ return UnquantizedMoeBackend.TRITON, backend_to_kernel_cls(
168
+ UnquantizedMoeBackend.TRITON
169
+ )
170
+
171
# NOTE: the kernels are selected in the following order.
172
AVAILABLE_BACKENDS = _get_priority_backends(moe_config)
173