modeling_ksbyte.py

from __future__ import annotations
from typing import Optional

import torch
import torch.nn.functional as F
from transformers.modeling_utils import PreTrainedModel
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import CausalLMOutput

from .configuration_ksbyte import KsByteConfig

# ---- Inlined ksbyte.config.py -------------------------------------------------
"""Single source of truth for every tunable knob in ks_byte_lm.

Mirrors the philosophy of the ks_diacritizer DiacConfig: nothing else in the
codebase hard-codes a hyper-parameter. Override fields from the CLI / Modal /
run_local by passing a dict to `.merge(...)`.

A frozen copy of this config is embedded in every checkpoint, so a model can be
rebuilt and resumed on any machine without the original launch command.
"""

from dataclasses import asdict, dataclass, field, replace
from typing import List, Optional

# ----- byte vocabulary -------------------------------------------------------
# 0..255 raw UTF-8 bytes, then three control ids. Stored as uint16 on disk.
BYTE_VOCAB = 256
BOS_ID = 256
EOS_ID = 257
PAD_ID = 258
VOCAB_SIZE = 259


@dataclass
class ByteLMConfig:
    # ------------------------------ dataset --------------------------------
    hf_dataset: str = "Omarrran/KS-PRET-5M_5_million_kashmiri_Pretrainning_LLM_dataset_12M_tokens_2026"
    hf_revision: Optional[str] = None
    text_col: str = "auto"             # "auto" => pick the longest string column
    local_text_file: Optional[str] = None   # bypass HF: a local UTF-8 .txt path
    min_ks_ratio: float = 0.90         # drop records below this script purity
    min_chars: int = 2
    max_chars: int = 100_000           # guard against pathological mega-rows
    keep_mixed_script: bool = True     # keep the ~18% non-Nastaliq content

    # ------------------------- normalization policy ------------------------
    zwnj_policy: str = "keep"          # keep | to_space | strip
    digit_policy: str = "keep_native"  # keep_native | to_ascii
    remove_diacritics: bool = False    # NEVER true for the foundation model

    # ------------------------------ split ----------------------------------
    val_frac: float = 0.05
    test_frac: float = 0.05
    split_seed: int = 13               # deterministic, content-hash based
    dedup: bool = True                 # exact-hash dedup before splitting

    # ----------------------------- data cache ------------------------------
    data_dir: str = "data"             # where {train,val,test}.bin + meta.json live
    rebuild_data: bool = False         # force re-running data_prep

    # ------------------------------ model ----------------------------------
    variant: str = "plain"             # "plain" (implemented) | "spacebyte" (P2)
    d_model: int = 384
    n_layers: int = 6                  # used by "plain"
    n_heads: int = 6
    n_kv_heads: int = 2                # GQA; must divide n_heads
    mlp_ratio: float = 2.6667          # SwiGLU hidden = round(mlp_ratio * d_model)
    ctx_len: int = 2048                # bytes per window
    rope_theta: float = 10_000.0
    dropout: float = 0.1               # ON: 48MB overfits a transformer fast
    qk_norm: bool = True               # LayerNorm on Q,K -> tames attention logits
    tie_embeddings: bool = True
    # SpaceByte-only (reserved; "spacebyte" variant lands in P2)
    n_local_in: int = 2
    n_global: int = 6
    n_local_out: int = 2
    max_patches: int = 320

    # ---------------------------- optimisation -----------------------------
    epochs: float = 4.0                # Muennighoff: <=4 epochs ~ free on repeats
    lr: float = 4e-4
    min_lr_ratio: float = 0.1          # cosine floor = lr * min_lr_ratio
    warmup_ratio: float = 0.03
    weight_decay: float = 0.1
    beta1: float = 0.9
    beta2: float = 0.95
    grad_clip: float = 1.0
    z_loss_weight: float = 1e-4        # auxiliary logit-norm regularizer
    label_smoothing: float = 0.0
    batch_size: int = 16               # sequences per micro-step
    grad_accum: int = 4                # effective batch = batch_size * grad_accum
    doc_attention_mask: bool = True    # block cross-document attention + RoPE reset

    # --------------------------- system / perf -----------------------------
    device: str = "auto"               # auto | cuda | cpu
    bf16: bool = True                  # L4 (Ada) supports bf16
    tf32: bool = True
    torch_compile: bool = False        # optional; can be flaky -> off by default
    num_workers: int = 4
    seed: int = 42

    # ------------------------------- eval ----------------------------------
    eval_interval: int = 500           # optimiser steps between evals
    eval_iters: int = 100              # micro-batches per eval estimate
    log_interval: int = 20
    generate_every: int = 1000         # emit a sample generation every N steps
    generate_tokens: int = 160
    early_stop_patience: int = 6       # evals without val-BPB improvement
    max_steps: Optional[int] = None    # hard cap (None => derive from epochs)

    # --------------------------- io / logging ------------------------------
    output_dir: str = "outputs"
    run_name: str = "ksbyte-plain-10m"
    resume: str = "auto"               # auto | never | <path-to-checkpoint>
    save_interval: int = 500           # also saves best-on-eval regardless
    save_total_limit: int = 3
    use_wandb: bool = False
    wandb_project: str = "ks-byte-lm"

    # ------------------------------ helpers --------------------------------
    @property
    def run_dir(self) -> str:
        return f"{self.output_dir.rstrip('/')}/{self.run_name}"

    @property
    def effective_batch(self) -> int:
        return self.batch_size * self.grad_accum

    @property
    def head_dim(self) -> int:
        if self.d_model % self.n_heads != 0:
            raise ValueError(f"d_model {self.d_model} not divisible by n_heads {self.n_heads}")
        return self.d_model // self.n_heads

    @property
    def ffn_hidden(self) -> int:
        return int(round(self.mlp_ratio * self.d_model))

    def validate(self) -> "ByteLMConfig":
        """Fail fast on inconsistent settings (called by train/data entrypoints)."""
        errs: List[str] = []
        if self.n_heads % self.n_kv_heads != 0:
            errs.append(f"n_heads {self.n_heads} not divisible by n_kv_heads {self.n_kv_heads}")
        if self.d_model % self.n_heads != 0:
            errs.append(f"d_model {self.d_model} not divisible by n_heads {self.n_heads}")
        if self.variant not in ("plain", "spacebyte"):
            errs.append(f"unknown variant {self.variant!r}")
        if self.zwnj_policy not in ("keep", "to_space", "strip"):
            errs.append(f"unknown zwnj_policy {self.zwnj_policy!r}")
        if self.digit_policy not in ("keep_native", "to_ascii"):
            errs.append(f"unknown digit_policy {self.digit_policy!r}")
        if not (0.0 <= self.val_frac + self.test_frac < 0.9):
            errs.append("val_frac + test_frac must be in [0, 0.9)")
        if self.ctx_len < 8:
            errs.append("ctx_len too small")
        if self.variant == "spacebyte":
            if self.n_local_in < 0 or self.n_global < 0 or self.n_local_out < 0:
                errs.append("spacebyte layer counts must be non-negative")
            if self.n_global < 1:
                errs.append("spacebyte requires at least one global block")
            if self.max_patches < 1:
                errs.append("spacebyte max_patches must be >= 1")
        if errs:
            raise ValueError("Invalid ByteLMConfig:\n  - " + "\n  - ".join(errs))
        return self

    def merge(self, overrides: Optional[dict]) -> "ByteLMConfig":
        """Return a copy with given keys replaced; warns (not fails) on unknown keys."""
        if not overrides:
            return replace(self)
        known = {k: v for k, v in overrides.items() if k in self.__dataclass_fields__}
        unknown = set(overrides) - set(known)
        if unknown:
            print(f"[config] WARNING ignoring unknown overrides: {sorted(unknown)}")
        return replace(self, **known)

    def to_dict(self) -> dict:
        return asdict(self)

    @classmethod
    def from_dict(cls, d: dict) -> "ByteLMConfig":
        known = {k: v for k, v in d.items() if k in cls.__dataclass_fields__}
        return cls(**known)


# ---- Inlined ksbyte.layers.py -------------------------------------------------
"""Modern Llama-style building blocks, isolated so architectures swap easily.

Components: RMSNorm (pre-norm), RoPE (position ids supplied by the loader so it
resets per document), GQA attention with optional QK-Norm and document-block
attention masking, and a SwiGLU MLP.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        dtype = x.dtype
        x = x.float()
        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return (x * self.weight).to(dtype)


class RotaryEmbedding(nn.Module):
    """Builds cos/sin tables from explicit position ids ([B,T])."""

    def __init__(self, head_dim: int, theta: float = 10_000.0):
        super().__init__()
        if head_dim % 2 != 0:
            raise ValueError("head_dim must be even for RoPE")
        inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    def forward(self, pos_ids: torch.Tensor):
        # pos_ids: [B,T] -> freqs [B,T,head_dim/2]
        freqs = pos_ids.float().unsqueeze(-1) * self.inv_freq.to(pos_ids.device)
        emb = torch.cat((freqs, freqs), dim=-1)         # [B,T,head_dim]
        return emb.cos().unsqueeze(1), emb.sin().unsqueeze(1)  # [B,1,T,head_dim]


def _rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


def apply_rope(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
    q = (q * cos) + (_rotate_half(q) * sin)
    k = (k * cos) + (_rotate_half(k) * sin)
    return q, k


class Attention(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.n_heads = cfg.n_heads
        self.n_kv_heads = cfg.n_kv_heads
        self.head_dim = cfg.head_dim
        self.dropout = cfg.dropout
        self.qk_norm = cfg.qk_norm
        d = cfg.d_model
        self.wq = nn.Linear(d, self.n_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(d, self.n_kv_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(d, self.n_kv_heads * self.head_dim, bias=False)
        self.wo = nn.Linear(self.n_heads * self.head_dim, d, bias=False)
        if self.qk_norm:
            self.q_norm = RMSNorm(self.head_dim)
            self.k_norm = RMSNorm(self.head_dim)

    def forward(self, x, cos, sin, attn_mask):
        B, T, _ = x.shape
        q = self.wq(x).view(B, T, self.n_heads, self.head_dim)
        k = self.wk(x).view(B, T, self.n_kv_heads, self.head_dim)
        v = self.wv(x).view(B, T, self.n_kv_heads, self.head_dim)
        if self.qk_norm:
            q = self.q_norm(q)
            k = self.k_norm(k)
        q = q.transpose(1, 2)  # [B,H,T,D]
        k = k.transpose(1, 2)  # [B,KV,T,D]
        v = v.transpose(1, 2)
        q, k = apply_rope(q, k, cos, sin)
        # GQA: expand kv heads to match query heads.
        if self.n_kv_heads != self.n_heads:
            rep = self.n_heads // self.n_kv_heads
            k = k.repeat_interleave(rep, dim=1)
            v = v.repeat_interleave(rep, dim=1)
        dp = self.dropout if self.training else 0.0
        if attn_mask is None:
            out = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=dp)
        else:
            # boolean mask [B,1,T,T]: True => attend.
            out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=dp)
        out = out.transpose(1, 2).contiguous().view(B, T, self.n_heads * self.head_dim)
        return self.wo(out)


class SwiGLU(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        d, h = cfg.d_model, cfg.ffn_hidden
        self.w_gate = nn.Linear(d, h, bias=False)
        self.w_up = nn.Linear(d, h, bias=False)
        self.w_down = nn.Linear(h, d, bias=False)

    def forward(self, x):
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))


class Block(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.norm1 = RMSNorm(cfg.d_model)
        self.attn = Attention(cfg)
        self.norm2 = RMSNorm(cfg.d_model)
        self.mlp = SwiGLU(cfg)
        self.drop = nn.Dropout(cfg.dropout)

    def forward(self, x, cos, sin, attn_mask):
        x = x + self.drop(self.attn(self.norm1(x), cos, sin, attn_mask))
        x = x + self.drop(self.mlp(self.norm2(x)))
        return x


def build_doc_attn_mask(seg_ids: torch.Tensor) -> torch.Tensor:
    """Boolean [B,1,T,T] mask: position i may attend to j iff j<=i (causal) and
    they share a document segment. True => attend."""
    B, T = seg_ids.shape
    causal = torch.tril(torch.ones(T, T, dtype=torch.bool, device=seg_ids.device))
    same_doc = seg_ids.unsqueeze(2) == seg_ids.unsqueeze(1)  # [B,T,T]
    return (causal.unsqueeze(0) & same_doc).unsqueeze(1)


# ---- Inlined ksbyte.model.py --------------------------------------------------
"""Byte-level decoder language models.

`variant="plain"` is the P1 small Llama-style causal decoder over raw UTF-8
bytes. `variant="spacebyte"` is the P2 word-boundary hierarchy: cheap local
byte blocks run everywhere, expensive global blocks run only on patch-boundary
positions, then the latest patch context is scattered back to every byte before
optional local output blocks.

forward() returns (logits, loss, parts) where `parts` breaks the loss into the
cross-entropy term and the z-loss term for logging.
"""

import math
from typing import Dict, Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F



# ASCII whitespace bytes that delimit Kashmiri words in the normalized corpus.
# UTF-8 non-ASCII punctuation is deliberately not treated as a patch boundary.
_SPACEBYTE_WHITESPACE = (9, 10, 11, 12, 13, 32)  # \t \n \v \f \r space


def build_spacebyte_boundary_mask(
    x: torch.Tensor,
    seg_ids: Optional[torch.Tensor] = None,
    spacelike_bytes: Tuple[int, ...] = _SPACEBYTE_WHITESPACE,
) -> torch.Tensor:
    """Return [B,T] positions promoted to SpaceByte global patches.

    A text patch starts at the first byte of each document/window and at a
    spacelike byte that is not itself preceded by another spacelike byte. BOS/EOS
    ids are also promoted because packed byte streams use them as hard document
    boundaries. Marking only the first byte in a run of spaces prevents wasting
    global compute on repeated whitespace.
    """
    if x.dim() != 2:
        raise ValueError(f"x must be [B,T], got shape {tuple(x.shape)}")
    B, T = x.shape
    is_space = torch.zeros_like(x, dtype=torch.bool)
    for byte in spacelike_bytes:
        is_space |= x == byte

    prev_space = torch.zeros_like(is_space)
    prev_space[:, 1:] = is_space[:, :-1]
    boundary = is_space & ~prev_space
    boundary |= (x == BOS_ID) | (x == EOS_ID)

    # Every sampled window needs at least one patch. When segment ids are known,
    # also seed each new document segment so document-aware global attention has
    # a patch before the first visible space.
    boundary[:, 0] = True
    if seg_ids is not None:
        new_doc = torch.zeros_like(boundary)
        new_doc[:, 1:] = seg_ids[:, 1:] != seg_ids[:, :-1]
        boundary |= new_doc
    return boundary


class ByteDecoder(nn.Module):
    def __init__(self, cfg: ByteLMConfig):
        super().__init__()
        cfg.validate()
        self.cfg = cfg
        self.embed = nn.Embedding(VOCAB_SIZE, cfg.d_model)
        self.drop = nn.Dropout(cfg.dropout)
        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layers)])
        self.norm_f = RMSNorm(cfg.d_model)
        self.lm_head = nn.Linear(cfg.d_model, VOCAB_SIZE, bias=False)
        self.rope = RotaryEmbedding(cfg.head_dim, cfg.rope_theta)
        if cfg.tie_embeddings:
            self.lm_head.weight = self.embed.weight

        self.apply(self._init_weights)
        # Scale residual-path projections by 1/sqrt(2*n_layers) for stable depth.
        scale = 1.0 / math.sqrt(2 * cfg.n_layers)
        for name, p in self.named_parameters():
            if name.endswith("wo.weight") or name.endswith("w_down.weight"):
                with torch.no_grad():
                    p.mul_(scale)

    @staticmethod
    def _init_weights(m: nn.Module):
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)

    def num_params(self, non_embedding: bool = False) -> int:
        n = sum(p.numel() for p in self.parameters())
        if non_embedding and not self.cfg.tie_embeddings:
            n -= self.embed.weight.numel()
        return n

    def forward(
        self,
        x: torch.Tensor,
        y: Optional[torch.Tensor] = None,
        seg_ids: Optional[torch.Tensor] = None,
        pos_ids: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Dict[str, float]]:
        B, T = x.shape
        if pos_ids is None:
            pos_ids = torch.arange(T, device=x.device).unsqueeze(0).expand(B, T)
        cos, sin = self.rope(pos_ids)
        attn_mask = build_doc_attn_mask(seg_ids) if seg_ids is not None else None

        h = self.drop(self.embed(x))
        for blk in self.blocks:
            h = blk(h, cos, sin, attn_mask)
        h = self.norm_f(h)
        logits = self.lm_head(h)

        loss, parts = _loss_parts(logits, y, self.cfg)
        return logits, loss, parts

    @torch.no_grad()
    def generate(self, idx: torch.Tensor, max_new_tokens: int,
                 temperature: float = 1.0, top_k: Optional[int] = None,
                 eos_id: Optional[int] = None) -> torch.Tensor:
        """Autoregressive sampling (no doc-mask; single growing sequence)."""
        self.eval()
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.cfg.ctx_len:]
            logits, _, _ = self(idx_cond)
            logits = logits[:, -1, :] / max(temperature, 1e-6)
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = float("-inf")
            probs = F.softmax(logits, dim=-1)
            nxt = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, nxt), dim=1)
            if eos_id is not None and (nxt == eos_id).all():
                break
        return idx


class SpaceByteDecoder(nn.Module):
    """SpaceByte-style hierarchy over raw bytes.

    Local-in blocks process every byte cheaply. Global blocks process only patch
    boundary states (word/document boundaries), with causal/document masks over
    the patch sequence. The latest available global patch state is then scattered
    back to each byte and added as context before local-out blocks and the shared
    LM head. Setting n_local_in=n_local_out=0 and n_global=n_layers makes this
    reduce exactly to `ByteDecoder` when every position is a boundary.
    """

    def __init__(self, cfg: ByteLMConfig):
        super().__init__()
        cfg.validate()
        self.cfg = cfg
        self.embed = nn.Embedding(VOCAB_SIZE, cfg.d_model)
        self.drop = nn.Dropout(cfg.dropout)
        self.local_in_blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_local_in)])
        self.global_blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_global)])
        self.local_out_blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_local_out)])
        self.norm_f = RMSNorm(cfg.d_model)
        self.lm_head = nn.Linear(cfg.d_model, VOCAB_SIZE, bias=False)
        self.rope = RotaryEmbedding(cfg.head_dim, cfg.rope_theta)
        if cfg.tie_embeddings:
            self.lm_head.weight = self.embed.weight

        self.apply(ByteDecoder._init_weights)
        total_layers = max(1, cfg.n_local_in + cfg.n_global + cfg.n_local_out)
        scale = 1.0 / math.sqrt(2 * total_layers)
        for name, p in self.named_parameters():
            if name.endswith("wo.weight") or name.endswith("w_down.weight"):
                with torch.no_grad():
                    p.mul_(scale)

    def num_params(self, non_embedding: bool = False) -> int:
        n = sum(p.numel() for p in self.parameters())
        if non_embedding and not self.cfg.tie_embeddings:
            n -= self.embed.weight.numel()
        return n

    @staticmethod
    def gather_patches(
        h: torch.Tensor,
        boundary_mask: torch.Tensor,
        pos_ids: torch.Tensor,
        seg_ids: Optional[torch.Tensor],
        max_patches: int,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """Gather boundary states into a padded [B,P,D] patch sequence.

        Returns `(patch_h, patch_valid, patch_pos, patch_seg, patch_ids)`, where
        `patch_ids[b,t]` is the index of the latest patch at or before byte `t`.
        If a sequence has more than `max_patches` boundaries, later bytes reuse
        the last retained patch context; this keeps memory bounded and explicit.
        """
        if h.dim() != 3 or boundary_mask.dim() != 2:
            raise ValueError("h must be [B,T,D] and boundary_mask must be [B,T]")
        B, T, D = h.shape
        counts = boundary_mask.sum(dim=1).clamp_min(1)
        P = int(min(max_patches, counts.max().item()))
        device = h.device
        gather_idx = torch.zeros((B, P), dtype=torch.long, device=device)
        patch_valid = torch.zeros((B, P), dtype=torch.bool, device=device)
        patch_ids = torch.empty((B, T), dtype=torch.long, device=device)

        for b in range(B):
            idx = torch.nonzero(boundary_mask[b], as_tuple=False).flatten()
            if idx.numel() == 0:
                idx = torch.zeros(1, dtype=torch.long, device=device)
            keep = idx[:P]
            n = keep.numel()
            gather_idx[b, :n] = keep
            patch_valid[b, :n] = True
            ordinal = torch.cumsum(boundary_mask[b].to(torch.long), dim=0) - 1
            patch_ids[b] = ordinal.clamp(min=0, max=max(n - 1, 0))

        patch_h = h.gather(1, gather_idx.unsqueeze(-1).expand(B, P, D))
        patch_pos = pos_ids.gather(1, gather_idx)
        if seg_ids is None:
            patch_seg = torch.zeros((B, P), dtype=torch.long, device=device)
        else:
            patch_seg = seg_ids.gather(1, gather_idx)
        return patch_h, patch_valid, patch_pos, patch_seg, patch_ids

    @staticmethod
    def scatter_patches(patch_h: torch.Tensor, patch_ids: torch.Tensor) -> torch.Tensor:
        """Scatter latest patch states from [B,P,D] back to each byte [B,T,D]."""
        B, T = patch_ids.shape
        D = patch_h.size(-1)
        return patch_h.gather(1, patch_ids.unsqueeze(-1).expand(B, T, D))

    @staticmethod
    def _patch_attn_mask(patch_seg: torch.Tensor, patch_valid: torch.Tensor) -> torch.Tensor:
        """Causal/document patch mask that keeps padded query rows finite."""
        mask = build_doc_attn_mask(patch_seg)
        valid = patch_valid.unsqueeze(1)
        mask = mask & valid.unsqueeze(-1) & valid.unsqueeze(-2)
        eye = torch.eye(patch_seg.size(1), dtype=torch.bool, device=patch_seg.device)
        # Padded rows are ignored later, but SDPA still needs at least one True.
        mask = mask | (eye.unsqueeze(0).unsqueeze(0) & ~valid.unsqueeze(-1))
        return mask

    def forward(
        self,
        x: torch.Tensor,
        y: Optional[torch.Tensor] = None,
        seg_ids: Optional[torch.Tensor] = None,
        pos_ids: Optional[torch.Tensor] = None,
        boundary_mask: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Dict[str, float]]:
        B, T = x.shape
        if pos_ids is None:
            pos_ids = torch.arange(T, device=x.device).unsqueeze(0).expand(B, T)
        byte_cos, byte_sin = self.rope(pos_ids)
        byte_attn_mask = build_doc_attn_mask(seg_ids) if seg_ids is not None else None

        h = self.drop(self.embed(x))
        for blk in self.local_in_blocks:
            h = blk(h, byte_cos, byte_sin, byte_attn_mask)

        if boundary_mask is None:
            boundary_mask = build_spacebyte_boundary_mask(x, seg_ids=seg_ids)
        patch_h, patch_valid, patch_pos, patch_seg, patch_ids = self.gather_patches(
            h, boundary_mask, pos_ids, seg_ids, self.cfg.max_patches
        )
        patch_cos, patch_sin = self.rope(patch_pos)
        patch_attn_mask = self._patch_attn_mask(patch_seg, patch_valid)
        for blk in self.global_blocks:
            patch_h = blk(patch_h, patch_cos, patch_sin, patch_attn_mask)
        patch_context = self.scatter_patches(patch_h, patch_ids)
        if self.cfg.n_local_in == 0 and self.cfg.n_local_out == 0:
            # Degenerate hierarchy: the global stream is the whole residual
            # stream, so all-boundary inputs exactly match the plain decoder.
            h = patch_context
        else:
            h = h + patch_context

        for blk in self.local_out_blocks:
            h = blk(h, byte_cos, byte_sin, byte_attn_mask)
        h = self.norm_f(h)
        logits = self.lm_head(h)

        loss, parts = _loss_parts(logits, y, self.cfg)
        return logits, loss, parts

    @torch.no_grad()
    def generate(self, idx: torch.Tensor, max_new_tokens: int,
                 temperature: float = 1.0, top_k: Optional[int] = None,
                 eos_id: Optional[int] = None) -> torch.Tensor:
        """Autoregressive sampling with SpaceByte boundary detection per window."""
        self.eval()
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.cfg.ctx_len:]
            logits, _, _ = self(idx_cond)
            logits = logits[:, -1, :] / max(temperature, 1e-6)
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = float("-inf")
            probs = F.softmax(logits, dim=-1)
            nxt = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, nxt), dim=1)
            if eos_id is not None and (nxt == eos_id).all():
                break
        return idx


def _loss_parts(
    logits: torch.Tensor,
    y: Optional[torch.Tensor],
    cfg: ByteLMConfig,
) -> Tuple[Optional[torch.Tensor], Dict[str, float]]:
    """Shared CE + z-loss computation for both decoder variants."""
    if y is None:
        return None, {}
    ce = F.cross_entropy(
        logits.view(-1, VOCAB_SIZE), y.reshape(-1),
        label_smoothing=cfg.label_smoothing,
    )
    z = logits.logsumexp(dim=-1).pow(2).mean()
    loss = ce + cfg.z_loss_weight * z
    return loss, {"ce": ce.item(), "z": z.item(), "loss": loss.item()}


def build_model(cfg: ByteLMConfig) -> nn.Module:
    if cfg.variant == "plain":
        return ByteDecoder(cfg)
    if cfg.variant == "spacebyte":
        return SpaceByteDecoder(cfg)
    raise ValueError(f"unknown variant {cfg.variant!r}")



def _to_byte_lm_config(config: KsByteConfig) -> ByteLMConfig:
    keys = ByteLMConfig.__dataclass_fields__.keys()
    return ByteLMConfig.from_dict({k: getattr(config, k) for k in keys if hasattr(config, k)})


class KsByteForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = KsByteConfig
    base_model_prefix = "model"
    main_input_name = "input_ids"
    supports_gradient_checkpointing = False
    _supports_cache_class = False
    _tied_weights_keys = ["model.lm_head.weight"]
    # Some Transformers releases expect this to be a dict and call .keys();
    # others only need membership. Keep it dict-shaped for Colab compatibility.
    all_tied_weights_keys = {"model.lm_head.weight": "model.embed.weight"}

    def __init__(self, config: KsByteConfig):
        super().__init__(config)
        self.model = build_model(_to_byte_lm_config(config))

    def get_input_embeddings(self):
        return self.model.embed

    def set_input_embeddings(self, value):
        self.model.embed = value
        if getattr(self.config, "tie_embeddings", False):
            self.model.lm_head.weight = self.model.embed.weight

    def get_output_embeddings(self):
        return self.model.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.model.lm_head = new_embeddings

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        if input_ids.shape[1] > self.config.ctx_len:
            input_ids = input_ids[:, -self.config.ctx_len:]
        return {"input_ids": input_ids}

    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.Tensor] = None,
        labels: Optional[torch.LongTensor] = None,
        **kwargs,
    ):
        logits, _, _ = self.model(input_ids)
        loss = None
        if labels is not None:
            shift_logits = logits[:, :-1, :].contiguous()
            shift_labels = labels[:, 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
                ignore_index=-100,
            )
        return CausalLMOutput(loss=loss, logits=logits)

    @torch.no_grad()
    def generate_bytes(self, input_ids, max_new_tokens=200, temperature=0.8, top_k=50):
        return self.model.generate(
            input_ids,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_k=top_k,
            eos_id=self.config.eos_token_id,
        )