Meta: native LLM inference runtime in pytensor_ml (t... | pymc-devs/pytensor-ml

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

Meta: native LLM inference runtime in pytensor_ml (tracking)

Description