GLM 5.1 ARCHITECTURE

class GlmMoeDsaRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return self.weight * hidden_states.to(input_dtype)

def apply_rotary_pos_emb(x, cos, sin, unsqueeze_dim: int = 1) -> torch.Tensor:
    """
    This is the transformers equivalent of DeepSeek V3.2's `apply_rotary_emb(x, freqs_cis, interleaved)`.
    Instead of using complex-number `freqs_cis`, we use pre-split `(cos, sin)` tensors from RotaryEmbedding.
    """
    cos = cos.unsqueeze(unsqueeze_dim)
    sin = sin.unsqueeze(unsqueeze_dim)
    # Split-half (NeoX/Llama style): (x[:d/2], x[d/2:])
    return (x * cos) + (rotate_half(x) * sin)

class GlmMoeDsaRotaryEmbedding(nn.Module):
    @staticmethod
    def compute_default_rope_parameters(config, device=None, seq_len=None):
        base = config.rope_parameters["rope_theta"]
        partial_rotary_factor = config.rope_parameters.get("partial_rotary_factor", 1.0)
        head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
        dim = int(head_dim * partial_rotary_factor)
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(...) / dim))
        return inv_freq, 1.0

def eager_attention_forward(module, query, key, value, attention_mask, scaling, dropout=0.0, **kwargs):
    key_states   = repeat_kv(key, module.num_key_value_groups)    # no-op (groups=1)
    value_states = repeat_kv(value, module.num_key_value_groups)

    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
    if attention_mask is not None:
        attn_weights = attn_weights + attention_mask              # ← combined DSA+causal here
    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
    attn_output  = torch.matmul(attn_weights, value_states)
    return attn_output.transpose(1, 2).contiguous(), attn_weights

class GlmMoeDsaMLP(nn.Module):
    def __init__(self, config, intermediate_size=None):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size
        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.up_proj   = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
        self.act_fn = ACT2FN[config.hidden_act]   # silu

    def forward(self, x):
        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))

class GlmMoeDsaTopkRouter(nn.Module):
    def __init__(self, config: GlmMoeDsaConfig):
        super().__init__()
        self.config = config
        self.top_k = config.num_experts_per_tok                     # 8
        self.n_routed_experts = config.n_routed_experts             # 256
        self.routed_scaling_factor = config.routed_scaling_factor   # 2.5
        self.n_group   = config.n_group                              # 1
        self.topk_group = config.topk_group                          # 1
        self.norm_topk_prob = config.norm_topk_prob                  # True

        self.weight = nn.Parameter(torch.empty((self.n_routed_experts, config.hidden_size)))
        self.register_buffer("e_score_correction_bias",
                             torch.zeros((self.n_routed_experts), dtype=torch.float32))

    def forward(self, hidden_states):
        hidden_states = hidden_states.view(-1, self.config.hidden_size)
        router_logits = F.linear(hidden_states.type(torch.float32),
                                 self.weight.type(torch.float32))
        return router_logits

@use_experts_implementation
class GlmMoeDsaNaiveMoe(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.num_experts     = config.num_local_experts                # 256
        self.hidden_dim      = config.hidden_size                       # 6144
        self.intermediate_dim = config.moe_intermediate_size            # 2048
        self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts,
                                  2 * self.intermediate_dim, self.hidden_dim))
        self.down_proj    = nn.Parameter(torch.empty(self.num_experts,
                                  self.hidden_dim, self.intermediate_dim))
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, hidden_states, top_k_index, top_k_weights):
        final_hidden_states = torch.zeros_like(hidden_states)
        with torch.no_grad():
            expert_mask = F.one_hot(top_k_index, num_classes=self.num_experts).permute(2,1,0)
            expert_hit  = torch.greater(expert_mask.sum(dim=(-1,-2)), 0).nonzero()

        for expert_idx in expert_hit:
            expert_idx = expert_idx[0]
            top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
            current_state = hidden_states[token_idx]
            gate, up = F.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1)
            current_hidden_states = self.act_fn(gate) * up
            current_hidden_states = F.linear(current_hidden_states, self.down_proj[expert_idx])
            current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
            final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(...))

        return final_hidden_states

class GlmMoeDsaMoE(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.experts = GlmMoeDsaNaiveMoe(config)
        self.gate    = GlmMoeDsaTopkRouter(config)
        self.shared_experts = GlmMoeDsaMLP(
            config=config,
            intermediate_size=config.moe_intermediate_size * config.n_shared_experts,  # 2048*1
        )

    def forward(self, hidden_states):
        residuals = hidden_states
        orig_shape = hidden_states.shape
        router_logits = self.gate(hidden_states)
        topk_indices, topk_weights = self.route_tokens_to_experts(router_logits)
        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
        hidden_states = self.experts(hidden_states, topk_indices, topk_weights).view(*orig_shape)
        hidden_states = hidden_states + self.shared_experts(residuals)
        return hidden_states

class GlmMoeDsaDecoderLayer(GradientCheckpointingLayer):
    def __init__(self, config, layer_idx):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.self_attn = GlmMoeDsaAttention(config, layer_idx)

        if config.mlp_layer_types[layer_idx] == "sparse":
            self.mlp = GlmMoeDsaMoE(config)
        else:
            self.mlp = GlmMoeDsaMLP(config)

        self.input_layernorm           = GlmMoeDsaRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.post_attention_layernorm  = GlmMoeDsaRMSNorm(config.hidden_size, config.rms_norm_eps)

    def forward(self, hidden_states, attention_mask=None, position_ids=None,
                past_key_values=None, use_cache=False, position_embeddings=None, **kwargs):
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        hidden_states, _ = self.self_attn(
            hidden_states=hidden_states, attention_mask=attention_mask,
            position_ids=position_ids, past_key_values=past_key_values,
            use_cache=use_cache, position_embeddings=position_embeddings, **kwargs)
        hidden_states = residual + hidden_states

        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states
        return hidden_states

class GlmMoeDsaModel(GlmMoeDsaPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
        self.layers = nn.ModuleList(
            [GlmMoeDsaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
        )
        self.norm = GlmMoeDsaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.rotary_emb = GlmMoeDsaRotaryEmbedding(config=config)
        self.gradient_checkpointing = False
        self.post_init()

    def forward(self, input_ids=None, attention_mask=None, position_ids=None,
                past_key_values=None, inputs_embeds=None, use_cache=None, **kwargs):
        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)
        if use_cache and past_key_values is None:
            past_key_values = DynamicCache(config=self.config)
        ...
        causal_mask = create_causal_mask(...)
        hidden_states = inputs_embeds
        position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)

        for decoder_layer in self.layers[:self.config.num_hidden_layers]:
            hidden_states = decoder_layer(
                hidden_states,
                attention_mask=causal_mask,
                position_embeddings=position_embeddings,
                position_ids=position_ids,
                past_key_values=past_key_values,
                use_cache=use_cache, **kwargs)

        hidden_states = self.norm(hidden_states)
        return BaseModelOutputWithPast(last_hidden_state=hidden_states,
                                       past_key_values=past_key_values)

class GlmMoeDsaForCausalLM(GlmMoeDsaPreTrainedModel, GenerationMixin):
    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}  # placeholder, tie_word_embeddings=False

    def __init__(self, config):
        super().__init__(config)
        self.model = GlmMoeDsaModel(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.post_init()

class GlmMoeDsaPreTrainedModel(PreTrainedModel):
    config: GlmMoeDsaConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["GlmMoeDsaDecoderLayer"]

    _supports_flash_attn = False  # flash-mla kernels need a bit more work...
    _supports_sdpa = True
    _supports_flex_attn = False
    _can_compile_fullgraph = True
    _supports_attention_backend = True

    # FP8 quantization uses _keep_in_fp32_modules to decide what NOT to convert
    _keep_in_fp32_modules         = ["indexer.weights_proj"]
    _keep_in_fp32_modules_strict  = ["e_score_correction_bias"]
    _keys_to_ignore_on_load_unexpected = [r"model\.layers\.78.*"]
    _compatible_flash_implementations  = ["kernels-community/flash-mla"]

attribute_map = {
    "num_local_experts": "n_routed_experts",   # for MoE generic interfaces
    "head_dim":          "qk_rope_head_dim",   # for RotaryEmbedding -- only the rope subspace
}

base_model_tp_plan = {
    "layers.*.self_attn.q_b_proj":             "colwise",
    "layers.*.self_attn.kv_a_proj_with_mqa":   "mla_kv_a_proj",
    "layers.*.self_attn.kv_b_proj":            "colwise",
    "layers.*.self_attn.o_proj":               "rowwise",
    "layers.*.mlp.experts.gate_up_proj":       "packed_colwise",
    "layers.*.mlp.experts.down_proj":          "rowwise",
    "layers.*.mlp.experts":                    "moe_tp_experts",
    "layers.*.mlp.shared_experts.gate_proj":   "colwise",
    "layers.*.mlp.shared_experts.up_proj":     "colwise",
    "layers.*.mlp.shared_experts.down_proj":   "rowwise",
    "layers.*.mlp.gate_proj":                  "colwise",   # dense layers
    "layers.*.mlp.up_proj":                    "colwise",
    "layers.*.mlp.down_proj":                  "rowwise",
}

def __post_init__(self, **kwargs):
    self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim    # 192 + 64 = 256

    # MLP layer types: first 3 dense, rest sparse
    if self.mlp_layer_types is None:
        self.mlp_layer_types = (
            ["dense"] * min(3, self.num_hidden_layers)
            + ["sparse"] * (self.num_hidden_layers - 3)
        )
    super().__post_init__(**kwargs)

attn_output, attn_weights = attention_interface(
    self,
    query_states,
    key_states,
    value_states,
    combined_mask,
    dropout=0.0 if not self.training else self.attention_dropout,
    scaling=self.scaling,
    indices=topk_indices,        # flash_mla_with_kvcache reads this
    **kwargs,
)

# Reset cache on prefill (new prompt) to avoid stale keys / batch-size mismatch
if seq_len > 1:
    self._cached_keys = None

if use_cache:
    if self._cached_keys is not None:
        k_cached = torch.cat([self._cached_keys, k], dim=1)  # [B, T, D]
    else:
        k_cached = k
    self._cached_keys = k_cached
else:
    k_cached = k

_keys_to_ignore_on_load_unexpected = [r"model\.layers\.78.*"]

# In GlmMoeDsaAttention.forward:
q_resid = self.q_a_layernorm(self.q_a_proj(hidden_states))   # [B, S, q_lora_rank]
query_states = self.q_b_proj(q_resid)
...
# Same q_resid is forwarded to the indexer:
topk_indices = self.indexer(
    hidden_states,
    q_resid,                                                  # ← reused!
    position_embeddings,
    indexer_mask,
    use_cache=past_key_values is not None,
)