OnnxTokenModule + load_onnx front door; repl slims to a front-end; delete token_model

- onnx_load.py gains OnnxTokenModule (token-I/O counterpart of
  OnnxHeadlessModule: same static-cache protocol, plus vocab
  tokenize/detokenize and an argmax generate loop ported from repl)
  and load_onnx(path), which dispatches on the sidecar format key.
- repl.py drops its private _Vocab/_Model/_load/generate quartet and
  loads via load_onnx (it now depends on torch via onnx_load; the old
  standalone/no-torch docstring claim is gone). run_once/run_repl
  signatures unchanged; root onnx_repl.py untouched.
- token_model.py (compile_token/CompiledToken) deleted — zero call
  sites in torchwright or torchwright_doom.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: physicsrob

tests/compile/forward/test_module.py

@@ -203,7 +203,7 @@ def test_token_onnx_decode_step_matches_full_prefill():
 
 
 def test_token_onnx_autoregressive_1digit():
-    from torchwright.compiler.repl import _load, generate
+    from torchwright.compiler.onnx_load import load_onnx
 
     output_node, pos_encoding, embedding = _build_1digit()
 
@@ -219,10 +219,10 @@ def test_token_onnx_autoregressive_1digit():
             verbose=False,
         )
 
-        model = _load(onnx_path)
+        model = load_onnx(onnx_path)
         test_cases = [("1+1\n", "2"), ("2+3\n", "5"), ("4+5\n", "9")]
         for input_str, expected in test_cases:
-            result = "".join(generate(model, input_str))
+            result = "".join(model.generate(input_str))
             assert (
                 result == expected
             ), f"{input_str}: expected {expected!r}, got {result!r}"
@@ -234,7 +234,7 @@ def test_token_onnx_autoregressive_1digit():
 
 
 def test_token_onnx_autoregressive_3digit():
-    from torchwright.compiler.repl import _load, generate
+    from torchwright.compiler.onnx_load import load_onnx
 
     output_node, pos_encoding, embedding = create_network_parts()
 
@@ -250,7 +250,7 @@ def test_token_onnx_autoregressive_3digit():
             verbose=False,
         )
 
-        model = _load(onnx_path)
+        model = load_onnx(onnx_path)
         test_cases = [
             ("1+2\n", "3"),
             ("12+34\n", "46"),
@@ -259,7 +259,7 @@ def test_token_onnx_autoregressive_3digit():
             ("456+123\n", "579"),
         ]
         for input_str, expected in test_cases:
-            result = "".join(generate(model, input_str))
+            result = "".join(model.generate(input_str))
             assert (
                 result == expected
             ), f"{input_str}: expected {expected!r}, got {result!r}"
@@ -271,7 +271,7 @@ def test_token_onnx_autoregressive_3digit():
 
 
 def test_token_onnx_sidecar_schema_and_metadata():
-    from torchwright.compiler.repl import _load
+    from torchwright.compiler.onnx_load import OnnxTokenModule, load_onnx
 
     output_node, pos_encoding, embedding = _build_1digit()
 
@@ -293,7 +293,8 @@ def test_token_onnx_sidecar_schema_and_metadata():
         assert meta["format"] == TOKEN_META_FORMAT
         assert meta["vocab"] == embedding.tokenizer.vocab
 
-        model = _load(onnx_path)
+        model = load_onnx(onnx_path)
+        assert isinstance(model, OnnxTokenModule)
         assert model.n_layers > 0
         assert len(model.per_layer_n_heads) == model.n_layers
         assert all(nh <= D // D_HEAD for nh in model.per_layer_n_heads)

torchwright/compiler/onnx_load.py

@@ -1,10 +1,18 @@
-"""Runtime loader for headless ONNX models.
+"""Runtime loaders for torchwright ONNX exports.
 
-Provides ``OnnxHeadlessModule`` — an ``onnxruntime``-backed callable
-that speaks the static-cache prefill/decode protocol produced by
-:func:`torchwright.compiler.export.compile_headless_to_onnx` — plus a
-:class:`HeadlessRuntime` :class:`typing.Protocol` that describes the
-shared interface with the in-memory
+:func:`load_onnx` is the front door: it reads the ``<stem>.meta.json``
+sidecar and dispatches on its format key to one of two loaders, both
+``onnxruntime``-backed callables speaking the static-cache
+prefill/decode protocol:
+
+- ``OnnxHeadlessModule`` — float I/O (``torchwright.headless.v1``),
+  produced by :func:`torchwright.compiler.export.compile_headless_to_onnx`.
+- ``OnnxTokenModule`` — token I/O (``torchwright.token.v1``), produced
+  by :func:`torchwright.compiler.export.compile_to_onnx`; adds the
+  vocab tokenizer and an argmax :meth:`OnnxTokenModule.generate` loop.
+
+A :class:`HeadlessRuntime` :class:`typing.Protocol` describes the
+interface shared with the in-memory
 :class:`torchwright.compiler.export.CompiledHeadless`.
 
 Two usage shapes:
@@ -23,19 +31,24 @@
 import json
 import os
 from dataclasses import dataclass
-from typing import List, Optional, Protocol, Tuple, Union
+from typing import Iterator, List, Optional, Protocol, Tuple, Union
 
 import numpy as np
 import torch
 
-from torchwright.compiler.export import HEADLESS_META_FORMAT, meta_path_for
+from torchwright.compiler.export import (
+    HEADLESS_META_FORMAT,
+    TOKEN_META_FORMAT,
+    meta_path_for,
+)
 
 
 def discover_cache_stride(inputs: dict, sidecar_stride, onnx_path) -> int:
     """Resolve the full static slot count ``S`` for a cached-protocol model.
 
-    Shared by all three loaders (``OnnxHeadlessModule``, ``repl``, and
-    torchwright_doom's ``OnnxTokenRuntime``).  Dual path:
+    Shared by all three loaders (``OnnxHeadlessModule``,
+    ``OnnxTokenModule``, and torchwright_doom's ``OnnxTokenRuntime``).
+    Dual path:
 
     - ``past_K_0`` first dim is a static int — a pre-bucketing export
       (old compile caches): use it.  Such models can only bind the full
@@ -287,3 +300,245 @@ def __call__(self, inputs: torch.Tensor) -> torch.Tensor:
 
     def eval(self) -> "OnnxHeadlessModule":
         return self
+
+
+class OnnxTokenModule:
+    """Loads a token-I/O cached ONNX model and exposes it as a callable.
+
+    The token counterpart of :class:`OnnxHeadlessModule`: same static-cache
+    protocol, but the sequence input is ``token_ids`` (int64) instead of a
+    float row, the sequence output is ``logits``, and the sidecar carries
+    the vocab — so the module can tokenize/detokenize and run an argmax
+    :meth:`generate` loop directly.
+
+    Args:
+        onnx_path: Path to the ``.onnx`` file.  A sidecar
+            ``<stem>.meta.json`` with format ``torchwright.token.v1``
+            must exist alongside it.
+        providers: ``onnxruntime`` execution providers list.  Defaults
+            to CPU.
+    """
+
+    def __init__(self, onnx_path: str, providers=None) -> None:
+        import onnxruntime as ort
+
+        meta_path = meta_path_for(onnx_path)
+        if not os.path.exists(meta_path):
+            raise FileNotFoundError(
+                f"Missing sidecar {meta_path}. Re-export with "
+                f"compile_to_onnx to produce it."
+            )
+        with open(meta_path) as f:
+            meta = json.load(f)
+
+        fmt = meta.get("format")
+        if fmt != TOKEN_META_FORMAT:
+            raise ValueError(
+                f"{meta_path}: unexpected format {fmt!r}, "
+                f"expected {TOKEN_META_FORMAT!r}"
+            )
+        self.vocab: List[str] = list(meta["vocab"])
+        self._token_to_id = {t: i for i, t in enumerate(self.vocab)}
+        self.metadata: dict = dict(meta.get("extra") or {})
+
+        self._session = ort.InferenceSession(
+            onnx_path,
+            providers=providers or ["CPUExecutionProvider"],
+        )
+
+        # KV cache topology discovery — identical to OnnxHeadlessModule
+        # (the two exports share the cached-protocol K/V plumbing).
+        inputs = {inp.name: inp for inp in self._session.get_inputs()}
+        self._n_layers = sum(1 for name in inputs if name.startswith("past_K_"))
+        assert (
+            self._n_layers > 0
+        ), f"{onnx_path}: no past_K_* inputs — is this a cached-protocol model?"
+        self._cache_stride = discover_cache_stride(
+            inputs, meta.get("cache_stride"), onnx_path
+        )
+        self._per_layer_n_heads = [
+            int(inputs[f"past_K_{i}"].shape[1]) for i in range(self._n_layers)
+        ]
+        self._d_head = int(inputs["past_K_0"].shape[2])
+
+        self._out_names = ["logits"]
+        for i in range(self._n_layers):
+            self._out_names += [f"delta_K_{i}", f"delta_V_{i}"]
+
+    @property
+    def cache_stride(self) -> int:
+        """The static slot count ``S`` baked into the loaded model."""
+        return self._cache_stride
+
+    @property
+    def n_layers(self) -> int:
+        return self._n_layers
+
+    @property
+    def per_layer_n_heads(self) -> List[int]:
+        return list(self._per_layer_n_heads)
+
+    @property
+    def d_head(self) -> int:
+        return self._d_head
+
+    def token_to_id(self, token: str) -> int:
+        return self._token_to_id.get(token, 0)  # 0 = <unk>
+
+    def id_to_token(self, token_id: int) -> str:
+        if 0 <= token_id < len(self.vocab):
+            return self.vocab[token_id]
+        return self.vocab[0]
+
+    def empty_past(self) -> OnnxPast:
+        """Full-S zero-filled sequence-major cache buffers, length 0.
+
+        Zero (not garbage) is load-bearing: slots beyond the committed
+        length are read by the attention with weight exactly 0.0, and
+        ``0 * NaN = NaN``.
+        """
+        S = self._cache_stride
+        past_K = tuple(
+            torch.zeros(S, nh, self._d_head) for nh in self._per_layer_n_heads
+        )
+        past_V = tuple(
+            torch.zeros(S, nh, self._d_head) for nh in self._per_layer_n_heads
+        )
+        return OnnxPast(k=past_K, v=past_V, length=0)
+
+    def step(
+        self,
+        token_ids: torch.Tensor,
+        past: OnnxPast,
+        past_len: Optional[int] = None,
+    ) -> Tuple[torch.Tensor, OnnxPast]:
+        """Run one cached-protocol call and return (logits, new_past).
+
+        Args:
+            token_ids: ``(n_new,)`` int64 tensor of token ids.
+            past: :class:`OnnxPast` from a prior step or
+                :meth:`empty_past`.
+            past_len: Optional explicit base position for the new rows.
+                When ``None`` (default), uses ``past.length``.  Must equal
+                the committed cache length (see
+                :meth:`OnnxHeadlessModule.step`).
+
+        Returns:
+            ``(logits, new_past)`` where ``logits`` is a
+            ``(n_new, vocab_size)`` torch tensor and ``new_past`` shares
+            the (in-place updated) buffers with ``past`` at the advanced
+            committed length.
+        """
+        if not isinstance(past, OnnxPast):
+            raise TypeError(
+                "OnnxTokenModule.step requires an OnnxPast from empty_past() "
+                f"(got {type(past).__name__}) — the static-cache protocol has "
+                "no growable-tuple representation"
+            )
+        assert len(past.k) == self._n_layers
+        assert len(past.v) == self._n_layers
+
+        base = past.length if past_len is None else int(past_len)
+        assert base == past.length, (
+            f"past_len {base} != committed length {past.length}: the static "
+            f"cache derives mask AND pos from cache_position; a trimmed cache "
+            f"with a larger absolute position is not expressible"
+        )
+        n_new = int(token_ids.shape[0])
+        if base + n_new > self._cache_stride:
+            raise RuntimeError(
+                f"static cache overrun: length {base} + n_new {n_new} exceeds "
+                f"cache_stride {self._cache_stride}; re-export with a larger "
+                f"cache_stride"
+            )
+
+        ids_np = token_ids.detach().cpu().numpy().astype(np.int64, copy=False)
+        feeds: dict = {
+            "token_ids": ids_np,
+            "cache_position": np.arange(base, base + n_new, dtype=np.int64),
+        }
+        for i in range(self._n_layers):
+            feeds[f"past_K_{i}"] = (
+                past.k[i].detach().cpu().numpy().astype(np.float32, copy=False)
+            )
+            feeds[f"past_V_{i}"] = (
+                past.v[i].detach().cpu().numpy().astype(np.float32, copy=False)
+            )
+
+        results = self._session.run(self._out_names, feeds)
+        logits = torch.from_numpy(results[0])
+
+        # Persist the per-layer deltas (new rows only) into the owned slots.
+        for i in range(self._n_layers):
+            past.k[i][base : base + n_new] = torch.from_numpy(results[1 + 2 * i])
+            past.v[i][base : base + n_new] = torch.from_numpy(results[1 + 2 * i + 1])
+
+        return logits, OnnxPast(k=past.k, v=past.v, length=base + n_new)
+
+    def __call__(self, token_ids: torch.Tensor) -> torch.Tensor:
+        """Convenience: stateless prefill that discards the cache.
+
+        Equivalent to ``self.step(token_ids, self.empty_past())[0]``.
+        """
+        logits, _ = self.step(token_ids, self.empty_past())
+        return logits
+
+    def generate(
+        self,
+        input_text: str,
+        max_new_tokens: int = 10,
+        bos_token: str = "<bos",
+        eos_token: str = "<eos>",
+    ) -> Iterator[str]:
+        """Autoregressive argmax generation over the cached protocol.
+
+        Prefills ``[bos_token] + list(input_text)``, then decodes one
+        token per :meth:`step`, yielding each generated token string as
+        it is produced.  Stops at ``eos_token`` or ``max_new_tokens``.
+        """
+        tokens = [bos_token] + list(input_text)
+        ids = torch.tensor(
+            [self.token_to_id(t) for t in tokens], dtype=torch.int64
+        )
+
+        logits, past = self.step(ids, self.empty_past())
+        for _ in range(max_new_tokens):
+            next_id = int(logits[-1].argmax())
+            next_token = self.id_to_token(next_id)
+            if next_token == eos_token:
+                break
+            yield next_token
+            logits, past = self.step(
+                torch.tensor([next_id], dtype=torch.int64), past
+            )
+
+    def eval(self) -> "OnnxTokenModule":
+        return self
+
+
+def load_onnx(
+    onnx_path: str, providers=None
+) -> Union[OnnxHeadlessModule, OnnxTokenModule]:
+    """Load any torchwright ONNX export, dispatching on its sidecar format.
+
+    Reads ``<stem>.meta.json`` next to ``onnx_path`` and returns an
+    :class:`OnnxHeadlessModule` (``torchwright.headless.v1``) or an
+    :class:`OnnxTokenModule` (``torchwright.token.v1``).
+    """
+    meta_path = meta_path_for(onnx_path)
+    if not os.path.exists(meta_path):
+        raise FileNotFoundError(
+            f"Missing sidecar {meta_path}. Re-export with compile_to_onnx / "
+            f"compile_headless_to_onnx to produce it."
+        )
+    with open(meta_path) as f:
+        meta = json.load(f)
+    fmt = meta.get("format")
+    if fmt == HEADLESS_META_FORMAT:
+        return OnnxHeadlessModule(onnx_path, providers=providers)
+    if fmt == TOKEN_META_FORMAT:
+        return OnnxTokenModule(onnx_path, providers=providers)
+    raise ValueError(
+        f"{meta_path}: unknown format {fmt!r}; expected "
+        f"{HEADLESS_META_FORMAT!r} or {TOKEN_META_FORMAT!r}"
+    )