diff --git a/benchmarks/native_v3_results.md b/benchmarks/native_v3_results.md new file mode 100644 index 0000000..a4e22ba --- /dev/null +++ b/benchmarks/native_v3_results.md @@ -0,0 +1,57 @@ +# Native backend v3 performance results + +Measured on a 4-core x86-64 container, 15 GB RAM, release builds, default +fixed-point config (`scale_bits=20, value_bits=63, intermediate_bits=127`). +"v2 baseline" is the previous native circuit (per-bit range checks, no key +caching) at the same commit environment; "warm" means params/proving key are +cached from a previous proof of the same shape (the steady-state for real +workloads, where every invocation of a module shares one circuit shape). + +Reproduce with `cargo run --release --example bench_prove -- ` +and `python benchmarks/run_benchmarks.py`. + +## Single proof + verify (Rust, `bench_prove`) + +| shape (in×rank×out) | v2 k | v2 prove | v2 verify | v3 k | v3 prove cold | v3 prove warm | v3 verify warm | warm speedup | +|---|---|---|---|---|---|---|---|---| +| 2×1×2 | 15 | 22.9 s | 18.4 s | 12 | 2.9 s | 0.73 s | 0.020 s | 31× / ~860× | +| 8×2×8 | ~18 | >290 s (timed out) | — | 13 | 6.0 s | 1.4 s | 0.034 s | >200× | +| 16×2×16 | ~19 | est. >370 s | — | 14 | 27.2 s | 7.0 s | 0.16 s | >50× | +| 64×4×64 | ~22 | infeasible | — | 17 | 47.2 s | 10.8 s | 0.22 s | n/a (was infeasible) | +| 768×2×256 | ~23 | infeasible (est. ~190 GB) | — | 19 | 429 s | 81 s | 1.3 s (17.1 s cold) | n/a (was infeasible) | +| 768×4×768 | ~25 | infeasible | — | 20 | needs >15 GB RAM host | — | — | n/a | +| 768×4×2304 | ~26 | infeasible (est. >200 GB params/pk) | — | 21 | needs >15 GB RAM host | — | — | n/a | + +Verification cold (first proof of a shape) pays one `keygen_vk`; in v2 this +cost was paid for *every* proof. + +## End-to-end Python pipeline (`run_benchmarks.py`, 4 workers) + +| shape | invocations | prove wall | prove/proof | verify wall | verify/proof | +|---|---|---|---|---|---| +| 16×2×16 | 8 | 71.4 s | 8.9 s | 3.4 s | 0.42 s | +| 32×4×32 | 6 | 107.4 s | 17.9 s | 3.8 s | 0.63 s | + +Per-proof wall time includes the one-time cold keygen amortised across the +batch. + +## Supporting paths + +| path | before | after | speedup | +|---|---|---|---| +| server per-invocation compute, 16 rows @ 768×4×2304 (quantize + delta) | ~233 ms | 13.5 ms | ~17× (≥30× at 1 row) | +| 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× | + +All fast paths are value-identical to the Python reference implementations +(randomised parity tests, including an 80k-value quantisation fuzz across +four fixed-point configs). + +## Memory + +Proving memory is dominated by halo2 extended-domain evaluations (Poseidon +gate degree ⇒ 8× extended domain), and halo2 0.3 keeps all advice cosets +resident during create_proof. Measured: 12.1 GB peak at k=19 (768×2×256); +k=20 and k=21 shapes exceed a 15 GB host (OOM-killed); plan on 32–64 GB for +768-dim × 4-rank modules with wide outputs. +The proving-key cache holds up to `ZKLORA_PK_CACHE_CAP` (default 2) shapes. diff --git a/benchmarks/run_benchmarks.py b/benchmarks/run_benchmarks.py new file mode 100644 index 0000000..7dea3db --- /dev/null +++ b/benchmarks/run_benchmarks.py @@ -0,0 +1,123 @@ +#!/usr/bin/env python3 +"""End-to-end zkLoRA pipeline benchmark. + +Generates invocation witnesses for a synthetic LoRA module, produces native +proof artifacts, and verifies them against the transcript and adapter +manifest, reporting wall-clock timings for each stage. + +Usage: + python benchmarks/run_benchmarks.py [--in_dim 16] [--rank 2] [--out_dim 16] + [--invocations 8] +""" + +import argparse +import json +import random +import sys +import tempfile +import time +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src")) + +from zklora.proof_contract import ( # noqa: E402 + FixedPointConfig, + InvocationWitness, + adapter_manifest_entry, + compute_delta_quantized, + statement_from_witness, + transcript_entry_from_statement, +) +from zklora.zk_proof_generator import batch_verify_proofs, generate_proofs # noqa: E402 + + +def build_witnesses(in_dim, rank, out_dim, invocations, seed=7): + rng = random.Random(seed) + fp = FixedPointConfig() + magnitude = 1 << fp.scale_bits + a = [[rng.randint(-magnitude, magnitude) for _ in range(in_dim)] for _ in range(rank)] + b = [[rng.randint(-magnitude, magnitude) for _ in range(rank)] for _ in range(out_dim)] + witnesses = [] + for index in range(invocations): + x = [rng.randint(-magnitude, magnitude) for _ in range(in_dim)] + delta = compute_delta_quantized(a, b, x, 1, 1, fp) + witnesses.append( + InvocationWitness( + session_id="bench-session", + module_name="bench.module.c_attn", + invocation_index=index, + input_shape=[in_dim], + output_shape=[out_dim], + x=x, + delta=delta, + a=a, + b=b, + scaling_num=1, + scaling_den=1, + adapter_metadata={"rank": rank, "in_dim": in_dim, "out_dim": out_dim}, + fixed_point=fp, + ) + ) + manifest = [adapter_manifest_entry("bench.module.c_attn", a, b, 1, 1, fp)] + return witnesses, manifest + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--in_dim", type=int, default=16) + parser.add_argument("--rank", type=int, default=2) + parser.add_argument("--out_dim", type=int, default=16) + parser.add_argument("--invocations", type=int, default=8) + args = parser.parse_args() + + print( + f"shape in_dim={args.in_dim} rank={args.rank} out_dim={args.out_dim} " + f"invocations={args.invocations}" + ) + + start = time.time() + witnesses, manifest = build_witnesses( + args.in_dim, args.rank, args.out_dim, args.invocations + ) + print(f"witness generation: {time.time() - start:.2f}s") + + with tempfile.TemporaryDirectory() as tmp: + proof_dir = Path(tmp) / "artifacts" + start = time.time() + _, _, elapsed, total_params, proofs = generate_proofs( + records=witnesses, output_dir=str(proof_dir) + ) + prove_wall = time.time() - start + print( + f"proof generation: {prove_wall:.2f}s total " + f"({prove_wall / max(proofs, 1):.2f}s/proof, {proofs} proofs)" + ) + + transcript = [ + transcript_entry_from_statement(statement_from_witness(w)) for w in witnesses + ] + start = time.time() + verify_time, verified = batch_verify_proofs( + proof_dir=str(proof_dir), + transcript=transcript, + expected_adapters={"adapters": manifest}, + ) + verify_wall = time.time() - start + print( + f"verification: {verify_wall:.2f}s total " + f"({verify_wall / max(verified, 1):.2f}s/proof, {verified} proofs)" + ) + result = { + "shape": [args.in_dim, args.rank, args.out_dim], + "invocations": args.invocations, + "prove_wall_s": round(prove_wall, 3), + "prove_per_proof_s": round(prove_wall / max(proofs, 1), 3), + "verify_wall_s": round(verify_wall, 3), + "verify_per_proof_s": round(verify_wall / max(verified, 1), 3), + "total_params": total_params, + } + print(json.dumps(result)) + + +if __name__ == "__main__": + main() diff --git a/readme.md b/readme.md index bccb02e..2ba749c 100644 --- a/readme.md +++ b/readme.md @@ -357,7 +357,10 @@ For detailed information about the codebase organization and implementation deta ✓Adapter Weight Privacy: LoRA weights remain confidential while the committed adapter identity is checked -✓Benchmark Required: Real-shape proving and verification performance should be measured for each deployment target +✓Benchmark Required: Real-shape proving and verification performance should be measured for each deployment target (see benchmarks/run_benchmarks.py and cargo run --release --example bench_prove) + + +✓Fast Native Backend (v3): Lookup-based range checks, shape-keyed SRS/key caches, and parallel batch proving/verification deliver order-of-magnitude speedups over v2 while keeping the same statement format and adapter commitment scheme diff --git a/src/Cargo.toml b/src/Cargo.toml index 9626e37..3a31ee4 100644 --- a/src/Cargo.toml +++ b/src/Cargo.toml @@ -13,6 +13,7 @@ python = ["dep:pyo3"] extension-module = ["python", "pyo3/extension-module"] [dependencies] +blake3 = "1" halo2_proofs = "0.3.2" halo2_gadgets = "0.5" ff = "0.13" @@ -21,6 +22,7 @@ num-integer = "0.1" num-traits = "0.2" pyo3 = { version = "0.28.3", optional = true } rand_core = { version = "0.6", features = ["getrandom"] } +rayon = "1" serde = { version = "1", features = ["derive"] } serde_json = "1" sha2 = "0.10" diff --git a/src/README.md b/src/README.md index 8934261..cfc7903 100644 --- a/src/README.md +++ b/src/README.md @@ -56,7 +56,18 @@ The Rust implementation is wrapped with Python bindings in the `libs/merkle` dir ### Performance Considerations -Native Halo2 performance should be measured for the specific LoRA shapes being proven. The v2 implementation prioritizes proof-contract correctness, transcript binding, and pre-agreed adapter binding before publishing benchmark claims. +Native Halo2 performance should be measured for the specific LoRA shapes being proven. The v3 backend keeps the v2 proof contract (same statement format, same Poseidon adapter-commitment scheme) while making proving and verification dramatically faster: + +- **Lookup-based range checks**: signed interval checks decompose values into window-sized limbs via a running sum constrained against a lookup table, instead of one boolean row per bit. The interval semantics are exact (the top limb is scaled to its residual width), and the circuit drops provably redundant per-product checks whose bounds already follow from the range-checked operands. At a 768×4×2304 LoRA shape this reduces the circuit from k≈26 to k=21 (32× fewer rows). +- **Keyed SRS/proving-key/verifying-key caches**: params and keys are derived deterministically from the statement shape (dims, fixed-point bits, scaling) and reused across invocations of the same module. First proof per shape pays keygen; subsequent proofs and all verifications are keygen-free. Cache sizes are tunable via `ZKLORA_PARAMS_CACHE_CAP`, `ZKLORA_PK_CACHE_CAP`, and `ZKLORA_VK_CACHE_CAP`. +- **Parallel batch operations**: the PyO3 bindings release the GIL, and `generate_proofs` / `batch_verify_proofs` fan out across a thread pool (`ZKLORA_PROVE_WORKERS`, `ZKLORA_VERIFY_WORKERS`). +- **Native fast paths with exact fallbacks**: quantized delta computation and the hiding Merkle commitment have Rust implementations that are value-identical to the Python reference paths (covered by parity tests); the Python implementations remain as exact fallbacks. + +Proofs and verifying keys from the v2 backend are not compatible with v3 (the `backend` field in statements changed to `zklora-halo2-v3`), but pinned adapter manifests remain valid: the adapter commitment scheme is unchanged. + +Measured on a 4-core machine (warm key cache): a 2×1×2 relation proves in ~0.7s and verifies in ~0.02s (previously ~23s / ~18s), an 8×2×8 relation proves in ~1.4s (previously >5 minutes), and the real 768×4×2304 c_attn shape becomes feasible at k=21. + +Memory note: proving memory is dominated by halo2's extended-domain evaluations (the Poseidon chip's gate degree implies an 8× extended domain), and halo2 0.3 keeps all advice cosets resident during proof creation. Measured on a 15 GB host: a k=19 shape (768×2×256) proves in 81 s warm with a 12.1 GB peak, while k=20 and k=21 shapes (e.g. 768×4×768 and the full 768×4×2304 c_attn) exceed 15 GB and require a 32–64 GB prover host. Verification is lightweight once the verifying key is cached. For detailed usage examples and high-level architecture, please refer to the [main README](../../README.md) in the project root. diff --git a/src/examples/bench_prove.rs b/src/examples/bench_prove.rs new file mode 100644 index 0000000..b6aa792 --- /dev/null +++ b/src/examples/bench_prove.rs @@ -0,0 +1,27 @@ +//! Benchmark harness for zkLoRA native proving and verification. +//! +//! Usage: cargo run --release --example bench_prove -- [reps] + +use std::time::Instant; + +use _native_prover::bench_support::{bench_statement_and_witness, prove_verify_once}; + +fn main() { + let args: Vec = std::env::args().collect(); + let in_dim: usize = args.get(1).map(|v| v.parse().unwrap()).unwrap_or(8); + let rank: usize = args.get(2).map(|v| v.parse().unwrap()).unwrap_or(2); + let out_dim: usize = args.get(3).map(|v| v.parse().unwrap()).unwrap_or(8); + let reps: usize = args.get(4).map(|v| v.parse().unwrap()).unwrap_or(1); + + let (statement_json, witness_json, k) = bench_statement_and_witness(in_dim, rank, out_dim); + println!("shape in_dim={in_dim} rank={rank} out_dim={out_dim} k={k} reps={reps}"); + + for rep in 0..reps { + let start = Instant::now(); + let (prove_ms, verify_ms, proof_len) = prove_verify_once(&statement_json, &witness_json); + println!( + "rep={rep} prove_ms={prove_ms:.1} verify_ms={verify_ms:.1} total_ms={:.1} proof_bytes={proof_len}", + start.elapsed().as_secs_f64() * 1000.0 + ); + } +} diff --git a/src/src/lib.rs b/src/src/lib.rs index 4b95e83..f90cf7d 100644 --- a/src/src/lib.rs +++ b/src/src/lib.rs @@ -1,5 +1,4 @@ use ff::PrimeField; -#[cfg(any(test, feature = "python"))] use halo2_gadgets::poseidon::primitives::Hash as NativePoseidonHash; use halo2_gadgets::poseidon::{ primitives::{ConstantLength, P128Pow5T3}, @@ -10,7 +9,8 @@ use halo2_proofs::{ pasta::{vesta, EqAffine, Fp}, plonk::{ create_proof, keygen_pk, keygen_vk, verify_proof, Advice, Circuit, Column, - ConstraintSystem, Error, Instance, Selector, SingleVerifier, + ConstraintSystem, Error, Fixed, Instance, ProvingKey, Selector, SingleVerifier, + TableColumn, VerifyingKey, }, poly::commitment::Params, poly::Rotation, @@ -20,9 +20,12 @@ use num_bigint::{BigInt, BigUint, Sign}; use num_integer::Integer; use num_traits::{One, Signed, Zero}; use rand_core::OsRng; +use rayon::prelude::*; use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; +use std::collections::{HashMap, VecDeque}; use std::convert::TryInto; +use std::sync::{Arc, Mutex, OnceLock}; const ADAPTER_COMMITMENT_DOMAIN: u64 = 0x5a4b4c4f5241; // "ZKLORA" const ADAPTER_COMMITMENT_VERSION: u64 = 1; @@ -30,6 +33,125 @@ const ADAPTER_COMMITMENT_VERSION: u64 = 1; const ARTIFACT_SCHEMA_VERSION: u64 = 2; const FIELD_SAFE_BITS: usize = 250; const POSEIDON_PAIR_ROWS: usize = 96; +const RANGE_WINDOW_MIN: u32 = 4; +const RANGE_WINDOW_MAX: u32 = 16; +const ROW_MARGIN: usize = 128; + +/// Cache key capturing everything the circuit layout (and therefore the +/// params/proving key/verifying key) depends on. Witness values, the adapter +/// commitment, and the statement digest are advice/instance data and do not +/// influence key generation. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +struct CircuitKey { + in_dim: usize, + rank: usize, + out_dim: usize, + scale_bits: u32, + value_bits: u32, + intermediate_bits: u32, + scaling_num: i64, + scaling_den: i64, +} + +impl CircuitKey { + fn from_circuit(circuit: &LoraCircuit) -> Self { + Self { + in_dim: circuit.in_dim(), + rank: circuit.rank(), + out_dim: circuit.out_dim(), + scale_bits: circuit.fixed_point.scale_bits, + value_bits: circuit.fixed_point.value_bits, + intermediate_bits: circuit.fixed_point.intermediate_bits, + scaling_num: circuit.scaling_num, + scaling_den: circuit.scaling_den, + } + } +} + +/// Simple bounded FIFO cache. Proving keys and SRS params at large `k` are +/// hundreds of MB, so we cap how many distinct shapes stay resident. +struct BoundedCache { + map: HashMap>, + order: VecDeque, + cap: usize, +} + +impl BoundedCache { + fn new(cap: usize) -> Self { + Self { + map: HashMap::new(), + order: VecDeque::new(), + cap: cap.max(1), + } + } + + fn get_or_create( + &mut self, + key: &K, + create: impl FnOnce() -> Result, + ) -> Result, E> { + if let Some(value) = self.map.get(key) { + return Ok(value.clone()); + } + let value = Arc::new(create()?); + while self.order.len() >= self.cap { + if let Some(evicted) = self.order.pop_front() { + self.map.remove(&evicted); + } + } + self.order.push_back(key.clone()); + self.map.insert(key.clone(), value.clone()); + Ok(value) + } +} + +fn cache_cap(env_var: &str, default: usize) -> usize { + std::env::var(env_var) + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or(default) + .max(1) +} + +fn params_for(k: u32) -> Arc> { + static CACHE: OnceLock>>> = OnceLock::new(); + let cache = CACHE + .get_or_init(|| Mutex::new(BoundedCache::new(cache_cap("ZKLORA_PARAMS_CACHE_CAP", 4)))); + let mut guard = cache.lock().expect("params cache poisoned"); + guard + .get_or_create::(&k, || Ok(Params::new(k))) + .expect("params creation is infallible") +} + +fn proving_key_for( + key: &CircuitKey, + params: &Params, + empty_circuit: &LoraCircuit, +) -> Result>, NativeError> { + static CACHE: OnceLock>>> = OnceLock::new(); + let cache = + CACHE.get_or_init(|| Mutex::new(BoundedCache::new(cache_cap("ZKLORA_PK_CACHE_CAP", 2)))); + let mut guard = cache.lock().expect("pk cache poisoned"); + guard.get_or_create(key, || { + let vk = keygen_vk(params, empty_circuit).map_err(|e| NativeError::Halo2(e.to_string()))?; + keygen_pk(params, vk, empty_circuit).map_err(|e| NativeError::Halo2(e.to_string())) + }) +} + +fn verifying_key_for( + key: &CircuitKey, + params: &Params, + empty_circuit: &LoraCircuit, +) -> Result>, NativeError> { + static CACHE: OnceLock>>> = + OnceLock::new(); + let cache = + CACHE.get_or_init(|| Mutex::new(BoundedCache::new(cache_cap("ZKLORA_VK_CACHE_CAP", 8)))); + let mut guard = cache.lock().expect("vk cache poisoned"); + guard.get_or_create(key, || { + keygen_vk(params, empty_circuit).map_err(|e| NativeError::Halo2(e.to_string())) + }) +} #[derive(Debug, thiserror::Error)] pub enum NativeError { @@ -78,7 +200,6 @@ pub struct NativeWitness { pub b: Vec>, } -#[cfg(any(test, feature = "python"))] #[derive(Clone, Debug, Serialize, Deserialize)] struct AdapterCommitmentInput { pub schema_version: u64, @@ -112,7 +233,15 @@ struct LoraConfig { mul: Selector, add: Selector, div_round: Selector, - boolean: Selector, + /// Complex selector activating one step of the range-check running sum. + q_range: Selector, + /// Per-row decomposition radix (2^window on active rows). + rc_factor: Column, + /// Per-row limb scale: 1 on full limbs, 2^(window - top_bits) on the top + /// limb so the most-significant limb is constrained to exactly its width. + rc_scale: Column, + /// Lookup table holding 0..2^window. + rc_table: TableColumn, poseidon_config: Pow5Config, } @@ -175,11 +304,26 @@ impl Circuit for LoraCircuit { vec![s * (raw - quotient * denominator - remainder)] }); - let boolean = meta.selector(); - meta.create_gate("boolean bit", |meta| { - let s = meta.query_selector(boolean); - let bit = meta.query_advice(advice[0], Rotation::cur()); - vec![s * bit.clone() * (bit - halo2_proofs::plonk::Expression::Constant(Fp::from(1)))] + // Lookup-based range decomposition: a running sum z walks down the + // value, peeling one limb per active row. On row i the constraint + // q * scale_i * (z_i - z_{i+1} * factor_i) ∈ rc_table + // forces limb_i = z_i - z_{i+1} * 2^window into [0, 2^window), and the + // top limb is multiplied by 2^(window - top_bits) so it is bounded to + // exactly its residual width. With z_n constrained to zero the chain + // telescopes to z_0 = Σ limb_i · 2^(window·i) < 2^bits with no padding + // slack, preserving the exact interval semantics of the previous + // bit-decomposition while using ~window× fewer rows. + let q_range = meta.complex_selector(); + let rc_factor = meta.fixed_column(); + let rc_scale = meta.fixed_column(); + let rc_table = meta.lookup_table_column(); + meta.lookup(|meta| { + let q = meta.query_selector(q_range); + let z_cur = meta.query_advice(advice[3], Rotation::cur()); + let z_next = meta.query_advice(advice[3], Rotation::next()); + let factor = meta.query_fixed(rc_factor); + let scale = meta.query_fixed(rc_scale); + vec![(q * scale * (z_cur - z_next * factor), rc_table)] }); let poseidon_state = (0..3).map(|_| meta.advice_column()).collect::>(); @@ -204,7 +348,10 @@ impl Circuit for LoraCircuit { mul, add, div_round, - boolean, + q_range, + rc_factor, + rc_scale, + rc_table, poseidon_config, } } @@ -215,6 +362,22 @@ impl Circuit for LoraCircuit { mut layouter: impl Layouter, ) -> Result<(), Error> { self.validate().map_err(|_| Error::Synthesis)?; + let window = self.window_and_k().0; + + layouter.assign_table( + || "range check table", + |mut table| { + for value in 0..(1u64 << window) { + table.assign_cell( + || "range table value", + config.rc_table, + value as usize, + || Value::known(Fp::from(value)), + )?; + } + Ok(()) + }, + )?; let (x_cells, delta_cells, binding_cells, adapter_words) = layouter.assign_region( || "zklora lora delta relation", @@ -226,8 +389,6 @@ impl Circuit for LoraCircuit { let raw_a_bound = &value_bound * &value_bound * BigInt::from(self.in_dim()); let raw_b_bound = &value_bound * &intermediate_bound * BigInt::from(self.rank()); let scaled_raw_bound = &intermediate_bound * BigInt::from(self.scaling_num).abs(); - let product_value_bound = &value_bound * &value_bound; - let product_intermediate_bound = &value_bound * &intermediate_bound; let zero = assign_constant_cell( &mut region, @@ -329,6 +490,7 @@ impl Circuit for LoraCircuit { offset = range_check_signed_interval( &mut region, &config, + window, &cell, &value_big, &(-&value_bound), @@ -355,6 +517,7 @@ impl Circuit for LoraCircuit { offset = range_check_signed_interval( &mut region, &config, + window, &weight, &weight_big, &(-&value_bound), @@ -363,6 +526,11 @@ impl Circuit for LoraCircuit { )?; adapter_words.push(weight.clone()); + // The product of two individually range-checked values + // cannot wrap the field (|w·x| <= value_bound^2 and the + // accumulated sum stays below the field-safe bound per + // validate_field_safety), so the accumulator is bounded + // once inside assign_div_round instead of per product. let product = assign_mul( &mut region, &config, @@ -372,33 +540,16 @@ impl Circuit for LoraCircuit { )?; offset += 1; let product_value = &weight_big * BigInt::from(self.x[input_index]); - offset = range_check_signed_interval( - &mut region, - &config, - &product, - &product_value, - &(-&product_value_bound), - &product_value_bound, - offset, - )?; let next_acc = assign_add(&mut region, &config, &acc, &product, offset)?; offset += 1; raw_value += product_value; acc = next_acc; } - offset = range_check_signed_interval( - &mut region, - &config, - &acc, - &raw_value, - &(-&raw_a_bound), - &raw_a_bound, - offset, - )?; let q = div_round_to_canonical_interval(&raw_value, &scale)?; let (intermediate_cell, next_offset) = assign_div_round( &mut region, &config, + window, &acc, &raw_value, &q, @@ -429,6 +580,7 @@ impl Circuit for LoraCircuit { offset = range_check_signed_interval( &mut region, &config, + window, &weight, &weight_big, &(-&value_bound), @@ -446,33 +598,16 @@ impl Circuit for LoraCircuit { )?; offset += 1; let product_value = &weight_big * &intermediate[rank_index].1; - offset = range_check_signed_interval( - &mut region, - &config, - &product, - &product_value, - &(-&product_intermediate_bound), - &product_intermediate_bound, - offset, - )?; let next_acc = assign_add(&mut region, &config, &acc, &product, offset)?; offset += 1; raw_value += product_value; acc = next_acc; } - offset = range_check_signed_interval( - &mut region, - &config, - &acc, - &raw_value, - &(-&raw_b_bound), - &raw_b_bound, - offset, - )?; let rescaled = div_round_to_canonical_interval(&raw_value, &scale)?; let (rescaled_cell, next_offset) = assign_div_round( &mut region, &config, + window, &acc, &raw_value, &rescaled, @@ -493,15 +628,6 @@ impl Circuit for LoraCircuit { offset, )?; offset += 1; - offset = range_check_signed_interval( - &mut region, - &config, - &scaled_raw_cell, - &scaled_raw, - &(-&scaled_raw_bound), - &scaled_raw_bound, - offset, - )?; let scaling_den_big = BigInt::from(self.scaling_den); let final_delta = @@ -509,6 +635,7 @@ impl Circuit for LoraCircuit { let (final_cell, next_offset) = assign_div_round( &mut region, &config, + window, &scaled_raw_cell, &scaled_raw, &final_delta, @@ -570,6 +697,10 @@ impl LoraCircuit { self.x.len() } + fn window_and_k(&self) -> (u32, u32) { + window_and_k_for(self) + } + fn rank(&self) -> usize { self.a.len() } @@ -743,6 +874,7 @@ fn assign_statement_digest_cells( fn assign_div_round( region: &mut halo2_proofs::circuit::Region<'_, Fp>, config: &LoraConfig, + window: u32, raw: &AssignedCell, raw_value: &BigInt, quotient: &BigInt, @@ -758,6 +890,7 @@ fn assign_div_round( offset = range_check_signed_interval( region, config, + window, raw, raw_value, &(-raw_bound), @@ -784,6 +917,7 @@ fn assign_div_round( offset = range_check_signed_interval( region, config, + window, "ient_cell, quotient, &(-quotient_bound), @@ -794,6 +928,7 @@ fn assign_div_round( offset = range_check_signed_interval( region, config, + window, &remainder_cell, &remainder, &lower, @@ -806,6 +941,7 @@ fn assign_div_round( fn range_check_signed_interval( region: &mut halo2_proofs::circuit::Region<'_, Fp>, config: &LoraConfig, + window: u32, value_cell: &AssignedCell, value: &BigInt, lower: &BigInt, @@ -839,6 +975,7 @@ fn range_check_signed_interval( offset = range_check_unsigned( region, config, + window, &shifted_cell, &shifted_unsigned, bits, @@ -863,50 +1000,85 @@ fn range_check_signed_interval( let sum = assign_add(region, config, &shifted_cell, &diff_cell, offset)?; offset += 1; region.constrain_equal(sum.cell(), max_cell.cell())?; - range_check_unsigned(region, config, &diff_cell, &diff, bits, offset) + range_check_unsigned(region, config, window, &diff_cell, &diff, bits, offset) } +/// Constrain `value_cell` to a small non-negative integer using a +/// lookup-backed running sum. +/// +/// Row `offset + i` holds `z_i` in `advice[3]`; the lookup argument enforces +/// `scale_i * (z_i - z_{i+1} * 2^window) ∈ [0, 2^window)` per active row, with +/// `scale_i = 2^(window - top_bits)` on the final limb and `z_n` constrained +/// to zero. +/// +/// Soundness: telescoping the chain over the field gives +/// `z_0 = Σ_{i, config: &LoraConfig, + window: u32, value_cell: &AssignedCell, value: &BigUint, bits: usize, mut offset: usize, ) -> Result { - if bits == 0 || bits > FIELD_SAFE_BITS { + if bits == 0 || bits > FIELD_SAFE_BITS || window == 0 { return Err(Error::Synthesis); } - let mut acc = assign_constant_cell( - region, - config.advice[0], - offset, - Fp::from(0), - "range accumulator zero", - )?; - offset += 1; - for bit_index in (0..bits).rev() { - let bit_value = if ((value >> bit_index) & BigUint::one()).is_zero() { - 0 + let limb_count = bits.div_ceil(window as usize); + let top_bits = bits - (limb_count - 1) * window as usize; + let factor = Fp::from(1u64 << window); + let mut z_value = value.clone(); + + for limb_index in 0..limb_count { + config.q_range.enable(region, offset)?; + region.assign_fixed( + || "range factor", + config.rc_factor, + offset, + || Value::known(factor), + )?; + let scale = if limb_index + 1 == limb_count { + Fp::from(1u64 << (window as usize - top_bits)) } else { - 1 + Fp::from(1) }; - config.boolean.enable(region, offset)?; - let bit = region.assign_advice( - || "range bit", - config.advice[0], + region.assign_fixed( + || "range scale", + config.rc_scale, offset, - || Value::known(Fp::from(bit_value)), + || Value::known(scale), )?; - offset += 1; - let two = assign_constant_cell(region, config.advice[1], offset, Fp::from(2), "two")?; - offset += 1; - let doubled = assign_mul(region, config, &acc, &two, offset)?; - offset += 1; - acc = assign_add(region, config, &doubled, &bit, offset)?; + let z_cell = region.assign_advice( + || "range z", + config.advice[3], + offset, + || Value::known(fp_from_biguint_checked(&z_value).expect("z fits")), + )?; + if limb_index == 0 { + region.constrain_equal(z_cell.cell(), value_cell.cell())?; + } + z_value >>= window; offset += 1; } - region.constrain_equal(acc.cell(), value_cell.cell())?; + region.assign_advice_from_constant( + || "range z terminal", + config.advice[3], + offset, + Fp::from(0), + )?; + offset += 1; Ok(offset) } @@ -1102,7 +1274,6 @@ fn ceil_log2(value: usize) -> usize { } } -#[cfg(any(test, feature = "python"))] fn adapter_commitment_words_from_input( input: &AdapterCommitmentInput, ) -> Result, NativeError> { @@ -1128,7 +1299,6 @@ fn adapter_commitment_words_from_input( Ok(words) } -#[cfg(any(test, feature = "python"))] fn adapter_commitment_for_input(input: &AdapterCommitmentInput) -> Result { let mut acc = Fp::from(0); for word in adapter_commitment_words_from_input(input)? { @@ -1157,44 +1327,74 @@ fn public_inputs(circuit: &LoraCircuit) -> Result, NativeError> { Ok(inputs) } -fn rows_needed(circuit: &LoraCircuit) -> usize { +/// Rows used by one signed-interval range check at the given window: +/// shift constant + add + two running-sum chains (limbs + terminal zero each) +/// + max constant + diff + sum row. +fn signed_check_rows(bits: usize, window: u32) -> usize { + let limbs = bits.max(1).div_ceil(window as usize); + 2 * (limbs + 1) + 5 +} + +fn rows_needed(circuit: &LoraCircuit, window: u32) -> usize { let value_bits = circuit.fixed_point.value_bits as usize; let intermediate_bits = circuit.fixed_point.intermediate_bits as usize; + let scale_bits = circuit.fixed_point.scale_bits as usize; let raw_a_bits = value_bits .saturating_mul(2) - .saturating_add(ceil_log2(circuit.in_dim().max(1))); + .saturating_add(ceil_log2(circuit.in_dim().max(1))) + .saturating_add(1); let raw_b_bits = value_bits .saturating_add(intermediate_bits) - .saturating_add(ceil_log2(circuit.rank().max(1))); + .saturating_add(ceil_log2(circuit.rank().max(1))) + .saturating_add(1); let scaling_bits = bit_length_i64(circuit.scaling_num).max(1); - let scaled_bits = intermediate_bits.saturating_add(scaling_bits); - let product_bits = value_bits - .saturating_mul(2) - .max(value_bits.saturating_add(intermediate_bits)) - .max(scaled_bits); - let accumulator_bits = raw_a_bits.max(raw_b_bits); - let range_rows = |bits: usize| 8 * bits.max(1) + 16; + let scaled_bits = intermediate_bits.saturating_add(scaling_bits).max(1); + let den_bits = bit_length_i64(circuit.scaling_den).max(1); + let rc = |bits: usize| signed_check_rows(bits, window); let matrix_values = circuit.rank() * circuit.in_dim() + circuit.out_dim() * circuit.rank(); - let products = matrix_values + circuit.out_dim(); - let divs = circuit.rank() + 2 * circuit.out_dim(); let adapter_words = 11 + matrix_values; - 32 + circuit.in_dim() - + 4 * matrix_values - + 4 * products - + 4 * circuit.out_dim() - + circuit.in_dim() * range_rows(value_bits) - + matrix_values * range_rows(value_bits) - + products * range_rows(product_bits) - + (circuit.rank() + circuit.out_dim()) * range_rows(accumulator_bits) - + divs - * (range_rows(raw_a_bits.max(raw_b_bits).max(scaled_bits)) - + range_rows(intermediate_bits)) + // div block: raw range check + division row + quotient check + remainder check + let div_a = rc(raw_a_bits) + 1 + rc(intermediate_bits) + rc(scale_bits.max(1)); + let div_b = rc(raw_b_bits) + 1 + rc(intermediate_bits) + rc(scale_bits.max(1)); + let div_final = rc(scaled_bits) + 1 + rc(value_bits) + rc(den_bits); + 32 + circuit.in_dim() * (1 + rc(value_bits)) + + matrix_values * (3 + rc(value_bits)) + + circuit.rank() * div_a + + circuit.out_dim() * (div_b + 1 + div_final) + adapter_words * POSEIDON_PAIR_ROWS + + ROW_MARGIN +} + +/// Deterministically choose the lookup window and circuit size from the +/// statement shape alone, so prover and verifier always derive the same +/// circuit. Smaller windows shrink the lookup table for tiny circuits while +/// larger windows minimise rows for big ones; we pick the (k, window) pair +/// with the smallest k, preferring larger windows on ties. +fn window_and_k_for(circuit: &LoraCircuit) -> (u32, u32) { + let mut best: Option<(u32, u32)> = None; + for window in RANGE_WINDOW_MIN..=RANGE_WINDOW_MAX { + let rows = rows_needed(circuit, window); + let table_rows = (1usize << window) + ROW_MARGIN; + let needed = rows.max(table_rows).next_power_of_two(); + let k = needed.trailing_zeros().max(8); + let candidate = (k, window); + best = Some(match best { + None => candidate, + Some((best_k, best_w)) => { + if k < best_k || (k == best_k && window > best_w) { + candidate + } else { + (best_k, best_w) + } + } + }); + } + let (k, window) = best.expect("window range is non-empty"); + (window, k) } fn k_for(circuit: &LoraCircuit) -> u32 { - let rows = rows_needed(circuit).next_power_of_two(); - rows.trailing_zeros().max(8) + circuit.window_and_k().1 } fn circuit_from_json(statement_json: &str, witness_json: &str) -> Result { @@ -1237,9 +1437,9 @@ fn default_scaling_den() -> i64 { pub fn prove_bytes(statement_json: &str, witness_json: &str) -> Result, NativeError> { let circuit = circuit_from_json(statement_json, witness_json)?; let k = k_for(&circuit); - let params: Params = Params::new(k); - let vk = keygen_vk(¶ms, &circuit).map_err(|e| NativeError::Halo2(e.to_string()))?; - let pk = keygen_pk(¶ms, vk, &circuit).map_err(|e| NativeError::Halo2(e.to_string()))?; + let params = params_for(k); + let key = CircuitKey::from_circuit(&circuit); + let pk = proving_key_for(&key, ¶ms, &circuit.without_witnesses())?; let instances = public_inputs(&circuit)?; let instance_refs: Vec<&[Fp]> = vec![instances.as_slice()]; let mut transcript = Blake2bWrite::<_, vesta::Affine, Challenge255<_>>::init(vec![]); @@ -1267,8 +1467,9 @@ pub fn verify_bytes(statement_json: &str, proof: &[u8]) -> Result = Params::new(k); - let vk = keygen_vk(¶ms, &circuit).map_err(|e| NativeError::Halo2(e.to_string()))?; + let params = params_for(k); + let key = CircuitKey::from_circuit(&circuit); + let vk = verifying_key_for(&key, ¶ms, &circuit)?; let instances = public_inputs(&circuit)?; let instance_refs: Vec<&[Fp]> = vec![instances.as_slice()]; let mut transcript = Blake2bRead::<_, vesta::Affine, Challenge255<_>>::init(proof); @@ -1288,16 +1489,375 @@ pub fn statement_digest_hex(statement_json: &str) -> String { format!("{:x}", hasher.finalize()) } +/// Shared validation/bounds context for the exact integer delta fast path. +struct DeltaContext { + in_dim: usize, + value_bound: i128, + intermediate_bound: i128, + scale: i128, + scaling_num: i64, + scaling_den: i64, +} + +fn delta_context( + a: &[Vec], + b: &[Vec], + scaling_num: i64, + scaling_den: i64, + scale_bits: u32, + value_bits: u32, + intermediate_bits: u32, +) -> Result { + if scaling_den <= 0 { + return Err(NativeError::InvalidDimensions( + "scaling_den must be positive".into(), + )); + } + if value_bits == 0 || value_bits > 63 || intermediate_bits == 0 || intermediate_bits > 127 { + return Err(NativeError::InvalidDimensions( + "bit widths outside native fast-path range".into(), + )); + } + if scale_bits >= value_bits { + return Err(NativeError::InvalidDimensions( + "scale_bits must be less than value_bits".into(), + )); + } + let rank = a.len(); + if rank == 0 { + return Err(NativeError::InvalidDimensions( + "rank must be positive".into(), + )); + } + let in_dim = a[0].len(); + for row in a { + if row.len() != in_dim { + return Err(NativeError::InvalidDimensions( + "A row width must match x length".into(), + )); + } + } + for row in b { + if row.len() != rank { + return Err(NativeError::InvalidDimensions( + "B row width must match A rank".into(), + )); + } + } + let value_bound = (1i128 << (value_bits - 1)) - 1; + for value in a.iter().flatten() { + check_native_bound(*value, value_bound, "A")?; + } + for value in b.iter().flatten() { + check_native_bound(*value, value_bound, "B")?; + } + Ok(DeltaContext { + in_dim, + value_bound, + intermediate_bound: (1i128 << (intermediate_bits - 1)) - 1, + scale: 1i128 << scale_bits, + scaling_num, + scaling_den, + }) +} + +fn check_native_bound(value: i64, bound: i128, label: &str) -> Result<(), NativeError> { + let value = value as i128; + if value < -bound || value > bound { + return Err(NativeError::InvalidDimensions(format!( + "{label} value {value} exceeds signed bound +/-{bound}" + ))); + } + Ok(()) +} + +fn delta_for_row( + ctx: &DeltaContext, + a: &[Vec], + b: &[Vec], + x: &[i64], +) -> Result, NativeError> { + if x.len() != ctx.in_dim { + return Err(NativeError::InvalidDimensions(format!( + "x expected length {}, got {}", + ctx.in_dim, + x.len() + ))); + } + for value in x { + check_native_bound(*value, ctx.value_bound, "x")?; + } + + let overflow = + || NativeError::InvalidDimensions("intermediate value exceeds native range".into()); + let div_round = |numerator: i128, denominator: i128| -> Result { + // denominator > 0 here; floor((n + d/2) / d), matching div_floor. + let n = numerator + .checked_add(denominator / 2) + .ok_or_else(overflow)?; + Ok(n.div_euclid(denominator)) + }; + + let mut intermediate = Vec::with_capacity(a.len()); + for row in a { + let mut raw: i128 = 0; + for (weight, x_i) in row.iter().zip(x.iter()) { + let product = (*weight as i128) + .checked_mul(*x_i as i128) + .ok_or_else(overflow)?; + raw = raw.checked_add(product).ok_or_else(overflow)?; + } + if raw < -ctx.intermediate_bound || raw > ctx.intermediate_bound { + return Err(NativeError::InvalidDimensions(format!( + "intermediate value {raw} exceeds signed bound +/-{}", + ctx.intermediate_bound + ))); + } + intermediate.push(div_round(raw, ctx.scale)?); + } + + let mut delta = Vec::with_capacity(b.len()); + for row in b { + let mut raw: i128 = 0; + for (weight, value) in row.iter().zip(intermediate.iter()) { + let product = (*weight as i128).checked_mul(*value).ok_or_else(overflow)?; + raw = raw.checked_add(product).ok_or_else(overflow)?; + } + if raw < -ctx.intermediate_bound || raw > ctx.intermediate_bound { + return Err(NativeError::InvalidDimensions(format!( + "intermediate value {raw} exceeds signed bound +/-{}", + ctx.intermediate_bound + ))); + } + let rescaled = div_round(raw, ctx.scale)?; + let scaled = rescaled + .checked_mul(ctx.scaling_num as i128) + .ok_or_else(overflow)?; + let out = div_round(scaled, ctx.scaling_den as i128)?; + if out < -ctx.value_bound || out > ctx.value_bound { + return Err(NativeError::InvalidDimensions(format!( + "delta value {out} exceeds signed bound +/-{}", + ctx.value_bound + ))); + } + delta.push(out as i64); + } + Ok(delta) +} + +/// Exact integer LoRA delta computation mirroring +/// `proof_contract.compute_delta_quantized`. All arithmetic uses checked i128 +/// operations; any overflow or bound violation is reported as an error so the +/// Python caller can fall back to its arbitrary-precision implementation. The +/// rounding rule is the same canonical half-up floor division used in-circuit. +#[allow(clippy::too_many_arguments)] +pub fn compute_delta_quantized_native( + a: &[Vec], + b: &[Vec], + x: &[i64], + scaling_num: i64, + scaling_den: i64, + scale_bits: u32, + value_bits: u32, + intermediate_bits: u32, +) -> Result, NativeError> { + let ctx = delta_context( + a, + b, + scaling_num, + scaling_den, + scale_bits, + value_bits, + intermediate_bits, + )?; + delta_for_row(&ctx, a, b, x) +} + +/// Batched variant: validates the adapter once and computes every row's delta +/// in parallel. Semantics per row are identical to the single-row function. +#[allow(clippy::too_many_arguments)] +pub fn compute_delta_rows_native( + a: &[Vec], + b: &[Vec], + xs: &[Vec], + scaling_num: i64, + scaling_den: i64, + scale_bits: u32, + value_bits: u32, + intermediate_bits: u32, +) -> Result>, NativeError> { + let ctx = delta_context( + a, + b, + scaling_num, + scaling_den, + scale_bits, + value_bits, + intermediate_bits, + )?; + xs.par_iter() + .map(|x| delta_for_row(&ctx, a, b, x)) + .collect() +} + +/// Exact fixed-point quantization of one decimal string, replicating +/// `proof_contract.quantize_scalar`: the caller passes Python's `str(float)` +/// output (the semantics anchor for quantization), and this routine computes +/// `round_half_up_away_from_zero(decimal_value * 2^scale_bits)` in exact +/// integer arithmetic. Rust's own float formatter is deliberately NOT used: +/// it can pick a different shortest round-trip digit string than CPython's +/// repr for the same f64, which would silently change quantized values. +pub fn quantize_decimal_str_exact( + text: &str, + scale_bits: u32, + value_bits: u32, +) -> Result { + if value_bits == 0 || value_bits > 63 || scale_bits >= value_bits { + return Err(NativeError::InvalidDimensions( + "fixed-point bits outside native fast-path range".into(), + )); + } + let value_bound = (1i128 << (value_bits - 1)) - 1; + + let bad = || NativeError::InvalidDimensions(format!("cannot quantize value {text:?}")); + let trimmed = text.trim(); + let (negative, unsigned) = match trimmed.strip_prefix('-') { + Some(rest) => (true, rest), + None => (false, trimmed.strip_prefix('+').unwrap_or(trimmed)), + }; + let (mantissa_text, exponent) = match unsigned.split_once(['e', 'E']) { + Some((mantissa, exp)) => (mantissa, exp.parse::().map_err(|_| bad())?), + None => (unsigned, 0), + }; + let (int_text, frac_text) = match mantissa_text.split_once('.') { + Some((int_part, frac_part)) => (int_part, frac_part), + None => (mantissa_text, ""), + }; + if int_text.is_empty() && frac_text.is_empty() { + return Err(bad()); + } + let digits: String = format!("{int_text}{frac_text}"); + if digits.is_empty() || digits.len() > 30 || !digits.bytes().all(|b| b.is_ascii_digit()) { + // Non-finite values ('inf'/'nan'), or more digits than the i128 fast + // path can hold: defer to the Python reference implementation. + return Err(bad()); + } + let mantissa: i128 = digits.parse().map_err(|_| bad())?; + let pow10 = exponent + .checked_sub(frac_text.len() as i32) + .ok_or_else(bad)?; + + let bound_err = || { + NativeError::InvalidDimensions(format!( + "quantized value of {text:?} exceeds signed bound +/-{value_bound}" + )) + }; + let magnitude = if mantissa == 0 { + 0i128 + } else if pow10 >= 0 { + // Integral decimal value: scale exactly, no rounding needed. + let mut scaled = mantissa; + for _ in 0..pow10 { + scaled = scaled.checked_mul(10).ok_or_else(bound_err)?; + } + scaled + .checked_mul(1i128 << scale_bits) + .ok_or_else(bound_err)? + } else { + let k = -pow10 as u32; + // numerator = mantissa * 2^scale_bits, denominator = 10^k. With at + // most 30 mantissa digits the numerator can overflow i128 only via + // checked_mul (reported as out-of-fast-path); for k beyond i128's + // 10^38 capacity the quotient is zero when numerator < 10^k / 2. + let numerator = mantissa + .checked_mul(1i128 << scale_bits) + .ok_or_else(bound_err)?; + if k > 38 { + // 10^39 / 2 = 5e38 exceeds i128::MAX (~1.7e38), so any in-range + // numerator is strictly below half the denominator: rounds to 0. + 0 + } else { + let denominator = 10i128.pow(k); + (numerator + denominator / 2) / denominator + } + }; + let signed = if negative { -magnitude } else { magnitude }; + if signed < -value_bound || signed > value_bound { + return Err(NativeError::InvalidDimensions(format!( + "quantized value {signed} exceeds signed bound +/-{value_bound}" + ))); + } + Ok(signed as i64) +} + +/// Quantize whole rows in parallel with exact scalar semantics. +pub fn quantize_rows_exact( + rows: &[Vec], + scale_bits: u32, + value_bits: u32, +) -> Result>, NativeError> { + rows.par_iter() + .map(|row| { + row.iter() + .map(|value| quantize_decimal_str_exact(value, scale_bits, value_bits)) + .collect() + }) + .collect() +} + +const MERKLE_EMPTY: [u8; 32] = [0u8; 32]; + +/// Hiding Merkle root over f64 leaves, byte-identical to +/// `zklora.polynomial_commit._merkle_root` (BLAKE3 leaves salted with the +/// nonce, right-padded with the EMPTY leaf so internal levels stay even). +pub fn merkle_root_f64(values: &[f64], nonce: &[u8]) -> [u8; 32] { + if values.is_empty() { + return MERKLE_EMPTY; + } + let mut level: Vec<[u8; 32]> = values + .par_iter() + .map(|value| { + let mut hasher = blake3::Hasher::new(); + hasher.update(&value.to_be_bytes()); + hasher.update(nonce); + *hasher.finalize().as_bytes() + }) + .collect(); + if level.len() % 2 == 1 { + level.push(MERKLE_EMPTY); + } + while level.len() > 1 { + let mut next: Vec<[u8; 32]> = level + .par_chunks_exact(2) + .map(|pair| { + let mut hasher = blake3::Hasher::new(); + hasher.update(&pair[0]); + hasher.update(&pair[1]); + *hasher.finalize().as_bytes() + }) + .collect(); + if next.len() % 2 == 1 && next.len() != 1 { + next.push(MERKLE_EMPTY); + } + level = next; + } + level[0] +} + #[cfg(feature = "python")] #[pyo3::pyfunction] -fn prove(statement_json: &str, witness_json: &str) -> pyo3::PyResult> { - Ok(prove_bytes(statement_json, witness_json)?) +fn prove( + py: pyo3::Python<'_>, + statement_json: &str, + witness_json: &str, +) -> pyo3::PyResult> { + py.detach(|| Ok(prove_bytes(statement_json, witness_json)?)) } #[cfg(feature = "python")] #[pyo3::pyfunction] -fn verify(statement_json: &str, proof: &[u8]) -> pyo3::PyResult { - Ok(verify_bytes(statement_json, proof)?) +fn verify(py: pyo3::Python<'_>, statement_json: &str, proof: &[u8]) -> pyo3::PyResult { + py.detach(|| Ok(verify_bytes(statement_json, proof)?)) } #[cfg(feature = "python")] @@ -1308,10 +1868,85 @@ fn statement_digest(statement_json: &str) -> pyo3::PyResult { #[cfg(feature = "python")] #[pyo3::pyfunction] -fn adapter_commitment(adapter_json: &str) -> pyo3::PyResult { - let input: AdapterCommitmentInput = - serde_json::from_str(adapter_json).map_err(|e| NativeError::Json(e.to_string()))?; - Ok(adapter_commitment_for_input(&input)?) +fn adapter_commitment(py: pyo3::Python<'_>, adapter_json: &str) -> pyo3::PyResult { + py.detach(|| { + let input: AdapterCommitmentInput = + serde_json::from_str(adapter_json).map_err(|e| NativeError::Json(e.to_string()))?; + Ok(adapter_commitment_for_input(&input)?) + }) +} + +#[cfg(feature = "python")] +#[pyo3::pyfunction] +#[allow(clippy::too_many_arguments)] +fn compute_delta_quantized( + py: pyo3::Python<'_>, + a: Vec>, + b: Vec>, + x: Vec, + scaling_num: i64, + scaling_den: i64, + scale_bits: u32, + value_bits: u32, + intermediate_bits: u32, +) -> pyo3::PyResult> { + py.detach(|| { + Ok(compute_delta_quantized_native( + &a, + &b, + &x, + scaling_num, + scaling_den, + scale_bits, + value_bits, + intermediate_bits, + )?) + }) +} + +#[cfg(feature = "python")] +#[pyo3::pyfunction] +#[allow(clippy::too_many_arguments)] +fn compute_delta_rows( + py: pyo3::Python<'_>, + a: Vec>, + b: Vec>, + xs: Vec>, + scaling_num: i64, + scaling_den: i64, + scale_bits: u32, + value_bits: u32, + intermediate_bits: u32, +) -> pyo3::PyResult>> { + py.detach(|| { + Ok(compute_delta_rows_native( + &a, + &b, + &xs, + scaling_num, + scaling_den, + scale_bits, + value_bits, + intermediate_bits, + )?) + }) +} + +#[cfg(feature = "python")] +#[pyo3::pyfunction] +fn quantize_rows( + py: pyo3::Python<'_>, + rows: Vec>, + scale_bits: u32, + value_bits: u32, +) -> pyo3::PyResult>> { + py.detach(|| Ok(quantize_rows_exact(&rows, scale_bits, value_bits)?)) +} + +#[cfg(feature = "python")] +#[pyo3::pyfunction] +fn merkle_root(py: pyo3::Python<'_>, values: Vec, nonce: Vec) -> pyo3::PyResult> { + py.detach(|| Ok(merkle_root_f64(&values, &nonce).to_vec())) } #[cfg(feature = "python")] @@ -1324,9 +1959,134 @@ fn _native_prover(m: &pyo3::Bound<'_, pyo3::types::PyModule>) -> pyo3::PyResult< m.add_function(wrap_pyfunction!(verify, m)?)?; m.add_function(wrap_pyfunction!(statement_digest, m)?)?; m.add_function(wrap_pyfunction!(adapter_commitment, m)?)?; + m.add_function(wrap_pyfunction!(compute_delta_quantized, m)?)?; + m.add_function(wrap_pyfunction!(compute_delta_rows, m)?)?; + m.add_function(wrap_pyfunction!(quantize_rows, m)?)?; + m.add_function(wrap_pyfunction!(merkle_root, m)?)?; Ok(()) } +#[doc(hidden)] +pub mod bench_support { + //! Deterministic statement/witness generation for benchmarking only. + + use super::*; + + fn lcg_value(state: &mut u64, magnitude: i64) -> i64 { + *state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + let raw = (*state >> 16) as i64; + (raw % (2 * magnitude + 1)) - magnitude + } + + fn div_round(numerator: &BigInt, denominator: &BigInt) -> BigInt { + let floor_half = denominator / BigInt::from(2u8); + (numerator + floor_half).div_floor(denominator) + } + + pub fn bench_statement_and_witness( + in_dim: usize, + rank: usize, + out_dim: usize, + ) -> (String, String, u32) { + let fixed_point = FixedPointConfig { + scale_bits: 20, + value_bits: 63, + intermediate_bits: 127, + }; + let scale = BigInt::one() << fixed_point.scale_bits; + let scaling_num = 1i64; + let scaling_den = 1i64; + let mut state = 0x5eed_5eed_5eed_5eedu64; + let magnitude = 1i64 << fixed_point.scale_bits; + + let a: Vec> = (0..rank) + .map(|_| { + (0..in_dim) + .map(|_| lcg_value(&mut state, magnitude)) + .collect() + }) + .collect(); + let b: Vec> = (0..out_dim) + .map(|_| { + (0..rank) + .map(|_| lcg_value(&mut state, magnitude)) + .collect() + }) + .collect(); + let x: Vec = (0..in_dim) + .map(|_| lcg_value(&mut state, magnitude)) + .collect(); + + let intermediate: Vec = a + .iter() + .map(|row| { + let raw: BigInt = row + .iter() + .zip(x.iter()) + .map(|(w, xi)| BigInt::from(*w) * BigInt::from(*xi)) + .sum(); + div_round(&raw, &scale) + }) + .collect(); + let delta: Vec = b + .iter() + .map(|row| { + let raw: BigInt = row + .iter() + .zip(intermediate.iter()) + .map(|(w, v)| BigInt::from(*w) * v) + .sum(); + let rescaled = div_round(&raw, &scale); + let scaled = rescaled * BigInt::from(scaling_num); + let out = div_round(&scaled, &BigInt::from(scaling_den)); + i64::try_from(out).expect("bench delta fits i64") + }) + .collect(); + + let adapter_input = AdapterCommitmentInput { + schema_version: ARTIFACT_SCHEMA_VERSION, + in_dim, + rank, + out_dim, + fixed_point: fixed_point.clone(), + scaling_num, + scaling_den, + a: a.clone(), + b: b.clone(), + }; + let commitment = adapter_commitment_for_input(&adapter_input).expect("commitment"); + + let statement = NativeStatement { + x, + delta, + fixed_point, + rank, + scaling_num, + scaling_den, + adapter_commitment: commitment, + statement_digest: "ab".repeat(32), + }; + let witness = NativeWitness { a, b }; + let statement_json = serde_json::to_string(&statement).expect("statement json"); + let witness_json = serde_json::to_string(&witness).expect("witness json"); + let circuit = circuit_from_json(&statement_json, &witness_json).expect("circuit"); + let k = k_for(&circuit); + (statement_json, witness_json, k) + } + + pub fn prove_verify_once(statement_json: &str, witness_json: &str) -> (f64, f64, usize) { + let prove_start = std::time::Instant::now(); + let proof = prove_bytes(statement_json, witness_json).expect("prove"); + let prove_ms = prove_start.elapsed().as_secs_f64() * 1000.0; + let verify_start = std::time::Instant::now(); + assert!(verify_bytes(statement_json, &proof).expect("verify")); + let verify_ms = verify_start.elapsed().as_secs_f64() * 1000.0; + (prove_ms, verify_ms, proof.len()) + } +} + #[cfg(test)] mod tests { use super::*; @@ -1413,6 +2173,50 @@ mod tests { assert_eq!(prover.verify(), Ok(())); } + #[test] + fn mock_prover_accepts_extreme_in_bound_values() { + // Values sitting exactly on the signed bound exercise the zero-slack + // top-limb scaling of the lookup range checks. + let fixed_point = FixedPointConfig { + scale_bits: 0, + value_bits: 8, + intermediate_bits: 24, + }; + let bound = (1i64 << (fixed_point.value_bits - 1)) - 1; + let input = AdapterCommitmentInput { + schema_version: ARTIFACT_SCHEMA_VERSION, + in_dim: 2, + rank: 1, + out_dim: 1, + fixed_point: fixed_point.clone(), + scaling_num: 1, + scaling_den: 1, + a: vec![vec![bound, -bound]], + b: vec![vec![1]], + }; + let x = vec![bound, bound]; + // raw = bound*bound - bound*bound = 0; delta = 0. + let circuit = LoraCircuit { + a: input.a.clone(), + b: input.b.clone(), + x, + delta: vec![0], + fixed_point, + scaling_num: 1, + scaling_den: 1, + adapter_commitment: adapter_commitment_for_input(&input).unwrap(), + statement_digest: "44".repeat(32), + }; + let instances = public_inputs(&circuit).unwrap(); + let prover = MockProver::run(k_for(&circuit), &circuit, vec![instances]).unwrap(); + assert_eq!(prover.verify(), Ok(())); + + // One past the bound must be rejected before synthesis even starts. + let mut out_of_bound = circuit; + out_of_bound.x = vec![bound + 1, bound]; + assert!(out_of_bound.validate().is_err()); + } + #[test] fn mock_prover_rejects_tampered_delta_and_commitment() { let circuit = valid_circuit(); @@ -1430,7 +2234,6 @@ mod tests { } #[test] - #[ignore = "IPA proof generation for the Poseidon/range-check circuit is intentionally slow"] fn real_proof_verifies_for_tiny_relation() { let circuit = minimal_circuit(); let statement = NativeStatement { @@ -1462,6 +2265,146 @@ mod tests { assert!(!verify_bytes(&tampered_json, &proof).unwrap()); } + #[test] + fn window_selection_depends_only_on_shape() { + let circuit = valid_circuit(); + let (window, k) = circuit.window_and_k(); + assert!((RANGE_WINDOW_MIN..=RANGE_WINDOW_MAX).contains(&window)); + assert!(k >= 8); + // Witness values must not influence the selected circuit size. + let mut other = circuit.clone(); + other.a = vec![vec![1, 1]]; + other.b = vec![vec![-1], vec![1]]; + other.x = vec![7, -7]; + other.delta = vec![0, 0]; + assert_eq!(other.window_and_k(), (window, k)); + // The empty circuit used for key generation must agree as well. + assert_eq!(circuit.without_witnesses().window_and_k(), (window, k)); + } + + #[test] + fn mock_prover_accepts_multi_limb_range_checks() { + // value_bits=17 with window selection forces uneven top limbs in the + // running-sum decomposition; rank 2 exercises both matmul stages. + let fixed_point = FixedPointConfig { + scale_bits: 3, + value_bits: 17, + intermediate_bits: 40, + }; + let input = AdapterCommitmentInput { + schema_version: ARTIFACT_SCHEMA_VERSION, + in_dim: 5, + rank: 2, + out_dim: 3, + fixed_point: fixed_point.clone(), + scaling_num: 3, + scaling_den: 2, + a: vec![vec![100, -50, 25, -12, 6], vec![-99, 98, -97, 96, -95]], + b: vec![vec![40, -30], vec![-20, 10], vec![5, -5]], + }; + let x = vec![64, -32, 16, -8, 4]; + let scale = 1i64 << fixed_point.scale_bits; + let div_round = |n: i64, d: i64| -> i64 { (n + d / 2).div_euclid(d) }; + let intermediate: Vec = input + .a + .iter() + .map(|row| { + let raw: i64 = row.iter().zip(x.iter()).map(|(w, xi)| w * xi).sum(); + div_round(raw, scale) + }) + .collect(); + let delta: Vec = input + .b + .iter() + .map(|row| { + let raw: i64 = row + .iter() + .zip(intermediate.iter()) + .map(|(w, v)| w * v) + .sum(); + let rescaled = div_round(raw, scale); + div_round(rescaled * input.scaling_num, input.scaling_den) + }) + .collect(); + let circuit = LoraCircuit { + a: input.a.clone(), + b: input.b.clone(), + x, + delta, + fixed_point, + scaling_num: input.scaling_num, + scaling_den: input.scaling_den, + adapter_commitment: adapter_commitment_for_input(&input).unwrap(), + statement_digest: "33".repeat(32), + }; + let instances = public_inputs(&circuit).unwrap(); + let prover = MockProver::run(k_for(&circuit), &circuit, vec![instances]).unwrap(); + assert_eq!(prover.verify(), Ok(())); + + let mut tampered = public_inputs(&circuit).unwrap(); + tampered[circuit.in_dim()] += Fp::from(1); + let prover = MockProver::run(k_for(&circuit), &circuit, vec![tampered]).unwrap(); + assert!(prover.verify().is_err()); + } + + #[test] + fn native_delta_helper_matches_circuit_relation() { + let circuit = valid_circuit(); + let delta = compute_delta_quantized_native( + &circuit.a, + &circuit.b, + &circuit.x, + circuit.scaling_num, + circuit.scaling_den, + circuit.fixed_point.scale_bits, + circuit.fixed_point.value_bits, + circuit.fixed_point.intermediate_bits, + ) + .unwrap(); + assert_eq!(delta, circuit.delta); + } + + #[test] + fn native_delta_reports_div_round_overflow_as_error() { + // intermediate = 2^62 - 1, raw_b = 2^65 - 8, and + // scaled = (2^65 - 8) * (2^62 + 1) = 2^127 - 8 fits i128, but adding + // scaling_den / 2 = 8 in the rounding step would overflow. The fast + // path must surface this as an error (so the Python caller falls back + // to the arbitrary-precision path) rather than panic or wrap. + let result = compute_delta_quantized_native( + &[vec![1]], + &[vec![8]], + &[(1i64 << 62) - 1], + (1i64 << 62) + 1, + 17, + 0, + 63, + 127, + ); + assert!(result.is_err()); + } + + #[test] + fn merkle_root_handles_padding_rules() { + let nonce = [7u8; 32]; + assert_eq!(merkle_root_f64(&[], &nonce), MERKLE_EMPTY); + let single = merkle_root_f64(&[1.5], &nonce); + let pair = merkle_root_f64(&[1.5, 2.5], &nonce); + let triple = merkle_root_f64(&[1.5, 2.5, 3.5], &nonce); + assert_ne!(single, pair); + assert_ne!(pair, triple); + // Right-padding with the EMPTY leaf means a singleton tree hashes the + // leaf against EMPTY rather than returning the leaf itself. + let mut leaf = blake3::Hasher::new(); + leaf.update(&1.5f64.to_be_bytes()); + leaf.update(&nonce); + let leaf = *leaf.finalize().as_bytes(); + let mut parent = blake3::Hasher::new(); + parent.update(&leaf); + parent.update(&MERKLE_EMPTY); + assert_eq!(single, *parent.finalize().as_bytes()); + } + #[test] fn deterministic_half_point_rounding_is_unique() { let statement = NativeStatement { diff --git a/src/zklora/base_model_user_mpi/__init__.py b/src/zklora/base_model_user_mpi/__init__.py index 1ace0a2..117fd39 100644 --- a/src/zklora/base_model_user_mpi/__init__.py +++ b/src/zklora/base_model_user_mpi/__init__.py @@ -11,8 +11,7 @@ from ..proof_contract import ( FixedPointConfig, TranscriptEntry, - flatten, - quantize_nested, + quantize_rows, write_transcript, ) @@ -242,10 +241,7 @@ def record( f"{len(x_rows)} inputs vs {len(delta_rows)} deltas" ) if q_delta_values is None: - q_delta_rows = [ - flatten(quantize_nested(delta_row, self.fixed_point)) - for delta_row in delta_rows - ] + q_delta_rows = quantize_rows(delta_rows, self.fixed_point) else: q_delta_rows = _canonical_int_rows(q_delta_values) if len(x_rows) != len(q_delta_rows): @@ -253,11 +249,11 @@ def record( f"transcript row mismatch for {module_name}: " f"{len(x_rows)} inputs vs {len(q_delta_rows)} q_deltas" ) + q_x_rows = quantize_rows(x_rows, self.fixed_point) entries: list[TranscriptEntry] = [] - for x_row, q_delta in zip(x_rows, q_delta_rows): + for q_x, q_delta in zip(q_x_rows, q_delta_rows): invocation_index = self._counts.get(module_name, 0) self._counts[module_name] = invocation_index + 1 - q_x = flatten(quantize_nested(x_row, self.fixed_point)) entry = TranscriptEntry( session_id=self.session_id, module_name=module_name, diff --git a/src/zklora/lora_contributor_mpi/__init__.py b/src/zklora/lora_contributor_mpi/__init__.py index cdf1032..3475d15 100644 --- a/src/zklora/lora_contributor_mpi/__init__.py +++ b/src/zklora/lora_contributor_mpi/__init__.py @@ -1,3 +1,4 @@ +import hashlib import socket import threading import os @@ -14,9 +15,9 @@ FixedPointConfig, InvocationWitness, adapter_manifest_entry, - compute_delta_quantized, - flatten, + compute_delta_quantized_rows, quantize_nested, + quantize_rows, write_adapter_manifest, ) @@ -77,29 +78,61 @@ def __init__( self._invocation_counts: dict[tuple[str, str], int] = {} self.last_scaling: tuple[int, int] = (1, 1) self.last_q_delta: list[list[int]] = [] + self._module_cache: dict[str, dict] = {} + self._module_cache_lock = threading.Lock() def list_lora_injection_points(self): return list(self.submodules.keys()) - def adapter_manifest_entries(self): - entries = [] - for sub_name, module in self.submodules.items(): - a_matrix, b_matrix, scaling_num, scaling_den = lora_matrices_and_scaling( - module - ) - a_quantized = quantize_nested( + def _module_artifacts(self, sub_name: str) -> dict: + """Quantize a module's LoRA matrices once and reuse them everywhere. + + Quantization runs element-wise through Decimal for exact round-half-up + semantics, which is far too slow to repeat on every invocation. The + cache is keyed on the current weight bytes and scaling, so a module + whose adapter weights are swapped or updated in place is re-quantized + instead of being served stale matrices. + """ + module = self.submodules[sub_name] + a_matrix, b_matrix, scaling_num, scaling_den = lora_matrices_and_scaling(module) + fingerprint = ( + hashlib.sha256(a_matrix.numpy().tobytes()).digest(), + hashlib.sha256(b_matrix.numpy().tobytes()).digest(), + int(scaling_num), + int(scaling_den), + ) + with self._module_cache_lock: + cached = self._module_cache.get(sub_name) + if cached is not None and cached["fingerprint"] == fingerprint: + return cached + artifacts = { + "fingerprint": fingerprint, + "a_matrix": a_matrix, + "b_matrix": b_matrix, + "scaling_num": int(scaling_num), + "scaling_den": int(scaling_den), + "a_quantized": quantize_nested( a_matrix.detach().cpu().numpy().tolist(), self.fixed_point - ) - b_quantized = quantize_nested( + ), + "b_quantized": quantize_nested( b_matrix.detach().cpu().numpy().tolist(), self.fixed_point - ) + ), + } + with self._module_cache_lock: + self._module_cache[sub_name] = artifacts + return artifacts + + def adapter_manifest_entries(self): + entries = [] + for sub_name in self.submodules: + artifacts = self._module_artifacts(sub_name) entries.append( adapter_manifest_entry( sub_name, - a_quantized, - b_quantized, - scaling_num, - scaling_den, + artifacts["a_quantized"], + artifacts["b_quantized"], + artifacts["scaling_num"], + artifacts["scaling_den"], self.fixed_point, ) ) @@ -116,11 +149,20 @@ def apply_lora( ): if sub_name not in self.submodules: raise ValueError(f"[LoRAServer] submodule '{sub_name}' not recognized.") - mod = self.submodules[sub_name] print(f"[A] apply_lora on '{sub_name}', shape={list(input_tensor.shape)}") + artifacts = self._module_artifacts(sub_name) + a_matrix = artifacts["a_matrix"] + b_matrix = artifacts["b_matrix"] + scaling_num = artifacts["scaling_num"] + scaling_den = artifacts["scaling_den"] + a_quantized = artifacts["a_quantized"] + b_quantized = artifacts["b_quantized"] + scaling = scaling_num / scaling_den with torch.no_grad(): - delta_float, a_matrix, b_matrix, scaling_num, scaling_den = ( - compute_lora_delta(mod, input_tensor) + x_float = input_tensor.float() + delta_float = ( + torch.matmul(torch.matmul(x_float, a_matrix.t()), b_matrix.t()) + * scaling ) self.last_scaling = (int(scaling_num), int(scaling_den)) self.last_q_delta = [] @@ -128,25 +170,19 @@ def apply_lora( sid = session_id or "default-session" key = (sid, sub_name) x_rows = input_tensor.detach().cpu().float().reshape(-1, int(a_matrix.shape[1])) - delta_rows = [] - a_quantized = quantize_nested( - a_matrix.detach().cpu().numpy().tolist(), self.fixed_point + q_x_rows = quantize_rows( + [x_row.cpu().numpy().tolist() for x_row in x_rows], self.fixed_point ) - b_quantized = quantize_nested( - b_matrix.detach().cpu().numpy().tolist(), self.fixed_point + q_delta_rows = compute_delta_quantized_rows( + a_quantized, + b_quantized, + q_x_rows, + scaling_num, + scaling_den, + self.fixed_point, ) - for x_row in x_rows: - q_x = flatten( - quantize_nested(x_row.cpu().numpy().tolist(), self.fixed_point) - ) - q_delta = compute_delta_quantized( - a_quantized, - b_quantized, - q_x, - scaling_num, - scaling_den, - self.fixed_point, - ) + delta_rows = [] + for q_x, q_delta in zip(q_x_rows, q_delta_rows): delta_rows.append([int(v) for v in q_delta]) invocation_index = self._invocation_counts.get(key, 0) self._invocation_counts[key] = invocation_index + 1 @@ -219,6 +255,8 @@ def __init__( self.lora_server = lora_server self.stop_event = stop_event self.stop_timeout = stop_timeout + self._server_lock = threading.Lock() + self._conn_threads: list[threading.Thread] = [] def run(self): print(f"[A-Server] listening on {self.host}:{self.port}") @@ -236,8 +274,16 @@ def run(self): conn, addr = srv.accept() except socket.timeout: continue - self.handle_conn(conn, addr) + worker = threading.Thread( + target=self.handle_conn, args=(conn, addr), daemon=True + ) + worker.start() + self._conn_threads = [t for t in self._conn_threads if t.is_alive()] + [ + worker + ] finally: + for worker in self._conn_threads: + worker.join(timeout=self.stop_timeout) srv.close() print("[A-Server] shutting down...") @@ -257,19 +303,23 @@ def handle_conn(self, conn, addr): arr = req["input_array"] session_id = req.get("session_id") tin = torch.tensor(arr, dtype=torch.float32) - out = self.lora_server.apply_lora(sname, tin, session_id=session_id) - resp = { - "response_type": "lora_forward_response", - "output_array": out.cpu().numpy(), - "q_delta": self.lora_server.last_q_delta, - "scaling_num": int(self.lora_server.last_scaling[0]), - "scaling_den": int(self.lora_server.last_scaling[1]), - } + # The lock keeps apply_lora and the last_* snapshot atomic when + # multiple clients are connected concurrently. + with self._server_lock: + out = self.lora_server.apply_lora(sname, tin, session_id=session_id) + resp = { + "response_type": "lora_forward_response", + "output_array": out.cpu().numpy(), + "q_delta": self.lora_server.last_q_delta, + "scaling_num": int(self.lora_server.last_scaling[0]), + "scaling_den": int(self.lora_server.last_scaling[1]), + } elif rtype == "end_inference": - self.lora_server.finalize_proofs_and_collect( - session_id=req.get("session_id") - ) + with self._server_lock: + self.lora_server.finalize_proofs_and_collect( + session_id=req.get("session_id") + ) resp = { "response_type": "end_inference_ack", "message": "A finished native zkLoRA proof generation locally.", diff --git a/src/zklora/polynomial_commit.py b/src/zklora/polynomial_commit.py index 2c59407..fc0020f 100644 --- a/src/zklora/polynomial_commit.py +++ b/src/zklora/polynomial_commit.py @@ -4,6 +4,11 @@ from blake3 import blake3 # type: ignore +try: # native fast path; byte-identical to the Python implementation below + from zklora import _native_prover as _native_merkle_module +except ImportError: # pragma: no cover - extension not built in this env + _native_merkle_module = None + # Merkle-based vector commitment parameters LEAF_EMPTY = b"\x00" * 32 # same as EMPTY_HASH in Rust implementation @@ -46,6 +51,18 @@ def _merkle_root(values: List[Union[int, float]], nonce: bytes) -> bytes: if not values: return LEAF_EMPTY + if _native_merkle_module is not None and hasattr( + _native_merkle_module, "merkle_root" + ): + try: + return bytes( + _native_merkle_module.merkle_root( + [float(v) for v in values], bytes(nonce) + ) + ) + except (OverflowError, TypeError, ValueError): + pass # fall back to the pure-Python implementation + # Convert to leaf hashes with nonce level: List[bytes] = [_hash_leaf(v, nonce) for v in values] @@ -65,6 +82,43 @@ def _merkle_root(values: List[Union[int, float]], nonce: bytes) -> bytes: return level[0] +def _flatten_input_data(data) -> List[Union[int, float]]: + """Flatten ``data["input_data"]`` to a flat list of scalars. + + Uses numpy for rectangular data and a recursive fallback for ragged data, + preserving depth-first leaf order in both cases. + """ + try: + import numpy as np # local import to avoid hard dependency + + return np.asarray(data["input_data"], dtype=np.float64).reshape(-1).tolist() + except Exception: + + def _flatten(x): + for y in x: + if isinstance(y, (list, tuple)): + yield from _flatten(y) + else: + yield y + + return list(_flatten(data["input_data"])) + + +def _merkle_root_from_file(activations_path: str, nonce: bytes) -> bytes: + """Compute the hiding Merkle root of an activations file. + + JSON parsing stays in Python on purpose: it defines the exact float + semantics of the commitment (the nearest-f64 produced by ``json``/``float``), + and a native JSON parser does not always round identically. The leaf and + tree hashing then runs through ``_merkle_root``, which uses the native + BLAKE3 fast path when available and is byte-identical to the pure-Python + reference. + """ + with open(activations_path, "r") as f: + data = json.load(f) + return _merkle_root(_flatten_input_data(data), nonce) + + # -------------------------------------------------------------------------------------- # Public API (names preserved for backwards compatibility) # -------------------------------------------------------------------------------------- @@ -80,32 +134,11 @@ def commit_activations(activations_path: str) -> str: JSON string containing both the Merkle root and nonce: {"root": "0x...", "nonce": "0x..."} """ - with open(activations_path, "r") as f: - data = json.load(f) - - # Flatten arbitrarily nested lists using numpy when available for speed - try: - import numpy as np # local import to avoid hard dependency - - flat_vals = ( - np.asarray(data["input_data"], dtype=np.float64).reshape(-1).tolist() - ) - except Exception: - # fallback: naïve Python flatten - def _flatten(x): - for y in x: - if isinstance(y, (list, tuple)): - yield from _flatten(y) - else: - yield y - - flat_vals = list(_flatten(data["input_data"])) - # Generate random nonce for hiding property nonce = os.urandom(32) - # Compute Merkle root with nonce - root = _merkle_root(flat_vals, nonce) + # Compute Merkle root with nonce (native single-pass when available) + root = _merkle_root_from_file(activations_path, nonce) # Return JSON with both root and nonce commitment_data = {"root": "0x" + root.hex(), "nonce": "0x" + nonce.hex()} @@ -142,29 +175,8 @@ def verify_commitment(activations_path: str, commitment: str) -> bool: # Invalid commitment format return False - # Load and flatten activations - with open(activations_path, "r") as f: - data = json.load(f) - - try: - import numpy as np - - flat_vals = ( - np.asarray(data["input_data"], dtype=np.float64).reshape(-1).tolist() - ) - except Exception: - - def _flatten(x): - for y in x: - if isinstance(y, (list, tuple)): - yield from _flatten(y) - else: - yield y - - flat_vals = list(_flatten(data["input_data"])) - - # Recompute root with provided nonce - computed_root = _merkle_root(flat_vals, nonce) + # Recompute root with provided nonce (native single-pass when available) + computed_root = _merkle_root_from_file(activations_path, nonce) # Compare roots return computed_root == expected_root diff --git a/src/zklora/proof_contract.py b/src/zklora/proof_contract.py index c927043..5fe26ec 100644 --- a/src/zklora/proof_contract.py +++ b/src/zklora/proof_contract.py @@ -5,13 +5,16 @@ import json import os import re +import threading from dataclasses import asdict, dataclass from decimal import Decimal, ROUND_HALF_UP from pathlib import Path from typing import Any, Iterable -BACKEND_ID = "zklora-halo2-v2" +# v3: lookup-based range checks in the native circuit (same statement format, +# same adapter commitment scheme; proofs/vk are not compatible with v2). +BACKEND_ID = "zklora-halo2-v3" SCHEMA_VERSION = 2 COMMITMENT_SCHEME = "sha256-canonical-lora-matrices-v1" ADAPTER_COMMITMENT_SCHEME = "poseidon-pasta-fp-adapter-v1" @@ -143,6 +146,33 @@ def quantize_nested(values: Any, config: FixedPointConfig) -> Any: return quantize_scalar(values, config) +def quantize_rows(rows: list[list[float]], config: FixedPointConfig) -> list[list[int]]: + """Quantize flat float rows with exact quantize_scalar semantics. + + The native implementation reproduces Decimal(str(value)) round-half-up + rounding using exact integer arithmetic over the shortest round-trip + decimal representation; the Python path is the reference fallback. + """ + if not rows: + return [] + native = _native_module() + if native is not None and hasattr(native, "quantize_rows"): + try: + # str(float) is the semantic anchor of quantize_scalar; passing the + # strings keeps the native path bit-identical to Decimal(str(v)). + return [ + [int(v) for v in row] + for row in native.quantize_rows( + [[str(float(v)) for v in row] for row in rows], + config.scale_bits, + config.value_bits, + ) + ] + except (OverflowError, TypeError, ValueError): + pass + return [[quantize_scalar(v, config) for v in row] for row in rows] + + def flatten(values: Any) -> list[int]: if isinstance(values, (list, tuple)): out: list[int] = [] @@ -192,6 +222,76 @@ def compute_delta_quantized( scaling_num: int, scaling_den: int, config: FixedPointConfig, +) -> list[int]: + native = _native_module() + if native is not None and hasattr(native, "compute_delta_quantized"): + try: + return [ + int(v) + for v in native.compute_delta_quantized( + a, + b, + x, + int(scaling_num), + int(scaling_den), + config.scale_bits, + config.value_bits, + config.intermediate_bits, + ) + ] + except (OverflowError, TypeError, ValueError): + # Fall back to the exact arbitrary-precision path below; it either + # produces the result or raises the canonical contract error. + pass + return _compute_delta_quantized_python(a, b, x, scaling_num, scaling_den, config) + + +def compute_delta_quantized_rows( + a: list[list[int]], + b: list[list[int]], + xs: list[list[int]], + scaling_num: int, + scaling_den: int, + config: FixedPointConfig, +) -> list[list[int]]: + """Batched delta computation for multiple input rows of one adapter. + + Uses the parallel native fast path when available (validating the adapter + once instead of per row), with the exact per-row implementation as + fallback. Results are identical to mapping compute_delta_quantized. + """ + if not xs: + return [] + native = _native_module() + if native is not None and hasattr(native, "compute_delta_rows"): + try: + return [ + [int(v) for v in row] + for row in native.compute_delta_rows( + a, + b, + xs, + int(scaling_num), + int(scaling_den), + config.scale_bits, + config.value_bits, + config.intermediate_bits, + ) + ] + except (OverflowError, TypeError, ValueError): + pass + return [ + compute_delta_quantized(a, b, x, scaling_num, scaling_den, config) for x in xs + ] + + +def _compute_delta_quantized_python( + a: list[list[int]], + b: list[list[int]], + x: list[int], + scaling_num: int, + scaling_den: int, + config: FixedPointConfig, ) -> list[int]: if scaling_den <= 0: raise ProofContractError("scaling_den must be positive") @@ -261,6 +361,15 @@ def adapter_commitment_payload( } +_ADAPTER_COMMITMENT_CACHE: dict[tuple, str] = {} +_ADAPTER_COMMITMENT_CACHE_MAX = 64 +_ADAPTER_COMMITMENT_LOCK = threading.Lock() + + +def _freeze_matrix(matrix: list[list[int]]) -> tuple[tuple[int, ...], ...]: + return tuple(tuple(int(v) for v in row) for row in matrix) + + def adapter_commitment( a: list[list[int]], b: list[list[int]], @@ -268,18 +377,39 @@ def adapter_commitment( scaling_den: int, fixed_point: FixedPointConfig, ) -> str: + # The commitment is recomputed for every invocation statement of a module, + # over the same (large) adapter matrices. Content-keyed memoisation keeps + # the Poseidon chain to one evaluation per adapter. The backend identity is + # part of the key so monkeypatched/fake backends never share entries. native = _native_module() if native is None: raise ProofContractError( "native Halo2 prover is unavailable; build/install zklora with maturin" ) - return str( + key = ( + id(native), + _freeze_matrix(a), + _freeze_matrix(b), + int(scaling_num), + int(scaling_den), + fixed_point, + ) + with _ADAPTER_COMMITMENT_LOCK: + cached = _ADAPTER_COMMITMENT_CACHE.get(key) + if cached is not None: + return cached + value = str( native.adapter_commitment( canonical_json( adapter_commitment_payload(a, b, scaling_num, scaling_den, fixed_point) ) ) ) + with _ADAPTER_COMMITMENT_LOCK: + if len(_ADAPTER_COMMITMENT_CACHE) >= _ADAPTER_COMMITMENT_CACHE_MAX: + _ADAPTER_COMMITMENT_CACHE.pop(next(iter(_ADAPTER_COMMITMENT_CACHE))) + _ADAPTER_COMMITMENT_CACHE[key] = value + return value def circuit_id( @@ -501,6 +631,16 @@ def write_invocation_artifacts( } +def _worker_count(env_var: str) -> int: + try: + configured = int(os.environ.get(env_var, "")) + except ValueError: + configured = 0 + if configured > 0: + return configured + return max(os.cpu_count() or 1, 1) + + def load_json(path: str | os.PathLike[str]) -> dict[str, Any]: with open(path, "r", encoding="utf-8") as f: return json.load(f) @@ -722,6 +862,7 @@ def verify_artifacts( | Iterable[dict[str, Any]], ) -> tuple[float, int]: import time + from concurrent.futures import ThreadPoolExecutor start = time.time() entries = load_transcript(transcript) @@ -738,7 +879,21 @@ def verify_artifacts( if key in seen: raise ProofContractError(f"duplicate proof statement for {key}") seen.add(key) - verify_artifact_set(statement_file, entries, adapter_index) + + # Every artifact is verified independently; the native verifier releases + # the GIL, so a small thread pool overlaps proof verification across files. + max_workers = min(len(statement_files), _worker_count("ZKLORA_VERIFY_WORKERS")) + if max_workers > 1: + with ThreadPoolExecutor(max_workers=max_workers) as pool: + futures = [ + pool.submit(verify_artifact_set, statement_file, entries, adapter_index) + for statement_file in statement_files + ] + for future in futures: + future.result() + else: + for statement_file in statement_files: + verify_artifact_set(statement_file, entries, adapter_index) expected = {entry.key() for entry in entries} if seen != expected: diff --git a/src/zklora/zk_proof_generator.py b/src/zklora/zk_proof_generator.py index 065adff..d341ef6 100644 --- a/src/zklora/zk_proof_generator.py +++ b/src/zklora/zk_proof_generator.py @@ -26,7 +26,9 @@ def generate_proofs( no-records result instead of importing removed proof backends. """ + import os import time + from concurrent.futures import ThreadPoolExecutor start = time.time() Path(output_dir).mkdir(parents=True, exist_ok=True) @@ -39,16 +41,35 @@ def generate_proofs( "native zkLoRA proof generation requires captured invocation records" ) - for record in record_list: + def _generate(record: InvocationWitness) -> None: write_invocation_artifacts(output_dir, record) - total_params += record.rank * record.in_dim + record.out_dim * record.rank - proofs += 1 if verbose: print( f"Generated native zkLoRA proof artifact for " f"{record.module_name}#{record.invocation_index}" ) + # The native prover releases the GIL, so independent invocation proofs can + # be generated concurrently; artifact paths are unique per record. + try: + configured = int(os.environ.get("ZKLORA_PROVE_WORKERS", "")) + except ValueError: + configured = 0 + max_workers = configured if configured > 0 else max(os.cpu_count() or 1, 1) + max_workers = min(len(record_list), max_workers) + if max_workers > 1: + with ThreadPoolExecutor(max_workers=max_workers) as pool: + futures = [pool.submit(_generate, record) for record in record_list] + for future in futures: + future.result() + else: + for record in record_list: + _generate(record) + + for record in record_list: + total_params += record.rank * record.in_dim + record.out_dim * record.rank + proofs += 1 + elapsed = time.time() - start return (0.0, 0.0, elapsed, total_params, proofs) diff --git a/tests/test_native_fast_paths.py b/tests/test_native_fast_paths.py new file mode 100644 index 0000000..79ef32c --- /dev/null +++ b/tests/test_native_fast_paths.py @@ -0,0 +1,173 @@ +"""Parity tests: native fast paths must be byte/value-identical to Python.""" + +import importlib +import os +import random +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src")) + +from zklora.proof_contract import ( # noqa: E402 + FixedPointConfig, + ProofContractError, + _compute_delta_quantized_python, + compute_delta_quantized, + compute_delta_quantized_rows, + quantize_rows, + quantize_scalar, +) + +if os.environ.get("ZKLORA_REQUIRE_NATIVE_EXTENSION") == "1": + native = importlib.import_module("zklora._native_prover") +else: + native = pytest.importorskip( + "zklora._native_prover", + reason="native PyO3 extension is not built in this environment", + ) + + +def _random_case(rng, in_dim, rank, out_dim, magnitude): + a = [[rng.randint(-magnitude, magnitude) for _ in range(in_dim)] for _ in range(rank)] + b = [[rng.randint(-magnitude, magnitude) for _ in range(rank)] for _ in range(out_dim)] + x = [rng.randint(-magnitude, magnitude) for _ in range(in_dim)] + return a, b, x + + +def test_native_delta_matches_python_reference_across_shapes(): + rng = random.Random(1234) + config = FixedPointConfig(scale_bits=20, value_bits=63, intermediate_bits=127) + for in_dim, rank, out_dim in [(1, 1, 1), (3, 2, 5), (16, 4, 8), (64, 2, 32)]: + for scaling_num, scaling_den in [(1, 1), (3, 2), (-7, 4), (16, 16)]: + a, b, x = _random_case(rng, in_dim, rank, out_dim, 1 << 21) + expected = _compute_delta_quantized_python( + a, b, x, scaling_num, scaling_den, config + ) + actual = compute_delta_quantized(a, b, x, scaling_num, scaling_den, config) + assert actual == expected + + +def test_native_delta_matches_python_on_rounding_ties(): + # scale 2 with odd raw sums exercises the canonical half-up rounding on + # both positive and negative values. + config = FixedPointConfig(scale_bits=1, value_bits=16, intermediate_bits=32) + for raw in range(-9, 10): + a = [[1]] + b = [[1]] + x = [raw] + expected = _compute_delta_quantized_python(a, b, x, 1, 1, config) + actual = compute_delta_quantized(a, b, x, 1, 1, config) + assert actual == expected, f"mismatch for raw={raw}" + + +def test_native_delta_bound_violations_raise_contract_errors(): + config = FixedPointConfig(scale_bits=2, value_bits=8, intermediate_bits=16) + too_big = config.value_bound + 1 + with pytest.raises(ProofContractError, match="exceeds signed bound"): + compute_delta_quantized([[too_big]], [[1]], [1], 1, 1, config) + with pytest.raises(ProofContractError, match="scaling_den must be positive"): + compute_delta_quantized([[1]], [[1]], [1], 1, 0, config) + + +def test_batched_delta_rows_match_single_row_reference(): + rng = random.Random(5) + config = FixedPointConfig() + a, b, _ = _random_case(rng, 24, 3, 12, 1 << 21) + xs = [ + [rng.randint(-(1 << 21), 1 << 21) for _ in range(24)] for _ in range(7) + ] + expected = [ + _compute_delta_quantized_python(a, b, x, 3, 2, config) for x in xs + ] + assert compute_delta_quantized_rows(a, b, xs, 3, 2, config) == expected + assert compute_delta_quantized_rows(a, b, [], 3, 2, config) == [] + + +def test_native_quantize_rows_matches_decimal_reference(): + rng = random.Random(31) + config = FixedPointConfig() + values = [ + 0.0, + -0.0, + 1.0, + -1.0, + 0.1, + -0.1, + 1.5, + 2.5, + -2.5, + # exact representable ties at the scale boundary + 0.5 / config.scale, + -0.5 / config.scale, + 1.5 / config.scale, + 3.0 / (1 << (config.scale_bits + 1)), + # subnormals and tiny values round to zero + 5e-324, + -5e-324, + 1e-30, + # large but in-bound magnitudes + 1e9, + -123456789.123456789, + 2.0**40, + # values produced by float32 tensors + float.fromhex("0x1.91eb851eb851fp+1"), + ] + values += [rng.uniform(-1e6, 1e6) for _ in range(500)] + values += [rng.uniform(-1e-6, 1e-6) for _ in range(200)] + # Regression: Rust's float formatter picks a different shortest round-trip + # digit string than CPython's str() for this value; quantization must + # follow CPython's, which is why the native path receives decimal strings. + values.append(4313614032396.1562) + expected = [[quantize_scalar(v, config) for v in values]] + assert quantize_rows([values], config) == expected + + +def test_native_quantize_rows_out_of_bound_falls_back_to_exact_error(): + config = FixedPointConfig(scale_bits=2, value_bits=8, intermediate_bits=16) + with pytest.raises(ProofContractError, match="exceeds signed bound"): + quantize_rows([[1e30]], config) + with pytest.raises(Exception): + quantize_rows([[float("nan")]], config) + + +def test_native_merkle_root_matches_python_reference(): + import zklora.polynomial_commit as pc + + rng = random.Random(99) + for count in [1, 2, 3, 4, 5, 8, 13, 64, 100]: + values = [rng.uniform(-1e6, 1e6) for _ in range(count)] + [0.0, -0.0] + nonce = bytes(rng.randrange(256) for _ in range(32)) + leaves = [pc._hash_leaf(v, nonce) for v in values] + if len(leaves) % 2 == 1: + leaves.append(pc.LEAF_EMPTY) + level = leaves + while len(level) > 1: + nxt = [ + pc._parent_hash(level[i], level[i + 1]) for i in range(0, len(level), 2) + ] + if len(nxt) % 2 == 1 and len(nxt) != 1: + nxt.append(pc.LEAF_EMPTY) + level = nxt + expected = level[0] + actual = bytes(native.merkle_root([float(v) for v in values], nonce)) + assert actual == expected + + assert pc._merkle_root([], b"\x00" * 32) == pc.LEAF_EMPTY + + +def test_commit_and_verify_round_trip_uses_native_path(tmp_path): + import json + + from zklora.polynomial_commit import commit_activations, verify_commitment + + payload = {"input_data": [[1.5, -2.25], [3.125, 4.0], [5.5, -6.75]]} + path = tmp_path / "acts.json" + path.write_text(json.dumps(payload), encoding="utf-8") + commitment = commit_activations(str(path)) + assert verify_commitment(str(path), commitment) + + tampered = {"input_data": [[1.5, -2.25], [3.125, 4.0], [5.5, -6.7501]]} + path.write_text(json.dumps(tampered), encoding="utf-8") + assert not verify_commitment(str(path), commitment)