diff --git a/tests/lora/test_trtllm_mxfp4_lora_moe.py b/tests/lora/test_trtllm_mxfp4_lora_moe.py new file mode 100644 index 0000000000..2165fcf95d --- /dev/null +++ b/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) diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py index 1c5698932a..cdbdd7d50f 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py +++ b/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) diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py index 6b62e8ba9c..2adaadc982 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py +++ b/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 diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index fccbd1a9d8..7f9cb06167 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/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: diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index ae79298eea..8052bb6cdb 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/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 )