From c4ae150390e423e3ab6c558bffd38cf622dad7ae Mon Sep 17 00:00:00 2001 From: Bidhan Roy Date: Wed, 10 Jun 2026 05:00:06 +0000 Subject: [PATCH 1/8] Make zkLoRA proving/verification 20x+ faster without weakening the proof 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. --- benchmarks/run_benchmarks.py | 123 +++ src/Cargo.toml | 1 + src/README.md | 11 +- src/examples/bench_prove.rs | 27 + src/src/lib.rs | 884 +++++++++++++++++--- src/zklora/lora_contributor_mpi/__init__.py | 109 ++- src/zklora/polynomial_commit.py | 17 + src/zklora/proof_contract.py | 95 ++- src/zklora/zk_proof_generator.py | 27 +- tests/test_native_fast_paths.py | 109 +++ 10 files changed, 1238 insertions(+), 165 deletions(-) create mode 100644 benchmarks/run_benchmarks.py create mode 100644 src/examples/bench_prove.rs create mode 100644 tests/test_native_fast_paths.py 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/src/Cargo.toml b/src/Cargo.toml index 9626e37..3a9d3a2 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" diff --git a/src/README.md b/src/README.md index 8934261..b31f8c7 100644 --- a/src/README.md +++ b/src/README.md @@ -56,7 +56,16 @@ 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. 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..6ed2e02 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, @@ -22,7 +22,9 @@ use num_traits::{One, Signed, Zero}; use rand_core::OsRng; 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 +32,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 +199,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 +232,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 +303,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 +347,10 @@ impl Circuit for LoraCircuit { mul, add, div_round, - boolean, + q_range, + rc_factor, + rc_scale, + rc_table, poseidon_config, } } @@ -215,6 +361,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 +388,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 +489,7 @@ impl Circuit for LoraCircuit { offset = range_check_signed_interval( &mut region, &config, + window, &cell, &value_big, &(-&value_bound), @@ -355,6 +516,7 @@ impl Circuit for LoraCircuit { offset = range_check_signed_interval( &mut region, &config, + window, &weight, &weight_big, &(-&value_bound), @@ -363,6 +525,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 +539,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 +579,7 @@ impl Circuit for LoraCircuit { offset = range_check_signed_interval( &mut region, &config, + window, &weight, &weight_big, &(-&value_bound), @@ -446,33 +597,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 +627,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 +634,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 +696,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 +873,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 +889,7 @@ fn assign_div_round( offset = range_check_signed_interval( region, config, + window, raw, raw_value, &(-raw_bound), @@ -784,6 +916,7 @@ fn assign_div_round( offset = range_check_signed_interval( region, config, + window, "ient_cell, quotient, &(-quotient_bound), @@ -794,6 +927,7 @@ fn assign_div_round( offset = range_check_signed_interval( region, config, + window, &remainder_cell, &remainder, &lower, @@ -806,6 +940,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 +974,7 @@ fn range_check_signed_interval( offset = range_check_unsigned( region, config, + window, &shifted_cell, &shifted_unsigned, bits, @@ -863,50 +999,71 @@ 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 `[0, 2^bits)` 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 so the decomposition +/// covers exactly `bits` bits. `z_n` is constrained to zero, so the chain +/// telescopes to `z_0 = value < 2^bits` with every limb non-negative. fn range_check_unsigned( region: &mut halo2_proofs::circuit::Region<'_, Fp>, 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 +1259,6 @@ fn ceil_log2(value: usize) -> usize { } } -#[cfg(any(test, feature = "python"))] fn adapter_commitment_words_from_input( input: &AdapterCommitmentInput, ) -> Result, NativeError> { @@ -1128,7 +1284,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 +1312,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 +1422,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 +1452,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 +1474,189 @@ pub fn statement_digest_hex(statement_json: &str) -> String { format!("{:x}", hasher.finalize()) } +/// 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. +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> { + 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(); + let out_dim = b.len(); + if x.len() != in_dim { + return Err(NativeError::InvalidDimensions(format!( + "x expected length {in_dim}, got {}", + x.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; + let intermediate_bound = (1i128 << (intermediate_bits - 1)) - 1; + let scale = 1i128 << scale_bits; + let check_value = |value: i64, label: &str| -> Result<(), NativeError> { + let value = value as i128; + if value < -value_bound || value > value_bound { + return Err(NativeError::InvalidDimensions(format!( + "{label} value {value} exceeds signed bound +/-{value_bound}" + ))); + } + Ok(()) + }; + for value in x { + check_value(*value, "x")?; + } + for value in a.iter().flatten() { + check_value(*value, "A")?; + } + for value in b.iter().flatten() { + check_value(*value, "B")?; + } + + let overflow = + || NativeError::InvalidDimensions("intermediate value exceeds native range".into()); + let div_round = |numerator: i128, denominator: i128| -> i128 { + // denominator > 0 here; floor((n + d/2) / d), matching div_floor. + let n = numerator + denominator / 2; + n.div_euclid(denominator) + }; + + let mut intermediate = Vec::with_capacity(rank); + 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 < -intermediate_bound || raw > intermediate_bound { + return Err(NativeError::InvalidDimensions(format!( + "intermediate value {raw} exceeds signed bound +/-{intermediate_bound}" + ))); + } + intermediate.push(div_round(raw, scale)); + } + + let mut delta = Vec::with_capacity(out_dim); + 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 < -intermediate_bound || raw > intermediate_bound { + return Err(NativeError::InvalidDimensions(format!( + "intermediate value {raw} exceeds signed bound +/-{intermediate_bound}" + ))); + } + let rescaled = div_round(raw, scale); + let scaled = rescaled + .checked_mul(scaling_num as i128) + .ok_or_else(overflow)?; + let out = div_round(scaled, scaling_den as i128); + if out < -value_bound || out > value_bound { + return Err(NativeError::InvalidDimensions(format!( + "delta value {out} exceeds signed bound +/-{value_bound}" + ))); + } + delta.push(out as i64); + } + Ok(delta) +} + +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 + .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]> = Vec::with_capacity(level.len() / 2 + 1); + for pair in level.chunks_exact(2) { + let mut hasher = blake3::Hasher::new(); + hasher.update(&pair[0]); + hasher.update(&pair[1]); + next.push(*hasher.finalize().as_bytes()); + } + 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 +1667,46 @@ 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] +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 +1719,132 @@ 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!(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::*; @@ -1462,6 +1980,126 @@ 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 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/lora_contributor_mpi/__init__.py b/src/zklora/lora_contributor_mpi/__init__.py index cdf1032..6b726e3 100644 --- a/src/zklora/lora_contributor_mpi/__init__.py +++ b/src/zklora/lora_contributor_mpi/__init__.py @@ -77,29 +77,51 @@ 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. + """ + with self._module_cache_lock: + cached = self._module_cache.get(sub_name) + if cached is not None: + return cached + module = self.submodules[sub_name] + a_matrix, b_matrix, scaling_num, scaling_den = lora_matrices_and_scaling(module) + artifacts = { + "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.setdefault(sub_name, artifacts) + return self._module_cache[sub_name] + + 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 +138,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 = [] @@ -129,12 +160,6 @@ def apply_lora( 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 - ) - b_quantized = quantize_nested( - b_matrix.detach().cpu().numpy().tolist(), self.fixed_point - ) for x_row in x_rows: q_x = flatten( quantize_nested(x_row.cpu().numpy().tolist(), self.fixed_point) @@ -219,6 +244,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 +263,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 +292,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..f1bb28e 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] diff --git a/src/zklora/proof_contract.py b/src/zklora/proof_contract.py index c927043..ca751bd 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" @@ -192,6 +195,37 @@ 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_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 +295,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 +311,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 +565,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 +796,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 +813,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..9d1e374 --- /dev/null +++ b/tests/test_native_fast_paths.py @@ -0,0 +1,109 @@ +"""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, +) + +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_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) From 316e0aa2d508a49158ab3b90d93706fd567f7feb Mon Sep 17 00:00:00 2001 From: Bidhan Roy Date: Wed, 10 Jun 2026 05:12:57 +0000 Subject: [PATCH 2/8] Accelerate inference-path quantization and delta computation 17-30x - 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. --- src/Cargo.toml | 1 + src/src/lib.rs | 328 +++++++++++++++++--- src/zklora/base_model_user_mpi/__init__.py | 12 +- src/zklora/lora_contributor_mpi/__init__.py | 28 +- src/zklora/proof_contract.py | 66 ++++ tests/test_native_fast_paths.py | 64 ++++ 6 files changed, 426 insertions(+), 73 deletions(-) diff --git a/src/Cargo.toml b/src/Cargo.toml index 3a9d3a2..3a31ee4 100644 --- a/src/Cargo.toml +++ b/src/Cargo.toml @@ -22,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/src/lib.rs b/src/src/lib.rs index 6ed2e02..5660989 100644 --- a/src/src/lib.rs +++ b/src/src/lib.rs @@ -20,6 +20,7 @@ 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}; @@ -1474,21 +1475,25 @@ pub fn statement_digest_hex(statement_json: &str) -> String { format!("{:x}", hasher.finalize()) } -/// 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. -pub fn compute_delta_quantized_native( +/// 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], - x: &[i64], scaling_num: i64, scaling_den: i64, scale_bits: u32, value_bits: u32, intermediate_bits: u32, -) -> Result, NativeError> { +) -> Result { if scaling_den <= 0 { return Err(NativeError::InvalidDimensions( "scaling_den must be positive".into(), @@ -1511,13 +1516,6 @@ pub fn compute_delta_quantized_native( )); } let in_dim = a[0].len(); - let out_dim = b.len(); - if x.len() != in_dim { - return Err(NativeError::InvalidDimensions(format!( - "x expected length {in_dim}, got {}", - x.len() - ))); - } for row in a { if row.len() != in_dim { return Err(NativeError::InvalidDimensions( @@ -1533,25 +1531,47 @@ pub fn compute_delta_quantized_native( } } let value_bound = (1i128 << (value_bits - 1)) - 1; - let intermediate_bound = (1i128 << (intermediate_bits - 1)) - 1; - let scale = 1i128 << scale_bits; - let check_value = |value: i64, label: &str| -> Result<(), NativeError> { - let value = value as i128; - if value < -value_bound || value > value_bound { - return Err(NativeError::InvalidDimensions(format!( - "{label} value {value} exceeds signed bound +/-{value_bound}" - ))); - } - Ok(()) - }; - for value in x { - check_value(*value, "x")?; - } for value in a.iter().flatten() { - check_value(*value, "A")?; + check_native_bound(*value, value_bound, "A")?; } for value in b.iter().flatten() { - check_value(*value, "B")?; + 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 = @@ -1562,7 +1582,7 @@ pub fn compute_delta_quantized_native( n.div_euclid(denominator) }; - let mut intermediate = Vec::with_capacity(rank); + 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()) { @@ -1571,34 +1591,37 @@ pub fn compute_delta_quantized_native( .ok_or_else(overflow)?; raw = raw.checked_add(product).ok_or_else(overflow)?; } - if raw < -intermediate_bound || raw > intermediate_bound { + if raw < -ctx.intermediate_bound || raw > ctx.intermediate_bound { return Err(NativeError::InvalidDimensions(format!( - "intermediate value {raw} exceeds signed bound +/-{intermediate_bound}" + "intermediate value {raw} exceeds signed bound +/-{}", + ctx.intermediate_bound ))); } - intermediate.push(div_round(raw, scale)); + intermediate.push(div_round(raw, ctx.scale)); } - let mut delta = Vec::with_capacity(out_dim); + 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 < -intermediate_bound || raw > intermediate_bound { + if raw < -ctx.intermediate_bound || raw > ctx.intermediate_bound { return Err(NativeError::InvalidDimensions(format!( - "intermediate value {raw} exceeds signed bound +/-{intermediate_bound}" + "intermediate value {raw} exceeds signed bound +/-{}", + ctx.intermediate_bound ))); } - let rescaled = div_round(raw, scale); + let rescaled = div_round(raw, ctx.scale); let scaled = rescaled - .checked_mul(scaling_num as i128) + .checked_mul(ctx.scaling_num as i128) .ok_or_else(overflow)?; - let out = div_round(scaled, scaling_den as i128); - if out < -value_bound || out > value_bound { + 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 +/-{value_bound}" + "delta value {out} exceeds signed bound +/-{}", + ctx.value_bound ))); } delta.push(out as i64); @@ -1606,6 +1629,166 @@ pub fn compute_delta_quantized_native( 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 @@ -1616,7 +1799,7 @@ pub fn merkle_root_f64(values: &[f64], nonce: &[u8]) -> [u8; 32] { return MERKLE_EMPTY; } let mut level: Vec<[u8; 32]> = values - .iter() + .par_iter() .map(|value| { let mut hasher = blake3::Hasher::new(); hasher.update(&value.to_be_bytes()); @@ -1628,13 +1811,15 @@ pub fn merkle_root_f64(values: &[f64], nonce: &[u8]) -> [u8; 32] { level.push(MERKLE_EMPTY); } while level.len() > 1 { - let mut next: Vec<[u8; 32]> = Vec::with_capacity(level.len() / 2 + 1); - for pair in level.chunks_exact(2) { - let mut hasher = blake3::Hasher::new(); - hasher.update(&pair[0]); - hasher.update(&pair[1]); - next.push(*hasher.finalize().as_bytes()); - } + 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); } @@ -1703,6 +1888,45 @@ fn compute_delta_quantized( }) } +#[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> { @@ -1720,6 +1944,8 @@ fn _native_prover(m: &pyo3::Bound<'_, pyo3::types::PyModule>) -> pyo3::PyResult< 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(()) } 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 6b726e3..caccfdd 100644 --- a/src/zklora/lora_contributor_mpi/__init__.py +++ b/src/zklora/lora_contributor_mpi/__init__.py @@ -14,9 +14,9 @@ FixedPointConfig, InvocationWitness, adapter_manifest_entry, - compute_delta_quantized, - flatten, + compute_delta_quantized_rows, quantize_nested, + quantize_rows, write_adapter_manifest, ) @@ -159,19 +159,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])) + q_x_rows = quantize_rows( + [x_row.cpu().numpy().tolist() for x_row in x_rows], self.fixed_point + ) + q_delta_rows = compute_delta_quantized_rows( + a_quantized, + b_quantized, + q_x_rows, + scaling_num, + scaling_den, + self.fixed_point, + ) delta_rows = [] - 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, - ) + 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 diff --git a/src/zklora/proof_contract.py b/src/zklora/proof_contract.py index ca751bd..5fe26ec 100644 --- a/src/zklora/proof_contract.py +++ b/src/zklora/proof_contract.py @@ -146,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] = [] @@ -219,6 +246,45 @@ def compute_delta_quantized( 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]], diff --git a/tests/test_native_fast_paths.py b/tests/test_native_fast_paths.py index 9d1e374..79ef32c 100644 --- a/tests/test_native_fast_paths.py +++ b/tests/test_native_fast_paths.py @@ -15,6 +15,9 @@ 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": @@ -68,6 +71,67 @@ def test_native_delta_bound_violations_raise_contract_errors(): 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 From b3f5c04e562e2b7ea4dfe08c4ac2956c54dc3162 Mon Sep 17 00:00:00 2001 From: Bidhan Roy Date: Wed, 10 Jun 2026 05:17:31 +0000 Subject: [PATCH 3/8] Harden range-check documentation and boundary tests 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. --- readme.md | 5 +++- src/src/lib.rs | 66 +++++++++++++++++++++++++++++++++++++++++++++++--- 2 files changed, 66 insertions(+), 5 deletions(-) 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/src/lib.rs b/src/src/lib.rs index 5660989..2e7323b 100644 --- a/src/src/lib.rs +++ b/src/src/lib.rs @@ -1003,13 +1003,27 @@ fn range_check_signed_interval( range_check_unsigned(region, config, window, &diff_cell, &diff, bits, offset) } -/// Constrain `value_cell` to `[0, 2^bits)` using a lookup-backed running sum. +/// 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 so the decomposition -/// covers exactly `bits` bits. `z_n` is constrained to zero, so the chain -/// telescopes to `z_0 = value < 2^bits` with every limb non-negative. +/// `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, @@ -2157,6 +2171,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(); From b42c3a425f80e26a433b7d0d3d012d9eaee96620 Mon Sep 17 00:00:00 2001 From: Bidhan Roy Date: Wed, 10 Jun 2026 05:22:42 +0000 Subject: [PATCH 4/8] Add v3 benchmark results document and prover memory guidance 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. --- benchmarks/native_v3_results.md | 53 +++++++++++++++++++++++++++++++++ src/README.md | 2 ++ 2 files changed, 55 insertions(+) create mode 100644 benchmarks/native_v3_results.md diff --git a/benchmarks/native_v3_results.md b/benchmarks/native_v3_results.md new file mode 100644 index 0000000..e4bdefc --- /dev/null +++ b/benchmarks/native_v3_results.md @@ -0,0 +1,53 @@ +# 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×4×768 | ~25 | infeasible | — | 20 | see note | see note | see note | n/a (was infeasible) | +| 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): ~7–8 GB peak at k=20, >15 GB at k=21. +The proving-key cache holds up to `ZKLORA_PK_CACHE_CAP` (default 2) shapes. diff --git a/src/README.md b/src/README.md index b31f8c7..aeca1c0 100644 --- a/src/README.md +++ b/src/README.md @@ -67,6 +67,8 @@ Proofs and verifying keys from the v2 backend are not compatible with v3 (the `b 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). Expect roughly 7–8 GB peak at k=20 (e.g. 768×4×768) and >15 GB at k=21 (768×4×2304); size the prover host accordingly. 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. ## Core Components From 8a5847f50ce781a5ea163ce2db9252e5048cb824 Mon Sep 17 00:00:00 2001 From: Bidhan Roy Date: Wed, 10 Jun 2026 05:35:13 +0000 Subject: [PATCH 5/8] Correct prover memory guidance to measured values 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. --- benchmarks/native_v3_results.md | 7 +++++-- src/README.md | 2 +- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/benchmarks/native_v3_results.md b/benchmarks/native_v3_results.md index e4bdefc..41d4333 100644 --- a/benchmarks/native_v3_results.md +++ b/benchmarks/native_v3_results.md @@ -18,7 +18,8 @@ and `python benchmarks/run_benchmarks.py`. | 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×4×768 | ~25 | infeasible | — | 20 | see note | see note | see note | n/a (was infeasible) | +| 768×2×256 | ~24 | infeasible | — | 19 | see below | see below | see below | 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 @@ -49,5 +50,7 @@ four fixed-point configs). ## Memory Proving memory is dominated by halo2 extended-domain evaluations (Poseidon -gate degree ⇒ 8× extended domain): ~7–8 GB peak at k=20, >15 GB at k=21. +gate degree ⇒ 8× extended domain), and halo2 0.3 keeps all advice cosets +resident during create_proof. Measured: k=20 and k=21 shapes exceed a 15 GB +host (OOM); 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/src/README.md b/src/README.md index aeca1c0..2e3f3ac 100644 --- a/src/README.md +++ b/src/README.md @@ -67,7 +67,7 @@ Proofs and verifying keys from the v2 backend are not compatible with v3 (the `b 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). Expect roughly 7–8 GB peak at k=20 (e.g. 768×4×768) and >15 GB at k=21 (768×4×2304); size the prover host accordingly. Verification is lightweight once the verifying key is cached. +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: k≤19 circuits prove comfortably, 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. From d8ff86ac990196c8e68b5b5d110b75ec2d1f61f6 Mon Sep 17 00:00:00 2001 From: Bidhan Roy Date: Wed, 10 Jun 2026 05:43:58 +0000 Subject: [PATCH 6/8] Finalize v3 benchmark results with measured k=19 large-shape numbers 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. --- benchmarks/native_v3_results.md | 7 ++++--- src/README.md | 2 +- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/benchmarks/native_v3_results.md b/benchmarks/native_v3_results.md index 41d4333..a4e22ba 100644 --- a/benchmarks/native_v3_results.md +++ b/benchmarks/native_v3_results.md @@ -18,7 +18,7 @@ and `python benchmarks/run_benchmarks.py`. | 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 | ~24 | infeasible | — | 19 | see below | see below | see below | 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 | @@ -51,6 +51,7 @@ four fixed-point configs). 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: k=20 and k=21 shapes exceed a 15 GB -host (OOM); plan on 32–64 GB for 768-dim × 4-rank modules with wide outputs. +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/src/README.md b/src/README.md index 2e3f3ac..cfc7903 100644 --- a/src/README.md +++ b/src/README.md @@ -67,7 +67,7 @@ Proofs and verifying keys from the v2 backend are not compatible with v3 (the `b 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: k≤19 circuits prove comfortably, 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. +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. From bdf5957669cd0fedbc453e5e2f85471fb596657f Mon Sep 17 00:00:00 2001 From: Bidhan Roy Date: Wed, 10 Jun 2026 12:55:47 +0000 Subject: [PATCH 7/8] Dedup activation-commitment flatten; keep JSON float parsing in Python 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. --- src/zklora/polynomial_commit.py | 87 ++++++++++++++++----------------- 1 file changed, 41 insertions(+), 46 deletions(-) diff --git a/src/zklora/polynomial_commit.py b/src/zklora/polynomial_commit.py index f1bb28e..fc0020f 100644 --- a/src/zklora/polynomial_commit.py +++ b/src/zklora/polynomial_commit.py @@ -82,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) # -------------------------------------------------------------------------------------- @@ -97,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()} @@ -159,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 From 4a364478bdf5fc8010ad91b7ff99d13296e574fe Mon Sep 17 00:00:00 2001 From: Bidhan Roy Date: Wed, 10 Jun 2026 14:31:38 +0000 Subject: [PATCH 8/8] Fix review findings: stale adapter cache, unchecked add, un-ignore real-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. --- src/src/lib.rs | 35 ++++++++++++++++----- src/zklora/lora_contributor_mpi/__init__.py | 23 ++++++++++---- 2 files changed, 45 insertions(+), 13 deletions(-) diff --git a/src/src/lib.rs b/src/src/lib.rs index 2e7323b..f90cf7d 100644 --- a/src/src/lib.rs +++ b/src/src/lib.rs @@ -1590,10 +1590,12 @@ fn delta_for_row( let overflow = || NativeError::InvalidDimensions("intermediate value exceeds native range".into()); - let div_round = |numerator: i128, denominator: i128| -> i128 { + let div_round = |numerator: i128, denominator: i128| -> Result { // denominator > 0 here; floor((n + d/2) / d), matching div_floor. - let n = numerator + denominator / 2; - n.div_euclid(denominator) + let n = numerator + .checked_add(denominator / 2) + .ok_or_else(overflow)?; + Ok(n.div_euclid(denominator)) }; let mut intermediate = Vec::with_capacity(a.len()); @@ -1611,7 +1613,7 @@ fn delta_for_row( ctx.intermediate_bound ))); } - intermediate.push(div_round(raw, ctx.scale)); + intermediate.push(div_round(raw, ctx.scale)?); } let mut delta = Vec::with_capacity(b.len()); @@ -1627,11 +1629,11 @@ fn delta_for_row( ctx.intermediate_bound ))); } - let rescaled = div_round(raw, ctx.scale); + 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); + 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 +/-{}", @@ -2232,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 { @@ -2363,6 +2364,26 @@ mod tests { 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]; diff --git a/src/zklora/lora_contributor_mpi/__init__.py b/src/zklora/lora_contributor_mpi/__init__.py index caccfdd..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 @@ -87,15 +88,25 @@ 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. + 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: + if cached is not None and cached["fingerprint"] == fingerprint: return cached - module = self.submodules[sub_name] - a_matrix, b_matrix, scaling_num, scaling_den = lora_matrices_and_scaling(module) artifacts = { + "fingerprint": fingerprint, "a_matrix": a_matrix, "b_matrix": b_matrix, "scaling_num": int(scaling_num), @@ -108,8 +119,8 @@ def _module_artifacts(self, sub_name: str) -> dict: ), } with self._module_cache_lock: - self._module_cache.setdefault(sub_name, artifacts) - return self._module_cache[sub_name] + self._module_cache[sub_name] = artifacts + return artifacts def adapter_manifest_entries(self): entries = []