[Feature][LoRA][MoE] Expert LoRA for FlashInfer TRTLLM MXFP4 experts (#61979)

Signed-off-by: aoshen524 <aoshen524@gmail.com>
Co-authored-by: aoshen524 <aoshen524@gmail.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

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

@@ -1,8 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""LoRA-aware FlashInfer TRT-LLM MoE experts (BF16).
+"""LoRA-aware FlashInfer TRT-LLM MoE experts (BF16, MXFP4 x BF16).
 
-Reuses the routed API + ``gemm1_lora_delta`` path from FlashInfer PR #3153:
+Reuses the routed API + ``gemm1_lora_delta`` path from FlashInfer PR #3153
+(BF16) and PR #3987 (FP4 block-scale):
 
   * The W13 (gate_up) LoRA delta is passed directly as ``gemm1_lora_delta`` to
     the routed kernel, which fuses it into FC1 before SwiGLU (BiasType::Mn).
@@ -13,11 +14,12 @@ Reuses the routed API + ``gemm1_lora_delta`` path from FlashInfer PR #3153:
 
 Constraints (matching the PR support matrix; final gating lives in the oracle):
   * SM100+ (Blackwell), gated SwiGLU, shuffled weights only;
-  * BF16 only;
+  * BF16 weights, or MXFP4 weights with BF16 activations;
   * routing must be computed outside the MoE (the Modular path satisfies this).
 
 """
 
+import functools
 from abc import abstractmethod
 
 import torch
@@ -37,9 +39,17 @@ from vllm.model_executor.layers.fused_moe.experts.lora_experts_mixin import (
 from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import (
     view_as_block_major_k,
 )
+from vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe import (
+    TrtLlmMxfp4ExpertsBase,
+)
 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
     TopKWeightAndReduceNoOP,
 )
+from vllm.model_executor.layers.fused_moe.utils import fi_moe_largest_bucket
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    QuantKey,
+    kMxfp4Static,
+)
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
@@ -224,6 +234,15 @@ class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
         """
         raise NotImplementedError
 
+    def prepare_gemm1_lora_delta(
+        self,
+        gemm1_lora_delta: torch.Tensor,
+        lora_context: MoELoRAContext,
+        topk_ids: torch.Tensor,
+    ) -> torch.Tensor:
+        """Hook to rewrite the W13 delta before it is handed to the kernel."""
+        return gemm1_lora_delta
+
     def apply(
         self,
         output: torch.Tensor,
@@ -247,7 +266,8 @@ class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
         num_tokens = hidden_states.size(0)
         top_k = self.topk
         intermediate_size = self.intermediate_size_per_partition
-        K = output.size(1)
+        # LoRA buffers use the padded hidden size; `output` may be narrower.
+        K = self.hidden_dim
 
         # ---- Base-model fast path ----
         # When no token in the batch selects a LoRA adapter, skip the LoRA machinery
@@ -288,8 +308,7 @@ class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
         w13_meta = (None, None, None, None)
 
         # zeros (not empty): under EP the punica expand kernel only writes
-        # slots whose expert is local to this rank; non-local (token, top_k)
-        # slots must stay 0 so they contribute no bias when fed to flashinfer.
+        # slots whose expert is local to this rank; the kernel skips the rest.
         gemm1_lora_delta = torch.zeros(
             num_tokens,
             top_k,
@@ -326,6 +345,9 @@ class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
             add_inputs=False,
             swap_w13_slices=True,
         )
+        gemm1_lora_delta = self.prepare_gemm1_lora_delta(
+            gemm1_lora_delta, lora_context, topk_ids
+        )
 
         # ---- 2) Call the routed flashinfer kernel ----
         ret = self.invoke_routed_moe(
@@ -466,7 +488,9 @@ class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
         (``expert_weights`` in expanded order, ``idx_map < 0`` dropped), scale by
         ``scale``, and add the already-weighted ``w2_delta`` reduced over top_k.
         """
-        K = gemm2_permuted.size(1)
+        # Inputs may carry hidden padding that `output` does not.
+        K = output.size(1)
+        assert gemm2_permuted.size(1) >= K and w2_delta.size(2) >= K
         BLOCK_K = 512
         grid = (num_tokens, triton.cdiv(K, BLOCK_K))
         _finalize_lora_kernel[grid](
@@ -550,3 +574,158 @@ class TrtLlmBf16LoRAExperts(_TrtLlmLoRAExpertsBase):
             return list(ret)
         # do_finalize=True finalized directly into `output`.
         return [output]
+
+
+# MXFP4 weights x BF16 activations trtllm MoE + LoRA
+class TrtLlmMxfp4LoRAExperts(_TrtLlmLoRAExpertsBase, TrtLlmMxfp4ExpertsBase):
+    """MXFP4-weight, BF16-activation trtllm MoE + LoRA. Not MXFP8 activations:
+    the kernel then returns the FC1 activation quantized, without its scales."""
+
+    def __init__(
+        self,
+        moe_config: FusedMoEConfig,
+        quant_config: FusedMoEQuantConfig,
+    ):
+        _TrtLlmLoRAExpertsBase.__init__(self, moe_config, quant_config)
+        # gemm1_alpha / gemm1_beta / gemm1_clamp_limit, hidden_dim_unpadded.
+        TrtLlmMxfp4ExpertsBase.__init__(self, moe_config, quant_config)
+
+    @staticmethod
+    def _supports_quant_scheme(
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        return weight_key == kMxfp4Static and activation_key is None
+
+    @staticmethod
+    def _supports_activation(activation: MoEActivation) -> bool:
+        return activation in (MoEActivation.SWIGLUOAI, MoEActivation.SILU)
+
+    @staticmethod
+    def _supports_routing_method(routing_method, weight_key, activation_key) -> bool:
+        return True  # routing is done outside the kernel
+
+    @staticmethod
+    def _supports_parallel_config(
+        moe_parallel_config: FusedMoEParallelConfig,
+    ) -> bool:
+        # DP gathers tokens, but the no-LoRA fast path reads the local batch.
+        p = moe_parallel_config
+        return p.dp_size == 1 and not p.use_all2all_kernels and not p.enable_eplb
+
+    def workspace_shapes(
+        self,
+        M,
+        N,
+        K,
+        topk,
+        global_num_experts,
+        local_num_experts,
+        expert_tokens_meta,
+        activation,
+    ):
+        # flashinfer manages its own workspace and drops the hidden padding.
+        return (0,), (0,), (M, self.hidden_dim_unpadded)
+
+    def output_hidden_dim(self, hidden_dim: int) -> int:
+        # FusedMoEExpertsModular precedes TrtLlmMxfp4ExpertsBase in the MRO.
+        return self.hidden_dim_unpadded
+
+    @functools.cached_property
+    def _w13_bias_up_gate_index(self) -> torch.Tensor:
+        """Columns of the kernel-format FC1 bias in [up | gate] order."""
+        from flashinfer.fused_moe.core import get_w2_permute_indices_with_cache
+
+        # Kernel format is the (up_i, gate_i) bias permuted for the tile.
+        assert self.w1_bias is not None
+        bias = self.w1_bias.reshape(self.local_num_experts, -1)[0]
+        permute = get_w2_permute_indices_with_cache({}, bias.reshape(-1, 1), 128)
+        interleaved = torch.argsort(permute.to(bias.device))
+        return torch.cat([interleaved[0::2], interleaved[1::2]])
+
+    def prepare_gemm1_lora_delta(
+        self,
+        gemm1_lora_delta: torch.Tensor,
+        lora_context: MoELoRAContext,
+        topk_ids: torch.Tensor,
+    ) -> torch.Tensor:
+        """Bring the W13 delta to [up | gate] halves and fold in the FC1 bias."""
+        intermediate_size = gemm1_lora_delta.size(-1) // 2
+        # View as (I, 2) pairs of the first and second w13 half. SWIGLUOAI
+        # weights, and the deltas that follow them, are interleaved.
+        if self.moe_config.activation == MoEActivation.SWIGLUOAI:
+            pairs = gemm1_lora_delta.unflatten(-1, (intermediate_size, 2))
+        else:
+            pairs = gemm1_lora_delta.unflatten(-1, (2, intermediate_size)).mT
+        # Per-expert deltas arrive as (up, gate); fused 3D ones follow the
+        # checkpoint (gate, up) order.
+        if lora_context.w13_num_slices != 2:
+            pairs = pairs.flip(-1)
+        delta = pairs.mT.flatten(-2)  # [up | gate]
+
+        # The kernel has a single FC1 bias slot, shared by the static bias and
+        # the per-token delta, so the static bias rides along in the delta.
+        if self.w1_bias is not None:
+            bias = self.w1_bias.reshape(self.local_num_experts, -1)
+            bias = bias[:, self._w13_bias_up_gate_index].float()
+            local_ids = topk_ids.long() - self.ep_rank * self.local_num_experts
+            local_ids = local_ids.clamp_(0, self.local_num_experts - 1)
+            delta = (delta.float() + bias[local_ids]).to(torch.bfloat16)
+        return delta.contiguous()
+
+    def invoke_routed_moe(
+        self,
+        *,
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_ids_and_weights: tuple[torch.Tensor, torch.Tensor],
+        gemm1_lora_delta: torch.Tensor | None,
+        global_num_experts: int,
+        a1q_scale: torch.Tensor | None,
+        output: torch.Tensor,
+    ) -> list[torch.Tensor]:
+        from flashinfer import trtllm_fp4_block_scale_routed_moe
+
+        assert a1q_scale is None and hidden_states.dtype == torch.bfloat16
+        topk_ids, topk_weights = topk_ids_and_weights
+        # See TrtLlmBf16LoRAExperts.invoke_routed_moe for the return contract.
+        do_finalize = gemm1_lora_delta is None
+        ret = trtllm_fp4_block_scale_routed_moe(
+            topk_ids=(topk_ids.to(torch.int32), topk_weights),
+            routing_bias=None,
+            hidden_states=hidden_states,
+            hidden_states_scale=None,
+            gemm1_weights=w1,
+            gemm1_weights_scale=self.w1_scale,
+            # With a delta the bias is folded into it (prepare_gemm1_lora_delta).
+            gemm1_bias=self.w1_bias if do_finalize else None,
+            gemm1_alpha=self.gemm1_alpha,
+            gemm1_beta=self.gemm1_beta,
+            gemm1_clamp_limit=self.gemm1_clamp_limit,
+            gemm2_weights=w2,
+            gemm2_weights_scale=self.w2_scale,
+            gemm2_bias=self.w2_bias,
+            output1_scale_scalar=None,
+            output1_scale_gate_scalar=None,
+            output2_scale_scalar=None,
+            num_experts=global_num_experts,
+            top_k=self.topk,
+            n_group=None,
+            topk_group=None,
+            intermediate_size=self.intermediate_size_per_partition,
+            local_expert_offset=self.ep_rank * self.local_num_experts,
+            local_num_experts=self.local_num_experts,
+            routed_scaling_factor=None,
+            # Routing is already done; Renormalize is a safe default.
+            routing_method_type=RoutingMethodType.Renormalize,
+            do_finalize=do_finalize,
+            enable_pdl=True,
+            activation_type=self._flashinfer_activation_type(
+                self.moe_config.activation
+            ),
+            output=output if do_finalize else None,
+            tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
+            gemm1_lora_delta=gemm1_lora_delta,
+        )
+        return [output] if do_finalize else list(ret)

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

@@ -122,6 +122,9 @@ class TrtLlmMxfp4ExpertsBase:
         p = current_platform
         return p.is_cuda() and p.is_device_capability_family(100) and has_flashinfer()
 
+    def output_hidden_dim(self, hidden_dim: int) -> int:
+        return self.hidden_dim_unpadded
+
     @staticmethod
     def _supports_no_act_and_mul() -> bool:
         return False

vllm/model_executor/layers/fused_moe/modular_kernel.py

@@ -843,6 +843,11 @@ class FusedMoEExpertsModular(FusedMoEExperts):
         """Workspace type: The dtype to use for the workspace tensors."""
         return act_dtype
 
+    def output_hidden_dim(self, hidden_dim: int) -> int:
+        """Hidden size of the finalized output. Experts that drop the hidden
+        padding return the unpadded size."""
+        return hidden_dim
+
     @abstractmethod
     def workspace_shapes(
         self,
@@ -1321,7 +1326,8 @@ class FusedMoEKernelModularImpl:
         # low-latency kernels are always batched and can never run into
         # the tensor.numel() == 0 case.
         if M_full == 0:
-            return torch.empty_like(a1q, dtype=in_dtype)
+            out_dim = self.fused_experts.output_hidden_dim(a1q.size(-1))
+            return a1q.new_empty((*a1q.shape[:-1], out_dim), dtype=in_dtype)
 
         workspace13, workspace2, fused_out, fused_out_is_specialized = (
             self._allocate_buffers(
@@ -1505,7 +1511,12 @@ class FusedMoEKernelModularImpl:
             a fused consumer.
 
         """
-        output = torch.empty_like(hidden_states)
+        output = hidden_states.new_empty(
+            (
+                *hidden_states.shape[:-1],
+                self.fused_experts.output_hidden_dim(hidden_states.size(-1)),
+            )
+        )
 
         local_num_experts = w1.shape[0]
         if global_num_experts == -1:

vllm/model_executor/layers/fused_moe/oracle/mxfp4.py

@@ -535,6 +535,30 @@ def _make_log_unsupported(backend: Mxfp4MoeBackend, reason: str | None) -> str:
     return f"{base} since {reason}." if reason else f"{base}."
 
 
+def _requested_kernel_cls(
+    backend: Mxfp4MoeBackend,
+    config: FusedMoEConfig,
+) -> list[type[mk.FusedMoEExperts]]:
+    """Kernel classes for an explicitly requested backend. The LoRA-aware
+    TRT-LLM experts are opt-in: "auto" keeps its current LoRA selection."""
+    candidates = backend_to_kernel_cls(backend)
+    if not config.is_lora_enabled or backend not in TRTLLM_BACKENDS:
+        return candidates
+    if backend == Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8:
+        raise ValueError(
+            _make_log_unsupported(
+                backend,
+                "expert LoRA is not supported with MXFP8 activations; use "
+                "moe_backend='flashinfer_trtllm'",
+            )
+        )
+    from vllm.model_executor.layers.fused_moe.experts.trtllm_lora_moe import (
+        TrtLlmMxfp4LoRAExperts,
+    )
+
+    return [TrtLlmMxfp4LoRAExperts, *candidates]
+
+
 def _return_or_raise(
     backend: Mxfp4MoeBackend,
     config: FusedMoEConfig,
@@ -544,7 +568,7 @@ def _return_or_raise(
     scope: Literal["process", "global", "local"] = "local",
 ) -> tuple[Mxfp4MoeBackend, type[mk.FusedMoEExperts]]:
     reason: str | None = None
-    for k_cls in backend_to_kernel_cls(backend):
+    for k_cls in _requested_kernel_cls(backend, config):
         supported, reason = k_cls.is_supported_config(
             k_cls, config, weight_key, activation_key, activation_format
         )