Meta: native LLM inference runtime in pytensor_ml (tracking)

[cetagostini]

Goal

Make pytensor_ml able to load and run a local LLM end to end from Python — load → tokenize → prefill → cached decode → sample → stream — with one symbolic model definition compiling through C, Numba, JAX and MLX, and a persistent KV cache instead of prefix replay.

Reference point is llama.cpp's shape (immutable model, mutable context, batch/decode, sampler chain), not its throughput or model coverage. The differentiator stays composability: an inspectable graph inside the scientific Python stack.

Build on what already exists

This is the important constraint. The repo already has the pieces below, and this work extends them in place rather than adding a parallel stack:

Already here	Consequence
serialize_graph / deserialize_graph, type + op registries, OpFromGraph / SymbolicOp / Scan codecs	No second codec. A new op serializes for free if it is a LayerOp with complete __props__; qualname becomes on-disk format
save_network / load_network, InputKind, GRAPH_FORMAT_VERSION	Architecture persistence already solved; schema changes bump the version
from_pretrained(source_format="huggingface") → NotImplementedError	This stub is the loading entry point (PR 10), not a new loader
save_state / load_state(name_map=...) with key/shape/dtype validation	HF renaming uses the existing name_map; only lazy reads, sharding and telemetry are missing
StatefulOp.update_map + non_trainable + collect_non_trainable_updates + function(updates=...)	This is the KV-cache mechanism — BatchNormLayer already proves it. No bespoke session-mutation layer
AttentionLayer bottom-right causal alignment (k_idx <= q_idx + (sk - sq)) + _repeat_kv	Decode-shaped masking and GQA are already correct and reused as-is
predict_db / rewrite_for_prediction, dispatch/ meta-path hook, fused MLX/JAX attention	Quantized kernels and new op lowering register through these
Model.compile_train(loss=...) (#52)	LM training objective already supported — this work is inference-only

Explicitly not in scope because it already exists: a new archive format or reader, a weight-store abstraction, a model-directory layout, a second HF detection path, a rewrite pass manager, a backend-registration mechanism, a SiLU activation (Swish(beta=1)), and Scan serialization.

How to read the list

Each block below is one PR. Boxes inside a block land together; blocks land in listed order within a tier. F-nn are stable feature ids so a PR can say "closes F-19, F-20" without a second issue.

PR bundling rules:

• A format change (GRAPH_FORMAT_VERSION, __props__ of a serialized op) ships with its rejection test in the same PR — never split.
• An op and the rewrite/lowering that selects it ship together; a kernel no compiled graph selects is not a capability.
• Correctness fixtures land before the thing they gate.
• Existing suites (test_pretrained, test_checkpoint, test_serialize, test_attention, rewriting/, dispatch/) must keep passing untouched unless the PR is a deliberate, version-bumped contract change.

Tier A — native Python runtime (release blocker)

PR 1 — Runtime contracts (docs only)

• F-63 load_llm / LLM / Session / Engine roles, error taxonomy, result schema, plan-key schema, import direction
• Record which existing modules are extended vs what is new; freeze Model and from_pretrained signatures

No code. Unblocks every other PR by fixing names and boundaries.

PR 2 — Conformance fixtures

• F-60 Tiny offline decoder fixture: pinned config, weights, tokenizer, prompt tokens, full-prefix logits, cached-step logits
• F-60 Independent oracle that may not import production model/RoPE/mask/cache helpers; mutation tests that must fail on Q/K orientation, RoPE pairing, mask polarity, cache frontier

Reuses tests/test_serialize.py::assert_outputs_roundtrip and tests/test_attention.py conventions.

PR 3 — Missing decoder primitives

• F-13 RMSNorm layer + RMSNormLayer op beside the existing LayerNormLayer
• F-14 RoPE layer/op with position inputs and scaling variants
• F-15 Gated (SwiGLU) MLP beside FeedForward — Swish(beta=1) already provides SiLU
• F-16a Tied embedding/unembedding sharing one parameter
• F-16b LM head with optional logit softcap and fp32-safe logits (sole owner of softcap)
• F-17 Sliding-window/local masking added to the existing AttentionLayer
• F-18 Llama-style block: pre-RMSNorm + GQA + gated MLP
• F-68 GRAPH_FORMAT_VERSION bump, stale-config rejection test, __props__ payload migration, qualname-stability test — required because F-17 changes a serialized op

All new ops are LayerOps with complete __props__, so they serialize as leaves for free. Attention, GQA and decode-aligned causal masking already exist and are reused.

PR 4 — Artifact manifest, config, architecture registry

• F-10 Manifest: local-only resolution, hashes, sizes, dtype/shape/offset bounds, provenance
• F-09a Parse model.safetensors.index.json and validate exact shard/tensor coverage (no tensor reads)
• F-06 HuggingFace config.json → normalized decoder config
• F-07 Architecture builder dispatch table keyed by the already-detected model_type/architectures
• F-08a HF → pytensor_ml parameter-name map as data for load_state(name_map=...)

Extends _looks_like_huggingface/_detect_format; HF detection is not re-implemented.

PR 5 — Tokenizer and chat templates

• F-01 Encode/decode, batch, special tokens, vocab, stop-token metadata
• F-02 Incremental streaming detokenizer with UTF-8/byte-fallback buffering
• F-03 Chat-template rendering, generation prompt, prefill/continuation
• F-04 HF tokenizer artifact adapter (tokenizer.json, tokenizer_config.json, chat_template.jinja)

Nothing tokenizer-related exists today. Chat template is the sole owner of BOS/EOS insertion.

PR 6 — Backend capabilities, compiled plans, resource accounting

• F-28 Capability declaration; unsupported model/backend/dtype combinations fail before token 1
• F-26 Compiled-plan cache keyed by architecture/backend/dtype/shape bucket, extending compile_predict
• F-27 Shape-bucketed prefill and fixed-shape decode compilation
• F-69 Plans are never serialized; keys recomputed over the pre-compilation graph. Test fails if a fused Composite reaches the codec
• F-25 Aggregate host/device byte accounting for weights, plan buffers, cache allocation
• F-29 LLM entries registered into the existing predict_db

Uses the existing pytensorf.function / compile_predict / predict_db / dispatch extension points.

PR 7 — Large-model weight residency

• F-65 Sole owner of variable→value binding and residency/copy telemetry, plus the tensor-handle contract
• F-24 Lazy per-tensor byte materialization beside load_state's whole-archive read
• F-09b Shard iteration for multi-file archives
• F-08b Apply the HF name map via load_state(name_map=...), reporting unmapped/unexpected keys

Extends checkpoint.py / pretrained.py. No new archive format, directory layout, or weight store. The safetensors NumPy path copies bytes, so no zero-copy claim.

PR 8 — KV cache and prefill/decode plans

• F-67 One cache-state interface (logical read/append/gather/mask) that dense and later paged storage both implement
• F-19 KVCacheLayer op declaring update_map() — write-back rides the existing function(..., updates=...) path
• F-20 Cache buffers as named non_trainable state with capacity and length counters
• F-21 Per-step position_ids / cache_length / seq_id, produced once and consumed by RoPE, masking, append
• F-22 Prefill and decode graph builders over one architecture definition
• F-23 Exclude cache buffers from _weight_variables / save_pretrained
• F-71 Cache updates survive rewrite_for_prediction; test fails if updates are dropped after rewriting

Mechanism already exists: StatefulOp.update_map + non_trainable + collect_non_trainable_updates, as BatchNormLayer demonstrates. Dense storage only; no wrap or slot reuse. Bounded append writes, stable compile count.

PR 9 — Llama-family reference architecture

• F-47 Llama architecture module and weight schema composed from PR 3 primitives

Emits distinct prefill/decode graphs; no generation, loader, or scheduler code.

PR 10 — Implement HuggingFace loading

• F-11 Replace the NotImplementedError in from_pretrained(source_format="huggingface") by wiring PR 4 config/registry/name-map and PR 7 lazy loading

This is the existing stub. Real pinned safetensors model loads, binds, prefills, and decodes with all-logit oracle parity.

PR 11 — Generation policy

• F-30 GenerationConfig, greedy and sampled selection
• F-31 Logits processors: temperature, top-k, top-p, min-p, repetition/presence/frequency penalties
• F-32 Stop criteria: EOS set, stop tokens, stop strings, max tokens, finish-reason precedence
• F-70 Session RNG ownership: creation, isolation from the config-carried InputKind.RNG seed, snapshot/restore
• F-33 Per-request seeded draws over that state

Pure policy — no weights, compilation, or cache allocation. Uses existing RNG threading in pytensorf.function.

PR 12 — Public API, sessions, streaming, engine

• F-34 Autoregressive loop driving prefill/decode plans and cache updates
• F-35 Streaming iterator with cancellation and bounded buffering
• F-36 Batched generation with independent per-sequence state
• F-37 Multi-turn session/chat reusing a live cache
• F-38 Structured result: tokens, text, finish reason, usage, timings, compile/copy counters

First real load_llm → generate / stream / chat path. No oracle in the normal import or call path.

PR 13 — Dense cache backend lowering

• F-45 Per-backend lowering of dense cache append/read with bounded write evidence

Follows the existing dispatch/{mlx,jax}/attention.py pattern. A whole-capacity where update cannot be advertised as supported.

PR 14 — Gemma 3n text port

• F-50 Gemma 3n architecture: AltUp, LAuReL, per-layer embeddings, sparse/dense MLP; consumes the F-16b softcap
• F-66 End-to-end native integration on the public API

Replaces the prototype's per-prefix recompilation, per-layer streaming, host-side vocabulary chunking, and always-on oracle. Tier A may materialize affine-4 to dense once at load, metered.

Tier B — performance architecture

PR 15 — GGUF loading

• F-39 GGUF container reader: header, metadata, tensor directory validation, checked block/tail arithmetic
• F-40 Packed handles implementing the PR 7 handle contract, mmap-backed until consumption
• F-41 Packed-block → dense array conversion for the correctness baseline
• F-05 GGUF tokenizer-metadata adapter

Values still bind through load_state. Untrusted binary input: bounded preflight before any third-party parser allocation.

PR 16 — CPU quantized kernels

• F-42 Packed matmul/embedding ops plus capability-gated rewrite selection, no whole-tensor dequantization

Registers into the existing rewriting/ DB and dispatch/. Oracle is an independently produced ggml reference, not the production decoder.

PR 17 — Device quantized kernels

• F-43 MLX/JAX packed kernels
• F-44 MLX affine-4 (Gemma) packed layout

One-time repack is allowed if cached and metered; no zero-copy claim across host→device.

PR 18 — Benchmark and telemetry harness

• F-61 Raw-JSON benchmarks: prompt processing, decode, TTFT, ITL, throughput, peak RSS/device memory, compile count, copied bytes

llama-bench excludes tokenization/sampling, so comparisons match that scope or report both.

PR 19 — Paged/ring KV and prefix reuse

• F-48 Ring/paged storage, block tables, allocator as a backend of the PR 8 interface
• F-46 Per-backend lowering of ring/paged storage
• F-49 Prefix reuse with copy-on-write and a full canonical key

Logical results must be storage-layout independent. Prefix reuse is optional and never a scheduler prerequisite.

PR 20 — Continuous batching

• F-51 Scheduler: dynamic admission, mixed prefill/decode, backpressure, fairness, cancellation and reclamation

Two-phase cancellation: detach logically, quarantine in-flight blocks until the backend fence completes.

Tier C — optional interfaces (never blocks A or B)

PR 21 — Constrained decoding

• F-52 Grammar/JSON constraints with per-sequence state and bounded compile/step work

Optional. Does not require batching.

PR 22 — Speculative decoding

• F-53 Draft/target verification with exact cache and RNG rollback via F-70

Optional. Greedy output must equal ordinary decode.

PR 23 — CLI

• F-54 Thin CLI over the public API with streaming and signal cancellation

Optional. No generation logic of its own.

PR 24 — Server and metrics

• F-55 Async OpenAI-compatible completions/chat with SSE, bounded admission, loopback default
• F-56 Prometheus metrics derived from the F-38 counters

Optional. Engine never imports the server; importing pytensor_ml must not import a web stack.

PR 25 — Pooling and reranking

• F-57 Pooling modes, embedding and rerank results

Optional. No generation cache required.

PR 26 — LoRA

• F-58 Adapter identity, validation, immutable-base application
• F-59 Prefix-cache and scheduler integration

Optional. Adapter identity enters plan and prefix keys.

Governance

PR 27 — Release gates and upstream follow-ups

• F-62 Machine-readable required-cell manifests per tier; release jobs fail on skip, missing artifact, or stale evidence
• F-64 Evidence-led PyTensor upstream requests, starting with a supported backend-dispatch registration hook to replace the sys.meta_path workaround

Tier C never blocks Tier A/B.

Acceptance rules that apply to every PR

• No prefix replay: prefill/decode compile once per plan key or documented shape bucket, never per token.
• Decode consumes token(s), positions, sequence ids and persistent KV — it never re-evaluates the full prefix.
• A one-token append does not write over the whole cache capacity.
• No file read, per-layer dequantization or transpose inside the token loop.
• Cached output matches full-prefix output and an independent oracle on all logits at every tested step — not just plausible text.
• The oracle is test-only: it never appears in the normal import or call path.
• Batch > 1, independent sessions, seeded RNG, stop state, streaming and cancellation are runtime behaviour, not report wrappers.
• Unsupported model/backend/quant/cache combinations fail at load or compile with the missing capability named, not at token 1.

Non-goals

Replacing llama.cpp or matching its architecture/backend matrix; multimodal encoders; distributed or tensor-parallel execution; training APIs; auth/TLS; a web UI; exact CLI flag parity; executing model repo code (trust_remote_code=False by default); and any silent fallback to full-prefix replay, whole-model dequantization or another framework.

---

_{Detailed per-item design, dependency DAG and adversarial-review notes are kept in .dev/planning/ alongside the branch work; this issue is the public follow-up surface to reference from PRs.}

[jessegrabowski]

We need to pull out the sub-issues from this that are in scope for this repo, because I don't think the entire llm hosting/optimization chain is. I envision this repo as something like jax.nn/flax/torch/tinygrad, that has all the building blocks for defining, training, and analyzing deep learning-based modeling. This plan includes a number of highly specialized items that go beyond that scope.

[cetagostini]

Agreed — and the split falls out cleanly if we use your analogy as the actual test rather than as a vibe. Two questions per item:

1. Does it ship inside the installed core package of ≥2 of jax.nn / flax / torch.nn / tinygrad? Same-repo-but-not-installed (examples/, extra/, benchmarks/) counts as out — that is precisely how those projects express "not a building block".
2. Is it a node in the graph — an op, layer or rewrite PyTensor can compile, differentiate, serialize and inspect — or is it Python orchestration around the graph (loops, sockets, request state, schedulers, artifact registries)?

Plus three tie-breakers: does it serve more than one model family, does an external project already own it (tokenizers, jinja, gguf-py, huggingface_hub, xgrammar, vLLM), and does it cost a new hard dep / event loop / network listener / obligation to track third-party model releases. Current runtime dep budget is pytensor, numpy, safetensors, which is the constraint the line has to respect.

Ran that against all 74 F- rows. 19 survive, and every one is a layer, an op, a rewrite, a checkpoint primitive or a serialization contract. They group into 10 sub-issues:

In scope — the 19

• RMSNorm — F-13. Belongs on Implement Normalization Layers #11 next to the existing LayerNorm box. Core in three of the four: torch.nn.RMSNorm (2.4.0), flax.linen/nnx.RMSNorm, tinygrad.nn.RMSNorm.
• Rotary position embeddings — F-14. RotaryEmbeddingLayer op + layer taking an explicit position_ids input (not an implicit 0..T-1), base/scaling variants as props. Precedent flax.nnx.RoPE; torchtune/keras-hub carry it one layer out. Used well beyond text now (ViT, ASR, time series).
• Gated (SwiGLU) feed-forward + pluggable block internals — F-15, F-18. A gate= option on the existing FeedForward (Swish() is already SiLU), and TransformerBlock taking its norm and MLP as arguments — so a "Llama block" is a configuration, not a class in the package. Note torch.nn.GLU is the sigmoid-gated unit, not this.
• Sliding-window / local attention masking — F-17, F-68. Exact precedent: jax.nn.dot_product_attention(..., local_window_size=...); torch's flex_attention does it as a mask_mod. Because it changes AttentionLayer.__props__ (currently ("is_causal", "scale")), the format contract ships in the same PR: GRAPH_FORMAT_VERSION bump, prior-version rejection test, qualname-stability test, and a guard that a fused Composite never reaches the codec.
• Incremental (KV-cached) attention as non-trainable layer state — F-19, F-20, F-21, F-23, F-45. This one is genuinely a building block, not runtime: flax.linen.MultiHeadDotProductAttention(decode=True) keeps cached_key/cached_value in a cache collection, and flax.nnx.MultiHeadAttention exposes init_cache() + nnx.Cache. We already have the entire mechanism — StatefulOp.update_map, params.non_trainable, collect_non_trainable_updates, function(updates=...) — and BatchNormLayer already uses it for running stats. _sdpa_graph's bottom-right causal alignment already lets a short query block attend to a long prefix, so the masking side is done. Scope is: the cache op + buffers + shared per-step inputs + keeping cache state out of _weight_variables/save_pretrained + MLX/JAX lowering. The caller writes the loop — same as flax.
• Weight tying between embedding and output projection — F-16a. Needs framework support only because one shared parameter has to round-trip through save_pretrained/load_state as one tensor, not two equal copies.
• Sharded + lazily-read safetensors checkpoints — F-09a, F-09b, F-24, F-08b. load_state currently reads a whole archive via safetensors.numpy.load_file and pretrained assumes a single WEIGHTS_FILENAME. Parse a shard index and validate coverage before any read, materialize per tensor via safe_open, and report unmapped/unexpected keys instead of dying on the first mismatch (torch.load_state_dict returns missing_keys/unexpected_keys; tinygrad.nn.state.safe_load is mmap-backed). This is a large-tensor problem, not an LLM problem.
• Named error when a marker op has no backend lowering — F-28. Today that fails somewhere inside the backend; it should fail at compile time naming op + backend. Trimmed hard: no model/dtype capability matrix.
• Replace the sys.meta_path dispatch workaround with a supported PyTensor hook — F-64. dispatch/__init__.py installs a MetaPathFinder to register funcify handlers after pytensor.link.{jax,mlx}.dispatch loads, because PyTensor has no plugin hook. Worth upstreaming regardless of anything LLM-shaped; the constraint to keep is that importing pytensor_ml must not import jax/mlx.
• LoRA adapter layer — F-58. flax.nnx.LoRA/LoRALinear is core in flax, and PEFT is training, which your scope names explicitly. Strictly the layer + parameter identity: frozen base, trainable A/B, alpha scaling, base weights bitwise unchanged after a step.

Five more rows (F-29, F-60, F-67, F-69, F-71) aren't separate issues — they're acceptance criteria that ride the ones above: each new op ships its own predict_db rewrite in the same PR, each lands with numerical parity against an independent reference (mutation tests that fail on flipped RoPE pairing, mask polarity, cache frontier), the cache interface stays dense-only, and there's a regression test that declared cache updates survive rewrite_for_prediction.

Out of scope — the other 48

Tokenizers and streaming detokenizers, chat templates, HF/GGUF tokenizer adapters, HF config normalization and the model_type → builder dispatch table, artifact manifests with hashes/provenance, remote snapshot downloads, GenerationConfig and greedy/sampled selection, logits processors, stop criteria, session RNG ownership, the autoregressive loop, streaming iterators, batched generation, chat sessions, structured results, prefill/decode plan builders and shape-bucketed plan caches, host/device byte accounting, GGUF container reading and packed handles, CPU/device packed kernels and MLX affine-4, paged/ring KV storage and its lowering, prefix reuse, the continuous-batching scheduler, constrained decoding, speculative decoding, pooling/reranking result APIs, LoRA × prefix-cache × scheduler integration, the Gemma 3n architecture (AltUp, LAuReL, per-layer embeddings) and its e2e integration, the benchmark/telemetry harness, release-gate manifests, and the load_llm/Session/Engine API contract itself.

In every one of those cases the analogy set puts it in a sibling project (huggingface/transformers, huggingface/tokenizers, pytorch/ao, torchtune, Orbax, google-deepmind/gemma, MaxText, vLLM/SGLang) or in a non-installed examples//extra/ tree — and most fail the graph test outright. Two are worth calling out specifically: quantization has a decisive precedent, since torch's own kernels left core for pytorch/ao and torch.ao.quantization is slated for deletion; and paged KV is explicitly experimental in both jax (pallas.ops.tpu.ragged_paged_attention) and torch (nn.attention.experimental._paged_attention, documented "no BC/FC guarantees").

One honest counter-example, since it's the strongest argument against my own conclusion: tinygrad v0.13.0 does ship tinygrad.llm inside the installed package — KV cache, GGUF plus every dequant kernel, a from-scratch BPE tokenizer, chat templating, a generation loop, and an OpenAI-compatible HTTP server. But it's four files, originating as "target is 200 line ollama", it exists to exercise the compiler, tinygrad's own sz.py lists tinygrad/llm under NONCORE_DIRS, and even it has no scheduler, no paged KV, no LoRA and no speculative decoding. So the most permissive reading of the analogy still lands inside Tier A/B and excludes all of Tier C.

Proposal: this issue becomes the tracking link, and the runtime body moves to a separate repo that depends on pytensor_ml.

Parked — 2 to decide later, not blocking the 19

• from_pretrained(source_format="huggingface") (F-11, with F-06/F-07/F-08a contingent on it). pretrained.py already advertises Format = Literal["auto", "pytensor", "huggingface"] and raises NotImplementedError, so "out of scope" has a cost either way: either drop the promise (keep HF detection only to raise a clear "this is a HuggingFace directory, not a pytensor_ml one"), or implement it for exactly one architecture with the builder table capped at one entry. My read is that a model_type → builder table is a model zoo and none of jax.nn/flax/torch.nn ships one, but it's your call and it doesn't need making now. Either way the sharded/lazy checkpoint work stays in scope — that's about large tensors, not about HuggingFace.
• One decoder architecture as an examples/ entry (F-47), feeding Create an example gallery #4. Precedent is unanimous — flax examples/gemma, ExecuTorch examples/models/llama/, tinygrad extra/models/llama.py: an example proves the primitives compose without putting a zoo in the package. Only worth deciding once RMSNorm/RoPE/gated-FFN/KV-cache have landed.

Full per-row triage for all 74 rows — gates, the where-does-it-actually-live evidence matrix across the four references, verdict and reason per row, plus what evidence would flip each verdict — is in .dev/planning/notes/scope-triage-issue-53.md alongside the branch work. Happy to open the 10 issues above (and the two parked ones as needs-decision) once you're happy with the cut.