[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>

tests/lora/test_punica_ops.py

@@ -482,3 +482,127 @@ def test_kernels_hidden_size(
             seq_length=128,
             add_inputs=True,
         )
+
+
+@pytest.mark.parametrize("device", DEVICES)
+def test_add_lora_fused_moe_early_exit(device):
+    """
+    Ensures add_lora_fused_moe does not invoke the LoRA kernel or
+    modify the output tensor when no_lora_flag_cpu is True
+    """
+    from types import SimpleNamespace
+
+    from vllm.lora.punica_wrapper.punica_gpu import PunicaWrapperGPU
+
+    torch.set_default_device(device)
+    torch.accelerator.set_device_index(device)
+
+    max_loras, num_tokens = 4, 16
+    num_experts, top_k, max_lora_rank = 8, 2, 16
+    K, N = 256, 128
+
+    # build PunicaWrapperGPU with minimal lora_config mock
+    lora_config = SimpleNamespace(
+        max_loras=max_loras,
+        specialize_active_lora=False,
+    )
+    wrapper = PunicaWrapperGPU(
+        max_num_batched_tokens=num_tokens,
+        max_batches=num_tokens,
+        device=device,
+        lora_config=lora_config,
+    )
+
+    # simulate a prior LoRA batch so the internal mapping is
+    # populated with stale LoRA IDs
+    lora_mapping = torch.zeros(
+        num_tokens,
+        dtype=torch.int32,
+        device=device,
+    )
+    lora_mapping[:8] = 1
+    lora_mapping[8:] = 2
+    wrapper.token_mapping_meta.prepare_tensors(lora_mapping)
+
+    # simulate a base-model batch (all -1)
+    base_mapping = torch.full(
+        (num_tokens,),
+        -1,
+        dtype=torch.int32,
+        device=device,
+    )
+    wrapper.token_mapping_meta.prepare_tensors(base_mapping)
+
+    assert wrapper.token_mapping_meta.no_lora_flag_cpu[0].item() is True
+
+    # dummy tensors for add_lora_fused_moe
+    y = torch.rand(num_tokens, top_k, N, dtype=torch.bfloat16, device=device)
+    y_snapshot = y.clone()
+    x = torch.rand(num_tokens, K, dtype=torch.bfloat16, device=device)
+
+    lora_a_stacked = (
+        torch.rand(
+            max_loras,
+            num_experts,
+            max_lora_rank,
+            K,
+            dtype=torch.bfloat16,
+            device=device,
+        ),
+    )
+    lora_b_stacked = (
+        torch.rand(
+            max_loras,
+            num_experts,
+            N,
+            max_lora_rank,
+            dtype=torch.bfloat16,
+            device=device,
+        ),
+    )
+    topk_weights = torch.ones(
+        num_tokens,
+        top_k,
+        dtype=torch.float32,
+        device=device,
+    )
+    adapter_enabled = torch.ones(
+        max_loras + 1,
+        dtype=torch.int32,
+        device=device,
+    )
+    shrink_config = expand_config = {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 32,
+        "BLOCK_SIZE_K": 64,
+        "GROUP_SIZE_M": 1,
+        "NUM_WARPS": 4,
+        "NUM_STAGES": 3,
+        "SPLIT_K": 1,
+    }
+
+    # call add_lora_fused_moe - the early exit should prevent any
+    # modification to the output
+    wrapper.add_lora_fused_moe(
+        y=y,
+        x=x,
+        lora_a_stacked=lora_a_stacked,
+        lora_b_stacked=lora_b_stacked,
+        topk_weights=topk_weights,
+        sorted_token_ids=None,
+        expert_ids=torch.zeros(
+            num_tokens * top_k,
+            dtype=torch.int32,
+            device=device,
+        ),
+        num_tokens_post_padded=None,
+        max_lora_rank=max_lora_rank,
+        top_k_num=top_k,
+        shrink_config=shrink_config,
+        expand_config=expand_config,
+        adapter_enabled=adapter_enabled,
+    )
+
+    assert torch.equal(y, y_snapshot), (
+        "add_lora_fused_moe modified output tensor despite no_lora_flag_cpu=True"
+    )

vllm/lora/punica_wrapper/punica_gpu.py

@@ -446,11 +446,17 @@ class PunicaWrapperGPU(PunicaWrapperBase):
             _,
             _,
             lora_ids,
-            _,
+            no_lora_flag,
             num_active_loras,
         ) = self.token_mapping_meta.meta_args(
             x.size(0), self.lora_config.specialize_active_lora
         )
+
+        assert no_lora_flag.numel() == 1
+        if no_lora_flag.item():
+            # None of the inputs require LoRA.
+            return
+
         if token_lora_mapping is None:
             token_lora_mapping = token_lora_mapping_meta
         fused_moe_lora(

vllm/model_executor/layers/fused_moe/experts/lora_context.py

@@ -59,3 +59,10 @@ class MoELoRAContext:
     # None means no dispatch happened (non-EP path), in which case callers
     # fall back to punica_wrapper.token_mapping_meta.
     local_token_lora_mapping: torch.Tensor | None = None
+
+    # Original unquantized hidden states, stashed by the modular kernel
+    # before the prepare step potentially quantizes them. Used by
+    # apply_w13_lora so the LoRA kernel sees correct-magnitude activations
+    # instead of raw quantized values that are missing the activation scale.
+    # Set per forward pass; None until the modular kernel writes it.
+    original_hidden_states: torch.Tensor | None = None

vllm/model_executor/layers/fused_moe/experts/triton_moe.py

@@ -77,6 +77,16 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
     def activation_format() -> mk.FusedMoEActivationFormat:
         return mk.FusedMoEActivationFormat.Standard
 
+    @property
+    def expects_unquantized_inputs(self) -> bool:
+        # Defer activation quantization to apply() only when LoRA is active AND
+        # tokens are dispatched across ranks (DP+EP all2all).
+        return (
+            self._lora_context is not None
+            and self.quant_dtype is not None
+            and self.moe_config.moe_parallel_config.use_all2all_kernels
+        )
+
     @staticmethod
     def _supports_current_device() -> bool:
         return current_platform.is_cuda_alike() or current_platform.is_xpu()
@@ -223,6 +233,25 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
             torch.float8_e4m3fnuz,
         ]
 
+        # We declared expects_unquantized_inputs (LoRA + DP/EP all2all), so the
+        # prepare step deferred activation quantization to this kernel:
+        # `hidden_states` arrives unquantized. Keep the unquantized tensor for
+        # the LoRA shrink input and quantize a copy here for the base GEMM
+        # (mirrors what the prepare step would have done, but after the
+        # all-gather so the layout matches the gathered topk_ids / token map).
+        lora_unquantized_hidden_states: torch.Tensor | None = None
+        if self.expects_unquantized_inputs:
+            assert a1q_scale is None
+            lora_unquantized_hidden_states = hidden_states
+            hidden_states, a1q_scale = moe_kernel_quantize_input(
+                hidden_states,
+                self.a1_scale,
+                self.quant_dtype,
+                self.per_act_token_quant,
+                self.block_shape,
+                quantization_emulation=self.quantization_emulation,
+            )
+
         E, num_tokens, N, K, top_k_num = self.moe_problem_size(
             hidden_states, w1, w2, topk_ids
         )
@@ -280,12 +309,28 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
         # GEMM on the default stream and the LoRA fast-path on aux_stream;
         # the LoRA writes its delta into a fresh zero buffer (add_inputs=
         # False) and we sum it into intermediate_cache1 after both finish.
-
+        #
+        # The LoRA shrink kernel needs unquantized, gathered-layout
+        # activations. When activation quant was deferred to this kernel
+        # (expects_unquantized_inputs), the input we quantized above is exactly
+        # that, so use it directly. Otherwise fall back to the context stash
+        # (e.g. weight-only quant), guarding on a row-count match so a
+        # DP-gathered layout never indexes a local stash out of bounds.
         sorted_token_ids_lora = None
         expert_ids_lora = None
         num_tokens_post_padded_lora = None
         token_lora_mapping = None
         lora_context = self._lora_context
+        if lora_unquantized_hidden_states is not None:
+            lora_x = lora_unquantized_hidden_states
+        elif (
+            lora_context is not None
+            and lora_context.original_hidden_states is not None
+            and lora_context.original_hidden_states.shape[0] == hidden_states.shape[0]
+        ):
+            lora_x = lora_context.original_hidden_states
+        else:
+            lora_x = hidden_states
 
         def _base_w13_fn():
             invoke_fused_moe_triton_kernel(
@@ -322,7 +367,7 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
                 return self.apply_w13_lora(
                     lora_context,
                     y=lora_delta_w13,
-                    x=hidden_states,
+                    x=lora_x,
                     topk_ids=topk_ids,
                     topk_weights=topk_weights,
                     expert_map=expert_map,
@@ -359,7 +404,7 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
                 ) = self.apply_w13_lora(
                     lora_context,
                     y=intermediate_cache1,
-                    x=hidden_states,
+                    x=lora_x,
                     topk_ids=topk_ids,
                     topk_weights=topk_weights,
                     expert_map=expert_map,

vllm/model_executor/layers/fused_moe/modular_kernel.py

@@ -1407,6 +1407,13 @@ class FusedMoEKernelModularImpl:
             apply_router_weight_on_input,
         )
 
+        # Stash the original unquantized hidden states on the LoRA context
+        # so apply_w13_lora sees correct-magnitude activations instead of
+        # the potentially quantized values produced by _prepare().
+        lora_ctx = getattr(self.fused_experts, "_lora_context", None)
+        if lora_ctx is not None:
+            lora_ctx.original_hidden_states = hidden_states
+
         fused_out = self._fused_experts(
             in_dtype=hidden_states.dtype,
             a1q=a1q,
@@ -1424,6 +1431,9 @@ class FusedMoEKernelModularImpl:
             output_alias=output,
         )
 
+        if lora_ctx is not None:
+            lora_ctx.original_hidden_states = None
+
         return self._finalize(
             output,
             fused_out,