lightseekorg--tokenspeed/python/tokenspeed/runtime/layers/moe/expert.py

# Copyright (c) 2026 LightSeek Foundation
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.


import tokenspeed_kernel
import torch

from tokenspeed.runtime.distributed.process_group_manager import (
    process_group_manager as pg_manager,
)
from tokenspeed.runtime.layers.activation import SwigluArg
from tokenspeed.runtime.layers.moe.topk import TopKOutput, TopKOutputFormat
from tokenspeed.runtime.layers.moe.types import MoELayerSpec
from tokenspeed.runtime.layers.moe.utils import (
    RoutingMethodType,
    get_all2all_backend,
    get_moe_backend,
)
from tokenspeed.runtime.layers.moe.weights import create_layer_weights
from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig
from tokenspeed.runtime.layers.quantization.utils import (
    should_exclude_quant_module,
    should_ignore_quant_layer,
)
from tokenspeed.runtime.utils.env import global_server_args_dict
from tokenspeed.runtime.utils.pdl import pdl_enabled


class MoELayer(torch.nn.Module):
    def __init__(
        self,
        top_k: int,
        num_experts: int,
        hidden_size: int,
        intermediate_size: int,
        quant_config: QuantizationConfig,
        layer_index: int,
        prefix: str = "",
        tp_rank: int | None = None,
        tp_size: int | None = None,
        ep_rank: int | None = None,
        ep_size: int | None = None,
        zero_expert_type: str = "",
        activation: str = "silu",
        activation_alpha=None,
        swiglu_limit=None,
        swiglu_beta: float | None = None,
        w13_input_layout: str = "concatenated",
        with_bias=False,
        routing_config: dict = {},
    ):
        super().__init__()
        self.layer_index = layer_index
        self.prefix = prefix
        self.top_k = top_k
        self.num_experts = num_experts
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.quant_config = quant_config
        self.ep_num_redundant_experts = global_server_args_dict[
            "ep_num_redundant_experts"
        ]
        self.zero_expert_type = zero_expert_type
        self.activation = activation
        self.swiglu_arg = None
        if self.activation == "swiglu":
            self.swiglu_arg = SwigluArg(alpha=activation_alpha, limit=swiglu_limit)
        # Per-model knobs the MoE backend reads in process_weights_after_loading.
        # ``swiglu_beta``: gpt-oss uses silu(α·gate)·(up + 1) and sets 1.0;
        # standard SwiGLU (e.g. deepseek-v4) leaves it None.
        # ``w13_input_layout``: "interleaved" for HF gpt-oss-style row layout
        # ([w1_0, w3_0, w1_1, w3_1, ...]); "concatenated" (default) for the
        # shared MoE checkpoint loader's [w1_all | w3_all] block layout.
        self.swiglu_beta = swiglu_beta
        if w13_input_layout not in {"interleaved", "concatenated"}:
            raise ValueError(
                f"w13_input_layout must be 'interleaved' or 'concatenated', "
                f"got {w13_input_layout!r}"
            )
        self.w13_input_layout = w13_input_layout

        if tp_rank is None:
            assert tp_size is None
            tp_rank, tp_size = 0, 1
        self.tp_rank, self.tp_size = tp_rank, tp_size
        self.moe_tp_size = self.tp_size
        if ep_rank is None:
            assert ep_size is None
            ep_rank, ep_size = 0, 1
        self.ep_rank, self.ep_size = ep_rank, ep_size

        if tp_size > 1 and ep_size > 1:
            raise ValueError("Mixed TP and EP is not supported yet.")

        num_local_experts = num_experts // self.ep_size

        self.num_local_experts = num_local_experts
        self._spec = MoELayerSpec(
            top_k=top_k,
            num_experts=num_experts,
            num_local_experts=num_local_experts,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            activation=activation,
            tp_rank=self.tp_rank,
            tp_size=self.tp_size,
            ep_rank=self.ep_rank,
            ep_size=self.ep_size,
            prefix=prefix,
            a2a_backend=get_all2all_backend().value,
        )

        # Routing config
        self.routing_config = routing_config
        self._correction_bias = routing_config.get("correction_bias", None)
        self._routing_method_type = routing_config.get(
            "routing_method_type", RoutingMethodType.DeepSeekV3
        )
        self._routing_logits_dtype = torch.bfloat16
        if self._routing_method_type in (
            RoutingMethodType.DeepSeekV3,
            RoutingMethodType.MiniMax2,
        ):
            self._routing_logits_dtype = torch.float32
        self._n_group = routing_config.get("n_group", 0)
        self._topk_group = routing_config.get("topk_group", 0)
        self._routed_scaling_factor = routing_config.get("routed_scaling_factor", 1.0)
        self._normalize_topk_weights = routing_config.get(
            "normalize_topk_weights", True
        )

        # Quantization config. ignored_layers (compressed-tensors) keys the MoE
        # block; exclude_modules (ModelOpt) keys the fused experts.
        self._quant_kind = "unquant"
        if (
            quant_config is not None
            and not should_ignore_quant_layer(self.prefix, quant_config.ignored_layers)
            and not should_exclude_quant_module(
                f"{self.prefix}.experts", quant_config.exclude_modules
            )
        ):
            self._quant_kind = quant_config.moe_weight_dtype()

        fp8_scale_block_shape = None
        internal_activation_dtype = "input"
        if self._quant_kind == "fp8":
            fp8_scale_block_shape = tuple(self.quant_config.weight_block_size)
        if self._quant_kind == "mxfp4":
            if self.quant_config.is_w4a8_fp8:
                internal_activation_dtype = "fp8"
            elif getattr(self.quant_config, "use_dynamic_mxfp4_activations", False):
                internal_activation_dtype = "input"

        input_dtype = torch.get_default_dtype()
        if input_dtype not in {torch.float16, torch.bfloat16}:
            input_dtype = torch.float16

        deepep_group = None
        if self._spec.use_deepep:
            mapping = global_server_args_dict["mapping"]
            deepep_group = pg_manager.get_process_group(
                "nccl",
                mapping.moe.tp_ep_group,
            )

        # Moe Backend plan
        moe_backend = get_moe_backend().value
        moe_backend = None if moe_backend == "auto" else moe_backend
        self.plan = tokenspeed_kernel.moe_plan(
            self._quant_kind,
            input_dtype=input_dtype,
            activation=self.activation,
            a2a_backend=self._spec.a2a_backend,
            ep_size=self.ep_size,
            ispp=self.intermediate_size // self.tp_size,
            fp8_scale_block_shape=fp8_scale_block_shape,
            internal_activation_dtype=internal_activation_dtype,
            with_bias=with_bias,
            deepep_group=deepep_group,
            solution=moe_backend,
        )

        create_layer_weights(
            self._spec,
            self,
            self._quant_kind,
            self.quant_config,
            with_bias=with_bias,
            solution=self.plan["solution"],
        )
        self._weights_processed = False

    def process_weights_after_loading(self, module) -> None:
        if self._weights_processed:
            return

        tokenspeed_kernel.moe_process_weights(self.plan, module)
        self._weights_processed = True

    @property
    def support_routing(self) -> bool:
        return self.plan["support_routing"]

    @property
    def topk_output_format(self):
        if self.support_routing:
            return TopKOutputFormat.BYPASSED
        return TopKOutputFormat.STANDARD

    @property
    def supports_deferred_finalize(self) -> bool:
        return self.plan["supports_deferred_finalize"]

    def forward_zero_experts(self, topk_output):
        zero_expert_limit = self.num_experts
        if self.ep_num_redundant_experts is not None:
            zero_expert_limit = zero_expert_limit - self.ep_num_redundant_experts

        normal_expert_mask = topk_output.topk_ids >= zero_expert_limit
        topk_output.topk_ids[normal_expert_mask] = -1
        if self.zero_expert_type == "copy":
            topk_output.topk_weights[normal_expert_mask] = 1.0
        if self.zero_expert_type == "drop":
            topk_output.topk_weights[normal_expert_mask] = 0.0

    def forward(
        self,
        hidden_states: torch.Tensor,
        topk_output: TopKOutput,
        num_global_tokens: int,
        max_num_tokens_per_gpu: int,
        do_finalize: bool = True,
    ):
        if not do_finalize and not self.supports_deferred_finalize:
            raise AssertionError("MoELayer does not support do_finalize=False")

        if self.support_routing:
            return tokenspeed_kernel.moe_apply(
                self.plan,
                hidden_states,
                self,
                topk_output.router_logits,
                num_tokens_global=num_global_tokens,
                max_num_tokens_per_gpu=max_num_tokens_per_gpu,
                do_finalize=do_finalize,
                enable_pdl=pdl_enabled(),
            )
        else:
            return tokenspeed_kernel.moe_apply(
                self.plan,
                hidden_states,
                self,
                topk_output.router_logits,
                topk_weights=topk_output.topk_weights,
                topk_ids=topk_output.topk_ids,
                num_tokens_global=num_global_tokens,
                max_num_tokens_per_gpu=max_num_tokens_per_gpu,
                do_finalize=do_finalize,
                enable_pdl=pdl_enabled(),
            )