sgl-project--sglang/python/sglang/srt/models/minimax_m3.py

# Copyright 2023-2024 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

# Adapted from DeepSeek and Mixtral implementation
"""Inference-only MiniMax M3 model compatible with HuggingFace weights."""

import logging
from contextlib import nullcontext
from typing import Iterable, List, Optional, Set, Tuple, Union

import torch
from torch import nn
from transformers import PretrainedConfig

from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo
from sglang.srt.configs.model_config import (
    get_minimax_sparse_disable_value_layer_ids,
    get_minimax_sparse_layer_ids,
)
from sglang.srt.distributed import (
    get_pp_group,
    tensor_model_parallel_all_reduce,
)
from sglang.srt.environ import envs
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo
from sglang.srt.layers.activation import SiluAndMul
from sglang.srt.layers.communicator import (
    LayerCommunicator,
    LayerScatterModes,
    ScatterMode,
    enable_moe_dense_fully_dp,
)
from sglang.srt.layers.dp_attention import is_dp_attention_enabled
from sglang.srt.layers.layernorm import GemmaRMSNorm, RMSNorm
from sglang.srt.layers.linear import (
    MergedColumnParallelLinear,
    QKVParallelLinear,
    ReplicatedLinear,
    RowParallelLinear,
)
from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
from sglang.srt.layers.moe.topk import TopK
from sglang.srt.layers.moe.utils import get_moe_a2a_backend
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.layers.rotary_embedding import get_rope
from sglang.srt.layers.utils import PPMissingLayer
from sglang.srt.layers.utils.common import get_layer_id
from sglang.srt.layers.vocab_parallel_embedding import (
    ParallelLMHead,
    VocabParallelEmbedding,
)
from sglang.srt.model_executor.cuda_graph_config import (
    Backend,
    Phase,
    check_cuda_graph_backend,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_executor.forward_context import (
    get_forward_context,
    has_forward_context,
)
from sglang.srt.model_loader.weight_utils import (
    default_weight_loader,
    maybe_remap_kv_scale_name,
)
from sglang.srt.models.minimax_m2 import MiniMaxM2RMSNormTP
from sglang.srt.runtime_context import get_parallel, get_server_args
from sglang.srt.utils import (
    add_prefix,
    get_device_sm,
    is_cuda,
    is_hip,
    log_info_on_rank0,
    make_layers,
)
from sglang.srt.utils.hf_transformers_utils import get_rope_config

_is_cuda = is_cuda()
_is_hip = is_hip()
_device_sm = get_device_sm()

_FP8_KV_DTYPES = (
    torch.float8_e4m3fn,
    torch.float8_e5m2,
    torch.float8_e4m3fnuz,
)

# rotary_dim required by the fused qknorm+rope JIT kernel: rotary_dim/2 must
# equal the CUDA warp size (32) so each warp norms+ropes one head in one pass.
_M3_FUSED_QKNORM_ROPE_ROTARY_DIM = 64

_has_rocm_qk_norm_rope = False
if _is_hip:
    try:
        from sglang.jit_kernel.minimax_m3.qk_norm_rope import (
            qk_gemma_rmsnorm_rope,
            sparse_qk_index_gemma_rmsnorm_rope,
            sparse_qk_index_gemma_rmsnorm_rope_cache,
        )

        _has_rocm_qk_norm_rope = True
    except ImportError:
        _has_rocm_qk_norm_rope = False

logger = logging.getLogger(__name__)


class MultiHeadRMSNorm(nn.Module):
    def __init__(
        self,
        num_heads: int,
        head_dim: int,
        eps: float = 1e-6,
        apply_layernorm_1p: bool = False,
    ) -> None:
        super().__init__()
        self.tp_world = get_parallel().attn_tp_size
        self.tp_rank = get_parallel().attn_tp_rank
        self.num_heads = num_heads
        self.num_heads_per_tp = num_heads // self.tp_world
        self.head_dim = head_dim
        self.weight = nn.Parameter(
            torch.ones(self.num_heads_per_tp, self.head_dim, dtype=torch.float32)
        )
        self.weight.weight_loader = self.weight_loader
        self.apply_layernorm_1p = apply_layernorm_1p
        self.variance_epsilon = eps

    @staticmethod
    def weight_loader(
        param: nn.Parameter,
        loaded_weight: torch.Tensor,
    ) -> None:
        tp_world = get_parallel().attn_tp_size
        tp_rank = get_parallel().attn_tp_rank

        shard_size = loaded_weight.shape[0] // tp_world
        shard = slice(tp_rank * shard_size, (tp_rank + 1) * shard_size)
        param.data.copy_(loaded_weight[shard].reshape_as(param))

    def forward(
        self,
        x: torch.Tensor,
    ) -> torch.Tensor:
        orig_dtype = x.dtype
        x = x.view(-1, self.num_heads_per_tp, self.head_dim).to(torch.float32)
        variance = x.pow(2).mean(dim=-1, keepdim=True, dtype=torch.float32)
        x = x * torch.rsqrt(variance + self.variance_epsilon)
        if self.apply_layernorm_1p:
            x = x * (self.weight + 1)[None, ...]
        else:
            x = x * self.weight[None, ...]
        x = x.view(-1, self.num_heads_per_tp * self.head_dim)
        return x.to(orig_dtype)


class _FusedQKVIndexProj(nn.Module):
    def __init__(
        self,
        quant_method,
        weight: torch.Tensor,
        weight_scale_inv: Optional[torch.Tensor],
        input_size_per_partition: int,
        logical_widths: List[int],
        orig_dtype: torch.dtype,
    ) -> None:
        super().__init__()
        # Named ``_qm`` (not ``quant_method``) so the loader's post-process loop
        # skips this module; the backend scale layout is derived once below.
        self._qm = quant_method
        self.register_parameter("weight", nn.Parameter(weight, requires_grad=False))
        self.input_size_per_partition = input_size_per_partition
        self.output_size_per_partition = weight.shape[0]
        self.logical_widths = logical_widths
        self.orig_dtype = orig_dtype
        self.input_scale = None
        if weight_scale_inv is not None:
            self.register_parameter(
                "weight_scale_inv", nn.Parameter(weight_scale_inv, requires_grad=False)
            )
            self.weight_scale_inv.format_ue8m0 = True
            # Must derive the backend scale layout here: the loader skips this
            # module (see ``_qm``), so it won't run process_weights_after_loading.
            quant_method._process_mxfp8_linear_weight_scale(self)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self._qm.apply(self, x, None)


def build_minimax_fused_qkv_index(model: nn.Module) -> None:
    for module in model.modules():
        if isinstance(module, MiniMaxM3Attention):
            module.maybe_build_fused_qkv_index()


class MiniMaxM3MLP(nn.Module):
    def __init__(
        self,
        config: PretrainedConfig,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
        reduce_results: bool = True,
        intermediate_size: int = None,
        tp_rank: Optional[int] = None,
        tp_size: Optional[int] = None,
    ) -> None:
        super().__init__()
        hidden_size = config.hidden_size
        hidden_act = config.hidden_act

        self.gate_up_proj = MergedColumnParallelLinear(
            hidden_size,
            [intermediate_size] * 2,
            bias=False,
            quant_config=quant_config,
            prefix=add_prefix("gate_up_proj", prefix),
            tp_rank=tp_rank,
            tp_size=tp_size,
        )
        self.down_proj = RowParallelLinear(
            intermediate_size,
            hidden_size,
            bias=False,
            quant_config=quant_config,
            reduce_results=reduce_results,
            prefix=add_prefix("down_proj", prefix),
            tp_rank=tp_rank,
            tp_size=tp_size,
        )
        if hidden_act == "silu":
            self.act_fn = SiluAndMul()
        elif hidden_act == "swigluoai":
            from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import (
                swiglu_no_interleaved_with_alpha_and_limit,
            )

            self.act_fn = lambda x: swiglu_no_interleaved_with_alpha_and_limit(
                x, config.swiglu_alpha, config.swiglu_limit
            )
        else:
            raise ValueError(
                f"Unsupported activation: {hidden_act}. Only silu is supported for now."
            )

    def forward(
        self,
        x,
        should_allreduce_fusion: bool = False,
        use_reduce_scatter: bool = False,
    ):
        gate_up, _ = self.gate_up_proj(x)
        x = self.act_fn(gate_up)
        x, _ = self.down_proj(
            x,
            skip_all_reduce=should_allreduce_fusion or use_reduce_scatter,
        )
        return x


class MiniMaxM3MoE(nn.Module):
    """MiniMax MoE implementation using DeepEP for Expert Parallel support."""

    def __init__(
        self,
        config: PretrainedConfig,
        layer_id: int,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ):
        super().__init__()
        self.tp_size = get_parallel().tp_size
        self.n_shared_experts = getattr(config, "n_shared_experts", None)
        self.num_fused_shared_experts = (
            0
            if get_server_args().disable_shared_experts_fusion
            else config.n_shared_experts
        )

        if self.tp_size > config.num_local_experts:
            raise ValueError(
                f"Tensor parallel size {self.tp_size} is greater than "
                f"the number of experts {config.num_local_experts}."
            )
        self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
        self.use_routing_bias = getattr(config, "use_routing_bias", False)
        if self.use_routing_bias:
            self.e_score_correction_bias = nn.Parameter(
                torch.empty(config.num_local_experts, dtype=torch.float32)
            )
            self.e_score_correction_bias.weight_loader = (
                MiniMaxM3MoE.ebias_weight_loader
            )
        else:
            self.e_score_correction_bias = None

        self.experts = get_moe_impl_class(quant_config)(
            num_experts=config.num_local_experts
            + self.num_fused_shared_experts
            + get_server_args().ep_num_redundant_experts,
            num_fused_shared_experts=self.num_fused_shared_experts,
            top_k=config.num_experts_per_tok + self.num_fused_shared_experts,
            hidden_size=config.hidden_size,
            intermediate_size=config.intermediate_size,
            layer_id=layer_id,
            quant_config=quant_config,
            activation="silu",
            is_gated=True,
            gemm1_alpha=config.swiglu_alpha,
            gemm1_clamp_limit=config.swiglu_limit,
            prefix=add_prefix("experts", prefix),
            gate_up_interleaved=False,
        )
        self.topk = TopK(
            top_k=config.num_experts_per_tok + self.num_fused_shared_experts,
            renormalize=True,
            layer_id=layer_id,
            scoring_func=config.scoring_func,
            correction_bias=self.e_score_correction_bias,
            num_fused_shared_experts=self.num_fused_shared_experts,
            routed_scaling_factor=self.routed_scaling_factor,
            apply_routed_scaling_factor_on_output=True,
        )

        if self.n_shared_experts is not None and self.num_fused_shared_experts == 0:
            intermediate_size = config.intermediate_size * self.n_shared_experts
            # DeepEP all-gathers (not all-reduces) the layer output, so a TP-sharded
            # shared MLP would leave an unreduced partial; replicate (tp_size=1), like GLM4 / DSV2.
            shared_experts_tp1 = get_moe_a2a_backend().is_deepep()
            self.shared_experts = MiniMaxM3MLP(
                config=config,
                quant_config=quant_config,
                prefix=add_prefix("shared_experts", prefix),
                reduce_results=False,
                intermediate_size=intermediate_size,
                **(dict(tp_rank=0, tp_size=1) if shared_experts_tp1 else {}),
            )
        else:
            self.shared_experts = None

        self.bf16_router_gemm = envs.SGLANG_OPT_USE_BF16_ROUTER_GEMM.get()
        self.gate = ReplicatedLinear(
            config.hidden_size,
            config.num_local_experts,
            bias=False,
            params_dtype=torch.bfloat16 if self.bf16_router_gemm else torch.float32,
            quant_config=None,
            prefix=add_prefix("gate", prefix),
        )

        self.layer_id = layer_id

        if get_moe_a2a_backend().is_deepep():
            self.ep_size = get_parallel().moe_ep_size
            self.top_k = config.num_experts_per_tok

    @staticmethod
    def ebias_weight_loader(param: nn.Parameter, loaded_weight: torch.Tensor) -> None:
        assert param.size() == loaded_weight.size()
        param.data.copy_(loaded_weight.to(torch.float32))

    def forward(
        self,
        hidden_states: torch.Tensor,
        forward_batch: ForwardBatch,
        should_allreduce_fusion: bool = False,
        use_reduce_scatter: bool = False,
    ) -> torch.Tensor:
        if get_moe_a2a_backend().is_deepep():
            return self.forward_deepep(hidden_states, forward_batch)
        else:
            return self.forward_normal(
                hidden_states, should_allreduce_fusion, use_reduce_scatter
            )

    def forward_normal(
        self,
        hidden_states: torch.Tensor,
        should_allreduce_fusion: bool = False,
        use_reduce_scatter: bool = False,
    ) -> torch.Tensor:
        if hidden_states.shape[0] > 0:
            shared_output = self._forward_shared_experts(hidden_states)
            router_logits = self._compute_router_logits(hidden_states)
            topk_output = self.topk(hidden_states, router_logits)
        else:
            shared_output = None
            topk_output = self.topk.empty_topk_output(hidden_states.device)

        final_hidden_states = self.experts(hidden_states, topk_output)

        if shared_output is not None:
            final_hidden_states = final_hidden_states + shared_output
        if self.tp_size > 1 and not should_allreduce_fusion and not use_reduce_scatter:
            final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)

        return final_hidden_states

    def forward_deepep(
        self, hidden_states: torch.Tensor, forward_batch: ForwardBatch
    ) -> torch.Tensor:
        shared_output = None
        if hidden_states.shape[0] > 0:
            shared_output = self._forward_shared_experts(hidden_states)
            router_logits = self._compute_router_logits(hidden_states)
            topk_output = self.topk(
                hidden_states,
                router_logits,
                num_token_non_padded=forward_batch.num_token_non_padded,
                expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new(
                    layer_id=self.layer_id,
                ),
            )
        else:
            topk_output = self.topk.empty_topk_output(hidden_states.device)

        # DeepEP returns the complete per-token routed result (no TP all-reduce here);
        # shared experts are replicated (tp_size=1), so both add directly.
        final_hidden_states = self.experts(hidden_states, topk_output)

        if shared_output is not None:
            final_hidden_states = final_hidden_states + shared_output

        return final_hidden_states

    def _compute_router_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
        if self.bf16_router_gemm:
            return torch.mm(
                hidden_states, self.gate.weight.t(), out_dtype=torch.float32
            )
        router_logits, _ = self.gate(hidden_states.to(torch.float32))
        return router_logits

    def _forward_shared_experts(self, hidden_states: torch.Tensor):
        if (hidden_states.shape[0] > 0) and (self.num_fused_shared_experts == 0):
            return self.shared_experts(hidden_states)
        else:
            return None


class MiniMaxM3Attention(nn.Module):
    def __init__(
        self,
        config: PretrainedConfig,
        layer_id: int = 0,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
        is_sparse_attention_layer: bool = False,
        disable_index_value: bool = False,
    ) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
        self.is_sparse_attention_layer = is_sparse_attention_layer
        self.disable_index_value = is_sparse_attention_layer and disable_index_value

        attn_tp_rank = get_parallel().attn_tp_rank
        attn_tp_size = get_parallel().attn_tp_size
        self.attn_tp_size = attn_tp_size
        self.attn_tp_rank = attn_tp_rank

        self.total_num_heads = config.num_attention_heads
        assert self.total_num_heads % attn_tp_size == 0
        self.num_heads = self.total_num_heads // attn_tp_size
        self.total_num_kv_heads = config.num_key_value_heads

        if self.total_num_kv_heads >= attn_tp_size:
            # Number of KV heads is greater than TP size, so we partition
            # the KV heads across multiple tensor parallel GPUs.
            assert self.total_num_kv_heads % attn_tp_size == 0
        else:
            # Number of KV heads is less than TP size, so we replicate
            # the KV heads across multiple tensor parallel GPUs.
            assert attn_tp_size % self.total_num_kv_heads == 0
        self.num_kv_heads = max(1, self.total_num_kv_heads // attn_tp_size)

        self.head_dim = getattr(
            config, "head_dim", self.hidden_size // self.total_num_heads
        )
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim
        self.scaling = self.head_dim**-0.5

        self.rope_theta, self.rope_scaling = get_rope_config(config)
        self.max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
        self.rotary_dim = getattr(config, "rotary_dim", self.head_dim)

        self.qk_norm_type = getattr(config, "qk_norm_type", "per_layer")
        self.use_gemma_norm = getattr(config, "use_gemma_norm", False)

        if self.is_sparse_attention_layer:
            assert self.qk_norm_type == "per_head", (
                f"sparse attention only supports qk_norm_type='per_head', "
                f"got {self.qk_norm_type!r}"
            )
            sparse_cfg = config.sparse_attention_config
            self.total_idx_heads = sparse_cfg["sparse_num_index_heads"]
            self.idx_head_dim = sparse_cfg["sparse_index_dim"]
            # idx_replica_size ranks share one idx head; pre-divide idx_o on the activation
            # (not the weight) so the TP all-reduce sums right and stays FP8-quant-safe.
            if self.total_idx_heads >= attn_tp_size:
                assert self.total_idx_heads % attn_tp_size == 0
            else:
                assert attn_tp_size % self.total_idx_heads == 0
            self.idx_head_tp_size = min(attn_tp_size, self.total_idx_heads)
            self.idx_replica_size = attn_tp_size // self.idx_head_tp_size
            self.idx_head_rank = attn_tp_rank // self.idx_replica_size
            self.num_idx_heads = self.total_idx_heads // self.idx_head_tp_size

        self.qkv_proj = QKVParallelLinear(
            self.hidden_size,
            self.head_dim,
            self.total_num_heads,
            self.total_num_kv_heads,
            bias=False,
            quant_config=quant_config,
            tp_rank=attn_tp_rank,
            tp_size=attn_tp_size,
            prefix=add_prefix("qkv_proj", prefix),
        )

        self.o_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            self.hidden_size,
            bias=False,
            reduce_results=False,
            quant_config=quant_config,
            tp_rank=attn_tp_rank,
            tp_size=attn_tp_size,
            prefix=add_prefix("o_proj", prefix),
        )

        self.rotary_emb = get_rope(
            self.head_dim,
            rotary_dim=self.rotary_dim,
            max_position=self.max_position_embeddings,
            base=self.rope_theta,
            rope_scaling=self.rope_scaling,
        )

        if self.is_sparse_attention_layer:
            self.index_qkv_proj = QKVParallelLinear(
                self.hidden_size,
                self.idx_head_dim,
                self.total_idx_heads,
                total_num_kv_heads=1,
                bias=False,
                quant_config=quant_config,
                v_head_size=(0 if self.disable_index_value else self.idx_head_dim),
                tp_rank=self.idx_head_rank,
                tp_size=self.idx_head_tp_size,
                prefix=add_prefix("index_qkv_proj", prefix),
            )

            if self.disable_index_value:
                self.index_o_proj = None
            else:
                self.index_o_proj = RowParallelLinear(
                    self.total_idx_heads * self.idx_head_dim,
                    self.hidden_size,
                    bias=False,
                    input_is_parallel=True,
                    reduce_results=False,
                    quant_config=quant_config,
                    prefix=add_prefix("index_o_proj", prefix),
                    tp_rank=self.idx_head_rank,
                    tp_size=self.idx_head_tp_size,
                )
            self.index_rotary_emb = self.rotary_emb

        if self.qk_norm_type == "per_layer":
            if attn_tp_size > 1:
                self.q_norm = MiniMaxM2RMSNormTP(
                    self.total_num_heads * self.head_dim, eps=config.rms_norm_eps
                )
                self.k_norm = MiniMaxM2RMSNormTP(
                    self.total_num_kv_heads * self.head_dim, eps=config.rms_norm_eps
                )
            else:
                self.q_norm = RMSNorm(
                    self.total_num_heads * self.head_dim, eps=config.rms_norm_eps
                )
                self.k_norm = RMSNorm(
                    self.total_num_kv_heads * self.head_dim, eps=config.rms_norm_eps
                )
        elif self.qk_norm_type == "per_head":
            if self.use_gemma_norm:
                self.q_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
                self.k_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
            else:
                self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
                self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
            if self.is_sparse_attention_layer:
                if self.use_gemma_norm:
                    self.index_q_norm = GemmaRMSNorm(
                        self.idx_head_dim, eps=config.rms_norm_eps
                    )
                    self.index_k_norm = GemmaRMSNorm(
                        self.idx_head_dim, eps=config.rms_norm_eps
                    )
                else:
                    self.index_q_norm = RMSNorm(
                        self.idx_head_dim, eps=config.rms_norm_eps
                    )
                    self.index_k_norm = RMSNorm(
                        self.idx_head_dim, eps=config.rms_norm_eps
                    )
        elif self.qk_norm_type == "multi_head":
            self.q_norm = MultiHeadRMSNorm(
                self.total_num_heads,
                self.head_dim,
                eps=config.rms_norm_eps,
                apply_layernorm_1p=self.use_gemma_norm,
            )
            self.k_norm = MultiHeadRMSNorm(
                self.total_num_kv_heads,
                self.head_dim,
                eps=config.rms_norm_eps,
                apply_layernorm_1p=self.use_gemma_norm,
            )
        else:
            raise ValueError(f"Invalid qk_norm_type: {self.qk_norm_type}")

        self.attn = RadixAttention(
            self.num_heads,
            self.head_dim,
            self.scaling,
            num_kv_heads=self.num_kv_heads,
            layer_id=layer_id,
            quant_config=quant_config,
            prefix=add_prefix("attn", prefix),
        )

        self._use_fused_qknorm_rope = (
            _is_cuda
            and self.qk_norm_type == "per_head"
            and self.use_gemma_norm
            and self.head_dim == 128
            and self.rotary_dim == _M3_FUSED_QKNORM_ROPE_ROTARY_DIM
            and getattr(self.rotary_emb, "is_neox_style", False)
        )

        self._fuse_qkv_index_enabled = self.is_sparse_attention_layer and (
            _is_cuda or _is_hip
        )
        self._fused_qkv_index = None
        self._fused_main_size = self.q_size + 2 * self.kv_size

        self._combined_qknorm_ok = (
            self.is_sparse_attention_layer
            and self._use_fused_qknorm_rope
            and self.idx_head_dim == 128
            and self.rotary_emb.cos_sin_cache.dtype == torch.float32
        )
        if self.is_sparse_attention_layer:
            off_iq = self.num_heads + 2 * self.num_kv_heads
            off_ik = off_iq + self.num_idx_heads
            self._qknorm_group_meta = (
                (0, self.num_heads),
                (self.num_heads, self.num_kv_heads),
                (off_iq, self.num_idx_heads),
                (off_ik, 1),
            )

        self._can_use_rocm_qk_norm_rope_static = (
            _has_rocm_qk_norm_rope
            and self.qk_norm_type == "per_head"
            and self.use_gemma_norm
            and self.q_norm.variance_epsilon == self.k_norm.variance_epsilon
            and hasattr(self.rotary_emb, "cos_sin_cache")
            and self.rotary_emb.rotary_dim == self.rotary_dim
            and self.rotary_dim <= self.head_dim
        )
        self._can_use_rocm_index_qk_norm_rope_static = (
            self.is_sparse_attention_layer
            and _has_rocm_qk_norm_rope
            and self.use_gemma_norm
            and self.index_q_norm.variance_epsilon == self.index_k_norm.variance_epsilon
            and hasattr(self.index_rotary_emb, "cos_sin_cache")
            and self.index_rotary_emb.rotary_dim == self.rotary_dim
            and self.rotary_dim <= self.idx_head_dim
        )
        self._can_use_rocm_sparse_qk_index_norm_rope_static = (
            self.is_sparse_attention_layer
            and self._can_use_rocm_qk_norm_rope_static
            and self._can_use_rocm_index_qk_norm_rope_static
            and self.idx_head_dim == self.head_dim
            and self.index_q_norm.variance_epsilon == self.q_norm.variance_epsilon
            and self.index_k_norm.variance_epsilon == self.q_norm.variance_epsilon
            and self.index_rotary_emb is self.rotary_emb
        )

    def _can_use_rocm_qk_norm_rope(
        self, positions: torch.Tensor, q: torch.Tensor, k: torch.Tensor
    ) -> bool:
        return (
            self._can_use_rocm_qk_norm_rope_static
            and positions.dim() == 1
            and q.dim() == 2
            and k.dim() == 2
            and q.dtype in (torch.bfloat16, torch.float16)
            and k.dtype == q.dtype
        )

    def _qk_norm_rope(
        self, positions: torch.Tensor, q: torch.Tensor, k: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        if self._can_use_rocm_qk_norm_rope(positions, q, k):
            return qk_gemma_rmsnorm_rope(
                q,
                k,
                self.q_norm.weight.data,
                self.k_norm.weight.data,
                positions,
                self.rotary_emb.cos_sin_cache,
                self.q_norm.variance_epsilon,
                self.head_dim,
                self.rotary_dim,
                self.rotary_emb.is_neox_style,
            )
        q, k = self._qk_norm(q, k)
        return self.rotary_emb(positions, q, k)

    def _qk_norm(
        self, q: torch.Tensor, k: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        if self.qk_norm_type == "per_layer":
            if self.attn_tp_size > 1:
                q, k = MiniMaxM2RMSNormTP.forward_qk(
                    self.q_norm, self.k_norm, q.contiguous(), k.contiguous()
                )
            else:
                q = self.q_norm(q.contiguous())
                k = self.k_norm(k.contiguous())
        elif self.qk_norm_type == "per_head":
            q_shape = q.shape
            k_shape = k.shape
            q = q.reshape(-1, self.head_dim).contiguous()
            k = k.reshape(-1, self.head_dim).contiguous()
            q = self.q_norm(q).reshape(q_shape)
            k = self.k_norm(k).reshape(k_shape)
        elif self.qk_norm_type == "multi_head":
            q = self.q_norm(q.contiguous())
            k = self.k_norm(k.contiguous())
        else:
            raise ValueError(f"Invalid qk_norm_type: {self.qk_norm_type}")
        return q, k

    def _index_qk_norm_rope(
        self, positions: torch.Tensor, idx_q: torch.Tensor, idx_k: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        if (
            self._can_use_rocm_index_qk_norm_rope_static
            and positions.dim() == 1
            and idx_q.dim() == 2
            and idx_k.dim() == 2
            and idx_q.dtype in (torch.bfloat16, torch.float16)
            and idx_k.dtype == idx_q.dtype
        ):
            return qk_gemma_rmsnorm_rope(
                idx_q,
                idx_k,
                self.index_q_norm.weight.data,
                self.index_k_norm.weight.data,
                positions,
                self.index_rotary_emb.cos_sin_cache,
                self.index_q_norm.variance_epsilon,
                self.idx_head_dim,
                self.rotary_dim,
                self.index_rotary_emb.is_neox_style,
            )
        idx_q, idx_k = self._index_qk_norm(idx_q, idx_k)
        return self.index_rotary_emb(positions, idx_q, idx_k)

    def _index_qk_norm(
        self, idx_q: torch.Tensor, idx_k: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        idx_q_shape = idx_q.shape
        idx_k_shape = idx_k.shape
        idx_q = idx_q.reshape(-1, self.idx_head_dim)
        idx_k = idx_k.reshape(-1, self.idx_head_dim)
        idx_q = self.index_q_norm(idx_q).reshape(idx_q_shape)
        idx_k = self.index_k_norm(idx_k).reshape(idx_k_shape)
        return idx_q, idx_k

    def _split_index_qkv(
        self, idx_qkv: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
        q_size = self.num_idx_heads * self.idx_head_dim
        if self.disable_index_value:
            idx_q, idx_k = idx_qkv.split([q_size, self.idx_head_dim], dim=-1)
            idx_v = None
        else:
            idx_q, idx_k, idx_v = idx_qkv.split(
                [q_size, self.idx_head_dim, self.idx_head_dim], dim=-1
            )
        return idx_q, idx_k, idx_v

    def maybe_build_fused_qkv_index(self) -> None:
        if not self._fuse_qkv_index_enabled or self._fused_qkv_index is not None:
            return
        from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod

        qp, ip = self.qkv_proj, self.index_qkv_proj
        qm = qp.quant_method
        if type(ip.quant_method) is not type(qm):
            return

        # gfx942 converts MXFP8->block-fp8 in process_weights_after_loading; the
        # fused module skips that pass, so keep two separate (converted) GEMMs.
        if getattr(qm, "convert_mxfp8_to_block", False):
            return

        is_unquant = isinstance(qm, UnquantizedLinearMethod)
        use_mxfp8 = getattr(qm, "use_mxfp8", False) and hasattr(qp, "weight_scale_inv")
        if not (is_unquant or use_mxfp8):
            return

        weight = torch.cat([qp.weight.data, ip.weight.data], dim=0).contiguous()
        if is_unquant:
            scale = None
        else:
            scale = torch.cat(
                [qp.weight_scale_inv.data, ip.weight_scale_inv.data], dim=0
            ).contiguous()

        holder = _FusedQKVIndexProj(
            qm,
            weight,
            scale,
            getattr(qp, "input_size_per_partition", qp.input_size),
            [qp.output_size_per_partition, ip.output_size_per_partition],
            getattr(qp, "orig_dtype", qp.params_dtype),
        )
        self.add_module("fused_qkv_index_proj", holder)
        self._fused_qkv_index = holder

        # Free the dead originals and drop their quant_method so the loader's
        # post-process loop ignores them (see ``_qm``).
        for m in (qp, ip):
            m.quant_method = None
            for attr in ("weight", "weight_scale_inv"):
                p = getattr(m, attr, None)
                if isinstance(p, nn.Parameter):
                    p.data = torch.empty(0, dtype=p.dtype, device=p.data.device)

    def _qknorm_groups(self):
        weights = (
            self.q_norm.weight,
            self.k_norm.weight,
            self.index_q_norm.weight,
            self.index_k_norm.weight,
        )
        return [
            (w, off, cnt) for w, (off, cnt) in zip(weights, self._qknorm_group_meta)
        ]

    def _can_use_rocm_sparse_qk_index_norm_rope(
        self,
        positions: torch.Tensor,
        q: torch.Tensor,
        k: torch.Tensor,
        idx_q: torch.Tensor,
        idx_k: torch.Tensor,
    ) -> bool:
        return (
            self._can_use_rocm_sparse_qk_index_norm_rope_static
            and positions.dim() == 1
            and q.dim() == 2
            and k.dim() == 2
            and idx_q.dim() == 2
            and idx_k.dim() == 2
            and q.dtype in (torch.bfloat16, torch.float16)
            and k.dtype == q.dtype
            and idx_q.dtype == q.dtype
            and idx_k.dtype == q.dtype
        )

    def _sparse_qk_index_norm_rope(
        self,
        positions: torch.Tensor,
        q: torch.Tensor,
        k: torch.Tensor,
        idx_q: torch.Tensor,
        idx_k: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        if self._can_use_rocm_sparse_qk_index_norm_rope(positions, q, k, idx_q, idx_k):
            return sparse_qk_index_gemma_rmsnorm_rope(
                q,
                k,
                idx_q,
                idx_k,
                self.q_norm.weight.data,
                self.k_norm.weight.data,
                self.index_q_norm.weight.data,
                self.index_k_norm.weight.data,
                positions,
                self.rotary_emb.cos_sin_cache,
                self.q_norm.variance_epsilon,
                self.head_dim,
                self.rotary_dim,
                self.rotary_emb.is_neox_style,
            )
        q, k = self._qk_norm_rope(positions, q, k)
        idx_q, idx_k = self._index_qk_norm_rope(positions, idx_q, idx_k)
        return q, k, idx_q, idx_k

    @staticmethod
    def _mark_sparse_kv_cached_by_fusion(
        forward_batch: ForwardBatch, layer_id: int
    ) -> None:
        layer_ids = forward_batch.minimax_m3_precached_sparse_layers
        if layer_ids is None:
            layer_ids = set()
            forward_batch.minimax_m3_precached_sparse_layers = layer_ids
        layer_ids.add(layer_id)

    @staticmethod
    def _get_sparse_kv_pool():
        if not has_forward_context():
            return None
        attn_backend = get_forward_context().attn_backend
        sparse_backend = getattr(attn_backend, "sparse", None)
        return getattr(sparse_backend, "kv_pool", None)

    def _sparse_qk_index_norm_rope_cache(
        self,
        positions: torch.Tensor,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        idx_q: torch.Tensor,
        idx_k: torch.Tensor,
        idx_v: Optional[torch.Tensor],
        forward_batch: ForwardBatch,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        kv_pool = self._get_sparse_kv_pool()
        # The fused kernel writes normed bf16 K/V straight into the paged cache, so an
        # fp8 main K/V cache (--kv-cache-dtype fp8_*) can't use it; fall back to norm+rope.
        main_kv_is_fp8 = kv_pool is not None and kv_pool.dtype in _FP8_KV_DTYPES
        can_use_cache_fusion = (
            not main_kv_is_fp8
            and idx_v is None
            and self._can_use_rocm_sparse_qk_index_norm_rope(
                positions, q, k, idx_q, idx_k
            )
            and getattr(forward_batch, "out_cache_loc", None) is not None
            and v.dim() == 2
            and v.dtype == q.dtype
            and v.shape == k.shape
        )
        if can_use_cache_fusion and kv_pool is not None:
            layer_id = self.attn.layer_id
            k_cache, v_cache = kv_pool.get_kv_buffer(layer_id)
            idx_k_cache = kv_pool.get_index_k_buffer(layer_id)
            q, k, idx_q, idx_k = sparse_qk_index_gemma_rmsnorm_rope_cache(
                q,
                k,
                v,
                idx_q,
                idx_k,
                k_cache,
                v_cache,
                idx_k_cache,
                forward_batch.out_cache_loc,
                self.q_norm.weight.data,
                self.k_norm.weight.data,
                self.index_q_norm.weight.data,
                self.index_k_norm.weight.data,
                positions,
                self.rotary_emb.cos_sin_cache,
                self.q_norm.variance_epsilon,
                self.head_dim,
                self.rotary_dim,
                self.rotary_emb.is_neox_style,
            )
            self._mark_sparse_kv_cached_by_fusion(forward_batch, layer_id)
            return q, k, idx_q, idx_k
        return self._sparse_qk_index_norm_rope(positions, q, k, idx_q, idx_k)

    def forward_prepare(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        forward_batch: ForwardBatch,
    ):
        fused_out = None
        if self._fused_qkv_index is not None:
            fused_out = self.fused_qkv_index_proj(hidden_states)
            qkv = fused_out[:, : self._fused_main_size]

            if self._combined_qknorm_ok:
                from sglang.jit_kernel.minimax_qknorm_rope import (
                    minimax_qknorm_rope_grouped,
                )

                minimax_qknorm_rope_grouped(
                    fused_out,
                    self._qknorm_groups(),
                    self.rotary_emb.cos_sin_cache,
                    positions,
                    self.q_norm.variance_epsilon,
                )
                q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
                idx_qkv = fused_out[:, self._fused_main_size :]
                idx_q, idx_k, idx_v = self._split_index_qkv(idx_qkv)
                inner_state = (q, k, v, idx_q, idx_k, idx_v, forward_batch)
                return None, forward_batch, inner_state
        else:
            qkv, _ = self.qkv_proj(hidden_states)

        if self._use_fused_qknorm_rope:
            from sglang.jit_kernel.minimax_qknorm_rope import minimax_qknorm_rope

            minimax_qknorm_rope(
                qkv,
                self.q_norm.weight,
                self.k_norm.weight,
                self.rotary_emb.cos_sin_cache,
                positions,
                self.num_heads,
                self.num_kv_heads,
                self.num_kv_heads,
                self.q_norm.variance_epsilon,
            )
            q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
            main_qk_already_normed = True
        else:
            q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
            main_qk_already_normed = False

        if self.is_sparse_attention_layer:
            if fused_out is not None:
                idx_qkv = fused_out[:, self._fused_main_size :]
            else:
                idx_qkv, _ = self.index_qkv_proj(hidden_states)

            if main_qk_already_normed:
                use_fused_index_norm_rope = (
                    self._use_fused_qknorm_rope
                    and self.idx_head_dim == 128
                    and self.index_rotary_emb.cos_sin_cache.dtype == torch.float32
                )
                if use_fused_index_norm_rope:
                    from sglang.jit_kernel.minimax_qknorm_rope import (
                        minimax_qknorm_rope,
                    )

                    minimax_qknorm_rope(
                        idx_qkv,
                        self.index_q_norm.weight,
                        self.index_k_norm.weight,
                        self.index_rotary_emb.cos_sin_cache,
                        positions,
                        self.num_idx_heads,
                        1,
                        0 if self.disable_index_value else 1,
                        self.index_q_norm.variance_epsilon,
                    )
                    idx_q, idx_k, idx_v = self._split_index_qkv(idx_qkv)
                else:
                    idx_q, idx_k, idx_v = self._split_index_qkv(idx_qkv)
                    idx_q, idx_k = self._index_qk_norm_rope(positions, idx_q, idx_k)
            else:
                idx_q, idx_k, idx_v = self._split_index_qkv(idx_qkv)
                q, k, idx_q, idx_k = self._sparse_qk_index_norm_rope_cache(
                    positions, q, k, v, idx_q, idx_k, idx_v, forward_batch
                )

            inner_state = (q, k, v, idx_q, idx_k, idx_v, forward_batch)
        else:
            if not main_qk_already_normed:
                q, k = self._qk_norm_rope(positions, q, k)
            inner_state = (q, k, v, forward_batch)
        return None, forward_batch, inner_state

    def forward_core(self, intermediate_state):
        _, _, inner_state = intermediate_state

        if self.is_sparse_attention_layer:
            q, k, v, idx_q, idx_k, idx_v, forward_batch = inner_state
            q = q.view(q.shape[0], self.num_heads, self.head_dim)
            k = k.view(k.shape[0], self.num_kv_heads, self.head_dim)
            v = v.view(v.shape[0], self.num_kv_heads, self.head_dim)
            idx_q = idx_q.reshape(idx_q.shape[0], self.num_idx_heads, self.idx_head_dim)
            idx_k = idx_k.reshape(idx_k.shape[0], 1, self.idx_head_dim)
            if idx_v is not None:
                idx_v = idx_v.reshape(idx_v.shape[0], 1, self.idx_head_dim)
            idx_o, attn_output = self.attn(
                q, k, v, forward_batch, idx_q=idx_q, idx_k=idx_k, idx_v=idx_v
            )
            output, _ = self.o_proj(attn_output)
            if self.disable_index_value:
                return output
            # idx_replica_size ranks produce identical idx_o; pre-divide idx_o (not the
            # o_proj weight) so the TP all-reduce sums right and stays FP8-quant-safe.
            if self.idx_replica_size > 1:
                idx_o = idx_o / self.idx_replica_size
            idx_output, _ = self.index_o_proj(idx_o)
            return output + idx_output

        q, k, v, forward_batch = inner_state
        attn_output = self.attn(q, k, v, forward_batch)
        output, _ = self.o_proj(attn_output)
        return output

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        forward_batch: ForwardBatch,
    ) -> torch.Tensor:
        s = self.forward_prepare(
            positions=positions,
            hidden_states=hidden_states,
            forward_batch=forward_batch,
        )
        return self.forward_core(s)


class MiniMaxM3DecoderLayer(nn.Module):
    def __init__(
        self,
        config: PretrainedConfig,
        layer_id: int,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
        self.layer_id = layer_id

        sparse_attention_config = getattr(config, "sparse_attention_config", None)
        if sparse_attention_config is not None:
            _, sparse_layer_ids = get_minimax_sparse_layer_ids(sparse_attention_config)
            is_sparse_attention_layer = layer_id in sparse_layer_ids
            disable_value_layer_ids = set(
                get_minimax_sparse_disable_value_layer_ids(sparse_attention_config)
            )
            disable_index_value = layer_id in disable_value_layer_ids
        else:
            is_sparse_attention_layer = False
            disable_index_value = False

        self.self_attn = MiniMaxM3Attention(
            config=config,
            layer_id=layer_id,
            quant_config=quant_config,
            prefix=add_prefix("self_attn", prefix),
            is_sparse_attention_layer=is_sparse_attention_layer,
            disable_index_value=disable_index_value,
        )

        moe_layer_freq = getattr(config, "moe_layer_freq", None)
        # Means "MLP is a sparse MoE", not attention sparsity. Kept as ``is_layer_sparse``
        # because LayerCommunicator / LayerScatterModes / other models read this attr.
        self.is_layer_sparse = (
            moe_layer_freq[layer_id] != 0 if moe_layer_freq is not None else True
        )

        if self.is_layer_sparse:
            self.mlp = MiniMaxM3MoE(
                config=config,
                layer_id=layer_id,
                quant_config=quant_config,
                prefix=add_prefix("mlp", prefix),
            )
        else:
            if enable_moe_dense_fully_dp():
                mlp_tp_rank, mlp_tp_size = 0, 1
            else:
                mlp_tp_rank, mlp_tp_size = None, None
            self.mlp = MiniMaxM3MLP(
                config=config,
                quant_config=quant_config,
                prefix=add_prefix("mlp", prefix),
                intermediate_size=config.dense_intermediate_size,
                tp_rank=mlp_tp_rank,
                tp_size=mlp_tp_size,
            )

        self.use_gemma_norm = getattr(config, "use_gemma_norm", False)
        if self.use_gemma_norm:
            self.input_layernorm = GemmaRMSNorm(
                config.hidden_size, eps=config.rms_norm_eps
            )
            self.post_attention_layernorm = GemmaRMSNorm(
                config.hidden_size, eps=config.rms_norm_eps
            )
        else:
            self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
            self.post_attention_layernorm = RMSNorm(
                config.hidden_size, eps=config.rms_norm_eps
            )

        def _is_layer_sparse(lid):
            if moe_layer_freq is None:
                return True
            if lid < 0 or lid >= config.num_hidden_layers:
                return True
            return moe_layer_freq[lid] != 0

        is_previous_layer_sparse = _is_layer_sparse(layer_id - 1)
        is_next_layer_sparse = _is_layer_sparse(layer_id + 1)
        self.layer_scatter_modes = LayerScatterModes.init_new(
            layer_id=layer_id,
            num_layers=config.num_hidden_layers,
            is_layer_sparse=self.is_layer_sparse,
            is_previous_layer_sparse=is_previous_layer_sparse,
            is_next_layer_sparse=is_next_layer_sparse,
        )

        self.layer_communicator = LayerCommunicator(
            layer_scatter_modes=self.layer_scatter_modes,
            input_layernorm=self.input_layernorm,
            post_attention_layernorm=self.post_attention_layernorm,
            allow_reduce_scatter=True,
        )

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        forward_batch: ForwardBatch,
        residual: Optional[torch.Tensor],
        captured_last_layer_outputs: Optional[List[torch.Tensor]] = None,
        **kwargs,
    ) -> torch.Tensor:
        hidden_states, residual = (
            self.layer_communicator.prepare_attn_and_capture_last_layer_outputs(
                hidden_states,
                residual,
                forward_batch,
                captured_last_layer_outputs=captured_last_layer_outputs,
                **kwargs,
            )
        )

        if hidden_states.shape[0] != 0:
            hidden_states = self.self_attn(
                positions=positions,
                hidden_states=hidden_states,
                forward_batch=forward_batch,
            )

        hidden_states, residual = self.layer_communicator.prepare_mlp(
            hidden_states, residual, forward_batch
        )

        should_allreduce_fusion = (
            self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
                forward_batch
            )
        )
        if self.is_layer_sparse and get_parallel().tp_size > 1:
            # Sparse MoE outputs are TP-partial; deferring their all-reduce into the next
            # layer's fusion re-triggers the M3 no-EOS runaway. Force immediate all-reduce.
            should_allreduce_fusion = False

        use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
            forward_batch
        )

        if self.is_layer_sparse or hidden_states.shape[0] != 0:
            hidden_states = self.mlp(
                hidden_states,
                should_allreduce_fusion,
                use_reduce_scatter,
            )

        if should_allreduce_fusion:
            hidden_states._sglang_needs_allreduce_fusion = True
        else:
            hidden_states, residual = self.layer_communicator.postprocess_layer(
                hidden_states, residual, forward_batch
            )

        return hidden_states, residual


class MiniMaxM3Model(nn.Module):
    """MiniMax Model implementation."""

    fall_back_to_pt_during_load = False

    def __init__(
        self,
        config: PretrainedConfig,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
        super().__init__()

        self.padding_idx = getattr(config, "pad_token_id", 0)
        self.vocab_size = config.vocab_size
        self.pp_group = get_pp_group()
        self.use_gemma_norm = getattr(config, "use_gemma_norm", False)

        if self.pp_group.is_first_rank:
            self.embed_tokens = VocabParallelEmbedding(
                config.vocab_size,
                config.hidden_size,
                use_attn_tp_group=is_dp_attention_enabled(),
                prefix=add_prefix("embed_tokens", prefix),
            )
        else:
            self.embed_tokens = PPMissingLayer()

        def layer_fn(idx, prefix: str) -> nn.Module:
            return MiniMaxM3DecoderLayer(
                config=config,
                layer_id=idx,
                quant_config=quant_config,
                prefix=prefix,
            )

        self.layers, self.start_layer, self.end_layer = make_layers(
            config.num_hidden_layers,
            layer_fn,
            pp_rank=self.pp_group.rank_in_group,
            pp_size=self.pp_group.world_size,
            prefix=add_prefix("layers", prefix),
        )
        if self.pp_group.is_last_rank:
            if self.use_gemma_norm:
                self.norm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
            else:
                self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        else:
            self.norm = PPMissingLayer(return_tuple=True)

        self.layers_to_capture = []

    def get_input_embeddings(self) -> torch.Tensor:
        return self.embed_tokens

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        forward_batch: ForwardBatch,
        input_embeds: torch.Tensor = None,
        pp_proxy_tensors: Optional[PPProxyTensors] = None,
    ) -> Union[torch.Tensor, PPProxyTensors, Tuple[torch.Tensor, list[torch.Tensor]]]:
        if self.pp_group.is_first_rank:
            if input_embeds is None:
                embeds = self.get_input_embeddings()
                hidden_states = embeds(input_ids)
            else:
                hidden_states = input_embeds
            residual = None
        else:
            assert pp_proxy_tensors is not None
            hidden_states = pp_proxy_tensors["hidden_states"]
            residual = pp_proxy_tensors["residual"]

        aux_hidden_states = []
        if forward_batch.can_run_tbo:
            hidden_states, residual = model_forward_maybe_tbo(
                layers=self.layers,
                enable_tbo=True,
                input_data_scatter_mode=ScatterMode.model_input_output(),
                positions=positions,
                forward_batch=forward_batch,
                hidden_states=hidden_states,
                residual=residual,
            )
        else:
            for i in range(self.start_layer, self.end_layer):
                # NOTE: torch dynamo does not support graph break in context manager
                ctx = (
                    nullcontext()
                    if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
                    else get_global_expert_distribution_recorder().with_current_layer(i)
                )
                with ctx:
                    layer = self.layers[i]
                    hidden_states, residual = layer(
                        positions=positions,
                        forward_batch=forward_batch,
                        hidden_states=hidden_states,
                        residual=residual,
                        captured_last_layer_outputs=(
                            aux_hidden_states
                            if getattr(layer, "_is_layer_to_capture", False)
                            else None
                        ),
                    )

        if not self.pp_group.is_last_rank:
            return PPProxyTensors(
                {"hidden_states": hidden_states, "residual": residual}
            )
        if hidden_states.shape[0] != 0:
            if residual is not None:
                hidden_states, _ = self.norm(hidden_states, residual)
            else:
                hidden_states = self.norm(hidden_states)

        if len(aux_hidden_states) == 0:
            return hidden_states
        return hidden_states, aux_hidden_states


class MiniMaxM3SparseForCausalLM(nn.Module):
    def __init__(
        self,
        config: PretrainedConfig,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
        super().__init__()

        self.config = config
        self.quant_config = quant_config
        self.pp_group = get_pp_group()

        self.num_fused_shared_experts = 0
        self.determine_num_fused_shared_experts()

        self.model = MiniMaxM3Model(
            config, quant_config, prefix=add_prefix("model", prefix)
        )

        if self.pp_group.is_last_rank:
            self.lm_head = ParallelLMHead(
                config.vocab_size,
                config.hidden_size,
                quant_config=quant_config,
                prefix=add_prefix("lm_head", prefix),
                use_attn_tp_group=get_server_args().enable_dp_lm_head,
            )

            self.logits_processor = LogitsProcessor(config)
        else:
            self.lm_head = PPMissingLayer()

        self.capture_aux_hidden_states = False

    def get_input_embeddings(self):
        return self.model.get_input_embeddings()

    def determine_num_fused_shared_experts(self):
        if get_server_args().disable_shared_experts_fusion:
            return

        disable_reason = None
        if not getattr(self.config, "n_shared_experts", None):
            disable_reason = "No shared experts are defined in the config."
        elif not _is_cuda:
            disable_reason = "Shared experts fusion currently requires CUDA devices."
        elif _is_cuda and (_device_sm is not None) and (_device_sm < 80):
            disable_reason = "Shared experts fusion requires SM80 or newer GPUs."
        elif get_parallel().moe_ep_size > 1:
            disable_reason = "Shared experts fusion is not supported together with expert parallelism yet."
        elif get_moe_a2a_backend().is_deepep():
            disable_reason = "Shared experts fusion is not supported when Deepep MoE backend is enabled."

        if disable_reason is not None:
            from sglang.srt.arg_groups.overrides import declare_load_time_override

            declare_load_time_override(
                "MiniMaxM3ForCausalLM.determine_num_fused_shared_experts",
                {"disable_shared_experts_fusion": True},
            )
            log_info_on_rank0(
                logger,
                f"{disable_reason} Shared experts fusion optimization is disabled.",
            )
            return

        self.num_fused_shared_experts = self.config.n_shared_experts
        assert (
            self.num_fused_shared_experts == 1
        ), "Only 1 fused shared expert is supported for MiniMax-M3"
        log_info_on_rank0(logger, "Shared experts fusion optimization enabled.")

    def set_eagle3_layers_to_capture(self, layer_ids: Optional[list[int]] = None):
        if not self.pp_group.is_last_rank:
            return

        self.capture_aux_hidden_states = True
        if layer_ids is None:
            num_layers = self.config.num_hidden_layers
            self.model.layers_to_capture = [
                2,
                num_layers // 2,
                num_layers - 3,
            ]
        else:
            self.model.layers_to_capture = [val + 1 for val in layer_ids]

    def get_embed_and_head(self):
        return self.model.embed_tokens.weight, self.lm_head.weight

    @torch.no_grad()
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        forward_batch: ForwardBatch,
        input_embeds: torch.Tensor = None,
        pp_proxy_tensors: Optional[PPProxyTensors] = None,
    ) -> torch.Tensor:
        hidden_states = self.model(
            input_ids, positions, forward_batch, input_embeds, pp_proxy_tensors
        )

        aux_hidden_states = None
        if self.capture_aux_hidden_states:
            hidden_states, aux_hidden_states = hidden_states

        if self.pp_group.is_last_rank:
            return self.logits_processor(
                input_ids, hidden_states, self.lm_head, forward_batch, aux_hidden_states
            )
        else:
            return hidden_states

    @property
    def start_layer(self):
        return self.model.start_layer

    @property
    def end_layer(self):
        return self.model.end_layer

    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
        """Load model weights with proper mapping for MiniMax architecture."""

        stacked_params_mapping = [
            # Leading "." on ".qkv_proj" prevents it from falsely matching the sparse
            # index_q/k/v_proj weights (remapped separately below).
            (".qkv_proj", ".q_proj", "q"),
            (".qkv_proj", ".k_proj", "k"),
            (".qkv_proj", ".v_proj", "v"),
            (".gate_up_proj", ".gate_proj", 0),
            (".gate_up_proj", ".up_proj", 1),
        ]

        # Value-disabled layers have no ".index_v_proj" weight, so that entry never matches.
        if getattr(self.config, "sparse_attention_config", None) is not None:
            stacked_params_mapping += [
                (".index_qkv_proj", ".index_q_proj", "q"),
                (".index_qkv_proj", ".index_k_proj", "k"),
                (".index_qkv_proj", ".index_v_proj", "v"),
            ]

        expert_params_mapping = FusedMoE.make_expert_params_mapping(
            ckpt_gate_proj_name="w1",
            ckpt_down_proj_name="w2",
            ckpt_up_proj_name="w3",
            num_experts=self.config.num_local_experts + self.num_fused_shared_experts,
        )

        params_dict = dict(self.named_parameters())
        loaded_params: Set[str] = set()
        for name, loaded_weight in weights:
            layer_id = get_layer_id(name)
            if layer_id is not None and (
                layer_id < self.model.start_layer or layer_id >= self.model.end_layer
            ):
                continue

            name = name.replace(".block_sparse_moe", ".mlp")

            if self.num_fused_shared_experts > 0 and "mlp.shared_experts" in name:
                name = name.replace(
                    "mlp.shared_experts",
                    f"mlp.experts.{self.config.num_local_experts}",
                )
                name = name.replace("gate_proj", "w1")
                name = name.replace("down_proj", "w2")
                name = name.replace("up_proj", "w3")

            if "rotary_emb.inv_freq" in name:
                continue

            spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
            if spec_layer is not None:
                continue

            for param_name, weight_name, shard_id in stacked_params_mapping:
                if weight_name not in name:
                    continue
                # Must skip experts before the name.replace below, else gate_proj ->
                # gate_up_proj -> gate_gate_up_proj double-remap breaks load.
                if "mlp.experts." in name:
                    continue
                name = name.replace(weight_name, param_name)
                if name.endswith(".bias") and name not in params_dict:
                    continue
                if name not in params_dict:
                    continue

                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
                is_expert_weight = False

                for mapping in expert_params_mapping:
                    param_name, weight_name, expert_id, shard_id = mapping
                    if weight_name not in name:
                        continue

                    is_expert_weight = True

                    name = name.replace(weight_name, param_name)
                    if name not in params_dict:
                        continue

                    param = params_dict[name]
                    weight_loader = param.weight_loader
                    weight_loader(
                        param,
                        loaded_weight,
                        name,
                        shard_id=shard_id,
                        expert_id=expert_id,
                    )
                    break
                else:
                    if is_expert_weight:
                        continue

                    if name.endswith(".bias") and name not in params_dict:
                        continue

                    name = maybe_remap_kv_scale_name(name, params_dict)
                    if name is None:
                        continue

                    if name in params_dict:
                        param = params_dict[name]
                        weight_loader = getattr(
                            param, "weight_loader", default_weight_loader
                        )
                        try:
                            weight_loader(param, loaded_weight)
                        except Exception as e:
                            logger.warning(f"Error loading weight {name}: {e}")
                            continue
                    else:
                        logger.warning(f"Parameter {name} not found in params_dict")
            loaded_params.add(name)

        # Run before the loader's process pass: the raw fp8 weight + uint8 scale are
        # final here (mxfp8 post-process only derives the packed scale, not these).
        build_minimax_fused_qkv_index(self)
        return loaded_params

    @classmethod
    def get_model_config_for_expert_location(cls, config):
        from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation

        return ModelConfigForExpertLocation(
            num_layers=config.num_hidden_layers,
            num_logical_experts=config.num_local_experts,
            num_groups=None,
        )


def get_spec_layer_idx_from_weight_name(
    config: PretrainedConfig, weight_name: str
) -> Optional[int]:
    # M3 checkpoints emit MTP weights as model.mtp.layers.{i}.*; skip them here
    # (no NextN module is built for the main model).
    if hasattr(config, "num_mtp_modules") and (config.num_mtp_modules > 0):
        for i in range(config.num_mtp_modules):
            if weight_name.startswith(f"model.mtp.layers.{i}."):
                return config.num_hidden_layers + i
    return None


EntryClass = [MiniMaxM3SparseForCausalLM]