diff --git a/tests/lora/test_punica_ops.py b/tests/lora/test_punica_ops.py index 7706d0e2aa..be87847262 100644 --- a/tests/lora/test_punica_ops.py +++ b/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" + ) diff --git a/vllm/lora/punica_wrapper/punica_gpu.py b/vllm/lora/punica_wrapper/punica_gpu.py index ccf95eb684..18272354b4 100644 --- a/vllm/lora/punica_wrapper/punica_gpu.py +++ b/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( diff --git a/vllm/model_executor/layers/fused_moe/experts/lora_context.py b/vllm/model_executor/layers/fused_moe/experts/lora_context.py index 404457bb34..117f744aee 100644 --- a/vllm/model_executor/layers/fused_moe/experts/lora_context.py +++ b/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 diff --git a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py index d81458b375..0d9b43658f 100644 --- a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py +++ b/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, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index e80224be70..0e55e827c2 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/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,