import math from dataclasses import dataclass import torch import torch.nn.functional as F from torch import nn from transformers import PretrainedConfig, PreTrainedModel from transformers.modeling_outputs import CausalLMOutputWithPast from transformers.models.deepseek_v3.modeling_deepseek_v3 import DeepseekV3MoE from transformers.models.qwen2_moe.modeling_qwen2_moe import Qwen2MoeMLP try: from transformers.models.diffusion_gemma.generation_diffusion_gemma import ( DiffusionGemmaGenerationMixin, ) except ImportError: class DiffusionGemmaGenerationMixin: pass from transformers.cache_utils import DynamicCache from .attentions import KairosLiZAttention2, KairosNorm, KairosRotaryEmbedding class KairosConfig(PretrainedConfig): """modality_scales defaults every modality id up to num_modalities to scale 0 so.""" model_type = "kairos" def __init__( self, d_model=768, n_heads=12, n_layers=12, vocab_size=259, intermediate_size=2048, window_size=128, stride=5, **kwargs, ): super().__init__(**kwargs) self.hidden_size = d_model self.num_attention_heads = n_heads self.num_hidden_layers = n_layers self.vocab_size = vocab_size self.num_modalities = kwargs.get("num_modalities", 8) self.text_modality_id = kwargs.get("text_modality_id", 0) self.num_scales = kwargs.get("num_scales", 4) default_scales = {0: [0, 1], 1: [1, 2], 2: [2, 3]} for m in range(self.num_modalities): default_scales.setdefault(m, [0]) self.modality_scales = kwargs.get("modality_scales", default_scales) assert d_model % n_heads == 0, "hidden_size must be divisible by n_heads" self.stride = stride self.sliding_window_size = window_size self.num_key_value_heads = n_heads self.head_dim = d_model // n_heads self.attention_dropout = 0.0 self.rope_theta = 10000.0 self.max_position_embeddings = 4096 self.linear_num_value_heads = kwargs.get("linear_num_value_heads", n_heads) self.linear_num_key_heads = kwargs.get("linear_num_key_heads", n_heads) self.linear_key_head_dim = kwargs.get("linear_key_head_dim", self.head_dim) self.linear_value_head_dim = kwargs.get("linear_value_head_dim", self.head_dim) self.linear_conv_kernel_dim = kwargs.get("linear_conv_kernel_dim", 4) self.hidden_act = kwargs.get("hidden_act", "silu") self.rms_norm_eps = kwargs.get("rms_norm_eps", 1e-6) self.time_step_min = 0.001 self.time_step_max = 0.1 self.time_step_floor = 1e-4 self.A_init_range = (1.0, 16.0) self.initializer_range = kwargs.get("initializer_range", 0.02) self.intermediate_size = intermediate_size # only num_local_experts is real; n_routed_experts is a property alias below self.num_local_experts = kwargs.get("num_local_experts", kwargs.get("n_routed_experts", 8)) self.num_experts_per_tok = kwargs.get("num_experts_per_tok", 2) self.moe_intermediate_size = kwargs.get("moe_intermediate_size", intermediate_size) self.n_shared_experts = kwargs.get("n_shared_experts", 1) self.routed_scaling_factor = kwargs.get("routed_scaling_factor", 1.0) self.n_group = kwargs.get("n_group", 1) self.topk_group = kwargs.get("topk_group", 1) self.norm_topk_prob = kwargs.get("norm_topk_prob", False) self.use_moe = kwargs.get("use_moe", False) self.use_memory_gate = kwargs.get("use_memory_gate", False) self.layers_config = kwargs.get("layers_config", ["ld"] * n_layers) self.slw_wsize = kwargs.get("slw_wsize", -1) # v3 Block-AttnRes: windows prior layer outputs self.attnres_block_size = kwargs.get("attnres_block_size", 1) @property def n_routed_experts(self): """Alias for num_local_experts, the field DeepseekV3Experts actually reads.""" return self.num_local_experts @n_routed_experts.setter def n_routed_experts(self, value): self.num_local_experts = value class KairosCache(DynamicCache): """Cache for block-diffusion inference: `.clone()` before each denoising step to avoid state.""" def __init__(self, config): super().__init__() self.config = config self.conv_caches = [] self.ssm_caches = [] self._key_cache = {} self._value_cache = {} for idx, layer_type in enumerate(config.layers_config): if "l" in layer_type or "d" in layer_type: self._key_cache[idx] = None self._value_cache[idx] = None self.conv_caches.append(None) self.ssm_caches.append(None) self.window_size = config.sliding_window_size self.layers_config = config.layers_config self.past_length = [0 for _ in range(len(config.layers_config))] def update(self, k, v, layer_idx): added_len = k.size(1) k_cache = self._key_cache[layer_idx] v_cache = self._value_cache[layer_idx] if k_cache is None: k_cache, v_cache = k, v else: k_cache = torch.cat([k_cache, k], dim=1) v_cache = torch.cat([v_cache, v], dim=1) self._key_cache[layer_idx] = k_cache self._value_cache[layer_idx] = v_cache self.past_length[layer_idx] += added_len return k_cache, v_cache def trim(self, layer_idx): if "l" not in self.layers_config[layer_idx]: return window = min(self.window_size, self.config.slw_wsize) if self.config.slw_wsize > 0 else self.window_size k = self._key_cache[layer_idx] v = self._value_cache[layer_idx] if k is not None and k.size(1) > window: self._key_cache[layer_idx] = k[:, -window:, ...].contiguous() self._value_cache[layer_idx] = v[:, -window:, ...].contiguous() def get_ssm_cache(self, layer_idx): return (self.conv_caches[layer_idx], self.ssm_caches[layer_idx]) def get_total_seen(self, layer_idx): return self.past_length[layer_idx] def clone(self): new_cache = KairosCache(self.config) new_cache.conv_caches = [c.clone() if c is not None else None for c in self.conv_caches] new_cache.ssm_caches = [c.clone() if c is not None else None for c in self.ssm_caches] new_cache._key_cache = {k: v.clone() if v is not None else None for k, v in self._key_cache.items()} new_cache._value_cache = {k: v.clone() if v is not None else None for k, v in self._value_cache.items()} new_cache.past_length = self.past_length.copy() return new_cache class KairosMultiCache(DynamicCache): def __init__(self, config): super().__init__() self.config = config self.caches = [KairosCache(config) for _ in range(config.num_scales)] def get(self, idx): return self.caches[idx] def clone(self): out = KairosMultiCache.__new__(KairosMultiCache) out.config = self.config out.caches = [c.clone() for c in self.caches] return out class KairosMemoryGate(nn.Module): """Cross-attention gate over a low-rank bottleneck; blends state_t with a memory bank.""" def __init__(self, state_dim, bottleneck_dim=None): super().__init__() self.bottleneck_dim = bottleneck_dim or max(8, round(math.sqrt(state_dim))) self.down = nn.Linear(state_dim, self.bottleneck_dim) self.up = nn.Linear(self.bottleneck_dim, state_dim) self.context_attn = nn.MultiheadAttention(self.bottleneck_dim, num_heads=1, batch_first=True) def forward(self, state_t, memory=None): """state_t: (B, D). memory: (M, D) or None. Returns (B, D).""" if memory is None or memory.size(0) == 0: return state_t if memory.size(0) == 1: return memory[0].expand_as(state_t).contiguous() q = self.down(state_t).unsqueeze(1) kv = torch.cat([q, self.down(memory).unsqueeze(0).expand(state_t.size(0), -1, -1)], dim=1) out, _ = self.context_attn(q, kv, kv) return self.up(out.squeeze(1)) def gate_memory_bank(model, memory_caches: list, batch_size: int) -> "KairosMultiCache": """Gates a zero state_t against memory_caches to seed ssm_caches; no-op layers stay None.""" new_cache = KairosMultiCache(model.config) gate = model.memory_gate if gate is None: return new_cache for scale_idx, backbone in enumerate(model.backbones): for layer_idx in backbone.deltanet_layer_indices: parts = [] per_row_shape = None for c in memory_caches: s = c.caches[scale_idx].ssm_caches[layer_idx] if s is not None: per_row_shape = s.shape[1:] parts.append(s.reshape(s.shape[0], -1)) if not parts: continue memory = torch.cat(parts, dim=0) state_t = memory.new_zeros(batch_size, memory.shape[1]) blended = gate(state_t, memory) new_cache.caches[scale_idx].ssm_caches[layer_idx] = blended.reshape(batch_size, *per_row_shape) return new_cache class KairosFFN(Qwen2MoeMLP): pass class KairosMoE(DeepseekV3MoE): """DeepseekV3MoE's expert weights are raw torch.empty(), never initialized; fixed here.""" def __init__(self, config): super().__init__(config) std = getattr(config, "initializer_range", 0.02) self.experts.gate_up_proj.data.normal_(mean=0.0, std=std) self.experts.down_proj.data.normal_(mean=0.0, std=std) self.gate.weight.data.normal_(mean=0.0, std=std) # was torch.zeros() at construction; fine class DiffusionBlock(nn.Module): def __init__(self, config, layer_idx, use_moe=False): super().__init__() self.norm1 = KairosNorm(config.hidden_size) self.norm2 = KairosNorm(config.hidden_size) self.attn = KairosLiZAttention2(config, layer_idx) self.ffn = KairosMoE(config) if use_moe else KairosFFN(config) def forward(self, x, position_embeddings=None, cache_params=None, attention_mask=None, position_ids=None): x = x + self.attn( self.norm1(x), position_embeddings=position_embeddings, cache_params=cache_params, attention_mask=attention_mask, position_ids=position_ids, ) x = x + self.ffn(self.norm2(x)) return x class KairosCastingNorm(nn.RMSNorm): def forward(self, x): w = self.weight if self.weight.dtype == x.dtype else self.weight.to(x.dtype) return F.rms_norm(x, self.normalized_shape, w, self.eps) class KairosAttnRes(nn.Module): def __init__(self, n_embd): super().__init__() self.w = nn.Parameter(torch.zeros(n_embd)) self.key_norm = KairosCastingNorm(n_embd) def forward(self, prior_values): V = torch.stack(prior_values, dim=0) K = self.key_norm(V) logits = torch.einsum("d,lbtd->lbt", self.w, K) weights = F.softmax(logits, dim=0) return (weights.unsqueeze(-1) * V).sum(dim=0) class KairosDiffusionBackbone(nn.Module): """v3 Block-AttnRes: prior layer outputs are windowed into blocks before aggregation.""" def __init__(self, config, use_moe=False): super().__init__() self.layers = nn.ModuleList([DiffusionBlock(config, i, use_moe) for i in range(config.num_hidden_layers)]) self.norm = KairosNorm(config.hidden_size) self.aggregator = KairosAttnRes(config.hidden_size) self.attnres_block_size = max(1, getattr(config, "attnres_block_size", 1)) self.deltanet_layer_indices = [i for i, lt in enumerate(config.layers_config) if "d" in lt] def forward(self, x, position_embeddings=None, cache_params=None, attention_mask=None, position_ids=None): emb = x completed = [] # finalized block-sums of prior layer partial = None # running sum of the current in_block = 0 S = self.attnres_block_size def sources(): return [emb] + completed + ([partial] if partial is not None else []) for layer in self.layers: h = self.aggregator(sources()) x = layer( h, position_embeddings=position_embeddings, cache_params=cache_params, attention_mask=attention_mask, position_ids=position_ids, ) partial = x if partial is None else partial + x in_block += 1 if in_block == S: completed.append(partial) partial = None in_block = 0 return self.norm(x) class KairosEmbedding(nn.Module): def __init__(self, vocab_size: int, num_modalities: int, d_model: int): super().__init__() self.token_embed = nn.Embedding(vocab_size, d_model) self.modality_embed = nn.Embedding(num_modalities, d_model) self.fusion_proj = nn.Linear(d_model * 2, d_model) self.scale = d_model**0.5 def forward(self, token_ids, modality_ids): tok = self.token_embed(token_ids) mod = self.modality_embed(modality_ids) h = self.fusion_proj(torch.cat([tok, mod], dim=-1)) h = h * self.scale return h class OutputHead(nn.Module): def __init__(self, embedding: KairosEmbedding): super().__init__() d_model = embedding.token_embed.embedding_dim self.vocab_size = embedding.token_embed.num_embeddings self.num_modalities = embedding.modality_embed.num_embeddings self.token_head = nn.Linear(d_model, self.vocab_size, bias=False) self.modality_head = nn.Linear(d_model, self.num_modalities, bias=False) self.token_head.weight = embedding.token_embed.weight self.modality_head.weight = embedding.modality_embed.weight def forward(self, h): return self.token_head(h), self.modality_head(h) class KairosScaleRouter(nn.Module): """Gathers active positions per scale into a padded batch, runs the backbone.""" def __init__(self, modality_scales): super().__init__() self.modality_scales = modality_scales def build_active_mask(self, modality_ids, scale_len, scale_idx): device = modality_ids.device allowed = [m for m, scales in self.modality_scales.items() if scale_idx in scales] if not allowed: return torch.zeros(modality_ids.shape[0], scale_len, dtype=torch.bool, device=device) allowed_t = torch.tensor(allowed, device=device) active_full = torch.isin(modality_ids, allowed_t) pooled = F.adaptive_max_pool1d(active_full.float().unsqueeze(1), scale_len).squeeze(1) return pooled > 0.5 @staticmethod def gather_active(x, active_mask): _, _, D = x.shape lengths = active_mask.sum(dim=1) max_len = int(lengths.max().item()) if lengths.numel() > 0 else 0 if max_len == 0: return None, None, None order = torch.argsort((~active_mask).long(), dim=1, stable=True) positions = order[:, :max_len] gathered = torch.gather(x, 1, positions.unsqueeze(-1).expand(-1, -1, D)) arange = torch.arange(max_len, device=x.device).unsqueeze(0) pad_mask = arange < lengths.unsqueeze(1) return gathered, pad_mask, positions @staticmethod def scatter_active(output, chunk, pad_mask, positions): D = output.shape[-1] idx = positions.unsqueeze(-1).expand(-1, -1, D) current = torch.gather(output, 1, idx) values = torch.where(pad_mask.unsqueeze(-1), chunk.to(output.dtype), current) return output.scatter(1, idx, values) @dataclass class CodecOutput: scales: list length: int class PyramidalConvCodec(nn.Module): """Parallel multi-scale convolutional codec with modality routing.""" def __init__(self, d_model, stride=5, num_scales=4): super().__init__() self.stride = stride self.num_scales = num_scales self.encoders = nn.ModuleList() self.decoders = nn.ModuleList() for level in range(num_scales): scale_stride = stride ** (level + 1) kernel_size = scale_stride // 2 kernel_size += kernel_size % 2 == 0 padding = kernel_size // 2 self.encoders.append( nn.Conv1d( d_model, d_model, kernel_size=kernel_size, stride=scale_stride, padding=padding, groups=d_model ) ) self.decoders.append( nn.ConvTranspose1d( d_model, d_model, kernel_size=kernel_size, stride=scale_stride, padding=padding, output_padding=max(scale_stride - 1, 0), groups=d_model, ) ) self.norm = KairosNorm(d_model * num_scales) self.fusion = nn.Linear(d_model * num_scales, d_model) def encode(self, x): h = x.transpose(1, 2) scales = [encoder(h).transpose(1, 2) for encoder in self.encoders] return CodecOutput(scales=scales, length=x.shape[1]) def decode(self, encoded): scales = encoded.scales length = encoded.length reconstructed = [] for scale, decoder in zip(scales, self.decoders): h = decoder(scale.transpose(1, 2)) reconstructed.append(h.transpose(1, 2)) padded = [] for r in reconstructed: if r.shape[1] < length: pad_amount = length - r.shape[1] r = F.pad(r.transpose(1, 2), (0, pad_amount), mode="replicate").transpose(1, 2) padded.append(r[:, :length]) h = torch.cat(padded, dim=-1) h = self.norm(h) return self.fusion(h) @dataclass class KairosOutput(CausalLMOutputWithPast): encoder_last_hidden_state: torch.FloatTensor = None modality_logits: torch.FloatTensor = None class KairosDiffusionLLM(PreTrainedModel, DiffusionGemmaGenerationMixin): def __init__(self, config, vocab_size=None, use_moe=None): super().__init__(config) if use_moe is None: use_moe = config.use_moe self.codec = PyramidalConvCodec(d_model=config.hidden_size, stride=config.stride, num_scales=config.num_scales) self.router = KairosScaleRouter(config.modality_scales) if vocab_size is None: vocab_size = config.vocab_size self.embedding = KairosEmbedding( vocab_size=vocab_size, num_modalities=config.num_modalities, d_model=config.hidden_size ) self.backbones = nn.ModuleList( [KairosDiffusionBackbone(config=config, use_moe=use_moe) for _ in range(self.codec.num_scales)] ) if getattr(config, "use_memory_gate", False): head_dim = config.hidden_size // config.num_attention_heads state_dim = config.num_attention_heads * head_dim * 2 * head_dim self.memory_gate = KairosMemoryGate(state_dim=state_dim) else: self.memory_gate = None self.rotary = KairosRotaryEmbedding(config, config.head_dim) self.norm = KairosNorm(config.hidden_size) self.lm_head = OutputHead(self.embedding) self.post_init() # triggers _init_weights on every submodule/parameter def _init_weights(self, module): """Every PreTrainedModel subclass must define this (the base class default is a.""" 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) def forward( self, input_ids=None, decoder_input_ids=None, modality_ids=None, attention_mask=None, self_conditioning_logits=None, cache_params=None, **kwargs, ): x = decoder_input_ids if decoder_input_ids is not None else input_ids if x is None: raise ValueError("either input_ids or decoder_input_ids must be provided") if modality_ids is None: modality_ids = torch.full_like(x, self.config.text_modality_id) h = self.embedding(token_ids=x, modality_ids=modality_ids) if self_conditioning_logits is not None: probs = torch.softmax(self_conditioning_logits, dim=-1) h = h + (probs @ self.embedding.token_embed.weight) encoded = self.codec.encode(h) features = [] for scale_idx, (scale, backbone) in enumerate(zip(encoded.scales, self.backbones)): output = scale.clone() local_cache = cache_params.get(scale_idx) if cache_params is not None else None active_mask = self.router.build_active_mask(modality_ids, scale.shape[1], scale_idx) if attention_mask is not None: pad_pool = F.adaptive_max_pool1d(attention_mask.float().unsqueeze(1), scale.shape[1]).squeeze(1) active_mask = active_mask & (pad_pool > 0.5) gathered, pad_mask, positions = self.router.gather_active(scale, active_mask) if gathered is not None: cache_offset = local_cache.get_total_seen(0) if local_cache is not None else 0 position_ids = positions + cache_offset cos, sin = self.rotary(scale, position_ids, max_position=None) chunk = backbone( gathered, position_embeddings=(cos, sin), cache_params=local_cache, attention_mask=pad_mask, position_ids=position_ids, ) output = self.router.scatter_active(output, chunk, pad_mask, positions) features.append(output) decoded = CodecOutput(scales=features, length=encoded.length) h = self.codec.decode(decoded) h = self.norm(h) token_logits, modality_logits = self.lm_head(h) return KairosOutput(logits=token_logits, modality_logits=modality_logits, past_key_values=cache_params)