from __future__ import annotations import os from dataclasses import dataclass from typing import Optional import torch import torch.nn as nn import torch.nn.functional as F def _detect_cpu_bf16() -> bool: try: with open("/proc/cpuinfo") as f: return "avx512_bf16" in f.read() except Exception: return False _CPU_HAS_BF16: bool = _detect_cpu_bf16() try: import pyfftw pyfftw.interfaces.cache.enable() pyfftw.interfaces.cache.set_keepalive_time(60.0) except ImportError: pass try: if os.environ.get("NO_CPP_EXT"): raise ImportError("C++ extensions disabled via NO_CPP_EXT") import torch.utils.cpp_extension as _cpp_ext _ext_path = os.path.join(os.path.dirname(__file__), "csrc") _gla_ext = ( _cpp_ext.load( name="gla_scan_cpu", sources=[os.path.join(_ext_path, "gla_scan.cpp")], extra_cflags=["-O3", "-fopenmp", "-march=native"], extra_ldflags=["-fopenmp"], verbose=False, ) if os.path.exists(os.path.join(_ext_path, "gla_scan.cpp")) else None ) except Exception: _gla_ext = None try: if os.environ.get("NO_CPP_EXT"): raise ImportError("C++ extensions disabled via NO_CPP_EXT") import torch.utils.cpp_extension as _fno_cpp_ext _fno_ext_src = os.path.join(os.path.dirname(__file__), "csrc", "fno_conv.cpp") _fno_ext = ( _fno_cpp_ext.load( name="fno_conv_cpu", sources=[_fno_ext_src], extra_cflags=["-O3", "-fopenmp", "-march=native"], extra_ldflags=["-fopenmp"], verbose=False, ) if os.path.exists(_fno_ext_src) else None ) except Exception: _fno_ext = None try: from flashfftconv import FlashFFTConv as _FlashFFTConv _flash_fft_available = True except ImportError: _FlashFFTConv = None _flash_fft_available = False try: from fla.ops.gla import chunk_gla as _fla_chunk_gla _fla_available = True try: from fla.ops import chunk_gated_delta_rule as _fla_chunk_gdr from fla.layers import GatedDeltaNet as _FlaGatedDeltaNet except Exception: _fla_chunk_gdr = None _FlaGatedDeltaNet = None except ImportError: _fla_chunk_gla = None _fla_chunk_gdr = None _FlaGatedDeltaNet = None _fla_available = False def _fft_dtype_to_torch(fft_dtype: str): return {"fp32": torch.float32, "bf16": torch.bfloat16, "fp16": torch.float16}[ fft_dtype ] def _gla_scan_py(q_k, k_k, v_f, chunk_gate, CS): B, H, T, D = q_k.shape n_chunks = T // CS bf16 = q_k.dtype == torch.bfloat16 q_c = q_k.reshape(B, H, n_chunks, CS, D) k_c = k_k.reshape(B, H, n_chunks, CS, D) v_c = v_f.reshape(B, H, n_chunks, CS, D) state = torch.zeros(B, H, D, D, device=q_k.device, dtype=torch.float32) z_norm = torch.zeros(B, H, D, device=q_k.device, dtype=torch.float32) chunks = [] for c in range(n_chunks): if bf16: num = q_c[:, :, c] @ state.bfloat16() den = (q_c[:, :, c] @ z_norm.bfloat16().unsqueeze(-1)).clamp(min=1.0) chunks.append(num / den) g = chunk_gate[:, :, c, None, None] state = g * state + (k_c[:, :, c].transpose(-2, -1) @ v_c[:, :, c]).float() z_norm = chunk_gate[:, :, c, None] * z_norm + k_c[:, :, c].sum(-2).float() else: num = q_c[:, :, c] @ state den = (q_c[:, :, c] @ z_norm.unsqueeze(-1)).clamp(min=1.0) chunks.append(num / den) g = chunk_gate[:, :, c, None, None] state = g * state + k_c[:, :, c].transpose(-2, -1) @ v_c[:, :, c] z_norm = chunk_gate[:, :, c, None] * z_norm + k_c[:, :, c].sum(dim=-2) return torch.cat(chunks, dim=2) @torch.compiler.disable class _GLAScanFn(torch.autograd.Function): @staticmethod def forward(ctx, q_k, k_k, v_f, chunk_gate, CS): ctx.CS = CS with torch.no_grad(): out, states, z_norms = _gla_ext.gla_scan_fwd( q_k.contiguous(), k_k.contiguous(), v_f.contiguous(), chunk_gate.contiguous(), CS, ) ctx.save_for_backward(q_k, k_k, v_f, chunk_gate, states, z_norms) return out @staticmethod def backward(ctx, grad_out): q_k, k_k, v_f, chunk_gate, states, z_norms = ctx.saved_tensors CS = ctx.CS d_q, d_k, d_v, d_gate = _gla_ext.gla_scan_backward( grad_out.float().contiguous(), q_k.contiguous(), k_k.contiguous(), v_f.contiguous(), chunk_gate.contiguous(), states, z_norms, CS, ) return (d_q, d_k, d_v, d_gate, None) @dataclass class CPUGPTConfig: vocab_size: int = 50257 seq_len: int = 1024 n_layer: int = 12 n_embd: int = 768 n_head: int = 12 fno_modes: int = 256 gla_chunk: int = 64 ffn_hidden: int = 2048 layer_pattern: str = "SSSL" bias: bool = False dropout: float = 0.0 fft_dtype: str = "fp32" use_flash_fft: bool = False gla_delta: bool = False sliding_window: int = 0 swa_window: int = 0 swa_fused_window: int = 0 attn_layer_every: int = 0 attn_full: bool = False attn_window: int = 512 landmark_layer_every: int = 0 landmark_chunk: int = 32 landmark_max: int = 64 gdn_expand_v: float = 1.0 gdn_head_dim: int = 0 def __post_init__(self): assert self.n_embd % 16 == 0, ( f"n_embd={self.n_embd} must be divisible by 16 for AMX BF16" ) assert self.ffn_hidden % 16 == 0, ( f"ffn_hidden={self.ffn_hidden} must be divisible by 16 for AMX BF16" ) def _layer_is_gla(layer_idx: int, n_layer: int, pattern: str) -> bool: if pattern == "SSSL": return layer_idx % 4 == 3 elif pattern in ("SGSG", "SLSL"): return layer_idx % 2 == 1 elif pattern == "SSLL": return layer_idx % 4 >= 2 elif pattern == "SLLL": return layer_idx % 4 >= 1 elif pattern == "GLA": return True elif pattern == "FNO": return False return False def rms_norm(x: torch.Tensor) -> torch.Tensor: return F.rms_norm(x, (x.size(-1),)) class SwiGLU(nn.Module): def __init__(self, cfg: CPUGPTConfig): super().__init__() h = cfg.ffn_hidden d = cfg.n_embd self.gate = nn.Linear(d, h, bias=cfg.bias) self.up = nn.Linear(d, h, bias=cfg.bias) self.down = nn.Linear(h, d, bias=cfg.bias) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.down(F.silu(self.gate(x)) * self.up(x)) @torch.compiler.disable class _FNOConvFn(torch.autograd.Function): @staticmethod def forward(ctx, x, filter_td, T): x_f = x.float().contiguous() f_f = filter_td.float().detach().contiguous() y, Xf, Hf = _fno_ext.fno_conv_fwd(x_f, f_f, T) ctx.save_for_backward(Xf, Hf) ctx.n_use = min(filter_td.shape[1], T) ctx.M = filter_td.shape[1] return y.to(x.dtype) @staticmethod def backward(ctx, grad_out): Xf, Hf = ctx.saved_tensors d_x, d_filter = _fno_ext.fno_conv_backward( grad_out.float().contiguous(), Xf, Hf, ctx.n_use, ctx.M ) return (d_x, d_filter, None) class FNOSeqMixer(nn.Module): def __init__(self, cfg: CPUGPTConfig): super().__init__() C = cfg.n_embd M = cfg.fno_modes self.filter_td = nn.Parameter(torch.empty(C, M)) self.out_scale = nn.Linear(C, C, bias=cfg.bias) self.register_buffer("_hf_cache", None, persistent=False) self._filter_version: int = -1 self._fft_dtype = _fft_dtype_to_torch(cfg.fft_dtype) self._use_flash_fft = cfg.use_flash_fft and _flash_fft_available self._flash_fft_conv = None self._flash_fft_T: int = -1 nn.init.normal_(self.filter_td, std=0.02) def _flash_conv(self, x: torch.Tensor, T: int, C: int, n_use: int) -> torch.Tensor: if self._flash_fft_conv is None or self._flash_fft_T != T: self._flash_fft_conv = _FlashFFTConv(2 * T, dtype=self._fft_dtype).to( x.device ) self._flash_fft_T = T h = self.filter_td.new_zeros(C, 2 * T) h[:, :n_use] = self.filter_td[:, :n_use] xp = F.pad(x.transpose(1, 2), (0, T)) y = self._flash_fft_conv(xp.to(self._fft_dtype), h.to(self._fft_dtype)) return y[:, :, :T].transpose(1, 2).to(x.dtype) def _rfft_conv(self, x: torch.Tensor, T: int, C: int, n_use: int) -> torch.Tensor: dt = self._fft_dtype h = self.filter_td.new_zeros(2 * T, C).to(dt) h[:n_use] = self.filter_td[:, :n_use].T.to(dt) xp = F.pad(x, (0, 0, 0, T)).to(dt) Xf = torch.fft.rfft(xp, dim=1) Hf = torch.fft.rfft(h, dim=0).unsqueeze(0) Xr, Xi = (Xf.real, Xf.imag) Hr, Hi = (Hf.real, Hf.imag) YHf = torch.view_as_complex( torch.stack([Xr * Hr - Xi * Hi, Xr * Hi + Xi * Hr], dim=-1).contiguous() ) return torch.fft.irfft(YHf, n=2 * T, dim=1)[:, :T] def forward(self, x: torch.Tensor) -> torch.Tensor: B, T, C = x.shape n_use = min(self.filter_td.shape[1], T) if _fno_ext is not None: y = _FNOConvFn.apply(x, self.filter_td, T) elif self._use_flash_fft and x.is_cuda: y = self._flash_conv(x, T, C, n_use) elif not self.training: y = self._forward_eval_cached(x, T, C, n_use) else: y = self._rfft_conv(x, T, C, n_use) return self.out_scale(y.to(x.dtype)) @torch.compiler.disable def _forward_eval_cached( self, x: torch.Tensor, T: int, C: int, n_use: int ) -> torch.Tensor: ver = self.filter_td._version if ( self._hf_cache is None or self._filter_version != ver or self._hf_cache.shape[0] != T + 1 ): h = self.filter_td.new_zeros(2 * T, C) h[:n_use] = self.filter_td[:, :n_use].T.detach() self._hf_cache = torch.fft.rfft(h.float(), dim=0).contiguous() self._filter_version = ver dt = self._fft_dtype xp = F.pad(x, (0, 0, 0, T)).to(dt) Xf = torch.fft.rfft(xp, dim=1) H = self._hf_cache.unsqueeze(0) Xr, Xi, Hr, Hi = (Xf.real, Xf.imag, H.real, H.imag) YHf = torch.view_as_complex( torch.stack([Xr * Hr - Xi * Hi, Xr * Hi + Xi * Hr], dim=-1).contiguous() ) return torch.fft.irfft(YHf, n=2 * T, dim=1)[:, :T] def prepare_inference(self): if hasattr(self, "_h_td"): return with torch.no_grad(): pass def _build_h_td(self, T: int) -> None: if hasattr(self, "_h_td"): return C = self.filter_td.shape[0] M = min(self.filter_td.shape[1], T) with torch.no_grad(): h_td = self.filter_td[:, :M].T.float().contiguous() self.register_buffer("_h_td", h_td.contiguous()) def step( self, x: torch.Tensor, buf: torch.Tensor, pos: int ) -> tuple[torch.Tensor, int]: M = buf.shape[1] self._build_h_td(M) slot = pos % M buf[:, slot, :] = x.detach() indices = torch.arange(M, device=x.device) ordered_idx = (slot + 1 + indices) % M ordered = buf[:, ordered_idx, :] h_flip = self._h_td.flip(0) y = (h_flip.unsqueeze(0) * ordered).sum(1) return (self.out_scale(y.to(x.dtype)), pos + 1) def forward_chunk(self, h: torch.Tensor, ctx): B, L, C = h.shape M = self.filter_td.shape[1] if ctx is None: ctx = h.new_zeros(B, M - 1, C) full = torch.cat([ctx, h], dim=1) P = full.shape[1] n_use = min(M, P) dt = self._fft_dtype hk = self.filter_td.new_zeros(2 * P, C).to(dt) hk[:n_use] = self.filter_td[:, :n_use].T.to(dt) xp = F.pad(full, (0, 0, 0, P)).to(dt) Xf = torch.fft.rfft(xp, dim=1) Hf = torch.fft.rfft(hk, dim=0).unsqueeze(0) Xr, Xi, Hr, Hi = (Xf.real, Xf.imag, Hf.real, Hf.imag) YHf = torch.view_as_complex( torch.stack([Xr * Hr - Xi * Hi, Xr * Hi + Xi * Hr], dim=-1).contiguous() ) y = torch.fft.irfft(YHf, n=2 * P, dim=1)[:, M - 1 : P] return (self.out_scale(y.to(h.dtype)), full[:, -(M - 1) :, :]) class GLAMixer(nn.Module): def __init__(self, cfg: CPUGPTConfig): super().__init__() H = cfg.n_head D = cfg.n_embd // H C = cfg.n_embd CS = cfg.gla_chunk assert C % H == 0, "n_embd must be divisible by n_head" assert cfg.seq_len % CS == 0, ( f"seq_len {cfg.seq_len} must be divisible by gla_chunk {CS}" ) self.n_head = H self.d_head = D self.chunk = CS self.gla_delta = bool(getattr(cfg, "gla_delta", False)) self.sliding_window = int(getattr(cfg, "sliding_window", 0)) self.swa_window = int(getattr(cfg, "swa_window", 0)) if self.swa_window > 0: self.swa_q = nn.Linear(C, C, bias=False) self.swa_k = nn.Linear(C, C, bias=False) self.swa_v = nn.Linear(C, C, bias=False) self.swa_o = nn.Linear(C, C, bias=False) self.swa_fused_window = int(getattr(cfg, "swa_fused_window", 0)) if self.swa_fused_window > 0: self.swa_fused_q = nn.Linear(C, C, bias=False) self.swa_fused_k = nn.Linear(C, C, bias=False) self.swa_fused_v = nn.Linear(C, C, bias=False) self.swa_fused_o = nn.Linear(C, C, bias=False) nn.init.zeros_(self.swa_fused_o.weight) self.gdn_expand_v = float(getattr(cfg, "gdn_expand_v", 1.0)) self.gdn_head_dim = int(getattr(cfg, "gdn_head_dim", 0)) or D if not (self.gla_delta and _FlaGatedDeltaNet is not None): self.q_proj = nn.Linear(C, C, bias=False) self.k_proj = nn.Linear(C, C, bias=False) self.v_proj = nn.Linear(C, C, bias=False) self.g_proj = nn.Linear(C, H, bias=True) self.out_proj = nn.Linear(C, C, bias=False) self.gla_scale = nn.Parameter(torch.ones(H) * 0.5) nn.init.constant_(self.g_proj.bias, -4.0) self._swa_mask = None if self.gla_delta and _FlaGatedDeltaNet is not None: self.gdn = _FlaGatedDeltaNet( hidden_size=C, num_heads=H, head_dim=self.gdn_head_dim, expand_v=self.gdn_expand_v, mode="chunk", ) def _python_scan( self, q_k: torch.Tensor, k_k: torch.Tensor, v: torch.Tensor, chunk_gate: torch.Tensor, ) -> torch.Tensor: B, H, T, D = q_k.shape CS = self.chunk n_chunks = T // CS q_c = q_k.view(B, H, n_chunks, CS, D) k_c = k_k.view(B, H, n_chunks, CS, D) v_c = v.view(B, H, n_chunks, CS, D) state = torch.zeros(B, H, D, D, device=q_k.device, dtype=q_k.dtype) z_norm = torch.zeros(B, H, D, device=q_k.device, dtype=q_k.dtype) out = torch.zeros(B, H, T, D, device=q_k.device, dtype=q_k.dtype) for c in range(n_chunks): num = q_c[:, :, c] @ state den = (q_c[:, :, c] @ z_norm.unsqueeze(-1)).clamp(min=1.0) out[:, :, c * CS : (c + 1) * CS] = num / den g = chunk_gate[:, :, c, None, None] state = g * state + k_c[:, :, c].transpose(-2, -1) @ v_c[:, :, c] z_norm = chunk_gate[:, :, c, None] * z_norm + k_c[:, :, c].sum(dim=-2) return out def _swa(self, q, k, v): B, H, T, D = q.shape w = self.sliding_window pad = (w - T % w) % w if pad: q, k, v = (F.pad(t, (0, 0, 0, pad)) for t in (q, k, v)) Tp = T + pad nb = Tp // w qb = q.reshape(B, H, nb, w, D) kb = k.reshape(B, H, nb, w, D) vb = v.reshape(B, H, nb, w, D) kc = torch.cat([F.pad(kb, (0, 0, 0, 0, 1, 0))[:, :, :-1], kb], dim=3) vc = torch.cat([F.pad(vb, (0, 0, 0, 0, 1, 0))[:, :, :-1], vb], dim=3) if ( self._swa_mask is None or self._swa_mask.shape[-1] != 2 * w or self._swa_mask.device != q.device ): i = torch.arange(w, device=q.device) j = torch.arange(2 * w, device=q.device) self._swa_mask = (j[None, :] < w) | (j[None, :] - w <= i[:, None]) out = F.scaled_dot_product_attention(qb, kc, vc, attn_mask=self._swa_mask) return out.reshape(B, H, Tp, D)[:, :, :T] def _swa_exact(self, x: torch.Tensor) -> torch.Tensor: B, T, C = x.shape H, D, w = (self.n_head, self.d_head, self.swa_window) q = self.swa_q(x).reshape(B, T, H, D).transpose(1, 2) k = self.swa_k(x).reshape(B, T, H, D).transpose(1, 2) v = self.swa_v(x).reshape(B, T, H, D).transpose(1, 2) i = torch.arange(T, device=x.device)[:, None] j = torch.arange(T, device=x.device)[None, :] mask = (j <= i) & (j > i - w) o = F.scaled_dot_product_attention(q, k, v, attn_mask=mask) o = o.transpose(1, 2).reshape(B, T, C) return self.swa_o(o) def _swa_fused(self, x: torch.Tensor) -> torch.Tensor: B, T, C = x.shape H, D, w = (self.n_head, self.d_head, self.swa_fused_window) q = self.swa_fused_q(x).reshape(B, T, H, D).transpose(1, 2) k = self.swa_fused_k(x).reshape(B, T, H, D).transpose(1, 2) v = self.swa_fused_v(x).reshape(B, T, H, D).transpose(1, 2) i = torch.arange(T, device=x.device)[:, None] j = torch.arange(T, device=x.device)[None, :] mask = (j <= i) & (j > i - w) o = F.scaled_dot_product_attention(q, k, v, attn_mask=mask) o = o.transpose(1, 2).reshape(B, T, C) return x + self.swa_fused_o(o) def forward(self, x: torch.Tensor) -> torch.Tensor: B, T, C = x.shape H, D, CS = (self.n_head, self.d_head, self.chunk) if self.gla_delta and hasattr(self, "gdn") and x.is_cuda: gdn_in = self._swa_fused(x) if self.swa_fused_window > 0 else x o = self.gdn(gdn_in) o = o[0] if isinstance(o, tuple) else o if self.swa_window > 0: o = o + self._swa_exact(x) return o q = self.q_proj(x).reshape(B, T, H, D).transpose(1, 2) k = self.k_proj(x).reshape(B, T, H, D).transpose(1, 2) v = self.v_proj(x).reshape(B, T, H, D).transpose(1, 2) if x.is_cuda and _fla_available: log_g = -F.softplus(self.g_proj(x)) q_t = q.transpose(1, 2).contiguous() k_t = k.transpose(1, 2).contiguous() v_t = v.transpose(1, 2).contiguous() if self.gla_delta and _fla_chunk_gdr is not None: beta = self.beta_proj(x).float() g_delta = self.a_proj(x).float() y_fla, _ = _fla_chunk_gdr( q_t, k_t, v_t, g=g_delta, beta=beta, scale=D ** (-0.5), output_final_state=False, use_beta_sigmoid_in_kernel=True, use_qk_l2norm_in_kernel=True, ) else: g_t = log_g.unsqueeze(-1).expand(B, T, H, D).contiguous().float() y_fla, _ = _fla_chunk_gla( q_t, k_t, v_t, g=g_t, scale=D ** (-0.5), output_final_state=False ) y = y_fla.to(x.dtype).transpose(1, 2).reshape(B, T, C) if self.sliding_window > 0: y = y + self._swa(q, k, v).transpose(1, 2).reshape(B, T, C) return self.out_proj(y) n_chunks = T // CS BH = B * H q_l = q.reshape(BH, n_chunks, CS, D) k_l = k.reshape(BH, n_chunks, CS, D) v_l = v.reshape(BH, n_chunks, CS, D) y_local = F.scaled_dot_product_attention(q_l, k_l, v_l, is_causal=True).reshape( B, H, T, D ) if _CPU_HAS_BF16 and q.dtype == torch.bfloat16: q_k = F.elu(q) + 1.0 k_k = F.elu(k) + 1.0 v_f = v else: q_k = F.elu(q.float()) + 1.0 k_k = F.elu(k.float()) + 1.0 v_f = v.float() log_g = -F.softplus(self.g_proj(x).float()) log_g = log_g.transpose(1, 2) chunk_log_g = log_g.reshape(B, H, n_chunks, CS).sum(-1) chunk_gate = chunk_log_g.exp() if _gla_ext is not None and q_k.dtype == torch.float32 and (not q_k.is_cuda): y_gla = _GLAScanFn.apply(q_k, k_k, v_f, chunk_gate, CS) else: y_gla = _gla_scan_py(q_k, k_k, v_f, chunk_gate, CS) scale = self.gla_scale.reshape(1, H, 1, 1).to(x.dtype) y = (y_local + scale * y_gla.to(x.dtype)).transpose(1, 2).reshape(B, T, C) return self.out_proj(y) def step( self, x: torch.Tensor, state: torch.Tensor, z_norm: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: B = x.shape[0] H, D = (self.n_head, self.d_head) q = self.q_proj(x).view(B, H, D) k = self.k_proj(x).view(B, H, D) v = self.v_proj(x).view(B, H, D) log_g = -F.softplus(self.g_proj(x).float()) scale = D ** (-0.5) if x.is_cuda and _fla_available: st = ( state if torch.is_tensor(state) and state.dtype == torch.float32 else None ) g_t = log_g.view(B, 1, H, 1).expand(B, 1, H, D).contiguous() y, state = _fla_chunk_gla( q.view(B, 1, H, D), k.view(B, 1, H, D), v.view(B, 1, H, D), g=g_t, scale=scale, initial_state=st, output_final_state=True, ) y = y.reshape(B, H * D).to(x.dtype) return (self.out_proj(y), state, z_norm) gate = log_g.exp() qf, kf, vf = (q.float(), k.float(), v.float()) if not (torch.is_tensor(state) and state.dim() == 4): state = torch.zeros(B, H, D, D, device=x.device, dtype=torch.float32) state = gate.unsqueeze(-1).unsqueeze(-1) * state + torch.einsum( "bhd,bhe->bhde", kf, vf ) y_gla = torch.einsum("bhd,bhde->bhe", qf, state) * scale return (self.out_proj(y_gla.reshape(B, H * D).to(x.dtype)), state, z_norm) def forward_chunk(self, h: torch.Tensor, state): B, L, C = h.shape H, D = (self.n_head, self.d_head) q = self.q_proj(h).view(B, L, H, D) k = self.k_proj(h).view(B, L, H, D) v = self.v_proj(h).view(B, L, H, D) log_g = -F.softplus(self.g_proj(h).float()) g = log_g.unsqueeze(-1).expand(B, L, H, D).contiguous() st = state if torch.is_tensor(state) and state.dtype == torch.float32 else None y, state = _fla_chunk_gla( q, k, v, g=g, scale=D ** (-0.5), initial_state=st, output_final_state=True ) return (self.out_proj(y.reshape(B, L, C).to(h.dtype)), state) class AttnMixer(nn.Module): def __init__(self, cfg: CPUGPTConfig): super().__init__() C, H = (cfg.n_embd, cfg.n_head) assert C % H == 0 self.n_head = H self.d_head = C // H self.full = bool(getattr(cfg, "attn_full", False)) self.window = int(getattr(cfg, "attn_window", 512)) self.q_proj = nn.Linear(C, C, bias=False) self.k_proj = nn.Linear(C, C, bias=False) self.v_proj = nn.Linear(C, C, bias=False) self.o_proj = nn.Linear(C, C, bias=False) def forward(self, x: torch.Tensor) -> torch.Tensor: B, T, C = x.shape H, D = (self.n_head, self.d_head) q = self.q_proj(x).reshape(B, T, H, D).transpose(1, 2) k = self.k_proj(x).reshape(B, T, H, D).transpose(1, 2) v = self.v_proj(x).reshape(B, T, H, D).transpose(1, 2) if self.full: o = F.scaled_dot_product_attention(q, k, v, is_causal=True) else: w = self.window i = torch.arange(T, device=x.device)[:, None] j = torch.arange(T, device=x.device)[None, :] mask = (j <= i) & (j > i - w) o = F.scaled_dot_product_attention(q, k, v, attn_mask=mask) o = o.transpose(1, 2).reshape(B, T, C) return self.o_proj(o) def _layer_is_attn(layer_idx: int, cfg: CPUGPTConfig) -> bool: k = int(getattr(cfg, "attn_layer_every", 0)) if k <= 0: return False return (layer_idx + 1) % k == 0 class LandmarkMixer(nn.Module): def __init__(self, cfg: CPUGPTConfig): super().__init__() C, H = (cfg.n_embd, cfg.n_head) assert C % H == 0 self.n_head = H self.d_head = C // H self.chunk = max(1, int(getattr(cfg, "landmark_chunk", 32))) self.max_land = max(1, int(getattr(cfg, "landmark_max", 64))) self.q_proj = nn.Linear(C, C, bias=False) self.k_proj = nn.Linear(C, C, bias=False) self.v_proj = nn.Linear(C, C, bias=False) self.o_proj = nn.Linear(C, C, bias=False) self.sink = nn.Parameter(torch.zeros(1, 1, C)) def forward(self, x: torch.Tensor) -> torch.Tensor: B, T, C = x.shape H, D = (self.n_head, self.d_head) c = max(self.chunk, (T + self.max_land - 1) // self.max_land) nc = (T + c - 1) // c pad = nc * c - T xp = F.pad(x, (0, 0, 0, pad)) land = xp.view(B, nc, c, C).mean(2) land = torch.cat([self.sink.expand(B, 1, C), land], dim=1) q = self.q_proj(x).reshape(B, T, H, D).transpose(1, 2) k = self.k_proj(land).reshape(B, nc + 1, H, D).transpose(1, 2) v = self.v_proj(land).reshape(B, nc + 1, H, D).transpose(1, 2) tok_c = (torch.arange(T, device=x.device) // c)[:, None] land_i = torch.arange(nc, device=x.device)[None, :] sink_ok = torch.ones(T, 1, dtype=torch.bool, device=x.device) chunk_ok = land_i < tok_c mask = torch.cat([sink_ok, chunk_ok], dim=1).view(1, 1, T, nc + 1) o = F.scaled_dot_product_attention(q, k, v, attn_mask=mask) o = o.transpose(1, 2).reshape(B, T, C) return self.o_proj(o) def _layer_is_landmark(layer_idx: int, cfg: CPUGPTConfig) -> bool: k = int(getattr(cfg, "landmark_layer_every", 0)) return k > 0 and (layer_idx + 1) % k == 0 class CPUGPTBlock(nn.Module): def __init__(self, cfg: CPUGPTConfig, layer_idx: int): super().__init__() is_landmark = _layer_is_landmark(layer_idx, cfg) is_attn = not is_landmark and _layer_is_attn(layer_idx, cfg) is_gla = ( not is_landmark and (not is_attn) and _layer_is_gla(layer_idx, cfg.n_layer, cfg.layer_pattern) ) if is_landmark: self.mixer = LandmarkMixer(cfg) elif is_attn: self.mixer = AttnMixer(cfg) else: self.mixer = GLAMixer(cfg) if is_gla else FNOSeqMixer(cfg) self.ffn = SwiGLU(cfg) self.ln1 = nn.RMSNorm(cfg.n_embd) self.ln2 = nn.RMSNorm(cfg.n_embd) self.is_gla = is_gla self.is_attn = is_attn self.is_landmark = is_landmark def forward(self, x: torch.Tensor) -> torch.Tensor: x = x + self.mixer(self.ln1(x)) x = x + self.ffn(self.ln2(x)) return x def step(self, x: torch.Tensor, bstate: dict) -> tuple: h = self.ln1(x) mixer = self.mixer if isinstance(mixer, (AttnMixer, LandmarkMixer)): raise NotImplementedError( "AttnMixer/LandmarkMixer recurrent step() not implemented; use forward()" ) if isinstance(mixer, GLAMixer): h_out, bstate["gla_state"], bstate["z_norm"] = mixer.step( h, bstate["gla_state"], bstate["z_norm"] ) else: assert isinstance(mixer, FNOSeqMixer) h_out, bstate["pos"] = mixer.step(h, bstate["buf"], bstate["pos"]) x = x + h_out x = x + self.ffn(self.ln2(x)) return (x, bstate) def forward_chunk(self, x: torch.Tensor, bstate: dict) -> tuple: h = self.ln1(x) mixer = self.mixer if isinstance(mixer, (AttnMixer, LandmarkMixer)): raise NotImplementedError( "AttnMixer/LandmarkMixer forward_chunk() not implemented; use forward()" ) if isinstance(mixer, GLAMixer): h_out, bstate["gla_state"] = mixer.forward_chunk(h, bstate.get("gla_state")) else: assert isinstance(mixer, FNOSeqMixer) h_out, bstate["fno_ctx"] = mixer.forward_chunk(h, bstate.get("fno_ctx")) x = x + h_out x = x + self.ffn(self.ln2(x)) return (x, bstate) class _ChunkedCEFn(torch.autograd.Function): @staticmethod def forward(ctx, x_flat, weight, targets, chunk_size): N = x_flat.shape[0] use_bf16 = x_flat.dtype == torch.bfloat16 count = (targets != -1).sum().clamp(min=1) w_g = weight.bfloat16() if use_bf16 else weight.float() loss_sum = torch.zeros((), dtype=torch.float32, device=x_flat.device) for start in range(0, N, chunk_size): end = min(start + chunk_size, N) x_c = x_flat[start:end] t_c = targets[start:end] logits_c = (x_c @ w_g.T).float() logits_c = 15.0 * torch.tanh(logits_c * (1.0 / 15.0)) loss_sum = loss_sum + F.cross_entropy( logits_c, t_c, ignore_index=-1, reduction="sum" ) ctx.save_for_backward(x_flat, weight, targets) ctx.chunk_size = chunk_size ctx.count = int(count.item()) ctx.use_bf16 = use_bf16 return loss_sum / count @staticmethod def backward(ctx, grad_out): x_flat, weight, targets = ctx.saved_tensors N, C = x_flat.shape chunk_size = ctx.chunk_size use_bf16 = ctx.use_bf16 scale = grad_out.item() / ctx.count w_g = weight.bfloat16() if use_bf16 else weight.float() dx = torch.zeros(N, C, dtype=torch.float32, device=x_flat.device) dw = torch.zeros_like(weight, dtype=torch.float32) for start in range(0, N, chunk_size): end = min(start + chunk_size, N) x_c = x_flat[start:end] t_c = targets[start:end] logits_c = (x_c @ w_g.T).float() tanh_c = torch.tanh(logits_c * (1.0 / 15.0)) logits_cap = 15.0 * tanh_c probs_c = torch.softmax(logits_cap, dim=-1) valid = t_c != -1 probs_c[~valid] = 0.0 if valid.any(): probs_c[valid, t_c[valid]] -= 1.0 d_logits_c = probs_c * (1.0 - tanh_c * tanh_c) * scale if use_bf16: dx[start:end] = (d_logits_c.bfloat16() @ w_g).float() else: dx[start:end] = d_logits_c @ w_g if use_bf16: dw.add_((d_logits_c.bfloat16().T @ x_c).float()) else: dw.addmm_(d_logits_c.T, x_c) return (dx.to(x_flat.dtype), dw, None, None) def _chunked_cross_entropy( weight: torch.Tensor, x: torch.Tensor, targets: torch.Tensor, chunk: int = 512 ) -> torch.Tensor: B, T, C = x.shape return _ChunkedCEFn.apply( x.reshape(B * T, C), weight, targets.reshape(B * T), chunk ) class CPUGPT(nn.Module): def __init__(self, cfg: CPUGPTConfig): super().__init__() self.cfg = cfg self.wte = nn.Embedding(cfg.vocab_size, cfg.n_embd) self.blocks = nn.ModuleList([CPUGPTBlock(cfg, i) for i in range(cfg.n_layer)]) self.ln_out = nn.RMSNorm(cfg.n_embd) self.lm_head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False) self.lm_head.weight = self.wte.weight self._init_weights() def _init_weights(self): nn.init.normal_(self.wte.weight, std=0.02) for block in self.blocks: if isinstance(block.mixer, GLAMixer) and hasattr(block.mixer, "q_proj"): for proj in [ block.mixer.q_proj, block.mixer.k_proj, block.mixer.v_proj, block.mixer.out_proj, ]: nn.init.normal_(proj.weight, std=0.02) if isinstance(block.mixer, AttnMixer): for proj in [ block.mixer.q_proj, block.mixer.k_proj, block.mixer.v_proj, block.mixer.o_proj, ]: nn.init.normal_(proj.weight, std=0.02) nn.init.normal_(block.ffn.gate.weight, std=0.02) nn.init.normal_(block.ffn.up.weight, std=0.02) nn.init.zeros_(block.ffn.down.weight) def param_count(self) -> int: return sum((p.numel() for p in self.parameters())) def prepare_inference(self): for block in self.blocks: if not block.is_gla and (not getattr(block, "is_attn", False)): block.mixer.prepare_inference() def init_state(self, batch_size: int = 1, device=None) -> list: if device is None: device = next(self.parameters()).device H = self.cfg.n_head D = self.cfg.n_embd // H M = self.cfg.fno_modes C = self.cfg.n_embd states = [] for block in self.blocks: if getattr(block, "is_attn", False): states.append({}) elif block.is_gla: states.append( { "gla_state": torch.zeros(batch_size, H, D, D, device=device), "z_norm": torch.zeros(batch_size, H, D, device=device), } ) else: states.append( {"buf": torch.zeros(batch_size, M, C, device=device), "pos": 0} ) return states @torch.no_grad() def step(self, idx: torch.Tensor, states: list): x = self.wte(idx.unsqueeze(1) if idx.dim() == 1 else idx) if x.dim() == 3: x = x.squeeze(1) x = F.rms_norm(x, (x.size(-1),)) new_states = [] for block, bstate in zip(self.blocks, states): x, bstate = block.step(x, bstate) new_states.append(bstate) x = self.ln_out(x) logits = self.lm_head(x.unsqueeze(1)).squeeze(1).float() logits = 15.0 * torch.tanh(logits / 15.0) return (logits, new_states) @torch.no_grad() def prefill_chunked(self, idx: torch.Tensor, chunk_size: int = 512): B, T = idx.shape states = [dict() for _ in self.blocks] last_hidden = None for c0 in range(0, T, chunk_size): x = self.wte(idx[:, c0 : c0 + chunk_size]) x = F.rms_norm(x, (x.size(-1),)) for l, block in enumerate(self.blocks): x, states[l] = block.forward_chunk(x, states[l]) last_hidden = x[:, -1:, :] logits = self.lm_head(self.ln_out(last_hidden)).float() logits = 15.0 * torch.tanh(logits / 15.0) return (logits.squeeze(1), states) @torch.no_grad() def generate( self, prompt_ids: torch.Tensor, max_new_tokens: int = 200, temperature: float = 0.8, top_k: int = 50, ) -> torch.Tensor: self.prepare_inference() device = prompt_ids.device states = self.init_state(batch_size=1, device=device) logits = torch.zeros(prompt_ids.shape[0], self.cfg.vocab_size, device=device) for i in range(prompt_ids.shape[1]): tok = prompt_ids[:, i] logits, states = self.step(tok, states) generated = [] for _ in range(max_new_tokens): if temperature == 0.0: next_tok = logits.argmax(dim=-1) else: scaled = logits / temperature if top_k > 0: topk_vals, _ = torch.topk(scaled, top_k) scaled = scaled.masked_fill( scaled < topk_vals[:, -1:], float("-inf") ) probs = torch.softmax(scaled, dim=-1) next_tok = torch.multinomial(probs, num_samples=1).squeeze(1) generated.append(next_tok) logits, states = self.step(next_tok, states) return torch.stack(generated, dim=1) def forward( self, idx: torch.Tensor, targets: Optional[torch.Tensor] = None, reduction: str = "mean", ) -> torch.Tensor: B, T = idx.shape assert T <= self.cfg.seq_len, f"Sequence {T} > max {self.cfg.seq_len}" x = self.wte(idx) x = F.rms_norm(x, (x.size(-1),)) if _CPU_HAS_BF16: with torch.autocast(device_type="cpu", dtype=torch.bfloat16): for block in self.blocks: x = block(x) else: for block in self.blocks: x = block(x) x = self.ln_out(x) if targets is None: logits = self.lm_head(x).float() logits = 15.0 * torch.tanh(logits / 15.0) return logits return _chunked_cross_entropy(self.lm_head.weight, x, targets) def gpt2_small_config(**overrides) -> CPUGPTConfig: cfg = CPUGPTConfig( vocab_size=50257, seq_len=1024, n_layer=12, n_embd=768, n_head=12, fno_modes=256, gla_chunk=64, ffn_hidden=2048, layer_pattern="SSSL", ) for k, v in overrides.items(): setattr(cfg, k, v) return cfg def smoke_config(**overrides) -> CPUGPTConfig: cfg = CPUGPTConfig( vocab_size=50257, seq_len=256, n_layer=4, n_embd=256, n_head=4, fno_modes=64, gla_chunk=64, ffn_hidden=512, layer_pattern="SSSL", ) for k, v in overrides.items(): setattr(cfg, k, v) return cfg def gpt2_8b_config() -> CPUGPTConfig: return CPUGPTConfig( n_layer=32, n_embd=4096, n_head=32, ffn_hidden=14336, fno_modes=512, gla_chunk=256, seq_len=2048, layer_pattern="SSSL", ) def gpt2_1b_config() -> CPUGPTConfig: return CPUGPTConfig( n_layer=24, n_embd=2048, n_head=16, ffn_hidden=5632, fno_modes=512, gla_chunk=256, seq_len=2048, layer_pattern="SSSL", vocab_size=50257, ) def byte125m_config(**overrides) -> CPUGPTConfig: cfg = CPUGPTConfig( vocab_size=261, seq_len=16384, n_layer=12, n_embd=1024, n_head=16, fno_modes=256, gla_chunk=512, ffn_hidden=2816, layer_pattern="SSSL", fft_dtype="fp32", use_flash_fft=True, ) for k, v in overrides.items(): setattr(cfg, k, v) return cfg def gpt2_1b_optimized_config() -> CPUGPTConfig: return CPUGPTConfig( n_layer=24, n_embd=2048, n_head=16, ffn_hidden=5632, fno_modes=512, gla_chunk=512, seq_len=2048, layer_pattern="SSSL", vocab_size=50257, fft_dtype="bf16", use_flash_fft=True, ) QWEN_CODER_VOCAB = 152064 QWEN3_4B_VOCAB = 151936 def code_3b_config(**overrides) -> CPUGPTConfig: cfg = CPUGPTConfig( vocab_size=QWEN_CODER_VOCAB, seq_len=4096, n_layer=28, n_embd=3072, n_head=24, fno_modes=512, gla_chunk=256, ffn_hidden=8192, layer_pattern="SSSL", gla_delta=True, sliding_window=0, ) for k, v in overrides.items(): setattr(cfg, k, v) return cfg def code_3b_exact_config(**overrides) -> CPUGPTConfig: cfg = CPUGPTConfig( vocab_size=QWEN_CODER_VOCAB, seq_len=4096, n_layer=28, n_embd=3584, n_head=28, fno_modes=512, gla_chunk=256, ffn_hidden=18944, layer_pattern="SSSL", gla_delta=True, sliding_window=0, ) for k, v in overrides.items(): setattr(cfg, k, v) return cfg def code_4b_config(**overrides) -> CPUGPTConfig: cfg = CPUGPTConfig( vocab_size=QWEN3_4B_VOCAB, seq_len=4096, n_layer=36, n_embd=2560, n_head=20, fno_modes=512, gla_chunk=256, ffn_hidden=9728, layer_pattern="SSSL", gla_delta=True, sliding_window=0, ) for k, v in overrides.items(): setattr(cfg, k, v) return cfg def code_1b_config(**overrides) -> CPUGPTConfig: cfg = CPUGPTConfig( vocab_size=QWEN_CODER_VOCAB, seq_len=4096, n_layer=24, n_embd=2048, n_head=16, fno_modes=512, gla_chunk=256, ffn_hidden=5632, layer_pattern="SSSL", gla_delta=True, sliding_window=0, ) for k, v in overrides.items(): setattr(cfg, k, v) return cfg def code_1p5b_exact_config(**overrides) -> CPUGPTConfig: cfg = CPUGPTConfig( vocab_size=151936, seq_len=4096, n_layer=28, n_embd=1536, n_head=12, fno_modes=512, gla_chunk=256, ffn_hidden=8960, layer_pattern="SSSL", gla_delta=True, sliding_window=0, ) for k, v in overrides.items(): setattr(cfg, k, v) return cfg def arc_a_swa512_config(**overrides) -> CPUGPTConfig: cfg = code_3b_config(swa_window=512) for k, v in overrides.items(): setattr(cfg, k, v) return cfg def arc_c_samba_k6_config(**overrides) -> CPUGPTConfig: cfg = code_3b_config(attn_layer_every=6, attn_full=False, attn_window=512) for k, v in overrides.items(): setattr(cfg, k, v) return cfg def arc_c_samba_k3_config(**overrides) -> CPUGPTConfig: cfg = code_3b_config(attn_layer_every=3, attn_full=False, attn_window=512) for k, v in overrides.items(): setattr(cfg, k, v) return cfg def arc_ub_full4_config(**overrides) -> CPUGPTConfig: cfg = code_3b_config(attn_layer_every=7, attn_full=True) for k, v in overrides.items(): setattr(cfg, k, v) return cfg def arc_d_bigstate_config(**overrides) -> CPUGPTConfig: cfg = code_3b_config(gdn_expand_v=2.0) for k, v in overrides.items(): setattr(cfg, k, v) return cfg def arc_swadelta_config(**overrides) -> CPUGPTConfig: cfg = code_3b_config(swa_fused_window=512) for k, v in overrides.items(): setattr(cfg, k, v) return cfg def code_smoke_config(**overrides) -> CPUGPTConfig: cfg = CPUGPTConfig( vocab_size=QWEN_CODER_VOCAB, seq_len=256, n_layer=4, n_embd=256, n_head=4, fno_modes=64, gla_chunk=64, ffn_hidden=512, layer_pattern="SSSL", gla_delta=False, ) for k, v in overrides.items(): setattr(cfg, k, v) return cfg def get_config(name: str) -> CPUGPTConfig: configs = { "gpt2-small": gpt2_small_config, "smoke": smoke_config, "gpt2-8b": gpt2_8b_config, "gpt2-1b": gpt2_1b_config, "gpt2-1b-optimized": gpt2_1b_optimized_config, "byte-125m": byte125m_config, "code-3b": code_3b_config, "code-3b-exact": code_3b_exact_config, "code-4b": code_4b_config, "code-1b": code_1b_config, "code-1.5b-exact": code_1p5b_exact_config, "code-smoke": code_smoke_config, "arc-a-swa512": arc_a_swa512_config, "arc-swadelta": arc_swadelta_config, "arc-c-samba-k6": arc_c_samba_k6_config, "arc-c-samba-k3": arc_c_samba_k3_config, "arc-ub-full4": arc_ub_full4_config, "arc-d-bigstate": arc_d_bigstate_config, } if name not in configs: raise ValueError(f"Unknown config '{name}'. Available: {list(configs.keys())}") return configs[name]()