lightseekorg--tokenspeed/python/tokenspeed/runtime/models/glm5_nextn.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 GLM5 NextN speculative decoding."""

from __future__ import annotations

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

import torch
from torch import nn
from transformers import PretrainedConfig

from tokenspeed.runtime.distributed.mapping import Mapping
from tokenspeed.runtime.execution.context import (
    ForwardContext,
    report_collective_sizing,
)
from tokenspeed.runtime.layers.attention.dsa.utils import workspace_indices_to_kv_slots
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.quantization.base_config import QuantizationConfig
from tokenspeed.runtime.layers.quantization.utils import block_dequant
from tokenspeed.runtime.layers.utils import (
    CP_METADATA,
    ENABLE_CP,
    cp_all_gather_rerange_output,
    cp_split_and_rebuild_data,
)
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.glm5 import (
    GlmMoeDsaDecoderLayer,
    GlmMoeDsaForCausalLM,
    pad_fused_qkv_a_proj_weight_for_fp8_blockscale,
)

_NEXTN_SPEC_WEIGHT_NAMES = (
    "shared_head.norm",
    "eh_proj",
    "enorm",
    "hnorm",
)

_STACKED_PARAMS_MAPPING = (
    ("gate_up_proj", "gate_proj", 0),
    ("gate_up_proj", "up_proj", 1),
)


class GlmMoeDsaModelNextN(nn.Module):
    def __init__(
        self,
        config: PretrainedConfig,
        mapping: Mapping,
        quant_config: QuantizationConfig | None = None,
    ) -> None:
        super().__init__()
        self.mapping = mapping
        self.vocab_size = config.vocab_size

        self.embed_tokens = VocabParallelEmbedding(
            config.vocab_size,
            config.hidden_size,
            tp_rank=self.mapping.attn.tp_rank,
            tp_size=self.mapping.attn.tp_size,
            tp_group=self.mapping.attn.tp_group,
        )

        self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

        self.eh_proj = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False)

        self.alt_stream = torch.cuda.Stream()
        self.decoder = GlmMoeDsaDecoderLayer(
            config,
            0,
            mapping=self.mapping,
            quant_config=quant_config,
            is_nextn=True,
            alt_stream=self.alt_stream,
        )

        self.shared_head = nn.Module()
        self.shared_head.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, None]:
        if input_embeds is None:
            hidden_states = self.embed_tokens(input_ids)
        else:
            hidden_states = input_embeds
        hidden_states = torch.where(positions.unsqueeze(-1) == 0, 0, hidden_states)
        if captured_hidden_states is None:
            if not ctx.forward_mode.is_idle():
                raise ValueError("GLM5 NextN requires captured_hidden_states.")
            captured_hidden_states = hidden_states

        hidden_states = self.eh_proj(
            torch.cat(
                (
                    self.enorm(hidden_states),
                    self.hnorm(captured_hidden_states),
                ),
                dim=-1,
            )
        )

        residual = None
        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,
            )
        hidden_states, residual = self.decoder(
            positions,
            hidden_states,
            ctx,
            out_cache_loc,
            residual,
        )

        if not ctx.forward_mode.is_idle():
            if not ENABLE_CP:
                hidden_states, _ = self.decoder.comm_manager.final_norm(
                    hidden_states, residual, ctx, self.shared_head.norm
                )
            else:
                hidden_states, _ = self.shared_head.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, None


class GlmMoeDsaForCausalLMNextN(GlmMoeDsaForCausalLM):
    compute_dsa_topk_first_step = True

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

        if quant_config is not None and quant_config.get_name() == "nvfp4":
            quant_config = None

        self.quant_config = quant_config

        self.model = GlmMoeDsaModelNextN(
            config, mapping=self.mapping, quant_config=quant_config
        )

        if self.mapping.attn.has_dp:
            self.lm_head = ReplicatedLinear(
                config.hidden_size,
                config.vocab_size,
                bias=False,
            )
        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,
            )
        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,
        )

    @staticmethod
    def _apply_first_step_correction(ctx: ForwardContext) -> None:
        seq_lens_buf = ctx.draft_seq_lens_buf
        accept_lengths = ctx.accept_lengths
        if seq_lens_buf is None or accept_lengths is None:
            return
        num_extends = ctx.num_extends
        if num_extends >= ctx.bs:
            return
        correction = (
            ctx.attn_backend.spec_num_tokens - accept_lengths[num_extends:]
        ).to(seq_lens_buf.dtype)
        seq_lens_buf[num_extends : ctx.bs].sub_(correction).clamp_(min=1)
        ctx.attn_backend.advance_draft_forward_metadata(seq_lens_buf[: ctx.bs])

    @staticmethod
    def prepare_dsa_topk_for_mtp_decode(
        dsa_topk: tuple[Any | None, Any | None],
        gather_ids: torch.Tensor,
        *,
        num_prefill_rows: int = 0,
    ) -> tuple[Any | None, Any | None]:
        prefill_topk, decode_topk = dsa_topk
        if decode_topk is None:
            return dsa_topk
        topk_indices = decode_topk.topk_indices
        topk_lens = decode_topk.topk_lens
        if topk_indices.shape[0] == 0:
            return dsa_topk
        if num_prefill_rows <= 0 and topk_indices.shape[0] <= gather_ids.numel():
            return dsa_topk
        if num_prefill_rows <= 0:
            selected_indices = topk_indices.index_select(0, gather_ids)
            selected_lens = topk_lens.index_select(0, gather_ids)
        else:
            if prefill_topk is None:
                return dsa_topk
            num_prefill_rows = min(int(num_prefill_rows), gather_ids.numel())
            prefill_row_ids = gather_ids[:num_prefill_rows]
            decode_row_ids = gather_ids[num_prefill_rows:]
            selected_prefill_indices = workspace_indices_to_kv_slots(
                prefill_topk.workspace_indices.index_select(0, prefill_row_ids),
                prefill_topk.kv_workspace_slots,
            ).to(device=topk_indices.device, dtype=topk_indices.dtype)
            selected_prefill_lens = prefill_topk.topk_lens.index_select(
                0,
                prefill_row_ids,
            ).to(
                device=topk_lens.device,
                dtype=topk_lens.dtype,
            )
            if decode_row_ids.numel() > 0:
                selected_decode_indices = topk_indices.index_select(0, decode_row_ids)
                selected_decode_lens = topk_lens.index_select(0, decode_row_ids)
                selected_indices = torch.cat(
                    [selected_prefill_indices, selected_decode_indices],
                    dim=0,
                )
                selected_lens = torch.cat(
                    [selected_prefill_lens, selected_decode_lens],
                    dim=0,
                )
            else:
                selected_indices = selected_prefill_indices
                selected_lens = selected_prefill_lens
        selected_decode_topk = replace(
            decode_topk,
            topk_indices=selected_indices,
            topk_lens=selected_lens,
        )
        return prefill_topk, selected_decode_topk

    @torch.no_grad()
    def forward(
        self,
        ctx: ForwardContext,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        out_cache_loc: torch.Tensor,
        captured_hidden_states: torch.Tensor | None = None,
    ) -> torch.Tensor:
        with report_collective_sizing(ctx, ctx.bs, ctx.global_bs):
            hidden_states, _ = self.model(
                input_ids,
                positions,
                ctx,
                out_cache_loc,
                captured_hidden_states=captured_hidden_states,
            )
        self._apply_first_step_correction(ctx)
        logits_metadata = LogitsMetadata.from_forward_context(ctx)
        return self.logits_processor(
            input_ids, hidden_states, self.lm_head, logits_metadata
        )

    def get_hot_token_id(self) -> None:
        return None

    def _nextn_layer_prefix(self, name: str) -> str | None:
        if not hasattr(self.config, "num_nextn_predict_layers"):
            raise ValueError("num_nextn_predict_layers is not in the config")
        if self.config.num_nextn_predict_layers != 1:
            raise ValueError("Only 1 nextn layer is supported")

        if self.config.num_nextn_predict_layers == self.config.num_hidden_layers:
            prefix = "model.layers.0"
            return prefix if name.startswith(prefix) else None

        if not name.startswith("model.layers."):
            return None
        name_parts = name.split(".")
        if len(name_parts) < 3:
            return None
        try:
            layer_id = int(name_parts[2])
        except ValueError:
            return None
        if layer_id < self.config.num_hidden_layers:
            return None
        return f"model.layers.{layer_id}"

    def _map_checkpoint_name(self, raw_name: str) -> str | None:
        nextn_layer_prefix = self._nextn_layer_prefix(raw_name)
        if nextn_layer_prefix is None:
            return None
        if "shared_head.head" in raw_name or "embed_tokens" in raw_name:
            return None
        if "rotary_emb.inv_freq" in raw_name:
            return None

        if any(weight_name in raw_name for weight_name in _NEXTN_SPEC_WEIGHT_NAMES):
            return raw_name.replace(nextn_layer_prefix, "model")
        return raw_name.replace(nextn_layer_prefix, "model.decoder")

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> None:
        fuse_qkv_a_proj = hasattr(self.config, "q_lora_rank") and (
            self.config.q_lora_rank is not None
        )
        cached_a_proj: dict[str, torch.Tensor] | None = {} if fuse_qkv_a_proj else None

        params_dict = dict(self.named_parameters())
        modules_dict = dict(self.named_modules())
        pending_fp8_wk: dict[str, dict[str, torch.Tensor]] = {}
        loaded_fused_indexer_shards: dict[str, set[int]] = {}

        moe_loader = build_moe_checkpoint_loader(
            params_dict=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 raw_name, loaded_weight in weights:
            name = self._map_checkpoint_name(raw_name)
            if name is None:
                continue

            if ".indexer." in name:
                if name.endswith(".bias") and name not in params_dict:
                    continue
                param = self.get_param(params_dict, name)
                if param is not None:
                    weight_loader = getattr(
                        param, "weight_loader", default_weight_loader
                    )
                    weight_loader(param, loaded_weight)
                self._try_load_fused_indexer_projection(
                    name=name,
                    loaded_weight=loaded_weight,
                    params_dict=params_dict,
                    modules_dict=modules_dict,
                    pending_fp8_wk=pending_fp8_wk,
                    loaded_shards=loaded_fused_indexer_shards,
                )
                continue

            for param_name, weight_name, shard_id in _STACKED_PARAMS_MAPPING:
                if weight_name not in name:
                    continue
                if ("mlp.experts." in name) and name not in params_dict:
                    continue
                name = name.replace(weight_name, param_name)
                if name.endswith(".bias") and name not in params_dict:
                    continue
                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
                if name.endswith(".bias") and name not in params_dict:
                    continue
                if moe_loader.matches(name):
                    moe_loader.load(name, loaded_weight)
                    continue

                if cached_a_proj is not None and (
                    "q_a_proj" in name or "kv_a_proj_with_mqa" in name
                ):
                    cached_a_proj[name] = loaded_weight
                    q_a_proj_name = (
                        name
                        if "q_a_proj" in name
                        else name.replace("kv_a_proj_with_mqa", "q_a_proj")
                    )
                    kv_a_proj_name = (
                        name
                        if "kv_a_proj_with_mqa" in name
                        else name.replace("q_a_proj", "kv_a_proj_with_mqa")
                    )

                    if (
                        q_a_proj_name in cached_a_proj
                        and kv_a_proj_name in cached_a_proj
                    ):
                        q_a_proj_weight = cached_a_proj[q_a_proj_name]
                        kv_a_proj_weight = cached_a_proj[kv_a_proj_name]
                        fused_weight = torch.cat(
                            [q_a_proj_weight, kv_a_proj_weight], dim=0
                        )

                        if "q_a_proj" in name:
                            param_name = name.replace(
                                "q_a_proj", "fused_qkv_a_proj_with_mqa"
                            )
                        else:
                            param_name = name.replace(
                                "kv_a_proj_with_mqa", "fused_qkv_a_proj_with_mqa"
                            )
                        param = params_dict[param_name]

                        weight_loader = getattr(
                            param, "weight_loader", default_weight_loader
                        )
                        weight_loader(param, fused_weight)
                        cached_a_proj.pop(q_a_proj_name)
                        cached_a_proj.pop(kv_a_proj_name)
                else:
                    if ".mlp.experts." in name:
                        continue
                    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) -> None:
        self_attn = self.model.decoder.self_attn
        pad_fused_qkv_a_proj_weight_for_fp8_blockscale(self_attn)
        if (
            hasattr(self.quant_config, "weight_block_size")
            and (self.quant_config.weight_block_size is not None)
            and self_attn.kv_b_proj.weight.dtype
            in (
                torch.float8_e4m3fn,
                torch.float8_e4m3fnuz,
            )
        ):
            weight_block_size = self.quant_config.weight_block_size
            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

        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)
        self_attn.w_kc = w_kc.transpose(1, 2).contiguous().transpose(1, 2)
        self_attn.w_vc = w_vc.contiguous().transpose(1, 2)


EntryClass = [GlmMoeDsaForCausalLMNextN]