"""HuggingFace CausalLM wrapper for Full YatNMN GPT (YatNMN attn + MLP), with KV cache. KV cache for YatNMN attention caches: - xh_past: RoPE-applied heads (B, H, T_past, D) for pairwise scoring - v_past: value heads (B, H, T_past, D) for aggregation """ from __future__ import annotations import math from typing import Optional, Tuple import torch import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedModel from transformers.modeling_outputs import CausalLMOutputWithPast from transformers.generation import GenerationMixin try: from .configuration_yatnmn_full_gpt import YatFullGPTHfConfig from .yatnmn_full_gpt import YatFull_GPT, YatFullGPTConfig from .torch_gpt import rms_norm, apply_rotary_emb except ImportError: from torch_port.yat_full.configuration_yatnmn_full_gpt import YatFullGPTHfConfig from torch_port.yat_full.yatnmn_full_gpt import YatFull_GPT, YatFullGPTConfig from torch_port.torch_gpt import rms_norm, apply_rotary_emb def _kvcache_yat_attn( attn_module: nn.Module, x_norm: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, window_size: Tuple[int, int], past_xh: Optional[torch.Tensor], past_v: Optional[torch.Tensor], ): """YatNMN attention with KV cache.""" cfg = attn_module.config B, T_new, C = x_norm.shape n_head, n_kv_head = cfg.n_head, cfg.n_kv_head head_dim = C // n_head # New x_heads with RoPE x_heads_new = apply_rotary_emb(x_norm.reshape(B, T_new, n_head, head_dim), cos, sin) xh_new = x_heads_new.float().transpose(1, 2) # (B, H, T_new, D) # New values v_new = attn_module.c_v(x_norm).reshape(B, T_new, n_kv_head, head_dim) if n_kv_head < n_head: v_new = v_new.repeat_interleave(n_head // n_kv_head, dim=2) v_new_t = v_new.transpose(1, 2) # (B, H, T_new, D) # Concatenate with past if past_xh is not None: xh_all = torch.cat([past_xh, xh_new], dim=2) # (B, H, T_total, D) v_all = torch.cat([past_v, v_new_t], dim=2) else: xh_all = xh_new v_all = v_new_t T_total = xh_all.shape[2] new_xh = xh_all new_v = v_all # Compute pairwise scores: query = xh_new, key = xh_all # dots[b,h,i,j] = xh_new[b,h,i] . xh_all[b,h,j] dots = torch.matmul(xh_new, xh_all.transpose(-2, -1)) # (B, H, T_new, T_total) # distances q_sq = (xh_new ** 2).sum(dim=-1) # (B, H, T_new) k_sq = (xh_all ** 2).sum(dim=-1) # (B, H, T_total) dist_sq = torch.clamp(q_sq.unsqueeze(-1) + k_sq.unsqueeze(-2) - 2.0 * dots, min=0.0) b = F.softplus(attn_module.attn_bias_raw) eps = F.softplus(attn_module.attn_eps_raw) scores = (dots + b[None, :, None, None]) ** 2 / (dist_sq + eps[None, :, None, None]) # Strict causal: j < i (absolute positions) device = x_norm.device q_pos = torch.arange(T_total - T_new, T_total, device=device) # absolute positions of queries k_pos = torch.arange(T_total, device=device) # absolute positions of all keys causal = k_pos.unsqueeze(0) < q_pos.unsqueeze(1) # (T_new, T_total) window_left = window_size[0] if 0 < window_left < T_total: causal = causal & ((q_pos.unsqueeze(1) - k_pos.unsqueeze(0)) <= window_left) scores = torch.where(causal.unsqueeze(0).unsqueeze(0), scores, torch.zeros_like(scores)) scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-8) scores = scores.to(v_all.dtype) y = torch.matmul(scores, v_all) # (B, H, T_new, D) y = y.transpose(1, 2).contiguous().reshape(B, T_new, C) return attn_module.c_proj(y), new_xh, new_v class YatFullGPTForCausalLM(PreTrainedModel, GenerationMixin): config_class = YatFullGPTHfConfig base_model_prefix = "yatnmn_full_gpt" supports_gradient_checkpointing = False _no_split_modules = ["Block"] _supports_cache_class = False _supports_static_cache = False def _supports_default_dynamic_cache(self): return False def __init__(self, config: YatFullGPTHfConfig): super().__init__(config) inner = YatFullGPTConfig( sequence_len=config.sequence_len, vocab_size=config.vocab_size, n_layer=config.n_layer, n_head=config.n_head, n_kv_head=config.n_kv_head, n_embd=config.n_embd, window_pattern=config.window_pattern, tie_embeddings=config.tie_embeddings, rope_base=config.rope_base, pad_vocab_size_to=config.pad_vocab_size_to, mlp_type=config.mlp_type, scalar_bias=config.scalar_bias, softplus_bias=config.softplus_bias, learnable_epsilon=config.learnable_epsilon, epsilon_init=config.epsilon_init, constant_alpha=config.constant_alpha, ) self.inner_config = inner self.model = YatFull_GPT(inner) self.post_init() def get_input_embeddings(self): return self.model.wte def set_input_embeddings(self, v): self.model.wte = v def can_generate(self): return True def _forward_full(self, input_ids): return self.model(input_ids) def _forward_with_cache(self, input_ids_new, past_key_values, prev_token_embed=None): m = self.model cfg = m.config B, T_new = input_ids_new.shape past_len = 0 if past_key_values is None else past_key_values[0][0].shape[2] T_total = past_len + T_new cos_full, sin_full = m._get_rope(T_total, m.wte.weight.dtype, m.wte.weight.device) cos = cos_full[:, past_len:T_total] sin = sin_full[:, past_len:T_total] x_new = rms_norm(m.wte(input_ids_new)) if past_len == 0: if T_new >= 2: gate = m.smear_lambda * torch.sigmoid(m.smear_gate(x_new[:, 1:, :24])) x_smeared = x_new[:, 1:] + gate * x_new[:, :-1] x = torch.cat([x_new[:, :1], x_smeared], dim=1) else: x = x_new else: assert prev_token_embed is not None x_cat = torch.cat([prev_token_embed, x_new], dim=1) gate = m.smear_lambda * torch.sigmoid(m.smear_gate(x_cat[:, 1:, :24])) x = x_cat[:, 1:] + gate * x_cat[:, :-1] x0 = x backout_layer = cfg.n_layer // 2 x_backout = None new_past = [] for i, block in enumerate(m.blocks): x = m.resid_lambdas[i] * x + m.x0_lambdas[i] * x0 past_xh = past_key_values[i][0] if past_key_values is not None else None past_v = past_key_values[i][1] if past_key_values is not None else None x_norm = rms_norm(x) attn_out, new_xh, new_v = _kvcache_yat_attn( block.attn, x_norm, cos, sin, m.window_sizes[i], past_xh, past_v, ) new_past.append((new_xh, new_v)) x = x + attn_out x = x + block.mlp(rms_norm(x)) if i == backout_layer: x_backout = x if x_backout is not None: x = x - m.backout_lambda * x_backout x = rms_norm(x) softcap = 15.0 logits = x @ m.wte.weight.t() if m.tie_embeddings else m.lm_head(x) logits = logits[..., : cfg.vocab_size].to(torch.float32) logits = softcap * torch.tanh(logits / softcap) last_embed = x_new[:, -1:, :] return logits, new_past, last_embed def forward( self, input_ids=None, attention_mask=None, past_key_values=None, labels=None, use_cache=None, output_attentions=None, output_hidden_states=None, return_dict=None, cache_position=None, **kwargs, ) -> CausalLMOutputWithPast: if input_ids is None: raise ValueError("input_ids required") use_cache = bool(use_cache) if use_cache is not None else (past_key_values is not None) kv_list = None prev_embed = None if past_key_values is not None: kv_list, prev_embed = past_key_values if use_cache: logits, new_past, new_last_embed = self._forward_with_cache( input_ids, kv_list, prev_token_embed=prev_embed, ) pkv = (tuple(new_past), new_last_embed) else: logits = self._forward_full(input_ids) pkv = None 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, ) return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=pkv, hidden_states=None, attentions=None, ) def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **kwargs): if past_key_values is not None: input_ids = input_ids[:, -1:] return {"input_ids": input_ids, "past_key_values": past_key_values, "use_cache": True} __all__ = ["YatFullGPTHfConfig", "YatFullGPTForCausalLM"]