modeling_codva1.py

# modeling_codva1.py
"""
HuggingFace-compatible implementation of CodVa1 (CodVa-Dense-CodeMoE).
~1B total / ~500M active params, code-specialized sparse MoE, FIM-capable.

Supports:
  - AutoModelForCausalLM / trust_remote_code
  - model.generate() with KV-cache
  - MoE routed FFN (top-k + shared experts)
  - GQA attention with QK-norm and RoPE
  - Optional structural attention bias (full-sequence mode only, unused if
    rel_ids is never passed — matches original training setup)
"""

import math
from typing import Optional, Tuple, List

import torch
import torch.nn as nn
import torch.nn.functional as F

from transformers import PreTrainedModel
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast

from .configuration_codva1 import Codva1Config


# ==============================================================================
# Building blocks
# ==============================================================================

class Codva1RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.scale = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        return F.rms_norm(x, self.scale.shape, self.scale, self.eps)


class Codva1SwiGLU(nn.Module):
    def __init__(self, d_model, hidden):
        super().__init__()
        self.gate = nn.Linear(d_model, hidden, bias=False)
        self.up = nn.Linear(d_model, hidden, bias=False)
        self.down = nn.Linear(hidden, d_model, bias=False)

    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))


def precompute_freqs_cis(head_dim, max_len, theta=10_000.0, device=None):
    freqs = 1.0 / (theta ** (
        torch.arange(0, head_dim, 2, dtype=torch.float32, device=device) / head_dim))
    t = torch.arange(max_len, dtype=torch.float32, device=device)
    freqs = torch.outer(t, freqs)
    return torch.cos(freqs), torch.sin(freqs)


def apply_rope(xq, xk, cos, sin, start_pos: int = 0):
    L = xq.shape[1]
    cos_s = cos[start_pos:start_pos + L]
    sin_s = sin[start_pos:start_pos + L]
    c = torch.cat([cos_s, cos_s], -1).unsqueeze(0).unsqueeze(2)
    s = torch.cat([sin_s, sin_s], -1).unsqueeze(0).unsqueeze(2)

    def rot(x):
        x1, x2 = x[..., :x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
        return torch.cat([-x2, x1], dim=-1)

    xq_f, xk_f = xq.float(), xk.float()
    return ((xq_f * c + rot(xq_f) * s).to(xq.dtype),
            (xk_f * c + rot(xk_f) * s).to(xk.dtype))


class Codva1Attention(nn.Module):
    def __init__(self, config: Codva1Config):
        super().__init__()
        self.nh = config.n_heads
        self.nkv = config.n_kv_heads
        self.hd = config.d_model // config.n_heads
        self.ng = config.n_heads // config.n_kv_heads
        self.use_qk_norm = config.use_qk_norm

        D, Dkv = config.n_heads * self.hd, config.n_kv_heads * self.hd
        self.q = nn.Linear(config.d_model, D, bias=False)
        self.k = nn.Linear(config.d_model, Dkv, bias=False)
        self.v = nn.Linear(config.d_model, Dkv, bias=False)
        self.o = nn.Linear(D, config.d_model, bias=False)

        if self.use_qk_norm:
            self.q_norm = Codva1RMSNorm(self.hd)
            self.k_norm = Codva1RMSNorm(self.hd)

    def forward(self, x, cos, sin, start_pos=0, past_key_value=None,
                use_cache=False, attn_bias=None):
        B, L, _ = x.shape
        q = self.q(x).view(B, L, self.nh, self.hd)
        k = self.k(x).view(B, L, self.nkv, self.hd)
        v = self.v(x).view(B, L, self.nkv, self.hd)

        if self.use_qk_norm:
            q = self.q_norm(q)
            k = self.k_norm(k)

        q, k = apply_rope(q, k, cos, sin, start_pos=start_pos)

        q = q.transpose(1, 2)   # [B, nh,  L, hd]
        k = k.transpose(1, 2)   # [B, nkv, L, hd]
        v = v.transpose(1, 2)

        if past_key_value is not None:
            past_k, past_v = past_key_value
            k = torch.cat([past_k, k], dim=2)
            v = torch.cat([past_v, v], dim=2)

        present = (k, v) if use_cache else None

        k_rep = k.repeat_interleave(self.ng, dim=1)
        v_rep = v.repeat_interleave(self.ng, dim=1)

        Tq, Tk = q.shape[2], k_rep.shape[2]

        if attn_bias is not None and Tq == Tk:
            causal = torch.triu(
                torch.ones(Tq, Tk, device=x.device, dtype=torch.bool), diagonal=1)
            bias = attn_bias.masked_fill(causal.view(1, 1, Tq, Tk), float('-inf'))
            out = F.scaled_dot_product_attention(
                q, k_rep, v_rep, attn_mask=bias.float(), is_causal=False)
        elif Tq == Tk:
            out = F.scaled_dot_product_attention(q, k_rep, v_rep, is_causal=True)
        else:
            # incremental decode step(s) with KV cache
            q_pos = torch.arange(start_pos, start_pos + Tq, device=x.device)
            k_pos = torch.arange(0, Tk, device=x.device)
            mask_bool = k_pos[None, :] > q_pos[:, None]
            attn_mask = torch.zeros(Tq, Tk, device=x.device, dtype=torch.float32)
            attn_mask.masked_fill_(mask_bool, float('-inf'))
            out = F.scaled_dot_product_attention(
                q, k_rep, v_rep, attn_mask=attn_mask, is_causal=False)

        out = out.transpose(1, 2).contiguous().view(B, L, -1)
        return self.o(out), present


class Codva1StructuralBias(nn.Module):
    def __init__(self, n_heads, n_rel):
        super().__init__()
        self.n_rel = n_rel + 1
        self.rel_emb = nn.Parameter(torch.zeros(n_heads, self.n_rel))

    def forward(self, rel_ids):
        bias = self.rel_emb.t()[rel_ids]           # [B, L, L, nh]
        return bias.permute(0, 3, 1, 2).contiguous()


class Codva1MoEFFN(nn.Module):
    """Token-level sparse MoE FFN: shared experts + top-k routed experts."""

    def __init__(self, config: Codva1Config):
        super().__init__()
        self.d_model = config.d_model
        self.top_k = config.moe_top_k
        self.n_exp = config.moe_experts

        self.shared = nn.ModuleList(
            [Codva1SwiGLU(config.d_model, config.moe_hidden) for _ in range(config.moe_shared)])
        self.experts = nn.ModuleList(
            [Codva1SwiGLU(config.d_model, config.moe_hidden) for _ in range(config.moe_experts)])
        self.gate = nn.Linear(config.d_model, config.moe_experts, bias=False)

        self.register_buffer("balance_bias", torch.zeros(config.moe_experts))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, L, D = x.shape
        xf = x.reshape(-1, D)
        N = xf.shape[0]

        out = torch.zeros_like(xf)
        for s in self.shared:
            out = out + s(xf)

        logits = self.gate(xf) + self.balance_bias
        _, idx = torch.topk(logits, self.top_k, dim=-1)
        w = logits.softmax(-1).gather(1, idx)
        w = w / (w.sum(-1, keepdim=True) + 1e-9)

        flat_idx = idx.reshape(-1)
        flat_w = w.reshape(-1)
        token_ids = (torch.arange(N, device=x.device)
                     .unsqueeze(1).expand(-1, self.top_k).reshape(-1))

        sort_order = flat_idx.argsort(stable=True)
        sorted_exp = flat_idx[sort_order]
        sorted_tok = token_ids[sort_order]
        sorted_w = flat_w[sort_order]
        counts = torch.bincount(sorted_exp, minlength=self.n_exp)
        dispatched = xf[sorted_tok]

        routed_out = torch.zeros_like(dispatched)
        start = 0
        for e in range(self.n_exp):
            end = start + int(counts[e])
            if end > start:
                routed_out[start:end] = self.experts[e](dispatched[start:end])
            start = end

        routed_out = routed_out * sorted_w.unsqueeze(-1)
        unsort = sort_order.argsort()
        routed_out = routed_out[unsort].view(N, self.top_k, D)
        out = out + routed_out.sum(dim=1)

        return out.view(B, L, D)


class Codva1Block(nn.Module):
    def __init__(self, config: Codva1Config, layer_idx: int):
        super().__init__()
        self.attn_norm = Codva1RMSNorm(config.d_model)
        self.ffn_norm = Codva1RMSNorm(config.d_model)
        self.attn = Codva1Attention(config)
        self.is_moe = config.use_moe and (layer_idx % config.moe_every == 0)
        self.ffn = (Codva1MoEFFN(config) if self.is_moe
                    else Codva1SwiGLU(config.d_model, config.ffn_hidden))

    def forward(self, x, cos, sin, start_pos=0, past_key_value=None,
                use_cache=False, attn_bias=None):
        attn_out, present = self.attn(
            self.attn_norm(x), cos, sin, start_pos=start_pos,
            past_key_value=past_key_value, use_cache=use_cache, attn_bias=attn_bias)
        x = x + attn_out
        x = x + self.ffn(self.ffn_norm(x))
        return x, present


# ==============================================================================
# HF PreTrainedModel wrapper
# ==============================================================================

class Codva1PreTrainedModel(PreTrainedModel):
    config_class = Codva1Config
    base_model_prefix = "model"
    supports_gradient_checkpointing = False
    _no_split_modules = ["Codva1Block"]
    _supports_sdpa = True

    def _init_weights(self, module):
        std = self.config.initializer_range
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()


class Codva1ForCausalLM(Codva1PreTrainedModel, GenerationMixin):
    """
    Flat structure — attribute names (embed, layers, norm, lm_head, struct_bias)
    match the original training checkpoint 1:1, so weights load with zero
    key remapping needed.
    """

    def __init__(self, config: Codva1Config):
        super().__init__(config)
        self.embed = nn.Embedding(config.vocab_size, config.d_model)

        cos, sin = precompute_freqs_cis(
            config.head_dim, config.max_len, config.rope_theta)
        self.register_buffer("rope_cos", cos, persistent=False)
        self.register_buffer("rope_sin", sin, persistent=False)

        self.layers = nn.ModuleList(
            [Codva1Block(config, i) for i in range(config.n_layers)])
        self.norm = Codva1RMSNorm(config.d_model)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)

        # weight tying
        self.lm_head.weight = self.embed.weight

        self.struct_bias = (
            Codva1StructuralBias(config.n_heads, config.n_struct_rel)
            if config.use_structural_bias else None
        )

        # init weights (manual, skips HF tied weights detection)
        self.apply(self._init_weights)


    @property
    def all_tied_weights_keys(self):
        return {}

    # ---- embedding accessors (required by HF) ----
    def get_input_embeddings(self):
        return self.embed

    def set_input_embeddings(self, value):
        self.embed = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    # ---- forward ----
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        rel_ids: Optional[torch.LongTensor] = None,
        **kwargs,
    ):
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        use_cache = use_cache if use_cache is not None else self.config.use_cache

        if inputs_embeds is None:
            x = self.embed(input_ids)
        else:
            x = inputs_embeds

        past_length = 0
        if past_key_values is not None and past_key_values[0] is not None:
            past_length = past_key_values[0][0].shape[2]

        attn_bias = None
        if self.struct_bias is not None and rel_ids is not None:
            attn_bias = self.struct_bias(rel_ids)

        presents = () if use_cache else None
        all_hidden_states = () if output_hidden_states else None

        for i, block in enumerate(self.layers):
            if output_hidden_states:
                all_hidden_states += (x,)
            layer_past = past_key_values[i] if past_key_values is not None else None
            x, present = block(
                x, self.rope_cos, self.rope_sin,
                start_pos=past_length, past_key_value=layer_past,
                use_cache=use_cache, attn_bias=attn_bias)
            if use_cache:
                presents += (present,)

        x = self.norm(x)
        if output_hidden_states:
            all_hidden_states += (x,)

        logits = self.lm_head(x)

        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
                ignore_index=-100,
            )

        if not return_dict:
            output = (logits, presents, all_hidden_states)
            return ((loss,) + output) if loss is not None else output

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=presents,
            hidden_states=all_hidden_states,
            attentions=None,
        )

    # ---- generation helpers ----
    def prepare_inputs_for_generation(
        self, input_ids, past_key_values=None, attention_mask=None,
        inputs_embeds=None, **kwargs
    ):
        if past_key_values is not None:
            input_ids = input_ids[:, -1:]
        model_inputs = {
            "input_ids": input_ids,
            "past_key_values": past_key_values,
            "use_cache": kwargs.get("use_cache", True),
        }
        return model_inputs

    @staticmethod
    def _reorder_cache(past_key_values, beam_idx):
        reordered = ()
        for layer_past in past_key_values:
            reordered += (tuple(t.index_select(0, beam_idx) for t in layer_past),)
        return reordered

    # ---- param accounting (handy for debugging) ----
    def count_params(self):
        seen, total = set(), 0
        for p in self.parameters():
            if id(p) not in seen:
                seen.add(id(p))
                total += p.numel()
        return total

    def active_params(self):
        seen, total = set(), 0
        for name, p in self.named_parameters():
            if id(p) in seen:
                continue
            seen.add(id(p))
            if ".experts." in name:
                continue
            total += p.numel()
        for b in self.layers:
            if b.is_moe:
                per_exp = sum(p.numel() for p in b.ffn.experts[0].parameters())
                total += per_exp * (b.ffn.top_k + len(b.ffn.shared))
        return total


# ---- register for trust_remote_code ----
Codva1Config.register_for_auto_class()
Codva1ForCausalLM.register_for_auto_class("AutoModelForCausalLM")