Upload folder using huggingface_hub

Dockerfile

@@ -0,0 +1,24 @@
+FROM python:3.11-slim
+
+WORKDIR /app
+
+ENV PYTHONDONTWRITEBYTECODE=1 \
+    PYTHONUNBUFFERED=1 \
+    PYTHONPATH=/app/src \
+    HOST=0.0.0.0 \
+    PORT=7860
+
+RUN pip install --no-cache-dir \
+    torch --index-url https://download.pytorch.org/whl/cpu && \
+    pip install --no-cache-dir numpy pyyaml tiktoken
+
+COPY pyproject.toml README.md ./
+COPY configs/ ./configs/
+COPY src/ ./src/
+COPY data/tinystories/tiktoken_compact_4096.json ./data/tinystories/tiktoken_compact_4096.json
+COPY data/tinystories/val_sample_256.txt ./data/tinystories/val_sample_256.txt
+COPY runs/tinystories/benchmark_24h.json ./runs/tinystories/benchmark_24h.json
+
+EXPOSE 7860
+
+CMD ["python", "-m", "conveyer.web_sim", "--host", "0.0.0.0", "--port", "7860"]

configs/synthetic.yaml

@@ -0,0 +1,53 @@
+model:
+  vocab_size: 128
+  d_model: 256
+  n_heads: 4
+  n_layers: 6
+  d_ff: 688
+  max_seq_len: 512
+  dropout: 0.0
+  belt_size: 64
+  tau_min: 0.5
+  tau_max: 4.0
+  tau_gamma: 1.5
+  gate_bias_init: 1.0
+  spawn_noise_std: 0.5
+  thought_noise_std: 0.25
+  use_thought_highway: true
+  causal_belt_only: false
+
+train:
+  out_dir: "runs/synthetic"
+  seed: 42
+  batch_size: 64
+  lr: 0.0005
+  weight_decay: 0.01
+  grad_clip: 1.0
+  max_steps: 1200
+  warmup_steps: 80
+  eval_every: 100
+  log_every: 25
+  unroll_steps: 16
+  seed_sigma_max: 3.5
+  seed_sigma_eta: 1.25
+  loss_w_min: 0.05
+  loss_w_beta: 1.5
+  entropy_reg: 0.001
+  no_belt_bptt: false
+  no_thought_highway: false
+  ar_baseline: false
+
+infer:
+  prompt_overlap: 16
+  bake_steps: 64
+  temperature: 0.0
+  max_new_tokens: 64
+
+data:
+  dataset: "synthetic"
+  task: "all"
+  seq_len: 128
+  num_train: 24000
+  num_val: 1024
+  data_dir: "data/synthetic"
+  vocab_size: 128

pyproject.toml

@@ -0,0 +1,37 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "diffusion-conveyer-llm"
+version = "0.1.0"
+description = "64-Slot Dual-Channel Conveyor-Belt Diffusion Language Model"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "torch>=2.4.0",
+    "numpy>=1.26.0",
+    "pyyaml>=6.0",
+    "tqdm>=4.66.0",
+    "tokenizers>=0.19.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=8.0.0",
+]
+
+[project.scripts]
+conveyer-train = "conveyer.train:main"
+conveyer-eval = "conveyer.eval:main"
+conveyer-infer = "conveyer.infer:main"
+conveyer-sweep = "conveyer.sweep:main"
+conveyer-live = "conveyer.live:main"
+conveyer-web = "conveyer.web_sim:main"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+pythonpath = ["src"]

src/conveyer/__init__.py

@@ -0,0 +1,4 @@
+from conveyer.config import ConveyerConfig, load_config
+from conveyer.model import ConveyerDiffusionLM
+
+__all__ = ["ConveyerConfig", "load_config", "ConveyerDiffusionLM"]

src/conveyer/checkpoint.py

@@ -0,0 +1,151 @@
+from __future__ import annotations
+
+from dataclasses import asdict
+import os
+from pathlib import Path
+from typing import Any, Optional
+import urllib.request
+import torch
+
+from conveyer.config import ConveyerConfig, _merge_dataclass
+from conveyer.model import ConveyerDiffusionLM
+
+HF_DEFAULT_CKPT_URL = "https://huggingface.co/sumitk/diffusion-conveyer-llm/resolve/main/best.pt"
+
+
+def _move_to_cpu(obj: Any) -> Any:
+    if isinstance(obj, torch.Tensor):
+        return obj.detach().to(device="cpu", non_blocking=False)
+    if isinstance(obj, dict):
+        return {k: _move_to_cpu(v) for k, v in obj.items()}
+    if isinstance(obj, list):
+        return [_move_to_cpu(v) for v in obj]
+    if isinstance(obj, tuple):
+        return tuple(_move_to_cpu(v) for v in obj)
+    return obj
+
+
+def _atomic_write_payload(path: Path, payload: dict[str, Any]) -> None:
+    """
+    Crash-safe atomic checkpoint writer on Windows NTFS:
+      1. If `path` already exists, rotates `path -> path.prev` BEFORE writing so an abrupt BSOD/power loss
+         never leaves both the current and previous checkpoint with un-flushed NTFS data blocks.
+      2. Writes `payload` to `path.tmp`, calls `f.flush()` + `os.fsync(f.fileno())` to force NTFS data blocks
+         to physical NVMe storage BEFORE updating the directory metadata via `os.replace`.
+    """
+    path.parent.mkdir(parents=True, exist_ok=True)
+    tmp_p = path.with_suffix(path.suffix + ".tmp")
+    prev_p = path.with_suffix(path.suffix + ".prev")
+
+    with open(tmp_p, "wb") as f:
+        torch.save(payload, f)
+        f.flush()
+        os.fsync(f.fileno())
+
+    if path.exists():
+        try:
+            os.replace(path, prev_p)
+        except OSError:
+            pass
+
+    os.replace(tmp_p, path)
+
+
+def save_checkpoint(
+    path: str | Path,
+    model: ConveyerDiffusionLM,
+    cfg: ConveyerConfig,
+    step: int,
+    optimizer: Optional[torch.optim.Optimizer] = None,
+    metrics: Optional[dict[str, Any]] = None,
+    best_val_loss: Optional[float] = None,
+    elapsed_sec: float = 0.0,
+) -> None:
+    """
+    Copies `model` and `optimizer` state dicts to CPU RAM first (avoiding 350 MB of synchronous
+    WDDM GPU->CPU DMA transfers inside `torch.save` that can trigger `dxgmms2.sys` BSODs),
+    then writes atomically with `os.fsync` and `.prev` backup rotation.
+    """
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+
+    cpu_model_state = _move_to_cpu(model.state_dict())
+    payload: dict[str, Any] = {
+        "model_state": cpu_model_state,
+        "config": asdict(cfg),
+        "step": step,
+        "metrics": metrics or {},
+        "best_val_loss": best_val_loss,
+        "elapsed_sec": elapsed_sec,
+    }
+    if optimizer is not None:
+        payload["optimizer_state"] = _move_to_cpu(optimizer.state_dict())
+
+    _atomic_write_payload(Path(path), payload)
+
+
+def load_checkpoint_payload(path: str | Path, map_location: Any = "cpu") -> tuple[dict[str, Any], Path]:
+    """
+    Loads checkpoint dict from `path`, automatically falling back to `path.prev` or the latest
+    valid `ckpt_step_*.pt` in the same directory if `path` was corrupted by a power loss or BSOD.
+    """
+    p = Path(path)
+    candidates: list[Path] = [p, p.with_suffix(p.suffix + ".prev")]
+    if p.parent.exists():
+        milestones = sorted(p.parent.glob("ckpt_step_*.pt"), key=lambda x: x.stat().st_mtime, reverse=True)
+        for m in milestones:
+            if m not in candidates:
+                candidates.append(m)
+            m_prev = m.with_suffix(m.suffix + ".prev")
+            if m_prev.exists() and m_prev not in candidates:
+                candidates.append(m_prev)
+
+    last_err: Optional[Exception] = None
+    for cand in candidates:
+        if not cand.exists() or cand.stat().st_size == 0:
+            continue
+        try:
+            data = torch.load(cand, map_location=map_location, weights_only=False)
+            if isinstance(data, dict) and "model_state" in data:
+                return data, cand
+        except Exception as e:
+            last_err = e
+            print(f"[Checkpoint] Warning: skipping unreadable checkpoint {cand.name} ({e})", flush=True)
+
+    if p.name == "best.pt" and "tinystories" in str(p).replace("\\", "/"):
+        try:
+            p.parent.mkdir(parents=True, exist_ok=True)
+            tmp_dl = p.with_suffix(".pt.dl")
+            print(f"[Checkpoint] Downloading pretrained checkpoint from {HF_DEFAULT_CKPT_URL} ...", flush=True)
+            req = urllib.request.Request(HF_DEFAULT_CKPT_URL, headers={"User-Agent": "Mozilla/5.0"})
+            with urllib.request.urlopen(req, timeout=120) as resp, open(tmp_dl, "wb") as f:
+                while True:
+                    chunk = resp.read(4 * 1024 * 1024)
+                    if not chunk:
+                        break
+                    f.write(chunk)
+            os.replace(tmp_dl, p)
+            data = torch.load(p, map_location=map_location, weights_only=False)
+            if isinstance(data, dict) and "model_state" in data:
+                return data, p
+        except Exception as e:
+            last_err = e
+
+    raise RuntimeError(f"No valid checkpoint found for {p} (last error: {last_err})")
+
+
+def load_checkpoint(
+    path: str | Path,
+    device: Optional[torch.device] = None,
+) -> tuple[ConveyerDiffusionLM, ConveyerConfig, dict[str, Any]]:
+    if device is None:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    ckpt, used_path = load_checkpoint_payload(path, map_location=device)
+    if used_path != Path(path):
+        print(f"[Checkpoint] Recovered from fallback checkpoint {used_path}", flush=True)
+    cfg = ConveyerConfig()
+    if "config" in ckpt:
+        _merge_dataclass(cfg, ckpt["config"])
+    model = ConveyerDiffusionLM(cfg.model).to(device)
+    model.load_state_dict(ckpt["model_state"])
+    return model, cfg, ckpt

src/conveyer/config.py

@@ -0,0 +1,222 @@
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, field, fields, is_dataclass
+import json
+from pathlib import Path
+from typing import Any
+
+try:
+    import yaml  # type: ignore
+except ImportError:
+    yaml = None
+
+
+@dataclass
+class ModelConfig:
+    vocab_size: int = 128
+    d_model: int = 256
+    n_heads: int = 4
+    n_layers: int = 6
+    d_ff: int = 688
+    max_seq_len: int = 512
+    dropout: float = 0.0
+    rope_theta: float = 10000.0
+    # Conveyor belt parameters
+    belt_size: int = 64
+    front_conveyer_slots: int = 16  # Slots 0..15 shift left by 1 every tick; Slots 16..63 (75%) are Elastic Planning Zone
+    plan_horizon_stride: int = 2    # Back 75% planning slots span 2x longer horizon while sentences expand in middle
+    tau_min: float = 0.5
+    tau_max: float = 4.0
+    tau_gamma: float = 1.5
+    gate_bias_init: float = 1.0
+    spawn_noise_std: float = 0.5
+    thought_noise_std: float = 0.25
+    # Ablation & efficiency switches
+    use_thought_highway: bool = True
+    use_elastic_planning: bool = True  # Enables 75% paused/elastic planning belt with mid-belt expansion
+    causal_belt_only: bool = False     # If True, disables bidirectional belt attention (AR baseline mode)
+    grad_checkpointing: bool = False   # Enables activation checkpointing for long unroll horizons (U=24..32)
+
+
+@dataclass
+class TrainConfig:
+    out_dir: str = "runs/synthetic"
+    seed: int = 42
+    batch_size: int = 64
+    lr: float = 5e-4
+    weight_decay: float = 0.01
+    grad_clip: float = 1.0
+    max_steps: int = 1200
+    max_hours: float = 0.0           # Optional wall-clock training budget in hours (0.0 = use max_steps only)
+    warmup_steps: int = 80
+    eval_every: int = 100
+    save_every: int = 500            # Periodic crash-safe checkpoint interval (steps)
+    log_every: int = 25
+    unroll_steps: int = 16
+    prefix_lm_weight: float = 1.0    # Weight on causal next-token LM loss across the prompt/sequence
+    step_cooldown_ms: int = 0        # Optional GPU thermal pacing sleep per step in ms
+    max_gpu_temp_c: int = 80         # Automatic thermal throttle threshold in Celsius
+    # Corrupted-belt seeding & recovery parameters
+    seed_sigma_max: float = 3.5
+    seed_sigma_eta: float = 1.25
+    recovery_prob: float = 0.35      # DAgger / self-conditioned prefix + belt recovery probability
+    belt_corrupt_prob: float = 0.20  # Mid-unroll belt perturbation rate to train active belt recovery
+    # Horizon loss weights: w_k = w_min + (1 - w_min) * (1 - k / (B - 1))^w_beta
+    loss_w_min: float = 0.05
+    loss_w_beta: float = 1.5
+    entropy_reg: float = 0.001
+    # Ablations
+    no_belt_bptt: bool = False
+    no_thought_highway: bool = False
+    ar_baseline: bool = False
+    overfit_batches: int = 0
+
+
+@dataclass
+class InferConfig:
+    prompt_overlap: int = 12
+    bake_steps: int = 64
+    temperature: float = 0.0  # 0.0 = greedy argmax at slot 0
+    max_new_tokens: int = 64
+
+
+@dataclass
+class DataConfig:
+    dataset: str = "synthetic"  # "synthetic" or "tinystories"
+    task: str = "all"           # "reverse", "path", "constraint", "all", "lm"
+    seq_len: int = 192
+    num_train: int = 24000
+    num_val: int = 1024
+    data_dir: str = "data/synthetic"
+    vocab_size: int = 128
+
+
+@dataclass
+class ConveyerConfig:
+    model: ModelConfig = field(default_factory=ModelConfig)
+    train: TrainConfig = field(default_factory=TrainConfig)
+    infer: InferConfig = field(default_factory=InferConfig)
+    data: DataConfig = field(default_factory=DataConfig)
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+def _parse_scalar(val_str: str) -> Any:
+    s = val_str.strip()
+    if not s:
+        return ""
+    if (s.startswith('"') and s.endswith('"')) or (s.startswith("'") and s.endswith("'")):
+        return s[1:-1]
+    low = s.lower()
+    if low == "true":
+        return True
+    if low == "false":
+        return False
+    if low in ("null", "none"):
+        return None
+    try:
+        if "." in s or "e" in low:
+            return float(s)
+        return int(s)
+    except ValueError:
+        return s
+
+
+def _simple_yaml_parse(text: str) -> dict[str, Any]:
+    """Minimal dependency-free YAML parser for 2-level config files."""
+    if yaml is not None:
+        return yaml.safe_load(text) or {}
+    text_stripped = text.strip()
+    if text_stripped.startswith("{"):
+        return json.loads(text_stripped)
+    root: dict[str, Any] = {}
+    current_section: dict[str, Any] | None = None
+    for raw_line in text.splitlines():
+        line = raw_line.split("#", 1)[0].rstrip()
+        if not line.strip():
+            continue
+        indent = len(line) - len(line.lstrip(" "))
+        stripped = line.strip()
+        if ":" not in stripped:
+            continue
+        key, val_part = stripped.split(":", 1)
+        key = key.strip()
+        val_part = val_part.strip()
+        if indent == 0:
+            if val_part == "":
+                current_section = {}
+                root[key] = current_section
+            else:
+                root[key] = _parse_scalar(val_part)
+                current_section = None
+        else:
+            if current_section is None:
+                raise ValueError(f"Unexpected indentation in YAML line: {raw_line}")
+            current_section[key] = _parse_scalar(val_part)
+    return root
+
+
+def _merge_dataclass(dc_instance: Any, updates: dict[str, Any]) -> None:
+    valid_fields = {f.name: f for f in fields(dc_instance)}
+    for k, v in updates.items():
+        if k not in valid_fields:
+            continue
+        current_val = getattr(dc_instance, k)
+        if is_dataclass(current_val) and isinstance(v, dict):
+            _merge_dataclass(current_val, v)
+        else:
+            setattr(dc_instance, k, v)
+
+
+def _cast_value(raw: str, target_type: Any) -> Any:
+    if target_type is bool or raw.lower() in ("true", "false"):
+        return raw.lower() in ("true", "1", "yes")
+    val = _parse_scalar(raw)
+    if target_type is float and isinstance(val, int):
+        return float(val)
+    if target_type is int and isinstance(val, float):
+        return int(val)
+    if target_type is str:
+        return str(val)
+    return val
+
+
+def apply_overrides(cfg: ConveyerConfig, overrides: list[str]) -> ConveyerConfig:
+    for item in overrides:
+        if "=" not in item:
+            raise ValueError(f"Override must be in key=value format, got: {item}")
+        key_path, raw_val = item.split("=", 1)
+        parts = key_path.strip().split(".")
+        obj: Any = cfg
+        for p in parts[:-1]:
+            if not hasattr(obj, p):
+                raise KeyError(f"Unknown config section '{p}' in '{key_path}'")
+            obj = getattr(obj, p)
+        leaf = parts[-1]
+        if not hasattr(obj, leaf):
+            raise KeyError(f"Unknown config key '{leaf}' in '{key_path}'")
+        cur = getattr(obj, leaf)
+        val = _cast_value(raw_val.strip(), type(cur))
+        setattr(obj, leaf, val)
+    _sync_config(cfg)
+    return cfg
+
+
+def _sync_config(cfg: ConveyerConfig) -> None:
+    cfg.model.vocab_size = cfg.data.vocab_size
+    if cfg.train.no_thought_highway:
+        cfg.model.use_thought_highway = False
+    if cfg.train.ar_baseline:
+        cfg.model.causal_belt_only = True
+
+
+def load_config(path: str | Path | None = None, overrides: list[str] | None = None) -> ConveyerConfig:
+    cfg = ConveyerConfig()
+    if path is not None:
+        raw = _simple_yaml_parse(Path(path).read_text(encoding="utf-8"))
+        _merge_dataclass(cfg, raw)
+    if overrides:
+        apply_overrides(cfg, overrides)
+    _sync_config(cfg)
+    return cfg

src/conveyer/data/__init__.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from typing import Any, Protocol
+import torch
+from torch.utils.data import Dataset
+
+from conveyer.config import ConveyerConfig
+
+
+class TokenizerProtocol(Protocol):
+    vocab_size: int
+    pad_id: int
+    bos_id: int
+    eos_id: int
+
+    def encode(self, text: str) -> list[int]: ...
+    def decode(self, ids: list[int]) -> str: ...
+    def id_to_piece(self, token_id: int) -> str: ...
+
+
+def build_dataset_and_tokenizer(
+    cfg: ConveyerConfig,
+    infer_only: bool = False,
+) -> tuple[Dataset[dict[str, Any]], Dataset[dict[str, Any]], TokenizerProtocol]:
+    if cfg.data.dataset == "synthetic":
+        from conveyer.data.synthetic import build_synthetic_datasets
+
+        return build_synthetic_datasets(cfg)
+    elif cfg.data.dataset == "tinystories":
+        from conveyer.data.tinystories import build_tinystories_datasets
+
+        return build_tinystories_datasets(cfg, infer_only=infer_only)
+    else:
+        raise ValueError(f"Unknown dataset '{cfg.data.dataset}'. Choose 'synthetic' or 'tinystories'.")
+
+
+def collate_conveyer_batch(batch: list[dict[str, Any]], pad_id: int = 0) -> dict[str, torch.Tensor]:
+    max_len = max(item["input_ids"].size(0) for item in batch)
+    bsz = len(batch)
+    input_ids = torch.full((bsz, max_len), pad_id, dtype=torch.long)
+    prompt_lengths = torch.zeros(bsz, dtype=torch.long)
+    task_ids = torch.zeros(bsz, dtype=torch.long)
+
+    for i, item in enumerate(batch):
+        seq = item["input_ids"]
+        input_ids[i, : seq.size(0)] = seq
+        prompt_lengths[i] = int(item.get("prompt_len", 1))
+        task_ids[i] = int(item.get("task_id", 0))
+
+    return {
+        "input_ids": input_ids,
+        "prompt_lengths": prompt_lengths,
+        "task_ids": task_ids,
+    }

src/conveyer/data/synthetic.py

@@ -0,0 +1,276 @@
+from __future__ import annotations
+
+from typing import Any
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+from conveyer.config import ConveyerConfig
+
+
+class SyntheticTokenizer:
+    """
+    Compact deterministic symbolic + word-piece vocabulary (128 tokens) for synthetic lookahead planning tasks.
+    """
+
+    def __init__(self, vocab_size: int = 128) -> None:
+        self.vocab_size = vocab_size
+        self.pad_id = 0
+        self.bos_id = 1
+        self.eos_id = 2
+        self.sep_id = 3  # '=>' separator between prompt and output
+
+        special = ["<PAD>", "<BOS>", "<EOS>", "=>", "|", ":", ",", ";", "(", ")", "[", "]", "->", "+", "=", "?"]
+        task_tags = ["[PATH]", "[REV]", "[PLAN]", "[STORY]", "START", "GOAL", "EDGES", "SEQ", "SUM", "END", "STEP", "OK"]
+        nodes = [f"N{i}" for i in range(32)]
+        numbers = [f"{i}" for i in range(32)]
+        words = [
+            "the", "little", "robot", "walked", "across", "bridge", "to", "find", "golden", "key",
+            "locked", "door", "opened", "quietly", "inside", "shiny", "crystal", "glowed", "bright", "blue",
+            "cat", "chased", "mouse", "under", "wooden", "table", "until", "both", "fell", "asleep",
+            "wise", "owl", "watched", "stars", "shine", "above", "dark", "forest",
+        ]
+        all_pieces = special + task_tags + nodes + numbers + words
+        while len(all_pieces) < vocab_size:
+            all_pieces.append(f"TOK_{len(all_pieces)}")
+        self.pieces = all_pieces[:vocab_size]
+        self.piece_to_id = {p: i for i, p in enumerate(self.pieces)}
+
+    def id_to_piece(self, token_id: int) -> str:
+        if 0 <= token_id < len(self.pieces):
+            return self.pieces[token_id]
+        return f"<{token_id}>"
+
+    def encode(self, text: str) -> list[int]:
+        tokens: list[int] = [self.bos_id]
+        for raw in text.strip().split():
+            if raw in self.piece_to_id:
+                tokens.append(self.piece_to_id[raw])
+            else:
+                # Fallback hash into word range
+                idx = 64 + (hash(raw) % max(1, self.vocab_size - 64))
+                tokens.append(idx)
+        return tokens
+
+    def decode(self, ids: list[int]) -> str:
+        out: list[str] = []
+        for idx in ids:
+            if idx in (self.pad_id, self.bos_id):
+                continue
+            if idx == self.eos_id:
+                break
+            out.append(self.id_to_piece(idx))
+        return " ".join(out)
+
+
+def _make_path_planning_sample(
+    rng: np.random.Generator,
+    tok: SyntheticTokenizer,
+    seq_len: int,
+) -> tuple[torch.Tensor, int, int]:
+    """
+    Task 0 (Lookahead Graph Pathfinding - `[PATH]`):
+    A fork graph where START branches into two paths of length 6:
+      - One branch leads to a DEAD_END.
+      - The other branch leads to the GOAL node.
+    The edges are listed in random order, and the target trajectory is:
+      `=> N_start N_step1 N_step2 N_step3 N_step4 N_step5 N_goal <EOS> ...`
+    Crucially, we also require the model to emit the GOAL-conditioned path in reverse-verified
+    or lookahead form where choosing `N_step1` at Position 0 requires chaining backwards from `GOAL`
+    across the 6-hop chain. On the 64-slot conveyor belt, `N_goal` at the end of the path is anchored
+    while `N_step1 .. N_step5` crystallize simultaneously!
+    """
+    node_pool = rng.choice(32, size=12, replace=False)
+    start_node = int(node_pool[0])
+    good_chain = [int(x) for x in node_pool[1:6]]  # 5 intermediate hops
+    goal_node = int(node_pool[6])
+    bad_chain = [int(x) for x in node_pool[7:12]]  # 5 dead-end hops
+
+    edges: list[tuple[int, int]] = []
+    # Branch A (reaches goal_node)
+    edges.append((start_node, good_chain[0]))
+    for i in range(len(good_chain) - 1):
+        edges.append((good_chain[i], good_chain[i + 1]))
+    edges.append((good_chain[-1], goal_node))
+
+    # Branch B (dead-end branch from start_node)
+    edges.append((start_node, bad_chain[0]))
+    for i in range(len(bad_chain) - 1):
+        edges.append((bad_chain[i], bad_chain[i + 1]))
+
+    # Shuffle edge order so local greedy heuristics fail unless the model traces backward from GOAL
+    perm = rng.permutation(len(edges))
+    shuffled_edges = [edges[i] for i in perm]
+
+    prompt_pieces = ["<BOS>", "[PATH]", "START", f"N{start_node}", "GOAL", f"N{goal_node}", "EDGES"]
+    for u, v in shuffled_edges:
+        prompt_pieces.extend([f"N{u}", "->", f"N{v}", ";"])
+    prompt_pieces.append("=>")
+
+    full_path = [start_node] + good_chain + [goal_node]
+    # Emit the path twice separated by '|' and ending with GOAL confirmation so the full 64-slot belt is active
+    target_pieces: list[str] = []
+    for n in full_path:
+        target_pieces.append(f"N{n}")
+    target_pieces.append("|")
+    for n in reversed(full_path):
+        target_pieces.append(f"N{n}")
+    target_pieces.extend(["OK", f"N{goal_node}", "<EOS>"])
+
+    prompt_ids = [tok.piece_to_id[p] for p in prompt_pieces]
+    target_ids = [tok.piece_to_id[p] for p in target_pieces]
+    full_ids = prompt_ids + target_ids
+
+    out = torch.full((seq_len,), tok.pad_id, dtype=torch.long)
+    L = min(seq_len, len(full_ids))
+    out[:L] = torch.tensor(full_ids[:L], dtype=torch.long)
+    return out, len(prompt_ids), 0
+
+
+def _make_reverse_plan_sample(
+    rng: np.random.Generator,
+    tok: SyntheticTokenizer,
+    seq_len: int,
+) -> tuple[torch.Tensor, int, int]:
+    """
+    Task 1 (Delayed-Constraint / Reverse Block Planning - `[REV]`):
+    The prompt gives a sequence of 16 numbers `a_0 .. a_15`.
+    The output must emit pairs `(a_i + a_{15-i}) mod 32` followed by the reversed sequence `a_15 .. a_0`.
+    """
+    length = 12
+    nums = [int(x) for x in rng.integers(0, 16, size=length)]
+    prompt_pieces = ["<BOS>", "[REV]", "SEQ"] + [f"{x}" for x in nums] + ["=>"]
+
+    sums = [(nums[i] + nums[length - 1 - i]) % 32 for i in range(length)]
+    rev = list(reversed(nums))
+    target_pieces = (
+        ["SUM"]
+        + [f"{s}" for s in sums]
+        + ["|"]
+        + [f"{r}" for r in rev]
+        + ["OK", "<EOS>"]
+    )
+
+    prompt_ids = [tok.piece_to_id[p] for p in prompt_pieces]
+    target_ids = [tok.piece_to_id[p] for p in target_pieces]
+    full_ids = prompt_ids + target_ids
+
+    out = torch.full((seq_len,), tok.pad_id, dtype=torch.long)
+    L = min(seq_len, len(full_ids))
+    out[:L] = torch.tensor(full_ids[:L], dtype=torch.long)
+    return out, len(prompt_ids), 1
+
+
+def _make_target_anchored_chain_sample(
+    rng: np.random.Generator,
+    tok: SyntheticTokenizer,
+    seq_len: int,
+) -> tuple[torch.Tensor, int, int]:
+    """
+    Task 2 (Target-Anchored Arithmetic / Story Planning - `[PLAN]`):
+    Prompt specifies a target final value `END = T` and a sequence of deltas `d_1 .. d_10`.
+    The model must choose the unique initial value `x_0` at Position 0 such that after applying
+    `x_{k} = (x_{k-1} + d_k) mod 32` for 10 steps, the final token at Position 10 equals `T`!
+    Then it emits the full trajectory `x_0, x_1, ..., x_10 = T`.
+    """
+    steps = 10
+    deltas = [int(x) for x in rng.integers(1, 8, size=steps)]
+    target_end = int(rng.integers(0, 32))
+    total_delta = sum(deltas) % 32
+    x0 = (target_end - total_delta) % 32
+
+    traj = [x0]
+    for d in deltas:
+        traj.append((traj[-1] + d) % 32)
+    assert traj[-1] == target_end
+
+    prompt_pieces = (
+        ["<BOS>", "[PLAN]", "END", f"{target_end}", "STEP"]
+        + [f"+ {d}" for d in deltas]
+        # flatten "+ d" into two tokens
+    )
+    flat_prompt: list[str] = ["<BOS>", "[PLAN]", "END", f"{target_end}", "STEP"]
+    for d in deltas:
+        flat_prompt.extend(["+", f"{d}"])
+    flat_prompt.append("=>")
+
+    target_pieces = [f"{val}" for val in traj] + ["=", f"{target_end}", "OK", "<EOS>"]
+
+    prompt_ids = [tok.piece_to_id[p] for p in flat_prompt]
+    target_ids = [tok.piece_to_id[p] for p in target_pieces]
+    full_ids = prompt_ids + target_ids
+
+    out = torch.full((seq_len,), tok.pad_id, dtype=torch.long)
+    L = min(seq_len, len(full_ids))
+    out[:L] = torch.tensor(full_ids[:L], dtype=torch.long)
+    return out, len(prompt_ids), 2
+
+
+class SyntheticLookaheadDataset(Dataset[dict[str, Any]]):
+    def __init__(
+        self,
+        num_samples: int,
+        seq_len: int,
+        task: str = "all",
+        vocab_size: int = 128,
+        seed: int = 42,
+    ) -> None:
+        super().__init__()
+        self.num_samples = num_samples
+        self.seq_len = seq_len
+        self.task = task
+        self.tokenizer = SyntheticTokenizer(vocab_size=vocab_size)
+
+        rng = np.random.default_rng(seed)
+        self.samples: list[dict[str, Any]] = []
+
+        for i in range(num_samples):
+            if task == "path":
+                t_choice = 0
+            elif task == "reverse":
+                t_choice = 1
+            elif task == "constraint":
+                t_choice = 2
+            else:
+                t_choice = i % 3
+
+            if t_choice == 0:
+                seq, plen, tid = _make_path_planning_sample(rng, self.tokenizer, seq_len)
+            elif t_choice == 1:
+                seq, plen, tid = _make_reverse_plan_sample(rng, self.tokenizer, seq_len)
+            else:
+                seq, plen, tid = _make_target_anchored_chain_sample(rng, self.tokenizer, seq_len)
+
+            self.samples.append(
+                {
+                    "input_ids": seq,
+                    "prompt_len": plen,
+                    "task_id": tid,
+                }
+            )
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, idx: int) -> dict[str, Any]:
+        return self.samples[idx]
+
+
+def build_synthetic_datasets(
+    cfg: ConveyerConfig,
+) -> tuple[SyntheticLookaheadDataset, SyntheticLookaheadDataset, SyntheticTokenizer]:
+    train_ds = SyntheticLookaheadDataset(
+        num_samples=cfg.data.num_train,
+        seq_len=cfg.data.seq_len,
+        task=cfg.data.task,
+        vocab_size=cfg.data.vocab_size,
+        seed=cfg.train.seed,
+    )
+    val_ds = SyntheticLookaheadDataset(
+        num_samples=cfg.data.num_val,
+        seq_len=cfg.data.seq_len,
+        task=cfg.data.task,
+        vocab_size=cfg.data.vocab_size,
+        seed=cfg.train.seed + 10000,
+    )
+    return train_ds, val_ds, train_ds.tokenizer

src/conveyer/data/tinystories.py

@@ -0,0 +1,533 @@
+from __future__ import annotations
+
+from collections import Counter
+import json
+import os
+from pathlib import Path
+import time
+from typing import Any, Iterator, Optional
+import urllib.request
+import numpy as np
+import tiktoken
+import torch
+from torch.utils.data import Dataset
+
+from conveyer.config import ConveyerConfig
+
+TINYSTORIES_VALID_URL = (
+    "https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/TinyStoriesV2-GPT4-valid.txt"
+)
+TINYSTORIES_TRAIN_URL = (
+    "https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/TinyStoriesV2-GPT4-train.txt"
+)
+
+
+class TiktokenCompactTokenizer:
+    """
+    Backed by OpenAI `tiktoken` (`gpt2` BPE encoding), mapped into a compact vocabulary
+    of size `vocab_size` (e.g. 4096) covering the most frequent real-TinyStories BPE subwords
+    plus all 256 single-byte BPE tokens so ANY arbitrary English string encodes with ZERO `<UNK>`.
+    """
+
+    def __init__(
+        self,
+        gpt2_to_compact: dict[int, int],
+        compact_to_gpt2: list[int],
+        source: str = "real_tinystories_v2",
+    ) -> None:
+        self.enc = tiktoken.get_encoding("gpt2")
+        self.gpt2_to_compact = gpt2_to_compact
+        self.compact_to_gpt2 = compact_to_gpt2
+        self.source = source
+        self.vocab_size = len(compact_to_gpt2)
+        self.pad_id = 0
+        self.bos_id = 1
+        self.eos_id = 2
+        self.unk_id = 3
+
+        # Prebuild byte-level fallback table so any rare BPE token decomposes into exact byte BPE tokens
+        self.byte_to_compact: dict[int, int] = {}
+        for b in range(256):
+            gid = self.enc.encode_single_token(bytes([b]))
+            if gid in self.gpt2_to_compact:
+                self.byte_to_compact[b] = self.gpt2_to_compact[gid]
+
+        # Prebuild fast array lookup for GPT-2 token IDs (0..50256) -> compact ID (-1 if rare)
+        self._lut = np.full(self.enc.n_vocab, -1, dtype=np.int32)
+        for gid, cid in self.gpt2_to_compact.items():
+            if 0 <= gid < self.enc.n_vocab:
+                self._lut[gid] = cid
+
+    @classmethod
+    def build_from_texts(cls, texts: list[str], vocab_size: int = 4096) -> "TiktokenCompactTokenizer":
+        enc = tiktoken.get_encoding("gpt2")
+        # Special tokens occupy 0..3: <PAD>=0, <BOS>=1, <EOS>=2, <UNK>=3
+        compact_to_gpt2: list[int] = [-1, -2, -3, -4]
+        seen_gpt2: set[int] = set()
+
+        # 1. Always include all 256 single-byte BPE tokens so 100% of UTF-8 strings are representable
+        for b in range(256):
+            gid = enc.encode_single_token(bytes([b]))
+            if gid not in seen_gpt2 and len(compact_to_gpt2) < vocab_size:
+                seen_gpt2.add(gid)
+                compact_to_gpt2.append(gid)
+
+        # 2. Count BPE token frequencies across real TinyStories
+        counts: Counter[int] = Counter()
+        sample = texts[:50000]
+        for batch_list in enc.encode_ordinary_batch(sample, num_threads=8):
+            counts.update(batch_list)
+
+        for gid, _ in counts.most_common():
+            if len(compact_to_gpt2) >= vocab_size:
+                break
+            if gid not in seen_gpt2:
+                seen_gpt2.add(gid)
+                compact_to_gpt2.append(gid)
+
+        # Pad up to vocab_size with sequential gpt2 token IDs if needed
+        next_gid = 0
+        while len(compact_to_gpt2) < vocab_size:
+            if next_gid not in seen_gpt2:
+                seen_gpt2.add(next_gid)
+                compact_to_gpt2.append(next_gid)
+            next_gid += 1
+
+        gpt2_to_compact = {gid: cid for cid, gid in enumerate(compact_to_gpt2) if gid >= 0}
+        return cls(
+            gpt2_to_compact=gpt2_to_compact,
+            compact_to_gpt2=compact_to_gpt2,
+            source="real_tinystories_v2",
+        )
+
+    def save(self, path: str | Path) -> None:
+        p = Path(path)
+        p.parent.mkdir(parents=True, exist_ok=True)
+        payload = {
+            "source": self.source,
+            "compact_to_gpt2": self.compact_to_gpt2,
+        }
+        p.write_text(json.dumps(payload), encoding="utf-8")
+
+    @classmethod
+    def load(cls, path: str | Path) -> "TiktokenCompactTokenizer":
+        data = json.loads(Path(path).read_text(encoding="utf-8"))
+        compact_to_gpt2: list[int] = data["compact_to_gpt2"]
+        source: str = data.get("source", "legacy")
+        gpt2_to_compact = {gid: cid for cid, gid in enumerate(compact_to_gpt2) if gid >= 0}
+        return cls(gpt2_to_compact=gpt2_to_compact, compact_to_gpt2=compact_to_gpt2, source=source)
+
+    def id_to_piece(self, token_id: int) -> str:
+        if token_id == self.pad_id:
+            return "<PAD>"
+        if token_id == self.bos_id:
+            return "<BOS>"
+        if token_id == self.eos_id:
+            return "<EOS>"
+        if token_id == self.unk_id or not (0 <= token_id < len(self.compact_to_gpt2)):
+            return "<UNK>"
+        gid = self.compact_to_gpt2[token_id]
+        if gid < 0:
+            return "<SPEC>"
+        raw = self.enc.decode_single_token_bytes(gid).decode("utf-8", errors="replace")
+        return raw
+
+    def _map_gids_to_compact(self, gids: list[int], max_tokens: Optional[int] = None) -> list[int]:
+        ids: list[int] = [self.bos_id]
+        gpt2_map = self.gpt2_to_compact
+        byte_map = self.byte_to_compact
+        dec_single = self.enc.decode_single_token_bytes
+        unk = self.unk_id
+        for gid in gids:
+            cid = gpt2_map.get(gid)
+            if cid is not None:
+                ids.append(cid)
+            else:
+                for b in dec_single(gid):
+                    ids.append(byte_map.get(b, unk))
+            if max_tokens is not None and len(ids) >= max_tokens:
+                break
+        return ids
+
+    def encode(self, text: str) -> list[int]:
+        gids = self.enc.encode_ordinary(text)
+        return self._map_gids_to_compact(gids)
+
+    def encode_batch_into_matrix(
+        self,
+        texts: list[str],
+        out_matrix: np.ndarray,
+        out_lens: np.ndarray,
+        start_row: int,
+        seq_len: int,
+    ) -> int:
+        """
+        Multi-threaded Rust BPE batch encoding into `out_matrix[start_row : start_row + written]` (uint16).
+        Truncates input text to ~1400 chars before BPE since we only need `seq_len` (e.g. 160) tokens.
+        """
+        char_cap = max(seq_len * 8, 1024)
+        clipped = [t[:char_cap] if len(t) > char_cap else t for t in texts]
+        batch_gids = self.enc.encode_ordinary_batch(clipped, num_threads=8)
+
+        row = start_row
+        max_rows = out_matrix.shape[0]
+        eos = self.eos_id
+        pad = self.pad_id
+        max_body = seq_len - 1
+
+        for gids in batch_gids:
+            if row >= max_rows:
+                break
+            ids = self._map_gids_to_compact(gids, max_tokens=max_body)
+            if len(ids) < 16:
+                continue
+            if len(ids) < seq_len:
+                ids.append(eos)
+            L = min(seq_len, len(ids))
+            out_matrix[row, :L] = ids[:L]
+            if L < seq_len:
+                out_matrix[row, L:] = pad
+            out_lens[row] = L
+            row += 1
+        return row - start_row
+
+    def decode(self, ids: list[int]) -> str:
+        gids: list[int] = []
+        for cid in ids:
+            if cid in (self.pad_id, self.bos_id, self.unk_id):
+                continue
+            if cid == self.eos_id:
+                break
+            if 0 <= cid < len(self.compact_to_gpt2):
+                gid = self.compact_to_gpt2[cid]
+                if gid >= 0:
+                    gids.append(gid)
+        return self.enc.decode(gids)
+
+
+def _download_with_resume(url: str, dest: Path, min_expected_bytes: int = 1_000_000) -> Path:
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    existing = dest.stat().st_size if dest.exists() else 0
+    total_size = 0
+    try:
+        req_head = urllib.request.Request(url, method="HEAD", headers={"User-Agent": "Mozilla/5.0"})
+        with urllib.request.urlopen(req_head, timeout=30) as r:
+            total_size = int(r.headers.get("Content-Length", 0))
+    except Exception:
+        total_size = 0
+
+    if total_size > 0 and existing >= total_size:
+        return dest
+    if total_size == 0 and existing >= min_expected_bytes:
+        return dest
+
+    headers = {"User-Agent": "Mozilla/5.0"}
+    mode = "wb"
+    if existing > 0 and total_size > existing:
+        headers["Range"] = f"bytes={existing}-"
+        mode = "ab"
+    else:
+        existing = 0
+
+    print(
+        f"[TinyStories] Downloading {dest.name} "
+        f"({existing / (1024 * 1024):.1f} / {total_size / (1024 * 1024):.1f} MB)...",
+        flush=True,
+    )
+    req = urllib.request.Request(url, headers=headers)
+    t0 = time.time()
+    last_log = t0
+    downloaded = existing
+    with urllib.request.urlopen(req, timeout=120) as r, open(dest, mode) as f:
+        while True:
+            chunk = r.read(4 * 1024 * 1024)
+            if not chunk:
+                break
+            f.write(chunk)
+            downloaded += len(chunk)
+            now = time.time()
+            if now - last_log >= 5.0:
+                dt = now - t0
+                speed = (downloaded - existing) / (1024 * 1024) / max(dt, 1e-3)
+                pct = 100.0 * downloaded / max(total_size, 1)
+                print(
+                    f"  [{dest.name}] {downloaded / (1024 * 1024):.1f} / "
+                    f"{total_size / (1024 * 1024):.1f} MB ({pct:.1f}%) at {speed:.1f} MB/s",
+                    flush=True,
+                )
+                last_log = now
+    return dest
+
+
+def iter_clean_real_stories(txt_path: Path, max_stories: Optional[int] = None) -> Iterator[str]:
+    """
+    Streams clean real TinyStories from a `<|endoftext|>`-separated raw text file without
+    loading the entire 2.12 GB string list into memory at once.
+    Filters out byte-0 mid-sentence fragments and any synthetic templates.
+    """
+    yielded = 0
+    buf = ""
+    with open(txt_path, "r", encoding="utf-8", errors="ignore") as f:
+        while True:
+            chunk = f.read(16 * 1024 * 1024)
+            if not chunk:
+                break
+            buf += chunk
+            parts = buf.split("<|endoftext|>")
+            buf = parts[-1]
+            for raw in parts[:-1]:
+                s = " ".join(raw.strip().split())
+                if len(s) < 80:
+                    continue
+                if not (s[0].isupper() or s[0] in "\"'"):
+                    continue
+                if "walked into the " in s and "Together, they" in s:
+                    continue
+                yield s
+                yielded += 1
+                if max_stories is not None and yielded >= max_stories:
+                    return
+    s = " ".join(buf.strip().split())
+    if len(s) >= 80 and (s[0].isupper() or s[0] in "\"'"):
+        if not ("walked into the " in s and "Together, they" in s):
+            if max_stories is None or yielded < max_stories:
+                yield s
+
+
+def ensure_raw_tinystories_files(data_dir: str | Path, need_full_train: bool = True) -> tuple[Path, Path]:
+    root = Path(data_dir)
+    root.mkdir(parents=True, exist_ok=True)
+    val_path = _download_with_resume(TINYSTORIES_VALID_URL, root / "TinyStoriesV2-GPT4-valid.txt", min_expected_bytes=20_000_000)
+    train_path = root / "TinyStoriesV2-GPT4-train.txt"
+    if need_full_train or not train_path.exists():
+        _download_with_resume(TINYSTORIES_TRAIN_URL, train_path, min_expected_bytes=2_000_000_000)
+    return train_path, val_path
+
+
+def load_or_prepare_tinystories(data_dir: str | Path, max_stories: int = 42000) -> list[str]:
+    """
+    Returns up to `max_stories` 100% REAL TinyStories strings (zero templates).
+    """
+    root = Path(data_dir)
+    train_path, val_path = ensure_raw_tinystories_files(root, need_full_train=(max_stories > 25000))
+    source_path = train_path if train_path.exists() and train_path.stat().st_size > 20_000_000 else val_path
+    stories = list(iter_clean_real_stories(source_path, max_stories=max_stories))
+    return stories
+
+
+def _build_or_load_memmap_split(
+    split_name: str,
+    raw_txt_path: Path,
+    tokenizer: TiktokenCompactTokenizer,
+    data_dir: Path,
+    target_stories: int,
+    seq_len: int,
+) -> tuple[np.memmap, np.ndarray]:
+    """
+    Pre-tokenizes `target_stories` real TinyStories from `raw_txt_path` into a compact disk-backed
+    `uint16` numpy memmap of shape `(N, seq_len)` + `uint16` lengths array `(N,)`.
+    Loading an existing memmap takes < 10ms and uses near-zero RAM.
+    """
+    prefix = data_dir / f"real_{split_name}_v5_n{target_stories}_s{seq_len}_v{tokenizer.vocab_size}"
+    tokens_bin = Path(str(prefix) + ".tokens.bin")
+    lens_bin = Path(str(prefix) + ".lens.bin")
+    meta_json = Path(str(prefix) + ".meta.json")
+
+    if tokens_bin.exists() and lens_bin.exists() and meta_json.exists():
+        meta = json.loads(meta_json.read_text(encoding="utf-8"))
+        n_rows = int(meta["num_rows"])
+        if n_rows > 0 and tokens_bin.stat().st_size == n_rows * seq_len * 2:
+            mm = np.memmap(tokens_bin, dtype=np.uint16, mode="r", shape=(n_rows, seq_len))
+            lens = np.fromfile(lens_bin, dtype=np.uint16)
+            return mm, lens
+
+    print(
+        f"[TinyStories] Building binary token cache for '{split_name}' "
+        f"(target={target_stories:,} stories, seq_len={seq_len})...",
+        flush=True,
+    )
+    tmp_tokens = Path(str(tokens_bin) + ".tmp")
+    tmp_lens = Path(str(lens_bin) + ".tmp")
+
+    mm_write = np.memmap(tmp_tokens, dtype=np.uint16, mode="w+", shape=(target_stories, seq_len))
+    lens_arr = np.zeros(target_stories, dtype=np.uint16)
+
+    t0 = time.time()
+    last_log = t0
+    row = 0
+    batch_texts: list[str] = []
+    batch_chunk_size = 50000
+
+    for story in iter_clean_real_stories(raw_txt_path, max_stories=int(target_stories * 1.05) + 2000):
+        batch_texts.append(story)
+        if len(batch_texts) >= batch_chunk_size:
+            added = tokenizer.encode_batch_into_matrix(batch_texts, mm_write, lens_arr, row, seq_len)
+            row += added
+            batch_texts.clear()
+            now = time.time()
+            if now - last_log >= 4.0:
+                dt = now - t0
+                rate = row / max(dt, 1e-3)
+                print(
+                    f"  [{split_name}] Tokenized {row:,} / {target_stories:,} stories "
+                    f"({100.0 * row / max(target_stories, 1):.1f}%) at {rate:,.0f} stories/s",
+                    flush=True,
+                )
+                last_log = now
+            if row >= target_stories:
+                break
+
+    if batch_texts and row < target_stories:
+        added = tokenizer.encode_batch_into_matrix(batch_texts, mm_write, lens_arr, row, seq_len)
+        row += added
+        batch_texts.clear()
+
+    mm_write.flush()
+    del mm_write
+
+    # If actual valid stories in file < target_stories, truncate file to exact `row * seq_len * 2` bytes
+    if row < target_stories:
+        with open(tmp_tokens, "r+b") as f:
+            f.truncate(row * seq_len * 2)
+        lens_arr = lens_arr[:row]
+
+    lens_arr.tofile(tmp_lens)
+    os.replace(tmp_tokens, tokens_bin)
+    os.replace(tmp_lens, lens_bin)
+    meta_json.write_text(
+        json.dumps({"num_rows": int(row), "seq_len": int(seq_len), "vocab_size": int(tokenizer.vocab_size)}),
+        encoding="utf-8",
+    )
+    dt = time.time() - t0
+    print(f"[TinyStories] Saved '{split_name}' binary cache: {row:,} real stories in {dt:.1f}s", flush=True)
+
+    mm = np.memmap(tokens_bin, dtype=np.uint16, mode="r", shape=(row, seq_len))
+    return mm, lens_arr
+
+
+class TinyStoriesDataset(Dataset[dict[str, Any]]):
+    def __init__(
+        self,
+        stories: Optional[list[str]] = None,
+        tokenizer: Optional[TiktokenCompactTokenizer] = None,
+        seq_len: int = 160,
+        prompt_len: int = 28,
+        randomize_prompt_len: bool = False,
+        token_matrix: Optional[np.ndarray] = None,
+        valid_lens: Optional[np.ndarray] = None,
+    ) -> None:
+        super().__init__()
+        self.tokenizer = tokenizer
+        self.seq_len = seq_len
+        self.prompt_len = prompt_len
+        self.randomize_prompt_len = randomize_prompt_len
+
+        if token_matrix is not None and valid_lens is not None:
+            self.token_matrix = token_matrix
+            self.valid_lens = valid_lens
+        else:
+            assert stories is not None and tokenizer is not None
+            n_cap = len(stories)
+            mat = np.full((n_cap, seq_len), tokenizer.pad_id, dtype=np.uint16)
+            lens = np.zeros(n_cap, dtype=np.uint16)
+            written = tokenizer.encode_batch_into_matrix(stories, mat, lens, 0, seq_len)
+            self.token_matrix = mat[:written]
+            self.valid_lens = lens[:written]
+
+    def __len__(self) -> int:
+        return int(self.token_matrix.shape[0])
+
+    def __getitem__(self, idx: int) -> dict[str, Any]:
+        row = torch.from_numpy(np.asarray(self.token_matrix[idx], dtype=np.int64))
+        L = int(self.valid_lens[idx])
+        if not self.randomize_prompt_len:
+            p_len = max(4, min(self.prompt_len, L - 8))
+        else:
+            low = 6
+            high = max(low + 1, min(64, L - 8))
+            p_len = int(torch.randint(low, high, (1,)).item())
+        return {
+            "input_ids": row,
+            "prompt_len": p_len,
+            "task_id": 0,
+        }
+
+
+def build_tinystories_datasets(
+    cfg: ConveyerConfig,
+    infer_only: bool = False,
+) -> tuple[TinyStoriesDataset, TinyStoriesDataset, TiktokenCompactTokenizer]:
+    data_dir = Path(cfg.data.data_dir)
+    data_dir.mkdir(parents=True, exist_ok=True)
+    tok_path = data_dir / f"tiktoken_compact_{cfg.data.vocab_size}.json"
+    val_sample_path = data_dir / "val_sample_256.txt"
+
+    if infer_only and tok_path.exists():
+        tokenizer = TiktokenCompactTokenizer.load(tok_path)
+        cfg.data.vocab_size = tokenizer.vocab_size
+        cfg.model.vocab_size = tokenizer.vocab_size
+        if val_sample_path.exists():
+            sample_lines = [l.strip() for l in val_sample_path.read_text(encoding="utf-8").splitlines() if l.strip()]
+        else:
+            _, val_txt = ensure_raw_tinystories_files(data_dir, need_full_train=False)
+            sample_lines = list(iter_clean_real_stories(val_txt, max_stories=256))
+        val_ds = TinyStoriesDataset(
+            stories=sample_lines,
+            tokenizer=tokenizer,
+            seq_len=cfg.data.seq_len,
+            prompt_len=28,
+            randomize_prompt_len=False,
+        )
+        return val_ds, val_ds, tokenizer
+
+    train_txt, val_txt = ensure_raw_tinystories_files(data_dir, need_full_train=True)
+
+    tokenizer_opt: Optional[TiktokenCompactTokenizer] = None
+    if tok_path.exists():
+        loaded = TiktokenCompactTokenizer.load(tok_path)
+        if loaded.source == "real_tinystories_v2":
+            tokenizer_opt = loaded
+
+    if tokenizer_opt is None:
+        print("[TinyStories] Building compact 4096 BPE vocabulary from 50,000 real stories...", flush=True)
+        sample_stories = list(iter_clean_real_stories(train_txt, max_stories=50000))
+        tokenizer_opt = TiktokenCompactTokenizer.build_from_texts(sample_stories, vocab_size=cfg.data.vocab_size)
+        tokenizer_opt.save(tok_path)
+
+    tokenizer = tokenizer_opt
+    cfg.data.vocab_size = tokenizer.vocab_size
+    cfg.model.vocab_size = tokenizer.vocab_size
+
+    train_mm, train_lens = _build_or_load_memmap_split(
+        split_name="train",
+        raw_txt_path=train_txt,
+        tokenizer=tokenizer,
+        data_dir=data_dir,
+        target_stories=cfg.data.num_train,
+        seq_len=cfg.data.seq_len,
+    )
+    val_mm, val_lens = _build_or_load_memmap_split(
+        split_name="val",
+        raw_txt_path=val_txt,
+        tokenizer=tokenizer,
+        data_dir=data_dir,
+        target_stories=cfg.data.num_val,
+        seq_len=cfg.data.seq_len,
+    )
+
+    train_ds = TinyStoriesDataset(
+        tokenizer=tokenizer,
+        seq_len=cfg.data.seq_len,
+        randomize_prompt_len=True,
+        token_matrix=train_mm,
+        valid_lens=train_lens,
+    )
+    val_ds = TinyStoriesDataset(
+        tokenizer=tokenizer,
+        seq_len=cfg.data.seq_len,
+        prompt_len=28,
+        randomize_prompt_len=False,
+        token_matrix=val_mm,
+        valid_lens=val_lens,
+    )
+    return train_ds, val_ds, tokenizer

src/conveyer/eval.py

@@ -0,0 +1,557 @@
+from __future__ import annotations
+
+import argparse
+from collections import Counter
+import json
+import math
+from pathlib import Path
+from typing import Any
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+
+from conveyer.checkpoint import load_checkpoint
+from conveyer.config import ConveyerConfig
+from conveyer.data import build_dataset_and_tokenizer, collate_conveyer_batch
+from conveyer.infer import generate_conveyer
+from conveyer.losses import compute_unrolled_conveyer_loss
+from conveyer.model import ConveyerDiffusionLM
+
+
+# ---------------------------------------------------------------------------
+# Benchmark 3: Zero-Shot Commonsense / Story-Coherence Multiple Choice (2-way)
+# ---------------------------------------------------------------------------
+# Model scores log P(option | context) on unseen story pairs testing:
+#   - Physical commonsense (heavy vs light, hot vs cold, broken vs whole)
+#   - Emotional / social cause-and-effect (lost toy -> sad, gift -> happy)
+#   - Narrative entity & object permanence (tracking who holds what, where they are)
+COMMONSENSE_CLOZE_PAIRS: list[tuple[str, str, str, str]] = [
+    (
+        "emotion_cause",
+        "Once upon a time, Tom lost his favorite red ball in the dark woods. Because he lost his ball, Tom felt very",
+        " sad and started to cry.",
+        " happy and started to laugh.",
+    ),
+    (
+        "emotion_cause",
+        "Lily worked all morning to bake a sweet cake for her mom. When her mom smiled and hugged her, Lily felt very",
+        " proud and happy.",
+        " angry and mean.",
+    ),
+    (
+        "physical_weight",
+        "The little mouse tried to lift the huge heavy rock, but the rock was much too",
+        " heavy for the tiny mouse to move.",
+        " light and floated into the sky.",
+    ),
+    (
+        "physical_temperature",
+        "Ben went outside in the snow without his coat or gloves. Soon his hands felt very",
+        " cold and freezing.",
+        " hot and burning.",
+    ),
+    (
+        "physical_temperature",
+        "The soup had just come off the stove and was boiling. Mom told Mia to wait because the soup was too",
+        " hot to eat right now.",
+        " cold and frozen like ice.",
+    ),
+    (
+        "physical_fragile",
+        "The glass cup fell off the high table onto the hard stone floor. With a loud crash, the glass",
+        " broke into many small pieces.",
+        " bounced high like a rubber ball.",
+    ),
+    (
+        "object_use",
+        "The wooden door was locked tight. To unlock the door, Anna took out her small brass",
+        " key and turned it in the lock.",
+        " spoon and stirred the soup.",
+    ),
+    (
+        "object_use",
+        "It started to rain hard outside. Before walking to school, Max opened his big",
+        " umbrella to stay dry.",
+        " pillow to go to sleep.",
+    ),
+    (
+        "biological_need",
+        "The little puppy had been running in the park all day and had not eaten anything. His tummy rumbled because he was very",
+        " hungry and wanted some food.",
+        " sleepy and wanted a haircut.",
+    ),
+    (
+        "biological_need",
+        "After playing all day until night came and the stars shone in the sky, the tired children went to bed to",
+        " sleep until morning.",
+        " swim in the cold ocean.",
+    ),
+    (
+        "entity_tracking",
+        "Sam had a blue car and Lucy had a yellow doll. Sam gave his blue car to Lucy. Now Lucy had the yellow doll and the",
+        " blue car.",
+        " green bicycle.",
+    ),
+    (
+        "entity_tracking",
+        " Lily put her shiny coin inside the small red box and closed the lid. When she opened the red box later, she saw her",
+        " shiny coin inside.",
+        " big brown horse inside.",
+    ),
+    (
+        "danger_reaction",
+        "Lily and a little bunny walked into the green forest. Suddenly a big fierce lion jumped out and roared loudly. The little bunny was",
+        " scared and hid behind a tree.",
+        " bored and read a newspaper.",
+    ),
+    (
+        "social_repair",
+        "Tom accidentally knocked over Mia's tall block tower. Tom felt bad about the accident, so he said,",
+        ' "I am sorry, Mia. Let me help you build it again."',
+        ' "Thank you for the delicious birthday cake."',
+    ),
+    (
+        "spatial_logic",
+        "The bird wanted to reach the nest at the very top of the tall tree, so it flapped its wings and",
+        " flew up high into the branches.",
+        " dug a deep hole under the mud.",
+    ),
+    (
+        "weather_visibility",
+        "At midnight all the lights went out in the room. Because it was pitch dark, Ben could not",
+        " see anything around him.",
+        " hear the loud thunder outside.",
+    ),
+]
+
+
+NOVEL_STORY_PROMPTS: list[tuple[str, list[str]]] = [
+    (
+        "Once upon a time, Lily and a little bunny walked into the green forest to find the golden key near the locked treasure chest. Suddenly a big lion",
+        ["lion", "Lily", "bunny", "key", "chest", "forest", "scared", "roar", "said"],
+    ),
+    (
+        "Once upon a time, there was a little robot named Beep who lived in a noisy workshop. One rainy morning, Beep found a tiny lost kitten by the door. Beep",
+        ["Beep", "kitten", "robot", "door", "rain", "warm", "milk", "help"],
+    ),
+    (
+        "Sam and his sister Mia built a huge sandcastle on the beach near the waves. Suddenly, a big wave splashed toward their castle. Sam quickly",
+        ["Sam", "Mia", "sand", "castle", "wave", "beach", "water", "wall"],
+    ),
+    (
+        "In a quiet village, an old baker dropped his round loaf of bread and it rolled out the door down the hill. A hungry puppy saw the bread and",
+        ["bread", "puppy", "baker", "hill", "chased", "ran", "ate", "mouth"],
+    ),
+    (
+        "Lucy had a box of bright crayons, but she could not find the blue crayon to color the sky. She looked under her bed and",
+        ["Lucy", "crayon", "blue", "bed", "sky", "found", "draw", "paper"],
+    ),
+    (
+        "One winter day, Max made a snowman in the yard with a carrot nose. When the warm sun came out in the afternoon, the snowman began to",
+        ["Max", "snowman", "sun", "melt", "carrot", "snow", "warm", "water"],
+    ),
+]
+
+
+@torch.no_grad()
+def evaluate_zero_shot_commonsense(
+    model: ConveyerDiffusionLM,
+    tokenizer: Any,
+    device: torch.device,
+) -> dict[str, Any]:
+    """
+    Evaluates zero-shot 2-way multiple-choice accuracy (BLiMP / StoryCloze style) using both:
+      1. Causal LM log-probability per token over the continuation.
+      2. Conveyor Slot-0 first-token preference on the continuation.
+    """
+    model.eval()
+    correct_lm = 0
+    correct_s0 = 0
+    margin_sum = 0.0
+    details: list[dict[str, Any]] = []
+
+    for category, ctx, good_cont, bad_cont in COMMONSENSE_CLOZE_PAIRS:
+        ctx_ids = tokenizer.encode(ctx.strip())
+        good_ids = tokenizer.encode(ctx.strip() + good_cont)
+        bad_ids = tokenizer.encode(ctx.strip() + bad_cont)
+        plen = len(ctx_ids)
+
+        def _score_seq(full_ids: list[int]) -> float:
+            inp = torch.tensor([full_ids], dtype=torch.long, device=device)
+            logits, _ = model.forward_causal_and_cache(inp, prefix_len=len(full_ids))
+            log_probs = F.log_softmax(logits[0, :-1, :].float(), dim=-1)
+            targets = inp[0, 1:]
+            # Mean log-prob over the continuation tokens [plen-1 : len-1]
+            start_idx = max(0, plen - 1)
+            cont_lp = log_probs[start_idx:, :].gather(1, targets[start_idx:].unsqueeze(1)).squeeze(1)
+            return float(cont_lp.mean().item())
+
+        score_good = _score_seq(good_ids)
+        score_bad = _score_seq(bad_ids)
+        is_ok = score_good > score_bad
+        correct_lm += int(is_ok)
+        margin_sum += score_good - score_bad
+
+        # Also check Conveyor Slot-0 log-prob on the first continuation token
+        if len(good_ids) > plen and len(bad_ids) > plen:
+            p_tensor = torch.tensor([ctx_ids], dtype=torch.long, device=device)
+            kv = model.prefill_prefix(p_tensor)
+            belt = model.init_empty_belt(1, tick=plen, device=device, deterministic=True)
+            out = model.forward_belt(belt, kv_cache=kv)
+            s0_lp = F.log_softmax(out.slot0_logits[0].float(), dim=-1)
+            s0_good = float(s0_lp[good_ids[plen]].item())
+            s0_bad = float(s0_lp[bad_ids[plen]].item())
+            correct_s0 += int(s0_good > s0_bad)
+
+        details.append(
+            {
+                "category": category,
+                "correct": is_ok,
+                "logprob_good": round(score_good, 3),
+                "logprob_bad": round(score_bad, 3),
+            }
+        )
+
+    n = max(1, len(COMMONSENSE_CLOZE_PAIRS))
+    return {
+        "commonsense_lm_acc": correct_lm / n,
+        "commonsense_slot0_acc": correct_s0 / n,
+        "mean_logprob_margin": margin_sum / n,
+        "num_pairs": len(COMMONSENSE_CLOZE_PAIRS),
+        "details": details,
+    }
+
+
+@torch.no_grad()
+def evaluate_prefix_sensitivity_and_generation_quality(
+    model: ConveyerDiffusionLM,
+    tokenizer: Any,
+    device: torch.device,
+    max_new_tokens: int = 80,
+) -> dict[str, Any]:
+    """
+    Measures:
+      1. Prefix Sensitivity KL(P(slot0 | prompt) || P(slot0 | shuffled_prompt)):
+         In the broken template model this was 0.03 nats (ignored word order).
+         A real language model should have KL >> 0.5 nats.
+      2. Open-Ended Generation Quality on 6 unseen prompts:
+         - Distinct-2 / Distinct-3 n-gram diversity (1.0 = no repetition, <0.5 = looping)
+         - 4-gram repetition rate (0.0 = clean prose)
+         - Prompt entity/topic retention rate (% of continuations that reference the prompt's new entities/events)
+    """
+    model.eval()
+    kl_list: list[float] = []
+    dist2_list: list[float] = []
+    dist3_list: list[float] = []
+    rep4_list: list[float] = []
+    entity_hits = 0
+    samples: list[dict[str, str]] = []
+
+    g = torch.Generator().manual_seed(123)
+    for prompt_text, expected_keywords in NOVEL_STORY_PROMPTS:
+        ids = tokenizer.encode(prompt_text.strip())
+        p_true = torch.tensor([ids], dtype=torch.long, device=device)
+        kv_true = model.prefill_prefix(p_true)
+        belt_true = model.init_empty_belt(1, tick=len(ids), device=device, deterministic=True)
+        d_true = F.softmax(model.forward_belt(belt_true, kv_cache=kv_true).slot0_logits[0].float(), dim=-1)
+
+        perm = torch.randperm(len(ids) - 1, generator=g) + 1
+        shuf_ids = [ids[0]] + [ids[i] for i in perm.tolist()]
+        p_shuf = torch.tensor([shuf_ids], dtype=torch.long, device=device)
+        kv_shuf = model.prefill_prefix(p_shuf)
+        belt_shuf = model.init_empty_belt(1, tick=len(ids), device=device, deterministic=True)
+        d_shuf = F.softmax(model.forward_belt(belt_shuf, kv_cache=kv_shuf).slot0_logits[0].float(), dim=-1)
+
+        kl = float((d_true * (d_true.clamp_min(1e-9).log() - d_shuf.clamp_min(1e-9).log())).sum().item())
+        kl_list.append(kl)
+
+        out_ids, _ = generate_conveyer(
+            model=model,
+            prompt_tokens=p_true,
+            max_new_tokens=max_new_tokens,
+            prompt_overlap=10,
+            temperature=0.0,
+            eos_token_id=tokenizer.eos_id,
+            deterministic_spawn=True,
+        )
+        gen_toks = out_ids[0, len(ids) :].tolist()
+        gen_text = tokenizer.decode(gen_toks).strip()
+
+        words = [w.strip(".,!?;:\"'()").lower() for w in gen_text.split() if w.strip()]
+        if len(words) >= 4:
+            bigrams = list(zip(words[:-1], words[1:]))
+            trigrams = list(zip(words[:-2], words[1:-1], words[2:]))
+            fourgrams = list(zip(words[:-3], words[1:-2], words[2:-1], words[3:]))
+            dist2_list.append(len(set(bigrams)) / max(1, len(bigrams)))
+            dist3_list.append(len(set(trigrams)) / max(1, len(trigrams)))
+            c4 = Counter(fourgrams)
+            repeated_4 = sum(v - 1 for v in c4.values() if v > 1)
+            rep4_list.append(repeated_4 / max(1, len(fourgrams)))
+
+        low_gen = gen_text.lower()
+        hit = any(kw.lower() in low_gen for kw in expected_keywords[:4])
+        entity_hits += int(hit)
+        samples.append({"prompt": prompt_text, "continuation": gen_text})
+
+    return {
+        "prefix_word_order_kl": sum(kl_list) / max(1, len(kl_list)),
+        "distinct_2": sum(dist2_list) / max(1, len(dist2_list)),
+        "distinct_3": sum(dist3_list) / max(1, len(dist3_list)),
+        "repeat_4gram_rate": sum(rep4_list) / max(1, len(rep4_list)),
+        "prompt_entity_retention": entity_hits / max(1, len(NOVEL_STORY_PROMPTS)),
+        "samples": samples,
+    }
+
+
+@torch.no_grad()
+def evaluate_crystallization_curve(
+    model: ConveyerDiffusionLM,
+    cfg: ConveyerConfig,
+    val_loader: DataLoader[dict[str, torch.Tensor]],
+    device: torch.device,
+    pad_id: int = 0,
+    max_batches: int = 20,
+) -> dict[str, Any]:
+    """
+    Computes the Per-Slot Accuracy, Entropy, and Pause Rate across the 64-slot oven belt
+    plus held-out real TinyStories Perplexity (Causal LM & Conveyor Slot-0).
+    """
+    model.eval()
+    B = model.belt_size
+    acc_sum = torch.zeros(B, device=device)
+    ent_sum = torch.zeros(B, device=device)
+    s0_loss_sum = 0.0
+    lm_loss_sum = 0.0
+    lm_acc_sum = 0.0
+    pause_sum = 0.0
+    count = 0
+
+    for i, batch in enumerate(val_loader):
+        if i >= max_batches:
+            break
+        seq = batch["input_ids"].to(device)
+        plen = batch["prompt_lengths"].to(device)
+        use_amp = device.type == "cuda" and torch.cuda.is_bf16_supported()
+        with torch.autocast(device_type=device.type, dtype=torch.bfloat16 if use_amp else torch.float32, enabled=use_amp):
+            m = compute_unrolled_conveyer_loss(
+                model=model,
+                seq_tokens=seq,
+                cfg=cfg,
+                prompt_lengths=plen,
+                tick_start=int(plen.min().item()),
+                pad_token_id=pad_id,
+            )
+        acc_sum += m.per_slot_acc
+        ent_sum += m.per_slot_entropy
+        s0_loss_sum += m.slot0_loss
+        lm_loss_sum += m.prefix_lm_loss
+        lm_acc_sum += m.prefix_lm_acc
+        pause_sum += m.mean_plan_pause
+        count += 1
+
+    denom = max(count, 1)
+    acc_curve = (acc_sum / denom).cpu().tolist()
+    ent_curve = (ent_sum / denom).cpu().tolist()
+    s0_ce = s0_loss_sum / denom
+    lm_ce = lm_loss_sum / denom
+    return {
+        "slot0_ce": s0_ce,
+        "slot0_ppl": math.exp(min(s0_ce, 20.0)),
+        "slot0_acc": acc_curve[0],
+        "causal_lm_ce": lm_ce,
+        "causal_lm_ppl": math.exp(min(lm_ce, 20.0)),
+        "causal_lm_acc": lm_acc_sum / denom,
+        "mean_plan_pause": pause_sum / denom,
+        "per_slot_acc": acc_curve,
+        "per_slot_entropy": ent_curve,
+    }
+
+
+@torch.no_grad()
+def evaluate_closed_loop_generation(
+    model: ConveyerDiffusionLM,
+    val_ds: Any,
+    tokenizer: Any,
+    device: torch.device,
+    num_episodes: int = 32,
+    prompt_overlap: int = 12,
+    bake_steps: int = 64,
+    max_new_tokens: int = 24,
+) -> dict[str, float]:
+    model.eval()
+    n_eval = min(num_episodes, len(val_ds))
+    first_tok_correct = 0
+    total_tok_correct = 0
+    total_tok_count = 0
+    exact_match = 0
+
+    for i in range(n_eval):
+        item = val_ds[i]
+        seq = item["input_ids"].to(device)
+        plen = int(item["prompt_len"])
+        prompt = seq[:plen].unsqueeze(0)
+
+        gt_full = seq[plen:].tolist()
+        gt_target: list[int] = []
+        for tid in gt_full:
+            if tid == tokenizer.pad_id:
+                break
+            gt_target.append(tid)
+            if tid == tokenizer.eos_id:
+                break
+
+        if not gt_target:
+            continue
+
+        gen_len = min(max_new_tokens, len(gt_target))
+        full_out, _ = generate_conveyer(
+            model=model,
+            prompt_tokens=prompt,
+            max_new_tokens=gen_len,
+            prompt_overlap=prompt_overlap,
+            bake_steps=bake_steps,
+            temperature=0.0,
+            eos_token_id=tokenizer.eos_id,
+            deterministic_spawn=True,
+        )
+        pred_ids = full_out[0, plen : plen + gen_len].tolist()
+        gt_slice = gt_target[:gen_len]
+
+        if len(pred_ids) > 0 and pred_ids[0] == gt_slice[0]:
+            first_tok_correct += 1
+
+        matches = sum(1 for p, g in zip(pred_ids, gt_slice) if p == g)
+        total_tok_correct += matches
+        total_tok_count += len(gt_slice)
+        if pred_ids == gt_slice:
+            exact_match += 1
+
+    return {
+        "first_token_acc": first_tok_correct / max(1, n_eval),
+        "token_acc": total_tok_correct / max(1, total_tok_count),
+        "exact_seq_acc": exact_match / max(1, n_eval),
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Comprehensive Benchmark Suite for Conveyor-Belt Diffusion LLM")
+    parser.add_argument("--ckpt", type=str, default="runs/tinystories/best.pt", help="Path to checkpoint (.pt)")
+    parser.add_argument("--device", type=str, default="auto", help="'cpu', 'cuda', or 'auto'")
+    parser.add_argument("--episodes", type=int, default=32, help="Number of closed-loop episodes")
+    parser.add_argument("--bake-steps", type=int, nargs="+", default=[16, 64], help="Bake depths K to evaluate")
+    parser.add_argument("--overlaps", type=int, nargs="+", default=[0, 10], help="Prompt-overlap windows W_pre")
+    parser.add_argument("--out-json", type=str, default=None, help="Optional path to save evaluation JSON")
+    args = parser.parse_args()
+
+    if args.device == "auto":
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    else:
+        device = torch.device(args.device)
+
+    model, cfg, ckpt_meta = load_checkpoint(args.ckpt, device=device)
+    _, val_ds, tokenizer = build_dataset_and_tokenizer(cfg)
+    val_loader = DataLoader(
+        val_ds,
+        batch_size=cfg.train.batch_size,
+        shuffle=False,
+        collate_fn=lambda b: collate_conveyer_batch(b, pad_id=tokenizer.pad_id),
+    )
+
+    step_num = ckpt_meta.get("step", "?")
+    elapsed_h = float(ckpt_meta.get("elapsed_sec", 0.0)) / 3600.0
+
+    cryst = evaluate_crystallization_curve(model, cfg, val_loader, device, pad_id=tokenizer.pad_id)
+    acc_c = cryst["per_slot_acc"]
+    ent_c = cryst["per_slot_entropy"]
+    B = len(acc_c)
+
+    print("=" * 92)
+    print(
+        f"CONVEYOR-BELT DIFFUSION LLM BENCHMARK SUITE | Ckpt: {args.ckpt} (step {step_num}, {elapsed_h:.2f}h) "
+        f"| Params: {model.num_parameters()/1e6:.2f}M | Device: {device}"
+    )
+    print("=" * 92)
+    print(
+        f"[1] Held-Out Real TinyStories LM Metrics (10,000 unseen stories, 0% templates):\n"
+        f"    - Conveyor Slot-0 Perplexity : {cryst['slot0_ppl']:6.2f}  | Top-1 Next-Token Acc: {cryst['slot0_acc']*100:5.2f}%\n"
+        f"    - Causal Prefix LM Perplexity: {cryst['causal_lm_ppl']:6.2f}  | Top-1 Next-Token Acc: {cryst['causal_lm_acc']*100:5.2f}%\n"
+        f"    - Back-75% Planning Pause    : {cryst['mean_plan_pause']*100:5.1f}%"
+    )
+    print("-" * 92)
+    print("[2] Per-Slot Crystallization Curve along the 64-Slot Oven Belt (Back k=63 -> Front k=0):")
+    sample_slots = [B - 1, (3 * B) // 4, B // 2, B // 4, min(8, B - 1), min(3, B - 1), 0]
+    seen: set[int] = set()
+    for k in sample_slots:
+        if k in seen:
+            continue
+        seen.add(k)
+        zone = "Back (75% Plan)   " if k >= 16 else "Front (25% Stream)"
+        print(f"    Slot {k:2d} [{zone}] | Top-1 Acc: {acc_c[k]*100:6.2f}% | Entropy: {ent_c[k]:.3f} nats")
+
+    if cfg.data.dataset == "tinystories":
+        cs = evaluate_zero_shot_commonsense(model, tokenizer, device)
+        gen_q = evaluate_prefix_sensitivity_and_generation_quality(model, tokenizer, device)
+        print("-" * 92)
+        print(
+            f"[3] Zero-Shot StoryCloze / Commonsense 2-Way Multiple Choice ({cs['num_pairs']} unseen pairs, random=50%):\n"
+            f"    - Causal LM Choice Accuracy  : {cs['commonsense_lm_acc']*100:6.2f}% (mean log-prob margin: {cs['mean_logprob_margin']:+.3f} nats)\n"
+            f"    - Conveyor Slot-0 Choice Acc : {cs['commonsense_slot0_acc']*100:6.2f}%"
+        )
+        print("-" * 92)
+        print(
+            f"[4] Prompt Responsiveness & Open-Ended Generation Quality (6 novel story prompts):\n"
+            f"    - Word-Order Sensitivity KL  : {gen_q['prefix_word_order_kl']:6.3f} nats (broken model was 0.032; >0.50 = strong)\n"
+            f"    - Prompt Entity Retention    : {gen_q['prompt_entity_retention']*100:6.1f}%\n"
+            f"    - Lexical Diversity          : Distinct-2 = {gen_q['distinct_2']*100:5.1f}% | Distinct-3 = {gen_q['distinct_3']*100:5.1f}%\n"
+            f"    - 4-gram Repetition Rate     : {gen_q['repeat_4gram_rate']*100:5.2f}% (lower is better; 0% = zero loops)"
+        )
+        print("-" * 92)
+        print("[5] Live Novel Prompt Continutations:")
+        for s in gen_q["samples"][:3]:
+            print(f"    PROMPT : {s['prompt']}")
+            print(f"    OUTPUT : {s['continuation']}\n")
+    else:
+        cs = {}
+        gen_q = {}
+
+    print("-" * 92)
+    print(f"{'Bake Depth K':<15} | {'Prompt Overlap W':<18} | {'Pos-0 Acc':<12} | {'Token Acc':<12} | {'Exact Seq Acc':<14}")
+    print("-" * 92)
+
+    sweep_results: list[dict[str, Any]] = []
+    for k_bake in args.bake_steps:
+        for w_ov in args.overlaps:
+            res = evaluate_closed_loop_generation(
+                model=model,
+                val_ds=val_ds,
+                tokenizer=tokenizer,
+                device=device,
+                num_episodes=args.episodes,
+                prompt_overlap=w_ov,
+                bake_steps=k_bake,
+            )
+            print(
+                f"K = {k_bake:<11d} | W_pre = {w_ov:<10d} | "
+                f"{res['first_token_acc']*100:9.2f}%  | {res['token_acc']*100:9.2f}%  | {res['exact_seq_acc']*100:10.2f}%"
+            )
+            sweep_results.append({"bake_steps": k_bake, "prompt_overlap": w_ov, **res})
+
+    if args.out_json:
+        Path(args.out_json).parent.mkdir(parents=True, exist_ok=True)
+        Path(args.out_json).write_text(
+            json.dumps(
+                {
+                    "step": step_num,
+                    "crystallization": cryst,
+                    "commonsense": cs,
+                    "generation_quality": gen_q,
+                    "closed_loop": sweep_results,
+                },
+                indent=2,
+            ),
+            encoding="utf-8",
+        )
+
+
+if __name__ == "__main__":
+    main()

src/conveyer/infer.py

@@ -0,0 +1,319 @@
+from __future__ import annotations
+
+import argparse
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterator, Optional
+import torch
+import torch.nn.functional as F
+
+from conveyer.model import BeltState, BeltStepOutput, ConveyerDiffusionLM, PrefixKVCache
+
+
+@dataclass
+class TickSnapshot:
+    """Complete telemetry of a single conveyor-belt oven tick (used for live CLI & Web UI)."""
+    tick: int                         # absolute sequence index at Slot 0
+    is_warmup: bool                   # True during prompt-overlap pre-baking (t < prompt_len)
+    emitted_token: int                # token popped out at Slot 0 (clamped if is_warmup)
+    predicted_slot0_token: int        # model's raw top-1 prediction at Slot 0
+    slot_top1_tokens: list[int]       # (B,) top-1 token ID at each slot k=0..B-1
+    slot_top1_probs: list[float]      # (B,) top-1 probability at each slot k=0..B-1
+    slot_entropies: list[float]       # (B,) Shannon entropy at each slot k=0..B-1
+    slot_gate_means: list[float]      # (B,) mean thought-highway gate activation g_k
+    slot_shift_probs: list[float]     # (B,) elastic shift probability s_k (1.0 in front, <1 in back 75% planning)
+    slot_top3_tokens: list[list[int]] # (B, 3) top-3 candidate token IDs per slot
+    slot_top3_probs: list[list[float]]# (B, 3) top-3 candidate probabilities per slot
+
+
+@torch.no_grad()
+def sample_slot0_token(
+    logits: torch.Tensor,
+    temperature: float = 0.0,
+    recent_tokens: Optional[list[int]] = None,
+    repetition_penalty: float = 1.25,
+) -> torch.Tensor:
+    """
+    logits: (N, V)
+    Masks out special tokens (<PAD>=0, <BOS>=1, <UNK>=3) and applies mild repetition penalty
+    on recently emitted tokens before greedy or top-k sampling.
+    """
+    adj_logits = logits.clone()
+    if adj_logits.size(-1) > 4:
+        adj_logits[:, 0] = -1e9
+        adj_logits[:, 1] = -1e9
+        adj_logits[:, 3] = -1e9
+
+    if recent_tokens and repetition_penalty > 1.0:
+        for tid in set(recent_tokens[-6:]):
+            if 0 <= tid < adj_logits.size(-1):
+                val = adj_logits[:, tid]
+                adj_logits[:, tid] = torch.where(val > 0, val / repetition_penalty, val * repetition_penalty)
+        # Only block 3-in-a-row identical token stammers (so legitimate words after prompt or across clauses are never clipped!)
+        if len(recent_tokens) >= 2 and recent_tokens[-1] == recent_tokens[-2]:
+            last_tid = recent_tokens[-1]
+            if 0 <= last_tid < adj_logits.size(-1):
+                adj_logits[:, last_tid] -= 3.0
+
+    if temperature <= 1e-5:
+        return adj_logits.argmax(dim=-1)
+    scaled = adj_logits / temperature
+    k = min(24, scaled.size(-1))
+    top_vals, _ = torch.topk(scaled, k, dim=-1)
+    min_top = top_vals[:, -1].unsqueeze(-1)
+    scaled = torch.where(scaled < min_top, torch.full_like(scaled, -1e9), scaled)
+    probs = F.softmax(scaled, dim=-1)
+    return torch.multinomial(probs, num_samples=1).squeeze(-1)
+
+
+@torch.no_grad()
+def stream_conveyer_generation(
+    model: ConveyerDiffusionLM,
+    prompt_tokens: torch.Tensor,
+    max_new_tokens: int = 64,
+    prompt_overlap: int = 12,
+    bake_steps: Optional[int] = None,
+    temperature: float = 0.0,
+    eos_token_id: Optional[int] = None,
+    deterministic_spawn: bool = False,
+) -> Iterator[tuple[torch.Tensor, TickSnapshot]]:
+    """
+    Streams generation tick-by-tick using the `n - W_pre` Prompt-Overlap Pre-Baking protocol:
+      - Input at every tick t : [Committed Prefix x_{0..t-1}] + [64-Slot Conveyor/Planning Belt]
+      - Output at every tick t: [Popped-out Confirmed Token x_t at Slot 0] + [Updated 64-Slot Belt]
+    """
+    model.eval()
+    if prompt_tokens.dim() == 1:
+        prompt_tokens = prompt_tokens.unsqueeze(0)
+
+    bsz, n = prompt_tokens.shape
+    device = prompt_tokens.device
+    B = min(bake_steps, model.belt_size) if bake_steps is not None else model.belt_size
+    B = max(1, B)
+
+    max_overlap = min(model.front_slots, B - 1) if B > 1 else 0
+    w_pre = max(0, min(n - 1 if n > 1 else 0, prompt_overlap, max_overlap))
+    t_start = n - w_pre
+
+    # 1. Prefill prefix KV cache on x_{0 : t_start - 1}
+    kv_cache: PrefixKVCache = model.prefill_prefix(prompt_tokens[:, :t_start])
+
+    # 2. Initialize conveyor belt at t_start
+    belt: BeltState = model.init_empty_belt(
+        batch_size=bsz,
+        tick=t_start,
+        device=device,
+        deterministic=deterministic_spawn,
+    )
+    if B < model.belt_size:
+        belt = BeltState(
+            e_soft=belt.e_soft[:, :B, :],
+            m_thought=belt.m_thought[:, :B, :],
+            tick=t_start,
+        )
+
+    offsets = model.get_slot_target_offsets(active_slots=B).to(device)  # (B,)
+
+    # Clamp any belt slots at tick t_start that fall inside the known input prompt (t_start + offsets[k] < n)
+    if w_pre > 0:
+        slot_pos = t_start + offsets  # (B,)
+        in_prompt = slot_pos < n
+        clamped_pos = slot_pos.clamp(min=0, max=n - 1).unsqueeze(0).expand(bsz, -1)
+        prompt_gathered = torch.gather(prompt_tokens, dim=1, index=clamped_pos)
+        overlap_emb = model.tok_emb(prompt_gathered)
+        e_init = torch.where(in_prompt.view(1, B, 1), overlap_emb, belt.e_soft)
+        belt = BeltState(e_soft=e_init, m_thought=belt.m_thought, tick=t_start)
+
+    total_ticks = w_pre + max_new_tokens
+    recent_emitted: list[int] = []
+
+    for step_idx in range(total_ticks):
+        t = t_start + step_idx
+        is_warmup = t < n
+
+        # Single forward pass across all active slots on the belt
+        step_out: BeltStepOutput = model.forward_belt(belt=belt, kv_cache=kv_cache, active_slots=B)
+
+        raw_pred_tok = sample_slot0_token(
+            step_out.slot0_logits,
+            temperature=temperature,
+            recent_tokens=recent_emitted if not is_warmup else None,
+            repetition_penalty=1.08 if temperature > 1e-5 else 1.0,
+        )
+        if is_warmup:
+            emitted_tok = prompt_tokens[:, t]
+        else:
+            emitted_tok = raw_pred_tok
+            recent_emitted.append(int(emitted_tok[0].item()))
+
+        probs0 = step_out.probs[0]  # (B, V)
+        p_clamped = probs0.clamp(min=1e-8)
+        entropies = -(p_clamped * p_clamped.log()).sum(dim=-1)  # (B,)
+        k_top = min(3, model.vocab_size)
+        top_probs, top_ids = probs0.topk(k_top, dim=-1)
+        gate_means = step_out.gate[0].mean(dim=-1)  # (B,)
+        shift_probs = step_out.shift_prob[0]        # (B,)
+
+        snapshot = TickSnapshot(
+            tick=t,
+            is_warmup=is_warmup,
+            emitted_token=int(emitted_tok[0].item()),
+            predicted_slot0_token=int(raw_pred_tok[0].item()),
+            slot_top1_tokens=top_ids[:, 0].tolist(),
+            slot_top1_probs=top_probs[:, 0].tolist(),
+            slot_entropies=entropies.tolist(),
+            slot_gate_means=gate_means.tolist(),
+            slot_shift_probs=shift_probs.tolist(),
+            slot_top3_tokens=top_ids.tolist(),
+            slot_top3_probs=top_probs.tolist(),
+        )
+
+        yield emitted_tok, snapshot
+
+        if (not is_warmup) and eos_token_id is not None and int(emitted_tok[0].item()) == eos_token_id:
+            break
+
+        # Commit popped-out token x_t to frozen prefix KV cache
+        kv_cache = model.append_committed_token(emitted_tok, kv_cache)
+
+        # Shift belt left (with elastic 75% back planning zone) and clamp any remaining prompt tokens at t + 1
+        next_t = t + 1
+        if next_t < n:
+            next_pos = next_t + offsets  # (B,)
+            clamp_mask = (next_pos < n).unsqueeze(0).expand(bsz, -1)  # (N, B)
+            clamped_idx = next_pos.clamp(min=0, max=n - 1).unsqueeze(0).expand(bsz, -1)
+            clamp_tokens = torch.gather(prompt_tokens, dim=1, index=clamped_idx)
+        else:
+            clamp_tokens = None
+            clamp_mask = None
+
+        belt = model.shift_belt_left(
+            step_out=step_out,
+            current_tick=t,
+            deterministic_spawn=deterministic_spawn,
+            clamp_tokens=clamp_tokens,
+            clamp_mask=clamp_mask,
+        )
+
+
+@torch.no_grad()
+def generate_conveyer(
+    model: ConveyerDiffusionLM,
+    prompt_tokens: torch.Tensor,
+    max_new_tokens: int = 64,
+    prompt_overlap: int = 12,
+    bake_steps: Optional[int] = None,
+    temperature: float = 0.0,
+    eos_token_id: Optional[int] = None,
+    deterministic_spawn: bool = False,
+) -> tuple[torch.Tensor, list[TickSnapshot]]:
+    if prompt_tokens.dim() == 1:
+        prompt_tokens = prompt_tokens.unsqueeze(0)
+
+    generated: list[torch.Tensor] = []
+    snapshots: list[TickSnapshot] = []
+
+    for tok, snap in stream_conveyer_generation(
+        model=model,
+        prompt_tokens=prompt_tokens,
+        max_new_tokens=max_new_tokens,
+        prompt_overlap=prompt_overlap,
+        bake_steps=bake_steps,
+        temperature=temperature,
+        eos_token_id=eos_token_id,
+        deterministic_spawn=deterministic_spawn,
+    ):
+        snapshots.append(snap)
+        if not snap.is_warmup:
+            generated.append(tok.unsqueeze(1))
+
+    if len(generated) > 0:
+        gen_tensor = torch.cat(generated, dim=1)
+        full_seq = torch.cat([prompt_tokens, gen_tensor], dim=1)
+    else:
+        full_seq = prompt_tokens.clone()
+
+    return full_seq, snapshots
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="CLI Test & Tick-by-Tick Trace for Conveyor-Belt Diffusion LLM")
+    parser.add_argument("--ckpt", type=str, default=None, help="Path to model checkpoint (.pt)")
+    parser.add_argument("--prompt", type=str, default=None, help="Initial text prompt to continue")
+    parser.add_argument("--tokens", type=int, default=32, help="Number of new tokens to bake and emit")
+    parser.add_argument("--overlap", type=int, default=12, help="Prompt-overlap pre-baking window (n - W)")
+    parser.add_argument("--bake-steps", type=int, default=64, help="Active conveyor belt length K (1..64)")
+    parser.add_argument("--temp", type=float, default=0.0, help="Sampling temperature at Slot 0 (0 = greedy, 0.7 = creative)")
+    args = parser.parse_args()
+
+    from conveyer.checkpoint import load_checkpoint
+    from conveyer.data import build_dataset_and_tokenizer
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    ckpt_path = args.ckpt
+    if ckpt_path is None:
+        ckpt_path = (
+            "runs/synthetic/best.pt"
+            if (not Path("runs/tinystories/best.pt").exists() and Path("runs/synthetic/best.pt").exists())
+            else "runs/tinystories/best.pt"
+        )
+
+    model, cfg, _ = load_checkpoint(ckpt_path, device=device)
+    _, val_ds, tokenizer = build_dataset_and_tokenizer(cfg, infer_only=True)
+
+    if args.prompt is not None:
+        # A trailing space becomes a lone ' ' BPE token that never ends a prefix in training.
+        prompt_ids = tokenizer.encode(args.prompt.strip())
+        prompt_tensor = torch.tensor([prompt_ids], dtype=torch.long, device=device)
+    else:
+        sample = val_ds[0]
+        p_len = int(sample["prompt_len"])
+        prompt_tensor = sample["input_ids"][:p_len].unsqueeze(0).to(device)
+
+    prompt_decoded = tokenizer.decode(prompt_tensor[0].tolist())
+    print("=" * 108)
+    print(f"Checkpoint : {ckpt_path} ({model.num_parameters()/1e6:.2f}M params, device={device})")
+    print(f"Input Text : {prompt_decoded}")
+    print("-" * 108)
+    print(
+        f"{'tick':>4} | {'mode':<6} | {'pop(Slot 0)':<14} | {'front conveyor [1..4]':<30} | "
+        f"{'mid expand [15..18]':<24} | {'back 75% plan [58..63] (pause%)'}"
+    )
+    print("-" * 108)
+
+    emitted_ids: list[int] = []
+    for tok, snap in stream_conveyer_generation(
+        model=model,
+        prompt_tokens=prompt_tensor,
+        max_new_tokens=args.tokens,
+        prompt_overlap=args.overlap,
+        bake_steps=args.bake_steps,
+        temperature=args.temp,
+        eos_token_id=tokenizer.eos_id,
+    ):
+        if not snap.is_warmup:
+            emitted_ids.append(int(tok[0].item()))
+        mode_str = "WARMUP" if snap.is_warmup else "EMIT  "
+        pop_str = repr(tokenizer.id_to_piece(snap.emitted_token).strip() or tokenizer.id_to_piece(snap.emitted_token))
+        B = len(snap.slot_top1_tokens)
+        front_str = " ".join(tokenizer.id_to_piece(t).strip() for t in snap.slot_top1_tokens[1 : min(5, B)])
+        mid_str = " ".join(tokenizer.id_to_piece(t).strip() for t in snap.slot_top1_tokens[min(15, B - 1) : min(19, B)])
+        back_str = " ".join(tokenizer.id_to_piece(t).strip() for t in snap.slot_top1_tokens[max(0, B - 5) : B])
+        F_slots = min(model.front_slots, B)
+        back_pause = (
+            100.0 * (1.0 - sum(snap.slot_shift_probs[F_slots:]) / max(1, B - F_slots))
+            if B > F_slots
+            else 0.0
+        )
+        print(
+            f"{snap.tick:4d} | {mode_str} | {pop_str:<14} | {front_str[:30]:<30} | "
+            f"{mid_str[:24]:<24} | {back_str[:22]:<22} ({back_pause:2.0f}% paused)"
+        )
+
+    print("-" * 108)
+    print(f"FULL OUTPUT: {prompt_decoded}{tokenizer.decode(emitted_ids)}")
+    print("=" * 108)
+
+
+if __name__ == "__main__":
+    main()

src/conveyer/live.py

@@ -0,0 +1,164 @@
+from __future__ import annotations
+
+import argparse
+import sys
+import time
+import torch
+
+from conveyer.checkpoint import load_checkpoint
+from conveyer.config import load_config
+from conveyer.data import build_dataset_and_tokenizer
+from conveyer.infer import TickSnapshot, stream_conveyer_generation
+from conveyer.model import ConveyerDiffusionLM
+
+RESET = "\033[0m"
+BOLD = "\033[1m"
+DIM = "\033[2m"
+GREEN = "\033[92m"
+CYAN = "\033[96m"
+YELLOW = "\033[93m"
+MAGENTA = "\033[95m"
+GRAY = "\033[90m"
+
+
+def _format_slot_badge(slot_idx: int, piece: str, prob: float) -> str:
+    clean = piece.replace("\n", " ")[:9]
+    if slot_idx == 0:
+        return f"{BOLD}{GREEN}[0:{clean} {prob*100:2.0f}%]{RESET}"
+    elif slot_idx < 15:
+        return f"{GREEN}{clean}({prob*100:2.0f}%){RESET}"
+    elif slot_idx < 40:
+        return f"{YELLOW}{clean}({prob*100:2.0f}%){RESET}"
+    else:
+        return f"{GRAY}{clean}{RESET}"
+
+
+def render_frame(
+    prompt_str: str,
+    committed_pieces: list[str],
+    snap: TickSnapshot,
+    tokenizer: Any,
+    clear_screen: bool = True,
+) -> str:
+    B = len(snap.slot_top1_tokens)
+    lines: list[str] = []
+    if clear_screen:
+        lines.append("\033[2J\033[H")
+    lines.append("=" * 98)
+    mode_tag = (
+        f"{CYAN}PRE-BAKING WARMUP (Clamping Input -> Preparing Conveyor Belt){RESET}"
+        if snap.is_warmup
+        else f"{GREEN}LIVE CONVEYOR BAKING (Serving Slot 0 -> Shifting 1..63 -> Spawning 63){RESET}"
+    )
+    lines.append(f"{BOLD}64-SLOT DUAL-CHANNEL CONVEYOR OVEN LLM{RESET} | Tick t={snap.tick:3d} | {mode_tag}")
+    lines.append("-" * 98)
+    lines.append(f"{BOLD}Input Prompt :{RESET} {prompt_str}")
+    committed_str = " ".join(committed_pieces) if committed_pieces else "(warming up belt...)"
+    lines.append(f"{BOLD}Served Output:{RESET} {BOLD}{GREEN}{committed_str}{RESET}")
+    lines.append("-" * 98)
+
+    # Zone 1: Front of Belt (0 .. 14) - Final Grammatical Polish
+    front_items = [
+        _format_slot_badge(k, tokenizer.id_to_piece(snap.slot_top1_tokens[k]), snap.slot_top1_probs[k])
+        for k in range(min(15, B))
+    ]
+    lines.append(f"{BOLD}{GREEN}FRONT [00..14] (Final Polish, alpha->0):{RESET}")
+    lines.append("  <- POP " + " | ".join(front_items[:8]))
+    if len(front_items) > 8:
+        lines.append("         " + " | ".join(front_items[8:15]))
+
+    # Zone 2: Middle of Belt (15 .. 39) - Firming Up Phrasing
+    if B > 15:
+        mid_indices = list(range(15, min(40, B), 2))
+        mid_items = [
+            _format_slot_badge(k, tokenizer.id_to_piece(snap.slot_top1_tokens[k]), snap.slot_top1_probs[k])
+            for k in mid_indices[:10]
+        ]
+        lines.append(f"{BOLD}{YELLOW}MIDDLE [15..39] (Firming Phrasing, Dual-Channel):{RESET}")
+        lines.append("         " + " | ".join(mid_items))
+
+    # Zone 3: Back of Belt (40 .. 63) - Rough Ideas & Fresh Raw Dough
+    if B > 40:
+        back_indices = list(range(40, B, 2))
+        if (B - 1) not in back_indices:
+            back_indices.append(B - 1)
+        back_items = [
+            f"[{k}:{_format_slot_badge(k, tokenizer.id_to_piece(snap.slot_top1_tokens[k]), snap.slot_top1_probs[k])}]"
+            for k in back_indices[-8:]
+        ]
+        lines.append(f"{BOLD}{MAGENTA}BACK   [40..63] (Rough Ideas <- Slot 63 Raw Dough Spawn):{RESET}")
+        lines.append("         " + " ".join(back_items) + " <- [NEW DOUGH]")
+
+    # Telemetry bar: Entropy & Gate Retention
+    ent_front = sum(snap.slot_entropies[: min(8, B)]) / max(1, min(8, B))
+    ent_mid = sum(snap.slot_entropies[min(15, B - 1) : min(40, B)]) / max(1, min(40, B) - min(15, B - 1))
+    ent_back = sum(snap.slot_entropies[max(0, B - 16) : B]) / max(1, B - max(0, B - 16))
+    gate_avg = sum(snap.slot_gate_means) / max(1, B)
+    lines.append("-" * 98)
+    lines.append(
+        f"Entropy (Front/Mid/Back): {ent_front:.2f} / {ent_mid:.2f} / {ent_back:.2f} nats | "
+        f"Mean Thought Gate g_k: {gate_avg:.3f}"
+    )
+    lines.append("=" * 98)
+    return "\n".join(lines)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Live Terminal Visualizer for the 64-Slot Conveyor-Belt Diffusion LLM")
+    parser.add_argument("--ckpt", type=str, default=None, help="Checkpoint path (.pt). If omitted, initializes from config")
+    parser.add_argument("--config", type=str, default="configs/synthetic.yaml", help="Config path if --ckpt is not provided")
+    parser.add_argument("--prompt", type=str, default=None, help="Custom input prompt")
+    parser.add_argument("--tokens", type=int, default=24, help="Number of output tokens to bake")
+    parser.add_argument("--overlap", type=int, default=12, help="Prompt-overlap pre-baking window (n - W)")
+    parser.add_argument("--bake-steps", type=int, default=64, help="Conveyor belt size (1..64)")
+    parser.add_argument("--hz", type=float, default=10.0, help="Refresh rate in ticks/sec (use 0 for no delay)")
+    parser.add_argument("--no-clear", action="store_true", help="Print frames sequentially without ANSI screen clear")
+    args = parser.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    if args.ckpt is not None:
+        model, cfg, _ = load_checkpoint(args.ckpt, device=device)
+    else:
+        cfg = load_config(args.config)
+        model = ConveyerDiffusionLM(cfg.model).to(device)
+
+    _, val_ds, tokenizer = build_dataset_and_tokenizer(cfg)
+
+    if args.prompt is not None:
+        prompt_ids = tokenizer.encode(args.prompt)
+        prompt_tensor = torch.tensor([prompt_ids], dtype=torch.long, device=device)
+    else:
+        sample = val_ds[0]
+        p_len = int(sample["prompt_len"])
+        prompt_tensor = sample["input_ids"][:p_len].unsqueeze(0).to(device)
+
+    prompt_str = tokenizer.decode(prompt_tensor[0].tolist())
+    committed: list[str] = []
+    delay = (1.0 / args.hz) if args.hz > 0 else 0.0
+
+    for tok, snap in stream_conveyer_generation(
+        model=model,
+        prompt_tokens=prompt_tensor,
+        max_new_tokens=args.tokens,
+        prompt_overlap=args.overlap,
+        bake_steps=args.bake_steps,
+        temperature=0.0,
+        eos_token_id=tokenizer.eos_id,
+    ):
+        if not snap.is_warmup:
+            committed.append(tokenizer.id_to_piece(int(tok[0].item())))
+        frame = render_frame(
+            prompt_str=prompt_str,
+            committed_pieces=committed,
+            snap=snap,
+            tokenizer=tokenizer,
+            clear_screen=(not args.no_clear),
+        )
+        sys.stdout.write(frame + "\n")
+        sys.stdout.flush()
+        if delay > 0:
+            time.sleep(delay)
+
+
+if __name__ == "__main__":
+    main()

src/conveyer/losses.py

@@ -0,0 +1,406 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+import torch
+import torch.nn.functional as F
+
+from conveyer.config import ConveyerConfig
+from conveyer.model import BeltState, ConveyerDiffusionLM, PrefixKVCache
+
+
+@dataclass
+class UnrollMetrics:
+    loss: torch.Tensor
+    slot0_loss: float
+    slot0_acc: float
+    belt_mean_loss: float
+    belt_mean_acc: float
+    per_slot_acc: torch.Tensor      # (B,) accuracy at each belt position k=0..B-1
+    per_slot_entropy: torch.Tensor  # (B,) entropy at each belt position k=0..B-1
+    mean_gate: float
+    mean_plan_pause: float = 0.0    # Mean pause rate (1 - s_k) on the back 75% planning belt
+    prefix_lm_loss: float = 0.0     # Causal next-token CE loss across the prompt/full sequence
+    prefix_lm_acc: float = 0.0      # Causal next-token accuracy across the prompt/full sequence
+
+
+def build_horizon_weights(
+    belt_size: int,
+    w_min: float,
+    w_beta: float,
+    device: torch.device,
+) -> torch.Tensor:
+    k = torch.arange(belt_size, device=device, dtype=torch.float32)
+    ratio = k / max(belt_size - 1, 1)
+    w = w_min + (1.0 - w_min) * ((1.0 - ratio) ** w_beta)
+    return w / w.sum()
+
+
+def gather_belt_targets(
+    seq_tokens: torch.Tensor,
+    tick: int,
+    offsets: torch.Tensor,
+    pad_token_id: int = 0,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Gathers target tokens `x_{tick + offsets[k]}` for each belt slot `k = 0..B-1`:
+      - Front 25% Conveyor Slots (k = 0 .. F-1): `offsets[k] = k` (immediate token stream `t + k`)
+      - Back 75% Planning Slots (k = F .. B-1): `offsets[k] = F + (k - F) * stride` (longer-range sentence plan!)
+    Returns `(targets, in_bounds_mask)` of shape `(N, B)`.
+    """
+    bsz, seq_len = seq_tokens.shape
+    device = seq_tokens.device
+    pos = tick + offsets.to(device)  # (B,)
+    in_bounds = (pos >= 0) & (pos < seq_len)  # (B,)
+    clamped_pos = pos.clamp(min=0, max=seq_len - 1).unsqueeze(0).expand(bsz, -1)  # (N, B)
+    gathered = torch.gather(seq_tokens, dim=1, index=clamped_pos)
+    pad_tensor = torch.full_like(gathered, pad_token_id)
+    targets = torch.where(in_bounds.unsqueeze(0), gathered, pad_tensor)
+    return targets, in_bounds.unsqueeze(0).expand(bsz, -1)
+
+
+def seed_corrupted_belt(
+    model: ConveyerDiffusionLM,
+    target_tokens: torch.Tensor,
+    tick_start: int,
+    sigma_max: float = 3.5,
+    sigma_eta: float = 1.25,
+    cold_drop_prob: float = 0.85,
+    prompt_lengths: Optional[torch.Tensor] = None,
+) -> BeltState:
+    bsz, B = target_tokens.shape
+    device = target_tokens.device
+    V = model.vocab_size
+
+    k = torch.arange(B, device=device, dtype=torch.float32)
+    ratio = k / max(B - 1, 1)
+
+    signal_strength = 3.0 * ((1.0 - ratio) ** (sigma_eta * 2.0))
+    signal_strength[-1] = 0.0
+    sigma_k = 0.5 + sigma_max * (ratio ** sigma_eta)
+
+    one_hot = F.one_hot(target_tokens.clamp(min=0, max=V - 1), num_classes=V).float()
+    noise = torch.randn((bsz, B, V), device=device, dtype=torch.float32)
+    tau = model.tau_schedule[:B].view(1, B, 1)
+
+    noisy_logits = signal_strength.view(1, B, 1) * one_hot + sigma_k.view(1, B, 1) * noise
+    tilde_p = F.softmax(noisy_logits / tau, dim=-1).type_as(model.tok_emb.weight)
+    e_soft = torch.matmul(tilde_p, model.tok_emb.weight)
+
+    _, m_thought = model.spawn_raw_dough(batch_size=bsz, num_slots=B, device=device)
+
+    effective_cold_prob = 1.0 if (not model.training and cold_drop_prob > 0.0) else cold_drop_prob
+    if effective_cold_prob > 0.0:
+        if effective_cold_prob >= 1.0:
+            cold_mask = torch.ones((bsz, 1, 1), device=device, dtype=torch.bool)
+        else:
+            cold_mask = torch.rand((bsz, 1, 1), device=device) < effective_cold_prob
+        e_cold, m_cold = model.spawn_raw_dough(batch_size=bsz, num_slots=B, device=device)
+        e_soft = torch.where(cold_mask, e_cold, e_soft)
+        m_thought = torch.where(cold_mask, m_cold, m_thought)
+        uniform_p = torch.full_like(tilde_p, 1.0 / V)
+        tilde_p = torch.where(cold_mask, uniform_p, tilde_p)
+
+    if prompt_lengths is not None:
+        offsets = model.get_slot_target_offsets(active_slots=B).to(device)
+        pos = tick_start + offsets.unsqueeze(0)  # (1, B)
+        is_prompt = pos < prompt_lengths.unsqueeze(1)  # (N, B)
+        prompt_emb = model.tok_emb(target_tokens.clamp(min=0, max=V - 1))
+        e_soft = torch.where(is_prompt.unsqueeze(-1), prompt_emb, e_soft)
+
+    return BeltState(
+        e_soft=e_soft,
+        m_thought=m_thought,
+        tick=tick_start,
+        probs=tilde_p,
+        logits=noisy_logits,
+    )
+
+
+def apply_belt_recovery_perturbation(
+    model: ConveyerDiffusionLM,
+    belt: BeltState,
+    corrupt_rate: float = 0.20,
+    prompt_mask: Optional[torch.Tensor] = None,
+) -> BeltState:
+    """
+    Belt Recovery Training:
+    Randomly corrupts a subset of slots on the active conveyor belt mid-unroll (replacing their
+    soft embeddings with random vocabulary embeddings or noise), forcing the next forward pass
+    to actively repair and re-synchronize the belt from surrounding slots and the prefix!
+    """
+    if corrupt_rate <= 0.0 or not model.training:
+        return belt
+    bsz, B, _ = belt.e_soft.shape
+    device = belt.e_soft.device
+    # Select ~corrupt_rate of slots per batch item to perturb
+    slot_drop = torch.rand((bsz, B), device=device) < corrupt_rate
+    if prompt_mask is not None:
+        slot_drop = slot_drop & (~prompt_mask)
+
+    rand_toks = torch.randint(0, model.vocab_size, (bsz, B), device=device)
+    rand_emb = model.tok_emb(rand_toks)
+    e_noise, m_noise = model.spawn_raw_dough(batch_size=bsz, num_slots=B, device=device)
+    corrupted_e = 0.5 * rand_emb + 0.5 * e_noise
+
+    mask_exp = slot_drop.unsqueeze(-1)
+    e_repaired = torch.where(mask_exp, corrupted_e, belt.e_soft)
+    m_repaired = torch.where(mask_exp, m_noise, belt.m_thought)
+    return BeltState(
+        e_soft=e_repaired,
+        m_thought=m_repaired,
+        tick=belt.tick,
+        probs=belt.probs,
+        logits=belt.logits,
+        gates=belt.gates,
+        shift_probs=belt.shift_probs,
+    )
+
+
+def compute_unrolled_conveyer_loss(
+    model: ConveyerDiffusionLM,
+    seq_tokens: torch.Tensor,
+    cfg: ConveyerConfig,
+    prompt_lengths: Optional[torch.Tensor] = None,
+    tick_start: Optional[int] = None,
+    pad_token_id: int = 0,
+) -> UnrollMetrics:
+    bsz, seq_len = seq_tokens.shape
+    device = seq_tokens.device
+    B = model.belt_size
+    F_slots = min(model.front_slots, B)
+    U = min(cfg.train.unroll_steps, max(1, seq_len - F_slots))
+    offsets = model.get_slot_target_offsets(active_slots=B)
+
+    if tick_start is None:
+        max_t0 = max(1, seq_len - F_slots - U + 1)
+        if prompt_lengths is not None:
+            min_prompt = int(prompt_lengths.min().item())
+            max_prompt = int(prompt_lengths.max().item())
+            low = max(1, min_prompt - min(U - 2, cfg.infer.prompt_overlap))
+            if torch.rand((), device=device).item() < 0.5:
+                high = max(low + 1, min(max_t0 + 1, max_prompt + 16))
+            else:
+                high = max(low + 1, max_t0 + 1)
+            tick_start = int(torch.randint(low, high, (1,), device=device).item())
+        else:
+            tick_start = int(torch.randint(1, max_t0 + 1, (1,), device=device).item())
+
+    # 1. Causal pass over the full sequence + extract causal KV cache
+    prefix_lm_weight = getattr(cfg.train, "prefix_lm_weight", 1.0)
+    full_kv: Optional[PrefixKVCache] = None
+    if prefix_lm_weight > 0.0 and seq_len > 1:
+        causal_logits, full_kv = model.forward_causal_and_cache(seq_tokens, prefix_len=seq_len)
+        shift_logits = causal_logits[:, :-1, :].contiguous()
+        shift_targets = seq_tokens[:, 1:].contiguous()
+        prefix_lm_loss = F.cross_entropy(
+            shift_logits.view(-1, model.vocab_size),
+            shift_targets.view(-1),
+            ignore_index=pad_token_id,
+        )
+        with torch.no_grad():
+            valid_tok = shift_targets != pad_token_id
+            n_valid = valid_tok.float().sum().clamp(min=1.0)
+            pred_tok = shift_logits.argmax(dim=-1)
+            prefix_lm_acc_t = ((pred_tok == shift_targets) & valid_tok).float().sum() / n_valid
+        kv_cache = PrefixKVCache(
+            k_cache=[k[:, :, :tick_start, :] for k in full_kv.k_cache],
+            v_cache=[v[:, :, :tick_start, :] for v in full_kv.v_cache],
+            seq_len=tick_start,
+        )
+    else:
+        prefix_tokens = seq_tokens[:, :tick_start]
+        kv_cache = model.prefill_prefix(prefix_tokens)
+        prefix_lm_loss = torch.tensor(0.0, device=device)
+        prefix_lm_acc_t = torch.tensor(0.0, device=device)
+
+    # 2. Seed initial 64-slot belt at tick_start using slot target offsets
+    init_targets, _ = gather_belt_targets(seq_tokens, tick_start, offsets, pad_token_id=pad_token_id)
+
+    belt = seed_corrupted_belt(
+        model=model,
+        target_tokens=init_targets,
+        tick_start=tick_start,
+        sigma_max=cfg.train.seed_sigma_max,
+        sigma_eta=cfg.train.seed_sigma_eta,
+        cold_drop_prob=0.85,
+        prompt_lengths=prompt_lengths,
+    )
+
+    horizon_w = build_horizon_weights(
+        belt_size=B,
+        w_min=cfg.train.loss_w_min,
+        w_beta=cfg.train.loss_w_beta,
+        device=device,
+    )
+
+    total_loss = torch.tensor(0.0, device=device)
+    prev_logits: Optional[torch.Tensor] = None
+    slot0_loss_sum_t = torch.tensor(0.0, device=device)
+    slot0_acc_sum_t = torch.tensor(0.0, device=device)
+    s0_valid_ticks_t = torch.tensor(0.0, device=device)
+    belt_loss_sum_t = torch.tensor(0.0, device=device)
+    belt_acc_sum_t = torch.tensor(0.0, device=device)
+    gate_sum_t = torch.tensor(0.0, device=device)
+    pause_sum_t = torch.tensor(0.0, device=device)
+    valid_ticks = 0
+    dagger_diverged = False
+
+    per_slot_correct = torch.zeros(B, device=device)
+    per_slot_count = torch.zeros(B, device=device)
+    per_slot_entropy_sum = torch.zeros(B, device=device)
+
+    for u in range(U):
+        t = tick_start + u
+        if t >= seq_len:
+            break
+
+        targets, in_bounds = gather_belt_targets(seq_tokens, t, offsets, pad_token_id=pad_token_id)
+        slot_positions = t + offsets.to(device).unsqueeze(0)  # (1, B)
+
+        # Mid-unroll Belt Recovery Training: occasionally inject wrong/noisy thoughts onto the belt
+        # so the next forward pass learns to actively repair the belt!
+        if u > 0 and model.training and cfg.train.belt_corrupt_prob > 0.0:
+            if torch.rand(()) < 0.35:
+                p_mask = (slot_positions < prompt_lengths.unsqueeze(1)) if prompt_lengths is not None else None
+                belt = apply_belt_recovery_perturbation(
+                    model=model,
+                    belt=belt,
+                    corrupt_rate=cfg.train.belt_corrupt_prob,
+                    prompt_mask=p_mask,
+                )
+
+        step_out = model.forward_belt(belt=belt, kv_cache=kv_cache)
+
+        valid_mask = in_bounds & (targets != pad_token_id)  # (N, B)
+        if prompt_lengths is not None:
+            is_gen = slot_positions >= prompt_lengths.unsqueeze(1)
+            valid_mask = valid_mask & is_gen
+
+        if cfg.model.causal_belt_only:
+            ar_mask = torch.zeros_like(valid_mask)
+            ar_mask[:, 0] = valid_mask[:, 0]
+            valid_mask = ar_mask
+
+        ce_per_token = F.cross_entropy(
+            step_out.logits.reshape(-1, model.vocab_size),
+            targets.reshape(-1),
+            reduction="none",
+        ).view(bsz, B)
+
+        w_broadcast = horizon_w.unsqueeze(0) * valid_mask.float()
+        w_denom = w_broadcast.sum().clamp(min=1e-6)
+        tick_ce_loss = (ce_per_token * w_broadcast).sum() / w_denom
+
+        probs_clamped = step_out.probs.clamp(min=1e-8)
+        entropy = -(probs_clamped * probs_clamped.log()).sum(dim=-1)
+        if cfg.train.entropy_reg > 0.0 and not cfg.model.causal_belt_only:
+            back_weights = (torch.arange(B, device=device, dtype=torch.float32) / max(B - 1, 1)).unsqueeze(0)
+            ent_bonus = (entropy * back_weights * valid_mask.float()).sum() / valid_mask.float().sum().clamp(min=1.0)
+            tick_loss = tick_ce_loss - cfg.train.entropy_reg * ent_bonus
+        else:
+            tick_loss = tick_ce_loss
+
+        # Shift-Consistency on the Front Conveyor Stream (0 .. F_slots - 2)
+        if u > 0 and prev_logits is not None and not cfg.model.causal_belt_only and F_slots > 1:
+            p_prev = F.softmax(prev_logits[:, 1:F_slots, :].detach(), dim=-1)
+            log_p_curr = F.log_softmax(step_out.logits[:, : F_slots - 1, :], dim=-1)
+            kl_shift = F.kl_div(log_p_curr, p_prev, reduction="none").sum(dim=-1)
+            cons_mask = valid_mask[:, : F_slots - 1].float()
+            cons_loss = (kl_shift * cons_mask).sum() / cons_mask.sum().clamp(min=1.0)
+            tick_loss = tick_loss + 0.08 * cons_loss
+
+        prev_logits = step_out.logits
+        total_loss = total_loss + tick_loss
+        valid_ticks += 1
+
+        with torch.no_grad():
+            valid_f = valid_mask.float()
+            preds = step_out.logits.argmax(dim=-1)
+            correct = (preds == targets).float() * valid_f
+
+            s0_valid = valid_f[:, 0]
+            s0_sum = s0_valid.sum()
+            has_s0 = (s0_sum > 0.0).float()
+            s0_denom = s0_sum.clamp(min=1.0)
+            slot0_loss_sum_t = slot0_loss_sum_t + has_s0 * ((ce_per_token[:, 0] * s0_valid).sum() / s0_denom)
+            slot0_acc_sum_t = slot0_acc_sum_t + has_s0 * (correct[:, 0].sum() / s0_denom)
+            s0_valid_ticks_t = s0_valid_ticks_t + has_s0
+
+            all_denom = valid_f.sum().clamp(min=1.0)
+            belt_loss_sum_t = belt_loss_sum_t + (ce_per_token * valid_f).sum() / all_denom
+            belt_acc_sum_t = belt_acc_sum_t + correct.sum() / all_denom
+            gate_sum_t = gate_sum_t + step_out.gate.mean()
+            if B > F_slots:
+                pause_sum_t = pause_sum_t + (1.0 - step_out.shift_prob[:, F_slots:]).mean()
+
+            per_slot_correct += correct.sum(dim=0)
+            per_slot_count += valid_f.sum(dim=0)
+            per_slot_entropy_sum += entropy.mean(dim=0)
+
+        if u < U - 1 and (t + 1) < seq_len:
+            # Prefix Recovery (DAgger-style Scheduled Sampling on final ticks, or fast slice from full_kv):
+            use_dagger = (
+                model.training
+                and (u >= max(0, U - 4) or full_kv is None)
+                and (torch.rand(()) < cfg.train.recovery_prob)
+            )
+            if use_dagger:
+                pred_s0 = step_out.slot0_logits.argmax(dim=-1)
+                if prompt_lengths is not None:
+                    is_p = t < prompt_lengths
+                    committed_x_t = torch.where(is_p, seq_tokens[:, t], pred_s0)
+                else:
+                    committed_x_t = pred_s0
+                dagger_diverged = True
+
+            if not dagger_diverged and full_kv is not None:
+                next_len = t + 1
+                kv_cache = PrefixKVCache(
+                    k_cache=[k[:, :, :next_len, :] for k in full_kv.k_cache],
+                    v_cache=[v[:, :, :next_len, :] for v in full_kv.v_cache],
+                    seq_len=next_len,
+                )
+            else:
+                if not use_dagger:
+                    committed_x_t = seq_tokens[:, t]
+                kv_cache = model.append_committed_token(committed_x_t, kv_cache)
+
+            next_targets, _ = gather_belt_targets(seq_tokens, t + 1, offsets, pad_token_id=pad_token_id)
+            if prompt_lengths is not None:
+                next_pos = (t + 1) + offsets.to(device).unsqueeze(0)
+                clamp_mask = next_pos < prompt_lengths.unsqueeze(1)
+            else:
+                clamp_mask = None
+
+            belt = model.shift_belt_left(
+                step_out=step_out,
+                current_tick=t,
+                deterministic_spawn=False,
+                clamp_tokens=next_targets if clamp_mask is not None else None,
+                clamp_mask=clamp_mask,
+            )
+            if cfg.train.no_belt_bptt:
+                belt = belt.detach()
+                kv_cache = kv_cache.detach()
+
+    denom = max(valid_ticks, 1)
+    s0_denom_ticks = float(s0_valid_ticks_t.item()) if float(s0_valid_ticks_t.item()) > 0.0 else 1.0
+    per_slot_acc = per_slot_correct / per_slot_count.clamp(min=1.0)
+    per_slot_entropy = per_slot_entropy_sum / denom
+    conveyer_loss = total_loss / denom
+    combined_loss = conveyer_loss + prefix_lm_weight * prefix_lm_loss
+
+    return UnrollMetrics(
+        loss=combined_loss,
+        slot0_loss=float(slot0_loss_sum_t.item()) / s0_denom_ticks,
+        slot0_acc=float(slot0_acc_sum_t.item()) / s0_denom_ticks,
+        belt_mean_loss=float(belt_loss_sum_t.item()) / denom,
+        belt_mean_acc=float(belt_acc_sum_t.item()) / denom,
+        per_slot_acc=per_slot_acc,
+        per_slot_entropy=per_slot_entropy,
+        mean_gate=float(gate_sum_t.item()) / denom,
+        mean_plan_pause=float(pause_sum_t.item()) / denom,
+        prefix_lm_loss=float(prefix_lm_loss.item()),
+        prefix_lm_acc=float(prefix_lm_acc_t.item()),
+    )

src/conveyer/model.py

@@ -0,0 +1,739 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint
+
+from conveyer.config import ModelConfig
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, d_model: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(d_model))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        norm = x.float().pow(2).mean(dim=-1, keepdim=True)
+        x_normed = x.float() * torch.rsqrt(norm + self.eps)
+        return (self.weight * x_normed).type_as(x)
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, head_dim: int, max_seq_len: int = 2048, theta: float = 10000.0) -> None:
+        super().__init__()
+        assert head_dim % 2 == 0, "head_dim must be even for RoPE"
+        self.head_dim = head_dim
+        self.max_seq_len = max_seq_len
+        self.theta = theta
+        inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self._build_cache(max_seq_len)
+
+    def _build_cache(self, seq_len: int) -> None:
+        self.max_seq_len = max(self.max_seq_len, seq_len)
+        t = torch.arange(self.max_seq_len, device=self.inv_freq.device, dtype=torch.float32)
+        freqs = torch.outer(t, self.inv_freq)  # (seq_len, head_dim // 2)
+        cos = torch.cos(freqs)
+        sin = torch.sin(freqs)
+        self.register_buffer("cos_cached", cos, persistent=False)
+        self.register_buffer("sin_cached", sin, persistent=False)
+
+    def get_cos_sin(self, positions: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """positions: (L,) or (N, L) integer positions."""
+        max_pos = int(positions.max().item()) + 1 if positions.numel() > 0 else 1
+        if max_pos > self.cos_cached.size(0):
+            self._build_cache(max_pos + 256)
+        cos = self.cos_cached[positions]
+        sin = self.sin_cached[positions]
+        return cos, sin
+
+
+def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    """
+    x: (N, H, L, D)
+    cos, sin: (L, D//2) or (N, L, D//2)
+    """
+    if cos.dim() == 2:
+        cos = cos.unsqueeze(0).unsqueeze(0)  # (1, 1, L, D//2)
+        sin = sin.unsqueeze(0).unsqueeze(0)
+    elif cos.dim() == 3:
+        cos = cos.unsqueeze(1)  # (N, 1, L, D//2)
+        sin = sin.unsqueeze(1)
+    d2 = x.shape[-1] // 2
+    x1 = x[..., :d2]
+    x2 = x[..., d2:]
+    out1 = x1 * cos - x2 * sin
+    out2 = x1 * sin + x2 * cos
+    return torch.cat([out1, out2], dim=-1).type_as(x)
+
+
+class SwiGLUFFN(nn.Module):
+    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.0) -> None:
+        super().__init__()
+        self.w_gate = nn.Linear(d_model, d_ff, bias=False)
+        self.w_up = nn.Linear(d_model, d_ff, bias=False)
+        self.w_down = nn.Linear(d_ff, d_model, bias=False)
+        self.drop = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.drop(self.w_down(F.silu(self.w_gate(x)) * self.w_up(x)))
+
+
+@dataclass
+class PrefixKVCache:
+    """Stores committed prefix Keys and Values per layer: list of (K, V) each (N, H, P, D_h)."""
+    k_cache: list[torch.Tensor]
+    v_cache: list[torch.Tensor]
+    seq_len: int
+
+    @classmethod
+    def empty(cls, n_layers: int, batch_size: int, n_heads: int, head_dim: int, device: torch.device, dtype: torch.dtype) -> "PrefixKVCache":
+        k_cache = [torch.empty((batch_size, n_heads, 0, head_dim), device=device, dtype=dtype) for _ in range(n_layers)]
+        v_cache = [torch.empty((batch_size, n_heads, 0, head_dim), device=device, dtype=dtype) for _ in range(n_layers)]
+        return cls(k_cache=k_cache, v_cache=v_cache, seq_len=0)
+
+    def clone(self) -> "PrefixKVCache":
+        return PrefixKVCache(
+            k_cache=[k.clone() for k in self.k_cache],
+            v_cache=[v.clone() for v in self.v_cache],
+            seq_len=self.seq_len,
+        )
+
+    def detach(self) -> "PrefixKVCache":
+        return PrefixKVCache(
+            k_cache=[k.detach() for k in self.k_cache],
+            v_cache=[v.detach() for v in self.v_cache],
+            seq_len=self.seq_len,
+        )
+
+
+@dataclass
+class BeltState:
+    """Holds the 64-slot Dual-Channel + Elastic Planning dough state on the conveyor belt."""
+    e_soft: torch.Tensor      # (N, B, D) lexical anchor soft embeddings
+    m_thought: torch.Tensor   # (N, B, D) continuous thought highway vectors
+    tick: int                 # current absolute sequence position of Slot 0
+    probs: Optional[torch.Tensor] = None       # (N, B, V) last probability distributions
+    logits: Optional[torch.Tensor] = None      # (N, B, V) last raw logits
+    gates: Optional[torch.Tensor] = None       # (N, B, D) last update gate activations
+    shift_probs: Optional[torch.Tensor] = None # (N, B) elastic shift vs. pause probability s_k
+
+    def detach(self) -> "BeltState":
+        return BeltState(
+            e_soft=self.e_soft.detach(),
+            m_thought=self.m_thought.detach(),
+            tick=self.tick,
+            probs=self.probs.detach() if self.probs is not None else None,
+            logits=self.logits.detach() if self.logits is not None else None,
+            gates=self.gates.detach() if self.gates is not None else None,
+            shift_probs=self.shift_probs.detach() if self.shift_probs is not None else None,
+        )
+
+
+@dataclass
+class BeltStepOutput:
+    logits: torch.Tensor          # (N, B, V) raw vocabulary logits at all 64 slots
+    probs: torch.Tensor           # (N, B, V) temperature-scaled probabilities p_k
+    e_soft: torch.Tensor          # (N, B, D) expected vocab embeddings p_k^T E_vocab
+    cand_thought: torch.Tensor    # (N, B, D) candidate thought vectors c_k
+    gate: torch.Tensor            # (N, B, D) update gate g_k
+    m_updated: torch.Tensor       # (N, B, D) updated thought vectors before index shift
+    shift_prob: torch.Tensor      # (N, B) elastic shift probability s_k (1.0 in front conveyor, <1 in back 75% planning)
+    expand_cand: torch.Tensor     # (N, D) mid-belt expansion thought vector for forming new sentence clauses
+    slot0_logits: torch.Tensor    # (N, V) crisp logits at Position 0
+
+
+class ConveyerAttention(nn.Module):
+    def __init__(self, cfg: ModelConfig) -> None:
+        super().__init__()
+        assert cfg.d_model % cfg.n_heads == 0
+        self.n_heads = cfg.n_heads
+        self.head_dim = cfg.d_model // cfg.n_heads
+        self.d_model = cfg.d_model
+        self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False)
+        self.proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
+        self.dropout = cfg.dropout
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+        layer_kv: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+        return_new_kv: bool = False,
+        is_causal: bool = False,
+    ) -> tuple[torch.Tensor, Optional[tuple[torch.Tensor, torch.Tensor]]]:
+        """
+        x: (N, L, D)
+        layer_kv: optional (K_prev, V_prev) each (N, H, P, D_h)
+        """
+        bsz, seq_len, _ = x.shape
+        qkv = self.qkv(x)
+        q, k, v = qkv.chunk(3, dim=-1)
+        q = q.view(bsz, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+        k = k.view(bsz, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+        v = v.view(bsz, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+
+        q = apply_rope(q, cos, sin)
+        k = apply_rope(k, cos, sin)
+
+        new_kv = (k, v) if return_new_kv else None
+
+        if layer_kv is not None and layer_kv[0].size(2) > 0:
+            k = torch.cat([layer_kv[0].type_as(k), k], dim=2)
+            v = torch.cat([layer_kv[1].type_as(v), v], dim=2)
+
+        out = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=None if is_causal else attn_mask,
+            dropout_p=self.dropout if self.training else 0.0,
+            is_causal=is_causal,
+        )
+        out = out.transpose(1, 2).contiguous().view(bsz, seq_len, self.d_model)
+        return self.proj(out), new_kv
+
+
+class ConveyerBlock(nn.Module):
+    def __init__(self, cfg: ModelConfig) -> None:
+        super().__init__()
+        self.norm1 = RMSNorm(cfg.d_model)
+        self.attn = ConveyerAttention(cfg)
+        self.norm2 = RMSNorm(cfg.d_model)
+        self.ffn = SwiGLUFFN(cfg.d_model, cfg.d_ff, cfg.dropout)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+        layer_kv: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+        return_new_kv: bool = False,
+        is_causal: bool = False,
+    ) -> tuple[torch.Tensor, Optional[tuple[torch.Tensor, torch.Tensor]]]:
+        attn_out, new_kv = self.attn(
+            self.norm1(x),
+            cos=cos,
+            sin=sin,
+            attn_mask=attn_mask,
+            layer_kv=layer_kv,
+            return_new_kv=return_new_kv,
+            is_causal=is_causal,
+        )
+        x = x + attn_out
+        x = x + self.ffn(self.norm2(x))
+        return x, new_kv
+
+
+class ConveyerDiffusionLM(nn.Module):
+    """
+    64-Slot Dual-Channel Elastic Conveyor-Belt Diffusion Language Model.
+
+    Input / Output Contract at every tick t:
+      - Input : [All Committed Text Prefix x_{0..t-1}] + [64-Slot Conveyor/Planning Belt]
+      - Output: [Popped-out Confirmed Token x_t at Slot 0] + [Updated 64-Slot Conveyor/Planning Belt]
+
+    Two-Zone Human-Style Thinking & Speaking Belt:
+      - Front 25% (Slots 0 .. F-1, e.g. 0..15): Immediate Articulation Stream.
+        Shifts left by 1 slot every tick (s_k = 1.0), polishing immediate grammar toward Slot 0 (alpha_0 = 0).
+      - Back 75% (Slots F .. B-1, e.g. 16..63): Elastic Sentence-Planning Horizon.
+        Can pause in-place (s_k < 1.0) to hold long-range sentence goals in mind while new sentence clauses
+        bud and expand in the middle of the belt (around Slot F).
+    """
+
+    def __init__(self, cfg: ModelConfig) -> None:
+        super().__init__()
+        self.cfg = cfg
+        self.belt_size = cfg.belt_size
+        self.front_slots = min(cfg.front_conveyer_slots, cfg.belt_size)
+        self.d_model = cfg.d_model
+        self.vocab_size = cfg.vocab_size
+
+        # Token embedding & weight-tied LM head
+        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+        self.lm_head.weight = self.tok_emb.weight  # weight tying
+
+        # Learned oven slot embeddings for positions k = 0 .. B-1 on the conveyor belt
+        self.slot_emb = nn.Embedding(cfg.belt_size, cfg.d_model)
+        # Learned initial thought vector for newly spawned dough at Slot B-1
+        self.m_init = nn.Parameter(torch.zeros(cfg.d_model))
+
+        # Dough input RMSNorm
+        self.dough_norm = RMSNorm(cfg.d_model)
+
+        # Transformer trunk
+        self.rope = RotaryEmbedding(
+            head_dim=cfg.d_model // cfg.n_heads,
+            max_seq_len=cfg.max_seq_len,
+            theta=cfg.rope_theta,
+        )
+        self.blocks = nn.ModuleList([ConveyerBlock(cfg) for _ in range(cfg.n_layers)])
+        self.final_norm = RMSNorm(cfg.d_model)
+
+        # Channel 2: Continuous Thought Highway heads (RecurrentCubeLM style)
+        self.thought_cand = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
+        self.thought_gate = nn.Linear(cfg.d_model, cfg.d_model, bias=True)
+
+        # Elastic Planning & Mid-Belt Sentence Expansion Heads
+        # shift_head predicts whether each back-75% planning slot shifts left or pauses in place
+        self.shift_head = nn.Linear(cfg.d_model, 1, bias=True)
+        # expand_head synthesizes a new phrase/clause embedding at the boundary between the
+        # paused 75% planning zone and the moving 25% front conveyor
+        self.expand_head = nn.Linear(2 * cfg.d_model, cfg.d_model, bias=False)
+
+        # Precompute positional temperature schedule tau_k, thought mixing alpha_k, and target horizon offsets
+        B = cfg.belt_size
+        F_slots = self.front_slots
+        k_idx = torch.arange(B, dtype=torch.float32)
+        ratio = k_idx / max(B - 1, 1)
+        tau_schedule = cfg.tau_min + (ratio ** cfg.tau_gamma) * (cfg.tau_max - cfg.tau_min)
+        alpha_schedule = ratio.clone()  # alpha_0 = 0.0, alpha_{B-1} = 1.0
+
+        # Target sequence offsets delta_k from current tick t:
+        # - Front slots k = 0 .. F-1 map 1-to-1 to immediate tokens t + k
+        # - Back 75% planning slots k = F .. B-1 span a wider lookahead horizon (stride >= 1)
+        stride = max(1, cfg.plan_horizon_stride) if cfg.use_elastic_planning else 1
+        target_offsets = torch.arange(B, dtype=torch.long)
+        if B > F_slots and stride > 1:
+            back_idx = torch.arange(B - F_slots, dtype=torch.long)
+            target_offsets[F_slots:] = F_slots + back_idx * stride
+
+        # Prior bias for shift_head: front slots always shift (+6.0), back 75% slots default to ~0.35 shift (65% pause)
+        shift_prior = torch.zeros(B, dtype=torch.float32)
+        shift_prior[:F_slots] = 6.0
+        if B > F_slots:
+            back_ratio = torch.linspace(0.0, 1.0, B - F_slots)
+            shift_prior[F_slots:] = -0.3 - 1.2 * back_ratio  # sigmoid(-0.3..-1.5) ~ 0.42 down to 0.18
+
+        self.register_buffer("tau_schedule", tau_schedule, persistent=False)
+        self.register_buffer("alpha_schedule", alpha_schedule, persistent=False)
+        self.register_buffer("target_offsets", target_offsets, persistent=False)
+        self.register_buffer("shift_prior", shift_prior, persistent=False)
+
+        self._init_weights()
+
+    def _init_weights(self) -> None:
+        nn.init.normal_(self.tok_emb.weight, mean=0.0, std=0.02)
+        nn.init.normal_(self.slot_emb.weight, mean=0.0, std=0.02)
+        nn.init.normal_(self.m_init, mean=0.0, std=0.02)
+        for m in self.modules():
+            if isinstance(m, nn.Linear) and m is not self.lm_head:
+                nn.init.normal_(m.weight, mean=0.0, std=0.02)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+        nn.init.constant_(self.thought_gate.bias, self.cfg.gate_bias_init)
+        nn.init.zeros_(self.shift_head.weight)
+        nn.init.zeros_(self.shift_head.bias)
+
+    def num_parameters(self, trainable_only: bool = True) -> int:
+        return sum(p.numel() for p in self.parameters() if (p.requires_grad or not trainable_only))
+
+    def get_slot_target_offsets(self, active_slots: Optional[int] = None) -> torch.Tensor:
+        B = active_slots if active_slots is not None else self.belt_size
+        return self.target_offsets[:B]
+
+    def init_prefix_cache(self, batch_size: int, device: torch.device, dtype: Optional[torch.dtype] = None) -> PrefixKVCache:
+        if dtype is None:
+            dtype = self.tok_emb.weight.dtype
+        return PrefixKVCache.empty(
+            n_layers=self.cfg.n_layers,
+            batch_size=batch_size,
+            n_heads=self.cfg.n_heads,
+            head_dim=self.cfg.d_model // self.cfg.n_heads,
+            device=device,
+            dtype=dtype,
+        )
+
+    def prefill_prefix(
+        self,
+        prefix_tokens: torch.Tensor,
+        kv_cache: Optional[PrefixKVCache] = None,
+    ) -> PrefixKVCache:
+        """
+        Encodes committed prefix tokens x_{t_start .. t_end} into the causal PrefixKVCache.
+        prefix_tokens: (N, P) long tensor.
+        """
+        bsz, plen = prefix_tokens.shape
+        device = prefix_tokens.device
+        if kv_cache is None:
+            kv_cache = self.init_prefix_cache(bsz, device)
+        if plen == 0:
+            return kv_cache
+
+        start_pos = kv_cache.seq_len
+        positions = torch.arange(start_pos, start_pos + plen, device=device, dtype=torch.long)
+        cos, sin = self.rope.get_cos_sin(positions)
+
+        slot0_bias = self.slot_emb.weight[0].view(1, 1, -1)
+        x = self.dough_norm(self.tok_emb(prefix_tokens) + slot0_bias)
+
+        total_k = start_pos + plen
+        use_is_causal = start_pos == 0 and plen > 1
+        if use_is_causal or plen == 1:
+            attn_mask = None
+        else:
+            causal_mask = torch.ones((plen, total_k), device=device, dtype=torch.bool).tril(diagonal=start_pos)
+            attn_mask = causal_mask.unsqueeze(0).unsqueeze(0)
+
+        new_k_list: list[torch.Tensor] = []
+        new_v_list: list[torch.Tensor] = []
+        for i, block in enumerate(self.blocks):
+            layer_kv = (kv_cache.k_cache[i], kv_cache.v_cache[i])
+            x, new_kv = block(
+                x,
+                cos=cos,
+                sin=sin,
+                attn_mask=attn_mask,
+                layer_kv=layer_kv,
+                return_new_kv=True,
+                is_causal=use_is_causal,
+            )
+            assert new_kv is not None
+            new_k_list.append(torch.cat([kv_cache.k_cache[i].type_as(new_kv[0]), new_kv[0]], dim=2))
+            new_v_list.append(torch.cat([kv_cache.v_cache[i].type_as(new_kv[1]), new_kv[1]], dim=2))
+
+        return PrefixKVCache(k_cache=new_k_list, v_cache=new_v_list, seq_len=total_k)
+
+    def forward_causal_and_cache(
+        self,
+        seq_tokens: torch.Tensor,
+        prefix_len: int,
+    ) -> tuple[torch.Tensor, PrefixKVCache]:
+        """
+        Runs a full causal pass over `seq_tokens` (N, L) through the shared transformer blocks:
+          1. Computes next-token logits `(N, L, V)` at every position so all ~160 tokens per story
+             receive direct causal language-modeling supervision (10x more signal per step).
+          2. Extracts the exact causal KV cache for `seq_tokens[:, :prefix_len]` (invariant to tokens >= prefix_len)
+             to seed the conveyor belt unroll without running a duplicate prefill pass.
+        """
+        bsz, seq_len = seq_tokens.shape
+        device = seq_tokens.device
+        positions = torch.arange(0, seq_len, device=device, dtype=torch.long)
+        cos, sin = self.rope.get_cos_sin(positions)
+
+        slot0_bias = self.slot_emb.weight[0].view(1, 1, -1)
+        x = self.dough_norm(self.tok_emb(seq_tokens) + slot0_bias)
+
+        k_list: list[torch.Tensor] = []
+        v_list: list[torch.Tensor] = []
+        for block in self.blocks:
+            x, new_kv = block(
+                x,
+                cos=cos,
+                sin=sin,
+                attn_mask=None,
+                layer_kv=None,
+                return_new_kv=True,
+                is_causal=(seq_len > 1),
+            )
+            assert new_kv is not None
+            k_list.append(new_kv[0][:, :, :prefix_len, :].contiguous())
+            v_list.append(new_kv[1][:, :, :prefix_len, :].contiguous())
+
+        h = self.final_norm(x)
+        causal_logits = self.lm_head(h)
+        kv_cache = PrefixKVCache(k_cache=k_list, v_cache=v_list, seq_len=prefix_len)
+        return causal_logits, kv_cache
+
+    def append_committed_token(
+        self,
+        token_ids: torch.Tensor,
+        kv_cache: PrefixKVCache,
+    ) -> PrefixKVCache:
+        if token_ids.dim() == 1:
+            token_ids = token_ids.unsqueeze(1)
+        return self.prefill_prefix(token_ids, kv_cache=kv_cache)
+
+    def spawn_raw_dough(
+        self,
+        batch_size: int,
+        num_slots: int,
+        device: torch.device,
+        deterministic: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        mean_emb = self.tok_emb.weight.mean(dim=0, keepdim=True).unsqueeze(0)
+        m_base = self.m_init.view(1, 1, -1)
+        if deterministic or self.cfg.spawn_noise_std <= 0.0:
+            e_new = mean_emb.expand(batch_size, num_slots, self.d_model).clone()
+            m_new = m_base.expand(batch_size, num_slots, self.d_model).clone()
+        else:
+            eps_e = torch.randn((batch_size, num_slots, self.d_model), device=device, dtype=mean_emb.dtype)
+            eps_m = torch.randn((batch_size, num_slots, self.d_model), device=device, dtype=mean_emb.dtype)
+            e_new = mean_emb + self.cfg.spawn_noise_std * 0.02 * eps_e
+            m_new = m_base + self.cfg.thought_noise_std * 0.02 * eps_m
+        return e_new, m_new
+
+    def init_empty_belt(
+        self,
+        batch_size: int,
+        tick: int,
+        device: torch.device,
+        deterministic: bool = False,
+    ) -> BeltState:
+        e_soft, m_thought = self.spawn_raw_dough(
+            batch_size=batch_size,
+            num_slots=self.belt_size,
+            device=device,
+            deterministic=deterministic,
+        )
+        return BeltState(e_soft=e_soft, m_thought=m_thought, tick=tick)
+
+    def compose_dough_inputs(self, belt: BeltState, active_slots: Optional[int] = None) -> torch.Tensor:
+        B = active_slots if active_slots is not None else self.belt_size
+        e_soft = belt.e_soft[:, :B, :]
+        m_thought = belt.m_thought[:, :B, :]
+        slots = torch.arange(B, device=e_soft.device, dtype=torch.long)
+        s_emb = self.slot_emb(slots).unsqueeze(0)
+
+        if self.cfg.use_thought_highway:
+            alpha = self.alpha_schedule[:B].view(1, B, 1)  # alpha_0 == 0.0
+            fused = e_soft + alpha * m_thought + s_emb
+        else:
+            fused = e_soft + s_emb
+
+        return self.dough_norm(fused)
+
+    def _build_belt_mask(
+        self,
+        belt_len: int,
+        prefix_len: int,
+        device: torch.device,
+    ) -> Optional[torch.Tensor]:
+        total_k = prefix_len + belt_len
+        if not self.cfg.causal_belt_only:
+            return torch.ones((1, 1, belt_len, total_k), device=device, dtype=torch.bool)
+        mask = torch.ones((belt_len, total_k), device=device, dtype=torch.bool).tril(diagonal=prefix_len)
+        return mask.unsqueeze(0).unsqueeze(0)
+
+    def forward_belt(
+        self,
+        belt: BeltState,
+        kv_cache: Optional[PrefixKVCache] = None,
+        active_slots: Optional[int] = None,
+    ) -> BeltStepOutput:
+        """
+        Single forward pass across all active slots on the belt:
+        Input : [Committed Prefix KV Cache x_{0..t-1}] + [Belt Dough 0..B-1]
+        Output: [Slot 0 Token Logits + Updated Belt State (Lexical + Thought + Elastic Shift Gate)]
+        """
+        B = active_slots if active_slots is not None else self.belt_size
+        dough = self.compose_dough_inputs(belt, active_slots=B)
+        device = dough.device
+
+        prefix_len = kv_cache.seq_len if kv_cache is not None else 0
+        offsets = self.get_slot_target_offsets(active_slots=B)
+        positions = belt.tick + offsets
+        cos, sin = self.rope.get_cos_sin(positions)
+
+        attn_mask = (
+            None
+            if not self.cfg.causal_belt_only
+            else self._build_belt_mask(belt_len=B, prefix_len=prefix_len, device=device)
+        )
+
+        h = dough
+        use_ckpt = self.training and self.cfg.grad_checkpointing
+        for i, block in enumerate(self.blocks):
+            layer_kv = (kv_cache.k_cache[i], kv_cache.v_cache[i]) if kv_cache is not None else None
+            if use_ckpt:
+                def _run_block(inp: torch.Tensor, blk: ConveyerBlock = block, lkv: Optional[tuple[torch.Tensor, torch.Tensor]] = layer_kv) -> torch.Tensor:
+                    out_h, _ = blk(inp, cos=cos, sin=sin, attn_mask=attn_mask, layer_kv=lkv, return_new_kv=False)
+                    return out_h
+                h = checkpoint(_run_block, h, use_reentrant=False)
+            else:
+                h, _ = block(
+                    h,
+                    cos=cos,
+                    sin=sin,
+                    attn_mask=attn_mask,
+                    layer_kv=layer_kv,
+                    return_new_kv=False,
+                )
+
+        h = self.final_norm(h)  # (N, B, D)
+
+        # Channel 1: Lexical Anchor (Soft Probability Projection)
+        logits = self.lm_head(h)  # (N, B, V)
+        tau = self.tau_schedule[:B].view(1, B, 1).type_as(logits)
+        probs = F.softmax((logits / tau).float(), dim=-1).type_as(self.tok_emb.weight)
+        e_soft = torch.matmul(probs, self.tok_emb.weight)  # (N, B, D)
+
+        # Channel 2: Continuous Thought Highway (Gated Residual Update)
+        cand = self.thought_cand(h)
+        gate = torch.sigmoid(self.thought_gate(h))
+        m_prev = belt.m_thought[:, :B, :]
+        if self.cfg.use_thought_highway:
+            m_updated = gate * m_prev + (1.0 - gate) * cand
+        else:
+            m_updated = torch.zeros_like(m_prev)
+
+        # Elastic Planning Shift/Pause Gate & Mid-Belt Sentence Expansion
+        F_slots = min(self.front_slots, B)
+        if self.cfg.use_elastic_planning and B > F_slots:
+            prior = self.shift_prior[:B].view(1, B).type_as(h)
+            raw_shift = self.shift_head(h).squeeze(-1) + prior  # (N, B)
+            shift_prob = torch.sigmoid(raw_shift)
+            # Front conveyor slots (0 .. F-1) always shift left at full speed (1.0)
+            front_ones = torch.ones((h.size(0), F_slots), device=device, dtype=shift_prob.dtype)
+            shift_prob = torch.cat([front_ones, shift_prob[:, F_slots:]], dim=1)
+            # Mid-belt expansion thought synthesized from the boundary between Front (F-1) and Planning (F)
+            boundary_pair = torch.cat([h[:, F_slots - 1, :], h[:, F_slots, :]], dim=-1)
+            expand_cand = self.expand_head(boundary_pair)  # (N, D)
+        else:
+            shift_prob = torch.ones((h.size(0), B), device=device, dtype=h.dtype)
+            expand_cand = h[:, -1, :]
+
+        return BeltStepOutput(
+            logits=logits,
+            probs=probs,
+            e_soft=e_soft,
+            cand_thought=cand,
+            gate=gate,
+            m_updated=m_updated,
+            shift_prob=shift_prob,
+            expand_cand=expand_cand,
+            slot0_logits=logits[:, 0, :],
+        )
+
+    def shift_belt_left(
+        self,
+        step_out: BeltStepOutput,
+        current_tick: int,
+        deterministic_spawn: bool = False,
+        clamp_tokens: Optional[torch.Tensor] = None,
+        clamp_mask: Optional[torch.Tensor] = None,
+    ) -> BeltState:
+        """
+        Updates the belt for tick t + 1:
+          1. Front Conveyor Zone (Slots 0 .. F-2):
+             Always shifts left by 1 slot (`1..F-1 -> 0..F-2`) so fully-baked tokens flow smoothly out of Slot 0.
+          2. Mid-Belt Sentence Expansion Zone (Slot F-1, e.g. Slot 15):
+             Blends the advancing plan from Slot F with the newly expanded clause (`expand_cand`),
+             allowing new sentences and words to form in the middle while the back 75% pauses!
+          3. Back 75% Elastic Planning Zone (Slots F .. B-2, e.g. 16..62):
+             Each slot k updates elastically using its learned shift probability `s_k`:
+               e'_{k} = s_k * e_{k+1} + (1 - s_k) * e_k
+               m'_{k} = s_k * m_{k+1} + (1 - s_k) * m_k
+             When `s_k` is low (~0.25), the back 75% of the belt pauses in-place holding the long-range
+             sentence plan steady while the front conveyor articulates the current sentence!
+          4. Slot B-1 (Slot 63):
+             Blends fresh raw dough into the tail of the planning belt.
+        """
+        bsz, B, _ = step_out.e_soft.shape
+        device = step_out.e_soft.device
+        F_slots = min(self.front_slots, B)
+
+        e_spawn, m_spawn = self.spawn_raw_dough(
+            batch_size=bsz,
+            num_slots=1,
+            device=device,
+            deterministic=deterministic_spawn,
+        )
+
+        if B <= 1:
+            e_next = e_spawn
+            m_next = m_spawn
+            probs_next = torch.full((bsz, 1, self.vocab_size), 1.0 / self.vocab_size, device=device, dtype=step_out.probs.dtype)
+            logits_next = torch.zeros((bsz, 1, self.vocab_size), device=device, dtype=step_out.logits.dtype)
+            gates_next = torch.ones((bsz, 1, self.d_model), device=device, dtype=step_out.gate.dtype)
+        elif (not self.cfg.use_elastic_planning) or B <= F_slots:
+            # Strict 1-by-1 shift across all B slots
+            e_next = torch.cat([step_out.e_soft[:, 1:, :], e_spawn], dim=1)
+            m_next = torch.cat([step_out.m_updated[:, 1:, :], m_spawn], dim=1)
+            probs_next = torch.cat(
+                [
+                    step_out.probs[:, 1:, :],
+                    torch.full((bsz, 1, self.vocab_size), 1.0 / self.vocab_size, device=device, dtype=step_out.probs.dtype),
+                ],
+                dim=1,
+            )
+            logits_next = torch.cat(
+                [
+                    step_out.logits[:, 1:, :],
+                    torch.zeros((bsz, 1, self.vocab_size), device=device, dtype=step_out.logits.dtype),
+                ],
+                dim=1,
+            )
+            gates_next = torch.cat(
+                [
+                    step_out.gate[:, 1:, :],
+                    torch.ones((bsz, 1, self.d_model), device=device, dtype=step_out.gate.dtype),
+                ],
+                dim=1,
+            )
+        else:
+            # Elastic Two-Zone Update:
+            # (a) Front slots 0 .. F-2 shift strictly from 1 .. F-1
+            e_front = step_out.e_soft[:, 1:F_slots, :]
+            m_front = step_out.m_updated[:, 1:F_slots, :]
+
+            # (b) Mid-Belt Expansion Slot (F-1): where paused back-belt plans expand into new active words
+            s_mid = step_out.shift_prob[:, F_slots : F_slots + 1].unsqueeze(-1)  # (N, 1, 1)
+            e_FromPlan = step_out.e_soft[:, F_slots : F_slots + 1, :]
+            m_FromPlan = step_out.m_updated[:, F_slots : F_slots + 1, :]
+            m_expand = step_out.expand_cand.unsqueeze(1)  # (N, 1, D)
+            e_mid = e_FromPlan
+            m_mid = s_mid * m_FromPlan + (1.0 - s_mid) * m_expand
+
+            # (c) Back 75% Planning Slots (F .. B-2): elastic interpolation between staying (k) and shifting (k+1)
+            s_back = step_out.shift_prob[:, F_slots : B - 1].unsqueeze(-1)  # (N, B - 1 - F, 1)
+            e_stay = step_out.e_soft[:, F_slots : B - 1, :]
+            e_adv = step_out.e_soft[:, F_slots + 1 : B, :]
+            e_back = s_back * e_adv + (1.0 - s_back) * e_stay
+
+            m_stay = step_out.m_updated[:, F_slots : B - 1, :]
+            m_adv = step_out.m_updated[:, F_slots + 1 : B, :]
+            m_back = s_back * m_adv + (1.0 - s_back) * m_stay
+
+            # (d) Slot B-1 (Tail of planning belt): elastic blend between holding Slot B-1 and fresh raw dough
+            s_tail = step_out.shift_prob[:, B - 1 : B].unsqueeze(-1)
+            e_tail = s_tail * e_spawn + (1.0 - s_tail) * step_out.e_soft[:, B - 1 : B, :]
+            m_tail = s_tail * m_spawn + (1.0 - s_tail) * step_out.m_updated[:, B - 1 : B, :]
+
+            e_next = torch.cat([e_front, e_mid, e_back, e_tail], dim=1)
+            m_next = torch.cat([m_front, m_mid, m_back, m_tail], dim=1)
+
+            probs_next = torch.cat(
+                [
+                    step_out.probs[:, 1:, :],
+                    torch.full((bsz, 1, self.vocab_size), 1.0 / self.vocab_size, device=device, dtype=step_out.probs.dtype),
+                ],
+                dim=1,
+            )
+            logits_next = torch.cat(
+                [
+                    step_out.logits[:, 1:, :],
+                    torch.zeros((bsz, 1, self.vocab_size), device=device, dtype=step_out.logits.dtype),
+                ],
+                dim=1,
+            )
+            gates_next = torch.cat(
+                [
+                    step_out.gate[:, 1:, :],
+                    torch.ones((bsz, 1, self.d_model), device=device, dtype=step_out.gate.dtype),
+                ],
+                dim=1,
+            )
+
+        if clamp_tokens is not None and clamp_mask is not None:
+            clamped_emb = self.tok_emb(clamp_tokens)
+            mask_exp = clamp_mask.unsqueeze(-1)
+            e_next = torch.where(mask_exp, clamped_emb, e_next)
+
+        return BeltState(
+            e_soft=e_next,
+            m_thought=m_next,
+            tick=current_tick + 1,
+            probs=probs_next,
+            logits=logits_next,
+            gates=gates_next,
+            shift_probs=step_out.shift_prob,
+        )

src/conveyer/sweep.py

@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import Any
+import torch
+
+from conveyer.checkpoint import load_checkpoint
+from conveyer.config import load_config
+from conveyer.data import build_dataset_and_tokenizer
+from conveyer.eval import evaluate_closed_loop_generation
+from conveyer.train import train_model
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Head-to-Head Benchmark: 1-Pass AR Baseline vs. 64-Pass Conveyor Diffusion LLM & Ablations"
+    )
+    parser.add_argument("--config", type=str, default="configs/synthetic.yaml", help="Base config file")
+    parser.add_argument("--steps", type=int, default=400, help="Training steps per variant")
+    parser.add_argument("--episodes", type=int, default=64, help="Closed-loop evaluation episodes")
+    parser.add_argument("--out-dir", type=str, default="runs/sweep", help="Output directory for sweep")
+    args = parser.parse_args()
+
+    variants = [
+        {
+            "name": "1_ar_baseline",
+            "label": "1-Pass Causal AR Baseline (B=1)",
+            "overrides": ["train.ar_baseline=true", "train.unroll_steps=8"],
+            "bake_steps": 1,
+            "overlap": 0,
+        },
+        {
+            "name": "2_conveyer_soft_only",
+            "label": "64-Slot Conveyor (Soft-Vocab Only, alpha_k=0)",
+            "overrides": ["train.no_thought_highway=true"],
+            "bake_steps": 64,
+            "overlap": 16,
+        },
+        {
+            "name": "3_conveyer_nobptt",
+            "label": "64-Slot Conveyor (Detached W=1, No Belt BPTT)",
+            "overrides": ["train.no_belt_bptt=true"],
+            "bake_steps": 64,
+            "overlap": 16,
+        },
+        {
+            "name": "4_conveyer_dual_full",
+            "label": "64-Slot Dual-Channel Conveyor (Soft-Vocab + Gated Thought + BPTT)",
+            "overrides": [],
+            "bake_steps": 64,
+            "overlap": 16,
+        },
+    ]
+
+    summary: list[dict[str, Any]] = []
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    for var in variants:
+        run_dir = Path(args.out_dir) / var["name"]
+        overrides = [
+            f"train.max_steps={args.steps}",
+            f"train.eval_every={max(50, args.steps // 2)}",
+            f"train.out_dir={run_dir.as_posix()}",
+        ] + var["overrides"]
+
+        print("\n" + "=" * 90)
+        print(f"Training Variant: {var['label']}")
+        print("=" * 90)
+        cfg = load_config(args.config, overrides)
+        val_m = train_model(cfg)
+
+        ckpt_path = run_dir / "best.pt"
+        model, loaded_cfg, _ = load_checkpoint(ckpt_path, device=device)
+        _, val_ds, tokenizer = build_dataset_and_tokenizer(loaded_cfg)
+
+        cl_metrics = evaluate_closed_loop_generation(
+            model=model,
+            val_ds=val_ds,
+            tokenizer=tokenizer,
+            device=device,
+            num_episodes=args.episodes,
+            prompt_overlap=var["overlap"],
+            bake_steps=var["bake_steps"],
+        )
+        row = {
+            "name": var["name"],
+            "label": var["label"],
+            "val_slot0_acc": val_m.get("val_slot0_acc", 0.0),
+            "val_slot0_ppl": val_m.get("val_slot0_ppl", 0.0),
+            **cl_metrics,
+        }
+        summary.append(row)
+
+    print("\n" + "=" * 100)
+    print(f"{'Model / Ablation Variant':<56} | {'Pos-0 Acc':<10} | {'Token Acc':<10} | {'Exact Seq':<10}")
+    print("-" * 100)
+    for r in summary:
+        print(
+            f"{r['label']:<56} | {r['first_token_acc']*100:8.2f}% | "
+            f"{r['token_acc']*100:8.2f}% | {r['exact_seq_acc']*100:8.2f}%"
+        )
+    print("=" * 100)
+
+    out_path = Path(args.out_dir) / "sweep_summary.json"
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    out_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")
+
+
+if __name__ == "__main__":
+    main()

src/conveyer/train.py

@@ -0,0 +1,472 @@
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import os
+from pathlib import Path
+import subprocess
+import time
+from typing import Any, Optional
+import torch
+from torch.utils.data import DataLoader, Subset
+
+from conveyer.checkpoint import load_checkpoint_payload, save_checkpoint
+from conveyer.config import ConveyerConfig, load_config
+from conveyer.data import TokenizerProtocol, build_dataset_and_tokenizer, collate_conveyer_batch
+from conveyer.infer import generate_conveyer
+from conveyer.losses import compute_unrolled_conveyer_loss
+from conveyer.model import ConveyerDiffusionLM
+
+
+def _query_gpu_temp_c() -> Optional[int]:
+    try:
+        out = subprocess.check_output(
+            ["nvidia-smi", "--query-gpu=temperature.gpu", "--format=csv,noheader,nounits"],
+            timeout=2.0,
+            stderr=subprocess.DEVNULL,
+        )
+        return int(out.decode("utf-8", errors="ignore").strip().splitlines()[0])
+    except Exception:
+        return None
+
+
+def get_lr(step: int, max_steps: int, base_lr: float, warmup_steps: int) -> float:
+    if step < warmup_steps:
+        return base_lr * float(step + 1) / float(max(1, warmup_steps))
+    progress = float(step - warmup_steps) / float(max(1, max_steps - warmup_steps))
+    progress = min(max(progress, 0.0), 1.0)
+    return base_lr * (0.1 + 0.9 * 0.5 * (1.0 + math.cos(math.pi * progress)))
+
+
+@torch.no_grad()
+def evaluate_validation(
+    model: ConveyerDiffusionLM,
+    val_loader: DataLoader[dict[str, torch.Tensor]],
+    cfg: ConveyerConfig,
+    device: torch.device,
+    pad_id: int = 0,
+    max_batches: int = 25,
+) -> dict[str, Any]:
+    model.eval()
+    B = model.belt_size
+    loss_sum = 0.0
+    s0_loss_sum = 0.0
+    s0_acc_sum = 0.0
+    lm_loss_sum = 0.0
+    lm_acc_sum = 0.0
+    belt_acc_sum = 0.0
+    per_slot_acc_sum = torch.zeros(B, device=device)
+    per_slot_ent_sum = torch.zeros(B, device=device)
+    count = 0
+
+    for i, batch in enumerate(val_loader):
+        if i >= max_batches:
+            break
+        seq = batch["input_ids"].to(device)
+        plen = batch["prompt_lengths"].to(device)
+        t_eval = max(1, int(plen.min().item()) - min(4, cfg.train.unroll_steps // 2))
+        use_amp = device.type == "cuda" and torch.cuda.is_bf16_supported()
+        with torch.autocast(device_type=device.type, dtype=torch.bfloat16 if use_amp else torch.float32, enabled=use_amp):
+            m = compute_unrolled_conveyer_loss(
+                model=model,
+                seq_tokens=seq,
+                cfg=cfg,
+                prompt_lengths=plen,
+                tick_start=t_eval,
+                pad_token_id=pad_id,
+            )
+        loss_sum += float(m.loss.item())
+        s0_loss_sum += m.slot0_loss
+        s0_acc_sum += m.slot0_acc
+        lm_loss_sum += m.prefix_lm_loss
+        lm_acc_sum += m.prefix_lm_acc
+        belt_acc_sum += m.belt_mean_acc
+        per_slot_acc_sum += m.per_slot_acc
+        per_slot_ent_sum += m.per_slot_entropy
+        count += 1
+
+    model.train()
+    denom = max(count, 1)
+    s0_ce = s0_loss_sum / denom
+    lm_ce = lm_loss_sum / denom
+    return {
+        "val_loss": loss_sum / denom,
+        "val_slot0_ce": s0_ce,
+        "val_slot0_ppl": math.exp(min(s0_ce, 20.0)),
+        "val_slot0_acc": s0_acc_sum / denom,
+        "val_lm_ce": lm_ce,
+        "val_lm_ppl": math.exp(min(lm_ce, 20.0)),
+        "val_lm_acc": lm_acc_sum / denom,
+        "val_belt_acc": belt_acc_sum / denom,
+        "val_per_slot_acc": (per_slot_acc_sum / denom).cpu().tolist(),
+        "val_per_slot_entropy": (per_slot_ent_sum / denom).cpu().tolist(),
+    }
+
+
+def _prune_old_milestones(out_dir: Path, keep_last: int = 8) -> None:
+    milestones = sorted(out_dir.glob("ckpt_step_*.pt"), key=lambda p: p.stat().st_mtime)
+    if len(milestones) > keep_last:
+        for old_p in milestones[: len(milestones) - keep_last]:
+            try:
+                old_p.unlink()
+                prev_p = old_p.with_suffix(old_p.suffix + ".prev")
+                if prev_p.exists():
+                    prev_p.unlink()
+            except OSError:
+                pass
+
+
+@torch.no_grad()
+def _sample_preview(
+    model: ConveyerDiffusionLM,
+    tokenizer: TokenizerProtocol,
+    device: torch.device,
+    prompt_text: str,
+    max_new_tokens: int = 32,
+) -> str:
+    model.eval()
+    p_ids = torch.tensor([tokenizer.encode(prompt_text.strip())], dtype=torch.long, device=device)
+    out_ids, _ = generate_conveyer(
+        model=model,
+        prompt_tokens=p_ids,
+        max_new_tokens=max_new_tokens,
+        prompt_overlap=10,
+        temperature=0.0,
+        eos_token_id=tokenizer.eos_id,
+    )
+    gen_part = tokenizer.decode(out_ids[0, p_ids.size(1) :].tolist())
+    model.train()
+    return gen_part.replace("\n", " ").strip()
+
+
+def _sanitize_jsonl(log_path: Path) -> None:
+    """Strips any trailing NUL bytes left by an abrupt OS crash so the log remains valid JSONL."""
+    if not log_path.exists():
+        return
+    raw = log_path.read_bytes().replace(b"\x00", b"")
+    lines = [ln for ln in raw.splitlines() if ln.strip().startswith(b"{") and ln.strip().endswith(b"}")]
+    log_path.write_bytes(b"\n".join(lines) + (b"\n" if lines else b""))
+
+
+def train_model(
+    cfg: ConveyerConfig,
+    init_ckpt: str | Path | None = None,
+    resume_ckpt: str | Path | None = None,
+) -> dict[str, Any]:
+    torch.manual_seed(cfg.train.seed)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    train_ds, val_ds, tokenizer = build_dataset_and_tokenizer(cfg)
+    if cfg.train.overfit_batches > 0:
+        n_sub = min(len(train_ds), cfg.train.overfit_batches * cfg.train.batch_size)
+        train_ds = Subset(train_ds, list(range(n_sub)))
+        val_ds = train_ds
+
+    collate_fn = lambda b: collate_conveyer_batch(b, pad_id=tokenizer.pad_id)
+    train_loader = DataLoader(
+        train_ds,
+        batch_size=cfg.train.batch_size,
+        shuffle=(cfg.train.overfit_batches == 0),
+        collate_fn=collate_fn,
+        drop_last=False,
+    )
+    val_loader = DataLoader(
+        val_ds,
+        batch_size=cfg.train.batch_size,
+        shuffle=False,
+        collate_fn=collate_fn,
+        drop_last=False,
+    )
+
+    model = ConveyerDiffusionLM(cfg.model).to(device)
+    optimizer = torch.optim.AdamW(
+        model.parameters(),
+        lr=cfg.train.lr,
+        weight_decay=cfg.train.weight_decay,
+        betas=(0.9, 0.95),
+    )
+
+    out_dir = Path(cfg.train.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    log_jsonl = out_dir / "train_log.jsonl"
+    _sanitize_jsonl(log_jsonl)
+
+    start_step = 1
+    best_val_loss = float("inf")
+    prior_elapsed = 0.0
+
+    # Resolve resume vs warm-start initialization (with automatic fallback to .prev or ckpt_step_*.pt)
+    resume_target: Optional[Path] = None
+    if resume_ckpt is not None:
+        if str(resume_ckpt).lower() == "auto":
+            resume_target = out_dir / "last.pt"
+        else:
+            resume_target = Path(resume_ckpt)
+
+    if resume_target is not None:
+        try:
+            ckpt_data, used_path = load_checkpoint_payload(resume_target, map_location="cpu")
+            model.load_state_dict(ckpt_data["model_state"], strict=False)
+            if "optimizer_state" in ckpt_data:
+                try:
+                    optimizer.load_state_dict(ckpt_data["optimizer_state"])
+                except Exception as e:
+                    print(f"[ConveyerTrain] Warning: could not restore optimizer state ({e})", flush=True)
+            start_step = int(ckpt_data.get("step", 0)) + 1
+            best_val_loss = float(
+                ckpt_data.get("best_val_loss") or ckpt_data.get("metrics", {}).get("val_loss", float("inf"))
+            )
+            prior_elapsed = float(ckpt_data.get("elapsed_sec", 0.0))
+            print(
+                f"[ConveyerTrain] Resumed from {used_path} at step {start_step} "
+                f"(best_val_loss={best_val_loss:.4f}, prior_elapsed={prior_elapsed/3600:.2f}h)",
+                flush=True,
+            )
+            del ckpt_data
+            if device.type == "cuda":
+                torch.cuda.empty_cache()
+        except Exception as e:
+            print(f"[ConveyerTrain] No readable resume checkpoint in {out_dir} ({e}); starting from step 1.", flush=True)
+    elif init_ckpt is not None and Path(init_ckpt).exists():
+        ckpt_data, used_path = load_checkpoint_payload(init_ckpt, map_location="cpu")
+        model.load_state_dict(ckpt_data["model_state"], strict=False)
+        print(f"[ConveyerTrain] Warm-started weights from {used_path}", flush=True)
+        del ckpt_data
+
+    use_amp = device.type == "cuda" and torch.cuda.is_bf16_supported()
+    amp_dtype = torch.bfloat16 if use_amp else torch.float32
+
+    mode_desc = (
+        "1-Pass AR Baseline"
+        if cfg.model.causal_belt_only
+        else (
+            "64-Slot Conveyor (Soft-Vocab Only)"
+            if not cfg.model.use_thought_highway
+            else "64-Slot Dual-Channel Conveyor (Soft-Vocab + Gated Thought)"
+        )
+    )
+    max_hours = float(getattr(cfg.train, "max_hours", 0.0))
+    save_every = max(1, int(getattr(cfg.train, "save_every", 500)))
+    milestone_every = save_every * 2
+
+    print(
+        f"[ConveyerTrain] Mode: {mode_desc} | Params: {model.num_parameters()/1e6:.2f}M | "
+        f"Device: {device} | Train stories: {len(train_ds):,} | Val stories: {len(val_ds):,} | "
+        f"Unroll U={cfg.train.unroll_steps} | Belt B={cfg.model.belt_size} | "
+        f"max_steps={cfg.train.max_steps:,} | max_hours={max_hours:.1f}h | save_every={save_every}",
+        flush=True,
+    )
+
+    train_iter = iter(train_loader)
+    last_val_metrics: dict[str, Any] = {}
+    t0 = time.time()
+
+    model.train()
+    for step in range(start_step, cfg.train.max_steps + 1):
+        try:
+            batch = next(train_iter)
+        except StopIteration:
+            train_iter = iter(train_loader)
+            batch = next(train_iter)
+
+        lr = get_lr(step - 1, cfg.train.max_steps, cfg.train.lr, cfg.train.warmup_steps)
+        for pg in optimizer.param_groups:
+            pg["lr"] = lr
+
+        seq = batch["input_ids"].to(device)
+        plen = batch["prompt_lengths"].to(device)
+
+        optimizer.zero_grad(set_to_none=True)
+        with torch.autocast(device_type=device.type, dtype=amp_dtype, enabled=use_amp):
+            metrics = compute_unrolled_conveyer_loss(
+                model=model,
+                seq_tokens=seq,
+                cfg=cfg,
+                prompt_lengths=plen,
+                pad_token_id=tokenizer.pad_id,
+            )
+        metrics.loss.backward()
+        if cfg.train.grad_clip > 0.0:
+            torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.train.grad_clip)
+        optimizer.step()
+
+        cooldown_ms = int(getattr(cfg.train, "step_cooldown_ms", 0))
+        if cooldown_ms > 0:
+            if device.type == "cuda":
+                torch.cuda.synchronize()
+            time.sleep(cooldown_ms / 1000.0)
+
+        total_elapsed = prior_elapsed + (time.time() - t0)
+
+        if step % cfg.train.log_every == 0 or step == start_step:
+            acc_back = float(metrics.per_slot_acc[-1].item())
+            acc_mid = float(metrics.per_slot_acc[model.belt_size // 2].item())
+            steps_done = max(1, step - start_step + 1)
+            sec_per_step = (time.time() - t0) / steps_done
+            gpu_temp = _query_gpu_temp_c() if device.type == "cuda" else None
+            temp_str = f" | {gpu_temp}C" if gpu_temp is not None else ""
+            print(
+                f"step {step:6d}/{cfg.train.max_steps} | loss {metrics.loss.item():.4f} | "
+                f"lm_ce {metrics.prefix_lm_loss:.3f} (acc {metrics.prefix_lm_acc:.3f}) | "
+                f"s0_ce {metrics.slot0_loss:.3f} (acc {metrics.slot0_acc:.3f}) | "
+                f"belt(0/m/63) {metrics.slot0_acc:.2f}/{acc_mid:.2f}/{acc_back:.2f} | "
+                f"pause {metrics.mean_plan_pause*100:.0f}% | lr {lr:.2e} | "
+                f"{sec_per_step:.2f}s/it{temp_str} | {total_elapsed/3600:.2f}h",
+                flush=True,
+            )
+            max_temp = int(getattr(cfg.train, "max_gpu_temp_c", 80))
+            if gpu_temp is not None and max_temp > 0 and gpu_temp >= max_temp:
+                torch.cuda.synchronize()
+                while True:
+                    print(
+                        f"  [THERMAL GUARD] GPU at {gpu_temp}C (>= {max_temp}C limit); pausing 8s to let fans cool GPU...",
+                        flush=True,
+                    )
+                    time.sleep(8.0)
+                    gpu_temp = _query_gpu_temp_c()
+                    if gpu_temp is None or gpu_temp <= max_temp - 5:
+                        print(f"  [THERMAL GUARD] GPU cooled to {gpu_temp}C; resuming training.", flush=True)
+                        break
+
+        is_time_up = max_hours > 0.0 and (total_elapsed >= max_hours * 3600.0)
+        is_eval_step = (step % cfg.train.eval_every == 0) or (step == cfg.train.max_steps) or is_time_up
+        is_save_step = (step % save_every == 0) or is_eval_step
+
+        if is_eval_step:
+            if device.type == "cuda":
+                torch.cuda.empty_cache()
+            val_m = evaluate_validation(
+                model=model,
+                val_loader=val_loader,
+                cfg=cfg,
+                device=device,
+                pad_id=tokenizer.pad_id,
+            )
+            last_val_metrics = val_m
+            curve = val_m["val_per_slot_acc"]
+            B = len(curve)
+            print(
+                f"  [VAL step {step:6d} | {total_elapsed/3600:.2f}h] "
+                f"val_loss {val_m['val_loss']:.4f} | "
+                f"real_s0_ppl {val_m['val_slot0_ppl']:.2f} | real_s0_acc {val_m['val_slot0_acc']:.3f} | "
+                f"real_lm_ppl {val_m['val_lm_ppl']:.2f} | real_lm_acc {val_m['val_lm_acc']:.3f} | "
+                f"crystallization(63->32->8->0): "
+                f"{curve[B-1]:.2f}->{curve[B//2]:.2f}->{curve[min(8, B-1)]:.2f}->{curve[0]:.2f}",
+                flush=True,
+            )
+            if cfg.data.dataset == "tinystories":
+                sample_p = (
+                    "Once upon a time, Lily and a little bunny walked into the green forest "
+                    "to find the golden key near the locked treasure chest. Suddenly a big lion"
+                )
+                preview = _sample_preview(model, tokenizer, device, sample_p, max_new_tokens=28)
+                print(f"    [SAMPLE] ...Suddenly a big lion -> {preview}", flush=True)
+            else:
+                preview = ""
+
+            if val_m["val_loss"] <= best_val_loss:
+                best_val_loss = val_m["val_loss"]
+                save_checkpoint(
+                    out_dir / "best.pt",
+                    model,
+                    cfg,
+                    step=step,
+                    optimizer=None,  # Keep best.pt lightweight (~115 MB) for fast inference & fewer disk writes
+                    metrics=val_m,
+                    best_val_loss=best_val_loss,
+                    elapsed_sec=total_elapsed,
+                )
+
+            log_entry = {
+                "step": step,
+                "elapsed_hours": round(total_elapsed / 3600.0, 4),
+                "val_loss": round(val_m["val_loss"], 4),
+                "val_slot0_ce": round(val_m["val_slot0_ce"], 4),
+                "val_slot0_ppl": round(val_m["val_slot0_ppl"], 2),
+                "val_slot0_acc": round(val_m["val_slot0_acc"], 4),
+                "val_lm_ce": round(val_m["val_lm_ce"], 4),
+                "val_lm_ppl": round(val_m["val_lm_ppl"], 2),
+                "val_lm_acc": round(val_m["val_lm_acc"], 4),
+                "val_belt_acc": round(val_m["val_belt_acc"], 4),
+                "sample_lion": preview,
+            }
+            with open(log_jsonl, "a", encoding="utf-8") as lf:
+                lf.write(json.dumps(log_entry) + "\n")
+                lf.flush()
+                os.fsync(lf.fileno())
+
+        if is_save_step:
+            save_checkpoint(
+                out_dir / "last.pt",
+                model,
+                cfg,
+                step=step,
+                optimizer=optimizer,
+                metrics=last_val_metrics,
+                best_val_loss=best_val_loss,
+                elapsed_sec=total_elapsed,
+            )
+            if step % milestone_every == 0:
+                save_checkpoint(
+                    out_dir / f"ckpt_step_{step:06d}.pt",
+                    model,
+                    cfg,
+                    step=step,
+                    optimizer=optimizer,
+                    metrics=last_val_metrics,
+                    best_val_loss=best_val_loss,
+                    elapsed_sec=total_elapsed,
+                )
+                _prune_old_milestones(out_dir, keep_last=8)
+            if device.type == "cuda":
+                torch.cuda.empty_cache()
+
+        if is_time_up:
+            print(
+                f"[ConveyerTrain] Reached max_hours={max_hours:.2f}h budget at step {step}. "
+                f"Checkpoints saved to {out_dir}.",
+                flush=True,
+            )
+            break
+
+    return last_val_metrics
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Train Conveyor-Belt Diffusion LLM")
+    parser.add_argument("--config", type=str, default="configs/synthetic.yaml", help="Path to YAML config")
+    parser.add_argument("--init", type=str, default=None, help="Optional checkpoint (.pt) to warm-start weights from")
+    parser.add_argument(
+        "--resume",
+        type=str,
+        default=None,
+        help="Checkpoint path (.pt) or 'auto' to resume full training state (model + optimizer + step) after a crash",
+    )
+    parser.add_argument("--max_hours", type=float, default=None, help="Wall-clock training budget in hours (e.g. 24)")
+    parser.add_argument("--overfit", type=int, default=0, help="Overfit on N fixed batches for verification")
+    parser.add_argument("--no_thought_highway", action="store_true", help="Disable continuous thought highway (alpha_k = 0)")
+    parser.add_argument("--no_belt_bptt", action="store_true", help="Detach conveyor belt state between unrolled ticks (W=1)")
+    parser.add_argument("--ar_baseline", action="store_true", help="Train 1-pass causal Autoregressive baseline")
+    parser.add_argument("overrides", nargs="*", help="Config overrides in section.key=value format")
+    args = parser.parse_args()
+
+    cfg = load_config(args.config, args.overrides)
+    if args.max_hours is not None:
+        cfg.train.max_hours = args.max_hours
+    if args.overfit > 0:
+        cfg.train.overfit_batches = args.overfit
+    if args.no_thought_highway:
+        cfg.train.no_thought_highway = True
+        cfg.model.use_thought_highway = False
+    if args.no_belt_bptt:
+        cfg.train.no_belt_bptt = True
+    if args.ar_baseline:
+        cfg.train.ar_baseline = True
+        cfg.model.causal_belt_only = True
+
+    train_model(cfg, init_ckpt=args.init, resume_ckpt=args.resume)
+
+
+if __name__ == "__main__":
+    main()

src/conveyer/web_sim.py

@@ -0,0 +1,840 @@
+from __future__ import annotations
+
+import argparse
+from http.server import BaseHTTPRequestHandler, HTTPServer
+import json
+import os
+from pathlib import Path
+from typing import Any
+from urllib.parse import parse_qs, urlparse
+import torch
+from torch.utils.data import DataLoader
+
+from conveyer.checkpoint import load_checkpoint
+from conveyer.config import load_config
+from conveyer.data import build_dataset_and_tokenizer, collate_conveyer_batch
+from conveyer.eval import evaluate_crystallization_curve
+from conveyer.infer import generate_conveyer
+from conveyer.model import ConveyerDiffusionLM
+
+HTML_PAGE = """<!DOCTYPE html>
+<html lang="en">
+<head>
+<meta charset="UTF-8">
+<title>Conveyor-Belt Diffusion LLM — Elastic Planning & Oven Simulator</title>
+<style>
+  :root {
+    --bg: #0b0f17;
+    --panel: #151b26;
+    --panel-alt: #1b2332;
+    --border: #2a3447;
+    --text: #e6edf3;
+    --muted: #8b949e;
+    --front: #2ea043;
+    --front-glow: rgba(46, 160, 67, 0.25);
+    --mid: #d29922;
+    --back: #8957e5;
+  }
+  * { box-sizing: border-box; }
+  body {
+    margin: 0;
+    padding: 18px 24px 32px 24px;
+    background: var(--bg);
+    color: var(--text);
+    font-family: 'Segoe UI', system-ui, -apple-system, sans-serif;
+  }
+  .header-row {
+    display: flex;
+    justify-content: space-between;
+    align-items: flex-start;
+    gap: 16px;
+    margin-bottom: 14px;
+  }
+  h1 { margin: 0 0 4px 0; font-size: 22px; letter-spacing: -0.3px; }
+  .subtitle { color: var(--muted); font-size: 13px; }
+  .badge-row { display: flex; gap: 8px; flex-wrap: wrap; }
+  .metric-pill {
+    background: var(--panel-alt);
+    border: 1px solid var(--border);
+    border-radius: 6px;
+    padding: 6px 12px;
+    font-size: 12px;
+  }
+  .metric-pill strong { color: #3fb950; font-size: 14px; margin-left: 4px; }
+
+  .prompt-card {
+    background: var(--panel);
+    border: 1px solid var(--border);
+    border-radius: 10px;
+    padding: 14px 16px;
+    margin-bottom: 14px;
+  }
+  .prompt-top {
+    display: flex;
+    justify-content: space-between;
+    align-items: center;
+    margin-bottom: 8px;
+    flex-wrap: wrap;
+    gap: 8px;
+  }
+  .prompt-top label { font-size: 12px; font-weight: 600; color: var(--muted); text-transform: uppercase; }
+  .presets { display: flex; gap: 6px; flex-wrap: wrap; }
+  .preset-btn {
+    background: var(--panel-alt);
+    color: var(--text);
+    border: 1px solid var(--border);
+    border-radius: 5px;
+    padding: 4px 10px;
+    font-size: 12px;
+    cursor: pointer;
+  }
+  .preset-btn:hover { border-color: #58a6ff; }
+  .prompt-input-row {
+    display: flex;
+    gap: 10px;
+    align-items: center;
+    flex-wrap: wrap;
+  }
+  #promptInput {
+    flex: 1;
+    min-width: 320px;
+    background: #080b11;
+    color: var(--text);
+    border: 1px solid var(--border);
+    border-radius: 6px;
+    padding: 10px 12px;
+    font-size: 14px;
+  }
+  .knob {
+    display: flex;
+    flex-direction: column;
+    font-size: 11px;
+    color: var(--muted);
+    gap: 3px;
+  }
+  .knob input {
+    width: 72px;
+    background: #080b11;
+    color: var(--text);
+    border: 1px solid var(--border);
+    border-radius: 5px;
+    padding: 6px 8px;
+    font-size: 13px;
+  }
+  .btn-primary {
+    background: #238636;
+    color: #fff;
+    border: none;
+    border-radius: 6px;
+    padding: 10px 18px;
+    cursor: pointer;
+    font-weight: 600;
+    font-size: 14px;
+  }
+  .btn-primary:hover { background: #2ea043; }
+  .btn-secondary {
+    background: var(--panel-alt);
+    color: var(--text);
+    border: 1px solid var(--border);
+    border-radius: 6px;
+    padding: 8px 14px;
+    cursor: pointer;
+    font-weight: 600;
+    font-size: 13px;
+  }
+
+  .story-grid {
+    display: grid;
+    grid-template-columns: 2fr 1fr;
+    gap: 14px;
+    margin-bottom: 14px;
+  }
+  @media (max-width: 1000px) { .story-grid { grid-template-columns: 1fr; } }
+  .card {
+    background: var(--panel);
+    border: 1px solid var(--border);
+    border-radius: 10px;
+    padding: 14px 16px;
+  }
+  .card h3 {
+    margin: 0 0 10px 0;
+    font-size: 12px;
+    color: var(--muted);
+    text-transform: uppercase;
+    letter-spacing: 0.5px;
+    display: flex;
+    justify-content: space-between;
+    align-items: center;
+  }
+  .story-box {
+    font-size: 17px;
+    line-height: 1.65;
+    background: #080b11;
+    border: 1px solid var(--border);
+    border-radius: 8px;
+    padding: 14px 16px;
+    min-height: 112px;
+  }
+  .story-prompt { color: #a5d6ff; }
+  .story-gen { color: #3fb950; font-weight: 600; }
+  .story-latest {
+    background: rgba(63, 185, 80, 0.22);
+    border-bottom: 2px solid #3fb950;
+    padding: 1px 4px;
+    border-radius: 4px;
+  }
+  .story-upcoming { color: #8b949e; font-style: italic; font-size: 14px; }
+
+  .cryst-row {
+    display: flex;
+    justify-content: space-between;
+    font-size: 12px;
+    padding: 5px 0;
+    border-bottom: 1px solid #1f2735;
+  }
+  .cryst-bar-wrap {
+    width: 100px;
+    height: 8px;
+    background: #080b11;
+    border-radius: 4px;
+    overflow: hidden;
+    display: inline-block;
+    vertical-align: middle;
+    margin-left: 8px;
+  }
+  .cryst-bar { height: 100%; background: #3fb950; }
+
+  .oven-card {
+    background: var(--panel);
+    border: 2px solid var(--border);
+    border-radius: 12px;
+    padding: 16px;
+  }
+  .oven-controls {
+    display: flex;
+    align-items: center;
+    gap: 14px;
+    margin-bottom: 12px;
+    flex-wrap: wrap;
+  }
+  .conveyor-track-wrapper {
+    position: relative;
+    background: #070a0f;
+    border: 2px solid #283246;
+    border-radius: 10px;
+    padding: 14px 12px;
+    overflow-x: auto;
+  }
+  .conveyor-belt-row {
+    display: flex;
+    gap: 8px;
+    align-items: stretch;
+    min-width: max-content;
+  }
+  .dough-ball {
+    width: 108px;
+    flex-shrink: 0;
+    background: #101622;
+    border: 1px solid var(--border);
+    border-radius: 8px;
+    padding: 8px;
+  }
+  .dough-ball.slot-zero {
+    width: 128px;
+    background: rgba(46, 160, 67, 0.14);
+    border: 2px solid #3fb950;
+    box-shadow: 0 0 14px var(--front-glow);
+  }
+  .dough-ball.zone-front { border-top: 3px solid var(--front); }
+  .dough-ball.zone-mid   { border-top: 3px solid var(--mid); }
+  .dough-ball.zone-back  { border-top: 3px solid var(--back); }
+
+  .ball-top {
+    display: flex;
+    justify-content: space-between;
+    font-size: 10px;
+    color: var(--muted);
+    margin-bottom: 4px;
+  }
+  .ball-word {
+    font-size: 15px;
+    font-weight: 700;
+    color: var(--text);
+    white-space: nowrap;
+    overflow: hidden;
+    text-overflow: ellipsis;
+    margin-bottom: 5px;
+  }
+  .ball-cand {
+    font-size: 10px;
+    color: var(--muted);
+    white-space: nowrap;
+    overflow: hidden;
+    text-overflow: ellipsis;
+    line-height: 1.35;
+  }
+  .conf-bar {
+    height: 4px;
+    background: #1f2735;
+    border-radius: 2px;
+    margin-top: 6px;
+    overflow: hidden;
+  }
+  .conf-fill { height: 100%; background: #3fb950; }
+  .zone-legend {
+    display: flex;
+    gap: 18px;
+    font-size: 12px;
+    color: var(--muted);
+    margin-top: 10px;
+    flex-wrap: wrap;
+  }
+  .dot { display: inline-block; width: 10px; height: 10px; border-radius: 50%; margin-right: 5px; vertical-align: middle; }
+</style>
+</head>
+<body>
+
+  <div class="header-row">
+    <div>
+      <h1>64-Slot Dual-Channel Elastic Conveyor-Belt Diffusion LLM</h1>
+      <div class="subtitle" id="modelSubtitle">Loading model info...</div>
+    </div>
+    <div class="badge-row">
+      <div class="metric-pill">Model Params: <strong id="paramBadge">--</strong></div>
+      <div class="metric-pill">Slot-0 Val Accuracy: <strong id="s0AccBadge">--</strong></div>
+      <div class="metric-pill">Back 75% Planning Pause: <strong id="pauseBadge">--</strong></div>
+      <div class="metric-pill">Sample Token Match: <strong id="sampleAccBadge">--</strong></div>
+    </div>
+  </div>
+
+  <!-- 1. Prompt Input & Controls -->
+  <div class="prompt-card">
+    <div class="prompt-top">
+      <label>1. Enter Any Initial Text Prompt (Press Enter or Click "Bake &amp; Continue Text"):</label>
+      <div class="presets">
+        <button class="preset-btn" onclick="setPreset(0)">Preset 1: Lion in the Forest</button>
+        <button class="preset-btn" onclick="setPreset(1)">Preset 2: Robot &amp; Lost Kitten</button>
+        <button class="preset-btn" onclick="setPreset(2)">Preset 3: Melting Snowman</button>
+        <button class="preset-btn" onclick="setPreset(3)">Preset 4: Missing Blue Crayon</button>
+      </div>
+    </div>
+    <div class="prompt-input-row">
+      <input type="text" id="promptInput" onkeydown="if(event.key==='Enter'){runGenerate(-1);}" value="Once upon a time, Lily and a little bunny walked into the green forest to find the golden key near the locked treasure chest. Suddenly a big lion" />
+      <div class="knob">
+        <span>Pre-Bake Overlap (W)</span>
+        <input type="number" id="overlapInput" value="12" min="0" max="60" />
+      </div>
+      <div class="knob">
+        <span>Oven Slots (B)</span>
+        <input type="number" id="bakeInput" value="64" min="1" max="64" />
+      </div>
+      <div class="knob">
+        <span>Generate Tokens</span>
+        <input type="number" id="tokensInput" value="32" min="6" max="64" />
+      </div>
+      <div class="knob">
+        <span>Temp (&tau;)</span>
+        <input type="number" id="tempInput" value="0.0" step="0.1" min="0.0" max="1.5" />
+      </div>
+      <button class="btn-primary" id="bakeBtn" onclick="runGenerate(-1)">Bake &amp; Continue Text</button>
+    </div>
+  </div>
+
+  <!-- 2. Live Text Continuation + Accuracy Panel -->
+  <div class="story-grid">
+    <div class="card">
+      <h3>
+        <span>2. Live Text Continuation (Blue = Input Prompt &bull; Green = Popped Out of Slot 0)</span>
+        <span id="tickModeLabel" style="color:#3fb950;"></span>
+      </h3>
+      <div class="story-box">
+        <span id="storyPromptSpan" class="story-prompt"></span><span id="storyGenSpan" class="story-gen"></span><span id="storyLatestSpan" class="story-latest" style="display:none;"></span>
+        <span id="storyPreviewSpan" class="story-upcoming"></span>
+      </div>
+      <div id="groundTruthBox" style="margin-top:10px; font-size:12px; color:var(--muted);"></div>
+    </div>
+
+    <div class="card">
+      <h3>
+        <span>Oven Crystallization &amp; Accuracy</span>
+        <span style="font-size:11px; color:#58a6ff;">Back (k=63) &rarr; Front (k=0)</span>
+      </h3>
+      <div id="crystTable"></div>
+    </div>
+  </div>
+
+  <!-- 3. Bottom Conveyor Belt Oven -->
+  <div class="oven-card">
+    <div class="oven-controls">
+      <strong style="font-size:13px; text-transform:uppercase; color:var(--muted);">3. Live 64-Slot Conveyor &amp; Planning Belt (Bottom):</strong>
+      <button class="btn-secondary" onclick="stepTick(-1)">&larr; Prev Tick</button>
+      <button class="btn-primary" id="playBtn" onclick="togglePlay()" style="padding:7px 16px;">Pause</button>
+      <button class="btn-secondary" onclick="stepTick(1)">Next Tick &rarr;</button>
+      <label style="font-size:12px; color:var(--muted); display:flex; align-items:center; gap:8px; flex:1; min-width:220px;">
+        <span id="scrubberLabel">Tick 0/0</span>
+        <input type="range" id="tickSlider" min="0" max="0" value="0" style="flex:1;" oninput="onScrub(this.value)" />
+      </label>
+      <label style="font-size:12px; color:var(--muted);">
+        Speed:
+        <select id="speedSelect" onchange="updateSpeed()" style="background:#080b11; color:var(--text); border:1px solid var(--border); padding:4px; border-radius:4px;">
+          <option value="450">Slow (2.2 Hz)</option>
+          <option value="240" selected>Normal (4 Hz)</option>
+          <option value="110">Fast (9 Hz)</option>
+        </select>
+      </label>
+    </div>
+
+    <div class="conveyor-track-wrapper">
+      <div id="conveyorRow" class="conveyor-belt-row"></div>
+    </div>
+
+    <div class="zone-legend">
+      <span><i class="dot" style="background:#3fb950;"></i><strong>Slot 0 (Front Exit, &alpha;<sub>0</sub> = 0):</strong> Confirmed Token Popping Out</span>
+      <span><i class="dot" style="background:#2ea043;"></i><strong>Slots 1..15 (Front 25% Conveyor):</strong> Shifts Left Every Tick (100% Speed)</span>
+      <span><i class="dot" style="background:#d29922;"></i><strong>Slot 15..20 (Mid-Belt Expansion):</strong> New Clauses Form Here</span>
+      <span><i class="dot" style="background:#8957e5;"></i><strong>Slots 16..63 (Back 75% Elastic Planning):</strong> Pauses in Place Holding Long-Range Sentence Plan</span>
+    </div>
+  </div>
+
+<script>
+const PRESETS = [
+  "Once upon a time, Lily and a little bunny walked into the green forest to find the golden key near the locked treasure chest. Suddenly a big lion",
+  "Once upon a time, there was a little robot named Beep who lived in a noisy workshop. One rainy morning, Beep found a tiny lost kitten by the door. Beep",
+  "One winter day, Max made a snowman in the yard with a carrot nose. When the warm sun came out in the afternoon, the snowman began to",
+  "Lucy had a box of bright crayons, but she could not find the blue crayon to color the sky. She looked under her bed and"
+];
+
+let traceData = null;
+let curFrame = 0;
+let playing = true;
+let intervalMs = 240;
+let timerHandle = null;
+
+function setPreset(idx) {
+  const inp = document.getElementById('promptInput');
+  inp.value = PRESETS[idx];
+  runGenerate(-1);
+}
+
+async function runGenerate(sampleOverride = -1) {
+  const btn = document.getElementById('bakeBtn');
+  btn.textContent = "Baking...";
+  const promptText = document.getElementById('promptInput').value.trim();
+  const ov = document.getElementById('overlapInput').value;
+  const bs = document.getElementById('bakeInput').value;
+  const mt = document.getElementById('tokensInput').value;
+  const tp = document.getElementById('tempInput').value;
+
+  let url = `/api/trace?overlap=${ov}&bake_steps=${bs}&tokens=${mt}&temp=${tp}&_t=${Date.now()}`;
+  if (sampleOverride >= 0 || promptText === "") {
+    url += `&sample=${Math.max(0, sampleOverride)}`;
+  } else {
+    url += `&prompt=${encodeURIComponent(promptText)}`;
+  }
+
+  const resp = await fetch(url, { cache: "no-store" });
+  traceData = await resp.json();
+  btn.textContent = "Bake & Continue Text";
+
+  if (sampleOverride >= 0 || promptText === "") {
+    document.getElementById('promptInput').value = traceData.prompt;
+  }
+
+  document.getElementById('modelSubtitle').textContent =
+    `Checkpoint: ${traceData.ckpt_name} | Tokenizer: Tiktoken BPE (${traceData.vocab_size} vocab) | Input: [All Prefix + 64 Belt] -> Output: [1 Token + 64 Belt]`;
+  document.getElementById('paramBadge').textContent = `${traceData.params_m}M`;
+  document.getElementById('s0AccBadge').textContent = `${(traceData.val_slot0_acc * 100).toFixed(1)}%`;
+  document.getElementById('sampleAccBadge').textContent =
+    traceData.sample_token_acc !== null ? `${(traceData.sample_token_acc * 100).toFixed(1)}%` : "Custom Input";
+
+  renderCrystallizationTable(traceData.crystallization);
+
+  // Jump straight to the first live emitted token tick so the user immediately sees their prompt continuation
+  let firstEmitIdx = 0;
+  for (let i = 0; i < traceData.snapshots.length; i++) {
+    if (!traceData.snapshots[i].is_warmup) {
+      firstEmitIdx = i;
+      break;
+    }
+  }
+  curFrame = firstEmitIdx;
+  const slider = document.getElementById('tickSlider');
+  slider.max = Math.max(0, traceData.snapshots.length - 1);
+  slider.value = curFrame;
+  renderFrame(curFrame);
+}
+
+function renderCrystallizationTable(cryst) {
+  const container = document.getElementById('crystTable');
+  if (!cryst || !cryst.length) {
+    container.innerHTML = '<div style="color:var(--muted);font-size:12px;">No crystallization stats available.</div>';
+    return;
+  }
+  container.innerHTML = cryst.map(row => {
+    const pct = Math.round(row.acc * 100);
+    return `
+      <div class="cryst-row">
+        <span><strong>Slot ${row.slot}</strong> (${row.zone})</span>
+        <span>
+          H=${row.entropy.toFixed(2)} | <strong>${pct}%</strong>
+          <span class="cryst-bar-wrap"><span class="cryst-bar" style="width:${pct}%"></span></span>
+        </span>
+      </div>
+    `;
+  }).join('');
+}
+
+function onScrub(val) {
+  curFrame = parseInt(val, 10);
+  renderFrame(curFrame);
+}
+
+function stepTick(delta) {
+  if (!traceData || !traceData.snapshots.length) return;
+  playing = false;
+  document.getElementById('playBtn').textContent = 'Play';
+  curFrame = (curFrame + delta + traceData.snapshots.length) % traceData.snapshots.length;
+  document.getElementById('tickSlider').value = curFrame;
+  renderFrame(curFrame);
+}
+
+function togglePlay() {
+  playing = !playing;
+  document.getElementById('playBtn').textContent = playing ? 'Pause' : 'Play';
+}
+
+function updateSpeed() {
+  intervalMs = parseInt(document.getElementById('speedSelect').value, 10);
+  if (timerHandle) clearInterval(timerHandle);
+  timerHandle = setInterval(tickLoop, intervalMs);
+}
+
+function renderFrame(idx) {
+  if (!traceData || !traceData.snapshots.length) return;
+  const snap = traceData.snapshots[idx];
+  document.getElementById('scrubberLabel').textContent =
+    `Tick ${idx + 1}/${traceData.snapshots.length} (Pos t=${snap.tick})`;
+  document.getElementById('pauseBadge').textContent = `${Math.round(snap.back_pause_pct)}% Paused`;
+
+  document.getElementById('tickModeLabel').textContent = snap.is_warmup
+    ? `[PRE-BAKING WARMUP: Clamping Input "${snap.emitted_piece}" while Belt Plans Ahead]`
+    : `[LIVE BAKING: Slot 0 Popped "${snap.emitted_piece}" | Back 75% Belt ${Math.round(snap.back_pause_pct)}% Paused]`;
+
+  document.getElementById('storyPromptSpan').textContent = traceData.prompt;
+
+  const genPieces = [];
+  for (let i = 0; i <= idx; i++) {
+    if (!traceData.snapshots[i].is_warmup) {
+      genPieces.push(traceData.snapshots[i].emitted_piece);
+    }
+  }
+
+  const latestSpan = document.getElementById('storyLatestSpan');
+  if (genPieces.length > 0) {
+    const prevText = genPieces.slice(0, -1).join('');
+    const lastPiece = genPieces[genPieces.length - 1];
+    document.getElementById('storyGenSpan').textContent = prevText;
+    latestSpan.style.display = 'inline';
+    latestSpan.textContent = lastPiece;
+  } else {
+    document.getElementById('storyGenSpan').textContent = "";
+    latestSpan.style.display = 'none';
+  }
+
+  const previewWords = snap.slots.slice(1, 7).map(s => s.piece).join('');
+  document.getElementById('storyPreviewSpan').textContent = `  [next on belt: ${previewWords} ...]`;
+
+  const gtBox = document.getElementById('groundTruthBox');
+  if (traceData.ground_truth_continuation) {
+    gtBox.innerHTML = `<strong>Reference Continuation:</strong> ${traceData.ground_truth_continuation}`;
+  } else {
+    gtBox.innerHTML = `<strong>Dynamic Prompt Active:</strong> Continuing "${traceData.prompt.slice(0, 64)}..." in real time`;
+  }
+
+  const row = document.getElementById('conveyorRow');
+  row.innerHTML = '';
+  for (let k = 0; k < snap.slots.length; k++) {
+    const s = snap.slots[k];
+    const zoneClass = k === 0 ? 'slot-zero zone-front' : (k < 16 ? 'zone-front' : (k < 22 ? 'zone-mid' : 'zone-back'));
+    const pct = Math.round(s.prob * 100);
+    const pausePct = Math.round((1.0 - s.shift_prob) * 100);
+    const div = document.createElement('div');
+    div.className = `dough-ball ${zoneClass}`;
+    const label = k === 0 ? 'SLOT 0 (POP)' : (k === 15 ? 'SLOT 15 (EXPAND)' : (k >= 16 ? `Plan ${k}` : `Conv ${k}`));
+    const top2 = s.top3.slice(1, 3).map(c => `${c.w.trim() || c.w}(${Math.round(c.p*100)}%)`).join(', ');
+    const motionTag = k < 16 ? `Shift 100%` : `Pause ${pausePct}%`;
+    div.innerHTML = `
+      <div class="ball-top"><span>${label}</span><span>${pct}%</span></div>
+      <div class="ball-word">${s.piece.trim() || s.piece}</div>
+      <div class="ball-cand">2nd: ${top2 || '-'}</div>
+      <div class="ball-cand">${motionTag} | H=${s.entropy.toFixed(1)}</div>
+      <div class="conf-bar"><div class="conf-fill" style="width:${pct}%"></div></div>
+    `;
+    row.appendChild(div);
+  }
+}
+
+function tickLoop() {
+  if (playing && traceData && traceData.snapshots.length > 0) {
+    curFrame = (curFrame + 1) % traceData.snapshots.length;
+    document.getElementById('tickSlider').value = curFrame;
+    renderFrame(curFrame);
+  }
+}
+
+timerHandle = setInterval(tickLoop, intervalMs);
+runGenerate(-1);
+</script>
+</body>
+</html>
+"""
+
+
+class ConveyerWebApp:
+    def __init__(
+        self,
+        model: ConveyerDiffusionLM,
+        cfg: Any,
+        ckpt_path: str,
+        ckpt_meta: dict[str, Any],
+        val_ds: Any,
+        tokenizer: Any,
+        device: torch.device,
+    ) -> None:
+        self.model = model
+        self.cfg = cfg
+        self.ckpt_path = ckpt_path
+        self.ckpt_meta = ckpt_meta
+        self.val_ds = val_ds
+        self.tokenizer = tokenizer
+        self.device = device
+        self.cached_cryst = self._compute_initial_stats()
+
+    def _compute_initial_stats(self) -> list[dict[str, Any]]:
+        metrics = self.ckpt_meta.get("metrics", {})
+        acc_curve = metrics.get("val_per_slot_acc")
+        ent_curve = metrics.get("val_per_slot_entropy")
+        if not acc_curve or not ent_curve:
+            val_loader = DataLoader(
+                self.val_ds,
+                batch_size=min(32, len(self.val_ds)),
+                shuffle=False,
+                collate_fn=lambda b: collate_conveyer_batch(b, pad_id=self.tokenizer.pad_id),
+            )
+            c = evaluate_crystallization_curve(
+                self.model, self.cfg, val_loader, self.device, pad_id=self.tokenizer.pad_id, max_batches=2
+            )
+            acc_curve = c["per_slot_acc"]
+            ent_curve = c["per_slot_entropy"]
+
+        B = len(acc_curve)
+        slots_to_show = [0, 3, 8, 15, 24, 36, 48, B - 1]
+        rows: list[dict[str, Any]] = []
+        seen: set[int] = set()
+        for k in slots_to_show:
+            k = min(k, B - 1)
+            if k in seen:
+                continue
+            seen.add(k)
+            zone = (
+                "Slot 0 Pop"
+                if k == 0
+                else ("Front Conveyor" if k < 16 else ("Mid Expansion" if k == 15 else "Back 75% Plan"))
+            )
+            rows.append(
+                {
+                    "slot": k,
+                    "zone": zone,
+                    "acc": float(acc_curve[k]),
+                    "entropy": float(ent_curve[k]),
+                }
+            )
+        return rows
+
+    def compute_trace(
+        self,
+        prompt_text: str | None,
+        sample_idx: int,
+        overlap: int,
+        bake_steps: int,
+        max_tokens: int,
+        temperature: float = 0.0,
+    ) -> dict[str, Any]:
+        gt_continuation: str | None = None
+        gt_target_ids: list[int] = []
+
+        if prompt_text is not None and prompt_text.strip() != "":
+            encoded = self.tokenizer.encode(prompt_text.strip())
+            prompt_ids = torch.tensor([encoded], dtype=torch.long, device=self.device)
+            for i in range(min(300, len(self.val_ds))):
+                seq = self.val_ds[i]["input_ids"]
+                if seq[: len(encoded)].tolist() == encoded:
+                    rem = [
+                        int(t)
+                        for t in seq[len(encoded) :].tolist()
+                        if int(t) not in (self.tokenizer.pad_id, self.tokenizer.eos_id)
+                    ]
+                    gt_target_ids = rem[:max_tokens]
+                    gt_continuation = self.tokenizer.decode(gt_target_ids)
+                    break
+        else:
+            idx = max(0, min(sample_idx, len(self.val_ds) - 1))
+            item = self.val_ds[idx]
+            plen = int(item["prompt_len"])
+            seq = item["input_ids"]
+            prompt_ids = seq[:plen].unsqueeze(0).to(self.device)
+            rem = [int(t) for t in seq[plen:].tolist() if int(t) not in (self.tokenizer.pad_id, self.tokenizer.eos_id)]
+            gt_target_ids = rem[:max_tokens]
+            gt_continuation = self.tokenizer.decode(gt_target_ids)
+
+        full_out, snapshots = generate_conveyer(
+            model=self.model,
+            prompt_tokens=prompt_ids,
+            max_new_tokens=max_tokens,
+            prompt_overlap=overlap,
+            bake_steps=bake_steps,
+            temperature=temperature,
+            eos_token_id=self.tokenizer.eos_id,
+            deterministic_spawn=(temperature <= 1e-5),
+        )
+
+        plen_actual = prompt_ids.size(1)
+        gen_ids = full_out[0, plen_actual:].tolist()
+
+        sample_token_acc: float | None = None
+        if gt_target_ids:
+            comp_len = min(len(gen_ids), len(gt_target_ids))
+            if comp_len > 0:
+                matches = sum(1 for a, b in zip(gen_ids[:comp_len], gt_target_ids[:comp_len]) if a == b)
+                sample_token_acc = matches / comp_len
+
+        F_slots = min(self.model.front_slots, bake_steps)
+        serialized_snaps: list[dict[str, Any]] = []
+        for s in snapshots:
+            slots_info = []
+            B = len(s.slot_top1_tokens)
+            for k in range(B):
+                top3_list = [
+                    {"w": self.tokenizer.id_to_piece(tid), "p": round(float(pr), 4)}
+                    for tid, pr in zip(s.slot_top3_tokens[k], s.slot_top3_probs[k])
+                ]
+                slots_info.append(
+                    {
+                        "k": k,
+                        "piece": self.tokenizer.id_to_piece(s.slot_top1_tokens[k]),
+                        "prob": round(float(s.slot_top1_probs[k]), 4),
+                        "entropy": round(float(s.slot_entropies[k]), 3),
+                        "gate": round(float(s.slot_gate_means[k]), 3),
+                        "shift_prob": round(float(s.slot_shift_probs[k]), 3),
+                        "top3": top3_list,
+                    }
+                )
+            back_pause = (
+                100.0 * (1.0 - sum(s.slot_shift_probs[F_slots:]) / max(1, B - F_slots))
+                if B > F_slots
+                else 0.0
+            )
+            serialized_snaps.append(
+                {
+                    "tick": s.tick,
+                    "is_warmup": s.is_warmup,
+                    "emitted_piece": self.tokenizer.id_to_piece(s.emitted_token),
+                    "mean_gate": sum(s.slot_gate_means) / max(1, B),
+                    "back_pause_pct": back_pause,
+                    "slots": slots_info,
+                }
+            )
+
+        metrics = self.ckpt_meta.get("metrics", {})
+        val_s0_acc = float(metrics.get("val_slot0_acc", self.cached_cryst[0]["acc"] if self.cached_cryst else 0.0))
+
+        return {
+            "ckpt_name": self.ckpt_path,
+            "dataset": self.cfg.data.dataset,
+            "vocab_size": self.tokenizer.vocab_size,
+            "params_m": f"{self.model.num_parameters() / 1e6:.2f}",
+            "val_slot0_acc": val_s0_acc,
+            "sample_token_acc": sample_token_acc,
+            "prompt": self.tokenizer.decode(prompt_ids[0].tolist()),
+            "ground_truth_continuation": gt_continuation,
+            "crystallization": self.cached_cryst,
+            "snapshots": serialized_snaps,
+        }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Interactive Web Simulator for the 64-Slot Conveyor-Belt Diffusion LLM")
+    parser.add_argument("--ckpt", type=str, default=None, help="Path to checkpoint (.pt)")
+    parser.add_argument("--config", type=str, default="configs/tinystories.yaml", help="Fallback config if no checkpoint")
+    parser.add_argument("--host", type=str, default=os.environ.get("HOST", "127.0.0.1"))
+    parser.add_argument("--port", type=int, default=int(os.environ.get("PORT", "8766")))
+    args = parser.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    ckpt_path = args.ckpt
+    if ckpt_path is None:
+        if Path("runs/tinystories/best.pt").exists():
+            ckpt_path = "runs/tinystories/best.pt"
+        elif Path("runs/synthetic/best.pt").exists():
+            ckpt_path = "runs/synthetic/best.pt"
+        else:
+            ckpt_path = "runs/tinystories/best.pt"
+
+    try:
+        model, cfg, ckpt_meta = load_checkpoint(ckpt_path, device=device)
+    except Exception:
+        cfg = load_config(args.config)
+        model = ConveyerDiffusionLM(cfg.model).to(device)
+        ckpt_path = "untrained"
+        ckpt_meta = {}
+
+    _, val_ds, tokenizer = build_dataset_and_tokenizer(cfg, infer_only=True)
+    app = ConveyerWebApp(
+        model=model,
+        cfg=cfg,
+        ckpt_path=ckpt_path,
+        ckpt_meta=ckpt_meta,
+        val_ds=val_ds,
+        tokenizer=tokenizer,
+        device=device,
+    )
+
+    class Handler(BaseHTTPRequestHandler):
+        def do_GET(self) -> None:
+            parsed = urlparse(self.path)
+            if parsed.path == "/api/trace":
+                qs = parse_qs(parsed.query)
+                prompt_text = qs.get("prompt", [None])[0]
+                sample_idx = int(qs.get("sample", ["0"])[0])
+                overlap = int(qs.get("overlap", ["12"])[0])
+                bake_steps = int(qs.get("bake_steps", ["64"])[0])
+                tokens = int(qs.get("tokens", ["32"])[0])
+                temp = float(qs.get("temp", ["0.0"])[0])
+                payload = json.dumps(
+                    app.compute_trace(prompt_text, sample_idx, overlap, bake_steps, tokens, temperature=temp)
+                ).encode("utf-8")
+                self.send_response(200)
+                self.send_header("Content-Type", "application/json")
+                self.send_header("Cache-Control", "no-store, no-cache, must-revalidate")
+                self.send_header("Content-Length", str(len(payload)))
+                self.end_headers()
+                self.wfile.write(payload)
+            else:
+                body = HTML_PAGE.encode("utf-8")
+                self.send_response(200)
+                self.send_header("Content-Type", "text/html; charset=utf-8")
+                self.send_header("Cache-Control", "no-store")
+                self.send_header("Content-Length", str(len(body)))
+                self.end_headers()
+                self.wfile.write(body)
+
+        def log_message(self, format: str, *args: Any) -> None:
+            return
+
+    server = HTTPServer((args.host, args.port), Handler)
+    print(f"Conveyor-Belt Diffusion LLM Web Simulator running at http://{args.host}:{args.port}", flush=True)
+    server.serve_forever()
+
+
+if __name__ == "__main__":
+    main()

tests/test_conveyer.py

@@ -0,0 +1,254 @@
+from __future__ import annotations
+
+import torch
+
+from conveyer.config import ConveyerConfig, ModelConfig, load_config
+from conveyer.data.synthetic import SyntheticLookaheadDataset
+from conveyer.infer import generate_conveyer, stream_conveyer_generation
+from conveyer.losses import build_horizon_weights, compute_unrolled_conveyer_loss, seed_corrupted_belt
+from conveyer.model import BeltState, ConveyerDiffusionLM
+
+
+def _small_model_cfg(
+    belt_size: int = 64,
+    use_thought_highway: bool = True,
+    causal_belt_only: bool = False,
+) -> ModelConfig:
+    return ModelConfig(
+        vocab_size=64,
+        d_model=64,
+        n_heads=4,
+        n_layers=2,
+        d_ff=128,
+        max_seq_len=256,
+        dropout=0.0,
+        belt_size=belt_size,
+        tau_min=0.5,
+        tau_max=4.0,
+        tau_gamma=1.5,
+        gate_bias_init=1.0,
+        use_thought_highway=use_thought_highway,
+        causal_belt_only=causal_belt_only,
+    )
+
+
+def test_prefix_kv_cache_chunked_equivalence() -> None:
+    """
+    Verifies that encoding a prefix in one shot (`prefill_prefix(x_{0:P})`) produces the exact same
+    KV cache and downstream belt logits as incremental token-by-token `append_committed_token`.
+    """
+    torch.manual_seed(123)
+    cfg = _small_model_cfg(belt_size=16)
+    model = ConveyerDiffusionLM(cfg).eval()
+
+    prefix = torch.randint(1, 60, (2, 10), dtype=torch.long)
+    belt = model.init_empty_belt(batch_size=2, tick=10, device=prefix.device, deterministic=True)
+
+    # Full prefill in one shot
+    kv_full = model.prefill_prefix(prefix)
+    out_full = model.forward_belt(belt, kv_cache=kv_full)
+
+    # Incremental prefill: first 4 tokens, then 1 by 1
+    kv_inc = model.prefill_prefix(prefix[:, :4])
+    for t in range(4, 10):
+        kv_inc = model.append_committed_token(prefix[:, t], kv_inc)
+    out_inc = model.forward_belt(belt, kv_cache=kv_inc)
+
+    assert kv_full.seq_len == kv_inc.seq_len == 10
+    for l_idx in range(cfg.n_layers):
+        assert torch.allclose(kv_full.k_cache[l_idx], kv_inc.k_cache[l_idx], atol=1e-5)
+        assert torch.allclose(kv_full.v_cache[l_idx], kv_inc.v_cache[l_idx], atol=1e-5)
+    assert torch.allclose(out_full.logits, out_inc.logits, atol=1e-5)
+
+
+def test_slot0_lexical_anchoring_alpha_zero() -> None:
+    """
+    Verifies that alpha_0 == 0.0 at Position 0:
+    Perturbing `m_thought[:, 0, :]` at Slot 0 must have ZERO effect on the composed dough input at Slot 0,
+    forcing Slot 0 to rely 100% on its crisp lexical projection `e_soft[:, 0, :]` before popping out.
+    Meanwhile, perturbing `m_thought[:, k, :]` for k > 0 MUST change `dough[:, k, :]`.
+    """
+    torch.manual_seed(42)
+    cfg = _small_model_cfg(belt_size=64, use_thought_highway=True)
+    model = ConveyerDiffusionLM(cfg).eval()
+
+    belt1 = model.init_empty_belt(batch_size=2, tick=5, device=torch.device("cpu"), deterministic=True)
+    m_mod = belt1.m_thought.clone()
+    m_mod[:, 0, :] += 100.0  # Massive perturbation at Slot 0 thought vector
+    belt2 = BeltState(e_soft=belt1.e_soft.clone(), m_thought=m_mod, tick=5)
+
+    dough1 = model.compose_dough_inputs(belt1)
+    dough2 = model.compose_dough_inputs(belt2)
+
+    assert float(model.alpha_schedule[0].item()) == 0.0
+    assert float(model.alpha_schedule[-1].item()) == 1.0
+    assert torch.allclose(dough1[:, 0, :], dough2[:, 0, :], atol=1e-6)
+
+    # Now perturb Slot 63 thought vector -> dough[:, 63, :] must change
+    m_mod63 = belt1.m_thought.clone()
+    m_mod63[:, 63, :] += 10.0
+    belt3 = BeltState(e_soft=belt1.e_soft.clone(), m_thought=m_mod63, tick=5)
+    dough3 = model.compose_dough_inputs(belt3)
+    assert not torch.allclose(dough1[:, 63, :], dough3[:, 63, :], atol=1e-4)
+
+
+def test_bidirectional_belt_attention_vs_causal_ar() -> None:
+    """
+    Verifies that in Conveyor mode (`causal_belt_only=False`), Slot 0's output logits actively depend on
+    downstream dough balls at the back of the belt (e.g., Slot 63), whereas in AR baseline mode
+    (`causal_belt_only=True`), Slot 0 is completely blind to Slots 1..63.
+    """
+    torch.manual_seed(7)
+    # 1. Conveyor bidirectional belt mode
+    cfg_bi = _small_model_cfg(belt_size=64, causal_belt_only=False)
+    model_bi = ConveyerDiffusionLM(cfg_bi).eval()
+
+    belt_a = model_bi.init_empty_belt(batch_size=1, tick=4, device=torch.device("cpu"), deterministic=True)
+    belt_b = BeltState(e_soft=belt_a.e_soft.clone(), m_thought=belt_a.m_thought.clone(), tick=4)
+    belt_b.e_soft[:, 63, :] += 5.0  # Perturb back of the belt (Slot 63)
+
+    out_bi_a = model_bi.forward_belt(belt_a)
+    out_bi_b = model_bi.forward_belt(belt_b)
+    # In bidirectional conveyor mode, Slot 0 uses the mental anchor from Slot 63
+    assert not torch.allclose(out_bi_a.slot0_logits, out_bi_b.slot0_logits, atol=1e-5)
+
+    # 2. Causal AR baseline mode
+    cfg_ar = _small_model_cfg(belt_size=64, causal_belt_only=True)
+    model_ar = ConveyerDiffusionLM(cfg_ar).eval()
+    out_ar_a = model_ar.forward_belt(belt_a)
+    out_ar_b = model_ar.forward_belt(belt_b)
+    # In causal AR mode, Slot 0 cannot see Slot 63
+    assert torch.allclose(out_ar_a.slot0_logits, out_ar_b.slot0_logits, atol=1e-6)
+
+
+def test_left_shift_mechanics_and_elastic_planning_zone() -> None:
+    """
+    Verifies both:
+    1. Strict 1-by-1 conveyor shift (`use_elastic_planning=False`), where slots 1..63 shift into 0..62.
+    2. Elastic Two-Zone Planning Belt (`use_elastic_planning=True`), where the front 25% (`0..14`) shifts
+       strictly from `1..15`, Slot 15 blends mid-belt sentence expansion (`expand_cand`), and the back 75%
+       (`16..63`) pauses elastically via `shift_prob < 1.0`.
+    """
+    torch.manual_seed(99)
+    # 1. Strict conveyor mode
+    cfg_strict = _small_model_cfg(belt_size=64)
+    cfg_strict.use_elastic_planning = False
+    model_strict = ConveyerDiffusionLM(cfg_strict).eval()
+
+    belt = model_strict.init_empty_belt(batch_size=1, tick=0, device=torch.device("cpu"), deterministic=True)
+    step_out = model_strict.forward_belt(belt)
+    shifted = model_strict.shift_belt_left(step_out, current_tick=0, deterministic_spawn=True)
+
+    assert shifted.tick == 1
+    assert shifted.e_soft.shape == (1, 64, cfg_strict.d_model)
+    assert torch.allclose(shifted.e_soft[:, :63, :], step_out.e_soft[:, 1:, :], atol=1e-6)
+    assert torch.allclose(shifted.m_thought[:, :63, :], step_out.m_updated[:, 1:, :], atol=1e-6)
+
+    # 2. Elastic Two-Zone Planning mode (Front 25% = 16 slots, Back 75% = 48 slots)
+    cfg_elastic = _small_model_cfg(belt_size=64)
+    cfg_elastic.use_elastic_planning = True
+    cfg_elastic.front_conveyer_slots = 16
+    model_elastic = ConveyerDiffusionLM(cfg_elastic).eval()
+
+    step_el = model_elastic.forward_belt(belt)
+    shifted_el = model_elastic.shift_belt_left(step_el, current_tick=0, deterministic_spawn=True)
+
+    # Front 15 slots (0..14) shift 100% from (1..15)
+    assert torch.allclose(shifted_el.e_soft[:, :15, :], step_el.e_soft[:, 1:16, :], atol=1e-6)
+    assert torch.allclose(shifted_el.m_thought[:, :15, :], step_el.m_updated[:, 1:16, :], atol=1e-6)
+    # Back 75% planning slots have shift_prob < 0.5 (paused > 50% in-place)
+    assert float(step_el.shift_prob[:, 16:].mean().item()) < 0.50
+
+
+def test_prompt_overlap_prebaking_clamps_input_and_emits() -> None:
+    """
+    Verifies the `n - W_pre` cold-start solution:
+    When `prompt_overlap = 6` on a prompt of length `n = 10`:
+      - Warm-up runs for ticks `t = 4, 5, 6, 7, 8, 9` (6 warm-up ticks), clamping emitted tokens to `prompt[t]`.
+      - Generation runs for `max_new_tokens = 8` ticks (`t = 10..17`), returning a sequence of length `10 + 8 = 18`.
+    """
+    torch.manual_seed(11)
+    cfg = _small_model_cfg(belt_size=32)
+    model = ConveyerDiffusionLM(cfg).eval()
+
+    prompt = torch.arange(5, 15, dtype=torch.long).unsqueeze(0)  # length n = 10
+    full_seq, snapshots = generate_conveyer(
+        model=model,
+        prompt_tokens=prompt,
+        max_new_tokens=8,
+        prompt_overlap=6,
+        bake_steps=32,
+        temperature=0.0,
+        deterministic_spawn=True,
+    )
+
+    assert full_seq.shape == (1, 18)
+    assert torch.equal(full_seq[:, :10], prompt)
+    assert len(snapshots) == 6 + 8
+
+    # First 6 snapshots must be warm-up ticks at t = 4..9 with clamped prompt tokens
+    for i in range(6):
+        assert snapshots[i].is_warmup is True
+        assert snapshots[i].tick == 4 + i
+        assert snapshots[i].emitted_token == int(prompt[0, 4 + i].item())
+
+    # Next 8 snapshots must be live emissions at t = 10..17
+    for j in range(8):
+        snap = snapshots[6 + j]
+        assert snap.is_warmup is False
+        assert snap.tick == 10 + j
+        assert snap.emitted_token == int(full_seq[0, 10 + j].item())
+
+
+def test_unrolled_bptt_gradient_flow_and_slot63_zero_leakage() -> None:
+    """
+    Verifies that:
+    1. `seed_corrupted_belt` puts zero target signal at Slot 63 (pure noise).
+    2. `compute_unrolled_conveyer_loss` backpropagates non-zero gradients into both the lexical embeddings
+       and the continuous thought highway (`thought_cand`, `thought_gate`, `slot_emb`, `m_init`).
+    """
+    torch.manual_seed(21)
+    full_cfg = ConveyerConfig(model=_small_model_cfg(belt_size=16, use_thought_highway=True))
+    full_cfg.train.unroll_steps = 4
+    model = ConveyerDiffusionLM(full_cfg.model).train()
+
+    # Check Slot B-1 has zero target signal in seed_corrupted_belt
+    torch.manual_seed(55)
+    t_a = torch.ones((1, 16), dtype=torch.long) * 5
+    b_a = seed_corrupted_belt(model, t_a, tick_start=2, cold_drop_prob=0.0)
+    torch.manual_seed(55)
+    t_b = torch.ones((1, 16), dtype=torch.long) * 5
+    t_b[:, -1] = 42  # Change target at the last slot (B-1)
+    b_b = seed_corrupted_belt(model, t_b, tick_start=2, cold_drop_prob=0.0)
+    assert torch.allclose(b_a.e_soft[:, -1, :], b_b.e_soft[:, -1, :], atol=1e-6)
+
+    # Check gradient flow across unrolled BPTT
+    seq = torch.randint(4, 50, (3, 40), dtype=torch.long)
+    plen = torch.tensor([8, 8, 8], dtype=torch.long)
+    metrics = compute_unrolled_conveyer_loss(
+        model=model,
+        seq_tokens=seq,
+        cfg=full_cfg,
+        prompt_lengths=plen,
+        tick_start=6,
+    )
+    metrics.loss.backward()
+
+    assert model.tok_emb.weight.grad is not None and model.tok_emb.weight.grad.abs().sum() > 0
+    assert model.slot_emb.weight.grad is not None and model.slot_emb.weight.grad.abs().sum() > 0
+    assert model.thought_cand.weight.grad is not None and model.thought_cand.weight.grad.abs().sum() > 0
+    assert model.thought_gate.weight.grad is not None and model.thought_gate.weight.grad.abs().sum() > 0
+
+
+def test_synthetic_lookahead_dataset_and_config_loader() -> None:
+    cfg = load_config("configs/synthetic.yaml", ["train.max_steps=10", "data.num_train=60"])
+    assert cfg.train.max_steps == 10
+    assert cfg.data.num_train == 60
+
+    ds = SyntheticLookaheadDataset(num_samples=12, seq_len=96, task="all", vocab_size=128)
+    assert len(ds) == 12
+    for i in range(6):
+        item = ds[i]
+        assert item["input_ids"].shape == (96,)
+        assert 4 < item["prompt_len"] < 60
+        assert item["task_id"] in (0, 1, 2)