[Bugfix] Fix corrupt outputs in MoE FP8 LoRA responses and MoE base model responses when LoRAs are loaded (#42120)
Signed-off-by: Nicholas Edelman <nedelman@nvidia.com> Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
nv-nedelman-1 committed
Jun 18, 2026 at 20:26 UTC
9ea3a4015b412d146d38ee1b697aafe92979c6ae
5 files changed
+196
-4
tests/lora/test_punica_ops.py
+124
@@ -482,3 +482,127 @@ def test_kernels_hidden_size(
482
seq_length=128,
483
add_inputs=True,
484
)
485
+
486
+
487
+@pytest.mark.parametrize("device", DEVICES)
488
+def test_add_lora_fused_moe_early_exit(device):
489
+ """
490
+ Ensures add_lora_fused_moe does not invoke the LoRA kernel or
491
+ modify the output tensor when no_lora_flag_cpu is True
492
+ """
493
+ from types import SimpleNamespace
494
+
495
+ from vllm.lora.punica_wrapper.punica_gpu import PunicaWrapperGPU
496
+
497
+ torch.set_default_device(device)
498
+ torch.accelerator.set_device_index(device)
499
+
500
+ max_loras, num_tokens = 4, 16
501
+ num_experts, top_k, max_lora_rank = 8, 2, 16
502
+ K, N = 256, 128
503
+
504
+ # build PunicaWrapperGPU with minimal lora_config mock
505
+ lora_config = SimpleNamespace(
506
+ max_loras=max_loras,
507
+ specialize_active_lora=False,
508
+ )
509
+ wrapper = PunicaWrapperGPU(
510
+ max_num_batched_tokens=num_tokens,
511
+ max_batches=num_tokens,
512
+ device=device,
513
+ lora_config=lora_config,
514
+ )
515
+
516
+ # simulate a prior LoRA batch so the internal mapping is
517
+ # populated with stale LoRA IDs
518
+ lora_mapping = torch.zeros(
519
+ num_tokens,
520
+ dtype=torch.int32,
521
+ device=device,
522
+ )
523
+ lora_mapping[:8] = 1
524
+ lora_mapping[8:] = 2
525
+ wrapper.token_mapping_meta.prepare_tensors(lora_mapping)
526
+
527
+ # simulate a base-model batch (all -1)
528
+ base_mapping = torch.full(
529
+ (num_tokens,),
530
+ -1,
531
+ dtype=torch.int32,
532
+ device=device,
533
+ )
534
+ wrapper.token_mapping_meta.prepare_tensors(base_mapping)
535
+
536
+ assert wrapper.token_mapping_meta.no_lora_flag_cpu[0].item() is True
537
+
538
+ # dummy tensors for add_lora_fused_moe
539
+ y = torch.rand(num_tokens, top_k, N, dtype=torch.bfloat16, device=device)
540
+ y_snapshot = y.clone()
541
+ x = torch.rand(num_tokens, K, dtype=torch.bfloat16, device=device)
542
+
543
+ lora_a_stacked = (
544
+ torch.rand(
545
+ max_loras,
546
+ num_experts,
547
+ max_lora_rank,
548
+ K,
549
+ dtype=torch.bfloat16,
550
+ device=device,
551
+ ),
552
+ )
553
+ lora_b_stacked = (
554
+ torch.rand(
555
+ max_loras,
556
+ num_experts,
557
+ N,
558
+ max_lora_rank,
559
+ dtype=torch.bfloat16,
560
+ device=device,
561
+ ),
562
+ )
563
+ topk_weights = torch.ones(
564
+ num_tokens,
565
+ top_k,
566
+ dtype=torch.float32,
567
+ device=device,
568
+ )
569
+ adapter_enabled = torch.ones(
570
+ max_loras + 1,
571
+ dtype=torch.int32,
572
+ device=device,
573
+ )
574
+ shrink_config = expand_config = {
575
+ "BLOCK_SIZE_M": 16,
576
+ "BLOCK_SIZE_N": 32,
577
+ "BLOCK_SIZE_K": 64,
578
+ "GROUP_SIZE_M": 1,
579
+ "NUM_WARPS": 4,
580
+ "NUM_STAGES": 3,
581
+ "SPLIT_K": 1,
582
+ }
583
+
584
+ # call add_lora_fused_moe - the early exit should prevent any
585
+ # modification to the output
586
+ wrapper.add_lora_fused_moe(
587
+ y=y,
588
+ x=x,
589
+ lora_a_stacked=lora_a_stacked,
590
+ lora_b_stacked=lora_b_stacked,
591
+ topk_weights=topk_weights,
592
+ sorted_token_ids=None,
593
+ expert_ids=torch.zeros(
594
+ num_tokens * top_k,
595
+ dtype=torch.int32,
596
+ device=device,
597
+ ),
598
+ num_tokens_post_padded=None,
599
+ max_lora_rank=max_lora_rank,
600
+ top_k_num=top_k,
601
+ shrink_config=shrink_config,
602
+ expand_config=expand_config,
603
+ adapter_enabled=adapter_enabled,
604
+ )
605
+
606
+ assert torch.equal(y, y_snapshot), (
607
+ "add_lora_fused_moe modified output tensor despite no_lora_flag_cpu=True"
608
+ )
vllm/lora/punica_wrapper/punica_gpu.py
+7
-1
@@ -446,11 +446,17 @@ class PunicaWrapperGPU(PunicaWrapperBase):
446
_,
447
_,
448
lora_ids,
449
- _,
449
+ no_lora_flag,
450
num_active_loras,
451
) = self.token_mapping_meta.meta_args(
452
x.size(0), self.lora_config.specialize_active_lora
453
)
454
+
455
+ assert no_lora_flag.numel() == 1
456
+ if no_lora_flag.item():
457
+ # None of the inputs require LoRA.
458
+ return
459
+
460
if token_lora_mapping is None:
461
token_lora_mapping = token_lora_mapping_meta
462
fused_moe_lora(
vllm/model_executor/layers/fused_moe/experts/lora_context.py
+7
@@ -59,3 +59,10 @@ class MoELoRAContext:
59
# None means no dispatch happened (non-EP path), in which case callers
60
# fall back to punica_wrapper.token_mapping_meta.
61
local_token_lora_mapping: torch.Tensor | None = None
62
+
63
+ # Original unquantized hidden states, stashed by the modular kernel
64
+ # before the prepare step potentially quantizes them. Used by
65
+ # apply_w13_lora so the LoRA kernel sees correct-magnitude activations
66
+ # instead of raw quantized values that are missing the activation scale.
67
+ # Set per forward pass; None until the modular kernel writes it.
68
+ original_hidden_states: torch.Tensor | None = None
vllm/model_executor/layers/fused_moe/experts/triton_moe.py
+48
-3
@@ -77,6 +77,16 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
77
def activation_format() -> mk.FusedMoEActivationFormat:
78
return mk.FusedMoEActivationFormat.Standard
79
80
+ @property
81
+ def expects_unquantized_inputs(self) -> bool:
82
+ # Defer activation quantization to apply() only when LoRA is active AND
83
+ # tokens are dispatched across ranks (DP+EP all2all).
84
+ return (
85
+ self._lora_context is not None
86
+ and self.quant_dtype is not None
87
+ and self.moe_config.moe_parallel_config.use_all2all_kernels
88
+ )
89
+
90
@staticmethod
91
def _supports_current_device() -> bool:
92
return current_platform.is_cuda_alike() or current_platform.is_xpu()
@@ -223,6 +233,25 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
233
torch.float8_e4m3fnuz,
234
]
235
236
+ # We declared expects_unquantized_inputs (LoRA + DP/EP all2all), so the
237
+ # prepare step deferred activation quantization to this kernel:
238
+ # `hidden_states` arrives unquantized. Keep the unquantized tensor for
239
+ # the LoRA shrink input and quantize a copy here for the base GEMM
240
+ # (mirrors what the prepare step would have done, but after the
241
+ # all-gather so the layout matches the gathered topk_ids / token map).
242
+ lora_unquantized_hidden_states: torch.Tensor | None = None
243
+ if self.expects_unquantized_inputs:
244
+ assert a1q_scale is None
245
+ lora_unquantized_hidden_states = hidden_states
246
+ hidden_states, a1q_scale = moe_kernel_quantize_input(
247
+ hidden_states,
248
+ self.a1_scale,
249
+ self.quant_dtype,
250
+ self.per_act_token_quant,
251
+ self.block_shape,
252
+ quantization_emulation=self.quantization_emulation,
253
+ )
254
+
255
E, num_tokens, N, K, top_k_num = self.moe_problem_size(
256
hidden_states, w1, w2, topk_ids
257
)
@@ -280,12 +309,28 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
309
# GEMM on the default stream and the LoRA fast-path on aux_stream;
310
# the LoRA writes its delta into a fresh zero buffer (add_inputs=
311
# False) and we sum it into intermediate_cache1 after both finish.
283
-
312
+ #
313
+ # The LoRA shrink kernel needs unquantized, gathered-layout
314
+ # activations. When activation quant was deferred to this kernel
315
+ # (expects_unquantized_inputs), the input we quantized above is exactly
316
+ # that, so use it directly. Otherwise fall back to the context stash
317
+ # (e.g. weight-only quant), guarding on a row-count match so a
318
+ # DP-gathered layout never indexes a local stash out of bounds.
319
sorted_token_ids_lora = None
320
expert_ids_lora = None
321
num_tokens_post_padded_lora = None
322
token_lora_mapping = None
323
lora_context = self._lora_context
324
+ if lora_unquantized_hidden_states is not None:
325
+ lora_x = lora_unquantized_hidden_states
326
+ elif (
327
+ lora_context is not None
328
+ and lora_context.original_hidden_states is not None
329
+ and lora_context.original_hidden_states.shape[0] == hidden_states.shape[0]
330
+ ):
331
+ lora_x = lora_context.original_hidden_states
332
+ else:
333
+ lora_x = hidden_states
334
335
def _base_w13_fn():
336
invoke_fused_moe_triton_kernel(
@@ -322,7 +367,7 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
367
return self.apply_w13_lora(
368
lora_context,
369
y=lora_delta_w13,
325
- x=hidden_states,
370
+ x=lora_x,
371
topk_ids=topk_ids,
372
topk_weights=topk_weights,
373
expert_map=expert_map,
@@ -359,7 +404,7 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
404
) = self.apply_w13_lora(
405
lora_context,
406
y=intermediate_cache1,
362
- x=hidden_states,
407
+ x=lora_x,
408
topk_ids=topk_ids,
409
topk_weights=topk_weights,
410
expert_map=expert_map,
vllm/model_executor/layers/fused_moe/modular_kernel.py
+10
@@ -1407,6 +1407,13 @@ class FusedMoEKernelModularImpl:
1407
apply_router_weight_on_input,
1408
)
1409
1410
+ # Stash the original unquantized hidden states on the LoRA context
1411
+ # so apply_w13_lora sees correct-magnitude activations instead of
1412
+ # the potentially quantized values produced by _prepare().
1413
+ lora_ctx = getattr(self.fused_experts, "_lora_context", None)
1414
+ if lora_ctx is not None:
1415
+ lora_ctx.original_hidden_states = hidden_states
1416
+
1417
fused_out = self._fused_experts(
1418
in_dtype=hidden_states.dtype,
1419
a1q=a1q,
@@ -1424,6 +1431,9 @@ class FusedMoEKernelModularImpl:
1431
output_alias=output,
1432
)
1433
1434
+ if lora_ctx is not None:
1435
+ lora_ctx.original_hidden_states = None
1436
+
1437
return self._finalize(
1438
output,
1439
fused_out,