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

tests/lora/test_trtllm_mxfp4_lora_moe.py

@@ -0,0 +1,218 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for the LoRA-aware FlashInfer TRTLLM MXFP4 experts
+(`TrtLlmMxfp4LoRAExperts`) against a float64 reference."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tests.kernels.moe.utils import make_dummy_moe_config
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.config import mxfp4_w4a16_moe_quant_config
+from vllm.model_executor.layers.fused_moe.experts.trtllm_lora_moe import (
+    TrtLlmMxfp4LoRAExperts,
+)
+from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
+    Mxfp4MoeBackend,
+    _requested_kernel_cls,
+    backend_to_kernel_cls,
+    convert_weight_to_mxfp4_moe_kernel_format,
+)
+from vllm.platforms import current_platform
+from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
+from vllm.utils.torch_utils import set_random_seed
+
+# FlashInfer's TRTLLM MoE LoRA kernels only run on SM100-family GPUs.
+pytestmark = [
+    pytest.mark.skipif(
+        not current_platform.is_device_capability_family(100)
+        or not has_flashinfer_trtllm_fused_moe(),
+        reason="Requires an SM100-family GPU and FlashInfer TRTLLM fused MoE",
+    ),
+    pytest.mark.skip_global_cleanup,
+]
+
+E, N, K, T, TOPK = 16, 256, 512, 33, 4
+E2M1 = [0, 0.5, 1, 1.5, 2, 3, 4, 6]
+
+
+def _dequant_mxfp4(weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    """(E, R, K/2) packed nibbles + (E, R, K/32) e8m0 scales -> (E, R, K) f64."""
+    nibbles = torch.stack([weight & 0xF, weight >> 4], dim=-1).flatten(-2).long()
+    magnitude = torch.tensor(E2M1, dtype=torch.float64, device=weight.device)
+    values = magnitude[nibbles & 7] * torch.where(nibbles >> 3 == 1, -1.0, 1.0)
+    block_scale = torch.pow(2.0, scale.double() - 127.0)
+    return values * block_scale.repeat_interleave(32, dim=-1)
+
+
+def _make_experts(has_bias: bool, activation=MoEActivation.SILU):
+    w13 = torch.randint(0, 256, (E, 2 * N, K // 2), dtype=torch.uint8, device="cuda")
+    w2 = torch.randint(0, 256, (E, K, N // 2), dtype=torch.uint8, device="cuda")
+    s13 = torch.randint(118, 122, (E, 2 * N, K // 32), dtype=torch.uint8, device="cuda")
+    s2 = torch.randint(118, 122, (E, K, N // 32), dtype=torch.uint8, device="cuda")
+    b13 = torch.randn(E, 2 * N, device="cuda") / 4 if has_bias else None
+    b2 = torch.randn(E, K, device="cuda") / 4 if has_bias else None
+    k13, k2, ks13, ks2, kb13, kb2 = convert_weight_to_mxfp4_moe_kernel_format(
+        Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
+        None,
+        w13.clone(),
+        w2.clone(),
+        s13.clone(),
+        s2.clone(),
+        None if b13 is None else b13.clone(),
+        None if b2 is None else b2.clone(),
+        _cache_permute_indices={},
+    )
+    config = make_dummy_moe_config(
+        num_experts=E,
+        experts_per_token=TOPK,
+        hidden_dim=K,
+        intermediate_size=N,
+        activation=activation,
+    )
+    quant_config = mxfp4_w4a16_moe_quant_config(ks13, ks2, kb13, kb2)
+    experts = TrtLlmMxfp4LoRAExperts(config, quant_config)
+    reference = (_dequant_mxfp4(w13, s13), _dequant_mxfp4(w2, s2), b13, b2)
+    return experts, (k13, k2), reference
+
+
+def _make_inputs():
+    x = torch.randn(T, K, device="cuda").to(torch.bfloat16)
+    ids = torch.stack([torch.randperm(E, device="cuda")[:TOPK] for _ in range(T)])
+    weights = torch.rand(T, TOPK, device="cuda", dtype=torch.float32) + 0.1
+    return x, ids.to(torch.int32), weights
+
+
+def _reference(reference, x, ids, weights, gate_delta=None, up_delta=None):
+    """Float64 MoE; the checkpoint w13 layout is [gate | up]."""
+    w13, w2, b13, b2 = reference
+    ids = ids.long()
+    h = torch.einsum("tk,tjnk->tjn", x.double(), w13[ids])
+    if b13 is not None:
+        h = h + b13.double()[ids]
+    gate, up = h[..., :N], h[..., N:]
+    if gate_delta is not None:
+        gate, up = gate + gate_delta.double(), up + up_delta.double()
+    act = torch.nn.functional.silu(gate) * up
+    out = torch.einsum("tjn,tjkn->tjk", act, w2[ids])
+    if b2 is not None:
+        out = out + b2.double()[ids]
+    return act, (out * weights.double()[..., None]).sum(1)
+
+
+def _run_lora_path(experts, kernel_weights, x, ids, weights, delta):
+    """`delta` is (T, TOPK, 2I) in the [up | gate] order of apply_w13_lora."""
+    lora_context = SimpleNamespace(w13_num_slices=2)
+    gemm2, expert_weights, idx_map, act_permuted = experts.invoke_routed_moe(
+        hidden_states=x,
+        w1=kernel_weights[0],
+        w2=kernel_weights[1],
+        topk_ids_and_weights=(ids, weights),
+        gemm1_lora_delta=experts.prepare_gemm1_lora_delta(delta, lora_context, ids),
+        global_num_experts=E,
+        a1q_scale=None,
+        output=torch.empty(T, K, dtype=torch.bfloat16, device="cuda"),
+    )
+    assert act_permuted.dtype == torch.bfloat16
+    act = experts._unpermute_activation(act_permuted, idx_map, T, TOPK, N)
+    output = torch.empty(T, K, dtype=torch.bfloat16, device="cuda")
+    no_w2_delta = torch.zeros(T, TOPK, K, dtype=torch.bfloat16, device="cuda")
+    experts._finalize_with_w2_lora(
+        output, gemm2, expert_weights, idx_map, no_w2_delta, T, TOPK
+    )
+    return act.view(T, TOPK, N), output
+
+
+def _rel_err(actual: torch.Tensor, expected: torch.Tensor) -> float:
+    actual, expected = actual.double(), expected.double()
+    return float((actual - expected).norm() / expected.norm())
+
+
+@pytest.mark.parametrize("has_bias", [False, True])
+@torch.inference_mode()
+def test_gemm1_lora_delta_matches_reference(has_bias: bool):
+    set_random_seed(0)
+    experts, kernel_weights, reference = _make_experts(has_bias)
+    x, ids, weights = _make_inputs()
+    gate_delta = (torch.randn(T, TOPK, N, device="cuda") / 2).to(torch.bfloat16)
+    up_delta = (torch.randn(T, TOPK, N, device="cuda") / 2).to(torch.bfloat16)
+    ref_act, ref_out = _reference(reference, x, ids, weights, gate_delta, up_delta)
+    _, ref_base_out = _reference(reference, x, ids, weights)
+    # The delta has to matter for the layout check below to mean anything.
+    assert _rel_err(ref_out, ref_base_out) > 0.3
+
+    delta = torch.cat([up_delta, gate_delta], dim=-1)
+    act, out = _run_lora_path(experts, kernel_weights, x, ids, weights, delta)
+    assert _rel_err(act, ref_act) < 0.02
+    assert _rel_err(out, ref_out) < 0.02
+
+    swapped = torch.cat([gate_delta, up_delta], dim=-1)
+    _, out = _run_lora_path(experts, kernel_weights, x, ids, weights, swapped)
+    assert _rel_err(out, ref_out) > 0.3
+
+
+@pytest.mark.parametrize("has_bias", [False, True])
+@torch.inference_mode()
+def test_zero_delta_matches_base_path(has_bias: bool):
+    set_random_seed(0)
+    experts, kernel_weights, reference = _make_experts(has_bias)
+    x, ids, weights = _make_inputs()
+    base = torch.empty(T, K, dtype=torch.bfloat16, device="cuda")
+    experts.invoke_routed_moe(
+        hidden_states=x,
+        w1=kernel_weights[0],
+        w2=kernel_weights[1],
+        topk_ids_and_weights=(ids, weights),
+        gemm1_lora_delta=None,
+        global_num_experts=E,
+        a1q_scale=None,
+        output=base,
+    )
+    _, ref_out = _reference(reference, x, ids, weights)
+    assert _rel_err(base, ref_out) < 0.02
+
+    zero = torch.zeros(T, TOPK, 2 * N, dtype=torch.bfloat16, device="cuda")
+    _, out = _run_lora_path(experts, kernel_weights, x, ids, weights, zero)
+    assert _rel_err(out, base) < 0.02
+
+
+@pytest.mark.parametrize("w13_num_slices", [1, 2])
+@pytest.mark.parametrize("activation", [MoEActivation.SILU, MoEActivation.SWIGLUOAI])
+@torch.inference_mode()
+def test_prepare_gemm1_lora_delta_layout(activation, w13_num_slices: int):
+    set_random_seed(0)
+    experts, _, _ = _make_experts(False, activation)
+    ids = _make_inputs()[1]
+    gate = torch.randn(T, TOPK, N, device="cuda").to(torch.bfloat16)
+    up = torch.randn(T, TOPK, N, device="cuda").to(torch.bfloat16)
+    if w13_num_slices == 2:
+        # What apply_w13_lora(swap_w13_slices=True) writes for per-expert LoRA.
+        pair, halves = (up, gate), (up, gate)
+    else:
+        # A fused (3D) adapter follows the checkpoint w13 row order.
+        pair, halves = (gate, up), (gate, up)
+    if activation == MoEActivation.SWIGLUOAI:
+        delta = torch.stack(pair, dim=-1).flatten(-2)
+    else:
+        delta = torch.cat(halves, dim=-1)
+
+    lora_context = SimpleNamespace(w13_num_slices=w13_num_slices)
+    prepared = experts.prepare_gemm1_lora_delta(delta, lora_context, ids)
+    assert torch.equal(prepared, torch.cat([up, gate], dim=-1))
+
+
+def test_lora_experts_are_opt_in_for_bf16_activations():
+    config = make_dummy_moe_config(num_experts=E, hidden_dim=K, intermediate_size=N)
+    bf16 = Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16
+    mxfp8 = Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8
+    assert _requested_kernel_cls(bf16, config) == backend_to_kernel_cls(bf16)
+    assert _requested_kernel_cls(mxfp8, config) == backend_to_kernel_cls(mxfp8)
+
+    config.is_lora_enabled = True
+    assert _requested_kernel_cls(bf16, config)[0] is TrtLlmMxfp4LoRAExperts
+    with pytest.raises(ValueError, match="MXFP8 activations"):
+        _requested_kernel_cls(mxfp8, config)
+    marlin = Mxfp4MoeBackend.MARLIN
+    assert _requested_kernel_cls(marlin, config) == backend_to_kernel_cls(marlin)

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
         )