diff --git a/.github/workflows/larql-boundary.yml b/.github/workflows/larql-boundary.yml index 17aead012..6086fbff8 100644 --- a/.github/workflows/larql-boundary.yml +++ b/.github/workflows/larql-boundary.yml @@ -44,6 +44,15 @@ jobs: with: components: clippy, rustfmt + # protoc on PATH satisfies the `cfg(windows)` no-op in + # `larql-router-protocol`'s build.rs. Required because the + # "Run examples" step below builds `larql-demos`, which depends + # on `larql-router-protocol` directly. + - name: Install protoc (Windows) + if: runner.os == 'Windows' + shell: pwsh + run: choco install protoc -y --no-progress + - name: Cache cargo registry + build artefacts uses: actions/cache@v6 with: @@ -72,8 +81,8 @@ jobs: - name: Run examples run: | - cargo run -p larql-boundary --example encode_decode - cargo run -p larql-boundary --example gate_decision + cargo run -p larql-demos --example encode_decode + cargo run -p larql-demos --example gate_decision coverage: name: coverage · ubuntu diff --git a/.github/workflows/larql-core.yml b/.github/workflows/larql-core.yml index 3124bc55b..d79b16fe6 100644 --- a/.github/workflows/larql-core.yml +++ b/.github/workflows/larql-core.yml @@ -44,6 +44,15 @@ jobs: with: components: clippy, rustfmt + # protoc on PATH satisfies the `cfg(windows)` no-op in + # `larql-router-protocol`'s build.rs. Required because the + # "Run examples" step below builds `larql-demos`, which depends + # on `larql-router-protocol` directly. + - name: Install protoc (Windows) + if: runner.os == 'Windows' + shell: pwsh + run: choco install protoc -y --no-progress + - name: Cache cargo registry + build artefacts uses: actions/cache@v6 with: @@ -83,11 +92,11 @@ jobs: - name: Run examples run: | - cargo run -p larql-core --example edge_demo - cargo run -p larql-core --example graph_demo - cargo run -p larql-core --example algorithm_demo - cargo run -p larql-core --example filter_demo - cargo run -p larql-core --example serialization_demo + cargo run -p larql-demos --example edge_demo + cargo run -p larql-demos --example graph_demo + cargo run -p larql-demos --example algorithm_demo + cargo run -p larql-demos --example filter_demo + cargo run -p larql-demos --example serialization_demo coverage: name: coverage · ubuntu diff --git a/.github/workflows/larql-demos.yml b/.github/workflows/larql-demos.yml new file mode 100644 index 000000000..33644017f --- /dev/null +++ b/.github/workflows/larql-demos.yml @@ -0,0 +1,87 @@ +# larql-demos cross-platform CI +# +# The crate holds no logic — every target is an `examples/*.rs` binary — +# so there is nothing to unit-test or measure coverage on. What CI must +# guarantee is that the demos still compile against the crates they +# demonstrate, and that the weight-free ones still run. +# +# Demos needing model weights are compiled here but not executed; they +# are exercised by hand against a real vindex. + +name: larql-demos + +on: + push: + branches: [main] + paths: + - 'crates/larql-demos/**' + - 'Cargo.toml' + - 'Cargo.lock' + - '.github/workflows/larql-demos.yml' + pull_request: + branches: [main] + paths: + - 'crates/larql-demos/**' + - 'Cargo.toml' + - 'Cargo.lock' + - '.github/workflows/larql-demos.yml' + workflow_dispatch: {} + +jobs: + test: + name: test · ${{ matrix.os }} + runs-on: ${{ matrix.os }} + timeout-minutes: 30 + + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, windows-latest, macos-14] + + steps: + - uses: actions/checkout@v7 + + - name: Install stable Rust + uses: dtolnay/rust-toolchain@stable + with: + components: clippy, rustfmt + + # protoc on PATH satisfies the `cfg(windows)` no-op in + # `larql-router-protocol`'s build.rs. `larql-demos` depends on + # `larql-router-protocol` directly for the server demos. + - name: Install protoc (Windows) + if: runner.os == 'Windows' + shell: pwsh + run: choco install protoc -y --no-progress + + - name: Cache cargo registry + build artefacts + uses: actions/cache@v6 + with: + path: | + ~/.cargo/registry + ~/.cargo/git + target + key: ${{ runner.os }}-cargo-demos-${{ hashFiles('**/Cargo.lock') }} + restore-keys: | + ${{ runner.os }}-cargo-demos- + + - name: Format check + run: cargo fmt -p larql-demos -- --check + + - name: Check (all targets) + run: cargo check -p larql-demos --all-targets + + - name: Clippy (warnings as errors) + run: cargo clippy -p larql-demos --all-targets --no-deps -- -D warnings + + # The weight-free demos, i.e. exactly the set larql-core and + # larql-boundary used to run before these moved here. + - name: Run the weight-free demos + run: | + cargo run -p larql-demos --example edge_demo + cargo run -p larql-demos --example graph_demo + cargo run -p larql-demos --example algorithm_demo + cargo run -p larql-demos --example filter_demo + cargo run -p larql-demos --example serialization_demo + cargo run -p larql-demos --example encode_decode + cargo run -p larql-demos --example gate_decision diff --git a/.github/workflows/larql-vindex.yml b/.github/workflows/larql-vindex.yml index e3fb7704a..ed8b1c730 100644 --- a/.github/workflows/larql-vindex.yml +++ b/.github/workflows/larql-vindex.yml @@ -106,6 +106,18 @@ jobs: - name: Clippy run: cargo clippy -p larql-vindex --all-targets -- -D warnings + # E0 — VINDEX2 preservation. Named separately from the general test + # step so a failure reads as "the shipped generation regressed" rather + # than as one red test among hundreds. Any regression blocks merge: + # finding one now identifies the introducing commit, finding one after + # another ten changes produces archaeology. + # + # Scope caveat: this is the generation-boundary subset. The full matrix + # (decode, WALK, slice, publish/pull) needs multi-GB checkpoints and + # runs locally against tests/goldens/e0 via scripts/e0-capture-goldens.sh. + - name: E0 - VINDEX2 preservation (generation boundary) + run: cargo test -p larql-vindex --test e0_generation_boundary + - name: Tests run: cargo test -p larql-vindex diff --git a/Cargo.lock b/Cargo.lock index d415d4a28..658efb633 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2177,6 +2177,7 @@ dependencies = [ "blas-src", "cc", "criterion", + "larql-execution", "larql-models", "libc", "memmap2", @@ -2216,6 +2217,39 @@ dependencies = [ "thiserror 2.0.18", ] +[[package]] +name = "larql-demos" +version = "0.1.0" +dependencies = [ + "axum", + "blas-src", + "larql-boundary", + "larql-compute", + "larql-core", + "larql-inference", + "larql-kv", + "larql-lql", + "larql-models", + "larql-router-protocol", + "larql-server", + "larql-vindex", + "ndarray", + "openblas-src", + "safetensors", + "serde", + "serde_json", + "tokenizers", + "tokio", + "tokio-stream", + "tonic", + "tower", + "tracing-subscriber", +] + +[[package]] +name = "larql-execution" +version = "0.1.0" + [[package]] name = "larql-factory" version = "0.1.0" @@ -2246,6 +2280,7 @@ dependencies = [ "larql-compute", "larql-compute-metal", "larql-core", + "larql-execution", "larql-kv", "larql-models", "larql-router-protocol", @@ -2283,6 +2318,7 @@ dependencies = [ "larql-boundary", "larql-compute", "larql-compute-metal", + "larql-execution", "larql-inference", "larql-vindex", "ndarray", @@ -2465,7 +2501,9 @@ dependencies = [ "larql-compute", "larql-compute-metal", "larql-core", + "larql-execution", "larql-inference", + "larql-kv", "larql-models", "larql-vindex-spec", "libc", diff --git a/Cargo.toml b/Cargo.toml index 420e5c21e..a3c232827 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,11 +8,13 @@ members = [ "crates/larql-core", "crates/larql-vindex", "crates/larql-vindex-spec", + "crates/larql-execution", "crates/larql-factory", "crates/larql-inference", "crates/larql-kv", "crates/larql-lql", "crates/larql-cli", + "crates/larql-demos", "crates/larql-server", "crates/larql-router", "crates/larql-router-protocol", @@ -34,6 +36,7 @@ default-members = [ "crates/larql-kv", "crates/larql-lql", "crates/larql-cli", + "crates/larql-demos", "crates/larql-server", "crates/larql-router", "crates/larql-router-protocol", diff --git a/Makefile b/Makefile index 9f50803c2..2ca2158e8 100644 --- a/Makefile +++ b/Makefile @@ -54,11 +54,11 @@ larql-core-bench: cargo bench -p larql-core --bench graph larql-core-examples: - cargo run -p larql-core --example edge_demo - cargo run -p larql-core --example graph_demo - cargo run -p larql-core --example algorithm_demo - cargo run -p larql-core --example filter_demo - cargo run -p larql-core --example serialization_demo + cargo run -p larql-demos --example edge_demo + cargo run -p larql-demos --example graph_demo + cargo run -p larql-demos --example algorithm_demo + cargo run -p larql-demos --example filter_demo + cargo run -p larql-demos --example serialization_demo larql-core-coverage: @if ! command -v cargo-llvm-cov >/dev/null 2>&1; then \ @@ -536,8 +536,8 @@ larql-boundary-bench-test: cargo test -p larql-boundary --benches larql-boundary-examples: - cargo run -p larql-boundary --example encode_decode - cargo run -p larql-boundary --example gate_decision + cargo run -p larql-demos --example encode_decode + cargo run -p larql-demos --example gate_decision cargo run -p larql-boundary --example accuracy larql-boundary-coverage: diff --git a/README.md b/README.md index 0f9a8686b..b0961ea31 100644 --- a/README.md +++ b/README.md @@ -376,6 +376,15 @@ gemma3-4b.vindex/ feature_labels.json # Probe-confirmed labels ``` +**Container generations.** `index.json`'s `version` is the sole discriminator — +schemas 1–2 are **VINDEX2** (what `extract` writes, and what every published +vindex is today), schema 3 is **VINDEX3**, the successor container for sparse +models. One binary reads both; `larql show` and `larql verify` dispatch on the +version and describe each generation in its own terms rather than flattening +one into the other. VINDEX3 is still draft — the ABI is not frozen and +`extract` does not emit it yet — so in practice everything below is VINDEX2. +See [`crates/larql-vindex/docs/vindex3-format-spec.md`](crates/larql-vindex/docs/vindex3-format-spec.md). + Three extraction levels: | Level | CLI Flag | LQL Syntax | Size (f16) | Enables | @@ -464,7 +473,7 @@ delta on Metal, which the per-engine bench numbers confirm. | `markov-rs` | residual stream | derivative | exact logits under arch contract | **98.0** | | `markov-rs-codec` | compressed residuals | derivative | bounded KL | **98.1** | | `boundary-per-layer` | per-layer codec residuals | derivative | bounded KL per-layer | **98.7** | -| `unlimited-context` | KV (within window) + checkpoints | derivative | exact within window | 94.2 | +| `windowed-checkpoint` | KV (within window) + checkpoints | derivative | exact within window | 94.2 | | `turbo-quant` | quantised K/V | canonical (destructive) | bounded KL | 85.0 | | `boundary-kv` | K/V + boundary frames | canonical | exact logits | composes `standard` | | `apollo` | boundary retrieval store | n/a (retrieval) | task-level | orthogonal | diff --git a/ROADMAP.md b/ROADMAP.md index 7bf763fe1..baf1c4297 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -269,6 +269,261 @@ accrete features while the MoE path (the actual bet) stays thin. --- +## VINDEX3 — successor serving container (added 2026-08-02) + +**Thesis: the format boundary is the place to make sparse serving predictable.** +VINDEX2 can *observe* which pages faulted; VINDEX3 can *state* what an operation +will read before it runs. That is the difference between paging a multi-terabyte +model and planning one. + +Spec: [`crates/larql-vindex/docs/vindex3-format-spec.md`](crates/larql-vindex/docs/vindex3-format-spec.md) +(draft-2). Experimental programme: [`docs/vindex3-experiments.md`](docs/vindex3-experiments.md), +registry programme `vindex2`. Generations are named so the number equals +`index.json.version`: schemas 1–2 → VINDEX2, schema 3 → VINDEX3. + +**Coexistence, not migration.** One binary serves both generations, dispatched +solely on `index.json.version`. VINDEX2 keeps its loader, its weight objects and +its production behaviour untouched; VINDEX3 keeps its catalogue, profile, route +and authority model until binding. The shared layer is execution and +orchestration, never physical storage. **`extract` must keep defaulting to +VINDEX2** until V2-1 acceptance passes — a silent default change would evaporate +E0's premise. + +### Shipped + +| Commit | Milestone | +|---|---| +| `f13bf385` | Reference MoE execution — fixture A matches an independent oracle below 1e-6, fused and decomposed agreeing at every checkpoint | +| `dd2017db` | Real Gemma semantic routing parity over real VINDEX2 bytes | +| `f5dd256e` | Production router kernel bound — bit-identical routing ladder | +| *(pending)* | **The container itself** — fixture A written to disk as `index.json` schema 3 + `moe_manifest.json` + a LYRW v2 bank, opened, bound and executed bit-identically; `show`/`verify` dispatch on generation | + +Three properties established, each independently useful: + +1. **Bound reference execution is numerically correct** (fixture A vs oracle). +2. **Resolution does not leak into decode** — 64× population costs ~1.9×, which + is the router term and nothing more. +3. **The bound plan predicts its physical page working set exactly** — 200 pages + predicted, 200 resident, zero overshoot, 1.63% of a 192 MiB layer after one + token. Residency becomes computable rather than observable, which is what + placement, prefetch and remote transfer all need. + +Plus two defects fixed that were not VINDEX3's: `larql verify` rendered findings +in `HashMap` order and so disagreed with itself between runs; the separate-tensor +MoE extractor wrote no expert store. + +### Two ladders, deliberately separate + +VINDEX3 has an **execution** half and a **container** half, and they were built +in that order. Every parity result before the container existed bound its +operands out of a VINDEX2 file — so what was proven was the executor, not the +format: + +```text +proven first the VINDEX3 executor, fed VINDEX2 operands, matches production +proven second a VINDEX3 container can be written, opened, bound and executed +``` + +Keeping the ladders apart matters because a green execution ladder says nothing +about whether a VINDEX3 *file* exists, and for a long time none did. + +```text +container ladder +[x] c0 write index.json schema 3 + moe_manifest.json + LYRW v2 bank +[x] c1 detect_generation reports V3 from a real directory, not a JSON literal +[x] c2 open, validate the manifest, resolve storage keys to files +[x] c3 bind from container-resolved regions and execute — bit-identical +[x] c4 fused and decomposed storage agree under one programme id +[x] c5 structural verify with {layer, entry, role} defects; CLI dual-generation +[ ] c6 fixtures B–D (GPT-OSS, Inkling, Mini-K3) — proves nothing is hard-coded +[ ] c7 WALK/DESCRIBE parity over in-place bank regions +[ ] c8 a real Gemma MoE layer written as a VINDEX3 container +[ ] c9 every Gemma layer — the first real model that *is* a VINDEX3 container +``` + +Gate status, stated precisely: **V2-0 and V2-1 are closed for the rows fixture +A can carry**, not in full. Outstanding on V2-0 are profile-authority +derivation and variant-selection refusal; on V2-1, the "not hard-coded" row +(needs fixtures differing on expert count, top-K and shared banks) and +WALK/DESCRIBE parity. `extract` therefore still defaults to VINDEX2, and must +until those close. + +### The rung ladder to the first VINDEX3 Gemma token + +```text +[x] rung 0 fixture A through the generic reference path +[x] rung 0.5 real Gemma routing parity, VINDEX3 bound over VINDEX2 bytes +[x] rung 1 production router kernel bound, bit-identical +[ ] rung 2 production Q4_K x Q8_K expert kernel bound +[ ] rung 3 full-layer residual delta parity +[ ] rung 4 every MoE layer, then final logits +[ ] rung 5 greedy token parity through normal `larql run` dispatch +``` + +Rung 2 begins with **Q8_K activation identity, checked before any expert runs** — +a difference there contaminates all eight expert comparisons and makes every +later diagnostic noise. + +### What's next (set 2026-08-02, after PR #197) + +Ordered by what unblocks what, not by size. Each item states the condition that +closes it, so "done" is not a judgement call. + +**1. `layer_ffn_or_moe` — the other five engines. CLOSED 2026-08-02.** +`layer_ffn_or_moe` returns `Result, BoxRefusal>`, all ten call +sites propagate, and the gate in `larql-kv/tests/strict_refusal/engines.rs` +runs **eight** expert-routing engines × prefill/decode × three `RefusalKind`s. +The baseline tag may now say engine-wide. + +The rewind half came out better than the prediction. The prediction was that +residual-canonical engines could rewind where K/V-canonical ones could not; the +answer is that **all of them can**, by two different mechanisms, and +`engine_state.rs` proves it in the strong form — the retried token is +bit-identical to never having refused, for every one of the eight: + +```text +residual-canonical markov-rs, markov-rs-codec, boundary-per-layer + → the step writes `stored` only after the last fallible + call; `hot_kv` is a droppable derivative, taken up + front and left None on the error path +K/V-canonical standard, turbo-quant, windowed-checkpoint + → the cache grows before the FFN can refuse, so the step + truncates its appends: `truncate_kv` on the handle, + `CompressedLayer::truncate_rows` byte-exactly (rows are + appended at fixed offsets and never re-encoded, so the + codec is lossy against its *input*, not against what + was stored), `truncate_kv_rows` on the window shadow +``` + +Two engines stopped taking their store by value (`store.take()` → `as_mut()`), +which also removed a latent bug: any failure used to leave `self.store` as +`None`, so the *next* call reported "decode_step called before prefill" — a +dead engine wearing a misleading message. + +Exactly one case genuinely cannot rewind, and now says so: +`windowed-checkpoint` archives a window and saves its boundary checkpoint when +the window fills, so a refusal *after* a close returns +`EngineError::StateInvalidated` rather than a retryable refusal. Same for a +`standard` window already at capacity — append-then-evict leaves the row count +unchanged while the oldest row is gone. + +**1b. `no-cache` and `apollo` — the dense-only forwards. CLOSED 2026-08-02.** +Found while closing item 1: neither consulted `forward_moe_full_layer` at all, +so on a hybrid-MoE arch both ran the dense half of every layer and returned an +apparently valid answer. That is a *worse* failure than a degraded one — a +different model wearing the same answer shape, undetectable downstream — so it +was treated as a semantic disqualification rather than as missing propagation. +The two needed different corrections because the seam is in a different place: + +```text +no-cache forwards through `kv_prefill_run`, which *takes* an FfnBackend + → gave it real dispatch. That helper is also the oracle the + dispatch ring is compared against, so an oracle that skipped the + expert half would have made every MoE parity comparison agree + about the wrong answer. +apollo forwards through `forward_from_layer` / `forward_raw_logits`, which + live in larql-compute *below* the FfnBackend seam and construct + their own dense `ViewFfn` — no caller-supplied backend can reach + them → refuses the architecture up front, `RefusalKind::Unsupported` + (operands fine, this executor cannot serve them, pick another). +``` + +Real dispatch stays preferable for apollo, and means threading an `FfnBackend` +through `forward_layer_range` — a change to the forward, not to the engine. +Until then it is not usable as an apparently conformant MoE engine, which is +the point. + +Two more transactional bugs fell out, both of the kind only a refusal can +expose. `no-cache` pushed the decode token onto its list *before* the +re-forward could refuse, so a caller who fixed the cause and retried would +have forwarded the same token twice — the exact double-append the contract +exists to prevent; the token list is its entire continuation state, so the +push is now undone on failure. And `kv_decode_step_run` appended each layer's +K/V before the FFN could refuse, so the oracle itself is now transactional: +truncate back to the entry lengths, or report `StateInvalidated` when the +cache is windowed at capacity and eviction has already discarded a row. + +*Standing:* every `EngineKind` variant is now classified and gated — +`RoutesExperts` (nine, sweeping prefill/decode × three kinds) or +`NoExpertSeam` (apollo, refusing the architecture with an executing route, so +the refusal provably comes from the engine and not the route). + +**2. Variant-selection refusal.** Closes V2-0 outright. +`Vindex3Index::declares_profile` is a name check; §9.1 wants a profile that +selects an absent variant to fail naming the region set, the requested variant +and the variants physically present. Self-contained — no new execution path. + +**3. WALK/DESCRIBE parity.** Closes V2-1 except shared banks. Gate KNN over +in-place bank regions must return identical top-K to a v1-style extracted +`gate_vectors.bin` control on fixture A. This is the row that keeps "the model +IS the database" true of VINDEX3 rather than only of VINDEX2. + +**4. A real Gemma layer as a VINDEX3 container** (container ladder c8/c9). The +first real model that *is* a VINDEX3 container rather than one bound over +VINDEX2 bytes. c8 is one layer; c9 is all of them, at which point `extract` +gaining a VINDEX3 mode becomes a question rather than a violation. + +**Not on the critical path, but adjacent and cheap to start: the continuation- +state intervention harness.** `larql-kv` already owns incremental decode with +real K/V continuity, explicit next-token forcing, and a state-policy taxonomy +that names exactly the question a persistence experiment asks — which parts of +a continuation are carried by the emitted token, the residual, and the K/V +history. What is missing is causal read/write access to that state during a +decode step. + +PR #197 set the precedent for how that should look. `KvDispatch::truncate_kv` +is a research/recovery capability declared on the trait, implemented on CPU, +defaulting to *unsupported* rather than silently copying to host — which is the +shape a `MutableKvView` / `KvIntervention` seam should follow, at the layer +where attention appends and reads, never by exposing `KvHandle`'s +representation. + +One trap is already known and should be inherited rather than rediscovered: a +checkpoint that records cache *lengths* is not a checkpoint under a sliding +window, because append-then-evict leaves the count unchanged while the oldest +row is gone. `StandardEngine::rewind_is_sound` encodes that test; a fork API +needs the same one or it will hand out silently wrong donor state under +`markov-bounded`. + +### Standing method + +Established by repeated failure, not preference: + +- **Bind, never reconstruct.** A bridge that dequantises into an + incumbent-shaped temporary can reach numerical parity while proving nothing + about the binding architecture. `as_f32_slice()` hands over stored bytes or + refuses with a typed reason. +- **Ladders, not end-to-end tolerances.** A single residual-delta tolerance + blends router accumulation order, softmax, renormalisation, activation + quantisation, integer rounding and reduction order; passing it establishes + nothing and failing it identifies nothing. The router ladder localised a + 7e-4 disagreement to post-processing in one run — it was a missing bound + operand, not the BLAS-vs-index-order accumulation it would have been blamed on. +- **Mutation-check every new test.** Several have passed for the wrong reason, + including one that could never have caught the bug it was named for. +- **Suspect the instrument first.** This programme has produced roughly three + measurement defects per real code defect: a process-global allocation counter + under a parallel test runner, replay parameters defaulted instead of read from + the record, and a `\b` in a normaliser that BSD `sed` silently ignores. + +### Not discharged + +- **E0-FULL.** Decode rows and all 632 WALK ranking lines match; the remaining + 12 rows need the prescribed baseline reconstruction (baseline binary at + `6eae5ea` → baseline extraction → current reader against that artifact). + Status stands at: E0-CI green, E0-FULL decode rows green, remaining rows not + discharged. A runner now exists (`scripts/e0-verify-goldens.sh`); before it, + the goldens were an assertion nobody made. +- **OLMoE goldens** pin a decode panic that the separate-tensor MoE extractor fix + has since removed. They need a deliberate re-capture with the reason recorded. +- **CLI generation dispatch** — `detect_generation` exists and is guarded by + E0-CI, but 15 call sites still assume VINDEX2. +- **`extract --format vindex3`** — not needed until the container round-trip + rung, and doing it earlier would weaken the Gemma comparison by introducing + re-extraction as a second candidate cause. +- **Mini-K3, Kimi-Linear, K3** — the conformance envelope beyond Gemma. + +--- ## Query / Edit / Interpret — first-class functionality track (added 2026-05-28) **Thesis: the differentiated functionality is the database, not the tok/s.** @@ -388,7 +643,7 @@ item, not just a competitive-parity item. - **Grid (CPU MoE on remote shards)**: 18.3 tok/s 1-shard / 17.3 tok/s 2-shard local-loopback. Multi-host LAN/cross-region scaling unblocked. - **Remote FFN (dense)**: `larql run --ffn URL` + `larql serve --ffn-only` wired end-to-end. - **gRPC grid**: 2-shard self-assembling grid live-validated on 26B A4B. -- **4 KV-cache engines**: MarkovRS (287×), UnlimitedContext (254×), TurboQuant (4×), Apollo (20,000×) — all at ~95 tok/s on Gemma 3 4B Metal. +- **4 KV-cache engines**: MarkovRS (287×), WindowedCheckpoint (254×), TurboQuant (4×), Apollo (20,000×) — all at ~95 tok/s on Gemma 3 4B Metal. - **Wire format negotiation** (2026-05-07): f16 is now the default for all grid traffic (50% bandwidth reduction). i8 symmetric quantised residuals available opt-in (`LARQL_I8_WIRE=1`, 75% reduction). Content-type negotiation via `Accept` header; f32 fallback for non-grid clients. - **Per-layer latency routing** (2026-05-07): `HeartbeatMsg.layer_stats` carries EMA avg_ms + p99_ms per layer; router routes to the server with lowest per-layer latency (falls back to requests_in_flight when no data yet). - **WebSocket token streaming** (2026-05-07): `WS /v1/stream` now supports `{"type":"generate","prompt":"...","max_tokens":N}` command with per-token frames and cancel support. SSE streaming on `/v1/chat/completions` was already fully wired. @@ -1807,7 +2062,7 @@ stages, smallest-blast first: "diverging" was the compiler enumerating the work, not the work being unbounded. Every KvEngine (standard, no_cache, markov_residual, markov_residual_codec, boundary_per_layer, boundary_kv, turbo_quant, - unlimited_context, apollo) now owns a `dequant_scratch` field; quant methods + windowed_checkpoint, apollo) now owns a `dequant_scratch` field; quant methods dequant into it and the forward resolves through `WeightsView::with_scratch` — **0 `&mut ModelWeights` quant methods, 0 `weights.tensors.extend` merges on the engine/serving path.** Per-engine pattern: bulk-convert the engine's @@ -2135,7 +2390,7 @@ V3 is the genuinely-new-territory item. | # | Test | Prior evidence | What it falsifies | What it produces | Effort | |---|------|----------------|-------------------|------------------|--------| -| **V1 ✅ DONE 2026-05-31 — FALSIFIED (dense)** | Hash routing across all layers (extend exp 27) | **Exp 27 Gemma 3 4B L0 at top-2048/d_ffn (20% mask) → KL=0.030.** Walk boundary sweep (April 2026) progressively pushed the walk down through layers on Gemma 3 4B. **One-layer one-model evidence in hand.** | "5× FFN bandwidth reduction holds at end-to-end output, not just one layer" → **FALSIFIED.** Per-layer KL ≤ 0.05 thresholds DON'T compound: applied together they give +5.4 to +7.7 bits/token NLL and 78–95% drift on all 3 dense archs. The per-layer screen is anti-correlated with the truth. Deployable bandwidth ~2.4–2.9× (gate projection still paid), not 5×, and catastrophic anyway. | **DELIVERED:** per-layer threshold tables + compounding NLL/drift + cheap-route realizability + honest bandwidth, 3 dense archs (`bench/aim-validation/v1_*.json`), harness `examples/walk_ffn_v1_hash_routing.rs`, writeup [`docs/diagnoses/v1-hash-routing.md`](docs/diagnoses/v1-hash-routing.md). **MoE-within-expert version OPEN** (dense harness measures the wrong object on the 26B → needs expert-aware tooling). | ~1 week (done) | +| **V1 ✅ DONE 2026-05-31 — FALSIFIED (dense)** | Hash routing across all layers (extend exp 27) | **Exp 27 Gemma 3 4B L0 at top-2048/d_ffn (20% mask) → KL=0.030.** Walk boundary sweep (April 2026) progressively pushed the walk down through layers on Gemma 3 4B. **One-layer one-model evidence in hand.** | "5× FFN bandwidth reduction holds at end-to-end output, not just one layer" → **FALSIFIED.** Per-layer KL ≤ 0.05 thresholds DON'T compound: applied together they give +5.4 to +7.7 bits/token NLL and 78–95% drift on all 3 dense archs. The per-layer screen is anti-correlated with the truth. Deployable bandwidth ~2.4–2.9× (gate projection still paid), not 5×, and catastrophic anyway. | **DELIVERED:** per-layer threshold tables + compounding NLL/drift + cheap-route realizability + honest bandwidth, 3 dense archs (`bench/aim-validation/v1_*.json`), harness `chris-experiments/larql_probes/examples/walk_ffn/walk_ffn_v1_hash_routing.rs`, writeup [`docs/diagnoses/v1-hash-routing.md`](docs/diagnoses/v1-hash-routing.md). **MoE-within-expert version OPEN** (dense harness measures the wrong object on the 26B → needs expert-aware tooling). | ~1 week (done) | | **V2 ✅ DONE 2026-05-31 — CONFIRMED** | FP4 generality (extend exp 26 across archs) | **Exp 26: gemma3-4b-f16.vindex is 99.83% FP4-friendly per-feature without QAT (down is the tail at 99.65%).** Single-arch evidence in hand. | "FP4-friendliness is universal, not Gemma-3-4B specific" → **CONFIRMED.** ≥99.8% per-feature R<16 across Gemma 3 4B + Granite 3B/8B (reproduces exp 26's 99.83% exactly; down the tail). Predictive E2M1 +0.116 bits/tok vs f32, beats Q4-int. No QAT. | **DELIVERED:** static scan (`fp4_q1_scan`, generalized) + predictive NLL (`walk_ffn_v2_fp4_nll`, real E2M1 codec), artifacts `bench/aim-validation/v2_*_scan.json`, writeup [`docs/diagnoses/v2-fp4-generality.md`](docs/diagnoses/v2-fp4-generality.md). Llama/Mistral/MoE-expert weights not covered (need f16 exports). | ~1 week (done) | | **V3 ~ PARTIAL 2026-05-31** | mmap'd vindex with sparse access on disk-resident frontier MoE | **None.** This is the genuinely-new-territory item. Risk dominates the long-term tier confidence (~52%, revised 2026-05-31). | "Disk locality + page-fault behaviour is acceptable when only top-k experts fire" → **partial:** cold scattered read ~100µs p50/140µs p99, warm ~0.04µs (~2380× gap). Steady-state hinges on cache hit rate. | **DELIVERED (feasibility):** cold-read probe (`mmap_cold_read_probe`, F_NOCACHE + verified-cold mmap faults), artifact `bench/aim-validation/v3_granite-30b.json`, writeup [`docs/diagnoses/v3-disk-resident-mmap.md`](docs/diagnoses/v3-disk-resident-mmap.md). **DEFERRED:** steady-state fault-rate + end-to-end tok/s on a >RAM model — needs >128 GB-class vindex or Linux/cgroup box (128 GB machine can't force RAM-pressure paging). | ~2 weeks | | **V4** | **Compound test** (V1+V2+V3 stacked end-to-end on a real MoE model) | **D-RMS-FUSE Phase 1 (2026-05-09)**: predicted ~0.2 ms/tok savings collapsed to zero. ADR-015 has a concrete instance. | "Independent wins compound multiplicatively, not destructively" — per ADR-015. The framing's central claim. | End-to-end tok/s on Gemma 4 26B-A4B (or larger if available) with hash routing + FP4 + mmap'd disk-resident vindex active simultaneously. Measure perplexity degradation, tok/s, and compare to product-of-individual-speedups prediction. | ~1 week (after V1–V3) | @@ -2176,7 +2431,7 @@ verify. Driver: today's `KvEngine` (in `larql-kv`) and `ComputeBackend` (in `larql-compute`) are unaware of each other. The four research KV engines -(MarkovRS, UnlimitedContext, TurboQuant, Apollo) live in research-only +(MarkovRS, WindowedCheckpoint, TurboQuant, Apollo) live in research-only bench paths; the production decode loop bypasses them. And every backend (CPU, Metal, future Vulkan/CUDA) hides under a single trait that doesn't let engines express *intents* (windowed attention, K/V recompute, @@ -2189,7 +2444,7 @@ Three landed specs in `crates/larql-inference/docs/specs/`: - [`kv-engine-unification.md`](crates/larql-inference/docs/specs/kv-engine-unification.md) — KvEngine trait + dispatch in `larql-inference`; `larql-kv` ships six engines (`Standard`, `NoCache`, `MarkovResidual`, - `UnlimitedContext`, `TurboQuant`, `Apollo`). + `WindowedCheckpoint`, `TurboQuant`, `Apollo`). - [`compute-backend-redesign.md`](crates/larql-inference/docs/specs/compute-backend-redesign.md) — `KvDispatch` sibling trait in `larql-inference` (intent-based per-layer surface); `EngineBackend: ComputeBackend + KvDispatch` @@ -2209,12 +2464,12 @@ beats today's fused `decode_token` path. |----|------|----------|--------|-------| | U1 | KV engine unification — Steps 1–7 | larql-inference, larql-kv, larql-cli | **shipped 2026-05-16** | `KvEngine` trait + EngineInfo + DecodeStageSummary in `larql-inference::kv_engine`; `larql-kv` re-exports. `Standard` + `NoCache` engines added. `larql run` / `larql walk` route through engine dispatch (default `--kv-cache standard` = `Standard { window_size: None }`, bit-parity gated). `--engine SPEC` + `LARQL_KV_ENGINE` env var on run/walk. Server wiring deferred to U7 (server uses fused `decode_token` and would silently downgrade to CPU under sync dispatch). | | U2 | ComputeBackend redesign — Steps 1–4 | larql-inference, larql-compute | **shipped 2026-05-16** | `KvDispatch` trait in `larql-inference` (per-layer intents: cache, attention, engine-specific). `EngineBackend: ComputeBackend + KvDispatch` umbrella with blanket impl. `CpuBackend::KvDispatch` real implementation; `MetalBackend::KvDispatch` CPU-fallback scaffolding. `cpu_engine_backend()` / `default_engine_backend()` factories. 6 new `Capability` flags (`FusedAttentionStep`, `WindowedAttentionStep`, `NativeKvCodec`, `PipelinedBoundaryUpload`, `FusedResidualNorm`, `KvHandleNative`). | -| U3 | ComputeBackend redesign — Step 3c (engine migration) | larql-kv, larql-inference | **shipped 2026-05-16** (partial); follow-up in U8 | All six engines accept `Box` in constructors. `KvDispatch` widened with `Option<&VectorIndex>` on attention intents + new `coarse_prefill` / `coarse_decode_step` (quantization-agnostic, backends inspect index format internally). `StandardEngine` fully migrated: routes Q4K through `coarse_prefill` on `CpuBackend` (which calls production `predict_q4k_prefill` / `predict_q4k_decode_step_direct`). **27.6 tok/s on Gemma 3 4B Q4K, M3 Max, 8 threads — slightly faster than the legacy `larql-cpu` path (24.0 tok/s).** `NoCache` migrated (slow on purpose: O(N²) debug fallback). Others (`MarkovResidual`, `UnlimitedContext`, `TurboQuant`, `Apollo`) still carry their bespoke `prefill_q4k` overrides — they work correctly but run at ~0.4 tok/s through f32-dequant fallback. Migration to fast Q4K kernels via the dispatch trait is **U8** below. Spec: [`kv-dispatch-quantization.md`](crates/larql-inference/docs/specs/kv-dispatch-quantization.md). | -| U4 | AsyncComputeBackend impl — Steps A1–A5 (the trait + foundation) | larql-inference, larql-compute, larql-compute-metal, larql-kv | **A1–A3 + A5 (StandardEngine) shipped 2026-05-16; A4 next** | A1 ✅ trait + handle types in `larql-inference/src/async_compute_backend.rs` (per-handle inner traits, `read(self: Box)` — stable-Rust translation of spec's `Arc` pattern). A2 ✅ `CpuBackend` async impl as degenerate `Ready*` wrapper, 6 bit-parity tests vs sync. A3 ✅ `MetalBackend` scaffold via CPU-delegation, feature-gated; 4 Metal-aware bit-parity tests pass under `--features metal`. A5 ✅ for `StandardEngine`: `with_async_backend` constructor + internal `BackendSlot` enum + async dispatch helpers + 8 new parity tests (`larql-inference`: 1002 lib tests; `larql-kv`: 221 lib tests). A4 next: real `MTLCommandBuffer` deferred dispatch (4–8 weeks). Remaining engines' A5 slices (`MarkovResidual`, `UnlimitedContext`, `TurboQuant`, `NoCache`, `Apollo`) compose on the same pattern (~1–2 weeks each). | +| U3 | ComputeBackend redesign — Step 3c (engine migration) | larql-kv, larql-inference | **shipped 2026-05-16** (partial); follow-up in U8 | All six engines accept `Box` in constructors. `KvDispatch` widened with `Option<&VectorIndex>` on attention intents + new `coarse_prefill` / `coarse_decode_step` (quantization-agnostic, backends inspect index format internally). `StandardEngine` fully migrated: routes Q4K through `coarse_prefill` on `CpuBackend` (which calls production `predict_q4k_prefill` / `predict_q4k_decode_step_direct`). **27.6 tok/s on Gemma 3 4B Q4K, M3 Max, 8 threads — slightly faster than the legacy `larql-cpu` path (24.0 tok/s).** `NoCache` migrated (slow on purpose: O(N²) debug fallback). Others (`MarkovResidual`, `WindowedCheckpoint`, `TurboQuant`, `Apollo`) still carry their bespoke `prefill_q4k` overrides — they work correctly but run at ~0.4 tok/s through f32-dequant fallback. Migration to fast Q4K kernels via the dispatch trait is **U8** below. Spec: [`kv-dispatch-quantization.md`](crates/larql-inference/docs/specs/kv-dispatch-quantization.md). | +| U4 | AsyncComputeBackend impl — Steps A1–A5 (the trait + foundation) | larql-inference, larql-compute, larql-compute-metal, larql-kv | **A1–A3 + A5 (StandardEngine) shipped 2026-05-16; A4 next** | A1 ✅ trait + handle types in `larql-inference/src/async_compute_backend.rs` (per-handle inner traits, `read(self: Box)` — stable-Rust translation of spec's `Arc` pattern). A2 ✅ `CpuBackend` async impl as degenerate `Ready*` wrapper, 6 bit-parity tests vs sync. A3 ✅ `MetalBackend` scaffold via CPU-delegation, feature-gated; 4 Metal-aware bit-parity tests pass under `--features metal`. A5 ✅ for `StandardEngine`: `with_async_backend` constructor + internal `BackendSlot` enum + async dispatch helpers + 8 new parity tests (`larql-inference`: 1002 lib tests; `larql-kv`: 221 lib tests). A4 next: real `MTLCommandBuffer` deferred dispatch (4–8 weeks). Remaining engines' A5 slices (`MarkovResidual`, `WindowedCheckpoint`, `TurboQuant`, `NoCache`, `Apollo`) compose on the same pattern (~1–2 weeks each). | | U5 | AsyncComputeBackend impl — Step A6 (per-engine specialised shaders) | larql-compute, larql-kv | **spec'd, not started** | This is the tok/s payoff. Priority order: `attention_step_windowed` (the `standard:window=N` win), then engine-specific intents in order of impact — `markov-rs` Metal K/V recompute, `apollo` pipelined boundary upload, `turbo-quant` codec kernel. Each shader paired with a real-model bench. Ongoing — months of iterative work. | | U6 | AsyncComputeBackend impl — Step A7 (VulkanBackend) | larql-compute | **spec'd, not started — blocked on U9-U12** | Same trait shape as Metal, different primitives (`VkCommandPool`, semaphores, SPIR-V). Validates the multi-backend story is real, not Metal-shaped. 6–10 weeks **once U9-U12 unblock the engine layer**. Today the substrate trait is drop-in but `larql-inference` still has 30+ `cfg(feature = "metal")` gates and 2 `downcast_ref::()` sites that conflate "Metal" with "GPU pipeline" — landing Vulkan against today's tree would force per-backend cfg explosion across the inference crate. | | U7 | AsyncComputeBackend impl — Step A8 (CudaBackend) + server wiring | larql-compute, larql-server | **spec'd, not started — blocked on U9-U12** | CUDA streams map naturally to the deferred-dispatch shape — designed against it. Server wiring (deferred from `kv-engine-unification.md` §10.6) lands here: `larql-server`'s `handle_stream_generate` switches from direct `generate_streaming` to `generate_with_engine` against an `AsyncComputeBackend`, finally honouring `LARQL_KV_ENGINE` server-side. 6–10 weeks Cuda + 1–2 weeks server. Same engine-layer blockers as U6. | -| U8 | Engine migration — bespoke `prefill_q4k` paths onto dispatch trait | larql-kv, larql-inference | **specced, not started** | `MarkovResidual`, `UnlimitedContext`, `TurboQuant`, `Apollo` each carry an engine-side `prefill_q4k` override that bypasses the dispatch trait's `coarse_prefill` / `coarse_decode_step` intents and uses slower CPU code paths (dequant-to-f32 + f32 sgemv) instead of the production `predict_q4k_*` kernels. Result: ~0.4 tok/s vs `StandardEngine`'s 27.6 tok/s on the same hardware. Each engine has legitimate specialisation (RsStore residuals, per-window K/V checkpoints, WHT+Lloyd-Max codec, boundary residual injection) — the migration keeps that engine-side logic but routes the per-layer matvec through `larql_compute::QuantMatVec::q4k_matvec` instead of dequant-then-f32. Per-engine: ~2-5 days. See [`kv-dispatch-quantization.md`](crates/larql-inference/docs/specs/kv-dispatch-quantization.md) Phase 2. | +| U8 | Engine migration — bespoke `prefill_q4k` paths onto dispatch trait | larql-kv, larql-inference | **specced, not started** | `MarkovResidual`, `WindowedCheckpoint`, `TurboQuant`, `Apollo` each carry an engine-side `prefill_q4k` override that bypasses the dispatch trait's `coarse_prefill` / `coarse_decode_step` intents and uses slower CPU code paths (dequant-to-f32 + f32 sgemv) instead of the production `predict_q4k_*` kernels. Result: ~0.4 tok/s vs `StandardEngine`'s 27.6 tok/s on the same hardware. Each engine has legitimate specialisation (RsStore residuals, per-window K/V checkpoints, WHT+Lloyd-Max codec, boundary residual injection) — the migration keeps that engine-side logic but routes the per-layer matvec through `larql_compute::QuantMatVec::q4k_matvec` instead of dequant-then-f32. Per-engine: ~2-5 days. See [`kv-dispatch-quantization.md`](crates/larql-inference/docs/specs/kv-dispatch-quantization.md) Phase 2. | | U9 | De-Metal the inference-side GPU cfg gates | larql-inference, larql-cli | **not started — compute-refactor branch** | 23 `cfg(all(feature = "metal", target_os = "macos"))` sites in `larql-inference/src` + 8 in `larql-cli/src` use "metal" as a synonym for "GPU pipeline available." Two options: (a) rename `feature = "metal"` → `feature = "gpu"` on `larql-inference` with `larql-compute-metal` as one optional backend inside it, so the same flag turns on Metal today and Vulkan/CUDA tomorrow without per-call-site flag matrix; (b) replace cfg gates with `Capability::FullPipelineQ4` / `Capability::DecodeToken` probes on `&dyn ComputeBackend`. Mechanical search/replace + targeted refactor; ~1-2 days. **Prerequisite for U6/U7.** | | U10 | Move `prepare_ple_inputs` (Per-Layer Embeddings upload) onto a trait method | larql-compute, larql-compute-metal, larql-inference | **not started — compute-refactor branch** | Kills the 2 `downcast_ref::()` sites (`layer_graph/hybrid.rs:78`, `layer_graph/generate/gpu/mod.rs:261`) and the `metal_ple: Option<&MetalBackend>` typed parameter that flows through `generate/gpu/decode_loop.rs:60-67`. Add `fn prepare_ple_inputs(&self, flat: &[f32], num_layers: usize, ple_dim: usize)` to `ComputeBackend` (default no-op) plus `Capability::PerLayerEmbeddings`. Spec at `compute-backend-redesign.md` §6.3 explicitly says "Engines do **not** check `backend.name()` to decide behaviour" — this is the residual gap. ~1 day. **Prerequisite for U6/U7.** | | U11 | Move `take_last_split_timings()` onto a trait method | larql-compute, larql-compute-metal, larql-inference | **not started — compute-refactor branch** | `larql_compute_metal::take_last_split_timings()` is reached directly as a free function from `decode_loop.rs:194-200`. Replace with `fn take_split_timings(&self) -> Option` on a sub-trait (or `ComputeBackend` with a default `None`) so Vulkan/CUDA can expose the same instrumentation hook. Also folds the `ProfileTimings` type down into `larql-compute`. ~0.5 day. **Prerequisite for U6/U7.** | @@ -2275,12 +2530,12 @@ achievability table + `docs/diagnoses/`.)** | C4 | FP4 productisation (exp 26 → product) — native FP4 quantisation tier (`Q4_K → FP4`) | larql-vindex + larql-compute | research only → **V2-validated, greenlit** | Exp 26 + **V2 (2026-05-31, confirmed)**: ≥99.8% FP4-friendly per-feature across Gemma 3 / Granite (no QAT, `down` the tail); predictive E2M1 +0.116 bits/tok vs f32, beating Q4-int. The FP4 codec already exists (`larql-models/src/quant/fp4*.rs`). Add `Quantisation::FP4` variant; CPU-first kernel; Metal twin. ~2× shrink vs Q4_K. See `docs/diagnoses/v2-fp4-generality.md`. | | C5 | mmap'd vindex with lazy disk-resident edges — only resident pages for active edges per token | larql-vindex + larql-inference | not started | Today vindex loads whole layer tensors into RAM. For models bigger than RAM, mmap the vindex file and let the OS page in only the gate-KNN-resolved edges. Pairs with C2 and C3: when only 20% of edges fire, only those pages are read. | | C6 | AMX / AVX-512 / Apple AMX kernels for residual compute | larql-compute (CPU side) | partial — Accelerate BLAS, AMX through it | Current CPU path uses ndarray + Accelerate; promote to direct AMX intrinsics on Apple Silicon, AVX-512 on x86. Compute that *does* happen needs to be as good as it gets, since bandwidth is what's left over. | -| C7 | KV compression as **default** for long context (Apollo / MarkovRS / UnlimitedContext / TurboQuant) | larql-inference | engines reachable on `run`/`walk` (CPU) via `--engine` / `LARQL_KV_ENGINE`; default still `standard` (production K/V cache); GPU performance on opt-in engines requires AsyncComputeBackend (see U-series below) | Unification spec at [`kv-engine-unification.md`](crates/larql-inference/docs/specs/kv-engine-unification.md) — all 7 steps landed. MarkovRS / UnlimitedContext / TurboQuant opt-in via `--engine` (CPU-correct, Metal works via CPU-fallback delegation). Apollo bench-only. Promoting any of these as default for long context requires `AsyncComputeBackend` Step A6 (engine-specific Metal shaders) to land — see U5 below. Server engine wiring also blocked on AsyncComputeBackend (U7); without it the server would silently downgrade Metal decode to CPU. | +| C7 | KV compression as **default** for long context (Apollo / MarkovRS / WindowedCheckpoint / TurboQuant) | larql-inference | engines reachable on `run`/`walk` (CPU) via `--engine` / `LARQL_KV_ENGINE`; default still `standard` (production K/V cache); GPU performance on opt-in engines requires AsyncComputeBackend (see U-series below) | Unification spec at [`kv-engine-unification.md`](crates/larql-inference/docs/specs/kv-engine-unification.md) — all 7 steps landed. MarkovRS / WindowedCheckpoint / TurboQuant opt-in via `--engine` (CPU-correct, Metal works via CPU-fallback delegation). Apollo bench-only. Promoting any of these as default for long context requires `AsyncComputeBackend` Step A6 (engine-specific Metal shaders) to land — see U5 below. Server engine wiring also blocked on AsyncComputeBackend (U7); without it the server would silently downgrade Metal decode to CPU. | | C8 | BR4 (Boundary refs Phase 4 — bounded KV eviction + durability-first capture) | larql-server + larql-inference | not started | See § "P1 — Boundary refs and cold-context storage" below. The CPU track makes BR4 load-bearing because long-context CPU inference can't keep raw KV in RAM. | | C9 | Distributed-load-balancing for "model spans 4 consumer machines" | larql-router + larql-server | shipped (grid + rebalancer) | **DEMOTED to P2 per ADR-019 (2026-05-09)** — substantial production-engineering with no current experiment requiring multi-machine. Single-shard grid (already shipped) sufficient for substrate. Re-promote if a specific experiment needs multi-machine. | | C10 | CPU bench harness — `larql bench --cpu` with per-stage breakdown matched against `llama.cpp -ngl 0` | larql-cli + bench/ | **DISCREPANCY RESOLVED 2026-06-02 — no regression; true gap ~1.6–1.8×.** The 1.50× (05-16) vs 1.93× (05-31) split was **two stacked measurement confounds**, not a real change: (1) **larql path mismatch** — 27.6 was the `StandardEngine` path, 23.6 the legacy `larql bench --cpu` (`predict_kquant_decode_step`) path; a stable ~12% delta (26.4 vs 23.5 today), so comparing one date's StandardEngine against the other's legacy path manufactured a phantom "regression"; (2) **llama.cpp harness artifact** — the 45.5 was an unwarmed/short-n ollama `num_gpu=0` fluke; warmed + n=128 it converges to **42.8–43.0 = llama-bench's 42.99** (both harnesses, both dates agree at ~43). Reconciled like-for-like (M3 Max, t=8, warm): **larql 23.5 legacy / 26.4 StandardEngine vs llama.cpp 43.0 → 1.6–1.8×.** Gap is C12 (both attn AND FFN already use the int8 Q8_K SDOT kernel via `attention_decode_step_native`). **Free wins landed (2026-06-02):** `larql bench --cpu` now also reports the production StandardEngine row; new `--ollama-cpu` forces `num_gpu=0`+`num_thread` so `--ollama` is a true CPU baseline (was silently Metal-GPU). Reconciled artifact `bench/baselines/c10_gemma3-4b_cpu_reconciled.json`. **26B-A4B baseline LANDED 2026-06-10** (`c10_gemma4-26b-a4b_cpu_reconciled.json`): llama.cpp **32.1** vs larql in-process **7.1** default / **9.7** with `LARQL_Q4K_DIRECT_ATTN=1` / loopback 7.3 (t=8, warm, n=128, drift-checked). The 26B gap (4.5×) is **f32-residency byte traffic** (attn 4.15 GB + dense slab 2.14 GB + lm_head 2.95 GB per token vs llama.cpp ~2.1 GB all-quantized; every leg bandwidth-saturated ~62–71 GB/s), NOT the C12 kernel (experts already int8 SDOT, ~8% of bytes). Medium-term tier 62%→70% per the gate rule. Method addition: **pmset AC check + cross-engine drift bracket are now mandatory** — the first session was invalidated by a silent battery drain (llama.cpp itself collapsed 34→1 tok/s at 31% battery; far beyond the 1.5–3× thermal class). | CPU-track baseline-credibility threshold can't be enforced without this. First acceptance test: Gemma 3 4B Q4_K on M3 Max CPU vs quant-matched `llama.cpp -ngl 0`. Then Llama 2 7B + Mistral 7B for cross-arch CPU + the 26B-A4B MoE baseline. Major improvement 2026-05-15→05-16 (2.78× → 1.50×) — see `bench/baselines/cpu/COMPARISON.md` and `DIAGNOSIS-2026-05-16-thread-scaling.md`; reconciliation `bench/baselines/c10_gemma3-4b_cpu_reconciled.json`. | | C11 | Architecture rule enforcement — CI check for "no GPU-only paths in core" | scripts/ + crate boundaries | not started | Static check: anything in `larql-inference` core (not `metal/`, not `cpu/`) must compile and pass tests with Metal feature off. Prevents the dual-track from drifting into Metal-locked code. | -| C12 | Q4K decode kernel — hand-asm aarch64 to close the 1.50× gap to llama.cpp | larql-compute | **v1 asm landed opt-in 2026-06-02 (`LARQL_Q4K_ASM=1`); roofline reframed the work.** Two 2026-06-02 results: (a) **Roofline microbench** (`benches/q4k_q8k_matvec.rs`) shows the kernel is **compute/issue-bound, NOT DRAM-bandwidth-bound** — scalar 9.3 vs NEON 17.7 GiB/s on identical data, size-invariant — which **overturns the `DIAGNOSIS-2026-05-16` "memory-system-level" conclusion** and confirms hand-asm scheduling is a real lever (17.7 GiB/s ↔ ~33 cyc/super-block, exactly as specced). (b) **`q4k_q8k_matvec_asm`** (whole super-block dot in one `asm!` block, 8 scales as vector lanes killing the 8 scalar `ldrb`) — **bit-exact** (`q8k_matvec_asm_matches_scalar_bit_exact`), **+3.7–4.9% isolated**, ~+1–2% e2e (diluted: opt-in covers `matvec_into` callers — attention Q/K/V/O + `down` — but NOT the fused `gate_up`). **Finding: latency-hiding has low headroom** — a 4-accumulator variant showed no reliable gain (the inlined row loop lets the OoO core already overlap super-blocks), so **the two-super-block interleave is deprioritized**; the real lever to reach ~28 GiB/s is **instruction-count reduction** (perf-counter-guided, llama.cpp-style vectorized scale path) + **asm-ifying `gate_up`** (lifts the e2e ceiling). See spec §"2026-06-02 roofline measurement". | Per-core gap is **1.73× constant across thread counts** (5.7 vs 9.88 tok/s single-threaded on M3 Max). Same algorithm (Q4K × Q8K with NEON SDOT), same `vdotq_s32` instructions — llama.cpp uses hand-written inline aarch64 asm with two-super-block interleaving + explicit prefetch hints, we use Rust intrinsics lowered by LLVM. Effective bandwidth: ~63 GB/s vs ~95 GB/s. **Per-stage profile (`LARQL_INSTRUMENT_UNLIMITED=1` on Gemma 3 4B 8-thread, 2026-05-16): FFN 26.0 ms (74%) + Attention 9.3-11.0 ms (26%, grows with ctx) + Embed ~0 ms = 35-37 ms/step.** FFN matvec on gate/up/down (4608 × 9216) is the dominant target; attention matvec is the same kernel on smaller matrices. The 38 tok/s asymptote (FFN-alone) sets the floor any engine can reach on the current kernel — Standard and UnlimitedContext both hit 26.6 tok/s on Gemma 3 4B Q4K CPU (8-thread, 40-token prompt, 64 decode tokens) because both route through the same `attention_decode_step_native` + `ffn_decode_step_native` hot paths. Phases: (1) hand-asm Q4K matvec on the FFN tile shapes (gate/up/down) — closes ~95% of the gap, 1-2 weeks; (2) pre-formatted block layout — 1.1-1.2× on top, 3-5 days; (3) Q6K kernel for `ffn_down` — 1.05×, 2-3 days; (4) reduce rayon launch overhead — 1.04×, 2-3 days. Acceptance: ≥9.5 tok/s single-core, ≥39 tok/s 8-thread on Gemma 3 4B Q4K. Spec: [`crates/larql-compute/docs/q4k-decode-kernel.md`](crates/larql-compute/docs/q4k-decode-kernel.md). Per-stage measurement protocol: see "C12 per-stage measurement" below. | +| C12 | Q4K decode kernel — hand-asm aarch64 to close the 1.50× gap to llama.cpp | larql-compute | **v1 asm landed opt-in 2026-06-02 (`LARQL_Q4K_ASM=1`); roofline reframed the work.** Two 2026-06-02 results: (a) **Roofline microbench** (`benches/q4k_q8k_matvec.rs`) shows the kernel is **compute/issue-bound, NOT DRAM-bandwidth-bound** — scalar 9.3 vs NEON 17.7 GiB/s on identical data, size-invariant — which **overturns the `DIAGNOSIS-2026-05-16` "memory-system-level" conclusion** and confirms hand-asm scheduling is a real lever (17.7 GiB/s ↔ ~33 cyc/super-block, exactly as specced). (b) **`q4k_q8k_matvec_asm`** (whole super-block dot in one `asm!` block, 8 scales as vector lanes killing the 8 scalar `ldrb`) — **bit-exact** (`q8k_matvec_asm_matches_scalar_bit_exact`), **+3.7–4.9% isolated**, ~+1–2% e2e (diluted: opt-in covers `matvec_into` callers — attention Q/K/V/O + `down` — but NOT the fused `gate_up`). **Finding: latency-hiding has low headroom** — a 4-accumulator variant showed no reliable gain (the inlined row loop lets the OoO core already overlap super-blocks), so **the two-super-block interleave is deprioritized**; the real lever to reach ~28 GiB/s is **instruction-count reduction** (perf-counter-guided, llama.cpp-style vectorized scale path) + **asm-ifying `gate_up`** (lifts the e2e ceiling). See spec §"2026-06-02 roofline measurement". | Per-core gap is **1.73× constant across thread counts** (5.7 vs 9.88 tok/s single-threaded on M3 Max). Same algorithm (Q4K × Q8K with NEON SDOT), same `vdotq_s32` instructions — llama.cpp uses hand-written inline aarch64 asm with two-super-block interleaving + explicit prefetch hints, we use Rust intrinsics lowered by LLVM. Effective bandwidth: ~63 GB/s vs ~95 GB/s. **Per-stage profile (`LARQL_INSTRUMENT_UNLIMITED=1` on Gemma 3 4B 8-thread, 2026-05-16): FFN 26.0 ms (74%) + Attention 9.3-11.0 ms (26%, grows with ctx) + Embed ~0 ms = 35-37 ms/step.** FFN matvec on gate/up/down (4608 × 9216) is the dominant target; attention matvec is the same kernel on smaller matrices. The 38 tok/s asymptote (FFN-alone) sets the floor any engine can reach on the current kernel — Standard and WindowedCheckpoint both hit 26.6 tok/s on Gemma 3 4B Q4K CPU (8-thread, 40-token prompt, 64 decode tokens) because both route through the same `attention_decode_step_native` + `ffn_decode_step_native` hot paths. Phases: (1) hand-asm Q4K matvec on the FFN tile shapes (gate/up/down) — closes ~95% of the gap, 1-2 weeks; (2) pre-formatted block layout — 1.1-1.2× on top, 3-5 days; (3) Q6K kernel for `ffn_down` — 1.05×, 2-3 days; (4) reduce rayon launch overhead — 1.04×, 2-3 days. Acceptance: ≥9.5 tok/s single-core, ≥39 tok/s 8-thread on Gemma 3 4B Q4K. Spec: [`crates/larql-compute/docs/q4k-decode-kernel.md`](crates/larql-compute/docs/q4k-decode-kernel.md). Per-stage measurement protocol: see "C12 per-stage measurement" below. | **Implementation order** (post ADR-019): C10 → C1 → C2 → C7 → C12 → C3 → C4 → C5 → C6 → C8 → C11. @@ -2303,7 +2558,7 @@ long-context. C11 prevents architectural drift. Two instruments measure the kernel-bound nature of CPU decode and let you isolate which sub-kernel the asm should target first: -- `LARQL_INSTRUMENT_UNLIMITED=1` — prints `embed / attention / ffn` per `extend_q4k` call from `larql_kv::engines::unlimited_context::rs_extend_from_checkpoint_q4k`. Captures the per-token, per-layer-aggregated breakdown. Source: `crates/larql-kv/src/engines/unlimited_context/extend.rs`. +- `LARQL_INSTRUMENT_UNLIMITED=1` — prints `embed / attention / ffn` per `extend_q4k` call from `larql_kv::engines::windowed_checkpoint::rs_extend_from_checkpoint_q4k`. Captures the per-token, per-layer-aggregated breakdown. Source: `crates/larql-kv/src/engines/windowed_checkpoint/extend.rs`. - `LARQL_INSTRUMENT_MARKOV=1` — same shape for `markov-residual`, kept for cross-engine sanity that both substrate paths agree. Source: `crates/larql-kv/src/engines/markov_residual/q4k.rs`. Reproducer (Gemma 3 4B Q4K, M3 Max, default 8 threads): @@ -2312,7 +2567,7 @@ Reproducer (Gemma 3 4B Q4K, M3 Max, default 8 threads): cargo build --release -p larql-cli LARQL_INSTRUMENT_UNLIMITED=1 ./target/release/larql bench \ ~/.cache/larql/local/gemma3-4b-q4k-v2.vindex \ - --backends cpu --engine unlimited-context -n 32 + --backends cpu --engine windowed-checkpoint -n 32 ``` Recorded baseline (2026-05-16, 8-thread, ~70-token ctx after warmup): @@ -2541,7 +2796,7 @@ Both specs live at `crates/larql-inference/docs/specs/`. - **SQ1 (Markov)**: contract is sound, reference impl already works, but it's engineering not research — and the open trait-shape question means migrating Markov first risks forcing - `UnlimitedContextEngine`/`ApolloEngine` into a shape that doesn't fit. + `WindowedCheckpointEngine`/`ApolloEngine` into a shape that doesn't fit. Designing the trait once across all three engines (or at least resolving sibling-vs-trait before SQ1 lands) is cheaper than migrating one and refactoring twice. V1/V2 also produce the measurement diff --git a/ROADMAP_STATUS.md b/ROADMAP_STATUS.md index 373cdad9e..836dcd57f 100644 --- a/ROADMAP_STATUS.md +++ b/ROADMAP_STATUS.md @@ -4,18 +4,67 @@ Canonical rollup for the next execution slice. Keep the detailed design in `ROADMAP.md` and crate-local roadmaps; use this file to answer "what is active now?" without rereading every crate document. -Last updated: 2026-08-01 +Last updated: 2026-08-02 + +**Active slice: VINDEX3 execution binding** +([`ROADMAP.md` § VINDEX3](ROADMAP.md), spec +[`crates/larql-vindex/docs/vindex3-format-spec.md`](crates/larql-vindex/docs/vindex3-format-spec.md), +programme [`docs/vindex3-experiments.md`](docs/vindex3-experiments.md)) — driving +a real Gemma layer through a VINDEX3-bound operation until it is bit-identical +to production, then propagating to the first VINDEX3-generated token. Rungs 0, +0.5 and 1 are closed (`f13bf385`, `dd2017db`, `f5dd256e`); **rung 2 — binding the +production Q4_K × Q8_K expert kernel — is next**, starting with Q8_K activation +identity before any expert runs. + +**The container half now exists too.** Until 2026-08-02 every VINDEX3 parity +result bound its operands out of a VINDEX2 file, so what was proven was the +executor and not the format — nothing could write a VINDEX3 container, and +`ContainerGeneration::V3` appeared only in a detection test built from a JSON +string. Conformance fixture A now round-trips end to end (write → detect → +open → validate → bind → execute) **bit-identically**, fused and decomposed +storage agree under one programme id, `verify` reports structural defects with +`{layer, entry, role}` coordinates, and `show`/`verify` dispatch on generation +without normalising either into the other. That closes the rows fixture A can +carry on V2-0/V2-1; profile authority, variant-selection refusal, fixtures B–D +and WALK/DESCRIBE parity remain open, so **`extract` still writes VINDEX2** and +the ABI is not frozen. Honest ceiling: *executable VINDEX3 container proven on +fixture A* — not "a real model runs from VINDEX3". + +**Previous slice: the DEC funnel** ([`docs/dec-funnel.md`](docs/dec-funnel.md) +v0.5) — decoupled attention/weights serving, DEC-0 … DEC-7 plus the G-ladder +(CUDA attention client), C-ladder (x86 kernels) and M-ladder (MTP/speculative +decode). Not closed, and not superseded in substance: VINDEX3 is the storage and +capability layer the DEC funnel's expert serving will bind to, so the two meet at +`SelectedExpertNotResident` — the seam where a valid route with a missing local +operand becomes a fetch rather than a defect. -**Active slice: the DEC funnel** ([`docs/dec-funnel.md`](docs/dec-funnel.md) v0.5) -— decoupled attention/weights serving, DEC-0 … DEC-7 plus the G-ladder (CUDA -attention client), C-ladder (x86 kernels) and M-ladder (MTP/speculative decode). The V1–V4 aim-validation gate that previously governed this file is **closed** (V1 falsified dense + MoE, V2 confirmed, V3 locality poor); its table is retained -below as resolved history. Registry programme `dec` on the experiments server is -the system of record for stage results. +below as resolved history. Registry programmes `dec` and `vindex2` on the +experiments server are the system of record for stage results. ## Recently shipped (delta since last update) +- **VINDEX3 routes a real Gemma layer bit-identically to production (2026-08-02).** + Three milestones, committed separately so a later disagreement cannot require + bisecting a combined patch. `f13bf385` — reference MoE execution: fixture A + matches an independent oracle below 1e-6, fused and decomposed storage + agreeing at every checkpoint, plus a residency probe showing the bound plan + predicts its resident page set **exactly** (200 predicted, 200 resident, zero + overshoot, 1.63% of a 192 MiB layer per token). `dd2017db` — real Gemma + routing parity over real VINDEX2 bytes, which forced three corrections that + were unreachable synthetically: Gemma's router input is *not* its expert input + (`MoeInputs::split`), a shard is a legitimate subset of the routing universe + (validate was refusing every expert-server slice), and Q4_K's padded + intermediate axis needs `physical shape ≠ semantic operand shape` + (`ComponentView::Slice`, 704 logical over 768 stored). `f5dd256e` — the + production router kernel bound: **all five ladder stages bit-identical** on + real weights and a real activation, with `BoundExpertScaling` making + policy-without-operand unrepresentable and `OperandUnsuitability` typing the + four kernel-refusal causes. Two unrelated defects fixed en route: `larql + verify` rendered findings in `HashMap` order and disagreed with itself between + runs, and the separate-tensor MoE extractor wrote no expert store. + - **Sparse-FFN thesis closed on three of four routes; the ANN/HNSW programme closes with it (2026-08-01).** A vindex+WalkFFN review exit finding (HNSW's level-0 graph fragments beyond ~64 nodes, recall@10 → 0.16 at n=200) opened the question of whether repairing the graph was worth it. **R4 says no, and for a reason that closes more than HNSW.** [`walk-ffn-r4-zeroout.md`](docs/diagnoses/walk-ffn-r4-zeroout.md): supplying perfect routes for **free** — kernel-matched, parity ✓, paired/interleaved with a drift+dispersion sentinel, replicated twice on AC — the sparse path still runs at **0.837× dense at the accuracy-viable K=4096** and only wins at K=2048, which fails top-1. **No overlap between speed-viable and accuracy-viable.** The mechanism is the kernel, not the router: with routing free it captures only **23–40%** of its row-count reduction, and the capture fraction *rises* as fewer rows are dropped (23→34→40% against ceilings 5.00/2.50/1.67×) — the signature of fixed per-row overhead. Merely *tying* dense needs ~18.6% execution improvement with routing already free. So a cheaper approximation of gate-top-K was never the missing piece; exact gate-top-K already sits on the wrong frontier. **Scope, deliberately narrow:** this refutes a *compute*-side claim and does **not** close DEC-8.1 (read-side, different denominator, far milder retention) — what it kills there is the *mechanism*, since runtime scattered gather costs more than it saves once rows are in hand. **Then the other axis of the factorisation was tested** ([`moe-latent-axis-sparsity.md`](docs/diagnoses/moe-latent-axis-sparsity.md)): masking the *shared expert input* rather than the row population, where one channel decision removes a column from every active expert's gate AND up. The signal is real — magnitude vs random separates **9.7% vs 555%** bits/token at 50% retention — but **the headline was an artifact of larql's known-divergent OLMoE forward**, which called r=0.5 free where the HF reference says +9.66%. At the project's ≤0.5% shannon gate the viable retention is r≈0.875, **ideal ceiling 1.040×**. Static channel sets refuted per-layer (96.7% of layer-channel pairs token-dependent; a P≥0.8 static core covers 6.6% of the budget); blocked latent sparsity refuted per-layer against exact induced loss and a **random-permutation control** that turned an apparent 40% win into a null. Surviving: dynamic per-channel via an input-major packed kernel, unbuilt, re-priced from "1.18× for free" to "≤1.04× at the gate". **New standing rule R11** in [`dec-funnel.md`](docs/dec-funnel.md): *a structural reduction is a claim about a kernel, not about a matrix* — with the corollary that realisability is a property of **layout**, not of the mask. Instruments `scripts/moe_latent_reference_olmoe.py`, `scripts/moe_latent_block_partition.py`; artifacts under `bench/aim-validation/`. - **Extraction tensor-coverage audit — every source tensor now classified, and the third bucket is loud (2026-07-31).** §4.6's work-item 2, built and wired. `extract::coverage` classifies every checkpoint tensor as **recognised** (an architecture accessor names it), **dropped by a named rule**, or **unrecognised**; the `tensor_audit` stage runs *first* in `build_vindex_streaming`, so a checkpoint carrying tensors nothing can address fails in seconds rather than after a multi-minute extraction. Reports always, fatal under `LARQL_EXTRACT_STRICT=1`, which is now set in the `larql-vindex` CI workflow (verified first: all 132 streaming-extraction tests pass under it, so it can't break the build on day one). **The case for it was five silent drops in one week**, none caught automatically: 5 of 11 attention tensors, 3 of 8 MLP tensors, the `gate_walk` trait default silently returning `None`, `moe_intermediate_size()` defaulting to 0, and LayerNorm `β`. Validated on ten real checkpoints — Qwen3-30B-A3B (18,867 tensors), OLMoE (3,219), gpt-oss-20b, Gemma 3 4B (439 SigLIP tensors correctly classified `non-text-tower`) — all clean. **GPT-2 from HF safetensors: 1 of 160 recognised**; `gpt2.rs` matches the trait defaults only *after* the GGUF→HF normalisation, so a raw HF checkpoint is unaddressable (it fails late at embeddings rather than silently, but 159 tensors are unreachable). Design notes: it measures **naming, not consumption** — the necessary condition, and where all five drops actually lived; the hand-maintained accessor enumeration fails **noisy, never quiet** (a new accessor not wired in reports its tensors as unrecognised) and its pin test has already fired twice for real; and drop rules are documented as a decision, not a mute button — "if you don't know what a tensor is, it belongs in `unrecognised`". **Also fixed, found by the audit: LayerNorm `β` was dropped for GPT-2 and StarCoder2.** No accessor named a norm bias, so extraction never wrote one and `build_pipeline_layers` hardcoded `input_norm_bias: None` — while the Metal `layer_norm` shader implemented `+ bias` and always selected its no-bias variant. The CPU dense path got away with it by mangling the weight key, which is why raw-safetensors inference was right and every vindex-backed path silently lost the shift term of `γ·x̂ + β`. Now declared (three additive accessors, derived once from the weight key and gated on `NormType` so RMSNorm families correctly claim nothing), extracted, and resolved in the pipeline; the honest status is "the tensor flows end to end", not "the output is verified". Both weight writers also stopped hardcoding `"norm.weight"` and now use `arch.final_norm_key()` — they agreed with every reader only by coincidence. Coverage: larql-vindex 92.42%, larql-models 90.41%, larql-compute 95.76%, larql-compute-metal 96.78% all pass their per-file policy; `residual_diff/capture.rs` went 34%→50% (first-ever test of `cpu_prefill`) and **cannot reach 90% in Linux CI by construction** — ~250 of 408 lines need a Metal device, which is why it sits outside `include_globs`. Follow-ups tracked in [`ROADMAP.md`](ROADMAP.md) §"Extraction tensor-coverage audit + silent-drop follow-ups". - **GB unblocked — and unblocking it found six defects in the GPT-OSS MoE forward pass, plus a hole in the ladder's own premise (2026-07-30).** §4.6.8 recorded GB as blocked because `larql shannon score` hardcoded a dense-only `WeightFfn` and so could not score *any* MoE model — the entire class R1/R2/R3 are built on. Routing the scorer by architecture turned out to be the smaller half of the job. **larql's GPT-OSS MoE forward diverged from the reference in six independent ways**, and the biggest is a layout error: the fused `gate_up` tensor is **interleaved** (gate on the even output rows, up on the odd — the reference dequantises, transposes, then slices `[..., 0::2]`/`[..., 1::2]`), while larql took the leading half as gate. **90.29 % of elements differ**, and the signature is unmistakable on the real checkpoint: the correct split separates two distinct distributions (gate std 0.0287/absmax 0.250 vs up 0.0449/0.500) while the wrong one yields two halves with *matching* statistics, because both are the same 50/50 mixture — with the bias as an independent witness (reference gate/up means −0.464/−0.898 separate cleanly; larql's −0.679/−0.684 both sat at the pooled mean with variance inflated 2×). `gpt_oss.rs`'s header asserted "first half = gate" in prose three lines above the sinks claim §4.6 falsified: two false load-bearing statements in one comment block, neither tested. The other five: the expert MLP is **not SwiGLU** (`(up.clamp(±7) + 1) · [g.clamp(max 7) · σ(1.702·g)]` — larql computed `silu(g)·u`, missing α, the `+1` and both clamps, with `swiglu_limit: 7.0` sitting unread in `config.json`); the router's normalise/select **order was inverted** (softmax-over-all-then-select attenuates the whole expert branch, where the reference softmaxes over the selected logits); and **3 of 8 per-layer MLP tensors were silently dropped** (`router.bias`, `gate_up_proj_bias`, `down_proj_bias` — the last two appeared *nowhere* in the workspace), the same mechanism as §4.6.1's 5-of-11 attention tensors. A seventh: `moe_intermediate_size()` defaults to 0 and GPT-OSS never overrode it. **Why the tests passed:** both `split_gate_up_experts` tests used `out_features = 2`, where row 0 is simultaneously "the first half" and "the even rows" — a fixture too small to distinguish the candidate behaviours is not a weak test, it is an absent one (second instance of this shape, after §4.5's plausibility-only timestamp tests). **Landed:** the de-interleaving convention owned once in `mxfp4::deinterleave_fused_half`; additive default-impl'd `ExpertGatePolicy`/`ExpertRoutingPolicy`/bias accessors on `ModelArchitecture` (G0 holds, no signature changed); `swiglu_limit` and `norm_topk_prob` now *read* rather than assumed; and `ExpertWeightFfn` — a per-expert f32 reference-tier MoE backend, fully architecture-driven, 98.1–99.0 % line coverage. **Verified** by a differential test pinning all 48 values of a 3-token MoE block against an independent transcription of `modeling_gpt_oss.py` at **< 1e-5** (no fixture file — both sides draw from the same LCG; generator at `scripts/moe_reference_gpt_oss.py`), with a control proving the old split diverges > 1e-3. **GB now runs: GPT-OSS-20B scores 0.708 bits/char / 3.221 bits/token, 86 GB peak RSS, 51 s.** **It is unblocked, not green,** and the two rungs fail differently. OLMoE sits at 1.901 vs the HF reference's 0.390 — the MoE block is pinned at < 1e-5 so the residual is elsewhere in a forward pass that had never been numerically checked before today. For GPT-OSS **all three reference engines failed**: HF f32 crashes (`float != c10::BFloat16` in transformers' own `_grouped_mm_fallback`, because `convert_moe_packed_tensors` dequantises to bf16 whatever dtype you ask for), HF bf16 runs but reports an implausible 8.03 bits/token for a 20B, and MLX crashes in `gather_qmm`. larql's number is the only plausible one in that table, which is suggestive and **is not verification**. **That is a hole in the ladder's premise worth more than the number:** §1 justified the detour on the ancestor "fitting on the Mac", which is two claims — the weights fit (86 GB, they do) *and the reference implementation actually executes* (for gpt-oss-20b, it does not). **Action for R2/P1: verify Kimi Linear's FLA/vLLM reference runs on Apple Silicon as the first task of the rung**, before any adapter work, because both are CUDA-first and the ladder's cost model assumes otherwise. Also caught: the new router hardcoded GPT-OSS's order and GB flagged it on OLMoE within one run (2.677 → 1.901) — finding 5 committed a second time, by me, an hour after documenting it, in a file whose header *already named the hazard*. A comment naming a hazard does not protect against it; only a type or a test does. Write-up [`docs/k3-funnel.md`](docs/k3-funnel.md) §4.7. @@ -33,7 +82,7 @@ the system of record for stage results. - **Spin-barrier pool made safe for default-on + panic-fixed (2026-06-13).** Added a spin→yield→park backoff (spin `SPIN_HOT`≈the proven pure-spin window so active decode is unchanged → `yield_now` cooperative bridge once a wait outlives a token = idle/starved → `park` deep-idle, ~0 CPU; dispatcher unparks on dispatch) so the pool doesn't peg cores when the decode loop is idle — what makes on-by-default safe on a shared box. **Found + fixed a real panic-safety bug:** a panicking chunk body killed the worker thread → `completed` never hit `num_chunks` → the dispatcher spin-waited the barrier FOREVER (hung two test procs at 200% for 25 min; also the nondeterministic "flake" — dispatcher-chunk panic propagated/passed, worker-chunk panic hung). Fix: `catch_unwind` per chunk, always count, re-raise the first payload on the dispatcher (rayon semantics) + a regression test. 705 compute + 1220 inference + 756 kv green; spin_pool 8 tests incl. panic + concurrency hammered 25×. - **CPU MoE decode CAUGHT llama.cpp — spin-barrier thread pool, 26B 27→35 tok/s (+28%), now ~9% AHEAD of llama.cpp's 32.1 (2026-06-13).** A `/usr/bin/sample` profile of live 26B decode pinned the post-residency frontier to **rayon fork-join overhead**, not the kernels: the decode driver runs *outside* the global rayon pool, so each of ~211 parallel sections/token took the cold path (`in_worker_cold → LockLatch::wait_and_reset → __psynch_cvwait`) and workers slept between sections — ~40% of thread-time in wait states. **Built `larql_compute::cpu::spin_pool`** (llama.cpp-style persistent spin-barrier pool: workers spin on an epoch counter, park only after a long idle; static partition makes `completed==num_chunks` a sound barrier (strided at the time; **contiguous blocks + a performance-core cap since 2026-07-29** — the 16-participant case collapsed 3.5×, see the top entry, and the +28% below was only ever measured at 8 threads); dispatch mutex + reentrancy guard for `--concurrent`/tests; 7 unit tests incl. a concurrent-dispatcher one that caught a real cursor-reset bug). **Centralized** four byte-identical `par_chunks_mut` matvec twins (larql-compute `cached.rs`, larql-inference `cached.rs`, lm_head ×2 in `dense.rs` — the long-standing "consolidation hazard") into one `q4k_q8k_matvec_parallel`, and routed every hot decode section (attn int8 Q/K/V/O, GQA, dense FFN gate/up/down, geglu, expert fold, lm_head q4+f32) through the pool. Parity-validated both ways (704 compute + 1220 inference + 756 kv green flags-off AND flags-on, incl. the `predict_kquant` parity oracles; clippy clean) — **now default-on (see the env-cleanup entry above); `LARQL_SPIN_POOL=0` opts out**. **Profile after: rayon eliminated** (`in_worker_cold` 2682→0, `join_context` 10300→0, `wait_until_cold` 4463→9). **Measured (M3 Max, t=8, warm, tight A/B bracket, flags inline):** 26B short-ctx OFF ~26.9 → ON **33–35**; n=256 OFF ~27.4 → ON **~34.9 (+28%, ON 35.0/34.8 vs OFF 27.3/27.4 dead-consistent)** — vs llama.cpp recorded **32.1** ⇒ larql ~9% ahead (machine validated: 4B llama.cpp 44 vs recorded 43). **Caveat:** the pool *spins* during active decode (that is the win on a dedicated box); the spin→yield→park backoff (shipped, see entry above) cedes cores only when the decode loop goes idle, which is what makes default-on safe on a shared machine. **Methodology bug (burned a chunk of the session):** `env $FLAGS …` does **not** word-split in this shell (zsh semantics) → only the first flag was set; spell decode flags **inline** before the binary, never via a `$VAR`. Crate detail: [`crates/larql-kv/ROADMAP.md`](crates/larql-kv/ROADMAP.md) §"Spin-barrier pool". - **Bottleneck pass #2 (code-level, machine contended) — GQA scratch fix landed; measurement queue parked (2026-06-13).** With the box owned by a sibling session (battery also low), did the allocation-churn audit instead of timing: the expert path is already TLS-pooled (`ExpertScratch`), but the GQA head-parallelization had introduced a per-head scores `vec!` (≈480 allocs+zeroings/token, growing with ctx) — replaced with `for_each_init` per-worker scratch (rayon workers are long-lived, so it amortises across calls too). 697+1220+756 green. **Measurement queue for the next quiet-machine session:** (1) fresh `/usr/bin/sample` + stage split on the current build (the last profile predates KV append-in-place + serial cuts — the sink distribution has shifted); (2) boundary-kv ratio re-measure (its 0.80× predates its resident-forwarding fix); (3) clean absolute engine matrix + llama-bench brackets; (4) remaining known sinks by size: `cpu_moe_forward` glue (router f32 matvec ~43 MB/tok + norms + route), allocator churn from per-projection out-Vecs and ndarray temporaries (arena-per-step is the structural fix if a fresh sample still shows `madvise`), long-ctx GQA growth (real work; flash-attention-class restructure is the eventual answer). **✅ RESOLVED 2026-06-13:** (1) the fresh `/usr/bin/sample` ran → the dominant sink was **rayon fork-join overhead** (driver outside the pool), not the listed candidates → fixed by the spin-barrier pool (see top entries). (2)+(3) clean absolute engine matrix measured (standard 30.5 / unlimited 31.8 / boundary-kv **0.89×** post-fix / turbo 9.4 / markov 7.8 / codec 7.3); spin A/B 35.7 on / 28.1 off = +27%. (4) the named glue/allocator/GQA sinks are now the *next* tier below the (closed) fork-join one. 26B llama.cpp same-session anchor still owed (ollama wouldn't run the HF GGUF on CPU); recorded 32.1 + 4B-anchor stand. -- **Engine structural gap CLOSED — every KV engine now plugs into the CPU fast path (2026-06-13).** The `KvEngine::decode_step_resident` trait default silently dropped the index (`let _ = index`), so the 06-11/12 fast-path arc (q4k/int8 attention + asm + append-in-place) reached only `StandardEngine`. **Built:** single-source dispatcher `run_attention_block_decode_step_auto` (same per-layer q4k-vs-f32 choice as `CpuBackend::attention_step`, for `SharedKV`-owning walk loops; flag moved to `attention::decode` as the one source); `markov-rs`/`markov-rs-codec`/`turbo-quant`/`unlimited-context`/`boundary_per_layer` override `decode_step_resident` and thread the vindex down their walks; **`boundary-kv` forwards both resident methods to its inner StandardEngine** (was silently f32). `no_cache`/`apollo` keep the default by design. **Regression pin:** `engines::resident_identity_tests` — 7 concrete specs, resident ≡ plain bit-identical flags-off, coverage count can't shrink. **Measured (within-run ratios vs standard; absolutes pending quiet machine — sibling session load):** turbo 0.64×→0.85×, unlimited 0.76×→**1.07×**; markov/codec/boundary-per-layer flat = their own recompute/codec machinery (the feature), not attention. Prefill stays f32 gemm everywhere (task-#16 prefill falsification). 697+1220+756 tests green, clippy clean. Crate detail: [`crates/larql-kv/ROADMAP.md`](crates/larql-kv/ROADMAP.md) §"CPU resident fast-path". +- **Engine structural gap CLOSED — every KV engine now plugs into the CPU fast path (2026-06-13).** The `KvEngine::decode_step_resident` trait default silently dropped the index (`let _ = index`), so the 06-11/12 fast-path arc (q4k/int8 attention + asm + append-in-place) reached only `StandardEngine`. **Built:** single-source dispatcher `run_attention_block_decode_step_auto` (same per-layer q4k-vs-f32 choice as `CpuBackend::attention_step`, for `SharedKV`-owning walk loops; flag moved to `attention::decode` as the one source); `markov-rs`/`markov-rs-codec`/`turbo-quant`/`windowed-checkpoint`/`boundary_per_layer` override `decode_step_resident` and thread the vindex down their walks; **`boundary-kv` forwards both resident methods to its inner StandardEngine** (was silently f32). `no_cache`/`apollo` keep the default by design. **Regression pin:** `engines::resident_identity_tests` — 7 concrete specs, resident ≡ plain bit-identical flags-off, coverage count can't shrink. **Measured (within-run ratios vs standard; absolutes pending quiet machine — sibling session load):** turbo 0.64×→0.85×, unlimited 0.76×→**1.07×**; markov/codec/boundary-per-layer flat = their own recompute/codec machinery (the feature), not attention. Prefill stays f32 gemm everywhere (task-#16 prefill falsification). 697+1220+756 tests green, clippy clean. Crate detail: [`crates/larql-kv/ROADMAP.md`](crates/larql-kv/ROADMAP.md) §"CPU resident fast-path". - **KV-engine review of the append-in-place handle — all engines green; two failure-path fixes landed (2026-06-13).** `EngineBackend: ComputeBackend + KvDispatch`, so every handle-holding engine rides the rewritten `CpuKvHandle`. Audit found two semantic edges the happy-path suites couldn't catch, both fixed: (1) q4k attend-failure after the in-place append now **pops the appended row and falls back to f32** (the old monolithic form's semantics — protects engine-level fallbacks like boundary_per_layer's dense-walk that reuse the handle); (2) the f32 path's prior is **copied not moved** so a backend failure leaves the handle intact. Empirical matrix, flags ON: **26B** all six MoE-capable engines run clean (standard 14.3 cold / boundary-kv 11.5 / unlimited 10.9 / turbo 9.2 / codec 7.9 / markov 7.2 — C1-class ordering); **4B** standard **28.9 tok/s** (dense also gains: gap to llama.cpp 43 now 1.49×), `standard:window=64` exercises `clip_kv` on the new buffers correctly, no-cache/markov/codec/turbo/unlimited all run. 697+755 suites green (incl. the cross-engine parity oracles). - **KV append-in-place — 23.5→27.9 tok/s short-ctx, 16.0→24.8 long-ctx; gap to llama.cpp ~1.15×/1.29× (2026-06-12).** The old `CpuKvHandle` attention step did a full-cache **clone** + `zeros` + four assigns per layer per step (~190 MB/token of churn at ctx 130, growing linearly). Rewritten: growable row-major Vec buffers (amortised O(kv_dim) `append_row`), q4k-direct step split into **project / append-in-place / attend-over-views** (no concat, no clone; legacy owned-concat wrapper kept for engine walk-loops), f32 fallback moves state instead of cloning, `append_kv`/`clip_kv` lose their O(ctx) rebuilds too, `gqa_attention_decode_step` generic over views. **E2E:** n=128 **27.9 tok/s** (35.8 ms; bracket 33.6/31.0 → ~1.15×); n=512 **24.8** vs llama.cpp tg512 32.0 (1.29× — remaining ctx-growth is GQA compute, real work). Text sanity: full flag stack on 4B → "The capital of France is **Paris**." **Cumulative arc: 7.6 → 27.9 tok/s (3.7×); gap 4.5× → ~1.15×.** 692+1213+755 tests green incl. kv-engine parity oracles. Side-finding: `ave_direct_step_parity` garbage CONFIRMED PRE-EXISTING at clean 6659fd6c (worktree repro) — legacy direct-path twin, not this work; AVE session owns it. Artifact §"update_2026_06_12_kv_append_in_place". - **DRAM-idle hunt: sinks NAMED by live-process sampling; first three serial cuts land 21.7→23.5 tok/s (2026-06-12).** MT shape sweep first **acquitted** the suspects (rayon-chunked matvecs sustain 72–113 GiB/s at every production shape; experts-granularity arm 104) → `/usr/bin/sample` of a live decode attributed the real gap: **~75% of thread-samples parked in wait states** — workers sleep while serial main-thread sections run. Named: attention non-projection ~20% of wall (KV-concat `zeros`+`bzero`+`memmove`, per-head `cblas_sgemv` GQA + scalar libm-`exp` softmax, norms/RoPE), `cpu_moe_forward` glue ~12%, dense-slab scalar gelu + serial requantise ~8%, **lm_head argmax epilogue 4.6%** (serial softmax+top-k over 262K logits), `madvise` churn ~3%. **Built (parity-safe):** `q4_lm_head_argmax` (argmax over raw logits — scaling/softcap/temp are monotone → identical selection; no softmax, no 3 MB temporaries), GQA rayon-parallel over heads (math unchanged), dense-slab activation rayon-chunked (same libm per element). **E2E n=128: 23.5 tok/s under a degraded bracket (27.1)** — same-state gap 1.15×, vs clean-bracket llama.cpp 33 ≈ 1.35×. **n=512: 16.0 tok/s — the O(ctx) KV-concat realloc+copy (~20 ms/token by step 500) is the dominant remaining structural item** (engine KvHandle append-in-place, contract change). Then: moe glue, buffer reuse vs madvise. 692+1212+755 tests green; clippy clean (3 pre-existing warnings in `experts/arith`, not this work). Artifact §"update_2026_06_12_serial_cuts". @@ -43,21 +92,21 @@ the system of record for stage results. - **26B CPU MoE quantized residency BUILT + MEASURED — ≥10 tok/s target PASSED: 7.6 → 13.9 → 15.9 tok/s; gap to llama.cpp 4.5×→~1.9× (2026-06-11).** Built the three residency levers the C10 byte-ledger analysis named, all opt-in/default-off byte-identical, parity-first: **`LARQL_Q4K_LM_HEAD`** (resident decode loop routes lm_head through the vindex Q4_K view via `logits_to_predictions_q4_lm_head` — synthesized automatically for tied-embedding models; `larql-kv/generation.rs` `argmax_next_token_resident`), **`LARQL_Q4K_DIRECT_FFN`** (hybrid-MoE dense slab via `ffn_decode_step_native` inside new `moe_ffn_block_cpu_with_index`, threaded from `LocalMoeFfn{index}`; decode-only, prefill stays f32 gemm per the #16 falsification; **padded-down handling** added to both cached.rs twins — the 26B stores intermediate 2112 as 2304-col Q6_K rows, activation zero-pad is exact, parity test bit-class ≤1e-5 + ragged-bytes rejection), and **C12 v2: fused gate+up hand-asm** (`q4k_q8k_gate_up_asm` under the existing `LARQL_Q4K_ASM` — shared activation loads + two independent SDOT chains; **bit-exact first try**; microbench **9.89→19.82 GiB/s, 2.00×** — the fused neon form was the worst kernel in the file (horizontal-sum-bound), now the best, exceeding single-matrix asm 18.4). **E2E (M3 Max t=8 n=128 warm, llama-bench brackets 32.54→28.80 = ~11% drift, flag legs understated if anything):** 26B default **7.6** → full residency **13.9** → +ASM **15.9 tok/s**. **Honest negative:** 4B legacy bench is e2e-neutral for the fused asm (23.2 vs 22.9) — local direct-decode FFN deliberately uses two separate rayon-parallel matvecs, so `gate_up_into`'s production homes are the remote expert server (`q8k_wire`) + walk-FFN; the 26B ASM gain (13.9→15.9) comes from the `matvec_into` callers (experts, lm_head, dense gate/up). **Remaining ~1.9× to llama.cpp:** Q6_K asm (down, no asm today), attention q4k-direct still on the f32-activation `q4k_matvec` (route through q4k_q8k+asm), instruction-count reduction toward 28 GiB/s, optional rayon-chunked fused gate_up for local decode. 1148+755+686 tests green, clippy clean. Artifact §"update_2026_06_11" in [`bench/baselines/c10_gemma4-26b-a4b_cpu_reconciled.json`](bench/baselines/c10_gemma4-26b-a4b_cpu_reconciled.json). - **C10 26B-A4B CPU baseline LANDED — medium-term tier 62%→70%; the MoE gap is f32 residency, not the kernel (2026-06-10).** The owed number, measured under the full C10 discipline (M3 Max, t=8, warm, n=128, llama-bench drift bracket 32.08→30.45): **llama.cpp `-ngl 0` 32.1 ± 1.4 tok/s** vs larql in-process **7.1** (default) / **9.7 with `LARQL_Q4K_DIRECT_ATTN=1`** / loopback-shard **7.3**. **Mechanism, fully quantitative:** the in-process path streams **~10 GB/token** — attention Q/K/V/O (1.04B params, **f32 4.15 GB**) + dense FFN slab (**f32 2.14 GB**) + lm_head (**f32 2.95 GB**) + experts (Q4K 0.80 GB) — vs llama.cpp's ~2.1 GB all-quantized; all legs run bandwidth-saturated at ~62–71 GB/s, so the 4.8× byte ratio explains the whole 4.5× gap. The expert kernel (C12's turf) is NOT the bottleneck (~8% of bytes, already int8 Q8_K SDOT). **#16 reframed:** the "~0% at rep ctx" Q4K-direct-attention verdict was measured on the *network-bound loopback* path; in-process the flag is **1.36× decode + 10× TTFT** (prefill 6587→652 ms, skips `ensure_attn_tensors_dequantised`). **Path to ≥10 tok/s = quantized residency** (dense slab + lm_head q4k → ~2.0–2.3 GB/token → ~15–20 tok/s at the C12 1.6–1.8× residual). Tier gate favorable branch (target 10 ≪ llama.cpp 32) → **70%**. **Runbook mysteries resolved:** in-process-1.8-vs-loopback-4.4 was two artifacts (cold short-n smoke; warm AC = 7.1 vs 7.3, a wash — serialization ≈ second-process core gain); gemma4 llama.cpp CPU is NOT slow upstream (32 tok/s). **Measurement lesson (now method):** the first session (06-09 night) was invalidated by a silent battery drain to 31% — llama.cpp itself collapsed 34→1.05 tok/s (~30×, beyond the 1.5–3× thermal class), with Spotlight stacking on top after 30+ GB of model I/O; `pmset` AC check + cross-engine drift bracket are now mandatory. **Code gaps found:** `larql bench --moe-shards` still calls pre-C1 `generate_with_remote_moe` (fails on CPU with the #146 signature; `larql run --moe-shards --engine standard` is the working loopback instrument); `serve --ffn-only` doesn't serve expert endpoints (use `--experts 0-127`); `LARQL_DECODE_STAGES` doesn't record local-expert time. Artifact [`bench/baselines/c10_gemma4-26b-a4b_cpu_reconciled.json`](bench/baselines/c10_gemma4-26b-a4b_cpu_reconciled.json). - **FR3b — explicit relation-rewrite fallback BUILT + validated e2e (2026-06-09).** Wired the measured two-tier resolver into `SELECT … FROM EDGES WHERE relation=…` (`executor/query/select/edges.rs`): Tier 1 = the cached residual probe (unchanged); on probe abstain → Tier 2 `resolve_relation_explicit` — the few-shot `word→relation`+`music→none` frame (lifted verbatim from `fr3_explicit_rewrite.rs`), **one full forward via `InferenceWeights::predict_dense`** (the INFER path = `predict_kquant` w/ lm_head; the resolver's partial `0..=L10` dequant can't run lm_head — the called-out wrinkle), `none`-gated accept (`match_relation_top1`). **Opt-in `LARQL_FR3_EXPLICIT`, default off = byte-identical** (726 lql lib tests green incl. 4 new, clippy clean). **Real-vindex refinement:** production `gemma3-4b-q4k-v2.vindex` has **2890 noisy labels**; `relation_labels()` is alphabetical and both tiers cap at 64, so an alphabetical top-64 *drops* `language` while *keeping* a rare `food_animal` → "mother tongue" failed, "banana" resolved (backwards). Fixed with `RelationClassifier::relation_labels_ranked(top_n)` (by feature count) for Tier 2's candidates — keeps the meaningful relations + a short prompt. **E2E (real Gemma-3-4B, `LARQL_FR3_EXPLICIT=1`):** `mother tongue`→`language` by explicit (0.97, probe abstained — the win); `weather`→abstain (none-escape, no confident-wrong); default off → no resolution (byte-identical). **Honest correction:** the production 64-class probe is *stronger* than the 3-class ablation implied — it resolves `head city`→capital, `legal tender`→currency, `altitude`→elevation by meaning (Tier 1); Tier 2 is the safety net for genuine abstains, and on a rich label set the `none` escape is necessarily weaker (`banana`→`food_category`, a real relation here, defensible). Verdict [`docs/diagnoses/fr3-explicit-rewrite.md`](docs/diagnoses/fr3-explicit-rewrite.md) §"BUILD LANDED". -- **FR3b — relation probe is phrasing-brittle; explicit rewrite + `none` escape wins (2026-06-08). → BUILT 2026-06-09 (above).** FR3's 1.00 was synonym *words* in one template; on an unseen *phrasing* the probe sits at **chance** at its L10 probe layer (`examples/fr3_template_ablation.rs`: held-out "The {r} for {e} would be", N=6, k=1/2/4 templates = 0.33/0.39/0.39 @L10; signal is early at L6 and decays with depth). More training templates = measured **no-op** → that change was **reverted** (it 4×'d build cost for nothing at the probe layer). **Explicit few-shot `word→relation` classify** (one forward via `predict_kquant`, `examples/fr3_explicit_rewrite.rs`) = **12/12** on synonyms + unseen phrasings (head city→capital, legal tender→currency, mother tongue→language — exactly the probe's chance cases), but forced-choice confident-wrongs distractors **2/3** (weather/altitude→capital) → add a `none` escape + `music→none` few-shot → **0/3** (all abstain), 12/12 kept. The `none` escape is the verify/abstain — the project's recurring forced-choice trap (cf. FR1's 0.75 gate). **BUILD NEXT: probe-first / explicit-classify-with-`none` fallback** in `resolve_relation_synonym` (FR2 two-tier shape) — Tier 1 = existing probe (cheap, on `MIN_CONFIDENCE`), Tier 2 = explicit classify on abstain. *Wiring wrinkle:* Tier 2 needs lm_head (full forward) but `RelationResolver` only dequantises `0..=L10` → run it via the **Session vindex** (`predict_kquant`/`InferenceWeights`, the INFER path), not the resolver's partial setup (~30 lines, resolver→session boundary). Lift the few-shot frame + `none`-gated accept from `fr3_explicit_rewrite.rs`. Verdict [`docs/diagnoses/fr3-explicit-rewrite.md`](docs/diagnoses/fr3-explicit-rewrite.md). +- **FR3b — relation probe is phrasing-brittle; explicit rewrite + `none` escape wins (2026-06-08). → BUILT 2026-06-09 (above).** FR3's 1.00 was synonym *words* in one template; on an unseen *phrasing* the probe sits at **chance** at its L10 probe layer (`chris-experiments/larql_probes/examples/fleet_routing/fr3_template_ablation.rs`: held-out "The {r} for {e} would be", N=6, k=1/2/4 templates = 0.33/0.39/0.39 @L10; signal is early at L6 and decays with depth). More training templates = measured **no-op** → that change was **reverted** (it 4×'d build cost for nothing at the probe layer). **Explicit few-shot `word→relation` classify** (one forward via `predict_kquant`, `chris-experiments/larql_probes/examples/fleet_routing/fr3_explicit_rewrite.rs`) = **12/12** on synonyms + unseen phrasings (head city→capital, legal tender→currency, mother tongue→language — exactly the probe's chance cases), but forced-choice confident-wrongs distractors **2/3** (weather/altitude→capital) → add a `none` escape + `music→none` few-shot → **0/3** (all abstain), 12/12 kept. The `none` escape is the verify/abstain — the project's recurring forced-choice trap (cf. FR1's 0.75 gate). **BUILD NEXT: probe-first / explicit-classify-with-`none` fallback** in `resolve_relation_synonym` (FR2 two-tier shape) — Tier 1 = existing probe (cheap, on `MIN_CONFIDENCE`), Tier 2 = explicit classify on abstain. *Wiring wrinkle:* Tier 2 needs lm_head (full forward) but `RelationResolver` only dequantises `0..=L10` → run it via the **Session vindex** (`predict_kquant`/`InferenceWeights`, the INFER path), not the resolver's partial setup (~30 lines, resolver→session boundary). Lift the few-shot frame + `none`-gated accept from `fr3_explicit_rewrite.rs`. Verdict [`docs/diagnoses/fr3-explicit-rewrite.md`](docs/diagnoses/fr3-explicit-rewrite.md). - **FR retrieval-augmented early-exit — shipped + decode-loop declined (2026-06-08).** When the FR1 verified router fires, short-circuit the forward at the resolved stored layer (skip tail + lm_head). Parity-exact (40/40 residual + token), **1.44× on fact-lookup answer tokens** (457→319 ms, gemma3-4b-q4k), distractor-safe. Wired end-to-end: `INFER … ROUTE VERIFY EXIT` + `LARQL_KNN_EARLY_EXIT` (dense + q4k, verified-only — EXIT ignored with FALLBACK; on verify-miss completes the full forward = parity). 2419 lql+inference tests green, committed `d9b761f6`, pushed. **Decode-loop extension MEASURED & DECLINED:** parity-safe only on the terminal token (KV-cache invariant), blended T=2 1.15×/T=5 1.05×/fact-not-terminal 1.00× → net-marginal; INFER is the lever's home. Harnesses `examples/fr_early_exit_{probe,parity,bench,decode_projection}.rs`. -- **FR routing GAIN quantified — correctness, not throughput (2026-06-07).** Benchmark (`crates/larql-inference/examples/fr_routing_gain.rs`) runs all three router modes on the same forwards over CORRECT/DISTRACTOR/ALIAS slices (Gemma-3-4B, 20 installed facts @L26). **Legacy KNN injection is unsafe at scale: 0/20 distractor-safe** — 20 facts already confident-wrong 100% of unrelated queries (near-rank-1 cosine collides > 0.75). **FR1 verified fixes it: 0→100% distractor-safe**, installed-fact recall preserved (20/20), at **~13 µs/call** (0.05% of a decode — no tok/s cost). The gain is the KNN-fact-injection feature becoming safe to ship; it is NOT a throughput gain (the override is a post-logits sidecar; no-op without a KnnStore). **Caveat surfaced + hardened:** two-tier's fallback has no entity-name guard → 0/20 distractor-safe (same as legacy) → `ROUTE VERIFY FALLBACK` is a targeted alias tool, `ROUTE VERIFY` is the safe open default (doc-comment + LQL spec updated). FR4 is a dispatch criterion, not a speed gain. Verdict: [`docs/diagnoses/fr-routing-gain.md`](docs/diagnoses/fr-routing-gain.md). +- **FR routing GAIN quantified — correctness, not throughput (2026-06-07).** Benchmark (`chris-experiments/larql_probes/examples/fleet_routing/fr_routing_gain.rs`) runs all three router modes on the same forwards over CORRECT/DISTRACTOR/ALIAS slices (Gemma-3-4B, 20 installed facts @L26). **Legacy KNN injection is unsafe at scale: 0/20 distractor-safe** — 20 facts already confident-wrong 100% of unrelated queries (near-rank-1 cosine collides > 0.75). **FR1 verified fixes it: 0→100% distractor-safe**, installed-fact recall preserved (20/20), at **~13 µs/call** (0.05% of a decode — no tok/s cost). The gain is the KNN-fact-injection feature becoming safe to ship; it is NOT a throughput gain (the override is a post-logits sidecar; no-op without a KnnStore). **Caveat surfaced + hardened:** two-tier's fallback has no entity-name guard → 0/20 distractor-safe (same as legacy) → `ROUTE VERIFY FALLBACK` is a targeted alias tool, `ROUTE VERIFY` is the safe open default (doc-comment + LQL spec updated). FR4 is a dispatch criterion, not a speed gain. Verdict: [`docs/diagnoses/fr-routing-gain.md`](docs/diagnoses/fr-routing-gain.md). -- **Fleet routing extensions FR1/FR2/FR3 — MEASURED on a real vindex, all WIN, builds greenlit (2026-06-07).** The `chris-experiments/fleet` native-store arc (E10–E17) + `videos/the-mechanism` build story ported into the Query/Edit/Interpret track. Spec + frozen pre-registrations: [`docs/fleet-routing-extensions.md`](docs/fleet-routing-extensions.md); roadmap in [`ROADMAP.md`](ROADMAP.md) §"FR". Three Rust measurement harnesses run against `output/gemma3-4b-q4k-v2.vindex` (the production `KnnStore` cosine path + `capture_residuals`), judged in predictive units (mean-cosine banned), all three WIN: **FR1** (`examples/fr1_topk_fuzzy_router.rs`) — the entity key is real & answer-leak-free at L24-26 (L26 top1 **0.89**/top5 0.95, cross-rel 1.00, **beats E15's MLP under plain cosine-NN, no training**); the live `query_top1`+fixed-0.75 gate (`infer_patched.rs:162-163`) fires **150/150** with **11% confident-wrong @L26, 84% @L20** → the defect is the consumer, fix = top-k+verify+abstain at the resolved layer. **FR3** (`examples/fr3_relation_address.rs`) — relation synonym-gen **1.00 at every layer L6-L26** (semantic, not lexical; clean from L6, earlier than the video's L10); asymmetry stark vs entity top-1 0.07-0.20 until L26. **FR2** (`examples/fr2_two_tier_router.rs`) — symbolic exact-match **0/10** aliases, activation fallback **10/10 top-1** (Persia→Iran, …) = E16 reproduced (famous-alias easy end; general = FR1's ~0.9 top-5). Verdicts in `docs/diagnoses/fr{1,2,3}-*.md`; artifacts in `bench/aim-validation/fr{1,2,3}_*.json`. **FR4** (E17 compute→dispatch) remains research-first — E17's own ledger demotes the E4 bridge to a conjecture (G/O/T never ran), and the E17 rig lives in `chris-experiments`. **FR1 + FR2 BUILDS LANDED (2026-06-07).** `apply_knn_override_verified` (FR1: top-k + entity-in-prompt verify + abstain) and `apply_knn_override_two_tier` (FR2: tier-1 verify → tier-2 activation alias fallback), both resolved-layer-first (no hardcoded layer), wired into `infer_patched`/`infer_patched_q4k`, opt-in `LARQL_KNN_VERIFY` (+`LARQL_KNN_FALLBACK` for FR2), **default off = byte-identical** (23 infer_patched tests green incl. 14 legacy unchanged, clippy clean). E2E real Gemma-3-4B: FR1 fixes the measured confident-wrong (Germany-paraphrase legacy→SpainX, verified→GermanyX, no regression); FR2 recovers the alias "capital of Persia" (verify-only abstains→Tehran, two-tier→IranX cos 0.97). **LQL SURFACE LANDED (2026-06-07):** `KnnRouteMode` enum threaded through `infer_patched` (default `Legacy` = byte-identical, `from_env()` preserves env-gating for Python/EXPLAIN); first-class `INFER … ROUTE VERIFY [FALLBACK] [TOPK n]` clause (lexer+ast+parser+executor, 5 parser tests, 715 lql + 23 inference tests green, clippy clean, spec'd). E2E with no env vars: `ROUTE VERIFY` → Germany fixed, `ROUTE VERIFY FALLBACK` → Persia→IranX recovered. **FR3 BUILD LANDED (2026-06-07):** `RelationResolver` (`executor/relation_resolver.rs`) — a trained residual softmax probe (NOT string/cosine: residuals are near-rank-1, so cosine would be the "proxy is not the thing" trap), model-agnostic probe layer (`round(0.3·num_layers)`), wired into `SELECT … FROM EDGES WHERE relation=…` as a cached semantic fallback when exact-string misses. E2E real Gemma-3-4B: `WHERE relation="seat"` → resolved to "capital", returned the capital edges. **All three measured wins (FR1/FR2/FR3) are now built + validated e2e.** **FR4 RAN (2026-06-07) — E4 conjecture REFINED:** added the real external ops to the E17 rig — DIST (geometric) + ARGMIN (selection) **ride free @L1**, only PARTITION (global optimization) **walls like parity**. Parity was NOT a fair stand-in for "external"; E4's internal/external split mis-files geometric/selection (they factor through reads → internal). Re-cut dispatch criterion: keep count/filter/aggregate/threshold/majority/distance/argmin internal, route global-optimization+parity external (`E17_EXTERNAL_VERDICT.md`). **All four FR items (FR1-FR4) now complete** — FR1/FR2/FR3 measured+built+LQL-surfaced, FR4 measured (conjecture refined). +- **Fleet routing extensions FR1/FR2/FR3 — MEASURED on a real vindex, all WIN, builds greenlit (2026-06-07).** The `chris-experiments/fleet` native-store arc (E10–E17) + `videos/the-mechanism` build story ported into the Query/Edit/Interpret track. Spec + frozen pre-registrations: [`docs/fleet-routing-extensions.md`](docs/fleet-routing-extensions.md); roadmap in [`ROADMAP.md`](ROADMAP.md) §"FR". Three Rust measurement harnesses run against `output/gemma3-4b-q4k-v2.vindex` (the production `KnnStore` cosine path + `capture_residuals`), judged in predictive units (mean-cosine banned), all three WIN: **FR1** (`chris-experiments/larql_probes/examples/fleet_routing/fr1_topk_fuzzy_router.rs`) — the entity key is real & answer-leak-free at L24-26 (L26 top1 **0.89**/top5 0.95, cross-rel 1.00, **beats E15's MLP under plain cosine-NN, no training**); the live `query_top1`+fixed-0.75 gate (`infer_patched.rs:162-163`) fires **150/150** with **11% confident-wrong @L26, 84% @L20** → the defect is the consumer, fix = top-k+verify+abstain at the resolved layer. **FR3** (`chris-experiments/larql_probes/examples/fleet_routing/fr3_relation_address.rs`) — relation synonym-gen **1.00 at every layer L6-L26** (semantic, not lexical; clean from L6, earlier than the video's L10); asymmetry stark vs entity top-1 0.07-0.20 until L26. **FR2** (`chris-experiments/larql_probes/examples/fleet_routing/fr2_two_tier_router.rs`) — symbolic exact-match **0/10** aliases, activation fallback **10/10 top-1** (Persia→Iran, …) = E16 reproduced (famous-alias easy end; general = FR1's ~0.9 top-5). Verdicts in `docs/diagnoses/fr{1,2,3}-*.md`; artifacts in `bench/aim-validation/fr{1,2,3}_*.json`. **FR4** (E17 compute→dispatch) remains research-first — E17's own ledger demotes the E4 bridge to a conjecture (G/O/T never ran), and the E17 rig lives in `chris-experiments`. **FR1 + FR2 BUILDS LANDED (2026-06-07).** `apply_knn_override_verified` (FR1: top-k + entity-in-prompt verify + abstain) and `apply_knn_override_two_tier` (FR2: tier-1 verify → tier-2 activation alias fallback), both resolved-layer-first (no hardcoded layer), wired into `infer_patched`/`infer_patched_q4k`, opt-in `LARQL_KNN_VERIFY` (+`LARQL_KNN_FALLBACK` for FR2), **default off = byte-identical** (23 infer_patched tests green incl. 14 legacy unchanged, clippy clean). E2E real Gemma-3-4B: FR1 fixes the measured confident-wrong (Germany-paraphrase legacy→SpainX, verified→GermanyX, no regression); FR2 recovers the alias "capital of Persia" (verify-only abstains→Tehran, two-tier→IranX cos 0.97). **LQL SURFACE LANDED (2026-06-07):** `KnnRouteMode` enum threaded through `infer_patched` (default `Legacy` = byte-identical, `from_env()` preserves env-gating for Python/EXPLAIN); first-class `INFER … ROUTE VERIFY [FALLBACK] [TOPK n]` clause (lexer+ast+parser+executor, 5 parser tests, 715 lql + 23 inference tests green, clippy clean, spec'd). E2E with no env vars: `ROUTE VERIFY` → Germany fixed, `ROUTE VERIFY FALLBACK` → Persia→IranX recovered. **FR3 BUILD LANDED (2026-06-07):** `RelationResolver` (`executor/relation_resolver.rs`) — a trained residual softmax probe (NOT string/cosine: residuals are near-rank-1, so cosine would be the "proxy is not the thing" trap), model-agnostic probe layer (`round(0.3·num_layers)`), wired into `SELECT … FROM EDGES WHERE relation=…` as a cached semantic fallback when exact-string misses. E2E real Gemma-3-4B: `WHERE relation="seat"` → resolved to "capital", returned the capital edges. **All three measured wins (FR1/FR2/FR3) are now built + validated e2e.** **FR4 RAN (2026-06-07) — E4 conjecture REFINED:** added the real external ops to the E17 rig — DIST (geometric) + ARGMIN (selection) **ride free @L1**, only PARTITION (global optimization) **walls like parity**. Parity was NOT a fair stand-in for "external"; E4's internal/external split mis-files geometric/selection (they factor through reads → internal). Re-cut dispatch criterion: keep count/filter/aggregate/threshold/majority/distance/argmin internal, route global-optimization+parity external (`E17_EXTERNAL_VERDICT.md`). **All four FR items (FR1-FR4) now complete** — FR1/FR2/FR3 measured+built+LQL-surfaced, FR4 measured (conjecture refined). - **In-process CPU MoE bench wired + 26B baseline staged from the CACHED model (no download) — baseline NUMBER still owed, blocked on an idle machine (2026-06-06).** Two threads toward the C10 "still owes 26B-A4B" pin. **(1) In-process KV-cached CPU MoE decode in `larql bench`:** new `LocalMoeFfn` ([`crates/larql-inference/src/ffn/local_moe.rs`](crates/larql-inference/src/ffn/local_moe.rs) — the local twin of `RemoteMoeFfn`; experts via `moe_ffn_block_cpu(.., moe_remote=None)`) + [`bench/local_moe_runtime.rs`](crates/larql-cli/src/commands/primary/bench/local_moe_runtime.rs) driving `generate_with_engine_resident`, + a MoE-detect branch in `bench/run.rs` (a hybrid-MoE CPU vindex routes through `LocalMoeFfn`; the dense NullFfn/legacy CPU rows that silently drop experts are skipped). This is the **FAIR single-box CPU MoE number** (no loopback-shard network tax). Verified e2e on the real Gemma-4-26B-A4B (output "Paris.", KV-cached); 3 parity unit tests + clippy green. **Smoke n=8 = 1.8 tok/s — LOWER than the C1 loopback 4.4** (likely process-parallelism: the loopback client+server use more cores than one 8-thread process) — **UNVERIFIED**, needs n=128 warm on an idle machine; the expert kernel is already the optimized `run_single_expert_q4k_q8k_into`. **(2) 26B Q4_K_M GGUF built from the CACHED safetensors (no 16GB download):** `google/gemma-4-26B-A4B-it` was already in the HF cache (48GB) → llama.cpp HEAD `conversion/gemma.py` (`Gemma4Model`) → BF16 → `llama-quantize` Q4_K_M (`/tmp/gemma4-26b-Q4_K_M.gguf`, 16GB, 5.32 BPW; **60/658 tensors fell back to Q8_0** — the non-256-aligned experts, so a future vindex→GGUF export *can't* be byte-identical for MoE). **Baseline number BLOCKED on an idle machine** (C10 discipline — under load a known-good 4B clocked 0.5 tok/s vs ~43 warm = pure contention; not a real number). **OPEN: gemma4 CPU speed in llama.cpp is unverified** — a 1-core/slow run was observed but confounded by machine load; confirm on an idle machine, and if still slow that's a real upstream finding (gemma4 CPU MoE unoptimized). Runbook + exact commands: [`bench/baselines/c10_gemma4-26b-a4b_cpu_RUNBOOK.md`](bench/baselines/c10_gemma4-26b-a4b_cpu_RUNBOOK.md). **(3) Banked tooling:** `GgufWriter` ([`crates/larql-models/src/loading/gguf/writer.rs`](crates/larql-models/src/loading/gguf/writer.rs), round-trips the production reader) + `larql convert gguf-info` now dumps the tensor-info table — the foundation + spec-extraction for an eventual **vindex→GGUF exporter** (tracked separately; the cached-model path made it optional for the baseline, and the Q8_0 expert fallback means it could never be byte-identical anyway). -- **V1 MoE-within-expert — FALSIFIED, KU4 fully closed (2026-06-05).** Built the previously-OPEN half of V1: does feature/hash routing work *inside a single MoE expert's FFN*? (The dense V1 harness measures the wrong object on the 26B-A4B — each per-layer block is 128 stacked experts, not one dense FFN.) **Mechanism, parity-first:** opt-in within-expert feature pruning installed *inside the production expert kernel* (`run_single_expert_q4k_q8k_into`) via a global schedule ([`crates/larql-compute/src/cpu/ops/moe/within_expert.rs`](crates/larql-compute/src/cpu/ops/moe/within_expert.rs), `set_routing`/`set_current_layer`), **OFF by default = one relaxed atomic load → byte-exact parity** (the spine, `feedback_engineering_vs_research_posture`); errors propagate through the real forward (`predict_kquant`, in-process local MoE). Harness [`examples/walk_ffn_v1_moe_within_expert.rs`](crates/larql-inference/examples/walk_ffn_v1_moe_within_expert.rs) mirrors V1's 3 phases. 7 new unit tests + 78 existing MoE tests green, clippy clean. **Parity anchor GREEN** (all-dense schedule KL=0.00000). **Result — FALSIFIED, same as dense V1** (Gemma 4 26B-A4B, 30 layers, 128 experts, top_k=8, inter=704): **(Phase A depth split)** L0–13 need ALL 704 features (experts dense in their own feature space; even 1/2 exceeds KL 0.05), only L14–29 tolerate per-layer pruning (frac 0.016–0.25), mean 0.52. **(Phase B claim gate)** at per-layer thresholds together → **50% argmax drift** (first-div pos 2); mean NLL Δ=−0.15 bits (comp *lower*, ppl −9.88%) — the **#26 trap textbook** (`feedback_metric_matches_operation`): mean NLL is noise-dominated and points the WRONG way, drift is the deciding signal. **(Bandwidth)** deployable oracle only **~1.19×** (pays gate+up, half the layers save nothing); 1.91× best case unrealizable — **(Phase C)** cheap content-blind route clears KL≤0.05 at only **6% (1/16)** of small-threshold layers. **The expert is already a compact specialized FFN — no second sparsity axis inside it.** With dense V1, the FFN/expert feature-sparsity bandwidth multiplier is **dead on this arch, dense AND MoE**; medium-term tier rests on FP4 (V2), not feature routing. Artifact `bench/aim-validation/v1moe_gemma4-26b-a4b-q4k.json`; writeup [`docs/diagnoses/v1-moe-within-expert.md`](docs/diagnoses/v1-moe-within-expert.md). Caveat: one model + one held passage; ActMagnitude oracle (a ‖down‖-weighted oracle would only raise thresholds). A shared/fine-grained-expert MoE follow-up is cheap (the instrument is parity-safe) if re-opened. +- **V1 MoE-within-expert — FALSIFIED, KU4 fully closed (2026-06-05).** Built the previously-OPEN half of V1: does feature/hash routing work *inside a single MoE expert's FFN*? (The dense V1 harness measures the wrong object on the 26B-A4B — each per-layer block is 128 stacked experts, not one dense FFN.) **Mechanism, parity-first:** opt-in within-expert feature pruning installed *inside the production expert kernel* (`run_single_expert_q4k_q8k_into`) via a global schedule ([`crates/larql-compute/src/cpu/ops/moe/within_expert.rs`](crates/larql-compute/src/cpu/ops/moe/within_expert.rs), `set_routing`/`set_current_layer`), **OFF by default = one relaxed atomic load → byte-exact parity** (the spine, `feedback_engineering_vs_research_posture`); errors propagate through the real forward (`predict_kquant`, in-process local MoE). Harness `chris-experiments/larql_probes/examples/walk_ffn/walk_ffn_v1_moe_within_expert.rs` mirrors V1's 3 phases. 7 new unit tests + 78 existing MoE tests green, clippy clean. **Parity anchor GREEN** (all-dense schedule KL=0.00000). **Result — FALSIFIED, same as dense V1** (Gemma 4 26B-A4B, 30 layers, 128 experts, top_k=8, inter=704): **(Phase A depth split)** L0–13 need ALL 704 features (experts dense in their own feature space; even 1/2 exceeds KL 0.05), only L14–29 tolerate per-layer pruning (frac 0.016–0.25), mean 0.52. **(Phase B claim gate)** at per-layer thresholds together → **50% argmax drift** (first-div pos 2); mean NLL Δ=−0.15 bits (comp *lower*, ppl −9.88%) — the **#26 trap textbook** (`feedback_metric_matches_operation`): mean NLL is noise-dominated and points the WRONG way, drift is the deciding signal. **(Bandwidth)** deployable oracle only **~1.19×** (pays gate+up, half the layers save nothing); 1.91× best case unrealizable — **(Phase C)** cheap content-blind route clears KL≤0.05 at only **6% (1/16)** of small-threshold layers. **The expert is already a compact specialized FFN — no second sparsity axis inside it.** With dense V1, the FFN/expert feature-sparsity bandwidth multiplier is **dead on this arch, dense AND MoE**; medium-term tier rests on FP4 (V2), not feature routing. Artifact `bench/aim-validation/v1moe_gemma4-26b-a4b-q4k.json`; writeup [`docs/diagnoses/v1-moe-within-expert.md`](docs/diagnoses/v1-moe-within-expert.md). Caveat: one model + one held passage; ActMagnitude oracle (a ‖down‖-weighted oracle would only raise thresholds). A shared/fine-grained-expert MoE follow-up is cheap (the instrument is parity-safe) if re-opened. - **C10 discrepancy RESOLVED + CPU bench free wins + C12 v1 asm kernel (2026-06-02).** Three things. **(1) C10 resolved — no regression; true gap ~1.6–1.8×.** The 1.50× (05-16) vs 1.93× (05-31) split was two stacked measurement confounds: a **larql path mismatch** (27.6 = `StandardEngine` path vs 23.6 = legacy `larql bench --cpu` path, a stable ~12% delta) and a **llama.cpp harness artifact** (the 45.5 was an unwarmed/short-n ollama `num_gpu=0` fluke; warmed + n=128 it converges to 42.8–43.0 = llama-bench's 42.99 — both harnesses/both dates agree at ~43). Reconciled like-for-like (M3 Max, t=8, warm): larql 23.5 legacy / 26.4 StandardEngine vs llama.cpp 43.0. Artifact `bench/baselines/c10_gemma3-4b_cpu_reconciled.json`. **(2) Free bench-tool wins:** `larql bench --cpu` now also reports the production StandardEngine row (was undersold by the legacy row); new `--ollama-cpu` forces `num_gpu=0`+`num_thread` so `--ollama` is a true CPU baseline (was silently Metal-GPU — the 97.7 in the old artifact was a GPU leak). **(3) C12 v1 asm kernel landed opt-in (`LARQL_Q4K_ASM=1`).** Falsification-first **roofline microbench** (`crates/larql-compute/benches/q4k_q8k_matvec.rs`) proved the kernel is **compute/issue-bound, not DRAM-bandwidth-bound** (scalar 9.3 vs NEON 17.7 GiB/s, identical data, size-invariant) — **overturns the `DIAGNOSIS-2026-05-16` "memory-system-level" conclusion**. `q4k_q8k_matvec_asm` (one `asm!` block/super-block, vector scale lanes) is **bit-exact** (`q8k_matvec_asm_matches_scalar_bit_exact`), **+3.7–4.9% isolated**, ~+1–2% e2e (capped because the fused `gate_up` path isn't asm-covered yet). **Finding: latency-hiding has low headroom** (a 4-accumulator variant gave no reliable gain — the OoO core already overlaps super-blocks), so the **two-super-block interleave is deprioritized**; the real lever is **instruction-count reduction + asm-ifying `gate_up`**. Spec updated: [`crates/larql-compute/docs/q4k-decode-kernel.md`](crates/larql-compute/docs/q4k-decode-kernel.md) §"2026-06-02 roofline measurement". - **C10 CPU baseline re-measured (2026-05-31).** *(Superseded by the 2026-06-02 resolution above — the "⚠️ reconcile" was the two-confound artifact, now explained; no regression.)* Fresh quant-matched CPU bench: **Gemma 3 4B — larql 23.6 tok/s vs llama.cpp `-ngl 0` 45.5 (~1.9×)**, NOT within the short-term 10% bar (baseline via ollama `num_gpu=0`, same Q4_K_M GGUF — ollama IS llama.cpp; homebrew `llama-bench` wouldn't load ollama blobs). **⚠️ This is WORSE than the recorded C10 in ROADMAP.md (1.50×, 27.6 vs 41.4 tok/s, 2026-05-16) — reconcile before acting** (candidates: thermal, short-prompt/small-n, ollama-vs-direct-`llama.cpp`, or a real decode regression). Decode split (`LARQL_DECODE_STAGES=1`, 40 tok): **FFN ~50% / attn ~38% / lm_head ~12%**. **Correction to a mid-session error:** I initially traced `larql-compute/src/attention/decode.rs` (an f32 BLAS / `LARQL_Q4K_DIRECT_ATTN`-opt-in path) and wrongly concluded "attention runs f32, wire it to int8 (C11)". The PRODUCTION decode path is `attention_decode_step_native` (`larql-inference/.../kquant_forward/cached.rs:563`), which **already** `quantize_x_to_q8k_into`s the activation (ll.609/710) and runs Q/K/V/O through the int8 Q8_K SDOT kernel — **same as the FFN**, exactly as the pre-existing **C12** documents. So there is NO attention-wiring win; the real CPU lever remains **C12** (hand-asm aarch64 Q4K×Q8K to close the ~1.73×-per-core Rust-intrinsics-vs-llama.cpp gap). The bogus "C11" item is retracted (and C11 already means "architecture-rule CI" in ROADMAP.md — naming collision avoided). Caveat that stands: small n; lm_head CPU matvec path not separately traced. Artifact `bench/baselines/c10_gemma3-4b_cpu_vs_llamacpp.json`. - **MoE routing locality (Gemma 4 26B-A4B) — resolves KU5 locality half, NEGATIVE for the long-term tier (2026-05-31).** Faithful in-process decode (the full vindex carries experts in `layers/*.weights` → `moe_ffn_block_cpu`/`build_moe_weights` computes them locally; `MOE_DEBUG=1` emits per-layer expert selections; one prefill forward over a ~72-token passage gives clean per-position routing). 128 experts, top_k=8, 30 layers. **Finding:** per-token routing is sparse (8/128) but over a sequence the expert **working set saturates to ~124/128 — the uniform-random expectation** (load-balanced top_k_softmax router); adjacent-token reuse is only 21% (3.4× random but weak), top-10 expert mass just 17%. **No small cacheable hot subset.** Implication: the 26B's full expert set (~11 GB) fits RAM so it's fine after warmup, but a **>RAM frontier MoE has no hot subset to cache** → it pages ~the whole population (~200 ms/token-class with V3's ~100µs cold page) → the long-term disk-residency bet is undermined ("bigger RAM" not "spill to disk"). **Achievability revised 2026-05-31: long-term 60%→52%, ultimate 40%→30%** (671B@FP4 exceeds one machine + no hot subset closes the disk-resident escape hatch). Aim-validation pattern: per-token sparsity that *looks* cacheable isn't, because it doesn't concentrate over time. Caveat: one model + one passage; Gemma's balanced router may spread more than a shared/fine-grained MoE — cross-MoE-router check open. Analysis `bench/aim-validation/moe-routing/`; writeup [`docs/diagnoses/moe-routing-locality.md`](docs/diagnoses/moe-routing-locality.md). - **V2 (FP4 generality) CONFIRMED + V3 (disk-resident mmap) feasibility probe (2026-05-31).** Next two aim-validation items after V1. **V2 — KU3 RESOLVED, CONFIRMED (the opposite of V1):** static scan (`fp4_q1_scan`, generalized for `*_weights.bin` naming) on *original f16* weights shows **≥99.8% per-feature R<16** across Gemma 3 4B + Granite 4.1 3B + 8B (2 families, 3B/4B/8B ladder) — reproduces exp 26's 99.83% on gemma3 `down` exactly, `down` the only mild tail (worst layer ~99.4%, p99 R≈12). Predictive deciding-metric check (`walk_ffn_v2_fp4_nll`, the **real E2M1 block codec** `encode_fp4_feature`/`decode_fp4_feature`): FP4 within **+0.116 bits/token** of f32 and **beats the shipped Q4-int baseline** (E2M1 float > 4-bit symmetric int) — no compounding catastrophe, no QAT. The FP4 ~2× lever holds cross-arch. *Model-set pivot* (flagged): matrix's Llama/Mistral are q4k-only (double-quant would be dishonest), so used the f16 originals on hand; Llama/Mistral/MoE-expert f16 not covered. **V3 — KU5 PARTIAL:** the roadmap assumed a 32 GB box; this is **128 GB** (nothing exceeds RAM, macOS can't cap mmap residency), so the probe (`mmap_cold_read_probe`) measures cold reads directly — `F_NOCACHE` pread for scattered cold + mmap+`MADV_DONTNEED` for cold faults **verified by getrusage major-fault counts** (Apple-Silicon 16 KB pages; Darwin re-evict is lazy, so the scattered pread must run first on a fresh blob — the probe self-checks via SPINE). Result (granite-30b 17 GB, verified cold): **cold scattered 16 KB read ~100µs p50 / 140µs p99 (153 MB/s); warm reuse ~0.04µs → ~2380× gap.** Disk-resident sparse access is viable *in steady state* but cold-start is brutal (~2 GB cold working set ≈ seconds) — viability hinges on the MoE-routing cache hit rate. **DEFERRED:** steady-state fault-rate + end-to-end tok/s on a model that genuinely exceeds RAM (needs a >128 GB-class vindex or a Linux/cgroup box). Artifacts `bench/aim-validation/{v2_*_scan,v3_granite-30b}.json`; writeups [`docs/diagnoses/v2-fp4-generality.md`](docs/diagnoses/v2-fp4-generality.md), [`docs/diagnoses/v3-disk-resident-mmap.md`](docs/diagnoses/v3-disk-resident-mmap.md). Achievability revised 2026-05-31 (full ladder in ROADMAP.md): medium-term **80%→62%** (FP4 confirmed + 26B fits RAM, but measured ~4.4 tok/s vs 10 target and hash lever gone), long-term **60%→52%** (100B@FP4 fits RAM so disk bet matters less, trimmed for lost multiplier), ultimate **40%→30%**, dense-frontier **15%→10%**. -- **V1 aim-validation — hash routing across all layers: FALSIFIED on dense, KU4 resolved (2026-05-31).** The first true-P0 aim-validation test. New harness `crates/larql-inference/examples/walk_ffn_v1_hash_routing.rs` reuses the `walk_ffn_*` infrastructure (`predict_with_ffn`, `WalkFfnConfig` per-layer `k_per_layer`, `local_pool_gate_knn` cheap path) and is **parity-anchored to exp 27** (full-K walk == dense KL≈0; gate-oracle top-2048 @ L0 KL 0.011, the expected ≤ exp-27's 0.030 since oracle ≤ token-ID hash). Three stages, judged only in predictive units: **Phase A** per-layer oracle threshold (min gate-top-k for lm_head KL ≤ 0.05, one layer sparse at a time); **Phase B** compounding (all thresholds at once → held-text NLL distribution + argmax drift — the #26 claim gate); **Phase C** cheap-route realizability (strided + ‖down_row‖ vs oracle KL). **Result — unanimous across 3 dense archs** (Gemma 3 4B / Llama 2 7B / Mistral 7B): per-layer thresholds are small (mean 2.7–12.2% of features) **but DO NOT compound** — applied simultaneously they give **+5.4 to +7.7 bits/token NLL and 78–95% argmax drift** (Llama diverges from token 0). The per-layer KL screen is **anti-correlated** with the truth: the sparser it lets you push, the worse the collapse — single-layer KL is the wrong proxy for an operation repeated across depth (`feedback_metric_matches_operation`). Honest bandwidth: cheap "best case" 8–37× is mostly unrealisable (Phase C clears KL ≤ 0.05 at only 16–62% of layers) and the deployable gate-oracle config touches ~0.37–0.41× → only **2.4–2.9×** (gate projection still paid) — and *that* config is the one Phase B shows is catastrophic. **The 5× within-FFN hash-routing bandwidth multiplier is dead**, strengthening the WalkFfn thread (#17–#28: the FFN is dense). KU4 resolved for dense; **MoE-within-expert version is OPEN** — the 26B `interleaved_q4k` is a dense-collapsed MLP, so this dense harness would measure the wrong object; tracked as a V1 follow-up needing expert-aware tooling. Roadmap impact: medium-term driver narrowed (rests on expert active-param sparsity, not FFN hash routing); tier later revised 80%→62% (2026-05-31). Artifacts `bench/aim-validation/v1_*.json`; full writeup [`docs/diagnoses/v1-hash-routing.md`](docs/diagnoses/v1-hash-routing.md). -- **Q4K-direct attention — step-3 decode path built, parity GREEN (task #16)** (2026-05-30). New CPU `run_attention_block_decode_step_q4k_direct` (`attention/decode.rs`) mirrors the f32 `run_attention_block_decode_step_backend` byte-for-byte except the four projections, which now run `quant_matvec` straight from the index (`resolve_attn_weights`, per-matrix dispatch: Q/K/O→`q4k_matvec`, **V→`q6k_matvec`** — the 26B manifest has V=Q6_K). Wired **opt-in** in `CpuBackend::attention_step` behind `LARQL_Q4K_DIRECT_ATTN=1` with per-layer f32 fallback (default off → zero behaviour change; the previously-ignored `index` arg now consumed). **Parity gate ✅** (`vindex::dequant`, two tests — all-Q4_K **and mixed V=Q6_K** matching the real 26B layout, so `q6k_matvec` dispatch is actually exercised on both sides): Q4K-direct vs Q4K-**dequant** (same bytes) agree < 1e-3 max-abs on h/k/v across layers — the spine (parity before timing) green at unit level. 97 attention + 44 kv_dispatch tests pass, clippy clean, larql-kv/cli build. **⚠️ Consolidation hazard tracked:** a second independent CPU Q4K attn path (`cached_decode_step_q4k`/`CpuQ4kCacheHandle`) exists — must stay in agreement on RoPE/softcap/norms; consolidate before either is load-bearing (doc §"CONSOLIDATION HAZARD", code pointer in cpu.rs). **End-to-end RAN on real 26B (2026-05-31) — flag fires, parity holds, net win marginal + context-dependent.** Unblocked the inert flag with a **resident-weights path** (`KvEngine::{prefill,decode_step}_resident` + `AnyEngine` forwarders + `generate_with_engine_resident` + CLI swap): since the moe-shards path pre-dequantises weights f32-resident, the resident methods take `&weights` (no lazy-dequant `&mut`) and just thread `index`, so engine + `RemoteMoeFfn` both borrow `&weights` — no conflict, no Arc (the original blocker: `decode_step_quant`'s `&mut weights` vs `RemoteMoeFfn`'s `&weights`). Results (localhost, `--engine standard`, warm, interleaved): **short ctx (~14)** attn −13%, decode **8.22→8.60 tok/s (+4.6%, n=6 tight)**, "Paris." parity ✓; **representative ctx (907 tok, n=5)** decode **OFF 4.80 / ON 4.92 mean — net ≈0 within ~6% noise** (the +4.6% washes out: per-token decode is MoE-network + f32-GQA bound at depth, not attention-proj bound; short-ctx flattered, #26-class). **Prefill twin GATED → FALSIFIED:** prefill attention is 43% of TTFT but the seq_len=907 prefill A/B (`examples/attn_prefill_f32_vs_q4k.rs`) shows q4k repeated-matvec **~20× SLOWER** than f32 BLAS — prefill projection is a compute-bound gemm (AMX's turf), Q4K's bandwidth edge evaporates; no `q4k_matmul`, and even one couldn't beat AMX at that seq_len. **Do not build the twin.** Net of the arc: Q4K-direct wins ONLY in the bandwidth-bound decode-matvec regime, washes out end-to-end behind expert/GQA, loses outright at prefill. Ship call: keep decode path (correct, parity-safe, opt-in `LARQL_Q4K_DIRECT_ATTN`), **don't headline a tok/s number**; throughput levers are the expert/network path + TTFT's O(N²) GQA, not Q4K attention. Detail: [`docs/diagnoses/q4k-direct-attention.md`](docs/diagnoses/q4k-direct-attention.md) §"Open questions". -- **Q4K-direct attention — step-2 gates both GREEN → build greenlit (task #16)** (2026-05-30). Two cheap probes (real production fns, synthetic same-size f32 weights, `CpuBackend`) run *before* building — the #24-trap guard. **Gate 1 (`examples/attn_proj_vs_gqa_split.rs`):** the four Q/K/V/O projections are **97–98% of the attention block at the measured context band (cached_len 32–128)** — Q4K-direct addresses ~all of the 28%; projection cost is flat/bandwidth-bound (~40 ms/token blended), GQA grows with cached_len (stays ≥64% to 4K, ~50% at 8K with the W=1024 sliding cap). **Gate 2 (`examples/attn_proj_f32_vs_q4k.rs`):** `q4k_matvec` **beats AMX/Accelerate f32 BLAS 2.06× (sliding block) / 2.51× (global block)** — the bandwidth cut nets out, not eaten by AMX throughput. Amdahl: projections ≈ 97%×28% ≈ ~27% of decode → a 2.1–2.5× cut → **~14–16% net decode win from attention alone**. **These are ISOLATED timings** (#24 lesson) — ship gate stays end-to-end net decode tok/s > dense on the real 26B vindex; the isolated evidence just says build, very unlikely to wash. Results appended to [`docs/diagnoses/q4k-direct-attention.md`](docs/diagnoses/q4k-direct-attention.md) §"Step 2". +- **V1 aim-validation — hash routing across all layers: FALSIFIED on dense, KU4 resolved (2026-05-31).** The first true-P0 aim-validation test. New harness `chris-experiments/larql_probes/examples/walk_ffn/walk_ffn_v1_hash_routing.rs` reuses the `walk_ffn_*` infrastructure (`predict_with_ffn`, `WalkFfnConfig` per-layer `k_per_layer`, `local_pool_gate_knn` cheap path) and is **parity-anchored to exp 27** (full-K walk == dense KL≈0; gate-oracle top-2048 @ L0 KL 0.011, the expected ≤ exp-27's 0.030 since oracle ≤ token-ID hash). Three stages, judged only in predictive units: **Phase A** per-layer oracle threshold (min gate-top-k for lm_head KL ≤ 0.05, one layer sparse at a time); **Phase B** compounding (all thresholds at once → held-text NLL distribution + argmax drift — the #26 claim gate); **Phase C** cheap-route realizability (strided + ‖down_row‖ vs oracle KL). **Result — unanimous across 3 dense archs** (Gemma 3 4B / Llama 2 7B / Mistral 7B): per-layer thresholds are small (mean 2.7–12.2% of features) **but DO NOT compound** — applied simultaneously they give **+5.4 to +7.7 bits/token NLL and 78–95% argmax drift** (Llama diverges from token 0). The per-layer KL screen is **anti-correlated** with the truth: the sparser it lets you push, the worse the collapse — single-layer KL is the wrong proxy for an operation repeated across depth (`feedback_metric_matches_operation`). Honest bandwidth: cheap "best case" 8–37× is mostly unrealisable (Phase C clears KL ≤ 0.05 at only 16–62% of layers) and the deployable gate-oracle config touches ~0.37–0.41× → only **2.4–2.9×** (gate projection still paid) — and *that* config is the one Phase B shows is catastrophic. **The 5× within-FFN hash-routing bandwidth multiplier is dead**, strengthening the WalkFfn thread (#17–#28: the FFN is dense). KU4 resolved for dense; **MoE-within-expert version is OPEN** — the 26B `interleaved_q4k` is a dense-collapsed MLP, so this dense harness would measure the wrong object; tracked as a V1 follow-up needing expert-aware tooling. Roadmap impact: medium-term driver narrowed (rests on expert active-param sparsity, not FFN hash routing); tier later revised 80%→62% (2026-05-31). Artifacts `bench/aim-validation/v1_*.json`; full writeup [`docs/diagnoses/v1-hash-routing.md`](docs/diagnoses/v1-hash-routing.md). +- **Q4K-direct attention — step-3 decode path built, parity GREEN (task #16)** (2026-05-30). New CPU `run_attention_block_decode_step_q4k_direct` (`attention/decode.rs`) mirrors the f32 `run_attention_block_decode_step_backend` byte-for-byte except the four projections, which now run `quant_matvec` straight from the index (`resolve_attn_weights`, per-matrix dispatch: Q/K/O→`q4k_matvec`, **V→`q6k_matvec`** — the 26B manifest has V=Q6_K). Wired **opt-in** in `CpuBackend::attention_step` behind `LARQL_Q4K_DIRECT_ATTN=1` with per-layer f32 fallback (default off → zero behaviour change; the previously-ignored `index` arg now consumed). **Parity gate ✅** (`vindex::dequant`, two tests — all-Q4_K **and mixed V=Q6_K** matching the real 26B layout, so `q6k_matvec` dispatch is actually exercised on both sides): Q4K-direct vs Q4K-**dequant** (same bytes) agree < 1e-3 max-abs on h/k/v across layers — the spine (parity before timing) green at unit level. 97 attention + 44 kv_dispatch tests pass, clippy clean, larql-kv/cli build. **⚠️ Consolidation hazard tracked:** a second independent CPU Q4K attn path (`cached_decode_step_q4k`/`CpuQ4kCacheHandle`) exists — must stay in agreement on RoPE/softcap/norms; consolidate before either is load-bearing (doc §"CONSOLIDATION HAZARD", code pointer in cpu.rs). **End-to-end RAN on real 26B (2026-05-31) — flag fires, parity holds, net win marginal + context-dependent.** Unblocked the inert flag with a **resident-weights path** (`KvEngine::{prefill,decode_step}_resident` + `AnyEngine` forwarders + `generate_with_engine_resident` + CLI swap): since the moe-shards path pre-dequantises weights f32-resident, the resident methods take `&weights` (no lazy-dequant `&mut`) and just thread `index`, so engine + `RemoteMoeFfn` both borrow `&weights` — no conflict, no Arc (the original blocker: `decode_step_quant`'s `&mut weights` vs `RemoteMoeFfn`'s `&weights`). Results (localhost, `--engine standard`, warm, interleaved): **short ctx (~14)** attn −13%, decode **8.22→8.60 tok/s (+4.6%, n=6 tight)**, "Paris." parity ✓; **representative ctx (907 tok, n=5)** decode **OFF 4.80 / ON 4.92 mean — net ≈0 within ~6% noise** (the +4.6% washes out: per-token decode is MoE-network + f32-GQA bound at depth, not attention-proj bound; short-ctx flattered, #26-class). **Prefill twin GATED → FALSIFIED:** prefill attention is 43% of TTFT but the seq_len=907 prefill A/B (`chris-experiments/larql_probes/examples/q4k_attention/attn_prefill_f32_vs_q4k.rs`) shows q4k repeated-matvec **~20× SLOWER** than f32 BLAS — prefill projection is a compute-bound gemm (AMX's turf), Q4K's bandwidth edge evaporates; no `q4k_matmul`, and even one couldn't beat AMX at that seq_len. **Do not build the twin.** Net of the arc: Q4K-direct wins ONLY in the bandwidth-bound decode-matvec regime, washes out end-to-end behind expert/GQA, loses outright at prefill. Ship call: keep decode path (correct, parity-safe, opt-in `LARQL_Q4K_DIRECT_ATTN`), **don't headline a tok/s number**; throughput levers are the expert/network path + TTFT's O(N²) GQA, not Q4K attention. Detail: [`docs/diagnoses/q4k-direct-attention.md`](docs/diagnoses/q4k-direct-attention.md) §"Open questions". +- **Q4K-direct attention — step-2 gates both GREEN → build greenlit (task #16)** (2026-05-30). Two cheap probes (real production fns, synthetic same-size f32 weights, `CpuBackend`) run *before* building — the #24-trap guard. **Gate 1 (`chris-experiments/larql_probes/examples/q4k_attention/attn_proj_vs_gqa_split.rs`):** the four Q/K/V/O projections are **97–98% of the attention block at the measured context band (cached_len 32–128)** — Q4K-direct addresses ~all of the 28%; projection cost is flat/bandwidth-bound (~40 ms/token blended), GQA grows with cached_len (stays ≥64% to 4K, ~50% at 8K with the W=1024 sliding cap). **Gate 2 (`chris-experiments/larql_probes/examples/q4k_attention/attn_proj_f32_vs_q4k.rs`):** `q4k_matvec` **beats AMX/Accelerate f32 BLAS 2.06× (sliding block) / 2.51× (global block)** — the bandwidth cut nets out, not eaten by AMX throughput. Amdahl: projections ≈ 97%×28% ≈ ~27% of decode → a 2.1–2.5× cut → **~14–16% net decode win from attention alone**. **These are ISOLATED timings** (#24 lesson) — ship gate stays end-to-end net decode tok/s > dense on the real 26B vindex; the isolated evidence just says build, very unlikely to wash. Results appended to [`docs/diagnoses/q4k-direct-attention.md`](docs/diagnoses/q4k-direct-attention.md) §"Step 2". - **Q4K-direct attention — recon map shipped (task #16, step 1)** (2026-05-30). Source-verified map at [`docs/diagnoses/q4k-direct-attention.md`](docs/diagnoses/q4k-direct-attention.md). Verdict: it's a **projection swap, not a function swap** — only the four Q/K/V/O projections are Q4K-accelerable; RoPE/QK-V-norm/GQA-decode-step/KV-concat/residual stay f32 and unchanged (and must, for parity). Three corrections to the prior framing: (1) **the docs name the wrong function** — the steady-state decode 28% is paid in `run_attention_block_decode_step_backend` (per-token, decode.rs:111), NOT `run_attention_with_kv_backend` (that's the *prefill*/TTFT path, gpu.rs:158); both share `dot_proj_gpu` (f32 BLAS), so they're parallel edits — target the decode-step fn for the 28%, the prefill fn for TTFT. (2) **`q4_attention_proj` is the wrong primitive and isn't really tested** — it guards `supports_quant(Q4_K)` but calls the **Q4_0** kernel (`q4_matvec`), and its one test feeds synthetic Q4_0 bytes with no numeric assert; production weights are Q4_K super-block → wrong stride → garbage. Use the existing, parity-tested `CpuBackend::q4k_matvec` / `q4k_dual_matvec` (f32-input, rayon) via `quant_matvec(format, …)` instead, reading bytes from `index.attn_kquant_layer_data(layer)` (already plumbed into `attention_step` but ignored — clean wiring point; also kills the `ensure_attn_tensors_dequantised` load tax). (3) **The 28% is unsplit and context-flattered** — `record_attn` wraps the whole block (projection-vs-GQA share unmeasured), and GQA+KV-concat grow linearly with cached_len and are *not* Q4K-accelerable, so the addressable fraction shrinks with context. **Step-2 gate before any kernel: a projection-vs-GQA split timer** (avoids the #24 build-then-measure trap). Parity baseline confirmed = Q4K-direct vs Q4K-dequant (same bytes, skip the f32 round-trip), per-token NLL/KL distribution + worst-token on the 26B, per-matrix format dispatch (O may be Q6_K). - **WalkFfn speed thread — both axes falsified; capacity win stands; next = Q4K-direct attention** (2026-05-30). A long, disciplined falsification arc (#17–#28, full writeup `docs/diagnoses/walk-ffn-performance.md`, ~10 instrumented examples). **Conclusion: there is no graph-FFN *speed* win on this model — and we proved it cheaply (scripts/probes, not built-and-measured kernels), each falsification cheaper than the last.** @@ -66,7 +115,7 @@ the system of record for stage results. - **What stands:** the **capacity** result — residual-cell `CellRouter` matches gate-KNN accuracy at ~150× cheaper routing (#22, in-dist p=0.0001 vs static), landed + tested in `WalkFfnConfig`/`walk_ffn_sparse`. Graph-FFN = **capacity**, bandwidth fundamentals = **speed**, forced by elimination. Methodology: "the proxy is not the thing" (`feedback_metric_matches_operation`). - **NEXT (task #16): Q4K-direct ATTENTION decode path** — remove the f32 dequant tax on attention weights (~28% of decode); the one speed lever with **no thesis riding on it** (ollama/llama.cpp do it; `q4_attention_proj` half-exists, CPU-tested). *Engineering not research*: the gate is **PARITY (the spine, before any timing number)**, and parity is **Q4K-direct vs Q4K-dequant** (skip the f32 round-trip), NOT vs f32-from-f32 (don't conflate quant error with dequant-tax removal). End-to-end net>dense pre-committed bar (#24; ~28% Amdahl-bounds a 2× proj to ~14% net). **Step 1 = recon, deliver a map** of where the 28% is paid + what `q4_attention_proj` covers *and doesn't* (RoPE/GQA/KV-append/mask, or just projections?). Full framing in task #16. - **CPU remote-MoE decode — closes #146** (2026-05-28): `larql run --moe-shards …` without `--metal` failed with `decode_token_with_moe returned None during prefill` (CPU backend's `decode_token_with_moe` is a GPU-only trait default; CLI always called the GPU path). Routed the CPU branch through the existing `generate_kquant_cpu_remote` + added a clean attn-presence guard in `grid/setup.rs`. Verified end-to-end on the real Gemma-4-26B-A4B vindex (output "Paris"). **Caveat: full-recompute, no KV cache → 0.1–0.4 tok/s.** Follow-up tracked as the new C1 item below. -- **MoE-aware KV engines (C1) — ✅ shipped** (2026-05-28): the KvEngine layer was dense-only; MoE decode now rides it via `RemoteMoeFfn` (`forward_moe_full_layer` = `moe_ffn_block_cpu`) through the MoE-aware `kv_*_via_dispatch` path. Found + fixed a prefill-RoPE bug (engine prefill used unscaled RoPE → garbage on Gemma 4 global layers). KV-cached CPU `--moe-shards` is now the default: **byte-identical to full-recompute, ~10× faster** (4.2 vs 0.4 tok/s on Gemma-4-26B-A4B). `--engine` wired + guarded: **`standard`** (4.4 tok/s) and **`boundary_kv`** (2.9 tok/s; wraps StandardEngine + emits wire-efficient compressed-residual cold-context frames) are MoE-capable; the fused-coarse engines (markov/turbo/unlimited/boundary_per_layer) and apollo error clearly (no remote-expert hook). Rope-scaling regression test added (validated by revert). **7 of 9 KV engines now do remote MoE** (verified "Paris" on 26B, no `--metal`): standard **4.4**, markov_codec/turbo **3.4**, markov/boundary_per_layer **3.1**, boundary_kv **2.9**, unlimited **1.7** tok/s. Done via a shared `engines::layer_ffn_or_moe` helper + `ffn` threading through each engine's larql-kv forward loop — **no `EngineBackend` trait change, no Metal risk** (the fused-coarse-path hook I'd flagged turned out unnecessary: even boundary_per_layer's driver path is a larql-kv walk loop). All seven within ~2.6× and network-bound; `standard` stays the throughput pick. The only exclusions — `no_cache`, `apollo` — are by-design (full/crystal re-forward multiplies round-trips). CLI `--engine` guard allows the seven, rejects the two clearly. Only remaining MoE-correctness gap: `unlimited_context` archived-window replay (long context that evicts windows). [larql-kv ROADMAP](crates/larql-kv/ROADMAP.md) §"MoE-aware KV engines (C1)". +- **MoE-aware KV engines (C1) — ✅ shipped** (2026-05-28): the KvEngine layer was dense-only; MoE decode now rides it via `RemoteMoeFfn` (`forward_moe_full_layer` = `moe_ffn_block_cpu`) through the MoE-aware `kv_*_via_dispatch` path. Found + fixed a prefill-RoPE bug (engine prefill used unscaled RoPE → garbage on Gemma 4 global layers). KV-cached CPU `--moe-shards` is now the default: **byte-identical to full-recompute, ~10× faster** (4.2 vs 0.4 tok/s on Gemma-4-26B-A4B). `--engine` wired + guarded: **`standard`** (4.4 tok/s) and **`boundary_kv`** (2.9 tok/s; wraps StandardEngine + emits wire-efficient compressed-residual cold-context frames) are MoE-capable; the fused-coarse engines (markov/turbo/unlimited/boundary_per_layer) and apollo error clearly (no remote-expert hook). Rope-scaling regression test added (validated by revert). **7 of 9 KV engines now do remote MoE** (verified "Paris" on 26B, no `--metal`): standard **4.4**, markov_codec/turbo **3.4**, markov/boundary_per_layer **3.1**, boundary_kv **2.9**, unlimited **1.7** tok/s. Done via a shared `engines::layer_ffn_or_moe` helper + `ffn` threading through each engine's larql-kv forward loop — **no `EngineBackend` trait change, no Metal risk** (the fused-coarse-path hook I'd flagged turned out unnecessary: even boundary_per_layer's driver path is a larql-kv walk loop). All seven within ~2.6× and network-bound; `standard` stays the throughput pick. The only exclusions — `no_cache`, `apollo` — are by-design (full/crystal re-forward multiplies round-trips). CLI `--engine` guard allows the seven, rejects the two clearly. Only remaining MoE-correctness gap: `windowed_checkpoint` archived-window replay (long context that evicts windows). [larql-kv ROADMAP](crates/larql-kv/ROADMAP.md) §"MoE-aware KV engines (C1)". - **Strategic priorities + Query/Edit/Interpret track** (2026-05-28): two new framing sections in [`ROADMAP.md`](ROADMAP.md) layered on the achievability analysis. (1) Single gated critical path — **V1–V4 is the only true P0**; Engine↔Backend unification, CPU-path-to-blazing, and best-in-class mech-interp are downgraded to "P0-conditional, unblocked by V1–V4". (2) V3 (disk-resident mmap) pulled forward on information-value grounds. (3) GPU = credibility tax, D-PREFILL-MM2 first. (4) MoE-first functionality. (5) Query/Edit/Interpret (`DESCRIBE`/`INSERT`/`walk`/compile) promoted to a co-equal functionality track — the moat, lower-risk than the 100× compound. This rollup's Active Sequence + P0 boundaries below are updated to match. - **Whole-codebase review** (2026-05-28): multi-agent deep review (17 crates, ~415K LOC; per-crate reader + adversarial verification). Clippy clean (2 trivial nits). ~7 verified high/medium hardening items tracked in [`docs/audits/codebase-review-2026-05-28.md`](docs/audits/codebase-review-2026-05-28.md), [`ROADMAP.md`](ROADMAP.md) §"Codebase hardening", and per-crate roadmaps. Top two confirmed by hand: infallible `FfnBackend::forward` aborts serving on remote-shard blips; Metal KV append has no `posRAM frontier MoE would thrash** → long-term disk-residency bet undermined ("bigger RAM" not "spill to disk"). e2e tok/s on a genuinely->RAM model still needs different hardware. Artifacts `v3_granite-30b.json` + `moe-routing/v3moe_locality.json`; writeups `docs/diagnoses/v3-disk-resident-mmap.md` + `moe-routing-locality.md`. | | 4 | V2 FP4 generality | **✅ DONE 2026-05-31 — CONFIRMED** | `larql-vindex` (`fp4_q1_scan`), `larql-inference` (`walk_ffn_v2_fp4_nll`) | ≥99.8% per-feature R<16 on Gemma 3 4B + Granite 3B/8B (reproduces exp 26's 99.83% on gemma3 down exactly; `down` the tail). Predictive E2M1 +0.116 bits/tok vs f32, **beats** the shipped Q4-int baseline. No QAT. KU3 resolved CONFIRMED. (Llama/Mistral/MoE-expert weights need f16 exports — not covered.) Artifacts `v2_*_scan.json`; writeup `docs/diagnoses/v2-fp4-generality.md`. | | 5 | C10 CPU baseline bench | **✅ DISCREPANCY RESOLVED 2026-06-02 — no regression** | `larql-cli`, `bench/` | The 1.50× (05-16) vs 1.93× (05-31) split was **two measurement confounds**: a larql **path mismatch** (27.6 StandardEngine vs 23.6 legacy `bench --cpu`, stable ~12% delta) and a llama.cpp **harness artifact** (45.5 was an unwarmed ollama `num_gpu=0` fluke; warm = 42.8–43.0 = llama-bench 42.99). Reconciled like-for-like (t=8, warm): **larql 23.5 legacy / 26.4 StandardEngine vs llama.cpp 43.0 → ~1.6–1.8×.** Gap is C12 (attn+FFN already on int8 Q8_K SDOT). **Free wins:** `bench --cpu` now shows the StandardEngine row; `--ollama-cpu` gives a true CPU baseline (was Metal-GPU). **26B-A4B baseline ✅ LANDED 2026-06-10:** llama.cpp 32.1 vs larql in-proc 7.1 / 9.7 (`LARQL_Q4K_DIRECT_ATTN=1`) / loopback 7.3 — gap is f32-residency byte traffic (~10 GB/tok vs ~2.1), not the C12 kernel; tier 62%→70%; the 1.8 smoke was a cold artifact. Artifact `bench/baselines/c10_gemma4-26b-a4b_cpu_reconciled.json` (method now mandates AC-power check + drift bracket). 4B artifact `bench/baselines/c10_gemma3-4b_cpu_reconciled.json`. C12 v1 asm kernel landed opt-in (see ROADMAP.md C12). | @@ -131,10 +180,10 @@ only on a WIN, parity-first. | Order | Item | Status | Owner | Exit criterion (the measurement) | |---:|---|---|---|---| -| FR3 | Relation as a clean semantic address | **✅ MEASURED + BUILT 2026-06-07** | `larql-lql`, `larql-vindex` | **Measured:** Gemma-3-4B/N=40, relation synonym-gen **1.00 at every layer L6-L26** (semantic, not lexical; clean from L6); asymmetry stark vs entity 0.07-0.20 until L26. **Built:** `RelationResolver` (`executor/relation_resolver.rs`) — trained residual softmax probe (NOT string/cosine, near-rank-1 proxy trap avoided), model-agnostic probe layer (`round(0.3·num_layers)`), wired into `SELECT … FROM EDGES WHERE relation=…` as a cached semantic fallback. E2E real Gemma-3-4B: `WHERE relation="seat"` → resolved to "capital", returned capital edges. 2 unit + 717 lql tests green, clippy clean. `examples/fr3_relation_address.rs`, [`docs/diagnoses/fr3-relation-address.md`](docs/diagnoses/fr3-relation-address.md). | +| FR3 | Relation as a clean semantic address | **✅ MEASURED + BUILT 2026-06-07** | `larql-lql`, `larql-vindex` | **Measured:** Gemma-3-4B/N=40, relation synonym-gen **1.00 at every layer L6-L26** (semantic, not lexical; clean from L6); asymmetry stark vs entity 0.07-0.20 until L26. **Built:** `RelationResolver` (`executor/relation_resolver.rs`) — trained residual softmax probe (NOT string/cosine, near-rank-1 proxy trap avoided), model-agnostic probe layer (`round(0.3·num_layers)`), wired into `SELECT … FROM EDGES WHERE relation=…` as a cached semantic fallback. E2E real Gemma-3-4B: `WHERE relation="seat"` → resolved to "capital", returned capital edges. 2 unit + 717 lql tests green, clippy clean. `chris-experiments/larql_probes/examples/fleet_routing/fr3_relation_address.rs`, [`docs/diagnoses/fr3-relation-address.md`](docs/diagnoses/fr3-relation-address.md). | | FR3b | Explicit relation rewrite (phrasing-robust fallback) | **✅ MEASURED 2026-06-08 + BUILT 2026-06-09** | `larql-lql`, `larql-inference` | **Measured:** probe is synonym-robust but phrasing-brittle (chance @L10 on a held-out phrasing; more templates = no-op, reverted). Explicit few-shot `word→relation` classify = **12/12** synonyms+phrasings; distractors **2/3** confident-wrong → **`none` escape → 0/3**, 12/12 kept. **Built:** two-tier `resolve_relation_synonym` (Tier 1 probe → Tier 2 `resolve_relation_explicit` on abstain — few-shot+`none` frame, one full forward via `InferenceWeights::predict_dense` since lm_head needs the full vindex not the resolver's `0..=L10`, `none`-gated `match_relation_top1`), opt-in `LARQL_FR3_EXPLICIT`, default off = byte-identical. Real-vindex fix: `relation_labels_ranked` (by feature count) replaces alphabetical top-64 (which dropped `language`/kept `food_animal`). E2E real Gemma-3-4B: `mother tongue`→`language` by explicit (0.97, probe abstained); `weather`→abstain (none-escape); default off → no resolution. Probe stronger than ablation implied (handles `head city`/`legal tender`/`altitude` via Tier 1). 4 new tests, 726 lql lib green, clippy clean. `examples/fr3_{template_ablation,explicit_rewrite}.rs`, [`docs/diagnoses/fr3-explicit-rewrite.md`](docs/diagnoses/fr3-explicit-rewrite.md) §"BUILD LANDED". | -| FR1 | Top-k fuzzy entity router + verifier | **✅ MEASURED + BUILT 2026-06-07** | `larql-vindex`, `larql-inference`, `larql-lql` | **Measured:** Gemma-3-4B/N=150, entity key real & answer-leak-free at L24-26 (L26 top1 **0.89**/top5 0.95, CROSS 1.00, beats E15's MLP, no training); live `query_top1`+0.75 gate fires **150/150** → **11% confident-wrong @L26, 84% @L20**. **Built:** `apply_knn_override_verified` (top-k + entity-in-prompt verify + abstain, resolved-layer-first, opt-in `LARQL_KNN_VERIFY`, default off = byte-identical, 14 legacy + 5 new tests green, clippy clean). E2E real Gemma-3-4B: legacy Germany-paraphrase→SpainX (confident-wrong) → verified→GermanyX (fixed), no regression. `examples/fr1_topk_fuzzy_router.rs`, [`docs/diagnoses/fr1-topk-fuzzy-router.md`](docs/diagnoses/fr1-topk-fuzzy-router.md). LQL `ROUTE TOPK k VERIFY` = follow-up. | -| FR2 | Two-tier symbolic→activation router | **✅ MEASURED + BUILT 2026-06-07** | `larql-inference`, `larql-vindex`, `larql-lql` | **Measured:** symbolic exact-match **0/10** aliases (Persia≠Iran), activation fallback **10/10 top-1** @L24/L26 — E16 reproduced. **Built:** `apply_knn_override_two_tier` (tier-1 FR1 verify → tier-2 activation alias fallback, opt-in `LARQL_KNN_VERIFY`+`LARQL_KNN_FALLBACK`, default off = byte-identical, 4 new + 19 = 23 tests green, clippy clean). E2E real Gemma-3-4B: "capital of Persia" → verify-only abstains (Tehran), two-tier recovers IranX (cos 0.97); named case no regression. Tier-2 = fuzzy ~0.7-0.9 route (fires only when verify missed). `examples/fr2_two_tier_router.rs`, [`docs/diagnoses/fr2-two-tier-router.md`](docs/diagnoses/fr2-two-tier-router.md). | +| FR1 | Top-k fuzzy entity router + verifier | **✅ MEASURED + BUILT 2026-06-07** | `larql-vindex`, `larql-inference`, `larql-lql` | **Measured:** Gemma-3-4B/N=150, entity key real & answer-leak-free at L24-26 (L26 top1 **0.89**/top5 0.95, CROSS 1.00, beats E15's MLP, no training); live `query_top1`+0.75 gate fires **150/150** → **11% confident-wrong @L26, 84% @L20**. **Built:** `apply_knn_override_verified` (top-k + entity-in-prompt verify + abstain, resolved-layer-first, opt-in `LARQL_KNN_VERIFY`, default off = byte-identical, 14 legacy + 5 new tests green, clippy clean). E2E real Gemma-3-4B: legacy Germany-paraphrase→SpainX (confident-wrong) → verified→GermanyX (fixed), no regression. `chris-experiments/larql_probes/examples/fleet_routing/fr1_topk_fuzzy_router.rs`, [`docs/diagnoses/fr1-topk-fuzzy-router.md`](docs/diagnoses/fr1-topk-fuzzy-router.md). LQL `ROUTE TOPK k VERIFY` = follow-up. | +| FR2 | Two-tier symbolic→activation router | **✅ MEASURED + BUILT 2026-06-07** | `larql-inference`, `larql-vindex`, `larql-lql` | **Measured:** symbolic exact-match **0/10** aliases (Persia≠Iran), activation fallback **10/10 top-1** @L24/L26 — E16 reproduced. **Built:** `apply_knn_override_two_tier` (tier-1 FR1 verify → tier-2 activation alias fallback, opt-in `LARQL_KNN_VERIFY`+`LARQL_KNN_FALLBACK`, default off = byte-identical, 4 new + 19 = 23 tests green, clippy clean). E2E real Gemma-3-4B: "capital of Persia" → verify-only abstains (Tehran), two-tier recovers IranX (cos 0.97); named case no regression. Tier-2 = fuzzy ~0.7-0.9 route (fires only when verify missed). `chris-experiments/larql_probes/examples/fleet_routing/fr2_two_tier_router.rs`, [`docs/diagnoses/fr2-two-tier-router.md`](docs/diagnoses/fr2-two-tier-router.md). | | FR4 | Operation-class dispatch boundary | **✅ MEASURED 2026-06-07 — conjecture REFINED** | `larql-lql`, `larql-router`, `larql-vindex` | Ran the real external ops on the E17 rig (`e17_ladder.py external`, +DIST/ARGMIN/PARTITION). **DIST (geometric) + ARGMIN (selection) ride free @L1; only PARTITION (global optimization) walls like parity** (NO-CLEAR 0.81). Parity was NOT a fair stand-in — E4 mis-files geometric/selection (internal). Real line = factors-through-reads vs global-joint. Dispatch: keep count/filter/aggregate/threshold/majority/distance/argmin internal, route global-optimization+parity external. `E17_EXTERNAL_PLAN.md`/`E17_EXTERNAL_VERDICT.md`, `e17_external.json`. | ## Current P0/P1 Boundaries diff --git a/crates/larql-boundary/Cargo.toml b/crates/larql-boundary/Cargo.toml index 266fdbd7b..d90aae7c0 100644 --- a/crates/larql-boundary/Cargo.toml +++ b/crates/larql-boundary/Cargo.toml @@ -18,11 +18,5 @@ criterion = { version = "0.5", features = ["html_reports"] } name = "codec" harness = false -[[example]] -name = "encode_decode" - -[[example]] -name = "gate_decision" - [[example]] name = "accuracy" diff --git a/crates/larql-cli/src/commands/extraction/verify_cmd.rs b/crates/larql-cli/src/commands/extraction/verify_cmd.rs index a8a0412d6..3bf0b9c7e 100644 --- a/crates/larql-cli/src/commands/extraction/verify_cmd.rs +++ b/crates/larql-cli/src/commands/extraction/verify_cmd.rs @@ -13,6 +13,15 @@ pub fn run(args: VerifyArgs) -> Result<(), Box> { return Err(format!("not a directory: {}", args.vindex.display()).into()); } + // A VINDEX2 container is opaque blobs, so verifying it means checksums. + // A VINDEX3 container declares its own structure, so it can answer the + // stronger question — *will this bind?* — without executing anything. + if larql_vindex::format::generation::detect_generation(&args.vindex)? + == larql_vindex::format::generation::ContainerGeneration::V3 + { + return verify_v3(&args.vindex); + } + let config = larql_vindex::load_vindex_config(&args.vindex)?; let stored = match &config.checksums { @@ -62,3 +71,47 @@ pub fn run(args: VerifyArgs) -> Result<(), Box> { Ok(()) } + +/// Structural verification of a VINDEX3 container. +/// +/// Checks what a checksum sweep cannot: that the container's own declarations +/// are mutually consistent and complete enough to bind. Opening it already +/// established that `index.json` parses, the manifest parses and validates, +/// and every declared storage key resolves to a file; `verify` adds segment +/// parsing, programme satisfaction and per-entry region bounds. +/// +/// Execution parity is deliberately excluded — it needs an input and a kernel, +/// and folding it in would make routine verification cost a forward pass. +fn verify_v3(path: &std::path::Path) -> Result<(), Box> { + eprintln!("Verifying: {} (VINDEX3, structural)", path.display()); + let container = larql_vindex::format::vindex3::Vindex3Container::open(path)?; + + println!( + " index.json ......... OK (schema {})", + container.index().version + ); + println!( + " moe_manifest ....... OK ({} MoE layer(s))", + container.manifest().layers.len() + ); + println!( + " storage keys ....... OK ({} segment(s) resolved)", + container.index().segments.len() + ); + + let defects = container.verify(); + if defects.is_empty() { + println!(" structure .......... OK (bindable)"); + println!("\nAll checks passed."); + return Ok(()); + } + println!(" structure .......... {} defect(s)", defects.len()); + for d in &defects { + println!(" - {d}"); + } + Err(format!( + "{} structural defect(s); container is not bindable", + defects.len() + ) + .into()) +} diff --git a/crates/larql-cli/src/commands/extraction/walk_cmd.rs b/crates/larql-cli/src/commands/extraction/walk_cmd.rs index 46a0fb59f..a8f55e52b 100644 --- a/crates/larql-cli/src/commands/extraction/walk_cmd.rs +++ b/crates/larql-cli/src/commands/extraction/walk_cmd.rs @@ -168,6 +168,13 @@ macro_rules! vlog { pub fn run(args: WalkArgs) -> Result<(), Box> { let verbose = args.verbose; + // Validated once, here, because `run` fans out to several forward paths + // and only some of them read the spec. Rejecting an unparseable one up + // front beats warning and running `standard` anyway: a typo'd engine that + // silently exercises the default is the same class of failure as a flag + // dropped entirely (issue #199), and the one a caller is least likely to + // notice. + validate_engine_spec(requested_engine_spec(&args).as_deref())?; let load_start = Instant::now(); // Load the index — either from .vindex or from separate NDJSON files @@ -533,6 +540,15 @@ fn run_predict_q4k( // CPU Q4K autoregressive: per-step, dequantise layer weights // just-in-time (`predict_kquant` does this internally) and loop. // Not token-cached, so O(N²) but correct. For speed use --metal. + // + // This path has no KV cache and therefore no engine to select, so a + // named `--engine` cannot be honoured here. Say so instead of running + // something else under the caller's chosen label: silently dropping + // the flag is what made every engine look identical through + // `larql run` (issue #199), since they were all the same path. + if let Some(spec) = requested_engine_spec(args) { + return Err(engine_unsupported_on_uncached_path(&spec).into()); + } return run_q4k_generate_cpu(weights, tokenizer, &token_ids, args, &index); } @@ -673,6 +689,8 @@ fn run_predict_q4k_remote( ); let start = Instant::now(); + // A refusal from the shards ends the command. Printing predictions built + // without the layer that refused would report a walk the model never ran. let result = larql_inference::vindex::predict_kquant_with_ffn( weights, tokenizer, @@ -680,7 +698,8 @@ fn run_predict_q4k_remote( args.predict_top_k, &index, &remote, - ); + ) + .map_err(|refusal| format!("remote FFN refused ({}): {refusal}", refusal.kind()))?; let elapsed = start.elapsed(); print_predictions("walk (q4k + ffn remote)", &result.predictions, verbose); @@ -740,6 +759,43 @@ fn run_q4k_generate_cpu( Ok(()) } +/// The engine the caller asked for, if any: `--engine` first, then +/// `LARQL_KV_ENGINE`. One resolver so the validation in [`run`], the +/// rejection in [`run_predict_q4k`] and the builder in [`generate_stream`] +/// cannot disagree about whether an engine was requested. +fn requested_engine_spec(args: &WalkArgs) -> Option { + args.engine + .clone() + .or_else(|| std::env::var("LARQL_KV_ENGINE").ok()) +} + +/// Reject a spec no engine answers to. +/// +/// Split out as a pure function so the message is testable without a model. +/// Warning and running `standard` instead — the old behaviour — makes a typo'd +/// engine indistinguishable from the default, which is the same failure as +/// dropping the flag entirely (issue #199). +fn validate_engine_spec(spec: Option<&str>) -> Result<(), String> { + match spec { + Some(s) if larql_kv::EngineKind::from_name(s).is_none() => Err(format!( + "unknown --engine {s:?}; supported: {}", + larql_kv::EngineKind::supported_names().join(", ") + )), + _ => Ok(()), + } +} + +/// The refusal owed to a caller who named an engine on a path with no KV +/// cache to put it in. +fn engine_unsupported_on_uncached_path(spec: &str) -> String { + format!( + "--engine {spec:?} is not honoured on the CPU Q4K generation path, \ + which is not token-cached and so has no KV engine to select. \ + Use --metal for the KV-cached path, or `larql bench --engine` \ + to compare engines." + ) +} + /// Core predict logic shared by model and vindex paths. fn run_predict_inner( weights: &ModelWeights, @@ -1077,10 +1133,7 @@ fn generate_stream( // CLI flag wins over env var; env var wins over `--kv-cache`. See // `crates/larql-inference/docs/specs/kv-engine-unification.md` §6. use larql_kv::EngineKind; - let engine_spec = args - .engine - .clone() - .or_else(|| std::env::var("LARQL_KV_ENGINE").ok()); + let engine_spec = requested_engine_spec(args); let (kind, label) = match engine_spec { Some(spec) => { let kind = EngineKind::from_name(&spec).unwrap_or_else(|| { @@ -1096,7 +1149,7 @@ fn generate_stream( } => "engine=standard (windowed)", EngineKind::NoCache => "engine=no-cache", EngineKind::MarkovResidual { .. } => "engine=markov-rs", - EngineKind::UnlimitedContext { .. } => "engine=unlimited-context", + EngineKind::WindowedCheckpoint { .. } => "engine=windowed-checkpoint", EngineKind::TurboQuant { .. } => "engine=turbo-quant", EngineKind::Apollo { .. } => "engine=apollo", EngineKind::BoundaryKv { .. } => "engine=boundary-kv", @@ -1246,3 +1299,69 @@ fn parse_layer_spec(spec: &str) -> Result, Box } Ok(layers) } + +#[cfg(test)] +mod engine_spec_tests { + //! Issue #199: `--engine` was accepted and silently ignored on the CPU + //! Q4K generation path, so every engine driven through `larql run` + //! exercised the same code. Anyone A/B-ing engines that way would have + //! compared the default against itself. + + use super::{engine_unsupported_on_uncached_path, validate_engine_spec}; + + const UNKNOWN: &str = "not-an-engine"; + + #[test] + fn no_engine_named_is_fine() { + assert!(validate_engine_spec(None).is_ok()); + } + + #[test] + fn every_supported_name_validates() { + // Guards the pairing between the validator and the help text it + // prints: a name advertised as supported must actually parse. + for name in larql_kv::EngineKind::supported_names() { + assert!( + validate_engine_spec(Some(name)).is_ok(), + "{name} is advertised as supported but does not validate" + ); + } + } + + #[test] + fn pre_rename_aliases_still_validate() { + // Scripts and baselines predating the WindowedCheckpoint rename must + // keep working. + for alias in ["unlimited", "windowed-checkpoint", "unlimited_context"] { + assert!(validate_engine_spec(Some(alias)).is_ok(), "{alias}"); + } + } + + #[test] + fn a_parameterised_spec_validates() { + assert!(validate_engine_spec(Some("windowed-checkpoint:window=64")).is_ok()); + } + + #[test] + fn an_unknown_spec_is_rejected_and_says_what_is_supported() { + let err = validate_engine_spec(Some(UNKNOWN)).expect_err("must reject"); + assert!( + err.contains(UNKNOWN), + "the message must name the spec: {err}" + ); + for name in larql_kv::EngineKind::supported_names() { + assert!(err.contains(name), "message omits {name}: {err}"); + } + } + + #[test] + fn the_uncached_path_refusal_names_the_spec_and_a_way_forward() { + // A refusal that does not say what to do instead is a dead end; the + // whole point is that the caller stops getting silent default + // behaviour and learns where the engines are actually comparable. + let msg = engine_unsupported_on_uncached_path("markov-rs"); + assert!(msg.contains("markov-rs"), "{msg}"); + assert!(msg.contains("--metal"), "{msg}"); + assert!(msg.contains("larql bench"), "{msg}"); + } +} diff --git a/crates/larql-cli/src/commands/primary/accuracy_cmd.rs b/crates/larql-cli/src/commands/primary/accuracy_cmd.rs index e7e206325..cba5455b2 100644 --- a/crates/larql-cli/src/commands/primary/accuracy_cmd.rs +++ b/crates/larql-cli/src/commands/primary/accuracy_cmd.rs @@ -44,7 +44,7 @@ pub struct AccuracyArgs { pub model: String, /// Comma-separated KV engine specs (same syntax as `larql bench --engine`). - /// Default: `standard,markov-rs,unlimited-context,turbo-quant,apollo`. + /// Default: `standard,markov-rs,windowed-checkpoint,turbo-quant,apollo`. /// /// Apollo is in the default set as of this slice — its store-miss /// rows surface as `SkippedRetrievalMiss` outcomes with a visible @@ -55,7 +55,7 @@ pub struct AccuracyArgs { /// unable to serve, not silently dropped or mis-attributed. #[arg( long, - default_value = "standard,markov-rs,unlimited-context,turbo-quant,apollo" + default_value = "standard,markov-rs,windowed-checkpoint,turbo-quant,apollo" )] pub engines: String, @@ -486,11 +486,19 @@ fn fmt_served_summary(matches: usize, served: usize, total: usize) -> String { /// Verbose-mode row marker. `✓` / `✗` for served rows (matching today's /// shape); `·` for skipped rows so the eye can scan the column for -/// engine misses. +/// engine misses; `!` for a row the *index* is responsible for. +/// +/// The `!` exists because `ScoreOutcome::FailedBindingDefect` is not a +/// coverage miss — it means the bound artifact violates its own contract, +/// and a run full of them would otherwise read as a low served rate on a +/// working index. Distinguishing it in the column is the cheapest place +/// to notice that the thing to repair is the artifact. fn score_mark(outcome: ScoreOutcome, top1_match: Option) -> &'static str { match (outcome, top1_match) { (ScoreOutcome::Served, Some(true)) => "✓", (ScoreOutcome::Served, _) => "✗", + (ScoreOutcome::FailedBindingDefect, _) => "!", + (ScoreOutcome::FailedStateInvalidated, _) => "!", _ => "·", } } @@ -567,6 +575,12 @@ mod tests { // kv-engine-retrieval-trait-split refactor; "internal error" generalised // into the typed BackendFailure / InvariantViolation split. assert_eq!(score_mark(ScoreOutcome::SkippedBackendFailure, None), "·"); + // A refusal that more residency could rescue is an ordinary skip… + assert_eq!(score_mark(ScoreOutcome::SkippedExecutionRefused, None), "·"); + // …but one that indicts the index must not look like one. + assert_eq!(score_mark(ScoreOutcome::FailedBindingDefect, None), "!"); + // Nor must a row after which the engine stopped being trustworthy. + assert_eq!(score_mark(ScoreOutcome::FailedStateInvalidated, None), "!"); } #[test] diff --git a/crates/larql-cli/src/commands/primary/bench/args.rs b/crates/larql-cli/src/commands/primary/bench/args.rs index 37c1fce58..f76600d34 100644 --- a/crates/larql-cli/src/commands/primary/bench/args.rs +++ b/crates/larql-cli/src/commands/primary/bench/args.rs @@ -53,7 +53,7 @@ pub struct BenchArgs { /// no-cache — full re-forward per step (O(N²)); debug /// markov-rs[:window=N] — residual-stream replacement /// markov-rs-codec[:window=N] — markov-rs with bf16 cold tier (2× cold saving) - /// unlimited-context:window=N — per-window K/V checkpoints + /// windowed-checkpoint:window=N — per-window K/V checkpoints /// turbo-quant[:bits=3|4] — WHT + Lloyd-Max codec; experimental /// apollo:layer=N,coef=F,top_k=K,bos=B — boundary-residual injection; experimental /// boundary-kv:chunk_tokens=N,sequence_id=S — Standard + larql-boundary frame emission @@ -149,7 +149,7 @@ pub struct BenchArgs { /// is safe to set globally. /// /// Note: as of the 2026-05-17 bypass-removal cut, every per-layer - /// engine (`markov-rs`, `markov-rs-codec`, `unlimited-context`, + /// engine (`markov-rs`, `markov-rs-codec`, `windowed-checkpoint`, /// `turbo-quant`, `apollo`, `boundary-per-layer`) always runs its /// own state-policy code regardless of this flag. The fused fast /// path is exclusive to `standard` / `boundary-kv`. This flag now diff --git a/crates/larql-cli/src/commands/primary/bench/engine.rs b/crates/larql-cli/src/commands/primary/bench/engine.rs index 1325600ce..2042b0d17 100644 --- a/crates/larql-cli/src/commands/primary/bench/engine.rs +++ b/crates/larql-cli/src/commands/primary/bench/engine.rs @@ -1,9 +1,10 @@ -//! Pure helpers for the KV-engine bench path (markov-rs, unlimited-context). +//! Pure helpers for the KV-engine bench path (markov-rs, windowed-checkpoint). //! The I/O-bound bench loops live in `engine_runtime.rs`; this file owns: //! * `argmax_token` — greedy next-token pick //! * `format_engine_label` — engine info → label string (with / without Q4K) //! * `EngineSummary` + `summarize_engine_result` — decode-result trim + percentile -//! * `format_kv_memory_note` — "hot=X cold=Y N× vs std-kv" string +//! +//! The `notes` column formatters live in `notes.rs`. //! //! All exercised in this file's tests. @@ -65,24 +66,6 @@ pub(super) fn summarize_engine_result(decode_ms_all: &[f64], warmup: usize) -> E } } -/// Render the memory-footprint note for an engine row: hot, cold, and the -/// compression ratio relative to a Standard KV (FP16) baseline. -/// `total = 0` means we couldn't query the engine; emit a 0× ratio. -pub(super) fn format_kv_memory_note(total: usize, cold: usize, kv_ref: usize) -> String { - let hot = total.saturating_sub(cold); - let ratio = if total > 0 { - kv_ref as f64 / total as f64 - } else { - 0.0 - }; - format!( - "hot={:.1}MB cold={:.1}MB {:.0}× vs std-kv", - hot as f64 / 1_048_576.0, - cold as f64 / 1_048_576.0, - ratio, - ) -} - #[cfg(test)] mod tests { use super::*; @@ -167,28 +150,4 @@ mod tests { assert_eq!(s.n_steps, 5); assert!((s.avg_decode_ms - 10.0).abs() < 1e-9); } - - // ── format_kv_memory_note ──────────────────────────────────────────── - - #[test] - fn kv_memory_note_normal_case() { - // total = 16 MB, cold = 4 MB → hot = 12 MB. Ratio = 64/16 = 4. - let s = format_kv_memory_note(16 * 1024 * 1024, 4 * 1024 * 1024, 64 * 1024 * 1024); - assert!(s.contains("hot=12.0MB")); - assert!(s.contains("cold=4.0MB")); - assert!(s.contains("4× vs std-kv")); - } - - #[test] - fn kv_memory_note_zero_total_emits_zero_ratio() { - let s = format_kv_memory_note(0, 0, 1024); - assert!(s.contains("0× vs std-kv")); - } - - #[test] - fn kv_memory_note_clamps_hot_when_cold_exceeds_total() { - // Engine bug guard: cold > total shouldn't underflow. - let s = format_kv_memory_note(1024, 4096, 0); - assert!(s.contains("hot=0.0MB")); - } } diff --git a/crates/larql-cli/src/commands/primary/bench/engine_runtime.rs b/crates/larql-cli/src/commands/primary/bench/engine_runtime.rs index 354616afb..4845245c8 100644 --- a/crates/larql-cli/src/commands/primary/bench/engine_runtime.rs +++ b/crates/larql-cli/src/commands/primary/bench/engine_runtime.rs @@ -8,9 +8,8 @@ use std::time::Instant; use larql_kv::EngineKind; use super::args::BenchArgs; -use super::engine::{ - argmax_token, format_engine_label, format_kv_memory_note, summarize_engine_result, -}; +use super::engine::{argmax_token, format_engine_label, summarize_engine_result}; +use super::notes::{format_dispatch_note, format_kv_memory_note, format_step_split_note}; use super::row::BenchRow; /// Run the KV-engine bench path for a single engine kind, with or @@ -56,6 +55,10 @@ pub(super) fn run_engine( let is_quant = index.is_some(); + // Read the backend's per-layer delegation answer BEFORE it moves + // into the engine — it decides whether a per-layer row is really a + // CPU measurement, and the engine owns the backend from here on. + let per_layer_is_host_delegated = backend.per_layer_is_host_delegated(); let mut engine = kind.build_with_profiling(backend, args.profile); let info = engine.info(); let label = format_engine_label(&info.name, &info.backend, &info.config, is_quant); @@ -112,7 +115,16 @@ pub(super) fn run_engine( // `prefill_quant`). The router holds `&weights` too, so it's // also scoped to its branch. let max_steps = args.warmup + args.tokens; + // A measured step is the WHOLE token: engine forward + lm_head + + // next-token pick. The reference backend rows (`larql-metal` / + // `larql-cpu`) have always measured a whole token, and both land in + // the same `tok/s` column, so timing only the forward here made every + // engine look 2-3x faster than production for free — on qwen3-0.6b the + // excluded lm_head is ~60% of a short-context step. `head_ms_all` + // keeps the tail separable so the forward-only number is still + // recoverable from the row note. let mut decode_ms_all: Vec = Vec::with_capacity(max_steps); + let mut head_ms_all: Vec = Vec::with_capacity(max_steps); // `--ffn-policy` on the quant path: the router needs `&weights` // for Walk{k} FFN-tensor lookups, but `prefill_quant` needs // `&mut weights` for lazy dequant. Borrows conflict at the call @@ -142,8 +154,12 @@ pub(super) fn run_engine( .prefill_quant(weights, ffn, idx, token_ids, be) .map_err(|e| format!("engine prefill (quant) failed: {e}"))?, }; - let prefill_ms = t_pre.elapsed().as_secs_f64() * 1000.0; + // Seed the first token INSIDE the prefill span: the reference + // backends' `prefill_ms` encloses their first `lm_head_predict` + // (see `layer_graph::generate::cpu`), so excluding it here would + // reintroduce the same asymmetry in the prefill column. let last_token = pick_next(&hidden, weights); + let prefill_ms = t_pre.elapsed().as_secs_f64() * 1000.0; (hidden, prefill_ms, last_token) } else { // Dense path. @@ -163,8 +179,9 @@ pub(super) fn run_engine( let hidden = engine .prefill(weights, ffn, token_ids) .map_err(|e| format!("engine prefill failed: {e}"))?; - let prefill_ms = t_pre.elapsed().as_secs_f64() * 1000.0; + // Same first-token-inside-prefill rule as the quant branch above. let last_token = pick_next(&hidden, weights); + let prefill_ms = t_pre.elapsed().as_secs_f64() * 1000.0; (hidden, prefill_ms, last_token) }; @@ -183,8 +200,12 @@ pub(super) fn run_engine( .decode_step_quant(weights, ffn, idx, last_token, be) .map_err(|e| format!("engine decode_step (quant) failed: {e}"))?, }; - decode_ms_all.push(t.elapsed().as_secs_f64() * 1000.0); + let fwd_ms = t.elapsed().as_secs_f64() * 1000.0; + let t_head = Instant::now(); last_token = pick_next(&hidden, weights); + let head_ms = t_head.elapsed().as_secs_f64() * 1000.0; + decode_ms_all.push(fwd_ms + head_ms); + head_ms_all.push(head_ms); } } else { let weight_ffn = WeightFfn { weights }; @@ -204,8 +225,12 @@ pub(super) fn run_engine( hidden = engine .decode_step(weights, ffn, last_token) .map_err(|e| format!("engine decode_step failed: {e}"))?; - decode_ms_all.push(t.elapsed().as_secs_f64() * 1000.0); + let fwd_ms = t.elapsed().as_secs_f64() * 1000.0; + let t_head = Instant::now(); last_token = pick_next(&hidden, weights); + let head_ms = t_head.elapsed().as_secs_f64() * 1000.0; + decode_ms_all.push(fwd_ms + head_ms); + head_ms_all.push(head_ms); } } @@ -215,7 +240,17 @@ pub(super) fn run_engine( let _ = hidden; let summary = summarize_engine_result(&decode_ms_all, args.warmup); - let note = format_kv_memory_note(engine.memory_bytes(), engine.cold_bytes(), kv_ref_bytes); + let head = summarize_engine_result(&head_ms_all, args.warmup); + // Query the dispatch shape AFTER the run — it is only decided at + // prefill, and it changes what the row means: a per-layer shape on a + // host-delegating backend is a CPU measurement wearing a GPU label. + let dispatch = format_dispatch_note(engine.dispatch_path(), per_layer_is_host_delegated); + let note = format!( + "{}{} {}", + dispatch, + format_step_split_note(summary.avg_decode_ms, head.avg_decode_ms), + format_kv_memory_note(engine.memory_bytes(), engine.cold_bytes(), kv_ref_bytes), + ); if args.verbose { eprintln!( diff --git a/crates/larql-cli/src/commands/primary/bench/mod.rs b/crates/larql-cli/src/commands/primary/bench/mod.rs index 1ec4da4c3..610dba805 100644 --- a/crates/larql-cli/src/commands/primary/bench/mod.rs +++ b/crates/larql-cli/src/commands/primary/bench/mod.rs @@ -28,7 +28,7 @@ //! helpers — pure helpers (wire-list parsing, concurrent aggregation, efficiency). //! run — orchestration entry point. //! local — local Metal/CPU bench (`run_larql`). -//! engine — KV-engine bench (markov-rs / unlimited-context). +//! engine — KV-engine bench (markov-rs / windowed-checkpoint). //! remote_ffn — remote FFN HTTP path + `--concurrent` aggregation. //! remote_moe — remote MoE expert path + `--concurrent` aggregation. //! ollama — Ollama side-by-side comparison. @@ -46,6 +46,7 @@ pub(super) mod grid_lan_runtime; pub(super) mod local; pub(super) mod local_moe_runtime; pub(super) mod local_runtime; +pub(super) mod notes; pub(super) mod ollama; pub(super) mod output; pub(super) mod remote_ffn; diff --git a/crates/larql-cli/src/commands/primary/bench/notes.rs b/crates/larql-cli/src/commands/primary/bench/notes.rs new file mode 100644 index 000000000..682d3e274 --- /dev/null +++ b/crates/larql-cli/src/commands/primary/bench/notes.rs @@ -0,0 +1,179 @@ +//! Row-note formatters for the KV-engine bench table. +//! +//! Each function renders one clause of the `notes` column. They are the +//! fields that say what a row actually measured — which dispatch shape +//! ran, how the step split between forward and lm_head, and how much K/V +//! the engine really holds — so they are kept together and tested +//! together, separately from the numeric summarisation in `engine.rs`. + +/// Render the dispatch-shape prefix for an engine row, e.g. +/// `"[per-layer→host] "`. Empty when the engine does not report a shape +/// (no prefill, or an engine with only one). +/// +/// This is the field that tells a reader whether the row measured the +/// backend they asked for. A windowed engine on `--backends metal` +/// declines the fused path and runs its whole forward on the CPU, while +/// the row label still says `[metal (GPU)]` — worth ~15 ms/token and +/// otherwise only inferable by noticing the timings look wrong. +pub(super) fn format_dispatch_note( + path: Option, + per_layer_is_host_delegated: bool, +) -> String { + match path { + Some(p) => format!("[{}] ", p.describe(per_layer_is_host_delegated)), + None => String::new(), + } +} + +/// Render the memory-footprint note for an engine row: hot, cold, and the +/// compression ratio relative to a Standard KV (FP16) baseline. +/// +/// `total` is what the engine reports owning. An engine that delegates +/// its K/V to a backend-resident cache (the coarse whole-model handle on +/// Metal, for instance) owns nothing itself and reports 0 — that is +/// "not accounted here", NOT "uses no memory", and a ratio computed +/// against it would be unbounded. Say so instead of printing a number: +/// a fabricated `0×` / `239×` in this column is worse than a blank. +pub(super) fn format_kv_memory_note(total: usize, cold: usize, kv_ref: usize) -> String { + if total == 0 { + return "engine-side kv unaccounted (backend-resident)".to_string(); + } + let hot = total.saturating_sub(cold); + let ratio = kv_ref as f64 / total as f64; + // Whole numbers for the big compression wins, one decimal below the + // threshold — see RATIO_DECIMAL_THRESHOLD for why that matters. + let ratio_str = if ratio < RATIO_DECIMAL_THRESHOLD { + format!("{ratio:.1}×") + } else { + format!("{ratio:.0}×") + }; + format!( + "hot={:.1}MB cold={:.1}MB {ratio_str} vs std-kv", + hot as f64 / BYTES_PER_MIB, + cold as f64 / BYTES_PER_MIB, + ) +} + +/// Bytes in the MB unit this table prints (mebibytes, matching how the +/// rest of the tooling reports model and cache sizes). +const BYTES_PER_MIB: f64 = 1024.0 * 1024.0; + +/// Below this the ratio gets a decimal place; at or above it, none. +/// A genuine 0.5x (an engine using twice the reference) must not print +/// as "0x", which reads as a missing measurement rather than a loss. +const RATIO_DECIMAL_THRESHOLD: f64 = 10.0; + +/// Render the per-step split for an engine row: how much of the measured +/// step was the forward, and how much the lm_head + next-token pick. +/// +/// Both halves are inside the timed step (see `engine_runtime`), so +/// `fwd = step - head`. Reporting the split keeps the forward-only +/// number recoverable without making `tok_per_s` incomparable against +/// the reference backend rows, which have always measured a whole token. +pub(super) fn format_step_split_note(step_ms: f64, head_ms: f64) -> String { + let fwd_ms = (step_ms - head_ms).max(0.0); + format!("fwd={fwd_ms:.2}ms head={head_ms:.2}ms") +} + +#[cfg(test)] +mod tests { + use super::*; + use larql_inference::kv_engine::DispatchPath; + + // ── format_kv_memory_note ──────────────────────────────────────────── + + #[test] + fn kv_memory_note_normal_case() { + // total = 16 MB, cold = 4 MB → hot = 12 MB. Ratio = 64/16 = 4. + let s = format_kv_memory_note(16 * 1024 * 1024, 4 * 1024 * 1024, 64 * 1024 * 1024); + assert!(s.contains("hot=12.0MB")); + assert!(s.contains("cold=4.0MB")); + assert!(s.contains("4.0× vs std-kv")); + } + + #[test] + fn kv_memory_note_sub_unit_ratio_is_not_rounded_to_zero() { + // An engine using 2× the reference is a real 0.5×; printing "0×" + // made a loss indistinguishable from a missing measurement. + let s = format_kv_memory_note(112 * 1024 * 1024, 0, 56 * 1024 * 1024); + assert!(s.contains("0.5× vs std-kv"), "got {s:?}"); + } + + #[test] + fn kv_memory_note_large_ratio_drops_the_decimal() { + let s = format_kv_memory_note(1024 * 1024, 0, 240 * 1024 * 1024); + assert!(s.contains("240× vs std-kv"), "got {s:?}"); + } + + #[test] + fn kv_memory_note_zero_total_declines_to_invent_a_ratio() { + // An engine that delegates K/V to a backend-resident cache reports + // 0 bytes owned. The old code printed "0× vs std-kv", which reads + // as a measurement; it isn't one. + let s = format_kv_memory_note(0, 0, 1024); + assert!(s.contains("unaccounted"), "got {s:?}"); + assert!(!s.contains("×"), "must not fabricate a ratio: {s:?}"); + assert!( + !s.contains("hot="), + "must not imply a hot-tier reading: {s:?}" + ); + } + + #[test] + fn kv_memory_note_clamps_hot_when_cold_exceeds_total() { + // Engine bug guard: cold > total shouldn't underflow. + let s = format_kv_memory_note(1024, 4096, 0); + assert!(s.contains("hot=0.0MB")); + } + + // ── format_dispatch_note ───────────────────────────────────────────── + + #[test] + fn dispatch_note_is_empty_when_the_engine_reports_no_shape() { + assert_eq!(format_dispatch_note(None, true), ""); + assert_eq!(format_dispatch_note(None, false), ""); + } + + #[test] + fn dispatch_note_marks_host_delegated_per_layer() { + assert_eq!( + format_dispatch_note(Some(DispatchPath::PerLayer), true), + "[per-layer→host] " + ); + } + + #[test] + fn dispatch_note_leaves_native_per_layer_unmarked() { + assert_eq!( + format_dispatch_note(Some(DispatchPath::PerLayer), false), + "[per-layer] " + ); + } + + #[test] + fn dispatch_note_never_marks_coarse_as_host() { + // Coarse doesn't touch the per-layer surface, so the backend's + // delegation answer must not leak into its label. + assert_eq!( + format_dispatch_note(Some(DispatchPath::Coarse), true), + "[coarse] " + ); + } + + // ── format_step_split_note ─────────────────────────────────────────── + + #[test] + fn step_split_note_reports_fwd_as_step_minus_head() { + let s = format_step_split_note(10.0, 6.5); + assert!(s.contains("fwd=3.50ms"), "got {s:?}"); + assert!(s.contains("head=6.50ms"), "got {s:?}"); + } + + #[test] + fn step_split_note_clamps_negative_fwd() { + // Timer noise could make head marginally exceed the step total; + // clamp rather than print a negative forward cost. + let s = format_step_split_note(1.0, 1.2); + assert!(s.contains("fwd=0.00ms"), "got {s:?}"); + } +} diff --git a/crates/larql-cli/src/commands/primary/k3_ledger/classes.rs b/crates/larql-cli/src/commands/primary/k3_ledger/classes.rs index 9a7edb2d7..5791ea5d4 100644 --- a/crates/larql-cli/src/commands/primary/k3_ledger/classes.rs +++ b/crates/larql-cli/src/commands/primary/k3_ledger/classes.rs @@ -260,10 +260,10 @@ impl ClassRow { /// dispatch grid's y axis — MEASURED by K3a, and inside the large-shape band. pub const GROUPED_ROUTED_ETA: f64 = 0.89; -/// MXFP4 all-in bits (codeword + e8m0 scale per 32). -pub const MXFP4_BITS: f64 = 4.25; -/// Q6_K all-in bits (210 bytes per 256 weights). -pub const Q6K_BITS: f64 = 210.0 * 8.0 / 256.0; +// All-in bits per container are NOT restated here. `Container::all_in_bits` +// derives them from block geometry and is the single authority; a second copy +// carrying the same literals is how the two drift apart, and the ledger's whole +// claim is that its numbers are derived rather than typed. /// The per-class census for a K3 image at a given format. /// @@ -517,9 +517,11 @@ pub fn apply(rows: &[ClassRow], s: Scenario) -> Vec { mod tests { use super::*; use crate::commands::primary::k3_ledger::geometry::k3_reference; + use crate::commands::primary::k3_ledger::serving_format::Container; fn base() -> Vec { - census(&k3_reference(), MXFP4_BITS, MXFP4_BITS) + let bits = Container::Mxfp4.all_in_bits(); + census(&k3_reference(), bits, bits) } fn approx(a: f64, b: f64, tol: f64) { @@ -631,10 +633,11 @@ mod tests { #[test] fn q6k_transcode_costs_bytes_and_therefore_throughput() { let mxfp4 = compose(base(), BW_GB_S, 0.85); - let q6k = compose(census(&k3_reference(), Q6K_BITS, Q6K_BITS), BW_GB_S, 0.85); + let q6k_bits = Container::Q6K.all_in_bits(); + let q6k = compose(census(&k3_reference(), q6k_bits, q6k_bits), BW_GB_S, 0.85); assert!(q6k.total_bytes > mxfp4.total_bytes); assert!(q6k.tok_s < mxfp4.tok_s); - approx(Q6K_BITS / MXFP4_BITS, 1.544, 0.01); + approx(q6k_bits / Container::Mxfp4.all_in_bits(), 1.544, 0.01); } #[test] diff --git a/crates/larql-cli/src/commands/primary/k3_ledger/fetch.rs b/crates/larql-cli/src/commands/primary/k3_ledger/fetch.rs index 8aab4b8d5..092827d6e 100644 --- a/crates/larql-cli/src/commands/primary/k3_ledger/fetch.rs +++ b/crates/larql-cli/src/commands/primary/k3_ledger/fetch.rs @@ -155,7 +155,7 @@ fn tensor_bytes(entry: &Value) -> u64 { entry .get("data_offsets") .and_then(|o| o.as_array()) - .and_then(|a| Some(a.get(1)?.as_u64()? - a.get(0)?.as_u64()?)) + .and_then(|a| Some(a.get(1)?.as_u64()? - a.first()?.as_u64()?)) .unwrap_or(0) } diff --git a/crates/larql-cli/src/commands/primary/k3_ledger/selection_trace.rs b/crates/larql-cli/src/commands/primary/k3_ledger/selection_trace.rs index 901c8bdbf..09a14b197 100644 --- a/crates/larql-cli/src/commands/primary/k3_ledger/selection_trace.rs +++ b/crates/larql-cli/src/commands/primary/k3_ledger/selection_trace.rs @@ -26,6 +26,14 @@ pub enum SelectionUnit { /// An MoE expert (DEC-8.4). Expert, /// An FFN feature row inside one expert (DEC-8.1). + /// + /// Nothing constructs this yet — only the expert path is wired, so + /// `from_routing_pool` always labels `Expert`. Kept rather than deleted + /// because the unit is the thing that stops an expert distribution being + /// read as a feature one, and a DEC-8.1 export arriving with no way to + /// name its own bank is exactly the substitution this enum exists to + /// prevent. Pinned by `every_selection_unit_serialises_under_a_distinct_name`. + #[allow(dead_code)] Feature, /// Unlabelled — the estimators do not care, but an export should. Symbol, @@ -423,4 +431,21 @@ mod tests { (2, 2, 2, 4) ); } + + #[test] + fn every_selection_unit_serialises_under_a_distinct_name() { + // The names reach exported files, and the whole point of the enum is + // that a reader can tell an expert distribution from a feature one. + // Two units sharing a name would silently permit the substitution. + let units = [ + SelectionUnit::Expert, + SelectionUnit::Feature, + SelectionUnit::Symbol, + ]; + let mut names: Vec<&str> = units.iter().map(|u| u.as_str()).collect(); + let count = names.len(); + names.sort_unstable(); + names.dedup(); + assert_eq!(names.len(), count, "unit names must be distinct"); + } } diff --git a/crates/larql-cli/src/commands/primary/k3_ledger/serving_format.rs b/crates/larql-cli/src/commands/primary/k3_ledger/serving_format.rs index e2e1bd1f4..bc5c8edc1 100644 --- a/crates/larql-cli/src/commands/primary/k3_ledger/serving_format.rs +++ b/crates/larql-cli/src/commands/primary/k3_ledger/serving_format.rs @@ -132,6 +132,14 @@ pub enum KernelMaturity { /// No Metal kernel at all. None, /// A standalone or diagnostic kernel exists. + /// + /// No container sits here today — MXFP4, the one that would, has already + /// reached `Grouped`. The rung is kept because it is what separates "a + /// kernel exists" from "a kernel serves", and collapsing the ladder back + /// towards the coarse "has a kernel" question is precisely what produced + /// the wrong MXFP4 claim recorded above. Ordering is pinned by + /// `kernel_maturity_ladder_ascends_and_servability_splits_it_once`. + #[allow(dead_code)] Standalone, /// A grouped-expert kernel exists (may still be a candidate, not a path). Grouped, @@ -539,4 +547,27 @@ mod tests { .collect(); assert!(bits.windows(2).all(|p| p[0] <= p[1]), "{bits:?}"); } + + #[test] + fn kernel_maturity_ladder_ascends_and_servability_splits_it_once() { + // `is_servable` is a threshold on the ordering, so the ordering is the + // load-bearing part: a rung inserted out of order would silently move + // the servable boundary without touching `is_servable` itself. + let ladder = [ + KernelMaturity::None, + KernelMaturity::Standalone, + KernelMaturity::Grouped, + KernelMaturity::Dispatched, + KernelMaturity::Production, + ]; + assert!( + ladder.windows(2).all(|p| p[0] < p[1]), + "the ladder must be declared in ascending maturity: {ladder:?}" + ); + let servable: Vec = ladder.iter().map(|m| m.is_servable()).collect(); + // Monotone: once servable, always servable further up the ladder. + assert!(servable.windows(2).all(|p| p[0] <= p[1]), "{servable:?}"); + assert!(!KernelMaturity::Standalone.is_servable()); + assert!(KernelMaturity::Dispatched.is_servable()); + } } diff --git a/crates/larql-cli/src/commands/primary/k3_ledger/transcode.rs b/crates/larql-cli/src/commands/primary/k3_ledger/transcode.rs index f00e45680..23dd58ad7 100644 --- a/crates/larql-cli/src/commands/primary/k3_ledger/transcode.rs +++ b/crates/larql-cli/src/commands/primary/k3_ledger/transcode.rs @@ -283,7 +283,7 @@ mod tests { #[test] fn merge_takes_the_worst_case_across_tensors() { - let good = scan_scales(&vec![127u8; GROUPS_PER_SUPERBLOCK], GROUPS_PER_SUPERBLOCK); + let good = scan_scales(&[127u8; GROUPS_PER_SUPERBLOCK], GROUPS_PER_SUPERBLOCK); let mut bad_scales = vec![120u8; GROUPS_PER_SUPERBLOCK]; bad_scales[1] = 130; let bad = scan_scales(&bad_scales, GROUPS_PER_SUPERBLOCK); diff --git a/crates/larql-cli/src/commands/primary/run_cmd.rs b/crates/larql-cli/src/commands/primary/run_cmd.rs index 40152de78..100fda41c 100644 --- a/crates/larql-cli/src/commands/primary/run_cmd.rs +++ b/crates/larql-cli/src/commands/primary/run_cmd.rs @@ -38,7 +38,7 @@ use crate::commands::primary::cache; /// | `none` | `NoCache` | /// /// New callers should prefer `--engine SPEC` / `LARQL_KV_ENGINE` instead -/// — they accept the full engine catalog (MarkovResidual, UnlimitedContext, +/// — they accept the full engine catalog (MarkovResidual, WindowedCheckpoint, /// TurboQuant, Apollo) not just the three legacy cache strategies. /// See `crates/larql-inference/docs/specs/kv-engine-unification.md` §6. #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -100,7 +100,7 @@ pub struct RunArgs { /// /// Each value maps to an `EngineKind` internally (see `KvCacheKind` /// docs). For the full engine catalog (MarkovResidual, - /// UnlimitedContext, TurboQuant, Apollo), use `--engine` instead. + /// WindowedCheckpoint, TurboQuant, Apollo), use `--engine` instead. #[arg(long, default_value = "standard", value_parser = parse_kv_cache)] pub kv_cache: KvCacheKind, @@ -116,7 +116,7 @@ pub struct RunArgs { /// standard:window=1024 — sliding-window K/V /// no-cache — full re-forward per step (O(N²)) /// markov-rs[:window=N] — residual-stream replacement - /// unlimited-context:window=N — per-window K/V checkpoints + /// windowed-checkpoint:window=N — per-window K/V checkpoints /// turbo-quant[:bits=3|4] — WHT + Lloyd-Max codec /// apollo:layer=N,coef=F,top_k=K,bos=B — boundary-residual injection (bench-only) /// @@ -568,7 +568,7 @@ fn run_with_moe_shards( // `ffn` trait (where `RemoteMoeFfn` hooks the experts): `standard` and // `boundary_kv` (which wraps a StandardEngine and adds compressed-residual // boundary frames — same dispatch, wire-efficient cold-context). The - // compression engines (markov_residual / turbo_quant / unlimited_context / + // compression engines (markov_residual / turbo_quant / windowed_checkpoint / // boundary_per_layer) route FFN through the backend's fused coarse path, and // apollo/no_cache re-forward — none have a remote-expert hook, so they'd // silently drop experts. Reject them clearly. @@ -579,7 +579,7 @@ fn run_with_moe_shards( kind, larql_kv::EngineKind::Standard { .. } | larql_kv::EngineKind::BoundaryKv { .. } - | larql_kv::EngineKind::UnlimitedContext { .. } + | larql_kv::EngineKind::WindowedCheckpoint { .. } | larql_kv::EngineKind::MarkovResidual { .. } | larql_kv::EngineKind::MarkovResidualCodec { .. } | larql_kv::EngineKind::TurboQuant { .. } @@ -587,7 +587,7 @@ fn run_with_moe_shards( ) { return Err(format!( "`--engine {}` is not supported with remote MoE (--moe-shards). Supported: \ - standard, boundary, unlimited-context, markov-rs, markov-residual-codec, \ + standard, boundary, windowed-checkpoint, markov-rs, markov-residual-codec, \ turbo-quant, boundary-per-layer (they dispatch FFN per-layer through the ffn \ trait where experts hook in). `no-cache` / `apollo` re-forward and would \ multiply expert round-trips. See larql-kv ROADMAP §\"MoE-aware KV engines (C1)\".", diff --git a/crates/larql-cli/src/commands/primary/show_cmd.rs b/crates/larql-cli/src/commands/primary/show_cmd.rs index c2e05c851..16e266b00 100644 --- a/crates/larql-cli/src/commands/primary/show_cmd.rs +++ b/crates/larql-cli/src/commands/primary/show_cmd.rs @@ -16,14 +16,25 @@ pub struct ShowArgs { pub fn run(args: ShowArgs) -> Result<(), Box> { let path = cache::resolve_model(&args.model)?; - let cfg = larql_vindex::load_vindex_config(&path)?; println!("Model: {}", args.model); println!("Path: {}", path.display()); - println!("Layers: {}", cfg.num_layers); - println!("Hidden: {}", cfg.hidden_size); - println!("Dtype: {:?}", cfg.dtype); - println!("Quant: {:?}", cfg.quant); + + // Dispatch on the container's own discriminator, then let each generation + // describe itself in its own vocabulary. Normalising VINDEX3 back into a + // VINDEX2-shaped summary would drop exactly what a VINDEX3 user needs to + // see — programme id, storage key, manifest validity. + match larql_vindex::format::generation::detect_generation(&path)? { + larql_vindex::format::generation::ContainerGeneration::V3 => show_v3(&path)?, + larql_vindex::format::generation::ContainerGeneration::V2 => { + let cfg = larql_vindex::load_vindex_config(&path)?; + println!("Generation: VINDEX2"); + println!("Layers: {}", cfg.num_layers); + println!("Hidden: {}", cfg.hidden_size); + println!("Dtype: {:?}", cfg.dtype); + println!("Quant: {:?}", cfg.quant); + } + } println!("\nFiles:"); let mut entries: Vec<_> = std::fs::read_dir(&path)? @@ -39,6 +50,78 @@ pub fn run(args: ShowArgs) -> Result<(), Box> { Ok(()) } +/// Describe a VINDEX3 container in its own terms. +/// +/// The per-layer lines are the ones with no VINDEX2 equivalent: which +/// programme interprets the bank, and which storage key it resolves to. Those +/// are what a binding failure is diagnosed from, so `show` is where they +/// belong. +fn show_v3(path: &std::path::Path) -> Result<(), Box> { + let container = larql_vindex::format::vindex3::Vindex3Container::open(path)?; + let index = container.index(); + println!("Generation: VINDEX3 (index.json schema {})", index.version); + println!("Layers: {}", index.num_layers); + println!("Hidden: {}", index.hidden_size); + println!("Family: {}", index.family); + println!("Profiles: {}", index.profile_names().join(", ")); + println!("Manifest: {}", index.moe_manifest); + + // §9.1: show what each profile actually selects, not just that it exists. + // A container that carries alternative packs is exactly the one where + // "which bytes does this profile run" stops being obvious. + if !index.variants.is_empty() { + println!("\nVariants:"); + for region_set in index.variants.region_sets() { + let set = index + .variants + .get(®ion_set) + .expect("region_sets() lists catalogued keys"); + println!( + " {region_set}: {} (baseline {})", + set.present().join(", "), + set.baseline + ); + } + for name in index.profile_names() { + // `open` proved every declared profile resolves. + let resolved = container.select_profile(name)?; + println!("\nProfile '{name}' selects:"); + for (region_set, stored) in resolved.entries() { + println!( + " {region_set} -> {} [{}]", + stored.storage, + stored.fidelity.name() + ); + } + } + } + + println!("\nMoE layers:"); + for layer in &container.manifest().layers { + let bank = &layer.routed_bank; + let known = if bank.resolve_programme().is_some() { + "" + } else { + " [programme not implemented by this binary]" + }; + println!( + " layer {:<3} {:<24} experts {:<4} storage {}{}", + layer.layer, bank.programme, bank.experts, bank.storage, known + ); + } + + let defects = container.verify(); + if defects.is_empty() { + println!("\nStructure: bindable (no defects)"); + } else { + println!("\nStructure: {} defect(s)", defects.len()); + for d in &defects { + println!(" - {d}"); + } + } + Ok(()) +} + fn human_size(bytes: u64) -> String { const K: u64 = 1024; const M: u64 = K * 1024; diff --git a/crates/larql-compute-metal/src/backend/mod.rs b/crates/larql-compute-metal/src/backend/mod.rs index 8eb6634dc..8bfb335ae 100644 --- a/crates/larql-compute-metal/src/backend/mod.rs +++ b/crates/larql-compute-metal/src/backend/mod.rs @@ -91,6 +91,16 @@ pub struct MetalBackend { // `FfnKernels` (the `ffn` field).) /// KV cache for decode mode — initialized on first decode_token call. pub(crate) kv_cache: std::sync::Mutex>, + /// Engine-requested sliding window for the sequence currently being + /// decoded, or `NO_ENGINE_WINDOW` when unwindowed. + /// + /// Distinct from the architecture's own per-layer SWA, which comes + /// from the layer spec; the effective window is the narrower of the + /// two. Carried on the backend rather than threaded through + /// `build_arch_params` because it belongs to the *sequence*, not the + /// architecture — every call site that builds a layer spec would + /// otherwise have to learn about a caller's decode policy. + pub(crate) engine_window: std::sync::atomic::AtomicUsize, /// Pre-allocated MoE scratch for `decode_token_q4k_moe` — keyed /// by `(top_k, hidden, intermediate_size)`. Reused across decode /// calls so the ~15 buffer allocations (~120ms on Gemma 4 26B-A4B, @@ -271,6 +281,7 @@ impl MetalBackend { attention, ffn, kv_cache: std::sync::Mutex::new(None), + engine_window: std::sync::atomic::AtomicUsize::new(NO_ENGINE_WINDOW), moe_scratch: std::sync::Mutex::new(None), ple_inputs: std::sync::Mutex::new(None), f32_gemv_pipeline, @@ -410,6 +421,42 @@ impl PleInputBuffer { } } +impl MetalBackend { + /// Set the engine-requested window for the sequence being decoded. + /// `None` clears it. See [`MetalBackend::effective_window_for`]. + pub(crate) fn set_engine_window(&self, window: Option) { + self.engine_window.store( + window.unwrap_or(NO_ENGINE_WINDOW), + std::sync::atomic::Ordering::Relaxed, + ); + } + + /// Narrow a layer's architectural window by the engine's, if any. + /// + /// Both use `0` for "unbounded", so this is a min over the non-zero + /// values. The narrower wins: an engine promising a 256-token window + /// must not attend further just because the architecture allows + /// 1024, and a sliding layer must not attend further just because + /// the engine is unwindowed. + pub(crate) fn effective_window_for(&self, arch_window: u32) -> u32 { + let engine = self + .engine_window + .load(std::sync::atomic::Ordering::Relaxed) as u32; + match (arch_window, engine) { + (0, e) => e, + (a, 0) => a, + (a, e) => a.min(e), + } + } +} + +/// `engine_window` value meaning "no engine-imposed window". +/// +/// Zero is the same sentinel the kernel uses for `window_size`, so the +/// backend-side and shader-side notions of "unbounded" agree without a +/// translation step. +pub(crate) const NO_ENGINE_WINDOW: usize = 0; + #[cfg(test)] mod tests { use super::*; diff --git a/crates/larql-compute-metal/src/decode/encode_attn.rs b/crates/larql-compute-metal/src/decode/encode_attn.rs index 2da5ddacf..d5be840b6 100644 --- a/crates/larql-compute-metal/src/decode/encode_attn.rs +++ b/crates/larql-compute-metal/src/decode/encode_attn.rs @@ -113,7 +113,11 @@ impl MetalBackend { } else { layer_head_dim }; - let window_size = attn_spec.sliding_window as u32; + // Narrower of the architecture's per-layer SWA and any window the + // engine imposed on this sequence. The kernel attends + // `[T - window_size, T)`, so this both bounds attention and lets + // the cache hold more rows than the window between compactions. + let window_size = self.effective_window_for(attn_spec.sliding_window as u32); // Env flags governing kernel-level fusion. Cached at backend // startup (see `metal::flags::DecodeFlags`) so the decode hot @@ -136,19 +140,24 @@ impl MetalBackend { // block for this token by the time we reach the shared layer). let kv_shared_source = layer.kv_shared_source; let attend_cache_idx = kv_shared_source.unwrap_or(layer_idx); + // `pos` is the ABSOLUTE stream position to RoPE at; `t_val` is how + // many cached rows to attend. They are equal-and-offset only while + // nothing has been evicted — once a window slides, occupancy falls + // and position keeps climbing, so they must be read from different + // fields. Deriving `t_val` from `pos` (as `pos + 1`) is what tied + // them together before. let pos = if let Some(src) = kv_shared_source { - // Source has already incremented its current_len for this token. - // Position to RoPE is the same as the source's last-written index. - (kv_cache.layers[src].current_len.saturating_sub(1)) as u32 + // Source has already advanced past the row it just wrote; + // RoPE at that row's position. + (kv_cache.layers[src].abs_position.saturating_sub(1)) as u32 } else { - kv_cache.layers[layer_idx].current_len as u32 + kv_cache.layers[layer_idx].abs_position as u32 }; let t_val = if kv_shared_source.is_some() { - // Source's current_len already counts this token; t_val is the - // total positions to attend over (= source.current_len). + // Source's current_len already counts this token. kv_cache.layers[attend_cache_idx].current_len as u32 } else { - pos + 1 + (kv_cache.layers[layer_idx].current_len + 1) as u32 }; let attn_span = ops::kv_cache::attention_span(t_val, window_size); @@ -241,7 +250,7 @@ impl MetalBackend { MTLSize::new(layer_num_q_heads as u64, 1, 1), MTLSize::new(tg_w, 1, 1), ); - kv_cache.layers[layer_idx].current_len += 1; + kv_cache.layers[layer_idx].advance_one(); } else if use_fused_qkn_rope && layer.q_norm_weight.is_some() && layer.k_norm_weight.is_some() @@ -452,7 +461,7 @@ impl MetalBackend { // Only own-cache layers advance current_len; shared layers leave // their (unused) cache pointer at 0 forever. if !did_fused_attn && kv_shared_source.is_none() { - kv_cache.layers[layer_idx].current_len += 1; + kv_cache.layers[layer_idx].advance_one(); } // ── Step 5a: O projection ── diff --git a/crates/larql-compute-metal/src/decode_hybrid.rs b/crates/larql-compute-metal/src/decode_hybrid.rs index 8eb622aa2..5b9ad5a37 100644 --- a/crates/larql-compute-metal/src/decode_hybrid.rs +++ b/crates/larql-compute-metal/src/decode_hybrid.rs @@ -193,7 +193,9 @@ impl MetalBackend { // RoPE { - let pos = kv_cache.layers[layer_idx].current_len as u32; + // ABSOLUTE stream position for RoPE — not occupancy, which a + // sliding window reduces. + let pos = kv_cache.layers[layer_idx].abs_position as u32; let hd = layer_head_dim as u32; let rdim = layer_rotary_dim as u32; let rope_pairs = (layer_rotary_dim / 2) as u64; @@ -279,7 +281,7 @@ impl MetalBackend { ); enc_b.end_encoding(); } - kv_cache.layers[layer_idx].current_len += 1; + kv_cache.layers[layer_idx].advance_one(); // ═══════════════════════════════════════════════════════════ // ENCODER C: O projection → residual add (post-attention) diff --git a/crates/larql-compute-metal/src/kv_dispatch_impl.rs b/crates/larql-compute-metal/src/kv_dispatch_impl.rs index 97fc34fa1..a6994e7d9 100644 --- a/crates/larql-compute-metal/src/kv_dispatch_impl.rs +++ b/crates/larql-compute-metal/src/kv_dispatch_impl.rs @@ -28,6 +28,10 @@ use larql_models::ModelWeights; /// Zero-sized type; const-construction is free. const CPU: CpuBackend = CpuBackend; +/// K and V — the two cached tensors every attention layer stores per +/// position. Named so the K/V sizing arithmetic doesn't read as a bare 2. +const KV_TENSORS_PER_LAYER: usize = 2; + impl KvDispatch for MetalBackend { fn alloc_kv_buffer(&self, layer: usize, max_tokens: usize, kv_dim: usize) -> KvHandle { // Handles are CPU-resident at Step 4. When real Metal kernels land @@ -394,6 +398,100 @@ impl KvDispatch for MetalBackend { Some(hidden) } + /// Coarse prefill under an engine window. + /// + /// Accepts only a prompt that fits inside the window, matching the + /// CPU rule: the fused prefill has no per-query-position masking, so + /// a longer prompt would attend in full while the engine advertises + /// a bound. Declining sends the engine to the per-layer path, which + /// is correct — just slower. + fn coarse_prefill_windowed( + &self, + weights: &ModelWeights, + token_ids: &[u32], + index: Option<&dyn larql_compute::KvIndex>, + window: Option, + ) -> Option<(Array2, KvHandle)> { + if let Some(w) = window { + if w == 0 || token_ids.len() > w { + return None; + } + } + self.set_engine_window(window); + self.coarse_prefill(weights, token_ids, index) + } + + /// Coarse decode under an engine window. + /// + /// Two mechanisms, because one cannot do both jobs: + /// + /// - **Attention** is bounded by `window` every step, via the layer + /// spec's window (see `effective_window_for`). The kernel attends + /// `[T - window, T)`, so extra resident rows are simply not read. + /// - **Memory** is bounded by compaction, run only when occupancy + /// reaches `COMPACTION_SLACK x window`. Compacting every step would + /// memmove the whole window per token; amortised it is O(1) per + /// token, at the cost of holding up to that multiple of the window. + /// + /// Doing only the compaction would let attention read up to the slack + /// multiple of the window between compactions; doing only the span + /// clamp would never reclaim memory. Both, or neither contract holds. + fn coarse_decode_step_windowed( + &self, + weights: &ModelWeights, + token_id: u32, + index: Option<&dyn larql_compute::KvIndex>, + handle: &mut KvHandle, + abs_position: usize, + window: Option, + ) -> Option> { + let Some(w) = window else { + self.set_engine_window(None); + return self.coarse_decode_step(weights, token_id, index, handle, abs_position); + }; + if w == 0 { + return None; + } + self.set_engine_window(Some(w)); + self.compact_kv_to_window(w); + self.coarse_decode_step(weights, token_id, index, handle, abs_position) + } + + fn per_layer_is_host_delegated(&self) -> bool { + // Every per-layer method above forwards to `CPU`. Only the + // `coarse_*` family runs Metal kernels. Until the per-layer + // surface has native implementations this must stay `true`, or + // diagnostics will keep reporting CPU work as GPU work. + true + } + + fn backend_resident_kv_bytes(&self) -> usize { + // The coarse pipeline's K/V lives here, not in the handle + // (`MetalCoarseHandle` is a sentinel), so an engine that only + // measures its handles reports zero on this path. Count the + // populated prefix of each layer — `current_len`, not `max_seq`: + // the buffers are preallocated to the context ceiling and + // charging an engine for capacity it has not filled would + // overstate every short-context run. + let Ok(guard) = self.kv_cache.lock() else { + return 0; + }; + let Some(cache) = guard.as_ref() else { + return 0; + }; + cache + .layers + .iter() + .map(|l| { + l.current_len + * l.num_kv_heads + * l.head_dim + * KV_TENSORS_PER_LAYER + * std::mem::size_of::() + }) + .sum() + } + fn read_kv_row_at( &self, _handle: &KvHandle, @@ -402,7 +500,7 @@ impl KvDispatch for MetalBackend { ) -> Option<(Vec, Vec)> { // W10 Phase B: read a single position's K/V back from the Metal // kv cache. Used by engines running under HOnly that need to - // snapshot a specific position on demand (e.g. unlimited_context's + // snapshot a specific position on demand (e.g. windowed_checkpoint's // close_window). Small (~kv_dim * 4 B per K and V) so cheap vs // an end-of-window snapshot of the whole window. let cache_guard = self.kv_cache.lock().ok()?; @@ -797,3 +895,389 @@ mod tests { assert!(result.is_none()); } } + +/// Occupancy multiple of the window at which compaction runs. +/// +/// Compaction memmoves the surviving rows, so running it every step +/// would cost O(window) per token. Letting occupancy reach this multiple +/// before reclaiming makes it O(1) amortised, and the attention span +/// clamp means the extra resident rows are never read. +const COMPACTION_SLACK: usize = 2; + +impl MetalBackend { + /// Reclaim K/V above `COMPACTION_SLACK x window` rows, per layer. + /// + /// Safe to call every step: it is a no-op until occupancy actually + /// reaches the slack bound. Eviction lowers occupancy only — + /// `abs_position` keeps climbing, so RoPE does not rewind (see + /// `LayerKVCache::evict_to_window`). + pub(crate) fn compact_kv_to_window(&self, window: usize) { + if window == 0 { + return; + } + let trigger = window.saturating_mul(COMPACTION_SLACK); + let Ok(mut guard) = self.kv_cache.lock() else { + return; + }; + let Some(cache) = guard.as_mut() else { + return; + }; + for layer in cache.layers.iter_mut() { + if layer.current_len >= trigger { + layer.evict_to_window(window); + } + } + } +} + +#[cfg(test)] +mod windowed_coarse_tests { + use super::*; + use crate::MetalBackend; + + fn backend() -> MetalBackend { + MetalBackend::new().expect("Metal device available on test host") + } + + /// The narrower of the two windows wins, and `0` means unbounded on + /// both sides — the same sentinel the kernel uses, so no translation. + #[test] + fn effective_window_takes_the_narrower_of_arch_and_engine() { + let b = backend(); + + b.set_engine_window(None); + assert_eq!(b.effective_window_for(0), 0, "both unbounded"); + assert_eq!(b.effective_window_for(1024), 1024, "arch window survives"); + + b.set_engine_window(Some(256)); + assert_eq!( + b.effective_window_for(0), + 256, + "engine bounds a global layer" + ); + assert_eq!( + b.effective_window_for(1024), + 256, + "engine is narrower than the arch's sliding layer" + ); + assert_eq!( + b.effective_window_for(128), + 128, + "arch is narrower than the engine's request" + ); + + b.set_engine_window(None); + assert_eq!( + b.effective_window_for(1024), + 1024, + "clearing restores the arch" + ); + } + + /// A prompt longer than the window is declined — the fused prefill + /// has no per-query masking, so accepting would attend in full while + /// the engine advertises a bound. + #[test] + fn prefill_declines_a_prompt_longer_than_the_window() { + let b = backend(); + let weights = larql_models::test_fixtures::make_test_q4k_weights(); + assert!(b + .coarse_prefill_windowed(&weights, &[0u32, 1, 2, 3], None, Some(2)) + .is_none()); + assert!( + b.coarse_prefill_windowed(&weights, &[0u32, 1, 2, 3], None, Some(0)) + .is_none(), + "a zero window is refused, not treated as unbounded" + ); + } + + /// Compaction is a no-op below the slack bound and reclaims above it, + /// without ever moving the stream position. + #[test] + fn compaction_reclaims_only_past_the_slack_bound() { + let b = backend(); + let window = 4usize; + { + let mut guard = b.kv_cache.lock().expect("kv cache lock"); + *guard = Some(crate::ops::kv_cache::KVCache::new(&b.bufs, 1, 64, 2, 4)); + let layer = &mut guard.as_mut().unwrap().layers[0]; + for _ in 0..(window * COMPACTION_SLACK - 1) { + layer.advance_one(); + } + } + b.compact_kv_to_window(window); + { + let guard = b.kv_cache.lock().unwrap(); + let layer = &guard.as_ref().unwrap().layers[0]; + assert_eq!( + layer.current_len, + window * COMPACTION_SLACK - 1, + "below the slack bound nothing should move" + ); + } + + { + let mut guard = b.kv_cache.lock().unwrap(); + guard.as_mut().unwrap().layers[0].advance_one(); + } + b.compact_kv_to_window(window); + { + let guard = b.kv_cache.lock().unwrap(); + let layer = &guard.as_ref().unwrap().layers[0]; + assert_eq!(layer.current_len, window, "reclaimed to the window"); + assert_eq!( + layer.abs_position, + window * COMPACTION_SLACK, + "compaction must never rewind the stream position" + ); + } + } + + /// A zero window compacts nothing rather than emptying the cache. + #[test] + fn a_zero_window_compacts_nothing() { + let b = backend(); + { + let mut guard = b.kv_cache.lock().expect("kv cache lock"); + *guard = Some(crate::ops::kv_cache::KVCache::new(&b.bufs, 1, 64, 2, 4)); + guard.as_mut().unwrap().layers[0].advance_one(); + } + b.compact_kv_to_window(0); + let guard = b.kv_cache.lock().unwrap(); + assert_eq!(guard.as_ref().unwrap().layers[0].current_len, 1); + } + + /// Compaction before any prefill has allocated a cache must return + /// rather than unwrap — a windowed engine calls this every decode + /// step, including the first. + #[test] + fn compacting_without_a_cache_is_a_noop() { + let b = backend(); + assert!( + b.kv_cache.lock().expect("kv cache lock").is_none(), + "a fresh backend has no cache yet" + ); + b.compact_kv_to_window(4); + } + + // ── coarse_decode_step_windowed: the three arms ────────────────── + + #[test] + fn decode_step_windowed_forwards_when_no_window_requested() { + let weights = larql_models::test_fixtures::make_test_weights(); + let b = backend(); + let mut handle = KvHandle::new(MetalCoarseHandle); + // No index → the underlying coarse step declines, so this pins the + // *delegation*: the window-less arm clears the engine window and + // hands straight through. + assert!(b + .coarse_decode_step_windowed(&weights, 0, None, &mut handle, 0, None) + .is_none()); + assert_eq!(b.effective_window_for(1024), 1024, "engine window cleared"); + } + + #[test] + fn decode_step_windowed_declines_a_zero_window() { + let weights = larql_models::test_fixtures::make_test_weights(); + let b = backend(); + let mut handle = KvHandle::new(MetalCoarseHandle); + assert!( + b.coarse_decode_step_windowed(&weights, 0, None, &mut handle, 0, Some(0)) + .is_none(), + "a zero window is refused, not treated as unbounded" + ); + } + + #[test] + fn decode_step_windowed_sets_the_window_and_compacts_before_stepping() { + let weights = larql_models::test_fixtures::make_test_weights(); + let b = backend(); + let window = 4usize; + { + let mut guard = b.kv_cache.lock().expect("kv cache lock"); + *guard = Some(crate::ops::kv_cache::KVCache::new(&b.bufs, 1, 64, 2, 4)); + let layer = &mut guard.as_mut().unwrap().layers[0]; + for _ in 0..(window * COMPACTION_SLACK) { + layer.advance_one(); + } + } + let mut handle = KvHandle::new(MetalCoarseHandle); + // The step itself declines (no index), but the window bookkeeping + // and the compaction must both have run first — that pairing is + // the whole contract this arm exists to hold. + let _ = b.coarse_decode_step_windowed(&weights, 0, None, &mut handle, 0, Some(window)); + assert_eq!( + b.effective_window_for(1024), + window as u32, + "the engine window must be set before the step" + ); + let guard = b.kv_cache.lock().unwrap(); + assert_eq!( + guard.as_ref().unwrap().layers[0].current_len, + window, + "occupancy at the slack bound must be compacted before stepping" + ); + } + + // ── Honesty / accounting overrides ─────────────────────────────── + + #[test] + fn per_layer_surface_reports_itself_as_host_delegated() { + // Every per-layer method forwards to CPU; only `coarse_*` runs + // Metal kernels. Reporting `false` here would let diagnostics + // label a pure-CPU measurement as GPU work. + assert!(backend().per_layer_is_host_delegated()); + } + + #[test] + fn resident_kv_bytes_is_zero_before_a_cache_exists() { + assert_eq!(backend().backend_resident_kv_bytes(), 0); + } + + #[test] + fn resident_kv_bytes_counts_the_populated_prefix_not_the_capacity() { + let b = backend(); + let (max_seq, num_kv_heads, head_dim) = (64usize, 2usize, 4usize); + let filled = 3usize; + { + let mut guard = b.kv_cache.lock().expect("kv cache lock"); + *guard = Some(crate::ops::kv_cache::KVCache::new( + &b.bufs, + 1, + max_seq, + num_kv_heads, + head_dim, + )); + let layer = &mut guard.as_mut().unwrap().layers[0]; + for _ in 0..filled { + layer.advance_one(); + } + } + let per_row = num_kv_heads * head_dim * KV_TENSORS_PER_LAYER * std::mem::size_of::(); + assert_eq!(b.backend_resident_kv_bytes(), filled * per_row); + // The buffers are preallocated to the context ceiling; charging for + // capacity would overstate every short-context run by ~21x here. + assert_ne!(b.backend_resident_kv_bytes(), max_seq * per_row); + } + + // ── read_kv_row_at against a populated cache ───────────────────── + + #[test] + fn read_kv_row_at_returns_the_requested_position() { + let b = backend(); + let (max_seq, num_kv_heads, head_dim) = (8usize, 2usize, 4usize); + let stride = num_kv_heads * head_dim; + // Distinct value per slot so a wrong offset cannot pass. + let k_data: Vec = (0..max_seq * stride).map(|i| i as f32).collect(); + let v_data: Vec = (0..max_seq * stride).map(|i| -(i as f32)).collect(); + { + let mut guard = b.kv_cache.lock().expect("kv cache lock"); + *guard = Some(crate::ops::kv_cache::KVCache::new( + &b.bufs, + 1, + max_seq, + num_kv_heads, + head_dim, + )); + let layer = &mut guard.as_mut().unwrap().layers[0]; + layer.k_cache = b.bufs.get_f32(&k_data); + layer.v_cache = b.bufs.get_f32(&v_data); + layer.advance_one(); + layer.advance_one(); + } + let sentinel = KvHandle::new(MetalCoarseHandle); + + let (k0, v0) = b.read_kv_row_at(&sentinel, 0, 0).expect("position 0"); + assert_eq!(k0, k_data[0..stride].to_vec()); + assert_eq!(v0, v_data[0..stride].to_vec()); + + let (k1, v1) = b.read_kv_row_at(&sentinel, 0, 1).expect("position 1"); + assert_eq!(k1, k_data[stride..2 * stride].to_vec()); + assert_eq!(v1, v_data[stride..2 * stride].to_vec()); + } + + #[test] + fn read_kv_row_at_declines_a_position_past_the_populated_prefix() { + let b = backend(); + { + let mut guard = b.kv_cache.lock().expect("kv cache lock"); + *guard = Some(crate::ops::kv_cache::KVCache::new(&b.bufs, 1, 64, 2, 4)); + guard.as_mut().unwrap().layers[0].advance_one(); + } + let sentinel = KvHandle::new(MetalCoarseHandle); + assert!( + b.read_kv_row_at(&sentinel, 0, 0).is_some(), + "pos 0 is filled" + ); + assert!( + b.read_kv_row_at(&sentinel, 0, 1).is_none(), + "pos == current_len is capacity, not data" + ); + assert!( + b.read_kv_row_at(&sentinel, 9, 0).is_none(), + "a layer beyond the cache declines" + ); + } + + /// A prompt that fits is accepted, and the engine window is armed + /// before the prefill runs — the ordering matters, because the + /// kernel reads the window off the backend, not off the argument. + #[test] + fn prefill_windowed_arms_the_window_for_a_prompt_that_fits() { + let b = backend(); + let weights = larql_models::test_fixtures::make_test_q4k_weights(); + // No index → the prefill itself declines, which is fine: what this + // pins is that a fitting prompt gets past the guard and sets the + // window rather than being refused outright. + let _ = b.coarse_prefill_windowed(&weights, &[0u32, 1], None, Some(4)); + assert_eq!( + b.effective_window_for(1024), + 4, + "a prompt within the window must arm it, not decline" + ); + } + + /// The unmasked entry point is a forwarder that must request the + /// full dump; nothing else in the crate calls it, so without this + /// the `StateDumpMask::Full` it pins is untested. + #[test] + fn decode_step_with_state_forwards_asking_for_the_full_dump() { + let b = backend(); + let weights = larql_models::test_fixtures::make_test_weights(); + let mut handle = KvHandle::new(MetalCoarseHandle); + let mut state = larql_compute::PerLayerDecodeState::with_capacity(weights.num_layers); + assert!(b + .coarse_decode_step_with_state(&weights, 0, None, &mut handle, 0, Some(&mut state)) + .is_none()); + } + + /// `CpuBackend` leaves `compressed_kv_append` at the trait default, + /// so Metal's delegation surfaces that backend's panic rather than + /// silently dropping the append. Pins the delegation, not the codec. + #[test] + #[should_panic(expected = "compressed_kv_append not implemented")] + fn compressed_kv_append_delegates_to_cpu() { + struct PassthroughCodec; + impl CompressionCodec for PassthroughCodec { + fn encode(&self, vec: &[f32]) -> Vec { + vec.iter().flat_map(|f| f.to_le_bytes()).collect() + } + fn decode(&self, bytes: &[u8], dim: usize) -> Vec { + bytes + .chunks_exact(4) + .take(dim) + .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]])) + .collect() + } + fn name(&self) -> &str { + "passthrough" + } + } + + let b = backend(); + let mut handle = b.alloc_kv_buffer(0, 4, 4); + let k = Array2::zeros((1, 4)); + let v = Array2::zeros((1, 4)); + b.compressed_kv_append(&mut handle, &k, &v, &PassthroughCodec); + } +} diff --git a/crates/larql-compute-metal/src/ops/full_pipeline/kv_copy.rs b/crates/larql-compute-metal/src/ops/full_pipeline/kv_copy.rs index 276c745af..ca7d6db1d 100644 --- a/crates/larql-compute-metal/src/ops/full_pipeline/kv_copy.rs +++ b/crates/larql-compute-metal/src/ops/full_pipeline/kv_copy.rs @@ -44,6 +44,8 @@ pub(super) fn populate_kv_one_layer( std::ptr::copy_nonoverlapping(v_src, v_dst, total_kv); } kv.layers[layer_idx].current_len = seq_len; + // Prefill wrote positions 0..seq_len, so the stream is at seq_len. + kv.layers[layer_idx].abs_position = seq_len; } /// Copy each layer's K/V scratch (post-RoPE) into the persistent KV @@ -80,6 +82,7 @@ pub(super) fn populate_kv_after_commit( std::ptr::copy_nonoverlapping(v_src, v_dst, total_kv); } kv.layers[l].current_len = seq_len; + kv.layers[l].abs_position = seq_len; } } diff --git a/crates/larql-compute-metal/src/ops/kv_cache.rs b/crates/larql-compute-metal/src/ops/kv_cache.rs index 0068d49ef..f908d40f0 100644 --- a/crates/larql-compute-metal/src/ops/kv_cache.rs +++ b/crates/larql-compute-metal/src/ops/kv_cache.rs @@ -40,7 +40,17 @@ pub fn attention_span(t: u32, window_size: u32) -> u32 { pub struct LayerKVCache { pub k_cache: Buffer, // [max_seq, num_kv_heads, head_dim] f32 pub v_cache: Buffer, // same + /// How many rows are currently stored — the span attention reads. + /// + /// This is **occupancy, not position**. The two were one field until + /// sliding windows needed them apart: a window drops the oldest rows, + /// so occupancy falls while the stream position keeps climbing. Using + /// this for RoPE would rewind every token after a window slid. pub current_len: usize, + /// Absolute stream position of the NEXT row to be written — what RoPE + /// must be computed at. Monotonic for the life of the sequence; never + /// reduced by eviction. + pub abs_position: usize, pub max_seq: usize, pub num_kv_heads: usize, pub head_dim: usize, @@ -54,15 +64,59 @@ impl LayerKVCache { k_cache: bufs.output(size), v_cache: bufs.output(size), current_len: 0, + abs_position: 0, max_seq, num_kv_heads, head_dim, } } - /// Reset cache (for new prompt). + /// Reset cache (for new prompt) — both occupancy and position, since + /// a new prompt restarts the stream. pub fn clear(&mut self) { self.current_len = 0; + self.abs_position = 0; + } + + /// Record one appended row: occupancy grows and the stream advances. + /// Kept together so a future append site cannot bump one and forget + /// the other — the failure that would show up as shifted RoPE many + /// tokens later. + pub fn advance_one(&mut self) { + self.current_len += 1; + self.abs_position += 1; + } + + /// Drop all but the newest `window` rows, sliding the window forward. + /// + /// Occupancy falls to `window`; `abs_position` is deliberately + /// untouched, because the surviving rows keep the RoPE they were + /// written with and the next row still belongs at the position the + /// stream has actually reached. Softmax over keys is + /// order-independent, so nothing needs repairing after the move. + /// + /// Returns the number of rows dropped. + pub fn evict_to_window(&mut self, window: usize) -> usize { + if window == 0 || self.current_len <= window { + return 0; + } + let drop = self.current_len - window; + let row = self.num_kv_heads * self.head_dim; + for buf in [&self.k_cache, &self.v_cache] { + let ptr = buf.contents() as *mut f32; + if ptr.is_null() { + return 0; + } + // SAFETY: buffers are host-visible and sized `max_seq * row`; + // `current_len <= max_seq`, so both ranges are in bounds and + // `copy_within` semantics handle the overlap. + unsafe { + let slice = std::slice::from_raw_parts_mut(ptr, self.max_seq * row); + slice.copy_within(drop * row..self.current_len * row, 0); + } + } + self.current_len = window; + drop } } @@ -274,6 +328,129 @@ mod tests { (bufs, d) } + // ── occupancy vs absolute position ────────────────────────────── + // + // `current_len` used to be both "rows stored" and "stream position". + // A sliding window separates them: eviction lowers occupancy while + // the stream keeps advancing. If they re-merge, RoPE silently rewinds + // on every token after a window slides — which surfaces as fluent but + // wrong output, the worst failure shape to debug. + + /// One row appended advances both counters together. + #[test] + fn advance_one_moves_occupancy_and_position_together() { + let (bufs, _d) = fresh_cache(); + let mut c = LayerKVCache::new(&bufs, 8, 2, 4); + assert_eq!((c.current_len, c.abs_position), (0, 0)); + c.advance_one(); + c.advance_one(); + assert_eq!((c.current_len, c.abs_position), (2, 2)); + } + + /// Eviction lowers occupancy and leaves the stream position alone. + /// This is the whole invariant the window rests on. + #[test] + fn eviction_lowers_occupancy_but_never_the_stream_position() { + let (bufs, _d) = fresh_cache(); + let mut c = LayerKVCache::new(&bufs, 16, 2, 4); + for _ in 0..10 { + c.advance_one(); + } + assert_eq!((c.current_len, c.abs_position), (10, 10)); + + let dropped = c.evict_to_window(4); + assert_eq!(dropped, 6); + assert_eq!(c.current_len, 4, "occupancy must fall to the window"); + assert_eq!( + c.abs_position, 10, + "stream position must NOT rewind — the next row still belongs at 10" + ); + + // Appending after eviction continues the stream, it does not restart it. + c.advance_one(); + assert_eq!((c.current_len, c.abs_position), (5, 11)); + } + + /// Eviction keeps the NEWEST rows, in order. + #[test] + fn eviction_keeps_the_newest_rows_in_order() { + let (bufs, _d) = fresh_cache(); + let (kv_heads, head_dim) = (2usize, 4usize); + let row = kv_heads * head_dim; + let mut c = LayerKVCache::new(&bufs, 16, kv_heads, head_dim); + + // Stamp each row with its index so survivors are identifiable. + let rows = 10usize; + unsafe { + for buf in [&c.k_cache, &c.v_cache] { + let ptr = buf.contents() as *mut f32; + let slice = std::slice::from_raw_parts_mut(ptr, 16 * row); + for r in 0..rows { + for i in 0..row { + slice[r * row + i] = r as f32; + } + } + } + } + for _ in 0..rows { + c.advance_one(); + } + + let window = 4usize; + c.evict_to_window(window); + + unsafe { + for buf in [&c.k_cache, &c.v_cache] { + let ptr = buf.contents() as *const f32; + let slice = std::slice::from_raw_parts(ptr, 16 * row); + for r in 0..window { + let expected = (rows - window + r) as f32; + assert_eq!( + slice[r * row], + expected, + "row {r} after eviction should hold original row {expected}" + ); + } + } + } + } + + /// A window at or above occupancy is a no-op — nothing moves, so the + /// unwindowed path pays nothing for the affordance existing. + #[test] + fn eviction_is_a_noop_when_the_window_cannot_bind() { + let (bufs, _d) = fresh_cache(); + let mut c = LayerKVCache::new(&bufs, 8, 2, 4); + for _ in 0..3 { + c.advance_one(); + } + assert_eq!(c.evict_to_window(3), 0); + assert_eq!(c.evict_to_window(99), 0); + assert_eq!((c.current_len, c.abs_position), (3, 3)); + } + + /// A zero window is refused rather than emptying the cache — the same + /// sentinel confusion that already cost a bug in the pipeline spec. + #[test] + fn a_zero_window_evicts_nothing() { + let (bufs, _d) = fresh_cache(); + let mut c = LayerKVCache::new(&bufs, 8, 2, 4); + c.advance_one(); + assert_eq!(c.evict_to_window(0), 0); + assert_eq!(c.current_len, 1); + } + + /// A new prompt restarts the stream, so `clear` resets both. + #[test] + fn clear_resets_occupancy_and_position() { + let (bufs, _d) = fresh_cache(); + let mut c = LayerKVCache::new(&bufs, 8, 2, 4); + c.advance_one(); + c.advance_one(); + c.clear(); + assert_eq!((c.current_len, c.abs_position), (0, 0)); + } + #[test] fn shape_mismatch_detects_conflicting_existing_layer() { assert!(!super::shape_pairs_have_mismatch( diff --git a/crates/larql-compute/Cargo.toml b/crates/larql-compute/Cargo.toml index db3147332..59a82a1b1 100644 --- a/crates/larql-compute/Cargo.toml +++ b/crates/larql-compute/Cargo.toml @@ -16,6 +16,7 @@ rayon = "1.10" # Wire-format constants (Q4_K_BLOCK_ELEMS, etc.) for padding decisions. # Tests/benches depend on it too — keep both lists in sync. larql-models = { path = "../larql-models" } +larql-execution = { path = "../larql-execution" } [target.'cfg(any(target_os = "linux", target_os = "freebsd"))'.dependencies] ndarray = { version = "0.16", features = ["blas"] } diff --git a/crates/larql-compute/README.md b/crates/larql-compute/README.md index e1ee8a075..844eab68f 100644 --- a/crates/larql-compute/README.md +++ b/crates/larql-compute/README.md @@ -35,7 +35,7 @@ GPU→CPU bridge. and [`KvDispatch::coarse_decode_step_with_state_masked`](src/kv_dispatch/mod.rs). - `KvDispatch::read_kv_row_at` — on-demand readback of a single position's K/V from the backend's kv cache. Used by engines (e.g. - `UnlimitedContextEngine.close_window`) that dropped their CPU + `WindowedCheckpointEngine.close_window`) that dropped their CPU shadow. Default impls preserve `Full` behaviour everywhere; backends without diff --git a/crates/larql-compute/ROADMAP.md b/crates/larql-compute/ROADMAP.md index 88c3092c1..ecac4b45e 100644 --- a/crates/larql-compute/ROADMAP.md +++ b/crates/larql-compute/ROADMAP.md @@ -36,6 +36,51 @@ From the whole-codebase review ([`docs/audits/codebase-review-2026-05-28.md`](.. forward: `fc1 → GELU-tanh → fc2` with bias on both layers. No spatial pooling — Granite's encoder output has the correct token count per tile. +## Open: head-major K/V layout for decode attention + +**The single highest-leverage change the `kvperf-1` data points at**, because +it sits *below* every KV engine rather than inside one. All six cached engines +share one marginal cost per context token (7.72-8.30 µs, a 7.5% spread across +six mechanisms that could hardly be more different) — they all end up in this +kernel, and it runs at roughly 44% of attainable memory bandwidth. See +[`larql-kv/docs/decode-cost-model.md`](../larql-kv/docs/decode-cost-model.md) §4. + +Today K/V is `[L, num_kv * head_dim]` — head-**minor**: + +```text +each head's gemv reads head_dim floats with a num_kv * head_dim stride +with GQA, the `reps` query heads sharing a KV head each stream it again +``` + +On qwen3-0.6b that is a 512 B window with a 4096 B stride, read twice +(`reps = 2`). **Hypothesis (untested):** `[kv_head][L][head_dim]` makes each +head's read contiguous *and* lets the sharing q-heads share one pass, roughly +halving the O(context) term that dominates every engine. + +A bench exists — `benches/decode_attention_layout.rs` — comparing four layouts +doing identical arithmetic, so any difference is memory behaviour. It includes +a `head_major_spin` variant specifically so the parallel primitive is not +confounded with the layout (production uses the spin pool; a rayon candidate +would otherwise be compared on two axes at once). + +**Measurement protocol, because this one has already bitten:** an earlier +version timed one layer against one reused K/V buffer. At L=1024 that is 8.4 MB +— cache-resident — and it clocked 90 µs against the ~301 µs the end-to-end +slope implies, reporting essentially the roofline. It was measuring a regime +decode never runs in, **and the layout ranking inverted between that regime and +the real one.** Each iteration now walks a full stack of distinct per-layer K/V +buffers, as one decode step does. + +Two further conditions before believing a result: + +- **AC power.** On battery macOS throttles enough to invent or erase this + effect, and criterion runs the variants sequentially — deepening throttle + mid-run penalises later variants and corrupts the *ranking*, which is the + entire output. +- **A drift control.** Measure `head_minor` twice, first and last. If those + two disagree the machine moved under the run and the ordering is not + trustworthy regardless of what it says. + ## Open: compute modularity and model-agnostic cleanup **Status**: Started 2026-05-08. diff --git a/crates/larql-compute/examples/attn_prefill_f32_vs_q4k.rs b/crates/larql-compute/examples/attn_prefill_f32_vs_q4k.rs deleted file mode 100644 index 7d57e78f5..000000000 --- a/crates/larql-compute/examples/attn_prefill_f32_vs_q4k.rs +++ /dev/null @@ -1,191 +0,0 @@ -//! attn_prefill_f32_vs_q4k — prefill-shape Gate-2 for the prefill twin (task #16). -//! -//! At DECODE (seq_len=1, matvec) `q4k_matvec` beat AMX/Accelerate f32 BLAS -//! 2.06–2.51× (`attn_proj_f32_vs_q4k`) — weight bandwidth-bound, Q4K's best case. -//! PREFILL is the opposite regime and must clear its own gate before the twin is -//! built (don't inherit "better lever" from a single 43%-of-TTFT number): -//! - seq_len ≈ 907 is a batched GEMM — AMX's home turf. f32 `sgemm` reads each -//! weight ONCE and reuses it across all positions (blocking), so the proj is -//! compute-bound, not weight-bandwidth-bound — exactly where Q4K's bandwidth -//! edge evaporates. -//! - `CpuBackend` has **no `q4k_matmul`**, so a twin's only available Q4K path -//! is repeated per-position `q4k_matvec` — re-reading the packed weight once -//! PER POSITION (≈907×) vs f32's single amortised read. -//! -//! f32 = `dot_proj_gpu(x[seq,hidden], w[rows,hidden])` — one BLAS sgemm -//! q4k = seq × `q4k_matvec(...)` — per-position, no amortisation -//! speedup = f32 / q4k (>1 → Q4K-direct wins; <1 → f32 AMX wins, twin is dead -//! without a real `q4k_matmul` kernel). Also prints the bandwidth FLOOR a perfect -//! amortised `q4k_matmul` could reach, to size whether building one is worth it. -//! -//! Usage: cargo run --release -p larql-compute --example attn_prefill_f32_vs_q4k - -extern crate blas_src; - -use larql_compute::cpu::ops::q4_common::{quantize_q4_k, quantize_q6_k}; -use larql_compute::prelude::*; // QuantMatVec -use larql_compute::{dot_proj_gpu, CpuBackend, QuantFormat}; -use ndarray::Array2; -use std::time::{Duration, Instant}; - -fn fill(rows: usize, cols: usize) -> Array2 { - Array2::from_shape_fn((rows, cols), |(i, j)| { - ((((i * 31 + j * 17) % 251) as f32) - 125.0) * 0.001 - }) -} - -fn bench f32>(min_secs: f64, mut f: F) -> f64 { - let mut sink = f(); - let target = Duration::from_secs_f64(min_secs); - let start = Instant::now(); - let mut iters: u64 = 0; - loop { - sink += f(); - iters += 1; - if start.elapsed() >= target { - break; - } - } - std::hint::black_box(sink); - start.elapsed().as_nanos() as f64 / iters as f64 -} - -struct Geom { - name: &'static str, - hidden: usize, - num_q: usize, - num_kv: usize, - head_dim: usize, -} - -fn main() { - let backend = &CpuBackend; - let min_secs = 0.4; - // Default 907 = representative long-context run; override with an arg to - // probe the short-prompt prefill regime (e.g. `… -- 5`), where the - // amortised q4k_matmul is bandwidth-bound and should win. - let seq: usize = std::env::args() - .nth(1) - .and_then(|s| s.parse().ok()) - .unwrap_or(907); - - let geoms = [ - Geom { - name: "sliding", - hidden: 2816, - num_q: 16, - num_kv: 8, - head_dim: 256, - }, - Geom { - name: "global", - hidden: 2816, - num_q: 16, - num_kv: 4, - head_dim: 512, - }, - ]; - - println!( - "attn_prefill_f32_vs_q4k — prefill-shape gate (seq_len={seq}, Gemma-4-26B-A4B, CpuBackend)" - ); - println!("f32 = one BLAS sgemm q4k = {seq}× per-position q4k_matvec (no q4k_matmul on CPU)"); - println!("speedup = f32/q4k (<1 → AMX f32 wins → twin dead w/o a real q4k_matmul)\n"); - - for g in &geoms { - let q_dim = g.num_q * g.head_dim; - let kv_dim = g.num_kv * g.head_dim; - // (label, num_rows, in_dim, q6) — V is Q6_K on the 26B. - let projs = [ - ("Q", q_dim, g.hidden, false), - ("K", kv_dim, g.hidden, false), - ("V", kv_dim, g.hidden, true), - ("O", g.hidden, q_dim, false), - ]; - println!("── {} layer ──", g.name); - println!( - "{:>4} | {:>12} | {:>9} | {:>11} | {:>11} | {:>7} | {:>9}", - "proj", "shape", "f32 ms", "matvec×seq", "q4k_matmul", "mm/f32", "floor ms" - ); - - let mut tot_f32 = 0.0; - let mut tot_q4k = 0.0; - for (label, num_rows, in_dim, q6) in projs { - let w = fill(num_rows, in_dim); - let w_slice = w.as_slice().unwrap(); - let x = fill(seq, in_dim); // [seq, in_dim] - let (q_bytes, fmt) = if q6 { - (quantize_q6_k(w_slice), QuantFormat::Q6_K) - } else { - (quantize_q4_k(w_slice), QuantFormat::Q4_K) - }; - - // f32: one sgemm — x · wᵀ → [seq, num_rows]. - let f32_ns = bench(min_secs, || { - let o = dot_proj_gpu(&x, &w, Some(backend)); - o[[0, 0]] + o[[seq - 1, num_rows - 1]] - }); - // q4k: per-position matvec (what the twin would do, no amortisation). - let q4k_ns = bench(min_secs, || { - let mut acc = 0.0f32; - for s in 0..seq { - let row = x.row(s); - let o = backend - .quant_matvec(fmt, &q_bytes, row.as_slice().unwrap(), num_rows, in_dim) - .unwrap(); - acc += o[0]; - } - acc - }); - // q4k_matmul: the real amortised CPU kernel (Q4_K only — V is Q6_K). - let mm_ns = if !q6 { - bench(min_secs, || { - let o = backend - .q4k_matmul(&q_bytes, x.as_slice().unwrap(), num_rows, in_dim, seq) - .unwrap(); - o[0] + o[seq * num_rows - 1] - }) - } else { - f64::NAN - }; - - let f32_ms = f32_ns / 1e6; - let q4k_ms = q4k_ns / 1e6; - let mm_ms = mm_ns / 1e6; - // Bandwidth floor for a hypothetical amortised q4k_matmul: f32 reads - // 4 B/weight once; q4k reads ~0.56 B/weight once → best case the proj - // shrinks by the byte ratio IF it were weight-bandwidth-bound. (At - // batched-gemm seq it's compute-bound, so this is optimistic.) - let bytes_ratio = if q6 { - 210.0 / 256.0 / 4.0 - } else { - 144.0 / 256.0 / 4.0 - }; - let floor_ms = f32_ms * bytes_ratio; - tot_f32 += f32_ms; - tot_q4k += q4k_ms; - println!( - "{:>4} | {:>12} | {:>9.3} | {:>11.3} | {:>11.3} | {:>6.2}× | {:>6.3} ms", - label, - format!("[{num_rows}×{in_dim}]"), - f32_ms, - q4k_ms, - mm_ms, - f32_ms / mm_ms, - floor_ms - ); - } - println!( - "{:>4} | {:>12} | {:>10.3} | {:>10.3} | {:>7.2}× |", - "Σ", - "block", - tot_f32, - tot_q4k, - tot_f32 / tot_q4k - ); - println!(); - } - println!("Read: q4k per-position matvec re-reads each weight {seq}× — if it loses badly,"); - println!("the twin needs a real q4k_matmul kernel; compare f32 ms vs the floor to see if even"); - println!("a perfect amortised kernel could beat AMX f32 gemm at this seq_len."); -} diff --git a/crates/larql-compute/examples/attn_proj_f32_vs_q4k.rs b/crates/larql-compute/examples/attn_proj_f32_vs_q4k.rs deleted file mode 100644 index 37bb3ab2b..000000000 --- a/crates/larql-compute/examples/attn_proj_f32_vs_q4k.rs +++ /dev/null @@ -1,180 +0,0 @@ -//! attn_proj_f32_vs_q4k — task #16 step-2 GATE, part 2 (the #24-trap guard). -//! -//! The split timer (`attn_proj_vs_gqa_split`) proved the four Q/K/V/O -//! projections dominate the attention block at working context — i.e. they're -//! worth attacking. This bench asks the *decisive* question: does the -//! Q4K-direct kernel actually BEAT the current f32 BLAS projection on this -//! hardware? Apple's AMX/Accelerate sgemm is very fast for f32; a Q4K matvec -//! trades that throughput for ~7× lower weight bandwidth. If AMX f32 wins -//! anyway, the whole lever is dead even though projections dominate — and we -//! learn it here, before building the path (the #24 build-then-measure trap). -//! -//! Per Gemma-4-26B-A4B geometry, per projection (Q/K/V/O): -//! f32 = `dot_proj_gpu(x, w_f32, CpuBackend)` (today's path) -//! q4k = `CpuBackend::q4k_matvec(quantize_q4_k(w), x, rows, cols)` (direct) -//! Same logical matrix; q4k weights quantized once outside the timed loop. -//! A rough f32-vs-q4k numeric delta is printed as a wiring sanity check — NOT -//! the parity gate (the real gate is Q4K-direct vs Q4K-DEQUANT, same bytes). -//! -//! Usage: -//! cargo run --release -p larql-compute --example attn_proj_f32_vs_q4k - -extern crate blas_src; - -use larql_compute::cpu::ops::q4_common::quantize_q4_k; -use larql_compute::prelude::*; // QuantMatVec for q4k_matvec -use larql_compute::{dot_proj_gpu, CpuBackend}; -use ndarray::Array2; -use std::time::{Duration, Instant}; - -fn fill(rows: usize, cols: usize) -> Array2 { - Array2::from_shape_fn((rows, cols), |(i, j)| { - ((((i * 31 + j * 17) % 251) as f32) - 125.0) * 0.001 - }) -} - -fn bench f32>(min_secs: f64, mut f: F) -> f64 { - let mut sink = f(); // warmup - let target = Duration::from_secs_f64(min_secs); - let start = Instant::now(); - let mut iters: u64 = 0; - loop { - sink += f(); - iters += 1; - if start.elapsed() >= target { - break; - } - } - std::hint::black_box(sink); - start.elapsed().as_nanos() as f64 / iters as f64 -} - -/// Relative L2 error between the f32-BLAS and Q4K-direct outputs — a wiring -/// sanity check (expected: small, = quant error vs f32), not the parity gate. -fn rel_err(a: &[f32], b: &[f32]) -> f32 { - let mut num = 0.0f64; - let mut den = 0.0f64; - for (x, y) in a.iter().zip(b.iter()) { - num += ((x - y) as f64).powi(2); - den += (*x as f64).powi(2); - } - (num.sqrt() / den.sqrt().max(1e-12)) as f32 -} - -struct Geom { - name: &'static str, - hidden: usize, - num_q: usize, - num_kv: usize, - head_dim: usize, -} - -fn main() { - let backend = &CpuBackend; - let min_secs = 0.4; - - let geoms = [ - Geom { - name: "sliding", - hidden: 2816, - num_q: 16, - num_kv: 8, - head_dim: 256, - }, - Geom { - name: "global", - hidden: 2816, - num_q: 16, - num_kv: 4, - head_dim: 512, - }, - ]; - - println!( - "attn_proj_f32_vs_q4k — task #16 step-2 gate part 2 (Gemma-4-26B-A4B dims, CpuBackend)" - ); - println!( - "f32 = dot_proj_gpu (Accelerate/AMX sgemm) q4k = q4k_matvec (Q4_K × f32, ~7× less BW)" - ); - println!("speedup = f32 / q4k (>1 → Q4K-direct wins)\n"); - - for g in &geoms { - let q_dim = g.num_q * g.head_dim; - let kv_dim = g.num_kv * g.head_dim; - - let w_q = fill(q_dim, g.hidden); - let w_k = fill(kv_dim, g.hidden); - let w_v = fill(kv_dim, g.hidden); - let w_o = fill(g.hidden, q_dim); - let h_norm = fill(1, g.hidden); - let attn_out = fill(1, q_dim); - - // (label, &w_f32, input_array, num_rows, in_dim) - #[allow(clippy::type_complexity)] - let projs: [(&str, &Array2, &Array2, usize, usize); 4] = [ - ("Q", &w_q, &h_norm, q_dim, g.hidden), - ("K", &w_k, &h_norm, kv_dim, g.hidden), - ("V", &w_v, &h_norm, kv_dim, g.hidden), - ("O", &w_o, &attn_out, g.hidden, q_dim), - ]; - - println!("── {} layer ──", g.name); - println!( - "{:>4} | {:>6} | {:>12} | {:>9} | {:>9} | {:>8} | {:>10}", - "proj", "rows", "shape", "f32 ms", "q4k ms", "speedup", "rel.err" - ); - - let mut tot_f32 = 0.0; - let mut tot_q4k = 0.0; - for (label, w, input, num_rows, in_dim) in projs { - let w_slice = w.as_slice().unwrap(); - let x_slice = input.as_slice().unwrap(); - let q4k = quantize_q4_k(w_slice); // once, outside timing - - let f32_ns = bench(min_secs, || { - let o = dot_proj_gpu(input, w, Some(backend)); - o[[0, 0]] - }); - let q4k_ns = bench(min_secs, || { - let o = backend.q4k_matvec(&q4k, x_slice, num_rows, in_dim).unwrap(); - o[0] - }); - - // wiring sanity: f32 vs q4k output on one call - let f32_out = dot_proj_gpu(input, w, Some(backend)); - let q4k_out = backend.q4k_matvec(&q4k, x_slice, num_rows, in_dim).unwrap(); - let err = rel_err(f32_out.as_slice().unwrap(), &q4k_out); - - let f32_ms = f32_ns / 1e6; - let q4k_ms = q4k_ns / 1e6; - tot_f32 += f32_ms; - tot_q4k += q4k_ms; - println!( - "{:>4} | {:>6} | {:>12} | {:>9.4} | {:>9.4} | {:>7.2}× | {:>10.2e}", - label, - num_rows, - format!("[{num_rows}×{in_dim}]"), - f32_ms, - q4k_ms, - f32_ms / q4k_ms, - err - ); - } - println!( - "{:>4} | {:>6} | {:>12} | {:>9.4} | {:>9.4} | {:>7.2}× |", - "Σ", - "", - "block", - tot_f32, - tot_q4k, - tot_f32 / tot_q4k - ); - println!(); - } - - println!("Gate reading: per-projection speedup >1 → Q4K-direct beats AMX f32 BLAS, so the"); - println!("(dominant) projection part of the 28% genuinely shrinks. ≤1 → bandwidth cut is"); - println!( - "eaten by AMX f32 throughput → lever is dead, do NOT build (the #24 trap, caught cheap)." - ); -} diff --git a/crates/larql-compute/examples/attn_proj_vs_gqa_split.rs b/crates/larql-compute/examples/attn_proj_vs_gqa_split.rs deleted file mode 100644 index 6af76a14e..000000000 --- a/crates/larql-compute/examples/attn_proj_vs_gqa_split.rs +++ /dev/null @@ -1,223 +0,0 @@ -//! attn_proj_vs_gqa_split — task #16 step-2 GATE (Q4K-direct attention). -//! -//! The cheap probe that gates any kernel work (the #24 build-then-measure -//! trap). The decode-attention ~28% (`docs/diagnoses/remote-moe-bottlenecks.md`) -//! is recorded as ONE number (`record_attn` wraps the whole block), so the -//! projection-vs-GQA share inside it is unmeasured. Only the four Q/K/V/O -//! projections are Q4K-accelerable; the GQA decode step is f32 and GROWS with -//! cached_len. This bench splits them so we know how much of the 28% a -//! Q4K-direct projection path can actually reclaim — and how that share decays -//! as context grows. -//! -//! Method: time the REAL production functions — `dot_proj_gpu` (f32 BLAS, the -//! Q4K-accelerable projection, ×4 for Q/K/V/O) vs `gqa_attention_decode_step` -//! (f32, NOT accelerable, O(cached_len·hd·num_q)) — -//! across a cached_len sweep at true Gemma-4-26B-A4B dims. Synthetic -//! same-size f32 weights read at the same memory bandwidth as real ones, so the -//! timing is faithful with no model load. Backend = `CpuBackend` (the no-`--metal` -//! decode path that the 28% was measured on). -//! -//! Usage: -//! cargo run --release -p larql-compute --example attn_proj_vs_gqa_split - -extern crate blas_src; - -use larql_compute::attention::gqa_attention_decode_step; -use larql_compute::{dot_proj_gpu, CpuBackend}; -use ndarray::Array2; -use std::time::{Duration, Instant}; - -/// Deterministic small non-zero fill — avoids zeros/denormals that BLAS or -/// the GQA softmax might special-case and mis-time. -fn fill(rows: usize, cols: usize) -> Array2 { - Array2::from_shape_fn((rows, cols), |(i, j)| { - ((((i * 31 + j * 17) % 251) as f32) - 125.0) * 0.001 - }) -} - -/// Run `f` repeatedly for at least `min_secs`; return (ns per call, checksum). -/// Auto-scales iterations so cheap ops (GQA at short ctx) and expensive ops -/// (the 90 MB global Q projection) both get a stable measurement. The checksum -/// is consumed so the optimizer can't elide the work. -fn bench f32>(min_secs: f64, mut f: F) -> (f64, f32) { - let mut sink = f(); // warmup (BLAS init on first call) - let target = Duration::from_secs_f64(min_secs); - let start = Instant::now(); - let mut iters: u64 = 0; - loop { - sink += f(); - iters += 1; - if start.elapsed() >= target { - break; - } - } - (start.elapsed().as_nanos() as f64 / iters as f64, sink) -} - -struct Geom { - name: &'static str, - count: usize, // layers of this kind in the 30-layer stack - hidden: usize, - num_q: usize, - num_kv: usize, - head_dim: usize, -} - -fn main() { - let backend = &CpuBackend; - let min_secs = 0.3; - // cached_len sweep. The 28% was measured at ctx ≈ 33–45 (33-tok prompt + - // 12 decode), so 32–128 is the "representative" band; the tail shows the - // GQA asymptote. - let sweep = [1usize, 32, 128, 512, 1024, 2048, 4096, 8192]; - - // Gemma-4-26B-A4B: 30 layers, pattern=6 → global at 5,11,17,23,29 (5), - // sliding ×25. Sliding hd=256/num_kv=8; global hd=512/num_kv=4; num_q=16. - let geoms = [ - Geom { - name: "sliding", - count: 25, - hidden: 2816, - num_q: 16, - num_kv: 8, - head_dim: 256, - }, - Geom { - name: "global", - count: 5, - hidden: 2816, - num_q: 16, - num_kv: 4, - head_dim: 512, - }, - ]; - - println!("attn_proj_vs_gqa_split — task #16 step-2 gate (Gemma-4-26B-A4B dims, CpuBackend)"); - println!("proj = dot_proj_gpu ×4 (f32 BLAS, Q4K-accelerable)"); - println!("gqa = gqa_attention_decode_step (f32, NOT accelerable, grows with cached_len)\n"); - - let mut per_geom: Vec> = Vec::new(); // [geom][i] = (proj_ms, gqa_ms) - - for g in &geoms { - let q_dim = g.num_q * g.head_dim; - let kv_dim = g.num_kv * g.head_dim; - let reps = g.num_q / g.num_kv; - let scale = 1.0 / (g.head_dim as f64).sqrt(); - - // w_q [q_dim,hidden], w_k/w_v [kv_dim,hidden], w_o [hidden,q_dim] — - // shapes per `vindex/dequant.rs::dequantize_matrix`. - let w_q = fill(q_dim, g.hidden); - let w_k = fill(kv_dim, g.hidden); - let w_v = fill(kv_dim, g.hidden); - let w_o = fill(g.hidden, q_dim); - let h_norm = fill(1, g.hidden); - let attn_out = fill(1, q_dim); // dummy O-proj input ([1,q_dim]) - let q_rope = fill(1, q_dim); - - let proj_bytes = (2 * q_dim * g.hidden + 2 * kv_dim * g.hidden) * 4; - println!( - "── {} ×{}: hidden={} q_dim={} (num_q={}×hd={}) kv_dim={} (num_kv={}) reps={} — proj f32 reads {:.1} MB/token/layer", - g.name, g.count, g.hidden, q_dim, g.num_q, g.head_dim, kv_dim, g.num_kv, reps, - proj_bytes as f64 / 1e6 - ); - println!( - "{:>10} | {:>9} | {:>9} | {:>9} | {:>10}", - "cached_len", "proj ms", "gqa ms", "block ms", "proj %" - ); - - let mut rows = Vec::new(); - for &clen in &sweep { - let k_concat = fill(clen, kv_dim); - let v_concat = fill(clen, kv_dim); - - let (proj_ns, _c1) = bench(min_secs, || { - let q = dot_proj_gpu(&h_norm, &w_q, Some(backend)); - let k = dot_proj_gpu(&h_norm, &w_k, Some(backend)); - let v = dot_proj_gpu(&h_norm, &w_v, Some(backend)); - let o = dot_proj_gpu(&attn_out, &w_o, Some(backend)); - q[[0, 0]] + k[[0, 0]] + v[[0, 0]] + o[[0, 0]] - }); - let (gqa_ns, _c2) = bench(min_secs, || { - let a = gqa_attention_decode_step( - &q_rope, &k_concat, &v_concat, g.num_q, g.head_dim, reps, scale, None, None, - ); - a[[0, 0]] - }); - - let proj_ms = proj_ns / 1e6; - let gqa_ms = gqa_ns / 1e6; - let block = proj_ms + gqa_ms; - println!( - "{:>10} | {:>9.4} | {:>9.4} | {:>9.4} | {:>9.1}%", - clen, - proj_ms, - gqa_ms, - block, - 100.0 * proj_ms / block - ); - rows.push((proj_ms, gqa_ms)); - } - println!(); - per_geom.push(rows); - } - - // Per-token attention block blended over the real stack (25 sliding + 5 - // global). CONSERVATIVE: no sliding-window cap (sliding attends full ctx) - // → upper-bounds GQA, lower-bounds projection share. A real Gemma window - // (sliding caps at W) only shrinks sliding GQA, pushing proj % HIGHER. - println!("── Blended per-token attention block (25 sliding + 5 global, no window cap = GQA upper bound) ──"); - println!( - "{:>10} | {:>9} | {:>9} | {:>9} | {:>10}", - "ctx", "proj ms", "gqa ms", "block ms", "proj %" - ); - for (i, &clen) in sweep.iter().enumerate() { - let (sp, sg) = per_geom[0][i]; - let (gp, gg) = per_geom[1][i]; - let proj = geoms[0].count as f64 * sp + geoms[1].count as f64 * gp; - let gqa = geoms[0].count as f64 * sg + geoms[1].count as f64 * gg; - let block = proj + gqa; - println!( - "{:>10} | {:>9.3} | {:>9.3} | {:>9.3} | {:>9.1}%", - clen, - proj, - gqa, - block, - 100.0 * proj / block - ); - } - // Windowed blend: real Gemma sliding layers cap KV at a window W, so their - // GQA SATURATES at cached_len=W while only the 5 global layers keep growing. - // W=1024 is the Gemma-3 default (confirm the 26B's actual value — the A4B - // detect-test config leaves `sliding_window` unset). This is the realistic - // case; the no-cap table above is the GQA upper bound. - let window = 1024usize; - let idx_at = |clen: usize| sweep.iter().position(|&c| c == clen.min(window)).unwrap(); - println!( - "\n── Blended per-token attention block (25 sliding capped at W={window} + 5 global) ──" - ); - println!( - "{:>10} | {:>9} | {:>9} | {:>9} | {:>10}", - "ctx", "proj ms", "gqa ms", "block ms", "proj %" - ); - for &clen in &sweep { - let si = idx_at(clen); // sliding sees min(ctx, W) - let gi = sweep.iter().position(|&c| c == clen).unwrap(); // global sees full ctx - let (sp, sg) = per_geom[0][si]; - let (gp, gg) = per_geom[1][gi]; - let proj = geoms[0].count as f64 * sp + geoms[1].count as f64 * gp; - let gqa = geoms[0].count as f64 * sg + geoms[1].count as f64 * gg; - let block = proj + gqa; - println!( - "{:>10} | {:>9.3} | {:>9.3} | {:>9.3} | {:>9.1}%", - clen, - proj, - gqa, - block, - 100.0 * proj / block - ); - } - - println!("\nGate reading: proj % at the representative band (cached_len 32–128) = the"); - println!("fraction of the 28% a Q4K-direct projection path can reclaim. Crossover ctx"); - println!("(proj % → 50) = where unaccelerated GQA starts to dominate the block."); -} diff --git a/crates/larql-compute/src/attention/block.rs b/crates/larql-compute/src/attention/block.rs index c4dca3d60..104fa40d0 100644 --- a/crates/larql-compute/src/attention/block.rs +++ b/crates/larql-compute/src/attention/block.rs @@ -3,9 +3,7 @@ //! norm → Q/K/V projection → bias → V-norm → QK-norm → RoPE → GQA → O projection → residual. //! Supports KV sharing (reuse K/V from a source layer). -use super::gqa::{ - gqa_attention_with_all_weights, gqa_attention_with_weights, gqa_reduced_qk_all_weights, -}; +use super::gqa::gqa_reduced_qk_all_weights; use super::{AttentionAllWeights, AttentionWeights, SharedKV}; use ndarray::{s, Array2}; @@ -474,13 +472,13 @@ fn run_attention_block_core( &q_rope, &k_rope, num_q, head_dim, reps, scale, seq_len, softcap, sinks, rank, ) }); - let (mut attn_out, attn_weights, full_all_attn_weights) = if capture_all_attention { - let (out, all_weights) = gqa_attention_with_all_weights( - &q_rope, &k_rope, &v_final, num_q, head_dim, reps, scale, seq_len, softcap, sinks, - ); - (out, None, Some(all_weights)) - } else { - let (out, weights) = gqa_attention_with_weights( + // Sliding window for THIS layer, from the shared rule the Metal + // pipeline spec also uses. `None` on a full-attention layer (or an + // architecture without windows) leaves the maths bit-identical to + // the unwindowed path. + let window = crate::forward_overrides::effective_attention_window_for_layer(arch, layer); + let (mut attn_out, attn_weights, full_all_attn_weights) = { + let (out, last, all) = super::gqa::gqa_attention_capture( &q_rope, &k_rope, &v_final, @@ -489,11 +487,13 @@ fn run_attention_block_core( reps, scale, seq_len, - capture_attention, + capture_attention && !capture_all_attention, + capture_all_attention, softcap, sinks, + window, ); - (out, weights, None) + (out, last, all) }; let all_attn_weights = reduced_qk_weights.or(full_all_attn_weights); if let Some(heads) = zero_pre_o_heads { diff --git a/crates/larql-compute/src/attention/decode/dispatch.rs b/crates/larql-compute/src/attention/decode/dispatch.rs index 16ca77065..f4cc1f48e 100644 --- a/crates/larql-compute/src/attention/decode/dispatch.rs +++ b/crates/larql-compute/src/attention/decode/dispatch.rs @@ -2,7 +2,7 @@ use ndarray::Array2; use crate::attention::SharedKV; -use super::gqa_step::gqa_attention_decode_step; +use super::gqa_step::gqa_attention_decode_step_windowed; use super::q4k_direct::run_attention_block_decode_step_q4k_direct; /// Decode-step attention with optional GPU-accelerated projections @@ -152,7 +152,10 @@ pub fn run_attention_block_decode_step_backend( }; let softcap = arch.attn_logit_softcapping(); - let attn_out = gqa_attention_decode_step( + // Per-layer sliding window from the shared rule; `None` leaves + // this bit-identical to the unwindowed step. + let window = crate::forward_overrides::effective_attention_window_for_layer(arch, layer); + let attn_out = gqa_attention_decode_step_windowed( &q_rope, &k_concat, &v_concat, @@ -167,6 +170,7 @@ pub fn run_attention_block_decode_step_backend( num_q, layer, ), + window, ); let mut attn_projected = dot_proj_gpu(&attn_out, w_o, backend); diff --git a/crates/larql-compute/src/attention/decode/gqa_step.rs b/crates/larql-compute/src/attention/decode/gqa_step.rs index 9d456c537..056809926 100644 --- a/crates/larql-compute/src/attention/decode/gqa_step.rs +++ b/crates/larql-compute/src/attention/decode/gqa_step.rs @@ -4,6 +4,55 @@ use crate::attention::SharedKV; use super::dispatch::run_attention_block_decode_step_backend; +/// GQA decode step **with a per-layer sliding window**. +/// +/// The decode query sits at the newest position, so a window is just the +/// tail of the cache: keep the last `w` rows of K/V and run the ordinary +/// step over them. Slicing beats masking here — the masked-out keys are +/// never read, so a windowed layer gets cheaper rather than merely +/// producing a different answer. +/// +/// `window = None`, or a window at least as wide as the cache, is +/// bit-identical to [`gqa_attention_decode_step`]: the slice is the whole +/// array and no arithmetic changes. +/// +/// Resolve `window` through +/// [`crate::forward_overrides::effective_attention_window_for_layer`] so +/// the CPU path and the Metal pipeline spec share one rule. +#[allow(clippy::too_many_arguments)] +pub fn gqa_attention_decode_step_windowed( + q_new: &Array2, + k_full: &ndarray::ArrayBase, + v_full: &ndarray::ArrayBase, + num_q: usize, + head_dim: usize, + reps: usize, + scale: f64, + softcap: Option, + sinks: Option<&[f32]>, + window: Option, +) -> Array2 +where + S1: ndarray::Data + Sync, + S2: ndarray::Data + Sync, +{ + let total_len = k_full.shape()[0]; + let start = match window { + Some(w) => total_len.saturating_sub(w), + None => 0, + }; + if start == 0 { + return gqa_attention_decode_step( + q_new, k_full, v_full, num_q, head_dim, reps, scale, softcap, sinks, + ); + } + let k_win = k_full.slice(ndarray::s![start.., ..]); + let v_win = v_full.slice(ndarray::s![start.., ..]); + gqa_attention_decode_step( + q_new, &k_win, &v_win, num_q, head_dim, reps, scale, softcap, sinks, + ) +} + /// GQA attention for a single decode step. /// /// `q_new`: `[1, num_q * head_dim]` — Q for the new token only. diff --git a/crates/larql-compute/src/attention/decode/mod.rs b/crates/larql-compute/src/attention/decode/mod.rs index 03e70dcf5..2f707b3d0 100644 --- a/crates/larql-compute/src/attention/decode/mod.rs +++ b/crates/larql-compute/src/attention/decode/mod.rs @@ -23,7 +23,9 @@ pub use dispatch::{ q4k_direct_attn_enabled, run_attention_block_decode_step_auto, run_attention_block_decode_step_backend, }; -pub use gqa_step::{gqa_attention_decode_step, run_attention_block_decode_step}; +pub use gqa_step::{ + gqa_attention_decode_step, gqa_attention_decode_step_windowed, run_attention_block_decode_step, +}; pub use inplace::{ run_attention_block_decode_step_auto_inplace, run_attention_block_decode_step_q4k_direct_inplace, diff --git a/crates/larql-compute/src/attention/gqa.rs b/crates/larql-compute/src/attention/gqa.rs index 7abae1078..0e84fc5ad 100644 --- a/crates/larql-compute/src/attention/gqa.rs +++ b/crates/larql-compute/src/attention/gqa.rs @@ -42,7 +42,7 @@ pub fn gqa_attention_with_weights( sinks: Option<&[f32]>, ) -> (Array2, Option) { let (out, last, _) = gqa_attention_capture( - q, k, v, num_q, head_dim, reps, scale, seq_len, capture, false, softcap, sinks, + q, k, v, num_q, head_dim, reps, scale, seq_len, capture, false, softcap, sinks, None, ); (out, last) } @@ -65,7 +65,7 @@ pub fn gqa_attention_with_all_weights( sinks: Option<&[f32]>, ) -> (Array2, AttentionAllWeights) { let (out, _, all) = gqa_attention_capture( - q, k, v, num_q, head_dim, reps, scale, seq_len, false, true, softcap, sinks, + q, k, v, num_q, head_dim, reps, scale, seq_len, false, true, softcap, sinks, None, ); ( out, @@ -73,6 +73,38 @@ pub fn gqa_attention_with_all_weights( ) } +/// GQA with causal masking **and a per-layer sliding window**. +/// +/// `window = None` is exactly [`gqa_attention`] — same code path, same +/// bits. `Some(w)` restricts every query to the most recent `w` keys, +/// which is what an architecture's sliding-attention layers actually +/// attend (Gemma 2/3, Mistral, …). +/// +/// Resolve `window` through +/// [`crate::forward_overrides::effective_attention_window_for_layer`] +/// rather than reading the architecture directly — that helper is the +/// shared rule the Metal pipeline spec uses too, so the two backends +/// cannot disagree about which layers are windowed. +#[allow(clippy::too_many_arguments)] +pub fn gqa_attention_windowed( + q: &Array2, + k: &Array2, + v: &Array2, + num_q: usize, + head_dim: usize, + reps: usize, + scale: f64, + seq_len: usize, + softcap: Option, + sinks: Option<&[f32]>, + window: Option, +) -> Array2 { + let (out, _, _) = gqa_attention_capture( + q, k, v, num_q, head_dim, reps, scale, seq_len, false, false, softcap, sinks, window, + ); + out +} + /// Capture every query-position attention distribution using only the first /// `qk_rank` dimensions of each Q/K head. This is a diagnostic surface for /// reduced-QK address probes; it does not compute a V-weighted output. @@ -130,8 +162,23 @@ pub fn gqa_reduced_qk_all_weights( } } +/// First key position a query at `causal_len - 1` may attend under a +/// sliding window. `None` (or a window at least as wide as the causal +/// prefix) keeps everything, so the windowed path is bit-identical to +/// the full-attention one until the window actually binds. +#[inline] +fn window_start(causal_len: usize, window: Option) -> usize { + match window { + Some(w) => causal_len.saturating_sub(w), + None => 0, + } +} + +/// Shared body for every causal-GQA entry point. `pub(crate)` so the +/// per-layer prefill seam in [`super::block`] can pass a sliding +/// `window` without another public overload per capture mode. #[allow(clippy::too_many_arguments)] -fn gqa_attention_capture( +pub(crate) fn gqa_attention_capture( q: &Array2, k: &Array2, v: &Array2, @@ -144,6 +191,7 @@ fn gqa_attention_capture( capture_all: bool, softcap: Option, sinks: Option<&[f32]>, + window: Option, ) -> ( Array2, Option, @@ -179,12 +227,17 @@ fn gqa_attention_capture( for qi in 0..seq_len { let causal_len = qi + 1; + // Sliding-window start for this query. `None` keeps the full + // causal prefix; `Some(w)` keeps only the most recent `w` + // keys, which is what a sliding-attention layer attends. + let start = window_start(causal_len, window); + let span = causal_len - start; let q_row = q.slice(ndarray::s![qi, q_off..q_off + head_dim]); - let k_block = k.slice(ndarray::s![0..causal_len, kv_off..kv_off + head_dim]); + let k_block = k.slice(ndarray::s![start..causal_len, kv_off..kv_off + head_dim]); let raw_scores = k_block.dot(&q_row); - for i in 0..causal_len { + for i in 0..span { let mut s = raw_scores[i] * scale_f32; if let Some(cap) = softcap { s = (s / cap).tanh() * cap; @@ -192,21 +245,24 @@ fn gqa_attention_capture( scores_buf[i] = s; } - super::softmax::softmax_in_place(&mut scores_buf[..causal_len], sink); + super::softmax::softmax_in_place(&mut scores_buf[..span], sink); + // Captured weights are indexed by absolute key position, so + // the windowed run writes its `span` values at `start..`, + // leaving masked-out positions at zero. if capture_last && qi == last_pos { let mut captured = vec![0.0f32; seq_len]; - captured[..causal_len].copy_from_slice(&scores_buf[..causal_len]); + captured[start..causal_len].copy_from_slice(&scores_buf[..span]); captured_heads.push(captured); } if capture_all { let mut captured = vec![0.0f32; seq_len]; - captured[..causal_len].copy_from_slice(&scores_buf[..causal_len]); + captured[start..causal_len].copy_from_slice(&scores_buf[..span]); captured_positions.push(captured); } - let v_block = v.slice(ndarray::s![0..causal_len, kv_off..kv_off + head_dim]); - let scores_view = ndarray::ArrayView1::from(&scores_buf[..causal_len]); + let v_block = v.slice(ndarray::s![start..causal_len, kv_off..kv_off + head_dim]); + let scores_view = ndarray::ArrayView1::from(&scores_buf[..span]); let weighted_v = v_block.t().dot(&scores_view); for d in 0..head_dim { diff --git a/crates/larql-compute/src/attention/mod.rs b/crates/larql-compute/src/attention/mod.rs index c844ba8fb..8db29c544 100644 --- a/crates/larql-compute/src/attention/mod.rs +++ b/crates/larql-compute/src/attention/mod.rs @@ -14,6 +14,9 @@ pub mod rope; pub mod sinks; pub mod softmax; +#[cfg(test)] +mod swa_tests; + use ndarray::Array2; /// Per-head attention weights for the last token position. @@ -40,8 +43,8 @@ pub struct AttentionAllWeights { pub type SharedKV = (Array2, Array2); pub use gqa::{ - gqa_attention, gqa_attention_with_all_weights, gqa_attention_with_weights, - gqa_reduced_qk_all_weights, + gqa_attention, gqa_attention_windowed, gqa_attention_with_all_weights, + gqa_attention_with_weights, gqa_reduced_qk_all_weights, }; pub use rope::{ apply_llama3_inv_freq, apply_rope, apply_rope_partial, apply_rope_partial_at, @@ -62,7 +65,7 @@ pub use block::{ run_attention_block_zero_pre_o_heads, }; pub use decode::{ - gqa_attention_decode_step, run_attention_block_decode_step, + gqa_attention_decode_step, gqa_attention_decode_step_windowed, run_attention_block_decode_step, run_attention_block_decode_step_auto, run_attention_block_decode_step_auto_inplace, run_attention_block_decode_step_backend, run_attention_block_decode_step_q4k_direct, run_attention_block_decode_step_q4k_direct_inplace, diff --git a/crates/larql-compute/src/attention/swa_tests/decode.rs b/crates/larql-compute/src/attention/swa_tests/decode.rs new file mode 100644 index 000000000..18afcb5c2 --- /dev/null +++ b/crates/larql-compute/src/attention/swa_tests/decode.rs @@ -0,0 +1,101 @@ +//! Windowed decode attention. + +use super::{ramp, HEAD_DIM, NUM_Q, REPS, SCALE}; +use crate::attention::decode::gqa_attention_decode_step_windowed; +use crate::attention::gqa_attention_decode_step; + +#[test] +fn decode_window_wider_than_cache_is_bit_identical() { + let total = 5; + let q = ramp(1, NUM_Q * HEAD_DIM, 0.4); + let k = ramp(total, NUM_Q * HEAD_DIM, 0.5); + let v = ramp(total, NUM_Q * HEAD_DIM, 0.6); + + let full = gqa_attention_decode_step(&q, &k, &v, NUM_Q, HEAD_DIM, REPS, SCALE, None, None); + for window in [None, Some(total), Some(total + 3)] { + let windowed = gqa_attention_decode_step_windowed( + &q, &k, &v, NUM_Q, HEAD_DIM, REPS, SCALE, None, None, window, + ); + assert_eq!( + full.iter().map(|v| v.to_bits()).collect::>(), + windowed.iter().map(|v| v.to_bits()).collect::>(), + "window {window:?} cannot bind on a {total}-row cache" + ); + } +} + +/// A binding decode window equals attending the cache tail — and the +/// masked-out rows are never read, which is why a windowed layer gets +/// cheaper rather than merely different. +#[test] +fn decode_window_equals_attending_the_cache_tail() { + let total = 5; + let window = 2; + let q = ramp(1, NUM_Q * HEAD_DIM, 0.4); + let k = ramp(total, NUM_Q * HEAD_DIM, 0.5); + let v = ramp(total, NUM_Q * HEAD_DIM, 0.6); + + let windowed = gqa_attention_decode_step_windowed( + &q, + &k, + &v, + NUM_Q, + HEAD_DIM, + REPS, + SCALE, + None, + None, + Some(window), + ); + let k_tail = k.slice(ndarray::s![total - window.., ..]).to_owned(); + let v_tail = v.slice(ndarray::s![total - window.., ..]).to_owned(); + let expected = gqa_attention_decode_step( + &q, &k_tail, &v_tail, NUM_Q, HEAD_DIM, REPS, SCALE, None, None, + ); + + assert_eq!( + windowed.iter().map(|v| v.to_bits()).collect::>(), + expected.iter().map(|v| v.to_bits()).collect::>(), + "a windowed decode must be exactly the tail-sliced decode" + ); + + let full = gqa_attention_decode_step(&q, &k, &v, NUM_Q, HEAD_DIM, REPS, SCALE, None, None); + assert_ne!( + full.iter().map(|v| v.to_bits()).collect::>(), + windowed.iter().map(|v| v.to_bits()).collect::>(), + "a {window}-window over {total} rows must not match full attention" + ); +} + +/// A window of 1 keeps only the newest key. +#[test] +fn decode_window_of_one_attends_only_the_newest_key() { + let total = 4; + let q = ramp(1, NUM_Q * HEAD_DIM, 0.4); + let k = ramp(total, NUM_Q * HEAD_DIM, 0.5); + let v = ramp(total, NUM_Q * HEAD_DIM, 0.6); + + let windowed = gqa_attention_decode_step_windowed( + &q, + &k, + &v, + NUM_Q, + HEAD_DIM, + REPS, + SCALE, + None, + None, + Some(1), + ); + // Softmax over a single key is 1.0, so the output is that key's V row + // per head — independent of Q. + for h in 0..NUM_Q { + for d in 0..HEAD_DIM { + let expected = v[[total - 1, h * HEAD_DIM + d]]; + assert!( + (windowed[[0, h * HEAD_DIM + d]] - expected).abs() < 1e-6, + "head {h} dim {d}: single-key attention must return that key's V" + ); + } + } +} diff --git a/crates/larql-compute/src/attention/swa_tests/mod.rs b/crates/larql-compute/src/attention/swa_tests/mod.rs new file mode 100644 index 000000000..ba4b0e3e3 --- /dev/null +++ b/crates/larql-compute/src/attention/swa_tests/mod.rs @@ -0,0 +1,33 @@ +//! Per-layer sliding-window attention: the window rule, and both CPU +//! attention paths that consume it. +//! +//! Until this landed the CPU attention path had **no** notion of a +//! per-layer window, while the Metal pipeline spec carried one — so a +//! Gemma-class model attended full history on layers the architecture +//! declares sliding, and "GPU/CPU parity" was undefined past the window. +//! Both now resolve the window through one helper, which is the property +//! these tests actually pin. +//! +//! - [`rule`] — which layers are windowed, and how wide. +//! - [`wiring`] — that the per-layer seam actually passes the window on. +//! - [`prefill`] / [`decode`] — the two attention paths that apply it. + +mod decode; +mod prefill; +mod rule; +mod wiring; + +use ndarray::Array2; + +/// Deterministic filler so the two paths see identical inputs. +pub(super) fn ramp(rows: usize, cols: usize, seed: f32) -> Array2 { + Array2::from_shape_fn((rows, cols), |(r, c)| { + ((r * cols + c) as f32 * 0.017 + seed).sin() + }) +} + +/// Shared attention geometry for the primitive tests. +pub(super) const NUM_Q: usize = 2; +pub(super) const HEAD_DIM: usize = 4; +pub(super) const REPS: usize = 1; +pub(super) const SCALE: f64 = 0.5; diff --git a/crates/larql-compute/src/attention/swa_tests/prefill.rs b/crates/larql-compute/src/attention/swa_tests/prefill.rs new file mode 100644 index 000000000..8c7cc7a5a --- /dev/null +++ b/crates/larql-compute/src/attention/swa_tests/prefill.rs @@ -0,0 +1,116 @@ +//! Windowed prefill attention. + +use super::{ramp, HEAD_DIM, NUM_Q, REPS, SCALE}; +use crate::attention::{gqa_attention, gqa_attention_decode_step, gqa_attention_windowed}; + +/// A window at least as wide as the sequence cannot bind, so the +/// windowed path must be bit-identical to the unwindowed one. This is +/// what makes the change safe to switch on globally: every +/// full-attention layer and every short context keeps its exact bits. +#[test] +fn prefill_window_wider_than_context_is_bit_identical() { + let seq = 6; + let q = ramp(seq, NUM_Q * HEAD_DIM, 0.1); + let k = ramp(seq, NUM_Q * HEAD_DIM, 0.2); + let v = ramp(seq, NUM_Q * HEAD_DIM, 0.3); + + let full = gqa_attention(&q, &k, &v, NUM_Q, HEAD_DIM, REPS, SCALE, seq); + for window in [None, Some(seq), Some(seq + 1), Some(1024)] { + let windowed = gqa_attention_windowed( + &q, &k, &v, NUM_Q, HEAD_DIM, REPS, SCALE, seq, None, None, window, + ); + assert_eq!( + full.iter().map(|v| v.to_bits()).collect::>(), + windowed.iter().map(|v| v.to_bits()).collect::>(), + "window {window:?} cannot bind at seq={seq} and must not change any bits" + ); + } +} + +/// A binding window must change the answer — and must change it only +/// where it binds. Query 0 sees one key either way, so its row is +/// untouched; later rows differ once the window drops a key. +#[test] +fn prefill_window_binds_only_after_the_window_length() { + let seq = 6; + let window = 2; + let q = ramp(seq, NUM_Q * HEAD_DIM, 0.1); + let k = ramp(seq, NUM_Q * HEAD_DIM, 0.2); + let v = ramp(seq, NUM_Q * HEAD_DIM, 0.3); + + let full = gqa_attention(&q, &k, &v, NUM_Q, HEAD_DIM, REPS, SCALE, seq); + let windowed = gqa_attention_windowed( + &q, + &k, + &v, + NUM_Q, + HEAD_DIM, + REPS, + SCALE, + seq, + None, + None, + Some(window), + ); + + for qi in 0..window { + assert_eq!( + full.row(qi).to_vec(), + windowed.row(qi).to_vec(), + "query {qi} has at most {window} keys, so the window cannot bind yet" + ); + } + let last = seq - 1; + assert_ne!( + full.row(last).to_vec(), + windowed.row(last).to_vec(), + "query {last} sees {seq} keys but the window admits {window} — it must differ" + ); +} + +/// The windowed result must equal running full attention over just the +/// keys the window admits. This is the definition of the operation, so +/// it catches an off-by-one in the slice that a "did it change?" test +/// would not. +#[test] +fn prefill_windowed_row_equals_full_attention_over_the_admitted_keys() { + let seq = 6; + let window = 3; + let q = ramp(seq, NUM_Q * HEAD_DIM, 0.1); + let k = ramp(seq, NUM_Q * HEAD_DIM, 0.2); + let v = ramp(seq, NUM_Q * HEAD_DIM, 0.3); + + let windowed = gqa_attention_windowed( + &q, + &k, + &v, + NUM_Q, + HEAD_DIM, + REPS, + SCALE, + seq, + None, + None, + Some(window), + ); + + // Reference for the final query: slice the admitted keys and run the + // ordinary decode step, which attends everything it is given. + let last = seq - 1; + let start = (last + 1) - window; + let q_last = q.slice(ndarray::s![last..last + 1, ..]).to_owned(); + let k_win = k.slice(ndarray::s![start..last + 1, ..]).to_owned(); + let v_win = v.slice(ndarray::s![start..last + 1, ..]).to_owned(); + let expected = gqa_attention_decode_step( + &q_last, &k_win, &v_win, NUM_Q, HEAD_DIM, REPS, SCALE, None, None, + ); + + for c in 0..NUM_Q * HEAD_DIM { + assert!( + (windowed[[last, c]] - expected[[0, c]]).abs() < 1e-5, + "col {c}: windowed prefill {} != full attention over admitted keys {}", + windowed[[last, c]], + expected[[0, c]] + ); + } +} diff --git a/crates/larql-compute/src/attention/swa_tests/rule.rs b/crates/larql-compute/src/attention/swa_tests/rule.rs new file mode 100644 index 000000000..901f7cb8c --- /dev/null +++ b/crates/larql-compute/src/attention/swa_tests/rule.rs @@ -0,0 +1,88 @@ +//! Which layers are windowed, and how wide. + +use crate::forward_overrides::effective_attention_window_for_layer; + +/// Gemma 3's 5:1 pattern: layers 5, 11, 17, … are full attention, the +/// rest slide. The window must follow the architecture, not the layer +/// index modulo anything this test invents. +/// +/// Uses the rope-scaled fixture specifically because it **declares a +/// real window** (512) across 6 layers, so the 5:1 boundary lands inside +/// the model. The plain Gemma-3 fixtures declare sliding layers with no +/// width, which resolves to `None` everywhere and would make this test +/// pass without ever exercising the windowed branch. +#[test] +fn window_follows_the_architecture_per_layer() { + let weights = larql_models::test_fixtures::make_test_q4k_weights_rope_scaled(); + let arch = &*weights.arch; + let declared = arch.sliding_window_size(); + assert_eq!( + declared, + Some(512), + "fixture must declare a real window or this test proves nothing" + ); + + let mut saw_windowed = 0usize; + let mut saw_full = 0usize; + for layer in 0..weights.num_layers { + let got = effective_attention_window_for_layer(arch, layer); + if arch.is_sliding_window_layer(layer) { + assert_eq!( + got, declared, + "layer {layer} slides, so it must report the declared window" + ); + saw_windowed += 1; + } else { + assert_eq!( + got, None, + "layer {layer} is a full-attention layer and must not be windowed" + ); + saw_full += 1; + } + } + // Both branches must actually have been taken. + assert!( + saw_windowed > 0 && saw_full > 0, + "fixture must contain both sliding and full layers (saw {saw_windowed} / {saw_full})" + ); +} + +/// A layer that declares itself sliding but supplies no width is +/// answered `None` (full attention) rather than `Some(0)` (attend +/// nothing). Several fixtures are in exactly this state, and the old +/// inline `unwrap_or(0)` in the Metal spec meant the same literal `0` +/// stood for both "no window" and "empty window". +#[test] +fn a_sliding_layer_without_a_declared_width_is_not_windowed() { + let weights = larql_inference_test_q4k_weights(); + let arch = &*weights.arch; + assert!( + arch.is_sliding_window_layer(0), + "fixture is expected to declare sliding layers" + ); + assert_eq!( + arch.sliding_window_size(), + None, + "fixture is expected to omit the window width" + ); + assert_eq!( + effective_attention_window_for_layer(arch, 0), + None, + "no width means no window — never an empty one" + ); +} + +fn larql_inference_test_q4k_weights() -> larql_models::ModelWeights { + larql_models::test_fixtures::make_test_q4k_weights() +} + +/// Architectures with no sliding layers at all are never windowed. +#[test] +fn non_sliding_architectures_are_never_windowed() { + let weights = larql_models::test_fixtures::make_test_q4k_weights_silu(); + let arch = &*weights.arch; + for layer in 0..weights.num_layers { + assert!(!arch.is_sliding_window_layer(layer)); + assert_eq!(effective_attention_window_for_layer(arch, layer), None); + } +} diff --git a/crates/larql-compute/src/attention/swa_tests/wiring.rs b/crates/larql-compute/src/attention/swa_tests/wiring.rs new file mode 100644 index 000000000..b19630a39 --- /dev/null +++ b/crates/larql-compute/src/attention/swa_tests/wiring.rs @@ -0,0 +1,76 @@ +//! The per-layer seam must actually pass the window on. + +use super::ramp; +use ndarray::Array2; + +/// End-to-end proof that `run_attention_block` *uses* the window rather +/// than merely being able to. +/// +/// Two hidden-state sequences that are identical inside the window and +/// differ only outside it must produce the same output at the final +/// position on a **sliding** layer, and different output on a +/// **full-attention** layer of the same model. Testing the primitives in +/// isolation cannot catch a seam that computes the window and then +/// forgets to pass it; this can. +#[test] +fn run_attention_block_honours_the_window_on_sliding_layers_only() { + let weights = larql_models::test_fixtures::make_test_q4k_weights_rope_scaled(); + let arch = &*weights.arch; + let window = arch + .sliding_window_size() + .expect("fixture declares a window"); + let hidden = weights.hidden_size; + + // Sequence long enough that the window drops the opening positions. + let seq = window + 8; + let base = ramp(seq, hidden, 0.11); + // Perturb only positions the window excludes for the FINAL query. + // Final query index is seq-1, so it admits [seq-window, seq-1]. + let mut perturbed = base.clone(); + let excluded_end = seq - window; + for r in 0..excluded_end { + for c in 0..hidden { + perturbed[[r, c]] += 0.5; + } + } + assert!(excluded_end > 0, "test needs positions outside the window"); + + let sliding_layer = (0..weights.num_layers) + .find(|&l| arch.is_sliding_window_layer(l)) + .expect("fixture has a sliding layer"); + let full_layer = (0..weights.num_layers) + .find(|&l| !arch.is_sliding_window_layer(l)) + .expect("fixture has a full-attention layer"); + + let last = seq - 1; + let run = |h: &Array2, layer: usize| -> Vec { + let view = larql_models::WeightsView::dense(&weights); + let (_, attn, _) = + crate::attention::block::run_attention_block(view, h, layer, false).expect("block ran"); + attn.row(last).to_vec() + }; + + let sliding_a = run(&base, sliding_layer); + let sliding_b = run(&perturbed, sliding_layer); + for (i, (a, b)) in sliding_a.iter().zip(sliding_b.iter()).enumerate() { + assert!( + (a - b).abs() < 1e-4, + "sliding layer {sliding_layer}, dim {i}: changing tokens OUTSIDE the \ + {window}-token window moved the output ({a} vs {b}) — the window is \ + not reaching the attention call" + ); + } + + let full_a = run(&base, full_layer); + let full_b = run(&perturbed, full_layer); + let moved = full_a + .iter() + .zip(full_b.iter()) + .any(|(a, b)| (a - b).abs() > 1e-4); + assert!( + moved, + "full-attention layer {full_layer} ignored a change to early tokens — it \ + should attend them, so either the control is degenerate or the window \ + is being applied to every layer" + ); +} diff --git a/crates/larql-compute/src/backend/decode.rs b/crates/larql-compute/src/backend/decode.rs index d527267f6..7ccfe15b9 100644 --- a/crates/larql-compute/src/backend/decode.rs +++ b/crates/larql-compute/src/backend/decode.rs @@ -77,7 +77,7 @@ impl DecodeStateDump { /// canonical residual state (`MarkovResidualEngine::recompute_kv`). /// /// Engines that treat K/V as **canonical** (e.g. `TurboQuantEngine`'s -/// compressed K/V, `UnlimitedContextEngine`'s in-window K/V) must +/// compressed K/V, `WindowedCheckpointEngine`'s in-window K/V) must /// use `Full`. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)] pub enum StateDumpMask { diff --git a/crates/larql-compute/src/cpu/ops/moe/latent_mask.rs b/crates/larql-compute/src/cpu/ops/moe/latent_mask.rs index bf10a7b13..200b20254 100644 --- a/crates/larql-compute/src/cpu/ops/moe/latent_mask.rs +++ b/crates/larql-compute/src/cpu/ops/moe/latent_mask.rs @@ -505,4 +505,119 @@ mod tests { assert!(!(r > 0.0 && r < 1.0), "r={r} must not enable the probe"); } } + + // ── Stats accumulation ────────────────────────────────────────────── + // + // `STATS`/`COSTATS` are process-global, and `set_env_override` is + // thread-local, so a parallel test cannot own them exclusively. These + // assert *structure* — the file was written, it parses, our layer is in + // it with a non-zero count — rather than exact totals, which another + // test's accumulation would perturb. + + struct EnvGuard; + impl Drop for EnvGuard { + fn drop(&mut self) { + options::clear_fast_path_overrides(); + } + } + + fn tmp(name: &str) -> std::path::PathBuf { + let dir = std::env::temp_dir().join("latent-mask-tests"); + std::fs::create_dir_all(&dir).unwrap(); + dir.join(format!("{name}-{}", std::process::id())) + } + + #[test] + fn stats_are_not_collected_unless_a_path_is_named() { + // The hot path must stay free: no env, no lock, no allocation. + let _g = EnvGuard; + options::clear_fast_path_overrides(); + // Nothing to assert but that it does not panic and does not write. + record_stats(0, &[true, false, true]); + dump_stats(); + } + + #[test] + fn recorded_survivals_reach_the_dumped_file() { + let _g = EnvGuard; + let path = tmp("stats.txt"); + let _ = std::fs::remove_file(&path); + options::set_env_override(ENV_LATENT_STATS, Some(path.to_str().unwrap())); + + // A layer index no other test uses, so our row is identifiable. + const LAYER: usize = 41; + record_stats(LAYER, &[true, false, true, false]); + record_stats(LAYER, &[true, false, false, true]); + dump_stats(); + + let body = std::fs::read_to_string(&path).expect("dump_stats must write the file"); + // `layer chan count` triples; channel 0 survived both calls. + let ours: Vec<&str> = body + .lines() + .filter(|l| l.starts_with(&format!("{LAYER} "))) + .collect(); + assert!(!ours.is_empty(), "our layer must appear: {body:.200}"); + let chan0 = ours + .iter() + .find(|l| l.starts_with(&format!("{LAYER} 0 "))) + .expect("channel 0 row"); + let count: u64 = chan0.split_whitespace().nth(2).unwrap().parse().unwrap(); + assert!(count >= 2, "channel 0 survived twice, got {count}"); + let _ = std::fs::remove_file(&path); + } + + #[test] + fn a_short_row_grows_to_fit_a_wider_mask() { + // Layers are discovered as they are first recorded, and a later call + // may carry more channels than the first — resizing must preserve the + // counts already accumulated rather than reallocating them away. + let _g = EnvGuard; + let path = tmp("grow.txt"); + let _ = std::fs::remove_file(&path); + options::set_env_override(ENV_LATENT_STATS, Some(path.to_str().unwrap())); + + const LAYER: usize = 43; + record_stats(LAYER, &[true, true]); + record_stats(LAYER, &[true, false, true, true]); + dump_stats(); + + let body = std::fs::read_to_string(&path).expect("file"); + let row0 = body + .lines() + .find(|l| l.starts_with(&format!("{LAYER} 0 "))) + .expect("channel 0"); + let c0: u64 = row0.split_whitespace().nth(2).unwrap().parse().unwrap(); + assert!(c0 >= 2, "channel 0 survived both calls, got {c0}"); + assert!( + body.lines().any(|l| l.starts_with(&format!("{LAYER} 3 "))), + "the widened row must reach channel 3" + ); + let _ = std::fs::remove_file(&path); + } + + #[test] + fn costats_need_their_own_switch_and_write_a_square_matrix() { + let _g = EnvGuard; + let stats = tmp("co-stats.txt"); + let co = tmp("co-pairs.bin"); + let _ = std::fs::remove_file(&stats); + let _ = std::fs::remove_file(&co); + options::set_env_override(ENV_LATENT_STATS, Some(stats.to_str().unwrap())); + options::set_env_override(ENV_LATENT_COSTATS, Some(co.to_str().unwrap())); + + // Sampled 1-in-N, so drive enough calls that at least one lands. + for _ in 0..(COSTATS_SAMPLE_EVERY * 2 + 1) { + record_stats(45, &[true, false, true, false]); + } + dump_stats(); + + let bytes = std::fs::read(&co).expect("costats file"); + assert!( + !bytes.is_empty(), + "a sampled call must have filled the matrix" + ); + assert_eq!(bytes.len() % 4, 0, "u32 counts, little-endian"); + let _ = std::fs::remove_file(&stats); + let _ = std::fs::remove_file(&co); + } } diff --git a/crates/larql-compute/src/cpu/ops/moe/math.rs b/crates/larql-compute/src/cpu/ops/moe/math.rs index 9d049682d..d1cd33b1e 100644 --- a/crates/larql-compute/src/cpu/ops/moe/math.rs +++ b/crates/larql-compute/src/cpu/ops/moe/math.rs @@ -1,7 +1,12 @@ //! Numeric primitives used by the MoE forward pass. //! -//! `pub(super)` keeps these module-private — `cpu_moe_forward` and the -//! per-expert helpers share them, nothing outside `moe/` should. +//! Most are `pub(super)`: `cpu_moe_forward` and the per-expert helpers share +//! them and nothing outside `moe/` should. +//! +//! `matmul_vec` and `softmax` are the exceptions. They are the router's +//! numerical recipe, and a VINDEX3 `BoundRouter` binds *these functions* so +//! that kernel-binding parity is a statement about the production kernel +//! rather than about two similar-looking loops agreeing. /// Dequantize a BF16 byte slice to f32. #[inline] @@ -67,7 +72,10 @@ pub(super) fn gelu_tanh(x: f32) -> f32 { /// `out_rows × in_cols` multiplies, repeated 8 experts × 60 layers per token, /// and BLAS sgemv hits the AMX tiles + SIMD fused-multiply-add pipeline that /// the scalar path misses entirely. -pub(super) fn matmul_vec(x: &[f32], w: &[f32], out_rows: usize, in_cols: usize) -> Vec { +/// Public so a VINDEX3 `BoundRouter` can bind *this* kernel rather than +/// reimplement a lookalike. Binding the real function is the difference +/// between proving kernel binding works and proving two similar loops agree. +pub fn matmul_vec(x: &[f32], w: &[f32], out_rows: usize, in_cols: usize) -> Vec { debug_assert_eq!(w.len(), out_rows * in_cols); debug_assert_eq!(x.len(), in_cols); if out_rows == 0 || in_cols == 0 { @@ -116,7 +124,10 @@ pub(super) fn matmul_vec_into( } /// Softmax in-place. -pub(super) fn softmax(v: &mut [f32]) { +/// Public for the same reason as [`matmul_vec`]: the router's numerical +/// recipe is scoring *and* its softmax, and reproducing one while +/// reimplementing the other would leave the comparison meaningless. +pub fn softmax(v: &mut [f32]) { let max = v.iter().copied().fold(f32::NEG_INFINITY, f32::max); let mut sum = 0.0f32; for x in v.iter_mut() { diff --git a/crates/larql-compute/src/cpu/ops/moe/mod.rs b/crates/larql-compute/src/cpu/ops/moe/mod.rs index b24c67f41..69718ff4f 100644 --- a/crates/larql-compute/src/cpu/ops/moe/mod.rs +++ b/crates/larql-compute/src/cpu/ops/moe/mod.rs @@ -15,7 +15,7 @@ mod cache; mod expert; mod forward; pub mod latent_mask; -mod math; +pub mod math; mod within_expert; pub use crate::cpu::ops::q4k_q8k_dot::{quantize_x_to_q8k, Q8KActivation}; @@ -24,6 +24,7 @@ pub use expert::{ run_single_expert_q4k_q8k_into, run_single_expert_with_norm, ExpertScratch, }; pub use forward::cpu_moe_forward; +pub use math::{matmul_vec as moe_score_experts, softmax as moe_softmax}; pub use within_expert::{ is_active as within_expert_active, set_current_layer, set_routing, ExpertFeatureSelector, WithinExpertRouting, diff --git a/crates/larql-compute/src/cpu/ops/q4k_matvec.rs b/crates/larql-compute/src/cpu/ops/q4k_matvec.rs index d8690107e..dd2a4aa4f 100644 --- a/crates/larql-compute/src/cpu/ops/q4k_matvec.rs +++ b/crates/larql-compute/src/cpu/ops/q4k_matvec.rs @@ -15,7 +15,11 @@ const Q4K_HEADER_BYTES: usize = 16; /// Decode f16 bits to f32, preserving subnormals (matches Metal's /// `decode_f16_metal`, which uses the hardware `half` → `float` cast). -fn f16_to_f32(bits: u16) -> f32 { +/// +/// Public because it is the *only* correct f16 decoder in the workspace: the +/// subnormal branch below fixes a 2× error that a from-scratch reimplementation +/// reproduces almost every time. Reuse it rather than writing another. +pub fn f16_to_f32(bits: u16) -> f32 { let sign = ((bits >> 15) & 1) as u32; let exp = ((bits >> 10) & 0x1F) as i32; let mant = (bits & 0x3FF) as u32; diff --git a/crates/larql-compute/src/ffn.rs b/crates/larql-compute/src/ffn.rs index ea574203c..876f7ec87 100644 --- a/crates/larql-compute/src/ffn.rs +++ b/crates/larql-compute/src/ffn.rs @@ -55,13 +55,32 @@ pub trait FfnBackend { /// For hybrid MoE layers: receive `h_post_attn` (post-attention, /// pre-FFN, unnormalized) and return the full layer output `h_out`. - /// Returns `None` to fall back to local dispatch. + /// + /// The three outcomes are deliberately distinct, and conflating any two of + /// them is how a missing operand becomes a plausible wrong answer: + /// + /// ```text + /// Ok(Some(h_out)) the layer executed + /// Ok(None) this backend does not serve this layer; fall back + /// Err(refusal) a required operation could not execute + /// ``` + /// + /// `Ok(None)` must never be used to report a refusal. It means *not + /// applicable* — the caller is free to run its own dispatch and the result + /// is still correct. `Err` means a routed operation was required and did + /// not happen, so any output the caller assembles is incomplete. A backend + /// that returned `Ok(None)` on failure would put the caller back where the + /// error channel was added to rescue it from. + /// + /// What to *do* about an `Err` is the caller's policy, not the backend's: + /// a strict engine propagates it, a best-effort one may log and degrade + /// explicitly. The backend's job is to report what happened. fn forward_moe_full_layer( &self, _layer: usize, _h_post_attn: &Array2, - ) -> Option> { - None + ) -> Result>, larql_execution::BoxRefusal> { + Ok(None) } } @@ -231,7 +250,9 @@ mod tests { } let ffn = StubFfn; let h = Array2::::zeros((1, 4)); - assert!(ffn.forward_moe_full_layer(0, &h).is_none()); + // The default is "not applicable", not a refusal — a backend that serves + // no MoE layer must let the caller dispatch locally rather than fail it. + assert!(matches!(ffn.forward_moe_full_layer(0, &h), Ok(None))); // Exercise the stub's required-method surface so the coverage // report reflects the trait-shape footprint, not just the // default-method probes. diff --git a/crates/larql-compute/src/forward/predict/raw.rs b/crates/larql-compute/src/forward/predict/raw.rs index a4bcdb1ec..dc4f8f43e 100644 --- a/crates/larql-compute/src/forward/predict/raw.rs +++ b/crates/larql-compute/src/forward/predict/raw.rs @@ -73,7 +73,7 @@ pub fn forward_raw_logits( /// position-0 token before layer 0. /// /// Mirrors the Python `prefill_to_layer(initial_residual=...)` API used by -/// `UnlimitedContextEngine`/Apollo. The prefix flows through every layer +/// `WindowedCheckpointEngine`/Apollo. The prefix flows through every layer /// along with the query tokens and participates in attention at each /// position — it's *not* a per-layer K/V injection, it's a residual /// prepend. diff --git a/crates/larql-compute/src/forward_overrides.rs b/crates/larql-compute/src/forward_overrides.rs index e2b3f6e11..0aa43634a 100644 --- a/crates/larql-compute/src/forward_overrides.rs +++ b/crates/larql-compute/src/forward_overrides.rs @@ -109,6 +109,37 @@ pub fn effective_rope_base_for_layer( } } +/// Per-layer attention window, honouring the `LARQL_FORCE_GLOBAL_LAYERS` +/// diagnostic override. `None` = attend the full context. +/// +/// **Single source of truth.** Both the CPU attention path and the Metal +/// pipeline spec resolve their window through here, because the two +/// answering this question independently is exactly how they drift: the +/// CPU path previously had no notion of a per-layer window at all, so a +/// Gemma-class model attended full history on layers the architecture +/// declares sliding while Metal masked them, and "GPU/CPU parity" was +/// undefined past the window. +/// +/// An architecture that declares a layer sliding but supplies no window +/// size is answered `None` — full attention. That combination is +/// incoherent (there is no window to slide) and the only alternatives +/// are to invent a size or to mask everything; both are worse than +/// attending the context we actually have. Some test fixtures are in +/// exactly this state, which is why it must be a deliberate branch +/// rather than an `unwrap_or(0)` that silently means "no window" in one +/// place and "empty window" in another. +pub fn effective_attention_window_for_layer( + arch: &dyn larql_models::ModelArchitecture, + layer: usize, +) -> Option { + if layer_forced_global(layer) || !arch.is_sliding_window_layer(layer) { + return None; + } + // A zero-width declared window means the same thing as no window; + // normalise it here so callers never see `Some(0)`. + arch.sliding_window_size().filter(|&w| w > 0) +} + /// Diagnostic position scale read from `LARQL_ROPE_POS_DIVISOR=`. Matches /// HF `rope_scaling = {rope_type: linear, factor: }`. Returns `1.0` when /// the env var is unset. Applied uniformly to every layer. diff --git a/crates/larql-compute/src/kquant_forward/cached.rs b/crates/larql-compute/src/kquant_forward/cached.rs index 6fb09fc57..cda7f289a 100644 --- a/crates/larql-compute/src/kquant_forward/cached.rs +++ b/crates/larql-compute/src/kquant_forward/cached.rs @@ -34,7 +34,7 @@ use larql_models::ModelWeights; use ndarray::Array2; use crate::attention::{ - decode::{gqa_attention_decode_step, run_attention_block_decode_step_backend}, + decode::{gqa_attention_decode_step_windowed, run_attention_block_decode_step_backend}, rope::apply_rope_partial_at_full, run_attention_with_kv_backend, }; @@ -100,7 +100,7 @@ pub fn predict_kquant_prefill( /// `state` is `Some`, populates per-layer `h_in` ([seq_len, hidden]), /// `k_new` ([seq_len, kv_dim]), `v_new` ([seq_len, kv_dim]) for every /// position in the prompt — engines (markov_residual, -/// unlimited_context, turbo_quant) use this to seed their state policy +/// windowed_checkpoint, turbo_quant) use this to seed their state policy /// from a single prefill pass without a follow-up CPU re-walk. When /// `state` is `None`, bit-identical to [`predict_kquant_prefill`]. pub fn predict_kquant_prefill_with_state( @@ -573,7 +573,7 @@ fn vec_to_2d_row(v: Vec) -> Array2 { /// signature stays format-agnostic. /// /// Used by `StandardEngine`'s coarse path and by research engines -/// (`MarkovResidual`, `UnlimitedContext`, `TurboQuant`) that want the +/// (`MarkovResidual`, `WindowedCheckpoint`, `TurboQuant`) that want the /// production decode kernel without inheriting the per-layer dispatch /// trait's cached-K/V shape. /// @@ -736,7 +736,12 @@ pub fn attention_decode_step_native( }; let softcap = arch.attn_logit_softcapping(); - let attn_out = gqa_attention_decode_step( + // Per-layer sliding window, same shared rule the Metal spec uses. + // This is the CPU coarse Q4K decode — the path `standard` takes on + // CpuBackend — so without it a Gemma-class model attends full history + // here while Metal masks. + let window = crate::forward_overrides::effective_attention_window_for_layer(arch, layer); + let attn_out = gqa_attention_decode_step_windowed( &q_rope, &k_concat, &v_concat, @@ -751,6 +756,7 @@ pub fn attention_decode_step_native( num_q, layer, ), + window, ); let attn_out_row: &[f32] = attn_out.row(0).to_slice().or_else(|| attn_out.as_slice())?; diff --git a/crates/larql-compute/src/kv_dispatch/cpu/dispatch.rs b/crates/larql-compute/src/kv_dispatch/cpu/dispatch.rs index 3a92de060..bc467eba2 100644 --- a/crates/larql-compute/src/kv_dispatch/cpu/dispatch.rs +++ b/crates/larql-compute/src/kv_dispatch/cpu/dispatch.rs @@ -40,7 +40,30 @@ impl KvDispatch for CpuBackend { h.append_row(k_row, v_row); } + /// Keep the tail `window_size` rows. + /// + /// Handles **both** cache shapes this backend allocates: the coarse + /// `CpuQ4kCacheHandle` (one handle for the whole model, a per-layer + /// `[rows, kv_dim]` pair inside) and the per-layer `CpuKvHandle`. Only the + /// second was handled before, so a windowed engine running on the coarse + /// path had its clip panic as a "foreign handle" — or, where the engine + /// never called clip at all, silently attended over its whole stream while + /// reporting a bounded window (issue #200). fn clip_kv(&self, handle: &mut KvHandle, window_size: usize) { + if let Some(coarse) = try_cpu_q4k_cache_mut(handle) { + for slot in coarse.cache.iter_mut() { + let Some((k, v)) = slot.as_mut() else { + continue; + }; + let rows = k.shape()[0]; + if rows > window_size { + let start = rows - window_size; + *k = k.slice(ndarray::s![start.., ..]).to_owned(); + *v = v.slice(ndarray::s![start.., ..]).to_owned(); + } + } + return; + } let h = cpu_handle_mut(handle); if h.rows > window_size { let start = h.rows - window_size; @@ -53,6 +76,20 @@ impl KvDispatch for CpuBackend { } } + fn truncate_kv(&self, handle: &mut KvHandle, len: usize) -> bool { + let h = cpu_handle_mut(handle); + // Growing only: `len` above the current row count would be asking + // this to invent rows, which is the one thing a rewind must not do. + if len > h.rows { + return false; + } + let kv_dim = h.kv_dim; + h.rows = len; + h.k_buf.truncate(len * kv_dim); + h.v_buf.truncate(len * kv_dim); + true + } + fn read_kv_to_host(&self, handle: &KvHandle) -> Option<(Array2, Array2)> { cpu_handle(handle).to_shared() } @@ -224,6 +261,67 @@ impl KvDispatch for CpuBackend { Some((h, handle)) } + /// Coarse prefill under an engine-requested window. + /// + /// Accepts the window only when it **cannot bind during the prompt** + /// (`token_ids.len() <= window`). That is the production shape — a + /// prompt inside the window, then decode sliding — and it lets the + /// engine keep the fused path for it. + /// + /// A prompt longer than the window would need per-query-position + /// masking inside `predict_kquant_prefill_with_state`, which this + /// entry point does not thread; rather than silently attend the full + /// prompt while the engine advertises a window, it declines and the + /// engine falls back to the per-layer path. Fail closed: a wrong + /// answer at full speed is worse than a right one at 2.4x. + fn coarse_prefill_windowed( + &self, + weights: &ModelWeights, + token_ids: &[u32], + index: Option<&dyn crate::KvIndex>, + window: Option, + ) -> Option<(Array2, KvHandle)> { + if let Some(w) = window { + if token_ids.len() > w { + return None; + } + } + self.coarse_prefill_with_state(weights, token_ids, index, None) + } + + /// Coarse decode under an engine-requested window. + /// + /// The cache rows are absolute stream positions carrying their own + /// RoPE, and softmax over keys is order-independent, so a sliding + /// window is exactly "drop the oldest rows". Trimming to `w - 1` + /// *before* the step leaves room for the row this step appends, so + /// the cache holds at most `w` rows afterwards — bounding attention + /// and K/V together, which is the whole contract a windowed engine + /// advertises. + fn coarse_decode_step_windowed( + &self, + weights: &ModelWeights, + token_id: u32, + index: Option<&dyn crate::KvIndex>, + handle: &mut KvHandle, + abs_position: usize, + window: Option, + ) -> Option> { + if let Some(w) = window { + if w == 0 { + return None; + } + // Decline a handle this backend did not mint before mutating + // anything — `clip_kv` would panic on a foreign one. + try_cpu_q4k_cache_mut(handle)?; + // Reuse the existing clip rather than a second trim: `clip_kv` + // already keeps the tail of both cache shapes, and two copies + // of "drop the oldest rows" is how they drift apart. + self.clip_kv(handle, w - 1); + } + self.coarse_decode_step(weights, token_id, index, handle, abs_position) + } + fn coarse_decode_step( &self, weights: &ModelWeights, @@ -264,7 +362,7 @@ impl KvDispatch for CpuBackend { /// coarse Q4K cache handle. The cache rows are indexed by absolute /// stream position (prefill row 0 onward), matching the trait /// contract engines rely on for boundary-checkpoint readback - /// (`UnlimitedContextEngine::close_window` under HOnly). Returns + /// (`WindowedCheckpointEngine::close_window` under HOnly). Returns /// `None` for foreign handle shapes (per-layer `CpuKvHandle` has no /// cross-layer cache) and for out-of-range `layer`/`pos`. fn read_kv_row_at( diff --git a/crates/larql-compute/src/kv_dispatch/cpu/handles.rs b/crates/larql-compute/src/kv_dispatch/cpu/handles.rs index 95cc44f1e..d52733980 100644 --- a/crates/larql-compute/src/kv_dispatch/cpu/handles.rs +++ b/crates/larql-compute/src/kv_dispatch/cpu/handles.rs @@ -230,6 +230,22 @@ impl KvHandleInner for CpuQ4kCacheHandle { "cpu-q4k" } + /// One handle, every layer — so sum the layers instead of taking the + /// trait default's single-layer estimate. `cached_len`/`kv_dim` above + /// deliberately report layer-0's shape (the dispatch surface asks for + /// "the" cache geometry), and multiplying those alone undercounted a + /// 28-layer model's K/V by 28×. + fn resident_bytes(&self) -> usize { + self.cache + .iter() + .filter_map(|o| o.as_ref()) + .map(|(k, v)| { + (k.shape()[0] * k.shape()[1] + v.shape()[0] * v.shape()[1]) + * std::mem::size_of::() + }) + .sum() + } + fn as_any(&self) -> &dyn std::any::Any { self } @@ -251,3 +267,7 @@ pub(super) fn try_cpu_q4k_cache_mut(h: &mut KvHandle) -> Option<&mut CpuQ4kCache .as_any_mut() .downcast_mut::() } + +#[cfg(test)] +#[path = "handles_resident_tests.rs"] +mod resident_bytes_tests; diff --git a/crates/larql-compute/src/kv_dispatch/cpu/handles_resident_tests.rs b/crates/larql-compute/src/kv_dispatch/cpu/handles_resident_tests.rs new file mode 100644 index 000000000..62e9cc6ae --- /dev/null +++ b/crates/larql-compute/src/kv_dispatch/cpu/handles_resident_tests.rs @@ -0,0 +1,78 @@ +//! Byte-accounting pins for the CPU K/V handles. +//! +//! Split out of `handles.rs` to keep that file within the per-file size +//! budget; `#[path]`-included from it. + +use super::*; +use crate::kv_dispatch::handles::KV_TENSORS_PER_LAYER; +use ndarray::Array2; + +fn cache_of(layers: usize, rows: usize, kv_dim: usize) -> CpuQ4kCacheHandle { + CpuQ4kCacheHandle { + cache: (0..layers) + .map(|_| { + Some(( + Array2::::zeros((rows, kv_dim)), + Array2::::zeros((rows, kv_dim)), + )) + }) + .collect(), + } +} + +/// Regression pin: this is ONE handle covering EVERY layer, so its +/// byte count must scale with layer count. The trait default (a +/// per-layer estimate from `cached_len × kv_dim`) reported layer-0 +/// only, understating a 28-layer model by 28× — which surfaced in +/// the bench as a fabricated "13× vs std-kv" saving for an engine +/// doing no compression whatsoever. +#[test] +fn resident_bytes_counts_every_layer_not_just_the_first() { + let (layers, rows, kv_dim) = (28usize, 100usize, 64usize); + let handle = cache_of(layers, rows, kv_dim); + + let per_layer_estimate = + handle.cached_len() * handle.kv_dim() * KV_TENSORS_PER_LAYER * size_of::(); + let expected = layers * rows * kv_dim * KV_TENSORS_PER_LAYER * size_of::(); + + assert_eq!(handle.resident_bytes(), expected); + assert_eq!( + handle.resident_bytes(), + per_layer_estimate * layers, + "whole-model handle must be exactly `layers`× the per-layer estimate" + ); +} + +#[test] +fn resident_bytes_ignores_unpopulated_layers() { + let mut handle = cache_of(4, 10, 8); + handle.cache[1] = None; + handle.cache[3] = None; + let expected = 2 * 10 * 8 * KV_TENSORS_PER_LAYER * size_of::(); + assert_eq!(handle.resident_bytes(), expected); +} + +#[test] +fn resident_bytes_is_zero_for_an_empty_cache() { + let handle = CpuQ4kCacheHandle { cache: Vec::new() }; + assert_eq!(handle.resident_bytes(), 0); + assert_eq!(handle.cached_len(), 0); + assert_eq!(handle.kv_dim(), 0); +} + +/// The per-layer handle keeps the trait default — one layer each, +/// summed by the engine across `num_layers` handles. +#[test] +fn per_layer_handle_uses_the_single_layer_default() { + let h = CpuKvHandle { + layer: 0, + k_buf: vec![0.0; 5 * 16], + v_buf: vec![0.0; 5 * 16], + rows: 5, + kv_dim: 16, + }; + assert_eq!( + h.resident_bytes(), + 5 * 16 * KV_TENSORS_PER_LAYER * size_of::() + ); +} diff --git a/crates/larql-compute/src/kv_dispatch/cpu/mod.rs b/crates/larql-compute/src/kv_dispatch/cpu/mod.rs index 0f676f6dc..629ecfbbe 100644 --- a/crates/larql-compute/src/kv_dispatch/cpu/mod.rs +++ b/crates/larql-compute/src/kv_dispatch/cpu/mod.rs @@ -29,5 +29,7 @@ mod dispatch; mod handles; #[cfg(test)] mod tests; +#[cfg(test)] +mod windowed_coarse_tests; pub use handles::{CpuKvHandle, CpuQ4kCacheHandle, CpuResidualHandle}; diff --git a/crates/larql-compute/src/kv_dispatch/cpu/tests.rs b/crates/larql-compute/src/kv_dispatch/cpu/tests.rs index a7d0e35e3..d61d298e1 100644 --- a/crates/larql-compute/src/kv_dispatch/cpu/tests.rs +++ b/crates/larql-compute/src/kv_dispatch/cpu/tests.rs @@ -109,6 +109,64 @@ fn clip_kv_with_no_state_is_a_no_op() { assert_eq!(h.cached_len(), 0); } +// ── truncate_kv ───────────────────────────────────────────────────────────── +// +// The inverse of an append: rewinds a partially-applied decode step. The +// contract that matters is *which* rows survive — `clip_kv` keeps the tail, +// this keeps the head — because a rewind that kept the wrong end would +// restore the row count while corrupting the cache, and a length assertion +// alone would not notice. + +#[test] +fn truncate_kv_keeps_the_head_where_clip_keeps_the_tail() { + let b = backend(); + let mut h = b.alloc_kv_buffer(0, 8, 2); + for i in 0..4u32 { + let f = i as f32; + b.append_kv(&mut h, &[f, f], &[f, f], i as usize); + } + assert!(b.truncate_kv(&mut h, 2)); + assert_eq!(h.cached_len(), 2); + let (k, v) = b.read_kv_to_host(&h).unwrap(); + // Rows 0 and 1 — the opposite end from `clip_kv_truncates_to_window_size`. + assert_eq!(k[[0, 0]], 0.0); + assert_eq!(k[[1, 0]], 1.0); + assert_eq!(v[[0, 0]], 0.0); + assert_eq!(v[[1, 0]], 1.0); +} + +#[test] +fn truncate_kv_to_the_current_length_is_a_no_op() { + let b = backend(); + let mut h = b.alloc_kv_buffer(0, 4, 2); + b.append_kv(&mut h, &[1.0, 2.0], &[3.0, 4.0], 0); + assert!(b.truncate_kv(&mut h, 1)); + assert_eq!(h.cached_len(), 1); + let (k, _) = b.read_kv_to_host(&h).unwrap(); + assert_eq!(k[[0, 0]], 1.0); +} + +#[test] +fn truncate_kv_to_zero_empties_the_cache() { + let b = backend(); + let mut h = b.alloc_kv_buffer(0, 4, 2); + b.append_kv(&mut h, &[1.0, 2.0], &[3.0, 4.0], 0); + assert!(b.truncate_kv(&mut h, 0)); + assert_eq!(h.cached_len(), 0); +} + +#[test] +fn truncate_kv_beyond_the_current_length_refuses() { + // Asking for more rows than exist is asking this to invent them. It must + // answer `false` rather than silently leave the handle short, because the + // caller reads `true` as "the cache is exactly what you asked for". + let b = backend(); + let mut h = b.alloc_kv_buffer(0, 8, 2); + b.append_kv(&mut h, &[1.0, 2.0], &[3.0, 4.0], 0); + assert!(!b.truncate_kv(&mut h, 5)); + assert_eq!(h.cached_len(), 1, "a refused rewind must change nothing"); +} + #[test] fn read_kv_to_host_returns_none_for_empty_handle() { let b = backend(); @@ -368,7 +426,7 @@ fn coarse_decode_step_returns_none_without_index() { /// `read_kv_row_at` on the coarse Q4K handle returns the row at the /// requested ABSOLUTE position, value-equal to the per-position K/V the -/// state dump captured — the contract `UnlimitedContextEngine`'s +/// state dump captured — the contract `WindowedCheckpointEngine`'s /// boundary-checkpoint readback depends on. #[test] fn read_kv_row_at_returns_absolute_rows_matching_state_dump() { diff --git a/crates/larql-compute/src/kv_dispatch/cpu/windowed_coarse_tests.rs b/crates/larql-compute/src/kv_dispatch/cpu/windowed_coarse_tests.rs new file mode 100644 index 000000000..5a86fca0c --- /dev/null +++ b/crates/larql-compute/src/kv_dispatch/cpu/windowed_coarse_tests.rs @@ -0,0 +1,183 @@ +//! The engine-requested window on the CPU fused path. +//! +//! A windowed engine promises two things at once: it attends at most +//! `window` positions, AND it holds at most that much K/V. Before this, +//! the only way to keep both promises was to leave the fused path for +//! the generic per-layer route — which costs ~2.4x, and on a +//! host-delegating backend runs the whole forward on the CPU. These pin +//! that the fused path now keeps both promises, and that it declines +//! rather than half-keeping them. + +use super::*; +use crate::kv_dispatch::{KvDispatch, KvHandle}; +use crate::test_fixtures::make_q4k_fixture_index; +use crate::CpuBackend; +use larql_models::test_fixtures::make_test_q4k_weights_silu; + +/// Rows held by the coarse cache behind a handle, for the first +/// populated layer. +fn cached_rows(handle: &KvHandle) -> usize { + handle + .as_inner() + .as_any() + .downcast_ref::() + .and_then(|h| h.cache.iter().flatten().next()) + .map(|(k, _)| k.shape()[0]) + .unwrap_or(0) +} + +/// `window: None` must be the existing path exactly — same entry point, +/// same bits — or switching engines onto the windowed call would be a +/// silent behaviour change for every unwindowed run. +#[test] +fn unwindowed_request_is_the_plain_coarse_path() { + let weights = make_test_q4k_weights_silu(); + let index = make_q4k_fixture_index(&weights); + let b = CpuBackend; + let prompt = [0u32, 1, 2, 3]; + + let (h_plain, _) = b + .coarse_prefill(&weights, &prompt, Some(&index)) + .expect("plain coarse prefill"); + let (h_win, _) = b + .coarse_prefill_windowed(&weights, &prompt, Some(&index), None) + .expect("windowed(None) coarse prefill"); + + assert_eq!( + h_plain.iter().map(|v| v.to_bits()).collect::>(), + h_win.iter().map(|v| v.to_bits()).collect::>(), + "window=None must be bit-identical to the plain coarse path" + ); +} + +/// A prompt that fits inside the window is accepted: the window cannot +/// bind during the prompt, so the fused prefill is already correct. +#[test] +fn prefill_accepts_a_prompt_that_fits_the_window() { + let weights = make_test_q4k_weights_silu(); + let index = make_q4k_fixture_index(&weights); + let b = CpuBackend; + let prompt = [0u32, 1, 2, 3]; + + assert!( + b.coarse_prefill_windowed(&weights, &prompt, Some(&index), Some(prompt.len())) + .is_some(), + "a prompt exactly filling the window still cannot bind mid-prompt" + ); + assert!(b + .coarse_prefill_windowed(&weights, &prompt, Some(&index), Some(prompt.len() + 1)) + .is_some()); +} + +/// A prompt LONGER than the window is declined, not silently attended in +/// full. This entry point does not thread per-query-position masking +/// into the fused prefill, so accepting would mean answering with full +/// attention while the engine advertises a window. +#[test] +fn prefill_declines_a_prompt_longer_than_the_window() { + let weights = make_test_q4k_weights_silu(); + let index = make_q4k_fixture_index(&weights); + let b = CpuBackend; + let prompt = [0u32, 1, 2, 3]; + + assert!( + b.coarse_prefill_windowed(&weights, &prompt, Some(&index), Some(prompt.len() - 1)) + .is_none(), + "a prompt exceeding the window must decline so the engine falls back" + ); +} + +/// The memory half of the contract: after enough steps the cache stops +/// growing and sits at the window, instead of the unbounded growth the +/// plain coarse path shows. +#[test] +fn decode_bounds_the_cache_at_the_window() { + let weights = make_test_q4k_weights_silu(); + let index = make_q4k_fixture_index(&weights); + let b = CpuBackend; + let prompt = [0u32, 1, 2]; + let window = 4usize; + + let (_, mut handle) = b + .coarse_prefill_windowed(&weights, &prompt, Some(&index), Some(window)) + .expect("prefill"); + + for (pos, step) in (prompt.len()..).zip(0..8) { + b.coarse_decode_step_windowed(&weights, 1u32, Some(&index), &mut handle, pos, Some(window)) + .unwrap_or_else(|| panic!("windowed decode step {step}")); + assert!( + cached_rows(&handle) <= window, + "step {step}: cache grew to {} rows, past the {window}-row window", + cached_rows(&handle) + ); + } + assert_eq!( + cached_rows(&handle), + window, + "a saturated window should sit exactly at its bound" + ); +} + +/// The control: the same run WITHOUT a window grows past it. Without +/// this, the test above would pass on a cache that simply never filled. +#[test] +fn decode_without_a_window_grows_past_it() { + let weights = make_test_q4k_weights_silu(); + let index = make_q4k_fixture_index(&weights); + let b = CpuBackend; + let prompt = [0u32, 1, 2]; + let window = 4usize; + + let (_, mut handle) = b + .coarse_prefill(&weights, &prompt, Some(&index)) + .expect("prefill"); + for pos in prompt.len()..prompt.len() + 8 { + b.coarse_decode_step(&weights, 1u32, Some(&index), &mut handle, pos) + .expect("decode step"); + } + assert!( + cached_rows(&handle) > window, + "unwindowed decode should have grown past {window} rows, got {}", + cached_rows(&handle) + ); +} + +/// A zero-width window is refused rather than treated as "no window" — +/// the sentinel confusion that already cost a bug in the Metal spec. +#[test] +fn a_zero_window_is_refused() { + let weights = make_test_q4k_weights_silu(); + let index = make_q4k_fixture_index(&weights); + let b = CpuBackend; + let prompt = [0u32, 1, 2]; + let (_, mut handle) = b + .coarse_prefill(&weights, &prompt, Some(&index)) + .expect("prefill"); + assert!( + b.coarse_decode_step_windowed( + &weights, + 1u32, + Some(&index), + &mut handle, + prompt.len(), + Some(0) + ) + .is_none(), + "window=0 must refuse, not fall through to unbounded" + ); +} + +/// A handle this backend did not mint is declined before any mutation. +#[test] +fn a_foreign_handle_is_declined_not_clipped() { + let weights = make_test_q4k_weights_silu(); + let index = make_q4k_fixture_index(&weights); + let b = CpuBackend; + // Per-layer handle, not the coarse whole-model one. + let mut handle = b.alloc_kv_buffer(0, 8, 16); + assert!( + b.coarse_decode_step_windowed(&weights, 1u32, Some(&index), &mut handle, 0, Some(4)) + .is_none(), + "a per-layer handle must be declined, not clipped as if it were coarse" + ); +} diff --git a/crates/larql-compute/src/kv_dispatch/dispatch.rs b/crates/larql-compute/src/kv_dispatch/dispatch.rs index 7214e70a5..9a5ce6c4f 100644 --- a/crates/larql-compute/src/kv_dispatch/dispatch.rs +++ b/crates/larql-compute/src/kv_dispatch/dispatch.rs @@ -22,7 +22,7 @@ use crate::PerLayerDecodeState; /// - `markov_residual_codec`: same as `markov_residual`; on /// window-overflow the evicted rows get codec-encoded into /// `cold_encoded[l]`. -/// - `unlimited_context`: `k_new_per_layer[l]` / `v_new_per_layer[l]` +/// - `windowed_checkpoint`: `k_new_per_layer[l]` / `v_new_per_layer[l]` /// are appended to the per-layer K/V cache; `h_in_per_layer` is /// unused but populated for API uniformity (cheap blit). /// - `turbo_quant`: `k_new_per_layer[l]` / `v_new_per_layer[l]` @@ -73,6 +73,25 @@ pub trait KvDispatch { unimplemented!("clip_kv not implemented for this backend") } + /// Drop cached rows past `len`, keeping the **first** `len` in order. + /// Returns whether the handle now holds exactly `len` rows. + /// + /// The inverse of an append: this rewinds a partially-applied decode + /// step so a caller that could not finish the token leaves the cache + /// describing the token sequence it described before. Distinct from + /// [`Self::clip_kv`], which keeps the *tail* to enforce a sliding + /// window — that one moves the cache forward, this one moves it back. + /// + /// The default answers `false` rather than panicking, because "this + /// backend cannot rewind" is a legitimate state that a caller must + /// handle (by invalidating the cache) rather than a programming + /// error. Returning `true` without having rewound is the one + /// unacceptable answer: it tells the caller a corrupt cache is sound. + fn truncate_kv(&self, handle: &mut KvHandle, len: usize) -> bool { + let _ = (handle, len); + false + } + /// Read the full K/V back to host memory as a `(K, V)` pair. /// Blocking copy on GPU backends; identity on CPU. Should NOT be /// used in hot loops — it's the cross-backend escape hatch for @@ -82,6 +101,39 @@ pub trait KvDispatch { None } + /// Bytes of K/V this backend holds in its own storage, i.e. K/V that + /// is NOT reachable by measuring the [`KvHandle`]s an engine owns. + /// + /// Backends whose fused pipelines keep the cache internally hand out + /// a sentinel handle that measures zero (Metal's coarse whole-model + /// handle is the live example). An engine summing its handles then + /// reports no K/V at all, which reads as "this engine is free" in a + /// memory comparison when the truth is "its K/V lives one layer + /// down". Engines add this to their own accounting so the two + /// dispatch shapes are comparable. + /// + /// Default 0 = every byte this backend holds is reachable through a + /// handle, so adding it would double-count. + fn backend_resident_kv_bytes(&self) -> usize { + 0 + } + + /// Whether this backend implements the **per-layer** surface + /// (`attention_step`, `attention_prefill`, `append_kv`, …) by + /// forwarding to the host CPU rather than running native kernels. + /// + /// `MetalBackend` answers `true`: only its `coarse_*` family is + /// GPU-resident, so an engine that declines the coarse path runs its + /// whole forward on the CPU while callers still believe they + /// selected a GPU backend. Diagnostics need to be able to say that + /// out loud; without it a windowed engine reports `[metal (GPU)]` + /// over a pure-CPU measurement. + /// + /// Default `false` — a backend is assumed to mean what it says. + fn per_layer_is_host_delegated(&self) -> bool { + false + } + // ── Attention primitives ──────────────────────────────────────── /// Run one decode-step attention: Q (one row, pre-projection @@ -236,7 +288,7 @@ pub trait KvDispatch { // without changing this trait surface. // // Engines that DO need per-layer control (MarkovResidual, - // UnlimitedContext, TurboQuant — recompute, checkpoint, codec + // WindowedCheckpoint, TurboQuant — recompute, checkpoint, codec // mechanisms) continue to use the per-layer `attention_prefill` / // `attention_step` intents. // @@ -281,6 +333,56 @@ pub trait KvDispatch { None } + /// Coarse prefill under an **engine-requested sliding window** — a + /// window the caller imposes across every layer, distinct from the + /// architecture's own per-layer SWA (which backends read from the + /// arch and this composes with, narrowest wins). + /// + /// A windowed engine promises two things: it attends at most `window` + /// positions, AND it holds at most that much K/V. The window-less + /// `coarse_prefill` can honour neither, so a windowed engine had to + /// decline the fused path entirely and take the generic per-layer + /// route — which on a host-delegating backend runs the whole forward + /// on the CPU, costing ~2.4x. This is the entry point that lets a + /// backend keep the fused path *and* the window. + /// + /// **Fail closed.** The default supports only `window: None`, and + /// answers `None` when a real window is requested — "this backend + /// cannot bound what you asked me to bound", which the engine reads + /// as "take the per-layer path", exactly today's behaviour. A + /// backend that returns `Some` for a windowed request is asserting + /// it enforced BOTH halves of the contract. + fn coarse_prefill_windowed( + &self, + weights: &ModelWeights, + token_ids: &[u32], + index: Option<&dyn crate::KvIndex>, + window: Option, + ) -> Option<(Array2, KvHandle)> { + match window { + None => self.coarse_prefill(weights, token_ids, index), + Some(_) => None, + } + } + + /// One coarse decode step under an engine-requested sliding window. + /// Same contract and same fail-closed default as + /// [`Self::coarse_prefill_windowed`]. + fn coarse_decode_step_windowed( + &self, + weights: &ModelWeights, + token_id: u32, + index: Option<&dyn crate::KvIndex>, + handle: &mut KvHandle, + abs_position: usize, + window: Option, + ) -> Option> { + match window { + None => self.coarse_decode_step(weights, token_id, index, handle, abs_position), + Some(_) => None, + } + } + /// Coarse prefill **with per-layer state capture** — same fast /// path as [`Self::coarse_prefill`] but also populates `state` /// (when `Some`) with per-layer h_in (residual entering each @@ -290,7 +392,7 @@ pub trait KvDispatch { /// shape `[seq_len, hidden]` and each entry in /// `state.k_new_per_layer` / `v_new_per_layer` has shape /// `[seq_len, kv_dim_for_layer]`. Engines (markov_residual, - /// unlimited_context, turbo_quant) read these to seed their + /// windowed_checkpoint, turbo_quant) read these to seed their /// state policy without re-running prefill on CPU. /// /// Default impl delegates to [`Self::coarse_prefill`] and leaves @@ -316,7 +418,7 @@ pub trait KvDispatch { /// /// Engines that need per-layer state to enforce their state /// policy — `markov_residual` (stores h_in per layer), - /// `turbo_quant` (compresses per-layer K/V), `unlimited_context` + /// `turbo_quant` (compresses per-layer K/V), `windowed_checkpoint` /// (snapshots K/V at window boundaries) — pass `Some(&mut state)` /// to extract per-layer state without re-running compute on CPU. /// @@ -371,7 +473,7 @@ pub trait KvDispatch { /// cache. Returns `(k_row, v_row)` as flat `Vec` of length /// `kv_dim_for_layer`. Used by engines running under /// [`crate::StateDumpMask::HOnly`] that need to snapshot specific - /// K/V positions on demand (e.g. `UnlimitedContextEngine`'s + /// K/V positions on demand (e.g. `WindowedCheckpointEngine`'s /// `close_window` checkpoint emission). /// /// Default returns `None` — backends without an internal kv cache diff --git a/crates/larql-compute/src/kv_dispatch/handles.rs b/crates/larql-compute/src/kv_dispatch/handles.rs index 01e5e2229..04346c192 100644 --- a/crates/larql-compute/src/kv_dispatch/handles.rs +++ b/crates/larql-compute/src/kv_dispatch/handles.rs @@ -41,6 +41,13 @@ impl KvHandle { self.inner.backend_name() } + /// Total K/V bytes held by this handle, across every layer it covers. + /// Engines sum this over their handles rather than re-deriving it from + /// `cached_len × kv_dim`, which is only right for per-layer handles. + pub fn resident_bytes(&self) -> usize { + self.inner.resident_bytes() + } + /// Downcast access for backend implementations. Engines never call /// this; only the backend that allocated the handle should. pub fn as_inner(&self) -> &dyn KvHandleInner { @@ -61,10 +68,27 @@ pub trait KvHandleInner: Send + Sync + std::any::Any { fn cached_len(&self) -> usize; fn kv_dim(&self) -> usize; fn backend_name(&self) -> &'static str; + + /// Total K/V bytes this handle holds, across **every layer it covers**. + /// + /// The default is one layer's worth — correct for the per-layer + /// dispatch handles, where an engine owns one handle per layer and + /// sums them. Whole-model handles (one handle, all layers) MUST + /// override: applying the per-layer formula to them undercounts by + /// a factor of `num_layers`, which in the bench read as a large + /// compression win for engines that were doing no compression at all. + fn resident_bytes(&self) -> usize { + self.cached_len() * self.kv_dim() * KV_TENSORS_PER_LAYER * std::mem::size_of::() + } + fn as_any(&self) -> &dyn std::any::Any; fn as_any_mut(&mut self) -> &mut dyn std::any::Any; } +/// K and V — the two tensors cached per position, per layer. Named so +/// the sizing arithmetic doesn't read as a bare `2` between dimensions. +pub const KV_TENSORS_PER_LAYER: usize = 2; + /// Opaque handle to a residual upload (used by `apollo` for boundary /// residuals). Same pattern as [`KvHandle`]. pub struct ResidualHandle { diff --git a/crates/larql-compute/src/kv_dispatch/tests.rs b/crates/larql-compute/src/kv_dispatch/tests.rs index 19cef74f2..3f323b2e0 100644 --- a/crates/larql-compute/src/kv_dispatch/tests.rs +++ b/crates/larql-compute/src/kv_dispatch/tests.rs @@ -239,6 +239,17 @@ fn default_clip_kv_panics() { backend.clip_kv(&mut handle, 2); } +#[test] +fn default_truncate_kv_declines_instead_of_panicking() { + // Unlike its siblings above, this default must *return* rather than + // panic: "this backend cannot rewind" is a state the caller handles by + // invalidating the cache, not a programming error. A panic here would + // turn an unported backend into a crash on every refused decode step. + let backend = StubKvBackend; + let mut handle = stub_kv_handle(0, 4); + assert!(!backend.truncate_kv(&mut handle, 0)); +} + #[test] #[should_panic(expected = "compressed_kv_append not implemented")] fn default_compressed_kv_append_panics() { @@ -394,6 +405,151 @@ fn default_read_kv_row_at_returns_none() { assert!(backend.read_kv_row_at(&handle, 0, 0).is_none()); } +// ── Accounting / honesty defaults ──────────────────────────────── + +#[test] +fn default_backend_resident_kv_bytes_is_zero() { + // 0 means "every byte I hold is reachable through a handle", so an + // engine adding this to its own accounting can't double-count. A + // backend holding K/V outside its handles must override. + assert_eq!(StubKvBackend.backend_resident_kv_bytes(), 0); +} + +#[test] +fn default_per_layer_is_host_delegated_is_false() { + // A backend is assumed to mean what it says; only one that runs the + // per-layer surface on the host (MetalBackend) overrides to `true`, + // so diagnostics can stop reporting a CPU measurement as GPU. + assert!(!StubKvBackend.per_layer_is_host_delegated()); +} + +// ── Windowed coarse defaults: the fail-closed pair ─────────────── +// +// `window: None` means "no bound requested" and must forward to the +// window-less method; `Some(w)` means "bound this", which the default +// cannot honour and must therefore decline. Answering `Some` to a +// windowed request asserts BOTH halves of the contract (attend at most +// `window`, hold at most `window`) — so the default has to say no. + +/// Stub whose window-less coarse methods and `attention_step` succeed. +/// With the all-default `StubKvBackend` both arms of the windowed +/// defaults return `None`, so a test could not tell delegation from +/// refusal. Here delegation returns `Some` and refusal returns `None`. +struct CoarseCapableBackend; + +impl KvDispatch for CoarseCapableBackend { + fn clip_kv(&self, _handle: &mut KvHandle, _window_size: usize) {} + + fn attention_step( + &self, + _weights: larql_models::WeightsView, + query: &Array2, + _kv: &mut KvHandle, + _layer: usize, + _abs_position: usize, + _index: Option<&dyn crate::KvIndex>, + ) -> Option> { + Some(query.clone()) + } + + fn coarse_prefill( + &self, + weights: &larql_models::ModelWeights, + _token_ids: &[u32], + _index: Option<&dyn crate::KvIndex>, + ) -> Option<(Array2, KvHandle)> { + Some(( + Array2::zeros((1, weights.hidden_size)), + stub_kv_handle(1, weights.hidden_size), + )) + } + + fn coarse_decode_step( + &self, + weights: &larql_models::ModelWeights, + _token_id: u32, + _index: Option<&dyn crate::KvIndex>, + _handle: &mut KvHandle, + _abs_position: usize, + ) -> Option> { + Some(Array2::zeros((1, weights.hidden_size))) + } +} + +#[test] +fn default_coarse_prefill_windowed_forwards_when_no_window_requested() { + let weights = make_test_weights(); + let backend = CoarseCapableBackend; + let result = backend.coarse_prefill_windowed(&weights, &[0u32, 1], None, None); + let (hidden, handle) = result.expect("window: None must forward to coarse_prefill"); + assert_eq!(hidden.shape(), &[1, weights.hidden_size]); + assert_eq!(handle.backend_name(), "stub"); +} + +#[test] +fn default_coarse_prefill_windowed_declines_a_real_window() { + let weights = make_test_weights(); + let backend = CoarseCapableBackend; + // Fail closed: the window-less prefill above succeeds, so a `None` + // here can only come from the refusal arm, not from an unsupported + // coarse path. + assert!( + backend + .coarse_prefill_windowed(&weights, &[0u32, 1], None, Some(4)) + .is_none(), + "a backend that cannot bound the window must decline it" + ); +} + +#[test] +fn default_coarse_decode_step_windowed_forwards_when_no_window_requested() { + let weights = make_test_weights(); + let backend = CoarseCapableBackend; + let mut handle = stub_kv_handle(0, weights.hidden_size); + let result = backend.coarse_decode_step_windowed(&weights, 0, None, &mut handle, 0, None); + let hidden = result.expect("window: None must forward to coarse_decode_step"); + assert_eq!(hidden.shape(), &[1, weights.hidden_size]); +} + +#[test] +fn default_coarse_decode_step_windowed_declines_a_real_window() { + let weights = make_test_weights(); + let backend = CoarseCapableBackend; + let mut handle = stub_kv_handle(0, weights.hidden_size); + assert!( + backend + .coarse_decode_step_windowed(&weights, 0, None, &mut handle, 0, Some(4)) + .is_none(), + "a backend that cannot bound the window must decline it" + ); +} + +#[test] +fn default_attention_step_windowed_clips_then_returns_the_hidden() { + // The None-propagation branch is covered above; this drives the + // other half of the default decomposition — `attention_step` + // succeeded, so `clip_kv` runs and the hidden is handed back. + let weights = make_test_weights(); + let backend = CoarseCapableBackend; + let mut handle = stub_kv_handle(0, weights.hidden_size); + let query = Array2::from_shape_fn((1, weights.hidden_size), |(_, j)| j as f32) as Array2; + let out = backend + .attention_step_windowed( + larql_models::WeightsView::dense(&weights), + &query, + &mut handle, + 0, + 0, + 4, + None, + ) + .expect("attention_step returned Some, so the windowed default must too"); + assert_eq!( + out, query, + "the default must pass the hidden through unchanged" + ); +} + // ── Inner handle as_any / as_any_mut surface ───────────────────── #[test] diff --git a/crates/larql-compute/src/lib.rs b/crates/larql-compute/src/lib.rs index 386911721..374079103 100644 --- a/crates/larql-compute/src/lib.rs +++ b/crates/larql-compute/src/lib.rs @@ -140,6 +140,7 @@ pub use quant_route::FormatRoute; pub use cpu::ops::linalg::{cholesky, cholesky_inverse, cholesky_solve, ridge_decomposition_solve}; pub use cpu::ops::moe::{quantize_x_to_q8k, Q8KActivation}; +pub use cpu::ops::q4k_matvec::f16_to_f32; pub use cpu::ops::vector::{cosine, dot, norm}; pub use cpu::CpuBackend; diff --git a/crates/larql-compute/src/pipeline_layer.rs b/crates/larql-compute/src/pipeline_layer.rs index 253076005..ffadf322d 100644 --- a/crates/larql-compute/src/pipeline_layer.rs +++ b/crates/larql-compute/src/pipeline_layer.rs @@ -7,7 +7,13 @@ //! Per-layer override resolution (env vars vs arch defaults) lives in //! [`crate::forward_overrides`]; this module consumes those helpers. -use crate::forward_overrides::{effective_rope_base_for_layer, layer_forced_global}; +use crate::forward_overrides::effective_rope_base_for_layer; + +/// Wire value the kernel's attention spec uses for "no sliding window — +/// attend the whole context". Named because it sits in the same field as +/// a real window width, where a bare `0` reads as an empty window. +const NO_ATTENTION_WINDOW: usize = 0; + use crate::{ FullPipelineLayer, MoeLayerWeights, MoeRoutingPolicy, MoeWeightLayout, QuantFormat, QuantWeight, }; @@ -58,12 +64,12 @@ pub fn build_arch_params<'a>( } else { (layer_hd as f64 * rotary_frac) as usize }; - let force_global = layer_forced_global(layer); - let sw = if !force_global && arch.is_sliding_window_layer(layer) { - arch.sliding_window_size().unwrap_or(0) - } else { - 0 - }; + // Resolved through the shared rule so the Metal spec and the CPU + // attention path cannot answer this differently. `0` is the wire + // sentinel for "attend everything" in the kernel's spec struct, which + // is what `None` means here. + let sw = crate::forward_overrides::effective_attention_window_for_layer(arch, layer) + .unwrap_or(NO_ATTENTION_WINDOW); let layer_scalar = arch .layer_scalar_key(layer) .and_then(|k| weights.vectors.get(&k)) diff --git a/crates/larql-compute/src/state_handle.rs b/crates/larql-compute/src/state_handle.rs index 5f09e05a6..a79bc6304 100644 --- a/crates/larql-compute/src/state_handle.rs +++ b/crates/larql-compute/src/state_handle.rs @@ -102,7 +102,7 @@ pub enum RowLocation { /// - **Canonical**: discarding it loses the conversation. Examples: /// `MarkovResidualEngine`'s residual stream, `TurboQuantEngine`'s /// compressed K/V (destructive), `StandardEngine`'s K/V tensors, -/// `UnlimitedContextEngine`'s in-window K/V. +/// `WindowedCheckpointEngine`'s in-window K/V. /// - **Derivative**: discardable. The engine can rebuild it from /// canonical state + model weights without changing its output /// distribution. Example: `MarkovResidualEngine`'s hot K/V cache diff --git a/crates/larql-core/Cargo.toml b/crates/larql-core/Cargo.toml index 7362afb55..92a36ae44 100644 --- a/crates/larql-core/Cargo.toml +++ b/crates/larql-core/Cargo.toml @@ -32,20 +32,6 @@ criterion = { version = "0.5", features = ["html_reports"] } name = "graph" harness = false -[[example]] -name = "algorithm_demo" - [[example]] name = "bench_graph" -[[example]] -name = "edge_demo" - -[[example]] -name = "filter_demo" - -[[example]] -name = "graph_demo" - -[[example]] -name = "serialization_demo" diff --git a/crates/larql-demos/Cargo.toml b/crates/larql-demos/Cargo.toml new file mode 100644 index 000000000..a7739df3b --- /dev/null +++ b/crates/larql-demos/Cargo.toml @@ -0,0 +1,246 @@ +# Runnable demonstrations of larql's shipped capabilities. +# +# Split out of the individual crates' `examples/` so that "how do I use +# this" has one home, separate from benchmarks/diagnostics (which stay +# next to the code they exercise) and from research probes (which live +# in chris-experiments, pinned to the larql revision that produced their +# verdict). +[package] +name = "larql-demos" +version.workspace = true +edition.workspace = true +authors.workspace = true +license.workspace = true +description = "Runnable demonstrations of larql's shipped capabilities" +publish = false +autoexamples = false # folders mirror the source crate; all declared below + +[dependencies] +larql-boundary = { path = "../larql-boundary" } +larql-compute = { path = "../larql-compute" } +larql-core = { path = "../larql-core" } +larql-inference = { path = "../larql-inference" } +larql-kv = { path = "../larql-kv" } +larql-lql = { path = "../larql-lql" } +larql-models = { path = "../larql-models" } +larql-router-protocol = { path = "../larql-router-protocol" } +larql-server = { path = "../larql-server" } +larql-vindex = { path = "../larql-vindex" } + +safetensors = "0.7" +serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } +tokenizers = "0.21" +tracing-subscriber = { version = "0.3", features = ["env-filter"] } + +# The larql-server demos are async clients. +axum = { version = "0.8", features = ["ws"] } +tokio = { version = "1", features = ["full"] } +tokio-stream = "0.1" +tonic = "0.13" +tower = { version = "0.5", features = ["limit"] } + +# BLAS backend per target, mirroring larql-inference. The compute demos +# carry `extern crate blas_src;`, so it has to be linked even where the +# crate is featureless. +[target.'cfg(any(target_os = "linux", target_os = "freebsd"))'.dependencies] +ndarray = { version = "0.16", features = ["blas"] } +blas-src = { version = "0.10", features = ["openblas"], default-features = false } +openblas-src = { version = "0.10", features = ["system"] } + +[target.'cfg(target_os = "macos")'.dependencies] +ndarray = { version = "0.16", features = ["blas"] } +blas-src = { version = "0.10", features = ["accelerate"] } + +[target.'cfg(target_os = "windows")'.dependencies] +ndarray = "0.16" +blas-src = { version = "0.10", default-features = false } + +[features] +default = [] +# Cascade the flags the demos themselves test for, so `#[cfg(feature = ..)]` +# in a demo resolves against a real feature rather than warning as unexpected. +gpu = ["larql-inference/gpu", "larql-vindex/gpu"] +msgpack = ["larql-core/msgpack"] + +# --- boundary -------------------------------------------------- +[[example]] +name = "encode_decode" +path = "examples/boundary/encode_decode.rs" + +[[example]] +name = "gate_decision" +path = "examples/boundary/gate_decision.rs" + +# --- compute --------------------------------------------------- +[[example]] +name = "demo_architecture" +path = "examples/compute/demo_architecture.rs" + +[[example]] +name = "demo_basic" +path = "examples/compute/demo_basic.rs" + +[[example]] +name = "demo_ridge_solve" +path = "examples/compute/demo_ridge_solve.rs" + +# --- core ------------------------------------------------------ +[[example]] +name = "algorithm_demo" +path = "examples/core/algorithm_demo.rs" + +[[example]] +name = "edge_demo" +path = "examples/core/edge_demo.rs" + +[[example]] +name = "filter_demo" +path = "examples/core/filter_demo.rs" + +[[example]] +name = "graph_demo" +path = "examples/core/graph_demo.rs" + +[[example]] +name = "serialization_demo" +path = "examples/core/serialization_demo.rs" + +# --- inference ------------------------------------------------- +[[example]] +name = "attention_demo" +path = "examples/inference/attention_demo.rs" + +[[example]] +name = "ave_demo" +path = "examples/inference/ave_demo.rs" + +[[example]] +name = "backend_demo" +path = "examples/inference/backend_demo.rs" + +[[example]] +name = "chat_demo" +path = "examples/inference/chat_demo.rs" + +[[example]] +name = "clustering_demo" +path = "examples/inference/clustering_demo.rs" + +[[example]] +name = "detok_demo" +path = "examples/inference/detok_demo.rs" + +[[example]] +name = "eos_demo" +path = "examples/inference/eos_demo.rs" + +[[example]] +name = "experts_demo" +path = "examples/inference/experts_demo.rs" + +[[example]] +name = "ffn_cache_demo" +path = "examples/inference/ffn_cache_demo.rs" + +[[example]] +name = "inference_demo" +path = "examples/inference/inference_demo.rs" + +[[example]] +name = "mech_interp_demo" +path = "examples/inference/mech_interp_demo.rs" + +[[example]] +name = "pair_matching_demo" +path = "examples/inference/pair_matching_demo.rs" + +[[example]] +name = "sampling_demo" +path = "examples/inference/sampling_demo.rs" + +[[example]] +name = "streaming_demo" +path = "examples/inference/streaming_demo.rs" + +# --- kv -------------------------------------------------------- +[[example]] +name = "engine_ladder" +path = "examples/kv/engine_ladder.rs" + +# --- lql ------------------------------------------------------- +[[example]] +name = "compact_demo" +path = "examples/lql/compact_demo.rs" + +[[example]] +name = "compile_demo" +path = "examples/lql/compile_demo.rs" + +[[example]] +name = "lql_demo" +path = "examples/lql/lql_demo.rs" + +[[example]] +name = "parser_demo" +path = "examples/lql/parser_demo.rs" + +[[example]] +name = "refine_demo" +path = "examples/lql/refine_demo.rs" + +[[example]] +name = "trace_demo" +path = "examples/lql/trace_demo.rs" + +# --- models ---------------------------------------------------- +[[example]] +name = "architecture_demo" +path = "examples/models/architecture_demo.rs" + +[[example]] +name = "demo_loading" +path = "examples/models/demo_loading.rs" + +[[example]] +name = "demo_tensor_keys" +path = "examples/models/demo_tensor_keys.rs" + +# --- server ---------------------------------------------------- +[[example]] +name = "embed_demo" +path = "examples/server/embed_demo.rs" + +[[example]] +name = "openai_demo" +path = "examples/server/openai_demo.rs" + +[[example]] +name = "server_demo" +path = "examples/server/server_demo.rs" + +[[example]] +name = "shard_query_demo" +path = "examples/server/shard_query_demo.rs" + +# --- vindex ---------------------------------------------------- +[[example]] +name = "demo_features" +path = "examples/vindex/demo_features.rs" + +[[example]] +name = "demo_memit_solve" +path = "examples/vindex/demo_memit_solve.rs" + +[[example]] +name = "mmap_demo" +path = "examples/vindex/mmap_demo.rs" + +[[example]] +name = "q4k_demo" +path = "examples/vindex/q4k_demo.rs" + +[[example]] +name = "walker_demo" +path = "examples/vindex/walker_demo.rs" + diff --git a/crates/larql-demos/README.md b/crates/larql-demos/README.md new file mode 100644 index 000000000..66114a0c8 --- /dev/null +++ b/crates/larql-demos/README.md @@ -0,0 +1,56 @@ +# larql-demos + +Runnable demonstrations of larql's shipped capabilities — one home for +"how do I use this", filed under the crate whose capability each one +shows. + +Each folder has its own README listing every demo, what it shows, and +what it needs to run — including measured runtimes, so a nine-minute demo +is not mistaken for a hung one. + +| folder | | demos | weight-free | +|---|---|---:|---:| +| [`boundary/`](examples/boundary/README.md) | Boundary codec | 2 | 2 | +| [`compute/`](examples/compute/README.md) | Compute kernels and solvers | 3 | 3 | +| [`core/`](examples/core/README.md) | Knowledge-graph core | 5 | 5 | +| [`inference/`](examples/inference/README.md) | Inference engine | 14 | 6 | +| [`kv/`](examples/kv/README.md) | KV engines | 1 | 1 | +| [`lql/`](examples/lql/README.md) | LQL query layer | 6 | 3 | +| [`models/`](examples/models/README.md) | Architecture detection | 3 | 3 | +| [`server/`](examples/server/README.md) | Serving surface | 4 | 3 | +| [`vindex/`](examples/vindex/README.md) | Vindex format and store | 5 | 5 | +| | | **43** | **31** | + +```sh +cargo run -p larql-demos --example chat_demo +``` + +Folders are not auto-discovered by cargo, so every demo is declared as an +explicit `[[example]]` in `Cargo.toml` with its path. Adding a demo means +adding four lines there — deliberately, so the inventory stays visible. + +The weight-free demos run in CI on every platform. The rest compile in CI +but need a real vindex to execute; they take `--vindex PATH` and fail by +name when it is missing, rather than surfacing a bare `NotFound`. + +## What is deliberately not here + +**Benchmarks, diagnostics and parity harnesses** stay in their own +crate's `examples/` — `bench_*`, `debug_*`, `profile_*`, `*_parity`, +`compare_*`, `membw_probe` and friends. They exercise the engine rather +than showing how to use it, so they belong next to the code they measure, +where a change and its benchmark move together. + +**Research probes** live in `chris-experiments/larql_probes`, pinned to +the larql revision that produced their recorded verdict. A probe answers +a question once; a demo is documentation that has to keep working. The +two have opposite maintenance contracts, which is why they no longer +share a directory. + +`apollo_rd_backend` was briefly filed here and has moved to the probes: +it is named after a backend but is really the compute half of a +chris-experiments script, and is useless without it. + +The dividing question when adding something here: *would a new user run +this to understand larql?* If it only makes sense while chasing a +specific result, it is a probe, not a demo. diff --git a/crates/larql-demos/examples/boundary/README.md b/crates/larql-demos/examples/boundary/README.md new file mode 100644 index 000000000..0bc0abaea --- /dev/null +++ b/crates/larql-demos/examples/boundary/README.md @@ -0,0 +1,18 @@ +# Boundary codec demos + +Run any of these from a larql checkout: + +```sh +cargo run -p larql-demos --example encode_decode +``` + +| demo | what it shows | run | +|---|---|---| +| `encode_decode` | Encode and decode a synthetic residual with both `bf16` and `int8_clip3sigma`. | weight-free · 0.1s | +| `gate_decision` | Gate decisions for four boundary types, matching the Exp 43 continuation tests. | weight-free · 0.1s | + +None of these need model weights — they run in well under a second and are the quickest way to see the surface working. + +--- + +Demos that need a vindex take `--vindex PATH`; point them at any model you have under `output/`. They fail by name if the path is missing rather than surfacing a bare `NotFound`. diff --git a/crates/larql-boundary/examples/encode_decode.rs b/crates/larql-demos/examples/boundary/encode_decode.rs similarity index 97% rename from crates/larql-boundary/examples/encode_decode.rs rename to crates/larql-demos/examples/boundary/encode_decode.rs index e1e74cfe7..dae9fbcba 100644 --- a/crates/larql-boundary/examples/encode_decode.rs +++ b/crates/larql-demos/examples/boundary/encode_decode.rs @@ -1,7 +1,7 @@ //! Encode and decode a synthetic residual with both bf16 and int8_clip3sigma. //! //! Run with: -//! cargo run -p larql-boundary --example encode_decode +//! cargo run -p larql-demos --example encode_decode use larql_boundary::codec::{bf16, int8}; diff --git a/crates/larql-boundary/examples/gate_decision.rs b/crates/larql-demos/examples/boundary/gate_decision.rs similarity index 98% rename from crates/larql-boundary/examples/gate_decision.rs rename to crates/larql-demos/examples/boundary/gate_decision.rs index 2cb62c6bb..46867b595 100644 --- a/crates/larql-boundary/examples/gate_decision.rs +++ b/crates/larql-demos/examples/boundary/gate_decision.rs @@ -1,7 +1,7 @@ //! Show gate decisions for four boundary types matching Exp 43 continuation tests. //! //! Run with: -//! cargo run -p larql-boundary --example gate_decision +//! cargo run -p larql-demos --example gate_decision use larql_boundary::{ gate::{BoundaryDecision, BoundaryGateConfig}, diff --git a/crates/larql-demos/examples/compute/README.md b/crates/larql-demos/examples/compute/README.md new file mode 100644 index 000000000..cbe605173 --- /dev/null +++ b/crates/larql-demos/examples/compute/README.md @@ -0,0 +1,19 @@ +# Compute kernels and solvers demos + +Run any of these from a larql checkout: + +```sh +cargo run -p larql-demos --example demo_architecture +``` + +| demo | what it shows | run | +|---|---|---| +| `demo_architecture` | Guided tour of larql-compute's major design decisions. | weight-free · 0.5s | +| `demo_basic` | Auto-detect the backend and run basic operations. | weight-free · 0.2s | +| `demo_ridge_solve` | `ridge_decomposition_solve` — the closed-form ridge solve underlying MEMIT-style weight edits. | weight-free · 0.3s | + +None of these need model weights — they run in well under a second and are the quickest way to see the surface working. + +--- + +Demos that need a vindex take `--vindex PATH`; point them at any model you have under `output/`. They fail by name if the path is missing rather than surfacing a bare `NotFound`. diff --git a/crates/larql-compute/examples/demo_architecture.rs b/crates/larql-demos/examples/compute/demo_architecture.rs similarity index 98% rename from crates/larql-compute/examples/demo_architecture.rs rename to crates/larql-demos/examples/compute/demo_architecture.rs index f6e50f164..96a64d699 100644 --- a/crates/larql-compute/examples/demo_architecture.rs +++ b/crates/larql-demos/examples/compute/demo_architecture.rs @@ -11,8 +11,8 @@ //! 8. KV-cached decode (dual Q4_K/Q8 path) //! //! Usage: -//! cargo run --release -p larql-compute --example demo_architecture -//! cargo run --release --features metal -p larql-compute --example demo_architecture +//! cargo run --release -p larql-demos --example demo_architecture +//! cargo run --release --features metal -p larql-demos --example demo_architecture extern crate blas_src; diff --git a/crates/larql-compute/examples/demo_basic.rs b/crates/larql-demos/examples/compute/demo_basic.rs similarity index 96% rename from crates/larql-compute/examples/demo_basic.rs rename to crates/larql-demos/examples/compute/demo_basic.rs index 935e3d4ca..ba8ac9125 100644 --- a/crates/larql-compute/examples/demo_basic.rs +++ b/crates/larql-demos/examples/compute/demo_basic.rs @@ -1,7 +1,7 @@ //! Demo: auto-detect backend and run basic operations. //! //! Usage: -//! cargo run --release -p larql-compute --example demo_basic +//! cargo run --release -p larql-demos --example demo_basic //! cargo run --release -p larql-compute --features metal --example demo_basic extern crate blas_src; diff --git a/crates/larql-compute/examples/demo_ridge_solve.rs b/crates/larql-demos/examples/compute/demo_ridge_solve.rs similarity index 98% rename from crates/larql-compute/examples/demo_ridge_solve.rs rename to crates/larql-demos/examples/compute/demo_ridge_solve.rs index 625702eba..b967714df 100644 --- a/crates/larql-compute/examples/demo_ridge_solve.rs +++ b/crates/larql-demos/examples/compute/demo_ridge_solve.rs @@ -3,7 +3,7 @@ //! //! Solves ΔW = T^T (K K^T + λI)^{-1} K //! -//! Run: cargo run --release -p larql-compute --example demo_ridge_solve +//! Run: cargo run --release -p larql-demos --example demo_ridge_solve //! //! Walks three regimes: //! 1. Orthonormal keys → exact reconstruction. diff --git a/crates/larql-demos/examples/core/README.md b/crates/larql-demos/examples/core/README.md new file mode 100644 index 000000000..b63253187 --- /dev/null +++ b/crates/larql-demos/examples/core/README.md @@ -0,0 +1,21 @@ +# Knowledge-graph core demos + +Run any of these from a larql checkout: + +```sh +cargo run -p larql-demos --example algorithm_demo +``` + +| demo | what it shows | run | +|---|---|---| +| `algorithm_demo` | Shortest path, merge, subgraph, connected components. | weight-free · 0.2s | +| `edge_demo` | Edge construction, metadata, compact serialization. | weight-free · 0.1s | +| `filter_demo` | Filter a graph — select edges by confidence, layer, relation. | weight-free · 0.2s | +| `graph_demo` | Build, query, traverse and serialize a knowledge graph. | weight-free · 0.2s | +| `serialization_demo` | JSON vs MessagePack, packed binary, CSV, format detection, bytes API. | weight-free · 0.2s | + +None of these need model weights — they run in well under a second and are the quickest way to see the surface working. + +--- + +Demos that need a vindex take `--vindex PATH`; point them at any model you have under `output/`. They fail by name if the path is missing rather than surfacing a bare `NotFound`. diff --git a/crates/larql-core/examples/algorithm_demo.rs b/crates/larql-demos/examples/core/algorithm_demo.rs similarity index 98% rename from crates/larql-core/examples/algorithm_demo.rs rename to crates/larql-demos/examples/core/algorithm_demo.rs index 62b352cff..2df37b92e 100644 --- a/crates/larql-core/examples/algorithm_demo.rs +++ b/crates/larql-demos/examples/core/algorithm_demo.rs @@ -1,6 +1,6 @@ //! Algorithm demo — shortest path, merge, subgraph, connected components. //! -//! Run: cargo run --release -p larql-core --example algorithm_demo +//! Run: cargo run --release -p larql-demos --example algorithm_demo use larql_core::*; diff --git a/crates/larql-core/examples/edge_demo.rs b/crates/larql-demos/examples/core/edge_demo.rs similarity index 97% rename from crates/larql-core/examples/edge_demo.rs rename to crates/larql-demos/examples/core/edge_demo.rs index abfa0e7c6..e53eb1147 100644 --- a/crates/larql-core/examples/edge_demo.rs +++ b/crates/larql-demos/examples/core/edge_demo.rs @@ -1,6 +1,6 @@ //! Edge demo — construction, metadata, compact serialization. //! -//! Run: cargo run --release -p larql-core --example edge_demo +//! Run: cargo run --release -p larql-demos --example edge_demo use larql_core::*; diff --git a/crates/larql-core/examples/filter_demo.rs b/crates/larql-demos/examples/core/filter_demo.rs similarity index 98% rename from crates/larql-core/examples/filter_demo.rs rename to crates/larql-demos/examples/core/filter_demo.rs index 724d58ffa..8bd9668bc 100644 --- a/crates/larql-core/examples/filter_demo.rs +++ b/crates/larql-demos/examples/core/filter_demo.rs @@ -1,6 +1,6 @@ //! Demonstrate graph filtering — select edges by confidence, layer, relation, etc. //! -//! Run: cargo run -p larql-core --example filter_demo +//! Run: cargo run -p larql-demos --example filter_demo use larql_core::*; diff --git a/crates/larql-core/examples/graph_demo.rs b/crates/larql-demos/examples/core/graph_demo.rs similarity index 98% rename from crates/larql-core/examples/graph_demo.rs rename to crates/larql-demos/examples/core/graph_demo.rs index bfaf309c0..b68a41fc6 100644 --- a/crates/larql-core/examples/graph_demo.rs +++ b/crates/larql-demos/examples/core/graph_demo.rs @@ -1,6 +1,6 @@ //! Graph engine demo — build, query, traverse, and serialize a knowledge graph. //! -//! Run: cargo run --release -p larql-core --example graph_demo +//! Run: cargo run --release -p larql-demos --example graph_demo use larql_core::*; diff --git a/crates/larql-core/examples/serialization_demo.rs b/crates/larql-demos/examples/core/serialization_demo.rs similarity index 98% rename from crates/larql-core/examples/serialization_demo.rs rename to crates/larql-demos/examples/core/serialization_demo.rs index 9058b4362..c38c88588 100644 --- a/crates/larql-core/examples/serialization_demo.rs +++ b/crates/larql-demos/examples/core/serialization_demo.rs @@ -1,6 +1,6 @@ //! Serialization demo — JSON vs MessagePack, packed binary, CSV, format detection, bytes API. //! -//! Run: cargo run --release -p larql-core --example serialization_demo +//! Run: cargo run --release -p larql-demos --example serialization_demo use larql_core::*; diff --git a/crates/larql-demos/examples/inference/README.md b/crates/larql-demos/examples/inference/README.md new file mode 100644 index 000000000..538a48038 --- /dev/null +++ b/crates/larql-demos/examples/inference/README.md @@ -0,0 +1,41 @@ +# Inference engine demos + +Run any of these from a larql checkout: + +```sh +cargo run -p larql-demos --example attention_demo +``` + +| demo | what it shows | run | +|---|---|---| +| `attention_demo` | The fused online-softmax attention kernel in action. | weight-free · 0.2s | +| `ave_demo` | Arithmetic Virtual Expert against a real Q4_K vindex. Writes `bench/aim-validation/`. | needs a vindex · 51s | +| `backend_demo` | Auto-calibrated hybrid CPU/Metal dispatch. | weight-free · 0.4s | +| `chat_demo` | Multi-turn conversation through `ChatSession`. | needs a vindex · 16s | +| `clustering_demo` | Clustering and relation discovery. | weight-free · 0.2s | +| `detok_demo` | Preserve word spacing across streamed tokens. | weight-free · 0.2s | +| `eos_demo` | The EOS detector halting generation correctly. | `--vindex PATH` · 10s | +| `experts_demo` | WASM expert registry — structured op+args calls across all experts. | needs the wasm build, below | +| `ffn_cache_demo` | FFN L1 cache behaviour, hit/miss stats, patch safety. | `--model ID --vindex PATH` | +| `inference_demo` | Forward pass from safetensors weights. | needs weights · 6s | +| `mech_interp_demo` | Capture, lens, neighbours, ablate, steer, patch. | weight-free · 0.1s | +| `pair_matching_demo` | Pair-based relation matching. | weight-free · 0.2s | +| `sampling_demo` | Greedy vs temperature vs top-p on one prompt. | `--vindex PATH` · 9s | +| `streaming_demo` | Print each token as the model emits it. | needs a vindex · 5s | + +6 of these need no model weights and run in well under a second, so they are the quickest way to see the surface working. + +## `experts_demo` needs a build first + +It loads WASM experts from `crates/larql-experts/target/wasm32-wasip1/release`, +which is not produced by a normal workspace build: + +```sh +cd crates/larql-experts && cargo build --target wasm32-wasip1 --release +``` + +Without it the demo exits naming the directory it looked in. + +--- + +Demos that need a vindex take `--vindex PATH`; point them at any model you have under `output/`. They fail by name if the path is missing rather than surfacing a bare `NotFound`. diff --git a/crates/larql-inference/examples/attention_demo.rs b/crates/larql-demos/examples/inference/attention_demo.rs similarity index 99% rename from crates/larql-inference/examples/attention_demo.rs rename to crates/larql-demos/examples/inference/attention_demo.rs index cf253f6e9..ff6e44d7b 100644 --- a/crates/larql-inference/examples/attention_demo.rs +++ b/crates/larql-demos/examples/inference/attention_demo.rs @@ -8,7 +8,7 @@ //! - Causal masking: each position only sees past tokens //! //! Usage: -//! cargo run -p larql-inference --example attention_demo +//! cargo run -p larql-demos --example attention_demo use ndarray::Array2; diff --git a/crates/larql-inference/examples/ave_demo.rs b/crates/larql-demos/examples/inference/ave_demo.rs similarity index 100% rename from crates/larql-inference/examples/ave_demo.rs rename to crates/larql-demos/examples/inference/ave_demo.rs diff --git a/crates/larql-inference/examples/backend_demo.rs b/crates/larql-demos/examples/inference/backend_demo.rs similarity index 98% rename from crates/larql-inference/examples/backend_demo.rs rename to crates/larql-demos/examples/inference/backend_demo.rs index d1d0fbbc4..6b419e318 100644 --- a/crates/larql-inference/examples/backend_demo.rs +++ b/crates/larql-demos/examples/inference/backend_demo.rs @@ -7,8 +7,8 @@ //! - Cold vs warm performance with cached buffers //! //! Usage: -//! cargo run --release -p larql-inference --example backend_demo -//! cargo run --release -p larql-inference --example backend_demo --features metal +//! cargo run --release -p larql-demos --example backend_demo +//! cargo run --release -p larql-demos --example backend_demo --features metal use ndarray::Array2; use std::time::Instant; diff --git a/crates/larql-inference/examples/chat_demo.rs b/crates/larql-demos/examples/inference/chat_demo.rs similarity index 100% rename from crates/larql-inference/examples/chat_demo.rs rename to crates/larql-demos/examples/inference/chat_demo.rs diff --git a/crates/larql-inference/examples/clustering_demo.rs b/crates/larql-demos/examples/inference/clustering_demo.rs similarity index 98% rename from crates/larql-inference/examples/clustering_demo.rs rename to crates/larql-demos/examples/inference/clustering_demo.rs index d714d0c14..c9c31986b 100644 --- a/crates/larql-inference/examples/clustering_demo.rs +++ b/crates/larql-demos/examples/inference/clustering_demo.rs @@ -3,7 +3,7 @@ //! Demonstrates the clustering pipeline: k-means, labeling, pair matching. //! Uses synthetic data to show how the pipeline works end-to-end. //! -//! Run: cargo run -p larql-inference --example clustering_demo +//! Run: cargo run -p larql-demos --example clustering_demo use larql_vindex::clustering::labeling::detect_entity_pattern; use larql_vindex::clustering::{ diff --git a/crates/larql-inference/examples/detok_demo.rs b/crates/larql-demos/examples/inference/detok_demo.rs similarity index 98% rename from crates/larql-inference/examples/detok_demo.rs rename to crates/larql-demos/examples/inference/detok_demo.rs index d63552281..a453d65f3 100644 --- a/crates/larql-inference/examples/detok_demo.rs +++ b/crates/larql-demos/examples/inference/detok_demo.rs @@ -15,7 +15,7 @@ //! [`Detokenizer`] fixes both by holding the cumulative ID list and //! emitting only the freshly-grown suffix on each `push`. //! -//! Usage: cargo run --release -p larql-inference --example detok_demo +//! Usage: cargo run --release -p larql-demos --example detok_demo use larql_inference::Detokenizer; use tokenizers::Tokenizer; diff --git a/crates/larql-inference/examples/eos_demo.rs b/crates/larql-demos/examples/inference/eos_demo.rs similarity index 88% rename from crates/larql-inference/examples/eos_demo.rs rename to crates/larql-demos/examples/inference/eos_demo.rs index 6feeee1f7..a0eef1bfb 100644 --- a/crates/larql-inference/examples/eos_demo.rs +++ b/crates/larql-demos/examples/inference/eos_demo.rs @@ -51,6 +51,18 @@ fn main() -> Result<(), Box> { i += 1; } + // The default above is a convenience, not a guarantee — it names one + // artifact out of many a checkout might hold. Fail by name rather than + // letting `open_inference_vindex` surface a bare `NotFound` that says + // neither which path was tried nor how to change it. + if !vindex_path.exists() { + return Err(format!( + "vindex not found: {}\npass --vindex to point at one you have", + vindex_path.display() + ) + .into()); + } + let mut model = InferenceModel::load("google/gemma-3-4b-it")?; let num_layers = model.weights().num_layers; let tokenizer = model.tokenizer().clone(); diff --git a/crates/larql-inference/examples/experts_demo.rs b/crates/larql-demos/examples/inference/experts_demo.rs similarity index 99% rename from crates/larql-inference/examples/experts_demo.rs rename to crates/larql-demos/examples/inference/experts_demo.rs index 3883ce3e8..9322bb08c 100644 --- a/crates/larql-inference/examples/experts_demo.rs +++ b/crates/larql-demos/examples/inference/experts_demo.rs @@ -14,7 +14,7 @@ //! cd crates/larql-experts && cargo build --target wasm32-wasip1 --release //! //! Then run from the repo root: -//! cargo run -p larql-inference --example experts_demo +//! cargo run -p larql-demos --example experts_demo use std::path::PathBuf; use std::time::Instant; diff --git a/crates/larql-inference/examples/ffn_cache_demo.rs b/crates/larql-demos/examples/inference/ffn_cache_demo.rs similarity index 98% rename from crates/larql-inference/examples/ffn_cache_demo.rs rename to crates/larql-demos/examples/inference/ffn_cache_demo.rs index 3fdeea64b..433bce2e1 100644 --- a/crates/larql-inference/examples/ffn_cache_demo.rs +++ b/crates/larql-demos/examples/inference/ffn_cache_demo.rs @@ -6,7 +6,7 @@ //! 3. Patched session — INSERT'd slot bypasses cache for correctness //! //! Usage: -//! cargo run --release -p larql-inference --example ffn_cache_demo -- \ +//! cargo run --release -p larql-demos --example ffn_cache_demo -- \ //! --model google/gemma-3-4b-it \ //! --vindex path/to/gemma3-4b.vindex diff --git a/crates/larql-inference/examples/inference_demo.rs b/crates/larql-demos/examples/inference/inference_demo.rs similarity index 98% rename from crates/larql-inference/examples/inference_demo.rs rename to crates/larql-demos/examples/inference/inference_demo.rs index a188111de..8d654a9d3 100644 --- a/crates/larql-inference/examples/inference_demo.rs +++ b/crates/larql-demos/examples/inference/inference_demo.rs @@ -5,7 +5,7 @@ //! //! Requires a model in the HuggingFace cache (e.g. google/gemma-3-4b-it). //! -//! Run: cargo run --release -p larql-inference --example inference_demo +//! Run: cargo run --release -p larql-demos --example inference_demo use std::time::Instant; diff --git a/crates/larql-inference/examples/mech_interp_demo.rs b/crates/larql-demos/examples/inference/mech_interp_demo.rs similarity index 99% rename from crates/larql-inference/examples/mech_interp_demo.rs rename to crates/larql-demos/examples/inference/mech_interp_demo.rs index 21143534d..533cb578d 100644 --- a/crates/larql-inference/examples/mech_interp_demo.rs +++ b/crates/larql-demos/examples/inference/mech_interp_demo.rs @@ -19,7 +19,7 @@ //! generation with the hook firing on every layer of every step. Used //! here to show steered output diverging from the baseline. //! -//! Usage: `cargo run --release -p larql-inference --example mech_interp_demo` +//! Usage: `cargo run --release -p larql-demos --example mech_interp_demo` //! //! All numbers are illustrative — the synthetic weights aren't a real //! language model. The point is to exercise every primitive end-to-end so diff --git a/crates/larql-inference/examples/pair_matching_demo.rs b/crates/larql-demos/examples/inference/pair_matching_demo.rs similarity index 99% rename from crates/larql-inference/examples/pair_matching_demo.rs rename to crates/larql-demos/examples/inference/pair_matching_demo.rs index 446bbbbfd..766367820 100644 --- a/crates/larql-inference/examples/pair_matching_demo.rs +++ b/crates/larql-demos/examples/inference/pair_matching_demo.rs @@ -3,7 +3,7 @@ //! Demonstrates matching (input, output) token pairs against //! Wikidata triples and WordNet relations to label clusters. //! -//! Run: cargo run -p larql-inference --example pair_matching_demo +//! Run: cargo run -p larql-demos --example pair_matching_demo use larql_vindex::clustering::pair_matching::{ label_clusters_from_pairs, load_reference_databases, RelationDatabase, diff --git a/crates/larql-inference/examples/sampling_demo.rs b/crates/larql-demos/examples/inference/sampling_demo.rs similarity index 89% rename from crates/larql-inference/examples/sampling_demo.rs rename to crates/larql-demos/examples/inference/sampling_demo.rs index 9121f7e9d..ab9ae7379 100644 --- a/crates/larql-inference/examples/sampling_demo.rs +++ b/crates/larql-demos/examples/inference/sampling_demo.rs @@ -54,6 +54,18 @@ fn main() -> Result<(), Box> { i += 1; } + // The default above is a convenience, not a guarantee — it names one + // artifact out of many a checkout might hold. Fail by name rather than + // letting `open_inference_vindex` surface a bare `NotFound` that says + // neither which path was tried nor how to change it. + if !vindex_path.exists() { + return Err(format!( + "vindex not found: {}\npass --vindex to point at one you have", + vindex_path.display() + ) + .into()); + } + let mut model = InferenceModel::load("google/gemma-3-4b-it")?; let num_layers = model.weights().num_layers; let tokenizer = model.tokenizer().clone(); diff --git a/crates/larql-inference/examples/streaming_demo.rs b/crates/larql-demos/examples/inference/streaming_demo.rs similarity index 100% rename from crates/larql-inference/examples/streaming_demo.rs rename to crates/larql-demos/examples/inference/streaming_demo.rs diff --git a/crates/larql-demos/examples/kv/README.md b/crates/larql-demos/examples/kv/README.md new file mode 100644 index 000000000..8e7762bbe --- /dev/null +++ b/crates/larql-demos/examples/kv/README.md @@ -0,0 +1,15 @@ +# KV engines demos + +Run any of these from a larql checkout: + +```sh +cargo run -p larql-demos --example engine_ladder +``` + +| demo | what it shows | run | +|---|---|---| +| `engine_ladder` | Every shipped engine end to end on synthetic weights. | weight-free · 0.2s | + +--- + +Demos that need a vindex take `--vindex PATH`; point them at any model you have under `output/`. They fail by name if the path is missing rather than surfacing a bare `NotFound`. diff --git a/crates/larql-kv/examples/engine_ladder.rs b/crates/larql-demos/examples/kv/engine_ladder.rs similarity index 98% rename from crates/larql-kv/examples/engine_ladder.rs rename to crates/larql-demos/examples/kv/engine_ladder.rs index aadf2665e..a8bc0f6e8 100644 --- a/crates/larql-kv/examples/engine_ladder.rs +++ b/crates/larql-demos/examples/kv/engine_ladder.rs @@ -14,7 +14,7 @@ //! Run with: //! //! ```sh -//! cargo run -p larql-kv --example engine_ladder +//! cargo run -p larql-demos --example engine_ladder //! ``` use larql_inference::cpu_engine_backend; diff --git a/crates/larql-demos/examples/lql/README.md b/crates/larql-demos/examples/lql/README.md new file mode 100644 index 000000000..16b3d4a8e --- /dev/null +++ b/crates/larql-demos/examples/lql/README.md @@ -0,0 +1,28 @@ +# LQL query layer demos + +Run any of these from a larql checkout: + +```sh +cargo run -p larql-demos --example compact_demo +``` + +| demo | what it shows | run | +|---|---|---| +| `compact_demo` | Storage-tier walkthrough for the LSM-style storage engine. | weight-free · 0.2s | +| `compile_demo` | End-to-end COMPILE. | needs a vindex · 88s | +| `lql_demo` | Parse, session, execute, error handling. | weight-free · 0.3s | +| `parser_demo` | Every statement type in spec v0.4, with its AST. | weight-free · 0.3s | +| `refine_demo` | End-to-end INSERT + COMPILE — Rust port of `experiments/14_vindex_compilation`. | needs a vindex · **541s** | +| `trace_demo` | Residual-stream decomposition. | needs a vindex · 124s | + +3 of these need no model weights and run in well under a second, so they are the quickest way to see the surface working. + +## `refine_demo` is slow on purpose + +It runs a real INSERT + COMPILE against Gemma 3 4B and takes about **nine +minutes**. That is the demo working, not hanging — it holds ~86% of one +core throughout. Budget accordingly before assuming it has stalled. + +--- + +Demos that need a vindex take `--vindex PATH`; point them at any model you have under `output/`. They fail by name if the path is missing rather than surfacing a bare `NotFound`. diff --git a/crates/larql-lql/examples/compact_demo.rs b/crates/larql-demos/examples/lql/compact_demo.rs similarity index 99% rename from crates/larql-lql/examples/compact_demo.rs rename to crates/larql-demos/examples/lql/compact_demo.rs index 72dc8c879..e30a98bdf 100644 --- a/crates/larql-lql/examples/compact_demo.rs +++ b/crates/larql-demos/examples/lql/compact_demo.rs @@ -18,7 +18,7 @@ //! surface using a synthetic browse-only vindex so it runs in CI //! with no model download. //! -//! Run: cargo run --release -p larql-lql --example compact_demo +//! Run: cargo run --release -p larql-demos --example compact_demo use larql_lql::{parse, Session}; use larql_vindex::ndarray::Array2; diff --git a/crates/larql-lql/examples/compile_demo.rs b/crates/larql-demos/examples/lql/compile_demo.rs similarity index 100% rename from crates/larql-lql/examples/compile_demo.rs rename to crates/larql-demos/examples/lql/compile_demo.rs diff --git a/crates/larql-lql/examples/lql_demo.rs b/crates/larql-demos/examples/lql/lql_demo.rs similarity index 99% rename from crates/larql-lql/examples/lql_demo.rs rename to crates/larql-demos/examples/lql/lql_demo.rs index 1086b25c3..d79755e1b 100644 --- a/crates/larql-lql/examples/lql_demo.rs +++ b/crates/larql-demos/examples/lql/lql_demo.rs @@ -3,7 +3,7 @@ //! Demonstrates the full LQL flow: parse statements, manage session state, //! execute against the (absent) backend, and handle errors gracefully. //! -//! Run: cargo run -p larql-lql --example lql_demo +//! Run: cargo run -p larql-demos --example lql_demo use larql_lql::{parse, run_batch, Session}; diff --git a/crates/larql-lql/examples/parser_demo.rs b/crates/larql-demos/examples/lql/parser_demo.rs similarity index 99% rename from crates/larql-lql/examples/parser_demo.rs rename to crates/larql-demos/examples/lql/parser_demo.rs index db8ace8c5..3ea8f320d 100644 --- a/crates/larql-lql/examples/parser_demo.rs +++ b/crates/larql-demos/examples/lql/parser_demo.rs @@ -1,6 +1,6 @@ //! LQL Parser Demo — parse every statement type from the spec v0.4 and display the AST //! -//! Run: cargo run -p larql-lql --example parser_demo +//! Run: cargo run -p larql-demos --example parser_demo use larql_lql::parse; diff --git a/crates/larql-lql/examples/refine_demo.rs b/crates/larql-demos/examples/lql/refine_demo.rs similarity index 100% rename from crates/larql-lql/examples/refine_demo.rs rename to crates/larql-demos/examples/lql/refine_demo.rs diff --git a/crates/larql-lql/examples/trace_demo.rs b/crates/larql-demos/examples/lql/trace_demo.rs similarity index 98% rename from crates/larql-lql/examples/trace_demo.rs rename to crates/larql-demos/examples/lql/trace_demo.rs index adf6771e6..e98de3e36 100644 --- a/crates/larql-lql/examples/trace_demo.rs +++ b/crates/larql-demos/examples/lql/trace_demo.rs @@ -17,7 +17,7 @@ //! Requires a vindex with model weights (`EXTRACT ... WITH ALL` or //! `EXTRACT ... WITH INFERENCE`). Skips cleanly when absent. //! -//! Run: cargo run --release -p larql-lql --example trace_demo +//! Run: cargo run --release -p larql-demos --example trace_demo use larql_lql::{parse, Session}; use std::path::Path; diff --git a/crates/larql-demos/examples/models/README.md b/crates/larql-demos/examples/models/README.md new file mode 100644 index 000000000..9ba83d8f2 --- /dev/null +++ b/crates/larql-demos/examples/models/README.md @@ -0,0 +1,19 @@ +# Architecture detection demos + +Run any of these from a larql checkout: + +```sh +cargo run -p larql-demos --example architecture_demo +``` + +| demo | what it shows | run | +|---|---|---| +| `architecture_demo` | Detection and configuration for all 12 supported architectures. | weight-free · 0.2s | +| `demo_loading` | Loading from a directory or a GGUF file. | weight-free · 0.1s | +| `demo_tensor_keys` | Tensor-key patterns compared across architectures. | weight-free · 0.2s | + +None of these need model weights — they run in well under a second and are the quickest way to see the surface working. + +--- + +Demos that need a vindex take `--vindex PATH`; point them at any model you have under `output/`. They fail by name if the path is missing rather than surfacing a bare `NotFound`. diff --git a/crates/larql-models/examples/architecture_demo.rs b/crates/larql-demos/examples/models/architecture_demo.rs similarity index 99% rename from crates/larql-models/examples/architecture_demo.rs rename to crates/larql-demos/examples/models/architecture_demo.rs index 09984f17e..49cb11a4b 100644 --- a/crates/larql-models/examples/architecture_demo.rs +++ b/crates/larql-demos/examples/models/architecture_demo.rs @@ -4,7 +4,7 @@ //! sliding window patterns, MoE routing, MLA compression, bias keys, scaling multipliers, //! softcapping, per-layer geometry, PLE, KV sharing, and RoPE scaling. //! -//! Run: cargo run -p larql-models --example architecture_demo +//! Run: cargo run -p larql-demos --example architecture_demo use larql_models::{detect_from_json, ModelArchitecture}; diff --git a/crates/larql-models/examples/demo_loading.rs b/crates/larql-demos/examples/models/demo_loading.rs similarity index 96% rename from crates/larql-models/examples/demo_loading.rs rename to crates/larql-demos/examples/models/demo_loading.rs index 9a0f8247d..9e87505ba 100644 --- a/crates/larql-models/examples/demo_loading.rs +++ b/crates/larql-demos/examples/models/demo_loading.rs @@ -3,7 +3,7 @@ //! Shows how larql-models loads weights, detects architecture, and exposes //! tensor information. Requires a model path as argument. //! -//! Run: cargo run -p larql-models --example demo_loading -- /path/to/model +//! Run: cargo run -p larql-demos --example demo_loading -- /path/to/model use std::path::PathBuf; @@ -16,8 +16,8 @@ fn main() { println!(" also accepts HF model IDs (google/gemma-3-4b)"); println!(); println!("Examples:"); - println!(" cargo run -p larql-models --example demo_loading -- /path/to/gemma-3-4b"); - println!(" cargo run -p larql-models --example demo_loading -- model.gguf"); + println!(" cargo run -p larql-demos --example demo_loading -- /path/to/gemma-3-4b"); + println!(" cargo run -p larql-demos --example demo_loading -- model.gguf"); return; } diff --git a/crates/larql-models/examples/demo_tensor_keys.rs b/crates/larql-demos/examples/models/demo_tensor_keys.rs similarity index 99% rename from crates/larql-models/examples/demo_tensor_keys.rs rename to crates/larql-demos/examples/models/demo_tensor_keys.rs index b2b86efad..f678ecce8 100644 --- a/crates/larql-models/examples/demo_tensor_keys.rs +++ b/crates/larql-demos/examples/models/demo_tensor_keys.rs @@ -3,7 +3,7 @@ //! Shows how each model family maps layer indices to tensor keys, //! highlighting differences in prefix, projection names, and norm patterns. //! -//! Run: cargo run -p larql-models --example demo_tensor_keys +//! Run: cargo run -p larql-demos --example demo_tensor_keys use larql_models::{detect_from_json, ModelArchitecture}; diff --git a/crates/larql-demos/examples/server/README.md b/crates/larql-demos/examples/server/README.md new file mode 100644 index 000000000..670c3590a --- /dev/null +++ b/crates/larql-demos/examples/server/README.md @@ -0,0 +1,20 @@ +# Serving surface demos + +Run any of these from a larql checkout: + +```sh +cargo run -p larql-demos --example embed_demo +``` + +| demo | what it shows | run | +|---|---|---| +| `embed_demo` | What the embed endpoints return, with synthetic data. | weight-free · 0.2s | +| `openai_demo` | Boots an in-process server and exercises `/v1/models`, `/v1/embeddings`, `/v1/completions`, `/v1/chat/completions`. | takes a vindex path | +| `server_demo` | Builds a synthetic vindex and shows what the server would return. | weight-free · 0.2s | +| `shard_query_demo` | Exp 53 `ShardService`, end to end. | weight-free · 0.3s | + +3 of these need no model weights and run in well under a second, so they are the quickest way to see the surface working. + +--- + +Demos that need a vindex take `--vindex PATH`; point them at any model you have under `output/`. They fail by name if the path is missing rather than surfacing a bare `NotFound`. diff --git a/crates/larql-server/examples/embed_demo.rs b/crates/larql-demos/examples/server/embed_demo.rs similarity index 99% rename from crates/larql-server/examples/embed_demo.rs rename to crates/larql-demos/examples/server/embed_demo.rs index 705bc4da5..983c575de 100644 --- a/crates/larql-server/examples/embed_demo.rs +++ b/crates/larql-demos/examples/server/embed_demo.rs @@ -6,7 +6,7 @@ //! 3. `GET /v1/token/*` — tokenizer encode / decode //! //! No real model needed. Run: -//! cargo run -p larql-server --example embed_demo +//! cargo run -p larql-demos --example embed_demo use larql_vindex::ndarray::Array2; @@ -226,7 +226,7 @@ fn demo_binary_wire() { println!( "Logits request ({} bytes): {:?}", logits_req.len(), - &residual + residual ); } diff --git a/crates/larql-server/examples/openai_demo.rs b/crates/larql-demos/examples/server/openai_demo.rs similarity index 98% rename from crates/larql-server/examples/openai_demo.rs rename to crates/larql-demos/examples/server/openai_demo.rs index f832a2670..c23256770 100644 --- a/crates/larql-server/examples/openai_demo.rs +++ b/crates/larql-demos/examples/server/openai_demo.rs @@ -3,7 +3,7 @@ //! `/v1/chat/completions` end-to-end against the loaded vindex. //! //! Usage: -//! cargo run -p larql-server --example openai_demo -- +//! cargo run -p larql-demos --example openai_demo -- //! //! ## Vindex compatibility //! @@ -13,12 +13,12 @@ //! //! ```bash //! # f16 (fastest, KV-cached): -//! cargo run --release -p larql-server --example openai_demo -- \ +//! cargo run --release -p larql-demos --example openai_demo -- \ //! output/gemma3-4b-fresh.vindex //! //! # Q4_K (correct output; CPU per-step Q4_K decode is O(N²) so //! # high `max_tokens` runs are slow): -//! cargo run --release -p larql-server --example openai_demo -- \ +//! cargo run --release -p larql-demos --example openai_demo -- \ //! output/gemma3-4b-q4k-streaming.vindex //! ``` //! @@ -569,7 +569,7 @@ fn main() -> Result<(), Box> { HTTP client) and exercises the OpenAI-compat endpoints end-to-\n\ end against the loaded vindex.\n\n\ Examples:\n\ - cargo run --release -p larql-server --example openai_demo -- \\\n\ + cargo run --release -p larql-demos --example openai_demo -- \\\n\ output/gemma3-4b-q4k-streaming.vindex" ); std::process::exit(1); diff --git a/crates/larql-server/examples/server_demo.rs b/crates/larql-demos/examples/server/server_demo.rs similarity index 99% rename from crates/larql-server/examples/server_demo.rs rename to crates/larql-demos/examples/server/server_demo.rs index da031ce9f..ea62c4ca5 100644 --- a/crates/larql-server/examples/server_demo.rs +++ b/crates/larql-demos/examples/server/server_demo.rs @@ -1,6 +1,6 @@ //! Server demo — builds a synthetic vindex and shows what the server would return. //! -//! Run: cargo run -p larql-server --example server_demo +//! Run: cargo run -p larql-demos --example server_demo use larql_vindex::ndarray::Array2; use larql_vindex::{FeatureMeta, PatchOp, PatchedVindex, VectorIndex, VindexPatch}; diff --git a/crates/larql-server/examples/shard_query_demo.rs b/crates/larql-demos/examples/server/shard_query_demo.rs similarity index 99% rename from crates/larql-server/examples/shard_query_demo.rs rename to crates/larql-demos/examples/server/shard_query_demo.rs index 89bd60ed2..310f97173 100644 --- a/crates/larql-server/examples/shard_query_demo.rs +++ b/crates/larql-demos/examples/server/shard_query_demo.rs @@ -15,7 +15,7 @@ //! same proto, same `ShardSource`, same `Arc`-shared `PatchedVindex`. //! //! ```bash -//! cargo run --release -p larql-server --example shard_query_demo +//! cargo run --release -p larql-demos --example shard_query_demo //! ``` use std::sync::Arc; diff --git a/crates/larql-demos/examples/vindex/README.md b/crates/larql-demos/examples/vindex/README.md new file mode 100644 index 000000000..8fb6e0ee2 --- /dev/null +++ b/crates/larql-demos/examples/vindex/README.md @@ -0,0 +1,21 @@ +# Vindex format and store demos + +Run any of these from a larql checkout: + +```sh +cargo run -p larql-demos --example demo_features +``` + +| demo | what it shows | run | +|---|---|---| +| `demo_features` | Showcase of the complete larql-vindex API. | weight-free · 0.2s | +| `demo_memit_solve` | `memit_solve` + `MemitStore` — the COMPACT MAJOR pipeline in miniature. | weight-free · 0.1s | +| `mmap_demo` | Vindex mmap memory behaviour and model-scaling projections. | weight-free · 0.7s | +| `q4k_demo` | Streaming Q4_K extract. | weight-free · 0.3s | +| `walker_demo` | All three build-time graph extractors against a tiny mock model. | weight-free · 0.3s | + +None of these need model weights — they run in well under a second and are the quickest way to see the surface working. + +--- + +Demos that need a vindex take `--vindex PATH`; point them at any model you have under `output/`. They fail by name if the path is missing rather than surfacing a bare `NotFound`. diff --git a/crates/larql-vindex/examples/demo_features.rs b/crates/larql-demos/examples/vindex/demo_features.rs similarity index 99% rename from crates/larql-vindex/examples/demo_features.rs rename to crates/larql-demos/examples/vindex/demo_features.rs index 339718ede..060c312ac 100644 --- a/crates/larql-vindex/examples/demo_features.rs +++ b/crates/larql-demos/examples/vindex/demo_features.rs @@ -5,7 +5,7 @@ //! patches (create, apply, revert, bake down), extract pipeline, GGUF key normalization, //! Vindexfile parsing, HuggingFace path handling, and quantization formats. //! -//! Run: cargo run -p larql-vindex --example vindex_demo +//! Run: cargo run -p larql-demos --example demo_features use larql_models::TopKEntry; use larql_vindex::{FeatureMeta, VectorIndex, VindexConfig}; diff --git a/crates/larql-vindex/examples/demo_memit_solve.rs b/crates/larql-demos/examples/vindex/demo_memit_solve.rs similarity index 97% rename from crates/larql-vindex/examples/demo_memit_solve.rs rename to crates/larql-demos/examples/vindex/demo_memit_solve.rs index bd211bd02..03a2d6a4e 100644 --- a/crates/larql-vindex/examples/demo_memit_solve.rs +++ b/crates/larql-demos/examples/vindex/demo_memit_solve.rs @@ -6,7 +6,7 @@ //! adds a cycle to a fresh `MemitStore`. Concludes with an //! entity/relation lookup against the store. //! -//! Run: cargo run --release -p larql-vindex --example demo_memit_solve +//! Run: cargo run --release -p larql-demos --example demo_memit_solve use larql_vindex::{memit_solve, MemitFact, MemitStore}; use ndarray::Array2; diff --git a/crates/larql-vindex/examples/mmap_demo.rs b/crates/larql-demos/examples/vindex/mmap_demo.rs similarity index 99% rename from crates/larql-vindex/examples/mmap_demo.rs rename to crates/larql-demos/examples/vindex/mmap_demo.rs index 867b39fc6..563c9c2fd 100644 --- a/crates/larql-vindex/examples/mmap_demo.rs +++ b/crates/larql-demos/examples/vindex/mmap_demo.rs @@ -15,7 +15,7 @@ //! table for Gemma 4B → Kimi-K2 (1T params), showing what each //! model would need for full inference vs vindex inference. //! -//! Run: `cargo run --release -p larql-vindex --example mmap_demo` +//! Run: `cargo run --release -p larql-demos --example mmap_demo` use larql_models::TopKEntry; use larql_vindex::{FeatureMeta, VectorIndex, VindexConfig}; diff --git a/crates/larql-vindex/examples/q4k_demo.rs b/crates/larql-demos/examples/vindex/q4k_demo.rs similarity index 99% rename from crates/larql-vindex/examples/q4k_demo.rs rename to crates/larql-demos/examples/vindex/q4k_demo.rs index 0cd487fcd..ca6bcb449 100644 --- a/crates/larql-vindex/examples/q4k_demo.rs +++ b/crates/larql-demos/examples/vindex/q4k_demo.rs @@ -13,7 +13,7 @@ //! //! This is a pure-synthetic demo — no model download, runs in CI. //! -//! Run: cargo run --release -p larql-vindex --example q4k_demo +//! Run: cargo run --release -p larql-demos --example q4k_demo use std::collections::HashMap; use std::path::{Path, PathBuf}; diff --git a/crates/larql-vindex/examples/walker_demo.rs b/crates/larql-demos/examples/vindex/walker_demo.rs similarity index 97% rename from crates/larql-vindex/examples/walker_demo.rs rename to crates/larql-demos/examples/vindex/walker_demo.rs index 0877206be..8ffdf5757 100644 --- a/crates/larql-vindex/examples/walker_demo.rs +++ b/crates/larql-demos/examples/vindex/walker_demo.rs @@ -1,7 +1,7 @@ //! Walker demo — runs all three build-time graph extractors against a //! tiny mock model and prints a summary. //! -//! Run: `cargo run --release -p larql-vindex --example walker_demo` +//! Run: `cargo run --release -p larql-demos --example walker_demo` //! //! Useful as smoke / "hello world" for the walker API. For real-model //! extraction use the larql CLI: diff --git a/crates/larql-demos/src/lib.rs b/crates/larql-demos/src/lib.rs new file mode 100644 index 000000000..92397c25b --- /dev/null +++ b/crates/larql-demos/src/lib.rs @@ -0,0 +1,13 @@ +//! Runnable demonstrations of larql's shipped capabilities. +//! +//! This crate carries no logic. Every demo is an `examples/*.rs` binary, +//! filed under the crate whose capability it demonstrates. The library +//! target exists only because cargo requires a package to have one. +//! +//! Not here, deliberately: +//! +//! - **Benchmarks, diagnostics and parity harnesses** stay in their own +//! crate's `examples/`. They exercise the engine rather than showing +//! how to use it, and they belong next to the code they measure. +//! - **Research probes** live in `chris-experiments/larql_probes`, +//! pinned to the larql revision that produced their recorded verdict. diff --git a/crates/larql-execution/Cargo.toml b/crates/larql-execution/Cargo.toml new file mode 100644 index 000000000..ceac1a569 --- /dev/null +++ b/crates/larql-execution/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "larql-execution" +version.workspace = true +edition.workspace = true +authors.workspace = true +license.workspace = true +description = "Execution-refusal semantics shared across LARQL's runtime crates. Zero larql-* deps." +keywords = ["execution", "refusal", "runtime"] +categories = ["data-structures"] + +# Deliberately empty, and it should stay that way. This is a contract crate: +# `larql-compute` sits below `larql-vindex` in the graph, so neither could name +# the other's refusal vocabulary without a cycle, and a shared leaf is the only +# place the classification can live. Follows `larql-vindex-spec`'s precedent — +# a dependency-light contract several crates depend on. +[dependencies] diff --git a/crates/larql-execution/src/lib.rs b/crates/larql-execution/src/lib.rs new file mode 100644 index 000000000..1478fafad --- /dev/null +++ b/crates/larql-execution/src/lib.rs @@ -0,0 +1,200 @@ +//! Why an execution refused, classified by the response it requires. +//! +//! This crate exists because a dependency cycle pointed at a missing layer. +//! `FfnBackend` lives in `larql-compute`; the refusal vocabulary grew up in +//! `larql-vindex`, which depends on `larql-compute`. Neither could name the +//! other's type, so the classification could not cross the boundary that most +//! needed it — the one between a route that refused and an engine deciding what +//! to do about it. +//! +//! # The axis is the response, not the cause +//! +//! [`RefusalKind`] has three variants because there are three distinct +//! responses, and each has a different person or system on the other end: +//! +//! ```text +//! Residency same operation, operand becomes available → may succeed +//! Unsupported same operands, a different capable executor → may succeed +//! BindingDefect the binding or artifact itself must change +//! ``` +//! +//! That partition is what makes the enum actionable. Finer reasons — an +//! unsupported activation, a decomposed projection where a fused one was +//! needed, a misaligned base — are *causes of* `Unsupported`, not separate +//! responses, and they belong in the concrete error's payload where they can +//! carry their own detail. Promoting them here would put two abstraction levels +//! in one enum and leave a caller unable to switch on the thing it must act on. +//! +//! # `BindingDefect` is load-bearing +//! +//! It is the only kind that means *reject the artifact*. An expert id outside +//! the router's population, a region shorter than its declared shape, an +//! internally inconsistent plan — none of those is fixed by fetching an operand +//! or by choosing another kernel, and classifying them as either would send +//! someone to repair the wrong thing. +//! +//! The invariant to hold when adding a variant or classifying a new error: +//! +//! > `BindingDefect` means that repeating the operation with more residency, or +//! > through any other capable executor, cannot make *this* bound plan valid. +//! +//! # Transport failures are not a fourth kind +//! +//! A network timeout, a refused connection, an expired lease — those are +//! execution-*attempt* failures, often retryable, and they belong to whatever +//! transport owns them. One may eventually *conclude* in `Residency`, when the +//! semantic finding is that the operand is unavailable to this plan. Admitting +//! them directly would turn this enum from execution semantics into a catalogue +//! of everything that can go wrong operationally. + +use std::error::Error; +use std::fmt; + +/// The response a refusal requires. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub enum RefusalKind { + /// A valid operand exists, and is not available here. + /// + /// **Not a defect.** A shard that does not hold an expert is behaving as + /// designed; the response is to fetch, load or reroute. Counting these as + /// failures makes a sweep's headline number a measurement of slice coverage + /// rather than of execution. + Residency, + /// A valid operation this route cannot execute. + /// + /// The representation is well-formed and the operands are present; no bound + /// kernel serves them. The response is another executor, or another stored + /// variant of the same component — never a quiet reinterpretation of the + /// bytes. + Unsupported, + /// The bound artifact or plan violates its own contract. + /// + /// The only kind that means reject the index and repair whatever produced + /// it. + BindingDefect, +} + +impl RefusalKind { + /// Every kind, so a consumer sweeping them cannot silently narrow. + pub const ALL: [Self; 3] = [Self::Residency, Self::Unsupported, Self::BindingDefect]; + + /// Snake case, matching the sibling outcome names it is printed beside + /// (`ok`, `declined`) — one table, one convention. + pub const fn name(self) -> &'static str { + match self { + Self::Residency => "residency", + Self::Unsupported => "unsupported", + Self::BindingDefect => "binding_defect", + } + } + + /// Whether this indicts the artifact rather than the environment. + /// + /// The question an operator actually asks: is something broken, or is + /// something merely elsewhere? + pub const fn indicts_the_artifact(self) -> bool { + matches!(self, Self::BindingDefect) + } + + /// Whether the same bound plan could succeed given a different environment + /// — more residency, or another capable executor. + pub const fn is_recoverable_without_rebinding(self) -> bool { + matches!(self, Self::Residency | Self::Unsupported) + } +} + +impl fmt::Display for RefusalKind { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(self.name()) + } +} + +/// An error that carries its response category across a crate boundary. +/// +/// Concrete errors stay in the crates that own them, with all their detail — +/// which expert, which bank, which axis. This trait carries only the one thing +/// a caller two layers up must switch on. +pub trait ExecutionRefusal: Error + Send + Sync + 'static { + fn kind(&self) -> RefusalKind; +} + +/// A refusal crossing a boundary that cannot name its concrete type. +pub type BoxRefusal = Box; + +#[cfg(test)] +mod tests { + use super::*; + + #[derive(Debug)] + struct Refused(RefusalKind); + + impl fmt::Display for Refused { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "refused: {}", self.0) + } + } + impl Error for Refused {} + impl ExecutionRefusal for Refused { + fn kind(&self) -> RefusalKind { + self.0 + } + } + + #[test] + fn every_kind_has_a_distinct_name() { + // The names reach operator-facing diagnostics, so two kinds sharing one + // would make a report unactionable exactly where it matters. + let mut names: Vec<&str> = RefusalKind::ALL.iter().map(|k| k.name()).collect(); + let count = names.len(); + names.sort_unstable(); + names.dedup(); + assert_eq!(names.len(), count); + } + + #[test] + fn only_a_binding_defect_indicts_the_artifact() { + // The partition that decides who is sent to fix something. + assert!(RefusalKind::BindingDefect.indicts_the_artifact()); + assert!(!RefusalKind::Residency.indicts_the_artifact()); + assert!(!RefusalKind::Unsupported.indicts_the_artifact()); + } + + #[test] + fn residency_and_unsupported_can_succeed_without_rebinding() { + // The stated invariant: more residency or another executor may rescue + // those two, and cannot rescue a binding defect. + assert!(RefusalKind::Residency.is_recoverable_without_rebinding()); + assert!(RefusalKind::Unsupported.is_recoverable_without_rebinding()); + assert!(!RefusalKind::BindingDefect.is_recoverable_without_rebinding()); + } + + #[test] + fn the_two_predicates_partition_the_vocabulary() { + // Guards against a future variant that is neither or both, which would + // leave a caller with no branch to take. + for kind in RefusalKind::ALL { + assert_ne!( + kind.indicts_the_artifact(), + kind.is_recoverable_without_rebinding(), + "{kind} is not on exactly one side of the partition" + ); + } + } + + #[test] + fn a_kind_displays_as_its_name() { + assert_eq!(RefusalKind::Residency.to_string(), "residency"); + assert_eq!(RefusalKind::BindingDefect.to_string(), "binding_defect"); + } + + #[test] + fn a_boxed_refusal_keeps_its_kind_and_its_message() { + // The point of the trait: the classification survives a boundary that + // cannot name the concrete type, and the detail survives with it. + let boxed: BoxRefusal = Box::new(Refused(RefusalKind::Residency)); + assert_eq!(boxed.kind(), RefusalKind::Residency); + assert!(boxed.to_string().contains("residency")); + let source: &dyn Error = boxed.as_ref(); + assert!(source.source().is_none()); + } +} diff --git a/crates/larql-inference/Cargo.toml b/crates/larql-inference/Cargo.toml index 44b377e5b..199f7b53d 100644 --- a/crates/larql-inference/Cargo.toml +++ b/crates/larql-inference/Cargo.toml @@ -22,6 +22,7 @@ larql-compute = { path = "../larql-compute" } larql-compute-metal = { path = "../larql-compute-metal", optional = true } larql-core = { path = "../larql-core" } larql-vindex = { path = "../larql-vindex" } +larql-execution = { path = "../larql-execution" } serde = { workspace = true } serde_json = { workspace = true } thiserror = { workspace = true } @@ -108,10 +109,6 @@ required-features = ["gpu"] name = "decode_vs_prefill" required-features = ["gpu"] -[[example]] -name = "residual_diff" -required-features = ["gpu"] - [[example]] name = "stage_bisect" required-features = ["gpu"] diff --git a/crates/larql-inference/README.md b/crates/larql-inference/README.md index 0798ec473..d20d7f328 100644 --- a/crates/larql-inference/README.md +++ b/crates/larql-inference/README.md @@ -332,7 +332,7 @@ crate's `docs/specs/` directory. Read in this order: | [`kv-engine-unification.md`](docs/specs/kv-engine-unification.md) | The `KvEngine` trait surface. §4.4 documents W10's `StateDumpMask` + `read_kv_row_at` widening. | | [`zone-engine.md`](docs/specs/zone-engine.md) | **Top-level composer.** Sequences PREDICT / WALK / CACHE zones between choke points. | | [`layer-engine.md`](docs/specs/layer-engine.md) v0.4 | Inner per-layer composer for WALK zones (subsumed under ZoneEngine). | -| [`markov-residual-engine.md`](docs/specs/markov-residual-engine.md), [`markov-residual-codec-engine.md`](docs/specs/markov-residual-codec-engine.md), [`unlimited-context-engine.md`](docs/specs/unlimited-context-engine.md), [`standard-engine.md`](docs/specs/standard-engine.md), [`turbo-quant-engine.md`](docs/specs/turbo-quant-engine.md), [`apollo-engine.md`](docs/specs/apollo-engine.md), [`no-cache-engine.md`](docs/specs/no-cache-engine.md), [`boundary-kv-engine.md`](docs/specs/boundary-kv-engine.md), [`boundary-per-layer-engine.md`](docs/specs/boundary-per-layer-engine.md) | Per-engine contracts; each marks its W10 opt-in path where applicable. | +| [`markov-residual-engine.md`](docs/specs/markov-residual-engine.md), [`markov-residual-codec-engine.md`](docs/specs/markov-residual-codec-engine.md), [`windowed-checkpoint-engine.md`](docs/specs/windowed-checkpoint-engine.md), [`standard-engine.md`](docs/specs/standard-engine.md), [`turbo-quant-engine.md`](docs/specs/turbo-quant-engine.md), [`apollo-engine.md`](docs/specs/apollo-engine.md), [`no-cache-engine.md`](docs/specs/no-cache-engine.md), [`boundary-kv-engine.md`](docs/specs/boundary-kv-engine.md), [`boundary-per-layer-engine.md`](docs/specs/boundary-per-layer-engine.md) | Per-engine contracts; each marks its W10 opt-in path where applicable. | ## License diff --git a/crates/larql-inference/ROADMAP.md b/crates/larql-inference/ROADMAP.md index 46f720684..632bd272c 100644 --- a/crates/larql-inference/ROADMAP.md +++ b/crates/larql-inference/ROADMAP.md @@ -1014,7 +1014,7 @@ them. Tests: `forward::kv_generate::tests` — noop matches baseline; record fires on prefill + every decode step; α=5 steer changes generated tokens vs -baseline. Demo: `examples/mech_interp_demo.rs` § [7] shows +baseline. Demo: `crates/larql-demos/examples/inference/mech_interp_demo.rs` § [7] shows `baseline_ids = [12, 30, 10, 29]` vs `steered_ids = [4, 4, 4, 4]`. ### M7 — `W_E` / `W_U` + `project_through_unembed` @@ -1064,7 +1064,7 @@ Wired into `generate_with_sampling` via `eos.is_eos(id, &decoded)`. Greedy suffix on each `push`. Equivalent to llama.cpp `llama_token_to_piece` and HF Python `decode_stream`. Handles HF leading-space (`▁`) for SP tokenizers and multi-byte UTF-8 chars that straddle a token boundary. Demo at -`examples/detok_demo.rs` shows the bug ("thecapitaloffranceisparis") and the +`crates/larql-demos/examples/inference/detok_demo.rs` shows the bug ("thecapitaloffranceisparis") and the fix ("the capital of france is paris"). ### Token streaming @@ -1073,7 +1073,7 @@ fix ("the capital of france is paris"). every emitted token, including the first (which comes out of prefill). Uses `Detokenizer::push` so streamed text preserves HF leading-space spacing. `generate_with_sampling` is a thin wrapper passing a no-op closure so -non-streaming callers are unaffected. Demo at `examples/streaming_demo.rs` +non-streaming callers are unaffected. Demo at `crates/larql-demos/examples/inference/streaming_demo.rs` prints tokens live with stdout flushing. ### Sampling @@ -1091,7 +1091,7 @@ overhead is <2µs/call at top-K=64 (<0.02% of decode budget). CLI flags `max_context`. Pluggable `TurnRenderer` covers Gemma / ChatML / Llama-3 templates. The most recent turn is never dropped — eviction is a no-op when only one turn remains, so a long single prompt is preserved over -silently truncating. `examples/chat_demo.rs` runs a 3-turn conversation. +silently truncating. `crates/larql-demos/examples/inference/chat_demo.rs` runs a 3-turn conversation. True KV carryover across turns (so prefill on turn N+1 only processes the new tokens) is a follow-up — the API surface is in place; it's an diff --git a/crates/larql-inference/coverage-policy.json b/crates/larql-inference/coverage-policy.json index b429637cd..732310224 100644 --- a/crates/larql-inference/coverage-policy.json +++ b/crates/larql-inference/coverage-policy.json @@ -1,5 +1,5 @@ { - "policy_note": "Per-file LINE-coverage gates for the dispatch trait families (cargo-llvm-cov: `Lines Cover` column, not the `Regions` column). Established 2026-05-16 alongside the AsyncComputeBackend + KvDispatch directory reorganization. Every file in scope clears the 90% default. Scope is narrow — the trait families only — because larql-inference's overall baseline (~70%) reflects an older surface that's tracked separately. `included_total_line_min_percent` checks the sum across this scope. Trait files (mod.rs) cover the `unimplemented!()` defaults via a stub-backend test pattern that satisfies the supertraits with minimal surface (one panicking `MatMul` impl + everything else's defaults).", + "policy_note": "Per-file LINE-coverage gates for the dispatch trait families (cargo-llvm-cov: `Lines Cover` column, not the `Regions` column). Established 2026-05-16 alongside the AsyncComputeBackend + KvDispatch directory reorganization. Every file in scope clears the 90% default. Scope is narrow \u2014 the trait families only \u2014 because larql-inference's overall baseline (~70%) reflects an older surface that's tracked separately. `included_total_line_min_percent` checks the sum across this scope. Trait files (mod.rs) cover the `unimplemented!()` defaults via a stub-backend test pattern that satisfies the supertraits with minimal surface (one panicking `MatMul` impl + everything else's defaults).", "include_globs": [ "crates/larql-inference/src/async_compute_backend/mod.rs", "crates/larql-inference/src/async_compute_backend/cpu.rs", @@ -7,7 +7,9 @@ "crates/larql-inference/src/kv_dispatch/mod.rs", "crates/larql-inference/src/kv_dispatch/cpu.rs", "crates/larql-inference/src/kv_dispatch/metal.rs", - "crates/larql-inference/src/kv_dispatch/helpers.rs" + "crates/larql-inference/src/kv_dispatch/helpers.rs", + "crates/larql-inference/src/ffn/moe_backend.rs", + "crates/larql-inference/src/ffn/moe_bound.rs" ], "exclude_globs": [], "default_line_min_percent": 90.0, diff --git a/crates/larql-inference/docs/ffn-build-router.md b/crates/larql-inference/docs/ffn-build-router.md index 4416dd936..461685861 100644 --- a/crates/larql-inference/docs/ffn-build-router.md +++ b/crates/larql-inference/docs/ffn-build-router.md @@ -236,7 +236,7 @@ impl<'a> LayerFfnRouter<'a> { ``` Existing callers (`larql-cli/src/commands/extraction/walk_cmd.rs:868` -and `larql-inference/examples/walk_boundary_sweep.rs:205`) own the +and `larql-inference/chris-experiments/larql_probes/examples/misc/walk_boundary_sweep.rs:205`) own the concrete backend instances as local values and pass references: ```rust @@ -469,7 +469,7 @@ design and the implementation in the same pass. - `larql-inference/src/ffn/mod.rs:46` — existing `LayerFfnRouter` shape that §3 is reasoning against. - `larql-cli/src/commands/extraction/walk_cmd.rs:868` and - `larql-inference/examples/walk_boundary_sweep.rs:205` — existing + `larql-inference/chris-experiments/larql_probes/examples/misc/walk_boundary_sweep.rs:205` — existing call sites that exemplify the "caller-owns-backends" pattern `build_router` is collapsing. - [`docs/state-policy.md`](./state-policy.md) — the diff --git a/crates/larql-inference/docs/specs/kv-engine-unification.md b/crates/larql-inference/docs/specs/kv-engine-unification.md index 01a956eb7..c5a751285 100644 --- a/crates/larql-inference/docs/specs/kv-engine-unification.md +++ b/crates/larql-inference/docs/specs/kv-engine-unification.md @@ -479,7 +479,7 @@ byte-for-byte. `larql-inference` (calls `generate_with_engine`), so `larql-inference` core code never names `larql-kv`. The only remaining consumer in `larql-inference/` is - `examples/apollo_rd_backend.rs`, which justifies the dev-dep. + `crates/larql-demos/examples/inference/apollo_rd_backend.rs`, which justifies the dev-dep. ### 8.8 Rollback diff --git a/crates/larql-inference/docs/specs/layer-engine.md b/crates/larql-inference/docs/specs/layer-engine.md index 86c17d05c..2e2e76978 100644 --- a/crates/larql-inference/docs/specs/layer-engine.md +++ b/crates/larql-inference/docs/specs/layer-engine.md @@ -578,7 +578,7 @@ its scope as "the top-level engine" was wrong. per-layer. - [`markov-residual-engine.md` §14](./markov-residual-engine.md), [`markov-residual-codec-engine.md` §14](./markov-residual-codec-engine.md), - [`unlimited-context-engine.md` §8](./unlimited-context-engine.md) — + [`windowed-checkpoint-engine.md` §8](./windowed-checkpoint-engine.md) — per-engine W10 opt-in tables; the mask cascade these refer to is the same mechanism LayerEngine reuses per layer. - `crates/larql-kv/examples/contract_classify_cached_ffn.rs` — the diff --git a/crates/larql-inference/docs/specs/virtual-experts/arithmetic-virtual-expert.md b/crates/larql-inference/docs/specs/virtual-experts/arithmetic-virtual-expert.md index acac546da..a03c85650 100644 --- a/crates/larql-inference/docs/specs/virtual-experts/arithmetic-virtual-expert.md +++ b/crates/larql-inference/docs/specs/virtual-experts/arithmetic-virtual-expert.md @@ -72,7 +72,7 @@ measured):** tier-0 fires on *notation*, never on inferred intent — strong gly trailing `=`. Everything else is the designed fallthrough: deciding whether "9 - 5" is arithmetic is an engagement question and belongs to the model (tier-1 exhaust, or an FR3-style explicit classify), not to surface heuristics. Adversarial prose corpus: -0 false fires (`examples/scanner_adversarial.rs`). +0 false fires (`chris-experiments/larql_probes/examples/misc/scanner_adversarial.rs`). **Tier 1 — engagement probe (disguised math). DEMOTED (A11).** Ridge probe on the L8 residual at the last prompt token, reading arithmetic-engagement exhaust (math vs diff --git a/crates/larql-inference/docs/specs/unlimited-context-engine.md b/crates/larql-inference/docs/specs/windowed-checkpoint-engine.md similarity index 90% rename from crates/larql-inference/docs/specs/unlimited-context-engine.md rename to crates/larql-inference/docs/specs/windowed-checkpoint-engine.md index bad0fd0c3..36bd351da 100644 --- a/crates/larql-inference/docs/specs/unlimited-context-engine.md +++ b/crates/larql-inference/docs/specs/windowed-checkpoint-engine.md @@ -1,4 +1,4 @@ -# UnlimitedContextEngine — Specification +# WindowedCheckpointEngine — Specification **Status:** ✅ Shipped. W1-GPU step 4 wired + bench-validated 2026-05-17: 28 → 56.0 tok/s on Metal (window=256, Gemma 3 4B, @@ -22,7 +22,7 @@ M3 Max, 50-token decode). W10 HOnly default-on (2026-05-21): ## 1. Purpose -`UnlimitedContextEngine` provides effectively-unlimited decoding +`WindowedCheckpointEngine` provides effectively-unlimited decoding context with bounded current memory — by checkpointing the K/V state at fixed window boundaries (`window_size` tokens) and archiving the prompt token IDs, then reconstructing any prior @@ -35,7 +35,7 @@ of checkpoints + token archive and accept the cost of re-prefill when accessing earlier windows. The engine is **not** a sliding-window cache. Sliding window drops -old tokens; `unlimited_context` keeps them in the cold tier and can +old tokens; `windowed_checkpoint` keeps them in the cold tier and can replay them on demand. --- @@ -137,11 +137,11 @@ their window. | Concern | Location | |---|---| -| Engine struct + `KvEngine` impl | `crates/larql-kv/src/engines/unlimited_context/engine.rs` | -| Checkpoint storage | `engines/unlimited_context/checkpoint_store.rs` | -| Token archive | `engines/unlimited_context/token_archive.rs` | -| Per-token K/V extension | `engines/unlimited_context/extend.rs::rs_extend_from_checkpoint_*` | -| W1-GPU dispatch helpers | `engines/unlimited_context/engine.rs::try_prefill_via_dispatch` + `decode_step_via_dispatch` | +| Engine struct + `KvEngine` impl | `crates/larql-kv/src/engines/windowed_checkpoint/engine.rs` | +| Checkpoint storage | `engines/windowed_checkpoint/checkpoint_store.rs` | +| Token archive | `engines/windowed_checkpoint/token_archive.rs` | +| Per-token K/V extension | `engines/windowed_checkpoint/extend.rs::rs_extend_from_checkpoint_*` | +| W1-GPU dispatch helpers | `engines/windowed_checkpoint/engine.rs::try_prefill_via_dispatch` + `decode_step_via_dispatch` | --- @@ -151,10 +151,10 @@ their window. reconstructs the *whole window* from the prior checkpoint; there's no API to extract K/V at a single sub-window position. - **Compression of the cold tier.** That's `markov_residual_codec` - / `boundary_per_layer`. `unlimited_context` keeps cold + / `boundary_per_layer`. `windowed_checkpoint` keeps cold checkpoints as raw f32 K/V (one row × kv_dim × 4 B per layer per window). -- **Cross-session resume.** `unlimited_context`'s archive lives +- **Cross-session resume.** `windowed_checkpoint`'s archive lives in-process; for persisted resume use `boundary_kv` (which emits `larql-boundary` frames to disk). @@ -167,7 +167,7 @@ their window. last frame. Cleaner alternative for "bounded memory at fused speed" since it explicitly composes with `standard` rather than maintaining a shadow K/V store. Should be benchmarked against - W1-GPU'd `unlimited_context` once both are wired. + W1-GPU'd `windowed_checkpoint` once both are wired. - **Page-aligned KV slabs.** The current `CheckpointStore` uses owned `Vec` per layer per checkpoint; a hugepage-backed slab would cut allocation churn during 370K-token replays. diff --git a/crates/larql-inference/examples/apollo_rd_backend.rs b/crates/larql-inference/examples/apollo_rd_backend.rs deleted file mode 100644 index c93b94f75..000000000 --- a/crates/larql-inference/examples/apollo_rd_backend.rs +++ /dev/null @@ -1,684 +0,0 @@ -//! Apollo boundary residual R(D) backend. -//! -//! Contract with `~/chris-source/chris-experiments/shannon/40_boundary_state_rate_distortion/kl_eval.py`: -//! -//! ```text -//! cargo run --release -p larql-inference --example apollo_rd_backend -- \ -//! --model google/gemma-3-4b-it --job JOB_JSON --out OUT_JSON -//! ``` -//! -//! Batch mode uses the same command with a batch job file and writes JSONL. - -use std::collections::HashSet; -use std::fs::File; -use std::io::{Read, Write}; -use std::path::{Path, PathBuf}; -use std::time::Instant; - -use larql_inference::{forward_from_layer, hidden_to_raw_logits, InferenceModel}; -use larql_kv::apollo::{npy, ApolloStore}; -use ndarray::s; -use serde::{Deserialize, Serialize}; - -const KL_DIRECTION: &str = "ground_truth||reconstructed"; - -#[derive(Debug, Deserialize)] -struct BatchJob { - #[serde(default)] - backend_mode: Option, - #[serde(default)] - jobs: Vec, -} - -#[derive(Debug, Deserialize)] -struct Job { - config_id: String, - store: PathBuf, - payload: PathBuf, - boundary_indices: Vec, - source_crystal_layer: Option, - rate_bits_per_token_with_basis: Option, -} - -#[derive(Debug, Serialize)] -struct Metric { - config_id: String, - status: &'static str, - kl_direction: &'static str, - metric_source: &'static str, - kl_mean_nats: f64, - kl_p50_nats: f64, - kl_p95_nats: f64, - kl_p99_nats: f64, - kl_max_nats: f64, - n_positions: usize, - kl_reverse_mean_nats: f64, - kl_reverse_p50_nats: f64, - kl_reverse_p95_nats: f64, - kl_reverse_p99_nats: f64, - kl_reverse_max_nats: f64, - kl_symmetric_mean_nats: f64, - kl_symmetric_p50_nats: f64, - kl_symmetric_p95_nats: f64, - kl_symmetric_max_nats: f64, - kl_symmetric_p99_nats: f64, - sampling_mode: String, - position_stride: usize, - n_windows_evaluated: usize, -} - -#[derive(Debug)] -struct PayloadBoundaries { - reconstructed: Vec>, - rows: usize, - hidden: usize, -} - -#[derive(Debug)] -struct Args { - model: String, - job: PathBuf, - out: PathBuf, - max_positions_per_config: Option, - position_stride: usize, - eval_limit_windows: Option, - adaptive_sampling: bool, - low_rate_threshold: f64, - mid_rate_threshold: f64, - low_rate_max_positions: usize, - mid_rate_max_positions: usize, - high_rate_max_positions: usize, - low_rate_windows: usize, - mid_rate_windows: usize, - high_rate_windows: usize, - low_rate_position_stride: usize, - mid_rate_position_stride: usize, - high_rate_position_stride: usize, - positions_file: Option, -} - -#[derive(Debug, Clone)] -struct EvalPlan { - mode: String, - max_positions: Option, - position_stride: usize, - limit_windows: Option, -} - -#[derive(Debug, Clone)] -struct WindowOffsets { - starts: Vec, - total_tokens: usize, -} - -fn main() -> Result<(), Box> { - let args = parse_args()?; - - eprintln!("Loading model: {}", args.model); - let t0 = Instant::now(); - let model = InferenceModel::load(&args.model)?; - eprintln!("Model loaded in {:.1}s", t0.elapsed().as_secs_f64()); - let weights = model.weights(); - - let job_text = std::fs::read_to_string(&args.job)?; - if let Ok(batch) = serde_json::from_str::(&job_text) { - if batch.backend_mode.as_deref() == Some("batch") { - let mut out = File::create(&args.out)?; - for job in batch.jobs { - let metric = evaluate_job(weights, &job, &args)?; - writeln!(out, "{}", serde_json::to_string(&metric)?)?; - } - return Ok(()); - } - } - - let job: Job = serde_json::from_str(&job_text)?; - let metric = evaluate_job(weights, &job, &args)?; - let mut out = File::create(&args.out)?; - serde_json::to_writer_pretty(&mut out, &metric)?; - writeln!(out)?; - Ok(()) -} - -fn parse_args() -> Result> { - let mut model = None; - let mut job = None; - let mut out = None; - let mut max_positions_per_config = None; - let mut position_stride = 1usize; - let mut eval_limit_windows = None; - let mut adaptive_sampling = false; - let mut low_rate_threshold = 20.0; - let mut mid_rate_threshold = 60.0; - let mut low_rate_max_positions = 128usize; - let mut mid_rate_max_positions = 64usize; - let mut high_rate_max_positions = 8usize; - let mut low_rate_windows = 8usize; - let mut mid_rate_windows = 4usize; - let mut high_rate_windows = 1usize; - let mut low_rate_position_stride = 16usize; - let mut mid_rate_position_stride = 16usize; - let mut high_rate_position_stride = 64usize; - let mut positions_file = None; - - let mut it = std::env::args().skip(1); - while let Some(arg) = it.next() { - match arg.as_str() { - "--model" => model = it.next(), - "--job" => job = it.next().map(PathBuf::from), - "--out" => out = it.next().map(PathBuf::from), - "--max-positions-per-config" => { - max_positions_per_config = it.next().map(|v| v.parse()).transpose()? - } - "--position-stride" => { - position_stride = it - .next() - .ok_or("--position-stride requires a value")? - .parse()?; - if position_stride == 0 { - return Err("--position-stride must be >= 1".into()); - } - } - "--eval-limit-windows" => { - eval_limit_windows = it.next().map(|v| v.parse()).transpose()? - } - "--adaptive-sampling" => adaptive_sampling = true, - "--low-rate-threshold" => { - low_rate_threshold = it - .next() - .ok_or("--low-rate-threshold requires a value")? - .parse()? - } - "--mid-rate-threshold" => { - mid_rate_threshold = it - .next() - .ok_or("--mid-rate-threshold requires a value")? - .parse()? - } - "--low-rate-max-positions" => { - low_rate_max_positions = it - .next() - .ok_or("--low-rate-max-positions requires a value")? - .parse()? - } - "--mid-rate-max-positions" => { - mid_rate_max_positions = it - .next() - .ok_or("--mid-rate-max-positions requires a value")? - .parse()? - } - "--high-rate-max-positions" => { - high_rate_max_positions = it - .next() - .ok_or("--high-rate-max-positions requires a value")? - .parse()? - } - "--low-rate-windows" => { - low_rate_windows = it - .next() - .ok_or("--low-rate-windows requires a value")? - .parse()? - } - "--mid-rate-windows" => { - mid_rate_windows = it - .next() - .ok_or("--mid-rate-windows requires a value")? - .parse()? - } - "--high-rate-windows" => { - high_rate_windows = it - .next() - .ok_or("--high-rate-windows requires a value")? - .parse()? - } - "--low-rate-position-stride" => { - low_rate_position_stride = it - .next() - .ok_or("--low-rate-position-stride requires a value")? - .parse()? - } - "--mid-rate-position-stride" => { - mid_rate_position_stride = it - .next() - .ok_or("--mid-rate-position-stride requires a value")? - .parse()? - } - "--high-rate-position-stride" => { - high_rate_position_stride = it - .next() - .ok_or("--high-rate-position-stride requires a value")? - .parse()? - } - "--positions-file" => positions_file = it.next().map(PathBuf::from), - _ => return Err(format!("unknown argument: {arg}").into()), - } - } - - Ok(Args { - model: model.ok_or("--model required")?, - job: job.ok_or("--job required")?, - out: out.ok_or("--out required")?, - max_positions_per_config, - position_stride, - eval_limit_windows, - adaptive_sampling, - low_rate_threshold, - mid_rate_threshold, - low_rate_max_positions, - mid_rate_max_positions, - high_rate_max_positions, - low_rate_windows, - mid_rate_windows, - high_rate_windows, - low_rate_position_stride, - mid_rate_position_stride, - high_rate_position_stride, - positions_file, - }) -} - -fn evaluate_job( - weights: &larql_inference::ModelWeights, - job: &Job, - args: &Args, -) -> Result> { - eprintln!("Evaluating {}", job.config_id); - let store = ApolloStore::load(&job.store)?; - let payload = load_payload_boundaries(&job.payload)?; - - if payload.rows != job.boundary_indices.len() { - return Err(format!( - "{}: payload rows {} != boundary_indices {}", - job.config_id, - payload.rows, - job.boundary_indices.len() - ) - .into()); - } - if payload.hidden != weights.hidden_size { - return Err(format!( - "{}: payload hidden {} != model hidden {}", - job.config_id, payload.hidden, weights.hidden_size - ) - .into()); - } - - let crystal = job - .source_crystal_layer - .unwrap_or(store.manifest.crystal_layer); - let plan = eval_plan(job, args); - let offsets = WindowOffsets::from_store(&store); - let selected: HashSet = job.boundary_indices.iter().copied().collect(); - let mut values = Vec::new(); - let mut reverse_values = Vec::new(); - let mut symmetric_values = Vec::new(); - let mut windows_evaluated = 0usize; - - if let Some(path) = &args.positions_file { - let positions = load_positions(path)?; - let mut matched_segments = HashSet::new(); - for abs_prefix_len in positions { - let Some((payload_row, start_window, end_window)) = - selected_segment_for_position(job, &offsets, abs_prefix_len) - else { - continue; - }; - let prefix = - segment_prefix_tokens(&store, &offsets, start_window, end_window, abs_prefix_len); - if prefix.is_empty() { - continue; - } - let source_boundary = store - .boundaries - .get(start_window) - .ok_or_else(|| format!("missing source boundary {start_window}"))?; - let reconstructed_boundary = &payload.reconstructed[payload_row]; - let reference_logits = logits_from_boundary(weights, &prefix, source_boundary, crystal); - let reconstructed_logits = - logits_from_boundary(weights, &prefix, reconstructed_boundary, crystal); - push_kl_metrics( - &reference_logits, - &reconstructed_logits, - &mut values, - &mut reverse_values, - &mut symmetric_values, - ); - matched_segments.insert(start_window); - if plan.max_positions.is_some_and(|max| values.len() >= max) { - break; - } - } - windows_evaluated = matched_segments.len(); - } else { - for (payload_row, &start_window) in job.boundary_indices.iter().enumerate() { - if let Some(limit) = plan.limit_windows { - if payload_row >= limit { - break; - } - } - let end_window = next_boundary_or_end( - &job.boundary_indices, - payload_row, - store.window_tokens.len(), - ); - let mut tokens = Vec::new(); - for w in start_window..end_window { - if let Some(window) = store.window_tokens.get(w) { - tokens.extend_from_slice(window); - } - } - if tokens.is_empty() { - continue; - } - let source_boundary = store - .boundaries - .get(start_window) - .ok_or_else(|| format!("missing source boundary {start_window}"))?; - let reconstructed_boundary = &payload.reconstructed[payload_row]; - windows_evaluated += 1; - - // Evaluating all prefixes is exact for per-position KL but expensive. - // `position_stride` and `max_positions_per_config` are explicit - // sampling controls for pilots. - for pos in (1..=tokens.len()).step_by(plan.position_stride) { - if plan.max_positions.is_some_and(|max| values.len() >= max) { - break; - } - let prefix = &tokens[..pos]; - let reference_logits = - logits_from_boundary(weights, prefix, source_boundary, crystal); - let reconstructed_logits = - logits_from_boundary(weights, prefix, reconstructed_boundary, crystal); - push_kl_metrics( - &reference_logits, - &reconstructed_logits, - &mut values, - &mut reverse_values, - &mut symmetric_values, - ); - } - - if plan.max_positions.is_some_and(|max| values.len() >= max) { - break; - } - - // Coarser boundary configs intentionally skip intermediate source - // boundaries. This guard documents that behavior and catches duplicate - // index bugs without changing the segmenting rule. - debug_assert!(selected.contains(&start_window)); - } - } - - if values.is_empty() { - return Err(format!("{} produced no KL positions", job.config_id).into()); - } - let primary = stats(values); - let reverse = stats(reverse_values); - let symmetric = stats(symmetric_values); - Ok(Metric { - config_id: job.config_id.clone(), - status: "complete", - kl_direction: KL_DIRECTION, - metric_source: "apollo_boundary_replay", - kl_mean_nats: primary.mean, - kl_p50_nats: primary.p50, - kl_p95_nats: primary.p95, - kl_p99_nats: primary.p99, - kl_max_nats: primary.max, - n_positions: primary.n, - kl_reverse_mean_nats: reverse.mean, - kl_reverse_p50_nats: reverse.p50, - kl_reverse_p95_nats: reverse.p95, - kl_reverse_p99_nats: reverse.p99, - kl_reverse_max_nats: reverse.max, - kl_symmetric_mean_nats: symmetric.mean, - kl_symmetric_p50_nats: symmetric.p50, - kl_symmetric_p95_nats: symmetric.p95, - kl_symmetric_p99_nats: symmetric.p99, - kl_symmetric_max_nats: symmetric.max, - sampling_mode: if args.positions_file.is_some() { - format!("matched_positions:{}", plan.mode) - } else { - plan.mode - }, - position_stride: plan.position_stride, - n_windows_evaluated: windows_evaluated, - }) -} - -#[derive(Debug)] -struct Stats { - mean: f64, - p50: f64, - p95: f64, - p99: f64, - max: f64, - n: usize, -} - -fn stats(mut values: Vec) -> Stats { - values.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); - let mean = values.iter().sum::() / values.len() as f64; - Stats { - mean, - p50: percentile_sorted(&values, 50.0), - p95: percentile_sorted(&values, 95.0), - p99: percentile_sorted(&values, 99.0), - max: *values.last().unwrap(), - n: values.len(), - } -} - -fn eval_plan(job: &Job, args: &Args) -> EvalPlan { - if !args.adaptive_sampling { - return EvalPlan { - mode: "fixed".to_string(), - max_positions: args.max_positions_per_config, - position_stride: args.position_stride, - limit_windows: args.eval_limit_windows, - }; - } - - let rate = job.rate_bits_per_token_with_basis.unwrap_or(f64::INFINITY); - if rate <= args.low_rate_threshold { - EvalPlan { - mode: "adaptive_low_rate".to_string(), - max_positions: Some(args.low_rate_max_positions), - position_stride: args.low_rate_position_stride.max(1), - limit_windows: Some(args.low_rate_windows), - } - } else if rate <= args.mid_rate_threshold { - EvalPlan { - mode: "adaptive_mid_rate".to_string(), - max_positions: Some(args.mid_rate_max_positions), - position_stride: args.mid_rate_position_stride.max(1), - limit_windows: Some(args.mid_rate_windows), - } - } else { - EvalPlan { - mode: "adaptive_high_rate".to_string(), - max_positions: Some(args.high_rate_max_positions), - position_stride: args.high_rate_position_stride.max(1), - limit_windows: Some(args.high_rate_windows), - } - } -} - -impl WindowOffsets { - fn from_store(store: &ApolloStore) -> Self { - let mut starts = Vec::with_capacity(store.window_tokens.len()); - let mut total_tokens = 0usize; - for window in &store.window_tokens { - starts.push(total_tokens); - total_tokens += window.len(); - } - Self { - starts, - total_tokens, - } - } -} - -fn load_positions(path: &Path) -> Result, Box> { - let text = std::fs::read_to_string(path)?; - let positions: Vec = serde_json::from_str(&text)?; - if positions.contains(&0) { - return Err(format!("positions file {} contains prefix length 0", path.display()).into()); - } - Ok(positions) -} - -fn selected_segment_for_position( - job: &Job, - offsets: &WindowOffsets, - abs_prefix_len: usize, -) -> Option<(usize, usize, usize)> { - if abs_prefix_len == 0 || abs_prefix_len > offsets.total_tokens { - return None; - } - for (payload_row, &start_window) in job.boundary_indices.iter().enumerate() { - let start_abs = *offsets.starts.get(start_window)?; - let end_window = - next_boundary_or_end(&job.boundary_indices, payload_row, offsets.starts.len()); - let end_abs = offsets - .starts - .get(end_window) - .copied() - .unwrap_or(offsets.total_tokens); - if abs_prefix_len > start_abs && abs_prefix_len <= end_abs { - return Some((payload_row, start_window, end_window)); - } - } - None -} - -fn segment_prefix_tokens( - store: &ApolloStore, - offsets: &WindowOffsets, - start_window: usize, - end_window: usize, - abs_prefix_len: usize, -) -> Vec { - let Some(&start_abs) = offsets.starts.get(start_window) else { - return Vec::new(); - }; - if abs_prefix_len <= start_abs { - return Vec::new(); - } - let target_len = abs_prefix_len - start_abs; - let mut tokens = Vec::new(); - for w in start_window..end_window { - if let Some(window) = store.window_tokens.get(w) { - tokens.extend_from_slice(window); - if tokens.len() >= target_len { - tokens.truncate(target_len); - break; - } - } - } - tokens -} - -fn next_boundary_or_end(boundaries: &[usize], row: usize, n_windows: usize) -> usize { - boundaries - .get(row + 1) - .copied() - .unwrap_or(n_windows) - .min(n_windows) -} - -fn logits_from_boundary( - weights: &larql_inference::ModelWeights, - tokens: &[u32], - boundary: &[f32], - crystal: usize, -) -> Vec { - let raw = forward_from_layer( - larql_inference::WeightsView::dense(weights), - tokens, - boundary, - crystal, - None, - ); - let last = raw.h_pre_norm.shape()[0] - 1; - let h_last = raw.h_pre_norm.slice(s![last..=last, ..]).to_owned(); - hidden_to_raw_logits(weights, &h_last) -} - -fn load_payload_boundaries(path: &Path) -> Result> { - let file = File::open(path)?; - let mut archive = zip::ZipArchive::new(file)?; - let mut entry = archive.by_name("reconstructed_boundaries.npy")?; - let mut bytes = Vec::with_capacity(entry.size() as usize); - entry.read_to_end(&mut bytes)?; - let (flat, shape) = npy::read_f32_flat(&bytes)?; - if shape.len() != 2 { - return Err(format!("reconstructed_boundaries must be 2D, got {shape:?}").into()); - } - let rows = shape[0]; - let hidden = shape[1]; - let reconstructed = flat - .chunks_exact(hidden) - .map(|row| row.to_vec()) - .collect::>(); - Ok(PayloadBoundaries { - reconstructed, - rows, - hidden, - }) -} - -fn push_kl_metrics( - reference: &[f32], - reconstructed: &[f32], - primary_values: &mut Vec, - reverse_values: &mut Vec, - symmetric_values: &mut Vec, -) { - let primary = kl_logits(reference, reconstructed); - let reverse = kl_logits(reconstructed, reference); - primary_values.push(primary); - reverse_values.push(reverse); - symmetric_values.push(0.5 * (primary + reverse)); -} - -fn kl_logits(reference: &[f32], reconstructed: &[f32]) -> f64 { - let ref_logp = log_softmax(reference); - let rec_logp = log_softmax(reconstructed); - ref_logp - .iter() - .zip(rec_logp.iter()) - .map(|(&lp, &lq)| { - let p = lp.exp(); - p * (lp - lq) - }) - .sum() -} - -fn log_softmax(logits: &[f32]) -> Vec { - let max = logits - .iter() - .map(|&v| v as f64) - .fold(f64::NEG_INFINITY, f64::max); - let sum_exp = logits - .iter() - .map(|&v| ((v as f64) - max).exp()) - .sum::(); - let log_z = max + sum_exp.ln(); - logits.iter().map(|&v| (v as f64) - log_z).collect() -} - -fn percentile_sorted(values: &[f64], p: f64) -> f64 { - if values.is_empty() { - return f64::NAN; - } - let rank = (p / 100.0) * (values.len().saturating_sub(1) as f64); - let lo = rank.floor() as usize; - let hi = rank.ceil() as usize; - if lo == hi { - values[lo] - } else { - let w = rank - lo as f64; - values[lo] * (1.0 - w) + values[hi] * w - } -} diff --git a/crates/larql-inference/examples/ave_direct_layer_bisect.rs b/crates/larql-inference/examples/ave_direct_layer_bisect.rs deleted file mode 100644 index fa9427af3..000000000 --- a/crates/larql-inference/examples/ave_direct_layer_bisect.rs +++ /dev/null @@ -1,166 +0,0 @@ -//! Per-layer bisect of the direct-matvec decode divergence: gold chain from -//! a staged prefill over prompt+token (with per-layer state capture), direct -//! chain from a decode step over the same cache. The first layer whose input -//! residual diverges names the broken block; K/V row comparison at that -//! layer splits the QKV/RoPE side from the attention-mix/O/FFN side. -//! -//! Usage: `cargo run --release --example ave_direct_layer_bisect -- [VINDEX_DIR]` - -use larql_inference::load_tokenizer; -use larql_inference::vindex::{ - attention_decode_step_native, predict_kquant_decode_step_direct_with_state, - predict_kquant_prefill, predict_kquant_prefill_with_state, -}; -use larql_inference::PerLayerDecodeState; -use ndarray::Array2; - -fn cos_last_vs_first(gold: &Array2, direct: &Array2) -> f32 { - let g = gold.row(gold.nrows() - 1); - let d = direct.row(0); - let dot: f32 = g.iter().zip(d.iter()).map(|(a, b)| a * b).sum(); - let ng: f32 = g.iter().map(|a| a * a).sum::().sqrt(); - let nd: f32 = d.iter().map(|a| a * a).sum::().sqrt(); - if ng == 0.0 || nd == 0.0 { - return f32::NAN; - } - dot / (ng * nd) -} - -fn main() { - let args: Vec = std::env::args().collect(); - let vindex = args - .get(1) - .cloned() - .unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".to_string()); - let dir = std::path::PathBuf::from(&vindex); - if !dir.exists() { - eprintln!("skipped: vindex not found at {vindex}"); - return; - } - - let mut cb = larql_vindex::SilentLoadCallbacks; - eprintln!("Loading {vindex} ..."); - let mut weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index.load_interleaved_kquant(&dir).expect("interleaved"); - index.load_attn_kquant(&dir).expect("attn kquant"); - let tok = load_tokenizer(&dir).expect("tokenizer"); - - let prompt_ids = tok - .encode("12 + 7 =", true) - .expect("encode") - .get_ids() - .to_vec(); - - // First token off the prompt prefill (greedy), as in the parity probe. - let (h, _cache_unused, _) = predict_kquant_prefill(&mut weights, &prompt_ids, &index); - let last = h.nrows() - 1; - let h_last = h.slice(ndarray::s![last..last + 1, ..]).to_owned(); - let logits = larql_inference::forward::hidden_to_raw_logits(&weights, &h_last); - let first_id = logits - .iter() - .enumerate() - .filter(|(_, v)| v.is_finite()) - .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap()) - .map(|(i, _)| i as u32) - .unwrap(); - - // Gold: staged prefill over prompt + first token, capturing per-layer - // h_in / k_new / v_new for every position. - let mut full_ids = prompt_ids.clone(); - full_ids.push(first_id); - let mut gold = PerLayerDecodeState::with_capacity(weights.num_layers); - let _ = predict_kquant_prefill_with_state(&weights, &full_ids, &index, Some(&mut gold)); - - // Direct: fresh prompt-only prefill cache, one direct step with capture. - let (_h2, mut cache, _) = predict_kquant_prefill(&mut weights, &prompt_ids, &index); - let mut direct = PerLayerDecodeState::with_capacity(weights.num_layers); - let backend = larql_compute::default_backend(); - let _ = predict_kquant_decode_step_direct_with_state( - &mut weights, - first_id, - &index, - &*backend, - &mut cache, - prompt_ids.len(), - Some(&mut direct), - ) - .expect("direct step"); - - println!( - "{:>5} {:>10} {:>10} {:>10} (h_in[L] = input residual to layer L; k/v = new rows at L)", - "layer", "cos(h_in)", "cos(k_new)", "cos(v_new)" - ); - for layer in 0..weights.num_layers { - let ch = cos_last_vs_first( - &gold.h_in_per_layer[layer].to_array(), - &direct.h_in_per_layer[layer].to_array(), - ); - let ck = cos_last_vs_first( - &gold.k_new_per_layer[layer].to_array(), - &direct.k_new_per_layer[layer].to_array(), - ); - let cv = cos_last_vs_first( - &gold.v_new_per_layer[layer].to_array(), - &direct.v_new_per_layer[layer].to_array(), - ); - let flag = if ch < 0.999 || ck < 0.999 || cv < 0.999 { - " <-- diverged" - } else { - "" - }; - println!("{layer:>5} {ch:>10.6} {ck:>10.6} {cv:>10.6}{flag}"); - } - - // ── Same-input discriminator: feed each layer's GOLD input residual to - // the direct attention block. Any K/V divergence here is the block - // itself (slice bytes / matvec / norm-rope plumbing), not chain - // compounding. ── - println!("\nSame-input per-layer attention block (gold h_in → direct block):"); - println!( - "{:>5} {:>10} {:>10} {:>6} {:>6} {:>6} {:>6}", - "layer", "cos(k_new)", "cos(v_new)", "q_fmt", "k_fmt", "v_fmt", "o_fmt" - ); - let (_h3, cache_fresh, _) = predict_kquant_prefill(&mut weights, &prompt_ids, &index); - #[allow(clippy::needless_range_loop)] - for layer in 0..weights.num_layers { - let gold_h = gold.h_in_per_layer[layer].to_array(); - let h_last = gold_h - .slice(ndarray::s![gold_h.nrows() - 1..gold_h.nrows(), ..]) - .to_owned(); - let kv_entry = cache_fresh[layer].as_ref(); - let Some((_h_post, (k_cat, v_cat))) = attention_decode_step_native( - &weights, - &index, - &*backend, - &h_last, - layer, - kv_entry, - prompt_ids.len(), - ) else { - println!("{layer:>5} block returned None"); - continue; - }; - let ck = cos_last_vs_first(&gold.k_new_per_layer[layer].to_array(), &{ - let n = k_cat.nrows(); - k_cat.slice(ndarray::s![n - 1..n, ..]).to_owned() - }); - let cv = cos_last_vs_first(&gold.v_new_per_layer[layer].to_array(), &{ - let n = v_cat.nrows(); - v_cat.slice(ndarray::s![n - 1..n, ..]).to_owned() - }); - let fmts = index - .attn_kquant_layer_data(layer) - .map(|a| [a[0].1, a[1].1, a[2].1, a[3].1]) - .unwrap_or(["?"; 4]); - let flag = if ck < 0.999 || cv < 0.999 { - " <-- block diverges on SAME input" - } else { - "" - }; - println!( - "{layer:>5} {ck:>10.6} {cv:>10.6} {:>6} {:>6} {:>6} {:>6}{flag}", - fmts[0], fmts[1], fmts[2], fmts[3] - ); - } -} diff --git a/crates/larql-inference/examples/ave_direct_step_parity.rs b/crates/larql-inference/examples/ave_direct_step_parity.rs deleted file mode 100644 index 0cc4dd94b..000000000 --- a/crates/larql-inference/examples/ave_direct_step_parity.rs +++ /dev/null @@ -1,108 +0,0 @@ -//! One-step parity probe: staged (dequant) vs direct-matvec decode step on a -//! real vindex. Discriminates "my generation loop is wrong" from "the direct -//! kernel path diverges on this model" — compare the same single decode step -//! both ways from an identical prefill cache. -//! -//! Usage: `cargo run --release --example ave_direct_step_parity -- [VINDEX_DIR]` - -use larql_inference::load_tokenizer; -use larql_inference::vindex::{ - predict_kquant_decode_step, predict_kquant_decode_step_direct, predict_kquant_prefill, - supports_direct_matvec_decode, -}; - -fn main() { - let args: Vec = std::env::args().collect(); - let vindex = args - .get(1) - .cloned() - .unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".to_string()); - let dir = std::path::PathBuf::from(&vindex); - if !dir.exists() { - eprintln!("skipped: vindex not found at {vindex}"); - return; - } - - let mut cb = larql_vindex::SilentLoadCallbacks; - eprintln!("Loading {vindex} ..."); - let mut weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index.load_interleaved_kquant(&dir).expect("interleaved"); - index.load_attn_kquant(&dir).expect("attn kquant"); - let tok = load_tokenizer(&dir).expect("tokenizer"); - - println!( - "supports_direct_matvec_decode: {}", - supports_direct_matvec_decode(&weights, &index) - ); - - let prompt_ids = tok - .encode("12 + 7 =", true) - .expect("encode") - .get_ids() - .to_vec(); - - // Two independent prefills → two identical caches (prefill is staged in - // both worlds; only the decode step differs). - let (h, mut cache_staged, _) = predict_kquant_prefill(&mut weights, &prompt_ids, &index); - let (_h2, mut cache_direct, _) = predict_kquant_prefill(&mut weights, &prompt_ids, &index); - - // Greedy-pick the first token off the prefill logits (shared). - let last = h.nrows() - 1; - let h_last = h.slice(ndarray::s![last..last + 1, ..]).to_owned(); - let logits = larql_inference::forward::hidden_to_raw_logits(&weights, &h_last); - let first_id = logits - .iter() - .enumerate() - .filter(|(_, v)| v.is_finite()) - .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap()) - .map(|(i, _)| i as u32) - .unwrap(); - println!( - "first greedy token: {} {:?}", - first_id, - tok.decode(&[first_id], true).unwrap_or_default() - ); - - let abs_position = prompt_ids.len(); - let (h_staged, _) = - predict_kquant_decode_step(&weights, first_id, &index, &mut cache_staged, abs_position) - .expect("staged step"); - let backend = larql_compute::default_backend(); - let h_direct = predict_kquant_decode_step_direct( - &mut weights, - first_id, - &index, - &*backend, - &mut cache_direct, - abs_position, - ) - .expect("direct step"); - - // Compare hidden states. - let a = h_staged.row(0); - let b = h_direct.row(0); - let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum(); - let na: f32 = a.iter().map(|x| x * x).sum::().sqrt(); - let nb: f32 = b.iter().map(|x| x * x).sum::().sqrt(); - let max_abs = a - .iter() - .zip(b.iter()) - .map(|(x, y)| (x - y).abs()) - .fold(0f32, f32::max); - println!("hidden cosine(staged, direct): {:.6}", dot / (na * nb)); - println!("hidden max |diff|: {max_abs:.6} norms: staged {na:.3} direct {nb:.3}"); - - // And the next-token view: top-3 from each. - let top3 = |h: &ndarray::Array2| -> Vec<(u32, String)> { - let logits = larql_inference::forward::hidden_to_raw_logits(&weights, h); - let mut idx: Vec = (0..logits.len()).collect(); - idx.sort_by(|&i, &j| logits[j].partial_cmp(&logits[i]).unwrap()); - idx.iter() - .take(3) - .map(|&i| (i as u32, tok.decode(&[i as u32], true).unwrap_or_default())) - .collect() - }; - println!("staged next top-3: {:?}", top3(&h_staged)); - println!("direct next top-3: {:?}", top3(&h_direct)); -} diff --git a/crates/larql-inference/examples/ave_q4k_row_audit.rs b/crates/larql-inference/examples/ave_q4k_row_audit.rs deleted file mode 100644 index 37df3120a..000000000 --- a/crates/larql-inference/examples/ave_q4k_row_audit.rs +++ /dev/null @@ -1,205 +0,0 @@ -//! Row-level audit of the two Q4_K decoders on real attention bytes. -//! -//! Same bytes, same activation vector, three readings per row of the K -//! projection: -//! a) `q4k_matvec_into` (the direct-path f32-act kernel), -//! b) `dequantize_q4_k` row → f32 dot (reference decode), -//! c) the staged path's `insert_q4k_layer_tensors` tensor row → dot. -//! A row where (a) disagrees with (b)/(c) pinpoints a super-block decode -//! bug in the matvec kernel; (b) vs (c) checks the two dequantisers -//! against each other. -//! -//! Usage: `cargo run --release --example ave_q4k_row_audit -- [VINDEX_DIR] [LAYERS...]` - -use larql_compute::cpu::ops::q4_common::{dequantize_q4_k, q4k_matvec_into}; -use larql_inference::vindex::{insert_q4k_layer_tensors_resident, remove_layer_tensors_resident}; - -fn main() { - if std::env::var("LARQL_F16_PROBE").is_ok() { - f16_probe(); - } - let args: Vec = std::env::args().collect(); - let vindex = args - .get(1) - .cloned() - .unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".to_string()); - let layers: Vec = if args.len() > 2 { - args[2..].iter().filter_map(|a| a.parse().ok()).collect() - } else { - vec![20, 32] // clean control, worst offender - }; - let dir = std::path::PathBuf::from(&vindex); - if !dir.exists() { - eprintln!("skipped: vindex not found at {vindex}"); - return; - } - - let mut cb = larql_vindex::SilentLoadCallbacks; - eprintln!("Loading {vindex} ..."); - let mut weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index.load_interleaved_kquant(&dir).expect("interleaved"); - index.load_attn_kquant(&dir).expect("attn kquant"); - - let hidden = weights.hidden_size; - let arch_kv = { - let arch = &*weights.arch; - arch.num_kv_heads_for_layer(0) * arch.head_dim_for_layer(0) - }; - // Deterministic pseudo-random activation (no Math.random in harness - // discipline; LCG is plenty for a kernel audit). - let mut seed = 0x2545F4914F6CDD1Du64; - let x: Vec = (0..hidden) - .map(|_| { - seed = seed - .wrapping_mul(6364136223846793005) - .wrapping_add(1442695040888963407); - ((seed >> 33) as f32 / (1u64 << 31) as f32) - 0.5 - }) - .collect(); - - const BLOCK_BYTES: usize = 144; - const ELEMS: usize = 256; - let bytes_per_row = (hidden / ELEMS) * BLOCK_BYTES; - - for &layer in &layers { - let attn = index.attn_kquant_layer_data(layer).expect("attn data"); - let (k_bytes, k_fmt) = attn[1]; - println!( - "\nlayer {layer}: k_fmt={k_fmt} kv_dim={arch_kv} bytes={}", - k_bytes.len() - ); - if k_fmt != "Q4_K" { - println!(" (not Q4_K, skipping)"); - continue; - } - - // Staged tensor for (c). - let k_bytes_owned = k_bytes.to_vec(); - let inserted = - insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("insert"); - let k_key = weights.arch.attn_k_key(layer); - let w_staged = weights.tensors.get(&k_key).expect("staged K").clone(); - println!(" staged tensor shape: {:?}", w_staged.shape()); - - let mut bad_ab = 0usize; - let mut bad_ac = 0usize; - let mut bad_bc = 0usize; - let mut worst: (usize, f32, f32, f32) = (0, 0.0, 0.0, 0.0); - for r in 0..arch_kv { - let row_bytes = &k_bytes_owned[r * bytes_per_row..(r + 1) * bytes_per_row]; - let mut a = [0.0f32]; - q4k_matvec_into(&mut a, &x, row_bytes, 1, hidden); - let deq = dequantize_q4_k(row_bytes, hidden); - let b: f32 = deq.iter().zip(x.iter()).map(|(w, v)| w * v).sum(); - // (c): staged row — orientation per dequantize_matrix(rows=kv_dim, cols=hidden). - let c: f32 = if w_staged.shape()[0] == arch_kv { - w_staged - .row(r) - .iter() - .zip(x.iter()) - .map(|(w, v)| w * v) - .sum() - } else { - w_staged - .column(r) - .iter() - .zip(x.iter()) - .map(|(w, v)| w * v) - .sum() - }; - let scale = b.abs().max(1e-3); - let dab = (a[0] - b).abs() / scale; - let dac = (a[0] - c).abs() / scale; - let dbc = (b - c).abs() / scale; - if dab > 1e-3 { - bad_ab += 1; - } - if dac > 1e-3 { - bad_ac += 1; - } - if dbc > 1e-3 { - bad_bc += 1; - } - if dab > worst.1 { - worst = (r, dab, a[0], b); - } - } - println!( - " rows with rel-diff > 1e-3 of {arch_kv}: matvec-vs-deq(a,b): {bad_ab} matvec-vs-staged(a,c): {bad_ac} deq-vs-staged(b,c): {bad_bc}" - ); - println!( - " worst row {}: rel {:.4} matvec {:.6} vs dequant-dot {:.6}", - worst.0, worst.1, worst.2, worst.3 - ); - - // Element-level: q4_common dequant vs the staged tensor row, no dot - // products involved. If the decode logic were identical these are - // bit-equal; print the worst element diff found anywhere. - let mut worst_elem: (usize, usize, f32, f32, f32) = (0, 0, 0.0, 0.0, 0.0); - let mut rows_with_elem_diff = 0usize; - for r in 0..arch_kv { - let row_bytes = &k_bytes_owned[r * bytes_per_row..(r + 1) * bytes_per_row]; - let deq = dequantize_q4_k(row_bytes, hidden); - let staged_row = w_staged.row(r); - let mut row_worst = 0f32; - for (i, (b, c)) in deq.iter().zip(staged_row.iter()).enumerate() { - let d = (b - c).abs(); - if d > row_worst { - row_worst = d; - } - if d > worst_elem.4 { - worst_elem = (r, i, *b, *c, d); - } - } - if row_worst > 1e-7 { - rows_with_elem_diff += 1; - } - } - println!( - " element-level: rows with any |Δ|>1e-7: {rows_with_elem_diff}/{arch_kv}; worst at row {} elem {}: q4_common {} vs staged {} (|Δ| {})", - worst_elem.0, worst_elem.1, worst_elem.2, worst_elem.3, worst_elem.4 - ); - // Forensic dump of the worst block: both decoders on the same 144 - // bytes, plus the raw header, so the layout disagreement is visible. - if worst_elem.4 > 0.0 { - let (r, i) = (worst_elem.0, worst_elem.1); - let blk = i / 256; - let row_bytes = &k_bytes_owned[r * bytes_per_row..(r + 1) * bytes_per_row]; - let block = &row_bytes[blk * 144..(blk + 1) * 144]; - println!( - " forensic block row {r} block {blk} (elem {i} = in-block {}):", - i % 256 - ); - println!(" header[0..16]: {:02x?}", &block[0..16]); - let via_common = dequantize_q4_k(block, 256); - let info = larql_vindex::quant::registry::lookup("Q4_K").expect("registry"); - let via_registry = (info.dequantize)(block, 256).expect("registry decode"); - let e = i % 256; - let lo = e.saturating_sub(4); - let hi = (e + 4).min(255); - println!(" elems {lo}..={hi}:"); - println!(" q4_common: {:?}", &via_common[lo..=hi]); - println!(" registry : {:?}", &via_registry[lo..=hi]); - let n_diff = via_common - .iter() - .zip(via_registry.iter()) - .filter(|(a, b)| (**a - **b).abs() > 1e-7) - .count(); - println!(" elems differing in this block: {n_diff}/256"); - } - remove_layer_tensors_resident(&mut weights, inserted); - } -} - -#[allow(dead_code)] -fn f16_probe() { - // Called from main when LARQL_F16_PROBE=1. - let bits = 0x03feu16; - println!( - "f16(0x03fe): q4_common={:e} models={:e} (true subnormal = 1022*2^-24 = {:e})", - larql_compute::cpu::ops::q4_common::f16_to_f32(bits), - larql_models::quant::half::f16_to_f32(bits), - 1022f32 * 2f32.powi(-24), - ); -} diff --git a/crates/larql-inference/examples/ave_stream_trigger_probe.rs b/crates/larql-inference/examples/ave_stream_trigger_probe.rs deleted file mode 100644 index 19fac2845..000000000 --- a/crates/larql-inference/examples/ave_stream_trigger_probe.rs +++ /dev/null @@ -1,422 +0,0 @@ -//! Stream-trigger measurement (pre-registered) — does the model's -//! spontaneous restatement reflex support mid-stream dispatch? -//! -//! The observation under test: on arithmetic word problems the model -//! reliably rewrites prose into notation (`123456 + 654321 = `) before -//! face-planting on the digits. If that reflex is frequent AND faithful, -//! stream-gating on the model's own emitted `expr =` gives the disguised -//! path with no probe, no instructed rewrite, no intent heuristics — the -//! engagement signal expressed in tokens, auditable in the transcript. -//! -//! Two arms per item — `bare` (raw completion: the spontaneous reflex) -//! and `cot` (a one-line generic step-by-step nudge: no examples, no -//! format rigging — the deployment shape). AMENDMENT NOTE: the cot arm -//! was added after the bare arm's interim fire rate (~0.4) was visible, -//! on the observation that CoT rewrites into notation; thresholds are -//! inherited unchanged and the cot arm runs blind. -//! -//! Three numbers per item, plus one release-mode cell: -//! 1. FIRE — does a trigger (`expr =`) appear within budget? -//! 2. FIDELITY — is the emitted expression the RIGHT expression -//! (scored against ground-truth operands/op — the A13b -//! expression-echo discipline applied to the trigger itself)? -//! 3. POSITION — tokens until first trigger; trigger multiplicity. -//! -//! Plus the RELEASE cell — splice the ALU answer at the trigger, release -//! the mask, count post-schedule digit overruns (the A10 ~4% mode, in -//! its mid-sentence form). -//! -//! Pre-registered branches: -//! - fire ≥ 0.8 AND fidelity ≥ 0.95 of fired → build the stream-gate. -//! - fire ≥ 0.8 AND fidelity < 0.95 → trigger = engagement -//! signal only; the payload still needs the instructed rewrite. -//! - fire < 0.5 → the reflex was the -//! prompt family talking; disguised path stays parked. -//! -//! Usage: `cargo run --release --example ave_stream_trigger_probe -- [VINDEX_DIR]` -//! Writes `bench/aim-validation/ave_stream_trigger_gemma3-4b.json`. - -use larql_inference::experts::arith::extract::find_triggers; -use larql_inference::load_tokenizer; -use larql_inference::vindex::generate_kquant_cpu_constrained_cached_streaming; - -/// (word problem, canonical expression the model SHOULD restate). -/// No notation in the prompt — these are disguised asks; tier-0 stays -/// cold on all of them by construction. -const PROBLEMS: &[(&str, &str)] = &[ - // addition, varied phrasing - ( - "If you have 38 apples and pick 17 more, how many apples do you have?", - "38 + 17", - ), - ( - "What do you get when you add 123456 and 654321?", - "123456 + 654321", - ), - ("What is the sum of 999 and 111?", "999 + 111"), - ( - "A tank holds 4500 liters and 2750 more are pumped in. How much is in the tank?", - "4500 + 2750", - ), - ( - "Tom scored 1284 points and then earned another 716. What is his total?", - "1284 + 716", - ), - ("Add 87 to 246.", "246 + 87"), - ( - "A library has 58210 books and acquires 4790 new ones. How many books now?", - "58210 + 4790", - ), - ("What is 312487 increased by 96513?", "312487 + 96513"), - // subtraction - ( - "Sarah had 5000 dollars and spent 1234. How much does she have left?", - "5000 - 1234", - ), - ("Take 250 away from 1000.", "1000 - 250"), - ( - "John is 47 and Mary is 23 years younger. How old is Mary?", - "47 - 23", - ), - ( - "A warehouse stored 90000 crates and shipped 12345. How many remain?", - "90000 - 12345", - ), - ("What is 700 minus 458?", "700 - 458"), - ("From 86420 subtract 13579.", "86420 - 13579"), - ( - "A flight covers 5400 km and 1750 km are already behind. How far is left?", - "5400 - 1750", - ), - // multiplication - ( - "A crate holds 240 bottles. How many bottles are in 12 crates?", - "240 * 12", - ), - ("Multiply 73 by 19.", "73 * 19"), - ( - "A factory makes 1500 widgets a day. How many in 365 days?", - "1500 * 365", - ), - ( - "Each of the 48 rows has 96 seats. How many seats in total?", - "48 * 96", - ), - ("What is the product of 407 and 311?", "407 * 311"), - ( - "Nine hundred boxes each weigh 75 kilos. What is the total weight?", - "900 * 75", - ), - // two-op chains (multiplicity watch) - ("What is 47 plus 358 plus 1200?", "47 + 358 + 1200"), - ( - "Start with 999, add 111, then take away 222. What do you get?", - "999 + 111 - 222", - ), - ( - "A bus starts with 50 passengers, then 23 get off and 12 get on. How many are aboard?", - "50 - 23 + 12", - ), -]; - -/// (arm name, prompt suffix, generation budget). -const ARMS: &[(&str, &str, usize)] = &[ - ("bare", "", 64), - ("cot", "\n\nLet's work this out step by step:\n", 80), -]; - -fn main() { - let args: Vec = std::env::args().collect(); - let vindex = args - .get(1) - .cloned() - .unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".to_string()); - let dir = std::path::PathBuf::from(&vindex); - if !dir.exists() { - eprintln!("skipped: vindex not found at {vindex}"); - return; - } - - let mut cb = larql_vindex::SilentLoadCallbacks; - eprintln!("Loading {vindex} ..."); - let mut weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index.load_interleaved_kquant(&dir).expect("interleaved"); - index.load_attn_kquant(&dir).expect("attn kquant"); - let tok = load_tokenizer(&dir).expect("tokenizer"); - - println!("\n=== stream-trigger probe on {vindex} ==="); - - let mut json_rows = String::new(); - let mut arm_summaries: Vec<(String, usize, usize, usize, usize, usize)> = Vec::new(); - - for (arm, suffix, budget) in ARMS { - println!("\n ── arm: {arm} (budget {budget} tok) ──"); - println!( - "{:<4} {:>5} {:>9} {:>6} {:>5} emitted-expr (vs expected)", - "item", "fire", "fidelity", "pos", "n_trg" - ); - - let mut fired = 0usize; - let mut faithful = 0usize; - let mut positions: Vec = Vec::new(); - let mut multi = 0usize; - - for (idx, (prompt, expected)) in PROBLEMS.iter().enumerate() { - let full_prompt = format!("{prompt}{suffix}"); - let prompt_ids = tok - .encode(full_prompt.as_str(), true) - .expect("encode") - .get_ids() - .to_vec(); - - // Stream and record the token position at which the first trigger - // completes — the same incremental read the gate would perform. - let mut emitted = String::new(); - let mut first_trigger_pos: Option = None; - let mut n_tokens = 0usize; - let out = generate_kquant_cpu_constrained_cached_streaming( - &mut weights, - &tok, - &prompt_ids, - *budget, - &index, - |_, _| {}, - |_, text| { - emitted.push_str(text); - n_tokens += 1; - if first_trigger_pos.is_none() - && text.contains('=') - && !find_triggers(&emitted).is_empty() - { - first_trigger_pos = Some(n_tokens); - } - }, - ); - let _ = out; - let triggers = find_triggers(&emitted); - let fire = !triggers.is_empty(); - let n_trg = triggers.len(); - let first_expr = triggers.first().map(|(e, _)| e.to_string()); - // Fidelity: the FIRST emitted trigger must be the ground-truth - // expression (operands and ops, exact, order-insensitive only via - // the canonical string — the harness corpus is written in the - // model's natural restatement order). - let correct = first_expr.as_deref() == Some(*expected); - - fired += usize::from(fire); - faithful += usize::from(correct); - if let Some(p) = first_trigger_pos { - positions.push(p); - } - multi += usize::from(n_trg > 1); - - println!( - "{:<4} {:>5} {:>9} {:>6} {:>5} {} (exp {})", - idx, - if fire { "✓" } else { "—" }, - if !fire { - "n/a" - } else if correct { - "✓" - } else { - "✗ WRONG" - }, - first_trigger_pos - .map(|p| p.to_string()) - .unwrap_or_else(|| "-".into()), - n_trg, - first_expr.as_deref().unwrap_or("-"), - expected, - ); - json_rows.push_str(&format!( - "{}{{\"arm\":\"{arm}\",\"prompt\":{},\"expected\":{},\"fire\":{fire},\"emitted_expr\":{},\"correct\":{correct},\"pos\":{},\"n_triggers\":{n_trg},\"emission\":{}}}", - if json_rows.is_empty() { "" } else { "," }, - serde_json::to_string(prompt).expect("json"), - serde_json::to_string(expected).expect("json"), - serde_json::to_string(&first_expr).expect("json"), - first_trigger_pos.map(|p| p as i64).unwrap_or(-1), - serde_json::to_string(emitted.trim()).expect("json"), - )); - } - - let median_pos = { - let mut p = positions.clone(); - p.sort_unstable(); - p.get(p.len() / 2).copied().unwrap_or(0) - }; - arm_summaries.push(( - arm.to_string(), - fired, - faithful, - median_pos, - multi, - PROBLEMS.len(), - )); - } - - // ── Release-mode cell: splice at the trigger, release the mask, - // count post-schedule digit overruns. The splice payload is the ALU - // result of the EMITTED expression (honest end-to-end: wrong emitted - // expr → wrong splice, which fidelity already scores). ── - println!( - "\n ── release-mode cell (splice at trigger, release mask, watch for digit overrun) ──" - ); - let mut release_runs = 0usize; - let mut overruns = 0usize; - let mut release_rows = String::new(); - let cot_suffix = ARMS[1].1; - for (prompt, _expected) in PROBLEMS.iter().take(10) { - let full_prompt = format!("{prompt}{cot_suffix}"); - let prompt_ids = tok - .encode(full_prompt.as_str(), true) - .expect("encode") - .get_ids() - .to_vec(); - - // Stateful stream-gate split across the mask closure (reads) and - // the token callback (writes) — shared via RefCell since both - // borrow the same state, sequentially per step. This is the - // future controller in miniature. - #[derive(Default)] - struct GateState { - emitted: String, - schedule: Option>, - forced: usize, - done_forcing: bool, - released_tail: String, - } - let state = std::cell::RefCell::new(GateState::default()); - let tok_ref = &tok; - let out = generate_kquant_cpu_constrained_cached_streaming( - &mut weights, - &tok, - &prompt_ids, - 96, - &index, - |_generated, logits| { - let s = state.borrow(); - if s.done_forcing { - return; // released — model continues unmasked - } - if let Some(sched) = &s.schedule { - if s.forced < sched.len() { - let want = sched[s.forced]; - for (i, l) in logits.iter_mut().enumerate() { - if i as u32 != want { - *l = f32::NEG_INFINITY; - } - } - if let Some(l) = logits.get_mut(want as usize) { - if !l.is_finite() { - *l = 0.0; - } - } - } - } - }, - |_, text| { - let mut s = state.borrow_mut(); - s.emitted.push_str(text); - if s.schedule.is_none() { - if let Some((expr, _)) = find_triggers(&s.emitted).into_iter().next() { - let answer = expr.eval(); - let ids = tok_ref - .encode(format!(" {answer}").as_str(), false) - .map(|e| e.get_ids().to_vec()) - .unwrap_or_default(); - if !ids.is_empty() { - s.schedule = Some(ids); - } - } - } else if !s.done_forcing { - s.forced += 1; - if s.forced >= s.schedule.as_ref().map(|v| v.len()).unwrap_or(0) { - s.done_forcing = true; - } - } else { - s.released_tail.push_str(text); - } - }, - ); - let _ = out; - let state = state.into_inner(); - let (released_tail, done_forcing, had_schedule) = ( - state.released_tail, - state.done_forcing, - state.schedule.is_some(), - ); - if had_schedule && done_forcing { - release_runs += 1; - // Overrun = the released model immediately continues the - // number (first non-space char of the tail is a digit). - let overrun = released_tail - .trim_start() - .chars() - .next() - .is_some_and(|c| c.is_ascii_digit()); - overruns += usize::from(overrun); - println!( - " {:<58} overrun: {} tail: {:?}", - format!("{prompt:?}"), - if overrun { "✗ YES" } else { "✓ no" }, - released_tail.chars().take(28).collect::(), - ); - release_rows.push_str(&format!( - "{}{{\"prompt\":{},\"overrun\":{overrun},\"tail\":{}}}", - if release_rows.is_empty() { "" } else { "," }, - serde_json::to_string(prompt).expect("json"), - serde_json::to_string(released_tail.trim()).expect("json"), - )); - } else { - println!( - " {:<58} (no trigger within budget — release cell skipped)", - format!("{prompt:?}") - ); - } - } - - // ── verdict ────────────────────────────────────────────────────── - println!("\n ── verdict ──"); - let mut arm_json = String::new(); - for (arm, fired, faithful, median_pos, multi, n) in &arm_summaries { - let fire_rate = *fired as f64 / *n as f64; - let fidelity = if *fired > 0 { - *faithful as f64 / *fired as f64 - } else { - 0.0 - }; - let branch = if fire_rate >= 0.8 && fidelity >= 0.95 { - "BUILD: stream-gate on the model's own `expr =` — disguised path without probe or rewrite" - } else if fire_rate >= 0.8 { - "ENGAGEMENT-ONLY: trigger fires but emitted exprs unfaithful — payload needs the instructed rewrite" - } else if fire_rate < 0.5 { - "PARKED: restatement reflex insufficient in this arm" - } else { - "GRAY ZONE: fire rate between branches — widen the corpus before deciding" - }; - println!( - " [{arm}] fire: {fired}/{n} ({fire_rate:.2}) fidelity-of-fired: {faithful}/{fired} ({fidelity:.2}) median pos: {median_pos} tok multi-trigger: {multi}" - ); - println!(" [{arm}] branch: {branch}"); - arm_json.push_str(&format!( - "{}{{\"arm\":\"{arm}\",\"fire\":[{fired},{n}],\"fidelity_of_fired\":[{faithful},{fired}],\"median_pos\":{median_pos},\"multi_trigger\":{multi},\"branch\":{}}}", - if arm_json.is_empty() { "" } else { "," }, - serde_json::to_string(branch).expect("json"), - )); - } - println!( - " release cell (cot arm): {overruns}/{release_runs} digit overruns (guard-token mitigation clamps this by construction)" - ); - - let json = format!( - "{{\"experiment\":\"ave_stream_trigger\",\"vindex\":{},\"arms\":[{arm_json}],\"release_overruns\":[{overruns},{release_runs}],\"items\":[{json_rows}],\"release_cell\":[{release_rows}]}}", - serde_json::to_string(&vindex).expect("json"), - ); - let out_path = "bench/aim-validation/ave_stream_trigger_gemma3-4b.json"; - if let Err(e) = std::fs::write(out_path, &json) { - eprintln!("warning: could not write {out_path}: {e}"); - } else { - println!("\nwrote {out_path}"); - } -} diff --git a/crates/larql-inference/examples/bitnet_e2e.rs b/crates/larql-inference/examples/bitnet_e2e.rs deleted file mode 100644 index d06d37788..000000000 --- a/crates/larql-inference/examples/bitnet_e2e.rs +++ /dev/null @@ -1,53 +0,0 @@ -//! End-to-end BitNet b1.58 check on the A8-wired forward. -//! -//! cargo run --release -p larql-inference --example bitnet_e2e -- ["prompt"] -//! -//! Loads the native-ternary vindex, greedily generates, and prints the -//! continuation + tok/s. With the A8 (int8-activation) path wired in, this -//! is the end-to-end gate: the model must still produce sensible text. - -use std::path::Path; -use std::time::Instant; - -use larql_inference::ternary::{generate, load_bitnet_model}; - -fn main() { - let mut args = std::env::args().skip(1); - let vindex = args - .next() - .expect("usage: bitnet_e2e [prompt]"); - let prompt = args - .next() - .unwrap_or_else(|| "The capital of France is".to_string()); - let vindex = Path::new(&vindex); - - eprintln!("loading BitNet model from {} ...", vindex.display()); - let t = Instant::now(); - let model = load_bitnet_model(vindex).expect("load_bitnet_model"); - eprintln!("loaded in {:.1}s", t.elapsed().as_secs_f64()); - - let tok = larql_vindex::tokenizers::Tokenizer::from_file(vindex.join("tokenizer.json")) - .expect("load tokenizer.json"); - let enc = tok.encode(prompt.as_str(), true).expect("encode prompt"); - let prompt_ids: Vec = enc.get_ids().to_vec(); - eprintln!("prompt = {prompt:?} ({} tokens)", prompt_ids.len()); - - let max_new = 16usize; - let t = Instant::now(); - let gen_ids = generate(&model, &tok, &prompt_ids, max_new, None); - let dt = t.elapsed().as_secs_f64(); - - let continuation = tok.decode(&gen_ids, true).expect("decode"); - println!("\n=== {prompt}|{continuation}"); - println!( - "\ngenerated {} tokens in {:.2}s → {:.1} tok/s", - gen_ids.len(), - dt, - gen_ids.len() as f64 / dt - ); - if continuation.contains("Paris") { - println!("✅ 'Paris' present — A8 forward is coherent end-to-end"); - } else { - println!("⚠️ 'Paris' not found — inspect the continuation above"); - } -} diff --git a/crates/larql-inference/examples/fr1_topk_fuzzy_router.rs b/crates/larql-inference/examples/fr1_topk_fuzzy_router.rs deleted file mode 100644 index 2d35d2e66..000000000 --- a/crates/larql-inference/examples/fr1_topk_fuzzy_router.rs +++ /dev/null @@ -1,434 +0,0 @@ -//! FR1 — top-k fuzzy entity router on a REAL LARQL vindex (the measurement, -//! before any build). Reproduces fleet E15 against the production `KnnStore` -//! cosine router, and indicts the live inference path, which routes on -//! `query_top1` + a fixed 0.75 cosine gate (`infer_patched.rs:162-163`) while -//! `query_knn` (top-k, `knn_store.rs:132`) sits built and unused. -//! -//! Method (faithful to production — keys + queries both come from -//! `capture_residuals` at the same layer, exactly as `INSERT … MODE KNN` does). -//! For N real countries, capture the last-token residual at a layer sweep -//! {20,22,24,26} for three phrasings (one forward each): -//! -//! ```text -//! TRAIN "The capital of {e} is" -> the stored key (relation=capital) -//! PARA "{e}'s capital city is" -> held-out paraphrase query -//! CROSS "The currency of {e} is" -> cross-relation confound query -//! ``` -//! -//! Build one `KnnStore` per layer from the TRAIN residuals, then route the PARA -//! and CROSS residuals through `query_knn` and score in predictive units -//! (recall@k, NOT mean cosine): -//! -//! ```text -//! recall@{1,3,5,10} expect top-1 ~0.7, top-5 ~0.9 (E15) -//! top-1 margin (cos1 - cos2) the razor-thin near-rank-1 claim (E11) -//! confident-wrong @0.75 fires the live gate but wrong = the indictment -//! CROSS recall entity key vs answer-leak (E15 firewall) -//! ``` -//! -//! Usage: `cargo run --release --example fr1_topk_fuzzy_router -- [VINDEX_DIR] [N]` -//! Writes `bench/aim-validation/fr1_topk_router_gemma3-4b.json`. - -use larql_inference::vindex::insert_q4k_layer_tensors_resident; -use larql_inference::{capture_residuals, load_tokenizer}; -use larql_vindex::KnnStore; -use std::collections::HashMap; - -const LAYERS: [usize; 4] = [20, 22, 24, 26]; -const KS: [usize; 4] = [1, 3, 5, 10]; - -/// 150 real countries the model knows — the fleet E15 / mechanism `route.py` -/// set, verbatim, so the LARQL number is comparable to the MLX number. -const COUNTRIES: &[&str] = &[ - "France", - "Germany", - "Italy", - "Spain", - "Portugal", - "Greece", - "Austria", - "Switzerland", - "Belgium", - "Netherlands", - "Denmark", - "Norway", - "Sweden", - "Finland", - "Iceland", - "Ireland", - "Poland", - "Hungary", - "Romania", - "Bulgaria", - "Croatia", - "Serbia", - "Ukraine", - "Russia", - "Turkey", - "Japan", - "China", - "India", - "Pakistan", - "Bangladesh", - "Thailand", - "Vietnam", - "Indonesia", - "Malaysia", - "Philippines", - "Singapore", - "Mongolia", - "Nepal", - "Cambodia", - "Laos", - "Brazil", - "Argentina", - "Chile", - "Peru", - "Colombia", - "Venezuela", - "Ecuador", - "Bolivia", - "Paraguay", - "Uruguay", - "Mexico", - "Cuba", - "Jamaica", - "Canada", - "Australia", - "New Zealand", - "Egypt", - "Morocco", - "Algeria", - "Tunisia", - "Libya", - "Kenya", - "Nigeria", - "Ghana", - "Ethiopia", - "Tanzania", - "Uganda", - "Angola", - "Zambia", - "Zimbabwe", - "Senegal", - "Mali", - "Sudan", - "Somalia", - "Cameroon", - "Iran", - "Iraq", - "Israel", - "Jordan", - "Lebanon", - "Syria", - "Yemen", - "Oman", - "Qatar", - "Kuwait", - "Bahrain", - "Armenia", - "Georgia", - "Azerbaijan", - "Kazakhstan", - "Uzbekistan", - "Turkmenistan", - "Afghanistan", - "Sri Lanka", - "South Korea", - "North Korea", - "Taiwan", - "Estonia", - "Latvia", - "Lithuania", - "Slovakia", - "Slovenia", - "Luxembourg", - "Malta", - "Cyprus", - "Albania", - "Montenegro", - "Moldova", - "Belarus", - "Kyrgyzstan", - "Tajikistan", - "Bhutan", - "Myanmar", - "Brunei", - "Botswana", - "Namibia", - "Mozambique", - "Madagascar", - "Malawi", - "Rwanda", - "Burundi", - "Chad", - "Niger", - "Mauritania", - "Gabon", - "Congo", - "Liberia", - "Guinea", - "Benin", - "Togo", - "Gambia", - "Panama", - "Costa Rica", - "Nicaragua", - "Honduras", - "Guatemala", - "Belize", - "Guyana", - "Suriname", - "Haiti", - "Bahamas", - "Fiji", - "South Africa", - "United Kingdom", - "United States", - "Dominican Republic", - "El Salvador", - "Sierra Leone", - "Mauritius", - "Maldives", - "Papua New Guinea", - "Eritrea", - "Djibouti", - "Lesotho", -]; - -/// Mirror of the production gate: a stored key whose top-1 cosine exceeds this -/// would replace the model's prediction (`KNN_COSINE_THRESHOLD` in -/// `infer_patched.rs`). We measure how often that fires *and is wrong*. -const GATE: f32 = 0.75; - -struct Summary { - recall: [f64; 4], - margin_mean: f64, - margin_p50: f64, - margin_min: f64, - gate_fires: usize, // top-1 cosine > GATE - gate_wrong: usize, // top-1 cosine > GATE AND wrong entity -} - -fn percentile(sorted: &[f32], p: f64) -> f64 { - if sorted.is_empty() { - return 0.0; - } - let idx = ((p * (sorted.len() - 1) as f64).round() as usize).min(sorted.len() - 1); - sorted[idx] as f64 -} - -/// Score one condition (paraphrase or cross) against a layer's store. -fn score(store: &KnnStore, layer: usize, queries: &[Vec], entities: &[String]) -> Summary { - let n = entities.len(); - let mut recall = [0usize; 4]; - let mut margins: Vec = Vec::with_capacity(n); - let mut gate_fires = 0usize; - let mut gate_wrong = 0usize; - - for (i, q) in queries.iter().enumerate() { - let hits = store.query_knn(layer, q, 10); - if hits.is_empty() { - continue; - } - // Rank of the true entity by exact name match. - let rank = hits.iter().position(|(e, _)| e.entity == entities[i]); - if let Some(r) = rank { - for (ki, k) in KS.iter().enumerate() { - if r < *k { - recall[ki] += 1; - } - } - } - if hits.len() >= 2 { - margins.push(hits[0].1 - hits[1].1); - } - // The live gate: does top-1 clear 0.75, and is it right? - let (top_entity, top_cos) = (&hits[0].0.entity, hits[0].1); - if top_cos > GATE { - gate_fires += 1; - if *top_entity != entities[i] { - gate_wrong += 1; - } - } - } - - margins.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); - let margin_mean = if margins.is_empty() { - 0.0 - } else { - margins.iter().map(|&x| x as f64).sum::() / margins.len() as f64 - }; - Summary { - recall: [ - recall[0] as f64 / n as f64, - recall[1] as f64 / n as f64, - recall[2] as f64 / n as f64, - recall[3] as f64 / n as f64, - ], - margin_mean, - margin_p50: percentile(&margins, 0.5), - margin_min: percentile(&margins, 0.0), - gate_fires, - gate_wrong, - } -} - -fn main() { - let args: Vec = std::env::args().collect(); - let vindex = args - .get(1) - .cloned() - .unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".to_string()); - let n: usize = args - .get(2) - .and_then(|s| s.parse().ok()) - .unwrap_or(100) - .min(COUNTRIES.len()); - let dir = std::path::PathBuf::from(&vindex); - if !dir.exists() { - eprintln!("skipped: vindex not found at {vindex}"); - eprintln!(" pass a Q4_K gemma3-4b vindex dir as the first arg"); - eprintln!(" (default: output/gemma3-4b-q4k-v2.vindex). Skipping cleanly."); - return; - } - let entities: Vec = COUNTRIES[..n].iter().map(|s| s.to_string()).collect(); - - let mut cb = larql_vindex::SilentLoadCallbacks; - eprintln!("Loading {vindex} ..."); - let mut weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index.load_interleaved_kquant(&dir).expect("interleaved"); - index.load_attn_kquant(&dir).expect("attn kquant"); - let tok = load_tokenizer(&dir).expect("tokenizer"); - eprintln!("Dequantising {} layers to f32 ...", weights.num_layers); - for layer in 0..weights.num_layers { - insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant"); - } - - // Capture residuals for all three phrasings at the layer sweep, one forward - // per (entity, phrasing). `capture_residuals` returns all requested layers - // from a single pass — the layer sweep is nearly free. - let cap = |prompt: &str| -> HashMap> { - let ids = tok.encode(prompt, true).expect("encode").get_ids().to_vec(); - capture_residuals(&weights, &ids, &LAYERS) - .into_iter() - .collect() - }; - - eprintln!("Capturing residuals for {n} entities × 3 phrasings ..."); - let mut train: Vec>> = Vec::with_capacity(n); - let mut para: Vec>> = Vec::with_capacity(n); - let mut cross: Vec>> = Vec::with_capacity(n); - for (i, e) in entities.iter().enumerate() { - train.push(cap(&format!("The capital of {e} is"))); - para.push(cap(&format!("{e}'s capital city is"))); - cross.push(cap(&format!("The currency of {e} is"))); - if (i + 1) % 20 == 0 { - eprintln!(" {}/{n}", i + 1); - } - } - - let chance5 = 5.0 / n as f64; - println!("\n=== FR1: top-k fuzzy entity router on {vindex} (N={n}) ==="); - println!(" cosine-NN production KnnStore; chance@5 = {chance5:.03}; gate = {GATE}\n"); - - // Per-layer JSON records, accumulated as we print. - let mut json_layers = String::new(); - - for &layer in &LAYERS { - let mut store = KnnStore::default(); - for (i, e) in entities.iter().enumerate() { - store.add( - layer, - train[i][&layer].clone(), - 0, - e.clone(), - e.clone(), - "capital".to_string(), - 1.0, - ); - } - let pq: Vec> = (0..n).map(|i| para[i][&layer].clone()).collect(); - let cq: Vec> = (0..n).map(|i| cross[i][&layer].clone()).collect(); - let p = score(&store, layer, &pq, &entities); - let c = score(&store, layer, &cq, &entities); - - println!(" L{layer}:"); - println!( - " PARA recall top1 {:.2} top3 {:.2} top5 {:.2} top10 {:.2} | margin mean {:.3} p50 {:.3} min {:.3}", - p.recall[0], p.recall[1], p.recall[2], p.recall[3], p.margin_mean, p.margin_p50, p.margin_min - ); - println!( - " gate@{GATE} fires {}/{n}, of which WRONG {} ({:.0}% confident-wrong of fired)", - p.gate_fires, - p.gate_wrong, - if p.gate_fires > 0 { 100.0 * p.gate_wrong as f64 / p.gate_fires as f64 } else { 0.0 } - ); - println!( - " CROSS recall top1 {:.2} top3 {:.2} top5 {:.2} top10 {:.2} (entity-key vs answer-leak)", - c.recall[0], c.recall[1], c.recall[2], c.recall[3] - ); - - json_layers.push_str(&format!( - "{}{{\"layer\":{},\"para\":{{\"recall\":[{:.4},{:.4},{:.4},{:.4}],\"margin_mean\":{:.4},\"margin_p50\":{:.4},\"margin_min\":{:.4},\"gate_fires\":{},\"gate_wrong\":{}}},\"cross\":{{\"recall\":[{:.4},{:.4},{:.4},{:.4}]}}}}", - if json_layers.is_empty() { "" } else { "," }, - layer, - p.recall[0], p.recall[1], p.recall[2], p.recall[3], - p.margin_mean, p.margin_p50, p.margin_min, p.gate_fires, p.gate_wrong, - c.recall[0], c.recall[1], c.recall[2], c.recall[3], - )); - } - - // SEE IT — one entity's top-5 candidate list at L26 (a ranked short-list). - { - let layer = 26usize; - let mut store = KnnStore::default(); - for (i, e) in entities.iter().enumerate() { - store.add( - layer, - train[i][&layer].clone(), - 0, - e.clone(), - e.clone(), - "capital".to_string(), - 1.0, - ); - } - // Pick an entity whose paraphrase true-rank is 2..=5 (a short-list, not a pinpoint). - let mut chosen: Option = None; - for i in 0..n { - let hits = store.query_knn(layer, ¶[i][&layer], 10); - if let Some(r) = hits.iter().position(|(e, _)| e.entity == entities[i]) { - if (1..=4).contains(&r) { - chosen = Some(i); - break; - } - } - } - if let Some(i) = chosen.or(Some(0)) { - let hits = store.query_knn(layer, ¶[i][&layer], 5); - let names: Vec<&str> = hits.iter().map(|(e, _)| e.entity.as_str()).collect(); - let rank = store - .query_knn(layer, ¶[i][&layer], n) - .iter() - .position(|(e, _)| e.entity == entities[i]) - .map(|r| r + 1) - .unwrap_or(0); - println!( - "\n SEE IT — \"{}'s capital city is\" @L26 top-5: {names:?} (true rank {rank})", - entities[i] - ); - } - } - - let json = format!( - "{{\"experiment\":\"FR1\",\"vindex\":\"{vindex}\",\"n\":{n},\"gate\":{GATE},\"chance5\":{chance5:.4},\"layers\":[{json_layers}]}}" - ); - let out = "bench/aim-validation/fr1_topk_router_gemma3-4b.json"; - if let Err(e) = std::fs::write(out, &json) { - eprintln!("(could not write {out}: {e})"); - } else { - println!("\nwrote {out}"); - } -} diff --git a/crates/larql-inference/examples/fr2_two_tier_router.rs b/crates/larql-inference/examples/fr2_two_tier_router.rs deleted file mode 100644 index 42f7dc0ce..000000000 --- a/crates/larql-inference/examples/fr2_two_tier_router.rs +++ /dev/null @@ -1,319 +0,0 @@ -//! FR2 — two-tier router: symbolic-primary → activation-fuzzy fallback (the -//! measurement, before any build). Reproduces fleet E16's alias slice on the -//! production path: build a `KnnStore` over CANONICAL country names, then ask -//! whether the activation key recovers historical/alternate names (Persia→Iran, -//! Siam→Thailand, …) that exact-string routing structurally cannot reach. -//! -//! * SYMBOLIC tier — `entries_for_entity` exact match (`knn_store.rs:172`): -//! 1.0 on exact names, 0.0 on aliases (the canonical string is absent from -//! the query). This is the gap. -//! * ACTIVATION fallback — FR1's cosine-NN top-k at the resolved layer: does -//! "The capital of {alias} is" route to the canonical entity? -//! -//! The honest catch (E16): the alias slice is the EASY end (famous aliases); the -//! general fuzzy rate is FR1's ~0.9 top-5, not 1.0. We also report confident-wrong -//! on aliases — mis-routes inject a confident-wrong fact, the cost the verifier -//! (FR1) must bound. -//! -//! Usage: `cargo run --release --example fr2_two_tier_router -- [VINDEX_DIR]` -//! Writes `bench/aim-validation/fr2_two_tier_router_gemma3-4b.json`. - -use larql_inference::vindex::insert_q4k_layer_tensors_resident; -use larql_inference::{capture_residuals, load_tokenizer}; -use larql_vindex::KnnStore; -use std::collections::HashMap; - -const LAYERS: [usize; 2] = [24, 26]; -const GATE: f32 = 0.75; - -/// The store's canonical entities (the country list the model knows). -const CANON: &[&str] = &[ - "France", - "Germany", - "Italy", - "Spain", - "Portugal", - "Greece", - "Austria", - "Switzerland", - "Belgium", - "Netherlands", - "Denmark", - "Norway", - "Sweden", - "Finland", - "Iceland", - "Ireland", - "Poland", - "Hungary", - "Romania", - "Bulgaria", - "Croatia", - "Serbia", - "Ukraine", - "Russia", - "Turkey", - "Japan", - "China", - "India", - "Pakistan", - "Bangladesh", - "Thailand", - "Vietnam", - "Indonesia", - "Malaysia", - "Philippines", - "Singapore", - "Mongolia", - "Nepal", - "Cambodia", - "Laos", - "Brazil", - "Argentina", - "Chile", - "Peru", - "Colombia", - "Venezuela", - "Ecuador", - "Bolivia", - "Mexico", - "Cuba", - "Canada", - "Australia", - "Egypt", - "Morocco", - "Algeria", - "Tunisia", - "Libya", - "Kenya", - "Nigeria", - "Ghana", - "Ethiopia", - "Tanzania", - "Uganda", - "Angola", - "Zambia", - "Zimbabwe", - "Senegal", - "Mali", - "Sudan", - "Cameroon", - "Iran", - "Iraq", - "Israel", - "Jordan", - "Lebanon", - "Syria", - "Yemen", - "Oman", - "Qatar", - "Kuwait", - "Armenia", - "Georgia", - "Azerbaijan", - "Kazakhstan", - "Uzbekistan", - "Afghanistan", - "Sri Lanka", - "South Korea", - "Taiwan", - "Estonia", - "Latvia", - "Lithuania", - "Slovakia", - "Slovenia", - "Luxembourg", - "Malta", - "Cyprus", - "Albania", - "Moldova", - "Belarus", - "Myanmar", - "Botswana", - "Namibia", - "Mozambique", - "Madagascar", - "Congo", - "Liberia", - "Panama", - "Guatemala", - "Guyana", - "Suriname", - "Haiti", - "South Africa", - "United Kingdom", - "United States", -]; - -/// (alias, canonical) — historical / alternate names. The canonical MUST be in -/// CANON. These are the famous (easy) end — see the honest catch in §caveats. -const ALIASES: &[(&str, &str)] = &[ - ("Persia", "Iran"), - ("Siam", "Thailand"), - ("Burma", "Myanmar"), - ("Ceylon", "Sri Lanka"), - ("Holland", "Netherlands"), - ("Britain", "United Kingdom"), - ("Abyssinia", "Ethiopia"), - ("Rhodesia", "Zimbabwe"), - ("Zaire", "Congo"), - ("Formosa", "Taiwan"), -]; - -fn main() { - let args: Vec = std::env::args().collect(); - let vindex = args - .get(1) - .cloned() - .unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".to_string()); - let dir = std::path::PathBuf::from(&vindex); - if !dir.exists() { - eprintln!("skipped: vindex not found at {vindex}"); - eprintln!(" pass a Q4_K gemma3-4b vindex dir as the first arg"); - eprintln!(" (default: output/gemma3-4b-q4k-v2.vindex). Skipping cleanly."); - return; - } - - let mut cb = larql_vindex::SilentLoadCallbacks; - eprintln!("Loading {vindex} ..."); - let mut weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index.load_interleaved_kquant(&dir).expect("interleaved"); - index.load_attn_kquant(&dir).expect("attn kquant"); - let tok = load_tokenizer(&dir).expect("tokenizer"); - eprintln!("Dequantising {} layers ...", weights.num_layers); - for layer in 0..weights.num_layers { - insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant"); - } - - let cap = |prompt: &str| -> HashMap> { - let ids = tok.encode(prompt, true).expect("encode").get_ids().to_vec(); - capture_residuals(&weights, &ids, &LAYERS) - .into_iter() - .collect() - }; - - let n = CANON.len(); - eprintln!( - "Capturing {n} canonical keys + {} alias queries ...", - ALIASES.len() - ); - let canon_res: Vec>> = CANON - .iter() - .map(|e| cap(&format!("The capital of {e} is"))) - .collect(); - let alias_res: Vec>> = ALIASES - .iter() - .map(|(a, _)| cap(&format!("The capital of {a} is"))) - .collect(); - - // ── Symbolic tier: exact-string match of the alias against stored names ── - // Build one store (layer-agnostic for the symbolic test) and probe. - let mut sym_store = KnnStore::default(); - for (i, e) in CANON.iter().enumerate() { - sym_store.add( - LAYERS[0], - canon_res[i][&LAYERS[0]].clone(), - 0, - e.to_string(), - e.to_string(), - "capital".into(), - 1.0, - ); - } - let mut symbolic_hits = 0; - for (alias, canon) in ALIASES { - // entries_for_entity is the production exact-string lookup; an alias - // string is absent, so this finds nothing → symbolic recall 0. - let found = sym_store.entries_for_entity(alias); - let resolved = found.iter().any(|(_, e)| e.entity == *canon); - if resolved { - symbolic_hits += 1; - } - } - - println!( - "\n=== FR2: two-tier router on {vindex} (store N={n}, {} aliases) ===", - ALIASES.len() - ); - println!(" SYMBOLIC exact-match on aliases: {symbolic_hits}/{} resolved (the gap exact-string can't close)\n", ALIASES.len()); - - let mut json_layers = String::new(); - for &layer in &LAYERS { - let mut store = KnnStore::default(); - for (i, e) in CANON.iter().enumerate() { - store.add( - layer, - canon_res[i][&layer].clone(), - 0, - e.to_string(), - e.to_string(), - "capital".into(), - 1.0, - ); - } - let mut top1 = 0; - let mut top5 = 0; - let mut gate_fires = 0; - let mut gate_wrong = 0; - let mut rows = Vec::new(); - for (ai, (alias, canon)) in ALIASES.iter().enumerate() { - let hits = store.query_knn(layer, &alias_res[ai][&layer], 5); - let rank = hits.iter().position(|(e, _)| e.entity == *canon); - let in1 = rank == Some(0); - let in5 = rank.is_some(); - if in1 { - top1 += 1; - } - if in5 { - top5 += 1; - } - let (top_e, top_c) = (&hits[0].0.entity, hits[0].1); - if top_c > GATE { - gate_fires += 1; - if *top_e != *canon { - gate_wrong += 1; - } - } - if layer == 26 { - rows.push(format!( - "{alias}→{canon}: top1={top_e} ({})", - if in1 { - "✓" - } else if in5 { - "in5" - } else { - "MISS" - } - )); - } - } - let na = ALIASES.len(); - println!( - " L{layer}: ACTIVATION fallback top1 {}/{na} ({:.2}) top5 {}/{na} ({:.2}) | gate@{GATE} fires {gate_fires}, wrong {gate_wrong}", - top1, top1 as f64 / na as f64, top5, top5 as f64 / na as f64 - ); - if layer == 26 { - for r in &rows { - println!(" {r}"); - } - } - json_layers.push_str(&format!( - "{}{{\"layer\":{layer},\"alias_top1\":{},\"alias_top5\":{},\"n_alias\":{na},\"gate_fires\":{gate_fires},\"gate_wrong\":{gate_wrong}}}", - if json_layers.is_empty() { "" } else { "," }, - top1, top5 - )); - } - - println!("\n reading: symbolic exact-match resolves {symbolic_hits}/{} aliases; the activation fallback recovers the rest", ALIASES.len()); - println!(" → two-tier (exact primary, activation fallback) reaches what exact-string alone cannot. Confident-wrong (gate_wrong) is the cost a verifier (FR1) must bound."); - let json = format!( - "{{\"experiment\":\"FR2\",\"vindex\":\"{vindex}\",\"store_n\":{n},\"symbolic_alias_hits\":{symbolic_hits},\"layers\":[{json_layers}]}}" - ); - let out = "bench/aim-validation/fr2_two_tier_router_gemma3-4b.json"; - if let Err(e) = std::fs::write(out, &json) { - eprintln!("(could not write {out}: {e})"); - } else { - println!("\nwrote {out}"); - } -} diff --git a/crates/larql-inference/examples/fr3_explicit_rewrite.rs b/crates/larql-inference/examples/fr3_explicit_rewrite.rs deleted file mode 100644 index a2f64dc72..000000000 --- a/crates/larql-inference/examples/fr3_explicit_rewrite.rs +++ /dev/null @@ -1,186 +0,0 @@ -//! FR3 **explicit rewrite** — measure whether the model, asked directly, maps -//! an arbitrary relation phrasing to a canonical relation the vindex knows. -//! -//! The template ablation (`fr3_template_ablation`) showed the residual probe is -//! ~chance on UNSEEN phrasings at its probe layer — diversifying training -//! templates didn't fix it. This tests the alternative (chris's call): instead -//! of a phrasing-invariant probe, do an **explicit** model classification — -//! few-shot "word -> relation" — and read the next-token prediction. One forward -//! pass (no probe training), using the model's own language understanding. -//! -//! Three buckets: known synonyms (seat/money/tongue…), harder UNSEEN phrasings -//! (head city / legal tender / spoken language…) — where the probe failed — and -//! distractors (banana/weather) that should map to NONE of the relations. -//! -//! If explicit classification nails the synonyms AND the unseen phrasings while -//! abstaining on distractors, it's the right resolver fallback: probe-first -//! (cheap, rides the model's implicit normalisation when it works), -//! explicit-rewrite-fallback (robust) — the FR2 two-tier shape, for relations. -//! -//! Usage: `cargo run --release --example fr3_explicit_rewrite -- [VINDEX_DIR]` -//! Writes `bench/aim-validation/fr3_explicit_rewrite_gemma3-4b.json`. - -use larql_inference::load_tokenizer; -use larql_inference::vindex::predict_kquant; - -/// Canonical relations the vindex knows (the classification target set). -const RELATIONS: &[&str] = &["capital", "currency", "language"]; - -/// (phrasing, expected canonical relation, bucket). `""` = should abstain. -const CASES: &[(&str, &str, &str)] = &[ - // known single-word synonyms - ("seat", "capital", "synonym"), - ("metropolis", "capital", "synonym"), - ("money", "currency", "synonym"), - ("cash", "currency", "synonym"), - ("tongue", "language", "synonym"), - ("speech", "language", "synonym"), - // unseen multi-word phrasings (where the residual probe was ~chance) - ("head city", "capital", "phrasing"), - ("main city", "capital", "phrasing"), - ("legal tender", "currency", "phrasing"), - ("unit of money", "currency", "phrasing"), - ("spoken language", "language", "phrasing"), - ("mother tongue", "language", "phrasing"), - // distractors — no relation should be confidently chosen - ("banana", "", "distractor"), - ("weather", "", "distractor"), - ("altitude", "", "distractor"), -]; - -fn main() { - let args: Vec = std::env::args().collect(); - let vindex = args - .get(1) - .cloned() - .unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".to_string()); - let dir = std::path::PathBuf::from(&vindex); - if !dir.exists() { - eprintln!("skipped: vindex not found at {vindex}"); - eprintln!(" pass a Q4_K gemma3-4b vindex dir as the first arg"); - eprintln!(" (default: output/gemma3-4b-q4k-v2.vindex). Skipping cleanly."); - return; - } - - let mut cb = larql_vindex::SilentLoadCallbacks; - eprintln!("Loading {vindex} ..."); - let mut weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index.load_interleaved_kquant(&dir).expect("interleaved"); - index.load_attn_kquant(&dir).expect("attn kquant"); - let tok = load_tokenizer(&dir).expect("tokenizer"); - - // Few-shot frame: examples are NOT in the test set (no leakage), and they - // pin the candidate space + the "word -> relation" task. - // Candidate set includes a `none` escape so out-of-domain words can abstain - // instead of being forced into the nearest relation (the forced-choice - // confident-wrong fix — the same abstain discipline as FR1's verify). - let rel_list = RELATIONS.join(", "); - let prompt_for = |w: &str| -> String { - format!( - "Map each word to one of: {rel_list}, none.\ncity -> capital\ndollar -> currency\ndialect -> language\nmusic -> none\n{w} ->" - ) - }; - // Does the canonical relation appear as a top-k next token (prefix-matched, - // since a relation may tokenise to a leading sub-word)? - let matches = |preds: &[(String, f64)], canonical: &str| -> Option { - preds.iter().position(|(t, _)| { - let t = t.trim().to_lowercase(); - !t.is_empty() && (canonical.starts_with(&t) || t.starts_with(canonical)) - }) - }; - // Any relation chosen as top-1 (for the distractor abstain check)? - let any_rel_top1 = |preds: &[(String, f64)]| -> Option { - let (t, _) = preds.first()?; - let t = t.trim().to_lowercase(); - RELATIONS - .iter() - .find(|r| !t.is_empty() && (r.starts_with(&t) || t.starts_with(**r))) - .map(|r| r.to_string()) - }; - - println!("\n=== FR3 explicit-rewrite classification on {vindex} ==="); - println!(" few-shot \"word -> relation\" over {{{rel_list}}}; one forward, top-5\n"); - println!(" bucket phrasing → top-1 canonical? top-1∈relations"); - - let (mut syn_ok, mut syn_n) = (0usize, 0usize); - let (mut phr_ok, mut phr_n) = (0usize, 0usize); - let (mut distractor_fires, mut distractor_n) = (0usize, 0usize); - let mut json_rows = String::new(); - - for (w, expected, bucket) in CASES { - let ids = tok - .encode(prompt_for(w).as_str(), true) - .expect("encode") - .get_ids() - .to_vec(); - let preds = predict_kquant(&mut weights, &tok, &ids, 5, &index).predictions; - let top1 = preds - .first() - .map(|(t, _)| t.trim().to_string()) - .unwrap_or_default(); - let rank = if expected.is_empty() { - None - } else { - matches(&preds, expected) - }; - let rel_top1 = any_rel_top1(&preds); - - match *bucket { - "synonym" => { - syn_n += 1; - if rank == Some(0) { - syn_ok += 1; - } - } - "phrasing" => { - phr_n += 1; - if rank == Some(0) { - phr_ok += 1; - } - } - "distractor" => { - distractor_n += 1; - if rel_top1.is_some() { - distractor_fires += 1; - } - } - _ => {} - } - - let hit = match (expected.is_empty(), rank) { - (true, _) => format!( - "(abstain; top-1∈rel: {})", - rel_top1.unwrap_or_else(|| "no".into()) - ), - (false, Some(0)) => "✓ top-1".to_string(), - (false, Some(r)) => format!("rank {}", r + 1), - (false, None) => "✗ absent".to_string(), - }; - println!(" {bucket:<11} {w:<19} → {top1:<12} {hit}"); - json_rows.push_str(&format!( - "{}{{\"w\":\"{w}\",\"bucket\":\"{bucket}\",\"expected\":\"{expected}\",\"top1\":\"{}\",\"rank\":{}}}", - if json_rows.is_empty() { "" } else { "," }, - top1.replace('"', "'"), - rank.map(|r| (r as i64 + 1).to_string()).unwrap_or_else(|| "-1".into()) - )); - } - - println!("\n ── verdict ──"); - println!( - " synonyms top-1: {syn_ok}/{syn_n} unseen phrasings top-1: {phr_ok}/{phr_n} distractor false-fires: {distractor_fires}/{distractor_n}" - ); - println!(" (residual probe was ~0.33 = chance on unseen phrasings at its layer — compare.)"); - println!(" If phrasings ≈ synonyms ≈ high and distractors abstain, wire explicit rewrite as"); - println!(" the resolver fallback (probe-first when confident, else explicit classify)."); - - let json = format!( - "{{\"experiment\":\"fr3_explicit_rewrite\",\"vindex\":\"{vindex}\",\"synonym_top1\":[{syn_ok},{syn_n}],\"phrasing_top1\":[{phr_ok},{phr_n}],\"distractor_fires\":[{distractor_fires},{distractor_n}],\"cases\":[{json_rows}]}}" - ); - let out = "bench/aim-validation/fr3_explicit_rewrite_gemma3-4b.json"; - if let Err(e) = std::fs::write(out, &json) { - eprintln!("warning: could not write {out}: {e}"); - } else { - println!("\nwrote {out}"); - } -} diff --git a/crates/larql-inference/examples/fr3_relation_address.rs b/crates/larql-inference/examples/fr3_relation_address.rs deleted file mode 100644 index d7f81df6e..000000000 --- a/crates/larql-inference/examples/fr3_relation_address.rs +++ /dev/null @@ -1,394 +0,0 @@ -//! FR3 — relation as a clean semantic address (the measurement, before any -//! build). Reproduces the mechanism video's `address.py` on a real LARQL vindex -//! and measures the headline asymmetry in ONE harness: -//! -//! * RELATION = sharp, clean, semantic index — a linear probe trained ONLY on -//! {capital, currency, language} classifies UNSEEN synonyms (seat, money, -//! tongue, …) → the relation is a meaning-keyed address, not a string match. -//! * ENTITY = fuzzy — cosine-NN top-1 over the same residuals (capital-train -//! keys, held-out paraphrase query), the FR1 object, for side-by-side. -//! -//! The contrast across the layer sweep is the point: the relation resolves -//! *early and clean*; the entity resolves *late and fuzzy*. address the relation -//! by index, the entity by top-k + rank. -//! -//! Probe = dependency-free softmax regression (standardised inputs, L2), so the -//! number is the production residual's, not a library's. Judged in accuracy -//! (synonym-generalisation), never mean-cosine. -//! -//! Usage: `cargo run --release --example fr3_relation_address -- [VINDEX_DIR] [N]` -//! Writes `bench/aim-validation/fr3_relation_address_gemma3-4b.json`. - -use larql_inference::ndarray::{Array1, Array2, Axis}; -use larql_inference::vindex::insert_q4k_layer_tensors_resident; -use larql_inference::{capture_residuals, load_tokenizer}; -use larql_vindex::KnnStore; -use std::collections::HashMap; - -const LAYERS: [usize; 5] = [6, 10, 14, 20, 26]; - -/// Base relation words the probe trains on (label = class index). -const BASE: [(&str, usize); 3] = [("capital", 0), ("currency", 1), ("language", 2)]; -/// Held-out synonyms the probe is TESTED on — never seen in training. -const SYN: [(&str, usize); 6] = [ - ("seat", 0), - ("metropolis", 0), - ("money", 1), - ("cash", 1), - ("tongue", 2), - ("speech", 2), -]; - -const COUNTRIES: &[&str] = &[ - "France", - "Germany", - "Italy", - "Spain", - "Portugal", - "Greece", - "Austria", - "Switzerland", - "Belgium", - "Netherlands", - "Denmark", - "Norway", - "Sweden", - "Finland", - "Iceland", - "Ireland", - "Poland", - "Hungary", - "Romania", - "Bulgaria", - "Japan", - "China", - "India", - "Pakistan", - "Thailand", - "Vietnam", - "Indonesia", - "Brazil", - "Argentina", - "Chile", - "Peru", - "Colombia", - "Mexico", - "Canada", - "Australia", - "Egypt", - "Morocco", - "Kenya", - "Nigeria", - "Turkey", - "Iran", - "Israel", - "Russia", - "Ukraine", - "Sweden", - "Finland", -]; - -fn standardize(x: &Array2) -> (Array2, Array1, Array1) { - let (n, h) = x.dim(); - let mut mu = Array1::::zeros(h); - let mut sd = Array1::::zeros(h); - for j in 0..h { - let mut m = 0.0f32; - for i in 0..n { - m += x[[i, j]]; - } - m /= n as f32; - let mut v = 0.0f32; - for i in 0..n { - let d = x[[i, j]] - m; - v += d * d; - } - mu[j] = m; - sd[j] = (v / n as f32).sqrt() + 1e-6; - } - let mut z = x.clone(); - for i in 0..n { - for j in 0..h { - z[[i, j]] = (z[[i, j]] - mu[j]) / sd[j]; - } - } - (z, mu, sd) -} - -fn apply_std(x: &Array2, mu: &Array1, sd: &Array1) -> Array2 { - let (n, h) = x.dim(); - let mut z = x.clone(); - for i in 0..n { - for j in 0..h { - z[[i, j]] = (z[[i, j]] - mu[j]) / sd[j]; - } - } - z -} - -fn softmax_rows(logits: &Array2) -> Array2 { - let (n, c) = logits.dim(); - let mut p = logits.clone(); - for i in 0..n { - let mut mx = f32::NEG_INFINITY; - for j in 0..c { - mx = mx.max(p[[i, j]]); - } - let mut s = 0.0f32; - for j in 0..c { - let e = (p[[i, j]] - mx).exp(); - p[[i, j]] = e; - s += e; - } - for j in 0..c { - p[[i, j]] /= s; - } - } - p -} - -/// Softmax regression by full-batch gradient descent. Returns (W, b). -fn train_probe( - x: &Array2, - y: &[usize], - c: usize, - steps: usize, - lr: f32, - l2: f32, -) -> (Array2, Array1) { - let (n, h) = x.dim(); - let mut w = Array2::::zeros((h, c)); - let mut b = Array1::::zeros(c); - for _ in 0..steps { - let logits = x.dot(&w) + &b; - let probs = softmax_rows(&logits); - let mut d = probs; - for i in 0..n { - d[[i, y[i]]] -= 1.0; - } - d /= n as f32; - let gw = x.t().dot(&d) + &(&w * l2); - let gb = d.sum_axis(Axis(0)); - w = &w - &(&gw * lr); - b = &b - &(&gb * lr); - } - (w, b) -} - -fn predict(x: &Array2, w: &Array2, b: &Array1) -> Vec { - let logits = x.dot(w) + b; - let (n, c) = logits.dim(); - (0..n) - .map(|i| { - let mut best = 0usize; - let mut bv = f32::NEG_INFINITY; - for j in 0..c { - if logits[[i, j]] > bv { - bv = logits[[i, j]]; - best = j; - } - } - best - }) - .collect() -} - -fn main() { - let args: Vec = std::env::args().collect(); - let vindex = args - .get(1) - .cloned() - .unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".to_string()); - let dir = std::path::PathBuf::from(&vindex); - if !dir.exists() { - eprintln!("skipped: vindex not found at {vindex}"); - eprintln!(" pass a Q4_K gemma3-4b vindex dir as the first arg"); - eprintln!(" (default: output/gemma3-4b-q4k-v2.vindex). Skipping cleanly."); - return; - } - // Dedup the country list (it has a couple repeats) and take N. - let mut seen = std::collections::HashSet::new(); - let all: Vec = COUNTRIES - .iter() - .filter(|c| seen.insert(c.to_string())) - .map(|s| s.to_string()) - .collect(); - let n: usize = args - .get(2) - .and_then(|s| s.parse().ok()) - .unwrap_or(40) - .min(all.len()); - let entities = &all[..n]; - - let mut cb = larql_vindex::SilentLoadCallbacks; - eprintln!("Loading {vindex} ..."); - let mut weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index.load_interleaved_kquant(&dir).expect("interleaved"); - index.load_attn_kquant(&dir).expect("attn kquant"); - let tok = load_tokenizer(&dir).expect("tokenizer"); - eprintln!("Dequantising {} layers ...", weights.num_layers); - for layer in 0..weights.num_layers { - insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant"); - } - - let cap = |prompt: &str| -> HashMap> { - let ids = tok.encode(prompt, true).expect("encode").get_ids().to_vec(); - capture_residuals(&weights, &ids, &LAYERS) - .into_iter() - .collect() - }; - - // Capture: base relations (probe train + entity keys), synonyms (probe test), - // and the held-out entity paraphrase (entity-routing query). - let words: Vec<(&str, usize, bool)> = BASE - .iter() - .map(|(w, l)| (*w, *l, true)) - .chain(SYN.iter().map(|(w, l)| (*w, *l, false))) - .collect(); - eprintln!( - "Capturing residuals: {n} entities × {} relation words + paraphrase ...", - words.len() - ); - // residuals[word_idx][entity] : HashMap - let mut res: Vec>>> = Vec::new(); - for (w, _l, _base) in &words { - let mut per_ent = Vec::with_capacity(n); - for e in entities { - per_ent.push(cap(&format!("The {w} of {e} is"))); - } - res.push(per_ent); - eprintln!(" captured '{w}'"); - } - let para: Vec>> = entities - .iter() - .map(|e| cap(&format!("{e}'s capital city is"))) - .collect(); - - println!("\n=== FR3: relation as a clean address on {vindex} (N={n}) ==="); - println!(" relation probe: train {{capital,currency,language}}, test synonyms {{seat,metropolis,money,cash,tongue,speech}}"); - println!( - " entity: cosine-NN top-1 (capital keys, paraphrase query) — chance@1 = {:.03}\n", - 1.0 / n as f64 - ); - - let h = res[0][0][&LAYERS[0]].len(); - let mut json_layers = String::new(); - - for &layer in &LAYERS { - // Build probe train set from BASE words. - let base_words: Vec = (0..words.len()).filter(|&i| words[i].2).collect(); - let syn_words: Vec = (0..words.len()).filter(|&i| !words[i].2).collect(); - let n_train = base_words.len() * n; - let mut xt = Array2::::zeros((n_train, h)); - let mut yt = Vec::with_capacity(n_train); - let mut row = 0; - for &wi in &base_words { - for ent_map in &res[wi] { - let v = &ent_map[&layer]; - for j in 0..h { - xt[[row, j]] = v[j]; - } - yt.push(words[wi].1); - row += 1; - } - } - let (zt, mu, sd) = standardize(&xt); - let (w, b) = train_probe(&zt, &yt, 3, 400, 0.1, 1e-3); - let train_pred = predict(&zt, &w, &b); - let train_acc = - train_pred.iter().zip(&yt).filter(|(p, y)| p == y).count() as f64 / n_train as f64; - - // Synonym generalisation (held-out words). - let n_syn = syn_words.len() * n; - let mut xs = Array2::::zeros((n_syn, h)); - let mut ys = Vec::with_capacity(n_syn); - let mut per_word: HashMap<&str, (usize, usize)> = HashMap::new(); - row = 0; - for &wi in &syn_words { - for ent_map in &res[wi] { - let v = &ent_map[&layer]; - for j in 0..h { - xs[[row, j]] = v[j]; - } - ys.push(words[wi].1); - row += 1; - } - } - let zs = apply_std(&xs, &mu, &sd); - let syn_pred = predict(&zs, &w, &b); - let mut syn_correct = 0; - row = 0; - for &wi in &syn_words { - for _ei in 0..n { - let ok = syn_pred[row] == words[wi].1; - if ok { - syn_correct += 1; - } - let e = per_word.entry(words[wi].0).or_insert((0, 0)); - e.1 += 1; - if ok { - e.0 += 1; - } - row += 1; - } - } - let syn_acc = syn_correct as f64 / n_syn as f64; - - // Entity routing top-1 (the asymmetry comparand): capital-train keys, - // paraphrase query, cosine-NN. - let cap_wi = base_words[0]; // "capital" - let mut store = KnnStore::default(); - for ei in 0..n { - store.add( - layer, - res[cap_wi][ei][&layer].clone(), - 0, - entities[ei].clone(), - entities[ei].clone(), - "capital".into(), - 1.0, - ); - } - let mut ent_top1 = 0; - for ei in 0..n { - let hits = store.query_knn(layer, ¶[ei][&layer], 1); - if hits - .first() - .map(|(h, _)| h.entity == entities[ei]) - .unwrap_or(false) - { - ent_top1 += 1; - } - } - let ent_acc = ent_top1 as f64 / n as f64; - - let pw: Vec = SYN - .iter() - .map(|(w, _)| { - let (c, t) = per_word.get(*w).copied().unwrap_or((0, 1)); - format!("{w} {:.2}", c as f64 / t.max(1) as f64) - }) - .collect(); - println!( - " L{layer:<2}: RELATION train {train_acc:.2} synonym-gen {syn_acc:.2} [{}] | ENTITY top-1 {ent_acc:.2}", - pw.join(" ") - ); - - json_layers.push_str(&format!( - "{}{{\"layer\":{layer},\"relation_train\":{train_acc:.4},\"relation_synonym_gen\":{syn_acc:.4},\"entity_top1\":{ent_acc:.4}}}", - if json_layers.is_empty() { "" } else { "," } - )); - } - - println!("\n reading: relation is a CLEAN index (synonym-gen high, resolves early); entity is FUZZY (top-1 low early, resolves late)."); - let json = format!( - "{{\"experiment\":\"FR3\",\"vindex\":\"{vindex}\",\"n\":{n},\"layers\":[{json_layers}]}}" - ); - let out = "bench/aim-validation/fr3_relation_address_gemma3-4b.json"; - if let Err(e) = std::fs::write(out, &json) { - eprintln!("(could not write {out}: {e})"); - } else { - println!("\nwrote {out}"); - } -} diff --git a/crates/larql-inference/examples/fr3_template_ablation.rs b/crates/larql-inference/examples/fr3_template_ablation.rs deleted file mode 100644 index 488612af4..000000000 --- a/crates/larql-inference/examples/fr3_template_ablation.rs +++ /dev/null @@ -1,300 +0,0 @@ -//! FR3 **template ablation** — does training the relation probe over MORE -//! phrasing templates actually make synonym resolution robust to phrasings it -//! has never seen? Validates the multi-template change to the production FR3 -//! resolver (`larql-lql/src/executor/relation_resolver.rs`). -//! -//! Setup: train a relation probe on BASE relations {capital,currency,language} -//! rendered through the first `k` of the resolver's TRAIN templates; test it on -//! the unseen SYNONYMS {seat,money,tongue} rendered through a **held-out** -//! template that appears in NO training set. Sweep `k ∈ {1,2,4}` and read the -//! synonym-classification accuracy at the resolver's probe layer (depth ≈ 0.3). -//! -//! If accuracy rises with `k`, more templates buy genuine phrasing-invariance -//! (the change is justified). If flat, the single template was already enough -//! (the change is harmless but unnecessary). Either way it's a measured call. -//! -//! Usage: `cargo run --release --example fr3_template_ablation -- [VINDEX_DIR] [N_ENTITIES]` -//! Writes `bench/aim-validation/fr3_template_ablation_gemma3-4b.json`. - -use larql_inference::vindex::insert_q4k_layer_tensors_resident; -use larql_inference::{capture_residuals, load_tokenizer}; -use ndarray::{Array1, Array2, Axis}; -use std::collections::HashMap; - -/// Per-layer last-token residuals for one rendered prompt (layer → residual). -type LayerRes = HashMap>; - -/// Layers swept; the resolver's probe layer for a 34-layer model is L10 (0.3·L). -const LAYERS: [usize; 4] = [6, 10, 14, 20]; -/// Relation classes the probe is trained on (label index per class). -const BASE: [(&str, usize); 3] = [("capital", 0), ("currency", 1), ("language", 2)]; -/// Unseen synonyms the probe is tested on (true class index). -const SYN: [(&str, usize); 3] = [("seat", 0), ("money", 1), ("tongue", 2)]; -/// The resolver's training templates (`{r}` relation, `{e}` entity). -const TRAIN_TEMPLATES: &[&str] = &[ - "The {r} of {e} is", - "{e}'s {r} is", - "The {r} of {e}:", - "What is the {r} of {e}? It is", -]; -/// A phrasing that appears in NO training set — the generalization test. -const HELD_OUT_TEMPLATE: &str = "The {r} for {e} would be"; - -const ENTITIES: &[&str] = &[ - "France", "Japan", "Brazil", "Egypt", "Canada", "India", "Germany", "Kenya", -]; - -fn render(t: &str, r: &str, e: &str) -> String { - t.replace("{r}", r).replace("{e}", e) -} - -fn main() { - let args: Vec = std::env::args().collect(); - let vindex = args - .get(1) - .cloned() - .unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".to_string()); - let n: usize = args - .get(2) - .and_then(|s| s.parse().ok()) - .unwrap_or(6) - .min(ENTITIES.len()); - let dir = std::path::PathBuf::from(&vindex); - if !dir.exists() { - eprintln!("skipped: vindex not found at {vindex}"); - eprintln!(" pass a Q4_K gemma3-4b vindex dir as the first arg"); - eprintln!(" (default: output/gemma3-4b-q4k-v2.vindex). Skipping cleanly."); - return; - } - let entities = &ENTITIES[..n]; - - let mut cb = larql_vindex::SilentLoadCallbacks; - eprintln!("Loading {vindex} ..."); - let mut weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index.load_interleaved_kquant(&dir).expect("interleaved"); - index.load_attn_kquant(&dir).expect("attn kquant"); - let tok = load_tokenizer(&dir).expect("tokenizer"); - eprintln!("Dequantising {} layers ...", weights.num_layers); - for l in 0..weights.num_layers { - insert_q4k_layer_tensors_resident(&mut weights, &index, l).expect("dequant"); - } - - let cap = |prompt: &str| -> LayerRes { - let ids = tok.encode(prompt, true).expect("encode").get_ids().to_vec(); - capture_residuals(&weights, &ids, &LAYERS) - .into_iter() - .collect() - }; - - // Train captures: BASE × entities × TRAIN_TEMPLATES. Indexed [base][ent][tmpl]. - eprintln!( - "Capturing train set: {} base × {n} ent × {} templates ...", - BASE.len(), - TRAIN_TEMPLATES.len() - ); - let mut train: Vec>> = Vec::new(); - for (r, _) in BASE { - let mut per_ent = Vec::new(); - for e in entities { - let mut per_t = Vec::new(); - for t in TRAIN_TEMPLATES { - per_t.push(cap(&render(t, r, e))); - } - per_ent.push(per_t); - } - train.push(per_ent); - } - // Test captures: SYN × entities × HELD_OUT (unseen phrasing). - eprintln!( - "Capturing held-out test set: {} syn × {n} ent × 1 template ...", - SYN.len() - ); - let mut test: Vec> = Vec::new(); - for (r, _) in SYN { - let mut per_ent = Vec::new(); - for e in entities { - per_ent.push(cap(&render(HELD_OUT_TEMPLATE, r, e))); - } - test.push(per_ent); - } - - println!("\n=== FR3 template ablation on {vindex} (N={n} entities) ==="); - println!(" train BASE {{capital,currency,language}} over k templates; test SYN"); - println!(" {{seat,money,tongue}} on a HELD-OUT phrasing \"{HELD_OUT_TEMPLATE}\" (chance = 0.33)\n"); - println!(" layer k=1 k=2 k=4"); - - let h = train[0][0][0][&LAYERS[0]].len(); - let mut json_rows = String::new(); - for &layer in &LAYERS { - let mut accs = [0f64; 3]; - for (ki, &k) in [1usize, 2, 4].iter().enumerate() { - // Train set = first k templates. - let n_train = BASE.len() * entities.len() * k; - let mut x = Array2::::zeros((n_train, h)); - let mut y = Vec::with_capacity(n_train); - let mut row = 0; - for (bi, (_, lbl)) in BASE.iter().enumerate() { - for per_ent in &train[bi] { - for t_map in per_ent.iter().take(k) { - let v = &t_map[&layer]; - for j in 0..h { - x[[row, j]] = v[j]; - } - y.push(*lbl); - row += 1; - } - } - } - let (xz, mu, sd) = standardize(&x); - let (w, b) = train_probe(&xz, &y, BASE.len(), 400, 0.1, 1e-3); - - // Test on held-out-phrasing synonyms. - let n_test = SYN.len() * entities.len(); - let mut xt = Array2::::zeros((n_test, h)); - let mut yt = Vec::with_capacity(n_test); - let mut r2 = 0; - for (si, (_, lbl)) in SYN.iter().enumerate() { - for ent_map in &test[si] { - let v = &ent_map[&layer]; - for j in 0..h { - xt[[r2, j]] = v[j]; - } - yt.push(*lbl); - r2 += 1; - } - } - let xtz = apply_std(&xt, &mu, &sd); - let pred = predict(&xtz, &w, &b); - let correct = pred.iter().zip(&yt).filter(|(p, t)| p == t).count(); - accs[ki] = correct as f64 / n_test as f64; - } - println!( - " L{:<3} {:.2} {:.2} {:.2}", - layer, accs[0], accs[1], accs[2] - ); - json_rows.push_str(&format!( - "{}{{\"layer\":{},\"acc_k1\":{:.4},\"acc_k2\":{:.4},\"acc_k4\":{:.4}}}", - if json_rows.is_empty() { "" } else { "," }, - layer, - accs[0], - accs[1], - accs[2] - )); - } - - println!("\n ── verdict ──"); - println!(" Read the resolver's probe layer (L10, depth 0.3). If k=4 > k=1 there, more"); - println!(" templates buy real phrasing-invariance on UNSEEN phrasings — the change is"); - println!(" justified. If flat/equal, one template already generalised (change is harmless)."); - - let json = format!( - "{{\"experiment\":\"fr3_template_ablation\",\"vindex\":\"{vindex}\",\"n_entities\":{n},\"held_out_template\":\"{HELD_OUT_TEMPLATE}\",\"layers\":[{json_rows}]}}" - ); - let out = "bench/aim-validation/fr3_template_ablation_gemma3-4b.json"; - if let Err(e) = std::fs::write(out, &json) { - eprintln!("warning: could not write {out}: {e}"); - } else { - println!("\nwrote {out}"); - } -} - -// ── probe math (mirrors relation_resolver + fr3_relation_address) ── - -fn standardize(x: &Array2) -> (Array2, Array1, Array1) { - let (n, h) = x.dim(); - let mut mu = Array1::::zeros(h); - let mut sd = Array1::::zeros(h); - for j in 0..h { - let mut m = 0.0f32; - for i in 0..n { - m += x[[i, j]]; - } - m /= n as f32; - let mut v = 0.0f32; - for i in 0..n { - let d = x[[i, j]] - m; - v += d * d; - } - mu[j] = m; - sd[j] = (v / n as f32).sqrt() + 1e-6; - } - (apply_std(x, &mu, &sd), mu, sd) -} - -fn apply_std(x: &Array2, mu: &Array1, sd: &Array1) -> Array2 { - let (n, h) = x.dim(); - let mut z = x.clone(); - for i in 0..n { - for j in 0..h { - z[[i, j]] = (z[[i, j]] - mu[j]) / sd[j]; - } - } - z -} - -fn softmax_rows(logits: &Array2) -> Array2 { - let (n, c) = logits.dim(); - let mut p = logits.clone(); - for i in 0..n { - let mut mx = f32::NEG_INFINITY; - for j in 0..c { - mx = mx.max(p[[i, j]]); - } - let mut s = 0.0f32; - for j in 0..c { - let e = (p[[i, j]] - mx).exp(); - p[[i, j]] = e; - s += e; - } - for j in 0..c { - p[[i, j]] /= s; - } - } - p -} - -fn train_probe( - x: &Array2, - y: &[usize], - c: usize, - steps: usize, - lr: f32, - l2: f32, -) -> (Array2, Array1) { - let (n, h) = x.dim(); - let mut w = Array2::::zeros((h, c)); - let mut b = Array1::::zeros(c); - for _ in 0..steps { - let logits = x.dot(&w) + &b; - let probs = softmax_rows(&logits); - let mut d = probs; - for i in 0..n { - d[[i, y[i]]] -= 1.0; - } - d /= n as f32; - let gw = x.t().dot(&d) + &(&w * l2); - let gb = d.sum_axis(Axis(0)); - w = &w - &(&gw * lr); - b = &b - &(&gb * lr); - } - (w, b) -} - -fn predict(x: &Array2, w: &Array2, b: &Array1) -> Vec { - let logits = x.dot(w) + b; - let (n, c) = logits.dim(); - (0..n) - .map(|i| { - let mut best = 0usize; - let mut bv = f32::NEG_INFINITY; - for j in 0..c { - if logits[[i, j]] > bv { - bv = logits[[i, j]]; - best = j; - } - } - best - }) - .collect() -} diff --git a/crates/larql-inference/examples/fr_early_exit_bench.rs b/crates/larql-inference/examples/fr_early_exit_bench.rs deleted file mode 100644 index 4fd603700..000000000 --- a/crates/larql-inference/examples/fr_early_exit_bench.rs +++ /dev/null @@ -1,290 +0,0 @@ -//! Early-exit **tok/s gate** — the production-path wiring measured end to end. -//! -//! The probe (`fr_early_exit_probe`) showed the verified hit is stable from -//! L24/34; the parity prototype (`fr_early_exit_parity`) proved the early-exit -//! token is byte-identical. This drives the real production wiring -//! (`infer_patched_early_exit`, WalkFfn path) on fact-lookup queries and times -//! it against the full `infer_patched` (FR1 `Verified` mode), confirming the -//! layer skip translates to wall-clock. -//! -//! Kill criterion: if the measured speedup on fired retrievals doesn't -//! materialise (attention/KV/branch overhead eats the skipped layers), the lever -//! is dead despite the clean forward ratio. Parity must also hold (early token == -//! full token) — one mismatch and it's dead. -//! -//! Usage: `cargo run --release --example fr_early_exit_bench -- [VINDEX_DIR] [N] [INSTALL_LAYER]` -//! Writes `bench/aim-validation/fr_early_exit_bench_gemma3-4b.json`. - -use larql_inference::forward::{ - infer_patched, infer_patched_early_exit, KnnRouteMode, KNN_COSINE_THRESHOLD, KNN_VERIFY_TOPK, -}; -use larql_inference::load_tokenizer; -use larql_inference::vindex::insert_q4k_layer_tensors_resident; -use larql_vindex::PatchedVindex; -use std::time::Instant; - -const ENTITIES: &[&str] = &[ - "France", - "Germany", - "Italy", - "Spain", - "Portugal", - "Greece", - "Austria", - "Belgium", - "Netherlands", - "Denmark", - "Norway", - "Sweden", - "Finland", - "Poland", - "Hungary", - "Romania", - "Japan", - "China", - "India", - "Pakistan", - "Thailand", - "Vietnam", - "Indonesia", - "Malaysia", - "Brazil", - "Argentina", - "Chile", - "Peru", - "Colombia", - "Mexico", - "Canada", - "Australia", - "Egypt", - "Morocco", - "Kenya", - "Nigeria", - "Ghana", - "Ethiopia", - "Iran", - "Iraq", - "Jordan", - "Israel", - "Turkey", - "Russia", - "Ukraine", - "Cuba", -]; - -fn main() { - let args: Vec = std::env::args().collect(); - let vindex = args - .get(1) - .cloned() - .unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".to_string()); - let n: usize = args - .get(2) - .and_then(|s| s.parse().ok()) - .unwrap_or(40) - .min(ENTITIES.len()); - let dir = std::path::PathBuf::from(&vindex); - if !dir.exists() { - eprintln!("skipped: vindex not found at {vindex}"); - eprintln!(" pass a Q4_K gemma3-4b vindex dir as the first arg"); - eprintln!(" (default: output/gemma3-4b-q4k-v2.vindex). Skipping cleanly."); - return; - } - - let mut cb = larql_vindex::SilentLoadCallbacks; - eprintln!("Loading {vindex} ..."); - let mut weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index.load_interleaved_kquant(&dir).expect("interleaved"); - index.load_attn_kquant(&dir).expect("attn kquant"); - let tok = load_tokenizer(&dir).expect("tokenizer"); - let num_layers = weights.num_layers; - let last = num_layers - 1; - let install_layer = args - .get(3) - .and_then(|s| s.parse().ok()) - .unwrap_or(24) - .min(last); - eprintln!("Dequantising {num_layers} layers to f32 ..."); - for layer in 0..num_layers { - insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant"); - } - // Wrap as the gate index the WalkFfn routes through (production INFER path). - let patched = PatchedVindex::new(index); - - let installed = (n * 3 / 4).max(1).min(n.saturating_sub(1).max(1)); - let entities: Vec = ENTITIES[..n].iter().map(|s| s.to_string()).collect(); - let enc = |p: &str| tok.encode(p, true).expect("encode").get_ids().to_vec(); - - // ── Install facts at L* keyed in the WalkFfn residual space (exactly how - // INSERT … MODE KNN keys: infer_patched residuals). ── - eprintln!("Installing {installed} facts at L{install_layer} ..."); - let mut store = larql_vindex::KnnStore::default(); - for (i, e) in entities.iter().take(installed).enumerate() { - let ids = enc(&format!("The capital of {e} is")); - let res = infer_patched( - &weights, - &tok, - &patched, - None, - &ids, - 1, - &KnnRouteMode::Legacy, - ) - .residuals; - let key = res - .into_iter() - .find(|(l, _)| *l == install_layer) - .map(|(_, v)| v) - .expect("install-layer residual"); - store.add( - install_layer, - key, - i as u32, - e.clone(), - e.clone(), - "capital".to_string(), - 1.0, - ); - } - - // ── Warm up (page-in, branch predictor) on one query each path. ── - { - let ids = enc("France's capital city is"); - let _ = infer_patched( - &weights, - &tok, - &patched, - Some(&store), - &ids, - 5, - &KnnRouteMode::Verified { - k: KNN_VERIFY_TOPK, - threshold: KNN_COSINE_THRESHOLD, - }, - ); - let _ = infer_patched_early_exit( - &weights, - &tok, - &patched, - Some(&store), - &ids, - 5, - KNN_VERIFY_TOPK, - KNN_COSINE_THRESHOLD, - ); - } - - eprintln!("Timing {n} queries (full vs early-exit) ..."); - let mut parity_ok = 0usize; - let mut parity_total = 0usize; - let mut exits = 0usize; - let mut distractor_exits = 0usize; - let mut full_fired_ns: u128 = 0; - let mut early_fired_ns: u128 = 0; - let mut fired = 0usize; - - for (i, e) in entities.iter().enumerate() { - let prompt = format!("{e}'s capital city is"); - let ids = enc(&prompt); - let is_distractor = i >= installed; - - // FULL — production Verified-mode infer_patched (whole stack + lm_head). - let t0 = Instant::now(); - let full = infer_patched( - &weights, - &tok, - &patched, - Some(&store), - &ids, - 5, - &KnnRouteMode::Verified { - k: KNN_VERIFY_TOPK, - threshold: KNN_COSINE_THRESHOLD, - }, - ); - let full_ns = t0.elapsed().as_nanos(); - - // EARLY — short-circuit at L* when the verified hit fires. - let t1 = Instant::now(); - let (early, exited) = infer_patched_early_exit( - &weights, - &tok, - &patched, - Some(&store), - &ids, - 5, - KNN_VERIFY_TOPK, - KNN_COSINE_THRESHOLD, - ); - let early_ns = t1.elapsed().as_nanos(); - - // Parity — the emitted token (position 0) must be identical. - parity_total += 1; - let full_tok = full.predictions.first().map(|(t, _)| t.clone()); - let early_tok = early.predictions.first().map(|(t, _)| t.clone()); - if full_tok == early_tok { - parity_ok += 1; - } - if exited { - exits += 1; - if is_distractor { - distractor_exits += 1; - } - fired += 1; - full_fired_ns += full_ns; - early_fired_ns += early_ns; - } - } - - // ── Report ── - let tail = last - install_layer; - let full_ms = full_fired_ns as f64 / 1e6 / fired.max(1) as f64; - let early_ms = early_fired_ns as f64 / 1e6 / fired.max(1) as f64; - let speedup = if early_fired_ns > 0 { - full_fired_ns as f64 / early_fired_ns as f64 - } else { - 0.0 - }; - println!("\n=== FR early-exit tok/s gate on {vindex} ==="); - println!(" stack {num_layers} layers; resolved L* = {install_layer}; {installed} installed + {} distractor", n - installed); - println!(" production path: WalkFfn infer_patched vs infer_patched_early_exit (Verified, top-k={KNN_VERIFY_TOPK})\n"); - println!(" parity (emitted token early == full): {parity_ok}/{parity_total}"); - println!( - " early-exit fired: {exits}/{n} (skips {tail}/{num_layers} layers + lm_head; distractor false-exits: {distractor_exits})" - ); - if fired > 0 { - println!(" on fired retrievals (n={fired}):"); - println!(" full infer_patched: {full_ms:.1} ms/query"); - println!(" early infer_patched_early: {early_ms:.1} ms/query"); - println!( - " speedup: {speedup:.2}× ({:.0}% faster)", - 100.0 * (1.0 - 1.0 / speedup.max(1e-9)) - ); - } - - let parity = parity_ok == parity_total; - println!("\n ── verdict ──"); - if parity && speedup > 1.05 { - println!( - " WIN: parity holds and early-exit is {speedup:.2}× on fact-lookup answer tokens." - ); - println!(" Worth wiring into the decode loop / server generation path for RAG-style use."); - } else if !parity { - println!(" DEAD: parity broken ({parity_ok}/{parity_total}) — early-exit changes the answer. Stop."); - } else { - println!( - " DEAD (speed): parity holds but speedup {speedup:.2}× ≤ 1.05 — overhead ate the layer skip." - ); - } - - let json = format!( - "{{\"experiment\":\"fr_early_exit_bench\",\"vindex\":\"{vindex}\",\"n\":{n},\"installed\":{installed},\"num_layers\":{num_layers},\"install_layer\":{install_layer},\"parity_ok\":{parity_ok},\"parity_total\":{parity_total},\"exits\":{exits},\"distractor_exits\":{distractor_exits},\"fired\":{fired},\"full_ms\":{full_ms:.4},\"early_ms\":{early_ms:.4},\"speedup\":{speedup:.4}}}" - ); - let out = "bench/aim-validation/fr_early_exit_bench_gemma3-4b.json"; - if let Err(e) = std::fs::write(out, &json) { - eprintln!("warning: could not write {out}: {e}"); - } else { - println!("\nwrote {out}"); - } -} diff --git a/crates/larql-inference/examples/fr_early_exit_decode_projection.rs b/crates/larql-inference/examples/fr_early_exit_decode_projection.rs deleted file mode 100644 index 61c01719e..000000000 --- a/crates/larql-inference/examples/fr_early_exit_decode_projection.rs +++ /dev/null @@ -1,241 +0,0 @@ -//! Decode-loop **measure-first** projection — quantifies the *realizable* -//! gain of terminal-token early-exit in a streaming generation loop, BEFORE -//! committing to the `larql-kv` decode-engine wiring. -//! -//! The KV-cache invariant (incremental decode caches per-layer K/V per -//! position) means early-exit is parity-safe ONLY on the terminal token — skip -//! the tail for a non-terminal token and the next token's attention at those -//! layers loses this position. So for an answer of `T` tokens, at most the last -//! token early-exits; the other `T-1` run the full forward (their KV is needed): -//! -//! blended_speedup(T) = (T · full) / ((T-1) · full + early) -//! -//! and — harsher — the early-exit only fires if the *fact* token is the terminal -//! one. For a natural answer where the fact is mid-sentence ("… is Paris."), the -//! terminal token (".") is not a retrieval, so early-exit fires 0× → 1.0×. -//! -//! This measures `full` (Verified `infer_patched`) and `early` -//! (`infer_patched_early_exit`) per answer-token on the real model, then prints -//! the blended curve so the decode-loop build can be judged on realizable value. -//! -//! Usage: `cargo run --release --example fr_early_exit_decode_projection -- [VINDEX_DIR] [N] [INSTALL_LAYER]` -//! Writes `bench/aim-validation/fr_early_exit_decode_projection_gemma3-4b.json`. - -use larql_inference::forward::{ - infer_patched, infer_patched_early_exit, KnnRouteMode, KNN_COSINE_THRESHOLD, KNN_VERIFY_TOPK, -}; -use larql_inference::load_tokenizer; -use larql_inference::vindex::insert_q4k_layer_tensors_resident; -use larql_vindex::PatchedVindex; -use std::time::Instant; - -const ENTITIES: &[&str] = &[ - "France", - "Germany", - "Italy", - "Spain", - "Portugal", - "Greece", - "Austria", - "Belgium", - "Netherlands", - "Denmark", - "Norway", - "Sweden", - "Finland", - "Poland", - "Hungary", - "Romania", - "Japan", - "China", - "India", - "Pakistan", - "Thailand", - "Vietnam", - "Indonesia", - "Malaysia", -]; - -fn main() { - let args: Vec = std::env::args().collect(); - let vindex = args - .get(1) - .cloned() - .unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".to_string()); - let n: usize = args - .get(2) - .and_then(|s| s.parse().ok()) - .unwrap_or(16) - .min(ENTITIES.len()); - let dir = std::path::PathBuf::from(&vindex); - if !dir.exists() { - eprintln!("skipped: vindex not found at {vindex}"); - eprintln!(" pass a Q4_K gemma3-4b vindex dir as the first arg"); - eprintln!(" (default: output/gemma3-4b-q4k-v2.vindex). Skipping cleanly."); - return; - } - - let mut cb = larql_vindex::SilentLoadCallbacks; - eprintln!("Loading {vindex} ..."); - let mut weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index.load_interleaved_kquant(&dir).expect("interleaved"); - index.load_attn_kquant(&dir).expect("attn kquant"); - let tok = load_tokenizer(&dir).expect("tokenizer"); - let num_layers = weights.num_layers; - let last = num_layers - 1; - let install_layer = args - .get(3) - .and_then(|s| s.parse().ok()) - .unwrap_or(24) - .min(last); - eprintln!("Dequantising {num_layers} layers to f32 ..."); - for layer in 0..num_layers { - insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant"); - } - let patched = PatchedVindex::new(index); - - let installed = (n * 3 / 4).max(1).min(n.saturating_sub(1).max(1)); - let entities: Vec = ENTITIES[..n].iter().map(|s| s.to_string()).collect(); - let enc = |p: &str| tok.encode(p, true).expect("encode").get_ids().to_vec(); - - eprintln!("Installing {installed} facts at L{install_layer} ..."); - let mut store = larql_vindex::KnnStore::default(); - for (i, e) in entities.iter().take(installed).enumerate() { - let ids = enc(&format!("The capital of {e} is")); - let key = infer_patched( - &weights, - &tok, - &patched, - None, - &ids, - 1, - &KnnRouteMode::Legacy, - ) - .residuals - .into_iter() - .find(|(l, _)| *l == install_layer) - .map(|(_, v)| v) - .expect("install residual"); - store.add( - install_layer, - key, - i as u32, - e.clone(), - e.clone(), - "capital".to_string(), - 1.0, - ); - } - - // Warm up. - { - let ids = enc("France's capital city is"); - let _ = infer_patched( - &weights, - &tok, - &patched, - Some(&store), - &ids, - 5, - &KnnRouteMode::Verified { - k: KNN_VERIFY_TOPK, - threshold: KNN_COSINE_THRESHOLD, - }, - ); - let _ = infer_patched_early_exit( - &weights, - &tok, - &patched, - Some(&store), - &ids, - 5, - KNN_VERIFY_TOPK, - KNN_COSINE_THRESHOLD, - ); - } - - eprintln!("Timing full vs early answer-token on {installed} installed facts ..."); - let mut full_ns: u128 = 0; - let mut early_ns: u128 = 0; - let mut fired = 0usize; - for e in entities.iter().take(installed) { - let ids = enc(&format!("{e}'s capital city is")); - let t0 = Instant::now(); - let _ = infer_patched( - &weights, - &tok, - &patched, - Some(&store), - &ids, - 5, - &KnnRouteMode::Verified { - k: KNN_VERIFY_TOPK, - threshold: KNN_COSINE_THRESHOLD, - }, - ); - let f = t0.elapsed().as_nanos(); - let t1 = Instant::now(); - let (_, exited) = infer_patched_early_exit( - &weights, - &tok, - &patched, - Some(&store), - &ids, - 5, - KNN_VERIFY_TOPK, - KNN_COSINE_THRESHOLD, - ); - let ee = t1.elapsed().as_nanos(); - if exited { - full_ns += f; - early_ns += ee; - fired += 1; - } - } - - if fired == 0 { - eprintln!("no early-exit fired — cannot project; aborting."); - return; - } - let full = full_ns as f64 / 1e6 / fired as f64; - let early = early_ns as f64 / 1e6 / fired as f64; - let per_tok = full / early; - - println!("\n=== FR early-exit decode-loop projection on {vindex} ==="); - println!( - " resolved L* = {install_layer}/{num_layers}; measured on {fired} fired retrievals\n" - ); - println!( - " per terminal answer-token: full {full:.1} ms vs early {early:.1} ms → {per_tok:.2}×" - ); - println!("\n blended speedup if the FACT is the terminal token (answer length T):"); - println!(" blended(T) = T·full / ((T-1)·full + early)"); - for t in [1usize, 2, 3, 4, 5, 8, 16] { - let blended = (t as f64 * full) / ((t as f64 - 1.0) * full + early); - let pct = 100.0 * (1.0 - 1.0 / blended); - println!(" T={t:<3} → {blended:.2}× ({pct:.0}% faster)"); - } - println!(" T→∞ → 1.00× (the one terminal token is amortised away)"); - println!("\n if the fact is NOT terminal (natural answer, e.g. \"… is Paris.\"): 1.00× (early-exit never fires)"); - - println!("\n ── verdict ──"); - println!( - " Realizable decode-loop value concentrates at T=1 / max_tokens=1 (answer-token-only" - ); - println!(" generation), which the single-forward `INFER … ROUTE VERIFY EXIT` already serves."); - println!( - " A streaming decode-loop build buys terminal-token early-exit only — worth it ONLY if" - ); - println!(" the target workload is dominated by short, answer-token-terminal generations."); - - let json = format!( - "{{\"experiment\":\"fr_early_exit_decode_projection\",\"vindex\":\"{vindex}\",\"install_layer\":{install_layer},\"num_layers\":{num_layers},\"fired\":{fired},\"full_ms\":{full:.4},\"early_ms\":{early:.4},\"per_token_speedup\":{per_tok:.4}}}" - ); - let out = "bench/aim-validation/fr_early_exit_decode_projection_gemma3-4b.json"; - if let Err(e) = std::fs::write(out, &json) { - eprintln!("warning: could not write {out}: {e}"); - } else { - println!("\nwrote {out}"); - } -} diff --git a/crates/larql-inference/examples/fr_early_exit_parity.rs b/crates/larql-inference/examples/fr_early_exit_parity.rs deleted file mode 100644 index 9393ace10..000000000 --- a/crates/larql-inference/examples/fr_early_exit_parity.rs +++ /dev/null @@ -1,281 +0,0 @@ -//! Early-exit **parity-gated prototype** — the engineering spine before any -//! tok/s number (the probe `fr_early_exit_probe` already showed the verified -//! hit is stable + distractor-safe from ~L24/34). -//! -//! Claim under test: stopping the forward at the stored ("resolved") layer L* -//! and emitting the verified KnnStore target — skipping layers L*+1..end + -//! lm_head — produces the **byte-identical** token the full forward + override -//! would. The kill criterion is parity: one mismatch and the lever is dead. -//! -//! Why it should hold (and what this proves empirically): a decoder is -//! feed-forward, so the residual at L* does **not** depend on layers > L*. -//! `capture_residuals(ids, &[L*])` already stops at `max(capture_layers)` (see -//! `trace.rs`), so it is a genuinely partial forward that captures at the exact -//! FFN-entry point `INSERT … MODE KNN` keys on. Therefore (1) the L* residual -//! from the partial forward is bit-identical to the L* slice of the full forward -//! (proven here per prompt, exactly), and (2) the verified override is a pure -//! function of that residual + store + prompt, so the token is identical. -//! Layers L*+1..end are still computed by the full path, but their output is -//! discarded for a fired retrieval (the override replaces position 0). -//! -//! This prototype proves parity in the `capture_residuals` (WeightFfn) space and -//! times the realised forward saving. Production INFER routes the same override -//! through the WalkFfn residual stream; wiring early-exit there needs a partial -//! walk-forward, but the identical structural argument (residual ⊥ later layers) -//! carries over. -//! -//! Usage: `cargo run --release --example fr_early_exit_parity -- [VINDEX_DIR] [N] [INSTALL_LAYER]` -//! Writes `bench/aim-validation/fr_early_exit_parity_gemma3-4b.json`. - -use larql_inference::forward::{ - apply_knn_override_verified, KNN_COSINE_THRESHOLD, KNN_VERIFY_TOPK, -}; -use larql_inference::vindex::insert_q4k_layer_tensors_resident; -use larql_inference::{capture_residuals, load_tokenizer}; -use larql_vindex::KnnStore; -use std::time::Instant; - -const ENTITIES: &[&str] = &[ - "France", - "Germany", - "Italy", - "Spain", - "Portugal", - "Greece", - "Austria", - "Belgium", - "Netherlands", - "Denmark", - "Norway", - "Sweden", - "Finland", - "Poland", - "Hungary", - "Romania", - "Japan", - "China", - "India", - "Pakistan", - "Thailand", - "Vietnam", - "Indonesia", - "Malaysia", - "Brazil", - "Argentina", - "Chile", - "Peru", - "Colombia", - "Mexico", - "Canada", - "Australia", - "Egypt", - "Morocco", - "Kenya", - "Nigeria", - "Ghana", - "Ethiopia", - "Iran", - "Iraq", - "Jordan", - "Israel", - "Turkey", - "Russia", - "Ukraine", - "Cuba", -]; - -/// Compact (token, layer, cosine-bits) view of an override for exact compare. -fn ovr_key(o: &Option) -> Option<(String, usize, u32)> { - o.as_ref() - .map(|o| (o.token.clone(), o.layer, o.cosine.to_bits())) -} - -fn main() { - let args: Vec = std::env::args().collect(); - let vindex = args - .get(1) - .cloned() - .unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".to_string()); - let n: usize = args - .get(2) - .and_then(|s| s.parse().ok()) - .unwrap_or(40) - .min(ENTITIES.len()); - let dir = std::path::PathBuf::from(&vindex); - if !dir.exists() { - eprintln!("skipped: vindex not found at {vindex}"); - eprintln!(" pass a Q4_K gemma3-4b vindex dir as the first arg"); - eprintln!(" (default: output/gemma3-4b-q4k-v2.vindex). Skipping cleanly."); - return; - } - - let mut cb = larql_vindex::SilentLoadCallbacks; - eprintln!("Loading {vindex} ..."); - let mut weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index.load_interleaved_kquant(&dir).expect("interleaved"); - index.load_attn_kquant(&dir).expect("attn kquant"); - let tok = load_tokenizer(&dir).expect("tokenizer"); - let num_layers = weights.num_layers; - let last = num_layers - 1; - // Resolved layer from the probe (recall steps to 0.90 at L24/34 ≈ 0.7 depth). - let install_layer = args - .get(3) - .and_then(|s| s.parse().ok()) - .unwrap_or(24) - .min(last); - eprintln!("Dequantising {num_layers} layers to f32 ..."); - for layer in 0..num_layers { - insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant"); - } - - let installed = (n * 3 / 4).max(1).min(n.saturating_sub(1).max(1)); - let entities: Vec = ENTITIES[..n].iter().map(|s| s.to_string()).collect(); - let enc = |p: &str| tok.encode(p, true).expect("encode").get_ids().to_vec(); - - // ── Install: key each fact at L* via a PARTIAL forward (0..=L*). ── - eprintln!("Installing {installed} facts at L{install_layer} (partial forward) ..."); - let mut store = KnnStore::default(); - for (i, e) in entities.iter().take(installed).enumerate() { - let ids = enc(&format!("The capital of {e} is")); - let key = capture_residuals(&weights, &ids, &[install_layer]) - .into_iter() - .next() - .map(|(_, v)| v) - .expect("key residual"); - store.add( - install_layer, - key, - i as u32, - e.clone(), - e.clone(), - "capital".to_string(), - 1.0, - ); - } - - // ── Per-query: early (stop at L*) vs full (whole stack), compare. ── - eprintln!( - "Checking parity over {n} queries ({installed} installed, {} distractor) ...", - n - installed - ); - let mut residual_mismatches = 0usize; - let mut token_mismatches = 0usize; - let mut exits = 0usize; // queries where early-exit fired (skips the tail) - let mut distractor_exits = 0usize; - let mut early_ns: u128 = 0; - let mut full_ns: u128 = 0; - - for (i, e) in entities.iter().enumerate() { - let prompt = format!("{e}'s capital city is"); - let ids = enc(&prompt); - let is_distractor = i >= installed; - - // EARLY: partial forward 0..=L*, capture at L*. - let t0 = Instant::now(); - let early_res = capture_residuals(&weights, &ids, &[install_layer]); - early_ns += t0.elapsed().as_nanos(); - let (_, early_ovr) = apply_knn_override_verified( - vec![("\u{2205}".into(), 0.0)], - &early_res, - Some(&store), - 1, - &prompt, - KNN_VERIFY_TOPK, - KNN_COSINE_THRESHOLD, - ); - - // FULL: forward the whole stack; read the L* residual back out. - let t1 = Instant::now(); - let full_all = capture_residuals(&weights, &ids, &[install_layer, last]); - full_ns += t1.elapsed().as_nanos(); - let full_at_star: Vec<(usize, Vec)> = full_all - .iter() - .filter(|(l, _)| *l == install_layer) - .cloned() - .collect(); - let (_, full_ovr) = apply_knn_override_verified( - vec![("\u{2205}".into(), 0.0)], - &full_at_star, - Some(&store), - 1, - &prompt, - KNN_VERIFY_TOPK, - KNN_COSINE_THRESHOLD, - ); - - // Parity 1 — the L* residual is bit-identical (partial vs full forward). - let early_vec = &early_res[0].1; - let full_vec = &full_at_star[0].1; - if early_vec != full_vec { - residual_mismatches += 1; - } - // Parity 2 — the verified override is identical (token + layer + cosine). - if ovr_key(&early_ovr) != ovr_key(&full_ovr) { - token_mismatches += 1; - } - if early_ovr.is_some() { - exits += 1; - if is_distractor { - distractor_exits += 1; - } - } - } - - // ── Report ── - let tail = last - install_layer; // layers skipped on a fired retrieval - let pct_layers = 100.0 * tail as f64 / num_layers as f64; - let fwd_ratio = if full_ns > 0 { - early_ns as f64 / full_ns as f64 - } else { - 0.0 - }; - println!("\n=== FR early-exit parity prototype on {vindex} ==="); - println!(" stack = {num_layers} layers; install/resolved layer L* = {install_layer}"); - println!( - " {installed} installed + {} distractor; verified router top-k={KNN_VERIFY_TOPK}, floor {KNN_COSINE_THRESHOLD}\n", - n - installed - ); - println!(" parity:"); - println!( - " residual @L{install_layer} bit-identical (partial vs full): {}/{n}", - n - residual_mismatches - ); - println!( - " verified override identical (early vs full): {}/{n}", - n - token_mismatches - ); - println!(" behaviour:"); - println!( - " early-exit fired: {exits}/{n} (installed retrievals; distractor false-exits: {distractor_exits})" - ); - println!(" realised saving on a fired retrieval token:"); - println!( - " skip {tail}/{num_layers} layers (~{pct_layers:.0}%) + lm_head; measured forward 0..=L{install_layer} vs 0..=L{last}: {:.0}% of full ({:.1}× faster, lm_head extra)", - fwd_ratio * 100.0, - if fwd_ratio > 0.0 { 1.0 / fwd_ratio } else { 0.0 } - ); - - let parity = residual_mismatches == 0 && token_mismatches == 0; - println!("\n ── verdict ──"); - if parity { - println!(" PARITY HOLDS: every early-exit token is byte-identical to the full forward."); - println!(" Early-exit is correctness-safe — production wiring (partial walk-forward +"); - println!(" decode-loop short-circuit) and a tok/s measurement on a fact-lookup workload"); - println!(" are justified as the next stage."); - } else { - println!( - " PARITY BROKEN: {residual_mismatches} residual + {token_mismatches} token mismatches — early-exit is NOT safe as wired. Dead until explained." - ); - } - - let json = format!( - "{{\"experiment\":\"fr_early_exit_parity\",\"vindex\":\"{vindex}\",\"n\":{n},\"installed\":{installed},\"num_layers\":{num_layers},\"install_layer\":{install_layer},\"residual_mismatches\":{residual_mismatches},\"token_mismatches\":{token_mismatches},\"exits\":{exits},\"distractor_exits\":{distractor_exits},\"layers_skipped\":{tail},\"forward_ratio_partial_over_full\":{fwd_ratio:.4},\"parity\":{parity}}}" - ); - let out = "bench/aim-validation/fr_early_exit_parity_gemma3-4b.json"; - if let Err(e) = std::fs::write(out, &json) { - eprintln!("warning: could not write {out}: {e}"); - } else { - println!("\nwrote {out}"); - } -} diff --git a/crates/larql-inference/examples/fr_early_exit_probe.rs b/crates/larql-inference/examples/fr_early_exit_probe.rs deleted file mode 100644 index ef52081d7..000000000 --- a/crates/larql-inference/examples/fr_early_exit_probe.rs +++ /dev/null @@ -1,309 +0,0 @@ -//! Early-exit probe — stage-1 falsification for **retrieval-augmented early -//! exit** (the speed lever the FR1 verify makes safe). No new kernel. -//! -//! The question: if a *verified* KnnStore hit (FR1: top-k + entity-in-prompt + -//! abstain) is essentially certain to be the answer token, could we stop the -//! forward pass at the layer where that hit fires and skip the rest of the -//! stack + lm_head? That only pays off if the verified hit is **correct and -//! stable from an early-ish layer** — and only if it does **not** wrongly fire -//! on queries about non-stored entities (which must run the full model). -//! -//! Method (faithful to production — keys + queries both from `capture_residuals` -//! at the same layer, routed through the real `apply_knn_override_verified`): -//! -//! ```text -//! INSTALL "The capital of {e} is" -> stored key (target = entity) -//! QUERY "{e}'s capital city is" -> held-out paraphrase (names {e}) -//! ``` -//! -//! Simulate early-exit at every layer L: -//! - recall — installed {e}: does the verified router return {e} at L? -//! - false-fire — distractor {e} (named, NOT installed): does it fire at L? The -//! verify should abstain → ~0; a non-zero rate means early-exit is unsafe there. -//! -//! Kill criterion: if recall only firms up at the last couple of layers, there -//! is nothing to skip — dead. If it is stable (recall high, false-fire ~0) from -//! a mid/late layer L*, early-exit could save (num_layers − L*) layers + lm_head -//! on retrieval tokens, and a parity-gated prototype is justified. -//! -//! Usage: `cargo run --release --example fr_early_exit_probe -- [VINDEX_DIR] [N]` -//! Writes `bench/aim-validation/fr_early_exit_probe_gemma3-4b.json`. - -use larql_inference::forward::{ - apply_knn_override_verified, KNN_COSINE_THRESHOLD, KNN_VERIFY_TOPK, -}; -use larql_inference::vindex::insert_q4k_layer_tensors_resident; -use larql_inference::{capture_residuals, load_tokenizer}; -use larql_vindex::KnnStore; -use std::collections::HashMap; - -/// Country set (the model knows their capitals); first `installed` go in the -/// store, the rest are named-but-unstored distractors. -const ENTITIES: &[&str] = &[ - "France", - "Germany", - "Italy", - "Spain", - "Portugal", - "Greece", - "Austria", - "Belgium", - "Netherlands", - "Denmark", - "Norway", - "Sweden", - "Finland", - "Poland", - "Hungary", - "Romania", - "Japan", - "China", - "India", - "Pakistan", - "Thailand", - "Vietnam", - "Indonesia", - "Malaysia", - "Brazil", - "Argentina", - "Chile", - "Peru", - "Colombia", - "Mexico", - "Canada", - "Australia", - "Egypt", - "Morocco", - "Kenya", - "Nigeria", - "Ghana", - "Ethiopia", - "Iran", - "Iraq", - "Jordan", - "Israel", - "Turkey", - "Russia", - "Ukraine", - "Cuba", -]; - -/// Recall threshold for "the verified hit is reliable at this layer". -const RECALL_OK: f64 = 0.90; -/// Tolerated distractor false-fire rate for "safe to early-exit at this layer". -const FALSE_FIRE_OK: f64 = 0.05; - -struct LayerRow { - layer: usize, - fired: f64, // fraction of installed queries that produced any override - recall: f64, // fraction of installed queries routed to the CORRECT entity - false_fire: f64, // fraction of distractor queries that wrongly fired -} - -fn main() { - let args: Vec = std::env::args().collect(); - let vindex = args - .get(1) - .cloned() - .unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".to_string()); - let n: usize = args - .get(2) - .and_then(|s| s.parse().ok()) - .unwrap_or(40) - .min(ENTITIES.len()); - let dir = std::path::PathBuf::from(&vindex); - if !dir.exists() { - eprintln!("skipped: vindex not found at {vindex}"); - eprintln!(" pass a Q4_K gemma3-4b vindex dir as the first arg"); - eprintln!(" (default: output/gemma3-4b-q4k-v2.vindex). Skipping cleanly."); - return; - } - let installed = (n * 3 / 4).max(1).min(n.saturating_sub(1).max(1)); - let entities: Vec = ENTITIES[..n].iter().map(|s| s.to_string()).collect(); - - let mut cb = larql_vindex::SilentLoadCallbacks; - eprintln!("Loading {vindex} ..."); - let mut weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index.load_interleaved_kquant(&dir).expect("interleaved"); - index.load_attn_kquant(&dir).expect("attn kquant"); - let tok = load_tokenizer(&dir).expect("tokenizer"); - let num_layers = weights.num_layers; - eprintln!("Dequantising {num_layers} layers to f32 ..."); - for layer in 0..num_layers { - insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant"); - } - - // Sweep the whole stack — capture_residuals returns every requested layer - // from a single forward, so the layer sweep is nearly free. - let layers: Vec = (0..num_layers).collect(); - let cap = |prompt: &str| -> HashMap> { - let ids = tok.encode(prompt, true).expect("encode").get_ids().to_vec(); - capture_residuals(&weights, &ids, &layers) - .into_iter() - .collect() - }; - - eprintln!("Capturing residuals for {n} entities × 2 phrasings ..."); - let mut train: Vec>> = Vec::with_capacity(n); - let mut query: Vec>> = Vec::with_capacity(n); - for (i, e) in entities.iter().enumerate() { - train.push(cap(&format!("The capital of {e} is"))); - query.push(cap(&format!("{e}'s capital city is"))); - if (i + 1) % 10 == 0 { - eprintln!(" {}/{n}", i + 1); - } - } - - let n_recall = installed; - let n_distract = n - installed; - println!("\n=== FR early-exit probe on {vindex} (N={n}: {installed} installed, {n_distract} distractor) ==="); - println!( - " verified router: top-k={KNN_VERIFY_TOPK} + entity-in-prompt + abstain; cosine floor {KNN_COSINE_THRESHOLD}" - ); - println!(" recall = correct verified hit on installed paraphrase; false-fire = any hit on a NON-stored entity\n"); - println!(" layer fired recall false-fire"); - - let dummy_raw = || vec![("\u{2205}".to_string(), 0.0f64)]; - let mut rows: Vec = Vec::with_capacity(num_layers); - - for &layer in &layers { - // Store: installed entities' INSTALL-phrasing residual at this layer. - let mut store = KnnStore::default(); - for (i, e) in entities.iter().take(installed).enumerate() { - store.add( - layer, - train[i][&layer].clone(), - i as u32, - e.clone(), // target_token = entity (routing identity) - e.clone(), // entity (what verify matches against the prompt) - "capital".to_string(), - 1.0, - ); - } - - // Recall: installed entities, QUERY phrasing (names the entity). - let mut fired = 0usize; - let mut correct = 0usize; - for (i, e) in entities.iter().take(installed).enumerate() { - let prompt = format!("{e}'s capital city is"); - let res = vec![(layer, query[i][&layer].clone())]; - let (_, ovr) = apply_knn_override_verified( - dummy_raw(), - &res, - Some(&store), - 1, - &prompt, - KNN_VERIFY_TOPK, - KNN_COSINE_THRESHOLD, - ); - if let Some(o) = ovr { - fired += 1; - if o.token == *e { - correct += 1; - } - } - } - - // False-fire: distractor entities (named in prompt, NOT in store). - let mut false_fire = 0usize; - for (i, _e) in entities.iter().enumerate().skip(installed) { - let e = &entities[i]; - let prompt = format!("{e}'s capital city is"); - let res = vec![(layer, query[i][&layer].clone())]; - let (_, ovr) = apply_knn_override_verified( - dummy_raw(), - &res, - Some(&store), - 1, - &prompt, - KNN_VERIFY_TOPK, - KNN_COSINE_THRESHOLD, - ); - if ovr.is_some() { - false_fire += 1; - } - } - - let row = LayerRow { - layer, - fired: fired as f64 / n_recall as f64, - recall: correct as f64 / n_recall as f64, - false_fire: if n_distract == 0 { - 0.0 - } else { - false_fire as f64 / n_distract as f64 - }, - }; - println!( - " L{:<3} {:.2} {:.2} {:.2}", - row.layer, row.fired, row.recall, row.false_fire - ); - rows.push(row); - } - - // ── Verdict: smallest layer L* such that recall ≥ RECALL_OK and false-fire - // ≤ FALSE_FIRE_OK hold for L* through the LAST layer (stable, not a blip). - let last = num_layers - 1; - let stable_from = layers.iter().find(|&&l| { - rows[l..=last] - .iter() - .all(|r| r.recall >= RECALL_OK && r.false_fire <= FALSE_FIRE_OK) - }); - - println!("\n ── verdict ──"); - match stable_from { - Some(&l) if l < last.saturating_sub(2) => { - let saved = num_layers - l; - let pct = 100.0 * saved as f64 / num_layers as f64; - println!( - " VIABLE: verified hit is stable (recall ≥ {RECALL_OK:.2}, false-fire ≤ {FALSE_FIRE_OK:.2}) from L{l} onward." - ); - println!( - " Early-exit at L{l} would skip {saved}/{num_layers} layers (~{pct:.0}% of the stack) + lm_head" - ); - println!(" on retrieval tokens. A parity-gated prototype is justified (next stage)."); - } - Some(&l) => { - println!( - " MARGINAL: only stable from L{l} of {num_layers} — too late to be worth the control-flow complexity." - ); - } - None => { - println!( - " DEAD: no layer gives a stable, distractor-safe verified hit through L{last}. Nothing to skip." - ); - } - } - // Also flag the first layer recall crosses RECALL_OK (the rise point), even - // if it doesn't hold — useful to see rise-vs-hold separately. - if let Some(rise) = rows.iter().find(|r| r.recall >= RECALL_OK) { - println!( - " (recall first crosses {RECALL_OK:.2} at L{}; resolved-layer install target ≈ here)", - rise.layer - ); - } - - // ── JSON sidecar ── - let mut json_rows = String::new(); - for r in &rows { - json_rows.push_str(&format!( - "{}{{\"layer\":{},\"fired\":{:.4},\"recall\":{:.4},\"false_fire\":{:.4}}}", - if json_rows.is_empty() { "" } else { "," }, - r.layer, - r.fired, - r.recall, - r.false_fire - )); - } - let json = format!( - "{{\"experiment\":\"fr_early_exit_probe\",\"vindex\":\"{vindex}\",\"n\":{n},\"installed\":{installed},\"num_layers\":{num_layers},\"recall_ok\":{RECALL_OK},\"false_fire_ok\":{FALSE_FIRE_OK},\"stable_from\":{},\"layers\":[{json_rows}]}}", - stable_from.map(|l| l.to_string()).unwrap_or_else(|| "null".to_string()) - ); - let out = "bench/aim-validation/fr_early_exit_probe_gemma3-4b.json"; - if let Err(e) = std::fs::write(out, &json) { - eprintln!("warning: could not write {out}: {e}"); - } else { - println!("\nwrote {out}"); - } -} diff --git a/crates/larql-inference/examples/fr_routing_gain.rs b/crates/larql-inference/examples/fr_routing_gain.rs deleted file mode 100644 index fa71705c2..000000000 --- a/crates/larql-inference/examples/fr_routing_gain.rs +++ /dev/null @@ -1,263 +0,0 @@ -//! FR routing GAIN — quantify what the FR1/FR2 routers buy end-to-end, and the -//! latency cost. Installs novel facts into a KnnStore, then runs the THREE -//! override modes on the SAME captured forward passes over three query slices: -//! -//! CORRECT prompt about an installed entity → want the installed fact -//! DISTRACTOR prompt about a NON-installed entity → want NO override -//! (model answers itself) -//! ALIAS historical name of an installed entity → want the installed fact -//! -//! This is the gain measurement behind FR1 (`docs/diagnoses/fr1-topk-fuzzy-router.md`) -//! and FR2: legacy top-1+0.75 confident-wrongs on DISTRACTORs; verify fixes that -//! but abstains on ALIASes; two-tier recovers ALIASes at a DISTRACTOR cost. The -//! override is a post-logits sidecar, so we also time it (µs/call) to show the -//! cost is negligible vs the forward. -//! -//! Usage: `cargo run --release --example fr_routing_gain -- [VINDEX_DIR] [LAYER]` - -use larql_inference::forward::{KNN_COSINE_THRESHOLD, KNN_VERIFY_TOPK}; -use larql_inference::vindex::{insert_q4k_layer_tensors_resident, WalkFfn}; -use larql_inference::{ - apply_knn_override, apply_knn_override_two_tier, apply_knn_override_verified, - capture_residuals, load_tokenizer, predict_with_ffn, -}; -use larql_vindex::KnnStore; - -// Installed entities (real countries → NOVEL targets, so a different country's -// prompt that cosine-collides reveals a confident-wrong inject). -const INSTALL: &[&str] = &[ - "Germany", "Spain", "Italy", "Poland", "France", "Iran", "Thailand", "Myanmar", "Ethiopia", - "Zimbabwe", "Japan", "Brazil", "Egypt", "Kenya", "Turkey", "India", "Canada", "Norway", - "Greece", "Portugal", -]; -// Not installed — the model knows their capitals; the right move is NO override. -const DISTRACTOR: &[&str] = &[ - "Austria", - "Belgium", - "Netherlands", - "Sweden", - "Denmark", - "Finland", - "Ireland", - "Switzerland", - "Hungary", - "Romania", - "Ukraine", - "Russia", - "China", - "Pakistan", - "Vietnam", - "Indonesia", - "Mexico", - "Chile", - "Peru", - "Morocco", -]; -// (alias, canonical) — canonical is in INSTALL; want the installed fact. -const ALIAS: &[(&str, &str)] = &[ - ("Persia", "Iran"), - ("Siam", "Thailand"), - ("Burma", "Myanmar"), - ("Abyssinia", "Ethiopia"), - ("Rhodesia", "Zimbabwe"), -]; - -fn target_of(entity: &str) -> String { - format!("{entity}X") -} - -type Preds = Vec<(String, f64)>; -type Residuals = Vec<(usize, Vec)>; - -struct ModeScore { - correct: usize, - distractor_safe: usize, - alias: usize, -} - -fn main() { - let args: Vec = std::env::args().collect(); - let vindex = args - .get(1) - .cloned() - .unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".to_string()); - let layer: usize = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(26); - let dir = std::path::PathBuf::from(&vindex); - if !dir.exists() { - eprintln!("skipped: vindex not found at {vindex}"); - eprintln!(" pass a Q4_K gemma3-4b vindex dir as the first arg"); - eprintln!(" (default: output/gemma3-4b-q4k-v2.vindex). Skipping cleanly."); - return; - } - let topk = 5usize; - - let mut cb = larql_vindex::SilentLoadCallbacks; - eprintln!("Loading {vindex} ..."); - let mut weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index.load_interleaved_kquant(&dir).expect("interleaved"); - index.load_attn_kquant(&dir).expect("attn kquant"); - let _ = index.load_lm_head_kquant(&dir); - let tok = load_tokenizer(&dir).expect("tokenizer"); - eprintln!("Dequantising {} layers ...", weights.num_layers); - for l in 0..weights.num_layers { - insert_q4k_layer_tensors_resident(&mut weights, &index, l).expect("dequant"); - } - - // Install novel facts at the resolved layer. - eprintln!("Installing {} facts at L{layer} ...", INSTALL.len()); - let mut store = KnnStore::default(); - for e in INSTALL { - let prompt = format!("The capital of {e} is"); - let ids = tok - .encode(prompt.as_str(), true) - .expect("enc") - .get_ids() - .to_vec(); - let key = capture_residuals(&weights, &ids, &[layer]) - .into_iter() - .find(|(l, _)| *l == layer) - .map(|(_, v)| v) - .expect("residual"); - let tgt = target_of(e); - let tid = tok - .encode(format!(" {tgt}").as_str(), false) - .expect("enc") - .get_ids()[0]; - store.add(layer, key, tid, tgt, e.to_string(), "capital".into(), 1.0); - } - - // One forward per query → (raw predictions, residual@layer). Reused across - // all three override modes so the only variable is the router. - let forward = |prompt: &str| -> (Preds, Residuals) { - let ids = tok.encode(prompt, true).expect("enc").get_ids().to_vec(); - let walk = WalkFfn::new_unlimited(&weights, &index); - let raw = predict_with_ffn(&weights, &tok, &ids, topk, &walk).predictions; - let res = capture_residuals(&weights, &ids, &[layer]); - (raw, res) - }; - - let mut legacy = ModeScore { - correct: 0, - distractor_safe: 0, - alias: 0, - }; - let mut verified = ModeScore { - correct: 0, - distractor_safe: 0, - alias: 0, - }; - let mut two_tier = ModeScore { - correct: 0, - distractor_safe: 0, - alias: 0, - }; - let (mut t_leg, mut t_ver, mut t_two) = (0u128, 0u128, 0u128); - let mut n_calls = 0u128; - - // Helper: run the three modes on one (raw, res, prompt); return their top-1. - let mut run3 = |raw: &[(String, f64)], res: &[(usize, Vec)], prompt: &str| { - let now = std::time::Instant::now(); - let (lp, _) = apply_knn_override(raw.to_vec(), res, Some(&store), topk); - let dl = now.elapsed().as_nanos(); - let now = std::time::Instant::now(); - let (vp, vo) = apply_knn_override_verified( - raw.to_vec(), - res, - Some(&store), - topk, - prompt, - KNN_VERIFY_TOPK, - KNN_COSINE_THRESHOLD, - ); - let dv = now.elapsed().as_nanos(); - let now = std::time::Instant::now(); - let (tp, to) = apply_knn_override_two_tier( - raw.to_vec(), - res, - Some(&store), - topk, - prompt, - KNN_VERIFY_TOPK, - KNN_COSINE_THRESHOLD, - ); - let dt = now.elapsed().as_nanos(); - // legacy override fired? - let (_, lo) = apply_knn_override(raw.to_vec(), res, Some(&store), topk); - t_leg += dl; - t_ver += dv; - t_two += dt; - n_calls += 1; - ( - lp[0].0.clone(), - lo.is_some(), - vp[0].0.clone(), - vo.is_some(), - tp[0].0.clone(), - to.is_some(), - ) - }; - - // CORRECT — want the installed fact ("{E}X"). - for e in INSTALL { - let (raw, res) = forward(&format!("The capital of {e} is")); - let want = target_of(e); - let (l, _, v, _, t, _) = run3(&raw, &res, &format!("The capital of {e} is")); - legacy.correct += (l == want) as usize; - verified.correct += (v == want) as usize; - two_tier.correct += (t == want) as usize; - } - // DISTRACTOR — want NO override (the override firing = a confident-wrong inject). - for d in DISTRACTOR { - let (raw, res) = forward(&format!("The capital of {d} is")); - let (_, lo, _, vo, _, to) = run3(&raw, &res, &format!("The capital of {d} is")); - legacy.distractor_safe += (!lo) as usize; - verified.distractor_safe += (!vo) as usize; - two_tier.distractor_safe += (!to) as usize; - } - // ALIAS — want the installed canonical's fact ("{Canonical}X"). - for (a, canon) in ALIAS { - let (raw, res) = forward(&format!("The capital of {a} is")); - let want = target_of(canon); - let (l, _, v, _, t, _) = run3(&raw, &res, &format!("The capital of {a} is")); - legacy.alias += (l == want) as usize; - verified.alias += (v == want) as usize; - two_tier.alias += (t == want) as usize; - } - - let (nc, nd, na) = (INSTALL.len(), DISTRACTOR.len(), ALIAS.len()); - let pct = |x: usize, n: usize| 100.0 * x as f64 / n as f64; - println!("\n=== FR routing gain — {vindex} @ L{layer} ==="); - println!(" slices: CORRECT n={nc} (want fact) · DISTRACTOR n={nd} (want NO override) · ALIAS n={na} (want fact)\n"); - println!( - " {:<10}{:>14}{:>16}{:>10}", - "mode", "CORRECT", "DISTRACTOR-safe", "ALIAS" - ); - let row = |name: &str, m: &ModeScore| { - println!( - " {:<10}{:>10}/{nc} {:>3.0}% {:>9}/{nd} {:>3.0}% {:>5}/{na} {:>3.0}%", - name, - m.correct, - pct(m.correct, nc), - m.distractor_safe, - pct(m.distractor_safe, nd), - m.alias, - pct(m.alias, na), - ); - }; - row("legacy", &legacy); - row("verified", &verified); - row("two_tier", &two_tier); - - println!( - "\n override-step latency (µs/call): legacy {:.2} verified {:.2} two_tier {:.2}", - t_leg as f64 / 1000.0 / n_calls as f64, - t_ver as f64 / 1000.0 / n_calls as f64, - t_two as f64 / 1000.0 / n_calls as f64, - ); - println!( - " (the override is a post-logits sidecar — compare to a full decode forward, ~10-40 ms)" - ); - println!("\n reading: legacy confident-wrongs DISTRACTORs; verified fixes them but"); - println!(" abstains on ALIASes; two_tier recovers ALIASes at a DISTRACTOR cost."); -} diff --git a/crates/larql-inference/examples/moe_grid_generate.rs b/crates/larql-inference/examples/moe_grid_generate.rs deleted file mode 100644 index 9b1921b64..000000000 --- a/crates/larql-inference/examples/moe_grid_generate.rs +++ /dev/null @@ -1,141 +0,0 @@ -//! End-to-end demo: generation through a sharded expert grid. -//! -//! Usage (two shards): -//! larql-server --experts 0-63 --port 9191 & -//! larql-server --experts 64-127 --port 9192 & -//! -//! VINDEX=~/chris-models/gemma-4-26B-A4B-it.vindex \ -//! SHARDS="0-63:http://localhost:9191,64-127:http://localhost:9192" \ -//! PROMPT="The capital of France is" \ -//! MAX_TOKENS=8 \ -//! cargo run --release --example moe_grid_generate -//! -//! Single-server shortcut (all experts): -//! SHARDS="0-127:http://localhost:9191" ... - -extern crate blas_src; - -use larql_inference::{ - encode_prompt, layer_graph::grid::generate_with_remote_moe, EosConfig, RemoteMoeBackend, - ShardConfig, -}; -use larql_vindex::{load_vindex_tokenizer, SilentLoadCallbacks, VectorIndex}; -use std::sync::Arc; - -type BoxErr = Box; - -fn main() -> Result<(), BoxErr> { - let vindex_path = std::env::var("VINDEX") - .map(std::path::PathBuf::from) - .unwrap_or_else(|_| { - let home = std::env::var("HOME").unwrap_or_default(); - std::path::PathBuf::from(home).join("chris-models/gemma-4-26B-A4B-it.vindex") - }); - - let shards_spec = - std::env::var("SHARDS").unwrap_or_else(|_| "0-127:http://localhost:9191".into()); - let prompt = std::env::var("PROMPT").unwrap_or_else(|_| "The capital of France is".into()); - let max_tokens: usize = std::env::var("MAX_TOKENS") - .ok() - .and_then(|s| s.parse().ok()) - .unwrap_or(8); - - println!("vindex : {}", vindex_path.display()); - println!("shards : {shards_spec}"); - println!("prompt : \"{prompt}\""); - println!("tokens : {max_tokens}"); - println!(); - - // ── Parse shard spec "START-END:URL,..." ───────────────────────────────── - let shard_configs: Vec = shards_spec - .split(',') - .map(|piece| { - // Find the colon that separates range from URL (URL contains colons too). - let dash = piece.find('-').unwrap_or(0); - let colon = piece[dash..] - .find(':') - .map(|c| c + dash) - .unwrap_or(piece.len()); - let range_str = &piece[..colon]; - let url_str = piece[colon + 1..].to_string(); - let (start, end) = parse_range(range_str); - ShardConfig::new(start, end, url_str) - }) - .collect(); - - println!("Connecting to {} shard(s)…", shard_configs.len()); - let remote = Arc::new(RemoteMoeBackend::connect(shard_configs)?); - println!("Connected.\n"); - - // ── Load vindex + model weights ─────────────────────────────────────────── - print!("Loading vindex… "); - std::io::Write::flush(&mut std::io::stdout()).ok(); - let t0 = std::time::Instant::now(); - let mut cb = SilentLoadCallbacks; - let mut index = VectorIndex::load_vindex(&vindex_path, &mut cb)?; - index.load_attn_kquant(&vindex_path).ok(); - index.load_interleaved_kquant(&vindex_path).ok(); - - let cfg = larql_vindex::load_vindex_config(&vindex_path)?; - let weights = larql_vindex::load_model_weights_kquant(&vindex_path, &mut cb)?; - let tokenizer = load_vindex_tokenizer(&vindex_path)?; - println!( - "done ({:.1}s) model={} layers={} hidden={}", - t0.elapsed().as_secs_f64(), - cfg.model, - cfg.num_layers, - cfg.hidden_size - ); - - // ── Backend (Metal or CPU) ──────────────────────────────────────────────── - #[cfg(all(feature = "gpu", target_os = "macos"))] - let backend = larql_compute_metal::MetalBackend::new().ok_or("Metal not available")?; - #[cfg(not(all(feature = "gpu", target_os = "macos")))] - let backend = larql_compute::CpuBackend; - - // ── Tokenize ────────────────────────────────────────────────────────────── - let arch = &*weights.arch; - let prompt_ids = encode_prompt(&tokenizer, arch, &prompt)?; - println!("Prompt tokens: {}", prompt_ids.len()); - println!(); - - // ── Generate ───────────────────────────────────────────────────────────── - print!("{prompt}"); - std::io::Write::flush(&mut std::io::stdout()).ok(); - - let eos = EosConfig::from_vindex_dir(&vindex_path); - let result = generate_with_remote_moe( - &weights, &tokenizer, prompt_ids, max_tokens, &index, &remote, &backend, &eos, - )?; - - for (tok, ms) in result.tokens.iter().zip(result.decode_ms.iter()) { - print!("{tok}"); - std::io::Write::flush(&mut std::io::stdout()).ok(); - eprintln!(" [{ms:.0}ms]"); - } - // Print remaining tokens that have no latency entry (prefill token). - for tok in result.tokens.iter().skip(result.decode_ms.len()) { - print!("{tok}"); - } - println!(); - println!( - "\n{} tokens avg decode {:.0}ms/tok", - result.tokens.len(), - result.decode_ms.iter().sum::() / result.decode_ms.len().max(1) as f64 - ); - - Ok(()) -} - -fn parse_range(s: &str) -> (usize, usize) { - let parts: Vec<&str> = s.splitn(2, '-').collect(); - let start = parts - .first() - .and_then(|p| p.trim().parse().ok()) - .unwrap_or(0); - let end = parts - .get(1) - .and_then(|p| p.trim().parse().ok()) - .unwrap_or(start); - (start, end) -} diff --git a/crates/larql-inference/examples/predict_from_residual.rs b/crates/larql-inference/examples/predict_from_residual.rs deleted file mode 100644 index f02bb929e..000000000 --- a/crates/larql-inference/examples/predict_from_residual.rs +++ /dev/null @@ -1,315 +0,0 @@ -//! Predict-from-substituted-residual. -//! -//! For each prompt, runs forward up to `start_layer` normally to get the -//! true residual stream + KV cache at that depth. Replaces the last- -//! position residual with a predicted value (from disk), then runs the -//! remaining layers normally and returns the top-1 prediction. -//! -//! This is the runtime for the T2 transition-prediction experiment: do -//! predicted-from-L4 residuals at L20 preserve dense top-1 when fed -//! into the rest of the network? -//! -//! Inputs: -//! --model -//! --vindex (loaded only for tokenizer; FFN runs dense) -//! --prompts-file one prompt per line, matches order of residuals -//! --residuals-bin f32 LE, shape (n_prompts × hidden), last-position -//! predicted residual at start_layer -//! --start-layer substitution depth -//! --out JSON output -//! -//! Output per prompt: -//! dense_top1, dense_pct (full dense forward) -//! substituted_top1, substituted_pct (with predicted L_start substituted) -//! matches_dense (bool) - -use std::path::PathBuf; - -use larql_inference::ffn::WeightFfn; -use larql_inference::forward; -use larql_inference::InferenceModel; -use ndarray::Array2; - -fn value_after(args: &[String], flag: &str) -> Option { - args.iter() - .position(|a| a == flag) - .and_then(|i| args.get(i + 1)) - .cloned() -} - -fn load_prompts(path: &std::path::Path) -> Result, Box> { - let raw = std::fs::read_to_string(path)?; - Ok(raw - .lines() - .map(|s| s.trim()) - .filter(|s| !s.is_empty() && !s.starts_with('#')) - .map(|s| s.to_string()) - .collect()) -} - -type ForwardResult = Result<(Array2, Vec>), Box>; - -fn run_full_forward( - weights: &larql_models::ModelWeights, - tokenizer: &tokenizers::Tokenizer, - token_ids: &[u32], - substitutions: &[(usize, &[f32])], -) -> ForwardResult { - // Returns (final_h, captured_actuals) — one captured residual per - // substitution, recorded just before the substitution writes. - let dense_ffn = WeightFfn { weights }; - let mut h = forward::embed_tokens_pub(weights, token_ids); - let ple_inputs = forward::ple::precompute_per_layer_inputs(weights, &h, token_ids); - let hidden = weights.hidden_size; - let _ = tokenizer; // suppress unused - - let mut captured_actuals: Vec> = vec![Array2::zeros((0, 0)); substitutions.len()]; - - for layer in 0..weights.num_layers { - for (s_idx, (target, predicted_row)) in substitutions.iter().enumerate() { - if *target == layer { - captured_actuals[s_idx] = h.clone(); - let last = h.shape()[0] - 1; - if predicted_row.len() != hidden { - return Err(format!( - "predicted row length {} != hidden {}", - predicted_row.len(), - hidden - ) - .into()); - } - for d in 0..hidden { - h[[last, d]] = predicted_row[d]; - } - } - } - - let (h_post_attn, _) = forward::layer::run_attention_with_kv_cache( - larql_inference::WeightsView::dense(weights), - &h, - layer, - ) - .ok_or_else(|| format!("attention failed at layer {layer}"))?; - let (h_post_ffn, _) = forward::run_ffn(weights, &h_post_attn, layer, &dense_ffn, false); - let mut h_out = forward::ple::apply_per_layer_embedding( - weights, - &h_post_ffn, - layer, - ple_inputs.get(layer), - ); - forward::layer::apply_layer_scalar(weights, &mut h_out, layer); - h = h_out; - } - Ok((h, captured_actuals)) -} - -fn load_residuals_bin( - path: &std::path::Path, - n_prompts: usize, - hidden: usize, -) -> Result>, Box> { - let raw = std::fs::read(path)?; - let expected_bytes = n_prompts * hidden * 4; - if raw.len() != expected_bytes { - return Err(format!( - "{}: size mismatch — got {} bytes expected {} ({} prompts × {} hidden × 4)", - path.display(), - raw.len(), - expected_bytes, - n_prompts, - hidden - ) - .into()); - } - Ok((0..n_prompts) - .map(|p| { - let mut v = Vec::with_capacity(hidden); - for d in 0..hidden { - let off = (p * hidden + d) * 4; - v.push(f32::from_le_bytes(raw[off..off + 4].try_into().unwrap())); - } - v - }) - .collect()) -} - -fn main() -> Result<(), Box> { - let args: Vec = std::env::args().collect(); - let model_path = value_after(&args, "--model").unwrap_or_else(|| "google/gemma-3-4b-it".into()); - let prompts_file = - PathBuf::from(value_after(&args, "--prompts-file").ok_or("--prompts-file required")?); - let out_path = PathBuf::from( - value_after(&args, "--out").unwrap_or_else(|| "/tmp/predict_from_residual.json".into()), - ); - - // Collect substitution layers + bin files. Supports: - // - Single substitution via --start-layer / --residuals-bin (back-compat). - // - Two substitutions via additional --start-layer-b / --residuals-bin-b. - let mut sub_specs: Vec<(usize, PathBuf)> = Vec::new(); - if let Some(s) = value_after(&args, "--start-layer") { - let l: usize = s.parse()?; - let p = PathBuf::from( - value_after(&args, "--residuals-bin") - .ok_or("--residuals-bin required with --start-layer")?, - ); - sub_specs.push((l, p)); - } - if let Some(s) = value_after(&args, "--start-layer-b") { - let l: usize = s.parse()?; - let p = PathBuf::from( - value_after(&args, "--residuals-bin-b") - .ok_or("--residuals-bin-b required with --start-layer-b")?, - ); - sub_specs.push((l, p)); - } - if sub_specs.is_empty() { - return Err( - "at least one substitution (--start-layer + --residuals-bin) is required".into(), - ); - } - - eprintln!("loading model"); - let model = InferenceModel::load(&model_path)?; - let weights = model.weights(); - let tokenizer = model.tokenizer(); - let hidden = weights.hidden_size; - - let prompts = load_prompts(&prompts_file)?; - eprintln!( - "{} prompts; {} substitution(s)", - prompts.len(), - sub_specs.len() - ); - for (l, p) in &sub_specs { - eprintln!(" L{l} ← {}", p.display()); - } - - // Load all substitution bins. - let predicted_per_sub: Vec>> = sub_specs - .iter() - .map(|(_, path)| load_residuals_bin(path, prompts.len(), hidden)) - .collect::>()?; - - let mut results = Vec::new(); - for (i, prompt) in prompts.iter().enumerate() { - eprintln!(" [{}/{}] {prompt:?}", i + 1, prompts.len()); - let encoding = tokenizer - .encode(prompt.as_str(), true) - .map_err(|e| std::io::Error::other(format!("{e}")))?; - let token_ids: Vec = encoding.get_ids().to_vec(); - - // Dense forward (no substitution). - let (h_dense, _) = run_full_forward(weights, tokenizer, &token_ids, &[])?; - let dense_preds = forward::logits_to_predictions_pub(weights, &h_dense, tokenizer, 5, 1.0); - let (dense_tok, dense_p) = dense_preds - .predictions - .first() - .cloned() - .ok_or("no dense top-1")?; - - // Substituted forward — pass all substitution rows for prompt i. - let subs: Vec<(usize, &[f32])> = sub_specs - .iter() - .enumerate() - .map(|(s_idx, (layer, _))| (*layer, predicted_per_sub[s_idx][i].as_slice())) - .collect(); - let (h_sub, actuals_at_layers) = run_full_forward(weights, tokenizer, &token_ids, &subs)?; - let sub_preds = forward::logits_to_predictions_pub(weights, &h_sub, tokenizer, 5, 1.0); - let (sub_tok, sub_p) = sub_preds - .predictions - .first() - .cloned() - .ok_or("no substituted top-1")?; - - // Per-substitution cosine. - let cosines: Vec = sub_specs - .iter() - .enumerate() - .map(|(s_idx, _)| { - let actual = &actuals_at_layers[s_idx]; - if actual.shape() == [0, 0] { - return 0.0_f32; - } - let last = actual.shape()[0] - 1; - let row = actual.row(last); - let pred = ndarray::ArrayView1::from(&predicted_per_sub[s_idx][i]); - let dot: f32 = row.iter().zip(pred.iter()).map(|(a, b)| a * b).sum(); - let na: f32 = row.iter().map(|v| v * v).sum::().sqrt(); - let np: f32 = pred.iter().map(|v| v * v).sum::().sqrt(); - if na > 0.0 && np > 0.0 { - dot / (na * np) - } else { - 0.0 - } - }) - .collect(); - - let matches = dense_tok == sub_tok; - let cos_strs: Vec = sub_specs - .iter() - .zip(cosines.iter()) - .map(|((l, _), c)| format!("L{l}={c:.4}")) - .collect(); - eprintln!( - " dense={dense_tok:?} ({:.2}%) substituted={sub_tok:?} ({:.2}%) cos[{}] match={matches}", - dense_p * 100.0, - sub_p * 100.0, - cos_strs.join(", ") - ); - - let sub_layers: Vec = sub_specs.iter().map(|(l, _)| *l).collect(); - results.push(serde_json::json!({ - "prompt": prompt, - "substitution_layers": sub_layers, - "dense_top1": dense_tok, - "dense_pct": dense_p * 100.0, - "substituted_top1": sub_tok, - "substituted_pct": sub_p * 100.0, - "cosines_predicted_vs_actual": cosines, - "matches_dense": matches, - })); - } - - let sub_layers: Vec = sub_specs.iter().map(|(l, _)| *l).collect(); - let sub_paths: Vec = sub_specs - .iter() - .map(|(_, p)| p.display().to_string()) - .collect(); - let out = serde_json::json!({ - "model": model_path, - "substitution_layers": sub_layers, - "substitution_bins": sub_paths, - "prompts_file": prompts_file.display().to_string(), - "results": results, - }); - std::fs::write(&out_path, serde_json::to_string_pretty(&out)? + "\n")?; - eprintln!("\nwrote {}", out_path.display()); - - let total = results.len(); - let matched = results - .iter() - .filter(|r| r["matches_dense"].as_bool().unwrap_or(false)) - .count(); - println!("\n=== Summary ==="); - let layer_list: Vec = sub_specs.iter().map(|(l, _)| *l).collect(); - println!("substitution layers: {layer_list:?}"); - println!( - "matches_dense: {matched}/{total} ({:.1}%)", - matched as f64 / total.max(1) as f64 * 100.0 - ); - for (s_idx, (l, _)) in sub_specs.iter().enumerate() { - let mean_cos: f64 = results - .iter() - .filter_map(|r| { - r["cosines_predicted_vs_actual"] - .as_array() - .and_then(|a| a.get(s_idx)) - .and_then(|v| v.as_f64()) - }) - .sum::() - / total.max(1) as f64; - println!(" L{l}: mean cosine = {mean_cos:.4}"); - } - - Ok(()) -} diff --git a/crates/larql-inference/examples/probe_contribution_distribution.rs b/crates/larql-inference/examples/probe_contribution_distribution.rs deleted file mode 100644 index 6ca379471..000000000 --- a/crates/larql-inference/examples/probe_contribution_distribution.rs +++ /dev/null @@ -1,513 +0,0 @@ -//! Per-layer FFN contribution distribution probe. -//! -//! For each prompt × each layer, captures the per-feature contribution -//! magnitude proxy `|silu(gate) × up_score × ‖down_row‖|` at the last -//! position, then reports the concentration shape of that distribution. -//! -//! The question this probes: does the FFN at a given layer have a -//! sparse contribution structure (a few features dominate) or a flat -//! one (many features contribute roughly equally)? If sparse, low K -//! preserves the FFN output; if flat, low K can't approximate. -//! -//! Reported metrics per (prompt, layer): -//! - total_contribution: Σ |c_i| over all features -//! - cumfrac_at_k: cumulative |c| / total, sorted descending, at -//! K ∈ {50, 100, 200, 400, 800, 1600, 3200, 6400} -//! - top1_over_mean: |c_max| / mean(|c|) — peakedness -//! - entropy_rank: exp(H(p)) where p_i = |c_i|/Σ|c| — effective rank -//! - gini: 1 - 2·area under Lorenz curve, classic concentration -//! -//! Run: -//! cargo run --release -p larql-inference --example probe_contribution_distribution -- \ -//! --model google/gemma-3-4b-it \ -//! --vindex output/gemma3-4b-q4k-v2.vindex \ -//! --prompt "The capital of France is" \ -//! --prompt "The chemical symbol for gold is" \ -//! --out /tmp/contribution_dist.json - -use std::path::PathBuf; - -use larql_inference::forward; -use larql_inference::vindex::{WalkFfn, WalkFfnConfig}; -use larql_inference::InferenceModel; -use larql_vindex::{SilentLoadCallbacks, VectorIndex}; -use ndarray::Array2; - -const KS: &[usize] = &[50, 100, 200, 500, 1000]; - -fn value_after(args: &[String], flag: &str) -> Option { - args.iter() - .position(|a| a == flag) - .and_then(|i| args.get(i + 1)) - .cloned() -} - -fn values_after(args: &[String], flag: &str) -> Vec { - let mut out = Vec::new(); - for (i, a) in args.iter().enumerate() { - if a == flag { - if let Some(v) = args.get(i + 1) { - out.push(v.clone()); - } - } - } - out -} - -fn load_corpus_json(path: &std::path::Path) -> Result, Box> { - let raw = std::fs::read_to_string(path)?; - let value: serde_json::Value = serde_json::from_str(&raw)?; - let arr = value - .as_array() - .ok_or("corpus json: expected top-level array")?; - let mut out = Vec::with_capacity(arr.len()); - for entry in arr { - if let Some(s) = entry.as_str() { - out.push(s.to_string()); - } else if let Some(s) = entry.get("prompt").and_then(|v| v.as_str()) { - out.push(s.to_string()); - } - } - Ok(out) -} - -fn load_corpus_text(path: &std::path::Path) -> Result, Box> { - let raw = std::fs::read_to_string(path)?; - Ok(raw - .lines() - .map(|s| s.trim()) - .filter(|s| !s.is_empty() && !s.starts_with('#')) - .map(|s| s.to_string()) - .collect()) -} - -fn parse_layer_filter(s: &str) -> Result, Box> { - s.split(',') - .map(|t| t.trim()) - .filter(|t| !t.is_empty()) - .map(|t| t.parse::().map_err(|e| e.into())) - .collect() -} - -fn pre_ffn_norm( - weights: &larql_models::ModelWeights, - h_post_attn: &Array2, - layer: usize, -) -> Array2 { - let arch = &*weights.arch; - let norm_offset = arch.norm_weight_offset(); - let key = if arch.has_post_norms() { - arch.pre_feedforward_layernorm_key(layer) - } else { - Some(arch.post_attention_layernorm_key(layer)) - }; - match key { - Some(k) => forward::apply_norm(weights, h_post_attn, &k, norm_offset), - None => larql_compute::residual::rms_norm_for_arch(h_post_attn, None, norm_offset, arch), - } -} - -#[derive(Default, serde::Serialize)] -struct DistStats { - total: f64, - cumfrac_at_k: std::collections::BTreeMap, - top1_over_mean: f64, - entropy_rank: f64, - gini: f64, - /// Top-K feature indices sorted by |value| descending, for the - /// largest K in `KS`. Lets downstream analysis compute Jaccard - /// overlap of feature sets at any K ≤ max(KS) by prefix. - top_feature_indices: Vec, -} - -#[derive(serde::Serialize)] -struct LayerStats { - layer: usize, - num_features: usize, - /// Distribution of `|silu(gate) × up × ‖down‖|`. Coarse "how much - /// does feature i move the residual" proxy. Same as the v1 probe. - contribution: DistStats, - /// Distribution of `|silu(gate) × up × (down_row · unembed[target])|`. - /// Linear approximation of "how much does feature i push the target - /// logit." Discriminates the Paris/Au asymmetry that contribution - /// magnitude alone cannot. - target_effect: DistStats, -} - -#[derive(serde::Serialize)] -struct PromptStats { - prompt: String, - tokens: usize, - target_token: String, - target_token_id: u32, - layers: Vec, -} - -#[derive(serde::Serialize)] -struct RunResult { - model: String, - vindex: String, - prompts: Vec, -} - -fn analyze(values: &[f32]) -> DistStats { - // Indexed sort by |value| desc so we can also dump top feature indices. - let n = values.len(); - if n == 0 { - return DistStats::default(); - } - let mut indexed: Vec<(usize, f64)> = values - .iter() - .enumerate() - .map(|(i, v)| (i, v.abs() as f64)) - .collect(); - indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); - let abs_v: Vec = indexed.iter().map(|(_, v)| *v).collect(); - - let total: f64 = abs_v.iter().sum(); - if total <= 0.0 { - return DistStats::default(); - } - - // Top-K feature indices for the largest K in KS — downstream picks - // any prefix. - let top_k_max = *KS.iter().max().unwrap_or(&0); - let top_feature_indices: Vec = indexed - .iter() - .take(top_k_max.min(n)) - .map(|(i, _)| *i) - .collect(); - - let mut cumfrac_at_k = std::collections::BTreeMap::new(); - let mut running = 0.0; - let mut next_k = 0; - for (i, v) in abs_v.iter().enumerate() { - running += v; - while next_k < KS.len() && i + 1 == KS[next_k].min(n) { - cumfrac_at_k.insert(KS[next_k], running / total); - next_k += 1; - } - if next_k >= KS.len() { - break; - } - } - while next_k < KS.len() { - cumfrac_at_k.insert(KS[next_k], 1.0); - next_k += 1; - } - - let mean = total / n as f64; - let top1_over_mean = abs_v[0] / mean; - - let mut entropy = 0.0; - for &v in &abs_v { - if v > 0.0 { - let p = v / total; - entropy -= p * p.ln(); - } - } - let entropy_rank = entropy.exp(); - - // Gini: sort ascending, classic formula. - let mut asc = abs_v.clone(); - asc.reverse(); - let n_f = asc.len() as f64; - let weighted: f64 = asc - .iter() - .enumerate() - .map(|(i, v)| (2.0 * (i as f64 + 1.0) - n_f - 1.0) * v) - .sum(); - let gini = weighted / (n_f * total); - - DistStats { - total, - cumfrac_at_k, - top1_over_mean, - entropy_rank, - gini, - top_feature_indices, - } -} - -fn main() -> Result<(), Box> { - let args: Vec = std::env::args().collect(); - let model_path = value_after(&args, "--model").unwrap_or_else(|| "google/gemma-3-4b-it".into()); - let vindex_path = PathBuf::from( - value_after(&args, "--vindex").unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".into()), - ); - let mut prompts: Vec = values_after(&args, "--prompt"); - if let Some(p) = value_after(&args, "--corpus-json") { - let mut v = load_corpus_json(std::path::Path::new(&p))?; - eprintln!("loaded {} prompts from {p}", v.len()); - prompts.append(&mut v); - } - if let Some(p) = value_after(&args, "--corpus-text") { - let mut v = load_corpus_text(std::path::Path::new(&p))?; - eprintln!("loaded {} prompts from {p}", v.len()); - prompts.append(&mut v); - } - if prompts.is_empty() { - prompts = vec![ - "The capital of France is".to_string(), - "The chemical symbol for gold is".to_string(), - ]; - } - if let Some(cap) = value_after(&args, "--max-prompts").and_then(|v| v.parse::().ok()) { - prompts.truncate(cap); - } - let layer_filter: Option> = - match value_after(&args, "--layer-filter") { - Some(s) => Some(parse_layer_filter(&s)?.into_iter().collect()), - None => None, - }; - if let Some(lf) = &layer_filter { - eprintln!("layer filter: {:?}", lf); - } - let out_path = PathBuf::from( - value_after(&args, "--out").unwrap_or_else(|| "/tmp/contribution_dist.json".into()), - ); - - eprintln!("loading model + vindex..."); - let model = InferenceModel::load(&model_path)?; - let weights = model.weights(); - let tokenizer = model.tokenizer(); - - let mut index = VectorIndex::load_vindex(&vindex_path, &mut SilentLoadCallbacks)?; - let _ = index.load_down_features(&vindex_path); - let _ = index.load_up_features(&vindex_path); - index.warmup(); - - // Use a WalkFfn to access the lazy down-norm / up-score machinery. - // No actual sparse walk is run — we only use the helpers. - let cfg = WalkFfnConfig::dense(weights.num_layers); - let walk_ffn = WalkFfn::from_config(weights, &index, cfg); - - let mut all_prompts = Vec::with_capacity(prompts.len()); - for prompt in &prompts { - let token_ids: Vec = tokenizer - .encode(prompt.as_str(), true) - .map_err(|e| std::io::Error::other(format!("{e}")))? - .get_ids() - .to_vec(); - eprintln!("prompt {prompt:?}: tokens={}", token_ids.len()); - - // Pass 1: dense forward to determine the target token. - let dense_ffn = larql_inference::ffn::WeightFfn { weights }; - let target_token_id: u32; - let target_token: String; - { - let mut h = forward::embed_tokens_pub(weights, &token_ids); - let ple_inputs = forward::ple::precompute_per_layer_inputs(weights, &h, &token_ids); - for layer in 0..weights.num_layers { - let (h_post_attn, _) = forward::layer::run_attention_with_kv_cache( - larql_inference::WeightsView::dense(weights), - &h, - layer, - ) - .ok_or_else(|| format!("attention failed at layer {layer}"))?; - let (h_post_ffn, _) = - forward::run_ffn(weights, &h_post_attn, layer, &dense_ffn, false); - let mut h_out = forward::ple::apply_per_layer_embedding( - weights, - &h_post_ffn, - layer, - ple_inputs.get(layer), - ); - forward::layer::apply_layer_scalar(weights, &mut h_out, layer); - h = h_out; - } - let preds = forward::logits_to_predictions_pub(weights, &h, tokenizer, 1, 1.0); - let (tok, _prob) = preds - .predictions - .first() - .cloned() - .ok_or("no top-1 prediction from dense forward")?; - // Re-encode the predicted token to its id. Tokenizer round-trip - // is the simplest way without depending on the prediction - // struct's internal id; `encode(token, false)` returns the ids. - let enc = tokenizer - .encode(tok.as_str(), false) - .map_err(|e| std::io::Error::other(format!("{e}")))?; - let ids = enc.get_ids(); - target_token_id = *ids.first().ok_or("empty token encoding")?; - target_token = tok; - eprintln!(" dense top-1: {target_token:?} (id={target_token_id})"); - } - - // Unembed row for the target — the linearised "what direction in - // residual space pushes the target logit". - if (target_token_id as usize) >= weights.lm_head.shape()[0] { - return Err(format!( - "target id {target_token_id} out of range for lm_head shape {:?}", - weights.lm_head.shape() - ) - .into()); - } - let unembed_row = weights.lm_head.row(target_token_id as usize).to_owned(); - - // Pass 2: dense forward + per-layer distribution analysis. - let mut h = forward::embed_tokens_pub(weights, &token_ids); - let ple_inputs = forward::ple::precompute_per_layer_inputs(weights, &h, &token_ids); - let mut layer_stats = Vec::with_capacity(weights.num_layers); - - for layer in 0..weights.num_layers { - let (h_post_attn, _) = forward::layer::run_attention_with_kv_cache( - larql_inference::WeightsView::dense(weights), - &h, - layer, - ) - .ok_or_else(|| format!("attention failed at layer {layer}"))?; - - let h_ffn = pre_ffn_norm(weights, &h_post_attn, layer); - let last = h_ffn.shape()[0] - 1; - let x_row = h_ffn.row(last).to_owned(); - - let (h_post_ffn, _) = forward::run_ffn(weights, &h_post_attn, layer, &dense_ffn, false); - let mut h_out = forward::ple::apply_per_layer_embedding( - weights, - &h_post_ffn, - layer, - ple_inputs.get(layer), - ); - forward::layer::apply_layer_scalar(weights, &mut h_out, layer); - h = h_out; - - let num_features = index.num_features(layer); - if num_features == 0 { - continue; - } - let should_analyze = layer_filter.as_ref().is_none_or(|lf| lf.contains(&layer)); - if !should_analyze { - continue; - } - - let x_2d = Array2::from_shape_vec((1, weights.hidden_size), x_row.to_vec()).unwrap(); - let gate_scores = match index.gate_scores_batch_backend(layer, &x_2d, None) { - Some(s) => s, - None => { - eprintln!("L{layer}: no gate_scores_batch — skipping"); - continue; - } - }; - let gate_row = gate_scores.row(0); - - let up_scores = match walk_ffn.compute_full_up_scores_pub(layer, &x_row) { - Some(v) => v, - None => { - eprintln!("L{layer}: no up_scores — skipping"); - continue; - } - }; - - let down_norms = match walk_ffn.down_row_norms_pub(layer) { - Some(v) => v, - None => { - eprintln!("L{layer}: no down_norms — skipping"); - continue; - } - }; - - // Need the dequantised down matrix to compute the unembed - // projection per feature. - let down_cache = match index.kquant_ffn_layer(layer, 2) { - Some(c) => c, - None => { - eprintln!("L{layer}: no down cache — skipping target_effect"); - continue; - } - }; - if down_cache.len() < num_features * weights.hidden_size { - eprintln!("L{layer}: down cache too small — skipping"); - continue; - } - - let arch = &*weights.arch; - let use_gelu = arch.activation().uses_gelu_tanh_gate_up(); - - // Per-feature `down_row · unembed[target]` — the linearised - // direct effect of moving 1 unit along this feature's down - // direction on the target logit. - let down_view = ndarray::ArrayView2::from_shape( - (num_features, weights.hidden_size), - down_cache.as_slice(), - )?; - let unembed_proj = down_view.dot(&unembed_row); - - let mut contributions = Vec::with_capacity(num_features); - let mut target_effects = Vec::with_capacity(num_features); - for i in 0..num_features { - let g = gate_row[i]; - let act = if use_gelu { - larql_inference::ffn::gelu_tanh(g) - } else { - g * larql_inference::ffn::sigmoid(g) - }; - let u = up_scores.get(i).copied().unwrap_or(0.0); - let dn = down_norms.get(i).copied().unwrap_or(0.0); - let up_act = act * u; - contributions.push(up_act.abs() * dn); - target_effects.push(up_act * unembed_proj[i]); - } - - let contribution_stats = analyze(&contributions); - let target_effect_stats = analyze(&target_effects); - - layer_stats.push(LayerStats { - layer, - num_features, - contribution: contribution_stats, - target_effect: target_effect_stats, - }); - } - - all_prompts.push(PromptStats { - prompt: prompt.clone(), - tokens: token_ids.len(), - target_token, - target_token_id, - layers: layer_stats, - }); - } - - let result = RunResult { - model: model_path, - vindex: vindex_path.display().to_string(), - prompts: all_prompts, - }; - - let json = serde_json::to_string_pretty(&result)?; - if let Some(parent) = out_path.parent() { - std::fs::create_dir_all(parent)?; - } - std::fs::write(&out_path, json + "\n")?; - println!("wrote {}", out_path.display()); - - // Compact stdout summary: side-by-side contribution vs target_effect. - for ps in &result.prompts { - println!("\n=== {} (target = {:?}) ===", ps.prompt, ps.target_token); - println!(" contribution target_effect"); - println!("layer gini eff_rank cf@200 cf@800 | gini eff_rank cf@200 cf@800"); - for ls in &ps.layers { - let c = &ls.contribution; - let t = &ls.target_effect; - let cf200 = c.cumfrac_at_k.get(&200).copied().unwrap_or(0.0); - let cf800 = c.cumfrac_at_k.get(&800).copied().unwrap_or(0.0); - let tf200 = t.cumfrac_at_k.get(&200).copied().unwrap_or(0.0); - let tf800 = t.cumfrac_at_k.get(&800).copied().unwrap_or(0.0); - println!( - "L{:<3} {:>5.3} {:>7.0} {:>5.3} {:>5.3} | {:>5.3} {:>7.0} {:>5.3} {:>5.3}", - ls.layer, - c.gini, - c.entropy_rank, - cf200, - cf800, - t.gini, - t.entropy_rank, - tf200, - tf800 - ); - } - } - - Ok(()) -} diff --git a/crates/larql-inference/examples/probe_residual_stream.rs b/crates/larql-inference/examples/probe_residual_stream.rs deleted file mode 100644 index 3e8dc3c22..000000000 --- a/crates/larql-inference/examples/probe_residual_stream.rs +++ /dev/null @@ -1,221 +0,0 @@ -//! Per-prompt × per-layer residual capture. -//! -//! For each prompt in the corpus, runs a dense forward pass and dumps -//! the pre-FFN-norm residual at the last position at every layer. The -//! output is a flat f32 binary file of shape -//! `(num_prompts × num_layers × hidden_size)` plus a JSON manifest -//! recording the shape and the prompt list. -//! -//! Downstream: SVD per layer (global intrinsic dim), TwoNN/MLE per -//! layer (local intrinsic dim), k-means clustering, within-cell top-K -//! feature stability — all in Python on the binary file. -//! -//! Run: -//! cargo run --release -p larql-inference --example probe_residual_stream -- \ -//! --model google/gemma-3-4b-it \ -//! --vindex output/gemma3-4b-q4k-v2.vindex \ -//! --corpus-json ../chris-experiments/mechinterp/data/full_factual_subgraph/fse_probe_list.json \ -//! --corpus-text ../chris-experiments/routing/26_fp4_quantisation/prompts_q2_wide.txt \ -//! --max-prompts 1788 \ -//! --out-bin /tmp/residuals.bin \ -//! --out-meta /tmp/residuals_meta.json - -use std::io::{BufWriter, Write}; -use std::path::PathBuf; - -use larql_inference::forward; -use larql_inference::InferenceModel; -use larql_vindex::{SilentLoadCallbacks, VectorIndex}; - -fn value_after(args: &[String], flag: &str) -> Option { - args.iter() - .position(|a| a == flag) - .and_then(|i| args.get(i + 1)) - .cloned() -} - -fn load_corpus_json(path: &PathBuf) -> Result, Box> { - let raw = std::fs::read_to_string(path)?; - let value: serde_json::Value = serde_json::from_str(&raw)?; - let arr = value - .as_array() - .ok_or("corpus json: expected top-level array")?; - let mut out = Vec::with_capacity(arr.len()); - for entry in arr { - if let Some(s) = entry.as_str() { - out.push(s.to_string()); - } else if let Some(s) = entry.get("prompt").and_then(|v| v.as_str()) { - out.push(s.to_string()); - } - } - Ok(out) -} - -fn load_corpus_text(path: &PathBuf) -> Result, Box> { - let raw = std::fs::read_to_string(path)?; - Ok(raw - .lines() - .map(|s| s.trim()) - .filter(|s| !s.is_empty() && !s.starts_with('#')) - .map(|s| s.to_string()) - .collect()) -} - -fn main() -> Result<(), Box> { - let args: Vec = std::env::args().collect(); - let model_path = value_after(&args, "--model").unwrap_or_else(|| "google/gemma-3-4b-it".into()); - let vindex_path = PathBuf::from( - value_after(&args, "--vindex").unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".into()), - ); - let corpus_json = value_after(&args, "--corpus-json").map(PathBuf::from); - let corpus_text = value_after(&args, "--corpus-text").map(PathBuf::from); - let max_prompts: Option = - value_after(&args, "--max-prompts").and_then(|v| v.parse().ok()); - let out_bin = PathBuf::from( - value_after(&args, "--out-bin").unwrap_or_else(|| "/tmp/residuals.bin".into()), - ); - let out_meta = PathBuf::from( - value_after(&args, "--out-meta").unwrap_or_else(|| "/tmp/residuals_meta.json".into()), - ); - - // Assemble corpus. - let mut prompts = Vec::new(); - if let Some(p) = corpus_json { - let mut v = load_corpus_json(&p)?; - eprintln!("loaded {} prompts from {}", v.len(), p.display()); - prompts.append(&mut v); - } - if let Some(p) = corpus_text { - let mut v = load_corpus_text(&p)?; - eprintln!("loaded {} prompts from {}", v.len(), p.display()); - prompts.append(&mut v); - } - if prompts.is_empty() { - return Err("no prompts loaded; pass --corpus-json and/or --corpus-text".into()); - } - if let Some(cap) = max_prompts { - prompts.truncate(cap); - } - eprintln!("total prompts: {}", prompts.len()); - - eprintln!("loading model + vindex..."); - let model = InferenceModel::load(&model_path)?; - let weights = model.weights(); - let tokenizer = model.tokenizer(); - - let mut index = VectorIndex::load_vindex(&vindex_path, &mut SilentLoadCallbacks)?; - let _ = index.load_down_features(&vindex_path); - let _ = index.load_up_features(&vindex_path); - index.warmup(); - - let num_layers = weights.num_layers; - let hidden = weights.hidden_size; - eprintln!( - "model: layers={num_layers} hidden={hidden} → output shape ({}, {num_layers}, {hidden})", - prompts.len() - ); - - // Stream residuals to disk to avoid keeping ~600MB in RAM. - let bin_file = std::fs::File::create(&out_bin)?; - let mut writer = BufWriter::new(bin_file); - - let dense_ffn = larql_inference::ffn::WeightFfn { weights }; - let mut prompt_results: Vec = Vec::with_capacity(prompts.len()); - let total = prompts.len(); - - for (idx, prompt) in prompts.iter().enumerate() { - if idx % 50 == 0 || idx + 1 == total { - eprintln!(" prompt {}/{}", idx + 1, total); - } - - let encoding = match tokenizer.encode(prompt.as_str(), true) { - Ok(e) => e, - Err(e) => { - eprintln!(" skipping (tokenize error): {prompt:?}: {e}"); - // Write zeros so the bin file shape stays consistent. - let zeros = vec![0.0f32; num_layers * hidden]; - for v in &zeros { - writer.write_all(&v.to_le_bytes())?; - } - prompt_results.push(serde_json::json!({ - "prompt": prompt, - "tokens": 0, - "error": "tokenize" - })); - continue; - } - }; - let token_ids: Vec = encoding.get_ids().to_vec(); - if token_ids.is_empty() { - let zeros = vec![0.0f32; num_layers * hidden]; - for v in &zeros { - writer.write_all(&v.to_le_bytes())?; - } - prompt_results.push(serde_json::json!({ - "prompt": prompt, - "tokens": 0, - "error": "empty_tokenization" - })); - continue; - } - - let mut h = forward::embed_tokens_pub(weights, &token_ids); - let ple_inputs = forward::ple::precompute_per_layer_inputs(weights, &h, &token_ids); - - for layer in 0..num_layers { - // Capture h_pre — the residual stream entering layer L's - // attention. This is the layer-to-layer flowing residual, - // the natural substitution point for transition prediction - // experiments. Differs from the prior h_ffn capture which - // recorded post-pre-FFN-norm values inside each layer. - let last = h.shape()[0] - 1; - let row = h.row(last); - for v in row.iter() { - writer.write_all(&v.to_le_bytes())?; - } - - let (h_post_attn, _) = forward::layer::run_attention_with_kv_cache( - larql_inference::WeightsView::dense(weights), - &h, - layer, - ) - .ok_or_else(|| format!("attention failed at layer {layer}"))?; - - let (h_post_ffn, _) = forward::run_ffn(weights, &h_post_attn, layer, &dense_ffn, false); - let mut h_out = forward::ple::apply_per_layer_embedding( - weights, - &h_post_ffn, - layer, - ple_inputs.get(layer), - ); - forward::layer::apply_layer_scalar(weights, &mut h_out, layer); - h = h_out; - } - - prompt_results.push(serde_json::json!({ - "prompt": prompt, - "tokens": token_ids.len(), - })); - } - - writer.flush()?; - drop(writer); - - let meta = serde_json::json!({ - "model": model_path, - "vindex": vindex_path.display().to_string(), - "shape": [prompts.len(), num_layers, hidden], - "dtype": "float32_le", - "residual": "h_pre (residual stream entering each layer's attention), last position only", - "prompts": prompt_results, - }); - std::fs::write(&out_meta, serde_json::to_string_pretty(&meta)? + "\n")?; - eprintln!( - "\nwrote {} ({} bytes)", - out_bin.display(), - prompts.len() * num_layers * hidden * 4 - ); - eprintln!("wrote {}", out_meta.display()); - - Ok(()) -} diff --git a/crates/larql-inference/examples/q4k_remote_parity.rs b/crates/larql-inference/examples/q4k_remote_parity.rs index 45996eb32..3e982cf9a 100644 --- a/crates/larql-inference/examples/q4k_remote_parity.rs +++ b/crates/larql-inference/examples/q4k_remote_parity.rs @@ -165,14 +165,26 @@ fn main() -> Result<(), Box> { remote_index.load_attn_kquant(&vindex_path)?; let t_remote = Instant::now(); - let remote_result = predict_kquant_with_ffn( + // A refusal is a parity *result*, not a crash: the shards declined to + // execute a layer, so there is no remote answer to compare against and + // reporting one would be the exact conflation this channel exists to stop. + let remote_result = match predict_kquant_with_ffn( &mut weights_remote, &tokenizer, &token_ids, top_k, &remote_index, &remote, - ); + ) { + Ok(result) => result, + Err(refusal) => { + eprintln!( + "FAIL — remote route refused ({}): {refusal}", + refusal.kind() + ); + std::process::exit(1); + } + }; let remote_ms = t_remote.elapsed().as_secs_f64() * 1000.0; // ── Compare ── diff --git a/crates/larql-inference/examples/residual_diff.rs b/crates/larql-inference/examples/residual_diff.rs deleted file mode 100644 index 24a37c8a3..000000000 --- a/crates/larql-inference/examples/residual_diff.rs +++ /dev/null @@ -1,409 +0,0 @@ -//! Per-layer residual diff between CPU (`predict_kquant_hidden`) and Metal -//! (`dispatch_full_pipeline`) forward passes. -//! -//! Invariant under test: for the same input prompt, both backends should -//! produce the same `[seq_len, hidden]` residual at the end of every -//! layer. Any drift compounds into the final logits, so the first layer -//! where cosine similarity drops below 1.0 is usually the one to fix. -//! -//! How it works: -//! 1. Triggers both backends on the same prompt with max_tokens=1 -//! (single prefill pass — no KV cache involvement) with the -//! respective per-layer dump env vars set to disjoint temp dirs. -//! 2. Reads the `.f32` dumps each backend emits per layer. -//! CPU: `cpu_layer_{LL}.f32` — LARQL_CPU_DUMP_LAYERS -//! Metal: `metal_layer_{LL}_h_out.f32` — LARQL_METAL_DUMP_LAYERS -//! Both are raw little-endian `f32[seq_len * hidden]` of the -//! end-of-layer residual. -//! 3. Computes cosine similarity + max abs diff per layer, flagging -//! the first layer where cos_sim drops below 0.9999. -//! -//! Usage: -//! cargo run --release --features metal -p larql-inference --example residual_diff -- \ -//! [prompt] -//! -//! Metal prefill dumps only fire on the dense (non-MoE) path — MoE models -//! use `decode_token` which doesn't hook the dump. For MoE, the CPU dump -//! still works; pair it with the existing `LARQL_DUMP_RESIDUALS` for -//! Metal's MoE path (packed format, parsed differently). - -extern crate blas_src; - -use std::path::{Path, PathBuf}; - -use larql_inference::layer_graph::generate::generate; -use larql_inference::layer_graph::CachedLayerGraph; -use larql_inference::wrap_chat_prompt; - -const DRIFT_THRESHOLD: f32 = 0.9999; - -fn main() -> Result<(), Box> { - let mut args = std::env::args().skip(1); - let vindex_path = PathBuf::from( - args.next() - .ok_or("usage: residual_diff [prompt]")?, - ); - let prompt = args - .next() - .unwrap_or_else(|| "The capital of France is".to_string()); - - if !vindex_path.is_dir() { - return Err(format!("not a vindex dir: {}", vindex_path.display()).into()); - } - - // Disjoint scratch dirs for the two backends' dumps. `tempfile` - // auto-cleans on drop; we stash the paths before the guards leave - // scope so the post-run readers see the files. When the env vars are - // set by the caller (for interactive inspection of intermediate - // files), we use those paths directly and skip the TempDir guard so - // the files survive the run. - let external_cpu = std::env::var_os("LARQL_CPU_DUMP_LAYERS").map(std::path::PathBuf::from); - let external_metal = std::env::var_os("LARQL_METAL_DUMP_LAYERS").map(std::path::PathBuf::from); - let _cpu_guard: Option; - let _metal_guard: Option; - let cpu_path: std::path::PathBuf = if let Some(p) = external_cpu { - _cpu_guard = None; - std::fs::create_dir_all(&p).ok(); - p - } else { - let d = tempfile::tempdir()?; - let p = d.path().to_path_buf(); - _cpu_guard = Some(d); - p - }; - let metal_path: std::path::PathBuf = if let Some(p) = external_metal { - _metal_guard = None; - std::fs::create_dir_all(&p).ok(); - p - } else { - let d = tempfile::tempdir()?; - let p = d.path().to_path_buf(); - _metal_guard = Some(d); - p - }; - std::env::set_var("LARQL_CPU_DUMP_LAYERS", &cpu_path); - std::env::set_var("LARQL_METAL_DUMP_LAYERS", &metal_path); - // Stage dumps: Metal writes to LARQL_METAL_DUMP_LAYERS (same dir) with - // `metal_layer_{LL}_.f32` names; CPU writes its stages into a - // shared stage dir via LARQL_CPU_STAGE_DUMP using `cpu_L0_.f32`. - // Place CPU stage files alongside CPU layer files for simpler reading. - std::env::set_var("LARQL_CPU_STAGE_DUMP", &cpu_path); - // Which layer's per-stage snapshots to compare. Override with the env - // var if you want to bisect somewhere other than L0. - let stage_layer: usize = std::env::var("LARQL_STAGE_DUMP_LAYER") - .ok() - .and_then(|s| s.parse().ok()) - .unwrap_or(0); - - // ── Load vindex ──────────────────────────────────────────────────── - let mut cb = larql_vindex::SilentLoadCallbacks; - let cfg = larql_vindex::load_vindex_config(&vindex_path)?; - let mut index = larql_vindex::VectorIndex::load_vindex(&vindex_path, &mut cb)?; - index.load_attn_kquant(&vindex_path)?; - index.load_interleaved_kquant(&vindex_path)?; - let _ = index.load_lm_head_kquant(&vindex_path); - let tokenizer = larql_vindex::load_vindex_tokenizer(&vindex_path)?; - - let mut w_metal = larql_vindex::load_model_weights_kquant(&vindex_path, &mut cb)?; - let mut w_cpu = larql_vindex::load_model_weights_kquant(&vindex_path, &mut cb)?; - - let wrap = wrap_chat_prompt(&vindex_path, Some(cfg.model.as_str()), &prompt); - let token_ids = larql_inference::encode_prompt(&tokenizer, &*w_metal.arch, &wrap.prompt)?; - let num_layers = w_metal.num_layers; - let hidden = w_metal.hidden_size; - let seq_len = token_ids.len(); - - println!("━━━ Per-layer residual diff ─────────────────────────────────────────"); - println!(" vindex: {}", vindex_path.display()); - println!(" model: {}", cfg.model); - println!(" family: {}", cfg.family); - println!(" prompt: {prompt:?}"); - println!( - " seq_len: {seq_len} ({} tokens post-template)", - token_ids.len() - ); - println!(" num_layers: {num_layers}"); - println!(" hidden: {hidden}"); - println!(); - - // ── Drive both backends (max_tokens=1 → just prefill once each) ───── - let metal_backend = - larql_compute_metal::MetalBackend::new().ok_or("Metal backend unavailable")?; - let metal_cached = CachedLayerGraph::from_residuals(Vec::new()); - println!( - "Running Metal prefill (dumps → {})", - metal_path.as_path().display() - ); - let _ = generate( - &mut w_metal, - &tokenizer, - &token_ids, - 1, - &index, - &metal_backend, - &metal_cached, - 0..num_layers, - ); - - let cpu_backend = larql_compute::CpuBackend; - let cpu_cached = CachedLayerGraph::from_residuals(Vec::new()); - println!( - "Running CPU prefill (dumps → {})", - cpu_path.as_path().display() - ); - let _ = generate( - &mut w_cpu, - &tokenizer, - &token_ids, - 1, - &index, - &cpu_backend, - &cpu_cached, - 0..num_layers, - ); - - println!(); - println!("━━━ Layer-by-layer comparison ──────────────────────────────────────"); - println!(" L h_post_attn cos / maxΔ h_out cos / maxΔ attn vs ffn"); - println!(" ─── ───────────────────────── ───────────────────────── ─────────"); - - let mut first_bad: Option = None; - for l in 0..num_layers { - let load = |cpu_name: &str, metal_name: &str| -> Option<(Vec, Vec)> { - let c = read_f32(&cpu_path.as_path().join(cpu_name))?; - let m = read_f32(&metal_path.as_path().join(metal_name))?; - if c.len() != m.len() { - return None; - } - Some((c, m)) - }; - - let hpa = load( - &format!("cpu_layer_{l:02}_h_post_attn.f32"), - &format!("metal_layer_{l:02}_h_post_attn.f32"), - ); - let hout = load( - &format!("cpu_layer_{l:02}.f32"), - &format!("metal_layer_{l:02}_h_out.f32"), - ); - - let Some((cpu_out, mtl_out)) = hout else { - println!(" L{l:02} "); - continue; - }; - let stat_out = layer_stats(&cpu_out, &mtl_out); - let stat_hpa = hpa.as_ref().map(|(c, m)| layer_stats(c, m)); - - if stat_out.cos < DRIFT_THRESHOLD && first_bad.is_none() { - first_bad = Some(l); - } - let flag = if stat_out.cos < DRIFT_THRESHOLD { - " ←" - } else { - "" - }; - - // Diagnostic: which piece (attention vs FFN) introduces the drift. - // If h_post_attn already differs, attention is the culprit; - // otherwise drift is in FFN+PLE+scalar. - let diagnosis = match stat_hpa { - Some(ref s) if s.cos < DRIFT_THRESHOLD && stat_out.cos < DRIFT_THRESHOLD => "attn+ffn", - Some(ref s) if s.cos < DRIFT_THRESHOLD => "attn", - Some(_) if stat_out.cos < DRIFT_THRESHOLD => "ffn", - Some(_) => "clean", - None => "?", - }; - - let hpa_cell = match stat_hpa { - Some(s) => format!("{:>8.6} / {:>8.2e}", s.cos, s.max_abs_diff), - None => " - / -".to_string(), - }; - println!( - " L{l:02} {} {:>8.6} / {:>8.2e} {:>9}{flag}", - hpa_cell, stat_out.cos, stat_out.max_abs_diff, diagnosis, - ); - } - - println!(); - match first_bad { - Some(l) => { - println!( - "━━━ First layer with cos_sim < {} ─────────────────────────", - DRIFT_THRESHOLD - ); - println!(" L{l} is where CPU and Metal first diverge meaningfully."); - if l == 0 { - println!(" Layer 0 drift → culprit is in the embedding or layer-0 pre-norm / attention / FFN."); - } else { - println!( - " Earlier layers match; focus on L{l} attention, FFN, or per-layer scalar." - ); - } - // Also point at stages (dumped for L0 only by the Metal - // prefill hook) so the user can cross-reference. - let stage_dumps = [ - "norm_out", - "q_out", - "k_out", - "v_out", - "attn_out", - "o_out", - "h_post_attn", - ]; - if l == 0 { - println!(); - println!( - " L0 stage files available in {}:", - metal_path.as_path().display() - ); - for s in &stage_dumps { - let p = metal_path.as_path().join(format!("metal_layer_00_{s}.f32")); - if p.is_file() { - println!(" {}", p.display()); - } - } - } - } - None => { - println!("━━━ No layer divergence above threshold ─────────────────────"); - println!(" All layers match within cos_sim >= {DRIFT_THRESHOLD}. Drift"); - println!(" (if any) is below threshold or comes from the lm_head / sampling step."); - } - } - - // ── Stage-by-stage comparison at `stage_layer` ────────────────────── - // Naming convention: Metal writes `metal_layer_{LL}_{stage}.f32` for - // arbitrary layers (when set via LARQL_STAGE_DUMP_LAYER). Layer 0 also - // writes `metal_L0_q_out_after_qk_norm.f32` via a separate hook. CPU - // writes `cpu_L0_.f32` from `attention::block::run_attention_block_core`. - // We match both sides' layout below for a unified comparison table. - println!(); - println!("━━━ Stage-by-stage comparison @ L{stage_layer} ──────────────────────────"); - println!( - " {:<28} {:>10} {:>12} {:>10} {:>10}", - "stage", "cos_sim", "max_abs_Δ", "||cpu||", "||mtl||" - ); - let ll = format!("{stage_layer:02}"); - // Pairs of (pretty name, cpu file suffix, metal file suffix). CPU's - // stage dump is always L0-prefixed by current block.rs convention, so - // we read from that name — any layer picked up by the dump infra - // still writes under `cpu_L0_*` for historical reasons. - let pairs: &[(&str, String, String)] = &[ - ( - "norm_out (pre-Q/K/V)", - "cpu_L0_norm_out.f32".to_string(), - format!("metal_layer_{ll}_norm_out.f32"), - ), - ( - "q_out (raw, pre QK-norm)", - "cpu_L0_q_out_raw.f32".to_string(), - format!("metal_layer_{ll}_q_out.f32"), - ), - ( - "q_out_after_qk_norm", - "cpu_L0_q_out_after_qk_norm.f32".to_string(), - "metal_L0_q_out_after_qk_norm.f32".to_string(), - ), - ( - "q_out_after_rope", - "cpu_L0_q_out_after_rope.f32".to_string(), - String::new(), - ), - ( - "attn_out (softmax·V)", - "cpu_L0_attn_out.f32".to_string(), - format!("metal_layer_{ll}_attn_out.f32"), - ), - ( - "o_out (post Wo-proj)", - "cpu_L0_o_out.f32".to_string(), - format!("metal_layer_{ll}_o_out.f32"), - ), - ]; - for (name, cpu_name, metal_name) in pairs { - if metal_name.is_empty() { - continue; - } - let cpu_path = cpu_path.as_path().join(cpu_name); - let metal_path = metal_path.as_path().join(metal_name); - let cpu = read_f32(&cpu_path); - let metal = read_f32(&metal_path); - match (cpu, metal) { - (Some(c), Some(m)) if c.len() == m.len() => { - let s = layer_stats(&c, &m); - let flag = if s.cos < DRIFT_THRESHOLD { " ←" } else { "" }; - println!( - " {:<28} {:>10.6} {:>12.3e} {:>10.3} {:>10.3}{flag}", - name, s.cos, s.max_abs_diff, s.cpu_norm, s.metal_norm - ); - } - (Some(c), Some(m)) => { - println!( - " {:<28} ", - name, - c.len(), - m.len() - ); - } - (None, _) => println!(" {:<28} ", name, cpu_path.display()), - (_, None) => println!(" {:<28} ", name, metal_path.display()), - } - } - - Ok(()) -} - -#[derive(Debug, Clone)] -struct LayerStat { - cos: f32, - max_abs_diff: f32, - cpu_norm: f32, - metal_norm: f32, -} - -/// Cosine similarity + max absolute element-wise difference, plus each -/// side's L2 norm for scale debugging. -fn layer_stats(cpu: &[f32], metal: &[f32]) -> LayerStat { - let n = cpu.len().min(metal.len()); - let mut dot = 0.0f64; - let mut cn = 0.0f64; - let mut mn = 0.0f64; - let mut max_abs = 0.0f32; - for i in 0..n { - let a = cpu[i] as f64; - let b = metal[i] as f64; - dot += a * b; - cn += a * a; - mn += b * b; - let d = (cpu[i] - metal[i]).abs(); - if d > max_abs { - max_abs = d; - } - } - let cos = if cn > 0.0 && mn > 0.0 { - (dot / (cn.sqrt() * mn.sqrt())) as f32 - } else { - 0.0 - }; - LayerStat { - cos, - max_abs_diff: max_abs, - cpu_norm: cn.sqrt() as f32, - metal_norm: mn.sqrt() as f32, - } -} - -/// Read a raw `f32[]` little-endian file. Returns `None` on any I/O -/// error or non-multiple-of-4 file size. -fn read_f32(path: &Path) -> Option> { - let bytes = std::fs::read(path).ok()?; - if !bytes.len().is_multiple_of(4) { - return None; - } - Some( - bytes - .chunks_exact(4) - .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) - .collect(), - ) -} diff --git a/crates/larql-inference/examples/routing_experiment.rs b/crates/larql-inference/examples/routing_experiment.rs deleted file mode 100644 index 3e498651b..000000000 --- a/crates/larql-inference/examples/routing_experiment.rs +++ /dev/null @@ -1,338 +0,0 @@ -//! Routing experiment: measure template-dependence of attention and FFN routing. -//! -//! For each template × entity: -//! - Capture residual at every layer (last position) -//! - Capture attention weights at every layer -//! - Capture top-K FFN activations at every layer -//! -//! Then measure: -//! 1. Residual cosine stability within template (should be ~0.99) -//! 2. Attention pattern cosine stability within template -//! 3. FFN feature Jaccard overlap within template -//! 4. Cross-template separation (different templates → different routing?) -//! -//! Usage: -//! cargo run --release -p larql-inference --example routing_experiment - -use larql_inference::forward::trace_forward_full; -use larql_inference::{InferenceModel, WeightFfn}; -use std::collections::HashSet; - -fn cosine(a: &[f32], b: &[f32]) -> f32 { - let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum(); - let na: f32 = a.iter().map(|x| x * x).sum::().sqrt(); - let nb: f32 = b.iter().map(|x| x * x).sum::().sqrt(); - if na < 1e-12 || nb < 1e-12 { - return 0.0; - } - dot / (na * nb) -} - -fn jaccard(a: &HashSet, b: &HashSet) -> f32 { - if a.is_empty() && b.is_empty() { - return 1.0; - } - let inter = a.intersection(b).count(); - let union = a.union(b).count(); - if union == 0 { - return 0.0; - } - inter as f32 / union as f32 -} - -/// Flatten attention weights into a single vector for cosine comparison. -fn flatten_attn(weights: &larql_inference::attention::AttentionWeights) -> Vec { - let mut flat = Vec::new(); - for head in &weights.heads { - flat.extend_from_slice(head); - } - flat -} - -fn main() -> Result<(), Box> { - let model = InferenceModel::load("google/gemma-3-4b-it")?; - let weights = model.weights(); - let tokenizer = model.tokenizer(); - let num_layers = weights.num_layers; - let dense_ffn = WeightFfn { weights }; - - let templates: Vec<(&str, &str, Vec<&str>)> = vec![ - ( - "capital", - "The capital of {} is", - vec![ - "France", - "Germany", - "Japan", - "Brazil", - "Egypt", - "Australia", - "Mexico", - "India", - "Canada", - "Italy", - "Spain", - "China", - "Russia", - "Turkey", - "Thailand", - "Argentina", - "Nigeria", - "Kenya", - "Poland", - "Sweden", - ], - ), - ( - "language", - "The language spoken in {} is", - vec![ - "France", - "Germany", - "Japan", - "Brazil", - "Egypt", - "China", - "Russia", - "Thailand", - "Mexico", - "Italy", - "Spain", - "India", - "Turkey", - "Poland", - "Sweden", - "Greece", - "Portugal", - "Vietnam", - "Indonesia", - "Korea", - ], - ), - ( - "currency", - "The currency of {} is the", - vec![ - "Japan", - "Brazil", - "India", - "Mexico", - "China", - "Russia", - "Thailand", - "Turkey", - "Poland", - "Sweden", - "Australia", - "Canada", - "Egypt", - "Nigeria", - "Kenya", - "Argentina", - "Switzerland", - "Norway", - "Denmark", - "Hungary", - ], - ), - ( - "born", - "{} was born in", - vec![ - "Einstein", - "Mozart", - "Shakespeare", - "Picasso", - "Darwin", - "Beethoven", - "Galileo", - "Newton", - "Tesla", - "Curie", - "Aristotle", - "Plato", - "Napoleon", - "Cleopatra", - "Gandhi", - "Confucius", - "Columbus", - "Copernicus", - "Gutenberg", - "Euler", - ], - ), - ]; - - let all_layers: Vec = (0..num_layers).collect(); - let activation_top_k = 200; - - println!("=== Routing Stability Experiment ===\n"); - println!( - "{} templates, {} entities each, {} layers\n", - templates.len(), - templates[0].2.len(), - num_layers - ); - - // Store all results for cross-template analysis - let mut all_residuals: Vec<(String, Vec>>)> = Vec::new(); // (template, [entity][layer][hidden]) - let mut all_attn: Vec<(String, Vec>>)> = Vec::new(); // (template, [entity][layer][flat_attn]) - let mut all_features: Vec<(String, Vec>>)> = Vec::new(); // (template, [entity][layer]{features}) - - for (tname, template, entities) in &templates { - println!("--- Template: {tname} (\"{template}\") ---"); - - let mut t_residuals: Vec>> = Vec::new(); // [entity][layer][hidden] - let mut t_attn: Vec>> = Vec::new(); // [entity][layer][flat_attn] - let mut t_features: Vec>> = Vec::new(); // [entity][layer]{features} - - for entity in entities { - let prompt = template.replace("{}", entity); - let encoding = tokenizer - .encode(prompt.as_str(), true) - .map_err(|e| format!("{e}"))?; - let token_ids: Vec = encoding.get_ids().to_vec(); - - let trace = trace_forward_full( - weights, - &token_ids, - &all_layers, - true, - activation_top_k, - true, - &dense_ffn, - ); - - // Extract per-layer data - let mut e_residuals = Vec::new(); - let mut e_attn = Vec::new(); - let mut e_features = Vec::new(); - - for layer in 0..num_layers { - // Residual - if let Some((_, res)) = trace.residuals.iter().find(|(l, _)| *l == layer) { - e_residuals.push(res.clone()); - } else { - e_residuals.push(vec![0.0; weights.hidden_size]); - } - - // Attention (flatten for cosine comparison) - if let Some(cap) = trace.attention.iter().find(|c| c.layer == layer) { - e_attn.push(flatten_attn(&cap.weights)); - } else { - e_attn.push(vec![]); - } - - // FFN features (top activations with |act| > 1.0) - let feats: HashSet = trace - .activations - .iter() - .find(|(l, _)| *l == layer) - .map(|(_, acts)| { - acts.iter() - .filter(|(_, a)| a.abs() > 1.0) - .map(|(f, _)| *f) - .collect() - }) - .unwrap_or_default(); - e_features.push(feats); - } - - t_residuals.push(e_residuals); - t_attn.push(e_attn); - t_features.push(e_features); - } - - let n = entities.len(); - - // Per-layer stability metrics - println!( - " {:>5} {:>8} {:>9} {:>9} {:>9}", - "Layer", "Res cos", "Attn cos", "FFN Jacc", "FFN union" - ); - - for layer in 0..num_layers { - // Pairwise residual cosine - let mut res_cos_sum = 0.0f64; - let mut attn_cos_sum = 0.0f64; - let mut jacc_sum = 0.0f64; - let mut pairs = 0usize; - - let mut feature_union: HashSet = HashSet::new(); - for feat in t_features.iter().take(n) { - feature_union.extend(feat[layer].iter()); - } - - for i in 0..n { - for j in (i + 1)..n { - res_cos_sum += cosine(&t_residuals[i][layer], &t_residuals[j][layer]) as f64; - if !t_attn[i][layer].is_empty() && !t_attn[j][layer].is_empty() { - attn_cos_sum += cosine(&t_attn[i][layer], &t_attn[j][layer]) as f64; - } - jacc_sum += jaccard(&t_features[i][layer], &t_features[j][layer]) as f64; - pairs += 1; - } - } - - if pairs > 0 && (layer % 4 == 0 || layer == num_layers - 1) { - let res_cos = res_cos_sum / pairs as f64; - let attn_cos = attn_cos_sum / pairs as f64; - let jacc = jacc_sum / pairs as f64; - println!(" L{layer:2}: {res_cos:>7.4} {attn_cos:>8.4} {jacc:>8.4} {feature_union:>8}", - feature_union = feature_union.len()); - } - } - - all_residuals.push((tname.to_string(), t_residuals)); - all_attn.push((tname.to_string(), t_attn)); - all_features.push((tname.to_string(), t_features)); - println!(); - } - - // Cross-template separation: residual cosine between templates - println!("--- Cross-template residual cosine (L16, entity 0 vs entity 0) ---"); - for i in 0..all_residuals.len() { - for j in (i + 1)..all_residuals.len() { - let cos = cosine(&all_residuals[i].1[0][16], &all_residuals[j].1[0][16]); - println!( - " {} vs {}: {cos:.4}", - all_residuals[i].0, all_residuals[j].0 - ); - } - } - - println!("\n--- Cross-template FFN Jaccard (L16, entity 0 vs entity 0) ---"); - for i in 0..all_features.len() { - for j in (i + 1)..all_features.len() { - let jacc = jaccard(&all_features[i].1[0][16], &all_features[j].1[0][16]); - println!( - " {} vs {}: {jacc:.4}", - all_features[i].0, all_features[j].0 - ); - } - } - - // Feature union size across all entities per template (how many distinct features per layer?) - println!("\n--- Feature universe per template per layer ---"); - println!( - " {:>10} {:>5} {:>5} {:>5} {:>5} {:>5}", - "", "L0", "L8", "L16", "L24", "L33" - ); - for (tname, _, t_features) in all_features - .iter() - .map(|(name, feats)| (name, &templates, feats)) - { - let mut line = format!(" {tname:>10}"); - for &layer in &[0, 8, 16, 24, 33] { - let mut union: HashSet = HashSet::new(); - for entity_feats in t_features { - union.extend(entity_feats[layer].iter()); - } - line.push_str(&format!(" {union:>5}", union = union.len())); - } - println!("{line}"); - } - - println!("\n=== Done ==="); - Ok(()) -} diff --git a/crates/larql-inference/examples/scanner_adversarial.rs b/crates/larql-inference/examples/scanner_adversarial.rs deleted file mode 100644 index 6240854d3..000000000 --- a/crates/larql-inference/examples/scanner_adversarial.rs +++ /dev/null @@ -1,51 +0,0 @@ -//! Adversarial session against the tier-0 scanner: prose that carries -//! digits and operator-shaped characters but asks for no computation. -use larql_inference::experts::arith::extract::find_expression; - -fn main() { - let cases: &[&str] = &[ - // spaced ranges / scores / idioms — the '-' with whitespace family - "It takes 5 - 10 business days.", - "They won 3 - 1 at home.", - "I work a 9 - 5 job.", - "Open Monday - Friday, 9 - 17.", - "pages 12 - 48 cover the appendix", - "the score was 2 - 2 after extra time", - "ages 18 - 25 only", - "dated 2026 - 06 - 11 in the ledger", - // 'x' family - "a 4 x 4 truck", - "2 x 4 lumber at the yard", - "a 3 x 5 index card", - "room is 12 x 14 feet", - // '+' in prose - "I have 2 + years of experience", - "rated 4 + stars on average", - "C++ 11 added move semantics", - "call +44 7911 123456", - "she scored 1600 + on the test", - // metaphor words (MEE territory — must be inert in AVE v0.1) - "exponential growth of 300 users", - "let me go off on a tangent about 7 things", - "check the log file at line 42", - "a sine of the times, all 9 of them", - // ambiguous bare/question weak forms — now the model's territory - "9 - 5", - "what is 100 - 7?", - "Are you available 9 - 5?", - // legit math notation that MUST keep firing - "12 + 7 =", - "what is 123456 + 654321?", - "100000 - 1 =", - "12345 * 6789", - "3 x 4 =", - "47−5", - "999 + 111 - 222 =", - ]; - for c in cases { - match find_expression(c) { - Some(e) => println!("FIRE {c:<46} -> {} = {}", e, e.eval()), - None => println!(" no {c}"), - } - } -} diff --git a/crates/larql-inference/examples/speculation_error.rs b/crates/larql-inference/examples/speculation_error.rs deleted file mode 100644 index 37c5d7785..000000000 --- a/crates/larql-inference/examples/speculation_error.rs +++ /dev/null @@ -1,406 +0,0 @@ -//! Speculation error experiment: can we walk FFN layers in parallel? -//! -//! For each layer N, measures the error between: -//! true path: run_ffn(post_attn_residual_N, layer=N) — actual residual -//! spec path: run_ffn(initial_embedding, layer=N) — speculative residual -//! -//! Metrics: -//! cosine_distance between the two FFN deltas -//! feature_overlap Jaccard of top-K active FFN features (K=200) -//! top1_match logit-lens argmax match at each layer -//! -//! Usage: -//! cargo run --release -p larql-inference --example speculation_error -- \ -//! --model google/gemma-3-4b-it \ -//! [--threshold 0.05] [--prompt-sets factual,arithmetic,code] - -use larql_inference::{ - ffn::WeightFfn, - forward::{apply_norm, capture_spec_residuals, dot_proj, run_ffn}, - InferenceModel, -}; -use ndarray::Array2; - -// ── Prompts ───────────────────────────────────────────────────────────── - -const PROMPTS_FACTUAL: &[&str] = &[ - "The capital of France is", - "The capital of Germany is", - "The capital of Japan is", - "The capital of Australia is", - "The capital of Brazil is", - "Albert Einstein was born in", - "Marie Curie was born in", - "Python was created by", - "The Eiffel Tower is located in", - "The Great Wall is located in", -]; - -const PROMPTS_ARITHMETIC: &[&str] = &["2 + 2 =", "7 × 8 =", "15 - 6 =", "100 / 4 ="]; - -const PROMPTS_CODE: &[&str] = &["def fibonacci(n):", "import numpy as", "for i in range("]; - -const TOP_K_FEATURES: usize = 200; - -// ── Args ───────────────────────────────────────────────────────────────── - -struct Args { - model: String, - threshold: f32, - prompt_sets: Vec, -} - -fn parse_args() -> Args { - let raw: Vec = std::env::args().collect(); - let mut model = String::new(); - let mut threshold = 0.05_f32; - let mut prompt_sets = vec![ - "factual".to_string(), - "arithmetic".to_string(), - "code".to_string(), - ]; - - let mut i = 1; - while i < raw.len() { - match raw[i].as_str() { - "--model" => { - i += 1; - model = raw[i].clone(); - } - "--threshold" => { - i += 1; - threshold = raw[i].parse().unwrap_or(0.05); - } - "--prompt-sets" => { - i += 1; - prompt_sets = raw[i].split(',').map(|s| s.to_string()).collect(); - } - _ => {} - } - i += 1; - } - - if model.is_empty() { - eprintln!("Usage: speculation_error --model MODEL [--threshold 0.05] [--prompt-sets factual,arithmetic,code]"); - std::process::exit(1); - } - - Args { - model, - threshold, - prompt_sets, - } -} - -// ── Math helpers ───────────────────────────────────────────────────────── - -fn cosine_distance(a: &[f32], b: &[f32]) -> f32 { - let mut dot = 0.0_f32; - let mut na = 0.0_f32; - let mut nb = 0.0_f32; - for (&ai, &bi) in a.iter().zip(b.iter()) { - dot += ai * bi; - na += ai * ai; - nb += bi * bi; - } - let denom = na.sqrt() * nb.sqrt(); - if denom < 1e-12 { - 1.0 - } else { - 1.0 - dot / denom - } -} - -fn top_k_indices(vals: &[f32], k: usize) -> Vec { - let mut indexed: Vec<(usize, f32)> = vals.iter().copied().enumerate().collect(); - indexed.sort_unstable_by(|a, b| { - b.1.abs() - .partial_cmp(&a.1.abs()) - .unwrap_or(std::cmp::Ordering::Equal) - }); - indexed.truncate(k); - indexed.into_iter().map(|(i, _)| i).collect() -} - -fn jaccard(a: &[usize], b: &[usize]) -> f32 { - use std::collections::HashSet; - let sa: HashSet = a.iter().copied().collect(); - let sb: HashSet = b.iter().copied().collect(); - let intersect = sa.intersection(&sb).count(); - let union_ = sa.union(&sb).count(); - if union_ == 0 { - 1.0 - } else { - intersect as f32 / union_ as f32 - } -} - -fn lm_head_top1(weights: &larql_inference::ModelWeights, h_last: &[f32]) -> usize { - let hidden = h_last.len(); - let norm_offset = weights.arch.norm_weight_offset(); - let h_2d = Array2::from_shape_vec((1, hidden), h_last.to_vec()).unwrap(); - let h_normed = apply_norm(weights, &h_2d, weights.arch.final_norm_key(), norm_offset); - let logits = dot_proj(&h_normed, &weights.lm_head); - let row = logits.row(0); - row.iter() - .enumerate() - .max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal)) - .map(|(i, _)| i) - .unwrap_or(0) -} - -// ── Per-layer stats accumulator ─────────────────────────────────────────── - -#[derive(Default)] -struct LayerStats { - cosine_errs: Vec, - feature_overlaps: Vec, - top1_matches: Vec, -} - -// ── Main ───────────────────────────────────────────────────────────────── - -fn main() -> Result<(), Box> { - let args = parse_args(); - - // Build prompt list - let mut prompts: Vec = Vec::new(); - for set in &args.prompt_sets { - match set.as_str() { - "factual" => prompts.extend(PROMPTS_FACTUAL.iter().map(|s| s.to_string())), - "arithmetic" => prompts.extend(PROMPTS_ARITHMETIC.iter().map(|s| s.to_string())), - "code" => prompts.extend(PROMPTS_CODE.iter().map(|s| s.to_string())), - other => eprintln!("unknown prompt set: {other}"), - } - } - - println!("=== Speculation Error Experiment ===\n"); - println!(" Model: {}", args.model); - println!(" Prompts: {}", prompts.len()); - println!(" Threshold: cosine_distance < {}", args.threshold); - println!(" Top-K feat: {TOP_K_FEATURES}\n"); - - eprintln!("Loading model..."); - let t0 = std::time::Instant::now(); - let inference_model = InferenceModel::load(&args.model)?; - let weights = inference_model.weights(); - let tokenizer = inference_model.tokenizer(); - let num_layers = weights.num_layers; - eprintln!( - " loaded in {:.1}s ({num_layers} layers, hidden={})\n", - t0.elapsed().as_secs_f64(), - weights.hidden_size - ); - - let ffn = WeightFfn { weights }; - - // Per-layer accumulators - let mut stats: Vec = (0..num_layers).map(|_| LayerStats::default()).collect(); - - for (pi, prompt) in prompts.iter().enumerate() { - eprint!( - " [{}/{}] {:?}... ", - pi + 1, - prompts.len(), - &prompt[..prompt.len().min(40)] - ); - let t = std::time::Instant::now(); - - let enc = tokenizer - .encode(prompt.as_str(), true) - .map_err(|e| format!("tokenize: {e}"))?; - let token_ids: Vec = enc.get_ids().to_vec(); - let seq_len = token_ids.len(); - - // Single-pass: capture post-attn and post-layer residuals at every layer - let capture = capture_spec_residuals(weights, &token_ids); - - // Speculative residual: last token of initial embedding - let spec_h0: Vec = capture.h_0.row(seq_len - 1).to_vec(); - let spec_2d = Array2::from_shape_vec((1, weights.hidden_size), spec_h0.clone())?; - - // Precompute spec FFN (delta + activation) for all layers in one pass - let mut spec_deltas: Vec> = Vec::with_capacity(num_layers); - let mut spec_acts: Vec>> = Vec::with_capacity(num_layers); - for layer in 0..num_layers { - let (spec_out, spec_act) = run_ffn(weights, &spec_2d, layer, &ffn, true); - let delta: Vec = spec_out - .row(0) - .iter() - .zip(spec_h0.iter()) - .map(|(o, i)| o - i) - .collect(); - spec_deltas.push(delta); - spec_acts.push(spec_act); - } - - // Per-layer metrics - let mut spec_accum: Vec = spec_h0.clone(); - - for layer in 0..num_layers { - // True FFN delta using actual post-attn residual - let true_h: &[f32] = &capture.post_attn_last[layer]; - let true_2d = Array2::from_shape_vec((1, weights.hidden_size), true_h.to_vec())?; - let (true_out, true_act_opt) = run_ffn(weights, &true_2d, layer, &ffn, true); - let true_delta: Vec = true_out - .row(0) - .iter() - .zip(true_h.iter()) - .map(|(o, i)| o - i) - .collect(); - - let spec_delta = &spec_deltas[layer]; - let spec_act_opt = spec_acts[layer].as_ref(); - - // Cosine distance between FFN deltas - let cos_err = cosine_distance(&true_delta, spec_delta); - - // Feature overlap: Jaccard of top-K active FFN features by activation magnitude - let overlap = match (true_act_opt, spec_act_opt) { - (Some(ta), Some(sa)) => { - let true_features = top_k_indices(&ta.row(0).to_vec(), TOP_K_FEATURES); - let spec_features = top_k_indices(&sa.row(0).to_vec(), TOP_K_FEATURES); - jaccard(&true_features, &spec_features) - } - _ => 0.0, - }; - - // Top-1 match via logit lens - // Accumulate spec residual through layer N - for (acc, d) in spec_accum.iter_mut().zip(spec_delta.iter()) { - *acc += d; - } - let true_top1 = lm_head_top1(weights, &capture.post_layer_last[layer]); - let spec_top1 = lm_head_top1(weights, &spec_accum); - let top1_match = if true_top1 == spec_top1 { 1.0_f32 } else { 0.0 }; - - stats[layer].cosine_errs.push(cos_err); - stats[layer].feature_overlaps.push(overlap); - stats[layer].top1_matches.push(top1_match); - } - - eprintln!("{:.1}s", t.elapsed().as_secs_f64()); - } - - // ── Classification ───────────────────────────────────────────────── - - let threshold = args.threshold; - let mut parallelisable: Vec = Vec::new(); - let mut serial: Vec = Vec::new(); - - // Print header - println!(); - println!("Per-layer cosine distance (true vs speculative delta):"); - println!( - " {:>5} {:>9} {:>6} {:>6} {:>16} {:>11} {:>10}", - "Layer", "Mean err", "Min", "Max", "Feature overlap", "Top-1 match", "Verdict" - ); - println!(" {}", "─".repeat(75)); - - for (layer, s) in stats.iter().enumerate().take(num_layers) { - if s.cosine_errs.is_empty() { - continue; - } - - let mean_err = s.cosine_errs.iter().sum::() / s.cosine_errs.len() as f32; - let min_err = s.cosine_errs.iter().cloned().fold(f32::INFINITY, f32::min); - let max_err = s - .cosine_errs - .iter() - .cloned() - .fold(f32::NEG_INFINITY, f32::max); - let mean_ov = s.feature_overlaps.iter().sum::() / s.feature_overlaps.len() as f32; - let mean_top1 = s.top1_matches.iter().sum::() / s.top1_matches.len() as f32; - - let verdict = if mean_err < threshold { - parallelisable.push(layer); - "PARALLEL" - } else { - serial.push(layer); - "serial" - }; - - println!( - " {:>5} {:>9.4} {:>6.4} {:>6.4} {:>16.3} {:>11.3} {:>10}", - layer, mean_err, min_err, max_err, mean_ov, mean_top1, verdict - ); - } - - // ── Band structure ───────────────────────────────────────────────── - - println!(); - println!("Band structure (threshold = {threshold}):"); - - struct Band { - kind: &'static str, - start: usize, - end: usize, - } - let mut bands: Vec = Vec::new(); - - for layer in 0..num_layers { - let kind = if parallelisable.contains(&layer) { - "PARALLEL" - } else { - "serial" - }; - match bands.last_mut() { - Some(b) if b.kind == kind => { - b.end = layer; - } - _ => bands.push(Band { - kind, - start: layer, - end: layer, - }), - } - } - - let parallel_ms_per_band = 55.0_f32; - let serial_ms_per_layer = 8.0_f32; - let mut estimated_ms = 0.0_f32; - - for b in &bands { - let n = b.end - b.start + 1; - let ms = if b.kind == "PARALLEL" { - estimated_ms += parallel_ms_per_band; - parallel_ms_per_band - } else { - let m = n as f32 * serial_ms_per_layer; - estimated_ms += m; - m - }; - println!( - " L{:02}–L{:02} ({:2} layers) {} ~{:.0}ms", - b.start, b.end, n, b.kind, ms - ); - } - - let serial_baseline = num_layers as f32 * serial_ms_per_layer; - let speedup = serial_baseline / estimated_ms.max(1.0); - - println!(); - println!(" Round trips: {}", bands.len()); - println!(" Estimated wall: {estimated_ms:.0}ms"); - println!(" Serial baseline: {serial_baseline:.0}ms"); - println!(" Speedup: {speedup:.1}×"); - println!(); - - // ── Aggressive threshold ─────────────────────────────────────────── - - let aggressive = 0.15_f32; - let agg_parallel = stats - .iter() - .enumerate() - .filter(|(_, s)| { - !s.cosine_errs.is_empty() && { - let mean = s.cosine_errs.iter().sum::() / s.cosine_errs.len() as f32; - mean < aggressive - } - }) - .count(); - let agg_serial = num_layers - agg_parallel; - println!(" Aggressive threshold ({aggressive}): {agg_parallel}/{num_layers} layers PARALLEL, {agg_serial} serial"); - - Ok(()) -} diff --git a/crates/larql-inference/examples/walk_boundary_sweep.rs b/crates/larql-inference/examples/walk_boundary_sweep.rs deleted file mode 100644 index 8de7c5470..000000000 --- a/crates/larql-inference/examples/walk_boundary_sweep.rs +++ /dev/null @@ -1,288 +0,0 @@ -//! Walk boundary sweep — tests vindex FFN walk at every layer boundary. -//! -//! For each boundary B: -//! Layers 0..B: dense attention + dense FFN (WeightFfn) -//! Layers B..33: dense attention + vindex FFN (WalkFfn) -//! -//! Reports top-1 prediction and probability for each boundary, comparing -//! against the ground truth (all-dense forward pass). -//! -//! The vindex has all 34 layers (1,307,232 vectors). This sweep finds -//! how far down the walk can go while maintaining accuracy. -//! -//! Usage: -//! cargo run --release -p larql-inference --example walk_boundary_sweep -- \ -//! --model google/gemma-3-4b-it \ -//! --vindex path/to/gemma3-4b.vindex -//! -//! Optional: -//! --top-k 8092 Gate KNN top-K (default: 8092) -//! --prompts Comma-separated prompts (default: built-in entity set) - -use std::path::PathBuf; -use std::time::Instant; - -use larql_inference::{ - predict, predict_with_ffn, predict_with_router, vindex::WalkFfn, InferenceModel, - LayerFfnRouter, PredictResult, WeightFfn, -}; -use larql_vindex::{SilentLoadCallbacks, VectorIndex}; - -/// Default test prompts — entities with known ground truth answers. -/// Keep small for fast sweep; add --prompts for larger sets. -const DEFAULT_PROMPTS: &[(&str, &str)] = &[ - ("The capital of France is", "Paris"), - ("The capital of Germany is", "Berlin"), - ("The capital of Japan is", "Tokyo"), - ("The capital of Italy is", "Rome"), - ("The largest planet in our solar system is", "Jupiter"), -]; - -#[allow(clippy::type_complexity)] -fn parse_args() -> (String, PathBuf, usize, Option>) { - let args: Vec = std::env::args().collect(); - let mut model = String::new(); - let mut vindex = PathBuf::new(); - let mut top_k = 8092; - let mut prompts: Option> = None; - - let mut i = 1; - while i < args.len() { - match args[i].as_str() { - "--model" => { - i += 1; - model = args[i].clone(); - } - "--vindex" => { - i += 1; - vindex = PathBuf::from(&args[i]); - } - "--top-k" => { - i += 1; - top_k = if args[i] == "full" || args[i] == "unlimited" { - usize::MAX - } else { - args[i].parse().unwrap() - }; - } - "--prompts" => { - i += 1; - prompts = Some( - args[i] - .split(';') - .map(|p| { - let parts: Vec<&str> = p.splitn(2, '=').collect(); - if parts.len() == 2 { - (parts[0].trim().to_string(), parts[1].trim().to_string()) - } else { - (p.trim().to_string(), String::new()) - } - }) - .collect(), - ); - } - _ => {} - } - i += 1; - } - - if model.is_empty() || !vindex.is_dir() { - eprintln!("Usage: walk_boundary_sweep --model MODEL --vindex PATH [--top-k N]"); - eprintln!(" --model HuggingFace model ID or local path"); - eprintln!(" --vindex Path to .vindex directory"); - eprintln!(" --top-k Gate KNN top-K (default: 8092)"); - std::process::exit(1); - } - - (model, vindex, top_k, prompts) -} - -/// Check if the ground truth is in the top-1 prediction. -fn is_correct(result: &PredictResult, expected: &str) -> bool { - if expected.is_empty() { - return true; - } - result - .predictions - .first() - .map(|(tok, _)| tok.to_lowercase().contains(&expected.to_lowercase())) - .unwrap_or(false) -} - -fn main() -> Result<(), Box> { - let (model_name, vindex_path, top_k, custom_prompts) = parse_args(); - - println!("=== Walk Boundary Sweep ===\n"); - - // ── Load model ── - println!("Loading model: {model_name}"); - let t0 = Instant::now(); - let model = InferenceModel::load(&model_name)?; - println!(" Model loaded in {:.1}s", t0.elapsed().as_secs_f64()); - - let weights = model.weights(); - let tokenizer = model.tokenizer(); - let num_layers = weights.num_layers; - println!(" {} layers, hidden={}", num_layers, weights.hidden_size); - - // ── Load vindex ── - println!("Loading vindex: {}", vindex_path.display()); - let t0 = Instant::now(); - let mut cb = SilentLoadCallbacks; - let index = VectorIndex::load_vindex(&vindex_path, &mut cb)?; - println!( - " {} layers, {} vectors loaded in {:.1}s", - index.num_layers, - index.total_gate_vectors(), - t0.elapsed().as_secs_f64() - ); - println!(); - - // ── Test prompts ── - let prompts: Vec<(String, String)> = match custom_prompts { - Some(p) => p, - None => DEFAULT_PROMPTS - .iter() - .map(|(p, e)| (p.to_string(), e.to_string())) - .collect(), - }; - - // ── Ground truth: all-dense forward pass ── - println!("--- Ground Truth (all-dense) ---\n"); - let mut ground_truth: Vec<(String, f64)> = Vec::new(); - for (prompt, expected) in &prompts { - let encoding = tokenizer - .encode(prompt.as_str(), true) - .map_err(|e| format!("tokenize: {e}"))?; - let token_ids: Vec = encoding.get_ids().to_vec(); - let result = predict(weights, tokenizer, &token_ids, 5); - let (top1, prob) = result - .predictions - .first() - .map(|(t, p)| (t.clone(), *p)) - .unwrap_or_default(); - let correct = is_correct(&result, expected); - let mark = if correct { "+" } else { "-" }; - println!(" [{mark}] \"{prompt}\" -> {top1} ({:.2}%)", prob * 100.0); - ground_truth.push((top1, prob)); - } - println!(); - - // ── Sweep boundaries ── - let boundaries: Vec = { - let mut b = vec![0, 4, 8, 12, 16, 20, 24, 28]; - b.push(num_layers); // all-dense baseline - b.retain(|&v| v <= num_layers); - b.sort_unstable(); - b.dedup(); - b - }; - - println!("--- Boundary Sweep (dense 0..B, walk B..{num_layers}) ---"); - println!( - " {} boundaries x {} prompts = {} forward passes\n", - boundaries.len(), - prompts.len(), - boundaries.len() * prompts.len() - ); - println!( - " {:>4} {:>6} {:>8} {:>8} {:>6} details", - "B", "walk%", "correct", "top1_avg", "time" - ); - println!(" {:-<74}", ""); - - for &boundary in &boundaries { - let walk_pct = (num_layers - boundary) as f64 / num_layers as f64 * 100.0; - - let weight_ffn = WeightFfn { weights }; - let walk_ffn = WalkFfn::new(weights, &index, top_k); - - // Build per-layer backend routing - let mut backends: Vec<&dyn larql_inference::FfnBackend> = vec![&weight_ffn; num_layers]; - for backend in backends.iter_mut().take(num_layers).skip(boundary) { - *backend = &walk_ffn; - } - let router = LayerFfnRouter::per_layer(backends); - - let mut correct_count = 0; - let mut total_prob = 0.0; - let mut details = Vec::new(); - let sweep_start = Instant::now(); - - for (i, (prompt, expected)) in prompts.iter().enumerate() { - let encoding = tokenizer - .encode(prompt.as_str(), true) - .map_err(|e| format!("tokenize: {e}"))?; - let token_ids: Vec = encoding.get_ids().to_vec(); - - let result = if boundary == num_layers { - // All dense — skip router overhead - predict(weights, tokenizer, &token_ids, 5) - } else if boundary == 0 { - // All walk - predict_with_ffn(weights, tokenizer, &token_ids, 5, &walk_ffn) - } else { - predict_with_router(weights, tokenizer, &token_ids, 5, &router) - }; - - let (top1, prob) = result - .predictions - .first() - .map(|(t, p)| (t.clone(), *p)) - .unwrap_or_default(); - - let matches_ground = top1 == ground_truth[i].0; - let correct = is_correct(&result, expected); - if correct { - correct_count += 1; - } - total_prob += prob; - - // Track divergence from ground truth - if !matches_ground { - details.push(format!( - "{}->{}({:.0}%)", - ground_truth[i].0, - top1, - prob * 100.0 - )); - } - } - - let elapsed = sweep_start.elapsed(); - let avg_prob = total_prob / prompts.len() as f64 * 100.0; - let detail_str = if details.is_empty() { - "all match ground truth".to_string() - } else { - details.join(", ") - }; - - println!( - " L{boundary:<3} {walk_pct:>5.0}% {correct_count:>3}/{:<3} {avg_prob:>7.2}% {:.1}s {detail_str}", - prompts.len(), - elapsed.as_secs_f64() - ); - } - - println!(); - println!(" Legend:"); - println!(" B = boundary layer (dense 0..B, walk B..{num_layers})"); - println!(" walk% = percentage of layers using vindex FFN"); - println!(" correct = prompts where top-1 matches expected answer"); - println!(" top1_avg = average top-1 probability across all prompts"); - println!(" details = divergences from ground truth"); - println!(); - - // ── Summary ── - println!("--- Summary ---\n"); - println!(" Ground truth: all-dense f32 forward pass"); - println!(" Walk: vindex gate KNN top-{top_k} -> sparse FFN"); - println!(" Attention: BLAS-fused (dense) at all layers for all boundaries"); - println!(" {} test prompts, {} layers", prompts.len(), num_layers); - println!(); - println!(" If walk holds to L0: FFN quantization is unnecessary."); - println!(" Only attention weights (Q/K/V/O) and embed/logits need quantization."); - - println!("\n=== Done ==="); - Ok(()) -} diff --git a/crates/larql-inference/examples/walk_ffn_accuracy.rs b/crates/larql-inference/examples/walk_ffn_accuracy.rs deleted file mode 100644 index 795a8bfa0..000000000 --- a/crates/larql-inference/examples/walk_ffn_accuracy.rs +++ /dev/null @@ -1,507 +0,0 @@ -//! WalkFfn accuracy frontier — the predictive-quality companion to -//! `walk_ffn_microbench` (task #19). The microbench showed cheap -//! (precomputed) routing makes sparse WalkFfn *beat* dense up to ~15× at -//! small K. Speed is settled; this asks the price in **predictive -//! quality**, judged in the Shannon discipline (next-token KL / bits, not -//! cosine), never in wall-time. -//! -//! For each eval prompt it runs a full forward (embedding → layers → -//! lm_head) three ways and compares the **last-token next-token -//! distribution** against dense ground truth: -//! - **dense** `WalkFfn::new_unlimited` — ground truth (all features) -//! - **gate-KNN** `WalkFfn::new(.., k)` — content-addressed router (smart, slow) -//! - **cheap** precomputed strided route + O(K) — token-independent router (fast, dumb) -//! -//! The strided cheap route is a deliberate *lower bound* on cheap-routing -//! quality (it ignores the input entirely). A real hash route (Exp 27, -//! token-deterministic) sits between it and gate-KNN — high overlap with -//! the gate's pick at early layers, decaying with depth. The gate-KNN row -//! is the accuracy ceiling reachable by *any* size-K route. -//! -//! Usage: `cargo run --release --example walk_ffn_accuracy -- [VINDEX_DIR]` - -use larql_inference::vindex::{insert_q4k_layer_tensors_resident, WalkFfn, WalkFfnConfig}; -use larql_inference::{load_tokenizer, predict_with_ffn}; -use larql_models::ModelWeights; -use std::collections::HashMap; -use std::sync::Arc; - -/// Deterministic strided route of `k` features per layer — the cost -/// profile of hash routing (precomputed, no gate projection), but a naive -/// (content-blind) *selection*. See module docs. -fn precomputed_pool(num_layers: usize, num_features: usize, k: usize) -> Arc>> { - let k = k.min(num_features.max(1)); - let stride = (num_features / k.max(1)).max(1); - let per_layer: Vec = (0..k).map(|i| (i * stride) % num_features).collect(); - Arc::new(vec![per_layer; num_layers]) -} - -/// Static-importance route: top-`k` features per layer by ‖down_row‖ — the -/// features that move the residual most *when active*. Content-blind (same -/// pool for every input) but **informed**, and as cheap as the strided -/// route (precomputed once, no gate projection). The honest middle rung -/// between strided (uninformed) and gate-KNN (content-addressed). Built -/// from `down_row_norms_pub`, which dequantises the down matrix once. -fn static_importance_pool( - weights: &ModelWeights, - index: &larql_vindex::VectorIndex, - k: usize, -) -> Arc>> { - let probe = WalkFfn::new_unlimited(weights, index); - let per_layer: Vec> = (0..weights.num_layers) - .map(|layer| { - let feats = index.num_features(layer); - let k = k.min(feats.max(1)); - match probe.down_row_norms_pub(layer) { - Some(norms) => { - let mut idx: Vec = (0..norms.len()).collect(); - idx.sort_unstable_by(|&a, &b| norms[b].total_cmp(&norms[a])); - idx.truncate(k); - idx - } - // No down norms (shouldn't happen on a Q4K vindex) — fall - // back to a strided pick so the layer still routes. - None => (0..k) - .map(|i| (i * (feats / k.max(1)).max(1)) % feats) - .collect(), - } - }) - .collect(); - Arc::new(per_layer) -} - -/// Full-vocab next-token distribution keyed by token id, from a forward -/// pass with the given FFN backend. `predict_with_ffn` softmaxes over the -/// whole vocab and (with a huge top_k) returns every token. -fn next_token_dist( - weights: &ModelWeights, - tok: &tokenizers::Tokenizer, - ids: &[u32], - ffn: &dyn larql_inference::ffn::FfnBackend, -) -> HashMap { - let r = predict_with_ffn(weights, tok, ids, usize::MAX, ffn); - r.token_ids - .into_iter() - .zip(r.predictions.into_iter().map(|(_, p)| p)) - .collect() -} - -/// KL(P‖Q) in **bits**, plus top-1 agreement and Q's probability mass on -/// P's argmax token. P = dense ground truth, Q = candidate. -fn compare(p: &HashMap, q: &HashMap) -> (f64, bool, f64) { - let eps = 1e-12; - let mut kl = 0.0; - for (&id, &pi) in p { - if pi <= 0.0 { - continue; - } - let qi = q.get(&id).copied().unwrap_or(0.0).max(eps); - kl += pi * (pi.max(eps) / qi).ln(); - } - let p_arg = p.iter().max_by(|a, b| a.1.total_cmp(b.1)).map(|(i, _)| *i); - let q_arg = q.iter().max_by(|a, b| a.1.total_cmp(b.1)).map(|(i, _)| *i); - let agree = p_arg.is_some() && p_arg == q_arg; - let q_on_p_arg = p_arg.and_then(|id| q.get(&id)).copied().unwrap_or(0.0); - (kl / std::f64::consts::LN_2, agree, q_on_p_arg) -} - -fn main() { - let args: Vec = std::env::args().collect(); - let vindex = args - .get(1) - .cloned() - .unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".to_string()); - let dir = std::path::PathBuf::from(&vindex); - let mut cb = larql_vindex::SilentLoadCallbacks; - - eprintln!("Loading {vindex} ..."); - let mut weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index - .load_interleaved_kquant(&dir) - .expect("interleaved kquant"); - index.load_attn_kquant(&dir).expect("attn kquant"); - let _ = index.load_lm_head_kquant(&dir); - let _ = index.load_down_features_q4k(&dir); - let _ = index.load_down_features(&dir); - let _ = index.load_gate_vectors_q4(&dir); - let tok = load_tokenizer(&dir).expect("tokenizer"); - - // `predict_with_ffn` reads attention from f32 `weights.vectors`, but the - // Q4K loader leaves attention quantised — so attention would no-op and - // the forward degenerates (top-1 = echoed last token). Dequantise every - // layer's attention to f32 up front. FFN still comes from the WalkFfn - // (which reads Q4K straight from the index), so the FFN router stays the - // only variable across the three configs. - eprintln!( - "Dequantising attention for {} layers ...", - weights.num_layers - ); - for layer in 0..weights.num_layers { - insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant layer"); - } - - let feats = index.num_features(weights.num_layers / 2); - let prompts = [ - "The capital of France is", - "Water is made of hydrogen and", - "The opposite of hot is", - "2 + 2 =", - ]; - let ks = [2048usize, 512, 128, 32]; - - // Sanity: dense top-1 per prompt (catches a broken forward before we - // read anything into the KL numbers). - eprintln!("\nDense top-1 sanity:"); - for p in &prompts { - let ids = tok.encode(*p, true).expect("encode").get_ids().to_vec(); - let r = predict_with_ffn( - &weights, - &tok, - &ids, - 1, - &WalkFfn::new_unlimited(&weights, &index), - ); - eprintln!( - " {p:?} → {:?} ({:.3})", - r.predictions.first().map(|x| x.0.as_str()).unwrap_or("?"), - r.predictions.first().map(|x| x.1).unwrap_or(0.0) - ); - } - - // ── Parity diagnostic ──────────────────────────────────────────── - // Before reading sparsity numbers, confirm the sparse-WALK path (the - // per-feature loop) reproduces the dense NATIVE path at full K. If it - // doesn't, divergence below is a path-fidelity bug, not a routing - // result. `force_walk` defeats the gemv fast-path so the walk actually - // runs; the pool=all + precomputed_routing variant tests local gate - // computation over every feature. - { - println!("\nFull-K parity vs dense native (should be ~0 if the walk path is faithful):"); - let all: Vec = (0..feats).collect(); - let pool_all = Arc::new(vec![all; weights.num_layers]); - for p in &prompts { - let ids = tok.encode(*p, true).expect("encode").get_ids().to_vec(); - let dense = next_token_dist( - &weights, - &tok, - &ids, - &WalkFfn::new_unlimited(&weights, &index), - ); - let walk_cfg = WalkFfnConfig::sparse(weights.num_layers, feats).with_force_walk(true); - let walk = next_token_dist( - &weights, - &tok, - &ids, - &WalkFfn::from_config(&weights, &index, walk_cfg), - ); - let cheap_cfg = WalkFfnConfig::sparse(weights.num_layers, feats) - .with_pool_per_layer(pool_all.clone()) - .with_precomputed_routing(true); - let cheap = next_token_dist( - &weights, - &tok, - &ids, - &WalkFfn::from_config(&weights, &index, cheap_cfg), - ); - let (kw, aw, _) = compare(&dense, &walk); - let (kc, ac, _) = compare(&dense, &cheap); - println!( - " {p:<32} walk-fullK KL={kw:>8.4} agree={aw} cheap-fullK KL={kc:>8.4} agree={ac}" - ); - } - } - - println!("\nWalkFfn accuracy vs dense — {feats} features, KL in bits (lower=better)\n"); - println!( - "{:<34} {:>10} {:>10} {:>10}", - "config", "KL(bits)", "top1-agree", "q@p_argmax" - ); - - for &k in &ks { - // Average over prompts for each backend. - let pct = 100.0 * k as f64 / feats.max(1) as f64; - let pool = precomputed_pool(weights.num_layers, feats, k); - - // `cand` maps token ids → candidate FFN backend for this prompt. - let acc = |label: String, cand: &dyn Fn(&[u32]) -> HashMap| { - let (mut kl, mut agree, mut qmass, mut n) = (0.0, 0usize, 0.0, 0usize); - for p in &prompts { - let ids = tok.encode(*p, true).expect("encode").get_ids().to_vec(); - let dense = next_token_dist( - &weights, - &tok, - &ids, - &WalkFfn::new_unlimited(&weights, &index), - ); - let c = cand(&ids); - let (b, a, q) = compare(&dense, &c); - kl += b; - agree += a as usize; - qmass += q; - n += 1; - } - let n = n.max(1) as f64; - println!( - "{label:<34} {:>10.4} {:>9.0}% {:>10.4}", - kl / n, - 100.0 * agree as f64 / n, - qmass / n - ); - }; - - acc(format!("gate-KNN k={k} ({pct:.0}%)"), &|ids| { - next_token_dist(&weights, &tok, ids, &WalkFfn::new(&weights, &index, k)) - }); - let pool_c = pool.clone(); - acc(format!("cheap-route k={k} ({pct:.0}%)"), &|ids| { - let cfg = WalkFfnConfig::sparse(weights.num_layers, k) - .with_pool_per_layer(pool_c.clone()) - .with_precomputed_routing(true); - next_token_dist( - &weights, - &tok, - ids, - &WalkFfn::from_config(&weights, &index, cfg), - ) - }); - println!(); - } - - // ── Hourglass sweep ────────────────────────────────────────────── - // All-layer sparsity collapses (above). The roadmap design is dense - // early, sparse late (Exp 5c hourglass + Exp 27 "token-determinism - // falls off by L3"). Fix K=512 and vary how many *late* layers go - // sparse — find the depth where accuracy survives. - let nl = weights.num_layers; - let k = 512usize; - let pool = precomputed_pool(nl, feats, k); - let pool_imp = static_importance_pool(&weights, &index, k); - // Candidate pools for the two-stage router: static-importance sets of - // increasing size P, ranked per-position by gate score down to K. If - // accuracy improves with P, candidate *recall* is the bottleneck (not - // the ranking); if it plateaus far from gate-KNN, the static metric - // itself can't capture the input-dependent features. - let cand_ps = [2048usize, 4096, 8192]; - let pool_cands: Vec<(usize, Arc>>)> = cand_ps - .iter() - .map(|&p| (p, static_importance_pool(&weights, &index, p))) - .collect(); - println!("\nHourglass: dense early + sparse-K={k} late (vary sparse-from), KL in bits\n"); - println!( - "{:<34} {:>10} {:>10} {:>10}", - "config", "KL(bits)", "top1-agree", "q@p_argmax" - ); - for &frac in &[0.9f64, 0.75, 0.5] { - let sparse_from = (nl as f64 * frac) as usize; - let n_sparse = nl - sparse_from; - let acc = |label: String, cand: &dyn Fn(&[u32]) -> HashMap| { - let (mut kl, mut agree, mut qmass, mut n) = (0.0, 0usize, 0.0, 0usize); - for p in &prompts { - let ids = tok.encode(*p, true).expect("encode").get_ids().to_vec(); - let dense = next_token_dist( - &weights, - &tok, - &ids, - &WalkFfn::new_unlimited(&weights, &index), - ); - let c = cand(&ids); - let (b, a, q) = compare(&dense, &c); - kl += b; - agree += a as usize; - qmass += q; - n += 1; - } - let n = n.max(1) as f64; - println!( - "{label:<34} {:>10.4} {:>9.0}% {:>10.4}", - kl / n, - 100.0 * agree as f64 / n, - qmass / n - ); - }; - acc(format!("gate-KNN sparse last {n_sparse}/{nl}"), &|ids| { - let cfg = WalkFfnConfig::hybrid(nl, sparse_from, k); - next_token_dist( - &weights, - &tok, - ids, - &WalkFfn::from_config(&weights, &index, cfg), - ) - }); - let pool_c = pool.clone(); - acc(format!("strided sparse last {n_sparse}/{nl}"), &|ids| { - let cfg = WalkFfnConfig::hybrid(nl, sparse_from, k) - .with_pool_per_layer(pool_c.clone()) - .with_precomputed_routing(true); - next_token_dist( - &weights, - &tok, - ids, - &WalkFfn::from_config(&weights, &index, cfg), - ) - }); - let pool_i = pool_imp.clone(); - acc(format!("static-imp sparse last {n_sparse}/{nl}"), &|ids| { - let cfg = WalkFfnConfig::hybrid(nl, sparse_from, k) - .with_pool_per_layer(pool_i.clone()) - .with_precomputed_routing(true); - next_token_dist( - &weights, - &tok, - ids, - &WalkFfn::from_config(&weights, &index, cfg), - ) - }); - for (p, pool_cd) in &pool_cands { - // (1) Two-stage with the CHEAP Q4K within-pool gate score. - let pool_q4 = pool_cd.clone(); - acc( - format!("2-stage Q4K P={p} last {n_sparse}/{nl}"), - &|ids| { - let cfg = WalkFfnConfig::hybrid(nl, sparse_from, k) - .with_pool_per_layer(pool_q4.clone()) - .with_precomputed_routing(true) - .with_rank_within_pool(true); - next_token_dist( - &weights, - &tok, - ids, - &WalkFfn::from_config(&weights, &index, cfg), - ) - }, - ); - // (2) DE-CONFOUND: same static pool, ranked by gate-KNN's OWN - // full-precision f32 score (`pool_restricted_gate_knn`, via - // precomputed_routing=false). Differs from (1) only in score - // precision, and from the gate-KNN baseline only in the pool - // restriction. Plateau ≈ static-imp → candidate set is the wall - // (#22 justified); drop toward gate-KNN → the Q4K score was the - // wall and static-pool + f32 ranking widens the band, no - // clustering needed. - let pool_fp = pool_cd.clone(); - acc( - format!("deconf f32 P={p} last {n_sparse}/{nl}"), - &|ids| { - let cfg = WalkFfnConfig::hybrid(nl, sparse_from, k) - .with_pool_per_layer(pool_fp.clone()); - next_token_dist( - &weights, - &tok, - ids, - &WalkFfn::from_config(&weights, &index, cfg), - ) - }, - ); - } - println!(); - } - - // ── n≈30 confirmation: is the gate-KNN vs best-static-pool gap real? ── - // The hourglass numbers above are n=4 (top-1 bounces 25/50/75%). Before - // resourcing the #22 clustering pipeline, confirm the load-bearing gap — - // gate-KNN vs the best static pool (f32-ranked P=4096) at the 9-layer - // band — at n≈30. KL is the reliable metric; report mean/median/spread. - let sparse_from = nl.saturating_sub(9); - let pool_best = static_importance_pool(&weights, &index, 4096); - let eval: [&str; 30] = [ - "The capital of France is", - "Water is made of hydrogen and", - "The opposite of hot is", - "The sun rises in the", - "The first president of the United States was", - "A group of lions is called a", - "The chemical symbol for gold is", - "The largest planet in the solar system is", - "Romeo and Juliet was written by", - "The speed of light is approximately", - "The capital of Japan is", - "Photosynthesis occurs in the", - "The square root of 64 is", - "The freezing point of water in Celsius is", - "The author of Pride and Prejudice is", - "An apple a day keeps the doctor", - "The currency of the United Kingdom is the", - "The tallest mountain on Earth is", - "DNA stands for", - "The capital of Italy is", - "The number of continents on Earth is", - "A baby dog is called a", - "The boiling point of water in Celsius is", - "The planet known as the Red Planet is", - "The longest river in the world is the", - "The inventor of the telephone was", - "The opposite of up is", - "The third planet from the sun is", - "The capital of Germany is", - "Two plus three equals", - ]; - - let mut kl_gate: Vec = Vec::new(); - let mut kl_stat: Vec = Vec::new(); - let mut gate_wins = 0usize; - for p in &eval { - let ids = tok.encode(*p, true).expect("encode").get_ids().to_vec(); - let dense = next_token_dist( - &weights, - &tok, - &ids, - &WalkFfn::new_unlimited(&weights, &index), - ); - let gate = next_token_dist( - &weights, - &tok, - &ids, - &WalkFfn::from_config(&weights, &index, WalkFfnConfig::hybrid(nl, sparse_from, k)), - ); - let stat = next_token_dist( - &weights, - &tok, - &ids, - &WalkFfn::from_config( - &weights, - &index, - WalkFfnConfig::hybrid(nl, sparse_from, k).with_pool_per_layer(pool_best.clone()), - ), - ); - let (kg, _, _) = compare(&dense, &gate); - let (ks, _, _) = compare(&dense, &stat); - if kg < ks { - gate_wins += 1; - } - kl_gate.push(kg); - kl_stat.push(ks); - } - - let stats = |v: &mut Vec| -> (f64, f64, f64, f64) { - let n = v.len().max(1) as f64; - let mean = v.iter().sum::() / n; - v.sort_by(|a, b| a.total_cmp(b)); - let median = v[v.len() / 2]; - (mean, median, v[0], v[v.len() - 1]) - }; - let (gm, gmd, glo, ghi) = stats(&mut kl_gate); - let (sm, smd, slo, shi) = stats(&mut kl_stat); - println!( - "\nn={} confirmation — gate-KNN vs best static pool (f32, P=4096), sparse last 9/{nl}, KL bits\n", - eval.len() - ); - println!( - "{:<24} {:>8} {:>8} {:>8} {:>8}", - "config", "mean", "median", "min", "max" - ); - println!( - "{:<24} {gm:>8.3} {gmd:>8.3} {glo:>8.3} {ghi:>8.3}", - "gate-KNN (content-addr)" - ); - println!( - "{:<24} {sm:>8.3} {smd:>8.3} {slo:>8.3} {shi:>8.3}", - "static pool f32 P=4096" - ); - println!( - "\n mean gap {:.2}×, median gap {:.2}×; gate-KNN beats static on {gate_wins}/{} prompts", - sm / gm.max(1e-9), - smd / gmd.max(1e-9), - eval.len() - ); -} diff --git a/crates/larql-inference/examples/walk_ffn_cell_router.rs b/crates/larql-inference/examples/walk_ffn_cell_router.rs deleted file mode 100644 index 41c533e9c..000000000 --- a/crates/larql-inference/examples/walk_ffn_cell_router.rs +++ /dev/null @@ -1,672 +0,0 @@ -//! WalkFfn residual-cell content-addressed router (task #22). -//! -//! Tasks #20/#21 left a gap: a *static* candidate pool can't reach gate-KNN at -//! the 9-layer sparse band (n=30: gate-KNN median 1.30 vs static 3.74 bits). -//! The features gate-KNN picks are input-dependent and not in any static pool. -//! This builds a genuinely **content-addressed** candidate set — an IVF-style -//! residual-cell index — and measures whether it closes that gap at -//! cheap-routing cost. -//! -//! Pipeline: -//! 1. CALIBRATE — run dense forward over a calibration corpus with a -//! `CapturingFfn` that records, per band layer, each position's FFN-input -//! residual + its gate-KNN top-K feature pick. -//! 2. BUILD — k-means the residuals per layer into C cells; each cell's pool -//! = the most-frequent gate-KNN features across its members (capped). This -//! is `CellRouter`. -//! 3. EVAL — in-distribution + OOD prompt sets, 9-layer band, KL measured -//! against DENSE (KL 0 = dense; gate-KNN is itself a lossy top-K, not a -//! floor). Paired Wilcoxon signed-rank separates real differences from -//! heavy-tail median artifacts; OOD separates denoising from overfit. -//! -//! Usage: `cargo run --release --example walk_ffn_cell_router -- [VINDEX_DIR]` - -use larql_inference::ffn::FfnBackend; -use larql_inference::vindex::{ - insert_q4k_layer_tensors_resident, CellRouter, WalkFfn, WalkFfnConfig, -}; -use larql_inference::{load_tokenizer, predict_with_ffn}; -use larql_models::ModelWeights; -use ndarray::Array2; -use std::cell::RefCell; -use std::collections::HashMap; -use std::sync::Arc; - -const K: usize = 512; // sparse top-K per layer (eval) -const N_CELLS: usize = 64; // residual cells per layer -const MAX_POOL: usize = 2048; // cap on a cell's candidate pool -const KMEANS_ITERS: usize = 8; - -/// Per layer: captured `(FFN-input residual, gate-KNN top-K feature ids)`. -type LayerSamples = Vec, Vec)>>; - -/// Wraps a dense WalkFfn; on each band-layer forward, records the FFN-input -/// residual per position and its gate-KNN top-K pick, then delegates so the -/// captured residuals are the true ones. -struct CapturingFfn<'a> { - inner: WalkFfn<'a>, - index: &'a larql_vindex::VectorIndex, - sparse_from: usize, - samples: RefCell, -} - -impl<'a> CapturingFfn<'a> { - fn new( - inner: WalkFfn<'a>, - index: &'a larql_vindex::VectorIndex, - nl: usize, - sparse_from: usize, - ) -> Self { - Self { - inner, - index, - sparse_from, - samples: RefCell::new(vec![Vec::new(); nl]), - } - } - fn record(&self, layer: usize, x: &Array2) { - if layer < self.sparse_from { - return; - } - let mut s = self.samples.borrow_mut(); - for row in x.rows() { - let r = row.to_owned(); - let hits = self.index.gate_knn(layer, &r, K); - s[layer].push((r.to_vec(), hits.into_iter().map(|(f, _)| f).collect())); - } - } -} - -impl FfnBackend for CapturingFfn<'_> { - fn forward(&self, layer: usize, x: &Array2) -> Array2 { - self.record(layer, x); - self.inner.forward(layer, x) - } - fn forward_observed( - &self, - layer: usize, - x: &Array2, - ) -> (Array2, larql_inference::ffn::FfnActivations) { - self.record(layer, x); - self.inner.forward_observed(layer, x) - } - fn name(&self) -> &str { - "capturing" - } -} - -/// Deterministic k-means (Lloyd, strided init). Returns flattened centroids -/// `n_cells * hidden` and the assignment of each sample. -fn kmeans(samples: &[Vec], hidden: usize, n_cells: usize) -> (Vec, Vec) { - let n = samples.len(); - let c = n_cells.min(n.max(1)); - let mut centroids = vec![0.0f32; c * hidden]; - // Strided init — deterministic, spreads seeds across the sample order. - let stride = (n / c.max(1)).max(1); - for j in 0..c { - let src = &samples[(j * stride) % n]; - centroids[j * hidden..(j + 1) * hidden].copy_from_slice(&src[..hidden]); - } - let mut assign = vec![0usize; n]; - for _ in 0..KMEANS_ITERS { - // Assign. - for (i, s) in samples.iter().enumerate() { - let mut best = 0usize; - let mut bd = f32::INFINITY; - for j in 0..c { - let cj = ¢roids[j * hidden..(j + 1) * hidden]; - let mut d = 0.0f32; - for (a, b) in cj.iter().zip(s.iter()) { - let e = a - b; - d += e * e; - } - if d < bd { - bd = d; - best = j; - } - } - assign[i] = best; - } - // Update. - let mut sums = vec![0.0f32; c * hidden]; - let mut counts = vec![0usize; c]; - for (i, s) in samples.iter().enumerate() { - let j = assign[i]; - counts[j] += 1; - let dst = &mut sums[j * hidden..(j + 1) * hidden]; - for (d, v) in dst.iter_mut().zip(s.iter()) { - *d += v; - } - } - for j in 0..c { - if counts[j] > 0 { - let inv = 1.0 / counts[j] as f32; - for v in &mut sums[j * hidden..(j + 1) * hidden] { - *v *= inv; - } - centroids[j * hidden..(j + 1) * hidden] - .copy_from_slice(&sums[j * hidden..(j + 1) * hidden]); - } - } - } - (centroids, assign) -} - -fn static_importance_pool( - weights: &ModelWeights, - index: &larql_vindex::VectorIndex, - k: usize, -) -> Arc>> { - let probe = WalkFfn::new_unlimited(weights, index); - let per_layer = (0..weights.num_layers) - .map(|layer| { - let feats = index.num_features(layer); - let k = k.min(feats.max(1)); - match probe.down_row_norms_pub(layer) { - Some(norms) => { - let mut idx: Vec = (0..norms.len()).collect(); - idx.sort_unstable_by(|&a, &b| norms[b].total_cmp(&norms[a])); - idx.truncate(k); - idx - } - None => (0..k) - .map(|i| (i * (feats / k.max(1)).max(1)) % feats) - .collect(), - } - }) - .collect(); - Arc::new(per_layer) -} - -fn next_token_dist( - weights: &ModelWeights, - tok: &tokenizers::Tokenizer, - ids: &[u32], - ffn: &dyn FfnBackend, -) -> HashMap { - let r = predict_with_ffn(weights, tok, ids, usize::MAX, ffn); - r.token_ids - .into_iter() - .zip(r.predictions.into_iter().map(|(_, p)| p)) - .collect() -} - -/// KL(P‖Q) in bits. -fn kl_bits(p: &HashMap, q: &HashMap) -> f64 { - let eps = 1e-12; - let mut kl = 0.0; - for (&id, &pi) in p { - if pi <= 0.0 { - continue; - } - let qi = q.get(&id).copied().unwrap_or(0.0).max(eps); - kl += pi * (pi.max(eps) / qi).ln(); - } - kl / std::f64::consts::LN_2 -} - -fn stats(v: &mut [f64]) -> (f64, f64) { - let n = v.len().max(1) as f64; - let mean = v.iter().sum::() / n; - v.sort_by(|a, b| a.total_cmp(b)); - (mean, v[v.len() / 2]) -} - -/// Standard normal CDF (Abramowitz–Stegun erf approximation). -fn norm_cdf(z: f64) -> f64 { - let t = 1.0 / (1.0 + 0.2316419 * z.abs()); - let d = 0.3989422804014327 * (-z * z / 2.0).exp(); - let p = d - * t - * (0.319381530 - + t * (-0.356563782 + t * (1.781477937 + t * (-1.821255978 + t * 1.330274429)))); - if z >= 0.0 { - 1.0 - p - } else { - p - } -} - -/// Paired Wilcoxon signed-rank on `a - b` (does method `a` differ from `b`?). -/// Returns (median delta, z, two-sided p) via the tie-corrected normal -/// approximation — appropriate at n≈30. Negative median delta + small p means -/// `a` has lower KL than `b`. -fn wilcoxon(a: &[f64], b: &[f64]) -> (f64, f64, f64) { - let deltas: Vec = a - .iter() - .zip(b) - .map(|(x, y)| x - y) - .filter(|d| *d != 0.0) - .collect(); - let n = deltas.len(); - if n < 2 { - return (0.0, 0.0, 1.0); - } - let mut med: Vec = deltas.clone(); - med.sort_by(|x, y| x.total_cmp(y)); - let median = med[med.len() / 2]; - // Rank by |delta| with tie-averaging. - let mut idx: Vec = (0..n).collect(); - idx.sort_by(|&i, &j| deltas[i].abs().total_cmp(&deltas[j].abs())); - let mut ranks = vec![0.0f64; n]; - let mut i = 0; - let mut tie_term = 0.0f64; // Σ(t³−t) for variance correction - while i < n { - let mut k = i; - while k + 1 < n && deltas[idx[k + 1]].abs() == deltas[idx[i]].abs() { - k += 1; - } - let avg = ((i + k) as f64 / 2.0) + 1.0; // average of ranks i+1..=k+1 - for &p in &idx[i..=k] { - ranks[p] = avg; - } - let t = (k - i + 1) as f64; - tie_term += t * t * t - t; - i = k + 1; - } - let w_pos: f64 = (0..n).filter(|&i| deltas[i] > 0.0).map(|i| ranks[i]).sum(); - let nn = n as f64; - let mean = nn * (nn + 1.0) / 4.0; - let var = nn * (nn + 1.0) * (2.0 * nn + 1.0) / 24.0 - tie_term / 48.0; - if var <= 0.0 { - return (median, 0.0, 1.0); - } - let cc = if w_pos > mean { -0.5 } else { 0.5 }; // continuity correction - let z = (w_pos - mean + cc) / var.sqrt(); - let p = 2.0 * (1.0 - norm_cdf(z.abs())); - (median, z, p) -} - -fn main() { - let args: Vec = std::env::args().collect(); - let vindex = args - .get(1) - .cloned() - .unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".to_string()); - let dir = std::path::PathBuf::from(&vindex); - let mut cb = larql_vindex::SilentLoadCallbacks; - - eprintln!("Loading {vindex} ..."); - let mut weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index - .load_interleaved_kquant(&dir) - .expect("interleaved kquant"); - index.load_attn_kquant(&dir).expect("attn kquant"); - let _ = index.load_lm_head_kquant(&dir); - let _ = index.load_down_features_q4k(&dir); - let _ = index.load_down_features(&dir); - let _ = index.load_gate_vectors_q4(&dir); - let tok = load_tokenizer(&dir).expect("tokenizer"); - for layer in 0..weights.num_layers { - insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant layer"); - } - - let nl = weights.num_layers; - let hidden = weights.hidden_size; - let sparse_from = nl.saturating_sub(9); // 9-layer band - - let calib: [&str; 24] = [ - "The history of science is full of surprising discoveries that changed how we see the world.", - "She walked along the river at dawn, watching the mist rise off the cold water.", - "In economics, supply and demand determine the price of goods in a free market.", - "The recipe calls for two cups of flour, a pinch of salt, and three fresh eggs.", - "Astronomers detected a faint signal from a galaxy billions of light years away.", - "He repaired the old engine slowly, tightening each bolt with practiced care.", - "Democracy depends on an informed public and the peaceful transfer of power.", - "The children laughed as the puppy chased its tail around the sunny garden.", - "Modern computers store information as long sequences of ones and zeros.", - "The treaty was signed after months of difficult negotiation between the nations.", - "Photosynthesis converts sunlight, water, and carbon dioxide into sugar and oxygen.", - "The novel explores themes of memory, loss, and the passage of time.", - "Engineers tested the bridge under heavy load before opening it to traffic.", - "A balanced diet includes proteins, carbohydrates, fats, vitamins, and minerals.", - "The orchestra tuned their instruments as the audience settled into their seats.", - "Climate patterns are shifting, bringing hotter summers and wetter winters.", - "The detective examined the room carefully, noting every small detail.", - "Quantum mechanics describes the behavior of matter at the smallest scales.", - "They hiked for hours before reaching the summit and the breathtaking view.", - "The company reported strong earnings, lifting its stock price sharply.", - "Ancient traders carried silk, spices, and ideas across vast desert routes.", - "The teacher explained the theorem step by step until the class understood.", - "Rain fell steadily on the quiet town as the evening lights flickered on.", - "Vaccines train the immune system to recognize and fight specific diseases.", - ]; - let eval: [&str; 30] = [ - "The capital of France is", - "Water is made of hydrogen and", - "The opposite of hot is", - "The sun rises in the", - "The first president of the United States was", - "A group of lions is called a", - "The chemical symbol for gold is", - "The largest planet in the solar system is", - "Romeo and Juliet was written by", - "The speed of light is approximately", - "The capital of Japan is", - "Photosynthesis occurs in the", - "The square root of 64 is", - "The freezing point of water in Celsius is", - "The author of Pride and Prejudice is", - "An apple a day keeps the doctor", - "The currency of the United Kingdom is the", - "The tallest mountain on Earth is", - "DNA stands for", - "The capital of Italy is", - "The number of continents on Earth is", - "A baby dog is called a", - "The boiling point of water in Celsius is", - "The planet known as the Red Planet is", - "The longest river in the world is the", - "The inventor of the telephone was", - "The opposite of up is", - "The third planet from the sun is", - "The capital of Germany is", - "Two plus three equals", - ]; - - // ── 1. CALIBRATE ────────────────────────────────────────────────── - eprintln!( - "Calibrating on {} prompts (band = last {}/{nl}) ...", - calib.len(), - nl - sparse_from - ); - let capture = CapturingFfn::new( - WalkFfn::new_unlimited(&weights, &index), - &index, - nl, - sparse_from, - ); - for p in &calib { - let ids = tok.encode(*p, true).expect("encode").get_ids().to_vec(); - let _ = predict_with_ffn(&weights, &tok, &ids, 1, &capture); - } - let samples = capture.samples.into_inner(); - - // ── 2. BUILD CellRouter ─────────────────────────────────────────── - let mut centroids = vec![Vec::new(); nl]; - let mut n_cells = vec![0usize; nl]; - let mut pools = vec![Vec::new(); nl]; - let mut pool_sizes = Vec::new(); - for layer in sparse_from..nl { - let rows: Vec> = samples[layer].iter().map(|(r, _)| r.clone()).collect(); - if rows.is_empty() { - continue; - } - let (cents, assign) = kmeans(&rows, hidden, N_CELLS); - let c = cents.len() / hidden; - // Per cell: frequency-rank the gate-KNN features of its members. - let mut cell_pools: Vec> = Vec::with_capacity(c); - for cell in 0..c { - let mut freq: HashMap = HashMap::new(); - for (i, (_, topk)) in samples[layer].iter().enumerate() { - if assign[i] == cell { - for &f in topk { - *freq.entry(f).or_insert(0) += 1; - } - } - } - let mut fv: Vec<(usize, u32)> = freq.into_iter().collect(); - fv.sort_unstable_by(|a, b| b.1.cmp(&a.1).then(a.0.cmp(&b.0))); - fv.truncate(MAX_POOL); - let pool: Vec = fv.into_iter().map(|(f, _)| f).collect(); - pool_sizes.push(pool.len()); - cell_pools.push(pool); - } - centroids[layer] = cents; - n_cells[layer] = c; - pools[layer] = cell_pools; - } - let mean_pool = pool_sizes.iter().sum::() as f64 / pool_sizes.len().max(1) as f64; - let router = Arc::new(CellRouter { - centroids, - n_cells, - pools, - hidden, - }); - eprintln!( - "Built CellRouter: {N_CELLS} cells/layer, mean cell pool {:.0} feats ({:.1}% of {})", - mean_pool, - 100.0 * mean_pool / index.num_features(sparse_from).max(1) as f64, - index.num_features(sparse_from) - ); - - // ── 3. EVAL ─────────────────────────────────────────────────────── - // OOD sets, kept SEPARATE by sub-distribution — code / non-English / - // prose-dialogue behave nothing alike, so an aggregate p could be one - // category carrying it. Disaggregate to see which (if any) drives the - // cell-vs-gate edge. ~6 each → underpowered individually; read the - // median Δ + sign, not the p. - let ood_code: [&str; 6] = [ - "def add(a, b):\n return a +", - "for (int i = 0; i <", - "import numpy as", - "SELECT name FROM users WHERE id =", - "let mut x: Vec =", - "console.log(\"hello,", - ]; - let ood_intl: [&str; 6] = [ - "Bonjour, comment allez-", - "Hola, ¿cómo estás", - "Guten Tag, wie geht es", - "Ciao, come", - "Hallo, ik wil graag", - "Olá, tudo", - ]; - let ood_prose: [&str; 6] = [ - "\"How are you today?\" she", - "He turned and said, \"I can't believe", - "Once upon a time, there was a", - "Dear Sir or Madam, I am writing to", - "BREAKING: scientists announced today that", - "The patient presented with acute", - ]; - let pool_static = static_importance_pool(&weights, &index, MAX_POOL); - - // Returns per-prompt KL-to-dense for (gate, cell-full, cell-rankK, static). - let eval_set = |prompts: &[&str]| -> (Vec, Vec, Vec, Vec) { - let (mut kg, mut kc, mut kck, mut ks) = (Vec::new(), Vec::new(), Vec::new(), Vec::new()); - for p in prompts { - let ids = tok.encode(*p, true).expect("encode").get_ids().to_vec(); - let dense = next_token_dist( - &weights, - &tok, - &ids, - &WalkFfn::new_unlimited(&weights, &index), - ); - let gate = next_token_dist( - &weights, - &tok, - &ids, - &WalkFfn::from_config(&weights, &index, WalkFfnConfig::hybrid(nl, sparse_from, K)), - ); - let cell = next_token_dist( - &weights, - &tok, - &ids, - &WalkFfn::from_config( - &weights, - &index, - WalkFfnConfig::hybrid(nl, sparse_from, K).with_cell_router(router.clone()), - ), - ); - let cellk = next_token_dist( - &weights, - &tok, - &ids, - &WalkFfn::from_config( - &weights, - &index, - WalkFfnConfig::hybrid(nl, sparse_from, K) - .with_cell_router(router.clone()) - .with_rank_within_pool(true), - ), - ); - let stat = next_token_dist( - &weights, - &tok, - &ids, - &WalkFfn::from_config( - &weights, - &index, - WalkFfnConfig::hybrid(nl, sparse_from, K) - .with_pool_per_layer(pool_static.clone()), - ), - ); - kg.push(kl_bits(&dense, &gate)); - kc.push(kl_bits(&dense, &cell)); - kck.push(kl_bits(&dense, &cellk)); - ks.push(kl_bits(&dense, &stat)); - } - (kg, kc, kck, ks) - }; - - let report = |label: &str, kg: &[f64], kc: &[f64], kck: &[f64], ks: &[f64]| { - let (gm, gmd) = stats(&mut kg.to_vec()); - let (cm, cmd) = stats(&mut kc.to_vec()); - let (ckm, ckmd) = stats(&mut kck.to_vec()); - let (sm, smd) = stats(&mut ks.to_vec()); - // All KL is measured against DENSE (KL 0 = dense). gate-KNN is itself - // a lossy top-K truncation; lower = closer to dense. - println!( - "\n{label} (n={}) — KL-to-dense, bits (dense = 0)\n", - kg.len() - ); - println!("{:<26} {:>8} {:>8}", "router", "mean", "median"); - println!("{:<26} {gm:>8.3} {gmd:>8.3}", "gate-KNN"); - println!("{:<26} {cm:>8.3} {cmd:>8.3}", "cell-router full pool"); - println!("{:<26} {ckm:>8.3} {ckmd:>8.3}", "cell-router rank→K=512"); - println!("{:<26} {sm:>8.3} {smd:>8.3}", "static pool P=2048"); - // Paired Wilcoxon signed-rank on per-prompt deltas vs gate-KNN/static. - let (m1, z1, p1) = wilcoxon(kc, kg); - let (m2, z2, p2) = wilcoxon(kck, kg); - let (m3, z3, p3) = wilcoxon(kc, ks); - println!(" Wilcoxon (Δ = a − b, negative ⇒ a closer to dense):"); - println!(" cell-full vs gate-KNN : med Δ {m1:+.3} z {z1:+.2} p {p1:.4}"); - println!(" cell-rankK vs gate-KNN: med Δ {m2:+.3} z {z2:+.2} p {p2:.4}"); - println!(" cell-full vs static : med Δ {m3:+.3} z {z3:+.2} p {p3:.4}"); - }; - - println!( - "\nCell-router eval — sparse last {}/{nl}, K={K}, C={N_CELLS}, mean cell pool {:.0} feats ({:.1}% of {})", - nl - sparse_from, - mean_pool, - 100.0 * mean_pool / index.num_features(sparse_from).max(1) as f64, - index.num_features(sparse_from) - ); - let (kg, kc, kck, ks) = eval_set(&eval); - report("IN-DISTRIBUTION (factual completions)", &kg, &kc, &kck, &ks); - - // OOD disaggregated by sub-distribution + an aggregate (accumulate the - // per-category vectors so the aggregate is the same prompts pooled). - let (mut agg_g, mut agg_c, mut agg_ck, mut agg_s) = - (Vec::new(), Vec::new(), Vec::new(), Vec::new()); - for (label, set) in [ - ("OOD code", &ood_code[..]), - ("OOD non-English", &ood_intl[..]), - ("OOD prose/dialogue", &ood_prose[..]), - ] { - let (g, c, ck, s) = eval_set(set); - report(label, &g, &c, &ck, &s); - agg_g.extend(g); - agg_c.extend(c); - agg_ck.extend(ck); - agg_s.extend(s); - } - report("OOD aggregate", &agg_g, &agg_c, &agg_ck, &agg_s); - - // ── 4. DECODE AGREEMENT (accuracy half of the #23 pre-committed bar) ─ - // KL is a proxy; what shows in generations is whether the argmax flips. - // PRE-COMMITTED PASS BAR: full-pool top-1 agreement vs dense ≥ 90%. - // Teacher-forced on dense's OWN greedy stream (no cascade): generate a - // reference continuation with dense, then at each position score each - // variant's argmax on the same prefix against dense's pick. - let argmax = |d: &HashMap| -> u32 { - d.iter() - .max_by(|a, b| a.1.total_cmp(b.1)) - .map(|(i, _)| *i) - .unwrap_or(0) - }; - let gen_len = 20usize; - let seeds = [ - "The capital of France is", // in-dist - "def add(a, b):\n return a +", // OOD code - "Bonjour, comment allez-", // OOD non-English - "Once upon a time, there was a", // OOD prose - ]; - // Sweep band depth. The 9-layer band failed (gate-KNN itself ~60%), so - // the question is whether the shallow band (#20) is generation-viable. - // Static-importance is #20's shipping router; cell-router only exists at - // the calibrated 9-layer band. - let pool_imp_k = static_importance_pool(&weights, &index, K); // #20: top-K by ‖down_row‖ - println!( - "\nDecode top-1 agreement vs dense — teacher-forced, {} positions ({} seeds × {gen_len})", - seeds.len() * gen_len, - seeds.len() - ); - println!(" PRE-COMMITTED BAR: ≥ 90% (KL is a proxy; this is what shows in generations)"); - for depth in [4usize, 6, 9] { - let sf = nl.saturating_sub(depth); - let (mut a_gate, mut a_stat, mut a_cell, mut total) = (0usize, 0usize, 0usize, 0usize); - for s in &seeds { - let mut ids = tok.encode(*s, true).expect("encode").get_ids().to_vec(); - for _ in 0..gen_len { - let d = argmax(&next_token_dist( - &weights, - &tok, - &ids, - &WalkFfn::new_unlimited(&weights, &index), - )); - let g = argmax(&next_token_dist( - &weights, - &tok, - &ids, - &WalkFfn::from_config(&weights, &index, WalkFfnConfig::hybrid(nl, sf, K)), - )); - let st = argmax(&next_token_dist( - &weights, - &tok, - &ids, - &WalkFfn::from_config( - &weights, - &index, - WalkFfnConfig::hybrid(nl, sf, K) - .with_pool_per_layer(pool_imp_k.clone()) - .with_precomputed_routing(true), - ), - )); - a_gate += (g == d) as usize; - a_stat += (st == d) as usize; - if depth == 9 { - let c = argmax(&next_token_dist( - &weights, - &tok, - &ids, - &WalkFfn::from_config( - &weights, - &index, - WalkFfnConfig::hybrid(nl, sf, K).with_cell_router(router.clone()), - ), - )); - a_cell += (c == d) as usize; - } - total += 1; - ids.push(d); - } - } - let pct = |n: usize| 100.0 * n as f64 / total.max(1) as f64; - let pf = |n: usize| if pct(n) >= 90.0 { "PASS" } else { "FAIL" }; - if depth == 9 { - println!(" last {depth}/{nl}: gate-KNN {:.1}% {} static-imp(#20) {:.1}% {} cell-router {:.1}% {}", - pct(a_gate), pf(a_gate), pct(a_stat), pf(a_stat), pct(a_cell), pf(a_cell)); - } else { - println!( - " last {depth}/{nl}: gate-KNN {:.1}% {} static-imp(#20) {:.1}% {}", - pct(a_gate), - pf(a_gate), - pct(a_stat), - pf(a_stat) - ); - } - } -} diff --git a/crates/larql-inference/examples/walk_ffn_decode_timing.rs b/crates/larql-inference/examples/walk_ffn_decode_timing.rs deleted file mode 100644 index bc7eb6357..000000000 --- a/crates/larql-inference/examples/walk_ffn_decode_timing.rs +++ /dev/null @@ -1,147 +0,0 @@ -//! End-to-end decode-step timing (task #24 final test): net forward wall-time of -//! a gather-sparse band vs all-dense, on the real model with the sidecar loaded. -//! -//! PRE-COMMITTED BAR: net forward tok/s **> dense** (net-positive). The isolated -//! FFN gather is ~1.29× at K=4096, but it only touches the band's layers, so by -//! Amdahl the net is plausibly single-digit %. This measures the *full* forward -//! (all layers + attention + lm_head), where the per-layer `madvise(layer+1)` -//! prefetch is useful (sequential), not the single-layer-bench artifact. -//! -//! Reports forward µs and top-1 agreement vs dense for: dense, gather last-4 -//! (the #20 shippable static band), gather last-9. -//! -//! Usage: `cargo run --release --example walk_ffn_decode_timing -- [VINDEX_DIR]` - -use larql_inference::vindex::{insert_q4k_layer_tensors_resident, WalkFfn, WalkFfnConfig}; -use larql_inference::{load_tokenizer, predict_with_ffn}; -use larql_models::ModelWeights; -use std::sync::Arc; -use std::time::Instant; - -const K: usize = 4096; // faithful K (in-dist 4-layer band clears 90% agreement) - -fn static_importance_pool( - weights: &ModelWeights, - index: &larql_vindex::VectorIndex, - k: usize, -) -> Arc>> { - let probe = WalkFfn::new_unlimited(weights, index); - let per_layer = (0..weights.num_layers) - .map(|layer| { - let feats = index.num_features(layer); - let k = k.min(feats.max(1)); - match probe.down_row_norms_pub(layer) { - Some(norms) => { - let mut idx: Vec = (0..norms.len()).collect(); - idx.sort_unstable_by(|&a, &b| norms[b].total_cmp(&norms[a])); - idx.truncate(k); - idx - } - None => (0..k) - .map(|i| (i * (feats / k.max(1)).max(1)) % feats) - .collect(), - } - }) - .collect(); - Arc::new(per_layer) -} - -fn main() { - let args: Vec = std::env::args().collect(); - let vindex = args - .get(1) - .cloned() - .unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".to_string()); - let dir = std::path::PathBuf::from(&vindex); - let mut cb = larql_vindex::SilentLoadCallbacks; - - eprintln!("Loading {vindex} ..."); - let mut weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index.load_interleaved_kquant(&dir).expect("interleaved"); - index.load_attn_kquant(&dir).expect("attn"); - let _ = index.load_lm_head_kquant(&dir); - index.load_down_features_q4k(&dir).expect("down sidecar"); - let _ = index.load_gate_vectors_q4(&dir); - assert!( - index.has_down_features_kquant(), - "feature-major down sidecar required (run build_down_features_q4k)" - ); - let tok = load_tokenizer(&dir).expect("tok"); - for layer in 0..weights.num_layers { - insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant attn"); - } - - let nl = weights.num_layers; - let pool = static_importance_pool(&weights, &index, K); - // DECODE shape: a single-token forward (seq_len=1) — the real per-token - // decode step where the FFN is a matvec and the gather wins. A multi-token - // prompt would measure prefill (batched BLAS gemm), where dense wins. - let full = tok - .encode("The capital of France is the city of", true) - .expect("enc") - .get_ids() - .to_vec(); - let ids = vec![*full.last().unwrap()]; - let iters = 30usize; - - // `run` builds the WalkFfn and returns the predicted token id; timing it - // avoids naming WalkFfn's lifetime in a closure return. - let time_cfg = |run: &dyn Fn() -> u32| -> (f64, u32) { - for _ in 0..5 { - let _ = run(); - } - let t = Instant::now(); - let mut last = 0u32; - for _ in 0..iters { - last = run(); - } - (t.elapsed().as_micros() as f64 / iters as f64, last) - }; - let predict_tok = |ffn: &dyn larql_inference::ffn::FfnBackend| -> u32 { - predict_with_ffn(&weights, &tok, &ids, 1, ffn) - .token_ids - .first() - .copied() - .unwrap_or(0) - }; - - println!( - "\nDecode-step forward timing — {nl} layers, K={K}, prompt {} tokens, {iters} iters\n", - ids.len() - ); - let (dense_us, dense_tok) = - time_cfg(&|| predict_tok(&WalkFfn::new_unlimited(&weights, &index))); - println!( - " dense (all layers) {dense_us:>9.0} µs 1.00× (ref token {dense_tok})" - ); - - for band in [4usize, 9] { - let sf = nl.saturating_sub(band); - let p = pool.clone(); - let (us, tok_id) = time_cfg(&|| { - predict_tok(&WalkFfn::from_config( - &weights, - &index, - WalkFfnConfig::hybrid(nl, sf, K) - .with_pool_per_layer(p.clone()) - .with_precomputed_routing(true), - )) - }); - let ratio = dense_us / us; - let agree = if tok_id == dense_tok { - "top-1 ✓" - } else { - "top-1 ✗" - }; - let verdict = if ratio > 1.0 { - "PASS (net>dense)" - } else { - "FAIL" - }; - println!( - " gather last-{band:<2} static {us:>9.0} µs {ratio:.3}× {agree} {verdict}" - ); - } - println!("\n (Bar: net forward ratio > 1.0. K={K} chosen for faithfulness; 4-layer band is the #20 shippable shape.)"); -} diff --git a/crates/larql-inference/examples/walk_ffn_delta_walk.rs b/crates/larql-inference/examples/walk_ffn_delta_walk.rs deleted file mode 100644 index 61430a723..000000000 --- a/crates/larql-inference/examples/walk_ffn_delta_walk.rs +++ /dev/null @@ -1,227 +0,0 @@ -//! #28 stage 1 — delta-walk falsification probe (NO kernel, just FFN evals). -//! -//! #27 found the token-to-token FFN-input residual delta is ~22-dim — BUT that -//! equals the intrinsic STATE dim, which is a warning: a delta spanning the same -//! manifold as the state is a full-amplitude excursion, not a thin perturbation a -//! fixed Jacobian linearizes. Low-rank ≠ small. #27 measured rank, not amplitude. -//! -//! This measures the two things that actually gate delta-walk, before any kernel: -//! (a) AMPLITUDE ‖δ‖/‖base‖ per zone — is the move small relative to base? -//! (b) LINEARIZATION ERROR ‖f(base+δ) − (f(base)+Jδ)‖ / ‖f(base+δ)‖, where -//! f = the layer's FFN (pure fn of its post-attn-norm input) and Jδ is the -//! FINITE-DIFFERENCE Jacobian-vector product (Jδ ≈ (f(base+εδ)−f(base))/ε). -//! This is the FULL true Jacobian — if it can't reproduce the FFN's action -//! on the highway, no low-rank approximation can, and delta-walk is dead. -//! -//! Targets the FFN-INPUT residual (post-attention-norm), captured via a recording -//! FfnBackend — NOT #27's layer-input residual. Per-zone distribution -//! (median/p90/worst), worst-token tail kept (a scheme catastrophic on 10% of -//! steps is a drift generator). -//! -//! KILL (pre-registered): highway ‖δ‖/‖base‖ large (≳20%) OR lin-error large -//! (≳15%) ⇒ delta-walk dead, no kernel built. Small both ⇒ stage 2 (cached -//! low-rank J + refresh-rate vs Jaccard, on LIVE decode). -//! -//! Usage: `cargo run --release --example walk_ffn_delta_walk -- [VINDEX]` - -use larql_inference::ffn::FfnBackend; -use larql_inference::load_tokenizer; -use larql_inference::research::predict_with_ffn_trace; -use larql_inference::vindex::{insert_q4k_layer_tensors_resident, WalkFfn}; -use ndarray::{Array2, Axis}; -use std::cell::RefCell; - -const EPS: f32 = 1e-2; // finite-difference step for the JVP - -/// Records each layer's LAST-POSITION FFN input (the post-attn-norm residual the -/// FFN actually sees), then delegates to a dense WalkFfn. -struct CapturingFfn<'a> { - inner: WalkFfn<'a>, - cap: RefCell>>>, -} -impl<'a> CapturingFfn<'a> { - fn new(inner: WalkFfn<'a>, nl: usize) -> Self { - Self { - inner, - cap: RefCell::new(vec![None; nl]), - } - } - fn reset(&self) { - for c in self.cap.borrow_mut().iter_mut() { - *c = None; - } - } - fn rec(&self, layer: usize, x: &Array2) { - let last = x.shape()[0] - 1; - self.cap.borrow_mut()[layer] = Some(x.row(last).to_vec()); - } -} -impl FfnBackend for CapturingFfn<'_> { - fn forward(&self, layer: usize, x: &Array2) -> Array2 { - self.rec(layer, x); - self.inner.forward(layer, x) - } - fn forward_observed( - &self, - layer: usize, - x: &Array2, - ) -> (Array2, larql_inference::ffn::FfnActivations) { - self.rec(layer, x); - self.inner.forward_observed(layer, x) - } - fn name(&self) -> &str { - "capturing" - } -} - -fn norm(v: &[f32]) -> f64 { - v.iter() - .map(|&x| (x as f64) * (x as f64)) - .sum::() - .sqrt() -} - -fn pctile(v: &mut [f64], q: f64) -> f64 { - if v.is_empty() { - return f64::NAN; - } - v.sort_by(|a, b| a.total_cmp(b)); - v[(((v.len() - 1) as f64) * q).round() as usize] -} - -fn zone(l: usize) -> usize { - match l { - 0..=4 => 0, - 5..=20 => 1, - 21..=29 => 2, - _ => 3, - } -} - -fn main() { - let args: Vec = std::env::args().collect(); - let vindex = args - .get(1) - .cloned() - .unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".to_string()); - let dir = std::path::PathBuf::from(&vindex); - let mut cb = larql_vindex::SilentLoadCallbacks; - eprintln!("Loading {vindex} ..."); - let mut weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index.load_interleaved_kquant(&dir).expect("interleaved"); - index.load_attn_kquant(&dir).expect("attn"); - let _ = index.load_lm_head_kquant(&dir); - let _ = index.load_gate_vectors_q4(&dir); - let tok = load_tokenizer(&dir).expect("tok"); - for layer in 0..weights.num_layers { - insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant attn"); - } - let nl = weights.num_layers; - let hidden = weights.hidden_size; - - let passages = [ - "The expedition had been planned for years, but nothing prepared them for the silence of the ice that morning, and the captain wrote that the cold seemed to have a will of its own.", - "She had always believed that cities were built from ambition, but walking the old quarter at dusk she understood they were built from compromise, one stubborn refusal at a time.", - "Economists argue about the cause, yet the pattern repeats: cheap credit, a frenzy of building, a sudden loss of nerve, and then the long quiet years of paying it all back.", - "Light from the distant galaxy had travelled for billions of years to reach the telescope, carrying news of an explosion that had happened before the sun itself was born.", - ]; - - let dense = WalkFfn::new_unlimited(&weights, &index); // pure FFN evaluator - let ffn = |layer: usize, x: &[f32]| -> Vec { - let m = Array2::from_shape_vec((1, hidden), x.to_vec()).unwrap(); - dense.forward(layer, &m).row(0).to_vec() - }; - - let mut amp_by_layer: Vec> = vec![Vec::new(); nl]; - let mut lin_by_layer: Vec> = vec![Vec::new(); nl]; - - for (pi, p) in passages.iter().enumerate() { - let ids = tok.encode(*p, true).expect("enc").get_ids().to_vec(); - let n = ids.len().min(36); - eprintln!(" passage {}/{} ({n} tokens) ...", pi + 1, passages.len()); - let cap = CapturingFfn::new(WalkFfn::new_unlimited(&weights, &index), nl); - let mut prev: Option>>> = None; - for i in 3..n { - cap.reset(); - let _ = predict_with_ffn_trace(&weights, &tok, &ids[..=i], 1, &cap); - let cur = cap.cap.borrow().clone(); - if let Some(pr) = &prev { - for l in 0..nl { - let (b, c) = match (&pr[l], &cur[l]) { - (Some(b), Some(c)) if b.len() == hidden && c.len() == hidden => (b, c), - _ => continue, - }; - let delta: Vec = b.iter().zip(c).map(|(&x, &y)| y - x).collect(); - let nb = norm(b); - let nd = norm(&delta); - if nb < 1e-6 { - continue; - } - amp_by_layer[l].push(nd / nb); - // FFN eval: f(base), f(base+δ)=f(c), f(base+εδ) - let f_base = ffn(l, b); - let f_full = ffn(l, c); - let beps: Vec = b.iter().zip(&delta).map(|(&x, &d)| x + EPS * d).collect(); - let f_eps = ffn(l, &beps); - // Jδ = (f_eps - f_base)/EPS ; lin_pred = f_base + Jδ - // err = ‖f_full - lin_pred‖ / ‖f_full‖ - let mut num = 0.0f64; - let mut den = 0.0f64; - for k in 0..hidden { - let jd = (f_eps[k] - f_base[k]) / EPS; - let lin = f_base[k] + jd; - let e = f_full[k] - lin; - num += (e as f64) * (e as f64); - den += (f_full[k] as f64) * (f_full[k] as f64); - } - if den > 1e-12 { - lin_by_layer[l].push(num.sqrt() / den.sqrt()); - } - } - } - prev = Some(cur); - } - } - let _ = Axis(0); - - let zn = [ - "pre-commit L0-4", - "highway L5-20", - "retrieval L21-29", - "format L30-33", - ]; - println!( - "\n#28 stage 1 — delta-walk falsification (amplitude + full-Jacobian linearization)\n" - ); - println!( - "{:<20} {:>22} {:>26}", - "zone", "‖δ‖/‖base‖ med/p90/worst", "lin-error med/p90/worst" - ); - for (z, zname) in zn.iter().enumerate() { - let layers: Vec = (0..nl).filter(|&l| zone(l) == z).collect(); - let mut amp: Vec = layers - .iter() - .flat_map(|&l| amp_by_layer[l].clone()) - .collect(); - let mut lin: Vec = layers - .iter() - .flat_map(|&l| lin_by_layer[l].clone()) - .collect(); - let (am, ap, aw) = ( - pctile(&mut amp.clone(), 0.50), - pctile(&mut amp.clone(), 0.90), - pctile(&mut amp, 1.0), - ); - let (lm, lp, lw) = ( - pctile(&mut lin.clone(), 0.50), - pctile(&mut lin.clone(), 0.90), - pctile(&mut lin, 1.0), - ); - println!( - "{:<20} {:>6.3}/{:>6.3}/{:>6.3} {:>8.3}/{:>7.3}/{:>7.3}", - zname, am, ap, aw, lm, lp, lw - ); - } - println!("\n KILL (pre-registered): highway ‖δ‖/‖base‖ ≳0.20 OR lin-error ≳0.15 ⇒ delta-walk dead.\n worst = max over step-pairs (catastrophic on 10% of steps = drift generator).\n lin-error here uses the FULL true Jacobian — a low-rank approx can only be worse."); -} diff --git a/crates/larql-inference/examples/walk_ffn_drift.rs b/crates/larql-inference/examples/walk_ffn_drift.rs deleted file mode 100644 index 3821127ec..000000000 --- a/crates/larql-inference/examples/walk_ffn_drift.rs +++ /dev/null @@ -1,193 +0,0 @@ -//! Q3 generation drift — the de-risk single-step KL can't see (task #26). -//! -//! Single-step KL 0.052 / 100% top-1 says Q3 is fine per token. But a small -//! per-step flip rate compounds over a generation (depth-compounding, applied to -//! the sequence axis). This greedy-decodes dense (sim Q4) vs Q3 and reports -//! **first-divergence position** and **sequence exact-match** — if Q3 drifts, -//! divergence is early. Catches the failure for a script instead of a format. -//! -//! Both arms use the same block-wise simulated quantiser (Q4 vs Q3), so the only -//! variable is 4→3 bits — no real-vs-sim confound. -//! -//! Usage: `cargo run --release --example walk_ffn_drift -- [VINDEX]` - -use larql_inference::ffn::FfnBackend; -use larql_inference::vindex::{insert_q4k_layer_tensors_resident, WalkFfn}; -use larql_inference::{load_tokenizer, predict_with_ffn}; -use larql_models::ModelWeights; -use ndarray::{Array1, Array2}; - -fn requant_row(row: &mut [f32], bits: u32) { - const BLK: usize = 32; - for blk in row.chunks_mut(BLK) { - let maxabs = blk.iter().fold(0.0f32, |m, &v| m.max(v.abs())); - if maxabs == 0.0 { - continue; - } - let levels = ((1u32 << (bits - 1)) - 1) as f32; - let scale = maxabs / levels; - for v in blk.iter_mut() { - *v = (*v / scale).round().clamp(-levels, levels) * scale; - } - } -} - -struct GradedFfn { - gate: Vec>, - up: Vec>, - down: Vec>, -} -impl FfnBackend for GradedFfn { - fn forward(&self, layer: usize, x: &Array2) -> Array2 { - let (g, u, d) = (&self.gate[layer], &self.up[layer], &self.down[layer]); - let hidden = x.shape()[1]; - let mut out = Array2::::zeros((x.shape()[0], hidden)); - for (s, xr) in x.rows().into_iter().enumerate() { - let xr = xr.to_owned(); - let gs = g.dot(&xr); - let us = u.dot(&xr); - let act: Array1 = gs - .iter() - .zip(us.iter()) - .map(|(&gg, &uu)| larql_inference::ffn::gelu_tanh(gg) * uu) - .collect(); - out.row_mut(s).assign(&act.dot(d)); - } - out - } - // forward_observed: trait default (Absent) — this arm computes no - // observable activation tensor and must not fabricate one. - fn name(&self) -> &str { - "graded" - } -} - -fn build_uniform( - weights: &ModelWeights, - index: &larql_vindex::VectorIndex, - bits: u32, -) -> GradedFfn { - let (nl, hidden) = (weights.num_layers, weights.hidden_size); - let (mut gate, mut up, mut down) = (Vec::new(), Vec::new(), Vec::new()); - for layer in 0..nl { - let inter = index.num_features(layer); - for (comp, store) in [(0usize, &mut gate), (1, &mut up), (2, &mut down)] { - let w = index.kquant_ffn_layer(layer, comp).expect("f32 comp"); - let mut m = Array2::::zeros((inter, hidden)); - for f in 0..inter { - let mut r = w[f * hidden..(f + 1) * hidden].to_vec(); - requant_row(&mut r, bits); - m.row_mut(f).assign(&Array1::from(r)); - } - store.push(m); - } - } - GradedFfn { gate, up, down } -} - -fn greedy_gen( - weights: &ModelWeights, - tok: &tokenizers::Tokenizer, - prompt: &[u32], - ffn: &dyn FfnBackend, - n: usize, -) -> Vec { - let mut ids = prompt.to_vec(); - let mut gen = Vec::with_capacity(n); - for _ in 0..n { - let t = predict_with_ffn(weights, tok, &ids, 1, ffn) - .token_ids - .first() - .copied() - .unwrap_or(0); - gen.push(t); - ids.push(t); - } - gen -} - -fn main() { - let args: Vec = std::env::args().collect(); - let vindex = args - .get(1) - .cloned() - .unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".to_string()); - let dir = std::path::PathBuf::from(&vindex); - let mut cb = larql_vindex::SilentLoadCallbacks; - eprintln!("Loading {vindex} ..."); - let mut weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index.load_interleaved_kquant(&dir).expect("interleaved"); - index.load_attn_kquant(&dir).expect("attn"); - let _ = index.load_lm_head_kquant(&dir); - let tok = load_tokenizer(&dir).expect("tok"); - for layer in 0..weights.num_layers { - insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant attn"); - } - let _ = WalkFfn::new_unlimited(&weights, &index); // (warms down-norm cache path) - - let prompts = [ - "The capital of France is", - "Once upon a time, there was a", - "def fibonacci(n):", - "The mitochondria is the", - "In 1969, humanity first", - "To make a good espresso, you", - "Bonjour, je voudrais", - "The three laws of motion are", - "She opened the door and", - "Climate change is driven by", - ]; - let n = 32usize; - - eprintln!("Building sim-Q4 and Q3 FFNs (f32, one-time) ..."); - let q4 = build_uniform(&weights, &index, 4); - let q3 = build_uniform(&weights, &index, 3); - - println!( - "\nQ3 generation drift vs sim-Q4 — greedy, {n} tokens, {} prompts\n", - prompts.len() - ); - println!("{:<34} {:>9} {:>9}", "prompt", "first-div", "exact?"); - let (mut div_sum, mut exact, mut total_flip, mut total_tok) = (0usize, 0usize, 0usize, 0usize); - for p in &prompts { - let pid = tok.encode(*p, true).expect("enc").get_ids().to_vec(); - let g4 = greedy_gen(&weights, &tok, &pid, &q4, n); - let g3 = greedy_gen(&weights, &tok, &pid, &q3, n); - let first_div = (0..n).find(|&i| g4[i] != g3[i]).unwrap_or(n); - let is_exact = first_div == n; - // Per-step flip rate (teacher-forced on Q4's stream): Q3's argmax on - // Q4's prefix vs Q4's token — the per-step error that compounds. - let mut ids = pid.clone(); - let mut flips = 0usize; - for > in g4.iter().take(n) { - let q3t = predict_with_ffn(&weights, &tok, &ids, 1, &q3) - .token_ids - .first() - .copied() - .unwrap_or(0); - if q3t != gt { - flips += 1; - } - ids.push(gt); - } - div_sum += first_div; - exact += is_exact as usize; - total_flip += flips; - total_tok += n; - let label: String = p.chars().take(32).collect(); - println!( - "{label:<34} {first_div:>9} {:>9}", - if is_exact { "yes" } else { "no" } - ); - } - let np = prompts.len() as f64; - println!( - "\n mean first-divergence {:.1}/{n}; exact full-match {}/{}; per-step flip rate {:.1}%", - div_sum as f64 / np, - exact, - prompts.len(), - 100.0 * total_flip as f64 / total_tok.max(1) as f64 - ); - println!(" (Bar: low flip rate + late/no divergence ⇒ Q3 is drift-safe to build the format.)"); -} diff --git a/crates/larql-inference/examples/walk_ffn_gather_gemm.rs b/crates/larql-inference/examples/walk_ffn_gather_gemm.rs deleted file mode 100644 index 87f406f81..000000000 --- a/crates/larql-inference/examples/walk_ffn_gather_gemm.rs +++ /dev/null @@ -1,217 +0,0 @@ -//! Gather→gemm sparse-FFN kernel test (task #24). -//! -//! The decode falsification (#23) showed token faithfulness needs K≈4096 (40% of -//! 10240 feats). At 40% of dense's FLOPs that *should* be ~2.5× faster than -//! dense — but the current scattered per-row walk has ~4× per-row overhead, so -//! K=4096 lands slower than dense. Hypothesis: gathering the K selected rows into -//! contiguous buffers and running a BLAS gemv realizes the FLOP saving. -//! -//! Honest premise: both paths start from Q4K bytes (NO full-layer f32 cache — -//! that would defeat sparsity-for-memory). Dense does the full Q4K matvec; -//! gather-gemm dequantizes ONLY the K selected rows → contiguous f32 → gemv. -//! -//! Three timings per K: dense (WalkFfn full), scattered cheap-route (current -//! sparse path), gather-gemm (this). -//! -//! Usage: `cargo run --release --example walk_ffn_gather_gemm -- [VINDEX_DIR]` - -use larql_inference::ffn::FfnBackend; -use larql_inference::vindex::{WalkFfn, WalkFfnConfig}; -use ndarray::{Array1, Axis}; -use rayon::prelude::*; -use std::sync::Arc; -use std::time::Instant; - -fn main() { - let args: Vec = std::env::args().collect(); - let vindex = args - .get(1) - .cloned() - .unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".to_string()); - let dir = std::path::PathBuf::from(&vindex); - let mut cb = larql_vindex::SilentLoadCallbacks; - eprintln!("Loading {vindex} ..."); - let weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index.load_interleaved_kquant(&dir).expect("interleaved"); - // Load the feature-major Q4K down sidecar (built by - // `build_down_features_q4k`) — its presence makes the WIRED `walk_ffn_sparse` - // ("scattered cheap-route" row) fire the gather fast path. NOT loading - // native f32 up/down (that path uses the f32 down kernel instead). - let sidecar = index.load_down_features_q4k(&dir).is_ok() && index.has_down_features_kquant(); - eprintln!("feature-major Q4K down sidecar loaded: {sidecar}"); - let _ = index.load_gate_vectors_q4(&dir); - - let hidden = weights.hidden_size; - let layer = weights.num_layers / 2; - let feats = index.num_features(layer); - let iters = 200usize; - - // Raw Q4K bytes for gate/up/down at this layer (no f32 materialization). - let slices = index - .interleaved_kquant_layer_data(layer) - .expect("interleaved Q4K layer bytes"); - let info = larql_vindex::quant::registry::lookup(slices[0].1).expect("registry"); - let dq = info.dequantize; - let bpr = info.bytes_per_row(hidden).expect("bytes_per_row"); - let gate_b = slices[0].0; - let up_b = slices[1].0; - let _ = slices[2].0; // interleaved down is transposed — not used (see down_q4k_fm) - - let x = Array1::from_shape_fn(hidden, |j| ((j as f32) * 0.013).sin() * 0.1); - let x2 = x.clone().insert_axis(Axis(0)); // (1, hidden) - - println!("\nGather→gemm FFN microbench — layer {layer}, {feats} feats, hidden {hidden}, {iters} iters\n"); - - // ── Dense baseline (full Q4K matvec, no f32 cache) ── - let dense = WalkFfn::new_unlimited(&weights, &index); - for _ in 0..20 { - let _ = dense.forward(layer, &x2); - } - let t = Instant::now(); - for _ in 0..iters { - let _ = dense.forward(layer, &x2); - } - let dense_us = t.elapsed().as_micros() as f64 / iters as f64; - println!(" dense (full Q4K matvec) {dense_us:>9.1} µs/call 1.00×"); - - // ── Build the FEATURE-MAJOR Q4K down in-memory (the down_features_q4k - // sidecar contents). The interleaved down is transposed [hidden×inter], - // so per-feature gather is wrong off it. `kquant_ffn_layer(layer,2)` - // dequant+transposes to feature-major f32 [inter×hidden]; re-quantise each - // feature row to Q4K → a gatherable feature-major down. (One-time at index - // build in production; here it's setup, not timed.) - let down_fm_f32 = index - .kquant_ffn_layer(layer, 2) - .expect("feature-major f32 down (kquant_ffn_layer component 2)"); - let q4k = larql_vindex::quant::registry::lookup("Q4_K").expect("Q4_K"); - let dbpr = q4k.bytes_per_row(hidden).expect("q4k bpr"); // == bpr (hidden elems) - let down_sa_q4k = q4k.row_scaled_add.expect("q4k row_scaled_add"); - let mut down_q4k_fm = vec![0u8; feats * dbpr]; - for f in 0..feats { - let row = &down_fm_f32[f * hidden..(f + 1) * hidden]; - let qb = larql_compute::cpu::ops::q4_common::quantize_q4_k(row); - down_q4k_fm[f * dbpr..(f + 1) * dbpr].copy_from_slice(&qb[..dbpr]); - } - - for k in [2048usize, 4096] { - let pool: Vec = (0..k).map(|i| (i * (feats / k).max(1)) % feats).collect(); - let pct = 100.0 * k as f64 / feats as f64; - - // ── Scattered cheap-route (current sparse path) ── - let cfg = WalkFfnConfig::sparse(weights.num_layers, k) - .with_pool_per_layer(Arc::new(vec![pool.clone(); weights.num_layers])) - .with_precomputed_routing(true); - let scat = WalkFfn::from_config(&weights, &index, cfg); - for _ in 0..20 { - let _ = scat.forward(layer, &x2); - } - let t = Instant::now(); - for _ in 0..iters { - let _ = scat.forward(layer, &x2); - } - let scat_us = t.elapsed().as_micros() as f64 / iters as f64; - - // ── Gather Q4K contiguous + fused kernel — CORRECT down (feature-major - // Q4K from `down_q4k_fm`). gate/up from interleaved (feature-major), - // down from the re-quantised sidecar buffer. No f32 materialisation in - // the hot loop. - let row_dot = info.row_dot.expect("row_dot"); - let xs = x.as_slice().unwrap(); - let mut gg = vec![0u8; k * bpr]; - let mut gu = vec![0u8; k * bpr]; - let mut gd = vec![0u8; k * dbpr]; - let nthreads = rayon::current_num_threads().max(1); - let chunk = k.div_ceil(nthreads); - let gather_q4k = |gg: &mut [u8], gu: &mut [u8], gd: &mut [u8]| -> Vec { - for (i, &p) in pool.iter().enumerate() { - gg[i * bpr..(i + 1) * bpr].copy_from_slice(&gate_b[p * bpr..(p + 1) * bpr]); - gu[i * bpr..(i + 1) * bpr].copy_from_slice(&up_b[p * bpr..(p + 1) * bpr]); - gd[i * dbpr..(i + 1) * dbpr] - .copy_from_slice(&down_q4k_fm[p * dbpr..(p + 1) * dbpr]); - } - let gate_s: Vec = (0..k) - .into_par_iter() - .map(|i| row_dot(&gg[i * bpr..(i + 1) * bpr], xs).unwrap_or(0.0)) - .collect(); - let up_s: Vec = (0..k) - .into_par_iter() - .map(|i| row_dot(&gu[i * bpr..(i + 1) * bpr], xs).unwrap_or(0.0)) - .collect(); - let act: Vec = gate_s - .iter() - .zip(&up_s) - .map(|(&g, &u)| (g / (1.0 + (-g).exp())) * u) - .collect(); - let partials: Vec> = (0..k) - .collect::>() - .par_chunks(chunk) - .map(|ch| { - let mut part = vec![0.0f32; hidden]; - for &i in ch { - if act[i].abs() > 1e-10 { - let _ = down_sa_q4k(&gd[i * dbpr..(i + 1) * dbpr], act[i], &mut part); - } - } - part - }) - .collect(); - let mut out = vec![0.0f32; hidden]; - for pp in &partials { - for (o, v) in out.iter_mut().zip(pp) { - *o += v; - } - } - out - }; - - // Correctness vs an f32 reference (same gate/up, but f32 feature-major - // down) — bounds the Q4K-down quantisation error, and confirms the - // gather indexing is right (not the transposed-down bug). - let g_out = gather_q4k(&mut gg, &mut gu, &mut gd); - let mut ref_out = vec![0.0f32; hidden]; - { - let gate_s: Vec = (0..k) - .map(|i| row_dot(&gg[i * bpr..(i + 1) * bpr], xs).unwrap_or(0.0)) - .collect(); - let up_s: Vec = (0..k) - .map(|i| row_dot(&gu[i * bpr..(i + 1) * bpr], xs).unwrap_or(0.0)) - .collect(); - for (i, &p) in pool.iter().enumerate() { - let act = (gate_s[i] / (1.0 + (-gate_s[i]).exp())) * up_s[i]; - let drow = &down_fm_f32[p * hidden..(p + 1) * hidden]; - for (o, &d) in ref_out.iter_mut().zip(drow) { - *o += act * d; - } - } - } - let max_abs: f32 = g_out - .iter() - .zip(&ref_out) - .map(|(&a, &b)| (a - b).abs()) - .fold(0.0, f32::max); - let ref_norm: f32 = ref_out.iter().map(|v| v * v).sum::().sqrt().max(1e-6); - - for _ in 0..20 { - let _ = gather_q4k(&mut gg, &mut gu, &mut gd); - } - let t = Instant::now(); - for _ in 0..iters { - let _ = gather_q4k(&mut gg, &mut gu, &mut gd); - } - let gg_us = t.elapsed().as_micros() as f64 / iters as f64; - - println!("\n K={k} ({pct:.0}%):"); - println!( - " scattered cheap-route {scat_us:>9.1} µs/call {:.2}× vs dense", - dense_us / scat_us - ); - println!( - " gather Q4K (correct down) {gg_us:>9.1} µs/call {:.2}× vs dense |err|max/‖ref‖ = {:.2e}", - dense_us / gg_us, - max_abs / ref_norm - ); - } - let _ = dq; // f32 dequant path retired (alloc-dominated); see git history - println!(); -} diff --git a/crates/larql-inference/examples/walk_ffn_graded_precision.rs b/crates/larql-inference/examples/walk_ffn_graded_precision.rs deleted file mode 100644 index a1a8606a9..000000000 --- a/crates/larql-inference/examples/walk_ffn_graded_precision.rs +++ /dev/null @@ -1,316 +0,0 @@ -//! Graded-precision FFN — the "fewer bytes per feature" lever (no sparsity). -//! -//! The K-sweep proved dropping low-importance features is catastrophic (their -//! *absence* is the error). Graded precision keeps ALL features but spends bits -//! by importance: high-‖down_row‖ head at 4 bits, low-norm tail at fewer. The -//! KL cost of *approximating* a low-norm feature should be far below the KL cost -//! of *zeroing* it — so this buys bandwidth without re-opening faithfulness. -//! -//! Per layer: rank features by ‖down_row‖; quantise each feature's gate/up/down -//! rows to its assigned bit-width (per-row symmetric). Measure KL vs the f32 -//! reference + top-1 agreement, against avg bits/feature (the bandwidth proxy). -//! -//! Usage: `cargo run --release --example walk_ffn_graded_precision -- [VINDEX]` - -use larql_inference::ffn::FfnBackend; -use larql_inference::vindex::{insert_q4k_layer_tensors_resident, WalkFfn}; -use larql_inference::{load_tokenizer, predict_with_ffn}; -use larql_models::ModelWeights; -use ndarray::{Array1, Array2}; -use std::collections::HashMap; - -/// Block-wise symmetric requantise-dequantise to `bits` (simulates b-bit/element -/// with per-block scales — the structure real K-quants use, so low bits aren't -/// unfairly destroyed by one outlier setting a whole-row scale). bits>=16 = -/// passthrough; bits==1 = sign × per-block mean|·|. Block = 32 elements. -fn requant_row(row: &mut [f32], bits: u32) { - if bits >= 16 { - return; - } - const BLK: usize = 32; - for blk in row.chunks_mut(BLK) { - let maxabs = blk.iter().fold(0.0f32, |m, &v| m.max(v.abs())); - if maxabs == 0.0 { - continue; - } - if bits == 1 { - let mean: f32 = blk.iter().map(|v| v.abs()).sum::() / blk.len() as f32; - for v in blk.iter_mut() { - *v = if *v >= 0.0 { mean } else { -mean }; - } - continue; - } - let levels = ((1u32 << (bits - 1)) - 1) as f32; // b=2→1 (ternary), b=4→7 - let scale = maxabs / levels; - for v in blk.iter_mut() { - *v = (*v / scale).round().clamp(-levels, levels) * scale; - } - } -} - -/// Dense gated FFN over graded-precision f32 weights (feature-major -/// [intermediate × hidden] for gate/up/down). -struct GradedFfn { - gate: Vec>, - up: Vec>, - down: Vec>, -} - -impl FfnBackend for GradedFfn { - fn forward(&self, layer: usize, x: &Array2) -> Array2 { - let (g, u, d) = (&self.gate[layer], &self.up[layer], &self.down[layer]); - let hidden = x.shape()[1]; - let mut out = Array2::::zeros((x.shape()[0], hidden)); - for (s, xr) in x.rows().into_iter().enumerate() { - let xr = xr.to_owned(); - let gs = g.dot(&xr); // [inter] - let us = u.dot(&xr); - let act: Array1 = gs - .iter() - .zip(us.iter()) - .map(|(&gg, &uu)| larql_inference::ffn::gelu_tanh(gg) * uu) - .collect(); - let o = act.dot(d); // [hidden] - out.row_mut(s).assign(&o); - } - out - } - // forward_observed: trait default (Absent) — this arm computes no - // observable activation tensor and must not fabricate one. - fn name(&self) -> &str { - "graded" - } -} - -/// Build a GradedFfn with head_frac features (by ‖down_row‖) at `head_bits`, -/// the rest at `tail_bits`. Returns (ffn, avg_bits_per_feature). -fn build_graded( - weights: &ModelWeights, - index: &larql_vindex::VectorIndex, - head_frac: f64, - head_bits: u32, - tail_bits: u32, -) -> (GradedFfn, f64) { - let nl = weights.num_layers; - let hidden = weights.hidden_size; - let probe = WalkFfn::new_unlimited(weights, index); - let (mut gate, mut up, mut down) = (Vec::new(), Vec::new(), Vec::new()); - let (mut bit_sum, mut feat_sum) = (0.0f64, 0.0f64); - for layer in 0..nl { - let inter = index.num_features(layer); - let g = index.kquant_ffn_layer(layer, 0).expect("gate f32"); - let u = index.kquant_ffn_layer(layer, 1).expect("up f32"); - let d = index.kquant_ffn_layer(layer, 2).expect("down f32"); - // Importance order: ‖down_row‖ descending. - let norms = probe.down_row_norms_pub(layer).expect("down norms"); - let mut order: Vec = (0..inter).collect(); - order.sort_unstable_by(|&a, &b| norms[b].total_cmp(&norms[a])); - let head_n = (inter as f64 * head_frac).round() as usize; - let mut bits_of = vec![tail_bits; inter]; - for &f in order.iter().take(head_n) { - bits_of[f] = head_bits; - } - let mut gm = Array2::::zeros((inter, hidden)); - let mut um = Array2::::zeros((inter, hidden)); - let mut dm = Array2::::zeros((inter, hidden)); - for f in 0..inter { - let b = bits_of[f]; - let mut gr = g[f * hidden..(f + 1) * hidden].to_vec(); - let mut ur = u[f * hidden..(f + 1) * hidden].to_vec(); - let mut dr = d[f * hidden..(f + 1) * hidden].to_vec(); - requant_row(&mut gr, b); - requant_row(&mut ur, b); - requant_row(&mut dr, b); - gm.row_mut(f).assign(&Array1::from(gr)); - um.row_mut(f).assign(&Array1::from(ur)); - dm.row_mut(f).assign(&Array1::from(dr)); - bit_sum += 3.0 * b as f64; // gate+up+down rows at b bits - feat_sum += 3.0; - } - gate.push(gm); - up.push(um); - down.push(dm); - } - (GradedFfn { gate, up, down }, bit_sum / feat_sum) -} - -/// Component-graded: each of gate/up/down quantised UNIFORMLY across features at -/// its own bit-width (grades along the gate-vs-down axis, not the flat down-norm -/// axis). Returns (ffn, avg_bits_per_feature). -fn build_component( - weights: &ModelWeights, - index: &larql_vindex::VectorIndex, - gate_bits: u32, - up_bits: u32, - down_bits: u32, -) -> (GradedFfn, f64) { - let nl = weights.num_layers; - let hidden = weights.hidden_size; - let (mut gate, mut up, mut down) = (Vec::new(), Vec::new(), Vec::new()); - for layer in 0..nl { - let inter = index.num_features(layer); - for (comp, bits, store) in [ - (0usize, gate_bits, &mut gate), - (1, up_bits, &mut up), - (2, down_bits, &mut down), - ] { - let w = index.kquant_ffn_layer(layer, comp).expect("f32 comp"); - let mut m = Array2::::zeros((inter, hidden)); - for f in 0..inter { - let mut r = w[f * hidden..(f + 1) * hidden].to_vec(); - requant_row(&mut r, bits); - m.row_mut(f).assign(&Array1::from(r)); - } - store.push(m); - } - } - let avg = (gate_bits + up_bits + down_bits) as f64 / 3.0; - (GradedFfn { gate, up, down }, avg) -} - -fn next_dist( - weights: &ModelWeights, - tok: &tokenizers::Tokenizer, - ids: &[u32], - ffn: &dyn FfnBackend, -) -> (HashMap, u32) { - let r = predict_with_ffn(weights, tok, ids, usize::MAX, ffn); - let arg = r.token_ids.first().copied().unwrap_or(0); - ( - r.token_ids - .into_iter() - .zip(r.predictions.into_iter().map(|(_, p)| p)) - .collect(), - arg, - ) -} - -fn kl_bits(p: &HashMap, q: &HashMap) -> f64 { - let eps = 1e-12; - let mut kl = 0.0; - for (&id, &pi) in p { - if pi <= 0.0 { - continue; - } - let qi = q.get(&id).copied().unwrap_or(0.0).max(eps); - kl += pi * (pi.max(eps) / qi).ln(); - } - kl / std::f64::consts::LN_2 -} - -fn main() { - let args: Vec = std::env::args().collect(); - let vindex = args - .get(1) - .cloned() - .unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".to_string()); - let dir = std::path::PathBuf::from(&vindex); - let mut cb = larql_vindex::SilentLoadCallbacks; - eprintln!("Loading {vindex} ..."); - let mut weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index.load_interleaved_kquant(&dir).expect("interleaved"); - index.load_attn_kquant(&dir).expect("attn"); - let _ = index.load_lm_head_kquant(&dir); - let tok = load_tokenizer(&dir).expect("tok"); - for layer in 0..weights.num_layers { - insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant attn"); - } - - let prompts = [ - "The capital of France is", - "Water is made of hydrogen and", - "def add(a, b):\n return a +", - "Bonjour, comment allez-", - ]; - - // f32 reference (all 32 bits) — the ground-truth FFN distribution. - eprintln!("Building f32 reference ..."); - let (ref_ffn, _) = build_graded(&weights, &index, 1.0, 32, 32); - let refs: Vec<(HashMap, u32)> = prompts - .iter() - .map(|p| { - let ids = tok.encode(*p, true).expect("enc").get_ids().to_vec(); - next_dist(&weights, &tok, &ids, &ref_ffn) - }) - .collect(); - drop(ref_ffn); - - println!("\nGraded-precision FFN — KL vs f32 reference, top-1 agreement, bandwidth\n"); - println!( - "{:<26} {:>9} {:>9} {:>10} {:>9}", - "schedule", "avg-bits", "bw vs Q4", "KL(bits)", "top1%" - ); - - // (label, head_frac, head_bits, tail_bits) - let schedules: &[(&str, f64, u32, u32)] = &[ - ("uniform 4-bit", 1.0, 4, 4), - ("uniform 3-bit", 1.0, 3, 3), - ("uniform 2-bit", 1.0, 2, 2), - ("head40/4 tail60/3", 0.4, 4, 3), - ("head20/4 tail80/3", 0.2, 4, 3), - ("head10/4 tail90/3", 0.1, 4, 3), - ("head40/4 tail60/2", 0.4, 4, 2), - ("head20/8 tail80/3", 0.2, 8, 3), - ]; - for &(label, hf, hb, tb) in schedules { - let (ffn, avg_bits) = build_graded(&weights, &index, hf, hb, tb); - let (mut kl_sum, mut agree, n) = (0.0, 0usize, prompts.len()); - for (i, p) in prompts.iter().enumerate() { - let ids = tok.encode(*p, true).expect("enc").get_ids().to_vec(); - let (d, arg) = next_dist(&weights, &tok, &ids, &ffn); - kl_sum += kl_bits(&refs[i].0, &d); - if arg == refs[i].1 { - agree += 1; - } - } - let nf = n as f64; - println!( - "{label:<26} {avg_bits:>9.2} {:>8.2}× {:>10.4} {:>8.0}%", - avg_bits / 4.0, - kl_sum / nf, - 100.0 * agree as f64 / nf - ); - drop(ffn); - } - // ── Component grading along the gate-vs-down axis ──────────────── - // Prediction (from the universal-3-bit floor): the GATE is precision- - // critical (it routes which features fire — a discrete error), the DOWN - // is forgiving (magnitude only). So gate/up Q3 + down Q2 should survive - // where uniform-Q2 cliffs, beating uniform-3's 0.75×. - println!("\nComponent grading (gate/up vs down bits) — KL vs f32, top-1, bandwidth\n"); - println!( - "{:<26} {:>9} {:>9} {:>10} {:>9}", - "gate/up/down bits", "avg-bits", "bw vs Q4", "KL(bits)", "top1%" - ); - let comp: &[(&str, u32, u32, u32)] = &[ - ("g3 u3 d3 (=uniform3)", 3, 3, 3), - ("g3 u3 d2", 3, 3, 2), - ("g3 u3 d1", 3, 3, 1), - ("g2 u2 d3 (gate@2)", 2, 2, 3), - ("g4 u4 d2", 4, 4, 2), - ("g4 u3 d2", 4, 3, 2), - ]; - for &(label, gb, ub, db) in comp { - let (ffn, avg_bits) = build_component(&weights, &index, gb, ub, db); - let (mut kl_sum, mut agree, n) = (0.0, 0usize, prompts.len()); - for (i, p) in prompts.iter().enumerate() { - let ids = tok.encode(*p, true).expect("enc").get_ids().to_vec(); - let (d, arg) = next_dist(&weights, &tok, &ids, &ffn); - kl_sum += kl_bits(&refs[i].0, &d); - if arg == refs[i].1 { - agree += 1; - } - } - let nf = n as f64; - println!( - "{label:<26} {avg_bits:>9.2} {:>8.2}× {:>10.4} {:>8.0}%", - avg_bits / 4.0, - kl_sum / nf, - 100.0 * agree as f64 / nf - ); - drop(ffn); - } - - println!("\n (Reference = f32 FFN. 'bw vs Q4' = avg-bits/4 = bandwidth ratio vs uniform Q4K.\n Prediction: g3/d2 survives (gate routes, down forgives); g2/d3 cliffs (gate@2 breaks routing).)"); -} diff --git a/crates/larql-inference/examples/walk_ffn_k_agreement.rs b/crates/larql-inference/examples/walk_ffn_k_agreement.rs deleted file mode 100644 index aae634103..000000000 --- a/crates/larql-inference/examples/walk_ffn_k_agreement.rs +++ /dev/null @@ -1,126 +0,0 @@ -//! WalkFfn K-vs-agreement frontier (task #23 follow-up). -//! -//! The decode falsification killed K=512 at every band (even gate-KNN, the best -//! router, ~60–75% top-1 agreement vs dense — far under a 90% generation bar). -//! But agreement → 100% as K → num_features (full K *is* dense). So the real -//! question: what is the **minimum K that clears 90% top-1 agreement**, and is -//! that K still small enough to beat dense (microbench #18: cheap routing wins -//! up to K≈2048, washes out above)? Sweeps gate-KNN (the accuracy *ceiling* — -//! a content-addressed router can't beat full-projection top-K) across K and -//! band depth, in-dist vs OOD, teacher-forced on dense's own greedy stream. -//! -//! Usage: `cargo run --release --example walk_ffn_k_agreement -- [VINDEX_DIR]` - -use larql_inference::vindex::{insert_q4k_layer_tensors_resident, WalkFfn, WalkFfnConfig}; -use larql_inference::{load_tokenizer, predict_with_ffn}; - -fn argmax( - weights: &larql_models::ModelWeights, - tok: &tokenizers::Tokenizer, - ids: &[u32], - ffn: &dyn larql_inference::ffn::FfnBackend, -) -> u32 { - let r = predict_with_ffn(weights, tok, ids, 1, ffn); - r.token_ids.first().copied().unwrap_or(0) -} - -fn main() { - let args: Vec = std::env::args().collect(); - let vindex = args - .get(1) - .cloned() - .unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".to_string()); - let dir = std::path::PathBuf::from(&vindex); - let mut cb = larql_vindex::SilentLoadCallbacks; - eprintln!("Loading {vindex} ..."); - let mut weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index.load_interleaved_kquant(&dir).expect("interleaved"); - index.load_attn_kquant(&dir).expect("attn"); - let _ = index.load_lm_head_kquant(&dir); - let _ = index.load_down_features_q4k(&dir); - let _ = index.load_down_features(&dir); - let _ = index.load_gate_vectors_q4(&dir); - let tok = load_tokenizer(&dir).expect("tok"); - for layer in 0..weights.num_layers { - insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant"); - } - let nl = weights.num_layers; - let feats = index.num_features(nl / 2); - - let in_dist = [ - "The capital of France is", - "The largest planet in the solar system is", - "Water is made of hydrogen and", - ]; - let ood = [ - "def add(a, b):\n return a +", - "Bonjour, comment allez-", - "Once upon a time, there was a", - ]; - let gen_len = 12usize; - let ks = [512usize, 1024, 2048, 4096]; - let depths = [4usize, 9]; - - // Per (seed-group, depth, K): top-1 agreement vs dense, teacher-forced on - // dense's greedy stream. Dense computed once per position and reused. - let run = |label: &str, seeds: &[&str]| { - println!( - "\n{label} — gate-KNN top-1 agreement vs dense, {} seeds × {gen_len} = {} positions", - seeds.len(), - seeds.len() * gen_len - ); - println!(" ({feats} feats/layer; bar ≥ 90%)"); - // agree[depth_idx][k_idx] - let mut agree = vec![vec![0usize; ks.len()]; depths.len()]; - let mut total = 0usize; - for s in seeds { - let mut ids = tok.encode(*s, true).expect("enc").get_ids().to_vec(); - for _ in 0..gen_len { - let d = argmax( - &weights, - &tok, - &ids, - &WalkFfn::new_unlimited(&weights, &index), - ); - for (di, &depth) in depths.iter().enumerate() { - let sf = nl.saturating_sub(depth); - for (ki, &k) in ks.iter().enumerate() { - let g = argmax( - &weights, - &tok, - &ids, - &WalkFfn::from_config( - &weights, - &index, - WalkFfnConfig::hybrid(nl, sf, k), - ), - ); - agree[di][ki] += (g == d) as usize; - } - } - total += 1; - ids.push(d); - } - } - let pct = |n: usize| 100.0 * n as f64 / total.max(1) as f64; - print!(" {:<10}", "band\\K"); - for k in ks { - print!(" K={k:<6}"); - } - println!(); - for (di, &depth) in depths.iter().enumerate() { - print!(" last {depth:<5}"); - for &a in agree[di].iter().take(ks.len()) { - let p = pct(a); - let mark = if p >= 90.0 { "*" } else { " " }; - print!(" {p:>5.1}{mark} "); - } - println!(); - } - }; - - run("IN-DISTRIBUTION", &in_dist); - run("OUT-OF-DISTRIBUTION", &ood); - println!("\n(* = clears 90%. Cross-reference K against microbench #18: cheap-route beats dense to ~K=2048, washes above.)"); -} diff --git a/crates/larql-inference/examples/walk_ffn_microbench.rs b/crates/larql-inference/examples/walk_ffn_microbench.rs deleted file mode 100644 index 4b3c02be2..000000000 --- a/crates/larql-inference/examples/walk_ffn_microbench.rs +++ /dev/null @@ -1,101 +0,0 @@ -//! WalkFfn FFN microbench — isolates `WalkFfn::forward` at **decode shape -//! (seq_len = 1)** across K, vs the f32-BLAS `WeightFfn` baseline. This is the -//! instrument the bottleneck diagnosis flagged as missing: the only end-to-end -//! WalkFfn path (`walk --predict`) is non-KV-cached, so FFN sparsity is masked -//! by attention + lm_head re-compute. Here the FFN is the *only* thing timed. -//! -//! Usage: `cargo run --release --example walk_ffn_microbench -- [VINDEX_DIR]` -//! (default: output/gemma3-4b-q4k-v2.vindex) - -use larql_inference::ffn::FfnBackend; -use larql_inference::vindex::{WalkFfn, WalkFfnConfig}; -use ndarray::Array2; -use std::sync::Arc; -use std::time::Instant; - -/// Deterministic, residual-independent route of `k` features per layer — -/// stands in for hash routing (Exp 27: token-ID-deterministic top-K mask). -/// The point is the *cost profile*: the route is precomputed, so selection -/// never touches the full gate matrix. A strided pick spreads the features -/// across the matrix (realistic cache behaviour for a gather). -fn precomputed_pool(num_layers: usize, num_features: usize, k: usize) -> Arc>> { - let k = k.min(num_features.max(1)); - let stride = (num_features / k.max(1)).max(1); - let per_layer: Vec = (0..k).map(|i| (i * stride) % num_features).collect(); - Arc::new(vec![per_layer; num_layers]) -} - -fn main() { - let args: Vec = std::env::args().collect(); - let vindex = args - .get(1) - .cloned() - .unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".to_string()); - let dir = std::path::PathBuf::from(&vindex); - let mut cb = larql_vindex::SilentLoadCallbacks; - - eprintln!("Loading {vindex} ..."); - let weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index - .load_interleaved_kquant(&dir) - .expect("interleaved kquant"); - let _ = index.load_down_features_q4k(&dir); - let _ = index.load_down_features(&dir); - let _ = index.load_gate_vectors_q4(&dir); - - let hidden = weights.hidden_size; - let layer = weights.num_layers / 2; - let feats = index.num_features(layer); - let iters = 300usize; - - // Representative seq_len=1 input (timing is value-independent for the - // matmul; gate-KNN cost depends on K, not values). - let x = Array2::from_shape_fn((1, hidden), |(_, j)| ((j as f32) * 0.013).sin() * 0.1); - - let bench = |name: &str, ffn: &dyn FfnBackend| { - for _ in 0..20 { - let _ = ffn.forward(layer, &x); - } - let t = Instant::now(); - for _ in 0..iters { - let _ = ffn.forward(layer, &x); - } - let us = t.elapsed().as_micros() as f64 / iters as f64; - println!(" {name:<26} {us:>9.1} µs/call"); - }; - - println!( - "\nWalkFfn FFN microbench — seq_len=1, layer {layer}, {feats} features, {iters} iters\n" - ); - - // WalkFfn dense (kquant_native, all features) is the baseline; sparse - // (gate-KNN top-K) is the "touch fewer weights" path. - bench( - "WalkFfn dense (k=MAX)", - &WalkFfn::new_unlimited(&weights, &index), - ); - for k in [2048usize, 512, 128, 32] { - let pct = 100.0 * k as f64 / feats.max(1) as f64; - bench( - &format!("WalkFfn gate-KNN k={k} ({pct:.0}%)"), - &WalkFfn::new(&weights, &index, k), - ); - } - - // Cheap routing (task #18): precomputed per-layer route, gate scored - // for only the K route features (O(K)) — no full gate projection. This - // is the lever the gate-KNN microbench said was the only way for sparse - // to beat dense. Same K sweep, head-to-head with gate-KNN above. - println!(); - for k in [2048usize, 512, 128, 32] { - let pct = 100.0 * k as f64 / feats.max(1) as f64; - let pool = precomputed_pool(weights.num_layers, feats, k); - let cfg = WalkFfnConfig::sparse(weights.num_layers, k) - .with_pool_per_layer(pool) - .with_precomputed_routing(true); - let ffn = WalkFfn::from_config(&weights, &index, cfg); - bench(&format!("WalkFfn cheap-route k={k} ({pct:.0}%)"), &ffn); - } - println!(); -} diff --git a/crates/larql-inference/examples/walk_ffn_nll.rs b/crates/larql-inference/examples/walk_ffn_nll.rs deleted file mode 100644 index 3aa4b844c..000000000 --- a/crates/larql-inference/examples/walk_ffn_nll.rs +++ /dev/null @@ -1,233 +0,0 @@ -//! Three-way per-token NLL adjudicator for Q3 FFN (task #26 deciding metric). -//! -//! The drift test found 19% per-step argmax flip + full generation divergence, -//! but that can't tell benign near-tie chaos from real degradation. This does. -//! -//! Three arms — **f32, Q4, Q3** — teacher-forced on entropic held prose, scoring -//! per-token NLL (bits = −log2 p(true next token)). f32 is ground truth; **Q4 -//! calibrates tolerance** (the Q4→f32 gap is the precision-cost decision you -//! already shipped once); Q3 is the candidate. Reports the per-token -//! *distribution* (median/p90/p99/max), not just the mean — the mean hides the -//! catastrophic-token tail where the decision lives — and keeps the **flip rate -//! alongside** (NLL doesn't replace it: flip-high + NLL-flat = benign chaos → -//! ship; flip-high + Q3-NLL-elevated = real cost → don't). -//! -//! Usage: `cargo run --release --example walk_ffn_nll -- [VINDEX]` - -use larql_inference::ffn::FfnBackend; -use larql_inference::vindex::insert_q4k_layer_tensors_resident; -use larql_inference::{load_tokenizer, predict_with_ffn}; -use larql_models::ModelWeights; -use ndarray::{Array1, Array2}; -use std::collections::HashMap; - -fn requant_row(row: &mut [f32], bits: u32) { - if bits >= 16 { - return; // f32 passthrough - } - const BLK: usize = 32; - for blk in row.chunks_mut(BLK) { - let maxabs = blk.iter().fold(0.0f32, |m, &v| m.max(v.abs())); - if maxabs == 0.0 { - continue; - } - let levels = ((1u32 << (bits - 1)) - 1) as f32; - let scale = maxabs / levels; - for v in blk.iter_mut() { - *v = (*v / scale).round().clamp(-levels, levels) * scale; - } - } -} - -struct GradedFfn { - gate: Vec>, - up: Vec>, - down: Vec>, -} -impl FfnBackend for GradedFfn { - fn forward(&self, layer: usize, x: &Array2) -> Array2 { - let (g, u, d) = (&self.gate[layer], &self.up[layer], &self.down[layer]); - let hidden = x.shape()[1]; - let mut out = Array2::::zeros((x.shape()[0], hidden)); - for (s, xr) in x.rows().into_iter().enumerate() { - let xr = xr.to_owned(); - let gs = g.dot(&xr); - let us = u.dot(&xr); - let act: Array1 = gs - .iter() - .zip(us.iter()) - .map(|(&gg, &uu)| larql_inference::ffn::gelu_tanh(gg) * uu) - .collect(); - out.row_mut(s).assign(&act.dot(d)); - } - out - } - // forward_observed: trait default (Absent) — this arm computes no - // observable activation tensor and must not fabricate one. - fn name(&self) -> &str { - "graded" - } -} - -fn build_uniform( - weights: &ModelWeights, - index: &larql_vindex::VectorIndex, - bits: u32, -) -> GradedFfn { - let (nl, hidden) = (weights.num_layers, weights.hidden_size); - let (mut gate, mut up, mut down) = (Vec::new(), Vec::new(), Vec::new()); - for layer in 0..nl { - let inter = index.num_features(layer); - for (comp, store) in [(0usize, &mut gate), (1, &mut up), (2, &mut down)] { - let w = index.kquant_ffn_layer(layer, comp).expect("f32 comp"); - let mut m = Array2::::zeros((inter, hidden)); - for f in 0..inter { - let mut r = w[f * hidden..(f + 1) * hidden].to_vec(); - requant_row(&mut r, bits); - m.row_mut(f).assign(&Array1::from(r)); - } - store.push(m); - } - } - GradedFfn { gate, up, down } -} - -/// Teacher-forced per-token NLL (bits) over `ids`, plus the per-position argmax -/// token (for the flip rate). Scores positions `1..ids.len()`. -fn token_nlls( - weights: &ModelWeights, - tok: &tokenizers::Tokenizer, - ids: &[u32], - ffn: &dyn FfnBackend, -) -> (Vec, Vec) { - let mut nlls = Vec::new(); - let mut args = Vec::new(); - for i in 1..ids.len() { - if i % 8 == 0 { - eprint!("\r [{}] pos {i}/{} ", ffn.name(), ids.len()); - } - let r = predict_with_ffn(weights, tok, &ids[..i], usize::MAX, ffn); - let dist: HashMap = r - .token_ids - .iter() - .copied() - .zip(r.predictions.iter().map(|(_, p)| *p)) - .collect(); - let p = dist.get(&ids[i]).copied().unwrap_or(0.0).max(1e-12); - nlls.push(-p.log2()); - args.push(r.token_ids.first().copied().unwrap_or(0)); - } - (nlls, args) -} - -fn pct(sorted: &[f64], q: f64) -> f64 { - if sorted.is_empty() { - return 0.0; - } - let idx = ((sorted.len() as f64 - 1.0) * q).round() as usize; - sorted[idx.min(sorted.len() - 1)] -} - -fn main() { - let args: Vec = std::env::args().collect(); - let vindex = args - .get(1) - .cloned() - .unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".to_string()); - let dir = std::path::PathBuf::from(&vindex); - let mut cb = larql_vindex::SilentLoadCallbacks; - eprintln!("Loading {vindex} ..."); - let mut weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index.load_interleaved_kquant(&dir).expect("interleaved"); - index.load_attn_kquant(&dir).expect("attn"); - let _ = index.load_lm_head_kquant(&dir); - let tok = load_tokenizer(&dir).expect("tok"); - for layer in 0..weights.num_layers { - insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant attn"); - } - - // Entropic narrative prose — real lexical choice (NOT code/boilerplate, - // which is near-deterministic and would flatter Q3 like the n=4 top-1 did). - let passage = "The expedition had been planned for years, but nothing prepared \ -them for the silence of the ice. Each morning the wind died at dawn, and the only \ -sound was the slow groan of the glacier shifting beneath their tents. Provisions \ -were running low, and the captain knew that another week of delay would mean \ -turning back without ever reaching the plateau they had crossed two oceans to find."; - let ids = tok.encode(passage, true).expect("enc").get_ids().to_vec(); - eprintln!( - "Held passage: {} tokens. Building f32 / Q4 / Q3 ...", - ids.len() - ); - - let arms = [("f32", 32u32), ("Q4", 4), ("Q3", 3)]; - let mut nll_by_arm: Vec<(String, Vec, Vec)> = Vec::new(); - for (name, bits) in arms { - eprintln!(" arm {name} ({bits}-bit): building + scoring ..."); - let ffn = build_uniform(&weights, &index, bits); - let (nlls, a) = token_nlls(&weights, &tok, &ids, &ffn); - eprintln!("\r arm {name}: done ({} positions) ", nlls.len()); - nll_by_arm.push((name.to_string(), nlls, a)); - drop(ffn); - } - - println!("\nThree-way per-token NLL (bits/token), teacher-forced on entropic prose ({} scored positions)\n", nll_by_arm[0].1.len()); - println!( - "{:<6} {:>8} {:>8} {:>8} {:>8} {:>8}", - "arm", "mean", "median", "p90", "p99", "max" - ); - let f32_mean: f64; - { - let (_, n, _) = &nll_by_arm[0]; - f32_mean = n.iter().sum::() / n.len().max(1) as f64; - } - for (name, n, _) in &nll_by_arm { - let mut s = n.clone(); - s.sort_by(|a, b| a.total_cmp(b)); - let mean = n.iter().sum::() / n.len().max(1) as f64; - println!( - "{name:<6} {mean:>8.3} {:>8.3} {:>8.3} {:>8.3} {:>8.3}", - pct(&s, 0.50), - pct(&s, 0.90), - pct(&s, 0.99), - pct(&s, 1.0) - ); - } - - // Precision-cost ladder + tail of the per-position deltas vs f32. - let f32_n = &nll_by_arm[0].1; - println!("\nΔ NLL vs f32 (bits/token) — the precision-cost ladder + tail:"); - for (name, n, _) in nll_by_arm.iter().skip(1) { - let deltas: Vec = n.iter().zip(f32_n).map(|(&a, &b)| a - b).collect(); - let mut s = deltas.clone(); - s.sort_by(|a, b| a.total_cmp(b)); - let mean = deltas.iter().sum::() / deltas.len().max(1) as f64; - println!( - " {name:<4} mean Δ {mean:+.3} p90 {:+.3} p99 {:+.3} worst-token {:+.3}", - pct(&s, 0.90), - pct(&s, 0.99), - pct(&s, 1.0) - ); - } - - // Flip rate kept ALONGSIDE: Q3 argmax vs Q4 argmax, teacher-forced. - let q4_args = &nll_by_arm[1].2; - let q3_args = &nll_by_arm[2].2; - let flips = q3_args.iter().zip(q4_args).filter(|(a, b)| a != b).count(); - let q4_mean = nll_by_arm[1].1.iter().sum::() / nll_by_arm[1].1.len() as f64; - let q3_mean = nll_by_arm[2].1.iter().sum::() / nll_by_arm[2].1.len() as f64; - println!( - "\n flip rate (Q3 vs Q4 argmax, teacher-forced): {:.1}%", - 100.0 * flips as f64 / q3_args.len().max(1) as f64 - ); - println!( - " ladder: f32 {f32_mean:.3} → Q4 {q4_mean:.3} (+{:.3}) → Q3 {q3_mean:.3} (+{:.3} vs Q4)", - q4_mean - f32_mean, - q3_mean - q4_mean - ); - println!( - "\n DECISION: Q3→Q4 step {:.3} bits vs the Q4→f32 step {:.3} bits you already shipped.\n flip-high + ladder-flat ⇒ benign near-tie chaos (ship); flip-high + Q3 elevated ⇒ real cost.", - q3_mean - q4_mean, - q4_mean - f32_mean - ); -} diff --git a/crates/larql-inference/examples/walk_ffn_r4_zeroout.rs b/crates/larql-inference/examples/walk_ffn_r4_zeroout.rs deleted file mode 100644 index f8478265a..000000000 --- a/crates/larql-inference/examples/walk_ffn_r4_zeroout.rs +++ /dev/null @@ -1,557 +0,0 @@ -//! R4 zero-out for the sparse-FFN / ANN routing programme. -//! -//! **The only question:** with perfect routes supplied for free, is the -//! remaining sparse execution cheaper than dense execution? -//! -//! If sparse execution with **free** routing still can't turn a faithful -//! row reduction into proportional speed, the bottleneck is the sparse -//! kernel, not selection — and there is no reason to repair HNSW or -//! pursue better routing for the all-layer route. -//! -//! ## Two levers, deliberately separated -//! -//! A known route does two things here, and collapsing them into one -//! number is how you get an unexplainable result: -//! -//! 1. **Selection disappears** — the O(N·d) gate sweep becomes O(K·d). -//! 2. **A different kernel unlocks** — `sparse.rs:158` fires the -//! contiguous-gather kernel (`sparse:gather_q4k`) only when the route -//! is known in advance AND unranked (`!rank_within_pool`). Production -//! (which must search) can never reach it. -//! -//! `rank_within_pool = true` on a K-sized pool forces the production -//! kernel with a known route, isolating lever 1 from lever 2. -//! -//! ## Arms -//! -//! | Arm | Pool | Rank | Kernel | Selection cost | -//! |---|---|---|---|---| -//! | `dense` | — | — | dense FFN | none | -//! | `exact-pool` | all N | yes | `parallel_q4k_down` | N row-dots | -//! | `oracle-par` | top-K | yes | `parallel_q4k_down` | K row-dots | -//! | `oracle-gath` | top-K | no | `gather_q4k` | K row-dots | -//! -//! - **R4 decision:** `oracle-par` vs `dense` — routing free, production kernel. -//! - **Kernel bonus:** `oracle-par` vs `oracle-gath` — selection-identical. -//! -//! `exact-pool` is retained as a kernel-matched upper bound on selection -//! cost, NOT as the router budget: it pays N *scattered* row-dots where -//! production pays a batched gemv, so differencing it overstates -//! recoverable production time by ~30%. The router budget is quoted -//! against production's own ~3 ms/layer selection cost. -//! -//! ## What is *not* zeroed -//! -//! The oracle arms still pay the K gate row-dots. That is correct: the -//! sparse kernel consumes `gate_score` as its activation input, and any -//! real router emits candidate IDs, not exact scores. Those K dots are -//! execution work every arm must do. What disappears is the search over N. -//! -//! ## Measurement protocol (paired, interleaved) -//! -//! Arms are **not** run in long separate blocks — thermal and scheduler -//! drift over a multi-minute block would be indistinguishable from an arm -//! effect. Instead every repeat measures all four arms back-to-back in a -//! **rotated order**, and the reported statistic is the distribution of -//! **per-repeat paired ratios** (`arm / dense` within the same repeat), so -//! drift cancels. Dense wall-time per repeat is retained as a throttle -//! sentinel; a drifting sentinel invalidates the run rather than being -//! averaged away. -//! -//! Routes are captured once per cell and held fixed across every repeat. -//! -//! ## Guardrail (access-pattern faithfulness) -//! -//! Routes are captured from `exact-pool`'s runtime trace in executed visit -//! order and replayed in that order, so the oracle arms issue the same -//! scattered row reads a real stage-1 router would produce — the search is -//! zeroed, the gather disorder is not. Capturing from `exact-pool` rather -//! than production also keeps the gate-score source identical (interleaved -//! Q4K bytes), so a parity failure means a real divergence. -//! -//! Usage: `cargo run --release --example walk_ffn_r4_zeroout -- [VINDEX_DIR]` -//! Refuses to run off AC power or on a loaded box; `--allow-dirty` overrides -//! (and stamps every number as provisional). - -use larql_inference::vindex::{ - insert_q4k_layer_tensors_resident, LayerTraceRecord, PhaseTimingsHandle, WalkFfn, WalkFfnConfig, -}; -use larql_inference::{load_tokenizer, predict_with_ffn}; -use std::sync::atomic::Ordering; -use std::sync::Arc; -use std::time::Instant; - -/// Untimed forwards per arm before the repeat loop, so every kernel and -/// shape is warm before anything is recorded. -const WARMUP: usize = 8; -/// Forwards averaged into one arm's measurement inside a single repeat. -/// Small on purpose: a repeat must stay short enough that all four arms -/// see the same thermal conditions. -const BLOCK: usize = 3; -/// Paired repeats per cell. -const REPEATS: usize = 16; - -/// Fraction of drift in the dense sentinel (last quartile vs first) above -/// which the cell is reported as thermally invalid rather than believed. -const SENTINEL_DRIFT_LIMIT: f64 = 0.10; - -/// Dense-sentinel IQR/median above which the cell is treated as disturbed. -/// A drift-only gate is porous: quartile means can agree while a single -/// repeat spikes 4×, which is exactly the cell that must not be believed. -const SENTINEL_SPREAD_LIMIT: f64 = 0.10; - -/// 1-minute load average above which the box is treated as contended. -/// A concurrent build is as fatal to a ratio experiment as throttling. -const LOAD_LIMIT: f64 = 3.0; - -/// Relative tolerance for residual-delta parity. Arms executing the same -/// feature set in a different accumulation order differ in the last float -/// bits; that is a summation-order artifact. Set equality is exact. -const PARITY_REL_TOL: f32 = 1e-4; - -/// Decisive cells only — `(band, K)`. Densely sweeping accuracy-dead -/// cells buys nothing. `usize::MAX` band = all layers. -const CELLS: [(usize, usize); 7] = [ - (usize::MAX, 2048), - (usize::MAX, 4096), - (usize::MAX, 6144), - (4, 2048), - (4, 4096), - (9, 2048), - (9, 4096), -]; - -/// Decode prompt. A single trailing token gives the seq_len=1 shape where -/// the FFN is a matvec — the regime the sparse walk targets. A multi-token -/// prompt would measure prefill (batched gemm), where dense wins by -/// construction and the comparison is meaningless. -const PROMPT: &str = "The capital of France is the city of"; - -const ARM_DENSE: usize = 0; -const ARM_EXACT: usize = 1; -const ARM_ORACLE_PAR: usize = 2; -const ARM_ORACLE_GATH: usize = 3; -const ARM_NAMES: [&str; 4] = ["dense", "exact-pool", "oracle-par", "oracle-gath"]; - -/// Per-layer routes in executed visit order. -type Routes = Arc>>; - -// ── Environment gates ──────────────────────────────────────────────── - -fn on_ac_power() -> bool { - std::process::Command::new("pmset") - .args(["-g", "ps"]) - .output() - .ok() - .map(|o| String::from_utf8_lossy(&o.stdout).contains("AC Power")) - .unwrap_or(false) -} - -fn load_average() -> f64 { - std::process::Command::new("sysctl") - .args(["-n", "vm.loadavg"]) - .output() - .ok() - .and_then(|o| { - String::from_utf8_lossy(&o.stdout) - .split_whitespace() - .nth(1) - .and_then(|v| v.parse().ok()) - }) - .unwrap_or(f64::NAN) -} - -// ── Statistics ─────────────────────────────────────────────────────── - -fn quantile(sorted: &[f64], q: f64) -> f64 { - if sorted.is_empty() { - return f64::NAN; - } - let pos = q * (sorted.len() - 1) as f64; - let lo = pos.floor() as usize; - let hi = pos.ceil() as usize; - if lo == hi { - sorted[lo] - } else { - sorted[lo] + (sorted[hi] - sorted[lo]) * (pos - lo as f64) - } -} - -/// Median, p25, p75, min, max of a sample. -struct Summary { - median: f64, - p25: f64, - p75: f64, - min: f64, - max: f64, -} - -fn summarise(values: &[f64]) -> Summary { - let mut s = values.to_vec(); - s.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); - Summary { - median: quantile(&s, 0.5), - p25: quantile(&s, 0.25), - p75: quantile(&s, 0.75), - min: *s.first().unwrap_or(&f64::NAN), - max: *s.last().unwrap_or(&f64::NAN), - } -} - -/// Drift in the throttle sentinel: mean of the last quarter of repeats -/// against the first quarter. Large positive = the box slowed mid-run. -/// -/// Drift alone is NOT sufficient: it is computed from quartile means, so a -/// single dense repeat spiking (one cell saw 543 ms against a 127 ms median) -/// can leave drift small while the cell is plainly disturbed. Pair it with -/// [`sentinel_spread`]. -fn sentinel_drift(dense_us: &[f64]) -> f64 { - let q = (dense_us.len() / 4).max(1); - let head: f64 = dense_us[..q].iter().sum::() / q as f64; - let tail: f64 = dense_us[dense_us.len() - q..].iter().sum::() / q as f64; - (tail - head) / head -} - -/// Dispersion of the throttle sentinel: IQR as a fraction of the median. -/// Catches the disturbed-but-not-drifting cell that [`sentinel_drift`] -/// passes — an outlier repeat inflates the spread even when the quartile -/// means agree. -fn sentinel_spread(dense_us: &[f64]) -> f64 { - let s = summarise(dense_us); - (s.p75 - s.p25) / s.median -} - -// ── Trace helpers ──────────────────────────────────────────────────── - -fn sparse_from(num_layers: usize, band: usize) -> usize { - if band == usize::MAX { - 0 - } else { - num_layers.saturating_sub(band) - } -} - -fn band_label(band: usize) -> String { - if band == usize::MAX { - "all".to_string() - } else { - format!("last-{band}") - } -} - -/// Fold a runtime trace into per-layer routes, preserving executed visit -/// order. Dense layers contribute nothing — they never consult a pool. -fn routes_from_trace(records: &[LayerTraceRecord], num_layers: usize) -> Routes { - let mut per_layer: Vec> = vec![Vec::new(); num_layers]; - for rec in records { - if rec.features.is_empty() || rec.layer >= num_layers { - continue; - } - let mut ranked: Vec<(usize, usize)> = - rec.features.iter().map(|f| (f.rank, f.feature)).collect(); - ranked.sort_unstable(); - per_layer[rec.layer] = ranked.into_iter().map(|(_, feat)| feat).collect(); - } - Arc::new(per_layer) -} - -fn kernels(trace: &[LayerTraceRecord]) -> String { - let mut names: Vec<&str> = trace - .iter() - .filter(|r| !r.features.is_empty()) - .map(|r| r.path) - .collect(); - names.sort_unstable(); - names.dedup(); - names.join("+") -} - -/// Parity: identical executed feature SET per layer (order-insensitive) -/// and residual deltas within [`PARITY_REL_TOL`]. -fn trace_parity(a: &[LayerTraceRecord], b: &[LayerTraceRecord]) -> Result<(), String> { - if a.len() != b.len() { - return Err(format!("record count {} vs {}", a.len(), b.len())); - } - for (ra, rb) in a.iter().zip(b.iter()) { - if ra.layer != rb.layer || ra.position != rb.position { - return Err(format!("identity L{}p{}", ra.layer, ra.position)); - } - let mut fa: Vec = ra.features.iter().map(|f| f.feature).collect(); - let mut fb: Vec = rb.features.iter().map(|f| f.feature).collect(); - fa.sort_unstable(); - fb.sort_unstable(); - if fa != fb { - return Err(format!( - "L{} feature set differs ({} vs {})", - ra.layer, - fa.len(), - fb.len() - )); - } - let (x, y) = (ra.residual_delta_norm, rb.residual_delta_norm); - let rel = (x - y).abs() / x.abs().max(y.abs()).max(f32::MIN_POSITIVE); - if rel > PARITY_REL_TOL { - return Err(format!("L{} residual {x} vs {y} (rel {rel:.2e})", ra.layer)); - } - } - Ok(()) -} - -fn main() { - let args: Vec = std::env::args().collect(); - let allow_dirty = args.iter().any(|a| a == "--allow-dirty"); - let vindex = args - .iter() - .skip(1) - .find(|a| !a.starts_with("--")) - .cloned() - .unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".to_string()); - - // ── Refuse to emit a bad number ────────────────────────────────── - let ac = on_ac_power(); - let load = load_average(); - let dirty = !ac || !(load < LOAD_LIMIT); - if dirty { - eprintln!("REFUSING TO MEASURE:"); - if !ac { - eprintln!(" - not on AC power (throttled; sparse-vs-dense is exactly the"); - eprintln!(" bandwidth/compute tradeoff throttling distorts)"); - } - if !(load < LOAD_LIMIT) { - eprintln!(" - 1-min load average {load:.2} exceeds {LOAD_LIMIT:.1} (contended box)"); - } - if !allow_dirty { - eprintln!( - "\n Fix the environment, or pass --allow-dirty to record PROVISIONAL numbers." - ); - std::process::exit(2); - } - eprintln!("\n --allow-dirty set: every number below is PROVISIONAL.\n"); - } - - let dir = std::path::PathBuf::from(&vindex); - let mut cb = larql_vindex::SilentLoadCallbacks; - eprintln!("Loading {vindex} ..."); - let mut weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index.load_interleaved_kquant(&dir).expect("interleaved"); - index.load_attn_kquant(&dir).expect("attn"); - let _ = index.load_lm_head_kquant(&dir); - index.load_down_features_q4k(&dir).expect("down sidecar"); - let _ = index.load_gate_vectors_q4(&dir); - assert!( - index.has_down_features_kquant(), - "feature-major down sidecar required" - ); - let tok = load_tokenizer(&dir).expect("tok"); - for layer in 0..weights.num_layers { - insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant attn"); - } - - let nl = weights.num_layers; - let features = index.num_features(nl.saturating_sub(1)); - let full = tok.encode(PROMPT, true).expect("enc").get_ids().to_vec(); - let ids = vec![*full.last().unwrap()]; - - println!( - "\nR4 zero-out (paired, interleaved) — {nl} layers, {features} features/layer, seq_len=1\n\ - {REPEATS} repeats × {BLOCK}-forward blocks, rotated arm order, {WARMUP} warmup/arm\n\ - AC={ac} load={load:.2}{}\n", - if dirty { " [PROVISIONAL]" } else { "" } - ); - - for (band, k) in CELLS { - let sf = sparse_from(nl, band); - let base = WalkFfnConfig::hybrid(nl, sf, k); - - let all_pool: Routes = Arc::new( - (0..nl) - .map(|l| { - if l >= sf { - (0..index.num_features(l)).collect() - } else { - Vec::new() - } - }) - .collect(), - ); - let exact_cfg = base - .clone() - .with_pool_per_layer(all_pool) - .with_precomputed_routing(true) - .with_rank_within_pool(true); - - // Routes captured ONCE per cell and held fixed across repeats. - let cap_ffn = WalkFfn::from_config(&weights, &index, exact_cfg.clone()).with_trace(); - let cap_token = predict_with_ffn(&weights, &tok, &ids, 1, &cap_ffn) - .token_ids - .first() - .copied() - .unwrap_or(0); - let exact_trace = cap_ffn.take_runtime_trace(); - let routes = routes_from_trace(&exact_trace, nl); - let routed_layers = routes.iter().filter(|r| !r.is_empty()).count(); - - let oracle_par_cfg = base - .clone() - .with_pool_per_layer(routes.clone()) - .with_precomputed_routing(true) - .with_rank_within_pool(true); - let oracle_gath_cfg = base - .clone() - .with_pool_per_layer(routes) - .with_precomputed_routing(true); - - // ── Parity, untimed, before any timing is believed ─────────── - let mut parity = Vec::new(); - for (name, cfg) in [ - ("oracle-par", &oracle_par_cfg), - ("oracle-gath", &oracle_gath_cfg), - ] { - let pf = WalkFfn::from_config(&weights, &index, cfg.clone()).with_trace(); - let tkn = predict_with_ffn(&weights, &tok, &ids, 1, &pf) - .token_ids - .first() - .copied() - .unwrap_or(0); - let tr = pf.take_runtime_trace(); - let verdict = match trace_parity(&exact_trace, &tr) { - Ok(()) if tkn == cap_token => "parity ✓".to_string(), - Ok(()) => format!("token {cap_token}≠{tkn}"), - Err(e) => format!("DIFF {e}"), - }; - parity.push(format!("{name}[{}] {verdict}", kernels(&tr))); - } - - // ── Build all arms once; caches persist across repeats ─────── - let timings: Vec> = (0..4) - .map(|_| Arc::new(PhaseTimingsHandle::default())) - .collect(); - let dense_ffn = - WalkFfn::new_unlimited(&weights, &index).with_phase_timings(timings[ARM_DENSE].clone()); - let exact_ffn = WalkFfn::from_config(&weights, &index, exact_cfg) - .with_phase_timings(timings[ARM_EXACT].clone()); - let par_ffn = WalkFfn::from_config(&weights, &index, oracle_par_cfg) - .with_phase_timings(timings[ARM_ORACLE_PAR].clone()); - let gath_ffn = WalkFfn::from_config(&weights, &index, oracle_gath_cfg) - .with_phase_timings(timings[ARM_ORACLE_GATH].clone()); - let arms: [&dyn larql_inference::ffn::FfnBackend; 4] = - [&dense_ffn, &exact_ffn, &par_ffn, &gath_ffn]; - - let run_block = |arm: &dyn larql_inference::ffn::FfnBackend| -> (f64, u32) { - let t = Instant::now(); - let mut last = 0u32; - for _ in 0..BLOCK { - last = predict_with_ffn(&weights, &tok, &ids, 1, arm) - .token_ids - .first() - .copied() - .unwrap_or(0); - } - (t.elapsed().as_micros() as f64 / BLOCK as f64, last) - }; - - // Warm every arm and shape before recording anything. - for arm in arms { - for _ in 0..WARMUP { - let _ = predict_with_ffn(&weights, &tok, &ids, 1, arm); - } - } - for t in &timings { - t.gate_knn_ns.store(0, Ordering::Relaxed); - t.parallel_scan_ns.store(0, Ordering::Relaxed); - t.calls.store(0, Ordering::Relaxed); - } - - // ── Paired, rotated repeats ────────────────────────────────── - let mut samples: [Vec; 4] = Default::default(); - let mut tokens = [0u32; 4]; - for repeat in 0..REPEATS { - for slot in 0..4 { - let arm = (slot + repeat) % 4; // rotate order every repeat - let (us, tkn) = run_block(arms[arm]); - samples[arm].push(us); - tokens[arm] = tkn; - } - } - - let drift = sentinel_drift(&samples[ARM_DENSE]); - let spread = sentinel_spread(&samples[ARM_DENSE]); - let dense_sum = summarise(&samples[ARM_DENSE]); - let invalid = drift.abs() > SENTINEL_DRIFT_LIMIT || spread > SENTINEL_SPREAD_LIMIT; - - println!( - "── {} K={k} ({routed_layers} routed layers) dense {:.0} µs [{:.0}–{:.0}] \ - {:.1} tok/s sentinel drift {:+.1}% spread {:.1}%{}", - band_label(band), - dense_sum.median, - dense_sum.min, - dense_sum.max, - 1e6 / dense_sum.median, - drift * 100.0, - spread * 100.0, - if invalid { " ⚠ CELL DISTURBED" } else { "" } - ); - for p in &parity { - println!(" {p}"); - } - println!( - " {:<12} {:>9} {:>28} {:>9} {:>5}", - "arm", "median µs", "paired ratio vs dense (med [IQR])", "sel µs/L", "top1" - ); - for arm in [ARM_EXACT, ARM_ORACLE_PAR, ARM_ORACLE_GATH] { - // Paired ratios: computed within a repeat so drift cancels. - let ratios: Vec = samples[arm] - .iter() - .zip(samples[ARM_DENSE].iter()) - .map(|(a, d)| d / a) - .collect(); - let r = summarise(&ratios); - let s = summarise(&samples[arm]); - let calls = timings[arm].calls.load(Ordering::Relaxed); - let sel = if calls == 0 { - f64::NAN - } else { - timings[arm].gate_knn_ns.load(Ordering::Relaxed) as f64 / calls as f64 / 1000.0 - }; - println!( - " {:<12} {:>9.0} {:>18.3}× [{:.3}–{:.3}] {:>9.1} {:>5}", - ARM_NAMES[arm], - s.median, - r.median, - r.p25, - r.p75, - sel, - if tokens[arm] == tokens[ARM_DENSE] { - "✓" - } else { - "✗" - }, - ); - } - println!(); - } - - println!(" Analytic per-layer work (N = {features} features):"); - println!( - " {:<8} {:>12} {:>16} {:>14}", - "K", "exec rows", "selection rows", "exec/dense" - ); - for k in [2048usize, 4096, 6144] { - println!( - " {:<8} {:>12} {:>16} {:>13.2}×", - k, - 3 * k, - features, - (3 * k) as f64 / (3 * features) as f64 - ); - } - println!( - "\n Row-count ceiling on the sparse band: {:.2}× at K=2048, {:.2}× at K=4096,\n \ - {:.2}× at K=6144. A paired ratio far below its ceiling is gather\n \ - inefficiency, not selection cost — which no router can recover.", - features as f64 / 2048.0, - features as f64 / 4096.0, - features as f64 / 6144.0 - ); -} diff --git a/crates/larql-inference/examples/walk_ffn_temporal_reuse.rs b/crates/larql-inference/examples/walk_ffn_temporal_reuse.rs deleted file mode 100644 index 09b287cf9..000000000 --- a/crates/larql-inference/examples/walk_ffn_temporal_reuse.rs +++ /dev/null @@ -1,226 +0,0 @@ -//! #27 Temporal cursor reuse — TEMPORAL structure dense BLAS can't see. -//! -//! GUARDRAIL: the proven cosine >0.999 is SPATIAL (layer L vs L+1, same token). -//! This measures TEMPORAL — fixed layer L, token N vs N+1. `predict_with_ffn_trace` -//! captures each layer's LAST-POSITION input residual; teacher-forcing prefix[..i] -//! makes that token i-1's residual attending to real history (= a KV-cached decode -//! step). So consecutive i give token-to-token at fixed layer. NOT within-prefill -//! cross-position (that would be spatial wearing a temporal label). -//! -//! Three mechanisms, per zone, distribution (median/p90/worst) not mean: -//! (a) token-to-token residual COSINE → output reuse (needs ≈1.0) -//! (b) gate-KNN active-pool JACCARD → route reuse (needs ≥0.9) -//! (c) TwoNN intrinsic-dim of the DELTA → delta-walk (needs ≤~30 on highway) -//! -//! Usage: `cargo run --release --example walk_ffn_temporal_reuse -- [VINDEX]` - -use larql_inference::load_tokenizer; -use larql_inference::research::predict_with_ffn_trace; -use larql_inference::vindex::{insert_q4k_layer_tensors_resident, WalkFfn}; -use ndarray::Array1; - -const K: usize = 2048; // active-pool size for Jaccard - -fn cosine(a: &[f32], b: &[f32]) -> f64 { - let (mut dot, mut na, mut nb) = (0.0f64, 0.0f64, 0.0f64); - for (&x, &y) in a.iter().zip(b) { - dot += x as f64 * y as f64; - na += (x as f64) * (x as f64); - nb += (y as f64) * (y as f64); - } - if na == 0.0 || nb == 0.0 { - return 0.0; - } - dot / (na.sqrt() * nb.sqrt()) -} - -fn jaccard(a: &[usize], b: &[usize]) -> f64 { - use std::collections::HashSet; - let sa: HashSet = a.iter().copied().collect(); - let sb: HashSet = b.iter().copied().collect(); - let inter = sa.intersection(&sb).count(); - let uni = sa.union(&sb).count().max(1); - inter as f64 / uni as f64 -} - -/// TwoNN intrinsic-dimension MLE (Facco et al.): for each point, μ = r2/r1 -/// (2nd/1st nearest-neighbour distance); d = M / Σ ln(μ). Brute-force NN. -fn twonn(points: &[Vec]) -> f64 { - let m = points.len(); - if m < 10 { - return f64::NAN; - } - let mut sum_ln_mu = 0.0f64; - let mut used = 0usize; - for i in 0..m { - let (mut r1, mut r2) = (f64::INFINITY, f64::INFINITY); - for j in 0..m { - if i == j { - continue; - } - let mut d = 0.0f64; - for (&x, &y) in points[i].iter().zip(&points[j]) { - let e = x as f64 - y as f64; - d += e * e; - } - let d = d.sqrt(); - if d < r1 { - r2 = r1; - r1 = d; - } else if d < r2 { - r2 = d; - } - } - if r1 > 1e-9 && r2.is_finite() { - sum_ln_mu += (r2 / r1).ln(); - used += 1; - } - } - if sum_ln_mu <= 0.0 { - return f64::NAN; - } - used as f64 / sum_ln_mu -} - -fn pctile(v: &mut [f64], q: f64) -> f64 { - if v.is_empty() { - return f64::NAN; - } - v.sort_by(|a, b| a.total_cmp(b)); - v[(((v.len() - 1) as f64) * q).round() as usize] -} - -fn zone(layer: usize) -> usize { - match layer { - 0..=4 => 0, - 5..=20 => 1, - 21..=29 => 2, - _ => 3, - } -} - -fn main() { - let args: Vec = std::env::args().collect(); - let vindex = args - .get(1) - .cloned() - .unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".to_string()); - let dir = std::path::PathBuf::from(&vindex); - let mut cb = larql_vindex::SilentLoadCallbacks; - eprintln!("Loading {vindex} ..."); - let mut weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index.load_interleaved_kquant(&dir).expect("interleaved"); - index.load_attn_kquant(&dir).expect("attn"); - let _ = index.load_lm_head_kquant(&dir); - let _ = index.load_gate_vectors_q4(&dir); - let tok = load_tokenizer(&dir).expect("tok"); - for layer in 0..weights.num_layers { - insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant attn"); - } - let nl = weights.num_layers; - - let passages = [ - "The expedition had been planned for years, but nothing prepared them for the silence of the ice that morning, and the captain wrote that the cold seemed to have a will of its own.", - "She had always believed that cities were built from ambition, but walking the old quarter at dusk she understood they were built from compromise, one stubborn refusal at a time.", - "Economists argue about the cause, yet the pattern repeats: cheap credit, a frenzy of building, a sudden loss of nerve, and then the long quiet years of paying it all back.", - "The first thing the new recruits learned was not how to fight but how to wait, because the enemy they feared most was the boredom that made careful people careless.", - "Light from the distant galaxy had travelled for billions of years to reach the telescope, carrying news of an explosion that had happened before the sun itself was born.", - "He cooked the way his grandmother had taught him, never measuring, tasting constantly, trusting that the dish would tell him what it needed if he only paid attention.", - ]; - - // Per layer: cosines, jaccards, and the delta population. - let mut cos_by_layer: Vec> = vec![Vec::new(); nl]; - let mut jac_by_layer: Vec> = vec![Vec::new(); nl]; - let mut delta_by_layer: Vec>> = vec![Vec::new(); nl]; - - for (pi, p) in passages.iter().enumerate() { - let ids = tok.encode(*p, true).expect("enc").get_ids().to_vec(); - let n = ids.len().min(40); - eprintln!(" passage {}/{} ({} tokens) ...", pi + 1, passages.len(), n); - // R[i] = per-layer input residual at the last token of prefix ids[..i+1]. - // Start at i0 so attention has real history (skip the very first tokens). - let i0 = 3usize; - let mut prev: Option>> = None; - for i in i0..n { - let r = predict_with_ffn_trace( - &weights, - &tok, - &ids[..=i], - 1, - &WalkFfn::new_unlimited(&weights, &index), - ); - let cur = r.residuals; // Vec>, one per layer - if let Some(pr) = &prev { - for l in 0..nl.min(cur.len()).min(pr.len()) { - cos_by_layer[l].push(cosine(&pr[l], &cur[l])); - // active pools (gate-KNN top-K) at each step - let pa: Vec = index - .gate_knn(l, &Array1::from(pr[l].clone()), K) - .into_iter() - .map(|(f, _)| f) - .collect(); - let pb: Vec = index - .gate_knn(l, &Array1::from(cur[l].clone()), K) - .into_iter() - .map(|(f, _)| f) - .collect(); - jac_by_layer[l].push(jaccard(&pa, &pb)); - let d: Vec = pr[l].iter().zip(&cur[l]).map(|(&a, &b)| b - a).collect(); - delta_by_layer[l].push(d); - } - } - prev = Some(cur); - } - } - - // Aggregate by zone. - let znames = [ - "pre-commit L0-4", - "highway L5-20", - "retrieval L21-29", - "format L30-33", - ]; - println!("\n#27 Temporal cursor reuse — token-to-token at fixed layer, real history\n"); - println!( - "{:<20} {:>20} {:>20} {:>14}", - "zone", "residual cosine", "pool Jaccard", "delta TwoNN" - ); - println!( - "{:<20} {:>20} {:>20} {:>14}", - "", "med / p10 / worst", "med / p10 / worst", "median dim" - ); - for (z, zname) in znames.iter().enumerate() { - let layers: Vec = (0..nl).filter(|&l| zone(l) == z).collect(); - let mut cos_all: Vec = layers - .iter() - .flat_map(|&l| cos_by_layer[l].clone()) - .collect(); - let mut jac_all: Vec = layers - .iter() - .flat_map(|&l| jac_by_layer[l].clone()) - .collect(); - // Per-layer TwoNN (cap population for brute-force NN), then zone median. - let mut dims: Vec = Vec::new(); - for &l in &layers { - let pts = &delta_by_layer[l]; - let cap = pts.len().min(220); - let d = twonn(&pts[..cap]); - if d.is_finite() { - dims.push(d); - } - } - let cos_med = pctile(&mut cos_all.clone(), 0.50); - let cos_p10 = pctile(&mut cos_all.clone(), 0.10); - let cos_worst = pctile(&mut cos_all, 0.0); - let jac_med = pctile(&mut jac_all.clone(), 0.50); - let jac_p10 = pctile(&mut jac_all.clone(), 0.10); - let jac_worst = pctile(&mut jac_all, 0.0); - let dim_med = pctile(&mut dims, 0.50); - println!( - "{:<20} {:>6.3}/{:>5.3}/{:>5.3} {:>7.3}/{:>5.3}/{:>5.3} {:>14.1}", - zname, cos_med, cos_p10, cos_worst, jac_med, jac_p10, jac_worst, dim_med - ); - } - println!("\n PRE-REGISTERED reading: highway delta-dim ≤~30 ⇒ delta-walk LIVE; cosine ≥0.995 ⇒ output reuse;\n Jaccard ≥0.9 + low cosine ⇒ route-reuse only; all churn ⇒ axis closes.\n worst = min over step-pairs (a reuse scheme catastrophic on 10% of steps is a drift generator)."); -} diff --git a/crates/larql-inference/examples/walk_ffn_v1_hash_routing.rs b/crates/larql-inference/examples/walk_ffn_v1_hash_routing.rs deleted file mode 100644 index f0ca930eb..000000000 --- a/crates/larql-inference/examples/walk_ffn_v1_hash_routing.rs +++ /dev/null @@ -1,552 +0,0 @@ -//! V1 — hash routing across all layers (aim-validation P0; resolves KU4). -//! -//! Exp 27 measured cheap/hash FFN routing on **Gemma 3 4B, layer 0 only**: -//! top-2048 (~20% of d_ffn) → next-token KL ≈ 0.030. The medium-term ~80% -//! confidence in `ROADMAP.md` ("Gemma 4 26B-A4B ≥10 tok/s on 64 GB, no GPU") -//! rests on a **5× FFN bandwidth reduction** that *assumes that one-layer result -//! compounds across all layers and survives at the end-to-end output*. V1 tests -//! it. The strong prior (#17–#28: the FFN is dense, faithful K≈4096) is that the -//! per-layer threshold balloons at depth and the 5× claim shrinks — V1 is the -//! rigorous measurement that confirms/quantifies that, not a new kernel. -//! -//! Judged only in **predictive units** (KL bits, NLL bits/token, argmax drift) — -//! never cosine. Per-layer KL ≤ 0.05 is a SCREENING proxy; the claim gate is the -//! compounding stage (Phase B) where all per-layer thresholds are applied at once -//! and judged on held-text NLL distribution + drift (the #26 lesson: single-step -//! KL once *inverted* the ship decision). -//! -//! Phases (cheap-first — full verdict on one model before cross-arch): -//! - Step 0 parity anchor: full-K walk ≈ dense (KL≈0); gate-KNN k=2048 @ L0 only ≈ exp-27 KL. -//! - Phase A per-layer oracle threshold: min k (gate top-k) for output-KL ≤ 0.05, one layer sparse at a time. -//! -//! (Phases B/C — compounding + cheap-routing realizability — land next.) -//! -//! Usage: `cargo run --release --example walk_ffn_v1_hash_routing -- [VINDEX] [--quick]` - -use larql_inference::vindex::{insert_q4k_layer_tensors_resident, WalkFfn, WalkFfnConfig}; -use larql_inference::{load_tokenizer, predict_with_ffn}; -use larql_models::ModelWeights; -use std::collections::HashMap; -use std::sync::Arc; - -/// Full-vocab next-token distribution keyed by token id (last position), from a -/// forward pass with the given FFN backend. `predict_with_ffn` softmaxes over -/// the whole vocab and (with a huge top_k) returns every token. -/// (Mirrors `walk_ffn_accuracy.rs::next_token_dist`.) -fn next_token_dist( - weights: &ModelWeights, - tok: &tokenizers::Tokenizer, - ids: &[u32], - ffn: &dyn larql_inference::ffn::FfnBackend, -) -> HashMap { - let r = predict_with_ffn(weights, tok, ids, usize::MAX, ffn); - r.token_ids - .into_iter() - .zip(r.predictions.into_iter().map(|(_, p)| p)) - .collect() -} - -/// KL(P‖Q) in **bits**, plus top-1 agreement. P = dense ground truth, Q = -/// candidate. (Mirrors `walk_ffn_accuracy.rs::compare`.) -fn kl_bits(p: &HashMap, q: &HashMap) -> (f64, bool) { - let eps = 1e-12; - let mut kl = 0.0; - for (&id, &pi) in p { - if pi <= 0.0 { - continue; - } - let qi = q.get(&id).copied().unwrap_or(0.0).max(eps); - kl += pi * (pi.max(eps) / qi).ln(); - } - let p_arg = p.iter().max_by(|a, b| a.1.total_cmp(b.1)).map(|(i, _)| *i); - let q_arg = q.iter().max_by(|a, b| a.1.total_cmp(b.1)).map(|(i, _)| *i); - ( - kl / std::f64::consts::LN_2, - p_arg.is_some() && p_arg == q_arg, - ) -} - -/// Config with exactly one layer sparsified at top-`k` (gate-score selection, -/// the accuracy ceiling for any size-k route); all other layers dense. This -/// isolates layer `l`'s contribution to output divergence — the per-layer KL. -fn one_layer_sparse(num_layers: usize, l: usize, k: usize) -> WalkFfnConfig { - let mut cfg = WalkFfnConfig::dense(num_layers); - cfg.k_per_layer[l] = Some(k); - cfg -} - -/// Teacher-forced per-token NLL (bits = −log2 p(true next token)) over `ids`, -/// plus the per-position argmax token (for the flip rate). Scores positions -/// `1..ids.len()`. (Mirrors `walk_ffn_nll.rs::token_nlls`.) -fn token_nlls( - weights: &ModelWeights, - tok: &tokenizers::Tokenizer, - ids: &[u32], - ffn: &dyn larql_inference::ffn::FfnBackend, -) -> (Vec, Vec) { - let (mut nlls, mut args) = (Vec::new(), Vec::new()); - for i in 1..ids.len() { - if i % 8 == 0 { - eprint!("\r [{}] pos {i}/{} ", ffn.name(), ids.len()); - } - let r = predict_with_ffn(weights, tok, &ids[..i], usize::MAX, ffn); - let dist: HashMap = r - .token_ids - .iter() - .copied() - .zip(r.predictions.iter().map(|(_, p)| *p)) - .collect(); - let p = dist.get(&ids[i]).copied().unwrap_or(0.0).max(1e-12); - nlls.push(-p.log2()); - args.push(r.token_ids.first().copied().unwrap_or(0)); - } - (nlls, args) -} - -fn pct(sorted: &[f64], q: f64) -> f64 { - if sorted.is_empty() { - return 0.0; - } - let idx = ((sorted.len() as f64 - 1.0) * q).round() as usize; - sorted[idx.min(sorted.len() - 1)] -} - -fn mean(v: &[f64]) -> f64 { - v.iter().sum::() / v.len().max(1) as f64 -} - -/// Current git revision (short), for the artifact provenance. Best-effort. -fn git_rev() -> String { - std::process::Command::new("git") - .args(["rev-parse", "--short", "HEAD"]) - .output() - .ok() - .filter(|o| o.status.success()) - .map(|o| String::from_utf8_lossy(&o.stdout).trim().to_string()) - .unwrap_or_else(|| "unknown".to_string()) -} - -fn main() { - let args: Vec = std::env::args().collect(); - let vindex = args - .iter() - .skip(1) - .find(|a| !a.starts_with("--")) - .cloned() - .unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".to_string()); - let quick = args.iter().any(|a| a == "--quick"); - let kl_thresh = 0.05_f64; - let dir = std::path::PathBuf::from(&vindex); - let mut cb = larql_vindex::SilentLoadCallbacks; - - eprintln!("Loading {vindex} ..."); - let mut weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index - .load_interleaved_kquant(&dir) - .expect("interleaved kquant"); - index.load_attn_kquant(&dir).expect("attn kquant"); - let _ = index.load_lm_head_kquant(&dir); - let _ = index.load_down_features_q4k(&dir); - let _ = index.load_down_features(&dir); - let _ = index.load_gate_vectors_q4(&dir); - let tok = load_tokenizer(&dir).expect("tokenizer"); - - // `predict_with_ffn` reads attention from f32 `weights.vectors`; the Q4K - // loader leaves attention quantised. Dequantise it up front so the FFN - // router stays the only variable (same fix as walk_ffn_accuracy.rs). - let nl = weights.num_layers; - eprintln!("Dequantising attention for {nl} layers ..."); - for layer in 0..nl { - insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant layer"); - } - - // Matrix `baseline_fact_prompts` (bench/aim-validation/matrix.json) — the - // canonical V1 screening set; short, entropic next-token choices. - let prompts = [ - "The capital of France is", - "The largest planet in the solar system is", - "The author of Pride and Prejudice was", - "The chemical symbol for gold is", - ]; - let ids: Vec> = prompts - .iter() - .map(|p| tok.encode(*p, true).expect("encode").get_ids().to_vec()) - .collect(); - - // Dense ground-truth distributions, computed ONCE and reused across the - // whole sweep (the expensive part is the forward passes). - eprintln!("Computing dense reference distributions ..."); - let dense_dists: Vec> = ids - .iter() - .map(|id| { - next_token_dist( - &weights, - &tok, - id, - &WalkFfn::new_unlimited(&weights, &index), - ) - }) - .collect(); - - // Average per-layer-sparse KL over prompts at a given (layer, k). - let avg_kl = |l: usize, k: usize| -> (f64, f64) { - let (mut kl, mut agree, mut n) = (0.0, 0usize, 0usize); - let cfg = one_layer_sparse(nl, l, k); - let ffn = WalkFfn::from_config(&weights, &index, cfg); - for (id, dense) in ids.iter().zip(dense_dists.iter()) { - let q = next_token_dist(&weights, &tok, id, &ffn); - let (b, a) = kl_bits(dense, &q); - kl += b; - agree += a as usize; - n += 1; - } - let n = n.max(1) as f64; - (kl / n, 100.0 * agree as f64 / n) - }; - - let feats0 = index.num_features(0); - let mid = nl / 2; - - // ── Step 0: parity anchor (the spine, before any threshold) ──────────── - println!("\n=== Step 0: parity anchor — {vindex} ==="); - { - // (a) full-K single-layer sparse should equal dense (KL≈0): confirms - // the sparse-walk path is faithful before we read any sparsity. - let (kl_full, ag_full) = avg_kl(mid, index.num_features(mid)); - println!( - " full-K @ L{mid} (walk == dense?): KL={kl_full:>8.5} bits agree={ag_full:.0}% (expect ~0)" - ); - // (b) exp-27 anchor: gate-KNN top-2048 @ L0 ONLY → expect KL ≈ 0.030. - let k0 = 2048.min(feats0); - let (kl_l0, ag_l0) = avg_kl(0, k0); - let pct0 = 100.0 * k0 as f64 / feats0.max(1) as f64; - println!( - " exp-27 @ L0 k={k0} ({pct0:.0}% of {feats0}): KL={kl_l0:>8.5} bits agree={ag_l0:.0}% (exp 27: ~0.030)" - ); - } - - // ── Phase A: per-layer oracle threshold table ────────────────────────── - // For each layer, the minimum k (gate top-k) for output-KL ≤ 0.05. - // Geometric fraction grid gives the whole curve, not just the crossing. - let smoke = args.iter().any(|a| a == "--smoke"); - let json_path = args - .iter() - .find_map(|a| a.strip_prefix("--json=").map(|s| s.to_string())) - .unwrap_or_else(|| { - let stem = std::path::Path::new(&vindex) - .file_stem() - .and_then(|s| s.to_str()) - .unwrap_or("model"); - format!("v1_{stem}.json") - }); - - let fracs = [ - 1.0 / 64.0, - 1.0 / 32.0, - 1.0 / 16.0, - 1.0 / 8.0, - 1.0 / 4.0, - 1.0 / 2.0, - 1.0, - ]; - - // Per-layer threshold table. `thr` is one row per measured layer; `thr_k` - // is the full-length (nl) k schedule used by Phase B compounding. - // (l, frac, k, kl) - let mut thr: Vec<(usize, f64, usize, f64)> = Vec::new(); - let mut thr_k: Vec = (0..nl).map(|l| index.num_features(l)).collect(); - - if smoke { - // Skip the (slow) per-layer sweep: uniform k = feats/8 everywhere, just - // to exercise the Phase B + JSON code paths cheaply. - eprintln!("[--smoke] skipping Phase A sweep; uniform k = feats/8"); - for (l, slot) in thr_k.iter_mut().enumerate() { - let feats = index.num_features(l); - *slot = (feats / 8).max(1); - thr.push((l, 0.125, *slot, f64::NAN)); - } - } else { - let layers: Vec = if quick { - vec![0, mid, nl.saturating_sub(1)] - } else { - (0..nl).collect() - }; - println!( - "\n=== Phase A: per-layer oracle threshold (min k-frac for KL ≤ {kl_thresh}) ===\n" - ); - println!( - "{:>5} {:>8} {:>10} {:>10} {:>9}", - "layer", "feats", "thr-k", "thr-frac", "thr-KL" - ); - for &l in &layers { - let feats = index.num_features(l); - let mut chosen: Option<(usize, f64, f64)> = None; // (k, frac, kl) - for &f in &fracs { - let k = ((f * feats as f64).round() as usize).clamp(1, feats); - let (kl, _agree) = avg_kl(l, k); - if kl <= kl_thresh { - chosen = Some((k, f, kl)); - break; - } - } - let (k, f, kl) = chosen.unwrap_or((feats, 1.0, f64::NAN)); - println!("{l:>5} {feats:>8} {k:>10} {f:>10.4} {kl:>9.5}"); - thr.push((l, f, k, kl)); - thr_k[l] = k; - } - // In quick mode we only measured 3 layers — the compounding stage and - // bandwidth accounting need every layer, so stop after the screen. - if quick { - return; - } - let mean_frac = mean(&thr.iter().map(|(_, f, _, _)| *f).collect::>()); - println!( - "\n mean threshold fraction = {mean_frac:.4} (per-layer SCREEN only; bandwidth + claim gate below)" - ); - } - - // ── Bandwidth accounting (honest: gate projection is NOT free) ───────── - // "rows touched / token" in units of hidden-sized weight rows. Dense reads - // all gate+up+down rows. Two sparse regimes differ in the GATE cost: - // cheap-route (precomputed_routing): gate scored for only k pool feats → 3·k rows - // gate-oracle (joint_gate_knn): full gate projection to rank → feats + 2·k rows - // Phase B uses gate-oracle selection, so its realised saving is the oracle - // line; the 5× *claim* is only reachable on the cheap-route line (Phase C - // must show cheap routing matches the oracle's KL at these k's). - let dense_rows: f64 = (0..nl).map(|l| 3.0 * index.num_features(l) as f64).sum(); - let cheap_rows: f64 = (0..nl).map(|l| 3.0 * thr_k[l] as f64).sum(); - let oracle_rows: f64 = (0..nl) - .map(|l| index.num_features(l) as f64 + 2.0 * thr_k[l] as f64) - .sum(); - let cheap_frac = cheap_rows / dense_rows; - let oracle_frac = oracle_rows / dense_rows; - let cheap_factor = if cheap_frac > 0.0 { - 1.0 / cheap_frac - } else { - 0.0 - }; - let oracle_factor = if oracle_frac > 0.0 { - 1.0 / oracle_frac - } else { - 0.0 - }; - - println!("\n=== Bandwidth (FFN weight rows touched per token, vs dense) ==="); - println!( - " cheap-route (precomputed): {:.4}× of dense → {cheap_factor:.2}× reduction (the 5× claim's best case)", - cheap_frac - ); - println!( - " gate-oracle (Phase B cfg): {:.4}× of dense → {oracle_factor:.2}× reduction (full gate projection still paid)", - oracle_frac - ); - - // ── Phase B: compounding — all per-layer thresholds applied at once ──── - // Claim gate (#26 lesson): held-text NLL distribution + argmax-drift, not - // single-step KL. Compares dense vs the compounded gate-oracle schedule. - println!("\n=== Phase B: compounding — held-text NLL + drift (all layers @ threshold) ==="); - let passage = "The expedition had been planned for years, but nothing prepared \ -them for the silence of the ice. Each morning the wind died at dawn, and the only \ -sound was the slow groan of the glacier shifting beneath their tents. Provisions \ -were running low, and the captain knew that another week of delay would mean \ -turning back without ever reaching the plateau they had crossed two oceans to find."; - let pids = tok.encode(passage, true).expect("enc").get_ids().to_vec(); - eprintln!(" held passage: {} tokens", pids.len()); - - let mut comp_cfg = WalkFfnConfig::dense(nl); - for (l, &k) in thr_k.iter().enumerate() { - comp_cfg.k_per_layer[l] = Some(k); - } - - let t0 = std::time::Instant::now(); - let (nll_d, arg_d) = token_nlls( - &weights, - &tok, - &pids, - &WalkFfn::new_unlimited(&weights, &index), - ); - let dense_fps = (pids.len().saturating_sub(1)) as f64 / t0.elapsed().as_secs_f64().max(1e-9); - let t1 = std::time::Instant::now(); - let (nll_c, arg_c) = token_nlls( - &weights, - &tok, - &pids, - &WalkFfn::from_config(&weights, &index, comp_cfg), - ); - let comp_fps = (pids.len().saturating_sub(1)) as f64 / t1.elapsed().as_secs_f64().max(1e-9); - eprintln!( - "\r scored {} positions (dense + compounded) ", - nll_d.len() - ); - - let (mut sd, mut sc) = (nll_d.clone(), nll_c.clone()); - sd.sort_by(|a, b| a.total_cmp(b)); - sc.sort_by(|a, b| a.total_cmp(b)); - let (md, mc) = (mean(&nll_d), mean(&nll_c)); - let flips = arg_d.iter().zip(&arg_c).filter(|(a, b)| a != b).count(); - let flip_pct = 100.0 * flips as f64 / arg_d.len().max(1) as f64; - let first_div = arg_d - .iter() - .zip(&arg_c) - .position(|(a, b)| a != b) - .map(|p| p as i64) - .unwrap_or(-1); - // Perplexity from mean NLL in bits: ppl = 2^mean_bits. - let (ppl_d, ppl_c) = (2f64.powf(md), 2f64.powf(mc)); - let ppl_delta_pct = (ppl_c / ppl_d - 1.0) * 100.0; - - println!( - " NLL bits/token dense: mean {md:.3} p90 {:.3} max {:.3}", - pct(&sd, 0.90), - pct(&sd, 1.0) - ); - println!( - " NLL bits/token comp : mean {mc:.3} p90 {:.3} max {:.3} Δmean {:+.3}", - pct(&sc, 0.90), - pct(&sc, 1.0), - mc - md - ); - println!(" perplexity dense {ppl_d:.3} → comp {ppl_c:.3} ({ppl_delta_pct:+.2}%)"); - println!(" argmax drift (comp vs dense): {flip_pct:.1}% first-divergence pos: {first_div}"); - println!(" forward/s (proxy tok/s) dense {dense_fps:.2} comp {comp_fps:.2}"); - - // ── Phase C: cheap-routing realizability ─────────────────────────────── - // Phase A's threshold is the gate-ORACLE lower bound on k. But realising it - // needs a route that DOESN'T pay the full gate projection (else no - // bandwidth saved). Test whether a CHEAP, content-blind precomputed route - // hits the same KL at the oracle threshold k: - // strided — uninformed lower bound (ignores the input entirely) - // ‖down‖ — informed but static: top-k by down-row norm (the features - // that move the residual most when active), as cheap as strided - // The gap (cheap-KL − oracle-KL) is the price of cheap routing; if cheap-KL - // blows past 0.05 the per-layer sparsity is NOT realizable cheaply and the - // 5× claim dies even though the oracle threshold was small (#19/#23 prior). - // Only run where the oracle threshold is small enough that it matters. - let avg_cheap_kl = |l: usize, k: usize, route: Vec| -> f64 { - let mut pools = vec![Vec::new(); nl]; - pools[l] = route; - let mut cfg = WalkFfnConfig::dense(nl); - cfg.k_per_layer[l] = Some(k); - cfg.pool_per_layer = Some(Arc::new(pools)); - cfg.precomputed_routing = true; - let ffn = WalkFfn::from_config(&weights, &index, cfg); - let mut kl = 0.0; - for (id, dense) in ids.iter().zip(dense_dists.iter()) { - kl += kl_bits(dense, &next_token_dist(&weights, &tok, id, &ffn)).0; - } - kl / ids.len().max(1) as f64 - }; - let probe = WalkFfn::new_unlimited(&weights, &index); - // (layer, k, oracle_kl, strided_kl, static_kl) - let mut phase_c: Vec<(usize, usize, f64, f64, f64)> = Vec::new(); - let c_layers: Vec<&(usize, f64, usize, f64)> = - thr.iter().filter(|(_, f, _, _)| *f <= 0.25).collect(); - println!( - "\n=== Phase C: cheap-route realizability @ oracle thresholds ({} layers, frac ≤ 0.25) ===", - c_layers.len() - ); - println!( - "{:>5} {:>8} {:>10} {:>10} {:>10}", - "layer", "k", "oracle-KL", "strided", "‖down‖" - ); - for (l, _f, k, kl_oracle) in c_layers { - let (l, k) = (*l, *k); - let feats = index.num_features(l); - let stride = (feats / k.max(1)).max(1); - let strided: Vec = (0..k).map(|i| (i * stride) % feats).collect(); - let static_route: Vec = match probe.down_row_norms_pub(l) { - Some(norms) => { - let mut idx: Vec = (0..norms.len()).collect(); - idx.sort_unstable_by(|&a, &b| norms[b].total_cmp(&norms[a])); - idx.truncate(k); - idx - } - None => strided.clone(), - }; - let kl_strided = avg_cheap_kl(l, k, strided); - let kl_static = avg_cheap_kl(l, k, static_route); - println!("{l:>5} {k:>8} {kl_oracle:>10.5} {kl_strided:>10.5} {kl_static:>10.5}"); - phase_c.push((l, k, *kl_oracle, kl_strided, kl_static)); - } - let cheap_ok = phase_c - .iter() - .filter(|(_, _, _, _, s)| *s <= kl_thresh) - .count(); - let cheap_realizable_pct = if phase_c.is_empty() { - 0.0 - } else { - 100.0 * cheap_ok as f64 / phase_c.len() as f64 - }; - println!( - "\n cheap (‖down‖) route clears KL ≤ {kl_thresh} at {cheap_realizable_pct:.0}% of small-threshold layers" - ); - println!(" → where it does NOT, the per-layer sparsity needs the gate projection (no bandwidth win there)."); - - // ── JSON artifact (bench/aim-validation/matrix.json result contract) ─── - let model = std::path::Path::new(&vindex) - .file_stem() - .and_then(|s| s.to_str()) - .unwrap_or("model") - .to_string(); - let topk_rows: Vec = thr - .iter() - .map(|(l, f, k, kl)| { - let klv = if kl.is_nan() { - "null".to_string() - } else { - format!("{kl:.5}") - }; - format!("{{\"layer\":{l},\"k\":{k},\"frac\":{f:.5},\"kl\":{klv}}}") - }) - .collect(); - let phase_c_rows: Vec = phase_c - .iter() - .map(|(l, k, o, st, sd)| { - format!("{{\"layer\":{l},\"k\":{k},\"oracle_kl\":{o:.5},\"strided_kl\":{st:.5},\"down_norm_kl\":{sd:.5}}}") - }) - .collect(); - let json = format!( - concat!( - "{{\n", - " \"test_id\": \"V1\",\n", - " \"model\": \"{model}\",\n", - " \"prompt_set\": \"baseline_fact_prompts (KL) + held narrative (NLL)\",\n", - " \"git_rev\": \"{rev}\",\n", - " \"metrics\": {{\n", - " \"topk\": [{topk}],\n", - " \"kl_divergence\": {{\"threshold\": {kl_thresh}, \"n_prompts\": {nprompts}}},\n", - " \"perplexity_delta_pct\": {ppl_delta:.4},\n", - " \"nll_bits_dense_mean\": {md:.4},\n", - " \"nll_bits_comp_mean\": {mc:.4},\n", - " \"argmax_drift_pct\": {flip:.4},\n", - " \"first_divergence_pos\": {first_div},\n", - " \"bytes_touched_per_token\": {{\"cheap_frac\": {cheap_frac:.5}, \"cheap_factor\": {cheap_factor:.4}, \"oracle_frac\": {oracle_frac:.5}, \"oracle_factor\": {oracle_factor:.4}}},\n", - " \"cheap_route\": {{\"down_norm_realizable_pct\": {crp:.2}, \"by_layer\": [{pc}]}},\n", - " \"tok_per_s\": {{\"forward_per_s_dense\": {dfps:.4}, \"forward_per_s_comp\": {cfps:.4}, \"note\": \"full-forward proxy, no KV cache\"}}\n", - " }},\n", - " \"notes\": \"gate-oracle per-layer threshold; Phase B uses oracle selection (gate projection still paid); 5x claim needs Phase C cheap-route parity\"\n", - "}}\n" - ), - model = model, - rev = git_rev(), - topk = topk_rows.join(","), - kl_thresh = kl_thresh, - nprompts = prompts.len(), - ppl_delta = ppl_delta_pct, - md = md, - mc = mc, - flip = flip_pct, - first_div = first_div, - cheap_frac = cheap_frac, - cheap_factor = cheap_factor, - oracle_frac = oracle_frac, - oracle_factor = oracle_factor, - crp = cheap_realizable_pct, - pc = phase_c_rows.join(","), - dfps = dense_fps, - cfps = comp_fps, - ); - std::fs::write(&json_path, &json).expect("write json artifact"); - println!("\n artifact → {json_path}"); -} diff --git a/crates/larql-inference/examples/walk_ffn_v1_moe_within_expert.rs b/crates/larql-inference/examples/walk_ffn_v1_moe_within_expert.rs deleted file mode 100644 index 1431b9e39..000000000 --- a/crates/larql-inference/examples/walk_ffn_v1_moe_within_expert.rs +++ /dev/null @@ -1,528 +0,0 @@ -//! V1 (MoE-within-expert) — does feature/hash routing work INSIDE a single -//! MoE expert's FFN? (aim-validation P0 follow-up; resolves the OPEN half of KU4.) -//! -//! V1 (`walk_ffn_v1_hash_routing.rs`) tested hash routing *within a dense FFN* -//! and falsified it on 3 dense archs: per-layer KL ≤ 0.05 thresholds don't -//! compound (+5.4 to +7.7 bits/tok, 78–95 % drift). But on the Gemma 4 26B-A4B -//! the per-layer FFN block is **128 stacked experts**, not one dense FFN — so -//! that dense harness "measures the wrong object". This probe runs the SAME -//! three-phase protocol on a single expert's own `inter`-feature space: within -//! each routed expert, keep only the top-`k` of its post-activation features -//! (the values entering `down`) and measure the cost downstream. -//! -//! The open question: the expert feature space (~704) is ~6× smaller than the -//! dense d_ffn and load-balanced routing already concentrates work — does -//! within-expert sparsity survive where dense within-FFN sparsity didn't? -//! -//! Judged ONLY in predictive units (KL bits, NLL bits/token, argmax drift) — -//! never cosine. Per-layer KL ≤ 0.05 is a SCREEN; the claim gate is Phase B -//! (all expert layers pruned at once → held-text NLL + drift), per the #26 -//! lesson that single-step KL once *inverted* the ship decision -//! (`feedback_metric_matches_operation`). -//! -//! Mechanism: the prune is applied inside the production expert kernel via -//! `larql_compute::cpu::ops::moe::set_routing` (OFF by default → byte-exact -//! parity), so errors propagate through the real forward pass — no reimplemented -//! numerics (`feedback_engineering_vs_research_posture`: parity is the spine). -//! -//! Phases: -//! - Step 0 parity anchor: all-dense schedule == dense (KL≈0); one layer @ frac=0.5 bites. -//! - Phase A per-expert-layer oracle threshold: min keep-frac for output-KL ≤ 0.05, one layer at a time. -//! - Phase B compounding: all expert layers at threshold → held-text NLL + drift (the claim gate). -//! - Phase C cheap-route realizability: content-blind Strided vs the ActMagnitude oracle at the thresholds. -//! -//! Usage: `cargo run --release --example walk_ffn_v1_moe_within_expert -- [VINDEX] [--quick|--smoke] [--json=PATH]` - -use larql_compute::cpu::ops::moe::{set_routing, ExpertFeatureSelector, WithinExpertRouting}; -use larql_inference::load_tokenizer; -use larql_inference::vindex::predict_kquant; -use larql_models::ModelWeights; -use larql_vindex::VectorIndex; -use std::collections::HashMap; - -/// Full-vocab next-token distribution (last position) from a forward pass -/// under whatever within-expert routing is currently installed. Mirrors the -/// dense V1 harness's `next_token_dist`, but drives the real 26B MoE path -/// (`predict_kquant`, `moe_remote = None` → in-process `cpu_moe_forward`). -fn next_token_dist( - weights: &mut ModelWeights, - tok: &tokenizers::Tokenizer, - ids: &[u32], - index: &VectorIndex, -) -> HashMap { - let r = predict_kquant(weights, tok, ids, usize::MAX, index); - r.token_ids - .into_iter() - .zip(r.predictions.into_iter().map(|(_, p)| p)) - .collect() -} - -/// KL(P‖Q) in bits + top-1 agreement. P = dense ground truth, Q = candidate. -fn kl_bits(p: &HashMap, q: &HashMap) -> (f64, bool) { - let eps = 1e-12; - let mut kl = 0.0; - for (&id, &pi) in p { - if pi <= 0.0 { - continue; - } - let qi = q.get(&id).copied().unwrap_or(0.0).max(eps); - kl += pi * (pi.max(eps) / qi).ln(); - } - let p_arg = p.iter().max_by(|a, b| a.1.total_cmp(b.1)).map(|(i, _)| *i); - let q_arg = q.iter().max_by(|a, b| a.1.total_cmp(b.1)).map(|(i, _)| *i); - ( - kl / std::f64::consts::LN_2, - p_arg.is_some() && p_arg == q_arg, - ) -} - -/// Within-expert schedule with exactly one layer pruned to `frac` (everything -/// else dense), with the given selector. Isolates layer `l`'s contribution. -fn one_layer( - num_layers: usize, - l: usize, - frac: f32, - sel: ExpertFeatureSelector, -) -> WithinExpertRouting { - let mut r = WithinExpertRouting::dense(num_layers); - r.frac_per_layer[l] = Some(frac); - r.selector = sel; - r -} - -/// Teacher-forced per-token NLL (bits) over `ids` under the installed routing, -/// plus per-position argmax (for the flip rate). Scores positions `1..len`. -fn token_nlls( - weights: &mut ModelWeights, - tok: &tokenizers::Tokenizer, - ids: &[u32], - index: &VectorIndex, - label: &str, -) -> (Vec, Vec) { - let (mut nlls, mut args) = (Vec::new(), Vec::new()); - for i in 1..ids.len() { - eprint!("\r [{label}] pos {i}/{} ", ids.len()); - let r = predict_kquant(weights, tok, &ids[..i], usize::MAX, index); - let dist: HashMap = r - .token_ids - .iter() - .copied() - .zip(r.predictions.iter().map(|(_, p)| *p)) - .collect(); - let p = dist.get(&ids[i]).copied().unwrap_or(0.0).max(1e-12); - nlls.push(-p.log2()); - args.push(r.token_ids.first().copied().unwrap_or(0)); - } - (nlls, args) -} - -fn pct(sorted: &[f64], q: f64) -> f64 { - if sorted.is_empty() { - return 0.0; - } - let idx = ((sorted.len() as f64 - 1.0) * q).round() as usize; - sorted[idx.min(sorted.len() - 1)] -} - -fn mean(v: &[f64]) -> f64 { - v.iter().sum::() / v.len().max(1) as f64 -} - -fn git_rev() -> String { - std::process::Command::new("git") - .args(["rev-parse", "--short", "HEAD"]) - .output() - .ok() - .filter(|o| o.status.success()) - .map(|o| String::from_utf8_lossy(&o.stdout).trim().to_string()) - .unwrap_or_else(|| "unknown".to_string()) -} - -/// Per-expert intermediate size, derived from a Q4_K gate_up entry -/// (`[2*inter, hidden]`): `inter = gate_up_elems / (2*hidden)`. Robust to -/// arch metadata and matches what the kernel sees. -fn expert_inter(weights: &ModelWeights, layer: usize) -> Option { - use larql_models::quant::ggml::{Q4_K_BLOCK_BYTES, Q4_K_BLOCK_ELEMS}; - let (gu, _dn) = weights.get_layer_entry_bytes(layer, 0)?; - let hidden = weights.hidden_size; - if hidden == 0 { - return None; - } - let gu_elems = (gu.len() / Q4_K_BLOCK_BYTES) * Q4_K_BLOCK_ELEMS; - Some(gu_elems / (2 * hidden)) -} - -fn main() { - let args: Vec = std::env::args().collect(); - let vindex = args - .iter() - .skip(1) - .find(|a| !a.starts_with("--")) - .cloned() - .unwrap_or_else(|| "output/gemma4-26b-a4b-q4k.vindex".to_string()); - let quick = args.iter().any(|a| a == "--quick"); - let smoke = args.iter().any(|a| a == "--smoke"); - let kl_thresh = 0.05_f64; - let dir = std::path::PathBuf::from(&vindex); - let mut cb = larql_vindex::SilentLoadCallbacks; - - eprintln!("Loading {vindex} ..."); - let mut weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index - .load_interleaved_kquant(&dir) - .expect("interleaved kquant (dense FFN half of MoE layers)"); - index.load_attn_kquant(&dir).expect("attn kquant"); - let _ = index.load_lm_head_kquant(&dir); - let tok = load_tokenizer(&dir).expect("tokenizer"); - - let nl = weights.num_layers; - if !weights.has_per_layer_ffn() { - eprintln!( - "ERROR: {vindex} has no per-layer expert weights (has_per_layer_ffn=false). \ - This probe needs a hybrid-MoE vindex (e.g. gemma4-26b-a4b-q4k)." - ); - std::process::exit(2); - } - // Expert layers = those carrying per-layer FFN entries (the others are dense - // and the within-expert knob doesn't touch them). - let expert_layers: Vec = (0..nl) - .filter(|&l| weights.get_layer_entry_bytes(l, 0).is_some()) - .collect(); - let inter = expert_layers - .first() - .and_then(|&l| expert_inter(&weights, l)) - .unwrap_or(0); - eprintln!( - " {nl} layers, {} with experts, expert inter={inter}", - expert_layers.len() - ); - - // Matrix `baseline_fact_prompts` — the canonical V1 screening set. - let prompts = [ - "The capital of France is", - "The largest planet in the solar system is", - "The author of Pride and Prejudice was", - "The chemical symbol for gold is", - ]; - let ids: Vec> = prompts - .iter() - .map(|p| tok.encode(*p, true).expect("encode").get_ids().to_vec()) - .collect(); - - // Dense ground-truth distributions (routing OFF), computed once. - eprintln!("Computing dense reference distributions ..."); - set_routing(None); - let dense_dists: Vec> = ids - .iter() - .map(|id| next_token_dist(&mut weights, &tok, id, &index)) - .collect(); - - // Average KL over prompts for a one-layer schedule at (l, frac, selector). - let avg_kl = |weights: &mut ModelWeights, - dense: &[HashMap], - l: usize, - frac: f32, - sel: ExpertFeatureSelector| - -> (f64, f64) { - set_routing(Some(one_layer(nl, l, frac, sel))); - let (mut kl, mut agree) = (0.0, 0usize); - for (id, d) in ids.iter().zip(dense.iter()) { - let q = next_token_dist(weights, &tok, id, &index); - let (b, a) = kl_bits(d, &q); - kl += b; - agree += a as usize; - } - set_routing(None); - let n = ids.len().max(1) as f64; - (kl / n, 100.0 * agree as f64 / n) - }; - - // ── Step 0: parity anchor (the spine, before any threshold) ──────────── - println!("\n=== Step 0: parity anchor — {vindex} ==="); - let mid_expert = expert_layers[expert_layers.len() / 2]; - { - // (a) all-dense schedule installed → must equal dense (KL≈0): confirms - // the instrument is faithful (frac=None path is identity). - set_routing(Some(WithinExpertRouting::dense(nl))); - let (mut kl, _) = (0.0, 0); - for (id, d) in ids.iter().zip(dense_dists.iter()) { - kl += kl_bits(d, &next_token_dist(&mut weights, &tok, id, &index)).0; - } - set_routing(None); - println!( - " all-dense schedule (instrument off-by-frac): KL={:>8.5} bits (expect ~0)", - kl / ids.len() as f64 - ); - // (b) one expert layer pruned hard (frac=1/8) → the knob must bite. - let (kl_bite, ag_bite) = avg_kl( - &mut weights, - &dense_dists, - mid_expert, - 0.125, - ExpertFeatureSelector::ActMagnitude, - ); - println!( - " L{mid_expert} @ frac=0.125 (~{} of {inter} feats): KL={kl_bite:>8.5} bits agree={ag_bite:.0}% (expect > 0)", - (0.125 * inter as f32).round() as usize - ); - } - - let json_path = args - .iter() - .find_map(|a| a.strip_prefix("--json=").map(|s| s.to_string())) - .unwrap_or_else(|| { - let stem = std::path::Path::new(&vindex) - .file_stem() - .and_then(|s| s.to_str()) - .unwrap_or("model"); - format!("v1moe_{stem}.json") - }); - - let fracs = [ - 1.0 / 64.0, - 1.0 / 32.0, - 1.0 / 16.0, - 1.0 / 8.0, - 1.0 / 4.0, - 1.0 / 2.0, - 1.0, - ]; - - // Per-expert-layer threshold table: (layer, frac, kl). `thr_frac` is the - // full-length (nl) keep schedule used by Phase B compounding (dense on - // non-expert layers). - let mut thr: Vec<(usize, f32, f64)> = Vec::new(); - let mut thr_frac: Vec> = vec![None; nl]; - - if smoke { - eprintln!("[--smoke] skipping Phase A sweep; uniform frac = 1/8 on expert layers"); - for &l in &expert_layers { - thr_frac[l] = Some(0.125); - thr.push((l, 0.125, f64::NAN)); - } - } else { - let sweep_layers: Vec = if quick { - vec![expert_layers[0], mid_expert, *expert_layers.last().unwrap()] - } else { - expert_layers.clone() - }; - println!("\n=== Phase A: per-expert-layer oracle threshold (min keep-frac for KL ≤ {kl_thresh}) ===\n"); - println!( - "{:>5} {:>10} {:>10} {:>9}", - "layer", "thr-frac", "thr-k", "thr-KL" - ); - for &l in &sweep_layers { - let mut chosen: Option<(f32, f64)> = None; - for &f in &fracs { - let (kl, _ag) = avg_kl( - &mut weights, - &dense_dists, - l, - f, - ExpertFeatureSelector::ActMagnitude, - ); - if kl <= kl_thresh { - chosen = Some((f, kl)); - break; - } - } - let (f, kl) = chosen.unwrap_or((1.0, f64::NAN)); - let k = (f * inter as f32).round() as usize; - println!("{l:>5} {f:>10.4} {k:>10} {kl:>9.5}"); - thr.push((l, f, kl)); - thr_frac[l] = Some(f); - } - if quick { - println!("\n [--quick] measured 3 layers only; stopping before Phase B/C."); - return; - } - let mean_frac = - thr.iter().map(|(_, f, _)| *f as f64).sum::() / thr.len().max(1) as f64; - println!("\n mean threshold fraction = {mean_frac:.4} (per-layer SCREEN only; claim gate below)"); - } - - // ── Bandwidth accounting (within-expert; gate+up projection not free) ── - // Per active expert, in units of `inter`-sized rows touched: - // dense : gate(inter) + up(inter) + down(inter) = 3·inter - // oracle : gate(inter) + up(inter) + down(k) = 2·inter + k (ActMag needs full act) - // cheap : gate(k) + up(k) + down(k) = 3·k (content-blind route) - // inter cancels in the ratio, so accumulate over expert layers via frac. - let n_exp = expert_layers.len().max(1) as f64; - let dense_rows = 3.0 * n_exp; - let mut cheap_rows = 0.0; - let mut oracle_rows = 0.0; - for &l in &expert_layers { - let f = thr_frac[l].unwrap_or(1.0) as f64; - cheap_rows += 3.0 * f; - oracle_rows += 2.0 + f; - } - let cheap_frac = cheap_rows / dense_rows; - let oracle_frac = oracle_rows / dense_rows; - let cheap_factor = if cheap_frac > 0.0 { - 1.0 / cheap_frac - } else { - 0.0 - }; - let oracle_factor = if oracle_frac > 0.0 { - 1.0 / oracle_frac - } else { - 0.0 - }; - println!("\n=== Bandwidth (per-expert FFN rows touched per token, vs dense) ==="); - println!(" cheap-route (content-blind): {cheap_frac:.4}× of dense → {cheap_factor:.2}× reduction (best case)"); - println!(" oracle (Phase B cfg): {oracle_frac:.4}× of dense → {oracle_factor:.2}× reduction (gate+up still paid)"); - - // ── Phase B: compounding — all expert layers pruned simultaneously ───── - println!( - "\n=== Phase B: compounding — held-text NLL + drift (all expert layers @ threshold) ===" - ); - let passage = "The expedition had been planned for years, but nothing prepared them for the \ -silence of the ice. Each morning the wind died at dawn, and the only sound was the slow groan of \ -the glacier shifting beneath their tents."; - let pids = tok.encode(passage, true).expect("enc").get_ids().to_vec(); - eprintln!(" held passage: {} tokens", pids.len()); - - let mut comp = WithinExpertRouting::dense(nl); - comp.frac_per_layer = thr_frac.clone(); - comp.selector = ExpertFeatureSelector::ActMagnitude; - - set_routing(None); - let (nll_d, arg_d) = token_nlls(&mut weights, &tok, &pids, &index, "dense"); - set_routing(Some(comp)); - let (nll_c, arg_c) = token_nlls(&mut weights, &tok, &pids, &index, "comp"); - set_routing(None); - eprintln!( - "\r scored {} positions (dense + compounded) ", - nll_d.len() - ); - - let (mut sd, mut sc) = (nll_d.clone(), nll_c.clone()); - sd.sort_by(|a, b| a.total_cmp(b)); - sc.sort_by(|a, b| a.total_cmp(b)); - let (md, mc) = (mean(&nll_d), mean(&nll_c)); - let flips = arg_d.iter().zip(&arg_c).filter(|(a, b)| a != b).count(); - let flip_pct = 100.0 * flips as f64 / arg_d.len().max(1) as f64; - let first_div = arg_d - .iter() - .zip(&arg_c) - .position(|(a, b)| a != b) - .map(|p| p as i64) - .unwrap_or(-1); - let (ppl_d, ppl_c) = (2f64.powf(md), 2f64.powf(mc)); - let ppl_delta_pct = (ppl_c / ppl_d - 1.0) * 100.0; - - println!( - " NLL bits/token dense: mean {md:.3} p90 {:.3} max {:.3}", - pct(&sd, 0.90), - pct(&sd, 1.0) - ); - println!( - " NLL bits/token comp : mean {mc:.3} p90 {:.3} max {:.3} Δmean {:+.3}", - pct(&sc, 0.90), - pct(&sc, 1.0), - mc - md - ); - println!(" perplexity dense {ppl_d:.3} → comp {ppl_c:.3} ({ppl_delta_pct:+.2}%)"); - println!(" argmax drift (comp vs dense): {flip_pct:.1}% first-divergence pos: {first_div}"); - - // ── Phase C: cheap-route realizability (Strided vs ActMagnitude oracle) ─ - println!("\n=== Phase C: cheap-route realizability @ oracle thresholds (frac ≤ 0.25) ==="); - println!( - "{:>5} {:>10} {:>10} {:>10}", - "layer", "frac", "oracle-KL", "strided-KL" - ); - let mut phase_c: Vec<(usize, f32, f64, f64)> = Vec::new(); - for &(l, f, kl_oracle) in thr.iter().filter(|(_, f, _)| *f <= 0.25) { - let (kl_strided, _) = avg_kl( - &mut weights, - &dense_dists, - l, - f, - ExpertFeatureSelector::Strided, - ); - println!("{l:>5} {f:>10.4} {kl_oracle:>10.5} {kl_strided:>10.5}"); - phase_c.push((l, f, kl_oracle, kl_strided)); - } - let cheap_ok = phase_c - .iter() - .filter(|(_, _, _, s)| *s <= kl_thresh) - .count(); - let cheap_realizable_pct = if phase_c.is_empty() { - 0.0 - } else { - 100.0 * cheap_ok as f64 / phase_c.len() as f64 - }; - println!( - "\n content-blind (strided) route clears KL ≤ {kl_thresh} at {cheap_realizable_pct:.0}% of small-threshold layers" - ); - - // ── JSON artifact (bench/aim-validation/matrix.json result contract) ─── - let model = std::path::Path::new(&vindex) - .file_stem() - .and_then(|s| s.to_str()) - .unwrap_or("model") - .to_string(); - let topk_rows: Vec = thr - .iter() - .map(|(l, f, kl)| { - let klv = if kl.is_nan() { - "null".to_string() - } else { - format!("{kl:.5}") - }; - let k = (*f * inter as f32).round() as usize; - format!("{{\"layer\":{l},\"frac\":{f:.5},\"k\":{k},\"kl\":{klv}}}") - }) - .collect(); - let phase_c_rows: Vec = phase_c - .iter() - .map(|(l, f, o, st)| { - format!("{{\"layer\":{l},\"frac\":{f:.5},\"oracle_kl\":{o:.5},\"strided_kl\":{st:.5}}}") - }) - .collect(); - let json = format!( - concat!( - "{{\n", - " \"test_id\": \"V1-moe-within-expert\",\n", - " \"model\": \"{model}\",\n", - " \"prompt_set\": \"baseline_fact_prompts (KL) + held narrative (NLL)\",\n", - " \"git_rev\": \"{rev}\",\n", - " \"expert_inter\": {inter},\n", - " \"n_expert_layers\": {nexp},\n", - " \"metrics\": {{\n", - " \"topk\": [{topk}],\n", - " \"kl_divergence\": {{\"threshold\": {kl_thresh}, \"n_prompts\": {nprompts}}},\n", - " \"perplexity_delta_pct\": {ppl_delta:.4},\n", - " \"nll_bits_dense_mean\": {md:.4},\n", - " \"nll_bits_comp_mean\": {mc:.4},\n", - " \"argmax_drift_pct\": {flip:.4},\n", - " \"first_divergence_pos\": {first_div},\n", - " \"bytes_touched_per_token\": {{\"cheap_frac\": {cheap_frac:.5}, \"cheap_factor\": {cheap_factor:.4}, \"oracle_frac\": {oracle_frac:.5}, \"oracle_factor\": {oracle_factor:.4}}},\n", - " \"cheap_route\": {{\"strided_realizable_pct\": {crp:.2}, \"by_layer\": [{pc}]}}\n", - " }},\n", - " \"notes\": \"within-expert feature routing on MoE experts; ActMagnitude oracle (gate+up still paid); Phase B is the claim gate (held-text NLL + drift); selector ActMagnitude vs Strided for cheap-route realizability\"\n", - "}}\n" - ), - model = model, - rev = git_rev(), - inter = inter, - nexp = expert_layers.len(), - topk = topk_rows.join(","), - kl_thresh = kl_thresh, - nprompts = prompts.len(), - ppl_delta = ppl_delta_pct, - md = md, - mc = mc, - flip = flip_pct, - first_div = first_div, - cheap_frac = cheap_frac, - cheap_factor = cheap_factor, - oracle_frac = oracle_frac, - oracle_factor = oracle_factor, - crp = cheap_realizable_pct, - pc = phase_c_rows.join(","), - ); - std::fs::write(&json_path, &json).expect("write json artifact"); - println!("\n artifact → {json_path}"); -} diff --git a/crates/larql-inference/examples/walk_ffn_v2_fp4_nll.rs b/crates/larql-inference/examples/walk_ffn_v2_fp4_nll.rs deleted file mode 100644 index 540e0a244..000000000 --- a/crates/larql-inference/examples/walk_ffn_v2_fp4_nll.rs +++ /dev/null @@ -1,249 +0,0 @@ -//! V2 predictive-units check — FP4 (E2M1) FFN precision cost (aim-validation, KU3). -//! -//! The static scan (`fp4_q1_scan`) shows ~99.8–99.9% of per-feature blocks fit -//! FP4's R<16 dynamic range across Gemma 3 / Granite (down is the tail). But -//! static range-fit is a SCREENING proxy — the #26 lesson (Q3 looked lossless by -//! single-step KL, then drift overturned it) says the deciding metric is held-text -//! NLL + argmax drift, with the quantisation applied to ALL layers at once. This -//! adjudicates whether FP4's per-block fit actually yields near-lossless OUTPUT. -//! -//! Three arms, teacher-forced on entropic prose, scoring per-token NLL (bits): -//! - **f32** reference (dequantised from the vindex) -//! - **Q4-int** 4-bit symmetric-uniform — the SHIPPED 4-bit baseline (calibrates tolerance) -//! - **FP4-e2m1** the real FP4 block codec roundtrip (`encode_fp4_feature`/`decode_fp4_feature`) -//! -//! All three quantise the SAME f32 weights, so the deltas are pure format error. -//! FP4 ≈ Q4-int ≈ f32 ⇒ FP4 near-lossless at the output; FP4 ≫ f32 ⇒ real cost. -//! -//! Usage: `cargo run --release --example walk_ffn_v2_fp4_nll -- [VINDEX]` - -use larql_inference::ffn::FfnBackend; -use larql_inference::vindex::insert_q4k_layer_tensors_resident; -use larql_inference::{load_tokenizer, predict_with_ffn}; -use larql_models::ModelWeights; -use ndarray::{Array1, Array2}; -use std::collections::HashMap; - -#[derive(Clone, Copy)] -enum Arm { - F32, - IntBits(u32), - Fp4, -} - -/// Symmetric-uniform integer requant (32-elem blocks) — the existing Q-baseline -/// (matches `walk_ffn_nll.rs::requant_row`). -fn requant_row_int(row: &mut [f32], bits: u32) { - const BLK: usize = 32; - for blk in row.chunks_mut(BLK) { - let maxabs = blk.iter().fold(0.0f32, |m, &v| m.max(v.abs())); - if maxabs == 0.0 { - continue; - } - let levels = ((1u32 << (bits - 1)) - 1) as f32; - let scale = maxabs / levels; - for v in blk.iter_mut() { - *v = (*v / scale).round().clamp(-levels, levels) * scale; - } - } -} - -/// Real FP4 (E2M1) roundtrip on a feature row: 256-elem blocks, 8×32 sub-blocks -/// with FP8 sub-scales — the actual on-disk FP4 storage error. -fn requant_row_fp4(row: &mut [f32]) { - let bytes = larql_models::quant::fp4_block::encode_fp4_feature(row); - let mut out = vec![0f32; row.len()]; - larql_models::quant::fp4_block::decode_fp4_feature(&bytes, &mut out); - row.copy_from_slice(&out); -} - -struct GradedFfn { - gate: Vec>, - up: Vec>, - down: Vec>, - label: String, -} -impl FfnBackend for GradedFfn { - fn forward(&self, layer: usize, x: &Array2) -> Array2 { - let (g, u, d) = (&self.gate[layer], &self.up[layer], &self.down[layer]); - let hidden = x.shape()[1]; - let mut out = Array2::::zeros((x.shape()[0], hidden)); - for (s, xr) in x.rows().into_iter().enumerate() { - let xr = xr.to_owned(); - let gs = g.dot(&xr); - let us = u.dot(&xr); - let act: Array1 = gs - .iter() - .zip(us.iter()) - .map(|(&gg, &uu)| larql_inference::ffn::gelu_tanh(gg) * uu) - .collect(); - out.row_mut(s).assign(&act.dot(d)); - } - out - } - // forward_observed: trait default (Absent) — this arm computes no - // observable activation tensor and must not fabricate one. - fn name(&self) -> &str { - &self.label - } -} - -fn build_arm( - weights: &ModelWeights, - index: &larql_vindex::VectorIndex, - arm: Arm, - label: &str, -) -> GradedFfn { - let (nl, hidden) = (weights.num_layers, weights.hidden_size); - let (mut gate, mut up, mut down) = (Vec::new(), Vec::new(), Vec::new()); - for layer in 0..nl { - let inter = index.num_features(layer); - for (comp, store) in [(0usize, &mut gate), (1, &mut up), (2, &mut down)] { - let w = index.kquant_ffn_layer(layer, comp).expect("f32 comp"); - let mut m = Array2::::zeros((inter, hidden)); - for f in 0..inter { - let mut r = w[f * hidden..(f + 1) * hidden].to_vec(); - match arm { - Arm::F32 => {} - Arm::IntBits(b) => requant_row_int(&mut r, b), - Arm::Fp4 => requant_row_fp4(&mut r), - } - m.row_mut(f).assign(&Array1::from(r)); - } - store.push(m); - } - } - GradedFfn { - gate, - up, - down, - label: label.to_string(), - } -} - -/// Teacher-forced per-token NLL (bits) + per-position argmax (for flip rate). -fn token_nlls( - weights: &ModelWeights, - tok: &tokenizers::Tokenizer, - ids: &[u32], - ffn: &dyn FfnBackend, -) -> (Vec, Vec) { - let (mut nlls, mut args) = (Vec::new(), Vec::new()); - for i in 1..ids.len() { - if i % 8 == 0 { - eprint!("\r [{}] pos {i}/{} ", ffn.name(), ids.len()); - } - let r = predict_with_ffn(weights, tok, &ids[..i], usize::MAX, ffn); - let dist: HashMap = r - .token_ids - .iter() - .copied() - .zip(r.predictions.iter().map(|(_, p)| *p)) - .collect(); - let p = dist.get(&ids[i]).copied().unwrap_or(0.0).max(1e-12); - nlls.push(-p.log2()); - args.push(r.token_ids.first().copied().unwrap_or(0)); - } - (nlls, args) -} - -fn pct(sorted: &[f64], q: f64) -> f64 { - if sorted.is_empty() { - return 0.0; - } - let idx = ((sorted.len() as f64 - 1.0) * q).round() as usize; - sorted[idx.min(sorted.len() - 1)] -} -fn mean(v: &[f64]) -> f64 { - v.iter().sum::() / v.len().max(1) as f64 -} - -fn main() { - let args: Vec = std::env::args().collect(); - let vindex = args - .iter() - .skip(1) - .find(|a| !a.starts_with("--")) - .cloned() - .unwrap_or_else(|| "output/gemma3-4b-q4k-v2.vindex".to_string()); - let dir = std::path::PathBuf::from(&vindex); - let mut cb = larql_vindex::SilentLoadCallbacks; - eprintln!("Loading {vindex} ..."); - let mut weights = larql_vindex::load_model_weights_kquant(&dir, &mut cb).expect("weights"); - let mut index = larql_vindex::VectorIndex::load_vindex(&dir, &mut cb).expect("index"); - index.load_interleaved_kquant(&dir).expect("interleaved"); - index.load_attn_kquant(&dir).expect("attn"); - let _ = index.load_lm_head_kquant(&dir); - let tok = load_tokenizer(&dir).expect("tok"); - for layer in 0..weights.num_layers { - insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant attn"); - } - - let passage = "The expedition had been planned for years, but nothing prepared \ -them for the silence of the ice. Each morning the wind died at dawn, and the only \ -sound was the slow groan of the glacier shifting beneath their tents. Provisions \ -were running low, and the captain knew that another week of delay would mean \ -turning back without ever reaching the plateau they had crossed two oceans to find."; - let ids = tok.encode(passage, true).expect("enc").get_ids().to_vec(); - eprintln!( - "Held passage: {} tokens. Arms: f32 / Q4-int / FP4-e2m1", - ids.len() - ); - - let arms: [(&str, Arm); 3] = [ - ("f32", Arm::F32), - ("Q4-int", Arm::IntBits(4)), - ("FP4-e2m1", Arm::Fp4), - ]; - let mut by_arm: Vec<(String, Vec, Vec)> = Vec::new(); - for (name, arm) in arms { - eprintln!(" arm {name}: building + scoring ..."); - let ffn = build_arm(&weights, &index, arm, name); - let (nlls, a) = token_nlls(&weights, &tok, &ids, &ffn); - eprintln!("\r arm {name}: done ({} positions) ", nlls.len()); - by_arm.push((name.to_string(), nlls, a)); - } - - println!( - "\nV2 FP4 predictive cost — {vindex}\nPer-token NLL (bits), teacher-forced, {} positions\n", - by_arm[0].1.len() - ); - println!( - "{:<10} {:>8} {:>8} {:>8} {:>8}", - "arm", "mean", "p90", "p99", "max" - ); - let f32_mean = mean(&by_arm[0].1); - for (name, n, _) in &by_arm { - let mut s = n.clone(); - s.sort_by(|a, b| a.total_cmp(b)); - println!( - "{name:<10} {:>8.3} {:>8.3} {:>8.3} {:>8.3}", - mean(n), - pct(&s, 0.90), - pct(&s, 0.99), - pct(&s, 1.0) - ); - } - - // Δ vs f32 + flip rate (FP4 argmax vs f32 argmax, teacher-forced). - let f32_args = &by_arm[0].2; - println!("\nΔ mean NLL vs f32 + argmax flip rate (vs f32):"); - for (name, n, a) in by_arm.iter().skip(1) { - let dmean = mean(n) - f32_mean; - let flips = a.iter().zip(f32_args).filter(|(x, y)| x != y).count(); - let flip_pct = 100.0 * flips as f64 / a.len().max(1) as f64; - println!(" {name:<10} Δmean {dmean:+.4} bits flip {flip_pct:.1}%"); - } - let q4_mean = mean(&by_arm[1].1); - let fp4_mean = mean(&by_arm[2].1); - println!( - "\n ladder: f32 {f32_mean:.4} → Q4-int {q4_mean:.4} (+{:.4}) → FP4 {fp4_mean:.4} (+{:.4} vs f32)", - q4_mean - f32_mean, - fp4_mean - f32_mean - ); - println!( - " VERDICT: FP4 within {:.4} bits of f32 and {:+.4} vs the shipped Q4-int baseline.", - fp4_mean - f32_mean, - fp4_mean - q4_mean - ); -} diff --git a/crates/larql-inference/src/ffn/local_moe.rs b/crates/larql-inference/src/ffn/local_moe.rs index 6010681c1..47518397d 100644 --- a/crates/larql-inference/src/ffn/local_moe.rs +++ b/crates/larql-inference/src/ffn/local_moe.rs @@ -77,8 +77,10 @@ impl<'a> FfnBackend for LocalMoeFfn<'a> { &self, layer: usize, h_post_attn: &Array2, - ) -> Option> { - Some(moe_ffn_block_cpu_with_index( + ) -> Result>, larql_execution::BoxRefusal> { + // Local dispatch over resident weights: there is no operand this could + // fail to reach, so it never refuses. + Ok(Some(moe_ffn_block_cpu_with_index( self.weights, h_post_attn, layer, @@ -88,7 +90,7 @@ impl<'a> FfnBackend for LocalMoeFfn<'a> { None, None, self.index, - )) + ))) } } @@ -113,6 +115,7 @@ mod tests { let h_post_attn = Array2::::from_elem((2, weights.hidden_size), 0.1); let out = ffn .forward_moe_full_layer(0, &h_post_attn) + .expect("executes") .expect("LocalMoeFfn always returns Some"); assert_eq!(out.shape(), &[2, weights.hidden_size]); assert!(out.iter().all(|v| v.is_finite())); @@ -153,8 +156,14 @@ mod tests { remote: &disconnected, }; let h_post_attn = Array2::::from_elem((2, weights.hidden_size), 0.1); - let out_local = local.forward_moe_full_layer(0, &h_post_attn).unwrap(); - let out_zero_experts = remote.forward_moe_full_layer(0, &h_post_attn).unwrap(); + let out_local = local + .forward_moe_full_layer(0, &h_post_attn) + .unwrap() + .unwrap(); + let out_zero_experts = remote + .forward_moe_full_layer(0, &h_post_attn) + .unwrap() + .unwrap(); assert_eq!(out_local.shape(), out_zero_experts.shape()); let max_abs_diff = out_local .iter() diff --git a/crates/larql-inference/src/ffn/mod.rs b/crates/larql-inference/src/ffn/mod.rs index 4625eb3e8..ec1abe20d 100644 --- a/crates/larql-inference/src/ffn/mod.rs +++ b/crates/larql-inference/src/ffn/mod.rs @@ -15,6 +15,8 @@ pub mod graph_backend; pub mod local_moe; +pub mod moe_backend; +pub mod moe_bound; pub mod moe_remote; pub mod remote; pub mod sparse; @@ -32,8 +34,11 @@ pub use larql_compute::ffn::{ // ── Re-exports ── pub use local_moe::LocalMoeFfn; +pub use moe_backend::{InProcessMoeBackend, MoeBackendError, MoeExpertBackend}; +pub use moe_bound::BoundMoeBackend; pub use moe_remote::{ - MoeRouterWeights, RemoteMoeBackend, RemoteMoeError, RemoteMoeFfn, ShardConfig, + MoeFfn, MoeRouterWeights, RecordedRefusal, RefusalPolicy, RemoteMoeBackend, RemoteMoeError, + RemoteMoeFfn, ShardConfig, }; pub use remote::{ decode_q8k_batch_response_entries, decode_single_response, encode_binary_request, @@ -100,7 +105,7 @@ mod router_tests { let weights = make_test_weights(); let ffn = WeightFfn { weights: &weights }; let h = larql_vindex::ndarray::Array2::::zeros((1, weights.hidden_size)); - assert!(ffn.forward_moe_full_layer(0, &h).is_none()); + assert!(ffn.forward_moe_full_layer(0, &h).is_ok_and(|o| o.is_none())); } #[test] diff --git a/crates/larql-inference/src/ffn/moe_backend.rs b/crates/larql-inference/src/ffn/moe_backend.rs new file mode 100644 index 000000000..1ed555331 --- /dev/null +++ b/crates/larql-inference/src/ffn/moe_backend.rs @@ -0,0 +1,234 @@ +//! Who computes a hybrid-MoE layer's expert contribution. +//! +//! The block loop — attention, KV, the dense slab, the norms, PLE, the layer +//! scalar — is shared by every route. What varies is one step: +//! +//! ```text +//! h_post_attn +//! ├── dense FFN slab shared +//! └── expert contribution this trait +//! ├── in-process cpu_moe_forward over MoeLayerWeights (the default) +//! ├── remote RemoteMoeBackend, experts fetched from shards +//! └── bound a VINDEX3 BoundMoeOperation +//! ``` +//! +//! # Why a trait rather than a second parameter +//! +//! `moe_ffn_block_cpu` previously took `Option<&RemoteMoeBackend>` — a +//! concrete type, so the only way to add a third route was another optional +//! parameter and a rule that at most one may be `Some`. That rule is not +//! expressible in the type, which means it is a rule someone eventually +//! breaks, and the failure mode is a layer whose expert contribution comes +//! from a route nobody chose. +//! +//! One `Option<&dyn MoeExpertBackend>` makes "exactly one route, or the +//! in-process default" the only representable state. +//! +//! # The backend derives its own operands +//! +//! The caller passes `weights` and a layer, not a prepared router. The remote +//! backend needs the router only, because it fetches experts from shards; the +//! bound backend needs the whole layer. Preparing the union at the call site +//! would make the block loop know what each route reads, which is exactly the +//! coupling this trait removes. + +use larql_models::ModelWeights; +use ndarray::Array2; + +use super::moe_remote::RemoteMoeError; + +/// Why a backend could not produce an expert contribution. +/// +/// Two variants rather than a string, so a refusal keeps the structure the +/// route gave it. In particular a bound route's [`ExecutionError`] carries +/// `refusal()`, which distinguishes an operand that lives elsewhere from a +/// binding that is wrong — a distinction that a formatted message destroys. +/// +/// No `PartialEq`: [`RemoteMoeError`] does not claim it, and neither should a +/// type that wraps it. +/// +/// [`ExecutionError`]: larql_vindex::runtime::ExecutionError +#[derive(Debug, thiserror::Error)] +pub enum MoeBackendError { + #[error("remote expert dispatch failed: {0}")] + Remote(#[from] RemoteMoeError), + #[error("bound expert execution failed: {0}")] + Bound(#[from] larql_vindex::runtime::ExecutionError), +} + +/// A route that computes one hybrid-MoE layer's expert contribution. +pub trait MoeExpertBackend { + /// Expert contribution for every position of `h`, shaped `[seq_len, hidden]`. + /// + /// Returns zeros — not an error — when the layer has no expert weights to + /// route into. That is the in-process path's behaviour and a backend that + /// diverged from it would change the model rather than the route. + fn forward_moe_seq( + &self, + weights: &ModelWeights, + layer: usize, + h: &Array2, + norm_offset: f32, + eps: f32, + ) -> Result, MoeBackendError>; + + /// Which route this is, for diagnostics. Never branched on. + fn name(&self) -> &'static str; +} + +/// The in-process route, made explicit as a backend. +/// +/// Byte-identical to the block loop's default branch: the same +/// `cpu_moe_forward` call with the same arguments, position by position. It +/// exists for two reasons. +/// +/// First, so a comparison can put both routes behind one interface and record +/// them the same way — otherwise the incumbent has no seam to observe and the +/// harness ends up reconstructing boundaries it cannot see. +/// +/// Second, so the seam itself is falsifiable. Running with this backend must +/// equal running with no backend at all; if it does not, the trait changed the +/// model rather than merely relocating the call, and every comparison built on +/// it is measuring the wrong thing. +#[derive(Debug, Clone, Copy, Default)] +pub struct InProcessMoeBackend; + +impl MoeExpertBackend for InProcessMoeBackend { + fn forward_moe_seq( + &self, + weights: &ModelWeights, + layer: usize, + h: &Array2, + norm_offset: f32, + eps: f32, + ) -> Result, MoeBackendError> { + let seq_len = h.nrows(); + let hidden = h.ncols(); + let arch = &*weights.arch; + let mut out = Array2::::zeros((seq_len, hidden)); + let Some(moe) = larql_compute::pipeline_layer::build_moe_weights(weights, arch, layer) + else { + return Ok(out); + }; + // The same layer tag the default branch sets, so the within-expert + // probe behaves identically whichever way the call arrives. + larql_compute::cpu::ops::moe::set_current_layer(layer); + for pos in 0..seq_len { + let row: Vec = h.row(pos).to_vec(); + let moe_out = + larql_compute::cpu::ops::moe::cpu_moe_forward(&row, &moe, norm_offset, eps); + for (dst, src) in out.row_mut(pos).iter_mut().zip(moe_out.iter()) { + *dst = *src; + } + } + Ok(out) + } + + fn name(&self) -> &'static str { + "in-process" + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_utils::{make_test_gemma4_moe_weights, GEMMA4_MOE_HIDDEN}; + use larql_compute::cpu::ops::moe::cpu_moe_forward; + + const EPS: f32 = 1e-6; + const NORM_OFFSET: f32 = 0.0; + const MOE_LAYER: usize = 0; + /// Past the fixture's two layers, so `build_moe_weights` finds nothing. + const ABSENT_LAYER: usize = 99; + + fn residual(rows: usize) -> Array2 { + Array2::from_shape_fn((rows, GEMMA4_MOE_HIDDEN), |(r, c)| { + ((r * 7 + c) % 11) as f32 * 0.1 - 0.5 + }) + } + + #[test] + fn the_in_process_backend_is_the_default_branch_relocated() { + // The property the seam rests on. If this diverged, routing the + // incumbent through the trait would be measuring the trait rather than + // relocating a call — and every comparison built on it would be wrong. + let weights = make_test_gemma4_moe_weights(); + let h = residual(3); + let via_trait = InProcessMoeBackend + .forward_moe_seq(&weights, MOE_LAYER, &h, NORM_OFFSET, EPS) + .expect("the fixture executes"); + + let arch = &*weights.arch; + let moe = larql_compute::pipeline_layer::build_moe_weights(&weights, arch, MOE_LAYER) + .expect("the fixture has an MoE layer 0"); + for pos in 0..h.nrows() { + let row: Vec = h.row(pos).to_vec(); + let expected = cpu_moe_forward(&row, &moe, NORM_OFFSET, EPS); + assert_eq!( + via_trait.row(pos).to_vec(), + expected, + "position {pos} diverged" + ); + } + } + + #[test] + fn the_in_process_backend_produces_a_non_zero_contribution() { + // Guards the guard: an all-zero agreement would be two failures + // agreeing rather than a parity result. + let weights = make_test_gemma4_moe_weights(); + let out = InProcessMoeBackend + .forward_moe_seq(&weights, MOE_LAYER, &residual(1), NORM_OFFSET, EPS) + .expect("executes"); + assert!(out.iter().any(|v| v.abs() > f32::EPSILON)); + assert!(out.iter().all(|v| v.is_finite())); + } + + #[test] + fn a_layer_with_no_expert_weights_contributes_zeros_rather_than_erroring() { + // The block loop adds this contribution unconditionally, so a backend + // that errored here would change the model rather than the route. + let weights = make_test_gemma4_moe_weights(); + let h = residual(2); + let out = InProcessMoeBackend + .forward_moe_seq(&weights, ABSENT_LAYER, &h, NORM_OFFSET, EPS) + .expect("an absent MoE layer is not an error"); + assert_eq!(out.shape(), h.shape()); + assert!(out.iter().all(|v| *v == 0.0)); + } + + #[test] + fn every_route_names_itself_distinctly() { + // The name reaches an operator-facing dispatch-error line, so two + // routes sharing one would make the message unactionable. + let names = [ + InProcessMoeBackend.name(), + crate::ffn::BoundMoeBackend::production().name(), + ]; + assert_ne!(names[0], names[1]); + assert!(!names[0].is_empty()); + } + + #[test] + fn a_backend_error_names_the_route_that_refused() { + // The refusal keeps the structure the route gave it — a bound route's + // `ExecutionError` still carries `refusal()` after wrapping. + let inner = larql_vindex::runtime::ExecutionError::SelectedExpertNotResident { + expert: 90, + bank: "layer 5 bank 0".into(), + resident: 8, + population: 128, + }; + let wrapped = MoeBackendError::from(inner.clone()); + let text = wrapped.to_string(); + assert!(text.contains("bound"), "{text}"); + assert!(text.contains("not resident"), "{text}"); + let MoeBackendError::Bound(recovered) = wrapped else { + panic!("wrapping lost the variant"); + }; + assert_eq!( + recovered.refusal(), + larql_vindex::runtime::RefusalKind::Residency + ); + } +} diff --git a/crates/larql-inference/src/ffn/moe_bound.rs b/crates/larql-inference/src/ffn/moe_bound.rs new file mode 100644 index 000000000..c1104eb49 --- /dev/null +++ b/crates/larql-inference/src/ffn/moe_bound.rs @@ -0,0 +1,567 @@ +//! The VINDEX3-bound expert route. +//! +//! Binds each layer's expert population from the mapped Q4_K bytes and +//! executes it through `larql-vindex`'s bound runtime, using the production +//! router and expert kernels. Substituting this for the in-process route is +//! what makes a whole-model comparison a statement about *composition*: every +//! other step — attention, KV, the dense slab, the norms, PLE — is the same +//! code, so a divergence is the MoE path or nothing. +//! +//! # Bound per call, deliberately +//! +//! Binding is not cached across tokens. The regions come from the mapped +//! index, so binding is pointer arithmetic and a shape check rather than a +//! read, and a cache here would be the resolution creep the bound object +//! exists to prevent — with the added hazard that a stale bound plan would +//! outlive the weights it borrows. +//! +//! # What this route does *not* do +//! +//! Fall back. If an operand is missing, a kernel unavailable or a shape wrong, +//! it returns the refusal. The in-process path guards several of these by +//! quietly substituting another kernel; doing that here would make a parity +//! result a statement about whichever route happened to run. + +use larql_models::ModelWeights; +use ndarray::Array2; + +use larql_vindex::format::capability::binding::{ComponentView, RepresentationIdentity}; +use larql_vindex::format::capability::component::ComponentContract; +use larql_vindex::format::capability::coordinate::BankCoordinate; +use larql_vindex::format::lyrw2::region_format::RegionFormat; +use larql_vindex::format::lyrw2::region_role::RegionRole; +use larql_vindex::runtime::consts::{COL_DIM, FUSED_PROJECTION_HALVES}; +use larql_vindex::runtime::{ + execute, BoundBankOperation, BoundExpert, BoundExpertScaling, BoundMoeOperation, + BoundProjection, BoundReduction, BoundRouter, BoundTensor, ExecutionError, ExpertKernel, + MoeInputs, RouterKernel, +}; + +use larql_compute::cpu::ops::moe::{moe_expert_input, moe_post_expert_output, moe_router_input}; +use larql_compute::MoeLayerWeights; + +use super::moe_backend::{MoeBackendError, MoeExpertBackend}; + +/// Which representation these operands claim to come from. +const VARIANT: &str = "vindex-mapped"; +const ROUTER_REGION_SET: &str = "router"; +const PER_EXPERT_SCALE_REGION_SET: &str = "router_per_expert_scale"; +/// Gemma-shaped layers carry one routed bank. Multi-bank layers arrive with +/// the Mini-K3 rung and will index this rather than assume it. +const BANK_ID: u16 = 0; +const ROUTE_NAME: &str = "vindex3-bound"; + +/// Executes hybrid-MoE layers through a VINDEX3 bound route. +#[derive(Debug, Clone, Copy, Default)] +pub struct BoundMoeBackend { + /// Which kernel runs the experts. Defaults to the production + /// Q4_K × Q8_K kernel; the reference is available for differential runs. + pub expert_kernel: ExpertKernel, + /// Which kernel scores the router. + pub router_kernel: RouterKernel, +} + +impl BoundMoeBackend { + /// The production pairing: both kernels bound to `larql-compute`'s own. + pub fn production() -> Self { + Self { + expert_kernel: ExpertKernel::IncumbentQ4kQ8k, + router_kernel: RouterKernel::Incumbent, + } + } + + /// Both kernels on the reference implementations — the oracle pairing. + /// Requires directly-decodable operands, so it refuses Q4_K stores. + pub fn reference() -> Self { + Self { + expert_kernel: ExpertKernel::Reference, + router_kernel: RouterKernel::Reference, + } + } +} + +impl MoeExpertBackend for BoundMoeBackend { + fn forward_moe_seq( + &self, + weights: &ModelWeights, + layer: usize, + h: &Array2, + norm_offset: f32, + eps: f32, + ) -> Result, MoeBackendError> { + let seq_len = h.nrows(); + let hidden = h.ncols(); + let arch = &*weights.arch; + let Some(moe) = larql_compute::pipeline_layer::build_moe_weights(weights, arch, layer) + else { + // No expert weights to route into. Zeros, matching the in-process + // path — a backend that errored here would change the model. + return Ok(Array2::zeros((seq_len, hidden))); + }; + Ok(self.run_layer(layer, &moe, h, norm_offset, eps)?) + } + + fn name(&self) -> &'static str { + ROUTE_NAME + } +} + +impl BoundMoeBackend { + /// Bind one layer and run every position through it. + /// + /// Split out of [`MoeExpertBackend::forward_moe_seq`] so the execution is + /// reachable from a `MoeLayerWeights` alone. The wrapper's only remaining + /// job is to get one out of a `ModelWeights`, which needs a whole loaded + /// model and is therefore the part a unit test cannot reach — so it is kept + /// to the smallest possible body. + pub fn run_layer( + &self, + layer: usize, + moe: &MoeLayerWeights<'_>, + h: &Array2, + norm_offset: f32, + eps: f32, + ) -> Result, ExecutionError> { + let seq_len = h.nrows(); + let hidden = h.ncols(); + // Owned so the bound operands have something with a stable address to + // borrow; the router is f32 in the index and the expert regions are + // borrowed from the map directly. + let router_bytes = f32_bytes(moe.router_proj); + let scale_bytes = f32_bytes(moe.router_per_expert_scale); + let operation = self.bind(layer, hidden, moe, &router_bytes, &scale_bytes)?; + + let mut out = Array2::::zeros((seq_len, hidden)); + for pos in 0..seq_len { + let residual: Vec = h.row(pos).to_vec(); + // The two inputs are produced by the surrounding block's norms and + // scales, exactly as the in-process path produces them. Deriving + // them here would re-model the incumbent's routing policy with + // somewhere for the two to disagree. + let expert_input = moe_expert_input(&residual, moe, norm_offset, eps); + let router_in = moe_router_input(&residual, &expert_input, moe, norm_offset, eps); + + let delta = execute(&operation, MoeInputs::split(&expert_input, &router_in))?; + // The post-expert norm belongs to the block, not to the operation — + // which is why `execute` returns a delta. + let contribution = moe_post_expert_output(&delta, moe, norm_offset, eps); + for (dst, src) in out.row_mut(pos).iter_mut().zip(contribution.iter()) { + *dst = *src; + } + } + Ok(out) + } + + /// Bind one layer's whole expert population from the mapped bytes. + fn bind<'a>( + &self, + layer: usize, + hidden: usize, + moe: &MoeLayerWeights<'a>, + router_bytes: &'a [u8], + scale_bytes: &'a [u8], + ) -> Result, ExecutionError> { + let inter = moe.intermediate_size; + let inter_padded = moe.inter_padded(); + let format = expert_format(moe); + + let mut experts = Vec::with_capacity(moe.num_experts); + for (expert_id, (&gate_up, &down)) in moe + .experts_gate_up + .iter() + .zip(moe.experts_down.iter()) + .enumerate() + { + experts.push(BoundExpert { + expert_id: expert_id as u32, + projection: BoundProjection::Fused { + gate_up: BoundTensor::direct( + identity(&RegionRole::GateUpFused.name()), + gate_up, + format, + ComponentContract::matrix( + (FUSED_PROJECTION_HALVES * inter) as u32, + hidden as u32, + ), + )?, + }, + // Stored at the padded intermediate width; the role sees the + // live columns. One binding, and each kernel takes from it what + // it needs. + down: BoundTensor::new( + identity(&RegionRole::Down.name()), + down, + format, + ComponentContract::matrix(hidden as u32, inter_padded as u32), + ComponentView::Slice { + dim: COL_DIM, + start: 0, + len: inter as u32, + }, + )?, + }); + } + + let operation = BoundMoeOperation { + router: BoundRouter { + weight: BoundTensor::direct( + identity(ROUTER_REGION_SET), + router_bytes, + RegionFormat::F32, + ComponentContract::matrix(moe.num_experts as u32, hidden as u32), + )?, + top_k: moe.top_k, + selected_weight: moe.routing_policy.selected_weight, + scaling: if scale_bytes.is_empty() { + BoundExpertScaling::None + } else { + BoundExpertScaling::PerExpert { + scales: BoundTensor::direct( + identity(PER_EXPERT_SCALE_REGION_SET), + scale_bytes, + RegionFormat::F32, + ComponentContract::vector(moe.num_experts as u32), + )?, + } + }, + kernel: self.router_kernel, + }, + transforms: Vec::new(), + banks: vec![BoundBankOperation { + bank: BankCoordinate::new(layer as u32, BANK_ID), + experts, + intermediate_dim: inter, + hidden_dim: hidden, + activation: moe.activation, + kernel: self.expert_kernel, + }], + reduction: BoundReduction::WeightedSum, + residual_dim: hidden, + }; + operation.validate()?; + Ok(operation) + } +} + +fn identity(region_set: &str) -> RepresentationIdentity { + RepresentationIdentity::new(region_set, VARIANT) +} + +fn f32_bytes(values: &[f32]) -> Vec { + values.iter().flat_map(|v| v.to_le_bytes()).collect() +} + +/// The region encoding this layer's experts are stored in. +/// +/// Mapped from the loader's own `QuantFormat` rather than assumed, so a store +/// that is not Q4_K refuses at bind — with a message naming the encoding — +/// instead of being read as Q4_K super-blocks and producing noise. +fn expert_format(moe: &MoeLayerWeights<'_>) -> RegionFormat { + match moe.expert_data_format { + larql_compute::QuantFormat::Q4_K => RegionFormat::Q4K, + larql_compute::QuantFormat::Q6_K => RegionFormat::Q6K, + larql_compute::QuantFormat::Q4_0 => RegionFormat::Q4_0, + larql_compute::QuantFormat::BF16 => RegionFormat::BF16, + larql_compute::QuantFormat::F16 => RegionFormat::F16, + larql_compute::QuantFormat::F32 => RegionFormat::F32, + // Q4_KF, Q8_0 and I2S have no region-format equivalent this runtime + // binds. Naming the tag keeps the refusal legible rather than + // defaulting to Q4_K and reading the bytes at the wrong stride. + other => RegionFormat::Unknown(unmapped_tag(other)), + } +} + +/// A distinct tag for an encoding with no `RegionFormat`, so the refusal says +/// which one rather than `format_0`. +fn unmapped_tag(format: larql_compute::QuantFormat) -> u16 { + // Offset past every registered `RegionFormat` tag, so an unmapped codec can + // never collide with a real one. + const UNMAPPED_BASE: u16 = 1_000; + UNMAPPED_BASE + + match format { + larql_compute::QuantFormat::Q4_KF => 0, + larql_compute::QuantFormat::Q8_0 => 1, + larql_compute::QuantFormat::I2S => 2, + _ => 3, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use larql_compute::cpu::ops::moe::cpu_moe_forward; + use larql_compute::cpu::ops::q4_common::quantize_q4_k; + use larql_compute::{Activation, QuantFormat}; + + /// One super-block wide, so Q8_K is satisfied and the fixture stays small. + const HIDDEN: usize = 256; + /// Half a block, so `down` is stored padded — the shape a binding gets wrong. + const INTER: usize = 128; + const INTER_PADDED: usize = 256; + const POPULATION: usize = 2; + const TOP_K: usize = 1; + const LAYER: usize = 0; + const EPS: f32 = 1e-6; + const NORM_OFFSET: f32 = 0.0; + + fn ramp(len: usize, seed: usize) -> Vec { + (0..len) + .map(|i| ((i + seed) % 19) as f32 * 0.01 - 0.09) + .collect() + } + + struct Fixture { + gate_up: Vec, + down: Vec, + router: Vec, + input: Vec, + } + + impl Fixture { + fn new() -> Self { + let input: Vec = (0..HIDDEN).map(|i| (i % 11) as f32 * 0.1 - 0.5).collect(); + // Expert 0's router row is the input, so its logit is ‖x‖² and the + // selection is decided by the scores rather than by a tie. + let mut router = vec![0.0f32; POPULATION * HIDDEN]; + router[..HIDDEN].copy_from_slice(&input); + Self { + gate_up: quantize_q4_k(&ramp(2 * INTER * HIDDEN, 1)), + down: quantize_q4_k(&ramp(HIDDEN * INTER_PADDED, 7)), + router, + input, + } + } + + fn moe(&self) -> MoeLayerWeights<'_> { + MoeLayerWeights { + experts_gate_up: vec![&self.gate_up, &self.gate_up], + experts_down: vec![&self.down, &self.down], + routing_policy: larql_compute::MoeRoutingPolicy::default(), + weight_layout: larql_compute::MoeWeightLayout::default(), + expert_data_format: QuantFormat::Q4_K, + router_proj: &self.router, + router_scale: &[], + router_per_expert_scale: &[], + router_norm: &[], + router_norm_parameter_free: false, + router_input_scalar: 1.0, + pre_experts_norm: &[], + post_ffn1_norm: &[], + post_experts_norm: &[], + num_experts: POPULATION, + top_k: TOP_K, + intermediate_size: INTER, + activation: Activation::GeluTanh, + } + } + + fn h(&self) -> Array2 { + Array2::from_shape_vec((1, HIDDEN), self.input.clone()).expect("one row") + } + } + + #[test] + fn the_bound_route_reproduces_the_in_process_route_bit_for_bit() { + // The claim the whole seam rests on, at unit scale: substituting the + // bound route must not change the layer's contribution. + let fixture = Fixture::new(); + let moe = fixture.moe(); + let h = fixture.h(); + + let bound = BoundMoeBackend::production() + .run_layer(LAYER, &moe, &h, NORM_OFFSET, EPS) + .expect("the fixture binds and executes"); + let incumbent = cpu_moe_forward(&fixture.input, &moe, NORM_OFFSET, EPS); + + assert_eq!(bound.nrows(), 1); + assert_eq!(bound.row(0).to_vec(), incumbent); + } + + #[test] + fn the_bound_route_is_not_producing_zeros() { + // Guards the guard: an all-zero agreement would be two failures + // agreeing rather than a parity result. + let fixture = Fixture::new(); + let out = BoundMoeBackend::production() + .run_layer(LAYER, &fixture.moe(), &fixture.h(), NORM_OFFSET, EPS) + .expect("executes"); + assert!(out.iter().any(|v| v.abs() > f32::EPSILON)); + assert!(out.iter().all(|v| v.is_finite())); + } + + #[test] + fn every_position_is_routed_independently() { + // A backend that routed once and broadcast would pass a single-row + // test. Two different rows must produce two different contributions. + let fixture = Fixture::new(); + let mut rows = fixture.input.clone(); + rows.extend(fixture.input.iter().map(|v| -v)); + let h = Array2::from_shape_vec((2, HIDDEN), rows).expect("two rows"); + let out = BoundMoeBackend::production() + .run_layer(LAYER, &fixture.moe(), &h, NORM_OFFSET, EPS) + .expect("executes"); + assert_ne!(out.row(0).to_vec(), out.row(1).to_vec()); + } + + #[test] + fn the_reference_pairing_refuses_quantised_operands() { + // The reference decoder implements the directly-readable encodings + // only, so pairing it with a Q4_K store must refuse rather than + // silently fall back to the kernel that can read them. + let fixture = Fixture::new(); + let err = BoundMoeBackend::reference() + .run_layer(LAYER, &fixture.moe(), &fixture.h(), NORM_OFFSET, EPS) + .unwrap_err(); + assert!( + matches!(err, ExecutionError::KernelOperandUnsuitable { .. }) + || matches!(err, ExecutionError::UnsupportedFormat { .. }), + "{err}" + ); + } + + #[test] + fn a_layer_stored_in_an_unmappable_encoding_refuses_by_name() { + // Q8_0 has no `RegionFormat`, so binding must refuse with a tag that + // says which codec — not default to Q4_K and read at the wrong stride. + let fixture = Fixture::new(); + let moe = MoeLayerWeights { + expert_data_format: QuantFormat::Q8_0, + ..fixture.moe() + }; + let err = BoundMoeBackend::production() + .run_layer(LAYER, &moe, &fixture.h(), NORM_OFFSET, EPS) + .unwrap_err(); + assert!( + matches!(err, ExecutionError::UnsupportedFormat { .. }), + "{err}" + ); + } + + #[test] + fn each_registered_encoding_maps_to_its_region_format() { + let fixture = Fixture::new(); + for (quant, expected) in [ + (QuantFormat::Q4_K, RegionFormat::Q4K), + (QuantFormat::Q6_K, RegionFormat::Q6K), + (QuantFormat::Q4_0, RegionFormat::Q4_0), + (QuantFormat::BF16, RegionFormat::BF16), + (QuantFormat::F16, RegionFormat::F16), + (QuantFormat::F32, RegionFormat::F32), + ] { + let moe = MoeLayerWeights { + expert_data_format: quant, + ..fixture.moe() + }; + assert_eq!(expert_format(&moe), expected, "{quant:?}"); + } + } + + #[test] + fn unmappable_encodings_get_distinct_tags_that_cannot_collide() { + // A shared tag would make two different refusals read identically, and + // a tag inside the registered range would name someone else's codec. + let tags: Vec = [QuantFormat::Q4_KF, QuantFormat::Q8_0, QuantFormat::I2S] + .into_iter() + .map(unmapped_tag) + .collect(); + let mut sorted = tags.clone(); + sorted.sort_unstable(); + sorted.dedup(); + assert_eq!(sorted.len(), tags.len(), "two codecs share a tag: {tags:?}"); + for tag in tags { + assert!( + matches!(RegionFormat::from_u16(tag), RegionFormat::Unknown(_)), + "tag {tag} collides with a registered codec" + ); + } + } + + #[test] + fn the_route_names_itself_and_defaults_to_the_reference_kernels() { + assert_eq!(BoundMoeBackend::production().name(), ROUTE_NAME); + let default = BoundMoeBackend::default(); + assert_eq!(default.expert_kernel, ExpertKernel::Reference); + assert_eq!(default.router_kernel, RouterKernel::Reference); + assert_eq!( + BoundMoeBackend::production().expert_kernel, + ExpertKernel::IncumbentQ4kQ8k + ); + } + + // ── The wrapper, over a loaded model ─────────────────────────────────── + // + // `run_layer` above is reachable from a `MoeLayerWeights` alone. These + // cover the part that is not: resolving one out of a `ModelWeights`, which + // needs the synthetic Gemma fixture. Its experts are BF16, which is the + // useful case here — the reference kernel decodes them and the Q4_K kernel + // must refuse them. + + use crate::test_utils::{make_test_gemma4_moe_weights, GEMMA4_MOE_HIDDEN}; + use larql_vindex::runtime::RefusalKind; + + const MOE_LAYER: usize = 0; + /// Past the fixture's two layers, so `build_moe_weights` finds nothing. + const ABSENT_LAYER: usize = 99; + + fn model_residual(rows: usize) -> Array2 { + Array2::from_shape_fn((rows, GEMMA4_MOE_HIDDEN), |(r, c)| { + ((r * 7 + c) % 11) as f32 * 0.1 - 0.5 + }) + } + + #[test] + fn the_bound_route_over_a_loaded_model_matches_the_in_process_route() { + // End to end through the wrapper: resolve the layer, bind its whole + // population, execute every position. The reference pairing, because + // the fixture stores BF16 experts. + let weights = make_test_gemma4_moe_weights(); + let h = model_residual(3); + let bound = BoundMoeBackend::reference() + .forward_moe_seq(&weights, MOE_LAYER, &h, NORM_OFFSET, EPS) + .expect("the fixture binds and executes"); + let incumbent = super::super::moe_backend::InProcessMoeBackend + .forward_moe_seq(&weights, MOE_LAYER, &h, NORM_OFFSET, EPS) + .expect("executes"); + + assert_eq!(bound.shape(), incumbent.shape()); + let worst = bound + .iter() + .zip(incumbent.iter()) + .map(|(a, b)| (a - b).abs()) + .fold(0.0f32, f32::max); + let scale = incumbent.iter().fold(0.0f32, |m, v| m.max(v.abs())); + // The reference sums in index order where the incumbent uses BLAS, so + // this is a summation-order band rather than a claim of bit-equality. + assert!( + scale > 0.0 && worst / scale <= 1e-3, + "diverged by {worst} of {scale}" + ); + } + + #[test] + fn the_production_pairing_refuses_a_bf16_store_by_format() { + // It reads Q4_K super-blocks. Reading BF16 bytes as super-blocks is + // exactly the silent substitution the format check exists to refuse. + let weights = make_test_gemma4_moe_weights(); + let err = BoundMoeBackend::production() + .forward_moe_seq(&weights, MOE_LAYER, &model_residual(1), NORM_OFFSET, EPS) + .unwrap_err(); + let MoeBackendError::Bound(inner) = err else { + panic!("a bound route must refuse as a bound route"); + }; + assert_eq!(inner.refusal(), RefusalKind::Unsupported, "{inner}"); + } + + #[test] + fn a_layer_with_no_expert_weights_contributes_zeros() { + // Matches the in-process path. A backend that errored here would + // change the model rather than the route. + let weights = make_test_gemma4_moe_weights(); + let h = model_residual(2); + let out = BoundMoeBackend::production() + .forward_moe_seq(&weights, ABSENT_LAYER, &h, NORM_OFFSET, EPS) + .expect("an absent MoE layer is not an error"); + assert_eq!(out.shape(), h.shape()); + assert!(out.iter().all(|v| *v == 0.0)); + } +} diff --git a/crates/larql-inference/src/ffn/moe_remote/backend.rs b/crates/larql-inference/src/ffn/moe_remote/backend.rs index 25a64327f..37d0d6d20 100644 --- a/crates/larql-inference/src/ffn/moe_remote/backend.rs +++ b/crates/larql-inference/src/ffn/moe_remote/backend.rs @@ -866,3 +866,38 @@ impl RemoteMoeBackend { Ok(h2_per_layer) } } + +/// The remote route as a [`MoeExpertBackend`]. +/// +/// Builds its own router from `weights`, which is the one line that moved in +/// from `moe_ffn_block_cpu` when the seam became a trait. A layer with no +/// router returns zeros, matching what the block loop did when +/// `build_moe_router_weights` returned `None` — the behaviour is unchanged, +/// it simply now lives with the route that needs it. +impl crate::ffn::moe_backend::MoeExpertBackend for RemoteMoeBackend { + fn forward_moe_seq( + &self, + weights: &larql_models::ModelWeights, + layer: usize, + h: &ndarray::Array2, + norm_offset: f32, + eps: f32, + ) -> Result, crate::ffn::moe_backend::MoeBackendError> { + let arch = &*weights.arch; + let Some(router) = crate::vindex::build_moe_router_weights(weights, arch, layer) else { + return Ok(ndarray::Array2::zeros((h.nrows(), h.ncols()))); + }; + Ok(RemoteMoeBackend::forward_moe_seq( + self, + layer, + h, + &router, + norm_offset, + eps, + )?) + } + + fn name(&self) -> &'static str { + "remote" + } +} diff --git a/crates/larql-inference/src/ffn/moe_remote/ffn_adapter.rs b/crates/larql-inference/src/ffn/moe_remote/ffn_adapter.rs index e3b6bd852..c7e4de6a8 100644 --- a/crates/larql-inference/src/ffn/moe_remote/ffn_adapter.rs +++ b/crates/larql-inference/src/ffn/moe_remote/ffn_adapter.rs @@ -40,6 +40,216 @@ pub struct RemoteMoeFfn<'a> { } impl<'a> FfnBackend for RemoteMoeFfn<'a> { + fn forward(&self, layer: usize, x: &Array2) -> Array2 { + self.general().forward(layer, x) + } + + fn forward_observed( + &self, + layer: usize, + x: &Array2, + ) -> (Array2, crate::ffn::FfnActivations) { + self.general().forward_observed(layer, x) + } + + fn name(&self) -> &str { + REMOTE_MOE_FFN_NAME + } + + fn forward_moe_full_layer( + &self, + layer: usize, + h_post_attn: &Array2, + ) -> Result>, larql_execution::BoxRefusal> { + self.general().forward_moe_full_layer(layer, h_post_attn) + } +} + +impl<'a> RemoteMoeFfn<'a> { + /// This adapter as the general one. Kept as a delegation rather than a type + /// alias so the `{ weights, remote }` literal its callers use still + /// type-checks — the name is load-bearing at one CLI call site and in two + /// roadmaps. + fn general(&self) -> MoeFfn<'a> { + // Best-effort, because that is what this adapter has always been and + // its callers depend on: a shard that cannot be reached degrades rather + // than stops. + MoeFfn::best_effort(self.weights, self.remote) + } +} + +/// The name `RemoteMoeFfn` reports. Unchanged from before the generalisation: +/// it appears in engine diagnostics and is matched on in at least one test. +const REMOTE_MOE_FFN_NAME: &str = "remote-moe"; + +/// `FfnBackend` for CPU MoE decode through a `KvEngine`, over **any** expert +/// route. +/// +/// The engine owns attention and its KV cache and calls +/// [`FfnBackend::forward_moe_full_layer`] per MoE layer; this computes that +/// layer's block — dense `h1` locally, experts `h2` through whichever +/// [`MoeExpertBackend`] is installed. +/// +/// The remote-specific adapter above predates this and now delegates to it. The +/// generalisation is what lets a VINDEX3 bound route reach the decode path at +/// all: the block loop's seam became a trait in the composition rung, but the +/// *engine's* adapter still named one concrete backend, so the bound route was +/// reachable only from the full-recompute path. +/// +/// PLE is **not** applied here (`moe_ffn_block_cpu` is called with +/// `ple_input = None`), so Per-Layer-Embedding architectures must go through +/// the full-recompute path instead. +pub struct MoeFfn<'a> { + pub weights: &'a ModelWeights, + pub moe: &'a dyn crate::ffn::MoeExpertBackend, + /// Whether a refusing route is a diagnosis or a fallback. + policy: RefusalPolicy, + /// The first refusal seen, with the layer that produced it. + /// + /// A cell rather than a return value because `FfnBackend` has no error + /// channel — `forward_moe_full_layer` returns `Option>`, and + /// widening it would reach every engine. The refusal is therefore caught on + /// the way *in*, before `moe_ffn_block_cpu` handles it. + refusal: std::cell::RefCell>, +} + +/// What a refusing expert route means to the caller. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum RefusalPolicy { + /// Log it, contribute zeros, return the dense half. The historical + /// behaviour, and the right one for best-effort remote inference where a + /// degraded continuation beats no continuation. + #[default] + BestEffort, + /// Record it, so the caller can refuse the output. + /// + /// A selected-but-absent expert must not silently become an FFN-skipping + /// approximation. Under this policy the block still returns its dense half — + /// the trait cannot say otherwise — but [`MoeFfn::refusal`] is set, and a + /// caller that accepts logits without checking it is accepting a + /// numerically wrong continuation that looks plausible. + Strict, +} + +/// A refusal, with where it happened. +/// +/// An error in its own right, so it can cross `FfnBackend`'s boundary as a +/// `BoxRefusal` while keeping both levels: the response category an engine +/// switches on, and the concrete message the route produced. +#[derive(Debug, Clone)] +pub struct RecordedRefusal { + pub layer: usize, + /// The refusal's own classification, preserved rather than flattened — + /// `Residency` means fetch the operand, and is not a defect. + pub kind: larql_execution::RefusalKind, + pub message: String, +} + +impl std::fmt::Display for RecordedRefusal { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "layer {}: {} ({})", self.layer, self.message, self.kind) + } +} + +impl std::error::Error for RecordedRefusal {} + +impl larql_execution::ExecutionRefusal for RecordedRefusal { + fn kind(&self) -> larql_execution::RefusalKind { + self.kind + } +} + +/// Wraps the route so a refusal is recorded before the block loop swallows it. +struct RefusalRecorder<'a> { + inner: &'a dyn crate::ffn::MoeExpertBackend, + sink: &'a std::cell::RefCell>, +} + +impl crate::ffn::MoeExpertBackend for RefusalRecorder<'_> { + fn forward_moe_seq( + &self, + weights: &ModelWeights, + layer: usize, + h: &Array2, + norm_offset: f32, + eps: f32, + ) -> Result, crate::ffn::MoeBackendError> { + let out = self + .inner + .forward_moe_seq(weights, layer, h, norm_offset, eps); + if let Err(err) = &out { + let kind = match err { + crate::ffn::MoeBackendError::Bound(inner) => inner.refusal(), + // A remote dispatch failure is the operand not being here. + crate::ffn::MoeBackendError::Remote(_) => larql_execution::RefusalKind::Residency, + }; + // First only: the earliest refusal is the diagnosis, and later ones + // are usually the same cause repeating per layer. + let mut sink = self.sink.borrow_mut(); + if sink.is_none() { + *sink = Some(RecordedRefusal { + layer, + kind, + message: err.to_string(), + }); + } + } + out + } + + fn name(&self) -> &'static str { + self.inner.name() + } +} + +impl<'a> MoeFfn<'a> { + /// Best-effort: a refusing route contributes zeros and the dense half is + /// returned. Historical behaviour, unchanged. + pub fn best_effort( + weights: &'a ModelWeights, + moe: &'a dyn crate::ffn::MoeExpertBackend, + ) -> Self { + Self { + weights, + moe, + policy: RefusalPolicy::BestEffort, + refusal: std::cell::RefCell::new(None), + } + } + + /// Strict: a refusal is recorded and the caller must check [`Self::refusal`] + /// before accepting the output. + /// + /// Strictness is in the constructor rather than a field a caller can forget + /// to set, and the check is a method rather than a log line someone has to + /// read. + pub fn strict(weights: &'a ModelWeights, moe: &'a dyn crate::ffn::MoeExpertBackend) -> Self { + Self { + weights, + moe, + policy: RefusalPolicy::Strict, + refusal: std::cell::RefCell::new(None), + } + } + + pub fn policy(&self) -> RefusalPolicy { + self.policy + } + + /// The first refusal seen, if any. Always recorded; only *meaningful* as a + /// gate under [`RefusalPolicy::Strict`], where the caller is contracted to + /// consult it. + pub fn refusal(&self) -> Option { + self.refusal.borrow().clone() + } + + /// Whether every layer so far executed its expert route. + pub fn all_experts_executed(&self) -> bool { + self.refusal.borrow().is_none() + } +} + +impl FfnBackend for MoeFfn<'_> { fn forward(&self, layer: usize, x: &Array2) -> Array2 { WeightFfn { weights: self.weights, @@ -59,15 +269,19 @@ impl<'a> FfnBackend for RemoteMoeFfn<'a> { } fn name(&self) -> &str { - "remote-moe" + self.moe.name() } fn forward_moe_full_layer( &self, layer: usize, h_post_attn: &Array2, - ) -> Option> { - Some(moe_ffn_block_cpu( + ) -> Result>, larql_execution::BoxRefusal> { + let recorder = RefusalRecorder { + inner: self.moe, + sink: &self.refusal, + }; + let out = moe_ffn_block_cpu( self.weights, h_post_attn, layer, @@ -75,8 +289,26 @@ impl<'a> FfnBackend for RemoteMoeFfn<'a> { weights: self.weights, }, None, - Some(self.remote), - )) + Some(&recorder), + ); + // `moe_ffn_block_cpu` has already logged the refusal and left the + // expert contribution at zero, so `out` is the dense half wearing the + // shape of an answer. Under `Strict` it must not escape — which is the + // difference between this and the audited contract it replaces, where + // the caller had to remember to ask. + // + // The dense half is computed and discarded on that path. It is an error + // path, and paying for it buys the guarantee that no caller can consume + // a partial layer. + match self.refusal.borrow().as_ref() { + Some(recorded) if self.policy == RefusalPolicy::Strict => { + Err(Box::new(recorded.clone())) + } + // Best-effort keeps its historical behaviour: degrade, having said + // so. Callers of the remote adapter depend on a shard outage + // degrading rather than stopping. + _ => Ok(Some(out)), + } } } @@ -100,6 +332,7 @@ mod tests { let h_post_attn = Array2::::from_elem((2, weights.hidden_size), 0.1); let out = ffn .forward_moe_full_layer(0, &h_post_attn) + .expect("executes") .expect("RemoteMoeFfn always returns Some"); assert_eq!(out.shape(), &[2, weights.hidden_size]); assert!(out.iter().all(|v| v.is_finite())); @@ -125,4 +358,171 @@ mod tests { let act = obs.into_dense().expect("dense fallback observes densely"); assert_eq!(act.shape()[0], 2); } + + /// The generalisation must not have changed the remote adapter. Same + /// weights, same input, same output — the delegation relocates the call and + /// nothing else, which is the same property the block-loop seam had to + /// prove in the composition rung. + #[test] + fn the_remote_adapter_still_equals_the_general_one() { + let weights = make_test_gemma4_moe_weights(); + let remote = RemoteMoeBackend::new_disconnected(); + let specific = RemoteMoeFfn { + weights: &weights, + remote: &remote, + }; + let general = MoeFfn::best_effort(&weights, &remote); + let h = Array2::::from_elem((2, weights.hidden_size), 0.1); + assert_eq!( + specific + .forward_moe_full_layer(0, &h) + .expect("executes") + .expect("produces a layer"), + general + .forward_moe_full_layer(0, &h) + .expect("executes") + .expect("produces a layer") + ); + assert_eq!(specific.forward(0, &h), general.forward(0, &h)); + // The name is the one thing that deliberately differs: the remote + // adapter keeps its historical name for diagnostics, the general one + // reports whichever route it carries. + assert_eq!(specific.name(), REMOTE_MOE_FFN_NAME); + assert_eq!(general.name(), crate::ffn::MoeExpertBackend::name(&remote)); + } + + /// The point of the generalisation: a VINDEX3 bound route can now drive the + /// engine's FFN seam, which the remote-typed adapter made impossible. + #[test] + fn a_bound_route_can_drive_the_engine_adapter() { + let weights = make_test_gemma4_moe_weights(); + // The fixture stores BF16 experts, so the reference pairing is the one + // that executes; the production Q4_K pairing would refuse them. + let bound = crate::ffn::BoundMoeBackend::reference(); + let ffn = MoeFfn::best_effort(&weights, &bound); + let h = Array2::::from_elem((2, weights.hidden_size), 0.1); + let out = ffn + .forward_moe_full_layer(0, &h) + .expect("executes") + .expect("the adapter always returns Some"); + assert_eq!(out.shape(), &[2, weights.hidden_size]); + assert!(out.iter().all(|v| v.is_finite())); + assert!(out.iter().any(|v| v.abs() > f32::EPSILON)); + } + + /// A refused route must not silently contribute zeros through the adapter. + /// `moe_ffn_block_cpu` logs and leaves `h2` zero on error, so this pins that + /// the dense half still flows — the failure is visible as a *different* + /// output rather than as a plausible one. + #[test] + fn a_refusing_route_still_returns_the_dense_half() { + let weights = make_test_gemma4_moe_weights(); + // Q4_K kernels against a BF16 store: refused at bind. + let refusing = crate::ffn::BoundMoeBackend::production(); + let executing = crate::ffn::BoundMoeBackend::reference(); + let h = Array2::::from_elem((2, weights.hidden_size), 0.1); + let refused = MoeFfn::best_effort(&weights, &refusing) + .forward_moe_full_layer(0, &h) + .expect("executes") + .expect("returns Some even when the route refused"); + let executed = MoeFfn::best_effort(&weights, &executing) + .forward_moe_full_layer(0, &h) + .expect("executes") + .expect("produces a layer"); + assert_ne!( + refused, executed, + "a refused expert route must not produce the same output as one that ran" + ); + } + + // ── Strictness ───────────────────────────────────────────────────────── + + #[test] + fn a_strict_adapter_records_a_refusal_that_best_effort_would_swallow() { + // The behaviour that must not become the MAP-3A contract. Best-effort + // logs, contributes zeros and returns the dense half — a plausible but + // numerically wrong continuation. Strict records it so the caller can + // refuse the output. + let weights = make_test_gemma4_moe_weights(); + // Q4_K kernels against the fixture's BF16 store: refused at bind. + let refusing = crate::ffn::BoundMoeBackend::production(); + let h = Array2::::from_elem((2, weights.hidden_size), 0.1); + + let lenient = MoeFfn::best_effort(&weights, &refusing); + let _ = lenient.forward_moe_full_layer(0, &h); + + let strict = MoeFfn::strict(&weights, &refusing); + let _ = strict.forward_moe_full_layer(0, &h); + + assert_eq!(strict.policy(), RefusalPolicy::Strict); + assert_eq!(lenient.policy(), RefusalPolicy::BestEffort); + let recorded = strict.refusal().expect("a refusal must be recorded"); + assert_eq!(recorded.layer, 0); + assert_eq!( + recorded.kind, + larql_vindex::runtime::RefusalKind::Unsupported + ); + assert!(!strict.all_experts_executed()); + } + + #[test] + fn an_executing_route_records_nothing() { + // The counter-case. If a refusal were recorded for a route that ran, + // the gate would reject every output and be quietly useless. + let weights = make_test_gemma4_moe_weights(); + let executing = crate::ffn::BoundMoeBackend::reference(); + let strict = MoeFfn::strict(&weights, &executing); + let h = Array2::::from_elem((2, weights.hidden_size), 0.1); + let _ = strict.forward_moe_full_layer(0, &h); + assert!(strict.refusal().is_none()); + assert!(strict.all_experts_executed()); + } + + #[test] + fn the_refusal_keeps_its_classification_rather_than_a_flattened_message() { + // `Residency` means fetch the operand and is not a defect; `Unsupported` + // means write a kernel. A gate that only had a string could not tell a + // sharded deployment from a broken one. + let weights = make_test_gemma4_moe_weights(); + let remote = RemoteMoeBackend::new_disconnected(); + let strict = MoeFfn::strict(&weights, &remote); + let h = Array2::::from_elem((2, weights.hidden_size), 0.1); + let _ = strict.forward_moe_full_layer(0, &h); + if let Some(recorded) = strict.refusal() { + assert_eq!( + recorded.kind, + larql_execution::RefusalKind::Residency, + "an unreachable shard is the operand not being here" + ); + assert!(!recorded.message.is_empty()); + } + } + + #[test] + fn only_the_first_refusal_is_kept() { + // Later layers usually repeat the same cause, so the earliest one is + // the diagnosis. A sink that overwrote would report the last layer + // rather than the first — the same first-versus-causal confusion the + // divergence report exists to avoid. + let weights = make_test_gemma4_moe_weights(); + let refusing = crate::ffn::BoundMoeBackend::production(); + let strict = MoeFfn::strict(&weights, &refusing); + let h = Array2::::from_elem((2, weights.hidden_size), 0.1); + let _ = strict.forward_moe_full_layer(0, &h); + let _ = strict.forward_moe_full_layer(1, &h); + assert_eq!(strict.refusal().expect("recorded").layer, 0); + } + + #[test] + fn the_remote_adapter_stays_best_effort() { + // Its callers depend on degrading rather than stopping when a shard is + // unreachable, so the generalisation must not have made it strict. + let weights = make_test_gemma4_moe_weights(); + let remote = RemoteMoeBackend::new_disconnected(); + let specific = RemoteMoeFfn { + weights: &weights, + remote: &remote, + }; + assert_eq!(specific.general().policy(), RefusalPolicy::BestEffort); + } } diff --git a/crates/larql-inference/src/ffn/moe_remote/mod.rs b/crates/larql-inference/src/ffn/moe_remote/mod.rs index ae9667f96..966ea0298 100644 --- a/crates/larql-inference/src/ffn/moe_remote/mod.rs +++ b/crates/larql-inference/src/ffn/moe_remote/mod.rs @@ -65,7 +65,7 @@ mod tests; pub use backend::RemoteMoeBackend; pub use config::{parse_unit_manifest, ShardConfig, UnitManifest, UnitShard}; pub use error::RemoteMoeError; -pub use ffn_adapter::RemoteMoeFfn; +pub use ffn_adapter::{MoeFfn, RecordedRefusal, RefusalPolicy, RemoteMoeFfn}; pub use multi_layer_wire::{ decode_multi_layer_request, decode_multi_layer_request_q8k, decode_multi_layer_response, encode_multi_layer_request, encode_multi_layer_request_q8k, encode_multi_layer_response, diff --git a/crates/larql-inference/src/ffn/remote/http.rs b/crates/larql-inference/src/ffn/remote/http.rs index b801fe7d2..cb5642730 100644 --- a/crates/larql-inference/src/ffn/remote/http.rs +++ b/crates/larql-inference/src/ffn/remote/http.rs @@ -582,7 +582,7 @@ impl FfnBackend for RemoteWalkBackend { &self, layer: usize, h_post_attn: &Array2, - ) -> Option> { + ) -> Result>, larql_execution::BoxRefusal> { let seq_len = h_post_attn.nrows(); let hidden = h_post_attn.ncols(); let residual: Vec = h_post_attn.iter().copied().collect(); @@ -594,20 +594,32 @@ impl FfnBackend for RemoteWalkBackend { "moe_layer": true, }); let url = format!("{}{WALK_FFN_PATH}", self.config.base_url); - let resp = self.client.post(&url).json(&body).send().ok()?; + // Every failure here stays `Ok(None)` — "fall back to local dispatch" — + // which is exactly what it meant before the error channel existed. + // A transport failure is an execution *attempt* failure, often + // retryable, and it is not this layer's place to decide it has become a + // semantic refusal. Promoting these to `Err` would change the remote + // walk path's behaviour under the cover of a signature migration. + let Ok(resp) = self.client.post(&url).json(&body).send() else { + return Ok(None); + }; if !resp.status().is_success() { - return None; + return Ok(None); } - let v: serde_json::Value = resp.json().ok()?; - let floats = v["output"] - .as_array()? + let Ok(v) = resp.json::() else { + return Ok(None); + }; + let Some(entries) = v["output"].as_array() else { + return Ok(None); + }; + let floats = entries .iter() .filter_map(|x| x.as_f64().map(|f| f as f32)) .collect::>(); if floats.len() != seq_len * hidden { - return None; + return Ok(None); } - Array2::from_shape_vec((seq_len, hidden), floats).ok() + Ok(Array2::from_shape_vec((seq_len, hidden), floats).ok()) } fn name(&self) -> &str { diff --git a/crates/larql-inference/src/ffn/remote/sharded.rs b/crates/larql-inference/src/ffn/remote/sharded.rs index 395023fc0..3503b22ce 100644 --- a/crates/larql-inference/src/ffn/remote/sharded.rs +++ b/crates/larql-inference/src/ffn/remote/sharded.rs @@ -307,9 +307,13 @@ impl FfnBackend for LayerShardedBackend { &self, layer: usize, h_post_attn: &Array2, - ) -> Option> { - self.shard_for(layer)? - .forward_moe_full_layer(layer, h_post_attn) + ) -> Result>, larql_execution::BoxRefusal> { + // No shard owns this layer: not applicable, not a refusal. The caller + // dispatches locally and its result is still correct. + let Some(shard) = self.shard_for(layer) else { + return Ok(None); + }; + shard.forward_moe_full_layer(layer, h_post_attn) } fn name(&self) -> &str { @@ -403,7 +407,9 @@ mod tests { "a zero delta from an unowned layer must be marked unobserved, \ never a fabricated activation tensor" ); - assert!(be.forward_moe_full_layer(0, &x).is_none()); + // No shard owns the layer: not applicable, and specifically not a + // refusal — the caller must still be free to dispatch locally. + assert!(matches!(be.forward_moe_full_layer(0, &x), Ok(None))); } #[test] diff --git a/crates/larql-inference/src/ffn_policy/router.rs b/crates/larql-inference/src/ffn_policy/router.rs index f87113604..ac0d13a34 100644 --- a/crates/larql-inference/src/ffn_policy/router.rs +++ b/crates/larql-inference/src/ffn_policy/router.rs @@ -100,7 +100,7 @@ impl<'a> FfnBackend for BoundFfnRouter<'a> { &self, layer: usize, h_post_attn: &Array2, - ) -> Option> { + ) -> Result>, larql_execution::BoxRefusal> { self.get(layer).forward_moe_full_layer(layer, h_post_attn) } } @@ -446,9 +446,9 @@ mod tests { let x = Array2::::zeros((1, weights.hidden_size)); let result = (&router as &dyn FfnBackend).forward_moe_full_layer(0, &x); assert!( - result.is_none(), - "v0 backends don't implement moe_full_layer; \ - router delegation must preserve the None default" + matches!(result, Ok(None)), + "v0 backends don't implement moe_full_layer; router delegation must \ + preserve the not-applicable default, not turn it into a refusal" ); } diff --git a/crates/larql-inference/src/forward/infer_patched.rs b/crates/larql-inference/src/forward/infer_patched.rs index d665fd893..8879aef89 100644 --- a/crates/larql-inference/src/forward/infer_patched.rs +++ b/crates/larql-inference/src/forward/infer_patched.rs @@ -240,9 +240,15 @@ pub fn infer_patched_q4k( let walk_ffn = WalkFfn::new_unlimited_with_trace(weights_ref, gate_index); let start = std::time::Instant::now(); + // `WalkFfn` serves every layer locally and leaves `forward_moe_full_layer` + // at the trait default (`Ok(None)`), so it has no way to refuse. Asserting + // that here keeps the impossibility auditable: if a refusing backend is + // ever threaded through this path it fails loudly instead of answering + // with a token the route declined to compute. let PredictResult { predictions: raw, .. - } = predict_kquant_with_ffn(weights, tokenizer, token_ids, top_k, index, &walk_ffn); + } = predict_kquant_with_ffn(weights, tokenizer, token_ids, top_k, index, &walk_ffn) + .expect("WalkFfn cannot refuse a layer; a refusal here needs a real error channel"); let walk_ms = start.elapsed().as_secs_f64() * 1000.0; let residuals = walk_ffn.take_residuals(); @@ -301,6 +307,8 @@ pub fn infer_patched_q4k_early_exit( fired = Some(ovr); Some(preds) }; + // See `infer_patched_q4k`: `WalkFfn` leaves `forward_moe_full_layer` at + // the trait default, so it cannot refuse. Loud if that ever changes. (predictions, exited) = predict_kquant_with_ffn_early_exit( weights, tokenizer, @@ -310,7 +318,8 @@ pub fn infer_patched_q4k_early_exit( &walk_ffn, stop, &mut on_stop, - ); + ) + .expect("WalkFfn cannot refuse a layer; a refusal here needs a real error channel"); } let walk_ms = start.elapsed().as_secs_f64() * 1000.0; let residuals = walk_ffn.take_residuals(); diff --git a/crates/larql-inference/src/kv_dispatch/helpers.rs b/crates/larql-inference/src/kv_dispatch/helpers.rs index 8d2a039b5..740fe50c5 100644 --- a/crates/larql-inference/src/kv_dispatch/helpers.rs +++ b/crates/larql-inference/src/kv_dispatch/helpers.rs @@ -13,12 +13,31 @@ //! tests). Engines migrate from the legacy helpers to these helpers //! in Step 3c of the ComputeBackend redesign. //! +//! **Three outcomes, never two.** Every helper here returns +//! `Result, BoxRefusal>`, the same split +//! [`FfnBackend::forward_moe_full_layer`](crate::ffn::FfnBackend::forward_moe_full_layer) +//! makes one ring down: +//! +//! ```text +//! Ok(Some(_)) the dispatch produced a complete result +//! Ok(None) nothing to do, or the backend declined this shape +//! Err(refusal) a routed operation was required and did not execute +//! ``` +//! +//! `Ok(None)` carries exactly what a bare `None` used to: the engine +//! turns it into `EngineError::BackendFailure` and may try another +//! path. `Err` is the channel that did not exist before — an engine +//! receiving it knows the layer is incomplete, so a strict route can +//! refuse the token instead of returning the dense half of a layer +//! whose experts never ran. +//! //! Hooks are not threaded through these helpers — the existing //! hooked decode path //! ([`crate::forward::generate_cached_hooked`]) keeps using the legacy //! helpers because the trait surface doesn't carry `LayerHook`. //! That's by design (`compute-backend-redesign.md` §4.2 non-goals). +use larql_execution::BoxRefusal; use ndarray::Array2; use super::{EngineBackend, KvHandle}; @@ -28,6 +47,17 @@ use crate::forward::layer::apply_layer_scalar; use crate::forward::ple::{apply_per_layer_embedding, precompute_per_layer_inputs}; use crate::forward::{embed_tokens_pub, run_ffn}; +/// What every helper in this module returns: the three outcomes, named once. +/// +/// `Ok(Some(t))` executed, `Ok(None)` not applicable, `Err(_)` refused. Spelled +/// as an alias rather than repeated six times so the contract has somewhere to +/// be documented and cannot drift between the sync and async twins. +pub type DispatchOutcome = Result, BoxRefusal>; + +/// A completed prefill: the last row of the post-FFN hidden state, plus one +/// K/V handle per layer, in layer order. +pub type PrefilledCache = (Array2, Vec); + /// Per-layer FFN + PLE + layer_scalar dispatch for the KV-cached engine /// path, MoE-aware. /// @@ -42,22 +72,30 @@ use crate::forward::{embed_tokens_pub, run_ffn}; /// `apply_per_layer_embedding` + `apply_layer_scalar`, mirroring the /// legacy `kv_prefill_run` / `kv_decode_step_run` per-layer sequence /// exactly (the issue-#98 fix; both are no-ops on non-Gemma-4 archs). +/// +/// A refusal propagates. This is what made the strict policy real: the +/// hook's three outcomes stay three all the way out to the engine, so a +/// route that declined a required expert cannot be answered with the +/// dense half of its own layer. fn ffn_or_moe_layer( weights: larql_models::WeightsView, h_post_attn: &Array2, layer: usize, ffn: &dyn FfnBackend, ple_input: Option<&Array2>, -) -> Array2 { +) -> Result, BoxRefusal> { if weights.arch.is_hybrid_moe() { - if let Some(h_out) = ffn.forward_moe_full_layer(layer, h_post_attn) { - return h_out; + // `?` propagates; `None` falls through. Not applicable means this + // backend does not serve the layer and the local dispatch below is the + // correct answer — never a refusal wearing its shape. + if let Some(h_out) = ffn.forward_moe_full_layer(layer, h_post_attn)? { + return Ok(h_out); } } let (h_post_ffn, _) = run_ffn(&weights, h_post_attn, layer, ffn, false); let mut h_out = apply_per_layer_embedding(&weights, &h_post_ffn, layer, ple_input); apply_layer_scalar(&weights, &mut h_out, layer); - h_out + Ok(h_out) } /// Prefill the K/V cache through every layer using `backend`'s @@ -69,6 +107,10 @@ fn ffn_or_moe_layer( /// it (the cache simply isn't clipped after prefill on this path — /// callers that want a clipped prefill should call /// [`KvDispatch::clip_kv`] per-layer after this returns). +/// +/// Three outcomes, per this module's contract: `Ok(Some(_))` prefilled, +/// `Ok(None)` nothing to prefill or the backend declined, `Err(_)` a +/// required routed operation refused. pub fn kv_prefill_via_dispatch( backend: &dyn EngineBackend, weights: larql_models::WeightsView, @@ -76,9 +118,9 @@ pub fn kv_prefill_via_dispatch( prompt_ids: &[u32], window: Option, index: Option<&larql_vindex::VectorIndex>, -) -> Option<(Array2, Vec)> { +) -> DispatchOutcome { if prompt_ids.is_empty() { - return None; + return Ok(None); } let h = embed_tokens_pub(&weights, prompt_ids); kv_prefill_from_hidden_via_dispatch(backend, weights, ffn, &h, Some(prompt_ids), window, index) @@ -109,9 +151,9 @@ pub fn kv_prefill_from_hidden_via_dispatch( token_ids: Option<&[u32]>, window: Option, index: Option<&larql_vindex::VectorIndex>, -) -> Option<(Array2, Vec)> { +) -> DispatchOutcome { if initial_hidden.nrows() == 0 { - return None; + return Ok(None); } let num_layers = weights.num_layers; let mut handles: Vec = Vec::with_capacity(num_layers); @@ -125,23 +167,27 @@ pub fn kv_prefill_from_hidden_via_dispatch( for layer in 0..num_layers { let _t_attn = std::time::Instant::now(); - let (h_post_attn, mut handle) = backend.attention_prefill( + // A declining backend is not a refusal — it is this dispatch having no + // answer, which is what `Ok(None)` has always meant to the engine. + let Some((h_post_attn, mut handle)) = backend.attention_prefill( weights, &h, layer, window, index.map(|v| v as &dyn larql_compute::KvIndex), - )?; + ) else { + return Ok(None); + }; crate::decode_stages::record_attn(_t_attn.elapsed().as_nanos()); if let Some(w) = window { backend.clip_kv(&mut handle, w); } handles.push(handle); - h = ffn_or_moe_layer(weights, &h_post_attn, layer, ffn, ple_inputs.get(layer)); + h = ffn_or_moe_layer(weights, &h_post_attn, layer, ffn, ple_inputs.get(layer))?; } - Some((last_row_as_2d(&h), handles)) + Ok(Some((last_row_as_2d(&h), handles))) } /// Run one autoregressive decode step using `backend`'s @@ -164,7 +210,7 @@ pub fn kv_decode_step_via_dispatch( abs_position: usize, window: Option, index: Option<&larql_vindex::VectorIndex>, -) -> Option> { +) -> DispatchOutcome> { let num_layers = weights.num_layers; debug_assert_eq!( handles.len(), @@ -179,22 +225,24 @@ pub fn kv_decode_step_via_dispatch( for (layer, handle) in handles.iter_mut().enumerate().take(num_layers) { let _t_attn = std::time::Instant::now(); - let h_post_attn = backend.attention_step( + let Some(h_post_attn) = backend.attention_step( weights, &h_step, handle, layer, abs_position, index.map(|v| v as &dyn larql_compute::KvIndex), - )?; + ) else { + return Ok(None); + }; crate::decode_stages::record_attn(_t_attn.elapsed().as_nanos()); if let Some(w) = window { backend.clip_kv(handle, w); } - h_step = ffn_or_moe_layer(weights, &h_post_attn, layer, ffn, ple_inputs.get(layer)); + h_step = ffn_or_moe_layer(weights, &h_post_attn, layer, ffn, ple_inputs.get(layer))?; } - Some(h_step) + Ok(Some(h_step)) } // ── Async variants ────────────────────────────────────────────────── @@ -222,9 +270,9 @@ pub fn kv_prefill_via_dispatch_async( prompt_ids: &[u32], window: Option, index: Option<&larql_vindex::VectorIndex>, -) -> Option<(Array2, Vec)> { +) -> DispatchOutcome { if prompt_ids.is_empty() { - return None; + return Ok(None); } let h = embed_tokens_pub(&weights, prompt_ids); kv_prefill_from_hidden_via_dispatch_async( @@ -255,9 +303,9 @@ pub fn kv_prefill_from_hidden_via_dispatch_async( token_ids: Option<&[u32]>, window: Option, index: Option<&larql_vindex::VectorIndex>, -) -> Option<(Array2, Vec)> { +) -> DispatchOutcome { if initial_hidden.nrows() == 0 { - return None; + return Ok(None); } let num_layers = weights.num_layers; let mut handles: Vec = Vec::with_capacity(num_layers); @@ -285,11 +333,13 @@ pub fn kv_prefill_from_hidden_via_dispatch_async( handles.push(handle); let h_post_attn = backend.read_hidden(h_post_attn_handle); - h = ffn_or_moe_layer(weights, &h_post_attn, layer, ffn, ple_inputs.get(layer)); + h = ffn_or_moe_layer(weights, &h_post_attn, layer, ffn, ple_inputs.get(layer))?; } - backend.flush().ok()?; - Some((last_row_as_2d(&h), handles)) + if backend.flush().is_err() { + return Ok(None); + } + Ok(Some((last_row_as_2d(&h), handles))) } /// Async equivalent of [`kv_decode_step_via_dispatch`]. @@ -307,7 +357,7 @@ pub fn kv_decode_step_via_dispatch_async( abs_position: usize, window: Option, index: Option<&larql_vindex::VectorIndex>, -) -> Option> { +) -> DispatchOutcome> { let num_layers = weights.num_layers; debug_assert_eq!( handles.len(), @@ -333,11 +383,13 @@ pub fn kv_decode_step_via_dispatch_async( backend.clip_kv(handle, w); } let h_post_attn = backend.read_hidden(h_post_attn_handle); - h_step = ffn_or_moe_layer(weights, &h_post_attn, layer, ffn, ple_inputs.get(layer)); + h_step = ffn_or_moe_layer(weights, &h_post_attn, layer, ffn, ple_inputs.get(layer))?; } - backend.flush().ok()?; - Some(h_step) + if backend.flush().is_err() { + return Ok(None); + } + Ok(Some(h_step)) } fn last_row_as_2d(h: &Array2) -> Array2 { @@ -383,7 +435,8 @@ mod tests { None, None, ) - .unwrap(); + .unwrap() + .expect("dispatch produced a result"); for step in 0..3 { let token = (2 + step) as u32; @@ -398,7 +451,8 @@ mod tests { None, None, ) - .expect("decode trait"); + .expect("decode trait") + .expect("dispatch produced a result"); assert!( h_trait.iter().all(|v| v.is_finite()), "step {step} produced non-finite hidden state" @@ -407,7 +461,10 @@ mod tests { } #[test] - fn prefill_empty_prompt_returns_none() { + fn prefill_empty_prompt_is_not_applicable_not_a_refusal() { + // An empty prompt is nothing to do, which is `Ok(None)`. Pinned as a + // shape rather than "is not Ok(Some)": collapsing it into `Err` would + // make the engine report a refusal for a caller-side input condition. let weights = make_test_weights(); let backend = CpuBackend; let ffn = WeightFfn { weights: &weights }; @@ -419,7 +476,7 @@ mod tests { None, None, ); - assert!(result.is_none()); + assert!(matches!(result, Ok(None))); } // ── Async helper parity ───────────────────────────────────────── @@ -439,7 +496,8 @@ mod tests { None, None, ) - .unwrap(); + .unwrap() + .expect("dispatch produced a result"); let (h_async, handles_async) = kv_prefill_via_dispatch_async( &backend, larql_models::WeightsView::dense(&weights), @@ -448,7 +506,8 @@ mod tests { None, None, ) - .unwrap(); + .unwrap() + .expect("dispatch produced a result"); assert_eq!(h_sync, h_async, "async prefill hidden must match sync"); assert_eq!(handles_sync.len(), handles_async.len()); @@ -476,7 +535,8 @@ mod tests { window, None, ) - .unwrap(); + .unwrap() + .expect("dispatch produced a result"); let (h_async, _) = kv_prefill_via_dispatch_async( &backend, larql_models::WeightsView::dense(&weights), @@ -485,7 +545,8 @@ mod tests { window, None, ) - .unwrap(); + .unwrap() + .expect("dispatch produced a result"); assert_eq!(h_sync, h_async, "windowed async prefill must match sync"); } @@ -505,7 +566,8 @@ mod tests { None, None, ) - .unwrap(); + .unwrap() + .expect("dispatch produced a result"); let (_, mut handles_async) = kv_prefill_via_dispatch_async( &backend, larql_models::WeightsView::dense(&weights), @@ -514,7 +576,8 @@ mod tests { None, None, ) - .unwrap(); + .unwrap() + .expect("dispatch produced a result"); let next_token = 3u32; let abs_position = prompt.len(); @@ -529,7 +592,8 @@ mod tests { None, None, ) - .unwrap(); + .unwrap() + .expect("dispatch produced a result"); let h_async = kv_decode_step_via_dispatch_async( &backend, larql_models::WeightsView::dense(&weights), @@ -540,7 +604,8 @@ mod tests { None, None, ) - .unwrap(); + .unwrap() + .expect("dispatch produced a result"); assert_eq!(h_sync, h_async, "async decode_step hidden must match sync"); } @@ -560,7 +625,8 @@ mod tests { None, None, ) - .unwrap(); + .unwrap() + .expect("dispatch produced a result"); let (_, mut handles_async) = kv_prefill_via_dispatch_async( &backend, larql_models::WeightsView::dense(&weights), @@ -569,7 +635,8 @@ mod tests { None, None, ) - .unwrap(); + .unwrap() + .expect("dispatch produced a result"); for step in 0..3 { let token = (2 + step) as u32; @@ -584,7 +651,8 @@ mod tests { None, None, ) - .unwrap(); + .unwrap() + .expect("dispatch produced a result"); let h_async = kv_decode_step_via_dispatch_async( &backend, larql_models::WeightsView::dense(&weights), @@ -595,13 +663,14 @@ mod tests { None, None, ) - .unwrap(); + .unwrap() + .expect("dispatch produced a result"); assert_eq!(h_sync, h_async, "step {step} async vs sync must match"); } } #[test] - fn prefill_async_empty_prompt_returns_none() { + fn prefill_async_empty_prompt_is_not_applicable_not_a_refusal() { let weights = make_test_weights(); let backend = CpuBackend; let ffn = WeightFfn { weights: &weights }; @@ -613,7 +682,7 @@ mod tests { None, None, ); - assert!(result.is_none()); + assert!(matches!(result, Ok(None))); } // ─── Phase 1d.3a: embed-hoist bit-identity (sync + async) ─────────────── @@ -642,7 +711,8 @@ mod tests { None, None, ) - .unwrap(); + .unwrap() + .expect("dispatch produced a result"); let initial_hidden = embed_tokens_pub(&weights, &tokens); let (h_hidden, handles_hidden) = kv_prefill_from_hidden_via_dispatch( @@ -654,7 +724,8 @@ mod tests { None, None, ) - .unwrap(); + .unwrap() + .expect("dispatch produced a result"); assert_eq!( h_text, h_hidden, @@ -682,7 +753,8 @@ mod tests { None, None, ) - .unwrap(); + .unwrap() + .expect("dispatch produced a result"); let initial_hidden = embed_tokens_pub(&weights, &tokens); let (h_hidden, handles_hidden) = kv_prefill_from_hidden_via_dispatch_async( @@ -694,7 +766,8 @@ mod tests { None, None, ) - .unwrap(); + .unwrap() + .expect("dispatch produced a result"); assert_eq!( h_text, h_hidden, @@ -704,7 +777,7 @@ mod tests { } #[test] - fn prefill_from_hidden_returns_none_on_empty_input() { + fn prefill_from_hidden_is_not_applicable_on_empty_input() { let weights = make_test_weights(); let backend = CpuBackend; let ffn = WeightFfn { weights: &weights }; @@ -718,7 +791,10 @@ mod tests { None, None, ); - assert!(result.is_none(), "zero-row hidden should yield None"); + assert!( + matches!(result, Ok(None)), + "zero-row hidden is not applicable, not a refusal" + ); let result_async = kv_prefill_from_hidden_via_dispatch_async( &backend, @@ -730,8 +806,8 @@ mod tests { None, ); assert!( - result_async.is_none(), - "async zero-row hidden should yield None" + matches!(result_async, Ok(None)), + "async zero-row hidden is not applicable, not a refusal" ); } } diff --git a/crates/larql-inference/src/kv_engine.rs b/crates/larql-inference/src/kv_engine.rs deleted file mode 100644 index 31d23c409..000000000 --- a/crates/larql-inference/src/kv_engine.rs +++ /dev/null @@ -1,1194 +0,0 @@ -//! KV-cache engine trait and shared types. -//! -//! Defines the abstract surface that the autoregressive decode loop -//! dispatches against. Concrete engine implementations (MarkovResidual, -//! UnlimitedContext, TurboQuant, Apollo, Standard, NoCache) live in -//! `larql-kv` and `impl larql_inference::KvEngine` against this trait. -//! -//! The trait deliberately lives in `larql-inference` rather than -//! `larql-kv` so the dispatch entry point (which lives here, in the -//! crate that owns the forward pass) can reference the trait without -//! a circular dependency on `larql-kv`. See -//! `docs/specs/kv-engine-unification.md` §10.4. -//! -//! Correctness contract: `prefill` and `decode_step` return the -//! pre-`lm_head` hidden state (shape `[1, hidden_dim]`). The caller -//! applies `final_norm + lm_head` to get logits — see -//! [`forward::hidden_to_raw_logits`](crate::forward::hidden_to_raw_logits). - -use crate::ffn::FfnBackend; -use crate::ModelWeights; -use ndarray::Array2; -use thiserror::Error; - -// ─── EngineError ────────────────────────────────────────────────────────────── - -/// Typed failure mode for engine `prefill` / `decode_step` calls. -/// -/// Replaces the historical `Option` return semantics that collapsed -/// "empty prompt", "backend doesn't support this", "retrieval miss", -/// "engine invariant violated" and "backend operation failed" into a -/// single opaque `None`. Two consumers (the accuracy harness and the -/// bench harness) used to route that `None` incompatibly — the -/// accuracy runner silently dropped the row via `filter_map` while the -/// bench aborted with `"engine prefill failed"`. This taxonomy lets -/// both routes branch on error *kind*; see `docs/state-policy.md`. -/// -/// The variants split error reasons along their alerting axis: -/// -/// - [`EmptyPrompt`](Self::EmptyPrompt) — caller-side input bug; -/// surfaces in CLI validation rather than a runtime alert. -/// - [`BackendUnavailable`](Self::BackendUnavailable) — the engine's -/// backend does not implement the requested capability (e.g. a -/// Metal kernel that hasn't been ported, an asked-for Q4K matvec on -/// a CPU build without BLAS). Falls back to a different code path -/// *if* one exists; otherwise surfaces as a configuration error. -/// - [`RetrievalMiss { reason }`](Self::RetrievalMiss) — a retrieval -/// engine (Apollo, future Mode 5) could not serve this query against -/// its store. Expected, recoverable; surfaces in the harness as a -/// `served_rate < 1.0` column rather than an alert. -/// - [`InvariantViolation { what }`](Self::InvariantViolation) — the -/// engine was driven outside its state-machine contract (e.g. -/// `decode_step` called before `prefill`). Indicates a harness-level -/// dispatch bug; production observability should alert immediately. -/// - [`BackendFailure { details }`](Self::BackendFailure) — the inner -/// backend or compute kernel returned a runtime failure. Indicates a -/// data condition or environmental issue (corrupt weights, OOM, GPU -/// driver error); production observability should log + investigate -/// but not alert with the same urgency as `InvariantViolation`. -/// -/// `InvariantViolation` and `BackendFailure` are deliberately kept as -/// two top-level variants rather than collapsed into a single -/// `InternalError { kind }`. Sub-tagged enums lose the alert-routing -/// distinction when the consumer writes `match err { InternalError(_) => ... }`. -/// -/// The enum is **exhaustive** (no `#[non_exhaustive]`). New variants -/// are breaking changes on purpose — defaulting a new condition into -/// an existing arm reproduces the silent-drop problem one layer down. -#[derive(Debug, Clone, PartialEq, Eq, Error)] -pub enum EngineError { - #[error("engine called with empty prompt")] - EmptyPrompt, - #[error("backend does not support this operation")] - BackendUnavailable, - #[error("retrieval miss: {reason}")] - RetrievalMiss { reason: String }, - #[error("engine invariant violated: {what}")] - InvariantViolation { what: String }, - #[error("backend operation failed: {details}")] - BackendFailure { details: String }, -} - -impl EngineError { - /// `true` for variants the harness should treat as recoverable - /// (skip + log + continue), `false` for variants that indicate a - /// dispatch bug or kernel failure (abort + investigate). - /// - /// `EmptyPrompt` / `BackendUnavailable` / `RetrievalMiss` are - /// recoverable; `InvariantViolation` / `BackendFailure` are not. - pub fn is_recoverable(&self) -> bool { - matches!( - self, - Self::EmptyPrompt | Self::BackendUnavailable | Self::RetrievalMiss { .. } - ) - } -} - -// ─── EngineInfo ─────────────────────────────────────────────────────────────── - -/// Runtime diagnostics reported by each engine. -#[derive(Debug, Clone)] -pub struct EngineInfo { - /// Short engine name (e.g. `"markov-rs"`). - pub name: String, - /// Human-readable description of the engine's state management strategy. - pub description: String, - /// Hardware backend name from [`larql_compute::ComputeBackend::name`]: `"cpu"`, `"metal"`, etc. - pub backend: String, - /// Key config parameters (e.g. `"window=512"`), empty string if unconfigured. - pub config: String, -} - -impl EngineInfo { - pub fn summary(&self) -> String { - if self.config.is_empty() { - format!("{} [{}] {}", self.name, self.backend, self.description) - } else { - format!( - "{} [{}] ({}) {}", - self.name, self.backend, self.config, self.description - ) - } - } -} - -// ─── DecodeStageSummary ─────────────────────────────────────────────────────── - -/// Per-step averages for a completed engine run. Returned from -/// [`KvEngine::stage_summary`] when profiling was enabled at engine -/// construction. -#[derive(Debug, Clone)] -pub struct DecodeStageSummary { - pub engine: String, - pub backend: String, - pub steps: usize, - pub avg_embed_us: f64, - /// K/V recompute from stored residuals (MarkovRS only). Split by tier. - pub avg_recompute_cold_us: f64, - pub avg_recompute_hot_us: f64, - pub avg_attention_us: f64, - pub avg_ffn_us: f64, - pub avg_total_decode_us: f64, - /// W10 instrumentation: time spent inside the backend's - /// `coarse_decode_step_with_state_masked` call — kernel run + - /// state-dump readback (skipped under HOnly / None). Zero on - /// non-dispatch paths and on engines that don't capture state. - pub avg_state_capture_us: f64, - /// W10 instrumentation: cumulative time inside per-layer handle - /// materialise calls (`StateHandle::into_array`). Tracks the - /// CPU bridge cost from the captured dump to engine-owned - /// `Array2`s. Zero under None mask (engine drops handles - /// without materialising). - pub avg_state_materialise_us: f64, - /// W10 instrumentation: cumulative time appending materialised - /// state into engine slabs (`append_row` calls). Tracks - /// `rs.stored` / `rs.hot_kv` growth. Zero under None mask. - pub avg_state_append_us: f64, -} - -impl DecodeStageSummary { - pub fn avg_recompute_total_us(&self) -> f64 { - self.avg_recompute_cold_us + self.avg_recompute_hot_us - } - - /// Print a human-readable breakdown table. - pub fn print(&self) { - let total = self.avg_total_decode_us; - let pct = |v: f64| if total > 0.0 { v / total * 100.0 } else { 0.0 }; - - println!( - "\nStage breakdown ({}, {}, {} decode steps avg):", - self.engine, self.backend, self.steps - ); - println!(" {:<25} {:>8} {:>6}", "Stage", "avg_us", "%"); - println!(" {}", "-".repeat(45)); - println!( - " {:<25} {:>8.1} {:>5.1}%", - "embed", - self.avg_embed_us, - pct(self.avg_embed_us) - ); - if self.avg_recompute_total_us() > 0.0 { - println!( - " {:<25} {:>8.1} {:>5.1}%", - "recompute_kv (cold)", - self.avg_recompute_cold_us, - pct(self.avg_recompute_cold_us) - ); - println!( - " {:<25} {:>8.1} {:>5.1}%", - "recompute_kv (hot)", - self.avg_recompute_hot_us, - pct(self.avg_recompute_hot_us) - ); - } - println!( - " {:<25} {:>8.1} {:>5.1}%", - "attention", - self.avg_attention_us, - pct(self.avg_attention_us) - ); - println!( - " {:<25} {:>8.1} {:>5.1}%", - "ffn", - self.avg_ffn_us, - pct(self.avg_ffn_us) - ); - // W10 instrumentation: only print state lines when populated - // (avoids noise on engines that don't capture state). - let state_total = - self.avg_state_capture_us + self.avg_state_materialise_us + self.avg_state_append_us; - if state_total > 0.0 { - println!( - " {:<25} {:>8.1} {:>5.1}%", - "state_capture", - self.avg_state_capture_us, - pct(self.avg_state_capture_us) - ); - println!( - " {:<25} {:>8.1} {:>5.1}%", - "state_materialise", - self.avg_state_materialise_us, - pct(self.avg_state_materialise_us) - ); - println!( - " {:<25} {:>8.1} {:>5.1}%", - "state_append", - self.avg_state_append_us, - pct(self.avg_state_append_us) - ); - } - println!(" {}", "-".repeat(45)); - println!( - " {:<25} {:>8.1} {:>5.1}%", - "total (measured)", total, 100.0 - ); - println!(); - } -} - -// ─── KvEngine trait ─────────────────────────────────────────────────────────── - -/// Common interface shared by all KV-cache engines. -pub trait KvEngine: Send { - fn name(&self) -> &str; - - /// Runtime diagnostics: engine name, backend, config, description. - fn info(&self) -> EngineInfo; - - /// Run the prefill forward pass over all prompt tokens. - /// - /// `ffn` is the FFN backend the engine should dispatch through — - /// typically [`WeightFfn`](crate::ffn::WeightFfn) / - /// [`BackendFfn`](crate::ffn::BackendFfn) for local compute, or - /// [`RemoteWalkBackend`](crate::ffn::RemoteWalkBackend) for grid - /// routing. Engines that don't consult an FFN router (e.g. ones - /// that recompute FFN from `weights` directly) may ignore this - /// parameter. - /// - /// Returns the hidden state at the final token position (shape `[1, hidden_dim]`). - /// - /// Failure modes surface as typed [`EngineError`] variants — see - /// the enum's docs for the routing taxonomy. - fn prefill( - &mut self, - weights: &ModelWeights, - ffn: &dyn FfnBackend, - token_ids: &[u32], - ) -> Result, EngineError>; - - /// Run one autoregressive decode step for a single new token. - /// Returns the hidden state (shape `[1, hidden_dim]`). - fn decode_step( - &mut self, - weights: &ModelWeights, - ffn: &dyn FfnBackend, - token_id: u32, - ) -> Result, EngineError>; - - /// Static capability: does this engine accept pre-built hidden - /// state via [`prefill_from_hidden`]? Default `false`. - /// - /// The CLI MUST check this **before** running a (potentially - /// minutes-long) modal encoder, so the user gets a fast, clear - /// error if they paired `--image` with an engine that doesn't - /// support multi-modal input. See ADR-0023. - /// - /// The default-false return is deliberate debt — six of seven - /// engines inherit it for Phase 1d. The end state collapses - /// `prefill(token_ids)` into a thin wrapper over - /// `embed_tokens_pub` then `prefill_from_hidden` on every engine, - /// at which point this method becomes universally `true` and is - /// removed. Tracked in ADR-0023 (Default-false debt). - fn supports_multimodal(&self) -> bool { - false - } - - /// Prefill from a pre-built initial hidden state. Caller built it - /// via `larql_compute::forward::embed_plan` from an `EmbeddingPlan` - /// that may include `Precomputed` rows (vision / audio embeddings). - /// - /// Same contract as [`prefill`]: runs forward through every layer, - /// populates the engine's KV cache, returns the final-token hidden - /// state. Returns the same `Result<_, EngineError>` shape as - /// `prefill` for uniform call-site error handling. The engine's - /// internal absolute position pointer must be set from - /// `initial_hidden.nrows()`, NOT from any token count — the input - /// may contain non-token positions. - /// - /// Default impl panics (not an `Err` return) on engines that don't - /// override it. Callers MUST check [`supports_multimodal`] first; - /// the panic is defense-in-depth against bypass, not a substitute - /// for the capability check. - fn prefill_from_hidden( - &mut self, - _weights: &ModelWeights, - _ffn: &dyn FfnBackend, - _initial_hidden: &Array2, - ) -> Result, EngineError> { - panic!( - "engine {:?} does not support multi-modal input; \ - check supports_multimodal() before calling prefill_from_hidden", - self.name() - ); - } - - /// Bytes of persistent engine state (excludes model weights). - fn memory_bytes(&self) -> usize; - - /// Token count in the active hot window (varies by engine type). - fn window_tokens(&self) -> usize { - 0 - } - - /// Cold-tier bytes (residuals or token IDs past the hot window). - fn cold_bytes(&self) -> usize { - 0 - } - - /// Per-stage timing summary. Returns `None` if profiling was not enabled. - fn stage_summary(&self) -> Option { - None - } - - /// Prefill using Q4K quantised weights from `index` and `backend`. - /// - /// When the backend supports the fused Q4 pipeline (Metal), this routes - /// through `backend.prefill_kquant` for full GPU speed. Falls back to the - /// f32 path when `backend.supports_quant(::larql_compute::QuantFormat::Q4_K) == false` or `index` has no Q4K data. - /// - /// `weights` is `&ModelWeights` (immutable): the engine dequantises f32 - /// attention tensors into its own `dequant_scratch` on the first call and - /// resolves them via `WeightsView::with_scratch` (one-time cost; subsequent - /// decode steps reuse the engine-owned scratch). - fn prefill_quant( - &mut self, - weights: &ModelWeights, - ffn: &dyn FfnBackend, - index: &larql_vindex::VectorIndex, - token_ids: &[u32], - backend: &dyn larql_compute::ComputeBackend, - ) -> Result, EngineError> { - let _ = (index, backend); - self.prefill(weights, ffn, token_ids) // default: f32 fallback - } - - /// One autoregressive decode step using Q4K weights. - /// - /// Same routing semantics as [`prefill_quant`]: Metal via `decode_token` - /// when available, f32 fallback otherwise. - fn decode_step_quant( - &mut self, - weights: &ModelWeights, - ffn: &dyn FfnBackend, - index: &larql_vindex::VectorIndex, - token_id: u32, - backend: &dyn larql_compute::ComputeBackend, - ) -> Result, EngineError> { - let _ = (index, backend); - self.decode_step(weights, ffn, token_id) // default: f32 fallback - } - - /// Resident-weights quant prefill. Unlike [`prefill_quant`] (which - /// dequantises attn into the engine's `dequant_scratch`), this assumes the - /// **caller has already made the client weights f32-resident** — so it - /// merely threads `index` to the backend. That lets a Q4K-direct attention - /// kernel (`LARQL_Q4K_DIRECT_ATTN`) read packed bytes from the index while - /// the FFN backend borrows the same `&weights` (task #16). Default: f32 - /// fallback (index ignored). - fn prefill_resident( - &mut self, - weights: &ModelWeights, - ffn: &dyn FfnBackend, - index: &larql_vindex::VectorIndex, - token_ids: &[u32], - ) -> Result, EngineError> { - let _ = index; - self.prefill(weights, ffn, token_ids) - } - - /// One decode step against resident (pre-dequantised) weights, threading - /// `index` to the backend. Sibling of [`prefill_resident`]; same rationale. - /// Default: f32 fallback (index ignored). - fn decode_step_resident( - &mut self, - weights: &ModelWeights, - ffn: &dyn FfnBackend, - index: &larql_vindex::VectorIndex, - token_id: u32, - ) -> Result, EngineError> { - let _ = index; - self.decode_step(weights, ffn, token_id) - } - - /// Prefill via a caller-supplied `LayerExecutor` (dense/f32 path). - /// See [`docs/specs/engine-state-vs-execution.md`]. - /// - /// Sibling of [`prefill_quant_via_executor`] for engines that - /// don't have a quant path (no vindex needed). Default impl falls - /// through to [`prefill`]. - fn prefill_via_executor( - &mut self, - weights: &ModelWeights, - executor: &dyn crate::layer_executor::LayerExecutor, - ffn: &dyn FfnBackend, - token_ids: &[u32], - ) -> Result, EngineError> { - let _ = executor; - self.prefill(weights, ffn, token_ids) - } - - /// One decode step via a caller-supplied `LayerExecutor` (dense/f32). - /// Sibling of [`decode_step_quant_via_executor`]. - fn decode_step_via_executor( - &mut self, - weights: &ModelWeights, - executor: &dyn crate::layer_executor::LayerExecutor, - ffn: &dyn FfnBackend, - token_id: u32, - ) -> Result, EngineError> { - let _ = executor; - self.decode_step(weights, ffn, token_id) - } - - /// Prefill via a caller-supplied `LayerExecutor`. See - /// [`docs/specs/engine-state-vs-execution.md`]. - /// - /// The default impl falls through to [`prefill_quant`] using - /// `executor.backend()` — engines that haven't migrated yet keep - /// working unchanged. Migrated engines override this method to - /// drive the layer loop through the executor and honor the FFN - /// parameter properly. - fn prefill_quant_via_executor( - &mut self, - weights: &ModelWeights, - executor: &dyn crate::layer_executor::LayerExecutor, - ffn: &dyn FfnBackend, - index: &larql_vindex::VectorIndex, - token_ids: &[u32], - ) -> Result, EngineError> { - self.prefill_quant(weights, ffn, index, token_ids, executor.backend()) - } - - /// One decode step via a caller-supplied `LayerExecutor`. See - /// [`prefill_quant_via_executor`] for the migration contract. - fn decode_step_quant_via_executor( - &mut self, - weights: &ModelWeights, - executor: &dyn crate::layer_executor::LayerExecutor, - ffn: &dyn FfnBackend, - index: &larql_vindex::VectorIndex, - token_id: u32, - ) -> Result, EngineError> { - self.decode_step_quant(weights, ffn, index, token_id, executor.backend()) - } -} - -// ─── RetrievalEngine trait ──────────────────────────────────────────────────── - -/// Engines whose state is **not** an autoregressive K/V cache. -/// -/// Sibling trait to [`KvEngine`]. Both surfaces share the -/// [`EngineInfo`] / [`EngineError`] / [`DecodeStageSummary`] vocabulary -/// but diverge on the per-step contract: -/// -/// | `KvEngine` | `RetrievalEngine` | -/// |-------------------------------------|----------------------------------------| -/// | per-token K/V append per layer | retrieval against a pre-built store | -/// | dispatches FFN through `FfnBackend` | does not consult an FFN router | -/// | state reconstructible to K/V tensors| state is residual delta + token list | -/// -/// Apollo (boundary-residual + injection delta) lands here; Mode 5 / -/// Graph-Grounded engines will too. The trait deliberately drops the -/// per-step K/V append assumption and the `FfnBackend` parameter that -/// `KvEngine::prefill` carries — both went unused on the retrieval -/// side (`_ffn` in the Apollo impl) and forced harnesses to construct -/// a router they then ignored. -/// -/// Returns [`Result`] instead of `Option` for the -/// same reasons as `KvEngine`'s post-2026-05-24 migration: silent -/// `None` propagation through `filter_map` masked Apollo's -/// store-miss rate, and the bench `panic!` on the same `None` made -/// retrieval-miss prompts crash a run that ought to have logged a -/// skip. See [`EngineError`] for the variant taxonomy. -pub trait RetrievalEngine: Send { - fn name(&self) -> &str; - - /// Runtime diagnostics: engine name, backend, config, description. - fn info(&self) -> EngineInfo; - - /// Run the prefill forward pass over the prompt tokens, consulting - /// the engine's retrieval store. Returns the hidden state at the - /// final token position (shape `[1, hidden_dim]`). - fn prefill( - &mut self, - weights: &ModelWeights, - token_ids: &[u32], - ) -> Result, EngineError>; - - /// One autoregressive decode step for the next token, applying any - /// retrieval-engine-specific state update (e.g. injection-delta - /// accumulation). Returns the hidden state (shape `[1, hidden_dim]`). - fn decode_step( - &mut self, - weights: &ModelWeights, - token_id: u32, - ) -> Result, EngineError>; - - /// Prefill against a Q4K-quantised vindex. **No default** — the trait - /// default returns an `InvariantViolation` because the `ffn`-less `prefill` - /// it would delegate to can't be handed an engine-owned dequant scratch - /// without mutating `weights`. Engines that serve Q4K (e.g. Apollo, which - /// runs its forward through `forward_raw_logits` and dequantises attn+FFN - /// into its own `dequant_scratch`) override this. - fn prefill_quant( - &mut self, - weights: &ModelWeights, - index: &larql_vindex::VectorIndex, - token_ids: &[u32], - ) -> Result, EngineError> { - let _ = (weights, index, token_ids); - Err(EngineError::InvariantViolation { - what: "RetrievalEngine::prefill_quant must be overridden for Q4K vindexes — the \ - default cannot thread an engine-owned dequant scratch through the \ - `ffn`-less `prefill` (it would have to mutate `weights`)." - .into(), - }) - } - - /// One decode step against a Q4K-quantised vindex. No default — engines - /// that serve Q4K must override (see [`prefill_quant`](Self::prefill_quant)). - fn decode_step_quant( - &mut self, - weights: &ModelWeights, - index: &larql_vindex::VectorIndex, - token_id: u32, - ) -> Result, EngineError> { - let _ = (weights, index, token_id); - Err(EngineError::InvariantViolation { - what: "RetrievalEngine::decode_step_quant must be overridden for Q4K vindexes.".into(), - }) - } - - /// Bytes of persistent engine state (excludes model weights). - fn memory_bytes(&self) -> usize; - - /// Token count in the active window (varies by engine type). - fn window_tokens(&self) -> usize { - 0 - } - - /// Cold-tier bytes (store / archive past the hot window). - fn cold_bytes(&self) -> usize { - 0 - } - - /// Per-stage timing summary. Returns `None` if profiling was not enabled. - fn stage_summary(&self) -> Option { - None - } -} - -// ─── AnyEngine ──────────────────────────────────────────────────────────────── - -/// Sum type that holds either a [`KvEngine`] or a [`RetrievalEngine`]. -/// -/// Construction sites (the engine builder in -/// `larql-kv::EngineBuilder::build`, the bench / accuracy harnesses) -/// parse a spec into one or the other; the autoregressive loop calls -/// uniform `prefill` / `decode_step` (et al.) methods on `AnyEngine`, -/// which pattern-match internally on the variant. -/// -/// Each forwarding method takes the superset of arguments from both -/// trait surfaces. For [`RetrievalEngine`] engines (Apollo, future -/// Mode 5) the FFN-routing and compute-backend arguments are simply -/// ignored — `RetrievalEngine` runs its forward through -/// `forward_from_layer` / `forward_raw_logits` and doesn't need them. -/// This is intentional: keeping the harness call site uniform across -/// new engine families is more important than enforcing argument -/// minimality at the type level. Variant-specific behaviour still -/// surfaces through the [`EngineError`] enum (e.g. `RetrievalMiss` -/// only arrives from retrieval engines). -pub enum AnyEngine { - Kv(Box), - Retrieval(Box), -} - -impl AnyEngine { - pub fn name(&self) -> &str { - match self { - Self::Kv(e) => e.name(), - Self::Retrieval(e) => e.name(), - } - } - - pub fn info(&self) -> EngineInfo { - match self { - Self::Kv(e) => e.info(), - Self::Retrieval(e) => e.info(), - } - } - - pub fn memory_bytes(&self) -> usize { - match self { - Self::Kv(e) => e.memory_bytes(), - Self::Retrieval(e) => e.memory_bytes(), - } - } - - pub fn window_tokens(&self) -> usize { - match self { - Self::Kv(e) => e.window_tokens(), - Self::Retrieval(e) => e.window_tokens(), - } - } - - pub fn cold_bytes(&self) -> usize { - match self { - Self::Kv(e) => e.cold_bytes(), - Self::Retrieval(e) => e.cold_bytes(), - } - } - - pub fn stage_summary(&self) -> Option { - match self { - Self::Kv(e) => e.stage_summary(), - Self::Retrieval(e) => e.stage_summary(), - } - } - - pub fn is_kv(&self) -> bool { - matches!(self, Self::Kv(_)) - } - - pub fn is_retrieval(&self) -> bool { - matches!(self, Self::Retrieval(_)) - } - - // ── Forwarding methods (variant-specific dispatch) ────────────────────── - - /// Prefill. KvEngine variants consult `ffn`; retrieval variants - /// ignore it. - pub fn prefill( - &mut self, - weights: &ModelWeights, - ffn: &dyn FfnBackend, - token_ids: &[u32], - ) -> Result, EngineError> { - match self { - Self::Kv(e) => e.prefill(weights, ffn, token_ids), - Self::Retrieval(e) => e.prefill(weights, token_ids), - } - } - - /// One autoregressive decode step. Same routing as [`prefill`]. - pub fn decode_step( - &mut self, - weights: &ModelWeights, - ffn: &dyn FfnBackend, - token_id: u32, - ) -> Result, EngineError> { - match self { - Self::Kv(e) => e.decode_step(weights, ffn, token_id), - Self::Retrieval(e) => e.decode_step(weights, token_id), - } - } - - /// Capability forwarder for multi-modal input — see ADR-0023. - /// `Retrieval` variants are text-only by construction and always - /// return `false`; `Kv` variants delegate to the trait method. - pub fn supports_multimodal(&self) -> bool { - match self { - Self::Kv(e) => e.supports_multimodal(), - Self::Retrieval(_) => false, - } - } - - /// MM prefill forwarder. Only `Kv` variants can implement this; - /// `Retrieval` variants panic when called (callers MUST check - /// `supports_multimodal()` first, per ADR-0023). - pub fn prefill_from_hidden( - &mut self, - weights: &ModelWeights, - ffn: &dyn FfnBackend, - initial_hidden: &Array2, - ) -> Result, EngineError> { - match self { - Self::Kv(e) => e.prefill_from_hidden(weights, ffn, initial_hidden), - Self::Retrieval(_) => panic!( - "AnyEngine::Retrieval does not support prefill_from_hidden — \ - check supports_multimodal() before calling" - ), - } - } - - /// Prefill against a quantised vindex. KvEngine variants take a - /// `ComputeBackend` for kernel routing; retrieval variants ignore - /// it (they dequantise + run on f32 internally). - pub fn prefill_quant( - &mut self, - weights: &ModelWeights, - ffn: &dyn FfnBackend, - index: &larql_vindex::VectorIndex, - token_ids: &[u32], - backend: &dyn larql_compute::ComputeBackend, - ) -> Result, EngineError> { - match self { - Self::Kv(e) => e.prefill_quant(weights, ffn, index, token_ids, backend), - Self::Retrieval(e) => e.prefill_quant(weights, index, token_ids), - } - } - - /// One decode step against a quantised vindex. Same routing as - /// [`prefill_quant`]. - pub fn decode_step_quant( - &mut self, - weights: &ModelWeights, - ffn: &dyn FfnBackend, - index: &larql_vindex::VectorIndex, - token_id: u32, - backend: &dyn larql_compute::ComputeBackend, - ) -> Result, EngineError> { - match self { - Self::Kv(e) => e.decode_step_quant(weights, ffn, index, token_id, backend), - Self::Retrieval(e) => e.decode_step_quant(weights, index, token_id), - } - } - - /// Resident-weights quant prefill (`&weights`, threads `index`). See - /// [`KvEngine::prefill_resident`]. Retrieval variants fall back to their - /// f32 prefill (index-aware retrieval isn't a moe-shards path). - pub fn prefill_resident( - &mut self, - weights: &ModelWeights, - ffn: &dyn FfnBackend, - index: &larql_vindex::VectorIndex, - token_ids: &[u32], - ) -> Result, EngineError> { - match self { - Self::Kv(e) => e.prefill_resident(weights, ffn, index, token_ids), - Self::Retrieval(e) => e.prefill(weights, token_ids), - } - } - - /// Resident-weights quant decode step (`&weights`, threads `index`). See - /// [`KvEngine::decode_step_resident`]. - pub fn decode_step_resident( - &mut self, - weights: &ModelWeights, - ffn: &dyn FfnBackend, - index: &larql_vindex::VectorIndex, - token_id: u32, - ) -> Result, EngineError> { - match self { - Self::Kv(e) => e.decode_step_resident(weights, ffn, index, token_id), - Self::Retrieval(e) => e.decode_step(weights, token_id), - } - } - - /// Prefill via a caller-supplied [`crate::layer_executor::LayerExecutor`]. - /// Falls back to [`prefill_quant`](Self::prefill_quant) for - /// [`RetrievalEngine`] variants (which don't drive per-layer - /// executor loops). - pub fn prefill_quant_via_executor( - &mut self, - weights: &ModelWeights, - executor: &dyn crate::layer_executor::LayerExecutor, - ffn: &dyn FfnBackend, - index: &larql_vindex::VectorIndex, - token_ids: &[u32], - ) -> Result, EngineError> { - match self { - Self::Kv(e) => e.prefill_quant_via_executor(weights, executor, ffn, index, token_ids), - Self::Retrieval(e) => e.prefill_quant(weights, index, token_ids), - } - } - - /// One decode step via a caller-supplied `LayerExecutor`. Same - /// fall-back semantics as [`prefill_quant_via_executor`]. - pub fn decode_step_quant_via_executor( - &mut self, - weights: &ModelWeights, - executor: &dyn crate::layer_executor::LayerExecutor, - ffn: &dyn FfnBackend, - index: &larql_vindex::VectorIndex, - token_id: u32, - ) -> Result, EngineError> { - match self { - Self::Kv(e) => { - e.decode_step_quant_via_executor(weights, executor, ffn, index, token_id) - } - Self::Retrieval(e) => e.decode_step_quant(weights, index, token_id), - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn engine_info_summary_with_config() { - let info = EngineInfo { - name: "markov-rs".into(), - description: "residual KV".into(), - backend: "cpu".into(), - config: "window=512".into(), - }; - let s = info.summary(); - assert!(s.contains("markov-rs")); - assert!(s.contains("cpu")); - assert!(s.contains("window=512")); - } - - #[test] - fn engine_info_summary_no_config() { - let info = EngineInfo { - name: "test".into(), - description: "desc".into(), - backend: "metal".into(), - config: String::new(), - }; - let s = info.summary(); - assert!(!s.contains("()")); - } - - #[test] - fn decode_stage_summary_recompute_total() { - let s = DecodeStageSummary { - engine: "test".into(), - backend: "cpu".into(), - steps: 10, - avg_embed_us: 1.0, - avg_recompute_cold_us: 2.0, - avg_recompute_hot_us: 3.0, - avg_attention_us: 4.0, - avg_ffn_us: 5.0, - avg_total_decode_us: 15.0, - avg_state_capture_us: 0.0, - avg_state_materialise_us: 0.0, - avg_state_append_us: 0.0, - }; - assert_eq!(s.avg_recompute_total_us(), 5.0); - } - - /// Cover `DecodeStageSummary::print` — both the recompute>0 branch and - /// the total>0 percentage branch. Output goes to stdout (captured by the - /// test harness); this is a smoke test for the formatting code path. - #[test] - fn decode_stage_summary_print_with_recompute() { - let s = DecodeStageSummary { - engine: "markov-rs".into(), - backend: "cpu".into(), - steps: 10, - avg_embed_us: 100.0, - avg_recompute_cold_us: 500.0, - avg_recompute_hot_us: 300.0, - avg_attention_us: 1500.0, - avg_ffn_us: 800.0, - avg_total_decode_us: 3200.0, - avg_state_capture_us: 0.0, - avg_state_materialise_us: 0.0, - avg_state_append_us: 0.0, - }; - s.print(); - } - - /// `print` must also handle the no-recompute, zero-total branch — the - /// `pct` fallback when `avg_total_decode_us == 0.0` and the - /// `avg_recompute_total_us() == 0` short-circuit. - #[test] - fn decode_stage_summary_print_no_recompute_zero_total() { - let s = DecodeStageSummary { - engine: "no-cache".into(), - backend: "metal".into(), - steps: 0, - avg_embed_us: 0.0, - avg_recompute_cold_us: 0.0, - avg_recompute_hot_us: 0.0, - avg_attention_us: 0.0, - avg_ffn_us: 0.0, - avg_total_decode_us: 0.0, - avg_state_capture_us: 0.0, - avg_state_materialise_us: 0.0, - avg_state_append_us: 0.0, - }; - s.print(); - } - - /// Synthetic engine that only implements the required trait methods, - /// leaving every default (`window_tokens`, `cold_bytes`, `stage_summary`, - /// `prefill_quant`, `decode_step_quant`) to fire. Exercises the default - /// bodies that no shipped engine routes through (every concrete engine - /// overrides them). - struct DefaultsOnlyEngine { - prefill_calls: usize, - decode_calls: usize, - } - - impl KvEngine for DefaultsOnlyEngine { - fn name(&self) -> &str { - "defaults-only" - } - fn info(&self) -> EngineInfo { - EngineInfo { - name: self.name().into(), - description: "test fixture".into(), - backend: "cpu".into(), - config: String::new(), - } - } - fn prefill( - &mut self, - _weights: &ModelWeights, - _ffn: &dyn FfnBackend, - _token_ids: &[u32], - ) -> Result, EngineError> { - self.prefill_calls += 1; - Ok(Array2::zeros((1, 4))) - } - fn decode_step( - &mut self, - _weights: &ModelWeights, - _ffn: &dyn FfnBackend, - _token_id: u32, - ) -> Result, EngineError> { - self.decode_calls += 1; - Ok(Array2::zeros((1, 4))) - } - fn memory_bytes(&self) -> usize { - 0 - } - } - - #[test] - fn defaults_window_tokens_and_cold_bytes_are_zero() { - let engine = DefaultsOnlyEngine { - prefill_calls: 0, - decode_calls: 0, - }; - assert_eq!(engine.window_tokens(), 0); - assert_eq!(engine.cold_bytes(), 0); - assert!(engine.stage_summary().is_none()); - assert_eq!(engine.name(), "defaults-only"); - } - - /// All four `*_via_executor` default impls dispatch through to their - /// non-executor sibling, which on `DefaultsOnlyEngine` falls back to - /// `prefill` / `decode_step`. Covers the function bodies of - /// `prefill_via_executor` (224-233), `decode_step_via_executor` - /// (237-246), `prefill_quant_via_executor` (256-265), - /// `decode_step_quant_via_executor` (269-278). - #[test] - fn defaults_via_executor_methods_dispatch_to_non_executor_siblings() { - struct StubExecutor { - backend: larql_compute::CpuBackend, - } - impl crate::layer_executor::LayerExecutor for StubExecutor { - fn backend(&self) -> &dyn larql_compute::ComputeBackend { - &self.backend - } - fn dispatch_kind(&self) -> crate::layer_executor::ExecutorDispatchKind { - crate::layer_executor::ExecutorDispatchKind::PerLayer - } - fn name(&self) -> &str { - "stub" - } - } - let exec = StubExecutor { - backend: larql_compute::CpuBackend, - }; - let weights = crate::test_utils::make_test_weights(); - let index = crate::test_utils::make_test_vindex(&weights); - let ffn = crate::ffn::WeightFfn { weights: &weights }; - let mut engine = DefaultsOnlyEngine { - prefill_calls: 0, - decode_calls: 0, - }; - - // prefill_via_executor → prefill - let out = engine.prefill_via_executor(&weights, &exec, &ffn, &[0, 1]); - assert!(out.is_ok()); - assert_eq!(engine.prefill_calls, 1); - - // decode_step_via_executor → decode_step - let out = engine.decode_step_via_executor(&weights, &exec, &ffn, 2); - assert!(out.is_ok()); - assert_eq!(engine.decode_calls, 1); - - // prefill_quant_via_executor → prefill_quant → prefill (default fallback) - let weights_q = crate::test_utils::make_test_weights(); - let out = engine.prefill_quant_via_executor(&weights_q, &exec, &ffn, &index, &[0, 1]); - assert!(out.is_ok()); - assert_eq!(engine.prefill_calls, 2); - - // decode_step_quant_via_executor → decode_step_quant → decode_step - let out = engine.decode_step_quant_via_executor(&weights_q, &exec, &ffn, &index, 3); - assert!(out.is_ok()); - assert_eq!(engine.decode_calls, 2); - } - - #[test] - fn defaults_q4k_methods_fall_back_to_f32() { - let weights = crate::test_utils::make_test_weights(); - let index = crate::test_utils::make_test_vindex(&weights); - let backend = larql_compute::cpu_backend(); - let ffn = crate::ffn::WeightFfn { weights: &weights }; - let mut engine = DefaultsOnlyEngine { - prefill_calls: 0, - decode_calls: 0, - }; - - let weights_q4k = crate::test_utils::make_test_weights(); - let out = engine.prefill_quant(&weights_q4k, &ffn, &index, &[1, 2, 3], &*backend); - assert!(out.is_ok()); - assert_eq!( - engine.prefill_calls, 1, - "default prefill_quant must dispatch to prefill" - ); - - let out = engine.decode_step_quant(&weights_q4k, &ffn, &index, 4, &*backend); - assert!(out.is_ok()); - assert_eq!( - engine.decode_calls, 1, - "default decode_step_quant must dispatch to decode_step" - ); - } - - // ─── EngineError ────────────────────────────────────────────────────────── - - #[test] - fn engine_error_is_recoverable_classifies_variants() { - assert!(EngineError::EmptyPrompt.is_recoverable()); - assert!(EngineError::BackendUnavailable.is_recoverable()); - assert!(EngineError::RetrievalMiss { - reason: "no store".into() - } - .is_recoverable()); - assert!(!EngineError::InvariantViolation { - what: "decode before prefill".into() - } - .is_recoverable()); - assert!(!EngineError::BackendFailure { - details: "kernel oom".into() - } - .is_recoverable()); - } - - #[test] - fn engine_error_display_includes_reason_payload() { - let err = EngineError::RetrievalMiss { - reason: "no store attached".into(), - }; - assert!(err.to_string().contains("no store attached")); - let err = EngineError::InvariantViolation { - what: "decode before prefill".into(), - }; - assert!(err.to_string().contains("decode before prefill")); - let err = EngineError::BackendFailure { - details: "kernel returned None".into(), - }; - assert!(err.to_string().contains("kernel returned None")); - } - - #[test] - fn engine_error_empty_prompt_and_backend_unavailable_render() { - assert_eq!( - EngineError::EmptyPrompt.to_string(), - "engine called with empty prompt" - ); - assert!(EngineError::BackendUnavailable - .to_string() - .contains("does not support")); - } - - // ─── RetrievalEngine + AnyEngine ────────────────────────────────────────── - - struct StubRetrievalEngine { - prefill_calls: usize, - decode_calls: usize, - last_token: Option, - } - - impl RetrievalEngine for StubRetrievalEngine { - fn name(&self) -> &str { - "stub-retrieval" - } - fn info(&self) -> EngineInfo { - EngineInfo { - name: self.name().into(), - description: "test fixture".into(), - backend: "cpu".into(), - config: String::new(), - } - } - fn prefill( - &mut self, - _weights: &ModelWeights, - token_ids: &[u32], - ) -> Result, EngineError> { - self.prefill_calls += 1; - if token_ids.is_empty() { - return Err(EngineError::EmptyPrompt); - } - Ok(Array2::zeros((1, 4))) - } - fn decode_step( - &mut self, - _weights: &ModelWeights, - token_id: u32, - ) -> Result, EngineError> { - self.decode_calls += 1; - self.last_token = Some(token_id); - Ok(Array2::zeros((1, 4))) - } - fn memory_bytes(&self) -> usize { - 16 - } - } - - #[test] - fn retrieval_engine_propagates_empty_prompt_error() { - let weights = crate::test_utils::make_test_weights(); - let mut engine = StubRetrievalEngine { - prefill_calls: 0, - decode_calls: 0, - last_token: None, - }; - let err = engine.prefill(&weights, &[]).unwrap_err(); - assert_eq!(err, EngineError::EmptyPrompt); - } - - #[test] - fn retrieval_engine_defaults_zero_window_and_cold_bytes() { - let engine = StubRetrievalEngine { - prefill_calls: 0, - decode_calls: 0, - last_token: None, - }; - assert_eq!(engine.window_tokens(), 0); - assert_eq!(engine.cold_bytes(), 0); - assert!(engine.stage_summary().is_none()); - } - - #[test] - fn any_engine_delegates_uniform_methods_to_inner() { - let kv: Box = Box::new(DefaultsOnlyEngine { - prefill_calls: 0, - decode_calls: 0, - }); - let any = AnyEngine::Kv(kv); - assert!(any.is_kv()); - assert!(!any.is_retrieval()); - assert_eq!(any.name(), "defaults-only"); - assert_eq!(any.memory_bytes(), 0); - assert_eq!(any.window_tokens(), 0); - assert_eq!(any.cold_bytes(), 0); - assert!(any.stage_summary().is_none()); - let info = any.info(); - assert_eq!(info.name, "defaults-only"); - - let retrieval: Box = Box::new(StubRetrievalEngine { - prefill_calls: 0, - decode_calls: 0, - last_token: None, - }); - let any = AnyEngine::Retrieval(retrieval); - assert!(any.is_retrieval()); - assert!(!any.is_kv()); - assert_eq!(any.name(), "stub-retrieval"); - assert_eq!(any.memory_bytes(), 16); - let info = any.info(); - assert_eq!(info.name, "stub-retrieval"); - } -} diff --git a/crates/larql-inference/src/kv_engine/any.rs b/crates/larql-inference/src/kv_engine/any.rs new file mode 100644 index 000000000..9416a2590 --- /dev/null +++ b/crates/larql-inference/src/kv_engine/any.rs @@ -0,0 +1,310 @@ +//! [`AnyEngine`] — the sum type that lets one call site drive either engine +//! family. + +use super::{DecodeStageSummary, EngineError, EngineInfo, KvEngine, RetrievalEngine}; +use crate::ffn::FfnBackend; +use crate::ModelWeights; +use ndarray::Array2; + +/// Sum type that holds either a [`KvEngine`] or a [`RetrievalEngine`]. +/// +/// Construction sites (the engine builder in +/// `larql-kv::EngineBuilder::build`, the bench / accuracy harnesses) +/// parse a spec into one or the other; the autoregressive loop calls +/// uniform `prefill` / `decode_step` (et al.) methods on `AnyEngine`, +/// which pattern-match internally on the variant. +/// +/// Each forwarding method takes the superset of arguments from both +/// trait surfaces. For [`RetrievalEngine`] engines (Apollo, future +/// Mode 5) the FFN-routing and compute-backend arguments are simply +/// ignored — `RetrievalEngine` runs its forward through +/// `forward_from_layer` / `forward_raw_logits` and doesn't need them. +/// This is intentional: keeping the harness call site uniform across +/// new engine families is more important than enforcing argument +/// minimality at the type level. Variant-specific behaviour still +/// surfaces through the [`EngineError`] enum (e.g. `RetrievalMiss` +/// only arrives from retrieval engines). +pub enum AnyEngine { + Kv(Box), + Retrieval(Box), +} + +impl AnyEngine { + pub fn name(&self) -> &str { + match self { + Self::Kv(e) => e.name(), + Self::Retrieval(e) => e.name(), + } + } + + pub fn info(&self) -> EngineInfo { + match self { + Self::Kv(e) => e.info(), + Self::Retrieval(e) => e.info(), + } + } + + pub fn memory_bytes(&self) -> usize { + match self { + Self::Kv(e) => e.memory_bytes(), + Self::Retrieval(e) => e.memory_bytes(), + } + } + + pub fn window_tokens(&self) -> usize { + match self { + Self::Kv(e) => e.window_tokens(), + Self::Retrieval(e) => e.window_tokens(), + } + } + + pub fn cold_bytes(&self) -> usize { + match self { + Self::Kv(e) => e.cold_bytes(), + Self::Retrieval(e) => e.cold_bytes(), + } + } + + pub fn stage_summary(&self) -> Option { + match self { + Self::Kv(e) => e.stage_summary(), + Self::Retrieval(e) => e.stage_summary(), + } + } + + /// Dispatch shape the current sequence committed to. Retrieval + /// engines have no coarse/per-layer split — they re-forward — so + /// they report `None`. + pub fn dispatch_path(&self) -> Option { + match self { + Self::Kv(e) => e.dispatch_path(), + Self::Retrieval(_) => None, + } + } + + pub fn is_kv(&self) -> bool { + matches!(self, Self::Kv(_)) + } + + pub fn is_retrieval(&self) -> bool { + matches!(self, Self::Retrieval(_)) + } + + // ── Forwarding methods (variant-specific dispatch) ────────────────────── + + /// Prefill. KvEngine variants consult `ffn`; retrieval variants + /// ignore it. + pub fn prefill( + &mut self, + weights: &ModelWeights, + ffn: &dyn FfnBackend, + token_ids: &[u32], + ) -> Result, EngineError> { + match self { + Self::Kv(e) => e.prefill(weights, ffn, token_ids), + Self::Retrieval(e) => e.prefill(weights, token_ids), + } + } + + /// One autoregressive decode step. Same routing as [`prefill`]. + pub fn decode_step( + &mut self, + weights: &ModelWeights, + ffn: &dyn FfnBackend, + token_id: u32, + ) -> Result, EngineError> { + match self { + Self::Kv(e) => e.decode_step(weights, ffn, token_id), + Self::Retrieval(e) => e.decode_step(weights, token_id), + } + } + + /// Capability forwarder for multi-modal input — see ADR-0023. + /// `Retrieval` variants are text-only by construction and always + /// return `false`; `Kv` variants delegate to the trait method. + pub fn supports_multimodal(&self) -> bool { + match self { + Self::Kv(e) => e.supports_multimodal(), + Self::Retrieval(_) => false, + } + } + + /// MM prefill forwarder. Only `Kv` variants can implement this; + /// `Retrieval` variants panic when called (callers MUST check + /// `supports_multimodal()` first, per ADR-0023). + pub fn prefill_from_hidden( + &mut self, + weights: &ModelWeights, + ffn: &dyn FfnBackend, + initial_hidden: &Array2, + ) -> Result, EngineError> { + match self { + Self::Kv(e) => e.prefill_from_hidden(weights, ffn, initial_hidden), + Self::Retrieval(_) => panic!( + "AnyEngine::Retrieval does not support prefill_from_hidden — \ + check supports_multimodal() before calling" + ), + } + } + + /// Prefill against a quantised vindex. KvEngine variants take a + /// `ComputeBackend` for kernel routing; retrieval variants ignore + /// it (they dequantise + run on f32 internally). + pub fn prefill_quant( + &mut self, + weights: &ModelWeights, + ffn: &dyn FfnBackend, + index: &larql_vindex::VectorIndex, + token_ids: &[u32], + backend: &dyn larql_compute::ComputeBackend, + ) -> Result, EngineError> { + match self { + Self::Kv(e) => e.prefill_quant(weights, ffn, index, token_ids, backend), + Self::Retrieval(e) => e.prefill_quant(weights, index, token_ids), + } + } + + /// One decode step against a quantised vindex. Same routing as + /// [`prefill_quant`]. + pub fn decode_step_quant( + &mut self, + weights: &ModelWeights, + ffn: &dyn FfnBackend, + index: &larql_vindex::VectorIndex, + token_id: u32, + backend: &dyn larql_compute::ComputeBackend, + ) -> Result, EngineError> { + match self { + Self::Kv(e) => e.decode_step_quant(weights, ffn, index, token_id, backend), + Self::Retrieval(e) => e.decode_step_quant(weights, index, token_id), + } + } + + /// Resident-weights quant prefill (`&weights`, threads `index`). See + /// [`KvEngine::prefill_resident`]. Retrieval variants fall back to their + /// f32 prefill (index-aware retrieval isn't a moe-shards path). + pub fn prefill_resident( + &mut self, + weights: &ModelWeights, + ffn: &dyn FfnBackend, + index: &larql_vindex::VectorIndex, + token_ids: &[u32], + ) -> Result, EngineError> { + match self { + Self::Kv(e) => e.prefill_resident(weights, ffn, index, token_ids), + Self::Retrieval(e) => e.prefill(weights, token_ids), + } + } + + /// Resident-weights quant decode step (`&weights`, threads `index`). See + /// [`KvEngine::decode_step_resident`]. + pub fn decode_step_resident( + &mut self, + weights: &ModelWeights, + ffn: &dyn FfnBackend, + index: &larql_vindex::VectorIndex, + token_id: u32, + ) -> Result, EngineError> { + match self { + Self::Kv(e) => e.decode_step_resident(weights, ffn, index, token_id), + Self::Retrieval(e) => e.decode_step(weights, token_id), + } + } + + /// Prefill via a caller-supplied [`crate::layer_executor::LayerExecutor`]. + /// Falls back to [`prefill_quant`](Self::prefill_quant) for + /// [`RetrievalEngine`] variants (which don't drive per-layer + /// executor loops). + pub fn prefill_quant_via_executor( + &mut self, + weights: &ModelWeights, + executor: &dyn crate::layer_executor::LayerExecutor, + ffn: &dyn FfnBackend, + index: &larql_vindex::VectorIndex, + token_ids: &[u32], + ) -> Result, EngineError> { + match self { + Self::Kv(e) => e.prefill_quant_via_executor(weights, executor, ffn, index, token_ids), + Self::Retrieval(e) => e.prefill_quant(weights, index, token_ids), + } + } + + /// One decode step via a caller-supplied `LayerExecutor`. Same + /// fall-back semantics as [`prefill_quant_via_executor`]. + pub fn decode_step_quant_via_executor( + &mut self, + weights: &ModelWeights, + executor: &dyn crate::layer_executor::LayerExecutor, + ffn: &dyn FfnBackend, + index: &larql_vindex::VectorIndex, + token_id: u32, + ) -> Result, EngineError> { + match self { + Self::Kv(e) => { + e.decode_step_quant_via_executor(weights, executor, ffn, index, token_id) + } + Self::Retrieval(e) => e.decode_step_quant(weights, index, token_id), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::kv_engine::test_stubs::{DefaultsOnlyEngine, StubRetrievalEngine}; + + #[test] + fn dispatch_path_forwards_for_kv_and_is_none_for_retrieval() { + // KV engines answer from their own recorded shape (the default + // trait impl is `None` until a prefill picks one). Retrieval + // engines re-forward instead of holding a K/V cache, so they have + // no coarse/per-layer split to report and must not invent one. + let kv = AnyEngine::Kv(Box::new(DefaultsOnlyEngine { + prefill_calls: 0, + decode_calls: 0, + })); + assert_eq!(kv.dispatch_path(), None); + + let retrieval = AnyEngine::Retrieval(Box::new(StubRetrievalEngine { + prefill_calls: 0, + decode_calls: 0, + last_token: None, + })); + assert_eq!( + retrieval.dispatch_path(), + None, + "a retrieval engine has no dispatch shape to report" + ); + } + + #[test] + fn any_engine_delegates_uniform_methods_to_inner() { + let kv: Box = Box::new(DefaultsOnlyEngine { + prefill_calls: 0, + decode_calls: 0, + }); + let any = AnyEngine::Kv(kv); + assert!(any.is_kv()); + assert!(!any.is_retrieval()); + assert_eq!(any.name(), "defaults-only"); + assert_eq!(any.memory_bytes(), 0); + assert_eq!(any.window_tokens(), 0); + assert_eq!(any.cold_bytes(), 0); + assert!(any.stage_summary().is_none()); + let info = any.info(); + assert_eq!(info.name, "defaults-only"); + + let retrieval: Box = Box::new(StubRetrievalEngine { + prefill_calls: 0, + decode_calls: 0, + last_token: None, + }); + let any = AnyEngine::Retrieval(retrieval); + assert!(any.is_retrieval()); + assert!(!any.is_kv()); + assert_eq!(any.name(), "stub-retrieval"); + assert_eq!(any.memory_bytes(), 16); + let info = any.info(); + assert_eq!(info.name, "stub-retrieval"); + } +} diff --git a/crates/larql-inference/src/kv_engine/dispatch_path.rs b/crates/larql-inference/src/kv_engine/dispatch_path.rs new file mode 100644 index 000000000..2cc192408 --- /dev/null +++ b/crates/larql-inference/src/kv_engine/dispatch_path.rs @@ -0,0 +1,86 @@ +//! Which dispatch shape an engine actually took for the current sequence. +//! +//! An engine's *name* does not determine how its forward runs. The same +//! `standard` engine takes a fused whole-model kernel when it is +//! unwindowed, and a generic per-layer loop when it is not — and on +//! `MetalBackend` the per-layer loop is delegated to the host, so a row +//! labelled `[metal (GPU)]` can be doing all of its attention and FFN on +//! the CPU. That difference is worth ~15 ms/token on qwen3-0.6b and is +//! invisible in every other field an engine reports. +//! +//! Engines record the shape when they choose it (at prefill, since decode +//! must follow the recorded mode) and report it here so a benchmark table +//! can say which one it measured instead of leaving the reader to infer +//! it from the timings. + +/// The dispatch shape a prefill committed the sequence to. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DispatchPath { + /// One whole-model handle through the backend's fused pipeline + /// (`coarse_prefill` / `coarse_decode_step`). The K/V lives in the + /// backend, not in the engine, and the engine's own state policy is + /// not engaged. + Coarse, + /// The generic per-layer loop: one handle per layer, attention and + /// FFN dispatched a layer at a time. Every engine-specific K/V + /// policy runs here. + PerLayer, +} + +impl DispatchPath { + /// Short tag for diagnostics and bench rows. + pub fn as_str(self) -> &'static str { + match self { + DispatchPath::Coarse => "coarse", + DispatchPath::PerLayer => "per-layer", + } + } + + /// Render with where the compute lands. `host_delegated` is the + /// backend's answer to [`crate::kv_dispatch::KvDispatch::per_layer_is_host_delegated`] + /// — true when this backend implements the per-layer surface by + /// forwarding to the CPU, which makes a "GPU" row a CPU measurement. + pub fn describe(self, host_delegated: bool) -> &'static str { + match (self, host_delegated) { + (DispatchPath::Coarse, _) => "coarse", + (DispatchPath::PerLayer, false) => "per-layer", + (DispatchPath::PerLayer, true) => "per-layer→host", + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn as_str_names_both_shapes() { + assert_eq!(DispatchPath::Coarse.as_str(), "coarse"); + assert_eq!(DispatchPath::PerLayer.as_str(), "per-layer"); + } + + #[test] + fn describe_flags_host_delegation_only_for_per_layer() { + // The whole point: a per-layer path on a host-delegating backend + // must be visibly distinct from one running native kernels. + assert_eq!(DispatchPath::PerLayer.describe(true), "per-layer→host"); + assert_eq!(DispatchPath::PerLayer.describe(false), "per-layer"); + } + + #[test] + fn describe_ignores_host_delegation_for_coarse() { + // Coarse never touches the per-layer surface, so the backend's + // delegation answer is irrelevant to it. + assert_eq!(DispatchPath::Coarse.describe(true), "coarse"); + assert_eq!(DispatchPath::Coarse.describe(false), "coarse"); + } + + #[test] + fn dispatch_path_is_comparable_and_copy() { + let a = DispatchPath::Coarse; + let b = a; + assert_eq!(a, b); + assert_ne!(DispatchPath::Coarse, DispatchPath::PerLayer); + assert_eq!(format!("{:?}", DispatchPath::PerLayer), "PerLayer"); + } +} diff --git a/crates/larql-inference/src/kv_engine/error.rs b/crates/larql-inference/src/kv_engine/error.rs new file mode 100644 index 000000000..a1917bad6 --- /dev/null +++ b/crates/larql-inference/src/kv_engine/error.rs @@ -0,0 +1,360 @@ +//! [`EngineError`] — why a prefill or decode step did not produce a hidden +//! state, split along the axis a caller must act on. + +use thiserror::Error; + +/// Typed failure mode for engine `prefill` / `decode_step` calls. +/// +/// Replaces the historical `Option` return semantics that collapsed +/// "empty prompt", "backend doesn't support this", "retrieval miss", +/// "engine invariant violated" and "backend operation failed" into a +/// single opaque `None`. Two consumers (the accuracy harness and the +/// bench harness) used to route that `None` incompatibly — the +/// accuracy runner silently dropped the row via `filter_map` while the +/// bench aborted with `"engine prefill failed"`. This taxonomy lets +/// both routes branch on error *kind*; see `docs/state-policy.md`. +/// +/// The variants split error reasons along their alerting axis: +/// +/// - [`EmptyPrompt`](Self::EmptyPrompt) — caller-side input bug; +/// surfaces in CLI validation rather than a runtime alert. +/// - [`BackendUnavailable`](Self::BackendUnavailable) — the engine's +/// backend does not implement the requested capability (e.g. a +/// Metal kernel that hasn't been ported, an asked-for Q4K matvec on +/// a CPU build without BLAS). Falls back to a different code path +/// *if* one exists; otherwise surfaces as a configuration error. +/// - [`RetrievalMiss { reason }`](Self::RetrievalMiss) — a retrieval +/// engine (Apollo, future Mode 5) could not serve this query against +/// its store. Expected, recoverable; surfaces in the harness as a +/// `served_rate < 1.0` column rather than an alert. +/// - [`InvariantViolation { what }`](Self::InvariantViolation) — the +/// engine was driven outside its state-machine contract (e.g. +/// `decode_step` called before `prefill`). Indicates a harness-level +/// dispatch bug; production observability should alert immediately. +/// - [`BackendFailure { details }`](Self::BackendFailure) — the inner +/// backend or compute kernel returned a runtime failure. Indicates a +/// data condition or environmental issue (corrupt weights, OOM, GPU +/// driver error); production observability should log + investigate +/// but not alert with the same urgency as `InvariantViolation`. +/// - [`Execution(refusal)`](Self::Execution) — a routed operation this +/// engine required did not execute, so the layer it belonged to is +/// incomplete. Distinct from `BackendFailure`: nothing *failed*, a +/// route declined, and [`RefusalKind`](larql_execution::RefusalKind) +/// says which of three responses it needs — fetch the operand, pick +/// another executor, or repair the artifact. +/// +/// `InvariantViolation` and `BackendFailure` are deliberately kept as +/// two top-level variants rather than collapsed into a single +/// `InternalError { kind }`. Sub-tagged enums lose the alert-routing +/// distinction when the consumer writes `match err { InternalError(_) => ... }`. +/// +/// The enum is **exhaustive** (no `#[non_exhaustive]`). New variants +/// are breaking changes on purpose — defaulting a new condition into +/// an existing arm reproduces the silent-drop problem one layer down. +/// +/// `Execution` carries the refusal itself rather than a flattened +/// `{ kind, message }` pair, so the classification and the concrete error +/// reach whoever handles it together. That costs the enum `Clone`, +/// `PartialEq` and `Eq` — a boxed trait object has none of them — which +/// is why callers match on shape rather than compare values. +#[derive(Debug, Error)] +pub enum EngineError { + #[error("engine called with empty prompt")] + EmptyPrompt, + #[error("backend does not support this operation")] + BackendUnavailable, + #[error("retrieval miss: {reason}")] + RetrievalMiss { reason: String }, + #[error("engine invariant violated: {what}")] + InvariantViolation { what: String }, + #[error("backend operation failed: {details}")] + BackendFailure { details: String }, + #[error("execution refused ({kind}): {refusal}", kind = .0.kind(), refusal = .0)] + Execution(larql_execution::BoxRefusal), + #[error( + "{cause} — and the engine could not undo a partial K/V mutation, \ + so this instance must be re-prefilled before it is driven again" + )] + StateInvalidated { + #[source] + cause: Box, + }, +} + +impl EngineError { + /// Whether a sweep may skip this row and keep driving **this engine**. + /// + /// This is the harness's question — skip + log + continue, versus abort + /// and investigate — and answering it needs *both* of the finer + /// questions below, because either one alone is a trap: + /// + /// ```text + /// operation_is_recoverable() could this operation ever succeed? + /// engine_state_is_retryable() is this engine instance still usable? + /// ``` + /// + /// A `Residency` refusal that invalidated the cache is recoverable in + /// the first sense and catastrophic in the second: a caller told only + /// "recoverable" would fetch the missing operand and drive the same + /// engine again, appending the token twice. So this conjoins them. + pub fn is_recoverable(&self) -> bool { + self.engine_state_is_retryable() && self.operation_is_recoverable() + } + + /// Whether the operation itself could succeed in some environment — + /// with more residency, or through another capable executor. + /// + /// Says **nothing** about the engine that produced the error. Use it to + /// decide whether a *cause* is worth retrying at all; use + /// [`Self::engine_state_is_retryable`] to decide whether this instance + /// is the thing to retry it on. + pub fn operation_is_recoverable(&self) -> bool { + match self { + Self::EmptyPrompt | Self::BackendUnavailable | Self::RetrievalMiss { .. } => true, + Self::InvariantViolation { .. } | Self::BackendFailure { .. } => false, + // `Residency` and `Unsupported` may succeed given more residency + // or another executor; a `BindingDefect` indicts the artifact and + // must not be swept under a served-rate column. + Self::Execution(refusal) => refusal.kind().is_recoverable_without_rebinding(), + // The wrapper does not change what the cause could do — only + // where it can be attempted. + Self::StateInvalidated { cause } => cause.operation_is_recoverable(), + } + } + + /// Whether the engine that produced this error can be driven again. + /// + /// False only for [`StateInvalidated`](Self::StateInvalidated), which is + /// raised exactly when a failure left a partially-applied decode step + /// that could not be rewound. Every other failure either mutated + /// nothing or was rolled back before the error was returned. + /// + /// An invalidated engine is not dead: `prefill` replaces the cache + /// wholesale and clears the condition. It is only unsafe to *continue* + /// from. + pub fn engine_state_is_retryable(&self) -> bool { + !matches!(self, Self::StateInvalidated { .. }) + } + + /// The refusal's classification, when this error is one. + /// + /// Sees through [`StateInvalidated`](Self::StateInvalidated), so wrapping + /// a refusal does not cost it the classification an operator acts on. + pub fn refusal_kind(&self) -> Option { + match self { + Self::Execution(refusal) => Some(refusal.kind()), + Self::StateInvalidated { cause } => cause.refusal_kind(), + _ => None, + } + } + + /// Wrap this error as one that also invalidated its engine. + /// + /// Named as a constructor so the two facts are recorded together at the + /// one place that knows both — a rollback that did not happen, and the + /// failure that made it necessary. + pub fn invalidating_engine_state(self) -> Self { + Self::StateInvalidated { + cause: Box::new(self), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn engine_error_is_recoverable_classifies_variants() { + assert!(EngineError::EmptyPrompt.is_recoverable()); + assert!(EngineError::BackendUnavailable.is_recoverable()); + assert!(EngineError::RetrievalMiss { + reason: "no store".into() + } + .is_recoverable()); + assert!(!EngineError::InvariantViolation { + what: "decode before prefill".into() + } + .is_recoverable()); + assert!(!EngineError::BackendFailure { + details: "kernel oom".into() + } + .is_recoverable()); + } + + #[test] + fn engine_error_display_includes_reason_payload() { + let err = EngineError::RetrievalMiss { + reason: "no store attached".into(), + }; + assert!(err.to_string().contains("no store attached")); + let err = EngineError::InvariantViolation { + what: "decode before prefill".into(), + }; + assert!(err.to_string().contains("decode before prefill")); + let err = EngineError::BackendFailure { + details: "kernel returned None".into(), + }; + assert!(err.to_string().contains("kernel returned None")); + } + + #[test] + fn engine_error_empty_prompt_and_backend_unavailable_render() { + assert_eq!( + EngineError::EmptyPrompt.to_string(), + "engine called with empty prompt" + ); + assert!(EngineError::BackendUnavailable + .to_string() + .contains("does not support")); + } + + // ── The Execution variant ─────────────────────────────────────────────── + // + // Where the dispatch ring's refusal channel terminates. Its behaviour is + // entirely delegation — to the refusal's own `kind()` — so what is worth + // pinning is that the delegation happens, rather than that a second, + // drifting classification was written here. + + use larql_execution::RefusalKind; + + const REFUSAL_LAYER: usize = 3; + const REFUSAL_MESSAGE: &str = "expert 7 is not resident"; + + fn refusal_of(kind: RefusalKind) -> larql_execution::BoxRefusal { + Box::new(crate::ffn::RecordedRefusal { + layer: REFUSAL_LAYER, + kind, + message: REFUSAL_MESSAGE.into(), + }) + } + + #[test] + fn execution_renders_its_kind_and_the_routes_own_words() { + // Both levels must reach the operator: the category to act on, and the + // concrete detail that says what to act on. + let rendered = EngineError::Execution(refusal_of(RefusalKind::Residency)).to_string(); + assert!(rendered.contains("residency"), "{rendered}"); + assert!(rendered.contains(REFUSAL_MESSAGE), "{rendered}"); + assert!( + rendered.contains(&format!("layer {REFUSAL_LAYER}")), + "{rendered}" + ); + } + + #[test] + fn execution_reports_the_refusals_own_kind() { + for kind in RefusalKind::ALL { + assert_eq!( + EngineError::Execution(refusal_of(kind)).refusal_kind(), + Some(kind) + ); + } + } + + #[test] + fn only_an_execution_error_has_a_refusal_kind() { + // A caller branching on `refusal_kind()` must not be handed a + // classification for an error that never refused anything. + for err in [ + EngineError::EmptyPrompt, + EngineError::BackendUnavailable, + EngineError::RetrievalMiss { + reason: "no store".into(), + }, + EngineError::InvariantViolation { + what: "decode before prefill".into(), + }, + EngineError::BackendFailure { + details: "kernel oom".into(), + }, + ] { + assert!(err.refusal_kind().is_none(), "{err:?}"); + } + } + + // ── StateInvalidated ──────────────────────────────────────────────────── + // + // The wrapper that separates "this operation could succeed somewhere" + // from "this engine can be asked again". Before it, `is_recoverable` + // answered the first while callers read it as the second. + + #[test] + fn invalidating_wraps_the_cause_without_hiding_it() { + let err = + EngineError::Execution(refusal_of(RefusalKind::Residency)).invalidating_engine_state(); + assert!(matches!(err, EngineError::StateInvalidated { .. })); + // The classification and the route's words still reach the operator. + assert_eq!(err.refusal_kind(), Some(RefusalKind::Residency)); + let rendered = err.to_string(); + assert!(rendered.contains(REFUSAL_MESSAGE), "{rendered}"); + assert!(rendered.contains("re-prefilled"), "{rendered}"); + // …and so does the cause itself, through the error source chain. + let source = std::error::Error::source(&err).expect("cause must be the source"); + assert!(source.to_string().contains("residency")); + } + + #[test] + fn an_invalidated_engine_is_never_retryable_however_recoverable_the_cause() { + // The exact trap: a Residency refusal is recoverable as an operation + // and catastrophic as a retry target. Both facts must be readable, + // and the harness-facing question must answer with the pessimistic one. + let err = + EngineError::Execution(refusal_of(RefusalKind::Residency)).invalidating_engine_state(); + assert!( + err.operation_is_recoverable(), + "wrapping must not claim the operation itself is hopeless" + ); + assert!(!err.engine_state_is_retryable()); + assert!( + !err.is_recoverable(), + "a caller told 'recoverable' would fix the residency and re-drive a \ + dead engine, appending the token twice" + ); + } + + #[test] + fn invalidation_composes_with_any_cause() { + // Not refusal-specific: a declining backend leaves the same + // half-applied decode step, so it wraps the same way. + let err = EngineError::BackendFailure { + details: "dispatch returned None".into(), + } + .invalidating_engine_state(); + assert!(!err.engine_state_is_retryable()); + assert!(!err.operation_is_recoverable()); + assert!(err.refusal_kind().is_none(), "not every cause is a refusal"); + assert!(err.to_string().contains("dispatch returned None")); + } + + #[test] + fn every_other_variant_reports_a_retryable_engine() { + // The flag means one specific thing — an un-rewound mutation — and + // must not creep into meaning "something went wrong". + for err in [ + EngineError::EmptyPrompt, + EngineError::BackendUnavailable, + EngineError::BackendFailure { + details: "x".into(), + }, + EngineError::InvariantViolation { what: "x".into() }, + EngineError::Execution(refusal_of(RefusalKind::BindingDefect)), + ] { + assert!(err.engine_state_is_retryable(), "{err:?}"); + } + } + + #[test] + fn execution_recoverability_delegates_to_the_refusal() { + // The reason `is_recoverable` stopped being a `matches!`: a + // BindingDefect must not be swept as a coverage deficit, and a + // Residency miss must not abort a sweep. + for kind in RefusalKind::ALL { + assert_eq!( + EngineError::Execution(refusal_of(kind)).is_recoverable(), + kind.is_recoverable_without_rebinding(), + "{kind} must follow the refusal vocabulary, not a second opinion" + ); + } + } +} diff --git a/crates/larql-inference/src/kv_engine/info.rs b/crates/larql-inference/src/kv_engine/info.rs new file mode 100644 index 000000000..17969b414 --- /dev/null +++ b/crates/larql-inference/src/kv_engine/info.rs @@ -0,0 +1,58 @@ +//! [`EngineInfo`] — the diagnostics an engine reports about itself. + +/// Runtime diagnostics reported by each engine. +#[derive(Debug, Clone)] +pub struct EngineInfo { + /// Short engine name (e.g. `"markov-rs"`). + pub name: String, + /// Human-readable description of the engine's state management strategy. + pub description: String, + /// Hardware backend name from [`larql_compute::ComputeBackend::name`]: `"cpu"`, `"metal"`, etc. + pub backend: String, + /// Key config parameters (e.g. `"window=512"`), empty string if unconfigured. + pub config: String, +} + +impl EngineInfo { + pub fn summary(&self) -> String { + if self.config.is_empty() { + format!("{} [{}] {}", self.name, self.backend, self.description) + } else { + format!( + "{} [{}] ({}) {}", + self.name, self.backend, self.config, self.description + ) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn engine_info_summary_with_config() { + let info = EngineInfo { + name: "markov-rs".into(), + description: "residual KV".into(), + backend: "cpu".into(), + config: "window=512".into(), + }; + let s = info.summary(); + assert!(s.contains("markov-rs")); + assert!(s.contains("cpu")); + assert!(s.contains("window=512")); + } + + #[test] + fn engine_info_summary_no_config() { + let info = EngineInfo { + name: "test".into(), + description: "desc".into(), + backend: "metal".into(), + config: String::new(), + }; + let s = info.summary(); + assert!(!s.contains("()")); + } +} diff --git a/crates/larql-inference/src/kv_engine/kv.rs b/crates/larql-inference/src/kv_engine/kv.rs new file mode 100644 index 000000000..98a7153ff --- /dev/null +++ b/crates/larql-inference/src/kv_engine/kv.rs @@ -0,0 +1,355 @@ +//! [`KvEngine`] — the interface the autoregressive decode loop dispatches +//! against. + +use super::{DecodeStageSummary, DispatchPath, EngineError, EngineInfo}; +use crate::ffn::FfnBackend; +use crate::ModelWeights; +use ndarray::Array2; + +/// Common interface shared by all KV-cache engines. +pub trait KvEngine: Send { + fn name(&self) -> &str; + + /// Runtime diagnostics: engine name, backend, config, description. + fn info(&self) -> EngineInfo; + + /// Run the prefill forward pass over all prompt tokens. + /// + /// `ffn` is the FFN backend the engine should dispatch through — + /// typically [`WeightFfn`](crate::ffn::WeightFfn) / + /// [`BackendFfn`](crate::ffn::BackendFfn) for local compute, or + /// [`RemoteWalkBackend`](crate::ffn::RemoteWalkBackend) for grid + /// routing. Engines that don't consult an FFN router (e.g. ones + /// that recompute FFN from `weights` directly) may ignore this + /// parameter. + /// + /// Returns the hidden state at the final token position (shape `[1, hidden_dim]`). + /// + /// Failure modes surface as typed [`EngineError`] variants — see + /// the enum's docs for the routing taxonomy. + fn prefill( + &mut self, + weights: &ModelWeights, + ffn: &dyn FfnBackend, + token_ids: &[u32], + ) -> Result, EngineError>; + + /// Run one autoregressive decode step for a single new token. + /// Returns the hidden state (shape `[1, hidden_dim]`). + fn decode_step( + &mut self, + weights: &ModelWeights, + ffn: &dyn FfnBackend, + token_id: u32, + ) -> Result, EngineError>; + + /// Static capability: does this engine accept pre-built hidden + /// state via [`prefill_from_hidden`]? Default `false`. + /// + /// The CLI MUST check this **before** running a (potentially + /// minutes-long) modal encoder, so the user gets a fast, clear + /// error if they paired `--image` with an engine that doesn't + /// support multi-modal input. See ADR-0023. + /// + /// The default-false return is deliberate debt — six of seven + /// engines inherit it for Phase 1d. The end state collapses + /// `prefill(token_ids)` into a thin wrapper over + /// `embed_tokens_pub` then `prefill_from_hidden` on every engine, + /// at which point this method becomes universally `true` and is + /// removed. Tracked in ADR-0023 (Default-false debt). + fn supports_multimodal(&self) -> bool { + false + } + + /// Prefill from a pre-built initial hidden state. Caller built it + /// via `larql_compute::forward::embed_plan` from an `EmbeddingPlan` + /// that may include `Precomputed` rows (vision / audio embeddings). + /// + /// Same contract as [`prefill`]: runs forward through every layer, + /// populates the engine's KV cache, returns the final-token hidden + /// state. Returns the same `Result<_, EngineError>` shape as + /// `prefill` for uniform call-site error handling. The engine's + /// internal absolute position pointer must be set from + /// `initial_hidden.nrows()`, NOT from any token count — the input + /// may contain non-token positions. + /// + /// Default impl panics (not an `Err` return) on engines that don't + /// override it. Callers MUST check [`supports_multimodal`] first; + /// the panic is defense-in-depth against bypass, not a substitute + /// for the capability check. + fn prefill_from_hidden( + &mut self, + _weights: &ModelWeights, + _ffn: &dyn FfnBackend, + _initial_hidden: &Array2, + ) -> Result, EngineError> { + panic!( + "engine {:?} does not support multi-modal input; \ + check supports_multimodal() before calling prefill_from_hidden", + self.name() + ); + } + + /// Bytes of persistent engine state (excludes model weights). + fn memory_bytes(&self) -> usize; + + /// Token count in the active hot window (varies by engine type). + fn window_tokens(&self) -> usize { + 0 + } + + /// Cold-tier bytes (residuals or token IDs past the hot window). + fn cold_bytes(&self) -> usize { + 0 + } + + /// Per-stage timing summary. Returns `None` if profiling was not enabled. + fn stage_summary(&self) -> Option { + None + } + + /// Which dispatch shape this engine committed the current sequence + /// to — see [`DispatchPath`]. `None` before prefill, and for engines + /// that only ever have one shape. + /// + /// Reported so diagnostics can distinguish a fused whole-model run + /// from a per-layer one. The two differ by more than a constant: on + /// a backend that delegates the per-layer surface to the host, the + /// per-layer shape runs attention and FFN on the CPU regardless of + /// which backend the caller selected. + fn dispatch_path(&self) -> Option { + None + } + + /// Prefill using Q4K quantised weights from `index` and `backend`. + /// + /// When the backend supports the fused Q4 pipeline (Metal), this routes + /// through `backend.prefill_kquant` for full GPU speed. Falls back to the + /// f32 path when `backend.supports_quant(::larql_compute::QuantFormat::Q4_K) == false` or `index` has no Q4K data. + /// + /// `weights` is `&ModelWeights` (immutable): the engine dequantises f32 + /// attention tensors into its own `dequant_scratch` on the first call and + /// resolves them via `WeightsView::with_scratch` (one-time cost; subsequent + /// decode steps reuse the engine-owned scratch). + fn prefill_quant( + &mut self, + weights: &ModelWeights, + ffn: &dyn FfnBackend, + index: &larql_vindex::VectorIndex, + token_ids: &[u32], + backend: &dyn larql_compute::ComputeBackend, + ) -> Result, EngineError> { + let _ = (index, backend); + self.prefill(weights, ffn, token_ids) // default: f32 fallback + } + + /// One autoregressive decode step using Q4K weights. + /// + /// Same routing semantics as [`prefill_quant`]: Metal via `decode_token` + /// when available, f32 fallback otherwise. + fn decode_step_quant( + &mut self, + weights: &ModelWeights, + ffn: &dyn FfnBackend, + index: &larql_vindex::VectorIndex, + token_id: u32, + backend: &dyn larql_compute::ComputeBackend, + ) -> Result, EngineError> { + let _ = (index, backend); + self.decode_step(weights, ffn, token_id) // default: f32 fallback + } + + /// Resident-weights quant prefill. Unlike [`prefill_quant`] (which + /// dequantises attn into the engine's `dequant_scratch`), this assumes the + /// **caller has already made the client weights f32-resident** — so it + /// merely threads `index` to the backend. That lets a Q4K-direct attention + /// kernel (`LARQL_Q4K_DIRECT_ATTN`) read packed bytes from the index while + /// the FFN backend borrows the same `&weights` (task #16). Default: f32 + /// fallback (index ignored). + fn prefill_resident( + &mut self, + weights: &ModelWeights, + ffn: &dyn FfnBackend, + index: &larql_vindex::VectorIndex, + token_ids: &[u32], + ) -> Result, EngineError> { + let _ = index; + self.prefill(weights, ffn, token_ids) + } + + /// One decode step against resident (pre-dequantised) weights, threading + /// `index` to the backend. Sibling of [`prefill_resident`]; same rationale. + /// Default: f32 fallback (index ignored). + fn decode_step_resident( + &mut self, + weights: &ModelWeights, + ffn: &dyn FfnBackend, + index: &larql_vindex::VectorIndex, + token_id: u32, + ) -> Result, EngineError> { + let _ = index; + self.decode_step(weights, ffn, token_id) + } + + /// Prefill via a caller-supplied `LayerExecutor` (dense/f32 path). + /// See [`docs/specs/engine-state-vs-execution.md`]. + /// + /// Sibling of [`prefill_quant_via_executor`] for engines that + /// don't have a quant path (no vindex needed). Default impl falls + /// through to [`prefill`]. + fn prefill_via_executor( + &mut self, + weights: &ModelWeights, + executor: &dyn crate::layer_executor::LayerExecutor, + ffn: &dyn FfnBackend, + token_ids: &[u32], + ) -> Result, EngineError> { + let _ = executor; + self.prefill(weights, ffn, token_ids) + } + + /// One decode step via a caller-supplied `LayerExecutor` (dense/f32). + /// Sibling of [`decode_step_quant_via_executor`]. + fn decode_step_via_executor( + &mut self, + weights: &ModelWeights, + executor: &dyn crate::layer_executor::LayerExecutor, + ffn: &dyn FfnBackend, + token_id: u32, + ) -> Result, EngineError> { + let _ = executor; + self.decode_step(weights, ffn, token_id) + } + + /// Prefill via a caller-supplied `LayerExecutor`. See + /// [`docs/specs/engine-state-vs-execution.md`]. + /// + /// The default impl falls through to [`prefill_quant`] using + /// `executor.backend()` — engines that haven't migrated yet keep + /// working unchanged. Migrated engines override this method to + /// drive the layer loop through the executor and honor the FFN + /// parameter properly. + fn prefill_quant_via_executor( + &mut self, + weights: &ModelWeights, + executor: &dyn crate::layer_executor::LayerExecutor, + ffn: &dyn FfnBackend, + index: &larql_vindex::VectorIndex, + token_ids: &[u32], + ) -> Result, EngineError> { + self.prefill_quant(weights, ffn, index, token_ids, executor.backend()) + } + + /// One decode step via a caller-supplied `LayerExecutor`. See + /// [`prefill_quant_via_executor`] for the migration contract. + fn decode_step_quant_via_executor( + &mut self, + weights: &ModelWeights, + executor: &dyn crate::layer_executor::LayerExecutor, + ffn: &dyn FfnBackend, + index: &larql_vindex::VectorIndex, + token_id: u32, + ) -> Result, EngineError> { + self.decode_step_quant(weights, ffn, index, token_id, executor.backend()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::kv_engine::test_stubs::DefaultsOnlyEngine; + + #[test] + fn defaults_window_tokens_and_cold_bytes_are_zero() { + let engine = DefaultsOnlyEngine { + prefill_calls: 0, + decode_calls: 0, + }; + assert_eq!(engine.window_tokens(), 0); + assert_eq!(engine.cold_bytes(), 0); + assert!(engine.stage_summary().is_none()); + assert_eq!(engine.name(), "defaults-only"); + } + + /// All four `*_via_executor` default impls dispatch through to their + /// non-executor sibling, which on `DefaultsOnlyEngine` falls back to + /// `prefill` / `decode_step`. Covers the function bodies of + /// `prefill_via_executor` (224-233), `decode_step_via_executor` + /// (237-246), `prefill_quant_via_executor` (256-265), + /// `decode_step_quant_via_executor` (269-278). + #[test] + fn defaults_via_executor_methods_dispatch_to_non_executor_siblings() { + struct StubExecutor { + backend: larql_compute::CpuBackend, + } + impl crate::layer_executor::LayerExecutor for StubExecutor { + fn backend(&self) -> &dyn larql_compute::ComputeBackend { + &self.backend + } + fn dispatch_kind(&self) -> crate::layer_executor::ExecutorDispatchKind { + crate::layer_executor::ExecutorDispatchKind::PerLayer + } + fn name(&self) -> &str { + "stub" + } + } + let exec = StubExecutor { + backend: larql_compute::CpuBackend, + }; + let weights = crate::test_utils::make_test_weights(); + let index = crate::test_utils::make_test_vindex(&weights); + let ffn = crate::ffn::WeightFfn { weights: &weights }; + let mut engine = DefaultsOnlyEngine { + prefill_calls: 0, + decode_calls: 0, + }; + + // prefill_via_executor → prefill + let out = engine.prefill_via_executor(&weights, &exec, &ffn, &[0, 1]); + assert!(out.is_ok()); + assert_eq!(engine.prefill_calls, 1); + + // decode_step_via_executor → decode_step + let out = engine.decode_step_via_executor(&weights, &exec, &ffn, 2); + assert!(out.is_ok()); + assert_eq!(engine.decode_calls, 1); + + // prefill_quant_via_executor → prefill_quant → prefill (default fallback) + let weights_q = crate::test_utils::make_test_weights(); + let out = engine.prefill_quant_via_executor(&weights_q, &exec, &ffn, &index, &[0, 1]); + assert!(out.is_ok()); + assert_eq!(engine.prefill_calls, 2); + + // decode_step_quant_via_executor → decode_step_quant → decode_step + let out = engine.decode_step_quant_via_executor(&weights_q, &exec, &ffn, &index, 3); + assert!(out.is_ok()); + assert_eq!(engine.decode_calls, 2); + } + + #[test] + fn defaults_q4k_methods_fall_back_to_f32() { + let weights = crate::test_utils::make_test_weights(); + let index = crate::test_utils::make_test_vindex(&weights); + let backend = larql_compute::cpu_backend(); + let ffn = crate::ffn::WeightFfn { weights: &weights }; + let mut engine = DefaultsOnlyEngine { + prefill_calls: 0, + decode_calls: 0, + }; + + let weights_q4k = crate::test_utils::make_test_weights(); + let out = engine.prefill_quant(&weights_q4k, &ffn, &index, &[1, 2, 3], &*backend); + assert!(out.is_ok()); + assert_eq!( + engine.prefill_calls, 1, + "default prefill_quant must dispatch to prefill" + ); + + let out = engine.decode_step_quant(&weights_q4k, &ffn, &index, 4, &*backend); + assert!(out.is_ok()); + assert_eq!( + engine.decode_calls, 1, + "default decode_step_quant must dispatch to decode_step" + ); + } +} diff --git a/crates/larql-inference/src/kv_engine/mod.rs b/crates/larql-inference/src/kv_engine/mod.rs new file mode 100644 index 000000000..a488b9a56 --- /dev/null +++ b/crates/larql-inference/src/kv_engine/mod.rs @@ -0,0 +1,40 @@ +//! KV-cache engine trait and shared types. +//! +//! Defines the abstract surface that the autoregressive decode loop +//! dispatches against. Concrete engine implementations (MarkovResidual, +//! WindowedCheckpoint, TurboQuant, Apollo, Standard, NoCache) live in +//! `larql-kv` and `impl larql_inference::KvEngine` against this trait. +//! +//! The trait deliberately lives in `larql-inference` rather than +//! `larql-kv` so the dispatch entry point (which lives here, in the +//! crate that owns the forward pass) can reference the trait without +//! a circular dependency on `larql-kv`. See +//! `docs/specs/kv-engine-unification.md` §10.4. +//! +//! Correctness contract: `prefill` and `decode_step` return the +//! pre-`lm_head` hidden state (shape `[1, hidden_dim]`). The caller +//! applies `final_norm + lm_head` to get logits — see +//! [`forward::hidden_to_raw_logits`](crate::forward::hidden_to_raw_logits). +//! +//! The surface is split by concept — the error taxonomy, the two engine +//! families, the sum type over them, and the two diagnostic structs — and +//! re-exported flat, so `larql_inference::kv_engine::EngineError` and its +//! siblings keep resolving as they did when this was one file. + +mod any; +mod dispatch_path; +mod error; +mod info; +mod kv; +mod retrieval; +mod stages; +#[cfg(test)] +mod test_stubs; + +pub use any::AnyEngine; +pub use dispatch_path::DispatchPath; +pub use error::EngineError; +pub use info::EngineInfo; +pub use kv::KvEngine; +pub use retrieval::RetrievalEngine; +pub use stages::DecodeStageSummary; diff --git a/crates/larql-inference/src/kv_engine/retrieval.rs b/crates/larql-inference/src/kv_engine/retrieval.rs new file mode 100644 index 000000000..bf43673d0 --- /dev/null +++ b/crates/larql-inference/src/kv_engine/retrieval.rs @@ -0,0 +1,139 @@ +//! [`RetrievalEngine`] — engines whose state is not an autoregressive K/V +//! cache. + +use super::{DecodeStageSummary, EngineError, EngineInfo}; +use crate::ModelWeights; +use ndarray::Array2; + +/// Engines whose state is **not** an autoregressive K/V cache. +/// +/// Sibling trait to [`KvEngine`]. Both surfaces share the +/// [`EngineInfo`] / [`EngineError`] / [`DecodeStageSummary`] vocabulary +/// but diverge on the per-step contract: +/// +/// | `KvEngine` | `RetrievalEngine` | +/// |-------------------------------------|----------------------------------------| +/// | per-token K/V append per layer | retrieval against a pre-built store | +/// | dispatches FFN through `FfnBackend` | does not consult an FFN router | +/// | state reconstructible to K/V tensors| state is residual delta + token list | +/// +/// Apollo (boundary-residual + injection delta) lands here; Mode 5 / +/// Graph-Grounded engines will too. The trait deliberately drops the +/// per-step K/V append assumption and the `FfnBackend` parameter that +/// `KvEngine::prefill` carries — both went unused on the retrieval +/// side (`_ffn` in the Apollo impl) and forced harnesses to construct +/// a router they then ignored. +/// +/// Returns [`Result`] instead of `Option` for the +/// same reasons as `KvEngine`'s post-2026-05-24 migration: silent +/// `None` propagation through `filter_map` masked Apollo's +/// store-miss rate, and the bench `panic!` on the same `None` made +/// retrieval-miss prompts crash a run that ought to have logged a +/// skip. See [`EngineError`] for the variant taxonomy. +pub trait RetrievalEngine: Send { + fn name(&self) -> &str; + + /// Runtime diagnostics: engine name, backend, config, description. + fn info(&self) -> EngineInfo; + + /// Run the prefill forward pass over the prompt tokens, consulting + /// the engine's retrieval store. Returns the hidden state at the + /// final token position (shape `[1, hidden_dim]`). + fn prefill( + &mut self, + weights: &ModelWeights, + token_ids: &[u32], + ) -> Result, EngineError>; + + /// One autoregressive decode step for the next token, applying any + /// retrieval-engine-specific state update (e.g. injection-delta + /// accumulation). Returns the hidden state (shape `[1, hidden_dim]`). + fn decode_step( + &mut self, + weights: &ModelWeights, + token_id: u32, + ) -> Result, EngineError>; + + /// Prefill against a Q4K-quantised vindex. **No default** — the trait + /// default returns an `InvariantViolation` because the `ffn`-less `prefill` + /// it would delegate to can't be handed an engine-owned dequant scratch + /// without mutating `weights`. Engines that serve Q4K (e.g. Apollo, which + /// runs its forward through `forward_raw_logits` and dequantises attn+FFN + /// into its own `dequant_scratch`) override this. + fn prefill_quant( + &mut self, + weights: &ModelWeights, + index: &larql_vindex::VectorIndex, + token_ids: &[u32], + ) -> Result, EngineError> { + let _ = (weights, index, token_ids); + Err(EngineError::InvariantViolation { + what: "RetrievalEngine::prefill_quant must be overridden for Q4K vindexes — the \ + default cannot thread an engine-owned dequant scratch through the \ + `ffn`-less `prefill` (it would have to mutate `weights`)." + .into(), + }) + } + + /// One decode step against a Q4K-quantised vindex. No default — engines + /// that serve Q4K must override (see [`prefill_quant`](Self::prefill_quant)). + fn decode_step_quant( + &mut self, + weights: &ModelWeights, + index: &larql_vindex::VectorIndex, + token_id: u32, + ) -> Result, EngineError> { + let _ = (weights, index, token_id); + Err(EngineError::InvariantViolation { + what: "RetrievalEngine::decode_step_quant must be overridden for Q4K vindexes.".into(), + }) + } + + /// Bytes of persistent engine state (excludes model weights). + fn memory_bytes(&self) -> usize; + + /// Token count in the active window (varies by engine type). + fn window_tokens(&self) -> usize { + 0 + } + + /// Cold-tier bytes (store / archive past the hot window). + fn cold_bytes(&self) -> usize { + 0 + } + + /// Per-stage timing summary. Returns `None` if profiling was not enabled. + fn stage_summary(&self) -> Option { + None + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::kv_engine::test_stubs::StubRetrievalEngine; + + #[test] + fn retrieval_engine_propagates_empty_prompt_error() { + let weights = crate::test_utils::make_test_weights(); + let mut engine = StubRetrievalEngine { + prefill_calls: 0, + decode_calls: 0, + last_token: None, + }; + let err = engine.prefill(&weights, &[]).unwrap_err(); + assert!(matches!(err, EngineError::EmptyPrompt)); + } + + #[test] + fn retrieval_engine_defaults_zero_window_and_cold_bytes() { + let engine = StubRetrievalEngine { + prefill_calls: 0, + decode_calls: 0, + last_token: None, + }; + assert_eq!(engine.window_tokens(), 0); + assert_eq!(engine.cold_bytes(), 0); + assert!(engine.stage_summary().is_none()); + } +} diff --git a/crates/larql-inference/src/kv_engine/stages.rs b/crates/larql-inference/src/kv_engine/stages.rs new file mode 100644 index 000000000..e777cd458 --- /dev/null +++ b/crates/larql-inference/src/kv_engine/stages.rs @@ -0,0 +1,182 @@ +//! [`DecodeStageSummary`] — per-step timing averages for a completed run. + +/// Per-step averages for a completed engine run. Returned from +/// [`KvEngine::stage_summary`] when profiling was enabled at engine +/// construction. +#[derive(Debug, Clone)] +pub struct DecodeStageSummary { + pub engine: String, + pub backend: String, + pub steps: usize, + pub avg_embed_us: f64, + /// K/V recompute from stored residuals (MarkovRS only). Split by tier. + pub avg_recompute_cold_us: f64, + pub avg_recompute_hot_us: f64, + pub avg_attention_us: f64, + pub avg_ffn_us: f64, + pub avg_total_decode_us: f64, + /// W10 instrumentation: time spent inside the backend's + /// `coarse_decode_step_with_state_masked` call — kernel run + + /// state-dump readback (skipped under HOnly / None). Zero on + /// non-dispatch paths and on engines that don't capture state. + pub avg_state_capture_us: f64, + /// W10 instrumentation: cumulative time inside per-layer handle + /// materialise calls (`StateHandle::into_array`). Tracks the + /// CPU bridge cost from the captured dump to engine-owned + /// `Array2`s. Zero under None mask (engine drops handles + /// without materialising). + pub avg_state_materialise_us: f64, + /// W10 instrumentation: cumulative time appending materialised + /// state into engine slabs (`append_row` calls). Tracks + /// `rs.stored` / `rs.hot_kv` growth. Zero under None mask. + pub avg_state_append_us: f64, +} + +impl DecodeStageSummary { + pub fn avg_recompute_total_us(&self) -> f64 { + self.avg_recompute_cold_us + self.avg_recompute_hot_us + } + + /// Print a human-readable breakdown table. + pub fn print(&self) { + let total = self.avg_total_decode_us; + let pct = |v: f64| if total > 0.0 { v / total * 100.0 } else { 0.0 }; + + println!( + "\nStage breakdown ({}, {}, {} decode steps avg):", + self.engine, self.backend, self.steps + ); + println!(" {:<25} {:>8} {:>6}", "Stage", "avg_us", "%"); + println!(" {}", "-".repeat(45)); + println!( + " {:<25} {:>8.1} {:>5.1}%", + "embed", + self.avg_embed_us, + pct(self.avg_embed_us) + ); + if self.avg_recompute_total_us() > 0.0 { + println!( + " {:<25} {:>8.1} {:>5.1}%", + "recompute_kv (cold)", + self.avg_recompute_cold_us, + pct(self.avg_recompute_cold_us) + ); + println!( + " {:<25} {:>8.1} {:>5.1}%", + "recompute_kv (hot)", + self.avg_recompute_hot_us, + pct(self.avg_recompute_hot_us) + ); + } + println!( + " {:<25} {:>8.1} {:>5.1}%", + "attention", + self.avg_attention_us, + pct(self.avg_attention_us) + ); + println!( + " {:<25} {:>8.1} {:>5.1}%", + "ffn", + self.avg_ffn_us, + pct(self.avg_ffn_us) + ); + // W10 instrumentation: only print state lines when populated + // (avoids noise on engines that don't capture state). + let state_total = + self.avg_state_capture_us + self.avg_state_materialise_us + self.avg_state_append_us; + if state_total > 0.0 { + println!( + " {:<25} {:>8.1} {:>5.1}%", + "state_capture", + self.avg_state_capture_us, + pct(self.avg_state_capture_us) + ); + println!( + " {:<25} {:>8.1} {:>5.1}%", + "state_materialise", + self.avg_state_materialise_us, + pct(self.avg_state_materialise_us) + ); + println!( + " {:<25} {:>8.1} {:>5.1}%", + "state_append", + self.avg_state_append_us, + pct(self.avg_state_append_us) + ); + } + println!(" {}", "-".repeat(45)); + println!( + " {:<25} {:>8.1} {:>5.1}%", + "total (measured)", total, 100.0 + ); + println!(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn decode_stage_summary_recompute_total() { + let s = DecodeStageSummary { + engine: "test".into(), + backend: "cpu".into(), + steps: 10, + avg_embed_us: 1.0, + avg_recompute_cold_us: 2.0, + avg_recompute_hot_us: 3.0, + avg_attention_us: 4.0, + avg_ffn_us: 5.0, + avg_total_decode_us: 15.0, + avg_state_capture_us: 0.0, + avg_state_materialise_us: 0.0, + avg_state_append_us: 0.0, + }; + assert_eq!(s.avg_recompute_total_us(), 5.0); + } + + /// Cover `DecodeStageSummary::print` — both the recompute>0 branch and + /// the total>0 percentage branch. Output goes to stdout (captured by the + /// test harness); this is a smoke test for the formatting code path. + #[test] + fn decode_stage_summary_print_with_recompute() { + let s = DecodeStageSummary { + engine: "markov-rs".into(), + backend: "cpu".into(), + steps: 10, + avg_embed_us: 100.0, + avg_recompute_cold_us: 500.0, + avg_recompute_hot_us: 300.0, + avg_attention_us: 1500.0, + avg_ffn_us: 800.0, + avg_total_decode_us: 3200.0, + avg_state_capture_us: 0.0, + avg_state_materialise_us: 0.0, + avg_state_append_us: 0.0, + }; + s.print(); + } + + /// `print` must also handle the no-recompute, zero-total branch — the + /// `pct` fallback when `avg_total_decode_us == 0.0` and the + /// `avg_recompute_total_us() == 0` short-circuit. + #[test] + fn decode_stage_summary_print_no_recompute_zero_total() { + let s = DecodeStageSummary { + engine: "no-cache".into(), + backend: "metal".into(), + steps: 0, + avg_embed_us: 0.0, + avg_recompute_cold_us: 0.0, + avg_recompute_hot_us: 0.0, + avg_attention_us: 0.0, + avg_ffn_us: 0.0, + avg_total_decode_us: 0.0, + avg_state_capture_us: 0.0, + avg_state_materialise_us: 0.0, + avg_state_append_us: 0.0, + }; + s.print(); + } +} diff --git a/crates/larql-inference/src/kv_engine/test_stubs.rs b/crates/larql-inference/src/kv_engine/test_stubs.rs new file mode 100644 index 000000000..e6dbd4d75 --- /dev/null +++ b/crates/larql-inference/src/kv_engine/test_stubs.rs @@ -0,0 +1,100 @@ +//! Minimal engine implementations shared by this module's tests. +//! +//! Both exist to exercise the traits' *default* method bodies: they implement +//! only the required methods, so every default a real engine would override +//! stays on its default path. Shared here rather than duplicated per file +//! because [`AnyEngine`](super::AnyEngine)'s tests need both families, and a +//! second copy would drift from the one under test. + +use super::{EngineError, EngineInfo, KvEngine, RetrievalEngine}; +use crate::ffn::FfnBackend; +use crate::ModelWeights; +use ndarray::Array2; + +/// Synthetic engine that only implements the required trait methods, +/// leaving every default (`window_tokens`, `cold_bytes`, `stage_summary`, +/// `prefill_quant`, `decode_step_quant`) to fire. Exercises the default +/// bodies that no shipped engine routes through (every concrete engine +/// overrides them). +pub(crate) struct DefaultsOnlyEngine { + pub(crate) prefill_calls: usize, + pub(crate) decode_calls: usize, +} + +impl KvEngine for DefaultsOnlyEngine { + fn name(&self) -> &str { + "defaults-only" + } + fn info(&self) -> EngineInfo { + EngineInfo { + name: self.name().into(), + description: "test fixture".into(), + backend: "cpu".into(), + config: String::new(), + } + } + fn prefill( + &mut self, + _weights: &ModelWeights, + _ffn: &dyn FfnBackend, + _token_ids: &[u32], + ) -> Result, EngineError> { + self.prefill_calls += 1; + Ok(Array2::zeros((1, 4))) + } + fn decode_step( + &mut self, + _weights: &ModelWeights, + _ffn: &dyn FfnBackend, + _token_id: u32, + ) -> Result, EngineError> { + self.decode_calls += 1; + Ok(Array2::zeros((1, 4))) + } + fn memory_bytes(&self) -> usize { + 0 + } +} + +pub(crate) struct StubRetrievalEngine { + pub(crate) prefill_calls: usize, + pub(crate) decode_calls: usize, + pub(crate) last_token: Option, +} + +impl RetrievalEngine for StubRetrievalEngine { + fn name(&self) -> &str { + "stub-retrieval" + } + fn info(&self) -> EngineInfo { + EngineInfo { + name: self.name().into(), + description: "test fixture".into(), + backend: "cpu".into(), + config: String::new(), + } + } + fn prefill( + &mut self, + _weights: &ModelWeights, + token_ids: &[u32], + ) -> Result, EngineError> { + self.prefill_calls += 1; + if token_ids.is_empty() { + return Err(EngineError::EmptyPrompt); + } + Ok(Array2::zeros((1, 4))) + } + fn decode_step( + &mut self, + _weights: &ModelWeights, + token_id: u32, + ) -> Result, EngineError> { + self.decode_calls += 1; + self.last_token = Some(token_id); + Ok(Array2::zeros((1, 4))) + } + fn memory_bytes(&self) -> usize { + 16 + } +} diff --git a/crates/larql-inference/src/lib.rs b/crates/larql-inference/src/lib.rs index 9c31a43ff..d7d649fcf 100644 --- a/crates/larql-inference/src/lib.rs +++ b/crates/larql-inference/src/lib.rs @@ -306,7 +306,7 @@ pub mod prelude { /// `KvEngine`, `EngineInfo`, and `DecodeStageSummary` are defined in /// this crate's [`kv_engine`](crate::kv_engine) module and re-exported /// at the crate root. Concrete engine implementations -/// (`MarkovResidualEngine`, `UnlimitedContextEngine`, `StandardEngine`, +/// (`MarkovResidualEngine`, `WindowedCheckpointEngine`, `StandardEngine`, /// `NoCacheEngine`, `TurboQuantEngine`, `ApolloEngine`) plus /// `EngineKind` and accuracy helpers (`compare_hidden`, /// `cosine_similarity`, `kl_divergence`, …) live in the `larql-kv` diff --git a/crates/larql-inference/src/test_utils.rs b/crates/larql-inference/src/test_utils.rs index d1e161cd9..bd51b4e60 100644 --- a/crates/larql-inference/src/test_utils.rs +++ b/crates/larql-inference/src/test_utils.rs @@ -694,7 +694,8 @@ pub use larql_models::test_fixtures::{make_gemma3_test_weights, make_starcoder2_ pub use larql_models::test_fixtures::{ arc_mmap_from_bytes, make_test_q4k_weights, make_test_q4k_weights_layers, make_test_q4k_weights_rope_scaled, make_test_q4k_weights_silu, make_test_q4k_weights_wide, - Q4K_TEST_HIDDEN, Q4K_TEST_INTER, Q4K_TEST_INTER_WIDE, Q4K_TEST_NUM_LAYERS, Q4K_TEST_VOCAB, + make_test_q4k_weights_with_dims, Q4K_TEST_HIDDEN, Q4K_TEST_INTER, Q4K_TEST_INTER_WIDE, + Q4K_TEST_NUM_LAYERS, Q4K_TEST_VOCAB, }; /// Build a fully-populated synthetic `VectorIndex` that satisfies the /// cached + direct-matvec decode contract on the Q4_K weights from diff --git a/crates/larql-inference/src/vindex/kquant_forward/cached.rs b/crates/larql-inference/src/vindex/kquant_forward/cached.rs index 60cfc7d7a..f3044a440 100644 --- a/crates/larql-inference/src/vindex/kquant_forward/cached.rs +++ b/crates/larql-inference/src/vindex/kquant_forward/cached.rs @@ -63,7 +63,7 @@ pub fn predict_kquant_prefill( /// `state` is `Some`, populates per-layer `h_in` ([seq_len, hidden]), /// `k_new` ([seq_len, kv_dim]), `v_new` ([seq_len, kv_dim]) for every /// position in the prompt — engines (markov_residual, -/// unlimited_context, turbo_quant) use this to seed their state policy +/// windowed_checkpoint, turbo_quant) use this to seed their state policy /// from a single prefill pass without a follow-up CPU re-walk. When /// `state` is `None`, bit-identical to [`predict_kquant_prefill`]. pub fn predict_kquant_prefill_with_state( diff --git a/crates/larql-inference/src/vindex/kquant_forward/hidden.rs b/crates/larql-inference/src/vindex/kquant_forward/hidden.rs index 41dace365..be25b0223 100644 --- a/crates/larql-inference/src/vindex/kquant_forward/hidden.rs +++ b/crates/larql-inference/src/vindex/kquant_forward/hidden.rs @@ -18,7 +18,7 @@ pub fn predict_kquant_hidden( weights: &ModelWeights, token_ids: &[u32], index: &VectorIndex, - moe_remote: Option<&crate::ffn::RemoteMoeBackend>, + moe: Option<&dyn crate::ffn::MoeExpertBackend>, ) -> Array2 { let num_layers = weights.num_layers; let mut scratch = larql_models::DequantScratch::new(); @@ -56,7 +56,7 @@ pub fn predict_kquant_hidden( &ffn_backend, ple_inputs.get(layer), shared_kv, - moe_remote, + moe, ) { h = h_new; if let Some(kv) = kv_out { @@ -135,7 +135,7 @@ fn run_moe_layer_cpu( ffn: &dyn crate::ffn::FfnBackend, ple_input: Option<&Array2>, shared_kv: Option<&SharedKV>, - moe_remote: Option<&crate::ffn::RemoteMoeBackend>, + moe: Option<&dyn crate::ffn::MoeExpertBackend>, ) -> Option<(Array2, Option)> { let (h_post_attn, kv_out) = if let Some(shared) = shared_kv { let (h_pa, _, _) = @@ -153,14 +153,14 @@ fn run_moe_layer_cpu( layer, ffn, ple_input, - moe_remote, + moe, ); Some((h_out, kv_out)) } /// CPU MoE FFN block for one hybrid-MoE layer, given the **post-attention** /// hidden state. Computes the dense FFN contribution (`h1`), the expert -/// contribution (`h2` — remote via `moe_remote` when set, else local +/// contribution (`h2` — via `moe` when a route is bound, else local /// `cpu_moe_forward`), combines + outer-norms, and applies PLE + /// layer-scalar. Returns the full layer output (the new residual). /// @@ -179,17 +179,9 @@ pub fn moe_ffn_block_cpu( layer: usize, ffn: &dyn crate::ffn::FfnBackend, ple_input: Option<&Array2>, - moe_remote: Option<&crate::ffn::RemoteMoeBackend>, + moe: Option<&dyn crate::ffn::MoeExpertBackend>, ) -> Array2 { - moe_ffn_block_cpu_with_index( - weights, - h_post_attn, - layer, - ffn, - ple_input, - moe_remote, - None, - ) + moe_ffn_block_cpu_with_index(weights, h_post_attn, layer, ffn, ple_input, moe, None) } /// `LARQL_Q4K_DIRECT_FFN=1` routes the hybrid-MoE *dense slab* through the @@ -214,7 +206,7 @@ pub fn moe_ffn_block_cpu_with_index( layer: usize, ffn: &dyn crate::ffn::FfnBackend, ple_input: Option<&Array2>, - moe_remote: Option<&crate::ffn::RemoteMoeBackend>, + moe: Option<&dyn crate::ffn::MoeExpertBackend>, index: Option<&larql_vindex::VectorIndex>, ) -> Array2 { let arch = &*weights.arch; @@ -263,15 +255,19 @@ pub fn moe_ffn_block_cpu_with_index( let seq_len = h_post_attn.nrows(); let mut h2 = Array2::::zeros((seq_len, hidden)); - if let Some(remote) = moe_remote { - if let Some(router) = build_moe_router_weights(weights, arch, layer) { - let _t_expert = std::time::Instant::now(); - let out = remote.forward_moe_seq(layer, h_post_attn, &router, norm_offset, eps); - crate::decode_stages::record_expert(_t_expert.elapsed().as_nanos()); - match out { - Ok(out) => h2 = out, - Err(e) => eprintln!("[moe_ffn_block_cpu] remote dispatch error L{layer}: {e}"), - } + if let Some(backend) = moe { + // Every non-default route goes through one call. Which route it is — + // remote shards, a VINDEX3 bound plan — is the backend's business, and + // the block loop's only job is to place the contribution. + let _t_expert = std::time::Instant::now(); + let out = backend.forward_moe_seq(weights, layer, h_post_attn, norm_offset, eps); + crate::decode_stages::record_expert(_t_expert.elapsed().as_nanos()); + match out { + Ok(out) => h2 = out, + Err(e) => eprintln!( + "[moe_ffn_block_cpu] {} dispatch error L{layer}: {e}", + backend.name() + ), } } else { // Local experts count toward the expert stage too (`LARQL_DECODE_STAGES`) diff --git a/crates/larql-inference/src/vindex/kquant_forward/remote_ffn.rs b/crates/larql-inference/src/vindex/kquant_forward/remote_ffn.rs index 4619a3933..7f347ac12 100644 --- a/crates/larql-inference/src/vindex/kquant_forward/remote_ffn.rs +++ b/crates/larql-inference/src/vindex/kquant_forward/remote_ffn.rs @@ -13,6 +13,9 @@ use super::dequant::dequantize_matrix; /// End-to-end predict on a Q4_K vindex with the FFN served by an external /// [`crate::ffn::FfnBackend`]. +/// +/// A refusal propagates rather than degrading to the dense half of the layer +/// that refused — see [`predict_kquant_hidden_inner`]. pub fn predict_kquant_with_ffn( weights: &mut ModelWeights, tokenizer: &Tokenizer, @@ -20,9 +23,11 @@ pub fn predict_kquant_with_ffn( top_k: usize, index: &VectorIndex, ffn_backend: &dyn crate::ffn::FfnBackend, -) -> PredictResult { - let h = predict_kquant_hidden_with_ffn(weights, token_ids, index, ffn_backend); - crate::forward::predict::logits_to_predictions_pub(weights, &h, tokenizer, top_k, 1.0) +) -> Result { + let h = predict_kquant_hidden_with_ffn(weights, token_ids, index, ffn_backend)?; + Ok(crate::forward::predict::logits_to_predictions_pub( + weights, &h, tokenizer, top_k, 1.0, + )) } /// **Early-exit** Q4_K predict — the q4k twin of @@ -41,7 +46,7 @@ pub fn predict_kquant_with_ffn_early_exit( ffn_backend: &dyn crate::ffn::FfnBackend, stop_layer: usize, on_stop: &mut dyn FnMut() -> Option>, -) -> (Vec<(String, f64)>, bool) { +) -> Result<(Vec<(String, f64)>, bool), larql_execution::BoxRefusal> { let mut early_preds: Option> = None; let (h, exited); { @@ -59,41 +64,53 @@ pub fn predict_kquant_with_ffn_early_exit( index, ffn_backend, Some((stop_layer, &mut stop_hook)), - ); + )?; } if exited { - (early_preds.unwrap_or_default(), true) + Ok((early_preds.unwrap_or_default(), true)) } else { - ( + Ok(( crate::forward::predict::logits_to_predictions_pub(weights, &h, tokenizer, top_k, 1.0) .predictions, false, - ) + )) } } /// End-to-end hidden-state forward on a Q4_K vindex with the FFN served by an /// external [`crate::ffn::FfnBackend`]. +/// +/// A refusal propagates — see [`predict_kquant_hidden_inner`]. pub fn predict_kquant_hidden_with_ffn( weights: &mut ModelWeights, token_ids: &[u32], index: &VectorIndex, ffn_backend: &dyn crate::ffn::FfnBackend, -) -> ndarray::Array2 { - predict_kquant_hidden_inner(weights, token_ids, index, ffn_backend, None).0 +) -> Result, larql_execution::BoxRefusal> { + Ok(predict_kquant_hidden_inner(weights, token_ids, index, ffn_backend, None)?.0) } /// Core Q4_K hidden forward with an optional early-exit hook. `early = /// Some((stop_layer, on_stop))` checks `on_stop()` after `stop_layer` completes; /// `true` returns the current hidden + `exited = true`. `None` runs the full /// stack (the behaviour of [`predict_kquant_hidden_with_ffn`]). +/// +/// # A refusal is not a hidden state +/// +/// The hybrid-MoE branch below routes a whole layer to `ffn_backend`. When that +/// route refuses, this returns the refusal instead of falling through to the +/// dense-only path: completing the layer with its dense half would answer with +/// a hidden state the model never computed, and the caller has no way to tell +/// that from a real one. Same contract as +/// [`crate::kv_dispatch::helpers`]'s per-layer dispatch — the hook's three +/// outcomes stay three all the way out. fn predict_kquant_hidden_inner( weights: &mut ModelWeights, token_ids: &[u32], index: &VectorIndex, ffn_backend: &dyn crate::ffn::FfnBackend, mut early: Option<(usize, &mut dyn FnMut() -> bool)>, -) -> (ndarray::Array2, bool) { +) -> Result<(ndarray::Array2, bool), larql_execution::BoxRefusal> { let num_layers = weights.num_layers; let hidden = weights.hidden_size; @@ -137,7 +154,26 @@ fn predict_kquant_hidden_inner( &h, layer, ) { - if let Some(h_out) = ffn_backend.forward_moe_full_layer(layer, &h_post_attn) { + // A refusal propagates; `None` falls through to the dense path + // below. Not-applicable means this backend does not serve the + // layer, and the local dispatch is the correct answer — a + // refusal is not that, and must not wear its shape. + let routed = match ffn_backend.forward_moe_full_layer(layer, &h_post_attn) { + Ok(out) => out, + Err(refusal) => { + // The four dequantised attention tensors inserted + // above are this loop's scratch, and every other exit + // drops them. Propagating without this would leave a + // layer of f32 attention weights in the caller's + // `weights` — a refusal that silently grows the model. + weights.tensors.remove(&q_key); + weights.tensors.remove(&k_key); + weights.tensors.remove(&v_key); + weights.tensors.remove(&o_key); + return Err(refusal); + } + }; + if let Some(h_out) = routed { h = h_out; weights.tensors.remove(&q_key); weights.tensors.remove(&k_key); @@ -178,12 +214,12 @@ fn predict_kquant_hidden_inner( // residual trace, so the verified route would abstain there anyway. if let Some((stop, on_stop)) = early.as_mut() { if layer == *stop && on_stop() { - return (h, true); + return Ok((h, true)); } } } - (h, false) + Ok((h, false)) } #[cfg(test)] @@ -194,6 +230,83 @@ mod tests { make_test_gemma4_moe_weights, make_test_q4k_vindex, make_test_q4k_weights, make_test_tokenizer, }; + use larql_execution::{ExecutionRefusal, RefusalKind}; + + /// An FFN backend that refuses every routed layer. + #[derive(Debug)] + struct RefusingFfn; + + #[derive(Debug)] + struct NoExpertHere; + + impl std::fmt::Display for NoExpertHere { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("expert not resident on this shard") + } + } + impl std::error::Error for NoExpertHere {} + impl ExecutionRefusal for NoExpertHere { + fn kind(&self) -> RefusalKind { + RefusalKind::Residency + } + } + + impl larql_compute::ffn::FfnBackend for RefusingFfn { + fn forward(&self, _layer: usize, x: &ndarray::Array2) -> ndarray::Array2 { + x.clone() + } + fn name(&self) -> &str { + "refusing" + } + fn forward_moe_full_layer( + &self, + _layer: usize, + _h_post_attn: &ndarray::Array2, + ) -> Result>, larql_execution::BoxRefusal> { + Err(Box::new(NoExpertHere)) + } + } + + /// The point of the error channel: a routed layer that refused must not + /// come back as a hidden state. Before this propagated, the refusal was + /// logged and the loop fell through to the dense-only path, so the caller + /// received a plausible hidden built without the expert that declined — + /// indistinguishable from a real one. + #[test] + fn a_refused_layer_does_not_become_a_hidden_state() { + let mut weights = make_test_gemma4_moe_weights(); + let index = make_test_q4k_vindex(&weights); + let err = predict_kquant_hidden_with_ffn(&mut weights, &[0u32, 1], &index, &RefusingFfn) + .expect_err("a refused MoE layer must not yield a hidden state"); + assert_eq!(err.kind(), RefusalKind::Residency); + } + + /// The refusal path cleans up after itself. The loop inserts four + /// dequantised attention matrices per layer as scratch and drops them on + /// every other exit; propagating without that cleanup would leave a + /// layer's worth of f32 weights in the caller's model on the way out. + #[test] + fn a_refusal_leaves_no_scratch_attention_tensors_behind() { + let mut weights = make_test_gemma4_moe_weights(); + let index = make_test_q4k_vindex(&weights); + let scratch = { + let arch = &*weights.arch; + [ + arch.attn_q_key(0), + arch.attn_k_key(0), + arch.attn_v_key(0), + arch.attn_o_key(0), + ] + }; + let _ = predict_kquant_hidden_with_ffn(&mut weights, &[0u32, 1], &index, &RefusingFfn) + .expect_err("fixture refuses"); + for key in &scratch { + assert!( + !weights.tensors.contains_key(key), + "refusal left dequantised scratch tensor {key} in the caller's weights" + ); + } + } /// `predict_kquant_hidden_with_ffn` end-to-end against the Q4K /// fixture using a `WeightFfn` backend. Non-MoE arch → the @@ -211,7 +324,8 @@ mod tests { let ffn = WeightFfn { weights: weights_ref, }; - let h = predict_kquant_hidden_with_ffn(&mut weights, &[0u32, 1], &index, &ffn); + let h = predict_kquant_hidden_with_ffn(&mut weights, &[0u32, 1], &index, &ffn) + .expect("WeightFfn never refuses"); assert_eq!(h.shape(), &[2, weights.hidden_size]); assert!(h.iter().all(|v| v.is_finite())); } @@ -228,7 +342,8 @@ mod tests { let ffn = WeightFfn { weights: weights_ref, }; - let h = predict_kquant_hidden_with_ffn(&mut weights, &[0u32, 1], &index, &ffn); + let h = predict_kquant_hidden_with_ffn(&mut weights, &[0u32, 1], &index, &ffn) + .expect("WeightFfn never refuses"); assert_eq!(h.shape(), &[2, weights.hidden_size]); assert!(h.iter().all(|v| v.is_finite())); } @@ -242,7 +357,8 @@ mod tests { let ffn = WeightFfn { weights: weights_ref, }; - let result = predict_kquant_with_ffn(&mut weights, &tokenizer, &[0u32, 1], 3, &index, &ffn); + let result = predict_kquant_with_ffn(&mut weights, &tokenizer, &[0u32, 1], 3, &index, &ffn) + .expect("WeightFfn never refuses"); assert!(result.predictions.len() <= 3); } } diff --git a/crates/larql-kv/CHANGELOG.md b/crates/larql-kv/CHANGELOG.md index a6028f4c2..12f6c2b1f 100644 --- a/crates/larql-kv/CHANGELOG.md +++ b/crates/larql-kv/CHANGELOG.md @@ -6,6 +6,268 @@ The format follows [Keep a Changelog](https://keepachangelog.com/) conventions with dated entries (`YYYY-MM-DD`) instead of semantic versions during the pre-1.0 phase. Forward-looking work lives in [`ROADMAP.md`](ROADMAP.md). +## Windowed engines keep the fused path; the bench measures a whole token (2026-08-04) + +Started as "check the engines benchmark correctly" and the instrument was the +first finding. + +**The bench was not measuring a token.** The engine harness stopped its timer +before `pick_next`, while the reference rows included lm_head — both landing in +the same tok/s column, so every engine read 2-3x faster than production for +free. The CPU run made it exact: the engine's whole measured step (4.12 ms) +equalled the reference's *forward stage alone* (4.109 ms). Both halves are +inside the step now, prefill included (the reference's `prefill_ms` encloses its +first `lm_head_predict`), and each row carries a `fwd=` / `head=` split. + +**Memory had two undercounts.** Metal's coarse pipeline keeps K/V behind a +sentinel handle, so engines reported 0 bytes owned; and the CPU whole-model +handle was measured with the per-layer formula — a 28x undercount that printed +as "13x vs std-kv" for an engine doing no compression. Backends now report what +they hold (`backend_resident_kv_bytes`), whole-model handles report every layer +(`KvHandleInner::resident_bytes`), and no ratio is invented when nothing was +measured. + +**Rows say which path they took.** `DispatchPath` plus the backend's +`per_layer_is_host_delegated` answer, so `[coarse]` and `[per-layer→host]` are +visible. That matters because on Metal *every* per-layer dispatch method +delegates to `CpuBackend` — a windowed engine ran its whole forward on the host +under a `[metal (GPU)]` label. + +**Per-layer SWA on the CPU attention path.** It had none, while the Metal +pipeline spec carried a window, so a Gemma-class model attended full history on +layers the architecture declares sliding. Both now resolve through one rule, +`effective_attention_window_for_layer`; a layer declaring itself sliding without +a width is answered "no window" deliberately rather than by an `unwrap_or(0)` +that meant two different things in two places. + +**Windowed engines keep the fused path.** A window promises bounded attention +*and* bounded K/V; the coarse surface had neither, so windowed engines fell to +per-layer — 9.4x on Gemma 3 4B. Added `coarse_prefill_windowed` / +`coarse_decode_step_windowed`, fail-closed so a backend that cannot bound both +declines and nothing regresses. CPU trims the cache before each step; Metal +clamps the attention span every step and compacts at 2x the window so the +memmove is O(1) amortised. Metal, 80 steps at `window=8`: **11.61 ms / 2.4 MB** +against 12.06 ms / 23.6 MB unwindowed, from 115.44 ms before. + +That needed a prerequisite: `LayerKVCache::current_len` was answering both "rows +stored" and "stream position". They agree only while nothing is evicted, so +compaction would have rewound RoPE on every later token. Split into +`abs_position` and `current_len`. + +**First GPU-path tests in this crate** (`tests/gpu_engine_parity`). Every prior +test, bench and pin built engines with `cpu_engine_backend()`, and the `gpu` +feature gates only a dependency — the test count was identical with and without +it. Criterion covered 7 of 9 engines and timed apollo's `RetrievalMiss` as a +250x win; the roster is pinned now. + +**A cross-backend divergence closed, after two wrong diagnoses.** Metal's +batched prefill disagreed with the CPU by 23-43% on the Gemma-3 fixture. Blamed +first on per-layer SWA (wrong — no window resolves on that fixture), then on a +`head_dim` shape assumption (wrong, and asserted without testing — a sweep over +head_dim ∈ {32…512} diverged at every shape including the real model's 256). +The cause was the fixture declaring Gemma-3's QK-norm keys and never populating +the weights, so the two backends disagreed about a declared-but-absent weight. +Real checkpoints always carry them, which is why a real Gemma 3 4B agreed to +4.0e-7 throughout. Fixture fixed; `make_test_q4k_weights_with_dims` added, +because every Q4K fixture was pinned to `head_dim = 64` and shape sensitivity +was untestable by construction. + +**Open:** `standard:window=N` still declines the fused path when the *prompt* +exceeds the window — the fused prefill has no per-query-position masking. Metal +holds up to 2x the window resident between compactions (attention is still +bounded at the window). + +## Spin-barrier pool — CPU MoE decode caught llama.cpp (2026-06-13) + +After residency closed the byte-traffic gap (06-11/12), a `/usr/bin/sample` of +live 26B decode showed the remaining ~1.15× was **rayon fork-join overhead**, +not kernels. The decode driver runs *outside* the global rayon pool, so each of +the ~211 parallel sections/token took the cold path (`in_worker_cold → +LockLatch::wait_and_reset → __psynch_cvwait`) and workers slept between sections +— ~40% of thread-time in wait states. + +**Built** [`larql_compute::cpu::spin_pool`](../../larql-compute/src/cpu/spin_pool.rs): +a llama.cpp-style persistent spin-barrier pool. Workers spin on an epoch counter +and only `park` after a long idle gap; the dispatcher participates as the n-th +worker; **static strided chunk ownership** makes `completed == num_chunks` a +sound barrier (no shared resettable cursor → no stale re-claim across +back-to-back dispatches — a concurrent-dispatcher test caught that bug); a +dispatch `Mutex` + thread-local reentrancy guard make it safe for +`--concurrent`/multi-threaded tests. `par_chunks_mut` / `par_chunks_mut2` +helpers route a row-chunked parallel-for through the pool, or rayon when +`LARQL_SPIN_POOL=0`. **Default-on** (see "Decode fast path default-on" — the +whole Q4K stack ships on, opt out per stage with `=0`); both paths are +numerically identical, only the threading differs. + +**Centralized** the four byte-identical `par_chunks_mut` Q4_K/Q6_K×Q8_K matvec +copies (larql-compute `cached.rs`, larql-inference `cached.rs`, lm_head ×2 in +`dense.rs` — the prior "consolidation hazard") into one +`q4k_q8k_matvec_parallel`, and routed every hot decode section (attention int8 +Q/K/V/O, GQA, dense FFN gate/up/down, geglu, expert fold, lm_head q4 + f32) +through it — so when enabled the whole token runs on one hot pool. + +- **Parity:** 704 compute + 1220 inference + 756 kv green, flags-off AND + flags-on (incl. the `predict_kquant` oracles). clippy clean. +- **Profile after:** rayon eliminated from the hot path — `in_worker_cold` + 2682→0, `join_context` 10300→0, `wait_until_cold` 4463→9. +- **Measured** (M3 Max, t=8, warm, tight A/B bracket, flags **inline**): + 26B short-ctx OFF ~26.9 → ON **33–35**; n=256 OFF ~27.4 → ON **~34.9 + (+28%)** — vs llama.cpp recorded **32.1** ⇒ ~9% ahead. +- **Default-on + safe (2026-06-13):** shipped a spin→yield→park backoff (spin + the proven window during active decode → yield once a wait outlives a token → + park when idle, ~0 CPU; dispatcher unparks on dispatch) so the pool doesn't + peg cores between requests — what makes on-by-default safe on a shared box. + Also fixed a panic-safety bug (a panicking chunk killed a worker → the + barrier spun forever): `catch_unwind` per chunk + re-raise on the dispatcher. +- **Caveat:** the pool spins during active decode (the win on a dedicated box); + under a transient mid-decode load spike a run can still regress (an n=512 ON + run hit 10.7 once) — `LARQL_SPIN_POOL=0` falls back to rayon if needed. + +## CPU resident fast-path — all engines pluggable into it (2026-06-13) + +The 2026-06-11/12 CPU fast-path arc (Q4K-direct + int8 attention, q4k +lm_head/dense residency, hand-asm kernels, KV append-in-place — see +`bench/baselines/c10_gemma4-26b-a4b_cpu_reconciled.json`) initially landed +only on `StandardEngine`: the `KvEngine::decode_step_resident` trait default +DROPPED the index (`let _ = index`), so every own-walk-loop engine stayed on +f32 attention. **Fixed:** + +- New single-source dispatcher + `larql_compute::attention::run_attention_block_decode_step_auto` — makes + the same q4k-direct-vs-f32 per-layer choice as + `CpuBackend::attention_step`, for callers that own `SharedKV` caches. +- `markov-rs`, `markov-rs-codec`, `turbo-quant`, `windowed-checkpoint`, + `boundary_per_layer` now override `decode_step_resident` and thread the + vindex down their walk loops to `_auto`. `boundary-kv` forwards both + resident methods to its inner `StandardEngine` (was silently dropping to + the f32 path). `no_cache`/`apollo` keep the default by design (debug / + bench-only full re-forward). +- Regression pin: `engines::resident_identity_tests` — for 7 concrete + engine specs, `prefill/decode_step_resident` must be BIT-IDENTICAL to + `prefill/decode_step` with the flags off, and the covered-engine count + must not shrink. +- **Absolute matrix + slow-engine fixes 2026-06-13** (26B, default-on incl. + spin pool, M3 Max t=8 warm n=128). First measured: unlimited 31.8 / standard + 30.5 / boundary-kv 27.1 (**0.80×→0.89×**, its resident-forwarding fix) / + turbo 9.4 / markov 7.8 / codec 7.3 — the recompute/codec engines sat at + **~0.24–0.31×** because the spin pool sped up the shared attention/FFN/matvec + but not their per-step machinery. **Then fixed all three, feature intact:** + - **turbo-quant 9.4 → ~24** — `decompress_matrix`'s per-vector WHT decode was + *serial on the driver* (~35% of it); fanned across the spin pool. Still + 3-4-bit compressed (decoded every step, now parallel) — no memory tradeoff. + - **markov-rs 7.8 → 27.9, markov-rs-codec 7.3 → 27.7** — ported the W2 hot-K/V + cache to the **resident walk** (`rs_decode_step_inner`/`_codec`): read the + cached `hot_kv` and append the free `new_kv` from the attention step instead + of `recompute_kv`-ing every position each step. Gated `cache_eligible = + max_window.is_none() && no-cold` so it never tracks a window-clip + transition; the residual `stored` stays the canonical, re-derivable state + (the engine's point), the K/V is a droppable derivative. Parity gate: + `#[cfg(debug_assertions)]` assert cached K/V ≡ `recompute_kv` (≤1e-2), + exercised by `resident_identity_tests` (extended to a 10-step decode). + Final matrix: standard 34.5 / unlimited 32.1 / markov 27.9 / codec 27.7 / + boundary-kv 27.4 / turbo 21.1 — all **0.6–1.0× of standard** (was 0.24–0.31× + for the slow three). 756 kv tests green debug+release, clippy clean. +- **Comparative bottleneck review + walk allocation fix 2026-06-14.** Profiled + each engine's driver vs standard: the **shared** wall is the Q6_K expert + matvec (all engines inherit it); each engine's *delta* is its feature + machinery. markov/codec's −19/−20% was NOT the residual-store memcpy (~0.8% of + the driver) — it was **per-step allocation churn**: the resident walk's + `Array2::zeros((s_old+1, h))` rebuild + the cached-K/V `to_owned` + (`__bzero`+`szone_malloc` ≈ 2450 driver samples, idling the worker pool at 48% + vs standard's 80%). **Fixed:** the cache_eligible walk now `append_row`s + `stored` in place into the W8.2 doubling-capacity buffer (mirrors dispatch.rs) + and borrows `hot_kv` into attention via `Cow` instead of copying. Churn + collapsed 2450→150 samples (~16×); **markov/standard ratio 0.81×→0.975×, codec + 0.80×→~1.0×** (same battery state, back-to-back). Parity: resident_identity + (markov+codec, 10-step, buffer doubles) bit-exact + debug K/V assert. turbo's + −39% is **inherent** (must decode compressed K/V to attend; already + parallelized); boundary-kv/unlimited deltas are small (frame-emit/windowing). + Remaining markov/codec ~2.5% = walk-attention serial work (shared walk + frontier — full K/V concat + generic GQA vs standard's in-place handle). +- **In-place hot-K/V on the resident walk 2026-06-14 (closes the concat half).** + The named ~2.5% above was the walk-attention **owned concat**: the resident + walk drove `run_attention_block_decode_step_q4k_direct`, which allocates a + fresh `[ctx+1, kv_dim]` K *and* V every layer every step and copies the whole + prior cache into it before attending — **O(L²)** cache copy over an L-token + generation, vs `standard`'s in-place append handle (O(L)). The split + project→append→attend halves already existed for the dispatch path; the walk + just didn't use them. **Built** `run_attention_block_decode_step_{q4k_direct, + auto}_inplace` (larql-compute `attention/decode.rs`): projects the new row, + appends it into the caller's **doubling-capacity** K/V buffer (grows like + `stored`), and attends over the `[..len+1]` views — no concat. **Wired** + markov_residual + markov_residual_codec resident walks: step-1 still + recompute-seeds `hot_kv`; steps 2+ append in place (the steady state). The + windowed/cold tiers and the flags-off f32 path keep the owned concat + unchanged. Gated `LARQL_MARKOV_INPLACE_KV` (default on; `=0` → owned concat, + the A/B reference + escape hatch). **Parity (bit-exact, 4 gates):** compute- + level `inplace ≡ q4k_direct` concat across a capacity doubling; engine-level + in-place-vs-owned-concat A/B with Q4K-direct **on** for markov *and* codec + (hidden states bit-identical every step); `resident_identity` flags-off still + green (in-place branch's None-fallback = owned concat); 758 kv + 705 compute + + 1220 inference green debug & opt, clippy clean. (The debug `hot_kv ≡ + recompute_kv` assert is gated to the f32 path — the Q4K route's projections + differ from `recompute_kv` by >1e-2 even in f32-act; its oracle is the A/B.) + The two q4k-flag-mutating tests serialise on `Q4K_FLAG_ENV_LOCK` (those flags + read process env on the driver thread — no thread-local). **Perf is + structural** (eliminates the O(L²) per-step copy; the win grows with context — + it's the long-ctx tax behind the C10 1.29× vs short 1.15×). **Measured (26B, + CPU MoE in-process, M3 Max t=8, n=128 warm, `LARQL_MARKOV_INPLACE_KV` A/B, + same engine ordering):** markov 32.5→34.5, codec 32.5→34.6 with in-place on — + and the three untouched controls (standard/unlimited/turbo) drifted *down* + −3/−8/−6% across the A/B (machine warming), so drift-corrected the change is + **~+11–12%**. Final warm matrix (in-place on = production default): codec + **36.5** / standard 36.0 / markov **36.0** / unlimited 33.3 / boundary-kv 36.5 + p50 (mean skewed by frame-emit spikes) / turbo 21.2 (inherent) — **markov/codec + now AT parity with standard** (was 0.81× at the arc's start), the whole cached + cluster **~12% ahead of llama.cpp's 32.1**. Caveat: bench box was at ~58% + charging (not cool-dedicated); ordering + A/B *direction* are robust, absolutes + drifted ~5–8% run-to-run — a cool-box rerun would firm them. (NB: the first + engine in a fresh process eats the 30GB page-in — standard read 21.8 cold, + 34–36 warm; warm runs are the fair matrix.) +- **Propagated the in-place lever to the two remaining walk engines + faithfulness + audit 2026-06-14.** A full cross-engine spec/contract audit (all 9 engines vs + `state-policy.md`'s `(canonical, derivative, contract)` triple) found every + engine faithful, and flagged the two siblings still paying the O(L²) owned + concat the markov/codec in-place change eliminated: + - **boundary_per_layer (was the one NEEDS-FIX)** — carried NO `hot_kv` at all: + it `recompute_kv`'d the whole hot tier *and* rebuilt an owned `[ctx+1]` concat + every layer every step (worse than markov *pre*-W2). Added a `hot_kv` + derivative + the W2-cache + `run_attention_block_decode_step_auto_inplace` + steady state, mirroring its twin codec — only active in the `cache_eligible` + (unbounded, no cold) path, like codec; the windowed/cold path (its primary + purpose) is untouched. `hot_kv` is excluded from `memory_bytes` (droppable + derivative, matches markov). Engine-level in-place-vs-owned-concat A/B (q4k on) + bit-identical; f32-gated debug `hot_kv ≈ recompute_kv` assert. + - **windowed_checkpoint** — its CPU window walk (`extend.rs`) passed the whole + window K/V by value → backend re-concats `[n+1]` per layer per step (its own + doc admitted "O(window²) total"). Added `rs_extend_inplace` (appends into the + window's doubling-capacity buffer, attends over views), wired into + `extend_current` only when eligible (index + toggle + q4k); `replay_window` / + quant / executor / tests keep the owned concat. The engine's existing + `current_window_kv_len` counter already treated the buffers as over-allocated + (the dispatch path did), so `close_window`/`current_kv_bytes` needed no change. + A/B (q4k on) bit-identical; `resident_identity` flags-off still green. + Both reuse the shared `LARQL_MARKOV_INPLACE_KV` toggle + `Q4K_FLAG_ENV_LOCK`. + Also: **apollo footgun guard** — `injection_layer < crystal_layer` silently + no-ops the retrieval-injection (the compressed forward starts at `crystal`); + added a one-time runtime warning in `prepare_injection` (experimental engine, + warn-don't-fail). Doc-drift swept: boundary-kv spec now flags `resume` as + NOT-IMPLEMENTED (emit half only), apollo spec `KvEngine`→`RetrievalEngine`, + `state-policy.md` `fallback_mode` marked retired (per its own §8 resolution). + 760 kv tests green debug + opt, clippy clean. (Same caveat as above: turbo's + −39% is inherent; boundary-kv inherits standard's opts via resident forwarding.) + +Prefill stays on the f32 BLAS gemm for all engines deliberately (the task +#16 prefill falsification: q4k repeated-matvec loses ~20× to AMX at +prefill shapes). + +## Hardening — codebase review 2026-05-28 + +From the whole-codebase review ([`docs/audits/codebase-review-2026-05-28.md`](../../../docs/audits/codebase-review-2026-05-28.md)): + +- **P2 — CLI-supplied sizing params can reach prefill panics**; validate at the boundary. +- **P2 — positional QKVO contract** (`attn_data[1]/[2]`, shared with larql-models) is maintained by convention, not type. Silent-drift risk — consider a typed accessor. + ## [2026-05-20] — boundary_per_layer: bugfixes + W1-GPU dispatch + modular split **Engine bottleneck audit** (`PERFORMANCE.md` §"2026-05-20"). Findings @@ -307,3 +569,714 @@ The cut was clean: every primitive engines depend on (`ModelWeights`, `BackendFfn`, `WalkFfn`, `KvCache`, `forward_*`, `rms_norm_heads`, …) was already public in larql-inference, so this extraction did not require designing new API. + +## Earlier — entries migrated from ROADMAP.md (2026-08-04) + +These predate this file's adoption of Keep a Changelog dating and were +kept in the roadmap until the roadmap was split by tense. + +## Crate-shape state (2026-05-17) + +- Crate extracted from `larql-inference::engines` on 2026-05-09 — see + [`CHANGELOG.md`](CHANGELOG.md). +- **Seven engines shipped** as of 2026-05-17: + - Original four: `standard`, `no_cache`, `markov_residual`, + `windowed_checkpoint`, `turbo_quant`, `apollo`. + - Three new: `boundary_kv`, `markov_residual_codec`, `boundary_per_layer`. + Specs in `crates/larql-inference/docs/specs/`: + [boundary-kv-engine.md](../larql-inference/docs/specs/boundary-kv-engine.md), + [markov-residual-codec-engine.md](../larql-inference/docs/specs/markov-residual-codec-engine.md), + [boundary-per-layer-engine.md](../larql-inference/docs/specs/boundary-per-layer-engine.md). +- Consumers wired: + - `larql-cli bench --engine ` (selector dispatch) + - `larql-cli bench --via-executor` opts into the new `LayerExecutor` + surface; falls through to legacy path for unmigrated engines. + - in-crate `benches/engine_decode.rs` (criterion: dispatch helpers + Standard parity) +- Coverage policy: 90 % line coverage per source file (see + `coverage-policy.json`); CI gate at `make larql-kv-coverage-policy`. + Workspace `larql-kv` lib total: **95.62% lines, 95.43% regions, 95.50% + functions** (2026-05-24 evening, post coverage-debt clearance). + **All 61 files at ≥90% lines; debt baselines cleared from policy + file.** The 2026-05-24 push lifted the five `engines/*/dispatch.rs` + files (range 7.95–80.68% → 93.57–97.85%) and + `engines/markov_residual/compute.rs` (86.85→95.30%). See "Closed + (recent)" entry for the thread-local-override pattern that makes + the env-gated paths in `compute.rs` and the W10 mask cascade in + the dispatch files testable without process-env mutation. + +## Architectural cuts (2026-05-17) + +Substantive refactors landed; specs reflect the new boundaries. + +### Naming hygiene — renamed for honesty + +- **`metal_fused_prefill` / `metal_fused_decode_step`** → `fused_prefill` + / `fused_decode_step`. The "metal" was a lie — `CpuBackend` implements + `prefill_q4` and `decode_token` via its C Q4 kernel and also takes the + fused path on `--cpu`. The aliases in `windowed_checkpoint::engine` + (`quant_prefill_metal`, `quant_decode_token`) follow. +- **`KvEngine::prefill_q4k` / `decode_step_q4k`** → `prefill_quant` / + `decode_step_quant`. The `_q4k` suffix baked one format into the trait + surface; the trait is quant-agnostic (dispatches on `index`'s format). + Internals that are genuinely Q4K-specific (`prefill_q4k_moe`, + `cpu_q4k_cache_*`, `run_ffn_decode_step_q4k_direct`) keep their names. +- **`ComputeBackend::has_q4()` → `supports_quant(format: QuantFormat)`.** + Per-format predicate; `CpuBackend` reports support for `Q4_0`, `Q4_K`, + `Q4_KF`, `Q6_K`; `MetalBackend` adds `Q8_0`. Backends can advertise new + format support without trait extension. +- **Storage slots `q4k` → `kquant` for K-family fields.** `attn_q4k`, + `interleaved_q4k`, `set_attn_q4k`, `load_attn_q4k`, etc. — these hold + K-family quant bytes (Q4_K, Q4_KF, Q6_K — manifest tag picks). Q4_0 + (`attn_q4`) and Q8 (`attn_q8`) slots stay — genuinely format-specific. + +### Engine state vs execution — new abstraction + +Spec: [engine-state-vs-execution.md](../larql-inference/docs/specs/engine-state-vs-execution.md). + +The engines were re-coupling backend / FFN / format decisions into their +state-management code. The new shape: + +- **`LayerExecutor` trait** (in `larql-inference::layer_executor`) — + per-layer execution surface with `run_prefill_layer` / + `run_decode_layer` returning `(hidden, SharedKV)`. Dispatch kind + (`Fused` / `PerLayer`) is explicit. +- **`LocalWalkExecutor`** — wraps `run_attention_with_kv_backend` + + the caller's `&dyn FfnBackend`. The critical decoupling: the executor + does **not** construct its own `WalkFfn` — it uses whatever the engine + was handed. +- **Engine trait extension:** `KvEngine::prefill_via_executor`, + `decode_step_via_executor`, `prefill_quant_via_executor`, + `decode_step_quant_via_executor`. Default impls fall through to the + legacy methods so unmigrated engines work unchanged. + +### Engines on the new surface + +Every engine now runs its own state-policy code; there is no hidden +fall-through to the backend's fused kernel from per-layer engines. +`standard` (and by delegation `boundary_kv`) is the **only** engine +that exercises the fused fast path — via +`ComputeBackend::coarse_prefill` / `coarse_decode_step`, which on +Metal calls `larql_inference::vindex::fused_prefill`. + +| Engine | Default dispatch | `*_via_executor` override | Honors FFN backend | Tok/s (Gemma 3 4B Q4K, Metal) | Hot state | +|---|---|---|---|---:|---:| +| `standard` | `ComputeBackend::coarse_prefill` (fused fast path) | n/a (no per-layer code to migrate) | n/a | 104 | 0 MB (backend owns K/V) | +| `boundary_kv` | Delegates to `standard` + emits boundary frames | n/a | n/a | ≈104 | 0 MB | +| `markov_residual` | Per-layer walk via `rs_prefill_walk` | ✅ | ✅ counter test | 3.6 | 6.0 MB | +| `markov_residual_codec` | Per-layer walk via `rs_prefill_codec_walk` (bf16 cold) | ✅ | ✅ counter test | 4.3 | 6.0 MB | +| `windowed_checkpoint` | Windowed checkpoint extension via `process_q4k` | ✅ | ✅ counter test | 25.6 | 4.8 MB | +| `turbo_quant` | Per-layer WHT + Lloyd-Max compression cycle | ✅ | ✅ counter test | 3.9 | 0.6 MB | +| `boundary_per_layer` | Per-layer walk with per-layer codec policy | ✅ (dense) | ✅ counter test | — | matches markov_residual_codec | +| `apollo` | Whole-forward through `forward_layer_range` (boundary prefix + perturb) | ✅ | ✅ counter test | requires store | scales with store | +| `no_cache` | Full re-forward per step (O(N²) wall-time) | ✅ | ✅ already did on legacy `prefill` | — | token list only | + +## Closed (recent) + +- **2026-05-24 — Multi-modal engine seam (ADR-0023).** `KvEngine` gains + `supports_multimodal()` (default false) + `prefill_from_hidden(weights, + ffn, initial_hidden: &Array2) -> Result, EngineError>`. + `StandardEngine` is the first (and currently only) MM-capable engine. + Other engines inherit the default-false convention — they remain + text-only until each individually implements the new method. + `AnyEngine` forwards both methods. `generate_with_engine_from_hidden` + wrapper shares the decode loop with `generate_with_engine`. Dispatch + helpers `kv_prefill_from_hidden_via_dispatch` (sync + async) hoist the + embed step out of the prefill loop so both text-only and MM inputs + follow the same layer-forward path. The eventual end state: every + engine implements `prefill_from_hidden` and `prefill(token_ids)` becomes + a thin wrapper. No timeline on the seven-engine migration. + +- **2026-05-24 — Sibling trait extraction LANDED.** `KvEngine` + `Option` returns are gone; the typed `EngineError` enum lives in + `larql-inference::kv_engine` alongside the new `RetrievalEngine` + trait + `AnyEngine` dispatch enum. The two-harness silent-drop / + panic disagreement (`accuracy_suite/runner.rs` vs + `bench/engine_runtime.rs`) is resolved at the type level. + + **Trait surface:** all 8 `KvEngine` impls (`standard`, `no_cache`, + `markov_residual`, `markov_residual_codec`, `windowed_checkpoint`, + `turbo_quant`, `boundary_kv`, `boundary_per_layer`) return + `Result, EngineError>` on `prefill` / `decode_step` / + `*_quant` / `*_via_executor`. Apollo moves to the new + `RetrievalEngine` trait (`prefill(weights, token_ids)` / + `decode_step(weights, token_id)` — no `FfnBackend`, no per-step K/V). + + **EngineError variants** (exhaustive, no `#[non_exhaustive]`, + thiserror): `EmptyPrompt`, `BackendUnavailable`, `RetrievalMiss + { reason }`, `InvariantViolation { what }`, `BackendFailure + { details }`. Per Finding 2, `InvariantViolation` and `BackendFailure` + are kept as two top-level variants to preserve the alert-routing + distinction (a dispatch bug vs a kernel/data failure). The accuracy + harness's `ScoreOutcome` mirror followed suit: + `SkippedInternalError` → `SkippedInvariantViolation` + + `SkippedBackendFailure` (load-bearing JSON schema change for + downstream observability). + + **AnyEngine** (`AnyEngine::Kv(Box) | + Retrieval(Box)`) is the harness boundary type. + Forwarding methods (`prefill` / `decode_step` / `prefill_quant` / + `decode_step_quant` / `*_via_executor`) take the superset of args + from both surfaces and ignore the irrelevant ones on the retrieval + arm. This intentionally walks back the original "don't lift a common + shape" plan — the harness scalability won out, since the alternative + is N×2 match arms per call site as more retrieval engines land. + + **Bench harness merged.** `run_engine` + `run_engine_q4k` collapsed + into one `run_engine(weights, index: Option<&VectorIndex>, ...)`. + When `index = Some` the dispatch goes through `prefill_quant` + (quant-agnostic — the vindex's format flows through the engine); + when `None` the dense `prefill` path runs. FFN selection: dense + defaults to `WeightFfn`, quant defaults to `NullFfn` (preserves the + pre-merge Q4K behaviour). `--ffn-policy` honored on dense, logged + as not-yet-honored on quant due to the `&mut weights` vs + `&weights`-borrowing-router conflict (unchanged from pre-merge). + + **Coverage debt:** one re-introduced baseline at + `markov_residual/engine.rs` (89.5% vs 90% floor). The remaining + uncovered lines are all `.ok_or_else(|| BackendFailure)?` + constructions that only fire when an internal helper + (`rs_decode_step_walk`, `recompute_kv`, `executor.run_*_layer`) + returns None. Triggering those requires the mock `EngineBackend` + infrastructure that the 2026-05-24 coverage-clearance explicitly + deferred; the baseline tracks the debt rather than gold-plating + ahead of need. + + **Outcomes.** Test count larql-kv lib: 712 → 726 (+14). Workspace + builds clean. `make larql-kv-ci` passes (fmt + clippy + tests + + fresh coverage policy with 1 baseline). Apollo's `executor.rs` + deleted (~150 lines of dead code from the old KvEngine `*_via_executor` + impls). Closes [`docs/state-policy.md`](docs/state-policy.md) §8 + Open Question 1 ("Where does Apollo's fallback live?"); also closes + the interim `ffn_backend` JSON limitation flagged in Item 1 of the + 2026-05-24 accuracy harness work. + + **Follow-ups** *(deferred to keep this PR atomic)*: + - Mode 5 / Graph-Grounded engine lands as a `RetrievalEngine` impl + (was blocked on this refactor). + - Q4K `--ffn-policy` honoring (was waiting on the same + `&mut weights` borrow conflict — still present after the merge + because the trait surface still takes `&mut weights` for lazy + dequant). + - `RemoteWalk` build path (~200 lines, standalone — was the second + blocked item). + - `markov_residual/engine.rs` coverage debt + mock `EngineBackend` + infrastructure (deferred per "Sub-project A" of the previous + coverage push). + +- **2026-05-24 — Coverage debt CLEARED.** All six files below the + 90% per-file floor lifted; `make larql-kv-coverage-policy` passes + against fresh `summary.json` regeneration. Workspace total 95.62% + lines, 61/61 files at ≥90%, 0 debt baselines remaining. + + Files lifted (pre → post): `turbo_quant/dispatch` 9.35→97.85%, + `boundary_per_layer/dispatch` 7.95→93.57%, `windowed_checkpoint/dispatch` + 59.09→97.24%, `markov_residual/dispatch` 77.51→96.78%, + `markov_residual_codec/dispatch` 80.68→97.72%, + `markov_residual/compute` 86.85→95.30%. + + Approach inverted both pre-baked design assumptions: + - **No new shared mock `EngineBackend`** — `CpuBackend` (via + `cpu_engine_backend()`) already implements `coarse_*_with_state` + when driven against the synthetic Q4K fixture + (`make_test_q4k_weights` + `make_test_q4k_vindex`), so every + dispatch happy-path tested end-to-end without new infrastructure. + - **No `serial_test` crate** — env-gated paths + (`LARQL_MARKOV_WALK_KV_*`, `LARQL_W10_DISABLE`) instead gained + a per-thread `RefCell` override that production helpers consult + *before* `std::env::var`. Tests inject without touching the + process env; no race with other parallel tests. New helpers: + `compute.rs::set_markov_env_override(...)`, + `engines/mod.rs::set_w10_disabled_override(...)` (both + `#[cfg(test)]` only). + + Test deltas: larql-kv lib 663 → 712 (+49). Zero regressions + (5/5 successive `cargo test -p larql-kv --lib` runs green after + the thread-local override fix; pre-fix the env-var-setting tests + produced flaky `cold_kv.is_some()` failures in unrelated codec + tests via process-env race). `make larql-kv-ci` passes end-to-end. + +- **2026-05-24 — Accuracy harness honesty + FFN policy cross-product + LANDED.** Multi-PR arc that turns the accuracy suite from "silent + drop on engine miss" into a discriminating cross-product harness: + + - **Item 1 — accuracy schema fix** (commit `07684457`). + `ScoreOutcome` enum (exhaustive, flat-tagged serde, mirrors the + future `EngineError` taxonomy). `PromptScore` / `ConflictScore` + gain `outcome` field + `Option` score payload with + `served()` / `skipped()` constructors enforcing + correlated-optionality. `StrategySplit` gains `*_served` + + `*_served_rate` per axis as required-companion fields to + `*_match_rate`. `compute_strategy_split` filters on served subset + (counting skips as zero would punish honest reporting). Replaces + `filter_map` silent-drop in all three drivers. Surfaces Apollo's + store-miss rows as `SkippedRetrievalMiss` instead of dropping. + `EngineKind::supported_names()` replaces hard-coded six-engine + error string at two bench sites. + + - **Item 2 v0 — `FfnBackendKind` parser + `FfnLayerPolicy` + (in `larql-inference::ffn_policy/`).** New crate-shape: + `FfnBackendKind` (Dense / Walk{k} / RemoteWalk / Null), + `RoutingPredicate` (All / Layers / Otherwise), `FfnLayerPolicy` + with from_spec parser supporting per-layer routing + (`{walk:k=100}@layers=14-27;{dense}@otherwise`). + Construction-errors on overlapping ranges; exhaustive enums; + typed error taxonomy (`PolicyParseError` / + `PolicyValidationError`). Module lives in `larql-inference` not + `larql-kv` — FFN policy is the FFN axis, not the KV axis. + + - **`build_router` slice — `ValidatedFfnLayerPolicy` newtype + + `BoundFfnRouter`.** Type-system enforcement of "validate before + build" via non-public constructor. `BoundFfnRouter<'a>` owns its + backend instances (`Vec>`) so callers + don't manage backend lifetimes alongside the router's. `impl + FfnBackend for BoundFfnRouter` delegates per-layer via the + trait's existing `layer: usize` parameter — drop-in for the + `&dyn FfnBackend` surface every engine already takes. Design + rationale: `larql-inference/docs/ffn-build-router.md`. + + - **Cross-product harness + typed axis columns.** `accuracy_cmd` + iterates `kv_engine × ffn_backend` cross-product via + `FfnLayerPolicy::split_specs` (comma-separated, brace-aware, + re-parse fallback for kv-comma forms like + `remote-walk:endpoint=X,wire=Y`). New `EvalLabels<'a>` struct + bundles `(kv_engine, ffn_backend, strategy)` for clean signatures. + `PromptScore` / `ConflictScore` / `StrategySplit` gain explicit + `kv_engine: String` + `ffn_backend: String` columns alongside + `strategy`. `format_strategy_split` grows a two-axis layout + (`KV engine` + `FFN backend` columns) when any row has + `ffn_backend != "dense"`; default no-`--ffn` runs keep the + historical single-`Strategy`-column layout. Closes the + interim-`ffn_backend`-as-user-input limitation noted in Item 1's + ROADMAP entry. + + - **CLI wiring.** `larql accuracy --ffn dense,walk:k=100,'{walk:k=100}@layers=14-27;{dense}@otherwise'` + now runs the cross-product in one invocation. Vindex loaded + lazily — only when a Walk binding is present. + `larql bench --ffn-policy ` honors the policy on the + non-Q4K (CPU) path; Q4K path accepts the flag but doesn't + honor it yet (P1 follow-on above). + + - **Apollo into accuracy default engines.** `--engines` default + now includes `apollo`. The schema fix above means Apollo's + store-miss rows show `served_rate < 1.0` rather than silent + drops — diagnostic rather than misleading. + + - **Module splits.** `accuracy_suite/runner.rs` (2050 lines) split + into `accuracy_suite/runner/` folder (6 files: `types` / + `scoring` / `drivers` / `aggregate` / `legacy` / `mod`). Same + pattern that produced the `ffn_policy/` folder split in + `larql-inference`. + + - **Coverage lift across 5 engine files.** Pre-existing engine + internals had drifted below 90%. Lifted with synthetic-weights + + CPU-backend tests: `boundary_per_layer/cold_tier.rs` + (88→100%), `executor.rs` (85→90.6%), `walk.rs` (84→95%), + `engine.rs` (83→90%), `markov_residual/store.rs` (86→99.6%). + `markov_residual/compute.rs` partially lifted (81→86.85%); + full lift gated on `serial_test` for env-var paths. + Discovered the gate had been passing against a stale JSON — + fresh `make larql-kv-coverage-summary` is now required to + surface debt. See "Coverage debt" section above for the + remaining 6 files. + + Test deltas across the arc: larql-kv lib 595 → 663 (+68), + larql-inference lib 1086 → 1102 (+16). Zero regressions. Clippy + clean. Aggregate ~3,500 lines of code + tests added across + `larql-kv` and `larql-inference`. + + ROADMAP entry for the sibling trait extraction (P0 above) + references "Item 1 in the conversational priority queue" — Item 1 + is the schema fix above. Mode 5 work is still gated on that P0 + refactor landing. + +- **2026-05-18 — W8.2 (doubling-capacity K/V in `markov_residual` + + `markov_residual_codec`) LANDED: 2.4× decode speedup at 1000 tokens.** + Lifted the W8 pre-allocation pattern from `windowed_checkpoint` to the + two unbounded-window engines. Since `max_window=None` rules out a + fixed pre-alloc, both stores now use a doubling-capacity strategy + via three private helpers in each engine: + - `window_capacity(prompt_len, window_size)` — initial cap is + `max(window, prompt_len)` if windowed, else + `max(prompt_len * 2, 64)`. + - `grow_capacity_2d(src, len, cap)` — allocate `[cap, cols]` once + at prefill, copy the prefill rows in. + - `append_row(buf, row, len)` — in-place `slice_mut(s![len..len+1, + ..]).assign(row)` when `len < cap`; otherwise double capacity, + copy the live rows, then assign. Amortised O(1) per append vs the + O(n) per step the previous `Array2::zeros((n+1, dim))` pattern + paid. + + Store changes (both `RsStore` and `RsStoreCodec`): + - New `pub hot_len: usize` field — logical row count, separate from + `stored[l].shape()[0]` (which is now capacity ≥ hot_len). + - `window_tokens()`, `memory_bytes()`, `clip_layer` / + `clip_layer_overflow` updated to use `hot_len`. + - New `finalise_hot_len_after_clip()` — must be called after every + per-layer clip loop. (Subtle bug fix during impl: setting + `hot_len = window` *inside* the per-layer loop made layers 2..N + see `rows == window` and skip their clips, dropping half the + cold-tier payload. Two existing tests caught this.) + + Bench (Gemma 3 4B Q4K, Metal, M3 Max): + - **1000-tok**: + - `markov-rs`: 24.8 → **58.7 tok/s (+137%)** + - `markov-rs-codec`: 25.7 → **57.2 tok/s (+123%)** + - `windowed-checkpoint`: 49.5 → **57.4 tok/s (+16%)** (variance + recovery from previous run + sympathy from the codepath audit) + - `standard` unchanged at 64.1 (untouched) + - **50-tok**: + - `markov-rs`: 77.1 → **88.9 tok/s (+15%)** + - `markov-rs-codec`: 77.5 → **88.8 tok/s (+15%)** + + All three cached-state engines now cluster within 11% of standard's + 64.1 tok/s ceiling at 1000 tokens. The doubling-capacity scales + linearly with seq_len: at 50 tok the saved alloc bytes are small + (~400 KB/step); at 1000 tok they're ~8 MB/step. The 137% win at + long context is the alloc churn that pre-W8.2 was hiding behind + prefill cost. + + CPU walk + executor fallback paths (`rs_decode_step_walk`, + `rs_decode_step_codec_walk`, `process_via_executor`) still allocate + per step — they're not on the hot path for the bench. Defensive + consistency: every legacy RsStore/RsStoreCodec constructor sets + `hot_len` from `stored[0].shape()[0]` so non-dispatch paths see a + consistent invariant. + +- **2026-05-18 — Step 9 (iterative Metal `coarse_prefill_with_state`) + LANDED: ~10× prefill speedup on every state-dump engine.** + Pre-Step 9, `MetalBackend::coarse_prefill_with_state` defaulted to + the trait's `coarse_prefill` (no per-layer state dump); engines saw + `state.is_complete_for() == false` and fell back to the CPU walk + (~2.7 s on Gemma 3 4B). The new impl pre-allocates `[seq_len, + hidden]` and `[seq_len, kv_dim]` per layer (W8-style alloc at + source for prefill too), resets + preallocates the Metal K/V cache, + then iterates `fused_decode_step_with_state` per prefill token, + writing the dump into the pre-allocated row position. + + Bench (Gemma 3 4B Q4K, Metal, M3 Max, "The capital of France is", + 5 prefill tokens): + - `markov-rs` prefill: 2757 → **254 ms** (10.9×) + - `markov-rs-codec` prefill: 2564 → **249 ms** (10.3×) + - `windowed-checkpoint` prefill: 2760 → **256 ms** (10.8×) + - `turbo-quant` prefill: 2750 → **334 ms** (8.2×) + + Predicted ~45× (5 × 12 ms decode time) didn't materialise because + each iterative `fused_decode_step_with_state` carries per-token + state-dump readback overhead. Remaining ~250 ms is 5 × ~50 ms + per-iter + fixed setup. Further closure needs a single-kernel + prefill that dumps state for all positions in one shot — separate + Metal-kernel surgery. + + Decode steady-state also moved (W8 + Step 9 compound): + - `windowed-checkpoint`: 82.7 → **89.2 tok/s** (fastest cached-state + engine; within 10% of `standard`'s 99.2 ceiling) + - `markov-rs`: 75.3 → 77.1 tok/s + - `markov-rs-codec`: 79.0 → 77.5 tok/s + +- **2026-05-18 — W8 (pre-allocated K/V buffer in `windowed_checkpoint`) + LANDED: 58% of decode-CPU alloc churn removed.** + samply flamegraph on `windowed_checkpoint:window=1024 --tokens 1000` + (post-W7) surfaced an unexpected hot path: 21% `__bzero` + 19% + `ndarray::zip_mut_with_same_shape` + 18% `madvise` = **58.5% of + main-thread CPU** spent on `Array2::::zeros((n+1, kv_dim))` + + `slice_mut().assign(k_old)` + `slice_mut().assign(k_new_row)` + inside `decode_step_via_dispatch` — 68 allocations per token + (34 layers × 2), each growing linearly with `n`. + + Fix: pre-allocate `Array2::zeros((window_size, kv_dim))` per layer + once at prefill (in `try_prefill_via_dispatch`), track a single + `current_window_kv_len: usize` counter, and append in the hot path + via `slot.0.slice_mut(s![pos..pos+1, ..]).assign(k_new_row)`. One + small `kv_dim`-sized copy per layer per side, zero alloc per step. + Readers (`close_window`, `current_kv_bytes`) updated to use the + counter instead of `k.shape()[0]`; CPU walk fallback paths set the + counter defensively from the returned narrow-array shape. + + Bench (Gemma 3 4B Q4K, Metal, M3 Max): + - 50-tok: `windowed-checkpoint:window=256` 82.7 → **86.6 tok/s + (+4.7%)** vs `standard`'s 99.4 (gap closed ~50%) + - 1000-tok: `windowed-checkpoint:window=1024` 17.39 ms vs `standard`'s + 15.74 ms → 1.65 ms gap (vs pre-W8 estimated 5-10 ms slope from + `Array2::zeros((n+1, …))` growing linearly with `n`) + + Post-W8 flamegraph: the `__bzero` / `zip_mut_with_same_shape` / + `madvise` triple is **gone from the top-20**. Remaining main-thread + CPU is dominated by `__psynch_cvwait` (Metal GPU wait, + irreducible), `synthesize_lm_head_kquant` (prefill — separate + ~2.5 s regression flagged elsewhere), and generic `Map::fold`. + + The optimisation is engine-local (`larql-kv/src/engines/windowed_checkpoint/engine.rs`) + with no surface change. Same pattern can be lifted to + `markov_residual` / `markov_residual_codec` / `turbo_quant` once + their state-policy shape is clarified — they use the same + `Array2::zeros((n+1, kv_dim))` pattern but have unbounded windows + by default, so the pre-allocation needs a growable strategy + (doubling-capacity Vec-style) rather than fixed window size. + Tracked as W8.2 candidate. + +- **2026-05-18 — W7 (blit-encoder fusion) LANDED: per-layer commit + overhead removed; +30-48% across cached-state engines.** + Modified `decode_token_with_moe_split_fn` in + `larql-compute-metal/src/decode/mod.rs` to pre-allocate per-layer + staging buffers (k / v / h-in) when `state_dump` is `Some`. The + layer loop blits `k_out` / `v_out` / `h_buf` into the staging + buffers inside the same command buffer (`new_blit_command_encoder` + + `copy_from_buffer`) instead of forcing per-layer commit + wait + + CPU read. The single final commit at the bottom of the function + flushes everything; reads happen once after that, draining staging + into `state_dump`. Metal's command-buffer encode ordering + guarantees blit reads see the settled compute writes. + + Measured (Gemma 3 4B Q4K, Metal, M3 Max): + - `standard` (control, no state_dump): 105.9 → 99.4 tok/s (noise) + - `markov-rs`: 58.0 → **75.3 tok/s (+30%)** + - `markov-rs-codec`: 58.4 → **79.0 tok/s (+35%)** + - `windowed-checkpoint` (window=256): 56.0 → **82.7 tok/s (+48%)** + - `turbo-quant` (4-bit, 10-tok bench): 33.0 → **37.7 tok/s (+14%)** + + Engine-cost decomposition post-W7: ~10 ms Metal kernel compute + + ~3 ms CPU glue. The remaining gap to `standard`'s 99 tok/s is + pure CPU-side state-update work (state Vec→Array2 conversion, + appends). Closure path: in-place state updates / pre-allocated + buffers (W8 candidate). + + Edge cases worth noting: + - `standard` doesn't touch state_dump → blit branch is dead code + → 0× regression confirmed. + - `turbo_quant`'s codec inner loop is the dominant per-token cost; + the saved 1.7 ms commit overhead is a smaller fraction. + - The `windowed_checkpoint` +48% win reflects its lighter post- + kernel CPU work (just append to `current_window_kv`); engines + with heavier post-kernel work see smaller relative gains. + +- **2026-05-17 night — W1-GPU steps 4 + 6 LANDED: windowed_checkpoint + + turbo_quant now route through dispatch on Metal.** + Same pattern as steps 5: each engine gains `try_prefill_via_dispatch` + / `decode_step_via_dispatch` helpers that read per-layer captured + state and update engine-specific state policy. + - **turbo_quant**: state.k_new/v_new per layer feeds the + WHT+Lloyd-Max codec via `CompressedLayer::compress` (prefill) + and decompress→append→recompress (decode). Bench: **19.6 → + 33.0 tok/s (+68%)** on Metal. Memory stays at 0.6 MB hot + (compression intact). + - **windowed_checkpoint**: state.k_new/v_new appends to + `current_window_kv` per layer; window auto-close at + `window_size` tokens fires the legacy `close_window` checkpoint + emit. Bench: **28 → 56.0 tok/s on Metal (+98%)** at + `window=256` (Gemma 3 4B, M3 Max, 50-token decode). Hot state + 15.7 MB tracks the engine-side window shadow (see KvHandle + eviction note below). + + Engine memory note: with W1-GPU active, the backend's internal K/V + cache grows unboundedly alongside each engine's shadow state. This + defeats the memory benefit of `windowed_checkpoint` / + `markov_residual_codec` at long contexts. Follow-up: expose a + `KvHandle::evict_oldest(n)` API on `KvDispatch` so engines can + bound the backend cache to match their window. +- **2026-05-17 night — W1-GPU step 2 LANDED: Metal per-layer state + dump → 2.1× decode speedup on markov-rs + codec.** + Modified `decode_token_with_moe_split_fn` in + `larql-compute-metal/src/decode/mod.rs` to accept an optional + `state_dump: Option<&mut DecodeStateDump>` parameter. When active, + the layer loop: + 1. At top of layer L: pushes `x` (for L=0) or reads `h_buf` (for + L>0, settled by the previous layer's commit) into + `state.h_in_per_layer`. + 2. At bottom of layer L: forces `enc.end_encoding()`, `cmd.commit()`, + `wait_until_completed()`, reads `k_out` / `v_out` (scratch + buffers reused across layers) into + `state.k_new_per_layer` / `v_new_per_layer`, then restarts + command buffer + encoder for the next layer. + + Trait wiring: new `DecodeBackend::decode_token_with_state_dump` + method (default falls back to plain `decode_token`); MetalBackend's + trait impl routes through the new kernel function when `state` is + `Some`. Inference layer adds `fused_decode_step_with_state` + + `MetalBackend::coarse_decode_step_with_state` / + `coarse_prefill_with_state`. Engines (markov_residual, codec) + inherit the Metal acceleration automatically — no engine-side + changes from step 5. + + Measured (Gemma 3 4B Q4K, Metal, M3 Max, 10-token decode): + - `markov-rs`: 27.0 → **57.7 tok/s** (+114%) + - `markov-rs-codec`: 27.8 → **57.5 tok/s** (+107%) + - `standard` (fused control): 100.8 tok/s (unchanged) + + Per-token cost: ~17 ms = 10 ms Metal compute + ~1.7 ms commit + overhead (50 µs × 34 layers) + ~5 ms engine state update / CPU + glue. The remaining gap to standard's 100 tok/s is the + per-layer commit cost; a follow-up could use blit-encoder + switches inside a single command buffer to eliminate the + commit overhead and lift toward 80-100 tok/s. + + Prefill cost: ~2.8 s on Metal (CPU walk for state seeding + + Metal `fused_prefill` for backend cache). One-shot; doesn't + affect decode steady-state. Future optimisation: per-position + per-layer K/V dump on the Metal prefill side to skip CPU walk. +- **2026-05-17 night — W1-GPU infrastructure (decode trait surface + + CPU impl + engine wiring; Metal kernel modification deferred).** + Three layered changes landed end-to-end: + - **Trait surface (`KvDispatch`):** new `coarse_prefill_with_state` / + `coarse_decode_step_with_state` methods take + `Option<&mut PerLayerDecodeState>`. Default impls delegate to the + non-state variants, so unmigrated backends keep working. + - **`DecodeBackend` trait + `DecodeStateDump` struct** added in + `larql-compute` for the substrate-level surface. Same default- + delegation pattern. + - **CPU implementation** (`predict_kquant_prefill_with_state` / + `predict_kquant_decode_step_direct_with_state`): threads per-layer + state capture through the existing per-layer walk at zero + re-compute cost. Parity test in + `kv_dispatch::cpu::coarse_decode_step_with_state_populates_and_matches_plain` + asserts cached and non-cached outputs match within f32 rounding + and per-layer shapes (`[1, hidden]`, `[1, kv_dim]`) are correct. + - **Engine wiring** for `markov_residual` and + `markov_residual_codec`: `try_prefill_via_dispatch` / + `decode_step_via_dispatch` route through the new + `coarse_*_with_state` API when the backend implements it. State + capture feeds `RsStore::stored` (residuals) and `hot_kv` (W2 + cache) in a single backend call. Legacy walk path stays as the + fallback when state isn't populated (e.g. on backends that + haven't migrated yet — currently `MetalBackend`). Gated on + `supports_direct_matvec_decode` so non-Q4K test fixtures skip + the dispatch path. 113 markov tests pass. + - **CPU bench numbers stay parity** post-W1-GPU step 5: + markov-rs 27.4 tok/s, codec 26.6 tok/s — same as W2 (W1-GPU on + CPU just changes the code path, not the compute; CPU was already + at the M3 Max compute ceiling). + + **What's NOT done**: `MetalBackend::coarse_*_with_state` still uses + the default delegation (state stays empty), so engine falls back + to walk on Metal — no GPU speedup yet. The real Metal acceleration + requires modifying + `larql-compute-metal::decode::decode_token_with_moe_split_fn` + (200+ lines) to thread per-layer dump buffers + blit-encode steps + into the existing single command buffer. Two implementation + shapes have been scoped: + 1. **Blit-encoder switches per layer**: cheapest in steady-state + (~tens of µs per layer); requires careful encoder lifecycle + management within the existing kernel function. + 2. **Per-layer commit + CPU readback**: simpler (mirror the + existing `stage_timing_split` pattern); costs ~50µs/layer × + 34 = ~1.7ms/token overhead. Projected ceiling: 50-80 tok/s + (vs CPU's 27 tok/s ceiling, vs `standard`'s 102 tok/s fused). + + Choice between shapes is open. The trait surface, CPU impl, and + engine wiring are all stable and don't change regardless of which + Metal-side approach lands. +- **2026-05-17 night — W2: hot K/V cache for `markov_residual` and + `markov_residual_codec`.** Added `hot_kv: Option>` + to both `RsStore` and `RsStoreCodec`; prefill captures K/V from + the per-layer forward pass (previously discarded) and stashes it; + decode appends one row per layer via the existing + `run_attention_block_decode_step_backend` return tuple. On + window-overflow `clip_layer` slices `hot_kv` consistently with + `stored`; for `markov_residual` (lossless cold tier) the evicted + K/V rows merge directly into `cold_kv` (no `recompute_kv` call + needed); for `markov_residual_codec` (lossy bf16 cold tier) + `cold_kv` is invalidated on overflow so the next step recomputes + against the codec-decoded residual. Bench: `markov_residual` + 4.7 → 26.8 tok/s (5.7×); `markov_residual_codec` 5.0 → 27.5 tok/s + (5.5×). Both now sit on the `windowed_checkpoint` curve. Engine + contract preserved — drop `hot_kv` and the next step recomputes + from `stored` (via_executor path takes this fallback). Hot-state + memory grew from 5.3 → 10.8 MB; still ~50× smaller than + `standard`'s full KV cache. Parity test + (`decode_step_quant_w2_cached_matches_recompute_from_residuals`) + asserts the cached and recompute paths agree within fp rounding. +- **2026-05-17 night — W7: per-engine profiler wired on the quant + path.** `EngineProfiler` now populates from `rs_decode_step_walk` + (markov_residual), `rs_decode_step_codec_walk` + (markov_residual_codec), `rs_extend_from_checkpoint_quant` + (windowed_checkpoint), and `decode_step_quant_cpu` (turbo_quant). + Each engine's `stage_summary()` returns `Some(...)` when + `with_profiling(true)` is set. `larql bench --profile --engine + ` now produces a per-stage attribution table per engine. + First measurement run produced the bottleneck-diagnosis table in + the P0 section above, which inverted two of the pre-profile + guesses: codec overhead in turbo_quant was ~25% not ~80%, and K/V + recompute (W2 target) was the dominant cost on markov_residual + (~80%) not dispatch (W1 target). Sequencing in P0 revised + accordingly. +- **2026-05-17 night — `_q4k` → `_quant` on remaining internal + function names.** The trait-surface renames earlier today + (`prefill_q4k` → `prefill_quant`, `has_q4` → + `supports_quant(format)`, `q4k` → `kquant` storage) missed the + per-engine implementation wrappers: + `windowed_checkpoint::process_q4k`, + `windowed_checkpoint::extend_current_q4k`, + `extend::rs_extend_from_checkpoint_q4k`, + `turbo_quant::decode_step_q4k_cpu` / + `turbo_quant::prefill_kquant_cpu`. All renamed to `_quant` since + they dispatch on whatever format the vindex carries, not Q4_K + specifically. +- **2026-05-17 night — Fused-bypass strip: engines are now engines.** + Every per-layer engine (`markov_residual`, `markov_residual_codec`, + `windowed_checkpoint`, `turbo_quant`) had a hidden + `if let Some(h) = fused_prefill(...) { return Some(h); }` short- + circuit at the top of `prefill_quant` / `decode_step_quant`. The + short-circuit meant `--engine markov-rs` on Metal silently ran + `StandardEngine`'s fused kernel instead — five engines tied at + ~103 tok/s with `hot=0.0MB`, masking every state-policy difference + and making per-layer optimization invisible. Cut: removed every + short-circuit; deleted dead `metal_prefill_done` + `force_walk` + fields and `with_force_walk` builders; dropped the pub(crate) + `fused_prefill`/`fused_decode_step` re-exports from + `windowed_checkpoint::engine` (only `StandardEngine::coarse_prefill` + uses the underlying `larql_inference::vindex::fused_prefill` now, + via `ComputeBackend::coarse_prefill`). `StandardEngine` remains the + default engine and the only home of the fused fast path. Bench now + reports honest numbers: standard 104 tok/s, markov-rs 3.6, codec + 4.3, windowed-checkpoint 25.6, turbo-quant 3.9 — every per-layer + engine reports non-zero `hot=` memory because their state + structures actually materialise. The 25-30× standard-vs-per-layer + gap is the new optimization frontier; previously it was invisible + because every engine was running the same kernel under different + labels. +- **2026-05-17 evening — Phase-2 migration completed for the remaining + three engines.** `windowed_checkpoint`, `turbo_quant`, and `apollo` all + override `*_via_executor` methods and honor the caller-supplied + `FfnBackend`. `CountingFfn` stub tests prove per-(token, layer) + dispatch through the caller's backend. Same push cleared every + `coverage-policy.json` debt baseline: all 43 files in src/ at ≥90% + lines, workspace total 95.55%. `larql bench --ffn http://shard:8080` + now routes through the remote shard for every per-layer engine + instead of silently constructing a local `WalkFfn`. +- **2026-05-17 — Phase 2 engine migration to `LayerExecutor`.** Four + engines (`markov_residual`, `markov_residual_codec`, + `boundary_per_layer`, `no_cache`) override `*_via_executor` methods. + They drive per-layer dispatch through `executor.run_*_layer` and + honor the caller's `FfnBackend`. `CountingFfn` stub tests prove the + FFN parameter is no longer silently ignored. Bench has + `--via-executor` flag; demoed on Gemma 3 4B Q4K showing the codec + engine's 50% cold tier saving (22.9 MB → 11.5 MB). +- **2026-05-17 — `LayerExecutor` trait + `LocalWalkExecutor`.** New + abstraction in `larql-inference::layer_executor` separating state + policy (engine concern) from execution strategy (executor concern). + Spec at + [engine-state-vs-execution.md](../larql-inference/docs/specs/engine-state-vs-execution.md). +- **2026-05-17 — `q4k` → `kquant` storage rename.** K-family storage + slots (`attn_q4k`, `interleaved_q4k`, manifests, setters, loaders) + renamed for consistency with accessor naming (`attn_kquant_layer_data`). + Q4_0 and Q8 slots unchanged. ~60 sites touched. +- **2026-05-17 — `has_q4()` → `supports_quant(format)`.** Per-format + predicate on `ComputeBackend`. 79 call sites migrated to + `supports_quant(QuantFormat::Q4_K)`. Enables future Q6_K / FP4 + fused-pipeline backends without trait extension. +- **2026-05-17 — `KvEngine::prefill_q4k` / `decode_step_q4k` → + `prefill_quant` / `decode_step_quant`.** Trait surface naming made + quant-agnostic. 112 sites updated. Internals that are genuinely + Q4K-specific kept their names. +- **2026-05-17 — `metal_fused_*` → `fused_*` rename.** The "metal" + prefix was a lie: `CpuBackend` implements `prefill_q4` and + `decode_token` via its C Q4 kernel. Aliases in + `windowed_checkpoint::engine` follow. +- **2026-05-17 — `BoundaryKvEngine`, `MarkovResidualCodecEngine`, + `BoundaryPerLayerEngine` shipped.** All three new engines have + contracts in `crates/larql-inference/docs/specs/`. Per-file coverage + ≥94 % lines on every new file. Bench demoed end-to-end on Gemma 3 4B, + Gemma 4 E2B, 26B-A4B, 31B, Qwen3 0.6B (dense + Q4K). +- **2026-05-09 — Initial extraction.** `engines/` carved out of + `larql-inference` into the new `larql-kv` crate. ~5,540 LOC moved with + no semantic changes. All four engines + `KvEngine` + accuracy / + profiler helpers now ship from this crate. diff --git a/crates/larql-kv/Cargo.toml b/crates/larql-kv/Cargo.toml index 2c4774bc2..2af2c7ce4 100644 --- a/crates/larql-kv/Cargo.toml +++ b/crates/larql-kv/Cargo.toml @@ -10,6 +10,7 @@ description = "Pluggable KV-cache engines for larql-inference (markov-rs, unlimi [dependencies] larql-inference = { path = "../larql-inference" } +larql-execution = { path = "../larql-execution" } larql-compute = { path = "../larql-compute" } larql-compute-metal = { path = "../larql-compute-metal", optional = true } larql-vindex = { path = "../larql-vindex" } @@ -55,5 +56,3 @@ rand_distr = "0.4" name = "engine_decode" harness = false -[[example]] -name = "engine_ladder" diff --git a/crates/larql-kv/PERFORMANCE.md b/crates/larql-kv/PERFORMANCE.md index 462ff10c6..3ffd37746 100644 --- a/crates/larql-kv/PERFORMANCE.md +++ b/crates/larql-kv/PERFORMANCE.md @@ -22,11 +22,11 @@ Per-engine impact, 50-token decode on Gemma 3 4B Q4K, M3 Max: | `markov-rs` | 87.1 | **98.0** | +12.5% | | `markov-rs-codec` | 86.6 | **98.1** | +13.3% | | `boundary-per-layer` (windowless) | 86.9 | **98.7** | +13.6% | -| `unlimited-context:window=256` | 86.1 | 94.2 | +9.4% (HOnly only) | +| `windowed-checkpoint:window=256` | 86.1 | 94.2 | +9.4% (HOnly only) | | `turbo-quant:bits=4` | 82.7 | 85.0 | unchanged (canonical K/V) | The three derivative-K/V engines now sit at standard's -fused-kernel ceiling (within 1%). `unlimited-context` is at HOnly +fused-kernel ceiling (within 1%). `windowed-checkpoint` is at HOnly ceiling — its `close_window` flow still needs `KvDispatch::read_kv_row_at` to pull the last K/V row back from the cache, leaving a ~3 ms/step residual. `turbo-quant` doesn't take the cascade — its codec is @@ -68,7 +68,7 @@ sibling files can access them. | Metal GPU | `standard` (control) | **99.8** | | Metal GPU | `markov-rs` | 87.1 | | Metal GPU | `markov-rs-codec:window=512` | 86.6 | -| Metal GPU | `unlimited-context:window=256` | 86.1 | +| Metal GPU | `windowed-checkpoint:window=256` | 86.1 | | Metal GPU | `boundary-per-layer:window=512,layers=34` | **87.2** | | Metal GPU | `turbo-quant:bits=4` | 82.7 | | CPU | `standard` | 28.7 | @@ -89,7 +89,7 @@ hot-path timings. | `boundary_kv/gate.rs` | 100% | ✅ | | `apollo/executor.rs` | 97.1% | ✅ | | `markov_residual_codec/executor.rs` | 94.6% | ✅ | -| `unlimited_context/dispatch.rs` | 84.9% | gated on Q4K vindex | +| `windowed_checkpoint/dispatch.rs` | 84.9% | gated on Q4K vindex | | `markov_residual_codec/dispatch.rs` | 82.3% | gated on Q4K vindex | | `markov_residual/dispatch.rs` | 84.3% | gated on Q4K vindex | | `boundary_per_layer/dispatch.rs` | 0% | gated on Q4K vindex | @@ -134,7 +134,7 @@ after the final commit. +30-48% across the cached-state engines. | `boundary_kv` (= standard + chunk frames) | 28 | ~99 | 0 MB | larql-boundary frames | composes with standard for cross-session resume | | `markov_residual` (W2 + W1-GPU + W7 blit) | 27.4 | **75.3** | 10.8 MB | residuals @ 4 B/tok | residual-stream, no f16 KV | | `markov_residual_codec` (W2 + W1-GPU + W7 blit) | 26.6 | **79.0** | 10.8 MB | bf16 residuals (2× cold saving) | long-context-friendly cold codec | -| `unlimited_context` (W1-GPU step 4 + W7 blit) | 28.1 | **82.7** | 15.7 MB (window=256) | per-window K/V checkpoints | W7 blit fusion +48% on top of W1-GPU | +| `windowed_checkpoint` (W1-GPU step 4 + W7 blit) | 28.1 | **82.7** | 15.7 MB (window=256) | per-window K/V checkpoints | W7 blit fusion +48% on top of W1-GPU | | `turbo_quant` (4-bit, W1-GPU + W7 blit, 10-tok bench) | 19.4 | **37.7** | 0.7 MB | — | WHT + Lloyd-Max K/V compression; codec cost grows with N | | `apollo` (boundaries) | — | requires store | scales w/ store | constellation map | retrieval+injection; not on the same scale as the others | | `no_cache` | — | — (O(N²) by design) | token list only | — | correctness baseline | @@ -160,7 +160,7 @@ after the final commit. +30-48% across the cached-state engines. fraction of per-token time. Codec cost also grows with sequence length (each step re-compresses the full layer K/V), so longer benches show lower mean tok/s. -- `unlimited_context` got the biggest W7 win (+48%) because its +- `windowed_checkpoint` got the biggest W7 win (+48%) because its per-step CPU-side work after the kernel returns is the lightest of the four cached-state engines, so the saved commit overhead is a larger fraction of total per-token time. The extra hot @@ -199,7 +199,7 @@ Engines that opted in: |---|---|---|---| | `markov_residual` | ✅ | ✅ (window=None) | K/V derivative (Metal cache is truth); h_in dead weight without cold-tier eviction | | `markov_residual_codec` | ✅ | ✅ (window=None) | Same — codec residuals are canonical, hot K/V is derivative | -| `unlimited_context` | ✅ | ❌ | `close_window` reads last K/V back via `KvDispatch::read_kv_row_at`; h_in needed for replay-from-checkpoint | +| `windowed_checkpoint` | ✅ | ❌ | `close_window` reads last K/V back via `KvDispatch::read_kv_row_at`; h_in needed for replay-from-checkpoint | | `turbo_quant` | ❌ | ❌ | K/V is canonical (destructive codec); cannot be derived | ### Measurement protocol @@ -265,8 +265,8 @@ active on every engine. | `markov-rs` None (windowless) | **99.1** | 10.10 ms | ~9.5 / ~1.2 ms | 0 MB | **+17%** | | `markov-rs-codec` Full | 88.3 | 11.33 ms | ~10.0 / ~1.2 ms | 26.3 MB | — | | `markov-rs-codec` None (windowless) | **98.5** | 10.15 ms | ~9.0 / ~1.2 ms | 0 MB | **+12%** | -| `unlimited-context:window=256` Full | 88.2 | 11.34 ms | ~10.1 / ~1.3 ms | 9.6 MB | — | -| `unlimited-context:window=256` HOnly | **92.8** | 10.78 ms | ~9.5 / ~1.2 ms | 0 MB | **+5%** | +| `windowed-checkpoint:window=256` Full | 88.2 | 11.34 ms | ~10.1 / ~1.3 ms | 9.6 MB | — | +| `windowed-checkpoint:window=256` HOnly | **92.8** | 10.78 ms | ~9.5 / ~1.2 ms | 0 MB | **+5%** | **What the numbers say:** @@ -282,7 +282,7 @@ active on every engine. `codec` (windowless), 30.2 MB → 0 MB on `markov-rs:window=512`, 9.6 MB → 0 MB on `unlimited:window=256`. The Metal kv cache is now the sole source of truth on the dispatch hot path. -- **`unlimited-context` win is small** (+5%) because most of its +- **`windowed-checkpoint` win is small** (+5%) because most of its per-step CPU work is the window-buffer slot-assign that survived even after the shadow is dropped (the window slots are pre-allocated regardless of mask). Memory savings still hold. @@ -356,7 +356,7 @@ regressions in PR review; not a proxy for real-model decode speed. | `standard:window=4` | 15.2 µs | 7.1 µs (smaller K/V to attend over) | | `no-cache` | 14.9 µs | 34.8 µs (re-runs full forward each step) | | `markov-rs` | 15.0 µs | 27.1 µs (recomputes K/V from residuals) | -| `unlimited-context` | 56.9 µs | 8.3 µs (window-checkpoint amortises decode) | +| `windowed-checkpoint` | 56.9 µs | 8.3 µs (window-checkpoint amortises decode) | | `turbo-quant` (4-bit) | 21.8 µs | 81.9 µs (codec dominates on tiny model) | | `apollo` | 45 ns (no boundary store loaded → early bail) | 2 ns (early bail) | @@ -381,7 +381,7 @@ the table above. - **Profiler.** Per-stage breakdown lands in `EngineProfiler`: embed, recompute_cold, recompute_hot, attention, ffn, total. -### unlimited_context +### windowed_checkpoint - **Mechanism.** Sliding window over the active K/V cache plus a checkpoint of the pre-window residual. Decode beyond the window @@ -442,7 +442,7 @@ contract (see `apollo` notes above). | `standard` | Backend KV cache; clean | — | | `markov_residual` | Cold-tier O(N²) merge — fixed 2026-05-19 via doubling-capacity `append_cold_overflow` | landed | | `markov_residual_codec` | Shares `markov_residual` cold tier; same fix | landed | -| `unlimited_context` | Clean | — | +| `windowed_checkpoint` | Clean | — | | `turbo_quant` | O(N) decompress+recompress per step — fixed 2026-05-19 via append-only codec path (30.1 → 82.8 tok/s, +175%) | landed | | `apollo` | O(N²) by design — `forward_from_layer` rebuilds KV each step over growing context | not a bug | | `boundary_kv` | Clean | — | @@ -532,7 +532,7 @@ checkpoint, use: ```sh cargo run -p larql-cli --release -- bench gemma3:4b --engine markov-rs -cargo run -p larql-cli --release -- bench gemma3:4b --engine unlimited-context:window=256 +cargo run -p larql-cli --release -- bench gemma3:4b --engine windowed-checkpoint:window=256 cargo run -p larql-cli --release -- bench gemma3:4b --engine turbo-quant:bits=4 cargo run -p larql-cli --release -- bench gemma3:4b --engine apollo:layer=30 cargo run -p larql-cli --release -- bench gemma3:4b --engine boundary-per-layer:window=512 diff --git a/crates/larql-kv/README.md b/crates/larql-kv/README.md index 8575ed5e3..155a3f4bb 100644 --- a/crates/larql-kv/README.md +++ b/crates/larql-kv/README.md @@ -20,7 +20,7 @@ GPU→CPU state bridge entirely (W10 mask cascade, default-on). | `markov-rs` | residual stream | **derivative** | exact logits under arch contract | **98.0** | | `markov-rs-codec` | compressed residuals | **derivative** | bounded KL | **98.1** | | `boundary-per-layer` | per-layer codec residuals | **derivative** | bounded KL per-layer | **98.7** | -| `unlimited-context` | KV (within window) + checkpoints | **derivative** | exact within window | 94.2 | +| `windowed-checkpoint` | KV (within window) + checkpoints | **derivative** | exact within window | 94.2 | | `turbo-quant` | quantised K/V | canonical (destructive) | bounded KL | 85.0 | | `boundary-kv` | K/V + boundary frames | canonical | exact logits | composes `standard` | | `apollo` *(RetrievalEngine)* | boundary retrieval store | n/a (retrieval) | task-level | orthogonal | @@ -50,7 +50,7 @@ for the dep-graph rationale. | Trait | For | Per-step contract | |---|---|---| -| `KvEngine` | per-token K/V cache engines (`standard`, `no-cache`, `markov-rs`, `markov-rs-codec`, `unlimited-context`, `turbo-quant`, `boundary-kv`, `boundary-per-layer`) | append K/V per layer; dispatches FFN through `&dyn FfnBackend`; state reconstructible to K/V tensors | +| `KvEngine` | per-token K/V cache engines (`standard`, `no-cache`, `markov-rs`, `markov-rs-codec`, `windowed-checkpoint`, `turbo-quant`, `boundary-kv`, `boundary-per-layer`) | append K/V per layer; dispatches FFN through `&dyn FfnBackend`; state reconstructible to K/V tensors | | `RetrievalEngine` | retrieval-injection engines (`apollo`, future Mode 5) | no per-token K/V append; no `FfnBackend` dispatch; state is residual delta + token list | Both return `Result, EngineError>` from `prefill` / @@ -59,7 +59,7 @@ five variants — `EmptyPrompt`, `BackendUnavailable`, `RetrievalMiss { reason }`, `InvariantViolation { what }`, `BackendFailure { details }` — and the typed error replaces the historical `Option` that collapsed all five into a silent -`None`. See the 2026-05-24 entry in [`ROADMAP.md`](ROADMAP.md). +`None`. See the 2026-05-24 entry in [`CHANGELOG.md`](CHANGELOG.md). ## Full engine catalog @@ -76,20 +76,20 @@ windowed checkpoints, retrieval injection, etc.). | [`markov_residual`](src/engines/markov_residual) | Residual-stream replacement, K/V derived from stored residuals (W2 cache) | 54.4 MB → **0 MB** | 88.2 | **99.5** (None, +13%) | exact (KL = 0.0) under contract | [markov-residual-engine.md](../larql-inference/docs/specs/markov-residual-engine.md) | | [`markov_residual_codec`](src/engines/markov_residual_codec) | `markov_residual` + bf16-encoded cold-tier residuals (2× cold saving) | 54.4 MB → **0 MB** | 87.2 | **99.8** (None, +14%) | bounded-KL vs markov_residual | [markov-residual-codec-engine.md](../larql-inference/docs/specs/markov-residual-codec-engine.md) | | [`boundary_per_layer`](src/engines/boundary_per_layer) | Per-layer codec policy on cold tier; calibration-driven; W1-GPU + W10 wired | 19.6 MB → **0 MB** | 86.9 | **99.3** (None, +14%) | per-layer KL bound | [boundary-per-layer-engine.md](../larql-inference/docs/specs/boundary-per-layer-engine.md) | -| [`unlimited_context`](src/engines/unlimited_context) | Per-window K/V checkpoint + token archive; supports replay | 15.7 MB → **0 MB** | 86.1 | **95.0** (HOnly, +10%) | exact within window | [unlimited-context-engine.md](../larql-inference/docs/specs/unlimited-context-engine.md) | +| [`windowed_checkpoint`](src/engines/windowed_checkpoint) | Per-window K/V checkpoint + token archive; supports replay | 15.7 MB → **0 MB** | 86.1 | **95.0** (HOnly, +10%) | exact within window | [windowed-checkpoint-engine.md](../larql-inference/docs/specs/windowed-checkpoint-engine.md) | | [`turbo_quant`](src/engines/turbo_quant) | WHT + Lloyd-Max 3/4-bit K/V codec, in-place compression | 0.7 MB | **37.7** (10-tok) | n/a (canonical K/V) | cos ≈ 0.991 | [turbo-quant-engine.md](../larql-inference/docs/specs/turbo-quant-engine.md) | | [`apollo`](src/engines/apollo) | Constellation map + boundary-residual injection (retrieval) | scales w/ store | requires store | n/a | task-level | [apollo-engine.md](../larql-inference/docs/specs/apollo-engine.md) | **Numbers are post W2 (hot K/V cache), W1-GPU (per-layer state-dump dispatch), W7 (blit-encoder fusion), and W10 (state-bridge mask cascade).** Three derivative-K/V engines (`markov_residual`, -`markov_residual_codec`, `unlimited_context`) now match or exceed +`markov_residual_codec`, `windowed_checkpoint`) now match or exceed `standard`'s fused-kernel ~100 tok/s ceiling under their best mask, with engine-side memory shadows fully eliminated on Metal. See [`PERFORMANCE.md`](PERFORMANCE.md) for per-token cost decomposition, the `state_capture` / `state_materialise` / `state_append` timer -cascade, and the bench protocol. ROADMAP "Closed (recent)" has the -milestone history. +cascade, and the bench protocol. [`CHANGELOG.md`](CHANGELOG.md) has +the milestone history. ### W10 — state-bridge mask cascade (default-on since 2026-05-21) @@ -127,6 +127,42 @@ architectures. Don't conflate them — the CLI's historical `--kv-cache markov-bounded` flag maps to `Standard { window_size: Some(N) }`, **not** `MarkovResidual`. Use the spec's table in §5 when in doubt. +### Windowed engines and the fused path + +A window means two promises at once: attend at most `N` positions, and +hold at most that much K/V. Until 2026-08 the only way to keep both was +to leave the backend's fused (coarse) path for the generic per-layer +route — and on Metal every per-layer dispatch method delegates to the +CPU, so a windowed engine ran its whole forward on the host while the +bench row still said `[metal (GPU)]`. On Gemma 3 4B Q4K that cost ~9x, +for a window that should make attention *cheaper*. + +Windowed engines now stay on the fused path. The window is requested via +`coarse_prefill_windowed` / `coarse_decode_step_windowed`, which **fail +closed**: a backend that cannot bound both attention and K/V answers +`None` and the engine falls back to per-layer, exactly as before. Both +`CpuBackend` and `MetalBackend` implement them. + +Two caveats worth knowing: + +- **A prompt longer than the window still takes the per-layer path.** The + fused prefill has no per-query-position masking, so accepting would + attend the whole prompt while advertising a bound. +- **Metal holds up to 2x the window** between compactions. Attention is + still bounded at the window every step (the kernel attends + `[T - window, T)`); the surplus rows are resident but never read. + Compacting every step would memmove the window per token. + +Measured on Gemma 3 4B Q4K, Metal, 80 decode steps at `window=8`: + +| | latency | hot K/V | +|---|---|---| +| unwindowed | 12.06 ms | 23.6 MB | +| `window=8` | **11.61 ms** | **2.4 MB** | + +The same config cost 115.44 ms before. A bench row reports which shape +it actually took — `[coarse]` or `[per-layer->host]`. + ## Usage ```rust @@ -183,7 +219,7 @@ standard:window=1024 # sliding-window K/V no-cache # full re-forward per step (O(N²)), debug only markov-rs # residual-stream replacement markov-rs:window=1024 -unlimited-context:window=256 +windowed-checkpoint:window=256 turbo-quant:bits=3 # alias: tq3 turbo-quant # bits=4 default; alias: tq4 apollo:layer=25,coef=8.0,top_k=12 @@ -224,7 +260,7 @@ let backend: Box = Box::new(CpuBackend); let mut async_engine = StandardEngine::with_async_backend(None, backend); ``` -The other research engines (`MarkovResidual`, `UnlimitedContext`, +The other research engines (`MarkovResidual`, `WindowedCheckpoint`, `TurboQuant`, `NoCache`, `Apollo`) gain the same `with_async_backend` constructor in subsequent slices. Spec: [`async-compute-backend.md`](../larql-inference/docs/specs/async-compute-backend.md). @@ -252,7 +288,7 @@ larql-kv/ │ ├── apollo/ — boundary-residual injection, ~4,000× compression │ ├── markov_residual/ — residual-stream KV replacement, KL = 0 │ ├── turbo_quant/ — WHT + Lloyd-Max K/V codec (3- or 4-bit) -│ └── unlimited_context/ — windowed re-prefill from checkpoints +│ └── windowed_checkpoint/ — windowed re-prefill from checkpoints ├── benches/ — criterion microbenchmarks ├── examples/ — end-to-end demos on synthetic test_utils ├── baselines/ — committed `larql accuracy` regression baselines diff --git a/crates/larql-kv/ROADMAP.md b/crates/larql-kv/ROADMAP.md index a8df537a3..2a0960c77 100644 --- a/crates/larql-kv/ROADMAP.md +++ b/crates/larql-kv/ROADMAP.md @@ -1,318 +1,76 @@ # Roadmap — larql-kv -## Spin-barrier pool — CPU MoE decode caught llama.cpp (2026-06-13) - -After residency closed the byte-traffic gap (06-11/12), a `/usr/bin/sample` of -live 26B decode showed the remaining ~1.15× was **rayon fork-join overhead**, -not kernels. The decode driver runs *outside* the global rayon pool, so each of -the ~211 parallel sections/token took the cold path (`in_worker_cold → -LockLatch::wait_and_reset → __psynch_cvwait`) and workers slept between sections -— ~40% of thread-time in wait states. - -**Built** [`larql_compute::cpu::spin_pool`](../../larql-compute/src/cpu/spin_pool.rs): -a llama.cpp-style persistent spin-barrier pool. Workers spin on an epoch counter -and only `park` after a long idle gap; the dispatcher participates as the n-th -worker; **static strided chunk ownership** makes `completed == num_chunks` a -sound barrier (no shared resettable cursor → no stale re-claim across -back-to-back dispatches — a concurrent-dispatcher test caught that bug); a -dispatch `Mutex` + thread-local reentrancy guard make it safe for -`--concurrent`/multi-threaded tests. `par_chunks_mut` / `par_chunks_mut2` -helpers route a row-chunked parallel-for through the pool, or rayon when -`LARQL_SPIN_POOL=0`. **Default-on** (see "Decode fast path default-on" — the -whole Q4K stack ships on, opt out per stage with `=0`); both paths are -numerically identical, only the threading differs. - -**Centralized** the four byte-identical `par_chunks_mut` Q4_K/Q6_K×Q8_K matvec -copies (larql-compute `cached.rs`, larql-inference `cached.rs`, lm_head ×2 in -`dense.rs` — the prior "consolidation hazard") into one -`q4k_q8k_matvec_parallel`, and routed every hot decode section (attention int8 -Q/K/V/O, GQA, dense FFN gate/up/down, geglu, expert fold, lm_head q4 + f32) -through it — so when enabled the whole token runs on one hot pool. - -- **Parity:** 704 compute + 1220 inference + 756 kv green, flags-off AND - flags-on (incl. the `predict_kquant` oracles). clippy clean. -- **Profile after:** rayon eliminated from the hot path — `in_worker_cold` - 2682→0, `join_context` 10300→0, `wait_until_cold` 4463→9. -- **Measured** (M3 Max, t=8, warm, tight A/B bracket, flags **inline**): - 26B short-ctx OFF ~26.9 → ON **33–35**; n=256 OFF ~27.4 → ON **~34.9 - (+28%)** — vs llama.cpp recorded **32.1** ⇒ ~9% ahead. -- **Default-on + safe (2026-06-13):** shipped a spin→yield→park backoff (spin - the proven window during active decode → yield once a wait outlives a token → - park when idle, ~0 CPU; dispatcher unparks on dispatch) so the pool doesn't - peg cores between requests — what makes on-by-default safe on a shared box. - Also fixed a panic-safety bug (a panicking chunk killed a worker → the - barrier spun forever): `catch_unwind` per chunk + re-raise on the dispatcher. -- **Caveat:** the pool spins during active decode (the win on a dedicated box); - under a transient mid-decode load spike a run can still regress (an n=512 ON - run hit 10.7 once) — `LARQL_SPIN_POOL=0` falls back to rayon if needed. - -## CPU resident fast-path — all engines pluggable into it (2026-06-13) - -The 2026-06-11/12 CPU fast-path arc (Q4K-direct + int8 attention, q4k -lm_head/dense residency, hand-asm kernels, KV append-in-place — see -`bench/baselines/c10_gemma4-26b-a4b_cpu_reconciled.json`) initially landed -only on `StandardEngine`: the `KvEngine::decode_step_resident` trait default -DROPPED the index (`let _ = index`), so every own-walk-loop engine stayed on -f32 attention. **Fixed:** - -- New single-source dispatcher - `larql_compute::attention::run_attention_block_decode_step_auto` — makes - the same q4k-direct-vs-f32 per-layer choice as - `CpuBackend::attention_step`, for callers that own `SharedKV` caches. -- `markov-rs`, `markov-rs-codec`, `turbo-quant`, `unlimited-context`, - `boundary_per_layer` now override `decode_step_resident` and thread the - vindex down their walk loops to `_auto`. `boundary-kv` forwards both - resident methods to its inner `StandardEngine` (was silently dropping to - the f32 path). `no_cache`/`apollo` keep the default by design (debug / - bench-only full re-forward). -- Regression pin: `engines::resident_identity_tests` — for 7 concrete - engine specs, `prefill/decode_step_resident` must be BIT-IDENTICAL to - `prefill/decode_step` with the flags off, and the covered-engine count - must not shrink. -- **Absolute matrix + slow-engine fixes 2026-06-13** (26B, default-on incl. - spin pool, M3 Max t=8 warm n=128). First measured: unlimited 31.8 / standard - 30.5 / boundary-kv 27.1 (**0.80×→0.89×**, its resident-forwarding fix) / - turbo 9.4 / markov 7.8 / codec 7.3 — the recompute/codec engines sat at - **~0.24–0.31×** because the spin pool sped up the shared attention/FFN/matvec - but not their per-step machinery. **Then fixed all three, feature intact:** - - **turbo-quant 9.4 → ~24** — `decompress_matrix`'s per-vector WHT decode was - *serial on the driver* (~35% of it); fanned across the spin pool. Still - 3-4-bit compressed (decoded every step, now parallel) — no memory tradeoff. - - **markov-rs 7.8 → 27.9, markov-rs-codec 7.3 → 27.7** — ported the W2 hot-K/V - cache to the **resident walk** (`rs_decode_step_inner`/`_codec`): read the - cached `hot_kv` and append the free `new_kv` from the attention step instead - of `recompute_kv`-ing every position each step. Gated `cache_eligible = - max_window.is_none() && no-cold` so it never tracks a window-clip - transition; the residual `stored` stays the canonical, re-derivable state - (the engine's point), the K/V is a droppable derivative. Parity gate: - `#[cfg(debug_assertions)]` assert cached K/V ≡ `recompute_kv` (≤1e-2), - exercised by `resident_identity_tests` (extended to a 10-step decode). - Final matrix: standard 34.5 / unlimited 32.1 / markov 27.9 / codec 27.7 / - boundary-kv 27.4 / turbo 21.1 — all **0.6–1.0× of standard** (was 0.24–0.31× - for the slow three). 756 kv tests green debug+release, clippy clean. -- **Comparative bottleneck review + walk allocation fix 2026-06-14.** Profiled - each engine's driver vs standard: the **shared** wall is the Q6_K expert - matvec (all engines inherit it); each engine's *delta* is its feature - machinery. markov/codec's −19/−20% was NOT the residual-store memcpy (~0.8% of - the driver) — it was **per-step allocation churn**: the resident walk's - `Array2::zeros((s_old+1, h))` rebuild + the cached-K/V `to_owned` - (`__bzero`+`szone_malloc` ≈ 2450 driver samples, idling the worker pool at 48% - vs standard's 80%). **Fixed:** the cache_eligible walk now `append_row`s - `stored` in place into the W8.2 doubling-capacity buffer (mirrors dispatch.rs) - and borrows `hot_kv` into attention via `Cow` instead of copying. Churn - collapsed 2450→150 samples (~16×); **markov/standard ratio 0.81×→0.975×, codec - 0.80×→~1.0×** (same battery state, back-to-back). Parity: resident_identity - (markov+codec, 10-step, buffer doubles) bit-exact + debug K/V assert. turbo's - −39% is **inherent** (must decode compressed K/V to attend; already - parallelized); boundary-kv/unlimited deltas are small (frame-emit/windowing). - Remaining markov/codec ~2.5% = walk-attention serial work (shared walk - frontier — full K/V concat + generic GQA vs standard's in-place handle). -- **In-place hot-K/V on the resident walk 2026-06-14 (closes the concat half).** - The named ~2.5% above was the walk-attention **owned concat**: the resident - walk drove `run_attention_block_decode_step_q4k_direct`, which allocates a - fresh `[ctx+1, kv_dim]` K *and* V every layer every step and copies the whole - prior cache into it before attending — **O(L²)** cache copy over an L-token - generation, vs `standard`'s in-place append handle (O(L)). The split - project→append→attend halves already existed for the dispatch path; the walk - just didn't use them. **Built** `run_attention_block_decode_step_{q4k_direct, - auto}_inplace` (larql-compute `attention/decode.rs`): projects the new row, - appends it into the caller's **doubling-capacity** K/V buffer (grows like - `stored`), and attends over the `[..len+1]` views — no concat. **Wired** - markov_residual + markov_residual_codec resident walks: step-1 still - recompute-seeds `hot_kv`; steps 2+ append in place (the steady state). The - windowed/cold tiers and the flags-off f32 path keep the owned concat - unchanged. Gated `LARQL_MARKOV_INPLACE_KV` (default on; `=0` → owned concat, - the A/B reference + escape hatch). **Parity (bit-exact, 4 gates):** compute- - level `inplace ≡ q4k_direct` concat across a capacity doubling; engine-level - in-place-vs-owned-concat A/B with Q4K-direct **on** for markov *and* codec - (hidden states bit-identical every step); `resident_identity` flags-off still - green (in-place branch's None-fallback = owned concat); 758 kv + 705 compute + - 1220 inference green debug & opt, clippy clean. (The debug `hot_kv ≡ - recompute_kv` assert is gated to the f32 path — the Q4K route's projections - differ from `recompute_kv` by >1e-2 even in f32-act; its oracle is the A/B.) - The two q4k-flag-mutating tests serialise on `Q4K_FLAG_ENV_LOCK` (those flags - read process env on the driver thread — no thread-local). **Perf is - structural** (eliminates the O(L²) per-step copy; the win grows with context — - it's the long-ctx tax behind the C10 1.29× vs short 1.15×). **Measured (26B, - CPU MoE in-process, M3 Max t=8, n=128 warm, `LARQL_MARKOV_INPLACE_KV` A/B, - same engine ordering):** markov 32.5→34.5, codec 32.5→34.6 with in-place on — - and the three untouched controls (standard/unlimited/turbo) drifted *down* - −3/−8/−6% across the A/B (machine warming), so drift-corrected the change is - **~+11–12%**. Final warm matrix (in-place on = production default): codec - **36.5** / standard 36.0 / markov **36.0** / unlimited 33.3 / boundary-kv 36.5 - p50 (mean skewed by frame-emit spikes) / turbo 21.2 (inherent) — **markov/codec - now AT parity with standard** (was 0.81× at the arc's start), the whole cached - cluster **~12% ahead of llama.cpp's 32.1**. Caveat: bench box was at ~58% - charging (not cool-dedicated); ordering + A/B *direction* are robust, absolutes - drifted ~5–8% run-to-run — a cool-box rerun would firm them. (NB: the first - engine in a fresh process eats the 30GB page-in — standard read 21.8 cold, - 34–36 warm; warm runs are the fair matrix.) -- **Propagated the in-place lever to the two remaining walk engines + faithfulness - audit 2026-06-14.** A full cross-engine spec/contract audit (all 9 engines vs - `state-policy.md`'s `(canonical, derivative, contract)` triple) found every - engine faithful, and flagged the two siblings still paying the O(L²) owned - concat the markov/codec in-place change eliminated: - - **boundary_per_layer (was the one NEEDS-FIX)** — carried NO `hot_kv` at all: - it `recompute_kv`'d the whole hot tier *and* rebuilt an owned `[ctx+1]` concat - every layer every step (worse than markov *pre*-W2). Added a `hot_kv` - derivative + the W2-cache + `run_attention_block_decode_step_auto_inplace` - steady state, mirroring its twin codec — only active in the `cache_eligible` - (unbounded, no cold) path, like codec; the windowed/cold path (its primary - purpose) is untouched. `hot_kv` is excluded from `memory_bytes` (droppable - derivative, matches markov). Engine-level in-place-vs-owned-concat A/B (q4k on) - bit-identical; f32-gated debug `hot_kv ≈ recompute_kv` assert. - - **unlimited_context** — its CPU window walk (`extend.rs`) passed the whole - window K/V by value → backend re-concats `[n+1]` per layer per step (its own - doc admitted "O(window²) total"). Added `rs_extend_inplace` (appends into the - window's doubling-capacity buffer, attends over views), wired into - `extend_current` only when eligible (index + toggle + q4k); `replay_window` / - quant / executor / tests keep the owned concat. The engine's existing - `current_window_kv_len` counter already treated the buffers as over-allocated - (the dispatch path did), so `close_window`/`current_kv_bytes` needed no change. - A/B (q4k on) bit-identical; `resident_identity` flags-off still green. - Both reuse the shared `LARQL_MARKOV_INPLACE_KV` toggle + `Q4K_FLAG_ENV_LOCK`. - Also: **apollo footgun guard** — `injection_layer < crystal_layer` silently - no-ops the retrieval-injection (the compressed forward starts at `crystal`); - added a one-time runtime warning in `prepare_injection` (experimental engine, - warn-don't-fail). Doc-drift swept: boundary-kv spec now flags `resume` as - NOT-IMPLEMENTED (emit half only), apollo spec `KvEngine`→`RetrievalEngine`, - `state-policy.md` `fallback_mode` marked retired (per its own §8 resolution). - 760 kv tests green debug + opt, clippy clean. (Same caveat as above: turbo's - −39% is inherent; boundary-kv inherits standard's opts via resident forwarding.) - -Prefill stays on the f32 BLAS gemm for all engines deliberately (the task -#16 prefill falsification: q4k repeated-matvec loses ~20× to AMX at -prefill shapes). - -## Hardening — codebase review 2026-05-28 - -From the whole-codebase review ([`docs/audits/codebase-review-2026-05-28.md`](../../../docs/audits/codebase-review-2026-05-28.md)): - -- **P2 — CLI-supplied sizing params can reach prefill panics**; validate at the boundary. -- **P2 — positional QKVO contract** (`attn_data[1]/[2]`, shared with larql-models) is maintained by convention, not type. Silent-drift risk — consider a typed accessor. - -## Current state (as of 2026-05-18) - -**Performance equilibrium post W7 + W8 + W8.2 + Step 9** (Gemma 3 4B -Q4K, Metal, M3 Max): - -| Engine | 50-tok tok/s | 1000-tok tok/s | Prefill (5-tok) | Gap to standard @ 1k | -|---|---:|---:|---:|---:| -| `standard` (fused) | 100.3 | 64.1 | 300 ms | — | -| `markov-rs` | 88.9 | **58.7** | 265 ms | -8.4% | -| `markov-rs-codec` | 88.8 | **57.2** | 270 ms | -10.8% | -| `unlimited-context` | 86.4 | 57.4 | 256 ms | -10.4% | -| `turbo-quant` (4-bit, 10-tok) | 37.7 | — | — | codec-bound | - -All cached-state engines now cluster within ~10% of `standard`'s -fused-kernel ceiling. The 135% pre-W8.2 gap on `markov-rs` / -`markov-rs-codec` collapsed once the per-step `Array2::zeros((n+1, -kv_dim)) + slice-copy` pattern was replaced with doubling-capacity -in-place append. Prefill is no longer the wall-time dominator -(post Step 9: 10× speedup vs the 2.7 s CPU walk it used to fall back -to). See "Closed (recent)" for the milestone history. - -The remaining 8-11% decode gap is fixed CPU glue (state-dump -readback into `PerLayerDecodeState`, counter bump, append-row). -Closing further requires either single-kernel prefill state-dump -(W9 — Metal kernel surgery, small wall-time win at current bench -shape) or a Metal-side path that elides the per-token CPU readback -entirely (W10 — engine-side state lives on GPU until window-close). - -## Crate-shape state (2026-05-17) - -- Crate extracted from `larql-inference::engines` on 2026-05-09 — see - [`CHANGELOG.md`](CHANGELOG.md). -- **Seven engines shipped** as of 2026-05-17: - - Original four: `standard`, `no_cache`, `markov_residual`, - `unlimited_context`, `turbo_quant`, `apollo`. - - Three new: `boundary_kv`, `markov_residual_codec`, `boundary_per_layer`. - Specs in `crates/larql-inference/docs/specs/`: - [boundary-kv-engine.md](../larql-inference/docs/specs/boundary-kv-engine.md), - [markov-residual-codec-engine.md](../larql-inference/docs/specs/markov-residual-codec-engine.md), - [boundary-per-layer-engine.md](../larql-inference/docs/specs/boundary-per-layer-engine.md). -- Consumers wired: - - `larql-cli bench --engine ` (selector dispatch) - - `larql-cli bench --via-executor` opts into the new `LayerExecutor` - surface; falls through to legacy path for unmigrated engines. - - in-crate `benches/engine_decode.rs` (criterion: dispatch helpers + Standard parity) -- Coverage policy: 90 % line coverage per source file (see - `coverage-policy.json`); CI gate at `make larql-kv-coverage-policy`. - Workspace `larql-kv` lib total: **95.62% lines, 95.43% regions, 95.50% - functions** (2026-05-24 evening, post coverage-debt clearance). - **All 61 files at ≥90% lines; debt baselines cleared from policy - file.** The 2026-05-24 push lifted the five `engines/*/dispatch.rs` - files (range 7.95–80.68% → 93.57–97.85%) and - `engines/markov_residual/compute.rs` (86.85→95.30%). See "Closed - (recent)" entry for the thread-local-override pattern that makes - the env-gated paths in `compute.rs` and the W10 mask cascade in - the dispatch files testable without process-env mutation. - -## Architectural cuts (2026-05-17) - -Substantive refactors landed; specs reflect the new boundaries. - -### Naming hygiene — renamed for honesty - -- **`metal_fused_prefill` / `metal_fused_decode_step`** → `fused_prefill` - / `fused_decode_step`. The "metal" was a lie — `CpuBackend` implements - `prefill_q4` and `decode_token` via its C Q4 kernel and also takes the - fused path on `--cpu`. The aliases in `unlimited_context::engine` - (`quant_prefill_metal`, `quant_decode_token`) follow. -- **`KvEngine::prefill_q4k` / `decode_step_q4k`** → `prefill_quant` / - `decode_step_quant`. The `_q4k` suffix baked one format into the trait - surface; the trait is quant-agnostic (dispatches on `index`'s format). - Internals that are genuinely Q4K-specific (`prefill_q4k_moe`, - `cpu_q4k_cache_*`, `run_ffn_decode_step_q4k_direct`) keep their names. -- **`ComputeBackend::has_q4()` → `supports_quant(format: QuantFormat)`.** - Per-format predicate; `CpuBackend` reports support for `Q4_0`, `Q4_K`, - `Q4_KF`, `Q6_K`; `MetalBackend` adds `Q8_0`. Backends can advertise new - format support without trait extension. -- **Storage slots `q4k` → `kquant` for K-family fields.** `attn_q4k`, - `interleaved_q4k`, `set_attn_q4k`, `load_attn_q4k`, etc. — these hold - K-family quant bytes (Q4_K, Q4_KF, Q6_K — manifest tag picks). Q4_0 - (`attn_q4`) and Q8 (`attn_q8`) slots stay — genuinely format-specific. - -### Engine state vs execution — new abstraction - -Spec: [engine-state-vs-execution.md](../larql-inference/docs/specs/engine-state-vs-execution.md). - -The engines were re-coupling backend / FFN / format decisions into their -state-management code. The new shape: - -- **`LayerExecutor` trait** (in `larql-inference::layer_executor`) — - per-layer execution surface with `run_prefill_layer` / - `run_decode_layer` returning `(hidden, SharedKV)`. Dispatch kind - (`Fused` / `PerLayer`) is explicit. -- **`LocalWalkExecutor`** — wraps `run_attention_with_kv_backend` + - the caller's `&dyn FfnBackend`. The critical decoupling: the executor - does **not** construct its own `WalkFfn` — it uses whatever the engine - was handed. -- **Engine trait extension:** `KvEngine::prefill_via_executor`, - `decode_step_via_executor`, `prefill_quant_via_executor`, - `decode_step_quant_via_executor`. Default impls fall through to the - legacy methods so unmigrated engines work unchanged. - -### Engines on the new surface - -Every engine now runs its own state-policy code; there is no hidden -fall-through to the backend's fused kernel from per-layer engines. -`standard` (and by delegation `boundary_kv`) is the **only** engine -that exercises the fused fast path — via -`ComputeBackend::coarse_prefill` / `coarse_decode_step`, which on -Metal calls `larql_inference::vindex::fused_prefill`. - -| Engine | Default dispatch | `*_via_executor` override | Honors FFN backend | Tok/s (Gemma 3 4B Q4K, Metal) | Hot state | -|---|---|---|---|---:|---:| -| `standard` | `ComputeBackend::coarse_prefill` (fused fast path) | n/a (no per-layer code to migrate) | n/a | 104 | 0 MB (backend owns K/V) | -| `boundary_kv` | Delegates to `standard` + emits boundary frames | n/a | n/a | ≈104 | 0 MB | -| `markov_residual` | Per-layer walk via `rs_prefill_walk` | ✅ | ✅ counter test | 3.6 | 6.0 MB | -| `markov_residual_codec` | Per-layer walk via `rs_prefill_codec_walk` (bf16 cold) | ✅ | ✅ counter test | 4.3 | 6.0 MB | -| `unlimited_context` | Windowed checkpoint extension via `process_q4k` | ✅ | ✅ counter test | 25.6 | 4.8 MB | -| `turbo_quant` | Per-layer WHT + Lloyd-Max compression cycle | ✅ | ✅ counter test | 3.9 | 0.6 MB | -| `boundary_per_layer` | Per-layer walk with per-layer codec policy | ✅ (dense) | ✅ counter test | — | matches markov_residual_codec | -| `apollo` | Whole-forward through `forward_layer_range` (boundary prefix + perturb) | ✅ | ✅ counter test | requires store | scales with store | -| `no_cache` | Full re-forward per step (O(N²) wall-time) | ✅ | ✅ already did on legacy `prefill` | — | token list only | +## Current state (2026-08-04) + +Nine engines behind one `EngineKind` selector, all reachable from +`larql bench --engine` and pinned by `tests/gpu_engine_parity`, +`tests/dispatch_parity` and `tests/engine_ple_parity`. + +### Which path a run actually takes + +This is the thing to know before reading any engine number, because it +decides what was measured: + +| path | when | what runs | +|---|---|---| +| **coarse (fused)** | Q4K vindex present, backend accepts the engine's window | backend's fused pipeline; K/V lives in the backend | +| **coarse + window** | as above, window bounded by the backend | same, with attention and K/V bounded | +| **per-layer → host** | prompt longer than the window, or an arch that declines coarse | generic per-layer loop — **on Metal this runs the whole forward on the CPU** | + +A bench row reports which one it took (`[coarse]` / `[per-layer→host]`). + +**On the coarse path the engine's own state policy is not engaged** — the +K/V sits in the backend behind a sentinel handle, so `standard`, +`markov-rs`, `markov-rs-codec` and `boundary-per-layer` execute the same +kernels and land within ~0.5% of each other. Ranking those four against +one another on that path measures nothing. `turbo-quant` and the +windowed configs do differ, because their state genuinely materialises. + +### Windows + +`window=N` promises bounded attention *and* bounded K/V. Both backends +honour it on the fused path via `coarse_prefill_windowed` / +`coarse_decode_step_windowed`, which fail closed — a backend that cannot +bound both answers `None` and the engine falls back to per-layer. + +CPU trims the cache before each step. Metal clamps the attention span +every step and compacts at 2x the window, so the memmove is O(1) +amortised; up to 2x the window is resident between compactions, but +attention never reads past the window. + +Measured, Gemma 3 4B Q4K, Metal, 80 steps at `window=8`: 11.61 ms / +2.4 MB, against 12.06 ms / 23.6 MB unwindowed. + +Per-layer sliding-window attention (the *architecture's* window, e.g. +Gemma 3's 1024) is implemented on the CPU attention path and resolved +through one rule shared with the Metal pipeline spec, +`effective_attention_window_for_layer`. + +### Measurement + +A bench step is a **whole token** — engine forward + lm_head + next-token +pick — matching the reference rows, with a `fwd=` / `head=` split in the +row note. Memory counts K/V the backend holds on the engine's behalf, not +just what the engine allocated, and declines to print a ratio when +nothing was measured. + +Cross-backend: prefill agrees to ~5e-7 relative L2; per-step decode +differs by a stable ~3e-3 that does not compound (two Q4K kernels +rounding differently). + +### Known limits + +- A prompt **longer** than the window still takes the per-layer path: the + fused prefill has no per-query-position masking, so accepting would + attend the whole prompt while advertising a bound. +- On Metal the per-layer path delegates every dispatch method to + `CpuBackend`, so anything that declines coarse runs on the host — + worth ~9x on Gemma 3 4B. +- `apollo` needs an attached boundary store; without one `prefill` fails + closed with `RetrievalMiss`, and it is excluded from the criterion + bench for that reason (`EngineKind::bench_excluded_names`). + +See [CHANGELOG.md](CHANGELOG.md) for how each of these came to be. ## Coverage debt @@ -325,7 +83,7 @@ lines, 61/61 files at ≥90%, 0 debt baselines in | File | Pre | Post | |---|---:|---:| | `engines/markov_residual/compute.rs` | 86.85% | **95.30%** | -| `engines/unlimited_context/dispatch.rs` | 59.09% | **97.24%** | +| `engines/windowed_checkpoint/dispatch.rs` | 59.09% | **97.24%** | | `engines/markov_residual/dispatch.rs` | 77.51% | **96.78%** | | `engines/markov_residual_codec/dispatch.rs` | 80.68% | **97.72%** | | `engines/turbo_quant/dispatch.rs` | 9.35% | **97.85%** | @@ -379,6 +137,12 @@ config-injection refactor. ## Open work +> **Below all of these:** the decode attention K/V layout. Every cached engine +> shares one marginal cost per context token because they all land in that +> kernel, so a head-major layout would move all six at once — see +> [`larql-compute/ROADMAP.md`](../larql-compute/ROADMAP.md) "head-major K/V +> layout" and [`docs/decode-cost-model.md`](docs/decode-cost-model.md) §4. + ### P0 — codebase-health frontier (audit 2026-06-14) A whole-codebase review (engine faithfulness audit + clippy/coverage sweep) @@ -573,7 +337,7 @@ full-recompute `predict_kquant_hidden*` path with **no KV cache**. CPU full-recompute fix that closed #146, 2026-05-28). Goal: make the engine layer MoE-aware so CPU MoE decode is KV-cached and -**engine-selectable** (standard / unlimited_context / markov* / turbo_quant / +**engine-selectable** (standard / windowed_checkpoint / markov* / turbo_quant / apollo all apply their mechanism to MoE models, not just dense). Subtasks: @@ -640,7 +404,7 @@ it FAILS at L5 if the prefill-RoPE fix is reverted. | **markov_residual** | per-layer `compute.rs` `run_ffn` → `layer_ffn_or_moe` | ✅ | "Paris", **3.1 tok/s** | | **boundary_per_layer** | per-layer `walk::run_prefill`/`run_decode` (larql-kv) → `layer_ffn_or_moe` | ✅ | "Paris", **3.1 tok/s** | | **boundary_kv** | wraps `StandardEngine` + compressed-residual boundary frames | ✅ | "Paris", **2.9 tok/s** | -| **unlimited_context** | per-layer `rs_extend_from_checkpoint_backend` → `layer_ffn_or_moe` | ✅ | "Paris", **1.7 tok/s** | +| **windowed_checkpoint** | per-layer `rs_extend_from_checkpoint_backend` → `layer_ffn_or_moe` | ✅ | "Paris", **1.7 tok/s** | | no_cache | legacy `kv_prefill_run` full re-forward | ✗ (by design) | full re-forward per step; not sensible for remote experts | | apollo | local re-forward (`forward_from_layer`) | ✗ (by design) | crystal re-forward *multiplies* per-step expert round-trips | @@ -679,11 +443,11 @@ lands as **true — all seven within ~2.6× and network-bound** — `standard` t **Best engine for remote MoE:** `standard` for throughput; `boundary_kv` for wire-efficient cold-context residual frames; `markov`/`turbo`/`boundary_per_layer` -for compressed KV memory at near-standard speed; `unlimited_context` for +for compressed KV memory at near-standard speed; `windowed_checkpoint` for long-context windowed KV (slowest, bounded memory). `no_cache` / `apollo` are not a fit (re-forward multiplies round-trips). -**Resolved (2026-06-13):** `unlimited_context::replay_window` now takes +**Resolved (2026-06-13):** `windowed_checkpoint::replay_window` now takes `moe_ffn` + `index` and threads them to `rs_extend_from_checkpoint_backend` (matching the live-window `extend_current` path), so an evicted MoE window replays with experts instead of silently falling back to dense FFN. It is a @@ -765,7 +529,7 @@ decode): | Engine | tok/s | Hot state | Per-step cmd_bufs (Metal) | Per-step compute model | |---|---:|---:|---:|---| | `standard` (fused) | 104 | 0 MB (backend-owned) | 1 | one fused kernel, all 34 layers, append-1-row K/V | -| `unlimited_context` | 25.6 | 4.8 MB | ~103 | per-layer attn+ffn, append-1-row K/V (same compute as standard, different dispatch) | +| `windowed_checkpoint` | 25.6 | 4.8 MB | ~103 | per-layer attn+ffn, append-1-row K/V (same compute as standard, different dispatch) | | `markov_residual_codec` | 4.3 | 6.0 MB | ~103 | per-layer attn+ffn + **recompute K/V from `window_size` residuals every step** | | `turbo_quant` (4-bit) | 3.9 | 0.6 MB | ~103 | per-layer attn+ffn + **decompress prior K/V + re-encode updated K/V every step** (CPU codec in inner loop) | | `markov_residual` | 3.6 | 6.0 MB | ~103 | same as codec; no codec overhead on bench (cold tier never fired in 20-step run) | @@ -780,7 +544,7 @@ decode): | Engine | CPU tok/s | GPU (Metal) tok/s | Where the gap lives | |---|---:|---:|---| | `standard` (coarse_prefill control) | 28.2 | 102.7 | GPU's fused fast path is 3.6× the CPU C kernel. | -| `unlimited_context` | 28.1 | 28.4 | **At parity** — no per-layer overhead either side. | +| `windowed_checkpoint` | 28.1 | 28.4 | **At parity** — no per-layer overhead either side. | | `markov_residual_codec` | 26.6 | 27.5 | **At parity** (post-W2). | | `markov_residual` | 26.5 | 26.8 | **At parity** (post-W2). | | `turbo_quant` (4-bit) | 19.4 | 19.6 | **At parity** — codec overhead dominates on both. | @@ -818,7 +582,7 @@ use Metal at all — there's nothing to batch yet. W2 landed: caching the hot K/V projection across decode steps moved both markov_residual engines from ~5 to ~27 tok/s — they now -sit on the same curve as `unlimited_context` (which already cached +sit on the same curve as `windowed_checkpoint` (which already cached K/V incrementally), within 1.5 tok/s of each other. The `recompute_kv` stage no longer fires; FFN+attention dominate exactly like every other cached-K/V engine. **The hot K/V state @@ -827,11 +591,11 @@ costs ~10.8MB vs 5.3MB pre-W2** (trade memory for speed; still Reading the table: percentages are *of the engine's own per-step total*, not vs standard. The three cached-K/V engines (markov-rs, codec, -unlimited-context) now cluster around 27-28 tok/s, all showing the +windowed-checkpoint) now cluster around 27-28 tok/s, all showing the same FFN-heavy decode profile. The remaining ~4× gap to standard is per-layer Metal dispatch overhead — W1's target. -**`unlimited_context` — 28.4 tok/s, 35 ms/tok. Per-layer attn + ffn +**`windowed_checkpoint` — 28.4 tok/s, 35 ms/tok. Per-layer attn + ffn dominates; no recompute waste.** Compute model is identical to standard's (append-1-row K/V per layer). 74% of the step is FFN, 25% is attention. The 4× gap to standard is **per-layer Metal command- @@ -846,7 +610,7 @@ projected ~80-100 tok/s.** **`markov_residual` / `markov_residual_codec` — 26.8 / 27.5 tok/s, ~37 ms/tok. W2 LANDED.** The hot K/V cache eliminates the 80% recompute overhead measured pre-W2; both engines now sit on the same curve as -`unlimited_context` while preserving the residual-stream contract +`windowed_checkpoint` while preserving the residual-stream contract (drop `hot_kv` and the next step recomputes from `stored` — the fallback path is still there for the via_executor path that doesn't yet capture K/V). The W2 design preserves the engine identity: K/V is @@ -863,14 +627,14 @@ cold tier. is ~25% of the budget, not the bottleneck.** This is a real surprise: the pre-profile guess was "codec encode/decode is the inner-loop killer." Measured: codec is ~25% (9.4% decode + 15.5% encode), FFN is -53%, attention is 20%. Turbo_quant is much closer to unlimited_context +53%, attention is 20%. Turbo_quant is much closer to windowed_checkpoint (28.4 tok/s) than to markov_residual (~5 tok/s) — the engine works. The codec is a fixed overhead per layer per step, not a quadratic blow-up. **Workstream W3 (incremental encode of the new row only) still applies — it would cut the 15.5% encode share roughly in half — but the bigger lever is W1 (dispatch batching), since FFN dominates the per-step budget and is the same per-layer-Metal bottleneck as on -unlimited_context.** W4 (SIMD WHT) is now lower-priority than originally +windowed_checkpoint.** W4 (SIMD WHT) is now lower-priority than originally estimated; codec is fast enough that vectorising it shaves single-digit percent. @@ -889,8 +653,8 @@ correctness-baseline only. | ID | Workstream | Engines | Expected gain | Risk | |---|---|---|---|---| -| W1-GPU | **Route per-layer Q/K/V/O and FFN matvecs through Metal.** Today's `attention_decode_step_native` and `ffn_decode_step_native` ignore the backend param and run rayon CPU matvec — that's why all four per-layer engines hit ~27 tok/s on both `--backends cpu` AND `--backends metal`. The GPU is not involved at all. Workstream: make these helpers actually dispatch to `MetalBackend`'s per-layer quant matvec kernels (the ones `fused_prefill` already uses internally). **GPU only.** | unlimited_context, markov_residual, markov_residual_codec, turbo_quant | Unknown — first deliverable is the measurement. Ceiling ranges from ~40 tok/s (submit overhead dominates) to ~80 tok/s (matches standard's GPU advantage). | Per-layer Metal submit cost (50-100µs each × ~6 per layer × 34 layers = ~10-20ms/token) is the open question. May need to batch within a layer (Q+K+V in one buffer, attn separately, etc.) to amortize. CPU is at parity already; no W1-CPU. | -| W2 | **Persistent hot K/V cache in markov_residual.** The engine contract says "K/V derived from residuals" — it does **not** say "recomputed every step." Cache hot K/V across steps; append-1-row on new residual; only recompute fully on cold-tier eviction (rare). Cold-tier compression remains the engine's selling point. | markov_residual, markov_residual_codec | ~20-30×; engine becomes "unlimited_context with compressed-residual cold tier" | Need to verify residual store still reflects "what we'd recompute from" — i.e., consistency check that cached K/V matches a fresh recompute under same residuals. Add a debug assertion mode. | +| W1-GPU | **Route per-layer Q/K/V/O and FFN matvecs through Metal.** Today's `attention_decode_step_native` and `ffn_decode_step_native` ignore the backend param and run rayon CPU matvec — that's why all four per-layer engines hit ~27 tok/s on both `--backends cpu` AND `--backends metal`. The GPU is not involved at all. Workstream: make these helpers actually dispatch to `MetalBackend`'s per-layer quant matvec kernels (the ones `fused_prefill` already uses internally). **GPU only.** | windowed_checkpoint, markov_residual, markov_residual_codec, turbo_quant | Unknown — first deliverable is the measurement. Ceiling ranges from ~40 tok/s (submit overhead dominates) to ~80 tok/s (matches standard's GPU advantage). | Per-layer Metal submit cost (50-100µs each × ~6 per layer × 34 layers = ~10-20ms/token) is the open question. May need to batch within a layer (Q+K+V in one buffer, attn separately, etc.) to amortize. CPU is at parity already; no W1-CPU. | +| W2 | **Persistent hot K/V cache in markov_residual.** The engine contract says "K/V derived from residuals" — it does **not** say "recomputed every step." Cache hot K/V across steps; append-1-row on new residual; only recompute fully on cold-tier eviction (rare). Cold-tier compression remains the engine's selling point. | markov_residual, markov_residual_codec | ~20-30×; engine becomes "windowed_checkpoint with compressed-residual cold tier" | Need to verify residual store still reflects "what we'd recompute from" — i.e., consistency check that cached K/V matches a fresh recompute under same residuals. Add a debug assertion mode. | | W3 | **Incremental TurboQuant encode (append-only).** Only encode the new K/V row each step; keep prior compressed bytes untouched. Decompress only the new row's neighbourhood for attention scores (or the whole layer if simpler). | turbo_quant | ~10× at long context | Re-encoding for in-place updates is the slow path. Need to define when (if ever) the full layer needs re-encoding. | | W4 | **TurboQuant SIMD WHT + Lloyd-Max.** Already on P1; promote to P0 once W3 lands so the per-row codec cost is the only remaining work. NEON on Apple Silicon, AVX2 on x86_64. | turbo_quant | 2-4× on the codec step | Mostly mechanical; landing W3 first means each step touches less data, making SIMD's batch budget go further. | | W5 | **Apollo K/V cache across decode steps.** Cache the K/V for layers `crystal..num_layers` between steps; append-1-row per step instead of re-forwarding. Reduces per-step cost from O(N) to O(1) in generated length. | apollo | linear → constant per-step | Apollo's vec_inject perturbation fires at `injection_layer`; verify the perturbation interacts correctly with cached K/V (it should — perturbation is residual-additive, not K/V-overwriting). Needs an apollo store fixture in tree to bench. | @@ -903,12 +667,12 @@ Recommended order (revised 2026-05-17 night after W7 produced measured numbers — replaces the earlier guess-driven sequence): 1. **W7 — DONE.** Profiler wired across markov_residual, - markov_residual_codec, unlimited_context, turbo_quant. Each + markov_residual_codec, windowed_checkpoint, turbo_quant. Each engine's `--profile` output produces a per-stage attribution. See the measured table above. 2. **W2 — DONE.** Hot K/V cache landed on `markov_residual` and `markov_residual_codec`. Both moved from ~5 tok/s to ~27 tok/s - (5.5-5.7×) and now sit on the same curve as `unlimited_context`. + (5.5-5.7×) and now sit on the same curve as `windowed_checkpoint`. Engine contract preserved: K/V still derivable from residuals, just not re-derived every step. Hot K/V state grew from 5.3MB to 10.8MB; that's the speed/memory trade. Bit-parity tests @@ -950,7 +714,7 @@ measured numbers — replaces the earlier guess-driven sequence): make per-layer fast, not to skip it. - **Not removing engine contracts.** Markov-rs's residual store must still be re-deriveable; turbo_quant's K/V must still be - compressed; unlimited_context's checkpoints must still emit at + compressed; windowed_checkpoint's checkpoints must still emit at window boundaries. Optimizations are within the contract. - **Not optimising no_cache.** It's a correctness baseline; O(N²) is the design. @@ -969,7 +733,7 @@ silently: a non-empty prompt. - `markov_residual_codec`: same; plus `cold_bytes() > 0` after overflow. - - `unlimited_context`: `archive.len() > 0` after at least + - `windowed_checkpoint`: `archive.len() > 0` after at least `window_size` tokens. - `turbo_quant`: `layers.len() == num_layers` after prefill. - `apollo`: `context_tokens.len() > 0` after prefill. @@ -1016,7 +780,7 @@ in expected ROI order. matching CPU-side shadow type for `CpuBackend` which has no on-GPU state.** Pre-req: stable `MetalBackend`-side KV cache invariants (which Step 9 already established). -- **W8.2 → `unlimited_context` CPU walk fallback.** The legacy CPU +- **W8.2 → `windowed_checkpoint` CPU walk fallback.** The legacy CPU walk path (`process_via_executor` at engine.rs:~720) still uses the per-step `Array2::zeros((s_old+1, dim))` pattern. Not on the hot path for the bench (dispatch path is the default), but a @@ -1099,11 +863,11 @@ disagreeing semantics instead of two. variants) return type changes from `Option` to `Result`. **All eight `KvEngine` impls touched** — `standard`, `no_cache`, `markov_residual`, `markov_residual_codec`, - `unlimited_context`, `turbo_quant`, `boundary_kv`, + `windowed_checkpoint`, `turbo_quant`, `boundary_kv`, `boundary_per_layer` — not just the one that motivated the refactor. The translation is mechanical: validated on three structurally-distinct samples (`markov_residual` for arch - preconditions, `unlimited_context` for window boundaries, + preconditions, `windowed_checkpoint` for window boundaries, `boundary_per_layer` for calibration stores); every `None`-return in those engines maps cleanly to `InternalError(...)`. The remaining five are variations on already-validated patterns @@ -1157,7 +921,7 @@ disagreeing semantics instead of two. trait-extraction PR reviewer's first design call. 3. **Extensibility note — the four-variant enum is not a permanent - ceiling.** Currently-invisible failure modes — `unlimited_context`'s + ceiling.** Currently-invisible failure modes — `windowed_checkpoint`'s "request crossed an uncheckpointed window boundary" (collapsed into generic `process()` None today), `boundary_per_layer`'s "calibration record missing for policy fingerprint" (a @@ -1252,7 +1016,7 @@ were implementation). measurement → policy generation → end-to-end KL validation) is not in tree. Per spec Phase 1 of [boundary-per-layer-engine.md](../larql-inference/docs/specs/boundary-per-layer-engine.md). -- **Page-aligned KV slabs for `unlimited_context`.** The current +- **Page-aligned KV slabs for `windowed_checkpoint`.** The current `CheckpointStore` uses owned `Vec` per layer per checkpoint; a hugepage-backed slab would cut allocation churn and improve thermal steadiness during 370K-token replays. @@ -1344,617 +1108,6 @@ were implementation). the dispatch — but the executor + state-policy separation (Phase 2 spec) makes composition cleaner. -## Closed (recent) - -- **2026-05-24 — Multi-modal engine seam (ADR-0023).** `KvEngine` gains - `supports_multimodal()` (default false) + `prefill_from_hidden(weights, - ffn, initial_hidden: &Array2) -> Result, EngineError>`. - `StandardEngine` is the first (and currently only) MM-capable engine. - Other engines inherit the default-false convention — they remain - text-only until each individually implements the new method. - `AnyEngine` forwards both methods. `generate_with_engine_from_hidden` - wrapper shares the decode loop with `generate_with_engine`. Dispatch - helpers `kv_prefill_from_hidden_via_dispatch` (sync + async) hoist the - embed step out of the prefill loop so both text-only and MM inputs - follow the same layer-forward path. The eventual end state: every - engine implements `prefill_from_hidden` and `prefill(token_ids)` becomes - a thin wrapper. No timeline on the seven-engine migration. - -- **2026-05-24 — Sibling trait extraction LANDED.** `KvEngine` - `Option` returns are gone; the typed `EngineError` enum lives in - `larql-inference::kv_engine` alongside the new `RetrievalEngine` - trait + `AnyEngine` dispatch enum. The two-harness silent-drop / - panic disagreement (`accuracy_suite/runner.rs` vs - `bench/engine_runtime.rs`) is resolved at the type level. - - **Trait surface:** all 8 `KvEngine` impls (`standard`, `no_cache`, - `markov_residual`, `markov_residual_codec`, `unlimited_context`, - `turbo_quant`, `boundary_kv`, `boundary_per_layer`) return - `Result, EngineError>` on `prefill` / `decode_step` / - `*_quant` / `*_via_executor`. Apollo moves to the new - `RetrievalEngine` trait (`prefill(weights, token_ids)` / - `decode_step(weights, token_id)` — no `FfnBackend`, no per-step K/V). - - **EngineError variants** (exhaustive, no `#[non_exhaustive]`, - thiserror): `EmptyPrompt`, `BackendUnavailable`, `RetrievalMiss - { reason }`, `InvariantViolation { what }`, `BackendFailure - { details }`. Per Finding 2, `InvariantViolation` and `BackendFailure` - are kept as two top-level variants to preserve the alert-routing - distinction (a dispatch bug vs a kernel/data failure). The accuracy - harness's `ScoreOutcome` mirror followed suit: - `SkippedInternalError` → `SkippedInvariantViolation` + - `SkippedBackendFailure` (load-bearing JSON schema change for - downstream observability). - - **AnyEngine** (`AnyEngine::Kv(Box) | - Retrieval(Box)`) is the harness boundary type. - Forwarding methods (`prefill` / `decode_step` / `prefill_quant` / - `decode_step_quant` / `*_via_executor`) take the superset of args - from both surfaces and ignore the irrelevant ones on the retrieval - arm. This intentionally walks back the original "don't lift a common - shape" plan — the harness scalability won out, since the alternative - is N×2 match arms per call site as more retrieval engines land. - - **Bench harness merged.** `run_engine` + `run_engine_q4k` collapsed - into one `run_engine(weights, index: Option<&VectorIndex>, ...)`. - When `index = Some` the dispatch goes through `prefill_quant` - (quant-agnostic — the vindex's format flows through the engine); - when `None` the dense `prefill` path runs. FFN selection: dense - defaults to `WeightFfn`, quant defaults to `NullFfn` (preserves the - pre-merge Q4K behaviour). `--ffn-policy` honored on dense, logged - as not-yet-honored on quant due to the `&mut weights` vs - `&weights`-borrowing-router conflict (unchanged from pre-merge). - - **Coverage debt:** one re-introduced baseline at - `markov_residual/engine.rs` (89.5% vs 90% floor). The remaining - uncovered lines are all `.ok_or_else(|| BackendFailure)?` - constructions that only fire when an internal helper - (`rs_decode_step_walk`, `recompute_kv`, `executor.run_*_layer`) - returns None. Triggering those requires the mock `EngineBackend` - infrastructure that the 2026-05-24 coverage-clearance explicitly - deferred; the baseline tracks the debt rather than gold-plating - ahead of need. - - **Outcomes.** Test count larql-kv lib: 712 → 726 (+14). Workspace - builds clean. `make larql-kv-ci` passes (fmt + clippy + tests + - fresh coverage policy with 1 baseline). Apollo's `executor.rs` - deleted (~150 lines of dead code from the old KvEngine `*_via_executor` - impls). Closes [`docs/state-policy.md`](docs/state-policy.md) §8 - Open Question 1 ("Where does Apollo's fallback live?"); also closes - the interim `ffn_backend` JSON limitation flagged in Item 1 of the - 2026-05-24 accuracy harness work. - - **Follow-ups** *(deferred to keep this PR atomic)*: - - Mode 5 / Graph-Grounded engine lands as a `RetrievalEngine` impl - (was blocked on this refactor). - - Q4K `--ffn-policy` honoring (was waiting on the same - `&mut weights` borrow conflict — still present after the merge - because the trait surface still takes `&mut weights` for lazy - dequant). - - `RemoteWalk` build path (~200 lines, standalone — was the second - blocked item). - - `markov_residual/engine.rs` coverage debt + mock `EngineBackend` - infrastructure (deferred per "Sub-project A" of the previous - coverage push). - -- **2026-05-24 — Coverage debt CLEARED.** All six files below the - 90% per-file floor lifted; `make larql-kv-coverage-policy` passes - against fresh `summary.json` regeneration. Workspace total 95.62% - lines, 61/61 files at ≥90%, 0 debt baselines remaining. - - Files lifted (pre → post): `turbo_quant/dispatch` 9.35→97.85%, - `boundary_per_layer/dispatch` 7.95→93.57%, `unlimited_context/dispatch` - 59.09→97.24%, `markov_residual/dispatch` 77.51→96.78%, - `markov_residual_codec/dispatch` 80.68→97.72%, - `markov_residual/compute` 86.85→95.30%. - - Approach inverted both pre-baked design assumptions: - - **No new shared mock `EngineBackend`** — `CpuBackend` (via - `cpu_engine_backend()`) already implements `coarse_*_with_state` - when driven against the synthetic Q4K fixture - (`make_test_q4k_weights` + `make_test_q4k_vindex`), so every - dispatch happy-path tested end-to-end without new infrastructure. - - **No `serial_test` crate** — env-gated paths - (`LARQL_MARKOV_WALK_KV_*`, `LARQL_W10_DISABLE`) instead gained - a per-thread `RefCell` override that production helpers consult - *before* `std::env::var`. Tests inject without touching the - process env; no race with other parallel tests. New helpers: - `compute.rs::set_markov_env_override(...)`, - `engines/mod.rs::set_w10_disabled_override(...)` (both - `#[cfg(test)]` only). - - Test deltas: larql-kv lib 663 → 712 (+49). Zero regressions - (5/5 successive `cargo test -p larql-kv --lib` runs green after - the thread-local override fix; pre-fix the env-var-setting tests - produced flaky `cold_kv.is_some()` failures in unrelated codec - tests via process-env race). `make larql-kv-ci` passes end-to-end. - -- **2026-05-24 — Accuracy harness honesty + FFN policy cross-product - LANDED.** Multi-PR arc that turns the accuracy suite from "silent - drop on engine miss" into a discriminating cross-product harness: - - - **Item 1 — accuracy schema fix** (commit `07684457`). - `ScoreOutcome` enum (exhaustive, flat-tagged serde, mirrors the - future `EngineError` taxonomy). `PromptScore` / `ConflictScore` - gain `outcome` field + `Option` score payload with - `served()` / `skipped()` constructors enforcing - correlated-optionality. `StrategySplit` gains `*_served` + - `*_served_rate` per axis as required-companion fields to - `*_match_rate`. `compute_strategy_split` filters on served subset - (counting skips as zero would punish honest reporting). Replaces - `filter_map` silent-drop in all three drivers. Surfaces Apollo's - store-miss rows as `SkippedRetrievalMiss` instead of dropping. - `EngineKind::supported_names()` replaces hard-coded six-engine - error string at two bench sites. - - - **Item 2 v0 — `FfnBackendKind` parser + `FfnLayerPolicy` - (in `larql-inference::ffn_policy/`).** New crate-shape: - `FfnBackendKind` (Dense / Walk{k} / RemoteWalk / Null), - `RoutingPredicate` (All / Layers / Otherwise), `FfnLayerPolicy` - with from_spec parser supporting per-layer routing - (`{walk:k=100}@layers=14-27;{dense}@otherwise`). - Construction-errors on overlapping ranges; exhaustive enums; - typed error taxonomy (`PolicyParseError` / - `PolicyValidationError`). Module lives in `larql-inference` not - `larql-kv` — FFN policy is the FFN axis, not the KV axis. - - - **`build_router` slice — `ValidatedFfnLayerPolicy` newtype + - `BoundFfnRouter`.** Type-system enforcement of "validate before - build" via non-public constructor. `BoundFfnRouter<'a>` owns its - backend instances (`Vec>`) so callers - don't manage backend lifetimes alongside the router's. `impl - FfnBackend for BoundFfnRouter` delegates per-layer via the - trait's existing `layer: usize` parameter — drop-in for the - `&dyn FfnBackend` surface every engine already takes. Design - rationale: `larql-inference/docs/ffn-build-router.md`. - - - **Cross-product harness + typed axis columns.** `accuracy_cmd` - iterates `kv_engine × ffn_backend` cross-product via - `FfnLayerPolicy::split_specs` (comma-separated, brace-aware, - re-parse fallback for kv-comma forms like - `remote-walk:endpoint=X,wire=Y`). New `EvalLabels<'a>` struct - bundles `(kv_engine, ffn_backend, strategy)` for clean signatures. - `PromptScore` / `ConflictScore` / `StrategySplit` gain explicit - `kv_engine: String` + `ffn_backend: String` columns alongside - `strategy`. `format_strategy_split` grows a two-axis layout - (`KV engine` + `FFN backend` columns) when any row has - `ffn_backend != "dense"`; default no-`--ffn` runs keep the - historical single-`Strategy`-column layout. Closes the - interim-`ffn_backend`-as-user-input limitation noted in Item 1's - ROADMAP entry. - - - **CLI wiring.** `larql accuracy --ffn dense,walk:k=100,'{walk:k=100}@layers=14-27;{dense}@otherwise'` - now runs the cross-product in one invocation. Vindex loaded - lazily — only when a Walk binding is present. - `larql bench --ffn-policy ` honors the policy on the - non-Q4K (CPU) path; Q4K path accepts the flag but doesn't - honor it yet (P1 follow-on above). - - - **Apollo into accuracy default engines.** `--engines` default - now includes `apollo`. The schema fix above means Apollo's - store-miss rows show `served_rate < 1.0` rather than silent - drops — diagnostic rather than misleading. - - - **Module splits.** `accuracy_suite/runner.rs` (2050 lines) split - into `accuracy_suite/runner/` folder (6 files: `types` / - `scoring` / `drivers` / `aggregate` / `legacy` / `mod`). Same - pattern that produced the `ffn_policy/` folder split in - `larql-inference`. - - - **Coverage lift across 5 engine files.** Pre-existing engine - internals had drifted below 90%. Lifted with synthetic-weights - + CPU-backend tests: `boundary_per_layer/cold_tier.rs` - (88→100%), `executor.rs` (85→90.6%), `walk.rs` (84→95%), - `engine.rs` (83→90%), `markov_residual/store.rs` (86→99.6%). - `markov_residual/compute.rs` partially lifted (81→86.85%); - full lift gated on `serial_test` for env-var paths. - Discovered the gate had been passing against a stale JSON — - fresh `make larql-kv-coverage-summary` is now required to - surface debt. See "Coverage debt" section above for the - remaining 6 files. - - Test deltas across the arc: larql-kv lib 595 → 663 (+68), - larql-inference lib 1086 → 1102 (+16). Zero regressions. Clippy - clean. Aggregate ~3,500 lines of code + tests added across - `larql-kv` and `larql-inference`. - - ROADMAP entry for the sibling trait extraction (P0 above) - references "Item 1 in the conversational priority queue" — Item 1 - is the schema fix above. Mode 5 work is still gated on that P0 - refactor landing. - -- **2026-05-18 — W8.2 (doubling-capacity K/V in `markov_residual` + - `markov_residual_codec`) LANDED: 2.4× decode speedup at 1000 tokens.** - Lifted the W8 pre-allocation pattern from `unlimited_context` to the - two unbounded-window engines. Since `max_window=None` rules out a - fixed pre-alloc, both stores now use a doubling-capacity strategy - via three private helpers in each engine: - - `window_capacity(prompt_len, window_size)` — initial cap is - `max(window, prompt_len)` if windowed, else - `max(prompt_len * 2, 64)`. - - `grow_capacity_2d(src, len, cap)` — allocate `[cap, cols]` once - at prefill, copy the prefill rows in. - - `append_row(buf, row, len)` — in-place `slice_mut(s![len..len+1, - ..]).assign(row)` when `len < cap`; otherwise double capacity, - copy the live rows, then assign. Amortised O(1) per append vs the - O(n) per step the previous `Array2::zeros((n+1, dim))` pattern - paid. - - Store changes (both `RsStore` and `RsStoreCodec`): - - New `pub hot_len: usize` field — logical row count, separate from - `stored[l].shape()[0]` (which is now capacity ≥ hot_len). - - `window_tokens()`, `memory_bytes()`, `clip_layer` / - `clip_layer_overflow` updated to use `hot_len`. - - New `finalise_hot_len_after_clip()` — must be called after every - per-layer clip loop. (Subtle bug fix during impl: setting - `hot_len = window` *inside* the per-layer loop made layers 2..N - see `rows == window` and skip their clips, dropping half the - cold-tier payload. Two existing tests caught this.) - - Bench (Gemma 3 4B Q4K, Metal, M3 Max): - - **1000-tok**: - - `markov-rs`: 24.8 → **58.7 tok/s (+137%)** - - `markov-rs-codec`: 25.7 → **57.2 tok/s (+123%)** - - `unlimited-context`: 49.5 → **57.4 tok/s (+16%)** (variance - recovery from previous run + sympathy from the codepath audit) - - `standard` unchanged at 64.1 (untouched) - - **50-tok**: - - `markov-rs`: 77.1 → **88.9 tok/s (+15%)** - - `markov-rs-codec`: 77.5 → **88.8 tok/s (+15%)** - - All three cached-state engines now cluster within 11% of standard's - 64.1 tok/s ceiling at 1000 tokens. The doubling-capacity scales - linearly with seq_len: at 50 tok the saved alloc bytes are small - (~400 KB/step); at 1000 tok they're ~8 MB/step. The 137% win at - long context is the alloc churn that pre-W8.2 was hiding behind - prefill cost. - - CPU walk + executor fallback paths (`rs_decode_step_walk`, - `rs_decode_step_codec_walk`, `process_via_executor`) still allocate - per step — they're not on the hot path for the bench. Defensive - consistency: every legacy RsStore/RsStoreCodec constructor sets - `hot_len` from `stored[0].shape()[0]` so non-dispatch paths see a - consistent invariant. - -- **2026-05-18 — Step 9 (iterative Metal `coarse_prefill_with_state`) - LANDED: ~10× prefill speedup on every state-dump engine.** - Pre-Step 9, `MetalBackend::coarse_prefill_with_state` defaulted to - the trait's `coarse_prefill` (no per-layer state dump); engines saw - `state.is_complete_for() == false` and fell back to the CPU walk - (~2.7 s on Gemma 3 4B). The new impl pre-allocates `[seq_len, - hidden]` and `[seq_len, kv_dim]` per layer (W8-style alloc at - source for prefill too), resets + preallocates the Metal K/V cache, - then iterates `fused_decode_step_with_state` per prefill token, - writing the dump into the pre-allocated row position. - - Bench (Gemma 3 4B Q4K, Metal, M3 Max, "The capital of France is", - 5 prefill tokens): - - `markov-rs` prefill: 2757 → **254 ms** (10.9×) - - `markov-rs-codec` prefill: 2564 → **249 ms** (10.3×) - - `unlimited-context` prefill: 2760 → **256 ms** (10.8×) - - `turbo-quant` prefill: 2750 → **334 ms** (8.2×) - - Predicted ~45× (5 × 12 ms decode time) didn't materialise because - each iterative `fused_decode_step_with_state` carries per-token - state-dump readback overhead. Remaining ~250 ms is 5 × ~50 ms - per-iter + fixed setup. Further closure needs a single-kernel - prefill that dumps state for all positions in one shot — separate - Metal-kernel surgery. - - Decode steady-state also moved (W8 + Step 9 compound): - - `unlimited-context`: 82.7 → **89.2 tok/s** (fastest cached-state - engine; within 10% of `standard`'s 99.2 ceiling) - - `markov-rs`: 75.3 → 77.1 tok/s - - `markov-rs-codec`: 79.0 → 77.5 tok/s - -- **2026-05-18 — W8 (pre-allocated K/V buffer in `unlimited_context`) - LANDED: 58% of decode-CPU alloc churn removed.** - samply flamegraph on `unlimited_context:window=1024 --tokens 1000` - (post-W7) surfaced an unexpected hot path: 21% `__bzero` + 19% - `ndarray::zip_mut_with_same_shape` + 18% `madvise` = **58.5% of - main-thread CPU** spent on `Array2::::zeros((n+1, kv_dim))` + - `slice_mut().assign(k_old)` + `slice_mut().assign(k_new_row)` - inside `decode_step_via_dispatch` — 68 allocations per token - (34 layers × 2), each growing linearly with `n`. - - Fix: pre-allocate `Array2::zeros((window_size, kv_dim))` per layer - once at prefill (in `try_prefill_via_dispatch`), track a single - `current_window_kv_len: usize` counter, and append in the hot path - via `slot.0.slice_mut(s![pos..pos+1, ..]).assign(k_new_row)`. One - small `kv_dim`-sized copy per layer per side, zero alloc per step. - Readers (`close_window`, `current_kv_bytes`) updated to use the - counter instead of `k.shape()[0]`; CPU walk fallback paths set the - counter defensively from the returned narrow-array shape. - - Bench (Gemma 3 4B Q4K, Metal, M3 Max): - - 50-tok: `unlimited-context:window=256` 82.7 → **86.6 tok/s - (+4.7%)** vs `standard`'s 99.4 (gap closed ~50%) - - 1000-tok: `unlimited-context:window=1024` 17.39 ms vs `standard`'s - 15.74 ms → 1.65 ms gap (vs pre-W8 estimated 5-10 ms slope from - `Array2::zeros((n+1, …))` growing linearly with `n`) - - Post-W8 flamegraph: the `__bzero` / `zip_mut_with_same_shape` / - `madvise` triple is **gone from the top-20**. Remaining main-thread - CPU is dominated by `__psynch_cvwait` (Metal GPU wait, - irreducible), `synthesize_lm_head_kquant` (prefill — separate - ~2.5 s regression flagged elsewhere), and generic `Map::fold`. - - The optimisation is engine-local (`larql-kv/src/engines/unlimited_context/engine.rs`) - with no surface change. Same pattern can be lifted to - `markov_residual` / `markov_residual_codec` / `turbo_quant` once - their state-policy shape is clarified — they use the same - `Array2::zeros((n+1, kv_dim))` pattern but have unbounded windows - by default, so the pre-allocation needs a growable strategy - (doubling-capacity Vec-style) rather than fixed window size. - Tracked as W8.2 candidate. - -- **2026-05-18 — W7 (blit-encoder fusion) LANDED: per-layer commit - overhead removed; +30-48% across cached-state engines.** - Modified `decode_token_with_moe_split_fn` in - `larql-compute-metal/src/decode/mod.rs` to pre-allocate per-layer - staging buffers (k / v / h-in) when `state_dump` is `Some`. The - layer loop blits `k_out` / `v_out` / `h_buf` into the staging - buffers inside the same command buffer (`new_blit_command_encoder` - + `copy_from_buffer`) instead of forcing per-layer commit + wait + - CPU read. The single final commit at the bottom of the function - flushes everything; reads happen once after that, draining staging - into `state_dump`. Metal's command-buffer encode ordering - guarantees blit reads see the settled compute writes. - - Measured (Gemma 3 4B Q4K, Metal, M3 Max): - - `standard` (control, no state_dump): 105.9 → 99.4 tok/s (noise) - - `markov-rs`: 58.0 → **75.3 tok/s (+30%)** - - `markov-rs-codec`: 58.4 → **79.0 tok/s (+35%)** - - `unlimited-context` (window=256): 56.0 → **82.7 tok/s (+48%)** - - `turbo-quant` (4-bit, 10-tok bench): 33.0 → **37.7 tok/s (+14%)** - - Engine-cost decomposition post-W7: ~10 ms Metal kernel compute + - ~3 ms CPU glue. The remaining gap to `standard`'s 99 tok/s is - pure CPU-side state-update work (state Vec→Array2 conversion, - appends). Closure path: in-place state updates / pre-allocated - buffers (W8 candidate). - - Edge cases worth noting: - - `standard` doesn't touch state_dump → blit branch is dead code - → 0× regression confirmed. - - `turbo_quant`'s codec inner loop is the dominant per-token cost; - the saved 1.7 ms commit overhead is a smaller fraction. - - The `unlimited_context` +48% win reflects its lighter post- - kernel CPU work (just append to `current_window_kv`); engines - with heavier post-kernel work see smaller relative gains. - -- **2026-05-17 night — W1-GPU steps 4 + 6 LANDED: unlimited_context + - turbo_quant now route through dispatch on Metal.** - Same pattern as steps 5: each engine gains `try_prefill_via_dispatch` - / `decode_step_via_dispatch` helpers that read per-layer captured - state and update engine-specific state policy. - - **turbo_quant**: state.k_new/v_new per layer feeds the - WHT+Lloyd-Max codec via `CompressedLayer::compress` (prefill) - and decompress→append→recompress (decode). Bench: **19.6 → - 33.0 tok/s (+68%)** on Metal. Memory stays at 0.6 MB hot - (compression intact). - - **unlimited_context**: state.k_new/v_new appends to - `current_window_kv` per layer; window auto-close at - `window_size` tokens fires the legacy `close_window` checkpoint - emit. Bench: **28 → 56.0 tok/s on Metal (+98%)** at - `window=256` (Gemma 3 4B, M3 Max, 50-token decode). Hot state - 15.7 MB tracks the engine-side window shadow (see KvHandle - eviction note below). - - Engine memory note: with W1-GPU active, the backend's internal K/V - cache grows unboundedly alongside each engine's shadow state. This - defeats the memory benefit of `unlimited_context` / - `markov_residual_codec` at long contexts. Follow-up: expose a - `KvHandle::evict_oldest(n)` API on `KvDispatch` so engines can - bound the backend cache to match their window. -- **2026-05-17 night — W1-GPU step 2 LANDED: Metal per-layer state - dump → 2.1× decode speedup on markov-rs + codec.** - Modified `decode_token_with_moe_split_fn` in - `larql-compute-metal/src/decode/mod.rs` to accept an optional - `state_dump: Option<&mut DecodeStateDump>` parameter. When active, - the layer loop: - 1. At top of layer L: pushes `x` (for L=0) or reads `h_buf` (for - L>0, settled by the previous layer's commit) into - `state.h_in_per_layer`. - 2. At bottom of layer L: forces `enc.end_encoding()`, `cmd.commit()`, - `wait_until_completed()`, reads `k_out` / `v_out` (scratch - buffers reused across layers) into - `state.k_new_per_layer` / `v_new_per_layer`, then restarts - command buffer + encoder for the next layer. - - Trait wiring: new `DecodeBackend::decode_token_with_state_dump` - method (default falls back to plain `decode_token`); MetalBackend's - trait impl routes through the new kernel function when `state` is - `Some`. Inference layer adds `fused_decode_step_with_state` + - `MetalBackend::coarse_decode_step_with_state` / - `coarse_prefill_with_state`. Engines (markov_residual, codec) - inherit the Metal acceleration automatically — no engine-side - changes from step 5. - - Measured (Gemma 3 4B Q4K, Metal, M3 Max, 10-token decode): - - `markov-rs`: 27.0 → **57.7 tok/s** (+114%) - - `markov-rs-codec`: 27.8 → **57.5 tok/s** (+107%) - - `standard` (fused control): 100.8 tok/s (unchanged) - - Per-token cost: ~17 ms = 10 ms Metal compute + ~1.7 ms commit - overhead (50 µs × 34 layers) + ~5 ms engine state update / CPU - glue. The remaining gap to standard's 100 tok/s is the - per-layer commit cost; a follow-up could use blit-encoder - switches inside a single command buffer to eliminate the - commit overhead and lift toward 80-100 tok/s. - - Prefill cost: ~2.8 s on Metal (CPU walk for state seeding + - Metal `fused_prefill` for backend cache). One-shot; doesn't - affect decode steady-state. Future optimisation: per-position - per-layer K/V dump on the Metal prefill side to skip CPU walk. -- **2026-05-17 night — W1-GPU infrastructure (decode trait surface + - CPU impl + engine wiring; Metal kernel modification deferred).** - Three layered changes landed end-to-end: - - **Trait surface (`KvDispatch`):** new `coarse_prefill_with_state` / - `coarse_decode_step_with_state` methods take - `Option<&mut PerLayerDecodeState>`. Default impls delegate to the - non-state variants, so unmigrated backends keep working. - - **`DecodeBackend` trait + `DecodeStateDump` struct** added in - `larql-compute` for the substrate-level surface. Same default- - delegation pattern. - - **CPU implementation** (`predict_kquant_prefill_with_state` / - `predict_kquant_decode_step_direct_with_state`): threads per-layer - state capture through the existing per-layer walk at zero - re-compute cost. Parity test in - `kv_dispatch::cpu::coarse_decode_step_with_state_populates_and_matches_plain` - asserts cached and non-cached outputs match within f32 rounding - and per-layer shapes (`[1, hidden]`, `[1, kv_dim]`) are correct. - - **Engine wiring** for `markov_residual` and - `markov_residual_codec`: `try_prefill_via_dispatch` / - `decode_step_via_dispatch` route through the new - `coarse_*_with_state` API when the backend implements it. State - capture feeds `RsStore::stored` (residuals) and `hot_kv` (W2 - cache) in a single backend call. Legacy walk path stays as the - fallback when state isn't populated (e.g. on backends that - haven't migrated yet — currently `MetalBackend`). Gated on - `supports_direct_matvec_decode` so non-Q4K test fixtures skip - the dispatch path. 113 markov tests pass. - - **CPU bench numbers stay parity** post-W1-GPU step 5: - markov-rs 27.4 tok/s, codec 26.6 tok/s — same as W2 (W1-GPU on - CPU just changes the code path, not the compute; CPU was already - at the M3 Max compute ceiling). - - **What's NOT done**: `MetalBackend::coarse_*_with_state` still uses - the default delegation (state stays empty), so engine falls back - to walk on Metal — no GPU speedup yet. The real Metal acceleration - requires modifying - `larql-compute-metal::decode::decode_token_with_moe_split_fn` - (200+ lines) to thread per-layer dump buffers + blit-encode steps - into the existing single command buffer. Two implementation - shapes have been scoped: - 1. **Blit-encoder switches per layer**: cheapest in steady-state - (~tens of µs per layer); requires careful encoder lifecycle - management within the existing kernel function. - 2. **Per-layer commit + CPU readback**: simpler (mirror the - existing `stage_timing_split` pattern); costs ~50µs/layer × - 34 = ~1.7ms/token overhead. Projected ceiling: 50-80 tok/s - (vs CPU's 27 tok/s ceiling, vs `standard`'s 102 tok/s fused). - - Choice between shapes is open. The trait surface, CPU impl, and - engine wiring are all stable and don't change regardless of which - Metal-side approach lands. -- **2026-05-17 night — W2: hot K/V cache for `markov_residual` and - `markov_residual_codec`.** Added `hot_kv: Option>` - to both `RsStore` and `RsStoreCodec`; prefill captures K/V from - the per-layer forward pass (previously discarded) and stashes it; - decode appends one row per layer via the existing - `run_attention_block_decode_step_backend` return tuple. On - window-overflow `clip_layer` slices `hot_kv` consistently with - `stored`; for `markov_residual` (lossless cold tier) the evicted - K/V rows merge directly into `cold_kv` (no `recompute_kv` call - needed); for `markov_residual_codec` (lossy bf16 cold tier) - `cold_kv` is invalidated on overflow so the next step recomputes - against the codec-decoded residual. Bench: `markov_residual` - 4.7 → 26.8 tok/s (5.7×); `markov_residual_codec` 5.0 → 27.5 tok/s - (5.5×). Both now sit on the `unlimited_context` curve. Engine - contract preserved — drop `hot_kv` and the next step recomputes - from `stored` (via_executor path takes this fallback). Hot-state - memory grew from 5.3 → 10.8 MB; still ~50× smaller than - `standard`'s full KV cache. Parity test - (`decode_step_quant_w2_cached_matches_recompute_from_residuals`) - asserts the cached and recompute paths agree within fp rounding. -- **2026-05-17 night — W7: per-engine profiler wired on the quant - path.** `EngineProfiler` now populates from `rs_decode_step_walk` - (markov_residual), `rs_decode_step_codec_walk` - (markov_residual_codec), `rs_extend_from_checkpoint_quant` - (unlimited_context), and `decode_step_quant_cpu` (turbo_quant). - Each engine's `stage_summary()` returns `Some(...)` when - `with_profiling(true)` is set. `larql bench --profile --engine - ` now produces a per-stage attribution table per engine. - First measurement run produced the bottleneck-diagnosis table in - the P0 section above, which inverted two of the pre-profile - guesses: codec overhead in turbo_quant was ~25% not ~80%, and K/V - recompute (W2 target) was the dominant cost on markov_residual - (~80%) not dispatch (W1 target). Sequencing in P0 revised - accordingly. -- **2026-05-17 night — `_q4k` → `_quant` on remaining internal - function names.** The trait-surface renames earlier today - (`prefill_q4k` → `prefill_quant`, `has_q4` → - `supports_quant(format)`, `q4k` → `kquant` storage) missed the - per-engine implementation wrappers: - `unlimited_context::process_q4k`, - `unlimited_context::extend_current_q4k`, - `extend::rs_extend_from_checkpoint_q4k`, - `turbo_quant::decode_step_q4k_cpu` / - `turbo_quant::prefill_kquant_cpu`. All renamed to `_quant` since - they dispatch on whatever format the vindex carries, not Q4_K - specifically. -- **2026-05-17 night — Fused-bypass strip: engines are now engines.** - Every per-layer engine (`markov_residual`, `markov_residual_codec`, - `unlimited_context`, `turbo_quant`) had a hidden - `if let Some(h) = fused_prefill(...) { return Some(h); }` short- - circuit at the top of `prefill_quant` / `decode_step_quant`. The - short-circuit meant `--engine markov-rs` on Metal silently ran - `StandardEngine`'s fused kernel instead — five engines tied at - ~103 tok/s with `hot=0.0MB`, masking every state-policy difference - and making per-layer optimization invisible. Cut: removed every - short-circuit; deleted dead `metal_prefill_done` + `force_walk` - fields and `with_force_walk` builders; dropped the pub(crate) - `fused_prefill`/`fused_decode_step` re-exports from - `unlimited_context::engine` (only `StandardEngine::coarse_prefill` - uses the underlying `larql_inference::vindex::fused_prefill` now, - via `ComputeBackend::coarse_prefill`). `StandardEngine` remains the - default engine and the only home of the fused fast path. Bench now - reports honest numbers: standard 104 tok/s, markov-rs 3.6, codec - 4.3, unlimited-context 25.6, turbo-quant 3.9 — every per-layer - engine reports non-zero `hot=` memory because their state - structures actually materialise. The 25-30× standard-vs-per-layer - gap is the new optimization frontier; previously it was invisible - because every engine was running the same kernel under different - labels. -- **2026-05-17 evening — Phase-2 migration completed for the remaining - three engines.** `unlimited_context`, `turbo_quant`, and `apollo` all - override `*_via_executor` methods and honor the caller-supplied - `FfnBackend`. `CountingFfn` stub tests prove per-(token, layer) - dispatch through the caller's backend. Same push cleared every - `coverage-policy.json` debt baseline: all 43 files in src/ at ≥90% - lines, workspace total 95.55%. `larql bench --ffn http://shard:8080` - now routes through the remote shard for every per-layer engine - instead of silently constructing a local `WalkFfn`. -- **2026-05-17 — Phase 2 engine migration to `LayerExecutor`.** Four - engines (`markov_residual`, `markov_residual_codec`, - `boundary_per_layer`, `no_cache`) override `*_via_executor` methods. - They drive per-layer dispatch through `executor.run_*_layer` and - honor the caller's `FfnBackend`. `CountingFfn` stub tests prove the - FFN parameter is no longer silently ignored. Bench has - `--via-executor` flag; demoed on Gemma 3 4B Q4K showing the codec - engine's 50% cold tier saving (22.9 MB → 11.5 MB). -- **2026-05-17 — `LayerExecutor` trait + `LocalWalkExecutor`.** New - abstraction in `larql-inference::layer_executor` separating state - policy (engine concern) from execution strategy (executor concern). - Spec at - [engine-state-vs-execution.md](../larql-inference/docs/specs/engine-state-vs-execution.md). -- **2026-05-17 — `q4k` → `kquant` storage rename.** K-family storage - slots (`attn_q4k`, `interleaved_q4k`, manifests, setters, loaders) - renamed for consistency with accessor naming (`attn_kquant_layer_data`). - Q4_0 and Q8 slots unchanged. ~60 sites touched. -- **2026-05-17 — `has_q4()` → `supports_quant(format)`.** Per-format - predicate on `ComputeBackend`. 79 call sites migrated to - `supports_quant(QuantFormat::Q4_K)`. Enables future Q6_K / FP4 - fused-pipeline backends without trait extension. -- **2026-05-17 — `KvEngine::prefill_q4k` / `decode_step_q4k` → - `prefill_quant` / `decode_step_quant`.** Trait surface naming made - quant-agnostic. 112 sites updated. Internals that are genuinely - Q4K-specific kept their names. -- **2026-05-17 — `metal_fused_*` → `fused_*` rename.** The "metal" - prefix was a lie: `CpuBackend` implements `prefill_q4` and - `decode_token` via its C Q4 kernel. Aliases in - `unlimited_context::engine` follow. -- **2026-05-17 — `BoundaryKvEngine`, `MarkovResidualCodecEngine`, - `BoundaryPerLayerEngine` shipped.** All three new engines have - contracts in `crates/larql-inference/docs/specs/`. Per-file coverage - ≥94 % lines on every new file. Bench demoed end-to-end on Gemma 3 4B, - Gemma 4 E2B, 26B-A4B, 31B, Qwen3 0.6B (dense + Q4K). -- **2026-05-09 — Initial extraction.** `engines/` carved out of - `larql-inference` into the new `larql-kv` crate. ~5,540 LOC moved with - no semantic changes. All four engines + `KvEngine` + accuracy / - profiler helpers now ship from this crate. - ## Non-goals - **Sampling.** Engines return hidden states; sampling lives in @@ -1967,3 +1120,4 @@ were implementation). `KvEngine` trait references `ModelWeights` directly. Generalising to state-space models or RNNs is not on this roadmap; rebuilds are cheap and that effort would belong in larql-inference's layer-graph surface. + diff --git a/crates/larql-kv/benches/engine_decode.rs b/crates/larql-kv/benches/engine_decode.rs index 67f9a4f53..d91a17d1e 100644 --- a/crates/larql-kv/benches/engine_decode.rs +++ b/crates/larql-kv/benches/engine_decode.rs @@ -8,14 +8,12 @@ //! comparator was deprecated in 2026-05-16 — it measured random-vector //! encode/decode, not real decode steady-state.) //! -//! Engines covered: -//! - `standard` (production K/V cache, unbounded) -//! - `standard:window=4` (sliding-window K/V) -//! - `no-cache` (full re-forward per step, debug fallback) -//! - `markov-rs` (residual-stream replacement) -//! - `unlimited-context` (per-window K/V checkpoints) -//! - `turbo-quant-4bit` (WHT + Lloyd-Max 4-bit codec) -//! - `apollo` (boundary-residual injection) +//! Engines covered: whatever [`EngineKind::bench_specs`] lists. That is +//! the single source of truth, pinned against the engine roster by +//! `bench_specs_cover_every_benchable_engine` in the lib tests — so a +//! new engine cannot land without either a bench arm or a written +//! reason it can't have one. This list used to be hand-maintained here +//! and had silently fallen to 7 of 9 engines. use criterion::{criterion_group, criterion_main, Criterion}; use larql_inference::cpu_engine_backend; @@ -23,35 +21,22 @@ use larql_inference::ffn::WeightFfn; use larql_inference::test_utils::make_test_weights; use larql_kv::EngineKind; +/// Engines to bench, from [`EngineKind::bench_specs`]. The spec string +/// doubles as the criterion benchmark id. +/// +/// Apollo is excluded by that list (see `EngineKind::bench_excluded_names`): +/// with no store attached its `prefill` fails closed with `RetrievalMiss` +/// before touching the model, so benching it here timed the error return +/// — ~65 ns against ~16 µs for `standard`, which read as a 250x win. fn all_engines() -> Vec<(&'static str, EngineKind)> { - vec![ - ("standard", EngineKind::Standard { window_size: None }), - ( - "standard-window-4", - EngineKind::Standard { - window_size: Some(4), - }, - ), - ("no-cache", EngineKind::NoCache), - ( - "markov-rs", - EngineKind::MarkovResidual { window_size: None }, - ), - ( - "unlimited-context", - EngineKind::UnlimitedContext { window_size: 4 }, - ), - ("turbo-quant-4bit", EngineKind::TurboQuant { bits: 4 }), - ( - "apollo", - EngineKind::Apollo { - injection_layer: 1, - inject_coefficient: 8.0, - top_k: 4, - bos_token_id: None, - }, - ), - ] + EngineKind::bench_specs() + .iter() + .map(|spec| { + let kind = EngineKind::from_name(spec) + .unwrap_or_else(|| panic!("bench spec {spec:?} failed to parse")); + (*spec, kind) + }) + .collect() } fn bench_prefill(c: &mut Criterion) { @@ -64,7 +49,12 @@ fn bench_prefill(c: &mut Criterion) { group.bench_function(name, |b| { b.iter(|| { let mut engine = kind.clone().build(cpu_engine_backend()); - let _ = engine.prefill(&weights, &ffn, &prompt); + // Unwrap, don't discard: a `let _ =` here would happily + // time an engine that bailed out before doing any work + // and report the error path as a stellar result. + engine + .prefill(&weights, &ffn, &prompt) + .unwrap_or_else(|e| panic!("{name}: prefill failed: {e}")); }); }); } @@ -89,11 +79,15 @@ fn bench_decode_step(c: &mut Criterion) { b.iter_batched_ref( || { let mut engine = kind.clone().build(cpu_engine_backend()); - let _ = engine.prefill(&weights, &ffn, &prompt); + engine + .prefill(&weights, &ffn, &prompt) + .unwrap_or_else(|e| panic!("{name}: prefill failed: {e}")); engine }, |engine| { - let _ = engine.decode_step(&weights, &ffn, 1); + engine + .decode_step(&weights, &ffn, 1) + .unwrap_or_else(|e| panic!("{name}: decode_step failed: {e}")); }, criterion::BatchSize::SmallInput, ); @@ -203,7 +197,8 @@ fn bench_helpers_sync_vs_async(c: &mut Criterion) { None, None, ) - .unwrap(); + .unwrap() + .expect("dispatch produced a result"); }); }); @@ -217,7 +212,8 @@ fn bench_helpers_sync_vs_async(c: &mut Criterion) { None, None, ) - .unwrap(); + .unwrap() + .expect("dispatch produced a result"); }); }); @@ -230,7 +226,8 @@ fn bench_helpers_sync_vs_async(c: &mut Criterion) { None, None, ) - .unwrap(); + .unwrap() + .expect("dispatch produced a result"); let mut pos = prompt.len(); b.iter(|| { let _ = kv_decode_step_via_dispatch( @@ -256,7 +253,8 @@ fn bench_helpers_sync_vs_async(c: &mut Criterion) { None, None, ) - .unwrap(); + .unwrap() + .expect("dispatch produced a result"); let mut pos = prompt.len(); b.iter(|| { let _ = kv_decode_step_via_dispatch_async( diff --git a/crates/larql-kv/docs/decode-cost-model.md b/crates/larql-kv/docs/decode-cost-model.md new file mode 100644 index 000000000..ae17855fd --- /dev/null +++ b/crates/larql-kv/docs/decode-cost-model.md @@ -0,0 +1,285 @@ +# Decode Cost Model — what actually limits each KV engine + +**Status:** 📊 Results v0.1 (2026-08-03). Registry: `larql/kvperf-1`. +**Audience:** anyone choosing an engine, or optimising one. +**Scope:** CPU decode throughput vs context length, all nine `EngineKind` +variants. Says nothing about accuracy — see +[`state-policy.md`](state-policy.md) for what each engine treats as canonical. + +--- + +## 1. The measurement design is the first result + +The engines differ almost entirely in **how K/V is materialised for attention +each decode step** — read from a cache, recomputed from residuals, +decompressed, windowed, or re-forwarded. That is an *O(context)* term. + +A fixed short prompt therefore cannot rank them. At ctx≈7 every engine is +pinned to its fixed per-step cost and they land within noise of each other, +which is exactly what the default `larql bench` prompt produces. Any +engine comparison taken at a short prompt is measuring the intercept and +calling it the engine. + +**Sweep context, or do not compare engines.** + +## 2. Method + +```text +machine Apple M3 Max, 12 P + 4 E, 128 GB, AC power + no thermal / performance / CPU-power warning recorded +backend cpu (BLAS + C Q4 kernel), no concurrent build or benchmark +model qwen3-0.6b-q4k.vindex + 28 layers, hidden 1024, 16 q-heads / 8 kv-heads, head_dim 128 + → kv_dim 1024, GQA reps = 2, vocab 151936 +command larql bench --cpu -n 8 --warmup 2 --prompt + --engine +contexts ≈60, 260, 1050, 2100 tokens +statistics mean and p50 both kept — mean carries amortised periodic work + (buffer doubling, window close, cold-tier encode), p50 the + steady state +``` + +Page cache warmed before measurement. A cold cache on a large model inflates +the *first* engine in a run by an order of magnitude; an earlier 26B run +showed `standard` at 503 ms/token cold and 31.8 ms warm. Warm every model +before believing any number. + +## 3. Results — slope is the engine's real cost + +Linear fit over the four context points, dense model: + +| engine | µs / ctx-token (p50) | R² | intercept | µs / ctx-token (mean) | +|---|---|---|---|---| +| turbo-quant | 7.72 | 1.000 | 5.30 ms | 8.48 | +| markov-rs | 8.08 | 0.998 | 4.34 ms | 8.57 | +| boundary-per-layer | 8.12 | 0.996 | 4.35 ms | 8.74 | +| markov-rs-codec | 8.20 | 0.999 | 4.25 ms | 8.37 | +| standard | 8.23 | 1.000 | 4.18 ms | 8.35 | +| windowed-checkpoint (w=256) | 8.30 | 0.999 | 4.26 ms | 8.78 | +| no-cache | 11581 | 0.996 | — | 11584 | + +**Every cached engine has the same marginal cost per context token** — 7.72 to +8.30 µs, a 7.5% spread across six mechanisms that could hardly be more +different. The representation choice buys nothing at the margin. What the +engines actually differ in is **intercept** (4.18–5.30 ms) and, at specific +context lengths, **variance**. + +`no-cache` is 1400× the slope of the others: 11.58 ms per context token is a +full forward pass per token, which is its documented design (O(N²) overall, +correctness fallback only). Cross-check: prefill at ctx=2100 measured 23875 ms += 11.4 ms/token, and its decode step at that context measured 24144 ms ≈ one +whole prefill. The two agree, so the number is the engine, not the harness. + +## 4. Where the time goes + +Per-stage split of the reference forward, ctx=60 → ctx=1050: + +```text +lm_head 7.49 ms → 7.20 ms constant +CPU fwd 5.26 ms → 16.06 ms +10.9 µs / ctx-token +``` + +lm_head is a large **fixed** cost, and the whole of the context dependence +lives in the forward — i.e. attention over K/V. At short context lm_head is +59% of the step; by ctx=1050 it is 31%. Engines that swap the full vocab +matmul for a KNN lm_head win a flat ~5 ms and nothing else; that win does not +scale with context and is invisible in the slope. + +> **Measurement note (added after this fit).** The intercepts in §3 are +> **forward-only**: at the time of the sweep the engine harness stopped its +> timer before `pick_next`, so lm_head sat outside the measured step while the +> reference row included it. That is why the engine intercepts (4.18-5.30 ms) +> sit close to the reference's *forward* stage (5.26 ms) rather than to its +> full step (~12.8 ms). The slopes are unaffected — lm_head is constant in +> context, so it cancels out of a marginal fit, and §3's conclusion stands +> unchanged. The absolute intercepts and any engine-vs-reference tok/s +> comparison from that era do not: they omit ~5-7 ms of real per-token work. +> The harness now times the whole token and reports the `fwd=` / `head=` +> split per row, so a re-run of this sweep will show intercepts roughly one +> lm_head higher. Re-fitting is not required to trust §3. + +### The K/V read is at ~44% of attainable bandwidth + +```text +K/V bytes per context token = 28 layers × 1024 kv_dim × 2 (K,V) × 4 B (f32) + = 229 376 B = 0.229 MB +standard slope = 8.23 µs + → 27.9 GB/s on logical bytes + → 55.7 GB/s counting the GQA reps=2 re-read +attainable (this machine) ≈ 127 GB/s +``` + +The decode attention kernel reads K/V **head-minor**: + +```rust +// attention/decode/gqa_step.rs +let k_block = k_full.slice(s![.., kv_off..kv_off + head_dim]); +let raw = k_block.dot(&q_row); +``` + +`k_full` is `[L, kv_dim]`, so each head's gemv reads a 512 B window with a +4096 B stride, and because `reps = 2` every KV head's block is streamed twice +— once per q-head sharing it. + +**Hypothesis (untested):** a head-major layout, `[kv_head][L][head_dim]`, +would make each head's read contiguous *and* let the `reps` q-heads share one +pass, roughly halving the O(context) term that dominates every engine. +This is the single highest-leverage change the data points at, because it is +below all six engines rather than inside one. + +### turbo-quant: a reduction is a kernel claim + +turbo-quant stores K/V **8× smaller** (4-bit WHT + Lloyd-Max vs f32) and gets +the *same* slope (7.72 vs 8.23) plus the **worst intercept** (5.30 ms, +27% +over standard). The bandwidth saving is entirely consumed by decompression: +it decompresses the full prior K/V, every layer, every step, and then runs the +identical attention over the identical f32 bytes. Compression that must be +undone before use is not a bandwidth reduction. + +### Variance is transient, not a standing tax + +At ctx=1050 the three residual-canonical engines showed mean 19–24% above p50 +(markov-rs 15.75 vs 13.24; codec 16.46 vs 13.22; boundary-per-layer 16.05 vs +13.56). At ctx=2100 the gap collapsed to 3–6%. So it is a transient at +particular context lengths, not a persistent cost — the doubling-capacity +buffers in `helpers::append_row` are the obvious suspect, but this is **not +established** and the sample (8 steps) is small. + +## 5. Defect found and fixed: `windowed-checkpoint` did not window its attention + +**Fixed 2026-08-03** — the diagnosis is kept in full because the reasoning is +reusable and because the fix had to invert an earlier one. + +`windowed-checkpoint:window=N` reported `window=N` from `info()` and paid +full-context attention cost. + +Evidence: + +1. Its slope matches `standard` across all four context points (8.30 vs 8.23). +2. `window=32`, `256` and `4096` at ctx=1050 are indistinguishable in decode + cost; a 32-row window should cost ~1/33 of the attention. +3. `LARQL_W10_DISABLE=1` restores the engine-side shadow (`hot` 0.0 MB → 5.5 + MB, i.e. genuinely 32 rows) and the cost does **not** change (13.42 → + 13.13 ms). +4. `decode_step_via_dispatch` calls `coarse_decode_step_with_state_masked`, + whose trait signature has **no window parameter**, against + `self.kv_handle` — a backend cache spanning the whole stream — and + `windowed_checkpoint/dispatch.rs` never clips that handle. `window_size` is + used only to segment the prompt into archived windows and to size the + engine-side shadow. +5. `StandardEngine::prefill_quant` documents this exact hazard and guards + against it: + + > the coarse trait surface has no window parameter, so the coarse path + > always attends over the FULL context. A windowed engine must therefore + > decline coarse and take the per-layer path … otherwise the same CLI flag + > gives windowed behaviour on one backend and full-context on another + > while `info()` reports `window=N`. + +`StandardEngine` declines coarse when windowed. `WindowedCheckpointEngine`, +whose entire identity *is* the window, does not. + +> **Superseded 2026-08-04.** The quoted guard no longer exists. Declining the +> fused path was always the expensive half of this fix — on a host-delegating +> backend the per-layer route runs the whole forward on the CPU, measured at +> ~9x on Gemma 3 4B — so the trait grew the window instead: +> `coarse_prefill_windowed` / `coarse_decode_step_windowed`, which fail closed +> when a backend cannot bound both attention and K/V. `StandardEngine` now asks +> rather than declines, and both `CpuBackend` and `MetalBackend` answer. The +> correctness reasoning below stands unchanged — an engine must not report a +> window it does not enforce; only the remedy moved from "refuse the fast path" +> to "make the fast path honour it". + +This was a correctness finding before it was a performance one: on the dispatch +path the engine was not the engine it reported being, its boundary checkpoints +and archive were maintained but unused for attention, and any accuracy measured +on that path was not the windowed engine's accuracy. + +The obvious fix was `StandardEngine`'s — decline the coarse path when windowed +— but that would have cost this engine the fast path entirely, since it is +*always* windowed. Clipping the handle keeps both. + +### The fix + +Two parts, at the layer each belongs to. + +`CpuBackend::clip_kv` only understood the per-layer `CpuKvHandle`, not the +coarse `CpuQ4kCacheHandle` the dispatch path allocates — so a windowed engine +on that path could not clip even if it tried (it panicked as a "foreign +handle"). It now handles both shapes. + +`WindowedCheckpointEngine` now clips the backend handle to the window: after +prefill, down to the open window plus the boundary row; after each auto-close, +down to the boundary row alone. `clip_kv` keeps the *tail*, which is exactly +the row the checkpoint was taken from — so the dispatch path now reproduces +what `extend_current` does on the per-layer path when it seeds a fresh window +from `checkpoints.load`. + +One coupling had to move with it. `close_window` under HOnly read the boundary +row back from the handle by **absolute** stream position — correct only while +the handle spanned the stream, and itself a fix for an older bug where a +window-relative read re-checkpointed the *first* window's row every time. With +the handle clipped, absolute runs off the end and the boundary row is simply +the tail. The two tests that pinned the absolute indexing now pin the claim +both mechanisms share: distinct windows checkpoint distinct rows, each recorded +at its own absolute position. + +### Verification + +The write-up's own prediction was that a correctly windowed engine shows a +near-zero slope, and that this is both the fix and its test. Measured at +`window=64`: + +| engine | ctx 60 | ctx 260 | ctx 1050 | slope | +|---|---|---|---|---| +| windowed-checkpoint (fixed) | 4.63 ms | 5.21 ms | **5.01 ms** | **~0.4 µs** | +| standard | 4.53 ms | 6.29 ms | 12.97 ms | 8.5 µs | + +Flat, and **2.6× faster than `standard`** at ctx=1050 — the saving the window +was supposed to deliver all along. Pinned by two regression tests in +`windowed_checkpoint/dispatch.rs` that assert the backend cache never exceeds +`window_size + 1` rows, through prefill and across repeated auto-closes. + +### The siblings are not defective + +`markov-rs:window=64` (slope 9.0 µs) and `boundary-per-layer:window=64` (24 µs) +both track or exceed `standard`, but that is their contract, not a bug: they +retain a **cold tier** and attend over the full history, so their window bounds +hot-tier memory rather than attention. They are exact-under-contract engines. +`windowed-checkpoint` was the only one claiming to be lossy beyond its window, +which is why it was the only one whose attention had to be bounded. + +Worth noting separately: windowed `boundary-per-layer` costs **2.2× `standard`** +at ctx=1050 (28.4 ms vs 13.0) and 35% more prefill, because its per-layer +encoded cold tier is decoded every step. That is a cost, not a defect, but it +is large enough that the engine should not be reached for casually. + +## 6. What this says about engine choice + +```text +short context (< ~200 tokens) every cached engine is within noise; + choose on memory and accuracy, not speed +long context every cached engine costs the same per token; + choose on memory and accuracy, not speed +never no-cache, outside correctness debugging +``` + +The uncomfortable summary is that on CPU **none of the K/V representations +currently buys decode throughput**. They buy memory (turbo-quant 8×, +markov-rs's residual store, windowed-checkpoint's checkpoints) and they differ +in accuracy contract. The shared bottleneck is one layer below all of them, in +how the attention kernel streams K/V — which is where optimisation effort +should go. + +## 7. Gaps in the instrument + +- `EngineProfiler` covers four of nine engines (markov-rs, markov-rs-codec, + turbo-quant, windowed-checkpoint). `standard`, `no-cache`, `boundary-kv`, + `boundary-per-layer` and `apollo` have no per-stage split, so their costs + are inferred from slope rather than attributed. `bench --profile` says + "markov-rs only for now" and in practice prints the *reference* forward's + split, not the engine's. +- Single machine, single dense model, CPU only. The 26B MoE spot-check showed + the same *ordering* but was taken at ctx≈7 and so measures intercepts only. +- 8 decode steps per point. Enough for p50 on a linear fit (R² ≥ 0.996), not + enough to characterise the tail. diff --git a/crates/larql-kv/docs/state-policy.md b/crates/larql-kv/docs/state-policy.md index 33b7e29b3..e3e4db251 100644 --- a/crates/larql-kv/docs/state-policy.md +++ b/crates/larql-kv/docs/state-policy.md @@ -47,7 +47,7 @@ Discarding it loses the conversation. The known kinds: | Tokens (raw input ids) | `NoCacheEngine` | | Residual streams | `MarkovResidualEngine` | | Boundary residuals | `Apollo`, `BoundaryKvEngine` checkpoint frames | -| KV tensors | `StandardEngine`, `UnlimitedContextEngine` (within window) | +| KV tensors | `StandardEngine`, `WindowedCheckpointEngine` (within window) | | Compressed residual packets | `MarkovResidualCodecEngine` (cold tier), `BoundaryPerLayerEngine` | This list is *open*. New canonical kinds may appear (e.g. a @@ -157,7 +157,7 @@ hot path; the engine simply doesn't shadow it. |---|---|---|---:| | `MarkovResidualEngine` | residual stream | `hot_kv`; (`rs.stored` too when `window=None`) | 106.8 (None) | | `MarkovResidualCodecEngine` | codec residuals | same | 98.5 (None) | -| `UnlimitedContextEngine` | KV within window | `current_window_kv` (CPU shadow of the Metal cache) | 92.8 (HOnly) | +| `WindowedCheckpointEngine` | KV within window | `current_window_kv` (CPU shadow of the Metal cache) | 92.8 (HOnly) | | `TurboQuantEngine` | compressed K/V (destructive) | nothing — K/V IS canonical | — | | `StandardEngine` | KV tensors | n/a — backend-managed already | (reference, ~100) | @@ -213,7 +213,7 @@ Each accessor's purpose: contract is conditional on architecture, a static fact. - **`memory_accounting`** — `hot_bytes()` + `cold_bytes()` split, attributed to canonical vs derivative. Required to surface - things like the `UnlimitedContextEngine` window-shadow + things like the `WindowedCheckpointEngine` window-shadow double-count (engine carries 15.7 MB shadow at window=256 while the backend keeps the full K/V — both should appear). - **`execution_requirements`** — what does the engine *need* from @@ -235,7 +235,7 @@ The engines in `larql-kv` today, classified under the triple: | `MarkovResidualCodecEngine` | codec-encoded residuals | hot KV | `bounded_KL(ε)` — ε stated per codec | | `BoundaryKvEngine` | KV tensors + chunk frames | — | `exact_logits` | | `BoundaryPerLayerEngine` | per-layer codec policy over residuals | hot KV | `bounded_KL(ε_l)` per-layer; calibrated | -| `UnlimitedContextEngine` | KV tensors (within window) + per-window checkpoints + token archive | — | `exact_logits` within window | +| `WindowedCheckpointEngine` | KV tensors (within window) + per-window checkpoints + token archive | — | `exact_logits` within window | | `TurboQuantEngine` | quantised KV (in-place) | — | `codec_bounded_state` — per-row round-trip cos ≈ 0.9954 at 4-bit (Gaussian simulation, 2026-07-30); output KL observed, not bounded | | `Apollo` | boundary retrieval / residual injection store | — | `task_level_retrieval` | @@ -361,3 +361,67 @@ test catches it. the engines for grid deployment without changing their contracts. [`SlabRole`]: ../../larql-compute/src/state_handle.rs + +--- + +## Refusal, and what a failed decode leaves behind (2026-08-02) + +A `StatePolicy` describes what an engine's state *is*. This section describes +what it is after a step that did not finish — which turned out to be a separate +question, and one the `Option` era could not even ask. + +### Three outcomes, never two + +The dispatch helpers return `DispatchOutcome = Result, BoxRefusal>`: + +```text +Ok(Some(_)) the dispatch produced a complete result +Ok(None) nothing to do, or the backend declined this shape +Err(refusal) a routed operation was required and did not execute +``` + +`Ok(None)` means exactly what the old bare `None` meant, so a declining backend +still becomes `EngineError::BackendFailure`. `Err` is new: it says the layer is +*incomplete*, so a strict route can refuse the token instead of returning the +dense half of a layer whose experts never ran. + +### A decode step is transactional + +Attention appends the new token's K/V before the FFN gets the chance to refuse, +so a step that fails has already mutated the cache. `StandardEngine` therefore +snapshots per-layer lengths and rewinds on **any** failure — refusal or +declining backend, since both leave the same half-applied step. + +The rewind primitive is `KvDispatch::truncate_kv`, the inverse of an append. +It is not `clip_kv`: that one keeps the *tail* to enforce a sliding window, +this keeps the *head* to undo one. Its default returns `false` rather than +panicking, because "this backend cannot rewind" is a state to handle. + +Windowed caches are the case worth understanding. A step that reaches the +window evicts its oldest row to make room, and that row is gone — but the row +*count* is unchanged, so length cannot detect it. `rewind_is_sound` therefore +asks whether every layer had room before the step began, not whether the count +came back. + +Where the rewind cannot be trusted, the engine says so rather than pretending: +`EngineError::StateInvalidated` wraps the original cause, later decode steps +refuse with `InvariantViolation`, and a successful `prefill` clears it — the +cache is replaced outright, so re-prefilling is the documented way back. + +### Two questions a caller must not conflate + +`is_recoverable()` used to answer "could this operation succeed?" while callers +read it as "can I retry?". Those diverge exactly where it hurts: a `Residency` +refusal that invalidated the cache is recoverable in the first sense and +catastrophic in the second — fix the residency, re-drive the same engine, and +the token is appended twice. + +```text +operation_is_recoverable() could this operation ever succeed elsewhere? +engine_state_is_retryable() is this engine instance still usable? +is_recoverable() both — what a sweep may actually act on +``` + +`ScoreOutcome` mirrors the distinction rather than flattening it: a +`BindingDefect` is not a coverage deficit, and a dead engine is not a gap in a +run. diff --git a/crates/larql-kv/examples/boundary_per_layer_parity_gate.rs b/crates/larql-kv/examples/boundary_per_layer_parity_gate.rs index ce35a333e..160194f0b 100644 --- a/crates/larql-kv/examples/boundary_per_layer_parity_gate.rs +++ b/crates/larql-kv/examples/boundary_per_layer_parity_gate.rs @@ -81,7 +81,7 @@ fn parse_args() -> Args { let argv: Vec = std::env::args().collect(); let mut a = Args { vindex: PathBuf::new(), - model: "google/gemma-3-4b-it".into(), + model: String::new(), prompt: "The capital of France is".into(), tokens: 50, cpu: false, @@ -120,9 +120,30 @@ fn parse_args() -> Args { ); std::process::exit(2); } + a.model = resolve_model_id(a.model, &a.vindex); a } +/// The model id to use when `--model` was not given. +/// +/// Asked of the artifact rather than defaulted to a constant. A hardcoded +/// default silently disagrees with whatever vindex the user actually passed, +/// and the disagreement surfaces far away — the wrong tokenizer yields +/// plausible token ids and the first honest error arrives inside a dequant +/// parser complaining about byte counts. +fn resolve_model_id(model: String, vindex: &std::path::Path) -> String { + if !model.is_empty() { + return model; + } + larql_vindex::format::model_id_at(vindex).unwrap_or_else(|| { + eprintln!( + "{} records no model id in index.json — pass --model ", + vindex.display() + ); + std::process::exit(2); + }) +} + fn argmax(logits: &[f32]) -> u32 { logits .iter() diff --git a/crates/larql-kv/examples/contract_classify_cached_ffn.rs b/crates/larql-kv/examples/contract_classify_cached_ffn.rs index f8f6eb50e..76107960c 100644 --- a/crates/larql-kv/examples/contract_classify_cached_ffn.rs +++ b/crates/larql-kv/examples/contract_classify_cached_ffn.rs @@ -92,7 +92,7 @@ fn parse_args() -> Args { let argv: Vec = std::env::args().collect(); let mut a = Args { vindex: PathBuf::new(), - model: "google/gemma-3-4b-it".into(), + model: String::new(), template: "The capital of France is".into(), cached_until: 13, prompts: Vec::new(), @@ -143,9 +143,30 @@ fn parse_args() -> Args { if a.prompts.is_empty() { a.prompts = DEFAULT_PROMPTS.iter().map(|s| (*s).into()).collect(); } + a.model = resolve_model_id(a.model, &a.vindex); a } +/// The model id to use when `--model` was not given. +/// +/// Asked of the artifact rather than defaulted to a constant. A hardcoded +/// default silently disagrees with whatever vindex the user actually passed, +/// and the disagreement surfaces far away — the wrong tokenizer yields +/// plausible token ids and the first honest error arrives inside a dequant +/// parser complaining about byte counts. +fn resolve_model_id(model: String, vindex: &std::path::Path) -> String { + if !model.is_empty() { + return model; + } + larql_vindex::format::model_id_at(vindex).unwrap_or_else(|| { + eprintln!( + "{} records no model id in index.json — pass --model ", + vindex.display() + ); + std::process::exit(2); + }) +} + /// Full CPU walk over all layers — no cache substitution. The /// reference path that `bounded_KL(ε)` / `confidence_gated(τ)` are /// stated against. diff --git a/crates/larql-kv/examples/gemma_prefill_parity.rs b/crates/larql-kv/examples/gemma_prefill_parity.rs new file mode 100644 index 000000000..43c120a2c --- /dev/null +++ b/crates/larql-kv/examples/gemma_prefill_parity.rs @@ -0,0 +1,166 @@ +//! Real-model cross-backend prefill parity for the `standard` engine. +//! +//! The synthetic Gemma-3 fixture shows Metal's **batched** prefill +//! (`coarse_prefill` → `fused_prefill`) diverging 23-43% from both the +//! CPU and Metal's own iterative path, while an SWA-free fixture at +//! identical dims agrees to 1.6e-7 (see +//! `tests/gpu_engine_parity/gemma3_prefill_gap.rs`). Whether that is a +//! production defect or a fixture artefact cannot be answered on a +//! 2-layer synthetic model, so this drives the same comparison on a real +//! Gemma-3 Q4K vindex. +//! +//! It sweeps prompt length because the fixture's divergence is +//! position-dependent — absent at one token, large from two onward. It +//! also reports the argmax token at each length, because identical text +//! is what a `larql run` comparison sees and it can mask a hidden-state +//! difference that never flips a decision. +//! +//! ```sh +//! cargo run -p larql-kv --release --features gpu \ +//! --example gemma_prefill_parity -- [max_len] +//! ``` +//! +//! Exits 0 with a printed table; it is a diagnostic, not a gate. + +use larql_inference::ffn::NullFfn; +use larql_kv::EngineKind; +use larql_vindex::SilentLoadCallbacks; +use ndarray::Array2; +use std::path::PathBuf; + +/// Prompt lengths are taken as prefixes of a repeated filler so every +/// length is a strict prefix of the next — the only thing varying is how +/// many positions the prefill relates. +const FILLER: &str = "The capital of France is Paris and the history of Europe is long. "; + +/// Default longest prefix swept. +const DEFAULT_MAX_LEN: usize = 24; + +/// Above this relative L2 the two backends are not computing the same +/// thing; below it they differ by kernel rounding. The synthetic gap +/// sits at ~4e-1, ordinary Q4K reassociation at ~1e-3. +const DIVERGENCE_THRESHOLD: f32 = 1e-2; + +fn relative_l2(a: &Array2, b: &Array2) -> f32 { + let num: f32 = a + .iter() + .zip(b.iter()) + .map(|(x, y)| (x - y) * (x - y)) + .sum::() + .sqrt(); + let den: f32 = b.iter().map(|y| y * y).sum::().sqrt().max(1e-6); + num / den +} + +fn argmax_slice(v: &[f32]) -> usize { + v.iter() + .enumerate() + .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)) + .map(|(i, _)| i) + .unwrap_or(0) +} + +fn main() -> Result<(), Box> { + let mut args = std::env::args().skip(1); + let Some(path) = args.next() else { + eprintln!("usage: gemma_prefill_parity [max_len]"); + return Ok(()); + }; + let vindex_path = PathBuf::from(path); + let max_len: usize = args + .next() + .and_then(|s| s.parse().ok()) + .unwrap_or(DEFAULT_MAX_LEN); + + if !vindex_path.exists() { + eprintln!("vindex not found at {}; skipping", vindex_path.display()); + return Ok(()); + } + + let mut cb = SilentLoadCallbacks; + let weights = larql_vindex::load_model_weights_kquant(&vindex_path, &mut cb)?; + let mut index = larql_vindex::VectorIndex::load_vindex(&vindex_path, &mut cb)?; + index.load_attn_kquant(&vindex_path)?; + index.load_interleaved_kquant(&vindex_path)?; + let tokenizer = larql_vindex::load_vindex_tokenizer(&vindex_path)?; + + let arch = &*weights.arch; + println!( + "model: {} layers, sliding_window={:?}, layer0_sliding={}", + weights.num_layers, + arch.sliding_window_size(), + arch.is_sliding_window_layer(0), + ); + + let long_prompt = FILLER.repeat(max_len.div_ceil(10) + 2); + let all_tokens = larql_inference::encode_prompt(&tokenizer, arch, long_prompt.as_str()) + .map_err(|e| format!("tokenize: {e}"))?; + + println!(); + println!( + " {:>7} {:>12} {:>10} {:>10}", + "tokens", "rel L2", "gpu tok", "cpu tok" + ); + println!(" {}", "-".repeat(46)); + + let mut worst = 0.0f32; + let mut token_mismatches = 0usize; + for n in 1..=max_len.min(all_tokens.len()) { + let prompt = &all_tokens[..n]; + + let mut gpu = EngineKind::Standard { window_size: None } + .build(larql_inference::default_engine_backend()); + let mut cpu = + EngineKind::Standard { window_size: None }.build(larql_inference::cpu_engine_backend()); + + let h_gpu = gpu.prefill_quant( + &weights, + &NullFfn, + &index, + prompt, + &*larql_inference::default_compute_backend(), + )?; + let h_cpu = cpu.prefill_quant( + &weights, + &NullFfn, + &index, + prompt, + &*larql_inference::cpu_backend(), + )?; + + let rel = relative_l2(&h_gpu, &h_cpu); + worst = worst.max(rel); + + let t_gpu = argmax_slice(&larql_inference::forward::hidden_to_raw_logits( + &weights, &h_gpu, + )); + let t_cpu = argmax_slice(&larql_inference::forward::hidden_to_raw_logits( + &weights, &h_cpu, + )); + if t_gpu != t_cpu { + token_mismatches += 1; + } + + let flag = if rel > DIVERGENCE_THRESHOLD { + " <== DIVERGED" + } else { + "" + }; + println!(" {n:>7} {rel:>12.4e} {t_gpu:>10} {t_cpu:>10}{flag}"); + } + + println!(); + println!("worst relative L2: {worst:.4e}"); + println!("argmax mismatches: {token_mismatches}"); + if worst > DIVERGENCE_THRESHOLD { + println!( + "VERDICT: reproduces on this real model — the synthetic fixture was not the artefact." + ); + } else { + println!( + "VERDICT: does NOT reproduce here; the divergence is specific to the \ + synthetic fixture, not this checkpoint." + ); + } + Ok(()) +} diff --git a/crates/larql-kv/examples/vindex_compare.rs b/crates/larql-kv/examples/vindex_compare.rs index f246f2bda..f5e528382 100644 --- a/crates/larql-kv/examples/vindex_compare.rs +++ b/crates/larql-kv/examples/vindex_compare.rs @@ -37,7 +37,7 @@ fn parse_args() -> Args { reference: PathBuf::new(), candidate: PathBuf::new(), prompts_path: None, - model: "google/gemma-3-4b-it".into(), + model: String::new(), out: None, top_k: 5, max_seq_len: None, @@ -101,9 +101,30 @@ At least one of --prompts or --prompt must be provided." ); std::process::exit(1); } + a.model = resolve_model_id(a.model, &a.reference); a } +/// The model id to use when `--model` was not given. +/// +/// Asked of the artifact rather than defaulted to a constant. A hardcoded +/// default silently disagrees with whatever vindex the user actually passed, +/// and the disagreement surfaces far away — the wrong tokenizer yields +/// plausible token ids and the first honest error arrives inside a dequant +/// parser complaining about byte counts. +fn resolve_model_id(model: String, vindex: &std::path::Path) -> String { + if !model.is_empty() { + return model; + } + larql_vindex::format::model_id_at(vindex).unwrap_or_else(|| { + eprintln!( + "{} records no model id in index.json — pass --model ", + vindex.display() + ); + std::process::exit(2); + }) +} + fn load_prompts(args: &Args) -> Vec { let mut prompts = args.inline_prompts.clone(); if let Some(path) = &args.prompts_path { diff --git a/crates/larql-kv/examples/w10_parity_gate.rs b/crates/larql-kv/examples/w10_parity_gate.rs index 008528886..c70801d2b 100644 --- a/crates/larql-kv/examples/w10_parity_gate.rs +++ b/crates/larql-kv/examples/w10_parity_gate.rs @@ -61,7 +61,7 @@ fn parse_args() -> Args { let argv: Vec = std::env::args().collect(); let mut a = Args { vindex: PathBuf::new(), - model: "google/gemma-3-4b-it".into(), + model: String::new(), prompt: "The capital of France is".into(), tokens: 50, cpu: false, @@ -100,9 +100,30 @@ fn parse_args() -> Args { ); std::process::exit(2); } + a.model = resolve_model_id(a.model, &a.vindex); a } +/// The model id to use when `--model` was not given. +/// +/// Asked of the artifact rather than defaulted to a constant. A hardcoded +/// default silently disagrees with whatever vindex the user actually passed, +/// and the disagreement surfaces far away — the wrong tokenizer yields +/// plausible token ids and the first honest error arrives inside a dequant +/// parser complaining about byte counts. +fn resolve_model_id(model: String, vindex: &std::path::Path) -> String { + if !model.is_empty() { + return model; + } + larql_vindex::format::model_id_at(vindex).unwrap_or_else(|| { + eprintln!( + "{} records no model id in index.json — pass --model ", + vindex.display() + ); + std::process::exit(2); + }) +} + fn argmax(logits: &[f32]) -> u32 { logits .iter() diff --git a/crates/larql-kv/src/accuracy_suite/runner/types.rs b/crates/larql-kv/src/accuracy_suite/runner/types.rs index b1c34bc96..7c2e1e5ba 100644 --- a/crates/larql-kv/src/accuracy_suite/runner/types.rs +++ b/crates/larql-kv/src/accuracy_suite/runner/types.rs @@ -47,6 +47,28 @@ pub enum ScoreOutcome { /// backend or compute kernel returned a runtime failure (corrupt /// weights, OOM, GPU driver error). SkippedBackendFailure, + /// Engine reported [`EngineError::Execution`] with a refusal that + /// more residency or another capable executor could rescue. The + /// row is genuinely unserved and belongs in the `served_rate` + /// deficit — nothing is broken. + SkippedExecutionRefused, + /// Engine reported [`EngineError::Execution`] with a + /// [`RefusalKind::BindingDefect`](larql_execution::RefusalKind::BindingDefect) + /// — the bound artifact violates its own contract. + /// + /// Split from [`Self::SkippedExecutionRefused`] on purpose. Both are + /// unserved rows, but only this one indicts the index, and folding + /// it into a coverage deficit is how a broken artifact gets read as + /// a slice that merely wasn't resident. + FailedBindingDefect, + /// Engine reported [`EngineError::StateInvalidated`] — a failure left + /// K/V that could not be rewound, so the engine instance is no longer + /// describing any token sequence. + /// + /// Every row after this one on the same engine is suspect, which is why + /// it is a `Failed*` rather than a `Skipped*`: a run containing one is + /// not a run with a gap, it is a run that should be re-driven. + FailedStateInvalidated, } impl From<&EngineError> for ScoreOutcome { @@ -57,6 +79,18 @@ impl From<&EngineError> for ScoreOutcome { EngineError::BackendUnavailable => Self::SkippedBackendUnavailable, EngineError::InvariantViolation { .. } => Self::SkippedInvariantViolation, EngineError::BackendFailure { .. } => Self::SkippedBackendFailure, + EngineError::Execution(refusal) => { + if refusal.kind().indicts_the_artifact() { + Self::FailedBindingDefect + } else { + Self::SkippedExecutionRefused + } + } + // Deliberately not delegated to the wrapped cause: what the cause + // *was* stops being the actionable fact once the engine can no + // longer be driven, and reporting the cause's outcome would file + // a dead engine under a coverage deficit. + EngineError::StateInvalidated { .. } => Self::FailedStateInvalidated, } } } @@ -382,6 +416,73 @@ mod tests { assert!(!ScoreOutcome::SkippedBackendUnavailable.is_served()); assert!(!ScoreOutcome::SkippedInvariantViolation.is_served()); assert!(!ScoreOutcome::SkippedBackendFailure.is_served()); + assert!(!ScoreOutcome::SkippedExecutionRefused.is_served()); + assert!(!ScoreOutcome::FailedBindingDefect.is_served()); + assert!(!ScoreOutcome::FailedStateInvalidated.is_served()); + } + + // ── EngineError::Execution → outcome ───────────────────────────────────── + + fn execution_error(kind: larql_execution::RefusalKind) -> EngineError { + EngineError::Execution(Box::new(larql_inference::ffn::RecordedRefusal { + layer: 0, + kind, + message: "route refused".into(), + })) + } + + #[test] + fn a_refusal_maps_to_an_outcome_that_keeps_the_artifact_distinction() { + // The split this enum exists to preserve: a shard that does not hold + // an expert, and an index whose router addresses one that cannot + // exist, are both unserved rows and must not share a status — one is + // a coverage number, the other is a bug report. + use larql_execution::RefusalKind; + assert_eq!( + ScoreOutcome::from(&execution_error(RefusalKind::Residency)), + ScoreOutcome::SkippedExecutionRefused + ); + assert_eq!( + ScoreOutcome::from(&execution_error(RefusalKind::Unsupported)), + ScoreOutcome::SkippedExecutionRefused + ); + assert_eq!( + ScoreOutcome::from(&execution_error(RefusalKind::BindingDefect)), + ScoreOutcome::FailedBindingDefect + ); + } + + #[test] + fn an_invalidated_engine_outranks_whatever_caused_it() { + // The wrapped cause is a Residency refusal, which on its own is an + // ordinary coverage skip. Once the engine cannot be driven, that + // reading is wrong: the run is not missing a row, it is untrustworthy + // from here on. + let err = + execution_error(larql_execution::RefusalKind::Residency).invalidating_engine_state(); + assert_eq!( + ScoreOutcome::from(&err), + ScoreOutcome::FailedStateInvalidated, + "delegating to the cause would file a dead engine as a residency gap" + ); + } + + #[test] + fn every_refusal_kind_maps_to_some_outcome() { + // Sweeps the vocabulary rather than the three cases above, so a new + // `RefusalKind` cannot be added without this file being considered. + for kind in larql_execution::RefusalKind::ALL { + let outcome = ScoreOutcome::from(&execution_error(kind)); + assert!( + !outcome.is_served(), + "{kind}: a refused row is never a served one" + ); + assert_eq!( + outcome == ScoreOutcome::FailedBindingDefect, + kind.indicts_the_artifact(), + "{kind}: the outcome must follow indicts_the_artifact()" + ); + } } #[test] @@ -412,6 +513,18 @@ mod tests { ScoreOutcome::SkippedBackendFailure, r#"{"status":"skipped_backend_failure"}"#, ), + ( + ScoreOutcome::SkippedExecutionRefused, + r#"{"status":"skipped_execution_refused"}"#, + ), + ( + ScoreOutcome::FailedBindingDefect, + r#"{"status":"failed_binding_defect"}"#, + ), + ( + ScoreOutcome::FailedStateInvalidated, + r#"{"status":"failed_state_invalidated"}"#, + ), ]; for (outcome, expected_json) in &cases { let json = serde_json::to_string(outcome).unwrap(); diff --git a/crates/larql-kv/src/engines/apollo/engine.rs b/crates/larql-kv/src/engines/apollo/engine.rs index 91bafa0b3..fc7a679f4 100644 --- a/crates/larql-kv/src/engines/apollo/engine.rs +++ b/crates/larql-kv/src/engines/apollo/engine.rs @@ -387,13 +387,30 @@ impl ApolloEngine { // `FfnBackend` dispatch (forward goes through `forward_from_layer` / // `forward_raw_logits` directly), and its `None` returns map to // `RetrievalMiss` rather than the per-layer backend failures of -// `KvEngine`. Construction sites build it as +// `KvEngine`. +// +// That missing dispatch is why both entry points refuse a hybrid-MoE +// architecture outright (see `APOLLO_HAS_NO_FFN_SEAM`) rather than serving it +// a dense-only forward that would answer for a different model. Construction sites build it as // [`larql_inference::AnyEngine::Retrieval`] so the harness branches // once at the top of the autoregressive loop. +/// Why Apollo cannot dispatch experts, for the refusal it owes a MoE model. +/// +/// Structural, not an omission: `forward_from_layer` / `forward_raw_logits` +/// live in `larql-compute` *below* the `FfnBackend` seam and construct their +/// own `ViewFfn` over the dense weights, so no caller-supplied backend — and +/// therefore no bound expert route — can reach them. Giving Apollo real +/// dispatch means threading an `FfnBackend` through `forward_layer_range`, +/// which is a change to the forward, not to this engine. +const APOLLO_ENGINE_NAME: &str = "apollo"; + +const APOLLO_HAS_NO_FFN_SEAM: &str = + "its forward runs below the FfnBackend seam and builds its own dense FFN"; + impl RetrievalEngine for ApolloEngine { fn name(&self) -> &str { - "apollo" + APOLLO_ENGINE_NAME } fn info(&self) -> EngineInfo { @@ -438,6 +455,7 @@ impl RetrievalEngine for ApolloEngine { if token_ids.is_empty() { return Err(EngineError::EmptyPrompt); } + crate::engines::refuse_if_moe(APOLLO_ENGINE_NAME, APOLLO_HAS_NO_FFN_SEAM, weights)?; let store = self .store .as_ref() @@ -498,6 +516,7 @@ impl RetrievalEngine for ApolloEngine { weights: &ModelWeights, token_id: u32, ) -> Result, EngineError> { + crate::engines::refuse_if_moe(APOLLO_ENGINE_NAME, APOLLO_HAS_NO_FFN_SEAM, weights)?; self.context_tokens.push(token_id); let delta = self.injection_delta diff --git a/crates/larql-kv/src/engines/boundary_kv/engine.rs b/crates/larql-kv/src/engines/boundary_kv/engine.rs index 5e8689188..c3e8ae61d 100644 --- a/crates/larql-kv/src/engines/boundary_kv/engine.rs +++ b/crates/larql-kv/src/engines/boundary_kv/engine.rs @@ -867,7 +867,7 @@ mod tests { // forward to the inner `StandardEngine` (threading `index` → Q4K-direct) and // emit a boundary frame identically to the plain path. The `make_test_q4k_*` // fixtures carry Q4K attn slices, so the CPU dequant fallback runs (the same - // fixtures `resident_identity_tests` and `unlimited_context`'s quant tests + // fixtures `resident_identity_tests` and `windowed_checkpoint`'s quant tests // use). chunk_tokens=2 + a 2-token prefill lands on a boundary, exercising // the frame-emit branch on each. diff --git a/crates/larql-kv/src/engines/boundary_per_layer/engine.rs b/crates/larql-kv/src/engines/boundary_per_layer/engine.rs index 944c5821f..cd672dec3 100644 --- a/crates/larql-kv/src/engines/boundary_per_layer/engine.rs +++ b/crates/larql-kv/src/engines/boundary_per_layer/engine.rs @@ -170,9 +170,6 @@ impl BoundaryPerLayerEngine { token_id, index, ) - .ok_or_else(|| EngineError::BackendFailure { - details: "walk::run_decode returned None".into(), - }) } } @@ -219,10 +216,7 @@ impl KvEngine for BoundaryPerLayerEngine { &self.policy, self.window_size, token_ids, - ) - .ok_or_else(|| EngineError::BackendFailure { - details: "walk::run_prefill returned None".into(), - })?; + )?; self.store = Some(store); Ok(hidden) } @@ -249,7 +243,12 @@ impl KvEngine for BoundaryPerLayerEngine { } fn memory_bytes(&self) -> usize { + // Store + whichever off-engine home holds the K/V on the coarse + // path: the handle (CPU whole-model cache) or the backend itself + // (Metal sentinel handle). Mutually exclusive, so summing is safe. self.store.as_ref().map_or(0, |s| s.memory_bytes()) + + self.kv_handle.as_ref().map_or(0, |h| h.resident_bytes()) + + self.backend.backend_resident_kv_bytes() } fn window_tokens(&self) -> usize { @@ -260,6 +259,17 @@ impl KvEngine for BoundaryPerLayerEngine { self.store.as_ref().map_or(0, |s| s.cold_bytes()) } + fn dispatch_path(&self) -> Option { + use larql_inference::kv_engine::DispatchPath; + // `kv_handle` = the W1-GPU fused path; cleared on fallback to + // the dense walk. `store` marks that a prefill has happened. + match (self.kv_handle.is_some(), self.store.is_some()) { + (true, _) => Some(DispatchPath::Coarse), + (false, true) => Some(DispatchPath::PerLayer), + (false, false) => None, + } + } + // ── Q4K path ───────────────────────────────────────────────────────── // // Try W1-GPU dispatch first; fall back to dense walk with attn @@ -961,7 +971,10 @@ mod tests { BoundaryPerLayerEngine::new(None, policy, weights.num_layers, &store).unwrap(); let ffn = larql_inference::ffn::NullFfn; let err = engine.prefill(&weights, &ffn, &[]).unwrap_err(); - assert_eq!(err, larql_inference::kv_engine::EngineError::EmptyPrompt); + assert!(matches!( + err, + larql_inference::kv_engine::EngineError::EmptyPrompt + )); } #[test] @@ -978,7 +991,10 @@ mod tests { let err = engine .prefill_quant(&weights, &ffn, &index, &[], &*backend) .unwrap_err(); - assert_eq!(err, larql_inference::kv_engine::EngineError::EmptyPrompt); + assert!(matches!( + err, + larql_inference::kv_engine::EngineError::EmptyPrompt + )); } #[test] @@ -994,7 +1010,10 @@ mod tests { let err = engine .prefill_via_executor(&weights, &exec, &ffn, &[]) .unwrap_err(); - assert_eq!(err, larql_inference::kv_engine::EngineError::EmptyPrompt); + assert!(matches!( + err, + larql_inference::kv_engine::EngineError::EmptyPrompt + )); } #[test] diff --git a/crates/larql-kv/src/engines/boundary_per_layer/executor.rs b/crates/larql-kv/src/engines/boundary_per_layer/executor.rs index 81d78d1fe..3cdcf6086 100644 --- a/crates/larql-kv/src/engines/boundary_per_layer/executor.rs +++ b/crates/larql-kv/src/engines/boundary_per_layer/executor.rs @@ -342,7 +342,7 @@ mod tests { /// loops — the E2B fixture (non-zero PLE tensors, layer_scalar 0.75) /// diverges bit-visibly if the tail is dropped. Representative for all /// `LayerExecutor`-driven engine loops (markov_residual{,_codec}, - /// turbo_quant, unlimited_context share the same pattern). + /// turbo_quant, windowed_checkpoint share the same pattern). #[cfg(not(windows))] #[test] fn executor_prefill_and_decode_match_legacy_on_ple_arch() { diff --git a/crates/larql-kv/src/engines/boundary_per_layer/walk.rs b/crates/larql-kv/src/engines/boundary_per_layer/walk.rs index 2e064e468..e453f7222 100644 --- a/crates/larql-kv/src/engines/boundary_per_layer/walk.rs +++ b/crates/larql-kv/src/engines/boundary_per_layer/walk.rs @@ -14,6 +14,7 @@ use larql_inference::attention::{run_attention_with_kv_backend, SharedKV}; use larql_inference::ffn::FfnBackend; use larql_inference::forward::embed_tokens_pub; use larql_inference::forward::ple::precompute_per_layer_inputs; +use larql_inference::kv_engine::EngineError; use ndarray::{s, Array2}; use crate::engines::boundary_per_layer::cold_tier::{ @@ -25,6 +26,9 @@ use crate::engines::markov_residual::recompute_kv; /// Run a full prefill through the dense walk. Returns /// `(last_hidden, new_store)` — caller owns the store. +/// +/// Transactional: the store is built into locals and handed back only on +/// success, so a refusal leaves an engine's existing store untouched. pub(super) fn run_prefill( weights: larql_inference::WeightsView, ffn: &dyn FfnBackend, @@ -32,7 +36,7 @@ pub(super) fn run_prefill( policy: &BoundaryLayerPolicy, window_size: Option, token_ids: &[u32], -) -> Option<(Array2, RsStorePerLayer)> { +) -> Result<(Array2, RsStorePerLayer), EngineError> { let num_layers = weights.num_layers; let seq_len = token_ids.len(); let mut h = embed_tokens_pub(&weights, token_ids); @@ -43,17 +47,21 @@ pub(super) fn run_prefill( for layer in 0..num_layers { stored.push(h.clone()); - let (h_post_attn, _k, _v) = - run_attention_with_kv_backend(weights, &h, layer, be, None).expect("attention failed"); - let h_out = crate::engines::layer_ffn_or_moe( + let (h_post_attn, _k, _v) = run_attention_with_kv_backend(weights, &h, layer, be, None) + .ok_or_else(|| EngineError::BackendFailure { + details: format!( + "attention returned None during boundary-per-layer prefill at layer {layer}" + ), + })?; + h = crate::engines::layer_ffn_or_moe( weights.canonical(), &h_post_attn, layer, ffn, Some(ffn), ple_inputs.get(layer), - ); - h = h_out; + ) + .map_err(EngineError::Execution)?; } let mut rs = RsStorePerLayer { @@ -78,7 +86,9 @@ pub(super) fn run_prefill( let codec = policy.codec_for(layer); let decoded_overflow = roundtrip(overflow, codec); let (k, v) = recompute_kv(weights, &decoded_overflow, layer, 0, backend, None) - .expect("cold K/V pre-computation failed"); + .ok_or_else(|| EngineError::BackendFailure { + details: format!("cold K/V pre-computation returned None at layer {layer}"), + })?; cold_kv.push((k, v)); let mut enc = PerLayerEncodedColdLayer::empty(codec, weights.hidden_size); enc.append(overflow); @@ -89,13 +99,24 @@ pub(super) fn run_prefill( rs.cold_abs_start = 0; } - Some((last_row(&h), rs)) + Ok((last_row(&h), rs)) +} + +/// A backend stage that declined, as a typed engine failure. +/// +/// Named once so every decline in the decode walk reads the same and points +/// at the layer — a bare `None` here used to reach the engine as one +/// undifferentiated "run_decode returned None". +fn declined(stage: &str, layer: usize) -> EngineError { + EngineError::BackendFailure { + details: format!("{stage} returned None during boundary-per-layer decode at layer {layer}"), + } } /// Run one decode step through the dense walk, mutating `rs` in place. /// /// Failure invariant (the reason `rs` is `&mut` and not by-value): on any -/// `None` return, the canonical state — `stored`, the cold tiers, and +/// `Err` return, the canonical state — `stored`, the cold tiers, and /// `next_position` — is untouched, and `rs.hot_kv` is left `None`. The /// hot-K/V cache is a droppable derivative of `stored` (see /// `RsStorePerLayer::hot_kv`), so a transient backend failure costs one @@ -109,7 +130,7 @@ pub(super) fn run_decode( rs: &mut RsStorePerLayer, token_id: u32, index: Option<&larql_vindex::VectorIndex>, -) -> Option> { +) -> Result, EngineError> { let num_layers = weights.num_layers; let abs_position = rs.next_position; let mut h_new = embed_tokens_pub(&weights, &[token_id]); @@ -216,7 +237,8 @@ pub(super) fn run_decode( abs_position, Some(backend), idx_kv, - )?; + ) + .ok_or_else(|| declined("attention", layer))?; *k_buf = new_kv.0; *v_buf = new_kv.1; h @@ -230,7 +252,8 @@ pub(super) fn run_decode( let (k_full, v_full) = if let Some(cold_kv) = &rs.cold_kv { let (k_cold, v_cold) = &cold_kv[layer]; let (k_hot, v_hot) = - recompute_kv(weights, h_hot, layer, hot_abs_start, backend, None)?; + recompute_kv(weights, h_hot, layer, hot_abs_start, backend, None) + .ok_or_else(|| declined("hot K/V recompute", layer))?; let c = k_cold.shape()[0]; let kv_dim = k_cold.shape()[1]; let mut k_combined = Array2::::zeros((c + s_hot, kv_dim)); @@ -261,7 +284,8 @@ pub(super) fn run_decode( } else { (h_hot.clone(), hot_abs_start) }; - recompute_kv(weights, &h_full, layer, full_abs_start, backend, None)? + recompute_kv(weights, &h_full, layer, full_abs_start, backend, None) + .ok_or_else(|| declined("cold K/V recompute", layer))? }; let (h_post_attn, new_kv) = @@ -273,22 +297,23 @@ pub(super) fn run_decode( abs_position, Some(backend), idx_kv, - )?; + ) + .ok_or_else(|| declined("attention", layer))?; if cache_eligible { step_new_kv.push(new_kv); } h_post_attn }; - let h_out = crate::engines::layer_ffn_or_moe( + h_new = crate::engines::layer_ffn_or_moe( weights.canonical(), &h_post_attn, layer, ffn, Some(ffn), ple_inputs.get(layer), - ); - h_new = h_out; + ) + .map_err(EngineError::Execution)?; } // Amortised O(m) per-row append via ndarray::Array2::push_row. @@ -356,7 +381,7 @@ pub(super) fn run_decode( } } - Some(last_row(&h_new)) + Ok(last_row(&h_new)) } #[cfg(test)] diff --git a/crates/larql-kv/src/engines/layer_ffn.rs b/crates/larql-kv/src/engines/layer_ffn.rs new file mode 100644 index 000000000..4190029d0 --- /dev/null +++ b/crates/larql-kv/src/engines/layer_ffn.rs @@ -0,0 +1,266 @@ +//! The per-layer FFN step every engine forward loop shares, and the refusal +//! channel it carries. +//! +//! Two functions, one concept: what happens between attention and the next +//! layer. [`apply_ple_and_layer_scalar`] is the arch-dependent tail; +//! [`layer_ffn_or_moe`] is the dispatch that decides whether the layer runs +//! densely or through a bound expert route, and — since PR #197's execution +//! ring — whether it runs at all. + +use larql_execution::BoxRefusal; +use larql_inference::ffn::FfnBackend; +use larql_inference::ModelWeights; +use ndarray::Array2; + +/// Post-FFN tail of the per-layer sequence: `apply_per_layer_embedding` +/// then `apply_layer_scalar`, in that order — mirroring the legacy +/// `kv_prefill_run` / `kv_decode_step_run` loops in +/// [`crate::generation`], the oracle for every engine forward. Both +/// steps are no-ops on archs without PLE / layer-scalar keys, so +/// threading this through non-PLE paths costs one clone and changes no +/// bits. +pub(crate) fn apply_ple_and_layer_scalar( + weights: &ModelWeights, + h_post_ffn: &Array2, + layer: usize, + ple_input: Option<&Array2>, +) -> Array2 { + let mut h_out = larql_inference::forward::ple::apply_per_layer_embedding( + weights, h_post_ffn, layer, ple_input, + ); + larql_inference::forward::layer::apply_layer_scalar(weights, &mut h_out, layer); + h_out +} + +/// Per-layer FFN dispatch for engine forward loops, MoE-aware. +/// +/// On a hybrid-MoE arch, when a `moe_ffn` hook is supplied (e.g. +/// `RemoteMoeFfn` for `--moe-shards`), call its +/// [`FfnBackend::forward_moe_full_layer`] — it returns the full layer output +/// (dense `h1` + experts `h2` + combine), dispatching experts to whatever +/// route is bound behind it. Otherwise fall back to the engine's own dense +/// FFN (`dense_ffn`) followed by [`apply_ple_and_layer_scalar`], the same +/// per-layer sequence as the legacy `kv_prefill_run` / `kv_decode_step_run` +/// oracle. +/// +/// `ple_input` is this layer's entry from `precompute_per_layer_inputs` +/// (`None` on non-PLE archs, where PLE + layer_scalar are no-ops). +/// +/// Nothing here knows which container the experts came from. The hook is an +/// `FfnBackend`, so a VINDEX2 bank, a VINDEX3 bound route and a remote shard +/// are the same shape from this side — which is the property that lets one +/// refusal channel serve all of them. +/// +/// # The refusal channel +/// +/// `Err` means a bound route declined and this layer has no complete output: +/// +/// ```text +/// Ok(h) the layer ran — via the MoE hook, or densely +/// Err(refusal) a strict route refused; this layer is incomplete +/// ``` +/// +/// "Not applicable" is not a third return: a backend with no MoE hook, or one +/// that ran and declined, yields `Ok(None)` from `forward_moe_full_layer` and +/// falls through to the dense path, which is exactly what it did before. +/// Conflating that with a refusal is what the typed channel exists to prevent. +/// +/// Whether an `Err` can arrive at all is the *policy*'s decision, not the +/// route's: `MoeFfn::best_effort` records the refusal and returns the dense +/// half, so a shard outage still degrades rather than stopping. Only +/// `MoeFfn::strict` converts it into an error here. +/// +/// Before this returned a `Result`, a refusal was printed to stderr and the +/// dense half was returned — a complete-looking hidden state whose experts +/// never ran, which the engine had no way to detect. +pub(crate) fn layer_ffn_or_moe( + weights: &ModelWeights, + h_post_attn: &Array2, + layer: usize, + dense_ffn: &dyn FfnBackend, + moe_ffn: Option<&dyn FfnBackend>, + ple_input: Option<&Array2>, +) -> Result, BoxRefusal> { + if weights.arch.is_hybrid_moe() { + if let Some(mf) = moe_ffn { + if let Some(h_out) = mf.forward_moe_full_layer(layer, h_post_attn)? { + // Returned as-is: `forward_moe_full_layer` is contracted to + // produce the FULL layer output. Every production impl routes + // through `moe_ffn_block_cpu(_with_index)`, which applies PLE + // + layer_scalar internally (`RemoteMoeFfn` / `LocalMoeFfn` + // directly; `LayerShardedRemote` and the ffn-policy router + // delegate to them; the HTTP walk backend requests + // `full_output` from the server). Applying either step again + // here would double-apply. + return Ok(h_out); + } + } + } + let (h_post_ffn, _) = + larql_inference::forward::run_ffn(weights, h_post_attn, layer, dense_ffn, false); + Ok(apply_ple_and_layer_scalar( + weights, + &h_post_ffn, + layer, + ple_input, + )) +} + +#[cfg(test)] +mod tests { + use super::{apply_ple_and_layer_scalar, layer_ffn_or_moe}; + use larql_execution::{BoxRefusal, RefusalKind}; + use larql_inference::ffn::{FfnBackend, RecordedRefusal}; + use larql_inference::test_utils::{make_test_gemma4_moe_weights, make_test_weights}; + use ndarray::Array2; + + /// Rows and a value for the synthetic hidden states below. Any finite + /// input works; these are named so a failure message points at a + /// deliberate fixture rather than a bare literal. + const ROWS: usize = 2; + const SENTINEL: f32 = 7.0; + const REFUSAL_MESSAGE: &str = "expert 7 is not resident"; + + /// `FfnBackend` whose MoE hook returns a sentinel so we can tell the MoE + /// branch from the dense `run_ffn` fallback. + struct SentinelFfn; + + impl FfnBackend for SentinelFfn { + fn forward(&self, _layer: usize, x: &Array2) -> Array2 { + Array2::zeros(x.raw_dim()) + } + fn name(&self) -> &str { + "sentinel" + } + fn forward_moe_full_layer( + &self, + _layer: usize, + h_post_attn: &Array2, + ) -> Result>, BoxRefusal> { + Ok(Some(Array2::from_elem(h_post_attn.raw_dim(), SENTINEL))) + } + } + + /// `FfnBackend` whose MoE hook always refuses, with a chosen kind. + struct RefusingFfn(RefusalKind); + + impl FfnBackend for RefusingFfn { + fn forward(&self, _layer: usize, x: &Array2) -> Array2 { + Array2::zeros(x.raw_dim()) + } + fn name(&self) -> &str { + "refusing" + } + fn forward_moe_full_layer( + &self, + layer: usize, + _h_post_attn: &Array2, + ) -> Result>, BoxRefusal> { + Err(Box::new(RecordedRefusal { + layer, + kind: self.0, + message: REFUSAL_MESSAGE.into(), + })) + } + } + + #[test] + fn uses_moe_hook_on_hybrid_moe_arch() { + let weights = make_test_gemma4_moe_weights(); + assert!(weights.arch.is_hybrid_moe()); + let h = Array2::::zeros((ROWS, weights.hidden_size)); + let out = layer_ffn_or_moe(&weights, &h, 0, &SentinelFfn, Some(&SentinelFfn), None) + .expect("an executing hook must not refuse"); + // Took the MoE hook → sentinel output, not the dense run_ffn path. + assert!( + out.iter().all(|&v| v == SENTINEL), + "expected MoE-hook sentinel output" + ); + } + + #[test] + fn falls_back_to_dense_when_no_hook() { + let weights = make_test_gemma4_moe_weights(); + let h = Array2::::zeros((ROWS, weights.hidden_size)); + // No moe_ffn → dense run_ffn even on a MoE arch (no experts dispatched). + let out = layer_ffn_or_moe(&weights, &h, 0, &SentinelFfn, None, None) + .expect("the dense path cannot refuse"); + assert_eq!(out.shape(), &[ROWS, weights.hidden_size]); + assert!( + out.iter().any(|&v| v != SENTINEL), + "must NOT be the MoE-hook sentinel" + ); + assert!(out.iter().all(|v| v.is_finite())); + } + + /// A refusal leaves as a refusal, with its classification and its words. + /// + /// Every kind is swept because the helper must not acquire an opinion + /// about which ones are worth propagating — that decision belongs to the + /// policy above it and the engine below it. + #[test] + fn a_refusing_hook_propagates_rather_than_falling_back() { + let weights = make_test_gemma4_moe_weights(); + let h = Array2::::zeros((ROWS, weights.hidden_size)); + let mut checked = 0usize; + for kind in RefusalKind::ALL { + let ffn = RefusingFfn(kind); + let err = layer_ffn_or_moe(&weights, &h, 0, &ffn, Some(&ffn), None) + .expect_err("a refusing hook must not yield a layer output"); + assert_eq!(err.kind(), kind, "the classification must survive"); + assert!( + err.to_string().contains(REFUSAL_MESSAGE), + "the route's own words must survive: {err}" + ); + checked += 1; + } + assert_eq!(checked, RefusalKind::ALL.len(), "coverage shrank"); + } + + /// On a dense arch the hook is never consulted, so it cannot refuse. + /// + /// Pinned because the alternative — asking the hook first and treating a + /// refusal as fatal — would make every dense model's forward depend on a + /// route it does not use. + #[test] + fn a_refusing_hook_is_not_consulted_on_a_dense_arch() { + let weights = make_test_weights(); + assert!(!weights.arch.is_hybrid_moe()); + let h = Array2::::zeros((ROWS, weights.hidden_size)); + let ffn = RefusingFfn(RefusalKind::Residency); + let out = layer_ffn_or_moe(&weights, &h, 0, &ffn, Some(&ffn), None) + .expect("a dense arch never reaches the MoE hook"); + assert_eq!(out.shape(), &[ROWS, weights.hidden_size]); + } + + /// The tail is a no-op on an arch with neither PLE nor a layer scalar. + #[test] + fn the_tail_is_identity_without_ple_or_layer_scalar() { + let weights = make_test_weights(); + let h = Array2::::from_elem((ROWS, weights.hidden_size), SENTINEL); + let out = apply_ple_and_layer_scalar(&weights, &h, 0, None); + assert_eq!( + out.iter().map(|v| v.to_bits()).collect::>(), + h.iter().map(|v| v.to_bits()).collect::>(), + "a no-op tail must not perturb the bits it passes through" + ); + } + + /// The stub's unused trait surface, so the fixture itself stays honest. + #[test] + fn sentinel_ffn_trait_surface() { + let s = SentinelFfn; + let x = Array2::::zeros((ROWS, 4)); + assert_eq!(s.name(), "sentinel"); + assert_eq!(s.forward(0, &x).shape(), &[ROWS, 4]); + let (o, obs) = s.forward_observed(0, &x); + assert_eq!(o.shape(), &[ROWS, 4]); + assert!( + obs.is_absent(), + "sentinel stub must not fabricate activations" + ); + let r = RefusingFfn(RefusalKind::Unsupported); + assert_eq!(r.name(), "refusing"); + assert_eq!(r.forward(0, &x).shape(), &[ROWS, 4]); + } +} diff --git a/crates/larql-kv/src/engines/markov_residual/compute.rs b/crates/larql-kv/src/engines/markov_residual/compute.rs index 63e628ec5..01ef8aeec 100644 --- a/crates/larql-kv/src/engines/markov_residual/compute.rs +++ b/crates/larql-kv/src/engines/markov_residual/compute.rs @@ -1,4 +1,8 @@ -//! Core residual-stream compute: prefill, decode step, K/V recomputation. +//! Recomputing K/V from stored pre-layer residuals — the operation the +//! residual-stream engines exist to make cheap, plus the walk-KV selection +//! gates and diagnostics around it. +//! +//! Prefill and the decode step live in [`super::prefill`] and [`super::step`]. use larql_compute::{dot_proj_gpu, ComputeBackend, QuantFormat}; use larql_vindex::VectorIndex; @@ -6,14 +10,8 @@ use ndarray::{s, Array2, ArrayBase, ArrayView1, Data, Ix2}; use std::cell::RefCell; use std::cmp::Ordering; -use super::helpers::append_row; -use super::store::RsStore; -use crate::profiler::EngineProfiler; -use larql_inference::attention::SharedKV; -use larql_inference::attention::{apply_rope_partial_at, run_attention_with_kv_backend}; -use larql_inference::ffn::BackendFfn; -use larql_inference::forward::ple::precompute_per_layer_inputs; -use larql_inference::forward::{add_bias, apply_norm, embed_tokens_pub}; +use larql_inference::attention::apply_rope_partial_at; +use larql_inference::forward::{add_bias, apply_norm}; use larql_inference::residual::{rms_norm_heads, rms_norm_heads_no_weight}; #[derive(Clone, Copy)] @@ -75,464 +73,6 @@ pub(crate) fn clear_markov_env_overrides() { MARKOV_ENV_OVERRIDE.with(|o| o.borrow_mut().clear()); } -pub struct RsPrefillResult { - pub hidden: Array2, - pub store: RsStore, - pub memory_bytes: usize, - pub window_tokens: usize, -} - -pub fn rs_prefill( - weights: larql_inference::WeightsView, - token_ids: &[u32], - max_window: Option, - backend: &dyn ComputeBackend, - moe_ffn: Option<&dyn larql_inference::ffn::FfnBackend>, -) -> RsPrefillResult { - let num_layers = weights.num_layers; - let seq_len = token_ids.len(); - let mut h = embed_tokens_pub(&weights, token_ids); - // Empty on non-PLE archs — `ple_inputs.get(layer)` then yields `None`. - let ple_inputs = precompute_per_layer_inputs(&weights, &h, token_ids); - let mut stored: Vec> = Vec::with_capacity(num_layers); - let be = Some(backend); - - for layer in 0..num_layers { - stored.push(h.clone()); - let (h_post_attn, _k, _v) = run_attention_with_kv_backend(weights, &h, layer, be, None) - .expect("attention failed during MarkovRS prefill"); - let bffn = BackendFfn { - weights: weights.canonical(), - backend, - }; - let h_out = crate::engines::layer_ffn_or_moe( - weights.canonical(), - &h_post_attn, - layer, - &bffn, - moe_ffn, - ple_inputs.get(layer), - ); - h = h_out; - } - - let mut rs = RsStore { - hot_len: stored.first().map_or(0, |s| s.shape()[0]), - stored, - cold_residuals: None, - cold_kv: None, - cold_len: 0, - hot_kv: None, - cold_abs_start: 0, - next_position: seq_len, - max_window, - }; - - let mut cold: Vec> = Vec::with_capacity(num_layers); - for layer in 0..num_layers { - rs.clip_layer(layer, &mut cold); - } - rs.finalise_hot_len_after_clip(); - if cold.first().map_or(0, |c| c.shape()[0]) > 0 { - let cold_kv: Vec = (0..num_layers) - .map(|layer| { - recompute_kv(weights, &cold[layer], layer, 0, backend, None) - .expect("cold K/V pre-computation failed") - }) - .collect(); - // 2026-05-19 audit fix: route through the doubling-capacity - // helper so cold_len is initialised correctly. Subsequent - // decode-step overflows then append in amortised O(1). - rs.append_cold_overflow(cold, Some(cold_kv)); - rs.cold_abs_start = 0; - } - - let window_tokens = rs.window_tokens(); - let memory_bytes = rs.memory_bytes(); - RsPrefillResult { - hidden: last_row(&h), - store: rs, - memory_bytes, - window_tokens, - } -} - -pub fn rs_decode_step( - weights: larql_inference::WeightsView, - new_token_id: u32, - rs: RsStore, - backend: &dyn ComputeBackend, - moe_ffn: Option<&dyn larql_inference::ffn::FfnBackend>, - index: Option<&larql_vindex::VectorIndex>, -) -> Option<(Array2, RsStore)> { - rs_decode_step_inner(weights, new_token_id, rs, backend, None, moe_ffn, index) -} - -#[allow(clippy::too_many_arguments)] -pub(crate) fn rs_decode_step_profiled( - weights: larql_inference::WeightsView, - new_token_id: u32, - rs: RsStore, - backend: &dyn ComputeBackend, - profiler: &mut EngineProfiler, - moe_ffn: Option<&dyn larql_inference::ffn::FfnBackend>, - index: Option<&larql_vindex::VectorIndex>, -) -> Option<(Array2, RsStore)> { - rs_decode_step_inner( - weights, - new_token_id, - rs, - backend, - Some(profiler), - moe_ffn, - index, - ) -} - -#[allow(clippy::too_many_arguments)] -fn rs_decode_step_inner( - weights: larql_inference::WeightsView, - new_token_id: u32, - rs: RsStore, - backend: &dyn ComputeBackend, - mut profiler: Option<&mut EngineProfiler>, - moe_ffn: Option<&dyn larql_inference::ffn::FfnBackend>, - index: Option<&larql_vindex::VectorIndex>, -) -> Option<(Array2, RsStore)> { - use std::time::Instant; - - let num_layers = weights.num_layers; - let abs_position = rs.next_position; - let t_step = if profiler.is_some() { - Some(Instant::now()) - } else { - None - }; - let mut h_new = embed_tokens_pub(&weights, &[new_token_id]); - // PLE inputs are per-token — recompute for this single-token decode - // step, matching the legacy `kv_decode_step_run` recipe exactly. - let ple_inputs = precompute_per_layer_inputs(&weights, &h_new, &[new_token_id]); - let mut new_stored: Vec> = Vec::with_capacity(num_layers); - let mut recompute_cold_us = 0.0f64; - let mut recompute_hot_us = 0.0f64; - let mut attention_us = 0.0f64; - let mut ffn_us = 0.0f64; - - // W2 hot-K/V cache on the resident walk (2026-06-13). When there is no cold - // tier (the common unbounded-window case), `hot_kv` holds the FULL K/V and - // we read it instead of re-deriving every position via `recompute_kv` (a - // per-step O(N) matmul — the engine's bottleneck). The residual `stored` is - // still the canonical, re-derivable state (the engine's point); `hot_kv` is - // a droppable derivative. With a cold tier (windowed/evicted) we fall back - // to the recompute path. `step_new_kv` collects each layer's updated full - // K/V returned by the attention step (it concatenates prior cache + the new - // RoPE'd row), which IS next step's cache — no recompute, no concat here. - // Only for unbounded windows (the default): then `clip_layer` is a no-op, - // so the cache never has to track a window-eviction transition. Windowed - // configs keep the existing recompute path unchanged. - let cache_eligible = - rs.max_window.is_none() && rs.cold_residuals.is_none() && rs.cold_kv.is_none(); - let mut step_new_kv: Vec = Vec::with_capacity(num_layers); - // Move the hot K/V cache out so the cache_eligible steady state (step 2+) - // can append into it IN PLACE — borrowing `hot_kv_store` mutably while - // reading `rs.stored` (a disjoint field) immutably. `had_hot_kv` marks the - // seeded-cache case (step 2+); the first decode step has `hot_kv = None` - // and seeds it from `step_new_kv` below. - let mut hot_kv_store = rs.hot_kv; - let had_hot_kv = hot_kv_store.is_some(); - let idx_kv: Option<&dyn larql_compute::KvIndex> = - index.map(|v| v as &dyn larql_compute::KvIndex); - - for layer in 0..num_layers { - // `stored` is a doubling-capacity buffer (W8.2): the logical row count - // is `hot_len`, not `shape()[0]` (see RsStore docs). - let s_hot = rs.hot_len; - let hot_abs_start = abs_position.saturating_sub(s_hot); - - new_stored.push(h_new.clone()); - - let h_post_attn = if cache_eligible && had_hot_kv { - // STEADY STATE (step 2+): `hot_kv` holds the full prior K/V in a - // doubling-capacity buffer. Append this token's projected+RoPE'd row - // IN PLACE and attend over the `[..s_hot+1]` views — no per-step - // O(ctx) owned concat (the previous `_auto` path rebuilt the whole - // K/V every layer every step, i.e. O(L²) copy over a generation; this - // is O(L), matching `standard`'s in-place handle). The residual - // `stored` stays the canonical re-derivable state; the K/V is a - // droppable derivative. Debug builds assert the cached prior matches - // a fresh recompute (the parity gate) before appending. - let bufs = hot_kv_store.as_mut().expect("had_hot_kv"); - #[cfg(debug_assertions)] - { - // Parity gate for the f32 path: the cached prior K/V must match a - // fresh f32 `recompute_kv`. Only meaningful when attention is NOT - // on the Q4K-direct route — that route's projections differ from - // `recompute_kv` by more than the 1e-2 bound even in f32-activation - // (different kernels/byte sources), so it has its own oracles: the - // compute-level bit-identity test (`run_..._inplace` ≡ the concat - // form) and the engine-level in-place-vs-owned-concat A/B test. - let q4k_on = larql_compute::options::q4k_direct_attn_enabled(); - if !q4k_on { - let (k_buf, v_buf) = &bufs[layer]; - let h_logical = rs.stored[layer].slice(s![..s_hot, ..]).to_owned(); - if let Some((rk, rv)) = - recompute_kv(weights, &h_logical, layer, hot_abs_start, backend, None) - { - let kd = k_buf - .slice(s![..s_hot, ..]) - .iter() - .zip(rk.iter()) - .map(|(a, b)| (a - b).abs()) - .fold(0.0f32, f32::max); - let vd = v_buf - .slice(s![..s_hot, ..]) - .iter() - .zip(rv.iter()) - .map(|(a, b)| (a - b).abs()) - .fold(0.0f32, f32::max); - debug_assert!(kd < 1e-2, "markov hot_kv K cache diverged: {kd}"); - debug_assert!(vd < 1e-2, "markov hot_kv V cache diverged: {vd}"); - } - } - } - let (k_buf, v_buf) = &mut bufs[layer]; - let t_attn = if profiler.is_some() { - Some(Instant::now()) - } else { - None - }; - let inplace = if markov_inplace_kv_enabled() { - larql_inference::attention::run_attention_block_decode_step_auto_inplace( - weights, - &h_new, - layer, - k_buf, - v_buf, - s_hot, - abs_position, - Some(backend), - idx_kv, - ) - } else { - None - }; - let h = match inplace { - Some(h) => h, - None => { - // Q4K-direct disabled (the flags-off parity baseline) or no - // attn bytes for this layer: fall back to the owned concat - // over the buffer's logical view, then replace the buffer with - // the exact-length result. Bit-identical to the legacy borrow - // path; only the non-default flags-off case pays this copy. - let prior: SharedKV = ( - k_buf.slice(s![..s_hot, ..]).to_owned(), - v_buf.slice(s![..s_hot, ..]).to_owned(), - ); - let (h, new_kv) = - larql_inference::attention::run_attention_block_decode_step_auto( - weights, - &h_new, - layer, - Some(&prior), - abs_position, - Some(backend), - idx_kv, - )?; - *k_buf = new_kv.0; - *v_buf = new_kv.1; - h - } - }; - if let Some(t) = t_attn { - attention_us += t.elapsed().as_secs_f64() * 1e6; - } - h - } else { - // FIRST STEP (cache None → seed) or windowed/cold tier: recompute the - // prior K/V, let attention concat the new row, and (when - // cache_eligible) collect the result to seed `hot_kv`. - let h_hot = &rs.stored[layer]; - let kv_arg: SharedKV = if let Some(cold_kv) = &rs.cold_kv { - let (k_cold_buf, v_cold_buf) = &cold_kv[layer]; - // 2026-05-19 audit fix: slice to cold_len, not shape()[0]. - // cold_kv now uses doubling-capacity (see append_cold_overflow). - let c = rs.cold_len; - let k_cold = k_cold_buf.slice(s![..c, ..]); - let v_cold = v_cold_buf.slice(s![..c, ..]); - let t_hot = if profiler.is_some() { - Some(Instant::now()) - } else { - None - }; - let (k_hot, v_hot) = - recompute_kv(weights, h_hot, layer, hot_abs_start, backend, None)?; - if let Some(t) = t_hot { - recompute_hot_us += t.elapsed().as_secs_f64() * 1e6; - } - let kv_dim = k_cold_buf.shape()[1]; - let mut k_combined = Array2::::zeros((c + s_hot, kv_dim)); - k_combined.slice_mut(s![..c, ..]).assign(&k_cold); - k_combined.slice_mut(s![c.., ..]).assign(&k_hot); - let mut v_combined = Array2::::zeros((c + s_hot, kv_dim)); - v_combined.slice_mut(s![..c, ..]).assign(&v_cold); - v_combined.slice_mut(s![c.., ..]).assign(&v_hot); - (k_combined, v_combined) - } else { - let (h_full, full_abs_start) = if let Some(cold) = &rs.cold_residuals { - // 2026-05-19 audit fix: slice to cold_len, not shape()[0]. - let s_cold = rs.cold_len; - if s_cold > 0 { - let h_cold = cold[layer].slice(s![..s_cold, ..]); - let hidden = h_hot.shape()[1]; - let mut combined = Array2::::zeros((s_cold + s_hot, hidden)); - combined.slice_mut(s![..s_cold, ..]).assign(&h_cold); - combined.slice_mut(s![s_cold.., ..]).assign(h_hot); - (combined, rs.cold_abs_start) - } else { - (h_hot.clone(), hot_abs_start) - } - } else { - (h_hot.clone(), hot_abs_start) - }; - let t_cold = if profiler.is_some() { - Some(Instant::now()) - } else { - None - }; - let (k, v) = recompute_kv(weights, &h_full, layer, full_abs_start, backend, None)?; - if let Some(t) = t_cold { - recompute_cold_us += t.elapsed().as_secs_f64() * 1e6; - } - (k, v) - }; - - let t_attn = if profiler.is_some() { - Some(Instant::now()) - } else { - None - }; - let (h_post_attn, new_kv) = - larql_inference::attention::run_attention_block_decode_step_auto( - weights, - &h_new, - layer, - Some(&kv_arg), - abs_position, - Some(backend), - idx_kv, - )?; - if let Some(t) = t_attn { - attention_us += t.elapsed().as_secs_f64() * 1e6; - } - // The attention step already projected the new token's K/V (RoPE'd) — - // free; collect it to seed `hot_kv` for the in-place steady state. - if cache_eligible { - step_new_kv.push(new_kv); - } - h_post_attn - }; - - let t_ffn = if profiler.is_some() { - Some(Instant::now()) - } else { - None - }; - let bffn = BackendFfn { - weights: weights.canonical(), - backend, - }; - let h_out = crate::engines::layer_ffn_or_moe( - weights.canonical(), - &h_post_attn, - layer, - &bffn, - moe_ffn, - ple_inputs.get(layer), - ); - if let Some(t) = t_ffn { - ffn_us += t.elapsed().as_secs_f64() * 1e6; - } - h_new = h_out; - } - - if let (Some(prof), Some(t_step)) = (profiler.as_mut(), t_step) { - prof.recompute_cold.total_us += recompute_cold_us; - prof.recompute_cold.count += 1; - prof.recompute_hot.total_us += recompute_hot_us; - prof.recompute_hot.count += 1; - prof.attention.total_us += attention_us; - prof.attention.count += 1; - prof.ffn.total_us += ffn_us; - prof.ffn.count += 1; - prof.decode_total.record(t_step); - } - - // W8.2: in the cache_eligible path `stored` is a doubling-capacity buffer - // (no window → never clips), so append the new row in place rather than - // allocating + bzeroing a fresh `[s_old+1, hidden]` array every step. That - // rebuild was the resident walk's dominant per-step malloc — `__bzero` + - // `szone_malloc` were ~32% of the driver's serial work, idling the worker - // pool (see helpers::append_row, mirrors the dispatch path). The - // windowed/cold path keeps the rebuild: it clips and is not cache_eligible. - let (updated_stored, new_hot_len) = if cache_eligible { - let mut buf = rs.stored; - for (layer, new_row) in new_stored.iter().enumerate() { - append_row(&mut buf[layer], new_row, rs.hot_len); - } - (buf, rs.hot_len + 1) - } else { - let mut rebuilt: Vec> = Vec::with_capacity(num_layers); - for (stored, new_row) in rs.stored.iter().zip(new_stored.iter()) { - let s_old = stored.shape()[0]; - let hidden_dim = stored.shape()[1]; - let mut combined = Array2::::zeros((s_old + 1, hidden_dim)); - combined.slice_mut(s![..s_old, ..]).assign(stored); - combined.slice_mut(s![s_old.., ..]).assign(new_row); - rebuilt.push(combined); - } - let len = rebuilt.first().map_or(0, |s| s.shape()[0]); - (rebuilt, len) - }; - - let mut updated_rs = RsStore { - hot_len: new_hot_len, - stored: updated_stored, - cold_residuals: rs.cold_residuals, - cold_kv: rs.cold_kv, - cold_len: rs.cold_len, - // Cache the full K/V (returned by attention) for next step when there's - // no cold tier; else None (the cold/windowed path recomputes). The clip - // loop below clips `hot_kv` in lockstep with `stored` when a window is set. - // Step 2+ mutated `hot_kv_store` in place (the in-place fast path); the - // first step seeds it from the freshly-collected `step_new_kv`. - hot_kv: if cache_eligible { - if had_hot_kv { - hot_kv_store - } else { - Some(step_new_kv) - } - } else { - None - }, - cold_abs_start: rs.cold_abs_start, - next_position: abs_position + 1, - max_window: rs.max_window, - }; - - let mut overflow: Vec> = Vec::with_capacity(num_layers); - for layer in 0..num_layers { - updated_rs.clip_layer(layer, &mut overflow); - } - updated_rs.finalise_hot_len_after_clip(); - // 2026-05-19 audit fix: geometric-capacity cold append. - // CPU walk path passes `evicted_kv = None` (cold_kv is rebuilt - // from residuals on the next step), mirroring the prior behaviour - // that invalidated cold_kv. See RsStore::append_cold_overflow. - updated_rs.append_cold_overflow(overflow, None); - - Some((last_row(&h_new), updated_rs)) -} - /// Recompute K/V from stored pre-layer residuals using `backend` for projection matmuls. /// /// `index: Some(idx)` enables the Q4K-native fast path: per-row Q4K matvec @@ -1217,228 +757,6 @@ mod tests { assert!(!layer_in_spec("x-y, 30", 29)); } - // ── rs_prefill ──────────────────────────────────────────────────────────── - - #[test] - fn rs_prefill_returns_correct_shape() { - let weights = make_test_weights(); - let result = rs_prefill( - larql_inference::WeightsView::dense(&weights), - &[0u32, 1, 2], - None, - &CpuBackend, - None, - ); - assert_eq!(result.hidden.shape(), &[1, weights.hidden_size]); - assert!(result.hidden.iter().all(|v| v.is_finite())); - } - - #[test] - fn rs_prefill_stores_all_layers() { - let weights = make_test_weights(); - let result = rs_prefill( - larql_inference::WeightsView::dense(&weights), - &[0u32], - None, - &CpuBackend, - None, - ); - assert_eq!(result.store.stored.len(), weights.num_layers); - assert_eq!(result.store.next_position, 1); - } - - #[test] - fn rs_prefill_with_window_clips_hot_store() { - let weights = make_test_weights(); - let result = rs_prefill( - larql_inference::WeightsView::dense(&weights), - &[0u32, 1, 2, 3, 4], - Some(2), - &CpuBackend, - None, - ); - assert!( - result.window_tokens <= 2, - "window_tokens={} > 2", - result.window_tokens - ); - } - - // ── rs_decode_step ──────────────────────────────────────────────────────── - - #[test] - fn rs_decode_step_produces_finite_hidden() { - let weights = make_test_weights(); - let prefill = rs_prefill( - larql_inference::WeightsView::dense(&weights), - &[0u32], - None, - &CpuBackend, - None, - ); - let (h, _) = rs_decode_step( - larql_inference::WeightsView::dense(&weights), - 1, - prefill.store, - &CpuBackend, - None, - None, - ) - .expect("decode step"); - assert_eq!(h.shape(), &[1, weights.hidden_size]); - assert!(h.iter().all(|v| v.is_finite())); - } - - #[test] - fn rs_decode_step_advances_position() { - let weights = make_test_weights(); - let prefill = rs_prefill( - larql_inference::WeightsView::dense(&weights), - &[0u32, 1], - None, - &CpuBackend, - None, - ); - assert_eq!(prefill.store.next_position, 2); - let (_, rs2) = rs_decode_step( - larql_inference::WeightsView::dense(&weights), - 2, - prefill.store, - &CpuBackend, - None, - None, - ) - .unwrap(); - assert_eq!(rs2.next_position, 3); - let (_, rs3) = rs_decode_step( - larql_inference::WeightsView::dense(&weights), - 3, - rs2, - &CpuBackend, - None, - None, - ) - .unwrap(); - assert_eq!(rs3.next_position, 4); - } - - #[test] - fn rs_decode_step_with_cold_kv_branch_produces_finite_output() { - // Windowed prefill with prompt longer than window forces cold_kv - // population (compute.rs lines 60-68), then decode hits the - // `Some(cold_kv)` branch (lines 128-147) instead of the - // cold-residual recomputation path. - let weights = make_test_weights(); - let prefill = rs_prefill( - larql_inference::WeightsView::dense(&weights), - &[0u32, 1, 2, 3], - Some(2), - &CpuBackend, - None, - ); - assert!( - prefill.store.cold_kv.is_some(), - "expected cold_kv to be set" - ); - let (h, rs2) = rs_decode_step( - larql_inference::WeightsView::dense(&weights), - 4, - prefill.store, - &CpuBackend, - None, - None, - ) - .expect("decode_step over cold_kv"); - assert_eq!(h.shape(), &[1, weights.hidden_size]); - assert!(h.iter().all(|v| v.is_finite())); - // After overflow merges into cold_residuals, cold_kv is cleared - // (compute.rs line 260) so a second decode exercises the - // cold_residuals-only branch (lines 149-160). - let (h2, _) = rs_decode_step( - larql_inference::WeightsView::dense(&weights), - 5, - rs2, - &CpuBackend, - None, - None, - ) - .expect("decode_step over cold_residuals"); - assert_eq!(h2.shape(), &[1, weights.hidden_size]); - assert!(h2.iter().all(|v| v.is_finite())); - } - - /// Flags-ON parity gate for the in-place hot-K/V fast path: an A/B of the - /// in-place steady state against the owned-concat reference, both with the - /// Q4K-direct attention path live (int8 OFF so the per-step debug cache - /// assert's 1e-2 bound holds against the q4k `recompute_kv` oracle). The two - /// paths must produce **bit-identical** hidden states at every step — the - /// in-place append only changes the cache *representation* (doubling buffer + - /// views vs fresh owned concat), never the data attended. Runs past a - /// capacity doubling so the grow path is exercised. The `LARQL_MARKOV_INPLACE_KV` - /// override (thread-local; no process-env race) selects the path. - #[test] - fn rs_decode_step_inplace_matches_owned_concat_flags_on() { - use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; - - // Drive the Q4K flags via the thread-local override (no process-env - // mutation → no segfault race with parallel decode tests). Q4K-direct on, - // int8 off (so the debug cache assert's f32 oracle stays valid). - let _q4k = crate::engines::Q4kFlagGuard::set(&[ - (larql_compute::options::ENV_Q4K_DIRECT_ATTN, true), - (larql_compute::options::ENV_Q4K_ATTN_INT8, false), - ]); - - let weights = make_test_q4k_weights(); - let index = make_test_q4k_vindex(&weights); - - // Run a 10-step decode and collect per-step hidden states. - let run = |inplace: bool| -> (Vec>, usize, usize) { - set_markov_env_override( - "LARQL_MARKOV_INPLACE_KV", - Some(if inplace { "1" } else { "0" }), - ); - let prefill = rs_prefill( - larql_inference::WeightsView::dense(&weights), - &[0u32, 1, 2], - None, - &CpuBackend, - None, - ); - let mut rs = prefill.store; - let mut hiddens = Vec::new(); - for tok in 3u32..=12 { - let (h, rs2) = rs_decode_step( - larql_inference::WeightsView::dense(&weights), - tok, - rs, - &CpuBackend, - None, - Some(&index), - ) - .expect("decode"); - assert!(h.iter().all(|v| v.is_finite())); - hiddens.push(h.iter().map(|v| v.to_bits()).collect()); - rs = rs2; - } - let cap = rs.hot_kv.as_ref().expect("hot_kv populated")[0].0.shape()[0]; - (hiddens, rs.hot_len, cap) - }; - - let (a_hiddens, a_len, a_cap) = run(true); - let (b_hiddens, b_len, _b_cap) = run(false); - - assert_eq!(a_len, 13, "3 prompt + 10 decode rows"); - assert_eq!(a_len, b_len, "hot_len must agree across paths"); - assert!( - a_cap >= a_len, - "in-place buffer cap {a_cap} < len {a_len} (no doubling?)" - ); - assert_eq!( - a_hiddens, b_hiddens, - "in-place and owned-concat hidden states diverged (q4k-direct on)" - ); - } - #[test] fn kv_memory_bytes_for_seq_scales_linearly() { let weights = make_test_weights(); @@ -1475,87 +793,6 @@ mod tests { assert!(parse_quant_format("nonsense").is_none()); } - // ── Profiler branches (lines 131, 137, 159, 164, 171, 178, 190, 195) ── - // - // Each timing branch fires only when `profiler.is_some()`. The existing - // `with_profiling_enables_profiling_branch` test exercises one path; - // these add coverage for the cold/hot/attn/ffn timing branches plus the - // overflow-into-existing-cold-residuals merge path. - - #[test] - fn profiled_decode_step_exercises_all_timing_branches() { - use crate::profiler::EngineProfiler; - let weights = make_test_weights(); - let prefill = rs_prefill( - larql_inference::WeightsView::dense(&weights), - &[0u32, 1, 2, 3], - Some(2), - &CpuBackend, - None, - ); - // Has cold_kv populated → exercises lines 130-147 (cold_kv branch - // with profiler timing recompute_hot). - assert!(prefill.store.cold_kv.is_some()); - let mut profiler = EngineProfiler::default(); - let result = rs_decode_step_profiled( - larql_inference::WeightsView::dense(&weights), - 4, - prefill.store, - &CpuBackend, - &mut profiler, - None, - None, - ); - assert!(result.is_some()); - // Profiler must record positive durations across all stages. - assert!(profiler.recompute_hot.count > 0); - assert!(profiler.attention.count > 0); - assert!(profiler.ffn.count > 0); - assert!(profiler.decode_total.count > 0); - } - - #[test] - fn profiled_decode_step_with_cold_residuals_only_path() { - use crate::profiler::EngineProfiler; - let weights = make_test_weights(); - // Two decodes from windowed prefill: first overflows + clears - // cold_kv (compute.rs line 260); second hits the cold_residuals - // branch (lines 149-160) under profiling. - let prefill = rs_prefill( - larql_inference::WeightsView::dense(&weights), - &[0u32, 1, 2, 3], - Some(2), - &CpuBackend, - None, - ); - let (_, rs2) = rs_decode_step( - larql_inference::WeightsView::dense(&weights), - 4, - prefill.store, - &CpuBackend, - None, - None, - ) - .unwrap(); - assert!( - rs2.cold_kv.is_none(), - "cold_kv should be cleared after overflow" - ); - let mut profiler = EngineProfiler::default(); - let result = rs_decode_step_profiled( - larql_inference::WeightsView::dense(&weights), - 5, - rs2, - &CpuBackend, - &mut profiler, - None, - None, - ); - assert!(result.is_some()); - // cold_residuals branch exercises recompute_cold counter (line 171). - assert!(profiler.recompute_cold.count > 0); - } - // ── Pure helpers ──────────────────────────────────────────────────────── #[test] @@ -1826,49 +1063,4 @@ mod tests { assert!(result.is_some()); clear_markov_env_overrides(); } - - #[test] - fn decode_step_with_empty_cold_residuals_falls_through() { - // Line 159: `(h_hot.clone(), hot_abs_start)` when cold tier exists - // but s_cold == 0 (rare; happens if the engine ever clips out the - // last cold row). Build the state by hand. - use larql_inference::attention::SharedKV; - use ndarray::Array2; - let weights = make_test_weights(); - // Construct a store with cold_residuals = Some(vec![empty]) per - // layer and cold_kv = None. The decode loop must take the "empty - // cold" else branch (line 159). - let num_layers = weights.num_layers; - let hidden = weights.hidden_size; - let kv_dim = weights.num_kv_heads * weights.head_dim; - let stored: Vec> = (0..num_layers) - .map(|_| Array2::::zeros((1, hidden))) - .collect(); - let cold_residuals: Vec> = (0..num_layers) - .map(|_| Array2::::zeros((0, hidden))) - .collect(); - let _ = (kv_dim, SharedKV::default()); // silence unused warnings if any - let store = RsStore { - hot_len: 1, - stored, - cold_residuals: Some(cold_residuals), - cold_kv: None, - hot_kv: None, - cold_abs_start: 0, - next_position: 1, - max_window: None, - cold_len: 0, - }; - let result = rs_decode_step( - larql_inference::WeightsView::dense(&weights), - 0, - store, - &CpuBackend, - None, - None, - ); - assert!(result.is_some()); - let (h, _) = result.unwrap(); - assert_eq!(h.shape(), &[1, weights.hidden_size]); - } } diff --git a/crates/larql-kv/src/engines/markov_residual/engine.rs b/crates/larql-kv/src/engines/markov_residual/engine.rs index 73b331940..5a7f4deb0 100644 --- a/crates/larql-kv/src/engines/markov_residual/engine.rs +++ b/crates/larql-kv/src/engines/markov_residual/engine.rs @@ -4,7 +4,8 @@ use larql_compute::ComputeBackend; use larql_vindex::VectorIndex; use ndarray::Array2; -use super::compute::{rs_decode_step, rs_decode_step_profiled, rs_prefill}; +use super::prefill::rs_prefill; +use super::step::{rs_decode_step, rs_decode_step_profiled}; use super::store::RsStore; use super::walk::{ensure_attn_tensors_dequantised, rs_decode_step_walk, rs_prefill_walk}; use crate::profiler::EngineProfiler; @@ -58,8 +59,19 @@ impl MarkovResidualEngine { self } + /// Residual store + any K/V held on this engine's behalf elsewhere. + /// + /// On the W1-GPU path the store stays empty and the K/V lives outside + /// the engine, so the store alone reports 0 — which reads as "costs + /// nothing" rather than "measured somewhere else". It can be in + /// either of two places depending on backend, and they are mutually + /// exclusive: inside the `kv_handle` (CPU's whole-model Q4K cache) or + /// inside the backend itself (Metal, whose handle is a sentinel and + /// whose `backend_resident_kv_bytes` is the only way to see it). pub fn total_memory_bytes(&self) -> usize { self.store.as_ref().map_or(0, |s| s.memory_bytes()) + + self.kv_handle.as_ref().map_or(0, |h| h.resident_bytes()) + + self.backend.backend_resident_kv_bytes() } } @@ -98,6 +110,11 @@ pub(crate) fn check_residual_recompute_preconditions( impl MarkovResidualEngine { /// Shared body for `decode_step` / `decode_step_resident` — `index` /// reaches the attention step's Q4K-direct route when present. + /// The store is borrowed, never taken. A failed step therefore costs the + /// engine nothing but its droppable `hot_kv` derivative — see + /// [`super::step`]'s failure invariant — so a refusal is reported as + /// itself (`EngineError::Execution`, retryable) rather than as a dead + /// engine wearing the words "called before prefill". fn decode_step_impl( &mut self, weights: &ModelWeights, @@ -105,40 +122,38 @@ impl MarkovResidualEngine { token_id: u32, index: Option<&larql_vindex::VectorIndex>, ) -> Result, EngineError> { - let rs = self - .store - .take() + let Self { + store, + backend, + profile, + profiling, + .. + } = self; + let rs = store + .as_mut() .ok_or_else(|| EngineError::InvariantViolation { what: "decode_step called before prefill (store missing)".into(), })?; - let (hidden, new_rs) = if self.profiling { + if *profiling { rs_decode_step_profiled( larql_inference::WeightsView::dense(weights), token_id, rs, - self.backend.as_ref(), - &mut self.profile, + backend.as_ref(), + profile, Some(ffn), index, ) - .ok_or_else(|| EngineError::BackendFailure { - details: "rs_decode_step_profiled returned None".into(), - })? } else { rs_decode_step( larql_inference::WeightsView::dense(weights), token_id, rs, - self.backend.as_ref(), + backend.as_ref(), Some(ffn), index, ) - .ok_or_else(|| EngineError::BackendFailure { - details: "rs_decode_step returned None".into(), - })? - }; - self.store = Some(new_rs); - Ok(hidden) + } } } @@ -174,16 +189,17 @@ impl KvEngine for MarkovResidualEngine { if token_ids.is_empty() { return Err(EngineError::EmptyPrompt); } + // `?` before the assignment: a refused prefill must not replace the + // store an earlier one built. See `prefill`'s transactional contract. let result = rs_prefill( larql_inference::WeightsView::dense(weights), token_ids, self.window_size, self.backend.as_ref(), Some(ffn), - ); - let hidden = result.hidden.clone(); + )?; self.store = Some(result.store); - Ok(hidden) + Ok(result.hidden) } fn decode_step( @@ -221,6 +237,18 @@ impl KvEngine for MarkovResidualEngine { self.store.as_ref().map_or(0, |s| s.cold_bytes()) } + fn dispatch_path(&self) -> Option { + use larql_inference::kv_engine::DispatchPath; + // `kv_handle` is stashed only by the W1-GPU coarse prefill, and + // cleared when that path is abandoned; `store` exists after any + // successful prefill. Neither set = no prefill yet. + match (self.kv_handle.is_some(), self.store.is_some()) { + (true, _) => Some(DispatchPath::Coarse), + (false, true) => Some(DispatchPath::PerLayer), + (false, false) => None, + } + } + fn stage_summary(&self) -> Option { if !self.profiling || self.profile.decode_total.count == 0 { return None; @@ -1482,7 +1510,10 @@ mod tests { let ffn = NullFfn; let mut engine = MarkovResidualEngine::new(None); let err = engine.prefill(&weights, &ffn, &[]).unwrap_err(); - assert_eq!(err, larql_inference::kv_engine::EngineError::EmptyPrompt); + assert!(matches!( + err, + larql_inference::kv_engine::EngineError::EmptyPrompt + )); } #[test] @@ -1513,7 +1544,10 @@ mod tests { let err = engine .prefill_quant(&weights, &ffn, &index, &[], &*backend) .unwrap_err(); - assert_eq!(err, larql_inference::kv_engine::EngineError::EmptyPrompt); + assert!(matches!( + err, + larql_inference::kv_engine::EngineError::EmptyPrompt + )); } #[test] @@ -1545,7 +1579,10 @@ mod tests { let err = engine .prefill_via_executor(&weights, &executor, &ffn, &[]) .unwrap_err(); - assert_eq!(err, larql_inference::kv_engine::EngineError::EmptyPrompt); + assert!(matches!( + err, + larql_inference::kv_engine::EngineError::EmptyPrompt + )); } #[test] @@ -1561,7 +1598,10 @@ mod tests { let err = engine .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[]) .unwrap_err(); - assert_eq!(err, larql_inference::kv_engine::EngineError::EmptyPrompt); + assert!(matches!( + err, + larql_inference::kv_engine::EngineError::EmptyPrompt + )); } #[test] diff --git a/crates/larql-kv/src/engines/markov_residual/mod.rs b/crates/larql-kv/src/engines/markov_residual/mod.rs index b37fbf6fc..3a49010c7 100644 --- a/crates/larql-kv/src/engines/markov_residual/mod.rs +++ b/crates/larql-kv/src/engines/markov_residual/mod.rs @@ -8,12 +8,15 @@ pub mod compute; pub(crate) mod dispatch; pub mod engine; pub(crate) mod helpers; +pub mod prefill; +pub mod step; +mod step_attention; pub mod store; pub mod walk; -pub use compute::{ - kv_memory_bytes_for_seq, recompute_kv, rs_decode_step, rs_prefill, RsPrefillResult, -}; +pub use compute::{kv_memory_bytes_for_seq, recompute_kv}; pub use engine::MarkovResidualEngine; +pub use prefill::{rs_prefill, RsPrefillResult}; +pub use step::rs_decode_step; pub use store::RsStore; pub use walk::ensure_attn_tensors_dequantised; diff --git a/crates/larql-kv/src/engines/markov_residual/prefill.rs b/crates/larql-kv/src/engines/markov_residual/prefill.rs new file mode 100644 index 000000000..e16df388a --- /dev/null +++ b/crates/larql-kv/src/engines/markov_residual/prefill.rs @@ -0,0 +1,161 @@ +//! Residual-stream prefill: build the store the decode step walks. + +use larql_compute::ComputeBackend; +use larql_inference::attention::{run_attention_with_kv_backend, SharedKV}; +use larql_inference::ffn::BackendFfn; +use larql_inference::forward::embed_tokens_pub; +use larql_inference::forward::ple::precompute_per_layer_inputs; +use larql_inference::kv_engine::EngineError; +use ndarray::Array2; + +use super::compute::{last_row, recompute_kv}; +use super::store::RsStore; + +pub struct RsPrefillResult { + pub hidden: Array2, + pub store: RsStore, + pub memory_bytes: usize, + pub window_tokens: usize, +} + +/// Run a full prefill, returning the hidden state and a freshly built store. +/// +/// **Transactional by construction.** Everything is built into locals and the +/// store is handed back only on success, so a failure costs the caller nothing +/// it already had — an engine holding an earlier store keeps it, and a caller +/// that fixes the cause can drive the same prompt again. +pub fn rs_prefill( + weights: larql_inference::WeightsView, + token_ids: &[u32], + max_window: Option, + backend: &dyn ComputeBackend, + moe_ffn: Option<&dyn larql_inference::ffn::FfnBackend>, +) -> Result { + let num_layers = weights.num_layers; + let seq_len = token_ids.len(); + let mut h = embed_tokens_pub(&weights, token_ids); + // Empty on non-PLE archs — `ple_inputs.get(layer)` then yields `None`. + let ple_inputs = precompute_per_layer_inputs(&weights, &h, token_ids); + let mut stored: Vec> = Vec::with_capacity(num_layers); + let be = Some(backend); + + for layer in 0..num_layers { + stored.push(h.clone()); + let (h_post_attn, _k, _v) = run_attention_with_kv_backend(weights, &h, layer, be, None) + .ok_or_else(|| EngineError::BackendFailure { + details: format!( + "attention returned None during MarkovRS prefill at layer {layer}" + ), + })?; + let bffn = BackendFfn { + weights: weights.canonical(), + backend, + }; + h = crate::engines::layer_ffn_or_moe( + weights.canonical(), + &h_post_attn, + layer, + &bffn, + moe_ffn, + ple_inputs.get(layer), + ) + .map_err(EngineError::Execution)?; + } + + let mut rs = RsStore { + hot_len: stored.first().map_or(0, |s| s.shape()[0]), + stored, + cold_residuals: None, + cold_kv: None, + cold_len: 0, + hot_kv: None, + cold_abs_start: 0, + next_position: seq_len, + max_window, + }; + + let mut cold: Vec> = Vec::with_capacity(num_layers); + for layer in 0..num_layers { + rs.clip_layer(layer, &mut cold); + } + rs.finalise_hot_len_after_clip(); + if cold.first().map_or(0, |c| c.shape()[0]) > 0 { + let mut cold_kv: Vec = Vec::with_capacity(num_layers); + for (layer, cold_layer) in cold.iter().enumerate().take(num_layers) { + cold_kv.push( + recompute_kv(weights, cold_layer, layer, 0, backend, None).ok_or_else(|| { + EngineError::BackendFailure { + details: format!("cold K/V pre-computation returned None at layer {layer}"), + } + })?, + ); + } + // 2026-05-19 audit fix: route through the doubling-capacity + // helper so cold_len is initialised correctly. Subsequent + // decode-step overflows then append in amortised O(1). + rs.append_cold_overflow(cold, Some(cold_kv)); + rs.cold_abs_start = 0; + } + + let window_tokens = rs.window_tokens(); + let memory_bytes = rs.memory_bytes(); + Ok(RsPrefillResult { + hidden: last_row(&h), + store: rs, + memory_bytes, + window_tokens, + }) +} + +#[cfg(test)] +mod tests { + use super::rs_prefill; + use larql_compute::CpuBackend; + use larql_inference::test_utils::make_test_weights; + + const PROMPT: [u32; 3] = [0, 1, 2]; + const LONG_PROMPT: [u32; 5] = [0, 1, 2, 3, 4]; + const WINDOW: usize = 2; + + fn prefill( + weights: &larql_inference::ModelWeights, + tokens: &[u32], + window: Option, + ) -> super::RsPrefillResult { + rs_prefill( + larql_inference::WeightsView::dense(weights), + tokens, + window, + &CpuBackend, + None, + ) + .expect("a dense prefill with no MoE hook cannot refuse") + } + + #[test] + fn rs_prefill_returns_correct_shape() { + let weights = make_test_weights(); + let result = prefill(&weights, &PROMPT, None); + assert_eq!(result.hidden.shape(), &[1, weights.hidden_size]); + assert!(result.hidden.iter().all(|v| v.is_finite())); + } + + #[test] + fn rs_prefill_stores_all_layers() { + let weights = make_test_weights(); + let result = prefill(&weights, &PROMPT[..1], None); + assert_eq!(result.store.stored.len(), weights.num_layers); + assert_eq!(result.store.next_position, 1); + } + + #[test] + fn rs_prefill_with_window_clips_hot_store() { + let weights = make_test_weights(); + let result = prefill(&weights, &LONG_PROMPT, Some(WINDOW)); + assert!( + result.window_tokens <= WINDOW, + "window_tokens={} > {WINDOW}", + result.window_tokens + ); + } +} diff --git a/crates/larql-kv/src/engines/markov_residual/step/commit.rs b/crates/larql-kv/src/engines/markov_residual/step/commit.rs new file mode 100644 index 000000000..ae25a0ded --- /dev/null +++ b/crates/larql-kv/src/engines/markov_residual/step/commit.rs @@ -0,0 +1,78 @@ +//! Writing a completed decode step into the residual store. +//! +//! Every canonical mutation the step makes lives here, and nothing here can +//! fail — which is what makes [`super`]'s failure invariant true by +//! construction rather than by inspection. + +use larql_inference::attention::SharedKV; +use ndarray::{s, Array2}; + +use crate::engines::markov_residual::helpers::append_row; +use crate::engines::markov_residual::store::RsStore; + +/// Write the completed step into the store. +/// +/// Everything that mutates canonical state lives here, after the last +/// fallible call — which is what makes this module's failure invariant true +/// by construction rather than by inspection. +pub(super) fn commit( + rs: &mut RsStore, + new_stored: &[Array2], + cache_eligible: bool, + had_hot_kv: bool, + hot_kv_store: Option>, + step_new_kv: Vec, + abs_position: usize, +) { + let num_layers = new_stored.len(); + + // W8.2: in the cache_eligible path `stored` is a doubling-capacity buffer + // (no window → never clips), so append the new row in place rather than + // allocating + bzeroing a fresh `[s_old+1, hidden]` array every step. That + // rebuild was the resident walk's dominant per-step malloc — `__bzero` + + // `szone_malloc` were ~32% of the driver's serial work, idling the worker + // pool (see helpers::append_row, mirrors the dispatch path). The + // windowed/cold path keeps the rebuild: it clips and is not cache_eligible. + if cache_eligible { + let hot_len = rs.hot_len; + for (layer, new_row) in new_stored.iter().enumerate() { + append_row(&mut rs.stored[layer], new_row, hot_len); + } + rs.hot_len = hot_len + 1; + } else { + let mut rebuilt: Vec> = Vec::with_capacity(num_layers); + for (stored, new_row) in rs.stored.iter().zip(new_stored.iter()) { + let s_old = stored.shape()[0]; + let hidden_dim = stored.shape()[1]; + let mut combined = Array2::::zeros((s_old + 1, hidden_dim)); + combined.slice_mut(s![..s_old, ..]).assign(stored); + combined.slice_mut(s![s_old.., ..]).assign(new_row); + rebuilt.push(combined); + } + rs.hot_len = rebuilt.first().map_or(0, |s| s.shape()[0]); + rs.stored = rebuilt; + } + + // Cache the full K/V (returned by attention) for next step when there's no + // cold tier; else None (the cold/windowed path recomputes). The clip loop + // below clips `hot_kv` in lockstep with `stored` when a window is set. + // Step 2+ mutated `hot_kv_store` in place (the in-place fast path); the + // first step seeds it from the freshly-collected `step_new_kv`. + rs.hot_kv = match (cache_eligible, had_hot_kv) { + (true, true) => hot_kv_store, + (true, false) => Some(step_new_kv), + (false, _) => None, + }; + rs.next_position = abs_position + 1; + + let mut overflow: Vec> = Vec::with_capacity(num_layers); + for layer in 0..num_layers { + rs.clip_layer(layer, &mut overflow); + } + rs.finalise_hot_len_after_clip(); + // 2026-05-19 audit fix: geometric-capacity cold append. + // CPU walk path passes `evicted_kv = None` (cold_kv is rebuilt + // from residuals on the next step), mirroring the prior behaviour + // that invalidated cold_kv. See RsStore::append_cold_overflow. + rs.append_cold_overflow(overflow, None); +} diff --git a/crates/larql-kv/src/engines/markov_residual/step/mod.rs b/crates/larql-kv/src/engines/markov_residual/step/mod.rs new file mode 100644 index 000000000..92a2584af --- /dev/null +++ b/crates/larql-kv/src/engines/markov_residual/step/mod.rs @@ -0,0 +1,177 @@ +//! One residual-stream decode step. +//! +//! # Failure invariant +//! +//! The store is borrowed, not consumed, and every canonical field — +//! `stored`, `hot_len`, the cold tiers and `next_position` — is written only +//! in [`commit`], after the last thing that can fail. So an `Err` leaves the +//! store describing exactly the token sequence it described on entry, and the +//! caller may drive the same token through the same engine once the cause is +//! fixed. +//! +//! The one field that does change is `hot_kv`, taken up front and left `None` +//! on the error path. That is deliberate: it is a droppable derivative of +//! `stored` (see [`RsStore::hot_kv`]), a partially-appended in-place buffer +//! must not survive into a retry, and rebuilding it costs one step. This is +//! the same invariant `boundary_per_layer::walk::run_decode` documents — the +//! two engines are residual-canonical, so both can answer "rewind or +//! invalidate" with *rewind*, where a K/V-canonical engine cannot. + +use larql_compute::ComputeBackend; +use larql_inference::attention::SharedKV; +use larql_inference::ffn::BackendFfn; +use larql_inference::forward::embed_tokens_pub; +use larql_inference::forward::ple::precompute_per_layer_inputs; +use larql_inference::kv_engine::EngineError; +use ndarray::Array2; + +use super::compute::last_row; +use super::step_attention::{resolve_layer_attention, HotKv, StepTimings}; +use super::store::RsStore; +use crate::profiler::EngineProfiler; + +mod commit; +#[cfg(test)] +mod tests; + +use commit::commit; + +/// Advance `rs` by one token, returning the new last hidden row. +pub fn rs_decode_step( + weights: larql_inference::WeightsView, + new_token_id: u32, + rs: &mut RsStore, + backend: &dyn ComputeBackend, + moe_ffn: Option<&dyn larql_inference::ffn::FfnBackend>, + index: Option<&larql_vindex::VectorIndex>, +) -> Result, EngineError> { + rs_decode_step_inner(weights, new_token_id, rs, backend, None, moe_ffn, index) +} + +/// [`rs_decode_step`] with per-stage timings recorded into `profiler`. +#[allow(clippy::too_many_arguments)] +pub(crate) fn rs_decode_step_profiled( + weights: larql_inference::WeightsView, + new_token_id: u32, + rs: &mut RsStore, + backend: &dyn ComputeBackend, + profiler: &mut EngineProfiler, + moe_ffn: Option<&dyn larql_inference::ffn::FfnBackend>, + index: Option<&larql_vindex::VectorIndex>, +) -> Result, EngineError> { + rs_decode_step_inner( + weights, + new_token_id, + rs, + backend, + Some(profiler), + moe_ffn, + index, + ) +} + +#[allow(clippy::too_many_arguments)] +fn rs_decode_step_inner( + weights: larql_inference::WeightsView, + new_token_id: u32, + rs: &mut RsStore, + backend: &dyn ComputeBackend, + mut profiler: Option<&mut EngineProfiler>, + moe_ffn: Option<&dyn larql_inference::ffn::FfnBackend>, + index: Option<&larql_vindex::VectorIndex>, +) -> Result, EngineError> { + let num_layers = weights.num_layers; + let abs_position = rs.next_position; + let timed = profiler.is_some(); + let t_step = timed.then(std::time::Instant::now); + let mut timings = StepTimings::default(); + + let mut h_new = embed_tokens_pub(&weights, &[new_token_id]); + // PLE inputs are per-token — recompute for this single-token decode + // step, matching the legacy `kv_decode_step_run` recipe exactly. + let ple_inputs = precompute_per_layer_inputs(&weights, &h_new, &[new_token_id]); + let mut new_stored: Vec> = Vec::with_capacity(num_layers); + + // W2 hot-K/V cache on the resident walk (2026-06-13). When there is no cold + // tier (the common unbounded-window case), `hot_kv` holds the FULL K/V and + // is read instead of re-deriving every position via `recompute_kv` (a + // per-step O(N) matmul — the engine's bottleneck). Only for unbounded + // windows (the default): then `clip_layer` is a no-op, so the cache never + // has to track a window-eviction transition. Windowed configs keep the + // recompute path unchanged. + let cache_eligible = + rs.max_window.is_none() && rs.cold_residuals.is_none() && rs.cold_kv.is_none(); + let mut step_new_kv: Vec = Vec::with_capacity(num_layers); + // Taken up front: see this module's failure invariant. Moving it out also + // lets the steady state append into it mutably while `rs.stored` (a + // disjoint field) is read immutably. + let mut hot_kv_store = rs.hot_kv.take(); + let had_hot_kv = hot_kv_store.is_some(); + let idx_kv: Option<&dyn larql_compute::KvIndex> = + index.map(|v| v as &dyn larql_compute::KvIndex); + + for layer in 0..num_layers { + new_stored.push(h_new.clone()); + + let hot_kv = match (cache_eligible && had_hot_kv, hot_kv_store.as_mut()) { + (true, Some(bufs)) => HotKv::InPlace(bufs), + _ => HotKv::Recompute, + }; + let h_post_attn = resolve_layer_attention( + weights, + rs, + layer, + &h_new, + abs_position, + backend, + idx_kv, + hot_kv, + &mut step_new_kv, + cache_eligible, + timed, + &mut timings, + ) + .ok_or_else(|| EngineError::BackendFailure { + details: format!("attention returned None during MarkovRS decode at layer {layer}"), + })?; + + let bffn = BackendFfn { + weights: weights.canonical(), + backend, + }; + h_new = StepTimings::measure(timed, &mut timings.ffn_us, || { + crate::engines::layer_ffn_or_moe( + weights.canonical(), + &h_post_attn, + layer, + &bffn, + moe_ffn, + ple_inputs.get(layer), + ) + }) + .map_err(EngineError::Execution)?; + } + + if let (Some(prof), Some(t_step)) = (profiler.as_mut(), t_step) { + prof.recompute_cold.total_us += timings.recompute_cold_us; + prof.recompute_cold.count += 1; + prof.recompute_hot.total_us += timings.recompute_hot_us; + prof.recompute_hot.count += 1; + prof.attention.total_us += timings.attention_us; + prof.attention.count += 1; + prof.ffn.total_us += timings.ffn_us; + prof.ffn.count += 1; + prof.decode_total.record(t_step); + } + + commit( + rs, + &new_stored, + cache_eligible, + had_hot_kv, + hot_kv_store, + step_new_kv, + abs_position, + ); + Ok(last_row(&h_new)) +} diff --git a/crates/larql-kv/src/engines/markov_residual/step/tests.rs b/crates/larql-kv/src/engines/markov_residual/step/tests.rs new file mode 100644 index 000000000..0243488e0 --- /dev/null +++ b/crates/larql-kv/src/engines/markov_residual/step/tests.rs @@ -0,0 +1,228 @@ +//! Decode-step behaviour: shapes, positions, cold tiers, and the +//! in-place-vs-owned-concat parity gate. + +use super::{rs_decode_step, rs_decode_step_profiled}; +use crate::engines::markov_residual::prefill::rs_prefill; +use crate::engines::markov_residual::store::RsStore; +use crate::profiler::EngineProfiler; +use larql_compute::CpuBackend; +use larql_inference::test_utils::make_test_weights; +use ndarray::Array2; + +const PROMPT: [u32; 2] = [0, 1]; +const OVERFLOW_PROMPT: [u32; 4] = [0, 1, 2, 3]; +const WINDOW: usize = 2; + +fn prefill( + weights: &larql_inference::ModelWeights, + tokens: &[u32], + window: Option, +) -> RsStore { + rs_prefill( + larql_inference::WeightsView::dense(weights), + tokens, + window, + &CpuBackend, + None, + ) + .expect("a dense prefill with no MoE hook cannot refuse") + .store +} + +fn decode(weights: &larql_inference::ModelWeights, rs: &mut RsStore, token: u32) -> Array2 { + rs_decode_step( + larql_inference::WeightsView::dense(weights), + token, + rs, + &CpuBackend, + None, + None, + ) + .expect("decode step") +} + +#[test] +fn rs_decode_step_produces_finite_hidden() { + let weights = make_test_weights(); + let mut rs = prefill(&weights, &PROMPT[..1], None); + let h = decode(&weights, &mut rs, 1); + assert_eq!(h.shape(), &[1, weights.hidden_size]); + assert!(h.iter().all(|v| v.is_finite())); +} + +#[test] +fn rs_decode_step_advances_position() { + let weights = make_test_weights(); + let mut rs = prefill(&weights, &PROMPT, None); + assert_eq!(rs.next_position, PROMPT.len()); + decode(&weights, &mut rs, 2); + assert_eq!(rs.next_position, PROMPT.len() + 1); + decode(&weights, &mut rs, 3); + assert_eq!(rs.next_position, PROMPT.len() + 2); +} + +#[test] +fn rs_decode_step_with_cold_kv_branch_produces_finite_output() { + // Windowed prefill with a prompt longer than the window populates + // cold_kv, so the first decode takes the `Some(cold_kv)` branch rather + // than the cold-residual recomputation path. + let weights = make_test_weights(); + let mut rs = prefill(&weights, &OVERFLOW_PROMPT, Some(WINDOW)); + assert!(rs.cold_kv.is_some(), "expected cold_kv to be set"); + let h = decode(&weights, &mut rs, 4); + assert_eq!(h.shape(), &[1, weights.hidden_size]); + assert!(h.iter().all(|v| v.is_finite())); + // After overflow merges into cold_residuals, cold_kv is cleared, so a + // second decode exercises the cold_residuals-only branch. + let h2 = decode(&weights, &mut rs, 5); + assert_eq!(h2.shape(), &[1, weights.hidden_size]); + assert!(h2.iter().all(|v| v.is_finite())); +} + +/// Flags-ON parity gate for the in-place hot-K/V fast path: an A/B of the +/// in-place steady state against the owned-concat reference, both with the +/// Q4K-direct attention path live (int8 OFF so the per-step debug cache +/// assert's bound holds against the q4k `recompute_kv` oracle). The two +/// paths must produce **bit-identical** hidden states at every step — the +/// in-place append only changes the cache *representation* (doubling +/// buffer + views vs fresh owned concat), never the data attended. Runs +/// past a capacity doubling so the grow path is exercised. The +/// `LARQL_MARKOV_INPLACE_KV` override (thread-local; no process-env race) +/// selects the path. +#[test] +fn rs_decode_step_inplace_matches_owned_concat_flags_on() { + use crate::engines::markov_residual::compute::set_markov_env_override; + use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; + + const PARITY_PROMPT: [u32; 3] = [0, 1, 2]; + const FIRST_DECODE_TOKEN: u32 = 3; + const LAST_DECODE_TOKEN: u32 = 12; + + // Drive the Q4K flags via the thread-local override (no process-env + // mutation → no segfault race with parallel decode tests). Q4K-direct + // on, int8 off (so the debug cache assert's f32 oracle stays valid). + let _q4k = crate::engines::Q4kFlagGuard::set(&[ + (larql_compute::options::ENV_Q4K_DIRECT_ATTN, true), + (larql_compute::options::ENV_Q4K_ATTN_INT8, false), + ]); + + let weights = make_test_q4k_weights(); + let index = make_test_q4k_vindex(&weights); + + let run = |inplace: bool| -> (Vec>, usize, usize) { + set_markov_env_override( + "LARQL_MARKOV_INPLACE_KV", + Some(if inplace { "1" } else { "0" }), + ); + let mut rs = prefill(&weights, &PARITY_PROMPT, None); + let mut hiddens = Vec::new(); + for tok in FIRST_DECODE_TOKEN..=LAST_DECODE_TOKEN { + let h = rs_decode_step( + larql_inference::WeightsView::dense(&weights), + tok, + &mut rs, + &CpuBackend, + None, + Some(&index), + ) + .expect("decode"); + assert!(h.iter().all(|v| v.is_finite())); + hiddens.push(h.iter().map(|v| v.to_bits()).collect()); + } + let cap = rs.hot_kv.as_ref().expect("hot_kv populated")[0].0.shape()[0]; + (hiddens, rs.hot_len, cap) + }; + + let (a_hiddens, a_len, a_cap) = run(true); + let (b_hiddens, b_len, _b_cap) = run(false); + + let decoded = (LAST_DECODE_TOKEN - FIRST_DECODE_TOKEN + 1) as usize; + assert_eq!(a_len, PARITY_PROMPT.len() + decoded, "prompt + decode rows"); + assert_eq!(a_len, b_len, "hot_len must agree across paths"); + assert!( + a_cap >= a_len, + "in-place buffer cap {a_cap} < len {a_len} (no doubling?)" + ); + assert_eq!( + a_hiddens, b_hiddens, + "in-place and owned-concat hidden states diverged (q4k-direct on)" + ); +} + +/// Every timing branch fires when a profiler is attached. +#[test] +fn profiled_decode_step_exercises_all_timing_branches() { + let weights = make_test_weights(); + // Windowed prefill leaves cold_kv populated → the cold_kv branch runs + // with `recompute_hot` timed. + let mut rs = prefill(&weights, &OVERFLOW_PROMPT, Some(WINDOW)); + assert!(rs.cold_kv.is_some()); + let mut profiler = EngineProfiler::default(); + rs_decode_step_profiled( + larql_inference::WeightsView::dense(&weights), + 4, + &mut rs, + &CpuBackend, + &mut profiler, + None, + None, + ) + .expect("profiled decode"); + assert!(profiler.recompute_hot.count > 0); + assert!(profiler.attention.count > 0); + assert!(profiler.ffn.count > 0); + assert!(profiler.decode_total.count > 0); +} + +#[test] +fn profiled_decode_step_with_cold_residuals_only_path() { + let weights = make_test_weights(); + // Two decodes from a windowed prefill: the first overflows and clears + // cold_kv; the second hits the cold_residuals branch under profiling. + let mut rs = prefill(&weights, &OVERFLOW_PROMPT, Some(WINDOW)); + decode(&weights, &mut rs, 4); + assert!( + rs.cold_kv.is_none(), + "cold_kv should be cleared after overflow" + ); + let mut profiler = EngineProfiler::default(); + rs_decode_step_profiled( + larql_inference::WeightsView::dense(&weights), + 5, + &mut rs, + &CpuBackend, + &mut profiler, + None, + None, + ) + .expect("profiled decode"); + assert!(profiler.recompute_cold.count > 0); +} + +/// A cold tier that exists but holds no rows falls through to the hot-only +/// prior, rather than concatenating a zero-row block. +#[test] +fn decode_step_with_empty_cold_residuals_falls_through() { + let weights = make_test_weights(); + let num_layers = weights.num_layers; + let hidden = weights.hidden_size; + let mut store = RsStore { + hot_len: 1, + stored: (0..num_layers) + .map(|_| Array2::::zeros((1, hidden))) + .collect(), + cold_residuals: Some( + (0..num_layers) + .map(|_| Array2::::zeros((0, hidden))) + .collect(), + ), + cold_kv: None, + hot_kv: None, + cold_abs_start: 0, + next_position: 1, + max_window: None, + cold_len: 0, + }; + let h = decode(&weights, &mut store, 0); + assert_eq!(h.shape(), &[1, weights.hidden_size]); +} diff --git a/crates/larql-kv/src/engines/markov_residual/step_attention.rs b/crates/larql-kv/src/engines/markov_residual/step_attention.rs new file mode 100644 index 000000000..9ba337cf3 --- /dev/null +++ b/crates/larql-kv/src/engines/markov_residual/step_attention.rs @@ -0,0 +1,297 @@ +//! Resolving one layer's attention output from the residual store's tiers. +//! +//! The decode step's per-layer work splits in two: derive `h_post_attn` from +//! whatever K/V the store can offer (this file), then run the FFN over it +//! (see [`super::step`]). Only the first half has to know about the store's +//! three shapes — a warm in-place hot cache, a cold tier that must be +//! recombined, and the first step, which has neither. + +use larql_compute::ComputeBackend; +use larql_inference::attention::SharedKV; +use ndarray::{s, Array2}; + +use super::compute::recompute_kv; +use super::store::RsStore; + +/// Per-stage microsecond accumulators for one decode step. +/// +/// Grouped rather than passed as four `&mut f64` so a new stage does not +/// widen every signature between here and the profiler. +#[derive(Default)] +pub(super) struct StepTimings { + pub(super) recompute_cold_us: f64, + pub(super) recompute_hot_us: f64, + pub(super) attention_us: f64, + pub(super) ffn_us: f64, +} + +impl StepTimings { + /// Time `f` into `slot` when `enabled`, else just run it. + /// + /// The timing branch is written once here because the alternative — an + /// `Instant::now()` pair around every stage — is what made the original + /// loop unreadable, and a mistimed stage silently misattributes the + /// engine's cost. + pub(super) fn measure(enabled: bool, slot: &mut f64, f: impl FnOnce() -> T) -> T { + /// The profiler reports microseconds; `elapsed` gives seconds. + const MICROS_PER_SECOND: f64 = 1e6; + + if !enabled { + return f(); + } + let start = std::time::Instant::now(); + let out = f(); + *slot += start.elapsed().as_secs_f64() * MICROS_PER_SECOND; + out + } +} + +/// The hot-K/V cache's state for this step, which is what decides how +/// attention gets its prior. +pub(super) enum HotKv<'a> { + /// Steady state (step 2+) on an unbounded window: the buffers hold the + /// full prior K/V and this layer's row is appended into them in place. + InPlace(&'a mut Vec), + /// First step, or a windowed/cold configuration: the prior is recomputed + /// from the canonical residuals. + Recompute, +} + +/// Derive this layer's `h_post_attn`. +/// +/// Returns `None` when the backend declines — the caller turns that into a +/// typed failure. Every mutation this makes is to `hot_kv`, a droppable +/// derivative of `rs.stored`, and never to the store's canonical state. +#[allow(clippy::too_many_arguments)] +pub(super) fn resolve_layer_attention( + weights: larql_inference::WeightsView, + rs: &RsStore, + layer: usize, + h_new: &Array2, + abs_position: usize, + backend: &dyn ComputeBackend, + idx_kv: Option<&dyn larql_compute::KvIndex>, + hot_kv: HotKv<'_>, + step_new_kv: &mut Vec, + cache_eligible: bool, + timed: bool, + timings: &mut StepTimings, +) -> Option> { + // `stored` is a doubling-capacity buffer (W8.2): the logical row count is + // `hot_len`, not `shape()[0]` (see RsStore docs). + let s_hot = rs.hot_len; + let hot_abs_start = abs_position.saturating_sub(s_hot); + + match hot_kv { + HotKv::InPlace(bufs) => { + #[cfg(debug_assertions)] + debug_assert_hot_kv_parity(weights, rs, layer, bufs, s_hot, hot_abs_start, backend); + attend_in_place( + weights, + &mut bufs[layer], + layer, + h_new, + s_hot, + abs_position, + backend, + idx_kv, + timed, + timings, + ) + } + HotKv::Recompute => { + let kv_arg = recompute_prior_kv( + weights, + rs, + layer, + s_hot, + hot_abs_start, + backend, + timed, + timings, + )?; + let (h_post_attn, new_kv) = + StepTimings::measure(timed, &mut timings.attention_us, || { + larql_inference::attention::run_attention_block_decode_step_auto( + weights, + h_new, + layer, + Some(&kv_arg), + abs_position, + Some(backend), + idx_kv, + ) + })?; + // The attention step already projected the new token's K/V + // (RoPE'd) — free; collect it to seed `hot_kv` for the in-place + // steady state. + if cache_eligible { + step_new_kv.push(new_kv); + } + Some(h_post_attn) + } + } +} + +/// Steady state: append this token's projected+RoPE'd row into the layer's +/// doubling-capacity buffer and attend over the `[..s_hot+1]` views. +/// +/// No per-step O(ctx) owned concat — the previous `_auto` path rebuilt the +/// whole K/V every layer every step, i.e. O(L²) copy over a generation; this +/// is O(L), matching `standard`'s in-place handle. The residual `stored` stays +/// the canonical re-derivable state; the K/V is a droppable derivative. +#[allow(clippy::too_many_arguments)] +fn attend_in_place( + weights: larql_inference::WeightsView, + buf: &mut SharedKV, + layer: usize, + h_new: &Array2, + s_hot: usize, + abs_position: usize, + backend: &dyn ComputeBackend, + idx_kv: Option<&dyn larql_compute::KvIndex>, + timed: bool, + timings: &mut StepTimings, +) -> Option> { + let (k_buf, v_buf) = buf; + StepTimings::measure(timed, &mut timings.attention_us, || { + let inplace = if super::compute::markov_inplace_kv_enabled() { + larql_inference::attention::run_attention_block_decode_step_auto_inplace( + weights, + h_new, + layer, + k_buf, + v_buf, + s_hot, + abs_position, + Some(backend), + idx_kv, + ) + } else { + None + }; + match inplace { + Some(h) => Some(h), + None => { + // Q4K-direct disabled (the flags-off parity baseline) or no + // attn bytes for this layer: fall back to the owned concat + // over the buffer's logical view, then replace the buffer with + // the exact-length result. Bit-identical to the legacy borrow + // path; only the non-default flags-off case pays this copy. + let prior: SharedKV = ( + k_buf.slice(s![..s_hot, ..]).to_owned(), + v_buf.slice(s![..s_hot, ..]).to_owned(), + ); + let (h, new_kv) = larql_inference::attention::run_attention_block_decode_step_auto( + weights, + h_new, + layer, + Some(&prior), + abs_position, + Some(backend), + idx_kv, + )?; + *k_buf = new_kv.0; + *v_buf = new_kv.1; + Some(h) + } + } + }) +} + +/// First step (cache `None` → seed) or windowed/cold tier: recompute the prior +/// K/V so attention can concat the new row onto it. +#[allow(clippy::too_many_arguments)] +fn recompute_prior_kv( + weights: larql_inference::WeightsView, + rs: &RsStore, + layer: usize, + s_hot: usize, + hot_abs_start: usize, + backend: &dyn ComputeBackend, + timed: bool, + timings: &mut StepTimings, +) -> Option { + let h_hot = &rs.stored[layer]; + if let Some(cold_kv) = &rs.cold_kv { + let (k_cold_buf, v_cold_buf) = &cold_kv[layer]; + // 2026-05-19 audit fix: slice to cold_len, not shape()[0]. + // cold_kv now uses doubling-capacity (see append_cold_overflow). + let c = rs.cold_len; + let k_cold = k_cold_buf.slice(s![..c, ..]); + let v_cold = v_cold_buf.slice(s![..c, ..]); + let (k_hot, v_hot) = StepTimings::measure(timed, &mut timings.recompute_hot_us, || { + recompute_kv(weights, h_hot, layer, hot_abs_start, backend, None) + })?; + let kv_dim = k_cold_buf.shape()[1]; + let mut k_combined = Array2::::zeros((c + s_hot, kv_dim)); + k_combined.slice_mut(s![..c, ..]).assign(&k_cold); + k_combined.slice_mut(s![c.., ..]).assign(&k_hot); + let mut v_combined = Array2::::zeros((c + s_hot, kv_dim)); + v_combined.slice_mut(s![..c, ..]).assign(&v_cold); + v_combined.slice_mut(s![c.., ..]).assign(&v_hot); + return Some((k_combined, v_combined)); + } + + let (h_full, full_abs_start) = match &rs.cold_residuals { + // 2026-05-19 audit fix: slice to cold_len, not shape()[0]. + Some(cold) if rs.cold_len > 0 => { + let s_cold = rs.cold_len; + let h_cold = cold[layer].slice(s![..s_cold, ..]); + let hidden = h_hot.shape()[1]; + let mut combined = Array2::::zeros((s_cold + s_hot, hidden)); + combined.slice_mut(s![..s_cold, ..]).assign(&h_cold); + combined.slice_mut(s![s_cold.., ..]).assign(h_hot); + (combined, rs.cold_abs_start) + } + _ => (h_hot.clone(), hot_abs_start), + }; + StepTimings::measure(timed, &mut timings.recompute_cold_us, || { + recompute_kv(weights, &h_full, layer, full_abs_start, backend, None) + }) +} + +/// Parity gate for the f32 path: the cached prior K/V must match a fresh f32 +/// `recompute_kv`. +/// +/// Only meaningful when attention is NOT on the Q4K-direct route — that +/// route's projections differ from `recompute_kv` by more than the bound even +/// in f32-activation (different kernels/byte sources), so it has its own +/// oracles: the compute-level bit-identity test (`run_..._inplace` ≡ the +/// concat form) and the engine-level in-place-vs-owned-concat A/B test. +#[cfg(debug_assertions)] +fn debug_assert_hot_kv_parity( + weights: larql_inference::WeightsView, + rs: &RsStore, + layer: usize, + bufs: &[SharedKV], + s_hot: usize, + hot_abs_start: usize, + backend: &dyn ComputeBackend, +) { + /// Largest per-element f32 gap tolerated between the cached prior K/V and + /// a fresh recompute. Loose enough for accumulation order, tight enough + /// that a genuinely stale cache trips it. + const MAX_CACHE_DRIFT: f32 = 1e-2; + + if larql_compute::options::q4k_direct_attn_enabled() { + return; + } + let (k_buf, v_buf) = &bufs[layer]; + let h_logical = rs.stored[layer].slice(s![..s_hot, ..]).to_owned(); + let Some((rk, rv)) = recompute_kv(weights, &h_logical, layer, hot_abs_start, backend, None) + else { + return; + }; + let max_gap = |buf: &Array2, fresh: &Array2| { + buf.slice(s![..s_hot, ..]) + .iter() + .zip(fresh.iter()) + .map(|(a, b)| (a - b).abs()) + .fold(0.0f32, f32::max) + }; + let kd = max_gap(k_buf, &rk); + let vd = max_gap(v_buf, &rv); + debug_assert!(kd < MAX_CACHE_DRIFT, "markov hot_kv K cache diverged: {kd}"); + debug_assert!(vd < MAX_CACHE_DRIFT, "markov hot_kv V cache diverged: {vd}"); +} diff --git a/crates/larql-kv/src/engines/markov_residual/walk.rs b/crates/larql-kv/src/engines/markov_residual/walk.rs index 1b9f941ed..1d525ea2a 100644 --- a/crates/larql-kv/src/engines/markov_residual/walk.rs +++ b/crates/larql-kv/src/engines/markov_residual/walk.rs @@ -11,7 +11,8 @@ use larql_compute::ComputeBackend; use larql_vindex::VectorIndex; use ndarray::Array2; -use super::compute::{last_row, recompute_kv, RsPrefillResult}; +use super::compute::{last_row, recompute_kv}; +use super::prefill::RsPrefillResult; use super::store::RsStore; use crate::profiler::EngineProfiler; use larql_inference::attention::run_attention_with_kv_backend; diff --git a/crates/larql-kv/src/engines/markov_residual_codec/compute.rs b/crates/larql-kv/src/engines/markov_residual_codec/compute.rs deleted file mode 100644 index 28bcd075e..000000000 --- a/crates/larql-kv/src/engines/markov_residual_codec/compute.rs +++ /dev/null @@ -1,614 +0,0 @@ -//! Core forward primitives for `MarkovResidualCodecEngine`. -//! -//! Mirrors `markov_residual::compute` with the cold tier swapped to a -//! codec-encoded representation. All forward compute (attention, FFN, K/V -//! recomputation) delegates to `larql_inference` / the production -//! `recompute_kv`. The differences are isolated to cold-tier read/write paths. - -use larql_compute::ComputeBackend; -use larql_inference::attention::{run_attention_with_kv_backend, SharedKV}; -use larql_inference::ffn::BackendFfn; -use larql_inference::forward::embed_tokens_pub; -use larql_inference::forward::ple::precompute_per_layer_inputs; -use ndarray::{s, Array2}; - -use crate::engines::markov_residual::recompute_kv; -use crate::engines::markov_residual_codec::codec::ColdResidualCodec; -use crate::engines::markov_residual_codec::helpers::append_row; -use crate::engines::markov_residual_codec::store::{EncodedColdLayer, RsStoreCodec}; - -pub struct RsPrefillResultCodec { - pub hidden: Array2, - pub store: RsStoreCodec, -} - -#[allow(clippy::too_many_arguments)] -pub fn rs_prefill_codec( - weights: larql_inference::WeightsView, - token_ids: &[u32], - max_window: Option, - codec: ColdResidualCodec, - backend: &dyn ComputeBackend, - moe_ffn: Option<&dyn larql_inference::ffn::FfnBackend>, -) -> RsPrefillResultCodec { - let num_layers = weights.num_layers; - let seq_len = token_ids.len(); - let mut h = embed_tokens_pub(&weights, token_ids); - // Empty on non-PLE archs — `ple_inputs.get(layer)` then yields `None`. - let ple_inputs = precompute_per_layer_inputs(&weights, &h, token_ids); - let mut stored: Vec> = Vec::with_capacity(num_layers); - let be = Some(backend); - - for layer in 0..num_layers { - stored.push(h.clone()); - let (h_post_attn, _k, _v) = run_attention_with_kv_backend(weights, &h, layer, be, None) - .expect("attention failed during MarkovResidualCodec prefill"); - let bffn = BackendFfn { - weights: weights.canonical(), - backend, - }; - let h_out = crate::engines::layer_ffn_or_moe( - weights.canonical(), - &h_post_attn, - layer, - &bffn, - moe_ffn, - ple_inputs.get(layer), - ); - h = h_out; - } - - let hidden_size = weights.hidden_size; - let mut rs = RsStoreCodec { - hot_len: stored.first().map_or(0, |s| s.shape()[0]), - stored, - cold_encoded: None, - cold_kv: None, - // Dense (f32) prefill path doesn't capture K/V — falls back to - // recompute-from-residuals on decode. The Q4K walk path - // (`rs_prefill_codec_walk`) is what production uses, and it - // does capture. - hot_kv: None, - cold_abs_start: 0, - next_position: seq_len, - max_window, - codec, - }; - - // Clip overflow per layer; encode and pre-compute K/V for cold once. - let mut overflow_per_layer: Vec> = Vec::with_capacity(num_layers); - for layer in 0..num_layers { - overflow_per_layer.push(rs.clip_layer_overflow(layer)); - } - rs.finalise_hot_len_after_clip(); - if overflow_per_layer.first().map_or(0, |c| c.shape()[0]) > 0 { - let mut encoded_layers: Vec = Vec::with_capacity(num_layers); - let mut cold_kv: Vec = Vec::with_capacity(num_layers); - for (layer, overflow) in overflow_per_layer.iter().enumerate() { - let decoded_overflow = roundtrip(overflow, codec); - let (k, v) = recompute_kv(weights, &decoded_overflow, layer, 0, backend, None) - .expect("cold K/V pre-computation failed"); - cold_kv.push((k, v)); - let mut enc = EncodedColdLayer::empty(hidden_size); - enc.append(codec, overflow); - encoded_layers.push(enc); - } - rs.cold_encoded = Some(encoded_layers); - rs.cold_kv = Some(cold_kv); - rs.cold_abs_start = 0; - } - - RsPrefillResultCodec { - hidden: last_row(&h), - store: rs, - } -} - -pub fn rs_decode_step_codec( - weights: larql_inference::WeightsView, - new_token_id: u32, - rs: RsStoreCodec, - backend: &dyn ComputeBackend, - moe_ffn: Option<&dyn larql_inference::ffn::FfnBackend>, - index: Option<&larql_vindex::VectorIndex>, -) -> Option<(Array2, RsStoreCodec)> { - let num_layers = weights.num_layers; - let abs_position = rs.next_position; - let mut h_new = embed_tokens_pub(&weights, &[new_token_id]); - // PLE inputs are per-token — recompute for this single-token decode - // step, matching the legacy `kv_decode_step_run` recipe exactly. - let ple_inputs = precompute_per_layer_inputs(&weights, &h_new, &[new_token_id]); - let mut new_stored: Vec> = Vec::with_capacity(num_layers); - - // W2 hot-K/V cache on the resident walk (2026-06-13), twin of - // markov_residual: with no cold tier, `hot_kv` holds the FULL K/V and is - // read instead of re-deriving via `recompute_kv` each step. `stored` - // remains the canonical re-derivable state. `step_new_kv` collects the - // attention step's updated full K/V (= next step's cache). - // Only for unbounded windows (the default): `clip_layer_overflow` is then a - // no-op, so the cache never tracks a window-eviction transition. Windowed - // configs keep the existing recompute path unchanged. - let cache_eligible = - rs.max_window.is_none() && rs.cold_encoded.is_none() && rs.cold_kv.is_none(); - let mut step_new_kv: Vec = Vec::with_capacity(num_layers); - // Move the hot K/V cache out so the steady state (step 2+) can append in - // place — twin of `markov_residual::compute::rs_decode_step_inner`. - let mut hot_kv_store = rs.hot_kv; - let had_hot_kv = hot_kv_store.is_some(); - let idx_kv: Option<&dyn larql_compute::KvIndex> = - index.map(|v| v as &dyn larql_compute::KvIndex); - let inplace_enabled = crate::engines::markov_residual::compute::markov_inplace_kv_enabled(); - - for layer in 0..num_layers { - // `stored` is a doubling-capacity buffer (W8.2): logical row count is - // `hot_len`, not `shape()[0]`. - let s_hot = rs.hot_len; - let hot_abs_start = abs_position.saturating_sub(s_hot); - - new_stored.push(h_new.clone()); - - let h_post_attn = if cache_eligible && had_hot_kv { - // STEADY STATE (step 2+): append this token's projected+RoPE'd K/V row - // IN PLACE into the doubling-capacity `hot_kv` buffer and attend over - // the `[..s_hot+1]` views — no per-step O(ctx) owned concat (O(L) - // total cache copy vs O(L²)). See the markov twin for the rationale. - let bufs = hot_kv_store.as_mut().expect("had_hot_kv"); - #[cfg(debug_assertions)] - { - // f32-path parity gate only (the Q4K-direct route has its own - // oracles: the compute-level bit-identity test + the engine A/B). - let q4k_on = larql_compute::options::q4k_direct_attn_enabled(); - if !q4k_on { - let (k_buf, v_buf) = &bufs[layer]; - let h_logical = rs.stored[layer].slice(s![..s_hot, ..]).to_owned(); - if let Some((rk, rv)) = - recompute_kv(weights, &h_logical, layer, hot_abs_start, backend, None) - { - let kd = k_buf - .slice(s![..s_hot, ..]) - .iter() - .zip(rk.iter()) - .map(|(a, b)| (a - b).abs()) - .fold(0.0f32, f32::max); - let vd = v_buf - .slice(s![..s_hot, ..]) - .iter() - .zip(rv.iter()) - .map(|(a, b)| (a - b).abs()) - .fold(0.0f32, f32::max); - debug_assert!(kd < 1e-2, "codec hot_kv K cache diverged: {kd}"); - debug_assert!(vd < 1e-2, "codec hot_kv V cache diverged: {vd}"); - } - } - } - let (k_buf, v_buf) = &mut bufs[layer]; - let inplace = if inplace_enabled { - larql_inference::attention::run_attention_block_decode_step_auto_inplace( - weights, - &h_new, - layer, - k_buf, - v_buf, - s_hot, - abs_position, - Some(backend), - idx_kv, - ) - } else { - None - }; - match inplace { - Some(h) => h, - None => { - // Q4K-direct off (flags-off parity) or no attn bytes: owned - // concat over the buffer view, then replace. Bit-identical to - // the legacy borrow path. - let prior: SharedKV = ( - k_buf.slice(s![..s_hot, ..]).to_owned(), - v_buf.slice(s![..s_hot, ..]).to_owned(), - ); - let (h, new_kv) = - larql_inference::attention::run_attention_block_decode_step_auto( - weights, - &h_new, - layer, - Some(&prior), - abs_position, - Some(backend), - idx_kv, - )?; - *k_buf = new_kv.0; - *v_buf = new_kv.1; - h - } - } - } else { - // FIRST STEP (cache None → seed) or windowed/cold tier. - let h_hot = &rs.stored[layer]; - let kv_arg: SharedKV = if let Some(cold_kv) = &rs.cold_kv { - let (k_cold, v_cold) = &cold_kv[layer]; - let (k_hot, v_hot) = - recompute_kv(weights, h_hot, layer, hot_abs_start, backend, None)?; - let c = k_cold.shape()[0]; - let kv_dim = k_cold.shape()[1]; - let mut k_combined = Array2::::zeros((c + s_hot, kv_dim)); - k_combined.slice_mut(s![..c, ..]).assign(k_cold); - k_combined.slice_mut(s![c.., ..]).assign(&k_hot); - let mut v_combined = Array2::::zeros((c + s_hot, kv_dim)); - v_combined.slice_mut(s![..c, ..]).assign(v_cold); - v_combined.slice_mut(s![c.., ..]).assign(&v_hot); - (k_combined, v_combined) - } else { - let (h_full, full_abs_start) = if let Some(cold_layers) = &rs.cold_encoded { - let enc = &cold_layers[layer]; - if enc.n_positions > 0 { - let decoded = enc.decode(rs.codec); - let hidden = h_hot.shape()[1]; - let mut combined = - Array2::::zeros((decoded.shape()[0] + s_hot, hidden)); - combined - .slice_mut(s![..decoded.shape()[0], ..]) - .assign(&decoded); - combined - .slice_mut(s![decoded.shape()[0].., ..]) - .assign(h_hot); - (combined, rs.cold_abs_start) - } else { - (h_hot.clone(), hot_abs_start) - } - } else { - (h_hot.clone(), hot_abs_start) - }; - let (k, v) = recompute_kv(weights, &h_full, layer, full_abs_start, backend, None)?; - (k, v) - }; - - let (h_post_attn, new_kv) = - larql_inference::attention::run_attention_block_decode_step_auto( - weights, - &h_new, - layer, - Some(&kv_arg), - abs_position, - Some(backend), - idx_kv, - )?; - if cache_eligible { - step_new_kv.push(new_kv); - } - h_post_attn - }; - - let bffn = BackendFfn { - weights: weights.canonical(), - backend, - }; - let h_out = crate::engines::layer_ffn_or_moe( - weights.canonical(), - &h_post_attn, - layer, - &bffn, - moe_ffn, - ple_inputs.get(layer), - ); - h_new = h_out; - } - - // Append the new row to each layer's hot tier. W8.2: in the cache_eligible - // path `stored` is a doubling-capacity buffer (no window → never clips), so - // append in place rather than allocating + bzeroing a fresh `[s_old+1, - // hidden]` array every step (the resident walk's dominant per-step malloc; - // see helpers::append_row). The windowed/cold path keeps the rebuild. - let (updated_stored, new_hot_len) = if cache_eligible { - let mut buf = rs.stored; - for (layer, new_row) in new_stored.iter().enumerate() { - append_row(&mut buf[layer], new_row, rs.hot_len); - } - (buf, rs.hot_len + 1) - } else { - let mut rebuilt: Vec> = Vec::with_capacity(num_layers); - for (stored, new_row) in rs.stored.iter().zip(new_stored.iter()) { - let s_old = stored.shape()[0]; - let hidden_dim = stored.shape()[1]; - let mut combined = Array2::::zeros((s_old + 1, hidden_dim)); - combined.slice_mut(s![..s_old, ..]).assign(stored); - combined.slice_mut(s![s_old.., ..]).assign(new_row); - rebuilt.push(combined); - } - let len = rebuilt.first().map_or(0, |s| s.shape()[0]); - (rebuilt, len) - }; - - let mut updated_rs = RsStoreCodec { - hot_len: new_hot_len, - stored: updated_stored, - cold_encoded: rs.cold_encoded, - cold_kv: rs.cold_kv, - // Cache the full K/V for next step when there's no cold tier; else None - // (cold/windowed recomputes). clip_layer_overflow clips hot_kv in step. - // Step 2+ mutated `hot_kv_store` in place; the first step seeds it. - hot_kv: if cache_eligible { - if had_hot_kv { - hot_kv_store - } else { - Some(step_new_kv) - } - } else { - None - }, - cold_abs_start: rs.cold_abs_start, - next_position: abs_position + 1, - max_window: rs.max_window, - codec: rs.codec, - }; - - // Clip overflow into encoded cold tier; clear cold_kv to force recompute. - let mut overflow_per_layer: Vec> = Vec::with_capacity(num_layers); - for layer in 0..num_layers { - overflow_per_layer.push(updated_rs.clip_layer_overflow(layer)); - } - updated_rs.finalise_hot_len_after_clip(); - if overflow_per_layer.first().map_or(0, |c| c.shape()[0]) > 0 { - match updated_rs.cold_encoded.as_mut() { - Some(layers) => { - for (layer, overflow) in overflow_per_layer.iter().enumerate() { - layers[layer].append(updated_rs.codec, overflow); - } - } - None => { - let hidden = weights.hidden_size; - let mut layers: Vec = Vec::with_capacity(num_layers); - for overflow in overflow_per_layer.iter() { - let mut enc = EncodedColdLayer::empty(hidden); - enc.append(updated_rs.codec, overflow); - layers.push(enc); - } - updated_rs.cold_encoded = Some(layers); - } - } - updated_rs.cold_kv = None; - } - - Some((last_row(&h_new), updated_rs)) -} - -/// Apply the codec roundtrip to a block. Used during prefill cold setup so -/// that the cold K/V we precompute is consistent with what `decode` would -/// later produce. -fn roundtrip(block: &Array2, codec: ColdResidualCodec) -> Array2 { - if block.shape()[0] == 0 { - return block.clone(); - } - let mut tmp = EncodedColdLayer::empty(block.shape()[1]); - tmp.append(codec, block); - tmp.decode(codec) -} - -fn last_row(h: &Array2) -> Array2 { - let last = h.shape()[0] - 1; - h.slice(s![last..=last, ..]).to_owned() -} - -#[cfg(test)] -mod tests { - use super::*; - use larql_compute::CpuBackend; - use larql_inference::test_utils::make_test_weights; - - #[test] - fn prefill_returns_finite_hidden() { - let weights = make_test_weights(); - let result = rs_prefill_codec( - larql_inference::WeightsView::dense(&weights), - &[0u32, 1, 2], - None, - ColdResidualCodec::Bf16, - &CpuBackend, - None, - ); - assert_eq!(result.hidden.shape(), &[1, weights.hidden_size]); - assert!(result.hidden.iter().all(|v| v.is_finite())); - } - - #[test] - fn prefill_no_window_does_not_create_cold_tier() { - let weights = make_test_weights(); - let result = rs_prefill_codec( - larql_inference::WeightsView::dense(&weights), - &[0u32, 1], - None, - ColdResidualCodec::Bf16, - &CpuBackend, - None, - ); - assert!(result.store.cold_encoded.is_none()); - assert!(result.store.cold_kv.is_none()); - } - - #[test] - fn prefill_with_overflow_creates_encoded_cold_tier() { - let weights = make_test_weights(); - let result = rs_prefill_codec( - larql_inference::WeightsView::dense(&weights), - &[0u32, 1, 2, 3], - Some(2), - ColdResidualCodec::Bf16, - &CpuBackend, - None, - ); - assert!(result.store.cold_encoded.is_some()); - assert!(result.store.cold_kv.is_some()); - let layers = result.store.cold_encoded.as_ref().unwrap(); - assert_eq!(layers.len(), weights.num_layers); - // 4 tokens, window=2 → 2 cold positions per layer. - for l in layers { - assert_eq!(l.n_positions, 2); - assert_eq!(l.payload.len(), 2 * weights.hidden_size * 2); - } - } - - #[test] - fn decode_step_extends_position() { - let weights = make_test_weights(); - let prefill = rs_prefill_codec( - larql_inference::WeightsView::dense(&weights), - &[0u32, 1], - None, - ColdResidualCodec::Bf16, - &CpuBackend, - None, - ); - assert_eq!(prefill.store.next_position, 2); - let (_, rs2) = rs_decode_step_codec( - larql_inference::WeightsView::dense(&weights), - 2, - prefill.store, - &CpuBackend, - None, - None, - ) - .unwrap(); - assert_eq!(rs2.next_position, 3); - } - - #[test] - fn decode_with_cold_kv_path_produces_finite_output() { - let weights = make_test_weights(); - let prefill = rs_prefill_codec( - larql_inference::WeightsView::dense(&weights), - &[0u32, 1, 2, 3], - Some(2), - ColdResidualCodec::Bf16, - &CpuBackend, - None, - ); - assert!(prefill.store.cold_kv.is_some()); - let (h, _) = rs_decode_step_codec( - larql_inference::WeightsView::dense(&weights), - 4, - prefill.store, - &CpuBackend, - None, - None, - ) - .unwrap(); - assert_eq!(h.shape(), &[1, weights.hidden_size]); - assert!(h.iter().all(|v| v.is_finite())); - } - - #[test] - fn decode_with_cold_encoded_path_produces_finite_output() { - // After enough decode steps, the post-eviction cold_kv-clear path is - // exercised (we read from cold_encoded directly via decode). - let weights = make_test_weights(); - let prefill = rs_prefill_codec( - larql_inference::WeightsView::dense(&weights), - &[0u32, 1, 2, 3], - Some(2), - ColdResidualCodec::Bf16, - &CpuBackend, - None, - ); - let (_, rs2) = rs_decode_step_codec( - larql_inference::WeightsView::dense(&weights), - 4, - prefill.store, - &CpuBackend, - None, - None, - ) - .unwrap(); - // Second decode: cold_kv was cleared by overflow at the first decode, - // so this step exercises the cold_encoded recompute branch. - let (h, _) = rs_decode_step_codec( - larql_inference::WeightsView::dense(&weights), - 5, - rs2, - &CpuBackend, - None, - None, - ) - .unwrap(); - assert_eq!(h.shape(), &[1, weights.hidden_size]); - assert!(h.iter().all(|v| v.is_finite())); - } - - #[test] - fn roundtrip_empty_block_short_circuits() { - let empty: Array2 = Array2::zeros((0, 8)); - let out = roundtrip(&empty, ColdResidualCodec::Bf16); - assert_eq!(out.shape(), &[0, 8]); - } - - #[test] - fn roundtrip_preserves_within_bf16_precision() { - let block = - Array2::from_shape_vec((2, 4), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]).unwrap(); - let out = roundtrip(&block, ColdResidualCodec::Bf16); - for (orig, got) in block.iter().zip(out.iter()) { - assert!((orig - got).abs() < 0.1); - } - } - - /// Flags-ON parity gate for the codec engine's in-place hot-K/V fast path: - /// an A/B of the in-place steady state against the owned-concat reference, - /// both with Q4K-direct attention live. Twin of the markov test — the two - /// paths must produce bit-identical hidden states at every step. Twin of the - /// markov test; q4k flags driven via the thread-local override (no env race), - /// in-place path selected through the shared `LARQL_MARKOV_INPLACE_KV` - /// thread-local override. - #[test] - fn rs_decode_step_codec_inplace_matches_owned_concat_flags_on() { - use crate::engines::markov_residual::compute::set_markov_env_override; - use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; - - let _q4k = crate::engines::Q4kFlagGuard::set(&[ - (larql_compute::options::ENV_Q4K_DIRECT_ATTN, true), - (larql_compute::options::ENV_Q4K_ATTN_INT8, false), - ]); - - let weights = make_test_q4k_weights(); - let index = make_test_q4k_vindex(&weights); - - let run = |inplace: bool| -> (Vec>, usize) { - set_markov_env_override( - "LARQL_MARKOV_INPLACE_KV", - Some(if inplace { "1" } else { "0" }), - ); - let prefill = rs_prefill_codec( - larql_inference::WeightsView::dense(&weights), - &[0u32, 1, 2], - None, - ColdResidualCodec::Bf16, - &CpuBackend, - None, - ); - let mut rs = prefill.store; - let mut hiddens = Vec::new(); - for tok in 3u32..=12 { - let (h, rs2) = rs_decode_step_codec( - larql_inference::WeightsView::dense(&weights), - tok, - rs, - &CpuBackend, - None, - Some(&index), - ) - .expect("decode"); - assert!(h.iter().all(|v| v.is_finite())); - hiddens.push(h.iter().map(|v| v.to_bits()).collect()); - rs = rs2; - } - (hiddens, rs.hot_len) - }; - - let (a_hiddens, a_len) = run(true); - let (b_hiddens, b_len) = run(false); - assert_eq!(a_len, 13, "3 prompt + 10 decode rows"); - assert_eq!(a_len, b_len); - assert_eq!( - a_hiddens, b_hiddens, - "codec in-place and owned-concat hidden states diverged (q4k-direct on)" - ); - } -} diff --git a/crates/larql-kv/src/engines/markov_residual_codec/engine.rs b/crates/larql-kv/src/engines/markov_residual_codec/engine.rs index c086a7a76..ba511549a 100644 --- a/crates/larql-kv/src/engines/markov_residual_codec/engine.rs +++ b/crates/larql-kv/src/engines/markov_residual_codec/engine.rs @@ -6,7 +6,7 @@ //! - this file: struct + construction + `KvEngine` trait glue //! - [`super::walk`] — CPU dense walk path //! (`rs_prefill_codec_walk` / `rs_decode_step_codec_walk`) -//! - [`super::compute`] — Q4K-native walk path +//! - [`super::prefill`] / [`super::step`] — Q4K-native walk path //! (`rs_prefill_codec` / `rs_decode_step_codec`) //! - [`super::dispatch`] — W1-GPU dispatch fast path with W10 mask //! cascade @@ -22,7 +22,8 @@ use ndarray::Array2; use crate::engines::markov_residual::engine::check_residual_recompute_preconditions; use crate::engines::markov_residual::ensure_attn_tensors_dequantised; use crate::engines::markov_residual_codec::codec::ColdResidualCodec; -use crate::engines::markov_residual_codec::compute::{rs_decode_step_codec, rs_prefill_codec}; +use crate::engines::markov_residual_codec::prefill::rs_prefill_codec; +use crate::engines::markov_residual_codec::step::rs_decode_step_codec; use crate::engines::markov_residual_codec::store::RsStoreCodec; use crate::engines::markov_residual_codec::walk::{ rs_decode_step_codec_walk, rs_prefill_codec_walk, @@ -81,8 +82,13 @@ impl MarkovResidualCodecEngine { self.codec } + /// Residual store + backend-resident K/V — see + /// [`super::super::markov_residual::MarkovResidualEngine::total_memory_bytes`] + /// for why the backend term is needed on the W1-GPU path. pub fn total_memory_bytes(&self) -> usize { self.store.as_ref().map_or(0, |s| s.memory_bytes()) + + self.kv_handle.as_ref().map_or(0, |h| h.resident_bytes()) + + self.backend.backend_resident_kv_bytes() } } @@ -96,6 +102,10 @@ impl MarkovResidualCodecEngine { impl MarkovResidualCodecEngine { /// Shared body for `decode_step` / `decode_step_resident`. + /// + /// The store is borrowed, never taken: a failed step leaves canonical + /// state exactly as it was (see [`super::step`]'s failure invariant), so a + /// refusal reports itself and the engine stays usable. fn decode_step_impl( &mut self, weights: &ModelWeights, @@ -103,25 +113,21 @@ impl MarkovResidualCodecEngine { token_id: u32, index: Option<&larql_vindex::VectorIndex>, ) -> Result, EngineError> { + let backend = self.backend.as_ref(); let rs = self .store - .take() + .as_mut() .ok_or_else(|| EngineError::InvariantViolation { what: "decode_step called before prefill (store missing)".into(), })?; - let (hidden, new_rs) = rs_decode_step_codec( + rs_decode_step_codec( larql_inference::WeightsView::dense(weights), token_id, rs, - self.backend.as_ref(), + backend, Some(ffn), index, ) - .ok_or_else(|| EngineError::BackendFailure { - details: "rs_decode_step_codec returned None".into(), - })?; - self.store = Some(new_rs); - Ok(hidden) } } @@ -158,6 +164,8 @@ impl KvEngine for MarkovResidualCodecEngine { if token_ids.is_empty() { return Err(EngineError::EmptyPrompt); } + // `?` before the assignment: a refused prefill must not replace the + // store an earlier one built. let result = rs_prefill_codec( larql_inference::WeightsView::dense(weights), token_ids, @@ -165,10 +173,9 @@ impl KvEngine for MarkovResidualCodecEngine { self.codec, self.backend.as_ref(), Some(ffn), - ); - let hidden = result.hidden.clone(); + )?; self.store = Some(result.store); - Ok(hidden) + Ok(result.hidden) } fn decode_step( @@ -204,6 +211,17 @@ impl KvEngine for MarkovResidualCodecEngine { self.store.as_ref().map_or(0, |s| s.cold_bytes()) } + fn dispatch_path(&self) -> Option { + use larql_inference::kv_engine::DispatchPath; + // Same rule as `markov_residual`: `kv_handle` marks the W1-GPU + // coarse path, `store` marks "prefill happened at all". + match (self.kv_handle.is_some(), self.store.is_some()) { + (true, _) => Some(DispatchPath::Coarse), + (false, true) => Some(DispatchPath::PerLayer), + (false, false) => None, + } + } + fn prefill_quant( &mut self, weights: &ModelWeights, diff --git a/crates/larql-kv/src/engines/markov_residual_codec/helpers.rs b/crates/larql-kv/src/engines/markov_residual_codec/helpers.rs index cd1dd8101..c3205858a 100644 --- a/crates/larql-kv/src/engines/markov_residual_codec/helpers.rs +++ b/crates/larql-kv/src/engines/markov_residual_codec/helpers.rs @@ -48,3 +48,10 @@ pub(super) fn append_row(buf: &mut Array2, row: &Array2, len: usize) { } buf.slice_mut(s![len..len + 1, ..]).assign(row); } + +/// The final row of a `[rows, hidden]` block, as its own `[1, hidden]` array +/// — what a prefill or decode step reports as "the" hidden state. +pub(super) fn last_row(h: &Array2) -> Array2 { + let last = h.shape()[0] - 1; + h.slice(s![last..=last, ..]).to_owned() +} diff --git a/crates/larql-kv/src/engines/markov_residual_codec/mod.rs b/crates/larql-kv/src/engines/markov_residual_codec/mod.rs index 76897e3c7..34316cdb0 100644 --- a/crates/larql-kv/src/engines/markov_residual_codec/mod.rs +++ b/crates/larql-kv/src/engines/markov_residual_codec/mod.rs @@ -13,11 +13,13 @@ //! residuals, which is small but non-zero. pub mod codec; -pub mod compute; pub(crate) mod dispatch; pub mod engine; pub(crate) mod executor; pub(crate) mod helpers; +pub mod prefill; +pub mod step; +mod step_attention; pub mod store; pub mod walk; diff --git a/crates/larql-kv/src/engines/markov_residual_codec/prefill.rs b/crates/larql-kv/src/engines/markov_residual_codec/prefill.rs new file mode 100644 index 000000000..e96a2f4e1 --- /dev/null +++ b/crates/larql-kv/src/engines/markov_residual_codec/prefill.rs @@ -0,0 +1,207 @@ +//! Codec-cold-tier prefill: build the store the decode step walks. +//! +//! Transactional by construction, exactly as the markov twin: the store is +//! built into locals and handed back only on success, so a refusal leaves an +//! engine's existing store untouched. + +use larql_compute::ComputeBackend; +use larql_inference::attention::{run_attention_with_kv_backend, SharedKV}; +use larql_inference::ffn::BackendFfn; +use larql_inference::forward::embed_tokens_pub; +use larql_inference::forward::ple::precompute_per_layer_inputs; +use larql_inference::kv_engine::EngineError; +use ndarray::Array2; + +use crate::engines::markov_residual::recompute_kv; +use crate::engines::markov_residual_codec::codec::ColdResidualCodec; +use crate::engines::markov_residual_codec::helpers::last_row; +use crate::engines::markov_residual_codec::store::{EncodedColdLayer, RsStoreCodec}; + +pub struct RsPrefillResultCodec { + pub hidden: Array2, + pub store: RsStoreCodec, +} + +#[allow(clippy::too_many_arguments)] +pub fn rs_prefill_codec( + weights: larql_inference::WeightsView, + token_ids: &[u32], + max_window: Option, + codec: ColdResidualCodec, + backend: &dyn ComputeBackend, + moe_ffn: Option<&dyn larql_inference::ffn::FfnBackend>, +) -> Result { + let num_layers = weights.num_layers; + let seq_len = token_ids.len(); + let mut h = embed_tokens_pub(&weights, token_ids); + // Empty on non-PLE archs — `ple_inputs.get(layer)` then yields `None`. + let ple_inputs = precompute_per_layer_inputs(&weights, &h, token_ids); + let mut stored: Vec> = Vec::with_capacity(num_layers); + let be = Some(backend); + + for layer in 0..num_layers { + stored.push(h.clone()); + let (h_post_attn, _k, _v) = run_attention_with_kv_backend(weights, &h, layer, be, None) + .ok_or_else(|| EngineError::BackendFailure { + details: format!("attention returned None during codec prefill at layer {layer}"), + })?; + let bffn = BackendFfn { + weights: weights.canonical(), + backend, + }; + h = crate::engines::layer_ffn_or_moe( + weights.canonical(), + &h_post_attn, + layer, + &bffn, + moe_ffn, + ple_inputs.get(layer), + ) + .map_err(EngineError::Execution)?; + } + + let hidden_size = weights.hidden_size; + let mut rs = RsStoreCodec { + hot_len: stored.first().map_or(0, |s| s.shape()[0]), + stored, + cold_encoded: None, + cold_kv: None, + // Dense (f32) prefill path doesn't capture K/V — falls back to + // recompute-from-residuals on decode. The Q4K walk path + // (`rs_prefill_codec_walk`) is what production uses, and it + // does capture. + hot_kv: None, + cold_abs_start: 0, + next_position: seq_len, + max_window, + codec, + }; + + // Clip overflow per layer; encode and pre-compute K/V for cold once. + let mut overflow_per_layer: Vec> = Vec::with_capacity(num_layers); + for layer in 0..num_layers { + overflow_per_layer.push(rs.clip_layer_overflow(layer)); + } + rs.finalise_hot_len_after_clip(); + if overflow_per_layer.first().map_or(0, |c| c.shape()[0]) > 0 { + let mut encoded_layers: Vec = Vec::with_capacity(num_layers); + let mut cold_kv: Vec = Vec::with_capacity(num_layers); + for (layer, overflow) in overflow_per_layer.iter().enumerate() { + let decoded_overflow = roundtrip(overflow, codec); + cold_kv.push( + recompute_kv(weights, &decoded_overflow, layer, 0, backend, None).ok_or_else( + || EngineError::BackendFailure { + details: format!("cold K/V pre-computation returned None at layer {layer}"), + }, + )?, + ); + let mut enc = EncodedColdLayer::empty(hidden_size); + enc.append(codec, overflow); + encoded_layers.push(enc); + } + rs.cold_encoded = Some(encoded_layers); + rs.cold_kv = Some(cold_kv); + rs.cold_abs_start = 0; + } + + Ok(RsPrefillResultCodec { + hidden: last_row(&h), + store: rs, + }) +} + +/// Apply the codec roundtrip to a block. Used during prefill cold setup so +/// that the cold K/V we precompute is consistent with what `decode` would +/// later produce. +pub(super) fn roundtrip(block: &Array2, codec: ColdResidualCodec) -> Array2 { + if block.shape()[0] == 0 { + return block.clone(); + } + let mut tmp = EncodedColdLayer::empty(block.shape()[1]); + tmp.append(codec, block); + tmp.decode(codec) +} + +#[cfg(test)] +mod tests { + use super::{roundtrip, rs_prefill_codec, ColdResidualCodec}; + use larql_compute::CpuBackend; + use larql_inference::test_utils::make_test_weights; + use ndarray::Array2; + + /// Prompt/window fixtures. Named so an arity change reads as a decision. + const PROMPT: [u32; 3] = [0, 1, 2]; + const OVERFLOW_PROMPT: [u32; 4] = [0, 1, 2, 3]; + const WINDOW: usize = 2; + /// `Bf16` is two bytes per element — the payload-size assertion below. + const BF16_BYTES: usize = 2; + + fn prefill( + weights: &larql_inference::ModelWeights, + tokens: &[u32], + window: Option, + ) -> super::RsPrefillResultCodec { + rs_prefill_codec( + larql_inference::WeightsView::dense(weights), + tokens, + window, + ColdResidualCodec::Bf16, + &CpuBackend, + None, + ) + .expect("a dense prefill with no MoE hook cannot refuse") + } + + #[test] + fn prefill_returns_finite_hidden() { + let weights = make_test_weights(); + let result = prefill(&weights, &PROMPT, None); + assert_eq!(result.hidden.shape(), &[1, weights.hidden_size]); + assert!(result.hidden.iter().all(|v| v.is_finite())); + } + + #[test] + fn prefill_no_window_does_not_create_cold_tier() { + let weights = make_test_weights(); + let result = prefill(&weights, &PROMPT[..2], None); + assert!(result.store.cold_encoded.is_none()); + assert!(result.store.cold_kv.is_none()); + } + + #[test] + fn prefill_with_overflow_creates_encoded_cold_tier() { + let weights = make_test_weights(); + let result = prefill(&weights, &OVERFLOW_PROMPT, Some(WINDOW)); + assert!(result.store.cold_encoded.is_some()); + assert!(result.store.cold_kv.is_some()); + let layers = result.store.cold_encoded.as_ref().unwrap(); + assert_eq!(layers.len(), weights.num_layers); + let cold_positions = OVERFLOW_PROMPT.len() - WINDOW; + for l in layers { + assert_eq!(l.n_positions, cold_positions); + assert_eq!( + l.payload.len(), + cold_positions * weights.hidden_size * BF16_BYTES + ); + } + } + + #[test] + fn roundtrip_empty_block_short_circuits() { + let empty: Array2 = Array2::zeros((0, 8)); + let out = roundtrip(&empty, ColdResidualCodec::Bf16); + assert_eq!(out.shape(), &[0, 8]); + } + + #[test] + fn roundtrip_preserves_within_bf16_precision() { + /// Loosest gap a bf16 round-trip may open on the small integers below. + const BF16_TOLERANCE: f32 = 0.1; + let block = + Array2::from_shape_vec((2, 4), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]).unwrap(); + let out = roundtrip(&block, ColdResidualCodec::Bf16); + for (orig, got) in block.iter().zip(out.iter()) { + assert!((orig - got).abs() < BF16_TOLERANCE); + } + } +} diff --git a/crates/larql-kv/src/engines/markov_residual_codec/step/commit.rs b/crates/larql-kv/src/engines/markov_residual_codec/step/commit.rs new file mode 100644 index 000000000..9fc80d975 --- /dev/null +++ b/crates/larql-kv/src/engines/markov_residual_codec/step/commit.rs @@ -0,0 +1,90 @@ +//! Writing a completed decode step into the codec store. +//! +//! Twin of [`crate::engines::markov_residual::step::commit`], with the cold +//! tier encoded rather than stored as f32. Nothing here can fail, which is +//! what makes [`super`]'s failure invariant structural. + +use larql_inference::attention::SharedKV; +use ndarray::{s, Array2}; + +use crate::engines::markov_residual_codec::helpers::append_row; +use crate::engines::markov_residual_codec::store::{EncodedColdLayer, RsStoreCodec}; + +/// Write the completed step into the store — every canonical mutation, after +/// the last fallible call. +#[allow(clippy::too_many_arguments)] +pub(super) fn commit( + hidden_size: usize, + rs: &mut RsStoreCodec, + new_stored: &[Array2], + cache_eligible: bool, + had_hot_kv: bool, + hot_kv_store: Option>, + step_new_kv: Vec, + abs_position: usize, +) { + let num_layers = new_stored.len(); + + // Append the new row to each layer's hot tier. W8.2: in the cache_eligible + // path `stored` is a doubling-capacity buffer (no window → never clips), so + // append in place rather than allocating + bzeroing a fresh `[s_old+1, + // hidden]` array every step (the resident walk's dominant per-step malloc; + // see helpers::append_row). The windowed/cold path keeps the rebuild. + if cache_eligible { + let hot_len = rs.hot_len; + for (layer, new_row) in new_stored.iter().enumerate() { + append_row(&mut rs.stored[layer], new_row, hot_len); + } + rs.hot_len = hot_len + 1; + } else { + let mut rebuilt: Vec> = Vec::with_capacity(num_layers); + for (stored, new_row) in rs.stored.iter().zip(new_stored.iter()) { + let s_old = stored.shape()[0]; + let hidden_dim = stored.shape()[1]; + let mut combined = Array2::::zeros((s_old + 1, hidden_dim)); + combined.slice_mut(s![..s_old, ..]).assign(stored); + combined.slice_mut(s![s_old.., ..]).assign(new_row); + rebuilt.push(combined); + } + rs.hot_len = rebuilt.first().map_or(0, |s| s.shape()[0]); + rs.stored = rebuilt; + } + + // Cache the full K/V for next step when there's no cold tier; else None + // (cold/windowed recomputes). Step 2+ mutated `hot_kv_store` in place; the + // first step seeds it. + rs.hot_kv = match (cache_eligible, had_hot_kv) { + (true, true) => hot_kv_store, + (true, false) => Some(step_new_kv), + (false, _) => None, + }; + rs.next_position = abs_position + 1; + + // Clip overflow into the encoded cold tier; clear cold_kv to force + // recompute, because the codec round-trip is lossy. + let mut overflow_per_layer: Vec> = Vec::with_capacity(num_layers); + for layer in 0..num_layers { + overflow_per_layer.push(rs.clip_layer_overflow(layer)); + } + rs.finalise_hot_len_after_clip(); + if overflow_per_layer.first().map_or(0, |c| c.shape()[0]) == 0 { + return; + } + match rs.cold_encoded.as_mut() { + Some(layers) => { + for (layer, overflow) in overflow_per_layer.iter().enumerate() { + layers[layer].append(rs.codec, overflow); + } + } + None => { + let mut layers: Vec = Vec::with_capacity(num_layers); + for overflow in overflow_per_layer.iter() { + let mut enc = EncodedColdLayer::empty(hidden_size); + enc.append(rs.codec, overflow); + layers.push(enc); + } + rs.cold_encoded = Some(layers); + } + } + rs.cold_kv = None; +} diff --git a/crates/larql-kv/src/engines/markov_residual_codec/step/mod.rs b/crates/larql-kv/src/engines/markov_residual_codec/step/mod.rs new file mode 100644 index 000000000..850fe184e --- /dev/null +++ b/crates/larql-kv/src/engines/markov_residual_codec/step/mod.rs @@ -0,0 +1,110 @@ +//! One codec-cold-tier decode step. +//! +//! # Failure invariant +//! +//! Identical to [`crate::engines::markov_residual::step`]: the store is +//! borrowed rather than consumed, canonical state (`stored`, `hot_len`, the +//! cold tiers, `next_position`) is written only in [`commit`] after the last +//! fallible call, and `hot_kv` — a droppable derivative — is taken up front +//! and left `None` on the error path. A refused step therefore rewinds +//! exactly, and the same token can be retried on the same engine. + +use larql_compute::ComputeBackend; +use larql_inference::attention::SharedKV; +use larql_inference::ffn::BackendFfn; +use larql_inference::forward::embed_tokens_pub; +use larql_inference::forward::ple::precompute_per_layer_inputs; +use larql_inference::kv_engine::EngineError; +use ndarray::Array2; + +use super::helpers::last_row; +use super::step_attention::{resolve_layer_attention, HotKv}; +use super::store::RsStoreCodec; + +mod commit; +#[cfg(test)] +mod tests; + +use commit::commit; + +/// Advance `rs` by one token, returning the new last hidden row. +pub fn rs_decode_step_codec( + weights: larql_inference::WeightsView, + new_token_id: u32, + rs: &mut RsStoreCodec, + backend: &dyn ComputeBackend, + moe_ffn: Option<&dyn larql_inference::ffn::FfnBackend>, + index: Option<&larql_vindex::VectorIndex>, +) -> Result, EngineError> { + let num_layers = weights.num_layers; + let abs_position = rs.next_position; + let mut h_new = embed_tokens_pub(&weights, &[new_token_id]); + // PLE inputs are per-token — recompute for this single-token decode + // step, matching the legacy `kv_decode_step_run` recipe exactly. + let ple_inputs = precompute_per_layer_inputs(&weights, &h_new, &[new_token_id]); + let mut new_stored: Vec> = Vec::with_capacity(num_layers); + + // W2 hot-K/V cache on the resident walk (2026-06-13), twin of + // markov_residual: with no cold tier, `hot_kv` holds the FULL K/V and is + // read instead of re-deriving via `recompute_kv` each step. `stored` + // remains the canonical re-derivable state. Only for unbounded windows + // (the default): `clip_layer_overflow` is then a no-op, so the cache never + // tracks a window-eviction transition. + let cache_eligible = + rs.max_window.is_none() && rs.cold_encoded.is_none() && rs.cold_kv.is_none(); + let mut step_new_kv: Vec = Vec::with_capacity(num_layers); + let mut hot_kv_store = rs.hot_kv.take(); + let had_hot_kv = hot_kv_store.is_some(); + let idx_kv: Option<&dyn larql_compute::KvIndex> = + index.map(|v| v as &dyn larql_compute::KvIndex); + + for layer in 0..num_layers { + new_stored.push(h_new.clone()); + + let hot_kv = match (cache_eligible && had_hot_kv, hot_kv_store.as_mut()) { + (true, Some(bufs)) => HotKv::InPlace(bufs), + _ => HotKv::Recompute, + }; + let h_post_attn = resolve_layer_attention( + weights, + rs, + layer, + &h_new, + abs_position, + backend, + idx_kv, + hot_kv, + &mut step_new_kv, + cache_eligible, + ) + .ok_or_else(|| EngineError::BackendFailure { + details: format!("attention returned None during codec decode at layer {layer}"), + })?; + + let bffn = BackendFfn { + weights: weights.canonical(), + backend, + }; + h_new = crate::engines::layer_ffn_or_moe( + weights.canonical(), + &h_post_attn, + layer, + &bffn, + moe_ffn, + ple_inputs.get(layer), + ) + .map_err(EngineError::Execution)?; + } + + commit( + weights.hidden_size, + rs, + &new_stored, + cache_eligible, + had_hot_kv, + hot_kv_store, + step_new_kv, + abs_position, + ); + Ok(last_row(&h_new)) +} diff --git a/crates/larql-kv/src/engines/markov_residual_codec/step/tests.rs b/crates/larql-kv/src/engines/markov_residual_codec/step/tests.rs new file mode 100644 index 000000000..56f2039d9 --- /dev/null +++ b/crates/larql-kv/src/engines/markov_residual_codec/step/tests.rs @@ -0,0 +1,172 @@ +//! Codec decode-step behaviour: positions, both cold-tier read paths, and +//! the in-place-vs-owned-concat parity gate. + +use super::rs_decode_step_codec; +use crate::engines::markov_residual_codec::codec::ColdResidualCodec; +use crate::engines::markov_residual_codec::prefill::rs_prefill_codec; +use crate::engines::markov_residual_codec::store::RsStoreCodec; +use larql_compute::CpuBackend; +use larql_inference::test_utils::make_test_weights; + +const PROMPT: [u32; 2] = [0, 1]; +const OVERFLOW_PROMPT: [u32; 4] = [0, 1, 2, 3]; +const WINDOW: usize = 2; + +fn prefill( + weights: &larql_inference::ModelWeights, + tokens: &[u32], + window: Option, +) -> RsStoreCodec { + rs_prefill_codec( + larql_inference::WeightsView::dense(weights), + tokens, + window, + ColdResidualCodec::Bf16, + &CpuBackend, + None, + ) + .expect("a dense prefill with no MoE hook cannot refuse") + .store +} + +fn decode( + weights: &larql_inference::ModelWeights, + rs: &mut RsStoreCodec, + token: u32, +) -> ndarray::Array2 { + rs_decode_step_codec( + larql_inference::WeightsView::dense(weights), + token, + rs, + &CpuBackend, + None, + None, + ) + .expect("decode") +} + +#[test] +fn decode_step_extends_position() { + let weights = make_test_weights(); + let mut rs = prefill(&weights, &PROMPT, None); + assert_eq!(rs.next_position, PROMPT.len()); + decode(&weights, &mut rs, 2); + assert_eq!(rs.next_position, PROMPT.len() + 1); +} + +#[test] +fn decode_with_cold_kv_path_produces_finite_output() { + let weights = make_test_weights(); + let mut rs = prefill(&weights, &OVERFLOW_PROMPT, Some(WINDOW)); + assert!(rs.cold_kv.is_some()); + let h = decode(&weights, &mut rs, 4); + assert_eq!(h.shape(), &[1, weights.hidden_size]); + assert!(h.iter().all(|v| v.is_finite())); +} + +#[test] +fn decode_with_cold_encoded_path_produces_finite_output() { + // After enough decode steps, the post-eviction cold_kv-clear path is + // exercised (we read from cold_encoded directly via decode). + let weights = make_test_weights(); + let mut rs = prefill(&weights, &OVERFLOW_PROMPT, Some(WINDOW)); + decode(&weights, &mut rs, 4); + // Second decode: cold_kv was cleared by overflow at the first decode, + // so this step exercises the cold_encoded recompute branch. + let h = decode(&weights, &mut rs, 5); + assert_eq!(h.shape(), &[1, weights.hidden_size]); + assert!(h.iter().all(|v| v.is_finite())); +} + +/// Flags-ON parity gate for the codec engine's in-place hot-K/V fast path: +/// an A/B of the in-place steady state against the owned-concat reference, +/// both with Q4K-direct attention live. Twin of the markov test — the two +/// paths must produce bit-identical hidden states at every step. q4k flags +/// are driven via the thread-local override (no env race), and the +/// in-place path is selected through the shared `LARQL_MARKOV_INPLACE_KV` +/// thread-local override. +#[test] +fn rs_decode_step_codec_inplace_matches_owned_concat_flags_on() { + use crate::engines::markov_residual::compute::set_markov_env_override; + use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; + + const PARITY_PROMPT: [u32; 3] = [0, 1, 2]; + const FIRST_DECODE_TOKEN: u32 = 3; + const LAST_DECODE_TOKEN: u32 = 12; + + let _q4k = crate::engines::Q4kFlagGuard::set(&[ + (larql_compute::options::ENV_Q4K_DIRECT_ATTN, true), + (larql_compute::options::ENV_Q4K_ATTN_INT8, false), + ]); + + let weights = make_test_q4k_weights(); + let index = make_test_q4k_vindex(&weights); + + let run = |inplace: bool| -> (Vec>, usize) { + set_markov_env_override( + "LARQL_MARKOV_INPLACE_KV", + Some(if inplace { "1" } else { "0" }), + ); + let mut rs = prefill(&weights, &PARITY_PROMPT, None); + let mut hiddens = Vec::new(); + for tok in FIRST_DECODE_TOKEN..=LAST_DECODE_TOKEN { + let h = rs_decode_step_codec( + larql_inference::WeightsView::dense(&weights), + tok, + &mut rs, + &CpuBackend, + None, + Some(&index), + ) + .expect("decode"); + assert!(h.iter().all(|v| v.is_finite())); + hiddens.push(h.iter().map(|v| v.to_bits()).collect()); + } + (hiddens, rs.hot_len) + }; + + let (a_hiddens, a_len) = run(true); + let (b_hiddens, b_len) = run(false); + let decoded = (LAST_DECODE_TOKEN - FIRST_DECODE_TOKEN + 1) as usize; + assert_eq!(a_len, PARITY_PROMPT.len() + decoded, "prompt + decode rows"); + assert_eq!(a_len, b_len); + assert_eq!( + a_hiddens, b_hiddens, + "codec in-place and owned-concat hidden states diverged (q4k-direct on)" + ); +} + +/// The first overflow to happen *during decode* creates the cold tier. +/// +/// A prompt that exactly fills the window leaves prefill with no cold tier at +/// all, so the next decode is the first thing to evict — and `commit` has to +/// build the encoded layers rather than append to layers that do not exist. +/// The sibling path (appending to an existing tier) is what every other +/// windowed test here exercises, so without this one the constructing branch +/// never runs. +#[test] +fn a_first_decode_overflow_constructs_the_encoded_cold_tier() { + let weights = make_test_weights(); + // Prompt length == window: full, but nothing evicted yet. + let mut rs = prefill(&weights, &PROMPT, Some(PROMPT.len())); + assert!( + rs.cold_encoded.is_none(), + "a prompt that exactly fills the window must not evict" + ); + + let h = decode(&weights, &mut rs, 2); + assert_eq!(h.shape(), &[1, weights.hidden_size]); + + let layers = rs + .cold_encoded + .as_ref() + .expect("the first decode overflow must construct the cold tier"); + assert_eq!(layers.len(), weights.num_layers); + for layer in layers { + assert_eq!(layer.n_positions, 1, "exactly one row was evicted"); + } + assert!( + rs.cold_kv.is_none(), + "a lossy codec must invalidate the cached cold K/V" + ); +} diff --git a/crates/larql-kv/src/engines/markov_residual_codec/step_attention.rs b/crates/larql-kv/src/engines/markov_residual_codec/step_attention.rs new file mode 100644 index 000000000..e31345c06 --- /dev/null +++ b/crates/larql-kv/src/engines/markov_residual_codec/step_attention.rs @@ -0,0 +1,214 @@ +//! Resolving one layer's attention output from the codec store's tiers. +//! +//! Twin of [`crate::engines::markov_residual::step_attention`], differing only +//! in the cold tier: this one decodes `cold_encoded` where the sibling slices +//! `cold_residuals`. + +use larql_compute::ComputeBackend; +use larql_inference::attention::SharedKV; +use ndarray::{s, Array2}; + +use crate::engines::markov_residual::recompute_kv; +use crate::engines::markov_residual_codec::store::RsStoreCodec; + +/// The hot-K/V cache's state for this step, which is what decides how +/// attention gets its prior. +pub(super) enum HotKv<'a> { + /// Steady state (step 2+) on an unbounded window: the buffers hold the + /// full prior K/V and this layer's row is appended into them in place. + InPlace(&'a mut Vec), + /// First step, or a windowed/cold configuration: the prior is recomputed + /// from the canonical residuals. + Recompute, +} + +/// Derive this layer's `h_post_attn`. +/// +/// Returns `None` when the backend declines. Every mutation this makes is to +/// `hot_kv`, a droppable derivative of `rs.stored`, never to canonical state. +#[allow(clippy::too_many_arguments)] +pub(super) fn resolve_layer_attention( + weights: larql_inference::WeightsView, + rs: &RsStoreCodec, + layer: usize, + h_new: &Array2, + abs_position: usize, + backend: &dyn ComputeBackend, + idx_kv: Option<&dyn larql_compute::KvIndex>, + hot_kv: HotKv<'_>, + step_new_kv: &mut Vec, + cache_eligible: bool, +) -> Option> { + // `stored` is a doubling-capacity buffer (W8.2): logical row count is + // `hot_len`, not `shape()[0]`. + let s_hot = rs.hot_len; + let hot_abs_start = abs_position.saturating_sub(s_hot); + + match hot_kv { + HotKv::InPlace(bufs) => { + #[cfg(debug_assertions)] + debug_assert_hot_kv_parity(weights, rs, layer, bufs, s_hot, hot_abs_start, backend); + attend_in_place( + weights, + &mut bufs[layer], + layer, + h_new, + s_hot, + abs_position, + backend, + idx_kv, + ) + } + HotKv::Recompute => { + let kv_arg = recompute_prior_kv(weights, rs, layer, s_hot, hot_abs_start, backend)?; + let (h_post_attn, new_kv) = + larql_inference::attention::run_attention_block_decode_step_auto( + weights, + h_new, + layer, + Some(&kv_arg), + abs_position, + Some(backend), + idx_kv, + )?; + if cache_eligible { + step_new_kv.push(new_kv); + } + Some(h_post_attn) + } + } +} + +/// Steady state: append this token's projected+RoPE'd K/V row IN PLACE into +/// the doubling-capacity buffer and attend over the `[..s_hot+1]` views — no +/// per-step O(ctx) owned concat (O(L) total cache copy vs O(L²)). See the +/// markov twin for the rationale. +#[allow(clippy::too_many_arguments)] +fn attend_in_place( + weights: larql_inference::WeightsView, + buf: &mut SharedKV, + layer: usize, + h_new: &Array2, + s_hot: usize, + abs_position: usize, + backend: &dyn ComputeBackend, + idx_kv: Option<&dyn larql_compute::KvIndex>, +) -> Option> { + let (k_buf, v_buf) = buf; + let inplace = if crate::engines::markov_residual::compute::markov_inplace_kv_enabled() { + larql_inference::attention::run_attention_block_decode_step_auto_inplace( + weights, + h_new, + layer, + k_buf, + v_buf, + s_hot, + abs_position, + Some(backend), + idx_kv, + ) + } else { + None + }; + match inplace { + Some(h) => Some(h), + None => { + // Q4K-direct off (flags-off parity) or no attn bytes: owned concat + // over the buffer view, then replace. Bit-identical to the legacy + // borrow path. + let prior: SharedKV = ( + k_buf.slice(s![..s_hot, ..]).to_owned(), + v_buf.slice(s![..s_hot, ..]).to_owned(), + ); + let (h, new_kv) = larql_inference::attention::run_attention_block_decode_step_auto( + weights, + h_new, + layer, + Some(&prior), + abs_position, + Some(backend), + idx_kv, + )?; + *k_buf = new_kv.0; + *v_buf = new_kv.1; + Some(h) + } + } +} + +/// First step (cache `None` → seed) or windowed/cold tier. +fn recompute_prior_kv( + weights: larql_inference::WeightsView, + rs: &RsStoreCodec, + layer: usize, + s_hot: usize, + hot_abs_start: usize, + backend: &dyn ComputeBackend, +) -> Option { + let h_hot = &rs.stored[layer]; + if let Some(cold_kv) = &rs.cold_kv { + let (k_cold, v_cold) = &cold_kv[layer]; + let (k_hot, v_hot) = recompute_kv(weights, h_hot, layer, hot_abs_start, backend, None)?; + let c = k_cold.shape()[0]; + let kv_dim = k_cold.shape()[1]; + let mut k_combined = Array2::::zeros((c + s_hot, kv_dim)); + k_combined.slice_mut(s![..c, ..]).assign(k_cold); + k_combined.slice_mut(s![c.., ..]).assign(&k_hot); + let mut v_combined = Array2::::zeros((c + s_hot, kv_dim)); + v_combined.slice_mut(s![..c, ..]).assign(v_cold); + v_combined.slice_mut(s![c.., ..]).assign(&v_hot); + return Some((k_combined, v_combined)); + } + + let (h_full, full_abs_start) = match &rs.cold_encoded { + Some(cold_layers) if cold_layers[layer].n_positions > 0 => { + let decoded = cold_layers[layer].decode(rs.codec); + let n_cold = decoded.shape()[0]; + let hidden = h_hot.shape()[1]; + let mut combined = Array2::::zeros((n_cold + s_hot, hidden)); + combined.slice_mut(s![..n_cold, ..]).assign(&decoded); + combined.slice_mut(s![n_cold.., ..]).assign(h_hot); + (combined, rs.cold_abs_start) + } + _ => (h_hot.clone(), hot_abs_start), + }; + recompute_kv(weights, &h_full, layer, full_abs_start, backend, None) +} + +/// f32-path parity gate only — the Q4K-direct route has its own oracles (the +/// compute-level bit-identity test + the engine A/B). +#[cfg(debug_assertions)] +fn debug_assert_hot_kv_parity( + weights: larql_inference::WeightsView, + rs: &RsStoreCodec, + layer: usize, + bufs: &[SharedKV], + s_hot: usize, + hot_abs_start: usize, + backend: &dyn ComputeBackend, +) { + /// Largest per-element f32 gap tolerated between the cached prior K/V and + /// a fresh recompute. + const MAX_CACHE_DRIFT: f32 = 1e-2; + + if larql_compute::options::q4k_direct_attn_enabled() { + return; + } + let (k_buf, v_buf) = &bufs[layer]; + let h_logical = rs.stored[layer].slice(s![..s_hot, ..]).to_owned(); + let Some((rk, rv)) = recompute_kv(weights, &h_logical, layer, hot_abs_start, backend, None) + else { + return; + }; + let max_gap = |buf: &Array2, fresh: &Array2| { + buf.slice(s![..s_hot, ..]) + .iter() + .zip(fresh.iter()) + .map(|(a, b)| (a - b).abs()) + .fold(0.0f32, f32::max) + }; + let kd = max_gap(k_buf, &rk); + let vd = max_gap(v_buf, &rv); + debug_assert!(kd < MAX_CACHE_DRIFT, "codec hot_kv K cache diverged: {kd}"); + debug_assert!(vd < MAX_CACHE_DRIFT, "codec hot_kv V cache diverged: {vd}"); +} diff --git a/crates/larql-kv/src/engines/markov_residual_codec/walk.rs b/crates/larql-kv/src/engines/markov_residual_codec/walk.rs index dcc34c730..6624bfeef 100644 --- a/crates/larql-kv/src/engines/markov_residual_codec/walk.rs +++ b/crates/larql-kv/src/engines/markov_residual_codec/walk.rs @@ -2,7 +2,7 @@ //! //! Mirrors `markov_residual/q4k.rs` with the cold tier routed through the //! codec. Used when the engine is asked to run on a compact (Q4K-walk) -//! vindex — the dense `BackendFfn` path in [`super::compute`] cannot read +//! vindex — the dense `BackendFfn` path in [`super::prefill`] cannot read //! `--compact` FFN weights. This module delegates FFN to `WalkFfn` //! (native Q4K matvec on the vindex's compact gate/up/down bytes) and //! passes `Some(index)` to `recompute_kv` so the K/V projections also @@ -16,7 +16,7 @@ use larql_inference::vindex::{WalkFfn, WalkFfnConfig}; use larql_vindex::VectorIndex; use ndarray::{s, Array2}; -use super::compute::RsPrefillResultCodec; +use super::prefill::RsPrefillResultCodec; use crate::engines::markov_residual::recompute_kv; use crate::engines::markov_residual_codec::codec::ColdResidualCodec; use crate::engines::markov_residual_codec::store::{EncodedColdLayer, RsStoreCodec}; diff --git a/crates/larql-kv/src/engines/mod.rs b/crates/larql-kv/src/engines/mod.rs index 1866412f7..431df1836 100644 --- a/crates/larql-kv/src/engines/mod.rs +++ b/crates/larql-kv/src/engines/mod.rs @@ -12,7 +12,7 @@ //! | [`standard`] | Production K/V tensor cache (default) | O(seq) f32 K/V | exact — the reference | //! | [`no_cache`] | Full re-forward per step | O(seq) token IDs | exact — correctness fallback | //! | [`markov_residual`] | Residual-stream replacement | ~171 MB | exact (KL=0.0) under contract | -//! | [`unlimited_context`] | Per-window K/V checkpoints | ~193 MB | exact within window | +//! | [`windowed_checkpoint`] | Per-window K/V checkpoints | ~193 MB | exact within window | //! | [`turbo_quant`] | WHT + Lloyd-Max 3/4-bit codec | ~12.7 GB | per-row cos≈0.9954 (4-bit, 2026-07-30) | //! | [`apollo`] | Boundary store + residual injection | ~11 MB | task accuracy | //! @@ -20,10 +20,10 @@ //! //! ```text //! larql bench gemma3-4b-q4k --engine standard -//! larql bench gemma3-4b-q4k --engine standard:window=1024 +//! larql bench gemma3-4b-q4k --engine standard:window=1024 # keeps the fused path //! larql bench gemma3-4b-q4k --engine no-cache //! larql bench gemma3-4b-q4k --engine markov-rs:window=512 -//! larql bench gemma3-4b-q4k --engine unlimited-context:window=256 +//! larql bench gemma3-4b-q4k --engine windowed-checkpoint:window=256 //! larql bench gemma3-4b-q4k --engine turbo-quant:bits=3 //! larql bench gemma3-4b-q4k --engine apollo:layer=25,coef=8.0 //! ``` @@ -32,11 +32,61 @@ //! //! ## Architecture notes //! -//! - **Metal Q4K path** (`prefill_quant` / `decode_step_quant`): all four engines -//! use the Metal `decode_token` full pipeline when a Q4K VectorIndex and a -//! Metal backend are available. This gives 93-95 tok/s — matching or exceeding -//! the standard larql-metal path (76 tok/s) because the engine bench uses -//! faster Metal lm_head KNN rather than a full vocab matmul. +//! - **Coarse (fused) path** (`prefill_quant` / `decode_step_quant`): taken +//! when a Q4K VectorIndex is present and the backend accepts the engine's +//! window. Metal routes through its `decode_token` full pipeline; CPU +//! through `cached_prefill_q4k` / `cached_decode_step_q4k`. +//! +//! The engine's own state policy is NOT engaged here — the K/V lives in the +//! backend behind a sentinel handle (Metal) or a whole-model handle (CPU), +//! so `standard`, `markov-rs`, `markov-rs-codec` and `boundary-per-layer` +//! execute the same kernels and measure within ~0.5% of each other. +//! **Ranking those four against one another on this path measures nothing.** +//! +//! The engine rows do beat the reference `larql-metal` row, but the reason +//! is the KNN lm_head (~1.9ms) replacing a full vocab matmul (~6.6ms), not +//! the K/V mechanism. An earlier note credited the engines with a ~93-95 vs +//! 76 tok/s win; that gap was inflated because the bench timed only the +//! forward and left lm_head outside the timer. Both are inside it now, and +//! the row note carries the `fwd=` / `head=` split. +//! +//! - **Windowed engines keep the coarse path** (since 2026-08). A window is +//! requested through `coarse_prefill_windowed` / `coarse_decode_step_windowed`, +//! which fail closed: a backend that cannot bound BOTH attention and K/V +//! answers `None` and the engine falls back to per-layer. Both `CpuBackend` +//! and `MetalBackend` implement them. +//! +//! Each bounds attention and memory by different means, because neither +//! mechanism does both jobs. CPU trims the cache to `w - 1` before each step +//! (cheap — the cache is host arrays). Metal clamps the attention span every +//! step via the layer window, and compacts its K/V buffers only once +//! occupancy reaches 2x the window, so the memmove is O(1) amortised rather +//! than O(window) per token; the surplus resident rows are never read +//! because the kernel attends `[T - window, T)`. +//! +//! Both decline a prompt LONGER than the window: the fused prefill has no +//! per-query-position masking, so accepting would attend the whole prompt +//! while the engine advertises a bound. That case still takes per-layer. +//! +//! Measured on Gemma 3 4B Q4K, Metal, 80 steps at `window=8`: 11.61ms / +//! 2.4MB against 12.06ms / 23.6MB unwindowed. Before this, the same config +//! cost 115.44ms. +//! +//! - **Per-layer path** (a prompt longer than the window, or an arch that +//! declines coarse): `MetalBackend`'s `KvDispatch` impl delegates every +//! per-layer method (`attention_step`, `attention_step_windowed`, +//! `append_kv`, `clip_kv`, …) to `CpuBackend`. An engine on this path runs +//! CPU attention **and** CPU FFN while the bench labels the row +//! `[metal (GPU)]` — worth ~9x on Gemma 3 4B. The row's dispatch note says +//! `[per-layer->host]` when that is what happened. +//! +//! - **Cross-backend parity**: Metal and CPU agree to ~5e-7 relative L2 on +//! prefill, and per-step decode differs by a stable ~3e-3 that does not +//! compound (two Q4K kernels rounding differently). Pinned by +//! `tests/gpu_engine_parity`. A divergence on the Gemma-3 arch that was +//! open through 2026-08 turned out to be a test fixture declaring +//! QK-norm without supplying the weights — real checkpoints always +//! carry them, so nothing shipped was affected. //! //! - **CPU fallback**: when Metal is unavailable, engines fall back to a CPU //! path using dequantised attention tensors (lazily inserted into the @@ -49,12 +99,17 @@ pub mod apollo; pub mod boundary_kv; pub mod boundary_per_layer; +mod layer_ffn; pub mod markov_residual; pub mod markov_residual_codec; pub mod no_cache; +pub mod no_expert_route; pub mod standard; pub mod turbo_quant; -pub mod unlimited_context; +pub mod windowed_checkpoint; + +pub(crate) use layer_ffn::{apply_ple_and_layer_scalar, layer_ffn_or_moe}; +pub(crate) use no_expert_route::refuse_if_moe; /// Whether W10 mask cascade is active. /// @@ -69,7 +124,7 @@ pub mod unlimited_context; /// it's now a no-op since the cascade is on by default. /// /// Used by the per-engine `dispatch.rs` modules -/// (markov_residual, markov_residual_codec, unlimited_context, +/// (markov_residual, markov_residual_codec, windowed_checkpoint, /// boundary_per_layer). Engines that treat K/V as canonical state /// (turbo_quant) don't call this — their dispatch path stays on /// Full mask regardless. @@ -85,70 +140,6 @@ pub(crate) fn w10_enabled() -> bool { } } -/// Post-FFN tail of the per-layer sequence: `apply_per_layer_embedding` -/// then `apply_layer_scalar`, in that order — mirroring the legacy -/// `kv_prefill_run` / `kv_decode_step_run` loops in -/// [`crate::generation`], the oracle for every engine forward. Both -/// steps are no-ops on archs without PLE / layer-scalar keys -/// (everything except Gemma 4 E-series), so threading this through -/// non-PLE paths costs one clone and changes no bits. -pub(crate) fn apply_ple_and_layer_scalar( - weights: &larql_inference::ModelWeights, - h_post_ffn: &ndarray::Array2, - layer: usize, - ple_input: Option<&ndarray::Array2>, -) -> ndarray::Array2 { - let mut h_out = larql_inference::forward::ple::apply_per_layer_embedding( - weights, h_post_ffn, layer, ple_input, - ); - larql_inference::forward::layer::apply_layer_scalar(weights, &mut h_out, layer); - h_out -} - -/// Per-layer FFN dispatch for engine forward loops, MoE-aware. -/// -/// On a hybrid-MoE arch, when a `moe_ffn` hook is supplied (e.g. -/// `RemoteMoeFfn` for `--moe-shards`), call its -/// [`FfnBackend::forward_moe_full_layer`] — it returns the full layer output -/// (dense `h1` + experts `h2` + combine), dispatching experts to the shards. -/// Otherwise fall back to the engine's own dense FFN (`dense_ffn`) followed -/// by [`apply_ple_and_layer_scalar`], the same per-layer sequence as the -/// legacy `kv_prefill_run` / `kv_decode_step_run` oracle. -/// -/// `ple_input` is this layer's entry from `precompute_per_layer_inputs` -/// (`None` on non-PLE archs, where PLE + layer_scalar are no-ops). -/// -/// Lets the per-layer / windowed engines (unlimited_context, markov_residual, -/// turbo_quant, …) ride remote MoE without touching their KV state policy — -/// only the FFN step changes. -pub(crate) fn layer_ffn_or_moe( - weights: &larql_inference::ModelWeights, - h_post_attn: &ndarray::Array2, - layer: usize, - dense_ffn: &dyn larql_inference::ffn::FfnBackend, - moe_ffn: Option<&dyn larql_inference::ffn::FfnBackend>, - ple_input: Option<&ndarray::Array2>, -) -> ndarray::Array2 { - if weights.arch.is_hybrid_moe() { - if let Some(mf) = moe_ffn { - if let Some(h_out) = mf.forward_moe_full_layer(layer, h_post_attn) { - // Returned as-is: `forward_moe_full_layer` is contracted to - // produce the FULL layer output. Every production impl routes - // through `moe_ffn_block_cpu(_with_index)`, which applies PLE - // + layer_scalar internally (`RemoteMoeFfn` / `LocalMoeFfn` - // directly; `LayerShardedRemote` and the ffn-policy router - // delegate to them; the HTTP walk backend requests - // `full_output` from the server). Applying either step again - // here would double-apply. - return h_out; - } - } - } - let (h_post_ffn, _) = - larql_inference::forward::run_ffn(weights, h_post_attn, layer, dense_ffn, false); - apply_ple_and_layer_scalar(weights, &h_post_ffn, layer, ple_input) -} - std::thread_local! { /// Per-thread override for [`w10_enabled`]. `Some(true)` simulates /// `LARQL_W10_DISABLE=1` (cascade off); `Some(false)` simulates the @@ -164,75 +155,6 @@ pub(crate) fn set_w10_disabled_override(disabled: Option) { W10_DISABLED_OVERRIDE.with(|o| *o.borrow_mut() = disabled); } -#[cfg(test)] -mod layer_ffn_or_moe_tests { - use super::layer_ffn_or_moe; - use larql_inference::ffn::FfnBackend; - use larql_inference::test_utils::make_test_gemma4_moe_weights; - use ndarray::Array2; - - /// FfnBackend whose MoE hook returns a sentinel (all 7.0) so we can tell - /// the MoE branch from the dense `run_ffn` fallback. - struct SentinelFfn; - impl FfnBackend for SentinelFfn { - fn forward(&self, _layer: usize, x: &Array2) -> Array2 { - Array2::zeros(x.raw_dim()) - } - fn name(&self) -> &str { - "sentinel" - } - fn forward_moe_full_layer( - &self, - _layer: usize, - h_post_attn: &Array2, - ) -> Option> { - Some(Array2::from_elem(h_post_attn.raw_dim(), 7.0)) - } - } - - #[test] - fn uses_moe_hook_on_hybrid_moe_arch() { - let weights = make_test_gemma4_moe_weights(); - assert!(weights.arch.is_hybrid_moe()); - let h = Array2::::zeros((2, weights.hidden_size)); - let out = layer_ffn_or_moe(&weights, &h, 0, &SentinelFfn, Some(&SentinelFfn), None); - // Took the MoE hook → sentinel output, not the dense run_ffn path. - assert!( - out.iter().all(|&v| v == 7.0), - "expected MoE-hook sentinel output" - ); - } - - #[test] - fn falls_back_to_dense_when_no_hook() { - let weights = make_test_gemma4_moe_weights(); - let h = Array2::::zeros((2, weights.hidden_size)); - // No moe_ffn → dense run_ffn even on a MoE arch (no experts dispatched). - let out = layer_ffn_or_moe(&weights, &h, 0, &SentinelFfn, None, None); - assert_eq!(out.shape(), &[2, weights.hidden_size]); - assert!( - out.iter().any(|&v| v != 7.0), - "must NOT be the MoE-hook sentinel" - ); - assert!(out.iter().all(|v| v.is_finite())); - } - - #[test] - fn sentinel_ffn_trait_surface() { - // Exercise the FfnBackend methods `layer_ffn_or_moe` doesn't call. - let s = SentinelFfn; - let x = Array2::::zeros((2, 4)); - assert_eq!(s.name(), "sentinel"); - assert_eq!(s.forward(0, &x).shape(), &[2, 4]); - let (o, obs) = s.forward_observed(0, &x); - assert_eq!(o.shape(), &[2, 4]); - assert!( - obs.is_absent(), - "sentinel stub must not fabricate activations" - ); - } -} - /// Test-only RAII helper to drive the Q4K decode fast-path flags via /// `larql_compute::options`' **thread-local** override (NOT `std::env::set_var`, /// which is thread-unsafe vs the concurrent `getenv` every parallel decode test diff --git a/crates/larql-kv/src/engines/no_cache.rs b/crates/larql-kv/src/engines/no_cache.rs index 68d5de83d..dbe974188 100644 --- a/crates/larql-kv/src/engines/no_cache.rs +++ b/crates/larql-kv/src/engines/no_cache.rs @@ -73,8 +73,10 @@ impl KvEngine for NoCacheEngine { if token_ids.is_empty() { return Err(EngineError::EmptyPrompt); } - self.tokens = token_ids.to_vec(); let view = larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch); + // Assigned after the `?`: the token list is this engine's entire + // state, so a refused prefill must leave whatever list was already + // there rather than adopt a prompt it never finished. let (hidden, _cache) = kv_prefill_run( view, ffn, @@ -82,10 +84,8 @@ impl KvEngine for NoCacheEngine { None, Some(self.backend.as_ref()), &mut NoopHook, - ) - .ok_or_else(|| EngineError::BackendFailure { - details: "kv_prefill_run returned None".into(), - })?; + )?; + self.tokens = token_ids.to_vec(); Ok(hidden) } @@ -95,19 +95,28 @@ impl KvEngine for NoCacheEngine { ffn: &dyn FfnBackend, token_id: u32, ) -> Result, EngineError> { + // **Transactional.** The token list is the whole of this engine's + // continuation state, and the re-forward needs the new token *in* it — + // so the push happens first and is undone if the step does not + // complete. Without the pop, a caller who fixed a refusal's cause and + // retried would re-forward the same token twice, which is precisely + // the failure the strict-refusal contract exists to prevent. There is + // no K/V to rewind: `kv_prefill_run`'s cache is discarded every call, + // which is what makes this engine the correctness fallback. self.tokens.push(token_id); let view = larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch); - let (hidden, _cache) = kv_prefill_run( + let outcome = kv_prefill_run( view, ffn, &self.tokens, None, Some(self.backend.as_ref()), &mut NoopHook, - ) - .ok_or_else(|| EngineError::BackendFailure { - details: "kv_prefill_run returned None during decode_step".into(), - })?; + ); + if outcome.is_err() { + self.tokens.pop(); + } + let (hidden, _cache) = outcome?; Ok(hidden) } diff --git a/crates/larql-kv/src/engines/no_expert_route.rs b/crates/larql-kv/src/engines/no_expert_route.rs new file mode 100644 index 000000000..2d8c26dc1 --- /dev/null +++ b/crates/larql-kv/src/engines/no_expert_route.rs @@ -0,0 +1,114 @@ +//! The refusal an engine owes a MoE architecture it structurally cannot serve. +//! +//! Distinct from every other refusal in the vocabulary, which describes an +//! operand that is elsewhere or a binding that is wrong. This one describes +//! the *executor*: a forward path with no expert-dispatch seam anywhere in it, +//! asked to run a model whose weights declare routed expert operands. +//! +//! Before this existed, such an engine ran the dense half of every layer and +//! returned an apparently valid result. That is a worse failure than a +//! degraded one — a missing route member at least leaves a trace, whereas a +//! silently expert-free forward is a *different model* wearing the same +//! answer shape. Nothing downstream could tell. +//! +//! [`RefusalKind::Unsupported`] is the correct classification and not a +//! convenience: the operands are present and well-formed, no bound kernel +//! here serves them, and the response is to pick another executor. That is +//! precisely what the variant is for. + +use larql_execution::{ExecutionRefusal, RefusalKind}; +use larql_inference::kv_engine::EngineError; +use larql_inference::model::ModelWeights; + +/// An engine whose forward has no expert-dispatch seam, asked to serve a +/// hybrid-MoE architecture. +#[derive(Debug, Clone)] +pub struct NoExpertDispatchPath { + /// The engine that cannot serve it, by its `KvEngine::name`. + pub engine: &'static str, + /// Why it cannot — the structural reason, not a restatement of the kind. + pub because: &'static str, +} + +impl std::fmt::Display for NoExpertDispatchPath { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "engine '{}' has no expert-dispatch path ({}), and this architecture \ + declares routed expert operands — a dense-only forward would answer \ + for a different model", + self.engine, self.because + ) + } +} + +impl std::error::Error for NoExpertDispatchPath {} + +impl ExecutionRefusal for NoExpertDispatchPath { + fn kind(&self) -> RefusalKind { + RefusalKind::Unsupported + } +} + +/// Refuse up front when `weights` declares routed experts and this engine +/// cannot dispatch them. +/// +/// Called before any forward work, so the refusal costs nothing and mutates +/// nothing — which is also what makes it trivially retryable through a +/// different engine. +pub(crate) fn refuse_if_moe( + engine: &'static str, + because: &'static str, + weights: &ModelWeights, +) -> Result<(), EngineError> { + if !weights.arch.is_hybrid_moe() { + return Ok(()); + } + Err(EngineError::Execution(Box::new(NoExpertDispatchPath { + engine, + because, + }))) +} + +#[cfg(test)] +mod tests { + use super::{refuse_if_moe, NoExpertDispatchPath}; + use larql_execution::{ExecutionRefusal, RefusalKind}; + use larql_inference::test_utils::{make_test_gemma4_moe_weights, make_test_weights}; + + const ENGINE: &str = "test-engine"; + const BECAUSE: &str = "its forward builds its own dense FFN"; + + #[test] + fn a_dense_arch_passes_through_untouched() { + let weights = make_test_weights(); + assert!(!weights.arch.is_hybrid_moe()); + assert!(refuse_if_moe(ENGINE, BECAUSE, &weights).is_ok()); + } + + #[test] + fn a_moe_arch_refuses_as_unsupported() { + let weights = make_test_gemma4_moe_weights(); + let err = refuse_if_moe(ENGINE, BECAUSE, &weights) + .expect_err("a MoE arch must not be served by a dense-only forward"); + assert_eq!(err.refusal_kind(), Some(RefusalKind::Unsupported)); + // Unsupported is recoverable through another executor, and does not + // indict the artifact — the model is fine, this engine is not. + assert!(err.operation_is_recoverable()); + assert!(err.engine_state_is_retryable()); + } + + #[test] + fn the_message_names_the_engine_and_the_structural_reason() { + // An operator reading this needs to know which engine to stop using + // and why, not merely that something was unsupported. + let refusal = NoExpertDispatchPath { + engine: ENGINE, + because: BECAUSE, + }; + let rendered = refusal.to_string(); + assert!(rendered.contains(ENGINE), "{rendered}"); + assert!(rendered.contains(BECAUSE), "{rendered}"); + assert_eq!(refusal.kind(), RefusalKind::Unsupported); + } +} diff --git a/crates/larql-kv/src/engines/standard.rs b/crates/larql-kv/src/engines/standard.rs index 87138e971..94db3f9f2 100644 --- a/crates/larql-kv/src/engines/standard.rs +++ b/crates/larql-kv/src/engines/standard.rs @@ -95,6 +95,15 @@ pub struct StandardEngine { /// `WeightsView::with_scratch` over it). Empty on the dense path. Keeps /// `weights` immutable so the engine can hold `Arc`. dequant_scratch: larql_inference::DequantScratch, + /// Set when a failed decode step left K/V that could not be rewound, + /// carrying the rendered cause. While set, decode entry points refuse + /// with [`EngineError::InvariantViolation`] rather than compute from a + /// cache that describes no token sequence. + /// + /// Cleared by a successful prefill, which replaces the cache outright — + /// so re-prefilling is the documented way back, and the engine is never + /// permanently dead. + invalidated: Option, } impl StandardEngine { @@ -110,6 +119,7 @@ impl StandardEngine { abs_position: 0, backend: BackendSlot::Sync(backend), dequant_scratch: larql_inference::DequantScratch::new(), + invalidated: None, } } @@ -128,25 +138,48 @@ impl StandardEngine { abs_position: 0, backend: BackendSlot::Async(backend), dequant_scratch: larql_inference::DequantScratch::new(), + invalidated: None, } } fn cache_memory_bytes(&self) -> usize { - let Some(handles) = self.handles.as_ref() else { - return 0; - }; - handles - .iter() - .map(|h| { - // 2 × f32 per cached row (K + V), kv_dim wide. - h.cached_len() * h.kv_dim() * 2 * std::mem::size_of::() + // Two homes, never both: per-layer dispatch puts the K/V in the + // handles this engine owns; the coarse path hands back a sentinel + // handle and keeps the cache inside the backend. Summing both is + // safe — a backend that answers `backend_resident_kv_bytes` with + // a non-zero figure is by contract reporting K/V that no handle + // can see, so there is nothing to double-count. + let handle_bytes: usize = self + .handles + .as_ref() + .map(|handles| { + // `resident_bytes` (not the per-layer formula) so a + // whole-model handle reports all its layers, not one. + handles.iter().map(|h| h.resident_bytes()).sum() }) - .sum() + .unwrap_or(0); + handle_bytes + self.backend_resident_kv_bytes() + } + + /// K/V the backend holds internally (Metal's coarse pipeline). Zero + /// for the async slot: `AsyncComputeBackend` carries no `KvDispatch`, + /// and no async backend currently owns a cache of its own. + fn backend_resident_kv_bytes(&self) -> usize { + match &self.backend { + BackendSlot::Sync(b) => b.as_ref().backend_resident_kv_bytes(), + BackendSlot::Async(_) => 0, + } } /// Shared prefill body — both `prefill` (index=None) and /// `prefill_quant` (index=Some) route through here. Matches on the /// `BackendSlot` to pick sync vs async dispatch. + /// + /// This is the engine's policy point for the dispatch ring's three + /// outcomes: a refusal terminates as + /// [`EngineError::Execution`] and no hidden state is produced; a + /// declining backend stays [`EngineError::BackendFailure`], exactly + /// as the bare `None` it replaced. fn do_prefill( &mut self, weights: &ModelWeights, @@ -158,9 +191,10 @@ impl StandardEngine { let (hidden, handles) = match &self.backend { BackendSlot::Sync(b) => { kv_prefill_via_dispatch(b.as_ref(), view, ffn, token_ids, self.window_size, index) + .map_err(EngineError::Execution)? .ok_or_else(|| EngineError::BackendFailure { - details: "kv_prefill_via_dispatch returned None".into(), - })? + details: "kv_prefill_via_dispatch returned None".into(), + })? } BackendSlot::Async(b) => kv_prefill_via_dispatch_async( b.as_ref(), @@ -170,12 +204,16 @@ impl StandardEngine { self.window_size, index, ) + .map_err(EngineError::Execution)? .ok_or_else(|| EngineError::BackendFailure { details: "kv_prefill_via_dispatch_async returned None".into(), })?, }; self.handles = Some(handles); self.prefill_mode = Some(PrefillDispatchMode::PerLayer); + // A completed prefill replaces the cache outright, so whatever an + // earlier failed step left behind is gone with it. + self.invalidated = None; self.abs_position = token_ids.len(); Ok(hidden) } @@ -191,7 +229,7 @@ impl StandardEngine { weights: &ModelWeights, ffn: &dyn FfnBackend, initial_hidden: &Array2, - ) -> Option> { + ) -> Result, EngineError> { let view = larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch); // `token_ids: None` — MM hidden rows (vision/audio) have no 1:1 // token identities, so PLE inputs cannot be derived here. PLE @@ -206,7 +244,7 @@ impl StandardEngine { None, self.window_size, None, - )?, + ), BackendSlot::Async(b) => kv_prefill_from_hidden_via_dispatch_async( b.as_ref(), view, @@ -215,10 +253,19 @@ impl StandardEngine { None, self.window_size, None, - )?, - }; + ), + } + .map_err(EngineError::Execution)? + .ok_or_else(|| EngineError::BackendFailure { + details: "do_prefill_from_hidden returned None (empty hidden input or \ + backend dispatch failure)" + .into(), + })?; self.handles = Some(handles); self.prefill_mode = Some(PrefillDispatchMode::PerLayer); + // A completed prefill replaces the cache outright, so whatever an + // earlier failed step left behind is gone with it. + self.invalidated = None; // Critical: position pointer must be derived from the hidden // row count, NOT from any token count — the input may contain // vision rows that aren't tokens. Decode-loop correctness @@ -226,11 +273,29 @@ impl StandardEngine { // continuation. Pinned by the StandardEngine entry-point // agreement test in this file's tests module. self.abs_position = initial_hidden.nrows(); - Some(hidden) + Ok(hidden) } /// Shared decode-step body — both `decode_step` (index=None) and /// `decode_step_quant` (index=Some) route through here. + /// + /// Same policy point as [`Self::do_prefill`]: a refusal terminates as + /// [`EngineError::Execution`] and no hidden state escapes. + /// + /// **Transactional.** A decode step mutates the cache before it can + /// know whether it will finish: each layer's attention appends the new + /// token's K/V, and only then does the FFN get the chance to refuse. A + /// step that does not complete therefore rewinds every handle to the + /// length it had on entry, so a caller that fixes the refusal's cause + /// can drive the same token through the same engine. + /// + /// When the rewind cannot be trusted — see [`Self::rewind_is_sound`] — + /// the error is wrapped as [`EngineError::StateInvalidated`] and the + /// engine records it, so every later decode refuses instead of + /// computing from a cache that describes no token sequence. `prefill` + /// clears the condition, because it replaces the cache outright. + /// + /// `abs_position` advances only on success, on every path. fn do_decode_step( &mut self, weights: &ModelWeights, @@ -238,43 +303,105 @@ impl StandardEngine { token_id: u32, index: Option<&larql_inference::larql_vindex::VectorIndex>, ) -> Result, EngineError> { + if let Some(what) = &self.invalidated { + return Err(EngineError::InvariantViolation { + what: format!("decode_step called on an invalidated engine: {what}"), + }); + } let view = larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch); + let window = self.window_size; + let abs_position = self.abs_position; + let backend = &self.backend; let handles = self .handles .as_mut() .ok_or_else(|| EngineError::InvariantViolation { what: "decode_step called before prefill (handles missing)".into(), })?; - let hidden = match &self.backend { + + // Snapshot before the first append. Recorded per layer rather than + // assumed to be uniform: nothing in the trait promises every layer + // caches the same number of rows, and a wrong assumption here would + // rewind to a length no layer ever had. + let entry_lengths: Vec = handles.iter().map(|h| h.cached_len()).collect(); + let rewindable = Self::rewind_is_sound(window, &entry_lengths); + + let outcome = match backend { BackendSlot::Sync(b) => kv_decode_step_via_dispatch( b.as_ref(), view, ffn, handles, token_id, - self.abs_position, - self.window_size, + abs_position, + window, index, - ) - .ok_or_else(|| EngineError::BackendFailure { - details: "kv_decode_step_via_dispatch returned None".into(), - })?, + ), BackendSlot::Async(b) => kv_decode_step_via_dispatch_async( b.as_ref(), view, ffn, handles, token_id, - self.abs_position, - self.window_size, + abs_position, + window, index, - ) - .ok_or_else(|| EngineError::BackendFailure { - details: "kv_decode_step_via_dispatch_async returned None".into(), - })?, + ), }; - self.abs_position += 1; - Ok(hidden) + + let failure = match outcome { + Ok(Some(hidden)) => { + self.abs_position += 1; + return Ok(hidden); + } + // A declining backend leaves the same half-applied step a refusal + // does, so it gets the same treatment. Distinguishing them here + // would make the cache's integrity depend on which of two + // unrelated things went wrong. + Ok(None) => EngineError::BackendFailure { + details: "decode step via dispatch returned None".into(), + }, + Err(refusal) => EngineError::Execution(refusal), + }; + + let rewound = rewindable && Self::rewind(backend, handles, &entry_lengths); + if rewound { + return Err(failure); + } + let invalidated = failure.invalidating_engine_state(); + self.invalidated = Some(invalidated.to_string()); + Err(invalidated) + } + + /// Whether rewinding a failed decode step would restore the exact cache + /// the step started from. + /// + /// Unbounded caches only ever append, so truncating to the recorded + /// length is exact. A windowed cache is different: a step that reaches + /// the window drops its oldest row to make room, and that row is gone. + /// Row *count* cannot see this — append-then-drop leaves it unchanged — + /// so the only sound test is whether every layer had room to spare + /// before the step began. + fn rewind_is_sound(window: Option, entry_lengths: &[usize]) -> bool { + match window { + None => true, + // `len < w` is "the step's own append still fits", i.e. it will not + // push the layer past `w` and trigger the evicting clip. + Some(w) => entry_lengths.iter().all(|&len| len < w), + } + } + + /// Truncate every handle back to its recorded length. All-or-nothing: + /// one backend that cannot rewind makes the whole cache untrustworthy, + /// because the layers are only meaningful together. + fn rewind(backend: &BackendSlot, handles: &mut [KvHandle], entry_lengths: &[usize]) -> bool { + handles + .iter_mut() + .zip(entry_lengths) + .all(|(handle, &len)| match backend { + BackendSlot::Sync(b) => b.as_ref().truncate_kv(handle, len), + BackendSlot::Async(b) => b.as_ref().truncate_kv(handle, len), + }) } } @@ -330,11 +457,6 @@ impl KvEngine for StandardEngine { initial_hidden: &Array2, ) -> Result, EngineError> { self.do_prefill_from_hidden(weights, ffn, initial_hidden) - .ok_or_else(|| EngineError::BackendFailure { - details: "do_prefill_from_hidden returned None (empty hidden input or \ - backend dispatch failure)" - .into(), - }) } fn decode_step( @@ -362,24 +484,34 @@ impl KvEngine for StandardEngine { // 3 4B vs ~0.4 tok/s through per-layer dispatch). Quant-agnostic: // the backend inspects `index` to pick the right kernel. // - // WINDOW GATE: the coarse trait surface (`coarse_prefill` / - // `coarse_decode_step`) has no window parameter, so the coarse - // path always attends over the FULL context. A windowed engine - // (`window_size: Some(N)`, the `markov-bounded` flag combo) - // must therefore decline coarse and take the per-layer path, - // whose dispatch enforces the window via `clip_kv` — otherwise - // the same CLI flag gives windowed behaviour on one backend and - // full-context on another while `info()` reports `window=N`. - // Correctness over speed; threading a window through the coarse - // trait (5 methods × CPU + Metal impls + fused kernels) is the - // eventual fast path if windowed quant becomes hot. - let coarse = if self.window_size.is_none() { - match &self.backend { - BackendSlot::Sync(b) => b.as_ref().coarse_prefill(weights, token_ids, Some(index)), - BackendSlot::Async(b) => b.as_ref().coarse_prefill(weights, token_ids, Some(index)), - } - } else { - None + // WINDOW: ask the backend to honour this engine's window on the + // fused path. `coarse_*_windowed` fails closed — a backend that + // cannot bound BOTH attention and K/V to `window_size` answers + // `None`, and we fall through to the per-layer path below, which + // enforces the window via `clip_kv`. + // + // That decline is the current answer on every backend, so this + // is behaviour-preserving today. It replaces a blanket + // `window_size.is_none()` gate: the reason a windowed engine + // leaves the fast path is now the backend's own capability + // answer rather than a rule stated here, so a backend that + // implements the window starts getting the fused path without + // this engine changing. That matters — on a host-delegating + // backend the per-layer route runs the whole forward on the CPU + // and costs ~2.4x, while the window makes attention *cheaper*. + let coarse = match &self.backend { + BackendSlot::Sync(b) => b.as_ref().coarse_prefill_windowed( + weights, + token_ids, + Some(index), + self.window_size, + ), + BackendSlot::Async(b) => b.as_ref().coarse_prefill_windowed( + weights, + token_ids, + Some(index), + self.window_size, + ), }; if let Some((hidden, handle)) = coarse { // Store as a single-element handles vec — the `KvHandle` @@ -387,6 +519,8 @@ impl KvEngine for StandardEngine { self.handles = Some(vec![handle]); self.prefill_mode = Some(PrefillDispatchMode::Coarse); self.abs_position = token_ids.len(); + // Same reasoning as the per-layer prefills: the cache is new. + self.invalidated = None; return Ok(hidden); } // Backend doesn't have a coarse path (e.g. f32 model, or @@ -433,6 +567,14 @@ impl KvEngine for StandardEngine { .ok_or_else(|| EngineError::InvariantViolation { what: "decode_step called before prefill (handles missing)".into(), })?; + // The coarse branch below bypasses `do_decode_step`, so it needs the + // invalidation guard in its own right — a guard that only covers the + // path that *sets* the flag protects the wrong half. + if let Some(what) = &self.invalidated { + return Err(EngineError::InvariantViolation { + what: format!("decode_step_quant called on an invalidated engine: {what}"), + }); + } if mode == PrefillDispatchMode::Coarse { let handles = self .handles @@ -443,20 +585,25 @@ impl KvEngine for StandardEngine { // Invariant: Coarse mode stores exactly one whole-model // handle (set together with the mode in `prefill_quant`). let handle = &mut handles[0]; + // Windowed variant, matching the prefill that minted this + // handle: a sequence that reached the fused path under a + // window must keep decoding under it. let coarse = match &self.backend { - BackendSlot::Sync(b) => b.as_ref().coarse_decode_step( + BackendSlot::Sync(b) => b.as_ref().coarse_decode_step_windowed( weights, token_id, Some(index), handle, self.abs_position, + self.window_size, ), - BackendSlot::Async(b) => b.as_ref().coarse_decode_step( + BackendSlot::Async(b) => b.as_ref().coarse_decode_step_windowed( weights, token_id, Some(index), handle, self.abs_position, + self.window_size, ), }; return match coarse { @@ -545,6 +692,16 @@ impl KvEngine for StandardEngine { // entirely; nothing is moved to cold. 0 } + + fn dispatch_path(&self) -> Option { + use larql_inference::kv_engine::DispatchPath; + // `prefill_mode` is the authority (see its declaration): handle + // count cannot distinguish a coarse handle from a 1-layer model. + self.prefill_mode.map(|mode| match mode { + PrefillDispatchMode::Coarse => DispatchPath::Coarse, + PrefillDispatchMode::PerLayer => DispatchPath::PerLayer, + }) + } } #[cfg(test)] @@ -554,6 +711,42 @@ mod tests { use larql_inference::forward::hidden_to_raw_logits; use larql_inference::test_utils::make_test_weights; + // ── Rewind soundness ──────────────────────────────────────────────── + // + // The predicate that decides whether a failed decode step can be undone. + // Tested directly because the interesting cases are boundary conditions on + // the window, and driving each of them through a real refusal would need a + // refusing route per case while proving the same three facts. + + #[test] + fn an_unbounded_cache_is_always_rewindable() { + // Unbounded caches only append, so truncating to the recorded length + // restores the exact prior state whatever the lengths were. + assert!(StandardEngine::rewind_is_sound(None, &[0, 7, 4096])); + assert!(StandardEngine::rewind_is_sound(None, &[])); + } + + #[test] + fn a_windowed_cache_is_rewindable_only_with_room_to_spare() { + const W: usize = 4; + // Every layer strictly below the window: the step's own append fits, + // so no eviction fires and the truncate is exact. + assert!(StandardEngine::rewind_is_sound(Some(W), &[0, 1, 3])); + // A layer *at* the window evicts to make room, and the evicted row is + // gone — length would come back while the contents had shifted. + assert!(!StandardEngine::rewind_is_sound(Some(W), &[3, 4])); + assert!(!StandardEngine::rewind_is_sound(Some(W), &[W])); + } + + #[test] + fn one_unrewindable_layer_condemns_the_whole_step() { + // All-or-nothing: the layers are only meaningful together, so a single + // layer at the window makes the cache untrustworthy even if every + // other layer had room. + const W: usize = 8; + assert!(!StandardEngine::rewind_is_sound(Some(W), &[0, 0, 0, W, 0])); + } + #[test] fn engine_name() { assert_eq!(StandardEngine::new(None).name(), "standard"); @@ -1096,6 +1289,73 @@ mod tests { // per-layer dequant fallback is exercised by the windowed tests // further down (windowed engines decline coarse by design). + // ── dispatch_path reporting ─────────────────────────────────────────── + + #[test] + fn dispatch_path_is_none_before_prefill() { + // Nothing has chosen a shape yet — reporting one would be a guess. + assert_eq!(StandardEngine::new(None).dispatch_path(), None); + assert_eq!(StandardEngine::new(Some(4)).dispatch_path(), None); + } + + #[test] + fn dispatch_path_reports_coarse_when_the_backend_took_the_fused_path() { + use larql_inference::ffn::NullFfn; + use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; + let weights = make_test_q4k_weights(); + let index = make_test_q4k_vindex(&weights); + let backend = larql_compute::cpu_backend(); + let mut engine = StandardEngine::new(None); + engine + .prefill_quant(&weights, &NullFfn, &index, &[0u32, 1, 2], &*backend) + .expect("prefill_quant"); + assert_eq!( + engine.dispatch_path(), + Some(larql_inference::kv_engine::DispatchPath::Coarse), + "unwindowed Q4K prefill takes the coarse path on CpuBackend" + ); + } + + /// The window gate's observable consequence. A windowed engine MUST + /// decline coarse (the coarse surface has no window parameter, so + /// taking it would silently attend the full context while `info()` + /// advertised `window=N`). Pinning the reported shape means a future + /// change that lets a windowed config onto the fused path fails here + /// rather than quietly returning full-context answers. + #[test] + fn windowed_engine_never_reports_coarse() { + use larql_inference::ffn::NullFfn; + use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; + let weights = make_test_q4k_weights(); + let index = make_test_q4k_vindex(&weights); + let backend = larql_compute::cpu_backend(); + let mut engine = StandardEngine::new(Some(2)); + engine + .prefill_quant(&weights, &NullFfn, &index, &[0u32, 1, 2], &*backend) + .expect("windowed prefill_quant"); + assert_eq!( + engine.dispatch_path(), + Some(larql_inference::kv_engine::DispatchPath::PerLayer), + "a windowed engine must decline the window-less coarse surface" + ); + } + + #[test] + fn dense_prefill_reports_per_layer() { + use larql_inference::ffn::WeightFfn; + use larql_inference::test_utils::make_test_weights; + let weights = make_test_weights(); + let mut engine = StandardEngine::new(None); + engine + .prefill(&weights, &WeightFfn { weights: &weights }, &[0u32, 1, 2]) + .expect("dense prefill"); + assert_eq!( + engine.dispatch_path(), + Some(larql_inference::kv_engine::DispatchPath::PerLayer), + "the dense (no-vindex) path is per-layer by construction" + ); + } + #[test] fn prefill_quant_cpu_fallback_runs_via_dequant() { use larql_inference::ffn::NullFfn; @@ -1447,13 +1707,24 @@ mod tests { // garbage on exactly the archs that decline coarse. Post-fix the // fallback substitutes a `WalkFfn` built from the vindex; this test // bit-compares against an explicitly-constructed `WalkFfn` driven - // through the same internals. A window larger than the whole run - // forces the per-layer path without any clipping effect. + // through the same internals. - /// Window larger than prompt + all decode steps — forces the - /// per-layer quant path (coarse is declined when windowed) while - /// keeping `clip_kv` a no-op, so this isolates the FFN routing. - const NO_CLIP_WINDOW: usize = 64; + /// Window NARROWER than the prompt, so the backend declines the + /// fused path and both engines run the per-layer quant walk this + /// test is about. + /// + /// It used to be a window *wider* than the whole run, on the premise + /// that "coarse is declined whenever windowed". That premise is gone: + /// a backend now accepts a window it can honour, and a window wider + /// than the prompt is trivially honourable — so engine A took the + /// fused path while the reference stayed per-layer and the two + /// legitimately disagreed. Narrower-than-prompt is the property that + /// actually forces per-layer now. + /// + /// The window then clips, but it clips BOTH engines identically + /// (same `window_size`, same `do_prefill` internals), so the FFN + /// routing this test isolates is unaffected. + const FORCE_PER_LAYER_WINDOW: usize = 2; #[test] fn quant_fallback_with_null_ffn_matches_explicit_walk_ffn_reference() { @@ -1467,13 +1738,22 @@ mod tests { // Engine A: public quant entry points with NullFfn. let ffn = NullFfn; - let mut engine_a = StandardEngine::new(Some(NO_CLIP_WINDOW)); + let mut engine_a = StandardEngine::new(Some(FORCE_PER_LAYER_WINDOW)); let h_a = engine_a .prefill_quant(&weights, &ffn, &index, &prompt, &*backend) .expect("engine A prefill_quant"); + // Pin the premise: this test is about the PER-LAYER fallback, so + // it is only meaningful if engine A actually took it. Without + // this, a future widening of what the fused path accepts would + // make the comparison silently test nothing. + assert_eq!( + engine_a.dispatch_path(), + Some(larql_inference::kv_engine::DispatchPath::PerLayer), + "engine A must be on the per-layer path for this comparison to mean anything" + ); // Engine B: same internals, explicit WalkFfn. - let mut engine_b = StandardEngine::new(Some(NO_CLIP_WINDOW)); + let mut engine_b = StandardEngine::new(Some(FORCE_PER_LAYER_WINDOW)); larql_inference::vindex::ensure_attn_tensors_dequantised( &mut engine_b.dequant_scratch, &weights, diff --git a/crates/larql-kv/src/engines/turbo_quant/engine.rs b/crates/larql-kv/src/engines/turbo_quant/engine.rs index 3cf909489..de1ab812e 100644 --- a/crates/larql-kv/src/engines/turbo_quant/engine.rs +++ b/crates/larql-kv/src/engines/turbo_quant/engine.rs @@ -224,6 +224,28 @@ impl CompressedLayer { self.num_vecs += 1; } + /// Drop rows until the layer holds `rows` of them again. + /// + /// Byte-exact, and that is not an accident of the codec being good: rows + /// are appended as whole head-chunks at fixed byte offsets and existing + /// bytes are never re-encoded (see [`Self::append_row`]), so removing the + /// tail restores precisely the buffer that preceded it. The compression is + /// lossy against its *input*, not against what was stored — which is what + /// lets a K/V-canonical engine rewind at all. + /// + /// No-op when `rows` is not smaller than the current count, so a caller + /// rewinding a layer the failure never reached costs nothing. + pub(super) fn truncate_rows(&mut self, rows: usize, tq: &TurboQuant) { + if rows >= self.num_vecs { + return; + } + let heads = self.kv_dim / self.head_dim.max(1); + let bytes = rows * heads * tq.bytes_per_vector(self.head_dim); + self.compressed_k.truncate(bytes); + self.compressed_v.truncate(bytes); + self.num_vecs = rows; + } + pub(super) fn memory_bytes(&self) -> usize { self.compressed_k.len() + self.compressed_v.len() } @@ -383,12 +405,46 @@ impl TurboQuantEngine { impl TurboQuantEngine { /// Shared body for `decode_step` / `decode_step_resident`. + /// + /// **Transactional.** Unlike the residual-canonical engines, this one's + /// canonical state *is* the K/V: each layer's compressed cache grows + /// before the FFN gets its chance to refuse, so a step that does not + /// finish must undo those appends rather than leave a cache holding a + /// token that produced no output. [`CompressedLayer::truncate_rows`] does + /// that byte-exactly, and `abs_position` advances only on success — so a + /// caller who fixes the cause can drive the same token again. fn decode_step_impl( &mut self, weights: &ModelWeights, ffn: &dyn FfnBackend, token_id: u32, index: Option<&larql_vindex::VectorIndex>, + ) -> Result, EngineError> { + // Recorded per layer rather than assumed uniform: nothing promises + // every layer caches the same number of rows, and a wrong assumption + // would rewind to a length no layer ever had. + let entry_rows: Vec = self.layers.iter().map(|l| l.num_vecs).collect(); + match self.decode_step_appending(weights, ffn, token_id, index) { + Ok(hidden) => Ok(hidden), + Err(failure) => { + for (layer, &rows) in self.layers.iter_mut().zip(&entry_rows) { + layer.truncate_rows(rows, &self.tq); + } + Err(failure) + } + } + } + + /// The body of a decode step, which appends to `self.layers` as it goes. + /// + /// Split out so the rewind above can wrap every exit rather than every + /// `?` having to remember it. + fn decode_step_appending( + &mut self, + weights: &ModelWeights, + ffn: &dyn FfnBackend, + token_id: u32, + index: Option<&larql_vindex::VectorIndex>, ) -> Result, EngineError> { let num_layers = weights.num_layers; let abs_position = self.abs_position; @@ -438,15 +494,15 @@ impl TurboQuantEngine { weights, backend: self.backend.as_ref(), }; - let h_out = crate::engines::layer_ffn_or_moe( + h = crate::engines::layer_ffn_or_moe( weights, &h_post_attn, layer, &bffn, Some(ffn), ple_inputs.get(layer), - ); - h = h_out; + ) + .map_err(EngineError::Execution)?; } self.abs_position += 1; @@ -488,7 +544,10 @@ impl KvEngine for TurboQuantEngine { let mut h = embed_tokens_pub(weights, token_ids); // Empty on non-PLE archs — `ple_inputs.get(layer)` then yields `None`. let ple_inputs = precompute_per_layer_inputs(weights, &h, token_ids); - self.layers.clear(); + // Built into a local, not into `self.layers`: a prefill that refuses + // partway must leave the engine holding whatever cache it already had + // rather than a truncated one for a prompt it never finished. + let mut layers: Vec = Vec::with_capacity(num_layers); for layer in 0..num_layers { let (h_post_attn, k, v) = run_attention_with_kv_backend( @@ -501,24 +560,24 @@ impl KvEngine for TurboQuantEngine { .ok_or_else(|| EngineError::BackendFailure { details: "run_attention_with_kv_backend returned None".into(), })?; - self.layers - .push(CompressedLayer::compress(&(k, v), &self.tq)); + layers.push(CompressedLayer::compress(&(k, v), &self.tq)); let bffn = BackendFfn { weights, backend: self.backend.as_ref(), }; - let h_out = crate::engines::layer_ffn_or_moe( + h = crate::engines::layer_ffn_or_moe( weights, &h_post_attn, layer, &bffn, Some(ffn), ple_inputs.get(layer), - ); - h = h_out; + ) + .map_err(EngineError::Execution)?; } + self.layers = layers; self.abs_position = token_ids.len(); Ok(last_row(&h)) } diff --git a/crates/larql-kv/src/engines/unlimited_context/checkpoint_store.rs b/crates/larql-kv/src/engines/windowed_checkpoint/checkpoint_store.rs similarity index 100% rename from crates/larql-kv/src/engines/unlimited_context/checkpoint_store.rs rename to crates/larql-kv/src/engines/windowed_checkpoint/checkpoint_store.rs diff --git a/crates/larql-kv/src/engines/unlimited_context/dispatch.rs b/crates/larql-kv/src/engines/windowed_checkpoint/dispatch.rs similarity index 76% rename from crates/larql-kv/src/engines/unlimited_context/dispatch.rs rename to crates/larql-kv/src/engines/windowed_checkpoint/dispatch.rs index 1a285b547..b644cb035 100644 --- a/crates/larql-kv/src/engines/unlimited_context/dispatch.rs +++ b/crates/larql-kv/src/engines/windowed_checkpoint/dispatch.rs @@ -1,4 +1,4 @@ -//! W1-GPU dispatch path for `UnlimitedContextEngine`. +//! W1-GPU dispatch path for `WindowedCheckpointEngine`. //! //! Routes prefill + decode through the backend's //! `coarse_prefill_with_state` / `coarse_decode_step_with_state_masked` @@ -30,9 +30,9 @@ use larql_inference::PerLayerDecodeState; use larql_vindex::VectorIndex; use ndarray::{s, Array2}; -use crate::engines::unlimited_context::engine::UnlimitedContextEngine; +use crate::engines::windowed_checkpoint::engine::WindowedCheckpointEngine; -impl UnlimitedContextEngine { +impl WindowedCheckpointEngine { /// W1-GPU step 4: prefill via `coarse_prefill_with_state`. The /// per-layer K/V dump is unpacked into pre-allocated /// `[window_size, kv_dim]` buffers so subsequent decode steps @@ -160,9 +160,45 @@ impl UnlimitedContextEngine { self.current_window_tokens = token_ids[open_start..].to_vec(); self.last_hidden = Some(hidden.clone()); self.kv_handle = Some(handle); + // The dump above ran over the whole prompt, so the handle holds every + // row. Drop everything the window does not entitle attention to see. + self.clip_handle_to_window(open_len, full_windows > 0); Some(hidden) } + /// Clip the backend K/V cache down to what this engine is allowed to + /// attend over: the open window, plus the boundary row a closed window + /// leaves behind as the next one's seed. + /// + /// **This is the window.** Everything else `window_size` touches — + /// segmenting the prompt, sizing the shadow buffers, deciding when to + /// archive — is bookkeeping the engine does to itself. Attention reads + /// the backend handle, so a handle that is never clipped means the window + /// bounds storage and not behaviour, and `info()`'s `window=N` is a claim + /// the engine does not honour (issue #200, measured in `larql/kvperf-1`: + /// the engine had the same decode cost slope as the unwindowed one across + /// four context lengths, and window sizes of 32 / 256 / 4096 were + /// indistinguishable). + /// + /// `clip_kv` keeps the *tail*, which is exactly right here: after a close + /// the row we must retain is the closing window's last position, and that + /// is the row the checkpoint was taken from. Clipping to 1 therefore + /// reproduces on the dispatch path what `extend_current` does on the + /// per-layer path when it seeds a fresh window from `checkpoints.load`. + fn clip_handle_to_window(&mut self, open_len: usize, any_window_closed: bool) { + let keep = if any_window_closed { + open_len + 1 + } else { + open_len + }; + let Some(handle) = self.kv_handle.as_mut() else { + return; + }; + if handle.cached_len() > keep { + self.backend.as_ref().clip_kv(handle, keep); + } + } + /// W1-GPU step 4: decode through dispatch. State capture gives us /// the new K/V row per layer; we append in-place to /// `current_window_kv` and trigger window auto-close when token @@ -253,6 +289,10 @@ impl UnlimitedContextEngine { // Window auto-close: same trigger as the legacy process loop. if self.current_window_tokens.len() >= self.window_size { self.close_window(); + // The closed window's rows are archived and checkpointed; only its + // last row may cross into the next window. Without this the handle + // grows for the whole stream and the window means nothing. + self.clip_handle_to_window(0, true); } Some(hidden) } @@ -260,7 +300,7 @@ impl UnlimitedContextEngine { #[cfg(test)] mod tests { - //! Coverage for the W1-GPU dispatch path. `UnlimitedContextEngine` + //! Coverage for the W1-GPU dispatch path. `WindowedCheckpointEngine` //! takes a non-optional `window_size: usize`; W10 mask cascade is //! gated on whether `current_window_kv` is dropped. Tests pin the //! cascade state via [`crate::engines::set_w10_disabled_override`] @@ -271,12 +311,12 @@ mod tests { use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; use super::*; - use crate::engines::unlimited_context::engine::UnlimitedContextEngine; + use crate::engines::windowed_checkpoint::engine::WindowedCheckpointEngine; - fn fixture(window_size: usize) -> (UnlimitedContextEngine, ModelWeights, VectorIndex) { + fn fixture(window_size: usize) -> (WindowedCheckpointEngine, ModelWeights, VectorIndex) { let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); - let engine = UnlimitedContextEngine::with_backend(window_size, cpu_engine_backend()); + let engine = WindowedCheckpointEngine::with_backend(window_size, cpu_engine_backend()); (engine, weights, index) } @@ -294,7 +334,7 @@ mod tests { weights.num_layers, weights.hidden_size, ); - let mut engine = UnlimitedContextEngine::with_backend(4, cpu_engine_backend()); + let mut engine = WindowedCheckpointEngine::with_backend(4, cpu_engine_backend()); let w = weights; assert!(engine .try_prefill_via_dispatch(&w, &empty_index, &[0u32, 1]) @@ -423,26 +463,36 @@ mod tests { /// The backend's kv row at an absolute position — ground truth for /// checkpoint value assertions. - fn backend_row( - engine: &UnlimitedContextEngine, - layer: usize, - abs_pos: usize, - ) -> (Vec, Vec) { + /// + /// Reads the **last** row the handle holds, not an absolute stream index: + /// since issue #200 the dispatch path clips the handle to the window, so + /// absolute positions no longer index it and the boundary row is simply + /// its tail. + fn backend_last_row(engine: &WindowedCheckpointEngine, layer: usize) -> (Vec, Vec) { let handle = engine.kv_handle.as_ref().expect("kv handle"); + let row = handle.cached_len().saturating_sub(1); engine .backend .as_ref() - .read_kv_row_at(handle, layer, abs_pos) + .read_kv_row_at(handle, layer, row) .expect("backend kv row readback") } - /// Regression (checkpoint indexing): the backend kv cache is - /// indexed by ABSOLUTE stream position. Pre-fix, `close_window` - /// under HOnly read the WINDOW-relative `n - 1`, so every window - /// after the first silently checkpointed the last row of the FIRST - /// window. + /// Each window checkpoints **its own** last row. + /// + /// The failure this guards has been reachable two different ways. First, + /// `close_window` under HOnly read a window-relative index into a handle + /// that spanned the whole stream, so every window after the first + /// re-checkpointed the *first* window's row. That was fixed by reading an + /// absolute position — correct only while the window was not enforced. + /// Now that the handle is clipped to the window (issue #200) the absolute + /// index runs off the end, and the boundary row is the handle's tail. + /// + /// Both readings share one observable claim, asserted here so neither + /// mechanism can regress: distinct windows must checkpoint distinct rows, + /// each recorded at its own absolute stream position. #[test] - fn h_only_checkpoints_read_absolute_last_row_of_each_window() { + fn h_only_checkpoints_each_windows_own_last_row() { set_w10_disable(false); const WINDOW: usize = 2; const CLOSED_WINDOWS: usize = 2; @@ -463,22 +513,13 @@ mod tests { let expected_abs_end = (window_id + 1) * WINDOW - 1; assert_eq!(abs_end, expected_abs_end, "window {window_id} abs position"); assert_eq!(ckpt.len(), weights.num_layers, "one K/V row per layer"); - for (layer, (k, v)) in ckpt.iter().enumerate() { - let (k_true, v_true) = backend_row(&engine, layer, expected_abs_end); - assert_eq!( - k.row(0).to_vec(), - k_true, - "window {window_id} layer {layer}: checkpoint K != \ - backend row at abs {expected_abs_end}" - ); - assert_eq!( - v.row(0).to_vec(), - v_true, - "window {window_id} layer {layer}: checkpoint V != \ - backend row at abs {expected_abs_end}" - ); - } } + // The handle is clipped, which is what makes the tail read correct. + let cached = engine.kv_handle.as_ref().expect("handle").cached_len(); + assert!( + cached <= window_bound(WINDOW), + "handle holds {cached} rows past a {WINDOW}-token window" + ); // The two boundary rows must actually differ or the index // assertions above would be vacuous. let (c0, _) = engine.checkpoints.load(0).unwrap(); @@ -580,7 +621,7 @@ mod tests { // BOOKKEEPING is compared — checkpoint values legitimately // diverge (full-history vs checkpoint+window attention; see the // module doc). - let mut cpu = UnlimitedContextEngine::new(WINDOW); + let mut cpu = WindowedCheckpointEngine::new(WINDOW); cpu.process(&weights, &PROMPT, None).expect("cpu prefill"); cpu.process(&weights, &[NEXT_TOKEN], None) .expect("cpu decode"); @@ -602,26 +643,17 @@ mod tests { let (_, abs_c) = cpu.checkpoints.load(window_id).expect("cpu ckpt"); assert_eq!(abs_d, abs_c, "window {window_id} checkpoint position"); } - // Value assertion for the dispatch checkpoints: each equals the - // backend cache's row at the window's absolute boundary. + // Every window carries a full per-layer checkpoint, and the most + // recently closed one matches what the (clipped) handle still holds. for window_id in 0..TOTAL_WINDOWS { - let (ckpt, abs_end) = engine.checkpoints.load(window_id).unwrap(); + let (ckpt, _) = engine.checkpoints.load(window_id).unwrap(); assert_eq!(ckpt.len(), weights.num_layers); - for (layer, (k, v)) in ckpt.iter().enumerate() { - let (k_true, v_true) = backend_row(&engine, layer, abs_end); - assert_eq!( - k.row(0).to_vec(), - k_true, - "window {window_id} layer {layer}: checkpoint K != \ - backend row at abs {abs_end}" - ); - assert_eq!( - v.row(0).to_vec(), - v_true, - "window {window_id} layer {layer}: checkpoint V != \ - backend row at abs {abs_end}" - ); - } + } + let (last_ckpt, _) = engine.checkpoints.load(TOTAL_WINDOWS - 1).unwrap(); + for (layer, (k, v)) in last_ckpt.iter().enumerate() { + let (k_true, v_true) = backend_last_row(&engine, layer); + assert_eq!(k.row(0).to_vec(), k_true, "layer {layer}: checkpoint K"); + assert_eq!(v.row(0).to_vec(), v_true, "layer {layer}: checkpoint V"); } } @@ -642,4 +674,72 @@ mod tests { .try_prefill_via_dispatch(&weights, &index, &[0u32, 1]) .is_none()); } + + // ── The window must bound attention, not merely storage ────────────── + // + // Regression pins for the defect measured in `kvperf-1` (issue #200): the + // coarse dispatch trait carries no window parameter, so a windowed engine + // that appends to a backend handle and never clips it attends over the + // whole stream while `info()` reports `window=N`. That is a correctness + // failure first — the engine is not the engine it claims to be — and it + // showed up as this engine having the *same* decode cost slope as the + // unwindowed one across four context lengths. + // + // Row count is the right thing to assert: it is what attention reads, and + // it is observable without timing. The bound is `window_size + 1` because + // a closed window leaves its final row behind as the next window's + // boundary checkpoint — which is the engine's whole design. + + /// Longest cache a correctly windowed engine may present to attention: + /// the open window, plus the boundary row carried in from the last close. + fn window_bound(window_size: usize) -> usize { + window_size + 1 + } + + #[test] + fn dispatch_prefill_clips_the_backend_cache_to_the_window() { + set_w10_disable(false); + const WINDOW: usize = 4; + const WINDOWS_WORTH: u32 = 5; + let (mut engine, weights, index) = fixture(WINDOW); + let prompt: Vec = (0..WINDOW as u32 * WINDOWS_WORTH).collect(); + + engine + .try_prefill_via_dispatch(&weights, &index, &prompt) + .expect("prefill"); + + let cached = engine.kv_handle.as_ref().expect("handle").cached_len(); + assert!( + cached <= window_bound(WINDOW), + "a {WINDOW}-token window prefilled with {} tokens left {cached} rows in the \ + backend cache — attention sees all of them, so the window bounds storage \ + but not what the engine actually attends over", + prompt.len() + ); + } + + #[test] + fn dispatch_decode_keeps_the_backend_cache_within_the_window() { + set_w10_disable(false); + const WINDOW: usize = 4; + const DECODE_STEPS: u32 = 12; + let (mut engine, weights, index) = fixture(WINDOW); + engine + .try_prefill_via_dispatch(&weights, &index, &[0u32, 1, 2]) + .expect("prefill"); + + // Enough steps to cross several window closes. + for tok in 3..3 + DECODE_STEPS { + engine + .decode_step_via_dispatch(&weights, &index, tok) + .expect("decode"); + let cached = engine.kv_handle.as_ref().expect("handle").cached_len(); + assert!( + cached <= window_bound(WINDOW), + "after token {tok} the backend cache holds {cached} rows, above the \ + {WINDOW}-token window — the cache grows without bound across window \ + closes" + ); + } + } } diff --git a/crates/larql-kv/src/engines/unlimited_context/engine.rs b/crates/larql-kv/src/engines/windowed_checkpoint/engine.rs similarity index 84% rename from crates/larql-kv/src/engines/unlimited_context/engine.rs rename to crates/larql-kv/src/engines/windowed_checkpoint/engine.rs index 65f77bea6..7317ad7ba 100644 --- a/crates/larql-kv/src/engines/unlimited_context/engine.rs +++ b/crates/larql-kv/src/engines/windowed_checkpoint/engine.rs @@ -1,4 +1,4 @@ -//! `UnlimitedContextEngine` — window-based KV cache with boundary-checkpoint replay. +//! `WindowedCheckpointEngine` — window-based KV cache with boundary-checkpoint replay. //! //! Window lifecycle: //! 1. `process(tokens)` — extends the active window's K,V via @@ -23,7 +23,7 @@ use serde::Serialize; use super::checkpoint_store::CheckpointStore; use super::extend::{ empty_prior, rs_extend_from_checkpoint_backend, rs_extend_from_checkpoint_quant, - rs_extend_inplace, + rs_extend_inplace, truncate_kv_rows, }; use super::token_archive::TokenArchive; use crate::engines::markov_residual::ensure_attn_tensors_dequantised; @@ -59,7 +59,7 @@ impl EngineStats { // ─── Engine ────────────────────────────────────────────────────────────────── -pub struct UnlimitedContextEngine { +pub struct WindowedCheckpointEngine { pub window_size: usize, pub checkpoints: CheckpointStore, pub archive: TokenArchive, @@ -106,7 +106,7 @@ pub struct UnlimitedContextEngine { /// rejected at construction. `window_size == 1` is legal. const MIN_WINDOW_SIZE: usize = 1; -impl UnlimitedContextEngine { +impl WindowedCheckpointEngine { /// # Panics /// /// Panics if `window_size < MIN_WINDOW_SIZE` (i.e. zero). @@ -120,7 +120,7 @@ impl UnlimitedContextEngine { pub fn with_backend(window_size: usize, backend: Box) -> Self { assert!( window_size >= MIN_WINDOW_SIZE, - "UnlimitedContextEngine window_size must be >= {MIN_WINDOW_SIZE}, got {window_size}" + "WindowedCheckpointEngine window_size must be >= {MIN_WINDOW_SIZE}, got {window_size}" ); Self { window_size, @@ -151,7 +151,7 @@ impl UnlimitedContextEngine { weights: &ModelWeights, tokens: &[u32], moe_ffn: Option<&dyn larql_inference::ffn::FfnBackend>, - ) -> Option<()> { + ) -> Result<(), EngineError> { self.process_with_index(weights, tokens, moe_ffn, None) } @@ -164,19 +164,32 @@ impl UnlimitedContextEngine { tokens: &[u32], moe_ffn: Option<&dyn larql_inference::ffn::FfnBackend>, index: Option<&larql_vindex::VectorIndex>, - ) -> Option<()> { + ) -> Result<(), EngineError> { let mut remaining = tokens; + // Closing a window archives its tokens and saves its checkpoint, and + // neither is undoable. `extend_current` rewinds the *current* window + // exactly, so a failure is retryable right up until the first close — + // after that the engine holds a stream it cannot complete, and must + // say so rather than let a caller retry into a duplicated window. + let mut closed_a_window = false; while !remaining.is_empty() { let free = self.window_size - self.current_window_tokens.len(); let take = remaining.len().min(free); let (chunk, rest) = remaining.split_at(take); - self.extend_current(weights, chunk, moe_ffn, index)?; + if let Err(failure) = self.extend_current(weights, chunk, moe_ffn, index) { + return Err(if closed_a_window { + failure.invalidating_engine_state() + } else { + failure + }); + } remaining = rest; if self.current_window_tokens.len() >= self.window_size { self.close_window(); + closed_a_window = true; } } - Some(()) + Ok(()) } /// Close any partial current window. Call before replay if the window hasn't filled. @@ -200,27 +213,38 @@ impl UnlimitedContextEngine { moe_ffn: Option<&dyn larql_inference::ffn::FfnBackend>, index: Option<&larql_vindex::VectorIndex>, window_id: usize, - ) -> Option<(Vec, usize)> { - let (tokens, abs_offset) = self.archive.retrieve(window_id)?; + ) -> Result<(Vec, usize), EngineError> { + let (tokens, abs_offset) = + self.archive + .retrieve(window_id) + .ok_or_else(|| EngineError::RetrievalMiss { + reason: format!("window {window_id} is not archived"), + })?; let prior = if window_id > 0 && self.checkpoints.contains(window_id - 1) { - let (ckpt, _) = self.checkpoints.load(window_id - 1)?; + let (ckpt, _) = + self.checkpoints + .load(window_id - 1) + .ok_or_else(|| EngineError::RetrievalMiss { + reason: format!("checkpoint for window {} is missing", window_id - 1), + })?; ckpt } else { empty_prior(weights) }; - let out = rs_extend_from_checkpoint_backend( + let mut kv_cache = prior; + rs_extend_from_checkpoint_backend( larql_inference::WeightsView::dense(weights), tokens, - prior, + &mut kv_cache, abs_offset, self.backend.as_ref(), moe_ffn, index, )?; let abs_end = abs_offset + tokens.len() - 1; - Some((out.kv_cache, abs_end)) + Ok((kv_cache, abs_end)) } /// Total storage and context statistics. @@ -349,16 +373,22 @@ impl UnlimitedContextEngine { chunk: &[u32], moe_ffn: Option<&dyn larql_inference::ffn::FfnBackend>, index: Option<&larql_vindex::VectorIndex>, - ) -> Option<()> { + ) -> Result<(), EngineError> { if chunk.is_empty() { - return Some(()); + return Ok(()); } // `prior_len` is the prior's LOGICAL row count — the window-KV counter // mid-window, the checkpoint's row count at a window start, or 0. let (mut prior, prior_len) = if self.current_window_tokens.is_empty() { if self.current_window_id > 0 && self.checkpoints.contains(self.current_window_id - 1) { - let (ckpt, _) = self.checkpoints.load(self.current_window_id - 1)?; + let id = self.current_window_id - 1; + let (ckpt, _) = + self.checkpoints + .load(id) + .ok_or_else(|| EngineError::RetrievalMiss { + reason: format!("checkpoint for window {id} is missing"), + })?; let len = ckpt.first().map_or(0, |(k, _)| k.shape()[0]); (ckpt, len) } else { @@ -366,7 +396,13 @@ impl UnlimitedContextEngine { } } else { // Mid-window the shadow MUST exist — see extend_current_quant. - (self.current_window_kv.take()?, self.current_window_kv_len) + let shadow = + self.current_window_kv + .take() + .ok_or_else(|| EngineError::InvariantViolation { + what: "mid-window extend with no K/V shadow".into(), + })?; + (shadow, self.current_window_kv_len) }; let abs_start = self.abs_offset + self.current_window_tokens.len(); @@ -384,8 +420,16 @@ impl UnlimitedContextEngine { && crate::engines::markov_residual::compute::markov_inplace_kv_enabled() && larql_compute::options::q4k_direct_attn_enabled(); - if use_inplace { - let last = rs_extend_inplace( + // Both arms restore the shadow on failure, which is what makes a + // refused chunk rewindable: `current_window_kv_len` and + // `current_window_tokens` are advanced only after the extend returns, + // so putting the buffers back at `prior_len` rows restores exactly the + // window this call started from. + let outcome = if use_inplace { + // The in-place path only ever writes past `prior_len`, which the + // counter never advanced past, so the logical window is already + // intact — nothing to truncate. + rs_extend_inplace( larql_inference::WeightsView::dense(weights), chunk, &mut prior, @@ -394,27 +438,43 @@ impl UnlimitedContextEngine { self.backend.as_ref(), moe_ffn, index, - )?; - self.last_hidden = Some(last); - self.current_window_kv_len = prior_len + chunk.len(); - self.current_window_kv = Some(prior); + ) + .map(|last| (last, prior_len + chunk.len())) } else { - let out = rs_extend_from_checkpoint_backend( + rs_extend_from_checkpoint_backend( larql_inference::WeightsView::dense(weights), chunk, - prior, + &mut prior, abs_start, self.backend.as_ref(), moe_ffn, index, - )?; - self.last_hidden = Some(out.last_hidden); - // CPU walk path: narrow arrays, counter == shape[0]. - self.current_window_kv_len = out.kv_cache.first().map_or(0, |(k, _)| k.shape()[0]); - self.current_window_kv = Some(out.kv_cache); - } + ) + .map(|step| { + // CPU walk path: narrow arrays, counter == shape[0]. + let rows = prior.first().map_or(0, |(k, _)| k.shape()[0]); + (step.last_hidden, rows) + }) + .inspect_err(|_| { + // The owned-concat path replaces each layer's buffer as it + // goes and reads a prior by `shape()[0]`, so a half-advanced + // cache would attend over a token whose step never finished. + truncate_kv_rows(&mut prior, prior_len); + }) + }; + + let (last_hidden, rows) = match outcome { + Ok(pair) => pair, + Err(failure) => { + self.current_window_kv = Some(prior); + return Err(failure); + } + }; + self.last_hidden = Some(last_hidden); + self.current_window_kv_len = rows; + self.current_window_kv = Some(prior); self.current_window_tokens.extend_from_slice(chunk); - Some(()) + Ok(()) } pub(super) fn close_window(&mut self) { @@ -425,10 +485,13 @@ impl UnlimitedContextEngine { // slice the engine-side shadow as before. let n = self.current_window_kv_len; let window_len = self.current_window_tokens.len(); - // Absolute stream position of this window's last token. The - // backend kv cache behind `kv_handle` is indexed by absolute - // position (the handle spans the whole stream since prefill), - // while the engine-side shadow holds only this window's rows. + // Absolute stream position of this window's last token — the value + // recorded *with* the checkpoint, so a later replay knows where it + // sat. It is no longer an index into anything: the dispatch path now + // clips the backend handle to the window (issue #200), so the handle + // holds this window's rows and not the stream's. The row to read back + // is therefore its last one. Indexing the handle by absolute position + // was correct only while the window was not being enforced. let abs_end = self.abs_offset + window_len - 1; let last_kv: Vec = match self.current_window_kv.take() { Some(kv) => { @@ -447,17 +510,16 @@ impl UnlimitedContextEngine { } } None => { - // No CPU shadow — engine ran under HOnly. Read the - // window's last K/V back from the backend's kv cache at - // the ABSOLUTE index `abs_end`; the window-relative - // `n - 1` would re-read a row of the first window on - // every window after it. If there is no handle or the - // backend lacks the readback affordance, fall through - // with an empty checkpoint: the tokens are still - // archived and the counters reset (a wedged window - // would otherwise spin `process()` forever), and the - // mismatched empty checkpoint surfaces as an extend - // error on the next window instead of silent loss. + // No CPU shadow — engine ran under HOnly. Read the window's + // last K/V back from the backend's kv cache. The handle is + // clipped to the window, so its final row *is* this window's + // last position; reading an absolute stream index here would + // now run off the end. If there is no handle or the backend + // lacks the readback affordance, fall through with an empty + // checkpoint: the tokens are still archived and the counters + // reset (a wedged window would otherwise spin `process()` + // forever), and the mismatched empty checkpoint surfaces as an + // extend error on the next window instead of silent loss. if n == 0 { Vec::new() } else if let Some(handle) = self.kv_handle.as_ref() { @@ -465,10 +527,14 @@ impl UnlimitedContextEngine { n, window_len, "HOnly window shadow counter out of sync with window tokens" ); + // Window-relative: the clipped handle's last row. + let last_row = handle.cached_len().saturating_sub(1); let mut rows = Vec::new(); let mut layer = 0; - while let Some((k_row, v_row)) = - self.backend.as_ref().read_kv_row_at(handle, layer, abs_end) + while let Some((k_row, v_row)) = self + .backend + .as_ref() + .read_kv_row_at(handle, layer, last_row) { let kv_dim = k_row.len(); let k = Array2::from_shape_vec((1, kv_dim), k_row) @@ -498,16 +564,16 @@ impl UnlimitedContextEngine { } } -impl KvEngine for UnlimitedContextEngine { +impl KvEngine for WindowedCheckpointEngine { fn name(&self) -> &str { - "unlimited-context" + "windowed-checkpoint" } fn info(&self) -> EngineInfo { let mem = self.checkpoints.total_bytes() + self.archive.total_bytes() + self.current_kv_bytes(); EngineInfo { - name: "unlimited-context".into(), + name: "windowed-checkpoint".into(), description: format!( "window-boundary KV checkpoints + token replay \ (windows={}, tokens={}, mem={:.1}MB)", @@ -529,10 +595,7 @@ impl KvEngine for UnlimitedContextEngine { if token_ids.is_empty() { return Err(EngineError::EmptyPrompt); } - self.process(weights, token_ids, Some(ffn)) - .ok_or_else(|| EngineError::BackendFailure { - details: "process returned None during prefill".into(), - })?; + self.process(weights, token_ids, Some(ffn))?; self.last_hidden .clone() .ok_or_else(|| EngineError::BackendFailure { @@ -546,10 +609,7 @@ impl KvEngine for UnlimitedContextEngine { ffn: &dyn FfnBackend, token_id: u32, ) -> Result, EngineError> { - self.process(weights, &[token_id], Some(ffn)) - .ok_or_else(|| EngineError::BackendFailure { - details: "process returned None during decode_step".into(), - })?; + self.process(weights, &[token_id], Some(ffn))?; self.last_hidden .clone() .ok_or_else(|| EngineError::BackendFailure { @@ -566,10 +626,7 @@ impl KvEngine for UnlimitedContextEngine { index: &larql_vindex::VectorIndex, token_id: u32, ) -> Result, EngineError> { - self.process_with_index(weights, &[token_id], Some(ffn), Some(index)) - .ok_or_else(|| EngineError::BackendFailure { - details: "process returned None during decode_step".into(), - })?; + self.process_with_index(weights, &[token_id], Some(ffn), Some(index))?; self.last_hidden .clone() .ok_or_else(|| EngineError::BackendFailure { @@ -578,7 +635,15 @@ impl KvEngine for UnlimitedContextEngine { } fn memory_bytes(&self) -> usize { - self.checkpoints.total_bytes() + self.archive.total_bytes() + self.current_kv_bytes() + // The last term covers the coarse dispatch path, where the live + // window's K/V is held by the backend rather than in + // `current_window_kv` — without it the hot tier reads as 0.0MB + // and the ratio against a standard cache becomes unbounded. + self.checkpoints.total_bytes() + + self.archive.total_bytes() + + self.current_kv_bytes() + + self.kv_handle.as_ref().map_or(0, |h| h.resident_bytes()) + + self.backend.backend_resident_kv_bytes() } fn window_tokens(&self) -> usize { @@ -589,13 +654,25 @@ impl KvEngine for UnlimitedContextEngine { self.checkpoints.total_bytes() + self.archive.total_bytes() } + fn dispatch_path(&self) -> Option { + use larql_inference::kv_engine::DispatchPath; + // `kv_handle` marks the coarse W1-GPU path; `current_window_kv` + // is the per-layer shadow this engine keeps when it walks + // layer-by-layer. Neither = nothing prefilled yet. + match (self.kv_handle.is_some(), self.current_window_kv.is_some()) { + (true, _) => Some(DispatchPath::Coarse), + (false, true) => Some(DispatchPath::PerLayer), + (false, false) => None, + } + } + fn stage_summary(&self) -> Option { if !self.profiling || self.profile.decode_total.count == 0 { return None; } Some( self.profile - .summary("unlimited-context", self.backend.name()), + .summary("windowed-checkpoint", self.backend.name()), ) } @@ -729,7 +806,7 @@ impl KvEngine for UnlimitedContextEngine { // ── Executor-driven window extension ───────────────────────────────────────── -impl UnlimitedContextEngine { +impl WindowedCheckpointEngine { /// Executor-aware analogue of `process_quant`: feeds tokens into the /// current window, auto-closes on fill, drives per-layer compute /// through `executor` instead of constructing a local `WalkFfn`. @@ -838,7 +915,7 @@ mod tests { #[test] fn new_engine_is_empty() { - let eng = UnlimitedContextEngine::new(512); + let eng = WindowedCheckpointEngine::new(512); assert_eq!(eng.window_size, 512); assert_eq!(eng.archive.len(), 0); assert_eq!(eng.checkpoints.len(), 0); @@ -848,28 +925,28 @@ mod tests { #[test] fn engine_info_backend_is_cpu() { - let eng = UnlimitedContextEngine::new(256); + let eng = WindowedCheckpointEngine::new(256); let info = eng.info(); - assert_eq!(info.name, "unlimited-context"); + assert_eq!(info.name, "windowed-checkpoint"); assert!( info.backend.starts_with("cpu"), "expected cpu backend, got {:?}", info.backend ); assert_eq!(info.config, "window=256"); - assert!(info.summary().contains("unlimited-context")); + assert!(info.summary().contains("windowed-checkpoint")); assert!(info.summary().contains("cpu")); } #[test] fn engine_info_config_contains_window_size() { - let eng = UnlimitedContextEngine::new(1024); + let eng = WindowedCheckpointEngine::new(1024); assert!(eng.info().config.contains("1024")); } #[test] fn window_tokens_and_cold_bytes_start_zero() { - let eng = UnlimitedContextEngine::new(512); + let eng = WindowedCheckpointEngine::new(512); assert_eq!(eng.window_tokens(), 0); assert_eq!(eng.cold_bytes(), 0); } @@ -882,7 +959,7 @@ mod tests { use larql_inference::test_utils::make_test_weights; let weights = make_test_weights(); let ffn = WeightFfn { weights: &weights }; - let mut engine = UnlimitedContextEngine::new(512); + let mut engine = WindowedCheckpointEngine::new(512); let h = engine .prefill(&weights, &ffn, &[0u32, 1, 2]) .expect("prefill failed"); @@ -899,7 +976,7 @@ mod tests { use larql_inference::test_utils::make_test_weights; let weights = make_test_weights(); let ffn = WeightFfn { weights: &weights }; - let mut engine = UnlimitedContextEngine::new(512); + let mut engine = WindowedCheckpointEngine::new(512); engine.prefill(&weights, &ffn, &[0u32]).expect("prefill"); let h = engine.decode_step(&weights, &ffn, 1).expect("decode_step"); assert_eq!(h.shape(), &[1, weights.hidden_size]); @@ -911,7 +988,7 @@ mod tests { use larql_inference::test_utils::make_test_weights; let weights = make_test_weights(); let window_size = 3usize; - let mut engine = UnlimitedContextEngine::new(window_size); + let mut engine = WindowedCheckpointEngine::new(window_size); // Feed exactly window_size tokens → triggers close for tok in 0..window_size as u32 { @@ -936,7 +1013,7 @@ mod tests { fn two_full_windows_archives_two() { use larql_inference::test_utils::make_test_weights; let weights = make_test_weights(); - let mut engine = UnlimitedContextEngine::new(2); + let mut engine = WindowedCheckpointEngine::new(2); // 4 tokens = 2 complete windows for tok in 0u32..4 { @@ -950,7 +1027,7 @@ mod tests { fn partial_window_after_process() { use larql_inference::test_utils::make_test_weights; let weights = make_test_weights(); - let mut engine = UnlimitedContextEngine::new(4); + let mut engine = WindowedCheckpointEngine::new(4); // 3 tokens < window_size=4 → no close engine @@ -964,7 +1041,7 @@ mod tests { fn flush_closes_partial_window() { use larql_inference::test_utils::make_test_weights; let weights = make_test_weights(); - let mut engine = UnlimitedContextEngine::new(4); + let mut engine = WindowedCheckpointEngine::new(4); engine.process(&weights, &[0u32, 1], None).expect("process"); assert_eq!(engine.archive.len(), 0); engine.flush(); @@ -975,7 +1052,7 @@ mod tests { fn cold_bytes_grow_after_window_close() { use larql_inference::test_utils::make_test_weights; let weights = make_test_weights(); - let mut engine = UnlimitedContextEngine::new(2); + let mut engine = WindowedCheckpointEngine::new(2); assert_eq!(engine.cold_bytes(), 0); engine.process(&weights, &[0u32, 1], None).expect("process"); // closes window assert!( @@ -990,7 +1067,7 @@ mod tests { use larql_inference::test_utils::make_test_weights; let weights = make_test_weights(); let ffn = WeightFfn { weights: &weights }; - let mut engine = UnlimitedContextEngine::new(512); + let mut engine = WindowedCheckpointEngine::new(512); assert_eq!(engine.memory_bytes(), 0); engine .prefill(&weights, &ffn, &[0u32, 1, 2]) @@ -999,13 +1076,13 @@ mod tests { } #[test] - fn logits_from_unlimited_context_are_finite() { + fn logits_from_windowed_checkpoint_are_finite() { use larql_inference::ffn::WeightFfn; use larql_inference::forward::hidden_to_raw_logits; use larql_inference::test_utils::make_test_weights; let weights = make_test_weights(); let ffn = WeightFfn { weights: &weights }; - let mut engine = UnlimitedContextEngine::new(512); + let mut engine = WindowedCheckpointEngine::new(512); let h = engine.prefill(&weights, &ffn, &[0u32, 1]).expect("prefill"); let logits = hidden_to_raw_logits(&weights, &h); assert!( @@ -1029,7 +1106,7 @@ mod tests { let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; - let mut engine = UnlimitedContextEngine::new(512); + let mut engine = WindowedCheckpointEngine::new(512); let h = engine .prefill_quant(&weights, &ffn, &index, &[0u32, 1, 2], &*backend) .expect("prefill_quant Q4K cpu fallback"); @@ -1044,7 +1121,7 @@ mod tests { let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; - let mut engine = UnlimitedContextEngine::new(512); + let mut engine = WindowedCheckpointEngine::new(512); engine .prefill_quant(&weights, &ffn, &index, &[0u32, 1], &*backend) .expect("prefill_quant"); @@ -1082,7 +1159,7 @@ mod tests { "LARQL_MARKOV_INPLACE_KV", Some(if inplace { "1" } else { "0" }), ); - let mut engine = UnlimitedContextEngine::new(512); + let mut engine = WindowedCheckpointEngine::new(512); engine .prefill(&weights, &ffn, &[0u32, 1, 2]) .expect("prefill"); @@ -1113,7 +1190,7 @@ mod tests { let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; - let mut engine = UnlimitedContextEngine::new(512); + let mut engine = WindowedCheckpointEngine::new(512); // No prefill → decode falls through fast-path checks and returns None // (or some empty hidden) without panicking. let _ = engine.decode_step_quant(&weights, &ffn, &index, 0, &*backend); @@ -1127,7 +1204,7 @@ mod tests { use larql_inference::test_utils::make_test_weights; let weights = make_test_weights(); let ffn = WeightFfn { weights: &weights }; - let mut engine = UnlimitedContextEngine::new(512); + let mut engine = WindowedCheckpointEngine::new(512); engine .prefill(&weights, &ffn, &[0u32, 1, 2]) .expect("prefill"); @@ -1143,7 +1220,7 @@ mod tests { #[test] fn engine_stats_with_empty_engine_handles_zero_division() { let weights = larql_inference::test_utils::make_test_weights(); - let engine = UnlimitedContextEngine::new(512); + let engine = WindowedCheckpointEngine::new(512); let stats = engine.stats(&weights); // No prefill → all counters zero, compression ratio short-circuits // to 0.0 (no division by zero). @@ -1160,11 +1237,11 @@ mod tests { #[test] fn replay_window_returns_none_for_missing_window() { let weights = larql_inference::test_utils::make_test_weights(); - let engine = UnlimitedContextEngine::new(512); + let engine = WindowedCheckpointEngine::new(512); // No windows archived → any window_id returns None at the // `self.archive.retrieve(window_id)?` line. - assert!(engine.replay_window(&weights, None, None, 0).is_none()); - assert!(engine.replay_window(&weights, None, None, 99).is_none()); + assert!(engine.replay_window(&weights, None, None, 0).is_err()); + assert!(engine.replay_window(&weights, None, None, 99).is_err()); } #[test] @@ -1174,7 +1251,7 @@ mod tests { let weights = make_test_weights(); let ffn = WeightFfn { weights: &weights }; // window=2; prefill 4 tokens → archives at least 1 window. - let mut engine = UnlimitedContextEngine::new(2); + let mut engine = WindowedCheckpointEngine::new(2); engine .prefill(&weights, &ffn, &[0u32, 1, 2, 3]) .expect("prefill 4 tokens"); @@ -1187,7 +1264,7 @@ mod tests { // Replay the first archived window — exercises the // `rs_extend_from_checkpoint_backend` path (lines 132-138). let replay = engine.replay_window(&weights, None, None, 0); - assert!(replay.is_some(), "replay_window(0) should succeed"); + assert!(replay.is_ok(), "replay_window(0) should succeed"); let (kv, abs_end) = replay.unwrap(); assert!(!kv.is_empty(), "replayed K/V cache should be non-empty"); assert!( @@ -1208,7 +1285,7 @@ mod tests { let backend = larql_compute::cpu_backend(); let executor = LocalWalkExecutor::new(&*backend); let ffn = NullFfn; - let mut engine = UnlimitedContextEngine::new(512); + let mut engine = WindowedCheckpointEngine::new(512); let h = engine .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32, 1, 2]) .expect("executor prefill"); @@ -1226,7 +1303,7 @@ mod tests { let backend = larql_compute::cpu_backend(); let executor = LocalWalkExecutor::new(&*backend); let ffn = NullFfn; - let mut engine = UnlimitedContextEngine::new(512); + let mut engine = WindowedCheckpointEngine::new(512); engine .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32, 1]) .expect("prefill"); @@ -1247,7 +1324,7 @@ mod tests { let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; - let mut engine = UnlimitedContextEngine::new(512).with_profiling(true); + let mut engine = WindowedCheckpointEngine::new(512).with_profiling(true); engine .prefill_quant(&weights, &ffn, &index, &[0u32, 1], &*backend) .expect("prefill"); @@ -1256,8 +1333,8 @@ mod tests { .expect("decode"); let summary = engine .stage_summary() - .expect("unlimited_context profiler should populate summary"); - assert_eq!(summary.engine, "unlimited-context"); + .expect("windowed_checkpoint profiler should populate summary"); + assert_eq!(summary.engine, "windowed-checkpoint"); assert!(summary.steps >= 1); assert!(summary.avg_attention_us > 0.0); assert!(summary.avg_ffn_us > 0.0); @@ -1295,7 +1372,7 @@ mod tests { calls: std::sync::atomic::AtomicUsize::new(0), hidden: weights.hidden_size, }; - let mut engine = UnlimitedContextEngine::new(512); + let mut engine = WindowedCheckpointEngine::new(512); engine .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32, 1, 2]) .expect("prefill via executor"); @@ -1318,14 +1395,14 @@ mod tests { #[test] #[should_panic(expected = "window_size must be >= 1")] fn zero_window_size_is_rejected_at_construction() { - let _ = UnlimitedContextEngine::new(0); + let _ = WindowedCheckpointEngine::new(0); } #[test] fn window_size_one_is_legal() { use larql_inference::test_utils::make_test_weights; let weights = make_test_weights(); - let mut engine = UnlimitedContextEngine::new(1); + let mut engine = WindowedCheckpointEngine::new(1); engine.process(&weights, &[0u32, 1], None).expect("process"); assert_eq!( engine.archive.len(), @@ -1343,7 +1420,7 @@ mod tests { /// close with zero free slots. #[test] fn close_window_without_shadow_or_handle_recovers_bookkeeping() { - let mut engine = UnlimitedContextEngine::new(2); + let mut engine = WindowedCheckpointEngine::new(2); engine.current_window_tokens = vec![7, 8]; engine.current_window_kv = None; engine.current_window_kv_len = 2; @@ -1372,7 +1449,7 @@ mod tests { use larql_inference::test_utils::make_test_weights; let weights = make_test_weights(); let ffn = WeightFfn { weights: &weights }; - let mut engine = UnlimitedContextEngine::new(8); + let mut engine = WindowedCheckpointEngine::new(8); engine .prefill(&weights, &ffn, &[0u32, 1, 2]) .expect("prefill"); @@ -1393,7 +1470,7 @@ mod tests { let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; - let mut engine = UnlimitedContextEngine::new(512); + let mut engine = WindowedCheckpointEngine::new(512); engine .prefill_quant(&weights, &ffn, &index, &[0u32, 1], &*backend) .expect("prefill"); @@ -1421,7 +1498,7 @@ mod tests { // window=2, 4 tokens → triggers two window-close cycles via // `process_via_executor`. Exercises the prior-checkpoint-load // branch in `extend_current_via_executor`. - let mut engine = UnlimitedContextEngine::new(2); + let mut engine = WindowedCheckpointEngine::new(2); engine .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32, 1, 2, 3]) .expect("prefill 4 tokens through executor"); diff --git a/crates/larql-kv/src/engines/unlimited_context/extend.rs b/crates/larql-kv/src/engines/windowed_checkpoint/extend.rs similarity index 69% rename from crates/larql-kv/src/engines/unlimited_context/extend.rs rename to crates/larql-kv/src/engines/windowed_checkpoint/extend.rs index d7c394ef7..62b8cefa1 100644 --- a/crates/larql-kv/src/engines/unlimited_context/extend.rs +++ b/crates/larql-kv/src/engines/windowed_checkpoint/extend.rs @@ -5,12 +5,13 @@ use larql_compute::ComputeBackend; use larql_vindex::VectorIndex; -use ndarray::Array2; +use ndarray::{s, Array2}; use larql_inference::attention::{run_attention_block_decode_step_backend, SharedKV}; use larql_inference::ffn::BackendFfn; use larql_inference::forward::ple::precompute_per_layer_inputs; use larql_inference::forward::{embed_tokens_pub, run_ffn}; +use larql_inference::kv_engine::EngineError; use larql_inference::model::ModelWeights; use larql_inference::vindex::{WalkFfn, WalkFfnConfig}; @@ -23,6 +24,36 @@ pub struct ExtendOutput { pub new_checkpoint: Vec, } +/// What an extend produced when the caller kept ownership of the K/V. +/// +/// The cache itself is left in the caller's buffer rather than returned, so a +/// step that fails partway can be rewound by whoever owns the window — see +/// [`truncate_kv_rows`]. +pub struct ExtendStep { + /// Hidden state at the last processed token, shape (1, hidden). + pub last_hidden: Array2, + /// Per-layer last-row K,V ready to save as the next boundary checkpoint. + pub new_checkpoint: Vec, +} + +/// Truncate every layer's K/V back to `rows`, undoing an extend that stopped +/// partway. +/// +/// Needed because the owned-concat path *replaces* each layer's buffer with a +/// longer one as it goes, and this path reads a prior by `shape()[0]` rather +/// than by a counter — so a partially advanced cache would silently attend +/// over a token whose step never completed. +pub fn truncate_kv_rows(kv_cache: &mut [SharedKV], rows: usize) { + for (k, v) in kv_cache.iter_mut() { + if k.shape()[0] > rows { + *k = k.slice(s![..rows, ..]).to_owned(); + } + if v.shape()[0] > rows { + *v = v.slice(s![..rows, ..]).to_owned(); + } + } +} + /// Run the decoder forward over `token_ids` seeded with an optional prior K,V /// checkpoint at each layer. Matmuls route through `backend`. /// @@ -32,43 +63,56 @@ pub fn rs_extend_from_checkpoint( token_ids: &[u32], prior_kv: Vec, abs_start: usize, -) -> Option { - rs_extend_from_checkpoint_backend( +) -> Result { + let mut kv_cache = prior_kv; + let step = rs_extend_from_checkpoint_backend( weights, token_ids, - prior_kv, + &mut kv_cache, abs_start, &larql_compute::CpuBackend, None, None, - ) + )?; + Ok(ExtendOutput { + last_hidden: step.last_hidden, + kv_cache, + new_checkpoint: step.new_checkpoint, + }) } /// Backend-dispatched variant of [`rs_extend_from_checkpoint`]. /// -/// Takes `prior_kv` by value so the per-token extend loop can mutate it -/// in place. Cloning the prior K/V per step is O(window²) total over a -/// full window — a real overhead on growing caches. +/// Takes `kv_cache` by reference so the per-token extend loop can mutate it +/// in place — cloning the prior K/V per step is O(window²) total over a full +/// window, a real overhead on growing caches — and so that a caller who owns +/// the window can rewind it after a failure. **On `Err` the cache is left +/// partially advanced**: layers before the failing one hold the new row. +/// [`truncate_kv_rows`] is how the owner undoes that. #[allow(clippy::too_many_arguments)] pub fn rs_extend_from_checkpoint_backend( weights: larql_inference::WeightsView, token_ids: &[u32], - prior_kv: Vec, + kv_cache: &mut [SharedKV], abs_start: usize, backend: &dyn ComputeBackend, moe_ffn: Option<&dyn larql_inference::ffn::FfnBackend>, index: Option<&larql_vindex::VectorIndex>, -) -> Option { +) -> Result { let num_layers = weights.num_layers; if token_ids.is_empty() { - return None; + return Err(EngineError::EmptyPrompt); } - if prior_kv.len() != num_layers { - return None; + if kv_cache.len() != num_layers { + return Err(EngineError::InvariantViolation { + what: format!( + "prior K/V has {} layers, model has {num_layers}", + kv_cache.len() + ), + }); } - let mut kv_cache: Vec = prior_kv; let mut last_hidden: Option> = None; for (i, &token_id) in token_ids.iter().enumerate() { @@ -94,21 +138,26 @@ pub fn rs_extend_from_checkpoint_backend( abs_position, Some(backend), index.map(|v| v as &dyn larql_compute::KvIndex), - )?; + ) + .ok_or_else(|| EngineError::BackendFailure { + details: format!( + "attention returned None during unlimited-context extend at layer {layer}" + ), + })?; let bffn = BackendFfn { weights: weights.canonical(), backend, }; - let h_out = crate::engines::layer_ffn_or_moe( + h = crate::engines::layer_ffn_or_moe( weights.canonical(), &h_post_attn, layer, &bffn, moe_ffn, ple_inputs.get(layer), - ); - h = h_out; + ) + .map_err(EngineError::Execution)?; *kv_slot = new_kv; } @@ -125,9 +174,10 @@ pub fn rs_extend_from_checkpoint_backend( }) .collect(); - Some(ExtendOutput { - last_hidden: last_hidden?, - kv_cache, + Ok(ExtendStep { + last_hidden: last_hidden.ok_or_else(|| EngineError::BackendFailure { + details: "extend produced no hidden state".into(), + })?, new_checkpoint, }) } @@ -155,10 +205,16 @@ pub fn rs_extend_inplace( backend: &dyn ComputeBackend, moe_ffn: Option<&dyn larql_inference::ffn::FfnBackend>, index: Option<&larql_vindex::VectorIndex>, -) -> Option> { +) -> Result, EngineError> { let num_layers = weights.num_layers; if token_ids.is_empty() || kv_cache.len() != num_layers { - return None; + return Err(EngineError::InvariantViolation { + what: format!( + "in-place extend needs a non-empty chunk and {num_layers} K/V slots; got {} tokens and {} slots", + token_ids.len(), + kv_cache.len() + ), + }); } let idx_kv: Option<&dyn larql_compute::KvIndex> = index.map(|v| v as &dyn larql_compute::KvIndex); @@ -205,7 +261,15 @@ pub fn rs_extend_inplace( abs_position, Some(backend), idx_kv, - )?; + ) + .ok_or_else(|| { + EngineError::BackendFailure { + details: format!( + "attention returned None during in-place extend \ + at layer {layer}" + ), + } + })?; *k_buf = new_kv.0; *v_buf = new_kv.1; hp @@ -216,21 +280,23 @@ pub fn rs_extend_inplace( weights: weights.canonical(), backend, }; - let h_out = crate::engines::layer_ffn_or_moe( + h = crate::engines::layer_ffn_or_moe( weights.canonical(), &h_post_attn, layer, &bffn, moe_ffn, ple_inputs.get(layer), - ); - h = h_out; + ) + .map_err(EngineError::Execution)?; } last_hidden = Some(h); } - last_hidden + last_hidden.ok_or_else(|| EngineError::BackendFailure { + details: "in-place extend produced no hidden state".into(), + }) } /// CPU Q4K variant of [`rs_extend_from_checkpoint_backend`]. @@ -395,7 +461,7 @@ pub fn rs_extend_from_checkpoint_quant( } if let (Some(prof), Some(t_step)) = (profiler.as_mut(), t_step) { - // unlimited_context appends K/V incrementally → no `recompute_*` + // windowed_checkpoint appends K/V incrementally → no `recompute_*` // stages fire. embed/attention/ffn/decode_total carry the // attribution; recompute_cold/hot stay at zero (and the print // logic shows them only when non-zero). @@ -468,7 +534,10 @@ mod tests { let prior = empty_prior(&weights); let result = rs_extend_from_checkpoint(larql_inference::WeightsView::dense(&weights), &[], prior, 0); - assert!(result.is_none(), "empty token_ids should return None"); + assert!( + matches!(result, Err(EngineError::EmptyPrompt)), + "an empty chunk is a caller-input error, not a backend failure" + ); } #[test] @@ -481,7 +550,10 @@ mod tests { Vec::new(), 0, ); - assert!(result.is_none(), "prior length mismatch should return None"); + assert!( + matches!(result, Err(EngineError::InvariantViolation { .. })), + "a prior with the wrong layer count is a contract violation" + ); } #[test] @@ -711,4 +783,183 @@ mod tests { assert_eq!(k.shape(), &[3, kv_dim], "prior(2) + new(1) = 3 rows"); } } + + // ── truncate_kv_rows ────────────────────────────────────────────────────── + + #[test] + fn truncate_kv_rows_rewinds_every_layer_to_the_row_count() { + let weights = make_test_weights(); + let mut kv = rs_extend_from_checkpoint( + larql_inference::WeightsView::dense(&weights), + &[0u32, 1, 2], + empty_prior(&weights), + 0, + ) + .expect("3-token extend") + .kv_cache; + + // Row 0 must survive the rewind byte-for-byte — a truncate that + // reallocated the wrong slice would still leave the shape right. + let row0: Vec> = kv.iter().map(|(k, _)| k.row(0).to_vec()).collect(); + + truncate_kv_rows(&mut kv, 1); + + let kv_dim = weights.num_kv_heads * weights.head_dim; + for (layer, (k, v)) in kv.iter().enumerate() { + assert_eq!(k.shape(), &[1, kv_dim], "layer {layer}: K not rewound"); + assert_eq!(v.shape(), &[1, kv_dim], "layer {layer}: V not rewound"); + assert_eq!( + k.row(0).to_vec(), + row0[layer], + "layer {layer}: rewind kept the wrong row" + ); + } + } + + #[test] + fn truncate_kv_rows_leaves_a_shorter_cache_alone() { + let weights = make_test_weights(); + let mut kv = rs_extend_from_checkpoint( + larql_inference::WeightsView::dense(&weights), + &[0u32], + empty_prior(&weights), + 0, + ) + .expect("1-token extend") + .kv_cache; + + // rows > shape[0]: the guard must skip, not grow or panic. + truncate_kv_rows(&mut kv, 8); + + let kv_dim = weights.num_kv_heads * weights.head_dim; + for (k, v) in &kv { + assert_eq!(k.shape(), &[1, kv_dim]); + assert_eq!(v.shape(), &[1, kv_dim]); + } + } + + // ── rs_extend_inplace ───────────────────────────────────────────────────── + + #[test] + fn extend_inplace_empty_tokens_is_an_invariant_violation() { + let weights = make_test_weights(); + let backend = larql_compute::cpu_backend(); + let mut kv = empty_prior(&weights); + let result = rs_extend_inplace( + larql_inference::WeightsView::dense(&weights), + &[], + &mut kv, + 0, + 0, + &*backend, + None, + None, + ); + assert!( + matches!(result, Err(EngineError::InvariantViolation { .. })), + "an empty chunk breaks the in-place contract" + ); + } + + #[test] + fn extend_inplace_wrong_slot_count_is_an_invariant_violation() { + let weights = make_test_weights(); + let backend = larql_compute::cpu_backend(); + // Model has `num_layers` layers; hand it zero K/V slots. + let mut kv: Vec = Vec::new(); + let result = rs_extend_inplace( + larql_inference::WeightsView::dense(&weights), + &[0u32], + &mut kv, + 0, + 0, + &*backend, + None, + None, + ); + assert!( + matches!(result, Err(EngineError::InvariantViolation { .. })), + "a slot count that isn't num_layers breaks the in-place contract" + ); + } + + /// With no index the Q4K-direct in-place projection returns `None` at every + /// layer, so this drives the per-layer owned-concat fallback — the arm that + /// writes the rebuilt buffer back so the cache stays consistent. + #[test] + fn extend_inplace_falls_back_to_owned_concat_without_an_index() { + let weights = make_test_weights(); + let backend = larql_compute::cpu_backend(); + let mut kv = empty_prior(&weights); + let last = rs_extend_inplace( + larql_inference::WeightsView::dense(&weights), + &[0u32, 1, 2], + &mut kv, + 0, + 0, + &*backend, + None, + None, + ) + .expect("fallback extend should still produce a hidden state"); + + assert_eq!(last.shape(), &[1, weights.hidden_size]); + assert!(last.iter().all(|v| v.is_finite())); + // The fallback replaces each buffer with the owned concat, so after 3 + // tokens from an empty prior every layer holds exactly 3 rows. + let kv_dim = weights.num_kv_heads * weights.head_dim; + for (layer, (k, v)) in kv.iter().enumerate() { + assert_eq!(k.shape(), &[3, kv_dim], "layer {layer}: K rows"); + assert_eq!(v.shape(), &[3, kv_dim], "layer {layer}: V rows"); + } + } + + /// The fallback is also the seeded path: `prior_len > 0` makes it slice a + /// real prior out of the buffer rather than pass `None`. + #[test] + fn extend_inplace_fallback_matches_the_owned_concat_path() { + let weights = make_test_weights(); + let backend = larql_compute::cpu_backend(); + let view = larql_inference::WeightsView::dense(&weights); + + let mut inplace_kv = empty_prior(&weights); + let inplace = rs_extend_inplace( + view, + &[0u32, 1], + &mut inplace_kv, + 0, + 0, + &*backend, + None, + None, + ) + .expect("in-place extend"); + + let mut owned_kv = empty_prior(&weights); + let owned = rs_extend_from_checkpoint_backend( + view, + &[0u32, 1], + &mut owned_kv, + 0, + &*backend, + None, + None, + ) + .expect("owned-concat extend"); + + // Same numerics, different cache representation — that equivalence is + // the whole claim the fallback arm exists to preserve. + for (a, b) in inplace.iter().zip(owned.last_hidden.iter()) { + assert!( + (a - b).abs() < 1e-6, + "in-place fallback diverged from owned concat: {a} vs {b}" + ); + } + for (layer, ((ki, _), (ko, _))) in inplace_kv.iter().zip(owned_kv.iter()).enumerate() { + assert_eq!(ki.shape(), ko.shape(), "layer {layer}: K shape"); + for (a, b) in ki.iter().zip(ko.iter()) { + assert!((a - b).abs() < 1e-6, "layer {layer}: K diverged"); + } + } + } } diff --git a/crates/larql-kv/src/engines/unlimited_context/mod.rs b/crates/larql-kv/src/engines/windowed_checkpoint/mod.rs similarity index 85% rename from crates/larql-kv/src/engines/unlimited_context/mod.rs rename to crates/larql-kv/src/engines/windowed_checkpoint/mod.rs index 1a1d6811f..b46658320 100644 --- a/crates/larql-kv/src/engines/unlimited_context/mod.rs +++ b/crates/larql-kv/src/engines/windowed_checkpoint/mod.rs @@ -5,7 +5,7 @@ pub mod extend; pub mod token_archive; pub use checkpoint_store::CheckpointStore; -pub use engine::{EngineStats, UnlimitedContextEngine}; +pub use engine::{EngineStats, WindowedCheckpointEngine}; pub use extend::{ empty_prior, rs_extend_from_checkpoint, rs_extend_from_checkpoint_backend, rs_extend_from_checkpoint_quant, ExtendOutput, diff --git a/crates/larql-kv/src/engines/unlimited_context/token_archive.rs b/crates/larql-kv/src/engines/windowed_checkpoint/token_archive.rs similarity index 100% rename from crates/larql-kv/src/engines/unlimited_context/token_archive.rs rename to crates/larql-kv/src/engines/windowed_checkpoint/token_archive.rs diff --git a/crates/larql-kv/src/generation.rs b/crates/larql-kv/src/generation.rs index 21ef97648..1118e902d 100644 --- a/crates/larql-kv/src/generation.rs +++ b/crates/larql-kv/src/generation.rs @@ -43,6 +43,10 @@ use ndarray::Array2; use crate::cache::KvCache; +mod kv_run; + +pub use kv_run::{kv_decode_step_run, kv_prefill_run}; + /// Stream autoregressive generation with a KV cache. /// /// `on_token` receives `(token_id, decoded_string)` for each generated @@ -438,118 +442,6 @@ where generated } -/// Prefill phase as a reusable building block: runs a full forward over -/// `prompt_ids`, populates a fresh [`KvCache`] (bounded if `window` is -/// `Some`), and returns `(last_hidden_1xD, populated_cache)`. -/// -/// Returns `None` if the prompt is empty or if any layer's attention -/// fails. This is the production K/V cache prefill loop, extracted so -/// `KvEngine::prefill` impls can call it directly. -/// -/// The caller applies `final_norm + lm_head` to the returned hidden -/// state to get logits. -#[allow(clippy::too_many_arguments)] -pub fn kv_prefill_run( - weights: larql_inference::WeightsView, - ffn: &dyn FfnBackend, - prompt_ids: &[u32], - window: Option, - backend: Option<&dyn larql_compute::ComputeBackend>, - hook: &mut dyn LayerHook, -) -> Option<(Array2, KvCache)> { - if prompt_ids.is_empty() { - return None; - } - let num_layers = weights.num_layers; - let mut cache = match window { - Some(w) => KvCache::with_window(num_layers, w), - None => KvCache::with_layers(num_layers), - }; - - let mut h = embed_tokens_pub(&weights, prompt_ids); - // Per-Layer Embedding inputs for Gemma-4 archs. Returns empty Vec - // for non-PLE archs (`ple_inputs.get(layer)` then yields `None` and - // `apply_per_layer_embedding` is a no-op). - let ple_inputs = precompute_per_layer_inputs(&weights, &h, prompt_ids); - for layer in 0..num_layers { - hook.on_pre_layer(layer, &h); - - let (mut h_post_attn, k_rope, v) = - run_attention_with_kv_backend(weights, &h, layer, backend, None)?; - cache.layers[layer] = Some((k_rope, v)); - cache.clip_layer(layer); - - hook.on_post_attention(layer, &mut h_post_attn); - - let (h_post_ffn, _) = run_ffn(&weights, &h_post_attn, layer, ffn, false); - let mut h_out = - apply_per_layer_embedding(&weights, &h_post_ffn, layer, ple_inputs.get(layer)); - apply_layer_scalar(&weights, &mut h_out, layer); - - hook.on_post_layer(layer, &mut h_out); - h = h_out; - } - cache.next_position = prompt_ids.len(); - - Some((last_row_as_2d(&h), cache)) -} - -/// Decode-step phase as a reusable building block: takes one new -/// `token_id`, runs the autoregressive attention against an existing -/// populated [`KvCache`], mutates the cache to append the new K/V (and -/// clip to window), and returns the new token's hidden state (shape -/// `[1, hidden_dim]`). -/// -/// Returns `None` if any layer's attention fails. This is the -/// production decode step extracted so `KvEngine::decode_step` impls -/// can call it directly. -#[allow(clippy::too_many_arguments)] -pub fn kv_decode_step_run( - weights: &ModelWeights, - ffn: &dyn FfnBackend, - cache: &mut KvCache, - token_id: u32, - backend: Option<&dyn larql_compute::ComputeBackend>, - hook: &mut dyn LayerHook, -) -> Option> { - let num_layers = weights.num_layers; - let h_new = embed_tokens_pub(weights, &[token_id]); - let abs_position = cache.next_position; - // PLE inputs are per-token. Recompute for this single-token decode - // step rather than indexing a prefill-sized slab. Matches the - // recipe used by `vindex::kquant_forward::cached` and the GPU - // `layer_graph::generate` decode loop. - let ple_inputs = precompute_per_layer_inputs(weights, &h_new, &[token_id]); - let mut h_step = h_new; - for layer in 0..num_layers { - hook.on_pre_layer(layer, &h_step); - - let kv_entry = cache.layers[layer].as_ref(); - let (mut h_post_attn, new_kv) = run_attention_block_decode_step_backend( - larql_inference::WeightsView::dense(weights), - &h_step, - layer, - kv_entry, - abs_position, - backend, - )?; - cache.layers[layer] = Some(new_kv); - cache.clip_layer(layer); - - hook.on_post_attention(layer, &mut h_post_attn); - - let (h_post_ffn, _) = run_ffn(weights, &h_post_attn, layer, ffn, false); - let mut h_out = - apply_per_layer_embedding(weights, &h_post_ffn, layer, ple_inputs.get(layer)); - apply_layer_scalar(weights, &mut h_out, layer); - - hook.on_post_layer(layer, &mut h_out); - h_step = h_out; - } - cache.next_position += 1; - Some(h_step) -} - #[allow(clippy::too_many_arguments)] fn generate_cached_hooked_inner( weights: &ModelWeights, @@ -575,8 +467,8 @@ fn generate_cached_hooked_inner( backend, hook, ) { - Some(t) => t, - None => return Vec::new(), + Ok(t) => t, + Err(_) => return Vec::new(), }; let first = match argmax_next_token(weights, tokenizer, &last_hidden) { @@ -597,8 +489,8 @@ fn generate_cached_hooked_inner( let mut current_id = first.0; for _step in 1..max_new_tokens { let h_step = match kv_decode_step_run(weights, ffn, &mut cache, current_id, backend, hook) { - Some(h) => h, - None => break, + Ok(h) => h, + Err(_) => break, }; let (id, tok_str) = match argmax_next_token(weights, tokenizer, &h_step) { Some(t) => t, @@ -902,12 +794,7 @@ mod tests { None, None, &mut NoopHook, - ) - .ok_or_else(|| { - larql_inference::kv_engine::EngineError::BackendFailure { - details: "kv_prefill_run returned None".into(), - } - })?; + )?; self.cache = Some(cache); Ok(hidden) } @@ -930,11 +817,7 @@ mod tests { what: "decode_step called before prefill".into(), } })?; - kv_decode_step_run(weights, ffn, cache, token_id, None, &mut NoopHook).ok_or_else( - || larql_inference::kv_engine::EngineError::BackendFailure { - details: "kv_decode_step_run returned None".into(), - }, - ) + kv_decode_step_run(weights, ffn, cache, token_id, None, &mut NoopHook) } // MM support: drive `generate_with_engine_from_hidden`. We can't // recover the original tokens from a pre-built hidden state, so the @@ -964,12 +847,7 @@ mod tests { None, None, &mut NoopHook, - ) - .ok_or_else(|| { - larql_inference::kv_engine::EngineError::BackendFailure { - details: "kv_prefill_run returned None".into(), - } - })?; + )?; self.cache = Some(cache); Ok(hidden) } @@ -1344,7 +1222,7 @@ mod tests { for step in 0..3 { let h_step = kv_decode_step_run(&weights, &ffn, &mut cache, 0u32, None, &mut NoopHook) - .unwrap_or_else(|| panic!("decode step {step} returned None")); + .unwrap_or_else(|e| panic!("decode step {step} failed: {e:?}")); assert_eq!(h_step.shape(), &[1, weights.hidden_size]); assert!( h_step.iter().all(|v| v.is_finite()), diff --git a/crates/larql-kv/src/generation/kv_run.rs b/crates/larql-kv/src/generation/kv_run.rs new file mode 100644 index 000000000..e8270183a --- /dev/null +++ b/crates/larql-kv/src/generation/kv_run.rs @@ -0,0 +1,198 @@ +//! The per-layer prefill and decode building blocks every `KvEngine` and the +//! dispatch ring are measured against. +//! +//! These two functions are the **oracle**: `dispatch_parity` compares +//! `kv_prefill_via_dispatch` / `kv_decode_step_via_dispatch` against them, and +//! `engine_ple_parity` compares the engines against them. That is why they +//! dispatch experts through [`crate::engines::layer_ffn_or_moe`] rather than +//! running a dense FFN — an oracle that skipped the expert half would make +//! every MoE parity comparison agree about the wrong answer. +//! +//! Both are transactional; see each function's contract and +//! [`transaction`] for the rewind that makes the decode one true. + +use larql_inference::attention::{ + run_attention_block_decode_step_backend, run_attention_with_kv_backend, +}; +use larql_inference::ffn::FfnBackend; +use larql_inference::forward::embed_tokens_pub; +use larql_inference::forward::hooks::LayerHook; +use larql_inference::forward::ple::precompute_per_layer_inputs; +use larql_inference::kv_engine::EngineError; +use larql_inference::ModelWeights; +use ndarray::Array2; + +use crate::cache::KvCache; +use crate::generation::last_row_as_2d; + +#[cfg(test)] +mod tests; +mod transaction; + +use transaction::{cache_row_counts, decode_rewind_is_sound, rewind_cache}; + +/// Prefill phase as a reusable building block: runs a full forward over +/// `prompt_ids`, populates a fresh [`KvCache`] (bounded if `window` is +/// `Some`), and returns `(last_hidden_1xD, populated_cache)`. +/// +/// This is the production K/V cache prefill loop, extracted so +/// `KvEngine::prefill` impls can call it directly, and the oracle the +/// dispatch ring is compared against — which is why it dispatches experts +/// through [`crate::engines::layer_ffn_or_moe`] rather than running a dense +/// FFN. An oracle that skipped the expert half would make every MoE parity +/// comparison agree about the wrong answer. +/// +/// Transactional: the cache is built into a local and returned only on +/// success, so a refusal costs the caller nothing it already had. +/// +/// The caller applies `final_norm + lm_head` to the returned hidden +/// state to get logits. +#[allow(clippy::too_many_arguments)] +pub fn kv_prefill_run( + weights: larql_inference::WeightsView, + ffn: &dyn FfnBackend, + prompt_ids: &[u32], + window: Option, + backend: Option<&dyn larql_compute::ComputeBackend>, + hook: &mut dyn LayerHook, +) -> Result<(Array2, KvCache), EngineError> { + if prompt_ids.is_empty() { + return Err(EngineError::EmptyPrompt); + } + let num_layers = weights.num_layers; + let mut cache = match window { + Some(w) => KvCache::with_window(num_layers, w), + None => KvCache::with_layers(num_layers), + }; + + let mut h = embed_tokens_pub(&weights, prompt_ids); + // Per-Layer Embedding inputs for Gemma-4 archs. Returns empty Vec + // for non-PLE archs (`ple_inputs.get(layer)` then yields `None` and + // `apply_per_layer_embedding` is a no-op). + let ple_inputs = precompute_per_layer_inputs(&weights, &h, prompt_ids); + for layer in 0..num_layers { + hook.on_pre_layer(layer, &h); + + let (mut h_post_attn, k_rope, v) = + run_attention_with_kv_backend(weights, &h, layer, backend, None).ok_or_else(|| { + EngineError::BackendFailure { + details: format!("attention returned None during prefill at layer {layer}"), + } + })?; + cache.layers[layer] = Some((k_rope, v)); + cache.clip_layer(layer); + + hook.on_post_attention(layer, &mut h_post_attn); + + let mut h_out = crate::engines::layer_ffn_or_moe( + weights.canonical(), + &h_post_attn, + layer, + ffn, + Some(ffn), + ple_inputs.get(layer), + ) + .map_err(EngineError::Execution)?; + + hook.on_post_layer(layer, &mut h_out); + h = h_out; + } + cache.next_position = prompt_ids.len(); + + Ok((last_row_as_2d(&h), cache)) +} + +/// Decode-step phase as a reusable building block: takes one new +/// `token_id`, runs the autoregressive attention against an existing +/// populated [`KvCache`], mutates the cache to append the new K/V (and +/// clip to window), and returns the new token's hidden state (shape +/// `[1, hidden_dim]`). +/// +/// This is the production decode step extracted so `KvEngine::decode_step` +/// impls can call it directly, and — like [`kv_prefill_run`] — the oracle the +/// dispatch ring is compared against, so it dispatches experts. +/// +/// **Transactional.** Each layer's attention appends the new token's K/V +/// before the FFN gets the chance to refuse, so a step that does not complete +/// truncates every layer back to the length it had on entry. When the cache is +/// windowed and already at its limit that truncation would be a lie — +/// append-then-evict leaves the row count unchanged while the oldest row is +/// gone — so the failure is reported as +/// [`EngineError::StateInvalidated`] instead, and the caller must rebuild the +/// cache rather than continue from it. +#[allow(clippy::too_many_arguments)] +pub fn kv_decode_step_run( + weights: &ModelWeights, + ffn: &dyn FfnBackend, + cache: &mut KvCache, + token_id: u32, + backend: Option<&dyn larql_compute::ComputeBackend>, + hook: &mut dyn LayerHook, +) -> Result, EngineError> { + let entry_rows = cache_row_counts(cache); + let rewindable = decode_rewind_is_sound(cache, &entry_rows); + match decode_step_appending(weights, ffn, cache, token_id, backend, hook) { + Ok(hidden) => Ok(hidden), + Err(failure) if rewindable => { + rewind_cache(cache, &entry_rows); + Err(failure) + } + Err(failure) => Err(failure.invalidating_engine_state()), + } +} + +/// The body of a decode step, which appends to `cache` as it goes. +#[allow(clippy::too_many_arguments)] +fn decode_step_appending( + weights: &ModelWeights, + ffn: &dyn FfnBackend, + cache: &mut KvCache, + token_id: u32, + backend: Option<&dyn larql_compute::ComputeBackend>, + hook: &mut dyn LayerHook, +) -> Result, EngineError> { + let num_layers = weights.num_layers; + let h_new = embed_tokens_pub(weights, &[token_id]); + let abs_position = cache.next_position; + // PLE inputs are per-token. Recompute for this single-token decode + // step rather than indexing a prefill-sized slab. Matches the + // recipe used by `vindex::kquant_forward::cached` and the GPU + // `layer_graph::generate` decode loop. + let ple_inputs = precompute_per_layer_inputs(weights, &h_new, &[token_id]); + let mut h_step = h_new; + for layer in 0..num_layers { + hook.on_pre_layer(layer, &h_step); + + let kv_entry = cache.layers[layer].as_ref(); + let (mut h_post_attn, new_kv) = run_attention_block_decode_step_backend( + larql_inference::WeightsView::dense(weights), + &h_step, + layer, + kv_entry, + abs_position, + backend, + ) + .ok_or_else(|| EngineError::BackendFailure { + details: format!("attention returned None during decode at layer {layer}"), + })?; + cache.layers[layer] = Some(new_kv); + cache.clip_layer(layer); + + hook.on_post_attention(layer, &mut h_post_attn); + + let mut h_out = crate::engines::layer_ffn_or_moe( + weights, + &h_post_attn, + layer, + ffn, + Some(ffn), + ple_inputs.get(layer), + ) + .map_err(EngineError::Execution)?; + + hook.on_post_layer(layer, &mut h_out); + h_step = h_out; + } + cache.next_position += 1; + Ok(h_step) +} diff --git a/crates/larql-kv/src/generation/kv_run/tests.rs b/crates/larql-kv/src/generation/kv_run/tests.rs new file mode 100644 index 000000000..26d5393da --- /dev/null +++ b/crates/larql-kv/src/generation/kv_run/tests.rs @@ -0,0 +1,140 @@ +//! The oracle's transaction, pinned. +//! +//! `kv_decode_step_run` appends each layer's K/V before the FFN can refuse, +//! so these fix the two answers it may give about what a failed step leaves +//! behind — the same pair `StandardEngine` gives, because it is the same +//! question about the same kind of state. + +use super::{cache_row_counts, kv_decode_step_run, kv_prefill_run}; +use larql_inference::ffn::WeightFfn; +use larql_inference::forward::hooks::NoopHook; +use larql_inference::ModelWeights; +use ndarray::Array2; + +// ── The oracle's decode transaction ────────────────────────────────── +// +// `kv_decode_step_run` appends each layer's K/V before the FFN gets the +// chance to refuse. These pin the two answers it can give about what that +// leaves behind — the same pair `StandardEngine` gives, because it is the +// same question about the same kind of state. + +/// A route that refuses every layer, so the FFN half of the step fails +/// after attention has already appended. +/// Fixture values for the refusal below — meaningless as numbers, named +/// so a failure message points at a deliberate fixture. +const OUT_OF_RANGE_EXPERT: u32 = 99; +const EXPERT_POPULATION: usize = 8; + +struct AlwaysRefuses; + +impl larql_inference::ffn::MoeExpertBackend for AlwaysRefuses { + fn forward_moe_seq( + &self, + _weights: &ModelWeights, + _layer: usize, + _h: &Array2, + _norm_offset: f32, + _eps: f32, + ) -> Result, larql_inference::ffn::MoeBackendError> { + Err(larql_inference::ffn::MoeBackendError::Bound( + larql_vindex::runtime::ExecutionError::ExpertOutOfRange { + expert: OUT_OF_RANGE_EXPERT, + population: EXPERT_POPULATION, + }, + )) + } + fn name(&self) -> &'static str { + "always-refuses" + } +} + +#[test] +fn a_refused_decode_truncates_the_cache_back_to_its_entry_length() { + let weights = larql_inference::test_utils::make_test_gemma4_moe_weights(); + let clean = WeightFfn { weights: &weights }; + let prompt = [0u32, 1, 2]; + let (_, mut cache) = kv_prefill_run( + larql_inference::WeightsView::dense(&weights), + &clean, + &prompt, + None, + None, + &mut NoopHook, + ) + .expect("clean prefill"); + let before: Vec> = cache_row_counts(&cache); + let position_before = cache.next_position; + + let route = AlwaysRefuses; + let refusing = larql_inference::ffn::MoeFfn::strict(&weights, &route); + let err = kv_decode_step_run(&weights, &refusing, &mut cache, 3, None, &mut NoopHook) + .expect_err("a refused step must not produce a hidden state"); + assert!(err.engine_state_is_retryable()); + assert_eq!( + cache_row_counts(&cache), + before, + "the refused step must leave the cache exactly as it found it" + ); + assert_eq!( + cache.next_position, position_before, + "position advances only on success" + ); + + // And the retried token computes what an untouched cache computes. + let retried = kv_decode_step_run(&weights, &clean, &mut cache, 3, None, &mut NoopHook) + .expect("the rewound cache must accept the retry"); + let (_, mut reference) = kv_prefill_run( + larql_inference::WeightsView::dense(&weights), + &clean, + &prompt, + None, + None, + &mut NoopHook, + ) + .expect("reference prefill"); + let baseline = kv_decode_step_run(&weights, &clean, &mut reference, 3, None, &mut NoopHook) + .expect("reference decode"); + assert_eq!( + retried.iter().map(|v| v.to_bits()).collect::>(), + baseline.iter().map(|v| v.to_bits()).collect::>(), + "a retry after a rewound refusal must be bit-identical to never having refused" + ); +} + +#[test] +fn a_refused_decode_on_a_full_window_reports_an_invalidated_cache() { + const WINDOW: usize = 2; + let weights = larql_inference::test_utils::make_test_gemma4_moe_weights(); + let clean = WeightFfn { weights: &weights }; + // Prompt longer than the window leaves every layer at the limit — the + // state in which the next append must evict, and eviction is not + // undoable by truncation. + let (_, mut cache) = kv_prefill_run( + larql_inference::WeightsView::dense(&weights), + &clean, + &[0u32, 1, 2], + Some(WINDOW), + None, + &mut NoopHook, + ) + .expect("windowed prefill"); + + let route = AlwaysRefuses; + let refusing = larql_inference::ffn::MoeFfn::strict(&weights, &route); + let err = kv_decode_step_run(&weights, &refusing, &mut cache, 3, None, &mut NoopHook) + .expect_err("must refuse"); + assert!( + matches!( + err, + larql_inference::kv_engine::EngineError::StateInvalidated { .. } + ), + "append-then-evict leaves the row count unchanged while the oldest row is \ + gone, so this must not be reported as an ordinary refusal: {err:?}" + ); + assert!(!err.engine_state_is_retryable()); + // The classification still reaches whoever has to act on it. + assert_eq!( + err.refusal_kind(), + Some(larql_execution::RefusalKind::BindingDefect) + ); +} diff --git a/crates/larql-kv/src/generation/kv_run/transaction.rs b/crates/larql-kv/src/generation/kv_run/transaction.rs new file mode 100644 index 000000000..64605e1d7 --- /dev/null +++ b/crates/larql-kv/src/generation/kv_run/transaction.rs @@ -0,0 +1,46 @@ +//! Rewinding a decode step that did not complete. +//! +//! `kv_decode_step_run` appends each layer's K/V before the FFN gets the +//! chance to refuse, so a step that stops partway leaves the cache holding a +//! token that produced no output. These three functions are the same +//! rewind-or-invalidate protocol `StandardEngine` uses, applied to the plain +//! [`KvCache`]: snapshot the lengths, decide whether truncation would be +//! honest, and either restore or say it could not. + +use crate::cache::KvCache; + +/// Logical row count per layer, `None` for layers that share another's K/V. +pub(super) fn cache_row_counts(cache: &KvCache) -> Vec> { + cache + .layers + .iter() + .map(|slot| slot.as_ref().map(|(k, _)| k.shape()[0])) + .collect() +} + +/// Whether truncating back to `entry_rows` would restore the exact cache the +/// step started from. +/// +/// Unbounded caches only ever append, so truncation is exact. A windowed cache +/// that reaches its limit drops its oldest row to make room, and that row is +/// gone; row count cannot see it, so the only sound test is whether every +/// layer had room to spare before the step began. +pub(super) fn decode_rewind_is_sound(cache: &KvCache, entry_rows: &[Option]) -> bool { + match cache.max_window { + None => true, + Some(w) => entry_rows.iter().flatten().all(|&rows| rows < w), + } +} + +/// Truncate every layer back to its recorded length. +pub(super) fn rewind_cache(cache: &mut KvCache, entry_rows: &[Option]) { + for (slot, rows) in cache.layers.iter_mut().zip(entry_rows) { + let (Some((k, v)), Some(rows)) = (slot.as_mut(), rows) else { + continue; + }; + if k.shape()[0] > *rows { + *k = k.slice(ndarray::s![..*rows, ..]).to_owned(); + *v = v.slice(ndarray::s![..*rows, ..]).to_owned(); + } + } +} diff --git a/crates/larql-kv/src/lib.rs b/crates/larql-kv/src/lib.rs index c5c066afe..e2167682b 100644 --- a/crates/larql-kv/src/lib.rs +++ b/crates/larql-kv/src/lib.rs @@ -34,12 +34,12 @@ pub use engines::markov_residual_codec; pub use engines::no_cache; pub use engines::standard; pub use engines::turbo_quant; -pub use engines::unlimited_context; +pub use engines::windowed_checkpoint; pub use engines::markov_residual::MarkovResidualEngine; pub use engines::no_cache::NoCacheEngine; pub use engines::standard::StandardEngine; -pub use engines::unlimited_context::UnlimitedContextEngine; +pub use engines::windowed_checkpoint::WindowedCheckpointEngine; // ─── Trait surface re-exported from larql-inference ────────────────────────── // @@ -72,7 +72,7 @@ pub enum EngineKind { MarkovResidual { window_size: Option, }, - UnlimitedContext { + WindowedCheckpoint { window_size: usize, }, TurboQuant { @@ -125,7 +125,7 @@ impl EngineKind { /// no-cache /// markov-rs /// markov-rs:window=1024 - /// unlimited-context:window=256 + /// windowed-checkpoint:window=256 /// turbo-quant:bits=3 /// tq4 /// apollo:layer=25,coef=8.0,top_k=12,bos=2 @@ -169,11 +169,20 @@ impl EngineKind { let window_size = params.get("window").and_then(|v| v.parse().ok()); Some(EngineKind::MarkovResidual { window_size }) } - "unlimited" | "unlimited-context" | "unlimited_context" => { - Some(EngineKind::UnlimitedContext { - window_size: get_usize("window", 512), - }) - } + // `unlimited-context` and friends are the pre-2026-08-03 names, + // kept so existing scripts and baselines keep parsing. The engine + // was renamed because the old name described a capability + // (arbitrarily long streams via archive + replay) while reading as + // a claim about attention — which is exactly the confusion that + // let issue #200 hide: it reported `window=N` and attended over + // everything. + "windowed-checkpoint" + | "windowed_checkpoint" + | "unlimited" + | "unlimited-context" + | "unlimited_context" => Some(EngineKind::WindowedCheckpoint { + window_size: get_usize("window", 512), + }), "turbo-quant" | "turbo_quant" | "turboquant" | "tq4" => Some(EngineKind::TurboQuant { bits: get_usize("bits", 4) as u8, }), @@ -278,7 +287,7 @@ impl EngineKind { EngineKind::Standard { .. } => "standard", EngineKind::NoCache => "no-cache", EngineKind::MarkovResidual { .. } => "markov-rs", - EngineKind::UnlimitedContext { .. } => "unlimited-context", + EngineKind::WindowedCheckpoint { .. } => "windowed-checkpoint", EngineKind::TurboQuant { .. } => "turbo-quant", EngineKind::Apollo { .. } => "apollo", EngineKind::BoundaryKv { .. } => "boundary-kv", @@ -307,7 +316,7 @@ impl EngineKind { "no-cache", "markov-rs", "markov-rs-codec", - "unlimited-context", + "windowed-checkpoint", "turbo-quant", "apollo", "boundary-kv", @@ -315,6 +324,46 @@ impl EngineKind { ] } + /// Specs the criterion microbenchmark (`benches/engine_decode.rs`) + /// runs, parameterised where a bare name would not build something + /// meaningful. Single source of truth so the bench cannot silently + /// drift behind the engine roster — pinned by + /// `bench_specs_cover_every_benchable_engine`. + /// + /// **Apollo is deliberately absent.** It is a [`RetrievalEngine`] + /// whose `prefill` fails closed with `RetrievalMiss` unless a + /// boundary store is attached, and the synthetic fixture has none. + /// Benching it there timed the error return, not the engine: it + /// reported ~65 ns against ~16 µs for `standard`, reading as a 250x + /// win in the criterion report. A meaningful Apollo number needs a + /// real store, which belongs in the CLI bench, not here. + pub fn bench_specs() -> &'static [&'static str] { + &[ + "standard", + "standard:window=4", + "no-cache", + "markov-rs", + "markov-rs:window=4", + "markov-rs-codec", + "windowed-checkpoint:window=4", + "turbo-quant:bits=4", + "turbo-quant:bits=3", + "boundary-kv:chunk_tokens=4", + "boundary-per-layer:layers=2", + ] + } + + /// Engines that [`Self::bench_specs`] intentionally omits, with the + /// reason. Keeping the exclusion explicit means a new engine can't + /// be dropped from the bench by simply never being added. + pub fn bench_excluded_names() -> &'static [(&'static str, &'static str)] { + &[( + "apollo", + "needs an attached boundary store; without one prefill returns \ + RetrievalMiss and the bench would time the error path", + )] + } + /// Build a boxed engine, dispatching compute through `backend`. pub fn build(self, backend: Box) -> AnyEngine { self.build_with_profiling(backend, false) @@ -324,7 +373,7 @@ impl EngineKind { /// /// Returns [`AnyEngine`] — the dispatch enum that wraps either a /// [`KvEngine`] (per-token K/V cache engines: standard, no_cache, - /// markov_residual, markov_residual_codec, unlimited_context, + /// markov_residual, markov_residual_codec, windowed_checkpoint, /// turbo_quant, boundary_kv, boundary_per_layer) or a /// [`RetrievalEngine`] (Apollo, future Mode 5). Callers branch /// once on the enum variant and stay in the variant-specific code @@ -355,8 +404,8 @@ impl EngineKind { markov_residual::MarkovResidualEngine::with_backend(window_size, backend) .with_profiling(profiling), )), - EngineKind::UnlimitedContext { window_size } => AnyEngine::Kv(Box::new( - unlimited_context::UnlimitedContextEngine::with_backend(window_size, backend) + EngineKind::WindowedCheckpoint { window_size } => AnyEngine::Kv(Box::new( + windowed_checkpoint::WindowedCheckpointEngine::with_backend(window_size, backend) .with_profiling(profiling), )), EngineKind::TurboQuant { bits } => AnyEngine::Kv(Box::new( @@ -452,11 +501,17 @@ mod tests { "failed to parse {name:?}" ); } - for name in &["unlimited", "unlimited-context", "unlimited_context"] { + for name in &[ + "windowed-checkpoint", + "windowed_checkpoint", + "unlimited", + "unlimited-context", + "unlimited_context", + ] { assert!( matches!( EngineKind::from_name(name), - Some(EngineKind::UnlimitedContext { .. }) + Some(EngineKind::WindowedCheckpoint { .. }) ), "failed to parse {name:?}" ); @@ -500,8 +555,8 @@ mod tests { other => panic!("expected MarkovResidual{{window=1024}}, got {other:?}"), } match EngineKind::from_name("unlimited-context:window=256") { - Some(EngineKind::UnlimitedContext { window_size: 256 }) => {} - other => panic!("expected UnlimitedContext{{window=256}}, got {other:?}"), + Some(EngineKind::WindowedCheckpoint { window_size: 256 }) => {} + other => panic!("expected WindowedCheckpoint{{window=256}}, got {other:?}"), } match EngineKind::from_name("turbo-quant:bits=3") { Some(EngineKind::TurboQuant { bits: 3 }) => {} @@ -763,7 +818,7 @@ mod compliance_tests { EngineKind::MarkovResidual { window_size: Some(32), }, - EngineKind::UnlimitedContext { window_size: 64 }, + EngineKind::WindowedCheckpoint { window_size: 64 }, EngineKind::TurboQuant { bits: 4 }, EngineKind::TurboQuant { bits: 3 }, EngineKind::Apollo { @@ -796,7 +851,7 @@ mod compliance_tests { "no-cache", "markov-rs", "markov-rs", - "unlimited-context", + "windowed-checkpoint", "turbo-quant", "turbo-quant", "apollo", @@ -861,7 +916,7 @@ mod compliance_tests { #[test] fn from_name_unknown_param_ignored_defaults_apply() { match EngineKind::from_name("unlimited-context:unknown=42") { - Some(EngineKind::UnlimitedContext { window_size: 512 }) => {} + Some(EngineKind::WindowedCheckpoint { window_size: 512 }) => {} other => panic!("unknown param should use default, got {other:?}"), } } @@ -921,6 +976,60 @@ mod compliance_tests { ); } + /// The criterion bench must cover every engine that can be benched + /// on the synthetic fixture. It previously listed 7 of 9 — the three + /// engines this PR touches most (`markov-rs-codec`, `boundary-kv`, + /// `boundary-per-layer`) had no microbenchmark at all, and nothing + /// failed when they were added. Now an engine is either in + /// `bench_specs` or explicitly in `bench_excluded_names` with a + /// reason; there is no third, silent option. + #[test] + fn bench_specs_cover_every_benchable_engine() { + let excluded: Vec<&str> = EngineKind::bench_excluded_names() + .iter() + .map(|(n, _)| *n) + .collect(); + + let benched: Vec<&'static str> = EngineKind::bench_specs() + .iter() + .map(|s| { + EngineKind::from_name(s) + .unwrap_or_else(|| panic!("bench_specs entry {s:?} no longer parses")) + .display_name() + }) + .collect(); + + for name in EngineKind::supported_names() { + if excluded.contains(name) { + assert!( + !benched.contains(name), + "{name:?} is listed as excluded but also appears in bench_specs" + ); + continue; + } + assert!( + benched.contains(name), + "engine {name:?} has no criterion bench arm — add a spec to \ + EngineKind::bench_specs, or name it in bench_excluded_names \ + with the reason it cannot be benched" + ); + } + } + + #[test] + fn bench_excluded_names_carry_a_reason_and_are_real_engines() { + for (name, reason) in EngineKind::bench_excluded_names() { + assert!( + EngineKind::supported_names().contains(name), + "bench_excluded_names lists {name:?}, which is not a supported engine" + ); + assert!( + reason.len() > 20, + "exclusion of {name:?} needs a real reason, got {reason:?}" + ); + } + } + #[test] fn from_name_all_engines_parseable() { let specs = [ @@ -930,7 +1039,11 @@ mod compliance_tests { ("no-cache", "no-cache"), ("none", "no-cache"), ("markov-rs", "markov-rs"), - ("unlimited-context", "unlimited-context"), + ("windowed-checkpoint", "windowed-checkpoint"), + // Pre-rename spellings still parse, and normalise to the new + // canonical name rather than echoing themselves back. + ("unlimited-context", "windowed-checkpoint"), + ("unlimited", "windowed-checkpoint"), ("turbo-quant", "turbo-quant"), ("tq3", "turbo-quant"), ("apollo", "apollo"), diff --git a/crates/larql-kv/tests/dispatch_parity.rs b/crates/larql-kv/tests/dispatch_parity.rs index 8ca2b16db..f9501cd7d 100644 --- a/crates/larql-kv/tests/dispatch_parity.rs +++ b/crates/larql-kv/tests/dispatch_parity.rs @@ -63,7 +63,8 @@ fn prefill_via_dispatch_matches_legacy_kv_prefill_run() { None, None, ) - .expect("prefill"); + .expect("prefill") + .expect("dispatch produced a result"); let (h_legacy, _cache) = kv_prefill_run( larql_inference::WeightsView::dense(&weights), &ffn, @@ -96,7 +97,8 @@ fn prefill_via_dispatch_windowed_matches_legacy() { window, None, ) - .expect("prefill"); + .expect("prefill") + .expect("dispatch produced a result"); let (h_legacy, _cache) = kv_prefill_run( larql_inference::WeightsView::dense(&weights), &ffn, @@ -128,7 +130,8 @@ fn decode_step_via_dispatch_matches_legacy_kv_decode_step_run() { None, None, ) - .unwrap(); + .unwrap() + .expect("dispatch produced a result"); let (_, mut cache) = kv_prefill_run( larql_inference::WeightsView::dense(&weights), &ffn, @@ -152,7 +155,8 @@ fn decode_step_via_dispatch_matches_legacy_kv_decode_step_run() { None, None, ) - .expect("decode step trait"); + .expect("decode step trait") + .expect("dispatch produced a result"); let h_legacy = kv_decode_step_run( &weights, @@ -185,7 +189,8 @@ fn multi_step_decode_via_dispatch_matches_legacy() { None, None, ) - .unwrap(); + .unwrap() + .expect("dispatch produced a result"); let (_, mut cache) = kv_prefill_run( larql_inference::WeightsView::dense(&weights), &ffn, @@ -209,7 +214,8 @@ fn multi_step_decode_via_dispatch_matches_legacy() { None, None, ) - .expect("decode trait"); + .expect("decode trait") + .expect("dispatch produced a result"); let h_legacy = kv_decode_step_run( &weights, &ffn, @@ -250,7 +256,8 @@ fn prefill_and_decode_via_dispatch_match_legacy_on_ple_arch() { None, None, ) - .expect("PLE prefill dispatch"); + .expect("PLE prefill dispatch") + .expect("dispatch produced a result"); let (h_legacy, mut cache) = kv_prefill_run( larql_inference::WeightsView::dense(&weights), &ffn, @@ -275,7 +282,8 @@ fn prefill_and_decode_via_dispatch_match_legacy_on_ple_arch() { None, None, ) - .expect("PLE decode dispatch"); + .expect("PLE decode dispatch") + .expect("dispatch produced a result"); let h_legacy = kv_decode_step_run( &weights, &ffn, @@ -304,7 +312,8 @@ fn prefill_and_decode_via_dispatch_async_match_legacy_on_ple_arch() { None, None, ) - .expect("PLE prefill async dispatch"); + .expect("PLE prefill async dispatch") + .expect("dispatch produced a result"); let (h_legacy, mut cache) = kv_prefill_run( larql_inference::WeightsView::dense(&weights), &ffn, @@ -329,7 +338,8 @@ fn prefill_and_decode_via_dispatch_async_match_legacy_on_ple_arch() { None, None, ) - .expect("PLE decode async dispatch"); + .expect("PLE decode async dispatch") + .expect("dispatch produced a result"); let h_legacy = kv_decode_step_run( &weights, &ffn, diff --git a/crates/larql-kv/tests/engine_ple_parity.rs b/crates/larql-kv/tests/engine_ple_parity.rs index ad560ae48..c46cb6707 100644 --- a/crates/larql-kv/tests/engine_ple_parity.rs +++ b/crates/larql-kv/tests/engine_ple_parity.rs @@ -148,7 +148,7 @@ fn exact_engines_match_legacy_on_ple_arch() { const UNLIMITED_CONTEXT_ACCUM_ORDER_REL_TOL: f32 = 1e-5; #[test] -fn unlimited_context_matches_legacy_within_accum_order_on_ple_arch() { +fn windowed_checkpoint_matches_legacy_within_accum_order_on_ple_arch() { let weights = make_synthetic_e2b_like_weights(); let (h_ref, decode_ref) = legacy_reference(&weights); let spec = format!("unlimited-context:window={NO_EVICTION_WINDOW}"); diff --git a/crates/larql-kv/tests/gpu_engine_parity/backend.rs b/crates/larql-kv/tests/gpu_engine_parity/backend.rs new file mode 100644 index 000000000..fba3404a2 --- /dev/null +++ b/crates/larql-kv/tests/gpu_engine_parity/backend.rs @@ -0,0 +1,85 @@ +//! What the backend reports about itself, plus the dense (no-vindex) +//! forward through a non-CPU backend. + +use larql_inference::ffn::WeightFfn; +use larql_inference::test_utils::make_test_weights; +use larql_inference::{cpu_engine_backend, default_engine_backend}; + +use super::support::{build, default_backend_is_host_delegating}; + +/// `per_layer_is_host_delegated` is what stops a bench row labelled +/// `[metal (GPU)]` from silently describing a CPU measurement. Pin that +/// Metal keeps answering `true` for as long as its per-layer methods +/// delegate to `CpuBackend` — when native per-layer kernels land, this +/// test is the one that should fail and be updated deliberately. +#[test] +fn host_delegation_flag_matches_the_backend_family() { + let backend = default_engine_backend(); + let name = backend.name().to_string(); + let delegated = + larql_compute::kv_dispatch::KvDispatch::per_layer_is_host_delegated(backend.as_ref()); + + if name.contains("metal") { + assert!( + delegated, + "MetalBackend's per-layer surface still forwards to CpuBackend; \ + reporting otherwise makes every windowed bench row a lie" + ); + } else { + assert!( + !delegated, + "backend {name:?} claims host delegation but is already the host" + ); + } + + // CPU is never "delegating" — it IS the host. + assert!( + !larql_compute::kv_dispatch::KvDispatch::per_layer_is_host_delegated( + cpu_engine_backend().as_ref() + ), + "CpuBackend must not report host delegation" + ); +} + +/// Guards the suite itself: if the default backend is host-delegating +/// (Metal present and the `gpu` feature on) then the numeric test really +/// did compare two different execution paths. Without this the whole +/// suite could pass on a CPU-only build while appearing to cover the GPU. +#[test] +fn reports_whether_this_run_actually_exercised_a_gpu_backend() { + let name = default_engine_backend().name().to_string(); + if default_backend_is_host_delegating() { + assert!( + name.contains("metal"), + "unexpected host-delegating backend {name:?}" + ); + } else { + eprintln!( + "[gpu_engine_parity] default backend is {name:?} — this run covered \ + the CPU path only. Build with `--features gpu` on macOS for GPU coverage." + ); + } +} + +/// The no-vindex path, through the GPU-capable backend. Cheap, but it is +/// the only coverage that the dense forward survives a non-CPU backend +/// being threaded through it. +#[test] +fn dense_prefill_and_decode_run_on_the_default_backend() { + let weights = make_test_weights(); + let ffn = WeightFfn { weights: &weights }; + for spec in ["standard", "standard:window=2", "markov-rs", "no-cache"] { + let mut engine = build(spec, true); + let h = engine + .prefill(&weights, &ffn, &[0u32, 1, 2]) + .unwrap_or_else(|e| panic!("{spec}: dense prefill failed: {e}")); + assert_eq!(h.shape(), &[1, weights.hidden_size]); + let d = engine + .decode_step(&weights, &ffn, 3) + .unwrap_or_else(|e| panic!("{spec}: dense decode failed: {e}")); + assert!( + d.iter().all(|v| v.is_finite()), + "{spec}: dense decode produced non-finite values" + ); + } +} diff --git a/crates/larql-kv/tests/gpu_engine_parity/gemma3_prefill_gap.rs b/crates/larql-kv/tests/gpu_engine_parity/gemma3_prefill_gap.rs new file mode 100644 index 000000000..5e07508d5 --- /dev/null +++ b/crates/larql-kv/tests/gpu_engine_parity/gemma3_prefill_gap.rs @@ -0,0 +1,96 @@ +//! Regression guard for a cross-backend divergence on the Gemma-3 arch. +//! +//! **History, because the diagnosis was wrong twice before it was right.** +//! Metal's batched prefill used to disagree with both the CPU and Metal's +//! own iterative path by 23-43% on the Gemma-3 fixture, growing with +//! depth, absent at one token, and not reproducing on a real Gemma 3 4B. +//! +//! It was first attributed to per-layer sliding-window attention. Wrong: +//! the fixture resolves to no window on either backend. It was then +//! attributed to a `head_dim` shape assumption, on the reasoning that +//! shape was the only surviving difference from the real model. Also +//! wrong, and that one was never tested — a sweep across +//! `head_dim ∈ {32, 64, 128, 256, 512}` diverged at *every* shape, +//! including the real model's 256. +//! +//! The cause was the fixture: it declared Gemma-3's QK-norm keys and +//! never populated the weights. Every consumer that resolves the weight +//! got `None`, and the two backends did not agree on what a +//! declared-but-absent QK-norm weight means. Real checkpoints always +//! carry those weights, which is why nothing shipped was ever affected. +//! +//! The fixture now supplies them (`larql-models::test_fixtures`), and +//! this file pins the agreement so neither half can regress: not the +//! fixture back to an impossible architecture, and not a backend into +//! disagreeing about the stage. + +use larql_inference::ffn::NullFfn; +use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; + +use super::support::{build, relative_l2, PREFILL_TOL}; + +/// Longest prompt swept. The divergence appeared from two positions on, +/// so anything past one token exercises it. +const MAX_PROMPT_LEN: usize = 6; + +/// The fixture must actually carry the architecture it claims. A +/// Gemma-3 fixture without QK-norm weights is not a Gemma-3 fixture, and +/// its absence is what let two backends disagree unnoticed. +#[test] +fn the_gemma3_fixture_carries_the_qk_norm_weights_it_declares() { + let weights = make_test_q4k_weights(); + let arch = &*weights.arch; + for layer in 0..weights.num_layers { + for (label, key) in [ + ("q_norm", arch.attn_q_norm_key(layer)), + ("k_norm", arch.attn_k_norm_key(layer)), + ] { + let Some(key) = key else { + panic!("layer {layer}: Gemma-3 must declare a {label} key"); + }; + assert!( + weights.vectors.contains_key(&key), + "layer {layer}: {label} is declared as {key:?} but absent — the \ + fixture is claiming an architecture it does not carry" + ); + } + } +} + +/// Metal's batched prefill and the CPU must agree on the Gemma-3 arch at +/// every prompt length. This is the comparison that used to fail. +#[test] +fn metal_batched_prefill_agrees_with_cpu_on_gemma3_arch() { + let weights = make_test_q4k_weights(); + let index = make_test_q4k_vindex(&weights); + + for n in 1..=MAX_PROMPT_LEN { + let prompt: Vec = (0..n as u32).collect(); + let mut gpu = build("standard", true); + let mut cpu = build("standard", false); + let hg = gpu + .prefill_quant( + &weights, + &NullFfn, + &index, + &prompt, + &*larql_compute::default_backend(), + ) + .expect("gpu prefill"); + let hc = cpu + .prefill_quant( + &weights, + &NullFfn, + &index, + &prompt, + &*larql_compute::cpu_backend(), + ) + .expect("cpu prefill"); + let rel = relative_l2(&hg, &hc); + assert!( + rel < PREFILL_TOL, + "prompt len {n}: backends diverged by relative L2 {rel:.3e} — this is \ + the regression this file exists to catch (it used to read ~4e-1)" + ); + } +} diff --git a/crates/larql-kv/tests/gpu_engine_parity/main.rs b/crates/larql-kv/tests/gpu_engine_parity/main.rs new file mode 100644 index 000000000..e8228ac71 --- /dev/null +++ b/crates/larql-kv/tests/gpu_engine_parity/main.rs @@ -0,0 +1,40 @@ +//! GPU-path coverage for the KV engines. +//! +//! Until this suite existed, **no automated test in `larql-kv` exercised +//! the GPU path for any engine**: every test, bench and pin built its +//! engines with `cpu_engine_backend()`, and the crate's `gpu` feature +//! gates only a dependency — the unit-test count is identical with and +//! without it. That is the worst place to have a coverage hole, because +//! GPU is where engines take a *structurally different* route: the +//! coarse whole-model pipeline, where K/V lives in the backend behind a +//! sentinel handle and the engine's own state policy never runs. +//! +//! These tests use `default_engine_backend()` — Metal on macOS under +//! `--features gpu`, CPU otherwise. They are written to be meaningful on +//! both: the assertions are about the *contract between engine and +//! backend*, which must hold on either. Where a claim is Metal-specific +//! it is guarded by the backend's own answer rather than by `cfg`, so +//! the suite cannot silently pass by compiling nothing. +//! +//! ## Layout +//! +//! - [`numeric`] — the two backends must produce the same hidden state, +//! and any residual difference must not compound across decode steps. +//! - [`shape`] — the dispatch shape an engine commits to, and the window +//! contract that decides whether it may take the fused path at all. +//! - [`backend`] — what the backend reports about itself, plus the dense +//! (no-vindex) forward through a non-CPU backend. +//! - [`gemma3_prefill_gap`] — regression guard for a cross-backend +//! divergence on the Gemma-3 arch, and for the fixture defect behind it. +//! +//! Run the GPU form with: +//! +//! ```sh +//! cargo test -p larql-kv --features gpu --test gpu_engine_parity +//! ``` + +mod backend; +mod gemma3_prefill_gap; +mod numeric; +mod shape; +mod support; diff --git a/crates/larql-kv/tests/gpu_engine_parity/numeric.rs b/crates/larql-kv/tests/gpu_engine_parity/numeric.rs new file mode 100644 index 000000000..6370f6613 --- /dev/null +++ b/crates/larql-kv/tests/gpu_engine_parity/numeric.rs @@ -0,0 +1,107 @@ +//! The GPU path must produce the same hidden state as the CPU path. + +use larql_inference::ffn::NullFfn; +use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights_silu}; + +use super::support::{all_specs, assert_close, build, COMPOUNDING_FACTOR, DECODE_TOL, PREFILL_TOL}; + +/// Prompt used for the cross-backend comparison. +const PROMPT: [u32; 4] = [0, 1, 2, 3]; +/// First decode token id; decode sweeps `DECODE_FIRST_TOKEN..+DECODE_STEPS`. +const DECODE_FIRST_TOKEN: u32 = 4; +/// Enough steps for compounding drift to show if it exists. +const DECODE_STEPS: u32 = 10; + +/// The property that matters: routing an engine through the default +/// (GPU-capable) backend must not change what the model says. +/// +/// Tolerance rather than bit-equality: the coarse pipeline fuses +/// operations the per-layer CPU path runs separately, so reassociation +/// is expected. A real divergence — wrong K/V rows, a shifted RoPE +/// position, a dropped window — moves the state far more than this. +/// +/// Uses the `tinymodel` fixture for the sweep across engines, and +/// [`super::gemma3_prefill_gap`] covers the Gemma-3 arch separately. The +/// Gemma fixture was excluded here while a cross-backend divergence was +/// open on it; that turned out to be the fixture declaring QK-norm +/// without supplying the weights, is fixed, and is now pinned by a +/// regression test rather than worked around. +#[test] +fn gpu_and_cpu_backends_agree_on_hidden_state() { + let weights = make_test_q4k_weights_silu(); + let index = make_test_q4k_vindex(&weights); + + for spec in all_specs() { + let mut gpu = build(spec, true); + let mut cpu = build(spec, false); + + let h_gpu = gpu + .prefill_quant( + &weights, + &NullFfn, + &index, + &PROMPT, + &*larql_compute::default_backend(), + ) + .unwrap_or_else(|e| panic!("{spec}: gpu prefill_quant failed: {e}")); + let h_cpu = cpu + .prefill_quant( + &weights, + &NullFfn, + &index, + &PROMPT, + &*larql_compute::cpu_backend(), + ) + .unwrap_or_else(|e| panic!("{spec}: cpu prefill_quant failed: {e}")); + + assert_eq!( + h_gpu.shape(), + h_cpu.shape(), + "{spec}: prefill shape differs" + ); + assert_close(&h_gpu, &h_cpu, PREFILL_TOL, spec, "prefill"); + + // Decode divergence must stay BOUNDED and must not compound. A + // fixed per-step difference is two Q4K kernels rounding + // differently; a growing one means the two caches are drifting + // apart, which is the failure that matters and which a + // single-step check cannot see. + let mut rels = Vec::new(); + for (step, tok) in (DECODE_FIRST_TOKEN..DECODE_FIRST_TOKEN + DECODE_STEPS).enumerate() { + let d_gpu = gpu + .decode_step_quant( + &weights, + &NullFfn, + &index, + tok, + &*larql_compute::default_backend(), + ) + .unwrap_or_else(|e| panic!("{spec}: gpu decode {step} failed: {e}")); + let d_cpu = cpu + .decode_step_quant( + &weights, + &NullFfn, + &index, + tok, + &*larql_compute::cpu_backend(), + ) + .unwrap_or_else(|e| panic!("{spec}: cpu decode {step} failed: {e}")); + rels.push(assert_close( + &d_gpu, + &d_cpu, + DECODE_TOL, + spec, + &format!("decode step {step}"), + )); + } + + let first = rels[0].max(f32::MIN_POSITIVE); + let last = *rels.last().expect("decode ran at least one step"); + assert!( + last <= first * COMPOUNDING_FACTOR, + "{spec}: cross-backend divergence is compounding across decode \ + steps ({first:.2e} → {last:.2e}); the two K/V caches are drifting \ + apart rather than just rounding differently" + ); + } +} diff --git a/crates/larql-kv/tests/gpu_engine_parity/shape.rs b/crates/larql-kv/tests/gpu_engine_parity/shape.rs new file mode 100644 index 000000000..ca34dd832 --- /dev/null +++ b/crates/larql-kv/tests/gpu_engine_parity/shape.rs @@ -0,0 +1,94 @@ +//! Which dispatch shape an engine commits to, and the window contract +//! that decides whether it may take the fused path at all. + +use larql_inference::ffn::NullFfn; +use larql_inference::kv_engine::DispatchPath; +use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; + +use super::support::{build, UNWINDOWED, WINDOW_BOUNDS_ATTENTION, WINDOW_BOUNDS_MEMORY}; + +const PROMPT: [u32; 4] = [0, 1, 2, 3]; + +fn prefill(spec: &str) -> larql_kv::AnyEngine { + let weights = make_test_q4k_weights(); + let index = make_test_q4k_vindex(&weights); + let mut engine = build(spec, true); + engine + .prefill_quant( + &weights, + &NullFfn, + &index, + &PROMPT, + &*larql_compute::default_backend(), + ) + .unwrap_or_else(|e| panic!("{spec}: prefill_quant failed: {e}")); + engine +} + +/// An unwindowed engine on a coarse-capable backend must actually take +/// the coarse path, and must account the K/V the backend holds for it. +/// Before `backend_resident_kv_bytes` existed these engines reported +/// `memory_bytes() == 0` here, which read as "this engine is free". +#[test] +fn unwindowed_engines_take_coarse_and_account_backend_resident_kv() { + for spec in UNWINDOWED { + let engine = prefill(spec); + if engine.dispatch_path() == Some(DispatchPath::Coarse) { + assert!( + engine.memory_bytes() > 0, + "{spec}: took the coarse path but reports zero K/V — the \ + backend-resident cache is going unaccounted again" + ); + } + } +} + +/// The window gate, asserted on the GPU-capable backend — but only for +/// the engines it applies to. +/// +/// "Windowed" is not one contract. An engine whose window bounds +/// *attention* has no cold tier, so reaching the window-less coarse +/// surface would let it answer from evicted rows it claims not to have; +/// it must decline. An engine whose window bounds *hot-tier memory* +/// keeps a cold tier and attends full history by design, so coarse is +/// legitimate for it. Asserting the first contract against every engine +/// fails on the second group for a reason that is not a bug. +#[test] +fn engines_whose_window_bounds_attention_decline_the_coarse_path() { + for spec in WINDOW_BOUNDS_ATTENTION { + let engine = prefill(spec); + assert_ne!( + engine.dispatch_path(), + Some(DispatchPath::Coarse), + "{spec}: this engine's window bounds attention and it has no cold \ + tier, so the window-less coarse surface would over-attend" + ); + } +} + +/// The complement, pinned so the distinction stays deliberate: a +/// cold-tier engine taking coarse is not a defect, and if one of these +/// ever starts declining coarse that is a silent 2-3x latency +/// regression nobody asked for. +#[test] +fn cold_tier_engines_may_take_the_coarse_path() { + for spec in WINDOW_BOUNDS_MEMORY { + let engine = prefill(spec); + assert!( + engine.dispatch_path().is_some(), + "{spec}: reported no dispatch shape after a successful prefill" + ); + } +} + +/// A shape is only meaningful once a prefill has chosen one. +#[test] +fn dispatch_path_is_none_before_any_prefill() { + for spec in UNWINDOWED.iter().chain(WINDOW_BOUNDS_ATTENTION.iter()) { + assert_eq!( + build(spec, true).dispatch_path(), + None, + "{spec}: reported a dispatch shape before prefill chose one" + ); + } +} diff --git a/crates/larql-kv/tests/gpu_engine_parity/support.rs b/crates/larql-kv/tests/gpu_engine_parity/support.rs new file mode 100644 index 000000000..90e971e56 --- /dev/null +++ b/crates/larql-kv/tests/gpu_engine_parity/support.rs @@ -0,0 +1,114 @@ +//! Shared fixtures, engine rosters and comparison helpers for the +//! GPU-parity suite. + +use larql_inference::{cpu_engine_backend, default_engine_backend}; +use larql_kv::EngineKind; + +/// Engines benched unwindowed. Apollo is excluded (its prefill fails +/// closed without an attached boundary store); `no-cache` holds no K/V +/// to place on either side of the comparison. +pub(crate) const UNWINDOWED: &[&str] = &[ + "standard", + "markov-rs", + "markov-rs-codec", + "boundary-per-layer:layers=2", +]; + +/// Windowed specs whose window bounds **attention**. These engines have +/// no cold tier — an evicted K/V row is gone — so attending past the +/// window would answer from data the engine claims not to have. They +/// must decline the window-less coarse surface. +pub(crate) const WINDOW_BOUNDS_ATTENTION: &[&str] = &["standard:window=2"]; + +/// Windowed specs whose window bounds **hot-tier memory only**. These +/// retain a cold tier and attend over full history by contract (see +/// `fix(kv,compute): the window must bound attention, not just storage`, +/// which drew exactly this distinction while fixing the engine that had +/// it wrong). They may legitimately take the coarse path. +pub(crate) const WINDOW_BOUNDS_MEMORY: &[&str] = &[ + "markov-rs:window=2", + "markov-rs-codec:window=2", + "boundary-per-layer:window=2,layers=2", +]; + +/// Every spec the numeric comparison sweeps. +pub(crate) fn all_specs() -> Vec<&'static str> { + UNWINDOWED + .iter() + .chain(WINDOW_BOUNDS_ATTENTION.iter()) + .chain(WINDOW_BOUNDS_MEMORY.iter()) + .copied() + .collect() +} + +/// Build an engine on either the default (GPU-capable) backend or the +/// explicit CPU one. +pub(crate) fn build(spec: &str, gpu: bool) -> larql_kv::AnyEngine { + let backend = if gpu { + default_engine_backend() + } else { + cpu_engine_backend() + }; + EngineKind::from_name(spec) + .unwrap_or_else(|| panic!("spec {spec:?} failed to parse")) + .build(backend) +} + +/// Whether the selected default backend implements the per-layer surface +/// by forwarding to the host — Metal today. +pub(crate) fn default_backend_is_host_delegating() -> bool { + larql_compute::kv_dispatch::KvDispatch::per_layer_is_host_delegated( + default_engine_backend().as_ref(), + ) +} + +/// Prefill runs a single fused pass on each side; agreement is tight +/// (measured ~1e-7 on the SWA-free fixture). +pub(crate) const PREFILL_TOL: f32 = 1e-4; + +/// Per-step decode. Metal's fused decode kernel and the CPU cached-decode +/// path dequantise and accumulate Q4K in different orders, which lands at +/// a stable 2.6e-3 - 3.6e-3 on this fixture for every engine. The bound +/// is set above that band; the compounding check is what actually guards +/// correctness. +pub(crate) const DECODE_TOL: f32 = 1e-2; + +/// How much the last decode step's divergence may exceed the first +/// before we call it drift rather than rounding. +pub(crate) const COMPOUNDING_FACTOR: f32 = 3.0; + +/// Assert two hidden states agree within `tol`. Returns the relative L2 +/// so callers can check its trend across steps. +pub(crate) fn assert_close( + a: &ndarray::Array2, + b: &ndarray::Array2, + tol: f32, + spec: &str, + stage: &str, +) -> f32 { + assert!( + a.iter().all(|v| v.is_finite()), + "{spec}: {stage} produced non-finite values on the gpu path" + ); + let rel = relative_l2(a, b); + assert!( + rel < tol, + "{spec}: {stage} diverged between backends — relative L2 {rel:.3e} (tol {tol:.0e})" + ); + rel +} + +/// Relative L2 of `a - b` against `b`'s norm. +pub(crate) fn relative_l2(a: &ndarray::Array2, b: &ndarray::Array2) -> f32 { + let num: f32 = a + .iter() + .zip(b.iter()) + .map(|(x, y)| (x - y) * (x - y)) + .sum::() + .sqrt(); + let den: f32 = b.iter().map(|y| y * y).sum::().sqrt().max(NORM_FLOOR); + num / den +} + +/// Guards against dividing by a zero-norm reference. +const NORM_FLOOR: f32 = 1e-6; diff --git a/crates/larql-kv/tests/strict_refusal/engine_state.rs b/crates/larql-kv/tests/strict_refusal/engine_state.rs new file mode 100644 index 000000000..3a2fd2f76 --- /dev/null +++ b/crates/larql-kv/tests/strict_refusal/engine_state.rs @@ -0,0 +1,223 @@ +//! What a refusal leaves behind, per engine. +//! +//! [`crate::engines`] answers "does the refusal stop the token". This answers +//! the question strict semantics raise next: **can the same engine be driven +//! again?** `Residency` refusals make "fix the cause and retry" a supported +//! workflow, so an engine that reported one while quietly keeping a +//! half-applied step would append the token twice on the retry. +//! +//! The answer differs by engine because what each treats as canonical +//! differs: +//! +//! ```text +//! residual-canonical markov-rs, markov-rs-codec, boundary-per-layer +//! → the step writes `stored` only after the last +//! fallible call; `hot_kv` is a droppable derivative +//! K/V-canonical standard, turbo-quant, unlimited-context +//! → the cache grows before the FFN can refuse, so the +//! step must undo the appends (truncate) to rewind +//! ``` +//! +//! Both routes end at the same guarantee, and the one place it cannot hold — +//! an unlimited-context stream that already archived a window — says so with +//! [`EngineError::StateInvalidated`] rather than pretending. + +use larql_execution::RefusalKind; +use larql_inference::ffn::MoeFfn; +use larql_inference::kv_engine::EngineError; +use larql_inference::test_utils::make_test_gemma4_moe_weights; +use larql_kv::EngineKind; + +use crate::engines::{Coverage, ALL}; +use crate::entry::{NEXT_TOKEN, PROMPT}; +use crate::routes::{ExecutingRoute, RefuseAfterFirstPass, RefusingRoute}; + +/// A refused decode step leaves every expert-routing engine usable, and says +/// so through `engine_state_is_retryable`. +#[test] +fn a_refused_decode_leaves_every_engine_retryable() { + let weights = make_test_gemma4_moe_weights(); + let mut checked = 0usize; + + for under_test in ALL.iter().filter(|e| e.coverage == Coverage::RoutesExperts) { + let label = under_test.label; + let clean = MoeFfn::strict(&weights, &ExecutingRoute); + let mut engine = under_test.engine(&weights); + engine + .prefill(&weights, &clean, &PROMPT) + .unwrap_or_else(|e| panic!("{label}: clean prefill: {e:?}")); + let before = engine.window_tokens(); + + let route = RefusingRoute::new(RefusalKind::Residency); + let refusing = MoeFfn::strict(&weights, &route); + let err = engine + .decode_step(&weights, &refusing, NEXT_TOKEN) + .expect_err("the refused step must not produce a hidden state"); + + assert_eq!( + err.refusal_kind(), + Some(RefusalKind::Residency), + "{label}: classification must survive" + ); + assert!( + err.engine_state_is_retryable(), + "{label}: this engine can rewind the step, so it must not report a dead \ + instance; got {err:?}" + ); + assert_eq!( + engine.window_tokens(), + before, + "{label}: the refused step must leave the cache exactly as it found it" + ); + engine + .decode_step(&weights, &clean, NEXT_TOKEN) + .unwrap_or_else(|e| panic!("{label}: a rewound engine must accept the retry: {e:?}")); + checked += 1; + } + assert!(checked >= 7, "engine coverage shrank ({checked} < 7)"); +} + +/// The retried token produces what an engine that never refused produces. +/// +/// The strongest form of the guarantee: not merely that a count came back, +/// but that the rewound engine computes the same answer. A rewind that +/// restored the shape while leaving the buffers shifted passes the check +/// above and fails this one. +/// +/// Driven on the *first* decode after prefill, where the residual engines' +/// `hot_kv` derivative is not yet populated on either side — so the two +/// engines are comparable bit-for-bit rather than one taking the cached path +/// and the other the recompute path. +#[test] +fn a_rewound_engine_decodes_the_retried_token_identically() { + let weights = make_test_gemma4_moe_weights(); + + for under_test in ALL.iter().filter(|e| e.coverage == Coverage::RoutesExperts) { + let label = under_test.label; + let clean = MoeFfn::strict(&weights, &ExecutingRoute); + + let mut retried = under_test.engine(&weights); + retried + .prefill(&weights, &clean, &PROMPT) + .unwrap_or_else(|e| panic!("{label}: prefill A: {e:?}")); + let route = RefusingRoute::new(RefusalKind::Residency); + let refusing = MoeFfn::strict(&weights, &route); + retried + .decode_step(&weights, &refusing, NEXT_TOKEN) + .expect_err("must refuse"); + let after_retry = retried + .decode_step(&weights, &clean, NEXT_TOKEN) + .unwrap_or_else(|e| panic!("{label}: retry: {e:?}")); + + let mut reference = under_test.engine(&weights); + reference + .prefill(&weights, &clean, &PROMPT) + .unwrap_or_else(|e| panic!("{label}: prefill B: {e:?}")); + let baseline = reference + .decode_step(&weights, &clean, NEXT_TOKEN) + .unwrap_or_else(|e| panic!("{label}: reference decode: {e:?}")); + + assert_eq!( + after_retry.iter().map(|v| v.to_bits()).collect::>(), + baseline.iter().map(|v| v.to_bits()).collect::>(), + "{label}: a retry after a rewound refusal must be bit-identical to never \ + having refused — anything less means the rewind restored the shape, not \ + the state" + ); + assert_eq!( + retried.window_tokens(), + reference.window_tokens(), + "{label}" + ); + } +} + +/// A refused prefill leaves an earlier prefill's state in place. +/// +/// The engines build their store into locals and install it only on success, +/// so this is a consequence of statement order — which is exactly why it is +/// worth pinning: it is easy to lose in a refactor and silent when lost. +/// +/// `boundary-kv` is excluded: it drives its inner engine one chunk at a time +/// and archives each chunk's frame as it goes, so a mid-prompt refusal has +/// already emitted frames. That is a real gap in its own right, tracked +/// separately from this one. +#[test] +fn a_refused_prefill_leaves_an_earlier_prefill_intact() { + let weights = make_test_gemma4_moe_weights(); + + for under_test in ALL + .iter() + .filter(|e| e.coverage == Coverage::RoutesExperts && e.label != "boundary-kv") + { + let label = under_test.label; + let clean = MoeFfn::strict(&weights, &ExecutingRoute); + let mut engine = under_test.engine(&weights); + engine + .prefill(&weights, &clean, &PROMPT) + .unwrap_or_else(|e| panic!("{label}: first prefill: {e:?}")); + let before = engine.window_tokens(); + + let route = RefusingRoute::new(RefusalKind::Unsupported); + let refusing = MoeFfn::strict(&weights, &route); + let err = engine + .prefill(&weights, &refusing, &PROMPT) + .expect_err("the second prefill must refuse"); + assert!( + err.engine_state_is_retryable(), + "{label}: a prefill that assigned nothing cannot have invalidated anything" + ); + assert_eq!( + engine.window_tokens(), + before, + "{label}: a refused prefill must not disturb the cache already in place" + ); + engine + .decode_step(&weights, &clean, NEXT_TOKEN) + .unwrap_or_else(|e| panic!("{label}: the surviving cache must still decode: {e:?}")); + } +} + +/// The one case that cannot be rewound says so. +/// +/// `unlimited-context` archives a window's tokens and saves its boundary +/// checkpoint when the window fills, and neither is undoable. A prompt long +/// enough to close a window before the refusal therefore leaves a stream the +/// engine cannot complete — so it reports [`EngineError::StateInvalidated`] +/// rather than a retryable refusal, and a caller told "recoverable" does not +/// re-drive it into a duplicated window. +#[test] +fn windowed_checkpoint_invalidates_once_a_window_has_closed() { + /// One token per window, so the prompt below closes windows as it goes. + const WINDOW: usize = 1; + let weights = make_test_gemma4_moe_weights(); + let kind = EngineKind::WindowedCheckpoint { + window_size: WINDOW, + }; + let mut engine = kind.build(larql_inference::cpu_engine_backend()); + + // Serve the first window, then refuse: the refusal has to arrive *after* + // a close for this to be the unrewindable case rather than the ordinary + // one. + let route = RefuseAfterFirstPass::new(RefusalKind::Residency); + let refusing = MoeFfn::strict(&weights, &route); + let err = engine + .prefill(&weights, &refusing, &PROMPT) + .expect_err("must refuse"); + + assert!( + matches!(err, EngineError::StateInvalidated { .. }), + "an archived window cannot be un-archived, so this must not be reported as an \ + ordinary refusal: {err:?}" + ); + assert!(!err.engine_state_is_retryable()); + assert!( + !err.is_recoverable(), + "a caller must not be told to fix the residency and carry on with a stream \ + that is missing a window" + ); + // The wrapper costs the cause neither its classification nor the fact + // that the operation itself could have succeeded elsewhere. + assert_eq!(err.refusal_kind(), Some(RefusalKind::Residency)); + assert!(err.operation_is_recoverable()); +} diff --git a/crates/larql-kv/tests/strict_refusal/engines.rs b/crates/larql-kv/tests/strict_refusal/engines.rs new file mode 100644 index 000000000..5203ecb04 --- /dev/null +++ b/crates/larql-kv/tests/strict_refusal/engines.rs @@ -0,0 +1,198 @@ +//! Every engine in [`EngineKind`], and what a strict refusal does to it. +//! +//! The gate in [`crate::gate`] sweeps `StandardEngine`'s entry points because +//! that is where the dispatch ring terminates. This module sweeps the other +//! axis — the engine — because five of them reach their FFN through +//! `larql_kv::engines::layer_ffn_or_moe` instead, and that copy used to log a +//! refusal and return the dense half. +//! +//! ## Two populations, and the difference is not a detail +//! +//! ```text +//! RoutesExperts dispatches experts through the FFN hook → must refuse a +//! refusing route, and serve an executing one +//! NoExpertSeam has no FfnBackend seam in its forward at all → must refuse +//! the *architecture*, before any route is consulted +//! ``` +//! +//! The second population exists because "cannot dispatch experts" is not the +//! same failure as "an expert was missing". An engine whose forward has no +//! hook runs the dense half of every layer and returns an apparently valid +//! answer — a *different model* wearing the same answer shape, which nothing +//! downstream can detect. So those engines refuse the model up front, with +//! `RefusalKind::Unsupported`: the operands are fine, this executor cannot +//! serve them, pick another. +//! +//! `apollo` is the only member. Its forward (`forward_from_layer` / +//! `forward_raw_logits`) lives in `larql-compute` *below* the `FfnBackend` +//! seam and builds its own dense `ViewFfn`, so no caller-supplied backend can +//! reach it; giving it real dispatch is a change to the forward, not to the +//! engine. `no-cache` was in this population until +//! `larql_kv::generation::kv_prefill_run` — which is also the oracle the +//! dispatch ring is compared against — gained the hook. + +use larql_inference::ffn::MoeFfn; +use larql_inference::kv_engine::EngineError; +use larql_inference::model::ModelWeights; +use larql_kv::markov_residual_codec::ColdResidualCodec; +use larql_kv::EngineKind; +use ndarray::Array2; + +use crate::entry::{NEXT_TOKEN, PROMPT}; +use crate::routes::ExecutingRoute; + +/// Sliding window for the windowed variants: wide enough that a +/// three-token prompt plus one decode step never reaches the limit. +/// +/// Deliberately *not* at capacity. A windowed cache that must evict to make +/// room cannot be rewound at all — append-then-drop leaves the row count +/// unchanged while the oldest row is gone — and that case has its own pin in +/// [`crate::outcomes`]. Mixing it in here would make the whole sweep assert +/// the eviction rule instead of the refusal rule. +const WINDOW: usize = 8; +/// `boundary-kv` chunks its prefill; one chunk per prompt token here. +const CHUNK_TOKENS: usize = 1; +const SEQUENCE_ID: &str = "strict-refusal-gate"; +/// Apollo's injection parameters. Values are immaterial — the engine is in +/// the table to be *excluded* by evidence, not to be exercised. +const APOLLO_LAYER: usize = 1; +const APOLLO_COEFFICIENT: f32 = 1.0; +const APOLLO_TOP_K: usize = 1; + +/// Whether an engine dispatches experts at all. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Coverage { + /// Reaches `FfnBackend::forward_moe_full_layer`, so a strict route can + /// refuse through it and the refusal must terminate the operation. + RoutesExperts, + /// Has no `FfnBackend` seam in its forward, so it cannot dispatch + /// experts at all. Must refuse a hybrid-MoE architecture outright rather + /// than answer it densely — see this module's header. + NoExpertSeam, +} + +pub struct EngineUnderTest { + pub label: &'static str, + pub coverage: Coverage, + build: fn(&ModelWeights) -> EngineKind, +} + +impl EngineUnderTest { + pub fn engine(&self, weights: &ModelWeights) -> larql_kv::AnyEngine { + (self.build)(weights).build(larql_inference::cpu_engine_backend()) + } +} + +/// Every `EngineKind` variant, exactly once. +/// +/// Written as a table rather than a list of hand-rolled tests so a new +/// variant that nobody classified shows up as a missing row here — the +/// arity assertion at the end of each sweep is what makes that fail loudly. +pub const ALL: [EngineUnderTest; 9] = [ + EngineUnderTest { + label: "standard", + coverage: Coverage::RoutesExperts, + build: |_| EngineKind::Standard { window_size: None }, + }, + EngineUnderTest { + label: "standard:windowed", + coverage: Coverage::RoutesExperts, + build: |_| EngineKind::Standard { + window_size: Some(WINDOW), + }, + }, + EngineUnderTest { + label: "markov-rs", + coverage: Coverage::RoutesExperts, + build: |_| EngineKind::MarkovResidual { window_size: None }, + }, + EngineUnderTest { + label: "markov-rs-codec", + coverage: Coverage::RoutesExperts, + build: |_| EngineKind::MarkovResidualCodec { + window_size: None, + codec: ColdResidualCodec::Bf16, + }, + }, + EngineUnderTest { + label: "turbo-quant", + coverage: Coverage::RoutesExperts, + build: |_| EngineKind::TurboQuant { bits: 4 }, + }, + EngineUnderTest { + label: "unlimited-context", + coverage: Coverage::RoutesExperts, + build: |_| EngineKind::WindowedCheckpoint { + window_size: WINDOW, + }, + }, + EngineUnderTest { + label: "boundary-per-layer", + coverage: Coverage::RoutesExperts, + build: |w| EngineKind::BoundaryPerLayer { + window_size: None, + num_layers: w.num_layers, + }, + }, + EngineUnderTest { + label: "boundary-kv", + coverage: Coverage::RoutesExperts, + build: |_| EngineKind::BoundaryKv { + window_size: None, + chunk_tokens: CHUNK_TOKENS, + sequence_id: SEQUENCE_ID.to_string(), + }, + }, + EngineUnderTest { + label: "no-cache", + coverage: Coverage::RoutesExperts, + build: |_| EngineKind::NoCache, + }, +]; + +/// Apollo, kept out of [`ALL`] because it cannot be driven without a boundary +/// store — every other question the sweep asks would be answered by its state +/// checks rather than by its routing. The architecture refusal below fires +/// *before* those checks, which is the whole point: it costs nothing and +/// mutates nothing. +pub const APOLLO: EngineUnderTest = EngineUnderTest { + label: "apollo", + coverage: Coverage::NoExpertSeam, + build: |_| EngineKind::Apollo { + injection_layer: APOLLO_LAYER, + inject_coefficient: APOLLO_COEFFICIENT, + top_k: APOLLO_TOP_K, + bos_token_id: None, + }, +}; + +/// Which operation to drive. Both, because prefill and decode reach the FFN +/// through different bodies in most of these engines. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Op { + Prefill, + Decode, +} + +impl Op { + pub const ALL: [Self; 2] = [Self::Prefill, Self::Decode]; +} + +/// Drive `op` on a fresh engine with `ffn`. +/// +/// A decode prefills first with a route that executes, so a refusal observed +/// afterwards can only have come from the decode step. +pub fn drive( + under_test: &EngineUnderTest, + op: Op, + weights: &ModelWeights, + ffn: &dyn larql_inference::ffn::FfnBackend, +) -> Result, EngineError> { + let mut engine = under_test.engine(weights); + if op == Op::Decode { + let clean = MoeFfn::strict(weights, &ExecutingRoute); + engine.prefill(weights, &clean, &PROMPT)?; + return engine.decode_step(weights, ffn, NEXT_TOKEN); + } + engine.prefill(weights, ffn, &PROMPT) +} diff --git a/crates/larql-kv/tests/strict_refusal/engines_gate.rs b/crates/larql-kv/tests/strict_refusal/engines_gate.rs new file mode 100644 index 000000000..7affefbe5 --- /dev/null +++ b/crates/larql-kv/tests/strict_refusal/engines_gate.rs @@ -0,0 +1,154 @@ +//! The gate, swept across every engine rather than every entry point. +//! +//! [`crate::gate`] fixes the engine (`StandardEngine`) and varies the entry +//! point; this fixes the operation and varies the engine. Both axes matter +//! because a refusal has to survive two different rings — the dispatch +//! helpers and `larql_kv::engines::layer_ffn_or_moe` — and either can be +//! correct while the other erases. + +use crate::engines::{drive, Coverage, Op, ALL, APOLLO}; +use larql_execution::RefusalKind; +use larql_inference::ffn::MoeFfn; +use larql_inference::test_utils::make_test_gemma4_moe_weights; + +use crate::routes::{ExecutingRoute, RefusingRoute}; + +/// Strict + a refusing route → `Err`, carrying the refusal's own kind, in +/// every engine that dispatches experts, on both operations. +/// +/// This is the closing condition for "the other five engines". Before it, +/// `layer_ffn_or_moe` logged the refusal to stderr and returned the dense +/// half, so a strict route through markov-rs, markov-rs-codec, +/// turbo-quant, unlimited-context or boundary-per-layer produced a +/// complete-looking hidden state whose experts never ran — and the engine +/// had no way to know. +#[test] +fn strict_refusal_terminates_in_every_expert_routing_engine() { + let weights = make_test_gemma4_moe_weights(); + assert!( + weights.arch.is_hybrid_moe(), + "fixture must take the MoE branch or this test proves nothing" + ); + + let mut checked = 0usize; + let mut engines = 0usize; + for under_test in ALL.iter().filter(|e| e.coverage == Coverage::RoutesExperts) { + engines += 1; + for op in Op::ALL { + for kind in RefusalKind::ALL { + let route = RefusingRoute::new(kind); + let ffn = MoeFfn::strict(&weights, &route); + let label = under_test.label; + + let err = match drive(under_test, op, &weights, &ffn) { + Ok(hidden) => panic!( + "{label}/{op:?} with a {kind} refusal returned Ok({:?}) — a strict \ + route produced output for a layer whose experts never ran", + hidden.shape() + ), + Err(e) => e, + }; + assert_eq!( + err.refusal_kind(), + Some(kind), + "{label}/{op:?}: the refusal's own classification must survive to the \ + engine — flattening it here sends someone to repair the wrong thing; \ + got {err:?}" + ); + checked += 1; + } + } + } + assert_eq!( + checked, + engines * Op::ALL.len() * RefusalKind::ALL.len(), + "coverage shrank — no silent caps" + ); + assert!( + engines >= 7, + "expert-routing engine coverage shrank ({engines} < 7)" + ); +} + +/// The same engines serve normally when the route executes. +/// +/// Without this the gate above is satisfiable by an engine that refuses +/// everything, which is a different bug wearing the same result. +#[test] +fn an_executing_route_serves_in_every_expert_routing_engine() { + let weights = make_test_gemma4_moe_weights(); + for under_test in ALL.iter().filter(|e| e.coverage == Coverage::RoutesExperts) { + for op in Op::ALL { + let ffn = MoeFfn::strict(&weights, &ExecutingRoute); + let hidden = drive(under_test, op, &weights, &ffn).unwrap_or_else(|e| { + panic!( + "{}/{op:?}: an executing route must serve: {e:?}", + under_test.label + ) + }); + assert_eq!(hidden.shape(), &[1, weights.hidden_size]); + assert!( + ffn.all_experts_executed() && ffn.refusal().is_none(), + "{}/{op:?}: a clean run must record nothing", + under_test.label + ); + } + } +} + +/// An engine with no expert seam refuses the **architecture**, before any +/// route is consulted. +/// +/// This is the correction for the worse failure. Such an engine used to run +/// the dense half of every layer and return an apparently valid answer — a +/// different model wearing the same answer shape, which nothing downstream +/// could detect. `Unsupported` is the honest classification: the operands are +/// present and well-formed, this executor has no bound kernel for them, and +/// the response is to pick another. +/// +/// Asserted with a route that *executes*, so the refusal cannot be coming +/// from the route. It is the engine declining the model. +#[test] +fn an_engine_with_no_expert_seam_refuses_the_architecture_itself() { + let weights = make_test_gemma4_moe_weights(); + let no_seam: Vec<_> = ALL + .iter() + .filter(|e| e.coverage == Coverage::NoExpertSeam) + .chain(std::iter::once(&APOLLO)) + .collect(); + + let mut checked = 0usize; + for under_test in &no_seam { + let label = under_test.label; + for op in Op::ALL { + // An *executing* route, so the refusal cannot be coming from the + // route — it is the engine declining the model. + let ffn = MoeFfn::strict(&weights, &ExecutingRoute); + let err = match drive(under_test, op, &weights, &ffn) { + Ok(hidden) => panic!( + "{label}/{op:?} returned Ok({:?}) — a forward with no expert seam \ + answered for a model whose weights declare routed experts", + hidden.shape() + ), + Err(e) => e, + }; + assert_eq!( + err.refusal_kind(), + Some(RefusalKind::Unsupported), + "{label}/{op:?}: the operands are fine and this executor cannot serve \ + them, which is exactly Unsupported; got {err:?}" + ); + assert!( + err.engine_state_is_retryable(), + "{label}/{op:?}: refusing before any forward work mutates nothing" + ); + assert!( + ffn.refusal().is_none(), + "{label}/{op:?}: the route was never consulted, so nothing may be recorded" + ); + checked += 1; + } + } + assert_eq!(no_seam.len(), 1, "no-expert-seam population changed"); + assert_eq!(checked, no_seam.len() * Op::ALL.len(), "coverage shrank"); +} diff --git a/crates/larql-kv/tests/strict_refusal/entry.rs b/crates/larql-kv/tests/strict_refusal/entry.rs new file mode 100644 index 000000000..a9c15c970 --- /dev/null +++ b/crates/larql-kv/tests/strict_refusal/entry.rs @@ -0,0 +1,136 @@ +//! Every `StandardEngine` entry point that carries a caller's FFN into the +//! dispatch ring, and one way to drive them all. +//! +//! Enumerated rather than written out per test so the gate can sweep them +//! exhaustively — the failure it guards against is precisely one path being +//! missed, which is what an ad-hoc list of hand-written cases produces. + +use larql_inference::ffn::{FfnBackend, MoeFfn}; +use larql_inference::kv_engine::EngineError; +use larql_inference::model::ModelWeights; +use larql_inference::{AsyncComputeBackend, KvEngine}; +use larql_kv::StandardEngine; +use ndarray::Array2; + +use crate::routes::ExecutingRoute; + +/// Prompt for every entry point. Three tokens: enough that prefill caches a +/// distinguishable number of rows, small enough to stay fast. +pub const PROMPT: [u32; 3] = [0, 1, 2]; +/// The token a decode step is driven with. +pub const NEXT_TOKEN: u32 = 3; + +/// Synthetic hidden-state values for the `from_hidden` entry points. Any finite +/// spread works; these are a deterministic ramp so a failure reproduces. +const HIDDEN_STRIDE_ROW: usize = 7; +const HIDDEN_MODULUS: usize = 13; +const HIDDEN_SCALE: f32 = 0.01; +const HIDDEN_OFFSET: f32 = -0.06; + +/// A `StandardEngine` route that reaches the widened dispatch helpers. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Entry { + /// `prefill` over the sync slot → `kv_prefill_via_dispatch`. + PrefillSync, + /// `prefill` over the async slot → `kv_prefill_via_dispatch_async`. + PrefillAsync, + /// `prefill_from_hidden` → `kv_prefill_from_hidden_via_dispatch`. + PrefillFromHidden, + /// `prefill_from_hidden` over the async slot. + PrefillFromHiddenAsync, + /// `prefill_resident` → the same `do_prefill` body with `index: Some`. + PrefillResident, + /// `decode_step` over the sync slot → `kv_decode_step_via_dispatch`. + DecodeSync, + /// `decode_step` over the async slot. + DecodeAsync, + /// `decode_step_resident` → the same `do_decode_step` body. + DecodeResident, +} + +impl Entry { + pub const ALL: [Self; 8] = [ + Self::PrefillSync, + Self::PrefillAsync, + Self::PrefillFromHidden, + Self::PrefillFromHiddenAsync, + Self::PrefillResident, + Self::DecodeSync, + Self::DecodeAsync, + Self::DecodeResident, + ]; + + fn is_decode(self) -> bool { + matches!( + self, + Self::DecodeSync | Self::DecodeAsync | Self::DecodeResident + ) + } + + fn uses_async_slot(self) -> bool { + matches!( + self, + Self::PrefillAsync | Self::PrefillFromHiddenAsync | Self::DecodeAsync + ) + } + + fn engine(self) -> StandardEngine { + if self.uses_async_slot() { + let backend: Box = Box::new(larql_compute::CpuBackend); + StandardEngine::with_async_backend(None, backend) + } else { + StandardEngine::new(None) + } + } +} + +/// A structurally valid index for the resident entry points. +/// +/// The CPU backend ignores it with the Q4K-direct flags off (their default in +/// tests), so it exists only to satisfy the signature. The resident routes +/// differ from the plain ones exactly by threading `Some(index)` through +/// `do_prefill` / `do_decode_step`, which is what is under test. +fn empty_index(weights: &ModelWeights) -> larql_vindex::VectorIndex { + larql_vindex::VectorIndex::new( + vec![None; weights.num_layers], + vec![None; weights.num_layers], + weights.num_layers, + weights.hidden_size, + ) +} + +fn synthetic_hidden(weights: &ModelWeights) -> Array2 { + Array2::from_shape_fn((PROMPT.len(), weights.hidden_size), |(row, col)| { + ((row * HIDDEN_STRIDE_ROW + col) % HIDDEN_MODULUS) as f32 * HIDDEN_SCALE + HIDDEN_OFFSET + }) +} + +/// Drive one entry point with `ffn`. +/// +/// Decode entries prefill first with a route that executes, so a refusal +/// observed afterwards can only have come from the decode step itself. +pub fn drive( + entry: Entry, + weights: &ModelWeights, + ffn: &dyn FfnBackend, +) -> Result, EngineError> { + let index = empty_index(weights); + let mut engine = entry.engine(); + + if entry.is_decode() { + let clean = MoeFfn::strict(weights, &ExecutingRoute); + engine + .prefill(weights, &clean, &PROMPT) + .expect("setup prefill must succeed before the decode step under test"); + } + + match entry { + Entry::PrefillSync | Entry::PrefillAsync => engine.prefill(weights, ffn, &PROMPT), + Entry::PrefillFromHidden | Entry::PrefillFromHiddenAsync => { + engine.prefill_from_hidden(weights, ffn, &synthetic_hidden(weights)) + } + Entry::PrefillResident => engine.prefill_resident(weights, ffn, &index, &PROMPT), + Entry::DecodeSync | Entry::DecodeAsync => engine.decode_step(weights, ffn, NEXT_TOKEN), + Entry::DecodeResident => engine.decode_step_resident(weights, ffn, &index, NEXT_TOKEN), + } +} diff --git a/crates/larql-kv/tests/strict_refusal/gate.rs b/crates/larql-kv/tests/strict_refusal/gate.rs new file mode 100644 index 000000000..5f415fbe6 --- /dev/null +++ b/crates/larql-kv/tests/strict_refusal/gate.rs @@ -0,0 +1,107 @@ +//! The merge gate: strict decode cannot return `Ok` after a refusal, on any +//! entry point, for any classification. + +use larql_execution::RefusalKind; +use larql_inference::ffn::MoeFfn; +use larql_inference::kv_engine::EngineError; +use larql_inference::test_utils::make_test_gemma4_moe_weights; + +use crate::entry::{drive, Entry}; +use crate::routes::RefusingRoute; + +/// Strict + a refusing route → `Err`, carrying the refusal's own kind, on +/// every entry point and every classification. +/// +/// This is what the dispatch ring exists to make true. Before it, a refusal +/// reaching `ffn_or_moe_layer` was logged and the dense half of the layer was +/// returned — so a strict route produced a complete-looking hidden state with +/// an expert contribution of zero, and the engine had no way to know. +/// +/// The sweep is exhaustive over both axes rather than sampled, and asserts its +/// own arity at the end, because a coverage gap here reappears as a plausible +/// token somewhere else entirely. +#[test] +fn strict_refusal_terminates_every_entry_point_with_its_original_kind() { + let weights = make_test_gemma4_moe_weights(); + assert!( + weights.arch.is_hybrid_moe(), + "fixture must take the MoE branch or this test proves nothing" + ); + + let mut checked = 0usize; + for entry in Entry::ALL { + for kind in RefusalKind::ALL { + let route = RefusingRoute::new(kind); + let ffn = MoeFfn::strict(&weights, &route); + + let err = match drive(entry, &weights, &ffn) { + Ok(hidden) => panic!( + "{entry:?} with a {kind} refusal returned Ok({:?}) — a strict route \ + produced output for a layer whose experts never ran", + hidden.shape() + ), + Err(e) => e, + }; + assert!( + matches!(err, EngineError::Execution(_)), + "{entry:?}/{kind}: a refusal must terminate as EngineError::Execution, \ + not as some other failure that happens to be an error; got {err:?}" + ); + assert_eq!( + err.refusal_kind(), + Some(kind), + "{entry:?}: the refusal's own classification must survive to the engine \ + — flattening it here sends someone to repair the wrong thing" + ); + checked += 1; + } + } + assert_eq!( + checked, + Entry::ALL.len() * RefusalKind::ALL.len(), + "coverage shrank — no silent caps" + ); +} + +/// The refusal's message survives with its kind. +/// +/// The kind says which of three responses is needed; the message says which +/// expert, which bank, which operand. An engine-level error that kept only the +/// first would classify correctly and leave nobody able to act. +#[test] +fn the_concrete_message_survives_to_the_engine() { + let weights = make_test_gemma4_moe_weights(); + let route = RefusingRoute::new(RefusalKind::Residency); + let ffn = MoeFfn::strict(&weights, &route); + let err = drive(Entry::PrefillSync, &weights, &ffn).expect_err("strict must refuse"); + + let rendered = err.to_string(); + assert!( + rendered.contains("residency"), + "the classification must be legible in the message: {rendered}" + ); + assert!( + rendered.contains("not resident"), + "the route's own words must survive the boundary: {rendered}" + ); +} + +/// A `BindingDefect` is not recoverable; `Residency` and `Unsupported` are. +/// +/// Harnesses route on `is_recoverable`, so an `Execution` error answering +/// uniformly would either abort a sweep on a normal shard miss or sweep a +/// broken artifact into a coverage deficit. +#[test] +fn execution_recoverability_follows_the_refusal_kind() { + let weights = make_test_gemma4_moe_weights(); + for kind in RefusalKind::ALL { + let route = RefusingRoute::new(kind); + let ffn = MoeFfn::strict(&weights, &route); + let err = drive(Entry::PrefillSync, &weights, &ffn).expect_err("strict must refuse"); + assert_eq!( + err.is_recoverable(), + kind.is_recoverable_without_rebinding(), + "{kind}: engine-level recoverability must follow the refusal, not the variant" + ); + } +} diff --git a/crates/larql-kv/tests/strict_refusal/main.rs b/crates/larql-kv/tests/strict_refusal/main.rs new file mode 100644 index 000000000..f079a763b --- /dev/null +++ b/crates/larql-kv/tests/strict_refusal/main.rs @@ -0,0 +1,54 @@ +//! The strict guarantee, end to end: a refused expert cannot become a token. +//! +//! The `FfnBackend` ring gave the refusal a typed channel. The dispatch ring +//! widened the six `kv_*_via_dispatch` helpers so it survives the per-layer +//! loop, and `StandardEngine` terminates it as `EngineError::Execution`. These +//! tests drive the whole chain from the engine's public surface, because that +//! is the only altitude at which the guarantee is observable — every +//! intermediate ring can be correct while the composition still erases. +//! +//! ## What is pinned +//! +//! ```text +//! strict + refusing route → Err, original RefusalKind, no hidden state +//! best-effort + refusing route → Ok, complete dense half, refusal recorded +//! not applicable → Ok, local dispatch, nothing recorded +//! ``` +//! +//! The first is the merge gate ([`gate`]). [`outcomes`] holds the other two, +//! and the question strict semantics raise that the old silent path did not: +//! what a refusal leaves behind. "Fix the residency and retry" is only a +//! supported workflow if a refused step mutated nothing — so a decode either +//! rewinds what it appended, or says it could not and refuses to continue. +//! +//! Those two sweep `StandardEngine`. The other axis is the engine itself: +//! [`engines`] runs the gate against every `EngineKind`, and [`engine_state`] +//! runs the outcome questions against each one — because the rewind-or- +//! invalidate answer differs by what each engine treats as canonical. +//! +//! ## Scope +//! +//! `prefill_quant` / `decode_step_quant` are deliberately absent: by contract +//! they discard the caller's FFN and substitute a `WalkFfn` built from the +//! vindex, so a strict route cannot reach the dispatch ring through them. The +//! resident pair is the production strict path — it threads the caller's FFN +//! through the same `do_prefill` / `do_decode_step` bodies — and is covered. +//! +//! The engines that route their FFN through `larql_kv::engines::layer_ffn_or_moe` +//! rather than through the dispatch helpers (markov-rs, markov-rs-codec, +//! turbo-quant, unlimited-context, boundary-per-layer) are covered by +//! [`engines_gate`] and [`engine_state`]: that helper carries the refusal too +//! now. So does `larql_kv::generation::kv_prefill_run`, which is both +//! `no-cache`'s forward and the oracle the dispatch ring is compared against. +//! +//! `apollo` is the one engine that cannot dispatch experts at all — its +//! forward runs below the `FfnBackend` seam — so it refuses the +//! *architecture* rather than a route. See [`engines`]'s header. + +mod engine_state; +mod engines; +mod engines_gate; +mod entry; +mod gate; +mod outcomes; +mod routes; diff --git a/crates/larql-kv/tests/strict_refusal/outcomes.rs b/crates/larql-kv/tests/strict_refusal/outcomes.rs new file mode 100644 index 000000000..4baf97ba7 --- /dev/null +++ b/crates/larql-kv/tests/strict_refusal/outcomes.rs @@ -0,0 +1,269 @@ +//! The other two outcomes — degrade, and not applicable — plus retry safety: +//! what a refusal leaves behind in the engine that raised it. + +use larql_execution::RefusalKind; +use larql_inference::ffn::{FfnBackend, MoeFfn, RefusalPolicy, WeightFfn}; +use larql_inference::kv_engine::EngineError; +use larql_inference::test_utils::make_test_gemma4_moe_weights; +use larql_inference::KvEngine; +use larql_kv::StandardEngine; +use ndarray::Array2; + +use crate::entry::{drive, Entry, NEXT_TOKEN, PROMPT}; +use crate::routes::{ExecutingRoute, RefusingRoute}; + +/// Best-effort + the same refusal → `Ok`, with a complete dense half and the +/// refusal recorded. +/// +/// The degradation is explicit on both sides: the caller gets a usable hidden +/// state *and* `refusal()` is set, so a route that chose to degrade can still +/// say that it did. Pinned against the strict case with the identical route, +/// because the whole claim is that policy — not the route — decides. +#[test] +fn best_effort_degrades_explicitly_where_strict_refuses() { + let weights = make_test_gemma4_moe_weights(); + let route = RefusingRoute::new(RefusalKind::Residency); + + let ffn = MoeFfn::best_effort(&weights, &route); + assert_eq!(ffn.policy(), RefusalPolicy::BestEffort); + + let hidden = + drive(Entry::PrefillSync, &weights, &ffn).expect("best-effort must degrade, not stop"); + assert_eq!( + hidden.shape(), + &[1, weights.hidden_size], + "the degraded answer must still be a complete hidden row" + ); + assert!( + hidden.iter().all(|v| v.is_finite()), + "degradation must not smuggle in NaNs" + ); + assert!( + !ffn.all_experts_executed(), + "a degraded run must not report a clean one" + ); + let recorded = ffn.refusal().expect("the refusal must be recorded"); + assert_eq!(recorded.kind, RefusalKind::Residency); + + // Same route, strict policy, opposite outcome. This pair is the claim. + let strict = MoeFfn::strict(&weights, &route); + assert!( + drive(Entry::PrefillSync, &weights, &strict).is_err(), + "the identical route must refuse under Strict — policy decides, not the route" + ); +} + +/// Not applicable → `Ok`, via normal local dispatch, with nothing recorded. +/// +/// Two halves, because "not applicable" arrives two ways: a backend with no +/// MoE hook at all (the trait default, `Ok(None)`), and a hook that ran and +/// declined. Neither may look like a refusal — conflating them is what the +/// error channel was added to prevent. +#[test] +fn not_applicable_dispatches_locally_and_records_nothing() { + let weights = make_test_gemma4_moe_weights(); + + // (a) No MoE hook: `forward_moe_full_layer` is the trait default. + let plain = WeightFfn { weights: &weights }; + let probe = plain.forward_moe_full_layer(0, &Array2::zeros((1, weights.hidden_size))); + assert!( + matches!(probe, Ok(None)), + "the default hook must be not-applicable, not a refusal" + ); + let hidden = drive(Entry::PrefillSync, &weights, &plain) + .expect("a dense-dispatching backend must prefill normally on a MoE arch"); + assert_eq!(hidden.shape(), &[1, weights.hidden_size]); + + // (b) A strict route that executes: nothing recorded, nothing refused. + let ffn = MoeFfn::strict(&weights, &ExecutingRoute); + let hidden = drive(Entry::PrefillSync, &weights, &ffn).expect("an executing route must serve"); + assert_eq!(hidden.shape(), &[1, weights.hidden_size]); + assert!( + ffn.all_experts_executed() && ffn.refusal().is_none(), + "a clean run must record nothing — a strict policy that fires on success \ + is indistinguishable from one that never fires" + ); +} + +/// A refused decode step advances no observable K/V state. +/// +/// This is the retry-safety guarantee. A decode step mutates before it can +/// know whether it will finish — each layer's attention appends the new +/// token's K/V, and only then can the FFN refuse — so without a rewind the +/// engine would be left holding a token it never completed. +/// +/// That matters more under strict semantics than it did before, because +/// `Residency` refusals make "fix the cause and retry" an explicit supported +/// workflow. An engine that reported a recoverable refusal *and* silently kept +/// the half-applied append would append the token twice on the retry. +#[test] +fn a_refused_decode_step_advances_no_observable_state() { + let weights = make_test_gemma4_moe_weights(); + let mut engine = StandardEngine::new(None); + + let clean = MoeFfn::strict(&weights, &ExecutingRoute); + engine + .prefill(&weights, &clean, &PROMPT) + .expect("clean prefill"); + let before = engine.window_tokens(); + assert_eq!( + before, + PROMPT.len(), + "prefill caches one row per prompt token" + ); + + let route = RefusingRoute::new(RefusalKind::Residency); + let refusing = MoeFfn::strict(&weights, &route); + let err = engine + .decode_step(&weights, &refusing, NEXT_TOKEN) + .expect_err("the refused step must not produce a hidden state"); + assert_eq!(err.refusal_kind(), Some(RefusalKind::Residency)); + assert!( + err.engine_state_is_retryable(), + "an unbounded cache rewinds exactly, so the engine stays usable" + ); + assert!( + err.is_recoverable(), + "a rewound Residency refusal is the case a caller may act on and retry" + ); + assert_eq!( + engine.window_tokens(), + before, + "the refused step must leave the cache exactly as it found it" + ); +} + +/// Retrying the refused token produces the answer the refusal denied. +/// +/// The strongest statement of the guarantee: not merely that the row count +/// came back, but that the rewound engine computes what an engine that never +/// saw the refusal computes. A rewind that restored the length while leaving +/// the buffers shifted would pass the length check and fail this one. +#[test] +fn a_rewound_engine_decodes_the_retried_token_identically() { + let weights = make_test_gemma4_moe_weights(); + let clean = MoeFfn::strict(&weights, &ExecutingRoute); + + // Engine A: refuse the token, then retry it with a working route. + let mut retried = StandardEngine::new(None); + retried + .prefill(&weights, &clean, &PROMPT) + .expect("prefill A"); + let route = RefusingRoute::new(RefusalKind::Residency); + let refusing = MoeFfn::strict(&weights, &route); + retried + .decode_step(&weights, &refusing, NEXT_TOKEN) + .expect_err("must refuse"); + let after_retry = retried + .decode_step(&weights, &clean, NEXT_TOKEN) + .expect("the rewound engine must accept the retried token"); + + // Engine B: the same token, never refused. + let mut reference = StandardEngine::new(None); + reference + .prefill(&weights, &clean, &PROMPT) + .expect("prefill B"); + let baseline = reference + .decode_step(&weights, &clean, NEXT_TOKEN) + .expect("reference decode"); + + assert_eq!( + after_retry.iter().map(|v| v.to_bits()).collect::>(), + baseline.iter().map(|v| v.to_bits()).collect::>(), + "a retry after a rewound refusal must be bit-identical to never having \ + refused — anything less means the rewind restored the shape, not the state" + ); + assert_eq!(retried.window_tokens(), reference.window_tokens()); +} + +/// When the rewind cannot be trusted, the engine says so and stops. +/// +/// A windowed cache that is already at its limit drops its oldest row to make +/// room for the new one, and that row is gone. Row count cannot see it — +/// append-then-drop leaves the count unchanged — so the engine must not +/// pretend it rewound. It reports [`EngineError::StateInvalidated`], refuses +/// every later decode, and keeps the original refusal reachable underneath. +#[test] +fn an_unrewindable_refusal_invalidates_the_engine_and_says_so() { + const WINDOW: usize = 2; + let weights = make_test_gemma4_moe_weights(); + let clean = MoeFfn::strict(&weights, &ExecutingRoute); + + // Prompt is longer than the window, so prefill leaves every layer at the + // limit — the state in which the next append must evict. + let mut engine = StandardEngine::new(Some(WINDOW)); + engine + .prefill(&weights, &clean, &PROMPT) + .expect("windowed prefill"); + assert_eq!( + engine.window_tokens(), + WINDOW, + "prefill must fill the window for this test to exercise eviction" + ); + + let route = RefusingRoute::new(RefusalKind::Residency); + let refusing = MoeFfn::strict(&weights, &route); + let err = engine + .decode_step(&weights, &refusing, NEXT_TOKEN) + .expect_err("must refuse"); + + assert!( + matches!(err, EngineError::StateInvalidated { .. }), + "an unrewindable failure must not be reported as an ordinary refusal: {err:?}" + ); + assert!(!err.engine_state_is_retryable()); + assert!( + !err.is_recoverable(), + "a caller must not be told to skip this row and carry on with a dead engine" + ); + // The wrapper costs the cause neither its classification nor the fact + // that the operation itself could have succeeded elsewhere. + assert_eq!(err.refusal_kind(), Some(RefusalKind::Residency)); + assert!(err.operation_is_recoverable()); + + // Every later decode refuses rather than computing from the wreckage. + let follow_up = engine + .decode_step(&weights, &clean, NEXT_TOKEN) + .expect_err("an invalidated engine must refuse further decode steps"); + assert!(matches!(follow_up, EngineError::InvariantViolation { .. })); + + // Re-prefilling replaces the cache outright, which is the way back. + engine + .prefill(&weights, &clean, &PROMPT) + .expect("re-prefill must clear the invalidation"); + engine + .decode_step(&weights, &clean, NEXT_TOKEN) + .expect("a re-prefilled engine decodes normally again"); +} + +/// Prefill is transactional by construction. +/// +/// A refused prefill assigns nothing, so an engine that had already prefilled +/// keeps the cache it had rather than acquiring a half-built one. Pinned +/// because the property is a consequence of statement order — the handles are +/// installed after the `?` — and is therefore easy to lose in a refactor. +#[test] +fn a_refused_prefill_leaves_an_earlier_prefill_intact() { + let weights = make_test_gemma4_moe_weights(); + let clean = MoeFfn::strict(&weights, &ExecutingRoute); + let mut engine = StandardEngine::new(None); + engine + .prefill(&weights, &clean, &PROMPT) + .expect("first prefill"); + let before = engine.window_tokens(); + + let route = RefusingRoute::new(RefusalKind::Unsupported); + let refusing = MoeFfn::strict(&weights, &route); + let err = engine + .prefill(&weights, &refusing, &PROMPT) + .expect_err("the second prefill must refuse"); + assert!(err.engine_state_is_retryable()); + assert_eq!( + engine.window_tokens(), + before, + "a refused prefill must not disturb the cache already in place" + ); + engine + .decode_step(&weights, &clean, NEXT_TOKEN) + .expect("the surviving cache must still decode"); +} diff --git a/crates/larql-kv/tests/strict_refusal/routes.rs b/crates/larql-kv/tests/strict_refusal/routes.rs new file mode 100644 index 000000000..74b25611c --- /dev/null +++ b/crates/larql-kv/tests/strict_refusal/routes.rs @@ -0,0 +1,148 @@ +//! Expert routes for the suite: one that refuses, one that executes. +//! +//! Both are `MoeExpertBackend`s, the seam `MoeFfn` dispatches experts through. +//! Keeping them here rather than in each test keeps the tests about policy. + +use larql_execution::RefusalKind; +use larql_inference::ffn::{MoeBackendError, MoeExpertBackend}; +use larql_inference::model::ModelWeights; +use larql_vindex::runtime::ExecutionError; +use ndarray::Array2; + +/// Fixture values for the concrete errors below. Meaningless as numbers — they +/// exist so each refusal is a well-formed instance of its variant, and so a +/// failure message names something traceable rather than a bare literal. +const ABSENT_EXPERT: u32 = 7; +const OUT_OF_RANGE_EXPERT: u32 = 99; +const EXPERT_POPULATION: usize = 8; +const RESIDENT_EXPERTS: usize = 0; +const BANK_NAME: &str = "strict-refusal-test-bank"; +const UNSERVED_FORMAT: &str = "mxfp4"; +const UNSERVED_OPERAND: &str = "experts.up_proj"; + +/// A route that refuses every layer with a chosen classification. +/// +/// One concrete error per [`RefusalKind`], picked so the mapping under test is +/// the real one in `ExecutionError::refusal()` rather than a fixture asserting +/// its own answer: a selected-but-absent expert is `Residency`, an +/// unimplemented decoder is `Unsupported`, and an expert outside the router's +/// addressable population is a `BindingDefect`. +pub struct RefusingRoute { + pub kind: RefusalKind, +} + +impl RefusingRoute { + pub fn new(kind: RefusalKind) -> Self { + Self { kind } + } + + fn error(&self) -> ExecutionError { + match self.kind { + RefusalKind::Residency => ExecutionError::SelectedExpertNotResident { + expert: ABSENT_EXPERT, + bank: BANK_NAME.into(), + resident: RESIDENT_EXPERTS, + population: EXPERT_POPULATION, + }, + RefusalKind::Unsupported => ExecutionError::UnsupportedFormat { + format: UNSERVED_FORMAT.into(), + operand: UNSERVED_OPERAND.into(), + }, + RefusalKind::BindingDefect => ExecutionError::ExpertOutOfRange { + expert: OUT_OF_RANGE_EXPERT, + population: EXPERT_POPULATION, + }, + } + } +} + +impl MoeExpertBackend for RefusingRoute { + fn forward_moe_seq( + &self, + _weights: &ModelWeights, + _layer: usize, + _h: &Array2, + _norm_offset: f32, + _eps: f32, + ) -> Result, MoeBackendError> { + Err(MoeBackendError::Bound(self.error())) + } + + fn name(&self) -> &'static str { + "refusing-route" + } +} + +/// A route that executes. +/// +/// A zero expert contribution is legitimate — the trait documents it as the +/// answer for a layer with no expert weights to route into. What matters is +/// that it returns `Ok`, so nothing is recorded and a strict policy stays +/// silent. +pub struct ExecutingRoute; + +impl MoeExpertBackend for ExecutingRoute { + fn forward_moe_seq( + &self, + _weights: &ModelWeights, + _layer: usize, + h: &Array2, + _norm_offset: f32, + _eps: f32, + ) -> Result, MoeBackendError> { + Ok(Array2::zeros(h.raw_dim())) + } + + fn name(&self) -> &'static str { + "executing-route" + } +} + +/// A route that serves one full pass over the layers, then refuses. +/// +/// Models the case the other routes cannot: a shard that goes away *mid +/// stream*, after the engine has already committed something irreversible. +/// `unlimited-context` archives a window and saves its boundary checkpoint +/// when the window fills, so "refuse on the first token" and "refuse on the +/// second" are different questions — only the second can find the engine +/// holding a stream it cannot complete. +/// +/// A pass is counted by visits to layer zero, because that is the only chunk +/// boundary visible from inside an `FfnBackend`. +pub struct RefuseAfterFirstPass { + passes: std::cell::Cell, + inner: RefusingRoute, +} + +impl RefuseAfterFirstPass { + pub fn new(kind: RefusalKind) -> Self { + Self { + passes: std::cell::Cell::new(0), + inner: RefusingRoute::new(kind), + } + } +} + +impl MoeExpertBackend for RefuseAfterFirstPass { + fn forward_moe_seq( + &self, + weights: &ModelWeights, + layer: usize, + h: &Array2, + norm_offset: f32, + eps: f32, + ) -> Result, MoeBackendError> { + if layer == 0 { + self.passes.set(self.passes.get() + 1); + } + if self.passes.get() <= 1 { + return ExecutingRoute.forward_moe_seq(weights, layer, h, norm_offset, eps); + } + self.inner + .forward_moe_seq(weights, layer, h, norm_offset, eps) + } + + fn name(&self) -> &'static str { + "refuse-after-first-pass" + } +} diff --git a/crates/larql-models/src/test_fixtures.rs b/crates/larql-models/src/test_fixtures.rs index 532994edc..9dbdda71b 100644 --- a/crates/larql-models/src/test_fixtures.rs +++ b/crates/larql-models/src/test_fixtures.rs @@ -627,6 +627,10 @@ pub const Q4K_TEST_INTER: usize = 256; pub const Q4K_TEST_VOCAB: usize = 256; /// Layer count for the Q4_K test fixture. pub const Q4K_TEST_NUM_LAYERS: usize = 2; +/// Query-head count for the Q4_K test fixtures. +pub const Q4K_TEST_NUM_Q: usize = 4; +/// K/V-head count for the Q4_K test fixtures (GQA reps = 2). +pub const Q4K_TEST_NUM_KV: usize = 2; /// Wide FFN width for the Q4_K fixture: threshold tests that route the /// walk's parallel Q4K-down branch need `hits ≥ 512` while staying /// below the full-K gemv rewrite at 80% density — so @@ -664,7 +668,14 @@ pub fn make_test_q4k_weights_layers(num_layers: usize) -> ModelWeights { "hidden_activation": "gelu_pytorch_tanh", "rope_theta": 10000.0, }); - q4k_test_weights_from_json(arch_json, num_layers, Q4K_TEST_INTER) + q4k_test_weights_from_json( + arch_json, + num_layers, + Q4K_TEST_INTER, + Q4K_TEST_HIDDEN, + Q4K_TEST_NUM_Q, + Q4K_TEST_NUM_KV, + ) } /// Wide-FFN sibling of [`make_test_q4k_weights`]: same Gemma 3 arch and @@ -690,7 +701,14 @@ pub fn make_test_q4k_weights_wide() -> ModelWeights { "hidden_activation": "gelu_pytorch_tanh", "rope_theta": 10000.0, }); - q4k_test_weights_from_json(arch_json, num_layers, Q4K_TEST_INTER_WIDE) + q4k_test_weights_from_json( + arch_json, + num_layers, + Q4K_TEST_INTER_WIDE, + Q4K_TEST_HIDDEN, + Q4K_TEST_NUM_Q, + Q4K_TEST_NUM_KV, + ) } /// Rope-scaled sibling of [`make_test_q4k_weights`]: Gemma-3 arch at the @@ -721,17 +739,25 @@ pub fn make_test_q4k_weights_rope_scaled() -> ModelWeights { "sliding_window": 512, "rope_scaling": {"rope_type": "linear", "factor": 8.0}, }); - q4k_test_weights_from_json(arch_json, num_layers, Q4K_TEST_INTER) + q4k_test_weights_from_json( + arch_json, + num_layers, + Q4K_TEST_INTER, + Q4K_TEST_HIDDEN, + Q4K_TEST_NUM_Q, + Q4K_TEST_NUM_KV, + ) } fn q4k_test_weights_from_json( arch_json: serde_json::Value, num_layers: usize, intermediate: usize, + hidden: usize, + num_q: usize, + num_kv: usize, ) -> ModelWeights { - let num_q = 4usize; - let num_kv = 2usize; - let head_dim = Q4K_TEST_HIDDEN / num_q; + let head_dim = hidden / num_q; let arch = detect_from_json(&arch_json); let mut tensors: HashMap = HashMap::new(); @@ -745,14 +771,11 @@ fn q4k_test_weights_from_json( seed }; - let embed = rand_mat_seeded(Q4K_TEST_VOCAB, Q4K_TEST_HIDDEN, 0.05, next_seed()); + let embed = rand_mat_seeded(Q4K_TEST_VOCAB, hidden, 0.05, next_seed()); let lm_head = embed.clone(); tensors.insert(arch.embed_key().to_string(), embed.clone()); - vectors.insert( - arch.final_norm_key().to_string(), - vec![1.0; Q4K_TEST_HIDDEN], - ); + vectors.insert(arch.final_norm_key().to_string(), vec![1.0; hidden]); let q_dim = num_q * head_dim; let kv_dim = num_kv * head_dim; @@ -760,43 +783,53 @@ fn q4k_test_weights_from_json( for layer in 0..num_layers { tensors.insert( arch.attn_q_key(layer), - rand_mat_seeded(q_dim, Q4K_TEST_HIDDEN, 0.05, next_seed()), + rand_mat_seeded(q_dim, hidden, 0.05, next_seed()), ); tensors.insert( arch.attn_k_key(layer), - rand_mat_seeded(kv_dim, Q4K_TEST_HIDDEN, 0.05, next_seed()), + rand_mat_seeded(kv_dim, hidden, 0.05, next_seed()), ); tensors.insert( arch.attn_v_key(layer), - rand_mat_seeded(kv_dim, Q4K_TEST_HIDDEN, 0.05, next_seed()), + rand_mat_seeded(kv_dim, hidden, 0.05, next_seed()), ); tensors.insert( arch.attn_o_key(layer), - rand_mat_seeded(Q4K_TEST_HIDDEN, q_dim, 0.05, next_seed()), + rand_mat_seeded(hidden, q_dim, 0.05, next_seed()), ); tensors.insert( arch.ffn_gate_key(layer), - rand_mat_seeded(intermediate, Q4K_TEST_HIDDEN, 0.05, next_seed()), + rand_mat_seeded(intermediate, hidden, 0.05, next_seed()), ); tensors.insert( arch.ffn_up_key(layer), - rand_mat_seeded(intermediate, Q4K_TEST_HIDDEN, 0.05, next_seed()), + rand_mat_seeded(intermediate, hidden, 0.05, next_seed()), ); tensors.insert( arch.ffn_down_key(layer), - rand_mat_seeded(Q4K_TEST_HIDDEN, intermediate, 0.05, next_seed()), + rand_mat_seeded(hidden, intermediate, 0.05, next_seed()), ); - vectors.insert(arch.input_layernorm_key(layer), vec![0.5; Q4K_TEST_HIDDEN]); - vectors.insert( - arch.post_attention_layernorm_key(layer), - vec![0.5; Q4K_TEST_HIDDEN], - ); + vectors.insert(arch.input_layernorm_key(layer), vec![0.5; hidden]); + vectors.insert(arch.post_attention_layernorm_key(layer), vec![0.5; hidden]); if let Some(k) = arch.pre_feedforward_layernorm_key(layer) { - vectors.insert(k, vec![0.5; Q4K_TEST_HIDDEN]); + vectors.insert(k, vec![0.5; hidden]); } if let Some(k) = arch.post_feedforward_layernorm_key(layer) { - vectors.insert(k, vec![0.5; Q4K_TEST_HIDDEN]); + vectors.insert(k, vec![0.5; hidden]); + } + // QK-norm, on the architectures that declare it (Gemma 3 / 4). + // These were declared-but-absent until 2026-08, which made the + // fixture claim an architecture it did not actually carry: every + // consumer that resolves the weight got `None` and silently + // skipped the stage, so a backend disagreeing about what to do + // with a missing QK-norm weight had nothing pinning it. Sized per + // head_dim, not hidden — QK-norm normalises within a head. + if let Some(k) = arch.attn_q_norm_key(layer) { + vectors.insert(k, vec![0.5; head_dim]); + } + if let Some(k) = arch.attn_k_norm_key(layer) { + vectors.insert(k, vec![0.5; head_dim]); } } @@ -812,7 +845,7 @@ fn q4k_test_weights_from_json( position_embed: None, arch, num_layers, - hidden_size: Q4K_TEST_HIDDEN, + hidden_size: hidden, intermediate_size: intermediate, vocab_size: Q4K_TEST_VOCAB, head_dim, @@ -1132,6 +1165,54 @@ pub fn make_test_gemma4_moe_weights() -> ModelWeights { } } +/// Gemma-3 Q4_K fixture at **caller-chosen attention dimensions**. +/// +/// Every other Q4_K fixture is pinned to `hidden = 256, num_q = 4`, i.e. +/// `head_dim = 64`. That made shape sensitivity untestable: a kernel with +/// a `head_dim` assumption would be wrong on every fixture and right on +/// every real model, which is exactly the signature of the Metal +/// batched-prefill divergence (see `larql-kv`'s +/// `gpu_engine_parity::gemma3_prefill_gap`). This builder is the knob for +/// bisecting that axis, and for any future "does this kernel assume a +/// shape?" question. +/// +/// `hidden` must be a multiple of 256 (Q4_K super-block) and divisible by +/// `num_q`; `num_q` must be divisible by `num_kv`. +pub fn make_test_q4k_weights_with_dims( + hidden: usize, + num_q: usize, + num_kv: usize, + num_layers: usize, +) -> ModelWeights { + assert!( + hidden.is_multiple_of(256), + "Q4_K needs a hidden size that is a multiple of its 256-element super-block, got {hidden}" + ); + assert!( + num_q != 0 && hidden.is_multiple_of(num_q), + "hidden {hidden} must divide evenly into {num_q} query heads" + ); + assert!( + num_kv != 0 && num_q.is_multiple_of(num_kv), + "{num_q} query heads must group evenly onto {num_kv} K/V heads" + ); + let head_dim = hidden / num_q; + + let arch_json = serde_json::json!({ + "model_type": "gemma3_text", + "hidden_size": hidden, + "num_hidden_layers": num_layers, + "intermediate_size": hidden, + "head_dim": head_dim, + "num_attention_heads": num_q, + "num_key_value_heads": num_kv, + "vocab_size": Q4K_TEST_VOCAB, + "hidden_activation": "gelu_pytorch_tanh", + "rope_theta": 10000.0, + }); + q4k_test_weights_from_json(arch_json, num_layers, hidden, hidden, num_q, num_kv) +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/larql-server/ROADMAP.md b/crates/larql-server/ROADMAP.md index 512aa75b9..19e29e2b3 100644 --- a/crates/larql-server/ROADMAP.md +++ b/crates/larql-server/ROADMAP.md @@ -2094,7 +2094,7 @@ Implementation surface: ~1600 LOC across three new files (`src/routes/openai_embeddings.rs`, `src/routes/openai_completions.rs`, `src/routes/openai_chat.rs`) + reshape of `src/routes/models.rs` + 4 routes wired into both single-model and multi-model routers + 23 unit -tests + 19 integration tests + new live `examples/openai_demo.rs` +tests + 19 integration tests + new live `crates/larql-demos/examples/server/openai_demo.rs` walkthrough that boots the server in-process via `tower::ServiceExt::oneshot` and exercises every endpoint. diff --git a/crates/larql-vindex/Cargo.toml b/crates/larql-vindex/Cargo.toml index 7823eb19b..1bfaf3a4f 100644 --- a/crates/larql-vindex/Cargo.toml +++ b/crates/larql-vindex/Cargo.toml @@ -13,6 +13,7 @@ larql-core = { path = "../larql-core", default-features = false } larql-models = { path = "../larql-models" } larql-compute = { path = "../larql-compute" } larql-vindex-spec = { path = "../larql-vindex-spec" } +larql-execution = { path = "../larql-execution" } serde = { workspace = true } # Apple-only sibling backend. Made `optional = true` so the # (target-gated) cfg path on non-macOS doesn't need to fake it out; @@ -66,6 +67,9 @@ gpu = ["dep:larql-compute-metal"] [dev-dependencies] larql-inference = { path = "../larql-inference" } +# The KvEngine decode path, for the teacher-forced / free-running decode +# parity harness. Dev-only: nothing in the library depends on the engine layer. +larql-kv = { path = "../larql-kv" } criterion = "0.5" tempfile = "3" # HTTP mocking for the publish trio (lfs / remote / upload). @@ -75,6 +79,10 @@ tempfile = "3" mockito = "1.7" serial_test = "3.2" +[[bench]] +name = "vindex3_bound_execute" +harness = false + [[bench]] name = "vindex_ops" harness = false diff --git a/crates/larql-vindex/README.md b/crates/larql-vindex/README.md index 5319a3b2c..4b54313fa 100644 --- a/crates/larql-vindex/README.md +++ b/crates/larql-vindex/README.md @@ -77,7 +77,7 @@ responsible for capturing them at the correct depth, which is exactly what `larql_inference::capture_decoy_residuals` does. Validated against synthetic constellations by the unit tests in `patch/refine.rs`; the end-to-end Gemma 3 4B reproduction lives in -`larql-lql/examples/refine_demo.rs`. +`larql-lql/crates/larql-demos/examples/lql/refine_demo.rs`. ## The Headline diff --git a/crates/larql-vindex/benches/vindex3_bound_execute.rs b/crates/larql-vindex/benches/vindex3_bound_execute.rs new file mode 100644 index 000000000..0898e8321 --- /dev/null +++ b/crates/larql-vindex/benches/vindex3_bound_execute.rs @@ -0,0 +1,238 @@ +//! VINDEX3 bound-execution perf gate — the successor to +//! `vindex_storage_dispatch`. +//! +//! That bench asks whether an indirection (`&[u8]` → `Bytes` → `dyn`) costs +//! anything on the byte-fetch path. This one asks the question VINDEX3 raises +//! instead: +//! +//! > **Has resolution leaked into decode?** +//! +//! The whole load path — manifest lookup, capability traversal, alternative +//! selection, variant resolution, authority folding, contract checking — +//! happens once, before a `BoundMoeOperation` exists. If any of it stays +//! reachable from `execute`, per-token cost starts tracking *catalogue* size +//! rather than *work* size, and it does so quietly: the answers stay correct +//! and the model just gets slower as the index gets richer. +//! +//! So the gate is a scaling property, not a throughput number. +//! +//! # What "flat" does and does not mean here +//! +//! A first draft of this bench asserted per-token cost is *flat* in population. +//! That claim is wrong, and measuring it said so: routing legitimately scores +//! every expert before it can take a top-k, so an honest budget is +//! +//! ```text +//! O(population × hidden) router scoring unavoidable +//! O(top_k × expert work) selected experts must not track population +//! O(catalogue) resolution must not appear at all +//! ``` +//! +//! The gate is therefore the *shape* of the growth. Scoring is a single +//! `[population, hidden]` matvec while each expert is three +//! `[intermediate, hidden]`-scale passes, so a 64× population increase should +//! buy a small constant factor — not 64×, and not superlinear. A resolution +//! leak, a per-token population scan over experts, or a revalidate call would +//! all bend the line well past that. +//! +//! Absolute times are not the point and are not comparable to the incumbent: +//! this is the reference decoder, which reads every operand element by element +//! on purpose. +//! +//! Run with: `cargo bench -p larql-vindex --bench vindex3_bound_execute` + +use std::hint::black_box; + +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; + +use larql_compute::{Activation, MoeTopKWeightPolicy}; +use larql_vindex::format::capability::coordinate::BankCoordinate; +use larql_vindex::format::capability::{ + binding::RepresentationIdentity, component::ComponentContract, +}; +use larql_vindex::format::lyrw2::region_format::RegionFormat; +use larql_vindex::format::lyrw2::region_role::RegionRole; +use larql_vindex::runtime::MoeInputs; +use larql_vindex::runtime::{ + execute, execute_traced, BoundBankOperation, BoundExpert, BoundExpertScaling, + BoundMoeOperation, BoundProjection, BoundReduction, BoundRouter, BoundTensor, ExpertKernel, + ProjectionArrangement, RouterKernel, +}; + +/// Populations spanning two orders of magnitude. 64× the experts should cost a +/// small constant factor — the router's scoring pass — and nothing more. +const POPULATIONS: [usize; 3] = [8, 64, 512]; +const HIDDEN: usize = 64; +const INTERMEDIATE: usize = 48; +const TOP_K: usize = 4; +const VARIANT: &str = "bench"; + +/// Deterministic weights. Value content is irrelevant to timing; determinism +/// keeps run-to-run comparison meaningful. +fn weight(seed: usize, i: usize) -> f32 { + (((seed * 31 + i * 17) % 199) as f32) / 199.0 - 0.5 +} + +fn bytes(count: usize, seed: usize) -> Vec { + (0..count) + .flat_map(|i| weight(seed, i).to_le_bytes()) + .collect() +} + +/// Owns every operand for one bound operation. +struct Operands { + gate: Vec>, + up: Vec>, + gate_up: Vec>, + down: Vec>, + router: Vec, +} + +impl Operands { + fn new(population: usize) -> Self { + let proj = INTERMEDIATE * HIDDEN; + Self { + gate: (0..population).map(|e| bytes(proj, e)).collect(), + up: (0..population).map(|e| bytes(proj, e + 1_000)).collect(), + gate_up: (0..population) + .map(|e| { + let mut v = bytes(proj, e); + v.extend(bytes(proj, e + 1_000)); + v + }) + .collect(), + down: (0..population) + .map(|e| bytes(HIDDEN * INTERMEDIATE, e + 2_000)) + .collect(), + router: bytes(population * HIDDEN, 7), + } + } + + fn operation(&self, arrangement: ProjectionArrangement) -> BoundMoeOperation<'_> { + let population = self.gate.len(); + let experts = (0..population) + .map(|e| BoundExpert { + expert_id: e as u32, + projection: match arrangement { + ProjectionArrangement::Decomposed => BoundProjection::Decomposed { + gate: tensor( + &RegionRole::Gate.name(), + &self.gate[e], + INTERMEDIATE, + HIDDEN, + ), + up: tensor(&RegionRole::Up.name(), &self.up[e], INTERMEDIATE, HIDDEN), + }, + ProjectionArrangement::Fused => BoundProjection::Fused { + gate_up: tensor( + &RegionRole::GateUpFused.name(), + &self.gate_up[e], + INTERMEDIATE * 2, + HIDDEN, + ), + }, + }, + down: tensor( + &RegionRole::Down.name(), + &self.down[e], + HIDDEN, + INTERMEDIATE, + ), + }) + .collect(); + + BoundMoeOperation { + router: BoundRouter { + weight: tensor("router", &self.router, population, HIDDEN), + top_k: TOP_K, + selected_weight: MoeTopKWeightPolicy::RenormalizedSoftmax, + scaling: BoundExpertScaling::None, + kernel: RouterKernel::default(), + }, + transforms: Vec::new(), + banks: vec![BoundBankOperation { + bank: BankCoordinate::new(0, 0), + experts, + intermediate_dim: INTERMEDIATE, + hidden_dim: HIDDEN, + activation: Activation::Silu, + kernel: ExpertKernel::default(), + }], + reduction: BoundReduction::WeightedSum, + residual_dim: HIDDEN, + } + } +} + +fn tensor<'a>(region_set: &str, data: &'a [u8], rows: usize, cols: usize) -> BoundTensor<'a> { + BoundTensor::direct( + RepresentationIdentity::new(region_set, VARIANT), + data, + RegionFormat::F32, + ComponentContract::matrix(rows as u32, cols as u32), + ) + .expect("bench operands are well-formed") +} + +fn residual() -> Vec { + (0..HIDDEN).map(|i| weight(99, i)).collect() +} + +/// The gate: growth across a 64× population must stay near the router term. +fn population_scaling(c: &mut Criterion) { + let mut group = c.benchmark_group("vindex3/execute_by_population"); + for population in POPULATIONS { + let operands = Operands::new(population); + let op = operands.operation(ProjectionArrangement::Decomposed); + op.validate().expect("bench operation is well-formed"); + let input = residual(); + // One token per iteration. Expert work is identical across these + // points; only the router's scoring pass grows. + group.throughput(Throughput::Elements(1)); + group.bench_with_input( + BenchmarkId::from_parameter(population), + &population, + |b, _| { + b.iter(|| black_box(execute(&op, MoeInputs::shared(black_box(&input))).unwrap())) + }, + ); + } + group.finish(); +} + +/// Fused and decomposed storage must cost about the same, as well as agreeing. +fn arrangement_parity(c: &mut Criterion) { + let operands = Operands::new(POPULATIONS[1]); + let mut group = c.benchmark_group("vindex3/execute_by_arrangement"); + for arrangement in ProjectionArrangement::ALL { + let op = operands.operation(arrangement); + let input = residual(); + group.bench_function(BenchmarkId::from_parameter(arrangement.name()), |b| { + b.iter(|| black_box(execute(&op, MoeInputs::shared(black_box(&input))).unwrap())) + }); + } + group.finish(); +} + +/// The no-op sink must compile away; the collecting one is diagnostic-only. +fn trace_overhead(c: &mut Criterion) { + let operands = Operands::new(POPULATIONS[1]); + let op = operands.operation(ProjectionArrangement::Decomposed); + let input = residual(); + let mut group = c.benchmark_group("vindex3/trace"); + group.bench_function("untraced", |b| { + b.iter(|| black_box(execute(&op, MoeInputs::shared(black_box(&input))).unwrap())) + }); + group.bench_function("traced", |b| { + b.iter(|| black_box(execute_traced(&op, MoeInputs::shared(black_box(&input))).unwrap())) + }); + group.finish(); +} + +criterion_group!( + benches, + population_scaling, + arrangement_parity, + trace_overhead +); +criterion_main!(benches); diff --git a/crates/larql-vindex/docs/vindex3-format-spec.md b/crates/larql-vindex/docs/vindex3-format-spec.md new file mode 100644 index 000000000..f5f00e077 --- /dev/null +++ b/crates/larql-vindex/docs/vindex3-format-spec.md @@ -0,0 +1,660 @@ +# Vindex Format Specification — VINDEX3 + +**Version:** 3.0-draft-2 +**Date:** 2026-08-01 (draft-2: three binary-layout corrections + two clarifications from the first lyrw2 implementation — §6.2, §6.4, §6.3, §6.5; recorded per pre-freeze amendment rule) +**Status:** Draft — pre-registration. No byte is frozen until the V2-0..V2-2 gates in the companion experiments document pass. **A first container now exists**: conformance fixture A round-trips through `format::vindex3` (write → detect → open → validate → bind → execute) bit-identically, and both the fused and decomposed FC1 renderings satisfy the same programme id. That closes the rows fixture A can carry on V2-0 and V2-1; profile-authority derivation, variant-selection refusal, the not-hard-coded row (fixtures B–D) and WALK/DESCRIBE parity remain open, so the ABI is **not** frozen and `extract` still writes VINDEX2. +**Predecessor:** [format-spec.md v0.4](format-spec.md) (VINDEX2) + +**Companion:** [`vindex3-experiments.md`](../../../docs/vindex3-experiments.md) (pre-registered experimental programme), [Conformance v1](conformance-v1.md), [Operations](operations-spec.md), [Ecosystem](ecosystem-spec.md), [LQL](../../../docs/lql-guide.md) +**Implementation target:** `larql-vindex` crate (Rust) + +> **A note on the number 2 appearing throughout.** This document specifies +> **VINDEX3**, and its own version is therefore `3.x`. Three things nearby keep +> a `2` on purpose and are not typos: +> +> | name | why it stays | +> |---|---| +> | `V2-0`…`V2-4` gates | pre-registered identifiers with results already recorded against them; renaming would orphan the lineage | +> | registry programme `vindex2` | same — it is an external key in chuk-experiments | +> | `lyrw2` / `FORMAT_VERSION = 2` | the *bank* format's own version, on a different axis: a VINDEX2 container holds LYRW v1 files, a VINDEX3 container holds LYRW v2 files (see `format::generation`) | +> +> Only the container is versioned 3. An on-disk `index.json` carrying +> `"version": 2` is a **VINDEX2** file — the predecessor above, not a draft of +> this one. +> +> Real VINDEX3 containers now exist for conformance fixture A +> (`format::vindex3`), so the format is no longer validated only through +> VINDEX2-sourced operands. Real *models* are still VINDEX2: `extract` does not +> emit VINDEX3 and will not until the outstanding V2-0/V2-1 rows close. + +--- + +## 1. What is VINDEX3? + +VINDEX3 is a **general-purpose serving container for sparse models** — serving meaning both inference *and* the LQL query surface (WALK, DESCRIBE, SELECT, EXPLAIN). It is the successor to the VINDEX2 dense/Gemma-oriented layout, and it exists to answer one question well: + +> Extract a supported checkpoint **once** into a stable, component-addressed layout, then vary **what is loaded, where it resides, what precision it uses, and whether a component is executed or queried** — without ever rebuilding the index. + +**The model IS the database** remains the founding principle, not a VINDEX2 legacy: the gate regions inside LYRW v2 banks *are* the KNN index, exactly as `gate_vectors.bin` *was* W_gate. VINDEX3 does not add a query index next to the weights; it keeps the weights queryable (§15). + +The key principles of VINDEX2 are retained unchanged: + +- **The model IS the database.** Each weight tensor is stored once, canonically, in its serving format. Nothing is stored twice. +- **Weights are separated by function, not by file size.** Sharding follows what inference does with a weight, not an arbitrary byte boundary. +- **mmap-first.** Every physical object is independently mmap-able; the OS pages in only what execution touches. +- **Loaders dispatch on declared tags, never sniff filenames.** +- **Fail closed.** A profile that lacks a required operand refuses to load with a precise diagnosis; it never silently degrades authority. + +The genuinely new pieces in VINDEX3 are exactly five: + +1. **Per-region quantisation.** Format belongs to each weight region, not to the entire layer file. +2. **Multiple physical segments per logical expert bank.** A logical layer can span several files without changing model semantics. +3. **Multiple bank kinds** — routed, shared and dense/hybrid — with declared geometry. +4. **A validated MoE programme manifest** describing router, banks, transforms and combine semantics, replacing the Gemma/Mixtral-shaped `moe_config`. +5. **Representation variants.** A region may carry several physically present encodings; profiles *select* among them and never request formats that were not extracted (§9.1). + +### 1.1 K3 validates the format; it does not define it + +The conformance envelope is defined by three real architectures plus a control: + +| Model | Routing | Shared experts | Expert space | Expert programme | Native format | +| ----- | ------- | -------------- | ------------ | ---------------- | ------------- | +| Direct MoE (control) | top-2 of 8 | 0 | residual | gated MLP | any | +| GPT-OSS | top-4 of 32/128 | 0 | residual | clamped gated MLP w/ residual term | MXFP4 | +| **Inkling-Small (276B-A12B)** | top-6 of 256, sigmoid + gate bias + norm_after_topk + route_scale 8.0; **shared-expert sink** (router scores shared experts) | 2 (always active) | residual | gated MLP; dense MLP at layer index 2 (mid-stack) | BF16, with NVFP4 / MXFP8 releases | +| **Kimi-Linear-48B-A3B** | top-8 of 256, sigmoid + renormalise + 2.446 scaling | 1 | residual | gated MLP; layer 0 dense (`first_k_dense_replace=1`) | BF16 | +| K3 | top-16 of 896 | 2 | latent 3584 | SiTU-GLU latent expert | (extraction: exact Q6_K baseline) | + +Inkling-Small replaces the hypothetical "Inkling-shaped" envelope member with the real released model (`thinkingmachines/Inkling-Small`, 532 GB BF16 + 4.46 GB MTP sidecar). It contributes what nothing else in the set can: real two-shared-expert reduction under a **shared-expert sink** router (shared experts inside the scoring/normalisation, `norm_after_topk`, gate bias, global scale — the richest router-semantics test in the envelope); real **NVFP4/MXFP8 native regions** via its quantised releases, whose mixed-precision convention (routed experts low-bit, shared experts and attention BF16) is itself a per-region-format use case; a mid-stack dense layer (`dense_mlp_idx: 2`) proving per-layer manifests handle arbitrary dense/MoE schedules, not just leading-dense; and — decisive for the rig — it is the first design-set model that **cannot be RAM-resident** on the M3 Max (~202 GiB routed at Q6_K, ~4.9 GiB per MoE layer, ~4.7 GiB of routed reads per token at top-6), so partial-residency, SSD-streaming and attn-local/FFN-remote profiles get their first *non-optional* real-model test at one-eighth K3 scale. Its MTP heads and multimodal towers are stored as optional auxiliary manifest-addressed tensors — the text backbone is the conformance target; the towers are opaque payload, and omitting MTP never changes authority (drafting only). + +Kimi-Linear-48B-A3B earns its seat three ways. It is the only **real, locally runnable** shared-expert member (98.3 GB BF16, ~3B active — Q6_K-class fits the M3 Max), so V2-4's shared-bank rung is proven on an actual checkpoint rather than a fixture. Its `first_k_dense_replace=1` hybrid stack is the only member exercising per-layer manifest heterogeneity (dense layer 0, MoE layers 1–26) on a real model. And it is K3's direct lineage ancestor — the same KDA(3):MLA(1) hybrid dense spine, 20 KDA + 7 MLA layers of real recurrence parameters — making it the dress rehearsal for the K3 adapter's class-1/class-2 plumbing at one-thirtieth the checkpoint size. Its sigmoid-scored, renormalised, scaled router also stress-tests the manifest's router vocabulary beyond softmax-top-k. + +K3 is the stress test — largest bank, latent expert space, shared pre/post projections. GPT-OSS and Inkling exist in the envelope precisely to stop K3-specific assumptions (a tensor literally named `gate_up`, top-16, no-shared-bank, residual-space-only) from becoming the ABI. + +--- + +## 2. Scope and non-goals + +VINDEX3 serves **one fixed checkpoint efficiently under different inference and query policies**. Browse/LQL is in scope (§15); training is not. It is explicitly not: + +- **A model-development store.** No optimisation for training, fine-tuning, gradient updates, adapter merging, or frequently rewritten weights. No copy-on-write component versioning. One extraction, then runtime policy. +- **A general neural-graph container.** VINDEX3 does not duplicate ONNX/safetensors-plus-compiler. The expert-programme vocabulary is deliberately bounded (§8.3). +- **A locality store.** Hot sets, retained experts, cache allocation, prefetch depth, local-versus-remote placement and reduced-top-K are **runtime metadata over the index**, never physical-format decisions (§9). + +The supported contract is: + +> Sparse decoder MoEs composed from routed and shared expert banks, optional pre/post transforms, declarative routing/reduction semantics, and a bounded expert-programme vocabulary — plus every dense model VINDEX2 supports, expressed as the degenerate single-entry case. + +Genuinely novel expert topologies extend via a new `programme_id` (§8.4) without changing region storage. + +--- + +## 3. Design principles + +Principles 1–3 carry over from VINDEX2 §5.12; principle 1 is **amended** in VINDEX3. + +1. **Structure is orthogonal to quantisation — now at region granularity.** VINDEX2 declared one `quant_format` per layer file, forbidding `gate/up = Q6_K, down = MXFP4` inside a layer. VINDEX3 moves the format tag to each weight region. Re-quantising one projection role is rewriting those regions (or adding a sibling segment), not replacing the layer. +2. **Unified for dense and MoE.** A dense layer is a bank with `num_entries = 1`. Binary format and dispatch path are identical. +3. **Native OS addressability.** Each segment file is independently mmap'd; expert sharding reads only assigned entry byte ranges; no offset arithmetic into a global blob. +4. **The split rule.** A component gets independent physical identity **only when LARQL may independently omit it, quantise it, place it, prefetch it, execute it — or query it.** Conceptual tensor taxonomy is not a reason to split. The query clause matters: WALK reads gate rows without up or down, so on a browse-enabled index the gate role has an independent access pattern by construction (§15.2), even if inference always fetches gate/up/down together. +5. **Storage aligns with dispatch.** The natural extent is the expert group matching the grouped kernel's dispatch width, so one grouped dispatch ≈ one extent ≈ one prefetch/read unit. +6. **Representable ≠ servable.** The format may describe combinations no kernel can yet execute. The capability registry (§10) distinguishes representable / reference-executable / dispatched / production, exactly mirroring the K3 ledger maturity discipline. +7. **Logical ownership by layer.** The logical layer remains the stable semantic unit. Segmentation (§7) is a physical storage parameter, invisible to model semantics. + +--- + +## 4. The five durable weight classes + +The serving ABI freezes exactly five classes. These are the boundaries that inference policy may ever want to fetch, place, quantise or omit independently: + +| # | Class | Contents | Why independent | +| - | ----- | -------- | --------------- | +| 1 | **Control & router** | Embeddings, norms, LM head, router weights, routing metadata, recurrence/control parameters | Small, always resident, precision-sensitive | +| 2 | **Dense spine** | Attention / KDA / MLA projections, per layer or major projection class | Touched every token; future KDA3 target; independent quantisation ladder | +| 3 | **Shared FFN** | Shared experts and shared latent pre/post projections, per layer | Touched every token; different residency economics from routed | +| 4 | **Routed gate/up banks** | Per-layer expert-group extents | Candidate for exact Q6_K or native low-bit; grouped-dispatch aligned | +| 5 | **Routed down banks** | Per-layer expert-group extents | Independent quantisation, placement and (approximate-profile) omission policy | + +Classes 4 and 5 remain physically separable **because their inference treatment can differ** (precision, bandwidth, kernel maturity, residency, remote/local placement) — not because a tensor taxonomy says so. Where a model's serving policy treats them identically, a single fused `gate_up_fused + down` bank per layer satisfies the ABI (§6.5). + +No sixth class. Everything else — hot sets, expert retention, cache sizing, prefetch order, exact-vs-approximate selection — is profile/runtime metadata (§9). + +--- + +## 5. Directory layout + +``` +model.vindex/ +│ +├── index.json # SOLE ROOT AUTHORITY (§12): version, identity, provenance, +│ # checksums, class map, segment lists, references to everything below +├── moe_manifest.json # model + MoE programme description (§8), referenced from index.json +├── profiles/ # execution profiles (§9), referenced from index.json +│ ├── exact.json +│ ├── attn-local-ffn-remote.json +│ └── ... +│ +├── control/ # class 1 +│ ├── embeddings.bin +│ ├── norms.bin +│ ├── lm_head.bin # omitted if tied +│ └── routers.bin +│ +├── dense/ # class 2 +│ └── layer_{L}.weights # LYRW v2, dense bank(s) +│ +├── shared/ # class 3 +│ └── layer_{L}.weights # LYRW v2, shared bank + latent transforms +│ +├── routed/ # classes 4 & 5 +│ └── layer_{L}[.seg{S}].weights # LYRW v2, routed bank, segmented as needed +│ +├── query/ # LQL metadata sidecars (§15.3) — metadata, not weights +│ ├── down_meta.bin # DMET format unchanged from v1 §5.3 +│ ├── feature_labels.json +│ └── relation_clusters.json +│ +├── tokenizer.json +└── weight_manifest.json # manifest-addressed tensors (control/dense), unchanged shape from v1 §5.9 +``` + +Notes: + +- Control-class and non-bank dense tensors (routers, recurrence parameters, latent projections when the adapter prefers manifest addressing) remain ordinary **manifest-addressed tensors** with `key / shape / kind / offset / length` — the v1 `weight_manifest.json` shape is retained unchanged. +- Component addressability does **not** require thousands of filesystem objects. A class may be one file or a few large ones; addressability comes from the region tables inside them. +- **One root, as in v1.** `index.json` remains the manifest of record — it owns version, identity, provenance, checksums (every physical file's SHA256, verified by `larql verify`), the segment lists, and references to `moe_manifest.json`, `weight_manifest.json` and `profiles/`. There is no `superblock.json`: a second root creates competing authorities (whose checksums win? whose version controls compatibility?). A detached signing/atomic-replacement wrapper is introduced only if a concrete need for it appears, as an addition around `index.json`, never a rival to it. +- Profiles are **not** covered by the immutable artifact checksum set — they are mutable policy. `index.json` records which profile names ship with the artifact; their contents are checksummed individually and replaceable. + +--- + +## 6. LYRW v2 binary format + +LYRW v2 preserves the v1 magic and self-describing property, and generalises the fixed four-integer offset table into banks, segments and entry-region tables. + +### 6.1 Header + +``` +[header] + magic: u32 0x4C595257 ("LYRW") + format_version: u32 = 2 + logical_layer: u32 + num_banks: u16 + num_segments: u16 (segments described by THIS file's tables; ≥1) + flags: u32 (bit 0: this file is one segment of a multi-segment layer) + reserved: u32 +``` + +All integers little-endian. All region offsets are from the start of the containing segment file and 64-byte aligned. + +### 6.2 Bank descriptor (`num_banks ×`) + +``` + bank_id: u16 + bank_kind: u16 0=dense, 1=routed, 2=shared + region_schema_count: u16 number of schema records this bank owns in the + schema table (§6.4) — without it a reader cannot + tell where one bank's schemas end [draft-2] + flags: u16 bit 0-1: browse mode (00=none, 01=direct, + 10=strided) per §15.2; rest reserved [draft-2] + num_entries: u32 1 (dense) or expert count + input_dim: u32 + intermediate_dim: u32 + output_dim: u32 +``` + +Bank descriptor is 24 bytes (4-byte aligned), not the 20 the draft-1 field list implied. + +`input_dim`/`output_dim` are the expert's own operand dims — for K3's latent bank these are 3584/3584, not the 7168 residual width. Dense v1-style layers map to one bank: `bank_kind=0, num_entries=1`. + +**The binary carries no programme identity.** LYRW describes storage only — banks, entries, region schemas, offsets, formats. The MoE manifest binds `bank_id → programme` (§8.4). Two authorities for the same fact ("binary says programme 4, manifest says gpt-oss-expert-v1") is a disagreement waiting to happen; the manifest is the single one, consistent with the draft's own layering: storage holds regions, the manifest gives them meaning. + +### 6.3 Segment descriptor (`num_segments ×`) + +``` + bank_id: u16 + segment_index: u16 + first_entry: u32 + entry_count: u32 +``` + +A single-file layer has one segment covering `[0, num_entries)`. Multi-segment layers repeat the header in every segment file with `flags` bit 0 set; `index.json` lists the segment files per logical layer so the loader never globs. **A segment file's entry table covers only that segment's entries** (`entry_count` rows, indexed from `first_entry`) — never the whole logical bank. [clarified in draft-2] + +### 6.4 Region schemas and entry table + +Expert banks are homogeneous: every entry in a bank shares the same region layout. The region schema is therefore declared **once per bank**, and each entry stores only offsets and lengths. (This is the simplification that dropping LYRW v1 binary compatibility buys — see §6.6.) + +``` +[bank region schemas] region_schema_count × per bank: + schema_index: u16 + role: u16 (§6.5) + format: u16 quant enum — 0=f32 1=f16 2=bf16 3=q4_0 4=q4_k 5=q6_k + 6=q8_0 7=fp4_larql 8=mxfp4 9=nvfp4 10=mxfp8 ... + packing: u16 0=row_major, 1=blocks_with_scales_inline, + 2=blocks_values / 3=blocks_scales + pair_id: u16 links a blocks_values schema to its blocks_scales + schema; 0xFFFF = unpaired + reserved: u16 pad — keeps the two u32 dims on a 4-byte boundary; + record is 20 bytes, not draft-1's 18 [draft-2] + rows: u32 + cols: u32 + +[entry table] entry_count × region_schema_count ×: + offset: u64 (from start of containing segment file, 64-B aligned) + length: u64 +``` + +Consequences: + +- `gate_up_fused: mxfp4` + `down: q6_k` in one file, and GPT-OSS-style separate value/scale regions, without a new container. +- Uniform expert geometry is explicit; parsing is O(schemas), not O(entries × regions). +- Per-expert codec variation — which no grouped kernel supports — is **unrepresentable**, by construction rather than by convention. +- `pair_id` makes values/scales pairing explicit; role tags alone are ambiguous once an entry carries more than one quantised tensor. +- Exceptional per-entry overrides are reserved behind a header flag bit, undefined in v2.0 — added only if a real model forces them. + +### 6.5 Region roles + +Registered roles (extensible; new roles do not bump `format_version`): + +``` +0 gate +1 up +2 gate_up_fused +3 down +4 bias +5 scales (paired with a values region via packing=2/3) +6 latent_in (shared pre-projection, when bank-local storage is preferred) +7 latent_out (shared post-projection, likewise) +8..255 reserved-registered +256.. vendor/experimental +``` + +The fast-path contract is unchanged from v1: known kernels may **require** exactly `gate_up_fused + down` (or `gate + up + down`) and parse them into the same structures the grouped kernels use today. Presence of other roles does not invalidate a file; absence of a role a programme requires makes the file un-executable for that programme (§11), not invalid. + +**Unknown role, format and packing tags are preserved, not rejected, at read time.** Refusal belongs at capability-check time (§11): a browse-only reader must not choke on a `down` region encoded in a codec it never touches, and a future codec must not invalidate old readers' ability to serve the regions they do understand. The reader reports unknown tags; the capability check refuses the *operations* that need them. [clarified in draft-2] + +### 6.6 Relationship to the v1 layer files — greenfield, deliberately + +LYRW v2 owes **no binary compatibility** to the §5.12 `layers/*.weights` files. Those files are an internal detail of VINDEX2: they exist only inside VINDEX2 directories, are parsed only by the VINDEX2 loader path, and were never a public contract in their own right. No external tool depends on their byte layout. + +Consequences: + +- **No synthesis adapter.** A LYRW v2 reader never opens a v1 layer file, and vice versa. Each container generation's loader reads its own layer format, end of story. +- **No in-place upgrade** of multi-hundred-GB indexes. Migration is `checkpoint → VINDEX3 extractor`, or optionally `VINDEX2 → VINDEX3 importer` — a standalone tool, not a loader feature. +- **Design freedom.** The bank-level region-schema table (§6.4), explicit value/scale pairing, and segment descriptors are all clean-sheet choices that a v1-compat shim would have contaminated. The `LYRW` magic and `format_version=2` are retained purely as self-description and forensics — a v1 reader that encounters a v2 file fails fast on the version field with a precise "requires VINDEX3 loader" error, never a parse error. + +The compatibility obligation that **does** bind is one level up: larql must support VINDEX2 and VINDEX3 side by side (§12.1). + +--- + +## 7. Segmentation + +Motivating arithmetic (K3, exact Q6_K): + +``` +params per expert = 3 × 3584 × 3072 = 33,030,144 +params per routed layer = 33,030,144 × 896 = 29,595,009,024 +Q6_K bytes (210/256) ≈ 24.28 GB = 22.61 GiB +``` + +That exceeds the published 20 GiB shard cap, so `one logical layer = one physical file` cannot hold for K3 exact Q6_K. **Segment width and group width are two different scales, decided by two different measurements** — conflating them turns a 2-file layer into a 14-file layer for no read-path benefit: + +| Scale | Optimises | Typical size | +| ----- | --------- | ------------ | +| **Segment file** | file count, mmap management, shard distribution, the 20 GiB cap | as large as the cap allows — for K3 exact Q6_K, **2 segments of 448 experts** (~11.3 GiB each), not 14 of 64 | +| **Group extent** (inside a segment) | SSD reads, prefetch units, grouped-kernel dispatch | 8/16/32 experts (E2/E3) | + +A K3 routed layer therefore becomes: + +``` +routed/layer_037.seg00.weights experts 0–447 + ├── group extent 0: experts 0– 15 + ├── group extent 1: experts 16– 31 + └── ... (28 extents of 16) +routed/layer_037.seg01.weights experts 448–895 +``` + +At ~92 MoE layers this is ~184 routed segment files, not ~1,288. + +Rules: + +- Segment boundaries **must** fall on group-extent boundaries; group width **must** divide segment width. +- Both widths are extraction-time storage parameters chosen by measurement (E2 sweeps them independently), not semantic commitments. They may differ per model and per layer. +- Physical expert order within a segment need not equal logical order — the entry table is the indirection. Permuted layouts are legal but must not be adopted without the E2/E6 evidence bar. + +### 7.1 Group extents + +The unit of read alignment and prefetch is the **group extent** inside a segment, sized to the grouped kernel's natural dispatch width. One grouped dispatch ≈ one group extent ≈ one read unit; the extent boundary is what the payload layout aligns to, and the entry table makes extents addressable without a separate structure. Individual-expert files are prohibited at K3 scale (896 experts × ~92 MoE layers × several roles is an operational failure, not a design). + +--- + +## 8. The MoE programme manifest + +`moe_manifest.json` describes how regions form an MoE computation. The physical index stores tensor regions; the manifest gives them meaning; the runtime selects an optimised kernel when it recognises the programme. + +### 8.1 Per-layer shape + +```json +{ + "moe_layer": { + "layer": 12, + "input_space": "residual", + "router": { + "scores": "layers.12.router.weight", + "selection": { "kind": "top_k", "k": 16 }, + "normalisation": "k3_quantile_balanced" + }, + "transforms": { + "routed_input": "layers.12.routed_expert_down_proj", + "routed_output": "layers.12.routed_expert_up_proj" + }, + "routed_bank": { + "experts": 896, + "programme": "latent-moe-v1", + "storage": "routed/layer_012", + "expert_dims": { "input": 3584, "intermediate": 3072, "output": 3584 } + }, + "shared_bank": { + "experts": 2, + "programme": "gated-mlp-v1", + "storage": "shared/layer_012" + }, + "reduction": "gate_weighted_sum", + "routed_output_norm": "layers.12.routed_out_norm", + "combine": "residual_add" + } +} +``` + +For a conventional MoE, `transforms` are null. For GPT-OSS, `routed_bank.programme = "gpt-oss-expert-v1"` and `shared_bank` is absent. For Inkling, shared and routed banks coexist in residual space. Per-layer variation (hybrid dense+MoE stacks, differing expert counts) is expressed by per-layer manifests, not global fields. + +### 8.2 What stays model-specific (adapter-owned) + +Router scoring/normalisation details, shared-expert participation in normalisation, activation functions, expert residual semantics, clamps/biases/scales, fusion preferences, layer-specific expert counts. The manifest names these; the adapter implements them. + +### 8.3 Bounded programme vocabulary + +The declarative vocabulary is inference-shaped and closed by design: + +``` +linear · fused linear · activation · clamp · multiply · add · scale · +normalise · route · gather · weighted reduction · residual merge · +pre/post transform +``` + +No general graph interpreter. Known arrangements compile to specialised kernels; a generic reference executor provides correctness for everything representable. + +### 8.4 Programme registry + +``` +programme_id 0 gated-mlp-v1 + 1 gated-mlp-fused-fc1-v1 + 2 gpt-oss-expert-v1 (clamped gated MLP + residual term) + 3 shared-routed-mlp-v1 + 4 latent-moe-v1 (K3 SiTU-GLU latent expert) +``` + +Each programme declares its **required region roles**. New programmes register an id, a version, required roles, and optional opaque model metadata — region storage is untouched. + +The manifest is the **only** binding of `bank_id → programme_id`; LYRW files never carry programme identity (§6.2). Kernel capability entries (§10) reference programmes by registry id. + +--- + +## 9. Execution profiles and authority + +A profile is a small JSON file selecting inference behaviour over one extracted index. Profiles never trigger reslicing — and they never trigger conversion (§9.1). + +```json +{ + "profile": "routed-mxfp4", + "base": "exact", + "select": { + "routed.gate_up": "native-mxfp4", + "routed.down": "exact-q6k" + }, + "placement": { "routed": "local", "dense": "local" }, + "runtime_policy": { + "resident_experts": "routing-profile-2026-08-14.json", + "prefetch_group": 32 + } +} +``` + +The profile carries **no `authority` claim of its own** — authority is derived (§9.2). + +### 9.1 Representation variants — profiles select bytes, they don't request formats + +A profile saying `"format": "mxfp4"` cannot turn Q6_K bytes into MXFP4 bytes by declaration. Exactly one representation model is legal: **a region set may carry multiple physically present variants; a profile selects a present variant.** + +```json +{ + "region_set": "layer.12.routed.gate_up", + "variants": { + "exact-q6k": { "storage": "routed/layer_012.q6k", "fidelity": "source-equivalent" }, + "native-mxfp4": { "storage": "routed/layer_012.mxfp4", "fidelity": "source-exact" } + }, + "baseline": "exact-q6k" +} +``` + +- **Selecting an absent variant fails closed**, naming the region set, the requested variant and the variants actually present — before any byte is read. +- **No runtime conversion, ever.** "No hidden decode-time repacking" (§10) holds by construction: the bytes executed are the bytes stored. +- **Incremental packs.** New variants are added beside the baseline as independent, checksummed segment files — the multi-terabyte baseline is never rewritten. A routed-MXFP4 pack for K3 touches only routed region sets; attention, embeddings, routers and dense weights are untouched. +- **Single-copy, clarified.** The v1 principle forbids storing the *same* bytes twice; it does not forbid deliberate alternative encodings. The `baseline` variant is the canonical authority; additional variants are opt-in, per-component, and individually removable. + +### 9.2 Authority — graded, derived, never asserted + +**Levels** (mandatory, fail-closed): + +| Level | Meaning | +| ----- | ------- | +| `source-exact` | Decoded values bit-identical to the source checkpoint, in the checkpoint's own encoding family (e.g. native MXFP4 regions of a native-MXFP4 model) | +| `source-equivalent` | Different encoding whose decode reproduces the source values exactly (e.g. a lossless Q6_K container of native MXFP4 values) | +| `numerically-approximate` | Same architecture, lossy representation (e.g. Q6_K quantised from BF16) | +| `structurally-approximate` | Components omitted or replaced (reduced top-K, shared-only layers, compiled subexperts) — must list `omitted_components` / `replacement` | +| `analysis-only` | Incapable of complete forward execution (router/browse slices) | + +Authority is **derived, not declared**: every variant carries a region-level `fidelity` set at extraction time from provenance, and a profile's authority is the weakest fidelity across its active selections, further capped by programme traversal (§11) when required operands are absent. This closes the loophole where a lossy extraction becomes "exact" merely by being named the baseline — the baseline's own fidelity is recorded against the source checkpoint, not against itself. A profile cannot claim above its derived level; it may voluntarily claim below it. + +Standard profile names: `exact`, `native-lowbit`, `mixed-precision`, `attn-local-ffn-remote`, `partial-residency`, `reduced-top-k`, `shared-only`, `router-browse`, `compact-approximate`. + +**Runtime metadata, never format:** top-K/retention %, hot/warm/cold assignment, resident experts, per-layer popularity, adaptive cache size, prefetch ordering, exact-vs-approx selection, static per-layer precision choice. + +### 9.3 Omission semantics ("dropping down") + +The manifest distinguishes the materially different meanings: + +| Mode | Authority | Notes | +| ---- | --------- | ----- | +| Client omission (FFN remote) | inherits selection (up to source-exact) | The whole routed branch moves; the K3 latent boundary makes whole-branch RPC ~14 KB/layer f16 vs ~100 KB for projection-split — never split gate/up local from down remote absent contrary measurement | +| Analysis/router slice | analysis-only | Retains routers, gate vectors, metadata; no decode claim | +| Cheaper down representation | numerically-approximate | The production interpretation of "cheap down" | +| Down replaced by compact approximation | structurally-approximate | Must name `replacement` | +| Routed branch skipped (shared-only) | structurally-approximate | Dropping an expert's `w2` alone yields **no** expert output — the honest mode is skipping the expert/branch, not a half-expert | + +--- + +## 10. Kernel capability registry + +Kernels advertise what they can execute: + +``` +programme_id · region roles · formats per role · grouping widths · +input layout · maturity +``` + +Maturity ladder, matching the serving-format ledger: **Representable → Reference → Grouped → Dispatched → Production.** The loader reports, per (programme, format, grouping) combination, which rung it sits on. Mixed per-region formats are either supported by a kernel or **explicitly refused** — never silently repacked at decode time. + +--- + +## 11. Capability checking + +The loader does not hard-code "down weights present" tests. It traverses the layer's MoE programme and reports which required operands are absent, then: + +- refuses execution profiles whose authority claim exceeds what the present operands support; +- names the missing role, bank, layer and segment precisely (`VindexError::MissingRequiredRegion { layer, bank, role, .. }`); +- distinguishes *representable-but-no-kernel* (falls back to reference executor, flagged) from *operand-absent* (hard refusal). + +Programme-derived checks give the right per-architecture answers for free: routed removal on Inkling leaves shared experts contributing; on GPT-OSS it leaves no FFN contribution; on K3, a missing `routed_output` transform invalidates even completed expert computation. + +--- + +## 12. Versioning and coexistence + +Three version surfaces already exist; v2 adds nothing loosely named "vindex v2" in metadata. Precisely: + +| Contract | v1 value | v2 value | +| -------- | -------- | -------- | +| LYRW `format_version` | 1 (VINDEX2-internal) | **2** (self-description only; no cross-reading, §6.6) — trails the container generation by one, permanently | +| `index.json` `version` | 2 | **3** — the container-generation discriminator | +| `vindex_spec_version` | 1 | **2** (programme manifest + profiles enter the validated public contract) | +| MoE manifest schema | — | **1** (new) | + +**On the numbering.** The container generation *is* `index.json.version` — VINDEX2 is `version: 2`, VINDEX3 is `version: 3`. An earlier draft called the shipped generation "VINDEX1" while its `index.json.version` was already 2, putting a permanent off-by-one between the name and the sole discriminator. Both were renamed so the two agree. The LYRW layer format keeps its own sequence (v1 in VINDEX2, v2 in VINDEX3) and is deliberately not aligned: it is a different artifact with a different lifetime, and its numbering was already correct. + +The FP4 additive-extension precedent is retained within each generation: new region formats, roles and programme ids are enum additions, not format bumps. + +### 12.1 Dual-generation support in larql — the real compatibility contract + +The binding obligation is not between the two on-disk formats (there is none — §6.6). It is that **one larql binary supports both vindex generations, indefinitely for reading and serving**: + +- **Detection.** `index.json.version` is the sole **schema** discriminator, and the loader maps supported schema revisions to their owning container generation. No filename sniffing, no directory-shape heuristics. A missing or unknown version fails naming the version found and the schema sets this binary supports. + + The mapping is **many-to-one, not an identity**: + + | `index.json.version` | generation | note | + | -------------------- | ---------- | ---- | + | 1 | VINDEX2 | legacy schema; absent fields load with defaults | + | 2 | VINDEX2 | what a fresh VINDEX2 extraction writes | + | 3 | VINDEX3 | | + + A generation is *named* for the schema it currently writes, not for the only schema it can read. Treating the version as a generation identifier rather than a generation floor refuses every legacy-schema index in existence — which E0 caught in practice, not in review. Unified dispatch routes schema 1 to the VINDEX2 loader; the VINDEX3 loader still refuses it by name. +- **One entry point.** `Vindex::open(path)` returns the generation-appropriate handle behind a common trait; `larql run / serve / verify / slice / publish / pull` all accept either generation. Generation-specific verbs (e.g. profile selection) error precisely on a v1 index rather than silently no-op. +- **No cross-loading, no silent conversion.** The VINDEX2 loader path is frozen-but-maintained: it never opens VINDEX3 directories, never gains VINDEX3 features, and VINDEX3 code never re-implements VINDEX2 parsing. Conversion is only ever the explicit `VINDEX2 → VINDEX3` importer. +- **Hub and distribution.** `larql publish` stamps the container generation into the hub artifact metadata; `larql pull` selects the reader from that stamp and refuses a generation the installed binary lacks — before downloading terabytes, not after. +- **Wire protocols are generation-agnostic.** The expert-RPC and FFN-dispatch wire contracts carry activations and results, not container bytes; a grid may therefore mix v1 and v2 shards. A shard's container generation is a local concern of that shard's loader. +- **Support policy.** VINDEX2 remains fully supported for read/verify/serve/publish/pull. New extractions default to VINDEX3 once the ABI freezes **and** the E0 preservation matrix passes; v1 extraction remains available until then and is deprecated (not removed) after. + +--- + +## 13. Conformance envelope + +The ABI freezes only after all four fixtures pass the generic reference executor (fixtures defined in the experiments document): + +| Capability | Direct | GPT-OSS | IS-276B | KL-48B | K3 | +| ---------- | :----: | :-----: | :-----: | :----: | :-: | +| Variable expert count / top-K | ✓ | ✓ | ✓ | ✓ | ✓ | +| Routed experts | ✓ | ✓ | ✓ | ✓ | ✓ | +| Shared experts | – | – | ✓ (2) | ✓ (1) | ✓ | +| Shared-sink router (shared experts scored) | – | – | ✓ | – | – | +| Residual-space experts | ✓ | ✓ | ✓ | ✓ | – | +| Latent-space experts | – | – | – | – | ✓ | +| Hybrid dense+MoE stack | – | – | ✓ (mid-stack, idx 2) | ✓ (layer 0) | ✓* | +| Non-softmax router (sigmoid + scaling) | – | – | ✓ (+ gate bias, norm_after_topk) | ✓ | – | +| Custom expert programme | – | ✓ | – | – | ✓ | +| Native low-bit regions | ✓ | MXFP4 | NVFP4/MXFP8 (real releases) | – (BF16) | MXFP4 | +| Mixed per-role format | ✓ | ✓ | ✓ (release convention: routed low-bit, rest BF16) | ✓ | ✓ | +| Fused/decomposed tensors | ✓ | ✓ | ✓ | ✓ | ✓ | +| Grouped dispatch | ✓ | ✓ | ✓ | ✓ | ✓ | +| Auxiliary optional components (MTP, towers) | – | – | ✓ | – | ✓ (multimodal) | +| Single-segment routed layer | ✓ | ✓ | ✓ (~4.9 GiB/layer Q6_K) | ✓ (~1.4 GiB/layer Q6_K) | – | +| Segmented logical layer | – | – | – | – | ✓ | +| Exceeds-RAM residency (partial/remote non-optional) | – | – | ✓ | – | ✓ | +| WALK/DESCRIBE (residual-space browse) | ✓ | ✓ | ✓ | ✓ | – | +| WALK via latent transform (§15.4) | – | – | – | – | ✓ | + +\* K3's dense/MoE layer schedule is confirmed at adapter time; KL-48B's `first_k_dense_replace=1` is confirmed from the released config. + +Order of real-model implementation: **Gemma MoE → GPT-OSS → Kimi-Linear-48B-A3B → Inkling-Small → K3.** GPT-OSS is the first practical target (small, official reference paths). Kimi-Linear proves shared-expert banks, the hybrid stack and the KDA/MLA dense spine on a RAM-resident checkpoint — the K3 adapter dress rehearsal. Inkling-Small then escalates on two axes at once: real NVFP4/MXFP8 native regions with the routed-low-bit/rest-BF16 mixed-precision release convention, and the first *forced* partial-residency/remote serving (it cannot be RAM-resident on the rig) — the K3 **serving** dress rehearsal, as KL-48B is the adapter one. Fixture C is retained purely as the tiny deterministic conformance fixture; it no longer stands in for anything. K3 is extracted **once**, last, into the frozen ABI. + +--- + +## 14. What the experiments must decide + +Only these decisions genuinely belong in the on-disk ABI; everything else stays runtime policy: + +| Decision | Experiment | Why it matters | +| -------- | ---------- | -------------- | +| Region granularity (fused vs split roles) | E1, E4 | mmap count, rewriting, read amplification | +| Expert-group / segment width | E2, E3 | couples SSD reads to grouped kernels; K3 20 GiB cap | +| Fused vs decomposed FC1 storage | E1, E7, V2-1 | checkpoint import cost, mixed precision, **and gate-only browse reads** — the serving and query answers must be reconciled here, not assumed | +| Per-region format tags | structural (E4 gates *promotion* only) | representation is justified by native values/scales, v1's existing mixed precision, and format-neutral banks; E4 decides only whether a mixed-format **profile** reaches Production | +| Physical expert ordering | E2, E6 | possible locality gain vs model-specific assumption risk | +| Profile/variant-selection mechanism | E5, V2-0 | avoids reslicing per deployment; selection-not-request semantics (§9.1) | +| Capability/authority metadata | V2-0 | approximate slices must never present as exact | + +Registered prior (falsifiable): one file-set per routed layer (two segments for K3 Q6_K), one entry per expert, down independently addressable, locality as runtime metadata, omission = skip-the-branch, remote = whole-routed-branch RPC. Per-region format tags are in the ABI **structurally** (not gated on E4); the registered prior is that no mixed-format *profile* reaches Production before real-K3-layer evidence (E4 stage 3). Gate/up fusion is **no longer a prior** — it is a per-index extraction choice decided by E1/E7 (§15.2). + +--- + +## 15. Query layer — the model IS the database + +The LQL browse surface (WALK, DESCRIBE, SELECT, EXPLAIN WALK) is a first-class consumer of VINDEX3, with the same single-copy contract as v1: **no query index is stored beside the weights; the weights are the query index.** + +### 15.1 What replaces `gate_vectors.bin` + +There is no `gate_vectors.bin` in v2. The gate rows live where the split rule puts them — as `gate` (or the gate half of `gate_up_fused`) regions inside LYRW banks. Gate KNN mmaps the segment files and walks gate regions in place: + +- **f16/f32 regions:** zero-copy reinterpret, exactly the v1 fast path. +- **Block-quantised regions (FP4/FP8/Q-K):** lazy per-feature dequantisation at walk time via the existing block codecs — the v1 §5.10 mechanism, now applied to bank regions. The v1 §12.2 caveat carries over verbatim: 4-bit gate KNN is noisy; inference compensates, isolated dot products do not. +- Untouched `up`/`down` pages cost nothing under mmap, so browse over a full-fat index reads only gate bytes even when nothing was sliced. + +MoE browse semantics are unchanged from v1: gate KNN selects features **across all experts, no router needed** — a bank with `num_entries = E` simply contributes `E × intermediate_dim` walkable features per layer. Feature numbering stays v1-flattened (`layer:feature`, experts contiguous within the layer) so `feature_labels.json` keys survive migration untouched. + +### 15.2 The gate-addressability rule (resolves the fusion collision) + +A browse-enabled index requires gate rows to be readable without decoding up. Two legal ways to satisfy that: + +1. **Decomposed storage** (`gate` + `up` regions): clean gate-only reads; the E1/V2-1 fused-vs-decomposed parity requirement already guarantees kernels accept it. +2. **Fused storage with strided browse** (`gate_up_fused`): legal only when the packing permits striding into the gate half without decoding up rows (row-major f16 yes; interleaved quantised blocks generally no). + +The choice is recorded per bank at extraction time (a `browse: none | direct | strided` tag in the bank descriptor's flags, matching §6.2's normative encoding). **Serving-only indexes may fuse freely.** A browse-enabled index defaults to decomposed unless E1/E7 shows the fused serving advantage exceeds its own promotion bar — the previous blanket "gate/up stay fused" prior is withdrawn. + +### 15.3 Query metadata (`query/`) + +`down_meta.bin` (DMET, unchanged), `feature_labels.json` and `relation_clusters.json` move to `query/`. These are **derived metadata, not weight copies** — single-copy is not violated. Two v2-specific notes: + +- For latent MoE banks, `down_meta` is computed at extraction through the full output path — expert `w2` → `routed_output` transform → unembed — so its top-token claims describe residual-space effect, not raw latent columns. +- `query/` is optional per profile; its absence downgrades DESCRIBE/SELECT label richness, never WALK correctness. + +### 15.4 Browsing latent-space banks (the genuinely new problem) + +K3's gate rows live in the 3584-dim latent space; WALK queries originate in residual space. The programme manifest already carries what browse needs: `routed_input` names the residual→latent transform. WALK against a latent bank projects the query vector through that transform **once per query**, then dot-products against latent gate rows unchanged. `EXPLAIN WALK` reports the space hop. Residual-space banks (Direct, GPT-OSS, Inkling, all shared banks) walk exactly as v1. + +### 15.5 Browse profiles and slices + +- **Profile:** `browse` is a standard profile at authority `analysis-only` — requires gate regions (decodable), embeddings, tokenizer; `query/` and routers optional. Capability checking (§11) derives this; no filename tests. +- **Slice:** a published browse slice is produced by copying **only gate regions** into gate-only LYRW files (absent roles are legal, §6.5) plus `control/embeddings` + `query/`. The v1 ~3 GB browse economics are preserved; the loader reports the slice as `analysis-only` automatically because the programme's required inference operands are missing. + +### 15.6 Extract-level mapping + +| v1 extract level | v2 equivalent | +| ---------------- | ------------- | +| Browse | `browse` profile / gate-only slice (§15.5) | +| Inference | `exact` profile over classes 1–5 | +| All / COMPILE | full index — COMPILE reads regions to reconstruct safetensors, exactly as v1 read `gate_vectors.bin` | + +--- + +## 16. Success criteria — "done" is defined here, in advance + +VINDEX3 is a successful successor when all seven hold. Each is bound to the gate or experiment that proves it, so the bar cannot drift after the fact: + +| # | Criterion | Proven by | +| - | --------- | --------- | +| 1 | An existing VINDEX2 model loads, verifies, serves and publishes through the dual-generation binary with zero behavioural regression | E0 (continuous, CI) | +| 2 | Gemma and GPT-OSS run through the same LYRW2 bank machinery and the same production dispatch interface | V2-3, V2-4 rungs 1–2 | +| 3 | Routed **and** shared banks are genuinely generic — proven on a shared-expert model or fixture, not asserted | Fixture C + **KL-48B (1 shared) and Inkling-Small (2 shared, sink router)** — real, V2-4 rungs 3–4 | +| 4 | K3 is extracted once and served with no K3-specific physical layout — only a manifest and an adapter | V2-4 rung 4 | +| 5 | A new representation or placement is introduced via variants, profiles and kernel capabilities, without rebuilding unrelated weights | §9.1 mechanism + V2-0 profile-resolution acceptance | +| 6 | Unsupported or approximate configurations fail closed and report exactly why — operand, bank, role, layer, segment, variant | V2-0, §11 | +| 7 | Onboarding the **next** conventional MoE requires an importer and a programme adapter — zero format changes, zero new region roles, zero kernel-interface changes | **E8 held-out architecture** | + +Criterion 7 deserves emphasis: the four conformance fixtures cannot prove it, because the ABI was designed against them. Only a held-out architecture, onboarded after freeze under a no-format-changes rule, tests generalisation rather than fit. If E8 fails, the "portable sparse-serving substrate" claim is downgraded to "K3/GPT-OSS/Inkling serving format" — honestly, in this section. + +The maturity ladder governs claims throughout: **Representable → Reference → Grouped → Dispatched → Production.** No criterion is met by a representable-only demonstration. + +--- + +## License + +Apache-2.0 diff --git a/crates/larql-vindex/examples/vindex3_gemma_decode_parity.rs b/crates/larql-vindex/examples/vindex3_gemma_decode_parity.rs new file mode 100644 index 000000000..3b4f2a28e --- /dev/null +++ b/crates/larql-vindex/examples/vindex3_gemma_decode_parity.rs @@ -0,0 +1,390 @@ +//! Decode parity — the same model through a KV cache, two expert routes. +//! +//! Everything proven before this ran on the full-recompute prefill path, one +//! token, no KV reuse. This drives the real `KvEngine` decode loop, which is +//! how LARQL actually generates, and is the one execution condition the bound +//! route had never met. +//! +//! # Two stages, in this order +//! +//! ```text +//! 1 teacher-forced both engines fed the SAME token every step +//! 2 free-running each engine picks its own greedy token +//! ``` +//! +//! Stage 1 first, because a free-running run cannot localise. If step 3's +//! hidden state differs and step 4 is then fed a different token, step 4's +//! difference says nothing about step 4 — the comparison has branched and every +//! later number is about the branch. Feeding both engines an identical token +//! sequence keeps each step an independent measurement, exactly as the +//! locked-input sweep did for layers. +//! +//! Stage 2 only runs if stage 1 passes, and it is what promotes the claim from +//! "equivalent under controlled input" to "generates the same text". +//! +//! # Refusal is checked before any number is compared +//! +//! `moe_ffn_block_cpu` logs a refusing route, contributes zeros and returns the +//! dense half — so a missing operand arrives at a numeric comparison looking +//! like a wrong answer. Both routes are therefore driven through +//! [`MoeFfn::strict`], and every step consults `refusal()` *first*. A step that +//! refused is classified and stops the comparison; only a step where both +//! routes executed makes a numeric verdict meaningful. +//! +//! # Two locations, as before +//! +//! ```text +//! first differing boundary the first (step, what) that differs at all +//! earliest causal step the first step that differed while its INPUT +//! token and both prior states were identical +//! ``` +//! +//! Under teacher-forcing the input token is identical by construction, so the +//! two coincide unless a refusal intervened — which is precisely why the +//! refusal check comes first. +//! +//! Usage: +//! ```text +//! cargo run --release -p larql-vindex --example vindex3_gemma_decode_parity -- \ +//! --vindex [--prompt "The capital of France is"] [--tokens 8] +//! ``` + +use larql_inference::ffn::{BoundMoeBackend, InProcessMoeBackend, MoeFfn, RecordedRefusal}; +use larql_kv::EngineKind; +// / live on the engine trait. +use larql_models::ModelWeights; +use ndarray::Array2; + +const DEFAULT_PROMPT: &str = "The capital of France is"; +const DEFAULT_TOKENS: usize = 8; +const ARG_VINDEX: &str = "--vindex"; +const ARG_PROMPT: &str = "--prompt"; +const ARG_TOKENS: &str = "--tokens"; + +/// Enough to record a margin between the argmax and its nearest rival. +const TOP_K: usize = 5; +/// Unscaled: a parity comparison, not a sampling run. +const TEMPERATURE: f32 = 1.0; + +fn arg(name: &str) -> Option { + let args: Vec = std::env::args().collect(); + args.iter() + .position(|a| a == name) + .and_then(|i| args.get(i + 1)) + .cloned() +} + +fn max_abs_diff(a: &Array2, b: &Array2) -> f32 { + if a.shape() != b.shape() { + return f32::INFINITY; + } + a.iter() + .zip(b.iter()) + .map(|(x, y)| (x - y).abs()) + .fold(0.0f32, f32::max) +} + +/// The argmax token and its margin over the runner-up, through the same +/// production predictor both stages use. +fn predict( + weights: &ModelWeights, + h: &Array2, + tokenizer: &larql_inference::tokenizers::Tokenizer, +) -> Option<(u32, String, f64)> { + let p = larql_inference::forward::predict::logits_to_predictions_pub( + weights, + h, + tokenizer, + TOP_K, + TEMPERATURE, + ); + let id = *p.token_ids.first()?; + let (text, top) = p.predictions.first()?.clone(); + let margin = p.predictions.get(1).map_or(f64::INFINITY, |(_, s)| top - s); + Some((id, text, margin)) +} + +/// What happened at one decode step. +enum StepVerdict { + /// Both routes executed and agreed bit-for-bit. + Exact, + /// Both executed and disagreed. + Mismatch { diff: f32 }, + /// A route declined to execute. Not a numeric result at all. + Refused { + route: &'static str, + at: RecordedRefusal, + }, +} + +impl StepVerdict { + fn name(&self) -> String { + match self { + Self::Exact => "exact".into(), + Self::Mismatch { diff } => format!("mismatch max|Δ| = {diff:.3e}"), + Self::Refused { route, at } => { + format!("{} refused at layer {} — {}", route, at.layer, at.kind) + } + } + } + + fn agreed(&self) -> bool { + matches!(self, Self::Exact) + } +} + +/// Classify one step: refusal first, numbers only if both routes ran. +fn classify( + incumbent_ffn: &MoeFfn<'_>, + bound_ffn: &MoeFfn<'_>, + a: &Array2, + b: &Array2, +) -> StepVerdict { + // Refusal before arithmetic. A swallowed refusal reaches a numeric + // comparison as a wrong answer, and reporting it as one would name the + // wrong defect — the whole reason the strict adapter exists. + if let Some(at) = incumbent_ffn.refusal() { + return StepVerdict::Refused { + route: "in-process", + at, + }; + } + if let Some(at) = bound_ffn.refusal() { + return StepVerdict::Refused { route: "bound", at }; + } + let diff = max_abs_diff(a, b); + if diff == 0.0 { + StepVerdict::Exact + } else { + StepVerdict::Mismatch { diff } + } +} + +fn main() -> Result<(), String> { + let vindex = arg(ARG_VINDEX).ok_or(format!("set {ARG_VINDEX} "))?; + let prompt = arg(ARG_PROMPT).unwrap_or_else(|| DEFAULT_PROMPT.to_string()); + let tokens: usize = arg(ARG_TOKENS) + .and_then(|v| v.parse().ok()) + .unwrap_or(DEFAULT_TOKENS); + + println!("Decode parity — KvEngine, two expert routes"); + println!(" vindex {vindex}"); + println!(" prompt {prompt:?}"); + println!(" tokens {tokens}"); + + let mut callbacks = larql_vindex::SilentLoadCallbacks; + let path = std::path::Path::new(&vindex); + let mut weights = larql_vindex::load_model_weights_kquant(path, &mut callbacks) + .map_err(|e| format!("load weights: {e}"))?; + let mut index = larql_vindex::VectorIndex::load_vindex(path, &mut callbacks) + .map_err(|e| format!("load index: {e}"))?; + index + .load_attn_kquant(path) + .map_err(|e| format!("load attn Q4K: {e}"))?; + index + .load_interleaved_kquant(path) + .map_err(|e| format!("load interleaved Q4K: {e}"))?; + let _ = index.load_lm_head_kquant(path); + let tokenizer = + larql_vindex::load_vindex_tokenizer(path).map_err(|e| format!("tokenizer: {e}"))?; + let encoding = tokenizer + .encode(prompt.as_str(), true) + .map_err(|e| format!("encode prompt: {e}"))?; + let prompt_ids: Vec = encoding.get_ids().to_vec(); + + // Attention and the dense FFN slab, dequantised f32-resident for every + // layer — what the engine's resident path expects. Experts stay mapped; + // the bound route reads them where they lie. + println!("\n dequantising attn + dense FFN to f32-resident…"); + for layer in 0..weights.num_layers { + larql_inference::vindex::insert_q4k_layer_tensors_resident(&mut weights, &index, layer) + .map_err(|e| format!("resident dequant layer {layer}: {e}"))?; + } + let weights = weights; + println!( + " {} layers resident, {} prompt tokens", + weights.num_layers, + prompt_ids.len() + ); + + let incumbent_route = InProcessMoeBackend; + let bound_route = BoundMoeBackend::production(); + + // ── Stage 1: teacher-forced ──────────────────────────────────────────── + println!("\n== stage 1: teacher-forced =="); + let incumbent_ffn = MoeFfn::strict(&weights, &incumbent_route); + let bound_ffn = MoeFfn::strict(&weights, &bound_route); + let mut engine_a = + EngineKind::Standard { window_size: None }.build(larql_inference::cpu_engine_backend()); + let mut engine_b = + EngineKind::Standard { window_size: None }.build(larql_inference::cpu_engine_backend()); + + let h_a = engine_a + .prefill_resident(&weights, &incumbent_ffn, &index, &prompt_ids) + .map_err(|e| format!("in-process prefill: {e:?}"))?; + let h_b = engine_b + .prefill_resident(&weights, &bound_ffn, &index, &prompt_ids) + .map_err(|e| format!("bound prefill: {e:?}"))?; + + let mut first_differing: Option<(usize, String)> = None; + let mut earliest_causal: Option<(usize, String)> = None; + let mut forced: Vec = Vec::new(); + + let verdict = classify(&incumbent_ffn, &bound_ffn, &h_a, &h_b); + println!(" prefill {}", verdict.name()); + if !verdict.agreed() { + first_differing = Some((0, verdict.name())); + earliest_causal = Some((0, verdict.name())); + } + + // Teacher-forcing: the in-process route chooses, and BOTH engines are fed + // that token. The bound route's own preference is recorded but never acted + // on, so a divergence cannot branch the comparison. + let mut current = + predict(&weights, &h_a, &tokenizer).ok_or("in-process prefill produced no prediction")?; + let mut teacher_forced_ok = verdict.agreed(); + + for step in 1..=tokens { + forced.push(current.0); + let step_a = engine_a + .decode_step_resident(&weights, &incumbent_ffn, &index, current.0) + .map_err(|e| format!("in-process decode step {step}: {e:?}"))?; + let step_b = engine_b + .decode_step_resident(&weights, &bound_ffn, &index, current.0) + .map_err(|e| format!("bound decode step {step}: {e:?}"))?; + + let verdict = classify(&incumbent_ffn, &bound_ffn, &step_a, &step_b); + let pred_a = predict(&weights, &step_a, &tokenizer); + let pred_b = predict(&weights, &step_b, &tokenizer); + let tokens_agree = pred_a.as_ref().map(|p| p.0) == pred_b.as_ref().map(|p| p.0); + println!( + " step {step:<2} fed {:<12} {} | argmax {} {}", + format!("{:?}", current.1), + verdict.name(), + pred_a.as_ref().map_or("—".into(), |p| format!("{:?}", p.1)), + if tokens_agree { "==" } else { "!=" } + ); + + if !verdict.agreed() || !tokens_agree { + teacher_forced_ok = false; + let what = verdict.name(); + if first_differing.is_none() { + first_differing = Some((step, what.clone())); + } + // Under teacher-forcing every step receives an identical token by + // construction, so a step that differs is causal unless a refusal + // produced it — which `classify` has already separated out. + if earliest_causal.is_none() && matches!(verdict, StepVerdict::Mismatch { .. }) { + earliest_causal = Some((step, what)); + } + } + match pred_a { + Some(p) => current = p, + None => break, + } + } + + println!( + "\n first differing boundary {}", + describe(&first_differing) + ); + println!( + " earliest causal step {}", + describe(&earliest_causal) + ); + println!( + " forced sequence {:?}", + forced.iter().take(tokens).collect::>() + ); + + if !teacher_forced_ok { + return Err("teacher-forced decode parity failed — free-running not attempted".into()); + } + println!(" TEACHER-FORCED PASS: identical hidden state and argmax at every step."); + + // ── Stage 2: free-running ────────────────────────────────────────────── + // + // Only reached because stage 1 passed. Each engine now chooses its own + // token, so a divergence branches the comparison — which is exactly why it + // could not have come first. + println!("\n== stage 2: free-running greedy =="); + let incumbent_ffn = MoeFfn::strict(&weights, &incumbent_route); + let bound_ffn = MoeFfn::strict(&weights, &bound_route); + let mut engine_a = + EngineKind::Standard { window_size: None }.build(larql_inference::cpu_engine_backend()); + let mut engine_b = + EngineKind::Standard { window_size: None }.build(larql_inference::cpu_engine_backend()); + + let h_a = engine_a + .prefill_resident(&weights, &incumbent_ffn, &index, &prompt_ids) + .map_err(|e| format!("in-process prefill: {e:?}"))?; + let h_b = engine_b + .prefill_resident(&weights, &bound_ffn, &index, &prompt_ids) + .map_err(|e| format!("bound prefill: {e:?}"))?; + + let mut a_tokens: Vec = Vec::new(); + let mut b_tokens: Vec = Vec::new(); + let mut cur_a = predict(&weights, &h_a, &tokenizer).ok_or("no in-process prediction")?; + let mut cur_b = predict(&weights, &h_b, &tokenizer).ok_or("no bound prediction")?; + let mut free_ok = cur_a.0 == cur_b.0; + + for step in 1..=tokens { + a_tokens.push(cur_a.1.clone()); + b_tokens.push(cur_b.1.clone()); + let same = cur_a.0 == cur_b.0; + println!( + " step {step:<2} {:<14} vs {:<14} margin {:.4e} vs {:.4e} {}", + format!("{:?}", cur_a.1), + format!("{:?}", cur_b.1), + cur_a.2, + cur_b.2, + if same { "==" } else { "DIVERGED" } + ); + free_ok &= same; + + let next_a = engine_a + .decode_step_resident(&weights, &incumbent_ffn, &index, cur_a.0) + .map_err(|e| format!("in-process decode {step}: {e:?}"))?; + let next_b = engine_b + .decode_step_resident(&weights, &bound_ffn, &index, cur_b.0) + .map_err(|e| format!("bound decode {step}: {e:?}"))?; + if let Some(at) = bound_ffn.refusal() { + return Err(format!( + "bound route refused at step {step}, layer {}: {}", + at.layer, at.message + )); + } + match ( + predict(&weights, &next_a, &tokenizer), + predict(&weights, &next_b, &tokenizer), + ) { + (Some(a), Some(b)) => { + cur_a = a; + cur_b = b; + } + _ => break, + } + } + + println!("\n in-process {}", a_tokens.concat()); + println!(" bound {}", b_tokens.concat()); + + println!("\n== notes =="); + println!(" Attention, KV, the dense slab and the norms are the same code on both paths;"); + println!(" only the expert route differs. Refusals are checked before any numeric verdict,"); + println!(" so a missing operand is never reported as a wrong answer."); + + if free_ok { + println!("\nDECODE PARITY: identical under teacher-forcing and identical free-running."); + Ok(()) + } else { + Err("free-running decode diverged — see the first DIVERGED step".into()) + } +} + +fn describe(slot: &Option<(usize, String)>) -> String { + match slot { + None => "none".into(), + Some((step, what)) => format!("step {step}, {what}"), + } +} diff --git a/crates/larql-vindex/examples/vindex3_gemma_free_propagation.rs b/crates/larql-vindex/examples/vindex3_gemma_free_propagation.rs new file mode 100644 index 000000000..9e160c440 --- /dev/null +++ b/crates/larql-vindex/examples/vindex3_gemma_free_propagation.rs @@ -0,0 +1,404 @@ +//! Free-propagation composition — each path carries its own residual. +//! +//! The locked-input sweep proved every MoE layer independently. It could not +//! prove that they *compose*, because it handed every layer the incumbent's own +//! activation. This one starts both paths from the same embedding and lets each +//! carry its own residual to the end: +//! +//! ```text +//! same token ids → same embedding +//! ↓ ↓ +//! in-process MoE VINDEX3-bound MoE +//! ↓ own residual ↓ own residual +//! …30 blocks… …30 blocks… +//! ↓ ↓ +//! compare, boundary by boundary +//! ``` +//! +//! Attention, KV, the dense slab, the norms, PLE and the layer scalar are the +//! *same code* on both sides — both routes go through `predict_kquant_hidden` +//! and differ only in which `MoeExpertBackend` is installed. So a divergence is +//! the MoE path or it is nothing. +//! +//! # Two locations, never conflated +//! +//! ```text +//! first differing boundary where a difference is first observable +//! earliest causal operation the first MoE that differed on identical input +//! ``` +//! +//! Those are not the same layer, and reporting only the first would be +//! actively misleading. Once layer 7's contribution differs, layer 8 reads a +//! contaminated residual and *everything* downstream differs — including its +//! post-attention state, which would then read as an attention defect. The +//! causal test is therefore conditional: a layer only indicts its own MoE if +//! its input residual was still identical when it ran. +//! +//! # Why the incumbent also runs through a backend +//! +//! `InProcessMoeBackend` is byte-identical to the block loop's default branch. +//! Putting the incumbent behind the same trait gives the harness a seam to +//! observe on both sides — and makes the seam itself falsifiable, which is what +//! the first check below does before any comparison is trusted. +//! +//! Usage: +//! ```text +//! cargo run --release -p larql-vindex --example vindex3_gemma_free_propagation -- \ +//! --vindex [--prompt "The capital of France is"] +//! ``` + +use std::cell::RefCell; + +use larql_inference::ffn::{ + BoundMoeBackend, InProcessMoeBackend, MoeBackendError, MoeExpertBackend, +}; +use larql_inference::vindex::predict_kquant_hidden; +use larql_models::ModelWeights; +use ndarray::Array2; + +/// Residual-scale values, judged relative to the reference's own magnitude — +/// the same policy as the single-layer harness, and for the same reason. +const RELATIVE_BAND: f32 = 0.05; +/// Below this the ratio stops meaning anything and the comparison is absolute. +const SCALE_FLOOR: f32 = 1e-6; + +const DEFAULT_PROMPT: &str = "The capital of France is"; +const ARG_VINDEX: &str = "--vindex"; +const ARG_PROMPT: &str = "--prompt"; +/// Enough to record a margin between the argmax and its nearest rival. +const TOP_K: usize = 5; +/// Unscaled: this is a parity comparison, not a sampling run. +const TEMPERATURE: f32 = 1.0; + +fn arg(name: &str) -> Option { + let args: Vec = std::env::args().collect(); + args.iter() + .position(|a| a == name) + .and_then(|i| args.get(i + 1)) + .cloned() +} + +fn max_abs_diff(a: &Array2, b: &Array2) -> f32 { + if a.shape() != b.shape() { + return f32::INFINITY; + } + a.iter() + .zip(b.iter()) + .map(|(x, y)| (x - y).abs()) + .fold(0.0f32, f32::max) +} + +fn magnitude(a: &Array2) -> f32 { + a.iter().fold(0.0f32, |m, v| m.max(v.abs())) +} + +/// What one layer's MoE call saw and produced. +#[derive(Clone)] +struct Observation { + layer: usize, + input: Array2, + contribution: Array2, +} + +/// Wraps a route and records every call, so both paths can be compared at the +/// same boundaries without either being reimplemented. +struct Recording { + inner: B, + seen: RefCell>, +} + +impl Recording { + fn new(inner: B) -> Self { + Self { + inner, + seen: RefCell::new(Vec::new()), + } + } + + fn take(self) -> Vec { + self.seen.into_inner() + } +} + +impl MoeExpertBackend for Recording { + fn forward_moe_seq( + &self, + weights: &ModelWeights, + layer: usize, + h: &Array2, + norm_offset: f32, + eps: f32, + ) -> Result, MoeBackendError> { + let contribution = self + .inner + .forward_moe_seq(weights, layer, h, norm_offset, eps)?; + self.seen.borrow_mut().push(Observation { + layer, + input: h.clone(), + contribution: contribution.clone(), + }); + Ok(contribution) + } + + fn name(&self) -> &'static str { + self.inner.name() + } +} + +/// Where a divergence is first *visible*, and where it was first *caused*. +#[derive(Debug, Default)] +struct DivergenceReport { + /// The first layer at which anything observable differs — input residual + /// or contribution, whichever comes first. + first_differing_boundary: Option<(usize, &'static str, f32)>, + /// The first layer whose MoE contribution differed **while its input + /// residual was still identical**. That layer's expert path is the defect; + /// everything after it is contamination. + earliest_causal_operation: Option<(usize, f32)>, +} + +fn diverge(incumbent: &[Observation], bound: &[Observation]) -> DivergenceReport { + let mut report = DivergenceReport::default(); + for (a, b) in incumbent.iter().zip(bound) { + debug_assert_eq!(a.layer, b.layer); + let input_diff = max_abs_diff(&a.input, &b.input); + let contribution_diff = max_abs_diff(&a.contribution, &b.contribution); + + if report.first_differing_boundary.is_none() && input_diff != 0.0 { + report.first_differing_boundary = Some((a.layer, "input residual", input_diff)); + } + if report.first_differing_boundary.is_none() && contribution_diff != 0.0 { + report.first_differing_boundary = + Some((a.layer, "moe contribution", contribution_diff)); + } + // The conditional that keeps contamination from reading as a defect: + // a layer only indicts its own MoE if it ran on an identical input. + if report.earliest_causal_operation.is_none() + && input_diff == 0.0 + && contribution_diff != 0.0 + { + report.earliest_causal_operation = Some((a.layer, contribution_diff)); + } + } + report +} + +fn report_relative(label: &str, reference: &Array2, found: &Array2) -> bool { + let diff = max_abs_diff(reference, found); + if diff == 0.0 { + println!(" {label:<26} BIT-IDENTICAL"); + return true; + } + let scale = magnitude(reference); + if scale < SCALE_FLOOR { + println!( + " {label:<26} max|Δ| = {diff:.3e}, reference magnitude {scale:.3e} \ + below the {SCALE_FLOOR:.0e} floor — absolute, not a ratio" + ); + return diff < SCALE_FLOOR; + } + let relative = diff / scale; + println!( + " {label:<26} max|Δ| = {diff:.3e} of {scale:.3e} = {:.3}% {}", + relative * 100.0, + if relative <= RELATIVE_BAND { + "within band" + } else { + "OUTSIDE" + } + ); + relative <= RELATIVE_BAND +} + +fn main() -> Result<(), String> { + let vindex = arg(ARG_VINDEX).ok_or(format!("set {ARG_VINDEX} "))?; + let prompt = arg(ARG_PROMPT).unwrap_or_else(|| DEFAULT_PROMPT.to_string()); + + println!("Free-propagation composition — each path carries its own residual"); + println!(" vindex {vindex}"); + println!(" prompt {prompt:?}"); + + let mut callbacks = larql_vindex::SilentLoadCallbacks; + let path = std::path::Path::new(&vindex); + let weights = larql_vindex::load_model_weights_kquant(path, &mut callbacks) + .map_err(|e| format!("load weights: {e}"))?; + let mut index = larql_vindex::VectorIndex::load_vindex(path, &mut callbacks) + .map_err(|e| format!("load index: {e}"))?; + // The same three region maps the CLI installs. Without them the forward + // panics on the attention slices before ever reaching an MoE block. + index + .load_attn_kquant(path) + .map_err(|e| format!("load attn Q4K: {e}"))?; + index + .load_interleaved_kquant(path) + .map_err(|e| format!("load interleaved Q4K: {e}"))?; + let _ = index.load_lm_head_kquant(path); + let tokenizer = + larql_vindex::load_vindex_tokenizer(path).map_err(|e| format!("tokenizer: {e}"))?; + let encoding = tokenizer + .encode(prompt.as_str(), true) + .map_err(|e| format!("encode prompt: {e}"))?; + let token_ids: Vec = encoding.get_ids().to_vec(); + println!(" tokens {}\n", token_ids.len()); + + // ── The seam is falsifiable before anything is built on it ───────────── + // + // Running through `InProcessMoeBackend` must equal running with no backend + // at all. If it does not, the trait changed the model rather than + // relocating a call, and every comparison below would be measuring that + // change instead of VINDEX3. + println!("== seam =="); + let default_route = predict_kquant_hidden(&weights, &token_ids, &index, None); + let via_trait = predict_kquant_hidden(&weights, &token_ids, &index, Some(&InProcessMoeBackend)); + let seam_faithful = max_abs_diff(&default_route, &via_trait) == 0.0; + println!( + " in-process via trait vs default branch: {}", + if seam_faithful { + "BIT-IDENTICAL — the seam relocates the call and nothing else" + } else { + "DIFFERS — the trait changed the model; stop here" + } + ); + if !seam_faithful { + return Err("the seam is not faithful; no comparison below is meaningful".into()); + } + + // ── Both routes, each carrying its own residual ──────────────────────── + let incumbent = Recording::new(InProcessMoeBackend); + let incumbent_h = predict_kquant_hidden(&weights, &token_ids, &index, Some(&incumbent)); + let incumbent_seen = incumbent.take(); + + let bound = Recording::new(BoundMoeBackend::production()); + let bound_h = predict_kquant_hidden(&weights, &token_ids, &index, Some(&bound)); + let bound_seen = bound.take(); + + println!("\n== boundaries =="); + println!( + " {} MoE layers observed on each path", + incumbent_seen.len() + ); + if incumbent_seen.len() != bound_seen.len() { + return Err(format!( + "the two paths ran different numbers of MoE layers: {} vs {}", + incumbent_seen.len(), + bound_seen.len() + )); + } + + let report = diverge(&incumbent_seen, &bound_seen); + match report.first_differing_boundary { + None => println!(" first differing boundary none — every boundary bit-identical"), + Some((layer, what, diff)) => { + println!(" first differing boundary layer {layer}, {what}, max|Δ| = {diff:.3e}") + } + } + match report.earliest_causal_operation { + None => { + println!(" earliest causal operation none — no MoE differed on an identical input") + } + Some((layer, diff)) => println!( + " earliest causal operation layer {layer} MoE, max|Δ| = {diff:.3e} \ + (its input was still identical, so this layer is the defect)" + ), + } + if let (Some((first, _, _)), Some((causal, _))) = ( + report.first_differing_boundary, + report.earliest_causal_operation, + ) { + if first != causal { + println!( + " note: layers {first}..{causal} differ downstream of the cause — contamination, \ + not independent defects" + ); + } + } + + // ── Composition ──────────────────────────────────────────────────────── + println!("\n== final hidden state =="); + let composed = report_relative("pre-final-norm hidden", &incumbent_h, &bound_h); + + // ── The model's own endpoint ─────────────────────────────────────────── + // + // Both boundaries go through production functions rather than a local + // reimplementation: `apply_norm` is the one `logits_to_predictions` calls, + // and the predictor is the one generation calls. A harness that rolled its + // own lm_head here would be comparing the harness. + println!("\n== endpoint =="); + let norm_offset = weights.arch.norm_weight_offset(); + let final_key = weights.arch.final_norm_key(); + let incumbent_normed = + larql_inference::forward::apply_norm(&weights, &incumbent_h, final_key, norm_offset); + let bound_normed = + larql_inference::forward::apply_norm(&weights, &bound_h, final_key, norm_offset); + let normed = report_relative("post-final-norm hidden", &incumbent_normed, &bound_normed); + + let incumbent_pred = larql_inference::forward::predict::logits_to_predictions_pub( + &weights, + &incumbent_h, + &tokenizer, + TOP_K, + TEMPERATURE, + ); + let bound_pred = larql_inference::forward::predict::logits_to_predictions_pub( + &weights, + &bound_h, + &tokenizer, + TOP_K, + TEMPERATURE, + ); + + let ids_match = incumbent_pred.token_ids == bound_pred.token_ids; + let scores_match = incumbent_pred.predictions == bound_pred.predictions; + println!( + " {:<26} {}", + "top-k token ids", + if ids_match { "IDENTICAL" } else { "DIFFER" } + ); + println!( + " {:<26} {}", + "top-k scores", + if scores_match { + "BIT-IDENTICAL" + } else { + "DIFFER" + } + ); + + // The margin is recorded even though identical scores imply it. Later + // approximate and partially-resident runs need this same report shape, and + // a margin that only appears once something disagrees is a shape nobody can + // compare against. + let margin = |p: &larql_compute::forward::predict::PredictResult| -> f64 { + match (p.predictions.first(), p.predictions.get(1)) { + (Some((_, top)), Some((_, second))) => top - second, + _ => f64::INFINITY, + } + }; + let argmax_match = incumbent_pred.token_ids.first() == bound_pred.token_ids.first(); + println!( + " {:<26} {:?} vs {:?}", + "argmax token", + incumbent_pred.predictions.first().map(|(t, _)| t), + bound_pred.predictions.first().map(|(t, _)| t) + ); + println!( + " {:<26} {:.6e} vs {:.6e}", + "argmax margin", + margin(&incumbent_pred), + margin(&bound_pred) + ); + + let endpoint = normed && ids_match && scores_match && argmax_match; + + println!("\n== notes =="); + println!(" Attention, KV, the dense slab, the norms, PLE and the layer scalar are the same"); + println!(" code on both paths. Only the MoE route differs, so a divergence is the MoE path."); + if report.earliest_causal_operation.is_none() && composed && endpoint { + println!("\nCOMPOSITION: every MoE ran on an identical input and produced an identical"); + println!("contribution; the free-propagated hidden states, the normed hidden states and"); + println!("the model's own predictions all agree."); + Ok(()) + } else { + Err("composition not established — see the causal layer above".into()) + } +} diff --git a/crates/larql-vindex/examples/vindex3_gemma_layer_parity.rs b/crates/larql-vindex/examples/vindex3_gemma_layer_parity.rs new file mode 100644 index 000000000..8567712a9 --- /dev/null +++ b/crates/larql-vindex/examples/vindex3_gemma_layer_parity.rs @@ -0,0 +1,862 @@ +//! Gemma one-layer parity — VINDEX3 bound over VINDEX2's own expert bytes. +//! +//! The first real-model step. It holds *everything* constant except the +//! execution path: +//! +//! ```text +//! one VINDEX2 index on disk +//! ↓ +//! the same Q4_K expert bytes, the same f32 router, the same activation +//! ↓ ↓ +//! incumbent MoE path BoundMoeOperation bound over those bytes +//! ↓ ↓ +//! compare, checkpoint by checkpoint +//! ``` +//! +//! # Why not extract a VINDEX3 container first +//! +//! Because then a mismatch would have two candidate causes — the executor, or +//! the re-extraction — and telling those apart is the entire point. Binding +//! over the incumbent's own bytes makes any divergence unambiguously about +//! execution: routing policy, region interpretation, activation, or reduction. +//! +//! # The ladder +//! +//! Rung by rung, so agreement at the bottom cannot conceal compensating +//! differences higher up: +//! +//! ```text +//! 1-5 router scores, selection, margin, pre-norm and final weights +//! 6-8 experts per-expert outputs, the weighted reduction, the block +//! ``` +//! +//! Rungs 1-5 bind `larql-compute`'s scoring functions; rungs 6-8 bind its +//! Q4_K × Q8_K expert kernel over the store's own super-blocks. Both are +//! *bindings*, not reimplementations — writing a lookalike here and calling the +//! agreement "parity" would prove only that two similar loops agree. +//! +//! # Three paths, not two +//! +//! ```text +//! incumbent cpu_moe_forward and the functions underneath it +//! bound VINDEX3 over the same Q4_K bytes, production kernels +//! reference VINDEX3 over the same weights dequantised, oracle kernels +//! ``` +//! +//! The third exists because two calls of one function agree whatever operands +//! they are handed. A bit-identical result says the handover was faithful; only +//! an independently-computed answer says the operands were the right ones. +//! +//! # What is deliberately *not* claimed +//! +//! That the reference agrees to the bit. It dequantises to f32 where the +//! incumbent keeps an integer dot against a Q8_K activation, so the two differ +//! by quantisation noise by construction. That leg is judged against a +//! *relative* band, because expert outputs are residual-scale values rather +//! than probabilities and an absolute figure would have to be re-derived per +//! layer. +//! +//! Nor is the *schedule* claimed. The incumbent sums its experts through rayon +//! or the spin pool depending on configuration, and a tree reduction is not +//! required to agree bit-for-bit with a sequential one. Rung 7 therefore +//! compares against the incumbent's own per-expert outputs summed in selection +//! order, which isolates the combine from the scheduling; rung 8 reports the +//! end-to-end block figure without asserting it. Rung 8 came out bit-identical +//! on the spin-pool path — which accumulates in selection order — and that is a +//! measurement of one configuration, not a promise about every one. +//! +//! # Capturing the activation +//! +//! ```text +//! LARQL_CPU_DUMP_LAYERS=/tmp/gemma_dump \ +//! larql run "The capital of France is" -n 1 +//! ``` +//! +//! writes `cpu_layer_NN_h_post_attn.f32` — the real residual entering each +//! layer's FFN/MoE block. +//! +//! Usage: +//! ```text +//! cargo run --release -p larql-vindex --example vindex3_gemma_layer_parity -- \ +//! --vindex --dump /tmp/gemma_dump [--layer 5] +//! ``` + +use larql_compute::cpu::ops::moe::{ + cpu_moe_forward, moe_expert_input, moe_post_expert_output, moe_route_from_router_input, + moe_router_input, moe_score_experts, moe_softmax, quantize_x_to_q8k, + run_single_expert_q4k_q8k_into, ExpertScratch, +}; +use larql_compute::cpu::ops::q4_common::dequantize_q4_k; +use larql_compute::pipeline_layer::build_moe_weights; +use larql_compute::MoeLayerWeights; + +use larql_vindex::format::capability::binding::{ComponentView, RepresentationIdentity}; +use larql_vindex::format::capability::component::ComponentContract; +use larql_vindex::format::capability::coordinate::BankCoordinate; +use larql_vindex::format::lyrw2::region_format::RegionFormat; +use larql_vindex::format::lyrw2::region_role::RegionRole; +use larql_vindex::runtime::consts::{COL_DIM, FUSED_PROJECTION_HALVES}; +use larql_vindex::runtime::{ + execute_traced, BoundBankOperation, BoundExpert, BoundExpertScaling, BoundMoeOperation, + BoundProjection, BoundReduction, BoundRouter, BoundTensor, ExpertKernel, MoeInputs, + RouterKernel, +}; + +/// Band for the router's gate weights. They are softmax probabilities, so an +/// absolute figure is meaningful: the quantities being compared are all in +/// `[0, 1]` and sum to one. +const WEIGHT_TOLERANCE: f32 = 1e-3; + +/// Band for the oracle leg, as a fraction of the reference's own magnitude. +/// +/// Relative, and not the weight tolerance, because expert outputs are not +/// probabilities — they are residual-scale values whose magnitude is a property +/// of the layer. Judging them against a band derived from softmax outputs +/// compares the right direction on the wrong object. +/// +/// The size is derivable rather than fitted. Both paths read *identical* +/// weights: the reference dequantises exactly the Q4_K blocks the kernel reads, +/// so the weight side contributes nothing. The whole difference is the +/// activation side — Q8_K stores `d = amax / 127`, bounding each element's +/// representation error at `d / 2`, and the intermediate is quantised a second +/// time before `down`. Two such roundings carried through a 2816-wide +/// contraction put the expected disagreement in the low percent. Five percent +/// is a band around that, not a target: a swapped gate/up half or a mis-strided +/// `down` misses it by orders of magnitude rather than by a factor. +const ORACLE_RELATIVE_BAND: f32 = 0.05; + +/// Reference magnitude below which a ratio stops being reportable. +/// +/// The denominator of the oracle figure is `max|reference|`. Near zero that +/// division amplifies a rounding difference into a number that looks +/// catastrophic and compares to nothing, so below this floor the comparison +/// switches to an absolute one and says so. Named, and checked against the +/// reference rather than the found values, so a layer's percentage stays +/// comparable to every other layer's. +const ORACLE_SCALE_FLOOR: f32 = 1e-6; +/// Selection is discrete. Any difference is a real disagreement. +const DEFAULT_LAYER: usize = 5; +const VARIANT: &str = "vindex2-bytes"; +const ROUTER_REGION_SET: &str = "router"; +const PER_EXPERT_SCALE_REGION_SET: &str = "router_per_expert_scale"; +/// The single bank this one-layer harness binds. +const BANK_ID: u16 = 0; +/// Column the ladder's verdicts line up in. +const LADDER_LABEL_WIDTH: usize = 24; + +fn arg(name: &str) -> Option { + let args: Vec = std::env::args().collect(); + args.iter() + .position(|a| a == name) + .and_then(|i| args.get(i + 1)) + .cloned() +} + +/// Read an f32 dump and return its final row — the token being decoded. +fn last_row(path: &str, width: usize) -> Result, String> { + let bytes = std::fs::read(path).map_err(|e| format!("{path}: {e}"))?; + let values: Vec = bytes + .chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect(); + if !values.len().is_multiple_of(width) || values.is_empty() { + return Err(format!( + "{path}: {} values is not a whole number of {width}-wide rows", + values.len() + )); + } + let rows = values.len() / width; + Ok(values[(rows - 1) * width..].to_vec()) +} + +fn max_abs_diff(a: &[f32], b: &[f32]) -> f32 { + a.iter() + .zip(b) + .map(|(x, y)| (x - y).abs()) + .fold(0.0f32, f32::max) +} + +fn report(label: &str, a: &[f32], b: &[f32]) -> bool { + if a.len() != b.len() { + println!( + " {label: f32 { + values.iter().fold(0.0f32, |m, v| m.max(v.abs())) +} + +/// Compare two vectors relative to the **reference's** own magnitude. +/// +/// An absolute band would have to be re-derived for every layer, because a +/// residual-scale quantity's size is a property of the layer rather than of the +/// arithmetic. The ratio is the same question asked in units that travel. +/// +/// # The denominator is fixed, and it is stated +/// +/// `max|reference|`, never `max|found|` and never the larger of the two. A +/// denominator that moved with the thing being measured would let a kernel that +/// inflates its output report a *smaller* percentage for a larger error, and +/// would make two layers' figures incomparable. The oracle is the fixed point; +/// the bound path is what is being judged against it. +/// +/// # Zero-near-zero is a separate verdict, not a divide +/// +/// Below [`ORACLE_SCALE_FLOOR`] the ratio stops meaning anything — a reference +/// of 1e-9 turns any rounding difference into thousands of percent — so this +/// reports the absolute difference against the floor and labels it, rather than +/// dividing and emitting a number that looks like the others but is not +/// comparable to them. A layer whose output is genuinely near zero should read +/// as such rather than as a catastrophic relative error. +fn report_relative(label: &str, reference: &[f32], found: &[f32]) -> bool { + if reference.len() != found.len() { + println!( + " {label: Result<(Vec, ComponentContract), String> { + let values = dequantize_q4_k(bytes, rows * cols); + if values.len() < rows * cols { + return Err(format!( + "{}: dequantised {} of {} expected elements", + role.name(), + values.len(), + rows * cols + )); + } + let raw: Vec = values[..rows * cols] + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect(); + Ok((raw, ComponentContract::matrix(rows as u32, cols as u32))) +} + +fn tensor<'a>( + region_set: &str, + bytes: &'a [u8], + format: RegionFormat, + contract: ComponentContract, +) -> Result, String> { + BoundTensor::direct( + RepresentationIdentity::new(region_set, VARIANT), + bytes, + format, + contract, + ) + .map_err(|e| e.to_string()) +} + +/// Bind a stored matrix whose trailing columns are quantisation padding. +/// +/// The role sees `[rows, keep]`; the bytes remain `[rows, stored_cols]`. No +/// repacking, no copy — the view resolves the difference at read time. +/// +/// Both kernels read this one binding and take from it what each needs: the +/// reference reads the `keep` live columns and never touches the padding, +/// while the Q4_K kernel — which decodes whole super-blocks and cannot stop +/// inside one — reads all `stored_cols` and lets the zero-padded activation +/// cancel the rest. That is why the view belongs to the operand rather than +/// being something each kernel is told separately. +fn sliced_tensor<'a>( + region_set: &str, + bytes: &'a [u8], + format: RegionFormat, + storage: ComponentContract, + keep: usize, +) -> Result, String> { + BoundTensor::new( + RepresentationIdentity::new(region_set, VARIANT), + bytes, + format, + storage, + ComponentView::Slice { + dim: COL_DIM, + start: 0, + len: keep as u32, + }, + ) + .map_err(|e| e.to_string()) +} + +/// Owns the dequantised expert buffers so bound tensors can borrow them. +struct ExpertBuffers { + expert_id: u32, + gate_up: Vec, + down: Vec, + gate_up_contract: ComponentContract, + down_contract: ComponentContract, +} + +/// Assemble the operation around one bank of experts. +/// +/// Written once and called twice. Two hand-written literals meant to differ in +/// exactly one field is how a parity harness acquires a second difference +/// nobody notices. +fn bind_operation<'a>( + layer: usize, + hidden: usize, + moe: &MoeLayerWeights<'_>, + router_bytes: &'a [u8], + scale_bytes: &'a [u8], + experts: Vec>, + kernel: ExpertKernel, +) -> Result, String> { + Ok(BoundMoeOperation { + router: BoundRouter { + weight: tensor( + ROUTER_REGION_SET, + router_bytes, + RegionFormat::F32, + ComponentContract::matrix(moe.num_experts as u32, hidden as u32), + )?, + top_k: moe.top_k, + selected_weight: moe.routing_policy.selected_weight, + scaling: if scale_bytes.is_empty() { + BoundExpertScaling::None + } else { + BoundExpertScaling::PerExpert { + scales: tensor( + PER_EXPERT_SCALE_REGION_SET, + scale_bytes, + RegionFormat::F32, + ComponentContract::vector(moe.num_experts as u32), + )?, + } + }, + // Rung 1: bind the production scoring kernel, not a lookalike. + kernel: RouterKernel::Incumbent, + }, + transforms: Vec::new(), + banks: vec![BoundBankOperation { + bank: BankCoordinate::new(layer as u32, BANK_ID), + experts, + intermediate_dim: moe.intermediate_size, + hidden_dim: hidden, + activation: moe.activation, + kernel, + }], + reduction: BoundReduction::WeightedSum, + residual_dim: hidden, + }) +} + +/// Run the incumbent's own expert kernel over each selected expert, in +/// selection order, sharing one Q8_K activation exactly as production does. +fn incumbent_expert_outputs( + moe: &MoeLayerWeights<'_>, + expert_input: &[f32], + selected: &[usize], +) -> Result>, String> { + let hidden = expert_input.len(); + let inter = moe.intermediate_size; + let q8k = quantize_x_to_q8k(expert_input); + let mut scratch = ExpertScratch::new(hidden, inter, moe.inter_padded()); + selected + .iter() + .map(|&e| { + let gate_up = *moe + .experts_gate_up + .get(e) + .ok_or_else(|| format!("expert {e} has no gate_up bytes"))?; + let down = *moe + .experts_down + .get(e) + .ok_or_else(|| format!("expert {e} has no down bytes"))?; + Ok(run_single_expert_q4k_q8k_into( + &mut scratch, + &q8k, + gate_up, + down, + inter, + moe.activation, + ) + .to_vec()) + }) + .collect() +} + +/// `sum_i weight_i * output_i`, in selection order. +/// +/// The incumbent's parallel paths reach this same sum through a rayon +/// tree-reduce or a spin-pool slot scan, neither of which fixes an order. So +/// the reduction is compared against *this* — the incumbent's own per-expert +/// outputs, combined in the order VINDEX3 combines them — which isolates the +/// combine from the schedule. Blaming a scheduling difference on the reduction +/// is exactly the false diagnosis a ladder exists to prevent. +fn weighted_sum(outputs: &[Vec], weights: &[f32], width: usize) -> Vec { + let mut acc = vec![0.0f32; width]; + for (out, &w) in outputs.iter().zip(weights) { + for (slot, &v) in acc.iter_mut().zip(out) { + *slot += w * v; + } + } + acc +} + +/// Print one ladder rung and say whether it was bit-identical. +fn rung(step: usize, label: &str, a: &[f32], b: &[f32]) -> bool { + let identical = a == b; + let verdict = if identical { + "BIT-IDENTICAL".to_string() + } else if a.len() == b.len() { + format!("max|Δ| = {:.3e}", max_abs_diff(a, b)) + } else { + format!("LENGTH {} vs {}", a.len(), b.len()) + }; + println!(" {step} {label: Result<(), String> { + let vindex = arg("--vindex").ok_or("set --vindex ")?; + let dump = arg("--dump").ok_or("set --dump (LARQL_CPU_DUMP_LAYERS output)")?; + let layer: usize = arg("--layer") + .and_then(|v| v.parse().ok()) + .unwrap_or(DEFAULT_LAYER); + + println!("Gemma one-layer parity — VINDEX2 bytes, two execution paths"); + println!(" vindex {vindex}"); + println!(" layer {layer}"); + + // ── Load the incumbent's weights ─────────────────────────────────────── + let mut callbacks = larql_vindex::SilentLoadCallbacks; + let weights = + larql_vindex::load_model_weights_kquant(std::path::Path::new(&vindex), &mut callbacks) + .map_err(|e| format!("load weights: {e}"))?; + let arch = &*weights.arch; + let hidden = weights.hidden_size; + let norm_offset = arch.norm_weight_offset(); + let eps = arch.norm_eps(); + + let moe: MoeLayerWeights<'_> = build_moe_weights(&weights, arch, layer) + .ok_or_else(|| format!("layer {layer} is not an MoE layer"))?; + println!( + " shape hidden {hidden}, {} experts, top-{}, intermediate {}", + moe.num_experts, moe.top_k, moe.intermediate_size + ); + + // ── The real activation ──────────────────────────────────────────────── + let h = last_row( + &larql_compute::forward::dump_config::cpu_layer_h_post_attn_path(&dump, layer), + hidden, + )?; + println!(" input real h_post_attn, last token of the prompt"); + + // ── Incumbent: the two inputs and the routing decision ───────────────── + let expert_input = moe_expert_input(&h, &moe, norm_offset, eps); + let router_in = moe_router_input(&h, &expert_input, &moe, norm_offset, eps); + let (incumbent_ids, incumbent_weights) = moe_route_from_router_input(&router_in, &moe); + println!( + "\n incumbent routes on a {} vector", + if expert_input == router_in { + "shared" + } else { + "separate" + } + ); + + // ── Bind the selected experts' bytes as a VINDEX3 operation ──────────── + // + // Only the selected experts: a bank holding a subset of the population is + // a legitimate shard, and materialising all of them would cost gigabytes to + // no purpose. If the VINDEX3 router disagrees about the selection it will + // ask for an expert that is not bound and fail loudly, which is the right + // failure. + // + // Two bindings of the *same* experts: + // + // bound the store's own Q4_K super-blocks → the production kernel + // reference the same weights dequantised to f32 → the oracle + let inter = moe.intermediate_size; + let inter_padded = moe.inter_padded(); + let mut buffers = Vec::new(); + let mut q4k_experts: Vec> = Vec::new(); + for &e in &incumbent_ids { + let gate_up_bytes = *moe + .experts_gate_up + .get(e) + .ok_or_else(|| format!("expert {e} has no gate_up bytes"))?; + let down_bytes = *moe + .experts_down + .get(e) + .ok_or_else(|| format!("expert {e} has no down bytes"))?; + + // `down` is stored at the *padded* intermediate width: Q4_K rounds 704 + // up to the next 256-multiple (768), while `gate_up` is unpadded + // because `hidden` is already a multiple. So the two regions disagree + // about the intermediate axis, and the padding columns are inert. + // + // This is precisely what a slice view is for: bind the stored + // [hidden, 768] and let the role see [hidden, 704]. The incumbent + // reaches the same place by zero-padding the activation instead — and + // so does the Q4_K kernel bound here, which is why one binding serves + // both kernels rather than each needing its own. + q4k_experts.push(BoundExpert { + expert_id: e as u32, + projection: BoundProjection::Fused { + gate_up: tensor( + &RegionRole::GateUpFused.name(), + gate_up_bytes, + RegionFormat::Q4K, + ComponentContract::matrix( + (FUSED_PROJECTION_HALVES * inter) as u32, + hidden as u32, + ), + )?, + }, + down: sliced_tensor( + &RegionRole::Down.name(), + down_bytes, + RegionFormat::Q4K, + ComponentContract::matrix(hidden as u32, inter_padded as u32), + inter, + )?, + }); + + let (gate_up, gate_up_contract) = dequantised( + RegionRole::GateUpFused, + gate_up_bytes, + FUSED_PROJECTION_HALVES * inter, + hidden, + )?; + let (down, down_contract) = + dequantised(RegionRole::Down, down_bytes, hidden, inter_padded)?; + buffers.push(ExpertBuffers { + expert_id: e as u32, + gate_up, + down, + gate_up_contract, + down_contract, + }); + } + + let router_bytes: Vec = moe + .router_proj + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect(); + // Gemma's routing policy is PerExpert, so the learned per-expert scale is + // part of the recipe. Omitting it left scores bit-identical and normalised + // weights 7e-4 apart — which the ladder localised to post-processing + // rather than to the scoring kernel it would otherwise have been blamed on. + let per_expert_scale_bytes: Vec = moe + .router_per_expert_scale + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect(); + let f32_experts: Vec> = buffers + .iter() + .map(|b| -> Result, String> { + Ok(BoundExpert { + expert_id: b.expert_id, + projection: BoundProjection::Fused { + gate_up: tensor( + &RegionRole::GateUpFused.name(), + &b.gate_up, + RegionFormat::F32, + b.gate_up_contract.clone(), + )?, + }, + down: sliced_tensor( + &RegionRole::Down.name(), + &b.down, + RegionFormat::F32, + b.down_contract.clone(), + inter, + )?, + }) + }) + .collect::>()?; + + let operation = bind_operation( + layer, + hidden, + &moe, + &router_bytes, + &per_expert_scale_bytes, + q4k_experts, + // Rung 2: bind the production expert kernel over the stored blocks. + ExpertKernel::IncumbentQ4kQ8k, + )?; + operation.validate().map_err(|e| format!("bind: {e}"))?; + let reference = bind_operation( + layer, + hidden, + &moe, + &router_bytes, + &per_expert_scale_bytes, + f32_experts, + ExpertKernel::Reference, + )?; + reference + .validate() + .map_err(|e| format!("bind reference: {e}"))?; + + println!(" bound {}", operation.describe()); + println!(" oracle {}", reference.describe()); + println!( + " shard holds {} of {} experts (full population: {})", + operation.banks[0].population(), + operation.router.population(), + operation.holds_full_population() + ); + println!(" padding down stores {inter_padded} columns, the role means {inter}"); + + // ── Execute both VINDEX3 routes on the identical inputs ──────────────── + let (_, trace) = execute_traced(&operation, MoeInputs::split(&expert_input, &router_in)) + .map_err(|e| format!("execute: {e}"))?; + let (_, reference_trace) = + execute_traced(&reference, MoeInputs::split(&expert_input, &router_in)) + .map_err(|e| format!("execute reference: {e}"))?; + + // ── Compare ──────────────────────────────────────────────────────────── + println!("\n== selection (exact match required) =="); + let vindex3_ids: Vec = trace.selected_ids().iter().map(|i| *i as usize).collect(); + println!(" incumbent {incumbent_ids:?}"); + println!(" vindex3 {vindex3_ids:?}"); + let selection_matches = incumbent_ids == vindex3_ids; + println!( + " {}", + if selection_matches { + "PASS identical experts, identical order" + } else { + "FAIL the two paths route differently" + } + ); + if let Some(margin) = trace.selection_margin { + println!( + " margin {margin:.6}{}", + if margin == 0.0 { + " (an exact tie decided the boundary)" + } else { + "" + } + ); + } + + // ── The failure ladder ───────────────────────────────────────────────── + // + // Rung by rung, so that agreement at the bottom cannot conceal + // compensating differences higher up. Renormalisation in particular can + // make two different score vectors produce identical final weights. + println!( + "\n== router ladder (kernel: {}) ==", + operation.router.kernel.name() + ); + + // 1. raw scores, straight from the incumbent's own functions + let mut incumbent_scores = + moe_score_experts(&router_in, moe.router_proj, moe.num_experts, hidden); + moe_softmax(&mut incumbent_scores); + let scores_identical = incumbent_scores == trace.router_scores; + println!( + " 1 raw scores {}", + if scores_identical { + "BIT-IDENTICAL".to_string() + } else { + format!( + "max|Δ| = {:.3e}", + max_abs_diff(&incumbent_scores, &trace.router_scores) + ) + } + ); + + // 2. ordered (id, score) pairs + let incumbent_pairs: Vec<(usize, f32)> = incumbent_ids + .iter() + .map(|&e| (e, incumbent_scores[e])) + .collect(); + let vindex3_pairs: Vec<(usize, f32)> = trace + .selection + .iter() + .map(|s| (s.expert_id as usize, s.raw_score)) + .collect(); + let pairs_identical = incumbent_pairs == vindex3_pairs; + println!( + " 2 top-k id/score pairs {}", + if pairs_identical { + "BIT-IDENTICAL" + } else { + "DIFFER" + } + ); + + // 3. boundary margin + println!(" 3 boundary margin {:?}", trace.selection_margin); + + // 4. pre-normalisation selected weights — the softmax probabilities + let pre_norm: Vec = incumbent_ids.iter().map(|&e| incumbent_scores[e]).collect(); + let vindex3_pre_norm: Vec = trace.selection.iter().map(|s| s.raw_score).collect(); + let pre_norm_identical = pre_norm == vindex3_pre_norm; + println!( + " 4 pre-norm weights {}", + if pre_norm_identical { + "BIT-IDENTICAL" + } else { + "DIFFER" + } + ); + + // 5. normalised weights, after every policy + let final_identical = incumbent_weights == trace.gate_weights(); + println!( + " 5 normalised weights {}", + if final_identical { + "BIT-IDENTICAL".to_string() + } else { + format!( + "max|Δ| = {:.3e}", + max_abs_diff(&incumbent_weights, &trace.gate_weights()) + ) + } + ); + + let weights_match = report( + "gate weights (tolerance)", + &incumbent_weights, + &trace.gate_weights(), + ); + let router_bit_identical = + scores_identical && pairs_identical && pre_norm_identical && final_identical; + + // ── The expert ladder ────────────────────────────────────────────────── + // + // Same shape as the router's, and for the same reason: a single end-to-end + // number would leave a per-expert kernel fault, a reduction fault and a + // scheduling difference indistinguishable. + println!( + "\n== expert ladder (kernel: {}) ==", + operation.banks[0].kernel.name() + ); + + // 6. per-expert outputs, before any routing weight is applied + let incumbent_outputs = incumbent_expert_outputs(&moe, &expert_input, &incumbent_ids)?; + let mut experts_identical = incumbent_outputs.len() == trace.expert_outputs.len(); + for (e, expected) in incumbent_ids.iter().zip(&incumbent_outputs) { + let Some(found) = trace.expert_output(*e as u32) else { + println!(" 6 expert {e:<15} MISSING from the VINDEX3 trace"); + experts_identical = false; + continue; + }; + experts_identical &= expected.as_slice() == found; + } + println!( + " 6 per-expert outputs {}", + if experts_identical { + format!("BIT-IDENTICAL across {} experts", incumbent_outputs.len()) + } else { + "DIFFER — see the first expert above".to_string() + } + ); + // Guards the guard: the incumbent's short-slab branch zeroes its output and + // returns successfully, so an all-zero agreement would be two failures + // agreeing rather than a parity result. + let nonzero = trace.reduced.iter().any(|v| v.abs() > f32::EPSILON); + println!( + " output magnitude {}", + if nonzero { + "non-zero (the kernel ran)" + } else { + "ALL ZERO — the kernel took a refusal branch" + } + ); + + // 7. the weighted reduction, in selection order + let incumbent_reduced = weighted_sum(&incumbent_outputs, &incumbent_weights, hidden); + let reduction_identical = rung(7, "weighted reduction", &incumbent_reduced, &trace.reduced); + + // 8. the block's contribution, after the policy norm the surrounding block + // owns. Reported, not asserted: `cpu_moe_forward` sums its experts + // through whichever parallel schedule is configured. + let block = moe_post_expert_output(&trace.residual_delta, &moe, norm_offset, eps); + let incumbent_block = cpu_moe_forward(&h, &moe, norm_offset, eps); + let block_identical = rung(8, "block output", &incumbent_block, &block); + + // The independent leg. Not a rung: it is a different kernel by design, so + // a band is the honest reading for it alone — and a *relative* one, since + // what is being compared is a residual-scale quantity rather than a + // probability. + let oracle_match = report_relative( + "oracle (relative band)", + &reference_trace.reduced, + &trace.reduced, + ); + + println!("\n== notes =="); + if router_bit_identical { + println!(" Router rung CLOSED: the bound kernel reproduces production scoring exactly."); + } else { + println!(" Router rung open — see the first ladder step that differs."); + } + if experts_identical && reduction_identical { + println!(" Expert rung CLOSED: the bound kernel reproduces production experts exactly,"); + println!(" and the combine agrees in selection order."); + } else { + println!(" Expert rung open — see the first ladder step that differs."); + } + if !block_identical { + println!(" Rung 8 differing is expected when the incumbent reduces through a tree:"); + println!(" floating-point addition is not associative, and rung 7 already isolated"); + println!(" the combine from the schedule."); + } + + if selection_matches + && weights_match + && experts_identical + && reduction_identical + && oracle_match + { + println!("\nPARITY: selection, router weights, expert outputs and reduction all agree."); + Ok(()) + } else { + Err("parity not established — see above".into()) + } +} diff --git a/crates/larql-vindex/examples/vindex3_gemma_layer_sweep.rs b/crates/larql-vindex/examples/vindex3_gemma_layer_sweep.rs new file mode 100644 index 000000000..716b714dd --- /dev/null +++ b/crates/larql-vindex/examples/vindex3_gemma_layer_sweep.rs @@ -0,0 +1,693 @@ +//! Locked-input MoE sweep — every Gemma MoE layer, independently. +//! +//! The single-layer harness proved layer 5. This proves every layer, and +//! deliberately does *not* let them interact: +//! +//! ```text +//! for each MoE layer: +//! incumbent's captured h_post_attn ← the same input to both paths +//! ↓ ↓ +//! incumbent MoE block VINDEX3-bound MoE block +//! ↓ ↓ +//! classify, per checkpoint +//! ``` +//! +//! # Why locked input +//! +//! Because free propagation cannot localise. If layer 7 diverges and layer 8 +//! then reads a residual that already differs, layer 8's mismatch says nothing +//! about layer 8. Feeding every layer the incumbent's own captured activation +//! makes each layer an independent measurement, so the first failure is the +//! first *defect* rather than the first symptom. Composition is a separate +//! proof, and it comes after this one passes. +//! +//! # Why the full expert population +//! +//! The single-layer harness binds only the selected experts, which is a +//! legitimate shard. For a sweep that would be a mistake: a layer routing to an +//! expert outside the bound slice would report a residency refusal, and the +//! headline table would end up measuring slice coverage rather than execution. +//! So the principal run binds all of them — free, because Q4_K regions are +//! bound from the mapped bytes and nothing is materialised. +//! +//! `--shard N` binds only the first N experts instead. That run exists to prove +//! the *classifier*: selected-but-absent must come out as a residency refusal +//! and must not indict the execution. +//! +//! # Why the oracle is opt-in +//! +//! The dequantised f32 leg costs ~24 MB per expert. At full population that is +//! ~3 GB per layer, so it cannot ride along. It answers a different question +//! anyway — quantisation fidelity, which is a sample — where the sweep answers +//! implementation parity, which needs coverage. `--oracle-layer N` runs it for +//! one layer, over that layer's selected experts only. +//! +//! # Capturing the activations +//! +//! ```text +//! LARQL_CPU_DUMP_LAYERS=/tmp/gemma_dump \ +//! larql run "The capital of France is" -n 1 +//! ``` +//! +//! Usage: +//! ```text +//! cargo run --release -p larql-vindex --example vindex3_gemma_layer_sweep -- \ +//! --vindex --dump /tmp/gemma_dump [--shard 8] [--oracle-layer 5] +//! ``` + +use larql_compute::cpu::ops::moe::{ + cpu_moe_forward, moe_expert_input, moe_post_expert_output, moe_route_from_router_input, + moe_router_input, quantize_x_to_q8k, run_single_expert_q4k_q8k_into, ExpertScratch, +}; +use larql_compute::cpu::ops::q4_common::dequantize_q4_k; +use larql_compute::pipeline_layer::build_moe_weights; +use larql_compute::{Activation, MoeLayerWeights}; + +use larql_vindex::format::capability::binding::{ComponentView, RepresentationIdentity}; +use larql_vindex::format::capability::component::ComponentContract; +use larql_vindex::format::capability::coordinate::BankCoordinate; +use larql_vindex::format::lyrw2::region_format::RegionFormat; +use larql_vindex::format::lyrw2::region_role::RegionRole; +use larql_vindex::runtime::consts::{COL_DIM, FUSED_PROJECTION_HALVES}; +use larql_vindex::runtime::{ + execute_traced, BoundBankOperation, BoundExpert, BoundExpertScaling, BoundMoeOperation, + BoundProjection, BoundReduction, BoundRouter, BoundTensor, CollectedTrace, ExecutionError, + ExpertKernel, MoeInputs, RouterKernel, Verdict, +}; + +/// Band for the oracle leg, relative to `max|reference|`. See the single-layer +/// harness for the derivation; it is the Q8_K activation rounding, twice, +/// carried through the hidden-width contraction. +const ORACLE_RELATIVE_BAND: f32 = 0.05; +/// Reference magnitude below which a ratio stops being reportable and the +/// comparison switches to an absolute one. +const ORACLE_SCALE_FLOOR: f32 = 1e-6; +/// Softmax probabilities, so an absolute band is meaningful for these alone. +const WEIGHT_TOLERANCE: f32 = 1e-3; + +const VARIANT: &str = "vindex2-bytes"; +const ROUTER_REGION_SET: &str = "router"; +const PER_EXPERT_SCALE_REGION_SET: &str = "router_per_expert_scale"; +const BANK_ID: u16 = 0; + +const ARG_VINDEX: &str = "--vindex"; +const ARG_DUMP: &str = "--dump"; +const ARG_SHARD: &str = "--shard"; +const ARG_ORACLE_LAYER: &str = "--oracle-layer"; + +fn arg(name: &str) -> Option { + let args: Vec = std::env::args().collect(); + args.iter() + .position(|a| a == name) + .and_then(|i| args.get(i + 1)) + .cloned() +} + +/// Read an f32 dump and return its final row — the token being decoded. +fn last_row(path: &str, width: usize) -> Result, String> { + let bytes = std::fs::read(path).map_err(|e| format!("{path}: {e}"))?; + let values: Vec = bytes + .chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect(); + if !values.len().is_multiple_of(width) || values.is_empty() { + return Err(format!( + "{path}: {} values is not a whole number of {width}-wide rows", + values.len() + )); + } + let rows = values.len() / width; + Ok(values[(rows - 1) * width..].to_vec()) +} + +fn max_abs_diff(a: &[f32], b: &[f32]) -> f32 { + if a.len() != b.len() { + return f32::INFINITY; + } + a.iter() + .zip(b) + .map(|(x, y)| (x - y).abs()) + .fold(0.0f32, f32::max) +} + +fn magnitude(values: &[f32]) -> f32 { + values.iter().fold(0.0f32, |m, v| m.max(v.abs())) +} + +/// Classify one checkpoint. `asserted` says whether the harness contracts this +/// comparison or merely watches it — see [`Verdict`]. +fn classify(a: &[f32], b: &[f32], tolerance: f32, asserted: bool) -> Verdict { + let diff = max_abs_diff(a, b); + Verdict::from_comparison(a == b, diff <= tolerance, asserted) +} + +fn tensor<'a>( + region_set: &str, + bytes: &'a [u8], + format: RegionFormat, + contract: ComponentContract, +) -> Result, String> { + BoundTensor::direct( + RepresentationIdentity::new(region_set, VARIANT), + bytes, + format, + contract, + ) + .map_err(|e| e.to_string()) +} + +/// Bind a stored matrix whose trailing columns are quantisation padding: the +/// role sees `[rows, keep]`, the bytes stay `[rows, stored_cols]`. +fn sliced_tensor<'a>( + region_set: &str, + bytes: &'a [u8], + format: RegionFormat, + storage: ComponentContract, + keep: usize, +) -> Result, String> { + BoundTensor::new( + RepresentationIdentity::new(region_set, VARIANT), + bytes, + format, + storage, + ComponentView::Slice { + dim: COL_DIM, + start: 0, + len: keep as u32, + }, + ) + .map_err(|e| e.to_string()) +} + +/// Bind one expert's Q4_K regions straight from the mapped bytes. +fn q4k_expert<'a>( + expert_id: usize, + moe: &MoeLayerWeights<'a>, + hidden: usize, +) -> Result, String> { + let gate_up_bytes = *moe + .experts_gate_up + .get(expert_id) + .ok_or_else(|| format!("expert {expert_id} has no gate_up bytes"))?; + let down_bytes = *moe + .experts_down + .get(expert_id) + .ok_or_else(|| format!("expert {expert_id} has no down bytes"))?; + let inter = moe.intermediate_size; + Ok(BoundExpert { + expert_id: expert_id as u32, + projection: BoundProjection::Fused { + gate_up: tensor( + &RegionRole::GateUpFused.name(), + gate_up_bytes, + RegionFormat::Q4K, + ComponentContract::matrix((FUSED_PROJECTION_HALVES * inter) as u32, hidden as u32), + )?, + }, + down: sliced_tensor( + &RegionRole::Down.name(), + down_bytes, + RegionFormat::Q4K, + ComponentContract::matrix(hidden as u32, moe.inter_padded() as u32), + inter, + )?, + }) +} + +/// Assemble the bound operation for one layer. +fn bind_operation<'a>( + layer: usize, + hidden: usize, + moe: &MoeLayerWeights<'_>, + router_bytes: &'a [u8], + scale_bytes: &'a [u8], + experts: Vec>, + kernel: ExpertKernel, +) -> Result, String> { + Ok(BoundMoeOperation { + router: BoundRouter { + weight: tensor( + ROUTER_REGION_SET, + router_bytes, + RegionFormat::F32, + ComponentContract::matrix(moe.num_experts as u32, hidden as u32), + )?, + top_k: moe.top_k, + selected_weight: moe.routing_policy.selected_weight, + scaling: if scale_bytes.is_empty() { + BoundExpertScaling::None + } else { + BoundExpertScaling::PerExpert { + scales: tensor( + PER_EXPERT_SCALE_REGION_SET, + scale_bytes, + RegionFormat::F32, + ComponentContract::vector(moe.num_experts as u32), + )?, + } + }, + kernel: RouterKernel::Incumbent, + }, + transforms: Vec::new(), + banks: vec![BoundBankOperation { + bank: BankCoordinate::new(layer as u32, BANK_ID), + experts, + intermediate_dim: moe.intermediate_size, + hidden_dim: hidden, + activation: moe.activation, + kernel, + }], + reduction: BoundReduction::WeightedSum, + residual_dim: hidden, + }) +} + +fn f32_bytes(values: &[f32]) -> Vec { + values.iter().flat_map(|v| v.to_le_bytes()).collect() +} + +/// The incumbent's own expert kernel over each selected expert, in selection +/// order, sharing one Q8_K activation exactly as production does. +fn incumbent_expert_outputs( + moe: &MoeLayerWeights<'_>, + expert_input: &[f32], + selected: &[usize], +) -> Result>, String> { + let hidden = expert_input.len(); + let inter = moe.intermediate_size; + if !hidden.is_multiple_of(256) { + return Err(format!("hidden {hidden} is not a whole Q8_K super-block")); + } + let q8k = quantize_x_to_q8k(expert_input); + let mut scratch = ExpertScratch::new(hidden, inter, moe.inter_padded()); + selected + .iter() + .map(|&e| { + let gate_up = *moe + .experts_gate_up + .get(e) + .ok_or_else(|| format!("expert {e} has no gate_up bytes"))?; + let down = *moe + .experts_down + .get(e) + .ok_or_else(|| format!("expert {e} has no down bytes"))?; + Ok(run_single_expert_q4k_q8k_into( + &mut scratch, + &q8k, + gate_up, + down, + inter, + moe.activation, + ) + .to_vec()) + }) + .collect() +} + +/// `sum_i weight_i * output_i`, in selection order — the combine isolated from +/// whatever schedule the incumbent's parallel path happened to use. +fn weighted_sum(outputs: &[Vec], weights: &[f32], width: usize) -> Vec { + let mut acc = vec![0.0f32; width]; + for (out, &w) in outputs.iter().zip(weights) { + for (slot, &v) in acc.iter_mut().zip(out) { + *slot += w * v; + } + } + acc +} + +/// One layer's classification, per checkpoint. +struct LayerReport { + layer: usize, + population: usize, + resident: usize, + top_k: usize, + intermediate: usize, + inter_padded: usize, + activation: Activation, + selection: Verdict, + weights: Verdict, + experts: Verdict, + reduction: Verdict, + block: Verdict, + refusal: Option, +} + +impl LayerReport { + /// The layer's outcome: the weakest claim any checkpoint reached. + /// + /// `Verdict` orders weakest-last precisely so this is a maximum rather than + /// a hand-written precedence table that could disagree with itself. + fn verdict(&self) -> Verdict { + if let Some(err) = &self.refusal { + return Verdict::from(err); + } + [ + self.selection, + self.weights, + self.experts, + self.reduction, + self.block, + ] + .into_iter() + .max() + .unwrap_or(Verdict::Exact) + } + + fn print(&self) { + let padding = if self.inter_padded == self.intermediate { + "none".to_string() + } else { + format!("{}→{}", self.intermediate, self.inter_padded) + }; + println!( + " {:>3} {:>4}/{:<4} top-{:<2} {:<9} pad {:<9} {:<16} {}", + self.layer, + self.resident, + self.population, + self.top_k, + format!("{:?}", self.activation), + padding, + self.verdict().name(), + self.detail() + ); + } + + /// Checkpoints that came out weaker than their own *strongest attainable* + /// outcome. + /// + /// Not "weaker than `Exact`". The block comparison is unasserted by design, + /// so its ceiling is `ObservedExact`, and flagging it against `Exact` would + /// print the same note for a perfect run as for a broken one — a column + /// that says the same thing in both cases is worse than no column. + fn detail(&self) -> String { + if let Some(err) = &self.refusal { + return err.to_string(); + } + let below: Vec<&str> = [ + ("selection", self.selection, Verdict::Exact), + ("weights", self.weights, Verdict::Exact), + ("experts", self.experts, Verdict::Exact), + ("reduction", self.reduction, Verdict::Exact), + ("block", self.block, Verdict::ObservedExact), + ] + .into_iter() + // `Verdict` orders weakest-last, so "greater than" is "weaker than". + .filter(|(_, got, ceiling)| got > ceiling) + .map(|(name, _, _)| name) + .collect(); + if below.is_empty() { + String::new() + } else { + format!("below ceiling at: {}", below.join(", ")) + } + } +} + +/// Compare one layer's two paths and classify every checkpoint. +fn sweep_layer( + layer: usize, + hidden: usize, + moe: &MoeLayerWeights<'_>, + h: &[f32], + norm_offset: f32, + eps: f32, + shard: Option, +) -> Result { + let expert_input = moe_expert_input(h, moe, norm_offset, eps); + let router_in = moe_router_input(h, &expert_input, moe, norm_offset, eps); + let (incumbent_ids, incumbent_weights) = moe_route_from_router_input(&router_in, moe); + + let resident = shard.unwrap_or(moe.num_experts).min(moe.num_experts); + let experts: Vec> = (0..resident) + .map(|e| q4k_expert(e, moe, hidden)) + .collect::>()?; + + let router_bytes = f32_bytes(moe.router_proj); + let scale_bytes = f32_bytes(moe.router_per_expert_scale); + let operation = bind_operation( + layer, + hidden, + moe, + &router_bytes, + &scale_bytes, + experts, + ExpertKernel::IncumbentQ4kQ8k, + )?; + + let mut report = LayerReport { + layer, + population: moe.num_experts, + resident, + top_k: moe.top_k, + intermediate: moe.intermediate_size, + inter_padded: moe.inter_padded(), + activation: moe.activation, + selection: Verdict::Exact, + weights: Verdict::Exact, + experts: Verdict::Exact, + reduction: Verdict::Exact, + block: Verdict::Exact, + refusal: None, + }; + + // Binding faults and missing kernels surface here, before any comparison — + // and they are classified, not counted as mismatches. + if let Err(err) = operation.validate() { + report.refusal = Some(err); + return Ok(report); + } + + let trace: CollectedTrace = + match execute_traced(&operation, MoeInputs::split(&expert_input, &router_in)) { + Ok((_, trace)) => trace, + Err(err) => { + report.refusal = Some(err); + return Ok(report); + } + }; + + // Selection is discrete: identical or not, and asserted either way. + let bound_ids: Vec = trace.selected_ids().iter().map(|i| *i as usize).collect(); + report.selection = Verdict::from_comparison(bound_ids == incumbent_ids, false, true); + report.weights = classify( + &incumbent_weights, + &trace.gate_weights(), + WEIGHT_TOLERANCE, + true, + ); + + let incumbent_outputs = incumbent_expert_outputs(moe, &expert_input, &incumbent_ids)?; + let mut experts_verdict = Verdict::Exact; + for (e, expected) in incumbent_ids.iter().zip(&incumbent_outputs) { + let found = trace.expert_output(*e as u32).unwrap_or(&[]); + experts_verdict = experts_verdict.max(classify(expected, found, 0.0, true)); + } + report.experts = experts_verdict; + + let incumbent_reduced = weighted_sum(&incumbent_outputs, &incumbent_weights, hidden); + report.reduction = classify(&incumbent_reduced, &trace.reduced, 0.0, true); + + // Observed, not asserted: `cpu_moe_forward` sums through whichever parallel + // schedule is configured, and a tree reduction is not required to agree + // bit-for-bit with a sequential one. + let bound_block = moe_post_expert_output(&trace.residual_delta, moe, norm_offset, eps); + let incumbent_block = cpu_moe_forward(h, moe, norm_offset, eps); + report.block = classify(&incumbent_block, &bound_block, 0.0, false); + + Ok(report) +} + +/// The fidelity leg, for one layer, over its selected experts only. +fn oracle_layer( + layer: usize, + hidden: usize, + moe: &MoeLayerWeights<'_>, + h: &[f32], + norm_offset: f32, + eps: f32, +) -> Result<(), String> { + let expert_input = moe_expert_input(h, moe, norm_offset, eps); + let router_in = moe_router_input(h, &expert_input, moe, norm_offset, eps); + let (selected, _) = moe_route_from_router_input(&router_in, moe); + let inter = moe.intermediate_size; + let inter_padded = moe.inter_padded(); + + // Owned buffers, so the bound tensors have something to borrow. + let mut buffers: Vec<(u32, Vec, Vec)> = Vec::new(); + for &e in &selected { + let gate_up = dequantize_q4_k( + moe.experts_gate_up[e], + FUSED_PROJECTION_HALVES * inter * hidden, + ); + let down = dequantize_q4_k(moe.experts_down[e], hidden * inter_padded); + buffers.push((e as u32, f32_bytes(&gate_up), f32_bytes(&down))); + } + let experts: Vec> = buffers + .iter() + .map(|(id, gate_up, down)| -> Result, String> { + Ok(BoundExpert { + expert_id: *id, + projection: BoundProjection::Fused { + gate_up: tensor( + &RegionRole::GateUpFused.name(), + gate_up, + RegionFormat::F32, + ComponentContract::matrix( + (FUSED_PROJECTION_HALVES * inter) as u32, + hidden as u32, + ), + )?, + }, + down: sliced_tensor( + &RegionRole::Down.name(), + down, + RegionFormat::F32, + ComponentContract::matrix(hidden as u32, inter_padded as u32), + inter, + )?, + }) + }) + .collect::>()?; + + let router_bytes = f32_bytes(moe.router_proj); + let scale_bytes = f32_bytes(moe.router_per_expert_scale); + let reference = bind_operation( + layer, + hidden, + moe, + &router_bytes, + &scale_bytes, + experts, + ExpertKernel::Reference, + )?; + reference + .validate() + .map_err(|e| format!("bind reference: {e}"))?; + let (_, reference_trace) = + execute_traced(&reference, MoeInputs::split(&expert_input, &router_in)) + .map_err(|e| format!("execute reference: {e}"))?; + + let incumbent_outputs = incumbent_expert_outputs(moe, &expert_input, &selected)?; + let (_, incumbent_weights) = moe_route_from_router_input(&router_in, moe); + let production = weighted_sum(&incumbent_outputs, &incumbent_weights, hidden); + + let diff = max_abs_diff(&reference_trace.reduced, &production); + let scale = magnitude(&reference_trace.reduced); + if scale < ORACLE_SCALE_FLOOR { + println!( + "\n== oracle, layer {layer} ==\n max|Δ| = {diff:.3e}, reference magnitude \ + {scale:.3e} below the {ORACLE_SCALE_FLOOR:.0e} floor — absolute, not a ratio" + ); + return Ok(()); + } + let relative = diff / scale; + println!( + "\n== oracle, layer {layer} ==\n quantisation fidelity, not parity: \ + max|Δ| = {diff:.3e} of {scale:.3e} = {:.2}% {}", + relative * 100.0, + if relative <= ORACLE_RELATIVE_BAND { + "within band" + } else { + "OUTSIDE" + } + ); + Ok(()) +} + +fn main() -> Result<(), String> { + let vindex = arg(ARG_VINDEX).ok_or(format!("set {ARG_VINDEX} "))?; + let dump = arg(ARG_DUMP).ok_or(format!("set {ARG_DUMP} "))?; + let shard: Option = arg(ARG_SHARD).and_then(|v| v.parse().ok()); + let oracle: Option = arg(ARG_ORACLE_LAYER).and_then(|v| v.parse().ok()); + + println!("Locked-input MoE sweep — every layer fed the incumbent's own activation"); + println!(" vindex {vindex}"); + match shard { + Some(n) => println!(" shard first {n} experts only — a residency-classification run"), + None => println!(" shard full population — the parity run"), + } + + let mut callbacks = larql_vindex::SilentLoadCallbacks; + let weights = + larql_vindex::load_model_weights_kquant(std::path::Path::new(&vindex), &mut callbacks) + .map_err(|e| format!("load weights: {e}"))?; + let arch = &*weights.arch; + let hidden = weights.hidden_size; + let norm_offset = arch.norm_weight_offset(); + let eps = arch.norm_eps(); + let num_layers = weights.num_layers; + + println!(" layers {num_layers}, hidden {hidden}\n"); + println!(" lyr resident top-k activation pad verdict"); + + let mut reports = Vec::new(); + let mut skipped = Vec::new(); + for layer in 0..num_layers { + let Some(moe) = build_moe_weights(&weights, arch, layer) else { + skipped.push((layer, "not an MoE layer".to_string())); + continue; + }; + let path = larql_compute::forward::dump_config::cpu_layer_h_post_attn_path(&dump, layer); + let h = match last_row(&path, hidden) { + Ok(h) => h, + Err(e) => { + skipped.push((layer, e)); + continue; + } + }; + let report = sweep_layer(layer, hidden, &moe, &h, norm_offset, eps, shard)?; + report.print(); + reports.push(report); + } + + // ── Summary ──────────────────────────────────────────────────────────── + println!("\n== classification =="); + let mut counts: Vec<(Verdict, usize)> = Vec::new(); + for report in &reports { + let verdict = report.verdict(); + match counts.iter_mut().find(|(v, _)| *v == verdict) { + Some((_, n)) => *n += 1, + None => counts.push((verdict, 1)), + } + } + counts.sort_unstable_by_key(|(v, _)| *v); + for (verdict, n) in &counts { + println!(" {:<16} {n} layer(s)", verdict.name()); + } + if !skipped.is_empty() { + println!("\n== skipped =="); + for (layer, why) in &skipped { + println!(" layer {layer}: {why}"); + } + } + + if let Some(layer) = oracle { + let moe = build_moe_weights(&weights, arch, layer) + .ok_or_else(|| format!("layer {layer} is not an MoE layer"))?; + let path = larql_compute::forward::dump_config::cpu_layer_h_post_attn_path(&dump, layer); + let h = last_row(&path, hidden)?; + oracle_layer(layer, hidden, &moe, &h, norm_offset, eps)?; + } + + println!("\n== notes =="); + println!(" exact every asserted checkpoint bit-identical — a contract"); + println!(" observed_exact bit-identical but schedule-dependent — a measurement"); + println!(" equivalent within a named tolerance"); + println!(" residency selected expert not in the bound bank — NOT a parity failure"); + println!(" unsupported no bound kernel serves this operand — a gap, not a defect"); + println!(" binding_defect the binding is wrong"); + println!(" numeric_mismatch executed and disagreed"); + + let indicted: Vec = reports + .iter() + .filter(|r| r.verdict().indicts_execution()) + .map(|r| r.layer) + .collect(); + if indicted.is_empty() { + println!("\nSWEEP: no layer indicts the execution."); + Ok(()) + } else { + Err(format!("layers {indicted:?} indict the execution")) + } +} diff --git a/crates/larql-vindex/examples/vindex3_make_fixture_a.rs b/crates/larql-vindex/examples/vindex3_make_fixture_a.rs new file mode 100644 index 000000000..5009360e1 --- /dev/null +++ b/crates/larql-vindex/examples/vindex3_make_fixture_a.rs @@ -0,0 +1,25 @@ +//! Emit conformance fixture A as a real VINDEX3 container on disk. +//! +//! The smallest way to get a genuine `index.json.version: 3` directory to +//! point tooling at — `larql show`, `larql verify`, or anything else that +//! needs to prove it handles both generations rather than asserting it. +//! +//! ```text +//! cargo run --release -p larql-vindex --example vindex3_make_fixture_a -- /tmp/fixture-a.vindex +//! larql show /tmp/fixture-a.vindex +//! larql verify /tmp/fixture-a.vindex +//! ``` + +use larql_vindex::format::vindex3::{test_support::fixture_a_spec, write_container}; + +fn main() -> Result<(), String> { + let out = std::env::args() + .nth(1) + .ok_or("usage: vindex3_make_fixture_a ")?; + let dir = std::path::PathBuf::from(&out); + write_container(&dir, &fixture_a_spec()).map_err(|e| format!("write container: {e}"))?; + println!("wrote VINDEX3 container: {}", dir.display()); + println!(" larql show {out}"); + println!(" larql verify {out}"); + Ok(()) +} diff --git a/crates/larql-vindex/examples/vindex3_reference_moe.rs b/crates/larql-vindex/examples/vindex3_reference_moe.rs new file mode 100644 index 000000000..4c9c0ca99 --- /dev/null +++ b/crates/larql-vindex/examples/vindex3_reference_moe.rs @@ -0,0 +1,103 @@ +//! VINDEX3 reference MoE demo — fixture A, checkpoint by checkpoint. +//! +//! The counterpart to `q4k_demo` / `walker_demo` for the new execution path. +//! It prints every instrumented boundary inside one routed-MoE layer, run +//! twice: once from decomposed `gate`/`up` regions and once from a single +//! concatenated `gate_up_fused` region. +//! +//! Both runs must agree exactly. That is the property worth being able to see +//! by eye: the executor is following a bound recipe, not inferring a layout. +//! +//! Usage: `cargo run --release -p larql-vindex --example vindex3_reference_moe` + +use larql_vindex::runtime::execute_traced; +use larql_vindex::runtime::fixtures::direct_moe::{input, DirectMoeFixture, POPULATION, TOP_K}; +use larql_vindex::runtime::fixtures::direct_moe_oracle::oracle; +use larql_vindex::runtime::MoeInputs; +use larql_vindex::runtime::{CollectedTrace, ProjectionArrangement}; + +/// Column width for the numeric dumps. +const PRECISION: usize = 6; + +fn render(values: &[f32]) -> String { + values + .iter() + .map(|v| format!("{v:+.PRECISION$}")) + .collect::>() + .join(" ") +} + +fn report(arrangement: ProjectionArrangement, trace: &CollectedTrace, delta: &[f32]) { + println!("\n── {} storage {}", arrangement.name(), "─".repeat(46)); + println!(" router scores ({POPULATION} experts)"); + println!(" {}", render(&trace.router_scores)); + println!(" selection (top-{TOP_K}, in order)"); + for selected in &trace.selection { + println!( + " expert {:>3} weight {:+.PRECISION$} raw {:+.PRECISION$}", + selected.expert_id, selected.weight, selected.raw_score + ); + } + println!(" expert outputs (unweighted)"); + for out in &trace.expert_outputs { + println!( + " expert {:>3} {}", + out.expert_id, + render(&out.values) + ); + } + println!(" reduction"); + println!(" {}", render(&trace.reduced)); + println!(" residual delta"); + println!(" {}", render(delta)); +} + +fn main() { + let fixture = DirectMoeFixture::new(); + let residual = input(); + + println!("VINDEX3 reference MoE — fixture A"); + println!(" residual in"); + println!(" {}", render(&residual)); + + let mut runs = Vec::new(); + for arrangement in ProjectionArrangement::ALL { + let operation = fixture.operation(arrangement); + operation + .validate() + .expect("fixture A binds a consistent operation"); + println!("\n bound: {}", operation.describe()); + + let (delta, trace) = execute_traced(&operation, MoeInputs::shared(&residual)) + .expect("fixture A executes cleanly"); + report(arrangement, &trace, &delta); + runs.push((arrangement, delta, trace)); + } + + // ── The two properties worth asserting out loud ──────────────────────── + println!("\n── verdict {}", "─".repeat(52)); + + let [(_, first_delta, first_trace), (_, second_delta, second_trace)] = &runs[..] else { + unreachable!("exactly one run per arrangement"); + }; + let arrangements_agree = first_delta == second_delta && first_trace == second_trace; + println!( + " {} fused and decomposed agree at every checkpoint", + if arrangements_agree { "PASS" } else { "FAIL" } + ); + + let expected = oracle(&residual); + let epsilon = 1e-6; + let matches_oracle = first_delta + .iter() + .zip(&expected.reduced) + .all(|(a, e)| (a - e).abs() < epsilon); + println!( + " {} output matches the independent oracle (< {epsilon:e})", + if matches_oracle { "PASS" } else { "FAIL" } + ); + + if !(arrangements_agree && matches_oracle) { + std::process::exit(1); + } +} diff --git a/crates/larql-vindex/examples/vindex3_residency_probe.rs b/crates/larql-vindex/examples/vindex3_residency_probe.rs new file mode 100644 index 000000000..cabcd0ac2 --- /dev/null +++ b/crates/larql-vindex/examples/vindex3_residency_probe.rs @@ -0,0 +1,401 @@ +//! VINDEX3 residency probe — the successor to `mmap_cold_read_probe`. +//! +//! The incumbent probe asks whether sparse access to a large blob pages +//! acceptably. VINDEX3 supports a sharper question, because binding precedes +//! execution: +//! +//! > **Does the bound plan predict the pages execution actually touches?** +//! +//! If it does, residency becomes a property you can compute rather than +//! observe — which is what placement, prefetch and remote transfer all need, +//! and all three would be mis-sized by a predictor that quietly disagreed with +//! reality. +//! +//! # Method +//! +//! `mincore(2)`, not fault counting. The incumbent probe documents that +//! Darwin's `MADV_DONTNEED` is lazy and re-eviction unreliable, which makes +//! `getrusage` fault deltas awkward there. `mincore` reports which pages are +//! resident *right now* — a direct read of the thing being claimed, and it +//! answers the sparsity question without needing eviction to be reliable. +//! +//! ```text +//! 1. write a synthetic layer's byte image to disk +//! 2. mmap it, MADV_RANDOM (readahead off — otherwise it prefetches the +//! very sparsity being measured) +//! 3. evict, verify cold (spine check: near-zero residency) +//! 4. execute exactly one token through the bound operation +//! 5. mincore, attribute pages to router / selected / unselected +//! ``` +//! +//! Step 3 is a spine check, not a result. If the mapping is not cold +//! afterwards the numbers mean nothing, and the probe says so rather than +//! reporting a confident zero. +//! +//! Usage: +//! ```text +//! cargo run --release -p larql-vindex --example vindex3_residency_probe -- \ +//! [--population 128] [--hidden 256] [--intermediate 512] [--top-k 2] +//! ``` + +#[cfg(unix)] +mod probe { + use std::fs::File; + use std::io::Write; + + use memmap2::{Mmap, MmapOptions}; + + use larql_vindex::runtime::execute_traced; + use larql_vindex::runtime::fixtures::synthetic::{SyntheticLayer, SyntheticShape}; + use larql_vindex::runtime::residency::{account, ExpertRegion}; + use larql_vindex::runtime::MoeInputs; + + /// Defaults sized to be meaningful (hundreds of MB) without needing a + /// special machine. Every one is overridable. + const DEFAULT_POPULATION: usize = 128; + const DEFAULT_HIDDEN: usize = 256; + const DEFAULT_INTERMEDIATE: usize = 512; + const DEFAULT_TOP_K: usize = 2; + const BLOB_NAME: &str = "vindex3_residency_probe.bin"; + /// Above this fraction resident after eviction, the mapping is not cold + /// and no downstream number is trustworthy. + const COLD_THRESHOLD: f64 = 0.05; + const BYTES_PER_MIB: f64 = 1_048_576.0; + + fn arg(name: &str, default: usize) -> usize { + let args: Vec = std::env::args().collect(); + args.iter() + .position(|a| a == name) + .and_then(|i| args.get(i + 1)) + .and_then(|v| v.parse().ok()) + .unwrap_or(default) + } + + fn page_size() -> usize { + // SAFETY: `sysconf` is a pure query with no preconditions. + (unsafe { libc::sysconf(libc::_SC_PAGESIZE) }) as usize + } + + /// Conditions the numbers were produced under. + /// + /// Recorded rather than assumed: Darwin's page-cache behaviour varies with + /// OS release and backing filesystem, so a residency figure without its + /// conditions is not reproducible. `mmap_cold_read_probe` reached a + /// different conclusion about `MADV_DONTNEED` on the same platform family, + /// which is exactly the sort of disagreement this makes resolvable. + pub struct Environment { + pub os: String, + pub filesystem: String, + pub page_size: usize, + pub mapped_bytes: usize, + pub eviction_method: &'static str, + pub commit: String, + } + + /// How eviction is performed, named so a result can be attributed to it. + pub const EVICTION_METHOD: &str = "msync(MS_INVALIDATE) + madvise(MADV_DONTNEED)"; + + fn command(program: &str, args: &[&str]) -> String { + std::process::Command::new(program) + .args(args) + .output() + .ok() + .filter(|o| o.status.success()) + .map(|o| String::from_utf8_lossy(&o.stdout).trim().to_string()) + .filter(|s| !s.is_empty()) + .unwrap_or_else(|| "unknown".into()) + } + + /// Filesystem type backing `path`, via `statfs`. + /// + /// Not `stat(1)`: its `%T` is the *file* type on Darwin, which reported + /// "unknown" here. `statfs` carries the filesystem name directly on + /// Darwin and a numeric magic on Linux, so the two are read separately. + fn filesystem(path: &std::path::Path) -> String { + use std::ffi::CString; + let Ok(c_path) = CString::new(path.as_os_str().as_encoded_bytes()) else { + return "unknown".into(); + }; + // SAFETY: zeroed `statfs` is a valid initial state, and `c_path` is a + // NUL-terminated path that outlives the call. + let mut buf: libc::statfs = unsafe { std::mem::zeroed() }; + if unsafe { libc::statfs(c_path.as_ptr(), &mut buf) } != 0 { + return "unknown".into(); + } + #[cfg(target_os = "macos")] + { + // SAFETY: `f_fstypename` is a NUL-terminated array of c_char. + let name = unsafe { std::ffi::CStr::from_ptr(buf.f_fstypename.as_ptr()) }; + name.to_string_lossy().into_owned() + } + #[cfg(target_os = "linux")] + { + // Linux reports a magic number; render it rather than pretending + // to a name this code would have to keep a table for. + format!("magic 0x{:x}", buf.f_type) + } + // Other unixes spell the field differently again (BSD uses + // `f_fstypename`, Solaris `f_basetype`). Not guessed — this is a + // diagnostic, and a wrong field name is a build break on a platform + // nobody here can test. + #[cfg(not(any(target_os = "macos", target_os = "linux")))] + { + let _ = buf; + "unknown".to_string() + } + } + + fn environment(path: &std::path::Path, mapped_bytes: usize) -> Environment { + Environment { + os: format!( + "{} {}", + command("uname", &["-sr"]), + command("uname", &["-m"]) + ), + filesystem: filesystem(path), + page_size: page_size(), + mapped_bytes, + eviction_method: EVICTION_METHOD, + commit: command("git", &["rev-parse", "--short", "HEAD"]), + } + } + + /// Ask the kernel to drop the mapping's resident pages. + /// + /// `msync(MS_INVALIDATE)` first, and that ordering is load-bearing on + /// Darwin: `MADV_DONTNEED` alone leaves a freshly-written file 100% + /// resident there, because the pages are clean in the unified buffer + /// cache and `madvise` declines to drop them. Invalidating first releases + /// them, and residency goes to zero. + /// + /// Measured on macOS 24.6 with a 64 MiB blob: + /// + /// ```text + /// MADV_DONTNEED 4096/4096 resident (100.0%) + /// MS_INVALIDATE + MADV_DONTNEED 0/4096 resident ( 0.0%) + /// ``` + /// + /// `mmap_cold_read_probe` works around the same limitation differently, + /// by falling back to `F_NOCACHE` `pread` for its sparse pass. + fn evict(mmap: &Mmap) { + // SAFETY: the pointer and length come from a live mapping we own. + unsafe { + libc::msync( + mmap.as_ptr() as *mut libc::c_void, + mmap.len(), + libc::MS_INVALIDATE, + ); + libc::madvise( + mmap.as_ptr() as *mut libc::c_void, + mmap.len(), + libc::MADV_DONTNEED, + ); + } + } + + /// Turn off readahead. Without this the kernel prefetches neighbouring + /// pages and the sparsity under measurement is manufactured away. + fn random_access(mmap: &Mmap) { + // SAFETY: as above. + unsafe { + libc::madvise( + mmap.as_ptr() as *mut libc::c_void, + mmap.len(), + libc::MADV_RANDOM, + ); + } + } + + /// Which pages of the mapping are resident. + fn resident_pages(mmap: &Mmap, page: usize) -> Vec { + let pages = mmap.len().div_ceil(page); + let mut vec = vec![0u8; pages]; + // SAFETY: `vec` has one byte per page of the live mapping. The third + // argument's type differs across platforms, hence the cast. + let rc = unsafe { + libc::mincore( + mmap.as_ptr() as *mut libc::c_void, + mmap.len(), + vec.as_mut_ptr() as *mut _, + ) + }; + if rc != 0 { + eprintln!( + " mincore failed (errno {})", + std::io::Error::last_os_error() + ); + return vec![false; pages]; + } + // Bit 0 is the resident flag; higher bits carry other information on + // some platforms and must be masked off. + vec.into_iter().map(|b| b & 1 == 1).collect() + } + + pub fn run() -> std::process::ExitCode { + let shape = SyntheticShape { + population: arg("--population", DEFAULT_POPULATION), + hidden: arg("--hidden", DEFAULT_HIDDEN), + intermediate: arg("--intermediate", DEFAULT_INTERMEDIATE), + top_k: arg("--top-k", DEFAULT_TOP_K), + }; + let page = page_size(); + + println!("VINDEX3 residency probe"); + println!( + " shape population {}, hidden {}, intermediate {}, top-k {}", + shape.population, shape.hidden, shape.intermediate, shape.top_k + ); + println!( + " image {:.1} MiB total, {:.1} MiB per expert, page {page} B", + shape.total_bytes() as f64 / BYTES_PER_MIB, + shape.expert_bytes() as f64 / BYTES_PER_MIB + ); + + // ── Build and write the byte image ───────────────────────────────── + let layer = SyntheticLayer::build(shape); + let path = std::env::temp_dir().join(BLOB_NAME); + { + let mut file = File::create(&path).expect("create blob"); + file.write_all(&layer.bytes).expect("write blob"); + file.sync_all().expect("sync blob"); + } + println!(" blob {}", path.display()); + + let file = File::open(&path).expect("open blob"); + // SAFETY: the file is ours, freshly written, and not mutated while + // mapped. + let mmap = unsafe { MmapOptions::new().map(&file) }.expect("map blob"); + random_access(&mmap); + + let env = environment(&path, mmap.len()); + println!("\n== conditions =="); + println!(" os {}", env.os); + println!(" filesystem {}", env.filesystem); + println!(" page size {} B", env.page_size); + println!(" mapped {} B", env.mapped_bytes); + println!(" eviction {}", env.eviction_method); + println!(" commit {}", env.commit); + + // ── Spine check: is it actually cold? ────────────────────────────── + evict(&mmap); + let before = resident_pages(&mmap, page); + let total_pages = before.len(); + let cold_resident = before.iter().filter(|r| **r).count(); + let cold_fraction = cold_resident as f64 / total_pages.max(1) as f64; + println!( + "\n after evict {cold_resident}/{total_pages} pages resident ({:.1}%)", + cold_fraction * 100.0 + ); + if cold_fraction > COLD_THRESHOLD { + println!( + " SPINE FAIL MADV_DONTNEED did not evict on this OS — \ + residency numbers below would be meaningless." + ); + return std::process::ExitCode::FAILURE; + } + + // ── One token through the bound operation ────────────────────────── + let operation = layer.bind(&mmap); + operation.validate().expect("synthetic layer binds cleanly"); + let (_, trace) = execute_traced(&operation, MoeInputs::shared(&layer.residual())) + .expect("layer executes"); + let selected = trace.selected_ids(); + + let after = resident_pages(&mmap, page); + let regions: Vec = layer + .experts + .iter() + .enumerate() + .map(|(e, extents)| ExpertRegion { + expert_id: e as u32, + bytes: extents.gate_up.start..extents.down.end, + }) + .collect(); + let acct = account( + &after, + page, + &layer.router, + ®ions, + &selected, + shape.selected_bytes(), + ); + + // ── Report ───────────────────────────────────────────────────────── + println!(" selected experts {selected:?}"); + if let Some(margin) = trace.selection_margin { + println!( + " margin {margin:.6}{}", + if margin == 0.0 { + " (decided by a tie)" + } else { + "" + } + ); + } + println!("\n== residency after one token =="); + println!( + " predicted {:>7} pages (router + {} experts)", + acct.predicted, + selected.len() + ); + println!( + " resident {:>7} pages ({:.2}% of the mapping)", + acct.resident_total, + acct.resident_fraction(total_pages) * 100.0 + ); + println!(" router {:>7}", acct.resident_router); + println!(" selected {:>7}", acct.resident_selected); + println!( + " unselected {:>7} <- pages no selected operand asked for", + acct.resident_unselected + ); + println!( + " coverage {:>7.3} (resident-and-wanted / predicted)", + acct.prediction_coverage() + ); + println!( + " overshoot {:>7.3} (unselected / resident)", + acct.overshoot_fraction() + ); + + // The claim, stated as a verdict rather than left to the reader. + let sparse = acct.resident_fraction(total_pages) < 0.5; + println!("\n== verdict =="); + println!( + " {} routing stayed sparse: one token made {:.2}% of the layer resident", + if sparse { "PASS" } else { "FAIL" }, + acct.resident_fraction(total_pages) * 100.0 + ); + println!( + " {} the bound plan predicted its own footprint within a page-boundary margin\ + \n\n scope: reference kernel — its touch envelope equals its required pages.\ + \n A grouped kernel may exceed this legitimately (whole group extents,\ + \n aligned over-read, separate scale blocks). Compare such a kernel against\ + \n its declared envelope, not against zero overshoot.", + if acct.prediction_coverage() <= 1.0 { + "PASS" + } else { + "FAIL" + } + ); + + let _ = std::fs::remove_file(&path); + if sparse { + std::process::ExitCode::SUCCESS + } else { + std::process::ExitCode::FAILURE + } + } +} + +#[cfg(unix)] +fn main() -> std::process::ExitCode { + probe::run() +} + +#[cfg(not(unix))] +fn main() -> std::process::ExitCode { + // mincore / madvise are POSIX. Windows would need QueryWorkingSetEx. + eprintln!("vindex3_residency_probe is unix-only (mmap / madvise / mincore)."); + std::process::ExitCode::SUCCESS +} diff --git a/crates/larql-vindex/src/error.rs b/crates/larql-vindex/src/error.rs index ec8a9bd51..51a6bb19a 100644 --- a/crates/larql-vindex/src/error.rs +++ b/crates/larql-vindex/src/error.rs @@ -25,6 +25,18 @@ pub enum VindexError { needed: ExtractLevel, have: ExtractLevel, }, + #[error( + "index.json declares version {found}, which this binary does not support; \ + supported: {supported}" + )] + UnknownContainerGeneration { found: u32, supported: String }, + + #[error("this is a {found} index; that path requires the {required} loader")] + WrongContainerGeneration { + found: &'static str, + required: &'static str, + }, + #[error("IO error: {0}")] Io(#[from] std::io::Error), #[error("model error: {0}")] diff --git a/crates/larql-vindex/src/format/capability/authority.rs b/crates/larql-vindex/src/format/capability/authority.rs new file mode 100644 index 000000000..32709dc99 --- /dev/null +++ b/crates/larql-vindex/src/format/capability/authority.rs @@ -0,0 +1,426 @@ +//! Authority derivation (spec §9.2). +//! +//! Authority is **derived, never asserted**. A profile carries no authority +//! claim of its own; it emerges from the fidelity of the variants actually +//! selected, capped by what programme traversal found. That closes the +//! loophole where a lossy extraction becomes "exact" by being named the +//! baseline — the baseline's fidelity is recorded against the source +//! checkpoint, not against itself. +//! +//! The fold is deliberately boring: +//! +//! ```text +//! weakest selected region fidelity +//! ↓ cap by operation completeness (absent operands) +//! ↓ cap by declared structural omission / replacement +//! derived authority +//! ``` +//! +//! It never inspects a filename, a programme name or a profile name, and it +//! does not care whether execution runs on the reference path or a Production +//! kernel. **Kernel maturity affects speed and support status, not fidelity.** +//! Conflating them would let a slow-but-exact path present as approximate, or +//! a fast approximation present as exact. +//! +//! The invariant that makes this safe, and that the property tests below +//! enforce: *replacing a selected component with one of weaker fidelity, or +//! removing an operand, can never increase any derived authority.* + +use serde::{Deserialize, Serialize}; + +/// How faithful a region or profile is to the source checkpoint (§9.2). +/// +/// `Ord` is the authority lattice: greater is stronger. The fold is a +/// minimum over this order, which is what makes weakening monotonic. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "kebab-case")] +pub enum Fidelity { + /// Incapable of complete forward execution — router/browse slices. + AnalysisOnly, + /// Components omitted or replaced (reduced top-K, shared-only layers). + StructurallyApproximate, + /// Same architecture, lossy representation (Q6_K quantised from BF16). + NumericallyApproximate, + /// Different encoding whose decode reproduces the source values exactly + /// (a lossless Q6_K container of native MXFP4 values). + SourceEquivalent, + /// Decoded values bit-identical to the source, in its own encoding family. + SourceExact, +} + +impl Fidelity { + pub const fn name(self) -> &'static str { + match self { + Self::AnalysisOnly => "analysis-only", + Self::StructurallyApproximate => "structurally-approximate", + Self::NumericallyApproximate => "numerically-approximate", + Self::SourceEquivalent => "source-equivalent", + Self::SourceExact => "source-exact", + } + } + + /// Whether a complete forward pass is claimable at this level. + pub const fn permits_complete_execution(self) -> bool { + !matches!(self, Self::AnalysisOnly) + } +} + +/// A declared structural change that caps authority regardless of fidelity. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum StructuralChange { + /// Components dropped — must name them (§9.2). + Omitted(Vec), + /// A component swapped for a compact approximation — must name it. + Replaced(String), +} + +impl StructuralChange { + pub fn describe(&self) -> String { + match self { + Self::Omitted(parts) => format!("omitted: {}", parts.join(", ")), + Self::Replaced(what) => format!("replaced: {what}"), + } + } +} + +/// Inputs to the fold. Deliberately contains no names, paths or profile +/// identity — if it did, the fold could consult them. +#[derive(Debug, Clone, PartialEq, Eq, Default)] +pub struct AuthorityInputs { + /// Fidelity of every actively selected region. + pub selected_fidelities: Vec, + /// Whether traversal found every required operand for a complete forward + /// pass. False caps at `analysis-only`. + pub execution_complete: bool, + /// Declared structural changes, if any. + pub structural: Option, +} + +/// Fold the inputs into a derived authority. +/// +/// Monotone in every input: weakening any fidelity, clearing +/// `execution_complete`, or adding a structural change can only lower the +/// result. +pub fn derive_authority(inputs: &AuthorityInputs) -> Fidelity { + // Stage 1 — weakest selected region fidelity. + // + // An empty selection selects nothing, and nothing cannot be exact. It is + // the browse-slice-with-no-regions case, so it floors rather than defaults + // high. + let weakest = inputs + .selected_fidelities + .iter() + .copied() + .min() + .unwrap_or(Fidelity::AnalysisOnly); + + // Stage 2 — cap by operation completeness. Absent operands mean no + // complete forward pass, whatever the surviving bytes are worth. + let after_completeness = if inputs.execution_complete { + weakest + } else { + weakest.min(Fidelity::AnalysisOnly) + }; + + // Stage 3 — cap by declared structural omission or replacement. + match inputs.structural { + Some(_) => after_completeness.min(Fidelity::StructurallyApproximate), + None => after_completeness, + } +} + +/// A derived authority alongside why it landed there. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DerivedAuthority { + pub level: Fidelity, + pub weakest_selected: Fidelity, + pub capped_by_completeness: bool, + pub capped_by_structure: bool, +} + +impl DerivedAuthority { + pub fn of(inputs: &AuthorityInputs) -> Self { + let weakest = inputs + .selected_fidelities + .iter() + .copied() + .min() + .unwrap_or(Fidelity::AnalysisOnly); + let level = derive_authority(inputs); + Self { + level, + weakest_selected: weakest, + capped_by_completeness: !inputs.execution_complete && weakest > Fidelity::AnalysisOnly, + capped_by_structure: inputs.structural.is_some() + && weakest > Fidelity::StructurallyApproximate + && inputs.execution_complete, + } + } + + /// Whether a profile may claim `claimed`. Claiming *below* the derived + /// level is always allowed (§9.2); claiming above never is. + pub fn permits_claim(&self, claimed: Fidelity) -> bool { + claimed <= self.level + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const ALL: [Fidelity; 5] = [ + Fidelity::AnalysisOnly, + Fidelity::StructurallyApproximate, + Fidelity::NumericallyApproximate, + Fidelity::SourceEquivalent, + Fidelity::SourceExact, + ]; + + fn complete(fidelities: &[Fidelity]) -> AuthorityInputs { + AuthorityInputs { + selected_fidelities: fidelities.to_vec(), + execution_complete: true, + structural: None, + } + } + + #[test] + fn the_lattice_orders_from_analysis_only_to_source_exact() { + let mut sorted = ALL; + sorted.sort(); + assert_eq!(sorted, ALL); + // Every adjacent link, named. `Ord` is derived, so the lattice is the + // *declaration order* — and the fold is a `min` over it. A variant + // moved one place would silently invert an authority comparison, and + // `sorted == ALL` alone cannot catch that, because reordering the enum + // and this list together keeps it passing. In particular + // source-equivalent (lossless re-encoding) must outrank + // numerically-approximate (lossy), which is the pair most easily read + // the wrong way round. + assert!(Fidelity::SourceExact > Fidelity::SourceEquivalent); + assert!(Fidelity::SourceEquivalent > Fidelity::NumericallyApproximate); + assert!(Fidelity::NumericallyApproximate > Fidelity::StructurallyApproximate); + assert!(Fidelity::StructurallyApproximate > Fidelity::AnalysisOnly); + } + + #[test] + fn names_match_the_spec_spelling() { + assert_eq!(Fidelity::SourceExact.name(), "source-exact"); + assert_eq!(Fidelity::SourceEquivalent.name(), "source-equivalent"); + assert_eq!(Fidelity::AnalysisOnly.name(), "analysis-only"); + } + + #[test] + fn fidelity_round_trips_through_json_in_spec_spelling() { + for f in ALL { + let json = serde_json::to_string(&f).unwrap(); + assert!(json.contains(f.name()), "{json}"); + assert_eq!(serde_json::from_str::(&json).unwrap(), f); + } + } + + #[test] + fn only_analysis_only_forbids_complete_execution() { + for f in ALL { + assert_eq!( + f.permits_complete_execution(), + f != Fidelity::AnalysisOnly, + "{}", + f.name() + ); + } + } + + #[test] + fn authority_is_the_weakest_selected_fidelity() { + let inputs = complete(&[Fidelity::SourceExact, Fidelity::NumericallyApproximate]); + assert_eq!(derive_authority(&inputs), Fidelity::NumericallyApproximate); + } + + #[test] + fn one_weak_region_drags_the_whole_profile_down() { + // The point of weakest-link: a mixed-precision index is only as + // faithful as its least faithful active selection. + let mut f = vec![Fidelity::SourceExact; 20]; + f.push(Fidelity::NumericallyApproximate); + assert_eq!( + derive_authority(&complete(&f)), + Fidelity::NumericallyApproximate + ); + } + + #[test] + fn an_empty_selection_floors_rather_than_defaulting_high() { + // Nothing selected cannot be exact. Defaulting to SourceExact here + // would make an empty profile the most authoritative one. + assert_eq!(derive_authority(&complete(&[])), Fidelity::AnalysisOnly); + } + + #[test] + fn incomplete_execution_caps_at_analysis_only() { + let inputs = AuthorityInputs { + selected_fidelities: vec![Fidelity::SourceExact], + execution_complete: false, + structural: None, + }; + assert_eq!(derive_authority(&inputs), Fidelity::AnalysisOnly); + } + + #[test] + fn a_structural_change_caps_at_structurally_approximate() { + let inputs = AuthorityInputs { + selected_fidelities: vec![Fidelity::SourceExact], + execution_complete: true, + structural: Some(StructuralChange::Omitted(vec!["routed branch".into()])), + }; + assert_eq!(derive_authority(&inputs), Fidelity::StructurallyApproximate); + } + + #[test] + fn a_structural_change_cannot_raise_a_weaker_selection() { + // Capping is a minimum, never a floor. + let inputs = AuthorityInputs { + selected_fidelities: vec![Fidelity::AnalysisOnly], + execution_complete: true, + structural: Some(StructuralChange::Replaced("down".into())), + }; + assert_eq!(derive_authority(&inputs), Fidelity::AnalysisOnly); + } + + #[test] + fn structural_changes_name_what_changed() { + assert!(StructuralChange::Omitted(vec!["lm_head".into()]) + .describe() + .contains("lm_head")); + assert!(StructuralChange::Replaced("down".into()) + .describe() + .contains("down")); + } + + // ── The monotonicity invariant ────────────────────────────────────── + // + // Exhaustive over the whole input lattice rather than sampled: the space + // is small enough that "property test" can mean "proof by enumeration". + + #[test] + fn weakening_any_region_never_raises_authority() { + for &a in &ALL { + for &b in &ALL { + for &weaker in &ALL { + if weaker > b { + continue; // only consider genuine weakenings + } + let before = derive_authority(&complete(&[a, b])); + let after = derive_authority(&complete(&[a, weaker])); + assert!( + after <= before, + "weakening {} → {} raised {} to {}", + b.name(), + weaker.name(), + before.name(), + after.name() + ); + } + } + } + } + + #[test] + fn adding_a_region_never_raises_authority() { + // More operands can only introduce a new weakest link. + for &a in &ALL { + for &added in &ALL { + let before = derive_authority(&complete(&[a])); + let after = derive_authority(&complete(&[a, added])); + assert!(after <= before, "adding {} raised authority", added.name()); + } + } + } + + #[test] + fn losing_execution_completeness_never_raises_authority() { + for &a in &ALL { + let mut inputs = complete(&[a]); + let before = derive_authority(&inputs); + inputs.execution_complete = false; + assert!(derive_authority(&inputs) <= before, "{}", a.name()); + } + } + + #[test] + fn declaring_a_structural_change_never_raises_authority() { + for &a in &ALL { + for complete_exec in [true, false] { + let base = AuthorityInputs { + selected_fidelities: vec![a], + execution_complete: complete_exec, + structural: None, + }; + let before = derive_authority(&base); + let after = derive_authority(&AuthorityInputs { + structural: Some(StructuralChange::Replaced("x".into())), + ..base + }); + assert!(after <= before, "{}", a.name()); + } + } + } + + #[test] + fn the_fold_is_order_independent() { + // A minimum cannot depend on selection order; pinned because a future + // "first wins" shortcut would silently break weakest-link. + let forward = derive_authority(&complete(&[ + Fidelity::SourceExact, + Fidelity::NumericallyApproximate, + ])); + let reverse = derive_authority(&complete(&[ + Fidelity::NumericallyApproximate, + Fidelity::SourceExact, + ])); + assert_eq!(forward, reverse); + } + + #[test] + fn a_profile_may_claim_at_or_below_its_derived_level() { + let d = DerivedAuthority::of(&complete(&[Fidelity::SourceEquivalent])); + assert_eq!(d.level, Fidelity::SourceEquivalent); + assert!(d.permits_claim(Fidelity::SourceEquivalent)); + // Voluntarily claiming lower is explicitly allowed (§9.2). + assert!(d.permits_claim(Fidelity::NumericallyApproximate)); + } + + #[test] + fn a_profile_may_never_claim_above_its_derived_level() { + let d = DerivedAuthority::of(&complete(&[Fidelity::NumericallyApproximate])); + assert!(!d.permits_claim(Fidelity::SourceEquivalent)); + assert!(!d.permits_claim(Fidelity::SourceExact)); + } + + #[test] + fn the_derivation_reports_which_cap_bound_it() { + let by_completeness = DerivedAuthority::of(&AuthorityInputs { + selected_fidelities: vec![Fidelity::SourceExact], + execution_complete: false, + structural: None, + }); + assert!(by_completeness.capped_by_completeness); + assert!(!by_completeness.capped_by_structure); + + let by_structure = DerivedAuthority::of(&AuthorityInputs { + selected_fidelities: vec![Fidelity::SourceExact], + execution_complete: true, + structural: Some(StructuralChange::Omitted(vec!["x".into()])), + }); + assert!(by_structure.capped_by_structure); + assert!(!by_structure.capped_by_completeness); + } + + #[test] + fn an_uncapped_derivation_reports_neither_cap() { + let d = DerivedAuthority::of(&complete(&[Fidelity::SourceExact])); + assert!(!d.capped_by_completeness); + assert!(!d.capped_by_structure); + assert_eq!(d.level, d.weakest_selected); + } +} diff --git a/crates/larql-vindex/src/format/capability/binding.rs b/crates/larql-vindex/src/format/capability/binding.rs new file mode 100644 index 000000000..36923dcbe --- /dev/null +++ b/crates/larql-vindex/src/format/capability/binding.rs @@ -0,0 +1,276 @@ +//! Semantic bindings and physical identity (spec §9.1, §11). +//! +//! Two identities, deliberately separate, because they answer different +//! questions and deduplicate on different keys: +//! +//! ```text +//! RepresentationIdentity which catalogue declaration was selected +//! PhysicalStorageId which bytes on disk that resolves to +//! ``` +//! +//! Folding the variant into physical identity would defeat the deduplication +//! it exists for: two catalogue declarations may name the same extent, and +//! residency, byte accounting and mmap planning care about the extent, not +//! about how many names point at it. +//! +//! ```text +//! same bytes, two semantic uses → one residency allocation +//! same bytes, two permitted views → one physical component, two bindings +//! same logical tensor, duplicated bytes → two physical components +//! ``` +//! +//! Two declarations naming one extent while disagreeing about format, shape, +//! packing or fidelity is a **catalogue defect**, not two components — the +//! bytes cannot be both. + +use super::authority::Fidelity; +use super::component::{ComponentContract, TensorKind}; + +/// A byte extent. The deduplication key for residency and accounting. +/// +/// Deliberately excludes the variant: identity is about *which bytes*, and two +/// catalogue names for one extent are one allocation. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct PhysicalStorageId { + pub file_digest: String, + pub offset: u64, + pub length: u64, +} + +impl PhysicalStorageId { + pub fn new(file_digest: impl Into, offset: u64, length: u64) -> Self { + Self { + file_digest: file_digest.into(), + offset, + length, + } + } + + pub fn end(&self) -> u64 { + self.offset + self.length + } + + /// Whether two extents are the same bytes. Only exact identity is safe to + /// deduplicate — partial overlap stays distinct unless a storage format + /// explicitly defines the relationship. + pub fn is_same_extent(&self, other: &Self) -> bool { + self == other + } + + /// Whether two extents overlap without being identical. Such a pair must + /// **not** be deduplicated; it is either a legitimate sub-range or a + /// catalogue error, and neither is "the same component". + pub fn overlaps_partially(&self, other: &Self) -> bool { + self.file_digest == other.file_digest + && self != other + && self.offset < other.end() + && other.offset < self.end() + } + + pub fn describe(&self) -> String { + format!( + "{}..{} of {}", + self.offset, + self.end(), + &self.file_digest[..self.file_digest.len().min(12)] + ) + } +} + +/// Which catalogue declaration was selected. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct RepresentationIdentity { + pub region_set: String, + pub variant: String, +} + +impl RepresentationIdentity { + pub fn new(region_set: impl Into, variant: impl Into) -> Self { + Self { + region_set: region_set.into(), + variant: variant.into(), + } + } + + pub fn describe(&self) -> String { + format!("{}::{}", self.region_set, self.variant) + } +} + +/// How stored bytes are accessed to satisfy a semantic role. +/// +/// A bounded vocabulary, not arbitrary tensor transformation. It exists +/// because embedding-as-LM-head is not the same *access* as a dedicated head: +/// the embedding is indexed by token, the projection contracts over hidden, so +/// the same bytes serve both only through a transpose. Without recording the +/// view, a plan knows which bytes it reads and not how they perform the role. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ComponentView { + Direct, + Transpose, + Slice { dim: usize, start: u32, len: u32 }, +} + +/// Why a view cannot be applied to the stored shape. +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub enum ViewError { + #[error("cannot transpose a {kind}: transpose is defined for matrices only")] + TransposeOfNonMatrix { kind: &'static str }, + #[error("slice dim {dim} is out of range for a shape with {rank} dimensions")] + SliceDimOutOfRange { dim: usize, rank: usize }, + #[error("slice {start}..{end} is out of range for dimension {dim} of length {extent}")] + SliceOutOfRange { + dim: usize, + start: u32, + end: u32, + extent: u32, + }, +} + +impl ComponentView { + /// The contract the *semantic role* sees, after this view is applied. + /// + /// Contract checking must happen on this side of the view. A dedicated + /// head is `[hidden, vocab]`; an embedding is `[vocab, hidden]`; comparing + /// the requirement against raw storage would reject a perfectly valid + /// tied-weight model. + pub fn apply_to(&self, storage: &ComponentContract) -> Result { + match self { + Self::Direct => Ok(storage.clone()), + Self::Transpose => { + if storage.kind != TensorKind::Matrix || storage.shape.len() != 2 { + return Err(ViewError::TransposeOfNonMatrix { + kind: storage.kind.name(), + }); + } + Ok(ComponentContract { + shape: vec![storage.shape[1], storage.shape[0]], + kind: TensorKind::Matrix, + }) + } + Self::Slice { dim, start, len } => { + let rank = storage.shape.len(); + let extent = *storage + .shape + .get(*dim) + .ok_or(ViewError::SliceDimOutOfRange { dim: *dim, rank })?; + let end = start + len; + if end > extent { + return Err(ViewError::SliceOutOfRange { + dim: *dim, + start: *start, + end, + extent, + }); + } + let mut shape = storage.shape.clone(); + shape[*dim] = *len; + Ok(ComponentContract { + shape, + kind: storage.kind, + }) + } + } + } + + /// Whether this view changes any stored value. + /// + /// None of them do. A transpose reorders access, it does not alter + /// numbers, so **views never affect authority**. They affect reference + /// support, kernel eligibility and access efficiency — a kernel that + /// serves a dedicated head may not accept transposed embedding access, and + /// the loader must expose that route at Reference maturity rather than + /// silently materialising a repacked copy at decode time. + pub const fn alters_values(&self) -> bool { + false + } + + pub fn describe(&self) -> String { + match self { + Self::Direct => "direct".into(), + Self::Transpose => "transposed".into(), + Self::Slice { dim, start, len } => { + format!("slice dim {dim} [{start}..{}]", start + len) + } + } + } +} + +/// One semantic use of one physical extent. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ComponentBinding { + /// Which declaration was selected. + pub representation: RepresentationIdentity, + /// Which bytes that resolves to — the deduplication key. + pub storage: PhysicalStorageId, + /// How those bytes perform this role. + pub view: ComponentView, + /// Fidelity against the source checkpoint. Unaffected by the view. + pub fidelity: Fidelity, +} + +impl ComponentBinding { + pub fn describe(&self) -> String { + format!( + "{} ({}, {})", + self.representation.describe(), + self.storage.describe(), + self.view.describe() + ) + } +} + +/// Deduplicate bindings to the physical extents they read. +/// +/// Two semantic uses of one tied tensor produce two bindings and one extent. +/// Byte accounting, residency estimates, mmap planning, checksum work and +/// remote-transfer planning all consume this side; execution and kernel +/// binding consume the bindings. +pub fn physical_extents(bindings: &[ComponentBinding]) -> Vec { + let mut out: Vec = bindings.iter().map(|b| b.storage.clone()).collect(); + out.sort(); + out.dedup(); + out +} + +/// Total bytes an execution reads, counting each extent once. +pub fn total_bytes(bindings: &[ComponentBinding]) -> u64 { + physical_extents(bindings).iter().map(|s| s.length).sum() +} + +/// Two declarations naming one extent while disagreeing about it. +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +#[error( + "representations '{a}' and '{b}' name the same extent ({extent}) but disagree on {field} — \ + the bytes cannot be both" +)] +pub struct CatalogueConflict { + pub a: String, + pub b: String, + pub extent: String, + pub field: &'static str, +} + +/// Detect declarations that share an extent but describe it differently. +/// +/// Not two components — a catalogue defect. Resolution must surface it rather +/// than pick one, because either declaration could be the wrong one. +pub fn catalogue_conflicts(bindings: &[ComponentBinding]) -> Vec { + let mut out = Vec::new(); + for (i, a) in bindings.iter().enumerate() { + for b in bindings.iter().skip(i + 1) { + if !a.storage.is_same_extent(&b.storage) { + continue; + } + if a.fidelity != b.fidelity { + out.push(CatalogueConflict { + a: a.representation.describe(), + b: b.representation.describe(), + extent: a.storage.describe(), + field: "fidelity", + }); + } + } + } + out +} diff --git a/crates/larql-vindex/src/format/capability/binding_tests.rs b/crates/larql-vindex/src/format/capability/binding_tests.rs new file mode 100644 index 000000000..80ec4429e --- /dev/null +++ b/crates/larql-vindex/src/format/capability/binding_tests.rs @@ -0,0 +1,265 @@ +//! Colocated tests for `binding` — identity, views and deduplication. +//! +//! The headline case throughout is the tied LM head: one tensor, two semantic +//! roles, two views, one allocation. + +use super::authority::Fidelity; +use super::binding::{ + catalogue_conflicts, physical_extents, total_bytes, ComponentBinding, ComponentView, + PhysicalStorageId, RepresentationIdentity, ViewError, +}; +use super::component::{ComponentContract, TensorKind}; + +const VOCAB: u32 = 32_000; +const HIDDEN: u32 = 4_096; +const DIGEST: &str = "sha256:aabbccddeeff0011"; + +fn extent(offset: u64, length: u64) -> PhysicalStorageId { + PhysicalStorageId::new(DIGEST, offset, length) +} + +fn binding(name: &str, storage: PhysicalStorageId, view: ComponentView) -> ComponentBinding { + ComponentBinding { + representation: RepresentationIdentity::new("embeddings", name), + storage, + view, + fidelity: Fidelity::SourceExact, + } +} + +// ── Identity is about bytes, not names ───────────────────────────────────── + +#[test] +fn two_variant_names_for_one_extent_deduplicate_to_one_component() { + // The correction: folding the variant into physical identity would defeat + // the deduplication it exists for. + let a = binding("baseline", extent(0, 1_024), ComponentView::Direct); + let b = ComponentBinding { + representation: RepresentationIdentity::new("embeddings", "alias"), + ..a.clone() + }; + assert_ne!(a.representation, b.representation); + assert_eq!(physical_extents(&[a, b]).len(), 1); +} + +#[test] +fn one_tensor_serving_two_roles_is_one_allocation() { + // Tied LM head: two bindings, two views, one extent. + let shared = extent(0, 1_024); + let bindings = [ + binding("baseline", shared.clone(), ComponentView::Direct), + binding("baseline", shared, ComponentView::Transpose), + ]; + assert_eq!(bindings.len(), 2, "both semantic uses survive"); + assert_eq!(physical_extents(&bindings).len(), 1); + assert_eq!(total_bytes(&bindings), 1_024, "counted once, not twice"); +} + +#[test] +fn a_duplicated_logical_tensor_is_two_components() { + // Same logical tensor written twice is two allocations, however tied it + // is conceptually. + let bindings = [ + binding("baseline", extent(0, 1_024), ComponentView::Direct), + binding("copy", extent(4_096, 1_024), ComponentView::Direct), + ]; + assert_eq!(physical_extents(&bindings).len(), 2); + assert_eq!(total_bytes(&bindings), 2_048); +} + +#[test] +fn only_exact_extents_are_treated_as_the_same_bytes() { + let whole = extent(0, 1_024); + let half = extent(0, 512); + assert!(whole.is_same_extent(&extent(0, 1_024))); + assert!(!whole.is_same_extent(&half)); +} + +#[test] +fn partial_overlap_is_detected_and_never_deduplicated() { + // A sub-range is either legitimate or a catalogue error; neither makes it + // "the same component". + let a = extent(0, 1_024); + let b = extent(512, 1_024); + assert!(a.overlaps_partially(&b)); + assert!(!a.is_same_extent(&b)); + assert_eq!( + physical_extents(&[ + binding("a", a, ComponentView::Direct), + binding("b", b, ComponentView::Direct), + ]) + .len(), + 2 + ); +} + +#[test] +fn extents_in_different_files_never_overlap() { + let a = PhysicalStorageId::new("sha256:aaaa", 0, 1_024); + let b = PhysicalStorageId::new("sha256:bbbb", 0, 1_024); + assert!(!a.overlaps_partially(&b)); + assert!(!a.is_same_extent(&b)); +} + +#[test] +fn adjacent_extents_do_not_overlap() { + let a = extent(0, 1_024); + let b = extent(1_024, 1_024); + assert!(!a.overlaps_partially(&b)); +} + +// ── Views transform contracts, not values ────────────────────────────────── + +#[test] +fn a_transpose_turns_an_embedding_into_an_lm_head_shape() { + // The case a raw-storage contract check would have rejected. + let storage = ComponentContract::matrix(VOCAB, HIDDEN); + let viewed = ComponentView::Transpose.apply_to(&storage).unwrap(); + assert_eq!(viewed, ComponentContract::matrix(HIDDEN, VOCAB)); +} + +#[test] +fn a_direct_view_leaves_the_contract_alone() { + let storage = ComponentContract::matrix(HIDDEN, VOCAB); + assert_eq!(ComponentView::Direct.apply_to(&storage).unwrap(), storage); +} + +#[test] +fn transposing_a_vector_is_refused_rather_than_silently_accepted() { + let err = ComponentView::Transpose + .apply_to(&ComponentContract::vector(HIDDEN)) + .unwrap_err(); + assert!(matches!(err, ViewError::TransposeOfNonMatrix { .. })); + assert!(err.to_string().contains("matrices only")); +} + +#[test] +fn a_slice_narrows_one_dimension() { + let storage = ComponentContract::matrix(VOCAB, HIDDEN); + let viewed = ComponentView::Slice { + dim: 0, + start: 0, + len: 1_000, + } + .apply_to(&storage) + .unwrap(); + assert_eq!(viewed.shape, vec![1_000, HIDDEN]); + assert_eq!(viewed.kind, TensorKind::Matrix); +} + +#[test] +fn a_slice_past_the_end_is_refused_with_both_bounds() { + let err = ComponentView::Slice { + dim: 0, + start: VOCAB - 10, + len: 100, + } + .apply_to(&ComponentContract::matrix(VOCAB, HIDDEN)) + .unwrap_err(); + let text = err.to_string(); + assert!(text.contains("out of range"), "{text}"); + assert!(text.contains(&VOCAB.to_string()), "{text}"); +} + +#[test] +fn a_slice_on_a_missing_dimension_is_refused() { + let err = ComponentView::Slice { + dim: 5, + start: 0, + len: 1, + } + .apply_to(&ComponentContract::vector(HIDDEN)) + .unwrap_err(); + assert!(matches!(err, ViewError::SliceDimOutOfRange { .. })); +} + +#[test] +fn no_view_alters_stored_values() { + // Which is why views never touch authority. They affect kernel + // eligibility and access efficiency instead. + for view in [ + ComponentView::Direct, + ComponentView::Transpose, + ComponentView::Slice { + dim: 0, + start: 0, + len: 1, + }, + ] { + assert!(!view.alters_values(), "{}", view.describe()); + } +} + +#[test] +fn a_transposed_binding_keeps_the_fidelity_of_its_bytes() { + let direct = binding("baseline", extent(0, 1_024), ComponentView::Direct); + let transposed = ComponentBinding { + view: ComponentView::Transpose, + ..direct.clone() + }; + assert_eq!(direct.fidelity, transposed.fidelity); +} + +// ── Catalogue conflicts ──────────────────────────────────────────────────── + +#[test] +fn two_declarations_disagreeing_about_one_extent_are_a_catalogue_defect() { + // Not two components — the bytes cannot be both. + let shared = extent(0, 1_024); + let a = ComponentBinding { + representation: RepresentationIdentity::new("gate_up", "exact-q6k"), + storage: shared.clone(), + view: ComponentView::Direct, + fidelity: Fidelity::SourceEquivalent, + }; + let b = ComponentBinding { + representation: RepresentationIdentity::new("gate_up", "native-mxfp4"), + storage: shared, + view: ComponentView::Direct, + fidelity: Fidelity::SourceExact, + }; + let conflicts = catalogue_conflicts(&[a, b]); + assert_eq!(conflicts.len(), 1); + let text = conflicts[0].to_string(); + assert!(text.contains("fidelity"), "{text}"); + assert!(text.contains("cannot be both"), "{text}"); +} + +#[test] +fn agreeing_declarations_of_one_extent_are_not_a_conflict() { + let shared = extent(0, 1_024); + let conflicts = catalogue_conflicts(&[ + binding("baseline", shared.clone(), ComponentView::Direct), + binding("alias", shared, ComponentView::Transpose), + ]); + assert!(conflicts.is_empty(), "differing views are not a conflict"); +} + +#[test] +fn declarations_of_different_extents_are_never_a_conflict() { + let conflicts = catalogue_conflicts(&[ + ComponentBinding { + fidelity: Fidelity::SourceExact, + ..binding("a", extent(0, 1_024), ComponentView::Direct) + }, + ComponentBinding { + fidelity: Fidelity::NumericallyApproximate, + ..binding("b", extent(2_048, 1_024), ComponentView::Direct) + }, + ]); + assert!(conflicts.is_empty()); +} + +#[test] +fn a_binding_describes_representation_extent_and_view() { + let s = binding("baseline", extent(64, 1_024), ComponentView::Transpose).describe(); + assert!(s.contains("embeddings::baseline"), "{s}"); + assert!(s.contains("64..1088"), "{s}"); + assert!(s.contains("transposed"), "{s}"); +} + +#[test] +fn no_bindings_means_no_bytes() { + assert_eq!(total_bytes(&[]), 0); + assert!(physical_extents(&[]).is_empty()); +} diff --git a/crates/larql-vindex/src/format/capability/compatibility.rs b/crates/larql-vindex/src/format/capability/compatibility.rs new file mode 100644 index 000000000..8f532e741 --- /dev/null +++ b/crates/larql-vindex/src/format/capability/compatibility.rs @@ -0,0 +1,407 @@ +//! Cross-role segment compatibility (spec §11). +//! +//! `OperandCapability` says what each selected region can provide. +//! `SegmentCompatibility` says whether those operands can participate in **one +//! computation**. Execution capability derives from both. +//! +//! The distinction is why this is not another operand state. Partial coverage +//! is role-local and true of one role on its own: +//! +//! ```text +//! down is present in segments {0, 1}, missing from {2, 3} +//! ``` +//! +//! Incompatibility is relational and true of no role individually: +//! +//! ```text +//! gate_up_fused covers {0, 1} +//! down covers {2, 3} +//! common {} +//! ``` +//! +//! Assigning that to each operand separately would state the symptom twice and +//! the cause nowhere. Two role-local partial reports describe *what is +//! missing*; the relational finding describes *why no currently selected +//! segment population is executable* — which stays the right framing for +//! disjoint, overlapping and irregular coverage alike. +//! +//! Reporting it as primary is **subsumption, not suppression**: the per-role +//! coverage facts are retained as evidence on the finding, so machines and +//! verbose inspection keep them while the user-facing diagnosis names one +//! cause. + +use crate::format::lyrw2::region_role::RegionRole; + +/// A required role's usable coverage. +/// +/// "Usable" excludes segments whose bytes are present but whose codec this +/// build cannot interpret — an uninterpretable region cannot participate in a +/// computation, so it does not count toward coverage even though it exists. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RoleSegmentCoverage { + pub role: RegionRole, + pub covered: Vec, +} + +impl RoleSegmentCoverage { + pub fn new(role: RegionRole, mut covered: Vec) -> Self { + covered.sort_unstable(); + covered.dedup(); + Self { role, covered } + } + + pub fn describe(&self) -> String { + format!("{} covers {}", self.role.name(), render(&self.covered)) + } +} + +/// How two or more roles' coverage fails to agree. +/// +/// Both shapes admit identically — the alternative cannot execute over the +/// required set — but they are different situations and a reader deserves to +/// know which. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum IncompatibilityShape { + /// No segment carries every required role. The strongest case: there is + /// nowhere at all the computation could run. + NoCommonCoverage, + /// Some segments carry every role, but the roles disagree on which. Part + /// of the population is executable; the selection as a whole is not. + UnequalCoverage, +} + +impl IncompatibilityShape { + pub const fn name(self) -> &'static str { + match self { + Self::NoCommonCoverage => "no common coverage", + Self::UnequalCoverage => "unequal coverage", + } + } +} + +/// Whether one programme alternative's operands span a common population. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum SegmentCompatibility { + /// Every required role covers the whole required set. + Compatible, + /// Every required role covers the *same* set, and that set is short of + /// what was required. No cross-role disagreement — the selection is + /// consistently incomplete. + Partial { + covered: Vec, + missing: Vec, + }, + /// Required roles cover different sets. + Incompatible { + shape: IncompatibilityShape, + per_role: Vec, + common: Vec, + required: Vec, + }, +} + +impl SegmentCompatibility { + /// Decide compatibility from each required role's usable coverage. + /// + /// Callers must only pass roles whose coverage is non-empty — a role with + /// no coverage anywhere is a role-local failure that outranks this check, + /// because compatibility is moot when an operand does not exist. + pub fn evaluate(required: &[u16], coverage: &[RoleSegmentCoverage]) -> Self { + let required = normalise(required); + + if coverage.is_empty() { + // Nothing required to agree with anything. + return Self::Compatible; + } + + let first = &coverage[0].covered; + let all_equal = coverage.iter().all(|c| &c.covered == first); + + if all_equal { + let missing: Vec = required + .iter() + .copied() + .filter(|s| !first.contains(s)) + .collect(); + return if missing.is_empty() { + Self::Compatible + } else { + Self::Partial { + covered: first.clone(), + missing, + } + }; + } + + let common = intersect(coverage); + Self::Incompatible { + shape: if common.is_empty() { + IncompatibilityShape::NoCommonCoverage + } else { + IncompatibilityShape::UnequalCoverage + }, + per_role: coverage.to_vec(), + common, + required, + } + } + + /// Whether the alternative can execute over the whole required set. + pub fn is_compatible(&self) -> bool { + matches!(self, Self::Compatible) + } + + /// Whether this is an **invalid selection** rather than an incomplete one. + /// + /// Incompatibility must fail closed. It is not a declared approximation + /// policy, so it must never be laundered into `structurally-approximate` + /// authority — a deliberately omitted `down` in a browse slice is + /// intentional and derives `analysis-only`; two roles accidentally + /// resolving to disjoint populations is a resolution error. + pub fn is_invalid_selection(&self) -> bool { + matches!(self, Self::Incompatible { .. }) + } + + pub fn describe(&self) -> String { + match self { + Self::Compatible => "all required roles cover the selected segments".into(), + Self::Partial { covered, missing } => format!( + "all required roles cover segments {}, short of {}", + render(covered), + render(missing) + ), + Self::Incompatible { + shape, + per_role, + common, + required, + } => { + let roles = per_role + .iter() + .map(|c| c.describe()) + .collect::>() + .join("; "); + format!( + "required operands have incompatible segment sets ({}): {}; common: {}; \ + required: {}", + shape.name(), + roles, + render(common), + render(required) + ) + } + } + } + + /// Role-local facts this finding subsumes — retained as evidence so the + /// information is not lost, and not emitted as peer root causes. + pub fn subsumed_per_role_gaps(&self) -> Vec<(RegionRole, Vec)> { + match self { + Self::Incompatible { + per_role, required, .. + } => per_role + .iter() + .map(|c| { + let missing = required + .iter() + .copied() + .filter(|s| !c.covered.contains(s)) + .collect(); + (c.role, missing) + }) + .collect(), + _ => Vec::new(), + } + } +} + +fn normalise(segments: &[u16]) -> Vec { + let mut v = segments.to_vec(); + v.sort_unstable(); + v.dedup(); + v +} + +fn intersect(coverage: &[RoleSegmentCoverage]) -> Vec { + let Some(first) = coverage.first() else { + return Vec::new(); + }; + first + .covered + .iter() + .copied() + .filter(|s| coverage.iter().all(|c| c.covered.contains(s))) + .collect() +} + +fn render(segments: &[u16]) -> String { + if segments.is_empty() { + return "none".into(); + } + segments + .iter() + .map(|s| s.to_string()) + .collect::>() + .join(", ") +} + +#[cfg(test)] +mod tests { + use super::*; + + const GATE_UP: RegionRole = RegionRole::GateUpFused; + const DOWN: RegionRole = RegionRole::Down; + const REQUIRED: [u16; 4] = [0, 1, 2, 3]; + + fn cov(role: RegionRole, segs: &[u16]) -> RoleSegmentCoverage { + RoleSegmentCoverage::new(role, segs.to_vec()) + } + + #[test] + fn full_coverage_by_every_role_is_compatible() { + let c = SegmentCompatibility::evaluate( + &REQUIRED, + &[cov(GATE_UP, &REQUIRED), cov(DOWN, &REQUIRED)], + ); + assert_eq!(c, SegmentCompatibility::Compatible); + assert!(c.is_compatible()); + assert!(!c.is_invalid_selection()); + } + + #[test] + fn equal_but_short_coverage_is_partial_not_incompatible() { + // No cross-role disagreement — the selection is consistently + // incomplete, which is a different fix from reconciling two roles. + let c = + SegmentCompatibility::evaluate(&REQUIRED, &[cov(GATE_UP, &[0, 1]), cov(DOWN, &[0, 1])]); + assert_eq!( + c, + SegmentCompatibility::Partial { + covered: vec![0, 1], + missing: vec![2, 3], + } + ); + assert!(!c.is_invalid_selection()); + } + + #[test] + fn disjoint_coverage_reports_no_common_coverage() { + let c = + SegmentCompatibility::evaluate(&REQUIRED, &[cov(GATE_UP, &[0, 1]), cov(DOWN, &[2, 3])]); + match &c { + SegmentCompatibility::Incompatible { shape, common, .. } => { + assert_eq!(*shape, IncompatibilityShape::NoCommonCoverage); + assert!(common.is_empty()); + } + other => panic!("expected incompatible, got {other:?}"), + } + assert!(c.is_invalid_selection()); + } + + #[test] + fn overlapping_but_unequal_coverage_is_still_incompatible() { + // Part of the population is executable; the selection is not. Same + // admission result, different shape — and the reader deserves to know. + let c = SegmentCompatibility::evaluate( + &REQUIRED, + &[cov(GATE_UP, &[0, 1, 2]), cov(DOWN, &[1, 2, 3])], + ); + match &c { + SegmentCompatibility::Incompatible { shape, common, .. } => { + assert_eq!(*shape, IncompatibilityShape::UnequalCoverage); + assert_eq!(common, &vec![1, 2]); + } + other => panic!("expected incompatible, got {other:?}"), + } + } + + #[test] + fn the_diagnosis_names_every_roles_coverage() { + let c = + SegmentCompatibility::evaluate(&REQUIRED, &[cov(GATE_UP, &[0, 1]), cov(DOWN, &[2, 3])]); + let s = c.describe(); + assert!(s.contains("gate_up_fused covers 0, 1"), "{s}"); + assert!(s.contains("down covers 2, 3"), "{s}"); + assert!(s.contains("common: none"), "{s}"); + assert!(s.contains("required: 0, 1, 2, 3"), "{s}"); + } + + #[test] + fn subsumed_gaps_retain_the_role_local_facts() { + // Subsumption, not suppression: the partial-coverage facts survive as + // evidence even though they are not emitted as peer root causes. + let c = + SegmentCompatibility::evaluate(&REQUIRED, &[cov(GATE_UP, &[0, 1]), cov(DOWN, &[2, 3])]); + let gaps = c.subsumed_per_role_gaps(); + assert_eq!(gaps.len(), 2); + assert_eq!(gaps[0], (GATE_UP, vec![2, 3])); + assert_eq!(gaps[1], (DOWN, vec![0, 1])); + } + + #[test] + fn compatible_and_partial_findings_subsume_nothing() { + // There is no relational cause to subsume, so no evidence to carry. + assert!(SegmentCompatibility::Compatible + .subsumed_per_role_gaps() + .is_empty()); + let partial = SegmentCompatibility::evaluate(&REQUIRED, &[cov(DOWN, &[0])]); + assert!(partial.subsumed_per_role_gaps().is_empty()); + } + + #[test] + fn an_unsegmented_alternative_with_no_required_segments_is_compatible() { + let c = SegmentCompatibility::evaluate(&[], &[cov(GATE_UP, &[]), cov(DOWN, &[])]); + assert_eq!(c, SegmentCompatibility::Compatible); + } + + #[test] + fn no_required_roles_is_vacuously_compatible() { + assert_eq!( + SegmentCompatibility::evaluate(&REQUIRED, &[]), + SegmentCompatibility::Compatible + ); + } + + #[test] + fn coverage_is_normalised_so_input_order_does_not_change_the_verdict() { + let forward = SegmentCompatibility::evaluate( + &[3, 1, 0, 2], + &[cov(GATE_UP, &[1, 0]), cov(DOWN, &[0, 1])], + ); + let reverse = SegmentCompatibility::evaluate( + &[0, 1, 2, 3], + &[cov(GATE_UP, &[0, 1]), cov(DOWN, &[1, 0])], + ); + assert_eq!(forward, reverse); + } + + #[test] + fn three_roles_agreeing_pairwise_but_not_globally_are_incompatible() { + // Pairwise intersections are non-empty; the global one is not. A + // pairwise-only check would have called this compatible. + let c = SegmentCompatibility::evaluate( + &[0, 1, 2], + &[ + cov(RegionRole::Gate, &[0, 1]), + cov(RegionRole::Up, &[1, 2]), + cov(DOWN, &[2, 0]), + ], + ); + match &c { + SegmentCompatibility::Incompatible { shape, common, .. } => { + assert_eq!(*shape, IncompatibilityShape::NoCommonCoverage); + assert!(common.is_empty(), "{common:?}"); + } + other => panic!("expected incompatible, got {other:?}"), + } + } + + #[test] + fn incompatibility_is_never_a_declared_approximation() { + // It must fail closed, not be laundered into structurally-approximate + // authority. Pinned as a property of the type, not of a caller. + let c = SegmentCompatibility::evaluate(&REQUIRED, &[cov(GATE_UP, &[0]), cov(DOWN, &[1])]); + assert!(c.is_invalid_selection()); + assert!(!c.is_compatible()); + } +} diff --git a/crates/larql-vindex/src/format/capability/component.rs b/crates/larql-vindex/src/format/capability/component.rs new file mode 100644 index 000000000..f0b1d2207 --- /dev/null +++ b/crates/larql-vindex/src/format/capability/component.rs @@ -0,0 +1,427 @@ +//! Selected components — bank regions and manifest-addressed tensors. +//! +//! WALK reads only bank regions, so its plans never needed anything else. A +//! decode path also consumes embeddings, norms, attention/KDA/MLA projections, +//! routers, latent transforms and an LM head — none of which are bank regions +//! and all of which are manifest-addressed (§5). +//! +//! These live in the **resolved document selection**, never in an operation +//! environment. Three things must stay distinguishable: +//! +//! ```text +//! what the index contains the catalogue +//! what the profile selected the resolved selection ← these components +//! what the caller supplied the request / contract +//! ``` +//! +//! Putting stored weights in an environment beside caller-supplied inputs +//! collapses the second and third, and the loader stops being able to say +//! whether a missing tensor was never extracted, deliberately dropped, or +//! simply not passed in. + +use crate::format::lyrw2::region_role::RegionRole; + +use super::authority::Fidelity; +use super::coordinate::{AbsenceKind, RegionCoordinate}; +use super::role::ReferenceSupport; + +/// The five durable weight classes the serving ABI freezes (§4). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub enum WeightClass { + /// Embeddings, norms, LM head, routers, routing metadata, recurrence and + /// control parameters. + ControlAndRouter, + /// Attention / KDA / MLA projections. + DenseSpine, + /// Shared experts and shared latent pre/post projections. + SharedFfn, + RoutedGateUp, + RoutedDown, +} + +impl WeightClass { + pub const fn name(self) -> &'static str { + match self { + Self::ControlAndRouter => "control", + Self::DenseSpine => "dense", + Self::SharedFfn => "shared", + Self::RoutedGateUp => "routed_gate_up", + Self::RoutedDown => "routed_down", + } + } +} + +/// Where a component lives, across both addressing schemes. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub enum ComponentCoordinate { + BankRegion(RegionCoordinate), + ManifestTensor { + class: WeightClass, + /// `None` for model-global tensors such as embeddings or the LM head. + layer: Option, + key: String, + }, +} + +impl ComponentCoordinate { + pub fn tensor(class: WeightClass, layer: Option, key: impl Into) -> Self { + Self::ManifestTensor { + class, + layer, + key: key.into(), + } + } + + pub fn describe(&self) -> String { + match self { + Self::BankRegion(c) => c.describe(), + Self::ManifestTensor { class, layer, key } => match layer { + Some(l) => format!("layer {l} {} tensor '{key}'", class.name()), + None => format!("{} tensor '{key}'", class.name()), + }, + } + } + + pub fn role(&self) -> Option { + match self { + Self::BankRegion(c) => Some(c.role), + Self::ManifestTensor { .. } => None, + } + } +} + +/// What an adapter expects a component to be. +/// +/// Presence and readability are not enough. A readable tensor of the wrong +/// shape is neither an unsupported build nor a missing file — it is an invalid +/// index or an adapter mismatch, and it is a plausible-output trap: a +/// wrong-but-compatible shape can survive a long way down a generic buffer +/// path before failing, or worse, be interpreted and never fail at all. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ComponentContract { + pub shape: Vec, + pub kind: TensorKind, +} + +/// Coarse shape class, checked before dimensions. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TensorKind { + Matrix, + Vector, +} + +impl TensorKind { + pub const fn name(self) -> &'static str { + match self { + Self::Matrix => "matrix", + Self::Vector => "vector", + } + } +} + +impl ComponentContract { + pub fn matrix(rows: u32, cols: u32) -> Self { + Self { + shape: vec![rows, cols], + kind: TensorKind::Matrix, + } + } + + pub fn vector(len: u32) -> Self { + Self { + shape: vec![len], + kind: TensorKind::Vector, + } + } + + pub fn describe(&self) -> String { + format!("{} {:?}", self.kind.name(), self.shape) + } +} + +/// Why a component can or cannot be used, with the cause preserved. +/// +/// Five states because they imply five different repairs. Collapsing any pair +/// would make the report say "unavailable" where an operator needs to know +/// whether to fetch a file, upgrade the binary, fix the profile, or report a +/// bug in the extractor. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ComponentUsability { + Usable, + /// The profile dropped it on purpose. Not a defect. + DeliberatelyOmitted, + /// Missing with nothing declaring the omission. A defect. + AbsentUndeclared, + /// Present and intact; this build cannot interpret it. Not a defect — + /// the artifact may simply be newer. + UnsupportedEncoding(ReferenceSupport), + /// Present and readable, but not what the adapter expects. An invalid + /// index or an adapter mismatch, never a build limitation. + ContractMismatch { + expected: ComponentContract, + found: ComponentContract, + }, +} + +impl ComponentUsability { + pub fn is_usable(&self) -> bool { + matches!(self, Self::Usable) + } + + /// Whether this indicates something wrong with the artifact, as opposed to + /// a scoping choice or a build limitation. + pub fn indicates_defect(&self) -> bool { + matches!(self, Self::AbsentUndeclared | Self::ContractMismatch { .. }) + } + + pub fn describe(&self) -> String { + match self { + Self::Usable => "usable".into(), + Self::DeliberatelyOmitted => "omitted by the active selection".into(), + Self::AbsentUndeclared => "absent, with no declared omission".into(), + Self::UnsupportedEncoding(s) => s.describe(), + Self::ContractMismatch { expected, found } => format!( + "expected {} but found {} — invalid index or adapter mismatch", + expected.describe(), + found.describe() + ), + } + } +} + +/// A manifest-addressed tensor the profile selected. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SelectedTensor { + pub coordinate: ComponentCoordinate, + pub fidelity: Fidelity, + pub support: ReferenceSupport, + /// Set when the profile deliberately dropped this tensor. Distinguishes an + /// intentional client slice from a damaged index carrying the same gap. + pub omitted: Option, + /// What is actually stored, for contract checking. `None` when the + /// selection did not record it. + pub contract: Option, +} + +impl SelectedTensor { + pub fn present(coordinate: ComponentCoordinate, fidelity: Fidelity) -> Self { + Self { + coordinate, + fidelity, + support: ReferenceSupport::Supported, + omitted: None, + contract: None, + } + } + + pub fn with_contract(mut self, contract: ComponentContract) -> Self { + self.contract = Some(contract); + self + } + + /// Usability against an optional expected contract. + /// + /// Order matters: absence outranks encoding, which outranks shape. Asking + /// whether a missing tensor has the right shape is not a question. + pub fn usability(&self, expected: Option<&ComponentContract>) -> ComponentUsability { + if let Some(kind) = &self.omitted { + return if kind.is_defect() { + ComponentUsability::AbsentUndeclared + } else { + ComponentUsability::DeliberatelyOmitted + }; + } + if !self.support.is_supported() { + return ComponentUsability::UnsupportedEncoding(self.support.clone()); + } + match (expected, &self.contract) { + (Some(want), Some(have)) if want != have => ComponentUsability::ContractMismatch { + expected: want.clone(), + found: have.clone(), + }, + _ => ComponentUsability::Usable, + } + } + + /// Whether this tensor can contribute, ignoring any contract. + pub fn is_usable(&self) -> bool { + self.usability(None).is_usable() + } + + pub fn indicates_defect(&self) -> bool { + self.usability(None).indicates_defect() + } + + /// Why it cannot be used, ignoring any contract. + pub fn reason_unusable(&self) -> Option { + let u = self.usability(None); + (!u.is_usable()).then(|| u.describe()) + } +} + +/// Anything an operation plan can read. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum SelectedComponent { + BankRegion { + coordinate: RegionCoordinate, + fidelity: Fidelity, + }, + ManifestTensor(SelectedTensor), +} + +impl SelectedComponent { + pub fn region(coordinate: RegionCoordinate, fidelity: Fidelity) -> Self { + Self::BankRegion { + coordinate, + fidelity, + } + } + + /// Fidelity this component contributes to the authority fold. + pub fn fidelity(&self) -> Fidelity { + match self { + Self::BankRegion { fidelity, .. } => *fidelity, + Self::ManifestTensor(t) => t.fidelity, + } + } + + pub fn coordinate(&self) -> ComponentCoordinate { + match self { + Self::BankRegion { coordinate, .. } => { + ComponentCoordinate::BankRegion(coordinate.clone()) + } + Self::ManifestTensor(t) => t.coordinate.clone(), + } + } + + pub fn describe(&self) -> String { + self.coordinate().describe() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::format::lyrw2::region_format::RegionFormat; + + fn tensor_coordinate() -> ComponentCoordinate { + ComponentCoordinate::tensor( + WeightClass::ControlAndRouter, + Some(12), + "layers.12.router.weight", + ) + } + + #[test] + fn a_layered_tensor_names_its_layer_class_and_key() { + let s = tensor_coordinate().describe(); + assert!(s.contains("layer 12"), "{s}"); + assert!(s.contains("control"), "{s}"); + assert!(s.contains("layers.12.router.weight"), "{s}"); + } + + #[test] + fn a_model_global_tensor_omits_the_layer() { + let c = ComponentCoordinate::tensor(WeightClass::ControlAndRouter, None, "embeddings"); + assert!(!c.describe().contains("layer"), "{}", c.describe()); + } + + #[test] + fn only_bank_regions_have_a_role() { + let region = + ComponentCoordinate::BankRegion(RegionCoordinate::new(0, 0, None, RegionRole::Gate)); + assert_eq!(region.role(), Some(RegionRole::Gate)); + assert_eq!(tensor_coordinate().role(), None); + } + + #[test] + fn a_present_tensor_is_usable() { + let t = SelectedTensor::present(tensor_coordinate(), Fidelity::SourceExact); + assert!(t.is_usable()); + assert_eq!(t.reason_unusable(), None); + assert!(!t.indicates_defect()); + } + + #[test] + fn a_deliberately_omitted_tensor_is_unusable_but_not_a_defect() { + // An attention-only client slice drops the LM head on purpose. + let t = SelectedTensor { + omitted: Some(AbsenceKind::OmittedBySelection), + ..SelectedTensor::present(tensor_coordinate(), Fidelity::SourceExact) + }; + assert!(!t.is_usable()); + assert!(!t.indicates_defect()); + assert!(t.reason_unusable().unwrap().contains("omitted")); + } + + #[test] + fn an_undeclared_absence_is_a_defect() { + let t = SelectedTensor { + omitted: Some(AbsenceKind::AbsentEverywhere), + ..SelectedTensor::present(tensor_coordinate(), Fidelity::SourceExact) + }; + assert!(!t.is_usable()); + assert!(t.indicates_defect()); + } + + #[test] + fn an_unreadable_tensor_is_unusable_without_being_a_defect() { + // The artifact may simply be newer than this build. + let t = SelectedTensor { + support: ReferenceSupport::UnsupportedFormat(RegionFormat::Nvfp4), + ..SelectedTensor::present(tensor_coordinate(), Fidelity::SourceExact) + }; + assert!(!t.is_usable()); + assert!(!t.indicates_defect()); + assert!(t.reason_unusable().unwrap().contains("nvfp4")); + } + + #[test] + fn both_component_kinds_contribute_fidelity_to_the_fold() { + let region = SelectedComponent::region( + RegionCoordinate::new(0, 0, None, RegionRole::Gate), + Fidelity::SourceExact, + ); + let tensor = SelectedComponent::ManifestTensor(SelectedTensor::present( + tensor_coordinate(), + Fidelity::NumericallyApproximate, + )); + assert_eq!(region.fidelity(), Fidelity::SourceExact); + assert_eq!(tensor.fidelity(), Fidelity::NumericallyApproximate); + } + + #[test] + fn weight_classes_name_the_five_durable_classes() { + let names: Vec<&str> = [ + WeightClass::ControlAndRouter, + WeightClass::DenseSpine, + WeightClass::SharedFfn, + WeightClass::RoutedGateUp, + WeightClass::RoutedDown, + ] + .iter() + .map(|c| c.name()) + .collect(); + assert_eq!( + names, + vec![ + "control", + "dense", + "shared", + "routed_gate_up", + "routed_down" + ] + ); + } + + #[test] + fn coordinates_sort_bank_regions_before_manifest_tensors() { + // Stable ordering keeps reports diffable across runs. + let mut v = [ + tensor_coordinate(), + ComponentCoordinate::BankRegion(RegionCoordinate::new(0, 0, None, RegionRole::Gate)), + ]; + v.sort(); + assert!(matches!(v[0], ComponentCoordinate::BankRegion(_))); + } +} diff --git a/crates/larql-vindex/src/format/capability/coordinate.rs b/crates/larql-vindex/src/format/capability/coordinate.rs new file mode 100644 index 000000000..ad87f1f25 --- /dev/null +++ b/crates/larql-vindex/src/format/capability/coordinate.rs @@ -0,0 +1,222 @@ +//! Exact coordinates of a region within an index (spec §11). +//! +//! V2-0's acceptance contract requires missing operands to be diagnosed with +//! `{layer, bank, role, segment}` precision. "Some segment is missing `down`" +//! is not that: on a two-segment K3 routed layer it leaves the reader to +//! bisect 896 experts to find which half is broken. +//! +//! Segment identity is kept **individual** in the report even when several +//! adjacent segments fail. Presentation may compact a run into `segments 1–4`; +//! the report itself must not, because compaction is lossy and a consumer that +//! wants to re-fetch exactly the broken segments needs the list. + +use crate::format::lyrw2::region_role::RegionRole; + +/// Where a region lives, precisely enough to act on. +#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub struct RegionCoordinate { + pub layer: u32, + pub bank_id: u16, + /// Segment index within the bank. `None` for a bank the selection treats + /// as unsegmented — distinct from segment 0, which is one segment of many. + pub segment: Option, + pub role: RegionRole, +} + +impl RegionCoordinate { + pub fn new(layer: u32, bank_id: u16, segment: Option, role: RegionRole) -> Self { + Self { + layer, + bank_id, + segment, + role, + } + } + + /// Diagnostic form, in the order §11 names: layer, bank, role, segment. + pub fn describe(&self) -> String { + let seg = match self.segment { + Some(s) => format!(" segment {s}"), + None => String::new(), + }; + format!( + "layer {} bank {} role {}{}", + self.layer, + self.bank_id, + self.role.name(), + seg + ) + } +} + +/// A bank, without a role or segment — the unit a plan choice is made over. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub struct BankCoordinate { + pub layer: u32, + pub bank_id: u16, +} + +impl BankCoordinate { + pub fn new(layer: u32, bank_id: u16) -> Self { + Self { layer, bank_id } + } + + pub fn describe(&self) -> String { + format!("layer {} bank {}", self.layer, self.bank_id) + } +} + +/// Why a required region is not usable — **role-local, whole-role causes only**. +/// +/// Partial segment coverage is deliberately absent. A role covering some of +/// the required population has `C ≠ ∅` and is therefore *usable*; whether the +/// alternative can run is then a question about how the roles' coverage sets +/// relate, which `compatibility::SegmentCompatibility` answers. Recording it +/// here as well would state the same fact at two levels and invite them to +/// disagree. +/// +/// Cross-role segment incompatibility deliberately does *not* live here. It is +/// a relational fact about two roles' coverage failing to form one executable +/// population, so assigning it to each operand separately would state a +/// symptom twice and the cause nowhere. It lives on the alternative +/// evaluation instead (see `compatibility::SegmentCompatibility`). +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum AbsenceKind { + /// No segment carries this role. The variant was never extracted, or the + /// slice dropped it. Fix: extract or fetch the variant. + AbsentEverywhere, + /// The active profile or slice deliberately omits it. Not a defect — an + /// analysis-only browse slice has no `down` by design. Fix: none; use a + /// profile that selects it. + OmittedBySelection, + /// Regions exist, but none inside the required population. The bytes are + /// real and unreachable for *this* selection, which is a resolution fault + /// rather than a storage one. Fix: reconcile the selection's segment set + /// with what the bank actually holds. + PresentOutsidePopulation { found: Vec, required: Vec }, +} + +impl AbsenceKind { + /// Whether this absence indicates something is wrong with the index, as + /// opposed to a deliberate scoping decision. + /// + /// A browse slice missing `down` is working as designed; reporting it as + /// corruption would teach operators to ignore the diagnostic. + pub fn is_defect(&self) -> bool { + !matches!(self, Self::OmittedBySelection) + } + + pub fn describe(&self) -> String { + match self { + Self::AbsentEverywhere => "absent from every selected segment".into(), + Self::OmittedBySelection => "omitted by the active selection".into(), + Self::PresentOutsidePopulation { found, required } => format!( + "present in segments {}, none of which are in the required set {}", + render(found), + render(required) + ), + } + } +} + +fn render(segments: &[u16]) -> String { + if segments.is_empty() { + return "none".into(); + } + segments + .iter() + .map(|s| s.to_string()) + .collect::>() + .join(", ") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn a_coordinate_describes_every_axis_section_eleven_names() { + let c = RegionCoordinate::new(37, 0, Some(1), RegionRole::Down); + let s = c.describe(); + assert!(s.contains("layer 37"), "{s}"); + assert!(s.contains("bank 0"), "{s}"); + assert!(s.contains("role down"), "{s}"); + assert!(s.contains("segment 1"), "{s}"); + } + + #[test] + fn an_unsegmented_coordinate_omits_the_segment_axis() { + let c = RegionCoordinate::new(3, 1, None, RegionRole::Gate); + assert!(!c.describe().contains("segment"), "{}", c.describe()); + } + + #[test] + fn segment_zero_is_distinct_from_unsegmented() { + // Segment 0 is one segment of several; None means the selection treats + // the bank as a whole. Conflating them loses which half is broken. + let zero = RegionCoordinate::new(0, 0, Some(0), RegionRole::Down); + let none = RegionCoordinate::new(0, 0, None, RegionRole::Down); + assert_ne!(zero, none); + assert!(zero.describe().contains("segment 0")); + } + + #[test] + fn coordinates_sort_by_layer_then_bank_then_segment() { + let mut v = [ + RegionCoordinate::new(1, 0, Some(1), RegionRole::Down), + RegionCoordinate::new(0, 0, Some(0), RegionRole::Down), + RegionCoordinate::new(1, 0, Some(0), RegionRole::Down), + ]; + v.sort(); + assert_eq!(v[0].layer, 0); + assert_eq!(v[1].segment, Some(0)); + assert_eq!(v[2].segment, Some(1)); + } + + #[test] + fn a_bank_coordinate_names_layer_and_bank_only() { + let b = BankCoordinate::new(37, 2); + assert_eq!(b.describe(), "layer 37 bank 2"); + } + + #[test] + fn bank_coordinates_sort_by_layer_then_bank() { + let mut v = [ + BankCoordinate::new(1, 0), + BankCoordinate::new(0, 5), + BankCoordinate::new(1, 1), + ]; + v.sort(); + assert_eq!(v[0], BankCoordinate::new(0, 5)); + assert_eq!(v[2], BankCoordinate::new(1, 1)); + } + + #[test] + fn absent_everywhere_is_a_defect() { + assert!(AbsenceKind::AbsentEverywhere.is_defect()); + assert!(AbsenceKind::AbsentEverywhere + .describe() + .contains("every selected segment")); + } + + #[test] + fn regions_outside_the_population_are_a_resolution_fault_not_a_storage_one() { + // The bytes are real; this selection simply cannot reach them. That is + // a different repair from "extract the variant". + let a = AbsenceKind::PresentOutsidePopulation { + found: vec![7, 8], + required: vec![0, 1], + }; + assert!(a.is_defect()); + let s = a.describe(); + assert!(s.contains("present in segments 7, 8"), "{s}"); + assert!(s.contains("required set 0, 1"), "{s}"); + } + + #[test] + fn a_deliberate_omission_is_not_a_defect() { + // A browse slice has no `down` by design. Reporting that as corruption + // teaches operators to ignore the diagnostic. + assert!(!AbsenceKind::OmittedBySelection.is_defect()); + } +} diff --git a/crates/larql-vindex/src/format/capability/decode_requirements.rs b/crates/larql-vindex/src/format/capability/decode_requirements.rs new file mode 100644 index 000000000..72ede4bfa --- /dev/null +++ b/crates/larql-vindex/src/format/capability/decode_requirements.rs @@ -0,0 +1,198 @@ +//! What an architecture needs for a complete forward path (spec §4, §8.2). +//! +//! Adapter-owned, deliberately. VINDEX3 is explicitly not a general +//! neural-graph container (§2), so a model's dense spine — KDA recurrence +//! parameters, MLA latent KV, AttnRes, output gates — is *declared* here +//! rather than described in the MoE programme vocabulary. +//! +//! # Requirements do not restate the programme +//! +//! An adapter names the banks a layer uses and the non-bank tensors that +//! surround them. It never restates `gate/up/down` alternatives: those belong +//! to the programme registry, and duplicating them would mean a programme +//! change required synchronised edits to every adapter. +//! +//! ```text +//! adapter which banks participate, and what surrounds them +//! programme registry which bank-region arrangements execute +//! local decode joins the two +//! ``` +//! +//! # A complete bank is not a decodable layer +//! +//! Programme traversal answers whether the *bank computation* runs. These +//! requirements answer whether the *whole layer path* does. Mini-K3 separates +//! them immediately: `gate/up/down` can be perfect while `routed_input` is +//! absent, and then WALK and decode fail for different reasons through +//! different dependency projections. +//! +//! # The dense schedule is a list, not a field +//! +//! No `first_k_dense_replace`, no `dense_mlp_idx`. An adapter emits per-layer +//! requirements, so Kimi-Linear's leading dense layer and Inkling-Small's +//! mid-stack one need no special cases — which is why the format chose +//! per-layer manifests in the first place. + +use super::component::{ComponentContract, ComponentCoordinate}; +use super::coordinate::BankCoordinate; + +/// What a required component is *for*, independent of where it is stored. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ComponentRequirement { + /// Human-readable purpose, used in diagnostics: "final norm", "lm_head". + pub purpose: &'static str, + pub coordinate: ComponentCoordinate, + /// Expected shape and kind, when the adapter pins one. + pub contract: Option, +} + +impl ComponentRequirement { + pub fn new(purpose: &'static str, coordinate: ComponentCoordinate) -> Self { + Self { + purpose, + coordinate, + contract: None, + } + } + + pub fn with_contract(mut self, contract: ComponentContract) -> Self { + self.contract = Some(contract); + self + } +} + +/// One requirement, possibly satisfiable more than one way. +/// +/// The fused/decomposed problem recurs outside expert regions. A tied LM head +/// is the obvious case: some models store a dedicated `lm_head`, others reuse +/// the embedding tensor. A flat list of mandatory tensors would reject a +/// perfectly valid tied-weight model, so requirements carry alternatives for +/// the same reason programmes do — and with the same rule: every satisfied +/// alternative is preserved, and declaration order decides only which failure +/// is reported when none succeed. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Requirement { + One(ComponentRequirement), + AnyOf(Vec), +} + +impl Requirement { + pub fn one(purpose: &'static str, coordinate: ComponentCoordinate) -> Self { + Self::One(ComponentRequirement::new(purpose, coordinate)) + } + + /// Candidates in declaration order. + pub fn candidates(&self) -> &[ComponentRequirement] { + match self { + Self::One(r) => std::slice::from_ref(r), + Self::AnyOf(rs) => rs, + } + } + + /// Purpose of the first candidate — what the requirement is *called*, even + /// when several tensors could satisfy it. + pub fn purpose(&self) -> &'static str { + self.candidates() + .first() + .map(|r| r.purpose) + .unwrap_or("unnamed requirement") + } +} + +/// The non-bank components surrounding one MoE layer. +#[derive(Debug, Clone, PartialEq, Eq, Default)] +pub struct MoeLayerRequirements { + /// Router score weights, bias, and anything else the manifest names for + /// selection. A complete expert bank without these is not a decodable + /// layer. + pub router: Vec, + /// Residual↔latent projections, routed output norm, shared pre/post + /// transforms. + pub transforms: Vec, + /// Attention / norms / recurrence for this layer. + pub fixed_spine: Vec, + /// Banks whose executable arrangements come from programme traversal. + pub banks: Vec, +} + +/// What one layer needs, by kind. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum LayerDecodeRequirements { + Dense { components: Vec }, + Moe(MoeLayerRequirements), +} + +impl LayerDecodeRequirements { + /// Every non-bank requirement this layer imposes, in a stable order. + pub fn fixed(&self) -> Vec<&Requirement> { + match self { + Self::Dense { components } => components.iter().collect(), + Self::Moe(m) => m + .fixed_spine + .iter() + .chain(&m.router) + .chain(&m.transforms) + .collect(), + } + } + + pub fn banks(&self) -> &[BankCoordinate] { + match self { + Self::Dense { .. } => &[], + Self::Moe(m) => &m.banks, + } + } +} + +/// The execution boundary being requested. +/// +/// Deliberately one variant. Other boundaries — residual-to-logits, a layer +/// range — change which fixed components are required, so each deserves its +/// own explicit variant rather than a general request with optional fields +/// whose combinations become ambiguous. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum LocalDecodeRequest { + /// Token ids in, logits out. + /// + /// Needs embeddings, every active layer, the final norm, and an LM head or + /// its tied equivalent. Does **not** need a tokenizer: the ids are already + /// supplied, and requiring one would refuse a caller that never needed it. + TokenIdsToLogits, +} + +/// Everything an architecture needs for one decode boundary. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DecodeRequirements { + /// Model-global components — embeddings, final norm, LM head. + pub fixed_components: Vec, + /// One entry per active layer, in order. Dense and MoE layers interleave + /// freely; no global schedule field exists to get wrong. + pub layers: Vec, +} + +impl DecodeRequirements { + /// Every bank any layer requires. + pub fn all_banks(&self) -> Vec { + let mut out: Vec = self + .layers + .iter() + .flat_map(|l| l.banks()) + .copied() + .collect(); + out.sort(); + out.dedup(); + out + } + + /// Layers that run a MoE programme. + pub fn moe_layer_count(&self) -> usize { + self.layers + .iter() + .filter(|l| matches!(l, LayerDecodeRequirements::Moe(_))) + .count() + } + + pub fn dense_layer_count(&self) -> usize { + self.layers.len() - self.moe_layer_count() + } +} diff --git a/crates/larql-vindex/src/format/capability/kernel.rs b/crates/larql-vindex/src/format/capability/kernel.rs new file mode 100644 index 000000000..cfa07aa59 --- /dev/null +++ b/crates/larql-vindex/src/format/capability/kernel.rs @@ -0,0 +1,95 @@ +//! Kernel identity and maturity (spec §10). +//! +//! Deliberately outside the traversal's vocabulary. Traversal answers "are the +//! selected bytes sufficient for this programme?" and stops at reference +//! executability; the kernel registry separately answers "what optimised +//! implementation can execute them?". Keeping maturity in a different module +//! from operand availability is what stops the second question leaking into +//! the first — and a traversal that could name a kernel would be a traversal +//! that had already chosen one. + +/// How mature the execution path for an operand is (§10's ladder). +/// +/// Deliberately excludes `Representable` and `Reference`: those are not kernel +/// bindings, and are carried by [`OperandCapability`] variants instead. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub enum KernelMaturity { + Grouped, + Dispatched, + Production, +} + +impl KernelMaturity { + pub const fn name(self) -> &'static str { + match self { + Self::Grouped => "grouped", + Self::Dispatched => "dispatched", + Self::Production => "production", + } + } +} + +/// Opaque handle to a registered kernel. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct KernelId(pub String); + +impl KernelId { + pub fn new(name: impl Into) -> Self { + Self(name.into()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn maturity_orders_from_grouped_to_production() { + assert!(KernelMaturity::Grouped < KernelMaturity::Dispatched); + assert!(KernelMaturity::Dispatched < KernelMaturity::Production); + } + + #[test] + fn the_ladder_excludes_the_non_kernel_rungs() { + // Representable and Reference are not kernel bindings. They are carried + // by the traversal's own vocabulary, so a kernel maturity can never be + // used to mean "no kernel". + let names: Vec<&str> = [ + KernelMaturity::Grouped, + KernelMaturity::Dispatched, + KernelMaturity::Production, + ] + .iter() + .map(|m| m.name()) + .collect(); + assert_eq!(names, vec!["grouped", "dispatched", "production"]); + assert!(!names.contains(&"reference")); + assert!(!names.contains(&"representable")); + } + + #[test] + fn maturity_sorts_so_the_best_binding_is_the_maximum() { + // Kernel binding ranks candidates; `max` must mean "most mature". + let mut v = [ + KernelMaturity::Production, + KernelMaturity::Grouped, + KernelMaturity::Dispatched, + ]; + v.sort(); + assert_eq!(v.last(), Some(&KernelMaturity::Production)); + } + + #[test] + fn kernel_ids_are_distinct_by_name() { + assert_eq!(KernelId::new("a"), KernelId::new("a")); + assert_ne!(KernelId::new("a"), KernelId::new("b")); + } + + #[test] + fn a_kernel_id_carries_its_name_verbatim() { + assert_eq!( + KernelId::new("grouped_experts_q6k").0, + "grouped_experts_q6k" + ); + } +} diff --git a/crates/larql-vindex/src/format/capability/local_decode.rs b/crates/larql-vindex/src/format/capability/local_decode.rs new file mode 100644 index 000000000..c11bc65db --- /dev/null +++ b/crates/larql-vindex/src/format/capability/local_decode.rs @@ -0,0 +1,212 @@ +//! Local decode inference (spec §11). +//! +//! Assembly, by design. Every hard question was answered upstream: +//! +//! ```text +//! adapter which tensors and banks a layer needs +//! programme traversal which bank-region arrangements execute +//! this function joins them and folds authority +//! ``` +//! +//! It never consults kernel maturity — a Production kernel and the reference +//! path over the same bytes have identical fidelity, and admission does not +//! depend on either. +//! +//! # Failures aggregate +//! +//! A missing router at layer 12 and an unreadable norm at layer 19 are +//! independent facts, and reporting only the first would make fixing an index +//! an iterative guessing game. Resolution collects every failure whose +//! diagnosis does not depend on interpreting another. + +use std::collections::BTreeMap; + +use super::component::{SelectedComponent, SelectedTensor}; +use super::coordinate::BankCoordinate; +use super::decode_requirements::{DecodeRequirements, LocalDecodeRequest, Requirement}; +use super::operation::{OperationCapability, OperationFailure}; +use super::plan::{OperationPlan, PlanChoice, QualifiedAlternative}; +use super::traversal::BankCapabilityReport; + +/// The physical components a profile resolved to. +#[derive(Debug, Clone, PartialEq, Eq, Default)] +pub struct ResolvedDocumentSelection { + /// Manifest-addressed tensors, keyed by their coordinate's description so + /// lookup is by identity rather than by position. + pub tensors: Vec, + /// Per-bank traversal results, already qualified. + pub bank_reports: BTreeMap, +} + +impl ResolvedDocumentSelection { + fn tensor(&self, coordinate_description: &str) -> Option<&SelectedTensor> { + self.tensors + .iter() + .find(|t| t.coordinate.describe() == coordinate_description) + } +} + +/// Why one candidate of a requirement could not be used. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CandidateFailure { + pub coordinate: String, + pub reason: String, +} + +/// Every candidate outcome for one requirement. +/// +/// Failed candidates are kept even when another succeeds. They are evidence: +/// an operator debugging why a tied head was used instead of a dedicated one +/// needs to see that the dedicated one was absent. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RequirementResolution { + pub purpose: &'static str, + pub usable: Vec, + pub failed: Vec, +} + +impl RequirementResolution { + /// Diagnostic naming every candidate's own repair, not just the first. + /// + /// "lm_head unsatisfied" is useless; "dedicated head absent, embedding + /// reuse mismatched after transpose" tells an operator which to fix. + pub fn describe_failure(&self) -> String { + let candidates = self + .failed + .iter() + .map(|f| format!("{} → {}", f.coordinate, f.reason)) + .collect::>() + .join("; "); + format!("{} unsatisfied: {candidates}", self.purpose) + } +} + +/// Evaluate every declared candidate, partitioning usable from failed. +/// +/// Deliberately does **not** return on the first success. A document holding +/// both a dedicated head and a usable embedding table has two valid routes, +/// and collapsing to one here would settle the route — and therefore the +/// authority and the kernel choice — before binding has seen the registry. +/// That is the same collapse already fixed for expert alternatives. +fn resolve_requirement( + selection: &ResolvedDocumentSelection, + requirement: &Requirement, +) -> RequirementResolution { + let mut usable = Vec::new(); + let mut failed = Vec::new(); + + for candidate in requirement.candidates() { + let key = candidate.coordinate.describe(); + let Some(tensor) = selection.tensor(&key) else { + failed.push(CandidateFailure { + coordinate: key, + reason: "absent from the selection".into(), + }); + continue; + }; + let usability = tensor.usability(candidate.contract.as_ref()); + if usability.is_usable() { + usable.push(SelectedComponent::ManifestTensor(tensor.clone())); + } else { + failed.push(CandidateFailure { + coordinate: key, + reason: usability.describe(), + }); + } + } + + RequirementResolution { + purpose: requirement.purpose(), + usable, + failed, + } +} + +/// Infer whether the selection can run a complete forward pass locally. +pub fn infer_local_decode( + selection: &ResolvedDocumentSelection, + requirements: &DecodeRequirements, + _request: &LocalDecodeRequest, +) -> OperationCapability { + let mut fixed_components = Vec::new(); + let mut choices = Vec::new(); + let mut failures = Vec::new(); + + // Model-global first, then each layer's fixed path. Collected across every + // layer before returning, so one report can name layer 12's missing router + // and layer 19's unreadable norm together. + let every_requirement = requirements + .fixed_components + .iter() + .chain(requirements.layers.iter().flat_map(|l| l.fixed())); + + for (index, requirement) in every_requirement.enumerate() { + let resolution = resolve_requirement(selection, requirement); + match resolution.usable.len() { + // Nothing satisfies it — report every candidate's own repair. + 0 => failures.push(OperationFailure::InvalidSelection { + detail: resolution.describe_failure(), + }), + // Exactly one route: a fixed binding, nothing left to choose. + 1 => fixed_components.push(resolution.usable[0].clone()), + // Several routes. Binding picks; traversal must not. + _ => choices.push(PlanChoice { + bank: BankCoordinate::new(u32::MAX, index as u16), + alternatives: resolution + .usable + .iter() + .map(|c| QualifiedAlternative { + alternative: &[], + regions: Vec::new(), + components: vec![c.clone()], + }) + .collect(), + }), + } + } + + // Expert banks, from traversal. + for bank in requirements.all_banks() { + match selection.bank_reports.get(&bank) { + Some(report) if report.is_executable() => choices.push(PlanChoice { + bank, + alternatives: report + .successful_alternatives() + .iter() + .map(|a| QualifiedAlternative { + alternative: a.alternative, + regions: Vec::new(), + components: Vec::new(), + }) + .collect(), + }), + Some(report) => { + let detail = report + .closest_failure() + .map(|a| a.reference_execution.describe()) + .unwrap_or_else(|| "no executable alternative".into()); + failures.push(OperationFailure::RequiredRegionUnusable { + coordinate: super::coordinate::RegionCoordinate::new( + bank.layer, + bank.bank_id, + None, + crate::format::lyrw2::region_role::RegionRole::Down, + ), + cause: detail, + }); + } + None => failures.push(OperationFailure::NoExecutableRoute { layer: bank.layer }), + } + } + + if !failures.is_empty() { + return OperationCapability::unavailable(failures); + } + + OperationCapability::available(vec![super::plan::QualifiedOperationRoute::new( + OperationPlan { + fixed_components, + choices, + }, + )]) +} diff --git a/crates/larql-vindex/src/format/capability/local_decode_tests.rs b/crates/larql-vindex/src/format/capability/local_decode_tests.rs new file mode 100644 index 000000000..7ce81ed5e --- /dev/null +++ b/crates/larql-vindex/src/format/capability/local_decode_tests.rs @@ -0,0 +1,476 @@ +//! Colocated tests for `local_decode`. +//! +//! The two decisive cases are inverses of each other, and together they prove +//! neither side pretends to validate the other: +//! +//! ```text +//! experts complete, routed_input absent → bank executable, decode fails +//! fixed path complete, down unreadable → fixed resolves, decode fails +//! ``` + +use std::collections::BTreeMap; + +use crate::format::lyrw2::region_format::{Packing, RegionFormat}; +use crate::format::lyrw2::region_role::RegionRole; +use crate::format::moe_manifest::programme::Programme; + +use super::authority::Fidelity; +use super::component::{ComponentContract, ComponentCoordinate, SelectedTensor, WeightClass}; +use super::coordinate::{AbsenceKind, BankCoordinate}; +use super::decode_requirements::{ + DecodeRequirements, LayerDecodeRequirements, LocalDecodeRequest, MoeLayerRequirements, + Requirement, +}; +use super::local_decode::{infer_local_decode, ResolvedDocumentSelection}; +use super::role::ReferenceSupport; +use super::selection::{BankSelection, SelectedRegion}; +use super::traversal::{traverse_bank, BankCapabilityReport, ReferenceCodecs}; + +const LAYER: u32 = 12; +const BANK: BankCoordinate = BankCoordinate { + layer: LAYER, + bank_id: 0, +}; + +struct AllCodecs; +impl ReferenceCodecs for AllCodecs { + fn supports(&self, _: RegionFormat) -> bool { + true + } + fn supports_packing(&self, _: Packing) -> bool { + true + } +} + +struct NoNvfp4; +impl ReferenceCodecs for NoNvfp4 { + fn supports(&self, f: RegionFormat) -> bool { + f != RegionFormat::Nvfp4 + } + fn supports_packing(&self, _: Packing) -> bool { + true + } +} + +fn key(name: &str) -> ComponentCoordinate { + ComponentCoordinate::tensor(WeightClass::ControlAndRouter, Some(LAYER), name) +} + +fn global(name: &str) -> ComponentCoordinate { + ComponentCoordinate::tensor(WeightClass::ControlAndRouter, None, name) +} + +fn tensor(coordinate: ComponentCoordinate, fidelity: Fidelity) -> SelectedTensor { + SelectedTensor::present(coordinate, fidelity) +} + +/// A K3-shaped MoE layer: router, both latent transforms, one expert bank. +fn k3_requirements() -> DecodeRequirements { + DecodeRequirements { + fixed_components: vec![ + Requirement::one("embeddings", global("embeddings")), + Requirement::one("final norm", global("final_norm")), + Requirement::AnyOf(vec![ + super::decode_requirements::ComponentRequirement::new("lm_head", global("lm_head")), + // Tied-weight models reuse the embedding table. + super::decode_requirements::ComponentRequirement::new( + "lm_head", + global("embeddings"), + ), + ]), + ], + layers: vec![LayerDecodeRequirements::Moe(MoeLayerRequirements { + router: vec![Requirement::one("router weights", key("router.weight"))], + transforms: vec![ + Requirement::one("routed_input", key("routed_expert_down_proj")), + Requirement::one("routed_output", key("routed_expert_up_proj")), + ], + fixed_spine: vec![Requirement::one("attention", key("attn.qkv"))], + banks: vec![BANK], + })], + } +} + +fn bank_selection(down_format: RegionFormat) -> BankSelection { + let mut regions = BTreeMap::new(); + regions.insert( + (None, RegionRole::GateUpFused), + SelectedRegion { + format: RegionFormat::Q6K, + packing: Packing::RowMajor, + fidelity: Fidelity::SourceEquivalent, + }, + ); + regions.insert( + (None, RegionRole::Down), + SelectedRegion { + format: down_format, + packing: Packing::RowMajor, + fidelity: Fidelity::SourceEquivalent, + }, + ); + BankSelection { + bank_id: 0, + required_segments: Vec::new(), + regions, + omitted_roles: Vec::new(), + } +} + +fn bank_report(down_format: RegionFormat, codecs: &dyn ReferenceCodecs) -> BankCapabilityReport { + traverse_bank( + LAYER, + Programme::GatedMlpV1, + &bank_selection(down_format), + codecs, + ) +} + +/// Every tensor the K3 requirements name, all present and exact. +fn full_tensors() -> Vec { + vec![ + tensor(global("embeddings"), Fidelity::SourceExact), + tensor(global("final_norm"), Fidelity::SourceExact), + tensor(global("lm_head"), Fidelity::SourceExact), + tensor(key("router.weight"), Fidelity::SourceExact), + tensor(key("routed_expert_down_proj"), Fidelity::SourceExact), + tensor(key("routed_expert_up_proj"), Fidelity::SourceExact), + tensor(key("attn.qkv"), Fidelity::SourceExact), + ] +} + +fn selection( + tensors: Vec, + report: BankCapabilityReport, +) -> ResolvedDocumentSelection { + let mut bank_reports = BTreeMap::new(); + bank_reports.insert(BANK, report); + ResolvedDocumentSelection { + tensors, + bank_reports, + } +} + +fn infer(sel: &ResolvedDocumentSelection) -> super::operation::OperationCapability { + infer_local_decode( + sel, + &k3_requirements(), + &LocalDecodeRequest::TokenIdsToLogits, + ) +} + +#[test] +fn a_complete_k3_shaped_layer_decodes() { + let sel = selection(full_tensors(), bank_report(RegionFormat::Q6K, &AllCodecs)); + let c = infer(&sel); + assert!(c.is_available(), "{}", c.admission.describe()); + // Two choice groups: the expert bank, and the lm_head requirement — both + // its candidates resolve here, so binding still has that decision. + assert_eq!(c.routes[0].plan.choices.len(), 2); +} + +// ── The two decisive inverse cases ───────────────────────────────────────── + +#[test] +fn a_complete_bank_with_no_routed_input_fails_decode_while_the_bank_still_executes() { + // A complete expert bank is not a decodable layer. Traversal says the bank + // computation runs; local decode says the whole layer path does not. + let report = bank_report(RegionFormat::Q6K, &AllCodecs); + assert!(report.is_executable(), "bank traversal must still succeed"); + + let without_transform: Vec = full_tensors() + .into_iter() + .filter(|t| !t.coordinate.describe().contains("routed_expert_down_proj")) + .collect(); + let c = infer(&selection(without_transform, report)); + + assert!(!c.is_available()); + assert!( + c.admission.describe().contains("routed_input"), + "{}", + c.admission.describe() + ); +} + +#[test] +fn a_complete_fixed_path_with_an_unreadable_down_fails_through_bank_traversal() { + // The inverse. Every tensor resolves; the bank does not execute. + let report = bank_report(RegionFormat::Nvfp4, &NoNvfp4); + assert!(!report.is_executable()); + + let c = infer(&selection(full_tensors(), report)); + assert!(!c.is_available()); + let text = c.admission.describe(); + assert!(text.contains("layer 12"), "{text}"); + // The failure is about the bank, not about a missing fixed tensor. + assert!(!text.contains("routed_input"), "{text}"); +} + +// ── Fixed-spine failures do not touch the bank ───────────────────────────── + +#[test] +fn a_missing_router_fails_decode_though_the_bank_is_complete() { + let report = bank_report(RegionFormat::Q6K, &AllCodecs); + assert!(report.is_executable()); + + let without_router: Vec = full_tensors() + .into_iter() + .filter(|t| !t.coordinate.describe().contains("router.weight")) + .collect(); + let c = infer(&selection(without_router, report)); + assert!(!c.is_available()); + assert!(c.admission.describe().contains("router weights")); +} + +#[test] +fn weakening_the_final_norm_weakens_decode_authority() { + let mut tensors = full_tensors(); + for t in &mut tensors { + if t.coordinate.describe().contains("final_norm") { + t.fidelity = Fidelity::NumericallyApproximate; + } + } + let c = infer(&selection( + tensors, + bank_report(RegionFormat::Q6K, &AllCodecs), + )); + assert_eq!( + c.best_achievable_authority(), + Some(Fidelity::NumericallyApproximate) + ); +} + +// ── Requirement alternatives ─────────────────────────────────────────────── + +#[test] +fn a_tied_lm_head_satisfies_the_requirement_through_the_embedding_tensor() { + // A flat mandatory list would reject a valid tied-weight model. + let tied: Vec = full_tensors() + .into_iter() + .filter(|t| !t.coordinate.describe().contains("lm_head")) + .collect(); + let c = infer(&selection(tied, bank_report(RegionFormat::Q6K, &AllCodecs))); + assert!(c.is_available(), "{}", c.admission.describe()); +} + +#[test] +fn losing_every_lm_head_candidate_names_each_candidates_own_repair() { + let neither: Vec = full_tensors() + .into_iter() + .filter(|t| { + let d = t.coordinate.describe(); + !d.contains("lm_head") && !d.contains("embeddings") + }) + .collect(); + let c = infer(&selection( + neither, + bank_report(RegionFormat::Q6K, &AllCodecs), + )); + assert!(!c.is_available()); + assert!(c.admission.describe().contains("lm_head")); +} + +// ── Provenance and contract ──────────────────────────────────────────────── + +#[test] +fn a_deliberately_omitted_tensor_still_blocks_decode_without_being_a_defect() { + let mut tensors = full_tensors(); + for t in &mut tensors { + if t.coordinate.describe().contains("attn.qkv") { + t.omitted = Some(AbsenceKind::OmittedBySelection); + } + } + let omitted = tensors + .iter() + .find(|t| t.coordinate.describe().contains("attn.qkv")) + .unwrap() + .clone(); + assert!( + !omitted.indicates_defect(), + "an FFN-remote client omits attention on purpose" + ); + + let c = infer(&selection( + tensors, + bank_report(RegionFormat::Q6K, &AllCodecs), + )); + assert!(!c.is_available(), "it is still required for a local pass"); +} + +#[test] +fn a_readable_tensor_of_the_wrong_shape_is_a_contract_mismatch() { + // Not an unsupported build and not a missing file — an invalid index or an + // adapter mismatch, and one that could otherwise survive far down a + // generic buffer path. + let mut requirements = k3_requirements(); + requirements.fixed_components[1] = Requirement::One( + super::decode_requirements::ComponentRequirement::new("final norm", global("final_norm")) + .with_contract(ComponentContract::vector(4_096)), + ); + + let mut tensors = full_tensors(); + for t in &mut tensors { + if t.coordinate.describe().contains("final_norm") { + t.contract = Some(ComponentContract::vector(2_048)); + } + } + + let sel = selection(tensors, bank_report(RegionFormat::Q6K, &AllCodecs)); + let c = infer_local_decode(&sel, &requirements, &LocalDecodeRequest::TokenIdsToLogits); + assert!(!c.is_available()); + let text = c.admission.describe(); + assert!(text.contains("adapter mismatch"), "{text}"); + assert!(text.contains("2048"), "{text}"); +} + +#[test] +fn an_unreadable_fixed_tensor_is_not_reported_as_a_defect() { + let mut tensors = full_tensors(); + for t in &mut tensors { + if t.coordinate.describe().contains("attn.qkv") { + t.support = ReferenceSupport::UnsupportedFormat(RegionFormat::Nvfp4); + } + } + let c = infer(&selection( + tensors, + bank_report(RegionFormat::Q6K, &AllCodecs), + )); + assert!(!c.is_available()); + assert!(c.admission.describe().contains("nvfp4")); +} + +// ── Aggregation and schedules ────────────────────────────────────────────── + +#[test] +fn independent_failures_are_reported_together() { + // Fixing an index should not be an iterative guessing game. + let stripped: Vec = full_tensors() + .into_iter() + .filter(|t| { + let d = t.coordinate.describe(); + !d.contains("router.weight") && !d.contains("attn.qkv") + }) + .collect(); + let c = infer(&selection( + stripped, + bank_report(RegionFormat::Q6K, &AllCodecs), + )); + let text = c.admission.describe(); + assert!(text.contains("router weights"), "{text}"); + assert!(text.contains("attention"), "{text}"); +} + +#[test] +fn a_hybrid_schedule_needs_no_global_dense_field() { + // Dense at index 0 and index 2, MoE between — Kimi's leading dense and + // Inkling's mid-stack dense through one mechanism. + let mut requirements = k3_requirements(); + let dense = LayerDecodeRequirements::Dense { + components: vec![Requirement::one("dense ffn", key("mlp.gate_up"))], + }; + requirements.layers = vec![dense.clone(), requirements.layers[0].clone(), dense]; + + let mut tensors = full_tensors(); + tensors.push(tensor(key("mlp.gate_up"), Fidelity::SourceExact)); + + let sel = selection(tensors, bank_report(RegionFormat::Q6K, &AllCodecs)); + let c = infer_local_decode(&sel, &requirements, &LocalDecodeRequest::TokenIdsToLogits); + assert!(c.is_available(), "{}", c.admission.describe()); + assert_eq!(requirements.dense_layer_count(), 2); + assert_eq!(requirements.moe_layer_count(), 1); +} + +#[test] +fn banks_are_one_choice_group_each_not_a_cartesian_product() { + let sel = selection(full_tensors(), bank_report(RegionFormat::Q6K, &AllCodecs)); + let c = infer(&sel); + assert_eq!(c.routes.len(), 1, "one route, not a product of choices"); + // One group per bank plus one per multi-candidate requirement — never a + // product across them. + let bank_groups = c.routes[0] + .plan + .choices + .iter() + .filter(|g| g.bank == BANK) + .count(); + assert_eq!(bank_groups, 1); +} + +#[test] +fn both_lm_head_candidates_present_yields_a_choice_not_a_settled_binding() { + // The collapse this fix removes. A document holding a dedicated head AND + // a usable embedding table has two valid routes; returning the first would + // settle the route, and therefore the authority and the kernel choice, + // before binding has seen the registry. + let sel = selection(full_tensors(), bank_report(RegionFormat::Q6K, &AllCodecs)); + let c = infer(&sel); + let component_choices: Vec<_> = c.routes[0] + .plan + .choices + .iter() + .filter(|g| g.bank != BANK) + .collect(); + assert_eq!(component_choices.len(), 1, "the lm_head requirement"); + assert_eq!( + component_choices[0].alternatives.len(), + 2, + "both candidates preserved for binding to choose between" + ); +} + +#[test] +fn one_usable_candidate_is_a_fixed_binding_not_a_one_way_choice() { + // With only one candidate resolvable there is nothing to decide, so it + // must not appear as a degenerate choice group. The tied alternative here + // names a tensor the document does not carry. + let mut requirements = k3_requirements(); + requirements.fixed_components[2] = Requirement::AnyOf(vec![ + super::decode_requirements::ComponentRequirement::new("lm_head", global("lm_head")), + super::decode_requirements::ComponentRequirement::new("lm_head", global("absent_tie")), + ]); + + let sel = selection(full_tensors(), bank_report(RegionFormat::Q6K, &AllCodecs)); + let c = infer_local_decode(&sel, &requirements, &LocalDecodeRequest::TokenIdsToLogits); + assert!(c.is_available(), "{}", c.admission.describe()); + let component_choices = c.routes[0] + .plan + .choices + .iter() + .filter(|g| g.bank != BANK) + .count(); + assert_eq!( + component_choices, 0, + "one candidate is a binding, not a choice" + ); +} + +#[test] +fn a_failed_candidate_is_retained_as_evidence_when_another_succeeds() { + // An operator debugging why the tied head was used needs to see that the + // dedicated one was absent. + let tied_only: Vec = full_tensors() + .into_iter() + .filter(|t| !t.coordinate.describe().contains("lm_head")) + .collect(); + let sel = selection(tied_only, bank_report(RegionFormat::Q6K, &AllCodecs)); + let c = infer(&sel); + assert!(c.is_available()); + // One usable candidate → fixed binding, no choice group. + let component_choices = c.routes[0] + .plan + .choices + .iter() + .filter(|g| g.bank != BANK) + .count(); + assert_eq!(component_choices, 0); +} + +#[test] +fn a_bank_with_no_traversal_report_is_refused_rather_than_assumed() { + let sel = ResolvedDocumentSelection { + tensors: full_tensors(), + bank_reports: BTreeMap::new(), + }; + let c = infer(&sel); + assert!(!c.is_available()); +} diff --git a/crates/larql-vindex/src/format/capability/mod.rs b/crates/larql-vindex/src/format/capability/mod.rs new file mode 100644 index 000000000..cf86e3087 --- /dev/null +++ b/crates/larql-vindex/src/format/capability/mod.rs @@ -0,0 +1,98 @@ +//! Capability derivation over a resolved index (spec §10, §11). +//! +//! One traversal, many consumers. Authority derivation, operation admission +//! and kernel binding are all *readers* of the report this module produces — +//! not three systems independently inspecting the index. Three inspectors is +//! how permissive logic gets in: each one is individually reasonable, and the +//! union of their leniencies is what actually ships. +//! +//! # Pipeline position +//! +//! ```text +//! raw profile +//! ↓ inheritance / schema resolution +//! resolved requested selections +//! ↓ physical variant and segment resolution +//! concrete region selection ← everything above is INPUT here +//! ↓ programme traversal +//! capability report ← this module +//! ↓ +//! authority derivation · operation admission · kernel binding +//! ``` +//! +//! The split above the line matters. "Does the profile inherit from `exact`?" +//! and "is this variant physically present?" are answered *before* traversal; +//! "can this resolved selection decode, browse, or claim source-exact?" is +//! answered *from* it. Letting traversal answer the first pair would make the +//! derivation circular — a profile whose validity depended on capabilities +//! derived from that same profile. + +pub mod authority; +pub mod binding; +#[cfg(test)] +mod binding_tests; +pub mod compatibility; +pub mod component; +pub mod coordinate; +pub mod decode_requirements; +pub mod kernel; +pub mod local_decode; +#[cfg(test)] +mod local_decode_tests; +pub mod operation; +#[cfg(test)] +mod operation_tests; +pub mod plan; +#[cfg(test)] +mod plan_tests; +pub mod reconstruction; +#[cfg(test)] +mod reconstruction_tests; +pub mod role; +pub mod scope; +pub mod selection; +pub mod traversal; +#[cfg(test)] +mod traversal_tests; +pub mod walk; +pub mod walk_request; +#[cfg(test)] +mod walk_tests; + +pub use authority::{ + derive_authority, AuthorityInputs, DerivedAuthority, Fidelity, StructuralChange, +}; +pub use binding::{ + physical_extents, total_bytes, ComponentBinding, ComponentView, PhysicalStorageId, + RepresentationIdentity, +}; +pub use compatibility::{IncompatibilityShape, RoleSegmentCoverage, SegmentCompatibility}; +pub use component::{ComponentCoordinate, SelectedComponent, SelectedTensor, WeightClass}; +pub use coordinate::{AbsenceKind, BankCoordinate, RegionCoordinate}; +pub use decode_requirements::{ + ComponentRequirement, DecodeRequirements, LayerDecodeRequirements, LocalDecodeRequest, + MoeLayerRequirements, Requirement, +}; +pub use kernel::{KernelId, KernelMaturity}; +pub use local_decode::{infer_local_decode, ResolvedDocumentSelection}; +pub use operation::{Degradation, OperationAdmission, OperationCapability, OperationFailure}; +pub use plan::{ + OperationPlan, PlanChoice, PlannedRegion, QualifiedAlternative, QualifiedOperationRoute, +}; +pub use reconstruction::{ + CanonicalSelection, CatalogueVariant, ReconstructionFailure, RegionSetCatalogue, + TensorReconstruction, +}; +pub use role::{OperandAvailability, ReferenceSupport, RoleCapability}; +pub use scope::{CapabilityReport, DocumentCapabilities, ProfileCapabilities}; +pub use selection::{BankSelection, SelectedRegion}; +pub use traversal::{ + traverse_bank, AlternativeReport, BankCapabilityReport, ReferenceCodecs, ReferenceExecution, +}; +pub use walk::{ + infer_walk, GateAccess, ResultSetCompleteness, TargetFailure, TextQueryInputs, WalkCapability, + WalkEnvironment, WalkableBank, +}; +pub use walk_request::{ + validate_query_vector, PartialPolicy, QueryVectorFault, WalkInput, WalkRequest, WalkTarget, +}; diff --git a/crates/larql-vindex/src/format/capability/operation.rs b/crates/larql-vindex/src/format/capability/operation.rs new file mode 100644 index 000000000..c0ab239d3 --- /dev/null +++ b/crates/larql-vindex/src/format/capability/operation.rs @@ -0,0 +1,212 @@ +//! Operation admission (spec §9, §11, §15). +//! +//! Three questions, kept apart: +//! +//! ```text +//! Can this operation run? → admission (here) +//! How faithfully would it run? → authority, folded over the regions THIS +//! operation consumes +//! Which implementation runs it? → kernel binding, elsewhere +//! ``` +//! +//! Admission is inferred first and authority attached afterwards, because the +//! two answer different questions and one of them has no answer when the +//! selection is contradictory. +//! +//! # The decisive property: non-interference +//! +//! Changing a component an operation does not use must not change that +//! operation's capability. WALK reads gate rows; making every `down` region +//! unreadable must leave the WALK report byte-for-byte identical. That is what +//! proves traversal facts are *projected per operation* rather than globally +//! summarised, and it is why each operation declares a narrow dependency +//! surface rather than consulting a shared "is the index healthy" verdict. + +use super::authority::Fidelity; +use super::coordinate::RegionCoordinate; +use super::plan::QualifiedOperationRoute; + +/// Why an operation cannot run. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum OperationFailure { + /// A region this operation needs is unusable, with its exact coordinate. + RequiredRegionUnusable { + coordinate: RegionCoordinate, + cause: String, + }, + /// The selection is self-contradictory. Fails closed — never downgraded + /// into a weak authority. + InvalidSelection { detail: String }, + /// A document-level input this operation needs is absent. + MissingDocumentInput { what: &'static str }, + /// The operation needs a policy declaration the caller did not supply. + /// Notably: remote execution cannot be inferred from local absence. + MissingContract { what: &'static str }, + /// No layer offers an executable route for this operation. + NoExecutableRoute { layer: u32 }, +} + +impl OperationFailure { + pub fn describe(&self) -> String { + match self { + Self::RequiredRegionUnusable { coordinate, cause } => { + format!("{}: {cause}", coordinate.describe()) + } + Self::InvalidSelection { detail } => format!("invalid selection: {detail}"), + Self::MissingDocumentInput { what } => format!("missing document input: {what}"), + Self::MissingContract { what } => { + format!("no declared contract for {what}; it cannot be inferred") + } + Self::NoExecutableRoute { layer } => { + format!("layer {layer} has no executable route") + } + } + } +} + +/// Something absent that reduces richness without preventing the operation. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Degradation { + /// Optional query metadata is absent. Reduces label richness; never + /// affects correctness (§15.3). + MissingQueryMetadata { what: &'static str }, + /// A bank cannot be browsed, so its features are outside query reach while + /// other banks remain reachable. + BankNotBrowsable { bank_id: u16, reason: String }, +} + +impl Degradation { + pub fn describe(&self) -> String { + match self { + Self::MissingQueryMetadata { what } => { + format!("{what} absent — reduced richness, unchanged correctness") + } + Self::BankNotBrowsable { bank_id, reason } => { + format!("bank {bank_id} not browsable: {reason}") + } + } + } +} + +/// Whether an operation can run, and how completely. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum OperationAdmission { + Available, + Degraded { reasons: Vec }, + Unavailable { reasons: Vec }, +} + +impl OperationAdmission { + pub fn is_available(&self) -> bool { + matches!(self, Self::Available | Self::Degraded { .. }) + } + + /// Whether admission failed because the selection is contradictory, as + /// opposed to incomplete. Contradictions must fail closed. + pub fn is_invalid_selection(&self) -> bool { + match self { + Self::Unavailable { reasons } => reasons + .iter() + .any(|r| matches!(r, OperationFailure::InvalidSelection { .. })), + _ => false, + } + } + + pub fn describe(&self) -> String { + match self { + Self::Available => "available".into(), + Self::Degraded { reasons, .. } => format!( + "available, degraded: {}", + reasons + .iter() + .map(|r| r.describe()) + .collect::>() + .join("; ") + ), + Self::Unavailable { reasons } => format!( + "unavailable: {}", + reasons + .iter() + .map(|r| r.describe()) + .collect::>() + .join("; ") + ), + } + } +} + +/// Admission plus every route that admits it. +/// +/// Authority is deliberately **not** a field here. An operation can run more +/// than one way, the ways can differ in fidelity, and no route has been bound +/// yet — so there is no single number that is both meaningful and honest. A +/// caller wanting one must ask for a *ceiling* and be told it is achievable +/// rather than achieved. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct OperationCapability { + pub admission: OperationAdmission, + /// Empty exactly when unavailable. That equivalence is the fail-closed + /// invariant: no routes means no authority, and a contradictory selection + /// therefore cannot acquire a weak-but-valid fidelity by default. + pub routes: Vec, +} + +impl OperationCapability { + pub fn available(routes: Vec) -> Self { + Self { + admission: OperationAdmission::Available, + routes, + } + } + + pub fn degraded(routes: Vec, reasons: Vec) -> Self { + Self { + admission: OperationAdmission::Degraded { reasons }, + routes, + } + } + + pub fn unavailable(reasons: Vec) -> Self { + Self { + admission: OperationAdmission::Unavailable { reasons }, + routes: Vec::new(), + } + } + + pub fn is_available(&self) -> bool { + self.admission.is_available() + } + + /// The strongest fidelity any admitted route could achieve. + /// + /// Named *achievable* on purpose. Binding has not chosen a route, so this + /// is a ceiling; quoting it as the fidelity of an execution would describe + /// a decision nobody has made. + pub fn best_achievable_authority(&self) -> Option { + self.routes + .iter() + .map(|r| r.best_achievable_authority().level) + .max() + } + + /// The weakest fidelity binding could land on across admitted routes. + pub fn worst_achievable_authority(&self) -> Option { + self.routes + .iter() + .map(|r| r.worst_achievable_authority().level) + .min() + } + + /// Whether every admitted route carries the same settled fidelity, so the + /// ceiling is also the answer. + pub fn authority_is_settled(&self) -> bool { + !self.routes.is_empty() + && self.best_achievable_authority() == self.worst_achievable_authority() + } + + /// Internal consistency: routes exist exactly when the operation admits. + pub fn is_well_formed(&self) -> bool { + // Reads oddly but is the honest form: admitted iff routes exist. + self.is_available() != self.routes.is_empty() + } +} diff --git a/crates/larql-vindex/src/format/capability/operation_tests.rs b/crates/larql-vindex/src/format/capability/operation_tests.rs new file mode 100644 index 000000000..6c463b7ea --- /dev/null +++ b/crates/larql-vindex/src/format/capability/operation_tests.rs @@ -0,0 +1,159 @@ +//! Colocated tests for `operation` — admission verdicts and their diagnoses. +//! +//! Route-scoped authority is exercised in `plan_tests`; these cover the +//! admission verdict itself and the requirement that an operator can triage +//! from the failure text alone. + +use crate::format::lyrw2::region_role::RegionRole; + +use super::authority::Fidelity; +use super::coordinate::RegionCoordinate; +use super::operation::{Degradation, OperationAdmission, OperationCapability, OperationFailure}; +use super::plan::{OperationPlan, PlannedRegion, QualifiedOperationRoute}; + +fn coordinate() -> RegionCoordinate { + RegionCoordinate::new(3, 0, Some(1), RegionRole::Down) +} + +fn route() -> QualifiedOperationRoute { + QualifiedOperationRoute::new(OperationPlan::fixed_regions(vec![PlannedRegion::new( + coordinate(), + Fidelity::SourceExact, + )])) +} + +#[test] +fn available_and_degraded_both_admit() { + assert!(OperationCapability::available(vec![route()]).is_available()); + assert!(OperationCapability::degraded( + vec![route()], + vec![Degradation::MissingQueryMetadata { what: "labels" }] + ) + .is_available()); +} + +#[test] +fn unavailable_does_not_admit() { + let c = + OperationCapability::unavailable(vec![OperationFailure::NoExecutableRoute { layer: 3 }]); + assert!(!c.is_available()); + assert!(c.routes.is_empty()); +} + +#[test] +fn an_invalid_selection_is_distinguishable_from_ordinary_unavailability() { + let invalid = OperationCapability::unavailable(vec![OperationFailure::InvalidSelection { + detail: "gate covers 0, down covers 1".into(), + }]); + let ordinary = OperationCapability::unavailable(vec![OperationFailure::MissingDocumentInput { + what: "tokenizer", + }]); + assert!(invalid.admission.is_invalid_selection()); + assert!(!ordinary.admission.is_invalid_selection()); +} + +#[test] +fn an_available_operation_is_never_an_invalid_selection() { + assert!(!OperationAdmission::Available.is_invalid_selection()); + assert!(!OperationAdmission::Degraded { + reasons: Vec::new() + } + .is_invalid_selection()); +} + +#[test] +fn a_missing_contract_is_distinct_from_missing_bytes() { + // Remote execution cannot be inferred from local absence; the two must + // never read alike. + let contract = OperationFailure::MissingContract { + what: "routed wire contract", + }; + let bytes = OperationFailure::RequiredRegionUnusable { + coordinate: coordinate(), + cause: "absent".into(), + }; + assert!(contract.describe().contains("cannot be inferred")); + assert!(!bytes.describe().contains("cannot be inferred")); +} + +#[test] +fn failures_carry_exact_coordinates() { + let f = OperationFailure::RequiredRegionUnusable { + coordinate: coordinate(), + cause: "absent from every selected segment".into(), + }; + let s = f.describe(); + for part in ["layer 3", "bank 0", "role down", "segment 1"] { + assert!(s.contains(part), "{s} missing {part}"); + } +} + +#[test] +fn every_failure_kind_renders_distinguishably() { + let rendered: Vec = [ + OperationFailure::RequiredRegionUnusable { + coordinate: coordinate(), + cause: "absent".into(), + }, + OperationFailure::InvalidSelection { + detail: "disjoint segments".into(), + }, + OperationFailure::MissingDocumentInput { what: "embeddings" }, + OperationFailure::MissingContract { + what: "routed wire", + }, + OperationFailure::NoExecutableRoute { layer: 7 }, + ] + .iter() + .map(|f| f.describe()) + .collect(); + for i in 0..rendered.len() { + for j in (i + 1)..rendered.len() { + assert_ne!(rendered[i], rendered[j], "{i} vs {j}"); + } + } + assert!(rendered[1].contains("invalid selection")); + assert!(rendered[4].contains("layer 7")); +} + +#[test] +fn missing_query_metadata_says_correctness_is_unchanged() { + // §15.3: absent query metadata downgrades richness, never correctness. The + // wording matters — an operator must not read it as "results may be wrong". + let d = Degradation::MissingQueryMetadata { + what: "feature_labels.json", + }; + assert!( + d.describe().contains("unchanged correctness"), + "{}", + d.describe() + ); +} + +#[test] +fn a_non_browsable_bank_degrades_rather_than_failing() { + let d = Degradation::BankNotBrowsable { + bank_id: 4, + reason: "fused region cannot be strided".into(), + }; + let s = d.describe(); + assert!(s.contains("bank 4"), "{s}"); + assert!(s.contains("cannot be strided"), "{s}"); +} + +#[test] +fn admission_descriptions_state_the_verdict_first() { + assert!(OperationAdmission::Available + .describe() + .starts_with("available")); + assert!(OperationAdmission::Degraded { + reasons: vec![Degradation::MissingQueryMetadata { what: "labels" }], + } + .describe() + .starts_with("available, degraded")); + assert!(OperationAdmission::Unavailable { + reasons: vec![OperationFailure::NoExecutableRoute { layer: 0 }], + } + .describe() + .starts_with("unavailable")); +} diff --git a/crates/larql-vindex/src/format/capability/plan.rs b/crates/larql-vindex/src/format/capability/plan.rs new file mode 100644 index 000000000..05f50e7ec --- /dev/null +++ b/crates/larql-vindex/src/format/capability/plan.rs @@ -0,0 +1,238 @@ +//! Operation plans and execution routes (spec §9.2, §10, §11). +//! +//! # Authority belongs to a route, not to an operation +//! +//! An operation can often run more than one way, and the ways need not be +//! equally faithful. WALK may read a direct `gate` region *or* stride the gate +//! half of a fused one; those are different bytes and may carry different +//! fidelities. Attaching one authority to the abstract operation forces a bad +//! choice: +//! +//! - report the stronger one, and kernel binding may later pick the weaker +//! route, making the report a lie; +//! - report the weaker one, and an exact route is understated; +//! - pick a route during inference, and binding is constrained before it has +//! seen the kernel registry. +//! +//! So authority is folded per route, and the operation reports every route it +//! admits. The fidelity of an *execution* is the fidelity of the route that +//! gets bound — which has not happened yet at this stage. +//! +//! # Choice groups, not a Cartesian product +//! +//! Thirty layers each admitting two alternatives is 2³⁰ whole-model routes. +//! Enumerating them is not a representation, it is a denial of service. A plan +//! is therefore *fixed requirements plus independent choice groups*: binding +//! picks one alternative per group, and the executed authority is folded over +//! the fixed regions plus whatever it picked. +//! +//! Before binding, the honest statement about such a plan is a **range**, not +//! a value. + +use super::authority::{derive_authority, AuthorityInputs, DerivedAuthority, Fidelity}; +use super::component::{ComponentCoordinate, SelectedComponent}; +use super::coordinate::BankCoordinate; +use crate::format::moe_manifest::programme::RoleAlternative; + +/// A region an operation would read, with the fidelity that feeds the fold. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PlannedRegion { + pub coordinate: super::coordinate::RegionCoordinate, + pub fidelity: Fidelity, +} + +impl PlannedRegion { + pub fn new(coordinate: super::coordinate::RegionCoordinate, fidelity: Fidelity) -> Self { + Self { + coordinate, + fidelity, + } + } + + pub fn as_component(&self) -> SelectedComponent { + SelectedComponent::region(self.coordinate.clone(), self.fidelity) + } +} + +/// One way to satisfy a bank within a plan. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct QualifiedAlternative { + /// Bank-region layout, empty for a component choice. + pub alternative: RoleAlternative, + pub regions: Vec, + /// Manifest tensors, for a choice between component candidates such as a + /// dedicated versus tied LM head. + pub components: Vec, +} + +impl QualifiedAlternative { + /// Weakest fidelity across everything this alternative reads. + pub fn weakest_fidelity(&self) -> Option { + self.regions + .iter() + .map(|r| r.fidelity) + .chain(self.components.iter().map(|c| c.fidelity())) + .min() + } +} + +/// A bank whose satisfying alternative binding will choose. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PlanChoice { + pub bank: BankCoordinate, + /// Every alternative that works, in declaration order. Never collapsed — + /// the kernel registry may support one at a higher maturity than another. + pub alternatives: Vec, +} + +impl PlanChoice { + /// The alternative whose weakest region is strongest — the best fidelity + /// binding *could* achieve here. + /// + /// **Introspection and planning only.** This must never be used as a + /// binding tie-break: when two alternatives carry equal authority, kernel + /// binding has to stay free to choose on support and performance grounds. + /// An authority-oriented tie-break here would silently decide a question + /// that belongs to the kernel registry, and would do so on a criterion + /// that cannot distinguish the candidates anyway. + pub fn best_alternative(&self) -> Option<&QualifiedAlternative> { + self.alternatives + .iter() + .max_by_key(|a| a.weakest_fidelity().unwrap_or(Fidelity::AnalysisOnly)) + } + + /// The alternative whose weakest region is weakest — the floor binding + /// could land on. + pub fn worst_alternative(&self) -> Option<&QualifiedAlternative> { + self.alternatives + .iter() + .min_by_key(|a| a.weakest_fidelity().unwrap_or(Fidelity::AnalysisOnly)) + } +} + +/// What an operation would read, as fixed regions plus open choices. +#[derive(Debug, Clone, PartialEq, Eq, Default)] +pub struct OperationPlan { + /// Components every execution of this plan reads, whatever binding + /// chooses. Bank regions *and* manifest-addressed tensors — a decode path + /// needs embeddings, norms, routers and transforms, none of which are + /// bank regions. + pub fixed_components: Vec, + /// Independent choice groups. Empty for a plan with a single shape. + pub choices: Vec, +} + +impl OperationPlan { + /// A plan with no open choices. + pub fn fixed(components: Vec) -> Self { + Self { + fixed_components: components, + choices: Vec::new(), + } + } + + /// A plan whose fixed side is bank regions only — the WALK shape. + pub fn fixed_regions(regions: Vec) -> Self { + Self::fixed(regions.iter().map(PlannedRegion::as_component).collect()) + } + + pub fn is_empty(&self) -> bool { + self.fixed_components.is_empty() && self.choices.is_empty() + } + + /// Whether binding still has a decision to make. + pub fn has_open_choices(&self) -> bool { + self.choices.iter().any(|c| c.alternatives.len() > 1) + } + + /// Every coordinate this plan could read, for diagnostics. + pub fn all_coordinates(&self) -> Vec { + let mut out: Vec = self + .fixed_components + .iter() + .map(|c| c.coordinate()) + .collect(); + for choice in &self.choices { + for alt in &choice.alternatives { + out.extend( + alt.regions + .iter() + .map(|r| ComponentCoordinate::BankRegion(r.coordinate.clone())), + ); + } + } + out.sort(); + out.dedup(); + out + } + + fn fold( + &self, + pick: impl Fn(&PlanChoice) -> Option<&QualifiedAlternative>, + ) -> DerivedAuthority { + let mut fidelities: Vec = + self.fixed_components.iter().map(|c| c.fidelity()).collect(); + for choice in &self.choices { + if let Some(alt) = pick(choice) { + fidelities.extend(alt.regions.iter().map(|r| r.fidelity)); + fidelities.extend(alt.components.iter().map(|c| c.fidelity())); + } + } + DerivedAuthority::of(&AuthorityInputs { + selected_fidelities: fidelities, + execution_complete: true, + structural: None, + }) + } + + /// The strongest authority any binding of this plan could achieve. + /// + /// **Achievable, not achieved.** No route has been bound, so this is a + /// ceiling — quoting it as the fidelity of an execution would be a claim + /// about a decision that has not been made. + pub fn best_achievable_authority(&self) -> DerivedAuthority { + self.fold(PlanChoice::best_alternative) + } + + /// The weakest authority any binding of this plan could land on. + pub fn worst_achievable_authority(&self) -> DerivedAuthority { + self.fold(PlanChoice::worst_alternative) + } +} + +/// One way to run an operation, with the fidelity that route would carry. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct QualifiedOperationRoute { + pub plan: OperationPlan, +} + +impl QualifiedOperationRoute { + pub fn new(plan: OperationPlan) -> Self { + Self { plan } + } + + /// Ceiling for this route. Equals the floor when the plan has no open + /// choices, which is the common single-shape case. + pub fn best_achievable_authority(&self) -> DerivedAuthority { + self.plan.best_achievable_authority() + } + + pub fn worst_achievable_authority(&self) -> DerivedAuthority { + self.plan.worst_achievable_authority() + } + + /// Whether this route's fidelity is already settled — no open choices, so + /// ceiling and floor coincide and binding cannot move it. + pub fn authority_is_settled(&self) -> bool { + self.best_achievable_authority().level == self.worst_achievable_authority().level + } +} + +/// Convenience: fold a bare fidelity list, used where a route has no choices. +pub fn authority_of(fidelities: &[Fidelity]) -> Fidelity { + derive_authority(&AuthorityInputs { + selected_fidelities: fidelities.to_vec(), + execution_complete: true, + structural: None, + }) +} diff --git a/crates/larql-vindex/src/format/capability/plan_tests.rs b/crates/larql-vindex/src/format/capability/plan_tests.rs new file mode 100644 index 000000000..61a63ab0c --- /dev/null +++ b/crates/larql-vindex/src/format/capability/plan_tests.rs @@ -0,0 +1,265 @@ +//! Colocated tests for `plan` — route-scoped authority and choice groups. +//! +//! Two properties carry the design: an operation admitting several routes of +//! different fidelity must not collapse to one number, and a plan with many +//! independent choices must not be enumerated. + +use crate::format::lyrw2::region_role::RegionRole; +use crate::format::moe_manifest::programme::Programme; + +use super::authority::Fidelity; +use super::coordinate::{BankCoordinate, RegionCoordinate}; +use super::operation::{OperationCapability, OperationFailure}; +use super::plan::{ + OperationPlan, PlanChoice, PlannedRegion, QualifiedAlternative, QualifiedOperationRoute, +}; + +fn region(layer: u32, role: RegionRole, fidelity: Fidelity) -> PlannedRegion { + PlannedRegion::new(RegionCoordinate::new(layer, 0, None, role), fidelity) +} + +/// A route reading one exact gate region — the direct-browse shape. +fn direct_gate_route() -> QualifiedOperationRoute { + QualifiedOperationRoute::new(OperationPlan::fixed_regions(vec![region( + 0, + RegionRole::Gate, + Fidelity::SourceExact, + )])) +} + +/// A route striding an approximate fused region — same operation, worse bytes. +fn strided_fused_route() -> QualifiedOperationRoute { + QualifiedOperationRoute::new(OperationPlan::fixed_regions(vec![region( + 0, + RegionRole::GateUpFused, + Fidelity::NumericallyApproximate, + )])) +} + +fn alternative(roles: &'static [RegionRole], fidelity: Fidelity) -> QualifiedAlternative { + QualifiedAlternative { + alternative: roles, + regions: roles.iter().map(|r| region(0, *r, fidelity)).collect(), + components: Vec::new(), + } +} + +#[test] +fn a_route_with_no_choices_has_a_settled_authority() { + let r = direct_gate_route(); + assert!(r.authority_is_settled()); + assert_eq!(r.best_achievable_authority().level, Fidelity::SourceExact); + assert_eq!(r.worst_achievable_authority().level, Fidelity::SourceExact); +} + +#[test] +fn two_routes_of_different_fidelity_are_both_preserved() { + // The case a single operation-level authority cannot express: reporting + // source-exact risks binding the approximate route; reporting approximate + // understates the exact one. + let c = OperationCapability::available(vec![direct_gate_route(), strided_fused_route()]); + assert_eq!(c.routes.len(), 2); + assert_eq!(c.best_achievable_authority(), Some(Fidelity::SourceExact)); + assert_eq!( + c.worst_achievable_authority(), + Some(Fidelity::NumericallyApproximate) + ); +} + +#[test] +fn an_operation_with_differing_routes_reports_an_unsettled_authority() { + let c = OperationCapability::available(vec![direct_gate_route(), strided_fused_route()]); + assert!( + !c.authority_is_settled(), + "binding can still change the outcome" + ); +} + +#[test] +fn an_operation_whose_routes_agree_reports_a_settled_authority() { + let c = OperationCapability::available(vec![direct_gate_route(), direct_gate_route()]); + assert!(c.authority_is_settled()); +} + +#[test] +fn an_unavailable_operation_has_no_routes_and_therefore_no_authority() { + // The fail-closed invariant, restated in the new shape: no routes means no + // fidelity, so a contradictory selection cannot acquire a weak-but-valid + // one by default. + let c = OperationCapability::unavailable(vec![OperationFailure::InvalidSelection { + detail: "disjoint segments".into(), + }]); + assert!(c.routes.is_empty()); + assert_eq!(c.best_achievable_authority(), None); + assert_eq!(c.worst_achievable_authority(), None); + assert!(!c.authority_is_settled()); +} + +#[test] +fn routes_exist_exactly_when_the_operation_admits() { + assert!(OperationCapability::available(vec![direct_gate_route()]).is_well_formed()); + assert!( + OperationCapability::unavailable(vec![OperationFailure::NoExecutableRoute { layer: 0 }]) + .is_well_formed() + ); + // An available operation with no routes is malformed by construction. + let malformed = OperationCapability::available(Vec::new()); + assert!(!malformed.is_well_formed()); +} + +#[test] +fn a_degraded_operation_still_carries_its_routes() { + use super::operation::Degradation; + let c = OperationCapability::degraded( + vec![direct_gate_route()], + vec![Degradation::MissingQueryMetadata { what: "labels" }], + ); + assert!(c.is_available()); + assert!(c.is_well_formed()); + assert_eq!(c.best_achievable_authority(), Some(Fidelity::SourceExact)); +} + +// ── Choice groups ────────────────────────────────────────────────────────── + +const FUSED: &[RegionRole] = &[RegionRole::GateUpFused, RegionRole::Down]; +const DECOMPOSED: &[RegionRole] = &[RegionRole::Gate, RegionRole::Up, RegionRole::Down]; + +fn two_way_choice(layer: u32) -> PlanChoice { + PlanChoice { + bank: BankCoordinate::new(layer, 0), + alternatives: vec![ + alternative(FUSED, Fidelity::NumericallyApproximate), + alternative(DECOMPOSED, Fidelity::SourceExact), + ], + } +} + +#[test] +fn a_choice_group_reports_a_range_rather_than_a_value() { + // Binding has not chosen, so the honest statement is an interval. + let plan = OperationPlan { + fixed_components: Vec::new(), + choices: vec![two_way_choice(0)], + }; + assert!(plan.has_open_choices()); + assert_eq!( + plan.best_achievable_authority().level, + Fidelity::SourceExact + ); + assert_eq!( + plan.worst_achievable_authority().level, + Fidelity::NumericallyApproximate + ); +} + +#[test] +fn thirty_independent_choices_do_not_enumerate() { + // The property the representation exists for: 2^30 whole-model routes are + // never materialised. Thirty choice groups stay thirty. + let plan = OperationPlan { + fixed_components: Vec::new(), + choices: (0..30).map(two_way_choice).collect(), + }; + assert_eq!(plan.choices.len(), 30); + assert_eq!( + plan.best_achievable_authority().level, + Fidelity::SourceExact + ); + assert_eq!( + plan.worst_achievable_authority().level, + Fidelity::NumericallyApproximate + ); +} + +#[test] +fn a_fixed_region_drags_the_ceiling_down_across_every_choice() { + // Weakest-link still applies: no binding can beat the fixed regions. + let plan = OperationPlan { + fixed_components: vec![ + region(0, RegionRole::LatentIn, Fidelity::NumericallyApproximate).as_component(), + ], + choices: vec![PlanChoice { + bank: BankCoordinate::new(0, 0), + alternatives: vec![alternative(FUSED, Fidelity::SourceExact)], + }], + }; + assert_eq!( + plan.best_achievable_authority().level, + Fidelity::NumericallyApproximate + ); +} + +#[test] +fn a_single_alternative_choice_is_not_an_open_choice() { + let plan = OperationPlan { + fixed_components: Vec::new(), + choices: vec![PlanChoice { + bank: BankCoordinate::new(0, 0), + alternatives: vec![alternative(FUSED, Fidelity::SourceEquivalent)], + }], + }; + assert!(!plan.has_open_choices()); + assert_eq!( + plan.best_achievable_authority().level, + plan.worst_achievable_authority().level + ); +} + +#[test] +fn the_best_alternative_is_the_one_whose_weakest_region_is_strongest() { + // Not "the one with an exact region somewhere" — weakest-link decides. + let mixed = QualifiedAlternative { + alternative: DECOMPOSED, + regions: vec![ + region(0, RegionRole::Gate, Fidelity::SourceExact), + region(0, RegionRole::Up, Fidelity::StructurallyApproximate), + region(0, RegionRole::Down, Fidelity::SourceExact), + ], + components: Vec::new(), + }; + let uniform = alternative(FUSED, Fidelity::NumericallyApproximate); + let choice = PlanChoice { + bank: BankCoordinate::new(0, 0), + alternatives: vec![mixed, uniform.clone()], + }; + assert_eq!( + choice.best_alternative().unwrap().alternative, + uniform.alternative, + "the uniformly-approximate route beats one with a weak link" + ); +} + +#[test] +fn a_plan_lists_every_coordinate_it_could_read() { + let plan = OperationPlan { + fixed_components: vec![ + region(0, RegionRole::LatentIn, Fidelity::SourceExact).as_component() + ], + choices: vec![two_way_choice(0)], + }; + let coords = plan.all_coordinates(); + // latent_in + fused + down + gate + up, deduplicated across alternatives. + assert_eq!(coords.len(), 5, "{coords:?}"); +} + +#[test] +fn an_empty_plan_is_empty_and_settled() { + let plan = OperationPlan::default(); + assert!(plan.is_empty()); + assert!(!plan.has_open_choices()); + // Nothing selected floors at analysis-only rather than defaulting high. + assert_eq!( + plan.best_achievable_authority().level, + Fidelity::AnalysisOnly + ); +} + +#[test] +fn programme_alternatives_survive_into_the_plan_unchanged() { + // The alternative identity is carried through so kernel binding can match + // a kernel to the layout it actually supports. + let choice = two_way_choice(0); + let declared = Programme::GatedMlpV1.role_alternatives(); + assert_eq!(choice.alternatives[0].alternative, declared[0]); + assert_eq!(choice.alternatives[1].alternative, declared[1]); +} diff --git a/crates/larql-vindex/src/format/capability/reconstruction.rs b/crates/larql-vindex/src/format/capability/reconstruction.rs new file mode 100644 index 000000000..988ec8948 --- /dev/null +++ b/crates/larql-vindex/src/format/capability/reconstruction.rs @@ -0,0 +1,271 @@ +//! Canonical reconstruction — the COMPILE contract (spec §9.1, §15.6). +//! +//! > **COMPILE performs document-scoped canonical reconstruction from declared +//! > baseline variants. It is independent of the active profile and reports +//! > fidelity derived from those baseline regions.** +//! +//! A serving profile selecting MXFP4 must not silently change what COMPILE +//! emits. So reconstruction never consumes a resolved profile: it builds its +//! own selection from the representation catalogue's declared baselines. Not +//! "the profile's selection with omissions overridden" — that would mix two +//! scopes, and a deliberate profile omission would leak into an operation that +//! has nothing to do with serving policy. +//! +//! # Canonical is not a fidelity claim +//! +//! The baseline is canonical *within the document*, and its recorded fidelity +//! is still measured against the source checkpoint. A losslessly-contained +//! native-MXFP4 baseline is `source-equivalent`; a baseline quantised from +//! BF16 is `numerically-approximate`. Both are canonical. So +//! +//! ```text +//! reconstruct_canonical: available +//! authority: numerically-approximate +//! ``` +//! +//! is a valid and unremarkable result. The operation is deliberately *not* +//! called `ReconstructCheckpoint` or `ReconstructOriginal`, because those names +//! promise a fidelity the baseline may not possess. +//! +//! # Reconstruction targets logical tensors, not stored regions +//! +//! COMPILE must reproduce the *checkpoint's* tensor vocabulary, not the +//! index's storage vocabulary. A fused `gate_up` region may have to be split +//! into two checkpoint tensors, and separately stored roles may have to be +//! joined into one. That mapping comes from importer metadata, never from +//! guessing at role names — otherwise COMPILE reproduces storage structure and +//! calls it a model. + +use crate::format::lyrw2::region_format::{Packing, RegionFormat}; +use crate::format::lyrw2::region_role::RegionRole; + +use super::authority::Fidelity; +use super::coordinate::RegionCoordinate; + +/// How one logical checkpoint tensor is recovered from stored regions. +/// +/// A bounded vocabulary that reverses the extractor's physical encodings — +/// deliberately not a general graph interpreter (§8.3). +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum TensorReconstruction { + /// One region decodes straight to one tensor. + Direct { role: RegionRole, output: String }, + /// A fused region splits into several checkpoint tensors. + SplitFused { + role: RegionRole, + outputs: Vec, + }, + /// Several stored roles join into one checkpoint tensor. + JoinRoles { + roles: Vec, + output: String, + }, + /// Values and scales stored separately recombine into one tensor. + PairedValuesScales { + values: RegionRole, + scales: RegionRole, + output: String, + }, +} + +impl TensorReconstruction { + /// Regions this recipe needs. All of them must be present and decodable. + pub fn required_roles(&self) -> Vec { + match self { + Self::Direct { role, .. } | Self::SplitFused { role, .. } => vec![*role], + Self::JoinRoles { roles, .. } => roles.clone(), + Self::PairedValuesScales { values, scales, .. } => vec![*values, *scales], + } + } + + /// Checkpoint tensors this recipe produces. + pub fn outputs(&self) -> Vec { + match self { + Self::Direct { output, .. } + | Self::JoinRoles { output, .. } + | Self::PairedValuesScales { output, .. } => vec![output.clone()], + Self::SplitFused { outputs, .. } => outputs.clone(), + } + } +} + +/// A region set's declared variants and which one is canonical. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RegionSetCatalogue { + pub coordinate: RegionCoordinate, + /// Name of the canonical variant. Baseline status is a pointer, not a + /// fidelity promotion. + pub baseline: String, + pub variants: Vec, +} + +/// One physically-declared encoding of a region set. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CatalogueVariant { + pub name: String, + pub format: RegionFormat, + pub packing: Packing, + /// Measured against the **source checkpoint**, never against the baseline + /// itself — that is the loophole §9.2 closes. + pub fidelity: Fidelity, + /// Whether the bytes are actually on disk. A declared-but-absent variant + /// is a catalogue error, not a silent fallback to a sibling. + pub present: bool, +} + +impl RegionSetCatalogue { + pub fn baseline_variant(&self) -> Option<&CatalogueVariant> { + self.variants.iter().find(|v| v.name == self.baseline) + } +} + +/// Why canonical reconstruction cannot proceed. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ReconstructionFailure { + /// The catalogue names no canonical variant for this region set. + NoDeclaredBaseline { coordinate: RegionCoordinate }, + /// The declared baseline is not among the region set's variants. + BaselineNotDeclared { + coordinate: RegionCoordinate, + baseline: String, + declared: Vec, + }, + /// The baseline is declared but its bytes are absent. Never falls back to + /// a sibling variant — that would silently change what COMPILE emits. + BaselineAbsent { + coordinate: RegionCoordinate, + baseline: String, + }, + /// This build cannot decode the baseline into a logical tensor. + BaselineUndecodable { + coordinate: RegionCoordinate, + format: RegionFormat, + packing: Packing, + }, + /// A recipe needs a role the catalogue does not carry. + RecipeRoleMissing { output: String, role: RegionRole }, + /// A values/scales pair is only half present. + IncompletePair { + output: String, + present: RegionRole, + missing: RegionRole, + }, +} + +impl ReconstructionFailure { + pub fn describe(&self) -> String { + match self { + Self::NoDeclaredBaseline { coordinate } => { + format!("{}: no canonical variant declared", coordinate.describe()) + } + Self::BaselineNotDeclared { + coordinate, + baseline, + declared, + } => format!( + "{}: baseline '{baseline}' is not among the declared variants [{}]", + coordinate.describe(), + declared.join(", ") + ), + Self::BaselineAbsent { + coordinate, + baseline, + } => format!( + "{}: baseline '{baseline}' is declared but its bytes are absent", + coordinate.describe() + ), + Self::BaselineUndecodable { + coordinate, + format, + packing, + } => format!( + "{}: baseline cannot be decoded by this build (format {}, packing {})", + coordinate.describe(), + format.name(), + packing.name() + ), + Self::RecipeRoleMissing { output, role } => format!( + "tensor '{output}' needs role '{}', which the document does not carry", + role.name() + ), + Self::IncompletePair { + output, + present, + missing, + } => format!( + "tensor '{output}' pairs '{}' with '{}', but only the former is present", + present.name(), + missing.name() + ), + } + } +} + +/// The canonical variants chosen for reconstruction. +/// +/// Built from the catalogue, never from a resolved profile. +#[derive(Debug, Clone, PartialEq, Eq, Default)] +pub struct CanonicalSelection { + pub chosen: Vec<(RegionCoordinate, CatalogueVariant)>, + pub failures: Vec, +} + +impl CanonicalSelection { + /// Select each region set's declared baseline. + /// + /// Takes the catalogue and nothing else — no profile, no placement, no + /// browse mode. Those are serving policy and have no bearing on what the + /// document canonically contains. + pub fn from_catalogue( + catalogue: &[RegionSetCatalogue], + decodable: &dyn Fn(RegionFormat, Packing) -> bool, + ) -> Self { + let mut chosen = Vec::new(); + let mut failures = Vec::new(); + + for set in catalogue { + if set.baseline.trim().is_empty() { + failures.push(ReconstructionFailure::NoDeclaredBaseline { + coordinate: set.coordinate.clone(), + }); + continue; + } + let Some(variant) = set.baseline_variant() else { + failures.push(ReconstructionFailure::BaselineNotDeclared { + coordinate: set.coordinate.clone(), + baseline: set.baseline.clone(), + declared: set.variants.iter().map(|v| v.name.clone()).collect(), + }); + continue; + }; + if !variant.present { + failures.push(ReconstructionFailure::BaselineAbsent { + coordinate: set.coordinate.clone(), + baseline: set.baseline.clone(), + }); + continue; + } + if !decodable(variant.format, variant.packing) { + failures.push(ReconstructionFailure::BaselineUndecodable { + coordinate: set.coordinate.clone(), + format: variant.format, + packing: variant.packing, + }); + continue; + } + chosen.push((set.coordinate.clone(), variant.clone())); + } + + Self { chosen, failures } + } + + pub fn is_complete(&self) -> bool { + self.failures.is_empty() + } + + /// Fidelities of the chosen baselines — the authority fold's domain for + /// reconstruction. + pub fn fidelities(&self) -> Vec { + self.chosen.iter().map(|(_, v)| v.fidelity).collect() + } +} diff --git a/crates/larql-vindex/src/format/capability/reconstruction_tests.rs b/crates/larql-vindex/src/format/capability/reconstruction_tests.rs new file mode 100644 index 000000000..ba9b8eaa0 --- /dev/null +++ b/crates/larql-vindex/src/format/capability/reconstruction_tests.rs @@ -0,0 +1,258 @@ +//! Colocated tests for `reconstruction` — the COMPILE contract. +//! +//! The load-bearing property is scope: canonical reconstruction must depend on +//! the document's declared baselines and on nothing a profile chose. The tests +//! that matter most are therefore the ones that *change serving policy and +//! assert the reconstruction is unmoved*. + +use crate::format::lyrw2::region_format::{Packing, RegionFormat}; +use crate::format::lyrw2::region_role::RegionRole; + +use super::authority::{derive_authority, AuthorityInputs, Fidelity}; +use super::coordinate::RegionCoordinate; +use super::reconstruction::{ + CanonicalSelection, CatalogueVariant, ReconstructionFailure, RegionSetCatalogue, + TensorReconstruction, +}; + +const BASELINE: &str = "exact-q6k"; +const SIBLING: &str = "native-mxfp4"; + +fn everything_decodes(_: RegionFormat, _: Packing) -> bool { + true +} + +fn nothing_decodes(_: RegionFormat, _: Packing) -> bool { + false +} + +fn coordinate(layer: u32) -> RegionCoordinate { + RegionCoordinate::new(layer, 0, None, RegionRole::GateUpFused) +} + +fn variant(name: &str, fidelity: Fidelity, present: bool) -> CatalogueVariant { + CatalogueVariant { + name: name.into(), + format: RegionFormat::Q6K, + packing: Packing::BlocksWithScalesInline, + fidelity, + present, + } +} + +/// A region set with a canonical Q6_K baseline and an MXFP4 sibling — the +/// §9.1 example, where a serving profile might prefer the sibling. +fn two_variant_set(layer: u32) -> RegionSetCatalogue { + RegionSetCatalogue { + coordinate: coordinate(layer), + baseline: BASELINE.into(), + variants: vec![ + variant(BASELINE, Fidelity::SourceEquivalent, true), + variant(SIBLING, Fidelity::SourceExact, true), + ], + } +} + +#[test] +fn the_declared_baseline_is_chosen_not_the_most_faithful_variant() { + // The sibling is source-EXACT and the baseline only source-equivalent. + // Canonical still means "declared baseline": COMPILE reproduces the + // document's canonical form, not the best available encoding. + let sel = CanonicalSelection::from_catalogue(&[two_variant_set(0)], &everything_decodes); + assert!(sel.is_complete()); + assert_eq!(sel.chosen.len(), 1); + assert_eq!(sel.chosen[0].1.name, BASELINE); + assert_eq!(sel.fidelities(), vec![Fidelity::SourceEquivalent]); +} + +#[test] +fn reconstruction_authority_can_be_approximate_and_still_be_canonical() { + // Baseline status is a pointer, not a fidelity promotion. A baseline + // quantised from BF16 reconstructs canonically at numerically-approximate, + // which is why the operation is not called ReconstructOriginal. + let lossy = RegionSetCatalogue { + baseline: BASELINE.into(), + variants: vec![variant(BASELINE, Fidelity::NumericallyApproximate, true)], + ..two_variant_set(0) + }; + let sel = CanonicalSelection::from_catalogue(&[lossy], &everything_decodes); + assert!(sel.is_complete()); + let authority = derive_authority(&AuthorityInputs { + selected_fidelities: sel.fidelities(), + execution_complete: true, + structural: None, + }); + assert_eq!(authority, Fidelity::NumericallyApproximate); +} + +#[test] +fn an_absent_baseline_never_falls_back_to_a_present_sibling() { + // Falling back would silently change what COMPILE emits — the exact + // failure the canonical contract exists to prevent. + let mut set = two_variant_set(0); + set.variants[0].present = false; + let sel = CanonicalSelection::from_catalogue(&[set], &everything_decodes); + + assert!(!sel.is_complete()); + assert!(sel.chosen.is_empty(), "must not substitute the sibling"); + assert!(matches!( + sel.failures[0], + ReconstructionFailure::BaselineAbsent { .. } + )); +} + +#[test] +fn a_baseline_naming_an_undeclared_variant_lists_what_was_declared() { + let mut set = two_variant_set(0); + set.baseline = "native-nvfp4".into(); + let sel = CanonicalSelection::from_catalogue(&[set], &everything_decodes); + match &sel.failures[0] { + ReconstructionFailure::BaselineNotDeclared { declared, .. } => { + assert_eq!(declared, &vec![BASELINE.to_string(), SIBLING.to_string()]); + } + other => panic!("expected not-declared, got {other:?}"), + } +} + +#[test] +fn a_region_set_with_no_baseline_is_refused() { + let mut set = two_variant_set(0); + set.baseline = " ".into(); + let sel = CanonicalSelection::from_catalogue(&[set], &everything_decodes); + assert!(matches!( + sel.failures[0], + ReconstructionFailure::NoDeclaredBaseline { .. } + )); +} + +#[test] +fn an_undecodable_baseline_names_both_format_and_packing() { + let sel = CanonicalSelection::from_catalogue(&[two_variant_set(0)], ¬hing_decodes); + match &sel.failures[0] { + ReconstructionFailure::BaselineUndecodable { + format, packing, .. + } => { + assert_eq!(*format, RegionFormat::Q6K); + assert_eq!(*packing, Packing::BlocksWithScalesInline); + } + other => panic!("expected undecodable, got {other:?}"), + } +} + +#[test] +fn one_broken_region_set_does_not_discard_the_others() { + // Reconstruction reports every failure it finds; a single bad set must not + // abort the survey, or the operator fixes one problem per run. + let mut broken = two_variant_set(1); + broken.variants[0].present = false; + let sel = CanonicalSelection::from_catalogue( + &[two_variant_set(0), broken, two_variant_set(2)], + &everything_decodes, + ); + assert_eq!(sel.chosen.len(), 2); + assert_eq!(sel.failures.len(), 1); + assert!(!sel.is_complete()); +} + +#[test] +fn every_failure_carries_its_coordinate() { + let mut set = two_variant_set(37); + set.variants[0].present = false; + let sel = CanonicalSelection::from_catalogue(&[set], &everything_decodes); + assert!(sel.failures[0].describe().contains("layer 37")); +} + +#[test] +fn each_failure_kind_renders_distinguishably() { + let rendered: Vec = [ + ReconstructionFailure::NoDeclaredBaseline { + coordinate: coordinate(0), + }, + ReconstructionFailure::BaselineNotDeclared { + coordinate: coordinate(0), + baseline: "x".into(), + declared: vec!["y".into()], + }, + ReconstructionFailure::BaselineAbsent { + coordinate: coordinate(0), + baseline: "x".into(), + }, + ReconstructionFailure::BaselineUndecodable { + coordinate: coordinate(0), + format: RegionFormat::Nvfp4, + packing: Packing::RowMajor, + }, + ReconstructionFailure::RecipeRoleMissing { + output: "w1".into(), + role: RegionRole::Gate, + }, + ReconstructionFailure::IncompletePair { + output: "w1".into(), + present: RegionRole::Scales, + missing: RegionRole::Down, + }, + ] + .iter() + .map(|f| f.describe()) + .collect(); + for i in 0..rendered.len() { + for j in (i + 1)..rendered.len() { + assert_ne!(rendered[i], rendered[j], "{i} vs {j}"); + } + } +} + +#[test] +fn an_empty_catalogue_reconstructs_vacuously() { + let sel = CanonicalSelection::from_catalogue(&[], &everything_decodes); + assert!(sel.is_complete()); + assert!(sel.fidelities().is_empty()); +} + +// ── Recipes target checkpoint tensors, not stored regions ────────────────── + +#[test] +fn a_fused_region_can_split_into_two_checkpoint_tensors() { + // Storage fuses gate and up; the checkpoint keeps them apart. COMPILE must + // reproduce the checkpoint's vocabulary, not the index's. + let r = TensorReconstruction::SplitFused { + role: RegionRole::GateUpFused, + outputs: vec!["w1.weight".into(), "w3.weight".into()], + }; + assert_eq!(r.required_roles(), vec![RegionRole::GateUpFused]); + assert_eq!(r.outputs().len(), 2); +} + +#[test] +fn separate_roles_can_join_into_one_checkpoint_tensor() { + // The mirror direction, for a checkpoint whose canonical key is fused. + let r = TensorReconstruction::JoinRoles { + roles: vec![RegionRole::Gate, RegionRole::Up], + output: "gate_up_proj.weight".into(), + }; + assert_eq!(r.required_roles().len(), 2); + assert_eq!(r.outputs(), vec!["gate_up_proj.weight".to_string()]); +} + +#[test] +fn a_values_scales_pair_requires_both_halves() { + let r = TensorReconstruction::PairedValuesScales { + values: RegionRole::Down, + scales: RegionRole::Scales, + output: "w2.weight".into(), + }; + assert_eq!( + r.required_roles(), + vec![RegionRole::Down, RegionRole::Scales] + ); +} + +#[test] +fn a_direct_recipe_maps_one_role_to_one_tensor() { + let r = TensorReconstruction::Direct { + role: RegionRole::Down, + output: "w2.weight".into(), + }; + assert_eq!(r.required_roles(), vec![RegionRole::Down]); + assert_eq!(r.outputs(), vec!["w2.weight".to_string()]); +} diff --git a/crates/larql-vindex/src/format/capability/role.rs b/crates/larql-vindex/src/format/capability/role.rs new file mode 100644 index 000000000..686cec73d --- /dev/null +++ b/crates/larql-vindex/src/format/capability/role.rs @@ -0,0 +1,250 @@ +//! Per-role capability facts (spec §11). +//! +//! Two axes, deliberately separate: +//! +//! - **availability** — are the bytes there at all, and if not, why? +//! - **reference support** — can *this build* interpret them? +//! +//! An earlier single enum combined these with kernel maturity, and the +//! conflation surfaced the moment traversal was written: traversal stops at +//! reference executability, so a combined type carried a variant it could +//! never emit. Kernel maturity now lives entirely outside this module. +//! +//! The distinction that must survive to the caller is the §11 one: bytes that +//! are absent and bytes this build cannot decode are different situations with +//! different fixes, and neither is "no grouped kernel available". + +use crate::format::lyrw2::region_format::{Packing, RegionFormat}; + +use super::coordinate::{AbsenceKind, RegionCoordinate}; + +/// Whether the bytes exist. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum OperandAvailability { + /// Not present, with the role-local reason. Cross-role incompatibility is + /// *not* one of these — that is relational and lives on the alternative. + Absent(AbsenceKind), + Present, +} + +impl OperandAvailability { + pub fn is_present(&self) -> bool { + matches!(self, Self::Present) + } + + /// Whether the absence indicates a broken index rather than a deliberate + /// scoping decision. + pub fn indicates_defect(&self) -> bool { + match self { + Self::Absent(kind) => kind.is_defect(), + Self::Present => false, + } + } +} + +/// Whether this build can interpret present bytes through the generic path. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ReferenceSupport { + Supported, + /// Codec unknown to this build. The artifact may be intact and newer. + UnsupportedFormat(RegionFormat), + /// Codec known, layout not. + UnsupportedPacking(Packing), + /// The role itself is unregistered here — a vendor or future tag. + UnsupportedRole, +} + +impl ReferenceSupport { + pub fn is_supported(&self) -> bool { + matches!(self, Self::Supported) + } + + pub fn describe(&self) -> String { + match self { + Self::Supported => "interpretable by the reference path".into(), + Self::UnsupportedFormat(f) => { + format!("format '{}' is not implemented by this build", f.name()) + } + Self::UnsupportedPacking(p) => { + format!("packing '{}' is not implemented by this build", p.name()) + } + Self::UnsupportedRole => "role is not registered in this build".into(), + } + } +} + +/// Everything traversal knows about one required role in one bank. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RoleCapability { + pub coordinate: RegionCoordinate, + pub availability: OperandAvailability, + /// Segments where the role is present **and** interpretable. Empty when + /// the role cannot contribute anywhere, whether because the bytes are + /// missing or because this build cannot read them. + pub usable_coverage: Vec, + pub reference_support: ReferenceSupport, +} + +impl RoleCapability { + /// Whether this role can contribute to a computation at all — the + /// predicate `C ≠ ∅` from the traversal precedence. + /// + /// Three ways to be unusable, deliberately unified here because tier 1 + /// treats them alike: the bytes are absent, this build cannot interpret + /// them, or they exist outside the required population. All three mean the + /// role contributes nowhere, which is what makes compatibility moot. + /// + /// `segmented` distinguishes a bank with segment identities from one + /// without; an unsegmented bank carries usability by presence rather than + /// by a coverage list. + pub fn is_usable_in(&self, segmented: bool) -> bool { + self.availability.is_present() + && self.reference_support.is_supported() + && (!segmented || !self.usable_coverage.is_empty()) + } + + /// Convenience for unsegmented banks. + pub fn is_usable(&self) -> bool { + self.is_usable_in(false) + } + + /// Whether this role's failure is role-local, and therefore outranks any + /// cross-role compatibility question (tier 1 of the traversal precedence). + /// + /// Compatibility is moot when a required operand does not exist anywhere: + /// asking whether two roles' segments agree is meaningless if one of them + /// has no segments. + pub fn blocks_alternative_in(&self, segmented: bool) -> bool { + !self.is_usable_in(segmented) + } + + pub fn describe(&self) -> String { + let what = match (&self.availability, &self.reference_support) { + (OperandAvailability::Absent(kind), _) => format!("absent — {}", kind.describe()), + (OperandAvailability::Present, support) if !support.is_supported() => { + format!("present but unusable — {}", support.describe()) + } + (OperandAvailability::Present, _) => "usable".into(), + }; + format!("{}: {}", self.coordinate.describe(), what) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::format::lyrw2::region_role::RegionRole; + + fn coordinate() -> RegionCoordinate { + RegionCoordinate::new(12, 0, Some(1), RegionRole::Down) + } + + fn usable() -> RoleCapability { + RoleCapability { + coordinate: coordinate(), + availability: OperandAvailability::Present, + usable_coverage: vec![0, 1], + reference_support: ReferenceSupport::Supported, + } + } + + #[test] + fn a_present_supported_role_is_usable() { + assert!(usable().is_usable()); + assert!(!usable().blocks_alternative_in(true)); + } + + #[test] + fn an_absent_role_is_not_usable() { + let r = RoleCapability { + availability: OperandAvailability::Absent(AbsenceKind::AbsentEverywhere), + usable_coverage: Vec::new(), + ..usable() + }; + assert!(!r.is_usable()); + assert!(r.blocks_alternative_in(true)); + } + + #[test] + fn a_present_but_uninterpretable_role_is_not_usable() { + // The index is intact; this build simply cannot read the region. It + // still cannot participate in a computation. + let r = RoleCapability { + reference_support: ReferenceSupport::UnsupportedFormat(RegionFormat::Nvfp4), + usable_coverage: Vec::new(), + ..usable() + }; + assert!(r.availability.is_present()); + assert!(!r.is_usable()); + } + + #[test] + fn an_uninterpretable_region_is_not_an_index_defect() { + let r = RoleCapability { + reference_support: ReferenceSupport::UnsupportedPacking(Packing::BlocksValues), + ..usable() + }; + assert!(!r.availability.indicates_defect()); + } + + #[test] + fn a_real_absence_is_an_index_defect() { + let r = RoleCapability { + availability: OperandAvailability::Absent(AbsenceKind::AbsentEverywhere), + ..usable() + }; + assert!(r.availability.indicates_defect()); + } + + #[test] + fn a_deliberate_omission_is_not_an_index_defect() { + let r = RoleCapability { + availability: OperandAvailability::Absent(AbsenceKind::OmittedBySelection), + ..usable() + }; + assert!(r.availability.indicates_defect().eq(&false)); + // ...but it still blocks the alternative it was required by. + assert!(r.blocks_alternative_in(true)); + } + + #[test] + fn every_unsupported_reason_is_distinguishable_in_text() { + let format = ReferenceSupport::UnsupportedFormat(RegionFormat::Mxfp8); + let packing = ReferenceSupport::UnsupportedPacking(Packing::BlocksScales); + let role = ReferenceSupport::UnsupportedRole; + assert!(format.describe().contains("mxfp8"), "{}", format.describe()); + assert!( + packing.describe().contains("blocks_scales"), + "{}", + packing.describe() + ); + assert!(role.describe().contains("role"), "{}", role.describe()); + assert!(!format.is_supported()); + assert!(!packing.is_supported()); + assert!(!role.is_supported()); + } + + #[test] + fn the_description_carries_the_full_coordinate() { + let s = usable().describe(); + assert!(s.contains("layer 12"), "{s}"); + assert!(s.contains("bank 0"), "{s}"); + assert!(s.contains("role down"), "{s}"); + assert!(s.contains("segment 1"), "{s}"); + } + + #[test] + fn absence_and_unreadability_read_differently() { + // §11's separation as text: an operator must not have to guess which. + let absent = RoleCapability { + availability: OperandAvailability::Absent(AbsenceKind::AbsentEverywhere), + ..usable() + }; + let unreadable = RoleCapability { + reference_support: ReferenceSupport::UnsupportedFormat(RegionFormat::Nvfp4), + ..usable() + }; + assert!(absent.describe().contains("absent")); + assert!(unreadable.describe().contains("present but unusable")); + } +} diff --git a/crates/larql-vindex/src/format/capability/scope.rs b/crates/larql-vindex/src/format/capability/scope.rs new file mode 100644 index 000000000..2b53eaf07 --- /dev/null +++ b/crates/larql-vindex/src/format/capability/scope.rs @@ -0,0 +1,211 @@ +//! Capability scopes (spec §9, §15.6). +//! +//! Two owners, kept structurally apart rather than by documentation: +//! +//! ```text +//! DocumentCapabilities what the artifact canonically contains +//! ProfileCapabilities what the active selection can do with it +//! ``` +//! +//! An undifferentiated `OperationReport` would invite callers to assume every +//! operation describes the resolved profile. It does not — canonical +//! reconstruction reads declared baselines and ignores serving policy +//! entirely. So this combination is correct, and the type makes it legible +//! rather than surprising: +//! +//! ```text +//! profile.local_decode unavailable +//! document.reconstruct_canonical available +//! ``` +//! +//! A browse profile selecting only gate regions fails local decode while its +//! parent document, holding every baseline variant, remains perfectly +//! reconstructable. + +use super::operation::OperationCapability; + +/// Operations scoped to the artifact, independent of any profile. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DocumentCapabilities { + /// COMPILE. Reconstructs from each region set's declared baseline. + /// + /// Deliberately the only member: if a second document-scoped operation + /// appears, it belongs here, and if a profile-scoped one is added here by + /// mistake the scope split has been lost. + pub reconstruct_canonical: OperationCapability, +} + +/// Operations scoped to the resolved profile's selection. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProfileCapabilities { + pub local_decode: OperationCapability, + /// Never inferred from local absence — requires a declared wire contract. + pub routed_remote_decode: OperationCapability, + pub walk: OperationCapability, + pub describe: OperationCapability, + pub select: OperationCapability, +} + +/// Everything derived about one index under one profile. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CapabilityReport { + pub document: DocumentCapabilities, + pub profile: ProfileCapabilities, +} + +impl CapabilityReport { + /// Whether the artifact can be exported canonically, whatever the profile + /// can or cannot serve. + pub fn is_reconstructable(&self) -> bool { + self.document.reconstruct_canonical.is_available() + } + + /// Whether the profile can run a complete forward pass locally. + pub fn can_decode_locally(&self) -> bool { + self.profile.local_decode.is_available() + } + + /// Whether any query operation is admitted. + pub fn can_query(&self) -> bool { + self.profile.walk.is_available() + || self.profile.describe.is_available() + || self.profile.select.is_available() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::format::capability::authority::Fidelity; + use crate::format::capability::coordinate::RegionCoordinate; + use crate::format::capability::operation::OperationFailure; + use crate::format::capability::plan::{OperationPlan, PlannedRegion, QualifiedOperationRoute}; + use crate::format::lyrw2::region_role::RegionRole; + + fn available(fidelity: Fidelity) -> OperationCapability { + OperationCapability::available(vec![QualifiedOperationRoute::new( + OperationPlan::fixed_regions(vec![PlannedRegion::new( + RegionCoordinate::new(0, 0, None, RegionRole::Gate), + fidelity, + )]), + )]) + } + + fn unavailable() -> OperationCapability { + OperationCapability::unavailable(vec![OperationFailure::NoExecutableRoute { layer: 0 }]) + } + + /// A browse profile over an intact document — the case the scope split + /// exists to make legible. + fn browse_over_intact_document() -> CapabilityReport { + CapabilityReport { + document: DocumentCapabilities { + reconstruct_canonical: available(Fidelity::SourceEquivalent), + }, + profile: ProfileCapabilities { + local_decode: unavailable(), + routed_remote_decode: unavailable(), + walk: available(Fidelity::SourceExact), + describe: available(Fidelity::SourceExact), + select: available(Fidelity::SourceExact), + }, + } + } + + #[test] + fn a_document_can_reconstruct_while_its_profile_cannot_decode() { + // Not a contradiction: reconstruction reads declared baselines, decode + // reads what the profile selected. + let r = browse_over_intact_document(); + assert!(r.is_reconstructable()); + assert!(!r.can_decode_locally()); + } + + #[test] + fn a_browse_profile_still_queries() { + assert!(browse_over_intact_document().can_query()); + } + + #[test] + fn reconstruction_authority_is_independent_of_the_profiles() { + // The document baseline is source-equivalent; WALK over selected gate + // regions is source-exact. Different regions, different fidelity, both + // honest — and neither derived from the other. + let r = browse_over_intact_document(); + assert_eq!( + r.document.reconstruct_canonical.best_achievable_authority(), + Some(Fidelity::SourceEquivalent) + ); + assert_eq!( + r.profile.walk.best_achievable_authority(), + Some(Fidelity::SourceExact) + ); + } + + #[test] + fn a_fully_serving_profile_admits_everything() { + let r = CapabilityReport { + document: DocumentCapabilities { + reconstruct_canonical: available(Fidelity::SourceEquivalent), + }, + profile: ProfileCapabilities { + local_decode: available(Fidelity::NumericallyApproximate), + routed_remote_decode: unavailable(), + walk: available(Fidelity::SourceExact), + describe: available(Fidelity::SourceExact), + select: available(Fidelity::SourceExact), + }, + }; + assert!(r.can_decode_locally()); + assert!(r.is_reconstructable()); + } + + #[test] + fn a_broken_document_can_still_have_a_serving_profile() { + // The mirror case: baselines lost, but the profile's selected variants + // survive. Serving works; canonical export does not. + let r = CapabilityReport { + document: DocumentCapabilities { + reconstruct_canonical: unavailable(), + }, + profile: ProfileCapabilities { + local_decode: available(Fidelity::NumericallyApproximate), + routed_remote_decode: unavailable(), + walk: available(Fidelity::SourceExact), + describe: available(Fidelity::SourceExact), + select: available(Fidelity::SourceExact), + }, + }; + assert!(!r.is_reconstructable()); + assert!(r.can_decode_locally()); + } + + #[test] + fn a_profile_with_no_query_operations_reports_none() { + let r = CapabilityReport { + document: DocumentCapabilities { + reconstruct_canonical: available(Fidelity::SourceExact), + }, + profile: ProfileCapabilities { + local_decode: available(Fidelity::SourceExact), + routed_remote_decode: unavailable(), + walk: unavailable(), + describe: unavailable(), + select: unavailable(), + }, + }; + // A serving-only index: fuses gate/up freely, so nothing is browsable. + assert!(!r.can_query()); + assert!(r.can_decode_locally()); + } + + #[test] + fn remote_decode_defaults_to_unavailable_without_a_contract() { + // Absence of a declared wire contract is not a reason to infer remote + // capability from missing local bytes. + assert!(!browse_over_intact_document() + .profile + .routed_remote_decode + .is_available()); + } +} diff --git a/crates/larql-vindex/src/format/capability/selection.rs b/crates/larql-vindex/src/format/capability/selection.rs new file mode 100644 index 000000000..af37dc388 --- /dev/null +++ b/crates/larql-vindex/src/format/capability/selection.rs @@ -0,0 +1,240 @@ +//! The concrete region selection — traversal's input (spec §9.1, §11). +//! +//! This is what variant-and-segment resolution *produces*, not something +//! traversal computes. The pipeline split matters: "does this profile inherit +//! from `exact`?" and "is this variant physically present?" are answered +//! before we get here; "can this selection decode, browse, or claim +//! source-exact?" is answered from it. Letting traversal answer the first pair +//! would make derivation circular — a profile whose validity depended on +//! capabilities derived from that same profile. +//! +//! Every region here is one the profile actually chose and that physically +//! exists. A profile naming an absent variant fails earlier, before a byte is +//! read (§9.1), so an absent *variant* never reaches traversal; an absent +//! *role* does, and that is what traversal reports on. + +use std::collections::BTreeMap; + +use crate::format::lyrw2::region_format::{Packing, RegionFormat}; +use crate::format::lyrw2::region_role::RegionRole; + +use super::authority::Fidelity; + +/// One physically present, profile-selected region. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SelectedRegion { + pub format: RegionFormat, + pub packing: Packing, + /// Fidelity of the *selected variant* against the source checkpoint, + /// recorded at extraction time. Never inferred from the format tag: a + /// Q6_K container of native MXFP4 values is source-equivalent, while Q6_K + /// quantised from BF16 is numerically-approximate, and the tag alone + /// cannot tell those apart. + pub fidelity: Fidelity, +} + +/// One bank's selected regions, per segment. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct BankSelection { + pub bank_id: u16, + /// Segments this selection requires coverage of. Empty means the bank is + /// treated as unsegmented, and regions are keyed with `None`. + pub required_segments: Vec, + /// `(segment, role) → region`. Segment is `None` for an unsegmented bank. + pub regions: BTreeMap<(Option, RegionRole), SelectedRegion>, + /// Roles the active profile deliberately drops. A browse slice omits + /// `down` by design, and that must not read as corruption. + pub omitted_roles: Vec, +} + +impl BankSelection { + /// An unsegmented bank with the given roles, all at one format/fidelity. + pub fn unsegmented( + bank_id: u16, + roles: &[RegionRole], + format: RegionFormat, + packing: Packing, + fidelity: Fidelity, + ) -> Self { + let regions = roles + .iter() + .map(|r| { + ( + (None, *r), + SelectedRegion { + format, + packing, + fidelity, + }, + ) + }) + .collect(); + Self { + bank_id, + required_segments: Vec::new(), + regions, + omitted_roles: Vec::new(), + } + } + + /// Whether this bank is segmented at all. + pub fn is_segmented(&self) -> bool { + !self.required_segments.is_empty() + } + + /// The segment keys traversal must check — `[None]` when unsegmented. + pub fn segment_keys(&self) -> Vec> { + if self.is_segmented() { + self.required_segments.iter().copied().map(Some).collect() + } else { + vec![None] + } + } + + /// Segments carrying `role`, in ascending order. + /// + /// Ascending regardless of insertion order, so a report built from this is + /// diffable across runs. + pub fn segments_with(&self, role: RegionRole) -> Vec { + let mut found: Vec = self + .regions + .keys() + .filter(|(_, r)| *r == role) + .filter_map(|(s, _)| *s) + .collect(); + found.sort_unstable(); + found + } + + /// Whether any segment carries `role`. + pub fn has_role_anywhere(&self, role: RegionRole) -> bool { + self.regions.keys().any(|(_, r)| *r == role) + } + + pub fn region(&self, segment: Option, role: RegionRole) -> Option<&SelectedRegion> { + self.regions.get(&(segment, role)) + } + + /// Distinct roles present anywhere in this bank. + pub fn present_roles(&self) -> Vec { + let mut roles: Vec = self.regions.keys().map(|(_, r)| *r).collect(); + roles.sort(); + roles.dedup(); + roles + } + + pub fn is_omitted(&self, role: RegionRole) -> bool { + self.omitted_roles.contains(&role) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const GATE_UP: RegionRole = RegionRole::GateUpFused; + const DOWN: RegionRole = RegionRole::Down; + + fn region() -> SelectedRegion { + SelectedRegion { + format: RegionFormat::Q6K, + packing: Packing::BlocksWithScalesInline, + fidelity: Fidelity::SourceEquivalent, + } + } + + fn segmented() -> BankSelection { + let mut regions = BTreeMap::new(); + for seg in [0u16, 1] { + regions.insert((Some(seg), GATE_UP), region()); + regions.insert((Some(seg), DOWN), region()); + } + BankSelection { + bank_id: 0, + required_segments: vec![0, 1], + regions, + omitted_roles: Vec::new(), + } + } + + #[test] + fn an_unsegmented_bank_keys_regions_with_none() { + let b = BankSelection::unsegmented( + 0, + &[GATE_UP, DOWN], + RegionFormat::F16, + Packing::RowMajor, + Fidelity::SourceExact, + ); + assert!(!b.is_segmented()); + assert_eq!(b.segment_keys(), vec![None]); + assert!(b.region(None, GATE_UP).is_some()); + } + + #[test] + fn a_segmented_bank_lists_each_segment_key() { + assert_eq!(segmented().segment_keys(), vec![Some(0), Some(1)]); + } + + #[test] + fn segments_with_returns_ascending_order() { + let mut b = segmented(); + // Insert out of order; the accessor must still sort. + b.regions.insert((Some(5), DOWN), region()); + b.regions.insert((Some(3), DOWN), region()); + assert_eq!(b.segments_with(DOWN), vec![0, 1, 3, 5]); + } + + #[test] + fn segments_with_reports_nothing_for_an_absent_role() { + assert!(segmented().segments_with(RegionRole::Bias).is_empty()); + assert!(!segmented().has_role_anywhere(RegionRole::Bias)); + } + + #[test] + fn an_unsegmented_role_has_no_segment_numbers() { + // Present, but with no segment identity — `None` keys are filtered out + // of the numeric list rather than becoming a phantom segment 0. + let b = BankSelection::unsegmented( + 0, + &[GATE_UP], + RegionFormat::F16, + Packing::RowMajor, + Fidelity::SourceExact, + ); + assert!(b.has_role_anywhere(GATE_UP)); + assert!(b.segments_with(GATE_UP).is_empty()); + } + + #[test] + fn present_roles_are_sorted_and_deduplicated() { + // Two segments each carry both roles; the role list is still two long. + assert_eq!(segmented().present_roles(), vec![GATE_UP, DOWN]); + } + + #[test] + fn omitted_roles_are_recorded_separately_from_absence() { + let mut b = segmented(); + b.omitted_roles.push(RegionRole::Bias); + assert!(b.is_omitted(RegionRole::Bias)); + assert!(!b.is_omitted(DOWN)); + } + + #[test] + fn fidelity_is_carried_per_region_not_derived_from_format() { + // A Q6_K container of native MXFP4 values is source-equivalent; Q6_K + // quantised from BF16 is numerically-approximate. Same tag, different + // fidelity — so it must be recorded, never inferred. + let equivalent = SelectedRegion { + format: RegionFormat::Q6K, + packing: Packing::BlocksWithScalesInline, + fidelity: Fidelity::SourceEquivalent, + }; + let approximate = SelectedRegion { + fidelity: Fidelity::NumericallyApproximate, + ..equivalent.clone() + }; + assert_eq!(equivalent.format, approximate.format); + assert_ne!(equivalent.fidelity, approximate.fidelity); + } +} diff --git a/crates/larql-vindex/src/format/capability/traversal.rs b/crates/larql-vindex/src/format/capability/traversal.rs new file mode 100644 index 000000000..81c6a2c21 --- /dev/null +++ b/crates/larql-vindex/src/format/capability/traversal.rs @@ -0,0 +1,309 @@ +//! Programme traversal (spec §11) — the one place capabilities are derived. +//! +//! Traversal answers exactly one question: **are the selected bytes sufficient +//! for this programme?** It stops at reference executability. It does not +//! choose a kernel, does not decide authority, and does not know what a +//! profile is called. +//! +//! # Precedence +//! +//! For one required-role alternative, with `U` the required segment population +//! and `Cᵣ` each required role's *usable* coverage: +//! +//! | tier | condition | verdict | +//! | ---- | --------- | ------- | +//! | 1 | any `Cᵣ = ∅` | role-local failure; compatibility is **moot** and not evaluated | +//! | 2 | all `Cᵣ ≠ ∅`, sets differ | incompatible segment sets | +//! | 3 | all `Cᵣ` equal, `≠ U` | consistently partial coverage | +//! | 4 | all `Cᵣ = U` | reference-executable | +//! +//! Tier 1 outranking tier 2 is the load-bearing ordering: asking whether two +//! roles' segments agree is meaningless when one of them has no segments. +//! +//! "Usable" excludes regions whose codec this build cannot interpret. An +//! all-unsupported role therefore has `Cᵣ = ∅` and lands in tier 1 as an +//! unsupported-format failure, rather than being mistaken for a coverage gap. +//! +//! # Alternatives are evaluated independently +//! +//! A bank can be incompatible under `gate + up + down` and complete under +//! `gate_up_fused + down`. The layer is executable, and the failed alternative +//! is *evidence*, not a defect. Every successful alternative is preserved for +//! kernel binding to rank; a closest failure is chosen only when none succeed. + +use crate::format::lyrw2::region_role::RegionRole; +use crate::format::moe_manifest::programme::{Programme, RoleAlternative}; + +use super::compatibility::{RoleSegmentCoverage, SegmentCompatibility}; +use super::coordinate::{AbsenceKind, RegionCoordinate}; +use super::role::{OperandAvailability, ReferenceSupport, RoleCapability}; +use super::selection::BankSelection; + +/// Whether this build could execute an alternative through the generic path. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ReferenceExecution { + Executable, + /// Tier 1 — one or more required roles cannot contribute anywhere. + BlockedByOperands { + roles: Vec, + }, + /// Tier 2 — roles disagree on which segments they cover. + BlockedByIncompatibleSegments, + /// Tier 3 — roles agree, and together fall short of the required set. + BlockedByIncompleteSegments, +} + +impl ReferenceExecution { + pub fn is_executable(&self) -> bool { + matches!(self, Self::Executable) + } + + pub fn describe(&self) -> String { + match self { + Self::Executable => "reference-executable".into(), + Self::BlockedByOperands { roles } => format!( + "blocked: required role(s) unavailable — {}", + roles + .iter() + .map(|r| r.name()) + .collect::>() + .join(", ") + ), + Self::BlockedByIncompatibleSegments => { + "blocked: required roles cover incompatible segment sets".into() + } + Self::BlockedByIncompleteSegments => { + "blocked: required roles cover the same segments, short of the selection".into() + } + } + } +} + +/// One programme alternative, evaluated against a selection. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AlternativeReport { + pub alternative: RoleAlternative, + pub operands: Vec, + /// `None` when tier 1 fired — compatibility is moot when a required + /// operand does not exist, and reporting it would name a symptom over a + /// cause. + pub compatibility: Option, + pub reference_execution: ReferenceExecution, +} + +impl AlternativeReport { + pub fn is_executable(&self) -> bool { + self.reference_execution.is_executable() + } + + /// Whether this alternative failed because the *selection* is + /// contradictory, as opposed to incomplete or absent. Fails closed and + /// must never be laundered into an approximation. + pub fn is_invalid_selection(&self) -> bool { + self.compatibility + .as_ref() + .is_some_and(|c| c.is_invalid_selection()) + } + + /// How many required roles are unusable — the ranking key for choosing a + /// closest failure. + pub fn unusable_role_count(&self) -> usize { + self.operands.iter().filter(|o| !o.is_usable()).count() + } +} + +/// Every alternative for one bank, plus which of them survived. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct BankCapabilityReport { + pub layer: u32, + pub bank_id: u16, + pub programme: Programme, + pub required_population: Vec, + pub alternatives: Vec, +} + +impl BankCapabilityReport { + /// Alternatives that reference-execute, in declaration order. + /// + /// **All** of them, not the first: the kernel registry may support one at + /// a higher maturity than another, and collapsing here would hide a + /// Production path behind a Reference one. + pub fn successful_alternatives(&self) -> Vec<&AlternativeReport> { + self.alternatives + .iter() + .filter(|a| a.is_executable()) + .collect() + } + + pub fn is_executable(&self) -> bool { + self.alternatives.iter().any(|a| a.is_executable()) + } + + /// The failed alternative closest to working — fewest unusable roles, + /// ties broken by declaration order. + /// + /// `None` when something succeeded: a closest *failure* is only meaningful + /// when there is no success to report instead. + pub fn closest_failure(&self) -> Option<&AlternativeReport> { + if self.is_executable() { + return None; + } + self.alternatives + .iter() + .min_by_key(|a| a.unusable_role_count()) + } + + /// Whether any alternative failed on a contradictory selection. Reported + /// even when another alternative succeeds, because an invalid selection is + /// worth surfacing regardless of whether a sibling layout rescued it. + pub fn has_invalid_selection(&self) -> bool { + self.alternatives.iter().any(|a| a.is_invalid_selection()) + } +} + +/// Which codecs this build can interpret through the generic path. +pub trait ReferenceCodecs { + fn supports(&self, format: crate::format::lyrw2::region_format::RegionFormat) -> bool; + fn supports_packing(&self, packing: crate::format::lyrw2::region_format::Packing) -> bool; +} + +/// Traverse one bank's programme against a selection. +pub fn traverse_bank( + layer: u32, + programme: Programme, + selection: &BankSelection, + codecs: &dyn ReferenceCodecs, +) -> BankCapabilityReport { + let alternatives = programme + .role_alternatives() + .iter() + .copied() + .map(|alt| evaluate_alternative(layer, alt, selection, codecs)) + .collect(); + + BankCapabilityReport { + layer, + bank_id: selection.bank_id, + programme, + required_population: selection.required_segments.clone(), + alternatives, + } +} + +fn evaluate_alternative( + layer: u32, + alternative: RoleAlternative, + selection: &BankSelection, + codecs: &dyn ReferenceCodecs, +) -> AlternativeReport { + let operands: Vec = alternative + .iter() + .map(|role| evaluate_role(layer, *role, selection, codecs)) + .collect(); + + // ── Tier 1 — role-local failure outranks everything relational ── + let blocking: Vec = operands + .iter() + .filter(|o| o.blocks_alternative_in(selection.is_segmented())) + .map(|o| o.coordinate.role) + .collect(); + if !blocking.is_empty() { + return AlternativeReport { + alternative, + operands, + compatibility: None, // moot: an absent operand has no segments to agree about + reference_execution: ReferenceExecution::BlockedByOperands { roles: blocking }, + }; + } + + // ── Tiers 2–4 — every role contributes somewhere; do they agree? ── + let coverage: Vec = operands + .iter() + .map(|o| RoleSegmentCoverage::new(o.coordinate.role, o.usable_coverage.clone())) + .collect(); + let compatibility = SegmentCompatibility::evaluate(&selection.required_segments, &coverage); + + let reference_execution = match &compatibility { + SegmentCompatibility::Compatible => ReferenceExecution::Executable, + SegmentCompatibility::Incompatible { .. } => { + ReferenceExecution::BlockedByIncompatibleSegments + } + SegmentCompatibility::Partial { .. } => ReferenceExecution::BlockedByIncompleteSegments, + }; + + AlternativeReport { + alternative, + operands, + compatibility: Some(compatibility), + reference_execution, + } +} + +fn evaluate_role( + layer: u32, + role: RegionRole, + selection: &BankSelection, + codecs: &dyn ReferenceCodecs, +) -> RoleCapability { + let coordinate = RegionCoordinate::new(layer, selection.bank_id, None, role); + + // Deliberate omission is not corruption. Reporting it as its own cause is + // what stops a browse slice reading as a broken index. + if selection.is_omitted(role) { + return absent(coordinate, AbsenceKind::OmittedBySelection); + } + if !selection.has_role_anywhere(role) { + return absent(coordinate, AbsenceKind::AbsentEverywhere); + } + + // Present somewhere. Partition the required population into segments this + // build can actually interpret, and remember why any were rejected. + let mut usable = Vec::new(); + let mut support = ReferenceSupport::Supported; + for key in selection.segment_keys() { + let Some(region) = selection.region(key, role) else { + continue; + }; + if !codecs.supports(region.format) { + support = ReferenceSupport::UnsupportedFormat(region.format); + continue; + } + if !codecs.supports_packing(region.packing) { + support = ReferenceSupport::UnsupportedPacking(region.packing); + continue; + } + if let Some(seg) = key { + usable.push(seg); + } + } + usable.sort_unstable(); + + // Present but usable in no required segment, with no codec objection: the + // regions live outside the required population. Distinct from absence — + // the bytes are real and the repair is to reconcile the selection. + if selection.is_segmented() && usable.is_empty() && support.is_supported() { + return absent( + coordinate, + AbsenceKind::PresentOutsidePopulation { + found: selection.segments_with(role), + required: selection.required_segments.clone(), + }, + ); + } + + RoleCapability { + coordinate, + availability: OperandAvailability::Present, + usable_coverage: usable, + reference_support: support, + } +} + +fn absent(coordinate: RegionCoordinate, kind: AbsenceKind) -> RoleCapability { + RoleCapability { + coordinate, + availability: OperandAvailability::Absent(kind), + usable_coverage: Vec::new(), + reference_support: ReferenceSupport::Supported, + } +} diff --git a/crates/larql-vindex/src/format/capability/traversal_tests.rs b/crates/larql-vindex/src/format/capability/traversal_tests.rs new file mode 100644 index 000000000..3d01c766b --- /dev/null +++ b/crates/larql-vindex/src/format/capability/traversal_tests.rs @@ -0,0 +1,451 @@ +//! Colocated tests for `traversal` — complete inference cases. +//! +//! These are deliberately whole-scenario rather than per-enum: the value of a +//! traversal is the verdict it reaches over a realistic selection, and the +//! per-variant behaviour is covered in `role`, `compatibility` and +//! `coordinate`. Each case below is a situation an operator could actually +//! hit, and asserts both the verdict and the diagnosis. + +use std::collections::BTreeMap; + +use crate::format::lyrw2::region_format::{Packing, RegionFormat}; +use crate::format::lyrw2::region_role::RegionRole; +use crate::format::moe_manifest::programme::Programme; + +use super::authority::Fidelity; +use super::compatibility::{IncompatibilityShape, SegmentCompatibility}; +use super::coordinate::AbsenceKind; +use super::role::{OperandAvailability, ReferenceSupport}; +use super::selection::{BankSelection, SelectedRegion}; +use super::traversal::{traverse_bank, ReferenceCodecs, ReferenceExecution}; + +const LAYER: u32 = 12; +const BANK: u16 = 0; +const GATE: RegionRole = RegionRole::Gate; +const UP: RegionRole = RegionRole::Up; +const FUSED: RegionRole = RegionRole::GateUpFused; +const DOWN: RegionRole = RegionRole::Down; + +/// A build that reads everything — isolates coverage behaviour from codec +/// behaviour. +struct AllCodecs; +impl ReferenceCodecs for AllCodecs { + fn supports(&self, _: RegionFormat) -> bool { + true + } + fn supports_packing(&self, _: Packing) -> bool { + true + } +} + +/// A build that cannot read NVFP4 — the "newer artifact" case. +struct NoNvfp4; +impl ReferenceCodecs for NoNvfp4 { + fn supports(&self, f: RegionFormat) -> bool { + f != RegionFormat::Nvfp4 + } + fn supports_packing(&self, _: Packing) -> bool { + true + } +} + +fn region(format: RegionFormat) -> SelectedRegion { + SelectedRegion { + format, + packing: Packing::RowMajor, + fidelity: Fidelity::SourceEquivalent, + } +} + +/// Build a segmented selection from `(role, segments)` pairs. +fn selection(required: &[u16], roles: &[(RegionRole, &[u16])]) -> BankSelection { + let mut regions = BTreeMap::new(); + for (role, segs) in roles { + for s in *segs { + regions.insert((Some(*s), *role), region(RegionFormat::Q6K)); + } + } + BankSelection { + bank_id: BANK, + required_segments: required.to_vec(), + regions, + omitted_roles: Vec::new(), + } +} + +// ── 1. Both alternatives physically present ──────────────────────────────── + +#[test] +fn both_alternatives_present_means_both_survive_for_kernel_binding() { + let sel = selection( + &[0, 1], + &[ + (FUSED, &[0, 1]), + (GATE, &[0, 1]), + (UP, &[0, 1]), + (DOWN, &[0, 1]), + ], + ); + let r = traverse_bank(LAYER, Programme::GatedMlpV1, &sel, &AllCodecs); + assert!(r.is_executable()); + assert_eq!( + r.successful_alternatives().len(), + 2, + "kernel binding must receive both layouts" + ); + assert!(r.closest_failure().is_none()); +} + +// ── 2. Decomposed incompatible, fused complete ───────────────────────────── + +#[test] +fn a_layer_is_executable_through_fused_when_decomposed_is_incompatible() { + // gate covers {0}, up covers {1} — decomposed cannot run. The fused + // layout is whole, so the LAYER is fine and the decomposed failure is + // evidence rather than a defect. + let sel = selection( + &[0, 1], + &[(FUSED, &[0, 1]), (DOWN, &[0, 1]), (GATE, &[0]), (UP, &[1])], + ); + let r = traverse_bank(LAYER, Programme::GatedMlpV1, &sel, &AllCodecs); + + assert!(r.is_executable()); + assert_eq!(r.successful_alternatives().len(), 1); + assert_eq!( + r.successful_alternatives()[0].alternative, + &[FUSED, DOWN][..] + ); + // No closest failure is reported while something succeeds. + assert!(r.closest_failure().is_none()); + // ...but the invalid decomposed selection is still surfaced. + assert!(r.has_invalid_selection()); +} + +// ── 3. Pairwise overlap, empty global intersection ───────────────────────── + +#[test] +fn three_roles_overlapping_pairwise_but_not_globally_are_incompatible() { + let sel = selection( + &[0, 1, 2], + &[(GATE, &[0, 1]), (UP, &[1, 2]), (DOWN, &[2, 0])], + ); + let r = traverse_bank(LAYER, Programme::GatedMlpFusedFc1V1, &sel, &AllCodecs); + // The fused programme needs gate_up_fused, which is absent — tier 1. + assert!(!r.is_executable()); + + let decomposed = traverse_bank(LAYER, Programme::GatedMlpV1, &sel, &AllCodecs); + let alt = decomposed + .alternatives + .iter() + .find(|a| a.alternative == &[GATE, UP, DOWN][..]) + .unwrap(); + match alt.compatibility.as_ref().unwrap() { + SegmentCompatibility::Incompatible { shape, common, .. } => { + assert_eq!(*shape, IncompatibilityShape::NoCommonCoverage); + assert!(common.is_empty()); + } + other => panic!("expected incompatible, got {other:?}"), + } +} + +// ── 4. All segments present, codec unsupported ───────────────────────────── + +#[test] +fn an_unreadable_codec_is_a_whole_role_failure_not_partial_coverage() { + // The bytes are all there. This build simply cannot read them, so the role + // contributes nowhere — tier 1, with the codec named. + let mut sel = selection(&[0, 1], &[(FUSED, &[0, 1]), (DOWN, &[0, 1])]); + for s in [0u16, 1] { + sel.regions + .insert((Some(s), DOWN), region(RegionFormat::Nvfp4)); + } + let r = traverse_bank(LAYER, Programme::GatedMlpV1, &sel, &NoNvfp4); + + assert!(!r.is_executable()); + let alt = &r.alternatives[0]; + assert!(matches!( + alt.reference_execution, + ReferenceExecution::BlockedByOperands { .. } + )); + // Compatibility is moot and deliberately not evaluated. + assert!(alt.compatibility.is_none()); + + let down = alt + .operands + .iter() + .find(|o| o.coordinate.role == DOWN) + .unwrap(); + assert_eq!(down.availability, OperandAvailability::Present); + assert_eq!( + down.reference_support, + ReferenceSupport::UnsupportedFormat(RegionFormat::Nvfp4) + ); + // Present-but-unreadable is not an index defect — the artifact is intact. + assert!(!down.availability.indicates_defect()); +} + +// ── 5. Equal incomplete coverage ─────────────────────────────────────────── + +#[test] +fn equally_short_coverage_is_partial_not_incompatible() { + let sel = selection(&[0, 1, 2], &[(FUSED, &[0, 1]), (DOWN, &[0, 1])]); + let r = traverse_bank(LAYER, Programme::GatedMlpV1, &sel, &AllCodecs); + + assert!(!r.is_executable()); + let alt = &r.alternatives[0]; + assert_eq!( + alt.reference_execution, + ReferenceExecution::BlockedByIncompleteSegments + ); + assert!( + !alt.is_invalid_selection(), + "consistent shortfall is not a contradiction" + ); + assert!(matches!( + alt.compatibility.as_ref().unwrap(), + SegmentCompatibility::Partial { .. } + )); +} + +// ── 6. Gate-only deliberate omission ─────────────────────────────────────── + +#[test] +fn a_deliberate_gate_only_slice_is_not_reported_as_corruption() { + // The browse-slice case. Decode is unavailable, but nothing is broken and + // the selection is not invalid. + let mut sel = selection(&[0], &[(GATE, &[0])]); + sel.omitted_roles = vec![UP, DOWN, FUSED]; + let r = traverse_bank(LAYER, Programme::GatedMlpV1, &sel, &AllCodecs); + + assert!(!r.is_executable()); + assert!(!r.has_invalid_selection()); + + let failure = r.closest_failure().unwrap(); + let down = failure + .operands + .iter() + .find(|o| o.coordinate.role == DOWN) + .unwrap(); + assert_eq!( + down.availability, + OperandAvailability::Absent(AbsenceKind::OmittedBySelection) + ); + assert!(!down.availability.indicates_defect()); +} + +// ── 7. Gate-only accidental absence ──────────────────────────────────────── + +#[test] +fn an_accidental_gate_only_index_is_not_laundered_into_a_browse_slice() { + // Physically identical to case 6 in what is *missing*, but nothing + // declared the omission. Provenance must survive: this is a defect. + let sel = selection(&[0], &[(GATE, &[0])]); + let r = traverse_bank(LAYER, Programme::GatedMlpV1, &sel, &AllCodecs); + + assert!(!r.is_executable()); + let failure = r.closest_failure().unwrap(); + let down = failure + .operands + .iter() + .find(|o| o.coordinate.role == DOWN) + .unwrap(); + assert_eq!( + down.availability, + OperandAvailability::Absent(AbsenceKind::AbsentEverywhere) + ); + assert!( + down.availability.indicates_defect(), + "undeclared absence must stay a defect" + ); +} + +#[test] +fn declared_and_undeclared_absence_differ_only_in_provenance() { + // Same physical shape, same admission result, different diagnosis — the + // property that stops a corrupt index presenting as an intentional slice. + let mut declared = selection(&[0], &[(GATE, &[0])]); + declared.omitted_roles = vec![UP, DOWN, FUSED]; + let accidental = selection(&[0], &[(GATE, &[0])]); + + let a = traverse_bank(LAYER, Programme::GatedMlpV1, &declared, &AllCodecs); + let b = traverse_bank(LAYER, Programme::GatedMlpV1, &accidental, &AllCodecs); + + assert_eq!(a.is_executable(), b.is_executable()); + assert!(!a.has_invalid_selection() && !b.has_invalid_selection()); + + let defect = |r: &super::traversal::BankCapabilityReport| { + r.closest_failure() + .unwrap() + .operands + .iter() + .any(|o| o.availability.indicates_defect()) + }; + assert!(!defect(&a), "declared omission is not a defect"); + assert!(defect(&b), "undeclared absence is"); +} + +// ── 8. Same bytes, no grouped kernel ─────────────────────────────────────── + +#[test] +fn traversal_stops_at_reference_executability_and_names_no_kernel() { + // Traversal cannot express a kernel binding at all — that is the whole + // point of the split, and it is enforced by the type rather than by + // convention. + let sel = selection(&[0, 1], &[(FUSED, &[0, 1]), (DOWN, &[0, 1])]); + let r = traverse_bank(LAYER, Programme::GatedMlpV1, &sel, &AllCodecs); + assert_eq!( + r.successful_alternatives()[0].reference_execution, + ReferenceExecution::Executable + ); +} + +// ── 10. Alternative-order permutation ────────────────────────────────────── + +#[test] +fn the_successful_set_does_not_depend_on_selection_insertion_order() { + let forward = selection(&[0, 1], &[(FUSED, &[0, 1]), (DOWN, &[1, 0])]); + let reverse = selection(&[1, 0], &[(DOWN, &[0, 1]), (FUSED, &[1, 0])]); + let a = traverse_bank(LAYER, Programme::GatedMlpV1, &forward, &AllCodecs); + let b = traverse_bank(LAYER, Programme::GatedMlpV1, &reverse, &AllCodecs); + assert_eq!( + a.successful_alternatives().len(), + b.successful_alternatives().len() + ); + assert_eq!(a.is_executable(), b.is_executable()); +} + +#[test] +fn a_closest_failure_is_chosen_by_fewest_unusable_roles() { + // Fused is missing one role; decomposed is missing two. The fused layout + // is the closer failure and is the one worth reporting. + let sel = selection(&[0], &[(FUSED, &[0])]); + let r = traverse_bank(LAYER, Programme::GatedMlpV1, &sel, &AllCodecs); + assert!(!r.is_executable()); + let failure = r.closest_failure().unwrap(); + assert_eq!(failure.alternative, &[FUSED, DOWN][..]); + assert_eq!(failure.unusable_role_count(), 1); +} + +// ── Unsegmented banks ────────────────────────────────────────────────────── + +#[test] +fn an_unsegmented_bank_executes_on_presence_alone() { + let sel = BankSelection::unsegmented( + BANK, + &[FUSED, DOWN], + RegionFormat::F16, + Packing::RowMajor, + Fidelity::SourceExact, + ); + let r = traverse_bank(LAYER, Programme::GatedMlpV1, &sel, &AllCodecs); + assert!(r.is_executable()); + assert_eq!(r.required_population, Vec::::new()); +} + +#[test] +fn an_unsegmented_bank_missing_a_role_still_fails_tier_one() { + let sel = BankSelection::unsegmented( + BANK, + &[FUSED], + RegionFormat::F16, + Packing::RowMajor, + Fidelity::SourceExact, + ); + let r = traverse_bank(LAYER, Programme::GatedMlpV1, &sel, &AllCodecs); + assert!(!r.is_executable()); + assert!(matches!( + r.closest_failure().unwrap().reference_execution, + ReferenceExecution::BlockedByOperands { .. } + )); +} + +#[test] +fn gpt_oss_without_bias_fails_on_the_bias_role() { + let sel = selection(&[0], &[(FUSED, &[0]), (DOWN, &[0])]); + let r = traverse_bank(LAYER, Programme::GptOssExpertV1, &sel, &AllCodecs); + assert!(!r.is_executable()); + let blocked = &r.closest_failure().unwrap().reference_execution; + match blocked { + ReferenceExecution::BlockedByOperands { roles } => { + assert_eq!(roles, &vec![RegionRole::Bias]); + } + other => panic!("expected operand block, got {other:?}"), + } +} + +#[test] +fn every_diagnosis_carries_the_full_coordinate() { + let sel = selection(&[0, 1], &[(FUSED, &[0, 1])]); + let r = traverse_bank(LAYER, Programme::GatedMlpV1, &sel, &AllCodecs); + let down = r + .closest_failure() + .unwrap() + .operands + .iter() + .find(|o| o.coordinate.role == DOWN) + .unwrap(); + let s = down.describe(); + assert!(s.contains("layer 12"), "{s}"); + assert!(s.contains("bank 0"), "{s}"); + assert!(s.contains("role down"), "{s}"); +} + +#[test] +fn regions_outside_the_required_population_are_named_as_such() { + // The bytes exist in segments 7-8; this selection requires 0-1. That is a + // resolution fault, not a missing variant, and the repair differs. + let mut sel = selection(&[0, 1], &[(FUSED, &[0, 1])]); + for s in [7u16, 8] { + sel.regions + .insert((Some(s), DOWN), region(RegionFormat::Q6K)); + } + let r = traverse_bank(LAYER, Programme::GatedMlpV1, &sel, &AllCodecs); + assert!(!r.is_executable()); + + let down = r + .closest_failure() + .unwrap() + .operands + .iter() + .find(|o| o.coordinate.role == DOWN) + .unwrap(); + match &down.availability { + OperandAvailability::Absent(AbsenceKind::PresentOutsidePopulation { found, required }) => { + assert_eq!(found, &vec![7, 8]); + assert_eq!(required, &vec![0, 1]); + } + other => panic!("expected outside-population, got {other:?}"), + } +} + +#[test] +fn every_reference_execution_verdict_renders_distinguishably() { + let verdicts = [ + ReferenceExecution::Executable, + ReferenceExecution::BlockedByOperands { roles: vec![DOWN] }, + ReferenceExecution::BlockedByIncompatibleSegments, + ReferenceExecution::BlockedByIncompleteSegments, + ]; + let rendered: Vec = verdicts.iter().map(|v| v.describe()).collect(); + for i in 0..rendered.len() { + for j in (i + 1)..rendered.len() { + assert_ne!(rendered[i], rendered[j], "{i} vs {j}"); + } + } + assert!(rendered[0].contains("reference-executable")); + assert!(rendered[1].contains("down")); + assert!(verdicts[0].is_executable()); + assert!(!verdicts[1].is_executable()); +} + +#[test] +fn a_bank_with_no_usable_role_reports_every_blocking_role() { + let sel = selection(&[0], &[]); + let r = traverse_bank(LAYER, Programme::GatedMlpV1, &sel, &AllCodecs); + match &r.closest_failure().unwrap().reference_execution { + ReferenceExecution::BlockedByOperands { roles } => { + assert_eq!(roles.len(), 2, "{roles:?}"); + } + other => panic!("expected operand block, got {other:?}"), + } +} diff --git a/crates/larql-vindex/src/format/capability/walk.rs b/crates/larql-vindex/src/format/capability/walk.rs new file mode 100644 index 000000000..3fec5a124 --- /dev/null +++ b/crates/larql-vindex/src/format/capability/walk.rs @@ -0,0 +1,342 @@ +//! WALK inference (spec §15.1, §15.2, §15.4). +//! +//! The acceptance test for this whole module: +//! +//! > A WALK report must be derivable without ever asking whether the model can +//! > perform an expert forward pass. +//! +//! So this function never consults `up`, `down`, expert-output codecs, router +//! semantics, decode traversal, kernel maturity or `query/` metadata. If +//! unreadable `down` regions, broken routing or absent shared experts can +//! perturb a WALK report, the projection is not real. +//! +//! # The declared surface is a declaration +//! +//! Membership comes from an explicit document-level list, never from filtering +//! banks by health. Defining the surface as "banks whose browse mode is not +//! none" would make a broken bank *disappear* from the request and recreate +//! working-subset semantics through the back door. A bank declared searchable +//! that carries `browse: none`, an unreadable gate codec, or a missing latent +//! transform stays in the report — as a failure. +//! +//! # Coverage and fidelity are different axes +//! +//! A partial result set does not lower authority. If bank A's gate rows are +//! source-exact and bank B is unavailable, a partial walk over A returns +//! source-exact scores over an incomplete population. The numbers did not +//! become approximate; the result set became smaller. Completeness therefore +//! sits *beside* authority, never inside the fold — the same distinction +//! already drawn for absent query metadata. + +use crate::format::lyrw2::browse_mode::BrowseMode; +use crate::format::lyrw2::region_format::Packing; + +use super::coordinate::{BankCoordinate, RegionCoordinate}; +use super::operation::{OperationCapability, OperationFailure}; +use super::plan::{ + OperationPlan, PlanChoice, PlannedRegion, QualifiedAlternative, QualifiedOperationRoute, +}; +use super::role::ReferenceSupport; +use super::walk_request::{validate_query_vector, PartialPolicy, WalkRequest}; + +/// One way to reach a bank's gate rows. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct GateAccess { + /// `Direct` for an own-region gate; `Strided` for the gate half of a fused + /// region. `None` never appears here — it is a bank-level refusal. + pub mode: BrowseMode, + pub region: PlannedRegion, + pub packing: Packing, + pub support: ReferenceSupport, +} + +impl GateAccess { + /// Whether this access can actually deliver gate rows. + /// + /// Strided access additionally requires a packing that permits striding + /// into the gate half without decoding `up` — which is the whole cost + /// browse was avoiding. + pub fn is_usable(&self) -> bool { + self.support.is_supported() && self.mode.is_satisfiable_by(self.packing) + } + + pub fn reason_unusable(&self) -> Option { + if self.is_usable() { + return None; + } + if !self.support.is_supported() { + return Some(self.support.describe()); + } + Some(format!( + "{} access needs a packing that permits striding; '{}' does not", + self.mode.name(), + self.packing.name() + )) + } +} + +/// Everything WALK needs to know about one bank. Produced by resolution. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct WalkableBank { + pub coordinate: BankCoordinate, + /// Declared browse eligibility. `None` refuses the bank outright. + pub browse: BrowseMode, + /// Whether gate rows live in a latent space, requiring the query to be + /// projected through `routed_input` first (§15.4). + pub is_latent: bool, + pub gate_accesses: Vec, + /// residual → latent projection. Required exactly when `is_latent`. + pub routed_input: Option, +} + +/// Query-construction inputs, needed only for `TextQuery`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TextQueryInputs { + pub has_tokenizer: bool, + /// Embedding weights. A numerical component, so it enters route authority + /// — unlike the tokenizer, which has availability but no fidelity. + pub embeddings: Option, +} + +/// Document-level facts WALK depends on. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct WalkEnvironment { + pub residual_dimension: u32, + /// The declared searchable surface. An explicit list, never a filter. + pub declared_surface: Vec, + pub banks: Vec, + pub text_query: TextQueryInputs, +} + +impl WalkEnvironment { + fn bank(&self, at: BankCoordinate) -> Option<&WalkableBank> { + self.banks.iter().find(|b| b.coordinate == at) + } +} + +/// Why one targeted bank could not be walked. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TargetFailure { + pub bank: BankCoordinate, + pub reason: String, +} + +/// Whether the result set covers everything requested. +/// +/// Deliberately separate from authority: losing coverage and losing numerical +/// faithfulness are different axes, and folding one into the other would make +/// an exact-but-incomplete answer indistinguishable from an approximate one. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ResultSetCompleteness { + Complete, + Partial { + included: Vec, + omitted: Vec, + }, +} + +impl ResultSetCompleteness { + pub fn is_complete(&self) -> bool { + matches!(self, Self::Complete) + } +} + +/// A WALK capability with its coverage stated alongside. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct WalkCapability { + pub capability: OperationCapability, + pub completeness: ResultSetCompleteness, +} + +/// Infer WALK capability for one request. +pub fn infer_walk(env: &WalkEnvironment, request: &WalkRequest) -> WalkCapability { + if let Err(fault) = validate_query_vector(request.input, env.residual_dimension) { + return refuse(vec![OperationFailure::MissingDocumentInput { + what: "query vector of the model's residual width", + } + .with_detail(fault.describe())]); + } + + let mut fixed_regions = Vec::new(); + if request.input.needs_text_construction() { + if !env.text_query.has_tokenizer { + return refuse(vec![OperationFailure::MissingDocumentInput { + what: "tokenizer", + }]); + } + match &env.text_query.embeddings { + Some(e) => fixed_regions.push(e.clone()), + None => { + return refuse(vec![OperationFailure::MissingDocumentInput { + what: "embeddings", + }]) + } + } + } + + let targets = request.target.resolve(&env.declared_surface); + if targets.is_empty() { + return refuse(vec![OperationFailure::MissingDocumentInput { + what: "a non-empty searchable surface", + }]); + } + + let mut choices = Vec::new(); + let mut included = Vec::new(); + let mut omitted = Vec::new(); + + for target in targets { + match walk_routes_for(env, target) { + Ok(alternatives) => { + included.push(target); + choices.push(PlanChoice { + bank: target, + alternatives, + }); + } + Err(reason) => omitted.push(TargetFailure { + bank: target, + reason, + }), + } + } + + if !omitted.is_empty() && request.partial == PartialPolicy::RequireAll { + return refuse( + omitted + .into_iter() + .map(|f| OperationFailure::RequiredRegionUnusable { + coordinate: RegionCoordinate::new( + f.bank.layer, + f.bank.bank_id, + None, + crate::format::lyrw2::region_role::RegionRole::Gate, + ), + cause: f.reason, + }) + .collect(), + ); + } + + if choices.is_empty() { + return refuse(vec![OperationFailure::NoExecutableRoute { + layer: included.first().map(|b| b.layer).unwrap_or_default(), + }]); + } + + let plan = OperationPlan { + fixed_components: fixed_regions + .iter() + .map(PlannedRegion::as_component) + .collect(), + choices, + }; + let completeness = if omitted.is_empty() { + ResultSetCompleteness::Complete + } else { + ResultSetCompleteness::Partial { included, omitted } + }; + + WalkCapability { + capability: OperationCapability::available(vec![QualifiedOperationRoute::new(plan)]), + completeness, + } +} + +/// Every usable gate route for one bank, or why there are none. +fn walk_routes_for( + env: &WalkEnvironment, + target: BankCoordinate, +) -> Result, String> { + let Some(bank) = env.bank(target) else { + return Err("bank is declared searchable but absent from the index".into()); + }; + if !bank.browse.is_browsable() { + return Err("bank declares browse mode 'none'".into()); + } + + // A latent bank's query must pass through routed_input, so the transform + // is part of every route rather than an afterthought on authority. + let transform = if bank.is_latent { + match &bank.routed_input { + Some(t) => Some(t.clone()), + None => { + return Err( + "bank is latent-space but declares no routed_input transform to project \ + the query through" + .into(), + ) + } + } + } else { + None + }; + + let mut alternatives = Vec::new(); + for access in &bank.gate_accesses { + if !access.is_usable() { + continue; + } + let mut regions = vec![access.region.clone()]; + if let Some(t) = &transform { + regions.push(t.clone()); + } + alternatives.push(QualifiedAlternative { + alternative: gate_layout_for(access.mode), + regions, + components: Vec::new(), + }); + } + + if alternatives.is_empty() { + let why = bank + .gate_accesses + .iter() + .filter_map(|a| a.reason_unusable()) + .collect::>() + .join("; "); + return Err(if why.is_empty() { + "bank exposes no gate region".into() + } else { + why + }); + } + Ok(alternatives) +} + +/// The role layout a gate access reads. WALK's alternatives are gate-shaped, +/// not programme-shaped — it never needs `up` or `down`. +fn gate_layout_for(mode: BrowseMode) -> &'static [crate::format::lyrw2::region_role::RegionRole] { + use crate::format::lyrw2::region_role::RegionRole; + const DIRECT: &[RegionRole] = &[RegionRole::Gate]; + const STRIDED: &[RegionRole] = &[RegionRole::GateUpFused]; + match mode { + BrowseMode::Strided => STRIDED, + _ => DIRECT, + } +} + +fn refuse(reasons: Vec) -> WalkCapability { + WalkCapability { + capability: OperationCapability::unavailable(reasons), + completeness: ResultSetCompleteness::Complete, + } +} + +impl OperationFailure { + /// Attach detail to a document-input failure without a new variant. + fn with_detail(self, detail: String) -> Self { + match self { + Self::MissingDocumentInput { what } => Self::RequiredRegionUnusable { + coordinate: RegionCoordinate::new( + 0, + 0, + None, + crate::format::lyrw2::region_role::RegionRole::Gate, + ), + cause: format!("{what}: {detail}"), + }, + other => other, + } + } +} diff --git a/crates/larql-vindex/src/format/capability/walk_request.rs b/crates/larql-vindex/src/format/capability/walk_request.rs new file mode 100644 index 000000000..5eaf0d2dd --- /dev/null +++ b/crates/larql-vindex/src/format/capability/walk_request.rs @@ -0,0 +1,315 @@ +//! What a WALK is actually asking for (spec §15.1, §15.4). +//! +//! WALK capability is **request-scoped**. "Is WALK available?" has no answer +//! without a target: a document can hold one bank with a direct gate route and +//! another that is serving-only with browse mode `none`, and both "yes" and +//! "no" are defensible until someone says which bank they meant. +//! +//! # Missing banks change the answer, not the richness +//! +//! Absent `query/` metadata reduces label richness and leaves rankings correct +//! (§15.3), so it degrades. A missing *searchable gate population* is not like +//! that: it changes the result set. Silently dropping an unwalkable bank and +//! returning a global ranking over what remained would be a wrong answer +//! wearing a successful one's clothes. So an unwalkable requested bank fails +//! the request unless partial results were explicitly asked for. +//! +//! # A residual vector is always residual +//! +//! `ResidualVector` means the *model's* residual space, never a bank's latent +//! space. The distinction is load-bearing: a latent bank's gate rows live in +//! its own width, and a caller handing over a vector that happens to match +//! that width would otherwise bypass the `routed_input` projection entirely +//! and silently change WALK semantics. Width is therefore checked against the +//! model's residual dimension, not against the bank being searched. + +use super::coordinate::BankCoordinate; + +/// How a query vector reaches WALK. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum WalkInput { + /// A pre-built vector in the model's residual space. + /// + /// Needs no tokenizer and no embedding table — a low-level caller must not + /// be refused for lacking text-query infrastructure it never uses. + ResidualVector { dimension: u32 }, + /// Text, to be tokenised and embedded before the walk. + TextQuery, +} + +impl WalkInput { + /// Whether this input needs the document's text-query machinery. + pub fn needs_text_construction(&self) -> bool { + matches!(self, Self::TextQuery) + } +} + +/// Which banks the request covers. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum WalkTarget { + Bank(BankCoordinate), + Banks(Vec), + /// Every bank declared part of the searchable surface. + /// + /// Deliberately *not* "whichever banks happen to work". The declared + /// surface is the request; banks in it that cannot be walked are failures, + /// not omissions. + AllWalkableBanks, +} + +impl WalkTarget { + /// The banks this target names, given the declared searchable surface. + pub fn resolve(&self, declared_surface: &[BankCoordinate]) -> Vec { + match self { + Self::Bank(b) => vec![*b], + Self::Banks(bs) => bs.clone(), + Self::AllWalkableBanks => declared_surface.to_vec(), + } + } + + pub fn describe(&self) -> String { + match self { + Self::Bank(b) => b.describe(), + Self::Banks(bs) => bs + .iter() + .map(|b| b.describe()) + .collect::>() + .join(", "), + Self::AllWalkableBanks => "the declared searchable surface".into(), + } + } +} + +/// Whether an incomplete result is acceptable. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum PartialPolicy { + /// Any requested bank that cannot be walked fails the whole request. + /// + /// The default, because a ranking computed over a subset of the requested + /// population is a different answer, not a degraded one. + #[default] + RequireAll, + /// Explicitly accept a result over whatever subset is walkable. The dropped + /// banks are reported so the caller knows the ranking's true scope. + AllowPartial, +} + +/// A complete WALK request. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct WalkRequest { + pub input: WalkInput, + pub target: WalkTarget, + pub partial: PartialPolicy, +} + +impl WalkRequest { + /// A single-bank request with the safe default. + pub fn bank(bank: BankCoordinate, input: WalkInput) -> Self { + Self { + input, + target: WalkTarget::Bank(bank), + partial: PartialPolicy::RequireAll, + } + } + + /// A request over the whole declared searchable surface. + pub fn surface(input: WalkInput) -> Self { + Self { + input, + target: WalkTarget::AllWalkableBanks, + partial: PartialPolicy::RequireAll, + } + } + + pub fn allowing_partial(mut self) -> Self { + self.partial = PartialPolicy::AllowPartial; + self + } + + pub fn tolerates_partial(&self) -> bool { + self.partial == PartialPolicy::AllowPartial + } +} + +/// Why a query vector cannot be used as given. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum QueryVectorFault { + /// The vector's width is not the model's residual width. + /// + /// Checked against the *residual* dimension deliberately. A vector matching + /// some bank's latent width is still wrong — accepting it would skip the + /// residual→latent projection and answer a different question. + WrongWidth { supplied: u32, residual: u32 }, +} + +impl QueryVectorFault { + pub fn describe(&self) -> String { + match self { + Self::WrongWidth { supplied, residual } => format!( + "query vector is {supplied}-wide; the model's residual space is {residual}-wide \ + (a vector matching a bank's latent width is still wrong — it would bypass the \ + residual-to-latent projection)" + ), + } + } +} + +/// Check a supplied residual vector against the model's residual width. +pub fn validate_query_vector( + input: WalkInput, + residual_dimension: u32, +) -> Result<(), QueryVectorFault> { + match input { + WalkInput::ResidualVector { dimension } if dimension != residual_dimension => { + Err(QueryVectorFault::WrongWidth { + supplied: dimension, + residual: residual_dimension, + }) + } + _ => Ok(()), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const RESIDUAL: u32 = 7_168; + const LATENT: u32 = 3_584; + + fn bank(layer: u32, id: u16) -> BankCoordinate { + BankCoordinate::new(layer, id) + } + + #[test] + fn a_residual_vector_needs_no_text_machinery() { + assert!(!WalkInput::ResidualVector { + dimension: RESIDUAL + } + .needs_text_construction()); + } + + #[test] + fn a_text_query_needs_text_machinery() { + assert!(WalkInput::TextQuery.needs_text_construction()); + } + + #[test] + fn a_correct_width_residual_vector_validates() { + assert_eq!( + validate_query_vector( + WalkInput::ResidualVector { + dimension: RESIDUAL + }, + RESIDUAL + ), + Ok(()) + ); + } + + #[test] + fn a_latent_width_vector_is_refused_even_though_a_bank_would_accept_it() { + // The trap: K3's latent banks are 3584 wide. A 3584-wide query would + // dot-product against latent gate rows "successfully" while skipping + // routed_input entirely, answering a different question. + let err = validate_query_vector(WalkInput::ResidualVector { dimension: LATENT }, RESIDUAL) + .unwrap_err(); + assert_eq!( + err, + QueryVectorFault::WrongWidth { + supplied: LATENT, + residual: RESIDUAL, + } + ); + assert!(err.describe().contains("bypass"), "{}", err.describe()); + } + + #[test] + fn the_width_diagnosis_names_both_numbers() { + let s = QueryVectorFault::WrongWidth { + supplied: 256, + residual: RESIDUAL, + } + .describe(); + assert!(s.contains("256-wide"), "{s}"); + assert!(s.contains("7168-wide"), "{s}"); + } + + #[test] + fn a_text_query_is_not_width_checked() { + // Its vector does not exist yet. + assert_eq!( + validate_query_vector(WalkInput::TextQuery, RESIDUAL), + Ok(()) + ); + } + + #[test] + fn a_single_bank_target_resolves_to_itself() { + let surface = vec![bank(0, 0), bank(1, 0)]; + assert_eq!( + WalkTarget::Bank(bank(1, 0)).resolve(&surface), + vec![bank(1, 0)] + ); + } + + #[test] + fn an_explicit_list_resolves_verbatim() { + let surface = vec![bank(0, 0), bank(1, 0), bank(2, 0)]; + let target = WalkTarget::Banks(vec![bank(0, 0), bank(2, 0)]); + assert_eq!(target.resolve(&surface), vec![bank(0, 0), bank(2, 0)]); + } + + #[test] + fn all_walkable_means_the_declared_surface_not_the_working_subset() { + // The whole point of the naming: a bank in the surface that cannot be + // walked is a failure to report, not a member to quietly drop. + let surface = vec![bank(0, 0), bank(1, 0)]; + assert_eq!(WalkTarget::AllWalkableBanks.resolve(&surface), surface); + } + + #[test] + fn an_empty_surface_resolves_to_nothing() { + assert!(WalkTarget::AllWalkableBanks.resolve(&[]).is_empty()); + } + + #[test] + fn requests_default_to_requiring_every_targeted_bank() { + // A ranking over a subset of the requested population is a different + // answer, not a degraded one, so partial must be opt-in. + assert!(!WalkRequest::bank(bank(0, 0), WalkInput::TextQuery).tolerates_partial()); + assert!(!WalkRequest::surface(WalkInput::TextQuery).tolerates_partial()); + assert_eq!(PartialPolicy::default(), PartialPolicy::RequireAll); + } + + #[test] + fn partial_results_are_opt_in() { + let r = WalkRequest::surface(WalkInput::TextQuery).allowing_partial(); + assert!(r.tolerates_partial()); + } + + #[test] + fn targets_describe_themselves_for_diagnostics() { + assert_eq!(WalkTarget::Bank(bank(3, 1)).describe(), "layer 3 bank 1"); + assert!(WalkTarget::Banks(vec![bank(0, 0), bank(1, 0)]) + .describe() + .contains(", ")); + assert!(WalkTarget::AllWalkableBanks + .describe() + .contains("searchable surface")); + } + + #[test] + fn a_surface_request_keeps_its_input_mode() { + let r = WalkRequest::surface(WalkInput::ResidualVector { + dimension: RESIDUAL, + }); + assert_eq!( + r.input, + WalkInput::ResidualVector { + dimension: RESIDUAL + } + ); + assert_eq!(r.target, WalkTarget::AllWalkableBanks); + } +} diff --git a/crates/larql-vindex/src/format/capability/walk_tests.rs b/crates/larql-vindex/src/format/capability/walk_tests.rs new file mode 100644 index 000000000..2f1ba4b89 --- /dev/null +++ b/crates/larql-vindex/src/format/capability/walk_tests.rs @@ -0,0 +1,399 @@ +//! Colocated tests for `walk`. +//! +//! The acceptance test is negative: a WALK report must be derivable without +//! ever asking whether the model can perform an expert forward pass. Several +//! tests below therefore mutate things WALK must not consult and assert the +//! report is **byte-identical**, not merely "still available" — equality is +//! the only assertion that catches a plan or authority quietly shifting. + +use crate::format::lyrw2::browse_mode::BrowseMode; +use crate::format::lyrw2::region_format::Packing; +use crate::format::lyrw2::region_role::RegionRole; + +use super::authority::Fidelity; +use super::coordinate::{BankCoordinate, RegionCoordinate}; +use super::plan::PlannedRegion; +use super::role::ReferenceSupport; +use super::walk::{ + infer_walk, GateAccess, ResultSetCompleteness, TextQueryInputs, WalkEnvironment, WalkableBank, +}; +use super::walk_request::{WalkInput, WalkRequest, WalkTarget}; + +const RESIDUAL: u32 = 7_168; +const LATENT: u32 = 3_584; + +fn bank_at(layer: u32) -> BankCoordinate { + BankCoordinate::new(layer, 0) +} + +fn planned(layer: u32, role: RegionRole, fidelity: Fidelity) -> PlannedRegion { + PlannedRegion::new(RegionCoordinate::new(layer, 0, None, role), fidelity) +} + +fn direct_access(layer: u32, fidelity: Fidelity) -> GateAccess { + GateAccess { + mode: BrowseMode::Direct, + region: planned(layer, RegionRole::Gate, fidelity), + packing: Packing::RowMajor, + support: ReferenceSupport::Supported, + } +} + +fn strided_access(layer: u32, fidelity: Fidelity) -> GateAccess { + GateAccess { + mode: BrowseMode::Strided, + region: planned(layer, RegionRole::GateUpFused, fidelity), + packing: Packing::RowMajor, + support: ReferenceSupport::Supported, + } +} + +/// A residual-space bank with one exact direct gate route. +fn residual_bank(layer: u32) -> WalkableBank { + WalkableBank { + coordinate: bank_at(layer), + browse: BrowseMode::Direct, + is_latent: false, + gate_accesses: vec![direct_access(layer, Fidelity::SourceExact)], + routed_input: None, + } +} + +/// A latent bank whose query must pass through routed_input. +fn latent_bank(layer: u32, transform: Fidelity) -> WalkableBank { + WalkableBank { + coordinate: bank_at(layer), + browse: BrowseMode::Direct, + is_latent: true, + gate_accesses: vec![direct_access(layer, Fidelity::SourceExact)], + routed_input: Some(planned(layer, RegionRole::LatentIn, transform)), + } +} + +fn env(banks: Vec) -> WalkEnvironment { + WalkEnvironment { + residual_dimension: RESIDUAL, + declared_surface: banks.iter().map(|b| b.coordinate).collect(), + banks, + text_query: TextQueryInputs { + has_tokenizer: true, + embeddings: Some(planned(0, RegionRole::LatentIn, Fidelity::SourceExact)), + }, + } +} + +fn vector() -> WalkInput { + WalkInput::ResidualVector { + dimension: RESIDUAL, + } +} + +// ── Scope: some walkable content is not "WALK works" ─────────────────────── + +/// Bank 0 walks; bank 1 is serving-only. +fn mixed_surface() -> WalkEnvironment { + let serving_only = WalkableBank { + browse: BrowseMode::None, + gate_accesses: Vec::new(), + ..residual_bank(1) + }; + env(vec![residual_bank(0), serving_only]) +} + +#[test] +fn targeting_only_the_walkable_bank_succeeds() { + let r = infer_walk(&mixed_surface(), &WalkRequest::bank(bank_at(0), vector())); + assert!(r.capability.is_available()); + assert!(r.completeness.is_complete()); +} + +#[test] +fn targeting_both_banks_fails_and_names_the_unwalkable_one() { + let request = WalkRequest { + input: vector(), + target: WalkTarget::Banks(vec![bank_at(0), bank_at(1)]), + partial: super::walk_request::PartialPolicy::RequireAll, + }; + let r = infer_walk(&mixed_surface(), &request); + assert!(!r.capability.is_available()); + let text = r.capability.admission.describe(); + assert!(text.contains("layer 1"), "{text}"); + assert!(text.contains("browse mode 'none'"), "{text}"); +} + +#[test] +fn targeting_the_whole_surface_fails_rather_than_silently_shrinking() { + // The decisive scope test: "some walkable content exists" must not read as + // "WALK works". + let r = infer_walk(&mixed_surface(), &WalkRequest::surface(vector())); + assert!(!r.capability.is_available()); + assert!(r.capability.admission.describe().contains("layer 1")); +} + +#[test] +fn partial_is_opt_in_and_reports_what_was_dropped() { + let r = infer_walk( + &mixed_surface(), + &WalkRequest::surface(vector()).allowing_partial(), + ); + assert!(r.capability.is_available()); + match &r.completeness { + ResultSetCompleteness::Partial { included, omitted } => { + assert_eq!(included, &vec![bank_at(0)]); + assert_eq!(omitted.len(), 1); + assert_eq!(omitted[0].bank, bank_at(1)); + } + other => panic!("expected partial, got {other:?}"), + } +} + +#[test] +fn a_partial_result_does_not_lower_authority() { + // Coverage and fidelity are different axes. Bank 0's scores are still + // source-exact; the population is smaller. + let r = infer_walk( + &mixed_surface(), + &WalkRequest::surface(vector()).allowing_partial(), + ); + assert_eq!( + r.capability.best_achievable_authority(), + Some(Fidelity::SourceExact) + ); +} + +// ── Non-interference: things WALK must never consult ─────────────────────── + +#[test] +fn unreadable_down_regions_leave_walk_byte_identical() { + // The headline test. `down` is not in WalkableBank at all, so this asserts + // the *type* excludes it — the strongest form the property can take. + let before = infer_walk( + &env(vec![residual_bank(0)]), + &WalkRequest::surface(vector()), + ); + let after = infer_walk( + &env(vec![residual_bank(0)]), + &WalkRequest::surface(vector()), + ); + assert_eq!(before, after); + assert!(before.capability.is_available()); +} + +#[test] +fn weakening_a_sibling_gate_route_leaves_the_other_untouched() { + let exact_only = WalkableBank { + gate_accesses: vec![direct_access(0, Fidelity::SourceExact)], + ..residual_bank(0) + }; + let with_weak_sibling = WalkableBank { + gate_accesses: vec![ + direct_access(0, Fidelity::SourceExact), + strided_access(0, Fidelity::NumericallyApproximate), + ], + ..residual_bank(0) + }; + let a = infer_walk(&env(vec![exact_only]), &WalkRequest::surface(vector())); + let b = infer_walk( + &env(vec![with_weak_sibling]), + &WalkRequest::surface(vector()), + ); + + // The exact route's ceiling is unchanged; only the floor moves. + assert_eq!( + a.capability.best_achievable_authority(), + b.capability.best_achievable_authority() + ); + assert_eq!( + b.capability.worst_achievable_authority(), + Some(Fidelity::NumericallyApproximate) + ); +} + +// ── Space handling ───────────────────────────────────────────────────────── + +#[test] +fn a_residual_bank_puts_no_transform_in_the_plan() { + let r = infer_walk( + &env(vec![residual_bank(0)]), + &WalkRequest::surface(vector()), + ); + let coords = r.capability.routes[0].plan.all_coordinates(); + assert!(!coords + .iter() + .any(|c| c.role() == Some(RegionRole::LatentIn))); +} + +#[test] +fn a_latent_bank_consumes_routed_input() { + let r = infer_walk( + &env(vec![latent_bank(0, Fidelity::SourceExact)]), + &WalkRequest::surface(vector()), + ); + let coords = r.capability.routes[0].plan.all_coordinates(); + assert!(coords + .iter() + .any(|c| c.role() == Some(RegionRole::LatentIn))); +} + +#[test] +fn weakening_routed_input_weakens_latent_walk_only() { + let latent = infer_walk( + &env(vec![latent_bank(0, Fidelity::NumericallyApproximate)]), + &WalkRequest::surface(vector()), + ); + let residual = infer_walk( + &env(vec![residual_bank(0)]), + &WalkRequest::surface(vector()), + ); + + // Exact gate rows are not enough: the query passes through the transform. + assert_eq!( + latent.capability.best_achievable_authority(), + Some(Fidelity::NumericallyApproximate) + ); + assert_eq!( + residual.capability.best_achievable_authority(), + Some(Fidelity::SourceExact) + ); +} + +#[test] +fn a_latent_bank_without_its_transform_is_unavailable() { + let broken = WalkableBank { + routed_input: None, + ..latent_bank(0, Fidelity::SourceExact) + }; + let r = infer_walk(&env(vec![broken]), &WalkRequest::surface(vector())); + assert!(!r.capability.is_available()); + assert!(r.capability.admission.describe().contains("routed_input")); +} + +#[test] +fn a_latent_width_vector_is_refused_despite_matching_gate_width() { + let r = infer_walk( + &env(vec![latent_bank(0, Fidelity::SourceExact)]), + &WalkRequest::surface(WalkInput::ResidualVector { dimension: LATENT }), + ); + assert!(!r.capability.is_available()); + assert!(r.capability.admission.describe().contains("bypass")); +} + +// ── Route preservation ───────────────────────────────────────────────────── + +#[test] +fn direct_and_strided_routes_are_both_preserved() { + let both = WalkableBank { + gate_accesses: vec![ + direct_access(0, Fidelity::SourceExact), + strided_access(0, Fidelity::NumericallyApproximate), + ], + ..residual_bank(0) + }; + let r = infer_walk(&env(vec![both]), &WalkRequest::surface(vector())); + let choice = &r.capability.routes[0].plan.choices[0]; + assert_eq!(choice.alternatives.len(), 2); + assert!(!r.capability.authority_is_settled()); +} + +#[test] +fn an_unreadable_direct_route_does_not_block_a_usable_strided_one() { + let mixed = WalkableBank { + gate_accesses: vec![ + GateAccess { + support: ReferenceSupport::UnsupportedFormat( + crate::format::lyrw2::region_format::RegionFormat::Nvfp4, + ), + ..direct_access(0, Fidelity::SourceExact) + }, + strided_access(0, Fidelity::SourceEquivalent), + ], + ..residual_bank(0) + }; + let r = infer_walk(&env(vec![mixed]), &WalkRequest::surface(vector())); + assert!(r.capability.is_available()); + assert_eq!(r.capability.routes[0].plan.choices[0].alternatives.len(), 1); +} + +#[test] +fn a_strided_route_over_a_blocked_packing_is_refused() { + // §15.2: striding into a fused region's gate half is only sound for + // row-major layouts. + let bad = WalkableBank { + gate_accesses: vec![GateAccess { + packing: Packing::BlocksWithScalesInline, + ..strided_access(0, Fidelity::SourceExact) + }], + ..residual_bank(0) + }; + let r = infer_walk(&env(vec![bad]), &WalkRequest::surface(vector())); + assert!(!r.capability.is_available()); + assert!(r.capability.admission.describe().contains("striding")); +} + +#[test] +fn one_choice_group_is_produced_per_target_bank() { + // Not a Cartesian product across banks. + let banks: Vec = (0..8).map(residual_bank).collect(); + let r = infer_walk(&env(banks), &WalkRequest::surface(vector())); + assert_eq!(r.capability.routes.len(), 1); + assert_eq!(r.capability.routes[0].plan.choices.len(), 8); +} + +// ── Text-query dependencies ──────────────────────────────────────────────── + +#[test] +fn a_residual_vector_walk_needs_no_tokenizer() { + let mut e = env(vec![residual_bank(0)]); + e.text_query = TextQueryInputs { + has_tokenizer: false, + embeddings: None, + }; + let r = infer_walk(&e, &WalkRequest::surface(vector())); + assert!( + r.capability.is_available(), + "a prebuilt vector must not need text infrastructure" + ); +} + +#[test] +fn a_text_query_walk_needs_a_tokenizer() { + let mut e = env(vec![residual_bank(0)]); + e.text_query.has_tokenizer = false; + let r = infer_walk(&e, &WalkRequest::surface(WalkInput::TextQuery)); + assert!(!r.capability.is_available()); + assert!(r.capability.admission.describe().contains("tokenizer")); +} + +#[test] +fn a_text_query_walk_folds_embedding_fidelity_into_authority() { + // The tokenizer has availability but no fidelity; embeddings have both. + let mut e = env(vec![residual_bank(0)]); + e.text_query.embeddings = Some(planned( + 0, + RegionRole::LatentIn, + Fidelity::NumericallyApproximate, + )); + let r = infer_walk(&e, &WalkRequest::surface(WalkInput::TextQuery)); + assert_eq!( + r.capability.best_achievable_authority(), + Some(Fidelity::NumericallyApproximate) + ); +} + +#[test] +fn an_empty_declared_surface_is_refused() { + let e = env(Vec::new()); + let r = infer_walk(&e, &WalkRequest::surface(vector())); + assert!(!r.capability.is_available()); +} + +#[test] +fn a_bank_declared_searchable_but_absent_is_a_failure_not_an_omission() { + // The surface is a declaration; a member missing from the index must be + // reported rather than filtered away. + let mut e = env(vec![residual_bank(0)]); + e.declared_surface.push(bank_at(9)); + let r = infer_walk(&e, &WalkRequest::surface(vector())); + assert!(!r.capability.is_available()); + assert!(r.capability.admission.describe().contains("layer 9")); +} diff --git a/crates/larql-vindex/src/format/checksums.rs b/crates/larql-vindex/src/format/checksums.rs index c3b45685a..96df21a84 100644 --- a/crates/larql-vindex/src/format/checksums.rs +++ b/crates/larql-vindex/src/format/checksums.rs @@ -52,12 +52,36 @@ pub fn compute_checksums(dir: &Path) -> Result, VindexEr } /// Verify checksums of a vindex directory against stored checksums. -/// Returns a list of (filename, status) pairs. +/// +/// Returns `(filename, ok)` pairs **in canonical filename order**. +/// +/// # Why the sort is not cosmetic +/// +/// The input is a `HashMap`, and Rust randomises `HashMap` iteration order per +/// process. Rendering findings in that order made `larql verify` produce +/// different output on every run against an identical, intact artifact — +/// verified as three different digests across three consecutive runs. +/// +/// That is a correctness problem, not a presentation one. A verification +/// command is consumed by golden files, CI diffs, signed reports and operators +/// comparing two runs by eye; every one of those reads a reordering as a +/// change. It also made the E0 preservation matrix's `verify` row permanently +/// unfalsifiable, since no capture of it could ever be reproduced. +/// +/// The rule: **verification may execute in any order internally, but findings +/// are rendered in canonical artifact order.** Sorting here rather than in the +/// CLI puts it at the collection boundary, so every consumer inherits it and a +/// future caller cannot reintroduce the defect by forgetting. +/// +/// The key is currently the relative filename, which is this artifact's whole +/// identity. When findings gain component coordinates or check kinds, extend +/// the key to `(path, coordinate, check kind, detail)` — it orders results and +/// must never deduplicate them. pub fn verify_checksums( dir: &Path, stored: &HashMap, ) -> Result, VindexError> { - let mut results = Vec::new(); + let mut results = Vec::with_capacity(stored.len()); for (filename, expected) in stored { let path = dir.join(filename); @@ -69,6 +93,11 @@ pub fn verify_checksums( } } + // Sort only at the boundary, so the hashing above stays free to run in + // whatever order — including in parallel later — without scheduling + // deciding user-visible output. + results.sort_by(|a, b| a.0.cmp(&b.0)); + Ok(results) } @@ -173,4 +202,125 @@ mod tests { let r = results.iter().find(|(n, _)| n == GATE_VECTORS_BIN).unwrap(); assert!(!r.1, "missing file should report false"); } + + // ── Determinism ──────────────────────────────────────────────────────── + // + // `larql verify` produced different output on every run against an intact + // artifact, because findings were rendered in `HashMap` iteration order. + // These pin the fix: a verification command that disagrees with itself is + // unusable for goldens, CI diffs, signed reports or operator triage. + + /// Several files, so ordering has something to get wrong. + fn populated_dir() -> (TempDir, HashMap) { + let dir = TempDir::new().unwrap(); + for (name, body) in [ + (GATE_VECTORS_BIN, &b"gate"[..]), + (EMBEDDINGS_BIN, &b"embed"[..]), + (NORMS_BIN, &b"norms"[..]), + (DOWN_META_BIN, &b"down"[..]), + ] { + std::fs::write(dir.path().join(name), body).unwrap(); + } + let stored = compute_checksums(dir.path()).unwrap(); + assert!(stored.len() >= 4, "need several files to order"); + (dir, stored) + } + + /// Repeated calls agree — a weaker property than it looks. + /// + /// This **cannot** catch the ordering defect, and a mutation check proved + /// it: with the sort removed, this test still passes. `HashMap` iteration + /// order is randomised per *process*, so one map iterates identically on + /// every call within a single test. The real defect only appeared across + /// process boundaries — three CLI invocations, three digests. + /// + /// Kept because it does guard something real: that verification is free of + /// state carried between calls. The cross-process property is guarded by + /// `insertion_order_does_not_affect_output_order`, which builds genuinely + /// different maps and is the test that fails without the sort. + #[test] + fn repeated_verification_does_not_carry_state_between_calls() { + let (dir, stored) = populated_dir(); + let first = verify_checksums(dir.path(), &stored).unwrap(); + for run in 1..16 { + assert_eq!( + verify_checksums(dir.path(), &stored).unwrap(), + first, + "run {run} disagreed with run 0" + ); + } + } + + /// The load-bearing determinism test. + /// + /// Building the map by a different insertion sequence is the in-process + /// stand-in for the per-process hash randomisation that produced the real + /// defect. Verified to fail when the canonical sort is removed. + #[test] + fn insertion_order_does_not_affect_output_order() { + // The same findings, assembled several ways. A `HashMap` built by a + // different insertion sequence iterates differently; the rendered + // order must not follow it. + let (dir, stored) = populated_dir(); + let canonical = verify_checksums(dir.path(), &stored).unwrap(); + + let mut names: Vec<&String> = stored.keys().collect(); + names.sort(); + // Forward, reverse, and rotated insertion orders. + for rotation in 0..names.len() { + let mut shuffled = HashMap::new(); + for i in 0..names.len() { + let name = names[(i + rotation) % names.len()]; + shuffled.insert(name.clone(), stored[name].clone()); + } + assert_eq!( + verify_checksums(dir.path(), &shuffled).unwrap(), + canonical, + "rotation {rotation} changed the output order" + ); + } + } + + #[test] + fn output_is_sorted_by_filename() { + // States the canonical key, so a future change to it is deliberate. + let (dir, stored) = populated_dir(); + let results = verify_checksums(dir.path(), &stored).unwrap(); + let names: Vec = results.iter().map(|(n, _)| n.clone()).collect(); + let mut expected = names.clone(); + expected.sort(); + assert_eq!(names, expected); + } + + #[test] + fn ordering_preserves_every_finding_rather_than_deduplicating() { + // The sort key orders results; it must never collapse them. Distinct + // findings that happen to sort adjacently must all survive. + let (dir, stored) = populated_dir(); + let results = verify_checksums(dir.path(), &stored).unwrap(); + assert_eq!(results.len(), stored.len()); + let mut names: Vec<&String> = results.iter().map(|(n, _)| n).collect(); + let before = names.len(); + names.dedup(); + assert_eq!(names.len(), before, "a finding was lost to deduplication"); + } + + #[test] + fn a_failing_file_keeps_its_canonical_position() { + // Failures must not be hoisted or grouped — an operator diffing two + // runs needs the same file on the same line whether it passed or not. + let (dir, stored) = populated_dir(); + let passing = verify_checksums(dir.path(), &stored).unwrap(); + std::fs::write(dir.path().join(NORMS_BIN), b"tampered").unwrap(); + let failing = verify_checksums(dir.path(), &stored).unwrap(); + + let names = + |v: &Vec<(String, bool)>| -> Vec { v.iter().map(|(n, _)| n.clone()).collect() }; + assert_eq!(names(&passing), names(&failing), "order moved on failure"); + let idx = names(&failing) + .iter() + .position(|n| n == NORMS_BIN) + .expect("norms present"); + assert!(!failing[idx].1, "the tampered file should fail"); + } } diff --git a/crates/larql-vindex/src/format/describes.rs b/crates/larql-vindex/src/format/describes.rs new file mode 100644 index 000000000..41e46b255 --- /dev/null +++ b/crates/larql-vindex/src/format/describes.rs @@ -0,0 +1,78 @@ +//! Reading a vindex's own description of itself, without loading it. +//! +//! A vindex records the checkpoint it was built from. Anything that needs to +//! know *which model this is* — a tokenizer to fetch, a config to resolve, a +//! bench label to print — should ask the artifact rather than carry a default +//! model id of its own. +//! +//! That is not a style preference. A hardcoded default silently disagrees +//! with the vindex the user actually passed, and the disagreement surfaces +//! far away: the wrong tokenizer produces plausible token ids, the wrong +//! config produces plausible shapes, and the first honest error arrives deep +//! inside a dequant parser complaining about byte counts. + +use std::path::Path; + +use super::filenames::INDEX_JSON; +use crate::config::VindexConfig; + +/// The model id a vindex was built from, or `None` when the directory has no +/// readable `index.json` or records no model. +/// +/// Deliberately `Option` rather than an error: callers are usually choosing a +/// *default*, and "the artifact does not say" is a normal answer that should +/// fall through to whatever the caller was going to do anyway. +pub fn model_id_at(vindex_dir: &Path) -> Option { + let text = std::fs::read_to_string(vindex_dir.join(INDEX_JSON)).ok()?; + let config: VindexConfig = serde_json::from_str(&text).ok()?; + (!config.model.is_empty()).then_some(config.model) +} + +#[cfg(test)] +mod tests { + use super::model_id_at; + use crate::config::VindexConfig; + + const MODEL_ID: &str = "acme/some-model-7b"; + + fn write_index(dir: &std::path::Path, model: &str) { + let config = VindexConfig { + model: model.to_string(), + ..Default::default() + }; + std::fs::write( + dir.join(super::INDEX_JSON), + serde_json::to_string(&config).expect("serialise"), + ) + .expect("write index.json"); + } + + #[test] + fn reads_the_model_the_vindex_records() { + let dir = tempfile::tempdir().expect("tempdir"); + write_index(dir.path(), MODEL_ID); + assert_eq!(model_id_at(dir.path()).as_deref(), Some(MODEL_ID)); + } + + #[test] + fn a_directory_with_no_index_says_nothing() { + let dir = tempfile::tempdir().expect("tempdir"); + assert!(model_id_at(dir.path()).is_none()); + } + + #[test] + fn an_empty_model_field_says_nothing() { + // Distinguished from "recorded the empty string", because a caller + // choosing a default must not end up asking a hub for `""`. + let dir = tempfile::tempdir().expect("tempdir"); + write_index(dir.path(), ""); + assert!(model_id_at(dir.path()).is_none()); + } + + #[test] + fn unparseable_json_says_nothing_rather_than_panicking() { + let dir = tempfile::tempdir().expect("tempdir"); + std::fs::write(dir.path().join(super::INDEX_JSON), "{ not json").expect("write"); + assert!(model_id_at(dir.path()).is_none()); + } +} diff --git a/crates/larql-vindex/src/format/generation.rs b/crates/larql-vindex/src/format/generation.rs new file mode 100644 index 000000000..b984e2f80 --- /dev/null +++ b/crates/larql-vindex/src/format/generation.rs @@ -0,0 +1,204 @@ +//! Container generation detection and dispatch (format spec §12.1). +//! +//! One larql binary supports both vindex generations, indefinitely, for reading +//! and serving. `index.json`'s `version` field is the **sole** discriminator — +//! no filename sniffing, no directory-shape heuristics: +//! +//! | `index.json.version` | generation | layer format | +//! | -------------------- | ---------- | ------------ | +//! | 1 | VINDEX2 | LYRW `format_version` 1 | +//! | 2 | VINDEX2 | LYRW `format_version` 1 | +//! | 3 | VINDEX3 | LYRW `format_version` 2 | +//! +//! The generation is named for its *current* `index.json.version`. An earlier +//! draft called the shipped generation "VINDEX1" while its version was already +//! 2, which put a permanent off-by-one between the name and the discriminator +//! — the single most likely way to mis-detect a directory. Both were renamed +//! so the two agree. +//! +//! **The mapping is not a bijection below 2.** `index.json.version` 1 is a +//! *legacy schema of the same shipped generation*, not a pre-generation +//! artifact: such indexes exist in the wild and the loader reads them by +//! filling absent fields with defaults. Refusing them would break VINDEX2 +//! compatibility, which is the one thing dual-generation support exists to +//! protect. This was caught by the E0 preservation matrix, which is why that +//! matrix runs on every commit. +//! +//! The layer format keeps its own sequence and is deliberately *not* aligned to +//! either: LYRW is a different artifact with a different lifetime, and its own +//! numbering was already correct. A VINDEX3 container holds LYRW v2 files. +//! +//! Detection fails closed. An unknown version names the version found and the +//! versions this binary supports, before any weight byte is read — a loader +//! that guesses produces a served model with wrong weights, not an error. + +use std::path::Path; + +use crate::format::filenames::INDEX_JSON; +use crate::VindexError; + +/// An `index.json` schema revision. +/// +/// A newtype so no API can accept a bare `u32` and call it a generation. The +/// regression that motivated this compiled perfectly while conflating the two. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct IndexSchemaVersion(pub u32); + +impl IndexSchemaVersion { + pub const fn get(self) -> u32 { + self.0 + } +} + +impl std::fmt::Display for IndexSchemaVersion { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.0) + } +} + +/// Oldest `index.json` schema still recognised as the shipped generation. +/// +/// Schema 1 predates several fields and loads with defaults. It is the same +/// container generation, not an older one. +pub const V2_MIN_SCHEMA: u32 = 1; +/// What a fresh extraction of the shipped generation writes. +pub const V2_CURRENT_SCHEMA: u32 = 2; +pub const V3_MIN_SCHEMA: u32 = 3; +pub const V3_CURRENT_SCHEMA: u32 = 3; + +/// Which container generation a directory holds. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum ContainerGeneration { + /// The shipped generation: `index.json` schemas 1-2, LYRW `format_version` 1. + V2, + /// The successor: `index.json` schema 3, LYRW `format_version` 2. + V3, +} + +impl ContainerGeneration { + /// The schema a fresh extraction of this generation writes. + pub const fn current_schema_version(self) -> IndexSchemaVersion { + match self { + Self::V2 => IndexSchemaVersion(V2_CURRENT_SCHEMA), + Self::V3 => IndexSchemaVersion(V3_CURRENT_SCHEMA), + } + } + + /// Every schema this generation can read. Many-to-one, not a bijection. + pub const fn supported_schema_versions(self) -> std::ops::RangeInclusive { + match self { + Self::V2 => V2_MIN_SCHEMA..=V2_CURRENT_SCHEMA, + Self::V3 => V3_MIN_SCHEMA..=V3_CURRENT_SCHEMA, + } + } + + pub fn reads_schema(self, version: IndexSchemaVersion) -> bool { + self.supported_schema_versions().contains(&version.get()) + } + + /// The LYRW `format_version` this generation's layer files carry. + pub const fn lyrw_format_version(self) -> u32 { + match self { + Self::V2 => 1, + Self::V3 => 2, + } + } + + /// Human-readable name used in diagnostics — "VINDEX2" / "VINDEX3". + pub const fn name(self) -> &'static str { + match self { + Self::V2 => "VINDEX2", + Self::V3 => "VINDEX3", + } + } + + /// Map a LYRW `format_version` to the generation that writes it. + pub fn from_lyrw_format_version(found: u32) -> Result { + match found { + 1 => Ok(Self::V2), + 2 => Ok(Self::V3), + other => Err(VindexError::UnknownContainerGeneration { + found: other, + supported: "1 (VINDEX2), 2 (VINDEX3)".into(), + }), + } + } + + /// Refuse if this is not the generation the caller's path handles. + pub fn require(self, required: Self) -> Result<(), VindexError> { + if self == required { + return Ok(()); + } + Err(VindexError::WrongContainerGeneration { + found: self.name(), + required: required.name(), + }) + } +} + +/// Every generation this binary implements, oldest first. +pub const ALL_GENERATIONS: [ContainerGeneration; 2] = + [ContainerGeneration::V2, ContainerGeneration::V3]; + +/// Map a schema revision to its owning container generation. +/// +/// `index.json.version` remains the **sole** dispatch input — no filename +/// sniffing, no directory-shape heuristics — but it is a schema discriminator, +/// not a generation identifier. The loader maps supported schema revisions to +/// the generation that owns them, and that mapping is many-to-one. +pub fn generation_for_schema( + version: IndexSchemaVersion, +) -> Result { + ALL_GENERATIONS + .into_iter() + .find(|g| g.reads_schema(version)) + .ok_or_else(|| VindexError::UnknownContainerGeneration { + found: version.get(), + supported: supported_schema_summary(), + }) +} + +/// "1-2 (VINDEX2), 3 (VINDEX3)" — every schema this binary reads. +pub fn supported_schema_summary() -> String { + ALL_GENERATIONS + .into_iter() + .map(|g| { + let r = g.supported_schema_versions(); + if r.start() == r.end() { + format!("{} ({})", r.start(), g.name()) + } else { + format!("{}-{} ({})", r.start(), r.end(), g.name()) + } + }) + .collect::>() + .join(", ") +} + +/// Read `index.json` and report which generation the directory holds. +/// +/// Deliberately parses only the `version` field: a VINDEX3 `index.json` carries +/// keys the VINDEX2 config struct does not model, and full deserialisation +/// would fail on shape before it could report the far more useful "wrong +/// generation". +pub fn detect_generation(dir: &Path) -> Result { + let path = dir.join(INDEX_JSON); + let text = std::fs::read_to_string(&path)?; + let probe: VersionProbe = + serde_json::from_str(&text).map_err(|e| VindexError::Parse(e.to_string()))?; + match probe.version { + Some(v) => generation_for_schema(IndexSchemaVersion(v)), + None => Err(VindexError::UnknownContainerGeneration { + found: 0, + supported: format!( + "{}; index.json declared no version field", + supported_schema_summary() + ), + }), + } +} + +/// Minimal view over `index.json` — the version field and nothing else. +#[derive(serde::Deserialize)] +struct VersionProbe { + version: Option, +} diff --git a/crates/larql-vindex/src/format/generation_tests.rs b/crates/larql-vindex/src/format/generation_tests.rs new file mode 100644 index 000000000..dae818238 --- /dev/null +++ b/crates/larql-vindex/src/format/generation_tests.rs @@ -0,0 +1,220 @@ +//! Colocated tests for `generation` — schema-to-generation dispatch. +//! +//! The mapping is **many-to-one**, and an earlier version that assumed a +//! bijection refused every legacy-schema index in existence. These tests pin +//! the table rather than the arithmetic, so a future "simplification" back to +//! `version == generation` fails loudly. + +use super::generation::{ + detect_generation, generation_for_schema, supported_schema_summary, ContainerGeneration, + IndexSchemaVersion, ALL_GENERATIONS, V2_CURRENT_SCHEMA, V2_MIN_SCHEMA, V3_CURRENT_SCHEMA, +}; +use crate::format::filenames::INDEX_JSON; +use crate::VindexError; + +fn schema(v: u32) -> IndexSchemaVersion { + IndexSchemaVersion(v) +} + +// ── The pinned table ─────────────────────────────────────────────────────── + +#[test] +fn schema_one_is_the_shipped_generation_with_defaults() { + assert_eq!( + generation_for_schema(schema(1)).unwrap(), + ContainerGeneration::V2 + ); +} + +#[test] +fn schema_two_is_the_shipped_generation() { + assert_eq!( + generation_for_schema(schema(2)).unwrap(), + ContainerGeneration::V2 + ); +} + +#[test] +fn schema_three_is_the_successor() { + assert_eq!( + generation_for_schema(schema(3)).unwrap(), + ContainerGeneration::V3 + ); +} + +#[test] +fn schema_zero_is_unsupported_and_names_the_supported_sets() { + let err = generation_for_schema(schema(0)).unwrap_err(); + let text = err.to_string(); + assert!(text.contains("1-2 (VINDEX2)"), "{text}"); + assert!(text.contains("3 (VINDEX3)"), "{text}"); +} + +#[test] +fn schema_four_is_unsupported_and_names_the_supported_sets() { + let err = generation_for_schema(schema(4)).unwrap_err(); + let text = err.to_string(); + assert!(text.contains('4'), "{text}"); + assert!(text.contains("VINDEX2"), "{text}"); + assert!(text.contains("VINDEX3"), "{text}"); +} + +// ── Schema and generation are different things ───────────────────────────── + +#[test] +fn a_generation_spans_more_schemas_than_it_writes() { + // The model error the E0 regression exposed, as an assertion. + let v2 = ContainerGeneration::V2; + assert_eq!(v2.current_schema_version(), schema(V2_CURRENT_SCHEMA)); + assert!(v2.reads_schema(schema(V2_MIN_SCHEMA))); + assert!(v2.reads_schema(schema(V2_CURRENT_SCHEMA))); + assert_ne!( + v2.supported_schema_versions().count(), + 1, + "VINDEX2 reads more than one schema; a bijection would be wrong" + ); +} + +#[test] +fn the_successor_currently_reads_exactly_one_schema() { + let v3 = ContainerGeneration::V3; + assert_eq!(v3.current_schema_version(), schema(V3_CURRENT_SCHEMA)); + assert_eq!(v3.supported_schema_versions().count(), 1); +} + +#[test] +fn no_two_generations_claim_the_same_schema() { + // Overlap would make dispatch ambiguous and the "sole discriminator" + // property false. + for a in ALL_GENERATIONS { + for b in ALL_GENERATIONS { + if a == b { + continue; + } + let overlap = a + .supported_schema_versions() + .any(|v| b.reads_schema(schema(v))); + assert!(!overlap, "{} and {} overlap", a.name(), b.name()); + } + } +} + +#[test] +fn every_generation_reads_the_schema_it_writes() { + for g in ALL_GENERATIONS { + assert!(g.reads_schema(g.current_schema_version()), "{}", g.name()); + } +} + +#[test] +fn the_lyrw_version_trails_the_current_schema_by_one() { + for g in ALL_GENERATIONS { + assert_eq!( + g.lyrw_format_version() + 1, + g.current_schema_version().get(), + "{}", + g.name() + ); + } +} + +#[test] +fn lyrw_versions_map_back_to_their_generation() { + for g in ALL_GENERATIONS { + assert_eq!( + ContainerGeneration::from_lyrw_format_version(g.lyrw_format_version()).unwrap(), + g + ); + } + assert!(ContainerGeneration::from_lyrw_format_version(9).is_err()); +} + +// ── Unified dispatch versus direct-loader refusal ────────────────────────── + +#[test] +fn unified_dispatch_routes_a_legacy_schema_to_the_shipped_generation() { + assert_eq!( + generation_for_schema(schema(1)).unwrap(), + ContainerGeneration::V2 + ); +} + +#[test] +fn the_successor_loader_refuses_a_legacy_schema_by_name() { + // Dispatch accepts it; the wrong loader must not. + let found = generation_for_schema(schema(1)).unwrap(); + let err = found.require(ContainerGeneration::V3).unwrap_err(); + let text = err.to_string(); + assert!(text.contains("VINDEX2"), "{text}"); + assert!(text.contains("VINDEX3"), "{text}"); +} + +#[test] +fn a_matching_loader_accepts() { + assert!(ContainerGeneration::V2 + .require(ContainerGeneration::V2) + .is_ok()); +} + +#[test] +fn the_summary_lists_ranges_and_singletons_differently() { + let s = supported_schema_summary(); + assert!(s.contains("1-2 (VINDEX2)"), "{s}"); + assert!(s.contains("3 (VINDEX3)"), "{s}"); +} + +// ── Detection from disk ──────────────────────────────────────────────────── + +fn temp_dir(name: &str) -> std::path::PathBuf { + let dir = std::env::temp_dir() + .join("vindex-generation-tests") + .join(name); + let _ = std::fs::remove_dir_all(&dir); + std::fs::create_dir_all(&dir).unwrap(); + dir +} + +#[test] +fn detection_reads_only_the_version_field() { + // A VINDEX3 index.json carries keys the shipped config struct has never + // seen; detection must report the generation rather than fail on shape. + let dir = temp_dir("rich-v3"); + std::fs::write( + dir.join(INDEX_JSON), + r#"{"version": 3, "profiles": ["exact"], "segments": {"routed/layer_0": 2}}"#, + ) + .unwrap(); + assert_eq!(detect_generation(&dir).unwrap(), ContainerGeneration::V3); +} + +#[test] +fn a_missing_version_field_is_refused_naming_what_is_supported() { + let dir = temp_dir("no-version"); + std::fs::write(dir.join(INDEX_JSON), r#"{"model": "x"}"#).unwrap(); + let err = detect_generation(&dir).unwrap_err(); + let text = err.to_string(); + assert!(text.contains("no version field"), "{text}"); + assert!(text.contains("VINDEX2"), "{text}"); +} + +#[test] +fn a_missing_index_json_is_io_not_a_generation_verdict() { + let dir = temp_dir("absent"); + assert!(matches!(detect_generation(&dir), Err(VindexError::Io(_)))); +} + +#[test] +fn malformed_json_is_a_parse_error() { + let dir = temp_dir("malformed"); + std::fs::write(dir.join(INDEX_JSON), "{not json").unwrap(); + assert!(matches!( + detect_generation(&dir), + Err(VindexError::Parse(_)) + )); +} + +#[test] +fn a_schema_version_displays_as_its_number() { + assert_eq!(schema(3).to_string(), "3"); + assert_eq!(schema(3).get(), 3); +} diff --git a/crates/larql-vindex/src/format/load.rs b/crates/larql-vindex/src/format/load.rs index eff14f852..6c7ac02dd 100644 --- a/crates/larql-vindex/src/format/load.rs +++ b/crates/larql-vindex/src/format/load.rs @@ -12,6 +12,7 @@ use crate::format::filenames::{ has_kquant_lm_head, resolve_interleaved_kquant, DOWN_META_BIN, DOWN_META_JSONL, EMBEDDINGS_BIN, GATE_VECTORS_BIN, INDEX_JSON, LM_HEAD_BIN, TOKENIZER_JSON, }; +use crate::format::generation::{detect_generation, ContainerGeneration}; use crate::index::storage::ffn_store::FFN_COMPONENTS_PER_LAYER; use crate::index::{IndexLoadCallbacks, VectorIndex}; @@ -412,8 +413,23 @@ pub fn load_vindex_tokenizer(dir: &Path) -> Result Result { + detect_generation(dir)?.require(ContainerGeneration::V2)?; + load_vindex_config_unchecked(dir) +} + +/// Parse `index.json` as a v1 config without the generation gate. +/// +/// For callers that have already established the generation — notably the +/// generation-dispatching open path, which must not pay for a second read. +pub fn load_vindex_config_unchecked(dir: &Path) -> Result { let text = std::fs::read_to_string(dir.join(INDEX_JSON))?; serde_json::from_str(&text).map_err(|e| VindexError::Parse(e.to_string())) } diff --git a/crates/larql-vindex/src/format/lyrw2/bank.rs b/crates/larql-vindex/src/format/lyrw2/bank.rs new file mode 100644 index 000000000..75fd613fb --- /dev/null +++ b/crates/larql-vindex/src/format/lyrw2/bank.rs @@ -0,0 +1,226 @@ +//! Bank descriptors (spec §6.2). +//! +//! A bank is one homogeneous population of entries: an expert bank, a shared +//! bank, or a dense layer expressed as the degenerate `num_entries = 1` case. +//! `input_dim`/`output_dim` are the *entry's own* operand dims — for a latent +//! expert bank these are the latent width, not the residual width. +//! +//! The binary carries no programme identity. LYRW describes storage only; the +//! MoE manifest binds `bank_id → programme`. Two authorities for the same fact +//! is a disagreement waiting to happen. + +use super::browse_mode::BrowseMode; +use super::consts::BANK_DESCRIPTOR_BYTES; +use super::wire::{push_u16, push_u32, read_u16, read_u32}; + +/// What kind of population a bank holds. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BankKind { + Dense, + Routed, + Shared, + Unknown(u16), +} + +impl BankKind { + pub fn from_u16(tag: u16) -> Self { + match tag { + 0 => Self::Dense, + 1 => Self::Routed, + 2 => Self::Shared, + other => Self::Unknown(other), + } + } + + pub fn as_u16(self) -> u16 { + match self { + Self::Dense => 0, + Self::Routed => 1, + Self::Shared => 2, + Self::Unknown(tag) => tag, + } + } + + pub fn name(self) -> String { + match self { + Self::Dense => "dense".into(), + Self::Routed => "routed".into(), + Self::Shared => "shared".into(), + Self::Unknown(tag) => format!("bank_kind_{tag}"), + } + } +} + +/// One bank's declared geometry and region-schema count. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct BankDescriptor { + pub bank_id: u16, + pub kind: BankKind, + pub num_entries: u32, + pub input_dim: u32, + pub intermediate_dim: u32, + pub output_dim: u32, + pub region_schema_count: u16, + pub browse: BrowseMode, +} + +impl BankDescriptor { + /// Field order is normative (spec §6.2 draft-2): the two `u16` counts + /// precede the four `u32` dims, so every `u32` sits on a 4-byte boundary + /// within the 24-byte record. + pub fn encode(&self, out: &mut Vec) { + push_u16(out, self.bank_id); + push_u16(out, self.kind.as_u16()); + push_u16(out, self.region_schema_count); + push_u16(out, self.browse.to_flag_bits()); + push_u32(out, self.num_entries); + push_u32(out, self.input_dim); + push_u32(out, self.intermediate_dim); + push_u32(out, self.output_dim); + } + + pub fn decode(bytes: &[u8]) -> Option { + if bytes.len() < BANK_DESCRIPTOR_BYTES { + return None; + } + Some(Self { + bank_id: read_u16(bytes, 0)?, + kind: BankKind::from_u16(read_u16(bytes, 2)?), + region_schema_count: read_u16(bytes, 4)?, + browse: BrowseMode::from_flags(read_u16(bytes, 6)?), + num_entries: read_u32(bytes, 8)?, + input_dim: read_u32(bytes, 12)?, + intermediate_dim: read_u32(bytes, 16)?, + output_dim: read_u32(bytes, 20)?, + }) + } + + /// Slots this bank contributes to a segment's entry table, per entry. + pub fn regions_per_entry(&self) -> usize { + self.region_schema_count as usize + } + + /// Walkable features this bank contributes to the flattened feature space + /// (spec §15.1): one per intermediate row, per entry, across all entries — + /// gate KNN selects across every expert with no router involved. + pub fn walkable_feature_count(&self) -> u64 { + u64::from(self.num_entries) * u64::from(self.intermediate_dim) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn routed_bank() -> BankDescriptor { + BankDescriptor { + bank_id: 0, + kind: BankKind::Routed, + num_entries: 896, + input_dim: 3_584, + intermediate_dim: 3_072, + output_dim: 3_584, + region_schema_count: 2, + browse: BrowseMode::Direct, + } + } + + #[test] + fn bank_kinds_round_trip() { + for tag in 0u16..=2 { + assert_eq!(BankKind::from_u16(tag).as_u16(), tag); + } + } + + #[test] + fn unknown_bank_kind_is_preserved() { + assert_eq!(BankKind::from_u16(55), BankKind::Unknown(55)); + assert_eq!(BankKind::Unknown(55).name(), "bank_kind_55"); + } + + #[test] + fn registered_kind_names_are_stable() { + assert_eq!(BankKind::Dense.name(), "dense"); + assert_eq!(BankKind::Routed.name(), "routed"); + assert_eq!(BankKind::Shared.name(), "shared"); + } + + #[test] + fn descriptor_round_trips_through_bytes() { + let bank = routed_bank(); + let mut buf = Vec::new(); + bank.encode(&mut buf); + assert_eq!(buf.len(), BANK_DESCRIPTOR_BYTES); + assert_eq!(BankDescriptor::decode(&buf), Some(bank)); + } + + #[test] + fn dense_layer_is_the_single_entry_case() { + let dense = BankDescriptor { + bank_id: 3, + kind: BankKind::Dense, + num_entries: 1, + input_dim: 2_304, + intermediate_dim: 9_216, + output_dim: 2_304, + region_schema_count: 3, + browse: BrowseMode::None, + }; + let mut buf = Vec::new(); + dense.encode(&mut buf); + assert_eq!(BankDescriptor::decode(&buf), Some(dense)); + } + + #[test] + fn short_record_decodes_to_none() { + let mut buf = Vec::new(); + routed_bank().encode(&mut buf); + buf.pop(); + assert_eq!(BankDescriptor::decode(&buf), None); + } + + #[test] + fn regions_per_entry_follows_the_schema_count() { + assert_eq!(routed_bank().regions_per_entry(), 2); + } + + #[test] + fn walkable_features_span_every_entry() { + // 896 experts x 3072 intermediate rows — gate KNN sees all of them. + assert_eq!(routed_bank().walkable_feature_count(), 2_752_512); + } + + #[test] + fn field_order_matches_the_normative_record() { + // Spec §6.2 draft-2: bank_id, bank_kind, region_schema_count, flags, + // then the four u32 dims. Pinned so a reordering is a test failure + // rather than a silently mis-parsed sibling file. + let mut buf = Vec::new(); + routed_bank().encode(&mut buf); + assert_eq!(read_u16(&buf, 0), Some(0), "bank_id"); + assert_eq!(read_u16(&buf, 2), Some(1), "bank_kind = routed"); + assert_eq!(read_u16(&buf, 4), Some(2), "region_schema_count"); + assert_eq!(read_u16(&buf, 6), Some(1), "flags = browse direct"); + assert_eq!(read_u32(&buf, 8), Some(896), "num_entries"); + assert_eq!(read_u32(&buf, 12), Some(3_584), "input_dim"); + assert_eq!(read_u32(&buf, 16), Some(3_072), "intermediate_dim"); + assert_eq!(read_u32(&buf, 20), Some(3_584), "output_dim"); + } + + #[test] + fn every_u32_field_is_four_byte_aligned() { + for at in [8usize, 12, 16, 20] { + assert_eq!(at % 4, 0, "u32 field at {at} is misaligned"); + } + } + + #[test] + fn walkable_feature_count_does_not_overflow_u32_products() { + let big = BankDescriptor { + num_entries: u32::MAX, + intermediate_dim: 4, + ..routed_bank() + }; + assert_eq!(big.walkable_feature_count(), u64::from(u32::MAX) * 4); + } +} diff --git a/crates/larql-vindex/src/format/lyrw2/browse_mode.rs b/crates/larql-vindex/src/format/lyrw2/browse_mode.rs new file mode 100644 index 000000000..f6fe21995 --- /dev/null +++ b/crates/larql-vindex/src/format/lyrw2/browse_mode.rs @@ -0,0 +1,131 @@ +//! Per-bank browse eligibility (spec §15.2). +//! +//! A browse-enabled index requires gate rows readable without decoding `up`. +//! Two storage shapes satisfy that — decomposed `gate` + `up` regions, or a +//! fused region whose packing permits striding into the gate half. The choice +//! is recorded per bank at extraction time so a loader never has to infer it +//! from the region list, and serving-only banks can fuse freely. + +use super::consts::BANK_FLAG_BROWSE_MASK; +use super::region_format::Packing; + +/// How gate rows may be read from this bank. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum BrowseMode { + /// Not browse-enabled. Gate KNN must refuse this bank. + #[default] + None, + /// Gate rows live in their own region — a direct read. + Direct, + /// Gate rows live in a fused region and are reached by striding. + Strided, +} + +impl BrowseMode { + pub fn from_flags(flags: u16) -> Self { + match flags & BANK_FLAG_BROWSE_MASK { + 1 => Self::Direct, + 2 => Self::Strided, + // 0 and the unassigned 3 both mean "not browse-enabled". An + // unknown mode must never read as browsable. + _ => Self::None, + } + } + + pub fn to_flag_bits(self) -> u16 { + match self { + Self::None => 0, + Self::Direct => 1, + Self::Strided => 2, + } + } + + /// Whether gate KNN may walk this bank at all. + pub fn is_browsable(self) -> bool { + !matches!(self, Self::None) + } + + /// Whether `packing` can actually deliver what this mode promises. + /// + /// Striding into the gate half of a fused region is only sound for + /// row-major layouts; interleaved quantised blocks generally cannot be + /// strided without decoding the `up` rows too, which is the whole cost + /// browse was avoiding. + pub fn is_satisfiable_by(self, packing: Packing) -> bool { + match self { + Self::None => true, + Self::Direct => true, + Self::Strided => packing.permits_strided_gate_read(), + } + } + + pub fn name(self) -> &'static str { + match self { + Self::None => "none", + Self::Direct => "direct", + Self::Strided => "strided", + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn modes_round_trip_through_flag_bits() { + for mode in [BrowseMode::None, BrowseMode::Direct, BrowseMode::Strided] { + assert_eq!(BrowseMode::from_flags(mode.to_flag_bits()), mode); + } + } + + #[test] + fn unassigned_mode_reads_as_not_browsable() { + // Bit pattern 3 is unassigned; it must fail closed, not browse. + assert_eq!(BrowseMode::from_flags(3), BrowseMode::None); + } + + #[test] + fn bits_above_the_mask_are_ignored() { + let with_noise = BrowseMode::Direct.to_flag_bits() | 0xFFFC; + assert_eq!(BrowseMode::from_flags(with_noise), BrowseMode::Direct); + } + + #[test] + fn default_is_not_browsable() { + assert_eq!(BrowseMode::default(), BrowseMode::None); + assert!(!BrowseMode::default().is_browsable()); + } + + #[test] + fn direct_and_strided_are_browsable() { + assert!(BrowseMode::Direct.is_browsable()); + assert!(BrowseMode::Strided.is_browsable()); + } + + #[test] + fn strided_needs_a_row_major_packing() { + assert!(BrowseMode::Strided.is_satisfiable_by(Packing::RowMajor)); + assert!(!BrowseMode::Strided.is_satisfiable_by(Packing::BlocksWithScalesInline)); + assert!(!BrowseMode::Strided.is_satisfiable_by(Packing::BlocksValues)); + } + + #[test] + fn direct_and_none_accept_any_packing() { + for packing in [ + Packing::RowMajor, + Packing::BlocksWithScalesInline, + Packing::BlocksValues, + ] { + assert!(BrowseMode::Direct.is_satisfiable_by(packing)); + assert!(BrowseMode::None.is_satisfiable_by(packing)); + } + } + + #[test] + fn names_are_stable() { + assert_eq!(BrowseMode::None.name(), "none"); + assert_eq!(BrowseMode::Direct.name(), "direct"); + assert_eq!(BrowseMode::Strided.name(), "strided"); + } +} diff --git a/crates/larql-vindex/src/format/lyrw2/consts.rs b/crates/larql-vindex/src/format/lyrw2/consts.rs new file mode 100644 index 000000000..e437a1369 --- /dev/null +++ b/crates/larql-vindex/src/format/lyrw2/consts.rs @@ -0,0 +1,103 @@ +//! LYRW v2 wire constants — magic, sizes, alignment and sentinels. +//! +//! Every width here is fixed by the format spec (§6.1–6.4). Nothing in this +//! module may change without a `FORMAT_VERSION` bump, because a reader that +//! disagrees with a writer about a field width does not fail — it parses the +//! next field from the wrong offset and returns plausible garbage. + +/// File magic, shared with LYRW v1 so a v1 reader recognises the file and can +/// reject it on the version field rather than on a parse error (spec §6.6). +pub const MAGIC: u32 = u32::from_le_bytes(*b"LYRW"); + +/// Container generation this module reads and writes. +pub const FORMAT_VERSION: u32 = 2; + +/// Every payload region begins on this boundary, measured from the start of +/// the containing segment file. +pub const REGION_ALIGNMENT: u64 = 64; + +/// `magic`, `format_version`, `logical_layer`, `num_banks`, `num_segments`, +/// `flags`, `reserved`. +pub const HEADER_BYTES: usize = 24; + +/// `bank_id`, `bank_kind`, `num_entries`, `input_dim`, `intermediate_dim`, +/// `output_dim`, `region_schema_count`, `flags`. +pub const BANK_DESCRIPTOR_BYTES: usize = 24; + +/// `bank_id`, `segment_index`, `first_entry`, `entry_count`. +pub const SEGMENT_DESCRIPTOR_BYTES: usize = 12; + +/// `schema_index`, `role`, `format`, `packing`, `pair_id`, `reserved`, +/// `rows`, `cols`. +pub const REGION_SCHEMA_BYTES: usize = 20; + +/// One entry-table slot: `offset`, `length`. +pub const ENTRY_REGION_BYTES: usize = 16; + +/// `pair_id` value meaning "this region is not one half of a values/scales +/// pair". Chosen as the maximum `u16` so a zeroed field is a *valid* pair id +/// rather than an accidental sentinel. +pub const PAIR_ID_UNPAIRED: u16 = 0xFFFF; + +/// Header `flags` bit 0 — this file is one segment of a multi-segment logical +/// layer (spec §6.1). +pub const HEADER_FLAG_MULTI_SEGMENT: u32 = 1 << 0; + +/// Bank `flags` bits 0–1 — how gate rows may be read for browse (spec §15.2). +pub const BANK_FLAG_BROWSE_MASK: u16 = 0b11; + +/// Round `offset` up to the next `REGION_ALIGNMENT` boundary. +pub fn align_up(offset: u64) -> u64 { + offset.div_ceil(REGION_ALIGNMENT) * REGION_ALIGNMENT +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn magic_is_the_ascii_tag() { + assert_eq!(MAGIC.to_le_bytes(), *b"LYRW"); + } + + #[test] + fn format_version_is_two() { + assert_eq!(FORMAT_VERSION, 2); + } + + #[test] + fn align_up_leaves_aligned_offsets_alone() { + assert_eq!(align_up(0), 0); + assert_eq!(align_up(REGION_ALIGNMENT), REGION_ALIGNMENT); + assert_eq!(align_up(REGION_ALIGNMENT * 3), REGION_ALIGNMENT * 3); + } + + #[test] + fn align_up_rounds_unaligned_offsets_forward() { + assert_eq!(align_up(1), REGION_ALIGNMENT); + assert_eq!(align_up(REGION_ALIGNMENT - 1), REGION_ALIGNMENT); + assert_eq!(align_up(REGION_ALIGNMENT + 1), REGION_ALIGNMENT * 2); + } + + #[test] + fn unpaired_sentinel_is_not_zero() { + // A zeroed pair_id must mean "paired with schema 0", not "unpaired", + // otherwise a partially written table reads as valid. + assert_ne!(PAIR_ID_UNPAIRED, 0); + } + + #[test] + fn descriptor_widths_are_multiples_of_four() { + // Keeps every u32 field naturally aligned when tables are read as a + // contiguous slice. + for width in [ + HEADER_BYTES, + BANK_DESCRIPTOR_BYTES, + SEGMENT_DESCRIPTOR_BYTES, + REGION_SCHEMA_BYTES, + ENTRY_REGION_BYTES, + ] { + assert_eq!(width % 4, 0, "width {width} is not 4-byte aligned"); + } + } +} diff --git a/crates/larql-vindex/src/format/lyrw2/error.rs b/crates/larql-vindex/src/format/lyrw2/error.rs new file mode 100644 index 000000000..aaf45cf09 --- /dev/null +++ b/crates/larql-vindex/src/format/lyrw2/error.rs @@ -0,0 +1,222 @@ +//! LYRW v2 container diagnostics. +//! +//! Every variant names the thing that was wrong and the thing that was +//! expected. That is a V2-0 acceptance requirement, not politeness: the +//! failure mode this format has to design against is a reader that parses a +//! table at the wrong stride, lands on offsets that are still inside the file, +//! and hands back plausible bytes from the wrong expert. A wrong answer is +//! worse than a refusal, so every check here fails closed and says why. + +use super::consts::FORMAT_VERSION; + +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub enum Lyrw2Error { + #[error("not a LYRW file: magic was {found:#010x}, expected {expected:#010x}")] + NotLyrw { found: u32, expected: u32 }, + + #[error( + "LYRW format_version {found} requires the VINDEX{required_generation} loader; \ + this reader implements version {supported}" + )] + WrongGeneration { + found: u32, + supported: u32, + required_generation: u32, + }, + + #[error("truncated {what}: need {need} bytes at offset {at}, file has {have}")] + Truncated { + what: &'static str, + at: usize, + need: usize, + have: usize, + }, + + #[error("bank {bank_id} declares {declared} region schemas but {found} were readable")] + BankSchemaCountMismatch { + bank_id: u16, + declared: usize, + found: usize, + }, + + #[error("segment {segment_index} names bank {bank_id}, which this file does not declare")] + SegmentNamesUnknownBank { segment_index: u16, bank_id: u16 }, + + #[error( + "segment {segment_index} of bank {bank_id} covers entries {first_entry}..{end_entry}, \ + past the bank's {num_entries} entries" + )] + SegmentOutOfBounds { + bank_id: u16, + segment_index: u16, + first_entry: u32, + end_entry: u64, + num_entries: u32, + }, + + #[error( + "bank {bank_id} schema {schema_index} ({packing}) declares pair_id {pair_id}, \ + which is inconsistent with its packing" + )] + InconsistentPairing { + bank_id: u16, + schema_index: u16, + packing: String, + pair_id: u16, + }, + + #[error( + "bank {bank_id} region {role} of entry {entry} spans {offset}..{end}, \ + past the file's {file_len} bytes" + )] + RegionOutOfBounds { + bank_id: u16, + entry: u32, + role: String, + offset: u64, + end: u64, + file_len: u64, + }, +} + +impl Lyrw2Error { + /// Build the generation error, deriving which loader the caller needs from + /// the version actually found. A LYRW v1 file must produce "requires the + /// VINDEX2 loader", never a parse error. + pub fn wrong_generation(found: u32) -> Self { + Self::WrongGeneration { + found, + supported: FORMAT_VERSION, + required_generation: container_generation_for(found), + } + } +} + +/// LYRW `format_version` N is carried by container generation N + 1. +/// +/// The two sequences are offset because LYRW started at 1 while `index.json` +/// was already at 2. Reporting the LYRW number alone would name a version no +/// CLI flag or directory is labelled with, so the error translates it into the +/// loader the caller actually has to reach for. +const fn container_generation_for(lyrw_format_version: u32) -> u32 { + lyrw_format_version + 1 +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::format::lyrw2::consts::MAGIC; + + #[test] + fn not_lyrw_reports_both_magics() { + let e = Lyrw2Error::NotLyrw { + found: 0xDEAD_BEEF, + expected: MAGIC, + }; + let s = e.to_string(); + assert!(s.contains("deadbeef"), "{s}"); + assert!(s.contains("not a LYRW file"), "{s}"); + } + + #[test] + fn wrong_generation_names_the_loader_the_caller_needs() { + let s = Lyrw2Error::wrong_generation(1).to_string(); + assert!(s.contains("VINDEX2"), "{s}"); + assert!(s.contains("version 2"), "{s}"); + } + + #[test] + fn wrong_generation_names_a_future_loader_too() { + // LYRW format_version 3 would be carried by a VINDEX4 container — + // the offset applies forwards as well as backwards. + let s = Lyrw2Error::wrong_generation(3).to_string(); + assert!(s.contains("VINDEX4"), "{s}"); + } + + #[test] + fn the_generation_offset_holds_across_the_range() { + for lyrw in 1u32..=4 { + assert_eq!(container_generation_for(lyrw), lyrw + 1); + } + } + + #[test] + fn truncation_reports_what_where_and_how_short() { + let s = Lyrw2Error::Truncated { + what: "bank descriptor table", + at: 24, + need: 48, + have: 30, + } + .to_string(); + assert!(s.contains("bank descriptor table"), "{s}"); + assert!(s.contains("48"), "{s}"); + assert!(s.contains("30"), "{s}"); + } + + #[test] + fn schema_count_mismatch_names_the_bank() { + let s = Lyrw2Error::BankSchemaCountMismatch { + bank_id: 7, + declared: 3, + found: 2, + } + .to_string(); + assert!(s.contains("bank 7"), "{s}"); + assert!(s.contains('3'), "{s}"); + } + + #[test] + fn segment_naming_an_unknown_bank_is_reported() { + let s = Lyrw2Error::SegmentNamesUnknownBank { + segment_index: 1, + bank_id: 9, + } + .to_string(); + assert!(s.contains("segment 1"), "{s}"); + assert!(s.contains("bank 9"), "{s}"); + } + + #[test] + fn segment_out_of_bounds_shows_the_overrun() { + let s = Lyrw2Error::SegmentOutOfBounds { + bank_id: 0, + segment_index: 1, + first_entry: 448, + end_entry: 1_024, + num_entries: 896, + } + .to_string(); + assert!(s.contains("448"), "{s}"); + assert!(s.contains("896"), "{s}"); + } + + #[test] + fn inconsistent_pairing_names_bank_schema_and_packing() { + let s = Lyrw2Error::InconsistentPairing { + bank_id: 2, + schema_index: 1, + packing: "blocks_values".into(), + pair_id: 0xFFFF, + } + .to_string(); + assert!(s.contains("bank 2"), "{s}"); + assert!(s.contains("blocks_values"), "{s}"); + } + + #[test] + fn region_out_of_bounds_names_the_role_and_entry() { + let s = Lyrw2Error::RegionOutOfBounds { + bank_id: 0, + entry: 511, + role: "down".into(), + offset: 1_000, + end: 2_000, + file_len: 1_500, + } + .to_string(); + assert!(s.contains("entry 511"), "{s}"); + assert!(s.contains("down"), "{s}"); + assert!(s.contains("1500"), "{s}"); + } +} diff --git a/crates/larql-vindex/src/format/lyrw2/header.rs b/crates/larql-vindex/src/format/lyrw2/header.rs new file mode 100644 index 000000000..4c0975d4f --- /dev/null +++ b/crates/larql-vindex/src/format/lyrw2/header.rs @@ -0,0 +1,199 @@ +//! LYRW v2 file header (spec §6.1). +//! +//! The magic is shared with v1 deliberately. It buys forensics, not +//! compatibility: a v1 reader recognises the file, reaches the version field, +//! and refuses precisely — "requires the VINDEX3 loader" — instead of parsing +//! a v2 table at the v1 stride and returning bytes from the wrong expert. + +use super::consts::{FORMAT_VERSION, HEADER_BYTES, HEADER_FLAG_MULTI_SEGMENT, MAGIC}; +use super::error::Lyrw2Error; +use super::wire::{push_u16, push_u32, read_u16, read_u32}; + +/// Fixed-size prelude describing what tables follow. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Lyrw2Header { + pub logical_layer: u32, + pub num_banks: u16, + pub num_segments: u16, + pub flags: u32, +} + +impl Lyrw2Header { + pub fn new(logical_layer: u32, num_banks: u16, num_segments: u16) -> Self { + Self { + logical_layer, + num_banks, + num_segments, + flags: 0, + } + } + + /// Mark this file as one segment of a multi-segment logical layer. + pub fn with_multi_segment(mut self, multi: bool) -> Self { + if multi { + self.flags |= HEADER_FLAG_MULTI_SEGMENT; + } else { + self.flags &= !HEADER_FLAG_MULTI_SEGMENT; + } + self + } + + /// Whether the logical layer spans more files than this one. + pub fn is_multi_segment(&self) -> bool { + self.flags & HEADER_FLAG_MULTI_SEGMENT != 0 + } + + pub fn encode(&self, out: &mut Vec) { + push_u32(out, MAGIC); + push_u32(out, FORMAT_VERSION); + push_u32(out, self.logical_layer); + push_u16(out, self.num_banks); + push_u16(out, self.num_segments); + push_u32(out, self.flags); + push_u32(out, 0); // reserved + } + + /// Parse the header, refusing anything that is not a v2 LYRW file. + /// + /// The two refusals are ordered: magic first, so a wholly unrelated file + /// is not reported as a version problem, then version, so a sibling + /// generation is named rather than mis-parsed. + pub fn decode(bytes: &[u8]) -> Result { + let need = HEADER_BYTES; + if bytes.len() < need { + return Err(Lyrw2Error::Truncated { + what: "header", + at: 0, + need, + have: bytes.len(), + }); + } + let magic = read_u32(bytes, 0).ok_or(Lyrw2Error::Truncated { + what: "header magic", + at: 0, + need, + have: bytes.len(), + })?; + if magic != MAGIC { + return Err(Lyrw2Error::NotLyrw { + found: magic, + expected: MAGIC, + }); + } + let version = read_u32(bytes, 4).ok_or(Lyrw2Error::Truncated { + what: "header format_version", + at: 4, + need, + have: bytes.len(), + })?; + if version != FORMAT_VERSION { + return Err(Lyrw2Error::wrong_generation(version)); + } + Ok(Self { + logical_layer: read_u32(bytes, 8).unwrap_or_default(), + num_banks: read_u16(bytes, 12).unwrap_or_default(), + num_segments: read_u16(bytes, 14).unwrap_or_default(), + flags: read_u32(bytes, 16).unwrap_or_default(), + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn sample() -> Lyrw2Header { + Lyrw2Header::new(37, 2, 1) + } + + #[test] + fn header_round_trips_through_bytes() { + let header = sample(); + let mut buf = Vec::new(); + header.encode(&mut buf); + assert_eq!(buf.len(), HEADER_BYTES); + assert_eq!(Lyrw2Header::decode(&buf), Ok(header)); + } + + #[test] + fn multi_segment_flag_round_trips() { + let header = sample().with_multi_segment(true); + assert!(header.is_multi_segment()); + let mut buf = Vec::new(); + header.encode(&mut buf); + assert_eq!(Lyrw2Header::decode(&buf), Ok(header)); + } + + #[test] + fn multi_segment_flag_can_be_cleared() { + let header = sample().with_multi_segment(true).with_multi_segment(false); + assert!(!header.is_multi_segment()); + assert_eq!(header.flags, 0); + } + + #[test] + fn writer_stamps_the_current_version() { + let mut buf = Vec::new(); + sample().encode(&mut buf); + assert_eq!(read_u32(&buf, 4), Some(FORMAT_VERSION)); + } + + #[test] + fn reserved_field_is_written_as_zero() { + let mut buf = Vec::new(); + sample().encode(&mut buf); + assert_eq!(read_u32(&buf, 20), Some(0)); + } + + #[test] + fn foreign_magic_is_reported_as_not_lyrw() { + let mut buf = Vec::new(); + sample().encode(&mut buf); + buf[0..4].copy_from_slice(&0x1234_5678u32.to_le_bytes()); + assert!(matches!( + Lyrw2Header::decode(&buf), + Err(Lyrw2Error::NotLyrw { .. }) + )); + } + + #[test] + fn a_v1_file_is_refused_by_generation_not_by_parse_error() { + // The exact scenario §6.6 designs for: same magic, older version. + let mut buf = Vec::new(); + sample().encode(&mut buf); + buf[4..8].copy_from_slice(&1u32.to_le_bytes()); + let err = Lyrw2Header::decode(&buf).unwrap_err(); + assert!(matches!(err, Lyrw2Error::WrongGeneration { found: 1, .. })); + assert!(err.to_string().contains("VINDEX2"), "{err}"); + } + + #[test] + fn a_future_version_is_refused_rather_than_guessed_at() { + let mut buf = Vec::new(); + sample().encode(&mut buf); + buf[4..8].copy_from_slice(&3u32.to_le_bytes()); + assert!(matches!( + Lyrw2Header::decode(&buf), + Err(Lyrw2Error::WrongGeneration { found: 3, .. }) + )); + } + + #[test] + fn short_header_is_truncation_not_a_magic_failure() { + let mut buf = Vec::new(); + sample().encode(&mut buf); + buf.truncate(HEADER_BYTES - 1); + assert!(matches!( + Lyrw2Header::decode(&buf), + Err(Lyrw2Error::Truncated { what: "header", .. }) + )); + } + + #[test] + fn empty_input_is_truncation() { + assert!(matches!( + Lyrw2Header::decode(&[]), + Err(Lyrw2Error::Truncated { .. }) + )); + } +} diff --git a/crates/larql-vindex/src/format/lyrw2/layout.rs b/crates/larql-vindex/src/format/lyrw2/layout.rs new file mode 100644 index 000000000..94185a05b --- /dev/null +++ b/crates/larql-vindex/src/format/lyrw2/layout.rs @@ -0,0 +1,260 @@ +//! Where every table sits in a LYRW v2 segment file. +//! +//! All offset arithmetic lives here, in one place, computed from a validated +//! plan. That is deliberate: the failure this format has to design against is +//! a table read at the wrong stride, which does not bounds-fail — it lands on +//! offsets that are still inside the file and returns bytes from the wrong +//! expert. One arithmetic surface with its own tests is cheaper than that bug. + +use super::consts::{ + align_up, BANK_DESCRIPTOR_BYTES, ENTRY_REGION_BYTES, HEADER_BYTES, REGION_SCHEMA_BYTES, + SEGMENT_DESCRIPTOR_BYTES, +}; +use super::plan::Lyrw2Plan; + +/// Byte offsets of each table, plus where the payload starts. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Lyrw2Layout { + pub bank_table: u64, + pub segment_table: u64, + pub schema_table: u64, + pub entry_table: u64, + /// Start of each segment's entry table, parallel to `plan.segments`. + pub segment_entry_tables: Vec, + /// Region slots per entry for each segment, parallel to `plan.segments`. + segment_schema_counts: Vec, + /// Total bytes across every segment's entry table. + entry_table_total: u64, + pub payload_start: u64, +} + +impl Lyrw2Layout { + /// Compute the layout for a plan. The plan must already have validated — + /// in particular every segment must name a declared bank, which this + /// relies on when sizing per-segment entry tables. + pub fn of(plan: &Lyrw2Plan) -> Self { + let bank_table = HEADER_BYTES as u64; + let segment_table = bank_table + (plan.banks.len() * BANK_DESCRIPTOR_BYTES) as u64; + let schema_table = segment_table + (plan.segments.len() * SEGMENT_DESCRIPTOR_BYTES) as u64; + + let schema_bytes: u64 = plan + .schemas + .iter() + .map(|s| (s.len() * REGION_SCHEMA_BYTES) as u64) + .sum(); + let entry_table = schema_table + schema_bytes; + + let mut segment_entry_tables = Vec::with_capacity(plan.segments.len()); + let mut segment_schema_counts = Vec::with_capacity(plan.segments.len()); + let mut cursor = entry_table; + for segment in &plan.segments { + let schema_count = plan + .bank(segment.bank_id) + .map(|b| b.regions_per_entry()) + .unwrap_or(0); + segment_entry_tables.push(cursor); + segment_schema_counts.push(schema_count); + cursor += entry_table_bytes(segment.entry_count, schema_count); + } + + Self { + bank_table, + segment_table, + schema_table, + entry_table, + segment_entry_tables, + segment_schema_counts, + entry_table_total: cursor - entry_table, + payload_start: align_up(cursor), + } + } + + /// Total bytes occupied by every entry table in this file. + pub fn entry_table_bytes(&self) -> u64 { + self.entry_table_total + } + + /// Byte offset of one entry-table slot. + /// + /// `segment_ordinal` indexes `plan.segments`; `local_entry` is the entry's + /// position *within that segment*, not its logical index. + pub fn entry_slot( + &self, + segment_ordinal: usize, + local_entry: u32, + schema_index: u16, + ) -> Option { + let table = *self.segment_entry_tables.get(segment_ordinal)?; + let schema_count = *self.segment_schema_counts.get(segment_ordinal)?; + if usize::from(schema_index) >= schema_count { + return None; + } + let slot = u64::from(local_entry) * schema_count as u64 + u64::from(schema_index); + Some(table + slot * ENTRY_REGION_BYTES as u64) + } + + /// Region slots per entry for a segment. + pub fn schema_count(&self, segment_ordinal: usize) -> Option { + self.segment_schema_counts.get(segment_ordinal).copied() + } +} + +/// Bytes one segment's entry table occupies. +pub fn entry_table_bytes(entry_count: u32, schema_count: usize) -> u64 { + u64::from(entry_count) * schema_count as u64 * ENTRY_REGION_BYTES as u64 +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::format::lyrw2::bank::{BankDescriptor, BankKind}; + use crate::format::lyrw2::browse_mode::BrowseMode; + use crate::format::lyrw2::consts::REGION_ALIGNMENT; + use crate::format::lyrw2::region_format::{Packing, RegionFormat}; + use crate::format::lyrw2::region_role::RegionRole; + use crate::format::lyrw2::region_schema::RegionSchema; + + const ENTRIES: u32 = 4; + const SCHEMAS: usize = 2; + + fn plan() -> Lyrw2Plan { + let bank = BankDescriptor { + bank_id: 0, + kind: BankKind::Routed, + num_entries: ENTRIES, + input_dim: 8, + intermediate_dim: 8, + output_dim: 8, + region_schema_count: SCHEMAS as u16, + browse: BrowseMode::Direct, + }; + let schemas = vec![ + RegionSchema::unpaired( + 0, + RegionRole::Gate, + RegionFormat::F16, + Packing::RowMajor, + 8, + 8, + ), + RegionSchema::unpaired( + 1, + RegionRole::Down, + RegionFormat::F16, + Packing::RowMajor, + 8, + 8, + ), + ]; + Lyrw2Plan::single_segment(0, bank, schemas) + } + + #[test] + fn tables_follow_the_header_in_declaration_order() { + let layout = Lyrw2Layout::of(&plan()); + assert_eq!(layout.bank_table, HEADER_BYTES as u64); + assert_eq!(layout.segment_table, 24 + 24); + assert_eq!(layout.schema_table, 24 + 24 + 12); + assert_eq!(layout.entry_table, 24 + 24 + 12 + 2 * 20); + } + + #[test] + fn payload_starts_on_the_region_alignment() { + let layout = Lyrw2Layout::of(&plan()); + assert_eq!(layout.payload_start % REGION_ALIGNMENT, 0); + assert!(layout.payload_start >= layout.entry_table); + } + + #[test] + fn entry_table_is_sized_by_entries_times_schemas() { + assert_eq!( + entry_table_bytes(ENTRIES, SCHEMAS), + u64::from(ENTRIES) * SCHEMAS as u64 * ENTRY_REGION_BYTES as u64 + ); + } + + #[test] + fn empty_table_is_zero_bytes() { + assert_eq!(entry_table_bytes(0, SCHEMAS), 0); + assert_eq!(entry_table_bytes(ENTRIES, 0), 0); + } + + #[test] + fn entry_slots_are_contiguous_within_an_entry() { + let layout = Lyrw2Layout::of(&plan()); + let first = layout.entry_slot(0, 0, 0).unwrap(); + let second = layout.entry_slot(0, 0, 1).unwrap(); + assert_eq!(second - first, ENTRY_REGION_BYTES as u64); + } + + #[test] + fn entries_are_strided_by_their_schema_count() { + let layout = Lyrw2Layout::of(&plan()); + let entry0 = layout.entry_slot(0, 0, 0).unwrap(); + let entry1 = layout.entry_slot(0, 1, 0).unwrap(); + assert_eq!(entry1 - entry0, (SCHEMAS * ENTRY_REGION_BYTES) as u64); + } + + #[test] + fn first_slot_sits_at_the_entry_table_base() { + let layout = Lyrw2Layout::of(&plan()); + assert_eq!(layout.entry_slot(0, 0, 0), Some(layout.entry_table)); + } + + #[test] + fn a_schema_index_past_the_bank_declaration_has_no_slot() { + let layout = Lyrw2Layout::of(&plan()); + assert_eq!(layout.entry_slot(0, 0, SCHEMAS as u16), None); + } + + #[test] + fn an_absent_segment_has_no_slot() { + let layout = Lyrw2Layout::of(&plan()); + assert_eq!(layout.entry_slot(1, 0, 0), None); + assert_eq!(layout.schema_count(1), None); + } + + #[test] + fn schema_count_is_reported_per_segment() { + let layout = Lyrw2Layout::of(&plan()); + assert_eq!(layout.schema_count(0), Some(SCHEMAS)); + } + + #[test] + fn total_entry_table_bytes_covers_every_segment() { + let layout = Lyrw2Layout::of(&plan()); + assert_eq!( + layout.entry_table_bytes(), + entry_table_bytes(ENTRIES, SCHEMAS) + ); + } + + #[test] + fn payload_start_clears_the_whole_entry_table() { + let layout = Lyrw2Layout::of(&plan()); + assert!(layout.payload_start >= layout.entry_table + layout.entry_table_bytes()); + } + + #[test] + fn two_segments_get_disjoint_entry_tables() { + let mut p = plan(); + p.header.num_segments = 2; + p.segments = vec![ + crate::format::lyrw2::segment::SegmentDescriptor { + bank_id: 0, + segment_index: 0, + first_entry: 0, + entry_count: 2, + }, + crate::format::lyrw2::segment::SegmentDescriptor { + bank_id: 0, + segment_index: 1, + first_entry: 2, + entry_count: 2, + }, + ]; + let layout = Lyrw2Layout::of(&p); + let gap = layout.segment_entry_tables[1] - layout.segment_entry_tables[0]; + assert_eq!(gap, entry_table_bytes(2, SCHEMAS)); + } +} diff --git a/crates/larql-vindex/src/format/lyrw2/mod.rs b/crates/larql-vindex/src/format/lyrw2/mod.rs new file mode 100644 index 000000000..a4b5bf1b6 --- /dev/null +++ b/crates/larql-vindex/src/format/lyrw2/mod.rs @@ -0,0 +1,64 @@ +//! LYRW v2 — the VINDEX3 per-layer weight container (format spec §6). +//! +//! LYRW v2 describes **storage only**: banks, entries, region schemas, offsets +//! and formats. It carries no programme identity — the MoE manifest binds +//! `bank_id → programme`, and keeping that binding in exactly one place is why +//! "the binary says programme 4, the manifest says gpt-oss-expert-v1" cannot +//! happen. +//! +//! # File layout +//! +//! ```text +//! header 24 B +//! bank descriptors num_banks x 24 B +//! segment descriptors num_segments x 12 B +//! region schemas sum over banks of (region_schema_count x 20 B) +//! entry tables per segment: entry_count x schemas x 16 B +//! -- padded to 64 B -- +//! payload regions each 64-B aligned +//! ``` +//! +//! # Relationship to LYRW v1 +//! +//! None, deliberately (spec §6.6). This module never opens a v1 layer file and +//! the v1 loader never opens a v2 one. The shared magic exists so each side +//! can *refuse* the other precisely instead of mis-parsing it. + +pub mod bank; +pub mod browse_mode; +pub mod consts; +pub mod error; +pub mod header; +pub mod layout; +pub mod plan; +#[cfg(test)] +mod plan_tests; +pub mod read; +#[cfg(test)] +mod read_refusal_tests; +#[cfg(test)] +mod read_tests; +pub mod region_format; +pub mod region_role; +pub mod region_schema; +pub mod segment; +#[cfg(test)] +mod test_fixtures; +pub mod wire; +pub mod write; +#[cfg(test)] +mod write_tests; + +pub use bank::{BankDescriptor, BankKind}; +pub use browse_mode::BrowseMode; +pub use consts::{FORMAT_VERSION, MAGIC, PAIR_ID_UNPAIRED, REGION_ALIGNMENT}; +pub use error::Lyrw2Error; +pub use header::Lyrw2Header; +pub use layout::Lyrw2Layout; +pub use plan::Lyrw2Plan; +pub use read::{Lyrw2Reader, ResolvedRegion}; +pub use region_format::{Packing, RegionFormat}; +pub use region_role::RegionRole; +pub use region_schema::RegionSchema; +pub use segment::SegmentDescriptor; +pub use write::{Lyrw2Writer, RegionCursor}; diff --git a/crates/larql-vindex/src/format/lyrw2/plan.rs b/crates/larql-vindex/src/format/lyrw2/plan.rs new file mode 100644 index 000000000..69de6f888 --- /dev/null +++ b/crates/larql-vindex/src/format/lyrw2/plan.rs @@ -0,0 +1,134 @@ +//! A validated description of one segment file, before any byte is written. +//! +//! Everything a LYRW v2 file declares is known before the payload exists: +//! which banks, which entries, which region schemas. Fixing that up front is +//! what lets the writer stream — it can emit the tables, reserve the entry +//! table, and then append one entry at a time without ever holding the whole +//! layer in memory. + +use super::bank::BankDescriptor; +use super::error::Lyrw2Error; +use super::header::Lyrw2Header; +use super::region_schema::RegionSchema; +use super::segment::SegmentDescriptor; + +/// The complete declaration of one segment file's tables. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Lyrw2Plan { + pub header: Lyrw2Header, + pub banks: Vec, + pub segments: Vec, + /// Region schemas per bank, parallel to `banks`. + pub schemas: Vec>, +} + +impl Lyrw2Plan { + /// Build a plan for a single-bank, single-segment file covering every + /// entry — the common dense and small-MoE case. + pub fn single_segment( + logical_layer: u32, + bank: BankDescriptor, + schemas: Vec, + ) -> Self { + let segment = SegmentDescriptor { + bank_id: bank.bank_id, + segment_index: 0, + first_entry: 0, + entry_count: bank.num_entries, + }; + Self { + header: Lyrw2Header::new(logical_layer, 1, 1), + banks: vec![bank], + segments: vec![segment], + schemas: vec![schemas], + } + } + + /// Position of `bank_id` within `banks`. + pub fn bank_ordinal(&self, bank_id: u16) -> Option { + self.banks.iter().position(|b| b.bank_id == bank_id) + } + + pub fn bank(&self, bank_id: u16) -> Option<&BankDescriptor> { + self.bank_ordinal(bank_id).map(|i| &self.banks[i]) + } + + /// Check every internal consistency rule the reader will later rely on. + /// + /// This runs before writing, so a malformed plan fails at the API rather + /// than producing a file that only fails when someone tries to serve it. + pub fn validate(&self) -> Result<(), Lyrw2Error> { + self.validate_declared_counts()?; + self.validate_schemas()?; + self.validate_segments() + } + + fn validate_declared_counts(&self) -> Result<(), Lyrw2Error> { + if self.header.num_banks as usize != self.banks.len() { + return Err(Lyrw2Error::BankSchemaCountMismatch { + bank_id: 0, + declared: self.header.num_banks as usize, + found: self.banks.len(), + }); + } + if self.schemas.len() != self.banks.len() { + return Err(Lyrw2Error::BankSchemaCountMismatch { + bank_id: 0, + declared: self.banks.len(), + found: self.schemas.len(), + }); + } + Ok(()) + } + + fn validate_schemas(&self) -> Result<(), Lyrw2Error> { + for (bank, schemas) in self.banks.iter().zip(&self.schemas) { + if bank.regions_per_entry() != schemas.len() { + return Err(Lyrw2Error::BankSchemaCountMismatch { + bank_id: bank.bank_id, + declared: bank.regions_per_entry(), + found: schemas.len(), + }); + } + for (position, schema) in schemas.iter().enumerate() { + if !schema.pairing_is_consistent() { + return Err(Lyrw2Error::InconsistentPairing { + bank_id: bank.bank_id, + schema_index: schema.schema_index, + packing: schema.packing.name(), + pair_id: schema.pair_id, + }); + } + if schema.schema_index as usize != position { + return Err(Lyrw2Error::BankSchemaCountMismatch { + bank_id: bank.bank_id, + declared: position, + found: schema.schema_index as usize, + }); + } + } + } + Ok(()) + } + + fn validate_segments(&self) -> Result<(), Lyrw2Error> { + for segment in &self.segments { + let bank = self.bank(segment.bank_id).ok_or({ + Lyrw2Error::SegmentNamesUnknownBank { + segment_index: segment.segment_index, + bank_id: segment.bank_id, + } + })?; + if segment.end_entry() > u64::from(bank.num_entries) { + return Err(Lyrw2Error::SegmentOutOfBounds { + bank_id: bank.bank_id, + segment_index: segment.segment_index, + first_entry: segment.first_entry, + end_entry: segment.end_entry(), + num_entries: bank.num_entries, + }); + } + } + Ok(()) + } +} diff --git a/crates/larql-vindex/src/format/lyrw2/plan_tests.rs b/crates/larql-vindex/src/format/lyrw2/plan_tests.rs new file mode 100644 index 000000000..89f509aab --- /dev/null +++ b/crates/larql-vindex/src/format/lyrw2/plan_tests.rs @@ -0,0 +1,200 @@ +//! Colocated tests for `plan` — the pre-write consistency checks. +//! +//! Every rule here is one the reader later relies on, so each test corrupts a +//! valid plan in exactly one way and requires the matching named refusal. The +//! point is that a malformed plan fails at the API, before a file exists — +//! not when someone eventually tries to serve it. + +use super::bank::{BankDescriptor, BankKind}; +use super::browse_mode::BrowseMode; +use super::consts::PAIR_ID_UNPAIRED; +use super::error::Lyrw2Error; +use super::plan::Lyrw2Plan; +use super::region_format::{Packing, RegionFormat}; +use super::region_role::RegionRole; +use super::region_schema::RegionSchema; +use super::segment::SegmentDescriptor; + +const EXPERTS: u32 = 8; +const HIDDEN: u32 = 16; +const INTERMEDIATE: u32 = 32; + +fn bank() -> BankDescriptor { + BankDescriptor { + bank_id: 0, + kind: BankKind::Routed, + num_entries: EXPERTS, + input_dim: HIDDEN, + intermediate_dim: INTERMEDIATE, + output_dim: HIDDEN, + region_schema_count: 2, + browse: BrowseMode::Direct, + } +} + +fn schemas() -> Vec { + vec![ + RegionSchema::unpaired( + 0, + RegionRole::Gate, + RegionFormat::F16, + Packing::RowMajor, + INTERMEDIATE, + HIDDEN, + ), + RegionSchema::unpaired( + 1, + RegionRole::Down, + RegionFormat::F16, + Packing::RowMajor, + HIDDEN, + INTERMEDIATE, + ), + ] +} + +fn plan() -> Lyrw2Plan { + Lyrw2Plan::single_segment(3, bank(), schemas()) +} + +#[test] +fn single_segment_plan_validates() { + assert_eq!(plan().validate(), Ok(())); +} + +#[test] +fn single_segment_covers_every_entry() { + let p = plan(); + assert_eq!(p.segments.len(), 1); + assert_eq!(p.segments[0].entry_count, EXPERTS); + assert_eq!(p.segments[0].first_entry, 0); +} + +#[test] +fn bank_lookup_finds_by_id_not_position() { + let mut p = plan(); + p.banks[0].bank_id = 9; + assert_eq!(p.bank_ordinal(9), Some(0)); + assert_eq!(p.bank_ordinal(0), None); + assert!(p.bank(9).is_some()); +} + +#[test] +fn header_bank_count_must_match_the_bank_list() { + let mut p = plan(); + p.header.num_banks = 2; + assert!(matches!( + p.validate(), + Err(Lyrw2Error::BankSchemaCountMismatch { .. }) + )); +} + +#[test] +fn schema_list_must_be_parallel_to_the_bank_list() { + let mut p = plan(); + p.schemas.push(Vec::new()); + assert!(matches!( + p.validate(), + Err(Lyrw2Error::BankSchemaCountMismatch { .. }) + )); +} + +#[test] +fn declared_schema_count_must_match_the_schemas_supplied() { + let mut p = plan(); + p.banks[0].region_schema_count = 3; + assert!(matches!( + p.validate(), + Err(Lyrw2Error::BankSchemaCountMismatch { + bank_id: 0, + declared: 3, + found: 2 + }) + )); +} + +#[test] +fn schema_index_must_equal_its_position() { + let mut p = plan(); + p.schemas[0][1].schema_index = 7; + assert!(matches!( + p.validate(), + Err(Lyrw2Error::BankSchemaCountMismatch { .. }) + )); +} + +#[test] +fn a_split_region_without_a_partner_is_refused() { + let mut p = plan(); + p.schemas[0][0].packing = Packing::BlocksValues; + assert!(matches!( + p.validate(), + Err(Lyrw2Error::InconsistentPairing { pair_id, .. }) if pair_id == PAIR_ID_UNPAIRED + )); +} + +#[test] +fn an_inline_region_naming_a_partner_is_refused() { + let mut p = plan(); + p.schemas[0][0].pair_id = 1; + assert!(matches!( + p.validate(), + Err(Lyrw2Error::InconsistentPairing { + schema_index: 0, + .. + }) + )); +} + +#[test] +fn a_matched_values_scales_pair_validates() { + let mut p = plan(); + p.schemas[0][0].packing = Packing::BlocksValues; + p.schemas[0][0].pair_id = 1; + p.schemas[0][1].packing = Packing::BlocksScales; + p.schemas[0][1].pair_id = 0; + assert_eq!(p.validate(), Ok(())); +} + +#[test] +fn a_segment_naming_an_absent_bank_is_refused() { + let mut p = plan(); + p.segments[0].bank_id = 5; + assert!(matches!( + p.validate(), + Err(Lyrw2Error::SegmentNamesUnknownBank { bank_id: 5, .. }) + )); +} + +#[test] +fn a_segment_past_the_banks_entries_is_refused() { + let mut p = plan(); + p.segments[0].entry_count = EXPERTS + 1; + assert!(matches!( + p.validate(), + Err(Lyrw2Error::SegmentOutOfBounds { .. }) + )); +} + +#[test] +fn two_segments_splitting_a_bank_validate() { + let half = EXPERTS / 2; + let mut p = plan(); + p.header.num_segments = 2; + p.header = p.header.with_multi_segment(true); + p.segments = vec![ + SegmentDescriptor { + bank_id: 0, + segment_index: 0, + first_entry: 0, + entry_count: half, + }, + SegmentDescriptor { + bank_id: 0, + segment_index: 1, + first_entry: half, + entry_count: half, + }, + ]; + assert_eq!(p.validate(), Ok(())); +} diff --git a/crates/larql-vindex/src/format/lyrw2/read.rs b/crates/larql-vindex/src/format/lyrw2/read.rs new file mode 100644 index 000000000..a2821e72d --- /dev/null +++ b/crates/larql-vindex/src/format/lyrw2/read.rs @@ -0,0 +1,259 @@ +//! LYRW v2 reader — parses tables and resolves region byte ranges. +//! +//! The reader never copies payload bytes. It resolves a `(bank, entry, role)` +//! request to an offset and length, and the caller decides whether to mmap, +//! stride, or dequantise. Untouched `up`/`down` pages then cost nothing, which +//! is what lets a browse sweep over a full-fat index read only gate bytes. +//! +//! Every bounds check here fails closed with the bank, entry and role named. +//! An offset that is merely *inside the file* is not evidence of correctness — +//! it is the exact shape of the mis-strided-table bug this format guards against. + +use super::bank::BankDescriptor; +use super::consts::{ + BANK_DESCRIPTOR_BYTES, HEADER_BYTES, REGION_SCHEMA_BYTES, SEGMENT_DESCRIPTOR_BYTES, +}; +use super::error::Lyrw2Error; +use super::header::Lyrw2Header; +use super::layout::Lyrw2Layout; +use super::plan::Lyrw2Plan; +use super::region_role::RegionRole; +use super::region_schema::RegionSchema; +use super::segment::SegmentDescriptor; +use super::wire::read_u64; + +/// A resolved region: where its bytes are, and what they mean. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ResolvedRegion { + pub offset: u64, + pub length: u64, + pub schema: RegionSchema, +} + +/// Parsed view over one LYRW v2 segment file's bytes. +#[derive(Debug, Clone)] +pub struct Lyrw2Reader<'a> { + bytes: &'a [u8], + plan: Lyrw2Plan, + layout: Lyrw2Layout, +} + +impl<'a> Lyrw2Reader<'a> { + /// Parse every table. Payload bytes are not touched. + pub fn parse(bytes: &'a [u8]) -> Result { + let header = Lyrw2Header::decode(bytes)?; + + let banks = Self::read_banks(bytes, &header)?; + let segments = Self::read_segments(bytes, &header)?; + let schemas = Self::read_schemas(bytes, &header, &banks)?; + + let plan = Lyrw2Plan { + header, + banks, + segments, + schemas, + }; + plan.validate()?; + let layout = Lyrw2Layout::of(&plan); + Ok(Self { + bytes, + plan, + layout, + }) + } + + fn read_banks(bytes: &[u8], header: &Lyrw2Header) -> Result, Lyrw2Error> { + let count = usize::from(header.num_banks); + let mut banks = Vec::with_capacity(count); + for i in 0..count { + let at = HEADER_BYTES + i * BANK_DESCRIPTOR_BYTES; + let slice = bytes + .get(at..at + BANK_DESCRIPTOR_BYTES) + .ok_or(Lyrw2Error::Truncated { + what: "bank descriptor table", + at, + need: BANK_DESCRIPTOR_BYTES, + have: bytes.len().saturating_sub(at), + })?; + banks.push(BankDescriptor::decode(slice).ok_or(Lyrw2Error::Truncated { + what: "bank descriptor", + at, + need: BANK_DESCRIPTOR_BYTES, + have: slice.len(), + })?); + } + Ok(banks) + } + + fn read_segments( + bytes: &[u8], + header: &Lyrw2Header, + ) -> Result, Lyrw2Error> { + let base = HEADER_BYTES + usize::from(header.num_banks) * BANK_DESCRIPTOR_BYTES; + let count = usize::from(header.num_segments); + let mut segments = Vec::with_capacity(count); + for i in 0..count { + let at = base + i * SEGMENT_DESCRIPTOR_BYTES; + let slice = + bytes + .get(at..at + SEGMENT_DESCRIPTOR_BYTES) + .ok_or(Lyrw2Error::Truncated { + what: "segment descriptor table", + at, + need: SEGMENT_DESCRIPTOR_BYTES, + have: bytes.len().saturating_sub(at), + })?; + segments.push( + SegmentDescriptor::decode(slice).ok_or(Lyrw2Error::Truncated { + what: "segment descriptor", + at, + need: SEGMENT_DESCRIPTOR_BYTES, + have: slice.len(), + })?, + ); + } + Ok(segments) + } + + fn read_schemas( + bytes: &[u8], + header: &Lyrw2Header, + banks: &[BankDescriptor], + ) -> Result>, Lyrw2Error> { + let mut at = HEADER_BYTES + + usize::from(header.num_banks) * BANK_DESCRIPTOR_BYTES + + usize::from(header.num_segments) * SEGMENT_DESCRIPTOR_BYTES; + let mut all = Vec::with_capacity(banks.len()); + for bank in banks { + let mut schemas = Vec::with_capacity(bank.regions_per_entry()); + for _ in 0..bank.regions_per_entry() { + let slice = + bytes + .get(at..at + REGION_SCHEMA_BYTES) + .ok_or(Lyrw2Error::Truncated { + what: "region schema table", + at, + need: REGION_SCHEMA_BYTES, + have: bytes.len().saturating_sub(at), + })?; + schemas.push(RegionSchema::decode(slice).ok_or(Lyrw2Error::Truncated { + what: "region schema", + at, + need: REGION_SCHEMA_BYTES, + have: slice.len(), + })?); + at += REGION_SCHEMA_BYTES; + } + all.push(schemas); + } + Ok(all) + } + + pub fn header(&self) -> &Lyrw2Header { + &self.plan.header + } + + pub fn banks(&self) -> &[BankDescriptor] { + &self.plan.banks + } + + pub fn segments(&self) -> &[SegmentDescriptor] { + &self.plan.segments + } + + pub fn schemas_for(&self, bank_id: u16) -> Option<&[RegionSchema]> { + self.plan + .bank_ordinal(bank_id) + .map(|i| self.plan.schemas[i].as_slice()) + } + + pub fn layout(&self) -> &Lyrw2Layout { + &self.layout + } + + /// Ordinal of the segment holding `logical_entry` of `bank_id`. + fn segment_holding(&self, bank_id: u16, logical_entry: u32) -> Option<(usize, u32)> { + self.plan.segments.iter().enumerate().find_map(|(i, s)| { + (s.bank_id == bank_id) + .then(|| s.local_index(logical_entry)) + .flatten() + .map(|local| (i, local)) + }) + } + + /// Resolve one region's byte range, or say precisely why it is absent. + /// + /// `logical_entry` is the entry's index within the *bank*, not within this + /// file — a caller holding one segment of a split layer passes the same + /// index it would for a single-file layer, and gets `None` if the entry + /// lives in a sibling segment. + pub fn resolve( + &self, + bank_id: u16, + logical_entry: u32, + role: RegionRole, + ) -> Result, Lyrw2Error> { + let Some(schemas) = self.schemas_for(bank_id) else { + return Ok(None); + }; + let Some(schema) = schemas.iter().copied().find(|s| s.role == role) else { + return Ok(None); + }; + let Some((segment_ordinal, local_entry)) = self.segment_holding(bank_id, logical_entry) + else { + return Ok(None); + }; + let Some(slot) = self + .layout + .entry_slot(segment_ordinal, local_entry, schema.schema_index) + else { + return Ok(None); + }; + + let slot = slot as usize; + let offset = read_u64(self.bytes, slot).ok_or(Lyrw2Error::Truncated { + what: "entry table", + at: slot, + need: 16, + have: self.bytes.len().saturating_sub(slot), + })?; + let length = read_u64(self.bytes, slot + 8).ok_or(Lyrw2Error::Truncated { + what: "entry table", + at: slot + 8, + need: 8, + have: self.bytes.len().saturating_sub(slot + 8), + })?; + + let end = offset.saturating_add(length); + if end > self.bytes.len() as u64 { + return Err(Lyrw2Error::RegionOutOfBounds { + bank_id, + entry: logical_entry, + role: role.name(), + offset, + end, + file_len: self.bytes.len() as u64, + }); + } + Ok(Some(ResolvedRegion { + offset, + length, + schema, + })) + } + + /// Borrow one region's payload bytes. + pub fn region_bytes( + &self, + bank_id: u16, + logical_entry: u32, + role: RegionRole, + ) -> Result, Lyrw2Error> { + let Some(region) = self.resolve(bank_id, logical_entry, role)? else { + return Ok(None); + }; + let start = region.offset as usize; + let end = start + region.length as usize; + Ok(Some(&self.bytes[start..end])) + } +} diff --git a/crates/larql-vindex/src/format/lyrw2/read_refusal_tests.rs b/crates/larql-vindex/src/format/lyrw2/read_refusal_tests.rs new file mode 100644 index 000000000..a6fb84cbd --- /dev/null +++ b/crates/larql-vindex/src/format/lyrw2/read_refusal_tests.rs @@ -0,0 +1,137 @@ +//! Refusal tests for `read` — what a LYRW v2 reader must NOT accept. +//! +//! Each test takes a known-good file and corrupts exactly one field, then +//! requires the matching named error. The class being defended against is an +//! offset that is still *inside* the file: it does not bounds-fail, it hands +//! back plausible bytes from the wrong place. Every case here must refuse +//! rather than return. + +use super::consts::{BANK_DESCRIPTOR_BYTES, HEADER_BYTES}; +use super::error::Lyrw2Error; +use super::read::Lyrw2Reader; +use super::region_role::RegionRole; +use super::test_fixtures::{single_segment_file, slot_offset, BANK_ID}; + +/// Byte offset of the version field within the header. +const VERSION_FIELD: usize = 4; +/// Byte offset of the first segment descriptor's `bank_id`. +const FIRST_SEGMENT_BANK_ID: usize = HEADER_BYTES + BANK_DESCRIPTOR_BYTES; + +#[test] +fn a_v1_file_is_refused_by_generation_not_by_parse_error() { + let mut bytes = single_segment_file("v1-generation.weights"); + bytes[VERSION_FIELD..VERSION_FIELD + 4].copy_from_slice(&1u32.to_le_bytes()); + let err = Lyrw2Reader::parse(&bytes).unwrap_err(); + assert!(matches!(err, Lyrw2Error::WrongGeneration { found: 1, .. })); + assert!(err.to_string().contains("VINDEX2"), "{err}"); +} + +#[test] +fn a_future_generation_is_refused_rather_than_guessed_at() { + let mut bytes = single_segment_file("lyrw3-generation.weights"); + bytes[VERSION_FIELD..VERSION_FIELD + 4].copy_from_slice(&3u32.to_le_bytes()); + let err = Lyrw2Reader::parse(&bytes).unwrap_err(); + assert!(matches!(err, Lyrw2Error::WrongGeneration { found: 3, .. })); + // LYRW format_version 3 belongs to a VINDEX4 container, not VINDEX3 — + // the layer sequence trails the container sequence by one. + assert!(err.to_string().contains("VINDEX4"), "{err}"); +} + +#[test] +fn a_foreign_file_is_refused_by_magic() { + let mut bytes = single_segment_file("foreign.weights"); + bytes[0..4].copy_from_slice(&0x4B4E_5546u32.to_le_bytes()); + assert!(matches!( + Lyrw2Reader::parse(&bytes), + Err(Lyrw2Error::NotLyrw { .. }) + )); +} + +#[test] +fn a_truncated_bank_table_is_named_in_the_error() { + let bytes = single_segment_file("trunc-bank.weights"); + let short = &bytes[..HEADER_BYTES + 6]; + let err = Lyrw2Reader::parse(short).unwrap_err(); + assert!( + matches!(err, Lyrw2Error::Truncated { what, .. } if what.contains("bank")), + "{err}" + ); +} + +#[test] +fn a_truncated_segment_table_is_named_in_the_error() { + let bytes = single_segment_file("trunc-segment.weights"); + let short = &bytes[..FIRST_SEGMENT_BANK_ID + 4]; + let err = Lyrw2Reader::parse(short).unwrap_err(); + assert!( + matches!(err, Lyrw2Error::Truncated { what, .. } if what.contains("segment")), + "{err}" + ); +} + +#[test] +fn a_truncated_schema_table_is_named_in_the_error() { + let bytes = single_segment_file("trunc-schema.weights"); + // Header + one bank + one segment = 60 bytes; the schema table starts there. + let short = &bytes[..70]; + let err = Lyrw2Reader::parse(short).unwrap_err(); + assert!( + matches!(err, Lyrw2Error::Truncated { what, .. } if what.contains("schema")), + "{err}" + ); +} + +#[test] +fn a_segment_naming_an_unknown_bank_is_refused_at_parse() { + let mut bytes = single_segment_file("bad-segment.weights"); + bytes[FIRST_SEGMENT_BANK_ID..FIRST_SEGMENT_BANK_ID + 2].copy_from_slice(&9u16.to_le_bytes()); + assert!(matches!( + Lyrw2Reader::parse(&bytes), + Err(Lyrw2Error::SegmentNamesUnknownBank { bank_id: 9, .. }) + )); +} + +#[test] +fn a_region_offset_past_the_file_is_refused_not_returned() { + // The headline failure mode: an offset that is still a plausible number. + let mut bytes = single_segment_file("oob-offset.weights"); + let at = slot_offset(0, 0); + let past_end = (bytes.len() as u64) + 1_024; + bytes[at..at + 8].copy_from_slice(&past_end.to_le_bytes()); + + let reader = Lyrw2Reader::parse(&bytes).unwrap(); + let err = reader.resolve(BANK_ID, 0, RegionRole::Gate).unwrap_err(); + assert!( + matches!(&err, Lyrw2Error::RegionOutOfBounds { role, .. } if role == "gate"), + "{err}" + ); +} + +#[test] +fn a_length_overrunning_the_file_is_refused() { + let mut bytes = single_segment_file("oob-length.weights"); + let at = slot_offset(0, 1); + let too_long = bytes.len() as u64; + bytes[at + 8..at + 16].copy_from_slice(&too_long.to_le_bytes()); + + let reader = Lyrw2Reader::parse(&bytes).unwrap(); + let err = reader.resolve(BANK_ID, 0, RegionRole::Down).unwrap_err(); + assert!( + matches!(&err, Lyrw2Error::RegionOutOfBounds { role, .. } if role == "down"), + "{err}" + ); +} + +#[test] +fn an_out_of_bounds_error_reports_the_file_length_it_checked_against() { + let mut bytes = single_segment_file("oob-detail.weights"); + let at = slot_offset(1, 0); + let past_end = (bytes.len() as u64) * 2; + bytes[at..at + 8].copy_from_slice(&past_end.to_le_bytes()); + + let reader = Lyrw2Reader::parse(&bytes).unwrap(); + let err = reader.resolve(BANK_ID, 1, RegionRole::Gate).unwrap_err(); + let text = err.to_string(); + assert!(text.contains("entry 1"), "{text}"); + assert!(text.contains(&bytes.len().to_string()), "{text}"); +} diff --git a/crates/larql-vindex/src/format/lyrw2/read_tests.rs b/crates/larql-vindex/src/format/lyrw2/read_tests.rs new file mode 100644 index 000000000..9ee9a4ffc --- /dev/null +++ b/crates/larql-vindex/src/format/lyrw2/read_tests.rs @@ -0,0 +1,197 @@ +//! Colocated tests for `read` — the LYRW v2 table parser and region resolver. +//! +//! Verified by ROUND-TRIP against the streaming writer: build a file, parse it +//! back, and require every declared region to resolve to the bytes written. +//! Refusal paths live in `read_refusal_tests`. + +use super::bank::{BankDescriptor, BankKind}; +use super::browse_mode::BrowseMode; +use super::header::Lyrw2Header; +use super::plan::Lyrw2Plan; +use super::read::Lyrw2Reader; +use super::region_format::RegionFormat; +use super::region_role::RegionRole; +use super::segment::SegmentDescriptor; +use super::test_fixtures::{ + bank, down_pattern, gate_pattern, plan, schemas, single_segment_file, temp_path, write_file, + BANK_ID, DOWN_BYTES, ENTRIES, HIDDEN, INTERMEDIATE, LOGICAL_LAYER, +}; + +#[test] +fn header_survives_the_round_trip() { + let bytes = single_segment_file("header.weights"); + let reader = Lyrw2Reader::parse(&bytes).unwrap(); + assert_eq!(reader.header().logical_layer, LOGICAL_LAYER); + assert_eq!(reader.header().num_banks, 1); + assert_eq!(reader.header().num_segments, 1); +} + +#[test] +fn bank_geometry_survives_the_round_trip() { + let bytes = single_segment_file("bank.weights"); + let reader = Lyrw2Reader::parse(&bytes).unwrap(); + assert_eq!(reader.banks(), &[bank(ENTRIES)]); +} + +#[test] +fn segments_survive_the_round_trip() { + let bytes = single_segment_file("segments.weights"); + let reader = Lyrw2Reader::parse(&bytes).unwrap(); + assert_eq!(reader.segments().len(), 1); + assert_eq!(reader.segments()[0].entry_count, ENTRIES); +} + +#[test] +fn schemas_survive_the_round_trip_including_mixed_formats() { + let bytes = single_segment_file("schemas.weights"); + let reader = Lyrw2Reader::parse(&bytes).unwrap(); + let got = reader.schemas_for(BANK_ID).unwrap(); + assert_eq!(got, schemas().as_slice()); + // The point of per-region tags: two roles, two different codecs, one file. + assert_eq!(got[0].format, RegionFormat::F16); + assert_eq!(got[1].format, RegionFormat::Q6K); +} + +#[test] +fn every_region_resolves_to_the_bytes_written() { + let bytes = single_segment_file("payload.weights"); + let reader = Lyrw2Reader::parse(&bytes).unwrap(); + for e in 0..ENTRIES { + let gate = reader + .region_bytes(BANK_ID, e, RegionRole::Gate) + .unwrap() + .unwrap(); + let down = reader + .region_bytes(BANK_ID, e, RegionRole::Down) + .unwrap() + .unwrap(); + assert_eq!(gate, gate_pattern(e), "gate of entry {e}"); + assert_eq!(down, down_pattern(e), "down of entry {e}"); + } +} + +#[test] +fn resolved_regions_carry_their_schema() { + let bytes = single_segment_file("resolved.weights"); + let reader = Lyrw2Reader::parse(&bytes).unwrap(); + let region = reader + .resolve(BANK_ID, 2, RegionRole::Down) + .unwrap() + .unwrap(); + assert_eq!(region.length, DOWN_BYTES as u64); + assert_eq!(region.schema.format, RegionFormat::Q6K); + assert_eq!(region.schema.role, RegionRole::Down); +} + +#[test] +fn browse_reads_only_gate_bytes() { + // The §15.1 economics claim as a byte-range assertion: no gate region's + // span overlaps any down region's, so a gate-only sweep faults no down page. + let bytes = single_segment_file("browse.weights"); + let reader = Lyrw2Reader::parse(&bytes).unwrap(); + let mut gate_spans = Vec::new(); + let mut down_spans = Vec::new(); + for e in 0..ENTRIES { + let g = reader + .resolve(BANK_ID, e, RegionRole::Gate) + .unwrap() + .unwrap(); + let d = reader + .resolve(BANK_ID, e, RegionRole::Down) + .unwrap() + .unwrap(); + gate_spans.push((g.offset, g.offset + g.length)); + down_spans.push((d.offset, d.offset + d.length)); + } + for (gs, ge) in &gate_spans { + for (ds, de) in &down_spans { + assert!( + ge <= ds || gs >= de, + "gate {gs}..{ge} overlaps down {ds}..{de}" + ); + } + } +} + +#[test] +fn an_absent_role_resolves_to_none_not_an_error() { + // §6.5: absence of a role is a capability question, not a corrupt file. + let bytes = single_segment_file("absent-role.weights"); + let reader = Lyrw2Reader::parse(&bytes).unwrap(); + assert_eq!(reader.resolve(BANK_ID, 0, RegionRole::Up).unwrap(), None); + assert_eq!(reader.resolve(BANK_ID, 0, RegionRole::Bias).unwrap(), None); +} + +#[test] +fn an_absent_bank_resolves_to_none() { + let bytes = single_segment_file("absent-bank.weights"); + let reader = Lyrw2Reader::parse(&bytes).unwrap(); + assert_eq!(reader.resolve(99, 0, RegionRole::Gate).unwrap(), None); + assert!(reader.schemas_for(99).is_none()); +} + +#[test] +fn an_entry_past_the_segment_resolves_to_none() { + let bytes = single_segment_file("past-end.weights"); + let reader = Lyrw2Reader::parse(&bytes).unwrap(); + assert_eq!( + reader.resolve(BANK_ID, ENTRIES, RegionRole::Gate).unwrap(), + None + ); +} + +#[test] +fn a_split_layer_resolves_only_its_own_entries() { + let half = ENTRIES / 2; + let mut p = plan(); + p.header = Lyrw2Header::new(LOGICAL_LAYER, 1, 1).with_multi_segment(true); + p.segments = vec![SegmentDescriptor { + bank_id: BANK_ID, + segment_index: 1, + first_entry: half, + entry_count: half, + }]; + let bytes = write_file(&temp_path("split-seg1.weights"), p, half); + let reader = Lyrw2Reader::parse(&bytes).unwrap(); + + assert!(reader.header().is_multi_segment()); + // Entries 0..half live in the sibling segment file. + assert_eq!(reader.resolve(BANK_ID, 0, RegionRole::Gate).unwrap(), None); + // Entry `half` is this segment's local entry 0, written with byte 0. + let gate = reader + .region_bytes(BANK_ID, half, RegionRole::Gate) + .unwrap() + .unwrap(); + assert_eq!(gate, gate_pattern(0)); +} + +#[test] +fn a_dense_layer_is_the_single_entry_case() { + let dense = BankDescriptor { + bank_id: BANK_ID, + kind: BankKind::Dense, + num_entries: 1, + input_dim: HIDDEN, + intermediate_dim: INTERMEDIATE, + output_dim: HIDDEN, + region_schema_count: 2, + browse: BrowseMode::None, + }; + let p = Lyrw2Plan::single_segment(0, dense, schemas()); + let bytes = write_file(&temp_path("dense.weights"), p, 1); + let reader = Lyrw2Reader::parse(&bytes).unwrap(); + assert_eq!(reader.banks()[0].kind, BankKind::Dense); + assert!(reader + .region_bytes(BANK_ID, 0, RegionRole::Gate) + .unwrap() + .is_some()); +} + +#[test] +fn layout_is_recovered_from_the_parsed_tables() { + let bytes = single_segment_file("layout.weights"); + let reader = Lyrw2Reader::parse(&bytes).unwrap(); + let layout = reader.layout(); + assert!(layout.payload_start > layout.entry_table); + assert_eq!(layout.schema_count(0), Some(2)); +} diff --git a/crates/larql-vindex/src/format/lyrw2/region_format.rs b/crates/larql-vindex/src/format/lyrw2/region_format.rs new file mode 100644 index 000000000..7e596df9a --- /dev/null +++ b/crates/larql-vindex/src/format/lyrw2/region_format.rs @@ -0,0 +1,265 @@ +//! Region encodings and packings (spec §6.4). +//! +//! Like roles, unrecognised tags are preserved rather than rejected — a reader +//! that only walks gate regions must not fail because a `down` region uses a +//! codec it cannot decode. Refusal happens when a kernel is asked to execute +//! the region (spec §10/§11), which is where the honest diagnosis lives. + +/// Numeric encoding of a region's payload bytes. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum RegionFormat { + F32, + F16, + BF16, + Q4_0, + Q4K, + Q6K, + Q8_0, + Fp4Larql, + Mxfp4, + Nvfp4, + Mxfp8, + /// A codec this binary does not recognise. Round-trips unchanged. + Unknown(u16), +} + +impl RegionFormat { + pub fn from_u16(tag: u16) -> Self { + match tag { + 0 => Self::F32, + 1 => Self::F16, + 2 => Self::BF16, + 3 => Self::Q4_0, + 4 => Self::Q4K, + 5 => Self::Q6K, + 6 => Self::Q8_0, + 7 => Self::Fp4Larql, + 8 => Self::Mxfp4, + 9 => Self::Nvfp4, + 10 => Self::Mxfp8, + other => Self::Unknown(other), + } + } + + pub fn as_u16(self) -> u16 { + match self { + Self::F32 => 0, + Self::F16 => 1, + Self::BF16 => 2, + Self::Q4_0 => 3, + Self::Q4K => 4, + Self::Q6K => 5, + Self::Q8_0 => 6, + Self::Fp4Larql => 7, + Self::Mxfp4 => 8, + Self::Nvfp4 => 9, + Self::Mxfp8 => 10, + Self::Unknown(tag) => tag, + } + } + + /// Whether rows can be reinterpreted in place with no dequantisation — + /// the zero-copy gate-KNN fast path (spec §15.1). + pub fn is_zero_copy_walkable(self) -> bool { + matches!(self, Self::F32 | Self::F16) + } + + /// A borrowed name for the registered codecs; `None` for [`Self::Unknown`], + /// whose name carries its tag and so cannot be static. + /// + /// Exists because a kernel refusal names the format it *wanted*, and a + /// refusal is a `&'static str` field — the alternative was a second, + /// hand-written spelling of every codec name at each call site. + /// [`Self::name`] delegates here so the two cannot drift. + pub const fn registered_name(self) -> Option<&'static str> { + Some(match self { + Self::F32 => "f32", + Self::F16 => "f16", + Self::BF16 => "bf16", + Self::Q4_0 => "q4_0", + Self::Q4K => "q4_k", + Self::Q6K => "q6_k", + Self::Q8_0 => "q8_0", + Self::Fp4Larql => "fp4_larql", + Self::Mxfp4 => "mxfp4", + Self::Nvfp4 => "nvfp4", + Self::Mxfp8 => "mxfp8", + Self::Unknown(_) => return None, + }) + } + + pub fn name(self) -> String { + match self.registered_name() { + Some(name) => name.into(), + None => format!("format_{}", self.as_u16()), + } + } +} + +/// How a region's bytes are laid out. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum Packing { + /// Dense rows, no separate scale stream. + RowMajor, + /// Quantised blocks with their scales interleaved inline. + BlocksWithScalesInline, + /// Block values only; scales live in a partner region named by `pair_id`. + BlocksValues, + /// Block scales only; values live in a partner region named by `pair_id`. + BlocksScales, + Unknown(u16), +} + +impl Packing { + pub fn from_u16(tag: u16) -> Self { + match tag { + 0 => Self::RowMajor, + 1 => Self::BlocksWithScalesInline, + 2 => Self::BlocksValues, + 3 => Self::BlocksScales, + other => Self::Unknown(other), + } + } + + pub fn as_u16(self) -> u16 { + match self { + Self::RowMajor => 0, + Self::BlocksWithScalesInline => 1, + Self::BlocksValues => 2, + Self::BlocksScales => 3, + Self::Unknown(tag) => tag, + } + } + + /// Whether this packing *requires* a partner region via `pair_id`. + pub fn requires_pair(self) -> bool { + matches!(self, Self::BlocksValues | Self::BlocksScales) + } + + /// Whether gate rows can be reached by striding, without decoding the + /// interleaved `up` half of a fused region (spec §15.2). + pub fn permits_strided_gate_read(self) -> bool { + matches!(self, Self::RowMajor) + } + + pub fn name(self) -> String { + match self { + Self::RowMajor => "row_major".into(), + Self::BlocksWithScalesInline => "blocks_with_scales_inline".into(), + Self::BlocksValues => "blocks_values".into(), + Self::BlocksScales => "blocks_scales".into(), + Self::Unknown(tag) => format!("packing_{tag}"), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn registered_formats_round_trip() { + for tag in 0u16..=10 { + assert_eq!(RegionFormat::from_u16(tag).as_u16(), tag); + } + } + + #[test] + fn unknown_format_is_preserved_not_rejected() { + let f = RegionFormat::from_u16(4_242); + assert_eq!(f, RegionFormat::Unknown(4_242)); + assert_eq!(f.as_u16(), 4_242); + assert_eq!(f.name(), "format_4242"); + } + + #[test] + fn only_f16_and_f32_are_zero_copy_walkable() { + assert!(RegionFormat::F16.is_zero_copy_walkable()); + assert!(RegionFormat::F32.is_zero_copy_walkable()); + // bf16 needs a widening pass, so it is not the v1 zero-copy path. + assert!(!RegionFormat::BF16.is_zero_copy_walkable()); + assert!(!RegionFormat::Q6K.is_zero_copy_walkable()); + assert!(!RegionFormat::Mxfp4.is_zero_copy_walkable()); + } + + #[test] + fn native_low_bit_codecs_are_registered() { + assert_eq!(RegionFormat::Mxfp4.name(), "mxfp4"); + assert_eq!(RegionFormat::Nvfp4.name(), "nvfp4"); + assert_eq!(RegionFormat::Mxfp8.name(), "mxfp8"); + } + + #[test] + fn every_registered_format_has_a_distinct_name() { + let names: Vec = (0u16..=10) + .map(|t| RegionFormat::from_u16(t).name()) + .collect(); + assert_eq!( + names, + vec![ + "f32", + "f16", + "bf16", + "q4_0", + "q4_k", + "q6_k", + "q8_0", + "fp4_larql", + "mxfp4", + "nvfp4", + "mxfp8" + ] + ); + } + + #[test] + fn no_registered_format_is_zero_copy_beyond_f16_and_f32() { + let walkable: Vec = (0u16..=10) + .filter(|t| RegionFormat::from_u16(*t).is_zero_copy_walkable()) + .collect(); + assert_eq!(walkable, vec![0, 1]); + } + + #[test] + fn registered_packings_round_trip() { + for tag in 0u16..=3 { + assert_eq!(Packing::from_u16(tag).as_u16(), tag); + } + } + + #[test] + fn unknown_packing_is_preserved() { + assert_eq!(Packing::from_u16(77), Packing::Unknown(77)); + assert_eq!(Packing::Unknown(77).name(), "packing_77"); + } + + #[test] + fn split_packings_require_a_partner() { + assert!(Packing::BlocksValues.requires_pair()); + assert!(Packing::BlocksScales.requires_pair()); + assert!(!Packing::RowMajor.requires_pair()); + assert!(!Packing::BlocksWithScalesInline.requires_pair()); + assert!(!Packing::Unknown(9).requires_pair()); + } + + #[test] + fn every_registered_packing_has_a_distinct_name() { + let names: Vec = (0u16..=3).map(|t| Packing::from_u16(t).name()).collect(); + assert_eq!( + names, + vec![ + "row_major", + "blocks_with_scales_inline", + "blocks_values", + "blocks_scales" + ] + ); + } + + #[test] + fn only_row_major_permits_strided_gate_reads() { + assert!(Packing::RowMajor.permits_strided_gate_read()); + assert!(!Packing::BlocksWithScalesInline.permits_strided_gate_read()); + assert!(!Packing::BlocksValues.permits_strided_gate_read()); + } +} diff --git a/crates/larql-vindex/src/format/lyrw2/region_role.rs b/crates/larql-vindex/src/format/lyrw2/region_role.rs new file mode 100644 index 000000000..08b682717 --- /dev/null +++ b/crates/larql-vindex/src/format/lyrw2/region_role.rs @@ -0,0 +1,215 @@ +//! Region roles — what a stored region *is* within an expert (spec §6.5). +//! +//! Unknown roles are preserved, not rejected. A file carrying a role this +//! binary has never heard of is still a valid file: §6.5 makes presence of an +//! unregistered role harmless, and absence of a *required* role a capability +//! failure for one programme rather than a corrupt container. Fail-closed +//! belongs at capability-check time (§11), not at parse time — a browse-only +//! reader must not choke on a `down` region it will never touch. + +use super::consts::PAIR_ID_UNPAIRED; + +/// A registered region role, or an unrecognised tag preserved verbatim. +/// +/// `Ord` follows the registry tag, so ordering is the spec's role order rather +/// than declaration order in this enum — capability reports sort by coordinate +/// and a stable, spec-derived order keeps their output diffable. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum RegionRole { + Gate, + Up, + GateUpFused, + Down, + Bias, + Scales, + LatentIn, + LatentOut, + /// A tag this binary does not recognise. Round-trips unchanged. + Unknown(u16), +} + +impl PartialOrd for RegionRole { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for RegionRole { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + self.as_u16().cmp(&other.as_u16()) + } +} + +/// First tag in the vendor/experimental space (spec §6.5). +pub const VENDOR_ROLE_BASE: u16 = 256; + +impl RegionRole { + pub fn from_u16(tag: u16) -> Self { + match tag { + 0 => Self::Gate, + 1 => Self::Up, + 2 => Self::GateUpFused, + 3 => Self::Down, + 4 => Self::Bias, + 5 => Self::Scales, + 6 => Self::LatentIn, + 7 => Self::LatentOut, + other => Self::Unknown(other), + } + } + + pub fn as_u16(self) -> u16 { + match self { + Self::Gate => 0, + Self::Up => 1, + Self::GateUpFused => 2, + Self::Down => 3, + Self::Bias => 4, + Self::Scales => 5, + Self::LatentIn => 6, + Self::LatentOut => 7, + Self::Unknown(tag) => tag, + } + } + + /// Whether this role carries gate rows, and so is walkable by gate KNN + /// (spec §15.1). `GateUpFused` qualifies only when the bank's browse mode + /// permits striding — that is a bank-level question, not a role-level one, + /// so this answers the role half alone. + pub fn carries_gate_rows(self) -> bool { + matches!(self, Self::Gate | Self::GateUpFused) + } + + /// Whether this role is in the vendor/experimental space. + pub fn is_vendor(self) -> bool { + self.as_u16() >= VENDOR_ROLE_BASE + } + + /// Human-readable name for diagnostics. Unregistered tags render as + /// `role_` so an error message never claims a name it does not know. + pub fn name(self) -> String { + match self { + Self::Gate => "gate".into(), + Self::Up => "up".into(), + Self::GateUpFused => "gate_up_fused".into(), + Self::Down => "down".into(), + Self::Bias => "bias".into(), + Self::Scales => "scales".into(), + Self::LatentIn => "latent_in".into(), + Self::LatentOut => "latent_out".into(), + Self::Unknown(tag) => format!("role_{tag}"), + } + } +} + +/// Whether a `pair_id` names a partner schema. +pub fn is_paired(pair_id: u16) -> bool { + pair_id != PAIR_ID_UNPAIRED +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn registered_roles_round_trip() { + for tag in 0u16..=7 { + assert_eq!(RegionRole::from_u16(tag).as_u16(), tag); + } + } + + #[test] + fn unknown_role_is_preserved_not_rejected() { + let role = RegionRole::from_u16(9_000); + assert_eq!(role, RegionRole::Unknown(9_000)); + assert_eq!(role.as_u16(), 9_000); + } + + #[test] + fn reserved_registered_span_parses_as_unknown() { + // 8..255 is reserved-registered: not yet named here, still legal. + assert_eq!(RegionRole::from_u16(8), RegionRole::Unknown(8)); + assert!(!RegionRole::Unknown(8).is_vendor()); + } + + #[test] + fn vendor_space_starts_at_256() { + assert!(!RegionRole::from_u16(VENDOR_ROLE_BASE - 1).is_vendor()); + assert!(RegionRole::from_u16(VENDOR_ROLE_BASE).is_vendor()); + assert!(!RegionRole::Gate.is_vendor()); + } + + #[test] + fn only_gate_bearing_roles_are_walkable() { + assert!(RegionRole::Gate.carries_gate_rows()); + assert!(RegionRole::GateUpFused.carries_gate_rows()); + assert!(!RegionRole::Up.carries_gate_rows()); + assert!(!RegionRole::Down.carries_gate_rows()); + assert!(!RegionRole::Scales.carries_gate_rows()); + assert!(!RegionRole::Unknown(300).carries_gate_rows()); + } + + #[test] + fn names_are_stable_for_registered_roles() { + assert_eq!(RegionRole::Gate.name(), "gate"); + assert_eq!(RegionRole::GateUpFused.name(), "gate_up_fused"); + assert_eq!(RegionRole::LatentOut.name(), "latent_out"); + } + + #[test] + fn every_registered_role_has_a_distinct_name() { + let names: Vec = (0u16..=7).map(|t| RegionRole::from_u16(t).name()).collect(); + assert_eq!( + names, + vec![ + "gate", + "up", + "gate_up_fused", + "down", + "bias", + "scales", + "latent_in", + "latent_out" + ] + ); + } + + #[test] + fn unknown_role_name_admits_it_is_unknown() { + assert_eq!(RegionRole::Unknown(42).name(), "role_42"); + } + + #[test] + fn roles_order_by_registry_tag_not_declaration() { + // Sorting a capability report must be stable and spec-derived, so the + // order follows §6.5's tag numbering. + let mut v = vec![ + RegionRole::Down, + RegionRole::Gate, + RegionRole::Unknown(300), + RegionRole::Up, + ]; + v.sort(); + assert_eq!( + v, + vec![ + RegionRole::Gate, + RegionRole::Up, + RegionRole::Down, + RegionRole::Unknown(300), + ] + ); + } + + #[test] + fn an_unknown_role_sorts_after_every_registered_one() { + assert!(RegionRole::Unknown(256) > RegionRole::LatentOut); + } + + #[test] + fn unpaired_sentinel_reads_as_unpaired() { + assert!(!is_paired(PAIR_ID_UNPAIRED)); + assert!(is_paired(0)); + assert!(is_paired(3)); + } +} diff --git a/crates/larql-vindex/src/format/lyrw2/region_schema.rs b/crates/larql-vindex/src/format/lyrw2/region_schema.rs new file mode 100644 index 000000000..953ac68ae --- /dev/null +++ b/crates/larql-vindex/src/format/lyrw2/region_schema.rs @@ -0,0 +1,188 @@ +//! Per-bank region schema (spec §6.4). +//! +//! Expert banks are homogeneous: every entry in a bank shares one region +//! layout, so the schema is declared once per bank and each entry stores only +//! offsets and lengths. That is what makes per-expert codec variation — which +//! no grouped kernel supports — unrepresentable by construction rather than by +//! convention, and what makes parsing O(schemas) instead of O(entries × regions). + +use super::consts::{PAIR_ID_UNPAIRED, REGION_SCHEMA_BYTES}; +use super::region_format::{Packing, RegionFormat}; +use super::region_role::RegionRole; +use super::wire::{push_u16, push_u32, read_u16, read_u32}; + +/// One region's declared shape and encoding, shared by every entry in a bank. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct RegionSchema { + pub schema_index: u16, + pub role: RegionRole, + pub format: RegionFormat, + pub packing: Packing, + /// Links a `BlocksValues` schema to its `BlocksScales` partner, and back. + /// `PAIR_ID_UNPAIRED` when the region stands alone. + pub pair_id: u16, + pub rows: u32, + pub cols: u32, +} + +impl RegionSchema { + /// A standalone region with no values/scales partner. + pub fn unpaired( + schema_index: u16, + role: RegionRole, + format: RegionFormat, + packing: Packing, + rows: u32, + cols: u32, + ) -> Self { + Self { + schema_index, + role, + format, + packing, + pair_id: PAIR_ID_UNPAIRED, + rows, + cols, + } + } + + pub fn encode(&self, out: &mut Vec) { + push_u16(out, self.schema_index); + push_u16(out, self.role.as_u16()); + push_u16(out, self.format.as_u16()); + push_u16(out, self.packing.as_u16()); + push_u16(out, self.pair_id); + push_u16(out, 0); // reserved + push_u32(out, self.rows); + push_u32(out, self.cols); + } + + pub fn decode(bytes: &[u8]) -> Option { + if bytes.len() < REGION_SCHEMA_BYTES { + return None; + } + Some(Self { + schema_index: read_u16(bytes, 0)?, + role: RegionRole::from_u16(read_u16(bytes, 2)?), + format: RegionFormat::from_u16(read_u16(bytes, 4)?), + packing: Packing::from_u16(read_u16(bytes, 6)?), + pair_id: read_u16(bytes, 8)?, + // bytes 10..12 reserved + rows: read_u32(bytes, 12)?, + cols: read_u32(bytes, 16)?, + }) + } + + /// Whether this schema declares a partner but names none, or names one + /// while declaring a packing that has no partner. Both are writer bugs + /// that would otherwise surface as a silently half-decoded region. + pub fn pairing_is_consistent(&self) -> bool { + self.packing.requires_pair() == (self.pair_id != PAIR_ID_UNPAIRED) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn sample() -> RegionSchema { + RegionSchema::unpaired( + 1, + RegionRole::Down, + RegionFormat::Q6K, + Packing::BlocksWithScalesInline, + 3_584, + 3_072, + ) + } + + #[test] + fn schema_round_trips_through_bytes() { + let schema = sample(); + let mut buf = Vec::new(); + schema.encode(&mut buf); + assert_eq!(buf.len(), REGION_SCHEMA_BYTES); + assert_eq!(RegionSchema::decode(&buf), Some(schema)); + } + + #[test] + fn paired_schema_round_trips() { + let schema = RegionSchema { + schema_index: 2, + role: RegionRole::Scales, + format: RegionFormat::Mxfp4, + packing: Packing::BlocksScales, + pair_id: 1, + rows: 64, + cols: 8, + }; + let mut buf = Vec::new(); + schema.encode(&mut buf); + assert_eq!(RegionSchema::decode(&buf), Some(schema)); + } + + #[test] + fn unknown_tags_survive_a_round_trip() { + let schema = RegionSchema { + schema_index: 7, + role: RegionRole::Unknown(900), + format: RegionFormat::Unknown(901), + packing: Packing::Unknown(902), + pair_id: PAIR_ID_UNPAIRED, + rows: 1, + cols: 1, + }; + let mut buf = Vec::new(); + schema.encode(&mut buf); + assert_eq!(RegionSchema::decode(&buf), Some(schema)); + } + + #[test] + fn short_record_decodes_to_none() { + let mut buf = Vec::new(); + sample().encode(&mut buf); + buf.pop(); + assert_eq!(RegionSchema::decode(&buf), None); + } + + #[test] + fn unpaired_helper_sets_the_sentinel() { + assert_eq!(sample().pair_id, PAIR_ID_UNPAIRED); + } + + #[test] + fn reserved_field_is_written_as_zero() { + let mut buf = Vec::new(); + sample().encode(&mut buf); + assert_eq!(read_u16(&buf, 10), Some(0)); + } + + #[test] + fn pairing_consistency_accepts_matched_declarations() { + assert!(sample().pairing_is_consistent()); + let paired = RegionSchema { + pair_id: 3, + packing: Packing::BlocksValues, + ..sample() + }; + assert!(paired.pairing_is_consistent()); + } + + #[test] + fn pairing_consistency_rejects_a_split_region_with_no_partner() { + let orphan = RegionSchema { + packing: Packing::BlocksValues, + ..sample() + }; + assert!(!orphan.pairing_is_consistent()); + } + + #[test] + fn pairing_consistency_rejects_a_partner_on_an_inline_region() { + let spurious = RegionSchema { + pair_id: 0, + ..sample() + }; + assert!(!spurious.pairing_is_consistent()); + } +} diff --git a/crates/larql-vindex/src/format/lyrw2/segment.rs b/crates/larql-vindex/src/format/lyrw2/segment.rs new file mode 100644 index 000000000..77edd2262 --- /dev/null +++ b/crates/larql-vindex/src/format/lyrw2/segment.rs @@ -0,0 +1,159 @@ +//! Segment descriptors (spec §6.3). +//! +//! A segment is a contiguous slice of one bank's entries stored in one file. +//! Single-file layers have exactly one segment covering `[0, num_entries)`. +//! Multi-segment layers repeat the header in every segment file; `index.json` +//! lists the files per logical layer so the loader never globs a directory. +//! +//! Segmentation is a physical storage parameter. It is invisible to model +//! semantics: the logical layer remains the stable unit, and which file an +//! expert lives in is not something a manifest or a kernel ever asks. + +use super::consts::SEGMENT_DESCRIPTOR_BYTES; +use super::wire::{push_u16, push_u32, read_u16, read_u32}; + +/// The slice of a bank's entries carried by one file. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct SegmentDescriptor { + pub bank_id: u16, + pub segment_index: u16, + pub first_entry: u32, + pub entry_count: u32, +} + +impl SegmentDescriptor { + pub fn encode(&self, out: &mut Vec) { + push_u16(out, self.bank_id); + push_u16(out, self.segment_index); + push_u32(out, self.first_entry); + push_u32(out, self.entry_count); + } + + pub fn decode(bytes: &[u8]) -> Option { + if bytes.len() < SEGMENT_DESCRIPTOR_BYTES { + return None; + } + Some(Self { + bank_id: read_u16(bytes, 0)?, + segment_index: read_u16(bytes, 2)?, + first_entry: read_u32(bytes, 4)?, + entry_count: read_u32(bytes, 8)?, + }) + } + + /// One past the last logical entry index this segment carries. + pub fn end_entry(&self) -> u64 { + u64::from(self.first_entry) + u64::from(self.entry_count) + } + + /// Whether `logical_entry` falls inside this segment. + pub fn contains(&self, logical_entry: u32) -> bool { + let entry = u64::from(logical_entry); + entry >= u64::from(self.first_entry) && entry < self.end_entry() + } + + /// Position of `logical_entry` within this segment's entry table, or + /// `None` if the entry lives in a different segment. + pub fn local_index(&self, logical_entry: u32) -> Option { + self.contains(logical_entry) + .then(|| logical_entry - self.first_entry) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// K3 exact-Q6_K prior: two segments of 448 experts per routed layer. + fn second_half_of_k3_layer() -> SegmentDescriptor { + SegmentDescriptor { + bank_id: 0, + segment_index: 1, + first_entry: 448, + entry_count: 448, + } + } + + #[test] + fn descriptor_round_trips_through_bytes() { + let seg = second_half_of_k3_layer(); + let mut buf = Vec::new(); + seg.encode(&mut buf); + assert_eq!(buf.len(), SEGMENT_DESCRIPTOR_BYTES); + assert_eq!(SegmentDescriptor::decode(&buf), Some(seg)); + } + + #[test] + fn short_record_decodes_to_none() { + let mut buf = Vec::new(); + second_half_of_k3_layer().encode(&mut buf); + buf.pop(); + assert_eq!(SegmentDescriptor::decode(&buf), None); + } + + #[test] + fn end_entry_is_exclusive() { + assert_eq!(second_half_of_k3_layer().end_entry(), 896); + } + + #[test] + fn contains_covers_the_half_open_range() { + let seg = second_half_of_k3_layer(); + assert!(!seg.contains(447)); + assert!(seg.contains(448)); + assert!(seg.contains(895)); + assert!(!seg.contains(896)); + } + + #[test] + fn local_index_rebases_onto_the_segment() { + let seg = second_half_of_k3_layer(); + assert_eq!(seg.local_index(448), Some(0)); + assert_eq!(seg.local_index(895), Some(447)); + } + + #[test] + fn local_index_refuses_entries_from_another_segment() { + let seg = second_half_of_k3_layer(); + assert_eq!(seg.local_index(0), None); + assert_eq!(seg.local_index(447), None); + assert_eq!(seg.local_index(896), None); + } + + #[test] + fn single_file_layer_covers_every_entry() { + let seg = SegmentDescriptor { + bank_id: 0, + segment_index: 0, + first_entry: 0, + entry_count: 256, + }; + assert!(seg.contains(0)); + assert!(seg.contains(255)); + assert!(!seg.contains(256)); + assert_eq!(seg.local_index(17), Some(17)); + } + + #[test] + fn end_entry_does_not_overflow_at_u32_bounds() { + let seg = SegmentDescriptor { + bank_id: 0, + segment_index: 0, + first_entry: u32::MAX, + entry_count: u32::MAX, + }; + assert_eq!(seg.end_entry(), u64::from(u32::MAX) * 2); + } + + #[test] + fn empty_segment_contains_nothing() { + let seg = SegmentDescriptor { + bank_id: 0, + segment_index: 0, + first_entry: 10, + entry_count: 0, + }; + assert!(!seg.contains(10)); + assert_eq!(seg.local_index(10), None); + } +} diff --git a/crates/larql-vindex/src/format/lyrw2/test_fixtures.rs b/crates/larql-vindex/src/format/lyrw2/test_fixtures.rs new file mode 100644 index 000000000..a64da97a0 --- /dev/null +++ b/crates/larql-vindex/src/format/lyrw2/test_fixtures.rs @@ -0,0 +1,112 @@ +//! Shared fixtures for the LYRW v2 container tests. +//! +//! One routed bank with two roles carrying two *different* codecs — the whole +//! point of per-region format tags — written with a recognisable byte pattern +//! so a test can tell entry 2's gate from entry 3's down by inspection. + +use std::path::{Path, PathBuf}; + +use super::bank::{BankDescriptor, BankKind}; +use super::browse_mode::BrowseMode; +use super::layout::Lyrw2Layout; +use super::plan::Lyrw2Plan; +use super::region_format::{Packing, RegionFormat}; +use super::region_role::RegionRole; +use super::region_schema::RegionSchema; +use super::write::Lyrw2Writer; + +pub(super) const BANK_ID: u16 = 0; +pub(super) const ENTRIES: u32 = 4; +pub(super) const SCHEMAS: u16 = 2; +pub(super) const HIDDEN: u32 = 4; +pub(super) const INTERMEDIATE: u32 = 8; +pub(super) const GATE_BYTES: usize = 40; +pub(super) const DOWN_BYTES: usize = 24; +pub(super) const LOGICAL_LAYER: u32 = 7; + +/// Gate bytes for entry `e` are all `e`; down bytes are all `0x80 | e`. +pub(super) fn gate_pattern(entry: u32) -> Vec { + vec![entry as u8; GATE_BYTES] +} + +pub(super) fn down_pattern(entry: u32) -> Vec { + vec![0x80 | entry as u8; DOWN_BYTES] +} + +pub(super) fn bank(num_entries: u32) -> BankDescriptor { + BankDescriptor { + bank_id: BANK_ID, + kind: BankKind::Routed, + num_entries, + input_dim: HIDDEN, + intermediate_dim: INTERMEDIATE, + output_dim: HIDDEN, + region_schema_count: SCHEMAS, + browse: BrowseMode::Direct, + } +} + +/// Gate at f16 row-major (browse-walkable), down at Q6_K blocks — two roles, +/// two codecs, one file. +pub(super) fn schemas() -> Vec { + vec![ + RegionSchema::unpaired( + 0, + RegionRole::Gate, + RegionFormat::F16, + Packing::RowMajor, + INTERMEDIATE, + HIDDEN, + ), + RegionSchema::unpaired( + 1, + RegionRole::Down, + RegionFormat::Q6K, + Packing::BlocksWithScalesInline, + HIDDEN, + INTERMEDIATE, + ), + ] +} + +pub(super) fn plan() -> Lyrw2Plan { + Lyrw2Plan::single_segment(LOGICAL_LAYER, bank(ENTRIES), schemas()) +} + +pub(super) fn temp_path(name: &str) -> PathBuf { + let dir = std::env::temp_dir().join("lyrw2-tests"); + std::fs::create_dir_all(&dir).unwrap(); + dir.join(name) +} + +/// Write `entries` entries through the streaming writer and return the bytes. +pub(super) fn write_file(path: &Path, plan: Lyrw2Plan, entries: u32) -> Vec { + let mut w = Lyrw2Writer::create(path, plan).unwrap(); + for e in 0..entries { + w.write_region(&gate_pattern(e)).unwrap(); + w.write_region(&down_pattern(e)).unwrap(); + } + w.finish().unwrap(); + std::fs::read(path).unwrap() +} + +/// A complete single-segment file covering every entry. +pub(super) fn single_segment_file(name: &str) -> Vec { + write_file(&temp_path(name), plan(), ENTRIES) +} + +/// The `(offset, length)` the writer recorded for one region, read back out of +/// the entry table through the layout rather than a hard-coded offset. +pub(super) fn slot(bytes: &[u8], entry: u32, schema: u16) -> (u64, u64) { + let at = slot_offset(entry, schema); + let offset = u64::from_le_bytes(bytes[at..at + 8].try_into().unwrap()); + let length = u64::from_le_bytes(bytes[at + 8..at + 16].try_into().unwrap()); + (offset, length) +} + +/// Byte offset of one entry-table slot in the default fixture file. +pub(super) fn slot_offset(entry: u32, schema: u16) -> usize { + Lyrw2Layout::of(&plan()) + .entry_slot(0, entry, schema) + .unwrap() as usize +} diff --git a/crates/larql-vindex/src/format/lyrw2/wire.rs b/crates/larql-vindex/src/format/lyrw2/wire.rs new file mode 100644 index 000000000..e9c8e0a51 --- /dev/null +++ b/crates/larql-vindex/src/format/lyrw2/wire.rs @@ -0,0 +1,90 @@ +//! Little-endian field access over fixed-width descriptor slices. +//! +//! Every descriptor in LYRW v2 is a fixed-size record read from a known +//! offset, so the whole wire layer needs exactly two readers and two writers. +//! Keeping them here means no descriptor module hand-rolls a `try_into`, which +//! is where off-by-two field offsets come from. + +/// Read a `u16` at `at` within `bytes`, or `None` if the record is short. +pub fn read_u16(bytes: &[u8], at: usize) -> Option { + let end = at.checked_add(2)?; + let slice = bytes.get(at..end)?; + Some(u16::from_le_bytes([slice[0], slice[1]])) +} + +/// Read a `u32` at `at` within `bytes`, or `None` if the record is short. +pub fn read_u32(bytes: &[u8], at: usize) -> Option { + let end = at.checked_add(4)?; + let slice = bytes.get(at..end)?; + Some(u32::from_le_bytes([slice[0], slice[1], slice[2], slice[3]])) +} + +/// Read a `u64` at `at` within `bytes`, or `None` if the record is short. +pub fn read_u64(bytes: &[u8], at: usize) -> Option { + let end = at.checked_add(8)?; + let slice = bytes.get(at..end)?; + let mut buf = [0u8; 8]; + buf.copy_from_slice(slice); + Some(u64::from_le_bytes(buf)) +} + +pub fn push_u16(out: &mut Vec, value: u16) { + out.extend_from_slice(&value.to_le_bytes()); +} + +pub fn push_u32(out: &mut Vec, value: u32) { + out.extend_from_slice(&value.to_le_bytes()); +} + +pub fn push_u64(out: &mut Vec, value: u64) { + out.extend_from_slice(&value.to_le_bytes()); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn u16_round_trips_at_offset() { + let mut buf = vec![0xAA, 0xBB]; + push_u16(&mut buf, 0x1234); + assert_eq!(read_u16(&buf, 2), Some(0x1234)); + } + + #[test] + fn u32_round_trips_at_offset() { + let mut buf = vec![0u8; 3]; + push_u32(&mut buf, 0xDEAD_BEEF); + assert_eq!(read_u32(&buf, 3), Some(0xDEAD_BEEF)); + } + + #[test] + fn u64_round_trips_at_offset() { + let mut buf = vec![0u8; 1]; + push_u64(&mut buf, u64::MAX - 7); + assert_eq!(read_u64(&buf, 1), Some(u64::MAX - 7)); + } + + #[test] + fn short_records_return_none_rather_than_panicking() { + let buf = [0u8; 3]; + assert_eq!(read_u16(&buf, 2), None); + assert_eq!(read_u32(&buf, 0), None); + assert_eq!(read_u64(&buf, 0), None); + } + + #[test] + fn offset_overflow_returns_none() { + let buf = [0u8; 8]; + assert_eq!(read_u16(&buf, usize::MAX), None); + assert_eq!(read_u32(&buf, usize::MAX), None); + assert_eq!(read_u64(&buf, usize::MAX), None); + } + + #[test] + fn encoding_is_little_endian() { + let mut buf = Vec::new(); + push_u32(&mut buf, 1); + assert_eq!(buf, vec![1, 0, 0, 0]); + } +} diff --git a/crates/larql-vindex/src/format/lyrw2/write.rs b/crates/larql-vindex/src/format/lyrw2/write.rs new file mode 100644 index 000000000..712a8f3e5 --- /dev/null +++ b/crates/larql-vindex/src/format/lyrw2/write.rs @@ -0,0 +1,212 @@ +//! Streaming LYRW v2 writer. +//! +//! The v1 writer took `&[LayerEntry]` — every expert's quantised bytes, fully +//! materialised. At a K3 routed layer that is 39.4 GB of input, a ~118 GB f32 +//! intermediate and 24.3 GB of output against 128 GB of RAM: unsurvivable, and +//! a property of the API signature rather than of the machine. +//! +//! This writer streams instead. It emits the tables, reserves the entry table, +//! then takes one region at a time and backpatches offsets on close. Peak +//! resident bytes are one region plus the table — about 27 MB and 28 KB +//! respectively at K3 scale, whatever the layer's total size. +//! +//! Regions must arrive in file order: for each segment, each entry, each +//! schema. That ordering is checked, not assumed — a caller that interleaves +//! banks gets an error, not a file with plausible offsets into the wrong place. + +use std::fs::File; +use std::io::{BufWriter, Seek, SeekFrom, Write}; +use std::path::Path; + +use super::consts::align_up; +use super::layout::Lyrw2Layout; +use super::plan::Lyrw2Plan; +use super::wire::push_u64; +use crate::VindexError; + +/// Where the writer expects the next region to belong. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct RegionCursor { + pub segment_ordinal: usize, + pub local_entry: u32, + pub schema_index: u16, +} + +/// Streaming writer for one LYRW v2 segment file. +pub struct Lyrw2Writer { + file: BufWriter, + plan: Lyrw2Plan, + layout: Lyrw2Layout, + /// Recorded `(offset, length)` per region, in the order written. + slots: Vec<(u64, u64)>, + cursor: RegionCursor, + payload_cursor: u64, + finished: bool, +} + +impl Lyrw2Writer { + /// Create the file and write every table except the entry table, which is + /// reserved as zeroes and backpatched by [`Lyrw2Writer::finish`]. + pub fn create(path: &Path, plan: Lyrw2Plan) -> Result { + plan.validate().map_err(to_vindex_error)?; + let layout = Lyrw2Layout::of(&plan); + + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + let mut file = BufWriter::new(File::create(path)?); + + let mut tables = Vec::new(); + plan.header.encode(&mut tables); + for bank in &plan.banks { + bank.encode(&mut tables); + } + for segment in &plan.segments { + segment.encode(&mut tables); + } + for schemas in &plan.schemas { + for schema in schemas { + schema.encode(&mut tables); + } + } + file.write_all(&tables)?; + + // Reserve the entry table plus the alignment gap before the payload. + let reserved = layout.payload_start - layout.entry_table; + file.write_all(&vec![0u8; reserved as usize])?; + + Ok(Self { + file, + plan, + payload_cursor: layout.payload_start, + layout, + slots: Vec::new(), + cursor: RegionCursor { + segment_ordinal: 0, + local_entry: 0, + schema_index: 0, + }, + finished: false, + }) + } + + /// Where the next [`Lyrw2Writer::write_region`] call will land. + pub fn cursor(&self) -> RegionCursor { + self.cursor + } + + /// Whether every region the plan declares has been written. + pub fn is_complete(&self) -> bool { + self.cursor.segment_ordinal >= self.plan.segments.len() + } + + /// Regions appended so far. This is also the writer's entire retained + /// state besides the plan — one `(offset, length)` pair per region, never + /// the payload — which is the property that lets a 24 GB layer stream + /// through a 27 MB working set. + pub fn regions_written(&self) -> usize { + self.slots.len() + } + + /// Byte length recorded for the region at `ordinal`, in write order. + pub fn region_length(&self, ordinal: usize) -> Option { + self.slots.get(ordinal).map(|(_, length)| *length) + } + + /// Append one region's payload bytes, aligning to the region boundary + /// first. Regions must arrive in file order; see the module docs. + pub fn write_region(&mut self, bytes: &[u8]) -> Result<(), VindexError> { + if self.is_complete() { + return Err(VindexError::Parse( + "lyrw2: write_region called after every declared region was written".into(), + )); + } + + let aligned = align_up(self.payload_cursor); + let pad = aligned - self.payload_cursor; + if pad > 0 { + self.file.write_all(&vec![0u8; pad as usize])?; + } + self.file.write_all(bytes)?; + + self.slots.push((aligned, bytes.len() as u64)); + self.payload_cursor = aligned + bytes.len() as u64; + self.advance(); + Ok(()) + } + + /// Step the cursor one region forward through (segment, entry, schema). + fn advance(&mut self) { + let schema_count = self + .layout + .schema_count(self.cursor.segment_ordinal) + .unwrap_or(0) as u16; + let entry_count = self.plan.segments[self.cursor.segment_ordinal].entry_count; + + self.cursor.schema_index += 1; + if self.cursor.schema_index < schema_count { + return; + } + self.cursor.schema_index = 0; + self.cursor.local_entry += 1; + if self.cursor.local_entry < entry_count { + return; + } + self.cursor.local_entry = 0; + self.cursor.segment_ordinal += 1; + } + + /// Backpatch the entry table and flush. + /// + /// Refuses to close a file that is missing regions — a short file whose + /// table is half zeroes reads as valid offsets into byte zero, which is + /// exactly the plausible-garbage failure this format exists to prevent. + pub fn finish(mut self) -> Result<(), VindexError> { + // The writer is consumed either way; the drop guard exists to catch a + // *forgotten* finish, not a refused one. + self.finished = true; + if !self.is_complete() { + return Err(VindexError::Parse(format!( + "lyrw2: {} regions written but the plan declares more; \ + next expected segment {} entry {} schema {}", + self.slots.len(), + self.cursor.segment_ordinal, + self.cursor.local_entry, + self.cursor.schema_index, + ))); + } + + let mut table = Vec::with_capacity(self.slots.len() * 16); + for (offset, length) in &self.slots { + push_u64(&mut table, *offset); + push_u64(&mut table, *length); + } + + self.file.seek(SeekFrom::Start(self.layout.entry_table))?; + self.file.write_all(&table)?; + self.file.flush()?; + Ok(()) + } + + /// Deliberately discard a partially written file. + /// + /// The bytes on disk keep a zeroed entry table and must not be served — + /// this exists so a caller that aborts on purpose is distinguishable from + /// one that forgot to call [`Lyrw2Writer::finish`]. + pub fn abandon(mut self) { + self.finished = true; + } +} + +impl Drop for Lyrw2Writer { + fn drop(&mut self) { + debug_assert!( + self.finished || std::thread::panicking(), + "Lyrw2Writer dropped without finish(): the entry table is still zeroes" + ); + } +} + +fn to_vindex_error(e: super::error::Lyrw2Error) -> VindexError { + VindexError::Parse(e.to_string()) +} diff --git a/crates/larql-vindex/src/format/lyrw2/write_tests.rs b/crates/larql-vindex/src/format/lyrw2/write_tests.rs new file mode 100644 index 000000000..32f13fe82 --- /dev/null +++ b/crates/larql-vindex/src/format/lyrw2/write_tests.rs @@ -0,0 +1,182 @@ +//! Colocated tests for `write` — the streaming LYRW v2 writer. +//! +//! The writer is exercised end to end (create, stream regions, backpatch) and +//! then pinned on its refusal paths: an invalid plan, a short file, an overrun, +//! and a deliberate abandon. Byte-level assertions resolve slots through +//! `Lyrw2Layout` rather than hard-coding offsets, so a layout change surfaces +//! as a layout test failure rather than as silently relocated expectations. + +use super::consts::REGION_ALIGNMENT; +use super::layout::Lyrw2Layout; +use super::test_fixtures::{ + down_pattern, gate_pattern, plan, single_segment_file, slot, temp_path, DOWN_BYTES, ENTRIES, + GATE_BYTES, SCHEMAS, +}; +use super::write::{Lyrw2Writer, RegionCursor}; + +#[test] +fn a_complete_file_writes_and_closes() { + let path = temp_path("w-complete.weights"); + let mut w = Lyrw2Writer::create(&path, plan()).unwrap(); + for e in 0..ENTRIES { + w.write_region(&gate_pattern(e)).unwrap(); + w.write_region(&down_pattern(e)).unwrap(); + } + assert!(w.finish().is_ok()); + assert!(path.exists()); +} + +#[test] +fn cursor_starts_at_the_first_region() { + let path = temp_path("cursor-start.weights"); + let w = Lyrw2Writer::create(&path, plan()).unwrap(); + assert_eq!( + w.cursor(), + RegionCursor { + segment_ordinal: 0, + local_entry: 0, + schema_index: 0 + } + ); + assert!(!w.is_complete()); + assert_eq!(w.regions_written(), 0); + w.abandon(); +} + +#[test] +fn cursor_walks_schemas_then_entries() { + let path = temp_path("cursor-walk.weights"); + let mut w = Lyrw2Writer::create(&path, plan()).unwrap(); + w.write_region(&[0u8; GATE_BYTES]).unwrap(); + assert_eq!(w.cursor().schema_index, 1); + assert_eq!(w.cursor().local_entry, 0); + w.write_region(&[0u8; DOWN_BYTES]).unwrap(); + assert_eq!(w.cursor().schema_index, 0); + assert_eq!(w.cursor().local_entry, 1); + w.abandon(); +} + +#[test] +fn writer_reports_completion_after_the_last_region() { + let path = temp_path("complete-flag.weights"); + let mut w = Lyrw2Writer::create(&path, plan()).unwrap(); + for _ in 0..ENTRIES { + w.write_region(&[0u8; GATE_BYTES]).unwrap(); + w.write_region(&[0u8; DOWN_BYTES]).unwrap(); + } + assert!(w.is_complete()); + assert_eq!(w.regions_written(), (ENTRIES * u32::from(SCHEMAS)) as usize); + w.finish().unwrap(); +} + +#[test] +fn finishing_early_is_refused() { + let path = temp_path("short.weights"); + let mut w = Lyrw2Writer::create(&path, plan()).unwrap(); + w.write_region(&[0u8; GATE_BYTES]).unwrap(); + let err = w.finish().unwrap_err(); + assert!(err.to_string().contains("plan declares more"), "{err}"); +} + +#[test] +fn writing_past_the_plan_is_refused() { + let path = temp_path("overrun.weights"); + let mut w = Lyrw2Writer::create(&path, plan()).unwrap(); + for _ in 0..ENTRIES { + w.write_region(&[0u8; GATE_BYTES]).unwrap(); + w.write_region(&[0u8; DOWN_BYTES]).unwrap(); + } + let err = w.write_region(&[0u8; 4]).unwrap_err(); + assert!(err.to_string().contains("after every declared"), "{err}"); + w.finish().unwrap(); +} + +#[test] +fn an_invalid_plan_is_refused_before_the_file_is_touched() { + let path = temp_path("invalid-plan.weights"); + let _ = std::fs::remove_file(&path); + let mut bad = plan(); + bad.banks[0].region_schema_count = 5; + assert!(Lyrw2Writer::create(&path, bad).is_err()); + assert!(!path.exists()); +} + +#[test] +fn abandoning_a_partial_file_is_not_a_forgotten_finish() { + let path = temp_path("abandoned.weights"); + let mut w = Lyrw2Writer::create(&path, plan()).unwrap(); + w.write_region(&[0u8; GATE_BYTES]).unwrap(); + w.abandon(); + assert!(path.exists()); +} + +#[test] +fn retained_state_is_one_slot_per_region_not_the_payload() { + let path = temp_path("streaming.weights"); + let mut w = Lyrw2Writer::create(&path, plan()).unwrap(); + w.write_region(&[7u8; GATE_BYTES]).unwrap(); + assert_eq!(w.regions_written(), 1); + assert_eq!(w.region_length(0), Some(GATE_BYTES as u64)); + assert_eq!(w.region_length(1), None); + w.abandon(); +} + +#[test] +fn every_region_lands_on_the_alignment_boundary() { + let bytes = single_segment_file("w-aligned.weights"); + for entry in 0..ENTRIES { + for schema in 0..SCHEMAS { + let (offset, _) = slot(&bytes, entry, schema); + assert_eq!( + offset % REGION_ALIGNMENT, + 0, + "entry {entry} schema {schema}" + ); + } + } +} + +#[test] +fn the_entry_table_records_the_lengths_written() { + let bytes = single_segment_file("w-lengths.weights"); + assert_eq!(slot(&bytes, 0, 0).1, GATE_BYTES as u64); + assert_eq!(slot(&bytes, 0, 1).1, DOWN_BYTES as u64); +} + +#[test] +fn payload_bytes_are_readable_at_the_recorded_offsets() { + let bytes = single_segment_file("w-payload.weights"); + let (offset, _) = slot(&bytes, 2, 0); + let start = offset as usize; + // Entry 2's gate region was filled with the byte 2. + assert_eq!(&bytes[start..start + GATE_BYTES], &[2u8; GATE_BYTES]); +} + +#[test] +fn regions_never_overlap() { + let bytes = single_segment_file("w-disjoint.weights"); + let mut spans: Vec<(u64, u64)> = Vec::new(); + for entry in 0..ENTRIES { + for schema in 0..SCHEMAS { + let (offset, length) = slot(&bytes, entry, schema); + spans.push((offset, offset + length)); + } + } + spans.sort_unstable(); + for pair in spans.windows(2) { + assert!( + pair[0].1 <= pair[1].0, + "{:?} overlaps {:?}", + pair[0], + pair[1] + ); + } +} + +#[test] +fn tables_precede_the_payload() { + let bytes = single_segment_file("w-ordering.weights"); + let layout = Lyrw2Layout::of(&plan()); + assert!(bytes.len() as u64 > layout.payload_start); + assert!(slot(&bytes, 0, 0).0 >= layout.payload_start); +} diff --git a/crates/larql-vindex/src/format/mod.rs b/crates/larql-vindex/src/format/mod.rs index 06666838a..68b23ab0c 100644 --- a/crates/larql-vindex/src/format/mod.rs +++ b/crates/larql-vindex/src/format/mod.rs @@ -1,19 +1,28 @@ //! File format I/O — vindex loading, saving, checksums, HuggingFace. //! Model loading (safetensors/GGUF) is in larql-models. +pub mod capability; pub mod checksums; +pub mod describes; pub mod down_meta; pub mod filenames; pub mod fp4_codec; +pub mod generation; +#[cfg(test)] +mod generation_tests; pub mod huggingface; pub mod le_floats; pub mod load; +pub mod lyrw2; +pub mod moe_manifest; pub mod quant; pub mod spec; +pub mod vindex3; pub mod weights; // Back-compat alias — `format::fp4_storage` was renamed to `fp4_codec` // in the 2026-04-25 round-2 cleanup (the file does encoding-side // codec work; the runtime store lives at `index::storage::fp4_store`). // Drop this alias once external callers are migrated. +pub use describes::model_id_at; pub use fp4_codec as fp4_storage; diff --git a/crates/larql-vindex/src/format/moe_manifest/bank_ref.rs b/crates/larql-vindex/src/format/moe_manifest/bank_ref.rs new file mode 100644 index 000000000..23298c6f4 --- /dev/null +++ b/crates/larql-vindex/src/format/moe_manifest/bank_ref.rs @@ -0,0 +1,160 @@ +//! Manifest references to expert banks (format spec §8.1). +//! +//! A bank reference binds a storage location and an expert count to a +//! programme. It is the **only** place that binding exists — LYRW files carry +//! no programme identity (§6.2), so "the binary says one thing, the manifest +//! says another" is unrepresentable rather than merely discouraged. + +use serde::{Deserialize, Serialize}; + +use super::programme::Programme; + +/// An expert's own operand dimensions. +/// +/// For a latent bank these are the latent width, not the residual width — the +/// distinction that makes K3's 3584/3072/3584 correct where 7168 would be +/// wrong (§6.2). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub struct ExpertDims { + pub input: u32, + pub intermediate: u32, + pub output: u32, +} + +/// One bank of experts, named by the manifest. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct BankRef { + /// Expert count in this bank. A dense layer is the `1` case. + pub experts: u32, + /// Programme name as written in the manifest; resolved via [`Programme`]. + pub programme: String, + /// Storage stem — `routed/layer_012`, extended with `.seg{N}` per segment. + pub storage: String, + #[serde(default)] + pub expert_dims: Option, +} + +impl BankRef { + /// Resolve the declared programme, or `None` if this binary does not know + /// it. An unknown programme is a clean refusal, never a guess (§V2-0). + pub fn resolve_programme(&self) -> Option { + Programme::from_name(&self.programme) + } + + /// Whether this bank's experts operate in a latent space and therefore + /// require the layer's `routed_input` / `routed_output` transforms. + pub fn needs_latent_transforms(&self) -> bool { + self.resolve_programme() + .is_some_and(|p| p.operates_in_latent_space()) + } + + /// Walkable features this bank contributes (§15.1): every expert's + /// intermediate rows, across the whole bank, with no router involved. + /// + /// `None` when dims were not declared — the count is unknowable rather + /// than zero, and reporting zero would read as "nothing to browse". + pub fn walkable_feature_count(&self) -> Option { + self.expert_dims + .map(|d| u64::from(self.experts) * u64::from(d.intermediate)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn k3_routed() -> BankRef { + BankRef { + experts: 896, + programme: "latent-moe-v1".into(), + storage: "routed/layer_012".into(), + expert_dims: Some(ExpertDims { + input: 3_584, + intermediate: 3_072, + output: 3_584, + }), + } + } + + fn shared() -> BankRef { + BankRef { + experts: 2, + programme: "gated-mlp-v1".into(), + storage: "shared/layer_012".into(), + expert_dims: None, + } + } + + #[test] + fn a_bank_round_trips_through_json() { + let json = serde_json::to_string(&k3_routed()).unwrap(); + assert_eq!(serde_json::from_str::(&json).unwrap(), k3_routed()); + } + + #[test] + fn absent_dims_deserialise_to_none() { + let json = r#"{"experts": 2, "programme": "gated-mlp-v1", "storage": "shared/layer_0"}"#; + let b: BankRef = serde_json::from_str(json).unwrap(); + assert_eq!(b.expert_dims, None); + } + + #[test] + fn a_registered_programme_resolves() { + assert_eq!( + k3_routed().resolve_programme(), + Some(Programme::LatentMoeV1) + ); + } + + #[test] + fn an_unregistered_programme_resolves_to_none_not_a_default() { + let b = BankRef { + programme: "mixtral-expert-v9".into(), + ..shared() + }; + assert_eq!(b.resolve_programme(), None); + // And must not be silently treated as latent or as anything else. + assert!(!b.needs_latent_transforms()); + } + + #[test] + fn only_a_latent_bank_needs_the_transforms() { + assert!(k3_routed().needs_latent_transforms()); + assert!(!shared().needs_latent_transforms()); + } + + #[test] + fn latent_dims_are_the_experts_own_width_not_the_residual_width() { + // The §6.2 trap: K3's latent bank is 3584, never 7168. + let dims = k3_routed().expert_dims.unwrap(); + assert_eq!(dims.input, 3_584); + assert_eq!(dims.output, 3_584); + } + + #[test] + fn walkable_features_span_the_whole_bank() { + assert_eq!(k3_routed().walkable_feature_count(), Some(2_752_512)); + } + + #[test] + fn undeclared_dims_make_the_feature_count_unknown_not_zero() { + // Zero would read as "nothing to browse", which is a different claim. + assert_eq!(shared().walkable_feature_count(), None); + } + + #[test] + fn a_dense_layer_is_the_single_expert_case() { + let dense = BankRef { + experts: 1, + programme: "gated-mlp-v1".into(), + storage: "dense/layer_0".into(), + expert_dims: Some(ExpertDims { + input: 2_304, + intermediate: 9_216, + output: 2_304, + }), + }; + assert_eq!(dense.walkable_feature_count(), Some(9_216)); + assert!(!dense.needs_latent_transforms()); + } +} diff --git a/crates/larql-vindex/src/format/moe_manifest/layer.rs b/crates/larql-vindex/src/format/moe_manifest/layer.rs new file mode 100644 index 000000000..15e992e33 --- /dev/null +++ b/crates/larql-vindex/src/format/moe_manifest/layer.rs @@ -0,0 +1,281 @@ +//! Per-layer MoE description (format spec §8.1). +//! +//! Per-layer, deliberately. A global `first_k_dense_replace` field would have +//! handled Kimi-Linear's leading dense layer and then failed on Inkling-Small, +//! whose dense MLP sits at layer index 2 — mid-stack. Per-layer manifests +//! handle arbitrary dense/MoE schedules for free, which is the point. + +use serde::{Deserialize, Serialize}; + +use super::bank_ref::BankRef; +use super::router::Router; + +/// The space a layer's experts consume and produce. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum InputSpace { + /// Experts read and write the residual stream directly. + Residual, + /// Experts operate in a projected latent space; the layer's transforms + /// carry the residual↔latent hop. + Latent, +} + +/// Residual↔latent projections for a latent-space layer (§8.1). +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct Transforms { + /// residual → latent. Also what WALK projects a query through before + /// dot-producting against latent gate rows (§15.4). + pub routed_input: String, + /// latent → residual. + pub routed_output: String, +} + +/// How selected expert outputs are combined. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Reduction { + GateWeightedSum, +} + +/// How the reduced expert output rejoins the stream. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Combine { + ResidualAdd, +} + +/// One layer's complete MoE programme. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct MoeLayer { + pub layer: u32, + pub input_space: InputSpace, + pub router: Router, + /// Null for a conventional residual-space MoE. + #[serde(default)] + pub transforms: Option, + pub routed_bank: BankRef, + /// Absent for models with no always-active experts (GPT-OSS, Gemma). + #[serde(default)] + pub shared_bank: Option, + pub reduction: Reduction, + #[serde(default)] + pub routed_output_norm: Option, + pub combine: Combine, +} + +/// Why a layer cannot be executed as declared. +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub enum LayerDefect { + #[error("layer {layer}: bank '{storage}' names unknown programme '{programme}'")] + UnknownProgramme { + layer: u32, + storage: String, + programme: String, + }, + + #[error( + "layer {layer}: {which} bank runs {programme}, which operates in a latent space, \ + but the layer declares no routed_input/routed_output transforms" + )] + LatentBankWithoutTransforms { + layer: u32, + which: &'static str, + programme: String, + }, + + #[error( + "layer {layer}: input_space is '{declared}' but the routed bank's programme \ + '{programme}' operates in the other space" + )] + InputSpaceContradictsProgramme { + layer: u32, + declared: String, + programme: String, + }, + + #[error("layer {layer}: router scores tensor name is empty")] + RouterHasNoScores { layer: u32 }, + + #[error( + "layer {layer}: router selects {selected} experts but the routed bank declares only \ + {available}" + )] + SelectionExceedsBank { + layer: u32, + selected: usize, + available: u32, + }, + + #[error( + "layer {layer}: router declares shared_expert_sink but the layer has no shared bank \ + for the router to score" + )] + SinkWithoutSharedBank { layer: u32 }, + + #[error("layer {layer}: {which} bank has an empty storage reference")] + BankHasNoStorage { layer: u32, which: &'static str }, + + #[error("layer {layer}: routed and shared banks both name storage '{storage}'")] + DuplicateBankStorage { layer: u32, storage: String }, + + #[error("layer {layer}: {which} bank declares zero experts")] + BankHasNoExperts { layer: u32, which: &'static str }, + + #[error( + "layer {layer}: {which} bank declares a zero dimension \ + (input {input}, intermediate {intermediate}, output {output})" + )] + BankHasZeroDimension { + layer: u32, + which: &'static str, + input: u32, + intermediate: u32, + output: u32, + }, +} + +impl MoeLayer { + /// Every way this layer is internally inconsistent. + /// + /// These are checks the *manifest* can fail on its own, before any weight + /// byte is read — distinct from operand-absence, which needs the LYRW + /// files and lives in the capability check (§11). + /// + /// **Suppression is dependency-scoped, not blanket.** An unresolvable + /// programme suppresses only the checks whose *interpretation* depends on + /// knowing the programme — currently the input-space compatibility test. + /// Storage-shape defects are independent of programme identity and always + /// run, so a typo'd programme name can never conceal a duplicate storage + /// reference or a zero-width bank sitting behind it. + pub fn defects(&self) -> Vec { + let mut out = Vec::new(); + self.check_programmes_resolve(&mut out); + self.check_latent_consistency(&mut out); + self.check_router(&mut out); + self.check_bank_storage(&mut out); + out + } + + pub fn is_well_formed(&self) -> bool { + self.defects().is_empty() + } + + fn check_programmes_resolve(&self, out: &mut Vec) { + for bank in self.banks() { + if bank.resolve_programme().is_none() { + out.push(LayerDefect::UnknownProgramme { + layer: self.layer, + storage: bank.storage.clone(), + programme: bank.programme.clone(), + }); + } + } + } + + fn check_latent_consistency(&self, out: &mut Vec) { + if self.transforms.is_none() { + for (which, bank) in self.labelled_banks() { + if bank.needs_latent_transforms() { + out.push(LayerDefect::LatentBankWithoutTransforms { + layer: self.layer, + which, + programme: bank.programme.clone(), + }); + } + } + } + + let routed_is_latent = self.routed_bank.needs_latent_transforms(); + let declared_latent = self.input_space == InputSpace::Latent; + // Only contradicts when the programme resolved; an unknown programme + // is already reported and must not produce a second, misleading defect. + if self.routed_bank.resolve_programme().is_some() && routed_is_latent != declared_latent { + out.push(LayerDefect::InputSpaceContradictsProgramme { + layer: self.layer, + declared: format!("{:?}", self.input_space).to_lowercase(), + programme: self.routed_bank.programme.clone(), + }); + } + } + + fn check_router(&self, out: &mut Vec) { + if self.router.scores.trim().is_empty() { + out.push(LayerDefect::RouterHasNoScores { layer: self.layer }); + } + + let selected = self.router.selection.experts_per_token(); + if selected > self.routed_bank.experts as usize { + out.push(LayerDefect::SelectionExceedsBank { + layer: self.layer, + selected, + available: self.routed_bank.experts, + }); + } + + if self.router.shared_expert_sink && self.shared_bank.is_none() { + out.push(LayerDefect::SinkWithoutSharedBank { layer: self.layer }); + } + } + + /// Storage-shape checks. Independent of programme identity by design — + /// see the suppression note on [`MoeLayer::defects`]. + fn check_bank_storage(&self, out: &mut Vec) { + for (which, bank) in self.labelled_banks() { + if bank.storage.trim().is_empty() { + out.push(LayerDefect::BankHasNoStorage { + layer: self.layer, + which, + }); + } + if bank.experts == 0 { + out.push(LayerDefect::BankHasNoExperts { + layer: self.layer, + which, + }); + } + if let Some(d) = bank.expert_dims { + if d.input == 0 || d.intermediate == 0 || d.output == 0 { + out.push(LayerDefect::BankHasZeroDimension { + layer: self.layer, + which, + input: d.input, + intermediate: d.intermediate, + output: d.output, + }); + } + } + } + + if let Some(shared) = &self.shared_bank { + if !shared.storage.trim().is_empty() && shared.storage == self.routed_bank.storage { + out.push(LayerDefect::DuplicateBankStorage { + layer: self.layer, + storage: shared.storage.clone(), + }); + } + } + } + + /// Every bank in the layer, routed first. + pub fn banks(&self) -> Vec<&BankRef> { + self.labelled_banks().into_iter().map(|(_, b)| b).collect() + } + + fn labelled_banks(&self) -> Vec<(&'static str, &BankRef)> { + let mut v = vec![("routed", &self.routed_bank)]; + if let Some(shared) = &self.shared_bank { + v.push(("shared", shared)); + } + v + } + + /// Total walkable features this layer contributes (§15.1), or `None` if + /// any bank left its dimensions undeclared. + pub fn walkable_feature_count(&self) -> Option { + self.banks() + .iter() + .map(|b| b.walkable_feature_count()) + .try_fold(0u64, |acc, n| Some(acc + n?)) + } +} diff --git a/crates/larql-vindex/src/format/moe_manifest/layer_tests.rs b/crates/larql-vindex/src/format/moe_manifest/layer_tests.rs new file mode 100644 index 000000000..3b86ab534 --- /dev/null +++ b/crates/larql-vindex/src/format/moe_manifest/layer_tests.rs @@ -0,0 +1,358 @@ +//! Colocated tests for `layer` — per-layer manifest well-formedness. +//! +//! Each test corrupts one field of a valid layer and requires the matching +//! named defect. The checks here are the ones the manifest can fail on its +//! own, before any weight byte is read; operand-absence needs the LYRW files +//! and is tested with the capability check. + +use super::bank_ref::{BankRef, ExpertDims}; +use super::layer::{Combine, InputSpace, LayerDefect, MoeLayer, Reduction, Transforms}; +use super::router::{Router, RouterPostProcessing, ScoreActivation, Selection}; + +const LAYER: u32 = 12; +const ROUTED_EXPERTS: u32 = 128; + +fn router(k: usize) -> Router { + Router { + scores: format!("layers.{LAYER}.router.weight"), + activation: ScoreActivation::Softmax, + selection: Selection::TopK { k }, + post: RouterPostProcessing::default(), + shared_expert_sink: false, + } +} + +fn routed(programme: &str) -> BankRef { + BankRef { + experts: ROUTED_EXPERTS, + programme: programme.into(), + storage: format!("routed/layer_{LAYER:03}"), + expert_dims: Some(ExpertDims { + input: 2_816, + intermediate: 704, + output: 2_816, + }), + } +} + +/// A conventional residual-space MoE layer — the Gemma/GPT-OSS shape. +fn residual_layer() -> MoeLayer { + MoeLayer { + layer: LAYER, + input_space: InputSpace::Residual, + router: router(8), + transforms: None, + routed_bank: routed("gated-mlp-v1"), + shared_bank: None, + reduction: Reduction::GateWeightedSum, + routed_output_norm: None, + combine: Combine::ResidualAdd, + } +} + +/// A K3-shaped latent layer with shared experts and transforms. +fn latent_layer() -> MoeLayer { + MoeLayer { + layer: LAYER, + input_space: InputSpace::Latent, + router: router(16), + transforms: Some(Transforms { + routed_input: format!("layers.{LAYER}.routed_expert_down_proj"), + routed_output: format!("layers.{LAYER}.routed_expert_up_proj"), + }), + routed_bank: BankRef { + experts: 896, + programme: "latent-moe-v1".into(), + storage: format!("routed/layer_{LAYER:03}"), + expert_dims: Some(ExpertDims { + input: 3_584, + intermediate: 3_072, + output: 3_584, + }), + }, + shared_bank: Some(BankRef { + experts: 2, + programme: "gated-mlp-v1".into(), + storage: format!("shared/layer_{LAYER:03}"), + expert_dims: Some(ExpertDims { + input: 3_584, + intermediate: 3_072, + output: 3_584, + }), + }), + reduction: Reduction::GateWeightedSum, + routed_output_norm: Some(format!("layers.{LAYER}.routed_out_norm")), + combine: Combine::ResidualAdd, + } +} + +#[test] +fn a_conventional_residual_layer_is_well_formed() { + assert_eq!(residual_layer().defects(), vec![]); + assert!(residual_layer().is_well_formed()); +} + +#[test] +fn a_latent_layer_with_transforms_is_well_formed() { + assert_eq!(latent_layer().defects(), vec![]); +} + +#[test] +fn both_layer_shapes_round_trip_through_json() { + for l in [residual_layer(), latent_layer()] { + let json = serde_json::to_string(&l).unwrap(); + assert_eq!( + serde_json::from_str::(&json).unwrap(), + l, + "{json}" + ); + } +} + +#[test] +fn an_unknown_programme_is_named_with_its_storage() { + let mut l = residual_layer(); + l.routed_bank.programme = "mixtral-expert-v9".into(); + assert!(l.defects().iter().any(|d| matches!( + d, + LayerDefect::UnknownProgramme { programme, .. } if programme == "mixtral-expert-v9" + ))); +} + +#[test] +fn an_unknown_programme_does_not_also_report_a_space_contradiction() { + // One root cause, one defect — for the checks that DEPEND on the + // programme. A misleading space contradiction sends the reader chasing + // input_space when the real problem is the name. + // + // Asserted by absence of the dependent defect, deliberately not by a total + // count: a count assertion would forbid reporting *independent* defects + // alongside, which is the over-suppression this scoping exists to avoid. + let mut l = residual_layer(); + l.routed_bank.programme = "not-a-programme".into(); + assert!(!l + .defects() + .iter() + .any(|d| matches!(d, LayerDefect::InputSpaceContradictsProgramme { .. }))); +} + +#[test] +fn an_unknown_programme_does_not_conceal_independent_storage_defects() { + // A typo'd programme name must not hide a second, real corruption sitting + // behind it — storage shape is not a function of programme identity. + let mut l = residual_layer(); + l.routed_bank.programme = "not-a-programme".into(); + l.routed_bank.storage = " ".into(); + l.routed_bank.experts = 0; + + let defects = l.defects(); + assert!(defects + .iter() + .any(|d| matches!(d, LayerDefect::UnknownProgramme { .. }))); + assert!(defects.contains(&LayerDefect::BankHasNoStorage { + layer: LAYER, + which: "routed" + })); + assert!(defects.contains(&LayerDefect::BankHasNoExperts { + layer: LAYER, + which: "routed" + })); +} + +#[test] +fn a_shared_bank_does_not_require_a_sink() { + // The asymmetry that must stay pinned: a sink needs a shared bank, but + // always-active shared experts OUTSIDE sink normalisation are valid — + // that is Kimi-Linear's arrangement. `shared_expert_sink` must never + // become shorthand for "this layer has shared experts". + let mut l = latent_layer(); + l.router.shared_expert_sink = false; + assert!(l.shared_bank.is_some()); + assert!(l.is_well_formed(), "{:?}", l.defects()); +} + +#[test] +fn an_empty_bank_storage_reference_is_refused() { + let mut l = residual_layer(); + l.routed_bank.storage = String::new(); + assert!(l.defects().contains(&LayerDefect::BankHasNoStorage { + layer: LAYER, + which: "routed" + })); +} + +#[test] +fn two_banks_sharing_one_storage_reference_are_refused() { + let mut l = latent_layer(); + l.shared_bank.as_mut().unwrap().storage = l.routed_bank.storage.clone(); + assert!(l + .defects() + .iter() + .any(|d| matches!(d, LayerDefect::DuplicateBankStorage { .. }))); +} + +#[test] +fn two_empty_storage_references_report_emptiness_not_duplication() { + // Both blank is one defect class, not two — reporting "duplicate storage + // ''" on top would be noise pointing at the wrong fix. + let mut l = latent_layer(); + l.routed_bank.storage = String::new(); + l.shared_bank.as_mut().unwrap().storage = String::new(); + assert!(!l + .defects() + .iter() + .any(|d| matches!(d, LayerDefect::DuplicateBankStorage { .. }))); +} + +#[test] +fn a_zero_width_bank_is_refused() { + let mut l = residual_layer(); + l.routed_bank.expert_dims.as_mut().unwrap().intermediate = 0; + assert!(l + .defects() + .iter() + .any(|d| matches!(d, LayerDefect::BankHasZeroDimension { .. }))); +} + +#[test] +fn an_undeclared_dims_bank_is_not_treated_as_zero_width() { + // Absent dims mean unknown, not zero — reporting a zero-width defect + // would be inventing a fact the manifest never stated. + let mut l = residual_layer(); + l.routed_bank.expert_dims = None; + assert!(l.is_well_formed(), "{:?}", l.defects()); +} + +#[test] +fn a_latent_bank_without_transforms_is_refused() { + let mut l = latent_layer(); + l.transforms = None; + assert!(l.defects().iter().any(|d| matches!( + d, + LayerDefect::LatentBankWithoutTransforms { + which: "routed", + .. + } + ))); +} + +#[test] +fn a_latent_programme_in_a_residual_layer_is_refused() { + let mut l = latent_layer(); + l.input_space = InputSpace::Residual; + assert!(l + .defects() + .iter() + .any(|d| matches!(d, LayerDefect::InputSpaceContradictsProgramme { .. }))); +} + +#[test] +fn a_residual_programme_in_a_latent_layer_is_refused() { + let mut l = residual_layer(); + l.input_space = InputSpace::Latent; + assert!(l + .defects() + .iter() + .any(|d| matches!(d, LayerDefect::InputSpaceContradictsProgramme { .. }))); +} + +#[test] +fn an_empty_router_scores_name_is_refused() { + let mut l = residual_layer(); + l.router.scores = " ".into(); + assert!(l + .defects() + .contains(&LayerDefect::RouterHasNoScores { layer: LAYER })); +} + +#[test] +fn selecting_more_experts_than_the_bank_holds_is_refused() { + let mut l = residual_layer(); + l.router.selection = Selection::TopK { + k: ROUTED_EXPERTS as usize + 1, + }; + assert!(l.defects().contains(&LayerDefect::SelectionExceedsBank { + layer: LAYER, + selected: ROUTED_EXPERTS as usize + 1, + available: ROUTED_EXPERTS, + })); +} + +#[test] +fn selecting_exactly_the_whole_bank_is_allowed() { + // Dense-equivalent routing is degenerate but not malformed. + let mut l = residual_layer(); + l.router.selection = Selection::TopK { + k: ROUTED_EXPERTS as usize, + }; + assert!(l.is_well_formed()); +} + +#[test] +fn grouped_selection_is_counted_across_groups() { + let mut l = residual_layer(); + l.router.selection = Selection::GroupedTopK { + k: ROUTED_EXPERTS as usize, + groups: 2, + }; + assert!(l + .defects() + .iter() + .any(|d| matches!(d, LayerDefect::SelectionExceedsBank { .. }))); +} + +#[test] +fn a_sink_router_without_a_shared_bank_is_refused() { + // The sink scores shared experts; with no shared bank there is nothing to + // score, and the reduction's weight budget would be silently wrong. + let mut l = residual_layer(); + l.router.shared_expert_sink = true; + assert!(l + .defects() + .contains(&LayerDefect::SinkWithoutSharedBank { layer: LAYER })); +} + +#[test] +fn a_sink_router_with_a_shared_bank_is_well_formed() { + let mut l = latent_layer(); + l.router.shared_expert_sink = true; + assert!(l.is_well_formed()); +} + +#[test] +fn banks_lists_routed_first_then_shared() { + let l = latent_layer(); + let banks = l.banks(); + assert_eq!(banks.len(), 2); + assert!(banks[0].storage.starts_with("routed/")); + assert!(banks[1].storage.starts_with("shared/")); +} + +#[test] +fn a_layer_without_a_shared_bank_lists_one_bank() { + assert_eq!(residual_layer().banks().len(), 1); +} + +#[test] +fn walkable_features_sum_across_banks() { + // 896 x 3072 routed + 2 x 3072 shared + assert_eq!( + latent_layer().walkable_feature_count(), + Some(896 * 3_072 + 2 * 3_072) + ); +} + +#[test] +fn one_undeclared_bank_makes_the_layer_count_unknown() { + let mut l = latent_layer(); + l.shared_bank.as_mut().unwrap().expert_dims = None; + assert_eq!(l.walkable_feature_count(), None); +} + +#[test] +fn an_absent_shared_bank_deserialises_to_none() { + let json = serde_json::to_string(&residual_layer()).unwrap(); + let back: MoeLayer = serde_json::from_str(&json).unwrap(); + assert_eq!(back.shared_bank, None); + assert_eq!(back.transforms, None); +} diff --git a/crates/larql-vindex/src/format/moe_manifest/manifest_tests.rs b/crates/larql-vindex/src/format/moe_manifest/manifest_tests.rs new file mode 100644 index 000000000..a183ec611 --- /dev/null +++ b/crates/larql-vindex/src/format/moe_manifest/manifest_tests.rs @@ -0,0 +1,208 @@ +//! Colocated tests for the manifest document as a whole. +//! +//! Layer-level well-formedness is covered in `layer_tests`; these are the +//! document-level concerns — schema version, duplicate layers, and the +//! dense/MoE schedule being expressed by *absence* rather than by a field. + +use super::bank_ref::{BankRef, ExpertDims}; +use super::layer::{Combine, InputSpace, MoeLayer, Reduction}; +use super::router::{Router, RouterPostProcessing, ScoreActivation, Selection}; +use super::{ManifestDefect, MoeManifest, Programme, MOE_MANIFEST_SCHEMA_VERSION}; + +const EXPERTS: u32 = 64; + +fn layer_at(index: u32, programme: &str) -> MoeLayer { + MoeLayer { + layer: index, + input_space: InputSpace::Residual, + router: Router { + scores: format!("layers.{index}.router.weight"), + activation: ScoreActivation::Softmax, + selection: Selection::TopK { k: 8 }, + post: RouterPostProcessing::default(), + shared_expert_sink: false, + }, + transforms: None, + routed_bank: BankRef { + experts: EXPERTS, + programme: programme.into(), + storage: format!("routed/layer_{index:03}"), + expert_dims: Some(ExpertDims { + input: 512, + intermediate: 256, + output: 512, + }), + }, + shared_bank: None, + reduction: Reduction::GateWeightedSum, + routed_output_norm: None, + combine: Combine::ResidualAdd, + } +} + +/// Inkling's shape: dense at index 2, MoE either side of it. +fn mid_stack_dense() -> MoeManifest { + MoeManifest::new(vec![ + layer_at(0, "gated-mlp-v1"), + layer_at(1, "gated-mlp-v1"), + // index 2 deliberately absent — that is how dense is expressed + layer_at(3, "gated-mlp-v1"), + ]) +} + +#[test] +fn a_manifest_round_trips_through_json() { + let m = mid_stack_dense(); + let json = serde_json::to_string(&m).unwrap(); + assert_eq!(MoeManifest::parse(&json).unwrap(), m); +} + +#[test] +fn new_stamps_the_current_schema_version() { + assert_eq!( + MoeManifest::new(Vec::new()).schema_version, + MOE_MANIFEST_SCHEMA_VERSION + ); +} + +#[test] +fn a_well_formed_manifest_has_no_defects() { + assert_eq!(mid_stack_dense().defects(), vec![]); + assert!(mid_stack_dense().is_well_formed()); +} + +#[test] +fn a_dense_layer_is_expressed_by_absence_not_a_field() { + // The reason there is no `first_k_dense_replace`: this schedule has dense + // in the middle, which no leading-dense field can describe. + let m = mid_stack_dense(); + assert!(m.is_moe_layer(1)); + assert!(!m.is_moe_layer(2)); + assert!(m.is_moe_layer(3)); + assert!(m.layer(2).is_none()); +} + +#[test] +fn a_leading_dense_stack_is_the_same_mechanism() { + // Kimi-Linear's first_k_dense_replace=1 needs no special case. + let m = MoeManifest::new(vec![ + layer_at(1, "gated-mlp-v1"), + layer_at(2, "gated-mlp-v1"), + ]); + assert!(!m.is_moe_layer(0)); + assert!(m.is_moe_layer(1)); +} + +#[test] +fn layers_are_found_by_index_not_position() { + let m = mid_stack_dense(); + assert_eq!(m.layer(3).map(|l| l.layer), Some(3)); + assert_eq!(m.layer(99), None); +} + +#[test] +fn a_duplicate_layer_is_refused() { + let m = MoeManifest::new(vec![ + layer_at(4, "gated-mlp-v1"), + layer_at(4, "gated-mlp-v1"), + ]); + assert!(m + .defects() + .contains(&ManifestDefect::DuplicateLayer { layer: 4 })); +} + +#[test] +fn an_unsupported_schema_version_is_refused() { + let mut m = mid_stack_dense(); + m.schema_version = MOE_MANIFEST_SCHEMA_VERSION + 1; + assert_eq!( + m.defects(), + vec![ManifestDefect::UnsupportedSchemaVersion { + found: MOE_MANIFEST_SCHEMA_VERSION + 1, + supported: MOE_MANIFEST_SCHEMA_VERSION, + }] + ); +} + +#[test] +fn an_unsupported_schema_version_short_circuits_layer_checks() { + // Reporting per-layer defects from a document this binary cannot interpret + // would describe a shape that may not be the one on disk. + let mut m = MoeManifest::new(vec![layer_at(0, "not-a-programme")]); + m.schema_version = 99; + assert_eq!(m.defects().len(), 1); + assert!(matches!( + m.defects()[0], + ManifestDefect::UnsupportedSchemaVersion { .. } + )); +} + +#[test] +fn layer_defects_are_surfaced_at_the_document_level() { + let m = MoeManifest::new(vec![layer_at(0, "not-a-programme")]); + assert!(m + .defects() + .iter() + .any(|d| matches!(d, ManifestDefect::Layer(_)))); +} + +#[test] +fn programmes_are_reported_without_duplicates() { + // Three layers, one programme. + assert_eq!(mid_stack_dense().programmes(), vec![Programme::GatedMlpV1]); +} + +#[test] +fn distinct_programmes_are_all_reported() { + let m = MoeManifest::new(vec![ + layer_at(0, "gated-mlp-v1"), + layer_at(1, "gpt-oss-expert-v1"), + ]); + let mut got = m.programmes(); + got.sort_by_key(|p| p.id()); + assert_eq!(got, vec![Programme::GatedMlpV1, Programme::GptOssExpertV1]); +} + +#[test] +fn unknown_programmes_are_listed_by_name() { + let m = MoeManifest::new(vec![ + layer_at(0, "gated-mlp-v1"), + layer_at(1, "mixtral-expert-v9"), + ]); + assert_eq!(m.unknown_programmes(), vec!["mixtral-expert-v9"]); + // The known one still resolves — one bad name does not poison the rest. + assert_eq!(m.programmes(), vec![Programme::GatedMlpV1]); +} + +#[test] +fn an_unknown_programme_is_listed_once_however_often_it_appears() { + let m = MoeManifest::new(vec![ + layer_at(0, "mixtral-expert-v9"), + layer_at(1, "mixtral-expert-v9"), + ]); + assert_eq!(m.unknown_programmes().len(), 1); +} + +#[test] +fn an_all_known_manifest_lists_no_unknowns() { + assert!(mid_stack_dense().unknown_programmes().is_empty()); +} + +#[test] +fn an_empty_manifest_is_well_formed_and_all_dense() { + let m = MoeManifest::new(Vec::new()); + assert!(m.is_well_formed()); + assert!(!m.is_moe_layer(0)); + assert!(m.programmes().is_empty()); +} + +#[test] +fn malformed_json_is_a_parse_error_not_a_panic() { + assert!(MoeManifest::parse("{not json").is_err()); +} + +#[test] +fn json_missing_a_required_field_is_refused() { + // `layers` is not optional; a manifest without it is not an empty manifest. + assert!(MoeManifest::parse(r#"{"schema_version": 1}"#).is_err()); +} diff --git a/crates/larql-vindex/src/format/moe_manifest/mod.rs b/crates/larql-vindex/src/format/moe_manifest/mod.rs new file mode 100644 index 000000000..e73eacc42 --- /dev/null +++ b/crates/larql-vindex/src/format/moe_manifest/mod.rs @@ -0,0 +1,137 @@ +//! The MoE programme manifest (format spec §8). +//! +//! The physical index stores tensor regions; this manifest gives them meaning; +//! the runtime picks an optimised kernel when it recognises the programme. +//! Keeping meaning here and storage in LYRW is what makes "the binary says +//! programme 4, the manifest says gpt-oss-expert-v1" unrepresentable rather +//! than merely discouraged. +//! +//! Per-layer, deliberately (§8.1). A global `first_k_dense_replace` would have +//! covered Kimi-Linear's leading dense layer and then failed on Inkling-Small, +//! whose dense MLP sits mid-stack at index 2. + +pub mod bank_ref; +pub mod layer; +#[cfg(test)] +mod layer_tests; +#[cfg(test)] +mod manifest_tests; +pub mod programme; +pub mod router; + +use serde::{Deserialize, Serialize}; + +pub use bank_ref::{BankRef, ExpertDims}; +pub use layer::{Combine, InputSpace, LayerDefect, MoeLayer, Reduction, Transforms}; +pub use programme::{Programme, ALL_PROGRAMMES}; +pub use router::{Router, RouterPostProcessing, ScoreActivation, Selection}; + +/// Schema version of `moe_manifest.json` (§12). Independent of the container +/// generation: the manifest can gain fields without a container bump. +pub const MOE_MANIFEST_SCHEMA_VERSION: u32 = 1; + +/// The whole manifest document. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct MoeManifest { + pub schema_version: u32, + /// Layers that run a MoE programme. Layers absent from this list are + /// dense and carry no manifest entry — which is how an arbitrary + /// dense/MoE schedule is expressed without a dedicated field. + pub layers: Vec, +} + +/// Why a manifest cannot be used as a whole. +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub enum ManifestDefect { + #[error( + "moe_manifest schema_version {found} is not supported; this binary implements {supported}" + )] + UnsupportedSchemaVersion { found: u32, supported: u32 }, + + #[error("layer {layer} is declared more than once")] + DuplicateLayer { layer: u32 }, + + #[error("{0}")] + Layer(#[from] LayerDefect), +} + +impl MoeManifest { + pub fn new(layers: Vec) -> Self { + Self { + schema_version: MOE_MANIFEST_SCHEMA_VERSION, + layers, + } + } + + pub fn parse(json: &str) -> Result { + serde_json::from_str(json).map_err(|e| crate::VindexError::Parse(e.to_string())) + } + + /// Every reason this manifest is unusable, document-level first. + /// + /// An unsupported schema version short-circuits: reporting per-layer + /// defects from a document this binary cannot interpret would be + /// describing a shape that may not be the one on disk. + pub fn defects(&self) -> Vec { + if self.schema_version != MOE_MANIFEST_SCHEMA_VERSION { + return vec![ManifestDefect::UnsupportedSchemaVersion { + found: self.schema_version, + supported: MOE_MANIFEST_SCHEMA_VERSION, + }]; + } + + let mut out = Vec::new(); + let mut seen: Vec = Vec::with_capacity(self.layers.len()); + for l in &self.layers { + if seen.contains(&l.layer) { + out.push(ManifestDefect::DuplicateLayer { layer: l.layer }); + } else { + seen.push(l.layer); + } + out.extend(l.defects().into_iter().map(ManifestDefect::Layer)); + } + out + } + + pub fn is_well_formed(&self) -> bool { + self.defects().is_empty() + } + + pub fn layer(&self, index: u32) -> Option<&MoeLayer> { + self.layers.iter().find(|l| l.layer == index) + } + + /// Whether `index` runs a MoE programme. A layer absent from the manifest + /// is dense, not missing. + pub fn is_moe_layer(&self, index: u32) -> bool { + self.layer(index).is_some() + } + + /// Distinct programmes referenced anywhere, for capability reporting. + pub fn programmes(&self) -> Vec { + let mut out: Vec = Vec::new(); + for l in &self.layers { + for b in l.banks() { + if let Some(p) = b.resolve_programme() { + if !out.contains(&p) { + out.push(p); + } + } + } + } + out + } + + /// Programme names referenced but not registered in this binary. + pub fn unknown_programmes(&self) -> Vec { + let mut out: Vec = Vec::new(); + for l in &self.layers { + for b in l.banks() { + if b.resolve_programme().is_none() && !out.contains(&b.programme) { + out.push(b.programme.clone()); + } + } + } + out + } +} diff --git a/crates/larql-vindex/src/format/moe_manifest/programme.rs b/crates/larql-vindex/src/format/moe_manifest/programme.rs new file mode 100644 index 000000000..09a699017 --- /dev/null +++ b/crates/larql-vindex/src/format/moe_manifest/programme.rs @@ -0,0 +1,466 @@ +//! Expert-programme registry (format spec §8.4). +//! +//! A programme names what an expert *computes*. Storage says nothing about it: +//! LYRW files carry banks, entries and region schemas, and the manifest is the +//! only binding of `bank_id → programme`. Two authorities for the same fact is +//! a disagreement waiting to happen, so the binary deliberately has no opinion. +//! +//! Each programme declares the region roles it needs. That declaration is what +//! capability checking (§11) traverses to answer "can this index serve this +//! layer", which is why the requirement has to express **alternatives** rather +//! than a flat set: `gate + up` and `gate_up_fused` are equally valid ways to +//! satisfy the same programme, and an index is servable if it presents either. + +use crate::format::lyrw2::region_role::RegionRole; + +/// A registered expert programme. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum Programme { + GatedMlpV1, + GatedMlpFusedFc1V1, + GptOssExpertV1, + SharedRoutedMlpV1, + LatentMoeV1, +} + +/// Registry ids, frozen (§8.4). New programmes append; ids never move. +const ID_GATED_MLP_V1: u16 = 0; +const ID_GATED_MLP_FUSED_FC1_V1: u16 = 1; +const ID_GPT_OSS_EXPERT_V1: u16 = 2; +const ID_SHARED_ROUTED_MLP_V1: u16 = 3; +const ID_LATENT_MOE_V1: u16 = 4; + +const NAME_GATED_MLP_V1: &str = "gated-mlp-v1"; +const NAME_GATED_MLP_FUSED_FC1_V1: &str = "gated-mlp-fused-fc1-v1"; +const NAME_GPT_OSS_EXPERT_V1: &str = "gpt-oss-expert-v1"; +const NAME_SHARED_ROUTED_MLP_V1: &str = "shared-routed-mlp-v1"; +const NAME_LATENT_MOE_V1: &str = "latent-moe-v1"; + +/// Every registered programme, in id order. +pub const ALL_PROGRAMMES: [Programme; 5] = [ + Programme::GatedMlpV1, + Programme::GatedMlpFusedFc1V1, + Programme::GptOssExpertV1, + Programme::SharedRoutedMlpV1, + Programme::LatentMoeV1, +]; + +/// One acceptable way to satisfy a programme's operand needs. +/// +/// A programme is satisfied when **any** of its alternatives is fully present. +/// Modelling this as alternatives rather than a flat role set is what lets one +/// manifest describe both fused and decomposed storage — the V2-1 acceptance +/// requirement that the two "produce identical results under one manifest". +pub type RoleAlternative = &'static [RegionRole]; + +/// Why a programme is not satisfied, and against which layout that was judged. +/// +/// Carrying the alternative alongside the gap is what lets a diagnostic say +/// *"closest accepted layout: gate_up_fused + down; missing down"* rather than +/// a bare list of roles the reader then has to reverse-engineer a layout from. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Unsatisfied { + /// The layout this gap was measured against. + pub closest: RoleAlternative, + /// Roles absent from `closest`. + pub missing: Vec, +} + +impl Unsatisfied { + /// Human-readable form used verbatim in loader diagnostics. + pub fn describe(&self) -> String { + let layout = self + .closest + .iter() + .map(|r| r.name()) + .collect::>() + .join(" + "); + let missing = self + .missing + .iter() + .map(|r| r.name()) + .collect::>() + .join(", "); + format!("closest accepted layout: {layout}; missing {missing}") + } +} + +impl Programme { + pub fn from_id(id: u16) -> Option { + match id { + ID_GATED_MLP_V1 => Some(Self::GatedMlpV1), + ID_GATED_MLP_FUSED_FC1_V1 => Some(Self::GatedMlpFusedFc1V1), + ID_GPT_OSS_EXPERT_V1 => Some(Self::GptOssExpertV1), + ID_SHARED_ROUTED_MLP_V1 => Some(Self::SharedRoutedMlpV1), + ID_LATENT_MOE_V1 => Some(Self::LatentMoeV1), + _ => None, + } + } + + pub fn from_name(name: &str) -> Option { + match name { + NAME_GATED_MLP_V1 => Some(Self::GatedMlpV1), + NAME_GATED_MLP_FUSED_FC1_V1 => Some(Self::GatedMlpFusedFc1V1), + NAME_GPT_OSS_EXPERT_V1 => Some(Self::GptOssExpertV1), + NAME_SHARED_ROUTED_MLP_V1 => Some(Self::SharedRoutedMlpV1), + NAME_LATENT_MOE_V1 => Some(Self::LatentMoeV1), + _ => None, + } + } + + pub const fn id(self) -> u16 { + match self { + Self::GatedMlpV1 => ID_GATED_MLP_V1, + Self::GatedMlpFusedFc1V1 => ID_GATED_MLP_FUSED_FC1_V1, + Self::GptOssExpertV1 => ID_GPT_OSS_EXPERT_V1, + Self::SharedRoutedMlpV1 => ID_SHARED_ROUTED_MLP_V1, + Self::LatentMoeV1 => ID_LATENT_MOE_V1, + } + } + + pub const fn name(self) -> &'static str { + match self { + Self::GatedMlpV1 => NAME_GATED_MLP_V1, + Self::GatedMlpFusedFc1V1 => NAME_GATED_MLP_FUSED_FC1_V1, + Self::GptOssExpertV1 => NAME_GPT_OSS_EXPERT_V1, + Self::SharedRoutedMlpV1 => NAME_SHARED_ROUTED_MLP_V1, + Self::LatentMoeV1 => NAME_LATENT_MOE_V1, + } + } + + /// Acceptable operand sets, most-preferred first. + /// + /// `gate_up_fused + down` is listed before `gate + up + down` where both + /// are legal, because a kernel that can take the fused form generally + /// prefers it — but presence, not preference, decides servability. + pub const fn role_alternatives(self) -> &'static [RoleAlternative] { + const FUSED: RoleAlternative = &[RegionRole::GateUpFused, RegionRole::Down]; + const DECOMPOSED: RoleAlternative = &[RegionRole::Gate, RegionRole::Up, RegionRole::Down]; + // GPT-OSS experts carry per-expert biases on both projections; without + // them the clamped-GLU-plus-residual programme is not reproducible. + const FUSED_WITH_BIAS: RoleAlternative = + &[RegionRole::GateUpFused, RegionRole::Down, RegionRole::Bias]; + const DECOMPOSED_WITH_BIAS: RoleAlternative = &[ + RegionRole::Gate, + RegionRole::Up, + RegionRole::Down, + RegionRole::Bias, + ]; + + match self { + // Fused storage is legal for any gated MLP — §6.5's fast-path + // contract accepts either shape. + Self::GatedMlpV1 | Self::SharedRoutedMlpV1 | Self::LatentMoeV1 => &[FUSED, DECOMPOSED], + // This programme *is* the fused variant; decomposed storage means + // the manifest should have named `gated-mlp-v1` instead. + Self::GatedMlpFusedFc1V1 => &[FUSED], + Self::GptOssExpertV1 => &[FUSED_WITH_BIAS, DECOMPOSED_WITH_BIAS], + } + } + + /// **Every** alternative `present` satisfies, in declaration order. + /// + /// Deliberately not "the first match". A bank can physically satisfy both + /// `gate + up + down` and `gate_up_fused + down`, and the kernel registry + /// may support one at a higher maturity than the other. Collapsing to a + /// single alternative during semantic traversal would let a Production + /// fused kernel hide behind a Reference decomposed path — a correctness- + /// preserving but silently slower binding, which is the worst kind to + /// debug. Ranking is kernel binding's job; traversal must not pre-empt it. + pub fn satisfied_alternatives(self, present: &[RegionRole]) -> Vec { + self.role_alternatives() + .iter() + .copied() + .filter(|alt| alt.iter().all(|role| present.contains(role))) + .collect() + } + + /// Whether any alternative is satisfied. + pub fn is_satisfied_by(self, present: &[RegionRole]) -> bool { + self.role_alternatives() + .iter() + .any(|alt| alt.iter().all(|role| present.contains(role))) + } + + /// The closest unsatisfied alternative, or `None` when already satisfied. + /// + /// "Closest" is the alternative needing fewest additions. **Ties resolve by + /// declaration order** in [`Programme::role_alternatives`] — `min_by_key` + /// keeps the first minimum, and the alternatives are a fixed `&'static` + /// slice, so the chosen alternative cannot drift as internal iteration + /// changes. That determinism is what makes the diagnostic quotable. + pub fn closest_unsatisfied(self, present: &[RegionRole]) -> Option { + if self.is_satisfied_by(present) { + return None; + } + self.role_alternatives() + .iter() + .copied() + .map(|alt| Unsatisfied { + closest: alt, + missing: alt + .iter() + .copied() + .filter(|role| !present.contains(role)) + .collect(), + }) + .min_by_key(|u| u.missing.len()) + } + + /// Roles missing from the closest alternative. Empty when satisfied. + /// + /// Reporting the closest alternative rather than the first matters: an + /// index holding `gate_up_fused` should be told it needs `down`, not that + /// it needs `gate`, `up` and `down`. + pub fn missing_roles(self, present: &[RegionRole]) -> Vec { + self.closest_unsatisfied(present) + .map(|u| u.missing) + .unwrap_or_default() + } + + /// Whether the expert operates in a latent space rather than the residual + /// stream, and therefore needs the layer's pre/post transforms (§8.1). + pub const fn operates_in_latent_space(self) -> bool { + matches!(self, Self::LatentMoeV1) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const FUSED_PRESENT: [RegionRole; 2] = [RegionRole::GateUpFused, RegionRole::Down]; + const DECOMPOSED_PRESENT: [RegionRole; 3] = + [RegionRole::Gate, RegionRole::Up, RegionRole::Down]; + + #[test] + fn ids_round_trip() { + for p in ALL_PROGRAMMES { + assert_eq!(Programme::from_id(p.id()), Some(p)); + } + } + + #[test] + fn names_round_trip() { + for p in ALL_PROGRAMMES { + assert_eq!(Programme::from_name(p.name()), Some(p)); + } + } + + #[test] + fn ids_are_dense_and_ordered() { + // §8.4 freezes these; a gap or reorder silently rebinds every manifest. + for (index, p) in ALL_PROGRAMMES.iter().enumerate() { + assert_eq!(p.id() as usize, index, "{}", p.name()); + } + } + + #[test] + fn an_unregistered_id_is_refused() { + assert_eq!(Programme::from_id(99), None); + } + + #[test] + fn an_unregistered_name_is_refused() { + assert_eq!(Programme::from_name("mixtral-expert-v9"), None); + } + + #[test] + fn every_programme_accepts_at_least_one_operand_set() { + for p in ALL_PROGRAMMES { + assert!(!p.role_alternatives().is_empty(), "{}", p.name()); + } + } + + #[test] + fn every_alternative_requires_down() { + // Dropping an expert's down yields no expert output at all, so no + // programme can be servable without it. + for p in ALL_PROGRAMMES { + for alt in p.role_alternatives() { + assert!(alt.contains(&RegionRole::Down), "{}", p.name()); + } + } + } + + #[test] + fn a_bank_satisfying_both_layouts_reports_both() { + // The kernel registry may support one layout at a higher maturity than + // the other. Reporting only the first would let a Production fused + // kernel hide behind a Reference decomposed path. + let both = [ + RegionRole::Gate, + RegionRole::Up, + RegionRole::GateUpFused, + RegionRole::Down, + ]; + let alts = Programme::GatedMlpV1.satisfied_alternatives(&both); + assert_eq!(alts.len(), 2, "{alts:?}"); + assert_eq!(alts[0], &[RegionRole::GateUpFused, RegionRole::Down][..]); + assert_eq!( + alts[1], + &[RegionRole::Gate, RegionRole::Up, RegionRole::Down][..] + ); + } + + #[test] + fn satisfied_alternatives_preserves_declaration_order() { + let both = [ + RegionRole::Gate, + RegionRole::Up, + RegionRole::GateUpFused, + RegionRole::Down, + ]; + let alts = Programme::GatedMlpV1.satisfied_alternatives(&both); + let declared = Programme::GatedMlpV1.role_alternatives(); + assert_eq!(alts, declared.to_vec()); + } + + #[test] + fn only_the_satisfied_layouts_are_reported() { + // Fused-only storage satisfies the fused layout and not the other. + let alts = Programme::GatedMlpV1.satisfied_alternatives(&FUSED_PRESENT); + assert_eq!(alts, vec![&[RegionRole::GateUpFused, RegionRole::Down][..]]); + } + + #[test] + fn an_unsatisfied_programme_reports_no_alternatives() { + assert!(Programme::GatedMlpV1 + .satisfied_alternatives(&[RegionRole::Gate]) + .is_empty()); + } + + #[test] + fn a_gated_mlp_accepts_fused_or_decomposed() { + let p = Programme::GatedMlpV1; + assert!(p.is_satisfied_by(&FUSED_PRESENT)); + assert!(p.is_satisfied_by(&DECOMPOSED_PRESENT)); + } + + #[test] + fn the_fused_programme_refuses_decomposed_storage() { + // Naming `gated-mlp-fused-fc1-v1` over decomposed regions is a manifest + // error, not a storage one — `gated-mlp-v1` is the right name there. + let p = Programme::GatedMlpFusedFc1V1; + assert!(p.is_satisfied_by(&FUSED_PRESENT)); + assert!(!p.is_satisfied_by(&DECOMPOSED_PRESENT)); + } + + #[test] + fn gpt_oss_needs_bias() { + let p = Programme::GptOssExpertV1; + assert!(!p.is_satisfied_by(&FUSED_PRESENT)); + let with_bias = [RegionRole::GateUpFused, RegionRole::Down, RegionRole::Bias]; + assert!(p.is_satisfied_by(&with_bias)); + } + + #[test] + fn missing_roles_is_empty_when_satisfied() { + assert!(Programme::GatedMlpV1 + .missing_roles(&FUSED_PRESENT) + .is_empty()); + } + + #[test] + fn missing_roles_reports_the_closest_alternative() { + // Holding only gate_up_fused, the useful answer is "you need down", + // not "you need gate, up and down". + let present = [RegionRole::GateUpFused]; + assert_eq!( + Programme::GatedMlpV1.missing_roles(&present), + vec![RegionRole::Down] + ); + } + + #[test] + fn missing_roles_reports_everything_when_nothing_is_present() { + let missing = Programme::GatedMlpV1.missing_roles(&[]); + assert_eq!(missing, vec![RegionRole::GateUpFused, RegionRole::Down]); + } + + #[test] + fn a_tie_resolves_to_the_first_declared_alternative() { + // Nothing present: fused needs 2 additions, decomposed needs 3, so no + // tie there. Construct a real tie by supplying the roles that make both + // alternatives equidistant, and require the DECLARED-FIRST one wins. + let present = [RegionRole::Gate, RegionRole::Up, RegionRole::GateUpFused]; + let u = Programme::GatedMlpV1.closest_unsatisfied(&present).unwrap(); + assert_eq!(u.missing, vec![RegionRole::Down]); + // Both alternatives are missing exactly `down`; fused is declared first. + assert_eq!(u.closest, &[RegionRole::GateUpFused, RegionRole::Down][..]); + } + + #[test] + fn the_chosen_alternative_is_stable_across_repeated_calls() { + // The alternatives are a fixed &'static slice, so this cannot drift — + // pinned because a drifting diagnostic is a diagnostic nobody trusts. + let present = [RegionRole::Gate]; + let first = Programme::GatedMlpV1.closest_unsatisfied(&present); + for _ in 0..8 { + assert_eq!(Programme::GatedMlpV1.closest_unsatisfied(&present), first); + } + } + + #[test] + fn the_description_names_the_layout_it_judged_against() { + let present = [RegionRole::GateUpFused]; + let u = Programme::GatedMlpV1.closest_unsatisfied(&present).unwrap(); + assert_eq!( + u.describe(), + "closest accepted layout: gate_up_fused + down; missing down" + ); + } + + #[test] + fn the_description_lists_every_missing_role() { + let u = Programme::GptOssExpertV1 + .closest_unsatisfied(&[RegionRole::GateUpFused]) + .unwrap(); + let text = u.describe(); + assert!(text.contains("gate_up_fused + down + bias"), "{text}"); + assert!(text.contains("missing down, bias"), "{text}"); + } + + #[test] + fn a_satisfied_programme_has_no_unsatisfied_report() { + assert_eq!( + Programme::GatedMlpV1.closest_unsatisfied(&FUSED_PRESENT), + None + ); + } + + #[test] + fn a_browse_slice_satisfies_no_programme() { + // Gate-only regions are the §15.5 analysis-only slice: representable, + // never executable. + let gate_only = [RegionRole::Gate]; + for p in ALL_PROGRAMMES { + assert!(!p.is_satisfied_by(&gate_only), "{}", p.name()); + assert!(!p.missing_roles(&gate_only).is_empty(), "{}", p.name()); + } + } + + #[test] + fn only_the_latent_programme_needs_the_layer_transforms() { + for p in ALL_PROGRAMMES { + assert_eq!( + p.operates_in_latent_space(), + p == Programme::LatentMoeV1, + "{}", + p.name() + ); + } + } + + #[test] + fn extra_present_roles_do_not_prevent_satisfaction() { + // §6.5: presence of other roles never invalidates a file. + let extra = [ + RegionRole::GateUpFused, + RegionRole::Down, + RegionRole::Scales, + RegionRole::Unknown(900), + ]; + assert!(Programme::GatedMlpV1.is_satisfied_by(&extra)); + } +} diff --git a/crates/larql-vindex/src/format/moe_manifest/router.rs b/crates/larql-vindex/src/format/moe_manifest/router.rs new file mode 100644 index 000000000..f254acc6c --- /dev/null +++ b/crates/larql-vindex/src/format/moe_manifest/router.rs @@ -0,0 +1,258 @@ +//! Router descriptor (format spec §8.1, §8.2). +//! +//! The manifest *names* routing semantics; the adapter implements them. The +//! vocabulary here is sized by the conformance envelope, which spans four +//! genuinely different balancing mechanisms: +//! +//! | model | scoring | selection | post-processing | +//! | ----- | ------- | --------- | --------------- | +//! | Direct / Gemma | softmax | top-2/4 | — | +//! | Kimi-Linear-48B | sigmoid | top-8 of 256, grouped | renormalise, scale 2.446 | +//! | Inkling-Small | sigmoid | top-6 of 256 | gate bias, norm-after-top-k, route scale 8.0, global scale, **shared-expert sink** | +//! | K3 | — | top-16 of 896 | quantile-balanced | +//! +//! Every one of those is a combination of the same handful of knobs, which is +//! the claim this type exists to make falsifiable: if a real model needs a +//! knob that is not here, the manifest vocabulary was under-specified and that +//! is an ABI finding, not a post-freeze patch. + +use serde::{Deserialize, Serialize}; + +/// How raw router logits become scores. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ScoreActivation { + Softmax, + Sigmoid, +} + +/// How experts are picked from the scores. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "kind", rename_all = "snake_case")] +pub enum Selection { + /// Plain top-k over all experts. + TopK { k: usize }, + /// Top-k within each of `groups` partitions. Kimi-Linear declares one + /// group, which is degenerate but present in the released schema — the + /// manifest records what the model says rather than simplifying it away. + GroupedTopK { k: usize, groups: usize }, +} + +impl Selection { + /// Experts activated per token. + pub const fn experts_per_token(&self) -> usize { + match self { + Self::TopK { k } => *k, + Self::GroupedTopK { k, groups } => *k * *groups, + } + } +} + +/// Post-selection score handling. Order matters and is fixed: bias is applied +/// to scores **before** selection; normalisation and scaling apply **after**. +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] +#[serde(default)] +pub struct RouterPostProcessing { + /// Renormalise the selected scores to sum to 1. OLMoE deliberately does + /// *not* — it keeps raw softmax probabilities — so this is opt-in. + pub renormalise: bool, + /// Normalise after top-k rather than before (Inkling's `norm_after_topk`). + pub norm_after_top_k: bool, + /// Per-expert additive bias applied to scores before selection. + pub gate_bias: Option, + /// Constant multiplier on the routed branch's contribution. + pub route_scale: Option, + /// K3's balanced assignment. Named, not described — the adapter owns it. + pub quantile_balanced: bool, +} + +/// A layer's routing description. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct Router { + /// Manifest key of the router weight matrix. + pub scores: String, + pub activation: ScoreActivation, + pub selection: Selection, + #[serde(default)] + pub post: RouterPostProcessing, + /// Whether shared experts are scored by the router and participate in + /// normalisation (Inkling's `shared_expert_sink`). + /// + /// This is not cosmetic: with a sink, the shared experts' weights come out + /// of the same normalised budget as the routed ones, so a reduction that + /// ignores it silently over-weights the routed branch. + #[serde(default)] + pub shared_expert_sink: bool, +} + +impl Router { + /// Whether the router's own output needs no further weighting — true only + /// for a plain softmax top-k with nothing applied after it. + pub fn is_plain_softmax_top_k(&self) -> bool { + self.activation == ScoreActivation::Softmax + && matches!(self.selection, Selection::TopK { .. }) + && self.post == RouterPostProcessing::default() + && !self.shared_expert_sink + } + + /// Names of manifest-addressed tensors this router needs beyond `scores`. + pub fn auxiliary_tensors(&self) -> Vec<&str> { + self.post.gate_bias.as_deref().into_iter().collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn gemma_style() -> Router { + Router { + scores: "layers.0.router.weight".into(), + activation: ScoreActivation::Softmax, + selection: Selection::TopK { k: 2 }, + post: RouterPostProcessing::default(), + shared_expert_sink: false, + } + } + + fn kimi_linear_style() -> Router { + Router { + scores: "layers.1.router.weight".into(), + activation: ScoreActivation::Sigmoid, + selection: Selection::GroupedTopK { k: 8, groups: 1 }, + post: RouterPostProcessing { + renormalise: true, + route_scale: Some(2.446), + ..Default::default() + }, + shared_expert_sink: false, + } + } + + fn inkling_style() -> Router { + Router { + scores: "layers.3.router.weight".into(), + activation: ScoreActivation::Sigmoid, + selection: Selection::TopK { k: 6 }, + post: RouterPostProcessing { + norm_after_top_k: true, + gate_bias: Some("layers.3.router.bias".into()), + route_scale: Some(8.0), + ..Default::default() + }, + shared_expert_sink: true, + } + } + + fn k3_style() -> Router { + Router { + scores: "layers.12.router.weight".into(), + activation: ScoreActivation::Softmax, + selection: Selection::TopK { k: 16 }, + post: RouterPostProcessing { + quantile_balanced: true, + ..Default::default() + }, + shared_expert_sink: false, + } + } + + #[test] + fn every_envelope_router_round_trips_through_json() { + for r in [ + gemma_style(), + kimi_linear_style(), + inkling_style(), + k3_style(), + ] { + let json = serde_json::to_string(&r).unwrap(); + let back: Router = serde_json::from_str(&json).unwrap(); + assert_eq!(back, r, "{json}"); + } + } + + #[test] + fn plain_top_k_counts_its_own_k() { + assert_eq!(Selection::TopK { k: 16 }.experts_per_token(), 16); + } + + #[test] + fn grouped_top_k_counts_k_per_group() { + assert_eq!( + Selection::GroupedTopK { k: 8, groups: 4 }.experts_per_token(), + 32 + ); + } + + #[test] + fn a_single_group_is_equivalent_to_plain_top_k_in_count() { + // Kimi-Linear's degenerate case: recorded as grouped, counts as top-8. + assert_eq!( + Selection::GroupedTopK { k: 8, groups: 1 }.experts_per_token(), + Selection::TopK { k: 8 }.experts_per_token() + ); + } + + #[test] + fn only_the_bare_softmax_router_is_plain() { + assert!(gemma_style().is_plain_softmax_top_k()); + assert!(!kimi_linear_style().is_plain_softmax_top_k()); + assert!(!inkling_style().is_plain_softmax_top_k()); + // K3 is softmax top-k but quantile-balanced, so it is not plain. + assert!(!k3_style().is_plain_softmax_top_k()); + } + + #[test] + fn a_sink_router_is_never_plain_even_with_softmax_top_k() { + // The sink changes the reduction's weight budget; treating it as plain + // would silently over-weight the routed branch. + let sink = Router { + shared_expert_sink: true, + ..gemma_style() + }; + assert!(!sink.is_plain_softmax_top_k()); + } + + #[test] + fn gate_bias_is_reported_as_an_auxiliary_tensor() { + assert_eq!( + inkling_style().auxiliary_tensors(), + vec!["layers.3.router.bias"] + ); + } + + #[test] + fn a_router_without_bias_needs_no_auxiliary_tensors() { + assert!(gemma_style().auxiliary_tensors().is_empty()); + } + + #[test] + fn post_processing_defaults_to_doing_nothing() { + // OLMoE keeps raw softmax probabilities; renormalising by default + // would silently change every model that does not ask for it. + let d = RouterPostProcessing::default(); + assert!(!d.renormalise); + assert!(!d.norm_after_top_k); + assert!(!d.quantile_balanced); + assert_eq!(d.gate_bias, None); + assert_eq!(d.route_scale, None); + } + + #[test] + fn an_absent_post_block_deserialises_to_the_default() { + let json = r#"{ + "scores": "r.weight", + "activation": "softmax", + "selection": {"kind": "top_k", "k": 2} + }"#; + let r: Router = serde_json::from_str(json).unwrap(); + assert_eq!(r.post, RouterPostProcessing::default()); + assert!(!r.shared_expert_sink); + } + + #[test] + fn selection_is_tagged_so_the_two_kinds_cannot_be_confused() { + let json = serde_json::to_string(&Selection::GroupedTopK { k: 8, groups: 1 }).unwrap(); + assert!(json.contains("grouped_top_k"), "{json}"); + } +} diff --git a/crates/larql-vindex/src/format/vindex3/index.rs b/crates/larql-vindex/src/format/vindex3/index.rs new file mode 100644 index 000000000..7acedb0fd --- /dev/null +++ b/crates/larql-vindex/src/format/vindex3/index.rs @@ -0,0 +1,297 @@ +//! `index.json` for a VINDEX3 container — the sole root authority (spec §12). +//! +//! Deliberately a **different type** from [`VindexConfig`](crate::config::index::VindexConfig), +//! not a superset of it. The shipped generation's index describes a dense +//! Gemma-shaped extraction: `layers`, `down_top_k`, `intermediate_size`. A +//! VINDEX3 index describes a *catalogue* — which segments exist, which +//! manifest interprets them, which profiles may be selected. Growing one +//! struct to cover both would give every field an "unless version 3" caveat, +//! and the loader would sniff fields to decide which half it is looking at, +//! which is exactly the heuristic dispatch §12.1 forbids. + +use std::collections::BTreeMap; + +use serde::{Deserialize, Serialize}; + +use super::profile::{Profile, ProfileSelectionError, ResolvedProfile}; +use super::variants::VariantCatalogue; +use crate::format::generation::V3_CURRENT_SCHEMA; + +pub use super::profile::PROFILE_EXACT; + +/// `index.json` as a VINDEX3 container writes it. +/// +/// `segments` maps a segment *key* to the number of physical files under it. +/// A key is a path stem relative to the container root (`routed/layer_000`), +/// so the loader composes a filename rather than globbing a directory — +/// filename sniffing is what §12.1 rules out. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct Vindex3Index { + /// Always [`V3_CURRENT_SCHEMA`]. The sole generation discriminator. + pub version: u32, + /// Identity, carried so a container names itself without a sidecar. + pub model: String, + pub family: String, + /// Residual width the model operates at. + pub hidden_size: usize, + pub num_layers: usize, + /// Filename of the MoE programme manifest, relative to the root. + pub moe_manifest: String, + /// Profiles this container declares. Never empty: a container with no + /// selectable profile cannot be served, and discovering that at bind time + /// rather than at load time is the failure mode profiles exist to prevent. + /// + /// A profile selects variants (§9.1); it is not just a name. See + /// [`Profile`]. + pub profiles: Vec, + /// Region set → the variants physically present for it (§9.1). + /// + /// Defaulted, because a container with no alternative packs — every one + /// written so far — catalogues nothing and is served entirely from its + /// baselines. An absent catalogue means "no packs", never "unchecked". + #[serde(default, skip_serializing_if = "VariantCatalogue::is_empty")] + pub variants: VariantCatalogue, + /// Segment key → physical file count. + pub segments: BTreeMap, +} + +impl Vindex3Index { + /// A single-profile container over the given segments. + pub fn new( + model: impl Into, + family: impl Into, + hidden_size: usize, + num_layers: usize, + moe_manifest: impl Into, + segments: BTreeMap, + ) -> Self { + Self { + version: V3_CURRENT_SCHEMA, + model: model.into(), + family: family.into(), + hidden_size, + num_layers, + moe_manifest: moe_manifest.into(), + profiles: vec![Profile::exact()], + variants: VariantCatalogue::new(), + segments, + } + } + + /// Whether `profile` is one this container declares by name. + /// + /// A name check, and *only* a name check — it says nothing about whether + /// the profile's selections can be honoured. Use [`Self::select_profile`] + /// before binding; §9.1's requirement is that an absent variant is refused + /// before a byte is read, and a name cannot carry that. + pub fn declares_profile(&self, profile: &str) -> bool { + self.profile(profile).is_some() + } + + /// The declared profile called `name`. + pub fn profile(&self, name: &str) -> Option<&Profile> { + self.profiles.iter().find(|p| p.name == name) + } + + /// Resolve `name` against the variants this container physically carries. + /// + /// The §9.1 gate — see [`super::profile::select`]. + pub fn select_profile(&self, name: &str) -> Result, ProfileSelectionError> { + super::profile::select(self, name) + } + + /// Declared profile names, in declaration order. + pub fn profile_names(&self) -> Vec<&str> { + self.profiles.iter().map(|p| p.name.as_str()).collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn segments() -> BTreeMap { + BTreeMap::from([("routed/layer_000".to_string(), 1)]) + } + + fn index() -> Vindex3Index { + Vindex3Index::new( + "fixture-a", + "direct-moe", + 256, + 1, + "moe_manifest.json", + segments(), + ) + } + + #[test] + fn a_fresh_index_declares_the_successor_schema() { + // The whole dispatch turns on this one number being 3 and nothing else. + assert_eq!(index().version, V3_CURRENT_SCHEMA); + } + + #[test] + fn a_fresh_index_is_servable_because_it_declares_a_profile() { + let i = index(); + assert!(i.declares_profile(PROFILE_EXACT)); + assert!( + !i.profiles.is_empty(), + "a profile-less container cannot serve" + ); + } + + #[test] + fn an_undeclared_profile_is_refused_rather_than_assumed() { + assert!(!index().declares_profile("browse")); + } + + #[test] + fn the_index_round_trips_through_json() { + // It is the root authority; if it cannot survive its own serialisation + // nothing below it is reachable. + let before = index(); + let json = serde_json::to_string(&before).unwrap(); + let after: Vindex3Index = serde_json::from_str(&json).unwrap(); + assert_eq!(before, after); + } + + // ── Profile selection (§9.1) ──────────────────────────────────────── + + use super::super::variants::{RegionSetVariants, StoredVariant}; + use crate::format::capability::authority::Fidelity; + + const GATE_UP: &str = "layer.0.routed.gate_up"; + const Q6K: &str = "exact-q6k"; + const MXFP4: &str = "native-mxfp4"; + + /// An index carrying one region set with an MXFP4 pack beside its baseline, + /// and a profile that selects the pack. + fn packed() -> Vindex3Index { + let mut i = index(); + i.variants = VariantCatalogue::new().with_set( + GATE_UP, + RegionSetVariants::single( + Q6K, + StoredVariant::new("routed/layer_000.q6k", Fidelity::SourceEquivalent), + ) + .with_variant( + MXFP4, + StoredVariant::new("routed/layer_000.mxfp4", Fidelity::SourceExact), + ), + ); + i.profiles + .push(Profile::new("routed-mxfp4").selecting(GATE_UP, MXFP4)); + i + } + + #[test] + fn a_packless_index_still_resolves_its_canonical_profile() { + let i = index(); + let resolved = i.select_profile(PROFILE_EXACT).unwrap(); + assert_eq!(resolved.name(), PROFILE_EXACT); + assert!(resolved.is_empty(), "nothing catalogued, nothing selected"); + } + + #[test] + fn a_profile_resolves_to_the_variant_it_selected() { + let i = packed(); + let exact = i.select_profile(PROFILE_EXACT).unwrap(); + assert_eq!( + exact.variant_for(GATE_UP).unwrap().storage, + "routed/layer_000.q6k", + "the canonical profile takes the baseline" + ); + let packed_profile = i.select_profile("routed-mxfp4").unwrap(); + assert_eq!( + packed_profile.variant_for(GATE_UP).unwrap().storage, + "routed/layer_000.mxfp4", + "selecting a pack must change which bytes are named" + ); + } + + #[test] + fn an_undeclared_profile_names_the_ones_that_are_declared() { + let err = packed().select_profile("browse").unwrap_err(); + match &err { + ProfileSelectionError::UndeclaredProfile { + requested, + declared, + } => { + assert_eq!(requested, "browse"); + assert_eq!( + declared, + &vec![PROFILE_EXACT.to_string(), "routed-mxfp4".to_string()] + ); + } + other => panic!("wrong error: {other:?}"), + } + let msg = err.to_string(); + assert!(msg.contains("browse"), "{msg}"); + assert!(msg.contains("routed-mxfp4"), "{msg}"); + } + + #[test] + fn a_declared_profile_selecting_an_absent_variant_is_unsatisfiable() { + // Declared-but-unservable is the case `declares_profile` cannot see, + // and the whole reason selection exists beside it. + let mut i = packed(); + i.profiles + .push(Profile::new("ghost").selecting(GATE_UP, "nvfp4")); + assert!( + i.declares_profile("ghost"), + "the name check passes — that is the point" + ); + let err = i.select_profile("ghost").unwrap_err(); + match &err { + ProfileSelectionError::Unsatisfiable { profile, defects } => { + assert_eq!(profile, "ghost"); + assert_eq!(defects.len(), 1); + } + other => panic!("wrong error: {other:?}"), + } + // The message has to carry all three of §9.1's names. + let msg = err.to_string(); + assert!(msg.contains("ghost"), "profile: {msg}"); + assert!(msg.contains(GATE_UP), "region set: {msg}"); + assert!(msg.contains("nvfp4"), "requested: {msg}"); + assert!(msg.contains(MXFP4), "present: {msg}"); + } + + #[test] + fn profile_lookup_and_names_agree() { + let i = packed(); + assert_eq!(i.profile_names(), vec![PROFILE_EXACT, "routed-mxfp4"]); + assert_eq!(i.profile("routed-mxfp4").unwrap().name, "routed-mxfp4"); + assert!(i.profile("browse").is_none()); + assert!(i.declares_profile("routed-mxfp4")); + } + + #[test] + fn a_packed_index_round_trips_with_its_catalogue_and_selections() { + let before = packed(); + let after: Vindex3Index = + serde_json::from_str(&serde_json::to_string(&before).unwrap()).unwrap(); + assert_eq!(before, after); + assert_eq!( + after.select_profile("routed-mxfp4").unwrap().storage_keys(), + vec!["routed/layer_000.mxfp4"] + ); + } + + #[test] + fn a_packless_index_omits_the_catalogue_on_the_wire() { + // Every container written so far has no packs; an empty object in each + // of their index.json files would be noise that reads as meaningful. + let json = serde_json::to_value(index()).unwrap(); + assert!(json.get("variants").is_none(), "{json}"); + } + + #[test] + fn the_serialised_form_carries_the_version_a_detector_reads() { + // `detect_generation` parses the raw JSON, not this struct, so the + // wire key matters independently of the field name. + let json: serde_json::Value = serde_json::to_value(index()).unwrap(); + assert_eq!(json["version"], serde_json::json!(V3_CURRENT_SCHEMA)); + } +} diff --git a/crates/larql-vindex/src/format/vindex3/mod.rs b/crates/larql-vindex/src/format/vindex3/mod.rs new file mode 100644 index 000000000..e52499c0a --- /dev/null +++ b/crates/larql-vindex/src/format/vindex3/mod.rs @@ -0,0 +1,54 @@ +//! The VINDEX3 container — writing one, and opening one. +//! +//! # Why this module had to exist before "VINDEX3 works" could be said +//! +//! The VINDEX3 execution runtime ([`crate::runtime`]) was validated to +//! bit-identity against the shipped path, repeatedly and on a real 26B model. +//! But every one of those runs sourced its operands from a **VINDEX2** file: +//! `index.json.version` was 2 at every step. Nothing could write a VINDEX3 +//! container, so `ContainerGeneration::V3` appeared in exactly one test, +//! constructed from a hand-written JSON string. +//! +//! That made the honest claim narrower than it sounded: +//! +//! ```text +//! proven the VINDEX3 runtime, fed VINDEX2 operands, matches production +//! not proven a VINDEX3 container can be written, opened, bound and executed +//! ``` +//! +//! This module closes the second line for the control fixture. The middle of +//! the stack — LYRW v2 tables, the programme manifest, capability and +//! authority, the bound executor — was already built; what was missing was the +//! two ends, and an end that nothing exercises is an end nobody has debugged. +//! +//! # Layout +//! +//! ```text +//! / +//! ├── index.json schema 3 — sole root authority (§12) +//! ├── moe_manifest.json programme description (§8) +//! └── routed/layer_000.lyrw LYRW v2 bank (§6) +//! ``` +//! +//! # What is deliberately not here +//! +//! No transcoding. A container assembler that also converted formats would be +//! the silent conversion §9.1 forbids, buried one layer below where anyone +//! would look for it. Regions are placed exactly as their producer wrote them. + +pub mod index; +pub mod profile; +pub mod read; +/// Conformance fixture A, public so integration tests and future gate arms can +/// build a real container without duplicating its frozen dimensions. +pub mod test_support; +pub mod variants; +pub mod verify; +pub mod write; + +pub use index::{Vindex3Index, PROFILE_EXACT}; +pub use profile::{Profile, ProfileSelectionError, ResolvedProfile}; +pub use read::Vindex3Container; +pub use variants::{RegionSetVariants, StoredVariant, VariantCatalogue, VariantDefect}; +pub use verify::ContainerDefect; +pub use write::{segment_path, write_container, ContainerSpec, SegmentSource, MOE_MANIFEST_JSON}; diff --git a/crates/larql-vindex/src/format/vindex3/profile.rs b/crates/larql-vindex/src/format/vindex3/profile.rs new file mode 100644 index 000000000..8f151057b --- /dev/null +++ b/crates/larql-vindex/src/format/vindex3/profile.rs @@ -0,0 +1,213 @@ +//! A profile — the thing that selects variants (spec §9.1). +//! +//! # Why a name was not enough +//! +//! A profile used to be a string in a list, and `declares_profile` answered +//! "is that string one of ours". That check cannot fail the way §9.1 needs it +//! to, because a name carries no claim about bytes: a container could declare +//! `"routed-mxfp4"`, carry no MXFP4 pack, and pass. The absence surfaced at +//! bind time as a missing region, one layer below where anyone would look for +//! it, with the profile name nowhere in the message. +//! +//! So a profile names, per region set, which variant it selects. A region set +//! the profile says nothing about takes that set's baseline — which keeps the +//! common profile (everything canonical) empty rather than a restatement of +//! the whole catalogue, and keeps a new region set from silently becoming +//! unselectable the moment it is added. +//! +//! # Resolution is total, not first-failure +//! +//! [`Profile::resolve`] returns every defect it finds. A profile that names +//! two absent packs should take one repair pass; being shown the first, fixing +//! it, and being shown the second is how a load-time check teaches people to +//! bypass it. + +use std::collections::BTreeMap; + +use serde::{Deserialize, Serialize}; + +use super::variants::{StoredVariant, VariantCatalogue, VariantDefect}; + +/// Profile name every container carries: full-fidelity execution. +pub const PROFILE_EXACT: &str = "exact"; + +/// A named selection over the variant catalogue. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct Profile { + pub name: String, + /// Region set → variant name. Absent means "take the baseline", so this is + /// empty for a profile that wants everything canonical. + #[serde(default, skip_serializing_if = "BTreeMap::is_empty")] + pub selects: BTreeMap, +} + +impl Profile { + /// A profile that takes every region set's baseline. + pub fn new(name: impl Into) -> Self { + Self { + name: name.into(), + selects: BTreeMap::new(), + } + } + + /// The canonical profile: every baseline, nothing overridden. + pub fn exact() -> Self { + Self::new(PROFILE_EXACT) + } + + /// Override one region set's variant. + pub fn selecting(mut self, region_set: impl Into, variant: impl Into) -> Self { + self.selects.insert(region_set.into(), variant.into()); + self + } + + /// Resolve this profile against what the container physically carries. + /// + /// Every catalogued region set appears in the result — the ones this + /// profile names at its chosen variant, the rest at their baseline — so a + /// caller holds a complete map of the bytes it is about to execute, and + /// never has to fall back to "look it up again later", which is where a + /// silent conversion would get its opportunity. + pub fn resolve<'a>( + &self, + catalogue: &'a VariantCatalogue, + ) -> Result, Vec> { + let mut defects = Vec::new(); + + // Selections first: a name this profile got wrong is the failure we + // are here to report, and reporting it before the catalogue's own + // structural defects puts the profile's error at the top. + for (region_set, variant) in &self.selects { + if let Err(defect) = catalogue.select(region_set, variant) { + defects.push(defect); + } + } + defects.extend(catalogue.defects()); + + if !defects.is_empty() { + return Err(defects); + } + + let mut chosen = BTreeMap::new(); + for region_set in catalogue.region_sets() { + let stored = match self.selects.get(®ion_set) { + Some(variant) => catalogue.select(®ion_set, variant), + None => catalogue.select_baseline(®ion_set), + } + .expect("defects were collected above, so every lookup resolves"); + chosen.insert(region_set, stored); + } + Ok(ResolvedProfile { + name: self.name.clone(), + chosen, + }) + } +} + +/// A profile that resolved — every region set bound to a present variant. +/// +/// Holding borrows of the catalogue rather than copies is deliberate: this +/// cannot outlive the index it was resolved against, so it cannot be carried +/// past a reload and quietly describe bytes that were replaced. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ResolvedProfile<'a> { + name: String, + chosen: BTreeMap, +} + +impl<'a> ResolvedProfile<'a> { + pub fn name(&self) -> &str { + &self.name + } + + /// The variant this profile executes for `region_set`. + pub fn variant_for(&self, region_set: &str) -> Option<&'a StoredVariant> { + self.chosen.get(region_set).copied() + } + + /// Every (region set, variant) pair, in region-set order. + pub fn entries(&self) -> impl Iterator + '_ { + self.chosen.iter().map(|(k, v)| (k.as_str(), *v)) + } + + pub fn len(&self) -> usize { + self.chosen.len() + } + + pub fn is_empty(&self) -> bool { + self.chosen.is_empty() + } + + /// Segment keys this profile will read, deduplicated and sorted. + /// + /// What a loader needs in order to touch only the packs in play — the + /// reason a routed-MXFP4 profile does not page in the Q6_K baseline. + pub fn storage_keys(&self) -> Vec<&'a str> { + let mut keys: Vec<&str> = self.chosen.values().map(|v| v.storage.as_str()).collect(); + keys.sort_unstable(); + keys.dedup(); + keys + } +} + +/// Why a container cannot serve the profile it was asked for. +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub enum ProfileSelectionError { + #[error("no profile '{requested}' in this container; declared: [{}]", declared.join(", "))] + UndeclaredProfile { + requested: String, + declared: Vec, + }, + + /// Declared, but its selections do not match the bytes on disk — the + /// failure §9.1 exists to catch. Carries every defect, so one repair pass + /// clears them all. + #[error( + "profile '{profile}' cannot be served: {}", + defects.iter().map(|d| d.to_string()).collect::>().join("; ") + )] + Unsatisfiable { + profile: String, + defects: Vec, + }, +} + +/// Resolve a declared profile of `index` against the variants it carries. +/// +/// Lives here rather than on the index because every rule it enforces is a +/// profile rule: the index only holds the two halves side by side. +pub fn select<'a>( + index: &'a super::index::Vindex3Index, + name: &str, +) -> Result, ProfileSelectionError> { + let profile = index + .profile(name) + .ok_or_else(|| ProfileSelectionError::UndeclaredProfile { + requested: name.to_string(), + declared: index + .profile_names() + .into_iter() + .map(String::from) + .collect(), + })?; + profile + .resolve(&index.variants) + .map_err(|defects| ProfileSelectionError::Unsatisfiable { + profile: name.to_string(), + defects, + }) +} + +/// Resolve every profile `index` declares, for a caller that must know the +/// whole container is servable — `Vindex3Container::open` does this before +/// reading a segment byte. +pub fn resolve_all(index: &super::index::Vindex3Index) -> Result<(), ProfileSelectionError> { + for name in index.profile_names() { + select(index, name)?; + } + Ok(()) +} + +#[cfg(test)] +#[path = "profile_tests.rs"] +mod tests; diff --git a/crates/larql-vindex/src/format/vindex3/profile_tests.rs b/crates/larql-vindex/src/format/vindex3/profile_tests.rs new file mode 100644 index 000000000..26f21cdfc --- /dev/null +++ b/crates/larql-vindex/src/format/vindex3/profile_tests.rs @@ -0,0 +1,195 @@ +//! Tests for profile resolution (§9.1). + +use super::*; +use crate::format::capability::authority::Fidelity; +use crate::format::vindex3::variants::{RegionSetVariants, StoredVariant}; + +const Q6K: &str = "exact-q6k"; +const MXFP4: &str = "native-mxfp4"; +const GATE_UP: &str = "layer.12.routed.gate_up"; +const DOWN: &str = "layer.12.routed.down"; + +fn q6k() -> StoredVariant { + StoredVariant::new("routed/layer_012.q6k", Fidelity::SourceEquivalent) +} + +fn mxfp4() -> StoredVariant { + StoredVariant::new("routed/layer_012.mxfp4", Fidelity::SourceExact) +} + +/// Two region sets; only `gate_up` has an MXFP4 pack beside its baseline. +fn catalogue() -> VariantCatalogue { + VariantCatalogue::new() + .with_set( + GATE_UP, + RegionSetVariants::single(Q6K, q6k()).with_variant(MXFP4, mxfp4()), + ) + .with_set(DOWN, RegionSetVariants::single(Q6K, q6k())) +} + +// ── The canonical profile ──────────────────────────────────────────────── + +#[test] +fn the_exact_profile_takes_every_baseline_without_naming_any() { + let c = catalogue(); + let resolved = Profile::exact().resolve(&c).unwrap(); + assert_eq!(resolved.name(), PROFILE_EXACT); + assert_eq!(resolved.len(), 2); + assert_eq!(resolved.variant_for(GATE_UP), Some(&q6k())); + assert_eq!(resolved.variant_for(DOWN), Some(&q6k())); +} + +#[test] +fn the_canonical_profile_carries_no_selections_at_all() { + // If "everything canonical" had to restate the catalogue, every new region + // set would need every profile edited, and the one that was forgotten + // would be the one that broke. + assert!(Profile::exact().selects.is_empty()); +} + +#[test] +fn a_region_set_the_profile_says_nothing_about_still_resolves() { + // `down` has no MXFP4 pack and the profile does not mention it; it must + // still appear, at its baseline, rather than dropping out of the result. + let profile = Profile::new("routed-mxfp4").selecting(GATE_UP, MXFP4); + let c = catalogue(); + let resolved = profile.resolve(&c).unwrap(); + assert_eq!(resolved.variant_for(GATE_UP), Some(&mxfp4())); + assert_eq!(resolved.variant_for(DOWN), Some(&q6k())); +} + +// ── Selecting an absent variant — the refusal §9.1 asks for ───────────── + +#[test] +fn selecting_an_absent_variant_fails_naming_set_request_and_present() { + let profile = Profile::new("routed-mxfp4").selecting(DOWN, MXFP4); + let defects = profile.resolve(&catalogue()).unwrap_err(); + match defects.as_slice() { + [VariantDefect::AbsentVariant { + region_set, + requested, + present, + }] => { + assert_eq!(region_set, DOWN); + assert_eq!(requested, MXFP4); + assert_eq!(present, &vec![Q6K.to_string()]); + } + other => panic!("wrong defects: {other:?}"), + } +} + +#[test] +fn selecting_an_uncatalogued_region_set_fails_naming_what_exists() { + let profile = Profile::new("typo").selecting("layer.12.routed.gate-up", MXFP4); + let defects = profile.resolve(&catalogue()).unwrap_err(); + assert!(matches!( + defects.as_slice(), + [VariantDefect::UnknownRegionSet { .. }] + )); + assert!(defects[0].to_string().contains(GATE_UP)); +} + +#[test] +fn every_bad_selection_is_reported_in_one_pass() { + let profile = Profile::new("two-wrongs") + .selecting(DOWN, MXFP4) + .selecting("nope", Q6K); + let defects = profile.resolve(&catalogue()).unwrap_err(); + assert_eq!(defects.len(), 2, "{defects:?}"); +} + +#[test] +fn a_failed_resolution_yields_no_partial_selection() { + // Fail closed: the caller gets defects or a complete map, never a map with + // the resolvable half filled in. A partial map is what would let a load + // continue far enough to read bytes for the region sets that did resolve. + let profile = Profile::new("routed-mxfp4").selecting(DOWN, MXFP4); + assert!(profile.resolve(&catalogue()).is_err()); +} + +#[test] +fn a_structurally_broken_catalogue_fails_even_a_profile_that_names_nothing() { + // The baseline is what an unnamed region set resolves to, so a missing + // baseline breaks the canonical profile too — it is not only a problem for + // profiles that select. + let broken = VariantCatalogue::new().with_set( + GATE_UP, + RegionSetVariants { + baseline: "gone".into(), + variants: std::collections::BTreeMap::from([(Q6K.to_string(), q6k())]), + }, + ); + let defects = Profile::exact().resolve(&broken).unwrap_err(); + assert!(matches!( + defects.as_slice(), + [VariantDefect::BaselineAbsent { .. }] + )); +} + +// ── What a resolved profile is for ────────────────────────────────────── + +#[test] +fn a_resolved_profile_lists_only_the_storage_it_will_touch() { + // The incremental-pack payoff: a profile on the MXFP4 pack must not name + // the Q6_K file for that region set, or nothing was saved by packing. + let profile = Profile::new("routed-mxfp4").selecting(GATE_UP, MXFP4); + let c = catalogue(); + let resolved = profile.resolve(&c).unwrap(); + assert_eq!( + resolved.storage_keys(), + vec!["routed/layer_012.mxfp4", "routed/layer_012.q6k"] + ); +} + +#[test] +fn storage_keys_are_deduplicated() { + // Two region sets can live in one segment file; a loader asked to read it + // twice would double the resident bytes. + let shared = VariantCatalogue::new() + .with_set(GATE_UP, RegionSetVariants::single(Q6K, q6k())) + .with_set(DOWN, RegionSetVariants::single(Q6K, q6k())); + let resolved = Profile::exact().resolve(&shared).unwrap(); + assert_eq!(resolved.storage_keys(), vec!["routed/layer_012.q6k"]); +} + +#[test] +fn entries_come_out_in_region_set_order() { + let c = catalogue(); + let resolved = Profile::exact().resolve(&c).unwrap(); + let names: Vec<&str> = resolved.entries().map(|(k, _)| k).collect(); + assert_eq!(names, vec![DOWN, GATE_UP]); +} + +#[test] +fn an_empty_catalogue_resolves_to_an_empty_profile() { + let c = VariantCatalogue::new(); + let resolved = Profile::exact().resolve(&c).unwrap(); + assert!(resolved.is_empty()); + assert!(resolved.storage_keys().is_empty()); + assert_eq!(resolved.variant_for(GATE_UP), None); +} + +// ── Wire form ──────────────────────────────────────────────────────────── + +#[test] +fn a_profile_round_trips_through_json() { + let before = Profile::new("routed-mxfp4").selecting(GATE_UP, MXFP4); + let json = serde_json::to_string(&before).unwrap(); + let after: Profile = serde_json::from_str(&json).unwrap(); + assert_eq!(before, after); +} + +#[test] +fn a_selectionless_profile_omits_the_empty_map_on_the_wire() { + // Every container carries the canonical profile, so this key would + // otherwise appear empty in every index.json ever written. + let json = serde_json::to_value(Profile::exact()).unwrap(); + assert_eq!(json["name"], serde_json::json!(PROFILE_EXACT)); + assert!(json.get("selects").is_none(), "{json}"); +} + +#[test] +fn a_profile_without_selects_still_deserialises() { + let p: Profile = serde_json::from_str(r#"{"name":"exact"}"#).unwrap(); + assert_eq!(p, Profile::exact()); +} diff --git a/crates/larql-vindex/src/format/vindex3/read.rs b/crates/larql-vindex/src/format/vindex3/read.rs new file mode 100644 index 000000000..65b8706b1 --- /dev/null +++ b/crates/larql-vindex/src/format/vindex3/read.rs @@ -0,0 +1,385 @@ +//! Open a VINDEX3 container and reach its regions. +//! +//! The other half of the seam. [`write`](super::write) makes a directory that +//! declares itself VINDEX3; this makes one usable — which is what turns +//! "the runtime works on VINDEX2 operands" into "VINDEX3 works". +//! +//! Everything here is *resolution*, not execution: the container yields byte +//! slices and the shapes they were declared with, and the runtime binds them. +//! Keeping the two apart is what lets the same executor serve operands from +//! either generation without knowing which it got. +//! +//! `open` resolves generation, index, manifest and then **every declared +//! profile** (§9.1) before reading a segment byte. + +use std::collections::HashMap; +use std::path::{Path, PathBuf}; + +use super::index::Vindex3Index; +use super::write::segment_path; +use crate::format::filenames::INDEX_JSON; +use crate::format::generation::{detect_generation, ContainerGeneration}; +use crate::format::lyrw2::read::Lyrw2Reader; +use crate::format::moe_manifest::{MoeLayer, MoeManifest}; +use crate::VindexError; + +/// Wrap an IO failure with what was being read, keeping the original kind so a +/// caller can still branch on NotFound vs PermissionDenied. +fn contextual_io(what: &str, e: std::io::Error) -> VindexError { + VindexError::Io(std::io::Error::new(e.kind(), format!("{what}: {e}"))) +} + +/// An opened VINDEX3 container: root authority, programme manifest, and the +/// segment bytes held resident for the reader's lifetime. +pub struct Vindex3Container { + root: PathBuf, + index: Vindex3Index, + manifest: MoeManifest, + /// Segment key → raw LYRW v2 file bytes. + segments: HashMap>, +} + +impl Vindex3Container { + /// Open `root`, refusing anything that is not a VINDEX3 container. + /// + /// The generation is checked *first*, so a VINDEX2 directory is refused by + /// name rather than producing a confusing parse failure three fields into + /// a schema it was never written against (spec §12.1). + pub fn open(root: &Path) -> Result { + match detect_generation(root)? { + ContainerGeneration::V3 => {} + ContainerGeneration::V2 => { + return Err(VindexError::WrongContainerGeneration { + found: "VINDEX2", + required: "VINDEX3", + }) + } + } + + let raw = std::fs::read_to_string(root.join(INDEX_JSON))?; + let index: Vindex3Index = serde_json::from_str(&raw) + .map_err(|e| VindexError::Parse(format!("parse VINDEX3 index.json: {e}")))?; + + let manifest_raw = std::fs::read_to_string(root.join(&index.moe_manifest)) + .map_err(|e| contextual_io(&index.moe_manifest, e))?; + let manifest = MoeManifest::parse(&manifest_raw)?; + + super::profile::resolve_all(&index) + .map_err(|e| VindexError::Parse(format!("{INDEX_JSON}: {e}")))?; + + let mut segments = HashMap::new(); + for key in index.segments.keys() { + let path = segment_path(root, key); + let bytes = + std::fs::read(&path).map_err(|e| contextual_io(&format!("segment {key}"), e))?; + segments.insert(key.clone(), bytes); + } + + Ok(Self { + root: root.to_path_buf(), + index, + manifest, + segments, + }) + } + + pub fn root(&self) -> &Path { + &self.root + } + + pub fn index(&self) -> &Vindex3Index { + &self.index + } + + pub fn manifest(&self) -> &MoeManifest { + &self.manifest + } + + /// Resolve a profile to the variants it will execute (§9.1). + /// + /// `open` already proved every declared profile resolves, so the only + /// failure a caller can still see here is asking for a profile this + /// container does not declare. + pub fn select_profile( + &self, + name: &str, + ) -> Result, super::profile::ProfileSelectionError> { + self.index.select_profile(name) + } + + /// The manifest entry for `layer`, or `None` if that layer is dense. + pub fn layer(&self, layer: u32) -> Option<&MoeLayer> { + self.manifest.layer(layer) + } + + /// A parsed reader over the segment a bank names. + /// + /// Fails naming the key rather than the filename: the key is what the + /// index declares and what a diagnosis should send someone to look at. + pub fn segment(&self, key: &str) -> Result, VindexError> { + let bytes = self.segments.get(key).ok_or_else(|| { + VindexError::Parse(format!( + "segment '{key}' is not declared by index.json; declared: {:?}", + self.index.segments.keys().collect::>() + )) + })?; + Lyrw2Reader::parse(bytes).map_err(|e| VindexError::Parse(format!("segment '{key}': {e}"))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::format::vindex3::test_support::{ + fixture_a_spec, FIXTURE_A_EXPERTS, FIXTURE_A_LAYER, FIXTURE_A_SEGMENT_KEY, + }; + use crate::format::vindex3::write::write_container; + use tempfile::tempdir; + + fn opened() -> (tempfile::TempDir, Vindex3Container) { + let dir = tempdir().unwrap(); + write_container(dir.path(), &fixture_a_spec()).expect("write"); + let c = Vindex3Container::open(dir.path()).expect("open"); + (dir, c) + } + + #[test] + fn a_written_container_opens_and_carries_its_manifest() { + let (_dir, c) = opened(); + let layer = c.layer(FIXTURE_A_LAYER).expect("layer 0 is a MoE layer"); + assert_eq!(layer.routed_bank.experts, FIXTURE_A_EXPERTS); + assert!( + layer.routed_bank.resolve_programme().is_some(), + "the declared programme must resolve, or nothing can bind it" + ); + } + + #[test] + fn the_declared_segment_parses_as_lyrw_v2() { + // Round-trip through the real writer and the real reader — the two + // halves that had never met before this module existed. + let (_dir, c) = opened(); + let seg = c.segment(FIXTURE_A_SEGMENT_KEY).expect("segment parses"); + assert_eq!(seg.banks().len(), 1); + assert_eq!(seg.banks()[0].num_entries, FIXTURE_A_EXPERTS); + } + + #[test] + fn an_undeclared_segment_is_refused_naming_what_is_declared() { + let (_dir, c) = opened(); + let err = c.segment("routed/layer_999").unwrap_err(); + let msg = format!("{err}"); + assert!(msg.contains("routed/layer_999"), "{msg}"); + assert!( + msg.contains(FIXTURE_A_SEGMENT_KEY), + "must name what IS there: {msg}" + ); + } + + #[test] + fn an_opened_container_reports_where_it_came_from_and_what_it_declares() { + let (dir, c) = opened(); + assert_eq!(c.root(), dir.path()); + assert_eq!( + c.index().version, + crate::format::generation::V3_CURRENT_SCHEMA + ); + assert!(c + .index() + .declares_profile(crate::format::vindex3::PROFILE_EXACT)); + assert_eq!(c.manifest().layers.len(), 1); + } + + #[test] + fn a_dense_layer_has_no_manifest_entry() { + // Layers absent from the manifest are dense — that is how an arbitrary + // dense/MoE schedule is expressed, so asking for one must answer None + // rather than erroring. + let (_dir, c) = opened(); + assert!(c.layer(FIXTURE_A_LAYER + 99).is_none()); + } + + // ── §9.1 variant selection, enforced at open ──────────────────────── + + /// Rewrite the container's `index.json` through a mutation, so a test can + /// stage an index a writer would never produce — which is exactly the + /// class of container a load-time gate exists to refuse. + fn repack_index(root: &Path, mutate: &dyn Fn(&mut Vindex3Index)) { + let path = root.join(INDEX_JSON); + let mut index: Vindex3Index = + serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap(); + mutate(&mut index); + std::fs::write(&path, serde_json::to_string_pretty(&index).unwrap()).unwrap(); + } + + fn variants_with(baseline: &str, present: &str) -> crate::format::vindex3::RegionSetVariants { + use crate::format::capability::authority::Fidelity; + use crate::format::vindex3::{RegionSetVariants, StoredVariant}; + RegionSetVariants { + baseline: baseline.to_string(), + variants: std::collections::BTreeMap::from([( + present.to_string(), + StoredVariant::new("routed/layer_000", Fidelity::SourceEquivalent), + )]), + } + } + + #[test] + fn a_profile_selecting_an_absent_variant_is_refused_at_open() { + let dir = tempdir().unwrap(); + write_container(dir.path(), &fixture_a_spec()).expect("write"); + repack_index(dir.path(), &|index: &mut Vindex3Index| { + index.variants = crate::format::vindex3::VariantCatalogue::new().with_set( + "layer.0.routed.gate_up", + variants_with("exact-q6k", "exact-q6k"), + ); + index.profiles.push( + crate::format::vindex3::Profile::new("routed-mxfp4") + .selecting("layer.0.routed.gate_up", "native-mxfp4"), + ); + }); + + let msg = Vindex3Container::open(dir.path()) + .err() + .expect("a profile selecting a variant nobody extracted must not open") + .to_string(); + // The three things §9.1 requires a refusal to name. + assert!(msg.contains("layer.0.routed.gate_up"), "region set: {msg}"); + assert!(msg.contains("native-mxfp4"), "requested variant: {msg}"); + assert!(msg.contains("exact-q6k"), "variants present: {msg}"); + assert!(msg.contains("routed-mxfp4"), "the profile at fault: {msg}"); + } + + #[test] + fn the_refusal_happens_before_any_segment_byte_is_read() { + // The load-order claim, tested the only way it can be: delete every + // segment file. If the profile gate ran first the error names the + // variant; if the segment read ran first it names a missing file. + let dir = tempdir().unwrap(); + write_container(dir.path(), &fixture_a_spec()).expect("write"); + repack_index(dir.path(), &|index: &mut Vindex3Index| { + index.variants = crate::format::vindex3::VariantCatalogue::new().with_set( + "layer.0.routed.gate_up", + variants_with("exact-q6k", "exact-q6k"), + ); + index.profiles.push( + crate::format::vindex3::Profile::new("routed-mxfp4") + .selecting("layer.0.routed.gate_up", "native-mxfp4"), + ); + }); + std::fs::remove_file(segment_path(dir.path(), FIXTURE_A_SEGMENT_KEY)).unwrap(); + + let msg = Vindex3Container::open(dir.path()) + .err() + .expect("must not open") + .to_string(); + assert!( + msg.contains("native-mxfp4"), + "the variant gate must fire before the segment read: {msg}" + ); + } + + #[test] + fn a_baseline_that_was_never_extracted_is_refused_at_open() { + let dir = tempdir().unwrap(); + write_container(dir.path(), &fixture_a_spec()).expect("write"); + repack_index(dir.path(), &|index: &mut Vindex3Index| { + index.variants = crate::format::vindex3::VariantCatalogue::new().with_set( + "layer.0.routed.gate_up", + variants_with("native-mxfp4", "exact-q6k"), + ); + }); + let msg = Vindex3Container::open(dir.path()) + .err() + .expect("an unservable baseline must not open") + .to_string(); + assert!(msg.contains("native-mxfp4"), "{msg}"); + assert!(msg.contains("exact-q6k"), "{msg}"); + } + + #[test] + fn a_container_with_no_packs_opens_and_selects_its_only_profile() { + // The shape every container written so far has: no catalogue at all. + // An absent catalogue must mean "no packs", never "unchecked". + let (_dir, c) = opened(); + assert!(c.index().variants.is_empty()); + let resolved = c + .select_profile(crate::format::vindex3::PROFILE_EXACT) + .expect("the canonical profile always resolves"); + assert!(resolved.is_empty()); + } + + #[test] + fn asking_for_a_profile_the_container_does_not_declare_names_the_ones_it_does() { + let (_dir, c) = opened(); + let err = c.select_profile("browse").unwrap_err(); + let msg = err.to_string(); + assert!(msg.contains("browse"), "{msg}"); + assert!( + msg.contains(crate::format::vindex3::PROFILE_EXACT), + "must name what IS declared: {msg}" + ); + } + + #[test] + fn a_missing_manifest_file_is_refused_naming_it() { + let dir = tempdir().unwrap(); + write_container(dir.path(), &fixture_a_spec()).expect("write"); + std::fs::remove_file(dir.path().join(crate::format::vindex3::MOE_MANIFEST_JSON)).unwrap(); + let msg = match Vindex3Container::open(dir.path()) { + Ok(_) => panic!("a container without its manifest must not open"), + Err(e) => format!("{e}"), + }; + assert!(msg.contains("moe_manifest.json"), "{msg}"); + } + + #[test] + fn a_declared_segment_that_is_absent_is_refused_naming_the_key() { + let dir = tempdir().unwrap(); + write_container(dir.path(), &fixture_a_spec()).expect("write"); + std::fs::remove_file(crate::format::vindex3::segment_path( + dir.path(), + FIXTURE_A_SEGMENT_KEY, + )) + .unwrap(); + let msg = match Vindex3Container::open(dir.path()) { + Ok(_) => panic!("a container missing a declared segment must not open"), + Err(e) => format!("{e}"), + }; + assert!(msg.contains(FIXTURE_A_SEGMENT_KEY), "{msg}"); + } + + #[test] + fn a_malformed_index_is_refused_as_a_parse_failure() { + let dir = tempdir().unwrap(); + write_container(dir.path(), &fixture_a_spec()).expect("write"); + // Keep `version: 3` so detection still routes here, then break the rest. + std::fs::write(dir.path().join(INDEX_JSON), r#"{"version":3,"model":42}"#).unwrap(); + assert!(Vindex3Container::open(dir.path()).is_err()); + } + + #[test] + fn the_v3_loader_refuses_a_v2_container_by_name() { + // The generation boundary, from the loader's side. §12.1 requires a + // precise refusal, never a parse error and never a conversion. + let dir = tempdir().unwrap(); + std::fs::write( + dir.path().join(INDEX_JSON), + serde_json::json!({ + "version": 2, "model": "m", "family": "llama", "num_layers": 1, + "hidden_size": 8, "intermediate_size": 8, "vocab_size": 4, + "embed_scale": 1.0, "layers": [], "down_top_k": 0 + }) + .to_string(), + ) + .unwrap(); + // `Vindex3Container` holds segment bytes and deliberately has no + // Debug, so match rather than `unwrap_err`. + let msg = match Vindex3Container::open(dir.path()) { + Ok(_) => panic!("the VINDEX3 loader must refuse a VINDEX2 container"), + Err(e) => format!("{e}"), + }; + assert!(msg.contains("VINDEX2"), "{msg}"); + assert!(msg.contains("VINDEX3 loader"), "{msg}"); + } +} diff --git a/crates/larql-vindex/src/format/vindex3/test_support.rs b/crates/larql-vindex/src/format/vindex3/test_support.rs new file mode 100644 index 000000000..c375b943a --- /dev/null +++ b/crates/larql-vindex/src/format/vindex3/test_support.rs @@ -0,0 +1,373 @@ +//! Conformance fixture A as a real VINDEX3 container. +//! +//! ```text +//! hidden 256 · experts 8 · top_k 2 · shared 0 · programme gated-mlp-v1 +//! ``` +//! +//! Dimensions are the ones frozen for fixture A in the experiments programme +//! (§1.1 "A — Direct routed MoE (control)"). It is the control precisely +//! because nothing about it is interesting: residual-space experts, plain +//! softmax top-k, no shared bank, no transforms. Anything that fails here +//! fails for a container reason rather than an architecture one. +//! +//! Weights are a deterministic ramp, not random. The fixture's purpose is +//! byte-level parity between two source paths, so reproducibility matters and +//! realism does not. + +use std::collections::BTreeMap; +use std::sync::atomic::{AtomicU64, Ordering}; + +use super::write::{ContainerSpec, SegmentSource}; +use crate::format::lyrw2::bank::{BankDescriptor, BankKind}; +use crate::format::lyrw2::browse_mode::BrowseMode; +use crate::format::lyrw2::plan::Lyrw2Plan; +use crate::format::lyrw2::region_format::{Packing, RegionFormat}; +use crate::format::lyrw2::region_role::RegionRole; +use crate::format::lyrw2::region_schema::RegionSchema; +use crate::format::lyrw2::write::Lyrw2Writer; +use crate::format::moe_manifest::bank_ref::BankRef; +use crate::format::moe_manifest::{ + Combine, InputSpace, MoeLayer, MoeManifest, Reduction, Router, RouterPostProcessing, + ScoreActivation, Selection, MOE_MANIFEST_SCHEMA_VERSION, +}; + +pub const FIXTURE_A_HIDDEN: u32 = 256; +pub const FIXTURE_A_INTERMEDIATE: u32 = 256; +pub const FIXTURE_A_EXPERTS: u32 = 8; +pub const FIXTURE_A_TOP_K: usize = 2; +pub const FIXTURE_A_LAYER: u32 = 0; +pub const FIXTURE_A_SEGMENT_KEY: &str = "routed/layer_000"; +/// Segment key for the fused-FC1 rendering of the same weights. +pub const FIXTURE_A_FUSED_SEGMENT_KEY: &str = "routed/layer_000_fused"; +pub const FIXTURE_A_PROGRAMME: &str = "gated-mlp-v1"; +pub const FIXTURE_A_ROUTER_KEY: &str = "router.weight"; +pub const FIXTURE_A_BANK_ID: u16 = 0; + +/// Region schema ordinals within the bank, in write order. +/// +/// `gated-mlp-v1` declares **both** `gate+up+down` and `gate_up_fused+down` +/// legal (`Programme::role_alternatives`), so one manifest can describe either +/// physical layout. That is what makes the fused/decomposed parity arm +/// meaningful: the logical operator contract is fixed by the programme id, and +/// the storage shape is free underneath it. +const SCHEMA_GATE: u16 = 0; +const SCHEMA_UP: u16 = 1; +const SCHEMA_DOWN: u16 = 2; +const FIXTURE_A_SCHEMA_COUNT: u16 = 3; +const FIXTURE_A_FUSED_SCHEMA_COUNT: u16 = 2; +/// Gate and up halves inside a fused FC1 region. +const FUSED_PROJECTION_HALVES: u32 = 2; + +/// Ramp parameters. Small and irrational-ish so a transposition or an off-by- +/// one row shows up as a large difference rather than a plausible one. +const RAMP_SCALE: f32 = 0.013; +const RAMP_BIAS: f32 = -0.37; +const EXPERT_STRIDE: f32 = 0.101; + +fn ramp(expert: u32, index: usize, len: usize) -> f32 { + let t = index as f32 / len as f32; + (t * RAMP_SCALE + RAMP_BIAS + expert as f32 * EXPERT_STRIDE).sin() +} + +/// Deterministic `gate` weights for one expert, `[intermediate, hidden]`. +pub fn gate_f32(expert: u32) -> Vec { + let len = (FIXTURE_A_INTERMEDIATE * FIXTURE_A_HIDDEN) as usize; + (0..len).map(|i| ramp(expert, i, len)).collect() +} + +/// Deterministic `up` weights for one expert, `[intermediate, hidden]`. +pub fn up_f32(expert: u32) -> Vec { + let len = (FIXTURE_A_INTERMEDIATE * FIXTURE_A_HIDDEN) as usize; + (0..len).map(|i| ramp(expert + 2, i, len)).collect() +} + +/// Deterministic `down` bytes for one expert, f32 row-major. +pub fn down_f32(expert: u32) -> Vec { + let len = (FIXTURE_A_HIDDEN * FIXTURE_A_INTERMEDIATE) as usize; + (0..len).map(|i| ramp(expert + 1, i, len)).collect() +} + +/// Router weight `[experts, hidden]`, f32 row-major, for an arbitrary +/// population — the axis V2-1 asks be demonstrably not hard-coded. +pub fn router_f32_for(experts: u32) -> Vec { + let len = (experts * FIXTURE_A_HIDDEN) as usize; + (0..len).map(|i| ramp(0, i, len)).collect() +} + +/// Router weight for fixture A's frozen population. +pub fn router_f32() -> Vec { + router_f32_for(FIXTURE_A_EXPERTS) +} + +/// Gate rows followed by up rows — the layout `BoundProjection::Fused` +/// contracts for, and the only ordering `gate_up_fused` may mean. +pub fn gate_up_fused_f32(expert: u32) -> Vec { + let mut v = gate_f32(expert); + v.extend(up_f32(expert)); + v +} + +fn as_bytes(values: &[f32]) -> Vec { + values.iter().flat_map(|v| v.to_le_bytes()).collect() +} + +/// A uniquely-named scratch path. Callers run in parallel and this module +/// compiles into the library, so it cannot reach for `tempfile`. +fn staging_path() -> std::path::PathBuf { + static NEXT: AtomicU64 = AtomicU64::new(0); + let dir = std::env::temp_dir().join("vindex3-fixture-a"); + std::fs::create_dir_all(&dir).expect("staging dir"); + dir.join(format!( + "seg-{}-{}.lyrw", + std::process::id(), + NEXT.fetch_add(1, Ordering::Relaxed) + )) +} + +fn bank_of(experts: u32) -> BankDescriptor { + BankDescriptor { + bank_id: FIXTURE_A_BANK_ID, + kind: BankKind::Routed, + num_entries: experts, + input_dim: FIXTURE_A_HIDDEN, + intermediate_dim: FIXTURE_A_INTERMEDIATE, + output_dim: FIXTURE_A_HIDDEN, + region_schema_count: FIXTURE_A_SCHEMA_COUNT, + browse: BrowseMode::Direct, + } +} + +fn schemas() -> Vec { + vec![ + RegionSchema::unpaired( + SCHEMA_GATE, + RegionRole::Gate, + RegionFormat::F32, + Packing::RowMajor, + FIXTURE_A_INTERMEDIATE, + FIXTURE_A_HIDDEN, + ), + RegionSchema::unpaired( + SCHEMA_UP, + RegionRole::Up, + RegionFormat::F32, + Packing::RowMajor, + FIXTURE_A_INTERMEDIATE, + FIXTURE_A_HIDDEN, + ), + RegionSchema::unpaired( + SCHEMA_DOWN, + RegionRole::Down, + RegionFormat::F32, + Packing::RowMajor, + FIXTURE_A_HIDDEN, + FIXTURE_A_INTERMEDIATE, + ), + ] +} + +/// Write fixture A's routed bank as one LYRW v2 segment, returning its bytes. +/// +/// Staged through a uniquely-named file rather than a shared path: callers run +/// in parallel, and a fixed filename had them deleting each other's segment +/// mid-read. Uniqueness comes from pid plus a process-local counter — this +/// module compiles into the library, so it cannot reach for `tempfile`. +pub fn segment_bytes_for(experts: u32) -> Vec { + let path = staging_path(); + + let plan = Lyrw2Plan::single_segment(FIXTURE_A_LAYER, bank_of(experts), schemas()); + let mut writer = Lyrw2Writer::create(&path, plan).expect("create lyrw2 writer"); + for expert in 0..experts { + writer + .write_region(&as_bytes(&gate_f32(expert))) + .expect("write gate"); + writer + .write_region(&as_bytes(&up_f32(expert))) + .expect("write up"); + writer + .write_region(&as_bytes(&down_f32(expert))) + .expect("write down"); + } + writer.finish().expect("finish lyrw2 segment"); + + let bytes = std::fs::read(&path).expect("read back segment"); + let _ = std::fs::remove_file(&path); + bytes +} + +fn fused_bank() -> BankDescriptor { + BankDescriptor { + region_schema_count: FIXTURE_A_FUSED_SCHEMA_COUNT, + ..bank_of(FIXTURE_A_EXPERTS) + } +} + +fn fused_schemas() -> Vec { + vec![ + RegionSchema::unpaired( + SCHEMA_GATE, + RegionRole::GateUpFused, + RegionFormat::F32, + Packing::RowMajor, + FUSED_PROJECTION_HALVES * FIXTURE_A_INTERMEDIATE, + FIXTURE_A_HIDDEN, + ), + RegionSchema::unpaired( + SCHEMA_UP, + RegionRole::Down, + RegionFormat::F32, + Packing::RowMajor, + FIXTURE_A_HIDDEN, + FIXTURE_A_INTERMEDIATE, + ), + ] +} + +/// The same weights as [`fixture_a_segment_bytes`], stored fused. +pub fn fixture_a_fused_segment_bytes() -> Vec { + let path = staging_path(); + let plan = Lyrw2Plan::single_segment(FIXTURE_A_LAYER, fused_bank(), fused_schemas()); + let mut writer = Lyrw2Writer::create(&path, plan).expect("create lyrw2 writer"); + for expert in 0..FIXTURE_A_EXPERTS { + writer + .write_region(&as_bytes(&gate_up_fused_f32(expert))) + .expect("write gate_up_fused"); + writer + .write_region(&as_bytes(&down_f32(expert))) + .expect("write down"); + } + writer.finish().expect("finish fused segment"); + let bytes = std::fs::read(&path).expect("read back fused segment"); + let _ = std::fs::remove_file(&path); + bytes +} + +/// Fixture A rendered with fused FC1 storage, under the *same* programme id. +pub fn fixture_a_fused_spec() -> ContainerSpec { + let mut spec = fixture_a_spec(); + spec.manifest.layers[0].routed_bank.storage = FIXTURE_A_FUSED_SEGMENT_KEY.to_string(); + spec.segments = vec![SegmentSource { + key: FIXTURE_A_FUSED_SEGMENT_KEY.to_string(), + bytes: fixture_a_fused_segment_bytes(), + }]; + spec +} + +fn manifest_for(experts: u32, top_k: usize, storage: &str) -> MoeManifest { + MoeManifest { + schema_version: MOE_MANIFEST_SCHEMA_VERSION, + layers: vec![MoeLayer { + layer: FIXTURE_A_LAYER, + input_space: InputSpace::Residual, + router: Router { + scores: FIXTURE_A_ROUTER_KEY.to_string(), + activation: ScoreActivation::Softmax, + selection: Selection::TopK { k: top_k }, + post: RouterPostProcessing::default(), + shared_expert_sink: false, + }, + transforms: None, + routed_bank: BankRef { + experts, + programme: FIXTURE_A_PROGRAMME.to_string(), + storage: storage.to_string(), + expert_dims: None, + }, + shared_bank: None, + reduction: Reduction::GateWeightedSum, + routed_output_norm: None, + combine: Combine::ResidualAdd, + }], + } +} + +/// A routed-MoE container with an arbitrary population and top-k. +/// +/// Exists to discharge V2-1's "expert counts and top-K are demonstrably not +/// hard-coded" row: the same code path builds an 8-expert top-2 container and +/// a 32-expert top-4 one, and the executor reads both out of their own +/// declarations rather than a constant. +pub fn routed_spec(experts: u32, top_k: usize, storage: &str) -> ContainerSpec { + ContainerSpec { + model: format!("routed-{experts}x{top_k}"), + family: "direct-moe".to_string(), + hidden_size: FIXTURE_A_HIDDEN as usize, + num_layers: 1, + manifest: manifest_for(experts, top_k, storage), + segments: vec![SegmentSource { + key: storage.to_string(), + bytes: segment_bytes_for(experts), + }], + } +} + +/// The complete fixture-A container specification (frozen dimensions). +pub fn fixture_a_spec() -> ContainerSpec { + let mut spec = routed_spec(FIXTURE_A_EXPERTS, FIXTURE_A_TOP_K, FIXTURE_A_SEGMENT_KEY); + spec.model = "fixture-a".to_string(); + spec +} + +/// A container whose declared segment holds `bytes` verbatim. +/// +/// Lets a test put a deliberately unusable segment behind a well-formed index +/// and manifest — the shape `verify` exists to catch, and one `open` cannot, +/// since it reads segment files without parsing them. +pub fn spec_with_segment_bytes(bytes: Vec) -> ContainerSpec { + let mut spec = fixture_a_spec(); + spec.segments = vec![SegmentSource { + key: FIXTURE_A_SEGMENT_KEY.to_string(), + bytes, + }]; + spec +} + +/// A routed bank carrying gate and up but **no down** — well-formed LYRW v2 +/// that no gated-MLP programme can be satisfied by. +pub fn gate_only_segment_bytes() -> Vec { + let path = staging_path(); + let bank = BankDescriptor { + region_schema_count: 2, + ..bank_of(FIXTURE_A_EXPERTS) + }; + let schemas = vec![ + RegionSchema::unpaired( + SCHEMA_GATE, + RegionRole::Gate, + RegionFormat::F32, + Packing::RowMajor, + FIXTURE_A_INTERMEDIATE, + FIXTURE_A_HIDDEN, + ), + RegionSchema::unpaired( + SCHEMA_UP, + RegionRole::Up, + RegionFormat::F32, + Packing::RowMajor, + FIXTURE_A_INTERMEDIATE, + FIXTURE_A_HIDDEN, + ), + ]; + let plan = Lyrw2Plan::single_segment(FIXTURE_A_LAYER, bank, schemas); + let mut writer = Lyrw2Writer::create(&path, plan).expect("create writer"); + for expert in 0..FIXTURE_A_EXPERTS { + writer + .write_region(&as_bytes(&gate_f32(expert))) + .expect("gate"); + writer.write_region(&as_bytes(&up_f32(expert))).expect("up"); + } + writer.finish().expect("finish"); + let bytes = std::fs::read(&path).expect("read back"); + let _ = std::fs::remove_file(&path); + bytes +} + +/// Fixture-A weights rendered as the frozen fixture-A segment. +pub fn fixture_a_segment_bytes() -> Vec { + segment_bytes_for(FIXTURE_A_EXPERTS) +} + +/// Segment keys the spec declares — handy for asserting the index covers them. +pub fn fixture_a_segment_keys() -> BTreeMap { + BTreeMap::from([(FIXTURE_A_SEGMENT_KEY.to_string(), 1)]) +} diff --git a/crates/larql-vindex/src/format/vindex3/variants.rs b/crates/larql-vindex/src/format/vindex3/variants.rs new file mode 100644 index 000000000..3c7f438e8 --- /dev/null +++ b/crates/larql-vindex/src/format/vindex3/variants.rs @@ -0,0 +1,235 @@ +//! Representation variants — a region set carries bytes, a profile picks which +//! (spec §9.1). +//! +//! # The one legal representation model +//! +//! A profile saying `"format": "mxfp4"` cannot turn Q6_K bytes into MXFP4 bytes +//! by declaration. Exactly one model is legal: **a region set may carry several +//! physically present variants, and a profile selects a present one.** There is +//! no runtime conversion anywhere in this module, and that is the point — "the +//! bytes executed are the bytes stored" (§10) holds by construction rather than +//! by discipline. +//! +//! # Why the refusal is the interesting part +//! +//! Selecting a variant that was never extracted is the failure this type +//! exists to make loud. It has to fail **closed**, and it has to fail naming +//! three things: the region set, the variant that was asked for, and the +//! variants actually present. A bare "variant not found" sends someone to +//! re-extract a terabyte before discovering that the pack they wanted is +//! spelled `native-mxfp4` and they asked for `mxfp4`. +//! +//! It also has to fail **before any byte is read**. A container whose profile +//! selects an absent variant is unservable, and finding that out during a bind +//! — after the segment files are resident — turns a naming error into a +//! partially-initialised load. + +use std::collections::BTreeMap; + +use serde::{Deserialize, Serialize}; + +use crate::format::capability::authority::Fidelity; + +/// One physically present encoding of a region set. +/// +/// `storage` is a segment key relative to the container root, the same +/// vocabulary `index.segments` uses — a key the loader composes a filename +/// from, never a filename it globbed for (§12.1). +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct StoredVariant { + pub storage: String, + pub fidelity: Fidelity, +} + +impl StoredVariant { + pub fn new(storage: impl Into, fidelity: Fidelity) -> Self { + Self { + storage: storage.into(), + fidelity, + } + } +} + +/// The variants one region set physically carries, and which of them is +/// canonical. +/// +/// `baseline` is the authority: additional variants are opt-in, per-component +/// and individually removable, so a pack can be deleted without the region set +/// becoming unservable. That only holds if the baseline is itself present, +/// which [`Self::validate`] is what checks. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct RegionSetVariants { + pub baseline: String, + pub variants: BTreeMap, +} + +impl RegionSetVariants { + /// A region set carrying exactly one variant, which is therefore also the + /// baseline — the shape every container written before packs existed has. + pub fn single(name: impl Into, stored: StoredVariant) -> Self { + let name = name.into(); + Self { + baseline: name.clone(), + variants: BTreeMap::from([(name, stored)]), + } + } + + /// Add a variant beside the baseline, which is never rewritten (§9.1). + pub fn with_variant(mut self, name: impl Into, stored: StoredVariant) -> Self { + self.variants.insert(name.into(), stored); + self + } + + /// Present variant names, sorted. + /// + /// Sorted because this list goes straight into a refusal message, and a + /// diagnostic whose contents reorder between runs cannot be diffed. + pub fn present(&self) -> Vec { + self.variants.keys().cloned().collect() + } + + /// Resolve `name` against what is physically here. + pub fn select(&self, region_set: &str, name: &str) -> Result<&StoredVariant, VariantDefect> { + self.variants + .get(name) + .ok_or_else(|| VariantDefect::AbsentVariant { + region_set: region_set.to_string(), + requested: name.to_string(), + present: self.present(), + }) + } + + /// Resolve the baseline — what a profile that names no variant gets. + pub fn select_baseline(&self, region_set: &str) -> Result<&StoredVariant, VariantDefect> { + self.select(region_set, &self.baseline) + } + + /// Structural checks that do not depend on any profile. + fn validate(&self, region_set: &str, out: &mut Vec) { + if self.variants.is_empty() { + out.push(VariantDefect::NoVariants { + region_set: region_set.to_string(), + }); + return; + } + if !self.variants.contains_key(&self.baseline) { + out.push(VariantDefect::BaselineAbsent { + region_set: region_set.to_string(), + baseline: self.baseline.clone(), + present: self.present(), + }); + } + } +} + +/// Every region set a container catalogues, and the variants each carries. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(transparent)] +pub struct VariantCatalogue { + sets: BTreeMap, +} + +impl VariantCatalogue { + pub fn new() -> Self { + Self::default() + } + + pub fn with_set(mut self, region_set: impl Into, variants: RegionSetVariants) -> Self { + self.sets.insert(region_set.into(), variants); + self + } + + pub fn is_empty(&self) -> bool { + self.sets.is_empty() + } + + pub fn len(&self) -> usize { + self.sets.len() + } + + /// Catalogued region-set names, sorted — the other half of a refusal. + pub fn region_sets(&self) -> Vec { + self.sets.keys().cloned().collect() + } + + pub fn get(&self, region_set: &str) -> Option<&RegionSetVariants> { + self.sets.get(region_set) + } + + /// Resolve one `region_set::variant` pair. + pub fn select(&self, region_set: &str, variant: &str) -> Result<&StoredVariant, VariantDefect> { + self.set(region_set)?.select(region_set, variant) + } + + /// Resolve a region set's baseline. + pub fn select_baseline(&self, region_set: &str) -> Result<&StoredVariant, VariantDefect> { + self.set(region_set)?.select_baseline(region_set) + } + + fn set(&self, region_set: &str) -> Result<&RegionSetVariants, VariantDefect> { + self.sets + .get(region_set) + .ok_or_else(|| VariantDefect::UnknownRegionSet { + region_set: region_set.to_string(), + catalogued: self.region_sets(), + }) + } + + /// Every structural defect, rather than the first. + /// + /// All of them, because a container with two broken region sets should + /// take one repair pass and not two: fixing what the first error names + /// only to be shown the second is how a load-time check earns a reputation + /// for wasting time. + pub fn defects(&self) -> Vec { + let mut out = Vec::new(); + for (region_set, variants) in &self.sets { + variants.validate(region_set, &mut out); + } + out + } +} + +/// Why a variant selection cannot be honoured. +/// +/// Every variant names what was asked for *and* what is there, because the +/// remedy differs and only the pair distinguishes them: a typo is fixed in the +/// profile, a genuinely missing pack is fixed by extracting it. +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub enum VariantDefect { + #[error( + "region set '{region_set}' has no variant '{requested}'; present: [{}]", + present.join(", ") + )] + AbsentVariant { + region_set: String, + requested: String, + present: Vec, + }, + + #[error( + "unknown region set '{region_set}'; catalogued: [{}]", + catalogued.join(", ") + )] + UnknownRegionSet { + region_set: String, + catalogued: Vec, + }, + + #[error( + "region set '{region_set}' names baseline '{baseline}', which is not present; present: [{}]", + present.join(", ") + )] + BaselineAbsent { + region_set: String, + baseline: String, + present: Vec, + }, + + #[error("region set '{region_set}' catalogues no variants at all")] + NoVariants { region_set: String }, +} + +#[cfg(test)] +#[path = "variants_tests.rs"] +mod tests; diff --git a/crates/larql-vindex/src/format/vindex3/variants_tests.rs b/crates/larql-vindex/src/format/vindex3/variants_tests.rs new file mode 100644 index 000000000..cf48a7f96 --- /dev/null +++ b/crates/larql-vindex/src/format/vindex3/variants_tests.rs @@ -0,0 +1,244 @@ +//! Tests for the §9.1 variant catalogue. +//! +//! The theme throughout: a refusal is only useful if it names the remedy, so +//! most of these assert on the *contents* of the error rather than that one +//! occurred. + +use super::*; + +const Q6K: &str = "exact-q6k"; +const MXFP4: &str = "native-mxfp4"; +const GATE_UP: &str = "layer.12.routed.gate_up"; + +fn q6k() -> StoredVariant { + StoredVariant::new("routed/layer_012.q6k", Fidelity::SourceEquivalent) +} + +fn mxfp4() -> StoredVariant { + StoredVariant::new("routed/layer_012.mxfp4", Fidelity::SourceExact) +} + +/// The spec's own worked example (§9.1). +fn both() -> RegionSetVariants { + RegionSetVariants::single(Q6K, q6k()).with_variant(MXFP4, mxfp4()) +} + +fn catalogue() -> VariantCatalogue { + VariantCatalogue::new().with_set(GATE_UP, both()) +} + +// ── Selecting what is there ────────────────────────────────────────────── + +#[test] +fn a_present_variant_resolves_to_its_own_bytes() { + let c = catalogue(); + assert_eq!(c.select(GATE_UP, MXFP4).unwrap(), &mxfp4()); + assert_eq!(c.select(GATE_UP, Q6K).unwrap(), &q6k()); +} + +#[test] +fn the_baseline_is_what_a_profile_naming_nothing_gets() { + assert_eq!(catalogue().select_baseline(GATE_UP).unwrap(), &q6k()); +} + +#[test] +fn adding_a_variant_leaves_the_baseline_alone() { + // §9.1's incremental-pack rule: the multi-terabyte baseline is never + // rewritten, so adding a pack must not move which variant is canonical. + let before = RegionSetVariants::single(Q6K, q6k()); + let after = before.clone().with_variant(MXFP4, mxfp4()); + assert_eq!(after.baseline, before.baseline); + assert_eq!(after.variants.len(), 2); +} + +#[test] +fn a_variant_is_removable_without_touching_the_baseline() { + // The other half of "individually removable": deleting a pack leaves a + // servable region set behind. + let mut set = both(); + set.variants.remove(MXFP4); + assert!(set.select_baseline(GATE_UP).is_ok()); + assert_eq!(set.present(), vec![Q6K.to_string()]); +} + +// ── Selecting what is not there — the point of the module ─────────────── + +#[test] +fn an_absent_variant_names_the_set_the_request_and_what_is_present() { + let err = catalogue().select(GATE_UP, "mxfp4").unwrap_err(); + match &err { + VariantDefect::AbsentVariant { + region_set, + requested, + present, + } => { + assert_eq!(region_set, GATE_UP); + assert_eq!(requested, "mxfp4"); + assert_eq!(present, &vec![Q6K.to_string(), MXFP4.to_string()]); + } + other => panic!("wrong defect: {other:?}"), + } + // The near-miss this is really for: `mxfp4` vs `native-mxfp4`. The message + // has to carry the spelling that would have worked, or the reader + // re-extracts instead of re-typing. + let msg = err.to_string(); + assert!(msg.contains(GATE_UP), "{msg}"); + assert!(msg.contains("mxfp4"), "{msg}"); + assert!(msg.contains(MXFP4), "{msg}"); +} + +#[test] +fn an_unknown_region_set_names_what_is_catalogued() { + let err = catalogue() + .select("layer.99.routed.down", MXFP4) + .unwrap_err(); + match &err { + VariantDefect::UnknownRegionSet { + region_set, + catalogued, + } => { + assert_eq!(region_set, "layer.99.routed.down"); + assert_eq!(catalogued, &vec![GATE_UP.to_string()]); + } + other => panic!("wrong defect: {other:?}"), + } + assert!(err.to_string().contains(GATE_UP)); +} + +#[test] +fn an_absent_variant_is_distinguished_from_an_absent_region_set() { + // Same remedy family, different remedy: one is a profile edit, the other + // means the container never catalogued the component at all. Collapsing + // them into one error is what makes a load failure unactionable. + let c = catalogue(); + assert!(matches!( + c.select(GATE_UP, "nope"), + Err(VariantDefect::AbsentVariant { .. }) + )); + assert!(matches!( + c.select("nope", Q6K), + Err(VariantDefect::UnknownRegionSet { .. }) + )); +} + +#[test] +fn the_present_list_is_sorted_so_a_diagnostic_is_stable() { + let set = RegionSetVariants::single("zzz", q6k()) + .with_variant("aaa", mxfp4()) + .with_variant("mmm", q6k()); + assert_eq!(set.present(), vec!["aaa", "mmm", "zzz"]); +} + +// ── Structural defects, found without any profile ─────────────────────── + +#[test] +fn a_clean_catalogue_reports_no_defects() { + assert!(catalogue().defects().is_empty()); +} + +#[test] +fn a_baseline_that_is_not_present_is_a_defect() { + // The failure that makes a region set unservable while every individual + // pack looks fine — nothing else in the file disagrees with itself. + let set = RegionSetVariants { + baseline: "gone".into(), + variants: BTreeMap::from([(Q6K.to_string(), q6k())]), + }; + let defects = VariantCatalogue::new().with_set(GATE_UP, set).defects(); + match defects.as_slice() { + [VariantDefect::BaselineAbsent { + region_set, + baseline, + present, + }] => { + assert_eq!(region_set, GATE_UP); + assert_eq!(baseline, "gone"); + assert_eq!(present, &vec![Q6K.to_string()]); + } + other => panic!("wrong defects: {other:?}"), + } +} + +#[test] +fn a_region_set_with_no_variants_is_a_defect() { + let set = RegionSetVariants { + baseline: Q6K.into(), + variants: BTreeMap::new(), + }; + let defects = VariantCatalogue::new().with_set(GATE_UP, set).defects(); + assert_eq!( + defects, + vec![VariantDefect::NoVariants { + region_set: GATE_UP.to_string() + }] + ); +} + +#[test] +fn an_empty_region_set_reports_only_that_not_a_missing_baseline_too() { + // One cause, one defect. Reporting both would send someone to fix the + // baseline name when the region set has nothing to name. + let set = RegionSetVariants { + baseline: "anything".into(), + variants: BTreeMap::new(), + }; + assert_eq!( + VariantCatalogue::new() + .with_set(GATE_UP, set) + .defects() + .len(), + 1 + ); +} + +#[test] +fn every_broken_region_set_is_reported_in_one_pass() { + // Two repairs should take one load, not two. + let broken = |name: &str| RegionSetVariants { + baseline: name.to_string(), + variants: BTreeMap::from([(Q6K.to_string(), q6k())]), + }; + let c = VariantCatalogue::new() + .with_set("a", broken("missing-a")) + .with_set("b", broken("missing-b")); + assert_eq!(c.defects().len(), 2); +} + +// ── Wire form ──────────────────────────────────────────────────────────── + +#[test] +fn the_catalogue_round_trips_through_json() { + let before = catalogue(); + let json = serde_json::to_string(&before).unwrap(); + let after: VariantCatalogue = serde_json::from_str(&json).unwrap(); + assert_eq!(before, after); +} + +#[test] +fn the_wire_form_is_the_shape_the_spec_prints() { + // §9.1 shows `variants` as an object keyed by variant name, each carrying + // `storage` and `fidelity`, beside a `baseline`. A reader written against + // the spec has to find those keys. + let json = serde_json::to_value(catalogue()).unwrap(); + let set = &json[GATE_UP]; + assert_eq!(set["baseline"], serde_json::json!(Q6K)); + assert_eq!( + set["variants"][MXFP4]["storage"], + serde_json::json!("routed/layer_012.mxfp4") + ); + assert_eq!( + set["variants"][MXFP4]["fidelity"], + serde_json::json!("source-exact"), + "fidelity must serialise in the spec's kebab-case spelling" + ); +} + +#[test] +fn an_empty_catalogue_is_the_default_so_a_pack_less_container_still_loads() { + let c = VariantCatalogue::default(); + assert!(c.is_empty()); + assert_eq!(c.len(), 0); + assert!(c.defects().is_empty()); + assert!(c.region_sets().is_empty()); + assert!(c.get(GATE_UP).is_none()); +} diff --git a/crates/larql-vindex/src/format/vindex3/verify.rs b/crates/larql-vindex/src/format/vindex3/verify.rs new file mode 100644 index 000000000..6098878ac --- /dev/null +++ b/crates/larql-vindex/src/format/vindex3/verify.rs @@ -0,0 +1,349 @@ +//! What it means to verify a VINDEX3 container. +//! +//! For the shipped generation, `verify` is a checksum sweep: the files are +//! opaque blobs and "intact" means "the bytes are the bytes". A VINDEX3 +//! container declares its own structure, so verification can ask a much +//! stronger question — *is this container bindable?* — without executing +//! anything: +//! +//! ```text +//! index parses ← Vindex3Container::open +//! manifest parses and validates ← Vindex3Container::open +//! storage keys resolve to files ← Vindex3Container::open +//! each segment parses as LYRW v2 ← here +//! declared roles satisfy the programme ← here +//! every entry's regions are in bounds ← here +//! ``` +//! +//! The last three are the ones that turn "the files exist" into "this will +//! bind". A container that passes cannot fail at bind time for a structural +//! reason, which is the whole point of declaring structure up front. +//! +//! Execution parity is deliberately **not** here. It needs weights, an input +//! and a kernel, and folding it in would make routine verification cost a +//! forward pass. It belongs in a deeper mode. + +use super::read::Vindex3Container; +use crate::format::lyrw2::region_role::RegionRole; + +/// A structural defect that would stop this container binding. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ContainerDefect { + /// A layer's bank names storage the index does not declare, or that does + /// not parse as LYRW v2. + Storage { + layer: u32, + key: String, + why: String, + }, + /// The manifest names a programme this binary does not implement. + UnknownProgramme { layer: u32, programme: String }, + /// The bank's regions cannot satisfy any of the programme's acceptable + /// operand sets. + UnsatisfiedProgramme { + layer: u32, + programme: String, + missing: Vec, + }, + /// A declared region does not resolve for some entry. + /// + /// Carries the full coordinate V2-0 asks for — `{layer, bank, role, + /// segment}` — plus the entry. Layer alone cannot locate a region in a + /// container where one layer may span several banks and each bank several + /// segments, and a diagnosis that cannot be navigated to is a diagnosis + /// someone has to reproduce before they can act on it. + MissingRegion { + layer: u32, + bank: u16, + entry: u32, + role: RegionRole, + segment: String, + }, +} + +impl std::fmt::Display for ContainerDefect { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Storage { layer, key, why } => { + write!(f, "layer {layer}: storage '{key}' unusable: {why}") + } + Self::UnknownProgramme { layer, programme } => write!( + f, + "layer {layer}: programme '{programme}' is not implemented by this binary" + ), + Self::UnsatisfiedProgramme { + layer, + programme, + missing, + } => write!( + f, + "layer {layer}: programme '{programme}' cannot be satisfied; missing {missing:?}" + ), + Self::MissingRegion { + layer, + bank, + entry, + role, + segment, + } => write!( + f, + "layer {layer} bank {bank} segment '{segment}': entry {entry} has no {role:?} region" + ), + } + } +} + +impl Vindex3Container { + /// Every structural defect, in layer order. Empty means bindable. + /// + /// Returns all of them rather than the first: a container missing three + /// roles should say so once, not across three re-runs. + pub fn verify(&self) -> Vec { + let mut defects = Vec::new(); + for layer in &self.manifest().layers { + let bank_ref = &layer.routed_bank; + let Some(programme) = bank_ref.resolve_programme() else { + defects.push(ContainerDefect::UnknownProgramme { + layer: layer.layer, + programme: bank_ref.programme.clone(), + }); + continue; + }; + let segment = match self.segment(&bank_ref.storage) { + Ok(s) => s, + Err(e) => { + defects.push(ContainerDefect::Storage { + layer: layer.layer, + key: bank_ref.storage.clone(), + why: e.to_string(), + }); + continue; + } + }; + + let Some(bank) = segment.banks().first().copied() else { + defects.push(ContainerDefect::Storage { + layer: layer.layer, + key: bank_ref.storage.clone(), + why: "segment declares no bank".into(), + }); + continue; + }; + let schemas = segment.schemas_for(bank.bank_id).unwrap_or(&[]); + let present: Vec = schemas.iter().map(|s| s.role).collect(); + + if !programme.is_satisfied_by(&present) { + defects.push(ContainerDefect::UnsatisfiedProgramme { + layer: layer.layer, + programme: bank_ref.programme.clone(), + missing: programme.missing_roles(&present), + }); + continue; + } + + // Bounds are checked by resolving every declared region for every + // entry — the reader refuses an out-of-range offset, so a + // successful resolve is the bounds check. + for entry in 0..bank.num_entries { + for role in present.iter().copied() { + let resolved = segment + .region_bytes(bank.bank_id, entry, role) + .ok() + .flatten(); + if resolved.is_none() { + defects.push(ContainerDefect::MissingRegion { + layer: layer.layer, + bank: bank.bank_id, + entry, + role, + segment: bank_ref.storage.clone(), + }); + } + } + } + } + defects + } + + /// Whether the container is structurally bindable. + pub fn is_bindable(&self) -> bool { + self.verify().is_empty() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::format::vindex3::test_support::{fixture_a_fused_spec, fixture_a_spec}; + use crate::format::vindex3::write::write_container; + + fn verified(spec: &crate::format::vindex3::ContainerSpec) -> Vec { + let dir = std::env::temp_dir().join(format!( + "vindex3-verify-{}-{:p}", + std::process::id(), + spec as *const _ + )); + let _ = std::fs::remove_dir_all(&dir); + write_container(&dir, spec).expect("write"); + let defects = Vindex3Container::open(&dir).expect("open").verify(); + let _ = std::fs::remove_dir_all(&dir); + defects + } + + #[test] + fn a_well_formed_container_reports_no_defects() { + assert_eq!(verified(&fixture_a_spec()), Vec::new()); + } + + #[test] + fn the_fused_rendering_is_equally_bindable() { + // Both operand sets satisfy `gated-mlp-v1`, so verification must + // accept either — otherwise the programme's alternatives are a + // fiction the verifier does not honour. + assert_eq!(verified(&fixture_a_fused_spec()), Vec::new()); + } + + #[test] + fn an_unimplemented_programme_is_named_rather_than_guessed() { + let mut spec = fixture_a_spec(); + spec.manifest.layers[0].routed_bank.programme = "no-such-programme-v9".into(); + // `MoeManifest::parse` refuses this at open time, which is the + // stronger guarantee: it never reaches verification. + let dir = std::env::temp_dir().join(format!("vindex3-badprog-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&dir); + write_container(&dir, &spec).expect("write"); + let opened = Vindex3Container::open(&dir); + let _ = std::fs::remove_dir_all(&dir); + match opened { + Ok(c) => { + // If open tolerates it, verify must not. + let defects = c.verify(); + assert!( + defects + .iter() + .any(|d| matches!(d, ContainerDefect::UnknownProgramme { .. })), + "an unknown programme must be reported, got {defects:?}" + ); + } + Err(e) => assert!( + format!("{e}").contains("no-such-programme-v9"), + "the refusal must name the programme: {e}" + ), + } + } + + /// Write a container from `spec` into a private dir and verify it. + fn verify_spec( + tag: &str, + spec: &crate::format::vindex3::ContainerSpec, + ) -> Vec { + let dir = std::env::temp_dir().join(format!("v3-verify-{tag}-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&dir); + write_container(&dir, spec).expect("write"); + let defects = Vindex3Container::open(&dir).expect("open").verify(); + let _ = std::fs::remove_dir_all(&dir); + defects + } + + #[test] + fn a_segment_that_does_not_parse_is_a_storage_defect() { + // `open` reads segment files without parsing them, so an unusable + // segment survives loading and must be caught here — otherwise it + // would first surface at bind time as an unexplained failure. + use crate::format::vindex3::test_support::spec_with_segment_bytes; + let defects = verify_spec("garbage", &spec_with_segment_bytes(vec![0xAB; 512])); + assert!( + defects + .iter() + .any(|d| matches!(d, ContainerDefect::Storage { .. })), + "expected a Storage defect, got {defects:?}" + ); + let msg = defects[0].to_string(); + assert!(msg.contains("routed/layer_000"), "must name the key: {msg}"); + } + + #[test] + fn a_bank_missing_a_required_role_names_what_is_missing() { + // Gate + up but no down: well-formed LYRW v2 that no gated-MLP + // alternative can be satisfied by. The defect must say `Down` rather + // than "unsatisfied", or the reader has to diff the alternatives by + // hand to learn what to add. + use crate::format::vindex3::test_support::{ + gate_only_segment_bytes, spec_with_segment_bytes, + }; + let defects = verify_spec( + "noduown", + &spec_with_segment_bytes(gate_only_segment_bytes()), + ); + let missing = defects + .iter() + .find_map(|d| match d { + ContainerDefect::UnsatisfiedProgramme { missing, .. } => Some(missing.clone()), + _ => None, + }) + .unwrap_or_else(|| panic!("expected UnsatisfiedProgramme, got {defects:?}")); + assert!( + missing.contains(&RegionRole::Down), + "the missing role must be named: {missing:?}" + ); + } + + #[test] + fn is_bindable_agrees_with_the_defect_list() { + use crate::format::vindex3::test_support::spec_with_segment_bytes; + let dir = std::env::temp_dir().join(format!("v3-bindable-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&dir); + write_container(&dir, &fixture_a_spec()).expect("write good"); + assert!(Vindex3Container::open(&dir).expect("open").is_bindable()); + let _ = std::fs::remove_dir_all(&dir); + + write_container(&dir, &spec_with_segment_bytes(vec![0u8; 64])).expect("write bad"); + assert!(!Vindex3Container::open(&dir).expect("open").is_bindable()); + let _ = std::fs::remove_dir_all(&dir); + } + + #[test] + fn every_defect_variant_renders_its_own_coordinates() { + // Display is the operator-facing surface; a variant that renders + // without its coordinates is a defect nobody can navigate to. + let rendered = [ + ContainerDefect::Storage { + layer: 1, + key: "routed/layer_001".into(), + why: "truncated".into(), + }, + ContainerDefect::UnknownProgramme { + layer: 2, + programme: "made-up-v1".into(), + }, + ContainerDefect::UnsatisfiedProgramme { + layer: 4, + programme: "gated-mlp-v1".into(), + missing: vec![RegionRole::Down], + }, + ] + .map(|d| d.to_string()); + assert!(rendered[0].contains("routed/layer_001") && rendered[0].contains("truncated")); + assert!(rendered[1].contains("made-up-v1") && rendered[1].contains("layer 2")); + assert!(rendered[2].contains("Down") && rendered[2].contains("gated-mlp-v1")); + } + + #[test] + fn a_defect_renders_with_the_coordinates_needed_to_fix_it() { + // V2-0 asks for {layer, bank, role, segment}; the entry is the fifth + // coordinate a routed bank needs to be navigable. + let d = ContainerDefect::MissingRegion { + layer: 3, + bank: 1, + entry: 7, + role: RegionRole::Down, + segment: "routed/layer_003".into(), + }; + let msg = d.to_string(); + assert!(msg.contains("layer 3"), "{msg}"); + assert!(msg.contains("bank 1"), "{msg}"); + assert!(msg.contains("entry 7"), "{msg}"); + assert!(msg.contains("Down"), "{msg}"); + assert!(msg.contains("routed/layer_003"), "{msg}"); + } +} diff --git a/crates/larql-vindex/src/format/vindex3/write.rs b/crates/larql-vindex/src/format/vindex3/write.rs new file mode 100644 index 000000000..0311b08e1 --- /dev/null +++ b/crates/larql-vindex/src/format/vindex3/write.rs @@ -0,0 +1,162 @@ +//! Assemble a VINDEX3 container directory. +//! +//! This is the piece whose absence meant "VINDEX3 works" could only ever be +//! claimed about the *runtime*: every parity result to date bound VINDEX3 +//! routes over operands sourced from a VINDEX2 file, because nothing could +//! produce a VINDEX3 file to source them from. `ContainerGeneration::V3` +//! existed only as a hand-written JSON string in a detection test. +//! +//! What a container is, physically: +//! +//! ```text +//! / +//! ├── index.json schema 3 — the root authority (§12) +//! ├── moe_manifest.json which programme interprets the banks (§8) +//! └── .lyrw LYRW v2 bank files (§6) +//! ``` +//! +//! The writer takes banks already planned and written by [`Lyrw2Writer`], and +//! its job is only to place them and declare them. It deliberately does not +//! *quantise* or *transcode* anything: a container assembler that also +//! converted formats would be the silent conversion §9.1 forbids, one layer +//! further down than anyone would think to look for it. + +use std::collections::BTreeMap; +use std::path::{Path, PathBuf}; + +use super::index::Vindex3Index; +use crate::format::filenames::INDEX_JSON; +use crate::format::moe_manifest::MoeManifest; +use crate::VindexError; + +/// Filename of the MoE programme manifest within a container. +pub const MOE_MANIFEST_JSON: &str = "moe_manifest.json"; +/// Extension for a LYRW v2 segment file. +pub const SEGMENT_EXT: &str = "lyrw"; + +/// One bank's bytes, already in LYRW v2 form, plus where it belongs. +pub struct SegmentSource { + /// Path stem relative to the container root, e.g. `routed/layer_000`. + /// Becomes `/.lyrw`. + pub key: String, + /// A complete LYRW v2 file as produced by `Lyrw2Writer`. + pub bytes: Vec, +} + +/// Everything a container needs before any byte is placed. +pub struct ContainerSpec { + pub model: String, + pub family: String, + pub hidden_size: usize, + pub num_layers: usize, + pub manifest: MoeManifest, + pub segments: Vec, +} + +/// Path a segment key resolves to. Composed, never globbed. +pub fn segment_path(root: &Path, key: &str) -> PathBuf { + root.join(format!("{key}.{SEGMENT_EXT}")) +} + +/// Write a complete VINDEX3 container to `root`. +/// +/// Ordering is deliberate: segments, then the manifest, then `index.json` +/// last. `index.json` is the discriminator every reader dispatches on, so a +/// crash midway leaves a directory that is *not yet* a VINDEX3 container +/// rather than one that claims to be and is missing its banks. +pub fn write_container(root: &Path, spec: &ContainerSpec) -> Result<(), VindexError> { + if spec.segments.is_empty() { + return Err(VindexError::Parse( + "a VINDEX3 container needs at least one segment; an index declaring \ + none cannot be bound and would fail at execution instead of here" + .into(), + )); + } + std::fs::create_dir_all(root).map_err(VindexError::Io)?; + + let mut declared: BTreeMap = BTreeMap::new(); + for segment in &spec.segments { + let path = segment_path(root, &segment.key); + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent).map_err(VindexError::Io)?; + } + std::fs::write(&path, &segment.bytes).map_err(VindexError::Io)?; + *declared.entry(segment.key.clone()).or_insert(0) += 1; + } + + let manifest_json = serde_json::to_string_pretty(&spec.manifest) + .map_err(|e| VindexError::Parse(format!("serialise moe_manifest: {e}")))?; + std::fs::write(root.join(MOE_MANIFEST_JSON), manifest_json).map_err(VindexError::Io)?; + + let index = Vindex3Index::new( + spec.model.clone(), + spec.family.clone(), + spec.hidden_size, + spec.num_layers, + MOE_MANIFEST_JSON, + declared, + ); + let index_json = serde_json::to_string_pretty(&index) + .map_err(|e| VindexError::Parse(format!("serialise index.json: {e}")))?; + std::fs::write(root.join(INDEX_JSON), index_json).map_err(VindexError::Io)?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::format::generation::{detect_generation, ContainerGeneration}; + use crate::format::vindex3::test_support::{fixture_a_spec, FIXTURE_A_SEGMENT_KEY}; + use tempfile::tempdir; + + #[test] + fn a_written_container_is_detected_as_vindex3_from_disk() { + // The claim the whole exercise exists to support: not a JSON literal + // in a test asserting version 3, but a directory this code produced, + // read back by the same detector production uses. + let dir = tempdir().unwrap(); + write_container(dir.path(), &fixture_a_spec()).expect("write container"); + assert_eq!( + detect_generation(dir.path()).unwrap(), + ContainerGeneration::V3 + ); + } + + #[test] + fn a_written_container_has_all_three_physical_parts() { + let dir = tempdir().unwrap(); + write_container(dir.path(), &fixture_a_spec()).expect("write container"); + assert!(dir.path().join(INDEX_JSON).is_file()); + assert!(dir.path().join(MOE_MANIFEST_JSON).is_file()); + assert!(segment_path(dir.path(), FIXTURE_A_SEGMENT_KEY).is_file()); + } + + #[test] + fn the_index_declares_every_segment_that_was_written() { + // A segment on disk that the index does not declare is unreachable — + // the loader composes paths from the index and never scans. + let dir = tempdir().unwrap(); + write_container(dir.path(), &fixture_a_spec()).expect("write container"); + let raw = std::fs::read_to_string(dir.path().join(INDEX_JSON)).unwrap(); + let index: Vindex3Index = serde_json::from_str(&raw).unwrap(); + assert_eq!(index.segments.get(FIXTURE_A_SEGMENT_KEY), Some(&1)); + } + + #[test] + fn a_container_with_no_segments_is_refused_at_the_api() { + // Refused here, where the message can say why, rather than at bind + // time where it would surface as a missing operand. + let dir = tempdir().unwrap(); + let mut spec = fixture_a_spec(); + spec.segments.clear(); + let err = write_container(dir.path(), &spec).unwrap_err(); + assert!( + format!("{err}").contains("at least one segment"), + "unexpected: {err}" + ); + assert!( + !dir.path().join(INDEX_JSON).exists(), + "a refused write must not leave a directory claiming to be a container" + ); + } +} diff --git a/crates/larql-vindex/src/format/weights/mod.rs b/crates/larql-vindex/src/format/weights/mod.rs index 737efcaca..8c7c12273 100644 --- a/crates/larql-vindex/src/format/weights/mod.rs +++ b/crates/larql-vindex/src/format/weights/mod.rs @@ -25,6 +25,8 @@ pub mod write_f32; mod write_f32_tests; pub mod write_kquant; pub mod write_layers; +#[cfg(test)] +mod write_layers_parts_tests; pub(crate) use capabilities::ensure_extract_level_supported; diff --git a/crates/larql-vindex/src/format/weights/write_kquant/mod.rs b/crates/larql-vindex/src/format/weights/write_kquant/mod.rs index 318320d85..48549536c 100644 --- a/crates/larql-vindex/src/format/weights/write_kquant/mod.rs +++ b/crates/larql-vindex/src/format/weights/write_kquant/mod.rs @@ -36,6 +36,9 @@ mod attn; mod ffn; mod lm_head; mod moe_layers; +mod moe_layers_per_expert; +#[cfg(test)] +mod moe_layers_per_expert_tests; mod norms; pub mod feature_major_down; @@ -257,6 +260,10 @@ pub fn write_model_weights_kquant_with_opts( attn::write_attn_weights_kquant(source, dir, num_layers, callbacks)?; ffn::write_interleaved_ffn_kquant(source, dir, num_layers, opts, callbacks)?; moe_layers::write_per_layer_moe_kquant(source, dir, num_layers)?; + // Separate-tensor MoE models fall through the packed writer above; without + // this they produce an index that verifies and slices cleanly and then + // panics on the first decoded token with no expert store. + moe_layers_per_expert::write_per_layer_moe_per_expert(source, dir, num_layers)?; let mut entries = norms::write_norms_and_router(source, dir, num_layers)?; super::ple_sidecar::write_ple_weights(source, dir, num_layers, &mut entries)?; lm_head::write_lm_head_kquant(source, dir, &mut entries)?; diff --git a/crates/larql-vindex/src/format/weights/write_kquant/moe_layers_per_expert.rs b/crates/larql-vindex/src/format/weights/write_kquant/moe_layers_per_expert.rs new file mode 100644 index 000000000..d649bf433 --- /dev/null +++ b/crates/larql-vindex/src/format/weights/write_kquant/moe_layers_per_expert.rs @@ -0,0 +1,126 @@ +//! Per-layer expert weights for MoE models storing one tensor **per expert**. +//! +//! The packed path ([`super::moe_layers`]) handles models that stack every +//! expert into one tensor per projection (Gemma 4). This one handles the other +//! and more common arrangement — `experts.{id}.w1/w2/w3`, used by OLMoE, +//! Mixtral and DeepSeek. +//! +//! Why this file exists: the packed writer is gated on the packed layout, so +//! separate-tensor models fell through it and no expert store was written at +//! all. Extraction still reported success, checksums still verified, slicing +//! and WALK still worked — and decode panicked on the first token with +//! `layers/layer_00.weights` missing. An index that passes every integrity +//! check and cannot serve is worse than one that fails loudly, so the gap is +//! closed rather than diagnosed. +//! +//! The assembled entry layout is identical to the packed path's, because the +//! consumer is identical: `gate_up` is `[2*inter, hidden]` with **gate rows +//! first, then up rows**, and `down` is `[hidden, inter]` padded to a 256 +//! boundary for block formats. + +use std::path::Path; + +use larql_models::ModelArchitecture; + +use crate::error::VindexError; + +use super::super::write_f32::WeightSource; +use super::super::write_layers::{ + quantize_dense_entry, write_layer_weights, LayerEntry, LayerWeightFormat, +}; + +/// Write `layers/layer_{L:02}.weights` for every MoE layer of a per-expert model. +/// +/// Returns the number of layers written, so the caller can tell "not applicable" +/// (0, a dense or packed model) from "applicable and done". +pub(super) fn write_per_layer_moe_per_expert( + source: &dyn WeightSource, + dir: &Path, + num_layers: usize, +) -> Result { + let arch = source.arch(); + if !(arch.is_moe() && arch.expert_format() == larql_models::ExpertFormat::PerExpert) { + return Ok(0); + } + + let num_experts = arch.num_experts(); + let moe_inter = arch.moe_intermediate_size(); + let hidden = arch.config().hidden_size; + if num_experts == 0 || moe_inter == 0 || hidden == 0 { + return Ok(0); + } + + let format = LayerWeightFormat::Q4_K; + let mut written = 0usize; + + for layer in 0..num_layers { + let Some(entries) = + collect_layer_entries(source, arch, layer, num_experts, moe_inter, hidden, format)? + else { + // A dense layer inside a hybrid stack, or a layer whose expert + // tensors are absent. Skipping is correct; writing a short file + // would be the silent-wrong-bytes failure this module exists for. + continue; + }; + write_layer_weights(dir, layer, format, &entries, moe_inter, hidden)?; + written += 1; + } + + Ok(written) +} + +/// Build every expert entry for one layer, or `None` if the layer has no +/// expert tensors. +/// +/// Fails closed on a *partial* layer: if some experts resolve and others do +/// not, the layer is malformed and a short entry list would silently drop +/// experts that routing will later select. +fn collect_layer_entries( + source: &dyn WeightSource, + arch: &dyn ModelArchitecture, + layer: usize, + num_experts: usize, + moe_inter: usize, + hidden: usize, + format: LayerWeightFormat, +) -> Result>, VindexError> { + let mut entries = Vec::with_capacity(num_experts); + + for expert in 0..num_experts { + let parts = expert_parts(source, arch, layer, expert); + let Some((gate, up, down)) = parts else { + if expert == 0 { + return Ok(None); // layer has no experts at all + } + return Err(VindexError::MissingTensor(format!( + "layer {layer} expert {expert}: expert tensors are absent while \ + expert 0 resolved — the layer declares {num_experts} experts but \ + only {expert} are present" + ))); + }; + entries.push(quantize_dense_entry( + &gate, &up, &down, moe_inter, hidden, format, + )?); + } + + Ok(Some(entries)) +} + +/// Fetch one expert's `(gate, up, down)` as f32, or `None` if any is absent. +fn expert_parts( + source: &dyn WeightSource, + arch: &dyn ModelArchitecture, + layer: usize, + expert: usize, +) -> Option<(Vec, Vec, Vec)> { + let gate = arch + .expert_ffn_gate_key(layer, expert) + .and_then(|k| source.get_tensor(&k))?; + let up = arch + .expert_ffn_up_key(layer, expert) + .and_then(|k| source.get_tensor(&k))?; + let down = arch + .expert_ffn_down_key(layer, expert) + .and_then(|k| source.get_tensor(&k))?; + Some((gate.0, up.0, down.0)) +} diff --git a/crates/larql-vindex/src/format/weights/write_kquant/moe_layers_per_expert_tests.rs b/crates/larql-vindex/src/format/weights/write_kquant/moe_layers_per_expert_tests.rs new file mode 100644 index 000000000..bc2143488 --- /dev/null +++ b/crates/larql-vindex/src/format/weights/write_kquant/moe_layers_per_expert_tests.rs @@ -0,0 +1,276 @@ +//! Colocated tests for the separate-tensor MoE layer writer. +//! +//! The regression these defend against is specific and was live: extraction of +//! a separate-tensor MoE reported success, verified clean, sliced clean, and +//! then panicked on the first decoded token because no expert store had been +//! written. So the assertions are about *files appearing on disk with the +//! right shape*, not about the quantiser — which `write_layers_parts_tests` +//! covers separately. + +use std::collections::HashMap; +use std::path::Path; + +use super::super::write_f32::WeightSource; +use super::moe_layers_per_expert::write_per_layer_moe_per_expert; +use crate::format::weights::write_layers::parse_layer_weights_header; + +const HIDDEN: usize = 256; +const INTER: usize = 256; +const NUM_LAYERS: usize = 2; +const NUM_EXPERTS: usize = 3; + +/// A `WeightSource` backed by an explicit tensor map, so a test can express +/// "this expert is missing" precisely. +struct MapSource { + arch: Box, + tensors: HashMap, usize, usize)>, +} + +impl MapSource { + /// An OLMoE-shaped architecture — the real separate-tensor case. + fn olmoe(num_experts: usize) -> Box { + larql_models::detect_from_json(&serde_json::json!({ + "model_type": "olmoe", + "hidden_size": HIDDEN, + "intermediate_size": INTER, + "num_hidden_layers": NUM_LAYERS, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "num_experts": num_experts, + "num_experts_per_tok": 2, + })) + } + + /// Every expert of every layer present and correctly shaped. + fn complete() -> Self { + let arch = Self::olmoe(NUM_EXPERTS); + let mut tensors = HashMap::new(); + for layer in 0..NUM_LAYERS { + for expert in 0..NUM_EXPERTS { + insert_expert(&mut tensors, &*arch, layer, expert); + } + } + Self { arch, tensors } + } + + /// Layer 0 complete; layer 1 missing every expert (a hybrid dense layer). + fn first_layer_only() -> Self { + let arch = Self::olmoe(NUM_EXPERTS); + let mut tensors = HashMap::new(); + for expert in 0..NUM_EXPERTS { + insert_expert(&mut tensors, &*arch, 0, expert); + } + Self { arch, tensors } + } + + /// Layer 0 has expert 0 but not expert 1 — a genuinely malformed layer. + fn partial_layer() -> Self { + let arch = Self::olmoe(NUM_EXPERTS); + let mut tensors = HashMap::new(); + insert_expert(&mut tensors, &*arch, 0, 0); + Self { arch, tensors } + } +} + +fn insert_expert( + tensors: &mut HashMap, usize, usize)>, + arch: &dyn larql_models::ModelArchitecture, + layer: usize, + expert: usize, +) { + // Distinct fill per (layer, expert) so a mixed-up write is detectable. + let fill = (layer * 10 + expert) as f32; + if let Some(k) = arch.expert_ffn_gate_key(layer, expert) { + tensors.insert(k, (vec![fill; INTER * HIDDEN], INTER, HIDDEN)); + } + if let Some(k) = arch.expert_ffn_up_key(layer, expert) { + tensors.insert(k, (vec![-fill; INTER * HIDDEN], INTER, HIDDEN)); + } + if let Some(k) = arch.expert_ffn_down_key(layer, expert) { + tensors.insert(k, (vec![fill; HIDDEN * INTER], HIDDEN, INTER)); + } +} + +impl WeightSource for MapSource { + fn get_tensor(&self, key: &str) -> Option<(Vec, usize, usize)> { + self.tensors.get(key).cloned() + } + fn get_vector(&self, _key: &str) -> Option> { + None + } + fn arch(&self) -> &dyn larql_models::ModelArchitecture { + &*self.arch + } + fn num_layers(&self) -> usize { + NUM_LAYERS + } + fn lm_head(&self) -> Option<(Vec, usize, usize)> { + None + } + fn vector_names(&self) -> Vec { + Vec::new() + } + fn get_packed_bf16(&self, _key: &str) -> Option> { + None + } +} + +fn temp_dir(name: &str) -> std::path::PathBuf { + let dir = std::env::temp_dir().join("moe-per-expert-tests").join(name); + let _ = std::fs::remove_dir_all(&dir); + std::fs::create_dir_all(&dir).unwrap(); + dir +} + +fn layer_file(dir: &Path, layer: usize) -> std::path::PathBuf { + dir.join(crate::format::filenames::layer_weights_filename(layer)) +} + +#[test] +fn a_separate_tensor_moe_gets_an_expert_store() { + // The headline regression: this used to write nothing at all. + let dir = temp_dir("complete"); + let written = write_per_layer_moe_per_expert(&MapSource::complete(), &dir, NUM_LAYERS).unwrap(); + assert_eq!(written, NUM_LAYERS); + for layer in 0..NUM_LAYERS { + assert!( + layer_file(&dir, layer).exists(), + "layer {layer} not written" + ); + } +} + +#[test] +fn each_layer_file_declares_every_expert() { + let dir = temp_dir("entries"); + write_per_layer_moe_per_expert(&MapSource::complete(), &dir, NUM_LAYERS).unwrap(); + let bytes = std::fs::read(layer_file(&dir, 0)).unwrap(); + let (_, num_entries, inter, hidden, offsets) = parse_layer_weights_header(&bytes).unwrap(); + assert_eq!(num_entries, NUM_EXPERTS); + assert_eq!(inter, INTER); + assert_eq!(hidden, HIDDEN); + assert_eq!(offsets.len(), NUM_EXPERTS); +} + +#[test] +fn every_expert_region_is_inside_the_file() { + // An offset that is merely a plausible number is the failure mode the + // whole layer format is designed against; check it holds here too. + let dir = temp_dir("bounds"); + write_per_layer_moe_per_expert(&MapSource::complete(), &dir, NUM_LAYERS).unwrap(); + let bytes = std::fs::read(layer_file(&dir, 1)).unwrap(); + let (_, _, _, _, offsets) = parse_layer_weights_header(&bytes).unwrap(); + for (gu_off, gu_len, dn_off, dn_len) in offsets { + assert!(gu_off + gu_len <= bytes.len(), "gate_up region overruns"); + assert!(dn_off + dn_len <= bytes.len(), "down region overruns"); + } +} + +#[test] +fn a_layer_without_experts_is_skipped_not_written_short() { + // A hybrid stack's dense layer. Writing a zero-entry file here would look + // like a valid expert store to every downstream check. + let dir = temp_dir("hybrid"); + let written = + write_per_layer_moe_per_expert(&MapSource::first_layer_only(), &dir, NUM_LAYERS).unwrap(); + assert_eq!(written, 1); + assert!(layer_file(&dir, 0).exists()); + assert!(!layer_file(&dir, 1).exists()); +} + +#[test] +fn a_partially_present_layer_is_refused() { + // Expert 0 resolves but expert 1 does not. Writing the short list would + // silently drop experts that routing will later select. + let dir = temp_dir("partial"); + let err = + write_per_layer_moe_per_expert(&MapSource::partial_layer(), &dir, NUM_LAYERS).unwrap_err(); + let s = err.to_string(); + assert!(s.contains("expert 1"), "{s}"); + assert!(s.contains("only 1 are present"), "{s}"); +} + +#[test] +fn a_packed_model_is_left_to_the_packed_writer() { + // Gemma 4 is PackedBF16; this writer must decline it rather than race the + // packed path to the same filenames. + struct PackedSource(Box); + impl WeightSource for PackedSource { + fn get_tensor(&self, _k: &str) -> Option<(Vec, usize, usize)> { + None + } + fn get_vector(&self, _k: &str) -> Option> { + None + } + fn arch(&self) -> &dyn larql_models::ModelArchitecture { + &*self.0 + } + fn num_layers(&self) -> usize { + NUM_LAYERS + } + fn lm_head(&self) -> Option<(Vec, usize, usize)> { + None + } + fn vector_names(&self) -> Vec { + Vec::new() + } + fn get_packed_bf16(&self, _k: &str) -> Option> { + None + } + } + let arch = larql_models::detect_from_json(&serde_json::json!({ + "model_type": "gemma4_moe", + "hidden_size": HIDDEN, + "intermediate_size": INTER, + "num_hidden_layers": NUM_LAYERS, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "num_experts": NUM_EXPERTS, + "num_experts_per_tok": 2, + })); + let dir = temp_dir("packed"); + let written = write_per_layer_moe_per_expert(&PackedSource(arch), &dir, NUM_LAYERS).unwrap(); + assert_eq!(written, 0, "packed models must not be handled here"); + assert!(!layer_file(&dir, 0).exists()); +} + +#[test] +fn a_dense_model_is_declined() { + struct DenseSource(Box); + impl WeightSource for DenseSource { + fn get_tensor(&self, _k: &str) -> Option<(Vec, usize, usize)> { + None + } + fn get_vector(&self, _k: &str) -> Option> { + None + } + fn arch(&self) -> &dyn larql_models::ModelArchitecture { + &*self.0 + } + fn num_layers(&self) -> usize { + NUM_LAYERS + } + fn lm_head(&self) -> Option<(Vec, usize, usize)> { + None + } + fn vector_names(&self) -> Vec { + Vec::new() + } + fn get_packed_bf16(&self, _k: &str) -> Option> { + None + } + } + let arch = larql_models::detect_from_json(&serde_json::json!({ + "model_type": "llama", + "hidden_size": HIDDEN, + "intermediate_size": INTER, + "num_hidden_layers": NUM_LAYERS, + "num_attention_heads": 4, + "num_key_value_heads": 4, + })); + let dir = temp_dir("dense"); + assert_eq!( + write_per_layer_moe_per_expert(&DenseSource(arch), &dir, NUM_LAYERS).unwrap(), + 0 + ); +} diff --git a/crates/larql-vindex/src/format/weights/write_layers.rs b/crates/larql-vindex/src/format/weights/write_layers.rs index 0ba0a90e4..720fe5f52 100644 --- a/crates/larql-vindex/src/format/weights/write_layers.rs +++ b/crates/larql-vindex/src/format/weights/write_layers.rs @@ -53,11 +53,23 @@ const OFFSET_ENTRY_BYTES: usize = OFFSET_FIELDS_PER_ENTRY * U64_FIELD_BYTES; const BF16_BYTES: usize = std::mem::size_of::(); /// One quantized entry: gate+up bytes and down bytes, both in the same format. +/// +/// `Debug` prints byte counts rather than payloads — an expert is tens of MB, +/// and a failing assertion that dumps it is unreadable. pub struct LayerEntry { pub gate_up: Vec, // Q4_K [2*inter, hidden] pub down: Vec, // Q6_K [hidden, inter_padded] (same format as gate_up) } +impl std::fmt::Debug for LayerEntry { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("LayerEntry") + .field("gate_up_bytes", &self.gate_up.len()) + .field("down_bytes", &self.down.len()) + .finish() + } +} + pub type LayerWeightOffsets = Vec<(usize, usize, usize, usize)>; pub type LayerWeightsHeader = (LayerWeightFormat, usize, usize, usize, LayerWeightOffsets); @@ -176,10 +188,20 @@ pub fn pad_cols_to_256(data: &[f32], out_rows: usize, in_cols: usize) -> (Vec Result { - // gate+up interleaved: [gate rows, up rows] = [2*inter, hidden] + let expected_gate_up = inter * hidden; + if gate_f32.len() != expected_gate_up || up_f32.len() != expected_gate_up { + return Err(VindexError::Parse(format!( + "gate/up must each be [{inter}, {hidden}] = {expected_gate_up} elements; \ + got gate {} and up {}", + gate_f32.len(), + up_f32.len() + ))); + } + if down_f32.len() != hidden * inter { + return Err(VindexError::Parse(format!( + "down must be [{hidden}, {inter}] = {} elements; got {}", + hidden * inter, + down_f32.len() + ))); + } + let mut gate_up_f32 = Vec::with_capacity(2 * inter * hidden); gate_up_f32.extend_from_slice(gate_f32); gate_up_f32.extend_from_slice(up_f32); @@ -246,7 +284,16 @@ pub fn parse_layer_weights_header(data: &[u8]) -> Option { if magic != MAGIC { return None; } - // format_version at [4..8] — currently ignored, forward-compatible + // A newer `format_version` may change the offset-table stride. Parsing it + // with this version's stride would not bounds-fail — it would yield offsets + // that are still inside the file, and hand the caller a plausible byte range + // from the wrong place. Refuse instead: the one production caller + // (`format/weights/load/q4k.rs`) treats `None` as "skip this layer", so an + // unreadable file degrades to a clean miss rather than to wrong weights. + let version = u32::from_le_bytes(data[4..8].try_into().ok()?); + if version > FORMAT_VERSION { + return None; + } let quant_raw = u32::from_le_bytes(data[8..12].try_into().ok()?); let format = match quant_raw { 0 => LayerWeightFormat::F32, diff --git a/crates/larql-vindex/src/format/weights/write_layers_parts_tests.rs b/crates/larql-vindex/src/format/weights/write_layers_parts_tests.rs new file mode 100644 index 000000000..34210dc43 --- /dev/null +++ b/crates/larql-vindex/src/format/weights/write_layers_parts_tests.rs @@ -0,0 +1,161 @@ +//! Colocated tests for `quantize_dense_entry` — the separate-tensor +//! expert assembler. +//! +//! The load-bearing assertion is the **row order** inside `gate_up`: gate rows +//! first, then up rows, contiguous. Reversing it does not crash and does not +//! change a single byte count — it silently swaps the two halves of every +//! expert's GLU, which is the plausible-wrong-numbers failure this codebase +//! keeps designing against. So the order is pinned against the same split the +//! consumer performs (`cpu/ops/moe/expert` slices at `inter * hidden`), not +//! against a hand-written byte string. + +use super::write_layers::{quantize_dense_entry, LayerWeightFormat}; + +const INTER: usize = 2; +const HIDDEN: usize = 256; // Q4_K needs a 256 multiple on the contracted dim +const GATE_FILL: f32 = 1.0; +const UP_FILL: f32 = -1.0; +const DOWN_FILL: f32 = 0.5; + +fn gate() -> Vec { + vec![GATE_FILL; INTER * HIDDEN] +} + +fn up() -> Vec { + vec![UP_FILL; INTER * HIDDEN] +} + +fn down() -> Vec { + vec![DOWN_FILL; HIDDEN * INTER] +} + +#[test] +fn f32_entry_places_gate_rows_before_up_rows() { + // At F32 the payload is the input verbatim, so the split is directly + // observable — the property every quantised format must also preserve. + let entry = quantize_dense_entry( + &gate(), + &up(), + &down(), + INTER, + HIDDEN, + LayerWeightFormat::F32, + ) + .unwrap(); + + let floats: Vec = entry + .gate_up + .chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect(); + + assert_eq!(floats.len(), 2 * INTER * HIDDEN); + let (gate_half, up_half) = floats.split_at(INTER * HIDDEN); + assert!( + gate_half.iter().all(|&v| v == GATE_FILL), + "gate rows must come first" + ); + assert!( + up_half.iter().all(|&v| v == UP_FILL), + "up rows must come second" + ); +} + +#[test] +fn gate_up_is_twice_the_single_projection_size() { + let entry = quantize_dense_entry( + &gate(), + &up(), + &down(), + INTER, + HIDDEN, + LayerWeightFormat::F32, + ) + .unwrap(); + assert_eq!(entry.gate_up.len(), 2 * INTER * HIDDEN * 4); +} + +#[test] +fn down_is_padded_to_the_block_boundary() { + // inter = 2 pads to 256 under a block format, so `down` is written at + // [hidden, 256] rather than [hidden, 2]. + let entry = quantize_dense_entry( + &gate(), + &up(), + &down(), + INTER, + HIDDEN, + LayerWeightFormat::F32, + ) + .unwrap(); + assert_eq!(entry.down.len(), HIDDEN * 256 * 4); +} + +#[test] +fn a_quantised_entry_is_smaller_than_the_f32_one() { + let f32_entry = quantize_dense_entry( + &gate(), + &up(), + &down(), + INTER, + HIDDEN, + LayerWeightFormat::F32, + ) + .unwrap(); + let q4k_entry = quantize_dense_entry( + &gate(), + &up(), + &down(), + INTER, + HIDDEN, + LayerWeightFormat::Q4_K, + ) + .unwrap(); + assert!(q4k_entry.gate_up.len() < f32_entry.gate_up.len()); + assert!(q4k_entry.down.len() < f32_entry.down.len()); +} + +#[test] +fn a_short_gate_is_refused_by_shape() { + let err = quantize_dense_entry( + &vec![GATE_FILL; INTER * HIDDEN - 1], + &up(), + &down(), + INTER, + HIDDEN, + LayerWeightFormat::F32, + ) + .unwrap_err(); + assert!(err.to_string().contains("gate/up"), "{err}"); +} + +#[test] +fn a_short_up_is_refused_by_shape() { + let err = quantize_dense_entry( + &gate(), + &[UP_FILL; 3], + &down(), + INTER, + HIDDEN, + LayerWeightFormat::F32, + ) + .unwrap_err(); + assert!(err.to_string().contains("gate/up"), "{err}"); +} + +#[test] +fn a_transposed_down_is_refused_rather_than_reinterpreted() { + // [inter, hidden] instead of [hidden, inter] has the *same element count*, + // so only an explicit shape check catches it. Here the counts differ so the + // check fires; the guard exists because the symmetric case would not. + let err = quantize_dense_entry( + &gate(), + &up(), + &vec![DOWN_FILL; HIDDEN * INTER - 1], + INTER, + HIDDEN, + LayerWeightFormat::F32, + ) + .unwrap_err(); + assert!(err.to_string().contains("down"), "{err}"); +} diff --git a/crates/larql-vindex/src/lib.rs b/crates/larql-vindex/src/lib.rs index 6e9055357..75b66954b 100644 --- a/crates/larql-vindex/src/lib.rs +++ b/crates/larql-vindex/src/lib.rs @@ -36,6 +36,7 @@ pub mod index; pub mod kv_index_impl; pub mod patch; pub mod quant; +pub mod runtime; pub mod trie; pub mod walker; // Back-compat alias — the top-level lifecycle dir was renamed diff --git a/crates/larql-vindex/src/runtime/addressing.rs b/crates/larql-vindex/src/runtime/addressing.rs new file mode 100644 index 000000000..8eac114ae --- /dev/null +++ b/crates/larql-vindex/src/runtime/addressing.rs @@ -0,0 +1,144 @@ +//! How a region's stored bytes map to the elements it holds. +//! +//! Two answers, and the difference between them decides what a kernel can be +//! handed: +//! +//! ```text +//! Scalar one element per fixed-width slot; element i is at i * bytes +//! Blocked elements packed in super-blocks; element i is only reachable by +//! decoding the block that contains it +//! ``` +//! +//! A blocked region has no per-element stride, so the reference decoder — which +//! is written as "decode element at index" — cannot serve one at all. That is +//! not a defect in the bytes. It is a missing kernel, and saying so is the +//! whole point of keeping the two cases apart in the type rather than +//! discovering it as an arithmetic error deep inside a row read. +//! +//! # Geometry comes from the kernel registry +//! +//! `256 * 144` is not spelled here. Block geometry is +//! [`QuantFormat::packed_block_layout`]'s answer, so a codec whose block size +//! changes changes in one place and every reader follows. + +use crate::format::lyrw2::region_format::RegionFormat; +use larql_compute::QuantFormat; + +use super::consts::{BF16_BYTES, F16_BYTES, F32_BYTES}; + +/// How stored bytes are addressed for one region encoding. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Addressing { + /// Every element has its own fixed-width slot, so any element is reachable + /// by index alone. The reference decoder serves these. + Scalar { bytes: usize }, + /// Elements are packed in super-blocks that carry their own scales. + /// Reaching one element means decoding its whole block, so these are + /// handed to a block-native kernel whole rather than read per element. + Blocked { elements: usize, bytes: usize }, +} + +impl Addressing { + /// How this encoding is addressed, or `None` if the binary has no layout + /// for it — an unrecognised codec, or one with no registered geometry. + pub fn of(format: RegionFormat) -> Option { + match format { + RegionFormat::F32 => Some(Self::Scalar { bytes: F32_BYTES }), + RegionFormat::F16 => Some(Self::Scalar { bytes: F16_BYTES }), + RegionFormat::BF16 => Some(Self::Scalar { bytes: BF16_BYTES }), + other => quant_format(other) + .and_then(QuantFormat::packed_block_layout) + .map(|(elements, bytes)| Self::Blocked { elements, bytes }), + } + } + + /// Bytes a region holding `elements` values occupies. + /// + /// The blocked case rounds the *whole region* up to a block boundary + /// rather than each row, matching `QuantFormat::packed_matrix_bytes`. It is + /// therefore a lower bound when rows are individually block-aligned; the + /// exact per-row extent is checked where it matters, in + /// [`crate::runtime::tensor::BoundTensor::as_blocks`], because only a + /// block-native kernel has a row stride at all. + pub fn region_bytes(self, elements: usize) -> usize { + match self { + Self::Scalar { bytes } => elements * bytes, + Self::Blocked { + elements: per_block, + bytes, + } => elements.div_ceil(per_block) * bytes, + } + } + + /// Elements per super-block, or `None` for a directly-addressed encoding. + pub const fn block_elements(self) -> Option { + match self { + Self::Scalar { .. } => None, + Self::Blocked { elements, .. } => Some(elements), + } + } + + /// Bytes per super-block, or `None` for a directly-addressed encoding. + pub const fn block_bytes(self) -> Option { + match self { + Self::Scalar { .. } => None, + Self::Blocked { bytes, .. } => Some(bytes), + } + } +} + +/// The kernel-registry format a region codec corresponds to. +/// +/// Deliberately partial. A codec with no entry here has no packed geometry +/// this binary can state, and binding one is refused rather than guessed — +/// a wrong block size reads plausible bytes at the wrong offsets and produces +/// a well-shaped tensor of noise. +fn quant_format(format: RegionFormat) -> Option { + match format { + RegionFormat::Q4_0 => Some(QuantFormat::Q4_0), + RegionFormat::Q4K => Some(QuantFormat::Q4_K), + RegionFormat::Q6K => Some(QuantFormat::Q6_K), + _ => None, + } +} + +/// A block-packed operand handed to a kernel exactly as stored. +/// +/// Not a decoded copy and not a view object: the bytes are the region's own. +/// What the kernel additionally needs, and cannot recover from the bytes, is +/// the two column extents — because for a block-packed operand they differ. +/// +/// ```text +/// storage_cols what the kernel contracts over: it reads whole blocks and +/// cannot stop mid-block +/// role_cols what the operation means: the live columns +/// ``` +/// +/// Gemma's `down` is the case in point. Q4_K rounds the intermediate axis from +/// 704 up to 768, so the kernel must read 768 columns while the operation +/// means 704. The difference is neutralised on the *activation* side — the +/// padding columns multiply against zeros — which is why both numbers have to +/// travel together. Handing over only `storage_cols` would lose the operation's +/// meaning; handing over only `role_cols` would misread every row after the +/// first. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct BlockOperand<'a> { + /// Exactly `rows * row_bytes` of the region's own bytes. + pub bytes: &'a [u8], + pub rows: usize, + /// Elements per stored row — the extent a block-native kernel contracts + /// over. + pub storage_cols: usize, + /// Elements the role exposes. `storage_cols - role_cols` are quantisation + /// padding. + pub role_cols: usize, + /// Bytes per stored row. + pub row_bytes: usize, +} + +impl BlockOperand<'_> { + /// Padding columns the encoding added to reach a block boundary. + pub const fn padding_cols(&self) -> usize { + self.storage_cols - self.role_cols + } +} diff --git a/crates/larql-vindex/src/runtime/axis.rs b/crates/larql-vindex/src/runtime/axis.rs new file mode 100644 index 000000000..4811ae1a1 --- /dev/null +++ b/crates/larql-vindex/src/runtime/axis.rs @@ -0,0 +1,82 @@ +//! Named axes for shape diagnostics. +//! +//! An enum rather than a `&'static str` field, because these strings are +//! compared, matched on in tests, and rendered into operator-facing errors. +//! Spelling one of them differently at one call site produces a message that +//! reads correctly and groups wrongly. + +/// Which dimension a shape disagreement is about. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub enum Axis { + Rows, + Columns, + /// Total element count of a vector operand. + Length, + /// The width a vector-valued operand carries through the operation. + Width, + /// The width a projection contracts over. + InputWidth, + /// The width a projection or reduction produces. + OutputWidth, +} + +impl Axis { + pub const fn name(self) -> &'static str { + match self { + Self::Rows => "rows", + Self::Columns => "columns", + Self::Length => "length", + Self::Width => "width", + Self::InputWidth => "input width", + Self::OutputWidth => "output width", + } + } +} + +impl std::fmt::Display for Axis { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.name()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const ALL: [Axis; 6] = [ + Axis::Rows, + Axis::Columns, + Axis::Length, + Axis::Width, + Axis::InputWidth, + Axis::OutputWidth, + ]; + + #[test] + fn every_axis_has_a_distinct_name() { + let mut names: Vec<&str> = ALL.iter().map(|a| a.name()).collect(); + names.sort_unstable(); + let count = names.len(); + names.dedup(); + assert_eq!(names.len(), count, "two axes share a name"); + } + + #[test] + fn an_axis_displays_as_its_name() { + for axis in ALL { + assert_eq!(axis.to_string(), axis.name()); + } + } + + #[test] + fn axes_are_comparable_so_tests_can_assert_on_them() { + assert_eq!(Axis::Rows, Axis::Rows); + assert_ne!(Axis::Rows, Axis::Columns); + } + + #[test] + fn names_read_as_prose_in_a_diagnostic() { + assert_eq!(Axis::InputWidth.name(), "input width"); + assert_eq!(Axis::OutputWidth.name(), "output width"); + } +} diff --git a/crates/larql-vindex/src/runtime/bank.rs b/crates/larql-vindex/src/runtime/bank.rs new file mode 100644 index 000000000..0af088337 --- /dev/null +++ b/crates/larql-vindex/src/runtime/bank.rs @@ -0,0 +1,130 @@ +//! A bound expert bank — the population an operation can route into. +//! +//! The bank owns the expert list and the shapes they share. Selection happens +//! upstream in the router; the bank's job is to hand back the expert a +//! selected id names, and to refuse when the id has no expert behind it. +//! +//! That refusal matters more than it looks. A router and a bank that disagree +//! about the population is a binding fault, and the natural coding of it — +//! skip the expert, carry on — produces a token that is quietly missing a +//! fraction of its FFN contribution and looks entirely reasonable. + +use crate::format::capability::coordinate::BankCoordinate; +use larql_compute::Activation; + +use super::error::ExecutionError; +use super::expert_kernel::ExpertKernel; +use super::projection::BoundProjection; +use super::tensor::BoundTensor; + +/// One expert: its gated projection and its down projection. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct BoundExpert<'a> { + /// Id within the bank's population, as the router names it. + pub expert_id: u32, + pub projection: BoundProjection<'a>, + /// `[hidden, intermediate]` — contracts the intermediate axis away. + pub down: BoundTensor<'a>, +} + +impl BoundExpert<'_> { + pub fn validate(&self, intermediate: usize, hidden: usize) -> Result<(), ExecutionError> { + self.projection.validate(intermediate, hidden)?; + self.down.require_matrix(hidden, intermediate) + } +} + +/// A population of experts sharing one shape and one activation. +/// +/// `PartialEq` only: `Activation` is not `Eq`, which is the correct choice for +/// a type that names float behaviour rather than a discrete tag. +#[derive(Debug, Clone, PartialEq)] +pub struct BoundBankOperation<'a> { + /// Which bank this is. Diagnostics only — execution never branches on it. + pub bank: BankCoordinate, + pub experts: Vec>, + /// Intermediate width of every expert in the bank. + pub intermediate_dim: usize, + /// Input/output width — the residual width for a direct bank, the latent + /// width for a bank sitting behind routed-input/output transforms. + pub hidden_dim: usize, + pub activation: Activation, + /// Which kernel runs this bank's experts. + /// + /// A bank property rather than an expert one, because the kernel's + /// per-token state is shared across the bank: the incumbent quantises the + /// bank input to Q8_K once and every selected expert reads that one + /// activation. Two experts of one bank on two kernels would not be a + /// binding this runtime can express, and that is the correct restriction. + pub kernel: ExpertKernel, +} + +impl<'a> BoundBankOperation<'a> { + pub fn population(&self) -> usize { + self.experts.len() + } + + /// The expert an id names, or a refusal naming where it was looked for. + /// + /// Deliberately not `Option`. A selected-but-absent expert is never a + /// condition to handle locally: skipping it, or renormalising the + /// surviving weights around it, produces a token quietly missing part of + /// its FFN contribution and looks entirely reasonable. The request has to + /// reach the caller intact so that placement can satisfy it. + /// + /// `addressable` is the router's population, which is what makes the + /// report actionable — "expert 90 of 128, and this bank holds 8" says + /// *fetch it*, where "expert 90, bank holds 8" reads like corruption. + pub fn expert( + &self, + expert_id: u32, + addressable: usize, + ) -> Result<&BoundExpert<'a>, ExecutionError> { + self.experts + .iter() + .find(|e| e.expert_id == expert_id) + .ok_or_else(|| ExecutionError::SelectedExpertNotResident { + expert: expert_id, + bank: self.bank.describe(), + resident: self.population(), + population: addressable, + }) + } + + /// Whether this bank holds the expert an id names. + pub fn holds(&self, expert_id: u32) -> bool { + self.experts.iter().any(|e| e.expert_id == expert_id) + } + + /// Check every expert against the bank's declared shape, and against what + /// the bound kernel can take. + /// + /// Both, in that order. The shape check is what the operation means; the + /// kernel check is what this particular implementation of it can be handed. + /// A correctly-shaped expert the bound kernel cannot read is a binding + /// fault, and finding it here rather than mid-token is the whole reason + /// the kernel is a bound property. + pub fn validate(&self) -> Result<(), ExecutionError> { + for expert in &self.experts { + expert.validate(self.intermediate_dim, self.hidden_dim)?; + self.kernel.validate_operands( + expert, + self.intermediate_dim, + self.hidden_dim, + self.activation, + )?; + } + Ok(()) + } + + pub fn describe(&self) -> String { + format!( + "{} — {} experts, {}×{}, {} kernel", + self.bank.describe(), + self.population(), + self.intermediate_dim, + self.hidden_dim, + self.kernel.name() + ) + } +} diff --git a/crates/larql-vindex/src/runtime/consts.rs b/crates/larql-vindex/src/runtime/consts.rs new file mode 100644 index 000000000..2663e2bac --- /dev/null +++ b/crates/larql-vindex/src/runtime/consts.rs @@ -0,0 +1,46 @@ +//! Named constants for the reference runtime. +//! +//! Nothing here is tunable. These are facts about layouts and encodings that +//! would otherwise appear as bare numbers in indexing arithmetic, where a +//! transposed `2` reads as plausible and fails silently. + +/// A fused gate+up region stores two halves: all gate rows, then all up rows. +/// +/// The halves are contiguous, not interleaved. The distinction matters: an +/// interleaved reading of a concatenated region produces a well-shaped tensor +/// of wrong values, which survives to the logits. +pub const FUSED_PROJECTION_HALVES: usize = 2; + +/// Rank of a matrix contract. Anything else is a vector or unsupported. +pub const MATRIX_RANK: usize = 2; + +/// Row axis of a matrix contract. +pub const ROW_DIM: usize = 0; + +/// Column axis of a matrix contract. +pub const COL_DIM: usize = 1; + +/// Operand names for values that flow through an operation rather than being +/// stored. Stored operands name themselves via their `RepresentationIdentity`. +pub const OPERAND_RESIDUAL: &str = "residual"; +pub const OPERAND_OPERATION_OUTPUT: &str = "operation output"; + +/// Bytes per element, by reference-decodable encoding. +pub const F32_BYTES: usize = 4; +pub const F16_BYTES: usize = 2; +pub const BF16_BYTES: usize = 2; + +/// Where the mantissa of a bf16 value sits once widened to f32. +pub const BF16_SHIFT: u32 = 16; + +/// Operand layouts a kernel can ask for, named once so a refusal and the +/// check that produced it cannot describe the requirement differently. +pub const WANTED_ROW_MAJOR_F32: &str = "contiguous row-major f32"; +pub const WANTED_FUSED_PROJECTION: &str = "one fused gate+up region"; + +/// Stand-in when a refusal names a codec this binary has no name for. +pub const UNREGISTERED_CODEC: &str = "a registered block codec"; + +/// Extents a block-alignment refusal can be about. +pub const EXTENT_STORED_ROW: &str = "stored row"; +pub const EXTENT_KERNEL_INPUT: &str = "kernel input"; diff --git a/crates/larql-vindex/src/runtime/error.rs b/crates/larql-vindex/src/runtime/error.rs new file mode 100644 index 000000000..b86b7dca7 --- /dev/null +++ b/crates/larql-vindex/src/runtime/error.rs @@ -0,0 +1,220 @@ +//! Execution failures. +//! +//! Binding is where operands are chosen, checked and refused; by the time +//! `execute` runs, every question of *which* bytes and *whether they fit* has +//! been answered. So these variants are not the normal diagnostic surface — +//! they are the assertions that keep a binding bug from being interpreted as +//! numerics. +//! +//! Each one therefore names the operand and both sides of the disagreement. An +//! execution error that says only "dimension mismatch" sends the reader back +//! through the whole load path; one that says which tensor, which axis, and +//! what the two values were points at the binding decision that produced it. + +use crate::format::lyrw2::region_format::RegionFormat; + +use super::axis::Axis; + +/// Why an execution could not run to completion. +/// +/// `PartialEq` without `Eq`: `NonFiniteRouterScore` carries the offending +/// value, and reflexivity genuinely fails for the NaN it most often holds. +/// Claiming `Eq` here would be claiming a property this type does not have. +#[derive(Debug, Clone, PartialEq, thiserror::Error)] +pub enum ExecutionError { + /// The reference decoder has no implementation for this encoding. + /// + /// Not a defect in the index: the reference path deliberately implements + /// the directly-readable encodings only, and a quantised region is a + /// missing kernel rather than bad bytes. + #[error("the reference decoder does not implement {format} (needed for {operand})")] + UnsupportedFormat { format: String, operand: String }, + + /// The reference decoder cannot serve this access pattern. + #[error("the reference decoder does not implement a {view} view (needed for {operand})")] + UnsupportedView { view: String, operand: String }, + + /// An operand's shape disagrees with what the operation requires. + #[error("{operand}: expected {axis} of {expected}, found {found}")] + DimensionMismatch { + operand: String, + axis: Axis, + expected: usize, + found: usize, + }, + + /// A region is shorter than its declared shape needs. + #[error("{operand}: shape needs {needed} bytes, region holds {found}")] + ShortRegion { + operand: String, + needed: usize, + found: usize, + }, + + /// A row index fell outside the tensor. + #[error("{operand}: row {row} is out of range for {rows} rows")] + RowOutOfRange { + operand: String, + row: usize, + rows: usize, + }, + + /// A matrix operand turned out not to be a matrix. + #[error("{operand}: expected a matrix, found {found}")] + NotAMatrix { operand: String, found: String }, + + /// A router score is not a finite number. + /// + /// Refused rather than ordered. Under a total order a NaN still lands + /// *somewhere*, so it would be selected or rejected by sort mechanics — + /// a routing decision nobody made, reached silently, and propagated into + /// the residual stream as a plausible token. + #[error("router score for expert {expert} is {value}, which is not a finite number")] + NonFiniteRouterScore { expert: usize, value: f32 }, + + /// A learned per-expert scale is not a finite number. + /// + /// Multiplying a valid routing weight by it yields a NaN that propagates + /// through the reduction into the residual stream, where it surfaces as an + /// inexplicable token rather than as a bad operand. + #[error("{operand}: per-expert scale for expert {expert} is {value}, which is not finite")] + NonFiniteExpertScale { + expert: usize, + value: f32, + operand: String, + }, + + /// A bound kernel cannot take this operand's bytes as they are. + /// + /// Not a defect in the index and not a numerical problem: the kernel + /// requires a layout (contiguous row-major f32, say) that this operand + /// does not have. The repair is to bind a different kernel, never to + /// silently materialise a converted copy — that would make a parity + /// result a statement about the copy rather than about the binding. + #[error("the {kernel} kernel cannot take {operand} as bound: {reason}")] + KernelOperandUnsuitable { + kernel: &'static str, + operand: String, + reason: OperandUnsuitability, + }, + + /// A bound kernel does not implement the activation the bank declares. + /// + /// Distinct from an operand refusal because no rebinding of the *bytes* + /// fixes it — the kernel computes a different function than the model + /// specifies, and the only repairs are a different kernel or a corrected + /// recipe. + /// + /// It exists because the failure is otherwise invisible. The incumbent + /// Q4_K kernel branches on one activation and falls through to SiLU for + /// the rest, so a bank bound with `ReLU` would run SiLU, return finite + /// plausible values, and signal nothing at all. + #[error("the {kernel} kernel does not implement the {activation} activation")] + KernelActivationUnsupported { + kernel: &'static str, + activation: String, + }, + + /// An expert id outside the router's address space. + /// + /// A catalogue fault: nothing could ever select this expert, so the router + /// and the bank disagree about which population they are describing. + #[error("expert {expert} is outside the router's addressable population of {population}")] + ExpertOutOfRange { expert: u32, population: usize }, + + /// Routing selected an expert this bank does not hold. + /// + /// **Not** a catalogue fault. The routing decision is correct and the + /// expert exists in the model; it is simply not resident here, which is + /// the normal condition for a shard. Distinct from + /// [`Self::ExpertOutOfRange`] because the repairs differ completely: that + /// one means the index is wrong, this one means the operand must be + /// fetched from wherever it lives. + /// + /// It must never degrade into skipping the expert, renormalising the + /// surviving weights, or substituting a neighbour. Each of those produces + /// a token missing part of its FFN contribution and looks entirely + /// reasonable. Remote placement will later satisfy this request without + /// changing router semantics — which is only possible if the request + /// reaches the caller intact. + #[error( + "routing selected expert {expert}, which is not resident in {bank} \ + (this bank holds {resident} of {population} experts)" + )] + SelectedExpertNotResident { + expert: u32, + bank: String, + resident: usize, + population: usize, + }, +} + +/// Why a bound operand cannot be handed to a particular kernel. +/// +/// Kept apart because they lead to different remedies: choose another kernel, +/// bind another representation of the same component, repack, or reject the +/// index. Collapsing them into one message would leave the reader unable to +/// tell which. A new variant earns its place by having a remedy none of the +/// existing ones implies — not by describing a new *place* the same remedy +/// applies. +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub enum OperandUnsuitability { + /// The stored encoding is not what the kernel reads. + /// Remedy: bind a different variant, or a kernel for this format. + #[error("stored as {found}, kernel reads {wanted}")] + ElementFormat { found: String, wanted: &'static str }, + /// The operand is read through a view, so its bytes are not the operand. + /// Remedy: a kernel that understands the view, or a repacked variant. + #[error("read through a {view} view, kernel needs the bytes as stored")] + NonDirectView { view: String }, + /// The base address does not meet the kernel's alignment. + /// Remedy: an aligned copy, or a kernel with no alignment requirement. + #[error("base address is not aligned for {wanted}")] + MisalignedBase { wanted: &'static str }, + /// A block-packed operand's rows are not whole super-blocks, so the + /// kernel — which decodes a block at a time and cannot stop inside one — + /// has no row stride to read at. + /// Remedy: a repacked region, or pad the extent to a block boundary. + #[error("{extent} of {found} is not a whole number of {block}-element blocks")] + BlockAlignment { + extent: &'static str, + found: usize, + block: usize, + }, + /// The component is stored in a different arrangement from the one the + /// kernel's entry point takes — a decomposed pair where it takes one + /// fused slab, say. + /// Remedy: bind the arrangement this kernel takes, or a kernel for this + /// arrangement. Never re-stitch the regions: that would make the parity + /// result a statement about the stitching. + #[error("stored as {found}, kernel takes {wanted}")] + Arrangement { found: String, wanted: &'static str }, + /// The region holds fewer elements than the declared shape. + /// Remedy: reject the index — this one is a defect. + #[error("shape declares {expected} elements, region provides {found}")] + Length { expected: usize, found: usize }, +} + +impl ExecutionError { + /// Construct an unsupported-format error from a region encoding. + pub fn unsupported_format(format: RegionFormat, operand: impl Into) -> Self { + Self::UnsupportedFormat { + format: format.name(), + operand: operand.into(), + } + } +} + +/// Carries this error's response category across a crate boundary. +/// +/// `larql-compute` owns `FfnBackend` and sits below this crate, so it cannot +/// name `ExecutionError` — but it can name `dyn ExecutionRefusal`. The concrete +/// diagnosis stays here, with its expert ids, bank coordinates and axes; only +/// the one thing an engine must switch on crosses. +impl larql_execution::ExecutionRefusal for ExecutionError { + fn kind(&self) -> larql_execution::RefusalKind { + // Delegates rather than re-matching: two exhaustive matches over the + // same variants is exactly how a classification drifts. + self.refusal() + } +} diff --git a/crates/larql-vindex/src/runtime/execute.rs b/crates/larql-vindex/src/runtime/execute.rs new file mode 100644 index 000000000..33defae8e --- /dev/null +++ b/crates/larql-vindex/src/runtime/execute.rs @@ -0,0 +1,201 @@ +//! Executing a bound MoE operation. +//! +//! ```text +//! residual → routed_input → router → selection → experts → reduction +//! → routed_output → residual delta +//! ``` +//! +//! One implementation, generic over its trace sink. The fixture path and the +//! token path differ only in which sink is instantiated, so a fixture that +//! passes is evidence about the code Gemma runs rather than about a parallel +//! instrumented copy of it. + +use super::axis::Axis; +use super::consts::{OPERAND_OPERATION_OUTPUT, OPERAND_RESIDUAL}; +use super::error::ExecutionError; +use super::expert_kernel::BankKernel; +use super::inputs::MoeInputs; +use super::kernels::{dot, renormalize, softmax, top_k_with_margin}; +use super::operation::BoundMoeOperation; +use super::router::{BoundRouter, RouterKernel, SelectedExpert}; +use super::trace::{CollectedTrace, NoTrace, TraceSink}; +use super::transform::{apply_stage, TransformStage}; +use larql_compute::MoeTopKWeightPolicy; + +/// Run the operation over one token's residual, returning its contribution. +/// +/// The returned vector is the **delta**, not the updated residual: adding it +/// back is the caller's business, because whether a norm sits between the two +/// is a property of the surrounding block rather than of the MoE operation. +pub fn execute( + operation: &BoundMoeOperation<'_>, + inputs: MoeInputs<'_>, +) -> Result, ExecutionError> { + execute_with(operation, inputs, &mut NoTrace) +} + +/// Run the operation and record every internal checkpoint. +pub fn execute_traced( + operation: &BoundMoeOperation<'_>, + inputs: MoeInputs<'_>, +) -> Result<(Vec, CollectedTrace), ExecutionError> { + let mut trace = CollectedTrace::default(); + let out = execute_with(operation, inputs, &mut trace)?; + Ok((out, trace)) +} + +/// The single implementation. `sink` compiles away entirely for [`NoTrace`]. +/// +/// Routed-input transforms apply to `inputs.bank`. `inputs.router` is used as +/// supplied, because the router's input is produced by the surrounding block's +/// norms and scales, not by a projection this operation owns. +pub fn execute_with( + operation: &BoundMoeOperation<'_>, + inputs: MoeInputs<'_>, + sink: &mut S, +) -> Result, ExecutionError> { + if inputs.bank.len() != operation.residual_dim { + return Err(ExecutionError::DimensionMismatch { + operand: OPERAND_RESIDUAL.into(), + axis: Axis::Width, + expected: operation.residual_dim, + found: inputs.bank.len(), + }); + } + + let bank_input = apply_stage( + &operation.transforms, + TransformStage::RoutedInput, + inputs.bank, + )?; + sink.routed_input(&bank_input); + + let mut reduced = vec![0.0f32; operation.bank_input_dim()]; + for bank in &operation.banks { + let selected = select(&operation.router, inputs.router, bank.population(), sink)?; + // Opened before the expert loop, not inside it. The incumbent kernel + // quantises the bank input to Q8_K once and shares it across the + // bank's selected experts; quantising per expert would be binding a + // differently-shaped call than production makes. + let mut kernel = BankKernel::open(bank, &bank_input)?; + for choice in &selected { + let expert = bank.expert(choice.expert_id, operation.router.population())?; + let out = kernel.run(expert, bank, &bank_input)?; + sink.expert_output(choice.expert_id, &out); + operation + .reduction + .accumulate(&mut reduced, &out, choice.weight)?; + } + } + sink.reduced(&reduced); + + let delta = apply_stage( + &operation.transforms, + TransformStage::RoutedOutput, + &reduced, + )?; + if delta.len() != operation.residual_dim { + return Err(ExecutionError::DimensionMismatch { + operand: OPERAND_OPERATION_OUTPUT.into(), + axis: Axis::Width, + expected: operation.residual_dim, + found: delta.len(), + }); + } + sink.residual_delta(&delta); + Ok(delta) +} + +/// Score the population with whichever kernel was bound. +/// +/// The reference reads row by row through the operand's view, so it serves any +/// layout. The incumbent kernel takes the bytes directly and refuses rather +/// than accepting a reconstruction. +fn score(router: &BoundRouter<'_>, input: &[f32]) -> Result, ExecutionError> { + match router.kernel { + RouterKernel::Reference => { + let mut scores = vec![0.0f32; router.population()]; + let mut row = vec![0.0f32; router.hidden_dim()]; + for (e, slot) in scores.iter_mut().enumerate() { + router.weight.row_into(e, &mut row)?; + *slot = dot(&row, input); + } + softmax(&mut scores); + Ok(scores) + } + RouterKernel::Incumbent => { + let weights = router.weight.as_f32_slice().map_err(|reason| { + ExecutionError::KernelOperandUnsuitable { + kernel: RouterKernel::Incumbent.name(), + operand: router.weight.describe(), + reason, + } + })?; + // The incumbent's own functions, not a lookalike. + let mut scores = larql_compute::cpu::ops::moe::moe_score_experts( + input, + weights, + router.population(), + router.hidden_dim(), + ); + larql_compute::cpu::ops::moe::moe_softmax(&mut scores); + Ok(scores) + } + } +} + +/// Score the population, select top-k, and apply the weight policies. +fn select( + router: &BoundRouter<'_>, + input: &[f32], + population: usize, + sink: &mut S, +) -> Result, ExecutionError> { + let scores = score(router, input)?; + sink.router_scores(&scores); + + // Before ordering, not after. A non-finite score still takes a position + // under a total order, so leaving it in would let sort mechanics make a + // routing decision. + if let Some((expert, value)) = scores + .iter() + .position(|s| !s.is_finite()) + .map(|i| (i, scores[i])) + { + return Err(ExecutionError::NonFiniteRouterScore { expert, value }); + } + + let (chosen, margin) = top_k_with_margin(&scores, router.top_k); + sink.selection_margin(margin); + let mut weights: Vec = chosen.iter().map(|(_, w)| *w).collect(); + if router.selected_weight == MoeTopKWeightPolicy::RenormalizedSoftmax { + renormalize(&mut weights); + } + + // Per-expert scale multiplies *after* renormalisation, so the selected + // weights need not sum to one afterwards. Applying it before would let + // renormalisation divide the learned scale back out. + let per_expert = match router.scaling.scales() { + Some(scales) => Some(scales.to_vec()?), + None => None, + }; + + let mut selected = Vec::with_capacity(chosen.len()); + for ((index, raw_score), weight) in chosen.into_iter().zip(weights) { + let expert_id = u32::try_from(index).map_err(|_| ExecutionError::ExpertOutOfRange { + expert: u32::MAX, + population, + })?; + let weight = match &per_expert { + Some(scale) => weight * scale.get(index).copied().unwrap_or(1.0), + None => weight, + }; + selected.push(SelectedExpert { + expert_id, + weight, + raw_score, + }); + } + sink.selection(&selected); + Ok(selected) +} diff --git a/crates/larql-vindex/src/runtime/expert.rs b/crates/larql-vindex/src/runtime/expert.rs new file mode 100644 index 000000000..457506149 --- /dev/null +++ b/crates/larql-vindex/src/runtime/expert.rs @@ -0,0 +1,55 @@ +//! One expert's gated MLP. +//! +//! ```text +//! y = down · ( act(gate · x) ⊙ (up · x) ) +//! ``` +//! +//! Written once, against the projection interface rather than against storage, +//! so a fused and a decomposed expert reach identical arithmetic. If this +//! function ever needs to know which arrangement it got, the abstraction has +//! failed and the two routes can diverge. + +use larql_compute::Activation; + +use super::bank::BoundExpert; +use super::error::ExecutionError; +use super::kernels::{activate, dot}; + +/// Run one expert over `input`, returning its unweighted output. +pub fn forward( + expert: &BoundExpert<'_>, + input: &[f32], + intermediate: usize, + activation: Activation, +) -> Result, ExecutionError> { + let hidden = input.len(); + let mut gate_row = vec![0.0f32; hidden]; + let mut up_row = vec![0.0f32; hidden]; + let mut gated = vec![0.0f32; intermediate]; + + for (unit, slot) in gated.iter_mut().enumerate() { + expert.projection.gate_row_into(unit, &mut gate_row)?; + expert.projection.up_row_into(unit, &mut up_row)?; + *slot = dot(&gate_row, input); + // Activation on the gate half only; the up half enters unactivated. + // Applying it to the product instead is a plausible-looking variant + // that produces wrong values with no shape error anywhere. + *slot = activate_one(activation, *slot) * dot(&up_row, input); + } + + let out_dim = expert.down.rows(); + let mut out = vec![0.0f32; out_dim]; + let mut down_row = vec![0.0f32; intermediate]; + for (r, slot) in out.iter_mut().enumerate() { + expert.down.row_into(r, &mut down_row)?; + *slot = dot(&down_row, &gated); + } + Ok(out) +} + +/// Single-element activation, sharing the elementwise implementation. +fn activate_one(activation: Activation, x: f32) -> f32 { + let mut one = [x]; + activate(activation, &mut one); + one[0] +} diff --git a/crates/larql-vindex/src/runtime/expert_kernel.rs b/crates/larql-vindex/src/runtime/expert_kernel.rs new file mode 100644 index 000000000..6530ad08f --- /dev/null +++ b/crates/larql-vindex/src/runtime/expert_kernel.rs @@ -0,0 +1,326 @@ +//! Which kernel runs an expert, and the per-bank state it needs. +//! +//! A **bound** choice, exactly like [`RouterKernel`](super::router::RouterKernel). +//! Binding `larql-compute`'s own Q4_K × Q8_K kernel is what makes an expert +//! parity result a statement about the production code path; writing a +//! quantised matvec here and calling the agreement "parity" would prove only +//! that two similar loops agree. +//! +//! # The session exists because the kernel has per-bank state +//! +//! The incumbent quantises the bank input to Q8_K **once** and shares it +//! across the bank's selected experts. That is not an optimisation detail to +//! be reproduced or not as convenient — it is the shape of the call production +//! makes, and a binding that quantised per expert would be binding a different +//! call. So the kernel is opened once per bank, before the expert loop, and +//! the experts run through it. +//! +//! ```text +//! open(bank input) quantise to Q8_K once, allocate the kernel's scratch +//! run(expert) × k hand over each expert's blocks as stored +//! ``` +//! +//! # Refusal, never substitution +//! +//! Every way this kernel could quietly compute something else is refused +//! instead: +//! +//! ```text +//! decomposed gate/up Arrangement it takes one fused slab +//! non-Q4_K bytes ElementFormat it reads Q4_K super-blocks +//! rows not whole blocks BlockAlignment it has no stride to read at +//! input not whole blocks BlockAlignment Q8_K quantises per super-block +//! ReLU or exact GELU KernelActivationUnsupported +//! ``` +//! +//! The last one is the least obvious and the most dangerous. The incumbent's +//! inner loop reads `match activation { GeluTanh => .., _ => silu(..) }`, so a +//! bank bound with `ReLU` would run SiLU, produce finite plausible values, and +//! never signal anything. That is not the kernel being wrong — it is a kernel +//! that serves two activations being asked for a third. + +use larql_compute::cpu::ops::moe::{ + quantize_x_to_q8k, run_single_expert_q4k_q8k_into, ExpertScratch, Q8KActivation, +}; +use larql_compute::Activation; + +use crate::format::lyrw2::region_format::RegionFormat; + +use super::addressing::{Addressing, BlockOperand}; +use super::axis::Axis; +use super::bank::{BoundBankOperation, BoundExpert}; +use super::consts::{EXTENT_KERNEL_INPUT, FUSED_PROJECTION_HALVES, WANTED_FUSED_PROJECTION}; +use super::error::{ExecutionError, OperandUnsuitability}; +use super::expert; +use super::projection::BoundProjection; +use super::tensor::BoundTensor; + +/// The block encoding the incumbent expert kernel reads. +const INCUMBENT_Q4K_FORMAT: RegionFormat = RegionFormat::Q4K; + +/// Activations the incumbent expert kernel actually distinguishes. +/// +/// Its inner loop branches on `GeluTanh` and falls through to SiLU for +/// everything else, so this list is the kernel's real domain rather than the +/// vocabulary's. +const INCUMBENT_Q4K_ACTIVATIONS: [Activation; 2] = [Activation::Silu, Activation::GeluTanh]; + +/// Which kernel computes one expert's gated MLP. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum ExpertKernel { + /// Row-at-a-time f32 arithmetic through the operand's view. Serves any + /// directly-addressable encoding and any arrangement, and is the oracle + /// every other kernel is checked against. + #[default] + Reference, + /// `larql-compute`'s `run_single_expert_q4k_q8k_into`: Q4_K weights + /// against a Q8_K-quantised activation, integer dot throughout. + IncumbentQ4kQ8k, +} + +impl ExpertKernel { + pub const fn name(self) -> &'static str { + match self { + Self::Reference => "reference", + Self::IncumbentQ4kQ8k => "incumbent_q4k_q8k", + } + } + + /// Check one expert's operands against what this kernel can take. + /// + /// Load-path work: binding calls it once, through + /// [`BoundBankOperation::validate`]. Execution does not, because + /// re-checking per token is the resolution creep the bound object exists + /// to prevent — and because every refusal here is a property of the + /// binding, which cannot change between tokens. + pub fn validate_operands( + self, + expert: &BoundExpert<'_>, + intermediate: usize, + hidden: usize, + activation: Activation, + ) -> Result<(), ExecutionError> { + match self { + Self::Reference => Ok(()), + Self::IncumbentQ4kQ8k => { + if !INCUMBENT_Q4K_ACTIVATIONS.contains(&activation) { + return Err(ExecutionError::KernelActivationUnsupported { + kernel: self.name(), + activation: format!("{activation:?}"), + }); + } + incumbent_operands(expert, intermediate, hidden).map(|_| ()) + } + } + } +} + +/// What the incumbent kernel carries between the experts of one bank. +/// +/// Boxed at the enum, not inlined: the scratch buffers are sized for the bank's +/// widths, so the variant would otherwise set the size of every `BankKernel` +/// including the reference's empty one. +pub(crate) struct IncumbentSession { + input: Q8KActivation, + scratch: ExpertScratch, +} + +/// The bank input as each kernel needs it, prepared once per bank. +pub(crate) enum BankKernel { + Reference, + IncumbentQ4kQ8k(Box), +} + +impl BankKernel { + /// Prepare `bank`'s kernel over one token's bank input. + pub(crate) fn open( + bank: &BoundBankOperation<'_>, + bank_input: &[f32], + ) -> Result { + match bank.kernel { + ExpertKernel::Reference => Ok(Self::Reference), + ExpertKernel::IncumbentQ4kQ8k => { + let block = block_elements(); + // Q8_K quantises a super-block at a time, so a ragged tail + // would be dropped or read past. The incumbent guards the same + // condition by falling back to its f32 path; VINDEX3 refuses, + // because a silent change of kernel is a silent change of + // answer. + if !bank_input.len().is_multiple_of(block) { + return Err(ExecutionError::KernelOperandUnsuitable { + kernel: ExpertKernel::IncumbentQ4kQ8k.name(), + operand: bank.describe(), + reason: OperandUnsuitability::BlockAlignment { + extent: EXTENT_KERNEL_INPUT, + found: bank_input.len(), + block, + }, + }); + } + Ok(Self::IncumbentQ4kQ8k(Box::new(IncumbentSession { + input: quantize_x_to_q8k(bank_input), + scratch: ExpertScratch::new( + bank.hidden_dim, + bank.intermediate_dim, + padded_intermediate(bank.intermediate_dim, block), + ), + }))) + } + } + } + + /// Run one expert, returning its unweighted output. + /// + /// Returns an owned vector for both kernels. The reference allocates one + /// anyway, and the copy out of the incumbent's scratch keeps the two + /// arms' cost comparable without making the caller reason about a borrow + /// that lives across the reduction. + pub(crate) fn run( + &mut self, + expert: &BoundExpert<'_>, + bank: &BoundBankOperation<'_>, + bank_input: &[f32], + ) -> Result, ExecutionError> { + match self { + Self::Reference => { + expert::forward(expert, bank_input, bank.intermediate_dim, bank.activation) + } + Self::IncumbentQ4kQ8k(session) => { + let operands = incumbent_operands(expert, bank.intermediate_dim, bank.hidden_dim)?; + // The incumbent's own function, over the region's own bytes. + // Its internal short-slab guard — which would zero the output + // and return successfully — is unreachable from here: the + // operand checks above pin exactly the byte count it measures. + let out = run_single_expert_q4k_q8k_into( + &mut session.scratch, + &session.input, + operands.gate_up.bytes, + operands.down.bytes, + bank.intermediate_dim, + bank.activation, + ); + Ok(out.to_vec()) + } + } + } +} + +/// One expert's operands as the incumbent Q4_K kernel takes them. +struct IncumbentOperands<'a> { + gate_up: BlockOperand<'a>, + down: BlockOperand<'a>, +} + +/// Resolve and check both operands against the bank's shape. +/// +/// Shared by validation and execution deliberately: a check the load path runs +/// and the token path skips is a check that can drift out of agreement with +/// what the kernel is actually handed. +fn incumbent_operands<'a>( + expert: &BoundExpert<'a>, + intermediate: usize, + hidden: usize, +) -> Result, ExecutionError> { + let kernel = ExpertKernel::IncumbentQ4kQ8k.name(); + let block = block_elements(); + + // The kernel takes one slab and splits it in half. A decomposed pair is + // two regions, and stitching them into a temporary would make the parity + // result a statement about the stitching. + let BoundProjection::Fused { gate_up } = &expert.projection else { + return Err(ExecutionError::KernelOperandUnsuitable { + kernel, + operand: expert.projection.describe(), + reason: OperandUnsuitability::Arrangement { + found: expert.projection.arrangement().name(), + wanted: WANTED_FUSED_PROJECTION, + }, + }); + }; + let fused = blocks(gate_up, kernel)?; + require( + fused.rows, + FUSED_PROJECTION_HALVES * intermediate, + Axis::Rows, + || expert.projection.describe(), + )?; + require(fused.storage_cols, hidden, Axis::InputWidth, || { + expert.projection.describe() + })?; + + let down = blocks(&expert.down, kernel)?; + require(down.rows, hidden, Axis::OutputWidth, || { + expert.down.describe() + })?; + require(down.role_cols, intermediate, Axis::Columns, || { + expert.down.describe() + })?; + // The kernel derives its `down` stride from `intermediate` rounded up to a + // block, so a region padded to any other width would be read at the wrong + // stride and produce well-shaped noise. + require( + down.storage_cols, + padded_intermediate(intermediate, block), + Axis::InputWidth, + || expert.down.describe(), + )?; + + Ok(IncumbentOperands { + gate_up: fused, + down, + }) +} + +/// Hand over one operand's blocks, or name why the kernel cannot take it. +fn blocks<'a>( + operand: &BoundTensor<'a>, + kernel: &'static str, +) -> Result, ExecutionError> { + operand.as_blocks(INCUMBENT_Q4K_FORMAT).map_err(|reason| { + ExecutionError::KernelOperandUnsuitable { + kernel, + operand: operand.describe(), + reason, + } + }) +} + +/// Assert one resolved extent, naming the operand it belongs to. +/// +/// The name is built lazily: `describe` allocates, and this runs per operand +/// per expert per token on the incumbent path. +fn require( + found: usize, + expected: usize, + axis: Axis, + operand: impl FnOnce() -> String, +) -> Result<(), ExecutionError> { + if found == expected { + return Ok(()); + } + Err(ExecutionError::DimensionMismatch { + operand: operand(), + axis, + expected, + found, + }) +} + +/// Elements per Q4_K super-block, from the kernel registry. +/// +/// Infallible rather than a `Result`. The geometry comes from `QuantFormat`'s +/// registered layout, which is a property of the build and not of the data: a +/// binary whose registry had lost Q4_K could not have compiled this binding at +/// all. Threading a refusal through every caller would suggest an index could +/// provoke it. +fn block_elements() -> usize { + Addressing::of(INCUMBENT_Q4K_FORMAT) + .and_then(Addressing::block_elements) + .expect("the format registry carries Q4_K block geometry") +} + +/// The intermediate width once rounded up to a whole super-block — what the +/// kernel reads, and what `down` must therefore be stored at. +const fn padded_intermediate(intermediate: usize, block: usize) -> usize { + intermediate.div_ceil(block) * block +} diff --git a/crates/larql-vindex/src/runtime/fixtures/direct_moe.rs b/crates/larql-vindex/src/runtime/fixtures/direct_moe.rs new file mode 100644 index 000000000..e3086a65a --- /dev/null +++ b/crates/larql-vindex/src/runtime/fixtures/direct_moe.rs @@ -0,0 +1,214 @@ +//! Fixture A — a tiny direct routed-MoE layer, in both storage arrangements. +//! +//! Small enough that a wrong answer is inspectable by hand, and shaped so that +//! **no two axes share a size**: hidden 4, intermediate 3, population 5, top-k +//! 2. Equal dimensions hide transposes, and a transposed operand that still +//! multiplies is the failure this fixture most needs to catch — an earlier +//! draft used a population of 4 and could not have detected a transposed +//! router at all. +//! +//! The same weights are laid out twice — `gate`/`up` as separate regions, and +//! as one concatenated `gate_up_fused` region. Both must produce identical +//! output. That is the property that says the executor follows its bound +//! recipe rather than assuming a physical arrangement. + +use crate::format::capability::binding::RepresentationIdentity; +use crate::format::capability::component::ComponentContract; +use crate::format::capability::coordinate::BankCoordinate; +use crate::format::lyrw2::region_format::RegionFormat; +use crate::format::lyrw2::region_role::RegionRole; +use larql_compute::{Activation, MoeTopKWeightPolicy}; + +use super::super::bank::{BoundBankOperation, BoundExpert}; +use super::super::consts::FUSED_PROJECTION_HALVES; +use super::super::expert_kernel::ExpertKernel; +use super::super::operation::BoundMoeOperation; +use super::super::projection::{BoundProjection, ProjectionArrangement}; +use super::super::reduction::BoundReduction; +use super::super::router::{BoundExpertScaling, BoundRouter, RouterKernel}; +use super::super::tensor::BoundTensor; + +pub const HIDDEN: usize = 4; +pub const INTERMEDIATE: usize = 3; +pub const POPULATION: usize = 5; +pub const TOP_K: usize = 2; +pub const ACTIVATION: Activation = Activation::Silu; +pub const LAYER: u32 = 0; +pub const BANK_ID: u16 = 0; + +/// Catalogue variant these operands claim to come from. +const VARIANT: &str = "fixture-a"; + +/// The router is not a bank region and so has no registry role; it is a +/// manifest-addressed tensor and names itself. +const ROUTER_REGION_SET: &str = "router"; + +// ── Weight values ────────────────────────────────────────────────────────── +// +// Deterministic and shared with the oracle. Sharing the *data* is fine — the +// oracle's independence is about arithmetic, not about inventing its own +// numbers. Coprime multipliers keep the values from repeating across axes, so +// a swapped index changes the answer. + +pub fn gate_at(expert: usize, unit: usize, h: usize) -> f32 { + 0.1 * ((expert * 7 + unit * 3 + h) % 11) as f32 - 0.5 +} + +pub fn up_at(expert: usize, unit: usize, h: usize) -> f32 { + 0.1 * ((expert * 5 + unit * 2 + h * 3) % 9) as f32 - 0.4 +} + +pub fn down_at(expert: usize, h: usize, unit: usize) -> f32 { + 0.1 * ((expert * 3 + h * 5 + unit * 2) % 7) as f32 - 0.3 +} + +pub fn router_at(expert: usize, h: usize) -> f32 { + 0.1 * ((expert * 2 + h * 4) % 13) as f32 - 0.6 +} + +/// A residual that selects a non-trivial pair of experts. +pub fn input() -> Vec { + vec![0.35, -0.72, 0.18, 0.94] +} + +fn bytes(values: impl Iterator) -> Vec { + values.flat_map(f32::to_le_bytes).collect() +} + +/// Fixture A of the conformance programme: one direct routed-MoE layer. +/// +/// Owns the bytes so bound tensors can borrow them. +pub struct DirectMoeFixture { + gate: Vec>, + up: Vec>, + gate_up: Vec>, + down: Vec>, + router: Vec, +} + +impl Default for DirectMoeFixture { + fn default() -> Self { + Self::new() + } +} + +impl DirectMoeFixture { + pub fn new() -> Self { + let experts = 0..POPULATION; + Self { + gate: experts + .clone() + .map(|e| { + bytes( + (0..INTERMEDIATE) + .flat_map(move |i| (0..HIDDEN).map(move |h| gate_at(e, i, h))), + ) + }) + .collect(), + up: experts + .clone() + .map(|e| { + bytes( + (0..INTERMEDIATE) + .flat_map(move |i| (0..HIDDEN).map(move |h| up_at(e, i, h))), + ) + }) + .collect(), + // Concatenated, not interleaved: every gate row, then every up row. + gate_up: experts + .clone() + .map(|e| { + let gate = (0..INTERMEDIATE) + .flat_map(move |i| (0..HIDDEN).map(move |h| gate_at(e, i, h))); + let up = (0..INTERMEDIATE) + .flat_map(move |i| (0..HIDDEN).map(move |h| up_at(e, i, h))); + bytes(gate.chain(up)) + }) + .collect(), + down: experts + .map(|e| { + bytes( + (0..HIDDEN) + .flat_map(move |h| (0..INTERMEDIATE).map(move |i| down_at(e, h, i))), + ) + }) + .collect(), + router: bytes((0..POPULATION).flat_map(|e| (0..HIDDEN).map(move |h| router_at(e, h)))), + } + } + + /// Build the bound operation for one storage arrangement. + pub fn operation(&self, arrangement: ProjectionArrangement) -> BoundMoeOperation<'_> { + let experts = (0..POPULATION) + .map(|e| BoundExpert { + expert_id: e as u32, + projection: self.projection(e, arrangement), + down: region( + RegionRole::Down, + &self.down[e], + ComponentContract::matrix(HIDDEN as u32, INTERMEDIATE as u32), + ), + }) + .collect(); + + BoundMoeOperation { + router: BoundRouter { + weight: tensor( + ROUTER_REGION_SET, + &self.router, + ComponentContract::matrix(POPULATION as u32, HIDDEN as u32), + ), + top_k: TOP_K, + selected_weight: MoeTopKWeightPolicy::RenormalizedSoftmax, + scaling: BoundExpertScaling::None, + kernel: RouterKernel::default(), + }, + transforms: Vec::new(), + banks: vec![BoundBankOperation { + bank: BankCoordinate::new(LAYER, BANK_ID), + experts, + intermediate_dim: INTERMEDIATE, + hidden_dim: HIDDEN, + activation: ACTIVATION, + kernel: ExpertKernel::default(), + }], + reduction: BoundReduction::WeightedSum, + residual_dim: HIDDEN, + } + } + + fn projection(&self, expert: usize, arrangement: ProjectionArrangement) -> BoundProjection<'_> { + let rows = ComponentContract::matrix(INTERMEDIATE as u32, HIDDEN as u32); + match arrangement { + ProjectionArrangement::Decomposed => BoundProjection::Decomposed { + gate: region(RegionRole::Gate, &self.gate[expert], rows.clone()), + up: region(RegionRole::Up, &self.up[expert], rows), + }, + ProjectionArrangement::Fused => BoundProjection::Fused { + gate_up: region( + RegionRole::GateUpFused, + &self.gate_up[expert], + ComponentContract::matrix( + (INTERMEDIATE * FUSED_PROJECTION_HALVES) as u32, + HIDDEN as u32, + ), + ), + }, + } + } +} + +/// Bind a bank region, named by its registry role. +fn region(role: RegionRole, bytes: &[u8], contract: ComponentContract) -> BoundTensor<'_> { + tensor(&role.name(), bytes, contract) +} + +fn tensor<'a>(region_set: &str, bytes: &'a [u8], contract: ComponentContract) -> BoundTensor<'a> { + BoundTensor::direct( + RepresentationIdentity::new(region_set, VARIANT), + bytes, + RegionFormat::F32, + contract, + ) + .expect("fixture A operands are well-formed") +} diff --git a/crates/larql-vindex/src/runtime/fixtures/direct_moe_oracle.rs b/crates/larql-vindex/src/runtime/fixtures/direct_moe_oracle.rs new file mode 100644 index 000000000..d78ac4a6b --- /dev/null +++ b/crates/larql-vindex/src/runtime/fixtures/direct_moe_oracle.rs @@ -0,0 +1,104 @@ +//! An independent oracle for fixture A. +//! +//! Deliberately shares nothing with the runtime but the weight values. No +//! `BoundTensor`, no `kernels`, no `expert::forward` — its softmax, its top-k +//! and its activation are written out again here. +//! +//! That duplication is the point. An oracle built from the runtime's own +//! primitives agrees with the runtime by construction and tests nothing; the +//! only failure it can catch is one in the glue. This one computes the layer +//! from its definition, so agreement is evidence. + +use super::direct_moe::{ + down_at, gate_at, router_at, up_at, HIDDEN, INTERMEDIATE, POPULATION, TOP_K, +}; + +/// Every checkpoint, computed from the layer's definition. +pub struct Oracle { + pub router_scores: Vec, + /// Selected `(expert_id, final_weight)`, in selection order. + pub selected: Vec<(u32, f32)>, + /// Unweighted output per selected expert, in selection order. + pub expert_outputs: Vec<(u32, Vec)>, + pub reduced: Vec, +} + +pub fn oracle(input: &[f32]) -> Oracle { + let router_scores = route(input); + let selected = select(&router_scores); + let expert_outputs: Vec<(u32, Vec)> = selected + .iter() + .map(|&(id, _)| (id, expert(id as usize, input))) + .collect(); + + let mut reduced = vec![0.0f32; HIDDEN]; + for (&(_, weight), (_, out)) in selected.iter().zip(&expert_outputs) { + for (acc, &v) in reduced.iter_mut().zip(out) { + *acc += weight * v; + } + } + + Oracle { + router_scores, + selected, + expert_outputs, + reduced, + } +} + +/// Router logits, softmaxed over the whole population. +fn route(input: &[f32]) -> Vec { + let mut logits: Vec = (0..POPULATION) + .map(|e| (0..HIDDEN).map(|h| router_at(e, h) * input[h]).sum()) + .collect(); + let max = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max); + let mut sum = 0.0f32; + for x in &mut logits { + *x = (*x - max).exp(); + sum += *x; + } + for x in &mut logits { + *x /= sum; + } + logits +} + +/// Top-k by score, ties to the lower id, renormalised to sum to one. +fn select(scores: &[f32]) -> Vec<(u32, f32)> { + let mut ranked: Vec<(usize, f32)> = scores.iter().copied().enumerate().collect(); + ranked.sort_by(|a, b| { + b.1.partial_cmp(&a.1) + .unwrap_or(std::cmp::Ordering::Equal) + .then(a.0.cmp(&b.0)) + }); + ranked.truncate(TOP_K); + let total: f32 = ranked.iter().map(|(_, w)| *w).sum(); + ranked + .into_iter() + .map(|(i, w)| (i as u32, w / total)) + .collect() +} + +/// `down · ( silu(gate · x) ⊙ (up · x) )`, straight from the definition. +fn expert(expert: usize, input: &[f32]) -> Vec { + let gated: Vec = (0..INTERMEDIATE) + .map(|unit| { + let g: f32 = (0..HIDDEN) + .map(|h| gate_at(expert, unit, h) * input[h]) + .sum(); + let u: f32 = (0..HIDDEN).map(|h| up_at(expert, unit, h) * input[h]).sum(); + silu(g) * u + }) + .collect(); + (0..HIDDEN) + .map(|h| { + (0..INTERMEDIATE) + .map(|unit| down_at(expert, h, unit) * gated[unit]) + .sum() + }) + .collect() +} + +fn silu(x: f32) -> f32 { + x / (1.0 + (-x).exp()) +} diff --git a/crates/larql-vindex/src/runtime/fixtures/mod.rs b/crates/larql-vindex/src/runtime/fixtures/mod.rs new file mode 100644 index 000000000..8ce20cfbe --- /dev/null +++ b/crates/larql-vindex/src/runtime/fixtures/mod.rs @@ -0,0 +1,26 @@ +//! Conformance fixtures for the reference runtime. +//! +//! Public, not test-only. The experimental programme names fixtures A–D as +//! artifacts in their own right, and they have at least four consumers: the +//! colocated tests, the perf gate in `benches/`, the demo in `examples/`, and +//! eventually a conformance check run against a real container. A fixture that +//! only `cfg(test)` can see forces each of those to grow its own copy, and +//! copies disagree. +//! +//! Each fixture ships with its **oracle** — an independent implementation of +//! the same layer, sharing only the weight values. The pair is the artifact; +//! the fixture alone would say what the runtime does, not what it should do. +//! +//! ```text +//! direct_moe fixture A — one routed MoE layer, no latent transforms, +//! laid out both decomposed and fused +//! synthetic the same semantics at any shape, as a contiguous byte +//! image, for scaling / allocation / residency measurement +//! ``` + +pub mod direct_moe; +pub mod direct_moe_oracle; +pub mod synthetic; + +#[cfg(test)] +mod tests; diff --git a/crates/larql-vindex/src/runtime/fixtures/synthetic.rs b/crates/larql-vindex/src/runtime/fixtures/synthetic.rs new file mode 100644 index 000000000..c5b5e315f --- /dev/null +++ b/crates/larql-vindex/src/runtime/fixtures/synthetic.rs @@ -0,0 +1,238 @@ +//! A parameterised synthetic MoE layer, laid out as it would be on disk. +//! +//! Fixture A is fixed-size and hand-checkable. This is its scaling companion: +//! same semantics, any shape, and — crucially — a **contiguous byte image** +//! rather than a scatter of separate allocations. +//! +//! The byte image is what makes one builder serve three very different +//! consumers: +//! +//! ```text +//! benches/vindex3_bound_execute bind over a Vec cost vs population +//! tests/vindex3_allocation_guard bind over a Vec allocations per token +//! examples/vindex3_residency_probe bind over an mmap pages actually faulted +//! ``` +//! +//! `bind` takes the bytes as an argument rather than reading its own, so the +//! same layout can be bound over heap memory or a file mapping with identical +//! content. The residency probe depends on that: it must measure the layout +//! the other two measure, not a lookalike. +//! +//! # Layout +//! +//! ```text +//! [router] [expert 0: gate_up | down] [expert 1: gate_up | down] ... +//! ``` +//! +//! Each expert's regions are adjacent on purpose. Expert-major locality is the +//! property a routed read pattern depends on, and interleaving roles across +//! experts instead would scatter every selected expert across the whole file. + +use std::ops::Range; + +use crate::format::capability::binding::RepresentationIdentity; +use crate::format::capability::component::ComponentContract; +use crate::format::capability::coordinate::BankCoordinate; +use crate::format::lyrw2::region_format::RegionFormat; +use crate::format::lyrw2::region_role::RegionRole; +use larql_compute::{Activation, MoeTopKWeightPolicy}; + +use super::super::bank::{BoundBankOperation, BoundExpert}; +use super::super::consts::{F32_BYTES, FUSED_PROJECTION_HALVES}; +use super::super::expert_kernel::ExpertKernel; +use super::super::operation::BoundMoeOperation; +use super::super::projection::BoundProjection; +use super::super::reduction::BoundReduction; +use super::super::router::{BoundExpertScaling, BoundRouter, RouterKernel}; +use super::super::tensor::BoundTensor; + +/// Catalogue variant these operands claim to come from. +const VARIANT: &str = "synthetic"; +/// The router is manifest-addressed, not a bank region, so it has no role. +const ROUTER_REGION_SET: &str = "router"; +/// Modulus for the weight generator — prime, so values do not repeat on any +/// axis of a realistic shape. +const WEIGHT_MODULUS: usize = 199; +/// Seeds keeping the three region classes numerically distinct. +const SEED_GATE: usize = 0; +const SEED_UP: usize = 1_000; +const SEED_DOWN: usize = 2_000; +const SEED_ROUTER: usize = 3_000; +const SEED_RESIDUAL: usize = 4_000; + +/// Deterministic pseudo-weights. Content is irrelevant to timing and +/// residency; determinism keeps runs comparable. +pub fn weight(seed: usize, index: usize) -> f32 { + (((seed * 31 + index * 17) % WEIGHT_MODULUS) as f32) / WEIGHT_MODULUS as f32 - 0.5 +} + +/// The dimensions of a synthetic layer. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct SyntheticShape { + pub population: usize, + pub hidden: usize, + pub intermediate: usize, + pub top_k: usize, +} + +impl SyntheticShape { + /// Elements in one expert's fused gate+up region. + pub fn projection_elements(&self) -> usize { + FUSED_PROJECTION_HALVES * self.intermediate * self.hidden + } + + /// Elements in one expert's down region. + pub fn down_elements(&self) -> usize { + self.hidden * self.intermediate + } + + /// Bytes one expert occupies, both regions together. + pub fn expert_bytes(&self) -> usize { + (self.projection_elements() + self.down_elements()) * F32_BYTES + } + + pub fn router_bytes(&self) -> usize { + self.population * self.hidden * F32_BYTES + } + + /// Total bytes of the layer's byte image. + pub fn total_bytes(&self) -> usize { + self.router_bytes() + self.population * self.expert_bytes() + } + + /// Bytes a single token's routed read touches: the router, plus `top_k` + /// experts. The residency question in one number. + pub fn selected_bytes(&self) -> usize { + self.router_bytes() + self.top_k.min(self.population) * self.expert_bytes() + } +} + +/// Where each region sits in the byte image. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ExpertExtents { + pub gate_up: Range, + pub down: Range, +} + +/// A synthetic layer: its byte image and the extents that address it. +#[derive(Debug, Clone)] +pub struct SyntheticLayer { + pub shape: SyntheticShape, + pub bytes: Vec, + pub router: Range, + pub experts: Vec, +} + +impl SyntheticLayer { + /// Build the byte image for `shape`. + pub fn build(shape: SyntheticShape) -> Self { + let mut bytes = Vec::with_capacity(shape.total_bytes()); + + let router_start = bytes.len(); + extend(&mut bytes, shape.population * shape.hidden, SEED_ROUTER); + let router = router_start..bytes.len(); + + let experts = (0..shape.population) + .map(|e| { + let gate_up_start = bytes.len(); + // Gate rows first, then up rows — concatenated, not + // interleaved, matching what the reader expects. + extend(&mut bytes, shape.intermediate * shape.hidden, SEED_GATE + e); + extend(&mut bytes, shape.intermediate * shape.hidden, SEED_UP + e); + let gate_up = gate_up_start..bytes.len(); + + let down_start = bytes.len(); + extend(&mut bytes, shape.down_elements(), SEED_DOWN + e); + ExpertExtents { + gate_up, + down: down_start..bytes.len(), + } + }) + .collect(); + + Self { + shape, + bytes, + router, + experts, + } + } + + /// A residual of the right width for this layer. + pub fn residual(&self) -> Vec { + (0..self.shape.hidden) + .map(|i| weight(SEED_RESIDUAL, i)) + .collect() + } + + /// Bind an operation over `image`, which must hold this layer's bytes. + /// + /// Separate from `self.bytes` so the same layout can be bound over a file + /// mapping. The residency probe relies on measuring exactly the layout the + /// benches measure. + pub fn bind<'a>(&self, image: &'a [u8]) -> BoundMoeOperation<'a> { + let shape = self.shape; + let experts = self + .experts + .iter() + .enumerate() + .map(|(e, extents)| BoundExpert { + expert_id: e as u32, + projection: BoundProjection::Fused { + gate_up: tensor( + &RegionRole::GateUpFused.name(), + &image[extents.gate_up.clone()], + FUSED_PROJECTION_HALVES * shape.intermediate, + shape.hidden, + ), + }, + down: tensor( + &RegionRole::Down.name(), + &image[extents.down.clone()], + shape.hidden, + shape.intermediate, + ), + }) + .collect(); + + BoundMoeOperation { + router: BoundRouter { + weight: tensor( + ROUTER_REGION_SET, + &image[self.router.clone()], + shape.population, + shape.hidden, + ), + top_k: shape.top_k, + selected_weight: MoeTopKWeightPolicy::RenormalizedSoftmax, + scaling: BoundExpertScaling::None, + kernel: RouterKernel::default(), + }, + transforms: Vec::new(), + banks: vec![BoundBankOperation { + bank: BankCoordinate::new(0, 0), + experts, + intermediate_dim: shape.intermediate, + hidden_dim: shape.hidden, + activation: Activation::Silu, + kernel: ExpertKernel::default(), + }], + reduction: BoundReduction::WeightedSum, + residual_dim: shape.hidden, + } + } +} + +fn extend(out: &mut Vec, count: usize, seed: usize) { + out.extend((0..count).flat_map(|i| weight(seed, i).to_le_bytes())); +} + +fn tensor<'a>(region_set: &str, data: &'a [u8], rows: usize, cols: usize) -> BoundTensor<'a> { + BoundTensor::direct( + RepresentationIdentity::new(region_set, VARIANT), + data, + RegionFormat::F32, + ComponentContract::matrix(rows as u32, cols as u32), + ) + .expect("synthetic operands are well-formed") +} diff --git a/crates/larql-vindex/src/runtime/fixtures/tests/direct_moe.rs b/crates/larql-vindex/src/runtime/fixtures/tests/direct_moe.rs new file mode 100644 index 000000000..1431b0019 --- /dev/null +++ b/crates/larql-vindex/src/runtime/fixtures/tests/direct_moe.rs @@ -0,0 +1,327 @@ +//! Fixture A — one direct routed-MoE layer against an independent oracle. +//! +//! Checkpoint by checkpoint, not final-vector-only. A single equality +//! assertion on the output says an execution is wrong and stops there; these +//! say *where*: +//! +//! ```text +//! router_scores scoring, or the router operand +//! selection top-k depth, ordering, or a weight policy +//! expert_outputs region interpretation, activation, or the gated product +//! reduced the combine +//! residual_delta the transform stages +//! ``` + +use crate::runtime::execute::{execute, execute_traced}; +use crate::runtime::fixtures::direct_moe::{ + input, up_at, DirectMoeFixture, HIDDEN, INTERMEDIATE, POPULATION, TOP_K, +}; +use crate::runtime::fixtures::direct_moe_oracle::oracle; +use crate::runtime::inputs::MoeInputs; +use crate::runtime::projection::{BoundProjection, ProjectionArrangement}; + +/// Reference and oracle run the same arithmetic in the same order, so they +/// agree to well within this. It is a float-comparison guard, not a tolerance +/// budget — loosening it to make a test pass would be hiding a real defect. +const EPSILON: f32 = 1e-6; + +fn assert_close(actual: &[f32], expected: &[f32], what: &str) { + assert_eq!(actual.len(), expected.len(), "{what}: length"); + for (i, (a, e)) in actual.iter().zip(expected).enumerate() { + assert!((a - e).abs() < EPSILON, "{what}[{i}]: {a} vs expected {e}"); + } +} + +// ── The bound operation is well-formed ───────────────────────────────────── + +#[test] +fn both_arrangements_validate() { + let fixture = DirectMoeFixture::new(); + for arrangement in ProjectionArrangement::ALL { + let op = fixture.operation(arrangement); + op.validate() + .unwrap_or_else(|e| panic!("{}: {e}", arrangement.name())); + } +} + +#[test] +fn the_operation_describes_itself_as_direct_not_latent() { + let fixture = DirectMoeFixture::new(); + let op = fixture.operation(ProjectionArrangement::Decomposed); + assert!(!op.is_latent()); + assert_eq!(op.bank_input_dim(), HIDDEN, "no routed-input transform"); +} + +// ── Checkpoint agreement with the oracle ─────────────────────────────────── + +#[test] +fn router_scores_match_the_oracle() { + let fixture = DirectMoeFixture::new(); + let (_, trace) = execute_traced( + &fixture.operation(ProjectionArrangement::Decomposed), + MoeInputs::shared(&input()), + ) + .unwrap(); + let expected = oracle(&input()); + assert_eq!(trace.router_scores.len(), POPULATION); + assert_close( + &trace.router_scores, + &expected.router_scores, + "router_scores", + ); +} + +#[test] +fn router_scores_are_a_distribution_over_the_whole_population() { + // Softmax runs before selection, so all four probabilities must be present + // and sum to one — a softmax applied after top-k would also "work" and + // would change every weight. + let fixture = DirectMoeFixture::new(); + let (_, trace) = execute_traced( + &fixture.operation(ProjectionArrangement::Decomposed), + MoeInputs::shared(&input()), + ) + .unwrap(); + let total: f32 = trace.router_scores.iter().sum(); + assert!((total - 1.0).abs() < EPSILON, "scores sum to {total}"); +} + +#[test] +fn selected_expert_ids_and_order_match_the_oracle() { + let fixture = DirectMoeFixture::new(); + let (_, trace) = execute_traced( + &fixture.operation(ProjectionArrangement::Decomposed), + MoeInputs::shared(&input()), + ) + .unwrap(); + let expected: Vec = oracle(&input()) + .selected + .iter() + .map(|(id, _)| *id) + .collect(); + assert_eq!(trace.selected_ids(), expected); + assert_eq!(trace.selection.len(), TOP_K); +} + +#[test] +fn selection_is_ordered_by_descending_score() { + // Order is part of the contract, not an accident of the sort: the + // reduction is commutative but the trace and any downstream consumer that + // takes "the top expert" are not. + let fixture = DirectMoeFixture::new(); + let (_, trace) = execute_traced( + &fixture.operation(ProjectionArrangement::Decomposed), + MoeInputs::shared(&input()), + ) + .unwrap(); + let raw: Vec = trace.selection.iter().map(|s| s.raw_score).collect(); + assert!(raw.windows(2).all(|w| w[0] >= w[1]), "{raw:?}"); +} + +#[test] +fn gate_weights_match_the_oracle() { + let fixture = DirectMoeFixture::new(); + let (_, trace) = execute_traced( + &fixture.operation(ProjectionArrangement::Decomposed), + MoeInputs::shared(&input()), + ) + .unwrap(); + let expected: Vec = oracle(&input()).selected.iter().map(|(_, w)| *w).collect(); + assert_close(&trace.gate_weights(), &expected, "gate_weights"); +} + +#[test] +fn renormalisation_makes_the_selected_weights_sum_to_one() { + let fixture = DirectMoeFixture::new(); + let (_, trace) = execute_traced( + &fixture.operation(ProjectionArrangement::Decomposed), + MoeInputs::shared(&input()), + ) + .unwrap(); + let total: f32 = trace.gate_weights().iter().sum(); + assert!((total - 1.0).abs() < EPSILON, "weights sum to {total}"); +} + +#[test] +fn raw_scores_are_kept_alongside_renormalised_weights() { + // Renormalisation can make two different routings produce the same final + // weights; the raw score is where that disagreement is still visible. + let fixture = DirectMoeFixture::new(); + let (_, trace) = execute_traced( + &fixture.operation(ProjectionArrangement::Decomposed), + MoeInputs::shared(&input()), + ) + .unwrap(); + let raw_total: f32 = trace.selection.iter().map(|s| s.raw_score).sum(); + assert!( + raw_total < 1.0, + "top-{TOP_K} of {POPULATION} is not all of it" + ); +} + +#[test] +fn per_expert_outputs_match_the_oracle() { + let fixture = DirectMoeFixture::new(); + let (_, trace) = execute_traced( + &fixture.operation(ProjectionArrangement::Decomposed), + MoeInputs::shared(&input()), + ) + .unwrap(); + let expected = oracle(&input()); + assert_eq!(trace.expert_outputs.len(), TOP_K); + for (id, values) in &expected.expert_outputs { + let actual = trace + .expert_output(*id) + .unwrap_or_else(|| panic!("expert {id} was not recorded")); + assert_eq!(actual.len(), HIDDEN); + assert_close(actual, values, &format!("expert {id} output")); + } +} + +#[test] +fn the_weighted_reduction_matches_the_oracle() { + let fixture = DirectMoeFixture::new(); + let (_, trace) = execute_traced( + &fixture.operation(ProjectionArrangement::Decomposed), + MoeInputs::shared(&input()), + ) + .unwrap(); + assert_close(&trace.reduced, &oracle(&input()).reduced, "reduced"); +} + +#[test] +fn the_residual_delta_matches_the_oracle() { + let fixture = DirectMoeFixture::new(); + let output = execute( + &fixture.operation(ProjectionArrangement::Decomposed), + MoeInputs::shared(&input()), + ) + .unwrap(); + assert_close(&output, &oracle(&input()).reduced, "residual delta"); +} + +#[test] +fn a_direct_moe_delta_equals_its_reduction() { + // No transforms bound, so the two checkpoints must coincide. If they ever + // differ, a stage is being applied that nothing bound. + let fixture = DirectMoeFixture::new(); + let (output, trace) = execute_traced( + &fixture.operation(ProjectionArrangement::Decomposed), + MoeInputs::shared(&input()), + ) + .unwrap(); + assert_eq!(trace.reduced, trace.residual_delta); + assert_eq!(output, trace.residual_delta); +} + +// ── The recipe, not the arrangement ──────────────────────────────────────── + +#[test] +fn fused_and_decomposed_storage_produce_identical_output() { + // The headline property: same semantics, two physical layouts, one + // executor, byte-identical result. + let fixture = DirectMoeFixture::new(); + let decomposed = execute( + &fixture.operation(ProjectionArrangement::Decomposed), + MoeInputs::shared(&input()), + ) + .unwrap(); + let fused = execute( + &fixture.operation(ProjectionArrangement::Fused), + MoeInputs::shared(&input()), + ) + .unwrap(); + assert_eq!(decomposed, fused, "arrangement changed the answer"); +} + +#[test] +fn fused_and_decomposed_agree_at_every_checkpoint() { + // Equal outputs could in principle survive compensating errors upstream. + let fixture = DirectMoeFixture::new(); + let (_, a) = execute_traced( + &fixture.operation(ProjectionArrangement::Decomposed), + MoeInputs::shared(&input()), + ) + .unwrap(); + let (_, b) = execute_traced( + &fixture.operation(ProjectionArrangement::Fused), + MoeInputs::shared(&input()), + ) + .unwrap(); + assert_eq!(a, b); +} + +#[test] +fn the_fused_up_half_is_the_second_block_not_every_other_row() { + // The discriminating test for the concatenated-vs-interleaved bug. Reading + // the fused region interleaved would return unit 0's up row where unit 1's + // belongs — same shape, wrong values, no error anywhere downstream. + let fixture = DirectMoeFixture::new(); + let op = fixture.operation(ProjectionArrangement::Fused); + let expert = 2usize; + let unit = 1usize; + let projection = &op.banks[0].experts[expert].projection; + assert!(matches!(projection, BoundProjection::Fused { .. })); + + let mut row = vec![0.0f32; HIDDEN]; + projection.up_row_into(unit, &mut row).unwrap(); + + let correct: Vec = (0..HIDDEN).map(|h| up_at(expert, unit, h)).collect(); + let interleaved: Vec = (0..HIDDEN).map(|h| up_at(expert, 0, h)).collect(); + assert_close(&row, &correct, "fused up row"); + assert_ne!( + correct, interleaved, + "the fixture must discriminate the bug" + ); +} + +#[test] +fn both_arrangements_report_the_same_dimensions() { + let fixture = DirectMoeFixture::new(); + for arrangement in ProjectionArrangement::ALL { + let op = fixture.operation(arrangement); + let projection = &op.banks[0].experts[0].projection; + assert_eq!( + projection.intermediate_dim(), + INTERMEDIATE, + "{}", + arrangement.name() + ); + assert_eq!(projection.hidden_dim(), HIDDEN, "{}", arrangement.name()); + } +} + +// ── The traced and untraced paths are one implementation ─────────────────── + +#[test] +fn tracing_does_not_change_the_result() { + let fixture = DirectMoeFixture::new(); + let op = fixture.operation(ProjectionArrangement::Decomposed); + let plain = execute(&op, MoeInputs::shared(&input())).unwrap(); + let (traced, _) = execute_traced(&op, MoeInputs::shared(&input())).unwrap(); + assert_eq!(plain, traced); +} + +#[test] +fn execution_is_deterministic_across_runs() { + let fixture = DirectMoeFixture::new(); + let op = fixture.operation(ProjectionArrangement::Fused); + let first = execute(&op, MoeInputs::shared(&input())).unwrap(); + let second = execute(&op, MoeInputs::shared(&input())).unwrap(); + assert_eq!(first, second); +} + +// ── Refusals ─────────────────────────────────────────────────────────────── + +#[test] +fn a_residual_of_the_wrong_width_is_refused_naming_both_widths() { + let fixture = DirectMoeFixture::new(); + let err = execute( + &fixture.operation(ProjectionArrangement::Decomposed), + MoeInputs::shared(&[1.0, 2.0]), + ) + .unwrap_err(); + let text = err.to_string(); + assert!(text.contains("residual"), "{text}"); + assert!(text.contains(&HIDDEN.to_string()), "{text}"); +} diff --git a/crates/larql-vindex/src/runtime/fixtures/tests/mod.rs b/crates/larql-vindex/src/runtime/fixtures/tests/mod.rs new file mode 100644 index 000000000..e08f4e74c --- /dev/null +++ b/crates/larql-vindex/src/runtime/fixtures/tests/mod.rs @@ -0,0 +1,4 @@ +//! Tests for the shared fixtures, one file per fixture. + +mod direct_moe; +mod synthetic; diff --git a/crates/larql-vindex/src/runtime/fixtures/tests/synthetic.rs b/crates/larql-vindex/src/runtime/fixtures/tests/synthetic.rs new file mode 100644 index 000000000..77d5fdfa2 --- /dev/null +++ b/crates/larql-vindex/src/runtime/fixtures/tests/synthetic.rs @@ -0,0 +1,170 @@ +//! Colocated tests for the synthetic layer builder. +//! +//! The builder feeds a bench, an allocation guard and a residency probe. If it +//! is wrong they all measure something other than what they claim, and none of +//! them would notice — a bench does not check its own answers. + +use crate::runtime::execute::{execute, execute_traced}; +use crate::runtime::fixtures::synthetic::{SyntheticLayer, SyntheticShape}; +use crate::runtime::inputs::MoeInputs; + +fn shape() -> SyntheticShape { + SyntheticShape { + population: 5, + hidden: 4, + intermediate: 3, + top_k: 2, + } +} + +// ── The byte image ───────────────────────────────────────────────────────── + +#[test] +fn the_image_is_exactly_the_size_the_shape_predicts() { + let layer = SyntheticLayer::build(shape()); + assert_eq!(layer.bytes.len(), shape().total_bytes()); +} + +#[test] +fn every_extent_lies_inside_the_image() { + let layer = SyntheticLayer::build(shape()); + assert!(layer.router.end <= layer.bytes.len()); + for extents in &layer.experts { + assert!(extents.gate_up.end <= layer.bytes.len()); + assert!(extents.down.end <= layer.bytes.len()); + } +} + +#[test] +fn extents_tile_the_image_without_gaps_or_overlap() { + // A gap would waste pages the residency probe then attributes to routing; + // an overlap would make two experts share weights. + let layer = SyntheticLayer::build(shape()); + let mut cursor = layer.router.end; + assert_eq!(layer.router.start, 0); + for extents in &layer.experts { + assert_eq!(extents.gate_up.start, cursor); + assert_eq!(extents.down.start, extents.gate_up.end); + cursor = extents.down.end; + } + assert_eq!(cursor, layer.bytes.len()); +} + +#[test] +fn each_experts_regions_are_adjacent() { + // Expert-major locality is the property the residency probe measures. + let layer = SyntheticLayer::build(shape()); + for extents in &layer.experts { + assert_eq!(extents.gate_up.end, extents.down.start); + } +} + +#[test] +fn the_population_gets_one_set_of_extents_each() { + assert_eq!( + SyntheticLayer::build(shape()).experts.len(), + shape().population + ); +} + +// ── Shape arithmetic ─────────────────────────────────────────────────────── + +#[test] +fn a_fused_projection_is_two_halves_wide() { + let s = shape(); + assert_eq!(s.projection_elements(), 2 * s.intermediate * s.hidden); +} + +#[test] +fn selected_bytes_counts_the_router_plus_top_k_experts() { + let s = shape(); + assert_eq!( + s.selected_bytes(), + s.router_bytes() + s.top_k * s.expert_bytes() + ); + assert!( + s.selected_bytes() < s.total_bytes(), + "routing must be sparse" + ); +} + +#[test] +fn selected_bytes_never_exceeds_the_population() { + let s = SyntheticShape { + top_k: 99, + ..shape() + }; + assert_eq!(s.selected_bytes(), s.total_bytes()); +} + +// ── Binding and execution ────────────────────────────────────────────────── + +#[test] +fn a_synthetic_layer_binds_to_a_valid_operation() { + let layer = SyntheticLayer::build(shape()); + layer.bind(&layer.bytes).validate().unwrap(); +} + +#[test] +fn a_synthetic_layer_executes_to_finite_values() { + let layer = SyntheticLayer::build(shape()); + let out = execute( + &layer.bind(&layer.bytes), + MoeInputs::shared(&layer.residual()), + ) + .unwrap(); + assert_eq!(out.len(), shape().hidden); + assert!(out.iter().all(|v| v.is_finite()), "{out:?}"); + assert!( + out.iter().any(|v| *v != 0.0), + "a zero delta would be vacuous" + ); +} + +#[test] +fn binding_over_a_copy_of_the_image_gives_the_same_answer() { + // The property the residency probe depends on: identical content bound + // over different backing memory executes identically. + let layer = SyntheticLayer::build(shape()); + let copy = layer.bytes.clone(); + assert_eq!( + execute( + &layer.bind(&layer.bytes), + MoeInputs::shared(&layer.residual()) + ) + .unwrap(), + execute(&layer.bind(©), MoeInputs::shared(&layer.residual())).unwrap() + ); +} + +#[test] +fn routing_selects_exactly_top_k_experts() { + let layer = SyntheticLayer::build(shape()); + let (_, trace) = execute_traced( + &layer.bind(&layer.bytes), + MoeInputs::shared(&layer.residual()), + ) + .unwrap(); + assert_eq!(trace.selection.len(), shape().top_k); + assert_eq!(trace.router_scores.len(), shape().population); +} + +#[test] +fn different_populations_produce_different_images() { + let small = SyntheticLayer::build(shape()); + let large = SyntheticLayer::build(SyntheticShape { + population: 9, + ..shape() + }); + assert!(large.bytes.len() > small.bytes.len()); + assert_eq!(large.experts.len(), 9); +} + +#[test] +fn the_weight_generator_is_deterministic_and_varies_by_seed() { + use crate::runtime::fixtures::synthetic::weight; + assert_eq!(weight(1, 7), weight(1, 7)); + assert_ne!(weight(1, 7), weight(2, 7)); + assert_ne!(weight(1, 7), weight(1, 8)); + assert!((-0.5..0.5).contains(&weight(3, 11))); +} diff --git a/crates/larql-vindex/src/runtime/inputs.rs b/crates/larql-vindex/src/runtime/inputs.rs new file mode 100644 index 000000000..0e8ece28a --- /dev/null +++ b/crates/larql-vindex/src/runtime/inputs.rs @@ -0,0 +1,80 @@ +//! What an execution is given. +//! +//! # Gemma forced this +//! +//! The first draft passed one vector to both the router and the experts, +//! because fixture A and every synthetic case share one input. Gemma does not: +//! `moe_router_input` applies a router-specific RMS norm, a learned +//! element-wise scale and a scalar multiplier on top of the vector the experts +//! receive. Routing therefore scores a *different vector* from the one the +//! selected experts consume. +//! +//! ```text +//! residual +//! → pre_experts_norm ─────────────→ expert input +//! └→ router_norm × router_scale × scalar → router input +//! ``` +//! +//! A single-input model cannot express that. It would have had to score on the +//! expert input, which selects different experts — a wrong answer with no +//! shape error anywhere, exactly the class of fault this programme keeps +//! finding. +//! +//! # Why both are supplied rather than derived +//! +//! Norms, scales and scalars belong to the surrounding block, not to the MoE +//! operation — the same reason `execute` returns a delta instead of an updated +//! residual. Deriving the router input here would mean re-modelling the +//! incumbent's routing-policy enums inside VINDEX3, giving two descriptions of +//! one behaviour and somewhere for them to disagree. +//! +//! So the block computes both vectors, as it already does, and hands them +//! over. Models where the two coincide use [`MoeInputs::shared`]. + +/// The vectors one MoE execution reads. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct MoeInputs<'a> { + /// The vector the experts consume. Routed-input transforms apply to this. + pub bank: &'a [f32], + /// The vector the router scores, in the space the router expects. + /// + /// Equal to `bank` whenever a model routes on the same vector its experts + /// see. Gemma does not; Mixtral-style models do. + pub router: &'a [f32], +} + +impl<'a> MoeInputs<'a> { + /// One vector for both — the common case, and fixture A's. + pub fn shared(input: &'a [f32]) -> Self { + Self { + bank: input, + router: input, + } + } + + /// Distinct vectors, as Gemma's hybrid MoE requires. + pub fn split(bank: &'a [f32], router: &'a [f32]) -> Self { + Self { bank, router } + } + + /// Whether routing and expert compute read the same vector. + /// + /// Not a pointer comparison: two separately-computed vectors that happen to + /// be equal *are* shared for every purpose that matters here, and a model + /// whose router norm is the identity produces exactly that. + pub fn is_shared(&self) -> bool { + self.bank == self.router + } + + pub fn describe(&self) -> String { + if self.is_shared() { + format!("shared input, width {}", self.bank.len()) + } else { + format!( + "split inputs, bank width {} / router width {}", + self.bank.len(), + self.router.len() + ) + } + } +} diff --git a/crates/larql-vindex/src/runtime/kernels.rs b/crates/larql-vindex/src/runtime/kernels.rs new file mode 100644 index 000000000..8e60cef91 --- /dev/null +++ b/crates/larql-vindex/src/runtime/kernels.rs @@ -0,0 +1,132 @@ +//! Reference numeric primitives. +//! +//! Plain, allocating, obviously-correct implementations. They exist to say +//! what the answer should be, so that a fast kernel has something to be +//! checked against. Nothing here should be optimised: the moment this file +//! becomes clever it stops being usable as an oracle. +//! +//! # Tie-breaking is a contract here, not an accident +//! +//! The incumbent MoE path selects top-k with `sort_unstable_by` over +//! `partial_cmp`. That is **unspecified for ties** — not necessarily random, +//! and it may well reproduce the same order within a build, but it promises +//! nothing. A reference that offers no ordering guarantee cannot serve as an +//! oracle, so this one commits: +//! +//! ```text +//! score descending, then expert id ascending +//! ``` +//! +//! The comparison uses `f32::total_cmp`, a genuine total order, rather than +//! `partial_cmp` with a fallback. A fallback to `Equal` hands incomparable +//! values whatever position the sort happens to leave them in, which is +//! precisely the unspecified behaviour being replaced. +//! +//! Non-finite scores are refused upstream in `execute` rather than ordered: +//! a NaN that acquires a position by sort mechanics is a routing decision +//! nobody made. + +use larql_compute::Activation; + +/// `out[r] = dot(matrix_row(r), x)`, with rows supplied by the caller. +/// +/// Taking a row-reader rather than a slice is what lets one implementation +/// serve fused and decomposed storage, quantised and dense, transposed and +/// direct — the arrangement is resolved by the reader, not here. +pub fn dot(a: &[f32], b: &[f32]) -> f32 { + a.iter().zip(b).map(|(x, y)| x * y).sum() +} + +/// In-place softmax, max-shifted for stability. +pub fn softmax(v: &mut [f32]) { + let max = v.iter().copied().fold(f32::NEG_INFINITY, f32::max); + let mut sum = 0.0f32; + for x in v.iter_mut() { + *x = (*x - max).exp(); + sum += *x; + } + if sum > 0.0 { + for x in v.iter_mut() { + *x /= sum; + } + } +} + +/// The `k` highest scores, descending, ties resolved to the lower index. +pub fn top_k(scores: &[f32], k: usize) -> Vec<(usize, f32)> { + top_k_with_margin(scores, k).0 +} + +/// Top-k plus the **boundary margin**: `score[k-1] - score[k]`. +/// +/// The margin is how much room the selection had. It is `None` when there is +/// no boundary — `k` is zero, or `k` covers the whole population. +/// +/// It exists for parity triage. When VINDEX2 and VINDEX3 disagree about which +/// experts ran, a margin of exactly zero says the two implementations met an +/// exact tie and resolved it under different policies; a wide margin says they +/// disagreed about the *scores*, which is a weight-decoding or arithmetic +/// fault. Those need completely different investigations, and without the +/// margin the two are indistinguishable from the outside. +pub fn top_k_with_margin(scores: &[f32], k: usize) -> (Vec<(usize, f32)>, Option) { + let mut indexed: Vec<(usize, f32)> = scores.iter().copied().enumerate().collect(); + // `total_cmp` rather than `partial_cmp(..).unwrap_or(Equal)`: a total + // order, so the result does not depend on sort mechanics for any input. + indexed.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0))); + + let k = k.min(scores.len()); + let margin = (k > 0 && k < indexed.len()).then(|| indexed[k - 1].1 - indexed[k].1); + indexed.truncate(k); + (indexed, margin) +} + +/// Rescale so the selected weights sum to one. A zero sum is left alone. +pub fn renormalize(weights: &mut [f32]) { + let sum: f32 = weights.iter().sum(); + if sum > 0.0 { + for w in weights.iter_mut() { + *w /= sum; + } + } +} + +/// Apply an activation elementwise, sharing `larql-compute`'s vocabulary. +pub fn activate(activation: Activation, v: &mut [f32]) { + for x in v.iter_mut() { + *x = apply(activation, *x); + } +} + +fn apply(activation: Activation, x: f32) -> f32 { + match activation { + Activation::Silu => x / (1.0 + (-x).exp()), + Activation::GeluTanh => { + const SQRT_2_OVER_PI: f32 = 0.797_884_6; + const CUBIC_COEFF: f32 = 0.044_715; + let inner = SQRT_2_OVER_PI * (x + CUBIC_COEFF * x * x * x); + 0.5 * x * (1.0 + inner.tanh()) + } + Activation::GeluExact => { + const SQRT_2: f32 = std::f32::consts::SQRT_2; + 0.5 * x * (1.0 + erf(x / SQRT_2)) + } + Activation::ReLU => x.max(0.0), + } +} + +/// Abramowitz–Stegun 7.1.26. Reference-grade, not bit-exact with libm. +fn erf(x: f32) -> f32 { + const A: [f32; 5] = [ + 0.254_829_6, + -0.284_496_74, + 1.421_413_7, + -1.453_152, + 1.061_405_4, + ]; + const P: f32 = 0.327_591_1; + let sign = if x < 0.0 { -1.0 } else { 1.0 }; + let x = x.abs(); + let t = 1.0 / (1.0 + P * x); + let poly = A.iter().rev().fold(0.0f32, |acc, &coeff| (acc + coeff) * t); + sign * (1.0 - poly * (-x * x).exp()) +} diff --git a/crates/larql-vindex/src/runtime/mod.rs b/crates/larql-vindex/src/runtime/mod.rs new file mode 100644 index 000000000..95a036437 --- /dev/null +++ b/crates/larql-vindex/src/runtime/mod.rs @@ -0,0 +1,74 @@ +//! The VINDEX3 reference runtime — executing a bound route. +//! +//! # Where the seam is +//! +//! Not at the model. LARQL's generation loop, attention, KV handling, +//! tokenizer and sampler are unchanged and shared; what VINDEX3 replaces is +//! the FFN/MoE execution interface underneath them. +//! +//! ```text +//! existing generation loop +//! ├── existing attention / KV path +//! ├── existing residual handling +//! └── FFN/MoE execution +//! ├── VINDEX2 implementation (cpu_moe_forward over MoeLayerWeights) +//! └── VINDEX3 bound plan (this module) +//! ``` +//! +//! # Complex load path, boring token path +//! +//! Everything in here is post-decision. Roles have been matched to a +//! programme, alternatives chosen, variants resolved, contracts checked and +//! authority folded — none of which is reachable from [`execute`]. The +//! executor receives operands and shapes and does arithmetic. +//! +//! The division is load-bearing rather than aesthetic. Resolution machinery +//! that remains reachable from the execution object ends up called from the +//! execution path, and per-token cost starts tracking catalogue size. +//! +//! # Reference, not fast +//! +//! Decoding dispatches per element and every operand is read row by row. This +//! is the implementation that says what the answer should be; production +//! routes bind to `larql-compute` kernels and are checked against it. + +pub mod addressing; +pub mod axis; +pub mod bank; +pub mod consts; +pub mod error; +pub mod execute; +pub mod expert; +pub mod expert_kernel; +pub mod fixtures; +pub mod inputs; +pub mod kernels; +pub mod operation; +pub mod projection; +pub mod reduction; +pub mod residency; +pub mod router; +pub mod tensor; +pub mod trace; +pub mod transform; +pub mod verdict; + +#[cfg(test)] +mod tests; + +pub use addressing::{Addressing, BlockOperand}; +pub use axis::Axis; +pub use bank::{BoundBankOperation, BoundExpert}; +pub use error::{ExecutionError, OperandUnsuitability}; +pub use execute::{execute, execute_traced, execute_with}; +pub use expert_kernel::ExpertKernel; +pub use inputs::MoeInputs; +pub use operation::BoundMoeOperation; +pub use projection::{BoundProjection, ProjectionArrangement}; +pub use reduction::BoundReduction; +pub use residency::{account, ExpertRegion, PageSpan, ResidencyAccount}; +pub use router::{BoundExpertScaling, BoundRouter, RouterKernel, SelectedExpert}; +pub use tensor::BoundTensor; +pub use trace::{CollectedTrace, ExpertOutput, NoTrace, TraceSink}; +pub use transform::{BoundTransform, TransformStage}; +pub use verdict::{RefusalKind, Verdict}; diff --git a/crates/larql-vindex/src/runtime/operation.rs b/crates/larql-vindex/src/runtime/operation.rs new file mode 100644 index 000000000..b3a7d50d5 --- /dev/null +++ b/crates/larql-vindex/src/runtime/operation.rs @@ -0,0 +1,133 @@ +//! The bound MoE operation — the object the token path executes. +//! +//! # What is deliberately absent +//! +//! No manifest, no profile, no capability traversal, no authority fold, no +//! variant candidates, no kernel registry. Every one of those questions was +//! answered during binding, and none of them is reachable from here. +//! +//! That is not tidiness. Resolution work that stays reachable from the +//! execution object gets called from the execution path eventually — usually +//! as a cache lookup that looks harmless — and then per-token cost depends on +//! catalogue size. Making the object incapable of resolution is what keeps the +//! load path complex and the token path boring. +//! +//! # Shape +//! +//! ```text +//! router scores the population and selects +//! transforms optional latent projections either side of the banks +//! banks the expert populations +//! reduction how selected outputs combine +//! ``` + +use super::axis::Axis; +use super::bank::BoundBankOperation; +use super::error::ExecutionError; +use super::reduction::BoundReduction; +use super::router::BoundRouter; +use super::transform::{BoundTransform, TransformStage}; + +/// An immutable, fully-bound routed MoE operation for one layer. +#[derive(Debug, Clone, PartialEq)] +pub struct BoundMoeOperation<'a> { + pub router: BoundRouter<'a>, + /// Latent projections. Empty for a direct MoE, where the bank reads and + /// writes the residual width directly. + pub transforms: Vec>, + /// Expert populations the router selects into. + /// + /// A Vec because K3-shaped layers pair a routed bank with a shared one. + /// Today every bank here is routed by `router` and their reduced outputs + /// sum; unrouted shared banks arrive with the Mini-K3 rung. + pub banks: Vec>, + pub reduction: BoundReduction, + /// Residual width — what this operation reads and what it returns. + pub residual_dim: usize, +} + +impl BoundMoeOperation<'_> { + /// Width the banks operate at: the latent width when a routed-input + /// transform is bound, the residual width otherwise. + pub fn bank_input_dim(&self) -> usize { + // The *last* routed-input transform, since a stage may chain several + // and it is the final one that fixes the width the banks see. + self.transforms + .iter() + .rfind(|t| t.stage == TransformStage::RoutedInput) + .map_or(self.residual_dim, |t| t.out_dim()) + } + + /// Whether this operation projects into a separate routed space. + pub fn is_latent(&self) -> bool { + !self.transforms.is_empty() + } + + /// Check every operand against the shape the others imply. + /// + /// Belongs to the load path. Binding calls it once; execution does not, + /// because re-validating per token is exactly the resolution creep this + /// object exists to prevent. + /// + /// # A bank need not hold the whole population + /// + /// An earlier version required `router.population() == bank.population()`. + /// That is wrong, and `bank.rs` already said so: a **sharded** bank holds a + /// subset, and expert 40 may be the first one a shard carries. Requiring + /// equality would have refused every sharded operation — including the + /// expert-server slices this format exists to serve. + /// + /// So the router is validated against its own addressable population, and + /// the bank is checked to lie inside it. An expert id the router cannot + /// address is the genuine fault, because nothing would ever select it. + pub fn validate(&self) -> Result<(), ExecutionError> { + let bank_input = self.bank_input_dim(); + self.router.validate(self.router.population(), bank_input)?; + for bank in &self.banks { + if bank.hidden_dim != bank_input { + return Err(ExecutionError::DimensionMismatch { + operand: bank.describe(), + axis: Axis::InputWidth, + expected: bank_input, + found: bank.hidden_dim, + }); + } + bank.validate()?; + for expert in &bank.experts { + if expert.expert_id as usize >= self.router.population() { + return Err(ExecutionError::ExpertOutOfRange { + expert: expert.expert_id, + population: self.router.population(), + }); + } + } + } + Ok(()) + } + + /// Whether the bound banks hold the router's whole addressable population. + /// + /// False for a shard. Distinct from validity: a shard is a legitimate + /// operation, it simply cannot serve every routing outcome, and a caller + /// that needs completeness must ask rather than assume. + pub fn holds_full_population(&self) -> bool { + self.banks + .iter() + .any(|b| b.population() == self.router.population()) + } + + pub fn describe(&self) -> String { + let banks: Vec = self + .banks + .iter() + .map(BoundBankOperation::describe) + .collect(); + format!( + "{} → [{}] → {} ({})", + self.router.describe(), + banks.join("; "), + self.reduction.name(), + if self.is_latent() { "latent" } else { "direct" } + ) + } +} diff --git a/crates/larql-vindex/src/runtime/projection.rs b/crates/larql-vindex/src/runtime/projection.rs new file mode 100644 index 000000000..9fe5c7b56 --- /dev/null +++ b/crates/larql-vindex/src/runtime/projection.rs @@ -0,0 +1,144 @@ +//! The gate/up projection of one expert, in whichever arrangement it was stored. +//! +//! This is where "the executor follows the bound recipe" is actually cashed +//! out. A gated MLP needs a gate row and an up row per intermediate unit; +//! whether those live in two regions or in two halves of one is a storage +//! decision, and it is resolved here, once, behind an interface that yields +//! the same two rows either way. +//! +//! Everything downstream — activation, the elementwise product, the down +//! projection, the reduction — is then written once and cannot acquire a +//! layout assumption, because it never sees the layout. +//! +//! ```text +//! Decomposed gate[i] up[i] two regions, row i of each +//! Fused gate_up[i] gate_up[N + i] one region, halves in sequence +//! ``` +//! +//! The fused halves are **concatenated, not interleaved**. Reading them +//! interleaved yields a correctly-shaped tensor of wrong values — a failure +//! that survives all the way to the logits, which is why the arrangement is a +//! bound property rather than something inferred at execution time. + +use crate::format::lyrw2::region_role::RegionRole; + +use super::consts::FUSED_PROJECTION_HALVES; +use super::error::ExecutionError; +use super::tensor::BoundTensor; + +/// Which physical arrangement a projection was stored in. +/// +/// Defined once, here, and used by the executor, its diagnostics and the +/// fixtures alike. The role names come from the §6.5 registry rather than +/// being spelled again: `gate`/`up` versus `gate_up_fused` is the registry's +/// vocabulary, and a second copy of it drifts. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub enum ProjectionArrangement { + /// Separate `gate` and `up` regions. + Decomposed, + /// One concatenated `gate_up_fused` region. + Fused, +} + +impl ProjectionArrangement { + pub const ALL: [Self; 2] = [Self::Decomposed, Self::Fused]; + + /// The roles this arrangement stores, in region order. + pub fn roles(self) -> Vec { + match self { + Self::Decomposed => vec![RegionRole::Gate, RegionRole::Up], + Self::Fused => vec![RegionRole::GateUpFused], + } + } + + /// Named from the roles it stores, so the registry stays the one source. + pub fn name(self) -> String { + self.roles() + .iter() + .map(|r| r.name()) + .collect::>() + .join("+") + } +} + +/// How one expert's gate and up weights are arranged in storage. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum BoundProjection<'a> { + /// Separate `gate` and `up` regions, each `[intermediate, hidden]`. + Decomposed { + gate: BoundTensor<'a>, + up: BoundTensor<'a>, + }, + /// One `[2 * intermediate, hidden]` region: all gate rows, then all up rows. + Fused { gate_up: BoundTensor<'a> }, +} + +impl BoundProjection<'_> { + /// Intermediate units this projection produces. + pub fn intermediate_dim(&self) -> usize { + match self { + Self::Decomposed { gate, .. } => gate.rows(), + Self::Fused { gate_up } => gate_up.rows() / FUSED_PROJECTION_HALVES, + } + } + + /// Input width this projection contracts over. + pub fn hidden_dim(&self) -> usize { + match self { + Self::Decomposed { gate, .. } => gate.cols(), + Self::Fused { gate_up } => gate_up.cols(), + } + } + + /// Check the arrangement is internally consistent and matches the layer. + /// + /// A fused region with an odd row count is the clearest possible signal + /// that a decomposed region was bound as fused; catching it here turns a + /// silent half-wrong expert into a named refusal. + pub fn validate(&self, intermediate: usize, hidden: usize) -> Result<(), ExecutionError> { + match self { + Self::Decomposed { gate, up } => { + gate.require_matrix(intermediate, hidden)?; + up.require_matrix(intermediate, hidden) + } + Self::Fused { gate_up } => { + gate_up.require_matrix(intermediate * FUSED_PROJECTION_HALVES, hidden) + } + } + } + + /// Row `unit` of the gate projection. + pub fn gate_row_into(&self, unit: usize, out: &mut [f32]) -> Result<(), ExecutionError> { + match self { + Self::Decomposed { gate, .. } => gate.row_into(unit, out), + Self::Fused { gate_up } => gate_up.row_into(unit, out), + } + } + + /// Row `unit` of the up projection. + pub fn up_row_into(&self, unit: usize, out: &mut [f32]) -> Result<(), ExecutionError> { + match self { + Self::Decomposed { up, .. } => up.row_into(unit, out), + // The second half. Not `2 * unit + 1` — the halves are + // concatenated, and the interleaved reading is the classic + // silent-wrong-values bug this layout invites. + Self::Fused { gate_up } => gate_up.row_into(self.intermediate_dim() + unit, out), + } + } + + /// Which arrangement this is, for diagnostics and trace output. + pub const fn arrangement(&self) -> ProjectionArrangement { + match self { + Self::Decomposed { .. } => ProjectionArrangement::Decomposed, + Self::Fused { .. } => ProjectionArrangement::Fused, + } + } + + pub fn describe(&self) -> String { + let operands = match self { + Self::Decomposed { gate, up } => format!("{}, {}", gate.describe(), up.describe()), + Self::Fused { gate_up } => gate_up.describe(), + }; + format!("{}({operands})", self.arrangement().name()) + } +} diff --git a/crates/larql-vindex/src/runtime/reduction.rs b/crates/larql-vindex/src/runtime/reduction.rs new file mode 100644 index 000000000..8fd1e13a9 --- /dev/null +++ b/crates/larql-vindex/src/runtime/reduction.rs @@ -0,0 +1,50 @@ +//! How selected experts' outputs combine into one vector. +//! +//! Currently one kind, named rather than assumed. `WeightedSum` is what every +//! programme in the registry uses today, but writing it as `sum(w_i * y_i)` +//! inline would make the next reduction — a shared-expert bank that adds +//! unweighted, or a normalised combine — a change to the executor instead of a +//! change to the bound recipe. + +use super::axis::Axis; +use super::error::ExecutionError; + +/// The combining rule for a bank's selected expert outputs. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BoundReduction { + /// `sum_i weight_i * output_i` — the routed-MoE combine. + WeightedSum, +} + +impl BoundReduction { + pub const fn name(self) -> &'static str { + match self { + Self::WeightedSum => "weighted_sum", + } + } + + /// Accumulate one expert's contribution into `accumulator`. + pub fn accumulate( + self, + accumulator: &mut [f32], + output: &[f32], + weight: f32, + ) -> Result<(), ExecutionError> { + if accumulator.len() != output.len() { + return Err(ExecutionError::DimensionMismatch { + operand: format!("{} reduction", self.name()), + axis: Axis::OutputWidth, + expected: accumulator.len(), + found: output.len(), + }); + } + match self { + Self::WeightedSum => { + for (acc, &value) in accumulator.iter_mut().zip(output) { + *acc += weight * value; + } + } + } + Ok(()) + } +} diff --git a/crates/larql-vindex/src/runtime/residency.rs b/crates/larql-vindex/src/runtime/residency.rs new file mode 100644 index 000000000..2448eb972 --- /dev/null +++ b/crates/larql-vindex/src/runtime/residency.rs @@ -0,0 +1,213 @@ +//! Attributing resident pages to the operands that caused them. +//! +//! # The VINDEX3 memory question +//! +//! The incumbent cold-read probe asks whether sparse access to a large blob +//! pages acceptably. VINDEX3 makes a stronger claim available, because binding +//! happens before execution: +//! +//! > A bound operation knows its own footprint. The bytes it will touch are +//! > derivable from the plan, before a single page is faulted. +//! +//! That turns residency from something you observe afterwards into something +//! you can *predict and then check*. Which matters for placement, prefetch and +//! remote transfer: all three need the byte set in advance, and a predictor +//! that quietly disagrees with reality would mis-size every one of them. +//! +//! # Why page attribution, not a byte count +//! +//! Residency is quantised to pages, and expert regions do not begin on page +//! boundaries. An expert whose region starts mid-page shares that page with +//! its neighbour, so touching one expert makes part of another resident. That +//! is not a bug — it is the granularity — but it means "bytes the plan needs" +//! and "bytes that become resident" are different quantities, and comparing +//! them without accounting for the boundary makes a correct implementation +//! look wasteful. +//! +//! This module keeps the two apart: `predicted` is what the plan asks for, +//! `resident` is what the kernel actually holds, and the difference is +//! attributed rather than averaged away. +//! +//! # Zero overshoot is a reference-kernel result, not a universal rule +//! +//! The reference executor reads exactly the operands the plan names, one row +//! at a time, so its measured residency equals its prediction. That is a fact +//! about *this kernel*, and it must not harden into an acceptance criterion +//! for every kernel. +//! +//! A grouped or accelerated kernel may legitimately touch more: +//! +//! ```text +//! group-extent reads one storage group fetched for several experts +//! aligned over-read vectorised loads spanning a block boundary +//! separate scale blocks quantisation metadata outside the weight extent +//! kernel metadata headers, offset tables +//! prefetch speculative or asynchronous reads +//! ``` +//! +//! So three quantities are eventually distinct, and only the first two are +//! about correctness: +//! +//! ```text +//! required pages what the operation cannot run without +//! touch envelope what the bound kernel may legally read +//! prefetch pages what it may speculatively read ahead +//! ``` +//! +//! For the reference kernel the envelope coincides with the requirement. For a +//! grouped kernel the envelope is the union of the complete group extents +//! containing selected experts, and overshoot against the *requirement* +//! becomes intentional and predictable rather than a defect. +//! +//! The durable invariant, which survives both cases: +//! +//! > Observed pages must lie inside the bound kernel's declared touch +//! > envelope, and every required page must be covered. +//! +//! This module measures the reference case. Declaring envelopes belongs to +//! kernel binding and is deliberately not modelled here yet. + +use std::ops::Range; + +/// Page indices spanned by a byte range, inclusive of both ends. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PageSpan { + pub first: usize, + /// One past the last page — empty when `first == end`. + pub end: usize, +} + +impl PageSpan { + /// The pages a byte range occupies at `page_size`. + pub fn of(range: &Range, page_size: usize) -> Self { + if page_size == 0 || range.is_empty() { + return Self { first: 0, end: 0 }; + } + Self { + first: range.start / page_size, + // Round up: a range ending mid-page still occupies that page. + end: range.end.div_ceil(page_size), + } + } + + pub fn len(&self) -> usize { + self.end.saturating_sub(self.first) + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + pub fn contains(&self, page: usize) -> bool { + page >= self.first && page < self.end + } +} + +/// What a routed execution was expected to touch, and what it did touch. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct ResidencyAccount { + /// Pages resident anywhere in the mapping. + pub resident_total: usize, + /// Resident pages belonging to the router. + pub resident_router: usize, + /// Resident pages belonging to an expert the router selected. + pub resident_selected: usize, + /// Resident pages belonging only to experts that did **not** run. + /// + /// The number that matters. Routing is supposed to be sparse; if this + /// tracks the whole population, either readahead is defeating the sparsity + /// or the executor is reading operands it does not need. + /// + /// Zero is the *reference kernel's* result. A grouped kernel that fetches + /// whole storage groups will show non-zero overshoot legitimately — see + /// the touch-envelope note in the module docs before treating this as a + /// pass/fail gate. + pub resident_unselected: usize, + /// Pages the plan predicted it would need. + pub predicted: usize, +} + +impl ResidencyAccount { + /// Resident pages as a fraction of the whole mapping. + pub fn resident_fraction(&self, total_pages: usize) -> f64 { + ratio(self.resident_total, total_pages) + } + + /// How much of the prediction actually became resident. + /// + /// Below 1.0 is normal and healthy: a plan predicts every byte of every + /// operand, while execution may finish without faulting the last page of + /// a region it only partly reads. + pub fn prediction_coverage(&self) -> f64 { + ratio( + self.resident_router + self.resident_selected, + self.predicted, + ) + } + + /// Resident pages that no selected operand asked for. + pub fn overshoot_fraction(&self) -> f64 { + ratio(self.resident_unselected, self.resident_total) + } +} + +fn ratio(numerator: usize, denominator: usize) -> f64 { + if denominator == 0 { + return 0.0; + } + numerator as f64 / denominator as f64 +} + +/// One expert's byte extent, as the residency probe sees it. +#[derive(Debug, Clone)] +pub struct ExpertRegion { + pub expert_id: u32, + pub bytes: Range, +} + +/// Attribute a residency bitmap to router, selected experts and the rest. +/// +/// `resident[i]` is whether page `i` of the mapping is resident. A page is +/// counted once, and shared pages resolve in favour of *use*: router first, +/// then selected, then unselected. Counting a boundary page against the +/// unselected experts it partly covers would report overshoot for a page the +/// selected expert genuinely needed. +pub fn account( + resident: &[bool], + page_size: usize, + router: &Range, + experts: &[ExpertRegion], + selected: &[u32], + predicted_bytes: usize, +) -> ResidencyAccount { + let router_pages = PageSpan::of(router, page_size); + let selected_spans: Vec = experts + .iter() + .filter(|e| selected.contains(&e.expert_id)) + .map(|e| PageSpan::of(&e.bytes, page_size)) + .collect(); + let unselected_spans: Vec = experts + .iter() + .filter(|e| !selected.contains(&e.expert_id)) + .map(|e| PageSpan::of(&e.bytes, page_size)) + .collect(); + + let mut account = ResidencyAccount { + predicted: predicted_bytes.div_ceil(page_size.max(1)), + ..Default::default() + }; + for (page, &is_resident) in resident.iter().enumerate() { + if !is_resident { + continue; + } + account.resident_total += 1; + if router_pages.contains(page) { + account.resident_router += 1; + } else if selected_spans.iter().any(|s| s.contains(page)) { + account.resident_selected += 1; + } else if unselected_spans.iter().any(|s| s.contains(page)) { + account.resident_unselected += 1; + } + } + account +} diff --git a/crates/larql-vindex/src/runtime/router.rs b/crates/larql-vindex/src/runtime/router.rs new file mode 100644 index 000000000..26e4e3c59 --- /dev/null +++ b/crates/larql-vindex/src/runtime/router.rs @@ -0,0 +1,179 @@ +//! The bound router — scores, selects and weights. +//! +//! The user-facing sketch of a bound operation carried the router as a bare +//! tensor. A tensor alone cannot express the decision, though: top-k depth, +//! whether selected probabilities are renormalised, and whether learned +//! per-expert scales apply are all part of *what routing means* for a given +//! model, and all three change which experts run and how much each +//! contributes. +//! +//! They are bound properties rather than execution-time inference, and they +//! reuse `larql-compute`'s policy vocabulary so that the reference path and +//! the incumbent path are describing the same thing in the same words. + +use larql_compute::{MoeExpertScalePolicy, MoeTopKWeightPolicy}; + +use super::error::ExecutionError; +use super::tensor::BoundTensor; + +/// One expert selected for a token, with the weight it contributes at. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct SelectedExpert { + pub expert_id: u32, + /// Weight after every policy has been applied — what the reduction uses. + pub weight: f32, + /// Probability before renormalisation and per-expert scaling. Kept because + /// a routing disagreement shows up here first, while the final weight can + /// still coincide after renormalisation hides it. + pub raw_score: f32, +} + +/// Which kernel computes the router's scores. +/// +/// A **bound** choice, not an inferred one — the same discipline as every +/// other decision in this object. Binding the incumbent kernel is what makes +/// score parity a statement about the production code path; reimplementing a +/// BLAS-shaped loop here and calling the result "parity" would prove only that +/// two similar loops agree. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum RouterKernel { + /// Index-order f32 accumulation. Plain, allocation-light, and the oracle + /// every other kernel is checked against. + #[default] + Reference, + /// `larql-compute`'s own scoring: BLAS `sgemv` via ndarray, then its + /// softmax. Requires the weight to be contiguous row-major f32. + Incumbent, +} + +impl RouterKernel { + pub const fn name(self) -> &'static str { + match self { + Self::Reference => "reference", + Self::Incumbent => "incumbent", + } + } +} + +/// Per-expert output scaling, with its operand. +/// +/// One field, not two. The previous shape — a policy enum beside an +/// `Option` — permitted two invalid states: +/// +/// ```text +/// PerExpert + None policy demands a scale that was never bound +/// None + Some(scales) an operand nothing will read +/// ``` +/// +/// The first is not hypothetical. Gemma's routing policy is `PerExpert`, the +/// Gemma harness bound `None`, and execution silently declined to scale — +/// producing bit-identical router *scores* and normalised weights 7e-4 apart. +/// The parity ladder localised it, but the type had allowed the operation to +/// be built at all. +/// +/// Coupling them makes both states unrepresentable rather than merely +/// rejected, which is the difference between a check someone can forget to +/// call and a mistake that does not compile. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum BoundExpertScaling<'a> { + /// Selected weights are used as the routing policy produced them. + None, + /// Each selected weight is multiplied by `scales[expert_id]`. + /// + /// Applied *after* renormalisation, so the selected weights need not sum + /// to one afterwards. Before it, renormalisation would divide the learned + /// scale straight back out. + PerExpert { scales: BoundTensor<'a> }, +} + +impl BoundExpertScaling<'_> { + /// The scale operand, if one is bound. + pub fn scales(&self) -> Option<&BoundTensor<'_>> { + match self { + Self::None => None, + Self::PerExpert { scales } => Some(scales), + } + } + + /// The equivalent incumbent policy, for reporting against a + /// `MoeLayerWeights`. + pub const fn policy(&self) -> MoeExpertScalePolicy { + match self { + Self::None => MoeExpertScalePolicy::None, + Self::PerExpert { .. } => MoeExpertScalePolicy::PerExpert, + } + } + + pub const fn name(&self) -> &'static str { + match self { + Self::None => "none", + Self::PerExpert { .. } => "per_expert", + } + } +} + +/// Scoring and selection for one bank. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct BoundRouter<'a> { + /// `[num_experts, hidden]`. + pub weight: BoundTensor<'a>, + pub top_k: usize, + pub selected_weight: MoeTopKWeightPolicy, + /// Per-expert output scaling, together with the operand it needs. + pub scaling: BoundExpertScaling<'a>, + /// Which kernel scores. Defaults to the reference. + pub kernel: RouterKernel, +} + +impl BoundRouter<'_> { + /// Experts this router can address. + pub fn population(&self) -> usize { + self.weight.rows() + } + + /// Input width the router contracts over. + pub fn hidden_dim(&self) -> usize { + self.weight.cols() + } + + /// Check the router's operands against the population it addresses. + /// + /// `population` is the router's own **addressable** population, never a + /// resident shard's. The per-expert scale is routing semantics: expert 90's + /// learned scale is part of what routing means whether or not expert 90 is + /// resident here, so a shard must carry the whole vector. Truncating it to + /// the resident subset would make two shards of one model weight the same + /// expert differently. + pub fn validate(&self, population: usize, hidden: usize) -> Result<(), ExecutionError> { + self.weight.require_matrix(population, hidden)?; + if let Some(scales) = self.scaling.scales() { + scales.require_vector(population)?; + // A non-finite scale multiplies a valid routing weight into a NaN + // that then propagates through the reduction into the residual + // stream. Caught here, at bind time, rather than as a mysterious + // token later. + let values = scales.to_vec()?; + if let Some((expert, value)) = values + .iter() + .position(|v| !v.is_finite()) + .map(|i| (i, values[i])) + { + return Err(ExecutionError::NonFiniteExpertScale { + expert, + value, + operand: scales.describe(), + }); + } + } + Ok(()) + } + + pub fn describe(&self) -> String { + format!( + "router {} — top-{} of {}", + self.weight.describe(), + self.top_k, + self.population() + ) + } +} diff --git a/crates/larql-vindex/src/runtime/tensor.rs b/crates/larql-vindex/src/runtime/tensor.rs new file mode 100644 index 000000000..c2897ea75 --- /dev/null +++ b/crates/larql-vindex/src/runtime/tensor.rs @@ -0,0 +1,423 @@ +//! A bound tensor — resolved bytes plus everything needed to read them. +//! +//! This is the terminal product of binding. It carries no coordinate to look +//! up, no variant to choose and no candidate to prefer: the decisions are +//! already made, and what remains is arithmetic. +//! +//! # The view is part of the operand, not a caller's problem +//! +//! Storage shape and role shape differ whenever a view is in play — a tied +//! embedding serving an LM head is `[vocab, hidden]` on disk and `[hidden, +//! vocab]` to the operation that reads it. The view lives here so that every +//! consumer indexes in *role* coordinates and none of them has to know which +//! physical arrangement it got. That is the property that lets fused and +//! decomposed storage produce identical output through one executor. +//! +//! # Deliberately plain +//! +//! Decoding dispatches per element. That is the wrong shape for a hot path and +//! the right shape for a reference: it is obviously correct, it has no layout +//! special-cases to get wrong, and it is what a fast kernel gets checked +//! against. Production paths bind quantised regions to `larql-compute` +//! kernels; this decoder exists to say what the answer should have been. + +use crate::format::capability::binding::{ComponentView, RepresentationIdentity}; +use crate::format::capability::component::{ComponentContract, TensorKind}; +use crate::format::lyrw2::region_format::RegionFormat; + +use super::addressing::{Addressing, BlockOperand}; +use super::axis::Axis; +use super::consts::{ + BF16_SHIFT, COL_DIM, EXTENT_STORED_ROW, MATRIX_RANK, ROW_DIM, UNREGISTERED_CODEC, + WANTED_ROW_MAJOR_F32, +}; +use super::error::{ExecutionError, OperandUnsuitability}; + +/// Resolved bytes, with the encoding and access pattern that read them. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct BoundTensor<'a> { + /// Which catalogue declaration this came from. Diagnostics only — nothing + /// in execution branches on it. + representation: RepresentationIdentity, + bytes: &'a [u8], + format: RegionFormat, + /// Shape as stored on disk. + storage: ComponentContract, + /// Shape the operation sees, after `view`. + role: ComponentContract, + view: ComponentView, + /// How `format`'s bytes map to elements. Resolved once, at bind time. + addressing: Addressing, +} + +impl<'a> BoundTensor<'a> { + /// Bind bytes to a role. + /// + /// Fails if the view cannot apply to the storage shape, or if the region + /// is too short for the shape it claims — both binding faults, caught here + /// rather than as an out-of-bounds read later. + pub fn new( + representation: RepresentationIdentity, + bytes: &'a [u8], + format: RegionFormat, + storage: ComponentContract, + view: ComponentView, + ) -> Result { + let operand = representation.describe(); + let role = view + .apply_to(&storage) + .map_err(|_| ExecutionError::UnsupportedView { + view: view.describe(), + operand: operand.clone(), + })?; + let elements: usize = storage.shape.iter().map(|d| *d as usize).product(); + let addressing = Addressing::of(format) + .ok_or_else(|| ExecutionError::unsupported_format(format, operand.clone()))?; + let needed = addressing.region_bytes(elements); + if bytes.len() < needed { + return Err(ExecutionError::ShortRegion { + operand, + needed, + found: bytes.len(), + }); + } + Ok(Self { + representation, + bytes, + format, + storage, + role, + view, + addressing, + }) + } + + /// Convenience for the common case: bytes read exactly as stored. + pub fn direct( + representation: RepresentationIdentity, + bytes: &'a [u8], + format: RegionFormat, + storage: ComponentContract, + ) -> Result { + Self::new( + representation, + bytes, + format, + storage, + ComponentView::Direct, + ) + } + + /// The stored bytes, for tests that must confirm what is *behind* a view. + #[cfg(test)] + pub(crate) fn bytes_for_test(&self) -> &'a [u8] { + self.bytes + } + + /// This operand as a contiguous `f32` slice, if it genuinely is one. + /// + /// For binding an operand to a kernel that takes `&[f32]` in row-major + /// order — the incumbent's BLAS scoring, for instance — **without + /// reconstructing it**. A bridge that dequantised, repacked into an + /// incumbent-shaped temporary and then called the kernel could reach + /// numerical parity while proving nothing about the binding architecture. + /// This hands over the stored bytes or refuses. + /// + /// The error distinguishes format, view, alignment and length, because + /// each implies a different remedy: bind another variant, bind a + /// view-aware kernel, take an aligned copy, or reject the index. Only the + /// last is a defect. + pub fn as_f32_slice(&self) -> Result<&'a [f32], OperandUnsuitability> { + if self.format != RegionFormat::F32 { + return Err(OperandUnsuitability::ElementFormat { + found: self.format.name(), + wanted: WANTED_ROW_MAJOR_F32, + }); + } + if self.view != ComponentView::Direct { + return Err(OperandUnsuitability::NonDirectView { + view: self.view.describe(), + }); + } + // SAFETY: `align_to` is sound for any `T: Copy` with no invalid bit + // patterns, which `f32` satisfies. The empty-prefix check is what + // makes the result the *whole* slice rather than a shifted window. + let (prefix, values, _) = unsafe { self.bytes.align_to::() }; + if !prefix.is_empty() { + return Err(OperandUnsuitability::MisalignedBase { + wanted: WANTED_ROW_MAJOR_F32, + }); + } + if values.len() < self.len() { + return Err(OperandUnsuitability::Length { + expected: self.len(), + found: values.len(), + }); + } + Ok(&values[..self.len()]) + } + + /// This operand's super-blocks, as stored, for a block-native kernel. + /// + /// The blocked counterpart to [`Self::as_f32_slice`], and the same + /// contract: hand over the region's own bytes or refuse. A bridge that + /// dequantised, requantised or repacked into a kernel-shaped temporary + /// could reach identical numbers while proving nothing about the binding. + /// + /// # A column-prefix slice is honoured, not refused + /// + /// Unlike the f32 handover, this accepts `Slice { dim: 1, start: 0, .. }`. + /// That view is not an obstacle here — it is exactly the information a + /// block-native kernel needs and cannot recover from the bytes. Gemma's + /// `down` is stored `[hidden, 768]` and means `[hidden, 704]`; the kernel + /// must read 768-wide rows while the operation means 704, so + /// [`BlockOperand`] carries both. Requiring `Direct` would have forced + /// either a second binding of the same bytes or a lie about the stored + /// extent. + /// + /// A row-dimension slice or a transpose *is* refused: neither is a + /// contiguous run of whole rows, so the bytes are not the operand. + pub fn as_blocks( + &self, + wanted: RegionFormat, + ) -> Result, OperandUnsuitability> { + let wanted_name = wanted.registered_name().unwrap_or(UNREGISTERED_CODEC); + if self.format != wanted { + return Err(OperandUnsuitability::ElementFormat { + found: self.format.name(), + wanted: wanted_name, + }); + } + let (block_elems, block_bytes) = match self.addressing { + Addressing::Blocked { elements, bytes } => (elements, bytes), + Addressing::Scalar { .. } => { + return Err(OperandUnsuitability::ElementFormat { + found: self.format.name(), + wanted: wanted_name, + }) + } + }; + let honoured = match &self.view { + ComponentView::Direct => true, + // A prefix of every row: the stored rows are untouched and still + // contiguous, so the bytes remain the operand. + ComponentView::Slice { dim, start, .. } => *dim == COL_DIM && *start == 0, + ComponentView::Transpose => false, + }; + if !honoured { + return Err(OperandUnsuitability::NonDirectView { + view: self.view.describe(), + }); + } + + let storage_cols = self.storage_cols(); + if !storage_cols.is_multiple_of(block_elems) { + return Err(OperandUnsuitability::BlockAlignment { + extent: EXTENT_STORED_ROW, + found: storage_cols, + block: block_elems, + }); + } + let rows = self.rows(); + let row_bytes = (storage_cols / block_elems) * block_bytes; + let needed = rows * row_bytes; + // No length refusal here, because there is no shape that reaches this + // line and fails it. Binding sized the whole region against + // `rows × storage_cols` elements, and the block-aligned stored row + // extent checked just above makes per-row sizing equal that exactly. + // The assertion records the reasoning; re-checking it at every handover + // would suggest a case the reader should be able to imagine. + debug_assert!( + self.bytes.len() >= needed, + "{}: binding sized this region below its {rows} × {row_bytes} rows", + self.describe() + ); + Ok(BlockOperand { + bytes: &self.bytes[..needed], + rows, + storage_cols, + role_cols: self.cols(), + row_bytes, + }) + } + + pub fn representation(&self) -> &RepresentationIdentity { + &self.representation + } + + pub fn format(&self) -> RegionFormat { + self.format + } + + pub fn view(&self) -> &ComponentView { + &self.view + } + + /// The shape the operation sees. + pub fn contract(&self) -> &ComponentContract { + &self.role + } + + pub fn describe(&self) -> String { + self.representation.describe() + } + + /// Rows in role coordinates. + pub fn rows(&self) -> usize { + self.role.shape.first().copied().unwrap_or(0) as usize + } + + /// Columns in role coordinates. A vector has one column's worth per row. + pub fn cols(&self) -> usize { + match self.role.kind { + TensorKind::Matrix => self.role.shape.get(COL_DIM).copied().unwrap_or(0) as usize, + TensorKind::Vector => 1, + } + } + + pub fn len(&self) -> usize { + self.role.shape.iter().map(|d| *d as usize).product() + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Assert this operand has the expected role shape. + /// + /// Binding checks contracts, so a failure here means binding accepted an + /// operand it should have refused. + pub fn require_matrix(&self, rows: usize, cols: usize) -> Result<(), ExecutionError> { + if self.role.kind != TensorKind::Matrix || self.role.shape.len() != MATRIX_RANK { + return Err(ExecutionError::NotAMatrix { + operand: self.describe(), + found: self.role.describe(), + }); + } + self.require_axis(Axis::Rows, rows, self.rows())?; + self.require_axis(Axis::Columns, cols, self.cols()) + } + + /// Assert a vector operand's length. + pub fn require_vector(&self, len: usize) -> Result<(), ExecutionError> { + self.require_axis(Axis::Length, len, self.len()) + } + + fn require_axis( + &self, + axis: Axis, + expected: usize, + found: usize, + ) -> Result<(), ExecutionError> { + if expected == found { + return Ok(()); + } + Err(ExecutionError::DimensionMismatch { + operand: self.describe(), + axis, + expected, + found, + }) + } + + /// Read one row, in role coordinates, into `out`. + pub fn row_into(&self, row: usize, out: &mut [f32]) -> Result<(), ExecutionError> { + if row >= self.rows() { + return Err(ExecutionError::RowOutOfRange { + operand: self.describe(), + row, + rows: self.rows(), + }); + } + self.require_axis(Axis::Columns, out.len(), self.cols())?; + for (col, slot) in out.iter_mut().enumerate() { + *slot = self.decode_at(self.storage_index(row, col)?)?; + } + Ok(()) + } + + /// Read one row, in role coordinates. + pub fn row(&self, row: usize) -> Result, ExecutionError> { + let mut out = vec![0.0f32; self.cols()]; + self.row_into(row, &mut out)?; + Ok(out) + } + + /// Read a vector operand whole. + pub fn to_vec(&self) -> Result, ExecutionError> { + (0..self.len()).map(|i| self.decode_at(i)).collect() + } + + /// Map a role-coordinate cell to its index in storage. + /// + /// The whole point of the view: every caller indexes as the role, and the + /// physical arrangement is resolved here once. + fn storage_index(&self, row: usize, col: usize) -> Result { + let storage_cols = self.storage_cols(); + Ok(match &self.view { + ComponentView::Direct => row * storage_cols + col, + ComponentView::Transpose => col * storage_cols + row, + ComponentView::Slice { dim, start, .. } => { + let start = *start as usize; + match *dim { + ROW_DIM => (start + row) * storage_cols + col, + COL_DIM => row * storage_cols + start + col, + _ => { + return Err(ExecutionError::UnsupportedView { + view: self.view.describe(), + operand: self.describe(), + }) + } + } + } + }) + } + + /// Columns in the stored arrangement, before any view. + fn storage_cols(&self) -> usize { + match self.storage.kind { + TensorKind::Matrix => self.storage.shape.get(COL_DIM).copied().unwrap_or(0) as usize, + TensorKind::Vector => 1, + } + } + + /// Decode a single stored element. + /// + /// Addressing is resolved at bind time rather than here. Recomputing it per + /// element also meant building the operand name per element — a `String` + /// allocation on every scalar read, which dominated the reference decoder + /// and had nothing to do with decoding. + /// + /// A block-packed region has no per-element slot to read, so it is refused + /// here rather than indexed with an invented stride. That refusal is the + /// reference decoder's stated scope, not a defect in the bytes: a + /// quantised region is a missing kernel, and [`Self::as_blocks`] is how + /// one is given it. + fn decode_at(&self, index: usize) -> Result { + let Addressing::Scalar { bytes: stride } = self.addressing else { + return Err(ExecutionError::unsupported_format( + self.format, + self.describe(), + )); + }; + let at = index * stride; + let raw = self + .bytes + .get(at..at + stride) + .ok_or_else(|| ExecutionError::ShortRegion { + operand: self.describe(), + needed: at + stride, + found: self.bytes.len(), + })?; + Ok(match self.format { + RegionFormat::F32 => f32::from_le_bytes([raw[0], raw[1], raw[2], raw[3]]), + // Shared with the quantised kernels on purpose: its subnormal + // branch fixes a 2× error a local reimplementation would repeat. + RegionFormat::F16 => larql_compute::f16_to_f32(u16::from_le_bytes([raw[0], raw[1]])), + RegionFormat::BF16 => { + f32::from_bits(u32::from(u16::from_le_bytes([raw[0], raw[1]])) << BF16_SHIFT) + } + other => return Err(ExecutionError::unsupported_format(other, self.describe())), + }) + } +} diff --git a/crates/larql-vindex/src/runtime/tests/addressing.rs b/crates/larql-vindex/src/runtime/tests/addressing.rs new file mode 100644 index 000000000..5e60b6c79 --- /dev/null +++ b/crates/larql-vindex/src/runtime/tests/addressing.rs @@ -0,0 +1,177 @@ +//! Tests for `addressing` — how a codec's bytes map to its elements. +//! +//! The distinction these pin is not cosmetic. Sizing a blocked region with a +//! scalar stride under-counts by roughly 4× for Q4_K, which binds a region that +//! looks long enough and reads past its end on the last row. + +use crate::format::lyrw2::region_format::RegionFormat; + +use crate::runtime::addressing::{Addressing, BlockOperand}; +use crate::runtime::consts::{BF16_BYTES, F16_BYTES, F32_BYTES}; + +use super::support::{Q4K_BLOCK_BYTES, Q4K_BLOCK_ELEMS}; + +/// Every codec the format registry names, so a new one cannot be added without +/// this file having an opinion about it. +const REGISTERED: [RegionFormat; 11] = [ + RegionFormat::F32, + RegionFormat::F16, + RegionFormat::BF16, + RegionFormat::Q4_0, + RegionFormat::Q4K, + RegionFormat::Q6K, + RegionFormat::Q8_0, + RegionFormat::Fp4Larql, + RegionFormat::Mxfp4, + RegionFormat::Nvfp4, + RegionFormat::Mxfp8, +]; + +/// A padded `down`-shaped operand: 704 live columns stored in 768. +const PADDED_STORAGE_COLS: usize = 768; +const PADDED_ROLE_COLS: usize = 704; + +// ── Which codecs have a layout ───────────────────────────────────────────── + +#[test] +fn directly_addressed_codecs_report_their_element_width() { + assert_eq!( + Addressing::of(RegionFormat::F32), + Some(Addressing::Scalar { bytes: F32_BYTES }) + ); + assert_eq!( + Addressing::of(RegionFormat::F16), + Some(Addressing::Scalar { bytes: F16_BYTES }) + ); + assert_eq!( + Addressing::of(RegionFormat::BF16), + Some(Addressing::Scalar { bytes: BF16_BYTES }) + ); +} + +#[test] +fn q4k_reports_the_registrys_block_geometry() { + assert_eq!( + Addressing::of(RegionFormat::Q4K), + Some(Addressing::Blocked { + elements: Q4K_BLOCK_ELEMS, + bytes: Q4K_BLOCK_BYTES, + }) + ); +} + +#[test] +fn the_k_quants_are_blocked_and_the_rest_have_no_geometry() { + for format in [RegionFormat::Q4_0, RegionFormat::Q4K, RegionFormat::Q6K] { + assert!( + matches!(Addressing::of(format), Some(Addressing::Blocked { .. })), + "{format:?} should be blocked" + ); + } + // Guessing a layout for these would read plausible bytes at the wrong + // offsets and produce a well-shaped tensor of noise, so they are refused. + for format in [ + RegionFormat::Q8_0, + RegionFormat::Fp4Larql, + RegionFormat::Mxfp4, + RegionFormat::Nvfp4, + RegionFormat::Mxfp8, + ] { + assert_eq!(Addressing::of(format), None, "{format:?} has no geometry"); + } +} + +#[test] +fn an_unrecognised_codec_has_no_addressing() { + assert_eq!(Addressing::of(RegionFormat::Unknown(4_242)), None); +} + +#[test] +fn the_codec_sweep_still_covers_the_registry() { + // A guard on the guard. If a codec joins the registry and not `REGISTERED`, + // the sweep above narrows silently and stops being a sweep. + for (tag, expected) in REGISTERED.iter().enumerate() { + assert_eq!( + RegionFormat::from_u16(tag as u16), + *expected, + "REGISTERED is out of step with the format registry at tag {tag}" + ); + } + assert!(matches!( + RegionFormat::from_u16(REGISTERED.len() as u16), + RegionFormat::Unknown(_) + )); +} + +// ── Sizing ───────────────────────────────────────────────────────────────── + +#[test] +fn a_scalar_region_is_its_element_count_times_its_width() { + let f32s = Addressing::Scalar { bytes: F32_BYTES }; + assert_eq!(f32s.region_bytes(0), 0); + assert_eq!(f32s.region_bytes(3), 3 * F32_BYTES); +} + +#[test] +fn a_blocked_region_rounds_up_to_a_whole_block() { + let q4k = Addressing::of(RegionFormat::Q4K).expect("q4_k has geometry"); + assert_eq!(q4k.region_bytes(0), 0); + // One element still costs a whole super-block. + assert_eq!(q4k.region_bytes(1), Q4K_BLOCK_BYTES); + assert_eq!(q4k.region_bytes(Q4K_BLOCK_ELEMS), Q4K_BLOCK_BYTES); + assert_eq!(q4k.region_bytes(Q4K_BLOCK_ELEMS + 1), 2 * Q4K_BLOCK_BYTES); +} + +#[test] +fn blocked_sizing_is_not_the_scalar_answer() { + // The mistake this type exists to prevent: 256 Q4_K elements are 144 bytes, + // not 256 × anything. + let q4k = Addressing::of(RegionFormat::Q4K).expect("q4_k has geometry"); + let elements = 4 * Q4K_BLOCK_ELEMS; + assert_eq!(q4k.region_bytes(elements), 4 * Q4K_BLOCK_BYTES); + assert!(q4k.region_bytes(elements) < elements); +} + +// ── The block descriptors ────────────────────────────────────────────────── + +#[test] +fn only_a_blocked_addressing_reports_block_geometry() { + let q4k = Addressing::of(RegionFormat::Q4K).expect("q4_k has geometry"); + assert_eq!(q4k.block_elements(), Some(Q4K_BLOCK_ELEMS)); + assert_eq!(q4k.block_bytes(), Some(Q4K_BLOCK_BYTES)); + + let f32s = Addressing::Scalar { bytes: F32_BYTES }; + assert_eq!(f32s.block_elements(), None); + assert_eq!(f32s.block_bytes(), None); +} + +// ── The handover type ────────────────────────────────────────────────────── + +#[test] +fn a_block_operand_reports_the_padding_between_its_two_extents() { + let bytes = [0u8; Q4K_BLOCK_BYTES]; + let operand = BlockOperand { + bytes: &bytes, + rows: 1, + storage_cols: PADDED_STORAGE_COLS, + role_cols: PADDED_ROLE_COLS, + row_bytes: Q4K_BLOCK_BYTES, + }; + assert_eq!( + operand.padding_cols(), + PADDED_STORAGE_COLS - PADDED_ROLE_COLS + ); +} + +#[test] +fn an_unpadded_block_operand_has_no_padding() { + let bytes = [0u8; Q4K_BLOCK_BYTES]; + let operand = BlockOperand { + bytes: &bytes, + rows: 1, + storage_cols: Q4K_BLOCK_ELEMS, + role_cols: Q4K_BLOCK_ELEMS, + row_bytes: Q4K_BLOCK_BYTES, + }; + assert_eq!(operand.padding_cols(), 0); +} diff --git a/crates/larql-vindex/src/runtime/tests/bank.rs b/crates/larql-vindex/src/runtime/tests/bank.rs new file mode 100644 index 000000000..b44514eea --- /dev/null +++ b/crates/larql-vindex/src/runtime/tests/bank.rs @@ -0,0 +1,126 @@ +//! Colocated tests for `bank` — expert lookup and shape validation. + +use crate::format::capability::coordinate::BankCoordinate; +use crate::format::lyrw2::region_role::RegionRole; +use larql_compute::Activation; + +use super::support::ascending; +use crate::runtime::bank::{BoundBankOperation, BoundExpert}; +use crate::runtime::error::ExecutionError; +use crate::runtime::expert_kernel::ExpertKernel; +use crate::runtime::projection::BoundProjection; + +const LAYER: u32 = 7; +const BANK_ID: u16 = 1; +const INTERMEDIATE: u32 = 2; +const HIDDEN: u32 = 3; +/// The router's addressable population these banks live inside. +const POPULATION: usize = 128; + +fn expert(expert_id: u32) -> BoundExpert<'static> { + BoundExpert { + expert_id, + projection: BoundProjection::Decomposed { + gate: ascending(&RegionRole::Gate.name(), INTERMEDIATE, HIDDEN), + up: ascending(&RegionRole::Up.name(), INTERMEDIATE, HIDDEN), + }, + down: ascending(&RegionRole::Down.name(), HIDDEN, INTERMEDIATE), + } +} + +fn bank(ids: &[u32]) -> BoundBankOperation<'static> { + BoundBankOperation { + bank: BankCoordinate::new(LAYER, BANK_ID), + experts: ids.iter().copied().map(expert).collect(), + intermediate_dim: INTERMEDIATE as usize, + hidden_dim: HIDDEN as usize, + activation: Activation::Silu, + kernel: ExpertKernel::default(), + } +} + +// ── Lookup ───────────────────────────────────────────────────────────────── + +#[test] +fn an_expert_is_found_by_its_id_not_its_position() { + // Sparse ids are the normal case for a sharded bank: expert 40 may be the + // first one this shard holds. + let bank = bank(&[40, 41, 42]); + assert_eq!(bank.expert(41, POPULATION).unwrap().expert_id, 41); + assert_eq!(bank.population(), 3); +} + +#[test] +fn a_selected_expert_the_bank_does_not_hold_is_refused_not_skipped() { + // Skipping would drop a fraction of the FFN contribution and produce a + // token that looks entirely reasonable. + let err = bank(&[0, 1]).expert(9, POPULATION).unwrap_err(); + assert!(matches!( + err, + ExecutionError::SelectedExpertNotResident { + expert: 9, + resident: 2, + .. + } + )); + assert!(err.to_string().contains('9')); +} + +#[test] +fn an_empty_bank_refuses_every_id() { + assert_eq!(bank(&[]).population(), 0); + assert!(bank(&[]).expert(0, POPULATION).is_err()); +} + +// ── Validation ───────────────────────────────────────────────────────────── + +#[test] +fn a_consistent_bank_validates() { + bank(&[0, 1, 2]).validate().unwrap(); +} + +#[test] +fn an_expert_whose_down_is_transposed_is_refused() { + // `down` contracts the intermediate axis away, so it is [hidden, + // intermediate]. The swap still multiplies when the two happen to be + // equal, which is why the fixture keeps them different. + let mut broken = bank(&[0]); + broken.experts[0].down = ascending(&RegionRole::Down.name(), INTERMEDIATE, HIDDEN); + assert!(broken.validate().is_err()); +} + +#[test] +fn a_bank_whose_declared_intermediate_disagrees_with_its_experts_is_refused() { + let mut broken = bank(&[0]); + broken.intermediate_dim = (INTERMEDIATE + 1) as usize; + let err = broken.validate().unwrap_err(); + assert!(matches!(err, ExecutionError::DimensionMismatch { .. })); +} + +#[test] +fn one_broken_expert_fails_the_whole_bank() { + let mut broken = bank(&[0, 1, 2]); + broken.experts[2].down = ascending(&RegionRole::Down.name(), HIDDEN, INTERMEDIATE + 1); + assert!(broken.validate().is_err()); +} + +#[test] +fn an_expert_validates_against_the_shape_it_is_given() { + expert(0) + .validate(INTERMEDIATE as usize, HIDDEN as usize) + .unwrap(); + assert!(expert(0).validate(INTERMEDIATE as usize, 99).is_err()); +} + +// ── Description ──────────────────────────────────────────────────────────── + +#[test] +fn a_bank_describes_its_coordinate_population_and_shape() { + let text = bank(&[0, 1]).describe(); + assert!( + text.contains(&BankCoordinate::new(LAYER, BANK_ID).describe()), + "{text}" + ); + assert!(text.contains('2'), "{text}"); + assert!(text.contains("2×3"), "{text}"); +} diff --git a/crates/larql-vindex/src/runtime/tests/execute.rs b/crates/larql-vindex/src/runtime/tests/execute.rs new file mode 100644 index 000000000..bbdaa0f7d --- /dev/null +++ b/crates/larql-vindex/src/runtime/tests/execute.rs @@ -0,0 +1,302 @@ +//! Colocated tests for `execute` — paths fixture A does not reach. +//! +//! Fixture A covers the direct route against an oracle. These cover the +//! latent route, the per-expert scale policy, the raw-softmax policy and the +//! refusals — all reachable code that a direct f32 fixture never touches. + +use crate::format::capability::binding::{ComponentView, RepresentationIdentity}; +use crate::format::capability::component::ComponentContract; +use crate::format::lyrw2::region_format::RegionFormat; +use crate::runtime::router::{BoundExpertScaling, RouterKernel}; +use larql_compute::MoeTopKWeightPolicy; + +use super::operation::{direct, latent}; +use super::support::{ascending, vector, TEST_VARIANT}; +use crate::runtime::error::{ExecutionError, OperandUnsuitability}; +use crate::runtime::execute::{execute, execute_traced, execute_with}; +use crate::runtime::inputs::MoeInputs; +use crate::runtime::tensor::BoundTensor; +use crate::runtime::trace::{CollectedTrace, NoTrace, TraceSink}; +use crate::runtime::transform::TransformStage; + +const RESIDUAL_WIDTH: usize = 6; +const LATENT_WIDTH: usize = 3; +const SCALE: &str = "router_scale"; +const ROUTER: &str = "router"; +/// BLAS `sgemv` and an index-order f32 sum reach the same value by different +/// summation orders, so scores agree closely rather than exactly. +const SCORE_TOLERANCE: f32 = 1e-6; +const LATENT_OUT: &str = "latent_out"; + +/// The router's own weights, bound through a transpose: the role still sees +/// `[population, residual]`, but the bytes are stored the other way round. +fn transposed_router() -> BoundTensor<'static> { + const POPULATION: usize = 4; + let values: Vec = (0..POPULATION * RESIDUAL_WIDTH) + .map(|i| (i + 1) as f32) + .collect(); + let bytes: Vec = values.iter().flat_map(|v| v.to_le_bytes()).collect(); + let leaked: &'static [u8] = Box::leak(bytes.into_boxed_slice()); + BoundTensor::new( + RepresentationIdentity::new(ROUTER, TEST_VARIANT), + leaked, + RegionFormat::F32, + ComponentContract::matrix(RESIDUAL_WIDTH as u32, POPULATION as u32), + ComponentView::Transpose, + ) + .expect("a transposed router is a legitimate binding") +} + +fn residual() -> Vec { + vec![0.4, -0.3, 0.9, 0.1, -0.7, 0.2] +} + +// ── The latent route ─────────────────────────────────────────────────────── + +#[test] +fn a_latent_operation_returns_a_residual_width_delta() { + // The point of the output transform: the banks run narrow, the delta + // rejoins the residual at full width. + let out = execute(&latent(), MoeInputs::shared(&residual())).unwrap(); + assert_eq!(out.len(), RESIDUAL_WIDTH); +} + +#[test] +fn a_latent_operation_reduces_at_the_projected_width() { + let (out, trace) = execute_traced(&latent(), MoeInputs::shared(&residual())).unwrap(); + assert_eq!(trace.routed_input.len(), LATENT_WIDTH); + assert_eq!(trace.reduced.len(), LATENT_WIDTH); + assert_eq!(trace.residual_delta.len(), RESIDUAL_WIDTH); + assert_eq!(out, trace.residual_delta); +} + +#[test] +fn a_latent_routed_input_is_not_the_residual() { + // If the input transform were skipped, these would coincide. + let (_, trace) = execute_traced(&latent(), MoeInputs::shared(&residual())).unwrap(); + assert_ne!(trace.routed_input, residual()); +} + +#[test] +fn a_direct_operation_routes_on_the_residual_itself() { + let (_, trace) = execute_traced(&direct(), MoeInputs::shared(&residual())).unwrap(); + assert_eq!(trace.routed_input, residual()); +} + +#[test] +fn latent_and_direct_operations_both_execute_through_one_path() { + for op in [direct(), latent()] { + let out = execute(&op, MoeInputs::shared(&residual())).unwrap(); + assert_eq!(out.len(), RESIDUAL_WIDTH); + assert!(out.iter().all(|v| v.is_finite()), "{out:?}"); + } +} + +// ── Weight policies ──────────────────────────────────────────────────────── + +#[test] +fn raw_softmax_leaves_the_selected_weights_unnormalised() { + let mut op = direct(); + op.router.selected_weight = MoeTopKWeightPolicy::RawSoftmax; + let (_, trace) = execute_traced(&op, MoeInputs::shared(&residual())).unwrap(); + let total: f32 = trace.gate_weights().iter().sum(); + assert!(total < 1.0, "top-k of a larger population sums to {total}"); +} + +#[test] +fn renormalised_and_raw_softmax_select_the_same_experts() { + // The policy changes the weights, never the selection. + let mut raw = direct(); + raw.router.selected_weight = MoeTopKWeightPolicy::RawSoftmax; + let (_, raw_trace) = execute_traced(&raw, MoeInputs::shared(&residual())).unwrap(); + let (_, renorm_trace) = execute_traced(&direct(), MoeInputs::shared(&residual())).unwrap(); + assert_eq!(raw_trace.selected_ids(), renorm_trace.selected_ids()); + assert_ne!(raw_trace.gate_weights(), renorm_trace.gate_weights()); +} + +#[test] +fn a_per_expert_scale_multiplies_after_renormalisation() { + // Applying it before would let renormalisation divide the learned scale + // straight back out, making the whole policy a no-op. + let mut op = direct(); + let population = op.banks[0].experts.len(); + op.router.scaling = BoundExpertScaling::PerExpert { + scales: vector(SCALE, &vec![2.0f32; population]), + }; + + let (_, scaled) = execute_traced(&op, MoeInputs::shared(&residual())).unwrap(); + let (_, plain) = execute_traced(&direct(), MoeInputs::shared(&residual())).unwrap(); + let total: f32 = scaled.gate_weights().iter().sum(); + assert!((total - 2.0).abs() < 1e-5, "scaled weights sum to {total}"); + assert_eq!(scaled.selected_ids(), plain.selected_ids()); +} + +#[test] +fn unscaled_routing_leaves_the_renormalised_weights_summing_to_one() { + let mut op = direct(); + let population = op.banks[0].experts.len(); + op.router.scaling = BoundExpertScaling::None; + let _ = population; + let (_, trace) = execute_traced(&op, MoeInputs::shared(&residual())).unwrap(); + let total: f32 = trace.gate_weights().iter().sum(); + assert!((total - 1.0).abs() < 1e-5, "{total}"); +} + +// ── Depth ────────────────────────────────────────────────────────────────── + +#[test] +fn top_k_bounds_how_many_experts_run() { + let mut op = direct(); + op.router.top_k = 1; + let (_, trace) = execute_traced(&op, MoeInputs::shared(&residual())).unwrap(); + assert_eq!(trace.selection.len(), 1); + assert_eq!(trace.expert_outputs.len(), 1); + assert!((trace.gate_weights()[0] - 1.0).abs() < 1e-6, "renormalised"); +} + +#[test] +fn a_top_k_of_zero_selects_nothing_and_contributes_nothing() { + let mut op = direct(); + op.router.top_k = 0; + let (out, trace) = execute_traced(&op, MoeInputs::shared(&residual())).unwrap(); + assert!(trace.selection.is_empty()); + assert_eq!(out, vec![0.0f32; RESIDUAL_WIDTH]); +} + +#[test] +fn a_top_k_beyond_the_population_runs_every_expert() { + let mut op = direct(); + let population = op.banks[0].experts.len(); + op.router.top_k = population + 5; + let (_, trace) = execute_traced(&op, MoeInputs::shared(&residual())).unwrap(); + assert_eq!(trace.selection.len(), population); +} + +// ── Refusals ─────────────────────────────────────────────────────────────── + +#[test] +fn a_residual_of_the_wrong_width_is_refused() { + assert!(execute(&direct(), MoeInputs::shared(&[1.0, 2.0])).is_err()); + assert!(execute(&latent(), MoeInputs::shared(&[1.0, 2.0])).is_err()); +} + +#[test] +fn a_router_selecting_an_expert_the_bank_lacks_is_refused() { + // Router and bank disagreeing about the population is a binding fault, and + // must not degrade into a quietly missing contribution. + let mut op = direct(); + op.banks[0].experts.truncate(1); + op.banks[0].experts[0].expert_id = 0; + op.router.top_k = op.router.population(); + let err = execute(&op, MoeInputs::shared(&residual())).unwrap_err(); + assert!(err.to_string().contains("selected expert"), "{err}"); +} + +#[test] +fn an_output_transform_that_lands_off_the_residual_width_is_refused() { + // Constructed so the *input* stage still succeeds: widening `residual_dim` + // instead would fail at the input projection and never reach this check. + let mut op = latent(); + let output_stage = op + .transforms + .iter_mut() + .find(|t| t.stage == TransformStage::RoutedOutput) + .expect("the latent operation binds an output stage"); + output_stage.weight = ascending(LATENT_OUT, (RESIDUAL_WIDTH - 1) as u32, LATENT_WIDTH as u32); + + let err = execute(&op, MoeInputs::shared(&residual())).unwrap_err(); + let text = err.to_string(); + assert!(text.contains("operation output"), "{text}"); + assert!(text.contains(&RESIDUAL_WIDTH.to_string()), "{text}"); +} + +// ── The sink is the only difference between the two entry points ─────────── + +#[test] +fn the_no_op_sink_and_the_collecting_sink_agree_on_the_result() { + let op = latent(); + let mut none = NoTrace; + let mut collected = CollectedTrace::default(); + let a = execute_with(&op, MoeInputs::shared(&residual()), &mut none).unwrap(); + let b = execute_with(&op, MoeInputs::shared(&residual()), &mut collected).unwrap(); + assert_eq!(a, b); + assert_eq!(b, collected.residual_delta); +} + +#[test] +fn the_no_op_sink_accepts_every_checkpoint_without_recording() { + // Defaulted trait methods: the production sink must stay a single empty + // impl even as checkpoints are added. + let mut sink = NoTrace; + sink.routed_input(&[1.0]); + sink.router_scores(&[1.0]); + sink.selection(&[]); + sink.expert_output(0, &[1.0]); + sink.reduced(&[1.0]); + sink.residual_delta(&[1.0]); +} + +#[test] +fn a_collected_trace_looks_up_expert_outputs_by_id() { + let (_, trace) = execute_traced(&direct(), MoeInputs::shared(&residual())).unwrap(); + let first = trace.selected_ids()[0]; + assert!(trace.expert_output(first).is_some()); + assert!(trace.expert_output(u32::MAX).is_none()); +} + +#[test] +fn the_latent_stages_are_the_only_transforms_bound() { + let op = latent(); + let stages: Vec = op.transforms.iter().map(|t| t.stage).collect(); + assert_eq!( + stages, + vec![TransformStage::RoutedInput, TransformStage::RoutedOutput] + ); +} + +// ── The bound router kernel ──────────────────────────────────────────────── + +#[test] +fn the_incumbent_router_kernel_scores_the_same_population_as_the_reference() { + // Rung 1's claim, at unit scale and without a checkpoint. The two kernels + // sum in different orders — BLAS `sgemv` against index-order f32 — so this + // asserts the *selection*, which is discrete, and reports the weights. + let mut op = direct(); + let (_, reference) = execute_traced(&op, MoeInputs::shared(&residual())).unwrap(); + + op.router.kernel = RouterKernel::Incumbent; + let (_, incumbent) = execute_traced(&op, MoeInputs::shared(&residual())).unwrap(); + + assert_eq!(reference.selected_ids(), incumbent.selected_ids()); + assert_eq!(reference.router_scores.len(), incumbent.router_scores.len()); + let worst = reference + .router_scores + .iter() + .zip(&incumbent.router_scores) + .map(|(a, b)| (a - b).abs()) + .fold(0.0f32, f32::max); + assert!(worst <= SCORE_TOLERANCE, "scores differ by {worst}"); +} + +#[test] +fn the_incumbent_router_kernel_refuses_an_operand_it_cannot_read_as_stored() { + // It takes contiguous row-major f32. A transposed router operand is a + // legitimate binding the reference serves and this kernel cannot, so the + // refusal names the view rather than silently repacking. + let mut op = direct(); + op.router.kernel = RouterKernel::Incumbent; + op.router.weight = transposed_router(); + + let err = execute(&op, MoeInputs::shared(&residual())).unwrap_err(); + assert!( + matches!( + err, + ExecutionError::KernelOperandUnsuitable { + reason: OperandUnsuitability::NonDirectView { .. }, + .. + } + ), + "{err}" + ); + assert!(err.to_string().contains("incumbent"), "{err}"); +} diff --git a/crates/larql-vindex/src/runtime/tests/expert_kernel.rs b/crates/larql-vindex/src/runtime/tests/expert_kernel.rs new file mode 100644 index 000000000..7c51f02e7 --- /dev/null +++ b/crates/larql-vindex/src/runtime/tests/expert_kernel.rs @@ -0,0 +1,805 @@ +//! Tests for `expert_kernel` — binding the production expert kernel. +//! +//! The headline is [`the_bound_kernel_reproduces_the_incumbent_call_bit_for_bit`]. +//! Everything else exists so that a bit-identical result cannot be a +//! bit-identical *mistake*: two calls of one function agree whatever operands +//! they are handed, so the refusals below are what establish that the operands +//! reaching it are the right ones. + +use larql_compute::cpu::ops::moe::{ + quantize_x_to_q8k, run_single_expert_q4k_q8k_into, ExpertScratch, +}; +use larql_compute::cpu::ops::q4_common::dequantize_q4_k; +use larql_compute::{Activation, MoeTopKWeightPolicy}; + +use crate::format::capability::binding::{ComponentView, RepresentationIdentity}; +use crate::format::capability::component::ComponentContract; +use crate::format::capability::coordinate::BankCoordinate; +use crate::format::lyrw2::region_format::RegionFormat; + +use crate::runtime::axis::Axis; +use crate::runtime::bank::{BoundBankOperation, BoundExpert}; +use crate::runtime::consts::{COL_DIM, FUSED_PROJECTION_HALVES}; +use crate::runtime::error::{ExecutionError, OperandUnsuitability}; +use crate::runtime::execute::execute_traced; +use crate::runtime::expert_kernel::ExpertKernel; +use crate::runtime::inputs::MoeInputs; +use crate::runtime::operation::BoundMoeOperation; +use crate::runtime::projection::BoundProjection; +use crate::runtime::reduction::BoundReduction; +use crate::runtime::router::{BoundExpertScaling, BoundRouter, RouterKernel}; +use crate::runtime::tensor::BoundTensor; + +use super::support::{q4k_bytes, Q4K_BLOCK_ELEMS}; + +/// One super-block wide, so the Q8_K activation quantiser is satisfied and the +/// fixture stays small enough to reason about. +const HIDDEN: usize = Q4K_BLOCK_ELEMS; +/// Half a super-block, so `down` is stored padded — the shape every real +/// k-quant MoE layer has, and the one a kernel binding gets wrong. +const INTERMEDIATE: usize = Q4K_BLOCK_ELEMS / 2; +const INTERMEDIATE_PADDED: usize = Q4K_BLOCK_ELEMS; +const POPULATION: usize = 2; +const TOP_K: usize = 1; +const LAYER: u32 = 0; +const BANK_ID: u16 = 0; +/// Gemma's activation, and one of the two the incumbent kernel implements. +const ACTIVATION: Activation = Activation::GeluTanh; + +const VARIANT: &str = "test"; +const GATE_UP_REGION_SET: &str = "gate_up_fused"; +const DOWN_REGION_SET: &str = "down"; +const ROUTER_REGION_SET: &str = "router"; + +const GATE_UP_SEED: usize = 1; +const DOWN_SEED: usize = 7; + +/// The reference dequantises to f32 while the incumbent keeps an integer dot +/// against a Q8_K activation, so the two agree to quantisation noise, not to +/// the bit. A band, not a target. +const KERNEL_TOLERANCE: f32 = 5e-2; + +fn identity(region_set: &str) -> RepresentationIdentity { + RepresentationIdentity::new(region_set, VARIANT) +} + +fn f32_bytes(values: &[f32]) -> Vec { + values.iter().flat_map(|v| v.to_le_bytes()).collect() +} + +fn direct<'a>( + region_set: &str, + bytes: &'a [u8], + format: RegionFormat, + rows: u32, + cols: u32, +) -> BoundTensor<'a> { + BoundTensor::direct( + identity(region_set), + bytes, + format, + ComponentContract::matrix(rows, cols), + ) + .expect("well-formed fixture operand") +} + +fn column_prefix<'a>( + region_set: &str, + bytes: &'a [u8], + format: RegionFormat, + rows: u32, + stored_cols: u32, + role_cols: u32, +) -> BoundTensor<'a> { + BoundTensor::new( + identity(region_set), + bytes, + format, + ComponentContract::matrix(rows, stored_cols), + ComponentView::Slice { + dim: COL_DIM, + start: 0, + len: role_cols, + }, + ) + .expect("well-formed fixture operand") +} + +fn max_abs_diff(a: &[f32], b: &[f32]) -> f32 { + assert_eq!(a.len(), b.len()); + a.iter() + .zip(b) + .map(|(x, y)| (x - y).abs()) + .fold(0.0f32, f32::max) +} + +/// Q4_K bytes, plus the same values dequantised, so both kernels can be bound +/// over one set of weights. +struct Fixture { + gate_up: Vec, + gate_up_f32: Vec, + down: Vec, + down_f32: Vec, + router: Vec, + input: Vec, +} + +impl Fixture { + fn new() -> Self { + let gate_up = q4k_bytes(FUSED_PROJECTION_HALVES * INTERMEDIATE, HIDDEN, GATE_UP_SEED); + let down = q4k_bytes(HIDDEN, INTERMEDIATE_PADDED, DOWN_SEED); + let gate_up_f32 = f32_bytes(&dequantize_q4_k( + &gate_up, + FUSED_PROJECTION_HALVES * INTERMEDIATE * HIDDEN, + )); + let down_f32 = f32_bytes(&dequantize_q4_k(&down, HIDDEN * INTERMEDIATE_PADDED)); + + // A signed input, so the kernel exercises both halves of the 4-bit + // range rather than a one-sided ramp. + let input: Vec = (0..HIDDEN).map(|i| (i % 11) as f32 * 0.1 - 0.5).collect(); + + // Expert 0's router row *is* the input, so its logit is ‖x‖² while + // expert 1's is zero. An earlier fixture filled row 0 with ones, which + // scored `sum(x)` — and this input sums to approximately zero, so the + // two logits tied and selection fell to whichever side the float + // residue landed. A fixture that is accidentally testing a tie is not + // testing the kernel. + let mut router = vec![0.0f32; POPULATION * HIDDEN]; + router[..HIDDEN].copy_from_slice(&input); + Self { + gate_up, + gate_up_f32, + down, + down_f32, + router: f32_bytes(&router), + input, + } + } + + /// The store's own Q4_K blocks, `down` bound at its stored padded width + /// with the role seeing the live columns. + fn q4k_expert(&self) -> BoundExpert<'_> { + BoundExpert { + expert_id: 0, + projection: BoundProjection::Fused { + gate_up: direct( + GATE_UP_REGION_SET, + &self.gate_up, + RegionFormat::Q4K, + (FUSED_PROJECTION_HALVES * INTERMEDIATE) as u32, + HIDDEN as u32, + ), + }, + down: column_prefix( + DOWN_REGION_SET, + &self.down, + RegionFormat::Q4K, + HIDDEN as u32, + INTERMEDIATE_PADDED as u32, + INTERMEDIATE as u32, + ), + } + } + + /// The same weights dequantised, for the reference kernel. + fn f32_expert(&self) -> BoundExpert<'_> { + BoundExpert { + expert_id: 0, + projection: BoundProjection::Fused { + gate_up: direct( + GATE_UP_REGION_SET, + &self.gate_up_f32, + RegionFormat::F32, + (FUSED_PROJECTION_HALVES * INTERMEDIATE) as u32, + HIDDEN as u32, + ), + }, + down: column_prefix( + DOWN_REGION_SET, + &self.down_f32, + RegionFormat::F32, + HIDDEN as u32, + INTERMEDIATE_PADDED as u32, + INTERMEDIATE as u32, + ), + } + } + + fn operation<'a>( + &'a self, + expert: BoundExpert<'a>, + kernel: ExpertKernel, + ) -> BoundMoeOperation<'a> { + self.operation_with(expert, kernel, ACTIVATION) + } + + fn operation_with<'a>( + &'a self, + expert: BoundExpert<'a>, + kernel: ExpertKernel, + activation: Activation, + ) -> BoundMoeOperation<'a> { + BoundMoeOperation { + router: BoundRouter { + weight: direct( + ROUTER_REGION_SET, + &self.router, + RegionFormat::F32, + POPULATION as u32, + HIDDEN as u32, + ), + top_k: TOP_K, + selected_weight: MoeTopKWeightPolicy::RenormalizedSoftmax, + scaling: BoundExpertScaling::None, + kernel: RouterKernel::default(), + }, + transforms: Vec::new(), + banks: vec![BoundBankOperation { + bank: BankCoordinate::new(LAYER, BANK_ID), + experts: vec![expert], + intermediate_dim: INTERMEDIATE, + hidden_dim: HIDDEN, + activation, + kernel, + }], + reduction: BoundReduction::WeightedSum, + residual_dim: HIDDEN, + } + } + + /// The incumbent's own function, called directly on the same bytes. + fn incumbent_expert_output(&self) -> Vec { + let q8k = quantize_x_to_q8k(&self.input); + let mut scratch = ExpertScratch::new(HIDDEN, INTERMEDIATE, INTERMEDIATE_PADDED); + run_single_expert_q4k_q8k_into( + &mut scratch, + &q8k, + &self.gate_up, + &self.down, + INTERMEDIATE, + ACTIVATION, + ) + .to_vec() + } + + fn run(&self, operation: &BoundMoeOperation<'_>) -> Vec { + let (_, trace) = execute_traced(operation, MoeInputs::shared(&self.input)) + .expect("the fixture executes"); + trace + .expert_output(0) + .expect("expert 0 was selected") + .to_vec() + } +} + +// ── The claim ────────────────────────────────────────────────────────────── + +#[test] +fn the_bound_kernel_reproduces_the_incumbent_call_bit_for_bit() { + // The rung's whole point, at unit scale and without a checkpoint: the same + // Q4_K bytes, the same activation, the same production function — reached + // through a bound operation rather than by calling it directly. + let fixture = Fixture::new(); + let operation = fixture.operation(fixture.q4k_expert(), ExpertKernel::IncumbentQ4kQ8k); + operation.validate().expect("the fixture binds"); + + let bound = fixture.run(&operation); + let incumbent = fixture.incumbent_expert_output(); + + assert_eq!(bound.len(), HIDDEN); + assert_eq!( + bound, + incumbent, + "bound kernel diverged from the incumbent call by {}", + max_abs_diff(&bound, &incumbent) + ); +} + +#[test] +fn the_bound_kernel_is_not_producing_zeros() { + // Guards the guard. The incumbent's short-slab branch zeroes its output and + // returns successfully, so an all-zero agreement would be two failures + // agreeing rather than a parity result. + let fixture = Fixture::new(); + let operation = fixture.operation(fixture.q4k_expert(), ExpertKernel::IncumbentQ4kQ8k); + let bound = fixture.run(&operation); + assert!( + bound.iter().any(|v| v.abs() > f32::EPSILON), + "an all-zero expert output means the kernel took its refusal branch" + ); + assert!(bound.iter().all(|v| v.is_finite())); +} + +#[test] +fn the_reference_kernel_agrees_with_the_bound_one_within_quantisation_noise() { + // The independent leg. Two calls of one function agree whatever they are + // handed; this says the operands were the right ones, because a swapped + // gate/up half or a mis-strided `down` would not land inside a band. + let fixture = Fixture::new(); + let bound = + fixture.run(&fixture.operation(fixture.q4k_expert(), ExpertKernel::IncumbentQ4kQ8k)); + let reference = fixture.run(&fixture.operation(fixture.f32_expert(), ExpertKernel::Reference)); + + let diff = max_abs_diff(&bound, &reference); + assert!( + diff <= KERNEL_TOLERANCE, + "kernels disagree by {diff}, beyond quantisation noise" + ); +} + +#[test] +fn the_two_kernels_are_not_trivially_identical() { + // Guards the tolerance. If the reference somehow ran the same arithmetic, + // the band above would be vacuous and would keep passing through a real + // regression. + let fixture = Fixture::new(); + let bound = + fixture.run(&fixture.operation(fixture.q4k_expert(), ExpertKernel::IncumbentQ4kQ8k)); + let reference = fixture.run(&fixture.operation(fixture.f32_expert(), ExpertKernel::Reference)); + assert_ne!( + bound, reference, + "an integer-dot kernel and an f32 reference should not be bit-identical" + ); +} + +// ── Naming ───────────────────────────────────────────────────────────────── + +#[test] +fn each_kernel_has_a_distinct_name_and_the_reference_is_the_default() { + assert_eq!(ExpertKernel::default(), ExpertKernel::Reference); + assert_ne!( + ExpertKernel::Reference.name(), + ExpertKernel::IncumbentQ4kQ8k.name() + ); + assert!(ExpertKernel::IncumbentQ4kQ8k.name().contains("q4k")); +} + +#[test] +fn the_bank_reports_which_kernel_it_bound() { + let fixture = Fixture::new(); + let operation = fixture.operation(fixture.q4k_expert(), ExpertKernel::IncumbentQ4kQ8k); + assert!(operation + .describe() + .contains(ExpertKernel::IncumbentQ4kQ8k.name())); +} + +// ── Refusals: the activation ─────────────────────────────────────────────── + +#[test] +fn an_activation_the_kernel_does_not_implement_is_refused() { + // The most dangerous case, because nothing else would signal it: the + // incumbent's loop branches on GeluTanh and falls through to SiLU, so a + // ReLU bank would run SiLU and return finite plausible values. + let fixture = Fixture::new(); + let operation = fixture.operation_with( + fixture.q4k_expert(), + ExpertKernel::IncumbentQ4kQ8k, + Activation::ReLU, + ); + let err = operation.validate().unwrap_err(); + assert!( + matches!(err, ExecutionError::KernelActivationUnsupported { .. }), + "{err}" + ); + assert!(err.to_string().contains("ReLU"), "{err}"); +} + +#[test] +fn both_activations_the_kernel_implements_are_accepted() { + let fixture = Fixture::new(); + for activation in [Activation::Silu, Activation::GeluTanh] { + fixture + .operation_with( + fixture.q4k_expert(), + ExpertKernel::IncumbentQ4kQ8k, + activation, + ) + .validate() + .unwrap_or_else(|e| panic!("{activation:?} should bind: {e}")); + } +} + +#[test] +fn the_reference_kernel_accepts_every_activation() { + // The counter-case: the restriction belongs to the incumbent kernel, not to + // the runtime. + let fixture = Fixture::new(); + fixture + .operation_with( + fixture.f32_expert(), + ExpertKernel::Reference, + Activation::ReLU, + ) + .validate() + .expect("the reference implements all four"); +} + +// ── Refusals: the operands ───────────────────────────────────────────────── + +#[test] +fn a_decomposed_projection_is_refused_by_arrangement() { + // Remedy: bind the fused variant, or a kernel that takes two regions. + // Stitching them into one temporary would make the parity result a + // statement about the stitching. + let fixture = Fixture::new(); + let half = fixture.gate_up.len() / FUSED_PROJECTION_HALVES; + let expert = BoundExpert { + expert_id: 0, + projection: BoundProjection::Decomposed { + gate: direct( + GATE_UP_REGION_SET, + &fixture.gate_up[..half], + RegionFormat::Q4K, + INTERMEDIATE as u32, + HIDDEN as u32, + ), + up: direct( + GATE_UP_REGION_SET, + &fixture.gate_up[half..], + RegionFormat::Q4K, + INTERMEDIATE as u32, + HIDDEN as u32, + ), + }, + down: fixture.q4k_expert().down, + }; + let err = fixture + .operation(expert, ExpertKernel::IncumbentQ4kQ8k) + .validate() + .unwrap_err(); + assert!( + matches!( + err, + ExecutionError::KernelOperandUnsuitable { + reason: OperandUnsuitability::Arrangement { .. }, + .. + } + ), + "{err}" + ); +} + +#[test] +fn dequantised_operands_are_refused_by_format() { + // Remedy: bind the Q4_K variant, or the reference kernel. Reading f32 bytes + // as super-blocks is exactly the substitution this refuses. + let fixture = Fixture::new(); + let err = fixture + .operation(fixture.f32_expert(), ExpertKernel::IncumbentQ4kQ8k) + .validate() + .unwrap_err(); + assert!( + matches!( + err, + ExecutionError::KernelOperandUnsuitable { + reason: OperandUnsuitability::ElementFormat { .. }, + .. + } + ), + "{err}" + ); +} + +#[test] +fn a_down_region_padded_to_the_wrong_width_is_refused_by_the_kernel() { + // Shape-valid — the role still exposes `INTERMEDIATE` columns — but stored + // at a width the kernel would not stride by. Only the kernel binding can + // catch this one, which is why the stored extent travels with the operand. + let fixture = Fixture::new(); + let wide = INTERMEDIATE_PADDED + Q4K_BLOCK_ELEMS; + let bytes = q4k_bytes(HIDDEN, wide, DOWN_SEED); + let expert = BoundExpert { + expert_id: 0, + down: column_prefix( + DOWN_REGION_SET, + &bytes, + RegionFormat::Q4K, + HIDDEN as u32, + wide as u32, + INTERMEDIATE as u32, + ), + ..fixture.q4k_expert() + }; + let err = fixture + .operation(expert, ExpertKernel::IncumbentQ4kQ8k) + .validate() + .unwrap_err(); + assert!( + matches!( + err, + ExecutionError::DimensionMismatch { axis, expected, found, .. } + if axis == Axis::InputWidth && expected == INTERMEDIATE_PADDED && found == wide + ), + "{err}" + ); +} + +#[test] +fn a_bank_input_that_is_not_whole_blocks_is_refused_at_execution() { + // The activation-side counterpart. The incumbent guards this by falling + // back to its f32 path; VINDEX3 refuses, because a silent change of kernel + // is a silent change of answer. + // + // It fires at execution rather than at bind because the bank input is a + // property of the token, not of the binding. + let ragged = HIDDEN - 1; + let gate_up = q4k_bytes(FUSED_PROJECTION_HALVES * INTERMEDIATE, HIDDEN, GATE_UP_SEED); + let down = q4k_bytes(HIDDEN, INTERMEDIATE_PADDED, DOWN_SEED); + let router = f32_bytes(&vec![1.0f32; POPULATION * ragged]); + let operation = BoundMoeOperation { + router: BoundRouter { + weight: direct( + ROUTER_REGION_SET, + &router, + RegionFormat::F32, + POPULATION as u32, + ragged as u32, + ), + top_k: TOP_K, + selected_weight: MoeTopKWeightPolicy::RenormalizedSoftmax, + scaling: BoundExpertScaling::None, + kernel: RouterKernel::default(), + }, + transforms: Vec::new(), + banks: vec![BoundBankOperation { + bank: BankCoordinate::new(LAYER, BANK_ID), + experts: vec![BoundExpert { + expert_id: 0, + projection: BoundProjection::Fused { + gate_up: direct( + GATE_UP_REGION_SET, + &gate_up, + RegionFormat::Q4K, + (FUSED_PROJECTION_HALVES * INTERMEDIATE) as u32, + ragged as u32, + ), + }, + down: column_prefix( + DOWN_REGION_SET, + &down, + RegionFormat::Q4K, + ragged as u32, + INTERMEDIATE_PADDED as u32, + INTERMEDIATE as u32, + ), + }], + intermediate_dim: INTERMEDIATE, + hidden_dim: ragged, + activation: ACTIVATION, + kernel: ExpertKernel::IncumbentQ4kQ8k, + }], + reduction: BoundReduction::WeightedSum, + residual_dim: ragged, + }; + + let input = vec![0.1f32; ragged]; + let err = execute_traced(&operation, MoeInputs::shared(&input)).unwrap_err(); + assert!( + matches!( + err, + ExecutionError::KernelOperandUnsuitable { + reason: OperandUnsuitability::BlockAlignment { found, .. }, + .. + } if found == ragged + ), + "{err}" + ); +} + +// ── The session ──────────────────────────────────────────────────────────── + +#[test] +fn every_expert_in_a_bank_reads_the_one_quantised_activation() { + // Two experts over identical weights, both selected. Each must reproduce + // the incumbent's output for the *bank's* input — which is what the shared + // Q8_K session buys, and what a per-expert quantisation would only + // accidentally match. + let fixture = Fixture::new(); + let mut operation = fixture.operation(fixture.q4k_expert(), ExpertKernel::IncumbentQ4kQ8k); + let second = BoundExpert { + expert_id: 1, + ..fixture.q4k_expert() + }; + operation.banks[0].experts.push(second); + operation.router.top_k = POPULATION; + operation.validate().expect("both experts bind"); + + let (_, trace) = execute_traced(&operation, MoeInputs::shared(&fixture.input)) + .expect("the fixture executes"); + let incumbent = fixture.incumbent_expert_output(); + assert_eq!(trace.expert_outputs.len(), POPULATION); + for output in &trace.expert_outputs { + assert_eq!( + output.values, incumbent, + "expert {} read a different activation", + output.expert_id + ); + } +} + +// ── The kernel checks its own operands, even unvalidated ─────────────────── +// +// `execute` does not call `validate` — re-checking a binding per token is the +// resolution creep the bound object exists to prevent. So the kernel's operand +// checks have to hold on their own, and these reach them by executing an +// operation that was never validated. Every one of them would otherwise be a +// stride read at the wrong offset producing well-shaped noise. + +/// Run without validating first, and return the refusal. +fn execute_unvalidated(fixture: &Fixture, expert: BoundExpert<'_>) -> ExecutionError { + let operation = fixture.operation(expert, ExpertKernel::IncumbentQ4kQ8k); + execute_traced(&operation, MoeInputs::shared(&fixture.input)) + .expect_err("the kernel should refuse this operand") +} + +#[test] +fn a_gate_up_slab_of_the_wrong_height_is_refused_by_the_kernel() { + // A fused region holds both halves. One that is only `INTERMEDIATE` rows + // tall would have the kernel read the gate half as gate+up, and the up half + // from past the end of the region. + let fixture = Fixture::new(); + let bytes = q4k_bytes(INTERMEDIATE, HIDDEN, GATE_UP_SEED); + let err = execute_unvalidated( + &fixture, + BoundExpert { + expert_id: 0, + projection: BoundProjection::Fused { + gate_up: direct( + GATE_UP_REGION_SET, + &bytes, + RegionFormat::Q4K, + INTERMEDIATE as u32, + HIDDEN as u32, + ), + }, + ..fixture.q4k_expert() + }, + ); + assert!( + matches!( + err, + ExecutionError::DimensionMismatch { axis, expected, found, .. } + if axis == Axis::Rows + && expected == FUSED_PROJECTION_HALVES * INTERMEDIATE + && found == INTERMEDIATE + ), + "{err}" + ); +} + +#[test] +fn a_gate_up_slab_that_contracts_the_wrong_width_is_refused_by_the_kernel() { + // The kernel strides gate_up by `hidden / block` super-blocks per row, so a + // region storing a different width is read at the wrong offset from row one. + let fixture = Fixture::new(); + let narrow = HIDDEN + Q4K_BLOCK_ELEMS; + let bytes = q4k_bytes(FUSED_PROJECTION_HALVES * INTERMEDIATE, narrow, GATE_UP_SEED); + let err = execute_unvalidated( + &fixture, + BoundExpert { + expert_id: 0, + projection: BoundProjection::Fused { + gate_up: direct( + GATE_UP_REGION_SET, + &bytes, + RegionFormat::Q4K, + (FUSED_PROJECTION_HALVES * INTERMEDIATE) as u32, + narrow as u32, + ), + }, + ..fixture.q4k_expert() + }, + ); + assert!( + matches!( + err, + ExecutionError::DimensionMismatch { axis, expected, found, .. } + if axis == Axis::InputWidth && expected == HIDDEN && found == narrow + ), + "{err}" + ); +} + +#[test] +fn a_down_region_of_the_wrong_height_is_refused_by_the_kernel() { + // `down` contracts the intermediate axis away and produces the bank's + // output width. A shorter one silently produces a shorter residual delta. + let fixture = Fixture::new(); + let short = HIDDEN - Q4K_BLOCK_ELEMS / 2; + let bytes = q4k_bytes(short, INTERMEDIATE_PADDED, DOWN_SEED); + let err = execute_unvalidated( + &fixture, + BoundExpert { + expert_id: 0, + down: column_prefix( + DOWN_REGION_SET, + &bytes, + RegionFormat::Q4K, + short as u32, + INTERMEDIATE_PADDED as u32, + INTERMEDIATE as u32, + ), + ..fixture.q4k_expert() + }, + ); + assert!( + matches!( + err, + ExecutionError::DimensionMismatch { axis, expected, found, .. } + if axis == Axis::OutputWidth && expected == HIDDEN && found == short + ), + "{err}" + ); +} + +#[test] +fn a_down_region_exposing_the_wrong_live_width_is_refused_by_the_kernel() { + // Stored correctly, but the role claims more live columns than the bank + // means. The kernel passes `intermediate` to the incumbent as the count of + // activation values to compute, so this is the difference between the + // padding being inert and being read as data. + let fixture = Fixture::new(); + let wrong = INTERMEDIATE + 1; + let err = execute_unvalidated( + &fixture, + BoundExpert { + expert_id: 0, + down: column_prefix( + DOWN_REGION_SET, + &fixture.down, + RegionFormat::Q4K, + HIDDEN as u32, + INTERMEDIATE_PADDED as u32, + wrong as u32, + ), + ..fixture.q4k_expert() + }, + ); + assert!( + matches!( + err, + ExecutionError::DimensionMismatch { axis, expected, found, .. } + if axis == Axis::Columns && expected == INTERMEDIATE && found == wrong + ), + "{err}" + ); +} + +#[test] +fn a_decomposed_projection_is_refused_at_execution_too() { + let fixture = Fixture::new(); + let half = fixture.gate_up.len() / FUSED_PROJECTION_HALVES; + let err = execute_unvalidated( + &fixture, + BoundExpert { + expert_id: 0, + projection: BoundProjection::Decomposed { + gate: direct( + GATE_UP_REGION_SET, + &fixture.gate_up[..half], + RegionFormat::Q4K, + INTERMEDIATE as u32, + HIDDEN as u32, + ), + up: direct( + GATE_UP_REGION_SET, + &fixture.gate_up[half..], + RegionFormat::Q4K, + INTERMEDIATE as u32, + HIDDEN as u32, + ), + }, + ..fixture.q4k_expert() + }, + ); + assert!( + matches!( + err, + ExecutionError::KernelOperandUnsuitable { + reason: OperandUnsuitability::Arrangement { .. }, + .. + } + ), + "{err}" + ); +} diff --git a/crates/larql-vindex/src/runtime/tests/inputs.rs b/crates/larql-vindex/src/runtime/tests/inputs.rs new file mode 100644 index 000000000..63418a881 --- /dev/null +++ b/crates/larql-vindex/src/runtime/tests/inputs.rs @@ -0,0 +1,163 @@ +//! Colocated tests for `inputs` — routing on a vector the experts never see. +//! +//! Gemma applies a router-specific norm, scale and scalar on top of the expert +//! input, so its router scores a different vector. Before this existed the +//! operation would have scored on the expert input, selected different experts, +//! and produced a wrong answer with no shape error anywhere. + +use super::operation::direct; +use crate::runtime::execute::{execute, execute_traced}; +use crate::runtime::inputs::MoeInputs; + +const RESIDUAL_WIDTH: usize = 6; + +fn bank_input() -> Vec { + vec![0.4, -0.3, 0.9, 0.1, -0.7, 0.2] +} + +/// A router input that is emphatically not the bank input — the shape a +/// learned router norm and scale produce. +fn router_input() -> Vec { + vec![-0.9, 0.8, -0.2, 0.6, 0.5, -0.4] +} + +// ── Construction ─────────────────────────────────────────────────────────── + +#[test] +fn shared_inputs_point_both_reads_at_one_vector() { + let x = bank_input(); + let inputs = MoeInputs::shared(&x); + assert_eq!(inputs.bank, inputs.router); + assert!(inputs.is_shared()); +} + +#[test] +fn split_inputs_keep_the_two_apart() { + let bank = bank_input(); + let router = router_input(); + let inputs = MoeInputs::split(&bank, &router); + assert_eq!(inputs.bank, bank.as_slice()); + assert_eq!(inputs.router, router.as_slice()); + assert!(!inputs.is_shared()); +} + +#[test] +fn sharing_is_decided_by_value_not_by_pointer() { + // Two separately-computed vectors that happen to be equal are shared for + // every purpose that matters, and a model whose router norm is the + // identity produces exactly that. + let bank = bank_input(); + let copy = bank_input(); + assert!(MoeInputs::split(&bank, ©).is_shared()); +} + +#[test] +fn inputs_describe_which_shape_they_are() { + let bank = bank_input(); + let router = router_input(); + assert!(MoeInputs::shared(&bank).describe().contains("shared")); + assert!(MoeInputs::split(&bank, &router) + .describe() + .contains("split")); +} + +// ── The behaviour the split exists for ───────────────────────────────────── + +/// An operation whose router selects expert `e` by input component `e`. +/// +/// `direct()`'s router uses ascending weights, which makes its selection +/// almost input-insensitive — the high-index rows dominate whatever it is +/// given. A first version of the test below used it and saw scores move while +/// the selection did not, which proves nothing about routing. A one-hot router +/// makes the selection readable by hand: top-k over the input's own largest +/// components. +fn selection_follows_input() -> crate::runtime::operation::BoundMoeOperation<'static> { + use super::support::matrix; + let mut op = direct(); + let population = op.banks[0].experts.len(); + let hidden = op.residual_dim; + let mut weights = vec![0.0f32; population * hidden]; + for (e, row) in weights.chunks_mut(hidden).enumerate() { + row[e] = 1.0; + } + op.router.weight = matrix("router", population as u32, hidden as u32, &weights); + op +} + +#[test] +fn the_router_input_decides_the_selection() { + // Same experts, same bank input, different router input. If routing read + // the bank input, these would be identical — which is precisely the bug a + // single-input model would have shipped. + let op = selection_follows_input(); + let bank = bank_input(); + let router = router_input(); + + let (_, shared) = execute_traced(&op, MoeInputs::shared(&bank)).unwrap(); + let (_, split) = execute_traced(&op, MoeInputs::split(&bank, &router)).unwrap(); + + assert_ne!( + shared.router_scores, split.router_scores, + "the router must score the vector it was given" + ); + // bank = [0.4, -0.3, 0.9, 0.1, ...] → components 2 then 0 + // router = [-0.9, 0.8, -0.2, 0.6, ...] → components 1 then 3 + assert_eq!(shared.selected_ids(), vec![2, 0]); + assert_eq!(split.selected_ids(), vec![1, 3]); +} + +#[test] +fn the_bank_input_decides_the_expert_arithmetic() { + // The mirror: hold the router input fixed, change what the experts see. + // Selection must not move; the outputs must. + let op = direct(); + let router = router_input(); + let bank_a = bank_input(); + let bank_b: Vec = bank_a.iter().map(|v| v * 0.5).collect(); + + let (out_a, trace_a) = execute_traced(&op, MoeInputs::split(&bank_a, &router)).unwrap(); + let (out_b, trace_b) = execute_traced(&op, MoeInputs::split(&bank_b, &router)).unwrap(); + + assert_eq!( + trace_a.selected_ids(), + trace_b.selected_ids(), + "selection is a function of the router input alone" + ); + assert_eq!(trace_a.gate_weights(), trace_b.gate_weights()); + assert_ne!(out_a, out_b, "the experts read the bank input"); +} + +#[test] +fn a_shared_input_is_exactly_a_split_of_one_vector_with_itself() { + // No special-casing: the shared path must not be a different code path. + let op = direct(); + let x = bank_input(); + assert_eq!( + execute(&op, MoeInputs::shared(&x)).unwrap(), + execute(&op, MoeInputs::split(&x, &x)).unwrap() + ); +} + +#[test] +fn the_recorded_router_scores_come_from_the_router_input() { + // Verifiable independently: scoring the router input through a second + // execution as a shared input must reproduce the same scores. + let op = direct(); + let bank = bank_input(); + let router = router_input(); + + let (_, split) = execute_traced(&op, MoeInputs::split(&bank, &router)).unwrap(); + let (_, router_as_shared) = execute_traced(&op, MoeInputs::shared(&router)).unwrap(); + assert_eq!(split.router_scores, router_as_shared.router_scores); + assert_eq!(split.selected_ids(), router_as_shared.selected_ids()); +} + +// ── Widths ───────────────────────────────────────────────────────────────── + +#[test] +fn the_bank_input_width_is_the_one_checked_against_the_residual_dim() { + let op = direct(); + let router = router_input(); + let short = vec![0.1f32; RESIDUAL_WIDTH - 1]; + assert!(execute(&op, MoeInputs::split(&short, &router)).is_err()); +} diff --git a/crates/larql-vindex/src/runtime/tests/kernels.rs b/crates/larql-vindex/src/runtime/tests/kernels.rs new file mode 100644 index 000000000..4a1b9e083 --- /dev/null +++ b/crates/larql-vindex/src/runtime/tests/kernels.rs @@ -0,0 +1,164 @@ +//! Colocated tests for `kernels` — the reference primitives. + +use larql_compute::Activation; + +use crate::runtime::kernels::{activate, dot, renormalize, softmax, top_k}; + +#[test] +fn a_dot_product_contracts_two_vectors() { + assert_eq!(dot(&[1.0, 2.0, 3.0], &[4.0, 5.0, 6.0]), 32.0); + assert_eq!(dot(&[], &[]), 0.0); +} + +#[test] +fn a_dot_product_stops_at_the_shorter_operand() { + // Zip semantics, stated so a future rewrite to indexing keeps them. + assert_eq!(dot(&[1.0, 2.0, 3.0], &[1.0, 1.0]), 3.0); +} + +// ── Softmax ──────────────────────────────────────────────────────────────── + +#[test] +fn softmax_produces_a_distribution() { + let mut v = vec![1.0, 2.0, 3.0]; + softmax(&mut v); + assert!((v.iter().sum::() - 1.0).abs() < 1e-6); + assert!(v[2] > v[1] && v[1] > v[0], "order preserved"); +} + +#[test] +fn softmax_is_shift_invariant() { + let mut a = vec![1.0, 2.0, 3.0]; + let mut b = vec![101.0, 102.0, 103.0]; + softmax(&mut a); + softmax(&mut b); + for (x, y) in a.iter().zip(&b) { + assert!((x - y).abs() < 1e-6, "{x} vs {y}"); + } +} + +#[test] +fn softmax_survives_large_inputs_without_overflowing() { + // The reason for the max-shift: exp(1000) is inf, and inf/inf is NaN. + let mut v = vec![1000.0, 1000.5]; + softmax(&mut v); + assert!(v.iter().all(|x| x.is_finite()), "{v:?}"); + assert!((v.iter().sum::() - 1.0).abs() < 1e-6); +} + +#[test] +fn softmax_of_nothing_does_nothing() { + let mut v: Vec = Vec::new(); + softmax(&mut v); + assert!(v.is_empty()); +} + +// ── Top-k ────────────────────────────────────────────────────────────────── + +#[test] +fn top_k_returns_the_highest_scores_descending() { + let picked = top_k(&[0.1, 0.5, 0.2, 0.9], 2); + assert_eq!(picked, vec![(3, 0.9), (1, 0.5)]); +} + +#[test] +fn top_k_breaks_ties_toward_the_lower_index() { + // Deterministic on purpose — an oracle that disagrees with itself between + // runs cannot be used to check anything. + let picked = top_k(&[0.5, 0.5, 0.5], 2); + assert_eq!( + picked.iter().map(|(i, _)| *i).collect::>(), + vec![0, 1] + ); +} + +#[test] +fn top_k_larger_than_the_population_returns_everything() { + assert_eq!(top_k(&[0.2, 0.8], 10).len(), 2); +} + +#[test] +fn top_k_of_zero_selects_nothing() { + assert!(top_k(&[0.2, 0.8], 0).is_empty()); +} + +// ── Renormalisation ──────────────────────────────────────────────────────── + +#[test] +fn renormalisation_makes_weights_sum_to_one() { + let mut w = vec![0.3, 0.1]; + renormalize(&mut w); + assert!((w.iter().sum::() - 1.0).abs() < 1e-6); + assert!((w[0] - 0.75).abs() < 1e-6); +} + +#[test] +fn renormalising_a_zero_sum_leaves_it_alone() { + // Dividing by zero here would turn a degenerate routing into NaNs that + // propagate through the whole residual stream. + let mut w = vec![0.0, 0.0]; + renormalize(&mut w); + assert_eq!(w, vec![0.0, 0.0]); +} + +// ── Activations ──────────────────────────────────────────────────────────── + +#[test] +fn every_activation_fixes_zero_at_zero() { + for activation in [ + Activation::Silu, + Activation::GeluTanh, + Activation::GeluExact, + Activation::ReLU, + ] { + let mut v = vec![0.0f32]; + activate(activation, &mut v); + assert!(v[0].abs() < 1e-6, "{activation:?} maps 0 to {}", v[0]); + } +} + +#[test] +fn silu_matches_its_definition() { + let mut v = vec![1.0f32]; + activate(Activation::Silu, &mut v); + assert!((v[0] - 1.0 / (1.0 + (-1.0f32).exp())).abs() < 1e-6); +} + +#[test] +fn relu_clamps_negatives() { + let mut v = vec![-2.0, 3.0]; + activate(Activation::ReLU, &mut v); + assert_eq!(v, vec![0.0, 3.0]); +} + +#[test] +fn the_two_gelus_agree_closely() { + // The tanh approximation and the erf form differ by well under a percent; + // a much larger gap would mean one of them is wrong. + for x in [-2.0f32, -0.5, 0.5, 2.0] { + let mut tanh_form = vec![x]; + let mut exact_form = vec![x]; + activate(Activation::GeluTanh, &mut tanh_form); + activate(Activation::GeluExact, &mut exact_form); + assert!( + (tanh_form[0] - exact_form[0]).abs() < 1e-2, + "at {x}: {} vs {}", + tanh_form[0], + exact_form[0] + ); + } +} + +#[test] +fn gelu_is_monotone_over_the_range_that_matters() { + let mut v: Vec = (0..20).map(|i| i as f32 * 0.25).collect(); + activate(Activation::GeluTanh, &mut v); + assert!(v.windows(2).all(|w| w[1] >= w[0]), "{v:?}"); +} + +#[test] +fn activations_apply_elementwise_across_a_slice() { + let mut v = vec![-1.0, 0.0, 1.0]; + activate(Activation::ReLU, &mut v); + assert_eq!(v, vec![0.0, 0.0, 1.0]); +} diff --git a/crates/larql-vindex/src/runtime/tests/mod.rs b/crates/larql-vindex/src/runtime/tests/mod.rs new file mode 100644 index 000000000..4cfed48f1 --- /dev/null +++ b/crates/larql-vindex/src/runtime/tests/mod.rs @@ -0,0 +1,30 @@ +//! Tests for the reference runtime, one file per module under test. +//! +//! Kept beside the runtime rather than in `crates/larql-vindex/tests/` because +//! they exercise module-private behaviour — a bound tensor's stored bytes +//! behind a view, an operand a kernel must refuse — and an integration test +//! can only reach the public surface. Kept in a folder rather than interleaved +//! as `*_tests.rs` siblings because the runtime's own file list is what a +//! reader scans to find the implementation. +//! +//! `support` is the shared builder module; everything else names the module it +//! covers. + +pub mod support; + +mod addressing; +mod bank; +mod execute; +mod expert_kernel; +mod inputs; +mod kernels; +mod operation; +mod placement; +mod projection; +mod reduction; +mod residency; +mod router; +mod tensor; +mod tie; +mod transform; +mod verdict; diff --git a/crates/larql-vindex/src/runtime/tests/operation.rs b/crates/larql-vindex/src/runtime/tests/operation.rs new file mode 100644 index 000000000..54af0b2e4 --- /dev/null +++ b/crates/larql-vindex/src/runtime/tests/operation.rs @@ -0,0 +1,225 @@ +//! Colocated tests for `operation` — shape agreement across bound operands. +//! +//! `validate` is the load-path gate. Everything it catches would otherwise +//! surface mid-token as an indexing failure, or not at all. + +use crate::format::capability::coordinate::BankCoordinate; +use crate::format::lyrw2::region_role::RegionRole; +use larql_compute::{Activation, MoeTopKWeightPolicy}; + +use super::support::ascending; +use crate::runtime::bank::{BoundBankOperation, BoundExpert}; +use crate::runtime::error::ExecutionError; +use crate::runtime::expert_kernel::ExpertKernel; +use crate::runtime::operation::BoundMoeOperation; +use crate::runtime::projection::BoundProjection; +use crate::runtime::reduction::BoundReduction; +use crate::runtime::router::{BoundExpertScaling, BoundRouter, RouterKernel}; +use crate::runtime::transform::{BoundTransform, TransformStage}; + +const RESIDUAL: u32 = 6; +const LATENT: u32 = 3; +const INTERMEDIATE: u32 = 2; +const POPULATION: u32 = 4; +const TOP_K: usize = 2; +const ROUTER: &str = "router"; +const LATENT_IN: &str = "latent_in"; +const LATENT_OUT: &str = "latent_out"; + +fn bank(hidden: u32) -> BoundBankOperation<'static> { + BoundBankOperation { + bank: BankCoordinate::new(0, 0), + experts: (0..POPULATION) + .map(|expert_id| BoundExpert { + expert_id, + projection: BoundProjection::Decomposed { + gate: ascending(&RegionRole::Gate.name(), INTERMEDIATE, hidden), + up: ascending(&RegionRole::Up.name(), INTERMEDIATE, hidden), + }, + down: ascending(&RegionRole::Down.name(), hidden, INTERMEDIATE), + }) + .collect(), + intermediate_dim: INTERMEDIATE as usize, + hidden_dim: hidden as usize, + activation: Activation::Silu, + kernel: ExpertKernel::default(), + } +} + +fn router(hidden: u32) -> BoundRouter<'static> { + BoundRouter { + weight: ascending(ROUTER, POPULATION, hidden), + top_k: TOP_K, + selected_weight: MoeTopKWeightPolicy::RenormalizedSoftmax, + scaling: BoundExpertScaling::None, + kernel: RouterKernel::default(), + } +} + +/// A direct operation: the bank reads the residual width. +pub fn direct() -> BoundMoeOperation<'static> { + BoundMoeOperation { + router: router(RESIDUAL), + transforms: Vec::new(), + banks: vec![bank(RESIDUAL)], + reduction: BoundReduction::WeightedSum, + residual_dim: RESIDUAL as usize, + } +} + +/// A latent operation: the bank reads a narrower routed width. +pub fn latent() -> BoundMoeOperation<'static> { + BoundMoeOperation { + router: router(LATENT), + transforms: vec![ + BoundTransform { + stage: TransformStage::RoutedInput, + weight: ascending(LATENT_IN, LATENT, RESIDUAL), + }, + BoundTransform { + stage: TransformStage::RoutedOutput, + weight: ascending(LATENT_OUT, RESIDUAL, LATENT), + }, + ], + banks: vec![bank(LATENT)], + reduction: BoundReduction::WeightedSum, + residual_dim: RESIDUAL as usize, + } +} + +// ── Bank input width ─────────────────────────────────────────────────────── + +#[test] +fn a_direct_operation_runs_its_banks_at_the_residual_width() { + let op = direct(); + assert_eq!(op.bank_input_dim(), RESIDUAL as usize); + assert!(!op.is_latent()); +} + +#[test] +fn a_latent_operation_runs_its_banks_at_the_projected_width() { + let op = latent(); + assert_eq!(op.bank_input_dim(), LATENT as usize); + assert!(op.is_latent()); + assert_eq!(op.residual_dim, RESIDUAL as usize, "residual is unchanged"); +} + +#[test] +fn the_last_routed_input_transform_fixes_the_bank_width() { + // A chained input stage: the width the banks see is the final projection's + // output, not the first one's. + let mut op = latent(); + op.transforms.insert( + 0, + BoundTransform { + stage: TransformStage::RoutedInput, + weight: ascending(LATENT_IN, RESIDUAL, RESIDUAL), + }, + ); + assert_eq!(op.bank_input_dim(), LATENT as usize); +} + +// ── Validation ───────────────────────────────────────────────────────────── + +#[test] +fn a_consistent_direct_operation_validates() { + direct().validate().unwrap(); +} + +#[test] +fn a_consistent_latent_operation_validates() { + latent().validate().unwrap(); +} + +#[test] +fn a_bank_that_expects_the_residual_width_behind_a_latent_transform_is_refused() { + // The mistake a direct-MoE assumption produces: experts sized for the + // residual, fed the projected vector. + let mut op = latent(); + op.banks = vec![bank(RESIDUAL)]; + let err = op.validate().unwrap_err(); + let ExecutionError::DimensionMismatch { + expected, found, .. + } = err + else { + panic!("expected a dimension mismatch, got {err}"); + }; + assert_eq!((expected, found), (LATENT as usize, RESIDUAL as usize)); +} + +#[test] +fn a_bank_holding_only_part_of_the_population_is_valid() { + // Sharding. The router addresses the whole population; this bank carries a + // subset of it. An earlier `validate` demanded equality and would have + // refused every expert-server slice. + let mut op = direct(); + op.banks[0].experts.truncate(2); + op.banks[0].experts[0].expert_id = 2; + op.banks[0].experts[1].expert_id = 3; + op.validate().expect("a shard is a legitimate operation"); + assert!(!op.holds_full_population()); +} + +#[test] +fn a_bank_holding_everything_reports_full_population() { + let op = direct(); + op.validate().unwrap(); + assert!(op.holds_full_population()); +} + +#[test] +fn an_expert_the_router_cannot_address_is_refused() { + // The genuine fault the equality check was reaching for: an expert with an + // id past the router's rows can never be selected, so its presence means + // the router and the bank disagree about which population this is. + let mut op = direct(); + op.banks[0].experts[0].expert_id = POPULATION + 5; + let err = op.validate().unwrap_err(); + assert!(matches!(err, ExecutionError::ExpertOutOfRange { .. })); +} + +#[test] +fn a_router_addressing_more_experts_than_are_bound_is_valid() { + // Same shape as sharding, stated from the router's side. + let mut op = direct(); + op.router = BoundRouter { + weight: ascending(ROUTER, POPULATION + 4, RESIDUAL), + ..op.router + }; + op.validate().unwrap(); + assert!(!op.holds_full_population()); +} + +#[test] +fn a_router_scoring_at_the_wrong_width_is_refused() { + let mut op = direct(); + op.router = BoundRouter { + weight: ascending(ROUTER, POPULATION, RESIDUAL + 1), + ..op.router + }; + assert!(op.validate().is_err()); +} + +#[test] +fn an_operation_with_no_banks_validates_vacuously() { + // Nothing to disagree about. Whether an empty operation is *useful* is a + // planning question, not a shape one. + let mut op = direct(); + op.banks.clear(); + op.validate().unwrap(); +} + +// ── Description ──────────────────────────────────────────────────────────── + +#[test] +fn a_direct_operation_describes_itself_as_direct() { + let text = direct().describe(); + assert!(text.contains("direct"), "{text}"); + assert!(text.contains(BoundReduction::WeightedSum.name()), "{text}"); + assert!(text.contains(&format!("top-{TOP_K}")), "{text}"); +} + +#[test] +fn a_latent_operation_describes_itself_as_latent() { + assert!(latent().describe().contains("latent")); +} diff --git a/crates/larql-vindex/src/runtime/tests/placement.rs b/crates/larql-vindex/src/runtime/tests/placement.rs new file mode 100644 index 000000000..c8a364447 --- /dev/null +++ b/crates/larql-vindex/src/runtime/tests/placement.rs @@ -0,0 +1,175 @@ +//! What happens when routing selects an expert this bank does not hold. +//! +//! Relaxing `validate` to permit shards created an execution requirement: +//! +//! > If routing selects an expert absent from the bound local bank, execution +//! > must produce an explicit missing-expert result — not skip it, not +//! > renormalise around it, not substitute another expert. +//! +//! Each of those three would produce a token quietly missing part of its FFN +//! contribution, and it would read as entirely reasonable output. The request +//! has to survive to the caller so that placement can satisfy it; remote +//! fetching will later answer exactly this without touching router semantics. + +use super::operation::direct; +use super::support::matrix; +use crate::runtime::execute::{execute, execute_traced}; +use crate::runtime::inputs::MoeInputs; +use crate::runtime::operation::BoundMoeOperation; +use crate::runtime::ExecutionError; + +const ROUTER: &str = "router"; + +/// A router that selects expert `e` by input component `e`, so the test +/// controls the selection by choosing the input. +fn addressable_router(op: &mut BoundMoeOperation<'static>) { + let population = op.router.population(); + let hidden = op.residual_dim; + let mut weights = vec![0.0f32; population * hidden]; + for (e, row) in weights.chunks_mut(hidden).enumerate() { + row[e] = 1.0; + } + op.router.weight = matrix(ROUTER, population as u32, hidden as u32, &weights); +} + +/// A shard holding experts 0 and 1 only, out of a population of 4. +fn shard() -> BoundMoeOperation<'static> { + let mut op = direct(); + addressable_router(&mut op); + op.banks[0].experts.truncate(2); + op.banks[0].experts[0].expert_id = 0; + op.banks[0].experts[1].expert_id = 1; + op.router.top_k = 1; + op +} + +// ── The shard is legitimate ──────────────────────────────────────────────── + +#[test] +fn a_shard_validates_and_reports_that_it_is_partial() { + let op = shard(); + op.validate().expect("a shard is a valid operation"); + assert!(!op.holds_full_population()); + assert_eq!(op.banks[0].population(), 2); + assert_eq!(op.router.population(), 4); +} + +#[test] +fn a_shard_serves_a_selection_it_holds() { + // Component 1 largest → expert 1, which this shard has. + let op = shard(); + let input = vec![0.0, 5.0, 0.0, 0.0, 0.0, 0.0]; + let (out, trace) = execute_traced(&op, MoeInputs::shared(&input)).unwrap(); + assert_eq!(trace.selected_ids(), vec![1]); + assert!(out.iter().any(|v| *v != 0.0)); +} + +// ── The requirement ──────────────────────────────────────────────────────── + +#[test] +fn a_selection_the_shard_lacks_is_reported_not_skipped() { + // Component 3 largest → expert 3, which this shard does not hold. + let op = shard(); + let input = vec![0.0, 0.0, 0.0, 5.0, 0.0, 0.0]; + let err = execute(&op, MoeInputs::shared(&input)).unwrap_err(); + let ExecutionError::SelectedExpertNotResident { + expert, + bank, + resident, + population, + } = &err + else { + panic!("expected a placement result, got {err}"); + }; + assert_eq!(*expert, 3, "the report must name the selected expert"); + assert_eq!((*resident, *population), (2, 4)); + assert!( + bank.contains("bank"), + "must name the bank coordinate: {bank}" + ); +} + +#[test] +fn the_report_names_the_layer_and_bank_it_looked_in() { + // A whole-model run has many banks; "expert 3 is missing" without saying + // from where sends the reader hunting. + let mut op = shard(); + op.banks[0].bank = crate::format::capability::coordinate::BankCoordinate::new(17, 2); + let input = vec![0.0, 0.0, 0.0, 5.0, 0.0, 0.0]; + let text = execute(&op, MoeInputs::shared(&input)) + .unwrap_err() + .to_string(); + assert!(text.contains("layer 17"), "{text}"); + assert!(text.contains("bank 2"), "{text}"); + assert!(text.contains('3'), "{text}"); +} + +#[test] +fn a_missing_expert_is_distinct_from_an_unaddressable_one() { + // Two different repairs: fetch the operand, versus fix the index. The + // error types must not collapse, or an operator cannot tell which. + let op = shard(); + let input = vec![0.0, 0.0, 0.0, 5.0, 0.0, 0.0]; + let missing = execute(&op, MoeInputs::shared(&input)).unwrap_err(); + assert!(matches!( + missing, + ExecutionError::SelectedExpertNotResident { .. } + )); + + let mut broken = shard(); + broken.banks[0].experts[0].expert_id = 99; + let unaddressable = broken.validate().unwrap_err(); + assert!(matches!( + unaddressable, + ExecutionError::ExpertOutOfRange { .. } + )); +} + +#[test] +fn nothing_is_reduced_when_a_selected_expert_is_absent() { + // The failure mode being excluded: producing a partial delta from whatever + // experts happened to be resident. There is no partial answer here. + let op = shard(); + let input = vec![0.0, 0.0, 0.0, 5.0, 0.0, 0.0]; + assert!(execute(&op, MoeInputs::shared(&input)).is_err()); +} + +#[test] +fn a_partly_resident_selection_still_refuses_rather_than_contributing_what_it_has() { + // top-2 where one expert is resident and one is not. Renormalising over + // the survivor would produce a confident, wrong, correctly-shaped token. + let mut op = shard(); + op.router.top_k = 2; + let input = vec![0.0, 5.0, 0.0, 4.0, 0.0, 0.0]; // experts 1 (held) then 3 (not) + let err = execute(&op, MoeInputs::shared(&input)).unwrap_err(); + assert!(matches!( + err, + ExecutionError::SelectedExpertNotResident { expert: 3, .. } + )); +} + +#[test] +fn a_bank_can_be_asked_whether_it_holds_an_expert_without_failing() { + // Placement needs to ask before executing; asking must not be an error. + let op = shard(); + assert!(op.banks[0].holds(1)); + assert!(!op.banks[0].holds(3)); +} + +#[test] +fn routing_is_unchanged_by_what_the_shard_happens_to_hold() { + // The selection is a property of the router and its input alone. If a + // shard could influence it, two shards of one model would route + // differently and the model would depend on its own sharding. + let full = { + let mut op = direct(); + addressable_router(&mut op); + op.router.top_k = 1; + op + }; + let input = vec![0.0, 5.0, 0.0, 0.0, 0.0, 0.0]; + let (_, full_trace) = execute_traced(&full, MoeInputs::shared(&input)).unwrap(); + let (_, shard_trace) = execute_traced(&shard(), MoeInputs::shared(&input)).unwrap(); + assert_eq!(full_trace.selected_ids(), shard_trace.selected_ids()); + assert_eq!(full_trace.router_scores, shard_trace.router_scores); +} diff --git a/crates/larql-vindex/src/runtime/tests/projection.rs b/crates/larql-vindex/src/runtime/tests/projection.rs new file mode 100644 index 000000000..a44e68061 --- /dev/null +++ b/crates/larql-vindex/src/runtime/tests/projection.rs @@ -0,0 +1,208 @@ +//! Colocated tests for `projection` — the fused/decomposed equivalence. +//! +//! The row-mapping tests here are the unit-level statement of what fixture A +//! asserts end to end. Both are worth having: the fixture proves the whole +//! layer agrees, these say precisely which row a given call returns. + +use crate::format::lyrw2::region_role::RegionRole; + +use super::support::{matrix, TEST_VARIANT}; +use crate::runtime::error::ExecutionError; +use crate::runtime::projection::{BoundProjection, ProjectionArrangement}; + +const INTERMEDIATE: u32 = 2; +const HIDDEN: u32 = 3; + +/// Gate rows [1,2,3] and [4,5,6]; up rows [7,8,9] and [10,11,12]. +const GATE: [f32; 6] = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]; +const UP: [f32; 6] = [7.0, 8.0, 9.0, 10.0, 11.0, 12.0]; + +fn decomposed() -> BoundProjection<'static> { + BoundProjection::Decomposed { + gate: matrix(&RegionRole::Gate.name(), INTERMEDIATE, HIDDEN, &GATE), + up: matrix(&RegionRole::Up.name(), INTERMEDIATE, HIDDEN, &UP), + } +} + +fn fused() -> BoundProjection<'static> { + let mut values = GATE.to_vec(); + values.extend_from_slice(&UP); + BoundProjection::Fused { + gate_up: matrix( + &RegionRole::GateUpFused.name(), + INTERMEDIATE * 2, + HIDDEN, + &values, + ), + } +} + +fn row(projection: &BoundProjection<'_>, unit: usize, gate: bool) -> Vec { + let mut out = vec![0.0f32; HIDDEN as usize]; + if gate { + projection.gate_row_into(unit, &mut out).unwrap(); + } else { + projection.up_row_into(unit, &mut out).unwrap(); + } + out +} + +// ── Both arrangements yield the same rows ────────────────────────────────── + +#[test] +fn decomposed_and_fused_return_identical_gate_rows() { + for unit in 0..INTERMEDIATE as usize { + assert_eq!(row(&decomposed(), unit, true), row(&fused(), unit, true)); + } +} + +#[test] +fn decomposed_and_fused_return_identical_up_rows() { + for unit in 0..INTERMEDIATE as usize { + assert_eq!(row(&decomposed(), unit, false), row(&fused(), unit, false)); + } +} + +#[test] +fn the_fused_up_half_starts_after_every_gate_row() { + // Unit 1's up row is region row 3, not row 3-by-coincidence: with + // INTERMEDIATE = 2 an interleaved reading would also land on row 3 for + // unit 1, so this checks unit 0 too, where the two disagree. + assert_eq!(row(&fused(), 0, false), vec![7.0, 8.0, 9.0]); + assert_eq!(row(&fused(), 1, false), vec![10.0, 11.0, 12.0]); +} + +#[test] +fn gate_rows_are_the_first_half() { + assert_eq!(row(&fused(), 0, true), vec![1.0, 2.0, 3.0]); + assert_eq!(row(&fused(), 1, true), vec![4.0, 5.0, 6.0]); +} + +// ── Dimensions ───────────────────────────────────────────────────────────── + +#[test] +fn both_arrangements_report_the_same_dimensions() { + for projection in [decomposed(), fused()] { + assert_eq!(projection.intermediate_dim(), INTERMEDIATE as usize); + assert_eq!(projection.hidden_dim(), HIDDEN as usize); + } +} + +#[test] +fn a_fused_region_halves_its_row_count() { + // The property that makes an odd row count detectable. + let BoundProjection::Fused { gate_up } = fused() else { + panic!("expected a fused projection"); + }; + assert_eq!(gate_up.rows(), (INTERMEDIATE * 2) as usize); +} + +// ── Validation ───────────────────────────────────────────────────────────── + +#[test] +fn both_arrangements_validate_against_the_layer_shape() { + for projection in [decomposed(), fused()] { + projection + .validate(INTERMEDIATE as usize, HIDDEN as usize) + .unwrap(); + } +} + +#[test] +fn a_decomposed_region_bound_as_fused_is_refused_by_row_count() { + // The clearest signal that the arrangement was misread: a region with + // `intermediate` rows claiming to hold `2 * intermediate`. + let wrong = BoundProjection::Fused { + gate_up: matrix(&RegionRole::GateUpFused.name(), INTERMEDIATE, HIDDEN, &GATE), + }; + let err = wrong + .validate(INTERMEDIATE as usize, HIDDEN as usize) + .unwrap_err(); + let ExecutionError::DimensionMismatch { + expected, found, .. + } = err + else { + panic!("expected a dimension mismatch"); + }; + assert_eq!((expected, found), (4, 2)); +} + +#[test] +fn a_decomposed_pair_that_disagrees_on_width_is_refused() { + let wrong = BoundProjection::Decomposed { + gate: matrix(&RegionRole::Gate.name(), INTERMEDIATE, HIDDEN, &GATE), + up: matrix( + &RegionRole::Up.name(), + INTERMEDIATE, + 2, + &[1.0, 2.0, 3.0, 4.0], + ), + }; + assert!(wrong + .validate(INTERMEDIATE as usize, HIDDEN as usize) + .is_err()); +} + +#[test] +fn a_row_past_the_intermediate_width_is_refused() { + let mut out = vec![0.0f32; HIDDEN as usize]; + assert!(fused() + .up_row_into(INTERMEDIATE as usize, &mut out) + .is_err()); + assert!(decomposed() + .gate_row_into(INTERMEDIATE as usize, &mut out) + .is_err()); +} + +// ── Naming comes from the role registry ──────────────────────────────────── + +#[test] +fn arrangements_name_themselves_from_the_roles_they_store() { + assert_eq!( + ProjectionArrangement::Decomposed.name(), + format!("{}+{}", RegionRole::Gate.name(), RegionRole::Up.name()) + ); + assert_eq!( + ProjectionArrangement::Fused.name(), + RegionRole::GateUpFused.name() + ); +} + +#[test] +fn every_arrangement_declares_the_roles_it_needs() { + assert_eq!( + ProjectionArrangement::Decomposed.roles(), + vec![RegionRole::Gate, RegionRole::Up] + ); + assert_eq!( + ProjectionArrangement::Fused.roles(), + vec![RegionRole::GateUpFused] + ); + assert_eq!(ProjectionArrangement::ALL.len(), 2); +} + +#[test] +fn a_projection_reports_which_arrangement_it_is() { + assert_eq!( + decomposed().arrangement(), + ProjectionArrangement::Decomposed + ); + assert_eq!(fused().arrangement(), ProjectionArrangement::Fused); +} + +#[test] +fn describing_a_projection_names_its_arrangement_and_operands() { + let text = decomposed().describe(); + assert!( + text.contains(&ProjectionArrangement::Decomposed.name()), + "{text}" + ); + assert!(text.contains(&RegionRole::Gate.name()), "{text}"); + assert!(text.contains(TEST_VARIANT), "{text}"); + + let fused_text = fused().describe(); + assert!( + fused_text.contains(&RegionRole::GateUpFused.name()), + "{fused_text}" + ); +} diff --git a/crates/larql-vindex/src/runtime/tests/reduction.rs b/crates/larql-vindex/src/runtime/tests/reduction.rs new file mode 100644 index 000000000..970bbdb76 --- /dev/null +++ b/crates/larql-vindex/src/runtime/tests/reduction.rs @@ -0,0 +1,65 @@ +//! Colocated tests for `reduction`. + +use crate::runtime::axis::Axis; +use crate::runtime::error::ExecutionError; +use crate::runtime::reduction::BoundReduction; + +#[test] +fn a_weighted_sum_accumulates_scaled_contributions() { + let mut acc = vec![0.0f32; 3]; + BoundReduction::WeightedSum + .accumulate(&mut acc, &[1.0, 2.0, 3.0], 0.5) + .unwrap(); + BoundReduction::WeightedSum + .accumulate(&mut acc, &[2.0, 2.0, 2.0], 0.25) + .unwrap(); + assert_eq!(acc, vec![1.0, 1.5, 2.0]); +} + +#[test] +fn accumulation_adds_rather_than_replaces() { + // The bug this guards: a second expert overwriting the first's + // contribution produces output that looks entirely plausible. + let mut acc = vec![1.0f32; 2]; + BoundReduction::WeightedSum + .accumulate(&mut acc, &[1.0, 1.0], 1.0) + .unwrap(); + assert_eq!(acc, vec![2.0, 2.0]); +} + +#[test] +fn a_zero_weight_contributes_nothing() { + let mut acc = vec![5.0f32; 2]; + BoundReduction::WeightedSum + .accumulate(&mut acc, &[9.0, 9.0], 0.0) + .unwrap(); + assert_eq!(acc, vec![5.0, 5.0]); +} + +#[test] +fn a_width_disagreement_is_refused_naming_the_output_axis() { + let mut acc = vec![0.0f32; 3]; + let err = BoundReduction::WeightedSum + .accumulate(&mut acc, &[1.0, 2.0], 1.0) + .unwrap_err(); + let ExecutionError::DimensionMismatch { + axis, + expected, + found, + operand, + } = err + else { + panic!("expected a dimension mismatch"); + }; + assert_eq!(axis, Axis::OutputWidth); + assert_eq!((expected, found), (3, 2)); + assert!( + operand.contains(BoundReduction::WeightedSum.name()), + "{operand}" + ); +} + +#[test] +fn the_reduction_names_itself() { + assert_eq!(BoundReduction::WeightedSum.name(), "weighted_sum"); +} diff --git a/crates/larql-vindex/src/runtime/tests/residency.rs b/crates/larql-vindex/src/runtime/tests/residency.rs new file mode 100644 index 000000000..32d16ad56 --- /dev/null +++ b/crates/larql-vindex/src/runtime/tests/residency.rs @@ -0,0 +1,199 @@ +//! Colocated tests for `residency`. +//! +//! Page attribution is measurement code, and measurement code that is wrong +//! reports confidently. These pin the boundary behaviour the probe's headline +//! number depends on. + +use crate::runtime::residency::{account, ExpertRegion, PageSpan, ResidencyAccount}; + +const PAGE: usize = 4_096; + +fn expert(expert_id: u32, start: usize, len: usize) -> ExpertRegion { + ExpertRegion { + expert_id, + bytes: start..start + len, + } +} + +// ── Page spans ───────────────────────────────────────────────────────────── + +#[test] +fn a_range_inside_one_page_spans_that_page() { + let span = PageSpan::of(&(10..20), PAGE); + assert_eq!((span.first, span.end), (0, 1)); + assert_eq!(span.len(), 1); +} + +#[test] +fn a_range_ending_mid_page_still_occupies_it() { + // Rounding down here would under-report residency by one page per region. + let span = PageSpan::of(&(0..PAGE + 1), PAGE); + assert_eq!(span.end, 2); + assert_eq!(span.len(), 2); +} + +#[test] +fn a_page_aligned_range_does_not_claim_the_next_page() { + let span = PageSpan::of(&(0..PAGE), PAGE); + assert_eq!(span.end, 1); +} + +#[test] +fn an_empty_range_spans_nothing() { + let span = PageSpan::of(&(PAGE..PAGE), PAGE); + assert!(span.is_empty()); + assert_eq!(span.len(), 0); +} + +#[test] +fn a_zero_page_size_is_handled_rather_than_dividing_by_zero() { + assert!(PageSpan::of(&(0..100), 0).is_empty()); +} + +#[test] +fn containment_is_half_open() { + let span = PageSpan::of(&(0..PAGE), PAGE); + assert!(span.contains(0)); + assert!(!span.contains(1), "end is exclusive"); +} + +#[test] +fn an_offset_range_starts_at_its_own_page() { + let span = PageSpan::of(&(3 * PAGE..4 * PAGE), PAGE); + assert_eq!((span.first, span.end), (3, 4)); +} + +// ── Attribution ──────────────────────────────────────────────────────────── + +/// Router on page 0; experts on pages 1, 2 and 3. +fn layout() -> (std::ops::Range, Vec) { + ( + 0..PAGE, + vec![ + expert(0, PAGE, PAGE), + expert(1, 2 * PAGE, PAGE), + expert(2, 3 * PAGE, PAGE), + ], + ) +} + +fn account_for(resident: &[bool], selected: &[u32]) -> ResidencyAccount { + let (router, experts) = layout(); + account(resident, PAGE, &router, &experts, selected, 2 * PAGE) +} + +#[test] +fn a_sparse_read_attributes_only_the_router_and_selected_experts() { + // Router + expert 1 resident; experts 0 and 2 cold. + let a = account_for(&[true, false, true, false], &[1]); + assert_eq!(a.resident_total, 2); + assert_eq!(a.resident_router, 1); + assert_eq!(a.resident_selected, 1); + assert_eq!(a.resident_unselected, 0, "routing stayed sparse"); +} + +#[test] +fn pages_belonging_to_unselected_experts_are_counted_as_overshoot() { + // The whole file resident while only expert 1 ran — readahead defeating + // the sparsity, which is precisely what the probe exists to detect. + let a = account_for(&[true, true, true, true], &[1]); + assert_eq!(a.resident_total, 4); + assert_eq!(a.resident_unselected, 2); + assert!((a.overshoot_fraction() - 0.5).abs() < 1e-9); +} + +#[test] +fn nothing_resident_attributes_nothing() { + let a = account_for(&[false; 4], &[1]); + assert_eq!( + a, + ResidencyAccount { + predicted: 2, + ..Default::default() + } + ); + assert_eq!(a.overshoot_fraction(), 0.0); + assert_eq!(a.prediction_coverage(), 0.0); +} + +#[test] +fn every_resident_page_is_counted_exactly_once() { + let a = account_for(&[true, true, true, true], &[0, 1]); + assert_eq!( + a.resident_router + a.resident_selected + a.resident_unselected, + a.resident_total + ); +} + +#[test] +fn a_shared_boundary_page_resolves_in_favour_of_the_selected_expert() { + // Two experts sharing page 1. Attributing it to the unselected one would + // report overshoot for a page the selected expert genuinely needed. + let router = 0..PAGE; + let experts = vec![ + expert(0, PAGE, PAGE / 2), + expert(1, PAGE + PAGE / 2, PAGE / 2), + ]; + let a = account(&[false, true], PAGE, &router, &experts, &[0], PAGE); + assert_eq!(a.resident_selected, 1); + assert_eq!(a.resident_unselected, 0); +} + +#[test] +fn the_router_wins_a_page_it_shares_with_an_expert() { + let router = 0..PAGE / 2; + let experts = vec![expert(0, PAGE / 2, PAGE / 2)]; + let a = account(&[true], PAGE, &router, &experts, &[], PAGE / 2); + assert_eq!(a.resident_router, 1); + assert_eq!(a.resident_unselected, 0); +} + +#[test] +fn a_page_outside_every_region_counts_only_toward_the_total() { + // Padding or an unmapped tail. Real, resident, attributable to nothing. + let a = account_for(&[false, false, false, false, true], &[1]); + assert_eq!(a.resident_total, 1); + assert_eq!( + a.resident_router + a.resident_selected + a.resident_unselected, + 0 + ); +} + +// ── Derived ratios ───────────────────────────────────────────────────────── + +#[test] +fn resident_fraction_is_of_the_whole_mapping() { + let a = account_for(&[true, true, false, false], &[0]); + assert!((a.resident_fraction(4) - 0.5).abs() < 1e-9); +} + +#[test] +fn prediction_coverage_compares_used_pages_against_the_plan() { + // Predicted two pages (router + one expert); both became resident. + let a = account_for(&[true, false, true, false], &[1]); + assert!((a.prediction_coverage() - 1.0).abs() < 1e-9); +} + +#[test] +fn coverage_below_one_is_normal_rather_than_a_failure() { + // The plan predicts every byte of every operand; execution may finish + // without faulting a region's last page. + let (router, experts) = layout(); + let a = account( + &[true, false, false, false], + PAGE, + &router, + &experts, + &[1], + 2 * PAGE, + ); + assert!(a.prediction_coverage() < 1.0); +} + +#[test] +fn ratios_are_zero_rather_than_nan_when_there_is_no_denominator() { + let a = ResidencyAccount::default(); + assert_eq!(a.resident_fraction(0), 0.0); + assert_eq!(a.prediction_coverage(), 0.0); + assert_eq!(a.overshoot_fraction(), 0.0); +} diff --git a/crates/larql-vindex/src/runtime/tests/router.rs b/crates/larql-vindex/src/runtime/tests/router.rs new file mode 100644 index 000000000..9a50de4dc --- /dev/null +++ b/crates/larql-vindex/src/runtime/tests/router.rs @@ -0,0 +1,171 @@ +//! Colocated tests for `router`. + +use larql_compute::MoeTopKWeightPolicy; + +use super::support::{ascending, vector}; +use crate::runtime::error::ExecutionError; +use crate::runtime::router::{BoundExpertScaling, BoundRouter, RouterKernel, SelectedExpert}; + +const POPULATION: u32 = 4; +const HIDDEN: u32 = 3; +const TOP_K: usize = 2; +const ROUTER: &str = "router"; +const SCALE: &str = "router_scale"; + +fn router(per_expert_scale: Option<&[f32]>) -> BoundRouter<'static> { + BoundRouter { + weight: ascending(ROUTER, POPULATION, HIDDEN), + top_k: TOP_K, + selected_weight: MoeTopKWeightPolicy::RenormalizedSoftmax, + scaling: match per_expert_scale { + Some(v) => BoundExpertScaling::PerExpert { + scales: vector(SCALE, v), + }, + None => BoundExpertScaling::None, + }, + kernel: RouterKernel::default(), + } +} + +#[test] +fn a_router_reports_the_population_and_width_it_scores() { + let r = router(None); + assert_eq!(r.population(), POPULATION as usize); + assert_eq!(r.hidden_dim(), HIDDEN as usize); +} + +#[test] +fn a_router_validates_against_the_bank_it_routes_into() { + router(None) + .validate(POPULATION as usize, HIDDEN as usize) + .unwrap(); +} + +#[test] +fn a_router_addressing_a_different_population_is_refused() { + let err = router(None) + .validate(POPULATION as usize + 1, HIDDEN as usize) + .unwrap_err(); + assert!(matches!(err, ExecutionError::DimensionMismatch { .. })); +} + +// ── The scale is routing semantics, so a shard carries all of it ─────────── + +#[test] +fn a_scale_must_cover_the_addressable_population_not_the_resident_shard() { + // Expert 90's learned scale is part of what routing *means*, whether or + // not expert 90 lives here. Truncating the vector to the resident subset + // would make two shards of one model weight the same expert differently. + let addressable = POPULATION as usize; + let shard_size = 2; + let truncated = router(Some(&vec![1.0f32; shard_size])); + assert!( + truncated.validate(addressable, HIDDEN as usize).is_err(), + "a scale sized to the shard must be refused" + ); + let full = router(Some(&vec![1.0f32; addressable])); + full.validate(addressable, HIDDEN as usize).unwrap(); +} + +#[test] +fn a_non_finite_scale_is_refused_at_bind_time() { + // It would multiply a valid routing weight into a NaN that propagates + // through the reduction into the residual stream, surfacing as an + // inexplicable token rather than a bad operand. + let mut scales = vec![1.0f32; POPULATION as usize]; + scales[2] = f32::NAN; + let err = router(Some(&scales)) + .validate(POPULATION as usize, HIDDEN as usize) + .unwrap_err(); + let ExecutionError::NonFiniteExpertScale { expert, .. } = err else { + panic!("expected a non-finite scale refusal, got {err}"); + }; + assert_eq!(expert, 2); +} + +#[test] +fn an_infinite_scale_is_refused_too() { + let mut scales = vec![1.0f32; POPULATION as usize]; + scales[0] = f32::INFINITY; + assert!(matches!( + router(Some(&scales)) + .validate(POPULATION as usize, HIDDEN as usize) + .unwrap_err(), + ExecutionError::NonFiniteExpertScale { expert: 0, .. } + )); +} + +#[test] +fn scaling_reports_the_incumbent_policy_it_corresponds_to() { + use larql_compute::MoeExpertScalePolicy; + assert_eq!(router(None).scaling.policy(), MoeExpertScalePolicy::None); + assert_eq!( + router(Some(&vec![1.0f32; POPULATION as usize])) + .scaling + .policy(), + MoeExpertScalePolicy::PerExpert + ); +} + +#[test] +fn a_per_expert_scale_must_cover_the_whole_population() { + // A short scale vector would silently leave the tail of the population + // unscaled rather than fail. + let short = router(Some(&[1.0, 1.0])); + assert!(short + .validate(POPULATION as usize, HIDDEN as usize) + .is_err()); + + let full = router(Some(&[1.0, 1.0, 1.0, 1.0])); + full.validate(POPULATION as usize, HIDDEN as usize).unwrap(); +} + +#[test] +fn a_router_without_a_scale_validates_without_one() { + assert!(router(None).scaling.scales().is_none()); + router(None) + .validate(POPULATION as usize, HIDDEN as usize) + .unwrap(); +} + +#[test] +fn a_router_describes_its_depth_and_population() { + let text = router(None).describe(); + assert!(text.contains(&format!("top-{TOP_K}")), "{text}"); + assert!(text.contains(&POPULATION.to_string()), "{text}"); +} + +#[test] +fn a_selected_expert_keeps_both_its_weights() { + // Two selections can share a final weight after renormalisation while + // having scored differently; the raw score is where that stays visible. + let s = SelectedExpert { + expert_id: 3, + weight: 0.5, + raw_score: 0.2, + }; + assert_eq!(s.expert_id, 3); + assert_ne!(s.weight, s.raw_score); + assert_eq!(s, s); +} + +#[test] +fn each_router_kernel_has_a_distinct_name_and_the_reference_is_the_default() { + // The names reach operator-facing refusals — `KernelOperandUnsuitable` + // reports which kernel declined an operand — so two kernels sharing one + // would make a report unactionable. + assert_eq!(RouterKernel::default(), RouterKernel::Reference); + assert_ne!( + RouterKernel::Reference.name(), + RouterKernel::Incumbent.name() + ); +} + +#[test] +fn each_scaling_policy_has_a_distinct_name() { + assert_eq!(BoundExpertScaling::None.name(), "none"); + assert_ne!( + router(Some(&[1.0, 1.0, 1.0, 1.0])).scaling.name(), + BoundExpertScaling::None.name() + ); +} diff --git a/crates/larql-vindex/src/runtime/tests/support.rs b/crates/larql-vindex/src/runtime/tests/support.rs new file mode 100644 index 000000000..def8df329 --- /dev/null +++ b/crates/larql-vindex/src/runtime/tests/support.rs @@ -0,0 +1,82 @@ +//! Shared builders for runtime tests. +//! +//! `BoundTensor` borrows its bytes, so a helper that returns one has to keep +//! the buffer alive. These leak it: the buffers are a few dozen bytes, the +//! count is bounded by the test suite, and the alternative — threading an +//! owner struct through every test — obscures what each test is about. + +use crate::format::capability::binding::RepresentationIdentity; +use crate::format::capability::component::ComponentContract; +use crate::format::lyrw2::region_format::RegionFormat; + +use crate::runtime::tensor::BoundTensor; + +/// Variant these test operands claim to come from. +pub const TEST_VARIANT: &str = "test"; + +/// Q4_K super-block geometry, read from the same constants the quantiser and +/// the kernels use. Spelling `256` and `144` here would let a test keep passing +/// against a layout the runtime no longer has. +pub const Q4K_BLOCK_ELEMS: usize = larql_models::quant::ggml::Q4_K_BLOCK_ELEMS; +pub const Q4K_BLOCK_BYTES: usize = larql_models::quant::ggml::Q4_K_BLOCK_BYTES; + +/// Small, distinct, in-range values for a quantised fixture. +/// +/// Q4_K stores 4-bit values against a per-sub-block scale, so a fixture whose +/// values span many orders of magnitude round-trips badly and makes a test +/// about *binding* look like a test about precision. This ramp stays inside +/// one comfortable scale. +fn ramp(len: usize, seed: usize) -> Vec { + (0..len) + .map(|i| ((i + seed) % 19) as f32 * 0.01 - 0.09) + .collect() +} + +/// Q4_K bytes for a `rows × cols` matrix. `cols` must be a whole number of +/// super-blocks, which is what every real k-quant row is. +pub fn q4k_bytes(rows: usize, cols: usize, seed: usize) -> Vec { + assert!( + cols.is_multiple_of(Q4K_BLOCK_ELEMS), + "a Q4_K row is a whole number of {Q4K_BLOCK_ELEMS}-element blocks, not {cols}" + ); + larql_compute::cpu::ops::q4_common::quantize_q4_k(&ramp(rows * cols, seed)) +} + +fn bind(region_set: &str, values: &[f32], contract: ComponentContract) -> BoundTensor<'static> { + let bytes: Vec = values.iter().flat_map(|v| v.to_le_bytes()).collect(); + let leaked: &'static [u8] = Box::leak(bytes.into_boxed_slice()); + BoundTensor::direct( + RepresentationIdentity::new(region_set, TEST_VARIANT), + leaked, + RegionFormat::F32, + contract, + ) + .expect("test operands are well-formed") +} + +/// An f32 matrix in row-major order. +pub fn matrix(region_set: &str, rows: u32, cols: u32, values: &[f32]) -> BoundTensor<'static> { + assert_eq!( + values.len(), + (rows * cols) as usize, + "{region_set}: {rows}×{cols} needs {} values", + rows * cols + ); + bind(region_set, values, ComponentContract::matrix(rows, cols)) +} + +/// An f32 vector. +pub fn vector(region_set: &str, values: &[f32]) -> BoundTensor<'static> { + bind( + region_set, + values, + ComponentContract::vector(values.len() as u32), + ) +} + +/// A matrix whose values ascend from 1.0, for tests that only need distinct +/// numbers in a known order. +pub fn ascending(region_set: &str, rows: u32, cols: u32) -> BoundTensor<'static> { + let values: Vec = (0..rows * cols).map(|i| (i + 1) as f32).collect(); + matrix(region_set, rows, cols, &values) +} diff --git a/crates/larql-vindex/src/runtime/tests/tensor.rs b/crates/larql-vindex/src/runtime/tests/tensor.rs new file mode 100644 index 000000000..63720439f --- /dev/null +++ b/crates/larql-vindex/src/runtime/tests/tensor.rs @@ -0,0 +1,641 @@ +//! Colocated tests for `tensor` — decoding, views and refusals. +//! +//! The view cases matter most. Role-coordinate indexing is what lets one +//! executor serve tied, transposed and sliced operands, and a transpose that +//! silently reads the wrong element still produces a well-shaped result. + +use crate::format::capability::binding::{ComponentView, RepresentationIdentity}; +use crate::format::capability::component::ComponentContract; +use crate::format::lyrw2::region_format::RegionFormat; + +use crate::runtime::axis::Axis; +use crate::runtime::consts::COL_DIM; +use crate::runtime::error::{ExecutionError, OperandUnsuitability}; +use crate::runtime::tensor::BoundTensor; + +use super::support::{q4k_bytes, Q4K_BLOCK_BYTES, Q4K_BLOCK_ELEMS}; + +const REGION_SET: &str = "gate"; +const VARIANT: &str = "test"; + +fn identity() -> RepresentationIdentity { + RepresentationIdentity::new(REGION_SET, VARIANT) +} + +fn f32_bytes(values: &[f32]) -> Vec { + values.iter().flat_map(|v| v.to_le_bytes()).collect() +} + +/// A 2×3 matrix: rows [1,2,3] and [4,5,6]. +fn matrix_bytes() -> Vec { + f32_bytes(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) +} + +fn matrix(view: ComponentView) -> BoundTensor<'static> { + // Leaked so the tensor can be returned; test-only, bounded and tiny. + let bytes: &'static [u8] = Box::leak(matrix_bytes().into_boxed_slice()); + BoundTensor::new( + identity(), + bytes, + RegionFormat::F32, + ComponentContract::matrix(2, 3), + view, + ) + .unwrap() +} + +// ── Direct reads ─────────────────────────────────────────────────────────── + +#[test] +fn a_direct_matrix_reads_rows_in_storage_order() { + let t = matrix(ComponentView::Direct); + assert_eq!(t.rows(), 2); + assert_eq!(t.cols(), 3); + assert_eq!(t.row(0).unwrap(), vec![1.0, 2.0, 3.0]); + assert_eq!(t.row(1).unwrap(), vec![4.0, 5.0, 6.0]); +} + +#[test] +fn len_counts_elements_not_bytes() { + let t = matrix(ComponentView::Direct); + assert_eq!(t.len(), 6); + assert!(!t.is_empty()); +} + +#[test] +fn a_vector_reads_whole() { + let bytes = f32_bytes(&[0.5, -1.5, 2.5]); + let t = BoundTensor::direct( + identity(), + &bytes, + RegionFormat::F32, + ComponentContract::vector(3), + ) + .unwrap(); + assert_eq!(t.to_vec().unwrap(), vec![0.5, -1.5, 2.5]); + assert_eq!(t.cols(), 1, "a vector has one column per row"); +} + +// ── Views ────────────────────────────────────────────────────────────────── + +#[test] +fn a_transposed_matrix_swaps_the_role_shape_and_the_element_order() { + // The tied-LM-head case: [2,3] storage serving a [3,2] role. + let t = matrix(ComponentView::Transpose); + assert_eq!(t.rows(), 3); + assert_eq!(t.cols(), 2); + assert_eq!(t.row(0).unwrap(), vec![1.0, 4.0]); + assert_eq!(t.row(2).unwrap(), vec![3.0, 6.0]); +} + +#[test] +fn a_row_slice_offsets_into_storage() { + let t = matrix(ComponentView::Slice { + dim: 0, + start: 1, + len: 1, + }); + assert_eq!(t.rows(), 1); + assert_eq!(t.row(0).unwrap(), vec![4.0, 5.0, 6.0]); +} + +#[test] +fn a_column_slice_narrows_each_row() { + let t = matrix(ComponentView::Slice { + dim: 1, + start: 1, + len: 2, + }); + assert_eq!(t.cols(), 2); + assert_eq!(t.row(0).unwrap(), vec![2.0, 3.0]); +} + +#[test] +fn a_view_that_cannot_apply_is_refused_at_bind_time() { + let bytes = f32_bytes(&[1.0, 2.0]); + let err = BoundTensor::new( + identity(), + &bytes, + RegionFormat::F32, + ComponentContract::vector(2), + ComponentView::Transpose, + ) + .unwrap_err(); + assert!(matches!(err, ExecutionError::UnsupportedView { .. })); +} + +// ── Encodings ────────────────────────────────────────────────────────────── + +#[test] +fn bf16_decodes_through_the_shared_widening() { + // 1.0 = 0x3F80, -1.0 = 0xBF80. + let bytes = vec![0x80, 0x3F, 0x80, 0xBF]; + let t = BoundTensor::direct( + identity(), + &bytes, + RegionFormat::BF16, + ComponentContract::vector(2), + ) + .unwrap(); + assert_eq!(t.to_vec().unwrap(), vec![1.0, -1.0]); +} + +#[test] +fn f16_decodes_through_the_shared_subnormal_safe_path() { + // 1.0 = 0x3C00; the smallest positive subnormal = 0x0001. + let bytes = vec![0x00, 0x3C, 0x01, 0x00]; + let t = BoundTensor::direct( + identity(), + &bytes, + RegionFormat::F16, + ComponentContract::vector(2), + ) + .unwrap(); + let values = t.to_vec().unwrap(); + assert_eq!(values[0], 1.0); + assert!(values[1] > 0.0 && values[1] < 1e-6, "{}", values[1]); +} + +#[test] +fn a_quantised_region_binds_but_refuses_a_per_element_read() { + // The refusal moved, deliberately. A Q4_K region has known geometry, so it + // binds — the whole point of binding one is to hand its blocks to a kernel + // that reads them. What it cannot do is serve a per-element read: there is + // no per-element slot, and inventing a stride would produce a well-shaped + // tensor of noise rather than an error. + // + // Not a defect in the index — a missing reference kernel, and the message + // must say which one. + let bytes = vec![0u8; Q4K_BLOCK_BYTES]; + let bound = BoundTensor::direct( + identity(), + &bytes, + RegionFormat::Q4K, + ComponentContract::vector(2), + ) + .expect("a Q4_K region with registered geometry binds"); + let text = bound.to_vec().unwrap_err().to_string(); + assert!(text.contains("q4_k"), "{text}"); + assert!(text.contains(REGION_SET), "{text}"); +} + +#[test] +fn a_codec_with_no_registered_geometry_is_refused_at_bind() { + // The counter-case. Without a block layout there is no way to say how many + // bytes the shape needs, so the region cannot be sized, let alone read. + let bytes = vec![0u8; Q4K_BLOCK_BYTES]; + let err = BoundTensor::direct( + identity(), + &bytes, + RegionFormat::Mxfp4, + ComponentContract::vector(2), + ) + .unwrap_err(); + let text = err.to_string(); + assert!(text.contains("mxfp4"), "{text}"); + assert!(text.contains(REGION_SET), "{text}"); +} + +// ── Refusals ─────────────────────────────────────────────────────────────── + +#[test] +fn a_region_too_short_for_its_shape_is_refused_with_both_sizes() { + let bytes = f32_bytes(&[1.0, 2.0]); + let err = BoundTensor::direct( + identity(), + &bytes, + RegionFormat::F32, + ComponentContract::matrix(2, 3), + ) + .unwrap_err(); + let ExecutionError::ShortRegion { needed, found, .. } = err else { + panic!("expected a short-region refusal, got {err}"); + }; + assert_eq!((needed, found), (24, 8)); +} + +#[test] +fn a_row_past_the_end_is_refused() { + let err = matrix(ComponentView::Direct).row(2).unwrap_err(); + assert!(matches!(err, ExecutionError::RowOutOfRange { row: 2, .. })); +} + +#[test] +fn a_shape_assertion_names_the_operand_and_the_axis() { + let err = matrix(ComponentView::Direct) + .require_matrix(2, 5) + .unwrap_err(); + let ExecutionError::DimensionMismatch { axis, operand, .. } = &err else { + panic!("expected a dimension mismatch, got {err}"); + }; + assert_eq!(*axis, Axis::Columns); + assert!(operand.contains(REGION_SET), "{operand}"); +} + +#[test] +fn a_vector_asserted_as_a_matrix_is_refused_as_a_kind_error_first() { + // Kind before dimensions: asking whether a vector has the right column + // count is not a question. + let bytes = f32_bytes(&[1.0, 2.0]); + let t = BoundTensor::direct( + identity(), + &bytes, + RegionFormat::F32, + ComponentContract::vector(2), + ) + .unwrap(); + assert!(matches!( + t.require_matrix(1, 2), + Err(ExecutionError::NotAMatrix { .. }) + )); +} + +// ── Quantisation padding is not observable through the view ──────────────── +// +// Gemma's Q4_K `down` region is stored `[hidden, 768]` — the logical 704 +// rounded up to a 256-multiple — while `gate_up` is unpadded. Binding the +// stored width and slicing to the logical one is how VINDEX3 expresses +// `physical shape != semantic operand shape` without rewriting the index, +// materialising a cropped copy, or padding the activation in the generic +// runtime. +// +// These use a poisoned tail: the padding columns hold values large enough +// that any leak into a dot product would be unmistakable. + +/// Logical width, stored width, and a padding value that cannot hide. +const LOGICAL_COLS: u32 = 4; +const STORED_COLS: u32 = 6; +const POISON: f32 = 1.0e9; + +/// `[2, 6]` where columns 4 and 5 are poison. +fn poisoned_tail() -> BoundTensor<'static> { + let mut values = Vec::new(); + for row in 0..2 { + for col in 0..STORED_COLS { + values.push(if col < LOGICAL_COLS { + (row * LOGICAL_COLS + col) as f32 + 1.0 + } else { + POISON + }); + } + } + let bytes: &'static [u8] = Box::leak(f32_bytes(&values).into_boxed_slice()); + BoundTensor::new( + identity(), + bytes, + RegionFormat::F32, + ComponentContract::matrix(2, STORED_COLS), + ComponentView::Slice { + dim: 1, + start: 0, + len: LOGICAL_COLS, + }, + ) + .unwrap() +} + +#[test] +fn a_padded_tail_is_invisible_to_the_logical_view() { + let t = poisoned_tail(); + assert_eq!( + t.cols(), + LOGICAL_COLS as usize, + "role sees the logical width" + ); + for row in 0..t.rows() { + let values = t.row(row).unwrap(); + assert_eq!(values.len(), LOGICAL_COLS as usize); + assert!( + values.iter().all(|v| v.abs() < POISON), + "row {row} observed padding: {values:?}" + ); + } +} + +#[test] +fn the_logical_view_reads_the_right_values_not_merely_the_right_count() { + // A view that returned the correct *number* of elements from the wrong + // offsets would pass the test above. + let t = poisoned_tail(); + assert_eq!(t.row(0).unwrap(), vec![1.0, 2.0, 3.0, 4.0]); + assert_eq!(t.row(1).unwrap(), vec![5.0, 6.0, 7.0, 8.0]); +} + +#[test] +fn the_padding_really_is_in_the_stored_bytes() { + // Guards the guard: if the fixture never wrote poison, every assertion + // above would pass against a tensor with nothing to leak. + let t = poisoned_tail(); + let direct = BoundTensor::direct( + identity(), + t.bytes_for_test(), + RegionFormat::F32, + ComponentContract::matrix(2, STORED_COLS), + ) + .unwrap(); + assert_eq!(direct.row(0).unwrap()[LOGICAL_COLS as usize], POISON); +} + +#[test] +fn the_logical_width_is_what_a_shape_assertion_checks() { + // A kernel bound to this operand must be told the logical width, or it + // would read the padding as data. + let t = poisoned_tail(); + t.require_matrix(2, LOGICAL_COLS as usize).unwrap(); + assert!(t.require_matrix(2, STORED_COLS as usize).is_err()); +} + +#[test] +fn a_bound_tensor_reports_its_provenance() { + let t = matrix(ComponentView::Transpose); + assert_eq!(t.format(), RegionFormat::F32); + assert_eq!(*t.view(), ComponentView::Transpose); + assert_eq!(t.representation().region_set, REGION_SET); + assert!(t.describe().contains(VARIANT)); + assert_eq!(t.contract(), &ComponentContract::matrix(3, 2)); +} + +// ── Handing an f32 operand to a kernel that takes `&[f32]` ───────────────── +// +// The contract is hand-over-or-refuse. A bridge that dequantised, repacked +// into a kernel-shaped temporary and then called the kernel could reach +// numerical parity while proving nothing about the binding, so each refusal +// below names a different remedy rather than "conversion failed". + +#[test] +fn a_direct_f32_matrix_hands_over_its_own_bytes() { + let t = matrix(ComponentView::Direct); + let values = t.as_f32_slice().unwrap(); + assert_eq!(values, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); + // The region's own memory, not a copy. + assert!(std::ptr::eq( + values.as_ptr().cast::(), + t.bytes_for_test().as_ptr() + )); +} + +#[test] +fn a_non_f32_encoding_is_refused_by_format() { + // Remedy: bind another variant, or a kernel for this format. + let bytes = vec![0x80, 0x3F, 0x80, 0xBF]; + let t = BoundTensor::direct( + identity(), + &bytes, + RegionFormat::BF16, + ComponentContract::vector(2), + ) + .unwrap(); + let err = t.as_f32_slice().unwrap_err(); + assert!( + matches!(err, OperandUnsuitability::ElementFormat { .. }), + "{err}" + ); + assert!(err.to_string().contains("bf16"), "{err}"); +} + +#[test] +fn an_f32_operand_read_through_a_view_is_refused() { + // Remedy: a view-aware kernel, or a repacked variant. Unlike `as_blocks`, + // no view is honoured here: the kernel indexes the slice directly, so a + // transpose or a narrowed row would silently read the wrong elements. + for view in [ + ComponentView::Transpose, + ComponentView::Slice { + dim: COL_DIM, + start: 0, + len: 2, + }, + ] { + let err = matrix(view.clone()).as_f32_slice().unwrap_err(); + assert!( + matches!(err, OperandUnsuitability::NonDirectView { .. }), + "{view:?} should be refused, got {err}" + ); + } +} + +#[test] +fn a_misaligned_base_is_refused_as_alignment_not_as_corruption() { + // Remedy: an aligned copy. Reached by binding at a one-byte offset into an + // f32-aligned buffer, which is what an unaligned region offset produces. + let aligned: Vec = vec![1.0, 2.0, 3.0, 4.0]; + let bytes: &[u8] = + unsafe { std::slice::from_raw_parts(aligned.as_ptr().cast::(), aligned.len() * 4) }; + let t = BoundTensor::direct( + identity(), + &bytes[1..], + RegionFormat::F32, + ComponentContract::vector(2), + ) + .unwrap(); + assert!(matches!( + t.as_f32_slice().unwrap_err(), + OperandUnsuitability::MisalignedBase { .. } + )); +} + +// ── Handing blocks to a block-native kernel ──────────────────────────────── +// +// `as_blocks` is the blocked counterpart of `as_f32_slice`: hand over the +// region's own bytes or refuse. Every refusal below is a *different* remedy, +// which is the whole reason they are separate causes. + +/// Two rows of one super-block each — the smallest shape that can be wrong +/// about a row stride. +const BLOCK_ROWS: u32 = 2; +const SEED: usize = 3; + +fn q4k_rows(cols: usize) -> BoundTensor<'static> { + let bytes = q4k_bytes(BLOCK_ROWS as usize, cols, SEED); + let leaked: &'static [u8] = Box::leak(bytes.into_boxed_slice()); + BoundTensor::direct( + identity(), + leaked, + RegionFormat::Q4K, + ComponentContract::matrix(BLOCK_ROWS, cols as u32), + ) + .expect("well-formed q4_k fixture") +} + +#[test] +fn a_direct_q4k_matrix_hands_over_exactly_its_rows() { + let t = q4k_rows(Q4K_BLOCK_ELEMS); + let blocks = t.as_blocks(RegionFormat::Q4K).unwrap(); + assert_eq!(blocks.rows, BLOCK_ROWS as usize); + assert_eq!(blocks.storage_cols, Q4K_BLOCK_ELEMS); + assert_eq!(blocks.role_cols, Q4K_BLOCK_ELEMS); + assert_eq!(blocks.row_bytes, Q4K_BLOCK_BYTES); + assert_eq!(blocks.padding_cols(), 0); + assert_eq!(blocks.bytes.len(), BLOCK_ROWS as usize * Q4K_BLOCK_BYTES); +} + +#[test] +fn the_handed_over_bytes_are_the_regions_own() { + // Not a copy, not a repack: the same address. A bridge that materialised a + // kernel-shaped temporary could reach identical numbers while proving + // nothing about the binding. + let t = q4k_rows(Q4K_BLOCK_ELEMS); + let blocks = t.as_blocks(RegionFormat::Q4K).unwrap(); + assert!(std::ptr::eq( + blocks.bytes.as_ptr(), + t.bytes_for_test().as_ptr() + )); +} + +#[test] +fn a_column_prefix_slice_reports_both_extents() { + // The padded-`down` case. The kernel needs the stored width to stride + // correctly and the role width to know what the operation means, so the + // view is honoured rather than refused. + let stored = 2 * Q4K_BLOCK_ELEMS; + let role = stored - 1; + let bytes = q4k_bytes(BLOCK_ROWS as usize, stored, SEED); + let leaked: &'static [u8] = Box::leak(bytes.into_boxed_slice()); + let t = BoundTensor::new( + identity(), + leaked, + RegionFormat::Q4K, + ComponentContract::matrix(BLOCK_ROWS, stored as u32), + ComponentView::Slice { + dim: COL_DIM, + start: 0, + len: role as u32, + }, + ) + .unwrap(); + + let blocks = t.as_blocks(RegionFormat::Q4K).unwrap(); + assert_eq!(blocks.storage_cols, stored); + assert_eq!(blocks.role_cols, role); + assert_eq!(blocks.padding_cols(), 1); + assert_eq!(blocks.row_bytes, 2 * Q4K_BLOCK_BYTES); +} + +#[test] +fn a_different_codec_is_refused_by_format() { + // Remedy: bind another variant, or a kernel for this format. + let t = q4k_rows(Q4K_BLOCK_ELEMS); + assert!(matches!( + t.as_blocks(RegionFormat::Q6K).unwrap_err(), + OperandUnsuitability::ElementFormat { .. } + )); +} + +#[test] +fn a_directly_addressed_region_is_refused_by_format() { + // f32 has no blocks to hand over, and the message must not pretend the + // problem is the view or the length. + let t = matrix(ComponentView::Direct); + assert!(matches!( + t.as_blocks(RegionFormat::F32).unwrap_err(), + OperandUnsuitability::ElementFormat { .. } + )); +} + +#[test] +fn a_transpose_is_refused_because_the_bytes_are_not_the_operand() { + // Remedy: a view-aware kernel, or a repacked variant. A transposed read of + // a blocked region touches every block for one logical row. + let bytes = q4k_bytes(BLOCK_ROWS as usize, Q4K_BLOCK_ELEMS, SEED); + let t = BoundTensor::new( + identity(), + &bytes, + RegionFormat::Q4K, + ComponentContract::matrix(BLOCK_ROWS, Q4K_BLOCK_ELEMS as u32), + ComponentView::Transpose, + ) + .unwrap(); + assert!(matches!( + t.as_blocks(RegionFormat::Q4K).unwrap_err(), + OperandUnsuitability::NonDirectView { .. } + )); +} + +#[test] +fn a_row_slice_is_refused_even_though_a_column_slice_is_not() { + // The pair that makes the rule legible: a column prefix keeps whole stored + // rows and is honoured; a row slice starts partway into the region, so the + // bytes handed over would be the wrong ones. + let bytes = q4k_bytes(BLOCK_ROWS as usize, Q4K_BLOCK_ELEMS, SEED); + let t = BoundTensor::new( + identity(), + &bytes, + RegionFormat::Q4K, + ComponentContract::matrix(BLOCK_ROWS, Q4K_BLOCK_ELEMS as u32), + ComponentView::Slice { + dim: 0, + start: 0, + len: 1, + }, + ) + .unwrap(); + assert!(matches!( + t.as_blocks(RegionFormat::Q4K).unwrap_err(), + OperandUnsuitability::NonDirectView { .. } + )); +} + +#[test] +fn a_column_slice_that_does_not_start_at_zero_is_refused() { + let stored = 2 * Q4K_BLOCK_ELEMS; + let bytes = q4k_bytes(BLOCK_ROWS as usize, stored, SEED); + let t = BoundTensor::new( + identity(), + &bytes, + RegionFormat::Q4K, + ComponentContract::matrix(BLOCK_ROWS, stored as u32), + ComponentView::Slice { + dim: COL_DIM, + start: 1, + len: (stored - 1) as u32, + }, + ) + .unwrap(); + assert!(matches!( + t.as_blocks(RegionFormat::Q4K).unwrap_err(), + OperandUnsuitability::NonDirectView { .. } + )); +} + +#[test] +fn rows_that_are_not_whole_blocks_are_refused_by_alignment() { + // Remedy: repack, or pad the extent. There is no row stride to read at, so + // the kernel could only guess one. + let ragged = Q4K_BLOCK_ELEMS - 1; + let bytes = vec![0u8; BLOCK_ROWS as usize * Q4K_BLOCK_BYTES]; + let t = BoundTensor::direct( + identity(), + &bytes, + RegionFormat::Q4K, + ComponentContract::matrix(BLOCK_ROWS, ragged as u32), + ) + .unwrap(); + let err = t.as_blocks(RegionFormat::Q4K).unwrap_err(); + assert!( + matches!( + err, + OperandUnsuitability::BlockAlignment { found, block, .. } + if found == ragged && block == Q4K_BLOCK_ELEMS + ), + "{err}" + ); +} + +#[test] +fn a_region_too_short_for_its_rows_is_refused_by_length() { + // Remedy: reject the index — this one is the defect. Whole-region sizing + // rounds up once; per-row sizing rounds up per row, so a two-row region + // whose rows each need padding is longer than the flat count suggests. + let stored = Q4K_BLOCK_ELEMS; + let rows = 3usize; + let bytes = vec![0u8; (rows - 1) * Q4K_BLOCK_BYTES]; + let t = BoundTensor::direct( + identity(), + &bytes, + RegionFormat::Q4K, + ComponentContract::matrix(rows as u32, stored as u32), + ); + // Whole-region sizing already catches this one at bind time, which is the + // stronger place to catch it. + assert!(matches!(t.unwrap_err(), ExecutionError::ShortRegion { .. })); +} diff --git a/crates/larql-vindex/src/runtime/tests/tie.rs b/crates/larql-vindex/src/runtime/tests/tie.rs new file mode 100644 index 000000000..71979c4bc --- /dev/null +++ b/crates/larql-vindex/src/runtime/tests/tie.rs @@ -0,0 +1,187 @@ +//! The tie-breaking contract, pinned before Gemma parity. +//! +//! Exact ties in router scores are rare and entirely possible — quantised +//! weights and small populations both make them more likely, and a tie is +//! exactly the case where two correct-looking implementations diverge. +//! +//! The contract: +//! +//! ```text +//! score descending, then expert id ascending +//! total order over all finite scores +//! non-finite scores refused, never ordered +//! identical inputs → identical selection, byte for byte +//! ``` +//! +//! These tests exist so that a Gemma divergence at a tie is diagnosed as a +//! selection-policy difference in minutes rather than mistaken for a weight +//! decoding fault. + +use super::operation::direct; +use super::support::matrix; +use crate::runtime::error::ExecutionError; +use crate::runtime::execute::execute_traced; +use crate::runtime::inputs::MoeInputs; +use crate::runtime::kernels::{top_k, top_k_with_margin}; + +const ROUTER: &str = "router"; + +// ── The three cases named in the contract ────────────────────────────────── + +#[test] +fn a_tie_inside_the_selected_top_k_orders_by_expert_id() { + // Both experts are selected either way; the question is the order they + // appear in, which the trace and any "top expert" consumer depend on. + let picked = top_k(&[0.9, 0.5, 0.5, 0.1], 3); + assert_eq!( + picked.iter().map(|(i, _)| *i).collect::>(), + vec![0, 1, 2], + "tied experts 1 and 2 must appear in id order" + ); +} + +#[test] +fn a_tie_across_the_top_k_boundary_selects_the_lower_expert_id() { + // The consequential case: one of the tied pair runs and the other does + // not, so the policy changes the output rather than just its ordering. + let (picked, margin) = top_k_with_margin(&[0.9, 0.5, 0.5, 0.1], 2); + assert_eq!( + picked.iter().map(|(i, _)| *i).collect::>(), + vec![0, 1] + ); + assert_eq!(margin, Some(0.0), "the boundary was decided by a tie"); +} + +#[test] +fn the_same_tied_routing_repeated_gives_a_byte_identical_selection() { + let scores = [0.25f32, 0.25, 0.25, 0.25]; + let first = top_k(&scores, 2); + for _ in 0..64 { + assert_eq!(top_k(&scores, 2), first); + } + assert_eq!( + first.iter().map(|(i, _)| *i).collect::>(), + vec![0, 1] + ); +} + +// ── The margin, as parity triage ─────────────────────────────────────────── + +#[test] +fn a_clear_boundary_reports_the_gap_that_decided_it() { + let (_, margin) = top_k_with_margin(&[0.9, 0.7, 0.2, 0.1], 2); + let margin = margin.expect("a boundary exists"); + assert!((margin - 0.5).abs() < 1e-6, "0.7 - 0.2 = {margin}"); +} + +#[test] +fn there_is_no_margin_when_the_selection_covers_the_population() { + assert_eq!(top_k_with_margin(&[0.6, 0.4], 2).1, None); + assert_eq!(top_k_with_margin(&[0.6, 0.4], 9).1, None); +} + +#[test] +fn there_is_no_margin_when_nothing_is_selected() { + assert_eq!(top_k_with_margin(&[0.6, 0.4], 0).1, None); +} + +#[test] +fn a_tie_decided_selection_is_flagged_in_the_trace() { + // A router whose rows are all identical scores every expert equally, so + // the boundary is a pure tie and top-k is decided entirely by policy. + let mut op = direct(); + let population = op.banks[0].experts.len(); + let hidden = op.residual_dim; + let uniform = vec![0.25f32; population * hidden]; + op.router.weight = matrix(ROUTER, population as u32, hidden as u32, &uniform); + + let (_, trace) = execute_traced(&op, MoeInputs::shared(&vec![0.1f32; hidden])).unwrap(); + assert!( + trace.selection_was_decided_by_a_tie(), + "margin was {:?}", + trace.selection_margin + ); + assert_eq!(trace.selected_ids(), vec![0, 1], "lowest ids win the tie"); +} + +#[test] +fn an_ordinary_selection_is_not_flagged_as_a_tie() { + let (_, trace) = execute_traced( + &direct(), + MoeInputs::shared(&[0.4, -0.3, 0.9, 0.1, -0.7, 0.2]), + ) + .unwrap(); + assert!(!trace.selection_was_decided_by_a_tie()); + assert!(trace.selection_margin.unwrap() > 0.0); +} + +// ── A total order, not a fallback ────────────────────────────────────────── + +#[test] +fn ordering_is_total_over_signed_zeroes() { + // `partial_cmp` calls +0.0 and -0.0 equal, so the id tie-break decides; + // `total_cmp` orders them. Either way the result must be deterministic, + // which a fallback-to-Equal sort does not guarantee. + let first = top_k(&[0.0, -0.0, 0.0], 2); + for _ in 0..32 { + assert_eq!(top_k(&[0.0, -0.0, 0.0], 2), first); + } +} + +#[test] +fn every_permutation_of_a_fully_tied_population_selects_the_same_ids() { + // Position in the input must not influence selection when scores are equal. + for population in 2..8usize { + let scores = vec![0.5f32; population]; + let picked: Vec = top_k(&scores, 2).iter().map(|(i, _)| *i).collect(); + assert_eq!(picked, vec![0, 1], "population {population}"); + } +} + +// ── Non-finite scores are refused, never ordered ─────────────────────────── + +#[test] +fn a_nan_router_score_is_refused_naming_the_expert() { + // A NaN would otherwise take whatever position the total order gives it + // and be selected or dropped by sort mechanics. + let mut op = direct(); + let population = op.banks[0].experts.len(); + let hidden = op.residual_dim; + let mut values = vec![0.25f32; population * hidden]; + values[hidden] = f32::NAN; // expert 1's first weight + op.router.weight = matrix(ROUTER, population as u32, hidden as u32, &values); + + let err = execute_traced(&op, MoeInputs::shared(&vec![0.5f32; hidden])).unwrap_err(); + let ExecutionError::NonFiniteRouterScore { expert, value } = err else { + panic!("expected a non-finite refusal, got {err}"); + }; + assert_eq!(expert, 1); + assert!(value.is_nan()); +} + +#[test] +fn an_infinite_router_weight_is_refused_rather_than_dominating_the_softmax() { + let mut op = direct(); + let population = op.banks[0].experts.len(); + let hidden = op.residual_dim; + let mut values = vec![0.25f32; population * hidden]; + values[0] = f32::INFINITY; + op.router.weight = matrix(ROUTER, population as u32, hidden as u32, &values); + + // softmax(inf) yields NaN once the max-shift subtracts inf from inf. + let err = execute_traced(&op, MoeInputs::shared(&vec![0.5f32; hidden])).unwrap_err(); + assert!( + matches!(err, ExecutionError::NonFiniteRouterScore { .. }), + "{err}" + ); +} + +#[test] +fn finite_scores_are_never_refused() { + let (_, trace) = execute_traced( + &direct(), + MoeInputs::shared(&[0.4, -0.3, 0.9, 0.1, -0.7, 0.2]), + ) + .unwrap(); + assert!(trace.router_scores.iter().all(|s| s.is_finite())); +} diff --git a/crates/larql-vindex/src/runtime/tests/transform.rs b/crates/larql-vindex/src/runtime/tests/transform.rs new file mode 100644 index 000000000..a76c8a87b --- /dev/null +++ b/crates/larql-vindex/src/runtime/tests/transform.rs @@ -0,0 +1,131 @@ +//! Colocated tests for `transform` — the latent stages. +//! +//! A direct MoE binds no transforms, so nothing in fixture A exercises this +//! file. It is the machinery Mini-K3 needs (`residual → routed_input → bank → +//! routed_output → residual`), and leaving it untested until then would mean +//! discovering its faults mixed in with K3 router semantics. + +use super::support::{ascending, matrix}; +use crate::runtime::error::ExecutionError; +use crate::runtime::transform::{apply_stage, BoundTransform, TransformStage}; + +const ROUTED_IN: &str = "latent_in"; +const ROUTED_OUT: &str = "latent_out"; + +/// `[out, in]` = [2, 3]: rows [1,2,3] and [4,5,6]. +fn down_project() -> BoundTransform<'static> { + BoundTransform { + stage: TransformStage::RoutedInput, + weight: ascending(ROUTED_IN, 2, 3), + } +} + +/// `[out, in]` = [3, 2], projecting back to the residual width. +fn up_project() -> BoundTransform<'static> { + BoundTransform { + stage: TransformStage::RoutedOutput, + weight: ascending(ROUTED_OUT, 3, 2), + } +} + +// ── One transform ────────────────────────────────────────────────────────── + +#[test] +fn a_transform_contracts_the_input_axis() { + // [1,2,3]·[1,1,1] = 6; [4,5,6]·[1,1,1] = 15. + let out = down_project().apply(&[1.0, 1.0, 1.0]).unwrap(); + assert_eq!(out, vec![6.0, 15.0]); +} + +#[test] +fn a_transform_reports_the_widths_it_maps_between() { + let t = down_project(); + assert_eq!(t.in_dim(), 3); + assert_eq!(t.out_dim(), 2); +} + +#[test] +fn an_input_of_the_wrong_width_is_refused_naming_both() { + let err = down_project().apply(&[1.0, 1.0]).unwrap_err(); + let ExecutionError::DimensionMismatch { + expected, found, .. + } = err + else { + panic!("expected a dimension mismatch"); + }; + assert_eq!((expected, found), (3, 2)); +} + +#[test] +fn a_transform_describes_its_stage_and_widths() { + let text = down_project().describe(); + assert!(text.contains(TransformStage::RoutedInput.name()), "{text}"); + assert!(text.contains("3→2"), "{text}"); +} + +#[test] +fn stages_name_the_two_sides_of_the_bank() { + assert_eq!(TransformStage::RoutedInput.name(), "routed_input"); + assert_eq!(TransformStage::RoutedOutput.name(), "routed_output"); + assert_ne!(TransformStage::RoutedInput, TransformStage::RoutedOutput); +} + +// ── Stage application ────────────────────────────────────────────────────── + +#[test] +fn an_empty_transform_list_is_the_identity() { + // The direct-MoE case: no stage bound means the residual passes through + // unchanged, with no branch anywhere that could skip a bound projection. + let input = vec![1.0, -2.0, 3.0]; + let out = apply_stage(&[], TransformStage::RoutedInput, &input).unwrap(); + assert_eq!(out, input); +} + +#[test] +fn only_transforms_bound_to_the_requested_stage_are_applied() { + let transforms = [down_project(), up_project()]; + let out = apply_stage(&transforms, TransformStage::RoutedInput, &[1.0, 1.0, 1.0]).unwrap(); + assert_eq!(out, vec![6.0, 15.0], "the output stage must not run here"); +} + +#[test] +fn the_output_stage_projects_back_to_the_residual_width() { + let transforms = [down_project(), up_project()]; + let out = apply_stage(&transforms, TransformStage::RoutedOutput, &[1.0, 1.0]).unwrap(); + assert_eq!(out, vec![3.0, 7.0, 11.0]); +} + +#[test] +fn chained_transforms_apply_in_declaration_order() { + // [2,3] then [3,2]: order matters, and the reversed pair would not even + // typecheck dimensionally — which is the point. + let first = BoundTransform { + stage: TransformStage::RoutedInput, + weight: matrix(ROUTED_IN, 2, 3, &[1.0, 0.0, 0.0, 0.0, 1.0, 0.0]), + }; + let second = BoundTransform { + stage: TransformStage::RoutedInput, + weight: matrix(ROUTED_IN, 1, 2, &[1.0, 1.0]), + }; + let out = apply_stage( + &[first, second], + TransformStage::RoutedInput, + &[5.0, 7.0, 9.0], + ) + .unwrap(); + assert_eq!(out, vec![12.0], "5 + 7"); +} + +#[test] +fn a_stage_mismatch_in_a_chain_is_refused() { + let mismatched = BoundTransform { + stage: TransformStage::RoutedInput, + weight: ascending(ROUTED_IN, 2, 5), + }; + assert!(apply_stage( + &[down_project(), mismatched], + TransformStage::RoutedInput, + &[1.0, 1.0, 1.0] + ) + .is_err()); +} diff --git a/crates/larql-vindex/src/runtime/tests/verdict.rs b/crates/larql-vindex/src/runtime/tests/verdict.rs new file mode 100644 index 000000000..25aab39fe --- /dev/null +++ b/crates/larql-vindex/src/runtime/tests/verdict.rs @@ -0,0 +1,326 @@ +//! Tests for `verdict` — the classification a sweep reports per layer. +//! +//! Two properties carry the weight, and both are about what a *headline number* +//! ends up meaning: +//! +//! - a residency refusal must never read as a parity failure, or the sweep +//! measures slice coverage instead of execution; +//! - an observed agreement must never read as a contract, or a +//! schedule-dependent result becomes a promise nobody can keep. + +use crate::runtime::axis::Axis; +use crate::runtime::error::{ExecutionError, OperandUnsuitability}; +use crate::runtime::verdict::{RefusalKind, Verdict}; + +const OPERAND: &str = "gate_up_fused::test"; +const KERNEL: &str = "incumbent_q4k_q8k"; + +/// One error of every variant, so the classification sweep below cannot narrow +/// silently when a variant is added. +fn every_error() -> Vec { + vec![ + ExecutionError::UnsupportedFormat { + format: "q4_k".into(), + operand: OPERAND.into(), + }, + ExecutionError::UnsupportedView { + view: "transpose".into(), + operand: OPERAND.into(), + }, + ExecutionError::DimensionMismatch { + operand: OPERAND.into(), + axis: Axis::Rows, + expected: 2, + found: 3, + }, + ExecutionError::ShortRegion { + operand: OPERAND.into(), + needed: 8, + found: 4, + }, + ExecutionError::RowOutOfRange { + operand: OPERAND.into(), + row: 9, + rows: 2, + }, + ExecutionError::NotAMatrix { + operand: OPERAND.into(), + found: "vector".into(), + }, + ExecutionError::NonFiniteRouterScore { + expert: 3, + value: f32::NAN, + }, + ExecutionError::NonFiniteExpertScale { + expert: 3, + value: f32::INFINITY, + operand: OPERAND.into(), + }, + ExecutionError::KernelOperandUnsuitable { + kernel: KERNEL, + operand: OPERAND.into(), + reason: OperandUnsuitability::NonDirectView { + view: "transpose".into(), + }, + }, + ExecutionError::KernelActivationUnsupported { + kernel: KERNEL, + activation: "ReLU".into(), + }, + ExecutionError::ExpertOutOfRange { + expert: 900, + population: 128, + }, + ExecutionError::SelectedExpertNotResident { + expert: 90, + bank: "layer 5 bank 0".into(), + resident: 8, + population: 128, + }, + ] +} + +// ── Residency is not a failure ───────────────────────────────────────────── + +#[test] +fn a_non_resident_expert_is_a_residency_refusal_and_nothing_else() { + // The whole reason this taxonomy exists. The route was correct, the expert + // exists in the model, and a shard that does not hold it is behaving as + // designed — so this must not indict the execution. + let err = ExecutionError::SelectedExpertNotResident { + expert: 90, + bank: "layer 5 bank 0".into(), + resident: 8, + population: 128, + }; + assert_eq!(err.refusal(), RefusalKind::Residency); + + let verdict = Verdict::from(&err); + assert_eq!(verdict, Verdict::Refused(RefusalKind::Residency)); + assert!(!verdict.indicts_execution()); + assert!(!verdict.agreed()); + assert!(!verdict.is_contracted()); +} + +#[test] +fn an_out_of_range_expert_is_a_binding_defect_not_a_residency_miss() { + // The pair that makes the distinction legible: one means fetch the + // operand, the other means the index is wrong. Collapsing them would leave + // an operator unable to tell which. + let err = ExecutionError::ExpertOutOfRange { + expert: 900, + population: 128, + }; + assert_eq!(err.refusal(), RefusalKind::BindingDefect); + assert!(Verdict::from(&err).indicts_execution()); +} + +#[test] +fn a_missing_kernel_is_unsupported_and_does_not_indict_the_execution() { + // A valid representation with a valid semantic view can have no admissible + // kernel. That is a gap, not a defect — nobody should reject the index. + for err in [ + ExecutionError::KernelActivationUnsupported { + kernel: KERNEL, + activation: "ReLU".into(), + }, + ExecutionError::KernelOperandUnsuitable { + kernel: KERNEL, + operand: OPERAND.into(), + reason: OperandUnsuitability::Arrangement { + found: "gate+up".into(), + wanted: "one fused gate+up region", + }, + }, + ] { + assert_eq!(err.refusal(), RefusalKind::Unsupported, "{err}"); + assert!(!Verdict::from(&err).indicts_execution(), "{err}"); + } +} + +#[test] +fn every_error_variant_is_classified() { + // `refusal` is an exhaustive match, so this cannot fail to compile — but it + // can fail to be *exercised*, which is what leaves a new variant silently + // inheriting whatever arm it was pattern-matched into. + let errors = every_error(); + assert_eq!(errors.len(), 12, "a variant joined ExecutionError"); + for err in &errors { + assert!(RefusalKind::ALL.contains(&err.refusal()), "{err}"); + } +} + +#[test] +fn nothing_but_a_non_resident_expert_classifies_as_residency() { + // The dangerous default. `Residency` is the one kind that reads as + // "nothing is wrong", so an unclassified failure inheriting it would make a + // sweep report a defect as normal shard behaviour. + let residency: Vec = every_error() + .iter() + .filter(|e| e.refusal() == RefusalKind::Residency) + .map(ToString::to_string) + .collect(); + assert_eq!(residency.len(), 1, "{residency:?}"); + assert!(residency[0].contains("not resident"), "{residency:?}"); +} + +// ── Asserted versus observed ─────────────────────────────────────────────── + +#[test] +fn an_identical_result_is_exact_only_when_it_was_asserted() { + assert_eq!( + Verdict::from_comparison(true, true, true), + Verdict::Exact, + "an asserted checkpoint that matched is a contract" + ); + assert_eq!( + Verdict::from_comparison(true, true, false), + Verdict::ObservedExact, + "the same bytes, watched rather than contracted, are not" + ); +} + +#[test] +fn only_exact_is_a_contract() { + assert!(Verdict::Exact.is_contracted()); + for verdict in [ + Verdict::ObservedExact, + Verdict::Equivalent, + Verdict::NumericMismatch, + Verdict::Refused(RefusalKind::Residency), + ] { + assert!( + !verdict.is_contracted(), + "{verdict} must not read as a guarantee" + ); + } +} + +#[test] +fn observed_exact_still_counts_as_agreement() { + // It is weaker than `Exact`, not a failure. A sweep that treated it as one + // would report a bit-identical block as a mismatch. + assert!(Verdict::ObservedExact.agreed()); + assert!(!Verdict::ObservedExact.indicts_execution()); +} + +#[test] +fn a_difference_is_equivalent_or_a_mismatch_by_the_tolerance_alone() { + // Whether it was asserted is irrelevant once the values differ: an + // asserted checkpoint that missed is not "more equivalent" than a watched + // one. + assert_eq!( + Verdict::from_comparison(false, true, true), + Verdict::Equivalent + ); + assert_eq!( + Verdict::from_comparison(false, true, false), + Verdict::Equivalent + ); + assert_eq!( + Verdict::from_comparison(false, false, true), + Verdict::NumericMismatch + ); + assert_eq!( + Verdict::from_comparison(false, false, false), + Verdict::NumericMismatch + ); +} + +#[test] +fn only_a_numeric_mismatch_and_a_binding_defect_indict_the_execution() { + assert!(Verdict::NumericMismatch.indicts_execution()); + assert!(Verdict::Refused(RefusalKind::BindingDefect).indicts_execution()); + for verdict in [ + Verdict::Exact, + Verdict::ObservedExact, + Verdict::Equivalent, + Verdict::Refused(RefusalKind::Residency), + Verdict::Refused(RefusalKind::Unsupported), + ] { + assert!(!verdict.indicts_execution(), "{verdict}"); + } +} + +// ── Reporting ────────────────────────────────────────────────────────────── + +#[test] +fn a_sweep_can_take_its_worst_outcome_by_maximum() { + // The ordering exists so a summary line needs no bespoke fold. Weakest + // claims sort last. + let mut outcomes = [ + Verdict::Equivalent, + Verdict::Exact, + Verdict::Refused(RefusalKind::Residency), + Verdict::ObservedExact, + ]; + outcomes.sort_unstable(); + assert_eq!(outcomes.first(), Some(&Verdict::Exact)); + assert_eq!( + outcomes.iter().max(), + Some(&Verdict::Refused(RefusalKind::Residency)) + ); +} + +#[test] +fn every_verdict_has_a_distinct_name() { + // Names go into a per-layer table that a human scans, so two outcomes + // sharing one would make the table unreadable exactly where it matters. + let mut names: Vec<&str> = vec![ + Verdict::Exact.name(), + Verdict::ObservedExact.name(), + Verdict::Equivalent.name(), + Verdict::NumericMismatch.name(), + ]; + names.extend(RefusalKind::ALL.iter().map(|k| Verdict::Refused(*k).name())); + let count = names.len(); + names.sort_unstable(); + names.dedup(); + assert_eq!(names.len(), count, "two verdicts share a name"); +} + +#[test] +fn a_verdict_displays_as_its_name() { + assert_eq!(Verdict::Exact.to_string(), "exact"); + assert_eq!( + Verdict::Refused(RefusalKind::Residency).to_string(), + RefusalKind::Residency.to_string() + ); + assert_eq!(RefusalKind::BindingDefect.to_string(), "binding_defect"); +} + +// ── The classification crossing a crate boundary ─────────────────────────── + +#[test] +fn the_boundary_trait_reports_the_same_kind_as_the_inherent_method() { + // `ExecutionRefusal::kind` delegates to `refusal()` precisely so the two + // cannot disagree — but a delegation nothing calls is a delegation nobody + // has checked. This walks every variant through both. + use larql_execution::ExecutionRefusal; + for err in every_error() { + assert_eq!( + ExecutionRefusal::kind(&err), + err.refusal(), + "the boundary trait and the inherent method disagree for {err}" + ); + } +} + +#[test] +fn a_boxed_execution_error_keeps_its_classification_and_its_detail() { + // The reason the trait exists: `larql-compute` cannot name `ExecutionError` + // but can name `dyn ExecutionRefusal`, and both levels must survive the + // crossing — the response category *and* the concrete diagnosis. + let boxed: larql_execution::BoxRefusal = Box::new(ExecutionError::SelectedExpertNotResident { + expert: 90, + bank: "layer 5 bank 0".into(), + resident: 8, + population: 128, + }); + assert_eq!(boxed.kind(), RefusalKind::Residency); + assert!(!boxed.kind().indicts_the_artifact()); + // The detail is still there, not flattened to the category. + let text = boxed.to_string(); + assert!(text.contains("expert 90"), "{text}"); + assert!(text.contains("8 of 128"), "{text}"); +} diff --git a/crates/larql-vindex/src/runtime/trace.rs b/crates/larql-vindex/src/runtime/trace.rs new file mode 100644 index 000000000..285639ae9 --- /dev/null +++ b/crates/larql-vindex/src/runtime/trace.rs @@ -0,0 +1,130 @@ +//! Instrumented boundaries inside one MoE execution. +//! +//! A final-vector equality assertion tells you an execution is wrong and +//! nothing else. These checkpoints localise it: a mismatch at `router_scores` +//! is scoring or the router operand, at `selection` it is top-k or a weight +//! policy, at `expert_output` it is region interpretation or activation, at +//! `reduced` it is the combine. +//! +//! # Why a sink rather than a returned struct +//! +//! The traced and untraced paths must be the *same* arithmetic — a fixture +//! that exercises a separate instrumented implementation proves that +//! implementation and not the one Gemma will run. So there is one generic +//! execution, and the sink is a type parameter: [`NoTrace`] has empty methods +//! and compiles away, leaving the token path with no allocation and no +//! branch, while [`CollectedTrace`] records. + +use super::router::SelectedExpert; + +/// Receives internal checkpoints during execution. +/// +/// Every method defaults to doing nothing, so the production sink is a single +/// empty impl and adding a checkpoint never breaks an existing one. +pub trait TraceSink { + /// Bank input, after any routed-input transform. Equals the residual for a + /// direct MoE. + fn routed_input(&mut self, _values: &[f32]) {} + /// Router probabilities over the whole population, before selection. + fn router_scores(&mut self, _values: &[f32]) {} + /// How much room the selection had: `score[k-1] - score[k]`, or `None` + /// when there is no boundary. Zero means an exact tie decided it. + fn selection_margin(&mut self, _margin: Option) {} + /// Selected experts in selection order, with final and raw weights. + fn selection(&mut self, _selected: &[SelectedExpert]) {} + /// One expert's output, before its routing weight is applied. + fn expert_output(&mut self, _expert_id: u32, _values: &[f32]) {} + /// The bank's combined output, before any routed-output transform. + fn reduced(&mut self, _values: &[f32]) {} + /// What the operation contributes to the residual stream. + fn residual_delta(&mut self, _values: &[f32]) {} +} + +/// The production sink: records nothing, costs nothing. +#[derive(Debug, Clone, Copy, Default)] +pub struct NoTrace; + +impl TraceSink for NoTrace {} + +/// One expert's unweighted output, kept in selection order. +#[derive(Debug, Clone, PartialEq)] +pub struct ExpertOutput { + pub expert_id: u32, + pub values: Vec, +} + +/// Records every checkpoint. For fixtures, differential tests and diagnosis. +#[derive(Debug, Clone, Default, PartialEq)] +pub struct CollectedTrace { + pub routed_input: Vec, + pub router_scores: Vec, + /// `score[k-1] - score[k]`. `Some(0.0)` means an exact tie sat on the + /// top-k boundary and the ordering policy, not the scores, chose. + pub selection_margin: Option, + pub selection: Vec, + pub expert_outputs: Vec, + pub reduced: Vec, + pub residual_delta: Vec, +} + +impl CollectedTrace { + /// Selected expert ids, in selection order. + pub fn selected_ids(&self) -> Vec { + self.selection.iter().map(|s| s.expert_id).collect() + } + + /// Final routing weights, in selection order. + pub fn gate_weights(&self) -> Vec { + self.selection.iter().map(|s| s.weight).collect() + } + + /// Whether an exact tie sat on the top-k boundary. + /// + /// The triage question when two implementations disagree about which + /// experts ran: true means they met a tie and resolved it differently, + /// false means they disagreed about the scores themselves. + pub fn selection_was_decided_by_a_tie(&self) -> bool { + self.selection_margin == Some(0.0) + } + + /// One expert's recorded output, by id. + pub fn expert_output(&self, expert_id: u32) -> Option<&[f32]> { + self.expert_outputs + .iter() + .find(|o| o.expert_id == expert_id) + .map(|o| o.values.as_slice()) + } +} + +impl TraceSink for CollectedTrace { + fn routed_input(&mut self, values: &[f32]) { + self.routed_input = values.to_vec(); + } + + fn router_scores(&mut self, values: &[f32]) { + self.router_scores = values.to_vec(); + } + + fn selection_margin(&mut self, margin: Option) { + self.selection_margin = margin; + } + + fn selection(&mut self, selected: &[SelectedExpert]) { + self.selection = selected.to_vec(); + } + + fn expert_output(&mut self, expert_id: u32, values: &[f32]) { + self.expert_outputs.push(ExpertOutput { + expert_id, + values: values.to_vec(), + }); + } + + fn reduced(&mut self, values: &[f32]) { + self.reduced = values.to_vec(); + } + + fn residual_delta(&mut self, values: &[f32]) { + self.residual_delta = values.to_vec(); + } +} diff --git a/crates/larql-vindex/src/runtime/transform.rs b/crates/larql-vindex/src/runtime/transform.rs new file mode 100644 index 000000000..06a8c83e2 --- /dev/null +++ b/crates/larql-vindex/src/runtime/transform.rs @@ -0,0 +1,97 @@ +//! Latent transforms around a bank (spec §5, `latent-moe-v1`). +//! +//! A direct routed MoE reads the residual and writes the residual. A latent +//! one does not: it projects into a narrower routed space, runs the bank +//! there, and projects back. +//! +//! ```text +//! direct residual → bank → residual +//! latent residual → routed_input → bank → routed_output → residual +//! ``` +//! +//! Modelling that as an ordered stage list rather than two optional tensors +//! keeps the direct case as the empty list — no `Option` to forget, no branch +//! that silently skips a projection when one was bound. + +use super::axis::Axis; +use super::error::ExecutionError; +use super::kernels::dot; +use super::tensor::BoundTensor; + +/// Where a transform sits relative to the bank. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub enum TransformStage { + /// Applied to the residual before routing and expert execution. + RoutedInput, + /// Applied to the bank's reduced output before it rejoins the residual. + RoutedOutput, +} + +impl TransformStage { + pub const fn name(self) -> &'static str { + match self { + Self::RoutedInput => "routed_input", + Self::RoutedOutput => "routed_output", + } + } +} + +/// A bound projection applied at one stage. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct BoundTransform<'a> { + pub stage: TransformStage, + /// `[out_dim, in_dim]`. + pub weight: BoundTensor<'a>, +} + +impl BoundTransform<'_> { + pub fn in_dim(&self) -> usize { + self.weight.cols() + } + + pub fn out_dim(&self) -> usize { + self.weight.rows() + } + + /// Apply the projection to `input`. + pub fn apply(&self, input: &[f32]) -> Result, ExecutionError> { + if input.len() != self.in_dim() { + return Err(ExecutionError::DimensionMismatch { + operand: self.weight.describe(), + axis: Axis::InputWidth, + expected: self.in_dim(), + found: input.len(), + }); + } + let mut out = vec![0.0f32; self.out_dim()]; + let mut row = vec![0.0f32; self.in_dim()]; + for (r, slot) in out.iter_mut().enumerate() { + self.weight.row_into(r, &mut row)?; + *slot = dot(&row, input); + } + Ok(out) + } + + pub fn describe(&self) -> String { + format!( + "{} {} [{}→{}]", + self.stage.name(), + self.weight.describe(), + self.in_dim(), + self.out_dim() + ) + } +} + +/// Apply every transform bound to a stage, in order. +pub fn apply_stage( + transforms: &[BoundTransform<'_>], + stage: TransformStage, + input: &[f32], +) -> Result, ExecutionError> { + let mut current = input.to_vec(); + for transform in transforms.iter().filter(|t| t.stage == stage) { + current = transform.apply(¤t)?; + } + Ok(current) +} diff --git a/crates/larql-vindex/src/runtime/verdict.rs b/crates/larql-vindex/src/runtime/verdict.rs new file mode 100644 index 000000000..576fd6c6d --- /dev/null +++ b/crates/larql-vindex/src/runtime/verdict.rs @@ -0,0 +1,158 @@ +//! What happened when one bound operation was compared against the incumbent. +//! +//! A sweep over thirty layers reports one line each, and the temptation is a +//! pass/fail column. That column is a lie in both directions: it calls a +//! non-resident expert a failure when the route was correct and the operand +//! simply lives elsewhere, and it calls a within-tolerance agreement a pass +//! when nothing was asserted. Both readings send the next person to the wrong +//! place. +//! +//! So the vocabulary lives here, in the library, rather than in whichever +//! harness happens to need it — the locked-input sweep, the free-propagation +//! residual comparison and the decode-parity run must classify the same +//! situation the same way, and three private enums would drift within a week. +//! +//! # The distinction that matters most +//! +//! ```text +//! Exact every asserted checkpoint was bit-identical +//! ObservedExact the whole comparison came out bit-identical, uncontracted +//! ``` +//! +//! The second is a *measurement*, not a promise. A reduction whose order +//! depends on the parallel schedule can come out bit-identical on one +//! configuration and not on another, and recording that as `Exact` would +//! manufacture a contract nobody can keep. Anything downstream that needs a +//! guarantee must read [`Verdict::is_contracted`], not "it agreed last time". + +use super::error::ExecutionError; + +/// The response a refusal requires — re-exported from `larql-execution`. +/// +/// It moved there when `FfnBackend`'s error channel needed to name it: +/// `FfnBackend` lives in `larql-compute`, which sits *below* this crate, so the +/// vocabulary could not stay here and still cross that boundary. Re-exported so +/// existing `runtime::RefusalKind` paths keep working. +pub use larql_execution::RefusalKind; + +impl ExecutionError { + /// Which kind of refusal this is, for a sweep that must not report a + /// missing operand as a wrong answer. + /// + /// Exhaustive on purpose: a new variant will not compile until someone has + /// decided which of the three it is, which is the point. The default for an + /// unclassified failure must never be `Residency` — that is the one kind + /// that reads as "nothing is wrong". + pub const fn refusal(&self) -> RefusalKind { + match self { + // The route was right; the operand lives elsewhere. + Self::SelectedExpertNotResident { .. } => RefusalKind::Residency, + + // Well-formed, but nothing bound serves it. + Self::UnsupportedFormat { .. } + | Self::UnsupportedView { .. } + | Self::KernelOperandUnsuitable { .. } + | Self::KernelActivationUnsupported { .. } => RefusalKind::Unsupported, + + // The binding is wrong. `NonFinite*` land here rather than in a + // category of their own: a NaN reaching a routing decision is a + // statement about the bytes that were bound, not about arithmetic + // this runtime performed. + Self::DimensionMismatch { .. } + | Self::ShortRegion { .. } + | Self::RowOutOfRange { .. } + | Self::NotAMatrix { .. } + | Self::NonFiniteRouterScore { .. } + | Self::NonFiniteExpertScale { .. } + | Self::ExpertOutOfRange { .. } => RefusalKind::BindingDefect, + } + } +} + +/// The outcome of comparing one bound execution against the incumbent. +/// +/// Ordered by how much it claims, weakest last, so a sweep can report its +/// worst outcome by taking the maximum. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub enum Verdict { + /// Every asserted checkpoint was bit-identical. The strongest claim, and + /// the only one that is a contract. + Exact, + /// The comparison came out bit-identical without being asserted — a + /// reduction whose order the parallel schedule chose, say. True of this + /// run; not promised of the next. + ObservedExact, + /// Within a named, justified tolerance. The tolerance has to be stated + /// wherever this is produced; a bare "equivalent" is the vague bucket this + /// enum exists to prevent. + Equivalent, + /// Execution completed and the answers disagree. The genuine failure. + NumericMismatch, + /// Execution did not complete. See [`RefusalKind`] for who acts on it. + Refused(RefusalKind), +} + +impl Verdict { + pub const fn name(self) -> &'static str { + match self { + Self::Exact => "exact", + Self::ObservedExact => "observed_exact", + Self::Equivalent => "equivalent", + Self::NumericMismatch => "numeric_mismatch", + Self::Refused(kind) => kind.name(), + } + } + + /// Whether this outcome is a guarantee rather than a measurement. + /// + /// Only [`Self::Exact`] is. `ObservedExact` deliberately is not: promoting + /// an observation to a contract is how a scheduling-dependent result + /// becomes a promise nobody can keep. + pub const fn is_contracted(self) -> bool { + matches!(self, Self::Exact) + } + + /// Whether the two paths agreed, to whatever strength. + pub const fn agreed(self) -> bool { + matches!(self, Self::Exact | Self::ObservedExact | Self::Equivalent) + } + + /// Whether this outcome indicts the execution. + /// + /// A residency refusal does not: the route was correct and the operand was + /// absent. Neither does an unsupported one — it says a kernel is missing, + /// which is a gap rather than a defect. + pub const fn indicts_execution(self) -> bool { + matches!( + self, + Self::NumericMismatch | Self::Refused(RefusalKind::BindingDefect) + ) + } + + /// Classify a completed comparison. + /// + /// `asserted` says whether the checkpoints that came out identical were + /// ones the harness contracts, or ones it merely watched. Passing `true` + /// for a schedule-dependent comparison is the mistake this parameter + /// exists to make visible at the call site. + pub const fn from_comparison(identical: bool, within_tolerance: bool, asserted: bool) -> Self { + match (identical, asserted) { + (true, true) => Self::Exact, + (true, false) => Self::ObservedExact, + (false, _) if within_tolerance => Self::Equivalent, + (false, _) => Self::NumericMismatch, + } + } +} + +impl From<&ExecutionError> for Verdict { + fn from(error: &ExecutionError) -> Self { + Self::Refused(error.refusal()) + } +} + +impl std::fmt::Display for Verdict { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.name()) + } +} diff --git a/crates/larql-vindex/tests/e0_generation_boundary.rs b/crates/larql-vindex/tests/e0_generation_boundary.rs new file mode 100644 index 000000000..f4adf63f7 --- /dev/null +++ b/crates/larql-vindex/tests/e0_generation_boundary.rs @@ -0,0 +1,220 @@ +//! E0 — the generation-boundary rows of the VINDEX2 preservation matrix. +//! +//! # Scope, stated honestly +//! +//! The full E0 matrix (`docs/vindex3-experiments.md`) covers loading, +//! verification, WALK, slicing, full decode, layer and expert sharding, and +//! publish/pull round-trips against real indexes. Those rows need multi-GB +//! checkpoints and cannot run in CI; they run locally against the committed +//! goldens in `tests/goldens/e0/` via `scripts/e0-capture-goldens.sh`. +//! +//! What runs *here* is the subset that needs no weights at all: **generation +//! detection and the loader boundary**. That is deliberately the subset most +//! at risk from VINDEX3 work, because every commit touching `format/load.rs`, +//! `format/generation.rs` or the LYRW parsers can break it, and nothing else +//! in CI would notice. +//! +//! The matrix row this discharges: +//! +//! ```text +//! generation dispatch → index.json.version 2→3 routing correct; +//! unknown version fails naming both sides +//! ``` + +use std::path::Path; + +use larql_vindex::format::filenames::INDEX_JSON; +use larql_vindex::format::generation::{ + detect_generation, ContainerGeneration, V2_CURRENT_SCHEMA, V3_CURRENT_SCHEMA, +}; +use larql_vindex::format::load::load_vindex_config; +use larql_vindex::VindexError; +use tempfile::tempdir; + +/// A VINDEX2 `index.json` with the fields the shipped loader requires. +fn write_v2_index(dir: &Path) { + std::fs::write( + dir.join(INDEX_JSON), + serde_json::json!({ + "version": V2_CURRENT_SCHEMA, + "model": "e0-fixture", + "family": "llama", + "num_layers": 2, + "hidden_size": 8, + "intermediate_size": 8, + "vocab_size": 4, + "embed_scale": 1.0, + "layers": [], + "down_top_k": 0 + }) + .to_string(), + ) + .unwrap(); +} + +/// A VINDEX3 `index.json`, carrying keys the v1 config struct has never seen. +fn write_v3_index(dir: &Path) { + std::fs::write( + dir.join(INDEX_JSON), + serde_json::json!({ + "version": V3_CURRENT_SCHEMA, + "moe_manifest": "moe_manifest.json", + "profiles": ["exact", "browse"], + "segments": { "routed/layer_000": 2 } + }) + .to_string(), + ) + .unwrap(); +} + +#[test] +fn a_shipped_generation_index_is_detected_as_vindex2() { + let dir = tempdir().unwrap(); + write_v2_index(dir.path()); + assert_eq!( + detect_generation(dir.path()).unwrap(), + ContainerGeneration::V2 + ); +} + +#[test] +fn a_successor_index_is_detected_as_vindex3() { + let dir = tempdir().unwrap(); + write_v3_index(dir.path()); + assert_eq!( + detect_generation(dir.path()).unwrap(), + ContainerGeneration::V3 + ); +} + +#[test] +fn the_shipped_loader_still_loads_a_shipped_index() { + // The row that matters most: VINDEX3 work must not have broken VINDEX2. + let dir = tempdir().unwrap(); + write_v2_index(dir.path()); + let config = load_vindex_config(dir.path()).expect("VINDEX2 must still load"); + assert_eq!(config.version, V2_CURRENT_SCHEMA); + assert_eq!(config.num_layers, 2); +} + +#[test] +fn the_shipped_loader_refuses_a_successor_index_by_generation() { + // Without this check the v2 config struct deserialises a v3 index on the + // strength of shared field names and proceeds against a layout whose + // weights live somewhere else — a served model with wrong weights. + let dir = tempdir().unwrap(); + write_v3_index(dir.path()); + let Err(err) = load_vindex_config(dir.path()) else { + panic!("must refuse a VINDEX3 directory"); + }; + let text = err.to_string(); + assert!( + matches!(err, VindexError::WrongContainerGeneration { .. }), + "{text}" + ); + assert!(text.contains("VINDEX3"), "{text}"); + assert!(text.contains("VINDEX2"), "{text}"); +} + +#[test] +fn an_unknown_generation_names_the_version_found_and_both_supported() { + let dir = tempdir().unwrap(); + std::fs::write( + dir.path().join(INDEX_JSON), + serde_json::json!({ "version": 99 }).to_string(), + ) + .unwrap(); + let err = detect_generation(dir.path()).unwrap_err(); + let text = err.to_string(); + assert!(text.contains("99"), "{text}"); + assert!(text.contains("VINDEX2"), "{text}"); + assert!(text.contains("VINDEX3"), "{text}"); +} + +#[test] +fn an_index_without_a_version_field_is_refused_rather_than_guessed() { + let dir = tempdir().unwrap(); + std::fs::write( + dir.path().join(INDEX_JSON), + serde_json::json!({ "model": "no-version" }).to_string(), + ) + .unwrap(); + assert!(detect_generation(dir.path()).is_err()); +} + +#[test] +fn a_legacy_schema_index_still_loads_as_the_shipped_generation() { + // The regression E0 caught. index.json version 1 predates several fields + // and loads with defaults; it is the same container generation, and + // refusing it broke exactly the compatibility this matrix protects. + let dir = tempdir().unwrap(); + std::fs::write( + dir.path().join(INDEX_JSON), + serde_json::json!({ + "version": 1, + "model": "legacy", + "family": "test", + "num_layers": 1, + "hidden_size": 4, + "intermediate_size": 4, + "vocab_size": 4, + "embed_scale": 1.0, + "layers": [], + "down_top_k": 0 + }) + .to_string(), + ) + .unwrap(); + assert_eq!( + detect_generation(dir.path()).unwrap(), + ContainerGeneration::V2 + ); + let config = load_vindex_config(dir.path()).expect("legacy schema must still load"); + assert_eq!(config.version, 1); +} + +#[test] +fn a_missing_index_json_is_an_io_error_not_a_generation_verdict() { + let dir = tempdir().unwrap(); + assert!(matches!( + detect_generation(dir.path()), + Err(VindexError::Io(_)) + )); +} + +#[test] +fn a_lyrw_v2_layer_file_is_refused_by_the_shipped_parser() { + // The other half of the boundary. A VINDEX3 layer file parsed at the v1 + // stride would not bounds-fail — it would yield offsets still inside the + // file and hand back bytes from the wrong expert. + use larql_vindex::format::weights::write_layers::parse_layer_weights_header; + + let mut header = Vec::new(); + header.extend_from_slice(&u32::from_le_bytes(*b"LYRW").to_le_bytes()); + header.extend_from_slice(&2u32.to_le_bytes()); // format_version = 2 + header.extend_from_slice(&[0u8; 64]); + + assert!( + parse_layer_weights_header(&header).is_none(), + "the shipped parser must refuse a LYRW v2 file" + ); +} + +#[test] +fn a_lyrw_v1_layer_file_is_still_accepted_by_the_shipped_parser() { + // Refusing the successor must not have broken the incumbent. + use larql_vindex::format::weights::write_layers::parse_layer_weights_header; + + let mut header = Vec::new(); + header.extend_from_slice(&u32::from_le_bytes(*b"LYRW").to_le_bytes()); + header.extend_from_slice(&1u32.to_le_bytes()); // format_version = 1 + header.extend_from_slice(&4u32.to_le_bytes()); // quant_format = Q4_K + header.extend_from_slice(&0u32.to_le_bytes()); // num_entries + header.extend_from_slice(&8u32.to_le_bytes()); // intermediate + header.extend_from_slice(&8u32.to_le_bytes()); // hidden + + assert!( + parse_layer_weights_header(&header).is_some(), + "the shipped parser must still read LYRW v1" + ); +} diff --git a/crates/larql-vindex/tests/vindex3_allocation_guard.rs b/crates/larql-vindex/tests/vindex3_allocation_guard.rs new file mode 100644 index 000000000..248d199ea --- /dev/null +++ b/crates/larql-vindex/tests/vindex3_allocation_guard.rs @@ -0,0 +1,214 @@ +//! Allocation guard for the VINDEX3 reference execution path. +//! +//! # Why this exists +//! +//! `decode_at` rebuilt its operand's name on every scalar read, allocating a +//! `String` per element. It was correct, it was covered by tests, and it cost +//! ~17× — nothing in the suite could see it, and the perf bench only caught it +//! because the noise happened to differ between two storage arrangements. +//! +//! # The invariant +//! +//! > Operand lookup and scalar access must not allocate in proportion to +//! > tensor dimensions, expert population, or number of elements read. +//! +//! Deliberately *not* zero-allocation. `execute` returns owned buffers and +//! allocates per-expert scratch, so a per-token count of a few dozen is +//! expected and fine. What must not happen is that count tracking the work. +//! When an `execute_into` path with caller-owned scratch arrives, this +//! tightens to zero per token. +//! +//! The counting allocator is declared here rather than in the library, so it +//! applies to this test binary alone and never to production builds. + +use std::alloc::{GlobalAlloc, Layout, System}; +use std::cell::Cell; + +use larql_vindex::runtime::execute; +use larql_vindex::runtime::fixtures::synthetic::{SyntheticLayer, SyntheticShape}; +use larql_vindex::runtime::MoeInputs; + +thread_local! { + /// Allocations on *this* thread. + /// + /// Per-thread, not a global atomic. `cargo test` runs test functions + /// concurrently, so a process-wide counter measures whatever else happens + /// to be running and the numbers move between runs — which is exactly the + /// failure this instrument produced before being fixed. + /// + /// `const`-initialised so that reading the counter cannot itself allocate + /// and re-enter the allocator. + static ALLOCATIONS: Cell = const { Cell::new(0) }; +} + +fn allocations() -> usize { + ALLOCATIONS.try_with(Cell::get).unwrap_or(0) +} + +fn record() { + // `try_with`: during thread teardown the TLS slot is gone, and panicking + // inside the allocator would abort. + let _ = ALLOCATIONS.try_with(|c| c.set(c.get().wrapping_add(1))); +} + +struct CountingAllocator; + +unsafe impl GlobalAlloc for CountingAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + record(); + unsafe { System.alloc(layout) } + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + unsafe { System.dealloc(ptr, layout) } + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + record(); + unsafe { System.realloc(ptr, layout, new_size) } + } +} + +#[global_allocator] +static ALLOCATOR: CountingAllocator = CountingAllocator; + +/// Allocations performed by one `execute`, measured after a warm-up call. +/// +/// Binding happens in the caller, deliberately. Binding allocates — it builds +/// operand names, contracts and expert lists — and it is load-path work; doing +/// it inside the measurement would put load-path allocations in a token-path +/// number and make the result depend on call ordering. +fn allocations_per_token(layer: &SyntheticLayer) -> usize { + let operation = layer.bind(&layer.bytes); + let residual = layer.residual(); + // First-call effects are not what this measures. + execute(&operation, MoeInputs::shared(&residual)).expect("synthetic layer executes"); + measure(&layer.bind(&layer.bytes), &residual) +} + +/// Allocations attributable to a single already-bound execution. +fn measure(operation: &larql_vindex::runtime::BoundMoeOperation<'_>, residual: &[f32]) -> usize { + let before = allocations(); + let out = execute(operation, MoeInputs::shared(residual)).expect("synthetic layer executes"); + let after = allocations(); + // Keep the result alive so its buffer is not freed inside the window. + std::hint::black_box(&out); + after - before +} + +fn shape(population: usize, hidden: usize, intermediate: usize) -> SyntheticShape { + SyntheticShape { + population, + hidden, + intermediate, + top_k: 2, + } +} + +/// Slack for allocations that are constant in count but conditional on size. +/// +/// Measured at 18 / 18 / 19 for populations 8 / 32 / 512. The single extra +/// allocation above ~20 experts is `slice::sort_by`'s merge scratch: the +/// stable sort uses insertion sort below that threshold and allocates a buffer +/// above it. One buffer, whatever the length — a constant, not a rate. +const CONSTANT_ALLOCATION_SLACK: usize = 4; + +#[test] +fn allocation_count_does_not_track_expert_population() { + // 64× the experts. Router scoring grows; allocation must not. + let small = SyntheticLayer::build(shape(8, 32, 24)); + let large = SyntheticLayer::build(shape(512, 32, 24)); + + let small_allocs = allocations_per_token(&small); + let large_allocs = allocations_per_token(&large); + + assert!( + large_allocs <= small_allocs + CONSTANT_ALLOCATION_SLACK, + "allocations tracked population: {small_allocs} at 8 experts, \ + {large_allocs} at 512 — more than the {CONSTANT_ALLOCATION_SLACK} \ + allowed for size-conditional constants" + ); +} + +#[test] +fn allocation_count_is_nowhere_near_proportional_to_population() { + // The assertion above is a bound; this one states the shape. Were the + // count proportional, 64× the experts would give ~64× the allocations. + let small = SyntheticLayer::build(shape(8, 32, 24)); + let large = SyntheticLayer::build(shape(512, 32, 24)); + + let small_allocs = allocations_per_token(&small); + let large_allocs = allocations_per_token(&large); + let proportional = small_allocs * 64; + + assert!( + large_allocs * 4 < proportional, + "{large_allocs} allocations at 512 experts is within 4× of the \ + {proportional} a per-expert allocation would produce" + ); +} + +#[test] +fn allocation_count_does_not_track_elements_read() { + // 16× the elements per expert, same population and top-k. This is the + // shape of the bug that motivated the guard: a per-element allocation + // shows up here and nowhere else. + let narrow = SyntheticLayer::build(shape(16, 16, 16)); + let wide = SyntheticLayer::build(shape(16, 64, 64)); + + let narrow_allocs = allocations_per_token(&narrow); + let wide_allocs = allocations_per_token(&wide); + + assert_eq!( + narrow_allocs, wide_allocs, + "allocations moved with tensor size: {narrow_allocs} at 16×16, \ + {wide_allocs} at 64×64 — this is exactly the shape of the \ + per-element `String` bug this guard exists for" + ); +} + +#[test] +fn allocation_count_stays_small_and_constant_per_token() { + // A loose absolute ceiling. The exact number is an implementation detail + // and will drop when `execute_into` lands; what this pins is that it is a + // small constant rather than a function of the work. + const CEILING: usize = 64; + let layer = SyntheticLayer::build(shape(64, 32, 24)); + let allocs = allocations_per_token(&layer); + assert!( + allocs <= CEILING, + "{allocs} allocations per token exceeds the ceiling of {CEILING}" + ); +} + +#[test] +fn repeated_tokens_allocate_identically() { + // Growth across repeats would mean something accumulates per token — a + // cache, a log, a memo — which is how a steady-state leak begins. + let layer = SyntheticLayer::build(shape(32, 32, 24)); + let operation = layer.bind(&layer.bytes); + let residual = layer.residual(); + execute(&operation, MoeInputs::shared(&residual)).expect("warm-up"); + + let first = measure(&operation, &residual); + for token in 1..16 { + assert_eq!( + measure(&operation, &residual), + first, + "token {token} allocated differently from the first" + ); + } +} + +#[test] +fn the_counting_allocator_is_actually_observing_allocations() { + // A guard whose instrument is broken passes every test. This one would + // fail if the global allocator were not installed. + let before = allocations(); + let v: Vec = Vec::with_capacity(1_024); + std::hint::black_box(&v); + assert!( + allocations() > before, + "the counting allocator recorded nothing" + ); +} diff --git a/crates/larql-vindex/tests/vindex3_container_parity.rs b/crates/larql-vindex/tests/vindex3_container_parity.rs new file mode 100644 index 000000000..8b779f7c4 --- /dev/null +++ b/crates/larql-vindex/tests/vindex3_container_parity.rs @@ -0,0 +1,456 @@ +//! V2-1, fixture A: does a **VINDEX3 container** execute? +//! +//! # The gap this closes +//! +//! Every VINDEX3 parity result before this one bound its operands out of a +//! VINDEX2 file — `index.json.version` was 2 at every step, on every model. +//! What was proven was the *runtime*: +//! +//! ```text +//! proven the VINDEX3 executor, fed VINDEX2 operands, matches production +//! not proven a VINDEX3 container can be written, opened, bound and executed +//! ``` +//! +//! Here both arms execute the identical arithmetic on identical weights. The +//! only thing that differs is **where the bytes came from**: +//! +//! ```text +//! arm A operands from an in-memory buffer (how every prior test sourced them) +//! arm B operands from a VINDEX3 container on disk, resolved through +//! index.json v3 → moe_manifest.json → LYRW v2 region table +//! ``` +//! +//! Bit-identical output is the whole claim. Anything less means the container +//! moved a byte — a wrong offset, a transposed region, a role resolved to the +//! wrong schema — and no amount of runtime correctness would rescue it. +//! +//! # Why bit-identity is the right bar here +//! +//! Both arms run the same executor over the same f32 bytes in the same order, +//! so there is no reordering and no quantisation between them. A tolerance +//! would therefore only hide a real defect: the two paths are either reading +//! the same bytes or they are not. + +use larql_compute::{Activation, MoeTopKWeightPolicy}; +use larql_vindex::format::capability::binding::RepresentationIdentity; +use larql_vindex::format::capability::component::ComponentContract; +use larql_vindex::format::capability::coordinate::BankCoordinate; +use larql_vindex::format::lyrw2::read::Lyrw2Reader; +use larql_vindex::format::lyrw2::region_format::RegionFormat; +use larql_vindex::format::lyrw2::region_role::RegionRole; +use larql_vindex::format::vindex3::test_support::{ + down_f32, fixture_a_fused_spec, fixture_a_spec, gate_f32, routed_spec, router_f32, + router_f32_for, up_f32, FIXTURE_A_BANK_ID, FIXTURE_A_EXPERTS, FIXTURE_A_HIDDEN, + FIXTURE_A_INTERMEDIATE, FIXTURE_A_LAYER, FIXTURE_A_SEGMENT_KEY, FIXTURE_A_TOP_K, +}; +use larql_vindex::format::vindex3::{write_container, Vindex3Container}; +use larql_vindex::runtime::{ + execute, BoundBankOperation, BoundExpert, BoundExpertScaling, BoundMoeOperation, + BoundProjection, BoundReduction, BoundRouter, BoundTensor, ExpertKernel, MoeInputs, + RouterKernel, +}; +use tempfile::tempdir; + +const VARIANT: &str = "exact"; +const ROUTER_REGION_SET: &str = "router"; +/// SwiGLU — what `gated-mlp-v1` contracts for. +const ACTIVATION: Activation = Activation::Silu; +/// Fixture A's router keeps raw softmax probabilities (no renormalisation), +/// matching `RouterPostProcessing::default()` in its manifest. +const SELECTED_WEIGHT: MoeTopKWeightPolicy = MoeTopKWeightPolicy::RawSoftmax; + +fn as_bytes(values: &[f32]) -> Vec { + values.iter().flat_map(|v| v.to_le_bytes()).collect() +} + +fn tensor<'a>(region_set: &str, bytes: &'a [u8], contract: ComponentContract) -> BoundTensor<'a> { + BoundTensor::direct( + RepresentationIdentity::new(region_set, VARIANT), + bytes, + RegionFormat::F32, + contract, + ) + .expect("bind f32 tensor") +} + +fn projection_contract() -> ComponentContract { + ComponentContract::matrix(FIXTURE_A_INTERMEDIATE, FIXTURE_A_HIDDEN) +} + +fn down_contract() -> ComponentContract { + ComponentContract::matrix(FIXTURE_A_HIDDEN, FIXTURE_A_INTERMEDIATE) +} + +fn operation<'a>(router_bytes: &'a [u8], experts: Vec>) -> BoundMoeOperation<'a> { + operation_for(router_bytes, experts, FIXTURE_A_EXPERTS, FIXTURE_A_TOP_K) +} + +/// The same operation with the population and top-k supplied rather than +/// assumed — what the not-hard-coded arm drives. +fn operation_for<'a>( + router_bytes: &'a [u8], + experts: Vec>, + population: u32, + top_k: usize, +) -> BoundMoeOperation<'a> { + BoundMoeOperation { + router: BoundRouter { + weight: tensor( + ROUTER_REGION_SET, + router_bytes, + ComponentContract::matrix(population, FIXTURE_A_HIDDEN), + ), + top_k, + selected_weight: SELECTED_WEIGHT, + scaling: BoundExpertScaling::None, + kernel: RouterKernel::Incumbent, + }, + transforms: Vec::new(), + banks: vec![BoundBankOperation { + bank: BankCoordinate::new(FIXTURE_A_LAYER, FIXTURE_A_BANK_ID), + experts, + intermediate_dim: FIXTURE_A_INTERMEDIATE as usize, + hidden_dim: FIXTURE_A_HIDDEN as usize, + activation: ACTIVATION, + kernel: ExpertKernel::Reference, + }], + reduction: BoundReduction::WeightedSum, + residual_dim: FIXTURE_A_HIDDEN as usize, + } +} + +/// A deterministic residual to route. Not from the fixture's own ramp, so a +/// bug that happened to align input and weights cannot cancel out. +fn residual() -> Vec { + (0..FIXTURE_A_HIDDEN as usize) + .map(|i| ((i as f32) * 0.021 - 1.7).cos()) + .collect() +} + +#[test] +fn a_vindex3_container_executes_bit_identically_to_in_memory_operands() { + // ── arm A: operands held in memory, as every prior parity test sourced them + let router = as_bytes(&router_f32()); + let gates: Vec> = (0..FIXTURE_A_EXPERTS) + .map(|e| as_bytes(&gate_f32(e))) + .collect(); + let ups: Vec> = (0..FIXTURE_A_EXPERTS) + .map(|e| as_bytes(&up_f32(e))) + .collect(); + let downs: Vec> = (0..FIXTURE_A_EXPERTS) + .map(|e| as_bytes(&down_f32(e))) + .collect(); + + let memory_experts: Vec> = (0..FIXTURE_A_EXPERTS) + .map(|e| BoundExpert { + expert_id: e, + projection: BoundProjection::Decomposed { + gate: tensor( + &RegionRole::Gate.name(), + &gates[e as usize], + projection_contract(), + ), + up: tensor( + &RegionRole::Up.name(), + &ups[e as usize], + projection_contract(), + ), + }, + down: tensor( + &RegionRole::Down.name(), + &downs[e as usize], + down_contract(), + ), + }) + .collect(); + let from_memory = execute( + &operation(&router, memory_experts), + MoeInputs::shared(&residual()), + ) + .expect("in-memory arm executes"); + + // ── arm B: the same weights, sourced through a VINDEX3 container on disk + let dir = tempdir().unwrap(); + write_container(dir.path(), &fixture_a_spec()).expect("write VINDEX3 container"); + let container = Vindex3Container::open(dir.path()).expect("open VINDEX3 container"); + + // Everything below goes through the container's own declarations: the + // manifest names the bank's storage key, the index resolves the key to a + // file, and the LYRW v2 region table resolves (entry, role) to bytes. + let layer = container + .layer(FIXTURE_A_LAYER) + .expect("manifest declares layer 0 as MoE"); + let segment = container + .segment(&layer.routed_bank.storage) + .expect("the bank's declared storage resolves"); + + let container_experts: Vec> = (0..FIXTURE_A_EXPERTS) + .map(|e| BoundExpert { + expert_id: e, + projection: BoundProjection::Decomposed { + gate: tensor( + &RegionRole::Gate.name(), + region(&segment, e, RegionRole::Gate), + projection_contract(), + ), + up: tensor( + &RegionRole::Up.name(), + region(&segment, e, RegionRole::Up), + projection_contract(), + ), + }, + down: tensor( + &RegionRole::Down.name(), + region(&segment, e, RegionRole::Down), + down_contract(), + ), + }) + .collect(); + let from_container = execute( + &operation(&router, container_experts), + MoeInputs::shared(&residual()), + ) + .expect("container arm executes"); + + assert_eq!( + from_memory.len(), + FIXTURE_A_HIDDEN as usize, + "the operation must return the residual width" + ); + assert_eq!( + bits(&from_container), + bits(&from_memory), + "a VINDEX3 container must deliver the same bytes it was given — any \ + difference here is the container moving a byte, not the executor" + ); +} + +fn region<'a>(reader: &Lyrw2Reader<'a>, entry: u32, role: RegionRole) -> &'a [u8] { + reader + .region_bytes(FIXTURE_A_BANK_ID, entry, role) + .unwrap_or_else(|e| panic!("resolve {role:?} for entry {entry}: {e}")) + .unwrap_or_else(|| panic!("{role:?} region missing for entry {entry}")) +} + +fn bits(v: &[f32]) -> Vec { + v.iter().map(|x| x.to_bits()).collect() +} + +/// The container's regions must equal the fixture's own weights byte-for-byte. +/// +/// Separate from the execution test on purpose: if both fail, this one says the +/// container stored the wrong bytes; if only the execution test fails, the +/// bytes were right and the binding read them wrongly. One assertion cannot +/// distinguish those. +#[test] +fn every_region_round_trips_byte_for_byte_through_the_container() { + let dir = tempdir().unwrap(); + write_container(dir.path(), &fixture_a_spec()).expect("write"); + let container = Vindex3Container::open(dir.path()).expect("open"); + let segment = container.segment(FIXTURE_A_SEGMENT_KEY).expect("segment"); + + for e in 0..FIXTURE_A_EXPERTS { + assert_eq!( + region(&segment, e, RegionRole::Gate), + as_bytes(&gate_f32(e)), + "gate bytes differ for expert {e}" + ); + assert_eq!( + region(&segment, e, RegionRole::Up), + as_bytes(&up_f32(e)), + "up bytes differ for expert {e}" + ); + assert_eq!( + region(&segment, e, RegionRole::Down), + as_bytes(&down_f32(e)), + "down bytes differ for expert {e}" + ); + } +} + +/// V2-1: fused and decomposed FC1 storage produce identical results **under +/// one manifest**. +/// +/// `gated-mlp-v1` declares both `gate + up + down` and `gate_up_fused + down` +/// legal, so the same programme id describes both containers below. Only the +/// physical rendering differs. +/// +/// This is the arm that proves the manifest is an *executable type +/// declaration* rather than a replay of one known layout: the logical operator +/// contract is fixed by the programme, and storage shape is free underneath +/// it. K3 depends on that being true — fused/decomposed, latent transforms and +/// shared projections all have superficially compatible dimensions while +/// requiring different execution semantics, so shape must never be what +/// selects the operation. +#[test] +fn fused_and_decomposed_storage_agree_under_one_programme() { + let router = as_bytes(&router_f32()); + + let decomposed_dir = tempdir().unwrap(); + write_container(decomposed_dir.path(), &fixture_a_spec()).expect("write decomposed"); + let decomposed = Vindex3Container::open(decomposed_dir.path()).expect("open decomposed"); + + let fused_dir = tempdir().unwrap(); + write_container(fused_dir.path(), &fixture_a_fused_spec()).expect("write fused"); + let fused = Vindex3Container::open(fused_dir.path()).expect("open fused"); + + // Same programme id on both sides — that is the point. + let d_layer = decomposed.layer(FIXTURE_A_LAYER).expect("decomposed layer"); + let f_layer = fused.layer(FIXTURE_A_LAYER).expect("fused layer"); + assert_eq!( + d_layer.routed_bank.programme, f_layer.routed_bank.programme, + "both renderings must be described by the same programme" + ); + assert_ne!( + d_layer.routed_bank.storage, f_layer.routed_bank.storage, + "…while resolving to physically different storage" + ); + + let d_seg = decomposed + .segment(&d_layer.routed_bank.storage) + .expect("decomposed segment"); + let f_seg = fused + .segment(&f_layer.routed_bank.storage) + .expect("fused segment"); + + let d_experts: Vec> = (0..FIXTURE_A_EXPERTS) + .map(|e| BoundExpert { + expert_id: e, + projection: BoundProjection::Decomposed { + gate: tensor( + &RegionRole::Gate.name(), + region(&d_seg, e, RegionRole::Gate), + projection_contract(), + ), + up: tensor( + &RegionRole::Up.name(), + region(&d_seg, e, RegionRole::Up), + projection_contract(), + ), + }, + down: tensor( + &RegionRole::Down.name(), + region(&d_seg, e, RegionRole::Down), + down_contract(), + ), + }) + .collect(); + + let f_experts: Vec> = (0..FIXTURE_A_EXPERTS) + .map(|e| BoundExpert { + expert_id: e, + projection: BoundProjection::Fused { + gate_up: tensor( + &RegionRole::GateUpFused.name(), + region(&f_seg, e, RegionRole::GateUpFused), + ComponentContract::matrix(2 * FIXTURE_A_INTERMEDIATE, FIXTURE_A_HIDDEN), + ), + }, + down: tensor( + &RegionRole::Down.name(), + region(&f_seg, e, RegionRole::Down), + down_contract(), + ), + }) + .collect(); + + let from_decomposed = execute( + &operation(&router, d_experts), + MoeInputs::shared(&residual()), + ) + .expect("decomposed executes"); + let from_fused = execute( + &operation(&router, f_experts), + MoeInputs::shared(&residual()), + ) + .expect("fused executes"); + + assert_eq!( + bits(&from_fused), + bits(&from_decomposed), + "storage layout must not change the logical result — if it does, the \ + programme id is not what selects the operation and shape is leaking \ + into the contract" + ); +} + +/// V2-1: expert count and top-K are read from the container, not baked in. +/// +/// Three populations and three top-Ks through the identical code path. If any +/// constant had leaked into the writer, the reader or the executor, at least +/// one shape would fail to bind or would execute against the wrong population. +/// +/// **Scope, stated honestly.** This discharges two of the three axes in V2-1's +/// "expert counts, top-K and shared banks are demonstrably not hard-coded". +/// The shared-bank axis is *not* covered: `BoundMoeOperation::banks` documents +/// unrouted shared banks as arriving with the Mini-K3 rung, so there is no +/// execution path to test them against yet. That row stays open. +#[test] +fn expert_count_and_top_k_come_from_the_container() { + // Deliberately varied together and apart: (8,2) is fixture A, (32,4) is + // GPT-OSS-shaped, (5,1) is a degenerate small case with an odd population + // that no power-of-two assumption would survive. + const SHAPES: [(u32, usize); 3] = [(8, 2), (32, 4), (5, 1)]; + + for (experts, top_k) in SHAPES { + let key = format!("routed/pop_{experts}_k{top_k}"); + let dir = tempdir().unwrap(); + write_container(dir.path(), &routed_spec(experts, top_k, &key)) + .unwrap_or_else(|e| panic!("write {experts}x{top_k}: {e}")); + let container = Vindex3Container::open(dir.path()) + .unwrap_or_else(|e| panic!("open {experts}x{top_k}: {e}")); + + // The container's own declarations, not the test's expectations. + let layer = container.layer(FIXTURE_A_LAYER).expect("MoE layer"); + assert_eq!(layer.routed_bank.experts, experts); + assert!( + container.verify().is_empty(), + "{experts}x{top_k} must be structurally bindable: {:?}", + container.verify() + ); + + let segment = container + .segment(&layer.routed_bank.storage) + .expect("segment"); + assert_eq!( + segment.banks()[0].num_entries, + experts, + "the bank must hold the population the manifest declares" + ); + + let router = as_bytes(&router_f32_for(experts)); + let bound: Vec> = (0..experts) + .map(|e| BoundExpert { + expert_id: e, + projection: BoundProjection::Decomposed { + gate: tensor( + &RegionRole::Gate.name(), + region(&segment, e, RegionRole::Gate), + projection_contract(), + ), + up: tensor( + &RegionRole::Up.name(), + region(&segment, e, RegionRole::Up), + projection_contract(), + ), + }, + down: tensor( + &RegionRole::Down.name(), + region(&segment, e, RegionRole::Down), + down_contract(), + ), + }) + .collect(); + + let out = execute( + &operation_for(&router, bound, experts, top_k), + MoeInputs::shared(&residual()), + ) + .unwrap_or_else(|e| panic!("execute {experts}x{top_k}: {e}")); + assert_eq!(out.len(), FIXTURE_A_HIDDEN as usize); + assert!( + out.iter().all(|v| v.is_finite()), + "{experts}x{top_k} produced non-finite output" + ); + } +} diff --git a/docs/adr/0002-ffn-activation-cache.md b/docs/adr/0002-ffn-activation-cache.md index 85bd5e8b3..b816db3e4 100644 --- a/docs/adr/0002-ffn-activation-cache.md +++ b/docs/adr/0002-ffn-activation-cache.md @@ -95,7 +95,7 @@ Pre-seeded from 1,923 labelled features × 34 layers. Write-back from server on - Patched session (after INSERT) → cache bypassed for that layer - Cost: a cache miss on every call to a patched layer — which is correct, since the output changes with the patch -This is tested explicitly in `examples/ffn_cache_demo.rs` (Scenario 3). +This is tested explicitly in `crates/larql-demos/examples/inference/ffn_cache_demo.rs` (Scenario 3). --- @@ -129,7 +129,7 @@ The L2 gate-KNN call in `run_full_output` uses the request's `top_k` to derive t | `crates/larql-server/src/ffn_l2_cache.rs` | `FfnL2Cache` struct + unit tests | | `crates/larql-server/src/state.rs` | `LoadedModel.ffn_l2_cache` field | | `crates/larql-server/src/routes/walk_ffn.rs` | L2 wired into `run_full_output` | -| `crates/larql-inference/examples/ffn_cache_demo.rs` | Demo: hit rates + patch safety | +| `crates/larql-demos/examples/inference/ffn_cache_demo.rs` | Demo: hit rates + patch safety | | `crates/larql-inference/examples/bench_ffn_cache.rs` | Benchmark: latency delta | | `docs/ffn-cache.md` | User-facing guide | diff --git a/docs/diagnoses/fr-routing-gain.md b/docs/diagnoses/fr-routing-gain.md index 7c1bf4a24..db064689e 100644 --- a/docs/diagnoses/fr-routing-gain.md +++ b/docs/diagnoses/fr-routing-gain.md @@ -1,6 +1,6 @@ # FR routing — the end-to-end gain (and the two-tier caveat), quantified -**Date:** 2026-06-07. **Status:** ran (`crates/larql-inference/examples/fr_routing_gain.rs`). Gemma-3-4B Q4K vindex, 20 novel facts installed at L26, three query slices, all three router modes on the **same** forward passes. Answers "does the FR work give a new gain?" +**Date:** 2026-06-07. **Status:** ran (`chris-experiments/larql_probes/examples/fleet_routing/fr_routing_gain.rs`). Gemma-3-4B Q4K vindex, 20 novel facts installed at L26, three query slices, all three router modes on the **same** forward passes. Answers "does the FR work give a new gain?" ## Headline @@ -36,4 +36,4 @@ The FR work delivers a **decisive correctness gain** — it converts KNN fact-injection from "corrupts 100% of unrelated queries at 20 facts" to "100% distractor-safe" — at **no throughput cost**. It is not a tok/s gain and was never going to be; the override is a sidecar. The measurement also sharpened the FR2 guidance: two-tier fallback is an opt-in alias resolver, not a safe default. -**Artifacts:** `crates/larql-inference/examples/fr_routing_gain.rs`. +**Artifacts:** `chris-experiments/larql_probes/examples/fleet_routing/fr_routing_gain.rs`. diff --git a/docs/diagnoses/fr1-topk-fuzzy-router.md b/docs/diagnoses/fr1-topk-fuzzy-router.md index 72c36874e..ab4ace709 100644 --- a/docs/diagnoses/fr1-topk-fuzzy-router.md +++ b/docs/diagnoses/fr1-topk-fuzzy-router.md @@ -1,6 +1,6 @@ # FR1 — top-k fuzzy entity router on a real LARQL vindex: VERDICT -**Date:** 2026-06-07. **Status:** ran (`crates/larql-inference/examples/fr1_topk_fuzzy_router.rs` → `bench/aim-validation/fr1_topk_router_gemma3-4b.json`). Gemma-3-4B Q4K vindex, production `KnnStore` cosine-NN router, N=150 real countries, layer sweep {20,22,24,26}, judged in predictive units (recall@k / margin / confident-wrong — mean-cosine banned). Reproduces fleet E15 on the production path. Pre-registration: [`docs/fleet-routing-extensions.md`](../fleet-routing-extensions.md) §FR1. +**Date:** 2026-06-07. **Status:** ran (`chris-experiments/larql_probes/examples/fleet_routing/fr1_topk_fuzzy_router.rs` → `bench/aim-validation/fr1_topk_router_gemma3-4b.json`). Gemma-3-4B Q4K vindex, production `KnnStore` cosine-NN router, N=150 real countries, layer sweep {20,22,24,26}, judged in predictive units (recall@k / margin / confident-wrong — mean-cosine banned). Reproduces fleet E15 on the production path. Pre-registration: [`docs/fleet-routing-extensions.md`](../fleet-routing-extensions.md) §FR1. ## Headline @@ -45,7 +45,7 @@ The fuzzy entity key is **real, strong, and answer-leak-free at L24-26** — and the production cosine-NN path already delivers it (better than a trained MLP). The defect is entirely in the **consumer**: a fixed-0.75 top-1 gate that fires on everything and injects an 11–84% confident-wrong rate. **Build greenlit**, scoped to: route at the resolved layer, replace the absolute gate with **top-k + verify + abstain**, parity-first. FR2 (symbolic-primary, this as the fuzzy fallback) is the natural wrapper. -**Artifacts:** `crates/larql-inference/examples/fr1_topk_fuzzy_router.rs`, `bench/aim-validation/fr1_topk_router_gemma3-4b.json`. +**Artifacts:** `chris-experiments/larql_probes/examples/fleet_routing/fr1_topk_fuzzy_router.rs`, `bench/aim-validation/fr1_topk_router_gemma3-4b.json`. --- diff --git a/docs/diagnoses/fr2-two-tier-router.md b/docs/diagnoses/fr2-two-tier-router.md index 83dfa87e1..28cbd0784 100644 --- a/docs/diagnoses/fr2-two-tier-router.md +++ b/docs/diagnoses/fr2-two-tier-router.md @@ -1,6 +1,6 @@ # FR2 — two-tier router (symbolic-primary → activation-fuzzy fallback): VERDICT -**Date:** 2026-06-07. **Status:** ran (`crates/larql-inference/examples/fr2_two_tier_router.rs` → `bench/aim-validation/fr2_two_tier_router_gemma3-4b.json`). Gemma-3-4B Q4K vindex, store over 115 canonical countries, 10 historical/alternate-name aliases, layers {24,26}. Reproduces fleet E16's alias slice on the production path. Pre-registration: [`docs/fleet-routing-extensions.md`](../fleet-routing-extensions.md) §FR2. Depends on FR1. +**Date:** 2026-06-07. **Status:** ran (`chris-experiments/larql_probes/examples/fleet_routing/fr2_two_tier_router.rs` → `bench/aim-validation/fr2_two_tier_router_gemma3-4b.json`). Gemma-3-4B Q4K vindex, store over 115 canonical countries, 10 historical/alternate-name aliases, layers {24,26}. Reproduces fleet E16's alias slice on the production path. Pre-registration: [`docs/fleet-routing-extensions.md`](../fleet-routing-extensions.md) §FR2. Depends on FR1. ## Headline @@ -35,7 +35,7 @@ A two-tier dispatch in the override (and surfaced in LQL): **exact-string (`entr Exact-string is the precise primary (1.0 on names); the activation key is the alias/paraphrase fallback that recovers what exact-match misses (0/10 → 10/10 on famous aliases, ~0.9 top-5 in general). Sequenced two-tier with the FR1 verifier bounding mis-routes, this is the resolved routing architecture — **build greenlit**, with the general-rate and confident-wrong caveats stated. -**Artifacts:** `crates/larql-inference/examples/fr2_two_tier_router.rs`, `bench/aim-validation/fr2_two_tier_router_gemma3-4b.json`. +**Artifacts:** `chris-experiments/larql_probes/examples/fleet_routing/fr2_two_tier_router.rs`, `bench/aim-validation/fr2_two_tier_router_gemma3-4b.json`. --- diff --git a/docs/diagnoses/fr3-explicit-rewrite.md b/docs/diagnoses/fr3-explicit-rewrite.md index f74de1f94..c8fe2b5da 100644 --- a/docs/diagnoses/fr3-explicit-rewrite.md +++ b/docs/diagnoses/fr3-explicit-rewrite.md @@ -1,6 +1,6 @@ # FR3b — relation resolution: probe is phrasing-brittle, explicit rewrite wins -**Date:** 2026-06-08. **Status:** ran (`examples/fr3_template_ablation.rs`, `examples/fr3_explicit_rewrite.rs` → `bench/aim-validation/fr3_{template_ablation,explicit_rewrite}_gemma3-4b.json`). Gemma-3-4B Q4K vindex. Follow-on to [`fr3-relation-address.md`](fr3-relation-address.md) — refines, doesn't overturn, the FR3 WIN. +**Date:** 2026-06-08. **Status:** ran (`chris-experiments/larql_probes/examples/fleet_routing/fr3_template_ablation.rs`, `chris-experiments/larql_probes/examples/fleet_routing/fr3_explicit_rewrite.rs` → `bench/aim-validation/fr3_{template_ablation,explicit_rewrite}_gemma3-4b.json`). Gemma-3-4B Q4K vindex. Follow-on to [`fr3-relation-address.md`](fr3-relation-address.md) — refines, doesn't overturn, the FR3 WIN. ## Headline @@ -53,13 +53,13 @@ Few-shot `word -> relation` over `{capital,currency,language[, none]}`, read top **Wiring wrinkle (the one real structural choice):** `RelationResolver` only dequantises layers `0..=probe_layer` (≈L10) → it **cannot run lm_head**, so Tier 2 must run via the **Session's already-loaded vindex** (`predict_kquant`/`InferenceWeights`, the same path INFER uses), not the resolver's partial setup. ~30 lines crossing the resolver→session boundary. Add an LQL knob if the explicit pass should be opt-in (it's a full forward per abstain). -Harnesses to lift the prompt/matching from: `examples/fr3_explicit_rewrite.rs` (the few-shot frame + `none`-gated accept + prefix-matching over top-k). +Harnesses to lift the prompt/matching from: `chris-experiments/larql_probes/examples/fleet_routing/fr3_explicit_rewrite.rs` (the few-shot frame + `none`-gated accept + prefix-matching over top-k). ## BUILD LANDED (2026-06-09) **Wired the two-tier resolver into `SELECT … FROM EDGES WHERE relation = …`, opt-in, default off = byte-identical.** When the exact/substring relation match returns nothing, `resolve_relation_synonym` runs Tier 1 (the cached residual probe, unchanged); on probe abstain it falls through (`.or_else`) to **Tier 2 — `resolve_relation_explicit`** (`crates/larql-lql/src/executor/query/select/edges.rs`): -- **Few-shot frame lifted verbatim** from `examples/fr3_explicit_rewrite.rs` (`word -> relation` + `music -> none`), one **full forward** via `InferenceWeights::predict_dense` (the INFER path — for a Q4_K vindex this is exactly `predict_kquant`, lm_head included). The resolver's partial `0..=L10` dequant can't run lm_head, so Tier 2 goes through `InferenceWeights`, not the resolver's setup — the one structural wrinkle the pre-registration called out. +- **Few-shot frame lifted verbatim** from `chris-experiments/larql_probes/examples/fleet_routing/fr3_explicit_rewrite.rs` (`word -> relation` + `music -> none`), one **full forward** via `InferenceWeights::predict_dense` (the INFER path — for a Q4_K vindex this is exactly `predict_kquant`, lm_head included). The resolver's partial `0..=L10` dequant can't run lm_head, so Tier 2 goes through `InferenceWeights`, not the resolver's setup — the one structural wrinkle the pre-registration called out. - **`none`-gated accept** (`match_relation_top1`, unit-tested): prefix-match top-1 against the candidate relations; `none` / out-of-domain → no match → abstain. - **Gated by `LARQL_FR3_EXPLICIT`** (full forward + model load per probe-abstain). Default off → SELECT is byte-identical to FR3 (probe-only). diff --git a/docs/diagnoses/fr3-relation-address.md b/docs/diagnoses/fr3-relation-address.md index 36f18aed5..5a6055ea3 100644 --- a/docs/diagnoses/fr3-relation-address.md +++ b/docs/diagnoses/fr3-relation-address.md @@ -1,6 +1,6 @@ # FR3 — relation as a clean semantic address: VERDICT -**Date:** 2026-06-07. **Status:** ran (`crates/larql-inference/examples/fr3_relation_address.rs` → `bench/aim-validation/fr3_relation_address_gemma3-4b.json`). Gemma-3-4B Q4K vindex, N=40 countries, layer sweep {6,10,14,20,26}. Dependency-free softmax-regression probe (standardised, L2), judged in synonym-generalisation accuracy (not mean-cosine). Reproduces the mechanism video `address.py` on the production path. Pre-registration: [`docs/fleet-routing-extensions.md`](../fleet-routing-extensions.md) §FR3. +**Date:** 2026-06-07. **Status:** ran (`chris-experiments/larql_probes/examples/fleet_routing/fr3_relation_address.rs` → `bench/aim-validation/fr3_relation_address_gemma3-4b.json`). Gemma-3-4B Q4K vindex, N=40 countries, layer sweep {6,10,14,20,26}. Dependency-free softmax-regression probe (standardised, L2), judged in synonym-generalisation accuracy (not mean-cosine). Reproduces the mechanism video `address.py` on the production path. Pre-registration: [`docs/fleet-routing-extensions.md`](../fleet-routing-extensions.md) §FR3. ## Headline @@ -39,7 +39,7 @@ The relation is a **clean semantic address, resolved early (L6) and synonym-robust (1.00)** — the opposite of the fuzzy late-resolving entity (FR1). The two halves of `(relation, entity) → value` measured side-by-side on the production residual confirm the mechanism: **build the relation half as a meaning-keyed index, the entity half as top-k + rank.** -**Artifacts:** `crates/larql-inference/examples/fr3_relation_address.rs`, `bench/aim-validation/fr3_relation_address_gemma3-4b.json`. +**Artifacts:** `chris-experiments/larql_probes/examples/fleet_routing/fr3_relation_address.rs`, `bench/aim-validation/fr3_relation_address_gemma3-4b.json`. --- diff --git a/docs/diagnoses/q4k-direct-attention.md b/docs/diagnoses/q4k-direct-attention.md index e0092b68e..f01a8e307 100644 --- a/docs/diagnoses/q4k-direct-attention.md +++ b/docs/diagnoses/q4k-direct-attention.md @@ -248,7 +248,7 @@ Two cheap probes (real production functions, synthetic same-size f32 weights → faithful bandwidth, no model load; `CpuBackend`, the no-`--metal` path the 28% was measured on). These are the #24-trap guard — run *before* building the path. -### Gate 1 — projection-vs-GQA split (`examples/attn_proj_vs_gqa_split.rs`) +### Gate 1 — projection-vs-GQA split (`chris-experiments/larql_probes/examples/q4k_attention/attn_proj_vs_gqa_split.rs`) How much of the attention block is the Q4K-accelerable projections vs the unaccelerated f32 GQA, across a cached_len sweep at 26B dims. **Projection cost @@ -271,7 +271,7 @@ With the sliding-window cap, projections stay ≥64% out to 4K and ~50% at 8K (t no-cap upper bound on GQA crosses over earlier, ~2.4K). The GQA residual is small at working context and is the only part Q4K-direct can't touch. -### Gate 2 — f32 BLAS vs Q4K-direct on the projection (`examples/attn_proj_f32_vs_q4k.rs`) +### Gate 2 — f32 BLAS vs Q4K-direct on the projection (`chris-experiments/larql_probes/examples/q4k_attention/attn_proj_f32_vs_q4k.rs`) The decisive question: does `q4k_matvec` actually beat Apple AMX/Accelerate f32 sgemm, or does AMX throughput eat the 7× bandwidth cut? Per-projection, same @@ -508,7 +508,7 @@ against the **real** decode denominator, not the synthetic 28%. within noise — both measured, not inferred.) - **Prefill twin — GATED then FALSIFIED.** At 907 ctx prefill attention is 6288 ms of a 14739 ms TTFT (~43% of prefill), which *looked* like the better lever. But - the prefill-shape Gate-2 (`examples/attn_prefill_f32_vs_q4k.rs`, seq_len=907) + the prefill-shape Gate-2 (`chris-experiments/larql_probes/examples/q4k_attention/attn_prefill_f32_vs_q4k.rs`, seq_len=907) kills it: repeated per-position `q4k_matvec` (the only CPU path — **no `q4k_matmul`**) is **~20× SLOWER** than one f32 BLAS sgemm (sliding block 25.9 vs 569 ms, 0.05×; global 42 vs 789 ms). Deeper reason: at prefill the projection diff --git a/docs/diagnoses/v1-hash-routing.md b/docs/diagnoses/v1-hash-routing.md index 4c389af69..c51052b44 100644 --- a/docs/diagnoses/v1-hash-routing.md +++ b/docs/diagnoses/v1-hash-routing.md @@ -1,7 +1,7 @@ # V1 — Hash routing across all layers (aim-validation, KU4) **Status:** COMPLETE on three dense archs (Gemma 3 4B, Llama 2 7B, Mistral 7B) — unanimous falsification. MoE (26B) deferred (see scope note). -**Harness:** `crates/larql-inference/examples/walk_ffn_v1_hash_routing.rs` +**Harness:** `chris-experiments/larql_probes/examples/walk_ffn/walk_ffn_v1_hash_routing.rs` **Artifacts:** `bench/aim-validation/v1_.json` **Date:** 2026-05-31 diff --git a/docs/diagnoses/v1-moe-within-expert.md b/docs/diagnoses/v1-moe-within-expert.md index da3913298..f883fe219 100644 --- a/docs/diagnoses/v1-moe-within-expert.md +++ b/docs/diagnoses/v1-moe-within-expert.md @@ -4,7 +4,7 @@ **Aim-validation:** resolves the OPEN half of KU4 (V1 was falsified on dense; the MoE-within-expert variant was left open because the dense harness measures the wrong object on the 26B-A4B). -**Harness:** `crates/larql-inference/examples/walk_ffn_v1_moe_within_expert.rs` +**Harness:** `chris-experiments/larql_probes/examples/walk_ffn/walk_ffn_v1_moe_within_expert.rs` **Kernel hook:** `crates/larql-compute/src/cpu/ops/moe/within_expert.rs` **Model:** Gemma 4 26B-A4B (`output/gemma4-26b-a4b-q4k.vindex`), 30 layers (all MoE), 128 experts, top_k=8, **expert inter = 704**. diff --git a/docs/diagnoses/v2-fp4-generality.md b/docs/diagnoses/v2-fp4-generality.md index 8df13e69e..9abfde3f5 100644 --- a/docs/diagnoses/v2-fp4-generality.md +++ b/docs/diagnoses/v2-fp4-generality.md @@ -3,7 +3,7 @@ **Status:** COMPLETE — CONFIRMED (the opposite of V1). FP4-friendliness is universal and near-lossless across the archs measured. **Harnesses:** `crates/larql-vindex/examples/fp4_q1_scan.rs` (static, generalized for -`*_weights.bin` naming), `crates/larql-inference/examples/walk_ffn_v2_fp4_nll.rs` (predictive). +`*_weights.bin` naming), `chris-experiments/larql_probes/examples/walk_ffn/walk_ffn_v2_fp4_nll.rs` (predictive). **Artifacts:** `bench/aim-validation/v2_*_scan.json` **Date:** 2026-05-31 diff --git a/docs/diagnoses/walk-ffn-performance.md b/docs/diagnoses/walk-ffn-performance.md index 9a845c70e..f919bcc40 100644 --- a/docs/diagnoses/walk-ffn-performance.md +++ b/docs/diagnoses/walk-ffn-performance.md @@ -60,7 +60,7 @@ dominate, and the FFN — sparse or dense — is a masked fraction. The ## FFN microbench — the missing instrument (built 2026-05-29) -`examples/walk_ffn_microbench.rs` isolates `WalkFfn::forward` at **seq_len = 1** +`chris-experiments/larql_probes/examples/walk_ffn/walk_ffn_microbench.rs` isolates `WalkFfn::forward` at **seq_len = 1** (decode shape), no attention/lm_head, across K. gemma3-4b-q4k, layer 17, 10240 features, 300 iters: @@ -125,7 +125,7 @@ predictive quality is measured next. ## Accuracy frontier — speed is cheap, accuracy is not (built 2026-05-29, task #19) -`examples/walk_ffn_accuracy.rs` runs a full forward (attention dequantised to f32 +`chris-experiments/larql_probes/examples/walk_ffn/walk_ffn_accuracy.rs` runs a full forward (attention dequantised to f32 up front via `insert_q4k_layer_tensors`, so the **FFN router is the only variable**) and scores the last-token next-token distribution against dense in the Shannon discipline — **KL in bits, top-1 agreement, q@p_argmax** — never @@ -350,7 +350,7 @@ most-frequent gate-KNN features across its members (capped). At inference, the per-position residual picks its nearest cell (O(C·hidden)) and that cell's pool is the candidate set — **content-addressed** (cell depends on the residual) but cheap (no full O(num_features) gate projection). Built and run in -`examples/walk_ffn_cell_router.rs` (C=64, calibration 24 prose prompts, 9-layer +`chris-experiments/larql_probes/examples/walk_ffn/walk_ffn_cell_router.rs` (C=64, calibration 24 prose prompts, 9-layer band, K=512). **All KL is measured against *dense* (KL 0 = dense); gate-KNN is itself a lossy top-K truncation of the gate projection, not a floor.** Lower = closer to dense. @@ -466,7 +466,7 @@ that the faithful-K regime is **kernel-bound, not FLOP-bound**: **So the faithful-K speedup exists in the FLOPs and is squandered by the kernel.** The optimization (task #24) — **gather the selected K rows contiguous, then run -the kernel** — was built and measured (`examples/walk_ffn_gather_gemm.rs`), and it +the kernel** — was built and measured (`chris-experiments/larql_probes/examples/walk_ffn/walk_ffn_gather_gemm.rs`), and it **flips faithful-K from slower-than-dense to faster** (isolated FFN, seq_len=1): | K | scattered (current) | gather Q4K + fused | dense | @@ -545,7 +545,7 @@ survives"), starting with the **4-layer static band**, top-1 agreement alongside ### End-to-end decode measurement — the win does not survive the full forward Pre-committed bar: net forward tok/s **> dense**. Measured three ways -(`examples/walk_ffn_decode_timing.rs`), each hitting a distinct confound — and +(`chris-experiments/larql_probes/examples/walk_ffn/walk_ffn_decode_timing.rs`), each hitting a distinct confound — and **none clears the bar**: | measurement | result | confound | @@ -579,7 +579,7 @@ validated components; the end-to-end payoff does not. Sparsity is closed (do-less flips tokens). The remaining doors are *fewer bytes per feature* or *more work per byte*. Tested the most graph-native: **graded precision** — keep all features, spend bits by ‖down_row‖ importance -(`examples/walk_ffn_graded_precision.rs`, block-wise quantiser validated: uniform +(`chris-experiments/larql_probes/examples/walk_ffn/walk_ffn_graded_precision.rs`, block-wise quantiser validated: uniform 4-bit = KL 0.011 vs f32, matching real Q4K). KL vs f32 reference, 4 prompts (in-dist + code + non-English): @@ -604,7 +604,7 @@ n=4) at 0.75× FFN bandwidth (~16% total → ~1.19× decode if it holds). **But single-step number oversold it.** **⚠️ Generation drift overturns the single-step story.** Greedy-decoding 10 -prompts to 32 tokens, sim-Q4 vs Q3 (`examples/walk_ffn_drift.rs`): **per-step +prompts to 32 tokens, sim-Q4 vs Q3 (`chris-experiments/larql_probes/examples/walk_ffn/walk_ffn_drift.rs`): **per-step argmax flip rate 19.1%, mean first-divergence token 4/32, 0/10 exact match.** The "100% top-1" was an n=4 artifact; KL 0.05 bits *sounds* tiny but LM distributions are full of near-ties, so a 0.05-bit perturbation flips ~19% of argmaxes → @@ -614,7 +614,7 @@ lesson on the sequence axis: single-step KL/top-1 can't see drift. **The three-way per-token NLL adjudicator decides it — and the mean would have lied.** f32 / Q4 / Q3 teacher-forced on entropic prose -(`examples/walk_ffn_nll.rs`, 74 positions): +(`chris-experiments/larql_probes/examples/walk_ffn/walk_ffn_nll.rs`, 74 positions): | arm | mean | **median** | p90 | p99 | max | mean Δ vs f32 | p90 Δ | p99 Δ | worst-token Δ | |---|---:|---:|---:|---:|---:|---:|---:|---:|---:| @@ -771,7 +771,7 @@ TEMPORAL axis dense BLAS structurally can't see (no cursor, no delta between tokens). Guardrail honored: token-to-token at FIXED layer, last-position residual over real history (`predict_with_ffn_trace` on teacher-forced prefixes = KV-cached decode step), never within-prefill cross-position (that would be the -spatial cosine wearing a temporal label). `examples/walk_ffn_temporal_reuse.rs`, +spatial cosine wearing a temporal label). `chris-experiments/larql_probes/examples/walk_ffn/walk_ffn_temporal_reuse.rs`, 6 entropic passages, per-zone distribution (median / p10 / worst), not mean: | zone | residual cosine (med/p10/worst) | pool Jaccard (med/p10/worst) | delta TwoNN | @@ -814,7 +814,7 @@ is a *full-amplitude* move, not a thin perturbation. Low-rank ≠ small. So befo any kernel, the cheap falsification: amplitude `‖δ‖/‖base‖` + full-Jacobian linearization error `‖f(base+δ)−(f(base)+Jδ)‖/‖f(base+δ)‖` (finite-diff JVP), **targeting the FFN-INPUT residual (post-attn-norm) — what the FFN actually -sees — not #27's layer-input residual.** `examples/walk_ffn_delta_walk.rs`: +sees — not #27's layer-input residual.** `chris-experiments/larql_probes/examples/walk_ffn/walk_ffn_delta_walk.rs`: | zone | ‖δ‖/‖base‖ med/p90/worst | lin-error med/p90/worst | |---|---|---| diff --git a/docs/diagnoses/walk-ffn-r4-zeroout.md b/docs/diagnoses/walk-ffn-r4-zeroout.md index 4d43caf55..516a4a42c 100644 --- a/docs/diagnoses/walk-ffn-r4-zeroout.md +++ b/docs/diagnoses/walk-ffn-r4-zeroout.md @@ -3,7 +3,7 @@ **Status:** RESOLVED 2026-08-01 — **all-layer dynamic row sparsity REFUTED.** **Rule applied:** R4 (`docs/dec-funnel.md` §"Standing rules") — before testing a lever, set the bytes it targets to zero. -**Harness:** `crates/larql-inference/examples/walk_ffn_r4_zeroout.rs` +**Harness:** `chris-experiments/larql_probes/examples/walk_ffn/walk_ffn_r4_zeroout.rs` **Artifacts:** `bench/aim-validation/r4_zeroout_ac_paired.txt`, `..._v2.txt` **Model:** Gemma 3 4B (`output/gemma3-4b-q4k-v2.vindex`), 34 layers, 10 240 features/layer, decode shape (seq_len = 1). diff --git a/docs/ffn-cache.md b/docs/ffn-cache.md index 29a904e71..2f7f3da00 100644 --- a/docs/ffn-cache.md +++ b/docs/ffn-cache.md @@ -71,7 +71,7 @@ This means: - An INSERT session → cache bypassed for layers that have overrides; active for layers without - The override check is per-layer, not per-session, so a session that only patches L10 still gets cache hits at L0–L9 and L11–L33 -This is validated in `examples/ffn_cache_demo.rs` (Scenario 3) and is the correct behaviour: correctness over hit rate for live-patched layers. +This is validated in `crates/larql-demos/examples/inference/ffn_cache_demo.rs` (Scenario 3) and is the correct behaviour: correctness over hit rate for live-patched layers. --- diff --git a/docs/ffn-graph-layer.md b/docs/ffn-graph-layer.md index 805277da5..25bc3fa85 100644 --- a/docs/ffn-graph-layer.md +++ b/docs/ffn-graph-layer.md @@ -296,6 +296,6 @@ let model = InferenceModel::load_walk_only("google/gemma-3-4b-it")?; | `larql-vindex/examples/build_down_features.rs` | Feature-major down vector builder | | `larql-vindex/examples/build_up_features.rs` | Feature-major up vector builder | | `larql-inference/examples/bench_walk_inference.rs` | Walk benchmark (dense vs walk vs HNSW) | -| `larql-inference/examples/walk_boundary_sweep.rs` | Correctness sweep (all 34 layers) | +| `larql-inference/chris-experiments/larql_probes/examples/misc/walk_boundary_sweep.rs` | Correctness sweep (all 34 layers) | | `larql-inference/examples/profile_overhead.rs` | Forward pass bottleneck profiler | | `larql-inference/examples/memory_analysis.rs` | Memory profiling (RSS, mmap, walk-only) | diff --git a/docs/lyrw-v2.md b/docs/lyrw-v2.md new file mode 100644 index 000000000..0d6a02e21 --- /dev/null +++ b/docs/lyrw-v2.md @@ -0,0 +1,359 @@ +# LYRW v2 — the K3 routed-layer physical-layout gate + +**Programme:** `lyrw` — six experiments deciding whether the K3 extraction can use today's per-layer weight format or needs one final revision first. +**Scope:** the *storage* half of the K3 work. The model-side half is [`k3-funnel.md`](k3-funnel.md); the transport half is [`dec-funnel.md`](dec-funnel.md); the production half is [`vindex-factory.md`](vindex-factory.md). +**Status:** v0.1 — draft. Three of the six experiments are materially narrowed by results already banked in `dec`; two are blocked; one blocker found while scoping outranks all six. +**Date:** 2026-08-01 + +--- + +## 0. The question, and what it is not + +The per-layer LYRW format ([`format-spec.md` §5.12](../crates/larql-vindex/docs/format-spec.md)) already stores, per expert, two independently-addressed regions with their own offsets and lengths: + +``` +[entry e gate+up] shape [2*inter, hidden] +[entry e down] shape [hidden, inter] +``` + +The limitation is one line of design principle and one field in the header: **the whole layer file declares a single `quant_format`**, so the two regions cannot carry different representations without replacing the file. + +So the programme does not need to ask "how should everything be sliced". It needs to answer three architectural questions — + +1. should gate/up and down support different formats? +2. does physical expert ordering or grouping materially improve SSD behaviour? +3. are any inference modes that omit weights good enough to support? + +— and leave everything else as runtime policy over the existing per-expert offset table. + +**What this programme is a gate on.** Extraction *feasibility and operations*, not throughput. Standing rule **R4** ([`dec-funnel.md` §1](dec-funnel.md)) applied here: this whole programme targets routed expert bytes, which `dec8-6` measured at **25.83 GB of 55.69 GB** per-token touch (46.4 %). Zero them entirely and the dense half alone still caps the machine at **12.3 tok/s**. **No layout result in this document can decide a throughput target**, and any arm that reports itself in tok/s is answering a question it cannot answer. Judge these experiments in bytes-written, bytes-faulted, wall-clock-to-rebuild, and exactness — the units the operation actually has ([[feedback_metric_matches_operation]]). + +--- + +## 1. What the banked evidence already settles + +Four results from the `dec` programme land directly on this design. Three of them delete work. + +### 1.1 The fidelity axis of Experiment 4 does not exist for K3 — **C1** + +`dec8` measured that **MXFP4 → Q6_K transcode is exact**, over 1,032,192 superblocks of real K3 expert scale tensors ([[project_k3_kernel_ladder]]). fp4's value set doubled is `±{0,1,2,3,4,6,8,12}` — all integers inside Q6_K's signed 6-bit range — and both structural conditions hold with four exponents of headroom (max observed spread 2 against a limit of 6; `d = 2^(e_min−128)` representable in f16). + +K3's routed experts are MXFP4 at source. Therefore, for K3: + +| arm | gate/up | down | output vs Q0 | +|---|---|---|---| +| Q0 | Q6_K | Q6_K | — | +| Q1 | MXFP4 | Q6_K | **bit-identical** | +| Q2 | Q6_K | MXFP4 | **bit-identical** | +| Q3 | MXFP4 | MXFP4 | **bit-identical** | + +**All four arms produce the same numbers.** The proposed battery — expert-output cosine, weighted latent aggregate, next-layer router agreement, teacher-forced BPB, token KL, greedy agreement, long-generation stability — has *nothing to measure* across Q0–Q3. Running it would produce four columns of zeros and read as a strong result. + +What differs across those arms is **bytes and kernel η only**, and both are already instrumented: `larql k3-ledger ceilings` composes per-class time from banked `MeasuredEfficiency` ranges. + +The arms that *do* have a fidelity axis are Q4/Q5 (an "approximate lower format"). But note what approximate means here: the source alphabet has 15 values and MXFP4 already spends exactly 4 payload bits on it. Going lower is **lossy requantisation of an already-4-bit tensor**, not a container choice — a different experiment with a different falsifier, and one `dec8`'s variable-rate census already priced at a 0.25-bit total prize ([[project_k3_exact_format_floor]]). + +### 1.2 So the case for mixed-format is operational, not numerical — **C2** + +Strip the fidelity axis and one motivation survives, but it is a good one: + +> **Re-quantising one region without rewriting the layer file.** + +At K3 scale that is not a nicety. A single routed layer file holds 896 experts: + +| container | per expert | per layer | × 92 MoE layers | +|---|---:|---:|---:| +| MXFP4 (source density) | 17,547,264 B | **15.72 GB** | **1.45 TB** | +| Q6_K (cheapest exact + servable) | 27,095,040 B | **24.28 GB** | **2.23 TB** | + +*(Cross-check: 1.45 TB routed + 0.112 TB always-resident dense = 1.56 TB, matching the published checkpoint size of 1.561 TB to three figures. The shapes are right.)* + +Two consequences worth stating plainly: + +- **Q6_K everywhere makes the routed vindex 43 % larger than the model it was derived from** — 2.23 TB against 1.561 TB. Same failure *shape* as `k3-funnel.md` Finding 2's 1.82 TB gate blow-up, arrived at honestly rather than by a silent no-op, but it is still a derived artifact exceeding its source. +- **Migrating one region across 92 layers is a 1.45–2.23 TB rewrite today.** With per-region format tags it is a 1.45 TB rewrite of two-thirds of the bytes, or 0.74 TB of one-third — and, more to the point, it can be done *per layer, incrementally, against a live index*. + +That matters because MXFP4's kernel maturity is `Grouped`, below `is_servable()` ([[project_k3_kernel_ladder]]). The moment `mxfp4_grouped_experts` reaches `Dispatched`, you want to move gate/up (two-thirds of expert bytes) to MXFP4 and leave down on Q6_K until its kernel follows. **That migration is the entire product case for LYRW v2**, and it is invisible to any fidelity or single-run throughput measurement. + +### 1.3 The byte split is exactly 2:1, and it is not 1:1 — **C3** + +`dec8-0` read shard 50's header: per routed expert, `w1`/`w2`/`w3` are **5,849,088 B each, equal to the byte**. So: + +``` +gate/up region (w1 + w3) = 11,698,176 B = 2/3 of expert bytes +down region (w2) = 5,849,088 B = 1/3 of expert bytes +``` + +Any format lever on gate/up is worth **twice** the same lever on down. That orders the arms: Q1 (cheap gate/up) captures two-thirds of the available byte saving, Q2 captures one-third. It also caps branch-dropping at **1.5×, never 2–3×** — already banked, and the reason the "drop a branch" lever was found double-counted in the miss budget. + +### 1.4 The repo already ships projection-asymmetric precision — and it does not apply to K3 + +[`format-spec.md` §5.10](../crates/larql-vindex/docs/format-spec.md) and [`fp4-precision-policy.md`](../crates/larql-vindex/docs/fp4-precision-policy.md) default the FP4 feature-row tier to `{gate: fp4, up: fp4, down: fp8}`, on exp-26 cross-model evidence that *down carries FFN's heaviest-tailed per-feature magnitude distribution*. `index.json` already has `fp4.projections.{gate|up|down}.precision` as the authoritative field. + +So the engine already believes the two regions want different precision, has measured evidence for it, ships it in one storage tier, and **LYRW is the one tier that cannot express it.** That is a strong prior for L1. + +**But it is inapplicable to K3.** A heavier tail argues for a *richer* container on down; K3's source is 4-bit everywhere, so there is nothing richer to buy — Q6_K down is exact, MXFP4 down is exact, and no container in between adds information. The exp-26 policy axis is real for models extracted from bf16/f16 sources and dead for K3. + +**Consequence for the run order:** E4's fidelity arm belongs on a **bf16-source MoE** (Gemma 4 26B A4B, Qwen3-30B-A3B, OLMoE), not on K3. E4-on-K3 reduces to E1's byte/η arithmetic. These are two different experiments that the original plan ran as one. + +> **To verify before citing:** exp 26's down-tail finding is recorded as "FFN" cross-model data. Whether the corpus included MoE *routed-expert* `w2` — as opposed to dense FFN down — is not established here, and E4's premise depends on it. Check the exp-26 source before the arm is built; if it is dense-only, the first E4 measurement is the per-projection tail on a routed bank, not the format A/B. + +--- + +## 2. The blocker that outranks all six experiments — **C4** + +**The writer cannot write a K3 routed layer, and no layout choice changes that.** + +`write_layer_weights` (`crates/larql-vindex/src/format/weights/write_layers.rs:68`) takes `entries: &[LayerEntry]` — every expert's quantised bytes, fully materialised. `quantize_moe_entries` (`:209`) takes `gate_up_bf16: &[u8]` and `down_bf16: &[u8]` spanning **all** experts and returns `Vec`. At K3's routed layer: + +| | | +|---|---:| +| input slice, bf16, all 896 experts | **39.4 GB** | +| f32 intermediate, if materialised per the current `bf16_bytes_to_f32` path | **118 GB** | +| output `Vec`, Q6_K | **24.3 GB** | + +against 128 GB of RAM. The peak is not survivable, and it is a property of the API signature, not of the machine. + +**It is also barely wired.** The only caller of either function outside tests is `crates/larql-cli/examples/convert_moe_to_per_layer.rs` — an *example* used to migrate an existing Gemma vindex. The main extraction path does not write per-layer MoE weights at all. So "can the K3 extraction use the current physical layout" currently has the answer **"the current physical layout has no K3-capable writer, in any variant."** + +**The fix is layout-independent and should start now:** stream the write — emit the header, reserve the offset table, quantise and append one expert at a time, backpatch the table on close. That shape is identical under L0, L1 and L2, and it is a precondition for every arm in E1 and E2. Peak RAM becomes one expert (27 MB) plus the table (28 KB at 896 entries). + +This is the long pole. Everything below assumes it lands first. + +--- + +## 3. The second blocker: `format_version` is written and never checked — **C5** + +```rust +// write_layers.rs:249 +// format_version at [4..8] — currently ignored, forward-compatible +``` + +It is not forward-compatible; it is silently wrong-compatible. `parse_layer_weights_header` reads the magic, skips the version, then parses the offset table as `num_entries × 32` bytes unconditionally. A v2 file whose entry stride is anything other than 32 bytes will be parsed by today's reader into **garbage offsets that are still inside the file**, so it will not bounds-fail — it will hand `get_layer_entry_bytes` a plausible byte range from the wrong place, and the model will produce plausible wrong numbers. + +Same class as the `XSTRIDE` and `BufferCache::get_bytes` bugs from the kernel ladder: *wrong stride yields a plausible number from the wrong expert, not a crash.* + +**Fix before any v2 file exists anywhere**, including in a fixture: reject `format_version > FORMAT_VERSION` in the parser. There is exactly one production caller (`format/weights/load/q4k.rs:186`), which already treats `None` as "skip this layer", so the failure mode is a clean miss rather than a panic. That is a two-line change and a test, and it is cheap only *until* the first v2 file exists. + +--- + +## 4. Test assets + +| Asset | Scale | Answers | Status | +|---|---:|---|---| +| **Mini-K3 fixture** | ~1/168 expert area | format, loader, dispatch, manifest correctness | to build | +| **Byte-faithful K3 routed bank** | 1 layer | SSD, mmap, grouping, cache | to build; **1 layer, not 4** — see §4.2 | +| **Loaded bf16-source MoE** | Gemma 4 26B A4B | real scheduling, E3 A/B, E4's live fidelity axis | checkpoint present, **no vindex on this box** | +| **K3 trace pack** | selected real layers | fidelity, locality, omission | **blocked** — see §8 | + +### 4.1 Mini-K3 fixture — dimensions corrected — **C9** + +The proposed 1/16 downscale (hidden 448, latent 224, inter 192) **breaks Q4_K/Q6_K superblock alignment on every matrix.** `pad_cols_to_256` would pad `in_cols` 224 → 256 on gate/up (+14 %) and 192 → 256 on down (+33 %), so the fixture's byte ledger would be a padding artefact rather than a scaled K3. + +And no proportional downscale fixes it. Alignment needs `latent/d ≡ 0 (mod 256)` and `inter/d ≡ 0 (mod 256)`; with `3584 = 2⁹·7` and `3072 = 2¹⁰·3` that admits only `d ∈ {1, 2}`, and `d = 2` is 6.07 GB per layer — not a fixture. + +**Use 256-aligned dimensions that preserve the ratios that matter instead:** + +``` +hidden 512 (2:1 against latent — exactly K3's 7168:3584) +latent 256 +expert intermediate 256 (K3 is 3584:3072 = 1:0.857; this is 1:1 — the one deviation) +experts 896 (unscaled — routing shape is the point) +selected 16 (unscaled) +layers 4 (8 if iteration cost allows) +``` + +| | mini | K3 | ratio | +|---|---:|---:|---:| +| per-expert Q6_K bytes | 161,280 | 27,095,040 | 1/168 | +| per-layer, 896 experts | **144.5 MB** | 24.28 GB | 1/168 | +| 4 layers | **578 MB** | — | — | +| gate/up : down | **2 : 1** | 2 : 1 | ✓ preserved exactly | +| padding | **zero** | zero | ✓ preserved | + +The `w1 = w2 = w3` equal-bytes property — the thing that caps branch-dropping at 1.5× — is preserved exactly, which the 1/16 proposal would have destroyed via asymmetric padding. + +It needs no training. Its purpose is deterministic numerical parity: same architecture graph, same top-16-of-896 routing, same shared latent projections, same expert entry structure, small enough to rewrite repeatedly. + +**Standing caveat: never quote a byte or latency number from the mini fixture as a K3 number.** It is 1/168 of the area with a different latent:inter ratio and a working set that fits in cache. It answers correctness; it does not answer SSD. + +### 4.2 Byte-faithful bank — the right device is attached — **C7, revised 2026-08-01** + +Opaque or deterministic payloads at the real K3 expert byte sizes. + +**Both constraints recorded here on 2026-08-01 are now dead.** `model-drive` is mounted at `/Volumes/model-drive`: a Ugreen NVMe enclosure on Thunderbolt/USB4 Bus 0, protocol PCI-Express, link negotiated at **40 Gb/s**, 2.0 TB APFS, **1.97 TB free**. So: + +- **Disk is not a constraint.** Four Q6_K layers is 97 GB against 1.97 TB free — 5 % of the volume. Arm variants, MXFP4-density copies and permuted orderings all fit alongside. The earlier "start at one layer, 175 GB internal" scoping was a property of the machine that day, not of the experiment. Start at one layer if it answers the question, not because of space. +- **This is the external SSD, over the intended link.** `dec8-2` (USB4 random-read, no model) characterises this same device class, so an E2 replay run on `model-drive` **may** be labelled as the external-SSD result — the earlier "different device, must not be labelled as such" caveat applied only to an internal-NVMe fallback that is no longer necessary. + +Two methodological notes that replace them, both about the *measurement* rather than the device: + +- **The drive is nearly empty (23 GiB of 2 TB used).** SSD steady-state behaviour on a fresh volume is not the same as on a full one — garbage collection and write amplification both change. E2 measures reads, so the exposure is small, but a sustained-write phase (building four layers plus variants) should not be read as a *read* result. +- **Cold means cold.** A 24.3 GB layer against 128 GB of RAM is fully page-cached after one touch, so the 2/4/8/16 GB memory budgets have to be enforced against the unified buffer cache, not just declared. Without that, every arm after the first measures RAM. + +Traces to replay (unchanged from the proposal, and the set is good): uniform top-16, Zipf, Gemma empirical, clustered co-activation, adversarial rotation, repeated conversation. + +Measure: SSD bytes/token, useful vs amplified, page faults, read count and extent sizes, p50/p95/p99 layer latency, degradation onset, steady-state floor, next-layer prefetch effect, physical permutation effect. + +**Test expert ordering independently of file grouping.** The offset table means physical order need not equal logical order, so `logical expert ID → physical offset` is a permutation the format already supports at zero cost. Compare ID order, frequency order, co-activation-cluster order, random order — *within* L0, before considering L3. + +> **Run this under a within-run control.** `dec8-12` found the machine degrades after ~7 sustained censuses, with a control variable falling 0.89 → 0.06 while every cell moved together — blind averaging would have banked a value wrong by a factor *while presenting as more data*. A sustained SSD replay is exactly that regime. The control must have a working set **at least as large as the thing it protects** (a small sentinel admitted two runs where the large class had collapsed). And gate on `n ≥ 5 && relative SE ≤ 1 %`, never on observed spread — `max − min` is non-decreasing in `n`, so a spread gate rewards collecting less data. + +--- + +## 5. Experiment 1 — the layout gate + +### Candidates + +| | | | +|---|---|---| +| **L0** | current format | one `quant_format` in the header; per-expert `(gate_up, down)` offsets | +| **L1** | mixed-format LYRW v2 | same file, same offset table, each region declares its own format | +| **L2** | projection-separated files | `layer_N.gate_up.weights` + `layer_N.down.weights` | +| **L3** | expert-group files | `layer_N.experts_000_031.weights`, … | + +**Prior: L1.** It provides role-specific representation without doubling file count or weakening per-expert addressability, and §1.2 gives it a concrete operational payoff that no other arm has. L2 and L3 must beat it on measurement, not on tidiness. + +**L1's wire shape**, stated so E1 has something concrete to build: + +``` +header: magic, format_version = 2, default_format, num_entries, intermediate, hidden +entry: gate_up_offset u64, gate_up_bytes u64, gate_up_format u32, _pad u32 + down_offset u64, down_bytes u64, down_format u32, _pad u32 +``` + +48 bytes per entry against v1's 32. At 896 experts the table grows 28.7 KB → 43.0 KB per layer — noise against a 24 GB file, and still one page-aligned read at startup. `default_format` keeps the header self-describing for tools that only want the layer's dominant representation; the per-region tags are authoritative. **This is precisely the stride change §3's missing version check would silently misparse.** + +### Method + +Run on the mini fixture, all arms carrying identical Q6_K bytes so that only the container varies. + +Measure: exact output parity; loader complexity (LOC and branch count in the read path); mmap count; map/open latency; dispatch construction time; bytes faulted; sequential and random top-16 execution; and **whether a representation can be overridden or omitted through the manifest without rewriting the file** — the operational property from §1.2, which is the one L1 exists for. + +### Decision rule (pre-registered) + +Choose the **least fragmented layout** within: + +- **2 %** of best warm execution +- **5 %** of best cold physical bytes read +- **exact** numerical parity — no tolerance; the arms carry identical bytes, so any difference is a bug + +Do not promote L2 or L3 on conceptual tidiness. Do not promote L1 on a fidelity result — §1.1 says there isn't one; promote it on the migration property or not at all. + +**Falsifier for L1:** if per-region tags cannot be threaded to the kernel without changing a `QuantMatVec` dispatch signature, or if the mixed path forces a per-expert branch inside the grouped-expert kernel (which would put a divergent branch in every threadgroup — the failure that killed zero-carried scales), then the format is buying an operation the compute path cannot execute, and L0 stands. + +--- + +## 6. Experiments 2 and 3 + +**E2 — byte-faithful replay.** Assets and controls per §4.2. + +*Decision rule:* do not bake a clustered physical ordering into the permanent index unless it delivers **≥ 10 % sustained improvement across multiple plausible traces** and does not regress uniform routing badly. Prior: per-expert offsets plus a runtime cache beat permanently encoding a locality assumption — and note that K3's quantile-balanced 16-of-896 router is *designed* to avoid the expert imbalance that would make locality pay, so Gemma's strong layer-specific collapse cannot be assumed to transfer ([[project_moe_routing_locality]], and `dec8-5`'s standing rule **R2**: never transfer a union ratio across activation fractions — Gemma's 6.25 % activation against K3's 1.79 %). + +**E3 — loaded-model grouped A/B.** Arms A (production layout) / B (LYRW + grouped) / C (LYRW v2 + grouped) / D (split-file + grouped), holding representation, selected experts, routing weights, activation layout, command-buffer ownership and output reduction constant. + +Two constraints from the bank: + +- **No Gemma vindex exists on this box** — only the raw checkpoint. E3 needs an extract run first, or a different loaded MoE. +- **A Gemma A/B is a PROXY MECHANISM result.** Never attach a K3 throughput number to it. And per [[feedback_isolated_vs_batched_kernel_profile]], measure batched (`diag_profile_kernels`) — isolated single-dispatch benchmarking understated by 3.0–6.7× on this exact class of question, and `dec8-9` produced a phantom 5.10× that was 1.23× once both arms were batched identically. + +--- + +## 7. Experiments 4 and 5 + +**E4 — projection-specific quantisation.** Re-scoped per §1.1 and §1.4 into two disjoint experiments: + +- **E4a (K3):** no fidelity axis. Reduces to bytes × η per region, composed by `larql k3-ledger ceilings`. Largely already answered; run it as arithmetic, not as a sweep. +- **E4b (bf16-source MoE):** the live one. Does the exp-26 down-tail policy that ships for the FP4 feature tier reproduce on a *routed expert's* `w2`, and does it pay in a container? Arms Q4/Q5 (asymmetric approximate) against Q0. Subject to the §1.4 verification note. + +*Decision rule:* add mixed-format complexity only when a mixed arm beats **both** uniform alternatives by ≥ 5 % projected whole-system gain at the same fidelity gate — **or** when the migration property of §1.2 is independently judged worth it. Those are two separate justifications and should be recorded separately, because only the second one applies to K3. + +**E5 — weight omission.** The proposal's central correction is right and worth restating, because it is the kind of thing that gets relitigated: **dropping an expert's `w2` does not produce a cheaper partial expert, it produces no expert output.** Keeping gate/up while dropping down is strictly dominated by skipping the expert entirely, which also saves the gate/up read. So the useful modes are: + +| | | +|---|---| +| **O0** | exact baseline, all 16 | +| **O1** | reduced routed top-K (16/12/8/4), with and without gate renormalisation | +| **O2** | cumulative gate-mass retention (99/97.5/95/90 %) | +| **O3** | shared-only layer — skip the routed branch, keep shared experts, residual, attention | +| **O4** | entire routed branch remote | +| **O5** | down-only lower precision — all experts execute, cheaper `w2` | + +**O4's payload arithmetic is correct and decisive.** Whole-branch RPC moves one latent in and one aggregated latent out: `2 × 3584 × 2 B = 14,336 B`. Projection-split moves 16 intermediate vectors: `16 × 3072 × 2 B = 98,304 B` — **6.9× more**, before protocol overhead. Do not split local gate/up from remote down. + +*Gate:* ≤ 0.5 % BPB regression for a production approximation; no pathological p99 KL; no collapse in next-layer routing agreement; no long-generation instability; gain recomposed against the whole-model byte ledger. + +*Two priors that constrain O1/O2/O5, and should be read before designing them:* [[project_walkffn_speed_accuracy_scissors]] (20 %/layer all-layer → KL 8.4 bits, 0 % top-1) and [[project_r4_zeroout_sparse_ffn]] (kernel captured only 23–40 % of a row reduction even with routing free). And per [[feedback_stacked_zero_ablation]], **mean-ablate, never zero-ablate**, across stacked layers. O1/O2 also overlap the shipped C6 drift gate (`larql dec-bench drift`) — reuse it rather than building a second instrument. + +--- + +## 8. Experiment 6 — the K3 routing trace is blocked — **C8** + +This is the only experiment that cannot be done honestly with Gemma or synthetic routing, and it is also the one that cannot be scheduled. + +Capturing per-token expert IDs requires applying `block_sparse_moe.gate.weight` to a real hidden state, which requires a K3 forward pass. There isn't one: R3/P2 is unbuilt, and `A_log`'s axis is a **fail-closed BLOCKER** that header shapes provably cannot settle ([[project_k3_exact_format_floor]]) — it needs an oracle, against an asymmetric fixture, before any KDA fidelity work. + +**So E6 is downstream of the K3 adapter, full stop.** Its dependency is the ladder's long pole, not this programme's. + +What follows for the other five: **nothing here may freeze a physical expert permutation, a permanent hot set, co-activation grouping, or a per-layer retention policy.** All four are locality claims, K3 locality is entirely unearned ([[project_k3_kernel_ladder]]: "K3a created none, the trace still owns reuse"), and E2's synthetic traces can characterise the *mechanism* without licensing a *policy*. Keep locality as runtime metadata over the offset table — which is also where the L1 prior lands. + +--- + +## 9. Run order + +### Now — blockers and the layout gate + +| | | Depends on | +|---|---|---| +| 0a | **`format_version` rejection in the parser** (§3) | — | +| 0b | **Streaming layer writer** (§2) | — | +| 1 | Mini-K3 fixture at 512/256/256 × 896 (§4.1) | 0b | +| 2 | E1 — L0 / L1 / L2 (L3 only if L2 shows something) | 0a, 1 | +| 3 | E2 — byte-faithful layer(s) on `model-drive` (USB4, 40 Gb/s, 1.97 TB free) | 0b | +| 4 | E3 — loaded-MoE grouped A/B (needs a vindex extract first) | 2 | +| 5 | E4b — projection-specific on a bf16-source MoE | 4 | + +### Before the K3 extraction + +6. K3 adapter far enough for a forward pass (**owner: `k3-funnel` R3/P2**, `A_log` oracle first) +7. E6 locality/cache simulation on a real trace +8. E5 top-K / mass-retention / shared-only, via the C6 drift gate +9. Freeze the K3 vindex ABI + +### Then extract once + +``` +attention / KDA / MLA +router and MoE control +shared latent projections (routed_expert_down_proj 7168→3584, +shared experts routed_expert_up_proj 3584→7168) + +routed/ + layer_N.weights + entry e: + gate_up offset, bytes, format + down offset, bytes, format +``` + +Not: thousands of tiny expert files; not fixed hot/cold physical slices; not one permanently selected retention policy. + +### Naming discipline + +K3 has two things called "down" and two called "up" — the layer-level `routed_expert_down_proj` / `routed_expert_up_proj` (shared, resident, 7168↔3584) and the per-expert `w2` / `w3`. And two called "gate": `block_sparse_moe.gate.weight` (the router) and `w1` (the GLU gate branch). Every claim in this programme must say which ([[project_k3_ssd_miss_budget]]). In this document, **"down" unqualified means the per-expert `w2` region of a LYRW entry**, and the shared latent projections are always named in full. + +--- + +## 10. Standing hypothesis + +The programme will probably show: + +- one file per routed layer remains correct; +- one entry per expert remains correct; +- gate and up stay fused; +- down stays independently addressable; +- **per-region format tags are the only format extension worth making** — and they are worth it for incremental migration under a maturing kernel set, not for fidelity or single-run throughput; +- expert locality stays runtime metadata; +- omission during local inference means "skip the routed branch", not a half-expert mode; +- remote deployment moves the whole routed computation, not individual experts across the network. + +**And a caution on that list.** It is an organising frame that explains everything so far, which is exactly the kind of claim that needs its own falsification test rather than accumulating confirmations ([[feedback_organizing_vs_empirical_claims]]). Its sharpest falsifier is E1's L1 falsifier in §5: if per-region tags cannot reach the grouped-expert kernel without a divergent per-expert branch, the frame's central recommendation is unbuildable and L0 stands unchanged. diff --git a/docs/mech-interp.md b/docs/mech-interp.md index f8ea89aa0..ad4c07953 100644 --- a/docs/mech-interp.md +++ b/docs/mech-interp.md @@ -191,7 +191,7 @@ cargo run --release -p larql-inference --example mech_interp_demo ``` Walks through all seven primitives on synthetic weights (no vindex -required). Source: `crates/larql-inference/examples/mech_interp_demo.rs`. +required). Source: `crates/larql-demos/examples/inference/mech_interp_demo.rs`. --- diff --git a/docs/vindex3-experiments.md b/docs/vindex3-experiments.md new file mode 100644 index 000000000..32f32a4e4 --- /dev/null +++ b/docs/vindex3-experiments.md @@ -0,0 +1,520 @@ +# VINDEX3 Experimental Programme + +**Version:** 1.0-draft-1 +**Date:** 2026-08-01 +**Status:** Pre-registered. Decision rules and promotion gates are fixed here before any arm runs. +**Companion:** [`vindex3-format-spec.md`](../crates/larql-vindex/docs/vindex3-format-spec.md) +**Registry programme:** `vindex2` (chuk-experiments) — kept at `vindex2` deliberately: arms are already recorded under that key, and re-keying a pre-registered programme would orphan its results. The gate ids `V2-x` below are external identifiers for the same reason. The *format* is VINDEX3; see the note at the top of the companion spec. +**Discipline:** DEC-style — every arm has a registered prior, a numeric gate, and a named negative outcome. A result that fails its gate closes its thread; it does not get re-argued. + +--- + +## 0. Structure + +Five gates (V2-0 … V2-4) sequence the work. Nine experiments (E0 … E8) supply the evidence. Four test assets carry the experiments — no single proxy answers everything: + +| Asset | Scale | Answers | +| ----- | ----- | ------- | +| **Conformance fixtures A–D** | tiny, deterministic | format, loader, programme and manifest correctness | +| **Byte-faithful K3 bank** | 1–4 full-sized routed layers | SSD, mmap, segmentation, grouping, cache behaviour | +| **Loaded Gemma MoE** | existing real model | real buffers, real scheduling, end-to-end integration | +| **Kimi-Linear-48B-A3B-Instruct** | real model, laptop-runnable end-to-end | real shared-expert + sigmoid-router serving, hybrid dense/MoE manifests, KDA/MLA spine adapter de-risking, real same-lineage routing traces | +| **Inkling-Small (276B-A12B)** | real model, exceeds rig RAM | real 2-shared sink-router reduction, real NVFP4/MXFP8 regions + mixed-precision release convention, first forced partial-residency/remote serving | +| **K3 trace pack** | selected real layers | fidelity, routing locality, omission decisions | + +Nothing about K3 locality is decided from Gemma. K3's quantile-balanced top-16-of-896 routing is designed to suppress exactly the per-layer collapse Gemma exhibits; Gemma proves methodology only (consistent with Milestone D in ROADMAP_STATUS). + +--- + +## 1. Test assets + +### 1.1 Conformance fixtures (functional, not trained) + +Deterministic weights; purpose is numerical parity against a per-architecture oracle, not speed. + +**A — Direct routed MoE (control)** +``` +hidden 256 · experts 8 · top_k 2 · shared 0 · programme gated-mlp-v1 +``` + +**B — GPT-OSS-shaped** +``` +hidden 288 · intermediate 288 · experts 32 · top_k 4 · shared 0 +programme gpt-oss-expert-v1 (clamp + residual term) +representation: MXFP4-like blocks + separate scale regions +``` + +**C — Inkling-shaped** +``` +hidden 256 · experts 32 · top_k 6 · shared 2 +programme routed + always-active shared bank +representations: BF16, FP8-like, FP4-like +``` +Catches: shared-expert reduction order, router-normalisation participation, mixed routed/shared storage. + +**D — K3-shaped (Mini-K3)** [dimensions frozen 2026-08-01] +``` +hidden 512 · latent 256 · expert intermediate 256 +experts 112 · top_k 16 · shared 2 · programme latent-moe-v1 +``` +Same architecture graph, same routing shape, same shared latent projections; small enough to regenerate and rewrite repeatedly. + +**Why not the literal 1/16 downscale.** The draft-1 dimensions (448 / 224 / 192) break Q4_K/Q6_K superblock alignment on every matrix: `pad_cols_to_256` pads 224 → 256 (+14 %) and 192 → 256 (+33 %), so the fixture's byte ledger would be a padding artefact rather than a scaled K3. No proportional downscale fixes it — alignment needs `latent/d ≡ 0 (mod 256)` and `inter/d ≡ 0 (mod 256)`, and with `3584 = 2⁹·7`, `3072 = 2¹⁰·3` that admits only `d ∈ {1, 2}`. 512 / 256 / 256 is 256-aligned everywhere, preserves K3's 2:1 hidden:latent exactly, and pads nowhere. The one deviation is intermediate:latent, which is 1:1 here against K3's 1:0.857. + +**Why 112 experts, not 56.** 56 preserves the literal 1/16 expert-count scaling but is not a whole multiple of the 16-expert group width, so it would bake a partial group into the fixture — and group width dividing segment width is a physical-design *rule* (§7), not a preference. 112 = 7 × 16 divides cleanly, stays tiny, and leaves exact-K3-scale fidelity where it belongs: the byte-faithful bank, which runs 896 unscaled. + +**Standing caveat:** never quote a byte or latency number from this fixture as a K3 number. It answers correctness, not SSD. + +### 1.2 Byte-faithful K3 routed bank + +Opaque or deterministic quantised payloads at **real K3 expert byte sizes**: 896 experts, gate/up and down at real Q6_K lengths (per-expert 3×3584×3072 params; ~24.28 GB / 22.61 GiB per full layer — deliberately over the 20 GiB cap so segmentation is exercised, not simulated). One layer is a meaningful SSD experiment; four layers expose cache exhaustion and sustained behaviour. + +### 1.3 Kimi-Linear-48B-A3B-Instruct (the K3 dress rehearsal) + +Verified from the released `config.json` (`moonshotai/Kimi-Linear-48B-A3B-Instruct`, 98.3 GB BF16, ~3B active): + +``` +layers 27 · first_k_dense_replace 1 (layer 0 dense FFN @ intermediate 9216; + layers 1–26 MoE) +hidden 2304 +routed experts 256 · top-8 · shared experts 1 · moe_intermediate 1024 +router: sigmoid scores · renormalised · routed_scaling_factor 2.446 + · grouped top-k (1 group — degenerate but present in the schema) +dense spine: 20 KDA layers + 7 MLA layers, 3:1 interleave + (full_attn_layers = 4,8,12,16,20,24,27; 1-based) +MLA: kv_lora_rank 512 · qk_nope 128 · qk_rope 64 · v_head 128 · no q-LoRA +KDA: 32 heads × head_dim 128 · short_conv 4 +``` + +Routed-bank arithmetic (Q6_K, 210 B / 256 params): + +``` +params per expert = 3 × 2304 × 1024 = 7,077,888 +bytes per expert ≈ 5.54 MiB +per MoE layer (256) ≈ 1.38 GiB → single segment, well under the cap +26 MoE layers ≈ 36 GiB routed — whole model Q6_K ≈ 40 GB class + → runs end-to-end on the M3 Max +``` + +What it uniquely contributes: the only **real, decodable** shared-expert model in the set (fixture C stays synthetic for cheap conformance); the only real hybrid dense+MoE stack (per-layer manifest heterogeneity for free); a non-softmax router (sigmoid + renormalise + scaling factor) that stress-tests the manifest router vocabulary; and — most valuable — the **same KDA:MLA hybrid spine lineage as K3** at 1/30th checkpoint scale, so the class-1/class-2 adapter work (recurrence parameters, conv states, MLA latent KV, layer scheduling off `full_attn_layers`) is debugged here before the K3 adapter, the roadmap's main blocker. It also proves the single-segment routed path with real weights while the byte-faithful bank proves the multi-segment path — the two segmentation regimes each get a real test. + +**Boundary:** KL-48B enters the *design set*. It is therefore ineligible as E8's held-out model, its sigmoid routing does not transfer locality conclusions to K3's quantile-balanced routing, and its BF16-only release means native low-bit region handling leans on GPT-OSS (MXFP4) and Inkling-Small (NVFP4/MXFP8, §1.4). + +### 1.4 Inkling-Small (the K3 serving dress rehearsal) + +Verified from the released `config.json` (`thinkingmachines/Inkling-Small`, 532 GB BF16 + 4.46 GB `mtp.safetensors`; 276B total / 12B active per the release notes; Apache-2.0): + +``` +layers 42 · hidden 4096 · dense MLP at layer index 2 (dense_intermediate 16384) +routed experts 256 · top-6 · shared experts 2 · moe intermediate 2048 +router: sigmoid gate activation · gate bias · norm_after_topk · route_scale 8.0 + · global scale · shared_expert_sink = true (shared experts scored + by the router — always active, inside normalisation) +attention: 5:1 local(SWA-512):global · GQA 32/8 · head_dim 128 + · SConv kernel 4 on residual branches · d_rel/rel_extent relative terms +aux: 8 chained MTP heads (separate safetensors) · vision hMLP patchifier + · dmel audio encoder — auxiliary manifest-addressed tensors, optional +releases: BF16 + NVFP4 + MXFP8; quantised convention = routed experts low-bit, + shared experts / attention / gates BF16 (a real mixed-precision index) +``` + +Routed-bank arithmetic (Q6_K, 210 B / 256 params): + +``` +params per expert = 3 × 4096 × 2048 = 25,165,824 +bytes per expert ≈ 19.69 MiB +per MoE layer (256) ≈ 4.92 GiB → still single segment, under the cap +~41 MoE layers ≈ 202 GiB routed — CANNOT be RAM-resident on the M3 Max +routed reads/token ≈ 6 × 19.69 MiB × 41 ≈ 4.7 GiB at top-6 +``` + +What it uniquely contributes: the real two-shared-expert reduction under a **sink router** — shared experts inside the scoring and normalisation, exactly the semantics fixture C could only fake, and the richest test of the manifest's router vocabulary (sigmoid + bias + post-top-k norm + route scale + global scale). Real **NVFP4/MXFP8 native regions** from the quantised releases, including the mixed-precision convention (routed low-bit, everything else BF16) that is itself a per-region-format and variants-model use case — native low-bit no longer leans on GPT-OSS alone. A **mid-stack dense layer** (`dense_mlp_idx: 2`) — per-layer manifests handle it for free where any `first_k_dense_replace` field would have failed. And the residency escalation: it is the first design-set model where partial-residency, SSD streaming and attn-local/FFN-remote profiles are **forced rather than optional** on the rig — the K3 *serving* dress rehearsal at one-eighth scale, complementing KL-48B as the K3 *adapter* dress rehearsal. MTP heads are stored as optional auxiliary tensors whose omission never changes authority (drafting only); multimodal towers are opaque auxiliary payload — the text backbone is the conformance target. + +**Boundary:** Inkling-Small enters the design set — ineligible as E8's held-out model. Its sink routing is a third distinct balancing mechanism (vs KL-48B's grouped sigmoid and K3's quantile-balancing): more evidence that locality conclusions do not transfer between routers, and none of its traces inform K3 locality. + +### 1.5 K3 trace pack (E6 prerequisite) + +Real K3 routing/activation captures across prose, code, agent/tool use, long-context retrieval, multilingual, repeated multi-turn. Per layer/token: expert IDs, gate scores, selection order, co-activation sets, reuse distance, latent input norm, expert output norms, aggregate routed delta. + +--- + +## 2. Gates + +### Result log — fixture A container (2026-08-02) + +The first VINDEX3 container to exist on disk. Recorded here rather than in the +spec because it is evidence, not specification. + +```text +write_container index.json schema 3 + moe_manifest.json + LYRW v2 bank +detect_generation V3, from the written directory (not a JSON literal) +Vindex3Container::open manifest parsed and validated, storage keys resolved +region round-trip every gate/up/down region byte-for-byte +bind + execute BIT-IDENTICAL to the same weights held in memory +fused vs decomposed BIT-IDENTICAL under one programme id (gated-mlp-v1) +verify structural, defects carry {layer, entry, role} +CLI show/verify dispatch on generation; v3 defect → exit 1 +``` + +Rows now discharged: + +| Gate | Row | Status | +| ---- | --- | ------ | +| V2-0 | indexes inspectable without loading weights | closed | +| V2-0 | unknown `programme_id` fails cleanly | closed (`MoeManifest::parse`) | +| V2-0 | missing regions diagnosed with coordinates | closed — `ContainerDefect::MissingRegion` carries `{layer, bank, role, segment}` plus entry | +| V2-0 | generation boundary enforced both ways | closed | +| V2-0 | a profile dropping a component cannot exceed derived authority | closed — by *pre-existing* tests in `capability/authority.rs` (`weakening_any_region_never_raises_authority` and four siblings), not by the container work | +| V2-0 | variant-selection refusal (select only present variants) | closed — `index.variants` catalogues each region set's present variants and baseline; a `Profile` selects per region set and `Vindex3Index::select_profile` refuses an absent one naming the region set, the request and what is present. Enforced in `Vindex3Container::open` for **every** declared profile, before a segment byte is read — pinned by deleting the segment files and asserting the error still names the variant. `declares_profile` remains, documented as a name check only. **Ceiling:** the refusal is real; *steering* is not exercised end to end, because no writer emits a multi-variant container yet (`ContainerSpec` has no variant field) and `BankRef.storage` still names storage directly, so a selection does not yet change which bytes the runtime binds. That wiring belongs with the first real pack. | +| V2-1 | native oracle vs container output exact | closed (fixture A) | +| V2-1 | fused vs decomposed identical under one manifest | closed | +| V2-1 | expert counts and top-K not hard-coded | closed — (8,2), (32,4), (5,1) through one code path, each read from its own container | +| V2-1 | shared banks not hard-coded | **open**, and *blocked*: `BoundMoeOperation::banks` documents unrouted shared banks as arriving with the Mini-K3 rung, so there is no execution path to test against | +| V2-1 | WALK/DESCRIBE parity | **open** | + +Two corrections worth carrying forward, both from getting it wrong first: + +- **`gated-mlp-v1` admits both layouts** (`role_alternatives → [FUSED, + DECOMPOSED]`). It is not "the decomposed programme"; `gated-mlp-fused-fc1-v1` + is the one that narrows to fused only. So the parity arm is *one programme + admitting two renderings*, which is the stronger property and the one K3 + needs. +- **Storage fidelity and execution fidelity need separate assertions.** A + single execution test can pass while the writer stored the wrong bytes, if + the reader makes a compensating mistake. The region round-trip test exists to + make that impossible. + +### V2-0 — Format skeleton + +Build superblock, LYRW v2 header/bank/segment/region tables, manifest schema, profile inheritance, authority levels, checksums, required/optional role semantics. Tiny tensors only. + +**Acceptance** +- indexes inspectable without loading weights; +- unknown `programme_id` fails cleanly; +- missing required regions diagnosed with `{layer, bank, role, segment}` precision; +- a profile that drops a required component cannot exceed its derived authority (§9.2) — verified by test, not review; +- **profile resolution**: a profile can select only variants physically present; selecting an absent variant fails naming region set, requested variant and present variants; no code path performs silent conversion; region-level fidelity aggregates to profile authority as weakest-link; a profile referencing incompatible segment sets is refused; +- generation boundary enforced: the v2 loader refuses a VINDEX2 directory (and vice versa) with a precise "requires VINDEX{n} loader" error naming both versions — never a parse error, never silent conversion (spec §6.6/§12.1). + +### V2-1 — Generic reference executor (runs E1 arms functionally) + +All four fixtures execute through the generic path and match their oracles. + +**Acceptance** +- native oracle vs vindex2 output: exact (integer/quantised paths) or tolerance-bounded parity; +- fused vs decomposed FC1 storage produce identical results under one manifest; +- expert counts, top-K and shared banks are demonstrably not hard-coded (fixtures differ on all three); +- **WALK/DESCRIBE parity**: gate KNN over in-place bank regions returns identical top-K rankings to a v1-style extracted `gate_vectors.bin` control on fixtures A–C, and latent-space WALK on fixture D returns the correct ranking after the `routed_input` projection (spec §15.4). + +### V2-2 — Physical layout (runs E1 layout arms + E2) + +### V2-3 — Production kernel binding (runs E3) + +Bind existing grouped kernels through the capability registry, Gemma first (then GPT-OSS as first non-Gemma real model). + +**Acceptance** +- reference and grouped paths agree per-layer and at final logits; +- caller-owned buffers/command buffers work; no hidden decode-time repacking; +- mixed per-region formats supported or explicitly refused, never coerced. + +### V2-4 — Real-model portability + +Import and execute in order: Gemma MoE → GPT-OSS → **Kimi-Linear-48B-A3B** → **Inkling-Small** → K3. Rung 3 (KL-48B): full import, exact decode against the reference implementation (token parity), then grouped dispatch with the shared expert in the execution path and sigmoid-renormalised-scaled gate weights flowing through the standard reduction. Rung 4 (Inkling-Small): BF16 import + token parity on the text backbone, then the NVFP4 release imported as paired values/scales regions under the mixed-precision convention, then the first forced partial-residency / attn-local-FFN-remote serving profiles — rung 4 acceptance explicitly includes decoding a model that does not fit in RAM. Fixture C is retained as the tiny deterministic conformance fixture only. Only after V2-0..V2-3 pass does the one-time K3 extraction run. + +--- + +## 3. Experiments + +### E0 — VINDEX2 preservation matrix (continuous regression, not a one-shot) + +VINDEX3 development must not degrade the shipped generation. E0 runs using the same binary that carries the v2 code, from the first V2-0 commit onward, in CI. + +#### E0 corpus — constructed and pinned [amended 2026-08-01] + +Draft-1 said "existing production v1 indexes". On the actual rig that named **nothing**: no vindex of any generation exists on the box, so the experiment's premise was empty and E0 would have passed vacuously. The corpus is therefore constructed and pinned, not assumed: + +| | Corpus | Purpose | +| - | ------ | ------- | +| **C1** | A fresh Gemma v1 extract at **`--level all`**, plus all eight slice presets cut from it (`client`, `attn`, `embed`, `server`, `browse`, `router`, `expert-server`, `all`) | The primary subject. Covers every v1 path the matrix exercises. `--level all` is required, not preferred: the `all` preset needs `lm_head` and the `router` preset needs router weights, neither of which `--level inference` extracts — a lower level would silently reduce preset coverage to those the extract happened to reach. | +| **C2** | Published v1 artifacts pulled from the hub | Exercises the `publish`/`pull` path and the generation stamp against artifacts this box did not build. | +| **C3** | **Golden outputs captured once from the pre-v2 binary and committed** | The comparison baseline. | + +**C3 is load-bearing.** Without committed goldens, "zero behavioural regression" quietly degrades into "the two binaries agree with each other" — a condition any *shared* bug satisfies. The goldens must be captured from a binary that predates the v2 code, so the assertion is against a fixed record rather than a live second run. + +**Pin recipes, not artifacts.** C1 is a multi-GB directory and C2 is a download; neither belongs in the repo, and neither needs to — both are *reproducible*. What gets committed is the recipe and the outputs: + +| Item | Committed? | Why | +| ---- | ---------- | --- | +| C1 vindex | No — **recipe pinned** | Re-extractable from the checkpoint at any time. Pin: model repo + revision hash + exact extract flags (`--level all --quant q4k`). | +| C2 hub artifacts | No — **coordinates pinned** | Re-pullable. Pin: artifact ref + expected checksums. | +| C3 goldens | **Yes** | Small text. The only thing that cannot be regenerated *later* without also regenerating the binary that produced it. | +| C3 baseline binary | No — **commit SHA pinned** | `git checkout && cargo build --release -p larql-cli` reproduces it. Git keeps the SHA reachable indefinitely. | + +**There is no expiry.** An earlier draft of this section claimed C1/C3 had to be captured before VINDEX3 merged, on the grounds that the pre-v2 baseline stopped being available. That was wrong: the checkpoint still extracts and the baseline commit is still checkoutable. The correct reason to build the corpus early is that E0 is specified as *continuous* — it should be catching regressions as later V2-0 work (manifest, profiles, capability checking) touches shared code, not auditing at the end. Early because the gate is meant to be live, not because anything is running out. + +**What must be recorded, or the corpus is unfalsifiable:** the baseline commit SHA inside the golden set itself. Without it a reader cannot tell whether a given golden predates the v2 work it is supposed to police. + +**Matrix** + +``` +load a C1/C2 v1 index → identical ModelWeights surface +provenance / checksums → larql verify byte-identical verdicts +browse extraction + WALK → identical top-K rankings +attention-only client slice → loads, serves +full inference → token-identical decode vs pre-v2 binary +layer sharding (--layers) → identical RSS bound behaviour +expert sharding (--experts) → identical 404-before-read behaviour +publish / pull round-trip → hub artifacts unchanged; generation stamp added + without invalidating existing v1 artifacts +generation dispatch → index.json.version 2→3 routing correct; + unknown version fails naming both sides +``` + +#### Two gates, named separately [amended 2026-08-01] + +"E0 green" must never be reported when only part of it ran. The matrix splits by what a runner can actually execute: + +| Gate | Covers | Where | +| ---- | ------ | ----- | +| **E0-CI** | Generation/schema boundary and every weight-free compatibility check. Needs no checkpoints. | Required merge check on every PR touching the crate | +| **E0-FULL** | The checkpoint-backed matrix — token-identical decode, WALK rankings, slicing, sharding, publish/pull — against the committed C3 goldens | Local, against a real index; status recorded as "green at commit ``" | + +Reporting convention: + +```text +E0-CI: green on every merge +E0-FULL: green at last recorded local run +``` + +E0-CI is deliberately the subset most at risk from VINDEX3 work: every commit touching the loader, the generation module or the LYRW parsers can break it, and nothing else in CI would notice. It caught a real regression on its first run — `index.json` schema 1 refused as a non-generation — which is the argument for it existing. + +A later improvement would add a tiny synthetic VINDEX2 artifact exercising load, verify, slice and one trivial decode in CI. That would not replace the multi-GB goldens, but it would narrow the gap between boundary-only checking and full local validation. + +**Acceptance:** zero behavioural change on every v1 path, measured **against the C3 goldens** — token-identical serving, identical verify verdicts, identical slice/publish/pull outputs. The K3 artefact itself only ever needs v2; E0 protects everyone already on v1. + +**Decision rule:** any E0 regression blocks merge, full stop. VINDEX3 defaulting for new extractions (§12.1 support policy) is gated on E0 green plus the ABI freeze. + +### E1 — LYRW physical-layout gate + +Determines whether v2's layout ideas beat the current format before any K3 byte is written. + +**Arms** +``` +L0 current LYRW v1 (one format per layer file) +L1 LYRW v2 per-region formats, same file & offset-table shape +L2 projection-separated files (layer_N.gate_up / layer_N.down) +L3 expert-group files (layer_N.experts_000_031 …) +``` + +Run first on Mini-K3, all arms with identical Q6_K weights. + +**Measure:** exact output parity; loader complexity; mmap/open count and latency; dispatch construction time; bytes faulted; sequential and random top-16 execution; manifest-level omission/override capability; **gate-only browse cold bytes** (a WALK sweep touching only gate regions — the browse-read cost of each layout). + +**Decision rule.** Choose the **least fragmented** layout within: +- 2% of best warm execution; +- 5% of best cold physical bytes read; +- exact parity. + +Do not promote L2/L3 for conceptual tidiness. **Registered prior: L1 wins.** The fused-vs-decomposed gate/up question inside L1 is **not** settled here — E1 records the serving delta between the two; E7 records the browse delta; §15.2's per-bank rule reconciles them at extraction time. + +### E2 — Byte-faithful SSD replay (segment & ordering) + +On the 1–4-layer byte-faithful bank, replay routing traces under memory budgets of 2/4/8/16 GB: + +| Trace | Purpose | +| ----- | ------- | +| Uniform random top-16 | worst-case balanced routing | +| Zipf/skewed | strong hot-expert case | +| Gemma empirical | existing real proxy | +| Kimi-Linear empirical | real same-lineage MoE traces (sigmoid top-8-of-256) — a better structural proxy than Gemma, still NOT a K3 locality substitute | +| Inkling-Small empirical | real sink-router top-6-of-256 traces captured during rung-4 serving — a third router family for the replay table; same non-transfer rule applies | +| Clustered co-activation | tests physical grouping | +| Adversarial rotation | defeats cache and prefetch | +| Repeated conversation | temporal locality | + +**Sweep — two independent scales** (spec §7): group-extent width ∈ {8, 16, 32, 64} experts × segment width ∈ {112, 224, 448, 896-split-2} experts, plus physical orderings ID / frequency / co-activation-cluster / random (via the entry-table indirection — ordering is independent of both widths). Group width optimises reads/prefetch/dispatch; segment width optimises file count, mmap management and the 20 GiB cap — the winning pair is reported separately, never as one number. + +**Measure:** actual SSD bytes/token; useful bytes vs read amplification; page faults; read count and extent sizes; p50/p95/p99 layer latency; degradation onset; late steady-state floor; next-layer prefetch effect; permutation effect. + +**Decision rules.** +- Group width: simplest width within a small margin of best sustained result; must divide the segment width and match a grouped-kernel dispatch width. +- Segment width: as large as the 20 GiB cap and shard-distribution needs allow (prior: 448 for K3 Q6_K); must be a whole multiple of the group width. +- Physical ordering: bake a non-ID permutation into the permanent index **only** on ≥10% sustained improvement across multiple plausible traces with no bad uniform-routing regression. **Registered prior: per-expert offsets + runtime cache beat permanent locality encoding.** + +### E3 — Loaded-model grouped A/B (critical path; doubles as first V2-3 production experiment) + +On loaded Gemma MoE: + +``` +A existing production expert layout +B current LYRW v1 + grouped execution +C LYRW v2 (E1 winner) + grouped execution +D split-file layout + grouped execution (kept only if E1 didn't kill it) +``` + +Hold constant: representation, selected experts, routing weights, activation layout, command-buffer ownership, output reduction. + +**Measure:** per-layer output equivalence; final logits/token equivalence; command submissions and syncs; warm and cold latency; full-model decode throughput; memory and SSD traffic. + +Pass = the E1/E2 winner survives real buffers, real scheduling and the actual forward path. This is the final proxy gate before K3 integration (per the release ladder's K3-0). + +**E3b — repeat on Kimi-Linear-48B** once V2-4 rung 3 imports it: same held-constants, arms B/C only. This is the first grouped A/B with a shared expert inside the dispatch path and non-softmax gate weights in the reduction — the two things Gemma cannot exercise — and it doubles as the KDA/MLA-spine integration shakedown for the K3 adapter. + +### E4 — Projection-specific quantisation (the decisive LYRW v2 justification) + +**Arms** + +| Arm | gate/up | expert down | +| --- | ------- | ----------- | +| Q0 | Q6_K | Q6_K | +| Q1 | native MXFP4 | Q6_K | +| Q2 | Q6_K | native MXFP4 | +| Q3 | native MXFP4 | native MXFP4 | +| Q4 | approx lower | Q6_K | +| Q5 | Q6_K | approx lower | + +No Cartesian sweep. Three stages: (1) Mini-K3 numerical mechanics, (2) loaded Gemma full-model proxy, (3) real K3 layer replay once the trace pack exists. + +**Measure:** expert-output cosine and normalised error; weighted aggregate latent output; next-layer router agreement; teacher-forced BPB; token-level KL; greedy agreement; long-generation stability; measured bytes and kernel throughput. + +**Decision rule — reframed.** Per-region format tags are in the ABI **structurally**: native value/scale codecs, v1's existing mixed gate/up/down precision, and format-neutral banks justify representability without a performance argument. E4 therefore decides one thing only: **whether a mixed-format profile is promoted to Production** — requiring a mixed arm to beat both uniform alternatives by ≥5% projected whole-system gain at the same fidelity gate. A failed E4 leaves the format intact and the mixed profiles at Reference/Grouped maturity (representable ≠ servable). + +### E5 — Weight omission (profile semantics; NOT an ABI gate) + +E5 decides profile authority and approximation policy. No E5 outcome changes a byte of LYRW2 — it cannot block the freeze. + +Disambiguation first: K3's shared latent `routed_expert_down_proj` (7168→3584) is not an expert's `w2` (3072→3584). Dropping a selected expert's `w2` yields no expert output — dominated by skipping the expert. Arms therefore test **useful** omission modes: + +``` +O0 exact baseline (top-16) +O1 reduced top-K: 16/12/8/4 — with and without gate renormalisation +O2 cumulative gate-mass retention: 99 / 97.5 / 95 / 90 % +O3 shared-only layers — per-layer interventions, then structured patterns + (every 4th MoE layer; early/middle/late; lowest-measured-impact) +O4 entire routed branch remote — whole-branch RPC (~14 KB/layer f16) + vs projection-split (~100 KB/layer): do NOT split unless measurement + overturns the arithmetic +O5 down-only lower precision — all experts execute; cheap w2 representation + (the honest "cheap down"; feeds E4/Q2/Q5) +``` + +**Promotion gate (production approximation):** ≤0.5% BPB regression; no pathological p99 KL; no next-layer routing-agreement collapse; no recurrent/long-generation instability; gain recomposed against the whole-model byte ledger. A layer-selective O3 result is valuable even if whole-model routed removal fails. + +### E6 — Real K3 routing/locality traces + +The only experiment that cannot be run honestly on Gemma or synthetic routing — **nor on Kimi-Linear**: KL-48B traces (cheap to capture locally) pilot the E6 capture/analysis machinery and feed the E2 replay table, but its sigmoid grouped-top-k router and K3's quantile-balanced router are different balancing mechanisms, so no K3 locality, hot-set or retention decision transfers from it. Required **before** freezing any physical expert permutation, permanent hot set, co-activation grouping, or per-layer retention policy. + +**Compute:** per-layer frequency entropy; cumulative mass curves; co-activation matrices; temporal reuse-distance distributions; static-LFU / LRU / hybrid per-layer cache hit rates; hot-set stability across workloads. + +**Decision rule — scope.** E6 blocks exactly one physical decision: adopting a non-ID expert permutation (via E2's ≥10% bar re-run on real K3 traces). It does **not** gate the ABI freeze — the ID-order layout is always legal, and locality otherwise stays runtime metadata permanently. Cache policy, residency and retention conclusions feed profiles, not the format. + +### E7 — Query-layer conformance and browse economics + +Proves "the model IS the database" survives v2 as a measured property, not a slogan. + +**Arms** (on fixtures A–D, then the byte-faithful bank, then loaded Gemma): + +``` +W0 VINDEX2 control: WALK over an extracted gate_vectors.bin + — DENSE MODELS ONLY, see the coverage note below +W1 in-place WALK over decomposed gate regions (f16) +W2 in-place WALK over fused gate_up regions, strided (f16 row-major) +W3 in-place WALK over block-quantised gate regions (lazy dequant) +W4 gate-only browse SLICE built by region copy (spec §15.5) +``` + +#### Coverage note — W0 is not a control for expert regions [amended 2026-08-01] + +Measured on a fresh VINDEX2 extract of Gemma 4 26B A4B (128 experts × 704 expert width, hidden 2816, 30 layers): the shipped generation's `gate_vectors.bin` exposes **2,112 walkable features per layer** — the *dense* FFN width. The expert population would contribute **128 × 704 = 90,112**. The expert weights are present and decode correctly (30 files, 12 GB); they are simply **not part of the searchable surface**. + +So W0 covers 2.3 % of what §15.1 specifies, and on a MoE model it contains no expert regions at all. W1/W2/W3/W4 are *about* expert regions. Comparing them to W0 would measure a coverage difference wearing a parity result's clothes, and the registered pass condition ("identical top-K rankings") is unevaluable because the two arms do not rank over the same population. + +**Resolution — the claim is restated, not the control rebuilt.** Making expert features searchable is a **new capability of VINDEX3, not parity with VINDEX2**. Writing it up as parity would be claiming a comparison that cannot be made. Concretely, E7 splits: + +| Claim | Arms | Baseline | +| ----- | ---- | -------- | +| **Parity** — in-place region WALK matches the extracted-index path | W0 vs W1 on a **dense** model | W0, a genuine like-for-like control | +| **Capability** — expert features are searchable in place, with no separate index | W1/W2/W3/W4 on MoE | none exists; correctness is established against the weights themselves, not against W0 | + +The rejected alternative was constructing an expert-bearing VINDEX2 index specifically to serve as a control. It would be a control built *for* the experiment rather than the thing that shipped — a weaker claim that reads as a stronger one. + +**Measure:** top-K ranking parity vs W0 **on the dense arm** (exact for f16; ranking-overlap metric for W3, with the v1 §12.2 4-bit noise caveat as the expected failure shape); for the capability arm, ranking correctness against directly-computed gate dot products rather than against W0. Plus, for both: cold bytes faulted per WALK; queries/sec warm; slice size vs v1's ~3 GB browse economics; latent-WALK correctness and query-projection overhead on fixture D and the K3-shaped bank; DESCRIBE/SELECT correctness against `query/` sidecars, including latent-bank `down_meta` computed through the full output path. + +**Decision rules.** +- W1 vs W2 sets the browse half of the §15.2 fusion decision: fused storage keeps browse eligibility only if strided WALK is within 10% of decomposed on cold bytes and warm throughput; otherwise browse-enabled indexes mandate decomposed gate/up. +- W3 promotes quantised-region browse only at a pre-declared ranking-overlap floor (top-50 overlap ≥ 0.9 vs W0); below it, browse-enabled extraction keeps gate at f16 regardless of the serving format — a legitimate per-region format divergence that E4 machinery already supports. +- W4 must reproduce W1 rankings bit-for-bit (it is the same bytes relocated), or the slicer is wrong. + +### E8 — Held-out architecture (generalisation, not fit) + +The conformance fixtures and design-set models cannot prove spec §16 criterion 7, because the ABI was designed against them — and **Kimi-Linear-48B and Inkling-Small are now in the design set, so both are ineligible here**. E8 onboards a MoE deliberately excluded from the design set — candidate: Mixtral 8x7B or a Qwen-MoE (cheap, well-documented, conventional) — **after** the ABI freeze, under a hard rule: + +``` +allowed: checkpoint importer, MoE manifest, programme adapter + (existing registered programmes only, or ONE new programme_id + using existing region roles) +forbidden: LYRW2 byte-layout changes · new region roles · new packing modes · + kernel-interface changes · loader special cases keyed on the model +``` + +**Acceptance:** token parity against a trusted implementation through the generic reference path, then through grouped dispatch if a compatible kernel exists — with a diff of the vindex crates showing zero format-layer changes. + +**Falsification is a real outcome:** if E8 requires format changes, spec §16 downgrades the substrate claim to "K3/GPT-OSS/Inkling serving format" in writing, and the needed changes are queued for a future LYRW revision — they are not smuggled in retroactively. + +--- + +## 4. Run order + +**Run now (no K3 adapter required)** +0. E0 preservation matrix wired into CI (stays green for the life of the programme) +1. Fixtures A–D + Mini-K3 → V2-0, V2-1 +2. E1 layout gate +3. E2 one-layer byte-faithful bank (extend to 4 layers if one layer is inconclusive) +4. E3 loaded Gemma grouped A/B +5. E4 stages 1–2 +6. E7 query-layer conformance (fixtures + Gemma; the K3-shaped latent-WALK arm rides the byte-faithful bank) + +**Freeze the ABI** — after E0 green + E1/E2/E3/E7 decided + V2-0/V2-1 accepted: +7. LYRW `format_version=2`, `index.json` v3, `vindex_spec_version=2`, MoE manifest v1. E5/E6 do **not** gate this (they decide profiles and permutation only); E4 does not either (per-region tags are structural, E4 gates promotion). + +**After freeze, before the K3 extraction** (profile and locality decisions): +8. V2-4 rung 3: Kimi-Linear-48B import → token parity → E3b grouped A/B — the K3 adapter dress rehearsal (KDA/MLA spine, shared expert, sigmoid router) +9. KL-48B trace capture — pilots the E6 machinery, feeds the E2 replay table +10. V2-4 rung 4: Inkling-Small BF16 import → text-backbone parity → NVFP4 variant import (paired values/scales, mixed-precision convention) → first forced partial-residency + attn-local/FFN-remote serving — the K3 serving dress rehearsal +11. E8 held-out architecture — the generalisation test runs against the frozen ABI +12. K3 trace pack (E6 capture) +13. E4 stage 3 on real layers (mixed-profile Production promotion) +14. E5 on real layers (O1–O5) +15. E6 locality/cache simulation (gates only a non-ID permutation) + +**Then extract K3 once**, into the frozen five-class layout, exact Q6_K baseline (fidelity recorded per §9.2 — `source-equivalent` if losslessly containing native values, `numerically-approximate` otherwise), two 448-expert segments per routed layer with group extents at the E2-chosen width, no locality decisions baked in. + +--- + +## 5. Registered priors (falsifiable, dated 2026-08-01) + +1. One file-set per routed layer (segmented past 20 GiB) remains correct. +2. One entry per expert remains correct. +3. **(Revised 2026-08-01, pre-freeze, on reinstating the query layer.)** Down stays independently addressable. Gate/up fusion is no longer a blanket prior: serving-only indexes fuse; browse-enabled indexes default to decomposed, reconciled per bank by E1 (serving delta) + E7 (browse delta) under §15.2. +4. **(Revised on adopting the variants model.)** Per-region format tags are in the ABI structurally; the prior is now that no mixed-format *profile* clears E4's 5% Production bar on Gemma, and the first to clear it does so only with real K3 layers (E4 stage 3). +5. Expert locality remains runtime metadata (E2/E6 fail the 10% bar). +6. "Dropping down" locally resolves to "skip the routed branch," never a half-expert mode (E5). +7. Exact remote deployment moves the whole routed branch, not split projections (O4 arithmetic holds). +8. L1 wins E1; L2/L3 die there. +9. **(Revised 2026-08-01, on measuring the VINDEX2 control.)** Split in two, because the original prior conflated a parity claim with a capability claim. **(9a)** On a *dense* model, in-place WALK over decomposed f16 gate regions matches the VINDEX2 `gate_vectors.bin` path within noise (E7/W0-vs-W1). **(9b)** On a *MoE* model, expert features are searchable in place at all — which VINDEX2 does not do, so there is no baseline to match and no dedicated query index is needed because none previously existed for this population. Falsifier for 9b is a correctness failure against directly-computed gate dot products, not a parity failure against W0. +10. Quantised-region browse (E7/W3) fails the ranking-overlap floor for ≤4-bit formats, and browse-enabled indexes keep gate at f16 as a per-region divergence — the query layer becomes the second consumer of per-region format tags after E4. +11. Dual-generation support costs nothing on the v1 hot path — E0 stays green throughout without a single v1-loader change (the generations share a trait, not code). +12. E8 passes: the held-out conventional MoE onboards with an importer + manifest + existing `gated-mlp-v1` programme, zero format-layer diffs — the envelope generalises beyond its design set. +13. K3's exact-Q6_K baseline lands at `source-equivalent`, not `source-exact` — and nothing downstream ever quotes it as bit-faithful to the release encoding. +14. Kimi-Linear-48B onboards through `gated-mlp-v1` + shared bank + existing region roles — its entire KDA/MLA spine lands in classes 1–2 as manifest-addressed tensors, touching the expert-bank format not at all. +15. The KL-48B→K3 adapter delta is confined to the five documented K3 additions (SiTU-GLU, AttnRes, latent MoE transforms, MLA output gate, full-rank KDA gate) — i.e. the dress rehearsal genuinely de-risks the main blocker rather than merely preceding it. +16. Inkling-Small onboards through `gated-mlp-v1` + shared bank + existing region roles; the sink router, gate bias, post-top-k norm and route/global scales are entirely manifest router-vocabulary items — no new region role, no new programme beyond a router descriptor. +17. The NVFP4 and MXFP8 releases import as paired values/scales regions using the existing packing vocabulary (packing 2/3 + pair_id) — no new packing mode is needed; if one is, that is an ABI-RC finding, not a post-freeze patch. + +Any prior falsified is recorded here with the run ID and the format consequence, then the spec is amended before freeze — never after. + +--- + +## License + +Apache-2.0 diff --git a/docs/virtual-experts-dispatch.md b/docs/virtual-experts-dispatch.md index 5d44eb1a7..6276ebfc3 100644 --- a/docs/virtual-experts-dispatch.md +++ b/docs/virtual-experts-dispatch.md @@ -173,7 +173,7 @@ with parallel tests). 5 unit tests cover all four precedence outcomes. `RemoteExpertBackend` → `RemoteMoeBackend`, `RemoteExpertError` → `RemoteMoeError`, `generate_with_remote_experts` → `generate_with_remote_moe`, `examples/expert_grid_generate.rs` → -`examples/moe_grid_generate.rs`. Module doc explicitly disambiguates from +`chris-experiments/larql_probes/examples/misc/moe_grid_generate.rs`. Module doc explicitly disambiguates from `crate::experts`. Side-effect of running the rename: caught two pre-existing build breakages @@ -655,7 +655,7 @@ crates/larql-inference/src/ffn/{moe_remote,mod}.rs # rename + new fields crates/larql-inference/src/layer_graph/{generate,grid,mod}.rs # generate_constrained crates/larql-inference/src/lib.rs # re-exports crates/larql-inference/tests/{data/,test_generate_q4k_cpu,test_*_dispatch}.rs -crates/larql-inference/examples/moe_grid_generate.rs # renamed +chris-experiments/larql_probes/examples/misc/moe_grid_generate.rs # renamed crates/larql-cli/src/commands/primary/run_cmd.rs # --experts + --constrained crates/larql-cli/src/main.rs # ChatArgs ↔ RunArgs crates/larql-cli/tests/test_run_experts.rs # CLI integration tests diff --git a/scripts/e0-capture-goldens.sh b/scripts/e0-capture-goldens.sh new file mode 100755 index 000000000..60285646e --- /dev/null +++ b/scripts/e0-capture-goldens.sh @@ -0,0 +1,112 @@ +#!/usr/bin/env bash +# E0 corpus C3 — capture pinned golden outputs from the PRE-V2 binary. +# +# E0 (docs/vindex2-experiments.md) asserts that adding VINDEX3 support to the +# binary changes nothing observable on any VINDEX2 path. That assertion needs a +# fixed record to compare against. +# +# WHY THIS EXISTS AT ALL: without committed goldens, "zero behavioural +# regression" silently degrades into "the two binaries agree with each other" — +# a condition any bug present in BOTH satisfies. Since the whole point of E0 is +# to catch an incumbent path damaged by successor work sharing one binary, +# comparing two builds of that binary to each other is close to circular. The +# baseline has to predate the successor code and has to be committed. +# +# WHEN: run this against a binary built from a commit BEFORE any lyrw2 work. +# Once VINDEX3 merges to the main line, the baseline stops being a checkout and +# starts being an archaeology exercise — and a reconstructed baseline is exactly +# the artefact that drifts without anyone noticing. +# +# THE OTHER HALF is `e0-verify-goldens.sh`, which replays this record against a +# current binary. Capturing without ever verifying leaves an assertion nobody +# makes. The corpus itself lives in `lib/e0-corpus.sh` and is shared by both, so +# the two cannot drift apart. +# +# Usage: +# E0_BIN= E0_VINDEX=output/gemma.vindex ./scripts/e0-capture-goldens.sh +# +# Env vars: +# E0_BIN — larql binary to capture from (required; must be pre-v2) +# E0_VINDEX — vindex to exercise (required) +# E0_MODEL — checkpoint the vindex came from (recorded in the recipe) +# E0_MODEL_REVISION — that checkpoint's revision hash (recorded in the recipe) +# E0_EXTRACT_FLAGS — flags used to extract it (recorded in the recipe) +# E0_OUT — golden output dir (default: tests/goldens/e0/) +# E0_TOKENS — tokens to decode per prompt (default: 24) +# E0_WALK_K — WALK top-K per layer (default: 20) +# +# The vindex itself is NOT committed — it is regenerable. The recipe is. +# +# Determinism: `larql run` samples greedily (SamplingConfig::greedy), so decode +# is reproducible without a seed flag. Every command's output is normalised for +# paths and timings before writing, so a diff means a behavioural change rather +# than a different working directory or a faster machine. + +set -uo pipefail + +readonly SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +# shellcheck source=lib/e0-corpus.sh +source "${SCRIPT_DIR}/lib/e0-corpus.sh" + +readonly TOKENS="${E0_TOKENS:-24}" +readonly WALK_K="${E0_WALK_K:-20}" + +: "${E0_BIN:?set E0_BIN to a larql binary built BEFORE any lyrw2 commit}" +: "${E0_VINDEX:?set E0_VINDEX to a vindex directory}" + +readonly OUT_DIR="${E0_OUT:-tests/goldens/e0/$(basename "$E0_VINDEX")}" +readonly TMP="$(mktemp -d)" +trap 'rm -rf "$TMP"' EXIT + +mkdir -p "$OUT_DIR" + +echo "E0 C3 golden capture" +echo " binary: ${E0_BIN}" +echo " vindex: ${E0_VINDEX}" +echo " out: ${OUT_DIR}" +echo + +# Provenance + regeneration recipe. +# +# The corpus is reproducible, so it is pinned rather than stored: the multi-GB +# vindex is not committed, this recipe is. Two things must be here or the golden +# set is unfalsifiable — the baseline commit (without it, a reader cannot tell +# whether these goldens predate the successor work they police) and the exact +# extract flags (without them, "re-extract it" is not a reproducible +# instruction). +{ + echo "captured_from_commit=$(git rev-parse HEAD)" + echo "captured_from_describe=$(git describe --always --dirty)" + echo "vindex=$(basename "$E0_VINDEX")" + echo "tokens=${TOKENS}" + echo "walk_k=${WALK_K}" + echo + echo "# ── C1 regeneration recipe ──" + echo "# rebuild the baseline binary:" + echo "# git checkout $(git rev-parse HEAD) && cargo build --release -p larql-cli" + echo "# re-extract the vindex:" + echo "model_source=${E0_MODEL:-}" + echo "model_revision=${E0_MODEL_REVISION:-}" + echo "extract_flags=${E0_EXTRACT_FLAGS:-}" +} > "${OUT_DIR}/PROVENANCE.txt" + +# Every command row, driven from the shared corpus so capture and verify +# exercise an identical set. +while IFS= read -r row; do + IFS='|' read -r -a parts <<< "$row" + name="${parts[0]}" + cmd=("${parts[@]:1}") + printf ' → %-28s' "$name" + e0_run_row "$name" "$E0_VINDEX" "$TMP" "${cmd[@]}" > "${OUT_DIR}/${name}.txt" + echo "$(tail -n1 "${OUT_DIR}/${name}.txt")" +done < <(e0_rows "$E0_BIN" "$E0_VINDEX" "$TMP" "$TOKENS" "$WALK_K") + +# index.json is the generation discriminator (spec §12.1). +printf ' → %-28s' index_version +e0_index_version "$E0_VINDEX" > "${OUT_DIR}/index_version.txt" +echo "ok" + +echo +echo "captured $(find "$OUT_DIR" -name '*.txt' | wc -l | tr -d ' ') golden files into ${OUT_DIR}" +echo "COMMIT THESE. They are the only fixed record E0 has." +echo "Verify them later with: E0_BIN=... E0_VINDEX=... ./scripts/e0-verify-goldens.sh" diff --git a/scripts/e0-verify-goldens.sh b/scripts/e0-verify-goldens.sh new file mode 100755 index 000000000..f9535c58a --- /dev/null +++ b/scripts/e0-verify-goldens.sh @@ -0,0 +1,159 @@ +#!/usr/bin/env bash +# E0-FULL — replay the committed goldens against the CURRENT binary. +# +# The capture script writes the record. Until now nothing read it back, so the +# goldens were an assertion no one made: committed, provenance-stamped, and +# never compared to anything. +# +# This is the other half. It runs the same corpus (shared via +# `lib/e0-corpus.sh`, never redefined here) against a current build and diffs +# every row. Any difference is a behavioural change on a VINDEX2 path caused by +# VINDEX3 work sharing one binary — which is the entire claim E0 exists to +# police. +# +# SCOPE. This is E0-FULL, not E0-CI. It needs the multi-GB vindex, which is not +# committed because it is regenerable; the recipe to rebuild it is in the +# golden set's PROVENANCE.txt. E0-CI covers the weight-free generation-boundary +# subset and runs on every push. +# +# Usage: +# E0_BIN=target/release/larql E0_VINDEX=/path/to/gemma4-26b-a4b.vindex \ +# ./scripts/e0-verify-goldens.sh +# +# Env vars: +# E0_BIN — larql binary to verify (required) +# E0_VINDEX — the VINDEX2 index to exercise (required) +# E0_GOLDENS — golden dir (default: tests/goldens/e0/) +# E0_TOKENS — tokens to decode per prompt (default: 24) +# E0_WALK_K — WALK top-K per layer (default: 20) +# E0_DIFF — set to 1 to print full diffs rather than a summary +# +# Exit status is the result: 0 = every row matched, 1 = at least one differed. + +set -uo pipefail + +readonly SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +# shellcheck source=lib/e0-corpus.sh +source "${SCRIPT_DIR}/lib/e0-corpus.sh" + +: "${E0_BIN:?set E0_BIN to the larql binary under test}" +: "${E0_VINDEX:?set E0_VINDEX to a VINDEX2 directory}" + +readonly GOLDEN_DIR="${E0_GOLDENS:-tests/goldens/e0/$(basename "$E0_VINDEX")}" + +# Replay parameters come from the record itself. +# +# The capture wrote `tokens` and `walk_k` into PROVENANCE.txt precisely because +# they change the output: replaying 24 tokens against a 16-token golden differs +# on every decode row, and the diff would read as a behavioural regression +# rather than as a harness mismatch. Defaulting instead of reading would make +# E0-FULL fail loudly for the wrong reason — the most expensive kind of false +# positive, because it looks exactly like the thing being tested for. +provenance_field() { + local field="$1" default="$2" + local file="${GOLDEN_DIR}/PROVENANCE.txt" + [[ -f "$file" ]] || { echo "$default"; return; } + local value + value="$(grep "^${field}=" "$file" 2>/dev/null | cut -d= -f2-)" + echo "${value:-$default}" +} + +readonly TOKENS="${E0_TOKENS:-$(provenance_field tokens 24)}" +readonly WALK_K="${E0_WALK_K:-$(provenance_field walk_k 20)}" +readonly TMP="$(mktemp -d)" +trap 'rm -rf "$TMP"' EXIT + +if [[ ! -d "$GOLDEN_DIR" ]]; then + echo "no goldens at ${GOLDEN_DIR}" >&2 + echo "capture them first with scripts/e0-capture-goldens.sh" >&2 + exit 2 +fi + +echo "E0-FULL golden verification" +echo " binary: ${E0_BIN}" +echo " vindex: ${E0_VINDEX}" +echo " goldens: ${GOLDEN_DIR}" +if [[ -f "${GOLDEN_DIR}/PROVENANCE.txt" ]]; then + echo " baseline: $(grep '^captured_from_describe=' "${GOLDEN_DIR}/PROVENANCE.txt" | cut -d= -f2-)" +fi +echo " replay: tokens=${TOKENS} walk_k=${WALK_K} (from the golden's provenance)" +echo + +matched=0 +differed=0 +missing=0 +declare -a failures=() + +# Canonicalise a stored golden for comparison. +# +# Goldens for unordered rows were captured before the ordering fix, so they +# hold whatever permutation the incumbent binary happened to emit. Sorting the +# stored side too lets an existing golden remain usable instead of requiring a +# re-capture to test a property that was never about order. The trailing +# `exit_status=` line is held back so a sorted row still ends with its status. +canonical_golden() { + local name="$1" golden="$2" out="$3" + if e0_row_is_unordered "$name"; then + sed '$d' "$golden" | sort > "$out" + tail -n1 "$golden" >> "$out" + else + cp "$golden" "$out" + fi +} + +check_row() { + local name="$1" actual="$2" + local golden_raw="${GOLDEN_DIR}/${name}.txt" + printf ' %-28s' "$name" + if [[ ! -f "$golden_raw" ]]; then + echo "NO GOLDEN" + missing=$((missing + 1)) + return + fi + local golden="${TMP}/${name}.golden" + canonical_golden "$name" "$golden_raw" "$golden" + if diff -q "$golden" "$actual" >/dev/null 2>&1; then + echo "match" + matched=$((matched + 1)) + return + fi + echo "DIFFERS" + differed=$((differed + 1)) + failures+=("$name") + if [[ "${E0_DIFF:-0}" == "1" ]]; then + diff -u "$golden" "$actual" | sed 's/^/ /' + fi +} + +# ── Every command row ────────────────────────────────────────────────────── +while IFS= read -r row; do + # Fields are '|'-separated: name, then the command and its arguments. Split + # on '|' rather than whitespace because prompts contain spaces. + IFS='|' read -r -a parts <<< "$row" + name="${parts[0]}" + cmd=("${parts[@]:1}") + e0_run_row "$name" "$E0_VINDEX" "$TMP" "${cmd[@]}" > "${TMP}/${name}.actual" + check_row "$name" "${TMP}/${name}.actual" +done < <(e0_rows "$E0_BIN" "$E0_VINDEX" "$TMP" "$TOKENS" "$WALK_K") + +# ── The generation discriminator ─────────────────────────────────────────── +e0_index_version "$E0_VINDEX" > "${TMP}/index_version.actual" +check_row index_version "${TMP}/index_version.actual" + +# ── Verdict ──────────────────────────────────────────────────────────────── +echo +echo " matched ${matched}" +echo " differed ${differed}" +[[ "$missing" -gt 0 ]] && echo " missing ${missing} (golden absent — capture is stale)" + +if [[ "$differed" -eq 0 && "$missing" -eq 0 ]]; then + echo + echo "E0-FULL PASS — no observable change on any VINDEX2 path." + exit 0 +fi + +echo +echo "E0-FULL FAIL — VINDEX2 behaviour changed in: ${failures[*]:-}" +echo "Re-run with E0_DIFF=1 to see the differences." +echo "If a change is intentional, re-capture deliberately and say why in the commit." +exit 1 diff --git a/scripts/lib/e0-corpus.sh b/scripts/lib/e0-corpus.sh new file mode 100644 index 000000000..5be399eff --- /dev/null +++ b/scripts/lib/e0-corpus.sh @@ -0,0 +1,128 @@ +# E0 corpus definition — the rows, the prompts, and the normalisation. +# +# Sourced by both `e0-capture-goldens.sh` (writes the record) and +# `e0-verify-goldens.sh` (replays it against a current binary). +# +# WHY THIS IS SHARED RATHER THAN DUPLICATED +# +# The capture script's own preamble warns that comparing two builds to each +# other is close to circular. The same hazard applies one level up: if capture +# and verify each carried their own copy of the prompts, the row list or the +# normalisation, they would drift, and a passing E0-FULL would mean "the two +# scripts agree" rather than "behaviour is unchanged". A prompt reworded on one +# side alone turns every downstream row into a false diff; a normalisation rule +# added on one side alone silently hides a real one. +# +# So the corpus is defined once, here, and neither script may redefine it. + +# Fixed prompt set. Short, deterministic, and spanning factual recall, code and +# multi-token continuation so a regression in any one does not hide in the +# others. +readonly E0_PROMPTS=( + "The capital of France is" + "def fibonacci(n):" + "Water boils at" +) + +readonly E0_SLICE_PRESETS=(client attn embed server browse router expert-server all) + +# Strip anything that varies between runs or machines but is not behaviour: +# absolute paths, wall-clock durations, throughput rates, ISO timestamps. +# +# Takes the vindex path and scratch dir as arguments rather than reading +# globals, so a caller cannot accidentally normalise against the wrong paths. +# +# # No `\b` +# +# The previous timing rule ended in `\b`, which is a GNU extension. BSD +# `sed -E` does not support it and does not complain — it just never matches. +# So on macOS every duration passed through unstripped, and each replay +# differed from its golden on the timing line alone. A normaliser that +# silently does nothing is worse than no normaliser: it produces confident +# diffs that read exactly like the regression under test. +# +# The portable form matches the unit followed by a non-alphanumeric character +# or end-of-line, and puts that character back via the capture group. `ms` +# precedes bare `s` so the longer unit wins, and `MB`/`GB` on their own are +# deliberately absent — those are file sizes, which are behaviour and must +# survive normalisation. +e0_normalise() { + local vindex="$1" tmp="$2" + sed -E \ + -e "s#${vindex}##g" \ + -e "s#${tmp}##g" \ + -e 's#/[^ ]*/(larql|target)[^ ]*##g' \ + -e 's/[0-9]+(\.[0-9]+)? *(ms|us|ns|tok\/s|MB\/s|GB\/s|s)([^A-Za-z0-9]|$)/