Skip to content

Native backend v3: 20x+ faster proving, verification, and inference paths#60

Merged
mdvillagra merged 8 commits into
mainfrom
claude/loop-command-xcacg0
Jun 10, 2026
Merged

Native backend v3: 20x+ faster proving, verification, and inference paths#60
mdvillagra merged 8 commits into
mainfrom
claude/loop-command-xcacg0

Conversation

@BidhanRoy

@BidhanRoy BidhanRoy commented Jun 10, 2026

Copy link
Copy Markdown
Contributor

Summary

Makes zkLoRA dramatically faster end-to-end while preserving (and in places strengthening) the proof contract's correctness and security guarantees. The statement format and adapter-commitment scheme are unchanged, so pinned adapter manifests remain valid; BACKEND_ID is bumped to zklora-halo2-v3 so stale v2 proofs/keys fail closed.

All numbers below are measured on a 4-core / 15 GB host and reproducible via cargo run --release --example bench_prove and python benchmarks/run_benchmarks.py (see benchmarks/native_v3_results.md).

Speedups

Path Before After Speedup
Proving (2×1×2) 22.9 s 0.73 s warm 31×
Proving (8×2×8) >290 s 1.4 s warm >200×
Verification 18.4 s 0.02–1.3 s 14×–900×
Server per-invocation compute (16 rows @ 768×4×2304) ~233 ms 13.5 ms ~17×
Exact delta (16 rows @ 768×4×2304) 171 ms 7.4 ms 23×
Hiding Merkle root (200k leaves) 0.353 s 0.043 s 8.3×

Real LoRA shapes (768-dim c_attn) went from infeasible (k≈26, est. >100 GB of params/keys) to runnable: 768×2×256 proves at k=19 with a 12.1 GB peak; 768×4×2304 fits at k=21.

What changed

Circuit (Rust):

  • Replaced per-bit boolean range decomposition (~8 rows/bit) with lookup-backed running-sum limb decomposition (~1 row per 16-bit limb), shrinking circuits ~32× in rows. The signed-interval semantics are exactly preserved — the soundness argument (including why the paired shifted + diff = max construction is exact with no field wraparound) is documented in-code.
  • Dropped only provably redundant per-product range checks whose bounds already follow from the individually range-checked operands; the same bounds are still enforced at every division witness.
  • Deterministic, shape-only window/k selection so prover and verifier always derive identical circuits.
  • Shape-keyed caches for SRS params, proving keys, and verifying keys; verification no longer pays prover-grade keygen per proof.
  • Released the GIL in the PyO3 bindings; added exact-integer (i128) delta and parallel (rayon) batched delta / Merkle helpers.

Python:

  • Parallel proof generation and batch verification thread pools.
  • Content-keyed adapter-commitment memoisation.
  • LoRAServer quantizes adapter matrices once per module instead of per invocation; the socket server handles connections in threads with a lock.
  • Native exact quantizer anchored on Python's str(float) (Rust's own float formatter can pick a different shortest-repr digit string — regression covered by tests). All native fast paths fall back to the exact Python reference.

Correctness & security

  • Statement format and Poseidon adapter-commitment scheme unchanged; pinned manifests stay valid.
  • Every native fast path is value-identical to its Python reference, verified by parity tests including an 80k-value randomized quantization fuzz across four fixed-point configs.
  • A faster single-pass native JSON-parse path for activation commitments was investigated and rejected: serde_json parses ~7% of full-precision float literals to a different nearest-f64 than Python's float(), which would silently change commitment values. The reference float semantics (Python json/float) are kept; only the hashing is native.
  • New tests: lookup multi-limb range checks, exact signed-bound boundary values, native-delta/Merkle parity, and the previously perma-ignored real-proof test now runs in ~1.5 s.
  • Review follow-up: the LoRAServer per-module quantization cache is keyed on the current weight bytes (so swapped/updated adapters are never served stale), the native delta rounding step uses fully checked i128 arithmetic with an overflow regression test, and the end-to-end real-proof test is no longer #[ignore]d.

Tests

51 tests green (40 Python + 11 Rust), lint and format clean, cargo check --features extension-module passes.

Bidhan Roy added 8 commits June 10, 2026 05:00
…oof contract

Rust circuit (backend v3, same statement format and adapter commitment scheme):
- Replace per-bit boolean range decomposition (~8 rows/bit) with
  lookup-backed running-sum limb decomposition; exact interval semantics
  including residual-width top limbs, no padding slack.
- Drop provably redundant per-product/per-accumulator range checks whose
  bounds follow from the individually range-checked operands; the same
  bounds are still enforced at each division witness.
- Deterministic, shape-only window/k selection so prover and verifier
  always derive identical circuits.
- Cache SRS params and proving/verifying keys keyed by circuit shape;
  verification no longer pays prover-grade keygen per proof.
- Release the GIL in PyO3 bindings; add native exact-integer
  compute_delta_quantized (checked i128) and hiding Merkle root helpers.

Python layer:
- Parallel proof generation and batch verification thread pools.
- Content-keyed adapter commitment memoisation.
- LoRAServer quantizes adapter matrices once per module instead of per
  invocation; socket server handles connections in threads with a lock.
- Native fast paths fall back to the exact Python implementations.

Measured (4 cores, warm caches): 2x1x2 prove 22.9s->0.73s, verify
18.4s->0.02s; 8x2x8 prove >290s->1.4s; real 768x4x2304 shape drops from
k~26 (infeasible) to k=21. New parity and mock-prover tests cover the
fast paths and multi-limb range checks.
- Add batched native compute_delta_rows: validates the adapter once and
  computes row deltas in parallel with rayon; per-row semantics identical
  to the single-row exact path.
- Add native exact quantizer: Python supplies str(float) decimal strings
  (the semantic anchor of quantize_scalar) and Rust performs the
  round-half-up scaling in exact i128 arithmetic. Rust's own float
  formatter is deliberately not used: it can emit a different shortest
  round-trip digit string than CPython's repr (regression value covered
  in tests), which would silently change quantized values.
- Parallelize Merkle leaf/level hashing with rayon.
- apply_lora and the transcript recorder use the batched paths.
- 80k-value randomized parity fuzz across four fixed-point configs
  validates the quantizer against the Decimal reference; batched delta
  and Merkle parity tests added.

Server per-invocation compute (16 rows, 768x4x2304): ~233ms -> 13.5ms.
Document the precise soundness argument for the lookup running-sum
range check: the unsigned chain alone admits a factor-2 slack (a top
limb need not be a multiple of the scale divisor), and exactness of the
signed interval comes from the paired shifted/diff decompositions plus
the shifted + diff = max equation, which cannot wrap for bits <= 250.
Add a mock-prover test at exact signed-bound values and a README
performance summary.
Records measured prove/verify times for the v2 baseline vs the v3
backend across shapes, end-to-end pipeline timings, supporting fast-path
speedups, and the halo2 extended-domain memory profile (peak ~7-8 GB at
k=20, >15 GB at k=21). The 768-dim rows will be finalised when the
in-flight k=20 measurement completes.
k=20 and k=21 shapes exceed a 15 GB host (OOM-killed) because halo2 0.3
keeps all extended-domain advice cosets resident during create_proof;
recommend 32-64 GB hosts for wide-output 768-dim modules. A k=19
768x2x256 measurement is in flight to anchor the feasible-on-15GB row.
768x2x256 (k=19): cold prove 429s, warm prove 81s, warm verify 1.3s,
12.1 GB peak RSS on a 15 GB host. The v2 circuit at this shape would
have needed k~23 (~16x rows, est. ~190 GB) and was infeasible.
Investigated moving JSON parse+flatten+hash into a single native pass to
push the commitment path past 20x. Rejected it: serde_json parses ~7% of
full-precision float literals to a different nearest-f64 than Python's
float(), which would silently change commitment values for existing data
(verified by per-element parity testing). The reference float semantics
are defined by Python's json/float, so parsing stays in Python.

Kept a small dedup: the numpy/recursive flatten logic that was copied in
commit_activations and verify_commitment is now _flatten_input_data, and
both route through _merkle_root, which already uses the native
BLAKE3 fast path (8.3x on hashing, byte-identical to the reference).
Parity re-verified across rectangular, ragged, and empty inputs.
…al-proof test

Three fixes from a correctness/security review of the native backend v3:

- LoRAServer._module_artifacts cached quantized adapter matrices keyed only
  by submodule name for the server's lifetime, so adapter weights swapped or
  updated in place would silently keep proving against the first weights
  seen. The cache is now keyed on a digest of the current weight bytes and
  scaling, re-quantizing whenever the adapter actually changes.

- delta_for_row's rounding step used an unchecked i128 add (numerator +
  denominator / 2); adversarial near-i128::MAX scaled values that pass the
  intermediate-bound check could panic in debug builds, surfacing as a
  PanicException that escapes the ValueError-based fallback to the exact
  Python path. The add is now checked and reported as a NativeError, with a
  regression test pinning the exact overflow boundary.

- real_proof_verifies_for_tiny_relation was still #[ignore]d even though the
  v3 speedups make it run in ~1.5 s release / ~14 s debug; it is the only
  test exercising real keygen, the shape-keyed key caches, and
  prove/verify round-tripping, so it now runs by default.
@BidhanRoy
BidhanRoy force-pushed the claude/loop-command-xcacg0 branch from 77dc094 to 4a36447 Compare June 10, 2026 14:33
@mdvillagra
mdvillagra merged commit 8b53724 into main Jun 10, 2026
6 checks passed
@mdvillagra
mdvillagra deleted the claude/loop-command-xcacg0 branch June 10, 2026 20:23
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants