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


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

    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 == "SGSG":
        return (layer_idx % 2) == 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)
        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 = torch.zeros(T, C, dtype=torch.float32)
            h_td[:M] = self.filter_td[:, :M].T.float()

        self.register_buffer("_h_td", h_td.contiguous())

    def step(
        self, x: torch.Tensor, buf: torch.Tensor, pos: int
    ) -> tuple[torch.Tensor, int]:
        T = buf.shape[1]
        self._build_h_td(T)
        slot = pos % T
        buf[:, slot, :] = x.detach()

        indices = torch.arange(T, device=x.device)
        ordered_idx = (slot + 1 + indices) % T
        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


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.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)

    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 forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, C = x.shape
        H, D, CS = self.n_head, self.d_head, self.chunk

        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)

        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:
            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())
        gate = log_g.exp()

        q_k = F.elu(q.float()) + 1.0
        k_k = F.elu(k.float()) + 1.0
        v_f = v.float()

        g4 = gate.unsqueeze(-1).unsqueeze(-1)
        state = g4 * state + torch.einsum("bhd,bhe->bhde", k_k, v_f)

        g3 = gate.unsqueeze(-1)
        z_norm = g3 * z_norm + k_k

        num = torch.einsum("bhd,bhde->bhe", q_k, state)
        den = torch.einsum("bhd,bhd->bh", q_k, z_norm).clamp_(min=1.0).unsqueeze(-1)
        y_gla = (num / den).to(x.dtype)

        scale = self.gla_scale.view(1, H, 1)
        y = (y_gla * scale).reshape(B, H * D)

        return self.out_proj(y), state, z_norm


class CPUGPTBlock(nn.Module):
    def __init__(self, cfg: CPUGPTConfig, layer_idx: int):
        super().__init__()
        is_gla = _layer_is_gla(layer_idx, cfg.n_layer, cfg.layer_pattern)
        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

    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, 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


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):
                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)

            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:
                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
        T = self.cfg.seq_len
        C = self.cfg.n_embd
        states = []

        for block in self.blocks:
            if 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, T, 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 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 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,
    )


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,
    }
    if name not in configs:
        raise ValueError(f"Unknown config '{name}'. Available: {list(configs.keys())}")

    return configs[name]()