deepseek_v2.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# 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:
# https://github.com/vllm-project/vllm/blob/fb6af8bc086328ca6659e72d11ffd4309ce4de22/vllm/model_executor/models/deepseek_v2.py
"""Inference-only DeepseekV2 model."""

from __future__ import annotations

import logging
from contextlib import contextmanager, nullcontext
from typing import Any, Dict, Iterable, List, Optional, Tuple, Union

import torch
import torch.nn.functional as F
from torch import nn
from transformers import PretrainedConfig

from sglang.kernels.ops.attention.dsv4 import (
    silu_and_mul_clamp,
    silu_and_mul_contig_post_quant,
)
from sglang.kernels.ops.quantization.fp8_kernel import (
    create_per_token_group_quant_fp8_output_scale,
)
from sglang.srt.batch_overlap.single_batch_overlap import SboFlags, compute_overlap_args
from sglang.srt.batch_overlap.two_batch_overlap import (
    MaybeTboDeepEPDispatcher,
    model_forward_maybe_tbo,
)
from sglang.srt.configs.model_config import (
    compute_mla_mscale_scaling,
    dsa_layer_skips_topk,
    get_dsa_index_head_dim,
    get_dsa_index_kpool,
    get_dsa_index_n_heads,
    get_dsa_index_topk,
    is_deepseek_dsa,
    is_glm_moe_dsa,
)
from sglang.srt.distributed import (
    divide,
    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 import ModelConfigForExpertLocation
from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.activation import SiluAndMul
from sglang.srt.layers.amx_utils import PackWeightMethod
from sglang.srt.layers.attention.dsa.dsa_indexer import Indexer
from sglang.srt.layers.attention.dsa.dsa_indexer_kpool import IndexerKPool
from sglang.srt.layers.attention.index_topk_share import IndexTopKShareState
from sglang.srt.layers.aux_hidden_states import (
    AuxHiddenStateAccumulator,
    AuxHiddenStatePacker,
)
from sglang.srt.layers.communicator import (
    LayerCommunicator,
    LayerScatterModes,
    enable_moe_dense_fully_dp,
    get_attn_tp_context,
)
from sglang.srt.layers.communicator_dsa_cp import (
    DSACPLayerCommunicator,
    maybe_prefetch_next_full_attention_kv,
)
from sglang.srt.layers.cp.cp_decode_attn_tp import get_cp_decode_attn_tp_ctx
from sglang.srt.layers.dcp.planner import (
    prepare_decode_context_parallel_metadata,
)
from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.linear import (
    ColumnParallelLinear,
    MergedColumnParallelLinear,
    ReplicatedLinear,
    RowParallelLinear,
)
from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.moe import (
    get_moe_a2a_backend,
    get_moe_runner_backend,
    should_skip_post_experts_all_reduce,
    should_use_flashinfer_cutlass_moe_fp4_allgather,
)
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.hash_topk import HashTopK
from sglang.srt.layers.moe.kt_ep_wrapper import KTEPWrapperMethod
from sglang.srt.layers.moe.token_dispatcher.base import (
    BaseDispatcher,
    CombineInput,
    DispatchOutput,
)
from sglang.srt.layers.moe.topk import BypassedTopKOutput, TopK, TopKOutputFormat
from sglang.srt.layers.moe.utils import (
    RoutingMethodType,
    filter_moe_weight_param_global_expert,
    has_per_rank_fused_shared_slots,
    is_deepep_class_backend,
    is_sbo_enabled,
    is_shared_experts_fusion_disabled,
    is_tbo_enabled,
)
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.quantization.fp8 import Fp8Config
from sglang.srt.layers.quantization.fp8_utils import (
    emit_transposed_bpreshuffle_scale,
    materialize_bpreshuffle_fp8_scale,
    view_aiter_fused_rms_transposed_fp8_scale,
)
from sglang.srt.layers.quantization.mxfp4_flashinfer_trtllm_moe import (
    maybe_fuse_routed_scale_and_shared_add,
)
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.layers.rotary_embedding import get_rope_wrapper
from sglang.srt.layers.utils import PPMissingLayer
from sglang.srt.layers.vocab_parallel_embedding import (
    ParallelLMHead,
    VocabParallelEmbedding,
    get_embedding_tp_kwargs,
)
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_attn_backend
from sglang.srt.model_executor.runner import get_is_capture_mode
from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import (
    is_in_breakable_cuda_graph,
)
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
    get_tc_piecewise_forward_context,
    is_in_tc_piecewise_cuda_graph,
)
from sglang.srt.models.deepseek_common.attention_backend_handler import (
    AttentionBackendRegistry,
    resolve_rocm_forward_method,
)
from sglang.srt.models.deepseek_common.attention_forward_methods import (
    AttnForwardMethod,
    DeepseekMHAForwardMixin,
    DeepseekMHARocmForwardMixin,
    DeepseekMLACpuForwardMixin,
    DeepseekMLAForwardMixin,
    DeepseekMLAFusedRopeRocmForwardMixin,
    DeepseekMLARocmForwardMixin,
)
from sglang.srt.models.deepseek_common.deepseek_weight_loader import (
    DeepseekV2WeightLoaderMixin,
)
from sglang.srt.models.deepseek_common.utils import (
    _get_llama_4_scaling,
    _is_block_scale_fp8,
    _is_cpu,
    _is_cpu_amx_available,
    _is_cuda,
    _is_gfx95_supported,
    _is_hip,
    _is_musa,
    _is_npu,
    _is_xpu,
    _use_aiter,
    _use_aiter_bpreshuffle_gfx95,
    _use_aiter_gfx95,
    is_wint4afp8_or_wint4a16_config,
    quant_blocks_shared_experts_fusion,
    tiny_router_gemm_max_tokens,
)
from sglang.srt.runtime_context import (
    attention_backends,
    get_device,
    get_exec,
    get_forward,
    get_model,
    get_parallel,
    get_platform,
    get_spec,
)
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
from sglang.srt.utils import (
    BumpAllocator,
    LazyValue,
    add_prefix,
    is_non_idle_and_non_empty,
    is_sm90_supported,
    make_layers,
    use_intel_amx_backend,
)
from sglang.srt.utils.custom_op import register_custom_op

if _use_aiter:
    from sglang.srt.layers.rocm_linear_utils import aiter_dsv3_router_gemm

if _use_aiter_gfx95:
    from sglang.srt.layers.rocm_linear_utils import (
        get_dsv3_gemm_output_zero_allocator_size,
    )

if _use_aiter:
    pass

if _is_cuda:
    from sglang.kernels.ops.gemm.tiny_gemm import tiny_gemm_bf16
elif _is_npu:
    from sglang.srt.hardware_backend.npu.modules.deepseek_v2_attention_mla_npu import (
        forward_dsa_core_npu,
        forward_dsa_prepare_npu,
        forward_mha_core_npu,
        forward_mha_prepare_npu,
        forward_mla_core_npu,
        forward_mla_prepare_npu,
    )
else:
    pass

from sglang.kernels.ops.gemm.fused_a_gemm import (
    fused_a_gemm_weight_eligible,
    linear_with_fused_a_gemm,
)

logger = logging.getLogger(__name__)

# One-time SGLANG_OPT_MOE_QUANT_ONCE engagement log (see _moe_quant_once_enabled).
_moe_quant_once_logged = False

_enable_pcg_dsv2_dual_stream = (
    _is_cuda and envs.SGLANG_ENABLE_PCG_DSV2_DUAL_STREAM.get()
)


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

        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 not hasattr(self.gate_up_proj, "weight") and hasattr(
            self.gate_up_proj, "weight_packed"
        ):
            self.gate_up_proj.weight = self.gate_up_proj.weight_packed
        if not hasattr(self.down_proj, "weight") and hasattr(
            self.down_proj, "weight_packed"
        ):
            self.down_proj.weight = self.down_proj.weight_packed
        if hidden_act != "silu":
            raise ValueError(
                f"Unsupported activation: {hidden_act}. Only silu is supported for now."
            )
        self.act_fn = SiluAndMul()
        self.use_fused_clamp_act_mul = _is_hip
        self._fused_clamp_fp8_checked = False
        self._fused_clamp_use_fp8 = False

    def forward(
        self,
        x,
        forward_batch=None,
        gemm_output_zero_allocator: BumpAllocator = None,
        gateup_pre_quant: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
    ):
        if (self.tp_size == 1) and x.shape[0] == 0:
            return x

        if (
            getattr(self, "_enable_nvfp4_gemm_swiglu_fusion", False)
            and self.swiglu_limit is None
            and not isinstance(x, tuple)
        ):
            from sglang.kernels.ops.quantization.nvfp4_gemm_swiglu_nvfp4_quant import (
                nvfp4_gemm_swiglu_nvfp4_quant,
            )
            from sglang.srt.layers.quantization.fp4_utils import fp4_quantize

            x_fp4, x_scale = fp4_quantize(
                x, self.gate_up_proj.input_scale_inv, enable_pdl=True
            )
            out_fp4, out_scale = nvfp4_gemm_swiglu_nvfp4_quant(
                x_fp4,
                x_scale,
                self.gate_up_proj.weight_swiglu_interleaved,
                self.gate_up_proj.weight_scale_swiglu_interleaved,
                self.gate_up_proj.alpha,
                self.down_proj.input_scale_inv,
                enable_pdl=True,
            )
            out, _ = self.down_proj((out_fp4, out_scale))
            return out

        if gateup_pre_quant is not None:
            # SGLANG_OPT_MOE_QUANT_ONCE: reuse the caller's per-token-group-128
            # fp8 (q, scale) of x for the gate_up GEMM instead of re-quantizing
            # inside the fp8 linear method. q rows may be padded to a multiple
            # of 4; the caller slices the MLP output back.
            gate_up, _ = self.gate_up_proj(gateup_pre_quant)
        else:
            if (
                gemm_output_zero_allocator is not None
                and x.shape[0] <= 256
                and getattr(self.gate_up_proj, "weight", None) is not None
                and self.gate_up_proj.weight.dtype == torch.uint8
            ):
                y = gemm_output_zero_allocator.allocate(
                    x.shape[0] * self.gate_up_proj.output_size_per_partition
                ).view(x.shape[0], self.gate_up_proj.output_size_per_partition)
                x = (x, None, y)

            gate_up, _ = self.gate_up_proj(x)
        # Fast path: fused silu+clamp+fp8_quant+deepgemm when conditions met.
        # Only valid when down_proj does NOT need an all-reduce and its weights
        # are fp8 (uint8 storage with weight_scale_inv).
        if (
            self.swiglu_limit is not None
            and not self.down_proj.reduce_results
            and getattr(self.down_proj, "weight", None) is not None
            and self.down_proj.weight.dtype == torch.uint8
            and hasattr(self.down_proj, "weight_scale_inv")
        ):
            M, N = gate_up.shape
            down_input_fp8 = gate_up.new_empty((M, N // 2), dtype=torch.float8_e4m3fn)
            scale_block_size = 128
            down_input_scale = create_per_token_group_quant_fp8_output_scale(
                x_shape=(M, N // 2),
                device=gate_up.device,
                group_size=scale_block_size,
                column_major_scales=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
                scale_tma_aligned=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
                scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
            )
            silu_and_mul_contig_post_quant(
                input=gate_up,
                output=down_input_fp8,
                output_scale=down_input_scale,
                quant_group_size=scale_block_size,
                scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
                transposed=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
                swiglu_limit=float(self.swiglu_limit),
            )
            down_output = gate_up.new_empty(
                (M, self.down_proj.output_size), dtype=torch.bfloat16
            )
            deep_gemm_wrapper.gemm_nt_f8f8bf16(
                (down_input_fp8, down_input_scale),
                (self.down_proj.weight, self.down_proj.weight_scale_inv),
                down_output,
            )
            return down_output

        if self.use_fused_clamp_act_mul and self.swiglu_limit is not None:
            from aiter.ops.triton.fusions.fused_clamp_act_mul import (
                fused_clamp_act_mul,
            )

            if not self._fused_clamp_fp8_checked:
                from sglang.srt.layers.quantization.fp8 import Fp8LinearMethod

                qm = getattr(self.down_proj, "quant_method", None)
                self._fused_clamp_use_fp8 = (
                    isinstance(qm, Fp8LinearMethod) and qm.block_quant
                )
                self._fused_clamp_fp8_checked = True

            if self._fused_clamp_use_fp8:
                from aiter import dtypes

                _emit_bpre = emit_transposed_bpreshuffle_scale(
                    gate_up.shape[0], on_bpreshuffle_gfx95=_use_aiter_bpreshuffle_gfx95
                )
                x_fp8, x_scale = fused_clamp_act_mul(
                    gate_up,
                    swiglu_limit=self.swiglu_limit,
                    activation="silu",
                    dtype_quant=dtypes.fp8,
                    transpose_scale=_emit_bpre,
                )
                if _emit_bpre:
                    x_scale = view_aiter_fused_rms_transposed_fp8_scale(x_scale)
                elif _use_aiter_bpreshuffle_gfx95:
                    x_scale = materialize_bpreshuffle_fp8_scale(x_scale)
                x = (x_fp8, x_scale)
            else:
                x = fused_clamp_act_mul(
                    gate_up,
                    swiglu_limit=self.swiglu_limit,
                    activation="silu",
                )

        # Fallback: fused silu+clamp kernel (still faster than unfused)
        elif self.swiglu_limit is not None:
            if _is_npu:
                x = torch.ops.npu.npu_clipped_swiglu(
                    gate_up,
                    alpha=1,
                    limit=self.swiglu_limit,
                    bias=0,
                    interleaved=False,
                )
            else:
                M, N = gate_up.shape
                x = gate_up.new_empty((M, N // 2))
                silu_and_mul_clamp(gate_up, x, float(self.swiglu_limit))
        else:
            x = self.act_fn(gate_up)
        x, _ = self.down_proj(x)
        return x


class MoEGate(nn.Module):
    def __init__(
        self,
        config,
        quant_config,
        prefix: str = "",
        is_hash_moe: bool = False,
        is_deepseek_v4: bool = False,
    ):
        super().__init__()
        self.is_deepseek_v4 = is_deepseek_v4
        self.weight = nn.Parameter(
            torch.empty(
                (config.n_routed_experts, config.hidden_size),
                dtype=(
                    torch.float32
                    if getattr(config, "router_fp32", False)
                    else torch.get_default_dtype()
                ),
            )
        )
        if config.topk_method == "noaux_tc" and not is_hash_moe:
            correction_bias_dtype = torch.float32
            # GLM-5.2's bias sits at an offset where its spread is only a few bf16 ULPs
            # wide, so bf16 collapses it and reorders top-k routing. HF stores it fp32.
            if quant_config is not None and not is_glm_moe_dsa(config):
                if _use_aiter and quant_config.get_name() in (
                    "fp8",
                    "compressed_tensors",
                    "quark",
                ):
                    correction_bias_dtype = torch.bfloat16
            correction_bias = torch.empty(
                (config.n_routed_experts), dtype=correction_bias_dtype
            )
            if quant_config is not None and quant_config.get_name() == "expert_pack":
                correction_bias.zero_()
            self.e_score_correction_bias = nn.Parameter(correction_bias)
        else:
            self.e_score_correction_bias = None
        if _is_cpu and _is_cpu_amx_available:
            self.quant_method = PackWeightMethod(weight_names=["weight"])
        self.tiny_router_gemm_max_tokens = tiny_router_gemm_max_tokens(
            num_experts=config.n_routed_experts,
            hidden_size=config.hidden_size,
            weight_dtype=self.weight.dtype,
        )

    def forward(
        self,
        hidden_states,
        gemm_output_zero_allocator: BumpAllocator = None,
        forward_batch: ForwardBatch = None,
    ):
        if self.weight.dtype == torch.float32:
            return F.linear(hidden_states.float(), self.weight)

        if use_intel_amx_backend(self):
            return torch.ops.sgl_kernel.weight_packed_linear(
                hidden_states,
                self.weight,
                None,  # bias
                True,  # is_vnni
            )

        if get_exec().deterministic.enable_deterministic_inference:
            return F.linear(hidden_states, self.weight, None)

        if hidden_states.shape[0] <= self.tiny_router_gemm_max_tokens:
            logits = tiny_gemm_bf16(
                hidden_states,
                self.weight,
                out_dtype=torch.float32,
                max_m=self.tiny_router_gemm_max_tokens,
            )
        elif _use_aiter:
            logits = aiter_dsv3_router_gemm(hidden_states, self.weight)
        elif not _is_cuda:
            logits = F.linear(hidden_states, self.weight, None)
        else:
            # cuBLAS bf16 x bf16 -> fp32 GEMM (torch.mm's out_dtype kwarg is CUDA-only)
            from sglang.kernels.ops.attention.dsv4 import linear_bf16_fp32

            logits = linear_bf16_fp32(hidden_states, self.weight)

        return logits


class DeepseekV2MoE(nn.Module):
    def __init__(
        self,
        config: PretrainedConfig,
        layer_id: int,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
        alt_stream: Optional[torch.cuda.Stream] = None,
        is_nextn: bool = False,
        is_deepseek_v4: bool = False,
    ):
        super().__init__()
        self.tp_size = get_parallel().tp_size
        self.moe_ep_size = get_parallel().moe_ep_size
        self.routed_scaling_factor = config.routed_scaling_factor
        self.n_shared_experts = config.n_shared_experts

        n_shared_experts = (
            0 if config.n_shared_experts is None else int(config.n_shared_experts)
        )
        _fusion_disabled = is_shared_experts_fusion_disabled()

        # num_fused_shared_experts drives weight remapping in deepseek_weight_loader:
        # mlp.shared_experts → mlp.experts.256 when > 0.
        self.num_fused_shared_experts = 0 if _fusion_disabled else n_shared_experts

        # DeepEP and MegaMOE shared expert fusion: shared expert is fused into
        # the same MoE kernel as a local expert at each EP rank. Expert layout
        # is expanded from 256 routed to 256+EP_size (e.g. 272 for EP=16).
        _uses_per_rank_shared_slots = has_per_rank_fused_shared_slots(
            self.num_fused_shared_experts
        )

        if _uses_per_rank_shared_slots:
            # 256 routed + EP_size shared slots = 272 experts total (for EP=16)
            num_experts_for_moe = config.n_routed_experts + self.moe_ep_size
            top_k_for_moe = config.num_experts_per_tok + 1  # 8 routed + 1 shared
            # Interleaving for DeepEP/MegaMOE dispatch is handled by TopK internally.
        else:
            num_experts_for_moe = (
                config.n_routed_experts + self.num_fused_shared_experts
            )
            top_k_for_moe = config.num_experts_per_tok + self.num_fused_shared_experts

        self.config = config
        self.layer_id = layer_id
        self.alt_stream = alt_stream
        self.is_nextn = is_nextn

        n_hash_layers = getattr(config, "num_hash_layers", 0)
        self.is_hash = layer_id < n_hash_layers and not (is_deepseek_v4 and is_nextn)

        if self.tp_size > config.n_routed_experts:
            raise ValueError(
                f"Tensor parallel size {self.tp_size} is greater than "
                f"the number of experts {config.n_routed_experts}."
            )

        if config.hidden_act != "silu":
            raise ValueError(
                f"Unsupported activation: {config.hidden_act}. "
                "Only silu is supported for now."
            )

        self.gate = MoEGate(
            config=config,
            quant_config=quant_config,
            prefix=add_prefix("gate", prefix),
            is_hash_moe=self.is_hash,
            is_deepseek_v4=is_deepseek_v4,
        )

        # scaling factor for fused shared experts on AMD-platform.
        # DeepEP/MegaMOE doesn't need this: shared expert is only computed on home rank
        # (not all-reduced), so no 1/ep_size correction is needed.
        fused_shared_experts_scaling_factor = None
        if (
            self.moe_ep_size > 1
            and self.num_fused_shared_experts > 0
            and not _uses_per_rank_shared_slots
        ):
            # if enable_ep_moe tp_szie == ep_size, every gpu get shared experts gemm output
            # so we scale with 1 / self.moe_ep_size in ep mode which will make it equalation as in tp mode
            # with fused_shared_experts
            fused_shared_experts_scaling_factor = 1.0 / float(self.moe_ep_size)

        self.experts = get_moe_impl_class(quant_config)(
            num_experts=num_experts_for_moe + get_exec().moe.ep_num_redundant_experts,
            num_fused_shared_experts=self.num_fused_shared_experts,
            top_k=top_k_for_moe,
            hidden_size=config.hidden_size,
            intermediate_size=config.moe_intermediate_size,
            layer_id=self.layer_id,
            quant_config=quant_config,
            routed_scaling_factor=self.routed_scaling_factor,
            routing_method_type=getattr(
                config, "routing_method_type", RoutingMethodType.DeepSeekV3
            ),
            swiglu_limit=getattr(config, "swiglu_limit", None),
            prefix=add_prefix("experts", prefix),
        )

        if self.is_hash and not (is_nextn and is_deepseek_v4):
            self.topk = HashTopK(
                topk=config.num_experts_per_tok + self.num_fused_shared_experts,
                num_experts=config.n_routed_experts,
                num_fused_shared_experts=self.num_fused_shared_experts,
                vocab_size=config.vocab_size,
                scoring_func=config.scoring_func,
                routed_scaling_factor=self.routed_scaling_factor,
                apply_routed_scaling_factor_on_output=self.experts.should_fuse_routed_scaling_factor_in_topk,
                layer_id=self.layer_id,
            )
        else:
            # Default: grouped noaux_tc top-k. Covers V3/V3.2/GLM-5/Glm4MoeLite.
            topk_kwargs = dict(
                top_k=config.num_experts_per_tok + self.num_fused_shared_experts,
                layer_id=self.layer_id,
                renormalize=config.norm_topk_prob,
                use_grouped_topk=True,
                num_expert_group=config.n_group,
                num_fused_shared_experts=self.num_fused_shared_experts,
                topk_group=config.topk_group,
                scoring_func=config.scoring_func,
                correction_bias=self.gate.e_score_correction_bias,
                quant_config=quant_config,
                routed_scaling_factor=self.routed_scaling_factor,
                apply_routed_scaling_factor_on_output=self.experts.should_fuse_routed_scaling_factor_in_topk,
                fused_shared_experts_scaling_factor=fused_shared_experts_scaling_factor,
                # Some Fp4 MoE backends require the output format to be bypassed but the MTP layers are unquantized
                # and requires the output format to be standard (except trtllm). We use quant_config to determine the output format.
                output_format=(
                    TopKOutputFormat.STANDARD
                    if (quant_config is None)
                    and (not get_moe_runner_backend().is_flashinfer_trtllm())
                    else None
                ),
            )
            # DSV4 override: ungrouped sqrtsoftplus + fp4 expert layout flag.
            if is_deepseek_v4:
                topk_kwargs.update(
                    use_grouped_topk=False,
                    scoring_func=config.scoring_func,
                    is_fp4_experts=getattr(quant_config, "is_fp4_experts", False),
                    apply_routed_scaling_factor_on_output=(
                        True
                        if _use_aiter
                        else self.experts.should_fuse_routed_scaling_factor_in_topk
                    ),
                )
            self.topk = TopK(**topk_kwargs)

        self.shared_experts_is_int8 = False
        self.shared_experts_is_fp8 = False
        self.shared_experts_weight_block_size = None
        self._shared_expert_tp1 = False
        # Shared experts: skip when fused into MoE kernel
        # (self.num_fused_shared_experts > 0) or when DeepEP/MegaMOE fusion is enabled.
        if (
            config.n_shared_experts is not None
            and config.n_shared_experts > 0
            and self.num_fused_shared_experts == 0
            and not _uses_per_rank_shared_slots
        ):
            intermediate_size = config.moe_intermediate_size * config.n_shared_experts
            # Disable TP for shared experts for A2A/FP4 allgather paths, or when
            # explicitly requested for DSV4 checkpoints whose shared scales are
            # not divisible by the global TP size.
            _shared_expert_use_tp1 = (
                get_moe_a2a_backend().is_deepep()
                or get_moe_a2a_backend().is_pplx()
                or get_moe_a2a_backend().is_mooncake()
                or get_moe_a2a_backend().is_nixl()
                or get_moe_a2a_backend().is_mori()
                or get_moe_a2a_backend().is_ascend_fuseep()
                or get_moe_a2a_backend().is_flashinfer()
                or get_moe_a2a_backend().is_megamoe()
                or get_moe_a2a_backend().is_flashinfer_megamoe()
                or get_moe_a2a_backend().is_deepep_v2()
                or should_use_flashinfer_cutlass_moe_fp4_allgather()
                or envs.SGLANG_SHARED_EXPERT_TP1.get()
            )
            self.shared_experts = DeepseekV2MLP(
                hidden_size=config.hidden_size,
                intermediate_size=intermediate_size,
                hidden_act=config.hidden_act,
                quant_config=quant_config,
                reduce_results=False,
                swiglu_limit=getattr(config, "swiglu_limit", None),
                prefix=add_prefix("shared_experts", prefix),
                **(dict(tp_rank=0, tp_size=1) if _shared_expert_use_tp1 else {}),
            )
            # Flags must be set before weight load so
            # process_weights_after_loading sees them and builds the
            # [Up, Gate]-interleaved weight + scale.
            from sglang.srt.layers.quantization.modelopt_quant import (
                ModelOptFp4LinearMethod,
            )

            fc1_n = self.shared_experts.gate_up_proj.output_size_per_partition
            if (
                get_platform().is_sm100
                and isinstance(
                    self.shared_experts.gate_up_proj.quant_method,
                    ModelOptFp4LinearMethod,
                )
                and self.shared_experts.gate_up_proj.quant_method.quant_mode == "w4a4"
                and isinstance(
                    self.shared_experts.down_proj.quant_method,
                    ModelOptFp4LinearMethod,
                )
                and fc1_n % 128 == 0
                and self.shared_experts.swiglu_limit is None
                and not check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
            ):
                self.shared_experts.gate_up_proj._interleave_for_swiglu_fusion = True
                self.shared_experts._enable_nvfp4_gemm_swiglu_fusion = True
                self.shared_experts.down_proj._accepts_prequantized_fp4 = True
            self._shared_expert_tp1 = _shared_expert_use_tp1
            is_packed_weight = (
                hasattr(self.shared_experts.gate_up_proj.quant_method, "quant_config")
                and self.shared_experts.gate_up_proj.quant_method.quant_config.get_name()
                in {
                    "awq",
                    "awq_marlin",
                    "moe_wna16",
                }
            )
            shared_gate_up_weight = getattr(
                self.shared_experts.gate_up_proj, "weight", None
            )
            if shared_gate_up_weight is None:
                shared_gate_up_weight = getattr(
                    self.shared_experts.gate_up_proj, "qweight", None
                )
            if shared_gate_up_weight is None:
                raise ValueError(
                    "shared expert gate/up projection has no weight storage"
                )
            self.shared_experts_is_int8 = (
                not is_packed_weight and shared_gate_up_weight.dtype == torch.int8
            )
            self.shared_experts_is_fp8 = (
                not is_packed_weight
                and shared_gate_up_weight.dtype == torch.float8_e4m3fn
            )
            if self.shared_experts_is_fp8 and not _is_npu:
                if (
                    _use_aiter
                    and config.quantization_config.get("quant_method")
                    == "compressed-tensors"
                ):
                    # For compressed-tensors ptpc model, don't need to check the weight_block_size
                    pass
                else:
                    assert (
                        self.shared_experts.gate_up_proj.quant_method.quant_config.weight_block_size
                        == self.shared_experts.down_proj.quant_method.quant_config.weight_block_size
                    )
                    self.shared_experts_weight_block_size = self.shared_experts.gate_up_proj.quant_method.quant_config.weight_block_size

        self.top_k = config.num_experts_per_tok

        if (
            get_moe_a2a_backend().is_deepep()
            or get_moe_a2a_backend().is_mooncake()
            or get_moe_a2a_backend().is_nixl()
            or get_moe_a2a_backend().is_mori()
            or get_moe_a2a_backend().is_ascend_fuseep()
            or get_moe_a2a_backend().is_deepep_v2()
        ):
            # TODO: we will support tp < ep in the future
            self.ep_size = get_parallel().moe_ep_size
            self.num_experts = (
                config.n_routed_experts + get_exec().moe.ep_num_redundant_experts
            )
            self.renormalize = config.norm_topk_prob
            self.topk_group = config.topk_group
            self.num_expert_group = config.n_group
            self.correction_bias = (
                self.gate.e_score_correction_bias.data
                if self.gate.e_score_correction_bias is not None
                else None
            )

        self._enable_a2a_moe = (
            get_moe_a2a_backend().is_deepep()
            or get_moe_a2a_backend().is_mooncake()
            or get_moe_a2a_backend().is_nixl()
            or get_moe_a2a_backend().is_mori()
            or get_moe_a2a_backend().is_ascend_fuseep()
            or get_moe_a2a_backend().is_flashinfer()
            or get_moe_a2a_backend().is_deepep_v2()
        )
        self._fuse_shared_experts_inside_sbo = SboFlags.fuse_shared_experts_inside_sbo()
        # SGLANG_OPT_MOE_QUANT_ONCE eligibility, resolved lazily on first
        # forward (weights and runner are final by then). None = undecided.
        self._moe_quant_once: Optional[bool] = None

    def get_moe_weights(self):
        # EPLB only rebalances physical routed experts. Fused shared expert
        # slots live after each rank's routed slots and must stay stable.
        num_local_experts_for_eplb = (
            self.experts.num_local_experts - self.num_fused_shared_experts
        )

        return [
            x.data[:num_local_experts_for_eplb]
            for name, x in self.experts.named_parameters()
            if name not in ["correction_bias"]
            and filter_moe_weight_param_global_expert(
                name, x, self.experts.num_local_experts
            )
        ]

    def _can_dual_stream_graph(self, hidden_states: torch.Tensor) -> bool:
        return (
            _enable_pcg_dsv2_dual_stream
            and (is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph())
            and get_moe_runner_backend().is_flashinfer_trtllm()
            and self.alt_stream is not None
            and self.num_fused_shared_experts == 0
            and hidden_states.shape[0] > 0
            and hasattr(self, "shared_experts")
            and getattr(self.experts, "use_flashinfer_trtllm_moe", False)
            and not self._enable_a2a_moe
            and not self._fuse_shared_experts_inside_sbo
            and not getattr(self, "is_hash", False)
            and not get_exec().moe.enable_eplb
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        forward_batch: Optional[ForwardBatch] = None,
        gemm_output_zero_allocator: BumpAllocator = None,
        input_ids: Optional[torch.Tensor] = None,
        input_ids_global: Optional[torch.Tensor] = None,
        skip_shared_experts: bool = False,
    ) -> torch.Tensor:
        from sglang.srt.layers.moe.mega_moe import forward_mega_moe, should_use_mega_moe

        if should_use_mega_moe(self, hidden_states):
            return forward_mega_moe(
                self,
                hidden_states,
                forward_batch,
                input_ids_global=input_ids_global,
            )

        if not self._enable_a2a_moe:
            if self._can_dual_stream_graph(hidden_states):
                fwd = get_forward()
                return dsv2_flashinfer_moe_dual_stream_graph(
                    hidden_states,
                    self.layer_id,
                    fwd.fuse_mlp_allreduce,
                    fwd.mlp_reduce_scatter,
                )
            elif (
                self.alt_stream is not None
                and self.num_fused_shared_experts == 0
                and hidden_states.shape[0] > 0
                and get_is_capture_mode()
            ):
                return self.forward_normal_dual_stream(
                    hidden_states,
                    gemm_output_zero_allocator,
                    input_ids,
                    input_ids_global=input_ids_global,
                )
            else:
                return self.forward_normal(
                    hidden_states,
                    gemm_output_zero_allocator,
                    input_ids,
                    input_ids_global=input_ids_global,
                    skip_shared_experts=skip_shared_experts,
                )
        else:
            return self.forward_deepep(
                hidden_states, forward_batch, input_ids_global=input_ids_global
            )

    def forward_normal_dual_stream(
        self,
        hidden_states: torch.Tensor,
        gemm_output_zero_allocator: BumpAllocator = None,
        input_ids: Optional[torch.Tensor] = None,
        input_ids_global: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        # Note(kpham-sgl): issue order satisfies 3 constraints:
        # - no stream explosion: main (routed) issued before alt block -> capture reuses 1 alt stream;
        # - PDL overlap: routed is the last main-stream kernel (fuses w/ residual add);
        # - dispose_tensor: disabled during capture (CaptureFlags.disable_dispose_tensor) so the routed
        #   deep_gemm does not free hidden_states, which the shared expert reads on the alt stream.
        use_flashinfer_trtllm_bypass = get_forward().flashinfer_trtllm_bypass
        current_stream = torch.cuda.current_stream()
        # Quantize-once (SGLANG_OPT_MOE_QUANT_ONCE) must happen on the main
        # stream BEFORE the alt-stream fork so both consumers see it.
        pre_quant_input = (
            None
            if use_flashinfer_trtllm_bypass
            else self._maybe_quant_moe_input_once(hidden_states)
        )
        self.alt_stream.wait_stream(current_stream)
        has_shared_output = (
            hidden_states.shape[0] > 0 and self.num_fused_shared_experts == 0
        )
        dispatch_info = (
            ExpertLocationDispatchInfo.init_new(layer_id=self.layer_id)
            if get_exec().moe.enable_eplb and not self.is_nextn
            else None
        )
        # router_logits: (num_tokens, n_experts)
        router_logits = self.gate(hidden_states, gemm_output_zero_allocator)
        if use_flashinfer_trtllm_bypass:
            topk_output = BypassedTopKOutput(
                hidden_states=hidden_states,
                router_logits=router_logits,
                topk_config=self.topk.topk_config,
            )
        else:
            topk_kwargs = (
                {"input_ids": input_ids_global}
                if getattr(self, "is_hash", False)
                else {}
            )
            topk_output = self.topk(
                hidden_states,
                router_logits,
                expert_location_dispatch_info=dispatch_info,
                **topk_kwargs,
            )
        deferred_finalize = (
            has_shared_output
            and not self._shared_expert_tp1
            and topk_output.format == TopKOutputFormat.BYPASSED
            and self.experts.supports_deferred_finalize
        )
        if deferred_finalize:
            final_hidden_states = self.experts.forward_deferred_finalize(
                hidden_states, topk_output
            )
        elif use_flashinfer_trtllm_bypass:
            final_hidden_states = self.experts.forward_impl(hidden_states, topk_output)
        elif pre_quant_input is not None:
            final_hidden_states = self.experts(
                hidden_states, topk_output, pre_quant_input=pre_quant_input
            )
        else:
            final_hidden_states = self.experts(hidden_states, topk_output)
        if (
            not _is_cuda
            and not _is_musa
            and not _use_aiter
            or isinstance(self.experts.quant_method, KTEPWrapperMethod)
        ):
            final_hidden_states *= self.routed_scaling_factor

        # Shared expert on alt stream, issued AFTER the main (routed) branch. See note above.
        with torch.cuda.stream(self.alt_stream):
            shared_output = self._forward_shared_experts(
                hidden_states,
                gemm_output_zero_allocator,
                pre_quant_input=pre_quant_input,
            )

        current_stream.wait_stream(self.alt_stream)

        if deferred_finalize:
            from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
                finalize_flashinfer_trtllm_deferred_output,
            )

            final_hidden_states = finalize_flashinfer_trtllm_deferred_output(
                final_hidden_states,
                shared_output,
            )
        else:
            final_hidden_states = maybe_fuse_routed_scale_and_shared_add(
                self.experts,
                final_hidden_states,
                None if self._shared_expert_tp1 else shared_output,
                self.routed_scaling_factor,
            )

        if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
            is_tp_path=True,
        ):
            final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
        # TP1 shared experts are replicated, so add them after all-reduce to
        # avoid summing the same shared output once per TP rank.
        if self._shared_expert_tp1:
            final_hidden_states += shared_output
        return final_hidden_states

    def forward_normal(
        self,
        hidden_states: torch.Tensor,
        gemm_output_zero_allocator: BumpAllocator = None,
        input_ids: Optional[torch.Tensor] = None,
        input_ids_global: Optional[torch.Tensor] = None,
        skip_shared_experts: bool = False,
    ) -> torch.Tensor:
        if hasattr(self, "shared_experts") and use_intel_amx_backend(
            self.shared_experts.gate_up_proj
        ):
            return self.forward_cpu(hidden_states)
        dispatch_info = (
            ExpertLocationDispatchInfo.init_new(layer_id=self.layer_id)
            if get_exec().moe.enable_eplb and not self.is_nextn
            else None
        )
        defer_shared = not self.experts.moe_runner_config.inplace
        # PoC (SGLANG_DP_SHARED_EXPERT_LOCAL): shared expert is computed on the LOCAL
        # hidden in the decoder layer (before the dp gather) and added after the
        # reduce_scatterv. When set, never compute/add it here (on the global buffer).
        shared_output = None
        if hidden_states.shape[0] > 0:
            # Quantize-once (SGLANG_OPT_MOE_QUANT_ONCE): only worthwhile when
            # the shared expert also runs here on the same tensor.
            pre_quant_input = (
                None
                if skip_shared_experts
                else self._maybe_quant_moe_input_once(hidden_states)
            )
            if (
                not defer_shared
                and not self._fuse_shared_experts_inside_sbo
                and not skip_shared_experts
            ):
                shared_output = self._forward_shared_experts(
                    hidden_states,
                    gemm_output_zero_allocator,
                    pre_quant_input=pre_quant_input,
                )
            # router_logits: (num_tokens, n_experts)
            router_logits = self.gate(hidden_states, gemm_output_zero_allocator)
            topk_kwargs = (
                {"input_ids": input_ids_global}
                if getattr(self, "is_hash", False)
                else {}
            )
            topk_output = self.topk(
                hidden_states,
                router_logits,
                expert_location_dispatch_info=dispatch_info,
                **topk_kwargs,
            )
        else:
            pre_quant_input = None
            shared_output = None
            topk_output = self.topk.empty_topk_output(
                hidden_states.device, layer_id=self.layer_id
            )

        if self._fuse_shared_experts_inside_sbo and not skip_shared_experts:
            shared_output = None

            def _pre_combine_hook(
                dispatcher: BaseDispatcher, combine_input: CombineInput
            ):

                nonlocal shared_output
                self.alt_stream.wait_stream(torch.cuda.current_stream())
                with torch.cuda.stream(self.alt_stream):
                    shared_output = self._forward_shared_experts(
                        hidden_states, gemm_output_zero_allocator
                    )

                pre_combine_hook_handle.remove()

            def _post_combine_hook(
                dispatcher: BaseDispatcher, hidden_states: torch.Tensor
            ):
                nonlocal shared_output
                torch.cuda.current_stream().wait_stream(self.alt_stream)
                post_combine_hook_handle.remove()

            pre_combine_hook_handle = self.experts.dispatcher.register_pre_combine_hook(
                _pre_combine_hook
            )
            post_combine_hook_handle = (
                self.experts.dispatcher.register_post_combine_hook(_post_combine_hook)
            )

        if pre_quant_input is not None:
            final_hidden_states = self.experts(
                hidden_states,
                topk_output,
                pre_quant_input=pre_quant_input,
            )
        else:
            final_hidden_states = self.experts(
                hidden_states,
                topk_output,
            )
        if (
            not _is_cuda
            and not _is_musa
            and not _is_xpu
            and not _use_aiter
            or isinstance(self.experts.quant_method, KTEPWrapperMethod)
        ):
            # fused in biased_grouped_topk so we can skip here
            final_hidden_states *= self.routed_scaling_factor

        if (
            defer_shared
            and hidden_states.shape[0] > 0
            and not self._fuse_shared_experts_inside_sbo
            and not skip_shared_experts
        ):
            shared_output = self._forward_shared_experts(
                hidden_states,
                gemm_output_zero_allocator,
                pre_quant_input=pre_quant_input,
            )

        final_hidden_states = maybe_fuse_routed_scale_and_shared_add(
            self.experts,
            final_hidden_states,
            None if self._shared_expert_tp1 else shared_output,
            self.routed_scaling_factor,
        )

        if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
            is_tp_path=True,
        ):
            final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
        # TP1 shared experts are replicated, so add them after all-reduce to
        # avoid summing the same shared output once per TP rank.
        if shared_output is not None and self._shared_expert_tp1:
            final_hidden_states += shared_output
        return final_hidden_states

    def forward_cpu(
        self,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        # router_logits: (num_tokens, n_experts)
        router_logits = self.gate(hidden_states)
        topk_output = self.topk(hidden_states, router_logits)
        fused_experts_out = self.experts(
            hidden_states=hidden_states, topk_output=topk_output
        )

        assert use_intel_amx_backend(
            self.shared_experts.gate_up_proj
        ) == use_intel_amx_backend(self.shared_experts.down_proj)
        # [Note] inplace should be False in fused_experts.
        # If inplace is True in fused_experts (self.experts), hidden_states will be changed after fused_experts
        # While hidden_states is still needed in shared_expert.
        final_hidden_states = torch.ops.sgl_kernel.shared_expert_cpu(
            hidden_states,
            self.shared_experts.gate_up_proj.weight,
            self.shared_experts.down_proj.weight,
            fused_experts_out,
            self.routed_scaling_factor,
            True,  # inplace
            self.shared_experts_is_int8,  # use_int8_w8a8
            self.shared_experts_is_fp8,  # use_fp8_w8a16
            (
                self.shared_experts.gate_up_proj.weight_scale
                if self.shared_experts_is_int8
                else (
                    self.shared_experts.gate_up_proj.weight_scale_inv
                    if self.shared_experts_is_fp8
                    else None
                )
            ),  # w1_scale
            (
                self.shared_experts.down_proj.weight_scale
                if self.shared_experts_is_int8
                else (
                    self.shared_experts.down_proj.weight_scale_inv
                    if self.shared_experts_is_fp8
                    else None
                )
            ),  # w2_scale
            (
                self.shared_experts_weight_block_size
                if self.shared_experts_is_fp8
                else None
            ),  # block_size
            True,  # is_vnni
        )
        if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
            is_tp_path=True,
        ):
            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,
        input_ids_global: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        shared_output = None
        sbo_enabled_flag = self._fuse_shared_experts_inside_sbo and not self.is_nextn
        sbo_overlap_dispatch_flag = (
            sbo_enabled_flag and SboFlags.enable_dispatch_shared_one_stream_overlap()
        )
        sbo_overlap_combine_flag = (
            sbo_enabled_flag and SboFlags.enable_combine_shared_two_stream_overlap()
        )

        if hidden_states.shape[0] > 0:
            # router_logits: (num_tokens, n_experts)
            router_logits = self.gate(hidden_states, forward_batch=forward_batch)
            if not sbo_enabled_flag and self.num_fused_shared_experts == 0:
                if self.alt_stream is not None:
                    self.alt_stream.wait_stream(torch.cuda.current_stream())
                    with torch.cuda.stream(self.alt_stream):
                        shared_output = self._forward_shared_experts(hidden_states)
                        shared_output.record_stream(self.alt_stream)
                        shared_event = self.alt_stream.record_event()
                    if is_in_breakable_cuda_graph():
                        # The MoE call below is an eager break, so record
                        # and wait must share one capture; joining here means
                        # the shared experts overlap nothing. The alt stream
                        # is kept for record_stream: without that marking the
                        # allocator recycles shared_output across the break.
                        torch.cuda.current_stream().wait_event(shared_event)
                else:
                    shared_output = self._forward_shared_experts(hidden_states)
            topk_kwargs = (
                {"input_ids": input_ids_global}
                if getattr(self, "is_hash", False)
                else {}
            )
            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,
                    )
                    if not self.is_nextn
                    else None
                ),
                **topk_kwargs,
            )
        else:
            topk_output = self.topk.empty_topk_output(
                hidden_states.device, layer_id=self.layer_id
            )

        if sbo_overlap_dispatch_flag:
            shared_output = None

            def _deepep_dispatch_hook(dispatcher: BaseDispatcher):
                nonlocal shared_output
                shared_output = self._forward_shared_experts(hidden_states)
                for handle in deepep_dispatch_hook_handle:
                    handle.remove()

            def _post_dispatch_hook(
                dispatcher: BaseDispatcher, dispatch_output: DispatchOutput
            ):
                combine_overlap_args, down_gemm_overlap_args, meta_overlap_args = (
                    compute_overlap_args(dispatch_output, self.alt_stream)
                )
                dispatcher.set_overlap_args(
                    combine_overlap_args=combine_overlap_args,
                    meta_overlap_args=meta_overlap_args,
                )
                self.experts.set_overlap_args(
                    down_gemm_overlap_args=down_gemm_overlap_args,
                    meta_overlap_args=meta_overlap_args,
                )
                post_dispatch_hook_handle.remove()

            def _post_combine_hook(
                dispatcher: BaseDispatcher, hidden_states: torch.Tensor
            ):
                dispatcher.clear_overlap_args()
                self.experts.clear_overlap_args()
                post_combine_hook_handle.remove()

            assert isinstance(self.experts.dispatcher, MaybeTboDeepEPDispatcher)
            deepep_dispatch_hook_handle = (
                self.experts.dispatcher.register_deepep_dispatch_hook(
                    _deepep_dispatch_hook
                )
            )
            post_dispatch_hook_handle = (
                self.experts.dispatcher.register_post_dispatch_hook(_post_dispatch_hook)
            )
            post_combine_hook_handle = (
                self.experts.dispatcher.register_post_combine_hook(_post_combine_hook)
            )

        elif sbo_overlap_combine_flag:
            shared_output = None

            def _post_dispatch_hook(
                dispatcher: BaseDispatcher, dispatch_output: DispatchOutput
            ):

                combine_overlap_args, down_gemm_overlap_args, meta_overlap_args = (
                    compute_overlap_args(dispatch_output, self.alt_stream)
                )
                dispatcher.set_overlap_args(
                    combine_overlap_args=combine_overlap_args,
                    meta_overlap_args=meta_overlap_args,
                )
                self.experts.set_overlap_args(
                    down_gemm_overlap_args=down_gemm_overlap_args,
                    meta_overlap_args=meta_overlap_args,
                )

                post_dispatch_hook_handle.remove()

            def _pre_combine_hook(
                dispatcher: BaseDispatcher, combine_input: CombineInput
            ):

                nonlocal shared_output

                if (
                    e := dispatcher.meta_overlap_args.get("record_event_after_down")
                ) is not None:
                    e.record()

                # TODO reduce sm for non-deepgemm
                with deep_gemm_wrapper.configure_deep_gemm_num_sms(
                    dispatcher.meta_overlap_args["compute_num_sms"]
                ):
                    shared_output = self._forward_shared_experts(hidden_states)

                pre_combine_hook_handle.remove()

            def _post_combine_hook(
                dispatcher: BaseDispatcher, hidden_states: torch.Tensor
            ):
                dispatcher.clear_overlap_args()
                self.experts.clear_overlap_args()
                post_combine_hook_handle.remove()

            post_dispatch_hook_handle = (
                self.experts.dispatcher.register_post_dispatch_hook(_post_dispatch_hook)
            )
            pre_combine_hook_handle = self.experts.dispatcher.register_pre_combine_hook(
                _pre_combine_hook
            )
            post_combine_hook_handle = (
                self.experts.dispatcher.register_post_combine_hook(_post_combine_hook)
            )
        elif envs.SGLANG_BLACKWELL_OVERLAP_SHARED_EXPERTS_OUTSIDE_SBO.get():
            # On GB200: Shared experts overlapped on alt_stream, down gemm overlapped with DeepEP Combine

            def _post_dispatch_hook(
                dispatcher: BaseDispatcher, dispatch_output: DispatchOutput
            ):

                combine_overlap_args, down_gemm_overlap_args, meta_overlap_args = (
                    compute_overlap_args(dispatch_output, self.alt_stream)
                )
                dispatcher.set_overlap_args(
                    combine_overlap_args=combine_overlap_args,
                    meta_overlap_args=meta_overlap_args,
                )
                self.experts.set_overlap_args(
                    down_gemm_overlap_args=down_gemm_overlap_args,
                    meta_overlap_args=meta_overlap_args,
                )

                post_dispatch_hook_handle.remove()

            def _pre_combine_hook(
                dispatcher: BaseDispatcher, combine_input: CombineInput
            ):
                if (
                    e := dispatcher.meta_overlap_args.get("record_event_after_down")
                ) is not None:
                    e.record()
                pre_combine_hook_handle.remove()

            def _post_combine_hook(
                dispatcher: BaseDispatcher, hidden_states: torch.Tensor
            ):
                dispatcher.clear_overlap_args()
                self.experts.clear_overlap_args()
                post_combine_hook_handle.remove()

            post_dispatch_hook_handle = (
                self.experts.dispatcher.register_post_dispatch_hook(_post_dispatch_hook)
            )
            pre_combine_hook_handle = self.experts.dispatcher.register_pre_combine_hook(
                _pre_combine_hook
            )
            post_combine_hook_handle = (
                self.experts.dispatcher.register_post_combine_hook(_post_combine_hook)
            )

        final_hidden_states = self.experts(
            hidden_states=hidden_states,
            topk_output=topk_output,
        )

        if (
            hidden_states.shape[0] > 0
            and not sbo_enabled_flag
            and self.num_fused_shared_experts == 0
            and self.alt_stream is not None
            and not is_in_breakable_cuda_graph()
        ):
            torch.cuda.current_stream().wait_event(shared_event)

        if shared_output is not None:
            x = shared_output
            # aiter moe call will handle routed_scaling_factor in the function
            # so add _use_aiter condition to eliminate to use self.routed_scaling_factor in add_ call
            if self.experts.should_fuse_routed_scaling_factor_in_topk or _use_aiter:
                x.add_(final_hidden_states)
            else:
                x.add_(final_hidden_states, alpha=self.routed_scaling_factor)
            final_hidden_states = x
        else:
            if not (
                self.experts.should_fuse_routed_scaling_factor_in_topk or _use_aiter
            ):
                final_hidden_states *= self.routed_scaling_factor

        return final_hidden_states

    def _forward_shared_experts(
        self,
        hidden_states,
        gemm_output_zero_allocator: BumpAllocator = None,
        pre_quant_input: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
    ):
        if (hidden_states.shape[0] > 0) and (self.num_fused_shared_experts == 0):
            if pre_quant_input is not None:
                # SGLANG_OPT_MOE_QUANT_ONCE: (q, s) rows may be padded to a
                # multiple of 4; the padded rows flow through the MLP (all ops
                # are row-local) and are sliced off here.
                out = self.shared_experts(
                    hidden_states, gateup_pre_quant=pre_quant_input
                )
                return out[: hidden_states.shape[0]]
            return self.shared_experts(
                hidden_states, gemm_output_zero_allocator=gemm_output_zero_allocator
            )
        else:
            return None

    def _moe_quant_once_enabled(self) -> bool:
        """SGLANG_OPT_MOE_QUANT_ONCE: quantize the (dp-gathered) MoE input to
        per-token-group-128 fp8 once per layer and feed both the fused shared
        expert's fp8 GEMM (cutlass or deepgemm w8a8 linear) and the routed
        experts' triton fused runner, instead of quantizing the same
        [T, hidden] tensor twice with different scale layouts."""
        if self._moe_quant_once is None:
            self._moe_quant_once, reason = self._compute_moe_quant_once_enabled()
            global _moe_quant_once_logged
            if envs.SGLANG_OPT_MOE_QUANT_ONCE.get() and not _moe_quant_once_logged:
                _moe_quant_once_logged = True
                logger.info(
                    "SGLANG_OPT_MOE_QUANT_ONCE: %s (layer %s)",
                    "ENGAGED" if self._moe_quant_once else f"INELIGIBLE: {reason}",
                    self.layer_id,
                )
        return self._moe_quant_once

    def _compute_moe_quant_once_enabled(self) -> Tuple[bool, str]:
        """Returns (eligible, reason); reason names the first failing check."""
        from sglang.srt.layers.moe.token_dispatcher.standard import StandardDispatcher
        from sglang.srt.layers.quantization.fp8 import Fp8LinearMethod, Fp8MoEMethod
        from sglang.srt.layers.quantization.fp8_utils import (
            cutlass_w8a8_block_fp8_linear_with_fallback,
            deepgemm_w8a8_block_fp8_linear_with_fallback,
        )

        if not envs.SGLANG_OPT_MOE_QUANT_ONCE.get():
            return False, "env off"
        if not _is_cuda:
            return False, "not CUDA"
        if self._enable_a2a_moe or self._fuse_shared_experts_inside_sbo:
            return False, "a2a MoE or SBO shared-expert fusion"
        # Shared-expert side: fp8 block-128 weights served by a w8a8 linear
        # backend taught to accept a pre-quantized (q, scale) tuple: cutlass
        # or deepgemm (fp32 scales only, i.e. not UE8M0/Blackwell).
        if self.num_fused_shared_experts != 0 or not hasattr(self, "shared_experts"):
            return False, "no separate shared experts"
        if not self.shared_experts_is_fp8:
            return False, "shared experts not fp8"
        if self.shared_experts_weight_block_size != [128, 128]:
            return False, "shared weight block size != [128, 128]"
        gate_up = self.shared_experts.gate_up_proj
        if not isinstance(gate_up.quant_method, Fp8LinearMethod):
            return False, "shared gate_up quant method not Fp8LinearMethod"
        linear_fn = gate_up.quant_method.w8a8_block_fp8_linear
        if linear_fn is cutlass_w8a8_block_fp8_linear_with_fallback:
            if gate_up.weight.shape[0] % 128 != 0 or gate_up.weight.shape[1] % 128 != 0:
                return False, "gate_up weight shape unsupported by cutlass"
        elif linear_fn is deepgemm_w8a8_block_fp8_linear_with_fallback:
            from sglang.srt.layers import deep_gemm_wrapper

            if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
                return False, "DeepGEMM UE8M0 scales (Blackwell) unsupported"
            if gate_up.weight.shape[0] % 64 != 0 or gate_up.weight.shape[1] % 128 != 0:
                return False, "gate_up weight shape unsupported by deepgemm"
        else:
            return False, f"w8a8 linear backend {linear_fn.__name__} unsupported"
        # Routed side: standard dispatcher + triton fused func with dynamic
        # per-token-group-128 fp8 activation quant.
        experts = self.experts
        if not isinstance(experts, FusedMoE):
            return False, "experts not FusedMoE"
        quant_method = experts.quant_method
        if not isinstance(quant_method, Fp8MoEMethod):
            return False, "experts quant method not Fp8MoEMethod"
        if not quant_method.block_quant or quant_method.use_mxfp8:
            return False, "experts not block-quant fp8"
        if quant_method.quant_config.weight_block_size != [128, 128]:
            return False, "experts weight block size != [128, 128]"
        # Fp8MoEMethod only sets .runner for runner backends it drives itself.
        runner = getattr(quant_method, "runner", None)
        if runner is None or not runner.runner_backend.is_triton():
            return False, "MoE runner backend not triton"
        if runner.fused_func is None or runner.lora_enabled:
            return False, "triton fused func unavailable (or LoRA enabled)"
        if not isinstance(experts.dispatcher, StandardDispatcher):
            return False, "dispatcher not StandardDispatcher"
        if experts.moe_runner_config.apply_router_weight_on_input:
            return False, "apply_router_weight_on_input"
        if experts.w13_input_scale is not None:
            return False, "static w13 input scale"
        return True, "ok"

    def _maybe_quant_moe_input_once(
        self, hidden_states: torch.Tensor
    ) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
        """Quantize hidden_states once (per-token-group-128 fp8, rows padded to
        a multiple of 4, column-major scales) for both the shared-expert GEMM
        and the routed dispatch, or return None when ineligible."""
        if hidden_states.shape[0] == 0 or hidden_states.dtype != torch.bfloat16:
            return None
        if not self._moe_quant_once_enabled():
            return None
        if is_in_tc_piecewise_cuda_graph():
            # The piecewise MoE op quantizes internally; a pre-quant here
            # would be dead work.
            return None
        from sglang.kernels.ops.quantization.fp8_kernel import (
            sglang_per_token_group_quant_fp8_row_padded,
        )

        q, s = sglang_per_token_group_quant_fp8_row_padded(hidden_states, 128)
        return q, s

    def op_gate(self, state):
        if state.hidden_states_mlp_input.shape[0] > 0:
            # router_logits: (num_tokens, n_experts)
            state.router_logits = self.gate(state.hidden_states_mlp_input)
        else:
            state.router_logits = None

    def op_shared_experts(self, state):
        hidden_states_mlp_input = state.pop("hidden_states_mlp_input")
        if (self.num_fused_shared_experts == 0) and is_non_idle_and_non_empty(
            state.forward_batch.forward_mode, hidden_states_mlp_input
        ):
            state.shared_output = self.shared_experts(hidden_states_mlp_input)
        else:
            state.shared_output = None

    def op_select_experts(self, state):
        router_logits = state.pop("router_logits")
        hidden_states = state.hidden_states_mlp_input

        # Hash MoE layers (e.g. DeepSeek-V4) route on input_ids; forward_deepep
        # passes them as a topk kwarg. The per-ubatch forward_batch.input_ids is
        # already sliced+padded to match hidden_states rows (and equals the
        # global ids under EP dp-attention). No-op for non-hash models.
        topk_kwargs = {}
        if getattr(self, "is_hash", False):
            topk_kwargs["input_ids"] = state.forward_batch.input_ids

        if router_logits is not None:
            with get_global_expert_distribution_recorder().with_current_layer(
                self.layer_id
            ):
                state.topk_output = self.topk(
                    hidden_states=hidden_states,
                    router_logits=router_logits,
                    num_token_non_padded=state.forward_batch.num_token_non_padded,
                    expert_location_dispatch_info=(
                        ExpertLocationDispatchInfo.init_new(
                            layer_id=self.layer_id,
                        )
                        if not self.is_nextn
                        else None
                    ),
                    **topk_kwargs,
                )
        else:
            state.topk_output = self.topk.empty_topk_output(
                hidden_states.device, layer_id=self.layer_id
            )

    def op_dispatch_a(self, state):
        if self.ep_size > 1:
            self.experts.dispatcher.dispatch_a(
                hidden_states=state.hidden_states_mlp_input,
                topk_output=state.pop("topk_output"),
                tbo_subbatch_index=state.get("tbo_subbatch_index"),
            )

    def op_dispatch_b(self, state):
        if self.ep_size > 1:
            with get_global_expert_distribution_recorder().with_current_layer(
                self.layer_id
            ):
                state.dispatch_output = self.experts.dispatcher.dispatch_b(
                    tbo_subbatch_index=state.get("tbo_subbatch_index"),
                )

    def op_experts(self, state):
        state.combine_input = self.experts.run_moe_core(
            dispatch_output=state.dispatch_output,
        )

    def op_combine_a(self, state):
        if self.ep_size > 1:
            self.experts.dispatcher.combine_a(
                combine_input=state.pop("combine_input"),
                tbo_subbatch_index=state.get("tbo_subbatch_index"),
            )
            state.pop("dispatch_output")

    def op_combine_b(self, state):
        if self.ep_size > 1:
            state.hidden_states_after_combine = self.experts.dispatcher.combine_b(
                tbo_subbatch_index=state.get("tbo_subbatch_index"),
            )

    def op_output(self, state):
        final_hidden_states = state.pop("hidden_states_after_combine")

        if get_moe_a2a_backend().is_mori():
            num_tokens = state.pop("num_tokens")
            final_hidden_states = final_hidden_states[:num_tokens]

        if (shared_output := state.pop("shared_output")) is not None:
            x = shared_output
            if _use_aiter:
                x.add_(final_hidden_states)
            else:
                x.add_(final_hidden_states, alpha=self.routed_scaling_factor)
            final_hidden_states = x
        elif _use_aiter:
            # fused in aiter_biased_grouped_topk so we can skip here
            pass
        else:
            final_hidden_states *= self.routed_scaling_factor

        state.hidden_states_mlp_output = final_hidden_states


class DeepseekV2AttentionMLA(
    nn.Module,
    DeepseekMHAForwardMixin,
    DeepseekMHARocmForwardMixin,
    DeepseekMLAForwardMixin,
    DeepseekMLARocmForwardMixin,
    DeepseekMLAFusedRopeRocmForwardMixin,
    DeepseekMLACpuForwardMixin,
):
    def __init__(
        self,
        config: PretrainedConfig,
        hidden_size: int,
        num_heads: int,
        qk_nope_head_dim: int,
        qk_rope_head_dim: int,
        v_head_dim: int,
        q_lora_rank: int,
        kv_lora_rank: int,
        rope_theta: float = 10000,
        rope_scaling: Optional[Dict[str, Any]] = None,
        max_position_embeddings: int = 8192,
        quant_config: Optional[QuantizationConfig] = None,
        reduce_results: bool = True,
        layer_id: int = None,
        prefix: str = "",
        alt_stream: Optional[torch.cuda.Stream] = None,
        skip_rope: bool = False,
        is_nextn: bool = False,
    ) -> None:
        super().__init__()
        self.layer_id = layer_id
        self.hidden_size = hidden_size
        self.qk_nope_head_dim = qk_nope_head_dim
        self.qk_rope_head_dim = qk_rope_head_dim
        self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
        self.v_head_dim = v_head_dim
        self.q_lora_rank = q_lora_rank
        self.kv_lora_rank = kv_lora_rank
        self.quant_config = quant_config
        self.is_nextn = is_nextn
        attn_tp_rank = get_parallel().attn_tp_rank
        attn_tp_size = get_parallel().attn_tp_size
        self.use_dsa = is_deepseek_dsa(config)
        self.num_heads = num_heads
        assert num_heads % attn_tp_size == 0
        self.num_local_heads = num_heads // attn_tp_size
        self.scaling = self.qk_head_dim**-0.5
        self.rope_theta = rope_theta
        self.max_position_embeddings = max_position_embeddings
        self.kv_cache_dtype = get_model().kv_cache_dtype

        # NOTE modification to rope_scaling must be done early enough, b/c e.g. Indexer needs it
        if rope_scaling:
            rope_scaling["rope_type"] = "deepseek_yarn"

        # For tensor parallel attention
        if self.q_lora_rank is not None:
            self.fused_qkv_a_proj_with_mqa = ReplicatedLinear(
                self.hidden_size,
                self.q_lora_rank + self.kv_lora_rank + self.qk_rope_head_dim,
                bias=False,
                quant_config=quant_config,
                prefix=add_prefix("fused_qkv_a_proj_with_mqa", prefix),
            )
            self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)
            self.q_b_proj = ColumnParallelLinear(
                q_lora_rank,
                self.num_heads * self.qk_head_dim,
                bias=False,
                quant_config=self._get_q_b_proj_quant_config(quant_config),
                prefix=add_prefix("q_b_proj", prefix),
                tp_rank=attn_tp_rank,
                tp_size=attn_tp_size,
            )
        else:
            self.q_proj = ColumnParallelLinear(
                self.hidden_size,
                self.num_heads * self.qk_head_dim,
                bias=False,
                quant_config=quant_config,
                prefix=add_prefix("q_proj", prefix),
                tp_rank=attn_tp_rank,
                tp_size=attn_tp_size,
            )
            self.kv_a_proj_with_mqa = ReplicatedLinear(
                self.hidden_size,
                self.kv_lora_rank + self.qk_rope_head_dim,
                bias=False,
                quant_config=quant_config,
                prefix=add_prefix("kv_a_proj_with_mqa", prefix),
            )

        self.skip_topk = None
        self.next_skip_topk = None
        self.indexer = None
        if self.use_dsa:
            # Refer: https://arxiv.org/abs/2603.12201 for more details.
            # skip_topk: when True, this layer will skip computation and reuse previous layer's topk indices.
            # next_skip_topk: when True, the next layer will skip computation and reuse this layer's topk indices.
            if is_nextn:
                self.skip_topk = True
                self.next_skip_topk = True
            else:
                self.skip_topk = dsa_layer_skips_topk(config, layer_id)
                self.next_skip_topk = dsa_layer_skips_topk(config, layer_id + 1)

            if not self.skip_topk or is_nextn:
                is_neox_style = not getattr(config, "indexer_rope_interleave", False)
                indexer_cls = (
                    IndexerKPool if get_dsa_index_kpool(config) > 1 else Indexer
                )
                indexer_kwargs = dict(
                    hidden_size=hidden_size,
                    index_n_heads=get_dsa_index_n_heads(config),
                    index_head_dim=get_dsa_index_head_dim(config),
                    rope_head_dim=qk_rope_head_dim,
                    index_topk=get_dsa_index_topk(config),
                    q_lora_rank=q_lora_rank,
                    max_position_embeddings=max_position_embeddings,
                    rope_theta=rope_theta,
                    scale_fmt="ue8m0",
                    block_size=128,
                    rope_scaling=rope_scaling,
                    is_neox_style=is_neox_style,
                    prefix=add_prefix("indexer", prefix),
                    quant_config=quant_config,
                    layer_id=layer_id,
                    alt_stream=alt_stream,
                    config=config,
                )
                if indexer_cls is IndexerKPool:
                    indexer_kwargs["skip_rope"] = skip_rope
                self.indexer = indexer_cls(**indexer_kwargs)

        self.kv_b_proj = ColumnParallelLinear(
            self.kv_lora_rank,
            self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
            bias=False,
            quant_config=quant_config,
            prefix=add_prefix("kv_b_proj", prefix),
            tp_rank=attn_tp_rank,
            tp_size=attn_tp_size,
        )
        # O projection.
        self.o_proj = RowParallelLinear(
            self.num_heads * self.v_head_dim,
            self.hidden_size,
            bias=False,
            quant_config=quant_config,
            reduce_results=reduce_results,
            prefix=add_prefix("o_proj", prefix),
            tp_rank=attn_tp_rank,
            tp_size=attn_tp_size,
        )
        self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)

        if not skip_rope and qk_rope_head_dim > 0:
            is_neox_style = not getattr(config, "rope_interleave", True)
            self.rotary_emb = get_rope_wrapper(
                qk_rope_head_dim,
                rotary_dim=qk_rope_head_dim,
                max_position=max_position_embeddings,
                base=rope_theta,
                rope_scaling=rope_scaling,
                is_neox_style=is_neox_style,
                device=get_device().device,
            )

            if rope_scaling and rope_scaling.get("apply_yarn_scaling", True):
                self.scaling = compute_mla_mscale_scaling(rope_scaling, self.scaling)
        else:
            self.rotary_emb = None
        self.use_deepseek_yarn_rope = rope_scaling is not None

        self.attn_mqa = RadixAttention(
            self.num_local_heads,
            self.kv_lora_rank + self.qk_rope_head_dim,
            self.scaling,
            num_kv_heads=1,
            layer_id=layer_id,
            v_head_dim=self.kv_lora_rank,
            quant_config=quant_config,
            prefix=add_prefix("attn_mqa", prefix),
        )
        # use num_local_heads * dcp_world_size because q_nope, q_rope is all gathered from dcp ranks
        if get_parallel().dcp_enabled:
            self.attn_mqa_for_dcp_decode = RadixAttention(
                self.num_local_heads * get_parallel().attn_dcp_size,
                self.kv_lora_rank + self.qk_rope_head_dim,
                self.scaling,
                num_kv_heads=1,
                layer_id=layer_id,
                v_head_dim=self.kv_lora_rank,
                quant_config=quant_config,
                prefix=add_prefix("attn_mqa", prefix),
            )

        self.attn_mha = RadixAttention(
            self.num_local_heads,
            self.qk_nope_head_dim + self.qk_rope_head_dim,
            self.scaling,
            num_kv_heads=self.num_local_heads,
            layer_id=layer_id,
            v_head_dim=self.v_head_dim,
            quant_config=quant_config,
            prefix=add_prefix("attn_mha", prefix),
        )

        self.alt_stream = alt_stream
        self.attn_mha.kv_b_proj = None

        self.w_kc = None
        self.w_vc = None
        self.w_scale = 1.0

        # Full-head Q/absorb weights for --dcp-replicate-q-proj, gathered once
        # pre-CUDA-graph-capture by the model runner; None unless replicate is on.
        self.w_kc_qrep = None
        self.q_b_proj_qrep_weight = None

        self.w_scale_k = None
        self.w_scale_v = None
        self.use_deep_gemm_bmm = False

        self.current_attention_backend = (
            None  # Attention backend used by current forward batch
        )

        self.has_fused_proj = hasattr(self, "fused_qkv_a_proj_with_mqa")
        self.is_packed_weight = (
            self.has_fused_proj
            and hasattr(self.fused_qkv_a_proj_with_mqa.quant_method, "quant_config")
            and self.fused_qkv_a_proj_with_mqa.quant_method.quant_config.get_name()
            in {"awq", "awq_marlin", "moe_wna16"}
        )
        self._use_min_latency_fused_a_gemm: bool | None = None
        self.fused_a_gemm_backend = "auto"

        self.has_q_b_proj = hasattr(self, "q_b_proj")
        q_b_proj_verified_shapes = {(2048, 2048), (4096, 2048)}
        q_b_weight = (
            getattr(self.q_b_proj, "weight", None) if self.has_q_b_proj else None
        )
        self._q_b_proj_verified_shape = q_b_weight is not None and (
            tuple(q_b_weight.shape) in q_b_proj_verified_shapes
        )
        self._use_min_latency_q_b_gemm: bool | None = None

        self.init_mha_forward()
        self.init_mla_forward()
        self.init_mla_fused_rope_rocm_forward()
        self.init_mla_fused_rope_cpu_forward()

    @contextmanager
    def maybe_use_decode_attn_tp(self, forward_batch: ForwardBatch):
        if self.q_lora_rank is None:
            yield
            return
        tensor_attrs = [
            (self, "w_kc", 0),
            (self, "w_vc", 0),
            (self, "w_scale_k", 0),
            (self, "w_scale_v", 0),
        ]
        ctx = get_cp_decode_attn_tp_ctx()
        with ctx.maybe_use_decode_attn_tp(
            forward_batch,
            [self.q_b_proj, self.o_proj],
            tensor_attrs=tensor_attrs,
            radix_attn=self.attn_mqa,
        ):
            if ctx.use_decode_attn_tp:
                orig_num_local_heads = self.num_local_heads
                self.num_local_heads = self.num_heads // ctx.decode_tp_size
                try:
                    yield
                finally:
                    self.num_local_heads = orig_num_local_heads
            else:
                yield

    def dispatch_attn_forward_method(
        self, forward_batch: ForwardBatch
    ) -> AttnForwardMethod:
        # Determine attention backend name for current forward batch: prefer the
        # name stamped per-runner on the backend object, else resolve from server args.
        backend = get_attn_backend()
        default_prefill_str, default_decode_str = attention_backends()
        prefill_backend_str = (
            backend.prefill_attention_backend_str or default_prefill_str
        )
        decode_backend_str = backend.decode_attention_backend_str or default_decode_str
        if forward_batch.forward_mode.is_decode_or_idle():
            attention_backend = decode_backend_str
        elif (
            forward_batch.forward_mode.is_target_verify()
            or forward_batch.forward_mode.is_draft_extend_v2()
        ):
            # Use the specified backend for speculative operations (both verify and draft extend)
            if get_spec().speculative_attention_mode == "decode":
                attention_backend = decode_backend_str
            else:  # default to prefill
                attention_backend = prefill_backend_str
        else:
            attention_backend = prefill_backend_str
        self.current_attention_backend = attention_backend

        handler = AttentionBackendRegistry.get_handler(attention_backend)
        return resolve_rocm_forward_method(handler(self, forward_batch))

    def op_prepare(self, state):
        state.attn_intermediate_state = self.forward_prepare(
            positions=state.positions,
            hidden_states=state.pop("hidden_states_after_comm_pre_attn"),
            forward_batch=state.forward_batch,
            zero_allocator=state.zero_allocator,
        )

    def op_core(self, state):
        result = self.forward_core(state.pop("attn_intermediate_state"))
        # forward_core may return (hidden_states, topk_indices) for DSA models
        # with index cache enabled. In the TBO path, topk_indices is not
        # propagated between layers, so we discard it here.
        if isinstance(result, tuple):
            state.hidden_states_after_attn = result[0]
        else:
            state.hidden_states_after_attn = result

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        forward_batch: ForwardBatch,
        zero_allocator: BumpAllocator,
        layer_scatter_modes: LayerScatterModes = None,
        llama_4_scaling: Optional[torch.Tensor] = None,
        prev_topk_indices: Optional[torch.Tensor] = None,
    ):
        s = self.forward_prepare(
            positions=positions,
            hidden_states=hidden_states,
            forward_batch=forward_batch,
            zero_allocator=zero_allocator,
            layer_scatter_modes=layer_scatter_modes,
            llama_4_scaling=llama_4_scaling,
            prev_topk_indices=prev_topk_indices,
        )
        return self.forward_core(s)

    def forward_prepare(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        forward_batch: ForwardBatch,
        zero_allocator: BumpAllocator,
        layer_scatter_modes: LayerScatterModes = None,
        llama_4_scaling: Optional[torch.Tensor] = None,
        prev_topk_indices: Optional[torch.Tensor] = None,
    ):
        if self.attn_mha.kv_b_proj is None and hasattr(self, "kv_b_proj"):
            self.attn_mha.kv_b_proj = self.kv_b_proj

        # when hidden_states is a tuple of tensors, the tuple will include quantized weight and scale tensor
        if isinstance(hidden_states, tuple):
            if (
                not get_attn_tp_context().input_scattered
                and hidden_states[0].shape[0] == 0
            ):
                assert not self.o_proj.reduce_results, (
                    "short-circuiting allreduce will lead to hangs"
                )
                return hidden_states[0]
        else:
            if (
                not get_attn_tp_context().input_scattered
                and hidden_states.shape[0] == 0
            ):
                assert not self.o_proj.reduce_results, (
                    "short-circuiting allreduce will lead to hangs"
                )
                return hidden_states, None, forward_batch, None

        attn_forward_method = self.dispatch_attn_forward_method(forward_batch)
        if attn_forward_method == AttnForwardMethod.MHA:
            inner_state = self.forward_normal_prepare(
                positions, hidden_states, forward_batch, zero_allocator
            )
        elif attn_forward_method == AttnForwardMethod.MHA_CHUNKED_KV:
            inner_state = self.forward_normal_chunked_kv_prepare(
                positions, hidden_states, forward_batch, zero_allocator
            )
        elif attn_forward_method == AttnForwardMethod.MHA_ONE_SHOT:
            inner_state = self.forward_normal_one_shot_prepare(
                positions, hidden_states, forward_batch, zero_allocator
            )
        elif attn_forward_method == AttnForwardMethod.MLA:
            inner_state = self.forward_absorb_prepare(
                positions,
                hidden_states,
                forward_batch,
                zero_allocator,
                llama_4_scaling,
                prev_topk_indices,
            )
        elif attn_forward_method == AttnForwardMethod.MHA_ROCM:
            inner_state = self.forward_normal_rocm_prepare(
                positions, hidden_states, forward_batch, zero_allocator
            )
        elif attn_forward_method == AttnForwardMethod.MHA_ONE_SHOT_ROCM:
            inner_state = self.forward_normal_one_shot_rocm_prepare(
                positions, hidden_states, forward_batch, zero_allocator
            )
        elif attn_forward_method == AttnForwardMethod.MHA_CHUNKED_KV_ROCM:
            inner_state = self.forward_normal_chunked_kv_rocm_prepare(
                positions, hidden_states, forward_batch, zero_allocator
            )
        elif attn_forward_method == AttnForwardMethod.MLA_ROCM:
            inner_state = self.forward_absorb_rocm_prepare(
                positions,
                hidden_states,
                forward_batch,
                zero_allocator,
                llama_4_scaling,
                prev_topk_indices,
            )
        elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE_ROCM:
            inner_state = self.forward_absorb_fused_mla_rope_prepare(
                positions, hidden_states, forward_batch, zero_allocator
            )
        elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE_CPU:
            inner_state = self.forward_absorb_fused_mla_rope_cpu_prepare(
                positions, hidden_states, forward_batch, zero_allocator
            )
        elif attn_forward_method == AttnForwardMethod.MHA_NPU:
            inner_state = forward_mha_prepare_npu(
                self,
                positions,
                hidden_states,
                forward_batch,
                zero_allocator,
                layer_scatter_modes,
            )
        elif attn_forward_method == AttnForwardMethod.MLA_NPU:
            inner_state = forward_mla_prepare_npu(
                self,
                positions,
                hidden_states,
                forward_batch,
                zero_allocator,
                layer_scatter_modes,
            )
        elif attn_forward_method == AttnForwardMethod.DSA_NPU:
            inner_state = forward_dsa_prepare_npu(
                self,
                positions,
                hidden_states,
                forward_batch,
                zero_allocator,
                layer_scatter_modes,
                prev_topk_indices,
            )
        else:
            raise NotImplementedError
        return None, attn_forward_method, forward_batch, inner_state

    def forward_core(self, intermediate_state):
        hidden_states, attn_forward_method, forward_batch, inner_state = (
            intermediate_state
        )
        if inner_state is None:
            return hidden_states

        if attn_forward_method == AttnForwardMethod.MHA:
            return self.forward_normal_core(*inner_state)
        elif attn_forward_method == AttnForwardMethod.MHA_CHUNKED_KV:
            return self.forward_normal_chunked_kv_core(*inner_state)
        elif attn_forward_method == AttnForwardMethod.MHA_ONE_SHOT:
            return self.forward_normal_one_shot_core(*inner_state)
        elif attn_forward_method == AttnForwardMethod.MLA:
            return self.forward_absorb_core(*inner_state)
        elif attn_forward_method == AttnForwardMethod.MHA_ROCM:
            return self.forward_normal_core(*inner_state)
        elif attn_forward_method == AttnForwardMethod.MHA_ONE_SHOT_ROCM:
            return self.forward_normal_one_shot_core(*inner_state)
        elif attn_forward_method == AttnForwardMethod.MHA_CHUNKED_KV_ROCM:
            return self.forward_normal_chunked_kv_core(*inner_state)
        elif attn_forward_method == AttnForwardMethod.MLA_ROCM:
            return self.forward_absorb_rocm_core(*inner_state)
        elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE_ROCM:
            return self.forward_absorb_fused_mla_rope_core(*inner_state)
        elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE_CPU:
            return self.forward_absorb_fused_mla_rope_cpu_core(*inner_state)
        elif attn_forward_method == AttnForwardMethod.MHA_NPU:
            return forward_mha_core_npu(self, *inner_state)
        elif attn_forward_method == AttnForwardMethod.MLA_NPU:
            return forward_mla_core_npu(self, *inner_state)
        elif attn_forward_method == AttnForwardMethod.DSA_NPU:
            return forward_dsa_core_npu(self, *inner_state)
        else:
            raise NotImplementedError

    def prepare_qkv_latent(
        self, hidden_states: torch.Tensor, forward_batch: ForwardBatch
    ):
        assert self.q_lora_rank is not None
        if hasattr(self, "q_a_proj"):
            return torch.cat(
                (
                    self.q_a_proj(hidden_states)[0],
                    self.kv_a_proj_with_mqa(hidden_states)[0],
                ),
                dim=-1,
            )
        if self._use_min_latency_fused_a_gemm is None:
            self._use_min_latency_fused_a_gemm = (
                self.has_fused_proj
                and not get_exec().deterministic.enable_deterministic_inference
                and not self.is_packed_weight
                and fused_a_gemm_weight_eligible(self.fused_qkv_a_proj_with_mqa)
            )
        if self._use_min_latency_fused_a_gemm:
            return linear_with_fused_a_gemm(
                self.fused_qkv_a_proj_with_mqa,
                hidden_states,
                backend=self.fused_a_gemm_backend,
            )
        return self.fused_qkv_a_proj_with_mqa(hidden_states)[0]

    def q_b_proj_forward(self, q_lora: torch.Tensor) -> torch.Tensor:
        if self._use_min_latency_q_b_gemm is None:
            self._use_min_latency_q_b_gemm = (
                self._q_b_proj_verified_shape
                and fused_a_gemm_weight_eligible(self.q_b_proj)
            )
        if self._use_min_latency_q_b_gemm:
            q = linear_with_fused_a_gemm(
                self.q_b_proj, q_lora, backend=self.fused_a_gemm_backend
            )
        else:
            q = self.q_b_proj(q_lora)[0]
        return q.view(-1, self.num_local_heads, self.qk_head_dim)

    @staticmethod
    def _get_q_b_proj_quant_config(quant_config):
        if envs.SGLANG_NVFP4_CKPT_FP8_GEMM_IN_ATTN.get():
            # refer to real DeepSeek V3 quant config
            return Fp8Config(
                is_checkpoint_fp8_serialized=True,
                weight_block_size=[128, 128],
            )
        else:
            return quant_config


class DeepseekV2DecoderLayer(nn.Module):
    def __init__(
        self,
        config: PretrainedConfig,
        layer_id: int,
        quant_config: Optional[QuantizationConfig] = None,
        moe_quant_config_override: Optional[QuantizationConfig] = None,
        is_nextn: bool = False,
        prefix: str = "",
        alt_stream: Optional[torch.cuda.Stream] = None,
        skip_rope: bool = False,
    ) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
        self.config = config
        if hasattr(config, "rope_parameters"):
            rope_theta = config.rope_parameters["rope_theta"]
            assert rope_theta is not None, f"rope_theta not found in config: {config}"
            rope_type = config.rope_parameters.get("rope_type")
            rope_scaling = config.rope_parameters if rope_type != "default" else None
        else:
            rope_theta = config.rope_theta
            rope_scaling = config.rope_scaling
        max_position_embeddings = config.max_position_embeddings
        self.speculative_algorithm = SpeculativeAlgorithm.from_string(
            get_spec().speculative_algorithm
        )
        self.layer_id = layer_id
        self.is_nextn = is_nextn
        if is_nextn and getattr(config, "mla_nope", False):
            # The NextN draft must match the NoPE target layers, or its Q/K
            # and the KV it verifies against live in different spaces.
            skip_rope = True
        self.self_attn = DeepseekV2AttentionMLA(
            config=config,
            hidden_size=self.hidden_size,
            num_heads=config.num_attention_heads,
            qk_nope_head_dim=config.qk_nope_head_dim,
            qk_rope_head_dim=config.qk_rope_head_dim,
            v_head_dim=config.v_head_dim,
            q_lora_rank=(
                config.q_lora_rank if hasattr(config, "q_lora_rank") else None
            ),
            kv_lora_rank=config.kv_lora_rank,
            rope_theta=rope_theta,
            rope_scaling=rope_scaling,
            max_position_embeddings=max_position_embeddings,
            quant_config=quant_config,
            layer_id=layer_id,
            reduce_results=False,
            prefix=add_prefix("self_attn", prefix),
            alt_stream=alt_stream,
            skip_rope=skip_rope,
            is_nextn=is_nextn,
        )
        if not hasattr(config, "q_lora_rank") and envs.SGLANG_USE_AG_AFTER_QLORA.get():
            raise ValueError(
                "SGLANG_USE_AG_AFTER_QLORA only supports the model with q_lora_rank"
            )

        self.is_layer_sparse = self._is_layer_sparse(layer_id, is_nextn=is_nextn)
        is_previous_layer_sparse = self._is_layer_sparse(layer_id - 1, is_nextn=False)
        is_next_layer_sparse = self._is_layer_sparse(layer_id + 1, is_nextn=False)

        self.layer_scatter_modes = LayerScatterModes.init_new(
            layer_id=layer_id,
            num_layers=1 if is_nextn else 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,
        )

        if self.is_layer_sparse:
            self.mlp = DeepseekV2MoE(
                config=config,
                quant_config=moe_quant_config_override or quant_config,
                prefix=add_prefix("mlp", prefix),
                layer_id=self.layer_id,
                alt_stream=alt_stream,
                is_nextn=is_nextn,
            )
        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 = DeepseekV2MLP(
                hidden_size=config.hidden_size,
                intermediate_size=config.intermediate_size,
                hidden_act=config.hidden_act,
                quant_config=quant_config,
                prefix=add_prefix("mlp", prefix),
                tp_rank=mlp_tp_rank,
                tp_size=mlp_tp_size,
                swiglu_limit=getattr(config, "swiglu_limit", None),
            )

        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
        )

        self._gfx95_quant_format = self._detect_gfx95_quant_format()

        communicator_cls = (
            DSACPLayerCommunicator
            if get_parallel().enable_prefill_cp
            else LayerCommunicator
        )
        self.layer_communicator = communicator_cls(
            layer_scatter_modes=self.layer_scatter_modes,
            input_layernorm=self.input_layernorm,
            post_attention_layernorm=self.post_attention_layernorm,
            allow_reduce_scatter=True,
            is_last_layer=(
                is_nextn or (self.layer_id == self.config.num_hidden_layers - 1)
            ),
            qkv_latent_func=self.self_attn.prepare_qkv_latent,
        )

    def _detect_gfx95_quant_format(self) -> str:
        if not _is_gfx95_supported:
            return ""
        proj = getattr(self.self_attn, "fused_qkv_a_proj_with_mqa", None)
        weight = getattr(proj, "weight", None)
        if weight is None:
            return ""
        if weight.dtype == torch.uint8:
            return "mxfp4"
        if weight.dtype == getattr(torch, "float8_e4m3fn", None):
            # Use _is_block_scale_fp8 to distinguish block-scale fp8 (K/128 scale
            # cols, compatible with fused_rms_fp8_group_quant) from per-channel fp8
            # ([N, 1] scale, must use the plain bf16 path).
            # weight_scale may not be reshaped yet at __init__ time — return
            # "fp8_pending" so _resolve_gfx95_quant_format re-checks on first forward.
            weight_scale = getattr(proj, "weight_scale", None)
            if weight_scale is None:
                return "fp8_pending"
            return "fp8" if _is_block_scale_fp8(proj) else ""
        return ""

    def _resolve_gfx95_quant_format(self) -> str:
        """Re-evaluate after weights are loaded if still pending."""
        fmt = getattr(self, "_gfx95_quant_format", "")
        if fmt == "fp8_pending":
            fmt = self._detect_gfx95_quant_format()
            if fmt == "fp8_pending":
                # weight_scale still unavailable — default to bf16 (safe fallback).
                fmt = ""
            self._gfx95_quant_format = fmt
        return fmt

    def _is_layer_sparse(self, layer_id: int, is_nextn: bool) -> bool:
        return is_nextn or (
            self.config.n_routed_experts is not None
            and layer_id >= self.config.first_k_dense_replace
            and layer_id % self.config.moe_layer_freq == 0
        )

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        forward_batch: ForwardBatch,
        residual: Optional[torch.Tensor],
        zero_allocator: BumpAllocator,
        gemm_output_zero_allocator: BumpAllocator = None,
        llama_4_scaling: Optional[torch.Tensor] = None,
        prev_topk_indices: Optional[torch.Tensor] = None,
        captured_last_layer_outputs: Optional[AuxHiddenStateAccumulator] = None,
        next_full_attention_layer_id: Optional[int] = None,
    ) -> torch.Tensor:
        hidden_states_orig = hidden_states
        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,
                quant_format=self._resolve_gfx95_quant_format(),
            )
        )

        with self.self_attn.maybe_use_decode_attn_tp(forward_batch):
            hidden_states = self.self_attn(
                positions=positions,
                hidden_states=hidden_states,
                forward_batch=forward_batch,
                zero_allocator=zero_allocator,
                llama_4_scaling=llama_4_scaling,
                layer_scatter_modes=self.layer_scatter_modes,
                prev_topk_indices=prev_topk_indices,
            )
        if isinstance(hidden_states, tuple):
            hidden_states, topk_indices = hidden_states
        else:
            topk_indices = None
        get_attn_tp_context().clear_attn_inputs()

        maybe_prefetch_next_full_attention_kv(
            forward_batch, next_full_attention_layer_id
        )

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

        fuse_mlp_allreduce = (
            self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
                forward_batch
            )
        )

        # For DP with padding, reduce scatter can be used instead of all-reduce.
        mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
            forward_batch
        )

        if isinstance(self.mlp, DeepseekV2MLP):
            gemm_output_zero_allocator = None

        if (
            isinstance(self.mlp, DeepseekV2MoE)
            and not self.mlp.experts.moe_runner_config.inplace
            and not torch.compiler.is_compiling()
        ):
            from sglang.srt.layers.moe.moe_runner.base import moe_output_buffer_ctx

            _mlp_ctx = moe_output_buffer_ctx(hidden_states_orig)
        else:
            _mlp_ctx = nullcontext()

        with get_forward().scoped(
            fuse_mlp_allreduce=fuse_mlp_allreduce,
            mlp_reduce_scatter=mlp_reduce_scatter,
        ):
            with _mlp_ctx:
                hidden_states = self.mlp(
                    hidden_states,
                    forward_batch,
                    gemm_output_zero_allocator,
                )

        if fuse_mlp_allreduce:
            hidden_states._sglang_needs_allreduce_fusion = True

        if not fuse_mlp_allreduce:
            hidden_states, residual = self.layer_communicator.postprocess_layer(
                hidden_states, residual, forward_batch
            )

        return hidden_states, residual, topk_indices

    def op_comm_prepare_attn(
        self,
        state,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        forward_batch: ForwardBatch,
        residual: Optional[torch.Tensor],
        zero_allocator: BumpAllocator,
        tbo_subbatch_index: Optional[int] = None,
    ):
        state.hidden_states_after_comm_pre_attn, state.residual_after_input_ln = (
            self.layer_communicator.prepare_attn(hidden_states, residual, forward_batch)
        )
        if get_moe_a2a_backend().is_mori():
            state.num_tokens = hidden_states.shape[0]
        state.update(
            dict(
                forward_batch=forward_batch,
                positions=positions,
                zero_allocator=zero_allocator,
                tbo_subbatch_index=tbo_subbatch_index,
            )
        )

    def op_comm_prepare_mlp(self, state):
        state.hidden_states_mlp_input, state.residual_after_comm_pre_mlp = (
            self.layer_communicator.prepare_mlp(
                state.pop("hidden_states_after_attn"),
                state.pop("residual_after_input_ln"),
                state.forward_batch,
            )
        )

    def op_comm_postprocess_layer(self, state):
        hidden_states, residual = self.layer_communicator.postprocess_layer(
            state.pop("hidden_states_mlp_output"),
            state.pop("residual_after_comm_pre_mlp"),
            state.forward_batch,
        )

        output = dict(
            positions=state.positions,
            hidden_states=hidden_states,
            residual=residual,
            forward_batch=state.forward_batch,
            zero_allocator=state.zero_allocator,
            tbo_subbatch_index=state.tbo_subbatch_index,
        )

        state.clear(
            expect_keys={
                "positions",
                "forward_batch",
                "zero_allocator",
                "tbo_subbatch_index",
            }
        )
        return output


class DeepseekV2Model(nn.Module):
    fall_back_to_pt_during_load = False

    def __init__(
        self,
        config: PretrainedConfig,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.config = config
        self.use_dsa = is_deepseek_dsa(config)
        self.padding_id = config.pad_token_id
        self.vocab_size = config.vocab_size
        self.first_k_dense_replace = config.first_k_dense_replace
        self.pp_group = get_pp_group()

        if self.pp_group.is_first_rank or (_is_npu and self.pp_group.is_last_rank):
            self.embed_tokens = VocabParallelEmbedding(
                config.vocab_size,
                config.hidden_size,
                **get_embedding_tp_kwargs(),
            )
        else:
            self.embed_tokens = PPMissingLayer()

        self.alt_stream = (
            torch.cuda.Stream()
            if (
                _is_cuda
                or _is_musa
                or envs.SGLANG_NPU_USE_MULTI_STREAM.get()
                or envs.SGLANG_ROCM_USE_MULTI_STREAM.get()
            )
            else None
        )

        self.layers, self.start_layer, self.end_layer = make_layers(
            config.num_hidden_layers,
            lambda idx, prefix: DeepseekV2DecoderLayer(
                config=config,
                layer_id=idx,
                quant_config=quant_config,
                prefix=prefix,
                alt_stream=self.alt_stream,
                skip_rope=config.qk_rope_head_dim == 0,
            ),
            pp_rank=self.pp_group.rank_in_group,
            pp_size=self.pp_group.world_size,
            prefix=add_prefix("layers", prefix),
            offloader_kwargs=dict(
                submodule_accessor=lambda layer: (
                    layer.mlp.experts
                    if isinstance(layer.mlp, DeepseekV2MoE)
                    else layer.mlp
                ),
                whitelist_param_names_creator=lambda module: (
                    [
                        "w13_weight",
                        "w2_weight",
                        # only for nvfp4
                        *(
                            [
                                "w13_blockscale_swizzled",
                                "w2_blockscale_swizzled",
                            ]
                            if hasattr(module, "w13_blockscale_swizzled")
                            else []
                        ),
                    ]
                    if isinstance(module, FusedMoE)
                    else []
                ),
            ),
        )

        local_layer_ids = list(range(self.start_layer, self.end_layer))
        self.next_full_attention_layer_id = dict(
            zip(local_layer_ids, local_layer_ids[1:])
        )
        if self.pp_group.is_last_rank:
            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        else:
            self.norm = PPMissingLayer(return_tuple=True)

        self.gemm_output_zero_allocator_size = 0
        if (
            _use_aiter_gfx95
            and config.n_routed_experts == 256
            and self.embed_tokens.embedding_dim == 7168
        ):
            num_moe_layers = sum(
                [
                    1
                    for i in range(len(self.layers))
                    if isinstance(self.layers[i].mlp, DeepseekV2MoE)
                ]
            )

            allocate_size = 0
            for i in range(len(self.layers)):
                if isinstance(self.layers[i].mlp, DeepseekV2MoE):
                    # tp_size = get_parallel().tp_size
                    # Keep the original deepep-class scope here and only add DeepEP v2,
                    # so unrelated backends' allocator sizing is unchanged.
                    is_a2a_moe = (
                        is_deepep_class_backend()
                        or get_moe_a2a_backend().is_deepep_v2()
                    )
                    tp_size = 1 if is_a2a_moe else get_parallel().tp_size
                    intermediate_size = (
                        config.moe_intermediate_size * config.n_shared_experts
                    )
                    share_expert_output_size_per_partition = divide(
                        intermediate_size * 2, tp_size
                    )
                    allocate_size = share_expert_output_size_per_partition
                    break

            self.gemm_output_zero_allocator_size = (
                get_dsv3_gemm_output_zero_allocator_size(
                    config.n_routed_experts,
                    num_moe_layers,
                    allocate_size,
                    self.embed_tokens.embedding_dim,
                )
            )
        self.layers_to_capture = []
        self.enable_a2a_moe = (
            get_moe_a2a_backend().is_deepep()
            or get_moe_a2a_backend().is_mooncake()
            or get_moe_a2a_backend().is_deepep_v2()
        )

        # llama_4_scaling: for supporting Mistral-Large-3 model
        self.llama_4_scaling_config = getattr(config, "llama_4_scaling", None)

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

    def _dsa_forward_uses_topk(self) -> bool:
        if not self.use_dsa:
            return False
        backend = get_attn_backend()
        backend = getattr(backend, "primary", backend)
        return not getattr(backend, "use_mha", False)

    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]:
        total_num_layers = self.end_layer - self.start_layer
        dsa_forward_uses_topk = self._dsa_forward_uses_topk()
        if self.pp_group.is_first_rank:
            if input_embeds is None:
                hidden_states = self.embed_tokens(input_ids)
            else:
                hidden_states = input_embeds
            residual = None
            initial_topk_indices = None
        else:
            assert pp_proxy_tensors is not None
            hidden_states = pp_proxy_tensors["hidden_states"]
            residual = pp_proxy_tensors["residual"]
            initial_topk_indices = pp_proxy_tensors.tensors.get("topk_indices")
        index_topk_share = IndexTopKShareState(forward_batch, initial_topk_indices)
        if not self.pp_group.is_first_rank:
            assert not (
                not forward_batch.forward_mode.is_idle()
                and hidden_states.shape[0] != 0
                and self.use_dsa
                and dsa_forward_uses_topk
                and dsa_layer_skips_topk(self.config, self.start_layer)
                and index_topk_share.topk_indices is None
            ), (
                f"PP stage starting at layer {self.start_layer} requires DSA "
                "topk_indices from the previous stage."
            )
        device = hidden_states.device
        zero_allocator = BumpAllocator(
            buffer_size=total_num_layers * 2 * (2 if forward_batch.can_run_tbo else 1),
            dtype=torch.float32,
            device=device,
        )

        has_gemm_output_zero_allocator = hasattr(
            self, "gemm_output_zero_allocator_size"
        )

        gemm_output_zero_allocator = (
            BumpAllocator(
                buffer_size=self.gemm_output_zero_allocator_size,
                dtype=torch.float32,
                device=device,
            )
            if has_gemm_output_zero_allocator
            and self.gemm_output_zero_allocator_size > 0
            else None
        )

        # llama_4_scaling: for supporting Mistral-Large-3 model
        # Compute llama 4 scaling once per forward pass if enabled
        llama_4_scaling: Optional[torch.Tensor] = None
        if self.llama_4_scaling_config is not None:
            llama_4_scaling = _get_llama_4_scaling(
                original_max_position_embeddings=self.llama_4_scaling_config[
                    "original_max_position_embeddings"
                ],
                scaling_beta=self.llama_4_scaling_config["beta"],
                positions=positions,
            )

        normal_start_layer = self.start_layer
        normal_end_layer = self.end_layer
        if forward_batch.can_run_tbo:
            if (
                self.first_k_dense_replace > normal_start_layer
                and self.first_k_dense_replace < normal_end_layer
            ):
                normal_end_layer = self.first_k_dense_replace
            elif self.first_k_dense_replace < normal_start_layer:
                normal_end_layer = normal_start_layer = 0
        # Append-compatible, so the shared capture path below is unchanged.
        aux_hidden_states = AuxHiddenStatePacker(len(self.layers_to_capture))
        for i in range(normal_start_layer, normal_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, topk_indices = layer(
                    positions,
                    hidden_states,
                    forward_batch,
                    residual,
                    zero_allocator,
                    gemm_output_zero_allocator,
                    llama_4_scaling,
                    prev_topk_indices=index_topk_share.topk_indices,
                    captured_last_layer_outputs=(
                        aux_hidden_states if i in self.layers_to_capture else None
                    ),
                    next_full_attention_layer_id=self.next_full_attention_layer_id.get(
                        i
                    ),
                )
                index_topk_share.update(topk_indices)

        if normal_end_layer != self.end_layer:
            hidden_states, residual = model_forward_maybe_tbo(
                layers=self.layers[normal_end_layer : self.end_layer],
                enable_tbo=True,
                positions=positions,
                forward_batch=forward_batch,
                hidden_states=hidden_states,
                residual=residual,
                input_data_scatter_mode=self.layers[
                    normal_end_layer - 1
                ].layer_scatter_modes.layer_output_mode,
                zero_allocator=zero_allocator,
            )

        if not self.pp_group.is_last_rank:
            proxy_tensors = {
                "hidden_states": hidden_states,
                "residual": residual,
            }
            if (
                self.use_dsa
                and dsa_forward_uses_topk
                and self.end_layer < self.config.num_hidden_layers
                and dsa_layer_skips_topk(self.config, self.end_layer)
            ):
                topk_indices = index_topk_share.topk_indices
                if (
                    not forward_batch.forward_mode.is_idle()
                    and hidden_states.shape[0] != 0
                ):
                    assert topk_indices is not None, (
                        f"PP stage ending at layer {self.end_layer} must forward "
                        "DSA topk_indices because the next stage starts on a "
                        "skip-topk layer."
                    )
                if topk_indices is None:
                    topk_indices = hidden_states.new_empty(
                        (0, get_dsa_index_topk(self.config)), dtype=torch.int32
                    )
                proxy_tensors["topk_indices"] = topk_indices
            return PPProxyTensors(proxy_tensors)
        else:
            if not forward_batch.forward_mode.is_idle():
                if residual is None:
                    hidden_states = self.norm(hidden_states)
                else:
                    hidden_states, _ = self.norm(hidden_states, residual)

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


class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
    # for quark model load
    packed_modules_mapping = {}

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

        # for quark model load
        # Fuse q_a_proj and kv_a_proj_with_mqa along output dimension when q_lora_rank is not None
        self.fuse_qkv_a_proj = (
            hasattr(config, "q_lora_rank") and config.q_lora_rank is not None
        )
        if self.fuse_qkv_a_proj:
            self.packed_modules_mapping["fused_qkv_a_proj_with_mqa"] = [
                "q_a_proj",
                "kv_a_proj_with_mqa",
            ]

        # Quant configs like Quark may rely on the model to provide fused-module
        # mappings so exclusion checks can unfuse derived names back to the
        # checkpoint's source layer names.
        if quant_config is not None:
            quant_config.update_packed_modules_mapping(self.packed_modules_mapping)

        self.pp_group = get_pp_group()
        self.config = config
        self.tp_size = get_parallel().tp_size
        self.quant_config = quant_config
        self.determine_num_fused_shared_experts()
        self.use_dsa = is_deepseek_dsa(config)
        self.model = DeepseekV2Model(
            config, quant_config, prefix=add_prefix("model", prefix)
        )

        if self.pp_group.is_last_rank:
            if self.pp_group.world_size == 1 and config.tie_word_embeddings:
                self.lm_head = self.model.embed_tokens
            else:
                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_parallel().enable_dp_lm_head,
                )
        else:
            # ranks other than the last rank will have a placeholder layer
            self.lm_head = PPMissingLayer()
        self.logits_processor = LogitsProcessor(config)

        self._routed_experts_weights_of_layer = LazyValue(
            lambda: {
                layer_id: self.model.layers[layer_id].mlp.get_moe_weights()
                for layer_id in range(self.model.start_layer, self.model.end_layer)
                if isinstance(self.model.layers[layer_id].mlp, DeepseekV2MoE)
            }
        )
        self.capture_aux_hidden_states = False

        q_lora_rank = config.q_lora_rank if hasattr(config, "q_lora_rank") else None
        get_attn_tp_context().init_context(q_lora_rank, is_deepseek_dsa(config))

    @property
    def routed_experts_weights_of_layer(self):
        return self._routed_experts_weights_of_layer.value

    # The architecture this class fuses shared experts for; a subclass whose
    # checkpoint reports a different name (the NextN drafts, GLM's DSA variant)
    # overrides it.
    fused_shared_experts_architecture = "DeepseekV3ForCausalLM"

    @classmethod
    def shared_experts_fusion_disable_reason(cls, hf_config, quant_config):
        """Why this checkpoint cannot fuse its shared expert, or None.

        Evaluated by the loader once per runner, before any layer is built (see
        ``install_shared_experts_fusion_decision``), so it takes the config and
        quantization it is asked about rather than reading an instance.
        """
        # Need to disable if quant precision mismatch, even if
        # --enforce-shared-experts-fusion is specified
        if quant_blocks_shared_experts_fusion(quant_config):
            return (
                "Quantization keeps shared experts at a higher precision than the "
                "routed experts, so they cannot be fused into the quantized "
                "routed-expert path."
            )
        if get_exec().moe.enforce_shared_experts_fusion:
            return None
        if (
            quant_config is not None
            and quant_config.get_name() == "modelopt_fp4"
            and is_sm90_supported()
            and get_moe_runner_backend().is_marlin()
        ):
            return (
                "Hopper modelopt_fp4 with moe_runner_backend=marlin: "
                "fusion off by default until the shared-expert fused load path "
                "is validated."
            )
        if is_sbo_enabled() or is_tbo_enabled():
            return "SBO/TBO enabled: incompatible with fusing shared expert into MoE kernel."
        if is_deepep_class_backend():
            return "DeepEP: fusion off by default (use --enforce-shared-experts-fusion to enable)."
        if (
            hf_config.architectures[0] != cls.fused_shared_experts_architecture
            # Allow-list of n_routed_experts values that have been validated
            # for shared-experts fusion under this code path. Currently:
            #   256 -> DeepSeek-V3 / R1
            #   384 -> Kimi-K2.5, only when the checkpoint is Quark MXFP4
            #          (amd/Kimi-K2.5-MXFP4); the standard
            #          moonshotai/Kimi-K2.5 (compressed-tensors) checkpoint
            #          stores the shared expert loose and is NOT pre-fused,
            #          so the fused path silently mis-loads it.
            or hf_config.n_routed_experts not in (256, 384)
            or hf_config.n_shared_experts != 1
            or (
                hf_config.n_routed_experts == 384
                and (quant_config is None or quant_config.get_name() != "quark")
            )
        ):
            return "Config does not support fused shared expert(s)."
        if (
            (not _is_cuda or torch.cuda.get_device_capability("cuda") < (8, 0))
            and (not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4))
            and (not _is_musa or torch.musa.get_device_capability("musa") < (3, 1))
        ):
            return (
                "Only Deepseek V3/R1 on NV-platform with capability >= 80 "
                "or AMD-platform with capability >= gfx942(MI30x) can use shared experts fusion optimization."
                "or MT-platform with capability >= 31 can use shared experts fusion optimization."
            )
        if get_parallel().moe_ep_size > 1 and (
            not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4)
        ):
            return (
                "Only Deepseek V3/R1 on AMD-platform with capability >= gfx942(MI30x) "
                "can use shared experts fusion optimization under expert parallelism."
            )
        if is_wint4afp8_or_wint4a16_config(quant_config):
            return "Deepseek V3/R1 W4AFP8/W4A16 model uses different quant method for routed experts and shared experts."
        return None

    def determine_num_fused_shared_experts(self):
        # The decision was installed by the loader; this only reads it.
        self.num_fused_shared_experts = (
            0 if is_shared_experts_fusion_disabled() else self.config.n_shared_experts
        )

    def get_input_embeddings(self) -> nn.Embedding:
        return self.model.embed_tokens

    @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:
        with get_attn_tp_context().maybe_input_scattered(forward_batch):
            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]], is_nextn=False):
        self.do_load_weights(weights, is_nextn)

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

    def set_embed_and_head(self, embed, head):
        del self.model.embed_tokens.weight
        del self.lm_head.weight
        self.model.embed_tokens.weight = embed
        self.lm_head.weight = head
        torch.cuda.empty_cache()
        torch.cuda.synchronize()

    @classmethod
    def get_model_config_for_expert_location(cls, config):
        return ModelConfigForExpertLocation(
            num_layers=config.num_hidden_layers,
            num_logical_experts=config.n_routed_experts,
            num_groups=config.n_group,
        )

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

        if layer_ids is None:
            self.capture_aux_hidden_states = True
            num_layers = self.config.num_hidden_layers
            self.model.layers_to_capture = [2, num_layers // 2, num_layers - 3]
        else:
            self.capture_aux_hidden_states = True
            # TODO (Qiaolin-Yu): check if other draft models need similar layer id
            # adjustment
            if layer_ids and layer_ids[0] == 1:
                self.model.layers_to_capture = [val + 1 for val in layer_ids]
            else:
                self.model.layers_to_capture = list(layer_ids)

    def set_dflash_layers_to_capture(self, layer_ids: List[int]):
        if not self.pp_group.is_last_rank:
            return

        if layer_ids is None:
            raise ValueError(
                "DFLASH requires explicit layer_ids for aux hidden capture."
            )

        self.capture_aux_hidden_states = True
        self.model.layers_to_capture = [val + 1 for val in layer_ids]

    def prepare_context_parallel_metadata_for_dcp(
        self,
        seq_lens: torch.Tensor,
        extend_prefix_lens: torch.Tensor,
        extend_prefix_lens_cpu: torch.Tensor,
        extend_seq_lens: torch.Tensor,
        req_pool_indices: torch.Tensor,
        req_to_token: torch.Tensor,
        seq_lens_sum: int,
        kv_buffer_shape: torch.Size,
        kv_cache_dtype,
        kv_cache_device,
        create_chunked_prefix_cache_kv_indices_fn,
    ):
        return prepare_decode_context_parallel_metadata(
            seq_lens=seq_lens,
            extend_prefix_lens=extend_prefix_lens,
            extend_prefix_lens_cpu=extend_prefix_lens_cpu,
            extend_seq_lens=extend_seq_lens,
            req_pool_indices=req_pool_indices,
            req_to_token=req_to_token,
            seq_lens_sum=seq_lens_sum,
            kv_buffer_shape=kv_buffer_shape,
            kv_cache_dtype=kv_cache_dtype,
            kv_cache_device=kv_cache_device,
            create_chunked_prefix_cache_kv_indices_fn=create_chunked_prefix_cache_kv_indices_fn,
        )


class DeepseekV3ForCausalLM(DeepseekV2ForCausalLM):
    pass


class DeepseekV32ForCausalLM(DeepseekV2ForCausalLM):
    pass


@register_custom_op(out_shape="hidden_states")
def dsv2_flashinfer_moe_dual_stream_graph(
    hidden_states: torch.Tensor,
    layer_id: int,
    fuse_mlp_allreduce: bool,
    mlp_reduce_scatter: bool,
) -> torch.Tensor:
    forward_context = get_tc_piecewise_forward_context()
    assert forward_context is not None
    assert forward_context.moe_fusions is not None

    moe_fusion = forward_context.moe_fusions[layer_id]
    assert moe_fusion is not None
    # Custom-op execution happens outside the caller's Python scope under
    # torch.compile. Carry graph-varying control state as scalar operands and
    # republish it for the nested MoE/linear consumers.
    with get_forward().scoped(
        fuse_mlp_allreduce=fuse_mlp_allreduce,
        mlp_reduce_scatter=mlp_reduce_scatter,
        flashinfer_trtllm_bypass=True,
    ):
        return moe_fusion.forward_normal_dual_stream(hidden_states)


EntryClass = [DeepseekV2ForCausalLM, DeepseekV3ForCausalLM, DeepseekV32ForCausalLM]