"""Logos — looped decoder-only transformer with hybrid attention variants. Each block selects one attention mechanism per execution from: * ``kda`` — Kimi Delta Attention, with no snapshot memory branch. * ``swa`` — local sliding-window softmax attention. * ``csa`` — 4-token compressed sparse global attention. * ``hca`` — heavily compressed dense global attention. The model is partitioned into Entry -> Body -> Exit. Body blocks are shared across ``num_loops`` iterations, and their attention kind can vary per loop using a flattened loop-major ``body_attn_pattern``. """ from __future__ import annotations from dataclasses import dataclass from typing import Any, Dict, List, Optional, Tuple import torch import torch.nn as nn import torch.nn.functional as F import torch.utils.checkpoint as ckpt_utils from .baseline import ( RMSNorm, SwiGLU, MoELayer, combine_lm_and_aux_loss, init_moe_router_weights, count_parameters, model_summary, ) from .hybrid import ( HybridConfig, HybridAttentionLayer, expand_attention_pattern, normalize_attention_type, parse_attention_pattern, ) from .residual import BlockAttentionResidual from .lm_loss import ( chunked_linear_cross_entropy, chunked_token_superposition_cross_entropy, lm_cross_entropy_from_logits, standard_lm_cross_entropy, token_superposition_attention_mask, token_superposition_embeddings, ) def _legacy_kind(layer_idx: int, config: "LogosConfig") -> str: return "swa" if (layer_idx % config.swa_every) == config.swa_offset else "kda" def _default_entry_schedule(config: "LogosConfig") -> List[str]: return [_legacy_kind(i, config) for i in range(config.num_entry_layers)] def _default_body_schedule(config: "LogosConfig") -> List[str]: out: List[str] = [] body_offset = config.num_entry_layers for _ in range(config.num_loops): for r in range(config.num_body_layers): out.append(_legacy_kind(body_offset + r, config)) return out def _default_exit_schedule(config: "LogosConfig") -> List[str]: exit_offset = config.num_entry_layers + config.num_body_layers return [ _legacy_kind(exit_offset + i, config) for i in range(config.num_exit_layers) ] def _resolve_logos_attention_schedules( config: "LogosConfig", ) -> Tuple[List[str], List[str], List[str]]: n_entry = config.num_entry_layers n_body_exec = config.num_body_layers * config.num_loops n_exit = config.num_exit_layers section_patterns = ( config.entry_attn_pattern, config.body_attn_pattern, config.exit_attn_pattern, ) if config.attn_pattern and all(p is None for p in section_patterns): full = expand_attention_pattern( config.attn_pattern, n_entry + n_body_exec + n_exit, default="kda", ) entry = full[:n_entry] body = full[n_entry:n_entry + n_body_exec] exit_ = full[n_entry + n_body_exec:] return entry, body, exit_ entry = ( expand_attention_pattern(config.entry_attn_pattern, n_entry, default="kda") if config.entry_attn_pattern is not None else _default_entry_schedule(config) ) body = ( expand_attention_pattern(config.body_attn_pattern, n_body_exec, default="kda") if config.body_attn_pattern is not None else _default_body_schedule(config) ) exit_ = ( expand_attention_pattern(config.exit_attn_pattern, n_exit, default="kda") if config.exit_attn_pattern is not None else _default_exit_schedule(config) ) return entry, body, exit_ @dataclass class LogosConfig(HybridConfig): # Auto-derived from entry + body + exit in __post_init__. num_layers: int = 0 num_entry_layers: int = 2 num_body_layers: int = 4 num_exit_layers: int = 2 num_loops: int = 4 # Fine-grained attention schedules. ``body_attn_pattern`` is expanded to # ``num_loops * num_body_layers`` entries in loop-major order: # loop0.block0, loop0.block1, ..., loop1.block0, ... entry_attn_pattern: Optional[str] = None body_attn_pattern: Optional[str] = None exit_attn_pattern: Optional[str] = None entry_top_k: Optional[int] = None exit_top_k: Optional[int] = None gradient_checkpointing: bool = False ckpt_granularity: str = "per-block" def __post_init__(self): super().__post_init__() if self.num_body_layers <= 0 or self.num_loops <= 0: raise ValueError( "num_body_layers and num_loops must both be > 0" ) if self.num_entry_layers < 0 or self.num_exit_layers < 0: raise ValueError( "num_entry_layers and num_exit_layers must be >= 0" ) if self.ckpt_granularity not in ("per-block", "per-loop"): raise ValueError( "ckpt_granularity must be 'per-block' or 'per-loop', " f"got {self.ckpt_granularity!r}" ) for pattern in ( self.entry_attn_pattern, self.body_attn_pattern, self.exit_attn_pattern, ): parse_attention_pattern(pattern) self.num_layers = ( self.num_entry_layers + self.num_body_layers + self.num_exit_layers ) class LogosTransformerBlock(nn.Module): """A Logos parameter-block with selectable attention kind per call.""" def __init__( self, config: LogosConfig, attention_kinds: List[str], num_loops: int = 1, top_k: Optional[int] = None, ): super().__init__() self.use_moe = config.use_moe self.attention_kinds = [normalize_attention_type(k) for k in attention_kinds] isolate_res = getattr( config, "block_residual_isolate_softmax", False, ) self.attn_norm = RMSNorm(config.d_model, eps=config.norm_eps) self.attn = HybridAttentionLayer(config, self.attention_kinds) self.attn_res = BlockAttentionResidual( config.d_model, eps=config.norm_eps, isolate_softmax=isolate_res, ) self.ffn_norm = RMSNorm(config.d_model, eps=config.norm_eps) if config.use_moe: self.ffn = MoELayer(config, num_loops=num_loops, top_k=top_k) else: self.ffn = SwiGLU(config.d_model, config.d_ff) self.ffn_res = BlockAttentionResidual( config.d_model, eps=config.norm_eps, isolate_softmax=isolate_res, ) def forward( self, blocks: List[torch.Tensor], partial: Optional[torch.Tensor], attention_kind: str, attention_mask: Optional[torch.Tensor] = None, is_causal: bool = True, cache: Optional[Dict[str, Any]] = None, loop_idx: int = 0, position_offset: int = 0, ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], torch.Tensor]: h = self.attn_res(blocks, partial) attn_out, index_loss = self.attn( attention_kind, self.attn_norm(h), attention_mask=attention_mask, is_causal=is_causal, cache=cache, position_offset=position_offset, ) if partial is None: partial = attn_out else: partial = partial + attn_out h = self.ffn_res(blocks, partial) if self.use_moe: ffn_out, aux_loss, topk_indices = self.ffn(self.ffn_norm(h), loop_idx=loop_idx) partial = partial + ffn_out return partial, aux_loss, topk_indices, index_loss partial = partial + self.ffn(self.ffn_norm(h)) zero = torch.zeros((), device=partial.device, dtype=partial.dtype) return partial, zero, None, index_loss class LogosTransformer(nn.Module): def __init__(self, config: LogosConfig): super().__init__() self.config = config self.entry_attn_schedule, self.body_attn_schedule, self.exit_attn_schedule = ( _resolve_logos_attention_schedules(config) ) self.token_emb = nn.Embedding(config.vocab_size, config.d_model) self.entry = nn.ModuleList([ LogosTransformerBlock( config, attention_kinds=[self.entry_attn_schedule[i]], num_loops=1, top_k=config.entry_top_k, ) for i in range(config.num_entry_layers) ]) self.body = nn.ModuleList([ LogosTransformerBlock( config, attention_kinds=[ self.body_attn_schedule[l * config.num_body_layers + i] for l in range(config.num_loops) ], num_loops=config.num_loops, ) for i in range(config.num_body_layers) ]) self.exit = nn.ModuleList([ LogosTransformerBlock( config, attention_kinds=[self.exit_attn_schedule[i]], num_loops=1, top_k=config.exit_top_k, ) for i in range(config.num_exit_layers) ]) self.final_res = BlockAttentionResidual( config.d_model, eps=config.norm_eps, isolate_softmax=getattr( config, "block_residual_isolate_softmax", False, ), ) self.final_norm = RMSNorm(config.d_model, eps=config.norm_eps) self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False) self.lm_head.weight = self.token_emb.weight self._init_weights() def _init_weights(self): for module in self.modules(): if isinstance(module, nn.Linear): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) if module.bias is not None: torch.nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) init_moe_router_weights(self, self.config.router_init_std) for module in self.modules(): if isinstance(module, BlockAttentionResidual): nn.init.zeros_(module.proj) def _lm_loss( self, hidden: torch.Tensor, labels: torch.Tensor, logits: Optional[torch.Tensor] = None, token_superposition_bag_size: int = 1, ) -> torch.Tensor: chunk_size = int(getattr(self.config, "lm_head_chunk_size", 0) or 0) if int(token_superposition_bag_size) > 1: if chunk_size > 0 and logits is None: return chunked_token_superposition_cross_entropy( hidden, self.lm_head.weight, labels, int(token_superposition_bag_size), chunk_size=chunk_size, ignore_index=-100, ) if logits is None: logits = self.lm_head(hidden) return lm_cross_entropy_from_logits( logits, labels, token_superposition_bag_size=token_superposition_bag_size, ignore_index=-100, ) if chunk_size > 0 and logits is None: return chunked_linear_cross_entropy( hidden, self.lm_head.weight, labels, chunk_size=chunk_size, ignore_index=-100, ) if logits is None: logits = self.lm_head(hidden) return standard_lm_cross_entropy( logits, labels, ignore_index=-100, ) def forward( self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, is_causal: bool = True, token_superposition_bag_size: int = 1, past_key_values: Optional[Dict[str, Any]] = None, use_cache: bool = False, ) -> Dict[str, Any]: if use_cache and int(token_superposition_bag_size) != 1: raise ValueError("Logos cache inference requires token_superposition_bag_size=1") x = token_superposition_embeddings( self.token_emb, input_ids, token_superposition_bag_size, ) attention_mask = token_superposition_attention_mask( attention_mask, token_superposition_bag_size, ) if use_cache and attention_mask is not None and attention_mask.size(1) != x.size(1): attention_mask = attention_mask[:, -x.size(1):] cache_state: Optional[Dict[str, Any]] = None layer_caches: Optional[Dict[str, Dict[str, Any]]] = None position_offset = 0 if use_cache: cache_state = past_key_values if past_key_values is not None else {} layer_caches = cache_state.setdefault("layers", {}) position_offset = int(cache_state.get("seen_tokens", 0) or 0) aux_loss = torch.zeros((), device=input_ids.device, dtype=x.dtype) index_loss = torch.zeros((), device=input_ids.device, dtype=x.dtype) topk_indices_list: List[Optional[torch.Tensor]] = [] blocks: List[torch.Tensor] = [x] partial: Optional[torch.Tensor] = None use_ckpt = self.config.gradient_checkpointing and self.training per_loop_ckpt = ( use_ckpt and getattr(self.config, "ckpt_granularity", "per-block") == "per-loop" ) per_block_ckpt = use_ckpt and not per_loop_ckpt def _layer_cache(name: str) -> Optional[Dict[str, Any]]: if layer_caches is None: return None return layer_caches.setdefault(name, {}) def _call_block(block_module, blocks_in, partial_in, attention_kind, loop_idx, cache_name): return block_module( blocks_in, partial_in, attention_kind, attention_mask=attention_mask, is_causal=is_causal, cache=_layer_cache(cache_name), loop_idx=loop_idx, position_offset=position_offset, ) def _ckpt_block(block_module, blocks_in, partial_in, attention_kind, loop_idx, cache_name=None): return ckpt_utils.checkpoint( block_module, blocks_in, partial_in, attention_kind, attention_mask=attention_mask, is_causal=is_causal, cache=None, loop_idx=loop_idx, position_offset=position_offset, use_reentrant=False, ) for idx, layer in enumerate(self.entry): partial, layer_aux, layer_topk, layer_index = _call_block( layer, blocks, partial, self.entry_attn_schedule[idx], 0, f"entry.{idx}", ) aux_loss = aux_loss + layer_aux index_loss = index_loss + layer_index topk_indices_list.append(layer_topk) if self.config.num_entry_layers > 0: assert partial is not None, "entry produced no partial block" blocks = blocks + [partial] partial = None for loop_idx in range(self.config.num_loops): if per_loop_ckpt: _li = loop_idx def _body_loop(blks, p, loop_i=_li): aux_sum = torch.zeros((), device=blks[0].device, dtype=blks[0].dtype) index_sum = torch.zeros((), device=blks[0].device, dtype=blks[0].dtype) topks: List[Optional[torch.Tensor]] = [] for r, block in enumerate(self.body): kind = self.body_attn_schedule[loop_i * self.config.num_body_layers + r] p, la, lt, li = block( blks, p, kind, attention_mask=attention_mask, is_causal=is_causal, cache=None, loop_idx=loop_i, position_offset=position_offset, ) aux_sum = aux_sum + la index_sum = index_sum + li topks.append(lt) return p, aux_sum, topks, index_sum partial, loop_aux, loop_topks, loop_index = ckpt_utils.checkpoint( _body_loop, blocks, partial, use_reentrant=False, ) aux_loss = aux_loss + loop_aux index_loss = index_loss + loop_index topk_indices_list.extend(loop_topks) else: if per_block_ckpt: runner = _ckpt_block else: runner = _call_block for r, block in enumerate(self.body): kind = self.body_attn_schedule[ loop_idx * self.config.num_body_layers + r ] partial, layer_aux, layer_topk, layer_index = runner( block, blocks, partial, kind, loop_idx, f"body.{loop_idx}.{r}", ) aux_loss = aux_loss + layer_aux index_loss = index_loss + layer_index topk_indices_list.append(layer_topk) assert partial is not None, f"body loop {loop_idx} produced no partial block" blocks = blocks + [partial] partial = None for idx, layer in enumerate(self.exit): partial, layer_aux, layer_topk, layer_index = _call_block( layer, blocks, partial, self.exit_attn_schedule[idx], 0, f"exit.{idx}", ) aux_loss = aux_loss + layer_aux index_loss = index_loss + layer_index topk_indices_list.append(layer_topk) h_main = self.final_res(blocks, partial) x = self.final_norm(h_main) use_chunked_lm_loss = ( labels is not None and int(getattr(self.config, "lm_head_chunk_size", 0) or 0) > 0 ) logits = None if use_chunked_lm_loss else self.lm_head(x) lm_loss: Optional[torch.Tensor] = None if labels is not None: lm_loss = self._lm_loss( x, labels, logits=logits, token_superposition_bag_size=token_superposition_bag_size, ) loss = combine_lm_and_aux_loss( lm_loss, aux_loss if self.config.use_moe else None, self.training, ) if loss is not None and self.training: loss = loss + self.config.csa_indexer_loss_weight * index_loss if use_cache and cache_state is not None: cache_state["seen_tokens"] = int(position_offset + x.size(1)) return { "logits": logits, "loss": loss, "lm_loss": lm_loss, "aux_loss": aux_loss if self.config.use_moe else None, "indexer_loss": index_loss, "topk_indices": topk_indices_list if self.config.use_moe else None, "past_key_values": cache_state if use_cache else None, } def update_router_biases(self, topk_indices_list: List[Optional[torch.Tensor]]) -> None: if not self.config.use_moe: return n_entry = self.config.num_entry_layers n_body = self.config.num_body_layers n_loops = self.config.num_loops for i, layer in enumerate(self.entry): topk = topk_indices_list[i] if topk is not None and isinstance(layer.ffn, MoELayer): layer.ffn.update_bias(topk, loop_idx=0) body_offset = n_entry for r, block in enumerate(self.body): if not isinstance(block.ffn, MoELayer): continue topk_per_loop: List[torch.Tensor] = [] valid = True for l in range(n_loops): idx = body_offset + l * n_body + r topk = topk_indices_list[idx] if topk is None: valid = False break topk_per_loop.append(topk) if valid: block.ffn.update_bias_per_loop(topk_per_loop) exit_offset = n_entry + n_loops * n_body for i, layer in enumerate(self.exit): topk = topk_indices_list[exit_offset + i] if topk is not None and isinstance(layer.ffn, MoELayer): layer.ffn.update_bias(topk, loop_idx=0) @torch.no_grad() def get_balance_stats(self) -> Dict[str, float]: if not self.config.use_moe: return {} stats: Dict[str, float] = {} def _record(name: str, layer: nn.Module, kind: str) -> None: ffn = layer.ffn if not hasattr(ffn, "bias"): return bias = ffn.bias stats[f"{name}_{kind}_bias_mean"] = bias.abs().mean().item() stats[f"{name}_{kind}_bias_max"] = bias.abs().max().item() for idx, layer in enumerate(self.entry): _record(f"entry{idx}", layer, self.entry_attn_schedule[idx]) for idx, block in enumerate(self.body): kinds = sorted({ self.body_attn_schedule[l * self.config.num_body_layers + idx] for l in range(self.config.num_loops) }) _record(f"body{idx}", block, "+".join(kinds)) for idx, layer in enumerate(self.exit): _record(f"exit{idx}", layer, self.exit_attn_schedule[idx]) return stats @torch.no_grad() def generate( self, input_ids: torch.Tensor, max_new_tokens: int = 100, temperature: float = 1.0, top_k: Optional[int] = None, attention_mask: Optional[torch.Tensor] = None, eos_token_id: Optional[int] = None, use_cache: bool = True, ) -> torch.Tensor: self.train(False) batch_size = input_ids.size(0) generated = input_ids cache: Optional[Dict[str, Any]] = None model_input = input_ids model_attention_mask = attention_mask for _ in range(max_new_tokens): outputs = self.forward( model_input, attention_mask=model_attention_mask, is_causal=True, past_key_values=cache, use_cache=use_cache, ) cache = outputs.get("past_key_values") if use_cache else None logits = outputs["logits"][:, -1, :] / temperature if top_k is not None: v, _ = torch.topk(logits, min(top_k, logits.size(-1))) logits = logits.masked_fill(logits < v[:, [-1]], float("-inf")) probs = F.softmax(logits, dim=-1) next_token = torch.multinomial(probs, num_samples=1) generated = torch.cat([generated, next_token], dim=-1) if attention_mask is not None: attention_mask = torch.cat([ attention_mask, torch.ones( (batch_size, 1), device=attention_mask.device, dtype=attention_mask.dtype, ), ], dim=-1) if eos_token_id is not None and (next_token == eos_token_id).all(): break if use_cache: model_input = next_token model_attention_mask = ( attention_mask[:, -1:] if attention_mask is not None else None ) else: model_input = generated model_attention_mask = attention_mask return generated __all__ = [ "LogosConfig", "LogosTransformerBlock", "LogosTransformer", "count_parameters", "model_summary", ]