diff --git a/README.md b/README.md index be9038479..1338532bc 100644 --- a/README.md +++ b/README.md @@ -625,6 +625,8 @@ vs ollama gemma3:4b on the same machine: ~103 tok/s steady → **gap 1.17×**, w **CPU vs llama.cpp** (reconciled 2026-06-02, M3 Max, 8 threads, warm): larql **26.4** (StandardEngine) / 23.5 (legacy `bench --cpu`) vs **llama.cpp `-ngl 0` 43.0** tok/s → **gap ~1.6–1.8×**. The gap is per-core kernel quality — both attention and FFN already run the int8 Q8_K SDOT kernel; closing it is C12 (hand-asm; an opt-in `LARQL_Q4K_ASM=1` v1 lands +~4% isolated). `larql bench --cpu` now reports both the legacy and production-StandardEngine rows; `--ollama-cpu` forces a true CPU ollama baseline (default `--ollama` runs on Metal GPU). The earlier 1.5×/1.9× spread was two measurement confounds (path mismatch + an unwarmed-ollama artifact), not a regression — see `bench/baselines/c10_gemma3-4b_cpu_reconciled.json`. +**CPU prefill** (2026-06-22): the per-layer f32 dequant — long the dominant prefill cost (~2.7 s / ~2 tok/s on the 5-token prompt) — is gone. Q/K/V/O **and** gate/up/down now project straight from the Q4_K/Q6_K vindex bytes via amortised `q4k_matmul` / `q6k_matmul` (the Q6_K twin handles the default Q6_K `v_proj` / `down_proj`) with a hand-written aarch64 NEON inner dot. Gemma 3 4B Q4_K CPU prefill: **2746 ms → 233 ms (11.8×)**, closing the gap to llama.cpp `pp5` from ~55× to **~3×**; the NEON `q4k_matmul` at seq=5 beats f32 AMX sgemm while still skipping the dequant. See `bench/baselines/cpu/COMPARISON.md`. + **Cross-arch coverage (2026-05-09)**: Gemma 3, Gemma 4 31B dense, Llama 2 7B, Mistral 7B all dispatch correctly through Metal. Gemma 4 E2B currently falls back to CPU (Per-Layer Embeddings not yet in Metal — ROADMAP D-METAL-PLE). See [crates/larql-compute/docs/architecture-shader-map.md](crates/larql-compute/docs/architecture-shader-map.md) for the per-architecture shader dispatch table. CPU walk breakdown: diff --git a/ROADMAP.md b/ROADMAP.md index fc1920ea0..4e5eaa2a9 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -187,6 +187,7 @@ Current state (2026-05-15): | **GPU (Metal)** | Gemma 4 + MTP (when adopted) | 88 tok/s no-MTP | ~225 with MTP | ~2.6× behind | far over | | **CPU** | Gemma 3 4B Q4K decode | **30.9 tok/s** (residency default-on, same-session 2026-06-13; was 24.5) | llama.cpp Q4_K_M CPU ~43 | **~1.42× behind** (was 1.69×) | over — the Q4_K residency + int8 + asm + spin-pool stack is now **default-on** (2026-06-13); earlier kernels (KV-cache, direct Q4_K matvec, NEON Q4_K/Q6_K/f32_dot, Q4 lm_head, par_chunks_mut(32), Q4_K×Q8_K sdot, auto-t=8) landed 2026-05-15/16, ~86× over the original 0.36 baseline. See `bench/baselines/cpu/DIAGNOSIS.md` | | **CPU** | Gemma 4 26B-A4B decode | in-proc KV-cached MoE **~35 tok/s** (spin pool, default-on; M3 Max t=8 warm n=256, 2026-06-13) | llama.cpp Q4_K_M CPU **32.1** (recorded, drift-bracketed) | **larql ~9% AHEAD** | ✅ **CAUGHT** — arc 7.6 → 13.9 (residency) → 21.7 (int8 attn) → 27.9 (KV append-in-place) → **~35 (spin-barrier pool)**. The final ~1.15× was **rayon fork-join overhead** (decode driver ran outside the pool → ~211 cold-path sections/token, ~40% of thread-time in waits), *not* kernel quality — closed by the spin pool (effective-bandwidth/scheduling, exactly as the C12 roofline-crossover entry predicted), shipped **default-on**. Caveat: the 26B llama.cpp anchor is the recorded 32.1 (ollama wouldn't run the HF GGUF on CPU this session); machine validated via 4B llama.cpp 44 ≈ recorded 43. `bench/baselines/c10_gemma4-26b-a4b_cpu_reconciled.json`. | +| **CPU** | Gemma 3 4B Q4K **prefill** (5-tok) | **233 ms** (q4k/q6k-direct attn+FFN + NEON dot; standard engine, 2026-06-22; was 2746 ms / ~2 tok/s) | llama.cpp pp5 ~70 ms | **~3.3× behind** (was 55×) | closing — eliminated the per-layer f32 dequant: Q/K/V/O **and** gate/up/down project straight from the Q4_K/Q6_K vindex bytes via amortised `q4k_matmul` / `q6k_matmul` (the Q6_K twin, for the default Q6_K `v_proj`/`down_proj`) with a hand-written aarch64 NEON inner dot that *beats* f32 AMX sgemm at seq=5. A Q6_K-`down_proj` mis-decode (format-tag dispatch) was caught in review and fixed before this number. Remaining gap is matmul constant-factor + batched attention, not dequant. Also de-duplicated the larql-inference↔larql-compute Q4_K forward (one substrate copy). `bench/baselines/cpu/COMPARISON.md` | Items the threshold makes load-bearing (not optional) on the **GPU track**: - **D-ATTN-MTG** — flash attention; without it, attention-mechanism deltas are muddied by missing baseline. @@ -608,6 +609,420 @@ Themes, in leverage order (concrete first steps live in hardening items --- +## BitNet b1.58 integration hardening (added 2026-06-20) + +Native-ternary BitNet (`microsoft/bitnet-b1.58-2B-4T`) landed on +`feat/quant-ternary-a8`. The **W1.58·A8 kernel work in `larql-compute` +is the strongest part and fits the architecture cleanly**: ternary matvec +lives in `cpu/ops/ternary_matvec.rs` alongside the q4k/q6k kernels, follows +the `_into` allocation-free convention, and is parity-disciplined — +dequant-reference parity, **bit-exact NEON-vs-scalar across `cols % 16` +tails**, and shape-guard rejection tests. NEON gives ~12–13× the f32 +reference on BitNet shapes; x86_64 has the scalar-A8 ~2.4× today. + +The **system-level integration is a deliberate parallel stack** — a fresh +instance of the consolidation-track theme above (parallel path created to +avoid destabilising a parity-verified one). BitNet bypasses every shared +seam: `QuantFormat`/`FormatRoute` dispatch, the `KvEngine` trait, +`larql-kv::KvCache`, the models arch registry, and the vindex build +pipeline. This is documented as intentional pending the FormatRoute +roadmap and is a defensible MVP posture for a narrowly-scoped path. The +module comment that gated folding-in on "once the quantised-activation +kernel exists" now has its precondition met (this branch), so the +promotion-or-isolation decision the policy box demands is no longer +blocked — it is made explicit below (the G1–G4 structural items), with the +2026-06-20 hardening pass clearing the quick wins first. + +Per-crate status: + +- **`larql-compute`** — fits. Kernel + tests land in the right module with + the right discipline. The earlier doc inaccuracy (header implied the path + already routes through `FormatRoute`) is **reconciled**: `QuantFormat` + (`pipeline.rs:25-34`) still has no ternary variant, `from_registry_tag` + (`pipeline.rs:104-116`) maps no ternary tag, and the `QuantMatVec` dispatch + trait has no ternary arm — `BitLinearWeight` is reachable only by direct + call, and the docstring now says so plainly (dispatch integration is the + open structural item, not done). +- **`larql-inference`** — bespoke parallel stack. `BitnetModel` is not a + `KvEngine`; `BitnetKvCache` is a hand-rolled `Vec>` rather + than `larql-kv::KvCache`, so none of the KV append-in-place / windowing / + surgery work applies. Entry is direct (`load_bitnet_model` → `generate`), + no unified dispatch picks between dense and ternary. (The hot path now + quantises the shared activation once per Q/K/V and gate/up, and the header + comment reflects the now-met kernel precondition — the structural + KvEngine/shared-cache fold remains open.) +- **`larql-vindex`** — `bitnet_writer`/`bitnet_loader` write a `bitnet/` + sidecar (`*.i2s` + `scales.f32` + `bitnet_layout.json`) and patch + `index.json` with a `bitnet_layout` field, independent of the + `quant: QuantFormat` enum field (two parallel quant-tag mechanisms). + Writer is a post-build patch from `convert_cmd.rs`, not part of the build + pipeline. +- **`larql-models`** — was the fragile seam; **FIXED 2026-06-20**. + `"bitnet-*"` is now recognised explicitly in `detect/mod.rs` and routes to + a thin named `BitnetArch` (`architectures/bitnet.rs`, `family() == + "bitnet"`, Llama-style defaults, `norm_eps` honoured from config) instead + of silently collapsing to `GenericArch`. Native-ternary inference is still + served by the `larql-inference` ternary path, not this trait; `BitnetArch` + is the home for first-class overrides when BitNet graduates. Covered by + `test_detect_bitnet_is_explicit_not_generic`. + +### Completed — hardening pass (2026-06-20) + +The quick-win review items landed; all touched crates build clean and +clippy-clean (`--all-targets`), tests green (compute ternary 19/19, +inference ternary 28/28 incl. the FFN A8-vs-f32 parity gate, models detect +59/59): + +1. ✅ **[larql-models] Killed the silent `GenericArch` fallback** — explicit + `bitnet-*` recognition → thin named `BitnetArch`; `norm_eps` honoured; + `test_detect_bitnet_is_explicit_not_generic`. *(was P1)* +2. ✅ **[larql-compute] Reconciled the `ternary_matvec.rs` docstring** — no + longer implies the path routes through `FormatRoute`; states that dispatch + integration is the open item and the kernel is reached by direct call. +3. ✅ **[larql-inference] Reuse one activation quant** — Q/K/V and gate/up + quantise the shared activation once (`quantize_activation_i8` + + `matvec_i2s_a8_into`) across all five forward sites. Bit-exact (parity + tests unchanged), saves the repeat int8 quantise per projection. +4. ✅ **[larql-inference] Refreshed the `ternary.rs` header comment** — the + "fold in once the quant-activation kernel exists" precondition is now met; + the comment frames the fold as live roadmap work, not a missing dependency. +5. ✅ **[larql-compute] x86_64 gap documented** — verified already clear at + the dispatch entry (`matvec_i2s_a8_into`: "scalar int8 elsewhere — AVX2 + twin is the x86_64 follow-up") and the status block. + +Owed back to the user (not a code change): + +6. **[git hygiene] Split the `pipeline_layer.rs` refactor** — the + `attn_str_to_format`/`ffn_str_to_format` → `from_registry_tag` dedup is a + sound single-source-of-truth cleanup but is **orthogonal to BitNet** + (BitNet never flows through `resolve_ffn_weights`). Land it as its own + "refactor: dedupe tag→format mapping" commit, not inside the feature. + +### Remaining — graduation to first-class (status 2026-06-20 after scoping) + +Close contact with the code (two scoping passes) revised this list: only G1 +was cleanly doable on the current machine. G2 and G4 hit genuine blockers and +G3's framing was falsified. Detail per item: + +- **G1 — `QuantFormat` ternary variant + dispatch — ✅ DONE 2026-06-20.** + Added `QuantFormat::I2S` (+ `registry_tag`/`from_registry_tag` round-trip, + `is_ternary`), a dedicated `QuantMatVec::ternary_matvec` method (a + `BitLinearWeight` carries the per-channel scales the `&[u8]` `quant_matvec` + signature can't), and a `CpuBackend` impl on the best-available A8 kernel. + `quant_matvec` returns `None` for I2S (loud, like Q8_0); Metal panics (no + ternary shader). Registry-reachable now, not only by direct call. Tested + + clippy-clean. +- **G2 — `KvEngine` impl + shared cache — BLOCKED on a breaking trait change + (your call).** Scoping (kv_engine.rs / larql-kv) found `KvEngine::prefill` + and `decode_step` are typed `(&ModelWeights, &dyn FfnBackend, …)` — dense + f32 weights. `BitnetModel` holds ternary `BitLinearWeight`s, an incompatible + container. Making BitNet a first-class `KvEngine` needs EITHER a + workspace-wide generalisation of the weight parameter to a `dyn` trait (real + breaking change, hot-path dyn dispatch — heavy for an example-only feature) + OR a type-lie (accept `&ModelWeights`, ignore it, route to an owned + `BitnetModel`) — rejected as exactly the parallel-path anti-pattern the + policy box forbids. The cache-only sub-part (`BitnetKvCache` → + `larql-kv::KvCache`) is shape-compatible but marginal (the shared cache + doesn't append-in-place for this path either) and carries hot-path parity + risk, and it does NOT reach "first-class". → Decision: take the breaking + trait generalisation, or leave BitNet isolated-but-explicit (recommended + until a second consumer exists). +- **G3 — vindex quant-tag unification — GOAL FALSIFIED; struck.** Scoping + found BitNet vindexes are **mixed**: the dense scaffold (`embed`, `lm_head`, + `output_norm`) is f32 and loaded by the standard loader with + `skip_attn`/`skip_ffn`, while only attn/ffn are ternary. `quant: + QuantFormat::None` is therefore *correct* for the dense loader — setting + `quant: I2S` would mislead it into decoding the embedding as ternary. A + single `quant` tag cannot represent a mixed model; the two-field design + (`quant` for the dense scaffold + `bitnet_layout` manifest for the ternary + tensors) is the right shape. The only survivor is a modest mechanical + cleanup — move `bitnet_writer` from a `convert_cmd` read-modify-write + post-patch into the build pipeline so `index.json` is written once — low + payoff, invasive through the shared build path. Not pursued. +- **G4 — AVX2 `_mm256_sign_epi8` twin — BLOCKED on an x86 build/test box.** + Design is clear (decode trit codes → a `{-1,0,+1}` int8 control, one + `_mm256_sign_epi8(x, control)`, widen-accumulate; bit-identical to scalar). + But this aarch64 machine can neither runtime-validate the SIMD NOR even + compile-check it (cross-`check` fails in the C-FFI build script — no + `x86_64-linux-gnu-gcc`). Committing unbuilt, unvalidated intrinsics violates + the parity discipline. → Defer to an x86 dev box / Linux CI runner, where + the scalar-vs-AVX2 bit-exact test (already the pattern for the NEON twin) + can gate it. x86_64 keeps the correct scalar-A8 path (~2.4×) meanwhile. + +### Productization plan (decision: PRODUCTIZE, 2026-06-20) + +Direction chosen: make BitNet a real served path, not a validated experiment. +Scoping fixed the magnitude — BitNet has **zero CLI/server hookup today** +(`load_bitnet_model` is called only from the example); the dense run path is +`layer_graph::generate_streaming` over the engine dispatch; `run_cmd::run()` +is a chain of early-return mode branches (experts / ffn / moe / image). Three +stages, smallest-blast first: + +- **P-A — Serve BitNet from `larql run` (CLI) — ✅ BEHAVIOUR-VERIFIED + 2026-06-20.** `run_cmd::run()` branches on `config.bitnet_layout.is_some()` + and drives `ternary::generate_streaming_bitnet` (greedy stream + chat REPL), + bypassing the dense `walk_cmd` path. Smoke-tested against + `~/larql-vindex/bitnet-2b.vindex`: + `larql run "The capital of France is" -n 16` → + `" Paris. Paris is a city that is known for its rich history, culture,"` — + deterministic across runs. **This greedy output is the P-B regression + oracle** (saved local-only at `bench/oracles/bitnet_2b_capital_of_france.txt`; + not committed — depends on the >1 GB vindex; repro = the command above). + Bridges at the run layer; does NOT make BitNet a `KvEngine`. + *Remaining (deferred to AFTER P-B, deliberately): server stream-route wiring + + chat-template/sampling parity — wiring the server now would thread + `&ModelWeights` through the hot path B1 is about to strip; wire once, after.* +- **P-B — First-class `KvEngine` (the structural refactor).** Blast radius + measured: **8 production engine impls + ~171 `prefill`/`decode_step` call + sites + `EngineKind`/`AnyEngine`**. The one-way-door is the trait shape; + pick before the breaking change: + - **B1 (CHOSEN 2026-06-20): engines own their weights.** Move `&ModelWeights` + out of `prefill`/`decode_step` into engine construction (engines hold + `Arc`); `BitnetEngine` holds `Arc`. + **Read-only check (done):** dense `prefill`/`decode_step` and the + `*_resident` path take `&ModelWeights` (read-only — B1-clean); BitNet + weights are final (no mutation). BUT the **quant-resident path + (`prefill_quant`/`decode_step_quant`) takes `&mut ModelWeights`** — it + memoizes resident-quant buffers back into the struct. So B1 is NOT pure + mechanical churn; it bundles one design sub-decision for that path: + **(a)** relocate the resident-quant memoization out of `ModelWeights` into + engine-owned derivative state (recommended — lands on the StatePolicy + split: canonical weights immutable, derived caches are engine state), or + **(b)** `Arc` + interior mutability (`OnceCell`/`RwLock`) on + just the resident-quant fields (smaller, keeps derived state in the + canonical struct). Cost = ~171 mechanical sites + this sub-decision. + - **B2: `&dyn ModelSource` param.** New trait; `&ModelWeights` auto-coerces + so most call sites are untouched, but the trait must mirror the slice of + `ModelWeights` engines use, and BitNet panics on dense-only methods. + - **B3: `ModelWeights` gains a ternary representation.** Smallest type diff + but leaks ternary-awareness into the dense engines — rejected. + Do P-B as its own PR after P-A proves the path; hold parity per the 7-spec + `resident_identity_tests` discipline. + + **Grounded execution stages (B1a chosen, real-code scope 2026-06-20).** The + `&mut` has a single chokepoint: + `larql_inference::vindex::dequant::ensure_attn_tensors_dequantised(&mut weights, index)` + (`vindex/dequant.rs:35`) — it dequantises Q4K Q/K/V/O into `weights.tensors` + (a `HashMap`) keyed by `arch.attn_{q,k,v,o}_key(layer)`, idempotent, and the + forward reads them back from that map. Pure derivative state. Stages, each + compilable + checked against the captured greedy oracle: + - **P-B.1 — relocate the dequant cache (HOME LOCKED: engine, not + `ModelWeights`; concurrency evidence 2026-06-20).** Move the dequantised- + attention `HashMap` out of `weights.tensors` into **engine-owned** state and + consult it at the forward's tensor-read sites (resolver: engine cache → + canonical weights). Drops the `&mut` from `prefill_quant`/`decode_step_quant`. + *Why engine, not an interior-mutable `RwLock` field on `ModelWeights`:* the + scratch is **transient** (per-layer evicted for the memory bound) → per- + forward state, not a persistent cache. The server holds one + `weights: OnceLock>` and **serializes every generation + behind an exclusive write lock** (`state.rs:186 lock_weights_for_gen`, + used by all OpenAI gen routes) *specifically because* this dequant mutation + makes weights non-immutable ("concurrent reads block while a generation is + in flight"). An interior-mut `RwLock` field can't lift that — two forwards + sharing one `Arc` would clobber each other's evicting scratch, so gen would + still have to serialize (and the dense 117 tok/s path would pay a per-resolve + read-lock + `ArcArray` clone for a scratch that's always empty for it). + Engine-owned scratch makes `ModelWeights` **truly immutable** → + `Arc` shared across **concurrent** generations, each engine + its own cache (no lock, no race, no tax) — the actual payoff of the refactor. + Resolver threads as `&mut self.dequant_cache` from the engine (it's already + the `&mut self` forward context). Touches the Q4K residency path — + `resident_identity_tests` + the oracle guard it. *(A provisional + `RwLock`-in-`ModelWeights` impl was tried this session and reverted on this + evidence before the read-site/trait sweep could cement it.)* + + **P-B.1 status (2026-06-20): signature stages DONE+committed, relocation + set up + reverted-to-green.** Done behavior-identical: `WeightsView`/ + `DequantScratch` foundation; **Stage 1** (`run_attention_with_kv_backend` → + `WeightsView`, ~22 `dense()` wraps); **Stage 2a** (`dense_ffn_forward` → + `WeightsView`, `WeightFfn`/`BackendFfn` wrap `dense()` internally so the 326 + `WeightFfn` construction sites stay untouched). The workspace-spanning + cross-crate signature diff is banked, decoupled from any behavior change, + each proven byte-identical by parity tests. + + **Stage 2b (the relocation, behavior-changing) was reverted, and the reason + is categorical, not cost.** The first four blast-radius escalations + (RwLock→engine, cross-crate, 326-`WeightFfn`, decode reader) were all + *compiler-visible*: change a signature, the compiler enumerates callers. The + fifth is *type-system-invisible* — a reader that resolves `weights.tensors` + via `Deref` (canonical) while the scratch sits in an unconsulted + `DequantScratch` compiles clean, runs, and is wrong **only on the decode + path under a real Q4K vindex**. Holding that on a red tree across a session + boundary strands a miscompilation `cargo check` can't recover, so revert to + Stage 2a green was the only correctness-preserving move. + + **Silent-break closure = make the miss LOUD, not enumerate readers.** The + grep inventory (`tensors.get(&arch.attn_*/ffn_*`) is *current, not complete* + — blind to precomputed-key reads, prefix iteration, accessor methods that + `.get` internally. The design fix: for a quant model those dequant keys were + **never** in canonical `tensors` (they only ever existed as the forward-time + mutation target being relocated), so if the relocation inserts only into + scratch and leaves canonical untouched, a missed reader resolves `None` → + the existing `.unwrap_or_else(panic)` / `?`-bail fires on first decode, on + any vindex. **Design property to enforce: leave canonical genuinely empty of + dequant keys (not shadowed)** → misses are loud by construction. Grep scopes + the conversion; the runtime catches its misses. + + **Stage 2b entry conditions (all met — no upstream gap this time):** + (1) a Q4K vindex — **`~/larql-vindex/qwen3-0.6b-q4k.vindex` exists**; + (2) a **multi-token DECODE** oracle captured at Stage 2a (NOT prefill-only / + single-token — the decode reader is exactly the one Stage 1 missed, so a + prefill-heavy capture has a blind spot the shape of the bug); byte-identical + decode vs Stage 2a is the regression spine; (3) the canonical-empty shaping + above. With these, the reader conversion is mechanical and the silent-break + class is closed by construction. + + **Stage 2b progress + the reader-family finding (2026-06-20).** Done + + committed behavior-identical: all THREE "primary" quant-path readers now + take `WeightsView` — `run_attention_with_kv_backend` (Stage 1), + `dense_ffn_forward` (Stage 2a), `run_attention_block_decode_step_backend` + (Stage 2b-pre). The Q4K **decode** oracle is captured + (`bench/oracles/q4k_qwen3_history_of_computing.txt`, 24-token greedy on + `qwen3-0.6b-q4k.vindex`). The relocation proper (inserters→scratch, + `ViewFfn`, wire the cached prefill+decode loops, drop `&mut`) was drafted + and reverted to green when the **secondary** loops (`hidden.rs`, + `interventions.rs`) surfaced that the reader set is *still* expanding on + contact: they reach attention through `run_layer_with_ffn` → + `run_attention_inner` / `run_attention_with_kv_cache` → + `run_attention_block_core` (block.rs) + `run_attention_block_gpu` (gpu.rs) + — **un-converted readers the grep never surfaced**, exactly the + "current-not-complete" inventory. So the true relocation scope is "convert + the **whole attention-reader family**" (with_kv_backend✓ / decode✓ / + block_core / block_gpu / inner / with_kv_cache), each a Stage-1-style + cascade through a widely-used fn — several more passes, not one. The cached + decode path (the oracle path) wired cleanly; the secondary loops need the + rest of the family first. **Loud-break makes this safe to do incrementally** + (canonical empty of dequant keys → a missed reader gets `None` → the + existing `.unwrap_or_else(panic)` fires on first decode, loud not silent), + so each remaining reader can be converted + the loop wired + validated + against the oracle without a silent miscompilation risk. + + **✅ DONE (2026-06-20, `9650582e` + `f0da87cc`).** The whole-family + conversion + the relocation both landed. (1) `9650582e` converted the + entire attention-reader family to `WeightsView` — `block_core`, `block_gpu`, + `run_attention_inner`, `run_attention_with_kv_cache`, `run_layer_with_ffn`, + `run_layer_with_capture[_hooked]`, `run_attention_public` + the block.rs + family — ~100 callers `dense()`-wrapped across compute/inference/kv/cli/ + server/examples, behavior-identical (the compiler enumerated the family for + me; the cascade bottoming out *is* the proof the inventory is now complete). + (2) `f0da87cc` did the relocation: the production decode path + (`predict_kquant_prefill/decode_step` + `hidden` + `interventions`) + dequantises into a forward-local `DequantScratch` resolved via + `WeightsView::with_scratch` + `ViewFfn` — **`weights` is `&ModelWeights` + (immutable, Arc-able) on the decode path, `&mut` dropped.** Bulk + f32-fallback + dev drivers (KvEngine `*_quant` trait defaults, all larql-kv + quant-engine overrides, apollo, ov_rd CLI, the lql relation resolver, the + vision/image CLI, examples) keep in-`weights` behaviour via `*_resident` + shims (dequant → scratch → merge into `weights.tensors`). Validated: + workspace `--all-targets` green, clippy 0 warnings, 50 kquant + 13 dequant + + resident_identity tests pass, decode **byte-identical to the oracle** both + after the family conversion and after the relocation. **Follow-up:** the + `*_resident` bulk path is still `&mut` — dropping it needs engine-owned + scratch state (folds into P-B.2/P-B.3, not a blocker); loud-break guards it. + + **P-B.1b — "no shims" full sweep (scoped 2026-06-20; WIP stashed).** Going + for zero `weights.tensors.extend` shims surfaced a **second `kquant_forward` + implementation**: the production `larql run` decode dispatches via KvEngines + → `coarse_prefill` → **larql-compute's** `kquant_forward` (1005 lines), NOT + the larql-inference copy (1772 lines) that P-B.1 relocated. The + larql-inference copy serves the direct-`predict_kquant`/AVE/hidden paths and + is validated by the 50 kquant unit tests; **the e2e oracle actually + exercises larql-compute's copy** (so the family conversion — shared + `run_attention_*` — was oracle-validated, but the larql-inference relocation + was unit-test-validated, not oracle). The full no-shim change is large and + interconnected: + (1) relocate **larql-compute's** `kquant_forward` too (cached/decode loops → + forward-local scratch + `ViewFfn`) — DONE in the stash, the real oracle + path now no-shim; + (2) `KvDispatch` (5 methods) + the 7 dispatch helpers + `AsyncComputeBackend` + + cpu/metal impls → `WeightsView` — DONE in the stash; + (3) coarse path (`coarse_prefill`/`coarse_decode_step`) drops `&mut` (delegates + to the now-`&` `predict_kquant_*`) — DONE; + (4) `KvEngine`/`RetrievalEngine` trait quant methods → `&ModelWeights`; + `RetrievalEngine::prefill_quant` default → loud error (apollo overrides; + the `ffn`-less `prefill` can't thread a scratch) — DONE; + (5) **engine-scratch design** (validated on `StandardEngine`): each engine + owns a `dequant_scratch: DequantScratch`; `do_prefill`/`do_decode_step` + build `WeightsView::with_scratch(weights, &self.dequant_scratch)` — the + view borrows `self.dequant_scratch` while `self.handles`/`self.backend` + are borrowed disjointly, so **no take/restore dance**; `prefill_quant` + dequants into the field, no merge. StandardEngine compiles clean. + Remaining (the stash is mid-sweep, ~29 errors): the **other 7 engines** + (no_cache, boundary×2, markov×2, turbo, unlimited, apollo) each need the + same field + view-thread + `&mut`-drop, plus their forward helpers + (`kv_prefill_run` + the `generate_cached_*` loops in larql-kv/generation.rs, + apollo's `forward_raw_logits`) converted to `WeightsView`, then the dev + drivers (ov_rd/lql/vision/examples) + delete the `*_resident` shims. The + pattern is mechanical but keeps surfacing forward helpers on contact (the + reader-family expansion, now in the engine layer) — a focused dedicated pass. + `git stash list` → "no-shims WIP". + + **Convergence measurement (2026-06-21).** Resumed the sweep and pushed into + the engines. Additional shared decode/recompute readers converted to + `WeightsView` (these are real foundation, beyond the StandardEngine work): + `run_attention_block_decode_step_auto` + `_auto_inplace` (the resident-decode + switcher used by **5 engines** — q4k-direct branch reads native index bytes + via `.canonical()`, f32 branch threads the view), `kv_prefill_run` + (no_cache + standard), `recompute_kv` + `attn_kv_projection_weights` + (boundary + markov), and **NoCacheEngine** fully (field + view-thread + + `&mut` drop, clean). **But the larql-kv error count DIVERGED as I converted: + 28 → 45 → 67.** Each engine's `walk.rs`/`compute.rs` forward module is a deep + chain (engine → walk → recompute → projection → `weights.tensors`), and + converting one helper exposes its callers + their internal reads. This is the + reader-family expansion at its widest — **converting every engine's full + forward/recompute internals** (~30-50+ functions across 6 engine modules + + apollo's `forward_raw_logits` + the dev drivers). The diverging count is the + decision signal: this is a **staged multi-session refactor**, best done one + engine module at a time (convert its walk+compute internals, validate that + engine against a per-engine oracle, commit), not a single grind. The shared + helpers above are the foundation already laid; the per-engine internals are + the remaining bulk. WIP re-stashed. + + **✅ DONE (2026-06-21, `379885ed`).** The diverging count (28→45→67) + **converged to 0** as each engine module got the same template — the + "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 + 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 + `walk.rs`/`compute.rs`/`executor.rs`/`cold_tier.rs`/`dispatch.rs` to + `WeightsView` (canonical reads — `embed`/`run_ffn`/`layer_ffn_or_moe`/ + `BackendFfn`/`WalkFfn`/native-q4k — via `.canonical()`/`&weights`; attn reads + via the view), then the engine adds the field + threads `with_scratch` to its + forward calls + drops `&mut`. Also converted: `LayerExecutor` trait + + local_walk, `recompute_kv` + `attn_kv_projection_weights` (explicit lifetime), + `auto`/`auto_inplace` (5 engines), `kv_prefill_run`, `forward_raw_logits`/ + `forward_from_layer` + raw.rs internals (`ViewFfn`; `hidden_to_raw_logits`/ + `apply_logits_transform` stay `&ModelWeights` = lm_head canonical). The + `*_resident` helpers (`ensure_attn_..._resident` etc.) deliberately **remain** + for ~58 dev/research call sites (ov_rd CLI, lql resolver, vision CLI, + examples) that own a `&mut ModelWeights` and run one-off forwards against + canonical `weights.tensors` — a documented separate API, not serving-path + shims. Validated: workspace `--all-targets` green, clippy 0, **766 larql-kv + + 40 kquant + resident_identity + 4 dispatch_parity (cross-engine bit-parity)** + tests pass, decode **byte-identical to the oracle**, and + markov-rs/unlimited/turbo-quant/no-cache smoke-tested coherent at runtime. + **Engine/serving path is now fully `&ModelWeights` — P-B.2 (Arc-owned) is + unblocked with no remaining `&mut` to chase in the engines.** + - **P-B.2 — Arc-owned weights.** Every weight param is now `&ModelWeights`; + move it into engine construction (engines hold `Arc`) and drop + the param from prefill/decode/quant/resident/executor variants. ~171 call + sites (all production, 0 in test files) + `EngineKind`/`AnyEngine`. + Compiler-driven — the safe kind of large churn. + - **P-B.3 — `BitnetEngine` + dispatch.** New engine holding + `Arc`, `impl KvEngine` over the ternary forward; add the + `EngineKind`/`AnyEngine` arm so unified dispatch picks ternary vs dense. + - **P-B.4 — validate** against the oracle (greedy "Paris…" byte-identical) + + the engine parity suite. + Best run in an isolated worktree so `main` stays stable through the change. +- **P-C — G4 AVX2 + x86 CI.** Add a Linux-x86 CI job; land the AVX2 twin gated + by the scalar-vs-AVX2 bit-exact test (the NEON-twin pattern). Independent of + P-A/P-B; unblocks G4's environment blocker. + +--- + ## Demo narrative ### Act 1 — "The model is the database" diff --git a/bench/baselines/cpu/COMPARISON.md b/bench/baselines/cpu/COMPARISON.md index 2fe9a7642..b63197d3c 100644 --- a/bench/baselines/cpu/COMPARISON.md +++ b/bench/baselines/cpu/COMPARISON.md @@ -1,5 +1,27 @@ # larql vs llama.cpp — CPU decode on Gemma 3 4B Q4_K +> **Update 2026-06-22 — prefill gap largely closed.** The q4k-direct prefill +> work changed the picture: Q4_K/Q6_K attention (Q/K/V/O) and FFN (gate/up/down) +> projections now run straight from the vindex bytes with no per-layer f32 +> dequant — `q4k_matmul`/`q6k_matmul` (the Q6_K twin, used by the default Q6_K +> `down_proj` and `v_proj`), with a hand-written aarch64 NEON inner dot. +> Apple M3 Max, CPU only (`-t 8`), same model + prompt as below. +> +> | Metric | larql (standard) | llama.cpp | Ratio | +> |---|---:|---:|---:| +> | Decode (tg, tok/s) | ~42 | ~38 | **~1.1× ahead** | +> | Prefill (5-tok prompt, ms) | 233 | ~70 | **~3.3× behind** (was 55×) | +> | Prefill vs the May full-dequant path | 2746 → 233 ms | | **11.8× faster** | +> +> Decode is now at/ahead of llama.cpp; prefill went from 55× behind to ~3×. The +> NEON `q4k_matmul` at seq=5 actually *beats* f32 AMX sgemm (1.0–1.3×) while +> skipping the dequant. The remaining prefill gap is constant-factor kernel work +> (our matmul vs llama.cpp's hand-tuned asm) plus batched attention, not dequant. +> Numbers are same-session (machine warm from builds) — ratios hold; cold +> absolutes run a touch faster. The 2026-05-15 baseline below is kept for history. + +--- + Recorded 2026-05-15 on Apple M3 Max, 12 threads, BLAS / Accelerate enabled, no GPU. Both engines load the same model weights — `output/larql-gemma-3-4b-it.gguf` quantized to Q4_K_M for llama.cpp, the matching `output/gemma3-4b-q4k-v2.vindex` diff --git a/crates/larql-cli/src/commands/dev/ov_rd/basis.rs b/crates/larql-cli/src/commands/dev/ov_rd/basis.rs index 42e5740e2..b501581a9 100644 --- a/crates/larql-cli/src/commands/dev/ov_rd/basis.rs +++ b/crates/larql-cli/src/commands/dev/ov_rd/basis.rs @@ -259,8 +259,12 @@ pub(super) fn fit_z_pca_bases( for layer in 0..weights.num_layers { let inserted = insert_q4k_layer_tensors(weights, index, layer)?; if let Some(layer_heads) = heads_by_layer.get(&layer) { - let (_, pre_o) = run_attention_block_with_pre_o(weights, &h, layer) - .ok_or_else(|| format!("pre-W_O capture failed at layer {layer}"))?; + let (_, pre_o) = run_attention_block_with_pre_o( + larql_models::WeightsView::dense(weights), + &h, + layer, + ) + .ok_or_else(|| format!("pre-W_O capture failed at layer {layer}"))?; let head_dim = weights.arch.head_dim_for_layer(layer); for head in layer_heads { let basis = bases.get(head).expect("basis pre-created for PCA fit"); @@ -287,9 +291,15 @@ pub(super) fn fit_z_pca_bases( { let ffn = WeightFfn { weights }; - if let Some((h_new, _, _)) = - run_layer_with_ffn(weights, &h, layer, &ffn, false, ple_inputs.get(layer), None) - { + if let Some((h_new, _, _)) = run_layer_with_ffn( + larql_inference::WeightsView::dense(weights), + &h, + layer, + &ffn, + false, + ple_inputs.get(layer), + None, + ) { h = h_new; } } diff --git a/crates/larql-cli/src/commands/dev/ov_rd/capture.rs b/crates/larql-cli/src/commands/dev/ov_rd/capture.rs index 168c974dd..cd16be968 100644 --- a/crates/larql-cli/src/commands/dev/ov_rd/capture.rs +++ b/crates/larql-cli/src/commands/dev/ov_rd/capture.rs @@ -132,8 +132,12 @@ pub(super) fn run_capture(args: CaptureArgs) -> Result<(), Box Result<(), Box Result<(), Box Result<(), Box Result, Box> { - let (_, pre_o) = run_attention_block_with_pre_o(weights, &h, layer) - .ok_or_else(|| format!("pre-W_O capture failed at layer {layer}"))?; + let (_, pre_o) = run_attention_block_with_pre_o( + larql_models::WeightsView::dense(weights), + &h, + layer, + ) + .ok_or_else(|| format!("pre-W_O capture failed at layer {layer}"))?; let head_dim = weights.arch.head_dim_for_layer(layer); let start = head.head * head_dim; let end = start + head_dim; @@ -962,9 +976,15 @@ fn capture_head_position_sq_errors( } { let ffn = WeightFfn { weights }; - if let Some((h_new, _, _)) = - run_layer_with_ffn(weights, &h, layer, &ffn, false, ple_inputs.get(layer), None) - { + if let Some((h_new, _, _)) = run_layer_with_ffn( + larql_inference::WeightsView::dense(weights), + &h, + layer, + &ffn, + false, + ple_inputs.get(layer), + None, + ) { h = h_new; } } diff --git a/crates/larql-cli/src/commands/dev/ov_rd/probe_program_class.rs b/crates/larql-cli/src/commands/dev/ov_rd/probe_program_class.rs index b98cb7e0c..88057f54f 100644 --- a/crates/larql-cli/src/commands/dev/ov_rd/probe_program_class.rs +++ b/crates/larql-cli/src/commands/dev/ov_rd/probe_program_class.rs @@ -641,7 +641,10 @@ fn capture_target_features( .kv_shared_source_layer(layer) .and_then(|src| kv_cache.get(&src)); let (_, pre_o, all_weights) = run_attention_block_with_pre_o_and_all_attention_weights( - weights, &h, layer, shared_kv, + larql_models::WeightsView::dense(weights), + &h, + layer, + shared_kv, ) .ok_or_else(|| { format!( @@ -683,7 +686,7 @@ fn capture_target_features( .and_then(|src| kv_cache.get(&src)); let ffn = WeightFfn { weights }; run_layer_with_ffn( - weights, + larql_inference::WeightsView::dense(weights), &h, layer, &ffn, diff --git a/crates/larql-cli/src/commands/dev/ov_rd/runtime.rs b/crates/larql-cli/src/commands/dev/ov_rd/runtime.rs index a93463689..71a1de10d 100644 --- a/crates/larql-cli/src/commands/dev/ov_rd/runtime.rs +++ b/crates/larql-cli/src/commands/dev/ov_rd/runtime.rs @@ -6,11 +6,13 @@ pub(super) fn insert_q4k_layer_tensors( index: &VectorIndex, layer: usize, ) -> Result, Box> { - larql_inference::vindex::insert_q4k_layer_tensors(weights, index, layer).map_err(|err| { - Box::::from(std::io::Error::new(std::io::ErrorKind::Other, err)) - }) + larql_inference::vindex::insert_q4k_layer_tensors_resident(weights, index, layer).map_err( + |err| { + Box::::from(std::io::Error::new(std::io::ErrorKind::Other, err)) + }, + ) } pub(super) fn remove_layer_tensors(weights: &mut ModelWeights, keys: Vec) { - larql_inference::vindex::remove_layer_tensors(weights, keys); + larql_inference::vindex::remove_layer_tensors_resident(weights, keys); } diff --git a/crates/larql-cli/src/commands/dev/ov_rd/sanity.rs b/crates/larql-cli/src/commands/dev/ov_rd/sanity.rs index b12c8fc63..b0c4058ed 100644 --- a/crates/larql-cli/src/commands/dev/ov_rd/sanity.rs +++ b/crates/larql-cli/src/commands/dev/ov_rd/sanity.rs @@ -142,7 +142,7 @@ pub(super) fn run_sanity_check(args: SanityCheckArgs) -> Result<(), Box Result<(), Box Result<(), Box> { // 7. Full attention (end-to-end via run_attention_public) let start = Instant::now(); for _ in 0..iters { - let _ = larql_inference::forward::run_attention_public(weights, &x, layer); + let _ = larql_inference::forward::run_attention_public( + larql_inference::WeightsView::dense(weights), + &x, + layer, + ); } let full_attn_us = start.elapsed().as_micros() as f64 / iters as f64; diff --git a/crates/larql-cli/src/commands/primary/bench/local_moe_runtime.rs b/crates/larql-cli/src/commands/primary/bench/local_moe_runtime.rs index 9c9e50908..8ad8b1320 100644 --- a/crates/larql-cli/src/commands/primary/bench/local_moe_runtime.rs +++ b/crates/larql-cli/src/commands/primary/bench/local_moe_runtime.rs @@ -48,7 +48,7 @@ pub(super) fn run_local_moe( // for the whole generation (experts stay Q4K — read directly by // `build_moe_weights`). Matches the `larql run --moe-shards` CPU default. for layer in 0..weights.num_layers { - larql_inference::vindex::insert_q4k_layer_tensors(weights, index, layer) + larql_inference::vindex::insert_q4k_layer_tensors_resident(weights, index, layer) .map_err(|e| format!("failed to dequantize layer {layer} to f32: {e}"))?; } diff --git a/crates/larql-cli/src/commands/primary/bench/run.rs b/crates/larql-cli/src/commands/primary/bench/run.rs index 502023966..43cd6b2a3 100644 --- a/crates/larql-cli/src/commands/primary/bench/run.rs +++ b/crates/larql-cli/src/commands/primary/bench/run.rs @@ -219,8 +219,10 @@ pub fn run(mut args: BenchArgs) -> Result<(), Box> { the remote block; dense engine row skipped" ); } else { - let kv_ref_bytes = - larql_kv::markov_residual::kv_memory_bytes_for_seq(&weights, token_ids.len()); + let kv_ref_bytes = larql_kv::markov_residual::kv_memory_bytes_for_seq( + larql_inference::WeightsView::dense(&weights), + token_ids.len(), + ); // Parse + validate --ffn-policy once before the engine loop // (multi-engine sweep reuses the same validated policy). @@ -266,8 +268,10 @@ pub fn run(mut args: BenchArgs) -> Result<(), Box> { let token_ids = larql_inference::encode_prompt(&tokenizer, &*weights.arch, args.prompt.as_str()) .map_err(|e| format!("tokenize: {e}"))?; - let kv_ref_bytes = - larql_kv::markov_residual::kv_memory_bytes_for_seq(&weights, token_ids.len()); + let kv_ref_bytes = larql_kv::markov_residual::kv_memory_bytes_for_seq( + larql_inference::WeightsView::dense(&weights), + token_ids.len(), + ); let validated_policy = parse_ffn_policy(args.ffn_policy.as_deref(), weights.num_layers)?; diff --git a/crates/larql-cli/src/commands/primary/run_cmd.rs b/crates/larql-cli/src/commands/primary/run_cmd.rs index b843d0e2e..b7f8a1270 100644 --- a/crates/larql-cli/src/commands/primary/run_cmd.rs +++ b/crates/larql-cli/src/commands/primary/run_cmd.rs @@ -297,6 +297,17 @@ pub fn run(args: RunArgs) -> Result<(), Box> { return experts::run(&vindex_path, &args); } + // BitNet b1.58 native-ternary vindex: served by the ternary forward + // (`larql_inference::ternary`), not the dense engine dispatch — detect + // early and route, bypassing the walk_cmd path the dense run delegates to. + // A failed config load falls through so the dense path surfaces the error. + if larql_vindex::load_vindex_config(&vindex_path) + .map(|c| c.bitnet_layout.is_some()) + .unwrap_or(false) + { + return run_bitnet(&vindex_path, &args); + } + if let Some(ref ffn_url) = args.ffn { let prompt = args.prompt.as_deref().ok_or( "--ffn requires a prompt argument (chat mode not yet supported with --ffn-dispatch batch)", @@ -420,6 +431,91 @@ fn run_chat( } } +/// Serve a BitNet b1.58 native-ternary vindex. +/// +/// BitNet uses a separate ternary forward (`larql_inference::ternary`) rather +/// than the dense engine dispatch, so it bypasses the `walk_cmd` path that the +/// dense `run_once` / `run_chat` delegate to. Greedy streaming generation; +/// no prompt drops into a simple stdin REPL. EOS is sourced the same way as +/// the dense path (`EosConfig::from_vindex_dir`). +/// +/// MVP scope (P-A): raw prompt encode (no chat-template rendering yet) and +/// greedy sampling. Chat-template + sampling-flag parity with the dense path +/// are follow-ups; see ROADMAP "Productization plan" P-A. +fn run_bitnet( + vindex_path: &std::path::Path, + args: &RunArgs, +) -> Result<(), Box> { + use larql_inference::layer_graph::generate::{eos::EosConfig, SamplingConfig}; + use larql_inference::ternary; + + let model = ternary::load_bitnet_model(vindex_path)?; + let tokenizer = + larql_vindex::tokenizers::Tokenizer::from_file(vindex_path.join("tokenizer.json")) + .map_err(|e| format!("load tokenizer.json: {e}"))?; + let eos = EosConfig::from_vindex_dir(vindex_path); + let max_tokens = args.max_tokens; + let verbose = args.verbose; + + let generate_one = |prompt: &str| -> Result<(), Box> { + let enc = tokenizer + .encode(prompt, true) + .map_err(|e| format!("encode prompt: {e}"))?; + let mut stdout = io::stdout(); + let emitted = ternary::generate_streaming_bitnet( + &model, + &tokenizer, + enc.get_ids(), + max_tokens, + SamplingConfig::greedy(), + &eos, + |_id, delta, _ms| { + print!("{delta}"); + let _ = stdout.flush(); + }, + ); + println!(); + if verbose { + eprintln!("[bitnet] {emitted} tokens"); + } + Ok(()) + }; + + if let Some(prompt) = args.prompt.as_deref() { + return generate_one(prompt); + } + + // Chat REPL (single-turn, no multi-turn template — BitNet MVP). + eprintln!( + "larql chat (bitnet) — {} (Ctrl-D to exit)", + vindex_path + .file_name() + .and_then(|n| n.to_str()) + .unwrap_or("model") + ); + let stdin = io::stdin(); + loop { + eprint!("> "); + io::stderr().flush()?; + let mut line = String::new(); + match stdin.lock().read_line(&mut line) { + Ok(0) => { + eprintln!(); + return Ok(()); + } + Ok(_) => {} + Err(e) => return Err(Box::new(e)), + } + let prompt = line.trim(); + if prompt.is_empty() { + continue; + } + if let Err(e) = generate_one(prompt) { + eprintln!("Error: {e}"); + } + } +} + /// Build a `WalkArgs` with sensible defaults from the slim `RunArgs`. The /// fields we don't surface to end users get stable defaults here. fn build_walk_args( @@ -688,8 +784,12 @@ fn run_with_moe_shards( // Dequantize attn + dense FFN to f32 for every layer, kept resident // for the whole generation (experts are not loaded on the client). for layer in 0..weights.num_layers { - larql_inference::vindex::insert_q4k_layer_tensors(&mut weights, &index, layer) - .map_err(|e| format!("failed to dequantize layer {layer} to f32: {e}"))?; + larql_inference::vindex::insert_q4k_layer_tensors_resident( + &mut weights, + &index, + layer, + ) + .map_err(|e| format!("failed to dequantize layer {layer} to f32: {e}"))?; } let moe_ffn = larql_inference::ffn::RemoteMoeFfn { weights: &weights, diff --git a/crates/larql-cli/src/commands/primary/run_cmd_image.rs b/crates/larql-cli/src/commands/primary/run_cmd_image.rs index aaf823ecd..6a4b885e0 100644 --- a/crates/larql-cli/src/commands/primary/run_cmd_image.rs +++ b/crates/larql-cli/src/commands/primary/run_cmd_image.rs @@ -250,7 +250,7 @@ pub fn run_with_images( // Materialise f32 attention tensors into `weights.tensors` so the // engine's attention dispatch reads f32 even though the on-disk // tensors are Q4K. - larql_inference::vindex::ensure_attn_tensors_dequantised(&mut weights, &idx); + larql_inference::vindex::ensure_attn_tensors_dequantised_resident(&mut weights, &idx); Some(idx) } else { None diff --git a/crates/larql-cli/src/commands/primary/shannon_cmd.rs b/crates/larql-cli/src/commands/primary/shannon_cmd.rs index e7634bf7b..5d7b485cf 100644 --- a/crates/larql-cli/src/commands/primary/shannon_cmd.rs +++ b/crates/larql-cli/src/commands/primary/shannon_cmd.rs @@ -1431,7 +1431,7 @@ fn forward_hidden( .kv_shared_source_layer(layer) .and_then(|src| kv_cache.get(&src)); if let Some((h_new, _, kv_out)) = larql_inference::forward::run_layer_with_ffn( - weights, + larql_inference::WeightsView::dense(weights), &h, layer, &ffn, @@ -1469,7 +1469,7 @@ fn forward_hidden_all_layers( .kv_shared_source_layer(layer) .and_then(|src| kv_cache.get(&src)); if let Some((h_new, _, kv_out)) = larql_inference::forward::run_layer_with_ffn( - weights, + larql_inference::WeightsView::dense(weights), &h, layer, &ffn, diff --git a/crates/larql-compute-metal/src/async_compute_backend_impl.rs b/crates/larql-compute-metal/src/async_compute_backend_impl.rs index 5e180da80..13aa70f66 100644 --- a/crates/larql-compute-metal/src/async_compute_backend_impl.rs +++ b/crates/larql-compute-metal/src/async_compute_backend_impl.rs @@ -26,7 +26,6 @@ use larql_compute::async_compute_backend::{ use larql_compute::ffn::FfnBackend; use larql_compute::kv_dispatch::{KvHandle, ResidualHandle}; use larql_compute::CpuBackend; -use larql_models::ModelWeights; /// Convenience — the CPU backend instance every method delegates to. /// Zero-sized type; const-construction is free. @@ -35,7 +34,7 @@ const CPU: CpuBackend = CpuBackend; impl AsyncComputeBackend for MetalBackend { fn attention_step_async( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, query: &Array2, kv: &mut KvHandle, layer: usize, @@ -50,7 +49,7 @@ impl AsyncComputeBackend for MetalBackend { fn attention_step_windowed_async( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, query: &Array2, kv: &mut KvHandle, layer: usize, @@ -63,7 +62,7 @@ impl AsyncComputeBackend for MetalBackend { fn attention_prefill_async( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, tokens_embedded: &Array2, layer: usize, window: Option, @@ -85,7 +84,7 @@ impl AsyncComputeBackend for MetalBackend { fn forward_from_layer_async( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, ffn: &dyn FfnBackend, start_layer: usize, residuals: &ResidualHandle, @@ -127,11 +126,24 @@ mod tests { let m = backend(); let tokens = vec![0u32, 1, 2]; let h_in = larql_compute::forward::embed_tokens_pub(&weights, &tokens); - let (_h_prefill, mut kv) = m.attention_prefill_async(&weights, &h_in, 0, None, None); + let (_h_prefill, mut kv) = m.attention_prefill_async( + larql_models::WeightsView::dense(&weights), + &h_in, + 0, + None, + None, + ); let h_new = larql_compute::forward::embed_tokens_pub(&weights, &[3u32]); let abs_position = tokens.len(); let h_async = m - .attention_step_async(&weights, &h_new, &mut kv, 0, abs_position, None) + .attention_step_async( + larql_models::WeightsView::dense(&weights), + &h_new, + &mut kv, + 0, + abs_position, + None, + ) .read(); assert_eq!(h_async.ncols(), weights.hidden_size); assert_eq!(h_async.nrows(), 1); @@ -143,11 +155,17 @@ mod tests { let m = backend(); let tokens = vec![0u32, 1, 2]; let h_in = larql_compute::forward::embed_tokens_pub(&weights, &tokens); - let (_, mut kv) = m.attention_prefill_async(&weights, &h_in, 0, None, None); + let (_, mut kv) = m.attention_prefill_async( + larql_models::WeightsView::dense(&weights), + &h_in, + 0, + None, + None, + ); let h_new = larql_compute::forward::embed_tokens_pub(&weights, &[3u32]); let h_async = m .attention_step_windowed_async( - &weights, + larql_models::WeightsView::dense(&weights), &h_new, &mut kv, 0, @@ -165,7 +183,13 @@ mod tests { let m = backend(); let tokens = vec![0u32, 1, 2]; let h_in = larql_compute::forward::embed_tokens_pub(&weights, &tokens); - let (h_handle, kv) = m.attention_prefill_async(&weights, &h_in, 0, None, None); + let (h_handle, kv) = m.attention_prefill_async( + larql_models::WeightsView::dense(&weights), + &h_in, + 0, + None, + None, + ); let h = h_handle.read(); assert_eq!(h.nrows(), tokens.len()); assert_eq!(h.ncols(), weights.hidden_size); @@ -205,10 +229,22 @@ mod tests { let (_, residuals_m) = m.upload_boundary_residual_async(&residual); let (_, residuals_c) = cpu.upload_boundary_residual_async(&residual); let h_m = m - .forward_from_layer_async(&weights, &ffn, 1, &residuals_m, &tokens) + .forward_from_layer_async( + larql_models::WeightsView::dense(&weights), + &ffn, + 1, + &residuals_m, + &tokens, + ) .read(); let h_c = cpu - .forward_from_layer_async(&weights, &ffn, 1, &residuals_c, &tokens) + .forward_from_layer_async( + larql_models::WeightsView::dense(&weights), + &ffn, + 1, + &residuals_c, + &tokens, + ) .read(); assert_eq!(h_m, h_c, "Metal delegation must bit-match CpuBackend"); } diff --git a/crates/larql-compute-metal/src/kv_dispatch_impl.rs b/crates/larql-compute-metal/src/kv_dispatch_impl.rs index e3e57006f..97fc34fa1 100644 --- a/crates/larql-compute-metal/src/kv_dispatch_impl.rs +++ b/crates/larql-compute-metal/src/kv_dispatch_impl.rs @@ -50,7 +50,7 @@ impl KvDispatch for MetalBackend { fn attention_step( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, query: &Array2, kv: &mut KvHandle, layer: usize, @@ -64,7 +64,7 @@ impl KvDispatch for MetalBackend { fn attention_step_windowed( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, query: &Array2, kv: &mut KvHandle, layer: usize, @@ -77,7 +77,7 @@ impl KvDispatch for MetalBackend { fn attention_prefill( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, tokens_embedded: &Array2, layer: usize, window: Option, @@ -88,7 +88,7 @@ impl KvDispatch for MetalBackend { fn recompute_kv_from_residuals( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, residuals: &Array2, layer: usize, ) -> Option { @@ -113,7 +113,7 @@ impl KvDispatch for MetalBackend { fn forward_from_layer( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, start_layer: usize, residuals: &ResidualHandle, token_ids: &[u32], @@ -150,7 +150,7 @@ impl KvDispatch for MetalBackend { fn coarse_prefill( &self, - weights: &mut ModelWeights, + weights: &ModelWeights, token_ids: &[u32], index: Option<&dyn larql_compute::KvIndex>, ) -> Option<(Array2, KvHandle)> { @@ -161,7 +161,7 @@ impl KvDispatch for MetalBackend { fn coarse_prefill_with_state( &self, - weights: &mut ModelWeights, + weights: &ModelWeights, token_ids: &[u32], index: Option<&dyn larql_compute::KvIndex>, state: Option<&mut larql_compute::PerLayerDecodeState>, @@ -283,7 +283,7 @@ impl KvDispatch for MetalBackend { fn coarse_decode_step( &self, - weights: &mut ModelWeights, + weights: &ModelWeights, token_id: u32, index: Option<&dyn larql_compute::KvIndex>, _handle: &mut KvHandle, @@ -295,7 +295,7 @@ impl KvDispatch for MetalBackend { fn coarse_decode_step_with_state( &self, - weights: &mut ModelWeights, + weights: &ModelWeights, token_id: u32, index: Option<&dyn larql_compute::KvIndex>, handle: &mut KvHandle, @@ -315,7 +315,7 @@ impl KvDispatch for MetalBackend { fn coarse_decode_step_with_state_masked( &self, - weights: &mut ModelWeights, + weights: &ModelWeights, token_id: u32, index: Option<&dyn larql_compute::KvIndex>, _handle: &mut KvHandle, @@ -523,11 +523,24 @@ mod tests { let tokens = vec![0u32, 1, 2]; let h_in = larql_compute::forward::embed_tokens_pub(&weights, &tokens); let (_, mut kv) = m - .attention_prefill(&weights, &h_in, 0, None, None) + .attention_prefill( + larql_models::WeightsView::dense(&weights), + &h_in, + 0, + None, + None, + ) .expect("prefill"); let h_new = larql_compute::forward::embed_tokens_pub(&weights, &[3u32]); let h = m - .attention_step(&weights, &h_new, &mut kv, 0, tokens.len(), None) + .attention_step( + larql_models::WeightsView::dense(&weights), + &h_new, + &mut kv, + 0, + tokens.len(), + None, + ) .expect("attention_step"); assert_eq!(h.shape(), &[1, weights.hidden_size]); } @@ -539,11 +552,25 @@ mod tests { let tokens = vec![0u32, 1, 2]; let h_in = larql_compute::forward::embed_tokens_pub(&weights, &tokens); let (_, mut kv) = m - .attention_prefill(&weights, &h_in, 0, None, None) + .attention_prefill( + larql_models::WeightsView::dense(&weights), + &h_in, + 0, + None, + None, + ) .expect("prefill"); let h_new = larql_compute::forward::embed_tokens_pub(&weights, &[3u32]); let h = m - .attention_step_windowed(&weights, &h_new, &mut kv, 0, tokens.len(), 64, None) + .attention_step_windowed( + larql_models::WeightsView::dense(&weights), + &h_new, + &mut kv, + 0, + tokens.len(), + 64, + None, + ) .expect("windowed attention_step"); assert_eq!(h.shape(), &[1, weights.hidden_size]); } @@ -555,7 +582,13 @@ mod tests { let tokens = vec![0u32, 1, 2]; let h_in = larql_compute::forward::embed_tokens_pub(&weights, &tokens); let (h, kv) = m - .attention_prefill(&weights, &h_in, 0, None, None) + .attention_prefill( + larql_models::WeightsView::dense(&weights), + &h_in, + 0, + None, + None, + ) .expect("prefill"); assert_eq!(h.shape(), &[tokens.len(), weights.hidden_size]); assert_eq!(kv.cached_len(), tokens.len()); @@ -574,8 +607,16 @@ mod tests { let residuals = Array2::from_shape_vec((3, weights.hidden_size), vec![0.0; 3 * weights.hidden_size]) .unwrap(); - let m_result = m.recompute_kv_from_residuals(&weights, &residuals, 0); - let cpu_result = cpu.recompute_kv_from_residuals(&weights, &residuals, 0); + let m_result = m.recompute_kv_from_residuals( + larql_models::WeightsView::dense(&weights), + &residuals, + 0, + ); + let cpu_result = cpu.recompute_kv_from_residuals( + larql_models::WeightsView::dense(&weights), + &residuals, + 0, + ); assert_eq!( m_result.is_some(), cpu_result.is_some(), @@ -603,7 +644,12 @@ mod tests { .unwrap(); let handle = m.upload_boundary_residual(&residual).expect("upload"); let h = m - .forward_from_layer(&weights, 1, &handle, &[0u32, 1, 2]) + .forward_from_layer( + larql_models::WeightsView::dense(&weights), + 1, + &handle, + &[0u32, 1, 2], + ) .expect("forward_from_layer"); assert_eq!(h.ncols(), weights.hidden_size); } @@ -643,10 +689,10 @@ mod tests { #[test] fn coarse_decode_step_without_index_returns_none() { - let mut weights = make_test_weights(); + let weights = make_test_weights(); let m = backend(); let mut handle = KvHandle::new(MetalCoarseHandle); - let result = m.coarse_decode_step(&mut weights, 0u32, None, &mut handle, 0); + let result = m.coarse_decode_step(&weights, 0u32, None, &mut handle, 0); assert!(result.is_none()); } @@ -659,9 +705,9 @@ mod tests { use larql_compute::test_fixtures::make_q4k_fixture_index; use larql_models::test_fixtures::make_test_q4k_weights; let m = backend(); - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let idx = make_q4k_fixture_index(&weights); - let result = m.coarse_prefill(&mut weights, &[0u32, 1, 2], Some(&idx)); + let result = m.coarse_prefill(&weights, &[0u32, 1, 2], Some(&idx)); let (h, _handle) = result.expect("Metal Q4K prefill succeeds"); assert_eq!(h.shape(), &[1, weights.hidden_size]); } @@ -672,11 +718,11 @@ mod tests { use larql_compute::test_fixtures::make_q4k_fixture_index; use larql_models::test_fixtures::make_test_q4k_weights; let m = backend(); - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let idx = make_q4k_fixture_index(&weights); let mut state = larql_compute::PerLayerDecodeState::with_capacity(weights.num_layers); let result = - m.coarse_prefill_with_state(&mut weights, &[0u32, 1, 2], Some(&idx), Some(&mut state)); + m.coarse_prefill_with_state(&weights, &[0u32, 1, 2], Some(&idx), Some(&mut state)); let (h, _handle) = result.expect("Metal Q4K prefill-with-state succeeds"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert!(state.is_complete_for(weights.num_layers)); @@ -688,13 +734,13 @@ mod tests { use larql_compute::test_fixtures::make_q4k_fixture_index; use larql_models::test_fixtures::make_test_q4k_weights; let m = backend(); - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let idx = make_q4k_fixture_index(&weights); // Seed the KV cache via prefill. let (_h, mut handle) = m - .coarse_prefill(&mut weights, &[0u32, 1, 2], Some(&idx)) + .coarse_prefill(&weights, &[0u32, 1, 2], Some(&idx)) .expect("prefill seeds the cache"); - let result = m.coarse_decode_step(&mut weights, 4u32, Some(&idx), &mut handle, 3); + let result = m.coarse_decode_step(&weights, 4u32, Some(&idx), &mut handle, 3); let h = result.expect("Metal Q4K decode step returns Some"); assert_eq!(h.shape(), &[1, weights.hidden_size]); } @@ -707,10 +753,10 @@ mod tests { use larql_compute::test_fixtures::make_q4k_fixture_index; use larql_models::test_fixtures::make_test_q4k_weights; let m = backend(); - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let idx = make_q4k_fixture_index(&weights); let (_h, mut handle) = m - .coarse_prefill(&mut weights, &[0u32, 1, 2], Some(&idx)) + .coarse_prefill(&weights, &[0u32, 1, 2], Some(&idx)) .expect("prefill seeds the cache"); for mask in [ larql_compute::StateDumpMask::Full, @@ -719,7 +765,7 @@ mod tests { ] { let mut state = larql_compute::PerLayerDecodeState::with_capacity(weights.num_layers); let result = m.coarse_decode_step_with_state_masked( - &mut weights, + &weights, 5u32, Some(&idx), &mut handle, @@ -736,11 +782,11 @@ mod tests { #[test] fn coarse_decode_step_with_state_masked_without_index_returns_none() { - let mut weights = make_test_weights(); + let weights = make_test_weights(); let m = backend(); let mut handle = KvHandle::new(MetalCoarseHandle); let result = m.coarse_decode_step_with_state_masked( - &mut weights, + &weights, 0u32, None, &mut handle, diff --git a/crates/larql-compute-metal/src/stages/quant_matvec.rs b/crates/larql-compute-metal/src/stages/quant_matvec.rs index b5d016a20..c3ba2d859 100644 --- a/crates/larql-compute-metal/src/stages/quant_matvec.rs +++ b/crates/larql-compute-metal/src/stages/quant_matvec.rs @@ -219,6 +219,17 @@ pub fn encode( `Pipelines` struct to carry a Q8 kernel." ); } + larql_compute::QuantFormat::I2S => { + // BitNet ternary (I2_S) has no Metal shader yet — it is a + // CPU-only path (`CpuBackend::ternary_matvec` over a + // `BitLinearWeight`). Fail loudly rather than silently no-op so a + // mis-route is caught, mirroring the Q8_0 arm above. + panic!( + "metal::stages::quant_matvec::encode: I2_S (BitNet ternary) is \ + not implemented on Metal. Run ternary matvec on the CPU \ + backend (`ternary_matvec`), or add a Metal sign-select shader." + ); + } larql_compute::QuantFormat::BF16 | larql_compute::QuantFormat::F16 | larql_compute::QuantFormat::F32 => { diff --git a/crates/larql-compute/Cargo.toml b/crates/larql-compute/Cargo.toml index 0b8058dac..db3147332 100644 --- a/crates/larql-compute/Cargo.toml +++ b/crates/larql-compute/Cargo.toml @@ -86,6 +86,10 @@ harness = false name = "q4k_q8k_matvec" harness = false +[[bench]] +name = "ternary_matvec" +harness = false + # No external `cfg(cargo-clippy)` macros reach this crate after the # `metal-rs` / `objc` deps moved to larql-compute-metal, so no # crate-wide cfg-allowlist is needed. diff --git a/crates/larql-compute/benches/ternary_matvec.rs b/crates/larql-compute/benches/ternary_matvec.rs new file mode 100644 index 000000000..f42f2f024 --- /dev/null +++ b/crates/larql-compute/benches/ternary_matvec.rs @@ -0,0 +1,97 @@ +//! Ternary (BitNet I2_S) matrix-vector throughput: f32 reference vs the +//! W1.58·A8 paths (scalar int8 and NEON sign-select). +//! +//! Single-token (matrix-vector) over BitNet b1.58 2B shapes. Throughput is +//! counted on the packed I2_S weight stream (`rows * cols / 4` bytes), which +//! is the memory the kernel actually walks. +//! +//! Run: `cargo bench -p larql-compute --bench ternary_matvec` + +extern crate blas_src; + +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; +use larql_compute::cpu::ops::ternary_matvec::{ + matvec_i2s_f32_into, matvec_i2s_q8_into, quantize_activation_i8, BitLinearWeight, +}; + +/// Pack random ternary trits into I2_S bytes (contiguous layout). +fn build_weight(rows: usize, cols: usize, seed: u64) -> BitLinearWeight { + let mut s = seed; + let mut bytes = vec![0u8; rows * cols / 4]; + for byte in bytes.iter_mut() { + let mut bv = 0u8; + for slot in 0..4 { + s = s.wrapping_mul(6364136223846793005).wrapping_add(1); + let code = match (s >> 33) % 3 { + 0 => 0b00u8, + 1 => 0b01, + _ => 0b10, + }; + bv |= code << (2 * slot); + } + *byte = bv; + } + let scales: Vec = (0..rows).map(|r| 0.05 + (r % 7) as f32 * 0.01).collect(); + BitLinearWeight::new(rows, cols, bytes, scales).unwrap() +} + +fn synth(n: usize, seed: u64) -> Vec { + let mut s = seed; + (0..n) + .map(|_| { + s = s.wrapping_mul(6364136223846793005).wrapping_add(1); + ((s >> 33) as f32) / (u32::MAX as f32) * 2.0 - 1.0 + }) + .collect() +} + +fn bench_ternary_matvec(c: &mut Criterion) { + // (label, rows, cols) — BitNet b1.58 2B: hidden=2560, inter=6912. + let shapes = [ + ("attn_proj", 2560usize, 2560usize), + ("ffn_up", 6912, 2560), + ("ffn_down", 2560, 6912), + ]; + + let mut group = c.benchmark_group("ternary_matvec"); + group.sample_size(60); + + for (name, rows, cols) in shapes { + let w = build_weight(rows, cols, 99); + let x = synth(cols, 7); + let (x_i8, x_scale) = quantize_activation_i8(&x); + let mut y = vec![0.0f32; rows]; + + group.throughput(Throughput::Bytes((rows * cols / 4) as u64)); + + group.bench_with_input(BenchmarkId::new("f32", name), &(), |b, _| { + b.iter(|| { + matvec_i2s_f32_into(&w, &x, &mut y).unwrap(); + std::hint::black_box(y[0]); + }); + }); + + group.bench_with_input(BenchmarkId::new("a8_scalar", name), &(), |b, _| { + b.iter(|| { + matvec_i2s_q8_into(&w, &x_i8, x_scale, &mut y).unwrap(); + std::hint::black_box(y[0]); + }); + }); + + #[cfg(target_arch = "aarch64")] + group.bench_with_input(BenchmarkId::new("a8_neon", name), &(), |b, _| { + b.iter(|| { + larql_compute::cpu::ops::ternary_matvec::matvec_i2s_q8_neon_into( + &w, &x_i8, x_scale, &mut y, + ) + .unwrap(); + std::hint::black_box(y[0]); + }); + }); + } + + group.finish(); +} + +criterion_group!(benches, bench_ternary_matvec); +criterion_main!(benches); diff --git a/crates/larql-compute/examples/attn_prefill_f32_vs_q4k.rs b/crates/larql-compute/examples/attn_prefill_f32_vs_q4k.rs index 6fa160684..7d57e78f5 100644 --- a/crates/larql-compute/examples/attn_prefill_f32_vs_q4k.rs +++ b/crates/larql-compute/examples/attn_prefill_f32_vs_q4k.rs @@ -61,7 +61,13 @@ struct Geom { fn main() { let backend = &CpuBackend; let min_secs = 0.4; - let seq = 907; // matches the representative-context end-to-end run + // 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 { @@ -98,8 +104,8 @@ fn main() { ]; println!("── {} layer ──", g.name); println!( - "{:>4} | {:>12} | {:>10} | {:>10} | {:>8} | {:>14}", - "proj", "shape", "f32 ms", "q4k ms", "speedup", "q4k_matmul floor" + "{:>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; @@ -131,9 +137,21 @@ fn main() { } 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 @@ -147,12 +165,13 @@ fn main() { tot_f32 += f32_ms; tot_q4k += q4k_ms; println!( - "{:>4} | {:>12} | {:>10.3} | {:>10.3} | {:>7.2}× | {:>11.3} ms", + "{:>4} | {:>12} | {:>9.3} | {:>11.3} | {:>11.3} | {:>6.2}× | {:>6.3} ms", label, format!("[{num_rows}×{in_dim}]"), f32_ms, q4k_ms, - f32_ms / q4k_ms, + mm_ms, + f32_ms / mm_ms, floor_ms ); } diff --git a/crates/larql-compute/examples/ternary_a8_bench.rs b/crates/larql-compute/examples/ternary_a8_bench.rs new file mode 100644 index 000000000..7f743dcde --- /dev/null +++ b/crates/larql-compute/examples/ternary_a8_bench.rs @@ -0,0 +1,89 @@ +//! Throwaway microbench: ternary matvec — f32 reference vs scalar A8 (int8) +//! vs NEON sign-select. Single-token (matrix-vector) on BitNet-ish shapes. +//! +//! Run on AC, cool machine: +//! cargo run --release -p larql-compute --example ternary_a8_bench + +use larql_compute::cpu::ops::ternary_matvec::{ + matvec_i2s_f32_into, matvec_i2s_q8_into, quantize_activation_i8, BitLinearWeight, +}; +use std::time::Instant; + +fn synth(n: usize, seed: u64) -> Vec { + let mut s = seed; + (0..n) + .map(|_| { + s = s.wrapping_mul(6364136223846793005).wrapping_add(1); + ((s >> 33) as f32) / (u32::MAX as f32) * 2.0 - 1.0 + }) + .collect() +} + +fn build_weight(rows: usize, cols: usize) -> BitLinearWeight { + // pack random ternary trits into I2_S bytes (contiguous layout) + let mut s = 99u64; + let mut bytes = vec![0u8; rows * cols / 4]; + for byte in bytes.iter_mut() { + let mut bv = 0u8; + for slot in 0..4 { + s = s.wrapping_mul(6364136223846793005).wrapping_add(1); + let code = match (s >> 33) % 3 { + 0 => 0b00u8, + 1 => 0b01, + _ => 0b10, + }; + bv |= code << (2 * slot); + } + *byte = bv; + } + let scales: Vec = (0..rows).map(|r| 0.05 + (r % 7) as f32 * 0.01).collect(); + BitLinearWeight::new(rows, cols, bytes, scales).unwrap() +} + +fn bench(label: &str, rows: usize, cols: usize, iters: usize) { + let w = build_weight(rows, cols); + let x = synth(cols, 7); + let (x_i8, x_scale) = quantize_activation_i8(&x); + let mut y = vec![0.0f32; rows]; + + let time = |f: &mut dyn FnMut()| -> f64 { + for _ in 0..iters / 5 { + f(); + } + let t = Instant::now(); + for _ in 0..iters { + f(); + } + t.elapsed().as_secs_f64() * 1e6 / iters as f64 // µs/call + }; + + let f32_us = time(&mut || { + matvec_i2s_f32_into(&w, &x, &mut y).unwrap(); + }); + let q8_us = time(&mut || { + matvec_i2s_q8_into(&w, &x_i8, x_scale, &mut y).unwrap(); + }); + #[cfg(target_arch = "aarch64")] + let neon_us = time(&mut || { + larql_compute::cpu::ops::ternary_matvec::matvec_i2s_q8_neon_into( + &w, &x_i8, x_scale, &mut y, + ) + .unwrap(); + }); + #[cfg(not(target_arch = "aarch64"))] + let neon_us = f32_us; + + println!( + "{label:<18} {rows}x{cols} f32={f32_us:7.2}µs int8_scalar={q8_us:7.2}µs \ + neon={neon_us:7.2}µs neon speedup vs f32 = {:.2}x", + f32_us / neon_us + ); +} + +fn main() { + println!("ternary matvec microbench (single-token, µs/call)\n"); + bench("attn-proj", 2560, 2560, 2000); + bench("ffn-up", 6912, 2560, 1000); + bench("ffn-down", 2560, 6912, 1000); + bench("lm-head-ish", 8192, 2560, 800); +} diff --git a/crates/larql-compute/src/async_compute_backend/cpu.rs b/crates/larql-compute/src/async_compute_backend/cpu.rs index 28ad90bb0..edd4f337c 100644 --- a/crates/larql-compute/src/async_compute_backend/cpu.rs +++ b/crates/larql-compute/src/async_compute_backend/cpu.rs @@ -49,12 +49,11 @@ use ndarray::Array2; use super::{AsyncComputeBackend, AttentionHandle, ResidualUploadHandle}; use crate::ffn::FfnBackend; use crate::kv_dispatch::{KvDispatch, KvHandle, ResidualHandle}; -use larql_models::ModelWeights; impl AsyncComputeBackend for CpuBackend { fn attention_step_async( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, query: &Array2, kv: &mut KvHandle, layer: usize, @@ -76,7 +75,7 @@ impl AsyncComputeBackend for CpuBackend { fn attention_prefill_async( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, tokens_embedded: &Array2, layer: usize, window: Option, @@ -105,7 +104,7 @@ impl AsyncComputeBackend for CpuBackend { fn forward_from_layer_async( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, ffn: &dyn FfnBackend, start_layer: usize, residuals: &ResidualHandle, @@ -175,10 +174,22 @@ mod tests { // Populate two independent handles via the sync prefill — same // initial state for both sync and async decode-step paths. let (_, mut handle_sync) = backend - .attention_prefill(&weights, &h_in, 0, None, None) + .attention_prefill( + larql_models::WeightsView::dense(&weights), + &h_in, + 0, + None, + None, + ) .expect("attention_prefill"); let (_, mut handle_async) = backend - .attention_prefill(&weights, &h_in, 0, None, None) + .attention_prefill( + larql_models::WeightsView::dense(&weights), + &h_in, + 0, + None, + None, + ) .expect("attention_prefill"); let h_new = crate::forward::embed_tokens_pub(&weights, &[3u32]); @@ -186,7 +197,7 @@ mod tests { let h_sync = ::attention_step( &backend, - &weights, + larql_models::WeightsView::dense(&weights), &h_new, &mut handle_sync, 0, @@ -196,7 +207,14 @@ mod tests { .expect("sync attention_step"); let h_async = backend - .attention_step_async(&weights, &h_new, &mut handle_async, 0, abs_position, None) + .attention_step_async( + larql_models::WeightsView::dense(&weights), + &h_new, + &mut handle_async, + 0, + abs_position, + None, + ) .read(); assert_array_close( @@ -220,10 +238,21 @@ mod tests { let h_in = crate::forward::embed_tokens_pub(&weights, &tokens); let (h_sync, handle_sync) = backend - .attention_prefill(&weights, &h_in, 0, None, None) + .attention_prefill( + larql_models::WeightsView::dense(&weights), + &h_in, + 0, + None, + None, + ) .expect("sync prefill"); - let (h_async_handle, handle_async) = - backend.attention_prefill_async(&weights, &h_in, 0, None, None); + let (h_async_handle, handle_async) = backend.attention_prefill_async( + larql_models::WeightsView::dense(&weights), + &h_in, + 0, + None, + None, + ); let h_async = h_async_handle.read(); // BLAS on Windows runs successive matmuls with different @@ -275,12 +304,23 @@ mod tests { let handle_async = backend.upload_boundary_residual(&residual).unwrap(); let h_sync = backend - .forward_from_layer(&weights, 1, &handle_sync, &tokens) + .forward_from_layer( + larql_models::WeightsView::dense(&weights), + 1, + &handle_sync, + &tokens, + ) .expect("sync forward_from_layer"); let ffn = crate::ffn::NullFfn; let h_async = backend - .forward_from_layer_async(&weights, &ffn, 1, &handle_async, &tokens) + .forward_from_layer_async( + larql_models::WeightsView::dense(&weights), + &ffn, + 1, + &handle_async, + &tokens, + ) .read(); assert_eq!( @@ -300,24 +340,43 @@ mod tests { let h_in = crate::forward::embed_tokens_pub(&weights, &tokens); let (_, mut handle_step) = backend - .attention_prefill(&weights, &h_in, 0, None, None) + .attention_prefill( + larql_models::WeightsView::dense(&weights), + &h_in, + 0, + None, + None, + ) .expect("prefill"); let (_, mut handle_windowed) = backend - .attention_prefill(&weights, &h_in, 0, None, None) + .attention_prefill( + larql_models::WeightsView::dense(&weights), + &h_in, + 0, + None, + None, + ) .expect("prefill"); let h_new = crate::forward::embed_tokens_pub(&weights, &[5u32]); let abs_position = tokens.len(); let h_step = backend - .attention_step_async(&weights, &h_new, &mut handle_step, 0, abs_position, None) + .attention_step_async( + larql_models::WeightsView::dense(&weights), + &h_new, + &mut handle_step, + 0, + abs_position, + None, + ) .read(); // After step, manually clip to window=3. backend.clip_kv(&mut handle_step, 3); let h_windowed = backend .attention_step_windowed_async( - &weights, + larql_models::WeightsView::dense(&weights), &h_new, &mut handle_windowed, 0, @@ -350,7 +409,11 @@ mod tests { vec![0.0; weights.hidden_size], ) .unwrap(); - let _ = backend.recompute_kv_from_residuals_async(&weights, &residuals, 0); + let _ = backend.recompute_kv_from_residuals_async( + larql_models::WeightsView::dense(&weights), + &residuals, + 0, + ); // (recompute_kv_from_residuals_async signature unchanged at A3.) } diff --git a/crates/larql-compute/src/async_compute_backend/mod.rs b/crates/larql-compute/src/async_compute_backend/mod.rs index de9209070..2f1047d04 100644 --- a/crates/larql-compute/src/async_compute_backend/mod.rs +++ b/crates/larql-compute/src/async_compute_backend/mod.rs @@ -53,7 +53,6 @@ pub mod cpu; use crate::ffn::FfnBackend; use crate::kv_dispatch::{KvDispatch, KvHandle, ResidualHandle}; -use larql_models::ModelWeights; use ndarray::Array2; use std::error::Error; use std::fmt; @@ -300,7 +299,7 @@ pub trait AsyncComputeBackend: crate::ComputeBackend + KvDispatch + Send { /// [`KvDispatch::attention_step`]. fn attention_step_async( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, query: &Array2, kv: &mut KvHandle, layer: usize, @@ -319,7 +318,7 @@ pub trait AsyncComputeBackend: crate::ComputeBackend + KvDispatch + Send { #[allow(clippy::too_many_arguments)] fn attention_step_windowed_async( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, query: &Array2, kv: &mut KvHandle, layer: usize, @@ -337,7 +336,7 @@ pub trait AsyncComputeBackend: crate::ComputeBackend + KvDispatch + Send { /// `AttentionHandle` is pending until commit. fn attention_prefill_async( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, tokens_embedded: &Array2, layer: usize, window: Option, @@ -350,7 +349,7 @@ pub trait AsyncComputeBackend: crate::ComputeBackend + KvDispatch + Send { /// Async equivalent of [`KvDispatch::recompute_kv_from_residuals`]. fn recompute_kv_from_residuals_async( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, residuals: &Array2, layer: usize, ) -> KvHandle { @@ -375,7 +374,7 @@ pub trait AsyncComputeBackend: crate::ComputeBackend + KvDispatch + Send { /// Async equivalent of [`KvDispatch::forward_from_layer`]. fn forward_from_layer_async( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, ffn: &dyn FfnBackend, start_layer: usize, residuals: &ResidualHandle, @@ -458,7 +457,14 @@ mod tests { dim: weights.hidden_size, }); let query = Array2::zeros((1, weights.hidden_size)); - let _ = backend.attention_step_async(&weights, &query, &mut kv, 0, 0, None); + let _ = backend.attention_step_async( + larql_models::WeightsView::dense(&weights), + &query, + &mut kv, + 0, + 0, + None, + ); } #[test] @@ -467,7 +473,13 @@ mod tests { let backend = StubAsyncBackend; let weights = larql_models::test_fixtures::make_test_weights(); let tokens = Array2::zeros((2, weights.hidden_size)); - let _ = backend.attention_prefill_async(&weights, &tokens, 0, None, None); + let _ = backend.attention_prefill_async( + larql_models::WeightsView::dense(&weights), + &tokens, + 0, + None, + None, + ); } #[test] @@ -487,7 +499,13 @@ mod tests { shape: (1, weights.hidden_size), }); let ffn = crate::ffn::NullFfn; - let _ = backend.forward_from_layer_async(&weights, &ffn, 0, &residuals, &[0u32]); + let _ = backend.forward_from_layer_async( + larql_models::WeightsView::dense(&weights), + &ffn, + 0, + &residuals, + &[0u32], + ); } #[test] @@ -502,7 +520,15 @@ mod tests { }); let query = Array2::zeros((1, weights.hidden_size)); let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - let _ = backend.attention_step_windowed_async(&weights, &query, &mut kv, 0, 0, 4, None); + let _ = backend.attention_step_windowed_async( + larql_models::WeightsView::dense(&weights), + &query, + &mut kv, + 0, + 0, + 4, + None, + ); })); let err = result.expect_err("should panic via attention_step_async"); let msg = err diff --git a/crates/larql-compute/src/attention/block.rs b/crates/larql-compute/src/attention/block.rs index 094bbb466..b940adc7a 100644 --- a/crates/larql-compute/src/attention/block.rs +++ b/crates/larql-compute/src/attention/block.rs @@ -12,7 +12,7 @@ use ndarray::{s, Array2}; /// Run the full attention block. Returns (h_post_attn, attn_projected, optional_weights). #[allow(clippy::too_many_arguments)] pub fn run_attention_block( - weights: &larql_models::ModelWeights, + weights: larql_models::WeightsView, h: &Array2, layer: usize, capture_attention: bool, @@ -24,7 +24,7 @@ pub fn run_attention_block( #[allow(clippy::too_many_arguments)] #[allow(clippy::type_complexity)] pub fn run_attention_block_with_kv_out( - weights: &larql_models::ModelWeights, + weights: larql_models::WeightsView, h: &Array2, layer: usize, capture_attention: bool, @@ -55,7 +55,7 @@ pub fn run_attention_block_with_kv_out( /// Run attention with optional shared K/V (discards K/V output). #[allow(clippy::too_many_arguments)] pub fn run_attention_block_shared( - weights: &larql_models::ModelWeights, + weights: larql_models::WeightsView, h: &Array2, layer: usize, capture_attention: bool, @@ -81,7 +81,7 @@ pub fn run_attention_block_shared( /// Returns `(h_post_attn, pre_o)` where `pre_o` has shape `[seq, num_q * head_dim]`. /// This is the equivalent of Python's `o_proj.register_forward_pre_hook`. pub fn run_attention_block_with_pre_o( - weights: &larql_models::ModelWeights, + weights: larql_models::WeightsView, h: &Array2, layer: usize, ) -> Option<(Array2, Array2)> { @@ -97,7 +97,7 @@ pub fn run_attention_block_with_pre_o( /// This is the shared-KV-safe variant used by research/intervention adapters /// that need to inspect a pre-W_O head before deciding how to replace it. pub fn run_attention_block_shared_with_pre_o( - weights: &larql_models::ModelWeights, + weights: larql_models::WeightsView, h: &Array2, layer: usize, shared_kv: Option<&SharedKV>, @@ -115,7 +115,7 @@ pub fn run_attention_block_shared_with_pre_o( /// from normal attention capture because all-position weights are /// O(heads * seq^2) memory. pub fn run_attention_block_with_pre_o_and_all_attention_weights( - weights: &larql_models::ModelWeights, + weights: larql_models::WeightsView, h: &Array2, layer: usize, shared_kv: Option<&SharedKV>, @@ -133,7 +133,7 @@ pub fn run_attention_block_with_pre_o_and_all_attention_weights( /// weights use `qk_rank`, so this can test reduced address computation without /// changing the model forward path. pub fn run_attention_block_with_pre_o_and_reduced_qk_attention_weights( - weights: &larql_models::ModelWeights, + weights: larql_models::WeightsView, h: &Array2, layer: usize, shared_kv: Option<&SharedKV>, @@ -160,7 +160,7 @@ pub fn run_attention_block_with_pre_o_and_reduced_qk_attention_weights( /// Returns the post-attention residual and, when K/V were computed by this call, /// the K/V pair for cross-layer sharing. pub fn run_attention_block_zero_pre_o_heads( - weights: &larql_models::ModelWeights, + weights: larql_models::WeightsView, h: &Array2, layer: usize, heads: &[usize], @@ -191,7 +191,7 @@ pub fn run_attention_block_zero_pre_o_heads( /// /// `replacement` must have shape `[seq_len, head_dim]`. pub fn run_attention_block_replace_pre_o_head( - weights: &larql_models::ModelWeights, + weights: larql_models::WeightsView, h: &Array2, layer: usize, head: usize, @@ -225,7 +225,7 @@ pub fn run_attention_block_replace_pre_o_head( /// This is numerically equivalent to zeroing those pre-W_O heads, but it checks /// the head-to-W_O block indexing independently. pub fn run_attention_block_subtract_pre_o_heads( - weights: &larql_models::ModelWeights, + weights: larql_models::WeightsView, h: &Array2, layer: usize, heads: &[usize], @@ -260,7 +260,7 @@ pub fn run_attention_block_subtract_pre_o_heads( /// This is the Mode D validation surface: runtime lookup/add tables can bypass /// W_O entirely while the rest of the layer remains unchanged. pub fn run_attention_block_replace_head_residual_delta( - weights: &larql_models::ModelWeights, + weights: larql_models::WeightsView, h: &Array2, layer: usize, head: usize, @@ -292,7 +292,7 @@ pub fn run_attention_block_replace_head_residual_delta( #[allow(clippy::too_many_arguments)] #[allow(clippy::type_complexity)] fn run_attention_block_core( - weights: &larql_models::ModelWeights, + weights: larql_models::WeightsView, h: &Array2, layer: usize, capture_attention: bool, @@ -344,12 +344,12 @@ fn run_attention_block_core( // Input norm let h_norm = - crate::forward::apply_norm(weights, h, &arch.input_layernorm_key(layer), norm_offset); + crate::forward::apply_norm(&weights, h, &arch.input_layernorm_key(layer), norm_offset); dump_f32("norm_out", &h_norm); // Q projection (always from current hidden state) - let w_q = weights.tensors.get(&arch.attn_q_key(layer))?; - let w_o = weights.tensors.get(&arch.attn_o_key(layer)).unwrap(); + let w_q = weights.tensor(&arch.attn_q_key(layer))?; + let w_o = weights.tensor(&arch.attn_o_key(layer)).unwrap(); let mut q_full = dot_proj(&h_norm, w_q); if let Some(bias) = arch .attn_q_bias_key(layer) @@ -396,7 +396,7 @@ fn run_attention_block_core( let (k_rope, v_final) = if let Some((cached_k, cached_v)) = shared_kv { (cached_k.clone(), cached_v.clone()) } else { - let w_k = weights.tensors.get(&arch.attn_k_key(layer)).unwrap(); + let w_k = weights.tensor(&arch.attn_k_key(layer)).unwrap(); let mut k_full = dot_proj(&h_norm, w_k); if let Some(bias) = arch @@ -428,7 +428,7 @@ fn run_attention_block_core( // Fallback: when W_v is genuinely absent from the vindex (older // extracts with no v_proj tensor for `attention_k_eq_v` layers), // reuse `k_full` — matches pre-Q6K-V behaviour. - let v_full = if let Some(w_v) = weights.tensors.get(&arch.attn_v_key(layer)) { + let v_full = if let Some(w_v) = weights.tensor(&arch.attn_v_key(layer)) { let mut v = dot_proj(&h_norm, w_v); if let Some(bias) = arch .attn_v_bias_key(layer) @@ -555,7 +555,7 @@ fn run_attention_block_core( let res_mult = arch.residual_multiplier(); let h_post_attn = if arch.has_post_norms() { let normed = crate::forward::apply_norm( - weights, + &weights, &attn_projected, &arch.post_attention_layernorm_key(layer), norm_offset, @@ -606,7 +606,8 @@ mod tests { let weights = make_test_weights(); let h = hidden(3, weights.hidden_size); let (h_out, attn_proj, _) = - run_attention_block(&weights, &h, 0, false).expect("run_attention_block failed"); + run_attention_block(larql_models::WeightsView::dense(&weights), &h, 0, false) + .expect("run_attention_block failed"); assert_eq!(h_out.shape(), &[3, weights.hidden_size]); assert_eq!(attn_proj.shape()[0], 3); } @@ -615,7 +616,8 @@ mod tests { fn attention_block_output_finite() { let weights = make_test_weights(); let h = hidden(2, weights.hidden_size); - let (h_out, _, _) = run_attention_block(&weights, &h, 0, false).unwrap(); + let (h_out, _, _) = + run_attention_block(larql_models::WeightsView::dense(&weights), &h, 0, false).unwrap(); assert!(h_out.iter().all(|v| v.is_finite())); } @@ -623,7 +625,8 @@ mod tests { fn attention_block_single_token() { let weights = make_test_weights(); let h = hidden(1, weights.hidden_size); - let (h_out, attn_proj, _) = run_attention_block(&weights, &h, 0, false).unwrap(); + let (h_out, attn_proj, _) = + run_attention_block(larql_models::WeightsView::dense(&weights), &h, 0, false).unwrap(); assert_eq!(h_out.shape(), &[1, weights.hidden_size]); assert_eq!(attn_proj.shape()[0], 1); } @@ -634,7 +637,8 @@ mod tests { let h = hidden(2, weights.hidden_size); for layer in 0..weights.num_layers { assert!( - run_attention_block(&weights, &h, layer, false).is_some(), + run_attention_block(larql_models::WeightsView::dense(&weights), &h, layer, false) + .is_some(), "layer {layer} failed" ); } @@ -644,7 +648,13 @@ mod tests { fn attention_block_with_kv_out_returns_kv() { let weights = make_test_weights(); let h = hidden(3, weights.hidden_size); - let result = run_attention_block_with_kv_out(&weights, &h, 0, false, None); + let result = run_attention_block_with_kv_out( + larql_models::WeightsView::dense(&weights), + &h, + 0, + false, + None, + ); // Returns (h_post, attn_proj, attn_w, k_rope, v_final) — 5 elements let (h_out, _attn_proj, _attn_w, k_rope, v_final) = result.unwrap(); assert_eq!(h_out.shape(), &[3, weights.hidden_size]); @@ -656,7 +666,8 @@ mod tests { fn attention_block_capture_returns_per_head_weights() { let weights = make_test_weights(); let h = hidden(3, weights.hidden_size); - let (_, _, attn_w) = run_attention_block(&weights, &h, 0, true).unwrap(); + let (_, _, attn_w) = + run_attention_block(larql_models::WeightsView::dense(&weights), &h, 0, true).unwrap(); let aw = attn_w.expect("capture=true must yield weights"); assert_eq!(aw.heads.len(), weights.num_q_heads); } @@ -665,7 +676,9 @@ mod tests { fn attention_block_with_pre_o_returns_per_head_pre_projection() { let weights = make_test_weights(); let h = hidden(2, weights.hidden_size); - let (h_post, pre_o) = run_attention_block_with_pre_o(&weights, &h, 0).unwrap(); + let (h_post, pre_o) = + run_attention_block_with_pre_o(larql_models::WeightsView::dense(&weights), &h, 0) + .unwrap(); assert_eq!(h_post.shape(), &[2, weights.hidden_size]); // pre_o is `[seq, num_q * head_dim]`. assert_eq!(pre_o.shape(), &[2, weights.num_q_heads * weights.head_dim]); @@ -675,12 +688,22 @@ mod tests { fn attention_block_shared_with_pre_o_works_under_kv_share() { let weights = make_test_weights(); let h = hidden(2, weights.hidden_size); - let (_, shared) = - crate::attention::run_attention_block_with_kv_out(&weights, &h, 0, false, None) - .map(|(p, _, _, k, v)| (p, (k, v))) - .unwrap(); - let (h_post, pre_o) = - run_attention_block_shared_with_pre_o(&weights, &h, 1, Some(&shared)).unwrap(); + let (_, shared) = crate::attention::run_attention_block_with_kv_out( + larql_models::WeightsView::dense(&weights), + &h, + 0, + false, + None, + ) + .map(|(p, _, _, k, v)| (p, (k, v))) + .unwrap(); + let (h_post, pre_o) = run_attention_block_shared_with_pre_o( + larql_models::WeightsView::dense(&weights), + &h, + 1, + Some(&shared), + ) + .unwrap(); assert_eq!(h_post.shape(), &[2, weights.hidden_size]); assert_eq!(pre_o.shape()[1], weights.num_q_heads * weights.head_dim); } @@ -689,9 +712,13 @@ mod tests { fn attention_block_with_all_attention_weights_returns_per_position_dist() { let weights = make_test_weights(); let h = hidden(3, weights.hidden_size); - let (_, _, all) = - run_attention_block_with_pre_o_and_all_attention_weights(&weights, &h, 0, None) - .unwrap(); + let (_, _, all) = run_attention_block_with_pre_o_and_all_attention_weights( + larql_models::WeightsView::dense(&weights), + &h, + 0, + None, + ) + .unwrap(); assert_eq!(all.heads.len(), weights.num_q_heads); for head in &all.heads { assert_eq!(head.len(), 3); // one distribution per Q position @@ -703,7 +730,11 @@ mod tests { let weights = make_test_weights(); let h = hidden(2, weights.hidden_size); let (_, _, all) = run_attention_block_with_pre_o_and_reduced_qk_attention_weights( - &weights, &h, 0, None, /*qk_rank=*/ 4, // half of head_dim=8 + larql_models::WeightsView::dense(&weights), + &h, + 0, + None, + /*qk_rank=*/ 4, // half of head_dim=8 ) .unwrap(); assert_eq!(all.heads.len(), weights.num_q_heads); @@ -715,9 +746,18 @@ mod tests { fn zero_pre_o_heads_changes_output_when_head_active() { let weights = make_test_weights(); let h = hidden(2, weights.hidden_size); - let baseline = run_attention_block(&weights, &h, 0, false).unwrap().0; - let (zeroed, kv_out) = - run_attention_block_zero_pre_o_heads(&weights, &h, 0, &[0], None).unwrap(); + let baseline = + run_attention_block(larql_models::WeightsView::dense(&weights), &h, 0, false) + .unwrap() + .0; + let (zeroed, kv_out) = run_attention_block_zero_pre_o_heads( + larql_models::WeightsView::dense(&weights), + &h, + 0, + &[0], + None, + ) + .unwrap(); assert_eq!(zeroed.shape(), baseline.shape()); // KV is computed when shared_kv is None. assert!(kv_out.is_some()); @@ -732,12 +772,23 @@ mod tests { fn zero_pre_o_heads_under_shared_kv_omits_kv_output() { let weights = make_test_weights(); let h = hidden(2, weights.hidden_size); - let (_, shared) = - crate::attention::run_attention_block_with_kv_out(&weights, &h, 0, false, None) - .map(|(p, _, _, k, v)| (p, (k, v))) - .unwrap(); - let (_, kv_out) = - run_attention_block_zero_pre_o_heads(&weights, &h, 1, &[0], Some(&shared)).unwrap(); + let (_, shared) = crate::attention::run_attention_block_with_kv_out( + larql_models::WeightsView::dense(&weights), + &h, + 0, + false, + None, + ) + .map(|(p, _, _, k, v)| (p, (k, v))) + .unwrap(); + let (_, kv_out) = run_attention_block_zero_pre_o_heads( + larql_models::WeightsView::dense(&weights), + &h, + 1, + &[0], + Some(&shared), + ) + .unwrap(); assert!(kv_out.is_none(), "shared-KV path must not return KV"); } @@ -747,10 +798,20 @@ mod tests { let h = hidden(2, weights.hidden_size); // Replacement is `[seq, head_dim]` of all-zeros — equivalent to // zeroing that head, so output must differ from baseline. - let baseline = run_attention_block(&weights, &h, 0, false).unwrap().0; + let baseline = + run_attention_block(larql_models::WeightsView::dense(&weights), &h, 0, false) + .unwrap() + .0; let zero_head = Array2::::zeros((2, weights.head_dim)); - let (replaced, _) = - run_attention_block_replace_pre_o_head(&weights, &h, 0, 0, &zero_head, None).unwrap(); + let (replaced, _) = run_attention_block_replace_pre_o_head( + larql_models::WeightsView::dense(&weights), + &h, + 0, + 0, + &zero_head, + None, + ) + .unwrap(); let mut max_diff = 0.0f32; for (a, b) in baseline.iter().zip(replaced.iter()) { max_diff = max_diff.max((a - b).abs()); @@ -763,10 +824,22 @@ mod tests { // Both paths zero the head's W_O contribution — output must match. let weights = make_test_weights(); let h = hidden(2, weights.hidden_size); - let (zeroed, _) = - run_attention_block_zero_pre_o_heads(&weights, &h, 0, &[0], None).unwrap(); - let (subtracted, _) = - run_attention_block_subtract_pre_o_heads(&weights, &h, 0, &[0], None).unwrap(); + let (zeroed, _) = run_attention_block_zero_pre_o_heads( + larql_models::WeightsView::dense(&weights), + &h, + 0, + &[0], + None, + ) + .unwrap(); + let (subtracted, _) = run_attention_block_subtract_pre_o_heads( + larql_models::WeightsView::dense(&weights), + &h, + 0, + &[0], + None, + ) + .unwrap(); for (a, b) in zeroed.iter().zip(subtracted.iter()) { assert!( (a - b).abs() < 1e-4, @@ -782,11 +855,23 @@ mod tests { let weights = make_test_weights(); let h = hidden(2, weights.hidden_size); let zero_delta = Array2::::zeros((2, weights.hidden_size)); - let (with_delta, _) = - run_attention_block_replace_head_residual_delta(&weights, &h, 0, 0, &zero_delta, None) - .unwrap(); - let (zeroed, _) = - run_attention_block_zero_pre_o_heads(&weights, &h, 0, &[0], None).unwrap(); + let (with_delta, _) = run_attention_block_replace_head_residual_delta( + larql_models::WeightsView::dense(&weights), + &h, + 0, + 0, + &zero_delta, + None, + ) + .unwrap(); + let (zeroed, _) = run_attention_block_zero_pre_o_heads( + larql_models::WeightsView::dense(&weights), + &h, + 0, + &[0], + None, + ) + .unwrap(); for (a, b) in with_delta.iter().zip(zeroed.iter()) { assert!( (a - b).abs() < 1e-4, @@ -802,11 +887,27 @@ mod tests { let weights = make_test_weights(); let h = hidden(2, weights.hidden_size); let bogus_layer = weights.num_layers + 5; - assert!(run_attention_block(&weights, &h, bogus_layer, false).is_none()); - assert!(run_attention_block_with_pre_o(&weights, &h, bogus_layer).is_none()); - assert!( - run_attention_block_zero_pre_o_heads(&weights, &h, bogus_layer, &[0], None).is_none() - ); + assert!(run_attention_block( + larql_models::WeightsView::dense(&weights), + &h, + bogus_layer, + false + ) + .is_none()); + assert!(run_attention_block_with_pre_o( + larql_models::WeightsView::dense(&weights), + &h, + bogus_layer + ) + .is_none()); + assert!(run_attention_block_zero_pre_o_heads( + larql_models::WeightsView::dense(&weights), + &h, + bogus_layer, + &[0], + None + ) + .is_none()); } // ── Gemma3-arch fixture (post-norms, QK norm, gelu_tanh) ─────────── @@ -818,7 +919,8 @@ mod tests { // tinymodel never exercises. let weights = larql_models::test_fixtures::make_gemma3_test_weights(); let h = hidden(2, weights.hidden_size); - let (h_post, _, _) = run_attention_block(&weights, &h, 0, false).unwrap(); + let (h_post, _, _) = + run_attention_block(larql_models::WeightsView::dense(&weights), &h, 0, false).unwrap(); assert_eq!(h_post.shape(), &[2, weights.hidden_size]); assert!(h_post.iter().all(|v| v.is_finite())); } @@ -829,7 +931,8 @@ mod tests { // takes a different branch in run_attention_block_core. let weights = larql_models::test_fixtures::make_gemma3_test_weights(); let h = hidden(2, weights.hidden_size); - let (h_post, _, _) = run_attention_block(&weights, &h, 1, false).unwrap(); + let (h_post, _, _) = + run_attention_block(larql_models::WeightsView::dense(&weights), &h, 1, false).unwrap(); assert_eq!(h_post.shape(), &[2, weights.hidden_size]); assert!(h_post.iter().all(|v| v.is_finite())); } @@ -842,7 +945,8 @@ mod tests { // `add_bias` call site fires. let weights = larql_models::test_fixtures::make_starcoder2_test_weights(); let h = hidden(2, weights.hidden_size); - let (h_post, _, _) = run_attention_block(&weights, &h, 0, false).unwrap(); + let (h_post, _, _) = + run_attention_block(larql_models::WeightsView::dense(&weights), &h, 0, false).unwrap(); assert_eq!(h_post.shape(), &[2, weights.hidden_size]); assert!(h_post.iter().all(|v| v.is_finite())); } @@ -851,8 +955,14 @@ mod tests { fn attention_block_starcoder_with_kv_out_returns_finite_kv() { let weights = larql_models::test_fixtures::make_starcoder2_test_weights(); let h = hidden(3, weights.hidden_size); - let (_, _, _, k, v) = - run_attention_block_with_kv_out(&weights, &h, 0, false, None).unwrap(); + let (_, _, _, k, v) = run_attention_block_with_kv_out( + larql_models::WeightsView::dense(&weights), + &h, + 0, + false, + None, + ) + .unwrap(); assert_eq!(k.shape()[0], 3); assert_eq!(v.shape()[0], 3); assert!(k.iter().all(|x| x.is_finite())); diff --git a/crates/larql-compute/src/attention/decode.rs b/crates/larql-compute/src/attention/decode.rs index 268bfb059..09a116724 100644 --- a/crates/larql-compute/src/attention/decode.rs +++ b/crates/larql-compute/src/attention/decode.rs @@ -141,7 +141,14 @@ pub fn run_attention_block_decode_step( kv_entry: Option<&SharedKV>, abs_position: usize, ) -> Option<(Array2, SharedKV)> { - run_attention_block_decode_step_backend(weights, h_new, layer, kv_entry, abs_position, None) + run_attention_block_decode_step_backend( + larql_models::WeightsView::dense(weights), + h_new, + layer, + kv_entry, + abs_position, + None, + ) } /// Decode-step attention with optional GPU-accelerated projections @@ -152,7 +159,7 @@ pub fn run_attention_block_decode_step( #[allow(clippy::too_many_arguments)] #[allow(clippy::type_complexity)] pub fn run_attention_block_decode_step_backend( - weights: &larql_models::ModelWeights, + weights: larql_models::WeightsView, h_new: &Array2, layer: usize, kv_entry: Option<&SharedKV>, @@ -177,14 +184,14 @@ pub fn run_attention_block_decode_step_backend( let position = abs_position; let h_norm = crate::forward::apply_norm( - weights, + &weights, h_new, &arch.input_layernorm_key(layer), norm_offset, ); - let w_q = weights.tensors.get(&arch.attn_q_key(layer))?; - let w_o = weights.tensors.get(&arch.attn_o_key(layer))?; + let w_q = weights.tensor(&arch.attn_q_key(layer))?; + let w_o = weights.tensor(&arch.attn_o_key(layer))?; let mut q_full = dot_proj_gpu(&h_norm, w_q, backend); if let Some(bias) = arch .attn_q_bias_key(layer) @@ -223,12 +230,12 @@ pub fn run_attention_block_decode_step_backend( ); // New token's K, V — RoPE'd at `position`, then appended to cache. - let w_k = weights.tensors.get(&arch.attn_k_key(layer))?; - let v_from_k = !weights.tensors.contains_key(&arch.attn_v_key(layer)); + let w_k = weights.tensor(&arch.attn_k_key(layer))?; + let v_from_k = !weights.has_tensor(&arch.attn_v_key(layer)); let w_v = if v_from_k { w_k } else { - weights.tensors.get(&arch.attn_v_key(layer))? + weights.tensor(&arch.attn_v_key(layer))? }; let mut k_full_new = dot_proj_gpu(&h_norm, w_k, backend); @@ -306,7 +313,7 @@ pub fn run_attention_block_decode_step_backend( let res_mult = arch.residual_multiplier(); let h_post_attn = if arch.has_post_norms() { let normed = crate::forward::apply_norm( - weights, + &weights, &attn_projected, &arch.post_attention_layernorm_key(layer), norm_offset, @@ -346,7 +353,7 @@ pub fn q4k_direct_attn_enabled() -> bool { #[allow(clippy::too_many_arguments)] #[allow(clippy::type_complexity)] pub fn run_attention_block_decode_step_auto( - weights: &larql_models::ModelWeights, + weights: larql_models::WeightsView, h_new: &Array2, layer: usize, kv_entry: Option<&SharedKV>, @@ -356,8 +363,10 @@ pub fn run_attention_block_decode_step_auto( ) -> Option<(Array2, SharedKV)> { if q4k_direct_attn_enabled() { if let (Some(be), Some(idx)) = (backend, index) { + // Q4K-direct reads native packed bytes from `index` (not the + // dequant scratch) — canonical weights for norms/config. if let Some(r) = run_attention_block_decode_step_q4k_direct( - weights, + weights.canonical(), h_new, layer, kv_entry, @@ -831,7 +840,7 @@ pub fn run_attention_block_decode_step_q4k_direct_inplace( /// [`run_attention_block_decode_step_auto`]. #[allow(clippy::too_many_arguments)] pub fn run_attention_block_decode_step_auto_inplace( - weights: &larql_models::ModelWeights, + weights: larql_models::WeightsView, h_new: &Array2, layer: usize, k_cache: &mut Array2, @@ -844,7 +853,7 @@ pub fn run_attention_block_decode_step_auto_inplace( if q4k_direct_attn_enabled() { if let (Some(be), Some(idx)) = (backend, index) { return run_attention_block_decode_step_q4k_direct_inplace( - weights, + weights.canonical(), h_new, layer, k_cache, @@ -1230,7 +1239,8 @@ mod tests { // dequantised into `weights.tensors`. We dequantise the fixture's // Q4K attn slices into the weights (what `insert_q4k_layer_tensors` // does) and compare against `run_attention_block_decode_step_backend`. - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); + let mut scratch = larql_models::DequantScratch::new(); let idx = make_q4k_fixture_index(&weights); let backend = crate::CpuBackend; let h = Array2::from_elem((1, weights.hidden_size), 0.1f32); @@ -1242,11 +1252,21 @@ mod tests { // Dequant the index's layer-0 attn/ffn bytes into weights.tensors, // then run the f32 path against those same (dequantised) weights. - let inserted = crate::kquant_forward::insert_q4k_layer_tensors(&mut weights, &idx, 0) - .expect("dequant layer 0 tensors"); - let (h_dequant, _) = - run_attention_block_decode_step_backend(&weights, &h, 0, None, 0, Some(&backend)) - .expect("f32 dequant step"); + let inserted = + crate::kquant_forward::insert_q4k_layer_tensors(&mut scratch, &weights, &idx, 0) + .expect("dequant layer 0 tensors"); + // The dequantised attn tensors live in `scratch` (not `weights.tensors`); + // resolve them via `with_scratch` so the f32 path reads the SAME bytes + // the Q4K-direct path read from the index. + let (h_dequant, _) = run_attention_block_decode_step_backend( + larql_models::WeightsView::with_scratch(&weights, &scratch), + &h, + 0, + None, + 0, + Some(&backend), + ) + .expect("f32 dequant step"); let _ = inserted; assert_eq!(h_direct.shape(), h_dequant.shape()); @@ -1277,7 +1297,7 @@ mod tests { let h = Array2::from_elem((1, weights.hidden_size), 0.1f32); let (out, _kv) = run_attention_block_decode_step_auto( - &weights, + larql_models::WeightsView::dense(&weights), &h, 0, None, @@ -1293,18 +1313,20 @@ mod tests { fn auto_decode_step_falls_back_to_f32_when_flag_disabled() { let _guard = crate::options::FastPathGuard::set(&[(crate::options::ENV_Q4K_DIRECT_ATTN, false)]); - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); + let mut scratch = larql_models::DequantScratch::new(); let idx = make_q4k_fixture_index(&weights); - // The f32 backend path reads `weights.tensors`; dequant the fixture's - // Q4K layer-0 bytes into them so the fallback produces a real result. - crate::kquant_forward::insert_q4k_layer_tensors(&mut weights, &idx, 0) + // The f32 backend path reads its `WeightsView`; dequant the fixture's + // Q4K layer-0 bytes into `scratch` so the fallback produces a real + // result (resolved via `with_scratch` below). + crate::kquant_forward::insert_q4k_layer_tensors(&mut scratch, &weights, &idx, 0) .expect("dequant layer 0"); let backend = crate::CpuBackend; let h = Array2::from_elem((1, weights.hidden_size), 0.1f32); // Flag off → Q4K branch skipped → `run_attention_block_decode_step_backend`. let (out, _kv) = run_attention_block_decode_step_auto( - &weights, + larql_models::WeightsView::with_scratch(&weights, &scratch), &h, 0, None, @@ -1332,7 +1354,7 @@ mod tests { let h = Array2::from_elem((1, weights.hidden_size), 0.1f32); let out = run_attention_block_decode_step_auto_inplace( - &weights, + larql_models::WeightsView::dense(&weights), &h, 0, &mut k_cache, @@ -1368,7 +1390,7 @@ mod tests { // Flag off → returns None so the caller uses the owned-concat path. let out = run_attention_block_decode_step_auto_inplace( - &weights, + larql_models::WeightsView::dense(&weights), &h, 0, &mut k_cache, diff --git a/crates/larql-compute/src/attention/gpu.rs b/crates/larql-compute/src/attention/gpu.rs index 807b75712..329cf41b4 100644 --- a/crates/larql-compute/src/attention/gpu.rs +++ b/crates/larql-compute/src/attention/gpu.rs @@ -11,7 +11,7 @@ use ndarray::Array2; /// GPU-accelerated attention block. Same as `run_attention_block` but routes /// Q/K/V/O projections through the ComputeBackend (Metal, CUDA, or CPU). pub fn run_attention_block_gpu( - weights: &larql_models::ModelWeights, + weights: larql_models::WeightsView, h: &Array2, layer: usize, capture_attention: bool, @@ -35,17 +35,17 @@ pub fn run_attention_block_gpu( let norm_offset = arch.norm_weight_offset(); let h_norm = - crate::forward::apply_norm(weights, h, &arch.input_layernorm_key(layer), norm_offset); + crate::forward::apply_norm(&weights, h, &arch.input_layernorm_key(layer), norm_offset); - let w_q = weights.tensors.get(&arch.attn_q_key(layer))?; - let w_k = weights.tensors.get(&arch.attn_k_key(layer)).unwrap(); - let v_from_k = !weights.tensors.contains_key(&arch.attn_v_key(layer)); + let w_q = weights.tensor(&arch.attn_q_key(layer))?; + let w_k = weights.tensor(&arch.attn_k_key(layer)).unwrap(); + let v_from_k = !weights.has_tensor(&arch.attn_v_key(layer)); let w_v = if v_from_k { w_k } else { - weights.tensors.get(&arch.attn_v_key(layer)).unwrap() + weights.tensor(&arch.attn_v_key(layer)).unwrap() }; - let w_o = weights.tensors.get(&arch.attn_o_key(layer)).unwrap(); + let w_o = weights.tensor(&arch.attn_o_key(layer)).unwrap(); let mut q_full = dot_proj_gpu(&h_norm, w_q, backend); let mut k_full = dot_proj_gpu(&h_norm, w_k, backend); @@ -125,7 +125,7 @@ pub fn run_attention_block_gpu( let res_mult = arch.residual_multiplier(); let h_post_attn = if arch.has_post_norms() { let normed = crate::forward::apply_norm( - weights, + &weights, &attn_projected, &arch.post_attention_layernorm_key(layer), norm_offset, @@ -151,15 +151,25 @@ pub fn run_attention_with_kv( h: &Array2, layer: usize, ) -> Option<(Array2, Array2, Array2)> { - run_attention_with_kv_backend(weights, h, layer, None) + run_attention_with_kv_backend( + larql_models::WeightsView::dense(weights), + h, + layer, + None, + None, + ) } /// Run attention with optional compute backend for accelerated projections. pub fn run_attention_with_kv_backend( - weights: &larql_models::ModelWeights, + weights: larql_models::WeightsView, h: &Array2, layer: usize, backend: Option<&dyn crate::ComputeBackend>, + // When `Some`, project Q/K/V/O straight from the vindex's Q4_K/Q6_K bytes + // (q4k-direct prefill — skips the f32 dequant). When `None`, read the + // dequantised projection weights from `weights` (the f32 view path). + index: Option<&dyn crate::KvIndex>, ) -> Option<(Array2, Array2, Array2)> { use crate::forward::{add_bias, apply_norm}; use crate::residual::{rms_norm_heads, rms_norm_heads_no_weight}; @@ -177,22 +187,41 @@ pub fn run_attention_with_kv_backend( let seq_len = h.shape()[0]; let norm_off = arch.norm_weight_offset(); - let h_norm = apply_norm(weights, h, &arch.input_layernorm_key(layer), norm_off); - let wq = weights.tensors.get(&arch.attn_q_key(layer))?; - let wk = weights.tensors.get(&arch.attn_k_key(layer))?; - let v_from_k = !weights.tensors.contains_key(&arch.attn_v_key(layer)); - let wv = if v_from_k { - wk + let h_norm = apply_norm(&weights, h, &arch.input_layernorm_key(layer), norm_off); + + // q4k-direct: project from the raw Q4_K/Q6_K attn bytes (Q/K/O are Q4_K, + // V is Q6_K in a default vindex) when an index is supplied; else read the + // dequantised projection weights from the view. The O projection (below) + // dispatches the same way. + let attn_q4k: Option<[(&[u8], &str); 4]> = match index { + Some(idx) => Some(idx.attn_kquant_layer_data(layer)?), + None => None, + }; + let q_dim = nq * hd; + let kv_dim = nkv * hd; + let in_dim = h_norm.ncols(); + + let (mut q, mut k, mut v) = if let Some(attn) = attn_q4k { + ( + crate::ffn::weight::quant_proj(attn[0].0, attn[0].1, &h_norm, q_dim, in_dim, seq_len), + crate::ffn::weight::quant_proj(attn[1].0, attn[1].1, &h_norm, kv_dim, in_dim, seq_len), + crate::ffn::weight::quant_proj(attn[2].0, attn[2].1, &h_norm, kv_dim, in_dim, seq_len), + ) } else { - weights.tensors.get(&arch.attn_v_key(layer))? + let wq = weights.tensor(&arch.attn_q_key(layer))?; + let wk = weights.tensor(&arch.attn_k_key(layer))?; + let v_from_k = !weights.has_tensor(&arch.attn_v_key(layer)); + let wv = if v_from_k { + wk + } else { + weights.tensor(&arch.attn_v_key(layer))? + }; + ( + crate::dot_proj_gpu(&h_norm, wq, backend), + crate::dot_proj_gpu(&h_norm, wk, backend), + crate::dot_proj_gpu(&h_norm, wv, backend), + ) }; - let wo = weights.tensors.get(&arch.attn_o_key(layer))?; - - let (mut q, mut k, mut v) = ( - crate::dot_proj_gpu(&h_norm, wq, backend), - crate::dot_proj_gpu(&h_norm, wk, backend), - crate::dot_proj_gpu(&h_norm, wv, backend), - ); for (proj, bias_fn) in [ (&mut q, arch.attn_q_bias_key(layer) as Option), (&mut k, arch.attn_k_bias_key(layer)), @@ -265,7 +294,12 @@ pub fn run_attention_with_kv_backend( false, arch.attn_logit_softcapping(), ); - let mut o = crate::dot_proj_gpu(&attn_out, wo, backend); + let mut o = if let Some(attn) = attn_q4k { + crate::ffn::weight::quant_proj(attn[3].0, attn[3].1, &attn_out, in_dim, q_dim, seq_len) + } else { + let wo = weights.tensor(&arch.attn_o_key(layer))?; + crate::dot_proj_gpu(&attn_out, wo, backend) + }; if let Some(b) = arch .attn_o_bias_key(layer) .and_then(|k| weights.vectors.get(&k)) @@ -276,7 +310,7 @@ pub fn run_attention_with_kv_backend( let rm = arch.residual_multiplier(); let h_out = if arch.has_post_norms() { let n = apply_norm( - weights, + &weights, &o, &arch.post_attention_layernorm_key(layer), norm_off, @@ -341,8 +375,14 @@ mod tests { fn run_attention_block_gpu_no_backend_falls_back_to_cpu() { let weights = make_test_weights(); let input = h(2, weights.hidden_size); - let (h_post, attn_proj, attn_w) = - run_attention_block_gpu(&weights, &input, 0, false, None).unwrap(); + let (h_post, attn_proj, attn_w) = run_attention_block_gpu( + larql_models::WeightsView::dense(&weights), + &input, + 0, + false, + None, + ) + .unwrap(); assert_eq!(h_post.shape(), &[2, weights.hidden_size]); assert_eq!(attn_proj.shape()[0], 2); assert!(attn_w.is_none()); @@ -352,9 +392,22 @@ mod tests { fn run_attention_block_gpu_with_cpu_backend_matches_no_backend() { let weights = make_test_weights(); let input = h(2, weights.hidden_size); - let (h_no, _, _) = run_attention_block_gpu(&weights, &input, 0, false, None).unwrap(); - let (h_cpu, _, _) = - run_attention_block_gpu(&weights, &input, 0, false, Some(&crate::CpuBackend)).unwrap(); + let (h_no, _, _) = run_attention_block_gpu( + larql_models::WeightsView::dense(&weights), + &input, + 0, + false, + None, + ) + .unwrap(); + let (h_cpu, _, _) = run_attention_block_gpu( + larql_models::WeightsView::dense(&weights), + &input, + 0, + false, + Some(&crate::CpuBackend), + ) + .unwrap(); for (a, b) in h_no.iter().zip(h_cpu.iter()) { assert!( (a - b).abs() < 1e-4, @@ -367,7 +420,14 @@ mod tests { fn run_attention_block_gpu_capture_attention_returns_weights() { let weights = make_test_weights(); let input = h(3, weights.hidden_size); - let (_, _, attn_w) = run_attention_block_gpu(&weights, &input, 0, true, None).unwrap(); + let (_, _, attn_w) = run_attention_block_gpu( + larql_models::WeightsView::dense(&weights), + &input, + 0, + true, + None, + ) + .unwrap(); let aw = attn_w.expect("capture=true must yield weights"); assert_eq!(aw.heads.len(), weights.num_q_heads); } @@ -377,8 +437,14 @@ mod tests { let weights = make_test_weights(); let input = h(2, weights.hidden_size); for layer in 0..weights.num_layers { - let (h_out, _, _) = - run_attention_block_gpu(&weights, &input, layer, false, None).unwrap(); + let (h_out, _, _) = run_attention_block_gpu( + larql_models::WeightsView::dense(&weights), + &input, + layer, + false, + None, + ) + .unwrap(); assert!( h_out.iter().all(|v| v.is_finite()), "layer {layer} non-finite" @@ -391,7 +457,14 @@ mod tests { let weights = make_test_weights(); let input = h(2, weights.hidden_size); let bogus = weights.num_layers + 5; - assert!(run_attention_block_gpu(&weights, &input, bogus, false, None).is_none()); + assert!(run_attention_block_gpu( + larql_models::WeightsView::dense(&weights), + &input, + bogus, + false, + None + ) + .is_none()); } #[test] @@ -412,7 +485,14 @@ mod tests { // rms_norm_heads branches and post-norm residual path. let weights = larql_models::test_fixtures::make_gemma3_test_weights(); let input = h(2, weights.hidden_size); - let (h_post, _, _) = run_attention_block_gpu(&weights, &input, 0, false, None).unwrap(); + let (h_post, _, _) = run_attention_block_gpu( + larql_models::WeightsView::dense(&weights), + &input, + 0, + false, + None, + ) + .unwrap(); assert_eq!(h_post.shape(), &[2, weights.hidden_size]); assert!(h_post.iter().all(|v| v.is_finite())); } @@ -422,7 +502,14 @@ mod tests { // Starcoder2 has Q/K/V/O bias keys → all four `add_bias` arms fire. let weights = larql_models::test_fixtures::make_starcoder2_test_weights(); let input = h(2, weights.hidden_size); - let (h_post, _, _) = run_attention_block_gpu(&weights, &input, 0, false, None).unwrap(); + let (h_post, _, _) = run_attention_block_gpu( + larql_models::WeightsView::dense(&weights), + &input, + 0, + false, + None, + ) + .unwrap(); assert_eq!(h_post.shape(), &[2, weights.hidden_size]); assert!(h_post.iter().all(|v| v.is_finite())); } @@ -434,7 +521,14 @@ mod tests { // `run_attention_block_gpu`). let weights = larql_models::test_fixtures::make_gemma3_test_weights(); let input = h(2, weights.hidden_size); - let (h_out, k, v) = run_attention_with_kv_backend(&weights, &input, 0, None).unwrap(); + let (h_out, k, v) = run_attention_with_kv_backend( + larql_models::WeightsView::dense(&weights), + &input, + 0, + None, + None, + ) + .unwrap(); assert_eq!(h_out.shape(), &[2, weights.hidden_size]); let kv_dim = weights.num_kv_heads * weights.head_dim; assert_eq!(k.shape(), &[2, kv_dim]); @@ -458,11 +552,20 @@ mod tests { let weights = larql_models::test_fixtures::make_gemma3_rope_scaled_test_weights(); let input = h(5, weights.hidden_size); for layer in 0..weights.num_layers { - let engine = - run_attention_with_kv_backend(&weights, &input, layer, Some(&crate::CpuBackend)) - .expect("engine attention"); + let engine = run_attention_with_kv_backend( + larql_models::WeightsView::dense(&weights), + &input, + layer, + Some(&crate::CpuBackend), + None, + ) + .expect("engine attention"); let recompute = crate::attention::block::run_attention_block_with_kv_out( - &weights, &input, layer, false, None, + larql_models::WeightsView::dense(&weights), + &input, + layer, + false, + None, ) .expect("full-recompute attention"); for (a, b) in engine.0.iter().zip(recompute.0.iter()) { @@ -480,7 +583,14 @@ mod tests { // `run_attention_with_kv_backend`. let weights = larql_models::test_fixtures::make_starcoder2_test_weights(); let input = h(2, weights.hidden_size); - let (h_out, _, _) = run_attention_with_kv_backend(&weights, &input, 0, None).unwrap(); + let (h_out, _, _) = run_attention_with_kv_backend( + larql_models::WeightsView::dense(&weights), + &input, + 0, + None, + None, + ) + .unwrap(); assert_eq!(h_out.shape(), &[2, weights.hidden_size]); assert!(h_out.iter().all(|x| x.is_finite())); } @@ -492,9 +602,22 @@ mod tests { // post-norm + QK-norm branches. let weights = larql_models::test_fixtures::make_gemma3_test_weights(); let input = h(2, weights.hidden_size); - let (h_no, _, _) = run_attention_with_kv_backend(&weights, &input, 0, None).unwrap(); - let (h_cpu, _, _) = - run_attention_with_kv_backend(&weights, &input, 0, Some(&crate::CpuBackend)).unwrap(); + let (h_no, _, _) = run_attention_with_kv_backend( + larql_models::WeightsView::dense(&weights), + &input, + 0, + None, + None, + ) + .unwrap(); + let (h_cpu, _, _) = run_attention_with_kv_backend( + larql_models::WeightsView::dense(&weights), + &input, + 0, + Some(&crate::CpuBackend), + None, + ) + .unwrap(); for (a, b) in h_no.iter().zip(h_cpu.iter()) { assert!((a - b).abs() < 1e-4, "diverged: {a} vs {b}"); } @@ -525,9 +648,22 @@ mod tests { fn run_attention_with_kv_backend_matches_no_backend() { let weights = make_test_weights(); let input = h(2, weights.hidden_size); - let (h_no, k_no, v_no) = run_attention_with_kv_backend(&weights, &input, 0, None).unwrap(); - let (h_cpu, k_cpu, v_cpu) = - run_attention_with_kv_backend(&weights, &input, 0, Some(&crate::CpuBackend)).unwrap(); + let (h_no, k_no, v_no) = run_attention_with_kv_backend( + larql_models::WeightsView::dense(&weights), + &input, + 0, + None, + None, + ) + .unwrap(); + let (h_cpu, k_cpu, v_cpu) = run_attention_with_kv_backend( + larql_models::WeightsView::dense(&weights), + &input, + 0, + Some(&crate::CpuBackend), + None, + ) + .unwrap(); for (a, b) in h_no.iter().zip(h_cpu.iter()) { assert!((a - b).abs() < 1e-4); } diff --git a/crates/larql-compute/src/backend/quant_matvec.rs b/crates/larql-compute/src/backend/quant_matvec.rs index 3650df351..1391756a5 100644 --- a/crates/larql-compute/src/backend/quant_matvec.rs +++ b/crates/larql-compute/src/backend/quant_matvec.rs @@ -15,6 +15,7 @@ //! Adding a new quant format = `QuantFormat` variant + match arm in //! `quant_matvec` + per-format helper for the fast path. +use crate::cpu::ops::ternary_matvec::BitLinearWeight; use crate::QuantFormat; use larql_models::quant::ggml::LEGACY_BLOCK_ELEMS; @@ -74,6 +75,11 @@ pub trait QuantMatVec { let (q8_x, q8_scales) = crate::cpu::ops::q4_common::quantize_to_q8(x); self.q8_matvec(weights, &q8_x, &q8_scales, num_rows, hidden) } + // Ternary is served by [`Self::ternary_matvec`] with a + // `BitLinearWeight` — the per-channel scales can't ride this + // `&[u8]`-only signature, so `quant_matvec` returns `None` (loud + // missing capability), same as Q8_0's dedicated-kernel story. + QuantFormat::I2S => None, QuantFormat::BF16 | QuantFormat::F16 | QuantFormat::F32 => None, } } @@ -112,6 +118,10 @@ pub trait QuantMatVec { let x_f32 = dequantise_q8(q8_x, q8_scales); self.quant_matvec(format, weights, &x_f32, num_rows, hidden) } + // Ternary has its own int8-activation (A8) quantisation and weight + // container — it doesn't consume the legacy block-32 Q8 input. + // Use [`Self::ternary_matvec`] with a `BitLinearWeight`. + QuantFormat::I2S => None, QuantFormat::BF16 | QuantFormat::F16 | QuantFormat::F32 => None, } } @@ -292,6 +302,20 @@ pub trait QuantMatVec { None } + /// BitNet ternary (I2_S) matvec: `out[N] = W[N, K] · x[K]` where `W` is a + /// 1.58-bit ternary [`BitLinearWeight`] (packed trits + per-channel + /// scales) and `x` is f32. This is the dedicated ternary entry point — + /// `quant_matvec` returns `None` for `QuantFormat::I2S` because the + /// per-channel scales can't ride its `&[u8]`-only signature. + /// + /// Backends with a ternary kernel (the CPU A8 sign-select path) override; + /// the default returns `None`. The implementation owns the activation + /// quantisation (W1.58·A8: int8-quantise `x` internally), so callers pass + /// raw f32 — mirroring how `q4k_matvec` takes f32 input. + fn ternary_matvec(&self, _w: &BitLinearWeight, _x: &[f32]) -> Option> { + None + } + /// Whether this backend implements quantised matvec on `format`. /// /// "Implements" means at minimum a single-matrix matvec path for the diff --git a/crates/larql-compute/src/cpu/mod.rs b/crates/larql-compute/src/cpu/mod.rs index 9315d5941..597dec144 100644 --- a/crates/larql-compute/src/cpu/mod.rs +++ b/crates/larql-compute/src/cpu/mod.rs @@ -92,6 +92,23 @@ impl QuantMatVec for CpuBackend { Some(out) } + fn q4k_matmul( + &self, + q4k_data: &[u8], + x: &[f32], + num_rows: usize, + hidden: usize, + seq_len: usize, + ) -> Option> { + // Amortised Q4_K matmul: decode each weight super-block to f32 once + // and reuse it across all `seq_len` activation columns — reads the + // Q4_K weight a single time, vs `seq_len ×` for repeated + // `q4k_matvec` or 4× the bytes for the dequant→sgemm prefill path. + let mut out = vec![0.0f32; seq_len * num_rows]; + ops::q4_common::q4k_matmul_into(&mut out, x, q4k_data, num_rows, hidden, seq_len); + Some(out) + } + fn q6k_matvec( &self, q6k_data: &[u8], @@ -118,11 +135,28 @@ impl QuantMatVec for CpuBackend { Some((out_a, out_b)) } + fn ternary_matvec( + &self, + w: &ops::ternary_matvec::BitLinearWeight, + x: &[f32], + ) -> Option> { + // Best-available A8 kernel for the target (NEON sign-select on + // aarch64, scalar int8 elsewhere). The kernel int8-quantises `x` + // internally (W1.58·A8). Shape errors surface as `None`. + let mut y = vec![0.0f32; w.rows]; + ops::ternary_matvec::matvec_i2s_a8_f32_into(w, x, &mut y).ok()?; + Some(y) + } + fn supports_quant(&self, format: crate::QuantFormat) -> bool { use crate::QuantFormat; matches!( format, - QuantFormat::Q4_0 | QuantFormat::Q4_K | QuantFormat::Q4_KF | QuantFormat::Q6_K + QuantFormat::Q4_0 + | QuantFormat::Q4_K + | QuantFormat::Q4_KF + | QuantFormat::Q6_K + | QuantFormat::I2S ) } } @@ -188,6 +222,51 @@ mod cpu_backend_tests { // ── QuantMatVec ───────────────────────────────────────────────────── + #[test] + fn cpu_backend_ternary_matvec_matches_direct_kernel() { + use crate::cpu::ops::ternary_matvec::{matvec_i2s_a8_f32_into, BitLinearWeight}; + use crate::QuantFormat; + + // CpuBackend advertises and serves the ternary (I2_S) format. + assert!(CpuBackend.supports_quant(QuantFormat::I2S)); + + let (rows, cols) = (6usize, 64usize); + // Deterministic ternary trits packed 4/byte (codes 0/1/2 → 0/+1/-1). + let mut s = 0x1234_5678u64; + let mut bytes = vec![0u8; rows * cols / 4]; + for b in bytes.iter_mut() { + let mut bv = 0u8; + for slot in 0..4 { + s = s.wrapping_mul(6364136223846793005).wrapping_add(1); + let code = ((s >> 33) % 3) as u8; // 0,1,2 + bv |= code << (2 * slot); + } + *b = bv; + } + let scales: Vec = (0..rows).map(|r| 0.1 + 0.05 * r as f32).collect(); + let w = BitLinearWeight::new(rows, cols, bytes, scales).unwrap(); + let x: Vec = (0..cols).map(|i| (i as f32 * 0.013).sin()).collect(); + + // Backend dispatch must match calling the kernel directly. + let via_backend = CpuBackend + .ternary_matvec(&w, &x) + .expect("ternary_matvec Some"); + let mut direct = vec![0.0f32; rows]; + matvec_i2s_a8_f32_into(&w, &x, &mut direct).unwrap(); + assert_eq!(via_backend.len(), rows); + for (a, b) in via_backend.iter().zip(&direct) { + assert_eq!(a.to_bits(), b.to_bits(), "backend {a} vs direct {b}"); + } + } + + #[test] + fn cpu_backend_ternary_matvec_shape_error_is_none() { + use crate::cpu::ops::ternary_matvec::BitLinearWeight; + let w = BitLinearWeight::new(2, 8, vec![0u8; 4], vec![1.0, 1.0]).unwrap(); + // x length != cols → kernel errors → None (loud missing capability). + assert!(CpuBackend.ternary_matvec(&w, &[0.0; 4]).is_none()); + } + #[test] fn cpu_backend_q4_matvec_returns_some() { let rows = 4usize; diff --git a/crates/larql-compute/src/cpu/ops/q4_common.rs b/crates/larql-compute/src/cpu/ops/q4_common.rs index 87e932905..4912819c2 100644 --- a/crates/larql-compute/src/cpu/ops/q4_common.rs +++ b/crates/larql-compute/src/cpu/ops/q4_common.rs @@ -658,6 +658,282 @@ fn process_q4k_superblock(w: &[u8], x: &[f32], sum_x: &[f32], row_base: usize, s acc } +/// Decode one Q4_K super-block (256 elements, 144 bytes) of row `row_base` +/// into `wf` as full f32 weight values. Per element the dequant is +/// `d * scale[sb] * q - dmin * min[sb]` — identical arithmetic to +/// [`process_q4k_superblock`], but materialised per element so the decoded +/// weights can be reused across many activation columns instead of folded +/// into a single dot. Nibble packing mirrors the matvec: each of the 4 +/// 32-byte groups holds sub-block `2g` in the low nibble and `2g+1` in the +/// high nibble. +#[inline(always)] +fn decode_q4k_superblock_into(w: &[u8], row_base: usize, sb: usize, wf: &mut [f32; 256]) { + const BLOCK_BYTES: usize = 144; + let block = &w[row_base + sb * BLOCK_BYTES..row_base + (sb + 1) * BLOCK_BYTES]; + let d = f16_to_f32(u16::from_le_bytes([block[0], block[1]])); + let dmin = f16_to_f32(u16::from_le_bytes([block[2], block[3]])); + let p = &block[4..16]; + let mut scales = [0u8; 8]; + let mut mins = [0u8; 8]; + for j in 0..4 { + scales[j] = p[j] & 0x3F; + mins[j] = p[j + 4] & 0x3F; + scales[j + 4] = (p[j + 8] & 0x0F) | ((p[j] >> 6) << 4); + mins[j + 4] = (p[j + 8] >> 4) | ((p[j + 4] >> 6) << 4); + } + let quants = &block[16..144]; + for g in 0..4 { + let sb_lo = 2 * g; + let sb_hi = 2 * g + 1; + let sc_lo = d * scales[sb_lo] as f32; + let sc_hi = d * scales[sb_hi] as f32; + let mn_lo = dmin * mins[sb_lo] as f32; + let mn_hi = dmin * mins[sb_hi] as f32; + let chunk = &quants[g * 32..(g + 1) * 32]; + for i in 0..32 { + wf[sb_lo * 32 + i] = sc_lo * (chunk[i] & 0x0F) as f32 - mn_lo; + wf[sb_hi * 32 + i] = sc_hi * (chunk[i] >> 4) as f32 - mn_hi; + } + } +} + +/// Decode one Q6_K super-block (256 elements, 210 bytes) of row `row_base` into +/// `wf` as full f32 weight values — mirrors [`larql_models::quant::ggml`]'s +/// `dequantize_q6_k` per-block math: a 4-bit low nibble (`ql`) plus a 2-bit high +/// part (`qh`), biased by −32, times the per-16-element int8 scale and the f16 +/// super-block scale. +#[inline(always)] +fn decode_q6k_superblock_into(w: &[u8], row_base: usize, sb: usize, wf: &mut [f32; 256]) { + const BLOCK_BYTES: usize = 210; + let block = &w[row_base + sb * BLOCK_BYTES..row_base + (sb + 1) * BLOCK_BYTES]; + let ql = &block[0..128]; + let qh = &block[128..192]; + let scales = &block[192..208]; + let d = f16_to_f32(u16::from_le_bytes([block[208], block[209]])); + for (j, &sc_byte) in scales.iter().enumerate() { + let sc = d * (sc_byte as i8) as f32; + for i in 0..16 { + let idx = j * 16 + i; + let lo4 = if idx % 2 == 0 { + ql[idx / 2] & 0x0F + } else { + (ql[idx / 2] >> 4) & 0x0F + }; + let hi2 = (qh[idx / 4] >> ((idx % 4) * 2)) & 0x03; + let val = ((lo4 as i32) | ((hi2 as i32) << 4)) - 32; + wf[idx] = sc * val as f32; + } + } +} + +/// Dot of a decoded 256-element f32 weight block against a 256-element f32 +/// activation slice — the inner of the amortised k-quant matmul. Dispatches to +/// NEON on aarch64; portable multi-accumulator fallback elsewhere. +#[inline] +fn dot_256_f32(wf: &[f32; 256], xs: &[f32]) -> f32 { + debug_assert!(xs.len() >= 256); + #[cfg(target_arch = "aarch64")] + { + // SAFETY: NEON is in the aarch64 base ISA; `wf` is 256 f32 and `xs` has + // ≥256 contiguous f32 (the caller slices exactly one super-block). + unsafe { dot_256_f32_neon(wf, xs) } + } + #[cfg(not(target_arch = "aarch64"))] + { + dot_256_f32_scalar(wf, xs) + } +} + +/// Portable reference: 8 independent accumulators so the reduction isn't a +/// scalar fp-add chain (Rust f32 add isn't associative). Also the parity oracle +/// for the NEON path. On aarch64 it's reached only from tests (the NEON path +/// serves the lib), so allow it to be otherwise-unused there. +#[cfg_attr(target_arch = "aarch64", allow(dead_code))] +#[inline] +fn dot_256_f32_scalar(wf: &[f32; 256], xs: &[f32]) -> f32 { + let mut acc = [0.0f32; 8]; + for c in 0..256 / 8 { + for l in 0..8 { + acc[l] += wf[c * 8 + l] * xs[c * 8 + l]; + } + } + acc.iter().sum::() +} + +/// NEON: 8 `float32x4` accumulators (32 elems/iter × 8 iters = 256). Eight +/// independent accumulators keep the FMA units busy across M-series' ~4-cycle +/// FMA latency (one accumulator would serialise at ~25% of peak). +#[cfg(target_arch = "aarch64")] +#[inline] +unsafe fn dot_256_f32_neon(wf: &[f32; 256], xs: &[f32]) -> f32 { + use std::arch::aarch64::*; + let wp = wf.as_ptr(); + let xp = xs.as_ptr(); + let mut a0 = vdupq_n_f32(0.0); + let mut a1 = vdupq_n_f32(0.0); + let mut a2 = vdupq_n_f32(0.0); + let mut a3 = vdupq_n_f32(0.0); + let mut a4 = vdupq_n_f32(0.0); + let mut a5 = vdupq_n_f32(0.0); + let mut a6 = vdupq_n_f32(0.0); + let mut a7 = vdupq_n_f32(0.0); + let mut i = 0usize; + while i < 256 { + a0 = vfmaq_f32(a0, vld1q_f32(wp.add(i)), vld1q_f32(xp.add(i))); + a1 = vfmaq_f32(a1, vld1q_f32(wp.add(i + 4)), vld1q_f32(xp.add(i + 4))); + a2 = vfmaq_f32(a2, vld1q_f32(wp.add(i + 8)), vld1q_f32(xp.add(i + 8))); + a3 = vfmaq_f32(a3, vld1q_f32(wp.add(i + 12)), vld1q_f32(xp.add(i + 12))); + a4 = vfmaq_f32(a4, vld1q_f32(wp.add(i + 16)), vld1q_f32(xp.add(i + 16))); + a5 = vfmaq_f32(a5, vld1q_f32(wp.add(i + 20)), vld1q_f32(xp.add(i + 20))); + a6 = vfmaq_f32(a6, vld1q_f32(wp.add(i + 24)), vld1q_f32(xp.add(i + 24))); + a7 = vfmaq_f32(a7, vld1q_f32(wp.add(i + 28)), vld1q_f32(xp.add(i + 28))); + i += 32; + } + let s = vaddq_f32( + vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)), + vaddq_f32(vaddq_f32(a4, a5), vaddq_f32(a6, a7)), + ); + vaddvq_f32(s) +} + +/// Amortised Q4_K × f32 matmul: `out[s, r] = sum_k W[r, k] * X[s, k]`. +/// +/// `w` is `[rows, hidden]` Q4_K (144-byte super-blocks), `x` is +/// `[seq, hidden]` f32 row-major, `out` is `[seq, rows]` f32 row-major. +/// +/// The win over the alternatives prefill could use: +/// * vs `seq ×` [`q4k_matvec_into`]: each weight super-block is decoded to +/// f32 **once** and reused across all `seq` columns, so the Q4_K bytes +/// are read once instead of `seq` times (the per-position matvec loop +/// re-reads every weight `seq` times — ~50× slower than f32 sgemm at +/// long context). +/// * vs dequant-whole-layer + BLAS sgemm: never materialises the full f32 +/// weight matrix (4× the bytes of Q4_K), so it skips the per-layer +/// dequant that dominates short-prompt prefill. +/// +/// Shape errors (zero dims, `hidden` not a multiple of 256, short `w`) +/// zero the output rather than panic, mirroring [`q4k_matvec_into`]. +pub fn q4k_matmul_into( + out: &mut [f32], + x: &[f32], + w: &[u8], + rows: usize, + hidden: usize, + seq: usize, +) { + kquant_matmul_into( + out, + x, + w, + rows, + hidden, + seq, + 144, + decode_q4k_superblock_into, + ); +} + +/// Amortised Q6_K × f32 matmul — the Q6_K twin of [`q4k_matmul_into`], used for +/// the `down_proj` (Q6_K is the default down format in a q4k vindex). Same +/// `[seq, rows]` output contract; reads the Q6_K weight once instead of +/// dequantising the whole matrix to f32 first. +pub fn q6k_matmul_into( + out: &mut [f32], + x: &[f32], + w: &[u8], + rows: usize, + hidden: usize, + seq: usize, +) { + kquant_matmul_into( + out, + x, + w, + rows, + hidden, + seq, + 210, + decode_q6k_superblock_into, + ); +} + +/// Shared amortised k-quant matmul: `out[s, r] = sum_k W[r, k] * X[s, k]`, +/// `out` is `[seq, rows]` row-major. Each weight super-block is decoded to f32 +/// **once** via `decode` and FMA'd across all `seq` columns, so the quantised +/// weight is read once (not `seq×`) and never fully materialised as f32. q4k +/// and q6k differ only in `block_bytes` and the per-block `decode`. +/// +/// Shape errors (zero dims, `hidden` not a 256-multiple, short `w`) zero the +/// output (defensive, mirroring `q4k_matvec_into`); callers keep the dims valid +/// (the `use_q4k_ffn` gate, the down padding). +#[allow(clippy::too_many_arguments)] +fn kquant_matmul_into( + out: &mut [f32], + x: &[f32], + w: &[u8], + rows: usize, + hidden: usize, + seq: usize, + block_bytes: usize, + decode: impl Fn(&[u8], usize, usize, &mut [f32; 256]) + Sync, +) { + debug_assert_eq!(out.len(), seq * rows); + debug_assert_eq!(x.len(), seq * hidden); + const ELEMS_PER_BLOCK: usize = 256; + // Shape errors zero the output (defensive, mirrors `q4k_matvec_into` — see + // `q4k_matmul_rejects_non_multiple_of_256`); callers keep the dims valid. + if rows == 0 || seq == 0 || hidden == 0 || !hidden.is_multiple_of(ELEMS_PER_BLOCK) { + out.iter_mut().for_each(|v| *v = 0.0); + return; + } + let n_blocks = hidden / ELEMS_PER_BLOCK; + let row_bytes = n_blocks * block_bytes; + if w.len() < rows * row_bytes { + out.iter_mut().for_each(|v| *v = 0.0); + return; + } + + // Accumulate into a row-major scratch `[rows, seq]` (each row's `seq` + // outputs contiguous) so the hot loop is row-parallel and write-local; + // transpose to the `[seq, rows]` contract afterwards. The transpose + // touches `rows * seq` f32 once — cheap next to the matmul, and it + // never re-reads the weights. + let mut tmp = vec![0.0f32; rows * seq]; + const CHUNK_ROWS: usize = 32; + let x_ref = x; + let w_ref = w; + let decode_ref = &decode; + crate::cpu::spin_pool::par_chunks_mut(&mut tmp, CHUNK_ROWS * seq, |chunk_idx, chunk_slots| { + let row_base_chunk = chunk_idx * CHUNK_ROWS; + let mut wf = [0.0f32; ELEMS_PER_BLOCK]; + for local_r in 0..CHUNK_ROWS { + let r = row_base_chunk + local_r; + if r >= rows { + break; + } + let out_row = &mut chunk_slots[local_r * seq..local_r * seq + seq]; + let row_base = r * row_bytes; + for sb in 0..n_blocks { + decode_ref(w_ref, row_base, sb, &mut wf); + let x_off = sb * ELEMS_PER_BLOCK; + for (s, slot) in out_row.iter_mut().enumerate() { + let xs = &x_ref[s * hidden + x_off..s * hidden + x_off + ELEMS_PER_BLOCK]; + // `wf` was decoded once above; dot it against this column. + // NEON (8 f32x4 accumulators to hide FMA latency) on + // aarch64, portable multi-accumulator fallback elsewhere. + *slot += dot_256_f32(&wf, xs); + } + } + } + }); + + for r in 0..rows { + for s in 0..seq { + out[s * rows + r] = tmp[r * seq + s]; + } + } +} + /// Fused two-weight Q4_K matvec sharing one input vector. /// /// `out_a[N] = W_a[N, K] · x[K]`, `out_b[N] = W_b[N, K] · x[K]`. @@ -1523,6 +1799,200 @@ mod tests { assert!(max_diff < 5e-3, "multi-block diverged: max_diff={max_diff}"); } + /// Amortised matmul must match dequantise→matmul for multiple rows AND + /// multiple sequence positions across several super-blocks. Exercises + /// the row-stride, the per-seq accumulation, and the [rows,seq] → + /// [seq,rows] transpose. + #[test] + fn q4k_matmul_matches_dequant_then_matmul() { + let rows = 5; + let hidden = 512; // 2 super-blocks per row + let seq = 4; + let n_elem = rows * hidden; + let weights: Vec = (0..n_elem) + .map(|i| (i as f32 * 0.0007).sin() * 0.9) + .collect(); + let q4k = quantize_q4_k(&weights); + let dequant = dequantize_q4_k(&q4k, n_elem); + + let x: Vec = (0..seq * hidden) + .map(|i| (i as f32 * 0.013).cos() * 0.5) + .collect(); + + // Reference: dequant → row-major matmul, out[s, r]. + let mut reference = vec![0.0f32; seq * rows]; + for s in 0..seq { + for r in 0..rows { + let mut acc = 0.0f32; + for k in 0..hidden { + acc += dequant[r * hidden + k] * x[s * hidden + k]; + } + reference[s * rows + r] = acc; + } + } + + let mut got = vec![0.0f32; seq * rows]; + q4k_matmul_into(&mut got, &x, &q4k, rows, hidden, seq); + + let max_diff: f32 = reference + .iter() + .zip(&got) + .map(|(a, b)| (a - b).abs()) + .fold(0.0, f32::max); + assert!( + max_diff < 5e-3, + "q4k_matmul diverged from dequant→matmul: max_diff={max_diff}" + ); + } + + /// Each sequence row of the matmul must equal the single-vector + /// `q4k_matvec` for that activation. The two kernels share decode + /// arithmetic, so they must agree row-for-row — catches transpose / + /// offset bugs that a dequant reference could mask if both paths were + /// wrong the same way. + #[test] + fn q4k_matmul_rows_match_q4k_matvec() { + let rows = 6; + let hidden = 256; + let seq = 3; + let weights: Vec = (0..rows * hidden) + .map(|i| ((i as f32 * 0.002) - 1.0) * 0.3) + .collect(); + let q4k = quantize_q4_k(&weights); + let x: Vec = (0..seq * hidden) + .map(|i| (i as f32 * 0.017).sin()) + .collect(); + + let mut mm = vec![0.0f32; seq * rows]; + q4k_matmul_into(&mut mm, &x, &q4k, rows, hidden, seq); + + for s in 0..seq { + let mut mv = vec![0.0f32; rows]; + q4k_matvec_into( + &mut mv, + &x[s * hidden..(s + 1) * hidden], + &q4k, + rows, + hidden, + ); + for r in 0..rows { + let diff = (mm[s * rows + r] - mv[r]).abs(); + assert!( + diff < 1e-4, + "matmul row s={s} r={r} != matvec: {} vs {}", + mm[s * rows + r], + mv[r] + ); + } + } + } + + /// Defensive shape guard: `hidden` not a multiple of 256 → zeroed + /// output (mirrors `q4k_matvec_into`). + #[test] + fn q4k_matmul_rejects_non_multiple_of_256() { + let mut out = vec![1.0f32; 2 * 3]; // seq=2, rows=3, pre-filled to detect zeroing + let x = vec![0.5f32; 2 * 100]; + let w = vec![0u8; 3 * 144]; + q4k_matmul_into(&mut out, &x, &w, 3, 100, 2); + assert_eq!(out, vec![0.0f32; 6]); + } + + /// `dot_256_f32` (NEON on aarch64) must match the scalar reference and a + /// plain sequential dot. + #[test] + fn dot_256_f32_matches_reference() { + let wf: [f32; 256] = std::array::from_fn(|i| (i as f32 * 0.013).sin() * 0.5); + let xs: Vec = (0..256).map(|i| (i as f32 * 0.021).cos() * 0.7).collect(); + let reference: f64 = (0..256).map(|i| (wf[i] * xs[i]) as f64).sum(); + let got = dot_256_f32(&wf, &xs); + assert!( + (got as f64 - reference).abs() < 1e-3, + "dot_256_f32 diverged: {got} vs {reference}" + ); + #[cfg(target_arch = "aarch64")] + { + let neon = unsafe { dot_256_f32_neon(&wf, &xs) }; + let scalar = dot_256_f32_scalar(&wf, &xs); + assert!( + (neon - scalar).abs() < 1e-4, + "neon {neon} != scalar {scalar}" + ); + } + } + + /// Multi-chunk (rows > CHUNK_ROWS = 32, spans multiple parallel work + /// units) and seq = 1 — the two q4k_matmul paths the earlier tests missed. + #[test] + fn q4k_matmul_multi_chunk_and_seq1() { + for (rows, seq) in [(40usize, 3usize), (4, 1)] { + let hidden = 512; + let n_elem = rows * hidden; + let weights: Vec = (0..n_elem) + .map(|i| (i as f32 * 0.0005).sin() * 0.6) + .collect(); + let q4k = quantize_q4_k(&weights); + let dequant = dequantize_q4_k(&q4k, n_elem); + let x: Vec = (0..seq * hidden) + .map(|i| (i as f32 * 0.011).cos() * 0.4) + .collect(); + let mut reference = vec![0.0f32; seq * rows]; + for s in 0..seq { + for r in 0..rows { + let mut acc = 0.0f32; + for k in 0..hidden { + acc += dequant[r * hidden + k] * x[s * hidden + k]; + } + reference[s * rows + r] = acc; + } + } + let mut got = vec![0.0f32; seq * rows]; + q4k_matmul_into(&mut got, &x, &q4k, rows, hidden, seq); + let max: f32 = reference + .iter() + .zip(&got) + .map(|(a, b)| (a - b).abs()) + .fold(0.0, f32::max); + assert!(max < 5e-3, "rows={rows} seq={seq} diverged: {max}"); + } + } + + /// Direct `q6k_matmul_into` kernel parity vs dequantise → matmul. + #[test] + fn q6k_matmul_matches_dequant_then_matmul() { + let rows = 5; + let hidden = 512; + let seq = 3; + let n_elem = rows * hidden; + let weights: Vec = (0..n_elem) + .map(|i| (i as f32 * 0.0007).sin() * 0.5) + .collect(); + let q6k = quantize_q6_k(&weights); + let dq = crate::kquant_forward::dequant::dequantize_matrix(&q6k, "Q6_K", rows, hidden); + let dq = dq.as_slice().expect("contiguous"); + let x: Vec = (0..seq * hidden) + .map(|i| (i as f32 * 0.009).cos() * 0.5) + .collect(); + let mut reference = vec![0.0f32; seq * rows]; + for s in 0..seq { + for r in 0..rows { + let mut acc = 0.0f32; + for k in 0..hidden { + acc += dq[r * hidden + k] * x[s * hidden + k]; + } + reference[s * rows + r] = acc; + } + } + let mut got = vec![0.0f32; seq * rows]; + q6k_matmul_into(&mut got, &x, &q6k, rows, hidden, seq); + let max: f32 = reference + .iter() + .zip(&got) + .map(|(a, b)| (a - b).abs()) + .fold(0.0, f32::max); + assert!(max < 5e-3, "q6k_matmul diverged: {max}"); + } + /// Defensive: caller passes a malformed `cols` (not multiple of 256). /// We zero the output rather than reading past the buffer, mirroring /// `dequantize_q4_k`'s `Vec::new()` shape-error contract. diff --git a/crates/larql-compute/src/cpu/ops/ternary_matvec.rs b/crates/larql-compute/src/cpu/ops/ternary_matvec.rs index 4b5d2ea4f..5392f820c 100644 --- a/crates/larql-compute/src/cpu/ops/ternary_matvec.rs +++ b/crates/larql-compute/src/cpu/ops/ternary_matvec.rs @@ -11,18 +11,23 @@ //! per-channel scale) instead of once per *element* (the dense f16/ //! f32 path), which is the entire point of native BitNet inference. //! -//! ## Status: correct reference, not yet the optimized fast path +//! ## Status: A8 (int8-activation) is the production path; NEON live, AVX2 pending //! -//! This is a scalar implementation: it walks the packed trits and -//! does `acc += sign * activation` in f32 (the "no multiplies" -//! property above describes the *algorithm* — the current code still -//! branches/selects per element in plain f32, it is not yet the -//! vectorised form). It is validated against a dequant-and-matmul -//! reference and is correct. The optimized path — int8-quantized -//! activations + a sign-select NEON/AVX2 kernel dispatched through -//! the compute backend — is future work that slots into the -//! `FormatRoute` / quant-registry machinery; this reference kernel -//! is what the BitNet forward uses today. +//! Three matvec paths, all validated against a dequant-and-matmul reference: +//! - [`matvec_i2s_f32`] — f32 activations (the original reference; kept for +//! parity tests). +//! - [`matvec_i2s_q8_into`] — scalar **W1.58·A8**: int8-quantized activations +//! ([`quantize_activation_i8`]) + integer sign-select accumulation. BitNet's +//! intended inference precision (~2.4× the f32 path on its own). +//! - [`matvec_i2s_q8_neon_into`] — NEON sign-select (aarch64): bit-identical +//! to the scalar A8 path, **~12-13× the f32 reference** on BitNet shapes. +//! +//! [`matvec_i2s_a8_into`] dispatches NEON on aarch64 / scalar int8 elsewhere, +//! and the `larql-inference` BitNet forward now runs on it (via +//! [`matvec_i2s_a8_f32_into`]) — validated end-to-end on +//! `microsoft/bitnet-b1.58-2B-4T` (coherent generation; FFN forward tracks the +//! f32 reference within int8 tolerance). Remaining: an AVX2 `_mm256_sign_epi8` +//! twin so x86_64 gets the full SIMD win (it has the scalar A8 ~2.4× today). //! //! For Microsoft's BitNet b1.58 2 B 4 T (`general.architecture = //! "bitnet-b1.58"`) the saving is dramatic: the weight tensor stays @@ -32,14 +37,16 @@ //! layer). Compare to the 5+ GB f16-after-dequant heap profile //! observed in the production triage. //! -//! This module ships the kernel + a typed weight container, -//! validated against a naive dequant-and-matmul reference. Wiring -//! it into the `larql-inference` forward pass for actual BitLinear -//! layers is a separate piece (it requires a vindex-format change to -//! retain the I2_S bytes and per-channel scales rather than -//! materialising f16 at convert-time, plus a forward-dispatch hook -//! that selects this kernel for ternary tensors); both are tracked -//! as follow-up work in `BUG-infer-deadlock.md`. +//! This module ships the kernel + a typed weight container, validated +//! against a naive dequant-and-matmul reference. The wiring that was once +//! tracked as follow-up work has landed: the vindex-format change retains the +//! I2_S bytes + per-channel scales at convert-time +//! (`larql_vindex::extract::bitnet_writer` / `bitnet_loader`, written to a +//! `bitnet/` sidecar), and the `larql-inference` BitNet forward calls these +//! kernels directly via [`matvec_i2s_a8_f32_into`]. What remains is *dispatch* +//! integration: this path is reached by direct function call, not yet through +//! the `QuantFormat` / `FormatRoute` registry (which has no ternary variant +//! today) — see ROADMAP "BitNet b1.58 integration hardening". //! //! ## API //! @@ -222,6 +229,242 @@ pub fn matvec_i2s_f32_into( Ok(()) } +// ── A8 path: int8-activation ternary matvec (BitNet W1.58·A8) ──────────────── +// +// BitNet b1.58 is trained as W1.58·A8 — 1.58-bit ternary weights AND 8-bit +// activations. Quantising the activation to int8 turns the inner product into +// pure integer sign-select accumulation (add at +1, subtract at -1, skip 0) +// with a single scale at the end — no per-element f32 multiply. It also halves +// the activation footprint and is the form the explicit SIMD sign-select +// kernels (NEON/AVX2) consume. This scalar version is the parity reference for +// those; numerically it matches `matvec_i2s_f32_into` up to the int8 +// activation quantisation error (which is the intended inference precision). + +/// Symmetric per-tensor int8 quantisation of an activation vector — the "A8" +/// half of W1.58·A8. `scale = max|x| / 127`; `x_i8[i] = round(x[i] / scale)`, +/// clamped to `[-127, 127]`. Returns `(x_i8, scale)` such that +/// `x[i] ≈ x_i8[i] * scale`. An all-zero input yields `scale = 0` and all-zero +/// codes (the matvec then produces zeros, matching the f32 path). +pub fn quantize_activation_i8(x: &[f32]) -> (Vec, f32) { + let amax = x.iter().fold(0.0f32, |m, &v| m.max(v.abs())); + if amax == 0.0 { + return (vec![0i8; x.len()], 0.0); + } + let scale = amax / 127.0; + let inv = 1.0 / scale; + let q = x + .iter() + .map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8) + .collect(); + (q, scale) +} + +/// Ternary × int8 matvec — the W1.58·A8 path. `x_i8` / `x_scale` come from +/// [`quantize_activation_i8`] (quantise the activation once, reuse across the +/// Q/K/V/O or gate/up projections that share it). Integer sign-select +/// accumulation in i32; the only float work is the per-row +/// `channel_scales[r] * x_scale` applied at the very end. +/// +/// # Errors +/// `ComputeError::ShapeMismatch` if `x_i8.len() != w.cols` or `y.len() < w.rows`. +pub fn matvec_i2s_q8_into( + w: &BitLinearWeight, + x_i8: &[i8], + x_scale: f32, + y: &mut [f32], +) -> Result<(), ComputeError> { + if x_i8.len() != w.cols { + return Err(ComputeError::ShapeMismatch(format!( + "matvec_i2s_q8: x_i8.len() = {}, expected w.cols = {}", + x_i8.len(), + w.cols + ))); + } + if y.len() < w.rows { + return Err(ComputeError::ShapeMismatch(format!( + "matvec_i2s_q8: y.len() = {} < w.rows = {}", + y.len(), + w.rows + ))); + } + + let row_bytes = w.row_bytes(); + debug_assert_eq!(row_bytes * 4, w.cols); + + for (r, y_r) in y.iter_mut().enumerate().take(w.rows) { + let row = &w.i2s_bytes[r * row_bytes..(r + 1) * row_bytes]; + // code 0/3 → 0, 1 → +x, 2 → -x. Branch-free sign LUT; i32 accumulate. + const SIGN: [i32; 4] = [0, 1, -1, 0]; + let mut acc: i32 = 0; + for (b, &byte) in row.iter().enumerate() { + let base = b * 4; + acc += SIGN[(byte & 0b11) as usize] * x_i8[base] as i32; + acc += SIGN[((byte >> 2) & 0b11) as usize] * x_i8[base + 1] as i32; + acc += SIGN[((byte >> 4) & 0b11) as usize] * x_i8[base + 2] as i32; + acc += SIGN[((byte >> 6) & 0b11) as usize] * x_i8[base + 3] as i32; + } + *y_r = acc as f32 * w.channel_scales[r] * x_scale; + } + + Ok(()) +} + +/// Convenience: quantise `x` to int8 and run the A8 matvec, returning a fresh +/// `Vec`. Prefer [`quantize_activation_i8`] + [`matvec_i2s_q8_into`] +/// directly when one activation feeds several weight matrices. +/// +/// # Errors +/// `ComputeError::ShapeMismatch` if `x.len() != w.cols`. +pub fn matvec_i2s_a8(w: &BitLinearWeight, x: &[f32]) -> Result, ComputeError> { + let (x_i8, x_scale) = quantize_activation_i8(x); + let mut y = vec![0.0f32; w.rows]; + matvec_i2s_q8_into(w, &x_i8, x_scale, &mut y)?; + Ok(y) +} + +/// Drop-in A8 replacement for [`matvec_i2s_f32_into`] (same `(w, x_f32, y)` +/// signature): quantises `x` to int8 internally, then runs the best-available +/// A8 kernel. The internal quantise is `O(cols)` — negligible next to the +/// `O(rows·cols)` matvec. When one activation feeds several matrices (Q/K/V, +/// gate/up), [`quantize_activation_i8`] + [`matvec_i2s_a8_into`] saves the +/// repeat quantise, but for a single matrix this is the convenient form. +/// +/// # Errors +/// `ComputeError::ShapeMismatch` if `x.len() != w.cols` or `y.len() < w.rows`. +#[inline] +pub fn matvec_i2s_a8_f32_into( + w: &BitLinearWeight, + x: &[f32], + y: &mut [f32], +) -> Result<(), ComputeError> { + let (x_i8, x_scale) = quantize_activation_i8(x); + matvec_i2s_a8_into(w, &x_i8, x_scale, y) +} + +/// Best-available A8 ternary matvec for the current target: NEON sign-select +/// on aarch64, scalar int8 elsewhere (an AVX2 twin is the x86_64 follow-up). +/// `x_i8` / `x_scale` come from [`quantize_activation_i8`] — quantise the +/// activation once and feed it to every weight matrix that shares it +/// (Q/K/V/O, gate/up). +/// +/// # Errors +/// `ComputeError::ShapeMismatch` if `x_i8.len() != w.cols` or `y.len() < w.rows`. +#[inline] +pub fn matvec_i2s_a8_into( + w: &BitLinearWeight, + x_i8: &[i8], + x_scale: f32, + y: &mut [f32], +) -> Result<(), ComputeError> { + #[cfg(target_arch = "aarch64")] + { + matvec_i2s_q8_neon_into(w, x_i8, x_scale, y) + } + #[cfg(not(target_arch = "aarch64"))] + { + matvec_i2s_q8_into(w, x_i8, x_scale, y) + } +} + +// ── A8 path: NEON sign-select (aarch64) ───────────────────────────────────── + +/// NEON (aarch64) implementation of [`matvec_i2s_q8_into`]. Decodes 16 trits +/// per iteration and accumulates `±x` via masked add/subtract in SIMD; the +/// `cols % 16` tail runs the scalar sign-LUT. Because the inner product is +/// pure integer sign-select, the result is **bit-identical** to +/// [`matvec_i2s_q8_into`] (integer summation is order-independent) — verified +/// in tests. This is the path that turns the algorithm's "no multiplies" +/// property into actual throughput. +/// +/// # Errors +/// `ComputeError::ShapeMismatch` if `x_i8.len() != w.cols` or `y.len() < w.rows`. +#[cfg(target_arch = "aarch64")] +pub fn matvec_i2s_q8_neon_into( + w: &BitLinearWeight, + x_i8: &[i8], + x_scale: f32, + y: &mut [f32], +) -> Result<(), ComputeError> { + if x_i8.len() != w.cols { + return Err(ComputeError::ShapeMismatch(format!( + "matvec_i2s_q8_neon: x_i8.len() = {}, expected w.cols = {}", + x_i8.len(), + w.cols + ))); + } + if y.len() < w.rows { + return Err(ComputeError::ShapeMismatch(format!( + "matvec_i2s_q8_neon: y.len() = {} < w.rows = {}", + y.len(), + w.rows + ))); + } + + let row_bytes = w.row_bytes(); + for (r, y_r) in y.iter_mut().enumerate().take(w.rows) { + let row = &w.i2s_bytes[r * row_bytes..(r + 1) * row_bytes]; + // SAFETY: NEON is baseline on aarch64; indices stay in bounds + // (`row` is `cols/4` bytes, `x_i8` is `cols` long, loop chunks of 16). + let acc = unsafe { i2s_row_dot_q8_neon(row, x_i8, w.cols) }; + *y_r = acc as f32 * w.channel_scales[r] * x_scale; + } + Ok(()) +} + +/// One row's integer dot via NEON sign-select. Returns `Σ sign(trit)·x_i8`. +#[cfg(target_arch = "aarch64")] +#[target_feature(enable = "neon")] +unsafe fn i2s_row_dot_q8_neon(row: &[u8], x_i8: &[i8], cols: usize) -> i32 { + use std::arch::aarch64::*; + + // Duplicate each of the 4 source bytes across 4 lanes, then right-shift + // each lane by {0,2,4,6} and mask to recover the 16 per-element 2-bit codes. + let dup_idx_arr: [u8; 16] = [0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]; + let shift_arr: [i8; 16] = [0, -2, -4, -6, 0, -2, -4, -6, 0, -2, -4, -6, 0, -2, -4, -6]; + let dup_idx = vld1q_u8(dup_idx_arr.as_ptr()); + let shifts = vld1q_s8(shift_arr.as_ptr()); + let mask3 = vdupq_n_u8(0b11); + let one = vdupq_n_u8(1); + let two = vdupq_n_u8(2); + + let mut acc = vdupq_n_s32(0); + let n16 = cols / 16; + for chunk in 0..n16 { + let c = chunk * 16; + let b = chunk * 4; + // Load the 4 packed bytes (16 trits) into lanes 0..4. + let word = u32::from_le_bytes([row[b], row[b + 1], row[b + 2], row[b + 3]]); + let reg = vreinterpretq_u8_u32(vsetq_lane_u32::<0>(word, vdupq_n_u32(0))); + let bytes_dup = vqtbl1q_u8(reg, dup_idx); + let codes = vandq_u8(vshlq_u8(bytes_dup, shifts), mask3); + // sign-select: +x where code==1, -x where code==2, 0 otherwise. + let plus = vceqq_u8(codes, one); + let minus = vceqq_u8(codes, two); + let x = vld1q_s8(x_i8.as_ptr().add(c)); + let pos = vandq_s8(x, vreinterpretq_s8_u8(plus)); + let neg = vandq_s8(x, vreinterpretq_s8_u8(minus)); + let contrib = vsubq_s8(pos, neg); + // Widen + accumulate into i32 (no overflow: i32 holds the full sum). + acc = vpadalq_s16(acc, vpaddlq_s8(contrib)); + } + let mut total = vaddvq_s32(acc); + + // Scalar tail for the `cols % 16` remainder (a multiple of 4). + const SIGN: [i32; 4] = [0, 1, -1, 0]; + let mut c = n16 * 16; + let mut b = n16 * 4; + while c < cols { + let byte = row[b]; + total += SIGN[(byte & 0b11) as usize] * x_i8[c] as i32; + total += SIGN[((byte >> 2) & 0b11) as usize] * x_i8[c + 1] as i32; + total += SIGN[((byte >> 4) & 0b11) as usize] * x_i8[c + 2] as i32; + total += SIGN[((byte >> 6) & 0b11) as usize] * x_i8[c + 3] as i32; + c += 4; + b += 1; + } + total +} + // ── Tests ──────────────────────────────────────────────────────────────────── #[cfg(test)] @@ -484,4 +727,146 @@ mod tests { kernel[0] - reference ); } + + // ── A8 (int8-activation) path ─────────────────────────────────────────── + + 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 == 0.0 || nb == 0.0 { + return 1.0; + } + dot / (na * nb) + } + + #[test] + fn quantize_activation_i8_zero_input_yields_zero_scale() { + let (q, scale) = quantize_activation_i8(&[0.0; 16]); + assert_eq!(scale, 0.0); + assert!(q.iter().all(|&c| c == 0)); + } + + #[test] + fn quantize_activation_i8_puts_max_at_127() { + let (q, scale) = quantize_activation_i8(&[0.5, -1.0, 0.25, 0.0]); + // amax = 1.0 → scale = 1/127; the -1.0 maps to -127. + assert!((scale - 1.0 / 127.0).abs() < 1e-9); + assert_eq!(q[1], -127); + // round-trip stays within one quantum. + for (&qc, &orig) in q.iter().zip([0.5, -1.0, 0.25, 0.0].iter()) { + assert!((qc as f32 * scale - orig).abs() <= scale); + } + } + + #[test] + fn matvec_a8_is_exact_when_activation_is_int8_representable() { + // Activations drawn from {0, +a, -a} all quantise to {0, ±127} + // exactly (amax == a), so the A8 path reproduces the f32 path with + // no quantisation error — only fp rounding. + let (rows, cols) = (8usize, 256usize); + let a = 0.5f32; + let x: Vec = synth_ternary(cols, 11).iter().map(|t| t * a).collect(); + let mut bytes = Vec::new(); + let mut scales = Vec::new(); + for r in 0..rows { + bytes.extend(encode_row(&synth_ternary(cols, 100 + r as u64), 1.0)); + scales.push(0.3 + 0.1 * r as f32); + } + let w = BitLinearWeight::new(rows, cols, bytes, scales).unwrap(); + + let y_f32 = matvec_i2s_f32(&w, &x).unwrap(); + let y_a8 = matvec_i2s_a8(&w, &x).unwrap(); + for (f, q) in y_f32.iter().zip(y_a8.iter()) { + assert!((f - q).abs() < 1e-3, "f32={f} a8={q}"); + } + } + + #[test] + fn matvec_a8_matches_f32_within_int8_tolerance() { + // Arbitrary real activations: the A8 path carries int8 activation + // quantisation error but must stay tightly aligned with the f32 path. + let (rows, cols) = (8usize, 512usize); + let x = synth(cols, 19); + let mut bytes = Vec::new(); + let mut scales = Vec::new(); + for r in 0..rows { + bytes.extend(encode_row(&synth_ternary(cols, 200 + r as u64), 1.0)); + scales.push(0.5 + 0.05 * r as f32); + } + let w = BitLinearWeight::new(rows, cols, bytes, scales).unwrap(); + + let y_f32 = naive_dequant_matvec(&w, &x); + let y_a8 = matvec_i2s_a8(&w, &x).unwrap(); + + let cos = cosine(&y_f32, &y_a8); + assert!(cos > 0.999, "A8 vs f32 cosine {cos} below 0.999"); + // relative L2 error from int8 activation quantisation stays small. + let err: f32 = y_f32 + .iter() + .zip(&y_a8) + .map(|(f, q)| (f - q) * (f - q)) + .sum::() + .sqrt(); + let mag: f32 = y_f32.iter().map(|f| f * f).sum::().sqrt(); + assert!( + err / mag < 0.03, + "A8 relative L2 error {} too high", + err / mag + ); + } + + #[test] + fn matvec_a8_zero_weight_returns_zero() { + let w = BitLinearWeight::new(2, 8, vec![0u8; 4], vec![1.0, 1.0]).unwrap(); + let y = matvec_i2s_a8(&w, &synth(8, 5)).unwrap(); + assert_eq!(y, vec![0.0, 0.0]); + } + + #[test] + fn matvec_i2s_q8_into_rejects_bad_shapes() { + let w = BitLinearWeight::new(2, 8, vec![0u8; 4], vec![1.0, 1.0]).unwrap(); + let (x_i8, s) = quantize_activation_i8(&synth(8, 5)); + // wrong activation length + assert!(matvec_i2s_q8_into(&w, &x_i8[..4], s, &mut [0.0; 2]).is_err()); + // output too small + assert!(matvec_i2s_q8_into(&w, &x_i8, s, &mut [0.0; 1]).is_err()); + } + + #[cfg(target_arch = "aarch64")] + #[test] + fn matvec_q8_neon_is_bit_identical_to_scalar() { + // Integer sign-select accumulation is order-independent, so the NEON + // kernel must match the scalar A8 path bit-for-bit — across full + // 16-wide chunks AND `cols % 16 != 0` tails (20, 36, 260). + for &cols in &[16usize, 64, 256, 512, 20, 36, 260] { + let rows = 6usize; + let mut bytes = Vec::new(); + let mut scales = Vec::new(); + for r in 0..rows { + bytes.extend(encode_row(&synth_ternary(cols, 300 + r as u64), 1.0)); + scales.push(0.4 + 0.07 * r as f32); + } + let w = BitLinearWeight::new(rows, cols, bytes, scales).unwrap(); + let (x_i8, x_scale) = quantize_activation_i8(&synth(cols, 77)); + + let mut y_scalar = vec![0.0f32; rows]; + matvec_i2s_q8_into(&w, &x_i8, x_scale, &mut y_scalar).unwrap(); + let mut y_neon = vec![0.0f32; rows]; + matvec_i2s_q8_neon_into(&w, &x_i8, x_scale, &mut y_neon).unwrap(); + + for (s, n) in y_scalar.iter().zip(&y_neon) { + assert_eq!(s.to_bits(), n.to_bits(), "cols={cols} scalar={s} neon={n}"); + } + } + } + + #[cfg(target_arch = "aarch64")] + #[test] + fn matvec_q8_neon_rejects_bad_shapes() { + let w = BitLinearWeight::new(2, 8, vec![0u8; 4], vec![1.0, 1.0]).unwrap(); + let (x_i8, s) = quantize_activation_i8(&synth(8, 5)); + assert!(matvec_i2s_q8_neon_into(&w, &x_i8[..4], s, &mut [0.0; 2]).is_err()); + assert!(matvec_i2s_q8_neon_into(&w, &x_i8, s, &mut [0.0; 1]).is_err()); + } } diff --git a/crates/larql-compute/src/ffn.rs b/crates/larql-compute/src/ffn.rs index fb0d4ed68..500781cb1 100644 --- a/crates/larql-compute/src/ffn.rs +++ b/crates/larql-compute/src/ffn.rs @@ -9,7 +9,10 @@ pub mod weight; -pub use weight::{dense_ffn_forward, dense_ffn_forward_backend, BackendFfn, NullFfn, WeightFfn}; +pub use weight::{ + dense_ffn_forward, dense_ffn_forward_backend, BackendFfn, NullFfn, Q4kMatmulFfn, ViewFfn, + WeightFfn, +}; use ndarray::Array2; diff --git a/crates/larql-compute/src/ffn/weight.rs b/crates/larql-compute/src/ffn/weight.rs index 3e9eb0aa3..40e6ddf19 100644 --- a/crates/larql-compute/src/ffn/weight.rs +++ b/crates/larql-compute/src/ffn/weight.rs @@ -6,7 +6,56 @@ use ndarray::Array2; use super::{gelu_tanh, gelu_tanh_gate_up, sigmoid, silu_gate_up, FfnBackend}; use crate::forward::add_bias; -use larql_models::ModelWeights; +use larql_models::{ModelWeights, WeightsView}; + +/// Amortised quant matmul for formats with a direct CPU kernel (Q4_K / Q6_K): +/// `x[seq, cols] -> [seq, rows]`, reading the quantised bytes once and never +/// materialising the full f32 weight. Returns `None` for a format without a +/// kernel so the caller can dequantise + matmul. Shared by the q4k-direct FFN +/// and the attention Q/K/V/O projections. +pub(crate) fn quant_matmul( + bytes: &[u8], + fmt: &str, + x: &[f32], + rows: usize, + cols: usize, + seq: usize, +) -> Option> { + let mut out = vec![0.0f32; seq * rows]; + match fmt { + "Q4_K" => crate::cpu::ops::q4_common::q4k_matmul_into(&mut out, x, bytes, rows, cols, seq), + "Q6_K" => crate::cpu::ops::q4_common::q6k_matmul_into(&mut out, x, bytes, rows, cols, seq), + _ => return None, + } + Some(Array2::from_shape_vec((seq, rows), out).expect("quant_matmul output shape [seq, rows]")) +} + +/// Project `x[seq, in_dim] -> [seq, out_rows]` from quantised weight bytes; +/// `in_dim` must be a 256-multiple. Q4_K/Q6_K read the bytes once via the +/// amortised kernel; any other format dequantises then matmuls. Used for the +/// FFN gate/up and the attention Q/K/V/O projections. +pub(crate) fn quant_proj( + bytes: &[u8], + fmt: &str, + x: &Array2, + out_rows: usize, + in_dim: usize, + seq: usize, +) -> Array2 { + if let Some(r) = quant_matmul( + bytes, + fmt, + x.as_slice().expect("contiguous x"), + out_rows, + in_dim, + seq, + ) { + r + } else { + let w = crate::kquant_forward::dequant::dequantize_matrix(bytes, fmt, out_rows, in_dim); + dot_proj_gpu(x, &w, None) + } +} /// Dense FFN: follows the model architecture exactly (CPU BLAS). /// Gated: activation(x @ gate.T) * (x @ up.T) @ down.T + bias @@ -17,11 +66,11 @@ pub struct WeightFfn<'a> { impl<'a> FfnBackend for WeightFfn<'a> { fn forward(&self, layer: usize, x: &Array2) -> Array2 { - dense_ffn_forward(self.weights, layer, x).0 + dense_ffn_forward(WeightsView::dense(self.weights), layer, x).0 } fn forward_with_activation(&self, layer: usize, x: &Array2) -> (Array2, Array2) { - dense_ffn_forward(self.weights, layer, x) + dense_ffn_forward(WeightsView::dense(self.weights), layer, x) } fn name(&self) -> &str { @@ -29,6 +78,29 @@ impl<'a> FfnBackend for WeightFfn<'a> { } } +/// FFN backend over a [`WeightsView`] — the quant forward's scratch-aware FFN. +/// Identical math to [`WeightFfn`], but resolves gate/up/down through the view +/// (engine scratch first, then canonical), so the per-layer dequantised FFN +/// tensors are visible without mutating `ModelWeights`. The quant forward loops +/// construct this with a `with_scratch` view; everything else keeps `WeightFfn`. +pub struct ViewFfn<'a> { + pub view: WeightsView<'a>, +} + +impl FfnBackend for ViewFfn<'_> { + fn forward(&self, layer: usize, x: &Array2) -> Array2 { + dense_ffn_forward(self.view, layer, x).0 + } + + fn forward_with_activation(&self, layer: usize, x: &Array2) -> (Array2, Array2) { + dense_ffn_forward(self.view, layer, x) + } + + fn name(&self) -> &str { + "view" + } +} + /// Backend-dispatched dense FFN. Matmuls route through `ComputeBackend` when /// `backend` is `Some` — useful for prefill on Metal where gate/up/down /// projections are the dominant cost. @@ -39,11 +111,22 @@ pub struct BackendFfn<'a, 'b> { impl<'a, 'b> FfnBackend for BackendFfn<'a, 'b> { fn forward(&self, layer: usize, x: &Array2) -> Array2 { - dense_ffn_forward_backend(self.weights, layer, x, Some(self.backend)).0 + dense_ffn_forward_backend( + WeightsView::dense(self.weights), + layer, + x, + Some(self.backend), + ) + .0 } fn forward_with_activation(&self, layer: usize, x: &Array2) -> (Array2, Array2) { - dense_ffn_forward_backend(self.weights, layer, x, Some(self.backend)) + dense_ffn_forward_backend( + WeightsView::dense(self.weights), + layer, + x, + Some(self.backend), + ) } fn name(&self) -> &str { @@ -51,6 +134,150 @@ impl<'a, 'b> FfnBackend for BackendFfn<'a, 'b> { } } +/// FFN backend that runs gate/up/down on the vindex's quantised weight bytes +/// **without a per-layer full dequant**. Each component is dispatched on its +/// stored format tag (from [`KvIndex::interleaved_kquant_layer_data`]): +/// `Q4_K` and `Q6_K` (the default `down_proj`) each go through their amortised +/// matmul kernel ([`q4k_matmul_into`](crate::cpu::ops::q4_common::q4k_matmul_into) +/// / [`q6k_matmul_into`](crate::cpu::ops::q4_common::q6k_matmul_into)) — reads +/// the quantised bytes once, no f32 materialisation. Only a format without a +/// direct kernel falls back to dequantise + matmul. The FFN weights are ~4× the +/// attention weights, so avoiding their f32 staging is the bulk of the prefill +/// saving. +/// +/// The surrounding gating/activation/bias math mirrors [`dense_ffn_forward`]; +/// only the projection source differs. Callers confirm the interleaved FFN +/// bytes are present and `hidden` is a 256-multiple before constructing this. +pub struct Q4kMatmulFfn<'a> { + pub weights: &'a ModelWeights, + pub index: &'a dyn crate::KvIndex, +} + +impl Q4kMatmulFfn<'_> { + /// Bytes per 256-element super-block for a quant format. + #[inline] + fn block_bytes(fmt: &str) -> usize { + match fmt { + "Q4_K" => 144, + "Q6_K" => 210, + other => panic!("Q4kMatmulFfn: unsupported FFN quant format {other}"), + } + } + + /// gate/up projection: `x[seq, in_dim] -> [seq, out_rows]`, where `in_dim` + /// (= `hidden`) is a 256-multiple. Q4_K/Q6_K read the bytes once via the + /// amortised kernel; any other format dequantises then matmuls. + fn project( + bytes: &[u8], + fmt: &str, + x: &Array2, + out_rows: usize, + in_dim: usize, + seq: usize, + ) -> Array2 { + quant_proj(bytes, fmt, x, out_rows, in_dim, seq) + } + + /// down projection: `act[seq, intermediate] -> [seq, hidden]`. The stored + /// down weight pads `intermediate` up to a 256-multiple. Q4_K/Q6_K zero-pad + /// the activation to the stored width and use the amortised kernel; any + /// other format dequantises the padded weight and slices to `intermediate` + /// before the f32 matmul. + fn project_down( + bytes: &[u8], + fmt: &str, + act: &Array2, + hidden: usize, + intermediate: usize, + seq: usize, + ) -> Array2 { + let inter_padded = bytes.len() / hidden / Self::block_bytes(fmt) * 256; + let act_slice = act.as_slice().expect("contiguous activation"); + + if inter_padded == intermediate { + if let Some(r) = quant_matmul(bytes, fmt, act_slice, hidden, inter_padded, seq) { + return r; + } + let w = + crate::kquant_forward::dequant::dequantize_matrix(bytes, fmt, hidden, intermediate); + return dot_proj_gpu(act, &w, None); + } + + // `intermediate` isn't a 256-multiple: the stored weight is padded, so + // zero-pad the activation columns to match before projecting. + let mut padded = vec![0.0f32; seq * inter_padded]; + for s in 0..seq { + padded[s * inter_padded..s * inter_padded + intermediate] + .copy_from_slice(&act_slice[s * intermediate..(s + 1) * intermediate]); + } + if let Some(r) = quant_matmul(bytes, fmt, &padded, hidden, inter_padded, seq) { + return r; + } + let w_full = + crate::kquant_forward::dequant::dequantize_matrix(bytes, fmt, hidden, inter_padded); + let w_down = w_full.slice(ndarray::s![.., ..intermediate]).to_owned(); + dot_proj_gpu(act, &w_down, None) + } +} + +impl FfnBackend for Q4kMatmulFfn<'_> { + fn forward(&self, layer: usize, x: &Array2) -> Array2 { + self.forward_with_activation(layer, x).0 + } + + fn forward_with_activation(&self, layer: usize, x: &Array2) -> (Array2, Array2) { + let arch = &*self.weights.arch; + let seq = x.nrows(); + let hidden = x.ncols(); + let intermediate = self.index.num_features(layer); + let ffn = self + .index + .interleaved_kquant_layer_data(layer) + .expect("Q4kMatmulFfn requires interleaved FFN bytes for this layer"); + let (gate_bytes, gate_fmt) = ffn[0]; + let (up_bytes, up_fmt) = ffn[1]; + let (down_bytes, down_fmt) = ffn[2]; + + let activation = if arch.ffn_type() == larql_models::FfnType::Gated { + let gate = Self::project(gate_bytes, gate_fmt, x, intermediate, hidden, seq); + let up = Self::project(up_bytes, up_fmt, x, intermediate, hidden, seq); + match arch.activation() { + larql_models::Activation::GeluTanh => gelu_tanh_gate_up(&gate, &up), + _ => silu_gate_up(&gate, &up), + } + } else { + let mut projected = Self::project(up_bytes, up_fmt, x, intermediate, hidden, seq); + if let Some(bias) = arch + .ffn_up_bias_key(layer) + .and_then(|k| self.weights.vectors.get(&k)) + { + add_bias(&mut projected, bias); + } + match arch.activation() { + larql_models::Activation::GeluTanh | larql_models::Activation::Gelu => { + projected.mapv(gelu_tanh) + } + _ => projected.mapv(|v| v * sigmoid(v)), + } + }; + + let mut out = + Self::project_down(down_bytes, down_fmt, &activation, hidden, intermediate, seq); + if let Some(bias) = arch + .ffn_down_bias_key(layer) + .and_then(|k| self.weights.vectors.get(&k)) + { + add_bias(&mut out, bias); + } + + (out, activation) + } + + fn name(&self) -> &str { + "q4k-matmul" + } +} + /// Stub FFN that returns the input unchanged. Used by callers that must /// satisfy the `KvEngine::{prefill,decode_step}` signature but know the /// engine they're calling doesn't consult an FFN router (e.g. engines @@ -78,7 +305,7 @@ impl FfnBackend for NullFfn { /// Architecture-correct dense FFN — CPU BLAS path. pub fn dense_ffn_forward( - weights: &ModelWeights, + weights: WeightsView, layer: usize, x: &Array2, ) -> (Array2, Array2) { @@ -88,8 +315,12 @@ pub fn dense_ffn_forward( /// Architecture-correct dense FFN with optional backend dispatch. /// `backend = None` → plain ndarray BLAS (same as `dense_ffn_forward`). /// `backend = Some(be)` → gate/up/down matmuls through `be.matmul_transb`. +/// +/// Resolves FFN weights through [`WeightsView::tensor`] (engine scratch first, +/// then canonical) so the quant forward's dequantised FFN tensors are visible +/// without mutating `ModelWeights`. Dense callers pass a `dense()` view. pub fn dense_ffn_forward_backend( - weights: &ModelWeights, + weights: WeightsView, layer: usize, x: &Array2, backend: Option<&dyn ComputeBackend>, @@ -100,18 +331,15 @@ pub fn dense_ffn_forward_backend( (or re-extract without `--compact` if you need dense matmul)."; let w_up = weights - .tensors - .get(&arch.ffn_up_key(layer)) + .tensor(&arch.ffn_up_key(layer)) .unwrap_or_else(|| panic!("{compact_hint} (key: {})", arch.ffn_up_key(layer))); let w_down = weights - .tensors - .get(&arch.ffn_down_key(layer)) + .tensor(&arch.ffn_down_key(layer)) .unwrap_or_else(|| panic!("{compact_hint} (key: {})", arch.ffn_down_key(layer))); let activation = if arch.ffn_type() == larql_models::FfnType::Gated { let w_gate = weights - .tensors - .get(&arch.ffn_gate_key(layer)) + .tensor(&arch.ffn_gate_key(layer)) .unwrap_or_else(|| panic!("{compact_hint} (key: {})", arch.ffn_gate_key(layer))); let gate = dot_proj_gpu(x, w_gate, backend); let up = dot_proj_gpu(x, w_up, backend); @@ -166,7 +394,7 @@ mod tests { fn dense_ffn_forward_shape() { let weights = make_test_weights(); let input = x(3, weights.hidden_size); - let (out, act) = dense_ffn_forward(&weights, 0, &input); + let (out, act) = dense_ffn_forward(WeightsView::dense(&weights), 0, &input); assert_eq!(out.shape(), &[3, weights.hidden_size]); assert_eq!(act.shape(), &[3, weights.intermediate_size]); } @@ -175,7 +403,7 @@ mod tests { fn dense_ffn_forward_output_finite() { let weights = make_test_weights(); let input = x(2, weights.hidden_size); - let (out, act) = dense_ffn_forward(&weights, 0, &input); + let (out, act) = dense_ffn_forward(WeightsView::dense(&weights), 0, &input); assert!( out.iter().all(|v| v.is_finite()), "FFN output has non-finite values" @@ -191,8 +419,8 @@ mod tests { // backend=None should produce the same result as dense_ffn_forward let weights = make_test_weights(); let input = x(2, weights.hidden_size); - let (out1, act1) = dense_ffn_forward(&weights, 0, &input); - let (out2, act2) = dense_ffn_forward_backend(&weights, 0, &input, None); + let (out1, act1) = dense_ffn_forward(WeightsView::dense(&weights), 0, &input); + let (out2, act2) = dense_ffn_forward_backend(WeightsView::dense(&weights), 0, &input, None); assert_eq!( out1, out2, "output should match between dense_ffn_forward and backend(None)" @@ -205,7 +433,7 @@ mod tests { let weights = make_test_weights(); let input = x(1, weights.hidden_size); for layer in 0..weights.num_layers { - let (out, _) = dense_ffn_forward(&weights, layer, &input); + let (out, _) = dense_ffn_forward(WeightsView::dense(&weights), layer, &input); assert_eq!( out.shape(), &[1, weights.hidden_size], @@ -291,7 +519,7 @@ mod tests { // Edge case: one row at the smallest meaningful seq_len. let weights = make_test_weights(); let input = x(1, weights.hidden_size); - let (out, act) = dense_ffn_forward(&weights, 0, &input); + let (out, act) = dense_ffn_forward(WeightsView::dense(&weights), 0, &input); assert_eq!(out.shape(), &[1, weights.hidden_size]); assert_eq!(act.shape(), &[1, weights.intermediate_size]); } @@ -303,7 +531,7 @@ mod tests { // change to the gated path. let weights = make_test_weights(); let input = Array2::::zeros((2, weights.hidden_size)); - let (out, act) = dense_ffn_forward(&weights, 0, &input); + let (out, act) = dense_ffn_forward(WeightsView::dense(&weights), 0, &input); assert!(out.iter().all(|v| v.is_finite())); assert!(act.iter().all(|v| v.is_finite())); } @@ -315,9 +543,14 @@ mod tests { // output (within float noise). let weights = make_test_weights(); let input = x(2, weights.hidden_size); - let (out_none, act_none) = dense_ffn_forward_backend(&weights, 0, &input, None); - let (out_some, act_some) = - dense_ffn_forward_backend(&weights, 0, &input, Some(&crate::CpuBackend)); + let (out_none, act_none) = + dense_ffn_forward_backend(WeightsView::dense(&weights), 0, &input, None); + let (out_some, act_some) = dense_ffn_forward_backend( + WeightsView::dense(&weights), + 0, + &input, + Some(&crate::CpuBackend), + ); for (a, b) in out_none.iter().zip(out_some.iter()) { assert!((a - b).abs() < 1e-4, "out diverged: {a} vs {b}"); } @@ -334,7 +567,7 @@ mod tests { // the `else` branch (no gate matrix; just up + activation + down). let weights = larql_models::test_fixtures::make_starcoder2_test_weights(); let input = x(2, weights.hidden_size); - let (out, act) = dense_ffn_forward(&weights, 0, &input); + let (out, act) = dense_ffn_forward(WeightsView::dense(&weights), 0, &input); assert_eq!(out.shape(), &[2, weights.hidden_size]); assert!(out.iter().all(|v| v.is_finite())); // Non-gated activation has shape (seq, intermediate). @@ -348,7 +581,7 @@ mod tests { // bias)` calls fire. let weights = larql_models::test_fixtures::make_starcoder2_test_weights(); let input = x(1, weights.hidden_size); - let (out, _) = dense_ffn_forward(&weights, 0, &input); + let (out, _) = dense_ffn_forward(WeightsView::dense(&weights), 0, &input); assert!(out.iter().all(|v| v.is_finite())); } @@ -360,8 +593,109 @@ mod tests { // `gelu_tanh_gate_up` branch instead of the default silu. let weights = larql_models::test_fixtures::make_gemma3_test_weights(); let input = x(2, weights.hidden_size); - let (out, _) = dense_ffn_forward(&weights, 0, &input); + let (out, _) = dense_ffn_forward(WeightsView::dense(&weights), 0, &input); assert_eq!(out.shape(), &[2, weights.hidden_size]); assert!(out.iter().all(|v| v.is_finite())); } + + // ── q4k-direct FFN parity ────────────────────────────────────────── + + /// `Q4kMatmulFfn` must match dequantising the SAME Q4_K bytes and + /// running the dense FFN — both decode identical weights, so they agree + /// within fp summation noise. This is the prefill correctness contract: + /// swapping the FFN to q4k-direct must not change the output. + #[test] + fn q4k_matmul_ffn_matches_dequant_dense() { + use super::Q4kMatmulFfn; + use crate::test_fixtures::make_q4k_fixture_index; + use larql_models::test_fixtures::make_test_q4k_weights; + + let weights = make_test_q4k_weights(); + let index = make_q4k_fixture_index(&weights); + let input = x(3, weights.hidden_size); + + // Reference: dequant the layer's Q4_K FFN bytes into scratch, then + // run the dense FFN against those f32 tensors (the current path). + let mut scratch = larql_models::DequantScratch::new(); + crate::kquant_forward::insert_q4k_layer_tensors(&mut scratch, &weights, &index, 0) + .expect("dequant layer 0"); + let (ref_out, ref_act) = + dense_ffn_forward(WeightsView::with_scratch(&weights, &scratch), 0, &input); + + // q4k-direct: same bytes, no dequant. + let ffn = Q4kMatmulFfn { + weights: &weights, + index: &index, + }; + let (got_out, got_act) = ffn.forward_with_activation(0, &input); + + let max_out: f32 = ref_out + .iter() + .zip(&got_out) + .map(|(a, b)| (a - b).abs()) + .fold(0.0, f32::max); + let max_act: f32 = ref_act + .iter() + .zip(&got_act) + .map(|(a, b)| (a - b).abs()) + .fold(0.0, f32::max); + assert!( + max_out < 5e-3, + "q4k FFN output diverged: max_diff={max_out}" + ); + assert!( + max_act < 5e-3, + "q4k FFN activation diverged: max_diff={max_act}" + ); + assert_eq!(got_out.shape(), &[3, weights.hidden_size]); + } + + /// Regression for the Q6_K `down_proj` mis-decode: the *default* q4k vindex + /// stores FFN down at Q6_K (210 B/block), not Q4_K (144 B/block). The down + /// projection must dispatch on the format tag — decoding Q6_K bytes through + /// `q4k_matmul` produces garbage and a wrong derived column count. Compare + /// the Q6_K `project_down` against dequantise+matmul on the same bytes. + #[test] + fn project_down_dispatches_on_q6k_format() { + use super::Q4kMatmulFfn; + use crate::cpu::ops::q4_common::quantize_q6_k; + + let hidden = 256; + let intermediate = 512; // 2 super-blocks; 256-multiple (matches real gemma) + let seq = 3; + + let down: Vec = (0..hidden * intermediate) + .map(|i| (i as f32 * 0.0007).sin() * 0.6) + .collect(); + let down_q6 = quantize_q6_k(&down); + // Bounded activation: the amortised kernel sums per-super-block while the + // BLAS reference sums the whole row, so an unbounded/growing activation + // amplifies fp-reorder noise via cancellation. This is a correctness + // check (same decoded weights), so keep the inputs well-conditioned. + let act = Array2::from_shape_fn((seq, intermediate), |(s, j)| { + (((s * 13 + j) % 97) as f32 / 97.0 - 0.5) * 0.8 + }); + + // Reference: dequantise the Q6_K weight to f32, then act @ down.T. + let w = crate::kquant_forward::dequant::dequantize_matrix( + &down_q6, + "Q6_K", + hidden, + intermediate, + ); + let reference = crate::dot_proj_gpu(&act, &w, None); + + let got = Q4kMatmulFfn::project_down(&down_q6, "Q6_K", &act, hidden, intermediate, seq); + assert_eq!(got.shape(), &[seq, hidden]); + + let max: f32 = reference + .iter() + .zip(&got) + .map(|(a, b)| (a - b).abs()) + .fold(0.0, f32::max); + assert!( + max < 5e-3, + "Q6_K project_down diverged from dequant+matmul: {max}" + ); + } } diff --git a/crates/larql-compute/src/forward/layer.rs b/crates/larql-compute/src/forward/layer.rs index 91c3fb0c7..c38c7c278 100644 --- a/crates/larql-compute/src/forward/layer.rs +++ b/crates/larql-compute/src/forward/layer.rs @@ -9,12 +9,12 @@ use super::ple::apply_per_layer_embedding; use crate::attention::{AttentionWeights, SharedKV}; use crate::ffn::FfnBackend; use crate::residual::rms_norm_for_arch; -use larql_models::ModelWeights; +use larql_models::{ModelWeights, WeightsView}; use ndarray::Array2; /// Public wrapper for run_attention — used by diagnostic/capture tooling. pub fn run_attention_public( - weights: &ModelWeights, + weights: WeightsView, h: &Array2, layer: usize, ) -> Option> { @@ -22,14 +22,14 @@ pub fn run_attention_public( } /// Run attention for a single layer. Returns the post-attention residual. -pub fn run_attention(weights: &ModelWeights, h: &Array2, layer: usize) -> Option> { +pub fn run_attention(weights: WeightsView, h: &Array2, layer: usize) -> Option> { let (h_post_attn, _) = run_attention_inner(weights, h, layer, false, None)?; Some(h_post_attn) } /// Run attention with optional per-head weight capture and shared K/V. pub fn run_attention_inner( - weights: &ModelWeights, + weights: WeightsView, h: &Array2, layer: usize, capture_attention: bool, @@ -48,7 +48,7 @@ pub fn run_attention_inner( /// Run attention returning post-processed K/V for caching (KV sharing source layers). pub fn run_attention_with_kv_cache( - weights: &ModelWeights, + weights: WeightsView, h: &Array2, layer: usize, ) -> Option<(Array2, SharedKV)> { @@ -147,7 +147,7 @@ pub fn apply_layer_scalar(weights: &ModelWeights, h: &mut Array2, layer: us /// sequence as `predict_with_temperature` without duplicating logic. #[allow(clippy::type_complexity)] pub fn run_layer_with_ffn( - weights: &ModelWeights, + weights: WeightsView, h: &Array2, layer: usize, ffn: &dyn FfnBackend, @@ -175,9 +175,9 @@ pub fn run_layer_with_ffn( let path = crate::forward::dump_config::cpu_layer_h_post_attn_path(dir, layer); let _ = std::fs::write(&path, &bytes); } - let (h_post_ffn, activation) = run_ffn(weights, &h_post_attn, layer, ffn, capture_activation); - let mut h_out = apply_per_layer_embedding(weights, &h_post_ffn, layer, ple_input); - apply_layer_scalar(weights, &mut h_out, layer); + let (h_post_ffn, activation) = run_ffn(&weights, &h_post_attn, layer, ffn, capture_activation); + let mut h_out = apply_per_layer_embedding(&weights, &h_post_ffn, layer, ple_input); + apply_layer_scalar(&weights, &mut h_out, layer); Some((h_out, activation, kv_out)) } @@ -188,7 +188,7 @@ pub fn run_layer_with_ffn( #[allow(clippy::too_many_arguments)] #[allow(clippy::type_complexity)] pub fn run_layer_with_capture( - weights: &ModelWeights, + weights: WeightsView, h: &Array2, layer: usize, ffn: &dyn FfnBackend, @@ -224,7 +224,7 @@ pub fn run_layer_with_capture( #[allow(clippy::too_many_arguments)] #[allow(clippy::type_complexity)] pub fn run_layer_with_capture_hooked( - weights: &ModelWeights, + weights: WeightsView, h: &Array2, layer: usize, ffn: &dyn FfnBackend, @@ -261,13 +261,13 @@ pub fn run_layer_with_capture_hooked( } hook.on_post_attention(layer, &mut h_post_attn); - let (h_post_ffn, activation) = run_ffn(weights, &h_post_attn, layer, ffn, capture_activation); + let (h_post_ffn, activation) = run_ffn(&weights, &h_post_attn, layer, ffn, capture_activation); if let Some(ref act) = activation { hook.on_ffn_activation(layer, act); } - let mut h_out = apply_per_layer_embedding(weights, &h_post_ffn, layer, ple_input); - apply_layer_scalar(weights, &mut h_out, layer); + let mut h_out = apply_per_layer_embedding(&weights, &h_post_ffn, layer, ple_input); + apply_layer_scalar(&weights, &mut h_out, layer); hook.on_post_layer(layer, &mut h_out); Some((h_out, activation, attn_weights, kv_out)) @@ -315,7 +315,8 @@ mod tests { let weights = make_test_weights(); let ffn = StubFfn { weights: &weights }; let h = Array2::from_elem((2, weights.hidden_size), 0.5f32); - let result = run_layer_with_ffn(&weights, &h, 0, &ffn, false, None, None); + let result = + run_layer_with_ffn(WeightsView::dense(&weights), &h, 0, &ffn, false, None, None); assert!(result.is_some(), "run_layer_with_ffn returned None"); let (h_out, _act, _kv) = result.unwrap(); assert_eq!(h_out.shape(), h.shape()); @@ -326,7 +327,7 @@ mod tests { fn run_attention_public_returns_post_attention_residual() { let weights = make_test_weights(); let h = Array2::from_elem((2, weights.hidden_size), 0.5f32); - let h_post = run_attention_public(&weights, &h, 0) + let h_post = run_attention_public(WeightsView::dense(&weights), &h, 0) .expect("attention should return post-residual on standard weights"); assert_eq!(h_post.shape(), h.shape()); assert!(h_post.iter().all(|v| v.is_finite())); @@ -336,7 +337,8 @@ mod tests { fn run_attention_returns_same_shape_as_input() { let weights = make_test_weights(); let h = Array2::from_elem((3, weights.hidden_size), 0.1f32); - let h_post = run_attention(&weights, &h, 0).expect("attention should succeed"); + let h_post = + run_attention(WeightsView::dense(&weights), &h, 0).expect("attention should succeed"); assert_eq!(h_post.shape(), h.shape()); } @@ -344,9 +346,14 @@ mod tests { fn run_attention_inner_with_capture_returns_attention_weights() { let weights = make_test_weights(); let h = Array2::from_elem((2, weights.hidden_size), 0.1f32); - let (h_post, attn_w) = - run_attention_inner(&weights, &h, 0, /*capture_attention=*/ true, None) - .expect("inner attention"); + let (h_post, attn_w) = run_attention_inner( + WeightsView::dense(&weights), + &h, + 0, + /*capture_attention=*/ true, + None, + ) + .expect("inner attention"); assert_eq!(h_post.shape(), h.shape()); assert!( attn_w.is_some(), @@ -358,9 +365,14 @@ mod tests { fn run_attention_inner_without_capture_drops_attention_weights() { let weights = make_test_weights(); let h = Array2::from_elem((2, weights.hidden_size), 0.1f32); - let (h_post, attn_w) = - run_attention_inner(&weights, &h, 0, /*capture_attention=*/ false, None) - .expect("inner attention"); + let (h_post, attn_w) = run_attention_inner( + WeightsView::dense(&weights), + &h, + 0, + /*capture_attention=*/ false, + None, + ) + .expect("inner attention"); assert_eq!(h_post.shape(), h.shape()); assert!( attn_w.is_none(), @@ -372,8 +384,8 @@ mod tests { fn run_attention_with_kv_cache_returns_kv_pair() { let weights = make_test_weights(); let h = Array2::from_elem((3, weights.hidden_size), 0.1f32); - let (h_post, (k, v)) = - run_attention_with_kv_cache(&weights, &h, 0).expect("kv cache attention"); + let (h_post, (k, v)) = run_attention_with_kv_cache(WeightsView::dense(&weights), &h, 0) + .expect("kv cache attention"); assert_eq!(h_post.shape(), h.shape()); // K and V have the same shape (seq_len × kv_dim). assert_eq!(k.shape(), v.shape()); @@ -388,7 +400,7 @@ mod tests { let h = Array2::from_elem((2, weights.hidden_size), 0.1f32); let mut record = RecordHook::for_layers([0]); let result = run_layer_with_capture_hooked( - &weights, + WeightsView::dense(&weights), &h, /*layer=*/ 0, &ffn, @@ -423,10 +435,11 @@ mod tests { let ffn = StubFfn { weights: &weights }; let h = Array2::from_elem((2, weights.hidden_size), 0.1f32); // Run once to build a SharedKV. - let (_, fresh_kv) = run_attention_with_kv_cache(&weights, &h, 0).unwrap(); + let (_, fresh_kv) = + run_attention_with_kv_cache(WeightsView::dense(&weights), &h, 0).unwrap(); let mut hook = crate::forward::NoopHook; let result = run_layer_with_capture_hooked( - &weights, + WeightsView::dense(&weights), &h, 0, &ffn, @@ -448,7 +461,16 @@ mod tests { let weights = make_test_weights(); let ffn = StubFfn { weights: &weights }; let h = Array2::from_elem((2, weights.hidden_size), 0.1f32); - let result = run_layer_with_capture(&weights, &h, 0, &ffn, true, true, None, None); + let result = run_layer_with_capture( + WeightsView::dense(&weights), + &h, + 0, + &ffn, + true, + true, + None, + None, + ); let (h_out, act, attn_w, kv_out) = result.expect("non-hooked capture wrapper"); assert_eq!(h_out.shape(), h.shape()); assert!(act.is_some()); @@ -464,7 +486,9 @@ mod tests { let weights = make_test_weights(); let ffn = StubFfn { weights: &weights }; let h = Array2::from_elem((3, weights.hidden_size), 1.0f32); - let (h_out, _, _) = run_layer_with_ffn(&weights, &h, 0, &ffn, false, None, None).unwrap(); + let (h_out, _, _) = + run_layer_with_ffn(WeightsView::dense(&weights), &h, 0, &ffn, false, None, None) + .unwrap(); let differ = h .iter() .zip(h_out.iter()) diff --git a/crates/larql-compute/src/forward/predict/raw.rs b/crates/larql-compute/src/forward/predict/raw.rs index d1640d265..a4bcdb1ec 100644 --- a/crates/larql-compute/src/forward/predict/raw.rs +++ b/crates/larql-compute/src/forward/predict/raw.rs @@ -7,7 +7,6 @@ use super::super::layer::run_layer_with_ffn; use super::super::ple::precompute_per_layer_inputs; use super::super::{apply_norm, dot_proj}; use crate::attention::SharedKV; -use crate::ffn::WeightFfn; use larql_models::ModelWeights; use ndarray::Array2; @@ -63,7 +62,7 @@ pub fn hidden_to_raw_logits(weights: &ModelWeights, h_single: &Array2) -> V /// (since `run_layer_with_ffn` already collapses the block's output + /// skip, perturbing the post-block `h[-1]` is algebraically the same). pub fn forward_raw_logits( - weights: &ModelWeights, + weights: larql_models::WeightsView, token_ids: &[u32], perturb: Option<(usize, ndarray::ArrayView1)>, ) -> RawForward { @@ -89,7 +88,7 @@ pub fn forward_raw_logits( /// `initial_residual`, when `Some`, must be a slice of exactly /// `weights.hidden_size` floats. `token_ids` may not be empty. pub fn forward_raw_logits_with_prefix( - weights: &ModelWeights, + weights: larql_models::WeightsView, token_ids: &[u32], initial_residual: Option<&[f32]>, perturb: Option<(usize, ndarray::ArrayView1)>, @@ -114,7 +113,7 @@ pub fn forward_raw_logits_with_prefix( /// Layout: `h[0] = boundary`, `h[1..]` = query embeddings. /// The perturbation is applied at `target_layer` to the last row. pub fn forward_from_layer( - weights: &ModelWeights, + weights: larql_models::WeightsView, token_ids: &[u32], boundary_residual: &[f32], from_layer: usize, @@ -140,7 +139,7 @@ pub fn forward_from_layer( /// Layout when `prefix` is `None`: rows 0..q = query embeddings, /// total_len = q. fn forward_layer_range( - weights: &ModelWeights, + weights: larql_models::WeightsView, token_ids: &[u32], prefix: Option<&[f32]>, layer_range: Range, @@ -149,7 +148,7 @@ fn forward_layer_range( let hidden = weights.hidden_size; let q_len = token_ids.len(); - let q_embed = embed_tokens_pub(weights, token_ids); + let q_embed = embed_tokens_pub(&weights, token_ids); let (mut h, total_len) = if let Some(prefix) = prefix { assert_eq!( prefix.len(), @@ -184,8 +183,8 @@ fn forward_layer_range( } else { token_ids.to_vec() }; - let ple_inputs = precompute_per_layer_inputs(weights, &h, &ple_token_ids); - let ffn = WeightFfn { weights }; + let ple_inputs = precompute_per_layer_inputs(&weights, &h, &ple_token_ids); + let ffn = crate::ffn::ViewFfn { view: weights }; let mut kv_cache: std::collections::HashMap = std::collections::HashMap::new(); @@ -227,11 +226,11 @@ fn forward_layer_range( let h_pre_norm = h.clone(); let norm_offset = weights.arch.norm_weight_offset(); - let h_final = apply_norm(weights, &h, weights.arch.final_norm_key(), norm_offset); + let h_final = apply_norm(&weights, &h, weights.arch.final_norm_key(), norm_offset); let last_2d = h_final.slice(ndarray::s![total_len - 1..total_len, ..]); let logits_raw = dot_proj(&last_2d, &weights.lm_head); - let logits = apply_logits_transform(weights, logits_raw.row(0).as_slice().unwrap()); + let logits = apply_logits_transform(&weights, logits_raw.row(0).as_slice().unwrap()); RawForward { h_pre_norm, @@ -250,7 +249,11 @@ mod forward_from_layer_tests { #[test] fn forward_raw_logits_returns_vocab_logits() { let weights = make_test_weights(); - let raw = forward_raw_logits(&weights, &[0u32, 1, 2], None); + let raw = forward_raw_logits( + larql_models::WeightsView::dense(&weights), + &[0u32, 1, 2], + None, + ); assert_eq!( raw.logits.len(), weights.vocab_size, @@ -266,7 +269,7 @@ mod forward_from_layer_tests { #[test] fn forward_raw_logits_single_token() { let weights = make_test_weights(); - let raw = forward_raw_logits(&weights, &[5u32], None); + let raw = forward_raw_logits(larql_models::WeightsView::dense(&weights), &[5u32], None); assert_eq!(raw.logits.len(), weights.vocab_size); assert!( raw.logits.iter().all(|v| v.is_finite()), @@ -284,7 +287,13 @@ mod forward_from_layer_tests { let token_ids = &[1u32, 2]; let boundary = vec![0.0f32; weights.hidden_size]; - let from_layer = forward_from_layer(&weights, token_ids, &boundary, 0, None); + let from_layer = forward_from_layer( + larql_models::WeightsView::dense(&weights), + token_ids, + &boundary, + 0, + None, + ); // from_layer=0 with zero boundary: should have (1 boundary + 2 query) positions assert_eq!(from_layer.h_pre_norm.shape(), &[3, weights.hidden_size]); assert_eq!(from_layer.logits.len(), weights.vocab_size); @@ -299,8 +308,20 @@ mod forward_from_layer_tests { let token_ids = &[3u32]; let boundary = vec![0.1f32; weights.hidden_size]; - let from_0 = forward_from_layer(&weights, token_ids, &boundary, 0, None); - let from_1 = forward_from_layer(&weights, token_ids, &boundary, 1, None); + let from_0 = forward_from_layer( + larql_models::WeightsView::dense(&weights), + token_ids, + &boundary, + 0, + None, + ); + let from_1 = forward_from_layer( + larql_models::WeightsView::dense(&weights), + token_ids, + &boundary, + 1, + None, + ); // Outputs should differ (layer 0's transform changes the residual) let differ = from_0 @@ -319,7 +340,7 @@ mod forward_from_layer_tests { let weights = make_test_weights(); // 3 query tokens, from_layer=1: h has 4 rows (1 boundary + 3 query) let raw = forward_from_layer( - &weights, + larql_models::WeightsView::dense(&weights), &[0u32, 1, 2], &vec![0.0; weights.hidden_size], 1, @@ -333,7 +354,12 @@ mod forward_from_layer_tests { fn forward_raw_logits_with_prefix_shape() { let weights = make_test_weights(); let prefix = vec![0.5f32; weights.hidden_size]; - let raw = forward_raw_logits_with_prefix(&weights, &[0u32, 1], Some(&prefix), None); + let raw = forward_raw_logits_with_prefix( + larql_models::WeightsView::dense(&weights), + &[0u32, 1], + Some(&prefix), + None, + ); // prefix + 2 tokens = 3 positions assert_eq!(raw.h_pre_norm.shape(), &[3, weights.hidden_size]); assert_eq!(raw.logits.len(), weights.vocab_size); diff --git a/crates/larql-compute/src/kquant_forward/cached.rs b/crates/larql-compute/src/kquant_forward/cached.rs index 35b56bc13..c600862c7 100644 --- a/crates/larql-compute/src/kquant_forward/cached.rs +++ b/crates/larql-compute/src/kquant_forward/cached.rs @@ -7,9 +7,12 @@ //! per-step decode (single-row attention against the cache + 1-row //! FFN). Speedup scales linearly with decode length. //! -//! Per-step Q4_K → f32 dequant via `insert_q4k_layer_tensors` is -//! still paid for now; eliminating it is a follow-up (route Q/K/V/O -//! and gate/up/down through `backend.q4k_matvec` directly). +//! Prefill projects Q/K/V/O and gate/up/down straight from the vindex's +//! Q4_K/Q6_K bytes via the amortised q4k/q6k matmul — no per-layer f32 +//! dequant — when the projection dims are 256-aligned (see +//! `predict_kquant_prefill_with_state`'s `use_q4k_attn` / `use_q4k_ffn`). +//! Per-step decode still dequantises via `insert_q4k_layer_tensors` +//! (a smaller follow-up). //! //! Scope: dense architectures only. Hybrid-MoE (Gemma 4 26B A4B) //! and cross-layer KV sharing (Gemma 4 E2B) fall back to the slow @@ -35,7 +38,6 @@ use crate::attention::{ rope::apply_rope_partial_at_full, run_attention_with_kv_backend, }; -use crate::ffn::WeightFfn; use crate::forward::embed_tokens_pub; use crate::forward::layer::apply_layer_scalar; use crate::forward::ple::{apply_per_layer_embedding, precompute_per_layer_inputs}; @@ -43,7 +45,7 @@ use crate::forward::run_ffn; use crate::forward::{add_bias, apply_norm}; use crate::residual::{rms_norm_heads, rms_norm_heads_no_weight}; -use super::tensors::{insert_q4k_layer_tensors, remove_layer_tensors}; +use super::tensors::{insert_q4k_attn_tensors, insert_q4k_layer_tensors, remove_layer_tensors}; /// Per-layer K/V captured during prefill. One entry per layer; matches /// the [`crate::attention::decode::KvCache`] convention so future work @@ -86,7 +88,7 @@ pub fn supports_cached_decode(weights: &ModelWeights) -> bool { /// Returns the `[seq_len, hidden]` hidden state and the populated /// cache. Caller takes the last row for lm_head. pub fn predict_kquant_prefill( - weights: &mut ModelWeights, + weights: &ModelWeights, token_ids: &[u32], index: &dyn crate::KvIndex, ) -> (Array2, CpuKvCache, CachedTimings) { @@ -102,7 +104,7 @@ pub fn predict_kquant_prefill( /// 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( - weights: &mut ModelWeights, + weights: &ModelWeights, token_ids: &[u32], index: &dyn crate::KvIndex, mut state: Option<&mut crate::PerLayerDecodeState>, @@ -110,14 +112,38 @@ pub fn predict_kquant_prefill_with_state( let num_layers = weights.num_layers; let mut cache: CpuKvCache = vec![None; num_layers]; let mut timings = CachedTimings::default(); + // Forward-local dequant scratch — per-forward derived state; `weights` + // stays immutable. Readers resolve via with_scratch (scratch ∪ canonical). + let mut scratch = larql_models::DequantScratch::new(); let mut h = embed_tokens_pub(weights, token_ids); let ple_inputs = precompute_per_layer_inputs(weights, &h, token_ids); for layer in 0..num_layers { + // q4k-direct prefill: project straight from the vindex's Q4_K/Q6_K bytes + // (no per-layer f32 dequant). The FFN gate/up/down (`use_q4k_ffn`) and + // attention Q/K/V/O (`use_q4k_attn`) are gated independently — each needs + // its interleaved bytes present and the relevant dims 256-aligned (the + // matmul contraction is `hidden` for Q/K/V/gate/up and `q_dim` for O). + // The FFN weights are ~4× the attention weights, so the FFN path is the + // bulk of the saving; the attn path closes the rest. + const BLK: usize = larql_models::quant::ggml::K_QUANT_BLOCK_ELEMS; + let q_dim = + weights.arch.num_q_heads_for_layer(layer) * weights.arch.head_dim_for_layer(layer); + let use_q4k_ffn = weights.hidden_size.is_multiple_of(BLK) + && index.interleaved_kquant_layer_data(layer).is_some(); + let use_q4k_attn = weights.hidden_size.is_multiple_of(BLK) + && q_dim.is_multiple_of(BLK) + && index.attn_kquant_layer_data(layer).is_some(); + let t0 = std::time::Instant::now(); - let inserted = - insert_q4k_layer_tensors(weights, index, layer).unwrap_or_else(|err| panic!("{err}")); + // Dequant only what the q4k-direct paths won't read straight from bytes. + let inserted = match (use_q4k_attn, use_q4k_ffn) { + (true, true) => Ok(Vec::new()), + (false, true) => insert_q4k_attn_tensors(&mut scratch, weights, index, layer), + _ => insert_q4k_layer_tensors(&mut scratch, weights, index, layer), + } + .unwrap_or_else(|err| panic!("{err}")); timings.dequant_ms += t0.elapsed().as_secs_f64() * 1000.0; // Snapshot pre-attention residual for this layer if engine wants it. @@ -126,17 +152,23 @@ pub fn predict_kquant_prefill_with_state( .push(crate::state_handle::CpuStateHandle::boxed(h.clone())); } - // Attention with K/V capture. Backend stays None — we want the - // CPU BLAS path for the dequantised f32 tensors that - // `insert_q4k_layer_tensors` just placed in `weights.tensors`. - let (h_post_attn, k_rope, v_final) = - match run_attention_with_kv_backend(weights, &h, layer, None) { - Some(t) => t, - None => { - remove_layer_tensors(weights, inserted); - return (h, cache, timings); - } - }; + // Attention with K/V capture. When `use_q4k_attn`, Q/K/V/O project + // straight from the Q4_K/Q6_K bytes (empty scratch — norms still come + // from canonical weights); else the dequantised f32 view path. + let attn_index = if use_q4k_attn { Some(index) } else { None }; + let (h_post_attn, k_rope, v_final) = match run_attention_with_kv_backend( + larql_models::WeightsView::with_scratch(weights, &scratch), + &h, + layer, + None, + attn_index, + ) { + Some(t) => t, + None => { + remove_layer_tensors(&mut scratch, inserted); + return (h, cache, timings); + } + }; if let Some(s) = state.as_deref_mut() { // Prefill K/V for THIS layer = full seq_len × kv_dim. @@ -146,13 +178,20 @@ pub fn predict_kquant_prefill_with_state( .push(crate::state_handle::CpuStateHandle::boxed(v_final.clone())); } - let ffn = WeightFfn { weights }; - let (h_post_ffn, _) = run_ffn(weights, &h_post_attn, layer, &ffn, false); + let h_post_ffn = if use_q4k_ffn { + let ffn = crate::ffn::Q4kMatmulFfn { weights, index }; + run_ffn(weights, &h_post_attn, layer, &ffn, false).0 + } else { + let ffn = crate::ffn::ViewFfn { + view: larql_models::WeightsView::with_scratch(weights, &scratch), + }; + run_ffn(weights, &h_post_attn, layer, &ffn, false).0 + }; 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); - remove_layer_tensors(weights, inserted); + remove_layer_tensors(&mut scratch, inserted); cache[layer] = Some((k_rope, v_final)); h = h_out; @@ -169,7 +208,7 @@ pub fn predict_kquant_prefill_with_state( /// `prompt_len + steps_already_decoded`. The caller maintains this /// counter (typical: `prompt_len + step_index` starting at 0). pub fn predict_kquant_decode_step( - weights: &mut ModelWeights, + weights: &ModelWeights, token_id: u32, index: &dyn crate::KvIndex, cache: &mut CpuKvCache, @@ -180,6 +219,7 @@ pub fn predict_kquant_decode_step( return None; } let mut timings = CachedTimings::default(); + let mut scratch = larql_models::DequantScratch::new(); // 1-row embed + 1-row PLE for the new token. let mut h = embed_tokens_pub(weights, &[token_id]); @@ -187,13 +227,13 @@ pub fn predict_kquant_decode_step( for layer in 0..num_layers { let t0 = std::time::Instant::now(); - let inserted = - insert_q4k_layer_tensors(weights, index, layer).unwrap_or_else(|err| panic!("{err}")); + let inserted = insert_q4k_layer_tensors(&mut scratch, weights, index, layer) + .unwrap_or_else(|err| panic!("{err}")); timings.dequant_ms += t0.elapsed().as_secs_f64() * 1000.0; let kv_entry = cache[layer].as_ref(); let (h_post_attn, new_kv) = match run_attention_block_decode_step_backend( - weights, + larql_models::WeightsView::with_scratch(weights, &scratch), &h, layer, kv_entry, @@ -202,19 +242,21 @@ pub fn predict_kquant_decode_step( ) { Some(t) => t, None => { - remove_layer_tensors(weights, inserted); + remove_layer_tensors(&mut scratch, inserted); return None; } }; cache[layer] = Some(new_kv); - let ffn = WeightFfn { weights }; + let ffn = crate::ffn::ViewFfn { + view: larql_models::WeightsView::with_scratch(weights, &scratch), + }; 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); - remove_layer_tensors(weights, inserted); + remove_layer_tensors(&mut scratch, inserted); h = h_out; } @@ -919,7 +961,7 @@ fn run_ffn_decode_step_q4k_direct( /// if any layer has a format the direct-matvec path doesn't handle /// (caller falls back to [`predict_kquant_decode_step`]). pub fn predict_kquant_decode_step_direct( - weights: &mut ModelWeights, + weights: &ModelWeights, token_id: u32, index: &dyn crate::KvIndex, backend: &dyn ComputeBackend, @@ -946,7 +988,7 @@ pub fn predict_kquant_decode_step_direct( /// coarse_decode_step_with_state`. When `state` is `None` this is /// bit-identical to [`predict_kquant_decode_step_direct`]. pub fn predict_kquant_decode_step_direct_with_state( - weights: &mut ModelWeights, + weights: &ModelWeights, token_id: u32, index: &dyn crate::KvIndex, backend: &dyn ComputeBackend, diff --git a/crates/larql-compute/src/kquant_forward/dequant.rs b/crates/larql-compute/src/kquant_forward/dequant.rs index e5184f1a2..ff856be92 100644 --- a/crates/larql-compute/src/kquant_forward/dequant.rs +++ b/crates/larql-compute/src/kquant_forward/dequant.rs @@ -15,7 +15,7 @@ use ndarray::Array2; /// the k-quant super-block size. Unknown formats panic; callers have /// already dispatched on format via the `attn_kquant_layer_data` / /// `interleaved_kquant_layer_data` tag. -pub(super) fn dequantize_matrix( +pub(crate) fn dequantize_matrix( bytes: &[u8], format: &str, rows: usize, diff --git a/crates/larql-compute/src/kquant_forward/hooks.rs b/crates/larql-compute/src/kquant_forward/hooks.rs index a122c8bdc..b593b65c5 100644 --- a/crates/larql-compute/src/kquant_forward/hooks.rs +++ b/crates/larql-compute/src/kquant_forward/hooks.rs @@ -18,7 +18,7 @@ use super::tensors::{insert_q4k_layer_tensors, remove_layer_tensors}; /// behavior of `predict_kquant_hidden` while exposing pre-layer, post-attention, /// optional attention-weight/FFN-activation, and post-layer hook points. pub fn predict_kquant_hidden_hooked( - weights: &mut ModelWeights, + weights: &ModelWeights, token_ids: &[u32], index: &dyn crate::KvIndex, capture_activation: bool, @@ -31,19 +31,21 @@ pub fn predict_kquant_hidden_hooked( ); } + let mut scratch = larql_models::DequantScratch::new(); let mut h = embed_tokens_pub(weights, token_ids); let ple_inputs = precompute_per_layer_inputs(weights, &h, token_ids); let mut kv_cache: HashMap = HashMap::new(); for layer in 0..weights.num_layers { - let inserted = insert_q4k_layer_tensors(weights, index, layer)?; + let inserted = insert_q4k_layer_tensors(&mut scratch, weights, index, layer)?; let shared_kv = weights .arch .kv_shared_source_layer(layer) .and_then(|src| kv_cache.get(&src)); - let ffn_backend = crate::ffn::WeightFfn { weights }; + let view = larql_models::WeightsView::with_scratch(weights, &scratch); + let ffn_backend = crate::ffn::ViewFfn { view }; let step = run_layer_with_capture_hooked( - weights, + view, &h, layer, &ffn_backend, @@ -55,14 +57,14 @@ pub fn predict_kquant_hidden_hooked( ); let Some((h_new, _, _, kv_out)) = step else { - remove_layer_tensors(weights, inserted); + remove_layer_tensors(&mut scratch, inserted); return Err(format!("Q4K hooked forward failed at layer {layer}")); }; h = h_new; if let Some(kv) = kv_out { kv_cache.insert(layer, kv); } - remove_layer_tensors(weights, inserted); + remove_layer_tensors(&mut scratch, inserted); } Ok(h) diff --git a/crates/larql-compute/src/kquant_forward/mod.rs b/crates/larql-compute/src/kquant_forward/mod.rs index fbea2768f..517e06b49 100644 --- a/crates/larql-compute/src/kquant_forward/mod.rs +++ b/crates/larql-compute/src/kquant_forward/mod.rs @@ -13,7 +13,7 @@ //! `KvDispatch`'s CPU impl needs to call. mod cached; -mod dequant; +pub(crate) mod dequant; mod hooks; mod tensors; mod walk_ffn; @@ -165,18 +165,19 @@ mod tests { #[test] fn tensors_insert_q4k_layer_populates_dense_f32_keys() { - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); + let mut scratch = larql_models::DequantScratch::new(); let idx = make_q4k_fixture_index(&weights); - let keys = insert_q4k_layer_tensors(&mut weights, &idx, 0) + let keys = insert_q4k_layer_tensors(&mut scratch, &weights, &idx, 0) .expect("insert_q4k_layer_tensors must succeed on Q4K fixture"); // Q/K/V/O + gate/up/down = 7 keys per layer. assert_eq!(keys.len(), 7); for key in &keys { - assert!(weights.tensors.contains_key(key)); + assert!(scratch.contains_key(key)); } - remove_layer_tensors(&mut weights, keys.clone()); + remove_layer_tensors(&mut scratch, keys.clone()); for key in &keys { - assert!(!weights.tensors.contains_key(key)); + assert!(!scratch.contains_key(key)); } } @@ -186,8 +187,9 @@ mod tests { // `ok_or_else` branch in `insert_q4k_layer_tensors` fires. struct EmptyIdx; impl crate::KvIndex for EmptyIdx {} - let mut weights = make_test_q4k_weights(); - let result = insert_q4k_layer_tensors(&mut weights, &EmptyIdx, 0); + let weights = make_test_q4k_weights(); + let mut scratch = larql_models::DequantScratch::new(); + let result = insert_q4k_layer_tensors(&mut scratch, &weights, &EmptyIdx, 0); let err = result.expect_err("missing attn data must fail"); assert!(err.contains("attn")); } @@ -220,8 +222,9 @@ mod tests { dyn_idx.attn_kquant_layer_data(0).unwrap()[0].0.to_vec() }; let idx = AttnOnlyIdx { attn_bytes }; - let mut weights = make_test_q4k_weights(); - let result = insert_q4k_layer_tensors(&mut weights, &idx, 0); + let weights = make_test_q4k_weights(); + let mut scratch = larql_models::DequantScratch::new(); + let result = insert_q4k_layer_tensors(&mut scratch, &weights, &idx, 0); let err = result.expect_err("missing ffn data must fail"); assert!(err.contains("ffn")); } @@ -401,9 +404,10 @@ mod tests { #[test] fn predict_kquant_prefill_runs_end_to_end_on_cpu() { - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); + let _scratch = larql_models::DequantScratch::new(); let idx = make_q4k_fixture_index(&weights); - let (h, cache, _timings) = predict_kquant_prefill(&mut weights, &[0u32, 1, 2], &idx); + let (h, cache, _timings) = predict_kquant_prefill(&weights, &[0u32, 1, 2], &idx); assert_eq!(h.shape(), &[3, weights.hidden_size]); // One cache entry per layer, all populated by the prefill loop. assert_eq!(cache.len(), weights.num_layers); @@ -414,11 +418,12 @@ mod tests { #[test] fn predict_kquant_prefill_with_state_captures_per_layer_residuals() { - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); + let _scratch = larql_models::DequantScratch::new(); let idx = make_q4k_fixture_index(&weights); let mut state = crate::PerLayerDecodeState::with_capacity(weights.num_layers); let (h, _cache, _timings) = - predict_kquant_prefill_with_state(&mut weights, &[0u32, 1, 2], &idx, Some(&mut state)); + predict_kquant_prefill_with_state(&weights, &[0u32, 1, 2], &idx, Some(&mut state)); assert_eq!(h.shape(), &[3, weights.hidden_size]); // State captured for every layer. assert!(state.is_complete_for(weights.num_layers)); @@ -426,34 +431,37 @@ mod tests { #[test] fn predict_kquant_decode_step_uses_prefill_cache() { - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); + let _scratch = larql_models::DequantScratch::new(); let idx = make_q4k_fixture_index(&weights); // First do prefill to populate the cache. - let (_h, mut cache, _) = predict_kquant_prefill(&mut weights, &[0u32, 1, 2], &idx); + let (_h, mut cache, _) = predict_kquant_prefill(&weights, &[0u32, 1, 2], &idx); // Then decode one new token at abs_position = 3. - let result = predict_kquant_decode_step(&mut weights, 4u32, &idx, &mut cache, 3); + let result = predict_kquant_decode_step(&weights, 4u32, &idx, &mut cache, 3); let (h, _t) = result.expect("decode step succeeds with populated cache"); assert_eq!(h.shape(), &[1, weights.hidden_size]); } #[test] fn predict_kquant_decode_step_rejects_mismatched_cache() { - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); + let _scratch = larql_models::DequantScratch::new(); let idx = make_q4k_fixture_index(&weights); // Wrong-sized cache → early return None. let mut wrong = vec![None; weights.num_layers + 1]; - let result = predict_kquant_decode_step(&mut weights, 0u32, &idx, &mut wrong, 0); + let result = predict_kquant_decode_step(&weights, 0u32, &idx, &mut wrong, 0); assert!(result.is_none()); } #[test] fn predict_kquant_decode_step_direct_runs_with_q4k_fixture() { - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); + let _scratch = larql_models::DequantScratch::new(); let idx = make_q4k_fixture_index(&weights); - let (_h, mut cache, _) = predict_kquant_prefill(&mut weights, &[0u32, 1, 2], &idx); + let (_h, mut cache, _) = predict_kquant_prefill(&weights, &[0u32, 1, 2], &idx); let backend = crate::CpuBackend; let result = - predict_kquant_decode_step_direct(&mut weights, 4u32, &idx, &backend, &mut cache, 3); + predict_kquant_decode_step_direct(&weights, 4u32, &idx, &backend, &mut cache, 3); match result { Some(h) => assert_eq!(h.shape(), &[1, weights.hidden_size]), None => { @@ -465,13 +473,14 @@ mod tests { #[test] fn predict_kquant_decode_step_direct_with_state_captures_per_layer() { - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); + let _scratch = larql_models::DequantScratch::new(); let idx = make_q4k_fixture_index(&weights); - let (_h, mut cache, _) = predict_kquant_prefill(&mut weights, &[0u32, 1, 2], &idx); + let (_h, mut cache, _) = predict_kquant_prefill(&weights, &[0u32, 1, 2], &idx); let backend = crate::CpuBackend; let mut state = crate::PerLayerDecodeState::with_capacity(weights.num_layers); let _ = predict_kquant_decode_step_direct_with_state( - &mut weights, + &weights, 4u32, &idx, &backend, @@ -510,7 +519,7 @@ mod tests { /// hooks.rs. #[test] fn hooks_predict_kquant_hidden_hooked_errors_on_hybrid_moe_arch() { - let mut weights = larql_models::test_fixtures::make_test_gemma4_moe_weights(); + let weights = larql_models::test_fixtures::make_test_gemma4_moe_weights(); assert!(weights.arch.is_hybrid_moe()); // We don't need a real Q4K vindex — the function checks // is_hybrid_moe() before reading any tensor. @@ -518,7 +527,7 @@ mod tests { impl crate::KvIndex for EmptyIdx {} let mut hook = crate::forward::NoopHook; let result = - predict_kquant_hidden_hooked(&mut weights, &[0u32], &EmptyIdx, false, false, &mut hook); + predict_kquant_hidden_hooked(&weights, &[0u32], &EmptyIdx, false, false, &mut hook); let err = result.expect_err("MoE arch must early-return Err"); assert!(err.contains("dense FFN")); } @@ -590,13 +599,14 @@ mod tests { #[test] fn attention_decode_step_native_runs_on_q4k_fixture() { use ndarray::Array2; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); + let mut scratch = larql_models::DequantScratch::new(); let idx = make_q4k_fixture_index(&weights); let backend = crate::CpuBackend; // Need to insert the layer 0 tensors so weights.tensors holds // the Q/K/V/O dense f32 matrices (the helper reads them). - let inserted = - insert_q4k_layer_tensors(&mut weights, &idx, 0).expect("layer 0 q4k tensors insert"); + let inserted = insert_q4k_layer_tensors(&mut scratch, &weights, &idx, 0) + .expect("layer 0 q4k tensors insert"); let h_new = Array2::::from_shape_vec( (1, weights.hidden_size), vec![0.01; weights.hidden_size], @@ -608,7 +618,7 @@ mod tests { ); let _ = result; // Clean up so subsequent tests don't see stale tensors. - remove_layer_tensors(&mut weights, inserted); + remove_layer_tensors(&mut scratch, inserted); } /// `ffn_decode_step_native` end-to-end on Q4K weights — drives @@ -616,11 +626,12 @@ mod tests { #[test] fn ffn_decode_step_native_runs_on_q4k_fixture() { use ndarray::Array2; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); + let mut scratch = larql_models::DequantScratch::new(); let idx = make_q4k_fixture_index(&weights); let backend = crate::CpuBackend; - let inserted = - insert_q4k_layer_tensors(&mut weights, &idx, 0).expect("layer 0 q4k tensors insert"); + let inserted = insert_q4k_layer_tensors(&mut scratch, &weights, &idx, 0) + .expect("layer 0 q4k tensors insert"); let h_post_attn = Array2::::from_shape_vec( (1, weights.hidden_size), vec![0.01; weights.hidden_size], @@ -628,7 +639,7 @@ mod tests { .unwrap(); let result = ffn_decode_step_native(&weights, &idx, &backend, &h_post_attn, 0); let _ = result; - remove_layer_tensors(&mut weights, inserted); + remove_layer_tensors(&mut scratch, inserted); } /// `fused_prefill` short-circuits on `!supports_quant(Q4_K)` @@ -753,18 +764,13 @@ mod tests { // string starting with "predict_kquant_hidden_hooked currently // supports dense FFN" when the guard fires. Skip if the fixture // is dense — we cover the happy-path branch in the next test. - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); + let _scratch = larql_models::DequantScratch::new(); assert!(!weights.arch.is_hybrid_moe()); let idx = make_q4k_fixture_index(&weights); let mut hook = crate::forward::NoopHook; - let result = predict_kquant_hidden_hooked( - &mut weights, - &[0u32, 1, 2], - &idx, - false, - false, - &mut hook, - ); + let result = + predict_kquant_hidden_hooked(&weights, &[0u32, 1, 2], &idx, false, false, &mut hook); // The dense path completes — exact assertion on shape. let h = result.expect("dense Q4K hooked forward must succeed"); assert_eq!(h.shape()[1], weights.hidden_size); diff --git a/crates/larql-compute/src/kquant_forward/tensors.rs b/crates/larql-compute/src/kquant_forward/tensors.rs index 00b5087d6..e4bd606a4 100644 --- a/crates/larql-compute/src/kquant_forward/tensors.rs +++ b/crates/larql-compute/src/kquant_forward/tensors.rs @@ -1,15 +1,17 @@ -use larql_models::ModelWeights; +use larql_models::{DequantScratch, ModelWeights}; use super::dequant::dequantize_matrix; /// Insert one Q4_K/Q6_K vindex layer's attention and dense FFN tensors into -/// `weights.tensors` as dense f32 matrices. +/// the engine-owned `scratch` as dense f32 matrices (`weights` stays immutable; +/// readers resolve them via `WeightsView::with_scratch`). /// /// This is the shared research/intervention primitive behind Q4K CPU forward /// and OV/RD-style experiments. Call [`remove_layer_tensors`] with the returned /// keys after the layer has run to keep peak f32 memory bounded. pub fn insert_q4k_layer_tensors( - weights: &mut ModelWeights, + scratch: &mut DequantScratch, + weights: &ModelWeights, index: &dyn crate::KvIndex, layer: usize, ) -> Result, String> { @@ -37,27 +39,27 @@ pub fn insert_q4k_layer_tensors( let up_key = arch.ffn_up_key(layer); let down_key = arch.ffn_down_key(layer); - weights.tensors.insert( + scratch.insert( q_key.clone(), dequantize_matrix(attn[0].0, attn[0].1, q_dim, hidden).into_shared(), ); - weights.tensors.insert( + scratch.insert( k_key.clone(), dequantize_matrix(attn[1].0, attn[1].1, kv_dim, hidden).into_shared(), ); - weights.tensors.insert( + scratch.insert( v_key.clone(), dequantize_matrix(attn[2].0, attn[2].1, kv_dim, hidden).into_shared(), ); - weights.tensors.insert( + scratch.insert( o_key.clone(), dequantize_matrix(attn[3].0, attn[3].1, hidden, q_dim).into_shared(), ); - weights.tensors.insert( + scratch.insert( gate_key.clone(), dequantize_matrix(ffn[0].0, ffn[0].1, intermediate, hidden).into_shared(), ); - weights.tensors.insert( + scratch.insert( up_key.clone(), dequantize_matrix(ffn[1].0, ffn[1].1, intermediate, hidden).into_shared(), ); @@ -70,16 +72,63 @@ pub fn insert_q4k_layer_tensors( } else { dequantize_matrix(ffn[2].0, ffn[2].1, hidden, intermediate) }; - weights - .tensors - .insert(down_key.clone(), w_down.into_shared()); + scratch.insert(down_key.clone(), w_down.into_shared()); Ok(vec![q_key, k_key, v_key, o_key, gate_key, up_key, down_key]) } +/// Dequantise only a layer's **attention** Q/K/V/O Q4_K tensors into +/// `scratch`, skipping gate/up/down. The quant prefill loop pairs this with +/// [`crate::ffn::Q4kMatmulFfn`], which runs the FFN directly on the Q4_K +/// bytes via `q4k_matmul` — so materialising the (≈4× larger) f32 FFN +/// weights here would be pure waste and is the bulk of prefill's dequant +/// cost. Returns the inserted keys for [`remove_layer_tensors`]. +pub fn insert_q4k_attn_tensors( + scratch: &mut DequantScratch, + weights: &ModelWeights, + index: &dyn crate::KvIndex, + layer: usize, +) -> Result, String> { + let attn = index + .attn_kquant_layer_data(layer) + .ok_or_else(|| format!("attn Q4K slices missing for layer {layer}"))?; + + let arch = &*weights.arch; + let hidden = weights.hidden_size; + let num_q = arch.num_q_heads_for_layer(layer); + let num_kv = arch.num_kv_heads_for_layer(layer); + let head_dim = arch.head_dim_for_layer(layer); + let q_dim = num_q * head_dim; + let kv_dim = num_kv * head_dim; + + let q_key = arch.attn_q_key(layer); + let k_key = arch.attn_k_key(layer); + let v_key = arch.attn_v_key(layer); + let o_key = arch.attn_o_key(layer); + + scratch.insert( + q_key.clone(), + dequantize_matrix(attn[0].0, attn[0].1, q_dim, hidden).into_shared(), + ); + scratch.insert( + k_key.clone(), + dequantize_matrix(attn[1].0, attn[1].1, kv_dim, hidden).into_shared(), + ); + scratch.insert( + v_key.clone(), + dequantize_matrix(attn[2].0, attn[2].1, kv_dim, hidden).into_shared(), + ); + scratch.insert( + o_key.clone(), + dequantize_matrix(attn[3].0, attn[3].1, hidden, q_dim).into_shared(), + ); + + Ok(vec![q_key, k_key, v_key, o_key]) +} + /// Remove tensor keys previously returned by [`insert_q4k_layer_tensors`]. -pub fn remove_layer_tensors(weights: &mut ModelWeights, keys: Vec) { +pub fn remove_layer_tensors(scratch: &mut DequantScratch, keys: Vec) { for key in keys { - weights.tensors.remove(&key); + scratch.remove(&key); } } diff --git a/crates/larql-compute/src/kv_dispatch/cpu.rs b/crates/larql-compute/src/kv_dispatch/cpu.rs index f96b7d262..76a90968f 100644 --- a/crates/larql-compute/src/kv_dispatch/cpu.rs +++ b/crates/larql-compute/src/kv_dispatch/cpu.rs @@ -325,7 +325,7 @@ impl KvDispatch for CpuBackend { fn attention_step( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, query: &Array2, kv: &mut KvHandle, layer: usize, @@ -339,7 +339,7 @@ impl KvDispatch for CpuBackend { // back per layer to the f32 path when the index lacks Q4K attn bytes. if let Some(idx) = index.filter(|_| q4k_direct_attn_enabled()) { if let Some(proj) = crate::attention::decode::decode_step_project_q4k_direct( - weights, + weights.canonical(), query, layer, abs_position, @@ -358,7 +358,7 @@ impl KvDispatch for CpuBackend { h.append_row(k_row, v_row); let (k_all, v_all) = h.views().expect("non-empty after append"); match crate::attention::decode::decode_step_attend_q4k_direct( - weights, + weights.canonical(), query, layer, &proj.q_rope, @@ -400,7 +400,7 @@ impl KvDispatch for CpuBackend { fn attention_prefill( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, tokens_embedded: &Array2, layer: usize, _window: Option, @@ -408,7 +408,7 @@ impl KvDispatch for CpuBackend { ) -> Option<(Array2, KvHandle)> { // See `attention_step` doc for the `_index` convention. let (h_post_attn, k_rope, v) = - run_attention_with_kv_backend(weights, tokens_embedded, layer, Some(self))?; + run_attention_with_kv_backend(weights, tokens_embedded, layer, Some(self), None)?; let kv_dim = k_rope.shape()[1]; let mut handle = CpuKvHandle::new(layer, kv_dim); handle.replace_state((k_rope, v)); @@ -430,7 +430,7 @@ impl KvDispatch for CpuBackend { fn forward_from_layer( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, start_layer: usize, residuals: &ResidualHandle, token_ids: &[u32], @@ -460,7 +460,7 @@ impl KvDispatch for CpuBackend { fn coarse_prefill( &self, - weights: &mut ModelWeights, + weights: &ModelWeights, token_ids: &[u32], index: Option<&dyn crate::KvIndex>, ) -> Option<(Array2, KvHandle)> { @@ -469,7 +469,7 @@ impl KvDispatch for CpuBackend { fn coarse_prefill_with_state( &self, - weights: &mut ModelWeights, + weights: &ModelWeights, token_ids: &[u32], index: Option<&dyn crate::KvIndex>, state: Option<&mut crate::PerLayerDecodeState>, @@ -492,7 +492,7 @@ impl KvDispatch for CpuBackend { fn coarse_decode_step( &self, - weights: &mut ModelWeights, + weights: &ModelWeights, token_id: u32, index: Option<&dyn crate::KvIndex>, handle: &mut KvHandle, @@ -533,7 +533,7 @@ impl KvDispatch for CpuBackend { /// an engine asks for it on the indirect path). fn coarse_decode_step_with_state( &self, - weights: &mut ModelWeights, + weights: &ModelWeights, token_id: u32, index: Option<&dyn crate::KvIndex>, handle: &mut KvHandle, @@ -782,7 +782,7 @@ mod tests { // `cpu_residual` before `forward_from_layer` reads the shape). assert_eq!(ResidualHandleInner::shape(&OtherResidual), (1, 4)); let h = ResidualHandle::new(OtherResidual); - let _ = b.forward_from_layer(&weights, 0, &h, &[0u32]); + let _ = b.forward_from_layer(larql_models::WeightsView::dense(&weights), 0, &h, &[0u32]); } // ── Coarse default early-return branches ───────────────────────── @@ -793,8 +793,8 @@ mod tests { #[test] fn coarse_prefill_with_state_returns_none_on_empty_tokens() { let b = backend(); - let mut weights = make_test_weights(); - let result = b.coarse_prefill_with_state(&mut weights, &[], None, None); + let weights = make_test_weights(); + let result = b.coarse_prefill_with_state(&weights, &[], None, None); assert!(result.is_none()); } @@ -802,8 +802,8 @@ mod tests { #[test] fn coarse_prefill_with_state_returns_none_without_index() { let b = backend(); - let mut weights = make_test_weights(); - let result = b.coarse_prefill_with_state(&mut weights, &[0u32, 1], None, None); + let weights = make_test_weights(); + let result = b.coarse_prefill_with_state(&weights, &[0u32, 1], None, None); assert!(result.is_none()); } @@ -812,8 +812,8 @@ mod tests { #[test] fn coarse_prefill_delegates_to_with_state_variant() { let b = backend(); - let mut weights = make_test_weights(); - let result = b.coarse_prefill(&mut weights, &[], None); + let weights = make_test_weights(); + let result = b.coarse_prefill(&weights, &[], None); assert!(result.is_none()); } @@ -826,11 +826,11 @@ mod tests { use crate::test_fixtures::make_q4k_fixture_index; use larql_models::test_fixtures::make_test_q4k_weights; let b = backend(); - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let idx = make_q4k_fixture_index(&weights); let mut state = crate::PerLayerDecodeState::with_capacity(weights.num_layers); let result = - b.coarse_prefill_with_state(&mut weights, &[0u32, 1, 2], Some(&idx), Some(&mut state)); + b.coarse_prefill_with_state(&weights, &[0u32, 1, 2], Some(&idx), Some(&mut state)); let (h, handle) = result.expect("Q4K prefill succeeds"); assert_eq!(h.shape(), &[1, weights.hidden_size]); // Handle width is reported through the inner trait. @@ -846,12 +846,12 @@ mod tests { use crate::test_fixtures::make_q4k_fixture_index; use larql_models::test_fixtures::make_test_q4k_weights; let b = backend(); - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let idx = make_q4k_fixture_index(&weights); let (_h, mut handle) = b - .coarse_prefill(&mut weights, &[0u32, 1, 2], Some(&idx)) + .coarse_prefill(&weights, &[0u32, 1, 2], Some(&idx)) .expect("prefill seeds the handle"); - let result = b.coarse_decode_step(&mut weights, 4u32, Some(&idx), &mut handle, 3); + let result = b.coarse_decode_step(&weights, 4u32, Some(&idx), &mut handle, 3); let h = result.expect("decode step succeeds with populated handle"); assert_eq!(h.shape(), &[1, weights.hidden_size]); } @@ -866,14 +866,14 @@ mod tests { use crate::test_fixtures::make_q4k_fixture_index; use larql_models::test_fixtures::make_test_q4k_weights; let b = backend(); - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let idx = make_q4k_fixture_index(&weights); let (_h, mut handle) = b - .coarse_prefill(&mut weights, &[0u32, 1, 2], Some(&idx)) + .coarse_prefill(&weights, &[0u32, 1, 2], Some(&idx)) .expect("prefill seeds the handle"); let mut state = crate::PerLayerDecodeState::with_capacity(weights.num_layers); let result = b.coarse_decode_step_with_state( - &mut weights, + &weights, 4u32, Some(&idx), &mut handle, @@ -889,10 +889,9 @@ mod tests { #[test] fn coarse_decode_step_with_state_returns_none_without_index() { let b = backend(); - let mut weights = make_test_weights(); + let weights = make_test_weights(); let mut handle = KvHandle::new(CpuQ4kCacheHandle { cache: vec![None] }); - let result = - b.coarse_decode_step_with_state(&mut weights, 0u32, None, &mut handle, 0, None); + let result = b.coarse_decode_step_with_state(&weights, 0u32, None, &mut handle, 0, None); assert!(result.is_none()); } @@ -902,9 +901,9 @@ mod tests { #[test] fn coarse_decode_step_returns_none_without_index() { let b = backend(); - let mut weights = make_test_weights(); + let weights = make_test_weights(); let mut handle = KvHandle::new(CpuQ4kCacheHandle { cache: vec![None] }); - let result = b.coarse_decode_step(&mut weights, 0u32, None, &mut handle, 0); + let result = b.coarse_decode_step(&weights, 0u32, None, &mut handle, 0); assert!(result.is_none()); } @@ -915,7 +914,7 @@ mod tests { fn coarse_prefill_with_state_returns_none_on_unsupported_arch() { use crate::test_fixtures::make_q4k_fixture_index; let b = backend(); - let mut weights = larql_models::test_fixtures::make_test_gemma4_moe_weights(); + let weights = larql_models::test_fixtures::make_test_gemma4_moe_weights(); assert!( !crate::kquant_forward::supports_cached_decode(&weights), "MoE arch must not support cached decode" @@ -923,7 +922,7 @@ mod tests { // A real index is supplied so the `index?` gate passes and the // `supports_cached_decode` gate is the one that bites. let idx = make_q4k_fixture_index(&weights); - let result = b.coarse_prefill_with_state(&mut weights, &[0u32, 1], Some(&idx), None); + let result = b.coarse_prefill_with_state(&weights, &[0u32, 1], Some(&idx), None); assert!(result.is_none()); } @@ -955,14 +954,28 @@ mod tests { // Step 1 (empty prior KV) through the q4k-direct branch. let h1 = b - .attention_step(&weights, &query, &mut handle, 0, 0, Some(&idx)) + .attention_step( + larql_models::WeightsView::dense(&weights), + &query, + &mut handle, + 0, + 0, + Some(&idx), + ) .expect("q4k-direct attention_step returns Some on the Q4K fixture"); assert_eq!(h1.shape(), &[1, weights.hidden_size]); assert_eq!(handle.cached_len(), 1, "KV grew by 1"); // Step 2 (prior KV present) keeps using the q4k-direct branch. let h2 = b - .attention_step(&weights, &query, &mut handle, 0, 1, Some(&idx)) + .attention_step( + larql_models::WeightsView::dense(&weights), + &query, + &mut handle, + 0, + 1, + Some(&idx), + ) .expect("second q4k-direct attention_step succeeds"); assert_eq!(h2.shape(), &[1, weights.hidden_size]); assert_eq!(handle.cached_len(), 2, "KV grew to 2"); @@ -995,7 +1008,14 @@ mod tests { // f32-BLAS path on `weights.tensors` (the fixture carries f32 attn // tensors), which succeeds. let h = b - .attention_step(&weights, &query, &mut handle, 0, 0, Some(&EmptyIdx)) + .attention_step( + larql_models::WeightsView::dense(&weights), + &query, + &mut handle, + 0, + 0, + Some(&EmptyIdx), + ) .expect("f32 fallback succeeds when the index lacks Q4K attn bytes"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert_eq!(handle.cached_len(), 1); diff --git a/crates/larql-compute/src/kv_dispatch/mod.rs b/crates/larql-compute/src/kv_dispatch/mod.rs index 182ffb409..542c93dbb 100644 --- a/crates/larql-compute/src/kv_dispatch/mod.rs +++ b/crates/larql-compute/src/kv_dispatch/mod.rs @@ -264,7 +264,7 @@ pub trait KvDispatch { /// See `docs/specs/kv-dispatch-quantization.md`. fn attention_step( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, query: &Array2, kv: &mut KvHandle, layer: usize, @@ -288,7 +288,7 @@ pub trait KvDispatch { #[allow(clippy::too_many_arguments)] fn attention_step_windowed( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, query: &Array2, kv: &mut KvHandle, layer: usize, @@ -313,7 +313,7 @@ pub trait KvDispatch { /// `index` follows the same convention as [`Self::attention_step`]. fn attention_prefill( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, tokens_embedded: &Array2, layer: usize, window: Option, @@ -332,7 +332,7 @@ pub trait KvDispatch { /// [`crate::MatMul`] directly. fn recompute_kv_from_residuals( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, residuals: &Array2, layer: usize, ) -> Option { @@ -368,7 +368,7 @@ pub trait KvDispatch { /// pre-crystal layers when boundaries are available. fn forward_from_layer( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, start_layer: usize, residuals: &ResidualHandle, token_ids: &[u32], @@ -410,7 +410,7 @@ pub trait KvDispatch { /// shared state. fn coarse_prefill( &self, - weights: &mut ModelWeights, + weights: &ModelWeights, token_ids: &[u32], index: Option<&dyn crate::KvIndex>, ) -> Option<(Array2, KvHandle)> { @@ -422,7 +422,7 @@ pub trait KvDispatch { /// by a prior [`Self::coarse_prefill`] on the same backend. fn coarse_decode_step( &self, - weights: &mut ModelWeights, + weights: &ModelWeights, token_id: u32, index: Option<&dyn crate::KvIndex>, handle: &mut KvHandle, @@ -450,7 +450,7 @@ pub trait KvDispatch { /// CPU walk. fn coarse_prefill_with_state( &self, - weights: &mut ModelWeights, + weights: &ModelWeights, token_ids: &[u32], index: Option<&dyn crate::KvIndex>, state: Option<&mut PerLayerDecodeState>, @@ -482,7 +482,7 @@ pub trait KvDispatch { /// engine. fn coarse_decode_step_with_state( &self, - weights: &mut ModelWeights, + weights: &ModelWeights, token_id: u32, index: Option<&dyn crate::KvIndex>, handle: &mut KvHandle, @@ -506,7 +506,7 @@ pub trait KvDispatch { #[allow(clippy::too_many_arguments)] fn coarse_decode_step_with_state_masked( &self, - weights: &mut ModelWeights, + weights: &ModelWeights, token_id: u32, index: Option<&dyn crate::KvIndex>, handle: &mut KvHandle, @@ -747,7 +747,14 @@ mod tests { let mut handle = stub_kv_handle(0, weights.hidden_size); let query = Array2::zeros((1, weights.hidden_size)); assert!(backend - .attention_step(&weights, &query, &mut handle, 0, 0, None) + .attention_step( + larql_models::WeightsView::dense(&weights), + &query, + &mut handle, + 0, + 0, + None + ) .is_none()); } @@ -761,7 +768,15 @@ mod tests { let mut handle = stub_kv_handle(0, weights.hidden_size); let query = Array2::zeros((1, weights.hidden_size)); assert!(backend - .attention_step_windowed(&weights, &query, &mut handle, 0, 0, 4, None) + .attention_step_windowed( + larql_models::WeightsView::dense(&weights), + &query, + &mut handle, + 0, + 0, + 4, + None + ) .is_none()); } @@ -771,7 +786,13 @@ mod tests { let weights = larql_models::test_fixtures::make_test_weights(); let tokens = Array2::zeros((2, weights.hidden_size)); assert!(backend - .attention_prefill(&weights, &tokens, 0, None, None) + .attention_prefill( + larql_models::WeightsView::dense(&weights), + &tokens, + 0, + None, + None + ) .is_none()); } @@ -781,7 +802,7 @@ mod tests { let weights = larql_models::test_fixtures::make_test_weights(); let residuals = Array2::zeros((1, weights.hidden_size)); assert!(backend - .recompute_kv_from_residuals(&weights, &residuals, 0) + .recompute_kv_from_residuals(larql_models::WeightsView::dense(&weights), &residuals, 0) .is_none()); } @@ -798,7 +819,12 @@ mod tests { let weights = larql_models::test_fixtures::make_test_weights(); let residuals = stub_residual_handle(1, weights.hidden_size); assert!(backend - .forward_from_layer(&weights, 0, &residuals, &[0u32]) + .forward_from_layer( + larql_models::WeightsView::dense(&weights), + 0, + &residuals, + &[0u32] + ) .is_none()); } @@ -907,37 +933,37 @@ mod tests { #[test] fn default_coarse_prefill_returns_none() { - let mut weights = make_test_weights(); + let weights = make_test_weights(); let backend = StubKvBackend; - let result = backend.coarse_prefill(&mut weights, &[0u32, 1], None); + let result = backend.coarse_prefill(&weights, &[0u32, 1], None); assert!(result.is_none()); } #[test] fn default_coarse_decode_step_returns_none() { - let mut weights = make_test_weights(); + let weights = make_test_weights(); let backend = StubKvBackend; let mut handle = KvHandle::new(StubKvInner { len: 0, dim: weights.hidden_size, }); - let result = backend.coarse_decode_step(&mut weights, 0, None, &mut handle, 0); + let result = backend.coarse_decode_step(&weights, 0, None, &mut handle, 0); assert!(result.is_none()); } #[test] fn default_coarse_prefill_with_state_delegates_to_coarse_prefill() { - let mut weights = make_test_weights(); + let weights = make_test_weights(); let backend = StubKvBackend; let mut state = PerLayerDecodeState::with_capacity(weights.num_layers); let result = - backend.coarse_prefill_with_state(&mut weights, &[0u32, 1], None, Some(&mut state)); + backend.coarse_prefill_with_state(&weights, &[0u32, 1], None, Some(&mut state)); assert!(result.is_none()); } #[test] fn default_coarse_decode_step_with_state_delegates() { - let mut weights = make_test_weights(); + let weights = make_test_weights(); let backend = StubKvBackend; let mut handle = KvHandle::new(StubKvInner { len: 0, @@ -945,7 +971,7 @@ mod tests { }); let mut state = PerLayerDecodeState::with_capacity(weights.num_layers); let result = backend.coarse_decode_step_with_state( - &mut weights, + &weights, 0, None, &mut handle, @@ -957,7 +983,7 @@ mod tests { #[test] fn default_coarse_decode_step_with_state_masked_delegates() { - let mut weights = make_test_weights(); + let weights = make_test_weights(); let backend = StubKvBackend; let mut handle = KvHandle::new(StubKvInner { len: 0, @@ -970,7 +996,7 @@ mod tests { ] { let mut state = PerLayerDecodeState::with_capacity(weights.num_layers); let result = backend.coarse_decode_step_with_state_masked( - &mut weights, + &weights, 0, None, &mut handle, diff --git a/crates/larql-compute/src/pipeline.rs b/crates/larql-compute/src/pipeline.rs index 4bc779105..e29f8d117 100644 --- a/crates/larql-compute/src/pipeline.rs +++ b/crates/larql-compute/src/pipeline.rs @@ -31,6 +31,14 @@ pub enum QuantFormat { BF16, // raw bfloat16 (2 bytes per value, no quantization scales) F16, // raw float16 (2 bytes per value) F32, // raw float32 (4 bytes per value) + /// BitNet 1.58-bit ternary (GGML I2_S, type 36): 4 trits/byte packed + /// row-major (`cols/4` bytes per row) plus a separate per-channel f32 + /// scale array. Unlike the block-quant formats, the weight is NOT a + /// flat `&[u8]` block stream — it is carried by + /// [`crate::cpu::ops::ternary_matvec::BitLinearWeight`] (bytes + scales) + /// and served by the dedicated `ternary_matvec` dispatch, not the + /// block-quant `quant_matvec` path (which has no per-channel-scale input). + I2S, } impl QuantFormat { @@ -111,9 +119,37 @@ impl QuantFormat { "BF16" => Self::BF16, "F16" => Self::F16, "F32" => Self::F32, + // BitNet ternary (GGML type 36). The vindex bitnet sidecar tags + // its I2_S weight stream with this so the registry recognises it. + "I2_S" => Self::I2S, _ => return None, }) } + + /// Inverse of [`Self::from_registry_tag`] — the canonical registry tag + /// string for this format. `from_registry_tag(f.registry_tag()) == Some(f)` + /// for every variant. Used by writers that record the per-tensor format + /// tag into the weight manifest / index. + pub fn registry_tag(self) -> &'static str { + match self { + Self::Q4_0 => "Q4_0", + Self::Q4_K => "Q4_K", + Self::Q4_KF => "Q4_KF", + Self::Q6_K => "Q6_K", + Self::Q8_0 => "Q8_0", + Self::BF16 => "BF16", + Self::F16 => "F16", + Self::F32 => "F32", + Self::I2S => "I2_S", + } + } + + /// Whether this format is BitNet ternary (I2_S). Served by the dedicated + /// `ternary_matvec` path with a [`crate::cpu::ops::ternary_matvec::BitLinearWeight`], + /// never the block-quant `quant_matvec` dispatch. + pub fn is_ternary(self) -> bool { + matches!(self, Self::I2S) + } } /// A quantized weight matrix — raw bytes with format tag. @@ -815,6 +851,40 @@ mod tests { assert_eq!(Activation::from(false), Activation::Silu); } + #[test] + fn registry_tag_round_trips_every_variant() { + // registry_tag() is the exact inverse of from_registry_tag() for + // every variant — including the new ternary I2_S. + for f in [ + QuantFormat::Q4_0, + QuantFormat::Q4_K, + QuantFormat::Q4_KF, + QuantFormat::Q6_K, + QuantFormat::Q8_0, + QuantFormat::BF16, + QuantFormat::F16, + QuantFormat::F32, + QuantFormat::I2S, + ] { + assert_eq!(QuantFormat::from_registry_tag(f.registry_tag()), Some(f)); + } + } + + #[test] + fn i2s_tag_and_family_predicates() { + assert_eq!( + QuantFormat::from_registry_tag("I2_S"), + Some(QuantFormat::I2S) + ); + assert_eq!(QuantFormat::I2S.registry_tag(), "I2_S"); + assert!(QuantFormat::I2S.is_ternary()); + assert!(!QuantFormat::Q4_K.is_ternary()); + // I2_S is none of the block-quant families and has no flat block layout. + assert!(!QuantFormat::I2S.is_kquant_family()); + assert!(!QuantFormat::I2S.is_legacy_q8()); + assert!(QuantFormat::I2S.packed_block_layout().is_none()); + } + #[test] fn is_gated_matches_ffn_type() { let norms = [1.0f32; 4]; diff --git a/crates/larql-compute/src/pipeline_layer.rs b/crates/larql-compute/src/pipeline_layer.rs index 2250f69d5..497957d22 100644 --- a/crates/larql-compute/src/pipeline_layer.rs +++ b/crates/larql-compute/src/pipeline_layer.rs @@ -294,12 +294,13 @@ pub fn build_moe_weights<'a>( /// aliasing to Q4_K. Lifted from inside `resolve_attn_weights` so the /// mapping is unit-testable in isolation. fn attn_str_to_format(s: &str) -> QuantFormat { - match s { - "Q4_K" => QuantFormat::Q4_K, - "Q6_K" => QuantFormat::Q6_K, - other => panic!( - "resolve_attn_weights: registry tag {other:?} has no compute::QuantFormat mapping" - ), + // Source the tag→format mapping from `from_registry_tag` (single source + // of truth) but keep the attention surface's supported-subset guard: + // only Q4_K / Q6_K have a q8k attention matvec, so anything else still + // fails loudly rather than silently aliasing. + match QuantFormat::from_registry_tag(s) { + Some(f @ (QuantFormat::Q4_K | QuantFormat::Q6_K)) => f, + _ => panic!("resolve_attn_weights: registry tag {s:?} has no compute::QuantFormat mapping"), } } @@ -377,14 +378,15 @@ pub fn resolve_attn_weights<'a>( /// didn't emit per-matrix tags. Lifted from inside `resolve_ffn_weights` /// so the mapping is unit-testable in isolation. fn ffn_str_to_format(s: &str, fallback: QuantFormat) -> QuantFormat { - match s { - "Q4_K" => QuantFormat::Q4_K, - "Q6_K" => QuantFormat::Q6_K, - "Q4_0" => QuantFormat::Q4_0, - "" => fallback, - other => panic!( - "resolve_ffn_weights: registry tag {other:?} has no compute::QuantFormat mapping" - ), + // Empty tag = legacy uniform-stride writer (no per-matrix tags) → fallback. + if s.is_empty() { + return fallback; + } + // Mapping from `from_registry_tag` (single source of truth); the FFN + // surface supports Q4_K / Q6_K / Q4_0, anything else fails loudly. + match QuantFormat::from_registry_tag(s) { + Some(f @ (QuantFormat::Q4_K | QuantFormat::Q6_K | QuantFormat::Q4_0)) => f, + _ => panic!("resolve_ffn_weights: registry tag {s:?} has no compute::QuantFormat mapping"), } } diff --git a/crates/larql-inference/examples/apollo_rd_backend.rs b/crates/larql-inference/examples/apollo_rd_backend.rs index b20c24d08..c93b94f75 100644 --- a/crates/larql-inference/examples/apollo_rd_backend.rs +++ b/crates/larql-inference/examples/apollo_rd_backend.rs @@ -593,7 +593,13 @@ fn logits_from_boundary( boundary: &[f32], crystal: usize, ) -> Vec { - let raw = forward_from_layer(weights, tokens, boundary, crystal, None); + 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) diff --git a/crates/larql-inference/examples/ave_direct_layer_bisect.rs b/crates/larql-inference/examples/ave_direct_layer_bisect.rs index b64e99edf..fa9427af3 100644 --- a/crates/larql-inference/examples/ave_direct_layer_bisect.rs +++ b/crates/larql-inference/examples/ave_direct_layer_bisect.rs @@ -70,7 +70,7 @@ fn main() { 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(&mut weights, &full_ids, &index, Some(&mut gold)); + 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); diff --git a/crates/larql-inference/examples/ave_direct_step_parity.rs b/crates/larql-inference/examples/ave_direct_step_parity.rs index 5a52f3c55..0cc4dd94b 100644 --- a/crates/larql-inference/examples/ave_direct_step_parity.rs +++ b/crates/larql-inference/examples/ave_direct_step_parity.rs @@ -65,14 +65,9 @@ fn main() { ); let abs_position = prompt_ids.len(); - let (h_staged, _) = predict_kquant_decode_step( - &mut weights, - first_id, - &index, - &mut cache_staged, - abs_position, - ) - .expect("staged step"); + 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, diff --git a/crates/larql-inference/examples/ave_q4k_row_audit.rs b/crates/larql-inference/examples/ave_q4k_row_audit.rs index 5a1130b2f..37df3120a 100644 --- a/crates/larql-inference/examples/ave_q4k_row_audit.rs +++ b/crates/larql-inference/examples/ave_q4k_row_audit.rs @@ -12,7 +12,7 @@ //! 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, remove_layer_tensors}; +use larql_inference::vindex::{insert_q4k_layer_tensors_resident, remove_layer_tensors_resident}; fn main() { if std::env::var("LARQL_F16_PROBE").is_ok() { @@ -76,7 +76,8 @@ fn main() { // Staged tensor for (c). let k_bytes_owned = k_bytes.to_vec(); - let inserted = insert_q4k_layer_tensors(&mut weights, &index, layer).expect("insert"); + 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()); @@ -187,7 +188,7 @@ fn main() { .count(); println!(" elems differing in this block: {n_diff}/256"); } - remove_layer_tensors(&mut weights, inserted); + remove_layer_tensors_resident(&mut weights, inserted); } } diff --git a/crates/larql-inference/examples/bitnet_e2e.rs b/crates/larql-inference/examples/bitnet_e2e.rs new file mode 100644 index 000000000..d06d37788 --- /dev/null +++ b/crates/larql-inference/examples/bitnet_e2e.rs @@ -0,0 +1,53 @@ +//! 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/debug_layers.rs b/crates/larql-inference/examples/debug_layers.rs index f38c605c2..6d4649f66 100644 --- a/crates/larql-inference/examples/debug_layers.rs +++ b/crates/larql-inference/examples/debug_layers.rs @@ -64,7 +64,11 @@ fn main() -> Result<(), Box> { let mut h_cpu = h.clone(); for layer in 0..weights.num_layers { let (h_pa, _, _) = larql_inference::attention::run_attention_block_gpu( - weights, &h_cpu, layer, false, None, + larql_inference::WeightsView::dense(weights), + &h_cpu, + layer, + false, + None, ) .unwrap(); let (h_out, _) = diff --git a/crates/larql-inference/examples/decode_vs_prefill.rs b/crates/larql-inference/examples/decode_vs_prefill.rs index a7fa24365..c8b5aaae5 100644 --- a/crates/larql-inference/examples/decode_vs_prefill.rs +++ b/crates/larql-inference/examples/decode_vs_prefill.rs @@ -62,7 +62,7 @@ fn main() -> Result<(), Box> { // Separate weight handles so CPU's per-layer dequant inserts don't // race with Metal's forward on a shared ModelWeights. 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 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 prompt_ids = larql_inference::encode_prompt(&tokenizer, &*w_metal.arch, &wrap.prompt)?; @@ -137,7 +137,7 @@ fn main() -> Result<(), Box> { // twice. let t0 = Instant::now(); let cpu_hidden_full = - larql_inference::vindex::predict_kquant_hidden(&mut w_cpu, &appended_ids, &index, None); + larql_inference::vindex::predict_kquant_hidden(&w_cpu, &appended_ids, &index, None); let cpu_ms = t0.elapsed().as_secs_f64() * 1000.0; let cpu_last = cpu_hidden_full .row(cpu_hidden_full.nrows().saturating_sub(1)) @@ -216,9 +216,8 @@ fn main() -> Result<(), Box> { // Re-run CPU full-prefill with the layer-dump env var set so we can // walk the two paths side by side. Cheap relative to the Metal // prefill we already paid for. - let mut w_cpu2 = larql_vindex::load_model_weights_kquant(&vindex_path, &mut cb)?; - let _ = - larql_inference::vindex::predict_kquant_hidden(&mut w_cpu2, &appended_ids, &index, None); + let w_cpu2 = larql_vindex::load_model_weights_kquant(&vindex_path, &mut cb)?; + let _ = larql_inference::vindex::predict_kquant_hidden(&w_cpu2, &appended_ids, &index, None); println!( " B) Metal prefill({} tok) + decode(1 tok) took {:>5.1} + {:>5.1} ms", diff --git a/crates/larql-inference/examples/fr1_topk_fuzzy_router.rs b/crates/larql-inference/examples/fr1_topk_fuzzy_router.rs index 02f6f259c..2d35d2e66 100644 --- a/crates/larql-inference/examples/fr1_topk_fuzzy_router.rs +++ b/crates/larql-inference/examples/fr1_topk_fuzzy_router.rs @@ -29,7 +29,7 @@ //! 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; +use larql_inference::vindex::insert_q4k_layer_tensors_resident; use larql_inference::{capture_residuals, load_tokenizer}; use larql_vindex::KnnStore; use std::collections::HashMap; @@ -303,7 +303,7 @@ fn main() { 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(&mut weights, &index, layer).expect("dequant"); + insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant"); } // Capture residuals for all three phrasings at the layer sweep, one forward diff --git a/crates/larql-inference/examples/fr2_two_tier_router.rs b/crates/larql-inference/examples/fr2_two_tier_router.rs index 49e261839..42f7dc0ce 100644 --- a/crates/larql-inference/examples/fr2_two_tier_router.rs +++ b/crates/larql-inference/examples/fr2_two_tier_router.rs @@ -18,7 +18,7 @@ //! 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; +use larql_inference::vindex::insert_q4k_layer_tensors_resident; use larql_inference::{capture_residuals, load_tokenizer}; use larql_vindex::KnnStore; use std::collections::HashMap; @@ -183,7 +183,7 @@ fn main() { let tok = load_tokenizer(&dir).expect("tokenizer"); eprintln!("Dequantising {} layers ...", weights.num_layers); for layer in 0..weights.num_layers { - insert_q4k_layer_tensors(&mut weights, &index, layer).expect("dequant"); + insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant"); } let cap = |prompt: &str| -> HashMap> { diff --git a/crates/larql-inference/examples/fr3_relation_address.rs b/crates/larql-inference/examples/fr3_relation_address.rs index 62c77de44..d7f81df6e 100644 --- a/crates/larql-inference/examples/fr3_relation_address.rs +++ b/crates/larql-inference/examples/fr3_relation_address.rs @@ -20,7 +20,7 @@ //! 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; +use larql_inference::vindex::insert_q4k_layer_tensors_resident; use larql_inference::{capture_residuals, load_tokenizer}; use larql_vindex::KnnStore; use std::collections::HashMap; @@ -229,7 +229,7 @@ fn main() { let tok = load_tokenizer(&dir).expect("tokenizer"); eprintln!("Dequantising {} layers ...", weights.num_layers); for layer in 0..weights.num_layers { - insert_q4k_layer_tensors(&mut weights, &index, layer).expect("dequant"); + insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant"); } let cap = |prompt: &str| -> HashMap> { diff --git a/crates/larql-inference/examples/fr3_template_ablation.rs b/crates/larql-inference/examples/fr3_template_ablation.rs index 13bb7969d..488612af4 100644 --- a/crates/larql-inference/examples/fr3_template_ablation.rs +++ b/crates/larql-inference/examples/fr3_template_ablation.rs @@ -16,7 +16,7 @@ //! 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; +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; @@ -77,7 +77,7 @@ fn main() { let tok = load_tokenizer(&dir).expect("tokenizer"); eprintln!("Dequantising {} layers ...", weights.num_layers); for l in 0..weights.num_layers { - insert_q4k_layer_tensors(&mut weights, &index, l).expect("dequant"); + insert_q4k_layer_tensors_resident(&mut weights, &index, l).expect("dequant"); } let cap = |prompt: &str| -> LayerRes { diff --git a/crates/larql-inference/examples/fr_early_exit_bench.rs b/crates/larql-inference/examples/fr_early_exit_bench.rs index 61722f3d5..4fd603700 100644 --- a/crates/larql-inference/examples/fr_early_exit_bench.rs +++ b/crates/larql-inference/examples/fr_early_exit_bench.rs @@ -19,7 +19,7 @@ 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; +use larql_inference::vindex::insert_q4k_layer_tensors_resident; use larql_vindex::PatchedVindex; use std::time::Instant; @@ -107,7 +107,7 @@ fn main() { .min(last); eprintln!("Dequantising {num_layers} layers to f32 ..."); for layer in 0..num_layers { - insert_q4k_layer_tensors(&mut weights, &index, layer).expect("dequant"); + 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); diff --git a/crates/larql-inference/examples/fr_early_exit_decode_projection.rs b/crates/larql-inference/examples/fr_early_exit_decode_projection.rs index db65a408e..61c01719e 100644 --- a/crates/larql-inference/examples/fr_early_exit_decode_projection.rs +++ b/crates/larql-inference/examples/fr_early_exit_decode_projection.rs @@ -25,7 +25,7 @@ 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; +use larql_inference::vindex::insert_q4k_layer_tensors_resident; use larql_vindex::PatchedVindex; use std::time::Instant; @@ -91,7 +91,7 @@ fn main() { .min(last); eprintln!("Dequantising {num_layers} layers to f32 ..."); for layer in 0..num_layers { - insert_q4k_layer_tensors(&mut weights, &index, layer).expect("dequant"); + insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant"); } let patched = PatchedVindex::new(index); diff --git a/crates/larql-inference/examples/fr_early_exit_parity.rs b/crates/larql-inference/examples/fr_early_exit_parity.rs index 6493b0699..9393ace10 100644 --- a/crates/larql-inference/examples/fr_early_exit_parity.rs +++ b/crates/larql-inference/examples/fr_early_exit_parity.rs @@ -30,7 +30,7 @@ use larql_inference::forward::{ apply_knn_override_verified, KNN_COSINE_THRESHOLD, KNN_VERIFY_TOPK, }; -use larql_inference::vindex::insert_q4k_layer_tensors; +use larql_inference::vindex::insert_q4k_layer_tensors_resident; use larql_inference::{capture_residuals, load_tokenizer}; use larql_vindex::KnnStore; use std::time::Instant; @@ -126,7 +126,7 @@ fn main() { .min(last); eprintln!("Dequantising {num_layers} layers to f32 ..."); for layer in 0..num_layers { - insert_q4k_layer_tensors(&mut weights, &index, layer).expect("dequant"); + 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)); diff --git a/crates/larql-inference/examples/fr_early_exit_probe.rs b/crates/larql-inference/examples/fr_early_exit_probe.rs index ae9634910..ef52081d7 100644 --- a/crates/larql-inference/examples/fr_early_exit_probe.rs +++ b/crates/larql-inference/examples/fr_early_exit_probe.rs @@ -32,7 +32,7 @@ use larql_inference::forward::{ apply_knn_override_verified, KNN_COSINE_THRESHOLD, KNN_VERIFY_TOPK, }; -use larql_inference::vindex::insert_q4k_layer_tensors; +use larql_inference::vindex::insert_q4k_layer_tensors_resident; use larql_inference::{capture_residuals, load_tokenizer}; use larql_vindex::KnnStore; use std::collections::HashMap; @@ -131,7 +131,7 @@ fn main() { let num_layers = weights.num_layers; eprintln!("Dequantising {num_layers} layers to f32 ..."); for layer in 0..num_layers { - insert_q4k_layer_tensors(&mut weights, &index, layer).expect("dequant"); + insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant"); } // Sweep the whole stack — capture_residuals returns every requested layer diff --git a/crates/larql-inference/examples/fr_routing_gain.rs b/crates/larql-inference/examples/fr_routing_gain.rs index 4c5b5cd3c..fa71705c2 100644 --- a/crates/larql-inference/examples/fr_routing_gain.rs +++ b/crates/larql-inference/examples/fr_routing_gain.rs @@ -16,7 +16,7 @@ //! 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, WalkFfn}; +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, @@ -101,7 +101,7 @@ fn main() { let tok = load_tokenizer(&dir).expect("tokenizer"); eprintln!("Dequantising {} layers ...", weights.num_layers); for l in 0..weights.num_layers { - insert_q4k_layer_tensors(&mut weights, &index, l).expect("dequant"); + insert_q4k_layer_tensors_resident(&mut weights, &index, l).expect("dequant"); } // Install novel facts at the resolved layer. diff --git a/crates/larql-inference/examples/predict_from_residual.rs b/crates/larql-inference/examples/predict_from_residual.rs index a3da1559c..f02bb929e 100644 --- a/crates/larql-inference/examples/predict_from_residual.rs +++ b/crates/larql-inference/examples/predict_from_residual.rs @@ -84,8 +84,12 @@ fn run_full_forward( } } - let (h_post_attn, _) = forward::layer::run_attention_with_kv_cache(weights, &h, layer) - .ok_or_else(|| format!("attention failed at layer {layer}"))?; + 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, diff --git a/crates/larql-inference/examples/probe_contribution_distribution.rs b/crates/larql-inference/examples/probe_contribution_distribution.rs index 37c4377fe..ad8204b54 100644 --- a/crates/larql-inference/examples/probe_contribution_distribution.rs +++ b/crates/larql-inference/examples/probe_contribution_distribution.rs @@ -299,9 +299,12 @@ fn main() -> Result<(), Box> { 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(weights, &h, layer) - .ok_or_else(|| format!("attention failed at layer {layer}"))?; + 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( @@ -348,8 +351,12 @@ fn main() -> Result<(), Box> { 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(weights, &h, layer) - .ok_or_else(|| format!("attention failed at layer {layer}"))?; + 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; diff --git a/crates/larql-inference/examples/probe_residual_stream.rs b/crates/larql-inference/examples/probe_residual_stream.rs index 2dbca9107..3e8dc3c22 100644 --- a/crates/larql-inference/examples/probe_residual_stream.rs +++ b/crates/larql-inference/examples/probe_residual_stream.rs @@ -174,8 +174,12 @@ fn main() -> Result<(), Box> { writer.write_all(&v.to_le_bytes())?; } - let (h_post_attn, _) = forward::layer::run_attention_with_kv_cache(weights, &h, layer) - .ok_or_else(|| format!("attention failed at layer {layer}"))?; + 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( diff --git a/crates/larql-inference/examples/profile_faithful_walk.rs b/crates/larql-inference/examples/profile_faithful_walk.rs index 14c53534f..0c511a35e 100644 --- a/crates/larql-inference/examples/profile_faithful_walk.rs +++ b/crates/larql-inference/examples/profile_faithful_walk.rs @@ -92,8 +92,12 @@ fn run_profile_pass( for layer in 0..weights.num_layers { let t = Instant::now(); - let (h_post_attn, _) = forward::layer::run_attention_with_kv_cache(weights, &h, layer) - .ok_or_else(|| format!("attention failed at layer {layer}"))?; + 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}"))?; stats.layers[layer].attention_ms = t.elapsed().as_secs_f64() * 1000.0; if let Some(index) = index { diff --git a/crates/larql-inference/examples/profile_overhead.rs b/crates/larql-inference/examples/profile_overhead.rs index 58727129d..450416edb 100644 --- a/crates/larql-inference/examples/profile_overhead.rs +++ b/crates/larql-inference/examples/profile_overhead.rs @@ -127,7 +127,12 @@ fn main() -> Result<(), Box> { // ── GQA attention (one layer) ── let t0 = Instant::now(); for _ in 0..10 { - let _ = larql_inference::attention::run_attention_block(weights, &h, 0, false); + let _ = larql_inference::attention::run_attention_block( + larql_inference::WeightsView::dense(weights), + &h, + 0, + false, + ); } let attn_ms = t0.elapsed().as_secs_f64() * 1000.0 / 10.0; println!("Attention block: {attn_ms:.1}ms (proj + norm + RoPE + fused attn + residual)"); @@ -135,8 +140,13 @@ fn main() -> Result<(), Box> { // ── Full layer (attention + FFN) ── let t0 = Instant::now(); for _ in 0..10 { - let (h_post_attn, _, _) = - larql_inference::attention::run_attention_block(weights, &h, 0, false).unwrap(); + let (h_post_attn, _, _) = larql_inference::attention::run_attention_block( + larql_inference::WeightsView::dense(weights), + &h, + 0, + false, + ) + .unwrap(); let h_ffn = apply_norm( weights, &h_post_attn, diff --git a/crates/larql-inference/examples/profile_walk_accuracy.rs b/crates/larql-inference/examples/profile_walk_accuracy.rs index 645839122..a026a26e0 100644 --- a/crates/larql-inference/examples/profile_walk_accuracy.rs +++ b/crates/larql-inference/examples/profile_walk_accuracy.rs @@ -43,9 +43,14 @@ fn main() { let token_ids: Vec = encoding.get_ids().to_vec(); let mut h = larql_inference::forward::embed_tokens_pub(weights, &token_ids); for layer in 0..14 { - let (h_pa, _, _) = - larql_inference::attention::run_attention_block_gpu(weights, &h, layer, false, None) - .unwrap(); + let (h_pa, _, _) = larql_inference::attention::run_attention_block_gpu( + larql_inference::WeightsView::dense(weights), + &h, + layer, + false, + None, + ) + .unwrap(); let dense_ffn = larql_inference::WeightFfn { weights }; let (h_out, _) = larql_inference::forward::run_ffn(weights, &h_pa, layer, &dense_ffn, false); @@ -53,8 +58,14 @@ fn main() { } // Get the post-attention state at L14 - let (h_post_attn, _, _) = - larql_inference::attention::run_attention_block_gpu(weights, &h, 14, false, None).unwrap(); + let (h_post_attn, _, _) = larql_inference::attention::run_attention_block_gpu( + larql_inference::WeightsView::dense(weights), + &h, + 14, + false, + None, + ) + .unwrap(); // Dense FFN output (ground truth) let dense_ffn = larql_inference::WeightFfn { weights }; diff --git a/crates/larql-inference/examples/profile_walk_ffn.rs b/crates/larql-inference/examples/profile_walk_ffn.rs index c46cc2c24..7cba2bc16 100644 --- a/crates/larql-inference/examples/profile_walk_ffn.rs +++ b/crates/larql-inference/examples/profile_walk_ffn.rs @@ -67,9 +67,14 @@ fn main() { let token_ids: Vec = encoding.get_ids().to_vec(); let mut h = larql_inference::forward::embed_tokens_pub(weights, &token_ids); for layer in 0..14 { - let (h_post_attn, _, _) = - larql_inference::attention::run_attention_block_gpu(weights, &h, layer, false, None) - .unwrap(); + let (h_post_attn, _, _) = larql_inference::attention::run_attention_block_gpu( + larql_inference::WeightsView::dense(weights), + &h, + layer, + false, + None, + ) + .unwrap(); let dense_ffn = larql_inference::WeightFfn { weights }; let (h_out, _) = larql_inference::forward::run_ffn(weights, &h_post_attn, layer, &dense_ffn, false); diff --git a/crates/larql-inference/examples/walk_ffn_accuracy.rs b/crates/larql-inference/examples/walk_ffn_accuracy.rs index fd5f4113a..795a8bfa0 100644 --- a/crates/larql-inference/examples/walk_ffn_accuracy.rs +++ b/crates/larql-inference/examples/walk_ffn_accuracy.rs @@ -20,7 +20,7 @@ //! //! Usage: `cargo run --release --example walk_ffn_accuracy -- [VINDEX_DIR]` -use larql_inference::vindex::{insert_q4k_layer_tensors, WalkFfn, WalkFfnConfig}; +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; @@ -138,7 +138,7 @@ fn main() { weights.num_layers ); for layer in 0..weights.num_layers { - insert_q4k_layer_tensors(&mut weights, &index, layer).expect("dequant layer"); + insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant layer"); } let feats = index.num_features(weights.num_layers / 2); diff --git a/crates/larql-inference/examples/walk_ffn_cell_router.rs b/crates/larql-inference/examples/walk_ffn_cell_router.rs index 60e456156..7c63e9a6d 100644 --- a/crates/larql-inference/examples/walk_ffn_cell_router.rs +++ b/crates/larql-inference/examples/walk_ffn_cell_router.rs @@ -22,7 +22,9 @@ //! Usage: `cargo run --release --example walk_ffn_cell_router -- [VINDEX_DIR]` use larql_inference::ffn::FfnBackend; -use larql_inference::vindex::{insert_q4k_layer_tensors, CellRouter, WalkFfn, WalkFfnConfig}; +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; @@ -293,7 +295,7 @@ fn main() { 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(&mut weights, &index, layer).expect("dequant layer"); + insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant layer"); } let nl = weights.num_layers; diff --git a/crates/larql-inference/examples/walk_ffn_decode_timing.rs b/crates/larql-inference/examples/walk_ffn_decode_timing.rs index c4cffa5da..bc7eb6357 100644 --- a/crates/larql-inference/examples/walk_ffn_decode_timing.rs +++ b/crates/larql-inference/examples/walk_ffn_decode_timing.rs @@ -12,7 +12,7 @@ //! //! Usage: `cargo run --release --example walk_ffn_decode_timing -- [VINDEX_DIR]` -use larql_inference::vindex::{insert_q4k_layer_tensors, WalkFfn, WalkFfnConfig}; +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; @@ -69,7 +69,7 @@ fn main() { ); let tok = load_tokenizer(&dir).expect("tok"); for layer in 0..weights.num_layers { - insert_q4k_layer_tensors(&mut weights, &index, layer).expect("dequant attn"); + insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant attn"); } let nl = weights.num_layers; diff --git a/crates/larql-inference/examples/walk_ffn_delta_walk.rs b/crates/larql-inference/examples/walk_ffn_delta_walk.rs index ca359ded2..888873ec5 100644 --- a/crates/larql-inference/examples/walk_ffn_delta_walk.rs +++ b/crates/larql-inference/examples/walk_ffn_delta_walk.rs @@ -27,7 +27,7 @@ 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, WalkFfn}; +use larql_inference::vindex::{insert_q4k_layer_tensors_resident, WalkFfn}; use ndarray::{Array2, Axis}; use std::cell::RefCell; @@ -111,7 +111,7 @@ fn main() { 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(&mut weights, &index, layer).expect("dequant attn"); + insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant attn"); } let nl = weights.num_layers; let hidden = weights.hidden_size; diff --git a/crates/larql-inference/examples/walk_ffn_drift.rs b/crates/larql-inference/examples/walk_ffn_drift.rs index 88fd53f85..1835993f3 100644 --- a/crates/larql-inference/examples/walk_ffn_drift.rs +++ b/crates/larql-inference/examples/walk_ffn_drift.rs @@ -12,7 +12,7 @@ //! Usage: `cargo run --release --example walk_ffn_drift -- [VINDEX]` use larql_inference::ffn::FfnBackend; -use larql_inference::vindex::{insert_q4k_layer_tensors, WalkFfn}; +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}; @@ -125,7 +125,7 @@ fn main() { 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(&mut weights, &index, layer).expect("dequant attn"); + insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant attn"); } let _ = WalkFfn::new_unlimited(&weights, &index); // (warms down-norm cache path) diff --git a/crates/larql-inference/examples/walk_ffn_graded_precision.rs b/crates/larql-inference/examples/walk_ffn_graded_precision.rs index 6e54240a6..f3e25a30c 100644 --- a/crates/larql-inference/examples/walk_ffn_graded_precision.rs +++ b/crates/larql-inference/examples/walk_ffn_graded_precision.rs @@ -13,7 +13,7 @@ //! Usage: `cargo run --release --example walk_ffn_graded_precision -- [VINDEX]` use larql_inference::ffn::FfnBackend; -use larql_inference::vindex::{insert_q4k_layer_tensors, WalkFfn}; +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}; @@ -217,7 +217,7 @@ fn main() { 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(&mut weights, &index, layer).expect("dequant attn"); + insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant attn"); } let prompts = [ diff --git a/crates/larql-inference/examples/walk_ffn_k_agreement.rs b/crates/larql-inference/examples/walk_ffn_k_agreement.rs index b88de9c15..aae634103 100644 --- a/crates/larql-inference/examples/walk_ffn_k_agreement.rs +++ b/crates/larql-inference/examples/walk_ffn_k_agreement.rs @@ -11,7 +11,7 @@ //! //! Usage: `cargo run --release --example walk_ffn_k_agreement -- [VINDEX_DIR]` -use larql_inference::vindex::{insert_q4k_layer_tensors, WalkFfn, WalkFfnConfig}; +use larql_inference::vindex::{insert_q4k_layer_tensors_resident, WalkFfn, WalkFfnConfig}; use larql_inference::{load_tokenizer, predict_with_ffn}; fn argmax( @@ -43,7 +43,7 @@ fn main() { 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(&mut weights, &index, layer).expect("dequant"); + insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant"); } let nl = weights.num_layers; let feats = index.num_features(nl / 2); diff --git a/crates/larql-inference/examples/walk_ffn_nll.rs b/crates/larql-inference/examples/walk_ffn_nll.rs index 168ac66a6..cbbf1f344 100644 --- a/crates/larql-inference/examples/walk_ffn_nll.rs +++ b/crates/larql-inference/examples/walk_ffn_nll.rs @@ -15,7 +15,7 @@ //! Usage: `cargo run --release --example walk_ffn_nll -- [VINDEX]` use larql_inference::ffn::FfnBackend; -use larql_inference::vindex::insert_q4k_layer_tensors; +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}; @@ -147,7 +147,7 @@ fn main() { 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(&mut weights, &index, layer).expect("dequant attn"); + insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant attn"); } // Entropic narrative prose — real lexical choice (NOT code/boilerplate, diff --git a/crates/larql-inference/examples/walk_ffn_temporal_reuse.rs b/crates/larql-inference/examples/walk_ffn_temporal_reuse.rs index c8081009e..09b287cf9 100644 --- a/crates/larql-inference/examples/walk_ffn_temporal_reuse.rs +++ b/crates/larql-inference/examples/walk_ffn_temporal_reuse.rs @@ -16,7 +16,7 @@ use larql_inference::load_tokenizer; use larql_inference::research::predict_with_ffn_trace; -use larql_inference::vindex::{insert_q4k_layer_tensors, WalkFfn}; +use larql_inference::vindex::{insert_q4k_layer_tensors_resident, WalkFfn}; use ndarray::Array1; const K: usize = 2048; // active-pool size for Jaccard @@ -116,7 +116,7 @@ fn main() { 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(&mut weights, &index, layer).expect("dequant attn"); + insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant attn"); } let nl = weights.num_layers; diff --git a/crates/larql-inference/examples/walk_ffn_v1_hash_routing.rs b/crates/larql-inference/examples/walk_ffn_v1_hash_routing.rs index 4965118a4..f0ca930eb 100644 --- a/crates/larql-inference/examples/walk_ffn_v1_hash_routing.rs +++ b/crates/larql-inference/examples/walk_ffn_v1_hash_routing.rs @@ -23,7 +23,7 @@ //! //! Usage: `cargo run --release --example walk_ffn_v1_hash_routing -- [VINDEX] [--quick]` -use larql_inference::vindex::{insert_q4k_layer_tensors, WalkFfn, WalkFfnConfig}; +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; @@ -158,7 +158,7 @@ fn main() { let nl = weights.num_layers; eprintln!("Dequantising attention for {nl} layers ..."); for layer in 0..nl { - insert_q4k_layer_tensors(&mut weights, &index, layer).expect("dequant layer"); + insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant layer"); } // Matrix `baseline_fact_prompts` (bench/aim-validation/matrix.json) — the diff --git a/crates/larql-inference/examples/walk_ffn_v2_fp4_nll.rs b/crates/larql-inference/examples/walk_ffn_v2_fp4_nll.rs index c7b774033..f34b50d7b 100644 --- a/crates/larql-inference/examples/walk_ffn_v2_fp4_nll.rs +++ b/crates/larql-inference/examples/walk_ffn_v2_fp4_nll.rs @@ -18,7 +18,7 @@ //! Usage: `cargo run --release --example walk_ffn_v2_fp4_nll -- [VINDEX]` use larql_inference::ffn::FfnBackend; -use larql_inference::vindex::insert_q4k_layer_tensors; +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}; @@ -179,7 +179,7 @@ fn main() { 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(&mut weights, &index, layer).expect("dequant attn"); + insert_q4k_layer_tensors_resident(&mut weights, &index, layer).expect("dequant attn"); } let passage = "The expedition had been planned for years, but nothing prepared \ diff --git a/crates/larql-inference/src/ffn/mod.rs b/crates/larql-inference/src/ffn/mod.rs index 1056fc0c6..444e4d180 100644 --- a/crates/larql-inference/src/ffn/mod.rs +++ b/crates/larql-inference/src/ffn/mod.rs @@ -42,7 +42,7 @@ pub use sparse_compute::{ sparse_ffn_forward, sparse_ffn_forward_with_full_overrides, sparse_ffn_forward_with_overrides, FeatureSlotOverride, }; -pub use weight::{dense_ffn_forward_backend, BackendFfn, NullFfn, WeightFfn}; +pub use weight::{dense_ffn_forward_backend, BackendFfn, NullFfn, ViewFfn, WeightFfn}; // ── Per-layer backend selection ── diff --git a/crates/larql-inference/src/ffn/sparse_compute.rs b/crates/larql-inference/src/ffn/sparse_compute.rs index 43849abd8..696a4c239 100644 --- a/crates/larql-inference/src/ffn/sparse_compute.rs +++ b/crates/larql-inference/src/ffn/sparse_compute.rs @@ -98,7 +98,7 @@ fn sparse_ffn_forward_impl( // Fall back to dense when most features are selected if k * 5 >= intermediate * 4 && overrides.is_empty() { - return dense_ffn_forward(weights, layer, x); + return dense_ffn_forward(larql_models::WeightsView::dense(weights), layer, x); } let is_gated = arch.ffn_type() == larql_models::FfnType::Gated; @@ -589,7 +589,11 @@ mod tests { // Request all features to trigger that path. let all: Vec = (0..weights.intermediate_size).collect(); let (sparse_out, _) = sparse_ffn_forward(&weights, 0, &x, &all); - let (dense_out, _) = crate::ffn::weight::dense_ffn_forward(&weights, 0, &x); + let (dense_out, _) = crate::ffn::weight::dense_ffn_forward( + larql_models::WeightsView::dense(&weights), + 0, + &x, + ); for (s, d) in sparse_out.iter().zip(dense_out.iter()) { assert!((s - d).abs() < 1e-4, "sparse/dense mismatch: {s} vs {d}"); } diff --git a/crates/larql-inference/src/forward/layer.rs b/crates/larql-inference/src/forward/layer.rs index 8520dbe25..aed838c6f 100644 --- a/crates/larql-inference/src/forward/layer.rs +++ b/crates/larql-inference/src/forward/layer.rs @@ -57,8 +57,16 @@ mod tests { let weights = make_test_weights(); let ffn = WeightFfn { weights: &weights }; let input = h(3, weights.hidden_size); - let (h_out, _act, _kv) = run_layer_with_ffn(&weights, &input, 0, &ffn, false, None, None) - .expect("run_layer_with_ffn failed"); + let (h_out, _act, _kv) = run_layer_with_ffn( + larql_models::WeightsView::dense(&weights), + &input, + 0, + &ffn, + false, + None, + None, + ) + .expect("run_layer_with_ffn failed"); assert_eq!(h_out.shape(), &[3, weights.hidden_size]); } @@ -69,7 +77,16 @@ mod tests { let input = h(2, weights.hidden_size); for layer in 0..weights.num_layers { assert!( - run_layer_with_ffn(&weights, &input, layer, &ffn, false, None, None).is_some(), + run_layer_with_ffn( + larql_models::WeightsView::dense(&weights), + &input, + layer, + &ffn, + false, + None, + None + ) + .is_some(), "layer {layer} failed" ); } @@ -79,8 +96,10 @@ mod tests { fn run_attention_public_matches_inner() { let weights = make_test_weights(); let input = h(3, weights.hidden_size); - let pub_out = run_attention_public(&weights, &input, 0).unwrap(); - let inner_out = run_attention(&weights, &input, 0).unwrap(); + let pub_out = + run_attention_public(larql_models::WeightsView::dense(&weights), &input, 0).unwrap(); + let inner_out = + run_attention(larql_models::WeightsView::dense(&weights), &input, 0).unwrap(); assert_eq!(pub_out.shape(), inner_out.shape()); for (a, b) in pub_out.iter().zip(inner_out.iter()) { assert!( @@ -94,8 +113,14 @@ mod tests { fn run_attention_inner_with_capture_attention_returns_weights() { let weights = make_test_weights(); let input = h(3, weights.hidden_size); - let (out, weights_opt) = - run_attention_inner(&weights, &input, 0, /*capture=*/ true, None).unwrap(); + let (out, weights_opt) = run_attention_inner( + larql_models::WeightsView::dense(&weights), + &input, + 0, + /*capture=*/ true, + None, + ) + .unwrap(); assert_eq!(out.shape(), &[3, weights.hidden_size]); let aw = weights_opt.expect("attention weights should be captured"); // One distribution per Q-head, each with seq_len=3 entries (last position). @@ -110,7 +135,8 @@ mod tests { let weights = make_test_weights(); let input = h(2, weights.hidden_size); let (h_post_attn, (k, v)) = - run_attention_with_kv_cache(&weights, &input, 0).expect("attn-with-kv must succeed"); + run_attention_with_kv_cache(larql_models::WeightsView::dense(&weights), &input, 0) + .expect("attn-with-kv must succeed"); assert_eq!(h_post_attn.shape(), &[2, weights.hidden_size]); // K/V have shape (seq, num_kv_heads * head_dim). let kv_dim = weights.num_kv_heads * weights.head_dim; @@ -135,7 +161,14 @@ mod tests { let ffn = WeightFfn { weights: &weights }; let input = h(2, weights.hidden_size); let (h_out, _act, _attn, _kv) = run_layer_with_capture( - &weights, &input, 0, &ffn, false, /*capture_attention=*/ true, None, None, + larql_models::WeightsView::dense(&weights), + &input, + 0, + &ffn, + false, + /*capture_attention=*/ true, + None, + None, ) .expect("run_layer_with_capture must succeed"); assert_eq!(h_out.shape(), &[2, weights.hidden_size]); @@ -149,10 +182,19 @@ mod tests { let ffn = WeightFfn { weights: &weights }; let input = h(2, weights.hidden_size); // Capture KV from layer 0 first, then re-feed at layer 1 as shared. - let (_, shared) = run_attention_with_kv_cache(&weights, &input, 0).unwrap(); - let (h_out, _, kv_out) = - run_layer_with_ffn(&weights, &input, 1, &ffn, false, None, Some(&shared)) - .expect("layer with shared KV must succeed"); + let (_, shared) = + run_attention_with_kv_cache(larql_models::WeightsView::dense(&weights), &input, 0) + .unwrap(); + let (h_out, _, kv_out) = run_layer_with_ffn( + larql_models::WeightsView::dense(&weights), + &input, + 1, + &ffn, + false, + None, + Some(&shared), + ) + .expect("layer with shared KV must succeed"); assert_eq!(h_out.shape(), &[2, weights.hidden_size]); assert!(kv_out.is_none(), "shared-KV path must not return new KV"); } @@ -178,8 +220,16 @@ mod tests { let weights = larql_models::test_fixtures::make_gemma3_test_weights(); let ffn = WeightFfn { weights: &weights }; let input = h(2, weights.hidden_size); - let (h_out, _, kv) = - run_layer_with_ffn(&weights, &input, 0, &ffn, false, None, None).unwrap(); + let (h_out, _, kv) = run_layer_with_ffn( + larql_models::WeightsView::dense(&weights), + &input, + 0, + &ffn, + false, + None, + None, + ) + .unwrap(); assert_eq!(h_out.shape(), &[2, weights.hidden_size]); assert!(h_out.iter().all(|v| v.is_finite())); assert!(kv.is_some()); @@ -192,8 +242,16 @@ mod tests { let weights = larql_models::test_fixtures::make_starcoder2_test_weights(); let ffn = WeightFfn { weights: &weights }; let input = h(2, weights.hidden_size); - let (h_out, _, _) = - run_layer_with_ffn(&weights, &input, 0, &ffn, false, None, None).unwrap(); + let (h_out, _, _) = run_layer_with_ffn( + larql_models::WeightsView::dense(&weights), + &input, + 0, + &ffn, + false, + None, + None, + ) + .unwrap(); assert_eq!(h_out.shape(), &[2, weights.hidden_size]); assert!(h_out.iter().all(|v| v.is_finite())); } @@ -205,10 +263,12 @@ mod tests { let weights = make_test_weights(); let ffn = WeightFfn { weights: &weights }; let input = h(2, weights.hidden_size); - let (_, shared) = run_attention_with_kv_cache(&weights, &input, 0).unwrap(); + let (_, shared) = + run_attention_with_kv_cache(larql_models::WeightsView::dense(&weights), &input, 0) + .unwrap(); let mut hook = crate::forward::hooks::NoopHook; let (h_out, _, _, kv_out) = run_layer_with_capture_hooked( - &weights, + larql_models::WeightsView::dense(&weights), &input, 1, &ffn, diff --git a/crates/larql-inference/src/forward/layer_interventions.rs b/crates/larql-inference/src/forward/layer_interventions.rs index e909729d3..67c09babe 100644 --- a/crates/larql-inference/src/forward/layer_interventions.rs +++ b/crates/larql-inference/src/forward/layer_interventions.rs @@ -48,7 +48,11 @@ pub fn run_layer_with_zeroed_pre_o_heads( shared_kv: Option<&SharedKV>, ) -> Option<(Array2, Option)> { let (h_post_attn, kv_out) = crate::attention::run_attention_block_zero_pre_o_heads( - weights, h, layer, heads, shared_kv, + larql_models::WeightsView::dense(weights), + h, + layer, + heads, + shared_kv, )?; if let Some(dir) = crate::forward::dump_config::DumpConfig::get().layer_dir() { let slice = h_post_attn.as_slice().unwrap_or(&[]); @@ -79,7 +83,7 @@ pub fn run_layer_with_replaced_pre_o_head( shared_kv: Option<&SharedKV>, ) -> Option<(Array2, Option)> { let (h_post_attn, kv_out) = crate::attention::run_attention_block_replace_pre_o_head( - weights, + larql_models::WeightsView::dense(weights), h, layer, head, @@ -113,8 +117,12 @@ pub fn run_layer_with_mapped_pre_o_head( where F: FnMut(&Array2) -> Option>, { - let (_, pre_o) = - crate::attention::run_attention_block_shared_with_pre_o(weights, h, layer, shared_kv)?; + let (_, pre_o) = crate::attention::run_attention_block_shared_with_pre_o( + larql_models::WeightsView::dense(weights), + h, + layer, + shared_kv, + )?; let head_dim = weights.arch.head_dim_for_layer(layer); let start = head.checked_mul(head_dim)?; let end = start.checked_add(head_dim)?; @@ -159,8 +167,12 @@ pub fn run_layer_with_mapped_head_residual_delta( where F: FnMut(&Array2) -> Option>, { - let (_, pre_o) = - crate::attention::run_attention_block_shared_with_pre_o(weights, h, layer, shared_kv)?; + let (_, pre_o) = crate::attention::run_attention_block_shared_with_pre_o( + larql_models::WeightsView::dense(weights), + h, + layer, + shared_kv, + )?; let head_dim = weights.arch.head_dim_for_layer(layer); let start = head.checked_mul(head_dim)?; let end = start.checked_add(head_dim)?; @@ -200,8 +212,12 @@ pub fn run_layer_with_original_head_residual_delta( ple_input: Option<&Array2>, shared_kv: Option<&SharedKV>, ) -> Option<(Array2, Option)> { - let (_, pre_o) = - crate::attention::run_attention_block_shared_with_pre_o(weights, h, layer, shared_kv)?; + let (_, pre_o) = crate::attention::run_attention_block_shared_with_pre_o( + larql_models::WeightsView::dense(weights), + h, + layer, + shared_kv, + )?; let head_dim = weights.arch.head_dim_for_layer(layer); let start = head.checked_mul(head_dim)?; let end = start.checked_add(head_dim)?; @@ -239,7 +255,11 @@ pub fn run_layer_with_subtracted_pre_o_heads( shared_kv: Option<&SharedKV>, ) -> Option<(Array2, Option)> { let (h_post_attn, kv_out) = crate::attention::run_attention_block_subtract_pre_o_heads( - weights, h, layer, heads, shared_kv, + larql_models::WeightsView::dense(weights), + h, + layer, + heads, + shared_kv, )?; Some(( finish_layer_tail(weights, &h_post_attn, layer, ffn, ple_input), @@ -265,7 +285,7 @@ pub fn run_layer_with_replaced_head_residual_delta( shared_kv: Option<&SharedKV>, ) -> Option<(Array2, Option)> { let (h_post_attn, kv_out) = crate::attention::run_attention_block_replace_head_residual_delta( - weights, + larql_models::WeightsView::dense(weights), h, layer, head, @@ -322,8 +342,16 @@ mod tests { let weights = make_test_weights(); let ffn = WeightFfn { weights: &weights }; let input = h(3, weights.hidden_size); - let (baseline, _, _) = run_layer_with_ffn(&weights, &input, 0, &ffn, false, None, None) - .expect("baseline layer failed"); + let (baseline, _, _) = run_layer_with_ffn( + larql_models::WeightsView::dense(&weights), + &input, + 0, + &ffn, + false, + None, + None, + ) + .expect("baseline layer failed"); let (mapped, _) = run_layer_with_mapped_pre_o_head(&weights, &input, 0, &ffn, 0, None, None, |head| { Some(head.clone()) @@ -344,8 +372,16 @@ mod tests { let weights = make_test_weights(); let ffn = WeightFfn { weights: &weights }; let input = h(2, weights.hidden_size); - let (baseline, _, _) = run_layer_with_ffn(&weights, &input, 0, &ffn, false, None, None) - .expect("baseline layer failed"); + let (baseline, _, _) = run_layer_with_ffn( + larql_models::WeightsView::dense(&weights), + &input, + 0, + &ffn, + false, + None, + None, + ) + .expect("baseline layer failed"); let (intervened, _) = run_layer_with_original_head_residual_delta(&weights, &input, 0, &ffn, 0, None, None) .expect("original-delta layer failed"); @@ -365,8 +401,16 @@ mod tests { let weights = make_test_weights(); let ffn = WeightFfn { weights: &weights }; let input = h(2, weights.hidden_size); - let (baseline, _, _) = run_layer_with_ffn(&weights, &input, 0, &ffn, false, None, None) - .expect("baseline layer failed"); + let (baseline, _, _) = run_layer_with_ffn( + larql_models::WeightsView::dense(&weights), + &input, + 0, + &ffn, + false, + None, + None, + ) + .expect("baseline layer failed"); let (intervened, _) = run_layer_with_zeroed_pre_o_heads(&weights, &input, 0, &ffn, &[], None, None) .expect("zeroed layer failed"); @@ -383,8 +427,16 @@ mod tests { let weights = make_test_weights(); let ffn = WeightFfn { weights: &weights }; let input = h(2, weights.hidden_size); - let (baseline, _, _) = run_layer_with_ffn(&weights, &input, 0, &ffn, false, None, None) - .expect("baseline layer failed"); + let (baseline, _, _) = run_layer_with_ffn( + larql_models::WeightsView::dense(&weights), + &input, + 0, + &ffn, + false, + None, + None, + ) + .expect("baseline layer failed"); let (intervened, _) = run_layer_with_zeroed_pre_o_heads(&weights, &input, 0, &ffn, &[0], None, None) .expect("zeroed layer failed"); @@ -490,8 +542,16 @@ mod tests { let weights = make_test_weights(); let ffn = WeightFfn { weights: &weights }; let input = h(2, weights.hidden_size); - let (baseline, _, _) = run_layer_with_ffn(&weights, &input, 0, &ffn, false, None, None) - .expect("baseline layer failed"); + let (baseline, _, _) = run_layer_with_ffn( + larql_models::WeightsView::dense(&weights), + &input, + 0, + &ffn, + false, + None, + None, + ) + .expect("baseline layer failed"); let (mapped, _) = run_layer_with_mapped_head_residual_delta( &weights, &input, diff --git a/crates/larql-inference/src/forward/predict/dense.rs b/crates/larql-inference/src/forward/predict/dense.rs index f98a4b1f2..4f9ddb1fd 100644 --- a/crates/larql-inference/src/forward/predict/dense.rs +++ b/crates/larql-inference/src/forward/predict/dense.rs @@ -413,7 +413,7 @@ pub fn predict_with_temperature( .kv_shared_source_layer(layer) .and_then(|src| kv_cache.get(&src)); match run_layer_with_ffn( - weights, + larql_models::WeightsView::dense(weights), &h, layer, &ffn, @@ -481,7 +481,15 @@ pub fn predict_from_hidden_with_ffn( }; for layer in start_layer..num_layers { - h = match run_layer_with_ffn(weights, &h, layer, ffn, false, ple_inputs.get(layer), None) { + h = match run_layer_with_ffn( + larql_models::WeightsView::dense(weights), + &h, + layer, + ffn, + false, + ple_inputs.get(layer), + None, + ) { Some((h_new, _, _)) => h_new, None => continue, }; @@ -507,7 +515,15 @@ pub fn predict_with_ffn_trace( let last_pos = h.shape()[0] - 1; residuals.push(h.row(last_pos).to_vec()); - h = match run_layer_with_ffn(weights, &h, layer, ffn, false, ple_inputs.get(layer), None) { + h = match run_layer_with_ffn( + larql_models::WeightsView::dense(weights), + &h, + layer, + ffn, + false, + ple_inputs.get(layer), + None, + ) { Some((h_new, _, _)) => h_new, None => continue, }; diff --git a/crates/larql-inference/src/forward/predict/ffn.rs b/crates/larql-inference/src/forward/predict/ffn.rs index 4871ecf7f..852874b4d 100644 --- a/crates/larql-inference/src/forward/predict/ffn.rs +++ b/crates/larql-inference/src/forward/predict/ffn.rs @@ -30,7 +30,7 @@ pub fn predict_with_ffn( .and_then(|src| kv_cache.get(&src)); match run_layer_with_ffn( - weights, + larql_models::WeightsView::dense(weights), &h, layer, ffn, @@ -84,7 +84,7 @@ pub fn predict_with_ffn_early_exit( .and_then(|src| kv_cache.get(&src)); match run_layer_with_ffn( - weights, + larql_models::WeightsView::dense(weights), &h, layer, ffn, @@ -132,7 +132,7 @@ pub fn predict_with_ffn_attention( for layer in 0..num_layers { match run_layer_with_capture( - weights, + larql_models::WeightsView::dense(weights), &h, layer, ffn, @@ -174,7 +174,15 @@ pub fn predict_with_router( for layer in 0..num_layers { let ffn = router.get(layer); - h = match run_layer_with_ffn(weights, &h, layer, ffn, false, ple_inputs.get(layer), None) { + h = match run_layer_with_ffn( + larql_models::WeightsView::dense(weights), + &h, + layer, + ffn, + false, + ple_inputs.get(layer), + None, + ) { Some((h_new, _, _)) => h_new, None => continue, }; @@ -199,7 +207,7 @@ pub fn predict_with_strategy( match mode { LayerMode::Compute(ffn) => { h = match run_layer_with_ffn( - weights, + larql_models::WeightsView::dense(weights), &h, layer, *ffn, @@ -215,7 +223,9 @@ pub fn predict_with_strategy( h *= *gain; } LayerMode::AttentionOnly => { - if let Some(h_post_attn) = run_attention(weights, &h, layer) { + if let Some(h_post_attn) = + run_attention(larql_models::WeightsView::dense(weights), &h, layer) + { h = h_post_attn; } } diff --git a/crates/larql-inference/src/forward/target_delta.rs b/crates/larql-inference/src/forward/target_delta.rs index a895764ae..a70497ed7 100644 --- a/crates/larql-inference/src/forward/target_delta.rs +++ b/crates/larql-inference/src/forward/target_delta.rs @@ -354,7 +354,11 @@ pub fn optimise_target_delta( } // Baseline forward (no perturbation) for KL regulariser. - let baseline = crate::forward::predict::forward_raw_logits(weights, tokens, None); + let baseline = crate::forward::predict::forward_raw_logits( + larql_models::WeightsView::dense(weights), + tokens, + None, + ); let base_probs = softmax_1d(&baseline.logits); let baseline_loss = { let (l, _) = cross_entropy_and_grad(baseline.logits.view(), target_id); @@ -373,7 +377,7 @@ pub fn optimise_target_delta( let mut final_loss = f32::NAN; for step in 1..=opts.steps { let out = crate::forward::predict::forward_raw_logits( - weights, + larql_models::WeightsView::dense(weights), tokens, Some((install_layer, delta.view())), ); diff --git a/crates/larql-inference/src/forward/trace.rs b/crates/larql-inference/src/forward/trace.rs index 516701aad..a03f0ee1f 100644 --- a/crates/larql-inference/src/forward/trace.rs +++ b/crates/larql-inference/src/forward/trace.rs @@ -39,7 +39,8 @@ pub fn capture_spec_residuals(weights: &ModelWeights, token_ids: &[u32]) -> Spec let mut post_layer_last = Vec::with_capacity(weights.num_layers); for layer in 0..weights.num_layers { - let h_post_attn = match run_attention(weights, &h, layer) { + let h_post_attn = match run_attention(larql_models::WeightsView::dense(weights), &h, layer) + { Some(pa) => pa, None => h.clone(), }; @@ -89,13 +90,13 @@ pub fn trace_forward_attn_only_with_head_zero( for layer in 0..=max_layer { let heads_to_zero = head_zero_map.get(&layer).unwrap_or(&empty); let h_post_attn = if heads_to_zero.is_empty() { - match run_attention(weights, &h, layer) { + match run_attention(larql_models::WeightsView::dense(weights), &h, layer) { Some(p) => p, None => h.clone(), } } else { match crate::attention::run_attention_block_zero_pre_o_heads( - weights, + larql_models::WeightsView::dense(weights), &h, layer, heads_to_zero, @@ -134,15 +135,19 @@ pub fn trace_forward_attn_only_capture_pre_o( for layer in 0..=max_layer { if capture_layers.contains(&layer) { // Capture both the post-attention residual (advances h) and pre_o. - if let Some((h_post_attn, pre_o)) = - crate::attention::run_attention_block_with_pre_o(weights, &h, layer) - { + if let Some((h_post_attn, pre_o)) = crate::attention::run_attention_block_with_pre_o( + larql_models::WeightsView::dense(weights), + &h, + layer, + ) { captures.insert(layer, pre_o); h = h_post_attn; } else { // Skip layer with degenerate weights — leave h unchanged. } - } else if let Some(h_post_attn) = run_attention(weights, &h, layer) { + } else if let Some(h_post_attn) = + run_attention(larql_models::WeightsView::dense(weights), &h, layer) + { h = h_post_attn; } } @@ -161,7 +166,15 @@ pub fn forward_to_layer( let ple_inputs = precompute_per_layer_inputs(weights, &h, token_ids); for layer in 0..=stop_layer { - h = match run_layer_with_ffn(weights, &h, layer, &ffn, false, ple_inputs.get(layer), None) { + h = match run_layer_with_ffn( + larql_models::WeightsView::dense(weights), + &h, + layer, + &ffn, + false, + ple_inputs.get(layer), + None, + ) { Some((h_new, _, _)) => h_new, None => continue, }; @@ -239,7 +252,7 @@ pub fn capture_ffn_activation_matrix( // truncation that happens there. let need_activation = l == layer; let (h_new, activation, _, _) = crate::forward::layer::run_layer_with_capture( - weights, + larql_models::WeightsView::dense(weights), &h, l, &ffn, @@ -439,7 +452,7 @@ pub fn trace_forward_full_hooked( let need_attention = capture_attention && is_capture_layer; let (h_new, activation, attn_weights, _) = match run_layer_with_capture_hooked( - weights, + larql_models::WeightsView::dense(weights), &h, layer, ffn, diff --git a/crates/larql-inference/src/kv_dispatch/helpers.rs b/crates/larql-inference/src/kv_dispatch/helpers.rs index 66c549c7f..ca597e518 100644 --- a/crates/larql-inference/src/kv_dispatch/helpers.rs +++ b/crates/larql-inference/src/kv_dispatch/helpers.rs @@ -25,7 +25,6 @@ use super::{EngineBackend, KvHandle}; use crate::async_compute_backend::AsyncComputeBackend; use crate::ffn::FfnBackend; use crate::forward::{embed_tokens_pub, run_ffn}; -use crate::model::ModelWeights; /// Per-layer FFN dispatch for the KV-cached engine path, MoE-aware. /// @@ -37,7 +36,7 @@ use crate::model::ModelWeights; /// cache. When the hook declines (`None`) — or the model is dense — fall /// back to the standard dense FFN, preserving prior behaviour exactly. fn ffn_or_moe_layer( - weights: &ModelWeights, + weights: larql_models::WeightsView, h_post_attn: &Array2, layer: usize, ffn: &dyn FfnBackend, @@ -47,7 +46,7 @@ fn ffn_or_moe_layer( return h_out; } } - run_ffn(weights, h_post_attn, layer, ffn, false).0 + run_ffn(&weights, h_post_attn, layer, ffn, false).0 } /// Prefill the K/V cache through every layer using `backend`'s @@ -61,7 +60,7 @@ fn ffn_or_moe_layer( /// [`KvDispatch::clip_kv`] per-layer after this returns). pub fn kv_prefill_via_dispatch( backend: &dyn EngineBackend, - weights: &ModelWeights, + weights: larql_models::WeightsView, ffn: &dyn FfnBackend, prompt_ids: &[u32], window: Option, @@ -70,7 +69,7 @@ pub fn kv_prefill_via_dispatch( if prompt_ids.is_empty() { return None; } - let h = embed_tokens_pub(weights, prompt_ids); + let h = embed_tokens_pub(&weights, prompt_ids); kv_prefill_from_hidden_via_dispatch(backend, weights, ffn, &h, window, index) } @@ -87,7 +86,7 @@ pub fn kv_prefill_via_dispatch( /// of this module. pub fn kv_prefill_from_hidden_via_dispatch( backend: &dyn EngineBackend, - weights: &ModelWeights, + weights: larql_models::WeightsView, ffn: &dyn FfnBackend, initial_hidden: &Array2, window: Option, @@ -134,7 +133,7 @@ pub fn kv_prefill_from_hidden_via_dispatch( #[allow(clippy::too_many_arguments)] pub fn kv_decode_step_via_dispatch( backend: &dyn EngineBackend, - weights: &ModelWeights, + weights: larql_models::WeightsView, ffn: &dyn FfnBackend, handles: &mut [KvHandle], token_id: u32, @@ -148,7 +147,7 @@ pub fn kv_decode_step_via_dispatch( num_layers, "kv_decode_step_via_dispatch: handles.len() must equal weights.num_layers" ); - let h_new = embed_tokens_pub(weights, &[token_id]); + let h_new = embed_tokens_pub(&weights, &[token_id]); let mut h_step = h_new; for (layer, handle) in handles.iter_mut().enumerate().take(num_layers) { @@ -191,7 +190,7 @@ pub fn kv_decode_step_via_dispatch( /// the end so any deferred work clears before returning. pub fn kv_prefill_via_dispatch_async( backend: &dyn AsyncComputeBackend, - weights: &ModelWeights, + weights: larql_models::WeightsView, ffn: &dyn FfnBackend, prompt_ids: &[u32], window: Option, @@ -200,7 +199,7 @@ pub fn kv_prefill_via_dispatch_async( if prompt_ids.is_empty() { return None; } - let h = embed_tokens_pub(weights, prompt_ids); + let h = embed_tokens_pub(&weights, prompt_ids); kv_prefill_from_hidden_via_dispatch_async(backend, weights, ffn, &h, window, index) } @@ -214,7 +213,7 @@ pub fn kv_prefill_via_dispatch_async( /// CPU paths, MM vs text must agree when text is the input. pub fn kv_prefill_from_hidden_via_dispatch_async( backend: &dyn AsyncComputeBackend, - weights: &ModelWeights, + weights: larql_models::WeightsView, ffn: &dyn FfnBackend, initial_hidden: &Array2, window: Option, @@ -258,7 +257,7 @@ pub fn kv_prefill_from_hidden_via_dispatch_async( #[allow(clippy::too_many_arguments)] pub fn kv_decode_step_via_dispatch_async( backend: &dyn AsyncComputeBackend, - weights: &ModelWeights, + weights: larql_models::WeightsView, ffn: &dyn FfnBackend, handles: &mut [KvHandle], token_id: u32, @@ -272,7 +271,7 @@ pub fn kv_decode_step_via_dispatch_async( num_layers, "kv_decode_step_via_dispatch_async: handles.len() must equal weights.num_layers" ); - let h_new = embed_tokens_pub(weights, &[token_id]); + let h_new = embed_tokens_pub(&weights, &[token_id]); let mut h_step = h_new; for (layer, handle) in handles.iter_mut().enumerate().take(num_layers) { @@ -330,15 +329,22 @@ mod tests { let ffn = WeightFfn { weights: &weights }; let prompt = vec![0u32, 1]; - let (_, mut handles) = - kv_prefill_via_dispatch(&backend, &weights, &ffn, &prompt, None, None).unwrap(); + let (_, mut handles) = kv_prefill_via_dispatch( + &backend, + larql_models::WeightsView::dense(&weights), + &ffn, + &prompt, + None, + None, + ) + .unwrap(); for step in 0..3 { let token = (2 + step) as u32; let abs_position = prompt.len() + step; let h_trait = kv_decode_step_via_dispatch( &backend, - &weights, + larql_models::WeightsView::dense(&weights), &ffn, &mut handles, token, @@ -359,7 +365,14 @@ mod tests { let weights = make_test_weights(); let backend = CpuBackend; let ffn = WeightFfn { weights: &weights }; - let result = kv_prefill_via_dispatch(&backend, &weights, &ffn, &[], None, None); + let result = kv_prefill_via_dispatch( + &backend, + larql_models::WeightsView::dense(&weights), + &ffn, + &[], + None, + None, + ); assert!(result.is_none()); } @@ -372,10 +385,24 @@ mod tests { let ffn = WeightFfn { weights: &weights }; let prompt = vec![0u32, 1, 2, 3]; - let (h_sync, handles_sync) = - kv_prefill_via_dispatch(&backend, &weights, &ffn, &prompt, None, None).unwrap(); - let (h_async, handles_async) = - kv_prefill_via_dispatch_async(&backend, &weights, &ffn, &prompt, None, None).unwrap(); + let (h_sync, handles_sync) = kv_prefill_via_dispatch( + &backend, + larql_models::WeightsView::dense(&weights), + &ffn, + &prompt, + None, + None, + ) + .unwrap(); + let (h_async, handles_async) = kv_prefill_via_dispatch_async( + &backend, + larql_models::WeightsView::dense(&weights), + &ffn, + &prompt, + None, + None, + ) + .unwrap(); assert_eq!(h_sync, h_async, "async prefill hidden must match sync"); assert_eq!(handles_sync.len(), handles_async.len()); @@ -395,10 +422,24 @@ mod tests { let prompt = vec![0u32, 1, 2, 3, 4]; let window = Some(2); - let (h_sync, _) = - kv_prefill_via_dispatch(&backend, &weights, &ffn, &prompt, window, None).unwrap(); - let (h_async, _) = - kv_prefill_via_dispatch_async(&backend, &weights, &ffn, &prompt, window, None).unwrap(); + let (h_sync, _) = kv_prefill_via_dispatch( + &backend, + larql_models::WeightsView::dense(&weights), + &ffn, + &prompt, + window, + None, + ) + .unwrap(); + let (h_async, _) = kv_prefill_via_dispatch_async( + &backend, + larql_models::WeightsView::dense(&weights), + &ffn, + &prompt, + window, + None, + ) + .unwrap(); assert_eq!(h_sync, h_async, "windowed async prefill must match sync"); } @@ -410,17 +451,31 @@ mod tests { let ffn = WeightFfn { weights: &weights }; let prompt = vec![0u32, 1, 2]; - let (_, mut handles_sync) = - kv_prefill_via_dispatch(&backend, &weights, &ffn, &prompt, None, None).unwrap(); - let (_, mut handles_async) = - kv_prefill_via_dispatch_async(&backend, &weights, &ffn, &prompt, None, None).unwrap(); + let (_, mut handles_sync) = kv_prefill_via_dispatch( + &backend, + larql_models::WeightsView::dense(&weights), + &ffn, + &prompt, + None, + None, + ) + .unwrap(); + let (_, mut handles_async) = kv_prefill_via_dispatch_async( + &backend, + larql_models::WeightsView::dense(&weights), + &ffn, + &prompt, + None, + None, + ) + .unwrap(); let next_token = 3u32; let abs_position = prompt.len(); let h_sync = kv_decode_step_via_dispatch( &backend, - &weights, + larql_models::WeightsView::dense(&weights), &ffn, &mut handles_sync, next_token, @@ -431,7 +486,7 @@ mod tests { .unwrap(); let h_async = kv_decode_step_via_dispatch_async( &backend, - &weights, + larql_models::WeightsView::dense(&weights), &ffn, &mut handles_async, next_token, @@ -451,17 +506,31 @@ mod tests { let ffn = WeightFfn { weights: &weights }; let prompt = vec![0u32, 1]; - let (_, mut handles_sync) = - kv_prefill_via_dispatch(&backend, &weights, &ffn, &prompt, None, None).unwrap(); - let (_, mut handles_async) = - kv_prefill_via_dispatch_async(&backend, &weights, &ffn, &prompt, None, None).unwrap(); + let (_, mut handles_sync) = kv_prefill_via_dispatch( + &backend, + larql_models::WeightsView::dense(&weights), + &ffn, + &prompt, + None, + None, + ) + .unwrap(); + let (_, mut handles_async) = kv_prefill_via_dispatch_async( + &backend, + larql_models::WeightsView::dense(&weights), + &ffn, + &prompt, + None, + None, + ) + .unwrap(); for step in 0..3 { let token = (2 + step) as u32; let abs_position = prompt.len() + step; let h_sync = kv_decode_step_via_dispatch( &backend, - &weights, + larql_models::WeightsView::dense(&weights), &ffn, &mut handles_sync, token, @@ -472,7 +541,7 @@ mod tests { .unwrap(); let h_async = kv_decode_step_via_dispatch_async( &backend, - &weights, + larql_models::WeightsView::dense(&weights), &ffn, &mut handles_async, token, @@ -490,7 +559,14 @@ mod tests { let weights = make_test_weights(); let backend = CpuBackend; let ffn = WeightFfn { weights: &weights }; - let result = kv_prefill_via_dispatch_async(&backend, &weights, &ffn, &[], None, None); + let result = kv_prefill_via_dispatch_async( + &backend, + larql_models::WeightsView::dense(&weights), + &ffn, + &[], + None, + None, + ); assert!(result.is_none()); } @@ -512,13 +588,20 @@ mod tests { let ffn = WeightFfn { weights: &weights }; let tokens = vec![0u32, 1, 2, 3]; - let (h_text, handles_text) = - kv_prefill_via_dispatch(&backend, &weights, &ffn, &tokens, None, None).unwrap(); + let (h_text, handles_text) = kv_prefill_via_dispatch( + &backend, + larql_models::WeightsView::dense(&weights), + &ffn, + &tokens, + None, + None, + ) + .unwrap(); let initial_hidden = embed_tokens_pub(&weights, &tokens); let (h_hidden, handles_hidden) = kv_prefill_from_hidden_via_dispatch( &backend, - &weights, + larql_models::WeightsView::dense(&weights), &ffn, &initial_hidden, None, @@ -544,13 +627,20 @@ mod tests { let ffn = WeightFfn { weights: &weights }; let tokens = vec![0u32, 1, 2, 3]; - let (h_text, handles_text) = - kv_prefill_via_dispatch_async(&backend, &weights, &ffn, &tokens, None, None).unwrap(); + let (h_text, handles_text) = kv_prefill_via_dispatch_async( + &backend, + larql_models::WeightsView::dense(&weights), + &ffn, + &tokens, + None, + None, + ) + .unwrap(); let initial_hidden = embed_tokens_pub(&weights, &tokens); let (h_hidden, handles_hidden) = kv_prefill_from_hidden_via_dispatch_async( &backend, - &weights, + larql_models::WeightsView::dense(&weights), &ffn, &initial_hidden, None, @@ -573,7 +663,7 @@ mod tests { let empty_hidden = Array2::::zeros((0, weights.hidden_size)); let result = kv_prefill_from_hidden_via_dispatch( &backend, - &weights, + larql_models::WeightsView::dense(&weights), &ffn, &empty_hidden, None, @@ -583,7 +673,7 @@ mod tests { let result_async = kv_prefill_from_hidden_via_dispatch_async( &backend, - &weights, + larql_models::WeightsView::dense(&weights), &ffn, &empty_hidden, None, diff --git a/crates/larql-inference/src/kv_engine.rs b/crates/larql-inference/src/kv_engine.rs index e28ce7e6c..31d23c409 100644 --- a/crates/larql-inference/src/kv_engine.rs +++ b/crates/larql-inference/src/kv_engine.rs @@ -347,12 +347,13 @@ pub trait KvEngine: Send { /// 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 `&mut` so the engine can lazily insert dequantised f32 - /// attention tensors into `weights.tensors` on the first call (one-time - /// cost; subsequent decode steps reuse the cached tensors). + /// `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: &mut ModelWeights, + weights: &ModelWeights, ffn: &dyn FfnBackend, index: &larql_vindex::VectorIndex, token_ids: &[u32], @@ -368,7 +369,7 @@ pub trait KvEngine: Send { /// when available, f32 fallback otherwise. fn decode_step_quant( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, ffn: &dyn FfnBackend, index: &larql_vindex::VectorIndex, token_id: u32, @@ -378,13 +379,12 @@ pub trait KvEngine: Send { self.decode_step(weights, ffn, token_id) // default: f32 fallback } - /// Resident-weights quant prefill. Unlike [`prefill_quant`] (which takes - /// `&mut weights` to lazily dequantise attn into `weights.tensors`), this - /// assumes the **caller has already made the client weights f32-resident** - /// — so it takes `&weights` and 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` - /// immutably — no `&mut`/`&` borrow conflict (task #16). Default: f32 + /// 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, @@ -451,7 +451,7 @@ pub trait KvEngine: Send { /// parameter properly. fn prefill_quant_via_executor( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, executor: &dyn crate::layer_executor::LayerExecutor, ffn: &dyn FfnBackend, index: &larql_vindex::VectorIndex, @@ -464,7 +464,7 @@ pub trait KvEngine: Send { /// [`prefill_quant_via_executor`] for the migration contract. fn decode_step_quant_via_executor( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, executor: &dyn crate::layer_executor::LayerExecutor, ffn: &dyn FfnBackend, index: &larql_vindex::VectorIndex, @@ -525,38 +525,39 @@ pub trait RetrievalEngine: Send { token_id: u32, ) -> Result, EngineError>; - /// Prefill against a Q4K-quantised vindex. Default impl dequantises - /// the attention tensors into `weights.tensors` and delegates to - /// [`prefill`](Self::prefill); engines that also need the FFN - /// tensors dequantised (e.g. Apollo, which runs its forward through - /// `forward_raw_logits` rather than an `FfnBackend` router) override - /// to insert those too. - /// - /// `weights` is `&mut` because the dequant step lazily populates - /// `weights.tensors`. Production callers reuse the populated - /// tensors across subsequent decode steps; the cost is one-time per - /// engine session. + /// 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: &mut ModelWeights, + weights: &ModelWeights, index: &larql_vindex::VectorIndex, token_ids: &[u32], ) -> Result, EngineError> { - crate::vindex::ensure_attn_tensors_dequantised(weights, index); - self.prefill(weights, token_ids) + 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. Default impl - /// dequantises the attention tensors and delegates to - /// [`decode_step`](Self::decode_step). + /// 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: &mut ModelWeights, + weights: &ModelWeights, index: &larql_vindex::VectorIndex, token_id: u32, ) -> Result, EngineError> { - crate::vindex::ensure_attn_tensors_dequantised(weights, index); - self.decode_step(weights, token_id) + 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). @@ -716,7 +717,7 @@ impl AnyEngine { /// it (they dequantise + run on f32 internally). pub fn prefill_quant( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, ffn: &dyn FfnBackend, index: &larql_vindex::VectorIndex, token_ids: &[u32], @@ -732,7 +733,7 @@ impl AnyEngine { /// [`prefill_quant`]. pub fn decode_step_quant( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, ffn: &dyn FfnBackend, index: &larql_vindex::VectorIndex, token_id: u32, @@ -781,7 +782,7 @@ impl AnyEngine { /// executor loops). pub fn prefill_quant_via_executor( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, executor: &dyn crate::layer_executor::LayerExecutor, ffn: &dyn FfnBackend, index: &larql_vindex::VectorIndex, @@ -797,7 +798,7 @@ impl AnyEngine { /// fall-back semantics as [`prefill_quant_via_executor`]. pub fn decode_step_quant_via_executor( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, executor: &dyn crate::layer_executor::LayerExecutor, ffn: &dyn FfnBackend, index: &larql_vindex::VectorIndex, @@ -1006,13 +1007,13 @@ mod tests { assert_eq!(engine.decode_calls, 1); // prefill_quant_via_executor → prefill_quant → prefill (default fallback) - let mut weights_q = crate::test_utils::make_test_weights(); - let out = engine.prefill_quant_via_executor(&mut weights_q, &exec, &ffn, &index, &[0, 1]); + 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(&mut weights_q, &exec, &ffn, &index, 3); + let out = engine.decode_step_quant_via_executor(&weights_q, &exec, &ffn, &index, 3); assert!(out.is_ok()); assert_eq!(engine.decode_calls, 2); } @@ -1028,15 +1029,15 @@ mod tests { decode_calls: 0, }; - let mut weights_q4k = crate::test_utils::make_test_weights(); - let out = engine.prefill_quant(&mut weights_q4k, &ffn, &index, &[1, 2, 3], &*backend); + 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(&mut weights_q4k, &ffn, &index, 4, &*backend); + let out = engine.decode_step_quant(&weights_q4k, &ffn, &index, 4, &*backend); assert!(out.is_ok()); assert_eq!( engine.decode_calls, 1, diff --git a/crates/larql-inference/src/layer_executor/local_walk.rs b/crates/larql-inference/src/layer_executor/local_walk.rs index c55e9d518..ad8b97c6a 100644 --- a/crates/larql-inference/src/layer_executor/local_walk.rs +++ b/crates/larql-inference/src/layer_executor/local_walk.rs @@ -26,7 +26,6 @@ use crate::attention::{ }; use crate::ffn::FfnBackend; use crate::forward::run_ffn; -use crate::model::ModelWeights; use larql_compute::ComputeBackend; use super::{ExecutorDispatchKind, LayerExecutor}; @@ -61,7 +60,7 @@ impl<'a> LayerExecutor for LocalWalkExecutor<'a> { fn run_prefill_layer( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, layer: usize, hidden_in: &Array2, ffn: &dyn FfnBackend, @@ -70,18 +69,18 @@ impl<'a> LayerExecutor for LocalWalkExecutor<'a> { // projection matmuls; `run_attention_with_kv_backend` returns // `(h_post_attn, k_rope, v_final)`. let (h_post_attn, k, v) = - run_attention_with_kv_backend(weights, hidden_in, layer, Some(self.backend))?; + run_attention_with_kv_backend(weights, hidden_in, layer, Some(self.backend), None)?; // FFN through the caller-supplied dispatcher. This is the // critical decoupling: local FFN uses `WeightFfn` / `BackendFfn`, // remote FFN uses `RemoteWalkBackend`, MoE shards use // `RemoteMoeBackend`. The executor doesn't pick. - let (h_out, _activation) = run_ffn(weights, &h_post_attn, layer, ffn, false); + let (h_out, _activation) = run_ffn(&weights, &h_post_attn, layer, ffn, false); Some((h_out, (k, v))) } fn run_decode_layer( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, layer: usize, hidden_in: &Array2, prior_kv: &SharedKV, @@ -100,7 +99,7 @@ impl<'a> LayerExecutor for LocalWalkExecutor<'a> { abs_position, Some(self.backend), )?; - let (h_out, _activation) = run_ffn(weights, &h_post_attn, layer, ffn, false); + let (h_out, _activation) = run_ffn(&weights, &h_post_attn, layer, ffn, false); Some((h_out, new_kv)) } } @@ -109,6 +108,7 @@ impl<'a> LayerExecutor for LocalWalkExecutor<'a> { mod tests { use super::*; use crate::ffn::WeightFfn; + use crate::model::ModelWeights; use crate::test_utils::make_test_weights; use larql_compute::CpuBackend; @@ -142,7 +142,12 @@ mod tests { let exec = LocalWalkExecutor::new(&backend); let hidden_in = Array2::from_elem((3, weights.hidden_size), 0.1f32); let (h_out, (k, v)) = exec - .run_prefill_layer(&weights, 0, &hidden_in, &ffn) + .run_prefill_layer( + larql_models::WeightsView::dense(&weights), + 0, + &hidden_in, + &ffn, + ) .expect("prefill_layer"); assert_eq!(h_out.shape(), &[3, weights.hidden_size]); assert!(h_out.iter().all(|v| v.is_finite())); @@ -165,7 +170,7 @@ mod tests { let mut h = crate::forward::embed_tokens_pub(&weights, &[0u32, 1, 2]); for layer in 0..weights.num_layers { let (h_next, _kv) = exec - .run_prefill_layer(&weights, layer, &h, &ffn) + .run_prefill_layer(larql_models::WeightsView::dense(&weights), layer, &h, &ffn) .expect("layer prefill"); assert_eq!(h_next.shape(), &[3, weights.hidden_size]); assert!(h_next.iter().all(|v| v.is_finite())); @@ -184,13 +189,25 @@ mod tests { // Seed K/V from a 2-token prefill, then decode one step. let prefill_hidden = Array2::from_elem((2, weights.hidden_size), 0.1f32); let (_, prior_kv) = exec - .run_prefill_layer(&weights, 0, &prefill_hidden, &ffn) + .run_prefill_layer( + larql_models::WeightsView::dense(&weights), + 0, + &prefill_hidden, + &ffn, + ) .unwrap(); assert_eq!(prior_kv.0.shape()[0], 2); let new_token_hidden = Array2::from_elem((1, weights.hidden_size), 0.05f32); let (h_out, new_kv) = exec - .run_decode_layer(&weights, 0, &new_token_hidden, &prior_kv, 2, &ffn) + .run_decode_layer( + larql_models::WeightsView::dense(&weights), + 0, + &new_token_hidden, + &prior_kv, + 2, + &ffn, + ) .expect("decode_layer"); assert_eq!(h_out.shape(), &[1, weights.hidden_size]); assert!(h_out.iter().all(|v| v.is_finite())); @@ -209,7 +226,14 @@ mod tests { let empty_kv: SharedKV = (Array2::zeros((0, kv_dim)), Array2::zeros((0, kv_dim))); let h_in = Array2::from_elem((1, weights.hidden_size), 0.1f32); let (h_out, new_kv) = exec - .run_decode_layer(&weights, 0, &h_in, &empty_kv, 0, &ffn) + .run_decode_layer( + larql_models::WeightsView::dense(&weights), + 0, + &h_in, + &empty_kv, + 0, + &ffn, + ) .expect("decode_layer with empty prior"); assert_eq!(h_out.shape(), &[1, weights.hidden_size]); assert_eq!(new_kv.0.shape()[0], 1); @@ -229,12 +253,22 @@ mod tests { let ffn_real = WeightFfn { weights: &weights }; let (h_real, _) = exec - .run_prefill_layer(&weights, 0, &h_in, &ffn_real) + .run_prefill_layer( + larql_models::WeightsView::dense(&weights), + 0, + &h_in, + &ffn_real, + ) .unwrap(); let ffn_null = NullFfn; let (h_null, _) = exec - .run_prefill_layer(&weights, 0, &h_in, &ffn_null) + .run_prefill_layer( + larql_models::WeightsView::dense(&weights), + 0, + &h_in, + &ffn_null, + ) .unwrap(); // The two FFN backends produce different outputs; the executor diff --git a/crates/larql-inference/src/layer_executor/mod.rs b/crates/larql-inference/src/layer_executor/mod.rs index 130f4c5ec..082c7fd44 100644 --- a/crates/larql-inference/src/layer_executor/mod.rs +++ b/crates/larql-inference/src/layer_executor/mod.rs @@ -40,7 +40,6 @@ pub use local_walk::LocalWalkExecutor; use crate::attention::SharedKV; use crate::ffn::FfnBackend; -use crate::model::ModelWeights; use ndarray::Array2; /// Whether an executor owns its K/V state internally (`Fused`) or @@ -95,7 +94,7 @@ pub trait LayerExecutor { /// `None` by default — engines should not call this on `Fused`. fn run_prefill_layer( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, layer: usize, hidden_in: &Array2, ffn: &dyn FfnBackend, @@ -119,7 +118,7 @@ pub trait LayerExecutor { /// `None` by default. fn run_decode_layer( &self, - weights: &ModelWeights, + weights: larql_models::WeightsView, layer: usize, hidden_in: &Array2, prior_kv: &SharedKV, @@ -172,13 +171,22 @@ mod tests { let weights = crate::test_utils::make_test_weights(); let ffn = crate::ffn::WeightFfn { weights: &weights }; let hidden = Array2::::zeros((1, weights.hidden_size)); - assert!(exec.run_prefill_layer(&weights, 0, &hidden, &ffn).is_none()); + assert!(exec + .run_prefill_layer(larql_models::WeightsView::dense(&weights), 0, &hidden, &ffn) + .is_none()); // SharedKV = (K, V) as Array2: shape doesn't matter — the // default impl returns None before touching either tensor. let kv: SharedKV = (Array2::::zeros((1, 1)), Array2::::zeros((1, 1))); assert!(exec - .run_decode_layer(&weights, 0, &hidden, &kv, 0, &ffn) + .run_decode_layer( + larql_models::WeightsView::dense(&weights), + 0, + &hidden, + &kv, + 0, + &ffn + ) .is_none()); } diff --git a/crates/larql-inference/src/layer_graph/cached.rs b/crates/larql-inference/src/layer_graph/cached.rs index 3f14c7ac8..8e943dc74 100644 --- a/crates/larql-inference/src/layer_graph/cached.rs +++ b/crates/larql-inference/src/layer_graph/cached.rs @@ -172,8 +172,14 @@ impl AttentionCache { let mut ffn_inputs = Vec::with_capacity(layer_range.len()); for layer in layer_range { // Attention (exact) - let (h_post_attn, _, _) = - crate::attention::run_attention_block_gpu(weights, &h, layer, false, None).unwrap(); + let (h_post_attn, _, _) = crate::attention::run_attention_block_gpu( + larql_models::WeightsView::dense(weights), + &h, + layer, + false, + None, + ) + .unwrap(); // Capture FFN-normed input (last token) let pre_ffn_key = if arch.has_post_norms() { diff --git a/crates/larql-inference/src/layer_graph/dense.rs b/crates/larql-inference/src/layer_graph/dense.rs index 2df2c39f9..8c6579d69 100644 --- a/crates/larql-inference/src/layer_graph/dense.rs +++ b/crates/larql-inference/src/layer_graph/dense.rs @@ -23,7 +23,7 @@ impl<'a> LayerGraph for DenseLayerGraph<'a> { ) -> Option { // Attention: dense matmul (Q·K·V), optionally GPU-accelerated let (h_post_attn, _attn_proj, attn_weights) = crate::attention::run_attention_block_gpu( - weights, + larql_models::WeightsView::dense(weights), h, layer, self.capture_attention, diff --git a/crates/larql-inference/src/layer_graph/predict/honest.rs b/crates/larql-inference/src/layer_graph/predict/honest.rs index dfd17acde..e04a285e6 100644 --- a/crates/larql-inference/src/layer_graph/predict/honest.rs +++ b/crates/larql-inference/src/layer_graph/predict/honest.rs @@ -161,10 +161,11 @@ pub fn predict_honest( for (rel_idx, abs_layer) in layer_range.clone().enumerate() { let (h_post_attn, k_rope, v) = crate::attention::gpu::run_attention_with_kv_backend( - weights, + larql_models::WeightsView::dense(weights), &h_cpu, abs_layer, Some(backend), + None, ) .unwrap(); @@ -217,8 +218,14 @@ pub fn predict_honest( if !used_gpu { let walk_ffn = crate::vindex::WalkFfn::new_unlimited(weights, index); for layer in layer_range { - let (h_post_attn, _, _) = - crate::attention::run_attention_block_gpu(weights, &h, layer, false, None).unwrap(); + let (h_post_attn, _, _) = crate::attention::run_attention_block_gpu( + larql_models::WeightsView::dense(weights), + &h, + layer, + false, + None, + ) + .unwrap(); let (h_out, _) = crate::forward::run_ffn(weights, &h_post_attn, layer, &walk_ffn, false); h = h_out; diff --git a/crates/larql-inference/src/layer_graph/predict/split.rs b/crates/larql-inference/src/layer_graph/predict/split.rs index 98154c04f..8a61771fa 100644 --- a/crates/larql-inference/src/layer_graph/predict/split.rs +++ b/crates/larql-inference/src/layer_graph/predict/split.rs @@ -62,12 +62,17 @@ pub fn predict_split_pass( for layer in layer_range.clone() { // Run attention only (CPU BLAS, no FFN) - let (h_post_attn, _attn_proj, _) = - crate::attention::run_attention_block_gpu(weights, &h, layer, false, None) - .unwrap_or_else(|| { - // Fallback: identity (shouldn't happen with valid weights) - (h.clone(), h.clone(), None) - }); + let (h_post_attn, _attn_proj, _) = crate::attention::run_attention_block_gpu( + larql_models::WeightsView::dense(weights), + &h, + layer, + false, + None, + ) + .unwrap_or_else(|| { + // Fallback: identity (shouldn't happen with valid weights) + (h.clone(), h.clone(), None) + }); // Compute pre-FFN norm (this is the FFN input) let pre_ffn_key = if arch.has_post_norms() { diff --git a/crates/larql-inference/src/layer_graph/prefill.rs b/crates/larql-inference/src/layer_graph/prefill.rs index 503228f53..6ffe4c5d9 100644 --- a/crates/larql-inference/src/layer_graph/prefill.rs +++ b/crates/larql-inference/src/layer_graph/prefill.rs @@ -19,9 +19,14 @@ pub fn prefill_with_kv( let seq_len = token_ids.len(); for layer in layer_range { - let (h_post_attn, k_rope, v) = - crate::attention::gpu::run_attention_with_kv_backend(weights, &h, layer, Some(backend)) - .unwrap(); + let (h_post_attn, k_rope, v) = crate::attention::gpu::run_attention_with_kv_backend( + larql_models::WeightsView::dense(weights), + &h, + layer, + Some(backend), + None, + ) + .unwrap(); if backend.has_kv_cache() { let layer_hd = weights.arch.head_dim_for_layer(layer); diff --git a/crates/larql-inference/src/layer_graph/template.rs b/crates/larql-inference/src/layer_graph/template.rs index f09f27cac..b49956633 100644 --- a/crates/larql-inference/src/layer_graph/template.rs +++ b/crates/larql-inference/src/layer_graph/template.rs @@ -162,8 +162,12 @@ impl<'a> LayerGraph for GuidedWalkLayerGraph<'a> { layer: usize, ) -> Option { // Attention: dense matmul - let (h_post_attn, _attn_proj, _) = - crate::attention::run_attention_block(weights, h, layer, false)?; + let (h_post_attn, _attn_proj, _) = crate::attention::run_attention_block( + larql_models::WeightsView::dense(weights), + h, + layer, + false, + )?; // FFN: guided walk — score only template universe features let residual = guided_walk_ffn(weights, &h_post_attn, layer, self.universe, self.index); diff --git a/crates/larql-inference/src/layer_graph/walk.rs b/crates/larql-inference/src/layer_graph/walk.rs index 0ac28367a..fdfa93c7e 100644 --- a/crates/larql-inference/src/layer_graph/walk.rs +++ b/crates/larql-inference/src/layer_graph/walk.rs @@ -21,8 +21,13 @@ impl<'a> LayerGraph for WalkLayerGraph<'a> { h: &Array2, layer: usize, ) -> Option { - let (h_post_attn, _attn_proj, _) = - crate::attention::run_attention_block_gpu(weights, h, layer, false, self.backend)?; + let (h_post_attn, _attn_proj, _) = crate::attention::run_attention_block_gpu( + larql_models::WeightsView::dense(weights), + h, + layer, + false, + self.backend, + )?; let (h_out, _) = crate::forward::run_ffn(weights, &h_post_attn, layer, self.ffn, false); Some(LayerOutput { residual: h_out, @@ -67,8 +72,13 @@ impl<'a> LayerGraph for PipelinedLayerGraph<'a> { } // Attention: CPU BLAS (fast, no GPU overhead) - let (h_post_attn, _attn_proj, _) = - crate::attention::run_attention_block_gpu(weights, h, layer, false, None)?; + let (h_post_attn, _attn_proj, _) = crate::attention::run_attention_block_gpu( + larql_models::WeightsView::dense(weights), + h, + layer, + false, + None, + )?; // FFN: use WalkFfn which handles Q4 dispatch internally. // WalkFfn checks for Q4 interleaved data and routes to Metal Q4 diff --git a/crates/larql-inference/src/lib.rs b/crates/larql-inference/src/lib.rs index 8e81c0fdc..b9f9d1bf9 100644 --- a/crates/larql-inference/src/lib.rs +++ b/crates/larql-inference/src/lib.rs @@ -264,7 +264,7 @@ pub use layer_graph::{ TemplateUniverse, WalkLayerGraph, }; -pub use model::{load_model_dir, resolve_model_path, ModelWeights}; +pub use model::{load_model_dir, resolve_model_path, DequantScratch, ModelWeights, WeightsView}; pub use tokenizer::{decode_token, decode_token_raw, encode_prompt, load_tokenizer}; pub use trace::{ trace as trace_decomposed, trace_residuals, AnswerWaypoint, BoundaryStore, BoundaryWriter, diff --git a/crates/larql-inference/src/model.rs b/crates/larql-inference/src/model.rs index 750754fa3..f858aeac5 100644 --- a/crates/larql-inference/src/model.rs +++ b/crates/larql-inference/src/model.rs @@ -1,7 +1,7 @@ //! Model loading — imports from larql-models. -pub use larql_models::ModelWeights; pub use larql_models::{ load_model_dir, load_model_dir_validated, load_model_dir_walk_only, load_model_dir_walk_only_validated, resolve_model_path, }; +pub use larql_models::{DequantScratch, ModelWeights, WeightsView}; diff --git a/crates/larql-inference/src/ternary.rs b/crates/larql-inference/src/ternary.rs index 4e114a05e..7f7d5ecaf 100644 --- a/crates/larql-inference/src/ternary.rs +++ b/crates/larql-inference/src/ternary.rs @@ -74,17 +74,25 @@ //! self-contained single-shot-prefill + greedy/temperature decode //! so the BitNet path is verifiable on its own (it was qualified //! end-to-end against the real 2 B model: "capital of France" -> -//! Paris 94.5%). The optimized future path — int8-quantized -//! activations + a sign-select NEON/AVX2 kernel dispatched through -//! the backend, and a shared KV-cache — lines up with the -//! `FormatRoute`/quant-registry roadmap; when that lands, the -//! decode loop here should ride it rather than duplicate it. +//! Paris 94.5%). Status (branch feat/quant-ternary-a8): the +//! int8-quantized-activation (A8) kernel + NEON sign-select now +//! EXIST in `larql_compute::cpu::ops::ternary_matvec`, and this +//! forward already runs on them (`matvec_i2s_a8_f32_into`). What +//! remains is the *shared* path — dispatch through the +//! `QuantFormat`/`FormatRoute` registry (no ternary variant yet) +//! and a shared KV-cache (`larql_kv::KvCache`) in place of the +//! bespoke `BitnetKvCache`. //! //! Net: the legitimate BitNet-specific divergences are the two //! sub-norms and the ReLU² FFN over ternary projections. The //! reimplemented KV/sampling is a maintenance cost acknowledged -//! here, to be folded into the engine once the quantized-activation -//! kernel path exists. +//! here. Its blocking precondition — the quantized-activation kernel +//! — is now met, so folding the decode loop onto the forward-pass +//! spine + a `KvEngine` impl is live roadmap work (ROADMAP "BitNet +//! b1.58 integration hardening"), no longer blocked on a missing +//! kernel. Until a second consumer exists it may stay isolated under +//! the no-premature-extraction rule; the point is the decision is now +//! explicit, not deferred to a non-existent dependency. //! //! ## I2_S layout (dual representation — intentional) //! @@ -102,7 +110,9 @@ //! Conflating the two scrambles every weight; see the decode-fix PR //! and the format spec for the authoritative description. -use larql_compute::cpu::ops::ternary_matvec::{matvec_i2s_f32_into, BitLinearWeight}; +use larql_compute::cpu::ops::ternary_matvec::{ + matvec_i2s_a8_f32_into, matvec_i2s_a8_into, quantize_activation_i8, BitLinearWeight, +}; use ndarray::{Array1, Array2, ArrayView2}; /// One BitLinear-FFN block. Holds three ternary weight tensors @@ -194,8 +204,11 @@ impl BitNetFfn { // 2. gate = ternary(gate.weight) · x_norm // up = ternary(up.weight) · x_norm - matvec_i2s_f32_into(&self.gate, &x_norm, gate).expect("gate shape"); - matvec_i2s_f32_into(&self.up, &x_norm, up).expect("up shape"); + // Both projections share x_norm — quantise it to int8 once (A8) + // and feed both matvecs, instead of re-quantising per call. + let (x_i8, x_scale) = quantize_activation_i8(&x_norm); + matvec_i2s_a8_into(&self.gate, &x_i8, x_scale, gate).expect("gate shape"); + matvec_i2s_a8_into(&self.up, &x_i8, x_scale, up).expect("up shape"); // 3. Squared-ReLU activation (BitNet b1.58 spec) + // element-wise multiply with up. @@ -209,7 +222,7 @@ impl BitNetFfn { rmsnorm_into(hid, &self.ffn_sub_norm, self.eps, &mut hid_norm); // 5. y = ternary(down.weight) · hid_norm - matvec_i2s_f32_into(&self.down, &hid_norm, y).expect("down shape"); + matvec_i2s_a8_f32_into(&self.down, &hid_norm, y).expect("down shape"); } } @@ -260,6 +273,95 @@ mod tests { BitLinearWeight::new(rows, cols, bytes, scales).unwrap() } + /// Parity gate for the A8 wiring: the production FFN forward (now the + /// int8-activation A8 path) must track a full f32-activation reference + /// FFN within int8 tolerance. Validates that swapping `matvec_i2s_f32` + /// → `matvec_i2s_a8_f32` across the forward didn't shift the numerics + /// beyond the intended activation-quantisation error. + #[test] + fn ffn_a8_forward_matches_f32_reference_within_tolerance() { + use larql_compute::cpu::ops::ternary_matvec::matvec_i2s_f32_into; + + fn rand_w(rows: usize, cols: usize, seed: u64) -> BitLinearWeight { + let mut s = seed; + let mut bytes = vec![0u8; rows * cols / 4]; + for b in bytes.iter_mut() { + let mut bv = 0u8; + for slot in 0..4 { + s = s.wrapping_mul(6364136223846793005).wrapping_add(1); + let code = match (s >> 33) % 3 { + 0 => 0b00u8, + 1 => 0b01, + _ => 0b10, + }; + bv |= code << (2 * slot); + } + *b = bv; + } + let scales = (0..rows).map(|r| 0.05 + (r % 5) as f32 * 0.01).collect(); + BitLinearWeight::new(rows, cols, bytes, scales).unwrap() + } + fn synth(n: usize, seed: u64) -> Vec { + let mut s = seed; + (0..n) + .map(|_| { + s = s.wrapping_mul(6364136223846793005).wrapping_add(1); + ((s >> 33) as f32) / (u32::MAX as f32) * 2.0 - 1.0 + }) + .collect() + } + fn cosine(a: &[f32], b: &[f32]) -> f32 { + let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum(); + let na = a.iter().map(|x| x * x).sum::().sqrt(); + let nb = b.iter().map(|x| x * x).sum::().sqrt(); + dot / (na * nb) + } + + let (hidden, inter) = (256usize, 512usize); + let ffn = BitNetFfn { + gate: rand_w(inter, hidden, 1), + up: rand_w(inter, hidden, 2), + down: rand_w(hidden, inter, 3), + ffn_norm: vec![1.0; hidden], + ffn_sub_norm: vec![1.0; inter], + eps: 1e-5, + }; + let x = synth(hidden, 42); + + // Production (A8) forward. + let mut gate = vec![0.0; inter]; + let mut up = vec![0.0; inter]; + let mut hid = vec![0.0; inter]; + let mut y_a8 = vec![0.0; hidden]; + ffn.forward_into(&x, &mut gate, &mut up, &mut hid, &mut y_a8); + + // f32-activation reference forward (same math, f32 kernel). + let mut x_norm = vec![0.0; hidden]; + rmsnorm_into(&x, &ffn.ffn_norm, ffn.eps, &mut x_norm); + let mut g = vec![0.0; inter]; + let mut u = vec![0.0; inter]; + matvec_i2s_f32_into(&ffn.gate, &x_norm, &mut g).unwrap(); + matvec_i2s_f32_into(&ffn.up, &x_norm, &mut u).unwrap(); + let hid_f: Vec = g + .iter() + .zip(&u) + .map(|(gv, uv)| { + let r = gv.max(0.0); + r * r * uv + }) + .collect(); + let mut hid_norm = vec![0.0; inter]; + rmsnorm_into(&hid_f, &ffn.ffn_sub_norm, ffn.eps, &mut hid_norm); + let mut y_f32 = vec![0.0; hidden]; + matvec_i2s_f32_into(&ffn.down, &hid_norm, &mut y_f32).unwrap(); + + let cos = cosine(&y_a8, &y_f32); + assert!( + cos > 0.999, + "A8 FFN forward vs f32 reference cosine {cos} < 0.999" + ); + } + #[test] fn rmsnorm_zero_input_zero_output() { let x = vec![0.0f32; 8]; @@ -561,23 +663,28 @@ pub fn predict_bitnet( ); } - // b. Q/K/V projections via ternary matvec, per token. + // b. Q/K/V projections via ternary matvec, per token. Q/K/V share + // x_norm.row(i), so quantise it to int8 once (A8) and reuse. for i in 0..seq_len { - matvec_i2s_f32_into( + let (x_i8, x_scale) = quantize_activation_i8(x_norm.row(i).as_slice().unwrap()); + matvec_i2s_a8_into( &layer.attn_q, - x_norm.row(i).as_slice().unwrap(), + &x_i8, + x_scale, q.row_mut(i).as_slice_mut().unwrap(), ) .expect("attn_q shape"); - matvec_i2s_f32_into( + matvec_i2s_a8_into( &layer.attn_k, - x_norm.row(i).as_slice().unwrap(), + &x_i8, + x_scale, k.row_mut(i).as_slice_mut().unwrap(), ) .expect("attn_k shape"); - matvec_i2s_f32_into( + matvec_i2s_a8_into( &layer.attn_v, - x_norm.row(i).as_slice().unwrap(), + &x_i8, + x_scale, v.row_mut(i).as_slice_mut().unwrap(), ) .expect("attn_v shape"); @@ -609,7 +716,7 @@ pub fn predict_bitnet( model.eps, attn_pool_norm.row_mut(i).as_slice_mut().unwrap(), ); - matvec_i2s_f32_into( + matvec_i2s_a8_f32_into( &layer.attn_o, attn_pool_norm.row(i).as_slice().unwrap(), attn_out.row_mut(i).as_slice_mut().unwrap(), @@ -707,9 +814,10 @@ fn ffn_forward_after_input_norm( debug_assert_eq!(up.len(), inter); debug_assert_eq!(hid.len(), inter); - // gate / up projections. - matvec_i2s_f32_into(&ffn.gate, x_norm, gate).expect("gate shape"); - matvec_i2s_f32_into(&ffn.up, x_norm, up).expect("up shape"); + // gate / up projections. Both share x_norm — quantise to int8 once (A8). + let (x_i8, x_scale) = quantize_activation_i8(x_norm); + matvec_i2s_a8_into(&ffn.gate, &x_i8, x_scale, gate).expect("gate shape"); + matvec_i2s_a8_into(&ffn.up, &x_i8, x_scale, up).expect("up shape"); // Squared-ReLU activation. for ((g, u), h) in gate.iter().zip(up.iter()).zip(hid.iter_mut()) { @@ -722,7 +830,7 @@ fn ffn_forward_after_input_norm( rmsnorm_into(hid, &ffn.ffn_sub_norm, eps, &mut hid_norm); // Down projection. - matvec_i2s_f32_into(&ffn.down, &hid_norm, y).expect("down shape"); + matvec_i2s_a8_f32_into(&ffn.down, &hid_norm, y).expect("down shape"); } /// Causal-masked scaled-dot-product attention with GQA support. @@ -1072,10 +1180,11 @@ pub fn decode_step(model: &BitnetModel, cache: &mut BitnetKvCache, new_token: u3 &mut x_norm, ); - // b. Q/K/V projections. - matvec_i2s_f32_into(&layer.attn_q, &x_norm, &mut q).expect("attn_q shape"); - matvec_i2s_f32_into(&layer.attn_k, &x_norm, &mut k).expect("attn_k shape"); - matvec_i2s_f32_into(&layer.attn_v, &x_norm, &mut v).expect("attn_v shape"); + // b. Q/K/V projections. Q/K/V share x_norm — quantise once (A8). + let (x_i8, x_scale) = quantize_activation_i8(&x_norm); + matvec_i2s_a8_into(&layer.attn_q, &x_i8, x_scale, &mut q).expect("attn_q shape"); + matvec_i2s_a8_into(&layer.attn_k, &x_i8, x_scale, &mut k).expect("attn_k shape"); + matvec_i2s_a8_into(&layer.attn_v, &x_i8, x_scale, &mut v).expect("attn_v shape"); // c. RoPE on the new token's Q + K only. The cached K // already carries RoPE for positions 0..position-1. @@ -1127,7 +1236,8 @@ pub fn decode_step(model: &BitnetModel, cache: &mut BitnetKvCache, new_token: u3 model.eps, &mut attn_pool_norm, ); - matvec_i2s_f32_into(&layer.attn_o, &attn_pool_norm, &mut attn_out).expect("attn_o shape"); + matvec_i2s_a8_f32_into(&layer.attn_o, &attn_pool_norm, &mut attn_out) + .expect("attn_o shape"); // g. Residual + FFN + residual. for (dst, &src) in h.iter_mut().zip(attn_out.iter()) { @@ -1375,21 +1485,26 @@ fn run_full_forward( ); } for i in 0..seq_len { - matvec_i2s_f32_into( + // Q/K/V share x_norm.row(i) — quantise to int8 once (A8) and reuse. + let (x_i8, x_scale) = quantize_activation_i8(x_norm.row(i).as_slice().unwrap()); + matvec_i2s_a8_into( &layer.attn_q, - x_norm.row(i).as_slice().unwrap(), + &x_i8, + x_scale, q.row_mut(i).as_slice_mut().unwrap(), ) .expect("attn_q shape"); - matvec_i2s_f32_into( + matvec_i2s_a8_into( &layer.attn_k, - x_norm.row(i).as_slice().unwrap(), + &x_i8, + x_scale, k.row_mut(i).as_slice_mut().unwrap(), ) .expect("attn_k shape"); - matvec_i2s_f32_into( + matvec_i2s_a8_into( &layer.attn_v, - x_norm.row(i).as_slice().unwrap(), + &x_i8, + x_scale, v.row_mut(i).as_slice_mut().unwrap(), ) .expect("attn_v shape"); @@ -1425,7 +1540,7 @@ fn run_full_forward( model.eps, attn_pool_norm.row_mut(i).as_slice_mut().unwrap(), ); - matvec_i2s_f32_into( + matvec_i2s_a8_f32_into( &layer.attn_o, attn_pool_norm.row(i).as_slice().unwrap(), attn_out.row_mut(i).as_slice_mut().unwrap(), diff --git a/crates/larql-inference/src/trace/capture.rs b/crates/larql-inference/src/trace/capture.rs index 90c0c4b6b..5c342d981 100644 --- a/crates/larql-inference/src/trace/capture.rs +++ b/crates/larql-inference/src/trace/capture.rs @@ -77,7 +77,7 @@ pub fn trace_residuals( .and_then(|src| kv_cache.get(&src)); let mut hook = TraceLayerHook::default(); let Some((h_out, _, attn_weights, kv_out)) = run_layer_with_capture_hooked( - weights, + larql_models::WeightsView::dense(weights), &h, layer, ffn, @@ -290,7 +290,7 @@ mod tests { let node = t .node(w.num_layers as i32 - 1, tokens.len() - 1) .expect("final trace node"); - let raw = forward_raw_logits(w, tokens, None); + let raw = forward_raw_logits(larql_models::WeightsView::dense(w), tokens, None); let traced_h = Array2::from_shape_vec((1, w.hidden_size), node.residual.clone()).expect("trace row"); diff --git a/crates/larql-inference/src/vindex/dequant.rs b/crates/larql-inference/src/vindex/dequant.rs index 746db7315..bc1cc1a96 100644 --- a/crates/larql-inference/src/vindex/dequant.rs +++ b/crates/larql-inference/src/vindex/dequant.rs @@ -23,21 +23,26 @@ //! See `docs/specs/kv-dispatch-quantization.md`. use crate::model::ModelWeights; +use larql_models::DequantScratch; use larql_vindex::VectorIndex; use ndarray::Array2; -/// Dequantise attention Q4K weights (Q, K, V, O) for all layers into -/// `weights.tensors`. Idempotent — skips layers whose `attn_q_key` is -/// already present in `weights.tensors`. +/// Dequantise attention Q4K weights (Q, K, V, O) for all layers into the +/// engine-owned `scratch` (NOT `weights`, which stays immutable). Idempotent — +/// skips layers already resolvable (in `scratch` or canonical `weights.tensors`). /// /// No-op for layers where `index.attn_kquant_layer_data(layer)` returns /// `None` (i.e., a layer with non-Q4K attention or no Q4K data at all). -pub fn ensure_attn_tensors_dequantised(weights: &mut ModelWeights, index: &VectorIndex) { +pub fn ensure_attn_tensors_dequantised( + scratch: &mut DequantScratch, + weights: &ModelWeights, + index: &VectorIndex, +) { let num_layers = weights.num_layers; for layer in 0..num_layers { let arch = &*weights.arch; let q_key = arch.attn_q_key(layer); - if weights.tensors.contains_key(&q_key) { + if scratch.contains_key(&q_key) || weights.tensors.contains_key(&q_key) { continue; } let Some(attn) = index.attn_kquant_layer_data(layer) else { @@ -56,13 +61,27 @@ pub fn ensure_attn_tensors_dequantised(weights: &mut ModelWeights, index: &Vecto let w_k = dequantize_matrix(attn[1].0, attn[1].1, kv_dim, hidden); let w_v = dequantize_matrix(attn[2].0, attn[2].1, kv_dim, hidden); let w_o = dequantize_matrix(attn[3].0, attn[3].1, hidden, q_dim); - weights.tensors.insert(q_key, w_q.into_shared()); - weights.tensors.insert(k_key, w_k.into_shared()); - weights.tensors.insert(v_key, w_v.into_shared()); - weights.tensors.insert(o_key, w_o.into_shared()); + scratch.insert(q_key, w_q.into_shared()); + scratch.insert(k_key, w_k.into_shared()); + scratch.insert(v_key, w_v.into_shared()); + scratch.insert(o_key, w_o.into_shared()); } } +/// Resident variant of [`ensure_attn_tensors_dequantised`] — dequantises the +/// attention tensors **into `weights.tensors`** (mutating `weights`) rather than +/// an engine scratch. For the bulk f32-fallback drivers (the vision/image CLI +/// path, dev tests) that run their forward against canonical `weights.tensors` +/// and don't thread a [`WeightsView`]. The KvEngine decode path uses the +/// scratch form + `WeightsView::with_scratch` to keep `ModelWeights` immutable. +/// +/// [`WeightsView`]: larql_models::WeightsView +pub fn ensure_attn_tensors_dequantised_resident(weights: &mut ModelWeights, index: &VectorIndex) { + let mut scratch = DequantScratch::new(); + ensure_attn_tensors_dequantised(&mut scratch, weights, index); + weights.tensors.extend(scratch); +} + fn dequantize_matrix(bytes: &[u8], format: &str, rows: usize, cols: usize) -> Array2 { let n = rows * cols; let padded = n.div_ceil(256) * 256; @@ -131,7 +150,7 @@ mod tests { weights.tensors.remove(v); weights.tensors.remove(o); } - ensure_attn_tensors_dequantised(&mut weights, &index); + ensure_attn_tensors_dequantised_resident(&mut weights, &index); for (l, (q, k, v, o)) in keys.iter().enumerate() { assert!(weights.tensors.contains_key(q), "Q missing layer {l}"); assert!(weights.tensors.contains_key(k), "K missing layer {l}"); @@ -148,13 +167,13 @@ mod tests { let mut weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let q_key = weights.arch.attn_q_key(0); - ensure_attn_tensors_dequantised(&mut weights, &index); + ensure_attn_tensors_dequantised_resident(&mut weights, &index); let q_ptr_before = weights .tensors .get(&q_key) .expect("Q present after first dequant") .as_ptr(); - ensure_attn_tensors_dequantised(&mut weights, &index); + ensure_attn_tensors_dequantised_resident(&mut weights, &index); let q_ptr_after = weights.tensors.get(&q_key).unwrap().as_ptr(); assert_eq!( q_ptr_before, q_ptr_after, @@ -175,7 +194,7 @@ mod tests { ); let q_key = weights.arch.attn_q_key(0); weights.tensors.remove(&q_key); - ensure_attn_tensors_dequantised(&mut weights, &empty_index); + ensure_attn_tensors_dequantised_resident(&mut weights, &empty_index); assert!( !weights.tensors.contains_key(&q_key), "no Q4K data → no insert" @@ -268,7 +287,7 @@ mod tests { deq_weights.tensors.remove(v); deq_weights.tensors.remove(o); } - ensure_attn_tensors_dequantised(&mut deq_weights, index); + ensure_attn_tensors_dequantised_resident(&mut deq_weights, index); deq_weights } @@ -289,7 +308,7 @@ mod tests { for layer in 0..weights.num_layers { let (h_deq, (k_deq, v_deq)) = run_attention_block_decode_step_backend( - &deq_weights, + larql_models::WeightsView::dense(&deq_weights), &h_new, layer, None, @@ -384,7 +403,7 @@ mod tests { (((j + step) % 11) as f32 - 5.0) * 0.03 }); let (h_deq, new_deq) = run_attention_block_decode_step_backend( - &deq_weights, + larql_models::WeightsView::dense(&deq_weights), &h_new, layer, kv_deq.as_ref(), diff --git a/crates/larql-inference/src/vindex/kquant_forward/cached.rs b/crates/larql-inference/src/vindex/kquant_forward/cached.rs index 32ec47287..60cfc7d7a 100644 --- a/crates/larql-inference/src/vindex/kquant_forward/cached.rs +++ b/crates/larql-inference/src/vindex/kquant_forward/cached.rs @@ -1,51 +1,23 @@ -//! KV-cached CPU Q4_K decode. +//! KV-cached CPU Q4_K decode — the `VectorIndex`-typed adapter over the +//! substrate forward in [`larql_compute::kquant_forward`]. //! -//! `predict_kquant_hidden` (sibling module) reprocesses the entire -//! `token_ids` sequence at every decode step — O(N²) work where N -//! grows with each generated token. This module splits that into -//! prefill (full-sequence pass that captures K/V per layer) plus -//! per-step decode (single-row attention against the cache + 1-row -//! FFN). Speedup scales linearly with decode length. +//! Every function here coerces `&VectorIndex` to `&dyn larql_compute::KvIndex` +//! and delegates: larql-compute owns the single implementation of the Q4_K CPU +//! prefill/decode (ADR-0022), including the q4k-direct FFN prefill. This module +//! exists only to keep a stable, `VectorIndex`-typed API and the local +//! `CachedTimings` type for callers that hold a concrete `VectorIndex`. //! -//! Per-step Q4_K → f32 dequant via `insert_q4k_layer_tensors` is -//! still paid for now; eliminating it is a follow-up (route Q/K/V/O -//! and gate/up/down through `backend.q4k_matvec` directly). -//! -//! Scope: dense architectures only. Hybrid-MoE (Gemma 4 26B A4B) -//! and cross-layer KV sharing (Gemma 4 E2B) fall back to the slow -//! `predict_kquant_hidden` path — the caller decides via -//! [`supports_cached_decode`]. - -// `cache[layer]` indexing reads more naturally than the iterator -// equivalent and pairs cleanly with the explicit `layer` ID that's -// passed into `insert_q4k_layer_tensors` / `run_attention_block_*`. -// The `(Array2, (Array2, Array2))` return is the documented -// `(h_post_attn, (k_cache, v_cache))` shape used across the decode -// helpers; introducing a type alias would just spread the shape -// across two files. -#![allow(clippy::needless_range_loop, clippy::type_complexity)] - -use larql_compute::cpu::ops::q4k_q8k_dot::{quantize_x_to_q8k_into, Q8KActivation}; +//! The unique, non-delegated Q4_K paths — `predict_kquant_hidden`, the OV/RD +//! interventions, Metal, MoE, and streaming generation — live in the sibling +//! modules of `vindex::kquant_forward`, not here. + +#![allow(clippy::type_complexity)] + use larql_compute::ComputeBackend; use larql_models::ModelWeights; use larql_vindex::VectorIndex; use ndarray::Array2; -use crate::attention::{ - decode::{gqa_attention_decode_step, run_attention_block_decode_step_backend}, - rope::apply_rope_partial_at_full, - run_attention_with_kv_backend, -}; -use crate::ffn::WeightFfn; -use crate::forward::embed_tokens_pub; -use crate::forward::layer::apply_layer_scalar; -use crate::forward::ple::{apply_per_layer_embedding, precompute_per_layer_inputs}; -use crate::forward::run_ffn; -use crate::forward::{add_bias, apply_norm}; -use crate::residual::{rms_norm_heads, rms_norm_heads_no_weight}; - -use super::tensors::{insert_q4k_layer_tensors, remove_layer_tensors}; - /// Per-layer K/V captured during prefill. One entry per layer; matches /// the [`crate::attention::decode::KvCache`] convention so future work /// can swap in window clipping or surgery without churn here. @@ -71,15 +43,7 @@ impl CachedTimings { /// attention helper only knows the "this layer has its own K/V" case /// today). pub fn supports_cached_decode(weights: &ModelWeights) -> bool { - if weights.arch.is_hybrid_moe() { - return false; - } - for layer in 0..weights.num_layers { - if weights.arch.kv_shared_source_layer(layer).is_some() { - return false; - } - } - true + larql_compute::kquant_forward::supports_cached_decode(weights) } /// Prefill: run the full prompt through every layer once, capturing @@ -103,69 +67,28 @@ pub fn predict_kquant_prefill( /// 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( - weights: &mut ModelWeights, + weights: &ModelWeights, token_ids: &[u32], index: &VectorIndex, - mut state: Option<&mut crate::PerLayerDecodeState>, + state: Option<&mut crate::PerLayerDecodeState>, ) -> (Array2, CpuKvCache, CachedTimings) { - let num_layers = weights.num_layers; - let mut cache: CpuKvCache = vec![None; num_layers]; - let mut timings = CachedTimings::default(); - - let mut h = embed_tokens_pub(weights, token_ids); - let ple_inputs = precompute_per_layer_inputs(weights, &h, token_ids); - - for layer in 0..num_layers { - let t0 = std::time::Instant::now(); - let inserted = - insert_q4k_layer_tensors(weights, index, layer).unwrap_or_else(|err| panic!("{err}")); - timings.dequant_ms += t0.elapsed().as_secs_f64() * 1000.0; - - // Snapshot pre-attention residual for this layer if engine wants it. - if let Some(s) = state.as_deref_mut() { - s.h_in_per_layer - .push(larql_compute::state_handle::CpuStateHandle::boxed( - h.clone(), - )); - } - - // Attention with K/V capture. Backend stays None — we want the - // CPU BLAS path for the dequantised f32 tensors that - // `insert_q4k_layer_tensors` just placed in `weights.tensors`. - let (h_post_attn, k_rope, v_final) = - match run_attention_with_kv_backend(weights, &h, layer, None) { - Some(t) => t, - None => { - remove_layer_tensors(weights, inserted); - return (h, cache, timings); - } - }; - - if let Some(s) = state.as_deref_mut() { - // Prefill K/V for THIS layer = full seq_len × kv_dim. - s.k_new_per_layer - .push(larql_compute::state_handle::CpuStateHandle::boxed( - k_rope.clone(), - )); - s.v_new_per_layer - .push(larql_compute::state_handle::CpuStateHandle::boxed( - v_final.clone(), - )); - } - - let ffn = WeightFfn { weights }; - 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); - - remove_layer_tensors(weights, inserted); - - cache[layer] = Some((k_rope, v_final)); - h = h_out; - } - - (h, cache, timings) + // Delegate to the substrate copy in larql-compute — the single source of + // truth for the Q4_K CPU forward (ADR-0022). `VectorIndex` satisfies + // `larql_compute::KvIndex`; the only impedance is `CachedTimings`, a + // same-shape struct defined in each crate. + let (h, cache, timings) = larql_compute::kquant_forward::predict_kquant_prefill_with_state( + weights, + token_ids, + index as &dyn larql_compute::KvIndex, + state, + ); + ( + h, + cache, + CachedTimings { + dequant_ms: timings.dequant_ms, + }, + ) } /// Decode step: run a single new token through every layer using the @@ -176,57 +99,29 @@ pub fn predict_kquant_prefill_with_state( /// `prompt_len + steps_already_decoded`. The caller maintains this /// counter (typical: `prompt_len + step_index` starting at 0). pub fn predict_kquant_decode_step( - weights: &mut ModelWeights, + weights: &ModelWeights, token_id: u32, index: &VectorIndex, cache: &mut CpuKvCache, abs_position: usize, ) -> Option<(Array2, CachedTimings)> { - let num_layers = weights.num_layers; - if cache.len() != num_layers { - return None; - } - let mut timings = CachedTimings::default(); - - // 1-row embed + 1-row PLE for the new token. - let mut h = embed_tokens_pub(weights, &[token_id]); - let ple_inputs = precompute_per_layer_inputs(weights, &h, &[token_id]); - - for layer in 0..num_layers { - let t0 = std::time::Instant::now(); - let inserted = - insert_q4k_layer_tensors(weights, index, layer).unwrap_or_else(|err| panic!("{err}")); - timings.dequant_ms += t0.elapsed().as_secs_f64() * 1000.0; - - let kv_entry = cache[layer].as_ref(); - let (h_post_attn, new_kv) = match run_attention_block_decode_step_backend( - weights, - &h, - layer, - kv_entry, - abs_position, - None, - ) { - Some(t) => t, - None => { - remove_layer_tensors(weights, inserted); - return None; - } - }; - cache[layer] = Some(new_kv); - - let ffn = WeightFfn { weights }; - 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); - - remove_layer_tensors(weights, inserted); - - h = h_out; - } - - Some((h, timings)) + // Delegate to the substrate copy in larql-compute (ADR-0022); bridge the + // per-crate `CachedTimings`. + larql_compute::kquant_forward::predict_kquant_decode_step( + weights, + token_id, + index as &dyn larql_compute::KvIndex, + cache, + abs_position, + ) + .map(|(h, timings)| { + ( + h, + CachedTimings { + dequant_ms: timings.dequant_ms, + }, + ) + }) } impl CachedTimings { @@ -247,214 +142,48 @@ impl CachedTimings { // (CPU `q4k_matvec_into` / `q6k_matvec_into`), skipping the dequant // staging entirely. -/// Format-aware Q*K × Q8_K matvec used by the production decode path. -/// Uses NEON `sdot` (Q4_K) or `vmlal_s8` (Q6_K) under the hood — ~2-3× -/// the f32-FMA throughput of `backend.quant_matvec`. Returns `None` -/// for any unsupported format (caller falls through to dequant). -fn matvec_q4k_or_q6k_q8k( - bytes: &[u8], - format: &str, - x_q8k: &Q8KActivation, - rows: usize, - cols: usize, -) -> Option> { - if rows == 0 || cols == 0 { - return Some(vec![0.0f32; rows]); - } - const ELEMS_PER_BLOCK: usize = 256; - if !cols.is_multiple_of(ELEMS_PER_BLOCK) { - return None; - } - // Gate on the two kernel-backed formats and take the packed row length - // from the format helper rather than re-spelling `(cols/256)*144` (twin - // of `larql_compute::cpu::ops::q4k_q8k_dot` — kept in sync via - // `from_registry_tag`). `%256` is guarded above, so `packed_matrix_bytes` - // is exact. - let bytes_per_row = match larql_compute::QuantFormat::from_registry_tag(format) { - Some(f @ (larql_compute::QuantFormat::Q4_K | larql_compute::QuantFormat::Q6_K)) => { - f.packed_matrix_bytes(1, cols)? - } - _ => return None, - }; - if bytes.len() < rows * bytes_per_row { - return None; - } - - // `q4k_q8k_matvec_into` (larql-compute) is a single-threaded - // per-row loop. Wrap it with `par_chunks_mut(CHUNK_ROWS)` here so - // every Q4_K/Q6_K × Q8_K matvec on the decode path scales across - // the 11 perf cores on M3 Max — matching the rayon strategy of - // `q4k_matvec_into` in `q4_common.rs`. Without this, decode runs - // single-threaded and the sdot path actually regresses vs the - // (rayon-parallel) f32 path despite each row being faster. - let mut out = vec![0.0f32; rows]; - larql_compute::cpu::ops::q4k_q8k_dot::q4k_q8k_matvec_parallel( - &mut out, x_q8k, bytes, rows, cols, format, - ); - Some(out) -} - -/// True when every Q/K/V/O + gate/up/down slice for `layer` is in a -/// format the direct-matvec path knows how to handle. Used to gate -/// per-layer routing: the cached decode step prefers the direct -/// matvec when this returns true and falls back to the dequant path -/// otherwise (e.g. Q4_KF layers, padded down projections). -fn layer_supports_direct_matvec(index: &VectorIndex, layer: usize) -> bool { - let attn = match index.attn_kquant_layer_data(layer) { - Some(a) => a, - None => return false, - }; - for (_, fmt) in attn.iter() { - if !matches!(*fmt, "Q4_K" | "Q6_K") { - return false; - } - } - let ffn = match index.interleaved_kquant_layer_data(layer) { - Some(f) => f, - None => return false, - }; - for (_, fmt) in ffn.iter() { - if !matches!(*fmt, "Q4_K" | "Q6_K") { - return false; - } - } - // The down projection in the FFN is sometimes stored with a padded - // intermediate dim (rounded up to a 256-multiple). `q4k_matvec_into` - // rejects non-multiple `cols`, which would silently zero the - // output — refuse the direct path so the dequant fallback runs. - let intermediate = index.num_features(layer); - intermediate.is_multiple_of(larql_models::quant::ggml::Q4_K_BLOCK_ELEMS) -} - /// True when the whole model can run on the direct-matvec decode path. /// Same gating as [`supports_cached_decode`] plus a per-layer format /// check. Used by the bench labeler and as the cpu.rs routing key. pub fn supports_direct_matvec_decode(weights: &ModelWeights, index: &VectorIndex) -> bool { - if !supports_cached_decode(weights) { - return false; - } - for layer in 0..weights.num_layers { - if !layer_supports_direct_matvec(index, layer) { - return false; - } - } - true -} - -fn vec_to_2d_row(v: Vec) -> Array2 { - let n = v.len(); - Array2::from_shape_vec((1, n), v).expect("matvec output shape") + larql_compute::kquant_forward::supports_direct_matvec_decode( + weights, + index as &dyn larql_compute::KvIndex, + ) } -/// One-row attention block using direct Q4_K/Q6_K matvec on the -/// quantised attention slices. Mirrors -/// [`crate::attention::decode::run_attention_block_decode_step_backend`] -/// but reads weights from `index.attn_kquant_layer_data(layer)` instead of -/// dequantised f32 in `weights.tensors`. -#[allow(clippy::too_many_arguments)] -/// Metal-fused prefill: run the prompt through every layer via the -/// backend's `prefill_kquant` kernel (one command buffer per session), seed -/// the backend's internal K/V cache, return last-row hidden. -/// -/// Returns `None` if the backend doesn't have Q4 support -/// (`!backend.supports_quant(::larql_compute::QuantFormat::Q4_K)`), the vindex lacks Q4K/Q4_0 interleaved FFN -/// bytes, or the architecture isn't compatible with the fused pipeline. -/// CPU callers get `None` — they use [`predict_kquant_prefill`] instead. -/// -/// Public counterpart to [`predict_kquant_prefill`] for the Metal side. -/// Previously lived inline in `larql-kv/engines/unlimited_context/engine.rs`; -/// promoted here so [`crate::kv_dispatch::metal::MetalBackend::coarse_prefill`] -/// can use it without an `larql-inference → larql-kv` dep cycle. +/// Fused Q4_K prefill via the compute backend (Metal fast path). +/// Delegates to the larql-compute substrate copy (ADR-0022). pub fn fused_prefill( weights: &ModelWeights, index: &VectorIndex, token_ids: &[u32], backend: &dyn ComputeBackend, ) -> Option> { - use crate::layer_graph::pipeline_layer::{build_pipeline_layers, DEFAULT_GPU_KV_CACHE_MAX_SEQ}; - use larql_vindex::GateIndex; - - if !backend.supports_quant(::larql_compute::QuantFormat::Q4_K) { - return None; - } - - let gate_index: &dyn GateIndex = index; - let (q4_ffn_mmap, ffn_is_q4k) = if let Some(m) = gate_index.interleaved_kquant_mmap_ref() { - (m, true) - } else if let Some(m) = gate_index.interleaved_q4_mmap_ref() { - (m, false) - } else { - return None; - }; - index.attn_kquant_layer_data(0)?; - - let arch = &*weights.arch; - let hidden = weights.hidden_size; - let num_layers = weights.num_layers; - let intermediate = gate_index.num_features(0); - if intermediate == 0 { - return None; - } - - let ffn_format = if ffn_is_q4k { - larql_compute::QuantFormat::Q4_K - } else { - larql_compute::QuantFormat::Q4_0 - }; - let q4_ffn_per_matrix = ffn_format.packed_matrix_bytes(intermediate, hidden)?; - - let layers = build_pipeline_layers( + larql_compute::kquant_forward::fused_prefill( weights, - index, - 0..num_layers, - q4_ffn_mmap, - q4_ffn_per_matrix, - ffn_format, - ); - - let h_embed = crate::forward::embed_tokens_pub(weights, token_ids); - let x: Vec = h_embed.as_slice().unwrap_or(&[]).to_vec(); - - let seq_len = token_ids.len(); - let softcap = arch.attn_logit_softcapping().unwrap_or(0.0); - let qk_norm = arch.attn_q_norm_key(0).is_some(); - - backend.reset_kv_cache(); - { - let kv_shapes: Vec<(usize, usize)> = (0..num_layers) - .map(|l| (arch.num_kv_heads_for_layer(l), arch.head_dim_for_layer(l))) - .collect(); - backend.preallocate_kv_cache_per_layer(&kv_shapes, DEFAULT_GPU_KV_CACHE_MAX_SEQ); - } - - let h_vec = - backend.prefill_kquant(&layers, &x, hidden, intermediate, seq_len, qk_norm, softcap)?; - - let h_2d = Array2::from_shape_vec((seq_len, hidden), h_vec).ok()?; - let last = h_2d.shape()[0] - 1; - Some(h_2d.slice(ndarray::s![last..=last, ..]).to_owned()) + index as &dyn larql_compute::KvIndex, + token_ids, + backend, + ) } -/// Metal-fused single-token decode: run one token through all layers via -/// the backend's fused `decode_token` kernel, using the K/V cache -/// populated by a prior [`fused_prefill`] call on the same backend. -/// -/// Returns `None` for CPU backends (no fused `decode_token` impl) and -/// for vindex shapes the fused pipeline can't handle. Public counterpart -/// to [`predict_kquant_decode_step_direct`] for the Metal side. +/// Fused Q4_K decode step via the compute backend. Delegates to larql-compute. pub fn fused_decode_step( weights: &ModelWeights, index: &VectorIndex, token_id: u32, backend: &dyn ComputeBackend, ) -> Option> { - fused_decode_step_inner(weights, index, token_id, backend, None) + larql_compute::kquant_forward::fused_decode_step( + weights, + index as &dyn larql_compute::KvIndex, + token_id, + backend, + ) } -/// Variant of [`fused_decode_step`] that also captures per-layer state -/// via the backend's `decode_token_with_state_dump`. Engines pass -/// `Some(state)` to drive their state policy without a CPU re-walk. -/// `None` is bit-identical to the non-state variant. +/// Fused Q4_K decode step with a per-layer state dump. Delegates to larql-compute. pub fn fused_decode_step_with_state( weights: &ModelWeights, index: &VectorIndex, @@ -462,283 +191,40 @@ pub fn fused_decode_step_with_state( backend: &dyn ComputeBackend, state: &mut larql_compute::DecodeStateDump, ) -> Option> { - fused_decode_step_inner(weights, index, token_id, backend, Some(state)) -} - -fn fused_decode_step_inner( - weights: &ModelWeights, - index: &VectorIndex, - token_id: u32, - backend: &dyn ComputeBackend, - state: Option<&mut larql_compute::DecodeStateDump>, -) -> Option> { - use crate::layer_graph::pipeline_layer::build_pipeline_layers; - use larql_vindex::GateIndex; - - let gate_index: &dyn GateIndex = index; - let (q4_ffn_mmap, ffn_is_q4k) = if let Some(m) = gate_index.interleaved_kquant_mmap_ref() { - (m, true) - } else if let Some(m) = gate_index.interleaved_q4_mmap_ref() { - (m, false) - } else { - return None; - }; - - let hidden = weights.hidden_size; - let num_layers = weights.num_layers; - let intermediate = gate_index.num_features(0); - - let ffn_format = if ffn_is_q4k { - larql_compute::QuantFormat::Q4_K - } else { - larql_compute::QuantFormat::Q4_0 - }; - let q4_ffn_per_matrix = ffn_format.packed_matrix_bytes(intermediate, hidden)?; - - let layers = build_pipeline_layers( + larql_compute::kquant_forward::fused_decode_step_with_state( weights, - index, - 0..num_layers, - q4_ffn_mmap, - q4_ffn_per_matrix, - ffn_format, - ); - - let h_tok = crate::forward::embed_tokens_pub(weights, &[token_id]); - let x_dec: Vec = h_tok.row(0).to_vec(); - - let h_vec = - backend.decode_token_with_state_dump(&layers, &x_dec, hidden, intermediate, state)?; - Array2::from_shape_vec((1, hidden), h_vec).ok() + index as &dyn larql_compute::KvIndex, + token_id, + backend, + state, + ) } -/// Production-path attention decode step reading **quantised** weights -/// from the vindex (not f32 dequantised tensors). Same input/output -/// shape as -/// [`crate::attention::run_attention_block_decode_step_backend`], but -/// reads `index.attn_kquant_layer_data(layer)` directly and dispatches -/// the Q/K/V/O projections to the backend's native quantised matvec -/// (today Q4K / Q4_KF / Q6K via `q4k_matvec_q8_input`). Extending to -/// new quantised formats is internal to this function — the public -/// signature stays format-agnostic. -/// -/// Used by `StandardEngine`'s coarse path and by research engines -/// (`MarkovResidual`, `UnlimitedContext`, `TurboQuant`) that want the -/// production decode kernel without inheriting the per-layer dispatch -/// trait's cached-K/V shape. -/// -/// `h_new` must be a single-row residual (1 × hidden). Multi-row -/// prefill is handled by `predict_kquant_prefill` (separate shape; the -/// `q4k_` in that name is pre-existing debt — see ROADMAP U8/U9 for -/// the broader quant-agnostic rename of the kquant_forward module). -/// -/// Returns `None` if the layer has no quantised attention data in the -/// index or if the backend's matvec for the format is unavailable. +/// Dequant-free attention decode step (Q4_K/Q6_K x Q8_K direct matvec). +/// Delegates to the larql-compute substrate copy (ADR-0022). +#[allow(clippy::type_complexity)] pub fn attention_decode_step_native( weights: &ModelWeights, index: &VectorIndex, - // Kept on the helper signature for parity with the outer - // `predict_kquant_decode_step_direct` API and any future asm dispatch - // that wants runtime feature detection. - _backend: &dyn ComputeBackend, + backend: &dyn ComputeBackend, h_new: &Array2, layer: usize, kv_entry: Option<&(Array2, Array2)>, abs_position: usize, ) -> Option<(Array2, (Array2, Array2))> { - let arch = &*weights.arch; - let hidden = weights.hidden_size; - let head_dim = arch.head_dim_for_layer(layer); - let num_q = arch.num_q_heads_for_layer(layer); - let num_kv = arch.num_kv_heads_for_layer(layer); - let reps = num_q / num_kv; - let q_dim = num_q * head_dim; - let kv_dim = num_kv * head_dim; - let scale = if arch.attention_multiplier() != 1.0 { - arch.attention_multiplier() as f64 - } else { - arch.attention_scale_for_layer(layer) - }; - let norm_offset = arch.norm_weight_offset(); - - let h_norm = apply_norm( + larql_compute::kquant_forward::attention_decode_step_native( weights, + index as &dyn larql_compute::KvIndex, + backend, h_new, - &arch.input_layernorm_key(layer), - norm_offset, - ); - let h_norm_row: &[f32] = h_norm.row(0).to_slice().or_else(|| h_norm.as_slice())?; - - let attn = index.attn_kquant_layer_data(layer)?; - let (q_bytes, q_fmt) = attn[0]; - let (k_bytes, k_fmt) = attn[1]; - let (v_bytes, v_fmt) = attn[2]; - let (o_bytes, o_fmt) = attn[3]; - - // Q8_K-quantise `h_norm` once and reuse for Q / K / V projections. - // sdot int8 dot is ~2-3× the f32 FMA throughput of the - // `q4k_matvec_into` path; the quantisation step itself is O(hidden) - // and amortises across the three projections (and O after attn). - let mut h_norm_q8k = Q8KActivation::with_capacity(hidden); - quantize_x_to_q8k_into(&mut h_norm_q8k, h_norm_row); - - let q_vec = matvec_q4k_or_q6k_q8k(q_bytes, q_fmt, &h_norm_q8k, q_dim, hidden)?; - let mut q_full = vec_to_2d_row(q_vec); - if let Some(bias) = arch - .attn_q_bias_key(layer) - .and_then(|k| weights.vectors.get(&k)) - { - add_bias(&mut q_full, bias); - } - - let qk_offset = arch.qk_norm_weight_offset(); - let qk_norm_off = if qk_offset != 0.0 { - qk_offset - } else { - norm_offset - }; - let q_normed = match arch - .attn_q_norm_key(layer) - .and_then(|k| weights.vectors.get(&k)) - { - Some(norm_w) => rms_norm_heads(&q_full, norm_w, num_q, head_dim, qk_norm_off), - None => q_full, - }; - // RoPE must match the staged path / prefill exactly: override-aware - // base, the per-layer position divisor (Gemma 3 linear rope_scaling - // applies ÷factor on GLOBAL layers only), and llama3 frequency - // scaling. The unscaled `apply_rope_partial_at` here was the direct- - // path divergence on gemma3-4b (global-layer K/Q rope'd at 8× the - // position the prefill cache used — `ave_direct_step_parity`). - let layer_rope_base = crate::forward_overrides::effective_rope_base_for_layer(arch, layer); - let rotary_frac = arch.rotary_fraction_for_layer(layer); - let pos_divisor = - crate::forward_overrides::effective_rope_position_divisor_for_layer(arch, layer); - let llama3 = crate::forward_overrides::effective_llama3_rope_scaling(arch); - let q_rope = apply_rope_partial_at_full( - &q_normed, - num_q, - head_dim, - layer_rope_base, - rotary_frac, - abs_position, - pos_divisor, - llama3, - ); - - let k_vec = matvec_q4k_or_q6k_q8k(k_bytes, k_fmt, &h_norm_q8k, kv_dim, hidden)?; - let v_vec = matvec_q4k_or_q6k_q8k(v_bytes, v_fmt, &h_norm_q8k, kv_dim, hidden)?; - let mut k_full_new = vec_to_2d_row(k_vec); - let mut v_full_new = vec_to_2d_row(v_vec); - if let Some(bias) = arch - .attn_k_bias_key(layer) - .and_then(|k| weights.vectors.get(&k)) - { - add_bias(&mut k_full_new, bias); - } - if let Some(bias) = arch - .attn_v_bias_key(layer) - .and_then(|k| weights.vectors.get(&k)) - { - add_bias(&mut v_full_new, bias); - } - if arch.has_v_norm() { - v_full_new = rms_norm_heads_no_weight(&v_full_new, num_kv, head_dim); - } - let k_normed = match arch - .attn_k_norm_key(layer) - .and_then(|k| weights.vectors.get(&k)) - { - Some(norm_w) => rms_norm_heads(&k_full_new, norm_w, num_kv, head_dim, qk_norm_off), - None => k_full_new, - }; - let k_new_rope = apply_rope_partial_at_full( - &k_normed, - num_kv, - head_dim, - layer_rope_base, - rotary_frac, + layer, + kv_entry, abs_position, - pos_divisor, - llama3, - ); - - let (k_concat, v_concat) = match kv_entry { - Some((k_cached, v_cached)) => { - let total = k_cached.shape()[0] + 1; - let mut k_out = Array2::::zeros((total, kv_dim)); - let mut v_out = Array2::::zeros((total, kv_dim)); - k_out - .slice_mut(ndarray::s![..k_cached.shape()[0], ..]) - .assign(k_cached); - v_out - .slice_mut(ndarray::s![..v_cached.shape()[0], ..]) - .assign(v_cached); - k_out - .slice_mut(ndarray::s![k_cached.shape()[0].., ..]) - .assign(&k_new_rope); - v_out - .slice_mut(ndarray::s![v_cached.shape()[0].., ..]) - .assign(&v_full_new); - (k_out, v_out) - } - None => (k_new_rope, v_full_new), - }; - - let softcap = arch.attn_logit_softcapping(); - let attn_out = gqa_attention_decode_step( - &q_rope, &k_concat, &v_concat, num_q, head_dim, reps, scale, softcap, - ); - let attn_out_row: &[f32] = attn_out.row(0).to_slice().or_else(|| attn_out.as_slice())?; - - // Re-quantise the attention output for the O projection. Different - // input from Q/K/V (attn_out vs h_norm), so we need a fresh Q8_K. - let mut attn_out_q8k = Q8KActivation::with_capacity(q_dim); - quantize_x_to_q8k_into(&mut attn_out_q8k, attn_out_row); - let o_vec = matvec_q4k_or_q6k_q8k(o_bytes, o_fmt, &attn_out_q8k, hidden, q_dim)?; - let mut attn_projected = vec_to_2d_row(o_vec); - if let Some(bias) = arch - .attn_o_bias_key(layer) - .and_then(|k| weights.vectors.get(&k)) - { - add_bias(&mut attn_projected, bias); - } - - let res_mult = arch.residual_multiplier(); - let h_post_attn = if arch.has_post_norms() { - let normed = apply_norm( - weights, - &attn_projected, - &arch.post_attention_layernorm_key(layer), - norm_offset, - ); - if res_mult != 1.0 { - h_new + &(&normed * res_mult) - } else { - h_new + &normed - } - } else if res_mult != 1.0 { - h_new + &(&attn_projected * res_mult) - } else { - h_new + &attn_projected - }; - - Some((h_post_attn, (k_concat, v_concat))) + ) } -/// One-row gated FFN block using direct native-quantised matvec on -/// the vindex's compact bytes (Q4K / Q6K today). Mirrors -/// [`crate::ffn::weight::dense_ffn_forward_backend`] but reads gate/up/ -/// down from the vindex slices and avoids the f32 staging — same -/// production path that powers `larql run` / `larql bench --cpu` at -/// ~24 tok/s on Gemma 3 4B Q4K (M3 Max, 8 threads). -/// -/// Returns `None` if the vindex layer lacks compact FFN bytes or the -/// architecture isn't supported by the direct-matvec path. Engines -/// that get `None` fall back to whichever `FfnBackend` they have. -/// -/// `h_post_attn` must be a single-row residual (1 × hidden). Public -/// counterpart to [`attention_decode_step_native`] for the FFN side. +/// Dequant-free FFN decode step (gate/up/down via direct Q4_K matvec). +/// Delegates to the larql-compute substrate copy (ADR-0022). pub fn ffn_decode_step_native( weights: &ModelWeights, index: &VectorIndex, @@ -746,162 +232,13 @@ pub fn ffn_decode_step_native( h_post_attn: &Array2, layer: usize, ) -> Option> { - run_ffn_decode_step_q4k_direct(weights, index, backend, h_post_attn, layer) -} - -/// One-row gated FFN block using direct Q4_K/Q6_K matvec. Mirrors -/// [`crate::ffn::weight::dense_ffn_forward_backend`] but reads gate/up/ -/// down from the vindex slices and avoids the f32 staging. -fn run_ffn_decode_step_q4k_direct( - weights: &ModelWeights, - index: &VectorIndex, - _backend: &dyn ComputeBackend, - h_post_attn: &Array2, - layer: usize, -) -> Option> { - let arch = &*weights.arch; - let hidden = weights.hidden_size; - let intermediate = index.num_features(layer); - let norm_offset = arch.norm_weight_offset(); - - // Pre-FFN norm: same selection logic as `run_ffn` — when the arch - // uses post_norms, the pre-FFN key is `pre_feedforward_layernorm`; - // otherwise it reuses `post_attention_layernorm` as the FFN input - // norm. Falls back to weightless RMS when no key is set. - let pre_ffn_key = if arch.has_post_norms() { - arch.pre_feedforward_layernorm_key(layer) - } else { - Some(arch.post_attention_layernorm_key(layer)) - }; - let h_in = match pre_ffn_key { - Some(key) => apply_norm(weights, h_post_attn, &key, norm_offset), - None => crate::residual::rms_norm(h_post_attn, None, norm_offset), - }; - let h_in_row: &[f32] = h_in.row(0).to_slice().or_else(|| h_in.as_slice())?; - - let ffn = index.interleaved_kquant_layer_data(layer)?; - let (gate_bytes, gate_fmt) = ffn[0]; - let (up_bytes, up_fmt) = ffn[1]; - let (down_bytes, down_fmt) = ffn[2]; - - // Only Gated FFNs reach this path today (it's what predict_kquant_hidden - // currently dequantises). Non-gated archs route through the dequant - // fallback via the per-layer gate at the caller. - if arch.ffn_type() != larql_models::FfnType::Gated { - return None; - } - - // Q8_K-quantise `h_in` once and feed it to both gate and up via the - // sdot-based fused matvec. This is the int8-dot Q4_K × Q8_K path - // that closes the bandwidth gap to llama.cpp on M3 Max. - let mut h_in_q8k = Q8KActivation::with_capacity(hidden); - quantize_x_to_q8k_into(&mut h_in_q8k, h_in_row); - - // Two separate matvecs, each rayon-parallel inside - // `matvec_q4k_or_q6k_q8k`. The "fused gate+up" variant in - // `larql-compute` (`q4k_q8k_gate_up_into`) is single-threaded; - // the input vector (10 KB) stays in L1 across two sequential - // calls anyway, so we don't need explicit fusion to keep `x` - // hot. Splitting lets both matvecs run row-parallel. - let gate_vec = matvec_q4k_or_q6k_q8k(gate_bytes, gate_fmt, &h_in_q8k, intermediate, hidden)?; - let up_vec = matvec_q4k_or_q6k_q8k(up_bytes, up_fmt, &h_in_q8k, intermediate, hidden)?; - - // Element-wise activation: activation(gate) * up. Rayon-chunked — the - // per-element math (libm tanh/exp included) is unchanged, so the output - // is bit-identical to the serial loop; the decode sample showed this - // scalar pass serial on the main thread while the workers slept. - let mut activated = vec![0.0f32; intermediate]; - { - let gelu = matches!(arch.activation(), larql_models::Activation::GeluTanh); - let sqrt_2_over_pi = (2.0f32 / std::f32::consts::PI).sqrt(); - let gate_ref = &gate_vec[..]; - let up_ref = &up_vec[..]; - larql_compute::cpu::spin_pool::par_chunks_mut(&mut activated, 256, |ci, a_c| { - let start = ci * 256; - let g_c = &gate_ref[start..start + a_c.len()]; - let u_c = &up_ref[start..start + a_c.len()]; - if gelu { - for ((a, &x), &u) in a_c.iter_mut().zip(g_c.iter()).zip(u_c.iter()) { - let inner = sqrt_2_over_pi * (x + 0.044715 * x * x * x); - *a = 0.5 * x * (1.0 + inner.tanh()) * u; - } - } else { - // SiLU = x * sigmoid(x). Same shape as dense_ffn_forward_backend. - for ((a, &x), &u) in a_c.iter_mut().zip(g_c.iter()).zip(u_c.iter()) { - let sig = 1.0 / (1.0 + (-x).exp()); - *a = x * sig * u; - } - } - }); - } - - // down projection: out = activated @ W_down.T → [hidden]. - // Re-quantise the post-activation vector (`intermediate`-wide) for - // the down matvec — different input from gate/up. - // - // The stored down row width may be PADDED up to a 256-multiple when - // `intermediate` isn't one (e.g. the 26B-A4B hybrid-MoE dense slab: - // intermediate 2112 stored as 2304-col Q6_K rows). Derive the stored - // width from the byte length and zero-pad the activation to match — - // pad columns multiply zero activations, so the result is exact. - // Bytes per 256-element super-block, from the format helper (= 144 / 210). - let down_sb_bytes = match larql_compute::QuantFormat::from_registry_tag(down_fmt) { - Some(f @ (larql_compute::QuantFormat::Q4_K | larql_compute::QuantFormat::Q6_K)) => { - f.packed_matrix_bytes(1, 256)? - } - _ => return None, - }; - let down_bytes_per_row = down_bytes.len() / hidden; - if down_bytes_per_row == 0 || !down_bytes_per_row.is_multiple_of(down_sb_bytes) { - return None; - } - let stored_cols = - down_bytes_per_row / down_sb_bytes * larql_models::quant::ggml::Q4_K_BLOCK_ELEMS; - if stored_cols < intermediate { - return None; - } - let activated_padded: Vec; - let act_slice: &[f32] = if stored_cols != intermediate { - let mut p = vec![0.0f32; stored_cols]; - p[..intermediate].copy_from_slice(&activated); - activated_padded = p; - &activated_padded - } else { - &activated - }; - let mut activated_q8k = Q8KActivation::with_capacity(stored_cols); - quantize_x_to_q8k_into(&mut activated_q8k, act_slice); - let down_vec = - matvec_q4k_or_q6k_q8k(down_bytes, down_fmt, &activated_q8k, hidden, stored_cols)?; - let mut out = vec_to_2d_row(down_vec); - if let Some(bias) = arch - .ffn_down_bias_key(layer) - .and_then(|k| weights.vectors.get(&k)) - { - add_bias(&mut out, bias); - } - - // Post-FFN residual + optional post-FFN layernorm. Same selection - // logic as `run_ffn`: only fire when has_post_norms() AND the arch - // exposes a post-FFN norm key. - let res_mult = arch.residual_multiplier(); - let h_post_ffn = if arch.has_post_norms() { - let normed = match arch.post_feedforward_layernorm_key(layer) { - Some(key) => apply_norm(weights, &out, &key, norm_offset), - None => crate::residual::rms_norm(&out, None, norm_offset), - }; - if res_mult != 1.0 { - h_post_attn + &(&normed * res_mult) - } else { - h_post_attn + &normed - } - } else if res_mult != 1.0 { - h_post_attn + &(&out * res_mult) - } else { - h_post_attn + &out - }; - - Some(h_post_ffn) + larql_compute::kquant_forward::ffn_decode_step_native( + weights, + index as &dyn larql_compute::KvIndex, + backend, + h_post_attn, + layer, + ) } /// Dequant-free decode step. Same shape contract as @@ -944,59 +281,18 @@ pub fn predict_kquant_decode_step_direct_with_state( backend: &dyn ComputeBackend, cache: &mut CpuKvCache, abs_position: usize, - mut state: Option<&mut crate::PerLayerDecodeState>, + state: Option<&mut crate::PerLayerDecodeState>, ) -> Option> { - use ndarray::s; - let num_layers = weights.num_layers; - if cache.len() != num_layers { - return None; - } - - let mut h = embed_tokens_pub(weights, &[token_id]); - let ple_inputs = precompute_per_layer_inputs(weights, &h, &[token_id]); - - for layer in 0..num_layers { - if let Some(s) = state.as_deref_mut() { - s.h_in_per_layer - .push(larql_compute::state_handle::CpuStateHandle::boxed( - h.clone(), - )); - } - let kv_entry = cache[layer].as_ref(); - let (h_post_attn, new_kv) = attention_decode_step_native( - weights, - index, - backend, - &h, - layer, - kv_entry, - abs_position, - )?; - if let Some(s) = state.as_deref_mut() { - // new_kv is the full prior+new K/V; the new row is the - // last row. Engines that cache per-layer K/V (markov_rs - // hot_kv, turbo_quant compressed) consume this row. - let n = new_kv.0.shape()[0]; - s.k_new_per_layer - .push(larql_compute::state_handle::CpuStateHandle::boxed( - new_kv.0.slice(s![n - 1..n, ..]).to_owned(), - )); - s.v_new_per_layer - .push(larql_compute::state_handle::CpuStateHandle::boxed( - new_kv.1.slice(s![n - 1..n, ..]).to_owned(), - )); - } - cache[layer] = Some(new_kv); - - let h_post_ffn = - run_ffn_decode_step_q4k_direct(weights, index, backend, &h_post_attn, layer)?; - 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); - h = h_out; - } - - Some(h) + // Delegate to the substrate copy in larql-compute (ADR-0022). + larql_compute::kquant_forward::predict_kquant_decode_step_direct_with_state( + weights, + token_id, + index as &dyn larql_compute::KvIndex, + backend, + cache, + abs_position, + state, + ) } #[cfg(test)] @@ -1026,71 +322,8 @@ mod tests { ); } - #[test] - fn layer_supports_direct_matvec_is_true_for_q4k_synthetic() { - let weights = make_test_q4k_weights(); - let index = make_test_q4k_vindex(&weights); - for l in 0..weights.num_layers { - assert!( - layer_supports_direct_matvec(&index, l), - "layer {l} should support direct matvec" - ); - } - } - // ── matvec_q4k_or_q6k_q8k dispatcher ──────────────────────────────── - #[test] - fn matvec_q4k_or_q6k_q8k_q4k_format_produces_finite_output() { - use larql_compute::cpu::ops::q4_common::quantize_q4_k; - use larql_compute::cpu::ops::q4k_q8k_dot::quantize_x_to_q8k_into; - - let rows = 16usize; - let cols = 256usize; // one super-block - let weights: Vec = (0..rows * cols) - .map(|i| ((i as f32) * 0.001).sin() * 0.1) - .collect(); - let q4k_bytes = quantize_q4_k(&weights); - - let x: Vec = (0..cols) - .map(|j| ((j as f32) * 0.007).cos() * 0.5) - .collect(); - let mut x_q8k = Q8KActivation::with_capacity(cols); - quantize_x_to_q8k_into(&mut x_q8k, &x); - - let out = matvec_q4k_or_q6k_q8k(&q4k_bytes, "Q4_K", &x_q8k, rows, cols) - .expect("Q4_K dispatch must succeed"); - assert_eq!(out.len(), rows); - assert!(out.iter().all(|v| v.is_finite())); - assert!( - out.iter().any(|&v| v.abs() > 1e-6), - "matvec should produce non-zero output for non-degenerate input" - ); - } - - #[test] - fn matvec_q4k_or_q6k_q8k_unknown_format_returns_none() { - let x_q8k = Q8KActivation::with_capacity(256); - // Empty bytes + unknown format should drop through to None. - let out = matvec_q4k_or_q6k_q8k(&[], "F32", &x_q8k, 4, 256); - assert!(out.is_none(), "unsupported format must return None"); - } - - #[test] - fn matvec_q4k_or_q6k_q8k_zero_dims_returns_zeroed_vec() { - let x_q8k = Q8KActivation::with_capacity(0); - let out = matvec_q4k_or_q6k_q8k(&[], "Q4_K", &x_q8k, 0, 0); - assert_eq!(out.as_deref(), Some(&[][..])); - } - - #[test] - fn matvec_q4k_or_q6k_q8k_non_multiple_cols_returns_none() { - let x_q8k = Q8KActivation::with_capacity(255); - // 255 is not a multiple of 256 → reject (caller falls back to dequant path). - let out = matvec_q4k_or_q6k_q8k(&[], "Q4_K", &x_q8k, 4, 255); - assert!(out.is_none(), "non-256-multiple cols must be rejected"); - } - // ── predict_kquant_prefill / predict_kquant_decode_step ──────────────────── #[test] @@ -1109,9 +342,9 @@ mod tests { "hidden state must be finite" ); assert_eq!(cache.len(), fx.weights.num_layers); - for layer in 0..fx.weights.num_layers { + for entry in &cache { assert!( - cache[layer].is_some(), + entry.is_some(), "every layer should have K/V populated after prefill" ); } @@ -1129,7 +362,7 @@ mod tests { .collect(); let (h_new, _step_timings) = - predict_kquant_decode_step(&mut fx.weights, 4, &fx.index, &mut cache, token_ids.len()) + predict_kquant_decode_step(&fx.weights, 4, &fx.index, &mut cache, token_ids.len()) .expect("decode step must succeed on a populated cache"); assert_eq!(h_new.shape(), &[1, fx.weights.hidden_size]); @@ -1146,10 +379,10 @@ mod tests { #[test] fn predict_kquant_decode_step_rejects_mismatched_cache_length() { - let mut fx = Q4KTestFixtures::build(); + let fx = Q4KTestFixtures::build(); // Cache length doesn't match num_layers — function must return None. let mut bad_cache: CpuKvCache = vec![None; fx.weights.num_layers + 1]; - let result = predict_kquant_decode_step(&mut fx.weights, 1, &fx.index, &mut bad_cache, 0); + let result = predict_kquant_decode_step(&fx.weights, 1, &fx.index, &mut bad_cache, 0); assert!(result.is_none()); } @@ -1168,7 +401,7 @@ mod tests { let (_, mut cache_a, _) = predict_kquant_prefill(&mut fx_a.weights, &token_ids, &fx_a.index); let (h_staged, _) = - predict_kquant_decode_step(&mut fx_a.weights, 4, &fx_a.index, &mut cache_a, 3) + predict_kquant_decode_step(&fx_a.weights, 4, &fx_a.index, &mut cache_a, 3) .expect("staged step"); let mut fx_b = Q4KTestFixtures::build(); @@ -1238,7 +471,7 @@ mod tests { let (_, mut cache_a, _) = predict_kquant_prefill(&mut weights_a, &token_ids, &index); let (h_staged, _) = - predict_kquant_decode_step(&mut weights_a, next, &index, &mut cache_a, token_ids.len()) + predict_kquant_decode_step(&weights_a, next, &index, &mut cache_a, token_ids.len()) .expect("staged step"); let mut weights_b = make_test_q4k_weights_rope_scaled(); @@ -1382,12 +615,10 @@ mod branch_tests { make_test_q4k_vindex, make_test_q4k_weights, make_test_tokenizer, Q4K_TEST_HIDDEN, Q4K_TEST_INTER, Q4K_TEST_NUM_LAYERS, Q4K_TEST_VOCAB, }; - use larql_compute::cpu::ops::q4_common::quantize_q4_k; use larql_compute::CpuBackend; use larql_models::{detect_from_json, ModelWeights, WeightArray}; use ndarray::Array2; use std::collections::HashMap; - use std::sync::Arc; fn rand_mat(rows: usize, cols: usize, seed: u64) -> WeightArray { let mut state = seed; @@ -1531,101 +762,11 @@ mod branch_tests { let token_ids = vec![1u32, 2]; let (_, mut cache, _) = predict_kquant_prefill(&mut weights, &token_ids, &index); let (h_new, _) = - predict_kquant_decode_step(&mut weights, 3, &index, &mut cache, token_ids.len()) + predict_kquant_decode_step(&weights, 3, &index, &mut cache, token_ids.len()) .expect("SiLU dequant decode step must succeed"); assert!(h_new.iter().all(|v| v.is_finite())); } - /// Q4K format check — when a matrix's bytes are too short for the - /// claimed (rows, cols), the dispatcher should return None instead - /// of panicking on the out-of-range slice. - #[test] - fn matvec_q4k_or_q6k_q8k_short_bytes_returns_none() { - let x_q8k = Q8KActivation::with_capacity(256); - // Need at least 4 * (256/256) * 144 = 576 bytes for 4×256 Q4_K. - let too_short = vec![0u8; 144]; // only 1 row's worth - let out = matvec_q4k_or_q6k_q8k(&too_short, "Q4_K", &x_q8k, 4, 256); - assert!(out.is_none(), "short byte buffer must be rejected"); - } - - /// Cover the no-attention-data branch of - /// `layer_supports_direct_matvec`: a vindex with no attn_kquant_layer_data. - #[test] - fn layer_supports_direct_matvec_false_without_attn_data() { - let weights = make_test_q4k_weights(); - // Build an index with no Q4_K storage at all. - let bare = larql_vindex::VectorIndex::new( - vec![None; weights.num_layers], - vec![None; weights.num_layers], - weights.num_layers, - weights.hidden_size, - ); - assert!( - !supports_direct_matvec_decode(&weights, &bare), - "vindex with no Q4_K data can't support direct matvec" - ); - for l in 0..weights.num_layers { - assert!(!layer_supports_direct_matvec(&bare, l)); - } - } - - /// Cover the unsupported-format branch of - /// `layer_supports_direct_matvec`: vindex carries Q4_K storage but - /// one entry is tagged with a format the direct path doesn't - /// recognise (e.g., `Q4_KF`). - #[test] - fn layer_supports_direct_matvec_false_with_unsupported_format() { - let weights = make_test_q4k_weights(); - let arch = &*weights.arch; - let mut attn_payload = Vec::new(); - let mut attn_manifest: Vec<(usize, usize, String)> = Vec::new(); - let q_dim = weights.num_q_heads * weights.head_dim; - let kv_dim = weights.num_kv_heads * weights.head_dim; - for layer in 0..weights.num_layers { - for (idx, key) in [ - arch.attn_q_key(layer), - arch.attn_k_key(layer), - arch.attn_v_key(layer), - arch.attn_o_key(layer), - ] - .iter() - .enumerate() - { - let rows = if idx == 0 || idx == 3 { q_dim } else { kv_dim }; - let _ = rows; - let tensor = weights.tensors.get(key).unwrap(); - let bytes = quantize_q4_k(tensor.as_slice().unwrap()); - let off = attn_payload.len(); - let len = bytes.len(); - attn_payload.extend_from_slice(&bytes); - // Tag the first Q with an unsupported format. - let fmt = if layer == 0 && idx == 0 { - "Q4_KF".to_string() - } else { - "Q4_K".to_string() - }; - attn_manifest.push((off, len, fmt)); - } - } - // Also need interleaved Q4_K so the FFN branch passes through - // before we hit the attn check. Use the synth helper for that. - let mut index = larql_vindex::VectorIndex::new( - vec![None; weights.num_layers], - vec![None; weights.num_layers], - weights.num_layers, - weights.hidden_size, - ); - index.vocab_size = weights.vocab_size; - let mut anon = memmap2::MmapMut::map_anon(attn_payload.len()).expect("anon"); - anon.copy_from_slice(&attn_payload); - let mmap = Arc::new(anon.make_read_only().unwrap()); - Arc::make_mut(&mut index.storage).set_attn_kquant(mmap, Some(attn_manifest)); - assert!( - !layer_supports_direct_matvec(&index, 0), - "layer with Q4_KF format must not support the direct matvec path" - ); - } - /// `CachedTimings::merge` is private; verify the public `add` /// wrapper covers it (both should sum into `dequant_ms`). #[test] diff --git a/crates/larql-inference/src/vindex/kquant_forward/hidden.rs b/crates/larql-inference/src/vindex/kquant_forward/hidden.rs index 10fabda93..40744bd95 100644 --- a/crates/larql-inference/src/vindex/kquant_forward/hidden.rs +++ b/crates/larql-inference/src/vindex/kquant_forward/hidden.rs @@ -15,12 +15,13 @@ use super::tensors::{insert_q4k_layer_tensors, remove_layer_tensors}; /// vindex, dequantising attn + FFN one layer at a time. Returns the /// `[seq_len, hidden]` array; caller owns the lm_head step. pub fn predict_kquant_hidden( - weights: &mut ModelWeights, + weights: &ModelWeights, token_ids: &[u32], index: &VectorIndex, moe_remote: Option<&crate::ffn::RemoteMoeBackend>, ) -> Array2 { let num_layers = weights.num_layers; + let mut scratch = larql_models::DequantScratch::new(); let mut h = embed_tokens_pub(weights, token_ids); let ple_inputs = precompute_per_layer_inputs(weights, &h, token_ids); @@ -34,18 +35,19 @@ pub fn predict_kquant_hidden( } for layer in 0..num_layers { - let inserted = - insert_q4k_layer_tensors(weights, index, layer).unwrap_or_else(|err| panic!("{err}")); + let inserted = insert_q4k_layer_tensors(&mut scratch, weights, index, layer) + .unwrap_or_else(|err| panic!("{err}")); let shared_kv = weights .arch .kv_shared_source_layer(layer) .and_then(|src| kv_cache.get(&src)); let is_moe_layer = weights.arch.is_hybrid_moe(); - let ffn_backend = crate::ffn::WeightFfn { weights }; + let view = larql_models::WeightsView::with_scratch(weights, &scratch); + let ffn_backend = crate::ffn::ViewFfn { view }; if is_moe_layer { if let Some((h_new, kv_out)) = run_moe_layer_cpu( - weights, + view, &h, layer, &ffn_backend, @@ -59,7 +61,7 @@ pub fn predict_kquant_hidden( } } } else if let Some((h_new, _, kv_out)) = run_layer_with_ffn( - weights, + view, &h, layer, &ffn_backend, @@ -73,7 +75,7 @@ pub fn predict_kquant_hidden( } } - remove_layer_tensors(weights, inserted); + remove_layer_tensors(&mut scratch, inserted); if let Some(dir) = dump_dir { let slice = h.as_slice().unwrap_or(&[]); @@ -117,7 +119,7 @@ fn build_moe_router_weights<'a>( /// CPU forward for one hybrid-MoE layer (Gemma 4 26B A4B). fn run_moe_layer_cpu( - weights: &ModelWeights, + weights: larql_models::WeightsView, h: &Array2, layer: usize, ffn: &dyn crate::ffn::FfnBackend, @@ -135,7 +137,14 @@ fn run_moe_layer_cpu( (h_pa, Some((k_rope, v_final))) }; - let h_out = moe_ffn_block_cpu(weights, &h_post_attn, layer, ffn, ple_input, moe_remote); + let h_out = moe_ffn_block_cpu( + weights.canonical(), + &h_post_attn, + layer, + ffn, + ple_input, + moe_remote, + ); Some((h_out, kv_out)) } @@ -342,9 +351,9 @@ mod tests { /// hybrid-MoE arch fixture; this test covers the rest. #[test] fn predict_kquant_hidden_returns_shape_and_finite() { - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); - let h = predict_kquant_hidden(&mut weights, &[0u32, 1, 2], &index, None); + let h = predict_kquant_hidden(&weights, &[0u32, 1, 2], &index, None); assert_eq!(h.shape(), &[3, weights.hidden_size]); assert!( h.iter().all(|v| v.is_finite()), @@ -354,9 +363,9 @@ mod tests { #[test] fn predict_kquant_hidden_single_token() { - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); - let h = predict_kquant_hidden(&mut weights, &[5u32], &index, None); + let h = predict_kquant_hidden(&weights, &[5u32], &index, None); assert_eq!(h.shape(), &[1, weights.hidden_size]); } @@ -393,9 +402,9 @@ mod tests { #[test] fn predict_kquant_hidden_routes_through_moe_branch_on_gemma4_fixture() { use crate::test_utils::{make_test_gemma4_moe_weights, make_test_q4k_vindex}; - let mut weights = make_test_gemma4_moe_weights(); + let weights = make_test_gemma4_moe_weights(); let index = make_test_q4k_vindex(&weights); - let h = predict_kquant_hidden(&mut weights, &[0u32, 1], &index, None); + let h = predict_kquant_hidden(&weights, &[0u32, 1], &index, None); assert_eq!(h.shape(), &[2, weights.hidden_size]); assert!( h.iter().all(|v| v.is_finite()), @@ -414,7 +423,7 @@ mod tests { let mut weights = make_test_gemma4_moe_weights(); let index = make_test_q4k_vindex(&weights); weights.raw_bytes.clear(); // drop per-expert blobs → build_moe_weights None - let h = predict_kquant_hidden(&mut weights, &[0u32, 1], &index, None); + let h = predict_kquant_hidden(&weights, &[0u32, 1], &index, None); assert_eq!(h.shape(), &[2, weights.hidden_size]); assert!( h.iter().all(|v| v.is_finite()), @@ -444,20 +453,20 @@ mod tests { } } - let mut weights = make_test_gemma4_moe_weights(); + let weights = make_test_gemma4_moe_weights(); let index = make_test_q4k_vindex(&weights); let nl = weights.num_layers; let _reset = RoutingReset; set_routing(None); - let dense = predict_kquant_hidden(&mut weights, &[0u32, 1], &index, None); + let dense = predict_kquant_hidden(&weights, &[0u32, 1], &index, None); // Aggressively prune every expert layer's feature set. set_routing(Some(WithinExpertRouting { frac_per_layer: vec![Some(0.125); nl], selector: ExpertFeatureSelector::ActMagnitude, })); - let pruned = predict_kquant_hidden(&mut weights, &[0u32, 1], &index, None); + let pruned = predict_kquant_hidden(&weights, &[0u32, 1], &index, None); drop(_reset); // restore before asserting assert_eq!(pruned.shape(), dense.shape()); @@ -483,10 +492,10 @@ mod tests { fn predict_kquant_hidden_with_disconnected_remote_moe_backend_falls_back() { use crate::ffn::RemoteMoeBackend; use crate::test_utils::{make_test_gemma4_moe_weights, make_test_q4k_vindex}; - let mut weights = make_test_gemma4_moe_weights(); + let weights = make_test_gemma4_moe_weights(); let index = make_test_q4k_vindex(&weights); let remote = RemoteMoeBackend::new_disconnected(); - let h = predict_kquant_hidden(&mut weights, &[0u32, 1], &index, Some(&remote)); + let h = predict_kquant_hidden(&weights, &[0u32, 1], &index, Some(&remote)); assert_eq!(h.shape(), &[2, weights.hidden_size]); assert!( h.iter().all(|v| v.is_finite()), @@ -529,9 +538,9 @@ mod tests { ); use crate::test_utils::{make_test_gemma4_moe_weights, make_test_q4k_vindex}; - let mut weights = make_test_gemma4_moe_weights(); + let weights = make_test_gemma4_moe_weights(); let index = make_test_q4k_vindex(&weights); - let h = predict_kquant_hidden(&mut weights, &[0u32, 1], &index, None); + let h = predict_kquant_hidden(&weights, &[0u32, 1], &index, None); assert_eq!(h.shape(), &[2, weights.hidden_size]); // Embed dump must exist (written at line 33 unconditionally when diff --git a/crates/larql-inference/src/vindex/kquant_forward/interventions.rs b/crates/larql-inference/src/vindex/kquant_forward/interventions.rs index 09a6dbb7e..da34baef0 100644 --- a/crates/larql-inference/src/vindex/kquant_forward/interventions.rs +++ b/crates/larql-inference/src/vindex/kquant_forward/interventions.rs @@ -47,21 +47,23 @@ where )); } + let mut scratch = larql_models::DequantScratch::new(); let mut h = embed_tokens_pub(weights, token_ids); let ple_inputs = precompute_per_layer_inputs(weights, &h, token_ids); let mut kv_cache: HashMap = HashMap::new(); for layer in 0..weights.num_layers { - let inserted = insert_q4k_layer_tensors(weights, index, layer)?; + let inserted = insert_q4k_layer_tensors(&mut scratch, weights, index, layer)?; let shared_kv = weights .arch .kv_shared_source_layer(layer) .and_then(|src| kv_cache.get(&src)); - let ffn_backend = crate::ffn::WeightFfn { weights }; + let view = larql_models::WeightsView::with_scratch(weights, &scratch); + let ffn_backend = crate::ffn::ViewFfn { view }; let step = if layer == target_layer { run_target_layer( - weights, + view.canonical(), &h, layer, &ffn_backend, @@ -70,7 +72,7 @@ where )? } else { run_layer_with_ffn( - weights, + view, &h, layer, &ffn_backend, @@ -82,14 +84,14 @@ where }; let Some((h_new, kv_out)) = step else { - remove_layer_tensors(weights, inserted); + remove_layer_tensors(&mut scratch, inserted); return Err(format!("{label} failed at layer {layer}")); }; h = h_new; if let Some(kv) = kv_out { kv_cache.insert(layer, kv); } - remove_layer_tensors(weights, inserted); + remove_layer_tensors(&mut scratch, inserted); } Ok(h) diff --git a/crates/larql-inference/src/vindex/kquant_forward/mod.rs b/crates/larql-inference/src/vindex/kquant_forward/mod.rs index 8b786271f..c05f3b4b4 100644 --- a/crates/larql-inference/src/vindex/kquant_forward/mod.rs +++ b/crates/larql-inference/src/vindex/kquant_forward/mod.rs @@ -51,5 +51,8 @@ pub use metal::{ pub use remote_ffn::{ predict_kquant_hidden_with_ffn, predict_kquant_with_ffn, predict_kquant_with_ffn_early_exit, }; -pub use tensors::{insert_q4k_layer_tensors, remove_layer_tensors}; +pub use tensors::{ + insert_q4k_layer_tensors, insert_q4k_layer_tensors_resident, remove_layer_tensors, + remove_layer_tensors_resident, +}; pub use walk_ffn::{kquant_ffn_forward_layer, kquant_ffn_forward_layer_q8k}; 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 8735c78ad..4619a3933 100644 --- a/crates/larql-inference/src/vindex/kquant_forward/remote_ffn.rs +++ b/crates/larql-inference/src/vindex/kquant_forward/remote_ffn.rs @@ -132,7 +132,11 @@ fn predict_kquant_hidden_inner( // backend (attention already done locally; server handles dense-FFN + // expert dispatch + combine). Fall through to dense-only on None. if weights.arch.is_hybrid_moe() { - if let Some(h_post_attn) = crate::forward::run_attention_public(weights, &h, layer) { + if let Some(h_post_attn) = crate::forward::run_attention_public( + larql_models::WeightsView::dense(weights), + &h, + layer, + ) { if let Some(h_out) = ffn_backend.forward_moe_full_layer(layer, &h_post_attn) { h = h_out; weights.tensors.remove(&q_key); @@ -149,7 +153,7 @@ fn predict_kquant_hidden_inner( .kv_shared_source_layer(layer) .and_then(|src| kv_cache.get(&src)); if let Some((h_new, _, kv_out)) = run_layer_with_ffn( - weights, + larql_models::WeightsView::dense(weights), &h, layer, ffn_backend, diff --git a/crates/larql-inference/src/vindex/kquant_forward/tensors.rs b/crates/larql-inference/src/vindex/kquant_forward/tensors.rs index af68dcb54..0533d4ecd 100644 --- a/crates/larql-inference/src/vindex/kquant_forward/tensors.rs +++ b/crates/larql-inference/src/vindex/kquant_forward/tensors.rs @@ -1,16 +1,20 @@ -use larql_models::ModelWeights; +use larql_models::{DequantScratch, ModelWeights}; use larql_vindex::VectorIndex; use super::dequant::dequantize_matrix; /// Insert one Q4_K/Q6_K vindex layer's attention and dense FFN tensors into -/// `weights.tensors` as dense f32 matrices. +/// the engine-owned `scratch` as dense f32 matrices (`weights` stays +/// immutable; readers resolve them via [`WeightsView::with_scratch`]). +/// +/// [`WeightsView::with_scratch`]: larql_models::WeightsView::with_scratch /// /// This is the shared research/intervention primitive behind Q4K CPU forward /// and OV/RD-style experiments. Call [`remove_layer_tensors`] with the returned /// keys after the layer has run to keep peak f32 memory bounded. pub fn insert_q4k_layer_tensors( - weights: &mut ModelWeights, + scratch: &mut DequantScratch, + weights: &ModelWeights, index: &VectorIndex, layer: usize, ) -> Result, String> { @@ -38,27 +42,27 @@ pub fn insert_q4k_layer_tensors( let up_key = arch.ffn_up_key(layer); let down_key = arch.ffn_down_key(layer); - weights.tensors.insert( + scratch.insert( q_key.clone(), dequantize_matrix(attn[0].0, attn[0].1, q_dim, hidden).into_shared(), ); - weights.tensors.insert( + scratch.insert( k_key.clone(), dequantize_matrix(attn[1].0, attn[1].1, kv_dim, hidden).into_shared(), ); - weights.tensors.insert( + scratch.insert( v_key.clone(), dequantize_matrix(attn[2].0, attn[2].1, kv_dim, hidden).into_shared(), ); - weights.tensors.insert( + scratch.insert( o_key.clone(), dequantize_matrix(attn[3].0, attn[3].1, hidden, q_dim).into_shared(), ); - weights.tensors.insert( + scratch.insert( gate_key.clone(), dequantize_matrix(ffn[0].0, ffn[0].1, intermediate, hidden).into_shared(), ); - weights.tensors.insert( + scratch.insert( up_key.clone(), dequantize_matrix(ffn[1].0, ffn[1].1, intermediate, hidden).into_shared(), ); @@ -71,15 +75,43 @@ pub fn insert_q4k_layer_tensors( } else { dequantize_matrix(ffn[2].0, ffn[2].1, hidden, intermediate) }; - weights - .tensors - .insert(down_key.clone(), w_down.into_shared()); + scratch.insert(down_key.clone(), w_down.into_shared()); Ok(vec![q_key, k_key, v_key, o_key, gate_key, up_key, down_key]) } /// Remove tensor keys previously returned by [`insert_q4k_layer_tensors`]. -pub fn remove_layer_tensors(weights: &mut ModelWeights, keys: Vec) { +pub fn remove_layer_tensors(scratch: &mut DequantScratch, keys: Vec) { + for key in keys { + scratch.remove(&key); + } +} + +/// Resident variant of [`insert_q4k_layer_tensors`] — dequantises **into +/// `weights.tensors`** (mutating `weights`) rather than an engine scratch. +/// +/// For dev/research drivers (the `ov_rd` CLI tooling, the relation resolver, +/// standalone examples) that run their own per-layer forward against canonical +/// `weights.tensors` and don't thread a [`WeightsView`]. The production decode +/// path uses the scratch-based [`insert_q4k_layer_tensors`] + +/// [`WeightsView::with_scratch`] so `ModelWeights` stays immutable/`Arc`-able. +/// +/// [`WeightsView`]: larql_models::WeightsView +/// [`WeightsView::with_scratch`]: larql_models::WeightsView::with_scratch +pub fn insert_q4k_layer_tensors_resident( + weights: &mut ModelWeights, + index: &VectorIndex, + layer: usize, +) -> Result, String> { + let mut scratch = DequantScratch::new(); + let inserted = insert_q4k_layer_tensors(&mut scratch, weights, index, layer)?; + weights.tensors.extend(scratch); + Ok(inserted) +} + +/// Resident variant of [`remove_layer_tensors`] — removes the keys from +/// `weights.tensors` (the [`insert_q4k_layer_tensors_resident`] counterpart). +pub fn remove_layer_tensors_resident(weights: &mut ModelWeights, keys: Vec) { for key in keys { weights.tensors.remove(&key); } diff --git a/crates/larql-inference/src/vindex/mod.rs b/crates/larql-inference/src/vindex/mod.rs index 304cea5c6..4d2e3465f 100644 --- a/crates/larql-inference/src/vindex/mod.rs +++ b/crates/larql-inference/src/vindex/mod.rs @@ -11,7 +11,7 @@ mod loader; mod walk_config; mod walk_ffn; -pub use dequant::ensure_attn_tensors_dequantised; +pub use dequant::{ensure_attn_tensors_dequantised, ensure_attn_tensors_dequantised_resident}; pub(crate) use kquant_forward::generate_kquant_cpu_constrained_streaming_sampled_with_eos; pub use kquant_forward::{ attention_decode_step_native, ffn_decode_step_native, fused_decode_step, @@ -19,11 +19,11 @@ pub use kquant_forward::{ generate_kquant_cpu_constrained, generate_kquant_cpu_constrained_cached, generate_kquant_cpu_constrained_cached_streaming, generate_kquant_cpu_constrained_streaming, generate_kquant_cpu_constrained_streaming_sampled, generate_kquant_cpu_remote, - insert_q4k_layer_tensors, is_end_of_turn, kquant_ffn_forward_layer, - kquant_ffn_forward_layer_q8k, moe_ffn_block_cpu, moe_ffn_block_cpu_with_index, predict_kquant, - predict_kquant_decode_step, predict_kquant_decode_step_direct, - predict_kquant_decode_step_direct_with_state, predict_kquant_hidden, - predict_kquant_hidden_hooked, predict_kquant_hidden_with_ffn, + insert_q4k_layer_tensors, insert_q4k_layer_tensors_resident, is_end_of_turn, + kquant_ffn_forward_layer, kquant_ffn_forward_layer_q8k, moe_ffn_block_cpu, + moe_ffn_block_cpu_with_index, predict_kquant, predict_kquant_decode_step, + predict_kquant_decode_step_direct, predict_kquant_decode_step_direct_with_state, + predict_kquant_hidden, predict_kquant_hidden_hooked, predict_kquant_hidden_with_ffn, predict_kquant_hidden_with_mapped_head_residual_delta, predict_kquant_hidden_with_mapped_pre_o_head, predict_kquant_hidden_with_original_head_residual_delta, @@ -34,8 +34,8 @@ pub use kquant_forward::{ predict_kquant_metal_capture_pre_wo, predict_kquant_metal_hidden, predict_kquant_metal_with_replaced_head_residual_delta, predict_kquant_prefill, predict_kquant_prefill_with_state, predict_kquant_with_ffn, predict_kquant_with_ffn_early_exit, - remove_layer_tensors, supports_cached_decode, supports_direct_matvec_decode, CachedTimings, - CpuKvCache, + remove_layer_tensors, remove_layer_tensors_resident, supports_cached_decode, + supports_direct_matvec_decode, CachedTimings, CpuKvCache, }; pub use l1_cache::FfnL1Cache; pub use loader::{open_inference_vindex, ENV_VINDEX_PATH}; diff --git a/crates/larql-inference/tests/test_layer_graph_integration.rs b/crates/larql-inference/tests/test_layer_graph_integration.rs index 2b18d23a1..dfbc458cb 100644 --- a/crates/larql-inference/tests/test_layer_graph_integration.rs +++ b/crates/larql-inference/tests/test_layer_graph_integration.rs @@ -158,7 +158,7 @@ fn prefill_with_kv_matches_predict_q4k_hidden() { ); // predict_kquant_hidden dequantises layer-by-layer - let h_q4k = predict_kquant_hidden(&mut weights, &prompt_ids, &index, None); + let h_q4k = predict_kquant_hidden(&weights, &prompt_ids, &index, None); // The two paths use different FFN implementations — cosine similarity should // be > 0.95 at the last position (they differ mainly in FFN quantisation). diff --git a/crates/larql-inference/tests/test_q4k_cached_parity.rs b/crates/larql-inference/tests/test_q4k_cached_parity.rs index ac0b61e60..4880c0aed 100644 --- a/crates/larql-inference/tests/test_q4k_cached_parity.rs +++ b/crates/larql-inference/tests/test_q4k_cached_parity.rs @@ -79,7 +79,7 @@ fn cached_decode_matches_uncached_tokens() { let mut cb = SilentLoadCallbacks; let mut weights_a = load_model_weights_kquant(&vindex_path, &mut cb).expect("load weights A"); - let mut weights_b = load_model_weights_kquant(&vindex_path, &mut cb).expect("load weights B"); + let weights_b = load_model_weights_kquant(&vindex_path, &mut cb).expect("load weights B"); let mut index = VectorIndex::load_vindex(&vindex_path, &mut cb).expect("load index"); index.load_attn_kquant(&vindex_path).expect("load attn Q4K"); index @@ -107,7 +107,7 @@ fn cached_decode_matches_uncached_tokens() { for step in 1..STEPS { let abs_position = prompt_ids.len() + (step - 1); let (h_new, _) = - predict_kquant_decode_step(&mut weights_a, next_id, &index, &mut cache, abs_position) + predict_kquant_decode_step(&weights_a, next_id, &index, &mut cache, abs_position) .expect("cached decode step"); next_id = argmax_token(&weights_a, &tokenizer, &h_new); cached_ids.push(next_id); @@ -115,12 +115,12 @@ fn cached_decode_matches_uncached_tokens() { // ── Path B: uncached predict_kquant_hidden per step ────────────── let mut ids = prompt_ids.clone(); - let h_full = predict_kquant_hidden(&mut weights_b, &ids, &index, None); + let h_full = predict_kquant_hidden(&weights_b, &ids, &index, None); let mut next_id = argmax_token(&weights_b, &tokenizer, &h_full); let mut uncached_ids = vec![next_id]; ids.push(next_id); for _ in 1..STEPS { - let h_full = predict_kquant_hidden(&mut weights_b, &ids, &index, None); + let h_full = predict_kquant_hidden(&weights_b, &ids, &index, None); next_id = argmax_token(&weights_b, &tokenizer, &h_full); uncached_ids.push(next_id); ids.push(next_id); @@ -192,7 +192,7 @@ fn direct_matvec_decode_matches_dequant_path() { for step in 1..STEPS { let abs_position = prompt_ids.len() + (step - 1); let (h_new, _) = - predict_kquant_decode_step(&mut weights_b, next_id, &index, &mut cache_b, abs_position) + predict_kquant_decode_step(&weights_b, next_id, &index, &mut cache_b, abs_position) .expect("dequant decode step"); next_id = argmax_token(&weights_b, &tokenizer, &h_new); dequant_ids.push(next_id); @@ -245,14 +245,9 @@ fn direct_matvec_decode_matches_dequant_path() { prompt_ids.len(), ) .expect("direct step"); - let (h_dequant, _) = predict_kquant_decode_step( - &mut weights_b, - first, - &index, - &mut cache_b, - prompt_ids.len(), - ) - .expect("dequant step"); + let (h_dequant, _) = + predict_kquant_decode_step(&weights_b, first, &index, &mut cache_b, prompt_ids.len()) + .expect("dequant step"); let a = h_direct.row(0); let b = h_dequant.row(0); let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum(); diff --git a/crates/larql-kv/benches/engine_decode.rs b/crates/larql-kv/benches/engine_decode.rs index bb71c4cc6..25d1fe1ca 100644 --- a/crates/larql-kv/benches/engine_decode.rs +++ b/crates/larql-kv/benches/engine_decode.rs @@ -78,12 +78,24 @@ fn bench_decode_step(c: &mut Criterion) { let mut group = c.benchmark_group("decode_step"); for (name, kind) in all_engines() { group.bench_function(name, |b| { - // Pre-warm: prefill once, then time a single decode_step. - let mut engine = kind.clone().build(cpu_engine_backend()); - let _ = engine.prefill(&weights, &ffn, &prompt); - b.iter(|| { - let _ = engine.decode_step(&weights, &ffn, 1); - }); + // Measure ONE decode step at a fixed (prompt-length) context. A + // fresh prefilled engine per timed call keeps the K/V cache from + // growing across iterations — the unbounded `standard` engine + // otherwise appends ~N positions over criterion's N iterations, + // making the per-call cost non-stationary (and the result a + // function of iteration count, not single-step latency). Setup is + // untimed; only `decode_step` is measured. + b.iter_batched_ref( + || { + let mut engine = kind.clone().build(cpu_engine_backend()); + let _ = engine.prefill(&weights, &ffn, &prompt); + engine + }, + |engine| { + let _ = engine.decode_step(&weights, &ffn, 1); + }, + criterion::BatchSize::SmallInput, + ); }); } group.finish(); @@ -182,36 +194,73 @@ fn bench_helpers_sync_vs_async(c: &mut Criterion) { group.bench_function("prefill_sync", |b| { b.iter(|| { - let _ = kv_prefill_via_dispatch(&cpu, &weights, &ffn, &prompt, None, None).unwrap(); + let _ = kv_prefill_via_dispatch( + &cpu, + larql_inference::WeightsView::dense(&weights), + &ffn, + &prompt, + None, + None, + ) + .unwrap(); }); }); group.bench_function("prefill_async", |b| { b.iter(|| { - let _ = - kv_prefill_via_dispatch_async(&cpu, &weights, &ffn, &prompt, None, None).unwrap(); + let _ = kv_prefill_via_dispatch_async( + &cpu, + larql_inference::WeightsView::dense(&weights), + &ffn, + &prompt, + None, + None, + ) + .unwrap(); }); }); group.bench_function("decode_step_sync", |b| { - let (_, mut handles) = - kv_prefill_via_dispatch(&cpu, &weights, &ffn, &prompt, None, None).unwrap(); + let (_, mut handles) = kv_prefill_via_dispatch( + &cpu, + larql_inference::WeightsView::dense(&weights), + &ffn, + &prompt, + None, + None, + ) + .unwrap(); let mut pos = prompt.len(); b.iter(|| { - let _ = - kv_decode_step_via_dispatch(&cpu, &weights, &ffn, &mut handles, 1, pos, None, None); + let _ = kv_decode_step_via_dispatch( + &cpu, + larql_inference::WeightsView::dense(&weights), + &ffn, + &mut handles, + 1, + pos, + None, + None, + ); pos += 1; }); }); group.bench_function("decode_step_async", |b| { - let (_, mut handles) = - kv_prefill_via_dispatch_async(&cpu, &weights, &ffn, &prompt, None, None).unwrap(); + let (_, mut handles) = kv_prefill_via_dispatch_async( + &cpu, + larql_inference::WeightsView::dense(&weights), + &ffn, + &prompt, + None, + None, + ) + .unwrap(); let mut pos = prompt.len(); b.iter(|| { let _ = kv_decode_step_via_dispatch_async( &cpu, - &weights, + larql_inference::WeightsView::dense(&weights), &ffn, &mut handles, 1, diff --git a/crates/larql-kv/examples/contract_classify_cached_ffn.rs b/crates/larql-kv/examples/contract_classify_cached_ffn.rs index 8e669b03b..f8f6eb50e 100644 --- a/crates/larql-kv/examples/contract_classify_cached_ffn.rs +++ b/crates/larql-kv/examples/contract_classify_cached_ffn.rs @@ -163,8 +163,15 @@ fn forward_no_cache( .kv_shared_source_layer(layer) .and_then(|src| kv_cache.get(&src)); let walk_ffn = WalkFfn::new_unlimited(weights, index); - let (h_new, _, kv_out) = - run_layer_with_ffn(weights, &h, layer, &walk_ffn, false, None, shared_kv)?; + let (h_new, _, kv_out) = run_layer_with_ffn( + larql_inference::WeightsView::dense(weights), + &h, + layer, + &walk_ffn, + false, + None, + shared_kv, + )?; h = h_new; if let Some(kv) = kv_out { kv_cache.insert(layer, kv); @@ -197,8 +204,15 @@ fn forward_with_cache( } else { // Cache miss — real compute for this layer. let walk_ffn = WalkFfn::new_unlimited(weights, index); - let (h_new, _, kv_out) = - run_layer_with_ffn(weights, &h, layer, &walk_ffn, false, None, None)?; + let (h_new, _, kv_out) = run_layer_with_ffn( + larql_inference::WeightsView::dense(weights), + &h, + layer, + &walk_ffn, + false, + None, + None, + )?; h = h_new; if let Some(kv) = kv_out { kv_cache.insert(layer, kv); @@ -211,8 +225,15 @@ fn forward_with_cache( .kv_shared_source_layer(layer) .and_then(|src| kv_cache.get(&src)); let walk_ffn = WalkFfn::new_unlimited(weights, index); - let (h_new, _, kv_out) = - run_layer_with_ffn(weights, &h, layer, &walk_ffn, false, None, shared_kv)?; + let (h_new, _, kv_out) = run_layer_with_ffn( + larql_inference::WeightsView::dense(weights), + &h, + layer, + &walk_ffn, + false, + None, + shared_kv, + )?; h = h_new; if let Some(kv) = kv_out { kv_cache.insert(layer, kv); diff --git a/crates/larql-kv/src/engines/apollo/engine.rs b/crates/larql-kv/src/engines/apollo/engine.rs index 5f62659bd..6951e7e28 100644 --- a/crates/larql-kv/src/engines/apollo/engine.rs +++ b/crates/larql-kv/src/engines/apollo/engine.rs @@ -74,6 +74,12 @@ pub struct ApolloEngine { /// running all 34 layers — ~8.5× faster on Gemma 3 4B (crystal_layer=30 → 4 layers). pub(super) boundary_residual: Option>, pub(super) crystal_layer: usize, + /// Engine-owned f32 dequant scratch for the Q4K path — `prefill_quant`/ + /// `decode_step_quant` dequantise attn+FFN into it; `prefill`/`decode_step` + /// resolve `forward_raw_logits`/`forward_from_layer` through a + /// `WeightsView::with_scratch` over it. Empty on the dense path. Keeps + /// `weights` immutable (no `weights.tensors` mutation). + pub(super) dequant_scratch: larql_inference::DequantScratch, } impl ApolloEngine { @@ -86,6 +92,7 @@ impl ApolloEngine { injection_delta: None, boundary_residual: None, crystal_layer: 0, + dequant_scratch: larql_inference::DequantScratch::new(), } } @@ -287,9 +294,19 @@ impl ApolloEngine { let perturb = Some((self.config.injection_layer, delta.view())); let raw = if let Some(ref bnd) = boundary { // Compressed: skip layers 0..crystal, run only crystal..34 (~4 layers) - forward_from_layer(weights, query_ids, bnd, crystal, perturb) + forward_from_layer( + larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch), + query_ids, + bnd, + crystal, + perturb, + ) } else { - forward_raw_logits(weights, &context, perturb) + forward_raw_logits( + larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch), + &context, + perturb, + ) }; let (top1_id, top1_logit) = raw .logits @@ -394,9 +411,19 @@ impl RetrievalEngine for ApolloEngine { let raw = if let Some(ref bnd) = boundary { // Compressed: boundary residual acts as position-0; skip layers 0..crystal. - forward_from_layer(weights, token_ids, bnd, crystal, perturb) + forward_from_layer( + larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch), + token_ids, + bnd, + crystal, + perturb, + ) } else { - forward_raw_logits(weights, &context, perturb) + forward_raw_logits( + larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch), + &context, + perturb, + ) }; // Cache decode state. @@ -432,14 +459,18 @@ impl RetrievalEngine for ApolloEngine { let raw = if let Some(ref bnd) = self.boundary_residual { // Compressed: re-run only crystal_layer..num_layers over growing query. forward_from_layer( - weights, + larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch), &self.context_tokens, bnd, self.crystal_layer, perturb, ) } else { - forward_raw_logits(weights, &self.context_tokens, perturb) + forward_raw_logits( + larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch), + &self.context_tokens, + perturb, + ) }; let last = raw.h_pre_norm.shape()[0] - 1; @@ -452,26 +483,44 @@ impl RetrievalEngine for ApolloEngine { /// override. fn prefill_quant( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, index: &larql_inference::larql_vindex::VectorIndex, token_ids: &[u32], ) -> Result, EngineError> { - larql_inference::vindex::ensure_attn_tensors_dequantised(weights, index); + larql_inference::vindex::ensure_attn_tensors_dequantised( + &mut self.dequant_scratch, + weights, + index, + ); for layer in 0..weights.num_layers { - let _ = larql_inference::vindex::insert_q4k_layer_tensors(weights, index, layer); + let _ = larql_inference::vindex::insert_q4k_layer_tensors( + &mut self.dequant_scratch, + weights, + index, + layer, + ); } self.prefill(weights, token_ids) } fn decode_step_quant( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, index: &larql_inference::larql_vindex::VectorIndex, token_id: u32, ) -> Result, EngineError> { - larql_inference::vindex::ensure_attn_tensors_dequantised(weights, index); + larql_inference::vindex::ensure_attn_tensors_dequantised( + &mut self.dequant_scratch, + weights, + index, + ); for layer in 0..weights.num_layers { - let _ = larql_inference::vindex::insert_q4k_layer_tensors(weights, index, layer); + let _ = larql_inference::vindex::insert_q4k_layer_tensors( + &mut self.dequant_scratch, + weights, + index, + layer, + ); } self.decode_step(weights, token_id) } @@ -863,7 +912,7 @@ mod tests { fn prefill_quant_via_executor_returns_hidden_state() { use larql_inference::ffn::NullFfn; use larql_inference::layer_executor::LocalWalkExecutor; - let mut weights = larql_inference::test_utils::make_test_weights(); + let weights = larql_inference::test_utils::make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let executor = LocalWalkExecutor::new(&*backend); @@ -871,7 +920,7 @@ mod tests { let mut engine = crate::AnyEngine::Retrieval(Box::new(mk_apollo_for_synthetic_weights(&weights))); let h = engine - .prefill_quant_via_executor(&mut weights, &executor, &ffn, &index, &[0u32, 1u32]) + .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32, 1u32]) .expect("executor prefill"); assert_eq!(h.shape(), &[1, weights.hidden_size]); } @@ -884,15 +933,15 @@ mod tests { // AnyEngine's `*_via_executor` forwards Retrieval variants to // these). We assert directly against the trait surface so we // retain access to `engine.context_tokens` for the post-condition. - let mut weights = larql_inference::test_utils::make_test_weights(); + let weights = larql_inference::test_utils::make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let mut engine = mk_apollo_for_synthetic_weights(&weights); engine - .prefill_quant(&mut weights, &index, &[0u32]) + .prefill_quant(&weights, &index, &[0u32]) .expect("prefill"); let ctx_before = engine.context_tokens.len(); let h = engine - .decode_step_quant(&mut weights, &index, 1) + .decode_step_quant(&weights, &index, 1) .expect("decode"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert_eq!( @@ -906,7 +955,7 @@ mod tests { fn prefill_via_executor_uncompressed_path_when_no_boundaries() { use larql_inference::ffn::NullFfn; use larql_inference::layer_executor::LocalWalkExecutor; - let mut weights = larql_inference::test_utils::make_test_weights(); + let weights = larql_inference::test_utils::make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let executor = LocalWalkExecutor::new(&*backend); @@ -922,7 +971,7 @@ mod tests { apollo.build_routing_index().unwrap(); let mut engine = crate::AnyEngine::Retrieval(Box::new(apollo)); let h = engine - .prefill_quant_via_executor(&mut weights, &executor, &ffn, &index, &[0u32, 1u32]) + .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32, 1u32]) .expect("executor prefill uncompressed"); assert_eq!(h.shape(), &[1, weights.hidden_size]); } @@ -956,7 +1005,7 @@ mod tests { #[test] fn executor_path_honors_ffn_parameter() { use larql_inference::layer_executor::LocalWalkExecutor; - let mut weights = larql_inference::test_utils::make_test_weights(); + let weights = larql_inference::test_utils::make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let executor = LocalWalkExecutor::new(&*backend); @@ -979,7 +1028,7 @@ mod tests { hidden: weights.hidden_size, }; engine - .prefill_quant_via_executor(&mut weights, &executor, &ffn, &index, &[0u32, 1u32]) + .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32, 1u32]) .expect("prefill via executor"); let calls = ffn.calls.load(std::sync::atomic::Ordering::SeqCst); // Post retrieval/KV trait split: ApolloEngine is now a @@ -1000,7 +1049,7 @@ mod tests { fn prefill_via_executor_falls_back_when_no_store() { use larql_inference::ffn::NullFfn; use larql_inference::layer_executor::LocalWalkExecutor; - let mut weights = larql_inference::test_utils::make_test_weights(); + let weights = larql_inference::test_utils::make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let executor = LocalWalkExecutor::new(&*backend); @@ -1009,7 +1058,7 @@ mod tests { crate::AnyEngine::Retrieval(Box::new(ApolloEngine::new(InjectionConfig::default()))); // No store → prepare_injection returns None → executor path returns None. assert!(engine - .prefill_quant_via_executor(&mut weights, &executor, &ffn, &index, &[0u32]) + .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32]) .is_err()); } } diff --git a/crates/larql-kv/src/engines/boundary_kv/engine.rs b/crates/larql-kv/src/engines/boundary_kv/engine.rs index 4a9f4c1f6..1f887cc76 100644 --- a/crates/larql-kv/src/engines/boundary_kv/engine.rs +++ b/crates/larql-kv/src/engines/boundary_kv/engine.rs @@ -277,7 +277,7 @@ impl KvEngine for BoundaryKvEngine { fn prefill_quant( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, ffn: &dyn FfnBackend, index: &larql_inference::larql_vindex::VectorIndex, token_ids: &[u32], @@ -300,7 +300,7 @@ impl KvEngine for BoundaryKvEngine { fn decode_step_quant( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, ffn: &dyn FfnBackend, index: &larql_inference::larql_vindex::VectorIndex, token_id: u32, @@ -831,18 +831,18 @@ mod tests { fn prefill_and_decode_quant_forward_and_emit() { use larql_inference::ffn::NullFfn; use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut eng = BoundaryKvEngine::new(config("seq", 2)); let h = eng - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1], &*backend) .expect("prefill_quant"); assert!(h.iter().all(|v| v.is_finite())); assert_eq!(eng.archive().total_frames(), Some(1)); let h2 = eng - .decode_step_quant(&mut weights, &ffn, &index, 2, &*backend) + .decode_step_quant(&weights, &ffn, &index, 2, &*backend) .expect("decode_step_quant"); assert!(h2.iter().all(|v| v.is_finite())); assert_eq!(eng.abs_position(), 3); @@ -852,7 +852,7 @@ mod tests { fn resident_and_quant_reject_empty_prompt() { use larql_inference::ffn::NullFfn; use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; @@ -862,7 +862,7 @@ mod tests { Err(EngineError::EmptyPrompt) )); assert!(matches!( - eng.prefill_quant(&mut weights, &ffn, &index, &[], &*backend), + eng.prefill_quant(&weights, &ffn, &index, &[], &*backend), Err(EngineError::EmptyPrompt) )); } diff --git a/crates/larql-kv/src/engines/boundary_per_layer/cold_tier.rs b/crates/larql-kv/src/engines/boundary_per_layer/cold_tier.rs index 65f7e984b..08dc179e6 100644 --- a/crates/larql-kv/src/engines/boundary_per_layer/cold_tier.rs +++ b/crates/larql-kv/src/engines/boundary_per_layer/cold_tier.rs @@ -22,7 +22,6 @@ use larql_compute::ComputeBackend; use larql_inference::attention::SharedKV; -use larql_inference::model::ModelWeights; use ndarray::{s, Array2}; use crate::engines::boundary_per_layer::policy::BoundaryLayerPolicy; @@ -41,7 +40,7 @@ use crate::engines::markov_residual_codec::codec::ColdResidualCodec; /// overflow lands — caller MUST snapshot this BEFORE appending the /// overflow to `cold_encoded` (which would advance `n_positions`). pub(super) fn extend_cold_kv_with_overflow( - weights: &ModelWeights, + weights: larql_inference::WeightsView, backend: &dyn ComputeBackend, policy: &BoundaryLayerPolicy, rs: &mut RsStorePerLayer, @@ -176,7 +175,7 @@ mod tests { .map(|_| Array2::::zeros((0, weights.hidden_size))) .collect(); let result = extend_cold_kv_with_overflow( - &weights, + larql_inference::WeightsView::dense(&weights), &CpuBackend, &policy, &mut rs, @@ -199,7 +198,7 @@ mod tests { .map(|_| Array2::::from_elem((2, weights.hidden_size), 0.3f32)) .collect(); let result = extend_cold_kv_with_overflow( - &weights, + larql_inference::WeightsView::dense(&weights), &CpuBackend, &policy, &mut rs, @@ -230,15 +229,22 @@ mod tests { let first_overflow: Vec> = (0..weights.num_layers) .map(|_| Array2::::from_elem((2, weights.hidden_size), 0.5f32)) .collect(); - extend_cold_kv_with_overflow(&weights, &CpuBackend, &policy, &mut rs, &first_overflow, 0) - .unwrap(); + extend_cold_kv_with_overflow( + larql_inference::WeightsView::dense(&weights), + &CpuBackend, + &policy, + &mut rs, + &first_overflow, + 0, + ) + .unwrap(); // Now extend with another 3 rows — exercises the Some arm. let second_overflow: Vec> = (0..weights.num_layers) .map(|_| Array2::::from_elem((3, weights.hidden_size), 0.7f32)) .collect(); let result = extend_cold_kv_with_overflow( - &weights, + larql_inference::WeightsView::dense(&weights), &CpuBackend, &policy, &mut rs, diff --git a/crates/larql-kv/src/engines/boundary_per_layer/dispatch.rs b/crates/larql-kv/src/engines/boundary_per_layer/dispatch.rs index 81c2ee822..042c33440 100644 --- a/crates/larql-kv/src/engines/boundary_per_layer/dispatch.rs +++ b/crates/larql-kv/src/engines/boundary_per_layer/dispatch.rs @@ -36,7 +36,7 @@ use crate::engines::w10_enabled as w10_env_on; /// backend / vindex lacks the required support (caller falls back to /// `walk::run_prefill`). pub(super) fn try_prefill_via_dispatch( - weights: &mut ModelWeights, + weights: &ModelWeights, backend: &dyn EngineBackend, policy: &BoundaryLayerPolicy, window_size: Option, @@ -103,8 +103,15 @@ pub(super) fn try_prefill_via_dispatch( for (layer, overflow) in overflow_per_layer.iter().enumerate() { 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"); + let (k, v) = recompute_kv( + larql_inference::WeightsView::dense(weights), + &decoded_overflow, + layer, + 0, + backend, + None, + ) + .expect("cold K/V pre-computation failed"); cold_kv.push((k, v)); let mut enc = PerLayerEncodedColdLayer::empty(codec, weights.hidden_size); enc.append(overflow); @@ -123,7 +130,7 @@ pub(super) fn try_prefill_via_dispatch( /// updated store. `None` signals a state-dump failure — caller should /// clear its `kv_handle` and fall back to the dense walk. pub(super) fn decode_step_via_dispatch( - weights: &mut ModelWeights, + weights: &ModelWeights, backend: &dyn EngineBackend, policy: &BoundaryLayerPolicy, handle: &mut KvHandle, @@ -211,7 +218,7 @@ pub(super) fn decode_step_via_dispatch( } } extend_cold_kv_with_overflow( - weights, + larql_inference::WeightsView::dense(weights), backend, policy, &mut rs, @@ -263,9 +270,9 @@ mod tests { weights.hidden_size, ); let backend = cpu_engine_backend(); - let mut w = weights; + let w = weights; assert!(try_prefill_via_dispatch( - &mut w, + &w, backend.as_ref(), &bf16_policy(), Some(4), @@ -278,11 +285,11 @@ mod tests { #[test] fn try_prefill_via_dispatch_windowed_populates_store_under_w10_honly() { clear_w10_override(); - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = cpu_engine_backend(); let (h, rs, _handle) = try_prefill_via_dispatch( - &mut weights, + &weights, backend.as_ref(), &bf16_policy(), Some(4), @@ -302,11 +309,11 @@ mod tests { #[test] fn try_prefill_via_dispatch_windowless_drops_stored_under_w10_none_mask() { clear_w10_override(); - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = cpu_engine_backend(); let (_h, rs, _handle) = try_prefill_via_dispatch( - &mut weights, + &weights, backend.as_ref(), &bf16_policy(), None, @@ -325,11 +332,11 @@ mod tests { #[test] fn decode_step_via_dispatch_appends_h_in_under_honly() { clear_w10_override(); - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = cpu_engine_backend(); let (_h, rs, mut handle) = try_prefill_via_dispatch( - &mut weights, + &weights, backend.as_ref(), &bf16_policy(), Some(4), @@ -339,7 +346,7 @@ mod tests { .expect("prefill"); let rows_before = rs.stored[0].shape()[0]; let (h, rs) = decode_step_via_dispatch( - &mut weights, + &weights, backend.as_ref(), &bf16_policy(), &mut handle, @@ -357,11 +364,11 @@ mod tests { #[test] fn decode_step_via_dispatch_windowless_takes_none_mask_path() { clear_w10_override(); - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = cpu_engine_backend(); let (_h, rs, mut handle) = try_prefill_via_dispatch( - &mut weights, + &weights, backend.as_ref(), &bf16_policy(), None, @@ -370,7 +377,7 @@ mod tests { ) .expect("prefill (windowless)"); let (_h, rs) = decode_step_via_dispatch( - &mut weights, + &weights, backend.as_ref(), &bf16_policy(), &mut handle, @@ -389,13 +396,13 @@ mod tests { #[test] fn decode_step_via_dispatch_overflow_extends_cold_tier() { clear_w10_override(); - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = cpu_engine_backend(); // window=2: after prefilling 2 tokens, one decode crosses the // window and the dispatch eviction path runs. let (_h, rs, mut handle) = try_prefill_via_dispatch( - &mut weights, + &weights, backend.as_ref(), &bf16_policy(), Some(2), @@ -405,7 +412,7 @@ mod tests { .expect("prefill"); assert!(rs.cold_encoded.is_none(), "no overflow at prefill"); let (_h, rs) = decode_step_via_dispatch( - &mut weights, + &weights, backend.as_ref(), &bf16_policy(), &mut handle, @@ -421,7 +428,7 @@ mod tests { // Subsequent decode should extend an existing cold_encoded // (Some(layers) branch of the match). let (_h, rs) = decode_step_via_dispatch( - &mut weights, + &weights, backend.as_ref(), &bf16_policy(), &mut handle, 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 76aeeec3b..6de43e376 100644 --- a/crates/larql-kv/src/engines/boundary_per_layer/engine.rs +++ b/crates/larql-kv/src/engines/boundary_per_layer/engine.rs @@ -55,6 +55,9 @@ pub struct BoundaryPerLayerEngine { /// dense walk (e.g. backend lacks cached_decode support). pub(super) kv_handle: Option, pub(super) backend: Box, + /// Engine-owned f32 dequant scratch for the per-layer walk fallback + /// (see `MarkovResidualEngine::dequant_scratch`). Keeps `weights` immutable. + pub(super) dequant_scratch: larql_inference::DequantScratch, } impl BoundaryPerLayerEngine { @@ -123,6 +126,7 @@ impl BoundaryPerLayerEngine { store: None, kv_handle: None, backend, + dequant_scratch: larql_inference::DequantScratch::new(), }) } @@ -150,8 +154,9 @@ impl BoundaryPerLayerEngine { .ok_or_else(|| EngineError::InvariantViolation { what: "decode_step called before prefill (store missing)".into(), })?; + let view = larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch); let (hidden, new_rs) = walk::run_decode( - weights, + view, ffn, self.backend.as_ref(), &self.policy, @@ -200,7 +205,7 @@ impl KvEngine for BoundaryPerLayerEngine { return Err(EngineError::EmptyPrompt); } let (hidden, store) = walk::run_prefill( - weights, + larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch), ffn, self.backend.as_ref(), &self.policy, @@ -255,7 +260,7 @@ impl KvEngine for BoundaryPerLayerEngine { fn prefill_quant( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, ffn: &dyn FfnBackend, index: &larql_inference::larql_vindex::VectorIndex, token_ids: &[u32], @@ -278,13 +283,17 @@ impl KvEngine for BoundaryPerLayerEngine { } // Fall back to dense f32 walk (compact vindexes / CPU backend). self.kv_handle = None; - larql_inference::vindex::dequant::ensure_attn_tensors_dequantised(weights, index); + larql_inference::vindex::dequant::ensure_attn_tensors_dequantised( + &mut self.dequant_scratch, + weights, + index, + ); self.prefill(weights, ffn, token_ids) } fn decode_step_quant( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, ffn: &dyn FfnBackend, index: &larql_inference::larql_vindex::VectorIndex, token_id: u32, @@ -329,7 +338,11 @@ impl KvEngine for BoundaryPerLayerEngine { } } } - larql_inference::vindex::dequant::ensure_attn_tensors_dequantised(weights, index); + larql_inference::vindex::dequant::ensure_attn_tensors_dequantised( + &mut self.dequant_scratch, + weights, + index, + ); self.decode_step(weights, ffn, token_id) } @@ -355,7 +368,7 @@ impl KvEngine for BoundaryPerLayerEngine { return self.prefill(weights, ffn, token_ids); } let (hidden, store) = executor::run_prefill( - weights, + larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch), executor, ffn, &self.policy, @@ -386,12 +399,17 @@ impl KvEngine for BoundaryPerLayerEngine { .ok_or_else(|| EngineError::InvariantViolation { what: "decode_step_via_executor called before prefill (store missing)".into(), })?; - let (hidden, new_rs) = - executor::run_decode(weights, executor, ffn, &self.policy, rs, token_id).ok_or_else( - || EngineError::BackendFailure { - details: "executor::run_decode returned None".into(), - }, - )?; + let (hidden, new_rs) = executor::run_decode( + larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch), + executor, + ffn, + &self.policy, + rs, + token_id, + ) + .ok_or_else(|| EngineError::BackendFailure { + details: "executor::run_decode returned None".into(), + })?; self.store = Some(new_rs); Ok(hidden) } @@ -800,7 +818,7 @@ mod tests { // NullFfn instead — the dense walk's FFN dispatch through // NullFfn produces zero residuals, which is fine for shape // checks (we only assert the output shape, not values). - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let policy = BoundaryLayerPolicy::bf16_uniform("test", weights.num_layers); let store = store_with_record(&policy); @@ -809,7 +827,7 @@ mod tests { let backend = larql_compute::CpuBackend; let ffn = larql_inference::ffn::NullFfn; let h = engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1], &backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1], &backend) .expect("dispatch-None fall-through must succeed via walk"); assert_eq!(h.shape(), &[1, weights.hidden_size]); // kv_handle should be None on the fall-through path. @@ -821,7 +839,7 @@ mod tests { // After a fall-through prefill (kv_handle is None), decode_step_quant // takes the `self.kv_handle.is_some()` == false path → falls into // self.decode_step via the dense walk. - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let policy = BoundaryLayerPolicy::bf16_uniform("test", weights.num_layers); let store = store_with_record(&policy); @@ -830,11 +848,11 @@ mod tests { let backend = larql_compute::CpuBackend; let ffn = larql_inference::ffn::NullFfn; engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1], &backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1], &backend) .unwrap(); assert!(engine.kv_handle.is_none()); let h = engine - .decode_step_quant(&mut weights, &ffn, &index, 2, &backend) + .decode_step_quant(&weights, &ffn, &index, 2, &backend) .expect("decode fall-through must succeed"); assert_eq!(h.shape(), &[1, weights.hidden_size]); } @@ -940,7 +958,7 @@ mod tests { #[test] fn prefill_quant_returns_empty_prompt_error_on_empty_input() { use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let policy = BoundaryLayerPolicy::bf16_uniform("test", weights.num_layers); @@ -949,7 +967,7 @@ mod tests { BoundaryPerLayerEngine::new(None, policy, weights.num_layers, &store).unwrap(); let ffn = larql_inference::ffn::NullFfn; let err = engine - .prefill_quant(&mut weights, &ffn, &index, &[], &*backend) + .prefill_quant(&weights, &ffn, &index, &[], &*backend) .unwrap_err(); assert_eq!(err, larql_inference::kv_engine::EngineError::EmptyPrompt); } @@ -973,7 +991,7 @@ mod tests { #[test] fn decode_step_quant_returns_invariant_violation_before_prefill() { use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let policy = BoundaryLayerPolicy::bf16_uniform("test", weights.num_layers); @@ -982,7 +1000,7 @@ mod tests { BoundaryPerLayerEngine::new(None, policy, weights.num_layers, &store).unwrap(); let ffn = larql_inference::ffn::NullFfn; let err = engine - .decode_step_quant(&mut weights, &ffn, &index, 0, &*backend) + .decode_step_quant(&weights, &ffn, &index, 0, &*backend) .unwrap_err(); assert!(matches!( err, 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 9cda5e1a3..29009d404 100644 --- a/crates/larql-kv/src/engines/boundary_per_layer/executor.rs +++ b/crates/larql-kv/src/engines/boundary_per_layer/executor.rs @@ -15,7 +15,6 @@ use larql_inference::attention::SharedKV; use larql_inference::ffn::FfnBackend; use larql_inference::forward::embed_tokens_pub; use larql_inference::layer_executor::LayerExecutor; -use larql_inference::model::ModelWeights; use ndarray::{s, Array2}; use crate::engines::boundary_per_layer::cold_tier::{ @@ -29,7 +28,7 @@ use crate::engines::markov_residual::recompute_kv; /// `executor.dispatch_kind() != Fused` (engine glue falls back to /// `walk::run_prefill` in that case). pub(super) fn run_prefill( - weights: &ModelWeights, + weights: larql_inference::WeightsView, executor: &dyn LayerExecutor, ffn: &dyn FfnBackend, policy: &BoundaryLayerPolicy, @@ -39,7 +38,7 @@ pub(super) fn run_prefill( let backend = executor.backend(); let num_layers = weights.num_layers; let seq_len = token_ids.len(); - let mut h = embed_tokens_pub(weights, token_ids); + let mut h = embed_tokens_pub(&weights, token_ids); let mut stored: Vec> = Vec::with_capacity(num_layers); for layer in 0..num_layers { @@ -87,7 +86,7 @@ pub(super) fn run_prefill( /// Executor-driven decode step. Caller MUST have already checked that /// `executor.dispatch_kind() != Fused`. pub(super) fn run_decode( - weights: &ModelWeights, + weights: larql_inference::WeightsView, executor: &dyn LayerExecutor, ffn: &dyn FfnBackend, policy: &BoundaryLayerPolicy, @@ -97,7 +96,7 @@ pub(super) fn run_decode( let backend = executor.backend(); let num_layers = weights.num_layers; let abs_position = rs.next_position; - let mut h_new = embed_tokens_pub(weights, &[token_id]); + let mut h_new = embed_tokens_pub(&weights, &[token_id]); let mut new_stored: Vec> = Vec::with_capacity(num_layers); for layer in 0..num_layers { @@ -203,8 +202,15 @@ mod tests { let ffn = NullFfn; let policy = BoundaryLayerPolicy::bf16_uniform("test", weights.num_layers); let token_ids: Vec = vec![0, 1, 2]; - let (hidden, rs) = run_prefill(&weights, &executor, &ffn, &policy, None, &token_ids) - .expect("prefill should succeed with synthetic weights"); + let (hidden, rs) = run_prefill( + larql_inference::WeightsView::dense(&weights), + &executor, + &ffn, + &policy, + None, + &token_ids, + ) + .expect("prefill should succeed with synthetic weights"); assert_eq!(hidden.shape(), &[1, weights.hidden_size]); assert_eq!(rs.next_position, 3); assert!(rs.cold_encoded.is_none(), "no overflow → no cold_encoded"); @@ -226,8 +232,15 @@ mod tests { let ffn = NullFfn; let policy = BoundaryLayerPolicy::bf16_uniform("test", weights.num_layers); let token_ids: Vec = vec![0, 1, 2]; - let (_hidden, rs) = run_prefill(&weights, &executor, &ffn, &policy, Some(2), &token_ids) - .expect("prefill should succeed"); + let (_hidden, rs) = run_prefill( + larql_inference::WeightsView::dense(&weights), + &executor, + &ffn, + &policy, + Some(2), + &token_ids, + ) + .expect("prefill should succeed"); assert!( rs.cold_encoded.is_some(), "overflow path must populate cold_encoded" @@ -251,14 +264,29 @@ mod tests { let executor = LocalWalkExecutor::new(&backend); let ffn = NullFfn; let policy = BoundaryLayerPolicy::bf16_uniform("test", weights.num_layers); - let (_, rs) = run_prefill(&weights, &executor, &ffn, &policy, Some(4), &[0]).unwrap(); + let (_, rs) = run_prefill( + larql_inference::WeightsView::dense(&weights), + &executor, + &ffn, + &policy, + Some(4), + &[0], + ) + .unwrap(); assert!( rs.cold_encoded.is_none(), "no overflow expected after prefill" ); - let (hidden, rs_after) = - run_decode(&weights, &executor, &ffn, &policy, rs, 1).expect("decode should succeed"); + let (hidden, rs_after) = run_decode( + larql_inference::WeightsView::dense(&weights), + &executor, + &ffn, + &policy, + rs, + 1, + ) + .expect("decode should succeed"); assert_eq!(hidden.shape(), &[1, weights.hidden_size]); assert_eq!(rs_after.next_position, 2); for slab in &rs_after.stored { @@ -278,7 +306,15 @@ mod tests { let executor = LocalWalkExecutor::new(&backend); let ffn = NullFfn; let policy = BoundaryLayerPolicy::bf16_uniform("test", weights.num_layers); - let (_, rs) = run_prefill(&weights, &executor, &ffn, &policy, Some(2), &[0, 1, 2]).unwrap(); + let (_, rs) = run_prefill( + larql_inference::WeightsView::dense(&weights), + &executor, + &ffn, + &policy, + Some(2), + &[0, 1, 2], + ) + .unwrap(); assert!( rs.cold_encoded.is_some(), "prefill should have populated cold_encoded" @@ -290,7 +326,15 @@ mod tests { .unwrap_or(0); assert_eq!(initial_cold_rows, 1, "1 row in cold after prefill"); - let (_, rs_after) = run_decode(&weights, &executor, &ffn, &policy, rs, 3).unwrap(); + let (_, rs_after) = run_decode( + larql_inference::WeightsView::dense(&weights), + &executor, + &ffn, + &policy, + rs, + 3, + ) + .unwrap(); let after_cold_rows = rs_after .cold_encoded .as_ref() 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 56828ac42..ed32156bc 100644 --- a/crates/larql-kv/src/engines/boundary_per_layer/walk.rs +++ b/crates/larql-kv/src/engines/boundary_per_layer/walk.rs @@ -13,7 +13,6 @@ use larql_compute::ComputeBackend; 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::model::ModelWeights; use ndarray::{s, Array2}; use crate::engines::boundary_per_layer::cold_tier::{ @@ -26,7 +25,7 @@ use crate::engines::markov_residual::recompute_kv; /// Run a full prefill through the dense walk. Returns /// `(last_hidden, new_store)` — caller owns the store. pub(super) fn run_prefill( - weights: &ModelWeights, + weights: larql_inference::WeightsView, ffn: &dyn FfnBackend, backend: &dyn ComputeBackend, policy: &BoundaryLayerPolicy, @@ -35,15 +34,21 @@ pub(super) fn run_prefill( ) -> Option<(Array2, RsStorePerLayer)> { let num_layers = weights.num_layers; let seq_len = token_ids.len(); - let mut h = embed_tokens_pub(weights, token_ids); + let mut h = embed_tokens_pub(&weights, 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).expect("attention failed"); - let h_out = crate::engines::layer_ffn_or_moe(weights, &h_post_attn, layer, ffn, Some(ffn)); + run_attention_with_kv_backend(weights, &h, layer, be, None).expect("attention failed"); + let h_out = crate::engines::layer_ffn_or_moe( + weights.canonical(), + &h_post_attn, + layer, + ffn, + Some(ffn), + ); h = h_out; } @@ -87,7 +92,7 @@ pub(super) fn run_prefill( /// the new store alongside the hidden output. #[allow(clippy::too_many_arguments)] pub(super) fn run_decode( - weights: &ModelWeights, + weights: larql_inference::WeightsView, ffn: &dyn FfnBackend, backend: &dyn ComputeBackend, policy: &BoundaryLayerPolicy, @@ -97,7 +102,7 @@ pub(super) fn run_decode( ) -> Option<(Array2, RsStorePerLayer)> { let num_layers = weights.num_layers; let abs_position = rs.next_position; - let mut h_new = embed_tokens_pub(weights, &[token_id]); + let mut h_new = embed_tokens_pub(&weights, &[token_id]); let mut new_stored: Vec> = Vec::with_capacity(num_layers); // W2 hot-K/V cache (twin of markov_residual_codec). When unbounded with no @@ -259,7 +264,13 @@ pub(super) fn run_decode( h_post_attn }; - let h_out = crate::engines::layer_ffn_or_moe(weights, &h_post_attn, layer, ffn, Some(ffn)); + let h_out = crate::engines::layer_ffn_or_moe( + weights.canonical(), + &h_post_attn, + layer, + ffn, + Some(ffn), + ); h_new = h_out; } @@ -336,8 +347,15 @@ mod tests { let backend = CpuBackend; let ffn = NullFfn; let policy = BoundaryLayerPolicy::bf16_uniform("test", weights.num_layers); - let (hidden, rs) = run_prefill(&weights, &ffn, &backend, &policy, None, &[0, 1, 2]) - .expect("prefill should succeed"); + let (hidden, rs) = run_prefill( + larql_inference::WeightsView::dense(&weights), + &ffn, + &backend, + &policy, + None, + &[0, 1, 2], + ) + .expect("prefill should succeed"); assert_eq!(hidden.shape(), &[1, weights.hidden_size]); assert_eq!(rs.next_position, 3); assert!(rs.cold_encoded.is_none()); @@ -353,7 +371,15 @@ mod tests { let backend = CpuBackend; let ffn = NullFfn; let policy = BoundaryLayerPolicy::bf16_uniform("test", weights.num_layers); - let (_, rs) = run_prefill(&weights, &ffn, &backend, &policy, Some(2), &[0, 1, 2]).unwrap(); + let (_, rs) = run_prefill( + larql_inference::WeightsView::dense(&weights), + &ffn, + &backend, + &policy, + Some(2), + &[0, 1, 2], + ) + .unwrap(); assert!(rs.cold_encoded.is_some(), "overflow → cold_encoded"); assert!(rs.cold_kv.is_some(), "overflow → cold_kv pre-computed"); let cold_kv = rs.cold_kv.as_ref().unwrap(); @@ -368,11 +394,27 @@ mod tests { let backend = CpuBackend; let ffn = NullFfn; let policy = BoundaryLayerPolicy::bf16_uniform("test", weights.num_layers); - let (_, rs) = run_prefill(&weights, &ffn, &backend, &policy, Some(4), &[0]).unwrap(); + let (_, rs) = run_prefill( + larql_inference::WeightsView::dense(&weights), + &ffn, + &backend, + &policy, + Some(4), + &[0], + ) + .unwrap(); assert!(rs.cold_encoded.is_none()); - let (hidden, rs_after) = - run_decode(&weights, &ffn, &backend, &policy, rs, 1, None).unwrap(); + let (hidden, rs_after) = run_decode( + larql_inference::WeightsView::dense(&weights), + &ffn, + &backend, + &policy, + rs, + 1, + None, + ) + .unwrap(); assert_eq!(hidden.shape(), &[1, weights.hidden_size]); assert_eq!(rs_after.next_position, 2); for slab in &rs_after.stored { @@ -396,8 +438,15 @@ mod tests { let policy = BoundaryLayerPolicy::bf16_uniform("test", weights.num_layers); // First build a normal state with overflow. - let (_, mut rs) = - run_prefill(&weights, &ffn, &backend, &policy, Some(2), &[0, 1, 2]).unwrap(); + let (_, mut rs) = run_prefill( + larql_inference::WeightsView::dense(&weights), + &ffn, + &backend, + &policy, + Some(2), + &[0, 1, 2], + ) + .unwrap(); // Now wipe the pre-computed cold_kv. cold_encoded stays // populated. Decode should recompute K/V from the decoded // cold residuals. @@ -406,8 +455,16 @@ mod tests { assert!(rs.cold_encoded.as_ref().unwrap()[0].n_positions > 0); let _ = ColdResidualCodec::Bf16; // keep import live - let (hidden, _) = run_decode(&weights, &ffn, &backend, &policy, rs, 3, None) - .expect("decode should succeed without cold_kv"); + let (hidden, _) = run_decode( + larql_inference::WeightsView::dense(&weights), + &ffn, + &backend, + &policy, + rs, + 3, + None, + ) + .expect("decode should succeed without cold_kv"); assert_eq!(hidden.shape(), &[1, weights.hidden_size]); } @@ -438,12 +495,27 @@ mod tests { "LARQL_MARKOV_INPLACE_KV", Some(if inplace { "1" } else { "0" }), ); - let (_, mut rs) = - run_prefill(&weights, &ffn, &backend, &policy, None, &[0u32, 1, 2]).unwrap(); + let (_, mut rs) = run_prefill( + larql_inference::WeightsView::dense(&weights), + &ffn, + &backend, + &policy, + None, + &[0u32, 1, 2], + ) + .unwrap(); let mut hiddens = Vec::new(); for tok in 3u32..=12 { - let (h, rs2) = - run_decode(&weights, &ffn, &backend, &policy, rs, tok, Some(&index)).unwrap(); + let (h, rs2) = run_decode( + larql_inference::WeightsView::dense(&weights), + &ffn, + &backend, + &policy, + rs, + tok, + Some(&index), + ) + .unwrap(); assert!(h.iter().all(|v| v.is_finite())); hiddens.push(h.iter().map(|v| v.to_bits()).collect()); rs = rs2; @@ -469,7 +541,15 @@ mod tests { let backend = CpuBackend; let ffn = NullFfn; let policy = BoundaryLayerPolicy::bf16_uniform("test", weights.num_layers); - let (_, rs) = run_prefill(&weights, &ffn, &backend, &policy, Some(2), &[0, 1, 2]).unwrap(); + let (_, rs) = run_prefill( + larql_inference::WeightsView::dense(&weights), + &ffn, + &backend, + &policy, + Some(2), + &[0, 1, 2], + ) + .unwrap(); let initial = rs .cold_encoded .as_ref() @@ -477,7 +557,16 @@ mod tests { .unwrap_or(0); assert_eq!(initial, 1); - let (_, rs_after) = run_decode(&weights, &ffn, &backend, &policy, rs, 3, None).unwrap(); + let (_, rs_after) = run_decode( + larql_inference::WeightsView::dense(&weights), + &ffn, + &backend, + &policy, + rs, + 3, + None, + ) + .unwrap(); let after = rs_after .cold_encoded .as_ref() @@ -500,14 +589,31 @@ mod tests { let policy = BoundaryLayerPolicy::bf16_uniform("test", weights.num_layers); // window=2, exactly 2 prompt tokens → fills the window, no overflow. - let (_, rs) = run_prefill(&weights, &ffn, &backend, &policy, Some(2), &[0, 1]).unwrap(); + let (_, rs) = run_prefill( + larql_inference::WeightsView::dense(&weights), + &ffn, + &backend, + &policy, + Some(2), + &[0, 1], + ) + .unwrap(); assert!( rs.cold_encoded.is_none(), "a window-filling prefill must not overflow yet" ); // First decode overflows by one row → initialises cold_encoded. - let (_, rs_after) = run_decode(&weights, &ffn, &backend, &policy, rs, 2, None).unwrap(); + let (_, rs_after) = run_decode( + larql_inference::WeightsView::dense(&weights), + &ffn, + &backend, + &policy, + rs, + 2, + None, + ) + .unwrap(); assert!( rs_after.cold_encoded.is_some(), "first decode overflow must initialise the cold tier" diff --git a/crates/larql-kv/src/engines/markov_residual/compute.rs b/crates/larql-kv/src/engines/markov_residual/compute.rs index d189894f4..775781505 100644 --- a/crates/larql-kv/src/engines/markov_residual/compute.rs +++ b/crates/larql-kv/src/engines/markov_residual/compute.rs @@ -13,7 +13,6 @@ 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::{add_bias, apply_norm, embed_tokens_pub}; -use larql_inference::model::ModelWeights; use larql_inference::residual::{rms_norm_heads, rms_norm_heads_no_weight}; #[derive(Clone, Copy)] @@ -83,7 +82,7 @@ pub struct RsPrefillResult { } pub fn rs_prefill( - weights: &ModelWeights, + weights: larql_inference::WeightsView, token_ids: &[u32], max_window: Option, backend: &dyn ComputeBackend, @@ -91,16 +90,25 @@ pub fn rs_prefill( ) -> RsPrefillResult { let num_layers = weights.num_layers; let seq_len = token_ids.len(); - let mut h = embed_tokens_pub(weights, token_ids); + let mut h = embed_tokens_pub(&weights, 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) + 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, backend }; - let h_out = crate::engines::layer_ffn_or_moe(weights, &h_post_attn, layer, &bffn, moe_ffn); + 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, + ); h = h_out; } @@ -146,7 +154,7 @@ pub fn rs_prefill( } pub fn rs_decode_step( - weights: &ModelWeights, + weights: larql_inference::WeightsView, new_token_id: u32, rs: RsStore, backend: &dyn ComputeBackend, @@ -158,7 +166,7 @@ pub fn rs_decode_step( #[allow(clippy::too_many_arguments)] pub(crate) fn rs_decode_step_profiled( - weights: &ModelWeights, + weights: larql_inference::WeightsView, new_token_id: u32, rs: RsStore, backend: &dyn ComputeBackend, @@ -179,7 +187,7 @@ pub(crate) fn rs_decode_step_profiled( #[allow(clippy::too_many_arguments)] fn rs_decode_step_inner( - weights: &ModelWeights, + weights: larql_inference::WeightsView, new_token_id: u32, rs: RsStore, backend: &dyn ComputeBackend, @@ -196,7 +204,7 @@ fn rs_decode_step_inner( } else { None }; - let mut h_new = embed_tokens_pub(weights, &[new_token_id]); + let mut h_new = embed_tokens_pub(&weights, &[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; @@ -422,8 +430,17 @@ fn rs_decode_step_inner( } else { None }; - let bffn = BackendFfn { weights, backend }; - let h_out = crate::engines::layer_ffn_or_moe(weights, &h_post_attn, layer, &bffn, moe_ffn); + 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, + ); if let Some(t) = t_ffn { ffn_us += t.elapsed().as_secs_f64() * 1e6; } @@ -517,7 +534,7 @@ fn rs_decode_step_inner( /// the right kernel (Q4K today; Q6K / future formats slot in /// automatically). `None` keeps the f32 fallback for legacy callers. pub fn recompute_kv( - weights: &ModelWeights, + weights: larql_inference::WeightsView, h_stored: &Array2, layer: usize, abs_start: usize, @@ -536,7 +553,7 @@ pub fn recompute_kv( }; let h_norm = apply_norm( - weights, + &weights, h_stored, &arch.input_layernorm_key(layer), norm_offset, @@ -699,17 +716,17 @@ type AttnKvWeightPair<'a> = ( &'a ArrayBase, Ix2>, ); -fn attn_kv_projection_weights( - weights: &ModelWeights, +fn attn_kv_projection_weights<'a>( + weights: larql_inference::WeightsView<'a>, layer: usize, -) -> Option> { +) -> Option> { let arch = &*weights.arch; - let w_k = weights.tensors.get(&arch.attn_k_key(layer))?; - let v_from_k = !weights.tensors.contains_key(&arch.attn_v_key(layer)); + let w_k = weights.tensor(&arch.attn_k_key(layer))?; + let v_from_k = !weights.has_tensor(&arch.attn_v_key(layer)); let w_v = if v_from_k { w_k } else { - weights.tensors.get(&arch.attn_v_key(layer))? + weights.tensor(&arch.attn_v_key(layer))? }; Some((w_k, w_v)) } @@ -994,7 +1011,7 @@ fn parse_quant_format(fmt: &str) -> Option { } /// Equivalent Standard KV memory in bytes for `seq_len` tokens (FP16). -pub fn kv_memory_bytes_for_seq(weights: &ModelWeights, seq_len: usize) -> usize { +pub fn kv_memory_bytes_for_seq(weights: larql_inference::WeightsView, seq_len: usize) -> usize { let arch = &*weights.arch; (0..weights.num_layers) .map(|l| { @@ -1021,7 +1038,14 @@ mod tests { fn recompute_kv_returns_some_with_valid_weights() { let weights = make_test_weights(); let h = Array2::from_elem((3, weights.hidden_size), 0.5f32); - let result = recompute_kv(&weights, &h, 0, 0, &CpuBackend, None); + let result = recompute_kv( + larql_inference::WeightsView::dense(&weights), + &h, + 0, + 0, + &CpuBackend, + None, + ); assert!( result.is_some(), "recompute_kv should return Some with valid weights" @@ -1033,7 +1057,15 @@ mod tests { let weights = make_test_weights(); let seq_len = 4; let h = Array2::from_elem((seq_len, weights.hidden_size), 1.0f32); - let (k, v) = recompute_kv(&weights, &h, 0, 0, &CpuBackend, None).unwrap(); + let (k, v) = recompute_kv( + larql_inference::WeightsView::dense(&weights), + &h, + 0, + 0, + &CpuBackend, + None, + ) + .unwrap(); let kv_dim = weights.num_kv_heads * weights.head_dim; assert_eq!(k.shape(), &[seq_len, kv_dim], "K shape mismatch"); assert_eq!(v.shape(), &[seq_len, kv_dim], "V shape mismatch"); @@ -1043,7 +1075,15 @@ mod tests { fn recompute_kv_output_is_finite() { let weights = make_test_weights(); let h = Array2::from_elem((2, weights.hidden_size), 0.1f32); - let (k, v) = recompute_kv(&weights, &h, 0, 0, &CpuBackend, None).unwrap(); + let (k, v) = recompute_kv( + larql_inference::WeightsView::dense(&weights), + &h, + 0, + 0, + &CpuBackend, + None, + ) + .unwrap(); assert!( k.iter().all(|v| v.is_finite()), "K contains non-finite values" @@ -1059,8 +1099,24 @@ mod tests { let weights = make_test_weights(); let h = Array2::from_elem((1, weights.hidden_size), 0.5f32); // Different abs_start should produce different RoPE-applied K - let (k0, _) = recompute_kv(&weights, &h, 0, 0, &CpuBackend, None).unwrap(); - let (k5, _) = recompute_kv(&weights, &h, 0, 5, &CpuBackend, None).unwrap(); + let (k0, _) = recompute_kv( + larql_inference::WeightsView::dense(&weights), + &h, + 0, + 0, + &CpuBackend, + None, + ) + .unwrap(); + let (k5, _) = recompute_kv( + larql_inference::WeightsView::dense(&weights), + &h, + 0, + 5, + &CpuBackend, + None, + ) + .unwrap(); let diff: f32 = k0.iter().zip(k5.iter()).map(|(a, b)| (a - b).abs()).sum(); assert!( diff > 0.0, @@ -1158,7 +1214,13 @@ mod tests { #[test] fn rs_prefill_returns_correct_shape() { let weights = make_test_weights(); - let result = rs_prefill(&weights, &[0u32, 1, 2], None, &CpuBackend, None); + 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())); } @@ -1166,7 +1228,13 @@ mod tests { #[test] fn rs_prefill_stores_all_layers() { let weights = make_test_weights(); - let result = rs_prefill(&weights, &[0u32], None, &CpuBackend, None); + 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); } @@ -1174,7 +1242,13 @@ mod tests { #[test] fn rs_prefill_with_window_clips_hot_store() { let weights = make_test_weights(); - let result = rs_prefill(&weights, &[0u32, 1, 2, 3, 4], Some(2), &CpuBackend, None); + 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", @@ -1187,9 +1261,22 @@ mod tests { #[test] fn rs_decode_step_produces_finite_hidden() { let weights = make_test_weights(); - let prefill = rs_prefill(&weights, &[0u32], None, &CpuBackend, None); - let (h, _) = rs_decode_step(&weights, 1, prefill.store, &CpuBackend, None, None) - .expect("decode step"); + 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())); } @@ -1197,11 +1284,33 @@ mod tests { #[test] fn rs_decode_step_advances_position() { let weights = make_test_weights(); - let prefill = rs_prefill(&weights, &[0u32, 1], None, &CpuBackend, None); + 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(&weights, 2, prefill.store, &CpuBackend, None, None).unwrap(); + 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(&weights, 3, rs2, &CpuBackend, None, None).unwrap(); + let (_, rs3) = rs_decode_step( + larql_inference::WeightsView::dense(&weights), + 3, + rs2, + &CpuBackend, + None, + None, + ) + .unwrap(); assert_eq!(rs3.next_position, 4); } @@ -1212,20 +1321,40 @@ mod tests { // `Some(cold_kv)` branch (lines 128-147) instead of the // cold-residual recomputation path. let weights = make_test_weights(); - let prefill = rs_prefill(&weights, &[0u32, 1, 2, 3], Some(2), &CpuBackend, None); + 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(&weights, 4, prefill.store, &CpuBackend, None, None) - .expect("decode_step over cold_kv"); + 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(&weights, 5, rs2, &CpuBackend, None, None) - .expect("decode_step over cold_residuals"); + 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())); } @@ -1260,12 +1389,25 @@ mod tests { "LARQL_MARKOV_INPLACE_KV", Some(if inplace { "1" } else { "0" }), ); - let prefill = rs_prefill(&weights, &[0u32, 1, 2], None, &CpuBackend, None); + 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(&weights, tok, rs, &CpuBackend, None, Some(&index)) - .expect("decode"); + 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; @@ -1292,8 +1434,8 @@ mod tests { #[test] fn kv_memory_bytes_for_seq_scales_linearly() { let weights = make_test_weights(); - let one = kv_memory_bytes_for_seq(&weights, 1); - let ten = kv_memory_bytes_for_seq(&weights, 10); + let one = kv_memory_bytes_for_seq(larql_inference::WeightsView::dense(&weights), 1); + let ten = kv_memory_bytes_for_seq(larql_inference::WeightsView::dense(&weights), 10); assert!(one > 0); assert_eq!(ten, one * 10, "kv memory must scale linearly with seq len"); } @@ -1336,13 +1478,19 @@ mod tests { fn profiled_decode_step_exercises_all_timing_branches() { use crate::profiler::EngineProfiler; let weights = make_test_weights(); - let prefill = rs_prefill(&weights, &[0u32, 1, 2, 3], Some(2), &CpuBackend, None); + 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( - &weights, + larql_inference::WeightsView::dense(&weights), 4, prefill.store, &CpuBackend, @@ -1365,15 +1513,36 @@ mod tests { // 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(&weights, &[0u32, 1, 2, 3], Some(2), &CpuBackend, None); - let (_, rs2) = rs_decode_step(&weights, 4, prefill.store, &CpuBackend, None, None).unwrap(); + 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(&weights, 5, rs2, &CpuBackend, &mut profiler, None, None); + 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); @@ -1568,7 +1737,15 @@ mod tests { set_markov_env_override("LARQL_MARKOV_KV_FORCE_F32", Some("1")); let weights = make_test_weights(); let h = Array2::from_elem((2, weights.hidden_size), 0.5f32); - let (k, v) = recompute_kv(&weights, &h, 0, 0, &CpuBackend, None).unwrap(); + let (k, v) = recompute_kv( + larql_inference::WeightsView::dense(&weights), + &h, + 0, + 0, + &CpuBackend, + None, + ) + .unwrap(); let kv_dim = weights.num_kv_heads * weights.head_dim; assert_eq!(k.shape(), &[2, kv_dim]); assert_eq!(v.shape(), &[2, kv_dim]); @@ -1581,7 +1758,14 @@ mod tests { set_markov_env_override("LARQL_MARKOV_WALK_KV_TOPK", Some("2")); let weights = make_test_weights(); let h = Array2::from_elem((2, weights.hidden_size), 0.25f32); - let result = recompute_kv(&weights, &h, 0, 0, &CpuBackend, None); + let result = recompute_kv( + larql_inference::WeightsView::dense(&weights), + &h, + 0, + 0, + &CpuBackend, + None, + ); assert!(result.is_some()); clear_markov_env_overrides(); } @@ -1595,9 +1779,23 @@ mod tests { let h = Array2::from_elem((2, weights.hidden_size), 0.25f32); // Layer 0: should_cache_selection fires, populates // WALK_KV_SELECTION; layer 1: walk_project_cached_topk reads it. - let _ = recompute_kv(&weights, &h, 0, 0, &CpuBackend, None); + let _ = recompute_kv( + larql_inference::WeightsView::dense(&weights), + &h, + 0, + 0, + &CpuBackend, + None, + ); if weights.num_layers >= 2 { - let result = recompute_kv(&weights, &h, 1, 0, &CpuBackend, None); + let result = recompute_kv( + larql_inference::WeightsView::dense(&weights), + &h, + 1, + 0, + &CpuBackend, + None, + ); assert!(result.is_some()); } clear_markov_env_overrides(); @@ -1609,7 +1807,14 @@ mod tests { set_markov_env_override("LARQL_MARKOV_WALK_KV_DIAG", Some("1")); let weights = make_test_weights(); let h = Array2::from_elem((1, weights.hidden_size), 0.5f32); - let result = recompute_kv(&weights, &h, 0, 0, &CpuBackend, None); + let result = recompute_kv( + larql_inference::WeightsView::dense(&weights), + &h, + 0, + 0, + &CpuBackend, + None, + ); assert!(result.is_some()); clear_markov_env_overrides(); } @@ -1646,7 +1851,14 @@ mod tests { max_window: None, cold_len: 0, }; - let result = rs_decode_step(&weights, 0, store, &CpuBackend, None, None); + 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/dispatch.rs b/crates/larql-kv/src/engines/markov_residual/dispatch.rs index ab84e8caf..2aa2656bd 100644 --- a/crates/larql-kv/src/engines/markov_residual/dispatch.rs +++ b/crates/larql-kv/src/engines/markov_residual/dispatch.rs @@ -30,7 +30,7 @@ impl MarkovResidualEngine { /// doesn't (engine falls back to per-layer walk). pub(super) fn try_prefill_via_dispatch( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, index: &VectorIndex, token_ids: &[u32], ) -> Option> { @@ -133,7 +133,7 @@ impl MarkovResidualEngine { /// (W8.2 doubling-capacity in-place append). pub(super) fn decode_step_via_dispatch( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, index: &VectorIndex, token_id: u32, ) -> Option> { @@ -305,9 +305,9 @@ mod tests { weights.hidden_size, ); let mut engine = MarkovResidualEngine::with_backend(Some(4), cpu_engine_backend()); - let mut w = weights; + let w = weights; assert!(engine - .try_prefill_via_dispatch(&mut w, &empty_index, &[0u32, 1]) + .try_prefill_via_dispatch(&w, &empty_index, &[0u32, 1]) .is_none()); assert!(engine.store.is_none()); assert!(engine.kv_handle.is_none()); @@ -321,9 +321,9 @@ mod tests { // `stored[layer]` is a doubling-capacity buffer // (shape `[max(window,prompt_len), hidden]`); the logical row // count lives in `hot_len`. - let (mut engine, mut weights, index) = fixture(Some(8)); + let (mut engine, weights, index) = fixture(Some(8)); let h = engine - .try_prefill_via_dispatch(&mut weights, &index, &[0u32, 1, 2]) + .try_prefill_via_dispatch(&weights, &index, &[0u32, 1, 2]) .expect("prefill"); assert_eq!(h.shape(), &[1, weights.hidden_size]); let rs = engine.store.as_ref().expect("store populated"); @@ -345,10 +345,10 @@ mod tests { // window=None + default env → both drop_hot_kv_shadow and // drop_stored_shadow = true. The drop_stored branch replaces // each `stored[l]` with `Array2::::zeros((0, hidden))`. - let (mut engine, mut weights, index) = fixture(None); + let (mut engine, weights, index) = fixture(None); let hidden = weights.hidden_size; engine - .try_prefill_via_dispatch(&mut weights, &index, &[0u32, 1, 2]) + .try_prefill_via_dispatch(&weights, &index, &[0u32, 1, 2]) .expect("prefill (windowless)"); let rs = engine.store.as_ref().unwrap(); assert!(rs.hot_kv.is_none()); @@ -363,8 +363,8 @@ mod tests { // LARQL_W10_DISABLE=1 → drop_hot_kv_shadow=false; both shadows // populated (Full mask in decode). set_w10_disable(true); - let (mut engine, mut weights, index) = fixture(Some(8)); - let res = engine.try_prefill_via_dispatch(&mut weights, &index, &[0u32, 1, 2]); + let (mut engine, weights, index) = fixture(Some(8)); + let res = engine.try_prefill_via_dispatch(&weights, &index, &[0u32, 1, 2]); set_w10_disable(false); res.expect("prefill"); let rs = engine.store.as_ref().unwrap(); @@ -375,10 +375,10 @@ mod tests { #[test] fn decode_step_via_dispatch_without_prefill_returns_none() { set_w10_disable(false); - let (mut engine, mut weights, index) = fixture(Some(4)); + let (mut engine, weights, index) = fixture(Some(4)); // kv_handle is None → early return at `self.kv_handle.as_mut()?`. assert!(engine - .decode_step_via_dispatch(&mut weights, &index, 0) + .decode_step_via_dispatch(&weights, &index, 0) .is_none()); } @@ -387,13 +387,13 @@ mod tests { set_w10_disable(false); // window=Some + default env → mask=HOnly. hot_len grows by 1; // hot_kv stays None. - let (mut engine, mut weights, index) = fixture(Some(8)); + let (mut engine, weights, index) = fixture(Some(8)); engine - .try_prefill_via_dispatch(&mut weights, &index, &[0u32, 1]) + .try_prefill_via_dispatch(&weights, &index, &[0u32, 1]) .expect("prefill"); let hot_len_before = engine.store.as_ref().unwrap().hot_len; let h = engine - .decode_step_via_dispatch(&mut weights, &index, 2) + .decode_step_via_dispatch(&weights, &index, 2) .expect("decode"); assert_eq!(h.shape(), &[1, weights.hidden_size]); let rs = engine.store.as_ref().unwrap(); @@ -407,13 +407,13 @@ mod tests { set_w10_disable(false); // window=None + default env → mask=None; stored stays empty, // hot_len doesn't bump. - let (mut engine, mut weights, index) = fixture(None); + let (mut engine, weights, index) = fixture(None); let hidden = weights.hidden_size; engine - .try_prefill_via_dispatch(&mut weights, &index, &[0u32, 1]) + .try_prefill_via_dispatch(&weights, &index, &[0u32, 1]) .expect("prefill (windowless)"); engine - .decode_step_via_dispatch(&mut weights, &index, 2) + .decode_step_via_dispatch(&weights, &index, 2) .expect("decode (None mask)"); let rs = engine.store.as_ref().unwrap(); for slab in &rs.stored { @@ -428,12 +428,12 @@ mod tests { fn decode_step_via_dispatch_full_mask_appends_hot_kv_with_w10_disabled() { // Full mask path: appends to BOTH stored and hot_kv on every layer. set_w10_disable(true); - let (mut engine, mut weights, index) = fixture(Some(8)); + let (mut engine, weights, index) = fixture(Some(8)); engine - .try_prefill_via_dispatch(&mut weights, &index, &[0u32, 1]) + .try_prefill_via_dispatch(&weights, &index, &[0u32, 1]) .expect("prefill"); let hot_len_before = engine.store.as_ref().unwrap().hot_len; - let result = engine.decode_step_via_dispatch(&mut weights, &index, 2); + let result = engine.decode_step_via_dispatch(&weights, &index, 2); set_w10_disable(false); result.expect("decode"); let rs = engine.store.as_ref().unwrap(); @@ -444,13 +444,13 @@ mod tests { #[test] fn decode_step_via_dispatch_with_profiling_records_stages() { set_w10_disable(false); - let (engine, mut weights, index) = fixture(Some(8)); + let (engine, weights, index) = fixture(Some(8)); let mut engine = engine.with_profiling(true); engine - .try_prefill_via_dispatch(&mut weights, &index, &[0u32, 1]) + .try_prefill_via_dispatch(&weights, &index, &[0u32, 1]) .expect("prefill"); engine - .decode_step_via_dispatch(&mut weights, &index, 2) + .decode_step_via_dispatch(&weights, &index, 2) .expect("decode"); // HOnly mask path populates state_capture / state_materialise / // state_append (the loop that runs under HOnly). diff --git a/crates/larql-kv/src/engines/markov_residual/engine.rs b/crates/larql-kv/src/engines/markov_residual/engine.rs index 565e1ac61..6c2cea688 100644 --- a/crates/larql-kv/src/engines/markov_residual/engine.rs +++ b/crates/larql-kv/src/engines/markov_residual/engine.rs @@ -27,6 +27,12 @@ pub struct MarkovResidualEngine { /// Position counter used by `coarse_decode_step_with_state` for RoPE. /// Tracks `prompt_len + steps_already_decoded`. pub(super) abs_position: usize, + /// Engine-owned f32 dequant scratch for the per-layer walk fallback — + /// `prefill_quant`/`decode_step_quant` populate it; the walk resolves + /// attention through a `WeightsView::with_scratch` over it. Empty on the + /// W1-GPU/coarse path. Keeps `weights` immutable (no `weights.tensors` + /// mutation). + pub(super) dequant_scratch: larql_inference::DequantScratch, } impl MarkovResidualEngine { @@ -43,6 +49,7 @@ impl MarkovResidualEngine { profile: EngineProfiler::default(), kv_handle: None, abs_position: 0, + dequant_scratch: larql_inference::DequantScratch::new(), } } @@ -79,7 +86,7 @@ impl MarkovResidualEngine { })?; let (hidden, new_rs) = if self.profiling { rs_decode_step_profiled( - weights, + larql_inference::WeightsView::dense(weights), token_id, rs, self.backend.as_ref(), @@ -92,7 +99,7 @@ impl MarkovResidualEngine { })? } else { rs_decode_step( - weights, + larql_inference::WeightsView::dense(weights), token_id, rs, self.backend.as_ref(), @@ -140,7 +147,7 @@ impl KvEngine for MarkovResidualEngine { return Err(EngineError::EmptyPrompt); } let result = rs_prefill( - weights, + larql_inference::WeightsView::dense(weights), token_ids, self.window_size, self.backend.as_ref(), @@ -195,7 +202,7 @@ impl KvEngine for MarkovResidualEngine { fn prefill_quant( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, _ffn: &dyn FfnBackend, index: &VectorIndex, token_ids: &[u32], @@ -214,8 +221,9 @@ impl KvEngine for MarkovResidualEngine { if let Some(hidden) = self.try_prefill_via_dispatch(weights, index, token_ids) { return Ok(hidden); } - ensure_attn_tensors_dequantised(weights, index); - let result = rs_prefill_walk(weights, index, token_ids, self.window_size, backend); + ensure_attn_tensors_dequantised(&mut self.dequant_scratch, weights, index); + let view = larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch); + let result = rs_prefill_walk(view, index, token_ids, self.window_size, backend); let hidden = result.hidden.clone(); self.store = Some(result.store); self.kv_handle = None; // ensure dispatch path is not used for subsequent decode @@ -225,7 +233,7 @@ impl KvEngine for MarkovResidualEngine { fn decode_step_quant( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, _ffn: &dyn FfnBackend, index: &VectorIndex, token_id: u32, @@ -241,7 +249,8 @@ impl KvEngine for MarkovResidualEngine { details: "decode_step_via_dispatch returned None".into(), }); } - ensure_attn_tensors_dequantised(weights, index); + ensure_attn_tensors_dequantised(&mut self.dequant_scratch, weights, index); + let view = larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch); let rs = self .store .take() @@ -249,7 +258,7 @@ impl KvEngine for MarkovResidualEngine { what: "decode_step_quant called before prefill (store missing)".into(), })?; let prof = self.profiling.then_some(&mut self.profile); - let (hidden, new_rs) = rs_decode_step_walk(weights, index, token_id, rs, backend, prof) + let (hidden, new_rs) = rs_decode_step_walk(view, index, token_id, rs, backend, prof) .ok_or_else(|| EngineError::BackendFailure { details: "rs_decode_step_walk returned None".into(), })?; @@ -270,7 +279,7 @@ impl KvEngine for MarkovResidualEngine { fn prefill_quant_via_executor( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, executor: &dyn larql_inference::layer_executor::LayerExecutor, ffn: &dyn FfnBackend, index: &VectorIndex, @@ -294,7 +303,7 @@ impl KvEngine for MarkovResidualEngine { // Q4K attn weights need dequant once before the per-layer // executor can drive f32 attention against them. - ensure_attn_tensors_dequantised(weights, index); + ensure_attn_tensors_dequantised(&mut self.dequant_scratch, weights, index); let backend = executor.backend(); let num_layers = weights.num_layers; @@ -309,7 +318,12 @@ impl KvEngine for MarkovResidualEngine { // the layer's K/V — residual-stream contract recomputes K/V // per decode step from the stored residuals. let (h_out, _kv) = executor - .run_prefill_layer(weights, layer, &h, ffn) + .run_prefill_layer( + larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch), + layer, + &h, + ffn, + ) .ok_or_else(|| EngineError::BackendFailure { details: "executor.run_prefill_layer returned None".into(), })?; @@ -342,8 +356,15 @@ impl KvEngine for MarkovResidualEngine { 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, Some(index)) - .expect("cold K/V pre-computation failed") + recompute_kv( + larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch), + &cold[layer], + layer, + 0, + backend, + Some(index), + ) + .expect("cold K/V pre-computation failed") }) .collect(); // 2026-05-19 audit fix: doubling-capacity append. @@ -362,7 +383,7 @@ impl KvEngine for MarkovResidualEngine { fn decode_step_quant_via_executor( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, executor: &dyn larql_inference::layer_executor::LayerExecutor, ffn: &dyn FfnBackend, index: &VectorIndex, @@ -378,7 +399,7 @@ impl KvEngine for MarkovResidualEngine { return self.decode_step_quant(weights, ffn, index, token_id, executor.backend()); } - ensure_attn_tensors_dequantised(weights, index); + ensure_attn_tensors_dequantised(&mut self.dequant_scratch, weights, index); let backend = executor.backend(); let rs = self @@ -402,11 +423,17 @@ impl KvEngine for MarkovResidualEngine { // and runs the layer. let prior_kv: 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, Some(index)) - .ok_or_else(|| EngineError::BackendFailure { - details: "recompute_kv (hot) returned None".into(), - })?; + let (k_hot, v_hot) = recompute_kv( + larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch), + h_hot, + layer, + hot_abs_start, + backend, + Some(index), + ) + .ok_or_else(|| EngineError::BackendFailure { + details: "recompute_kv (hot) returned None".into(), + })?; let c = k_cold.shape()[0]; let kv_dim = k_cold.shape()[1]; let mut k_combined = Array2::::zeros((c + s_hot, kv_dim)); @@ -430,7 +457,7 @@ impl KvEngine for MarkovResidualEngine { _ => (h_hot.clone(), hot_abs_start), }; recompute_kv( - weights, + larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch), &h_full, layer, full_abs_start, @@ -445,7 +472,14 @@ impl KvEngine for MarkovResidualEngine { new_stored.push(h_new.clone()); // Run the layer through the executor. let (h_out, _new_kv) = executor - .run_decode_layer(weights, layer, &h_new, &prior_kv, abs_position, ffn) + .run_decode_layer( + larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch), + layer, + &h_new, + &prior_kv, + abs_position, + ffn, + ) .ok_or_else(|| EngineError::BackendFailure { details: "executor.run_decode_layer returned None".into(), })?; @@ -686,7 +720,7 @@ mod tests { #[test] fn prefill_q4k_cpu_fallback_runs_walk_path() { use larql_inference::ffn::NullFfn; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); // `NullFfn` satisfies the trait without borrowing `weights`, which is @@ -694,7 +728,7 @@ mod tests { let ffn = NullFfn; let mut engine = MarkovResidualEngine::new(None); let h = engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1, 2], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1, 2], &*backend) .expect("prefill_quant cpu fallback"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert!(engine.memory_bytes() > 0); @@ -703,17 +737,17 @@ mod tests { #[test] fn decode_step_q4k_cpu_fallback_extends_store() { use larql_inference::ffn::NullFfn; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = MarkovResidualEngine::new(None); engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1], &*backend) .expect("prefill_quant"); let mem_before = engine.memory_bytes(); let h = engine - .decode_step_quant(&mut weights, &ffn, &index, 2, &*backend) + .decode_step_quant(&weights, &ffn, &index, 2, &*backend) .expect("decode_step_quant cpu fallback"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert!( @@ -732,13 +766,13 @@ mod tests { #[test] fn prefill_quant_walk_with_window_populates_cold_kv() { use larql_inference::ffn::NullFfn; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = MarkovResidualEngine::new(Some(2)); engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1, 2, 3], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1, 2, 3], &*backend) .expect("prefill_quant with overflow"); // window=2 + 4 prompt tokens → cold tier populated → walk.rs // lines 67-75 fire. @@ -754,7 +788,7 @@ mod tests { #[test] fn decode_step_quant_w2_cached_matches_recompute_from_residuals() { use larql_inference::ffn::NullFfn; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; @@ -762,32 +796,32 @@ mod tests { // Cached path (W2 default): prefill captures K/V, decode reuses. let mut cached = MarkovResidualEngine::new(None); cached - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1, 2], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1, 2], &*backend) .expect("prefill cached"); let h_cached_1 = cached - .decode_step_quant(&mut weights, &ffn, &index, 3, &*backend) + .decode_step_quant(&weights, &ffn, &index, 3, &*backend) .expect("decode cached 1"); let h_cached_2 = cached - .decode_step_quant(&mut weights, &ffn, &index, 4, &*backend) + .decode_step_quant(&weights, &ffn, &index, 4, &*backend) .expect("decode cached 2"); // Recompute path: same engine, but force hot_kv = None after // prefill so the fallback recompute fires for every step. let mut recompute = MarkovResidualEngine::new(None); recompute - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1, 2], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1, 2], &*backend) .expect("prefill recompute"); if let Some(s) = recompute.store.as_mut() { s.hot_kv = None; } let h_recompute_1 = recompute - .decode_step_quant(&mut weights, &ffn, &index, 3, &*backend) + .decode_step_quant(&weights, &ffn, &index, 3, &*backend) .expect("decode recompute 1"); if let Some(s) = recompute.store.as_mut() { s.hot_kv = None; } let h_recompute_2 = recompute - .decode_step_quant(&mut weights, &ffn, &index, 4, &*backend) + .decode_step_quant(&weights, &ffn, &index, 4, &*backend) .expect("decode recompute 2"); // Bit-equivalence: both paths run the same projection matmuls @@ -817,7 +851,7 @@ mod tests { #[test] fn decode_step_quant_w2_cached_hot_and_cold_steady_state() { use larql_inference::ffn::NullFfn; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; @@ -825,7 +859,7 @@ mod tests { // populating cold_kv from the evicted hot_kv slice. let mut engine = MarkovResidualEngine::new(Some(2)); engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1, 2, 3], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1, 2, 3], &*backend) .expect("prefill with overflow"); let store = engine.store.as_ref().unwrap(); assert!(store.hot_kv.is_some()); @@ -836,7 +870,7 @@ mod tests { // merge into cold_kv via the W2 evicted-K/V flow. for tok in 4u32..8 { let h = engine - .decode_step_quant(&mut weights, &ffn, &index, tok, &*backend) + .decode_step_quant(&weights, &ffn, &index, tok, &*backend) .expect("decode"); assert_eq!(h.shape(), &[1, weights.hidden_size]); } @@ -861,18 +895,18 @@ mod tests { #[test] fn decode_step_quant_w2_falls_back_when_hot_kv_dropped() { use larql_inference::ffn::NullFfn; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = MarkovResidualEngine::new(Some(2)); engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1, 2, 3], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1, 2, 3], &*backend) .expect("prefill"); // Drop hot_kv — forces the recompute path that mirrors pre-W2. engine.store.as_mut().unwrap().hot_kv = None; let h = engine - .decode_step_quant(&mut weights, &ffn, &index, 4, &*backend) + .decode_step_quant(&weights, &ffn, &index, 4, &*backend) .expect("decode via fallback"); assert_eq!(h.shape(), &[1, weights.hidden_size]); } @@ -883,14 +917,14 @@ mod tests { #[test] fn decode_step_quant_w2_overflow_merges_into_cold_kv() { use larql_inference::ffn::NullFfn; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = MarkovResidualEngine::new(Some(2)); engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1], &*backend) .expect("prefill within window"); // After prefill: hot_kv populated (2 rows), no cold_kv. assert!(engine.store.as_ref().unwrap().hot_kv.is_some()); @@ -898,7 +932,7 @@ mod tests { // Decode a token → no overflow yet (still 2 rows after step // since window=2, the new row pushes the oldest out). let _ = engine - .decode_step_quant(&mut weights, &ffn, &index, 2, &*backend) + .decode_step_quant(&weights, &ffn, &index, 2, &*backend) .expect("decode 1"); // Overflow fired this step: oldest row evicted from hot_kv, // merged into cold_kv. @@ -919,21 +953,21 @@ mod tests { #[test] fn decode_step_q4k_walk_with_profiling_populates_summary() { use larql_inference::ffn::NullFfn; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = MarkovResidualEngine::new(Some(2)).with_profiling(true); engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1, 2, 3], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1, 2, 3], &*backend) .expect("prefill"); // First decode: cold_kv branch (hot recompute timing arm). engine - .decode_step_quant(&mut weights, &ffn, &index, 4, &*backend) + .decode_step_quant(&weights, &ffn, &index, 4, &*backend) .expect("decode 1"); // Second decode: cold_residuals branch (cold recompute timing arm). engine - .decode_step_quant(&mut weights, &ffn, &index, 5, &*backend) + .decode_step_quant(&weights, &ffn, &index, 5, &*backend) .expect("decode 2"); let summary = engine .stage_summary() @@ -954,21 +988,21 @@ mod tests { // Some(overflow)`. Fires when prefill didn't overflow (cold = None) // but the first decode does (window cap exceeded mid-decode). use larql_inference::ffn::NullFfn; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; // window=2, prefill=1 token → no overflow on prefill (cold=None). let mut engine = MarkovResidualEngine::new(Some(2)); engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32], &*backend) .expect("prefill_quant"); // Decode until hot exceeds window → first-time cold population. engine - .decode_step_quant(&mut weights, &ffn, &index, 1, &*backend) + .decode_step_quant(&weights, &ffn, &index, 1, &*backend) .expect("decode 1"); engine - .decode_step_quant(&mut weights, &ffn, &index, 2, &*backend) + .decode_step_quant(&weights, &ffn, &index, 2, &*backend) .expect("decode 2 — triggers first-overflow None branch"); // After overflow, cold tier is populated. assert!(engine.cold_bytes() > 0); @@ -977,23 +1011,23 @@ mod tests { #[test] fn decode_step_quant_walk_after_overflow_hits_cold_residuals_branch() { use larql_inference::ffn::NullFfn; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = MarkovResidualEngine::new(Some(2)); engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1, 2, 3], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1, 2, 3], &*backend) .expect("prefill_quant"); // First decode: exercises walk.rs cold_kv branch (lines 132-161). engine - .decode_step_quant(&mut weights, &ffn, &index, 4, &*backend) + .decode_step_quant(&weights, &ffn, &index, 4, &*backend) .expect("first decode_step_quant"); // Second decode: cold_kv was cleared by overflow at the first // decode (walk.rs line 309), so this hits the cold_residuals // recompute branch (lines 162-187). let h = engine - .decode_step_quant(&mut weights, &ffn, &index, 5, &*backend) + .decode_step_quant(&weights, &ffn, &index, 5, &*backend) .expect("second decode_step_quant"); assert_eq!(h.shape(), &[1, weights.hidden_size]); } @@ -1004,14 +1038,14 @@ mod tests { fn prefill_quant_via_executor_runs_through_local_walk() { use larql_inference::ffn::NullFfn; use larql_inference::layer_executor::LocalWalkExecutor; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let executor = LocalWalkExecutor::new(&*backend); let ffn = NullFfn; let mut engine = MarkovResidualEngine::new(None); let h = engine - .prefill_quant_via_executor(&mut weights, &executor, &ffn, &index, &[0u32, 1, 2]) + .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32, 1, 2]) .expect("executor prefill"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert!(engine.memory_bytes() > 0); @@ -1021,18 +1055,18 @@ mod tests { fn decode_step_quant_via_executor_extends_store() { use larql_inference::ffn::NullFfn; use larql_inference::layer_executor::LocalWalkExecutor; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let executor = LocalWalkExecutor::new(&*backend); let ffn = NullFfn; let mut engine = MarkovResidualEngine::new(None); engine - .prefill_quant_via_executor(&mut weights, &executor, &ffn, &index, &[0u32, 1]) + .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32, 1]) .expect("prefill"); let mem_before = engine.memory_bytes(); let h = engine - .decode_step_quant_via_executor(&mut weights, &executor, &ffn, &index, 2) + .decode_step_quant_via_executor(&weights, &executor, &ffn, &index, 2) .expect("decode"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert!(engine.memory_bytes() > mem_before); @@ -1070,7 +1104,7 @@ mod tests { // WalkFfn internally (the legacy bug we're fixing) the counter // stays at zero. use larql_inference::layer_executor::LocalWalkExecutor; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let executor = LocalWalkExecutor::new(&*backend); @@ -1081,7 +1115,7 @@ mod tests { }; let mut engine = MarkovResidualEngine::new(None); engine - .prefill_quant_via_executor(&mut weights, &executor, &ffn, &index, &[0u32, 1, 2]) + .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32, 1, 2]) .expect("prefill via executor"); let call_count = ffn.calls.load(std::sync::atomic::Ordering::SeqCst); @@ -1099,14 +1133,14 @@ mod tests { fn prefill_quant_via_executor_with_window_populates_cold_tier() { use larql_inference::ffn::NullFfn; use larql_inference::layer_executor::LocalWalkExecutor; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let executor = LocalWalkExecutor::new(&*backend); let ffn = NullFfn; let mut engine = MarkovResidualEngine::new(Some(2)); engine - .prefill_quant_via_executor(&mut weights, &executor, &ffn, &index, &[0u32, 1, 2, 3]) + .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32, 1, 2, 3]) .expect("prefill with overflow"); assert!(engine.window_tokens() <= 2); assert!(engine.cold_bytes() > 0); @@ -1120,18 +1154,18 @@ mod tests { fn decode_step_quant_via_executor_uses_cold_kv_branch() { use larql_inference::ffn::NullFfn; use larql_inference::layer_executor::LocalWalkExecutor; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let executor = LocalWalkExecutor::new(&*backend); let ffn = NullFfn; let mut engine = MarkovResidualEngine::new(Some(2)); engine - .prefill_quant_via_executor(&mut weights, &executor, &ffn, &index, &[0u32, 1, 2, 3]) + .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32, 1, 2, 3]) .expect("prefill overflow → cold_kv populated"); // First decode reads cold_kv branch (rs.cold_kv = Some(_)). let h = engine - .decode_step_quant_via_executor(&mut weights, &executor, &ffn, &index, 4) + .decode_step_quant_via_executor(&weights, &executor, &ffn, &index, 4) .expect("decode via cold_kv branch"); assert_eq!(h.shape(), &[1, weights.hidden_size]); } @@ -1144,23 +1178,23 @@ mod tests { fn decode_step_quant_via_executor_hits_cold_residuals_branch() { use larql_inference::ffn::NullFfn; use larql_inference::layer_executor::LocalWalkExecutor; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let executor = LocalWalkExecutor::new(&*backend); let ffn = NullFfn; let mut engine = MarkovResidualEngine::new(Some(2)); engine - .prefill_quant_via_executor(&mut weights, &executor, &ffn, &index, &[0u32, 1, 2, 3]) + .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32, 1, 2, 3]) .expect("prefill"); // First decode clears cold_kv via overflow. engine - .decode_step_quant_via_executor(&mut weights, &executor, &ffn, &index, 4) + .decode_step_quant_via_executor(&weights, &executor, &ffn, &index, 4) .expect("first decode"); // Second decode: cold_kv is None, exercises the recompute_kv // from cold_residuals branch. let h = engine - .decode_step_quant_via_executor(&mut weights, &executor, &ffn, &index, 5) + .decode_step_quant_via_executor(&weights, &executor, &ffn, &index, 5) .expect("decode via cold_residuals recompute"); assert_eq!(h.shape(), &[1, weights.hidden_size]); } @@ -1186,7 +1220,7 @@ mod tests { #[test] fn fused_executor_falls_back_to_legacy_quant_path() { use larql_inference::ffn::NullFfn; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let exec = FusedStubExecutor { backend: larql_compute::CpuBackend, @@ -1194,11 +1228,11 @@ mod tests { let ffn = NullFfn; let mut engine = MarkovResidualEngine::new(None); let h = engine - .prefill_quant_via_executor(&mut weights, &exec, &ffn, &index, &[0u32, 1]) + .prefill_quant_via_executor(&weights, &exec, &ffn, &index, &[0u32, 1]) .expect("fused fallback prefill"); assert_eq!(h.shape(), &[1, weights.hidden_size]); let h2 = engine - .decode_step_quant_via_executor(&mut weights, &exec, &ffn, &index, 2) + .decode_step_quant_via_executor(&weights, &exec, &ffn, &index, 2) .expect("fused fallback decode"); assert_eq!(h2.shape(), &[1, weights.hidden_size]); } @@ -1222,13 +1256,13 @@ mod tests { } use larql_inference::ffn::NullFfn; use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = MarkovResidualEngine::new(None); let h = engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1, 2], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1, 2], &*backend) .expect("dispatch-path prefill on Q4K vindex"); assert_eq!(h.shape(), &[1, weights.hidden_size]); // Dispatch path populates kv_handle + abs_position. @@ -1248,18 +1282,18 @@ mod tests { fn decode_via_dispatch_grows_buffers_in_place() { use larql_inference::ffn::NullFfn; use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = MarkovResidualEngine::new(None); engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1], &*backend) .expect("prefill dispatch"); let mem_after_prefill = engine.memory_bytes(); for tok in 2..6u32 { let h = engine - .decode_step_quant(&mut weights, &ffn, &index, tok, &*backend) + .decode_step_quant(&weights, &ffn, &index, tok, &*backend) .expect("dispatch decode step"); assert_eq!(h.shape(), &[1, weights.hidden_size]); } @@ -1278,19 +1312,19 @@ mod tests { } use larql_inference::ffn::NullFfn; use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = MarkovResidualEngine::new(None).with_profiling(true); engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1], &*backend) .expect("prefill"); engine - .decode_step_quant(&mut weights, &ffn, &index, 2, &*backend) + .decode_step_quant(&weights, &ffn, &index, 2, &*backend) .expect("decode 1"); engine - .decode_step_quant(&mut weights, &ffn, &index, 3, &*backend) + .decode_step_quant(&weights, &ffn, &index, 3, &*backend) .expect("decode 2"); // W10 instrumentation: dispatch path bumps state_capture + // state_materialise + state_append + decode_total on every step. @@ -1319,13 +1353,13 @@ mod tests { // populated from snapshot_evicted_hot_kv). use larql_inference::ffn::NullFfn; use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = MarkovResidualEngine::new(Some(2)); engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1, 2, 3], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1, 2, 3], &*backend) .expect("prefill with window"); let store = engine.store.as_ref().expect("store"); assert!(store.cold_residuals.is_some()); @@ -1342,17 +1376,17 @@ mod tests { // K/V already exists and the eviction concatenates. use larql_inference::ffn::NullFfn; use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = MarkovResidualEngine::new(Some(2)); engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1, 2], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1, 2], &*backend) .expect("prefill at window cap"); for tok in 3..6u32 { engine - .decode_step_quant(&mut weights, &ffn, &index, tok, &*backend) + .decode_step_quant(&weights, &ffn, &index, tok, &*backend) .expect("dispatch decode with overflow"); } let store = engine.store.as_ref().expect("store"); @@ -1405,13 +1439,13 @@ mod tests { fn prefill_quant_returns_empty_prompt_error_on_empty_input() { use larql_inference::ffn::NullFfn; use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = MarkovResidualEngine::new(None); let err = engine - .prefill_quant(&mut weights, &ffn, &index, &[], &*backend) + .prefill_quant(&weights, &ffn, &index, &[], &*backend) .unwrap_err(); assert_eq!(err, larql_inference::kv_engine::EngineError::EmptyPrompt); } @@ -1420,13 +1454,13 @@ mod tests { fn decode_step_quant_returns_invariant_violation_before_prefill() { use larql_inference::ffn::NullFfn; use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = MarkovResidualEngine::new(None); let err = engine - .decode_step_quant(&mut weights, &ffn, &index, 0, &*backend) + .decode_step_quant(&weights, &ffn, &index, 0, &*backend) .unwrap_err(); assert!(matches!( err, @@ -1452,14 +1486,14 @@ mod tests { fn prefill_quant_via_executor_returns_empty_prompt_error_on_empty_input() { use larql_inference::ffn::NullFfn; use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = larql_compute::CpuBackend; let executor = larql_inference::layer_executor::LocalWalkExecutor::new(&backend); let ffn = NullFfn; let mut engine = MarkovResidualEngine::new(None); let err = engine - .prefill_quant_via_executor(&mut weights, &executor, &ffn, &index, &[]) + .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[]) .unwrap_err(); assert_eq!(err, larql_inference::kv_engine::EngineError::EmptyPrompt); } @@ -1468,14 +1502,14 @@ mod tests { fn decode_step_quant_via_executor_returns_invariant_violation_before_prefill() { use larql_inference::ffn::NullFfn; use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = larql_compute::CpuBackend; let executor = larql_inference::layer_executor::LocalWalkExecutor::new(&backend); let ffn = NullFfn; let mut engine = MarkovResidualEngine::new(None); let err = engine - .decode_step_quant_via_executor(&mut weights, &executor, &ffn, &index, 0) + .decode_step_quant_via_executor(&weights, &executor, &ffn, &index, 0) .unwrap_err(); assert!(matches!( err, diff --git a/crates/larql-kv/src/engines/markov_residual/walk.rs b/crates/larql-kv/src/engines/markov_residual/walk.rs index 612519a2d..48a524fa4 100644 --- a/crates/larql-kv/src/engines/markov_residual/walk.rs +++ b/crates/larql-kv/src/engines/markov_residual/walk.rs @@ -17,7 +17,6 @@ use crate::profiler::EngineProfiler; use larql_inference::attention::run_attention_with_kv_backend; use larql_inference::attention::SharedKV; use larql_inference::forward::{embed_tokens_pub, run_ffn}; -use larql_inference::model::ModelWeights; use larql_inference::vindex::{WalkFfn, WalkFfnConfig}; /// Re-export — see [`larql_inference::vindex::dequant::ensure_attn_tensors_dequantised`]. @@ -25,7 +24,7 @@ pub use larql_inference::vindex::ensure_attn_tensors_dequantised; /// Prefill using `WalkFfn` (Q4K FFN) instead of `BackendFfn` (f32 FFN). pub(super) fn rs_prefill_walk( - weights: &ModelWeights, + weights: larql_inference::WeightsView, index: &VectorIndex, token_ids: &[u32], max_window: Option, @@ -33,7 +32,7 @@ pub(super) fn rs_prefill_walk( ) -> RsPrefillResult { let num_layers = weights.num_layers; let seq_len = token_ids.len(); - let mut h = embed_tokens_pub(weights, token_ids); + let mut h = embed_tokens_pub(&weights, token_ids); let mut stored: Vec> = Vec::with_capacity(num_layers); let be = Some(backend); @@ -41,8 +40,9 @@ pub(super) fn rs_prefill_walk( // this rebuilt the WalkFfn 34 times per prefill (once per layer); // now once total. WalkFfn carries no per-layer state — it's the // gate-index + backend pair, both stable across the loop. - let walk_ffn = WalkFfn::from_config(weights, index, WalkFfnConfig::dense(num_layers)) - .with_backend(backend); + let walk_ffn = + WalkFfn::from_config(weights.canonical(), index, WalkFfnConfig::dense(num_layers)) + .with_backend(backend); // Capture per-layer K/V from each layer's attention block. These // are *already computed* by the forward pass; previously discarded @@ -53,10 +53,10 @@ pub(super) fn rs_prefill_walk( let mut hot_kv_captured: Vec = Vec::with_capacity(num_layers); 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) + let (h_post_attn, k, v) = run_attention_with_kv_backend(weights, &h, layer, be, None) .expect("attention failed during MarkovRS Q4K prefill"); hot_kv_captured.push((k, v)); - let (h_out, _) = run_ffn(weights, &h_post_attn, layer, &walk_ffn, false); + let (h_out, _) = run_ffn(&weights, &h_post_attn, layer, &walk_ffn, false); h = h_out; } @@ -133,7 +133,7 @@ pub(super) fn rs_prefill_walk( /// path. Sibling of [`super::compute::rs_decode_step_inner`] for the /// Q4K side. pub(super) fn rs_decode_step_walk( - weights: &ModelWeights, + weights: larql_inference::WeightsView, index: &VectorIndex, new_token_id: u32, rs: RsStore, @@ -154,7 +154,7 @@ pub(super) fn rs_decode_step_walk( let abs_position = rs.next_position; let t_step = if timing { Some(Instant::now()) } else { None }; let t_embed_start = t_step; - let mut h_new = embed_tokens_pub(weights, &[new_token_id]); + let mut h_new = embed_tokens_pub(&weights, &[new_token_id]); let embed_us = t_embed_start .map(|t| t.elapsed().as_secs_f64() * 1e6) .unwrap_or(0.0); @@ -162,8 +162,9 @@ pub(super) fn rs_decode_step_walk( // Hoist WalkFfn out of the per-layer loop — see note in // `rs_prefill_walk`. Was 34× construction per decode step. - let walk_ffn = WalkFfn::from_config(weights, index, WalkFfnConfig::dense(num_layers)) - .with_backend(backend); + let walk_ffn = + WalkFfn::from_config(weights.canonical(), index, WalkFfnConfig::dense(num_layers)) + .with_backend(backend); // Per-stage accumulators. With W2 caching landed, both // `recompute_*` timings should be near zero for cached-path decode @@ -280,7 +281,7 @@ pub(super) fn rs_decode_step_walk( let t_attn = if timing { Some(Instant::now()) } else { None }; let kv_pair = (k_full, v_full); let native_result = larql_inference::vindex::attention_decode_step_native( - weights, + weights.canonical(), index, backend, &h_new, @@ -328,14 +329,14 @@ pub(super) fn rs_decode_step_walk( // the backend doesn't have native quant support or the layer // isn't direct-matvec-eligible. let h_out = larql_inference::vindex::ffn_decode_step_native( - weights, + weights.canonical(), index, backend, &h_post_attn, layer, ) .unwrap_or_else(|| { - let (h, _) = run_ffn(weights, &h_post_attn, layer, &walk_ffn, false); + let (h, _) = run_ffn(&weights, &h_post_attn, layer, &walk_ffn, false); h }); if let Some(t) = t_ffn { @@ -451,7 +452,13 @@ mod tests { fn prefill_walk_returns_finite_hidden_and_full_window_store() { let weights = make_test_weights(); let index = make_test_vindex(&weights); - let result = rs_prefill_walk(&weights, &index, &[0u32, 1, 2], None, &CpuBackend); + let result = rs_prefill_walk( + larql_inference::WeightsView::dense(&weights), + &index, + &[0u32, 1, 2], + None, + &CpuBackend, + ); assert_eq!(result.hidden.shape(), &[1, weights.hidden_size]); assert!(result.hidden.iter().all(|v| v.is_finite())); assert!(result.store.cold_residuals.is_none()); @@ -465,7 +472,13 @@ mod tests { fn prefill_walk_with_overflow_populates_cold_tier_from_evicted_hot_kv() { let weights = make_test_weights(); let index = make_test_vindex(&weights); - let result = rs_prefill_walk(&weights, &index, &[0u32, 1, 2, 3], Some(2), &CpuBackend); + let result = rs_prefill_walk( + larql_inference::WeightsView::dense(&weights), + &index, + &[0u32, 1, 2, 3], + Some(2), + &CpuBackend, + ); assert!(result.store.cold_residuals.is_some()); assert!(result.store.cold_kv.is_some()); // Window-clipped, but cold tier captured the two evicted rows. @@ -479,10 +492,23 @@ mod tests { fn decode_walk_extends_position_and_returns_finite() { let weights = make_test_weights(); let index = make_test_vindex(&weights); - let prefill = rs_prefill_walk(&weights, &index, &[0u32, 1], None, &CpuBackend); + let prefill = rs_prefill_walk( + larql_inference::WeightsView::dense(&weights), + &index, + &[0u32, 1], + None, + &CpuBackend, + ); assert_eq!(prefill.store.next_position, 2); - let (h, rs2) = - rs_decode_step_walk(&weights, &index, 2, prefill.store, &CpuBackend, None).unwrap(); + let (h, rs2) = rs_decode_step_walk( + larql_inference::WeightsView::dense(&weights), + &index, + 2, + prefill.store, + &CpuBackend, + None, + ) + .unwrap(); assert_eq!(rs2.next_position, 3); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert!(h.iter().all(|v| v.is_finite())); @@ -495,10 +521,16 @@ mod tests { // the timing accumulators on every per-stage block. let weights = make_test_weights(); let index = make_test_vindex(&weights); - let prefill = rs_prefill_walk(&weights, &index, &[0u32, 1, 2, 3], Some(2), &CpuBackend); + let prefill = rs_prefill_walk( + larql_inference::WeightsView::dense(&weights), + &index, + &[0u32, 1, 2, 3], + Some(2), + &CpuBackend, + ); let mut prof = EngineProfiler::default(); let (h, _) = rs_decode_step_walk( - &weights, + larql_inference::WeightsView::dense(&weights), &index, 4, prefill.store, @@ -522,12 +554,25 @@ mod tests { // from h_hot and concat with the cached cold K/V. let weights = make_test_weights(); let index = make_test_vindex(&weights); - let prefill = rs_prefill_walk(&weights, &index, &[0u32, 1, 2, 3], Some(2), &CpuBackend); + let prefill = rs_prefill_walk( + larql_inference::WeightsView::dense(&weights), + &index, + &[0u32, 1, 2, 3], + Some(2), + &CpuBackend, + ); let mut store = prefill.store; store.hot_kv = None; let mut prof = EngineProfiler::default(); - let (h, rs2) = - rs_decode_step_walk(&weights, &index, 4, store, &CpuBackend, Some(&mut prof)).unwrap(); + let (h, rs2) = rs_decode_step_walk( + larql_inference::WeightsView::dense(&weights), + &index, + 4, + store, + &CpuBackend, + Some(&mut prof), + ) + .unwrap(); assert_eq!(h.shape(), &[1, weights.hidden_size]); // hot_kv is repopulated on every decode step. assert!(rs2.hot_kv.is_some()); @@ -541,11 +586,25 @@ mod tests { // recomputing the K/V. let weights = make_test_weights(); let index = make_test_vindex(&weights); - let prefill = rs_prefill_walk(&weights, &index, &[0u32, 1, 2, 3], Some(2), &CpuBackend); + let prefill = rs_prefill_walk( + larql_inference::WeightsView::dense(&weights), + &index, + &[0u32, 1, 2, 3], + Some(2), + &CpuBackend, + ); let mut store = prefill.store; store.hot_kv = None; store.cold_kv = None; - let (h, _) = rs_decode_step_walk(&weights, &index, 4, store, &CpuBackend, None).unwrap(); + let (h, _) = rs_decode_step_walk( + larql_inference::WeightsView::dense(&weights), + &index, + 4, + store, + &CpuBackend, + None, + ) + .unwrap(); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert!(h.iter().all(|v| v.is_finite())); } @@ -558,10 +617,23 @@ mod tests { // initialised from the evicted block. let weights = make_test_weights(); let index = make_test_vindex(&weights); - let prefill = rs_prefill_walk(&weights, &index, &[0u32, 1], Some(2), &CpuBackend); + let prefill = rs_prefill_walk( + larql_inference::WeightsView::dense(&weights), + &index, + &[0u32, 1], + Some(2), + &CpuBackend, + ); assert!(prefill.store.cold_residuals.is_none()); - let (_, rs2) = - rs_decode_step_walk(&weights, &index, 2, prefill.store, &CpuBackend, None).unwrap(); + let (_, rs2) = rs_decode_step_walk( + larql_inference::WeightsView::dense(&weights), + &index, + 2, + prefill.store, + &CpuBackend, + None, + ) + .unwrap(); assert!(rs2.cold_residuals.is_some()); // 2026-05-19 audit fix: shape()[0] is doubling capacity. Use cold_len. assert_eq!(rs2.cold_len, 1); diff --git a/crates/larql-kv/src/engines/markov_residual_codec/compute.rs b/crates/larql-kv/src/engines/markov_residual_codec/compute.rs index 234b9e53d..d94b34ba3 100644 --- a/crates/larql-kv/src/engines/markov_residual_codec/compute.rs +++ b/crates/larql-kv/src/engines/markov_residual_codec/compute.rs @@ -9,7 +9,6 @@ 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::model::ModelWeights; use ndarray::{s, Array2}; use crate::engines::markov_residual::recompute_kv; @@ -24,7 +23,7 @@ pub struct RsPrefillResultCodec { #[allow(clippy::too_many_arguments)] pub fn rs_prefill_codec( - weights: &ModelWeights, + weights: larql_inference::WeightsView, token_ids: &[u32], max_window: Option, codec: ColdResidualCodec, @@ -33,16 +32,25 @@ pub fn rs_prefill_codec( ) -> RsPrefillResultCodec { let num_layers = weights.num_layers; let seq_len = token_ids.len(); - let mut h = embed_tokens_pub(weights, token_ids); + let mut h = embed_tokens_pub(&weights, 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) + 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, backend }; - let h_out = crate::engines::layer_ffn_or_moe(weights, &h_post_attn, layer, &bffn, moe_ffn); + 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, + ); h = h_out; } @@ -93,7 +101,7 @@ pub fn rs_prefill_codec( } pub fn rs_decode_step_codec( - weights: &ModelWeights, + weights: larql_inference::WeightsView, new_token_id: u32, rs: RsStoreCodec, backend: &dyn ComputeBackend, @@ -102,7 +110,7 @@ pub fn rs_decode_step_codec( ) -> 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]); + let mut h_new = embed_tokens_pub(&weights, &[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 @@ -264,8 +272,17 @@ pub fn rs_decode_step_codec( h_post_attn }; - let bffn = BackendFfn { weights, backend }; - let h_out = crate::engines::layer_ffn_or_moe(weights, &h_post_attn, layer, &bffn, moe_ffn); + 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, + ); h_new = h_out; } @@ -374,7 +391,7 @@ mod tests { fn prefill_returns_finite_hidden() { let weights = make_test_weights(); let result = rs_prefill_codec( - &weights, + larql_inference::WeightsView::dense(&weights), &[0u32, 1, 2], None, ColdResidualCodec::Bf16, @@ -389,7 +406,7 @@ mod tests { fn prefill_no_window_does_not_create_cold_tier() { let weights = make_test_weights(); let result = rs_prefill_codec( - &weights, + larql_inference::WeightsView::dense(&weights), &[0u32, 1], None, ColdResidualCodec::Bf16, @@ -404,7 +421,7 @@ mod tests { fn prefill_with_overflow_creates_encoded_cold_tier() { let weights = make_test_weights(); let result = rs_prefill_codec( - &weights, + larql_inference::WeightsView::dense(&weights), &[0u32, 1, 2, 3], Some(2), ColdResidualCodec::Bf16, @@ -426,7 +443,7 @@ mod tests { fn decode_step_extends_position() { let weights = make_test_weights(); let prefill = rs_prefill_codec( - &weights, + larql_inference::WeightsView::dense(&weights), &[0u32, 1], None, ColdResidualCodec::Bf16, @@ -434,8 +451,15 @@ mod tests { None, ); assert_eq!(prefill.store.next_position, 2); - let (_, rs2) = - rs_decode_step_codec(&weights, 2, prefill.store, &CpuBackend, None, None).unwrap(); + let (_, rs2) = rs_decode_step_codec( + larql_inference::WeightsView::dense(&weights), + 2, + prefill.store, + &CpuBackend, + None, + None, + ) + .unwrap(); assert_eq!(rs2.next_position, 3); } @@ -443,7 +467,7 @@ mod tests { fn decode_with_cold_kv_path_produces_finite_output() { let weights = make_test_weights(); let prefill = rs_prefill_codec( - &weights, + larql_inference::WeightsView::dense(&weights), &[0u32, 1, 2, 3], Some(2), ColdResidualCodec::Bf16, @@ -451,8 +475,15 @@ mod tests { None, ); assert!(prefill.store.cold_kv.is_some()); - let (h, _) = - rs_decode_step_codec(&weights, 4, prefill.store, &CpuBackend, None, None).unwrap(); + 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())); } @@ -463,18 +494,33 @@ mod tests { // exercised (we read from cold_encoded directly via decode). let weights = make_test_weights(); let prefill = rs_prefill_codec( - &weights, + larql_inference::WeightsView::dense(&weights), &[0u32, 1, 2, 3], Some(2), ColdResidualCodec::Bf16, &CpuBackend, None, ); - let (_, rs2) = - rs_decode_step_codec(&weights, 4, prefill.store, &CpuBackend, None, None).unwrap(); + 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(&weights, 5, rs2, &CpuBackend, None, None).unwrap(); + 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())); } @@ -522,7 +568,7 @@ mod tests { Some(if inplace { "1" } else { "0" }), ); let prefill = rs_prefill_codec( - &weights, + larql_inference::WeightsView::dense(&weights), &[0u32, 1, 2], None, ColdResidualCodec::Bf16, @@ -532,9 +578,15 @@ mod tests { let mut rs = prefill.store; let mut hiddens = Vec::new(); for tok in 3u32..=12 { - let (h, rs2) = - rs_decode_step_codec(&weights, tok, rs, &CpuBackend, None, Some(&index)) - .expect("decode"); + 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; diff --git a/crates/larql-kv/src/engines/markov_residual_codec/dispatch.rs b/crates/larql-kv/src/engines/markov_residual_codec/dispatch.rs index ebd0a6a7f..e686269b4 100644 --- a/crates/larql-kv/src/engines/markov_residual_codec/dispatch.rs +++ b/crates/larql-kv/src/engines/markov_residual_codec/dispatch.rs @@ -31,7 +31,7 @@ impl MarkovResidualCodecEngine { /// invalidate `cold_kv` because codec round-trip is lossy. pub(super) fn try_prefill_via_dispatch( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, index: &larql_inference::larql_vindex::VectorIndex, token_ids: &[u32], ) -> Option> { @@ -129,7 +129,7 @@ impl MarkovResidualCodecEngine { /// `cold_kv` so the next step recomputes against the decoded bytes. pub(super) fn decode_step_via_dispatch( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, index: &larql_inference::larql_vindex::VectorIndex, token_id: u32, ) -> Option> { @@ -320,9 +320,9 @@ mod tests { ColdResidualCodec::Bf16, cpu_engine_backend(), ); - let mut w = weights; + let w = weights; assert!(engine - .try_prefill_via_dispatch(&mut w, &empty_index, &[0u32, 1]) + .try_prefill_via_dispatch(&w, &empty_index, &[0u32, 1]) .is_none()); assert!(engine.store.is_none()); assert!(engine.kv_handle.is_none()); @@ -331,9 +331,9 @@ mod tests { #[test] fn try_prefill_via_dispatch_windowed_keeps_stored_under_w10_default() { set_w10_disable(false); - let (mut engine, mut weights, index) = fixture(Some(8)); + let (mut engine, weights, index) = fixture(Some(8)); let h = engine - .try_prefill_via_dispatch(&mut weights, &index, &[0u32, 1, 2]) + .try_prefill_via_dispatch(&weights, &index, &[0u32, 1, 2]) .expect("prefill"); assert_eq!(h.shape(), &[1, weights.hidden_size]); let rs = engine.store.as_ref().unwrap(); @@ -348,10 +348,10 @@ mod tests { #[test] fn try_prefill_via_dispatch_windowless_drops_stored_under_w10() { set_w10_disable(false); - let (mut engine, mut weights, index) = fixture(None); + let (mut engine, weights, index) = fixture(None); let hidden = weights.hidden_size; engine - .try_prefill_via_dispatch(&mut weights, &index, &[0u32, 1, 2]) + .try_prefill_via_dispatch(&weights, &index, &[0u32, 1, 2]) .expect("prefill (windowless)"); let rs = engine.store.as_ref().unwrap(); assert!(rs.hot_kv.is_none()); @@ -366,9 +366,9 @@ mod tests { // window=2 against 4 tokens → prefill clip evicts 2 positions // into `cold_encoded` via the bf16 codec. set_w10_disable(false); - let (mut engine, mut weights, index) = fixture(Some(2)); + let (mut engine, weights, index) = fixture(Some(2)); engine - .try_prefill_via_dispatch(&mut weights, &index, &[0u32, 1, 2, 3]) + .try_prefill_via_dispatch(&weights, &index, &[0u32, 1, 2, 3]) .expect("prefill (overflow)"); let rs = engine.store.as_ref().unwrap(); let cold = rs.cold_encoded.as_ref().expect("cold_encoded populated"); @@ -381,8 +381,8 @@ mod tests { #[test] fn try_prefill_via_dispatch_full_mask_with_w10_disabled() { set_w10_disable(true); - let (mut engine, mut weights, index) = fixture(Some(8)); - let res = engine.try_prefill_via_dispatch(&mut weights, &index, &[0u32, 1, 2]); + let (mut engine, weights, index) = fixture(Some(8)); + let res = engine.try_prefill_via_dispatch(&weights, &index, &[0u32, 1, 2]); set_w10_disable(false); res.expect("prefill"); let rs = engine.store.as_ref().unwrap(); @@ -393,22 +393,22 @@ mod tests { #[test] fn decode_step_via_dispatch_without_prefill_returns_none() { set_w10_disable(false); - let (mut engine, mut weights, index) = fixture(Some(4)); + let (mut engine, weights, index) = fixture(Some(4)); assert!(engine - .decode_step_via_dispatch(&mut weights, &index, 0) + .decode_step_via_dispatch(&weights, &index, 0) .is_none()); } #[test] fn decode_step_via_dispatch_windowed_appends_h_in_under_honly() { set_w10_disable(false); - let (mut engine, mut weights, index) = fixture(Some(8)); + let (mut engine, weights, index) = fixture(Some(8)); engine - .try_prefill_via_dispatch(&mut weights, &index, &[0u32, 1]) + .try_prefill_via_dispatch(&weights, &index, &[0u32, 1]) .expect("prefill"); let hot_len_before = engine.store.as_ref().unwrap().hot_len; let h = engine - .decode_step_via_dispatch(&mut weights, &index, 2) + .decode_step_via_dispatch(&weights, &index, 2) .expect("decode"); assert_eq!(h.shape(), &[1, weights.hidden_size]); let rs = engine.store.as_ref().unwrap(); @@ -420,13 +420,13 @@ mod tests { #[test] fn decode_step_via_dispatch_windowless_uses_none_mask() { set_w10_disable(false); - let (mut engine, mut weights, index) = fixture(None); + let (mut engine, weights, index) = fixture(None); let hidden = weights.hidden_size; engine - .try_prefill_via_dispatch(&mut weights, &index, &[0u32, 1]) + .try_prefill_via_dispatch(&weights, &index, &[0u32, 1]) .expect("prefill (windowless)"); engine - .decode_step_via_dispatch(&mut weights, &index, 2) + .decode_step_via_dispatch(&weights, &index, 2) .expect("decode (None mask)"); let rs = engine.store.as_ref().unwrap(); for slab in &rs.stored { @@ -439,12 +439,12 @@ mod tests { #[test] fn decode_step_via_dispatch_full_mask_with_w10_disabled() { set_w10_disable(true); - let (mut engine, mut weights, index) = fixture(Some(8)); + let (mut engine, weights, index) = fixture(Some(8)); engine - .try_prefill_via_dispatch(&mut weights, &index, &[0u32, 1]) + .try_prefill_via_dispatch(&weights, &index, &[0u32, 1]) .expect("prefill"); let hot_len_before = engine.store.as_ref().unwrap().hot_len; - let res = engine.decode_step_via_dispatch(&mut weights, &index, 2); + let res = engine.decode_step_via_dispatch(&weights, &index, 2); set_w10_disable(false); res.expect("decode"); let rs = engine.store.as_ref().unwrap(); @@ -458,13 +458,13 @@ mod tests { // position into cold_encoded (the None match arm constructs it). // Second decode extends an existing cold_encoded (Some arm). set_w10_disable(false); - let (mut engine, mut weights, index) = fixture(Some(2)); + let (mut engine, weights, index) = fixture(Some(2)); engine - .try_prefill_via_dispatch(&mut weights, &index, &[0u32, 1]) + .try_prefill_via_dispatch(&weights, &index, &[0u32, 1]) .expect("prefill"); assert!(engine.store.as_ref().unwrap().cold_encoded.is_none()); engine - .decode_step_via_dispatch(&mut weights, &index, 2) + .decode_step_via_dispatch(&weights, &index, 2) .expect("decode 1"); assert!(engine.store.as_ref().unwrap().cold_encoded.is_some()); let n_before = engine @@ -476,7 +476,7 @@ mod tests { .unwrap()[0] .n_positions; engine - .decode_step_via_dispatch(&mut weights, &index, 3) + .decode_step_via_dispatch(&weights, &index, 3) .expect("decode 2"); let n_after = engine .store @@ -495,13 +495,13 @@ mod tests { #[test] fn decode_step_via_dispatch_with_profiling_records_stages() { set_w10_disable(false); - let (engine, mut weights, index) = fixture(Some(8)); + let (engine, weights, index) = fixture(Some(8)); let mut engine = engine.with_profiling(true); engine - .try_prefill_via_dispatch(&mut weights, &index, &[0u32, 1]) + .try_prefill_via_dispatch(&weights, &index, &[0u32, 1]) .expect("prefill"); engine - .decode_step_via_dispatch(&mut weights, &index, 2) + .decode_step_via_dispatch(&weights, &index, 2) .expect("decode"); assert!(engine.profile.decode_total.count >= 1); assert!(engine.profile.state_capture.count >= 1); 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 fdd066464..e7e6ced4a 100644 --- a/crates/larql-kv/src/engines/markov_residual_codec/engine.rs +++ b/crates/larql-kv/src/engines/markov_residual_codec/engine.rs @@ -41,6 +41,9 @@ pub struct MarkovResidualCodecEngine { /// W1-GPU: see `MarkovResidualEngine::kv_handle`. pub(super) kv_handle: Option, pub(super) abs_position: usize, + /// Engine-owned f32 dequant scratch for the per-layer codec-walk fallback + /// (see `MarkovResidualEngine::dequant_scratch`). Keeps `weights` immutable. + pub(super) dequant_scratch: larql_inference::DequantScratch, } impl MarkovResidualCodecEngine { @@ -64,6 +67,7 @@ impl MarkovResidualCodecEngine { profile: EngineProfiler::default(), kv_handle: None, abs_position: 0, + dequant_scratch: larql_inference::DequantScratch::new(), } } @@ -105,7 +109,7 @@ impl MarkovResidualCodecEngine { what: "decode_step called before prefill (store missing)".into(), })?; let (hidden, new_rs) = rs_decode_step_codec( - weights, + larql_inference::WeightsView::dense(weights), token_id, rs, self.backend.as_ref(), @@ -153,7 +157,7 @@ impl KvEngine for MarkovResidualCodecEngine { return Err(EngineError::EmptyPrompt); } let result = rs_prefill_codec( - weights, + larql_inference::WeightsView::dense(weights), token_ids, self.window_size, self.codec, @@ -200,7 +204,7 @@ impl KvEngine for MarkovResidualCodecEngine { fn prefill_quant( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, _ffn: &dyn FfnBackend, index: &larql_inference::larql_vindex::VectorIndex, token_ids: &[u32], @@ -216,9 +220,10 @@ impl KvEngine for MarkovResidualCodecEngine { if let Some(hidden) = self.try_prefill_via_dispatch(weights, index, token_ids) { return Ok(hidden); } - ensure_attn_tensors_dequantised(weights, index); + ensure_attn_tensors_dequantised(&mut self.dequant_scratch, weights, index); + let view = larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch); let result = rs_prefill_codec_walk( - weights, + view, index, token_ids, self.window_size, @@ -234,7 +239,7 @@ impl KvEngine for MarkovResidualCodecEngine { fn decode_step_quant( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, _ffn: &dyn FfnBackend, index: &larql_inference::larql_vindex::VectorIndex, token_id: u32, @@ -247,7 +252,8 @@ impl KvEngine for MarkovResidualCodecEngine { details: "decode_step_via_dispatch returned None".into(), }); } - ensure_attn_tensors_dequantised(weights, index); + ensure_attn_tensors_dequantised(&mut self.dequant_scratch, weights, index); + let view = larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch); let rs = self .store .take() @@ -255,12 +261,10 @@ impl KvEngine for MarkovResidualCodecEngine { what: "decode_step_quant called before prefill (store missing)".into(), })?; let prof = self.profiling.then_some(&mut self.profile); - let (hidden, new_rs) = rs_decode_step_codec_walk( - weights, index, token_id, rs, backend, prof, - ) - .ok_or_else(|| EngineError::BackendFailure { - details: "rs_decode_step_codec_walk returned None".into(), - })?; + let (hidden, new_rs) = rs_decode_step_codec_walk(view, index, token_id, rs, backend, prof) + .ok_or_else(|| EngineError::BackendFailure { + details: "rs_decode_step_codec_walk returned None".into(), + })?; self.store = Some(new_rs); self.abs_position += 1; Ok(hidden) @@ -282,7 +286,7 @@ impl KvEngine for MarkovResidualCodecEngine { fn prefill_quant_via_executor( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, executor: &dyn larql_inference::layer_executor::LayerExecutor, ffn: &dyn FfnBackend, index: &larql_inference::larql_vindex::VectorIndex, @@ -306,7 +310,7 @@ impl KvEngine for MarkovResidualCodecEngine { fn decode_step_quant_via_executor( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, executor: &dyn larql_inference::layer_executor::LayerExecutor, ffn: &dyn FfnBackend, index: &larql_inference::larql_vindex::VectorIndex, @@ -524,13 +528,13 @@ mod tests { #[test] fn prefill_quant_cpu_fallback_runs_walk_path() { use larql_inference::ffn::NullFfn; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = MarkovResidualCodecEngine::new(None, ColdResidualCodec::Bf16); let h = engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1, 2], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1, 2], &*backend) .expect("prefill_quant cpu fallback"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert!(engine.memory_bytes() > 0); @@ -539,17 +543,17 @@ mod tests { #[test] fn decode_step_quant_cpu_fallback_extends_store() { use larql_inference::ffn::NullFfn; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = MarkovResidualCodecEngine::new(None, ColdResidualCodec::Bf16); engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1], &*backend) .expect("prefill_quant"); let mem_before = engine.memory_bytes(); let h = engine - .decode_step_quant(&mut weights, &ffn, &index, 2, &*backend) + .decode_step_quant(&weights, &ffn, &index, 2, &*backend) .expect("decode_step_quant cpu fallback"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert!( @@ -563,13 +567,13 @@ mod tests { // Drive the walk path with a window small enough to force overflow // into the codec-encoded cold tier (lines 149-152 of engine.rs). use larql_inference::ffn::NullFfn; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = MarkovResidualCodecEngine::new(Some(2), ColdResidualCodec::Bf16); engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1, 2, 3], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1, 2, 3], &*backend) .expect("prefill_quant with overflow"); assert!(engine.window_tokens() <= 2); assert!( @@ -581,7 +585,7 @@ mod tests { #[test] fn decode_step_quant_without_prefill_returns_none() { use larql_inference::ffn::NullFfn; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; @@ -589,7 +593,7 @@ mod tests { // No prefill → store is None → decode_step_quant returns // InvariantViolation at the `self.store.take()` guard. assert!(engine - .decode_step_quant(&mut weights, &ffn, &index, 0, &*backend) + .decode_step_quant(&weights, &ffn, &index, 0, &*backend) .is_err()); } @@ -643,7 +647,7 @@ mod tests { #[test] fn prefill_quant_via_executor_runs_and_honors_ffn() { use larql_inference::layer_executor::LocalWalkExecutor; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let executor = LocalWalkExecutor::new(&*backend); @@ -654,7 +658,7 @@ mod tests { }; let mut engine = MarkovResidualCodecEngine::new(None, ColdResidualCodec::Bf16); let h = engine - .prefill_quant_via_executor(&mut weights, &executor, &ffn, &index, &[0u32, 1, 2]) + .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32, 1, 2]) .expect("prefill via executor"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert_eq!( @@ -668,18 +672,18 @@ mod tests { fn decode_step_quant_via_executor_extends_store() { use larql_inference::ffn::NullFfn; use larql_inference::layer_executor::LocalWalkExecutor; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let executor = LocalWalkExecutor::new(&*backend); let ffn = NullFfn; let mut engine = MarkovResidualCodecEngine::new(None, ColdResidualCodec::Bf16); engine - .prefill_quant_via_executor(&mut weights, &executor, &ffn, &index, &[0u32, 1]) + .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32, 1]) .expect("prefill"); let mem_before = engine.memory_bytes(); let h = engine - .decode_step_quant_via_executor(&mut weights, &executor, &ffn, &index, 2) + .decode_step_quant_via_executor(&weights, &executor, &ffn, &index, 2) .expect("decode"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert!(engine.memory_bytes() > mem_before); @@ -689,7 +693,7 @@ mod tests { fn executor_path_populates_codec_cold_tier_under_window() { use larql_inference::ffn::NullFfn; use larql_inference::layer_executor::LocalWalkExecutor; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let executor = LocalWalkExecutor::new(&*backend); @@ -698,7 +702,7 @@ mod tests { // through the codec (bf16). let mut engine = MarkovResidualCodecEngine::new(Some(2), ColdResidualCodec::Bf16); engine - .prefill_quant_via_executor(&mut weights, &executor, &ffn, &index, &[0u32, 1, 2, 3]) + .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32, 1, 2, 3]) .expect("prefill with overflow"); assert!(engine.window_tokens() <= 2); assert!( @@ -714,13 +718,13 @@ mod tests { #[test] fn decode_step_quant_w2_codec_cached_hot_and_cold_steady_state() { use larql_inference::ffn::NullFfn; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = MarkovResidualCodecEngine::new(Some(2), ColdResidualCodec::Bf16); engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1, 2, 3], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1, 2, 3], &*backend) .expect("prefill with overflow"); assert!(engine.store.as_ref().unwrap().hot_kv.is_some()); assert!(engine.store.as_ref().unwrap().cold_kv.is_some()); @@ -732,7 +736,7 @@ mod tests { // recompute) — exercises both sides of the codec's // post-overflow flow. let _ = engine - .decode_step_quant(&mut weights, &ffn, &index, tok, &*backend) + .decode_step_quant(&weights, &ffn, &index, tok, &*backend) .expect("decode"); } } @@ -743,17 +747,17 @@ mod tests { #[test] fn decode_step_quant_w2_codec_falls_back_when_hot_kv_dropped() { use larql_inference::ffn::NullFfn; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = MarkovResidualCodecEngine::new(Some(2), ColdResidualCodec::Bf16); engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1, 2, 3], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1, 2, 3], &*backend) .expect("prefill"); engine.store.as_mut().unwrap().hot_kv = None; let h = engine - .decode_step_quant(&mut weights, &ffn, &index, 4, &*backend) + .decode_step_quant(&weights, &ffn, &index, 4, &*backend) .expect("decode via fallback"); assert_eq!(h.shape(), &[1, weights.hidden_size]); } @@ -764,17 +768,17 @@ mod tests { #[test] fn decode_step_codec_walk_with_profiling_populates_summary() { use larql_inference::ffn::NullFfn; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut eng = MarkovResidualCodecEngine::new(Some(2), ColdResidualCodec::Bf16).with_profiling(true); - eng.prefill_quant(&mut weights, &ffn, &index, &[0u32, 1, 2, 3], &*backend) + eng.prefill_quant(&weights, &ffn, &index, &[0u32, 1, 2, 3], &*backend) .expect("prefill"); - eng.decode_step_quant(&mut weights, &ffn, &index, 4, &*backend) + eng.decode_step_quant(&weights, &ffn, &index, 4, &*backend) .expect("decode 1"); - eng.decode_step_quant(&mut weights, &ffn, &index, 5, &*backend) + eng.decode_step_quant(&weights, &ffn, &index, 5, &*backend) .expect("decode 2"); let summary = eng .stage_summary() @@ -791,17 +795,17 @@ mod tests { fn decode_via_executor_uses_cold_kv_branch() { use larql_inference::ffn::NullFfn; use larql_inference::layer_executor::LocalWalkExecutor; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let executor = LocalWalkExecutor::new(&*backend); let ffn = NullFfn; let mut engine = MarkovResidualCodecEngine::new(Some(2), ColdResidualCodec::Bf16); engine - .prefill_quant_via_executor(&mut weights, &executor, &ffn, &index, &[0u32, 1, 2, 3]) + .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32, 1, 2, 3]) .expect("prefill overflow"); let h = engine - .decode_step_quant_via_executor(&mut weights, &executor, &ffn, &index, 4) + .decode_step_quant_via_executor(&weights, &executor, &ffn, &index, 4) .expect("decode via cold_kv"); assert_eq!(h.shape(), &[1, weights.hidden_size]); } @@ -813,20 +817,20 @@ mod tests { fn decode_via_executor_hits_cold_encoded_branch() { use larql_inference::ffn::NullFfn; use larql_inference::layer_executor::LocalWalkExecutor; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let executor = LocalWalkExecutor::new(&*backend); let ffn = NullFfn; let mut engine = MarkovResidualCodecEngine::new(Some(2), ColdResidualCodec::Bf16); engine - .prefill_quant_via_executor(&mut weights, &executor, &ffn, &index, &[0u32, 1, 2, 3]) + .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32, 1, 2, 3]) .expect("prefill"); engine - .decode_step_quant_via_executor(&mut weights, &executor, &ffn, &index, 4) + .decode_step_quant_via_executor(&weights, &executor, &ffn, &index, 4) .expect("first decode clears cold_kv"); let h = engine - .decode_step_quant_via_executor(&mut weights, &executor, &ffn, &index, 5) + .decode_step_quant_via_executor(&weights, &executor, &ffn, &index, 5) .expect("decode via cold_encoded recompute"); assert_eq!(h.shape(), &[1, weights.hidden_size]); } @@ -851,7 +855,7 @@ mod tests { #[test] fn fused_executor_falls_back_to_legacy_quant_path() { use larql_inference::ffn::NullFfn; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let exec = FusedStubExecutor { backend: larql_compute::CpuBackend, @@ -859,11 +863,11 @@ mod tests { let ffn = NullFfn; let mut engine = MarkovResidualCodecEngine::new(None, ColdResidualCodec::Bf16); let h = engine - .prefill_quant_via_executor(&mut weights, &exec, &ffn, &index, &[0u32, 1]) + .prefill_quant_via_executor(&weights, &exec, &ffn, &index, &[0u32, 1]) .expect("fused fallback prefill"); assert_eq!(h.shape(), &[1, weights.hidden_size]); let h2 = engine - .decode_step_quant_via_executor(&mut weights, &exec, &ffn, &index, 2) + .decode_step_quant_via_executor(&weights, &exec, &ffn, &index, 2) .expect("fused fallback decode"); assert_eq!(h2.shape(), &[1, weights.hidden_size]); } @@ -885,13 +889,13 @@ mod tests { } use larql_inference::ffn::NullFfn; use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = MarkovResidualCodecEngine::new(None, ColdResidualCodec::Bf16); let h = engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1, 2], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1, 2], &*backend) .expect("dispatch-path prefill on Q4K vindex"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert!(engine.kv_handle.is_some()); @@ -908,18 +912,18 @@ mod tests { fn decode_via_dispatch_grows_buffers_in_place() { use larql_inference::ffn::NullFfn; use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = MarkovResidualCodecEngine::new(None, ColdResidualCodec::Bf16); engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1], &*backend) .expect("prefill dispatch"); let mem_after_prefill = engine.memory_bytes(); for tok in 2..6u32 { let h = engine - .decode_step_quant(&mut weights, &ffn, &index, tok, &*backend) + .decode_step_quant(&weights, &ffn, &index, tok, &*backend) .expect("dispatch decode step"); assert_eq!(h.shape(), &[1, weights.hidden_size]); } @@ -937,21 +941,21 @@ mod tests { // logits. use larql_inference::ffn::NullFfn; use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = MarkovResidualCodecEngine::new(None, ColdResidualCodec::Bf16).with_profiling(true); engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1], &*backend) .expect("prefill"); let h = engine - .decode_step_quant(&mut weights, &ffn, &index, 2, &*backend) + .decode_step_quant(&weights, &ffn, &index, 2, &*backend) .expect("decode 1"); assert!(h.iter().all(|v| v.is_finite())); let h2 = engine - .decode_step_quant(&mut weights, &ffn, &index, 3, &*backend) + .decode_step_quant(&weights, &ffn, &index, 3, &*backend) .expect("decode 2"); assert!(h2.iter().all(|v| v.is_finite())); } @@ -965,13 +969,13 @@ mod tests { // from the decoded payload. use larql_inference::ffn::NullFfn; use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = MarkovResidualCodecEngine::new(Some(2), ColdResidualCodec::Bf16); engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1, 2, 3], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1, 2, 3], &*backend) .expect("prefill with window"); let store = engine.store.as_ref().expect("store"); assert!(store.cold_encoded.is_some()); @@ -992,17 +996,17 @@ mod tests { // optionally extend cold_kv. use larql_inference::ffn::NullFfn; use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = MarkovResidualCodecEngine::new(Some(2), ColdResidualCodec::Bf16); engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1, 2], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1, 2], &*backend) .expect("prefill at window cap"); for tok in 3..6u32 { engine - .decode_step_quant(&mut weights, &ffn, &index, tok, &*backend) + .decode_step_quant(&weights, &ffn, &index, tok, &*backend) .expect("dispatch decode with codec overflow"); } let store = engine.store.as_ref().expect("store"); diff --git a/crates/larql-kv/src/engines/markov_residual_codec/executor.rs b/crates/larql-kv/src/engines/markov_residual_codec/executor.rs index f649a329e..33daa3c61 100644 --- a/crates/larql-kv/src/engines/markov_residual_codec/executor.rs +++ b/crates/larql-kv/src/engines/markov_residual_codec/executor.rs @@ -24,13 +24,14 @@ impl MarkovResidualCodecEngine { /// `prefill_quant` in that case). pub(super) fn prefill_via_executor_impl( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, executor: &dyn LayerExecutor, ffn: &dyn FfnBackend, index: &larql_inference::larql_vindex::VectorIndex, token_ids: &[u32], ) -> Option> { - ensure_attn_tensors_dequantised(weights, index); + ensure_attn_tensors_dequantised(&mut self.dequant_scratch, weights, index); + let view = larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch); let backend = executor.backend(); let num_layers = weights.num_layers; @@ -41,7 +42,7 @@ impl MarkovResidualCodecEngine { for layer in 0..num_layers { stored.push(h.clone()); - let (h_out, _kv) = executor.run_prefill_layer(weights, layer, &h, ffn)?; + let (h_out, _kv) = executor.run_prefill_layer(view, layer, &h, ffn)?; h = h_out; } @@ -74,7 +75,7 @@ impl MarkovResidualCodecEngine { let mut tmp = EncodedColdLayer::empty(hidden_size); tmp.append(self.codec, overflow); let decoded = tmp.decode(self.codec); - let (k, v) = recompute_kv(weights, &decoded, layer, 0, backend, Some(index)) + let (k, v) = recompute_kv(view, &decoded, layer, 0, backend, Some(index)) .expect("cold K/V pre-computation failed"); cold_kv.push((k, v)); let mut enc = EncodedColdLayer::empty(hidden_size); @@ -98,13 +99,14 @@ impl MarkovResidualCodecEngine { /// that `executor.dispatch_kind() != Fused`. pub(super) fn decode_step_via_executor_impl( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, executor: &dyn LayerExecutor, ffn: &dyn FfnBackend, index: &larql_inference::larql_vindex::VectorIndex, token_id: u32, ) -> Option> { - ensure_attn_tensors_dequantised(weights, index); + ensure_attn_tensors_dequantised(&mut self.dequant_scratch, weights, index); + let view = larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch); let backend = executor.backend(); let rs = self.store.take()?; @@ -121,7 +123,7 @@ impl MarkovResidualCodecEngine { let prior_kv: 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, Some(index))?; + recompute_kv(view, h_hot, layer, hot_abs_start, backend, Some(index))?; let c = k_cold.shape()[0]; let kv_dim = k_cold.shape()[1]; let mut k_combined = Array2::::zeros((c + s_hot, kv_dim)); @@ -148,19 +150,12 @@ impl MarkovResidualCodecEngine { } _ => (h_hot.clone(), hot_abs_start), }; - recompute_kv( - weights, - &h_full, - layer, - full_abs_start, - backend, - Some(index), - )? + recompute_kv(view, &h_full, layer, full_abs_start, backend, Some(index))? }; new_stored.push(h_new.clone()); let (h_out, _new_kv) = - executor.run_decode_layer(weights, layer, &h_new, &prior_kv, abs_position, ffn)?; + executor.run_decode_layer(view, layer, &h_new, &prior_kv, abs_position, ffn)?; h_new = h_out; } 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 1e3f1b594..081ac349e 100644 --- a/crates/larql-kv/src/engines/markov_residual_codec/walk.rs +++ b/crates/larql-kv/src/engines/markov_residual_codec/walk.rs @@ -11,7 +11,6 @@ use larql_compute::ComputeBackend; use larql_inference::attention::{run_attention_with_kv_backend, SharedKV}; use larql_inference::forward::{embed_tokens_pub, run_ffn}; -use larql_inference::model::ModelWeights; use larql_inference::vindex::{WalkFfn, WalkFfnConfig}; use larql_vindex::VectorIndex; use ndarray::{s, Array2}; @@ -23,7 +22,7 @@ use crate::engines::markov_residual_codec::store::{EncodedColdLayer, RsStoreCode use crate::profiler::EngineProfiler; pub fn rs_prefill_codec_walk( - weights: &ModelWeights, + weights: larql_inference::WeightsView, index: &VectorIndex, token_ids: &[u32], max_window: Option, @@ -32,12 +31,13 @@ pub fn rs_prefill_codec_walk( ) -> RsPrefillResultCodec { let num_layers = weights.num_layers; let seq_len = token_ids.len(); - let mut h = embed_tokens_pub(weights, token_ids); + let mut h = embed_tokens_pub(&weights, token_ids); let mut stored: Vec> = Vec::with_capacity(num_layers); let be = Some(backend); - let walk_ffn = WalkFfn::from_config(weights, index, WalkFfnConfig::dense(num_layers)) - .with_backend(backend); + let walk_ffn = + WalkFfn::from_config(weights.canonical(), index, WalkFfnConfig::dense(num_layers)) + .with_backend(backend); // Capture per-layer K/V from each layer's attention block — same // pattern as `rs_prefill_walk` (W2). Decode reuses these instead @@ -45,10 +45,10 @@ pub fn rs_prefill_codec_walk( let mut hot_kv_captured: Vec = Vec::with_capacity(num_layers); 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) + let (h_post_attn, k, v) = run_attention_with_kv_backend(weights, &h, layer, be, None) .expect("attention failed during MarkovResidualCodec Q4K prefill"); hot_kv_captured.push((k, v)); - let (h_out, _) = run_ffn(weights, &h_post_attn, layer, &walk_ffn, false); + let (h_out, _) = run_ffn(&weights, &h_post_attn, layer, &walk_ffn, false); h = h_out; } @@ -102,7 +102,7 @@ pub fn rs_prefill_codec_walk( } pub fn rs_decode_step_codec_walk( - weights: &ModelWeights, + weights: larql_inference::WeightsView, index: &VectorIndex, new_token_id: u32, rs: RsStoreCodec, @@ -116,14 +116,15 @@ pub fn rs_decode_step_codec_walk( let num_layers = weights.num_layers; let abs_position = rs.next_position; let t_embed = t_step; - let mut h_new = embed_tokens_pub(weights, &[new_token_id]); + let mut h_new = embed_tokens_pub(&weights, &[new_token_id]); let embed_us = t_embed .map(|t| t.elapsed().as_secs_f64() * 1e6) .unwrap_or(0.0); let mut new_stored: Vec> = Vec::with_capacity(num_layers); - let walk_ffn = WalkFfn::from_config(weights, index, WalkFfnConfig::dense(num_layers)) - .with_backend(backend); + let walk_ffn = + WalkFfn::from_config(weights.canonical(), index, WalkFfnConfig::dense(num_layers)) + .with_backend(backend); let mut recompute_cold_us = 0.0f64; let mut recompute_hot_us = 0.0f64; @@ -213,7 +214,7 @@ pub fn rs_decode_step_codec_walk( // markov_residual::walk::rs_decode_step_walk). let t_attn = if timing { Some(Instant::now()) } else { None }; let native_result = larql_inference::vindex::attention_decode_step_native( - weights, + weights.canonical(), index, backend, &h_new, @@ -245,14 +246,14 @@ pub fn rs_decode_step_codec_walk( // Native Q4K FFN, then WalkFfn fallback. let t_ffn = if timing { Some(Instant::now()) } else { None }; let h_out = larql_inference::vindex::ffn_decode_step_native( - weights, + weights.canonical(), index, backend, &h_post_attn, layer, ) .unwrap_or_else(|| { - let (h, _) = run_ffn(weights, &h_post_attn, layer, &walk_ffn, false); + let (h, _) = run_ffn(&weights, &h_post_attn, layer, &walk_ffn, false); h }); if let Some(t) = t_ffn { @@ -363,7 +364,7 @@ mod tests { let weights = make_test_weights(); let index = make_test_vindex(&weights); let result = rs_prefill_codec_walk( - &weights, + larql_inference::WeightsView::dense(&weights), &index, &[0u32, 1, 2], None, @@ -379,7 +380,7 @@ mod tests { let weights = make_test_weights(); let index = make_test_vindex(&weights); let result = rs_prefill_codec_walk( - &weights, + larql_inference::WeightsView::dense(&weights), &index, &[0u32, 1, 2, 3], Some(2), @@ -395,7 +396,7 @@ mod tests { let weights = make_test_weights(); let index = make_test_vindex(&weights); let prefill = rs_prefill_codec_walk( - &weights, + larql_inference::WeightsView::dense(&weights), &index, &[0u32, 1], None, @@ -403,9 +404,15 @@ mod tests { &CpuBackend, ); assert_eq!(prefill.store.next_position, 2); - let (h, rs2) = - rs_decode_step_codec_walk(&weights, &index, 2, prefill.store, &CpuBackend, None) - .unwrap(); + let (h, rs2) = rs_decode_step_codec_walk( + larql_inference::WeightsView::dense(&weights), + &index, + 2, + prefill.store, + &CpuBackend, + None, + ) + .unwrap(); assert_eq!(rs2.next_position, 3); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert!(h.iter().all(|v| v.is_finite())); @@ -416,7 +423,7 @@ mod tests { let weights = make_test_weights(); let index = make_test_vindex(&weights); let prefill = rs_prefill_codec_walk( - &weights, + larql_inference::WeightsView::dense(&weights), &index, &[0u32, 1, 2, 3], Some(2), @@ -424,9 +431,15 @@ mod tests { &CpuBackend, ); assert!(prefill.store.cold_kv.is_some()); - let (h, _) = - rs_decode_step_codec_walk(&weights, &index, 4, prefill.store, &CpuBackend, None) - .unwrap(); + let (h, _) = rs_decode_step_codec_walk( + larql_inference::WeightsView::dense(&weights), + &index, + 4, + prefill.store, + &CpuBackend, + None, + ) + .unwrap(); assert_eq!(h.shape(), &[1, weights.hidden_size]); } @@ -435,20 +448,33 @@ mod tests { let weights = make_test_weights(); let index = make_test_vindex(&weights); let prefill = rs_prefill_codec_walk( - &weights, + larql_inference::WeightsView::dense(&weights), &index, &[0u32, 1, 2, 3], Some(2), ColdResidualCodec::Bf16, &CpuBackend, ); - let (_, rs2) = - rs_decode_step_codec_walk(&weights, &index, 4, prefill.store, &CpuBackend, None) - .unwrap(); + let (_, rs2) = rs_decode_step_codec_walk( + larql_inference::WeightsView::dense(&weights), + &index, + 4, + prefill.store, + &CpuBackend, + None, + ) + .unwrap(); // First decode clears cold_kv; second decode exercises the // cold_encoded path. - let (h, _) = - rs_decode_step_codec_walk(&weights, &index, 5, rs2, &CpuBackend, None).unwrap(); + let (h, _) = rs_decode_step_codec_walk( + larql_inference::WeightsView::dense(&weights), + &index, + 5, + rs2, + &CpuBackend, + None, + ) + .unwrap(); assert_eq!(h.shape(), &[1, weights.hidden_size]); } @@ -464,7 +490,7 @@ mod tests { let weights = make_test_weights(); let index = make_test_vindex(&weights); let prefill = rs_prefill_codec_walk( - &weights, + larql_inference::WeightsView::dense(&weights), &index, &[0u32, 1, 2, 3], Some(2), @@ -473,7 +499,7 @@ mod tests { ); let mut prof = EngineProfiler::default(); let (h, _) = rs_decode_step_codec_walk( - &weights, + larql_inference::WeightsView::dense(&weights), &index, 4, prefill.store, @@ -501,7 +527,7 @@ mod tests { let weights = make_test_weights(); let index = make_test_vindex(&weights); let prefill = rs_prefill_codec_walk( - &weights, + larql_inference::WeightsView::dense(&weights), &index, &[0u32, 1], Some(2), @@ -511,9 +537,15 @@ mod tests { assert!(prefill.store.cold_encoded.is_none()); assert!(prefill.store.cold_kv.is_none()); - let (_, rs2) = - rs_decode_step_codec_walk(&weights, &index, 2, prefill.store, &CpuBackend, None) - .unwrap(); + let (_, rs2) = rs_decode_step_codec_walk( + larql_inference::WeightsView::dense(&weights), + &index, + 2, + prefill.store, + &CpuBackend, + None, + ) + .unwrap(); // First overflow hits the None arm and initialises cold_encoded. assert!(rs2.cold_encoded.is_some()); assert_eq!(rs2.cold_encoded.as_ref().unwrap()[0].n_positions, 1); @@ -528,7 +560,7 @@ mod tests { let weights = make_test_weights(); let index = make_test_vindex(&weights); let prefill = rs_prefill_codec_walk( - &weights, + larql_inference::WeightsView::dense(&weights), &index, &[0u32, 1, 2, 3], Some(2), @@ -540,8 +572,15 @@ mod tests { // leaving only the codec-encoded cold tier behind. store.hot_kv = None; store.cold_kv = None; - let (h, rs2) = - rs_decode_step_codec_walk(&weights, &index, 4, store, &CpuBackend, None).unwrap(); + let (h, rs2) = rs_decode_step_codec_walk( + larql_inference::WeightsView::dense(&weights), + &index, + 4, + store, + &CpuBackend, + None, + ) + .unwrap(); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert!(h.iter().all(|v| v.is_finite())); // hot_kv is re-captured on every decode step. @@ -556,7 +595,7 @@ mod tests { let weights = make_test_weights(); let index = make_test_vindex(&weights); let prefill = rs_prefill_codec_walk( - &weights, + larql_inference::WeightsView::dense(&weights), &index, &[0u32, 1], None, @@ -568,8 +607,15 @@ mod tests { store.cold_kv = None; // No cold tier at all → exercises the `_` arm. store.cold_encoded = None; - let (h, _) = - rs_decode_step_codec_walk(&weights, &index, 2, store, &CpuBackend, None).unwrap(); + let (h, _) = rs_decode_step_codec_walk( + larql_inference::WeightsView::dense(&weights), + &index, + 2, + store, + &CpuBackend, + None, + ) + .unwrap(); assert_eq!(h.shape(), &[1, weights.hidden_size]); } } diff --git a/crates/larql-kv/src/engines/mod.rs b/crates/larql-kv/src/engines/mod.rs index a455d033b..11dad22f1 100644 --- a/crates/larql-kv/src/engines/mod.rs +++ b/crates/larql-kv/src/engines/mod.rs @@ -39,8 +39,8 @@ //! faster Metal lm_head KNN rather than a full vocab matmul. //! //! - **CPU fallback**: when Metal is unavailable, engines fall back to a CPU -//! path using dequantised attention tensors (lazily inserted into -//! `weights.tensors`) and `WalkFfn` for Q4K FFN. +//! path using dequantised attention tensors (lazily inserted into the +//! engine-owned `dequant_scratch`) and `WalkFfn` for Q4K FFN. //! //! - **Apollo compressed path**: when the store has boundary residuals captured //! at `crystal_layer` (default 30), `forward_from_layer` runs only diff --git a/crates/larql-kv/src/engines/no_cache.rs b/crates/larql-kv/src/engines/no_cache.rs index 3b711f002..68d5de83d 100644 --- a/crates/larql-kv/src/engines/no_cache.rs +++ b/crates/larql-kv/src/engines/no_cache.rs @@ -22,6 +22,11 @@ use larql_inference::{cpu_engine_backend, EngineBackend}; pub struct NoCacheEngine { tokens: Vec, backend: Box, + /// Engine-owned f32 dequant scratch for the Q4K path — `prefill_quant`/ + /// `decode_step_quant` populate it; `prefill`/`decode_step` resolve the + /// forward through a `WeightsView::with_scratch` over it. Empty on the + /// dense path. Keeps `weights` immutable (no `weights.tensors` mutation). + dequant_scratch: larql_inference::DequantScratch, } impl NoCacheEngine { @@ -33,6 +38,7 @@ impl NoCacheEngine { Self { tokens: Vec::new(), backend, + dequant_scratch: larql_inference::DequantScratch::new(), } } } @@ -68,8 +74,9 @@ impl KvEngine for NoCacheEngine { return Err(EngineError::EmptyPrompt); } self.tokens = token_ids.to_vec(); + let view = larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch); let (hidden, _cache) = kv_prefill_run( - weights, + view, ffn, token_ids, None, @@ -89,8 +96,9 @@ impl KvEngine for NoCacheEngine { token_id: u32, ) -> Result, EngineError> { self.tokens.push(token_id); + let view = larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch); let (hidden, _cache) = kv_prefill_run( - weights, + view, ffn, &self.tokens, None, @@ -105,18 +113,22 @@ impl KvEngine for NoCacheEngine { fn prefill_quant( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, _ffn: &dyn FfnBackend, index: &larql_inference::larql_vindex::VectorIndex, token_ids: &[u32], backend: &dyn larql_inference::ComputeBackend, ) -> Result, EngineError> { - // Phase-1 pattern: dequant Q4K attn tensors into `weights.tensors`, - // then run the f32 prefill path. Q4K FFN dispatches through a + // Phase-1 pattern: dequant Q4K attn tensors into the engine-owned + // scratch, then run the f32 prefill path. Q4K FFN dispatches through a // `WalkFfn` constructed from the vindex (the bench passes // `NullFfn` because Q4K FFN is engine-side; using `_ffn` would // silently skip the FFN). See `kv-dispatch-quantization.md`. - larql_inference::vindex::ensure_attn_tensors_dequantised(weights, index); + larql_inference::vindex::ensure_attn_tensors_dequantised( + &mut self.dequant_scratch, + weights, + index, + ); let walk_ffn = larql_inference::vindex::WalkFfn::from_config( weights, index, @@ -128,13 +140,17 @@ impl KvEngine for NoCacheEngine { fn decode_step_quant( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, _ffn: &dyn FfnBackend, index: &larql_inference::larql_vindex::VectorIndex, token_id: u32, backend: &dyn larql_inference::ComputeBackend, ) -> Result, EngineError> { - larql_inference::vindex::ensure_attn_tensors_dequantised(weights, index); + larql_inference::vindex::ensure_attn_tensors_dequantised( + &mut self.dequant_scratch, + weights, + index, + ); let walk_ffn = larql_inference::vindex::WalkFfn::from_config( weights, index, @@ -154,7 +170,7 @@ impl KvEngine for NoCacheEngine { fn prefill_quant_via_executor( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, _executor: &dyn larql_inference::layer_executor::LayerExecutor, ffn: &dyn FfnBackend, index: &larql_inference::larql_vindex::VectorIndex, @@ -163,19 +179,27 @@ impl KvEngine for NoCacheEngine { // No K/V cache so we don't need to drive the per-layer loop // through the executor; the existing prefill (which honors the // FFN parameter) is the right path. Just dequant first. - larql_inference::vindex::ensure_attn_tensors_dequantised(weights, index); + larql_inference::vindex::ensure_attn_tensors_dequantised( + &mut self.dequant_scratch, + weights, + index, + ); self.prefill(weights, ffn, token_ids) } fn decode_step_quant_via_executor( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, _executor: &dyn larql_inference::layer_executor::LayerExecutor, ffn: &dyn FfnBackend, index: &larql_inference::larql_vindex::VectorIndex, token_id: u32, ) -> Result, EngineError> { - larql_inference::vindex::ensure_attn_tensors_dequantised(weights, index); + larql_inference::vindex::ensure_attn_tensors_dequantised( + &mut self.dequant_scratch, + weights, + index, + ); self.decode_step(weights, ffn, token_id) } @@ -332,13 +356,13 @@ mod tests { #[test] fn prefill_quant_cpu_fallback_runs_end_to_end() { use larql_inference::ffn::NullFfn; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = NoCacheEngine::new(); let h = engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1, 2], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1, 2], &*backend) .expect("prefill_quant cpu fallback"); assert_eq!(h.shape(), &[1, weights.hidden_size]); } @@ -346,17 +370,17 @@ mod tests { #[test] fn decode_step_quant_cpu_fallback_appends_token() { use larql_inference::ffn::NullFfn; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = NoCacheEngine::new(); engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1], &*backend) .expect("prefill_quant"); let mem_before = engine.memory_bytes(); let h = engine - .decode_step_quant(&mut weights, &ffn, &index, 2, &*backend) + .decode_step_quant(&weights, &ffn, &index, 2, &*backend) .expect("decode_step_quant cpu fallback"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert!( @@ -381,14 +405,14 @@ mod tests { fn prefill_quant_via_executor_dequants_and_runs_prefill() { use larql_inference::ffn::NullFfn; use larql_inference::layer_executor::LocalWalkExecutor; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let executor = LocalWalkExecutor::new(&*backend); let ffn = NullFfn; let mut engine = NoCacheEngine::new(); let h = engine - .prefill_quant_via_executor(&mut weights, &executor, &ffn, &index, &[0u32, 1]) + .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32, 1]) .expect("prefill via executor"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert_eq!(engine.window_tokens(), 2); @@ -398,18 +422,18 @@ mod tests { fn decode_step_quant_via_executor_appends_token() { use larql_inference::ffn::NullFfn; use larql_inference::layer_executor::LocalWalkExecutor; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let executor = LocalWalkExecutor::new(&*backend); let ffn = NullFfn; let mut engine = NoCacheEngine::new(); engine - .prefill_quant_via_executor(&mut weights, &executor, &ffn, &index, &[0u32]) + .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32]) .expect("prefill"); let mem_before = engine.memory_bytes(); let h = engine - .decode_step_quant_via_executor(&mut weights, &executor, &ffn, &index, 1) + .decode_step_quant_via_executor(&weights, &executor, &ffn, &index, 1) .expect("decode via executor"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert!(engine.memory_bytes() > mem_before); diff --git a/crates/larql-kv/src/engines/standard.rs b/crates/larql-kv/src/engines/standard.rs index 4dbdb19cf..b36142d27 100644 --- a/crates/larql-kv/src/engines/standard.rs +++ b/crates/larql-kv/src/engines/standard.rs @@ -63,6 +63,12 @@ pub struct StandardEngine { /// its own `next_position` field; this engine tracks it directly. abs_position: usize, backend: BackendSlot, + /// Engine-owned f32 dequant scratch for the per-layer fallback Q4K path + /// (`prefill_quant`/`decode_step_quant` populate it; `do_prefill`/ + /// `do_decode_step` resolve attention/FFN through a + /// `WeightsView::with_scratch` over it). Empty on the dense path. Keeps + /// `weights` immutable so the engine can hold `Arc`. + dequant_scratch: larql_inference::DequantScratch, } impl StandardEngine { @@ -76,6 +82,7 @@ impl StandardEngine { handles: None, abs_position: 0, backend: BackendSlot::Sync(backend), + dequant_scratch: larql_inference::DequantScratch::new(), } } @@ -92,6 +99,7 @@ impl StandardEngine { handles: None, abs_position: 0, backend: BackendSlot::Async(backend), + dequant_scratch: larql_inference::DequantScratch::new(), } } @@ -118,21 +126,17 @@ impl StandardEngine { token_ids: &[u32], index: Option<&larql_inference::larql_vindex::VectorIndex>, ) -> Result, EngineError> { + let view = larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch); let (hidden, handles) = match &self.backend { - BackendSlot::Sync(b) => kv_prefill_via_dispatch( - b.as_ref(), - weights, - ffn, - token_ids, - self.window_size, - index, - ) - .ok_or_else(|| EngineError::BackendFailure { - details: "kv_prefill_via_dispatch returned None".into(), - })?, + BackendSlot::Sync(b) => { + kv_prefill_via_dispatch(b.as_ref(), view, ffn, token_ids, self.window_size, index) + .ok_or_else(|| EngineError::BackendFailure { + details: "kv_prefill_via_dispatch returned None".into(), + })? + } BackendSlot::Async(b) => kv_prefill_via_dispatch_async( b.as_ref(), - weights, + view, ffn, token_ids, self.window_size, @@ -159,10 +163,11 @@ impl StandardEngine { ffn: &dyn FfnBackend, initial_hidden: &Array2, ) -> Option> { + let view = larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch); let (hidden, handles) = match &self.backend { BackendSlot::Sync(b) => kv_prefill_from_hidden_via_dispatch( b.as_ref(), - weights, + view, ffn, initial_hidden, self.window_size, @@ -170,7 +175,7 @@ impl StandardEngine { )?, BackendSlot::Async(b) => kv_prefill_from_hidden_via_dispatch_async( b.as_ref(), - weights, + view, ffn, initial_hidden, self.window_size, @@ -197,6 +202,7 @@ impl StandardEngine { token_id: u32, index: Option<&larql_inference::larql_vindex::VectorIndex>, ) -> Result, EngineError> { + let view = larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch); let handles = self .handles .as_mut() @@ -206,7 +212,7 @@ impl StandardEngine { let hidden = match &self.backend { BackendSlot::Sync(b) => kv_decode_step_via_dispatch( b.as_ref(), - weights, + view, ffn, handles, token_id, @@ -219,7 +225,7 @@ impl StandardEngine { })?, BackendSlot::Async(b) => kv_decode_step_via_dispatch_async( b.as_ref(), - weights, + view, ffn, handles, token_id, @@ -306,7 +312,7 @@ impl KvEngine for StandardEngine { fn prefill_quant( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, ffn: &dyn FfnBackend, index: &larql_inference::larql_vindex::VectorIndex, token_ids: &[u32], @@ -332,14 +338,20 @@ impl KvEngine for StandardEngine { } // Backend doesn't have a coarse path (e.g. f32 model, or // hybrid-MoE / cross-layer-KV models that don't fit the cached - // shape). Fall back to per-layer dispatch with dequant. - larql_inference::vindex::ensure_attn_tensors_dequantised(weights, index); + // shape). Fall back to per-layer dispatch, dequantising attention + // into the engine-owned scratch (`do_prefill` resolves it through a + // `WeightsView::with_scratch`) — `weights` stays immutable. + larql_inference::vindex::ensure_attn_tensors_dequantised( + &mut self.dequant_scratch, + weights, + index, + ); self.do_prefill(weights, ffn, token_ids, Some(index)) } fn decode_step_quant( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, ffn: &dyn FfnBackend, index: &larql_inference::larql_vindex::VectorIndex, token_id: u32, @@ -377,8 +389,14 @@ impl KvEngine for StandardEngine { return Ok(h); } } - // Per-layer dispatch fallback. - larql_inference::vindex::ensure_attn_tensors_dequantised(weights, index); + // Per-layer dispatch fallback. Dequantise attention into the + // engine-owned scratch (idempotent — persists across decode steps); + // `do_decode_step` resolves it through a `WeightsView::with_scratch`. + larql_inference::vindex::ensure_attn_tensors_dequantised( + &mut self.dequant_scratch, + weights, + index, + ); self.do_decode_step(weights, ffn, token_id, Some(index)) } @@ -921,13 +939,13 @@ mod tests { fn prefill_quant_cpu_fallback_runs_via_dequant() { use larql_inference::ffn::NullFfn; use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = StandardEngine::new(None); let h = engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1, 2], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1, 2], &*backend) .expect("prefill_quant Q4K cpu fallback"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert!(engine.memory_bytes() > 0); @@ -937,17 +955,17 @@ mod tests { fn decode_step_quant_cpu_fallback_extends_cache() { use larql_inference::ffn::NullFfn; use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = StandardEngine::new(None); engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1], &*backend) .expect("prefill_quant"); let mem_before = engine.memory_bytes(); let h = engine - .decode_step_quant(&mut weights, &ffn, &index, 2, &*backend) + .decode_step_quant(&weights, &ffn, &index, 2, &*backend) .expect("decode_step_quant Q4K cpu fallback"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert!( @@ -960,7 +978,7 @@ mod tests { fn decode_step_quant_without_prefill_returns_none() { use larql_inference::ffn::NullFfn; use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; @@ -968,7 +986,7 @@ mod tests { // self.handles is None → decode_step_quant returns an // InvariantViolation error at the `self.handles.as_mut()` guard. assert!(engine - .decode_step_quant(&mut weights, &ffn, &index, 0, &*backend) + .decode_step_quant(&weights, &ffn, &index, 0, &*backend) .is_err()); } @@ -996,13 +1014,13 @@ mod tests { fn prefill_quant_empty_prompt_returns_empty_prompt_error() { use larql_inference::ffn::NullFfn; use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = StandardEngine::new(None); let err = engine - .prefill_quant(&mut weights, &ffn, &index, &[], &*backend) + .prefill_quant(&weights, &ffn, &index, &[], &*backend) .unwrap_err(); assert!( matches!(err, EngineError::EmptyPrompt), @@ -1104,14 +1122,14 @@ mod tests { use larql_inference::ffn::NullFfn; use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; use larql_inference::AsyncComputeBackend; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let async_backend: Box = Box::new(CpuBackend); let ffn = NullFfn; let mut engine = StandardEngine::with_async_backend(None, async_backend); let h = engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1, 2], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1, 2], &*backend) .expect("prefill_quant async-slot fallback"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert!(engine.memory_bytes() > 0); @@ -1122,18 +1140,18 @@ mod tests { use larql_inference::ffn::NullFfn; use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; use larql_inference::AsyncComputeBackend; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let async_backend: Box = Box::new(CpuBackend); let ffn = NullFfn; let mut engine = StandardEngine::with_async_backend(None, async_backend); engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1], &*backend) .expect("prefill_quant async-slot"); let mem_before = engine.memory_bytes(); let h = engine - .decode_step_quant(&mut weights, &ffn, &index, 2, &*backend) + .decode_step_quant(&weights, &ffn, &index, 2, &*backend) .expect("decode_step_quant async-slot fallback"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert!( diff --git a/crates/larql-kv/src/engines/turbo_quant/dispatch.rs b/crates/larql-kv/src/engines/turbo_quant/dispatch.rs index 8763a3541..996d68047 100644 --- a/crates/larql-kv/src/engines/turbo_quant/dispatch.rs +++ b/crates/larql-kv/src/engines/turbo_quant/dispatch.rs @@ -25,7 +25,7 @@ impl TurboQuantEngine { /// entries (one per model layer) for the engine's contract. pub(super) fn try_prefill_via_dispatch( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, index: &VectorIndex, token_ids: &[u32], ) -> Option> { @@ -67,7 +67,7 @@ impl TurboQuantEngine { /// buffer (append-only — 2026-05-19 perf fix). pub(super) fn decode_step_via_dispatch( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, index: &VectorIndex, token_id: u32, ) -> Option> { @@ -166,9 +166,9 @@ mod tests { #[test] fn prefill_via_dispatch_compresses_per_layer_kv() { - let (mut engine, mut weights, index) = fixture(); + let (mut engine, weights, index) = fixture(); let h = engine - .try_prefill_via_dispatch(&mut weights, &index, &[0u32, 1, 2]) + .try_prefill_via_dispatch(&weights, &index, &[0u32, 1, 2]) .expect("prefill via dispatch"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert_eq!(engine.layers.len(), weights.num_layers); @@ -194,9 +194,9 @@ mod tests { weights.hidden_size, ); let mut engine = TurboQuantEngine::with_backend(4, cpu_engine_backend()); - let mut w = weights; + let w = weights; assert!(engine - .try_prefill_via_dispatch(&mut w, &empty_index, &[0u32, 1]) + .try_prefill_via_dispatch(&w, &empty_index, &[0u32, 1]) .is_none()); assert!(engine.kv_handle.is_none()); assert!(engine.layers.is_empty()); @@ -206,22 +206,22 @@ mod tests { fn prefill_via_dispatch_returns_none_on_empty_prompt() { // Empty token slice routes through CpuBackend's empty-tokens // guard in `coarse_prefill_with_state` → None. - let (mut engine, mut weights, index) = fixture(); + let (mut engine, weights, index) = fixture(); assert!(engine - .try_prefill_via_dispatch(&mut weights, &index, &[]) + .try_prefill_via_dispatch(&weights, &index, &[]) .is_none()); assert!(engine.kv_handle.is_none()); } #[test] fn decode_step_via_dispatch_appends_one_position() { - let (mut engine, mut weights, index) = fixture(); + let (mut engine, weights, index) = fixture(); engine - .try_prefill_via_dispatch(&mut weights, &index, &[0u32, 1]) + .try_prefill_via_dispatch(&weights, &index, &[0u32, 1]) .expect("prefill"); let bytes_before: usize = engine.layers.iter().map(|l| l.compressed_k.len()).sum(); let h = engine - .decode_step_via_dispatch(&mut weights, &index, 2) + .decode_step_via_dispatch(&weights, &index, 2) .expect("decode"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert_eq!(engine.abs_position, 3); @@ -239,22 +239,22 @@ mod tests { #[test] fn decode_step_via_dispatch_without_prefill_returns_none() { - let (mut engine, mut weights, index) = fixture(); + let (mut engine, weights, index) = fixture(); // kv_handle is None → early return at `self.kv_handle.as_mut()?`. assert!(engine - .decode_step_via_dispatch(&mut weights, &index, 0) + .decode_step_via_dispatch(&weights, &index, 0) .is_none()); } #[test] fn decode_step_via_dispatch_with_profiling_records_stages() { - let (engine, mut weights, index) = fixture(); + let (engine, weights, index) = fixture(); let mut engine = engine.with_profiling(true); engine - .try_prefill_via_dispatch(&mut weights, &index, &[0u32, 1]) + .try_prefill_via_dispatch(&weights, &index, &[0u32, 1]) .expect("prefill"); engine - .decode_step_via_dispatch(&mut weights, &index, 2) + .decode_step_via_dispatch(&weights, &index, 2) .expect("decode"); // The decode_total / state_capture / recompute_hot timers // should all have at least one observation. diff --git a/crates/larql-kv/src/engines/turbo_quant/engine.rs b/crates/larql-kv/src/engines/turbo_quant/engine.rs index bebb83810..e5b29d4d4 100644 --- a/crates/larql-kv/src/engines/turbo_quant/engine.rs +++ b/crates/larql-kv/src/engines/turbo_quant/engine.rs @@ -287,6 +287,9 @@ pub struct TurboQuantEngine { /// when prefill routes through `coarse_prefill_with_state`. `None` /// means the engine took the legacy per-layer walk path. pub(super) kv_handle: Option, + /// Engine-owned f32 dequant scratch for the per-layer fallback (see + /// `MarkovResidualEngine::dequant_scratch`). Keeps `weights` immutable. + pub(super) dequant_scratch: larql_inference::DequantScratch, } impl TurboQuantEngine { @@ -303,6 +306,7 @@ impl TurboQuantEngine { profiling: false, profile: crate::profiler::EngineProfiler::default(), kv_handle: None, + dequant_scratch: larql_inference::DequantScratch::new(), } } @@ -337,7 +341,7 @@ impl TurboQuantEngine { // Decode step returns updated K/V (prior + new token). let (h_post_attn, updated_kv) = larql_inference::attention::run_attention_block_decode_step_auto( - weights, + larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch), &h, layer, Some(&prior_kv), @@ -428,10 +432,16 @@ impl KvEngine for TurboQuantEngine { self.layers.clear(); for layer in 0..num_layers { - let (h_post_attn, k, v) = run_attention_with_kv_backend(weights, &h, layer, be) - .ok_or_else(|| EngineError::BackendFailure { - details: "run_attention_with_kv_backend returned None".into(), - })?; + let (h_post_attn, k, v) = run_attention_with_kv_backend( + larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch), + &h, + layer, + be, + None, + ) + .ok_or_else(|| EngineError::BackendFailure { + details: "run_attention_with_kv_backend returned None".into(), + })?; self.layers .push(CompressedLayer::compress(&(k, v), &self.tq)); @@ -489,7 +499,7 @@ impl KvEngine for TurboQuantEngine { /// (`prefill_quant_cpu`) for backends without state-capture support. fn prefill_quant( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, _ffn: &dyn FfnBackend, index: &VectorIndex, token_ids: &[u32], @@ -513,7 +523,7 @@ impl KvEngine for TurboQuantEngine { fn decode_step_quant( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, _ffn: &dyn FfnBackend, index: &VectorIndex, token_id: u32, @@ -544,7 +554,7 @@ impl KvEngine for TurboQuantEngine { // stays here; only the per-layer compute is delegated. fn prefill_quant_via_executor( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, executor: &dyn larql_inference::layer_executor::LayerExecutor, ffn: &dyn FfnBackend, index: &VectorIndex, @@ -557,14 +567,19 @@ impl KvEngine for TurboQuantEngine { if matches!(executor.dispatch_kind(), ExecutorDispatchKind::Fused) { return self.prefill_quant(weights, ffn, index, token_ids, executor.backend()); } - ensure_attn_tensors_dequantised(weights, index); + ensure_attn_tensors_dequantised(&mut self.dequant_scratch, weights, index); let num_layers = weights.num_layers; let mut h = embed_tokens_pub(weights, token_ids); self.layers.clear(); for layer in 0..num_layers { let (h_out, kv) = executor - .run_prefill_layer(weights, layer, &h, ffn) + .run_prefill_layer( + larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch), + layer, + &h, + ffn, + ) .ok_or_else(|| EngineError::BackendFailure { details: "executor.run_prefill_layer returned None".into(), })?; @@ -578,7 +593,7 @@ impl KvEngine for TurboQuantEngine { fn decode_step_quant_via_executor( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, executor: &dyn larql_inference::layer_executor::LayerExecutor, ffn: &dyn FfnBackend, index: &VectorIndex, @@ -588,7 +603,7 @@ impl KvEngine for TurboQuantEngine { if matches!(executor.dispatch_kind(), ExecutorDispatchKind::Fused) { return self.decode_step_quant(weights, ffn, index, token_id, executor.backend()); } - ensure_attn_tensors_dequantised(weights, index); + ensure_attn_tensors_dequantised(&mut self.dequant_scratch, weights, index); let num_layers = weights.num_layers; let abs_position = self.abs_position; let mut h = embed_tokens_pub(weights, &[token_id]); @@ -596,7 +611,14 @@ impl KvEngine for TurboQuantEngine { for layer in 0..num_layers { let prior_kv = self.layers[layer].decompress(&self.tq); let (h_out, updated_kv) = executor - .run_decode_layer(weights, layer, &h, &prior_kv, abs_position, ffn) + .run_decode_layer( + larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch), + layer, + &h, + &prior_kv, + abs_position, + ffn, + ) .ok_or_else(|| EngineError::BackendFailure { details: "executor.run_decode_layer returned None".into(), })?; @@ -623,12 +645,12 @@ impl KvEngine for TurboQuantEngine { impl TurboQuantEngine { fn prefill_quant_cpu( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, index: &VectorIndex, token_ids: &[u32], backend: &dyn ComputeBackend, ) -> Option> { - ensure_attn_tensors_dequantised(weights, index); + ensure_attn_tensors_dequantised(&mut self.dequant_scratch, weights, index); let num_layers = weights.num_layers; let be = Some(backend); let mut h = embed_tokens_pub(weights, token_ids); @@ -639,7 +661,13 @@ impl TurboQuantEngine { .with_backend(backend); for layer in 0..num_layers { - let (h_post_attn, k, v) = run_attention_with_kv_backend(weights, &h, layer, be)?; + let (h_post_attn, k, v) = run_attention_with_kv_backend( + larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch), + &h, + layer, + be, + None, + )?; self.layers .push(CompressedLayer::compress(&(k, v), &self.tq)); @@ -664,13 +692,13 @@ impl TurboQuantEngine { fn decode_step_quant_cpu( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, index: &VectorIndex, token_id: u32, backend: &dyn ComputeBackend, ) -> Option> { use std::time::Instant; - ensure_attn_tensors_dequantised(weights, index); + ensure_attn_tensors_dequantised(&mut self.dequant_scratch, weights, index); let num_layers = weights.num_layers; let abs_position = self.abs_position; let timing = self.profiling; @@ -717,7 +745,7 @@ impl TurboQuantEngine { ) .or_else(|| { run_attention_block_decode_step_backend( - weights, + larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch), &h, layer, Some(&prior_kv), @@ -1119,13 +1147,13 @@ mod integration_tests { #[test] fn prefill_q4k_cpu_fallback_compresses_kv() { use larql_inference::ffn::NullFfn; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = TurboQuantEngine::new(4); let h = engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1, 2], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1, 2], &*backend) .expect("prefill_quant cpu fallback"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert_eq!( @@ -1139,17 +1167,17 @@ mod integration_tests { #[test] fn decode_step_quant_cpu_fallback_grows_compressed_cache() { use larql_inference::ffn::NullFfn; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = TurboQuantEngine::new(4); engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1], &*backend) .expect("prefill_quant"); let mem_before = engine.memory_bytes(); let h = engine - .decode_step_quant(&mut weights, &ffn, &index, 2, &*backend) + .decode_step_quant(&weights, &ffn, &index, 2, &*backend) .expect("decode_step_quant cpu fallback"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert!( @@ -1164,14 +1192,14 @@ mod integration_tests { fn prefill_quant_via_executor_compresses_kv() { use larql_inference::ffn::NullFfn; use larql_inference::layer_executor::LocalWalkExecutor; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let executor = LocalWalkExecutor::new(&*backend); let ffn = NullFfn; let mut engine = TurboQuantEngine::new(4); let h = engine - .prefill_quant_via_executor(&mut weights, &executor, &ffn, &index, &[0u32, 1, 2]) + .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32, 1, 2]) .expect("executor prefill"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert_eq!(engine.layers.len(), weights.num_layers); @@ -1182,18 +1210,18 @@ mod integration_tests { fn decode_step_quant_via_executor_grows_cache() { use larql_inference::ffn::NullFfn; use larql_inference::layer_executor::LocalWalkExecutor; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let executor = LocalWalkExecutor::new(&*backend); let ffn = NullFfn; let mut engine = TurboQuantEngine::new(4); engine - .prefill_quant_via_executor(&mut weights, &executor, &ffn, &index, &[0u32, 1]) + .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32, 1]) .expect("prefill"); let mem_before = engine.memory_bytes(); let h = engine - .decode_step_quant_via_executor(&mut weights, &executor, &ffn, &index, 2) + .decode_step_quant_via_executor(&weights, &executor, &ffn, &index, 2) .expect("decode"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert!(engine.memory_bytes() > mem_before); @@ -1204,16 +1232,16 @@ mod integration_tests { #[test] fn decode_step_quant_cpu_with_profiling_populates_summary() { use larql_inference::ffn::NullFfn; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = TurboQuantEngine::new(4).with_profiling(true); engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1], &*backend) .expect("prefill"); engine - .decode_step_quant(&mut weights, &ffn, &index, 2, &*backend) + .decode_step_quant(&weights, &ffn, &index, 2, &*backend) .expect("decode"); let summary = engine .stage_summary() @@ -1255,7 +1283,7 @@ mod integration_tests { #[test] fn executor_path_honors_ffn_parameter() { use larql_inference::layer_executor::LocalWalkExecutor; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let executor = LocalWalkExecutor::new(&*backend); @@ -1265,7 +1293,7 @@ mod integration_tests { }; let mut engine = TurboQuantEngine::new(4); engine - .prefill_quant_via_executor(&mut weights, &executor, &ffn, &index, &[0u32, 1, 2]) + .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32, 1, 2]) .expect("prefill via executor"); // Prefill runs FFN once per layer (single chunked sequence). let call_count = ffn.calls.load(std::sync::atomic::Ordering::SeqCst); @@ -1299,7 +1327,7 @@ mod integration_tests { #[test] fn fused_executor_short_circuits_prefill_to_legacy_path() { use larql_inference::ffn::NullFfn; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let executor = FusedStubExecutor { backend: larql_compute::CpuBackend, @@ -1307,7 +1335,7 @@ mod integration_tests { let ffn = NullFfn; let mut engine = TurboQuantEngine::new(4); let h = engine - .prefill_quant_via_executor(&mut weights, &executor, &ffn, &index, &[0u32, 1, 2]) + .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32, 1, 2]) .expect("fused-stub prefill should route through prefill_quant"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert_eq!(engine.layers.len(), weights.num_layers); @@ -1316,7 +1344,7 @@ mod integration_tests { #[test] fn fused_executor_short_circuits_decode_to_legacy_path() { use larql_inference::ffn::NullFfn; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let executor = FusedStubExecutor { backend: larql_compute::CpuBackend, @@ -1324,10 +1352,10 @@ mod integration_tests { let ffn = NullFfn; let mut engine = TurboQuantEngine::new(4); engine - .prefill_quant_via_executor(&mut weights, &executor, &ffn, &index, &[0u32, 1, 2]) + .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32, 1, 2]) .expect("prefill"); let h = engine - .decode_step_quant_via_executor(&mut weights, &executor, &ffn, &index, 3) + .decode_step_quant_via_executor(&weights, &executor, &ffn, &index, 3) .expect("fused-stub decode should route through decode_step_quant"); assert_eq!(h.shape(), &[1, weights.hidden_size]); } diff --git a/crates/larql-kv/src/engines/unlimited_context/dispatch.rs b/crates/larql-kv/src/engines/unlimited_context/dispatch.rs index 09f93b431..5645a55dc 100644 --- a/crates/larql-kv/src/engines/unlimited_context/dispatch.rs +++ b/crates/larql-kv/src/engines/unlimited_context/dispatch.rs @@ -28,7 +28,7 @@ impl UnlimitedContextEngine { /// append a single row in-place rather than re-allocating. pub(super) fn try_prefill_via_dispatch( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, index: &VectorIndex, token_ids: &[u32], ) -> Option> { @@ -92,7 +92,7 @@ impl UnlimitedContextEngine { /// count crosses `window_size`. pub(super) fn decode_step_via_dispatch( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, index: &VectorIndex, token_id: u32, ) -> Option> { @@ -197,9 +197,9 @@ mod tests { weights.hidden_size, ); let mut engine = UnlimitedContextEngine::with_backend(4, cpu_engine_backend()); - let mut w = weights; + let w = weights; assert!(engine - .try_prefill_via_dispatch(&mut w, &empty_index, &[0u32, 1]) + .try_prefill_via_dispatch(&w, &empty_index, &[0u32, 1]) .is_none()); assert!(engine.kv_handle.is_none()); assert!(engine.current_window_kv.is_none()); @@ -211,9 +211,9 @@ mod tests { // doesn't populate `current_window_kv`; Metal's kv cache is the // truth. set_w10_disable(false); - let (mut engine, mut weights, index) = fixture(4); + let (mut engine, weights, index) = fixture(4); let h = engine - .try_prefill_via_dispatch(&mut weights, &index, &[0u32, 1, 2]) + .try_prefill_via_dispatch(&weights, &index, &[0u32, 1, 2]) .expect("prefill"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert!(engine.current_window_kv.is_none()); @@ -228,8 +228,8 @@ mod tests { // W10 off → engine pre-allocates `[window_cap, kv_dim]` per // layer and copies the prefill K/V rows in. set_w10_disable(true); - let (mut engine, mut weights, index) = fixture(8); - let res = engine.try_prefill_via_dispatch(&mut weights, &index, &[0u32, 1, 2]); + let (mut engine, weights, index) = fixture(8); + let res = engine.try_prefill_via_dispatch(&weights, &index, &[0u32, 1, 2]); set_w10_disable(false); res.expect("prefill"); let kv = engine @@ -245,10 +245,10 @@ mod tests { #[test] fn decode_step_via_dispatch_without_prefill_returns_none() { set_w10_disable(false); - let (mut engine, mut weights, index) = fixture(4); + let (mut engine, weights, index) = fixture(4); // kv_handle is None → early return. assert!(engine - .decode_step_via_dispatch(&mut weights, &index, 0) + .decode_step_via_dispatch(&weights, &index, 0) .is_none()); } @@ -258,13 +258,13 @@ mod tests { // The dispatch decode runs through the `if !matches!(mask, HOnly)` // skip branch and still bumps the per-step counters. set_w10_disable(false); - let (mut engine, mut weights, index) = fixture(4); + let (mut engine, weights, index) = fixture(4); engine - .try_prefill_via_dispatch(&mut weights, &index, &[0u32, 1]) + .try_prefill_via_dispatch(&weights, &index, &[0u32, 1]) .expect("prefill"); let kv_len_before = engine.current_window_kv_len; let h = engine - .decode_step_via_dispatch(&mut weights, &index, 2) + .decode_step_via_dispatch(&weights, &index, 2) .expect("decode"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert_eq!(engine.current_window_kv_len, kv_len_before + 1); @@ -280,12 +280,12 @@ mod tests { // W10 off → Full mask. Each layer's K/V row blits into the // pre-allocated window_kv buffer at `current_window_kv_len`. set_w10_disable(true); - let (mut engine, mut weights, index) = fixture(8); + let (mut engine, weights, index) = fixture(8); engine - .try_prefill_via_dispatch(&mut weights, &index, &[0u32, 1]) + .try_prefill_via_dispatch(&weights, &index, &[0u32, 1]) .expect("prefill"); let kv_len_before = engine.current_window_kv_len; - let res = engine.decode_step_via_dispatch(&mut weights, &index, 2); + let res = engine.decode_step_via_dispatch(&weights, &index, 2); set_w10_disable(false); res.expect("decode"); assert_eq!(engine.current_window_kv_len, kv_len_before + 1); @@ -307,13 +307,13 @@ mod tests { // window=2: prefill 2 → token count == window → close_window // fires inside the dispatch decode and resets the window. set_w10_disable(false); - let (mut engine, mut weights, index) = fixture(2); + let (mut engine, weights, index) = fixture(2); engine - .try_prefill_via_dispatch(&mut weights, &index, &[0u32]) + .try_prefill_via_dispatch(&weights, &index, &[0u32]) .expect("prefill"); let cp_before = engine.checkpoints.len(); engine - .decode_step_via_dispatch(&mut weights, &index, 1) + .decode_step_via_dispatch(&weights, &index, 1) .expect("decode that triggers window-close"); // close_window() emits a checkpoint and resets // current_window_tokens (the legacy emit path). diff --git a/crates/larql-kv/src/engines/unlimited_context/engine.rs b/crates/larql-kv/src/engines/unlimited_context/engine.rs index 765d8fd8f..eb15f2ea8 100644 --- a/crates/larql-kv/src/engines/unlimited_context/engine.rs +++ b/crates/larql-kv/src/engines/unlimited_context/engine.rs @@ -95,6 +95,9 @@ pub struct UnlimitedContextEngine { /// prefill routes through `coarse_prefill_with_state`. `None` = /// legacy CPU walk path. pub(super) kv_handle: Option, + /// Engine-owned f32 dequant scratch for the per-layer walk fallback + /// (see `MarkovResidualEngine::dequant_scratch`). Keeps `weights` immutable. + pub(super) dequant_scratch: larql_inference::DequantScratch, } impl UnlimitedContextEngine { @@ -117,6 +120,7 @@ impl UnlimitedContextEngine { profiling: false, profile: crate::profiler::EngineProfiler::default(), kv_handle: None, + dequant_scratch: larql_inference::DequantScratch::new(), } } @@ -191,7 +195,7 @@ impl UnlimitedContextEngine { }; let out = rs_extend_from_checkpoint_backend( - weights, + larql_inference::WeightsView::dense(weights), tokens, prior, abs_offset, @@ -289,9 +293,9 @@ impl UnlimitedContextEngine { let abs_start = self.abs_offset + self.current_window_tokens.len(); let prof = self.profiling.then_some(&mut self.profile); - let out = rs_extend_from_checkpoint_quant( - weights, index, chunk, prior, abs_start, backend, prof, - )?; + let view = larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch); + let out = + rs_extend_from_checkpoint_quant(view, index, chunk, prior, abs_start, backend, prof)?; self.last_hidden = Some(out.last_hidden); // CPU walk path returns narrow `[n, kv_dim]` arrays — counter @@ -368,7 +372,7 @@ impl UnlimitedContextEngine { if use_inplace { let last = rs_extend_inplace( - weights, + larql_inference::WeightsView::dense(weights), chunk, &mut prior, prior_len, @@ -382,7 +386,7 @@ impl UnlimitedContextEngine { self.current_window_kv = Some(prior); } else { let out = rs_extend_from_checkpoint_backend( - weights, + larql_inference::WeightsView::dense(weights), chunk, prior, abs_start, @@ -587,7 +591,7 @@ impl KvEngine for UnlimitedContextEngine { /// from captured per-layer state (W1-GPU) or computed via walk. fn prefill_quant( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, _ffn: &dyn FfnBackend, index: &VectorIndex, token_ids: &[u32], @@ -600,7 +604,7 @@ impl KvEngine for UnlimitedContextEngine { return Ok(hidden); } self.kv_handle = None; - ensure_attn_tensors_dequantised(weights, index); + ensure_attn_tensors_dequantised(&mut self.dequant_scratch, weights, index); self.process_quant(weights, index, token_ids, backend) .ok_or_else(|| EngineError::BackendFailure { details: "process_quant returned None during prefill_quant".into(), @@ -614,7 +618,7 @@ impl KvEngine for UnlimitedContextEngine { fn decode_step_quant( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, _ffn: &dyn FfnBackend, index: &VectorIndex, token_id: u32, @@ -627,7 +631,7 @@ impl KvEngine for UnlimitedContextEngine { details: "decode_step_via_dispatch returned None".into(), }); } - ensure_attn_tensors_dequantised(weights, index); + ensure_attn_tensors_dequantised(&mut self.dequant_scratch, weights, index); self.process_quant(weights, index, &[token_id], backend) .ok_or_else(|| EngineError::BackendFailure { details: "process_quant returned None during decode_step_quant".into(), @@ -652,7 +656,7 @@ impl KvEngine for UnlimitedContextEngine { // owns per-layer compute; window state is engine state. fn prefill_quant_via_executor( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, executor: &dyn larql_inference::layer_executor::LayerExecutor, ffn: &dyn FfnBackend, index: &VectorIndex, @@ -668,7 +672,7 @@ impl KvEngine for UnlimitedContextEngine { if matches!(executor.dispatch_kind(), ExecutorDispatchKind::Fused) { return self.prefill_quant(weights, ffn, index, token_ids, executor.backend()); } - ensure_attn_tensors_dequantised(weights, index); + ensure_attn_tensors_dequantised(&mut self.dequant_scratch, weights, index); self.process_via_executor(weights, executor, ffn, token_ids) .ok_or_else(|| EngineError::BackendFailure { details: "process_via_executor returned None during prefill_quant_via_executor" @@ -683,7 +687,7 @@ impl KvEngine for UnlimitedContextEngine { fn decode_step_quant_via_executor( &mut self, - weights: &mut ModelWeights, + weights: &ModelWeights, executor: &dyn larql_inference::layer_executor::LayerExecutor, ffn: &dyn FfnBackend, index: &VectorIndex, @@ -693,7 +697,7 @@ impl KvEngine for UnlimitedContextEngine { if matches!(executor.dispatch_kind(), ExecutorDispatchKind::Fused) { return self.decode_step_quant(weights, ffn, index, token_id, executor.backend()); } - ensure_attn_tensors_dequantised(weights, index); + ensure_attn_tensors_dequantised(&mut self.dequant_scratch, weights, index); self.process_via_executor(weights, executor, ffn, &[token_id]) .ok_or_else(|| EngineError::BackendFailure { details: "process_via_executor returned None during decode_step_quant_via_executor" @@ -771,8 +775,14 @@ impl UnlimitedContextEngine { let mut h = embed_tokens_pub(weights, &[token_id]); for (layer, kv_slot) in kv_cache.iter_mut().enumerate() { - let (h_out, new_kv) = - executor.run_decode_layer(weights, layer, &h, kv_slot, abs_position, ffn)?; + let (h_out, new_kv) = executor.run_decode_layer( + larql_inference::WeightsView::with_scratch(weights, &self.dequant_scratch), + layer, + &h, + kv_slot, + abs_position, + ffn, + )?; h = h_out; *kv_slot = new_kv; } @@ -983,13 +993,13 @@ mod tests { fn prefill_quant_cpu_runs_via_dequant_path() { use larql_inference::ffn::NullFfn; use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = UnlimitedContextEngine::new(512); let h = engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1, 2], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1, 2], &*backend) .expect("prefill_quant Q4K cpu fallback"); assert_eq!(h.shape(), &[1, weights.hidden_size]); } @@ -998,16 +1008,16 @@ mod tests { fn decode_step_quant_cpu_extends_state() { use larql_inference::ffn::NullFfn; use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = UnlimitedContextEngine::new(512); engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1], &*backend) .expect("prefill_quant"); let h = engine - .decode_step_quant(&mut weights, &ffn, &index, 2, &*backend) + .decode_step_quant(&weights, &ffn, &index, 2, &*backend) .expect("decode_step_quant Q4K cpu fallback"); assert_eq!(h.shape(), &[1, weights.hidden_size]); } @@ -1067,14 +1077,14 @@ mod tests { fn decode_step_quant_without_prefill_returns_none() { use larql_inference::ffn::NullFfn; use larql_inference::test_utils::{make_test_q4k_vindex, make_test_q4k_weights}; - let mut weights = make_test_q4k_weights(); + let weights = make_test_q4k_weights(); let index = make_test_q4k_vindex(&weights); let backend = larql_compute::cpu_backend(); let ffn = NullFfn; let mut engine = UnlimitedContextEngine::new(512); // No prefill → decode falls through fast-path checks and returns None // (or some empty hidden) without panicking. - let _ = engine.decode_step_quant(&mut weights, &ffn, &index, 0, &*backend); + let _ = engine.decode_step_quant(&weights, &ffn, &index, 0, &*backend); } // ── Public utility methods (stats, replay_window, summary) ──────────── @@ -1161,14 +1171,14 @@ mod tests { use larql_inference::ffn::NullFfn; use larql_inference::layer_executor::LocalWalkExecutor; use larql_inference::test_utils::make_test_weights; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let executor = LocalWalkExecutor::new(&*backend); let ffn = NullFfn; let mut engine = UnlimitedContextEngine::new(512); let h = engine - .prefill_quant_via_executor(&mut weights, &executor, &ffn, &index, &[0u32, 1, 2]) + .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32, 1, 2]) .expect("executor prefill"); assert_eq!(h.shape(), &[1, weights.hidden_size]); assert!(engine.memory_bytes() > 0); @@ -1179,17 +1189,17 @@ mod tests { use larql_inference::ffn::NullFfn; use larql_inference::layer_executor::LocalWalkExecutor; use larql_inference::test_utils::make_test_weights; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let executor = LocalWalkExecutor::new(&*backend); let ffn = NullFfn; let mut engine = UnlimitedContextEngine::new(512); engine - .prefill_quant_via_executor(&mut weights, &executor, &ffn, &index, &[0u32, 1]) + .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32, 1]) .expect("prefill"); let h = engine - .decode_step_quant_via_executor(&mut weights, &executor, &ffn, &index, 2) + .decode_step_quant_via_executor(&weights, &executor, &ffn, &index, 2) .expect("decode"); assert_eq!(h.shape(), &[1, weights.hidden_size]); } @@ -1201,16 +1211,16 @@ mod tests { fn process_quant_with_profiling_populates_summary() { use larql_inference::ffn::NullFfn; use larql_inference::test_utils::make_test_weights; - let mut weights = make_test_weights(); + let weights = make_test_weights(); 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); engine - .prefill_quant(&mut weights, &ffn, &index, &[0u32, 1], &*backend) + .prefill_quant(&weights, &ffn, &index, &[0u32, 1], &*backend) .expect("prefill"); engine - .decode_step_quant(&mut weights, &ffn, &index, 2, &*backend) + .decode_step_quant(&weights, &ffn, &index, 2, &*backend) .expect("decode"); let summary = engine .stage_summary() @@ -1252,7 +1262,7 @@ mod tests { fn executor_path_honors_ffn_parameter() { use larql_inference::layer_executor::LocalWalkExecutor; use larql_inference::test_utils::make_test_weights; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let executor = LocalWalkExecutor::new(&*backend); @@ -1263,7 +1273,7 @@ mod tests { }; let mut engine = UnlimitedContextEngine::new(512); engine - .prefill_quant_via_executor(&mut weights, &executor, &ffn, &index, &[0u32, 1, 2]) + .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32, 1, 2]) .expect("prefill via executor"); let call_count = ffn.calls.load(std::sync::atomic::Ordering::SeqCst); @@ -1284,7 +1294,7 @@ mod tests { use larql_inference::ffn::NullFfn; use larql_inference::layer_executor::LocalWalkExecutor; use larql_inference::test_utils::make_test_weights; - let mut weights = make_test_weights(); + let weights = make_test_weights(); let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let executor = LocalWalkExecutor::new(&*backend); @@ -1294,7 +1304,7 @@ mod tests { // branch in `extend_current_via_executor`. let mut engine = UnlimitedContextEngine::new(2); engine - .prefill_quant_via_executor(&mut weights, &executor, &ffn, &index, &[0u32, 1, 2, 3]) + .prefill_quant_via_executor(&weights, &executor, &ffn, &index, &[0u32, 1, 2, 3]) .expect("prefill 4 tokens through executor"); let stats = engine.stats(&weights); assert!( diff --git a/crates/larql-kv/src/engines/unlimited_context/extend.rs b/crates/larql-kv/src/engines/unlimited_context/extend.rs index 90d584056..615b3726a 100644 --- a/crates/larql-kv/src/engines/unlimited_context/extend.rs +++ b/crates/larql-kv/src/engines/unlimited_context/extend.rs @@ -27,7 +27,7 @@ pub struct ExtendOutput { /// /// `abs_start` is the absolute position of the *first new token*. pub fn rs_extend_from_checkpoint( - weights: &ModelWeights, + weights: larql_inference::WeightsView, token_ids: &[u32], prior_kv: Vec, abs_start: usize, @@ -50,7 +50,7 @@ pub fn rs_extend_from_checkpoint( /// full window — a real overhead on growing caches. #[allow(clippy::too_many_arguments)] pub fn rs_extend_from_checkpoint_backend( - weights: &ModelWeights, + weights: larql_inference::WeightsView, token_ids: &[u32], prior_kv: Vec, abs_start: usize, @@ -72,7 +72,7 @@ pub fn rs_extend_from_checkpoint_backend( for (i, &token_id) in token_ids.iter().enumerate() { let abs_position = abs_start + i; - let mut h = embed_tokens_pub(weights, &[token_id]); + let mut h = embed_tokens_pub(&weights, &[token_id]); for (layer, kv_slot) in kv_cache.iter_mut().enumerate() { let kv_entry: Option<&SharedKV> = if kv_slot.0.shape()[0] > 0 { @@ -92,9 +92,17 @@ pub fn rs_extend_from_checkpoint_backend( index.map(|v| v as &dyn larql_compute::KvIndex), )?; - let bffn = BackendFfn { weights, backend }; - let h_out = - crate::engines::layer_ffn_or_moe(weights, &h_post_attn, layer, &bffn, moe_ffn); + 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, + ); h = h_out; *kv_slot = new_kv; } @@ -134,7 +142,7 @@ pub fn rs_extend_from_checkpoint_backend( /// form (engine-level A/B test), only the cache representation differs. #[allow(clippy::too_many_arguments)] pub fn rs_extend_inplace( - weights: &ModelWeights, + weights: larql_inference::WeightsView, token_ids: &[u32], kv_cache: &mut [SharedKV], prior_len: usize, @@ -156,7 +164,7 @@ pub fn rs_extend_inplace( // Logical row count of every buffer at the start of this token: the // prior window length plus the tokens already appended this call. let len = prior_len + i; - let mut h = embed_tokens_pub(weights, &[token_id]); + let mut h = embed_tokens_pub(&weights, &[token_id]); for (layer, (k_buf, v_buf)) in kv_cache.iter_mut().enumerate() { let h_post_attn = @@ -196,9 +204,17 @@ pub fn rs_extend_inplace( } }; - let bffn = BackendFfn { weights, backend }; - let h_out = - crate::engines::layer_ffn_or_moe(weights, &h_post_attn, layer, &bffn, moe_ffn); + 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, + ); h = h_out; } @@ -215,7 +231,7 @@ pub fn rs_extend_inplace( /// uses the dequantised f32 tensors already inserted by /// `ensure_attn_tensors_dequantised`. Call that before this function. pub fn rs_extend_from_checkpoint_quant( - weights: &ModelWeights, + weights: larql_inference::WeightsView, index: &VectorIndex, token_ids: &[u32], prior_kv: Vec, @@ -238,8 +254,9 @@ pub fn rs_extend_from_checkpoint_quant( // Hoist WalkFfn out of both loops. Previously this rebuilt the // WalkFfn once per (token, layer) — N×34 times per extend call. // It's now once total. WalkFfn carries no per-(token,layer) state. - let walk_ffn = WalkFfn::from_config(weights, index, WalkFfnConfig::dense(num_layers)) - .with_backend(backend); + let walk_ffn = + WalkFfn::from_config(weights.canonical(), index, WalkFfnConfig::dense(num_layers)) + .with_backend(backend); // Per-stage timing. `LARQL_INSTRUMENT_UNLIMITED=1` enables the // verbose stderr line; `profiler` is the structured channel used by @@ -264,7 +281,7 @@ pub fn rs_extend_from_checkpoint_quant( } else { None }; - let mut h = embed_tokens_pub(weights, &[token_id]); + let mut h = embed_tokens_pub(&weights, &[token_id]); if let Some(start) = t_embed_start { t_embed += start.elapsed().as_secs_f64() * 1e6; } @@ -284,7 +301,7 @@ pub fn rs_extend_from_checkpoint_quant( None }; let attn_native = larql_inference::vindex::attention_decode_step_native( - weights, + weights.canonical(), index, backend, &h, @@ -319,7 +336,7 @@ pub fn rs_extend_from_checkpoint_quant( None }; let ffn_native = larql_inference::vindex::ffn_decode_step_native( - weights, + weights.canonical(), index, backend, &h_post_attn, @@ -329,7 +346,7 @@ pub fn rs_extend_from_checkpoint_quant( t_ffn_helper_misses += 1; } let h_out = ffn_native.unwrap_or_else(|| { - let (h, _) = run_ffn(weights, &h_post_attn, layer, &walk_ffn, false); + let (h, _) = run_ffn(&weights, &h_post_attn, layer, &walk_ffn, false); h }); if let Some(start) = t_ffn_start { @@ -428,7 +445,8 @@ mod tests { fn extend_empty_tokens_returns_none() { let weights = make_test_weights(); let prior = empty_prior(&weights); - let result = rs_extend_from_checkpoint(&weights, &[], prior, 0); + let result = + rs_extend_from_checkpoint(larql_inference::WeightsView::dense(&weights), &[], prior, 0); assert!(result.is_none(), "empty token_ids should return None"); } @@ -436,7 +454,12 @@ mod tests { fn extend_wrong_prior_len_returns_none() { let weights = make_test_weights(); // prior has 0 layers but model has 2 — mismatch - let result = rs_extend_from_checkpoint(&weights, &[0u32], Vec::new(), 0); + let result = rs_extend_from_checkpoint( + larql_inference::WeightsView::dense(&weights), + &[0u32], + Vec::new(), + 0, + ); assert!(result.is_none(), "prior length mismatch should return None"); } @@ -444,8 +467,13 @@ mod tests { fn extend_single_token_from_empty_prior() { let weights = make_test_weights(); let prior = empty_prior(&weights); - let output = rs_extend_from_checkpoint(&weights, &[0u32], prior, 0) - .expect("single token extend should succeed"); + let output = rs_extend_from_checkpoint( + larql_inference::WeightsView::dense(&weights), + &[0u32], + prior, + 0, + ) + .expect("single token extend should succeed"); assert_eq!(output.last_hidden.shape(), &[1, weights.hidden_size]); assert!(output.last_hidden.iter().all(|v| v.is_finite())); } @@ -454,8 +482,13 @@ mod tests { fn extend_kv_cache_grows_with_each_token() { let weights = make_test_weights(); let prior = empty_prior(&weights); - let output = - rs_extend_from_checkpoint(&weights, &[0u32, 1, 2], prior, 0).expect("3-token extend"); + let output = rs_extend_from_checkpoint( + larql_inference::WeightsView::dense(&weights), + &[0u32, 1, 2], + prior, + 0, + ) + .expect("3-token extend"); // After 3 tokens from empty prior, K has 3 rows per layer let kv_dim = weights.num_kv_heads * weights.head_dim; for (k, v) in &output.kv_cache { @@ -468,8 +501,13 @@ mod tests { fn extend_checkpoint_is_last_row_of_kv_cache() { let weights = make_test_weights(); let prior = empty_prior(&weights); - let output = - rs_extend_from_checkpoint(&weights, &[0u32, 1], prior, 0).expect("2-token extend"); + let output = rs_extend_from_checkpoint( + larql_inference::WeightsView::dense(&weights), + &[0u32, 1], + prior, + 0, + ) + .expect("2-token extend"); // new_checkpoint should be the last row of each K/V for (layer, ((k_cache, v_cache), (k_ckpt, v_ckpt))) in output .kv_cache @@ -494,8 +532,20 @@ mod tests { fn extend_abs_start_shifts_rope() { let weights = make_test_weights(); let prior = empty_prior(&weights); - let out0 = rs_extend_from_checkpoint(&weights, &[0u32], prior.clone(), 0).unwrap(); - let out5 = rs_extend_from_checkpoint(&weights, &[0u32], prior, 5).unwrap(); + let out0 = rs_extend_from_checkpoint( + larql_inference::WeightsView::dense(&weights), + &[0u32], + prior.clone(), + 0, + ) + .unwrap(); + let out5 = rs_extend_from_checkpoint( + larql_inference::WeightsView::dense(&weights), + &[0u32], + prior, + 5, + ) + .unwrap(); // Different abs_start → different RoPE → different K let k0 = &out0.kv_cache[0].0; let k5 = &out5.kv_cache[0].0; @@ -510,7 +560,13 @@ mod tests { fn extend_output_logits_are_finite() { let weights = make_test_weights(); let prior = empty_prior(&weights); - let output = rs_extend_from_checkpoint(&weights, &[0u32], prior, 0).unwrap(); + let output = rs_extend_from_checkpoint( + larql_inference::WeightsView::dense(&weights), + &[0u32], + prior, + 0, + ) + .unwrap(); let logits = hidden_to_raw_logits(&weights, &output.last_hidden); assert!(logits.iter().all(|v| v.is_finite())); } @@ -520,10 +576,21 @@ mod tests { // Extending from a non-empty checkpoint should not panic and should be finite. let weights = make_test_weights(); let prior = empty_prior(&weights); - let first = rs_extend_from_checkpoint(&weights, &[0u32], prior, 0).unwrap(); + let first = rs_extend_from_checkpoint( + larql_inference::WeightsView::dense(&weights), + &[0u32], + prior, + 0, + ) + .unwrap(); // Use the checkpoint from the first extend as the prior for the second - let second = rs_extend_from_checkpoint(&weights, &[1u32], first.new_checkpoint.clone(), 1) - .expect("extend from non-empty prior"); + let second = rs_extend_from_checkpoint( + larql_inference::WeightsView::dense(&weights), + &[1u32], + first.new_checkpoint.clone(), + 1, + ) + .expect("extend from non-empty prior"); assert_eq!(second.last_hidden.shape(), &[1, weights.hidden_size]); assert!(second.last_hidden.iter().all(|v| v.is_finite())); } @@ -536,7 +603,15 @@ mod tests { let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let prior = empty_prior(&weights); - let out = rs_extend_from_checkpoint_quant(&weights, &index, &[], prior, 0, &*backend, None); + let out = rs_extend_from_checkpoint_quant( + larql_inference::WeightsView::dense(&weights), + &index, + &[], + prior, + 0, + &*backend, + None, + ); assert!(out.is_none(), "empty token_ids should return None"); } @@ -546,7 +621,7 @@ mod tests { let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let out = rs_extend_from_checkpoint_quant( - &weights, + larql_inference::WeightsView::dense(&weights), &index, &[0u32], Vec::new(), @@ -564,7 +639,7 @@ mod tests { let backend = larql_compute::cpu_backend(); let prior = empty_prior(&weights); let out = rs_extend_from_checkpoint_quant( - &weights, + larql_inference::WeightsView::dense(&weights), &index, &[0u32, 1, 2], prior, @@ -590,7 +665,7 @@ mod tests { let index = larql_inference::test_utils::make_test_vindex(&weights); let backend = larql_compute::cpu_backend(); let first = rs_extend_from_checkpoint_quant( - &weights, + larql_inference::WeightsView::dense(&weights), &index, &[0u32, 1], empty_prior(&weights), @@ -600,7 +675,7 @@ mod tests { ) .expect("first extend"); let second = rs_extend_from_checkpoint_quant( - &weights, + larql_inference::WeightsView::dense(&weights), &index, &[2u32], first.kv_cache.clone(), diff --git a/crates/larql-kv/src/generation.rs b/crates/larql-kv/src/generation.rs index 31f1c59b3..9907913e9 100644 --- a/crates/larql-kv/src/generation.rs +++ b/crates/larql-kv/src/generation.rs @@ -450,7 +450,7 @@ where /// state to get logits. #[allow(clippy::too_many_arguments)] pub fn kv_prefill_run( - weights: &ModelWeights, + weights: larql_inference::WeightsView, ffn: &dyn FfnBackend, prompt_ids: &[u32], window: Option, @@ -466,25 +466,25 @@ pub fn kv_prefill_run( None => KvCache::with_layers(num_layers), }; - let mut h = embed_tokens_pub(weights, prompt_ids); + 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); + 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)?; + 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 (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); + 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; @@ -526,7 +526,7 @@ pub fn kv_decode_step_run( let kv_entry = cache.layers[layer].as_ref(); let (mut h_post_attn, new_kv) = run_attention_block_decode_step_backend( - weights, + larql_inference::WeightsView::dense(weights), &h_step, layer, kv_entry, @@ -567,11 +567,17 @@ fn generate_cached_hooked_inner( } // ── Phase 1: prefill ── - let (last_hidden, mut cache) = - match kv_prefill_run(weights, ffn, prompt_ids, window, backend, hook) { - Some(t) => t, - None => return Vec::new(), - }; + let (last_hidden, mut cache) = match kv_prefill_run( + larql_inference::WeightsView::dense(weights), + ffn, + prompt_ids, + window, + backend, + hook, + ) { + Some(t) => t, + None => return Vec::new(), + }; let first = match argmax_next_token(weights, tokenizer, &last_hidden) { Some(t) => t, @@ -722,8 +728,13 @@ where let mut h = embed_tokens_pub(weights, prompt_ids); for layer in 0..num_layers { - let (h_post_attn, k_rope, v) = match run_attention_with_kv_backend(weights, &h, layer, None) - { + let (h_post_attn, k_rope, v) = match run_attention_with_kv_backend( + larql_inference::WeightsView::dense(weights), + &h, + layer, + None, + None, + ) { Some(t) => t, None => return Vec::new(), }; @@ -755,7 +766,7 @@ where for layer in 0..num_layers { let kv_entry = cache.layers[layer].as_ref(); let (h_post_attn, new_kv) = match run_attention_block_decode_step_backend( - weights, + larql_inference::WeightsView::dense(weights), &h_step, layer, kv_entry, @@ -872,12 +883,19 @@ mod tests { details: "test stub: fail_prefill set".into(), }); } - let (hidden, cache) = - kv_prefill_run(weights, ffn, token_ids, None, None, &mut NoopHook).ok_or_else( - || larql_inference::kv_engine::EngineError::BackendFailure { - details: "kv_prefill_run returned None".into(), - }, - )?; + let (hidden, cache) = kv_prefill_run( + larql_inference::WeightsView::dense(weights), + ffn, + token_ids, + 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) } @@ -927,10 +945,19 @@ mod tests { }); } let ids: Vec = (0..initial_hidden.nrows() as u32).collect(); - let (hidden, cache) = kv_prefill_run(weights, ffn, &ids, None, None, &mut NoopHook) - .ok_or_else(|| larql_inference::kv_engine::EngineError::BackendFailure { + let (hidden, cache) = kv_prefill_run( + larql_inference::WeightsView::dense(weights), + ffn, + &ids, + 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) } @@ -1265,9 +1292,15 @@ mod tests { let weights = larql_inference::test_utils::make_synthetic_e2b_like_weights(); let ffn = WeightFfn { weights: &weights }; let prompt = [0u32, 1, 2]; - let (last_hidden, cache) = - kv_prefill_run(&weights, &ffn, &prompt, None, None, &mut NoopHook) - .expect("PLE-arch prefill should not fail"); + let (last_hidden, cache) = kv_prefill_run( + larql_inference::WeightsView::dense(&weights), + &ffn, + &prompt, + None, + None, + &mut NoopHook, + ) + .expect("PLE-arch prefill should not fail"); assert_eq!(last_hidden.shape(), &[1, weights.hidden_size]); assert!( last_hidden.iter().all(|v| v.is_finite()), @@ -1287,9 +1320,15 @@ mod tests { let weights = larql_inference::test_utils::make_synthetic_e2b_like_weights(); let ffn = WeightFfn { weights: &weights }; let prompt = [0u32, 1]; - let (_h_prefill, mut cache) = - kv_prefill_run(&weights, &ffn, &prompt, None, None, &mut NoopHook) - .expect("PLE-arch prefill should not fail"); + let (_h_prefill, mut cache) = kv_prefill_run( + larql_inference::WeightsView::dense(&weights), + &ffn, + &prompt, + None, + None, + &mut NoopHook, + ) + .expect("PLE-arch prefill should not fail"); for step in 0..3 { let h_step = kv_decode_step_run(&weights, &ffn, &mut cache, 0u32, None, &mut NoopHook) diff --git a/crates/larql-kv/src/lib.rs b/crates/larql-kv/src/lib.rs index ec1c3e213..6f069771a 100644 --- a/crates/larql-kv/src/lib.rs +++ b/crates/larql-kv/src/lib.rs @@ -981,15 +981,15 @@ mod compliance_tests { }; // Build a separate &mut binding for the `prefill_quant` call. - let mut weights_for_q4k = larql_inference::test_utils::make_test_weights(); - let out = engine.prefill_quant(&mut weights_for_q4k, &ffn, &index, &[1, 2, 3], &*backend); + let weights_for_q4k = larql_inference::test_utils::make_test_weights(); + let out = engine.prefill_quant(&weights_for_q4k, &ffn, &index, &[1, 2, 3], &*backend); assert!(out.is_ok()); assert_eq!( engine.prefill_calls, 1, "default prefill_quant must call prefill" ); - let out = engine.decode_step_quant(&mut weights_for_q4k, &ffn, &index, 4, &*backend); + let out = engine.decode_step_quant(&weights_for_q4k, &ffn, &index, 4, &*backend); assert!(out.is_ok()); assert_eq!( engine.decode_calls, 1, diff --git a/crates/larql-kv/src/vindex_compare.rs b/crates/larql-kv/src/vindex_compare.rs index c466668c2..368641017 100644 --- a/crates/larql-kv/src/vindex_compare.rs +++ b/crates/larql-kv/src/vindex_compare.rs @@ -160,9 +160,15 @@ pub fn forward_to_logits_traced( // positions are processed. let walk_ffn = WalkFfn::new_unlimited(weights, index).with_dispatch_trace(); - if let Some((h_new, _, kv_out)) = - run_layer_with_ffn(weights, &h, layer, &walk_ffn, false, None, shared_kv) - { + if let Some((h_new, _, kv_out)) = run_layer_with_ffn( + larql_inference::WeightsView::dense(weights), + &h, + layer, + &walk_ffn, + false, + None, + shared_kv, + ) { h = h_new; if let Some(kv) = kv_out { kv_cache.insert(layer, kv); diff --git a/crates/larql-kv/tests/dispatch_parity.rs b/crates/larql-kv/tests/dispatch_parity.rs index 82e0933ec..15bb6217d 100644 --- a/crates/larql-kv/tests/dispatch_parity.rs +++ b/crates/larql-kv/tests/dispatch_parity.rs @@ -34,11 +34,24 @@ fn prefill_via_dispatch_matches_legacy_kv_prefill_run() { let ffn = WeightFfn { weights: &weights }; let prompt = vec![0u32, 1, 2, 3]; - let (h_trait, _handles) = - kv_prefill_via_dispatch(&backend, &weights, &ffn, &prompt, None, None).expect("prefill"); - let (h_legacy, _cache) = - kv_prefill_run(&weights, &ffn, &prompt, None, Some(&backend), &mut NoopHook) - .expect("legacy prefill"); + let (h_trait, _handles) = kv_prefill_via_dispatch( + &backend, + larql_inference::WeightsView::dense(&weights), + &ffn, + &prompt, + None, + None, + ) + .expect("prefill"); + let (h_legacy, _cache) = kv_prefill_run( + larql_inference::WeightsView::dense(&weights), + &ffn, + &prompt, + None, + Some(&backend), + &mut NoopHook, + ) + .expect("legacy prefill"); assert_eq!( h_trait, h_legacy, @@ -54,10 +67,17 @@ fn prefill_via_dispatch_windowed_matches_legacy() { let prompt = vec![0u32, 1, 2, 3, 4]; let window = Some(2); - let (h_trait, _handles) = - kv_prefill_via_dispatch(&backend, &weights, &ffn, &prompt, window, None).expect("prefill"); + let (h_trait, _handles) = kv_prefill_via_dispatch( + &backend, + larql_inference::WeightsView::dense(&weights), + &ffn, + &prompt, + window, + None, + ) + .expect("prefill"); let (h_legacy, _cache) = kv_prefill_run( - &weights, + larql_inference::WeightsView::dense(&weights), &ffn, &prompt, window, @@ -79,17 +99,31 @@ fn decode_step_via_dispatch_matches_legacy_kv_decode_step_run() { let ffn = WeightFfn { weights: &weights }; let prompt = vec![0u32, 1, 2]; - let (_, mut handles) = - kv_prefill_via_dispatch(&backend, &weights, &ffn, &prompt, None, None).unwrap(); - let (_, mut cache) = - kv_prefill_run(&weights, &ffn, &prompt, None, Some(&backend), &mut NoopHook).unwrap(); + let (_, mut handles) = kv_prefill_via_dispatch( + &backend, + larql_inference::WeightsView::dense(&weights), + &ffn, + &prompt, + None, + None, + ) + .unwrap(); + let (_, mut cache) = kv_prefill_run( + larql_inference::WeightsView::dense(&weights), + &ffn, + &prompt, + None, + Some(&backend), + &mut NoopHook, + ) + .unwrap(); let next_token = 3u32; let abs_position = prompt.len(); let h_trait = kv_decode_step_via_dispatch( &backend, - &weights, + larql_inference::WeightsView::dense(&weights), &ffn, &mut handles, next_token, @@ -122,17 +156,31 @@ fn multi_step_decode_via_dispatch_matches_legacy() { let ffn = WeightFfn { weights: &weights }; let prompt = vec![0u32, 1]; - let (_, mut handles) = - kv_prefill_via_dispatch(&backend, &weights, &ffn, &prompt, None, None).unwrap(); - let (_, mut cache) = - kv_prefill_run(&weights, &ffn, &prompt, None, Some(&backend), &mut NoopHook).unwrap(); + let (_, mut handles) = kv_prefill_via_dispatch( + &backend, + larql_inference::WeightsView::dense(&weights), + &ffn, + &prompt, + None, + None, + ) + .unwrap(); + let (_, mut cache) = kv_prefill_run( + larql_inference::WeightsView::dense(&weights), + &ffn, + &prompt, + None, + Some(&backend), + &mut NoopHook, + ) + .unwrap(); for step in 0..3 { let token = (2 + step) as u32; let abs_position = prompt.len() + step; let h_trait = kv_decode_step_via_dispatch( &backend, - &weights, + larql_inference::WeightsView::dense(&weights), &ffn, &mut handles, token, diff --git a/crates/larql-lql/src/executor/relation_resolver.rs b/crates/larql-lql/src/executor/relation_resolver.rs index 80ffc3fa1..dd56da5a2 100644 --- a/crates/larql-lql/src/executor/relation_resolver.rs +++ b/crates/larql-lql/src/executor/relation_resolver.rs @@ -117,10 +117,12 @@ impl RelationResolver { let layer = Self::probe_layer(weights.num_layers); // Only the layers up to the probe layer need to be f32 for the partial // forward — dequant just those (cheap). + let mut scratch = larql_inference::DequantScratch::new(); for l in 0..=layer { - larql_inference::vindex::insert_q4k_layer_tensors(&mut weights, &index, l) + larql_inference::vindex::insert_q4k_layer_tensors(&mut scratch, &weights, &index, l) .map_err(|e| LqlError::exec("relation resolver: dequant", e))?; } + weights.tensors.extend(scratch); // Training set: one residual per (relation, probe entity). let mut samples: Vec<(Vec, usize)> = diff --git a/crates/larql-models/src/architectures/bitnet.rs b/crates/larql-models/src/architectures/bitnet.rs new file mode 100644 index 000000000..1bfc31cdb --- /dev/null +++ b/crates/larql-models/src/architectures/bitnet.rs @@ -0,0 +1,56 @@ +//! BitNet b1.58 architecture (`general.architecture = "bitnet-b1.58"`, +//! HF `model_type = "bitnet"`). +//! +//! ## Why this is a thin, explicit entry rather than the generic fallback +//! +//! BitNet's *native-ternary* inference does NOT go through this trait. The +//! W1.58·A8 forward lives in `larql_inference::ternary` over the I2_S sidecar +//! written by `larql_vindex::extract::bitnet_writer`; that path is selected at +//! the convert boundary (`larql-vindex convert_cmd`, which reads the +//! `bitnet-b1.58.*` GGUF metadata directly), not by `detect_from_json`. +//! +//! This entry exists so that a BitNet config reaching the *generic* model +//! loader is **recognised explicitly** (`family() == "bitnet"`) instead of +//! silently collapsing to [`GenericArch`](super::generic::GenericArch) — the +//! latter masks the model behind a "generic" label and inherits Llama-style +//! defaults with no signal, the same silent-config class behind the earlier +//! forward-divergence fixes (`rms_norm_eps` / `rope_scaling`). +//! +//! BitNet's dense scaffold (token embedding, RMSNorm, RoPE, GQA with separate +//! Q/K/V, gated FFN naming) IS Llama-shaped, so the inherited defaults are +//! correct for the parts the generic loader actually materialises (embeddings, +//! norms, LM head — attention/FFN are skipped and served from the ternary +//! sidecar). `norm_eps()` reads `config.norm_eps` (parsed from +//! `rms_norm_eps`), so the epsilon is honoured, not hardcoded. +//! +//! The genuinely BitNet-specific divergences — the two sub-norms +//! (`attn_sub_norm`, `ffn_sub_norm`) and the squared-ReLU FFN over ternary +//! projections — are not expressible through this trait and are owned by the +//! `larql_inference::ternary` path. When BitNet graduates to a first-class +//! engine (see ROADMAP "BitNet b1.58 integration hardening"), its overrides +//! get a home here. + +use crate::config::{ModelArchitecture, ModelConfig}; + +/// BitNet b1.58 — thin, explicitly-named architecture entry. See the module +/// docs for why this mirrors [`GenericArch`](super::generic::GenericArch) +/// behaviour while reporting `family() == "bitnet"`. +pub struct BitnetArch { + config: ModelConfig, +} + +impl BitnetArch { + pub fn from_config(config: ModelConfig) -> Self { + Self { config } + } +} + +impl ModelArchitecture for BitnetArch { + fn family(&self) -> &str { + "bitnet" + } + + fn config(&self) -> &ModelConfig { + &self.config + } +} diff --git a/crates/larql-models/src/architectures/mod.rs b/crates/larql-models/src/architectures/mod.rs index 147aeb989..716d900a9 100644 --- a/crates/larql-models/src/architectures/mod.rs +++ b/crates/larql-models/src/architectures/mod.rs @@ -4,6 +4,7 @@ //! Every architecture implements [`ModelArchitecture`](crate::config::ModelArchitecture) //! and returns its own `model_type` from `family()`. +pub mod bitnet; pub mod deepseek; pub mod deepseek_v4; pub mod gemma2; diff --git a/crates/larql-models/src/detect/mod.rs b/crates/larql-models/src/detect/mod.rs index 6463e1baa..dae1bfffd 100644 --- a/crates/larql-models/src/detect/mod.rs +++ b/crates/larql-models/src/detect/mod.rs @@ -11,6 +11,7 @@ use std::path::Path; +use crate::architectures::bitnet::BitnetArch; use crate::architectures::deepseek::DeepSeekArch; use crate::architectures::deepseek_v4::DeepSeekV4Arch; use crate::architectures::gemma2::Gemma2Arch; @@ -134,6 +135,11 @@ pub fn detect_from_json(config: &serde_json::Value) -> Box Box::new(GraniteArch::from_config(model_config)), // TinyModel — research-scale decoder used for LARQL compile/walk work "tinymodel" => Box::new(TinyModelArch::from_config(model_config)), + // BitNet b1.58 (HF "bitnet", GGUF "bitnet-b1.58"). Recognised + // explicitly so a BitNet config can't silently collapse to the + // generic fallback; native-ternary inference is served by the + // larql-inference ternary path, not this trait. See BitnetArch docs. + t if t.starts_with("bitnet") => Box::new(BitnetArch::from_config(model_config)), // Unknown — generic fallback _ => Box::new(GenericArch::from_config(model_config)), } diff --git a/crates/larql-models/src/detect/tests.rs b/crates/larql-models/src/detect/tests.rs index 06f2e12fd..e3f0e73cb 100644 --- a/crates/larql-models/src/detect/tests.rs +++ b/crates/larql-models/src/detect/tests.rs @@ -364,6 +364,28 @@ fn test_detect_unknown_defaults_to_generic() { assert_eq!(arch.family(), "generic"); } +#[test] +fn test_detect_bitnet_is_explicit_not_generic() { + // BitNet must be recognised explicitly — never silently collapse to the + // generic fallback (which would mask wrong/default config behind a + // "generic" label). Both the HF and GGUF-derived model_type spellings. + for model_type in ["bitnet", "bitnet-b1.58", "bitnet_b1_58"] { + let config = serde_json::json!({ + "model_type": model_type, + "hidden_size": 2560, + "num_hidden_layers": 30, + "rms_norm_eps": 1e-5 + }); + let arch = detect_from_json(&config); + assert_eq!(arch.family(), "bitnet", "model_type={model_type}"); + // Epsilon is honoured from config, not hardcoded. + assert!( + (arch.norm_eps() - 1e-5).abs() < 1e-9, + "model_type={model_type}" + ); + } +} + #[test] fn test_tensor_keys() { let config = serde_json::json!({"model_type": "gemma3_text"}); diff --git a/crates/larql-models/src/lib.rs b/crates/larql-models/src/lib.rs index cb4b3992f..7d0d35ded 100644 --- a/crates/larql-models/src/lib.rs +++ b/crates/larql-models/src/lib.rs @@ -47,7 +47,7 @@ pub use vectors::{ COMPONENT_ATTN_QK, COMPONENT_EMBEDDINGS, COMPONENT_FFN_DOWN, COMPONENT_FFN_GATE, COMPONENT_FFN_UP, }; -pub use weights::{ModelWeights, WeightArray}; +pub use weights::{DequantScratch, ModelWeights, WeightArray, WeightsView}; pub use loading::{ is_ffn_tensor, load_gguf, load_gguf_validated, load_model_dir, load_model_dir_filtered, diff --git a/crates/larql-models/src/weights.rs b/crates/larql-models/src/weights.rs index 84b7a1728..81550c0af 100644 --- a/crates/larql-models/src/weights.rs +++ b/crates/larql-models/src/weights.rs @@ -9,6 +9,99 @@ use std::collections::{HashMap, HashSet}; /// Owned: from safetensors loading (heap). Shared: from mmap (zero-copy). pub type WeightArray = ArcArray2; +/// Engine-owned, per-forward dequantisation scratch: lazily-dequantised Q4K +/// layer tensors (attention / FFN → f32), keyed by the **same** tensor names +/// as [`ModelWeights::tensors`] (`arch.attn_q_key(layer)` etc.). +/// +/// Lives in the **engine** — per-forward derived state, the same category as +/// the KV cache — NOT in `ModelWeights`. This keeps `ModelWeights` truly +/// immutable so it can be shared as `Arc` across **concurrent** +/// generations (each engine its own scratch, no lock, no race). See ROADMAP +/// "P-B.1" for why the interior-mutable alternative was rejected on the +/// server-serialisation evidence. +/// +/// Eviction identity: the per-layer memory bound inserts and removes entries +/// under these same keys, and [`WeightsView::tensor`] resolves by them, so an +/// insert for layer L and its post-L evict cancel exactly. (Guarded by a +/// `resident_identity_tests` case that asserts L's entry is *gone* after L's +/// evict, not merely that the resolver still finds it.) +pub type DequantScratch = HashMap; + +/// Read-only resolver bundling canonical [`ModelWeights`] with an optional +/// engine-owned [`DequantScratch`]. The single read path for layer weight +/// tensors across the shared forward (dense and quant). +/// +/// Derefs to `ModelWeights`, so every non-tensor access (`view.hidden_size`, +/// `view.arch`, `view.embed`, …) works unchanged; only the layer-tensor reads +/// switch from `weights.tensors.get(key)` to [`tensor`](Self::tensor), which +/// resolves **scratch first, then canonical**. Returns a borrow (no clone, no +/// lock). +/// +/// `Copy` and two words wide (a ref + an `Option`), so threading it by +/// value through the forward is free. The dense / non-quant path constructs it +/// via [`dense`](Self::dense) (or `(&weights).into()`) — `scratch == None`, so +/// it's a transparent pass-through over canonical `tensors` and pays nothing. +/// The quant forward constructs it via [`with_scratch`](Self::with_scratch) +/// from the engine's per-forward `DequantScratch`. +#[derive(Clone, Copy)] +pub struct WeightsView<'a> { + weights: &'a ModelWeights, + scratch: Option<&'a DequantScratch>, +} + +impl<'a> WeightsView<'a> { + /// View with no dequant scratch — the dense f32 / non-quant path. Reads + /// resolve straight to canonical `tensors`. + pub fn dense(weights: &'a ModelWeights) -> Self { + Self { + weights, + scratch: None, + } + } + + /// View overlaying the engine's per-forward dequant scratch on canonical + /// weights. Reads resolve scratch-first-then-canonical. + pub fn with_scratch(weights: &'a ModelWeights, scratch: &'a DequantScratch) -> Self { + Self { + weights, + scratch: Some(scratch), + } + } + + /// Resolve a layer tensor by key: engine scratch first (when present), + /// then canonical `tensors`. The drop-in replacement for + /// `weights.tensors.get(key)` on the shared forward. + pub fn tensor(&self, key: &str) -> Option<&'a WeightArray> { + self.scratch + .and_then(|s| s.get(key)) + .or_else(|| self.weights.tensors.get(key)) + } + + /// Whether a layer tensor resolves via [`tensor`](Self::tensor). + pub fn has_tensor(&self, key: &str) -> bool { + self.scratch.is_some_and(|s| s.contains_key(key)) || self.weights.tensors.contains_key(key) + } + + /// The canonical (immutable) weights, for the rare caller that needs the + /// `&ModelWeights` directly rather than via `Deref`. + pub fn canonical(&self) -> &'a ModelWeights { + self.weights + } +} + +impl<'a> From<&'a ModelWeights> for WeightsView<'a> { + fn from(weights: &'a ModelWeights) -> Self { + Self::dense(weights) + } +} + +impl std::ops::Deref for WeightsView<'_> { + type Target = ModelWeights; + fn deref(&self) -> &ModelWeights { + self.weights + } +} + pub(crate) const PACKED_EXPERTS_GATE_UP_PROJ: &str = "experts.gate_up_proj"; pub(crate) const PACKED_EXPERTS_DOWN_PROJ: &str = "experts.down_proj"; @@ -281,3 +374,68 @@ pub fn per_layer_ffn_key(layer: usize, entry: usize, component: &str) -> String pub const PER_LAYER_FFN_GATE_UP: &str = "gate_up"; /// Component string for the down half of a per-layer FFN entry. pub const PER_LAYER_FFN_DOWN: &str = "down"; + +#[cfg(test)] +mod weights_view_tests { + use super::*; + use crate::test_fixtures::make_test_weights; + use ndarray::arr2; + + /// The resolver reads scratch-first-then-canonical, so an engine's + /// dequant scratch transparently shadows the canonical map for the same + /// key, and canonical-only keys still fall through. This is the contract + /// the 19 read sites migrate onto in P-B.1. + #[test] + fn weights_view_resolves_scratch_before_canonical() { + let mut weights = make_test_weights(); + let canonical = arr2(&[[1.0f32, 2.0]]).into_shared(); + let scratch_override = arr2(&[[9.0f32, 9.0]]).into_shared(); + let scratch_only = arr2(&[[3.0f32]]).into_shared(); + weights.tensors.insert("shared".into(), canonical.clone()); + weights + .tensors + .insert("canon_only".into(), canonical.clone()); + + let mut scratch = DequantScratch::new(); + scratch.insert("shared".into(), scratch_override.clone()); // shadows canonical + scratch.insert("scratch_only".into(), scratch_only.clone()); + + let view = WeightsView::with_scratch(&weights, &scratch); + + // scratch shadows canonical for the same key + assert_eq!(view.tensor("shared").unwrap(), &scratch_override); + // scratch-only resolves + assert_eq!(view.tensor("scratch_only").unwrap(), &scratch_only); + // canonical-only falls through + assert_eq!(view.tensor("canon_only").unwrap(), &canonical); + // absent everywhere + assert!(view.tensor("nope").is_none()); + + assert!(view.has_tensor("shared")); + assert!(view.has_tensor("scratch_only")); + assert!(view.has_tensor("canon_only")); + assert!(!view.has_tensor("nope")); + + // Deref: non-tensor fields reachable through the view unchanged. + assert_eq!(view.hidden_size, weights.hidden_size); + assert_eq!(view.canonical().num_layers, weights.num_layers); + } + + /// The dense / non-quant view (no scratch) is a transparent pass-through + /// over canonical `tensors` — the "dense pays nothing" property. `From` + /// builds the same thing. + #[test] + fn weights_view_dense_is_passthrough() { + let mut weights = make_test_weights(); + let w = arr2(&[[1.0f32]]).into_shared(); + weights.tensors.insert("k".into(), w.clone()); + + let view = WeightsView::dense(&weights); + assert_eq!(view.tensor("k").unwrap(), &w); + assert!(view.tensor("absent").is_none()); + assert!(view.has_tensor("k") && !view.has_tensor("absent")); + + let via_from: WeightsView = (&weights).into(); + assert_eq!(via_from.tensor("k").unwrap(), &w); + } +} diff --git a/crates/larql-server/tests/test_synthetic_q4k_smoke.rs b/crates/larql-server/tests/test_synthetic_q4k_smoke.rs index 8a987fcec..8cc473620 100644 --- a/crates/larql-server/tests/test_synthetic_q4k_smoke.rs +++ b/crates/larql-server/tests/test_synthetic_q4k_smoke.rs @@ -47,7 +47,9 @@ fn q4k_fixture_unblocks_insert_q4k_layer_tensors() { let patched = model.patched.blocking_read(); let index = patched.base(); - let inserted = larql_inference::vindex::insert_q4k_layer_tensors(weights, index, 0); + let mut scratch = larql_inference::DequantScratch::new(); + let inserted = + larql_inference::vindex::insert_q4k_layer_tensors(&mut scratch, weights, index, 0); assert!( inserted.is_ok(), "insert_q4k_layer_tensors must succeed on the Q4K fixture; got {inserted:?}"