"""Portable ARLoop inference: PyTorch SDPA and FP32 attention-residual routing. Adapted from ARLoop's model and Parcae building blocks; see LICENSE-MIT. No training repository, FLA, Triton, or flash-attn installation is needed. """ import torch from torch import nn from torch.nn import functional as F from transformers import PreTrainedModel, GenerationMixin from transformers.cache_utils import DynamicCache from transformers.modeling_outputs import CausalLMOutputWithPast from .configuration_arloop import ARLoopConfig class RMSNorm(nn.Module): def __init__(self, size, eps): super().__init__() self.weight = nn.Parameter(torch.ones(size)) self.eps = eps def forward(self, x): # Preserve the native final norm's rounding and weight promotion. y = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + self.eps) return y.to(x.dtype) * self.weight class Attention(nn.Module): def __init__(self, config, layer): super().__init__() self.config = config dim = config.hidden_size kv_dim = config.num_key_value_heads * config.head_dim self.c_q = nn.Linear(dim, dim, bias=False) self.c_k = nn.Linear(dim, kv_dim, bias=False) self.c_v = nn.Linear(dim, kv_dim, bias=False) self.c_proj = nn.Linear(dim, dim, bias=False) self.ve_gate = (nn.Linear(32, config.num_key_value_heads, bias=False) if layer % 2 == (config.physical_layers - 1) % 2 else None) def forward(self, x, positions, mask, value_embedding, cache, effective_layer): c = self.config batch, length, _ = x.shape q = self.c_q(x).view(batch, length, c.num_attention_heads, c.head_dim) k = self.c_k(x).view(batch, length, c.num_key_value_heads, c.head_dim) v = self.c_v(x).view_as(k) if self.ve_gate is not None: gate = 2 * torch.sigmoid(self.ve_gate(x[..., :32])) v = v + gate.unsqueeze(-1) * value_embedding.view_as(v).to(v.dtype) # Adjacent-pair rotary embedding, unlike Llama's split-half convention. freq = 1.0 / (c.rope_theta ** (torch.arange(0, c.head_dim, 2, device=x.device, dtype=torch.float32) / c.head_dim)) angles = positions.float().unsqueeze(-1) / c.rope_condense_ratio * freq cos, sin = angles.cos().unsqueeze(2), angles.sin().unsqueeze(2) def rotate(tensor): pairs = tensor.float().unflatten(-1, (-1, 2)) return torch.stack((pairs[..., 0] * cos - pairs[..., 1] * sin, pairs[..., 1] * cos + pairs[..., 0] * sin), -1).flatten(-2).to(tensor.dtype) q, k = rotate(q), rotate(k) if c.qk_norm: # Native Q/K normalization uses dtype-dependent default epsilon. q = F.rms_norm(q, (c.head_dim,)) k = F.rms_norm(k, (c.head_dim,)) q, k, v = (t.transpose(1, 2) for t in (q, k, v)) if cache is not None: k, v = cache.update(k, v, effective_layer) causal = mask is None and q.shape[-2] == k.shape[-2] if mask is None and q.shape[-2] != k.shape[-2] and q.shape[-2] != 1: rows = torch.arange(k.shape[-2] - q.shape[-2], k.shape[-2], device=x.device) mask = torch.arange(k.shape[-2], device=x.device)[None, :] <= rows[:, None] y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, is_causal=causal, enable_gqa=c.num_attention_heads != c.num_key_value_heads, dropout_p=0.0) return self.c_proj(y.transpose(1, 2).contiguous().view(batch, length, c.hidden_size)) class MLP(nn.Module): def __init__(self, config): super().__init__() self.fc = nn.Linear(config.hidden_size, 2 * config.intermediate_size, bias=False) self.proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) def forward(self, x): gate, value = self.fc(x).chunk(2, dim=-1) return self.proj(F.silu(gate) * value) class ARLoopBlock(nn.Module): def __init__(self, config, layer): super().__init__() self.attn_norm = RMSNorm(config.hidden_size, config.rms_norm_eps) self.attn = Attention(config, layer) self.mlp_norm = RMSNorm(config.hidden_size, config.rms_norm_eps) self.mlp = MLP(config) class ARLoopForCausalLM(PreTrainedModel, GenerationMixin): config_class = ARLoopConfig base_model_prefix = "" main_input_name = "input_ids" _no_split_modules = ["ARLoopBlock"] _supports_cache_class = True def __init__(self, config): super().__init__(config) self.transformer = nn.ModuleDict({ "wte": nn.Embedding(config.vocab_size, config.hidden_size), "h": nn.ModuleList(ARLoopBlock(config, i) for i in range(config.physical_layers)), "ln_f": RMSNorm(config.hidden_size, config.rms_norm_eps), }) self.shared_value_embedding = nn.Embedding(config.vocab_size, config.num_key_value_heads * config.head_dim) self.attnres_queries = nn.Parameter(torch.zeros(2 * config.num_hidden_layers, config.hidden_size)) self.attnres_norms = nn.ModuleList(RMSNorm(config.hidden_size, config.rms_norm_eps) for _ in range(2 * config.num_hidden_layers)) self.final_attnres_query = nn.Parameter(torch.zeros(config.hidden_size)) self.final_attnres_norm = RMSNorm(config.hidden_size, config.rms_norm_eps) self.post_init() def _init_weights(self, module): if isinstance(module, (nn.Linear, nn.Embedding)): nn.init.normal_(module.weight, std=0.02) elif isinstance(module, RMSNorm): nn.init.ones_(module.weight) def get_input_embeddings(self): return self.transformer.wte def set_input_embeddings(self, value): self.transformer.wte = value def _route(self, query, norm, blocks, partial, output_norm=None): values = blocks + [partial] if len(values) > self.config.sw_blocks + 1: values = [values[0], *values[-self.config.sw_blocks:]] # FLA accumulates routing scores, softmax, and the fused output norm in # FP32, rounding only the final result back to the residual dtype. with torch.autocast(device_type=partial.device.type, enabled=False): stacked = torch.stack(values, dim=-2).float() scores = (stacked * (query.float() * norm.weight.float())).sum(-1) scores = scores * torch.rsqrt(stacked.square().mean(-1) + norm.eps) weights = scores.softmax(dim=-1) y = torch.zeros_like(partial, dtype=torch.float32) for i, value in enumerate(values): y = y + weights[..., i, None] * value.float() if output_norm is not None: y = y * torch.rsqrt(y.square().mean(-1, keepdim=True) + norm.eps) y = y * output_norm.weight.float() return y.to(partial.dtype) def forward(self, input_ids=None, attention_mask=None, position_ids=None, past_key_values=None, inputs_embeds=None, labels=None, use_cache=None, return_dict=None, output_hidden_states=False, output_attentions=False, cache_position=None, logits_to_keep=0, **kwargs): if input_ids is None or inputs_embeds is not None: raise ValueError("ARLoop requires input_ids for token-dependent value embeddings") if output_attentions: raise ValueError("output_attentions is not implemented") use_cache = self.config.use_cache if use_cache is None else use_cache if labels is not None: use_cache = False if past_key_values is not None and not use_cache: raise ValueError("past_key_values requires use_cache=True") cache = past_key_values if use_cache else None if use_cache and cache is None: cache = DynamicCache() past_length = cache.get_seq_length() if cache is not None else 0 batch, length = input_ids.shape total = past_length + length if total > self.config.max_position_embeddings: raise ValueError("Input and cache exceed max_position_embeddings") if position_ids is None: if attention_mask is None: position_ids = torch.arange(past_length, total, device=input_ids.device)[None, :] else: position_ids = (attention_mask.long().cumsum(-1) - 1).clamp_min(0)[:, -length:] key_indices = torch.arange(total, device=input_ids.device) query_indices = torch.arange(past_length, total, device=input_ids.device) mask = None if attention_mask is not None: if attention_mask.shape != (batch, total): raise ValueError("attention_mask must cover both cached and new tokens") if not bool(attention_mask.all()): mask = (key_indices[None, :] <= query_indices[:, None])[None, None, :, :] mask = mask & attention_mask[:, None, None, :].bool() # Use BF16 activations; parameters follow the requested loading dtype. with torch.autocast(device_type=input_ids.device.type, dtype=torch.bfloat16): embedded = (self.transformer.wte(input_ids) * self.config.embedding_scale).to(torch.bfloat16) ve = self.shared_value_embedding(input_ids) blocks, partial = [embedded], embedded states = [embedded] if output_hidden_states else None for loop in range(self.config.loop_count): for i, block in enumerate(self.transformer.h): effective = loop * self.config.physical_layers + i h = self._route(self.attnres_queries[2 * effective], self.attnres_norms[2 * effective], blocks, partial, block.attn_norm) if effective > 0 and effective % (self.config.attnres_block_size // 2) == 0: blocks.append(partial) partial = None a = block.attn(h, position_ids, mask, ve, cache, effective) partial = a if partial is None else partial + a h = self._route(self.attnres_queries[2 * effective + 1], self.attnres_norms[2 * effective + 1], blocks, partial, block.mlp_norm) partial = partial + block.mlp(h) if states is not None: states.append(partial) h = self._route(self.final_attnres_query, self.final_attnres_norm, blocks, partial) # Native RMSNorm disables autocast around the final normalization. with torch.autocast(device_type=input_ids.device.type, enabled=False): h = self.transformer.ln_f(h) if logits_to_keep and labels is None: h = h[:, -logits_to_keep:, :] logits = F.linear(h, self.transformer.wte.weight).float() if self.config.logit_softcap is not None: cap = self.config.logit_softcap logits = cap * torch.tanh(logits / cap) loss = None if labels is not None: loss = F.cross_entropy(logits[:, :-1].contiguous().view(-1, logits.size(-1)), labels[:, 1:].contiguous().view(-1), ignore_index=-100) result = CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=cache, hidden_states=tuple(states) if states is not None else None) return result if return_dict is not False else result.to_tuple() def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **kwargs): if past_key_values is not None: seen = past_key_values.get_seq_length() if seen: input_ids = input_ids[:, seen:] if input_ids.shape[1] > seen else input_ids[:, -1:] return dict(input_ids=input_ids.contiguous(), past_key_values=past_key_values, attention_mask=attention_mask, use_cache=kwargs.get("use_cache", True), logits_to_keep=1)