Skip to content

Tlaloc capability matrix

What runs, where it ran, and what pins it. The README carries a one-row-per-area summary; this is the full picture.

Status means exactly this:

βœ… Certified an automated test pins it, and the row says where it ran
πŸ§ͺ Written the code exists and unit-tests pass, but the end-to-end path has never run β€” the reason is always stated
πŸ“ Designed a design document exists; no implementation
⬜ Not started planned, nothing written yet
❌ Not planned

The suite has 3,346 automated tests: 3,224 that ./gradlew test runs, and 122 from the vendored Maestro modules, which the root test task does not run (50 in maestro-tlaloc, 4 Tlaloc tests in maestro-common, 68 in maestro-server). All 3,224 were counted with every test re-executed (./gradlew test --rerun-tasks --continue) on 2026-10-02 with 0 failures; the 122 were counted earlier in each Maestro module's own test --rerun, 0 failures, and not re-run for this count. (Two GPU and CPU timing assertions fail now and then under machine load: that dispatch floor, whose timing the GPU tests running in parallel disturb, and LlamaDecoderPytorchBenchTest's "backward slower than forward", which skips itself when it sees concurrent load.) On the GB10 workstation where they were counted, 102 of them skipped by name: 88 MLIR round trips that need stablehlo-translate or sdy-opt, 12 TPU tests (5 smoke, 7 Mosaic kernels), and two timing runs that need TLALOC_PAGED_BENCH=1 or TLALOC_VMAP_BENCH=1. This is the one place the documentation states the count.

Automatic differentiation

Capability Status Notes
Reverse mode β€” grad, grad2, grad3, valueAndGrad* βœ… Compile-time, via the K2 plugin
Forward mode β€” jvp, jvp2, valueAndJvp* βœ…
Batching β€” vmap, vmap2, batchAxis (experimental) βœ… GB10 Compile-time, via the K2 plugin; the batch axis is Named<N, Sym> / Named<N, Bounded<B>> in the types. Composes with grad and jvp both ways and with itself. Rules for elementwise, broadcasting, reduction, shape and softmax ops, matmul (shared weights not copied), canonical named contract, if, constant-trip loops, and all linear algebra (the while-lowered solve, det, QR and eigh run one loop over all matrices). Checked against stacking the examples (interpreter F32/F64, GB10, plugin; batch 1, 7, 64 and a bounded axis), per-example gradients against loops of single-example gradients, JAX 0.10.0 vmap at F64 to 1e-12. Refused by name: conv, pooling, gathers/scatters at a run-time index, RNG, sparse, attention/serving ops, collectives. Design
vjp / jacobian / jacobianReverse / hessian (+ *2 forms) βœ… Arbitrary-arity models go through graph capture, not the fixed-arity intrinsics
Higher order β€” fwd-over-rev, rev-over-rev, nesting matrix βœ… Full nesting matrix certified
Custom derivatives β€” customVjp, customJvp, customVjpJvp βœ…
Control flow β€” loops and branches under grad βœ… Ο†-calculus coarsening β€” Shen et al., Coarsening Optimization for Differentiable Programming, OOPSLA 2021 (local copy) β€” with a Symja-backed closed-form engine
Readable reverse code βœ… dumpGradSource compiler flag + DxirFunction.toKotlinSource(); printed source compiles and runs. Scalar lambdas (Float and Double) render through the plugin; tensor grad {} lambdas render through the capture route (CapturedStep.gradSource()). ABS, RSQRT, GELU and SILU still have no :core tensor twin and refuse by name
One AD engine βœ… The runtime tape was deleted; every gradient β€” intrinsics, capture API, :nn training β€” comes from the same reverse transform (audit)

Tensors and operations

Capability Status Notes
Shape- and dtype-typed tensors βœ… DTensor<Rank2<Sym, Sym>, F32>; mismatches are compile errors
Named axes βœ… Named<N, A>; a contract over misaligned axis names does not compile
Bounded dimensions (experimental) β€” Bounded<B> with object B : DimBound(max) βœ… A constant size outside 1..max, two bounds on axes an elementwise operator aligns, and a bound below 1 are compile errors at the call's file, line and column (BoundedDimCompileCheckTest); a bound from another compiled module is checked at run time only. boundedProgram runs at any size in the interpreter, exact or padded to a bucket with valid-mask inputs, and equals a static-shape program at every size (BoundedProgramTest); valueAndGrad gives a loss and its gradients per bucket, equal to the analytic ones at every size. Design
Op surface β€” elementwise, broadcasting, reductions, shape ops, concat/slice/pad, where βœ… Full DiffKT parity, closed
NN ops β€” conv2d (incl. grouped/depthwise), pooling, softmax, embedding, losses, batch norm βœ…
Special functions β€” lgamma, digamma, polygamma, integral βœ…
Linear algebra β€” cholesky, triangularSolve, tril/triu, solveSpd, logDetSpd, invSpd βœ… GB10 Rank-2, F32 and F64. Differentiable to any order under grad {}, jvp {}, hessian {} and the capture API (hessian of logDetSpd certified through the plugin), at F32 and F64. Cholesky's rules are Murray's (2016), the solves' implicit differentiation; the SPD functions are lowered to cholesky and triangular solves. Lowered to stablehlo.cholesky and stablehlo.triangular_solve. Certified against Double finite differences (interpreter, plugin, capture API), JAX 0.10.0 values and gradients (LinalgJaxParityTest, harness/python/linalg_jax_goldens.py), and on the GB10: F32 graphs against the interpreter, F64 gradient graphs against F64 finite differences to 1e-7 (PjrtLinalgTest) and against DxirInterpreterF64 to 3e-16 (PjrtF64OpsTest). At F64 through the plugin: every case of LinalgGradientTest (1e-8 against finite differences; det's Hessian against the exact one), JAX float64 to 1e-11 (LinalgJaxParityTest, and F64HessianJaxTest against live JAX)
Linear algebra β€” solve and det (LU with partial pivoting), qrQ/qrR (Householder), eighValues/eighVectors (Jacobi) βœ… GB10 Same certification as the row above. StableHLO has no LU, QR or eigensolver, so each is emitted as a stablehlo.while loop (no custom call): correct on the GB10 (f64 within 1e-12 of the host kernels) but sequential, n to 20Β·n(nβˆ’1)/2 loop steps, far slower than vendor routines for large matrices. det's gradient is NaN at a singular matrix; eighVectors' is infinite at a repeated eigenvalue (as in JAX); QR needs rows β‰₯ columns and, for derivatives, full column rank; eigh runs a fixed 20 Jacobi sweeps. Eigenvectors are signed so each column's largest entry is positive (JAX's signs are arbitrary)
Stateless RNG β€” threefry-2x32, uniform/normal/cauchy/exponential/chiSquare βœ… Bit-exact against JAX's classic stream; reparameterized gradients
Sparse β€” rank-2 CSR, sparseΓ—dense matmul through grad { } βœ… Host + interpreter; no GPU emission by design (audit)
dtypes β€” F32, F64, I32, BF16 (and I8 for quantized serving weights) βœ… bf16 end to end incl. native PJRT BF16 buffers and mixed-precision training
F64 under grad {} β€” every intrinsic (grad, grad2, grad3, valueAndGrad*, jvp, jvp2, vjp, vjp2, jacobian, jacobianReverse, hessian and their two-argument forms, customVjp, customJvp) over DTensor<…, F64> and Double βœ… GB10 Host path: F64 versions of every op except embedding, sparse products and RNG draws, which refuse by name (DTYPE_UNSUPPORTED); F32 and F64 never mix implicitly (DTYPE_MISMATCH at the call). Certified against fourth-order central differences to 1e-8–1e-9 (F64TensorGradientTest, F64OpCoverageTest: 24 op families), analytic derivatives to 1e-13, JAX float64, and a test whose inputs, literals and captured Doubles F32 rounding would move. 65 plugin test classes and LinalgJaxParityTest run every case at F32 and F64. PJRT: F64 graphs through PjrtSession.runOnHost agree with DxirInterpreterF64 within 3e-16 of the largest entry (lgamma/digamma 1.1e-13). F64 dots are about 9Γ— slower than F32 (TF32) on the GB10 (a 2048Β² gradient: 183 against 20 ms). Not run on a TPU. The Tracer-capture API, :nn training, bounded programs and IREE stay F32 and refuse F64 by name
dtypes β€” F16, FP8, int8 tensors ❌ int8 exists for KV-cache quantization only

Models and training (:nn)

Capability Status Notes
Layers β€” Dense, Conv2d, MaxPool/AvgPool, BatchNorm, Dropout, Embedding, EmbeddingBag, GRU, Flatten βœ… Immutable/functional; a training step returns a new model
Transformer layers β€” LayerNorm, RMSNorm, RoPE, MultiHeadAttention (causal, grouped-query, q/k norm), SwiGLU, Mlp, TransformerBlock, CausalLM βœ… Two small Llama/Qwen3-style models against plain PyTorch on shared weights: identical loss, every gradient within 2.4e-6 of its largest element, 20 AdamW steps within 9e-7 (TransformerVsPytorchTest)
MixtureOfExperts β€” top-k routed SwiGLU experts, trainable, the experts as RAGGED_DOT (CausalLmConfig.numExperts) βœ… Value and every parameter gradient against central differences of the equations (MixtureOfExpertsTest); a CausalLM with experts on the GB10 within 1.8e-3 of the interpreter, trained by AdamW (PjrtCausalLmTrainingTest). Not checked against PyTorch; no load-balancing loss
Hugging Face checkpoints into and out of CausalLM (Llama, Qwen3) βœ… Qwen3-0.6B predicts transformers' 16 greedy tokens, top-20 logits within 3e-5; save and reload is bit-exact (HfCausalLmTest)
Hugging Face tokenizers β€” BPE encode, decode, streaming decode, chat templates (:tokenizer) βœ… HfTokenizer.load(checkpoint). Byte-level BPE (Qwen3, Muse Glimmer, GPT-2) and SentencePiece-style BPE with byte fallback (TinyLlama, Gemma 4). Equal to transformers 5.17 AutoTokenizer on every golden (GoldenTest: about 900 checks per family, including 80 random strings, special tokens as text, partial-id decodes and tokenizers' DecodeStream); follows transformers where it rebuilds the pipeline from the tokenizer class (TinyLlama's Metaspace). Chat templates for Qwen3, TinyLlama (Zephyr) and Muse Glimmer are Kotlin renderers equal to apply_chat_template; tools are not rendered. Unigram/WordPiece models, lstrip/rstrip added tokens, truncation and padding refuse by name
Fine-tuning a pretrained LLM on the GPU βœ… GB10 examples/fine-tune: Qwen3-0.6B, AdamW, one compiled step; PjrtCausalLmTrainingTest checks GPU gradients against the interpreter
ODE integration β€” fixed-step RK4 (rk4, rk4Trajectory) βœ… On capture-API tracers (differentiable with respect to the initial state and the parameters f reads, reverse and forward mode) and host tensors. Checked against RK4's exact amplification factor, finite differences of a damped oscillator, and a parameter fit by gradient descent (OdeTest). Not callable inside grad {}; there the loop is written in the lambda
Frozen parameters β€” Frozen, capture(..., frozen), optimizer.step(..., frozen) βœ… The trained parameters' gradients equal a full capture's bit for bit, the gradient function has no adjoint work for the frozen ones (fewer MATMULs), frozen tensors are bit-identical after 30 AdamW steps with weight decay, and no optimizer state is kept for them (FrozenParametersTest). A key or prefix that selects no parameter refuses by name
LoRA β€” Lora.apply, LoraConfig (rank, alpha, dropout, rsLoRA, target modules), Lora.merge βœ… Targets by Hugging Face module name (q_proj) or key path. At initialization the adapted model's logits equal the base model's bit for bit: a tiny Qwen3-shaped model, Qwen3-0.6B (112 adapted layers) and TinyLlama-1.1B (154). Adapter gradients against f64 central differences of an f64 implementation within 3.3e-7 relative (LoraTest); trains with the base frozen in f32 and in MIXED_BF16. Merged against unmerged logits: 7.4e-5 on Qwen3-0.6B with adapters on all 196 linear layers (LoraMergeTest). Dropout masks are keyed constants: a new mask needs a new capture
LoRA adapters in PEFT's format β€” HfLoraAdapter βœ… adapter_config.json + adapter_model.safetensors with PEFT's names and layouts. Against peft 0.21.0 / transformers 5.17 (HfLoraAdapterTest, skips by name without a Python that has peft): a Tlaloc adapter read by PeftModel.from_pretrained within 2.3e-6 on a tiny random Qwen3 and 7.9e-5 on Qwen3-0.6B; a PEFT-written adapter read by Tlaloc within 2.4e-6; the merged checkpoint read by transformers without PEFT within 3.6e-6. DoRA, rank_pattern, alpha_pattern, trained biases and modules_to_save refuse by name
LoRA fine-tuning on the GPU βœ… GB10 examples/lora-finetune: Qwen3-0.6B, rank 16 on every linear layer (1.67 % of the parameters), the frozen weights staged on the device once; 320-369 ms a step, loss 4.35 to below 0.02 in 6-7 steps; the merged model gives the adapted model's answers
Optimizers β€” SGD, Momentum, RMSprop, Adam, AdamW, FixedLearningRate βœ… Pure (params, grads, state) β†’ (params', state')
LR schedules β€” step, exponential, cosine, linear warmup βœ… Pure functions of the step count; agree with PyTorch's StepLR/ExponentialLR/CosineAnnealingLR to 2.8e-7 relative. CosineDecay deliberately CLAMPS past T_max where PyTorch's is periodic β€” certified as a difference
Gradient clipping β€” by global norm, by value βœ… Agrees with clip_grad_value_ exactly; by-norm differs from clip_grad_norm_ by torch's own +1e-6 denominator guard (~1.6e-7 relative) and sits closer to the exact ratio
Model persistence β€” save/load a model + optimizer state βœ… One safetensors file; round trip is BIT-IDENTICAL and a resumed run's next 15 steps match the uninterrupted ones bit for bit. Loading builds a NEW model β€” layers stay immutable
Training loop β€” capture once, train βœ… MLP converges; loss curve matches PyTorch to 7 decimals
Training on GPU βœ… The captured gradient graph compiles to StableHLO and trains on CUDA
Mixed precision (bf16 compute, f32 master weights) βœ… No loss scaling needed
Cross-entropy against I32 ids; the LM head's cross-entropy in chunks (captureCausalLm, lmHeadCrossEntropy) with rematerialization (withRemat) βœ… GB10 Ids give the one-hot loss and gradients, ignored positions included; the chunked loss equals the full one, and rematerialized gradients are bit-identical (LmHeadCrossEntropyTest). 4,096 tokens, a 32,768-token vocabulary: peak 1,684 MiB full, 385 MiB in chunks of 512 with remat on the GB10 (PjrtLmHeadLossMemoryTest); compiled for TPU, 1,162 and 312 MiB on a v6e
Distributed / multi-GPU training πŸ“ Design + marshalling done; never run (needs 2+ hosts)

Inference and serving

Capability Status Notes
Paged attention, KV-cache writes, decode bucketing βœ… Inference-only ops; they refuse differentiation by name. Paged attention takes an optional sliding_window (a row sees its last W positions, its own included, as transformers' mask does), checked against a dense walk for every window and length in the interpreter and against the interpreter on the GB10. Its two f32 dots are emitted with HIGHEST precision, so XLA does not run them in TF32
Bounded programs served without a compile per size βœ… GB10 BoundedProgramExport writes one body per bucket and tlaloc-bounded.json (tlaloc-bounded-v1), refused when padding changes the result in the interpreter. tlaloc_bounded.py on PJRT-CUDA: a masked mean, a masked softmax, a two-bound embedding and a training step at every size up to the bound against the interpreter (BoundedArtifactRunTest). Triton's bounded mode: two example models at every length 1..16 over HTTP and gRPC (verify.sh). 200 requests of mixed length through attention: 6 compiles against 163 at exact lengths, execution 13 % slower from padding. Bounded dynamic StableHLO is not used: the PJRT C API cannot create its buffers (design)
A training loop over batches of varying length on the GPU βœ… GB10 BoundedProgram.valueAndGrad run through a PjrtSession (BoundedExecutor): 40 SGD steps over lengths 1..64 compile 4 executables, one per bucket, where exact lengths compile 27 (PjrtBoundedTrainingTest)
In-process inference from the JVM β€” ServingModel (:runtime-pjrt) βœ… GB10 Loads a serving artifact into the calling JVM and runs it through PJRT: prefill and decode, greedy, batches, the KV cache kept on the device. Plain-Java signatures, compiled by javac -Xlint:all -Werror in the suite (ServingModelJavaApiTest). The same ids as run_llama_generate.py on the same artifact, logits within 5.2e-6 of the interpreter's (ServingModelTest). The fine-tuned Qwen3-0.6B from plain Java: about 15 ms a token after an 8.5 s first request (examples/java-inference); inside a Spark 4.2 mapPartitions on JDK 25, local mode only (examples/spark-inference). Windowed and quantized KV pools refuse by name
HuggingFace safetensors ingestion βœ… Kotlin parser; certified against torch reading the same bytes
safetensors WRITING βœ… F32/F64/I32/BF16, header padded and tensors ordered so every offset is naturally aligned; certified both ways against the reference safetensors library on raw bytes. Sharded and streamed output are not supported; F16/FP8 refuse by name
A real Llama serving end to end βœ… TinyLlama-1.1B, all 22 layers, on PJRT-CUDA β€” 6/6 generated token ids identical to HuggingFace transformers, driven directly and through vLLM
A real Qwen3 serving end to end βœ… Qwen3-0.6B (Apache-2.0), all 28 layers: the 16 greedy token ids HuggingFace transformers produces (float32, CPU, eager attention) for a plain prompt and a chat-template prompt, in the reference interpreter (logits within 2.6e-6 of the largest) and served by Triton with a prefill call and the sequence batcher (within 3.3e-4; TF32). On the GB10 through Triton: about 25 ms to prefill and 19 ms per generated token (52 tokens/s). The interpreter test checks 4 tokens per prompt by default and all 16 with TLALOC_QWEN3_FULL=1; both ran. The weights stay in the HuggingFace cache and the tests skip by name without them
HuggingFace model families βœ… HfModelFamily reads Llama, Qwen3 and Muse Glimmer checkpoints through one config parser, name mapping and graph builder, with one DecoderLayerSpec per layer. A config.json key the family does not know is refused by name at parse (inside and outside a multimodal file's text_config); Qwen3's per-head q/k RMSNorm is certified against an independent Kotlin forward pass and, with its negative control, against the real checkpoint
A real Muse Glimmer serving end to end (text only) βœ… meta-models/Muse-Glimmer-30B (Apache-2.0, revision a4e59da5): the text decoder, 52 layers, 27.8 billion parameters, its 56 GB of weights bf16 on the device; the 809 vision-encoder tensors are not read and the image and video placeholder ids are refused by name. Served by Triton on the GB10 with a prefill call and the sequence batcher, for a plain and a chat-template prompt, 16 greedy tokens each: all 32 ids equal transformers 5.17 run with the same arithmetic (bf16 weights, f32 activations, float64 RoPE tables; CPU, eager), with logits within twice that run's own float32-vs-float64 difference (at most 1.4e-2 of the largest logit, against 1.8e-2); against transformers in bfloat16, 31 of 32 ids are equal and the 32nd is the one position where the two transformers runs choose differently from each other (margins 0.31 and 0.12). --perturb fails. Served from the long-context export (below): about 540 ms to prefill a short prompt (the 512-token entry), 265 ms per generated token (3.8 tokens/s, close to reading 56 GB of weights at the GB10's memory bandwidth; 241 ms with the batcher's 1 ms queue delay, which the export raised to 20 ms so that concurrent sequences batch; the backend now batches them itself with no delay, row "Decode steps batched by the backend"), load about 3.5 minutes (16 compiles in 2 minutes, 82 s of weight upload), a PJRT memory fraction of 0.55 computed from the weights, 87 to 96 GiB peak in use system-wide (unified memory; other processes included). triton/verify.sh runs it with MUSE_GLIMMER=1. The weights stay in the HuggingFace cache
Muse Glimmer arithmetic against transformers' own code βœ… HfMuseGlimmerTest, a five-layer model (hidden 32, window 3, soft-cap 3) whose weights are a closed-form function of the tensor name, run through transformers' MuseGlimmerForConditionalGeneration and through the decode graph in the interpreter, as a decode loop and as prefill then decode: within 7.1e-8 of the largest logit when both build RoPE tables in float64, and the same token at every position against unmodified transformers (whose float32 tables differ by an ulp here and there, which moves the logits by up to 8e-3 on this 32-wide model once bf16 rounding of the projection inputs amplifies it). Negative controls: without the window, with RoPE on every layer, without soft-capping, with w instead of 1 + w norms and without the q scale, the logits move by more than 1e-2
Sliding-window layers, layers without RoPE, attention and MLP output norms, final-logit soft-capping, attention output gate, gainless q/k norms βœ… Implemented in the one graph builder and certified through Muse Glimmer (above). attn_logit_softcapping is still refused by name. A Qwen3 config with sliding layers now builds; no real sliding-window Qwen3 checkpoint has been run
KV pages sized to the window for sliding layers βœ… A model's sliding-window layers keep their KV in a second pool class, the windowed KV pool (WindowedKvPool; a tlaloc-serving-v3 artifact): each sequence holds a ring of at most ceil(window / blockSize) + 1 pages, block b on ring page b % ringPages, with its own block table and slot mapping, and a request is split into calls the ring can hold. The full-attention layers keep full-history pages. Certified in the reference interpreter (WindowedKvPoolTest: two sequences several windows long through prefill chunks, batched and padded decode steps give logits bit-identical to full-history pools; a ring one page short, or a call one token longer than the ring holds, changes them) and through Triton (verify.sh: a three-layer model with window 8 and pages of 4 holds at most 3 windowed pages while its full pages grow to 15, and matches the full-history model sent the same calls to 5.3e-7 of the largest logit). For Muse Glimmer the computed KV per sequence is 365 MiB instead of 832 MiB at 8,192 positions and 989 MiB instead of 3,328 MiB at 32,768 (docs/SERVING_ARCHITECTURE.md); Muse Glimmer served from a v3 artifact (its 39 sliding layers in the windowed pool, rings of 160 pages) greedy-decodes the same 32 ids as before, and its pools for four sequences of 32,768 positions take 4,109 MiB on the device against 13,312 MiB with full-history pages (verify.sh with MUSE_GLIMMER=1). The framework-free Python runtime and the vLLM plugin refuse a v3 artifact by its schema version
Long contexts: prefill in chunks, Muse Glimmer up to 32,768 positions and four sequences βœ… A prefill entry can take fewer tokens than its context (-PprefillChunk); a longer prompt runs as several calls, and the windowed ring is sized for a chunk past the window. The prefill graph gathers each sequence's pages once instead of once per token (PAGED_ATTENTION accepts one block table per group of rows): on the GB10 the new lowering gives the per-row lowering's output exactly. In the reference interpreter chunked prefill gives whole-context logits bit for bit, with full-history and windowed pools, and chunks through a ring one page short change them (BatchedPrefillTest); through Triton a windowed model with 6-token chunks matches the full-history model sent the same calls (verify.sh). Muse Glimmer exported for contexts 512 to 32,768, decode batches 1, 2 and 4 and 512-token prefill calls, served by Triton on the GB10 (verify.sh with MUSE_GLIMMER=1): a 2,305-token prompt whose fact lies 2,190 tokens before the question, outside the sliding window, gives the 16 ids transformers gives with the same arithmetic (its answer starts with the fact, " 4719."; logits within 7.7e-3 of the largest, 2e-2 required), and the two short prompts keep their 32 ids. Measured with triton/context_bench.py: a decode step of one sequence takes 265 ms at context 512 and 317 ms at 32,768; four sequences together make 16.7 tokens/s at 512 and 9.6 at 32,768; prefill runs at 766 tokens/s for 432 tokens and 142 tokens/s for 32,688 (230 s); the machine peaked at 87 to 90 GiB in use (table in docs/SERVING_ARCHITECTURE.md). Those numbers were taken with the GPU shared and before sliding layers read their ring (next row), which cut the 32,688-token prefill to about a quarter.
Long-context attention that reads only what it needs βœ… GB10 (prefill) A sliding layer reads its ring (PAGED_ATTENTION with ring = true: 160 pages, 2,560 positions for Muse Glimmer) instead of the bucket, in decode and prefill. A prefill call over a context of several key blocks attends one block of 2,048 keys at a time with a running max and sum, from the first block any row can see to the block of the call's last position, so no score tensor as wide as the context is written. Measured through Triton on an idle GPU: a 512-token Muse Glimmer call at positions 31,744 takes 1.39 s instead of 4.50 s, and a 31,744-token prompt 62.6 s instead of 222.8 s (measured again, three idle rounds in turn: 1.38 s against 4.39 s, 62.2 s against 225.2 s). The interpreter gives the same bits through a ring as through a full-width table (PagedAttentionTest, WindowedKvPoolTest); on the GB10 both forms agree with it within 6e-8 (PjrtBlockwisePagedAttentionTest), and the blockwise form agrees with the interpreter and the one-pass per-row form within 2.4e-7 at block boundaries, lengths 0 and 1, an empty call, a window crossing a block boundary and fully masked blocks, and verify.sh keeps every token id. Decode of the full-attention layers still reads the whole bucket (⬜), and the attention dots stay f32 (a fused kernel is not started ⬜)
Where serving time goes (Triton) βœ… GB10 triton/profile.sh times decode steps and prefill calls on an idle GPU (it waits for one, or labels the numbers contended) and, with MODE=nsys XLA_DUMP=1, captures them with Nsight Systems and ties every kernel to its HLO instruction. Decode is at the memory ceiling: Muse Glimmer's step spends 243 ms in kernels against 228 ms for its 53 GB at the 233 GB/s a kernel can stream on the GB10, plus a 20.5 ms queue delay and about 5 ms of launch and copy-back. A 512-token prefill call at context 32,768 spends 88% of its 4.4 s in attention: every layer, the sliding ones included, scores the whole bucket in f32 without tensor cores and writes 2 GiB score and probability tensors. That was before the change in the row "Long-context attention that reads only what it needs", which takes such a call to 0.9 to 1.4 s. Tables in docs/SERVING_ARCHITECTURE.md
Decode steps batched by the backend, no queue delay (Triton) βœ… GB10 Triton hands each request over at once and the backend forms the batches: a batch waits only for the sequences that decoded in the previous one (up to 20 ms), so a lone sequence does not wait and concurrent ones stay together. verify.sh (sequence_checks.py --batching, TinyLlama): one sequence waits 0.35 ms a step before its execution, four clients pausing 0 to 5 ms before each step run 60 steps in 16 executions, and the same model batched by Triton after a 1 ms delay fails both checks. Measured again with the old and new backends in turn on an idle GPU (three rounds): Qwen3-0.6B bf16 10.70 against 11.33 ms a token, f32 15.39 against 16.79, TinyLlama 21.66 against 22.86, Muse Glimmer 244.3 against 265.9 ms with the same server compute; first measured on a GPU the desktop kept 15 to 22% busy (contended), where four Muse Glimmer sequences still ran one execution per step. Inputs read from mapped host memory and inputs uploaded one by one give bit-identical logits across steps of interleaved sequences (with XLA's autotuning fixed; without it, two loads of a model differ by up to 0.017 in a logit). Step inputs are read in place from mapped host memory and the logits copy follows the execution; those changed server time by less than the noise. A smaller output than the full logits is not built (⬜)
bf16 weights in a serving artifact βœ… HfDecoderConfig.weightDType BF16 (Muse Glimmer's default; -PweightDType=bf16 for Llama and Qwen3) stages the weights as the file's bf16 bytes, copied a block at a time (the Linears transposed in bands, so the 2.7 GB embedding table and head never sit in one array); each projection rounds its f32 input to bf16 and multiplies into f32. The streaming writer is checked bit for bit against the in-memory staging with a 64-byte block. The Triton backend drops the page cache of each weight file after uploading it. Qwen3-0.6B with bf16 weights through Triton (verify.sh): 1136 MiB on the device against 2273 MiB, all 32 fixture ids equal HuggingFace's, logits within 3.0e-3 of the largest (6e-3 required; the f32 tolerance of 2e-3 fails, so the comparison tells the two apart), 11.3 ms a decode step against 15.4 ms
Int8 weights in a serving artifact (opt-in) βœ… GB10 -PweightQuant=int8 (HfDecoderConfig.weightQuant, default none) stores each layer's projection weights as int8 with one f32 scale per output channel (the row's largest absolute weight over 127, codes rounded half to even); the embedding table, the norms and the head keep their dtype. The graph widens the codes into the matmul (fused by XLA in decode) and scales the output columns. Chosen from a PJRT micro-benchmark on the GB10 (int8 and float8 both halve a Muse Glimmer projection, 1.18 to 0.63 ms; int4 0.43 ms but at 154 GB/s and needing group scales, not built ⬜; NVFP4 reachable through the block-scaled dot, 0.54 ms, activations rounded to 4 bits). In the interpreter the int8 graph equals the unquantized graph fed the dequantized weights to 4.7e-6, the codes and scales follow the formula and stream exactly, and reversed scales move the logits (HfMuseGlimmerTest). Quality through Triton (triton/quant_checks.py): Muse Glimmer keeps all 48 fixture ids of the bf16 model (needle answer " 4719."), wikitext-103 perplexity 5.610 β†’ 5.602, 99.0% argmax agreement over 480 tokens; Qwen3-0.6B keeps its 32 HuggingFace ids, perplexity +1.1%, 96% agreement; verify.sh requires Qwen3's 32 ids with int8 weights, and requires them to fail when the codes are quantized along the input axis (triton/int8_wrong_axis.py: all 32 differ, perplexity 3.4 million). Speed through Triton on an idle GPU (old and new served in turn, three rounds): a Muse Glimmer decode step 138 ms against 244 (7.2 tokens/s against 4.1; 29,143 MiB of weights against 53,128), a 512-token prefill call about 160 ms slower (the widening is not free in the prefill GEMMs, ⬜), Qwen3-0.6B 8.8 ms against 10.4 (bf16) and 10.2 against 15.7 (f32). examples/triton-llm/run.sh --quant int8
A tied head without a second copy of the embedding table βœ… A tied head is one MATMUL that contracts the hidden axis of the final state and of the table (dot_general, contracting_dims = [1] x [1], no transpose). In the interpreter it gives the logits of a transposed copy bit for bit (f32 and bf16 weights; a changed copy element is seen). Qwen3-0.6B through Triton: 310 weights, 2273 MiB against 2867 MiB, the same 32 ids, logits within 5.6e-4
Placeholder token ids refused when serving βœ… The manifest lists a multimodal checkpoint's image and video placeholder ids (model.refusedTokens); the Triton backend refuses a request holding one by name and leaves the sequence as it was, and so does tlaloc_serve.py. verify.sh: stand-in ids of the window models over HTTP and gRPC (the next logits equal a sequence that never sent the request; an unlisted id is accepted), and Muse Glimmer's 200091 and 200092 with MUSE_GLIMMER=1
A full KV page pool (Triton) βœ… Pages are reclaimed only from sequences idle past the reclaim rule (Triton has ended them), least recently active first and only as many as needed; otherwise the request is refused by name, listing the sequences that hold pages, and a sequence refused mid-generation keeps its KV. verify.sh (TinyLlama): the refused sequence, sent the same request after another ends, continues with its solo ids; the reclaim order and count are read from the server log; a sequence whose requests are refused still counts as active and is not reclaimed, and a refused second START frees the earlier sequence's KV (both checked, and both failed before the fix)
Preempting a live sequence (KV swap or recompute) πŸ“ Not built: a live sequence is never preempted
Framework-free serving runtime βœ… The serving process imports no JAX, no PyTorch, no NumPy β€” just a PJRT plugin .so and a driver (proven by an import blocker that raises on those modules while the path runs)
The serving runtime finding a plugin for itself βœ… tlaloc_pjrt.find_plugin searches the same roots as the JVM's PjrtBinaries, then the interpreter's own site-packages, and raises the full search report when nothing is found. Pinned on synthetic install trees by vllm_tlaloc_test.PluginDiscoveryTest; on the GB10, examples/gpu-inference/serve.py finds the same plugin as the JVM lane with no environment variable set
vLLM platform plugin βœ… vLLM 0.29.0's LLM.generate() runs a real TinyLlama from a Tlaloc artifact β€” the same 6 token ids as the direct driver and as HuggingFace; a new sequence's prompt is written by one prefill call when the artifact has an entry that holds it (the direct driver's use of it is asserted; the runner's is unit-tested). One sequence, greedy, prompt within the compiled context; vllm serve's HTTP layer is not yet run (audit, runbook)
NVIDIA Triton backend (libtriton_tlaloc.so) βœ… Triton 25.11 on the GB10 compiles Tlaloc StableHLO through the PJRT C API and serves it over HTTP and gRPC. triton/verify.sh certifies: a reverse-mode gradient bit-identical to the DXIR interpreter; FP32, FP64, FP16, BF16, INT8, INT32, INT64, UINT8 and BOOL over HTTP and gRPC (FP16, INT8 and UINT8 from a hand-written module; Tlaloc has no FP16 or UINT8, and uses I8 only for quantized weights); shape buckets; dynamic batching of the gradient model (512 concurrent one-row requests, each with exactly its own row, in about a third as many executions; about 5x the unbatched throughput); ragged batches of two widths grouped by the shape of their rows (216 requests in 74 to 76 executions against 119 to 134 when only consecutive requests of one width run together, the same rows); CUDA shared-memory inputs read in place through a PJRT view and outputs copied device to device, with the same bits as the host path (16 MiB each way: about 0.5 ms of server time against about 4.5 ms through the host); decode steps with backend-held KV pools; sequence mode with prefill and batched decode; load-time refusals by name; TinyLlama-1.1B greedy decoding with the same 6 token ids as HuggingFace; Qwen3-0.6B with the same 16 ids as HuggingFace for a plain and a chat-template prompt; and, opt-in (MUSE_GLIMMER=1), Muse Glimmer's text decoder with bf16 weights (row above) (README)
KV pools updated in place (Triton) βœ… Each serving body aliases every KV_POOL_OUT to its KV_POOL_IN (tf.aliasing_output); the Triton backend donates the pools, so XLA writes them where they are. Checked by device address after every run: verify.sh requires 100 of 100 runs in place for TinyLlama and Qwen3, and a control with donate_kv_pools: false must be seen copying in 100 of 100 runs with the same ids. The framework-free Python runtime and the vLLM plugin still move the pools through the host
Prompts of several sequences prefilled in one call (Triton) βœ… Prefill entries for every batch of the ladder (HfServingExport.specs; -PprefillMaxBatch caps them). Each sequence is a right-aligned row with its own positions, block table and slots; padding writes nowhere. In the reference interpreter (BatchedPrefillTest) four prompts prefilled together, three in a batch-4 entry with a padding row, and second chunks at different positions give their solo logits and continuations bit for bit, with full-history and windowed pools; padding that writes its KV, or a left-aligned row, changes them. The Triton backend runs the prompts of one batch that share a context bucket as one call. On the GB10 (verify.sh), 2 and 4 TinyLlama and Qwen3-0.6B prompts sent together ran in one execution and gave their solo argmax and the same 8 greedy ids, logits within 9.0e-4 of the largest (TF32; the bound is 5e-3). Four TinyLlama prompts prefill in about 35 ms against about 110 ms one after the other, and four Qwen3-0.6B prompts in about 61 ms against about 100 ms. Prompts whose request also carries END were not batched together by Triton's sequence batcher
Triton instances on more than one GPU πŸ§ͺ One PJRT client per GPU (visible_devices, ordinal checked), artifacts compiled per GPU named by instance_group. GPU 0 selection and the refusal of a GPU the machine does not have ran on the GB10; a second GPU has not, because the GB10 has one
Kotlin writes a Triton model repository βœ… TritonModelRepository turns a serving artifact into <model>/config.pbtxt + <model>/1/: request tensors from the manifest, staged weights loaded at model load, KV pools as instance state. Unit-tested against golden text; its output is the reference_decode, TinyLlama and Qwen3 models verify.sh serves
SGLang plugin πŸ“ Design recorded; reuses the same artifact
KV-cache quantization (int8) βœ… Derived error bound, not a guess

Compilation and runtimes

Capability Status Notes
StableHLO + Shardy (SDY) emission βœ… We don't write CUDA; we lower to MLIR
PJRT execution from Kotlin (pure FFM β€” no JNI, no Python) βœ… Certified on NVIDIA GB10
PJRT execution from Python (pure ctypes β€” no framework) βœ…
IREE runtime (CPU + CUDA) βœ…
Pattern recognition + coarsening β€” FlashAttention, GQA, RMSNorm, RoPE, SwiGLU, cross-entropy, LayerNorm βœ…
KPTX β€” a PTX DSL, parser, transpiler, and recognizer-driven kernel claiming βœ… Kernels attach to recognized ops automatically
KPTX paged-attention kernel β€” performance πŸ§ͺ On-device (GB10, floors over six sessions): 1.4–1.9Γ— faster than XLA's own lowering at Llama-3-8B-shaped decode points, 1.6–1.8Γ— slower at TinyLlama-shaped toy points, where the measurement's own dispatch floor is 12–76% of it. Not registered by default β€” opt in per shape, at shapes you measured (numbers and ranked fixes)
TPU Mosaic kernels β€” MOSAIC_KERNEL op, tpu_custom_call emission, Pallas-exported and Kotlin-emitted payloads πŸ§ͺ Never run on a TPU (no hardware). On the CPU host: the emitted programs parse with jaxlib's StableHLO and TPU dialects, and their HLO custom-call is identical to the one jax.export writes for the same Pallas kernel; the Kotlin RMSNorm emitter's Mosaic module equals Pallas's after cse and serializes to the checked-in payload byte for byte; each op's reference decomposition matches numpy in the interpreter and on XLA CPU. PjrtTpuMosaicKernelTest (7 tests) runs the payloads on a TPU against numpy and XLA (TPU_MEGAKERNELS.md)
Netflix Maestro orchestration β€” manifest, step type, pod-spec builder βœ… Unit-certified; a live K8s run has not been done

The public surface itself

These rows are about the contract, not a computation.

Capability Status Notes
Compile-time diagnostics carry a source position βœ… DiagnosticSourcePositionTest asserts that LAMBDA_NOT_LOWERABLE and NAMED_INDEX_MISMATCH report at the offending call's own file, line and column, and that no Tlaloc error is emitted without one
An IDE "red squiggle" πŸ§ͺ A K2-mode IDE runs the same FIR checkers in its own process, so the error should appear at the position above; nothing in this repository drives an IDE, and there is no IDE plugin
@ExperimentalTlalocApi β€” opt-in marker on the provisional surface βœ… @RequiresOptIn(ERROR) in :core on three surfaces (the four-worlds taxonomy, AllReduceAttrs, the kernel-choice and cost-model packages) with the criterion recorded in the annotation's own KDoc. Pinned twice: ExperimentalTlalocApiTest reads the marker and the marked/unmarked sets out of the class files (BINARY retention is invisible to reflection), and ExperimentalApiOptInTest runs a real K2JVMCompiler with no -opt-in and asserts the refusal, the @OptIn fix, and that a certified surface is not affected
Binary-compatibility baseline βœ… api/<module>.api committed for all twelve published modules that emit classes (the eleven libraries and the Gradle plugin), Java 25 ones included; checkKotlinAbi (the Kotlin Gradle plugin's ABI validation, experimental in 2.4.20) is wired into check. Negative-tested by adding a public function to :runtime-cuda and to :compiler-plugin and watching each fail with the diff. ./gradlew updateKotlinAbi re-baselines
Aggregated API reference βœ… ./gradlew apiDocs β†’ build/docs/api/index.html: one Dokka site over the eleven library modules. Not hosted; it emits about 200 unresolved-KDoc-link warnings, most in :ir
Hosted documentation site ⬜ No hosted API reference or documentation site

Hardware and platforms

Target Status Notes
NVIDIA GPU (CUDA, via PJRT) βœ… Everything marked βœ…-on-GPU was certified on a GB10 (Blackwell, aarch64). Other NVIDIA parts are expected to work but are not certified here
CPU βœ… Host interpreter + IREE-CPU. The interpreter is a correctness engine, not a fast CPU backend
Google TPU πŸ§ͺ The plugin lane and platform gating are written, and every device test (smoke, bf16, training, Qwen3 greedy parity) switches to libtpu with TLALOC_TEST_PJRT_TARGET=tpu; scripts/tpu/prepare.sh packs them for a VM that needs no build, rehearsed on CUDA. Nothing has executed on a TPU yet. TPU_BRINGUP.md is the runbook
Apple GPU (Metal, via PJRT) βœ… PjrtTarget.Metal runs DXIR programs through MetalHLO's PJRT plugin; the device suite runs on it with TLALOC_TEST_PJRT_TARGET=metal. Metal has no double, so f64 is computed in f32: f64-precision checks run at f32 tolerance, and the ones that need true doubles skip by name. See GETTING_STARTED.md
Cerebras Wafer-Scale Engine (via PJRT) πŸ§ͺ PjrtTarget.Cerebras runs DXIR programs through CerebrasStableHLO's PJRT plugin; the device suite runs with TLALOC_TEST_PJRT_TARGET=cerebras, on the Cerebras SDK's WSE simulator with CEREBRASSTABLEHLO_BACKEND=simulator. Measured there: 134 pass, 68 skip, 11 fail (all 11 typed "too large for one PE" errors: big KV caches, MoE weights, gated-delta states, CausalLm / Qwen3.5-MoE programs); since then GatedDeltaOps and BlockwisePagedAttention pass, with oversized computations spread over groups of PEs, leaving MoeExperts, Qwen35MoeParity and CausalLmTraining. Work in progress: CerebrasStableHLO and CerebrasHLO are not public yet. f64 is computed in f32, as on Metal. Simulator only; nothing has run on wafer hardware. See GETTING_STARTED.md
AMD / Trainium πŸ“ Named in the kernel registry's target matrix; no runtime lane
JVM βœ… The only build target declared today
Android / iOS / WASM ❌ The modules are KMP-structured (commonMain source sets), which makes these reachable later β€” but no such targets are declared or built, and nothing has been tested on them

Where the tests have been executed

Every βœ… above was certified on one machine, a GB10 (aarch64, JDK 25). This table lists where else the suite has run.

Lane Status Notes
GB10 (aarch64 Linux, JDK 25) βœ… The reference machine. Every GPU, PJRT, IREE, KPTX and cross-language-oracle row in this file was certified here
The whole suite, re-run clean-room βœ… On the GB10: ./gradlew test --rerun-tasks with 0 failures, scripts/onboarding-smoke.sh, all ten examples, quickstart shapeError failing as designed, a -Werror consumer compile, every published POM with its sources and javadoc jars, and both halves of examples/gpu-inference (a real TinyLlama-1.1B decoding ' Paris.\n\n2.' through /usr/bin/python3)
The library's suites on a JDK 21 βœ… -PtlalocTestJdk=21 over :core :ir :autograd :nn :stablehlo :maestro, green on OpenJDK 21.0.2 on the GB10, and in the CI lane below
The plugin under a foreign Kotlin compiler βœ… ForeignCompilerGuardTest, part of ./gradlew test, runs a real Kotlin 2.3.20 compiler (the release 0.1.0-alpha01 supports) with this build's plugin jar loaded: KotlinVersionGuard refuses by name as a compile error, and with unsafeAllowUnsupportedKotlin set the link failure that follows is a compile error naming both versions, not a crash. The same harness on Kotlin 2.4.20 compiles with no message. The plugin's source no longer compiles against 2.3.x (-PtlalocKotlinVersion=2.3.20): it reports through the Kotlin 2.4 diagnostic API
x86_64 Linux CI (ubuntu-latest) βœ… .github/workflows/build.yml, green on GitHub Actions at d0ca85b; later commits have not run there yet. A green runner means the platform-neutral subset passes: a runner has no GPU, no PJRT plugin, no IREE, no stablehlo-translate and no PyTorch oracle venv, and every test needing one self-skips by name
aarch64 Linux CI (ubuntu-24.04-arm) βœ… Same workflow, green at d0ca85b. The only aarch64 evidence from a machine other than the GB10
arm64 macOS CI (macos-15) βœ… Same workflow, green at d0ca85b. The MLIR round trips self-skip, because stablehlo-translate, sdy-opt and iree-compile are not installed
JDK 21 CI lane βœ… build.yml, job library-jdk21, green at d0ca85b
Next-Kotlin CI lane βœ… (probe) .github/workflows/kotlin-next.yml probes the newest Kotlin on Maven Central in a later feature release than the catalog's β€” 2.5.0-Beta1 while the catalog is 2.4.20. It runs with continue-on-error, so its checkmark is green by construction; read its step outcomes. Run locally, -PtlalocKotlinVersion=2.5.0-Beta1 compiles the plugin; its tests then stop at the version guard, which refuses 2.5 by design. With the guard bypassed in a local experiment, every plugin test except the guard's own passes under 2.5.0-Beta1