lfm2_moe.py

# North Star OS addition (2026-07-14), style-matched to the Apple BSD-3 tree it lives in.
#
# LFM2.5-MoE (LiquidAI LFM2.5-8B-A1B): the LFM2 hybrid (short-conv + GQA) with a
# 32-expert top-4 sparse SwiGLU MLP on all layers except the first
# `num_dense_layers`. The expert compute rides the SwitchGLU/SwitchLinear (GatherMM)
# primitive, so ONLY the routed experts' weight slabs are read per token. Router math
# matches HF Lfm2MoeTopKRouter EXACTLY (ground truth, since we load HF weights):
# SIGMOID(fp32) -> +expert_bias (SELECTION only) -> top-k -> gather sigmoid weights
# (bias excluded) -> L1 renorm -> routed_scaling_factor.

import torch
import torch.nn as nn
from typing_extensions import Self, override

from coreai_models.models.base import BaseForCausalLM
from coreai_models.models.macos.lfm2 import (
    ConvState,
    Lfm2DecoderLayer,
    Lfm2MLP,
    Lfm2Model,
    _fuse_lfm2_attention_weights,
    build_lfm2_reference_inputs,
)
from coreai_models.primitives.macos.cache import KVCache
from coreai_models.primitives.macos.switch import SwitchGLU


class Lfm2MoeSparseBlock(nn.Module):
    def __init__(self, config) -> None:
        super().__init__()
        dim = config.hidden_size
        self.top_k = config.num_experts_per_tok
        self.norm_topk_prob = getattr(config, "norm_topk_prob", True)
        self.use_expert_bias = getattr(config, "use_expert_bias", False)
        self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)

        self.gate = nn.Linear(dim, config.num_experts, bias=False)
        if self.use_expert_bias:
            self.expert_bias = nn.Parameter(torch.zeros(config.num_experts))
        self.switch_mlp = SwitchGLU(dim, config.moe_intermediate_size, config.num_experts)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Matches HF Lfm2MoeTopKRouter EXACTLY (ground truth, since we load HF weights):
        # SIGMOID gate (not softmax); expert_bias added ONLY to pick the top-k experts,
        # while the weights are the *sigmoid* scores gathered at those experts (bias
        # excluded from weights); then L1-renorm and routed_scaling_factor.
        routing_weights = self.gate(x).to(torch.float32).sigmoid()  # (B, S, E)
        scores_for_routing = (
            routing_weights + self.expert_bias.to(torch.float32)
            if self.use_expert_bias
            else routing_weights
        )
        _, indices = torch.topk(scores_for_routing, self.top_k, dim=-1)  # selection
        scores = torch.gather(routing_weights, dim=-1, index=indices)    # weights = sigmoid, no bias
        if self.norm_topk_prob:
            scores = scores / (torch.sum(scores, dim=-1, keepdim=True) + 1e-6)
        scores = (scores * self.routed_scaling_factor).to(x.dtype)

        y = self.switch_mlp(x, indices.to(torch.uint16))  # (B, S, top_k, D)
        y = y * scores.unsqueeze(-1)
        return torch.sum(y, dim=-2).to(x.dtype)


class Lfm2MoeDecoderLayer(Lfm2DecoderLayer):
    def _build_feed_forward(self, config, layer_idx: int) -> nn.Module:
        if layer_idx < getattr(config, "num_dense_layers", 0):
            # lfm2_moe stores the dense ff dim post-adjust — use it as-is.
            return Lfm2MLP(config, auto_adjust=False)
        return Lfm2MoeSparseBlock(config)


class Lfm2MoeModel(Lfm2Model):
    layer_cls = Lfm2MoeDecoderLayer


class Lfm2MoeForCausalLM(BaseForCausalLM):
    _HF_MODEL_CLASS = None  # loaded straight from safetensors; no HF class needed

    @override
    def _init_model(self, config) -> None:
        self.model = Lfm2MoeModel(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        if getattr(config, "tie_word_embeddings", False):
            self.lm_head.weight = self.model.embed_tokens.weight

    @BaseForCausalLM.cast_logits_bfloat16_to_float16
    def forward(
        self,
        input_ids: torch.Tensor,
        position_ids: torch.IntTensor,
        k_cache: torch.Tensor,
        v_cache: torch.Tensor,
        conv_state: torch.Tensor,
    ) -> torch.Tensor:
        cache = KVCache(k_cache, v_cache)
        conv = ConvState(conv_state)
        out = self.model(input_ids, position_ids, cache, conv)
        return self.lm_head(out)

    @override
    def _mutate_state_dict(self: Self, state_dict: dict[str, torch.Tensor]) -> None:
        _fuse_lfm2_attention_weights(self.model, state_dict)

        # Stack per-expert w1/w3/w2 into the SwitchGLU gate/up/down layout
        # (1, num_experts, out_dims, in_dims). Checkpoint keys:
        #   model.layers.{i}.feed_forward.experts.{e}.{w1,w2,w3}.weight
        proj_map = {"w1": "gate_proj", "w3": "up_proj", "w2": "down_proj"}
        for i in range(len(self.model.layers)):
            prefix = f"model.layers.{i}.feed_forward"
            if f"{prefix}.experts.0.w1.weight" not in state_dict:
                continue
            num_experts = 0
            while f"{prefix}.experts.{num_experts}.w1.weight" in state_dict:
                num_experts += 1
            for src, dst in proj_map.items():
                first = state_dict[f"{prefix}.experts.0.{src}.weight"]
                output = torch.empty(
                    (1, num_experts) + first.shape, dtype=first.dtype, device=first.device
                )
                for e in range(num_experts):
                    output[0, e] = state_dict.pop(f"{prefix}.experts.{e}.{src}.weight")
                state_dict[f"{prefix}.switch_mlp.{dst}.weight"] = output

    def load_state_dict(self, state_dict, strict: bool = True, assign: bool = False):
        result = super().load_state_dict(state_dict, strict=strict, assign=assign)
        if getattr(self.config, "tie_word_embeddings", False):
            self.lm_head.weight = self.model.embed_tokens.weight
        return result

    # ---- export hooks (consumed by export_macos_model when present) ----
    @staticmethod
    def state_names() -> tuple[str, ...]:
        return ("k_cache", "v_cache", "conv_state")

    @classmethod
    def build_reference_inputs(cls, config, target_dtype, max_context_length, trace_query_len, trace_offset, trace_kv_seq_len):
        return build_lfm2_reference_inputs(
            config, target_dtype, max_context_length, trace_query_len, trace_offset, trace_kv_seq_len
        )