lightseekorg--tokenspeed/python/tokenspeed/runtime/models/deepseek_v3.py

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

"""Inference-only DeepseekV3 model."""

# ruff: noqa: E402

from __future__ import annotations

import re
from collections.abc import Iterable
from dataclasses import replace
from typing import Any

import torch
import torch.nn.functional as F
from tokenspeed_kernel.ops.attention import attn_merge_state
from tokenspeed_kernel.ops.attention.tokenspeed_mla import mla_kv_pack_quantize_fp8
from tokenspeed_kernel.ops.embedding import apply_rope_mla
from tokenspeed_kernel.ops.gemm.cute_dsl import (
    nvfp4_gemm_swiglu_nvfp4_quant,
)
from tokenspeed_kernel.ops.gemm.trtllm import dsv3_fused_a_gemm
from tokenspeed_kernel.ops.quantization.flashinfer import fp4_quantize
from tokenspeed_kernel.ops.quantization.triton import fp8_quantize
from tokenspeed_kernel.platform import current_platform
from tokenspeed_kernel.thirdparty.cuda import dsv3_router_gemm, moe_finalize_fuse_shared
from torch import nn
from transformers import PretrainedConfig

from tokenspeed.runtime.configs.utils import get_rope_theta
from tokenspeed.runtime.layers.moe import (
    ExpertCheckpointSchema,
    build_moe_checkpoint_loader,
)
from tokenspeed.runtime.layers.utils import (
    CP_METADATA,
    ENABLE_CP,
    cp_all_gather_rerange_output,
    cp_split_and_rebuild_data,
    get_layer_id,
)

_platform = current_platform()
_is_amd = _platform.is_amd
_is_blackwell = _platform.is_blackwell
_is_hopper_plus = _platform.is_hopper_plus
_device_sm = _platform.arch_version.major * 10 + _platform.arch_version.minor

from tokenspeed.runtime.distributed import Mapping
from tokenspeed.runtime.distributed.comm_manager import CommManager
from tokenspeed.runtime.execution.breakable_cuda_graph import (
    break_point,
    scrub_padding_tail,
)
from tokenspeed.runtime.execution.context import (
    ForwardContext,
    report_collective_sizing,
)
from tokenspeed.runtime.execution.cuda_graph_wrapper import get_is_capture_mode
from tokenspeed.runtime.execution.forward_batch_info import ForwardMode
from tokenspeed.runtime.layers.activation import SiluAndMul
from tokenspeed.runtime.layers.dense.nvfp4 import Nvfp4LinearMethod
from tokenspeed.runtime.layers.layernorm import FusedRMSNorm, RMSNorm
from tokenspeed.runtime.layers.linear import (
    ColumnParallelLinear,
    MergedColumnParallelLinear,
    ReplicatedLinear,
    RowParallelLinear,
)
from tokenspeed.runtime.layers.logits_processor import LogitsProcessor
from tokenspeed.runtime.layers.moe.expert import MoELayer
from tokenspeed.runtime.layers.moe.topk import TopK
from tokenspeed.runtime.layers.moe.utils import RoutingMethodType
from tokenspeed.runtime.layers.paged_attention import PagedAttention
from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig
from tokenspeed.runtime.layers.quantization.nvfp4 import Nvfp4Config
from tokenspeed.runtime.layers.quantization.utils import (
    block_dequant,
    should_exclude_quant_module,
)
from tokenspeed.runtime.layers.rotary_embedding import get_rope
from tokenspeed.runtime.layers.vocab_parallel_embedding import (
    ParallelLMHead,
    VocabParallelEmbedding,
)
from tokenspeed.runtime.model_loader.weight_utils import (
    default_weight_loader,
    kv_cache_scales_loader,
)
from tokenspeed.runtime.models.base import BaseCausalLM
from tokenspeed.runtime.models.utils import (
    create_fused_set_kv_buffer_arg,
)
from tokenspeed.runtime.moe.distribution_recorder import (
    get_global_expert_distribution_recorder,
)
from tokenspeed.runtime.moe.expert_location import ModelConfigForExpertLocation
from tokenspeed.runtime.utils import LazyValue, add_prefix, get_colorful_logger
from tokenspeed.runtime.utils.cuda_stream import StreamFork
from tokenspeed.runtime.utils.env import envs, global_server_args_dict
from tokenspeed.runtime.utils.pdl import pdl_enabled

logger = get_colorful_logger(__name__)

_OPTIONAL_MISSING_WEIGHT_SUFFIXES = (
    ".k_scale",
    ".v_scale",
)


def _prepare_mla_kv_b_proj_weights(
    w: torch.Tensor, self_attn
) -> tuple[torch.Tensor, torch.Tensor]:
    w_kc, w_vc = w.unflatten(
        0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim)
    ).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1)
    if _is_amd:
        return w_kc.contiguous(), w_vc.transpose(1, 2).contiguous()
    return (
        w_kc.transpose(1, 2).contiguous().transpose(1, 2),
        w_vc.contiguous().transpose(1, 2),
    )


class DeepseekV3MLP(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
        mapping: Mapping,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
        is_shared_expert: bool = False,
    ) -> None:
        super().__init__()
        self.mapping = mapping
        if is_shared_expert:
            tp_rank = self.mapping.moe.tp_ep_rank
            tp_size = self.mapping.moe.tp_ep_size
            tp_group = self.mapping.moe.tp_ep_group
        else:
            tp_rank = self.mapping.dense.tp_rank
            tp_size = self.mapping.dense.tp_size
            tp_group = self.mapping.dense.tp_group

        self.gate_up_proj = MergedColumnParallelLinear(
            hidden_size,
            [intermediate_size] * 2,
            bias=False,
            tp_size=tp_size,
            tp_rank=tp_rank,
            tp_group=tp_group,
            quant_config=quant_config,
            prefix=add_prefix("gate_up_proj", prefix),
        )
        self.down_proj = RowParallelLinear(
            intermediate_size,
            hidden_size,
            bias=False,
            reduce_results=False,  # Communication is handled externally and manually controlled
            tp_size=tp_size,
            tp_rank=tp_rank,
            tp_group=tp_group,
            quant_config=quant_config,
            prefix=add_prefix("down_proj", prefix),
        )
        if hidden_act != "silu":
            raise ValueError(
                f"Unsupported activation: {hidden_act}. Only silu is supported for now."
            )
        self.act_fn = SiluAndMul()
        self._use_nvfp4_gemm_swiglu_nvfp4_quant = (
            envs.TOKENSPEED_NVFP4_GEMM_SWIGLU_NVFP4_QUANT.get()
            and _is_blackwell
            and isinstance(self.gate_up_proj.quant_method, Nvfp4LinearMethod)
            and isinstance(self.down_proj.quant_method, Nvfp4LinearMethod)
        )
        self.gate_up_proj.interleave_linear_and_gate = (
            self._use_nvfp4_gemm_swiglu_nvfp4_quant
        )

    def forward(self, x):
        if x.size(0) == 0:
            return x

        if self._use_nvfp4_gemm_swiglu_nvfp4_quant:
            x_fc1_fp4, x_fc1_scale = fp4_quantize(
                x, self.gate_up_proj.input_scale_inv, enable_pdl=pdl_enabled()
            )
            x_fp4, x_scale = nvfp4_gemm_swiglu_nvfp4_quant(
                x_fc1_fp4,
                x_fc1_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=pdl_enabled(),
            )
            x, _ = self.down_proj((x_fp4, x_scale))
            return x

        gate_up, _ = self.gate_up_proj(x)
        x = self.act_fn(gate_up)
        x, _ = self.down_proj(x)
        return x


class MoEGate(nn.Module):
    _DSV3_ROUTER_GEMM_EXPERTS = (256, 384, 768)
    _DSV3_ROUTER_GEMM_HIDDEN = (3072, 6144, 7168)

    def __init__(self, config, prefix: str = ""):
        super().__init__()
        self.weight = nn.Parameter(
            torch.empty((config.n_routed_experts, config.hidden_size))
        )
        if config.topk_method == "noaux_tc":
            self.e_score_correction_bias = nn.Parameter(
                torch.empty((config.n_routed_experts), dtype=torch.float32)
            )
        else:
            self.e_score_correction_bias = None

        self.use_dsv3_router_gemm = (
            _is_hopper_plus
            and self.weight.dtype in (torch.bfloat16, torch.float32)
            and config.n_routed_experts in self._DSV3_ROUTER_GEMM_EXPERTS
            and config.hidden_size in self._DSV3_ROUTER_GEMM_HIDDEN
        )

    def forward(self, hidden_states, comm_manager=None):
        if self.use_dsv3_router_gemm and hidden_states.size(0) > 0:
            logits = dsv3_router_gemm(
                hidden_states,
                self.weight,
                out_dtype=torch.float32,
                enable_pdl=pdl_enabled(),
            )
        else:
            logits = F.linear(hidden_states, self.weight, None)
        return logits


class DeepseekV3MoE(nn.Module):
    def __init__(
        self,
        config: PretrainedConfig,
        mapping: Mapping,
        quant_config: QuantizationConfig | None = None,
        layer_index: int = -1,
        prefix: str = "",
        alt_stream: torch.cuda.Stream | None = None,
    ):
        super().__init__()
        self.mapping = mapping
        self.layer_index = layer_index
        self.n_shared_experts = config.n_shared_experts
        self.routed_scaling_factor = config.routed_scaling_factor
        self.stream_fork = StreamFork(alt_stream)

        if self.mapping.moe.ep_size > config.n_routed_experts:
            raise ValueError(
                f"EP size {self.mapping.moe.ep_size} is greater than 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, prefix=add_prefix("gate", prefix))

        if config.n_shared_experts is not None:
            intermediate_size = config.moe_intermediate_size * config.n_shared_experts
            self.shared_experts = DeepseekV3MLP(
                hidden_size=config.hidden_size,
                intermediate_size=intermediate_size,
                hidden_act=config.hidden_act,
                mapping=self.mapping,
                quant_config=quant_config,
                prefix=add_prefix("shared_experts", prefix),
                is_shared_expert=True,
            )

        self.experts = MoELayer(
            top_k=config.num_experts_per_tok,
            num_experts=config.n_routed_experts
            + global_server_args_dict["ep_num_redundant_experts"],
            hidden_size=config.hidden_size,
            intermediate_size=config.moe_intermediate_size,
            quant_config=quant_config,
            layer_index=layer_index,
            prefix=prefix,
            tp_rank=self.mapping.moe.tp_rank,
            tp_size=self.mapping.moe.tp_size,
            ep_rank=self.mapping.moe.ep_rank,
            ep_size=self.mapping.moe.ep_size,
            routing_config={
                "n_group": getattr(config, "n_group", 0),
                "topk_group": getattr(config, "topk_group", 0),
                "routed_scaling_factor": getattr(config, "routed_scaling_factor", 1.0),
                "normalize_topk_weights": config.norm_topk_prob,
                "correction_bias": self.gate.e_score_correction_bias,
                "routing_method_type": RoutingMethodType.DeepSeekV3,
            },
        )

        self.topk = TopK(
            top_k=config.num_experts_per_tok,
            renormalize=config.norm_topk_prob,
            use_grouped_topk=True,
            num_expert_group=config.n_group,
            num_fused_shared_experts=0,
            topk_group=config.topk_group,
            correction_bias=self.gate.e_score_correction_bias,
            routed_scaling_factor=self.routed_scaling_factor,
            output_format=self.experts.topk_output_format,
        )

    def get_moe_routed_weights(self):
        return [
            x.data
            for name, x in self.experts.named_parameters()
            if name not in ["correction_bias"] and "shared_experts" not in name
        ]

    def forward(
        self,
        hidden_states: torch.Tensor,
        num_global_tokens: int,
        max_num_tokens_per_gpu: int,
    ) -> torch.Tensor:
        num_tokens = hidden_states.size(0)

        with self.stream_fork.scope(enable=get_is_capture_mode()) as fork:
            # router_logits: (num_tokens, n_experts)
            router_logits = self.gate(hidden_states)
            if num_tokens > 0:
                topk_output = self.topk(hidden_states, router_logits)
            else:
                topk_output = self.topk.empty_topk_output(
                    hidden_states.device,
                    hidden_states=hidden_states,
                    router_logits=router_logits,
                )

            deferred_finalize = self.experts.supports_deferred_finalize
            routed_expert_output = self.experts(
                hidden_states=hidden_states,
                topk_output=topk_output,
                num_global_tokens=num_global_tokens,
                max_num_tokens_per_gpu=max_num_tokens_per_gpu,
                do_finalize=not deferred_finalize,
            )

            shared_output = None
            with fork.branch():
                if self.n_shared_experts is not None and num_tokens > 0:
                    shared_output = self.shared_experts(hidden_states)

        if deferred_finalize:
            gemm2_out, expert_weights, expanded_idx = routed_expert_output
            final_hidden_states = moe_finalize_fuse_shared(
                gemm2_out,
                expanded_idx,
                expert_weights,
                shared_output,
                top_k=self.topk.topk_config.top_k,
                enable_pdl=pdl_enabled(),
            )
        else:
            final_hidden_states = (
                routed_expert_output + shared_output
                if shared_output is not None
                else routed_expert_output
            )
        return final_hidden_states


def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float:
    import math

    if scale <= 1:
        return 1.0
    return 0.1 * mscale * math.log(scale) + 1.0


class DeepseekV3FusedQkvAProjWithMqa(ReplicatedLinear):
    def __init__(
        self,
        input_size: int,
        output_size: int,
        bias: bool = True,
        skip_bias_add: bool = False,
        params_dtype: torch.dtype | None = None,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ):
        # ModelOpt NVFP4 checkpoints (e.g. DeepSeek-R1-0528-NVFP4-v2) keep the
        # q_a_proj / kv_a_proj_with_mqa weights as bf16 via exclude_modules.
        # exclude_modules matches by component name, not by the fused parent
        # prefix, so the fused layer would otherwise allocate an NVFP4-packed
        # buffer and crash when bf16 weights are copied in.
        if isinstance(quant_config, Nvfp4Config) and prefix:
            q_a_prefix = prefix.replace("fused_qkv_a_proj_with_mqa", "q_a_proj")
            kv_a_prefix = prefix.replace(
                "fused_qkv_a_proj_with_mqa", "kv_a_proj_with_mqa"
            )
            if should_exclude_quant_module(
                q_a_prefix, quant_config.exclude_modules
            ) or should_exclude_quant_module(kv_a_prefix, quant_config.exclude_modules):
                quant_config = None
        super().__init__(
            input_size,
            output_size,
            bias=bias,
            skip_bias_add=skip_bias_add,
            params_dtype=params_dtype,
            quant_config=quant_config,
            prefix=prefix,
        )
        self.use_min_latency = (
            self.bias is None
            and self.weight.dtype == torch.bfloat16
            and self.weight.size() == (2112, 7168)
            and current_platform().is_nvidia
            and _device_sm >= 90
            and _device_sm not in (120, 121)
        )

    def forward(
        self, x: torch.Tensor, block_scale=None, output_dtype=None
    ) -> torch.Tensor:
        if (
            self.use_min_latency
            and x.size(0) > 0
            and block_scale is None
            and (output_dtype is None or output_dtype == torch.bfloat16)
        ):
            return dsv3_fused_a_gemm(x, self.weight.T)

        return super().forward(x, block_scale=block_scale, output_dtype=output_dtype)[0]


class DeepseekV3AttentionMLA(nn.Module):
    # Backends that use non-absorbed MLA kernels (ragged prefill, paged KV decode).
    _MLA_KERNEL_BACKENDS = ("mla", "trtllm_mla", "tokenspeed_mla")
    # Backends that support chunked ragged prefill with prefix replay.
    _RAGGED_PREFILL_BACKENDS = ("mla", "trtllm_mla", "tokenspeed_mla")

    def __init__(
        self,
        config: PretrainedConfig,
        mapping: Mapping,
        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: dict[str, Any] | None = None,
        max_position_embeddings: int = 8192,
        quant_config: QuantizationConfig | None = None,
        layer_id=None,
        prefix: str = "",
        reduce_attn_results=True,
        alt_stream: torch.cuda.Stream | None = None,
        skip_rope: bool = False,
    ) -> None:
        super().__init__()
        self.mapping = mapping
        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.num_heads = num_heads
        if num_heads % self.mapping.attn.tp_size != 0:
            raise ValueError(
                f"num_heads={num_heads} must be divisible by attn_tp_size={self.mapping.attn.tp_size}."
            )
        self.num_local_heads = num_heads // self.mapping.attn.tp_size
        self.scaling = self.qk_head_dim**-0.5
        self.rope_theta = rope_theta
        self.max_position_embeddings = max_position_embeddings
        self.config = config
        self.alt_stream = alt_stream
        self.attention_backend = global_server_args_dict["attention_backend"]
        self.cli_factor = getattr(config, "cli_factor", 1)
        self.prefix = prefix

        # 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"

        if self.q_lora_rank is not None:
            self.fused_qkv_a_proj_with_mqa = DeepseekV3FusedQkvAProjWithMqa(
                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=quant_config,
                prefix=add_prefix("q_b_proj", prefix),
                tp_rank=self.mapping.attn.tp_rank,
                tp_size=self.mapping.attn.tp_size,
                tp_group=self.mapping.attn.tp_group,
            )
        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=self.mapping.attn.tp_rank,
                tp_size=self.mapping.attn.tp_size,
                tp_group=self.mapping.attn.tp_group,
            )

            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.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=self.mapping.attn.tp_rank,
            tp_size=self.mapping.attn.tp_size,
            tp_group=self.mapping.attn.tp_group,
        )
        # O projection.
        self.o_proj = RowParallelLinear(
            self.num_heads * self.v_head_dim,
            self.hidden_size,
            bias=False,
            reduce_results=reduce_attn_results,
            quant_config=quant_config,
            prefix=add_prefix("o_proj", prefix),
            tp_rank=self.mapping.attn.tp_rank,
            tp_size=self.mapping.attn.tp_size,
            tp_group=self.mapping.attn.tp_group,
        )
        self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)

        # Fusion layer
        if self.q_lora_rank is not None:
            self.fused_qk_layernorm = FusedRMSNorm(
                self.q_a_layernorm,
                self.kv_a_layernorm,
            )

        if not skip_rope:
            self.rotary_emb = get_rope(
                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=False,
            )

            if rope_scaling:
                mscale_all_dim = rope_scaling.get("mscale_all_dim", False)
                scaling_factor = rope_scaling["factor"]
                mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
                self.scaling = self.scaling * mscale * mscale
        else:
            self.rotary_emb = None

        # Fused RoPE+KV write kernel is incompatible with MLA: it assumes
        # K and V have the same head_dim, but MLA's KV cache is a single
        # [latent(512)|rope(64)] buffer where the two dimensions differ.
        # Passing this to the kernel causes thread overflow and silent
        # corruption of the latent cache.  All DeepSeek V2/V3 models use
        # MLA (kv_lora_rank > 0), so we unconditionally disable it here.
        self.use_fused_set_kv_buffer = False

        self.attn_mqa = PagedAttention(
            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,
        )

        self.attn_mha = PagedAttention(
            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,
        )

        self.w_kc = None
        self.w_vc = None

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        ctx: ForwardContext,
        out_cache_loc: torch.Tensor,
        comm_manager: CommManager,
        block_scale: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """MLA attention with a NARROW prefill-graph break.

        The token-shaped input/output projections (q/kv-down, layernorm,
        q_b_proj, o_proj) stay in the captured prefill graph; only the
        data-dependent attention -- KV write + varlen prefill / absorb decode
        kernels + the live prefill/decode split -- runs as the eager break
        (``_attn``). This keeps the big projection GEMMs graphed instead of
        dispatch-bound eager, collapsing the inter-segment bubbles a coarse
        whole-attention break leaves. Outside capture the ``@break_point`` is
        a direct call, so the eager path is unchanged.
        """
        if hidden_states.shape[0] == 0:
            return hidden_states
        q, latent_cache = self._project_q_latent(
            hidden_states, ctx, comm_manager, block_scale
        )
        attn_output = self._attn(positions, q, latent_cache, ctx, out_cache_loc)
        output, _ = self.o_proj(attn_output)
        return output

    def _project_q_latent(
        self,
        hidden_states: torch.Tensor,
        ctx: ForwardContext,
        comm_manager: CommManager,
        block_scale: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """QKV projection producing absorbed ``q`` and raw ``latent_cache``."""
        if self.q_lora_rank is not None:
            qkv = self.fused_qkv_a_proj_with_mqa(
                hidden_states, block_scale, torch.bfloat16
            )
            qkv = comm_manager.pre_attn_comm(qkv, ctx)
            q_a, latent_cache = qkv.split(
                [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
                dim=-1,
            )
            kv_a = latent_cache[..., : self.kv_lora_rank]
            q_norm = torch.empty_like(q_a)
            if q_a.size(0) > 0:
                self.fused_qk_layernorm(
                    input_q_a=q_a, input_kv_a=kv_a, output_q_a=q_norm
                )
            q = self.q_b_proj(q_norm)[0]
        else:
            hidden_states = comm_manager.pre_attn_comm(hidden_states, ctx)
            q = self.q_proj(hidden_states)[0]
            latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0]
            kv_a = latent_cache[..., : self.kv_lora_rank]
            self.kv_a_layernorm(kv_a, inplace=True)
        return q, latent_cache

    @break_point
    def _attn(
        self,
        positions: torch.Tensor,
        q: torch.Tensor,
        latent_cache: torch.Tensor,
        ctx: ForwardContext,
        out_cache_loc: torch.Tensor,
    ) -> torch.Tensor:
        """The eager break: KV write + varlen prefill / absorb decode attention.

        Prefill/decode dispatch over the full rows; subclasses override (the
        draft variant narrows to live rows, see ``DeepseekV3DraftAttentionMLA``).
        The split is recovered from LIVE state -- correct both in eager and
        under a prefill-graph replay, where ``ctx`` is the live ambient context
        but ``q`` is padded to the graph bucket (``q.size(0)`` is NOT the real
        token count). The decode token count comes from the live ctx; the real
        prefill token count from the live attention metadata (the same source
        the padding scrub uses). Padded tail rows produce discarded garbage.
        """
        spec = ctx.attn_backend.spec_num_tokens or 1
        num_decodes = max(ctx.bs - ctx.num_extends, 0)
        num_decode_tokens = num_decodes * spec
        if ctx.num_extends > 0:
            cmeta = ctx.attn_backend.chunked_prefill_metadata
            num_prefill_tokens = int(sum(cmeta.extend_seq_lens_cpu))
        else:
            num_prefill_tokens = 0
        real_total = num_prefill_tokens + num_decode_tokens
        attn_output = torch.empty(
            q.size(0),
            self.num_local_heads * self.v_head_dim,
            dtype=q.dtype,
            device=q.device,
        )

        if num_prefill_tokens > 0:
            prefill_ctx = replace(
                ctx,
                bs=max(ctx.bs - num_decodes, 1),
                num_extends=max(ctx.bs - num_decodes, 1),
                input_num_tokens=num_prefill_tokens,
                forward_mode=ForwardMode.EXTEND,
            )
            self.forward_normal_chunked(
                positions[:num_prefill_tokens],
                q[:num_prefill_tokens],
                latent_cache[:num_prefill_tokens],
                prefill_ctx,
                out_cache_loc[:num_prefill_tokens],
                attn_output[:num_prefill_tokens],
            )

        if num_decode_tokens > 0:
            decode_ctx = replace(
                ctx,
                bs=num_decodes,
                num_extends=0,
                input_num_tokens=num_decode_tokens,
                forward_mode=ForwardMode.DECODE,
            )
            self.forward_absorb(
                positions[num_prefill_tokens:real_total],
                q[num_prefill_tokens:real_total],
                latent_cache[num_prefill_tokens:real_total],
                decode_ctx,
                out_cache_loc[num_prefill_tokens:real_total],
                attn_output[num_prefill_tokens:real_total],
            )

        return attn_output

    def forward_absorb(
        self,
        positions: torch.Tensor,
        q: torch.Tensor,
        latent_cache: torch.Tensor,
        ctx: ForwardContext,
        out_cache_loc: torch.Tensor,
        output: torch.Tensor,
    ) -> torch.Tensor:
        Q, K = self.forward_absorb_qkv_proj(
            q,
            latent_cache,
            positions,
            ctx,
            out_cache_loc,
        )
        return self.forward_absorb_attn_v_proj(Q, K, ctx, out_cache_loc, output)

    def forward_absorb_qkv_proj(
        self,
        q: torch.Tensor,
        latent_cache: torch.Tensor,
        positions,
        ctx: ForwardContext,
        out_cache_loc: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        q = q.view(-1, self.num_local_heads, self.qk_head_dim)
        q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)

        Q = torch.empty(
            q_nope.size(0),
            self.num_local_heads,
            self.kv_lora_rank + self.qk_rope_head_dim,
            dtype=q_nope.dtype,
            device=q_nope.device,
        )
        # latent_cache contains normalized kv_a and k_pe before rotate.
        K = latent_cache.unsqueeze(1)
        q_nope_out_view = Q[..., : self.kv_lora_rank]
        if _is_amd:
            q_nope_projected = torch.bmm(
                q_nope.transpose(0, 1).contiguous(),
                self.w_kc.contiguous(),
            )
            q_nope_out_view.copy_(q_nope_projected.transpose(0, 1))
        else:
            torch.bmm(
                q_nope.transpose(0, 1),
                self.w_kc,
                out=q_nope_out_view.transpose(0, 1),
            )
        # Model-owned fused FP8 decode: RoPE + quantize + KV cache write
        # all done here, so backend only needs to do attention.
        k_scale = getattr(self.attn_mqa, "k_scale_float", 1.0)
        use_fused_fp8_decode = (
            self.attention_backend in self._MLA_KERNEL_BACKENDS
            and getattr(ctx.attn_backend, "data_type", None) == torch.float8_e4m3fn
            and self.rotary_emb is not None
            and k_scale == 1.0
        )

        if use_fused_fp8_decode:
            q_nope_absorbed = Q[..., : self.kv_lora_rank]
            k_nope_raw = K[..., : self.kv_lora_rank]
            k_pe_raw = K[..., self.kv_lora_rank :]

            query_fp8, key_fp8 = apply_rope_mla(
                positions=positions,
                q_rope=q_pe,
                k_rope=k_pe_raw,
                q_nope=q_nope_absorbed,
                k_nope=k_nope_raw,
                cos_sin_cache=self.rotary_emb.cos_sin_cache,
                is_neox=getattr(self.rotary_emb, "is_neox_style", True),
                quant_scale_q=1.0,
                quant_scale_kv=k_scale,
                enable_pdl=pdl_enabled(),
            )

            # Write FP8 KV cache (single write, no double-write)
            ctx.token_to_kv_pool.set_mla_kv_buffer(
                self.attn_mqa,
                out_cache_loc,
                cache_k_nope=key_fp8[..., : self.kv_lora_rank],
                cache_k_rope=key_fp8[..., self.kv_lora_rank :],
            )
            return query_fp8, key_fp8

        elif self.rotary_emb is not None and q_nope.size(0) > 0:
            # Apply RoPE directly on Q and K slices
            q_pe, k_pe = self.rotary_emb(
                positions,
                q_pe,
                K[..., self.kv_lora_rank :],
                fused_set_kv_buffer_arg=(
                    create_fused_set_kv_buffer_arg(
                        value=K[..., : self.kv_lora_rank],
                        layer=self.attn_mqa,
                        out_cache_loc=out_cache_loc,
                        token_to_kv_pool=ctx.token_to_kv_pool,
                    )
                    if self.use_fused_set_kv_buffer
                    else None
                ),
            )
            Q[..., self.kv_lora_rank :].copy_(q_pe)
            K[..., self.kv_lora_rank :].copy_(k_pe)
        else:
            Q[..., self.kv_lora_rank :] = q_pe

        # For MLA kernel backends, write KV cache here (model-owned) so the
        # backend never has to. This unifies the FP8 fused path (written above)
        # and the BF16 path into a single ownership model.
        if (
            self.attention_backend in self._MLA_KERNEL_BACKENDS
            and not self.use_fused_set_kv_buffer
        ):
            ctx.token_to_kv_pool.set_mla_kv_buffer(
                self.attn_mqa,
                out_cache_loc,
                cache_k_nope=K[..., : self.kv_lora_rank],
                cache_k_rope=K[..., self.kv_lora_rank :],
            )

        return Q, K

    def forward_absorb_attn_v_proj(
        self,
        Q,
        K,
        ctx: ForwardContext,
        out_cache_loc: torch.Tensor,
        output: torch.Tensor,
        record_kv_cache: bool | None = None,
    ) -> torch.Tensor:
        # MLA kernel backends: KV cache already written in forward_absorb_qkv_proj.
        # Other backends: write via fused_set_kv_buffer or let backend handle it.
        if self.attention_backend in self._MLA_KERNEL_BACKENDS:
            need_save_kv = False
        else:
            need_save_kv = not self.use_fused_set_kv_buffer

        attn_output = self.attn_mqa(
            Q,
            K,
            K[..., : self.kv_lora_rank],
            ctx,
            out_cache_loc,
            save_kv_cache=need_save_kv,
            record_kv_cache=record_kv_cache,
        )
        attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank)
        if _is_amd:
            projected = torch.bmm(
                attn_output.transpose(0, 1).contiguous(),
                self.w_vc.contiguous(),
            )
            output.copy_(projected.transpose(0, 1).reshape_as(output))
        else:
            output_view = output.view(-1, self.num_local_heads, self.v_head_dim)
            torch.bmm(
                attn_output.transpose(0, 1),
                self.w_vc,
                out=output_view.transpose(0, 1),
            )
        return output

    def forward_normal_chunked(
        self,
        positions: torch.Tensor,
        q: torch.Tensor,
        latent_cache: torch.Tensor,
        ctx: ForwardContext,
        out_cache_loc: torch.Tensor,
        output: torch.Tensor,
    ) -> torch.Tensor:
        # Prefill-graph padding contract: zero garbage rows the per-row projections
        # + FP8 quantize would otherwise touch (see scrub_padding_tail).
        ntok = sum(ctx.attn_backend.chunked_prefill_metadata.extend_seq_lens_cpu)
        scrub_padding_tail(ntok, q, latent_cache)
        q, k, v = self.forward_normal_chunked_kv_prepare(
            positions, q, latent_cache, ctx, out_cache_loc
        )
        return self.forward_normal_chunked_kv_core(q, k, v, ctx, out_cache_loc, output)

    def forward_normal_chunked_kv_prepare(
        self,
        positions: torch.Tensor,
        q: torch.Tensor,
        latent_cache: torch.Tensor,
        ctx: ForwardContext,
        out_cache_loc: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        kv_a, k_pe = latent_cache.split(
            [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
        )
        k_pe = k_pe.unsqueeze(1)

        q = q.view(-1, self.num_local_heads, self.qk_head_dim)
        q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
        kv = self.kv_b_proj(kv_a)[0]
        kv = kv.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim)
        k_nope = kv[..., : self.qk_nope_head_dim]
        v = kv[..., self.qk_nope_head_dim :]

        # FP8 prefill: fused RoPE + FP8 quantize, direct FP8 KV cache write.
        # Disabled when k_scale != 1.0; mla_fp8_utils.py documents the current limitation.
        k_scale = getattr(self.attn_mha, "k_scale_float", 1.0)
        use_fp8_prefill = (
            self.attention_backend in self._MLA_KERNEL_BACKENDS
            and getattr(ctx.attn_backend, "data_type", None) == torch.float8_e4m3fn
            and self.rotary_emb is not None
            and k_scale == 1.0
        )

        if use_fp8_prefill:
            # Expand k_pe from [tokens,1,rope] to [tokens,heads,rope] for GQA
            k_pe_expanded = k_pe.expand(-1, self.num_local_heads, -1)

            q_fp8, k_fp8 = apply_rope_mla(
                positions=positions,
                q_rope=q_pe,
                k_rope=k_pe_expanded,
                q_nope=q_nope,
                k_nope=k_nope,
                cos_sin_cache=self.rotary_emb.cos_sin_cache,
                is_neox=getattr(self.rotary_emb, "is_neox_style", True),
                quant_scale_q=1.0,
                quant_scale_kv=k_scale,
                enable_pdl=pdl_enabled(),
            )

            v_fp8 = fp8_quantize(v, enable_pdl=pdl_enabled())

            # Write FP8 KV cache directly (skip BF16→FP8 conversion in pool)
            k_pe_for_cache = k_fp8[:, 0:1, self.qk_nope_head_dim :]
            kv_a_fp8 = fp8_quantize(kv_a, enable_pdl=pdl_enabled())
            ctx.token_to_kv_pool.set_mla_kv_buffer(
                self.attn_mha,
                out_cache_loc,
                cache_k_nope=kv_a_fp8.unsqueeze(1),
                cache_k_rope=k_pe_for_cache,
            )

            return q_fp8, k_fp8, v_fp8

        # BF16 path: apply RoPE, assemble Q/K, write cache
        if self.rotary_emb is not None:
            q_pe, k_pe = self.rotary_emb(
                positions,
                q_pe,
                k_pe,
                fused_set_kv_buffer_arg=(
                    create_fused_set_kv_buffer_arg(
                        value=kv_a.unsqueeze(1),
                        layer=self.attn_mha,
                        out_cache_loc=out_cache_loc,
                        token_to_kv_pool=ctx.token_to_kv_pool,
                    )
                    if self.use_fused_set_kv_buffer
                    else None
                ),
            )

        q[..., self.qk_nope_head_dim :] = q_pe
        k = torch.empty_like(q)
        k[..., : self.qk_nope_head_dim] = k_nope
        k[..., self.qk_nope_head_dim :] = k_pe

        if not self.use_fused_set_kv_buffer:
            ctx.token_to_kv_pool.set_mla_kv_buffer(
                self.attn_mha,
                out_cache_loc,
                cache_k_nope=kv_a.unsqueeze(1),
                cache_k_rope=k_pe,
            )

        return q, k, v

    def forward_normal_chunked_kv_core(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        ctx: ForwardContext,
        out_cache_loc: torch.Tensor,
        output: torch.Tensor,
    ) -> torch.Tensor:
        attn_backend = ctx.attn_backend
        chunk_meta = attn_backend.chunked_prefill_metadata
        token_to_kv_pool = ctx.token_to_kv_pool

        # Scale compensation for FP8 prefill: bmm1_scale = k_scale * softmax_scale
        scaling = self.attn_mha.scaling
        k_scale = getattr(self.attn_mha, "k_scale_float", 1.0)
        if q.dtype == torch.float8_e4m3fn:
            scaling = k_scale * scaling

        # Causal self-attention over the new chunk tokens. q_lens == kv_lens ==
        # extend_seq_lens, so cum_seq_lens_q and cum_seq_lens_kv alias the same
        # cum_extend_seq_lens. Causal pass writes directly into output; each
        # chunk's merge accumulates in place via attn_merge_state(inplace=True).
        num_extends = chunk_meta.extend_seq_lens.size(0)
        output_view = output.view(-1, self.num_local_heads, self.v_head_dim)
        _, accum_lse = attn_backend.forward_extend_chunked(
            q,
            k,
            v,
            scaling,
            self.attn_mha.logit_cap,
            cum_seq_lens_q=chunk_meta.cum_extend_seq_lens,
            cum_seq_lens_kv=chunk_meta.cum_extend_seq_lens,
            max_q_len=chunk_meta.max_extend_seq_len,
            max_kv_len=chunk_meta.max_extend_seq_len,
            seq_lens=chunk_meta.extend_seq_lens,
            batch_size=num_extends,
            causal=True,
            out=output_view,
        )

        # Always read KV cache as BF16 for kv_b_proj (weight is BF16), even if Q is FP8.
        read_dtype = (
            q.dtype
            if q.dtype not in (torch.float8_e4m3fn, torch.float8_e5m2)
            else torch.bfloat16
        )

        for loop_idx in range(chunk_meta.chunked_loop_num):
            chunk_kv_indices = chunk_meta.chunk_kv_indices_list[loop_idx]

            kv_a_normed, k_pe = token_to_kv_pool.get_mla_kv_buffer(
                self.attn_mha, chunk_kv_indices, read_dtype
            )

            kv_a_normed = kv_a_normed.squeeze(1)
            kv = self.kv_b_proj(kv_a_normed)[0]
            kv = kv.view(
                -1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim
            )
            v = kv[..., self.qk_nope_head_dim :]
            k_nope = kv[..., : self.qk_nope_head_dim]

            if q.dtype == torch.float8_e4m3fn:
                # FP8 Attention
                k, v = mla_kv_pack_quantize_fp8(
                    k_nope, k_pe, v, k_scale_inv=1.0 / k_scale, enable_pdl=pdl_enabled()
                )
            else:
                # BF16 Attention
                k = torch.cat(
                    [k_nope, k_pe.expand(-1, self.num_local_heads, -1)], dim=-1
                )

            chunk_output, lse = attn_backend.forward_extend_chunked(
                q,
                k,
                v,
                scaling,
                self.attn_mha.logit_cap,
                cum_seq_lens_q=chunk_meta.cum_extend_seq_lens,
                cum_seq_lens_kv=chunk_meta.cu_chunked_seq_len[loop_idx],
                max_q_len=chunk_meta.max_extend_seq_len,
                max_kv_len=chunk_meta.max_chunk_len_per_loop[loop_idx],
                seq_lens=chunk_meta.chunked_seq_len[loop_idx],
                batch_size=num_extends,
                causal=False,
            )

            attn_merge_state(
                output_view,
                accum_lse,
                chunk_output,
                lse,
                inplace=True,
                enable_pdl=pdl_enabled(),
            )

        return output


class DeepseekV3DraftAttentionMLA(DeepseekV3AttentionMLA):
    """Draft variant of MLA shared by the NextN and Eagle3 MLA drafters.

    On the active first draft step the full ``latent_cache`` (N rows) is
    projected so every KV cache entry is written, but only the live query rows
    (``ctx.gather_ids``) run the absorbed decode attention, narrowing the output
    to ``[bs, H]``.  Multi-step decode and target paths delegate to the base.
    Single-layer only, so dropping the dead rows has no downstream consumer.
    """

    def _attn(
        self,
        positions: torch.Tensor,
        q: torch.Tensor,
        latent_cache: torch.Tensor,
        ctx: ForwardContext,
        out_cache_loc: torch.Tensor,
    ) -> torch.Tensor:
        if ctx.accept_lengths is None:
            return super()._attn(positions, q, latent_cache, ctx, out_cache_loc)

        self._apply_correction(ctx)

        # Full q/latent_cache write all KV cache rows; only the live rows
        # (ctx.gather_ids) run the absorbed decode attention, so the output is
        # narrowed to [bs, H] for o_proj / MLP / post-norms.
        decode_ctx = replace(ctx, forward_mode=ForwardMode.DECODE)
        Q, K = self.forward_absorb_qkv_proj(
            q,
            latent_cache,
            positions,
            decode_ctx,
            out_cache_loc,
        )
        Q = Q.index_select(0, ctx.gather_ids)
        attn_output = q.new_empty(ctx.bs, self.num_local_heads * self.v_head_dim)
        # gather_ids keeps one live row per request, so the decode runs on the
        # full bs -- the page table and seq lens must span the same rows. Drop
        # the [num_extends:] slice a MIXED target's first-step metadata sets up
        # (mirrors the multi-step drafter loop's override_num_extends(0)).
        with ctx.attn_backend.override_num_extends(0):
            self.forward_absorb_attn_v_proj(
                Q,
                K,
                decode_ctx,
                out_cache_loc,
                attn_output,
                # Real-mode record: decode_ctx would skip the PD cache-step here.
                record_kv_cache=not ctx.forward_mode.is_decode_or_idle(),
            )
        return attn_output

    def _apply_correction(self, ctx: ForwardContext) -> None:
        """Trim decode rows' cache_seqlens by ``spec_num_tokens - accept_lengths``."""
        seq_lens_buf = ctx.draft_seq_lens_buf
        if seq_lens_buf is None or ctx.accept_lengths is None:
            return
        num_extends = ctx.num_extends
        if num_extends >= ctx.bs:
            return
        correction = (
            ctx.attn_backend.spec_num_tokens - ctx.accept_lengths[num_extends:]
        ).to(seq_lens_buf.dtype)
        seq_lens_buf[num_extends : ctx.bs].sub_(correction).clamp_(min=1)


class DeepseekV3DecoderLayer(nn.Module):
    @property
    def attention_cls(self) -> type[nn.Module]:
        return DeepseekV3AttentionMLA

    def __init__(
        self,
        config: PretrainedConfig,
        layer_id: int,
        mapping: Mapping,
        quant_config: QuantizationConfig | None = None,
        is_nextn: bool = False,
        prefix: str = "",
        alt_stream: torch.cuda.Stream | None = None,
    ) -> None:
        super().__init__()
        self.mapping = mapping
        self.hidden_size = config.hidden_size
        rope_theta = get_rope_theta(config)
        rope_scaling = getattr(config, "rope_scaling", None)
        max_position_embeddings = getattr(config, "max_position_embeddings", 8192)

        self.self_attn = self.attention_cls(
            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=(
                None
                if "self_attn" in getattr(config, "disable_quant_module", [])
                else quant_config
            ),
            layer_id=layer_id,
            prefix=add_prefix("self_attn", prefix),
            reduce_attn_results=False,
            alt_stream=alt_stream,
            mapping=self.mapping,
        )

        self.layer_id = layer_id
        self.is_moe_layer = self._is_moe_layer(layer_id, is_nextn, config)
        if self.is_moe_layer:
            self.mlp = DeepseekV3MoE(
                config=config,
                mapping=self.mapping,
                quant_config=quant_config,
                layer_index=layer_id,
                prefix=add_prefix("mlp", prefix),
                alt_stream=alt_stream,
            )
        else:
            self.mlp = DeepseekV3MLP(
                hidden_size=config.hidden_size,
                intermediate_size=(
                    config.ffn_hidden_size
                    if hasattr(config, "ffn_hidden_size")
                    else config.intermediate_size
                ),
                hidden_act=config.hidden_act,
                mapping=self.mapping,
                quant_config=(
                    None
                    if "dense_mlp" in getattr(config, "disable_quant_module", [])
                    else quant_config
                ),
                prefix=add_prefix("mlp", prefix),
                is_shared_expert=False,
            )
        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.comm_manager = CommManager(
            mapping=self.mapping,
            layer_id=self.layer_id,
            is_moe=self.is_moe_layer,
            prev_is_moe=self._is_moe_layer(layer_id - 1, is_nextn, config),
            input_layernorm=self.input_layernorm,
            post_attn_layernorm=self.post_attention_layernorm,
        )

    @staticmethod
    def _is_moe_layer(layer_id: int, is_nextn: bool, config):
        if is_nextn:
            return True
        if (
            config.n_routed_experts is not None
            and layer_id >= config.first_k_dense_replace
            and layer_id % config.moe_layer_freq == 0
        ):
            return True
        return False

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        ctx: ForwardContext,
        out_cache_loc: torch.Tensor,
        residual: torch.Tensor | None,
    ) -> torch.Tensor:

        num_global_tokens, max_num_tokens_per_gpu = self.comm_manager.get_num_tokens(
            ctx
        )

        if not ctx.forward_mode.is_idle():
            hidden_states, residual = self.comm_manager.input_reduce_norm(
                hidden_states, residual
            )
            hidden_states = self.self_attn(
                positions=positions,
                hidden_states=hidden_states,
                ctx=ctx,
                out_cache_loc=out_cache_loc,
                comm_manager=self.comm_manager,
            )
            hidden_states, residual = self.comm_manager.post_attn_reduce_norm(
                hidden_states, residual, ctx
            )
            hidden_states = self.forward_mlp(
                hidden_states,
                residual,
                ctx,
                num_global_tokens,
                max_num_tokens_per_gpu,
            )
        else:
            hidden_states = self.forward_mlp(
                hidden_states,
                residual,
                ctx,
                num_global_tokens,
                max_num_tokens_per_gpu,
            )
        return hidden_states, residual

    def input_layer_norm_fn(self, hidden_states, residual):
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
            hidden_states, residual = self.input_layernorm(hidden_states, residual)
        return hidden_states, residual

    def forward_mlp(
        self,
        hidden_states,
        residual,
        ctx: ForwardContext,
        num_global_tokens,
        max_num_tokens_per_gpu,
    ):
        hidden_states = self.comm_manager.pre_mlp_comm(hidden_states, ctx)
        if self.is_moe_layer:
            hidden_states = self.mlp(
                hidden_states, num_global_tokens, max_num_tokens_per_gpu
            )
        else:
            hidden_states = self.mlp(hidden_states)
        hidden_states, residual = self.comm_manager.post_mlp_fused(
            hidden_states, residual, ctx
        )
        return hidden_states


class DeepseekV3Model(nn.Module):
    fall_back_to_pt_during_load = False

    def __init__(
        self,
        config: PretrainedConfig,
        mapping: Mapping,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.mapping = mapping
        self.padding_id = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.embed_tokens = VocabParallelEmbedding(
            config.vocab_size,
            config.hidden_size,
        )
        self.alt_stream = torch.cuda.Stream()
        # config.num_hidden_layers = 5; self.start_layer,self.end_layer = 0, 5
        self.layers = nn.ModuleList(
            [
                DeepseekV3DecoderLayer(
                    config,
                    layer_id,
                    mapping=self.mapping,
                    quant_config=quant_config,
                    prefix=add_prefix(f"layers.{layer_id}", prefix),
                    alt_stream=self.alt_stream,
                )
                for layer_id in range(config.num_hidden_layers)
            ]
        )
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        # For EAGLE3 support: set of layer indices whose *input* hidden states
        # are captured. Populated by set_eagle3_layers_to_capture().
        self.layers_to_capture: set = set()

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        ctx: ForwardContext,
        out_cache_loc: torch.Tensor,
        input_embeds: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, list[torch.Tensor] | None]:
        if input_embeds is not None:
            hidden_states = input_embeds
        else:
            hidden_states = self.embed_tokens(input_ids)
        if CP_METADATA:
            hidden_states = cp_split_and_rebuild_data(
                hidden_states,
                CP_METADATA.value.split_list,
                CP_METADATA.value.zigzag_index,
            )
            positions = cp_split_and_rebuild_data(
                positions, CP_METADATA.value.split_list, CP_METADATA.value.zigzag_index
            )
        residual = None
        aux_hidden_states = [] if self.layers_to_capture else None
        for i in range(len(self.layers)):
            if aux_hidden_states is not None and i in self.layers_to_capture:
                # Under RSAG the inter-layer hidden/residual are reduce-
                # scattered across the attn TP group; aux consumers (e.g. the
                # EAGLE3 drafter) expect full rows, so gather before capturing.
                aux = (
                    hidden_states + residual if residual is not None else hidden_states
                )
                gathered = self.layers[i].comm_manager.gather_residual(aux, ctx)
                aux_hidden_states.append(
                    gathered if gathered is aux else gathered.clone()
                )
            with get_global_expert_distribution_recorder().with_current_layer(i):
                layer = self.layers[i]
                hidden_states, residual = layer(
                    positions,
                    hidden_states,
                    ctx,
                    out_cache_loc,
                    residual,
                )
        if not ctx.forward_mode.is_idle():
            if not ENABLE_CP:
                hidden_states, _ = layer.comm_manager.final_norm(
                    hidden_states, residual, ctx, self.norm
                )
            else:
                hidden_states, _ = self.norm(hidden_states, residual)
        if CP_METADATA:
            hidden_states = cp_all_gather_rerange_output(
                hidden_states,
                CP_METADATA.value,
                self.mapping.attn.tp_rank,
                self.mapping.attn.tp_group,
            )
        return hidden_states, aux_hidden_states


class DeepseekV3ForCausalLM(BaseCausalLM):
    model_cls = DeepseekV3Model

    def __init__(
        self,
        config: PretrainedConfig,
        mapping: Mapping,
        model: DeepseekV3Model | None = None,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        self._model_override = model
        super().__init__(
            config=config,
            mapping=mapping,
            quant_config=quant_config,
            prefix=prefix,
        )

    def resolve_model(
        self,
        config: PretrainedConfig,
        mapping: Mapping,
        quant_config: QuantizationConfig | None,
        prefix: str,
    ) -> DeepseekV3Model:
        if self._model_override is not None:
            return self._model_override
        return self.model_cls(
            config,
            mapping=mapping,
            quant_config=quant_config,
            prefix=add_prefix("model", prefix),
        )

    def post_init(self) -> None:
        self._routed_experts_weights_of_layer = LazyValue(
            lambda: {
                layer_id: layer.mlp.get_moe_routed_weights()
                for layer_id, layer in enumerate(self.model.layers)
                if isinstance(layer.mlp, DeepseekV3MoE)
            }
        )

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

    def set_eagle3_layers_to_capture(self, layer_ids: list[int] | None = None):
        # layer_ids are 0-indexed from the external API; +1 because the capture
        # check runs *before* the layer forward, so index i captures layer i-1's output.
        if layer_ids is None:
            num_layers = self.config.num_hidden_layers
            self.model.layers_to_capture = {2, num_layers // 2, num_layers - 3}
        else:
            self.model.layers_to_capture = {val + 1 for val in layer_ids}

    def set_dflash_layers_to_capture(self, layer_ids: list[int]):
        # DFlash checkpoints name 0-indexed target layer outputs. The capture
        # check runs before layer i, so capture at i + 1 for layer i's output.
        num_layers = len(self.model.layers)
        if len(set(layer_ids)) != len(layer_ids):
            raise ValueError("DFLASH target_layer_ids must be unique.")

        invalid = [val for val in layer_ids if val < 0 or val + 1 >= num_layers]
        if invalid:
            raise ValueError(
                "DFLASH target_layer_ids must map to capturable target layer "
                f"outputs. Got invalid ids {invalid}; valid range is "
                f"[0, {num_layers - 2}] for {num_layers} target layers."
            )
        self.model.layers_to_capture = {val + 1 for val in layer_ids}

    def get_param(self, params_dict, name):
        if name in params_dict:
            return params_dict[name]

        if "language_model." in name:
            name = name.replace("language_model.", "")
            if name in params_dict:
                return params_dict[name]

        if name.endswith(_OPTIONAL_MISSING_WEIGHT_SUFFIXES):
            return None

        logger.warning("The %s is not in the model.", name)
        return None

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("gate_up_proj", "gate_proj", 0),
            ("gate_up_proj", "up_proj", 1),
        ]

        # Fuse q_a_proj and kv_a_proj_with_mqa along output dimension when q_lora_rank is not None
        fuse_qkv_a_proj = getattr(self.config, "q_lora_rank", None) is not None

        params_dict = dict(self.named_parameters())
        moe_params_dict = dict(params_dict)
        for param_name, param in params_dict.items():
            if param_name.startswith("model."):
                moe_params_dict.setdefault(
                    param_name.replace("model.", "model.language_model.", 1),
                    param,
                )
                moe_params_dict.setdefault(
                    param_name.replace("model.", "language_model.model.", 1),
                    param,
                )
        # MoE expert weights, scales, and activation scales are handled
        # by the checkpoint loader.
        moe_loader = build_moe_checkpoint_loader(
            params_dict=moe_params_dict,
            expert_schema=ExpertCheckpointSchema(
                gate_proj_name="gate_proj",
                down_proj_name="down_proj",
                up_proj_name="up_proj",
            ),
            num_experts=self.config.n_routed_experts,
            ep_rank=self.mapping.moe.ep_rank,
            ep_size=self.mapping.moe.ep_size,
        )
        for name, loaded_weight in weights:
            layer_id = get_layer_id(name)
            if (
                layer_id is not None
                and hasattr(self.model, "start_layer")
                and (
                    layer_id < self.model.start_layer
                    or layer_id >= self.model.end_layer
                )
            ):
                continue
            if hasattr(self.config, "num_nextn_predict_layers"):
                num_nextn_layers = self.config.num_nextn_predict_layers
                if num_nextn_layers > 0 and name.startswith("model.layers"):
                    name_list = name.split(".")
                    if (
                        len(name_list) >= 3
                        and int(name_list[2]) >= self.config.num_hidden_layers
                    ):
                        continue
            if "rotary_emb.inv_freq" in name:
                continue
            if ".indexer." in name:
                continue
            for param_name, weight_name, shard_id in stacked_params_mapping:
                # Skip non-stacked layers and experts (experts handled below).
                if weight_name not in name:
                    continue
                # We have mlp.experts[0].gate_proj in the checkpoint.
                # Since moe_loader handles the experts below,
                # we need to skip here BEFORE we update the name, otherwise
                # name will be updated to mlp.experts[0].gate_up_proj, which
                # will then be updated below by moe_loader
                # for mlp.experts[0].gate_gate_up_proj, which breaks load.
                if ("mlp.experts." in name) and name not in params_dict:
                    continue
                name = name.replace(weight_name, param_name)
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
                param = self.get_param(params_dict, name)
                if param is None:
                    continue
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
                if moe_loader.matches(name):
                    moe_loader.load(name, loaded_weight)
                    continue

                if fuse_qkv_a_proj and (
                    "q_a_proj" in name or "kv_a_proj_with_mqa" in name
                ):
                    quant_block_size = 1
                    # ``weight_block_size`` exists only on block-FP8 configs;
                    # elsewhere (e.g. compressed-tensors INT4) q/kv_a_proj is unquantized.
                    weight_block_size = getattr(
                        self.quant_config, "weight_block_size", None
                    )
                    if weight_block_size is not None:
                        quant_block_size = weight_block_size[0]
                    begin_size_mp = {
                        "q_a_proj": 0,
                        "kv_a_proj_with_mqa": self.config.q_lora_rank,
                    }
                    if "q_a_proj" in name:
                        param = self.get_param(
                            params_dict,
                            name.replace("q_a_proj", "fused_qkv_a_proj_with_mqa"),
                        )
                        weight_loader = param.weight_loader
                        begin_size = begin_size_mp["q_a_proj"]
                    elif "kv_a_proj_with_mqa" in name:
                        param = self.get_param(
                            params_dict,
                            name.replace(
                                "kv_a_proj_with_mqa", "fused_qkv_a_proj_with_mqa"
                            ),
                        )
                        weight_loader = param.weight_loader
                        begin_size = begin_size_mp["kv_a_proj_with_mqa"]
                    if "scale_inv" in name:
                        begin_size //= quant_block_size
                    weight_loader(param, loaded_weight, begin_size=begin_size)
                else:
                    # Owned-expert weights were already consumed by ``moe_loader.load(...)`` above (matches() == True branch).
                    # Anything reaching here that still looks like an expert weight is for an expert this rank does ot own under ep_size > 1.
                    if ".mlp.experts." in name:
                        continue
                    if "q_a_proj" in name and name not in params_dict:
                        name = name.replace("q_a_proj", "q_proj")
                    param = self.get_param(params_dict, name)
                    if param is None:
                        continue
                    weight_loader = getattr(
                        param, "weight_loader", default_weight_loader
                    )
                    weight_loader(param, loaded_weight)

        self.post_load_weights()

    def post_load_weights(self):
        for layer_id in range(self.config.num_hidden_layers):
            self_attn = self.model.layers[layer_id].self_attn
            if hasattr(
                self.quant_config, "weight_block_size"
            ) and self_attn.kv_b_proj.weight.dtype in (
                torch.float8_e4m3fn,
                torch.float8_e4m3fnuz,
            ):
                weight_block_size = self.quant_config.weight_block_size
                if weight_block_size is not None:
                    if not hasattr(self_attn.kv_b_proj, "weight_scale_inv"):
                        raise RuntimeError(
                            "kv_b_proj.weight_scale_inv is required for block FP8 dequant."
                        )
                    dtype = torch.get_default_dtype()
                    w = block_dequant(
                        self_attn.kv_b_proj.weight,
                        self_attn.kv_b_proj.weight_scale_inv,
                        weight_block_size,
                    ).to(dtype)
            else:
                w = self_attn.kv_b_proj.weight

            self_attn.w_kc, self_attn.w_vc = _prepare_mla_kv_b_proj_weights(
                w, self_attn
            )

    def load_kv_cache_scales(self, quantization_param_path: str) -> None:
        tp_size = self.mapping.attn.tp_size
        tp_rank = self.mapping.attn.tp_rank
        for layer_idx, scaling_factor in kv_cache_scales_loader(
            quantization_param_path,
            tp_rank,
            tp_size,
            self.config.num_hidden_layers,
            self.config.__class__.model_type,
        ):
            if not isinstance(self.model.layers[layer_idx], nn.Identity):
                self_attn = self.model.layers[layer_idx].self_attn
                # Set on both attn_mha (non-absorbed prefill) and attn_mqa (absorbed decode).
                for attn in (self_attn.attn_mha, self_attn.attn_mqa):
                    if attn is not None and hasattr(attn, "k_scale"):
                        attn.k_scale = scaling_factor
                        attn.k_scale_float = scaling_factor

    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,
        )


# ---------------------------------------------------------------------------
# Eagle3 MLA draft model
# ---------------------------------------------------------------------------


class Eagle3MlaDecoderLayer(nn.Module):
    """Single decoder layer for Eagle3 MLA draft model.

    The fused_qkv_a_proj_with_mqa is overridden to accept 2x hidden_size
    input (concatenated [embeds, hidden_states]) while keeping o_proj at
    the standard hidden_size output.
    """

    def __init__(
        self,
        config: PretrainedConfig,
        mapping: Mapping,
        layer_id: int = 0,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.mapping = mapping
        self.hidden_size = config.hidden_size
        self.layer_id = layer_id
        rope_theta = get_rope_theta(config)
        rope_scaling = getattr(config, "rope_scaling", None)
        max_position_embeddings = getattr(config, "max_position_embeddings", 8192)

        self.self_attn = DeepseekV3DraftAttentionMLA(
            config=config,
            mapping=self.mapping,
            hidden_size=self.hidden_size,
            num_heads=config.num_attention_heads,
            qk_nope_head_dim=getattr(config, "qk_nope_head_dim", 128),
            qk_rope_head_dim=getattr(config, "qk_rope_head_dim", 64),
            v_head_dim=getattr(config, "v_head_dim", 128),
            q_lora_rank=getattr(config, "q_lora_rank", None),
            kv_lora_rank=getattr(config, "kv_lora_rank", 512),
            rope_theta=rope_theta,
            rope_scaling=rope_scaling,
            max_position_embeddings=max_position_embeddings,
            quant_config=quant_config,
            layer_id=layer_id,
            prefix=add_prefix("self_attn", prefix),
            reduce_attn_results=False,
        )

        if hasattr(self.self_attn, "fused_qkv_a_proj_with_mqa"):
            q_lora_rank = getattr(config, "q_lora_rank", 0) or 0
            kv_lora_rank = getattr(config, "kv_lora_rank", 512)
            qk_rope_head_dim = getattr(config, "qk_rope_head_dim", 64)
            self.self_attn.fused_qkv_a_proj_with_mqa = DeepseekV3FusedQkvAProjWithMqa(
                2 * self.hidden_size,
                q_lora_rank + kv_lora_rank + qk_rope_head_dim,
                bias=False,
                quant_config=quant_config,
                prefix=add_prefix(
                    "fused_qkv_a_proj_with_mqa",
                    add_prefix("self_attn", prefix),
                ),
            )

        self.mlp = DeepseekV3MLP(
            hidden_size=config.hidden_size,
            intermediate_size=getattr(
                config, "intermediate_size", config.hidden_size * 4
            ),
            hidden_act=getattr(config, "hidden_act", "silu"),
            mapping=self.mapping,
            quant_config=quant_config,
            prefix=add_prefix("mlp", prefix),
        )

        self.hidden_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.fused_input_hidden_norm = FusedRMSNorm(
            self.input_layernorm,
            self.hidden_norm,
        )
        self.post_attention_layernorm = RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )

        self.comm_manager = CommManager(
            mapping=self.mapping,
            layer_id=self.layer_id,
            is_moe=False,
            prev_is_moe=False,
            post_attn_layernorm=self.post_attention_layernorm,
        )

    def forward(
        self,
        positions: torch.Tensor,
        embeds: torch.Tensor,
        hidden_states: torch.Tensor,
        ctx: ForwardContext,
        out_cache_loc: torch.Tensor,
        residual: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        residual = hidden_states

        if not ctx.forward_mode.is_idle():
            fused_norm_out = torch.empty(
                embeds.size(0),
                self.hidden_size * 2,
                dtype=embeds.dtype,
                device=embeds.device,
            )
            # FusedRMSNorm's q_a/kv_a kwargs are MLA-specific names.
            # Here embeds and hidden_states corresponds to q_a and kv_a, separately.
            self.fused_input_hidden_norm(
                input_q_a=embeds,
                input_kv_a=hidden_states,
                output_q_a=fused_norm_out[..., : self.hidden_size],
                output_kv_a=fused_norm_out[..., self.hidden_size :],
            )

            hidden_states = self.self_attn(
                positions=positions,
                hidden_states=fused_norm_out,
                ctx=ctx,
                out_cache_loc=out_cache_loc,
                comm_manager=self.comm_manager,
            )

            # Active first draft step narrows attn output to [bs, H]; align the
            # residual to the same live rows before the post-attn reduce-norm.
            if ctx.accept_lengths is not None:
                residual = residual.index_select(0, ctx.gather_ids)
            hidden_states, residual = self.comm_manager.post_attn_reduce_norm(
                hidden_states, residual, ctx
            )

        hidden_states = self.comm_manager.pre_mlp_comm(hidden_states, ctx)
        hidden_states = self.mlp(hidden_states)
        hidden_states, residual = self.comm_manager.post_mlp_fused(
            hidden_states, residual, ctx
        )

        return hidden_states, residual


class Eagle3MlaModel(nn.Module):
    @staticmethod
    def _get_eagle_layer_ids(config: PretrainedConfig):
        """Extract eagle aux hidden state layer IDs from config, or None if absent."""
        eagle_config = getattr(config, "eagle_config", None)
        if eagle_config is None:
            return getattr(config, "eagle_aux_hidden_state_layer_ids", None)
        if isinstance(eagle_config, dict):
            return eagle_config.get("eagle_aux_hidden_state_layer_ids", None)
        return getattr(eagle_config, "eagle_aux_hidden_state_layer_ids", None)

    def __init__(
        self,
        config: PretrainedConfig,
        mapping: Mapping,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.mapping = mapping
        self.config = config
        self.vocab_size = config.vocab_size

        self.embed_tokens = VocabParallelEmbedding(
            config.vocab_size,
            config.hidden_size,
            prefix=add_prefix("embed_tokens", prefix),
        )

        layer_ids = self._get_eagle_layer_ids(config)
        self.num_fc_input_dim = len(layer_ids) if layer_ids is not None else 3

        target_hidden_size = getattr(config, "target_hidden_size", config.hidden_size)
        fc_input_size = target_hidden_size * self.num_fc_input_dim

        self.fc = ColumnParallelLinear(
            fc_input_size,
            config.hidden_size,
            bias=False,
            gather_output=True,
            quant_config=quant_config,
            prefix=add_prefix("fc", prefix),
            tp_rank=self.mapping.attn.tp_rank,
            tp_size=self.mapping.attn.tp_size,
            tp_group=self.mapping.attn.tp_group,
        )

        self.midlayer = Eagle3MlaDecoderLayer(
            config,
            mapping=self.mapping,
            layer_id=0,
            quant_config=quant_config,
            prefix=prefix,
        )

        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        ctx: ForwardContext,
        out_cache_loc: torch.Tensor,
        input_embeds: torch.Tensor | None = None,
        captured_hidden_states: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, list[torch.Tensor]]:
        if captured_hidden_states is None:
            raise ValueError("Eagle3 MLA forward requires captured_hidden_states.")
        if input_embeds is None:
            embeds = self.embed_tokens(input_ids)
        else:
            embeds = input_embeds

        hidden_states = captured_hidden_states
        if hidden_states.size(-1) != embeds.size(-1):
            hidden_states, _ = self.fc(hidden_states)

        residual = None
        hidden_states, residual = self.midlayer(
            positions,
            embeds,
            hidden_states,
            ctx,
            out_cache_loc,
            residual,
        )

        comm_manager = self.midlayer.comm_manager
        if comm_manager.should_fuse(hidden_states.size(0)):
            hidden_states_to_logits, hidden_states_to_aux, *_ = (
                self.norm.forward_with_allreduce_fusion(
                    self.mapping.dense.tp_rank,
                    self.mapping.dense.tp_group,
                    hidden_states,
                    residual,
                )
            )
        else:
            hidden_states_to_logits, hidden_states_to_aux = self.norm(
                hidden_states, residual
            )
            hidden_states_to_logits, _ = comm_manager.post_final_norm_comm(
                hidden_states_to_logits, None, ctx
            )
            hidden_states_to_aux, _ = comm_manager.post_final_norm_comm(
                hidden_states_to_aux, None, ctx
            )
        return hidden_states_to_logits, [hidden_states_to_aux]


class Eagle3DeepseekV2ForCausalLM(DeepseekV3ForCausalLM):
    """Eagle3 MLA draft model for DeepSeek-V2/V3 / Kimi-K2 style architectures.

    Inherits weight-loading fusion logic from DeepseekV3ForCausalLM but uses
    Eagle3MlaModel internally with a single MLA decoder layer that accepts
    concatenated [embeds || hidden_states] as input.
    """

    def __init__(
        self,
        config: PretrainedConfig,
        mapping: Mapping,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        nn.Module.__init__(self)
        self.config = config
        self.mapping = mapping
        self.quant_config = quant_config

        if self.config.num_hidden_layers != 1:
            raise ValueError("Eagle3 MLA drafter currently only supports 1 layer")

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

        self.load_lm_head_from_target = False
        if self.config.tie_word_embeddings:
            self.lm_head = self.model.embed_tokens
        else:
            draft_vocab_size = (
                getattr(config, "draft_vocab_size", None) or config.vocab_size
            )
            if not hasattr(config, "draft_vocab_size"):
                self.load_lm_head_from_target = True
            self.lm_head = ParallelLMHead(
                draft_vocab_size,
                config.hidden_size,
                quant_config=quant_config,
                tp_rank=self.mapping.attn.tp_rank,
                tp_size=self.mapping.attn.tp_size,
                tp_group=self.mapping.attn.tp_group,
                prefix=add_prefix("lm_head", prefix),
            )

        self.logits_processor = LogitsProcessor(
            config,
            skip_all_gather=self.mapping.attn.has_dp,
            do_argmax=True,
            tp_rank=self.mapping.attn.tp_rank,
            tp_size=self.mapping.attn.tp_size,
            tp_group=self.mapping.attn.tp_group,
        )
        self.capture_aux_hidden_states = True
        self.hot_token_id = None

    def forward(
        self,
        ctx: ForwardContext,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        out_cache_loc: torch.Tensor,
        **kwargs,
    ) -> torch.Tensor:
        with report_collective_sizing(ctx, ctx.bs, ctx.global_bs):
            return super().forward(ctx, input_ids, positions, out_cache_loc, **kwargs)

    def prepare_model_kwargs(
        self, ctx: ForwardContext, input_ids: torch.Tensor, kwargs: dict
    ) -> dict:
        model_kwargs = super().prepare_model_kwargs(ctx, input_ids, kwargs)
        captured_hidden_states = kwargs.get("captured_hidden_states")
        if captured_hidden_states is not None:
            model_kwargs["captured_hidden_states"] = captured_hidden_states
        else:
            # During CUDA graph capture warmup, provide dummy hidden states.
            target_hidden_size = getattr(
                self.config, "target_hidden_size", self.config.hidden_size
            )
            num_fc = self.model.num_fc_input_dim
            model_kwargs["captured_hidden_states"] = torch.zeros(
                input_ids.size(0),
                target_hidden_size * num_fc,
                dtype=torch.bfloat16,
                device=input_ids.device,
            )
        return model_kwargs

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
        remapped = []
        for name, loaded_weight in weights:
            if "d2t" in name:
                self.hot_token_id = loaded_weight + torch.arange(loaded_weight.size(0))
                continue
            if "t2d" in name:
                continue

            new_name = re.sub(r"^layers\.0\.", "midlayer.", name)

            if "lm_head" not in new_name:
                new_name = f"model.{new_name}"
            else:
                self.load_lm_head_from_target = False
            remapped.append((new_name, loaded_weight))

        super().load_weights(remapped)

    def post_load_weights(self):
        self_attn = self.model.midlayer.self_attn
        if (
            self.quant_config is not None
            and hasattr(self.quant_config, "weight_block_size")
            and self_attn.kv_b_proj.weight.dtype
            in (torch.float8_e4m3fn, torch.float8_e4m3fnuz)
        ):
            weight_block_size = self.quant_config.weight_block_size
            if weight_block_size is not None:
                if not hasattr(self_attn.kv_b_proj, "weight_scale_inv"):
                    raise RuntimeError(
                        "kv_b_proj.weight_scale_inv is required for block FP8 dequant."
                    )
                dtype = torch.get_default_dtype()
                w = block_dequant(
                    self_attn.kv_b_proj.weight,
                    self_attn.kv_b_proj.weight_scale_inv,
                    weight_block_size,
                ).to(dtype)
            else:
                w = self_attn.kv_b_proj.weight
        else:
            w = self_attn.kv_b_proj.weight

        self_attn.w_kc, self_attn.w_vc = _prepare_mla_kv_b_proj_weights(w, self_attn)

    def get_hot_token_id(self):
        return self.hot_token_id

    def set_embed_and_head(self, embed, head):
        if (
            hasattr(self.config, "target_hidden_size")
            and self.config.target_hidden_size != self.config.hidden_size
        ):
            return
        del self.model.embed_tokens.weight
        self.model.embed_tokens.weight = embed
        if head is not None and self.load_lm_head_from_target:
            del self.lm_head.weight
            self.lm_head.weight = head
        torch.cuda.empty_cache()
        torch.cuda.synchronize()


EntryClass = [
    DeepseekV3ForCausalLM,
    Eagle3DeepseekV2ForCausalLM,
]