model_cpu_gpt2.py

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]()