vllm.model_executor.layers.fused_moe.flashinfer_moe_ep

Classes:

• FlashInferMoeEp –

Functions:

• flashinfer_moe_ep_unsupported_reasons –

  Deployment constraints of the megakernel beyond the framework's checks.

FlashInferMoeEp

Methods:

• expert_weight_views –

  Every per-expert tensor the kernel reads: weights plus epilogue.

• kernel_weights –

  Kernel-resident fc1 weight, fc1 scales, fc2 weight, fc2 scales.

Source code in vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py

class FlashInferMoeEp:
    def __init__(
        self,
        moe: FusedMoEConfig,
        weights: FlashInferMoeEpWeights,
        epilogue: FlashInferMoeEpEpilogue | None = None,
        *,
        apply_topk_in_fc1: bool,
    ) -> None:
        spec = flashinfer_moe_ep_backend_spec(moe.moe_backend)
        if spec.kernel == "deep_gemm":
            _expose_deep_gemm_to_flashinfer()
        api = _load_flashinfer_moe_ep_api()
        # MegaMoE's workspace mode is fixed, including empty collective calls.
        self._defer_finalize = moe.use_deferred_moe_finalize
        if self._defer_finalize and (
            spec.kernel != "cutedsl"
            or apply_topk_in_fc1
            or not supports_flashinfer_moe_ep_deferred_finalize()
        ):
            raise ValueError(
                "Deferred finalization requires FlashInfer CuTeDSL "
                "forward_unfinalized with output-side routing weights"
            )
        if epilogue is None:
            epilogue = FlashInferMoeEpEpilogue()
        ep_group = get_ep_group()
        bootstrap = api.BootstrapConfig(
            world_size=ep_group.world_size,
            rank=ep_group.rank_in_group,
            process_group=ep_group.device_group,
            device=torch.accelerator.current_device_index(),
        )
        fleet_params = api.FleetParams(
            num_experts=moe.num_experts,
            max_tokens_per_rank=moe.max_num_tokens,
            token_hidden_size=moe.hidden_dim,
        )
        weight_pack = api.MoEWeightPack(
            w13=weights.w13,
            w2=weights.w2,
            w13_scale=weights.w13_scale,
            w2_scale=weights.w2_scale,
        )
        if spec.kernel == "cutedsl":
            # In-kernel fc2 reduce keeps its default (off); FlashInfer 0.7.1rc1
            # renamed in_kernel_fc2_reduce to enable_in_kernel_fc2_reduce.
            megakernel = api.Nvfp4CutedslMegaMoeConfig(
                intermediate_size=moe.intermediate_size,
                top_k=moe.experts_per_token,
                gate_up_clamp=moe.swiglu_limit,
                fast_math=True,
                apply_topk_in_fc1=apply_topk_in_fc1,
                combine_dtype="bf16",
                input_norm_const=epilogue.input_norm_const,
                fc1_alpha=None,
                fc2_alpha=None,
                fc1_norm_const=None,
                knobs=None,
                **({"do_finalize": False} if self._defer_finalize else {}),
            )
        else:
            megakernel = api.DeepGemmMegaMoeConfig(
                intermediate_size=moe.intermediate_size,
                top_k=moe.experts_per_token,
                activation_clamp=moe.swiglu_limit,
                fast_math=True,
            )
        backend = api.MegaConfig(
            megakernel=megakernel,
            quantize_input=True,
            preprocess_weights=True,
        )

        self._moe_ep_tensors_cls = api.MoEEpTensors
        self._mega_layer: Any | None = api.MoEEpMegaLayer(
            bootstrap,
            fleet_params,
            weight_pack,
            backend,
        )
        self._device = weights.w13.device
        self._hidden_size = moe.hidden_dim
        self._max_num_tokens = moe.max_num_tokens
        self._num_experts = moe.num_experts
        self._top_k = moe.experts_per_token
        self._epilogue: FlashInferMoeEpEpilogue | None = epilogue
        self._return_workspace_view = moe.routing_method is RoutingMethodType.DeepseekV4
        self._sp_size = moe.sp_size
        self._decode_max_num_tokens = 0
        self._decode_workspace: Any | None = None
        self._is_separate_workspace_supported = (
            spec.kernel == "cutedsl"
            and moe.moe_parallel_config.pcp_size == 1
            and hasattr(self._mega_layer, "create_workspace")
        )
        if self._is_separate_workspace_supported:
            vllm_config = get_current_vllm_config()
            decode_tokens = cdiv(
                vllm_config.scheduler_config.max_num_seqs
                * (1 + vllm_config.num_speculative_tokens),
                self._sp_size,
            )
            if decode_tokens < self._max_num_tokens:
                self._decode_max_num_tokens = decode_tokens

    @property
    def can_overlap_shared_experts(self) -> bool:
        return False

    @property
    def output_is_reduced(self) -> bool:
        return True

    @property
    def topk_indices_dtype(self) -> torch.dtype:
        return torch.int32

    @property
    def is_monolithic(self) -> bool:
        return False

    def _select_workspace(self, num_tokens: int) -> Any | None:
        if not is_forward_context_available():
            return None
        metadata = get_forward_context().dp_metadata
        if metadata is None:
            return None
        # Every EP rank must use the same capacity: peer buffer offsets depend
        # on it. Local token counts can differ, including on idle ranks.
        if metadata.local_sizes is not None:
            max_tokens = max(metadata.local_sizes)
        else:
            max_tokens = cdiv(
                int(metadata.num_tokens_across_dp_cpu.max().item()), self._sp_size
            )
        if (
            self._decode_workspace is not None
            and max_tokens <= self._decode_max_num_tokens
        ):
            assert num_tokens <= self._decode_max_num_tokens
            return self._decode_workspace
        return None

    def _tensors(
        self,
        hidden_states: torch.Tensor,
        topk_ids: torch.Tensor,
        topk_weights: torch.Tensor,
    ) -> Any:
        epilogue = cast(FlashInferMoeEpEpilogue, self._epilogue)
        return self._moe_ep_tensors_cls(
            hidden_states=hidden_states,
            topk_ids=topk_ids,
            topk_weights=topk_weights,
            fc1_alpha=epilogue.fc1_alpha,
            fc2_alpha=epilogue.fc2_alpha,
            fc1_norm_const=epilogue.fc1_norm_const,
        )

    def __call__(
        self,
        hidden_states: torch.Tensor,
        topk_ids: torch.Tensor,
        topk_weights: torch.Tensor,
    ) -> torch.Tensor | UnfinalizedMoEOutput:
        if self._mega_layer is None:
            raise RuntimeError("FlashInfer MoE-EP layer was destroyed")
        if hidden_states.shape[0] > self._max_num_tokens:
            raise ValueError(
                f"FlashInfer MoE-EP got {hidden_states.shape[0]} tokens, but "
                f"the workspace supports at most {self._max_num_tokens}"
            )
        tensors = self._tensors(hidden_states, topk_ids, topk_weights)
        workspace = self._select_workspace(hidden_states.shape[0])
        workspace_kwargs = {"workspace": workspace} if workspace is not None else {}
        if self._defer_finalize:
            unfinalized = convert_flashinfer_moe_output(
                list(self._mega_layer.forward_unfinalized(tensors, **workspace_kwargs)),
                do_finalize=False,
                num_tokens=hidden_states.shape[0],
                top_k=topk_ids.shape[1],
            )
            assert isinstance(unfinalized, UnfinalizedMoEOutput)
            return unfinalized
        if self._return_workspace_view and getattr(
            self._mega_layer, "supports_output_view", False
        ):
            return self._mega_layer.forward(
                tensors, return_workspace_view=True, **workspace_kwargs
            )
        return self._mega_layer(tensors, **workspace_kwargs)

    @torch.inference_mode()
    def warmup(self) -> None:
        if self._mega_layer is None:
            return
        hidden_states = torch.zeros(
            1,
            self._hidden_size,
            dtype=torch.bfloat16,
            device=self._device,
        )
        topk_ids = (
            torch.arange(self._top_k, dtype=torch.int32, device=self._device)
            .mul_(self._num_experts // self._top_k)
            .view(1, self._top_k)
        )
        topk_weights = torch.full(
            (1, self._top_k),
            1.0 / self._top_k,
            dtype=torch.float32,
            device=self._device,
        )
        tensors = self._tensors(hidden_states, topk_ids, topk_weights)
        self._mega_layer.warmup(tensors)
        if self._is_separate_workspace_supported and self._decode_max_num_tokens:
            if self._decode_workspace is None:
                self._decode_workspace = self._mega_layer.create_workspace(
                    self._decode_max_num_tokens
                )
            self._mega_layer.warmup(tensors, workspace=self._decode_workspace)

    def kernel_weights(
        self,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """Kernel-resident fc1 weight, fc1 scales, fc2 weight, fc2 scales.

        Contiguous per-expert byte tensors that alias the memory the megakernel
        reads, so a caller may permute experts in place (EPLB).
        """
        if self._mega_layer is None:
            raise RuntimeError("FlashInfer MoE-EP layer was destroyed")
        (fc1_weight, fc1_sf), (fc2_weight, fc2_sf) = self._mega_layer._transformed
        return (
            _expert_major_bytes(fc1_weight),
            _expert_major_bytes(fc1_sf),
            _expert_major_bytes(fc2_weight),
            _expert_major_bytes(fc2_sf),
        )

    def expert_weight_views(self) -> tuple[torch.Tensor, ...]:
        """Every per-expert tensor the kernel reads: weights plus epilogue."""
        epilogue = cast(FlashInferMoeEpEpilogue, self._epilogue)
        vectors = (epilogue.fc1_alpha, epilogue.fc2_alpha, epilogue.fc1_norm_const)
        return (*self.kernel_weights(), *(v for v in vectors if v is not None))

    def destroy(self) -> None:
        mega_layer = self._mega_layer
        if mega_layer is None:
            return
        mega_layer.destroy()
        self._mega_layer = None
        self._decode_workspace = None
        self._epilogue = None

    def __del__(self) -> None:
        with contextlib.suppress(Exception):
            self.destroy()

expert_weight_views()

Every per-expert tensor the kernel reads: weights plus epilogue.

Source code in vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py

def expert_weight_views(self) -> tuple[torch.Tensor, ...]:
    """Every per-expert tensor the kernel reads: weights plus epilogue."""
    epilogue = cast(FlashInferMoeEpEpilogue, self._epilogue)
    vectors = (epilogue.fc1_alpha, epilogue.fc2_alpha, epilogue.fc1_norm_const)
    return (*self.kernel_weights(), *(v for v in vectors if v is not None))

kernel_weights()

Kernel-resident fc1 weight, fc1 scales, fc2 weight, fc2 scales.

Contiguous per-expert byte tensors that alias the memory the megakernel reads, so a caller may permute experts in place (EPLB).

Source code in vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py

def kernel_weights(
    self,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Kernel-resident fc1 weight, fc1 scales, fc2 weight, fc2 scales.

    Contiguous per-expert byte tensors that alias the memory the megakernel
    reads, so a caller may permute experts in place (EPLB).
    """
    if self._mega_layer is None:
        raise RuntimeError("FlashInfer MoE-EP layer was destroyed")
    (fc1_weight, fc1_sf), (fc2_weight, fc2_sf) = self._mega_layer._transformed
    return (
        _expert_major_bytes(fc1_weight),
        _expert_major_bytes(fc1_sf),
        _expert_major_bytes(fc2_weight),
        _expert_major_bytes(fc2_sf),
    )

_expert_major_bytes(tensor)

Contiguous storage of a per-expert kernel tensor, 1-byte dtypes as uint8.

The CuTeDSL kernel keeps (E, K, N) transpose views over (E, N, K) storage.

Source code in vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py

def _expert_major_bytes(tensor: torch.Tensor) -> torch.Tensor:
    """Contiguous storage of a per-expert kernel tensor, 1-byte dtypes as uint8.

    The CuTeDSL kernel keeps (E, K, N) transpose views over (E, N, K) storage.
    """
    if not tensor.is_contiguous():
        tensor = tensor.transpose(1, 2)
    if not tensor.is_contiguous():
        raise ValueError(f"expert tensor is not expert-major: {tensor.stride()}")
    return tensor.view(torch.uint8) if tensor.element_size() == 1 else tensor

flashinfer_moe_ep_unsupported_reasons(moe, weight_key, activation_key)

Deployment constraints of the megakernel beyond the framework's checks.

activation_key is None for MXFP4 checkpoints (the kernel quantizes bf16 activations itself) but marks a W4A16 recipe for NVFP4 weights.

Source code in vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py

def flashinfer_moe_ep_unsupported_reasons(
    moe: FusedMoEConfig,
    weight_key: QuantKey | None,
    activation_key: QuantKey | None,
) -> tuple[str, ...]:
    """Deployment constraints of the megakernel beyond the framework's checks.

    ``activation_key`` is ``None`` for MXFP4 checkpoints (the kernel quantizes
    bf16 activations itself) but marks a W4A16 recipe for NVFP4 weights.
    """
    spec = flashinfer_moe_ep_backend_spec(moe.moe_backend)
    unsupported: list[str] = []
    weight_format = _WEIGHT_FORMATS.get(weight_key)
    if weight_format is None:
        unsupported.append(f"weight scheme {weight_key}")
    elif weight_format not in spec.weight_formats:
        unsupported.append(f"{weight_format.upper()} weights")
    if weight_key == kNvfp4Static and activation_key is None:
        unsupported.append("A16 activations with NVFP4 weights")
    if moe.skip_final_all_reduce:
        unsupported.append("skip_final_all_reduce")
    if moe.in_dtype != torch.bfloat16:
        unsupported.append(f"activation dtype {moe.in_dtype}")
    if moe.has_bias:
        unsupported.append("expert bias")
    if moe.swiglu_alpha is not None or moe.swiglu_beta is not None:
        unsupported.append("custom SwiGLU alpha or beta")
    if (
        spec.kernel == "deep_gemm"
        and moe.routing_method is not RoutingMethodType.DeepseekV4
    ):
        unsupported.append(f"routing method {moe.routing_method.name}")

    vllm_config = get_current_vllm_config()
    if vllm_config.weight_transfer_config is not None:
        unsupported.append("runtime weight transfer")
    if vllm_config.parallel_config.enable_dbo:
        unsupported.append("dual batch overlap")
    if spec.kernel == "deep_gemm" and vllm_config.parallel_config.enable_eplb:
        unsupported.append("EPLB")
    return tuple(unsupported)