vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe ¶

class TrtLlmMxfp4ExpertsBase:
    """
    MXFP4 TRTLLM-Gen MoE kernels. Shared base for modular and monolithic.
    """

    def __init__(
        self,
        moe_config: FusedMoEConfig,
        quant_config: FusedMoEQuantConfig,
        **kwargs,
    ):
        self.moe_config = moe_config
        self.quant_config = quant_config

        self.routing_method_type = moe_config.routing_method
        self.topk = moe_config.experts_per_token
        self.intermediate_size_per_partition = (
            moe_config.intermediate_size_per_partition
        )
        self.hidden_dim = moe_config.hidden_dim
        self.hidden_dim_unpadded = (
            moe_config.hidden_dim_unpadded or moe_config.hidden_dim
        )
        self.local_num_experts = moe_config.num_local_experts
        self.ep_rank = moe_config.moe_parallel_config.ep_rank

        # MXFP4-specific TRTLLM parameters from quant_config
        device = torch.accelerator.current_device_index()
        if quant_config.gemm1_alpha is not None:
            self.gemm1_alpha = torch.tensor(
                [quant_config.gemm1_alpha] * self.local_num_experts,
                dtype=torch.float32,
                device=device,
            )
        else:
            self.gemm1_alpha = None

        if quant_config.gemm1_beta is not None:
            self.gemm1_beta = torch.tensor(
                [quant_config.gemm1_beta] * self.local_num_experts,
                dtype=torch.float32,
                device=device,
            )
        else:
            self.gemm1_beta = None

        if quant_config.gemm1_clamp_limit is not None:
            self.gemm1_clamp_limit = torch.tensor(
                [quant_config.gemm1_clamp_limit] * self.local_num_experts,
                dtype=torch.float32,
                device=device,
            )
        else:
            self.gemm1_clamp_limit = None

        # SITU (SituGLU) TRTLLM-Gen kernel computes
        #   left  = alpha * tanh(x0 / alpha) * sigmoid(x0)   # gate (x0)
        #   right = beta  * tanh(x1 / beta)                  # up   (x1)
        # which matches vLLM's situ_and_mul with (beta, linear_beta), so map
        # situ beta -> gatedActAlpha (gemm1_alpha) and situ linear_beta ->
        # gatedActBeta (gemm1_beta). Both must be > 0.
        if moe_config.activation == MoEActivation.SITU:
            situ_beta = moe_config.activation_situ_beta
            situ_linear_beta = moe_config.activation_situ_linear_beta
            assert situ_beta is not None and situ_beta > 0, (
                "SITU requires activation_situ_beta > 0"
            )
            assert situ_linear_beta is not None and situ_linear_beta > 0, (
                "TRTLLM SiTuGlu requires activation_situ_linear_beta > 0 "
                "(the private cubin has no up-passthrough path)"
            )
            self.gemm1_alpha = torch.full(
                (self.local_num_experts,),
                float(situ_beta),
                dtype=torch.float32,
                device=device,
            )
            self.gemm1_beta = torch.full(
                (self.local_num_experts,),
                float(situ_linear_beta),
                dtype=torch.float32,
                device=device,
            )
            self.gemm1_clamp_limit = None

    @staticmethod
    def _supports_current_device() -> bool:
        p = current_platform
        return p.is_cuda() and p.is_device_capability_family(100) and has_flashinfer()

    @staticmethod
    def _supports_no_act_and_mul() -> bool:
        return False

    @staticmethod
    def _supports_quant_scheme(
        weight_key: QuantKey | None,
        activation_key: QuantKey | None,
    ) -> bool:
        SUPPORTED_W_A = [
            (kMxfp4Static, None),
            (kMxfp4Static, kMxfp8Dynamic),
        ]
        return (weight_key, activation_key) in SUPPORTED_W_A

    @staticmethod
    def _supports_activation(activation: MoEActivation) -> bool:
        if activation == MoEActivation.SITU:
            return has_flashinfer_situ_activation()
        return activation in (
            MoEActivation.SWIGLUOAI,
            MoEActivation.SILU,
        )

    @staticmethod
    def _flashinfer_activation_type(activation: MoEActivation) -> int:
        return activation_to_flashinfer_int(activation)

    @staticmethod
    def activation_format() -> mk.FusedMoEActivationFormat:
        return mk.FusedMoEActivationFormat.Standard

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

class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModular):
    """
    Modular version of the MXFP4 TRTLLM kernel (just the experts).
    Wraps flashinfer.trtllm_fp4_block_scale_routed_moe().
    Moved from trtllm_moe.py.
    """

    @staticmethod
    def _supports_parallel_config(
        moe_parallel_config: FusedMoEParallelConfig,
    ) -> bool:
        return True

    @staticmethod
    def _supports_routing_method(
        routing_method: RoutingMethodType,
        weight_key: QuantKey | None,
        activation_key: QuantKey | None,
    ) -> bool:
        # Modular kernel handles only the expert computation;
        # routing is done externally, so accept any routing method.
        return True

    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
        return TopKWeightAndReduceNoOP()

    def workspace_shapes(
        self,
        M: int,
        N: int,
        K: int,
        topk: int,
        global_num_experts: int,
        local_num_experts: int,
        expert_tokens_meta: mk.ExpertTokensMetadata | None,
        activation: MoEActivation,
    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
        # The workspaces for this implementation are managed by flashinfer.
        workspace1 = (0,)
        workspace2 = (0,)
        output = (M, self.hidden_dim_unpadded)
        return (workspace1, workspace2, output)

    def _max_supported_tokens(self, top_k: int, global_num_experts: int) -> int:
        """Max tokens per kernel call before the batched-GEMM grid overflows.

        The TRTLLM-Gen batched GEMM launches a static grid whose batch (Y)
        dimension is ``getMaxNumCtasInBatchDim(num_tokens, top_k, num_experts,
        tileTokensDim)`` and must stay <= 65535. Solving that for num_tokens
        with the smallest tile the kernel may pick (tileTokensDim=8, the runner
        default) gives a bound that is safe regardless of the tactic selected.
        Without it, large batches (e.g. Kimi-K3 top_k=16, EP16 profiling with
        131072 gathered tokens) overflow the grid and the GEMM launch fails.
        """
        MAX_GRID_Y = 65535
        MIN_TILE_TOKENS_DIM = 8
        max_tokens = (MAX_GRID_Y - global_num_experts) * MIN_TILE_TOKENS_DIM // top_k
        return max(1, min(300000, max_tokens))

    def _invoke_kernel(
        self,
        output: torch.Tensor,
        x_quant: torch.Tensor,
        x_scale: torch.Tensor | None,
        topk_ids: torch.Tensor,
        topk_weights: torch.Tensor,
        w1: torch.Tensor,
        w2: torch.Tensor,
        activation: MoEActivation,
        global_num_experts: int,
        local_num_experts: int,
        local_expert_offset: int,
        topk: int,
    ) -> None:
        from flashinfer import trtllm_fp4_block_scale_routed_moe

        packed_tensor = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)
        trtllm_fp4_block_scale_routed_moe(
            topk_ids=packed_tensor,
            routing_bias=None,
            hidden_states=x_quant,
            hidden_states_scale=x_scale,
            gemm1_weights=w1,
            gemm1_weights_scale=self.w1_scale,
            gemm1_bias=self.w1_bias,
            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=topk,
            n_group=None,
            topk_group=None,
            intermediate_size=self.intermediate_size_per_partition,
            local_expert_offset=local_expert_offset,
            local_num_experts=local_num_experts,
            routed_scaling_factor=None,
            # Modular kernel receives pre-routed tokens, so routing is already
            # done. Use Renormalize as a safe default the TRTLLM kernel supports.
            routing_method_type=RoutingMethodType.Renormalize,
            do_finalize=True,
            enable_pdl=True,
            activation_type=self._flashinfer_activation_type(activation),
            output=output,
            tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
        )

    def apply(
        self,
        output: torch.Tensor,
        hidden_states: torch.Tensor,
        w1: torch.Tensor,
        w2: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        activation: MoEActivation,
        global_num_experts: int,
        expert_map: torch.Tensor | None,
        a1q_scale: torch.Tensor | None,
        a2_scale: torch.Tensor | None,
        workspace13: torch.Tensor,
        workspace2: torch.Tensor,
        expert_tokens_meta: mk.ExpertTokensMetadata | None,
        apply_router_weight_on_input: bool,
    ):
        topk = topk_ids.size(-1)
        local_num_experts = w1.size(0)
        local_expert_offset = self.moe_config.ep_rank * local_num_experts

        if a1q_scale is not None:
            x_quant = hidden_states
            x_scale = a1q_scale.view(torch.float8_e4m3fn)
        else:
            assert hidden_states.dtype == torch.bfloat16
            x_quant = hidden_states
            x_scale = None

        assert self.w1_scale is not None
        assert self.w2_scale is not None

        # Chunk tokens so the batched-GEMM grid stays within CUDA limits.
        M = x_quant.size(0)
        chunk_size = self._max_supported_tokens(topk, global_num_experts)
        for start in range(0, M, chunk_size):
            end = min(start + chunk_size, M)
            self._invoke_kernel(
                output[start:end],
                x_quant[start:end],
                None if x_scale is None else x_scale[start:end],
                topk_ids[start:end],
                topk_weights[start:end],
                w1,
                w2,
                activation,
                global_num_experts,
                local_num_experts,
                local_expert_offset,
                topk,
            )

        return output

class TrtLlmMxfp4ExpertsMonolithic(
    TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsMonolithic
):
    """
    Monolithic version of the MXFP4 TRTLLM kernel (router + experts).
    Wraps flashinfer.trtllm_fp4_block_scale_moe().
    """

    def supports_routing_replay_capture(self) -> bool:
        return True

    @staticmethod
    def _supports_parallel_config(
        moe_parallel_config: FusedMoEParallelConfig,
    ) -> bool:
        return (
            not moe_parallel_config.use_all2all_kernels
            and not moe_parallel_config.enable_eplb
            and moe_parallel_config.dp_size <= 1
        )

    @staticmethod
    def _supports_routing_method(
        routing_method: RoutingMethodType,
        weight_key: QuantKey | None,
        activation_key: QuantKey | None,
    ) -> bool:
        return routing_method in [
            RoutingMethodType.DeepSeekV3,
            RoutingMethodType.Renormalize,
            RoutingMethodType.RenormalizeNaive,
        ]

    @staticmethod
    def _supports_router_logits_dtype(
        router_logits_dtype: torch.dtype | None,
        routing_method: RoutingMethodType,
    ) -> bool:
        # Kernel converts to bfloat16 internally
        return True

    def apply(
        self,
        hidden_states: torch.Tensor,
        w1: torch.Tensor,
        w2: torch.Tensor,
        router_logits: torch.Tensor,
        activation: MoEActivation,
        global_num_experts: int,
        expert_map: torch.Tensor | None,
        a1q_scale: torch.Tensor | None,
        apply_router_weight_on_input: bool,
        # grouped topk + fused topk bias parameters
        num_expert_group: int | None = None,
        e_score_correction_bias: torch.Tensor | None = None,
        routed_scaling_factor: float | None = None,
        topk_group: int | None = None,
    ) -> torch.Tensor | UnfinalizedMoEOutput:
        from flashinfer import trtllm_fp4_block_scale_moe

        if a1q_scale is not None:
            x_quant = hidden_states
            x_scale = a1q_scale.view(torch.float8_e4m3fn)
        else:
            assert hidden_states.dtype == torch.bfloat16
            x_quant = hidden_states
            x_scale = None
        num_tokens = hidden_states.shape[0]
        defer = self.moe_config.should_defer_moe_finalize(num_tokens)
        finalized_output = None
        if not defer:
            finalized_output = torch.empty(
                *hidden_states.shape[:-1],
                self.hidden_dim_unpadded,
                dtype=torch.bfloat16,
                device=hidden_states.device,
            )

        routing_replay_out = self._maybe_make_routing_replay_buffer(
            num_tokens=num_tokens,
            device=hidden_states.device,
        )
        flashinfer_output = trtllm_fp4_block_scale_moe(
            routing_logits=router_logits,
            routing_bias=e_score_correction_bias,
            hidden_states=x_quant,
            hidden_states_scale=x_scale,
            gemm1_weights=w1,
            gemm1_weights_scale=self.w1_scale,
            gemm1_bias=self.w1_bias,
            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=(num_expert_group or 0),
            topk_group=(topk_group or 0),
            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=routed_scaling_factor,
            routing_method_type=self.routing_method_type,
            do_finalize=not defer,
            activation_type=self._flashinfer_activation_type(activation),
            tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
            output=finalized_output,
            routing_replay_out=routing_replay_out,
        )
        routed_output = convert_flashinfer_moe_output(
            flashinfer_output,
            do_finalize=not defer,
            num_tokens=num_tokens,
            top_k=self.topk,
            finalized_output=finalized_output,
        )
        self._maybe_dispatch_routing_replay(routing_replay_out, num_tokens=num_tokens)
        return routed_output