reference/code/Diffulex-mbd-lms/diffulex/model/llada2.py

from __future__ import annotations

import copy
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange, reduce, repeat

from diffulex.attention import Attention
from diffulex.layer.activation import SiluAndMul
from diffulex.layer.embed_head import ParallelLMHead, VocabParallelEmbedding
from diffulex.layer.layernorm import RMSNorm
from diffulex.layer.linear import ColumnParallelLinear, RowParallelLinear, divide
from diffulex.layer.rotary_embedding import get_rope
from diffulex.model.auto_model import AutoModelForDiffusionLM
from diffulex.moe.config import get_norm_topk_prob, is_moe_layer
from diffulex.moe.layer.base import FusedMoE
from diffulex.moe.layer.ep_impl import EPFusedMoE
from diffulex.moe.layer.tp_impl import TPFusedMoE
from diffulex.moe.layer.naive_impl import NaiveFusedMoE
from diffulex.moe.topk import GroupLimitedTopKRouter
from diffulex.distributed.parallel_state import fetch_parallel_state
from diffulex.utils.checkpoint import LoadContext, ResolvedWeight


def _llada2_use_reference_view_path() -> bool:
    return os.getenv("DIFFULEX_LLADA2_REFERENCE_VIEW_PATH", "0") == "1"


def _llada2_use_legacy_qkv_path() -> bool:
    return os.getenv("DIFFULEX_LLADA2_LEGACY_QKV_PATH", "0") == "1"


def _llada2_gate_use_fp32() -> bool:
    # Keep fp32 gate as default for alignment; allow explicit opt-out for perf probing.
    return os.getenv("DIFFULEX_LLADA2_GATE_FP32", "1") != "0"


class LLaDA2QKVParallelLinear(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        head_size: int,
        total_num_heads: int,
        total_num_kv_heads: int,
        *,
        bias: bool = False,
    ) -> None:
        super().__init__()
        self.hidden_size = hidden_size
        self.head_size = head_size
        self.total_num_heads = total_num_heads
        self.total_num_kv_heads = total_num_kv_heads
        parallel_state = fetch_parallel_state()
        self.tp_size = parallel_state.get_tp_world_size()
        self.tp_rank = parallel_state.get_tp_rank()
        self.num_heads = divide(total_num_heads, self.tp_size)
        self.num_kv_heads = divide(total_num_kv_heads, self.tp_size)
        self.q_size = self.num_heads * head_size
        self.kv_size = self.num_kv_heads * head_size
        self.weight = nn.Parameter(torch.empty(self.q_size + 2 * self.kv_size, hidden_size))
        self.weight.weight_loader = self.weight_loader
        if bias:
            self.bias = nn.Parameter(torch.empty(self.q_size + 2 * self.kv_size))
            self.bias.weight_loader = self.weight_loader
        else:
            self.register_parameter("bias", None)

    def _local_qkv(self, loaded_weight: torch.Tensor) -> torch.Tensor:
        q_total = self.total_num_heads * self.head_size
        kv_total = self.total_num_kv_heads * self.head_size
        q, k, v = loaded_weight.split((q_total, kv_total, kv_total), dim=0)
        q = q.chunk(self.tp_size, dim=0)[self.tp_rank]
        k = k.chunk(self.tp_size, dim=0)[self.tp_rank]
        v = v.chunk(self.tp_size, dim=0)[self.tp_rank]
        return torch.cat((q, k, v), dim=0).contiguous()

    def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor) -> None:
        param.data.copy_(self._local_qkv(loaded_weight))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return F.linear(x, self.weight, self.bias)


class LLaDA2Attention(nn.Module):
    def __init__(self, config, layer_idx: int) -> None:
        super().__init__()
        self.layer_idx = layer_idx
        parallel_state = fetch_parallel_state()
        tp_size = parallel_state.get_tp_world_size()
        self.total_num_heads = config.num_attention_heads
        self.total_num_kv_heads = config.num_key_value_heads or config.num_attention_heads
        self.num_heads = divide(self.total_num_heads, tp_size)
        self.num_kv_heads = divide(self.total_num_kv_heads, tp_size)
        self.head_dim = getattr(config, "head_dim", None) or config.hidden_size // self.total_num_heads
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim
        self.scaling = self.head_dim**-0.5

        self.query_key_value = LLaDA2QKVParallelLinear(
            config.hidden_size,
            self.head_dim,
            self.total_num_heads,
            self.total_num_kv_heads,
            bias=bool(getattr(config, "use_qkv_bias", False)),
        )
        self.query_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.key_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.dense = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            config.hidden_size,
            bias=bool(getattr(config, "use_bias", False)),
        )
        partial_rotary_factor = float(getattr(config, "partial_rotary_factor", 1.0))
        rotary_dim = int(self.head_dim * partial_rotary_factor)
        rotary_dim = int(getattr(config, "rotary_dim", rotary_dim) or rotary_dim)
        self.rotary_emb = get_rope(
            head_size=self.head_dim,
            rotary_dim=rotary_dim,
            max_position=config.max_position_embeddings,
            base=getattr(config, "rope_theta", 10000),
        )
        self.attn = Attention(
            self.num_heads,
            self.head_dim,
            self.scaling,
            self.num_kv_heads,
            attn_impl=getattr(config, "attn_impl", "triton"),
        )

    def _split_qkv_reference(
        self,
        qkv: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        tokens = qkv.size(0)
        qkv = qkv.view(tokens, self.num_heads + 2 * self.num_kv_heads, self.head_dim)
        q = qkv[:, : self.num_heads, :]
        k = qkv[:, self.num_heads : self.num_heads + self.num_kv_heads, :]
        v = qkv[:, self.num_heads + self.num_kv_heads :, :]
        return q, k, v

    def _split_qkv_default(
        self,
        qkv: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        q, k, v = qkv.split((self.q_size, self.kv_size, self.kv_size), dim=-1)
        q = rearrange(q, "token (head head_dim) -> token head head_dim", head=self.num_heads)
        k = rearrange(k, "token (head head_dim) -> token head head_dim", head=self.num_kv_heads)
        v = rearrange(v, "token (head head_dim) -> token head head_dim", head=self.num_kv_heads)
        return q, k, v

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        qkv = self.query_key_value(hidden_states)
        if _llada2_use_legacy_qkv_path():
            q, k, v = qkv.split((self.q_size, self.kv_size, self.kv_size), dim=-1)
            q = rearrange(
                self.query_layernorm(
                    rearrange(q, "token (head head_dim) -> token head head_dim", head=self.num_heads)
                ),
                "token head head_dim -> token (head head_dim)",
            )
            k = rearrange(
                self.key_layernorm(
                    rearrange(k, "token (head head_dim) -> token head head_dim", head=self.num_kv_heads)
                ),
                "token head head_dim -> token (head head_dim)",
            )
            q, k = self.rotary_emb(positions, q, k)
            return self.dense(self.attn(q, k, v, mask))
        if _llada2_use_reference_view_path():
            q, k, v = self._split_qkv_reference(qkv)
        else:
            q, k, v = self._split_qkv_default(qkv)
        q = self.query_layernorm(q)
        k = self.key_layernorm(k)
        q, k = self.rotary_emb(positions, q, k)
        attn_out = self.attn(q, k, v, mask)
        attn_out = self.dense(attn_out)
        return attn_out


class LLaDA2DenseMLP(nn.Module):
    def __init__(self, config, intermediate_size: int | None = None) -> None:
        super().__init__()
        intermediate_size = int(intermediate_size or config.intermediate_size)
        self.gate_proj = ColumnParallelLinear(config.hidden_size, intermediate_size, bias=False)
        self.up_proj = ColumnParallelLinear(config.hidden_size, intermediate_size, bias=False)
        self.down_proj = RowParallelLinear(intermediate_size, config.hidden_size, bias=False)
        if getattr(config, "hidden_act", "silu") != "silu":
            raise NotImplementedError("LLaDA2 dense MLP currently supports only silu.")
        self.act_fn = SiluAndMul()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(self.act_fn(torch.cat((self.gate_proj(x), self.up_proj(x)), dim=-1)))


class LLaDA2MoEMixin:
    def _init_llada2_moe(self, config) -> None:
        self.gate.register_buffer("expert_bias", torch.zeros((self.num_experts,), dtype=torch.float32))
        self.gate.forward = self._forward_llada2_gate
        self.router = GroupLimitedTopKRouter(
            top_k=self.top_k,
            num_experts=self.num_experts,
            n_group=int(getattr(config, "n_group", 0) or 0),
            topk_group=int(getattr(config, "topk_group", 0) or 0),
            routed_scaling_factor=float(getattr(config, "routed_scaling_factor", 1.0)),
            renormalize=get_norm_topk_prob(config),
            expert_bias_getter=lambda: self.gate.expert_bias,
        )

    def _forward_llada2_gate(self, hidden_states: torch.Tensor) -> torch.Tensor:
        if _llada2_gate_use_fp32():
            logits = F.linear(hidden_states.to(torch.float32), self.gate.weight.to(torch.float32))
        else:
            logits = F.linear(hidden_states, self.gate.weight)
        if os.getenv("DIFFULEX_DISABLE_EXPERT_BIAS", "0") == "1":
            # Debug toggle: keep checkpoint loading intact but ignore expert bias at runtime.
            self.gate.expert_bias.zero_()
        return logits

    def resolve_checkpoint_weight(self, suffix: str, ctx: LoadContext) -> ResolvedWeight | None:
        if suffix in {"gate.e_score_correction_bias", "gate.expert_bias"}:
            return ResolvedWeight(buffer=self.gate.expert_bias)

        return super().resolve_checkpoint_weight(suffix, ctx)


class LLaDA2NaiveMoE(LLaDA2MoEMixin, NaiveFusedMoE):
    @classmethod
    def from_config(cls, config) -> "LLaDA2NaiveMoE":
        num_shared_experts = int(getattr(config, "num_shared_experts", 0) or 0)
        module = cls(
            hidden_size=config.hidden_size,
            intermediate_size=int(config.moe_intermediate_size),
            num_experts=int(config.num_experts),
            top_k=int(config.num_experts_per_tok),
            hidden_act=getattr(config, "hidden_act", "silu"),
            norm_topk_prob=get_norm_topk_prob(config),
            moe_gemm_impl=getattr(config, "moe_gemm_impl", "triton"),
            num_shared_experts=num_shared_experts,
            shared_expert_intermediate_size=int(config.moe_intermediate_size) * num_shared_experts,
        )
        module._init_llada2_moe(config)
        return module


class LLaDA2TPMoE(LLaDA2MoEMixin, TPFusedMoE):
    @classmethod
    def from_config(cls, config) -> "LLaDA2TPMoE":
        num_shared_experts = int(getattr(config, "num_shared_experts", 0) or 0)
        module = cls(
            hidden_size=config.hidden_size,
            intermediate_size=int(config.moe_intermediate_size),
            num_experts=int(config.num_experts),
            top_k=int(config.num_experts_per_tok),
            hidden_act=getattr(config, "hidden_act", "silu"),
            norm_topk_prob=get_norm_topk_prob(config),
            moe_gemm_impl=getattr(config, "moe_gemm_impl", "triton"),
            num_shared_experts=num_shared_experts,
            shared_expert_intermediate_size=int(config.moe_intermediate_size) * num_shared_experts,
        )
        module._init_llada2_moe(config)
        return module


class LLaDA2EPMoE(LLaDA2MoEMixin, EPFusedMoE):
    @classmethod
    def from_config(cls, config) -> "LLaDA2EPMoE":
        num_shared_experts = int(getattr(config, "num_shared_experts", 0) or 0)
        module = cls(
            hidden_size=config.hidden_size,
            intermediate_size=int(config.moe_intermediate_size),
            num_experts=int(config.num_experts),
            top_k=int(config.num_experts_per_tok),
            hidden_act=getattr(config, "hidden_act", "silu"),
            norm_topk_prob=get_norm_topk_prob(config),
            moe_gemm_impl=getattr(config, "moe_gemm_impl", "triton"),
            dispatcher_backend=getattr(config, "moe_dispatcher_backend", "naive"),
            deepep_mode=getattr(config, "deepep_mode", "auto"),
            deepep_num_max_dispatch_tokens_per_rank=getattr(
                config,
                "deepep_num_max_dispatch_tokens_per_rank",
                256,
            ),
            num_shared_experts=num_shared_experts,
            shared_expert_intermediate_size=int(config.moe_intermediate_size) * num_shared_experts,
        )
        module._init_llada2_moe(config)
        return module


def build_llada2_mlp(config, layer_idx: int) -> nn.Module:
    if not is_moe_layer(config, layer_idx):
        return LLaDA2DenseMLP(config)
    parallel_state = fetch_parallel_state()
    if parallel_state.is_ep_enabled():
        module = LLaDA2EPMoE.from_config(config)
    elif parallel_state.is_tp_enabled():
        module = LLaDA2TPMoE.from_config(config)
    else:
        module = LLaDA2NaiveMoE.from_config(config)
    module._llada2_layer_idx = layer_idx
    return module


class LLaDA2DecoderLayer(nn.Module):
    def __init__(self, config, layer_idx: int) -> None:
        super().__init__()
        self.attention = LLaDA2Attention(config, layer_idx)
        self.mlp = build_llada2_mlp(config, layer_idx)
        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.layer_idx = layer_idx

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        hidden_states = residual + self.attention(positions, hidden_states, mask)
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        mlp_output = self.mlp(hidden_states)
        if isinstance(mlp_output, tuple):
            hidden_states, _router_logits = mlp_output
        else:
            hidden_states = mlp_output
        hidden_states = residual + hidden_states
        return hidden_states


class LLaDA2Model(nn.Module):
    def __init__(self, config) -> None:
        super().__init__()
        self.word_embeddings = VocabParallelEmbedding(config.vocab_size, config.hidden_size)
        self.mask_token_id = int(getattr(config, "mask_token_id", -1))
        if self.mask_token_id >= 0:
            self.register_buffer(
                "mask_token_id_tensor",
                torch.tensor([self.mask_token_id], dtype=torch.int64),
                persistent=False,
            )
        else:
            self.register_buffer("mask_token_id_tensor", None, persistent=False)
        self.layers = nn.ModuleList(
            [LLaDA2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
        )
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        hidden_states = self.word_embeddings(input_ids)
        hidden_states = self._maybe_apply_token_merging(hidden_states)
        for layer in self.layers:
            hidden_states = layer(positions, hidden_states, mask)
        hidden_states = self.norm(hidden_states)
        return hidden_states

    def _maybe_apply_token_merging(self, hidden_states: torch.Tensor) -> torch.Tensor:
        from diffulex.attention import fetch_attn_metadata

        attn_metadata = fetch_attn_metadata()

        if not attn_metadata.token_merge_enabled:
            return hidden_states

        merge_mask = attn_metadata.token_merge_mask
        topk_ids = attn_metadata.token_merge_topk_ids
        topk_probs = attn_metadata.token_merge_topk_probs
        residual_probs = attn_metadata.token_merge_residual_probs
        mask_token_id = attn_metadata.token_merge_mask_token_id
        merge_mode = attn_metadata.token_merge_mode
        if merge_mask is None or topk_ids is None or topk_probs is None or residual_probs is None or mask_token_id is None:
            raise RuntimeError("token_merge_enabled=True but token-merge metadata is incomplete.")
        if int(merge_mask.numel()) != int(hidden_states.shape[0]):
            raise RuntimeError(
                "Token-merge metadata length does not match hidden states: "
                f"merge_mask={merge_mask.numel()}, hidden_states={hidden_states.shape[0]}"
            )

        device = hidden_states.device
        dtype = hidden_states.dtype
        merge_mask = merge_mask.to(device=device, dtype=torch.bool)
        is_compiling = bool(getattr(torch.compiler, "is_compiling", lambda: False)())
        if not is_compiling and not torch.cuda.is_current_stream_capturing() and not bool(merge_mask.any().item()):
            return hidden_states

        topk_ids = topk_ids.to(device=device, dtype=torch.int64)
        topk_probs = topk_probs.to(device=device, dtype=torch.float32)
        residual_probs = residual_probs.to(device=device, dtype=torch.float32)

        flat_topk_ids = rearrange(topk_ids, "token topk -> (token topk)")
        topk_embeds = rearrange(
            self.word_embeddings(flat_topk_ids),
            "(token topk) hidden -> token topk hidden",
            token=topk_ids.shape[0],
            topk=topk_ids.shape[1],
        )
        merge_dtype = hidden_states.dtype
        topk_embeds_merge = topk_embeds.to(dtype=merge_dtype)
        topk_probs_merge = topk_probs.to(dtype=merge_dtype)
        residual_probs_merge = residual_probs.to(dtype=merge_dtype)
        topk_weighted = reduce(
            topk_embeds_merge * rearrange(topk_probs_merge, "token topk -> token topk 1"),
            "token topk hidden -> token hidden",
            "sum",
        )

        if merge_mode == "dmax_topk":
            if self.mask_token_id_tensor is not None and self.mask_token_id == int(mask_token_id):
                mask_embed = self.word_embeddings(self.mask_token_id_tensor).to(dtype=merge_dtype)
            else:
                if torch.cuda.is_current_stream_capturing():
                    raise RuntimeError(
                        "CUDA graph capture requires LLaDA2Model.mask_token_id_tensor to match "
                        f"attn_metadata mask token id (model={self.mask_token_id}, metadata={mask_token_id})."
                    )
                mask_id = torch.tensor([int(mask_token_id)], dtype=torch.int64, device=device)
                mask_embed = self.word_embeddings(mask_id).to(dtype=merge_dtype)
            # Match native generate_spd's embedding-dtype blend and only lift norm
            # calculations to float temporarily.
            soft_embeds = topk_weighted + mask_embed * residual_probs_merge
            
            if attn_metadata.token_merge_renormalize:
                current_norm = torch.linalg.vector_norm(
                    soft_embeds.float(), dim=-1, keepdim=True
                ).clamp_min(1e-12).to(dtype=merge_dtype)
                topk_norms = torch.linalg.vector_norm(
                    topk_embeds.float(), dim=-1
                ).to(dtype=merge_dtype)
                expected_topk_norm = (topk_norms * topk_probs_merge).sum(dim=-1, keepdim=True)
                expected_mask_norm = torch.linalg.vector_norm(
                    mask_embed.float(), dim=-1, keepdim=True
                ).to(dtype=merge_dtype) * residual_probs_merge
                target_norm = expected_topk_norm + expected_mask_norm
                soft_embeds = soft_embeds * (target_norm / current_norm)
        elif merge_mode == "iter_smooth_topk":
            merge_weight = float(attn_metadata.token_merge_weight)
            soft_embeds = hidden_states.to(torch.float32) + merge_weight * topk_weighted
        else:
            raise ValueError(f"Unsupported token_merge_mode: {merge_mode}")

        merge_mask_expanded = rearrange(merge_mask, "token -> token 1")
        return torch.where(merge_mask_expanded, soft_embeds.to(dtype=dtype), hidden_states)


def build_llada2_runtime_config(config):
    runtime_config = copy.copy(getattr(config, "hf_config", config))
    for name in (
        "moe_dispatcher_backend",
        "moe_gemm_impl",
        "deepep_mode",
        "deepep_num_max_dispatch_tokens_per_rank",
        "expert_parallel_size",
        "tensor_parallel_size",
        "data_parallel_size",
        "mask_token_id",
        "attn_impl",
    ):
        if hasattr(config, name):
            setattr(runtime_config, name, getattr(config, name))
    return runtime_config


@AutoModelForDiffusionLM.register("llada2", use_full_config=True)
@AutoModelForDiffusionLM.register("llada2_moe", use_full_config=True)
@AutoModelForDiffusionLM.register("llada2_mini", use_full_config=True)
@AutoModelForDiffusionLM.register("llada2dot1_mini", use_full_config=True)
class LLaDA2ForDiffusionLM(nn.Module):
    packed_modules_mapping = {}

    def __init__(self, config) -> None:
        super().__init__()
        runtime_config = build_llada2_runtime_config(config)
        self.model = LLaDA2Model(runtime_config)
        self.lm_head = ParallelLMHead(runtime_config.vocab_size, runtime_config.hidden_size)
        if getattr(runtime_config, "tie_word_embeddings", False):
            self.lm_head.weight.data = self.model.word_embeddings.weight.data

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        return self.model(input_ids, positions, mask)

    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.lm_head(hidden_states)


SparseMoEBlock = FusedMoE


__all__ = [
    "LLaDA2Attention",
    "LLaDA2DecoderLayer",
    "LLaDA2DenseMLP",
    "LLaDA2ForDiffusionLM",
    "LLaDA2Model",
    "LLaDA2NaiveMoE",
    "LLaDA2TPMoE",
    "LLaDA2EPMoE",
]