# OxiBonsai TODO

> Pure Rust 1-bit LLM inference engine for PrismML Bonsai models
> 619 source files, ~177,000 lines of Rust code, 4,667 tests passing across workspace — verified 2026-06-02 (0.2.0)

## Bonsai-Image Port (FLUX.2 Text-to-Image, Pure Rust) — COMPLETE

> **Status: ✅ DONE (shipped in 0.1.5)** — full pure-Rust text→image pipeline works end-to-end and is GPU-accelerated (CUDA steps=4 ≈31.7s on A4000-class, Metal ≈52–62s on Apple Silicon M3, cos≥0.999 throughout). Runs natively via `pipeline::text_to_image` and the `oxibonsai image` CLI subcommand.

**Overview** — Bonsai-Image 4B is PrismML's text-to-image diffusion model (same team as Bonsai LLM). It is based on Black Forest Labs' FLUX.2 Klein 4B (rectified-flow DiT). The DiT's matmul-heavy layers are quantized to ternary g128 `{-1,0,+1}` (128-weight groups with bf16 scale per group), which maps exactly to OxiBonsai's native `BlockTQ2_0_g128` format. The text encoder is Qwen3-4B (already supported by OxiBonsai — reused). The VAE is `AutoencoderKLFlux2` (convolution-based; the largest new subsystem).

### Confirmed Architecture (config.json)

- **DiT (`Flux2Transformer2DModel`)** — hidden 3072 = 24 heads × 128 head_dim; 5 double-stream + 20 single-stream blocks; `mlp_ratio` 3.0 (intermediate 9216); `in_channels` 128 (latent 32ch × 2×2 patch); `joint_attention_dim` 7680 (= 3×2560); `axes_dims_rope` [32,32,32,32] (4-axis RoPE = 128 total; differs from FLUX.1's 3-axis); `rope_theta` 2000; `guidance_embeds` false (no CFG → one forward per step); `timestep_guidance_channels` 256; `eps` 1e-6.
- **VAE (`AutoencoderKLFlux2`)** — `latent_channels` 32; in/out 3/3; `block_out_channels` [128,256,512,512]; 4× UpDecoderBlock2D; `layers_per_block` 2; `norm_num_groups` 32 (GroupNorm); `mid_block_add_attention` true; `patch_size` [2,2]; `act_fn` silu; `force_upcast` true (fp32 execution); spatial downsampling ×8 + patch ×2 = ×16 (512px → 32×32 = 1024 tokens).
- **Scheduler (`FlowMatchEulerDiscreteScheduler`)** — `num_train_timesteps` 1000; `shift` 3.0; `use_dynamic_shifting` true (resolution-dependent); `base_shift`/`max_shift` 0.5/1.15; `base/max_image_seq_len` 256/4096; `time_shift_type` exponential. Default inference: 4 denoising steps, guidance 1.0.
- **Text encoder** — standard Qwen3-4B (hidden 2560, 36 layers, 32 heads, 8 kv_heads GQA, head_dim 128, intermediate 9728, rope_theta 1e6). Existing Qwen3 implementation reused as a hidden-state extractor.

### Source Format

- Conversion source: `prism-ml/bonsai-image-ternary-4B-mlx-2bit` (all safetensors). PyTorch `.pt` (pickle) rejected per COOLJAPAN Pure Rust policy. Safetensors supported via existing dependency (Pure Rust).
- Quantization metadata (`quantization_config.json`): `format = "mlx-packed-affine"`, `bits 2`, `group_size 128`, `scale_dtype = "bfloat16"` (bf16, not FP16 — important for the converter).
- Skip patterns (kept bf16, < 5% params): `proj_out, x_embedder, context_embedder, time_text_embed, time_guidance_embed, norm_out, double_stream_modulation_{img,txt}, single_stream_modulation`. Quantized targets are attn/MLP linears in double/single blocks (matches whitepaper Figure 3).

### Phased Bring-Up Checklist (historical record)

- [x] **Phase 0 — Golden reference & harness** — DiT/VAE/scheduler/TE config confirmed; MLX ternary model downloaded; Python golden dumper created (text embeddings → per-step latents → decoded image saved as `.npy`).
- [x] **Phase 1 — Converter + single-tensor dequant parity [primary gate]** — MLX `mlx-packed-affine` 2-bit bit layout identified from mflux/MLX source → converter (MLX safetensors → GGUF; quantized layers = `BlockTQ2_0_g128`, skip/VAE/TE = bf16) → Rust dequant validated against Python dequant dump.
- [x] **Phase 2 — Text encoder (Qwen3-4B feature extraction)** — existing Qwen3 forward adapted for hidden-state extraction (LM head / sampling removed); 7680-wide conditioning vector assembled; TE weights loaded from MLX int4 affine format; hidden states compared against Python dump (te_parity cos≥0.999999).
- [x] **Phase 3 — DiT forward (1 block → 25 blocks)** — double-stream (×5): AdaLN → QKV (img/txt separate) → joint attention (concat→split) → out proj → AdaLN → MLP; single-stream (×20): AdaLN → fused QKV+MLP-in → attn ∥ MLP act → cat → fused out; 4-axis RoPE / RMSNorm / mlp_ratio 3.0; quantized linears reuse ternary GEMM; skip_patterns use bf16 dense; velocity prediction compared vs Python dump (dit_parity 59 taps cos≥0.999).
- [x] **Phase 4 — Flow-match sampler** — `FlowMatchEulerDiscreteScheduler` + dynamic shift (seq_len→μ, base 0.5/max 1.15, exponential); 4 steps, no CFG (`guidance_embeds=false`); initial bring-up injected Python latents to isolate cross-framework RNG differences.
- [x] **Phase 5 — VAE decoder (`AutoencoderKLFlux2`) [largest new subsystem]** — Conv2d (im2col+GEMM), GroupNorm(32), SiLU, ResNet blocks (layers_per_block 2), mid-block attention, 4-stage upsample (128/256/512/512), post_quant_conv, patch[2,2] unpack; latent(32ch)→RGB, fp32 (force_upcast); decoded image compared against Python dump (vae_parity 11 taps cos≥0.999).
- [x] **Phase 6 — PNG output + CLI (end-to-end)** — oxiarc-deflate Pure Rust PNG (COOLJAPAN compression policy; no flate2/zstd/zip); `oxibonsai image --prompt … --seed N --out x.png`; native Threefry RNG byte-matches MLX/mflux reference at `--seed 42`.
- [x] **Phase 7 — GPU acceleration & optimization** — Metal (default-on): DiT flash-attention 5.47×, ternary GEMM v10 1.89×, VAE GPU 3.2×; CUDA (native-cuda): total 3.2× (gemm_tq2 ~6×, flash-attn ~6.3×, stage0 context_embedder 59×); all cos≥0.999.

### Open Questions (resolved during bring-up)

All open questions from the planning phase were resolved during implementation:
- `joint_attention_dim=7680` = concatenated hidden states from 3 Qwen3 layers (×2560 each).
- 4-axis RoPE position IDs: `[0, h//32, w%32, 0]` for image tokens; `[0, 0, 0, j]` for text tokens.
- AdaLN modulation: double-stream blocks use 6 parameters (scale/shift/gate × 2 paths); single-stream blocks use 3.
- VAE `scaling_factor`/`shift_factor`: applied at pipeline level (latent packing/unpacking), not inside the VAE decoder.
- Latent token order: `noise[c, s] = packed[s, c]` where s = h*32+w (raster scan).

## Phase Status

| Phase | Status | Summary |
|-------|--------|---------|
| Phase 0: Foundation | ✅ Complete | Workspace, GGUF loader, metadata extraction |
| Phase 1: 1-Bit Kernels | ✅ Complete | dequant/gemv/gemm + property tests (20 kernel tests) |
| Phase 2: Transformer Engine | ✅ Complete | RMSNorm, SwiGLU, RoPE, GQA, 36-layer Qwen3 forward pass |
| Phase 3: Inference Runtime | ✅ Complete | KV cache, sampling, tokenizer, CLI (run/chat/serve/info), OpenAI-compatible server |
| Phase 4: Production Hardening | ✅ Complete | AVX2 SIMD, NEON, AVX-512, rayon, benchmarks, tracing, TOML config, proptests, integration tests, circuit breaker, health, metrics, recovery, Prometheus metrics, memory profiling, fuzz testing |
| Phase 5: Ecosystem Integration | ✅ Complete | All major subsystems implemented (see below) |
| Phase 6: Advanced Infrastructure | ✅ Complete | Multi-GPU, CUDA/Metal stubs, PagedAttention, GGUF v3, dynamic quant, profiler |
| Phase 7: Production Features | ✅ Complete | Model merging, checkpointing, flash decoding, attention sink, compression, ROUGE, MoD, adaptive sampling, oxibonsai-serve |
| Phase 8: Final Polish | ✅ Complete | K-quant formats, streaming GGUF, kernel tuning, comprehensive testing |
| Phase 9: Ternary Bonsai | ✅ Complete | TQ2_0_g128 block type, SIMD kernels, model variants, GGUF surface, export |
| Phase 10: Ternary CPU SIMD tiers | ✅ Complete | AVX2 / AVX-512 / NEON TQ2 GEMV with pos_mask+neg_mask decomposition |
| Phase 11: Metal TQ2 GEMV + dispatch | ✅ Complete | Per-kernel Metal TQ2_0_g128 dispatch in `gpu_backend/metal_*` |
| Phase 12: Native CUDA backend | ✅ Complete | NVRTC kernels, fused Q1 + TQ2 full-forward; ~21.9 tok/s on Ternary-Bonsai-1.7B |
| Phase 13.x: Fused Metal TQ2 full-forward | ✅ Complete | Single command buffer per token, ~13× speedup; ~50 tok/s on Ternary-Bonsai-1.7B |
| Phase 13.y: Ternary LM head on GPU | ✅ Complete | All 7 `OutputWeight::Ternary` guard sites closed (4 Metal + 3 CUDA); Metal ~54.5 tok/s avg on Ternary-Bonsai-1.7B |
| Phase 14: Runtime controllers (0.1.4) | ✅ Complete | `KvCachePolicy`, `AdaptiveLookahead`, `RequestRateTracker`, `RequestId`, expanded Prometheus surface |
| Phase 15: Advanced Inference Capabilities | ✅ Complete | FP8 quantization (E4M3/E5M2), FP8 kernels + `Fp8Kernel` trait, `AllowList`/`Sequence`/`LengthConstraint`, BNF grammar engine + Earley-based `GrammarConstraint` |
| Phase 15.x: FP8 SIMD + Grammar Cache | ✅ Complete | AVX2/AVX-512/NEON gather-based FP8 kernels; parallel FP8; LUT decode tables; Earley `allowed_tokens` LRU cache (O(1) on repeated state) |
| Phase 16: FP8 Export + Criterion Benchmarks | ✅ Complete | `ExportFormat::FP8E4M3/E5M2` in `oxibonsai-model`; 4 export tests; `bench_fp8_kernels` in `kernel_benchmarks.rs`; 14 benchmark cases |
| Phase 16A: FP8 Model Layers + Integration | ✅ Complete | `LinearFP8E4M3<'a>` + `LinearFP8E5M2<'a>` in `oxibonsai-model`; `OutputWeight::FP8E4M3/E5M2` variants; `ModelVariant::FP8Bonsai8B/4B/1_7B` with spec builders; FP8 block loaders in `weight_loaders.rs`; GPU paths return graceful Err; 6 integration tests |
| Phase 16B: Token-byte precomputation + first-byte index | ✅ Complete | `GrammarConstraint` precomputes all token byte sequences at construction; `first_byte_index: Box<[Vec<u32>; 256]>` eliminates per-step decode calls; `empty_token_ids` for EOS/special tokens; `vocab_size()` + `index_memory_bytes()` accessors; 5 new integration tests |
| Phase 17: FP8 KV cache + JSON Schema BNF + ARC/GSM8K + SmoothQuant + Unigram | ✅ Complete | `Fp8KvLayer`/`Fp8KvCache`/`KvCacheLevel::Fp8`; JSON Schema→Grammar compiler; ARC/GSM8K evaluators; SmoothQuant FP8 calibrator + channel-aware quant; Unigram tokenizer with HF format support |
| Phase 18: Standard GGUF formats + K-quant ext + WordPiece + WinoGrande/BoolQ + Regex→BNF | ✅ Complete | Q4_0/Q8_0 full stack; Q5_K/Q6_K K-quant layers; WordPiece tokenizer; WinoGrande/BoolQ evaluators; Regex→BNF compiler (Thompson NFA→Subset DFA→Grammar) |
| Phase 19: Q2_K/Q3_K/Q4_K/Q8_K + MMLU + GBNF + HellaSwag/TruthfulQA + Tool Calling | ✅ Complete | Full Q2_K/Q3_K/Q4_K/Q8_K block types, GEMV kernels, linear layers, model integration; MMLU evaluator with per-subject breakdown; GBNF parser (two-pass, `*`/`+`/`?` expansion); HellaSwag/TruthfulQA evaluators; Tool calling API (select_tool, build_tool_constraint, ToolRegistry, validate_tool_arguments) |
| Bonsai-Image Phases 0–7: text-to-image pipeline | ✅ Complete | oxibonsai-image crate; DiT TQ2 + VAE + TE 4-bit + PNG + MLX RNG; Metal ≈52–62s, CUDA ≈31.7s, steps=4 512² |
| Phase 33: Split block/types.rs (Policy Compliance Milestone) | ✅ Complete | Split `crates/oxibonsai-model/src/block/types.rs` (1853 lines, 147-line margin) into `block/types/` directory module with 10 sub-files (max 424 lines): `mod.rs` (35), `upload.rs` (41), `layer_stats.rs` (48), `scratch.rs` (67), `block_def.rs` (104, `pub struct TransformerBlock<'a>` + `new()`), `forward_stats.rs` (283), `helpers.rs` (283, layer_idx + try_full_layer_gpu/cuda), `forward_sw.rs` (330, sliding-window forward), `forward.rs` (419, main forward), `accessors.rs` (424, all ~85 quant-family accessor methods grouped by family). The single big `impl<'a> TransformerBlock<'a>` block split across 5 sub-files. Visibility widened: 19 `TransformerBlock` fields + 16 `ScratchBuffers` fields from private to `pub(super)` (siblings destructure them); 2 cfg-gated helpers (`try_full_layer_gpu`, `try_full_layer_cuda`) from `fn` to `pub(super) fn`. **Milestone**: after Phase 33, every Rust source file in the workspace has ≥ 227 lines of margin to the 2000-line policy ceiling. Top file is now `cuda_k_quant_prefill_kernels.rs` at 1773 (kernel-source-string container, no host code). 4625 tests pass; clippy clean. |
| Phase 32: Split forward_cuda.rs (Policy Compliance) | ✅ Complete | Split `crates/oxibonsai-model/src/model/types/forward_cuda.rs` (1801 lines, 199-line margin) into `forward_cuda/` directory module with 6 sub-files: `mod.rs` (34, cfg-gate + decls), `byte_helpers.rs` (119, 8 free byte-cast helpers), `q1.rs` (580, Q1 builders + 4 top-level dispatch entry points), `ternary.rs` (285, TQ2 prefill methods), `q_std.rs` (328, Q4_0/Q8_0 methods), `k_quant.rs` (424, K-quant methods). Max sub-file 580. The single big `impl<'a> BonsaiModel<'a>` block split into 5 sibling impl blocks; private inherent methods widened to `pub(super)` (sibling-sub-module helpers) or `pub(in super::super)` (entry points called from `types/mod.rs`). Inner-attribute `#![cfg(...)]` on `mod.rs` eliminates 19 repeated per-method cfg attributes. After Phase 32: every CUDA-related file at policy edge is now split (cuda_prefill, cuda_k_quant_prefill, forward_cuda all done). New top of file ranking is `block/types.rs` at 1853 (147-line margin). 4625 tests pass; clippy clean. |
| Phase 31: Cap-of-8 Batch-Kernel Regression Audit | ✅ Complete | New `tests/capof8_audit.rs` integration test in `oxibonsai-kernels` enumerates every batch-GEMM kernel by name and asserts both (a) the kernel entry-point appears in its source string and (b) the source contains the `col_base += 8` (CUDA) / `col_base += 8u` (MSL) outer-loop pattern. Audits 36 CUDA batch kernels (6 in `cuda_prefill_kernels` for Q1+TQ2 × {gemm, gemm_residual, fused_gate_up_swiglu}; 18 in `cuda_k_quant_prefill_kernels` for Q2_K/Q3_K/Q4_K/Q5_K/Q6_K/Q8_K × 3 variants; 6 in `cuda_fp8_prefill_kernels` for E4M3/E5M2 × 3 variants; 6 in `cuda_q_std_prefill_kernels` for Q4_0/Q8_0 × 3 variants) plus 10 Metal MSL batch kernels (4 in `kernel_sources::prefill` for Q1+TQ2; 6 in `kernel_sources::fp8_prefill` for E4M3/E5M2 × 3 variants). Two defensive host-only sanity tests confirm the assertion helper rejects sources missing the pattern or the kernel name. All 46 batch kernels currently audit-clean. Host-only; no GPU required; the assertions run on every CI build that enables `native-cuda` (Linux/Windows) or `metal` (macOS) features. Phase Status: 4625 workspace tests passing (4619 + 6). |
| Phase 30: Metal + Runtime Refactor (Policy Compliance) | ✅ Complete | Split `metal_graph.rs` (1948 lines → 7 files in `metal_graph/{mod, error, reformat, pipelines, buffers, graph, tests}.rs`, max 670 lines) and `constrained_decoding.rs` (1966 lines → 8 files in `constrained_decoding/{mod, error_trait, regex, json, sampler, allow_list, sequence, length}.rs`, max 724 lines). All public surface preserved: `MetalGraph`/`MetalGraphError`/`MetalWeightHandle` (gpu_backend re-export) plus the 11 constraint types (`ConstraintError`, `TokenConstraint`, `NoConstraint`, `RegexConstraint`, `JsonConstraint`, `JsonParseState`, `ConstrainedSampler`, `ConstrainedSamplerBuilder`, `AllowListConstraint`, `SequenceConstraint`, `LengthConstraint`). Internal helpers (`alloc_buf`, `div_ceil`, `set_scalar`, `upload_f32`, `download_f32`) re-exported with `pub(crate)`; NFA internals (`NfaState`, `RegexNfa`, `Fragment`) tightened from previously-private to `pub(super)` to satisfy nested test modules. After Phase 29 + Phase 30: zero files within 100 lines of the 2000-line policy; the top-ranked file is now 1853 (147 lines of margin). 4619 tests still passing; clippy clean. |
| Phase 29: CUDA Prefill Module Refactor (Policy Compliance) | ✅ Complete | Split `cuda_prefill.rs` (1989 lines → 6 files) and `cuda_k_quant_prefill.rs` (1941 lines → 5 files) into directory modules mirroring the `cuda_full_layer/` template. New layout: `cuda_prefill/{mod, state, launchers, encode_q1, encode_ternary, try_apis, tests}.rs` and `cuda_k_quant_prefill/{mod, state, launchers, encode, try_api, tests}.rs`. All public surface preserved (`try_cuda_prefill`, `try_cuda_prefill_ternary`, `init_prefill_modules`, `CudaPrefillBuffers`, `CudaPrefillModules`, `try_cuda_prefill_k_quant`, `CudaKQuantPrefillLayerParams`, `KQuantFormat`, `CudaKQuantPrefillModules`, `init_k_quant_prefill_modules`). Internal helpers gated to `pub(super)`. Largest CUDA host files now ≤ 1801 lines; the two files at policy risk (1989 / 1941) are gone. Zero changes to runtime behaviour: 4619+ workspace tests still green; clippy clean. |
| Phase 28: Metal FP8 Batch GEMM Primitives | ✅ Complete | 8 new MSL kernels (`MSL_GEMM_FP8_E4M3_V1`, `_RESIDUAL`, `MSL_FUSED_GATE_UP_SWIGLU_GEMM_FP8_E4M3_V1`, `MSL_GEMV_FP8_E4M3_PF_V1`, plus 4 E5M2 variants) in `kernel_sources/fp8_prefill.rs` mirroring Phase 26 CUDA design: column-major batch I/O, AoS 34-byte FP8 blocks, cap-of-8 outer loop, one simdgroup per row. New `metal_fp8_prefill.rs` module with own singleton (Device + CommandQueue + 6 lazy pipelines) exposes 6 public host fns: `metal_gemm_fp8_e4m3`/`e5m2`, `metal_gemm_fp8_e4m3_residual`/`e5m2_residual`, `metal_fused_gate_up_swiglu_fp8_e4m3`/`e5m2`. 7 CI-GPU-gated CPU↔Metal parity tests (incl. `batch_size=12` cap-of-8 discriminator + non-mul-of-8 row boundary) + host-only kernel-source-string assertions. Forward-path integration (5-stage encoder + LM-head dispatch) deferred to Phase 28.B. |
| Phase 27: Metal FP8 GEMV (E4M3/E5M2) | ✅ Complete | `MSL_GEMV_FP8_E4M3_V1` + `MSL_GEMV_FP8_E5M2_V1` MSL kernels (AoS 34B blocks, simdgroup-per-row reduction); `metal_fp8_kernels.rs` with `metal_gemv_fp8_e4m3`/`metal_gemv_fp8_e5m2` host functions (own singleton Device + CommandQueue, lazy pipeline init); `impl Fp8Kernel for KernelDispatcher` extended with Metal-first GPU dispatch (Metal → CUDA → CPU SIMD); `BonsaiModel::forward_prefill` and `forward_prefill_verify` skip fused Metal prefill for FP8 and fall through to per-token sequential (which now uses Metal GPU GEMV via dispatcher); 2 CI-GPU-gated parity tests + host-only MSL source-string assertions. |
| Phase 26: CUDA FP8 E4M3/E5M2 Batch Prefill | ✅ Complete | 8 NVRTC kernels (E4M3/E5M2 × gemm/gemm_residual/fused_gate_up_swiglu/gemv_pf); `try_cuda_prefill_fp8`; `CudaFP8PrefillLayerParams`; `forward_cuda_fp8.rs` with `try_cuda_prefill_with_lm_head_fp8`/`try_cuda_prefill_verify_fp8`; handle namespaces E4M3 26M-31M, E5M2 28M-33M; replaces sequential FP8 CUDA GEMV loop for prefill |
| Phase 25: K-quant CUDA Batch Prefill | ✅ Complete | 18 NVRTC kernels (Q2K/Q3K/Q4K/Q5K/Q6K/Q8K × gemm/gemm_residual/fused_gate_up_swiglu); `try_cuda_prefill_k_quant`; `KQuantFormat` enum; 42 K-quant block accessor methods; replaces sequential CUDA GEMV loop for all K-quant prefill |
| Phase 24   | CUDA Q4_0/Q8_0 Batch Prefill | ✅ Complete | Q4_0/Q8_0 batch GEMM prefill (CPU-only tests pass) |
| Phase 23: K-quant CUDA GEMV kernels | ✅ Complete | Six NVRTC kernels (Q2K/Q3K/Q4K/Q5K/Q6K/Q8K) added; `LinearQ*K::forward` dispatch to GPU; `BonsaiModel::forward`/`forward_prefill` K-quant early-dispatch added. |
| Phase 22: Q4_0/Q8_0 CUDA full-forward dispatch | ✅ Complete | `LinearQ4_0::forward`/`LinearQ8_0::forward` dispatch to `cuda_gemv_q4_0`/`cuda_gemv_q8_0` NVRTC kernels when CUDA available; `BonsaiModel::forward`/`forward_prefill` early-dispatch for Q4_0/Q8_0 skips the Q1 graph attempt. Compile-time `BlockQ4_0`/`BlockQ8_0` layout assertions added. |
| Phase 21: CUDA FP8 dispatch wiring + Q4_0/Q8_0 GEMV kernels | ✅ Complete | FP8 CUDA early-dispatch in `BonsaiModel::forward`/`forward_prefill` (skips Q1 graph, uses `KernelTier::Gpu` GEMV via block dispatch); 14 FP8 E4M3/E5M2 + 14 Q4_0/Q8_0 block accessor methods added to `TransformerBlock`; `cuda_gemv_q4_0`/`cuda_gemv_q8_0` NVRTC kernels (AoS 18B/34B blocks, warp-per-row, nibble decode for Q4_0, int8 for Q8_0). |
| Phase 20: CUDA Ternary Batch Prefill + FP8 GEMV | ✅ Complete | `try_cuda_prefill_ternary`: TQ2_0_g128 fused batch GEMM kernels (`gemm_tq2_g128_v7`, `gemm_tq2_g128_v7_residual`, `fused_gate_up_swiglu_gemm_tq2`) with cap-of-8 outer loop, SoA weight layout, per-token `encode_attn_phase_tq2`; replaces sequential CUDA prefill fallback for Ternary models. `cuda_gemv_fp8_e4m3`/`cuda_gemv_fp8_e5m2`: NVRTC FP8 GEMV kernels with AoS block decode, warp-per-row reduction; `KernelDispatcher::Gpu` tier dispatches to GPU path for both E4M3 and E5M2; host-only + CI-GPU-gated parity tests. |

## Phase 4 — Complete

- [x] **NEON (ARM) SIMD kernels** -- Implemented `simd_neon.rs` with dequant/gemv/gemm
- [x] **AVX-512 SIMD kernels** -- Implemented `simd_avx512.rs` (compile-checked, untested on ARM)
- [x] **Property-based tests** -- proptest roundtrip, linearity, cross-tier in `proptest_kernels.rs`
- [x] **Integration tests** -- `tests/integration_pipeline.rs` + `tests/cross_tier.rs` + `tests/feature_flags.rs`
- [x] **Tracing enhancement** -- EnvFilter + `#[instrument]` spans on generate, forward, block, server, sampling
- [x] **TOML configuration system** -- `OxiBonsaiConfig` with layered loading, `--config` CLI arg
- [x] **Circuit breaker** -- Overload protection with configurable thresholds
- [x] **Health checks** -- `/health` endpoint with readiness probes
- [x] **Inference metrics** -- tokens/sec, latency, request counts via `InferenceMetrics`
- [x] **Builder pattern API** -- `EngineBuilder`, `ConfigBuilder`, `SamplerBuilder`
- [x] **Sampling presets** -- Greedy, Balanced, Creative, Code presets
- [x] **Batch engine** -- Batched inference support
- [x] **Async engine** -- Tokio-based async inference
- [x] **Recovery strategies** -- Error classification and automatic recovery
- [x] **Prometheus metrics** -- Self-contained Prometheus-compatible metrics; `/metrics` endpoint via `InferenceMetrics` + `metrics_text()` renderer
- [x] **Memory profiling** -- `MemoryProfiler` + `MemorySnapshot`; RSS via Mach (macOS) / `/proc/self/statm` (Linux); peak tracking
- [x] **Fuzz testing** -- Property-based kernel + GGUF parsing fuzz tests; cargo-fuzz harnesses being added

## Phase 5 — Complete

- [x] **Streaming chat completions (SSE)** -- Implemented with `generate_streaming` + SSE handler
- [x] **WASM compilation target** -- Compiles to `wasm32-unknown-unknown`; `wasm_api.rs` provides WASM-safe generate API
- [x] **Bonsai-4B and Bonsai-1.7B model support** -- `Qwen3Config::bonsai_4b()` / `bonsai_1_7b()`, `ModelVariant::Bonsai4B` / `Bonsai1_7B`, `model_variants.rs` with full `ModelSpec` + `CapabilityProfile`
- [x] **OxiRAG integration** -- Pure Rust RAG pipeline (`oxibonsai-rag`): chunker, IdentityEmbedder, TfIdfEmbedder, VectorStore, Retriever, RagPipeline
- [x] **MeCrab/OxiTokenizer** -- Native BPE tokenizer stub (`oxibonsai-tokenizer`): Vocabulary, BpeMerges, OxiTokenizer, char-level fallback, ChatTemplate/chatml
- [x] **Advanced sampling** -- Mirostat v1/v2, Locally Typical, Eta, Min-P samplers; `SamplerChain` composable pipeline; `LcgRng`
- [x] **Speculative decoding** -- `SpeculativeDecoder` with draft/verify loop, acceptance rate tracking, `generate_speculative`
- [x] **LoRA adapter support** -- `LoraAdapter`, `LoraRegistry`, `BonsaiLoraSet`; apply/merge into weights
- [x] **Context window management** -- `ContextWindow` with TruncateLeft/Right/SlidingWindow/Summarize strategies; system-prompt preservation
- [x] **Beam search** -- `BeamSearchEngine` with configurable `beam_width`, length penalty, no-repeat n-gram, early stopping
- [x] **Token healing** -- `TokenHealer`, `TokenHealingConfig`, `HealingDecoder`
- [x] **Inference pipeline API** -- `InferencePipeline`, `PipelineBuilder`, `GenerationStrategy`, `StopReason`, `PipelineOutput`; convenience constructors (`chat_pipeline`, `code_pipeline`, `greedy_pipeline`)
- [x] **INT8 quantization** -- Per-tensor and per-channel INT8 quantization; `Int8Tensor` with `dequantize()` and `matvec()`
- [x] **GGUF writer** -- `GgufWriter` producing valid GGUF byte streams with metadata and tensor entries
- [x] **Model export** -- `export_to_gguf` supporting Float32, Q1_0_g128, Int8PerChannel formats; FP32 exception layers; `ExportStats`
- [x] **Prefix caching** -- `PrefixCache`, `CacheBlock`, `CacheSession`, `PrefixAwarePrefill` for KV reuse
- [x] **Semantic cache** -- `SemanticCache` with cosine-similarity lookup for repeated prompts
- [x] **MoE architecture** -- Mixture-of-Experts router and expert networks (`moe_router.rs`, `moe_expert.rs`)
- [x] **Tensor parallelism** -- `TensorParallelConfig`, column/row parallel linear; `AllReduceOp`
- [x] **Grammar-constrained decoding** -- `TokenConstraint` trait; `JsonConstraint`, `RegexConstraint`, `AllowListConstraint`, `SequenceConstraint`
- [x] **OpenAI API extensions** -- Function calling, logprobs, JSON mode, `n` completions, response format, frequency/presence penalties (`api_extensions.rs`, `api_types.rs`)
- [x] **Embeddings endpoint** -- `POST /v1/embeddings`; `EmbedderRegistry` (IdentityEmbedder + TfIdfEmbedder); float + base64 encoding; dimensions truncation
- [x] **Completions endpoint** -- `POST /v1/completions` with single/batch prompt, echo, logprobs stub
- [x] **RAG server endpoints** -- `/v1/rag/index`, `/v1/rag/query` (feature-gated)
- [x] **Rate limiting + middleware** -- `RateLimiter`, `MiddlewareConfig`; tower-based rate-limit and timeout layers
- [x] **Admin API** -- `/admin/status`, `/admin/config`, `/admin/reset-metrics`, `/admin/cache-stats`
- [x] **Evaluation harness** -- `oxibonsai-eval` crate: HellaSwag, MMLU, TruthfulQA benchmarks; `EvalRunner`, `EvalResult`
- [x] **Optimizer suite** -- SGD (with momentum, Nesterov, weight decay), Adam, AdamW; CosineAnnealing / WarmupConstant / Constant LR schedulers
- [x] **LoRA training scaffold** -- `LoraTrainer` with per-step Adam updates, gradient clipping, warmup; `TrainingStep` metrics
- [x] **Fine-grained gradient computation** -- `GradientContext`, `LayerGradients`, `backward_rms_norm`, `backward_swiglu`, `backward_linear`
- [x] **Continuous batching** -- `ContinuousBatchScheduler` for in-flight request management
- [x] **Model registry** -- `ModelVariant`, `ModelSpec`, `CapabilityProfile`; auto-detection from config dimensions
- [x] **Comprehensive integration tests** -- `tests/full_pipeline_tests.rs` (25 tests), `tests/quantization_pipeline_tests.rs` (10 tests), `tests/server_integration_tests.rs` (8 tests)

## Phase 17 — FP8 KV Cache + JSON Schema BNF + Eval + SmoothQuant + Unigram

- [x] **FP8 KV cache (Slice A)** — `Fp8KvLayer`/`Fp8KvFormat`/`Fp8KvCache` in `kv_cache_quant.rs`; per-row abs-max scale (`max|row|/448.0` E4M3, `/57344.0` E5M2); `push`, `get_key`, `get_value`, `get_keys_at`, `get_values_at`, `memory_bytes`, `memory_bytes_fp32_equivalent`, `clear`; `KvCacheLevel::Fp8` variant (ordinal 2, between Q8 and Q4, `memory_factor=0.5`, `tag="fp8"`); all `ordinal`/`tag`/`memory_factor`/`from_ordinal` match arms extended; 11 model tests + 3 runtime tests; clippy clean
- [x] **JSON Schema → BNF compiler (Slice B)** — `grammar/json_schema_compiler.rs` in `oxibonsai-runtime`; two-pass compiler (pre-allocate `$defs` NTs in pass 1, compile rules in pass 2); supports `string`/`integer`/`number`/`boolean`/`null`, `object` (required properties), `array` (items), `enum`, `anyOf`/`oneOf`/`allOf`, `$ref` to `#/$defs/...` and `#/definitions/...`, empty/boolean schemas; lazy shared NTs for digit/string rules; max depth 32; `JsonSchemaCompileError` (`InvalidSchema`, `UnsupportedKeyword`, `DanglingRef`, `DepthExceeded`, `InvalidJson`); 30 integration tests; re-exported from `grammar::` and `lib.rs`
- [x] **ARC + GSM8K evaluators (Slice C)** — `arc.rs` (ArcEvaluator with `ArcSplit::Easy/Challenge`, delegates to `McEvaluator`/`McLogitEvaluator`, `ArcResult`; handles 4- and 5-choice questions); `gsm8k.rs` (regex-free `extract_final_answer` scanning for `####` + trailing numeric, 1e-6 relative tolerance, `Gsm8kResult` with `accuracy_pct`/`no_answer_rate`; `gsm8k_example` builder); 10 ARC tests + 15 GSM8K tests; re-exported from `oxibonsai-eval::lib`
- [x] **SmoothQuant FP8 calibrator (Slice D)** — `smoothquant.rs` in `oxibonsai-model`: `SmoothQuantCalibrator` with per-channel `running_max_abs` accumulator, `record_activation`, `smooth_factors`, `layer_count`, `has_layer`; `SmoothQuantError` enum; `quantize_fp8_e4m3_smooth` / `quantize_fp8_e5m2_smooth` (apply `smooth_weights` then standard block quantize); `BlockFP8E4M3::quantize_with_channel_scales` / `BlockFP8E5M2::quantize_with_channel_scales` in `oxibonsai-core` (channel-wise divisor before block packing, zero-divisor guards); 12 tests; re-exported from model and core
- [x] **Unigram tokenizer (Slice E)** — `unigram.rs` in `oxibonsai-tokenizer`: `UnigramVocab` (Viterbi forward-pass lattice, `token_to_id` HashMap, UNK_PENALTY=-1e6, UTF-8 char-boundary safe, per-byte UNK fallback, `encode`/`decode`/`token_count`/`unk_id`); `UnigramError` (`EmptyVocab`, `UnkOutOfRange`, `DuplicateToken`); `OxiTokenizer::with_unigram` constructor + `is_unigram()` + `Option<UnigramVocab>` dispatch in `encode()`; `HfModelType` enum + `model.type` branch in `hf_format.rs` (`parse_bpe_model`/`parse_unigram_model` helpers, `into_tokenizer` branches on type); 14 inline tests + 8 integration tests + 4 hf_format tests

## Phase 16 — FP8 Full Stack + Grammar Acceleration

- [x] **FP8 export formats** — `ExportFormat::FP8E4M3/E5M2` in `oxibonsai-model/src/export.rs`; size estimate 34/32 bytes/weight using `div_ceil`; GGUF type ID 43/44 mapping; zero-copy `repr(C)` block serialization via raw pointer cast; 4 tests (roundtrip E4M3, roundtrip E5M2, size estimate, FP32 exceptions)
- [x] **FP8 kernel criterion benchmarks** — `bench_fp8_kernels` in `benches/kernel_benchmarks.rs`; 4 benchmark groups: dequant 8192w, GEMV 128×256, GEMV 512×512 (sequential vs parallel), GEMV dispatch sweep 32/256/1024 rows; throughput reporting in element-ops/sec; E4M3 and E5M2 coverage; scalar vs simd_auto vs rayon-parallel
- [x] **Token-byte precomputation + first-byte index (Phase 16B)** — `GrammarConstraint` now precomputes byte sequences for all vocab tokens at construction time; `first_byte_index: Box<[Vec<u32>; 256]>` accelerates `allowed_tokens` by skipping non-first-byte tokens without calling the decode fn or checking HashSet; `empty_token_ids` handles EOS/special tokens; `vocab_size()` accessor; `index_memory_bytes()` for observability; `advance()` uses precomputed `token_bytes` (no decode call per token); combined with Phase 15.x state cache for O(1) repeated-state hits; 5 new integration tests in `tests/grammar_cache_tests.rs` + 8 new unit tests in `src/grammar/constraint.rs`

## Phase 15 — Advanced Inference Capabilities

- [x] **FP8 E4M3FN quantization** — `BlockFP8E4M3` (32 weights × 1 FP8 byte + FP16 scale = 34 bytes/block); `fp8_e4m3_encode/decode` with bit-exact RNE rounding, saturation to ±448, NaN=0x7f; GGUF type ID 43 (PrismML extension); 56 tests in `oxibonsai-core`
- [x] **FP8 E5M2 quantization** — `BlockFP8E5M2` (same 34 bytes/block); `fp8_e5m2_encode/decode` with ±infinity support; GGUF type ID 44 (PrismML extension)
- [x] **FP8 GGUF surface** — `GgufTensorType::F8_E4M3/F8_E5M2`, `TensorType::F8_E4M3/F8_E5M2`, `ExtendedQuantType::F8_E4M3/F8_E5M2` (8.5 bits/weight), `is_fp8()` predicate
- [x] **FP8 reference kernels** — `dequant_fp8_e4m3/e5m2`, `gemv_fp8_e4m3/e5m2`, `gemm_fp8_e4m3/e5m2` (scalar reference); `Fp8Kernel` trait; `impl Fp8Kernel for KernelDispatcher`; 349 tests in `oxibonsai-kernels`
- [x] **FP8 SIMD kernels (Phase 15.x)** — `fp8_lut.rs` (OnceLock 256-entry decode LUTs); `simd_fp8_avx2.rs` (gather-based 8-wide, `_mm256_i32gather_ps`); `simd_fp8_avx512.rs` (16-wide, `_mm512_i32gather_ps`); `simd_fp8_neon.rs` (LUT + `vfmaq_f32`, 4-wide); tier-aware dispatch in `KernelDispatcher`; parallel FP8 GEMV/GEMM in `parallel.rs`; 18 parity tests in `tests/fp8_simd_parity.rs`
- [x] **Earley `allowed_tokens` cache (Phase 15.x)** — `grammar/cache.rs` `AllowedTokensCache` (LRU, 256-entry default, `Arc<[bool]>` masks); `state_hash()` on `EarleyRecognizer` (sorted chart items → `DefaultHasher`); `Mutex<AllowedTokensCache>` in `GrammarConstraint`; `with_cache_capacity` constructor; `cache_stats()` for observability; 12 integration tests in `tests/grammar_cache_tests.rs`
- [x] **`AllowListConstraint`** — finite-set token-sequence constraint for multiple-choice output; candidate activation bitmask with prefix tracking; 32 tests in `oxibonsai-runtime`
- [x] **`SequenceConstraint`** — exact token-sequence forcing constraint; unconstrained after sequence consumed
- [x] **`LengthConstraint`** — hard `[min_len, max_len]` bounds with optional stop token enforcement
- [x] **BNF grammar engine** (`oxibonsai-runtime/src/grammar/`):
  - `ast.rs` — `Grammar`, `Rule`, `Symbol` (byte-level terminals), `normalise_terminals()` single-byte chain expansion
  - `bnf_parser.rs` — hand-rolled two-phase BNF parser (lexer + parser); forward references, line continuation, escape sequences, `BnfParseError`
  - `earley.rs` — full Earley recognizer; FIRST set computation; predict/complete/scan; `next_byte_set()` lookahead; `clone_state()` for speculative token probing; left-recursive grammar support via set-based memoization
  - `constraint.rs` — `GrammarConstraint` implementing `TokenConstraint`; byte-level probe via `clone_state()` per candidate token; EOS handling
  - `examples.rs` — `arithmetic_grammar()`, `simple_ab_grammar()`, `csv_row_grammar()`, `json_lite_grammar()`, `palindrome_grammar()`
  - 93 tests in `oxibonsai-runtime`

## CI / Publish

- [x] Publish script (`scripts/publish.sh`) -- dry-run by default, `--for-real` for actual publish
- [x] Feature flag tests (`tests/feature_flags.rs`) -- compile-time validation of feature propagation
- [x] GitHub Actions CI pipeline -- `.github/workflows/ci.yml` (build, test, clippy, WASM, feature matrix)
- [ ] crates.io publish (awaiting approval)

## Phase 6 — Complete

- [x] **True multi-GPU inference** -- NCCL-style all-reduce/gather/scatter; `DeviceMesh` (tp×pp); rayon-simulated collectives; column/row weight partition + merge helpers; `multi_gpu.rs` + 27 tests
- [x] **CUDA/Metal backend** -- GPU kernel dispatch stubs: `GpuBackend` trait, `CpuBackend` (always-on), `CudaBackend` (feature = "cuda"), `MetalBackend` (feature = "metal"); `select_backend`, `gpu_matmul`; 18 integration tests; cuBLAS / Metal Performance Shaders left as future implementation
- [x] **Production GGUF with real model weights** -- Streaming `TensorChunkIter`; memory budget estimation (`estimate_memory_bytes`, `fits_in_budget`); lazy tensor metadata loading; `validate_gguf_file`; `gguf_loader.rs` + 25 tests
- [x] **Web UI** -- Chat interface; model management dashboard
- [x] **Distributed serving** -- Multi-node coordinator; consistent hashing for request routing
- [x] **Full MeCrab tokenizer integration** -- `NativeTokenizerBridge` added to `oxibonsai-runtime`; wraps `OxiTokenizer` with encode/decode/chat-format API; pure Rust, WASM-compatible, zero C/FFI; 10 integration tests all passing
- [x] **PagedAttention / vLLM-style KV management** -- Block-based KV cache for higher utilisation; `PagedKvCache`, `BlockPool`, `BlockTable`, `KvPage` in `oxibonsai-model/src/paged_kv_cache.rs`
- [x] **GGMLv4 / GGUF v3 support** -- Forward compat layer: `GgufVersion`, `ExtendedQuantType`, `GgufCompatReport`, `check_gguf_header`, `build_compat_report` in `gguf/compat.rs`; 20 integration tests all passing
- [x] **Sparse attention patterns** -- `SparseAttentionMask` with Local Window, Strided, Bigbird, Longformer, and Dilated patterns; `sparse_attention_forward`; memory reduction estimate; `sparse_attention.rs` + tests in `oxibonsai-model`
- [x] **Weight tying** -- `TiedEmbedding` with tied embedding/unembedding pair; `embed()`, `lm_head_logits()`, `set_embedding()` API; `TyingError` type; `weight_tying.rs` + tests in `oxibonsai-model`
- [x] **Dynamic quantization** -- `DynamicQuant` trait; `DynamicQ8_0`, `DynamicQ4_0`, `DynamicQ4_1` quantizers; `QuantizedTensor`, `quantize_model_dynamic`, `calibrate_scale`; `dynamic_quant.rs` + 21 tests
- [x] **Runtime profiler** -- `InferenceProfiler` with span tracking, `SpanRecord`, `ProfileReport`, stats (mean/min/max/p95/p99); flame-graph text renderer; `profiler.rs` + 20 tests
- [x] **Advanced RAG chunking** -- `SemanticChunker` with cosine-similarity boundary detection; `HierarchicalChunker` with multi-level chunking; `SlideWindowChunker`; 31 tests in `oxibonsai-rag`
- [x] **Tokenizer serialization** -- `save_json`/`load_json` for `OxiTokenizer`; vocabulary and BPE merge round-trip; `TokenizerSerializer`; 16 tests in `oxibonsai-tokenizer`
- [x] **GGUF model card** -- `ModelCard` with metadata (author, license, description, tags, datasets); `CardFormat` enum; `render_card`/`parse_card`; embedded in GGUF metadata via `oxibonsai-core`
- [x] **Token streaming metrics** -- `StreamingSession` with per-token timing; `TokenEvent`, `StreamingStats` (mean/p50/p95 TTFT + TPOT); `StreamingProfiler`; 22 tests in `oxibonsai-runtime`
- [x] **KV cache quantization** -- `QuantizedKvCache` with Q8 and Q4 formats; per-head scale factors; `KvCacheQuantizer`, `store_quantized`, `load_dequantized`; `kv_cache_quant.rs` + 18 tests

## Phase 7 — Complete

- [x] **KV cache quantization (INT8)** -- `QuantizedKvCache` with Q8/Q4 formats; per-head scale factors; `kv_cache_quant.rs`
- [x] **Model merging (SLERP, TIES, DARE, task vector)** -- `model_merge.rs` with multi-strategy merge support
- [x] **Checkpoint save/load (OXCK binary format)** -- `checkpoint.rs` with serialization/deserialization
- [x] **YaRN extended RoPE** -- `yarn_rope.rs` with YaRN-style RoPE scaling for extended context
- [x] **Flash Decoding** -- `flash_decode.rs` for optimized single-token decoding
- [x] **Attention Sink / StreamingLLM** -- `attention_sink.rs` with sink tokens for infinite context windows
- [x] **Dynamic activation quantization** -- `dynamic_quant.rs` with DynamicQ8_0/Q4_0/Q4_1 quantizers
- [x] **Inference profiler with FLOP counting** -- `profiler.rs` with span tracking, p95/p99 stats, flame-graph renderer
- [x] **GGUF model card generation** -- `ModelCard` with metadata; `render_card`/`parse_card` in `oxibonsai-core`
- [x] **Token streaming metrics (TTFT, TBT)** -- `StreamingSession`, `StreamingStats`, `StreamingProfiler` in `oxibonsai-runtime`
- [x] **Weight importance analysis and pruning** -- `pruning.rs` with structured/unstructured pruning
- [x] **PTQ calibration pipeline** -- `calibration.rs` with post-training quantization calibration
- [x] **BPE tokenizer training** -- `trainer.rs` in `oxibonsai-tokenizer` with BPE merge learning
- [x] **Tokenizer serialization** -- `save_json`/`load_json` for `OxiTokenizer`; vocabulary and BPE merge round-trip
- [x] **Advanced RAG chunking** -- Sentence, recursive, sliding window, markdown chunking in `oxibonsai-rag`
- [x] **Loss functions (CE, focal, distillation, contrastive)** -- `losses.rs` with full training loss suite
- [x] **Advanced LR schedulers (OneCycle, ReducePlateau, CyclicLR)** -- `lr_schedulers.rs` with advanced scheduling
- [x] **Sparse attention patterns (local window, BigBird, strided)** -- `sparse_attention.rs` with multiple patterns
- [x] **Weight tying (embedding/LM head sharing)** -- `weight_tying.rs` with `TiedEmbedding`
- [x] **Model compression pipeline (prune+quantize)** -- `compression.rs` with end-to-end compression
- [x] **ROUGE evaluation metrics (ROUGE-1/2/L/S)** -- `rouge.rs` in `oxibonsai-eval`
- [x] **Mixture of Depths (adaptive compute)** -- `mixture_of_depths.rs` with `ModRouter`
- [x] **Token budget tracking** -- `token_budget.rs` in `oxibonsai-runtime`
- [x] **RoPE scaling variants (linear, DynamicNTK, LLaMA 3.1, LongRoPE)** -- `rope_scaling.rs` with multiple strategies
- [x] **N-best hypothesis tracking** -- `nbest.rs` in `oxibonsai-runtime`
- [x] **Adaptive sampling (entropy cooling, repetition adaptation)** -- `adaptive_sampling.rs` with `AdaptiveChain`
- [x] **Generation quality metrics (repetition, diversity, BLEU)** -- `quality_metrics.rs` in `oxibonsai-runtime`
- [x] **oxibonsai-serve standalone binary crate** -- `crates/oxibonsai-serve` with CLI args, banner, OpenAI-compatible server
- [x] **Cross-attention layers** -- `cross_attention.rs` for encoder-decoder architectures
- [x] **Gradient checkpointing** -- `gradient_checkpoint.rs` for memory-efficient training
- [x] **Request deduplication** -- `dedup.rs` with FNV-1a hashing for duplicate request detection
- [x] **Model hot-reload** -- `hot_reload.rs` for live model swapping without restart
- [x] **Performance auto-tuning engine** -- `auto_tuner.rs` with CPU detection, SIMD tier selection, memory budget estimation, kernel benchmarking
- [x] **JSON Schema structured output constraint** -- `json_schema.rs` with schema parsing, validation state machine, constrained token generation
- [x] **Multi-model serving with LoRA routing** -- `multi_model.rs` / `lora_router` with dynamic model selection and LoRA adapter dispatch

## Phase 8 — Complete

- [x] **Q2_K and Q4_K quantization formats** — `BlockQ2K`/`BlockQ4K` with dequant/quantize in `oxibonsai-core/quant_k.rs`; 21 tests
- [x] **Streaming GGUF reader** — `GgufStreamParser` state machine for progressive network-loaded GGUF parsing; 22 tests
- [x] **Kernel performance tuning** — `PlatformProfile` detection, `TunedThresholds` auto-computation, cache-line aligned buffers, software prefetch hints; 29 tests
- [x] **Layer-level correctness tests** — RMSNorm, SwiGLU, RoPE, Attention, TransformerBlock reference comparisons; 23 tests
- [x] **Numerical stability tests** — Extreme inputs, overflow/underflow, long sequences, KV cache stress; 25 tests
- [x] **Full generate() pipeline integration tests** — End-to-end generation, determinism, sampling parameter effects, edge cases; 26 tests
- [x] **Sampling distribution statistical tests** — Chi-square goodness-of-fit, temperature/top-k/top-p/repetition penalty validation; 22 tests

## Phase 9 — Ternary Bonsai

- [x] **Quant block types** — `crates/oxibonsai-core/src/quant_ternary.rs`: `BlockTQ2_0_g128` (34B/128w), `BlockTQ2_0` (66B/256w), `TernaryCode`, scalar dequant/quantize, zero-copy `slice_from_bytes`, 11 tests
- [x] **GGUF surface** — `GgufTensorType::TQ2_0=37/TQ2_0_g128=42`, `writer::TensorType`, `KNOWN_QUANT_TYPES`, `compute_tensor_size_bytes`, `is_ternary()` helper, ID verification against ggml-quants.h, 6 tests
- [x] **Core re-exports** — `pub mod quant_ternary` + re-export `BlockTQ2_0{_g128}`, `TernaryCode`, constants from `oxibonsai_core`
- [x] **Scalar ternary kernels** — `dequant_ternary.rs`, `gemv_ternary.rs`, `gemm_ternary.rs` correctness reference; 9 tests per file
- [x] **SIMD ternary kernels** — NEON / AVX2 / AVX-512 `gemv_tq2_0_g128` + `dequant` + `gemm`; pos_mask+neg_mask decomposition; sign-canary tests
- [x] **`TernaryKernel` trait + dispatcher** — `traits.rs` + `dispatch.rs` impl; GPU tier falls back to scalar; 2 dispatcher tests
- [x] **`Qwen3Config::ternary_bonsai_*()` constructors** — three named constructors in `config.rs`, co-drift via 1-bit ctor call; 2 tests
- [x] **`ModelVariant::TernaryBonsai*` + detection** — three enum variants + `from_config_and_sample_tensor_type()`; all-ternary byte budget; 4 tests
- [x] **`ternary_bonsai_*_spec()` + capability profiles** — six-entry `all_specs()`, ±5 % byte budget tests; 4 tests
- [x] **`weight_loaders.rs` + `LinearTernary` layer** — `load_ternary_blocks`, `load_ternary_embedding`, `OutputWeight::Ternary`; sum-type for Attn/MLP; 3 tests
- [x] **`ExportFormat::TernaryG128`** — `quantize_ternary.rs`, norm-only FP32 exceptions, `default_fp32_exceptions_ternary_excludes_embedding` guard; 5 tests
- [x] **Phase 9 docs** — TODO.md Phase Status table + Memory/Performance rows + crate-local entries; README.md Ternary Bonsai subsection

## Phase 10 — Complete

- [x] **NEON TQ2 GEMV** — `simd_neon.rs` ternary GEMV with pos_mask + neg_mask decomposition over TQ2_0_g128 blocks
- [x] **AVX2 TQ2 GEMV** — `simd_avx2.rs` ternary GEMV over 128-weight groups with FP16 scale broadcast
- [x] **AVX-512 TQ2 GEMV** — `simd_avx512.rs` ternary GEMV with 512-bit lanes
- [x] **Ternary dispatch** — `TernaryKernel` routes via `KernelDispatcher::auto_detect()` across NEON / AVX2 / AVX-512 / scalar tiers
- [x] **Ternary dequant SIMD** — per-tier `dequant_ternary` for TQ2_0_g128 (mirrors the Q1 dequant layout)
- [x] **Apple Silicon CPU baseline** — NEON TQ2 path measured at ~7–8 tok/s end-to-end on Ternary-Bonsai-1.7B via `scripts/cli_ternary.sh`

## Phase 11 — Complete

- [x] **Metal TQ2_0_g128 GEMV kernel** — per-kernel TQ2 dispatch path in `crates/oxibonsai-kernels/src/gpu_backend/metal_dispatch.rs`
- [x] **`blocks_as_bytes_ternary` helper** — zero-copy TQ2 block upload to the Metal backend
- [x] **Per-kernel dispatch parity with Q1** — Metal TQ2 path validated against scalar reference so CPU and Metal produce byte-identical output at `--temperature 0 --seed 42`
- [x] **Metal graph groundwork** — `gpu_backend/metal_graph.rs` and `metal_prefill.rs` cover prefill + per-step GEMV composition ahead of the Phase 13.x fused path

## Phase 12 — Complete

- [x] **NVRTC kernel pipeline** — `gpu_backend/cuda_kernels.rs`, `cuda_attn_kernels.rs`, `cuda_prefill_kernels.rs`; runtime-compiled PTX in `kernel_sources/`
- [x] **CUDA graph capture** — `gpu_backend/cuda_graph/` groups every layer op (RMSNorm, gemv v7/v8/v9 with residual, SwiGLU, gate+up, Q1 AoS→SoA and TQ2 block→SoA reformat) into a reusable CUDA graph
- [x] **Fused Q1 + TQ2 full-forward** — `cuda_full_layer.rs` encodes a full Qwen3 layer into one graph launch for both quant families
- [x] **Prefill path** — `gpu_backend/cuda_prefill.rs` shares kernel sources with decode to keep the prefill/decode invariants identical
- [x] **`native-cuda` feature flag** — opt-in backend selected via `KernelDispatcher` alongside `cuda` (scirs2) and `metal`
- [x] **Measured** — ~21.9 tok/s on Ternary-Bonsai-1.7B end-to-end via `scripts/cli_ternary.sh cuda`

## Phase 13.x — Complete

- [x] **Fused Metal TQ2 full-forward** — `gpu_backend/metal_full_layer/` encodes the 14-dispatch sequence (pre-attn RMSNorm → fused QKV → QK-norm + RoPE → KV-store → batched attention → attn output + residual → FFN RMSNorm → fused gate+up → SwiGLU → down + residual) into a single GPU command buffer per token
- [x] **Fused Metal Q1 full-forward** — same command-buffer layout reused for 1-bit weights; shared with the TQ2 code path via a weight-family-generic encoder
- [x] **Command buffer reuse** — per-layer argument buffers cached in `metal_full_layer/types.rs` to avoid per-token allocations
- [x] **`scripts/bench_ternary.sh`** — 3-run CPU-vs-Metal averaging harness with best-of-N reporting
- [x] **`scripts/download_ternary.sh`** — HF `hf` CLI fetch + safetensors → GGUF conversion wrapper
- [x] **Measured** — ~50 tok/s (best ~57) on Ternary-Bonsai-1.7B on Apple Silicon (M-series), ~13× speedup vs the per-GEMV dispatch of Phase 11

## Phase 13.y — Complete (shipped in 0.1.3)

Closes all 7 `OutputWeight::Ternary(_)` CPU-fallback guard sites in `crates/oxibonsai-model/src/model/types.rs`.

### Milestone 1 — Metal ternary LM head (4 guard sites)

- [x] Pre-split `model/types.rs` (1857 lines → >2000 after M1 additions) via `splitrs` (completed 2026-04-24)
  - **Goal:** Keep every file under 2000 lines. Split into `types.rs` + `types_gpu_cache.rs` + `types_forward_metal.rs`.
  - **Files:** `crates/oxibonsai-model/src/model/types/mod.rs` (541), `gpu_cache.rs` (233), `forward_metal.rs` (1054), `forward_cuda.rs` (454).
  - **Tests:** `cargo nextest run -p oxibonsai-model --all-features` green after split.

- [x] Add `MetalGraph::encode_tail_and_commit_ternary` + update `encode_full_forward_ternary` routing + update `try_metal_full_forward_ternary` LM-head upload (completed 2026-04-24)
  - **Goal:** Kernel-layer ternary LM head support. `encode_tail_and_commit_ternary` mirrors `encode_tail_and_commit` but uses `dispatch_gemv_tq2` at the LM-head GEMV site. `encode_full_forward_ternary` routes to it. `try_metal_full_forward_ternary` uploads via `get_or_upload_tq2_weight_soa` with namespace `4_000_000+`.
  - **Files:** `crates/oxibonsai-kernels/src/gpu_backend/metal_full_layer/functions_2.rs`, `functions_3.rs`, `metal_graph.rs` (test).
  - **Tests:** `test_encode_tail_ternary_matches_reference` passes — bit-exact token match, logits within 1e-3.

- [x] Extend `CachedModelWeights` + remove short-circuit + wire ternary cache builder + remove 4 Metal guards (completed 2026-04-24)
  - **Goal:** Model-layer: `CachedModelWeights` extended with ternary fields; ternary cache builder wired; all 4 Metal Err guards replaced with `try_metal_*_ternary` routes; 3 new thin wrappers in `functions_3.rs`.
  - **Files:** `crates/oxibonsai-model/src/model/types/gpu_cache.rs`, `forward_metal.rs`; `crates/oxibonsai-kernels/src/gpu_backend/metal_full_layer/functions_3.rs`.
  - **Tests:** CPU vs Metal greedy-decode output byte-matches on Ternary-Bonsai-1.7B (temp=0, seed=42). Metal avg=54.5 tok/s (best=55.4) — improvement from ~50 tok/s baseline.

### Milestone 2 — CUDA ternary full-forward (3 guard sites)

- [x] Pre-split `cuda_full_layer.rs` (1687 lines → >2000 after M2 additions) via `splitrs` (completed 2026-04-24)
  - **Goal:** Split into `cuda_full_layer/mod.rs` + `encode_q1.rs` + `encode_ternary.rs` (new, M2 target).
  - **Files:** `crates/oxibonsai-kernels/src/gpu_backend/cuda_full_layer.rs` → directory.
  - **Tests:** `cargo check --features cuda` green (execution gated by CI-GPU).

- [x] Implement CUDA ternary full-forward path (completed 2026-04-24)
  - **Goal:** `encode_layer_into_ternary` (14 dispatches/layer, same split-dispatch shape as Metal Phase 13.x), `encode_full_forward_ternary`, `encode_lm_head_gemv_ternary` (single `launch_gemv_tq2_v1`), `try_cuda_full_forward_ternary` + `try_cuda_full_forward_ternary_with_gpu_lm_head`. Remove 3 CUDA guards at `try_cuda_full_forward_with_lm_head`, `try_cuda_prefill_with_lm_head`, `try_cuda_prefill_verify`.
  - **Files:** `crates/oxibonsai-kernels/src/gpu_backend/cuda_full_layer/encode_ternary.rs` (implemented); `crates/oxibonsai-model/src/model/types/forward_cuda.rs` (3 guards removed).
  - **Tests:** `test_encode_lm_head_gemv_ternary_matches_reference`, `test_encode_full_forward_ternary_matches_reference` (CI-GPU-gated). `cargo check --all-features` + clippy + 140 non-CUDA tests all green.
  - **Risk:** `CuGraphHolder` captured-graph pattern — capture+replay must be tested end-to-end, not just isolated kernel launches. Post-merge validation on CUDA hardware required.

## Phase 18 — Standard GGUF Formats + K-quant Extensions + WordPiece + WinoGrande/BoolQ + Regex→BNF

- [x] **Q4_0 + Q8_0 full stack** — `BlockQ4_0` (18 bytes/32w, nibble dequant) + `BlockQ8_0` (34 bytes/32w, i8 dequant) in `oxibonsai-core::quant_std`; scalar GEMV kernels in `oxibonsai-kernels::{gemv_q4_0, gemv_q8_0}`; `LinearQ4_0<'a>` + `LinearQ8_0<'a>` layers in `oxibonsai-model::layers::linear_standard`; `LinearLayer::{Q4_0, Q8_0}` + `OutputWeight::{Q4_0, Q8_0}` variants; Q4_0/Q8_0 paths in `weight_loaders.rs`; 25 block tests + 8 kernel tests each
- [x] **Q5_K + Q6_K K-quant extensions** — `BlockQ5K` (176 bytes/256w, 8 sub-blocks, 5-bit q, K-quant scale/min encoding) + `BlockQ6K` (210 bytes/256w, 16 sub-blocks, 6-bit q, symmetric centered) in `oxibonsai-core::quant_k_ext`; GEMV kernels in `oxibonsai-kernels::{gemv_q5k, gemv_q6k}`; `LinearQ5K<'a>` + `LinearQ6K<'a>` layers; full model integration; 12 integration tests in `tests/q5k_q6k_model_tests.rs`
- [x] **WordPiece tokenizer** — `WordPieceVocab` with greedy longest-match-first segmentation + `##`-prefixed continuation tokens; Unicode-safe char-boundary iteration; `WordPieceError`; `HfModelType::WordPiece` parsing branch in `hf_format.rs`; `OxiTokenizer::with_wordpiece` constructor + dispatch; 20 unit tests + 15 integration tests
- [x] **WinoGrande + BoolQ evaluators** — `WinoGrandeEvaluator` delegating to `McEvaluator`/`McLogitEvaluator` via `as_mc_dataset()` (1-based→0-based answer mapping); `BoolQEvaluator` with `extract_answer()` case-insensitive yes/no prefix matching; 12 WinoGrande + 16 BoolQ integration tests in `crates/oxibonsai-eval/tests/`
- [x] **Regex → BNF compiler** — `compile_regex(pattern) -> Result<Grammar, RegexCompileError>` via full Thompson NFA → Subset DFA (2048-state limit) → Grammar pipeline; `ByteSet` (256-bit bitset); supports literals, `.`, `[...]`/`[^...]`, `*`/`+`/`?`, `{n,m}`, `|`, grouping, `\d\w\s` classes, anchors (silently ignored); unsupported features (backreferences, lookahead, named groups) return `UnsupportedFeature`; re-exported from `grammar::` and `lib.rs`; 38 integration tests

## Phase 14 — Runtime Controllers (0.1.4)

- [x] **`KvCachePolicy`** — runtime KV-cache compression policy: FP16 → Q8 → Q4 transitions driven by EWMA pressure with hysteresis to prevent thrashing; configurable thresholds and `[min, max]` clamps; thread-safe via atomics
- [x] **`AdaptiveLookahead`** — speculative-decoding draft-length controller; updates `k` from a running EWMA of accepted tokens per step; cooldown to suppress oscillation; clamped to a configured `[min, max]`
- [x] **`RequestRateTracker` + `RequestRateAggregator`** — per-request EWMA tokens/sec, p50/p95 inter-token latency (TBT), queue-wait time; workload rollup aggregator with sliding-window ring buffer
- [x] **`RequestId`** — RFC 4122 UUIDv4-style 128-bit identifiers with thread-safe SplitMix64 generator (no external dependencies); hex / UUID / display formatting; round-trip parse
- [x] **Prometheus surface expansion** — five new gauges: `oxibonsai_request_tokens_per_second`, `oxibonsai_inter_token_latency_p50/p95_seconds`, `oxibonsai_queue_wait_seconds`, `oxibonsai_kv_cache_compression_level`
- [x] **`InferenceMetrics::update_request_rate(...)` and `update_kv_cache_level(...)` helpers** — wire the workload aggregator + KV policy directly into Prometheus output
- [x] **`SpeculativeDecoder::with_adaptive(...)`** — wires `AdaptiveLookahead` into the speculative loop; lookahead updates after every step
- [x] **CI re-enabled** — `.github/workflows/ci.yml` restored from `.disabled`; branches updated to `main` / `master` / `0.1.*`

## Performance Targets

| Platform | Target | Status |
|----------|--------|--------|
| x86-64 AVX2 (Zen4) | >= 20 t/s | 🔶 Kernels done, not benchmarked end-to-end |
| x86-64 AVX-512 | >= 30 t/s | 🔶 Kernels done, not benchmarked end-to-end |
| ARM NEON | >= 18 t/s | 🔶 Kernels done, not benchmarked end-to-end |
| WASM | >= 5 t/s | 🔶 Compiles; runtime throughput not benchmarked |
| ARM NEON (ternary) | >= 14 t/s | 🔶 Apple Silicon measured ~7–8 t/s end-to-end on Ternary-Bonsai-1.7B |
| x86-64 AVX2 (ternary) | >= 18 t/s | 🔶 Kernels done, not benchmarked end-to-end |
| x86-64 AVX-512 (ternary) | >= 26 t/s | 🔶 Kernels done, not benchmarked end-to-end |
| Metal fused TQ2 (Apple Silicon) | >= 30 t/s | ✅ Measured avg=54.5 t/s (best=55.4) on Ternary-Bonsai-1.7B (Phase 13.y, 2026-04-24; was ~50 t/s after Phase 13.x) |
| CUDA native TQ2 (NVIDIA GPU) | >= 20 t/s | ✅ Measured ~21.9 t/s on Ternary-Bonsai-1.7B |

## Memory Targets

| Metric | Target | Status |
|--------|--------|--------|
| Model weights (mmap) | 1.15 GB | ✅ mmap zero-copy implemented |
| KV cache (4k context) | ~256 MB | ✅ Implemented |
| Runtime overhead | < 50 MB | 🔶 Profiler available via `MemoryProfiler`; end-to-end not yet characterized |
| Total (4k context) | < 1.5 GB | 🔶 Not profiled end-to-end |
| Ternary-Bonsai-8B | ~1.75 GB | TQ2_0_g128 (34B/128w) |
| Ternary-Bonsai-4B | ~900 MB | TQ2_0_g128 (34B/128w) |
| Ternary-Bonsai-1.7B | ~390 MB | TQ2_0_g128 (34B/128w) |

## Deferred follow-ups from the 2026-06-01 external blog review

Source: technical review at https://nazquadri.com/blog/everyday-thoughts/oxibonsai-1-bit-pure-rust/

All 5 blog-flagged issues were verified against the code; the in-scope ones are fixed (temperature-discard bug, SIMD docs accuracy, CPU↔Metal parity guard, base-server engine pool — plus token_embd Arc-share + CPU default `min(4, cores)`). The 3 items below were deferred: none were raised by the blog, and each is a new feature or large architectural change whose effort + parity risk is disproportionate to a review fix.

### 1. De-singletonize `MetalGraph` for true GPU concurrency — LOW priority / high effort / high risk
- **What**: `GLOBAL_METAL_GRAPH` (`crates/oxibonsai-kernels/src/gpu_backend/metal_graph/graph.rs:23`) is a process-global singleton owning ONE `kv_cache` Mutex + shared scratch buffers, so the engine pool clamps GPU→1. True parallel GPU decode would need N command queues / N KV caches / N scratch-buffer sets.
- **Why deferred**: Core Metal-backend rewrite, not a fix. A single GPU serializes compute anyway → the only win is overlapping CPU-side glue / async submission, not matmul throughput (the CPU pool already covers real N-way deployment). Re-architecting KV/scratch risks the byte-identity determinism just validated. High cost, low marginal benefit on single-GPU hardware.

### 2. AVX-VNNI / NEON-UDOT INT8 dot-product kernel tier — LOW priority / medium effort
- **What**: All 1-bit/ternary kernels currently expand weights to ±scale and accumulate in FP32 FMA. An INT8 dot-product tier (AVX-VNNI on x86, NEON-UDOT on ARM) could speed up CPU decode.
- **Why deferred**: New kernel + dispatch entry + parity validation (INT8 accumulation differs numerically from FP32 FMA → risks the `cos≥0.999` gate), not a fix. The blog only asked for accurate SIMD docs (DONE). x86 VNNI can't be benchmarked on this ARM Mac (NEON-UDOT could). Noted as a possible future enhancement.
- **Best entry if pursued**: the NEON-UDOT path first (testable on this Mac).

### 3. Pool the `rag_server` `RagState` for concurrent RAG serving — LOW priority / low effort
- **What**: `crates/oxibonsai-runtime/src/rag_server.rs` has its own separate `RagState` Mutex (distinct from the base `/v1/chat/completions` path that was pooled). Apply the same `EnginePool` pattern.
- **Why deferred**: Separate subsystem, outside the blog's scope. Not a pure copy-paste — RAG adds a retrieval / vector-store stage with its own concurrency considerations. The same GPU-singleton clamp applies → CPU-only benefit on this Mac. Lowest-risk of the three; do it when RAG concurrency is actually needed.

## imagen CUDA GPU optimization follow-ups (2026-06-02)

Context: this session optimized the `oxibonsai image` (FLUX.2 Klein) CUDA path from steps=4 ≈101s → ≈31.7s (**3.2×**), cos=1.0 throughout (FP32, **no tensor cores**): TE attention AVX2 + 4-bit dequant parallelised; DiT `gemm_tq2` retiled 32×32/2×2 → 128×128/8×8 transposed-shared (~6× kernel); DiT flash-attention rewritten warp-cooperative (lane-split head_dim + `__shfl_xor_sync`, ~6.3× kernel); DiT stage0 `context_embedder` ported CPU→GPU (`encode_gemm_f32`, ~59× on that op). The two items below are **deferred levers**, each established by a code-read (no benchmark). Both are recorded so a future session need not re-derive them.

### TE persistent GPU weight cache — design-confirmed lever; implement only when multi-image throughput can be benchmarked
- **Static finding (code-read, no bench): the TE (Qwen3-4B 4-bit, ~2.1 GB) re-dequantises (CPU) AND re-uploads (HtoD) every Linear weight on EVERY image — nothing persists across `text_to_image` calls.**
  - `te::cuda_gpu::te_matmul_gpu` (`crates/oxibonsai-image/src/te/cuda_gpu.rs:93`): `get_or_upload_f32_weight(key, weight)` (`:115`) → `encode_gemm_f32` (`:116`) → **`evict_f32_weight(key)` (`:117`)**. The evict drops the device-weight cache entry after every GEMM, so the next call re-uploads. (The evict exists to fix a pointer-key collision — the 4-bit dequant yields transient host buffers whose address is recycled across Linears.) Called per Linear from `matmul` (`crates/oxibonsai-image/src/te/forward.rs:494`, GPU branch `:544`).
  - `TeWeights::get` (`crates/oxibonsai-image/src/te/weights.rs:157`), `Source::Mlx4bit` arm (`:174`): returns `Rc::new(model.load_tensor(name)?)` — a **fresh CPU dequant** (`mlx4bit::dequantize_mlx_4bit_affine`) on every call, explicitly **No-cache** (`:173`). `linear()` (`:182`) wraps `get()`. (The `NpyDir` source *does* cache (`:160`–`:170`); the 4-bit MLX source we actually use does not.)
  - Path is **per-image, not one-time init**: `pipeline::compute_cond` (`crates/oxibonsai-image/src/pipeline.rs:325`) opens `TeWeights::open_mlx_4bit` (`:334`) + `encoder.forward` (`:338`), invoked once per `text_to_image` generation (`:355`, cond arm `:417`); `TeWeights` is constructed fresh each call (no process-lifetime persistence).
- **Conclusion: YES — the lever structurally exists.** A process-lifetime cache of the dequantised-f32 + GPU-resident TE weights would let the **2nd-and-later image's TE skip both the CPU dequant and the HtoD upload (GEMM only)**.
- **Caveat / why not now**: **no gain for a single image** — each weight is used exactly once per forward, so it is uploaded once either way; the win is purely **multi-image throughput** (batch/server), which this single-image CLI cannot benchmark. A real implementation must also (a) restore stable cache keys — cache the dequantised f32 with a stable identity and drop the per-call `evict_f32_weight` — and (b) trade back the no-cache policy's ~16 GB → ~2.5 GB RAM frugality plus GPU residency of the f32 weights. **Implement when a multi-image throughput benchmark environment exists.**

### VAE back-stage residency — net-new, large; revisit only when VAE latency is the priority
- **2026-06-02 evaluation**: after the kernel + stage0 wins, nsys (steps=1) shows the imagen GPU time is now **transfer-bound** (cuMemcpyDtoH ≈5.0 s + HtoD ≈4.0 s ≫ all kernels ≈4.3 s). The VAE decode (~7.8 s) bounces its large 512×512 feature maps (up to ~400 MB) CPU↔GPU **per op**: the per-op `conv2d_gpu` / `groupnorm_gpu` / `silu_gpu` / `upsample_gpu` are all host-in/host-out (`crates/oxibonsai-image/src/vae/cuda_gpu.rs`). The conv itself already keeps im2col+GEMM device-resident (`CudaGraph::encode_conv2d_f32_im2col`), so the per-op cost is the input-HtoD + output-DtoH + host bias-scatter between ops.
- **Right fix = residency, NOT implicit-GEMM** (the conv compute / im2col is small): keep the map device-resident across the whole decode (upload latents + weights once, download RGB once). This is the discrete-GPU inversion of Metal's unified-memory wash — and unlike the DiT residual `h` (small ⇒ residency measured WASH this session), VAE maps are huge ⇒ a real ~4–5 s win.
- **Why deferred = net-new, no template**: both `vae/gpu.rs` (Metal) and `vae/cuda_gpu.rs` (CUDA) are per-op; there is **no resident VAE decode to mirror**. It needs new device kernels (conv-output transpose+bias scatter, residual-add) + device-in/out conv (tiled) / groupnorm / silu / upsample + a resnet/upsample/conv-norm/conv-out orchestration = the session's largest single implementation. ROI is **one-shot** (~12 % of steps=4: 7.8 → ~4 s). **Revisit only when VAE / low-step-count latency becomes the priority.**

### DiT — parity-safe (FP32) floor reached this session; negative results (do not re-derive)
- **DiT residual `h` / q/k/v residency = measured WASH.** The resident tensors are small (≪ the VAE's ~400 MB maps), so the PCIe bounce is cheap; the real stage0 glue was the CPU `context_embedder` GEMM (now on GPU, 59×). Full single-block residency moved steps=4 only **45.6 → 45.4 s** (within noise). **Do not re-attempt DiT residency** — the discrete-GPU residency lever pays off ∝ resident-tensor size, and the DiT's tensors are too small (the opposite of the VAE maps in the subsection above).
- **gemm_tq2 (6×) + flash-attention (6.3×) are at the FP32 ceiling** (~2.3 s kernels/step; DiT sample ~3.3 s/step, the remaining glue ~0.5 s/step is small). Further kernel speedup would require **TF32/FP16 tensor cores — deliberately NOT used** to hold cos≥0.999 parity. So ~3.3 s/step is the parity-safe DiT floor; **don't re-attempt FP32 kernel micro-opt** (mirrors the Metal v10→v11 "local optimum" result).

### Cross-backend note (macOS side)
- The TE per-image **re-dequant** (`crates/oxibonsai-image/src/te/weights.rs:173`, `Source::Mlx4bit` No-cache) is **backend-agnostic** — it re-runs on Metal too. So the *dequant-cache* half of the TE lever above would also help **Metal** multi-image throughput; only the *HtoD GPU-residency* half is CUDA-specific (on Metal the upload is free via unified memory, so Metal's multi-image win from this lever would be dequant-only).
