[Bugfix][LoRA] Fall back for high-rank MoE LoRA (#55161)
Signed-off-by: Raina Zhong <yuweiz@nvidia.com> Co-authored-by: Raina Zhong <yuweiz@nvidia.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
xiaoyu-xyz committed
Oct 2, 2026 at 18:59 UTC
8c60714da541992a883a138b4f08df9558018b76
2 files changed
+94
-11
tests/lora/test_fused_moe_lora_kernel.py
+77
@@ -362,6 +362,83 @@ def test_fused_moe_lora_kernel(
362
torch.testing.assert_close(output, output2, atol=1e-2, rtol=1e-2)
363
364
365
+@pytest.mark.parametrize("max_lora_rank", [128, 256])
366
+@pytest.mark.parametrize("add_inputs", [True, False])
367
+@pytest.mark.parametrize("num_slices", [1, 2])
368
+@pytest.mark.parametrize("device", DEVICES)
369
+def test_fused_moe_lora_kernel_large_rank_fallback(
370
+ max_lora_rank,
371
+ add_inputs,
372
+ num_slices,
373
+ device,
374
+):
375
+ """Ranks above the one-shot limit fall back to shrink/expand kernels."""
376
+ torch.set_default_device(device)
377
+ set_random_seed(42)
378
+
379
+ num_tokens = 8
380
+ num_sequences = 4
381
+ top_k_num = 2
382
+ num_experts = 4
383
+ max_loras = 2
384
+ N = 256
385
+ K = 256
386
+ block_size = 16
387
+ dtype = torch.bfloat16
388
+
389
+ topk_ids, topk_weights, token_lora_mapping, lora_ids = sample_data(
390
+ num_tokens, num_sequences, max_loras, num_experts, top_k_num
391
+ )
392
+ lora_a_stacked = [
393
+ torch.rand(
394
+ (max_loras, num_experts, max_lora_rank, K),
395
+ dtype=dtype,
396
+ )
397
+ for _ in range(num_slices)
398
+ ]
399
+ lora_b_stacked = [
400
+ torch.rand(
401
+ (max_loras, num_experts, N // num_slices, max_lora_rank),
402
+ dtype=dtype,
403
+ )
404
+ for _ in range(num_slices)
405
+ ]
406
+ hidden_states = torch.rand((num_tokens, K), dtype=dtype)
407
+
408
+ residual = torch.randn((num_tokens, top_k_num, N), dtype=dtype)
409
+ output = residual.clone()
410
+ use_fused_moe_lora_kernel(
411
+ topk_ids,
412
+ topk_weights,
413
+ token_lora_mapping,
414
+ max_lora_rank,
415
+ top_k_num,
416
+ lora_ids,
417
+ lora_a_stacked,
418
+ lora_b_stacked,
419
+ hidden_states,
420
+ output,
421
+ max_loras,
422
+ num_experts,
423
+ block_size,
424
+ add_inputs=add_inputs,
425
+ )
426
+
427
+ expected = use_torch(
428
+ hidden_states,
429
+ token_lora_mapping,
430
+ topk_ids,
431
+ lora_a_stacked,
432
+ lora_b_stacked,
433
+ top_k_num,
434
+ num_slices,
435
+ )
436
+ if add_inputs:
437
+ expected += residual
438
+
439
+ torch.testing.assert_close(output, expected, atol=2e-2, rtol=2e-2)
440
+
441
+
442
def use_fused_moe_lora_kernel_naive(
443
topk_ids,
444
topk_weights,
vllm/lora/ops/triton_ops/fused_moe_lora_op.py
+17
-11
@@ -106,14 +106,18 @@ def _get_c_ptrs(
106
107
_LORA_PTR_DICT: dict[tuple[int, ...], torch.tensor] = {}
108
109
+# The one-shot kernel keeps the rank-dimension intermediate in registers.
110
+# Larger ranks require rank tiling, so route them through the two-kernel path.
111
+_FUSED_MOE_LORA_ONE_SHOT_MAX_RANK = 128
112
+
113
114
# ---------------------------------------------------------------------------
115
# Fully-fused MoE-LoRA kernel (one-shot): shrink + expand combined into a single
116
# launch with the rank-dim intermediate kept in registers. Used by the fast
113
-# path of `_fused_moe_lora` for `fully_sharded=False`. The legacy two-kernel
114
-# path (`_fused_moe_lora_kernel` above) is retained for `fully_sharded=True`
115
-# because that path needs to materialise the intermediate cache for an
116
-# all_reduce / all_gather between shrink and expand.
117
+# path of `_fused_moe_lora` for supported ranks with `fully_sharded=False`.
118
+# The legacy two-kernel path (`_fused_moe_lora_kernel` above) is retained for
119
+# larger ranks and for `fully_sharded=True`, where the intermediate cache must
120
+# be materialised for an all_reduce / all_gather between shrink and expand.
121
# ---------------------------------------------------------------------------
122
123
@@ -349,8 +353,9 @@ def _run_fused_moe_lora_one_shot(
353
# regressed across all M (+8 to +40%): the (64,32) fp32 accumulator +
354
# widened B tile pushed register count past spill threshold, lowering
355
# occupancy by more than the MMA gain saved.
352
- assert rank <= 128, (
353
- f"fused_moe_lora_one_shot supports max_lora_rank<=128; got rank={rank}"
356
+ assert rank <= _FUSED_MOE_LORA_ONE_SHOT_MAX_RANK, (
357
+ "fused_moe_lora_one_shot supports "
358
+ f"max_lora_rank<={_FUSED_MOE_LORA_ONE_SHOT_MAX_RANK}; got rank={rank}"
359
)
360
BLOCK_R = max(triton.next_power_of_2(rank), 16)
361
@@ -1318,6 +1323,7 @@ def _fused_moe_lora_expand(
1323
offset: int = 0,
1324
use_gdc: bool = False,
1325
use_tma: bool = False,
1326
+ add_inputs: bool = True,
1327
) -> None:
1328
b_ptr = _get_ptr(lora_b_stacked, device)
1329
K = max_lora_rank
@@ -1408,7 +1414,7 @@ def _fused_moe_lora_expand(
1414
token_mapping_factor=1,
1415
naive_block_assignment=sorted_token_ids is None,
1416
MUL_ROUTED_WEIGHT=mul_routed_weight,
1411
- ADD_INPUTS=True,
1417
+ ADD_INPUTS=add_inputs,
1418
USE_B_L2_CACHE=True,
1419
sort_c=False,
1420
IS_PRIMARY=False,
@@ -1478,7 +1484,7 @@ def _fused_moe_lora(
1484
assert top_k_num == topk_weights.shape[1]
1485
1486
# Fast path: single fused kernel
1481
- if not fully_sharded:
1487
+ if not fully_sharded and max_lora_rank <= _FUSED_MOE_LORA_ONE_SHOT_MAX_RANK:
1488
M_pairs = topk_weights.numel()
1489
if (
1490
sorted_token_ids is None
@@ -1524,9 +1530,8 @@ def _fused_moe_lora(
1530
)
1531
return
1532
1527
- assert add_inputs, (
1528
- "fused_moe_lora(add_inputs=False) is only supported on the "
1529
- "fully_sharded=False fast path"
1533
+ assert add_inputs or not fully_sharded, (
1534
+ "fused_moe_lora(add_inputs=False) is not supported with fully_sharded=True"
1535
)
1536
1537
device = qcurr_hidden_states.device
@@ -1660,6 +1665,7 @@ def _fused_moe_lora(
1665
offset,
1666
use_gdc=use_gdc,
1667
use_tma=use_tma,
1668
+ add_inputs=add_inputs,
1669
)
1670
1671