lightseekorg--tokenspeed/python/tokenspeed/runtime/models/deepseek_v4_mtp.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 DeepSeek V4 MTP / NextN draft model."""

from __future__ import annotations

import logging
import re
from collections.abc import Iterable

import torch
from torch import nn
from transformers import PretrainedConfig

from tokenspeed.runtime.distributed.mapping import Mapping
from tokenspeed.runtime.execution.context import ForwardContext
from tokenspeed.runtime.layers.layernorm import RMSNorm
from tokenspeed.runtime.layers.linear import ReplicatedLinear
from tokenspeed.runtime.layers.logits_processor import LogitsMetadata, LogitsProcessor
from tokenspeed.runtime.layers.moe import (
    ExpertCheckpointSchema,
    build_moe_checkpoint_loader,
)
from tokenspeed.runtime.layers.moe.expert import MoELayer
from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig
from tokenspeed.runtime.layers.vocab_parallel_embedding import (
    ParallelLMHead,
    VocabParallelEmbedding,
)
from tokenspeed.runtime.model_loader.weight_utils import default_weight_loader
from tokenspeed.runtime.models.deepseek_v4 import (
    DeepseekV4Compressor,
    DeepseekV4DecoderLayer,
    DeepseekV4MegaMoEExperts,
    _deepseek_v4_swa_slot_mapping,
    hc_head,
    mhc_post,
)
from tokenspeed.runtime.utils import add_prefix

logger = logging.getLogger(__name__)


_EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$")


def _spec_layer_idx(config: PretrainedConfig, weight_name: str) -> int | None:
    if getattr(config, "num_nextn_predict_layers", 0) <= 0:
        return None
    start = config.num_hidden_layers
    for idx in range(start, start + config.num_nextn_predict_layers):
        if weight_name.startswith(f"model.layers.{idx}."):
            return idx
    return None


def _find_mtp_layer_idx(name: str) -> int:
    parts = name.split(".")
    if len(parts) > 1 and parts[0] == "mtp":
        try:
            return int(parts[1])
        except ValueError:
            pass
    for part in parts:
        try:
            return int(part)
        except ValueError:
            continue
    return 0


class DeepseekV4MTPSharedHead(nn.Module):
    def __init__(self, config: PretrainedConfig) -> None:
        super().__init__()
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)


class DeepseekV4MultiTokenPredictorLayer(nn.Module):
    def __init__(
        self,
        config: PretrainedConfig,
        mapping: Mapping,
        layer_id: int,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
        cache_layer_index: int | None = None,
    ) -> None:
        super().__init__()
        self.config = config
        self.layer_id = layer_id
        self.cache_layer_index = (
            layer_id if cache_layer_index is None else cache_layer_index
        )
        self.rms_norm_eps = config.rms_norm_eps
        self.hc_eps = config.hc_eps
        self.hc_mult = config.hc_mult

        self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.e_proj = ReplicatedLinear(
            config.hidden_size,
            config.hidden_size,
            bias=False,
            quant_config=quant_config,
            prefix=add_prefix("e_proj", prefix),
        )
        self.h_proj = ReplicatedLinear(
            config.hidden_size,
            config.hidden_size,
            bias=False,
            quant_config=quant_config,
            prefix=add_prefix("h_proj", prefix),
        )
        self.hc_head_fn = nn.Parameter(
            torch.empty(
                self.hc_mult,
                self.hc_mult * config.hidden_size,
                dtype=torch.float32,
            ),
            requires_grad=False,
        )
        self.hc_head_base = nn.Parameter(
            torch.empty(self.hc_mult, dtype=torch.float32),
            requires_grad=False,
        )
        self.hc_head_scale = nn.Parameter(
            torch.empty(1, dtype=torch.float32),
            requires_grad=False,
        )
        self.shared_head = DeepseekV4MTPSharedHead(config)
        self.mtp_block = DeepseekV4DecoderLayer(
            config,
            layer_id,
            mapping,
            quant_config,
            add_prefix("mtp_block", prefix),
            cache_layer_index=self.cache_layer_index,
        )

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        previous_hidden_states: torch.Tensor,
        ctx: ForwardContext,
        out_cache_loc: torch.Tensor,
        input_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor:
        if input_embeds is None:
            raise ValueError("DeepSeek V4 MTP requires input_embeds.")
        input_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, input_embeds)
        input_embeds = self.enorm(input_embeds)
        previous_hidden_states = previous_hidden_states.view(
            -1, self.hc_mult, self.config.hidden_size
        )
        previous_hidden_states = self.hnorm(previous_hidden_states)
        h_out, _ = self.h_proj(previous_hidden_states)
        e_out, _ = self.e_proj(input_embeds)
        hidden_states = h_out + e_out.unsqueeze(-2)

        swa_slot_mapping = _deepseek_v4_swa_slot_mapping(
            ctx,
            positions,
            out_cache_loc,
        )
        residual, x_def, post_def, comb_def = self.mtp_block(
            positions,
            hidden_states,
            ctx,
            out_cache_loc,
            input_ids,
            swa_slot_mapping,
        )
        return mhc_post(x_def, residual, post_def, comb_def)

    def compute_logits_hidden(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = hidden_states.view(-1, self.hc_mult, self.config.hidden_size)
        hidden_states = hc_head(
            hidden_states,
            self.hc_head_fn,
            self.hc_head_scale,
            self.hc_head_base,
            self.rms_norm_eps,
            self.hc_eps,
        )
        return self.shared_head.norm(hidden_states)


class DeepseekV4MultiTokenPredictor(nn.Module):
    def __init__(
        self,
        config: PretrainedConfig,
        mapping: Mapping,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.config = config
        self.mapping = mapping
        self.mtp_start_layer_idx = config.num_hidden_layers
        self.num_mtp_layers = config.num_nextn_predict_layers
        self.embed_tokens = VocabParallelEmbedding(
            config.vocab_size,
            config.hidden_size,
            tp_rank=mapping.attn.tp_rank,
            tp_size=mapping.attn.tp_size,
            tp_group=mapping.attn.tp_group,
            prefix=add_prefix("embed_tokens", prefix),
        )
        layers = {}
        for local_idx in range(self.num_mtp_layers):
            # Checkpoint layer ids remain global, while draft KV slots are compact.
            layer_idx = self.mtp_start_layer_idx + local_idx
            layers[str(layer_idx)] = DeepseekV4MultiTokenPredictorLayer(
                config,
                mapping,
                layer_idx,
                quant_config,
                add_prefix(f"layers.{layer_idx}", prefix),
                cache_layer_index=local_idx,
            )
        self.layers = nn.ModuleDict(layers)

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        previous_hidden_states: torch.Tensor,
        ctx: ForwardContext,
        out_cache_loc: torch.Tensor,
        input_embeds: torch.Tensor | None = None,
        spec_step_idx: int = 0,
    ) -> torch.Tensor:
        if input_embeds is None:
            input_embeds = self.embed_tokens(input_ids)
        current_step_idx = spec_step_idx % self.num_mtp_layers
        layer_idx = self.mtp_start_layer_idx + current_step_idx
        return self.layers[str(layer_idx)](
            input_ids,
            positions,
            previous_hidden_states,
            ctx,
            out_cache_loc,
            input_embeds,
        )

    def compute_logits_hidden(
        self,
        hidden_states: torch.Tensor,
        spec_step_idx: int = 0,
    ) -> torch.Tensor:
        current_step_idx = spec_step_idx % self.num_mtp_layers
        layer_idx = self.mtp_start_layer_idx + current_step_idx
        return self.layers[str(layer_idx)].compute_logits_hidden(hidden_states)


class DeepseekV4ForCausalLMNextN(nn.Module):
    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
        self.model = DeepseekV4MultiTokenPredictor(
            config,
            mapping=mapping,
            quant_config=quant_config,
            prefix=add_prefix("model", prefix),
        )
        if self.mapping.attn.has_dp:
            self.lm_head = ReplicatedLinear(
                config.hidden_size,
                config.vocab_size,
                bias=False,
                prefix=add_prefix("lm_head", prefix),
            )
        else:
            self.lm_head = ParallelLMHead(
                config.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,
        )

    def get_hot_token_id(self):
        return None

    def get_embed_and_head(self) -> tuple[torch.Tensor, torch.Tensor]:
        return self.model.embed_tokens.weight, self.lm_head.weight

    def set_embed_and_head(self, embed: torch.Tensor, head: torch.Tensor) -> None:
        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()

    @torch.no_grad()
    def forward(
        self,
        ctx: ForwardContext,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        out_cache_loc: torch.Tensor,
        input_embeds: torch.Tensor | None = None,
        captured_hidden_states: torch.Tensor | None = None,
        spec_step_idx: int = 0,
        **kwargs,
    ):
        del kwargs
        if captured_hidden_states is None:
            if not ctx.forward_mode.is_idle():
                raise ValueError("DeepSeek V4 MTP requires captured_hidden_states.")
            captured_hidden_states = torch.zeros(
                0,
                self.config.hc_mult * self.config.hidden_size,
                device=input_ids.device,
                dtype=self.model.embed_tokens.weight.dtype,
            )

        mtp_hidden_states = self.model(
            input_ids,
            positions,
            captured_hidden_states,
            ctx,
            out_cache_loc,
            input_embeds=input_embeds,
            spec_step_idx=spec_step_idx,
        ).flatten(1)
        logits_hidden_states = self.model.compute_logits_hidden(
            mtp_hidden_states,
            spec_step_idx,
        )
        logits_metadata = LogitsMetadata.from_forward_context(ctx)
        return self.logits_processor(
            input_ids,
            logits_hidden_states,
            self.lm_head,
            logits_metadata,
            aux_hidden_states=[mtp_hidden_states],
        )

    @staticmethod
    def _remap_weight_name(name: str) -> str:
        for old, new in {
            ".emb.tok_emb.weight": ".embed_tokens.weight",
            ".head.weight": ".shared_head.head.weight",
            ".norm.weight": ".shared_head.norm.weight",
        }.items():
            if old in name:
                name = name.replace(old, new)
        return name

    @staticmethod
    def _rewrite_spec_layer_name(spec_layer: int, name: str) -> str:
        spec_layer_weight_names = (
            "embed_tokens",
            "enorm",
            "hnorm",
            "h_proj",
            "e_proj",
            "shared_head",
            "hc_head_fn",
            "hc_head_base",
            "hc_head_scale",
        )
        shared_weight_names = ("embed_tokens",)
        is_spec_weight = any(
            weight_name in name for weight_name in spec_layer_weight_names
        )
        is_shared_weight = any(
            weight_name in name for weight_name in shared_weight_names
        )
        if not is_spec_weight:
            name = name.replace(
                f"model.layers.{spec_layer}.",
                f"model.layers.{spec_layer}.mtp_block.",
            )
        elif is_shared_weight:
            name = name.replace(f"model.layers.{spec_layer}.", "model.")
        return name

    def _map_checkpoint_name(self, raw_name: str) -> str | None:
        if raw_name.startswith("mtp."):
            mtp_layer_idx = _find_mtp_layer_idx(raw_name)
            raw_name = raw_name.replace(
                f"mtp.{mtp_layer_idx}.",
                f"model.layers.{self.config.num_hidden_layers + mtp_layer_idx}.",
                1,
            )
        spec_layer = _spec_layer_idx(self.config, raw_name)
        if spec_layer is None:
            return None
        name = self._remap_weight_name(raw_name)
        name = self._rewrite_spec_layer_name(spec_layer, name)
        if name.endswith(".shared_head.head.weight"):
            return None
        if name.endswith(".scale"):
            suffix = (
                ".weight_scale"
                if _EXPERT_SCALE_RE.search(name)
                else ".weight_scale_inv"
            )
            name = name.removesuffix(".scale") + suffix
        if ".shared_experts.w2" in name:
            name = name.replace(".shared_experts.w2", ".shared_experts.down_proj")
        if ".ffn.gate.bias" in name:
            name = name.replace(".ffn.gate.bias", ".ffn.gate.e_score_correction_bias")
        return name

    def get_stacked_params_mapping(self):
        return [
            ("gate_up_proj", "w1", 0),
            ("gate_up_proj", "w3", 1),
            ("attn.fused_wqa_wkv", "attn.wq_a", 0),
            ("attn.fused_wqa_wkv", "attn.wkv", 1),
            ("compressor.fused_wkv_wgate", "compressor.wkv", 0),
            ("compressor.fused_wkv_wgate", "compressor.wgate", 1),
        ]

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
        stacked_params_mapping = self.get_stacked_params_mapping()
        params_dict = dict(self.named_parameters())
        moe_loader = build_moe_checkpoint_loader(
            params_dict=params_dict,
            expert_schema=ExpertCheckpointSchema(
                gate_proj_name="w1",
                down_proj_name="w2",
                up_proj_name="w3",
            ),
            num_experts=self.config.n_routed_experts,
            ep_rank=self.mapping.moe.ep_rank,
            ep_size=self.mapping.moe.ep_size,
        )
        loaded_params: set[str] = set()
        for raw_name, loaded_weight in weights:
            name = self._map_checkpoint_name(raw_name)
            if name is None:
                continue
            for param_name, weight_name, shard_id in stacked_params_mapping:
                if weight_name not in name or ".experts." in name:
                    continue
                mapped_name = name.replace(weight_name, param_name)
                param = params_dict.get(mapped_name)
                if param is None:
                    break
                param.weight_loader(param, loaded_weight, shard_id)
                loaded_params.add(mapped_name)
                break
            else:
                if moe_loader.matches(name):
                    mapped_name = moe_loader.load(name, loaded_weight)
                    loaded_params.add(mapped_name)
                    continue
                param = params_dict.get(name)
                if param is None:
                    logger.debug("Skipping unmatched DeepSeek V4 MTP weight: %s", name)
                    continue
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
                weight_loader(param, loaded_weight)
                loaded_params.add(name)

        missing_layers = []
        for layer_idx in range(
            self.model.mtp_start_layer_idx,
            self.model.mtp_start_layer_idx + self.model.num_mtp_layers,
        ):
            if not any(f"model.layers.{layer_idx}." in name for name in loaded_params):
                missing_layers.append(layer_idx)
        if missing_layers:
            raise ValueError(
                "DeepSeek V4 MTP weights missing for speculative layer(s) "
                f"{missing_layers}. Use a checkpoint that includes `mtp.*` "
                "weights or disable NEXTN speculative decoding."
            )
        self.post_load_weights()
        return loaded_params

    def post_load_weights(self):
        for module in self.modules():
            if isinstance(module, DeepseekV4Compressor):
                module.process_weights_after_loading()
            elif isinstance(module, DeepseekV4MegaMoEExperts):
                module.finalize_weights()
            elif isinstance(module, MoELayer):
                module.process_weights_after_loading(module)


EntryClass = [DeepseekV4ForCausalLMNextN]