vllm.model_executor.layers.fused_moe.experts.fused_humming_moe

Fused MoE utilities for Humming.

Classes:

• BatchedHummingGroupedExperts –
• HummingExpertsBase –
• HummingGroupedExperts –
• HummingIndexedExperts –

BatchedHummingGroupedExperts

Bases: HummingExpertsBase

Methods:

• apply –

  Standard apply implementation for Humming batched grouped experts.

Source code in vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py

class BatchedHummingGroupedExperts(HummingExpertsBase):
    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
        return TopKWeightAndReduceDelegate()

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

    @staticmethod
    def humming_gemm_type() -> "HummingGemmType":
        from vllm.utils.humming import GemmType as HummingGemmType

        return HummingGemmType.GROUPED_MASKED

    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,
    ) -> None:
        """Standard apply implementation for Humming batched grouped experts.

        Note: Humming kernels handle weights internally through the layer
        object, so w1, w2, a2_scale are unused. a1q_scale is None on the usual
        path (Humming quantizes the w13 input itself); for block-FP8 activations
        it carries the scale computed before dispatch, which is forwarded to
        process_input so Humming skips the redundant w13 quantization. The
        output is written into workspace13 via the buffer management.
        """
        assert not apply_router_weight_on_input
        assert expert_tokens_meta is not None

        hidden_states = hidden_states.view(-1, hidden_states.size(-1))
        # Keep the (batched) block-FP8 scale row-aligned with the flattened
        # [num_experts * max_tokens, K] hidden states above.
        if a1q_scale is not None and a1q_scale.dim() == 3:
            a1q_scale = a1q_scale.view(-1, a1q_scale.size(-1))
        valid_shape_m = self.estimate_local_valid_shape_m(topk_ids)
        expert_num_tokens = expert_tokens_meta.expert_num_tokens

        buffers = self.prepare_buffers(
            workspace13,
            workspace2,
            topk_ids.size(0),
            topk_ids.size(1),
            activation,
        )
        buffers["down_output"] = output

        inputs, input_scale, input_scale_2 = self.process_input(
            "w13",
            inputs=hidden_states,
            input_scale=a1q_scale,
            quanted_input=buffers.get("quanted_gate_up_input", None),
            expert_tokens=expert_num_tokens,
        )

        self.humming_forward(
            "w13",
            inputs=inputs,
            weight=w1,
            input_scale=input_scale,
            input_scale_2=input_scale_2,
            outputs=buffers["gate_up_output"],
            valid_shape_m=valid_shape_m,
            expert_layout=expert_num_tokens,
            compute_config=self.compute_config_str,
            tuning_config=self.w13_tuning_config_str,
        )

        inputs, input_scale, input_scale_2 = self.process_input(
            "w2",
            inputs=buffers["gate_up_output"],
            quanted_input=buffers["quanted_down_input"],
            expert_tokens=expert_num_tokens,
            activation=activation,
        )

        self.humming_forward(
            "w2",
            inputs=inputs,
            weight=w2,
            input_scale=input_scale,
            input_scale_2=input_scale_2,
            outputs=output.view(-1, hidden_states.size(-1)),
            valid_shape_m=valid_shape_m,
            expert_layout=expert_num_tokens,
            compute_config=self.compute_config_str,
            tuning_config=self.w2_tuning_config_str,
        )

apply(output, hidden_states, w1, w2, topk_weights, topk_ids, activation, global_num_experts, expert_map, a1q_scale, a2_scale, workspace13, workspace2, expert_tokens_meta, apply_router_weight_on_input)

Standard apply implementation for Humming batched grouped experts.

Note: Humming kernels handle weights internally through the layer object, so w1, w2, a2_scale are unused. a1q_scale is None on the usual path (Humming quantizes the w13 input itself); for block-FP8 activations it carries the scale computed before dispatch, which is forwarded to process_input so Humming skips the redundant w13 quantization. The output is written into workspace13 via the buffer management.

Source code in vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py

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,
) -> None:
    """Standard apply implementation for Humming batched grouped experts.

    Note: Humming kernels handle weights internally through the layer
    object, so w1, w2, a2_scale are unused. a1q_scale is None on the usual
    path (Humming quantizes the w13 input itself); for block-FP8 activations
    it carries the scale computed before dispatch, which is forwarded to
    process_input so Humming skips the redundant w13 quantization. The
    output is written into workspace13 via the buffer management.
    """
    assert not apply_router_weight_on_input
    assert expert_tokens_meta is not None

    hidden_states = hidden_states.view(-1, hidden_states.size(-1))
    # Keep the (batched) block-FP8 scale row-aligned with the flattened
    # [num_experts * max_tokens, K] hidden states above.
    if a1q_scale is not None and a1q_scale.dim() == 3:
        a1q_scale = a1q_scale.view(-1, a1q_scale.size(-1))
    valid_shape_m = self.estimate_local_valid_shape_m(topk_ids)
    expert_num_tokens = expert_tokens_meta.expert_num_tokens

    buffers = self.prepare_buffers(
        workspace13,
        workspace2,
        topk_ids.size(0),
        topk_ids.size(1),
        activation,
    )
    buffers["down_output"] = output

    inputs, input_scale, input_scale_2 = self.process_input(
        "w13",
        inputs=hidden_states,
        input_scale=a1q_scale,
        quanted_input=buffers.get("quanted_gate_up_input", None),
        expert_tokens=expert_num_tokens,
    )

    self.humming_forward(
        "w13",
        inputs=inputs,
        weight=w1,
        input_scale=input_scale,
        input_scale_2=input_scale_2,
        outputs=buffers["gate_up_output"],
        valid_shape_m=valid_shape_m,
        expert_layout=expert_num_tokens,
        compute_config=self.compute_config_str,
        tuning_config=self.w13_tuning_config_str,
    )

    inputs, input_scale, input_scale_2 = self.process_input(
        "w2",
        inputs=buffers["gate_up_output"],
        quanted_input=buffers["quanted_down_input"],
        expert_tokens=expert_num_tokens,
        activation=activation,
    )

    self.humming_forward(
        "w2",
        inputs=inputs,
        weight=w2,
        input_scale=input_scale,
        input_scale_2=input_scale_2,
        outputs=output.view(-1, hidden_states.size(-1)),
        valid_shape_m=valid_shape_m,
        expert_layout=expert_num_tokens,
        compute_config=self.compute_config_str,
        tuning_config=self.w2_tuning_config_str,
    )

HummingExpertsBase

Bases: FusedMoEExpertsModular

Attributes:

• expects_unquantized_inputs (bool) –

  Whether the prepare/finalize step should defer input quantization to

Source code in vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py

class HummingExpertsBase(mk.FusedMoEExpertsModular):
    quant_config: "HummingMoEQuantConfig"

    def __init__(
        self,
        moe_config: FusedMoEConfig,
        quant_config: FusedMoEQuantConfig,
        max_num_tokens: int | None = None,
        num_dispatchers: int | None = None,
    ):
        humming_quant_config = cast("HummingMoEQuantConfig", quant_config)
        self.humming_configs: dict[str, HummingLayerConfig] = {
            "w13": humming_quant_config.w1_humming_config,
            "w2": humming_quant_config.w2_humming_config,
        }
        self.locks = torch.zeros(1024, dtype=torch.int32, device=moe_config.device)
        self.num_experts = moe_config.num_local_experts
        self.global_num_experts = moe_config.num_experts
        self.quant_config = humming_quant_config
        self.init_humming_moe()

        if self.is_batched():
            assert max_num_tokens is not None and num_dispatchers is not None

        super().__init__(
            moe_config=moe_config,
            quant_config=quant_config,
            max_num_tokens=max_num_tokens,
            num_dispatchers=num_dispatchers,
        )

    def init_humming_moe(self):
        from vllm.utils.humming import get_heuristics_config

        self.compute_config = {
            "use_batch_invariant": envs.VLLM_BATCH_INVARIANT,
            "use_f16_accum": envs.VLLM_HUMMING_USE_F16_ACCUM,
            "gemm_type": self.humming_gemm_type().value,
        }
        self.w13_tuning_config = get_heuristics_config(
            layer_config=self.humming_configs["w13"],
            use_f16_accum=envs.VLLM_HUMMING_USE_F16_ACCUM,
            use_batch_invariant=envs.VLLM_BATCH_INVARIANT,
            gemm_type=self.humming_gemm_type(),
        )
        self.w2_tuning_config = get_heuristics_config(
            layer_config=self.humming_configs["w2"],
            use_f16_accum=envs.VLLM_HUMMING_USE_F16_ACCUM,
            use_batch_invariant=envs.VLLM_BATCH_INVARIANT,
            gemm_type=self.humming_gemm_type(),
        )
        self.compute_config_str = json.dumps(self.compute_config)
        self.w13_tuning_config_str = json.dumps(self.w13_tuning_config)
        self.w2_tuning_config_str = json.dumps(self.w2_tuning_config)

    def process_input(
        self,
        sublayer_name: str,
        inputs: torch.Tensor,
        quanted_input: torch.Tensor | None,
        input_scale: torch.Tensor | None = None,
        expert_tokens: torch.Tensor | None = None,
        activation: MoEActivation | None = None,
        scatter_idx: torch.Tensor | None = None,
        num_valid_tokens: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
        from vllm.model_executor.layers.quantization.utils.humming.activation import (
            get_humming_activation,
        )
        from vllm.utils.humming import may_process_input

        quant_config = self.quant_config
        config = self.humming_configs[sublayer_name]
        if input_scale is not None:
            assert scatter_idx is None
            if activation is not None:
                raise ValueError("Cannot apply activation to prequantized input")
            return inputs, input_scale, None

        activation_kwargs = {}
        if activation is not None:
            activation_config = self.activation_config
            activation_kwargs = get_humming_activation(activation, activation_config)

        prefix = "w1" if sublayer_name == "w13" else "w2"
        scale = getattr(quant_config, f"{prefix}_input_scale")
        scale2 = getattr(quant_config, f"{prefix}_input_scale_2")
        mode = config.input_quant_mode
        layout = "grouped_mask" if self.is_batched() else "normal"
        if scatter_idx is not None:
            layout = "scatter"

        inputs, group_scales, token_scales = may_process_input(
            config,
            inputs=inputs,
            outputs=quanted_input,
            token_scales=scale2 if mode.has_secondary_scale else scale,
            hadamard_block_size=getattr(quant_config, f"{prefix}_hadamard_block_size"),
            layout=layout,
            expert_tokens=expert_tokens if layout == "grouped_mask" else None,
            scatter_idx=scatter_idx,
            num_valid_tokens=num_valid_tokens,
            **activation_kwargs,
        )
        input_scale = group_scales if mode.has_group_scale else token_scales
        input_scale_2 = token_scales if mode.has_secondary_scale else None
        return inputs, input_scale, input_scale_2

    def humming_forward(
        self,
        sublayer_name: str,
        inputs: torch.Tensor,
        weight: torch.Tensor,
        input_scale: torch.Tensor | None,
        outputs: torch.Tensor,
        input_scale_2: torch.Tensor | None = None,
        **kwargs: Any,
    ) -> torch.Tensor:
        from vllm.utils.humming import humming_forward

        index = 1 if sublayer_name == "w13" else 2
        return humming_forward(
            self.humming_configs[sublayer_name],
            inputs=inputs,
            weight=weight,
            weight_scale=getattr(self.quant_config, f"w{index}_scale"),
            zero_point=getattr(self.quant_config, f"w{index}_zp"),
            bias=getattr(self.quant_config, f"w{index}_bias"),
            weight_scale_2=getattr(self.quant_config, f"g{index}_alphas"),
            input_scale=input_scale,
            input_scale_2=input_scale_2,
            outputs=outputs,
            locks=self.locks,
            **kwargs,
        )

    def _get_permute_scratch(
        self, topk: int, indices_only: bool = False
    ) -> MoEPermuteScratch | None:
        if not moe_permute_unpermute_supported():
            return None

        max_expanded_rows = (
            self.moe_config.max_num_tokens
            * self.moe_config.dp_size
            * self.moe_config.experts_per_token
        )
        return get_moe_permute_scratch(
            max_num_tokens=math.ceil(max_expanded_rows / topk),
            topk=topk,
            num_experts=self.moe_config.num_experts,
            num_local_experts=self.moe_config.num_local_experts,
            device=torch.device(self.moe_config.device),
            hidden_size=None if indices_only else self.moe_config.hidden_dim,
            hidden_dtype=None if indices_only else self.moe_config.in_dtype,
        )

    def get_global_valid_shape_m(self, topk_ids: torch.Tensor):
        ctx = get_forward_context()
        if ctx.dp_metadata is not None:
            num_tokens = ctx.dp_metadata.num_tokens_across_dp_cpu.sum().item()
            return num_tokens * self.moe_config.experts_per_token

        return topk_ids.size(0) * topk_ids.size(1)

    def estimate_local_valid_shape_m(self, topk_ids: torch.Tensor):
        # estimate shape_m for kernel tuning
        global_valid_shape_m = self.get_global_valid_shape_m(topk_ids)
        num_experts = self.num_experts
        global_num_experts = self.global_num_experts
        return math.ceil(global_valid_shape_m * num_experts / global_num_experts)

    @staticmethod
    def humming_gemm_type() -> "HummingGemmType":
        raise NotImplementedError

    @classmethod
    def is_batched(cls) -> bool:
        return cls.activation_format() == mk.FusedMoEActivationFormat.BatchedExperts

    @staticmethod
    def _supports_quant_scheme(
        weight_key: QuantKey | None,
        activation_key: QuantKey | None,
    ) -> bool:
        if (
            activation_key is not None
            and activation_key.scale == kNvfp4Dynamic.scale
            and activation_key.dtype == kNvfp4Dynamic.dtype
            and activation_key.scale2 is not None
            and not activation_key.scale2.static
            and activation_key.scale2.group_shape == GroupShape.PER_TOKEN
            and activation_key.scale2.dtype == torch.float32
        ):
            activation_key = replace(activation_key, scale2=kStaticTensorScale)
        SUPPORTED_W_A = [
            (kMxfp4Static, None),
            (kMxfp4Static, kMxfp4Dynamic),
            (kMxfp4Static, kMxfp8Dynamic),
            (kMxfp4Static, kFp8DynamicTokenSym),
            # MXFP4 weight (group-32 e8m0) with block-FP8 activation
            # (group-128 float32). Runs via WGMMA software dequant, so it
            # works on Hopper (SM90/H200) as well as Blackwell.
            (kMxfp4Static, kFp8Dynamic128Sym),
            (kNvfp4Static, None),
            (kNvfp4Static, kFp8DynamicTokenSym),
            (kMxfp8Static, None),
            (kMxfp8Static, kFp8DynamicTokenSym),
            (kFp8StaticChannelSym, None),
            (kFp8StaticChannelSym, kFp8DynamicTokenSym),
            (kFp8Static128BlockSym, None),
            (kFp8Static128BlockSym, kFp8DynamicTokenSym),
            (kInt4Static, None),
            (kInt4Static, kFp8DynamicTokenSym),
            (kInt8Static, None),
            (kInt8Static, kFp8DynamicTokenSym),
            # Checkpoint-driven (weight, activation) pairs the dense/MoE oracles
            # pass. Humming defers input quant (see expects_unquantized_inputs),
            # so the activation key does not constrain support.
            # fp8 (compressed-tensors / native / modelopt)
            (kFp8StaticChannelSym, kFp8StaticTensorSym),
            (kFp8StaticChannelSym, kFp8Dynamic128Sym),
            (kFp8StaticTensorSym, None),
            (kFp8StaticTensorSym, kFp8DynamicTokenSym),
            (kFp8StaticTensorSym, kFp8StaticTensorSym),
            (kFp8StaticTensorSym, kFp8Dynamic128Sym),
            (kFp8Static128BlockSym, kFp8Dynamic128Sym),
            # int8 (compressed-tensors w8a8 / experts_int8)
            (kInt8StaticChannelSym, None),
            (kInt8StaticChannelSym, kInt8DynamicTokenSym),
            # nvfp4 (compressed-tensors / modelopt / quark)
            (kNvfp4Static, kNvfp4Dynamic),
            # mxfp8 (compressed-tensors / modelopt / online)
            (kMxfp8Static, kMxfp8Dynamic),
        ]
        return (weight_key, activation_key) in SUPPORTED_W_A or (
            activation_key in (None, kFp8DynamicTokenSym, kInt8DynamicTokenSym)
            and _is_supported_wna16_weight_key(weight_key)
        )

    def _prequantizes_dispatch_activation(self) -> bool:
        """Whether the prepare/finalize step should quantize activations before
        the (EP all-to-all) dispatch instead of leaving it to Humming.

        This is enabled only for block-FP8 (group-128) activations: quantizing
        to FP8 before dispatch sends FP8 rather than BF16 over the interconnect,
        and Humming then consumes the pre-quantized FP8 + scale directly (see
        the apply() methods, which forward the dispatch scale into
        HummingExpertsBase.process_input, which is a
        no-op when an input scale is already supplied). The scale layout
        produced by vLLM's block-FP8 quantization ([M, K // 128] float32,
        row-major) matches what the Humming WGMMA grouped GEMM expects.
        """
        quant_config = self.quant_config
        return (
            quant_config.is_block_quantized
            and quant_config.quant_dtype == current_platform.fp8_dtype()
        )

    @property
    def expects_unquantized_inputs(self) -> bool:
        """Whether the prepare/finalize step should defer input quantization to
        the experts (by setting defer_input_quant=True and passing unquantized
        inputs).

        Humming normally quantizes inputs internally via
        HummingExpertsBase.process_input() in apply(), so we defer quantization
        (return True) to avoid quantizing twice -- once in prepare and once in
        Humming's apply().

        The exception is block-FP8 (group-128) activations, which are quantized
        before the dispatch to save interconnect bandwidth (see
        _prequantizes_dispatch_activation): for those we must NOT defer.
        """
        return not self._prequantizes_dispatch_activation()

    @staticmethod
    def _supports_current_device() -> bool:
        platform = current_platform
        return (
            has_humming()
            and platform.is_cuda()
            and platform.has_device_capability((7, 5))
        )

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

    @staticmethod
    def _supports_activation(activation: MoEActivation) -> bool:
        return apply_moe_activation_supported(activation)

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

    @staticmethod
    def _supports_batch_invariance() -> bool:
        return True

    def moe_problem_size(
        self,
        a1: torch.Tensor,
        w1: torch.Tensor,
        w2: torch.Tensor,
        topk_ids: torch.Tensor,
    ) -> tuple[int, int, int, int, int]:
        meta1 = self.humming_configs["w13"]
        meta2 = self.humming_configs["w2"]

        assert meta1.num_experts == meta2.num_experts

        num_experts = meta1.num_experts
        top_k = topk_ids.size(1)
        assert w1.size(0) == num_experts
        assert w2.size(0) == num_experts

        if not self.is_batched():
            num_tokens = a1.size(0)
            assert topk_ids.size(0) == num_tokens
        else:
            assert a1.dim() == 3
            assert a1.size(0) == num_experts
            num_tokens = a1.size(1)

        intermediate_dim = meta2.shape_k - meta2.pad_shape_k
        hidden_dim = meta1.shape_k - meta1.pad_shape_k
        return meta1.num_experts, num_tokens, intermediate_dim, hidden_dim, top_k

    def get_buffer_metas(self, M: int, topk: int, activation: MoEActivation):
        from vllm.utils.humming import GemmType as HummingGemmType
        from vllm.utils.humming import dtypes

        num_experts = self.num_experts
        w13_config = self.humming_configs["w13"]
        w2_config = self.humming_configs["w2"]
        gate_up_dim = w13_config.shape_n - w13_config.pad_shape_n
        intermediate_dim = w2_config.shape_k - w2_config.pad_shape_k
        K = w13_config.shape_k - w13_config.pad_shape_k
        assert isinstance(num_experts, int)
        assert isinstance(gate_up_dim, int)
        assert isinstance(intermediate_dim, int)
        assert isinstance(K, int)
        assert intermediate_dim == self.adjust_N_for_activation(gate_up_dim, activation)

        # hidden_states
        # -> quanted_gate_up_input (Hadamard + optional quantization)
        # -> gate_up_output
        # -> quanted_down_input (activation + Hadamard + optional quantization)
        # -> down_output
        # (-> output) (if not is_batched)
        # Neighboring nodes are required to utilize distinct workspaces.
        # The final output buffer is supplied by the modular kernel and may
        # alias workspace1.

        output_shape: tuple[int, ...]
        if self.is_batched():
            max_num_tokens = self.max_num_tokens
            num_dispatchers = self.num_dispatchers
            assert max_num_tokens is not None and num_dispatchers is not None
            real_shape_m = num_experts * max_num_tokens * num_dispatchers
            input_shape_m = real_shape_m
            output_shape = (num_experts, max_num_tokens * num_dispatchers, K)
        else:
            input_shape_m = M
            if self.humming_gemm_type() != HummingGemmType.INDEXED:
                input_shape_m = M * topk
            real_shape_m = M * topk
            output_shape = (M, K)

        a_dtype = self.humming_configs["w13"].a_dtype
        c_dtype = self.humming_configs["w13"].c_dtype
        down_a_dtype = w2_config.a_dtype
        torch_dtype_map = {
            dtypes.float16: torch.float16,
            dtypes.bfloat16: torch.bfloat16,
            dtypes.float32: torch.float32,
            dtypes.float8e3m4: torch.uint8,
            dtypes.float8e4m3: torch.float8_e4m3fn,
            dtypes.float8e5m2: torch.float8_e5m2,
            dtypes.int8: torch.int8,
            dtypes.int4: torch.uint8,
            dtypes.float4e0m3: torch.uint8,
            dtypes.float4e2m1: torch.uint8,
        }

        buffer_metas = {
            "quanted_gate_up_input": {
                "shape": (input_shape_m, K),
                "dtype": torch_dtype_map[a_dtype],
            },
            "gate_up_output": {
                "shape": (real_shape_m, gate_up_dim),
                "dtype": torch_dtype_map[c_dtype],
            },
            "quanted_down_input": {
                "shape": (real_shape_m, intermediate_dim),
                "dtype": torch_dtype_map[down_a_dtype],
            },
            "down_output": {
                "shape": output_shape if self.is_batched() else (real_shape_m, K),
                "dtype": torch_dtype_map[c_dtype],
            },
            "output": {
                "shape": output_shape,
                "dtype": torch_dtype_map[c_dtype],
            },
        }

        for key in buffer_metas:
            meta = buffer_metas[key]
            input_dtype = down_a_dtype if key == "quanted_down_input" else a_dtype
            if "quanted" in key and input_dtype.num_bits == 4:
                last_dim = meta["shape"][-1]
                if last_dim % 2 != 0:
                    raise ValueError(
                        f"Int4 packing requires last dimension to be even, "
                        f"got {last_dim} for buffer '{key}'"
                    )
                meta["shape"] = meta["shape"][:-1] + (last_dim // 2,)

        required_buffers = [
            "quanted_gate_up_input",
            "gate_up_output",
            "quanted_down_input",
            "down_output",
        ]

        # batched moe use down_output as output
        if not self.is_batched():
            required_buffers.append("output")

        return buffer_metas, required_buffers

    def _workspace_shapes(self, M: int, topk: int, activation: MoEActivation):
        buffer_metas, required_buffers = self.get_buffer_metas(M, topk, activation)

        workspace1_nbytes = 0
        workspace2_nbytes = 0

        for index, name in enumerate(required_buffers[::-1]):
            buffer_meta = buffer_metas[name]
            nelement = math.prod(buffer_meta["shape"])
            nbytes = nelement * buffer_meta["dtype"].itemsize
            if index % 2 == 0:
                workspace1_nbytes = max(workspace1_nbytes, nbytes)
            else:
                workspace2_nbytes = max(workspace2_nbytes, nbytes)

        output_key = "down_output" if self.is_batched() else "output"
        output_shape = buffer_metas[output_key]["shape"]
        elem_size = self.moe_config.in_dtype.itemsize

        return (
            (workspace1_nbytes // elem_size,),
            (workspace2_nbytes // elem_size,),
            output_shape,
        )

    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, ...]]:
        return self._workspace_shapes(M, topk, activation)

    def make_workspaces(self, M: int, topk: int, activation: MoEActivation):
        shapes = self._workspace_shapes(M, topk, activation)
        workspace1_shape, workspace2_shape, output_shape = shapes
        torch_dtype = self.moe_config.in_dtype
        workspace1, workspace2 = current_workspace_manager().get_simultaneous(
            (workspace1_shape, torch_dtype),
            (workspace2_shape, torch_dtype),
        )
        output = _resize_cache(workspace1, output_shape)
        return workspace1, workspace2, output

    def prepare_buffers(
        self,
        workspace1: torch.Tensor,
        workspace2: torch.Tensor,
        M: int,
        topk: int,
        activation: MoEActivation,
    ) -> dict[str, torch.Tensor]:
        buffer_metas, required_buffers = self.get_buffer_metas(M, topk, activation)
        buffers = {}
        for index, name in enumerate(required_buffers[::-1]):
            buffer_meta = buffer_metas[name]
            workspace = workspace1 if index % 2 == 0 else workspace2
            workspace = workspace.view(buffer_meta["dtype"])
            buffers[name] = _resize_cache(workspace, buffer_meta["shape"])

        return buffers

    # Note: apply method is implemented by subclasses following the
    # standard FusedMoEExpertsModular.apply signature

    @staticmethod
    def is_supported_config(
        cls: type[mk.FusedMoEExperts],
        moe_config: FusedMoEConfig,
        weight_key: QuantKey | None,
        activation_key: QuantKey | None,
        activation_format: mk.FusedMoEActivationFormat,
    ) -> tuple[bool, str | None]:
        supported, reason = mk.FusedMoEExpertsModular.is_supported_config(
            cls,
            moe_config,
            weight_key,
            activation_key,
            activation_format,
        )

        if supported:
            assert hasattr(cls, "humming_gemm_type")
            gemm_type = cls.humming_gemm_type().value.lower()
            preferred_gemm_type = get_humming_moe_gemm_type(moe_config)
            supported = preferred_gemm_type.lower() == gemm_type
            if not supported:
                reason = (
                    f"preferred gemm type {preferred_gemm_type} != "
                    f"supported gemm type {gemm_type}"
                )

        return supported, reason

    def apply_activation(
        self,
        activation: MoEActivation,
        output: torch.Tensor,
        input: torch.Tensor,
        valid_token_counts: torch.Tensor | None = None,
    ) -> None:
        activation_kwargs: dict[str, Any] = dict(
            activation=activation,
            input=input,
            output=output,
        )
        if valid_token_counts is not None:
            activation_kwargs["valid_token_counts"] = valid_token_counts
        self.activation(**activation_kwargs)

expects_unquantized_inputs property

Whether the prepare/finalize step should defer input quantization to the experts (by setting defer_input_quant=True and passing unquantized inputs).

Humming normally quantizes inputs internally via HummingExpertsBase.process_input() in apply(), so we defer quantization (return True) to avoid quantizing twice -- once in prepare and once in Humming's apply().

The exception is block-FP8 (group-128) activations, which are quantized before the dispatch to save interconnect bandwidth (see _prequantizes_dispatch_activation): for those we must NOT defer.

_prequantizes_dispatch_activation()

Whether the prepare/finalize step should quantize activations before the (EP all-to-all) dispatch instead of leaving it to Humming.

This is enabled only for block-FP8 (group-128) activations: quantizing to FP8 before dispatch sends FP8 rather than BF16 over the interconnect, and Humming then consumes the pre-quantized FP8 + scale directly (see the apply() methods, which forward the dispatch scale into HummingExpertsBase.process_input, which is a no-op when an input scale is already supplied). The scale layout produced by vLLM's block-FP8 quantization ([M, K // 128] float32, row-major) matches what the Humming WGMMA grouped GEMM expects.

Source code in vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py

def _prequantizes_dispatch_activation(self) -> bool:
    """Whether the prepare/finalize step should quantize activations before
    the (EP all-to-all) dispatch instead of leaving it to Humming.

    This is enabled only for block-FP8 (group-128) activations: quantizing
    to FP8 before dispatch sends FP8 rather than BF16 over the interconnect,
    and Humming then consumes the pre-quantized FP8 + scale directly (see
    the apply() methods, which forward the dispatch scale into
    HummingExpertsBase.process_input, which is a
    no-op when an input scale is already supplied). The scale layout
    produced by vLLM's block-FP8 quantization ([M, K // 128] float32,
    row-major) matches what the Humming WGMMA grouped GEMM expects.
    """
    quant_config = self.quant_config
    return (
        quant_config.is_block_quantized
        and quant_config.quant_dtype == current_platform.fp8_dtype()
    )

HummingGroupedExperts

Bases: HummingExpertsBase

Methods:

• apply –

  Standard apply implementation for Humming grouped experts.

Source code in vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py

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

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

    @staticmethod
    def humming_gemm_type() -> "HummingGemmType":
        from vllm.utils.humming import GemmType as HummingGemmType

        return HummingGemmType.GROUPED_CONTIGUOUS

    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,
    ) -> None:
        """Standard apply implementation for Humming grouped experts.

        Note: Humming kernels handle weights internally through the layer
        object, so w1, w2, a2_scale are unused. a1q_scale is None on the usual
        path (Humming quantizes the w13 input itself); for block-FP8 activations
        it carries the scale computed before dispatch. It is permuted alongside
        the tokens by moe_permute and forwarded to may_quant_input so Humming
        skips the redundant w13 quantization. The output is written into
        workspace13 via the buffer management.
        """
        assert not apply_router_weight_on_input

        valid_shape_m = self.estimate_local_valid_shape_m(topk_ids)

        buffers = self.prepare_buffers(
            workspace13,
            workspace2,
            topk_ids.size(0),
            topk_ids.size(1),
            activation,
        )
        buffers["output"] = output

        if a1q_scale is None:
            scratch = self._get_permute_scratch(topk_ids.size(1), indices_only=True)
            assert scratch is not None
            scatter_metadata = moe_prepare_scatter(topk_ids, expert_map, scratch)
            expert_offsets, scatter_idx = scatter_metadata
            inv_perm = scatter_idx.flatten()
            num_valid_tokens = expert_offsets[-1:] if expert_map is not None else None
        else:
            hidden_states, a1q_scale, expert_offsets, inv_perm, _ = moe_permute(
                hidden_states=hidden_states,
                a1q_scale=a1q_scale,
                topk_ids=topk_ids,
                n_expert=global_num_experts,
                n_local_expert=self.num_experts,
                expert_map=expert_map,
                scratch=self._get_permute_scratch(topk_ids.size(1)),
            )
            scatter_idx = None
            num_valid_tokens = None

        inputs, input_scale, input_scale_2 = self.process_input(
            "w13",
            inputs=hidden_states,
            input_scale=a1q_scale,
            quanted_input=buffers.get("quanted_gate_up_input", None),
            scatter_idx=scatter_idx,
            num_valid_tokens=num_valid_tokens,
        )

        self.humming_forward(
            "w13",
            inputs=inputs,
            weight=w1,
            input_scale=input_scale,
            input_scale_2=input_scale_2,
            outputs=buffers["gate_up_output"],
            valid_shape_m=valid_shape_m,
            expert_layout=expert_offsets,
            compute_config=self.compute_config_str,
            tuning_config=self.w13_tuning_config_str,
        )

        inputs, input_scale, input_scale_2 = self.process_input(
            "w2",
            inputs=buffers["gate_up_output"],
            quanted_input=buffers["quanted_down_input"],
            activation=activation,
            num_valid_tokens=expert_offsets[-1:],
        )

        self.humming_forward(
            "w2",
            inputs=inputs,
            weight=w2,
            input_scale=input_scale,
            input_scale_2=input_scale_2,
            outputs=buffers["down_output"],
            valid_shape_m=valid_shape_m,
            expert_layout=expert_offsets,
            compute_config=self.compute_config_str,
            tuning_config=self.w2_tuning_config_str,
        )

        moe_unpermute(
            out=output,
            permuted_hidden_states=buffers["down_output"].view(*topk_ids.shape, -1),
            topk_weights=topk_weights,
            inv_permuted_idx=inv_perm,
            expert_first_token_offset=expert_offsets,
        )

apply(output, hidden_states, w1, w2, topk_weights, topk_ids, activation, global_num_experts, expert_map, a1q_scale, a2_scale, workspace13, workspace2, expert_tokens_meta, apply_router_weight_on_input)

Standard apply implementation for Humming grouped experts.

Note: Humming kernels handle weights internally through the layer object, so w1, w2, a2_scale are unused. a1q_scale is None on the usual path (Humming quantizes the w13 input itself); for block-FP8 activations it carries the scale computed before dispatch. It is permuted alongside the tokens by moe_permute and forwarded to may_quant_input so Humming skips the redundant w13 quantization. The output is written into workspace13 via the buffer management.

Source code in vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py

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,
) -> None:
    """Standard apply implementation for Humming grouped experts.

    Note: Humming kernels handle weights internally through the layer
    object, so w1, w2, a2_scale are unused. a1q_scale is None on the usual
    path (Humming quantizes the w13 input itself); for block-FP8 activations
    it carries the scale computed before dispatch. It is permuted alongside
    the tokens by moe_permute and forwarded to may_quant_input so Humming
    skips the redundant w13 quantization. The output is written into
    workspace13 via the buffer management.
    """
    assert not apply_router_weight_on_input

    valid_shape_m = self.estimate_local_valid_shape_m(topk_ids)

    buffers = self.prepare_buffers(
        workspace13,
        workspace2,
        topk_ids.size(0),
        topk_ids.size(1),
        activation,
    )
    buffers["output"] = output

    if a1q_scale is None:
        scratch = self._get_permute_scratch(topk_ids.size(1), indices_only=True)
        assert scratch is not None
        scatter_metadata = moe_prepare_scatter(topk_ids, expert_map, scratch)
        expert_offsets, scatter_idx = scatter_metadata
        inv_perm = scatter_idx.flatten()
        num_valid_tokens = expert_offsets[-1:] if expert_map is not None else None
    else:
        hidden_states, a1q_scale, expert_offsets, inv_perm, _ = moe_permute(
            hidden_states=hidden_states,
            a1q_scale=a1q_scale,
            topk_ids=topk_ids,
            n_expert=global_num_experts,
            n_local_expert=self.num_experts,
            expert_map=expert_map,
            scratch=self._get_permute_scratch(topk_ids.size(1)),
        )
        scatter_idx = None
        num_valid_tokens = None

    inputs, input_scale, input_scale_2 = self.process_input(
        "w13",
        inputs=hidden_states,
        input_scale=a1q_scale,
        quanted_input=buffers.get("quanted_gate_up_input", None),
        scatter_idx=scatter_idx,
        num_valid_tokens=num_valid_tokens,
    )

    self.humming_forward(
        "w13",
        inputs=inputs,
        weight=w1,
        input_scale=input_scale,
        input_scale_2=input_scale_2,
        outputs=buffers["gate_up_output"],
        valid_shape_m=valid_shape_m,
        expert_layout=expert_offsets,
        compute_config=self.compute_config_str,
        tuning_config=self.w13_tuning_config_str,
    )

    inputs, input_scale, input_scale_2 = self.process_input(
        "w2",
        inputs=buffers["gate_up_output"],
        quanted_input=buffers["quanted_down_input"],
        activation=activation,
        num_valid_tokens=expert_offsets[-1:],
    )

    self.humming_forward(
        "w2",
        inputs=inputs,
        weight=w2,
        input_scale=input_scale,
        input_scale_2=input_scale_2,
        outputs=buffers["down_output"],
        valid_shape_m=valid_shape_m,
        expert_layout=expert_offsets,
        compute_config=self.compute_config_str,
        tuning_config=self.w2_tuning_config_str,
    )

    moe_unpermute(
        out=output,
        permuted_hidden_states=buffers["down_output"].view(*topk_ids.shape, -1),
        topk_weights=topk_weights,
        inv_permuted_idx=inv_perm,
        expert_first_token_offset=expert_offsets,
    )

HummingIndexedExperts

Bases: HummingExpertsBase

Methods:

• apply –

  Standard apply implementation for Humming indexed experts.

Source code in vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py

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

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

    @staticmethod
    def humming_gemm_type() -> "HummingGemmType":
        from vllm.utils.humming import GemmType as HummingGemmType

        return HummingGemmType.INDEXED

    def prepare_humming_moe_kwargs(
        self,
        topk_ids: torch.Tensor,
        expert_map: torch.Tensor | None,
        expert_tokens_meta: mk.ExpertTokensMetadata | None,
    ) -> tuple[dict[str, Any], dict[str, Any], torch.Tensor | None]:
        valid_shape_m = self.estimate_local_valid_shape_m(topk_ids)

        moe_block_size = None
        for min_shape_m, max_shape_m, config in self.w13_tuning_config:
            if valid_shape_m > min_shape_m and valid_shape_m <= max_shape_m:
                moe_block_size = config["block_shape"][0]
                break

        if moe_block_size is None:
            logger.warning_once(
                "No tuning config found for shape %s, using default block_size=64",
                valid_shape_m,
            )
            moe_block_size = 64

        use_scatter = False
        if expert_map is not None:
            num_sms = num_compute_units(topk_ids.get_device())
            use_scatter = topk_ids.numel() > 2 * num_sms

        alignment = moe_align_block_size(
            topk_ids=topk_ids,
            block_size=moe_block_size,
            num_experts=self.global_num_experts,
            expert_map=expert_map,
            ignore_invalid_experts=True,
            return_scatter_idx=use_scatter,
        )
        sorted_ids, expert_ids, num_tokens_padded = alignment[:3]
        scatter_idx = alignment[3] if len(alignment) == 4 else None

        moe_common_kwargs = {
            "sorted_ids": sorted_ids,
            "expert_ids": expert_ids,
            "num_tokens_padded": num_tokens_padded,
            "compute_config": self.compute_config_str,
            "valid_shape_m": valid_shape_m,
        }

        top_k = topk_ids.size(1)
        moe_kwargs1 = {"top_k": top_k, "tuning_config": self.w13_tuning_config_str}
        moe_kwargs2 = {"top_k": 1, "tuning_config": self.w2_tuning_config_str}
        moe_kwargs1.update(moe_common_kwargs)
        moe_kwargs2.update(moe_common_kwargs)

        w2_block_size = 64
        for lower, upper, config in self.w2_tuning_config:
            if lower < valid_shape_m <= upper:
                w2_block_size = config["block_shape"][0]
                break
        if w2_block_size != moe_block_size:
            sorted_ids2, expert_ids2, num_tokens_padded2 = moe_align_block_size(
                topk_ids=topk_ids,
                block_size=w2_block_size,
                num_experts=self.global_num_experts,
                expert_map=expert_map,
                ignore_invalid_experts=True,
            )
            moe_kwargs2.update(
                sorted_ids=sorted_ids2,
                expert_ids=expert_ids2,
                num_tokens_padded=num_tokens_padded2,
            )

        return moe_kwargs1, moe_kwargs2, scatter_idx

    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,
    ) -> None:
        """Standard apply implementation for Humming indexed experts.

        Note: Humming kernels handle weights internally through the layer
        object, so w1, w2, a2_scale are unused. a1q_scale is None on the usual
        path (Humming quantizes the w13 input itself); for block-FP8 activations
        it carries the scale computed before dispatch, which is forwarded to
        process_input so Humming skips the redundant w13 quantization. The
        output is written into workspace13 via the buffer management.
        """
        assert not apply_router_weight_on_input

        hidden_states = hidden_states.view(-1, hidden_states.size(-1))
        buffers = self.prepare_buffers(
            workspace13,
            workspace2,
            topk_ids.size(0),
            topk_ids.size(1),
            activation,
        )
        buffers["output"] = output

        moe_kwargs1, moe_kwargs2, scatter_idx = self.prepare_humming_moe_kwargs(
            topk_ids=topk_ids,
            expert_map=expert_map,
            expert_tokens_meta=expert_tokens_meta,
        )

        inputs, input_scale, input_scale_2 = self.process_input(
            "w13",
            inputs=hidden_states,
            input_scale=a1q_scale,
            quanted_input=buffers.get("quanted_gate_up_input", None),
        )

        self.humming_forward(
            "w13",
            inputs=inputs,
            weight=w1,
            input_scale=input_scale,
            input_scale_2=input_scale_2,
            outputs=buffers["gate_up_output"],
            **moe_kwargs1,
        )

        inputs, input_scale, input_scale_2 = self.process_input(
            "w2",
            inputs=buffers["gate_up_output"],
            quanted_input=buffers["quanted_down_input"],
            activation=activation,
            scatter_idx=scatter_idx,
        )

        self.humming_forward(
            "w2",
            inputs=inputs,
            weight=w2,
            input_scale=input_scale,
            input_scale_2=input_scale_2,
            outputs=buffers["down_output"].view(-1, hidden_states.size(-1)),
            **moe_kwargs2,
        )

        # expert_map masks any non-local id; num_valid_tokens bounds the
        # persistent kernel to the real token rows [0, num_recv) so the padding
        # tail is never iterated (CUDA-graph-safe device scalar).
        valid_tokens = None
        if expert_tokens_meta and expert_tokens_meta.psum_recv_per_rank is not None:
            valid_tokens = expert_tokens_meta.psum_recv_per_rank[-1:]

        moe_fused_mul_sum(
            inputs=buffers["down_output"].view(*topk_ids.shape, -1),
            topk_weights=topk_weights,
            topk_ids=topk_ids,
            expert_map=expert_map,
            outputs=output,
            num_valid_tokens=valid_tokens,
        )

apply(output, hidden_states, w1, w2, topk_weights, topk_ids, activation, global_num_experts, expert_map, a1q_scale, a2_scale, workspace13, workspace2, expert_tokens_meta, apply_router_weight_on_input)

Standard apply implementation for Humming indexed experts.

Note: Humming kernels handle weights internally through the layer object, so w1, w2, a2_scale are unused. a1q_scale is None on the usual path (Humming quantizes the w13 input itself); for block-FP8 activations it carries the scale computed before dispatch, which is forwarded to process_input so Humming skips the redundant w13 quantization. The output is written into workspace13 via the buffer management.

Source code in vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py

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,
) -> None:
    """Standard apply implementation for Humming indexed experts.

    Note: Humming kernels handle weights internally through the layer
    object, so w1, w2, a2_scale are unused. a1q_scale is None on the usual
    path (Humming quantizes the w13 input itself); for block-FP8 activations
    it carries the scale computed before dispatch, which is forwarded to
    process_input so Humming skips the redundant w13 quantization. The
    output is written into workspace13 via the buffer management.
    """
    assert not apply_router_weight_on_input

    hidden_states = hidden_states.view(-1, hidden_states.size(-1))
    buffers = self.prepare_buffers(
        workspace13,
        workspace2,
        topk_ids.size(0),
        topk_ids.size(1),
        activation,
    )
    buffers["output"] = output

    moe_kwargs1, moe_kwargs2, scatter_idx = self.prepare_humming_moe_kwargs(
        topk_ids=topk_ids,
        expert_map=expert_map,
        expert_tokens_meta=expert_tokens_meta,
    )

    inputs, input_scale, input_scale_2 = self.process_input(
        "w13",
        inputs=hidden_states,
        input_scale=a1q_scale,
        quanted_input=buffers.get("quanted_gate_up_input", None),
    )

    self.humming_forward(
        "w13",
        inputs=inputs,
        weight=w1,
        input_scale=input_scale,
        input_scale_2=input_scale_2,
        outputs=buffers["gate_up_output"],
        **moe_kwargs1,
    )

    inputs, input_scale, input_scale_2 = self.process_input(
        "w2",
        inputs=buffers["gate_up_output"],
        quanted_input=buffers["quanted_down_input"],
        activation=activation,
        scatter_idx=scatter_idx,
    )

    self.humming_forward(
        "w2",
        inputs=inputs,
        weight=w2,
        input_scale=input_scale,
        input_scale_2=input_scale_2,
        outputs=buffers["down_output"].view(-1, hidden_states.size(-1)),
        **moe_kwargs2,
    )

    # expert_map masks any non-local id; num_valid_tokens bounds the
    # persistent kernel to the real token rows [0, num_recv) so the padding
    # tail is never iterated (CUDA-graph-safe device scalar).
    valid_tokens = None
    if expert_tokens_meta and expert_tokens_meta.psum_recv_per_rank is not None:
        valid_tokens = expert_tokens_meta.psum_recv_per_rank[-1:]

    moe_fused_mul_sum(
        inputs=buffers["down_output"].view(*topk_ids.shape, -1),
        topk_weights=topk_weights,
        topk_ids=topk_ids,
        expert_map=expert_map,
        outputs=output,
        num_valid_tokens=valid_tokens,
    )