jagged: group reshaped q prime commitments (#59)

hero78119 · kunxian-xia · Copilot · web-flow · commit 2222cf29bc7d · 2026-05-18T10:57:04.000+08:00
* wip jagged_commit * feat: jagged sumcheck (#32) * claude code impl plan * implement the jagged_sumcheck using time-space tradeoff sumcheck prover algorithm * remove debugging codes * ref to the original paper * add jagged sumcheck bench * #32 parallelize (#34) * par wip * check f(z) * g(z) matches expected evaluation * fix clippy * fix clippy: add #[cfg(test)] to test-only method and fix unused import/variable warnings Agent-Logs-Url: https://github.com/scroll-tech/gkr-backend/sessions/4a61316e-cf28-47b8-a43e-fb6ab432701e Co-authored-by: hero78119 <3962077+hero78119@users.noreply.github.com> * refactor test * apply functional programming style * avoid unnecessary BaseField-to-ExtensionField conversion in q_evals access Use E * BaseField multiplication directly instead of converting q_evals elements to extension field with .into() first. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * replace col_row binary search with incremental ColRowIter Add ColRowIter that does one binary search at construction and O(1) per step, replacing per-element binary searches in build_m_table, bind_and_materialize, compute_claimed_sum, and final_evaluations_slow. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: hero78119 <3962077+hero78119@users.noreply.github.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> * extend jagged sumcheck benchmark to cover n=25..31 * switch jagged sumcheck benchmark to BabyBearExt4 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * remove jagged sumcheck plan doc Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: hero78119 <3962077+hero78119@users.noreply.github.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> * feat: make streaming epoch schedule configurable in jagged_sumcheck_prove (#38) * claude code impl plan * implement the jagged_sumcheck using time-space tradeoff sumcheck prover algorithm * remove debugging codes * ref to the original paper * add jagged sumcheck bench * #32 parallelize (#34) * par wip * check f(z) * g(z) matches expected evaluation * fix clippy * fix clippy: add #[cfg(test)] to test-only method and fix unused import/variable warnings Agent-Logs-Url: https://github.com/scroll-tech/gkr-backend/sessions/4a61316e-cf28-47b8-a43e-fb6ab432701e Co-authored-by: hero78119 <3962077+hero78119@users.noreply.github.com> * refactor test * apply functional programming style * avoid unnecessary BaseField-to-ExtensionField conversion in q_evals access Use E * BaseField multiplication directly instead of converting q_evals elements to extension field with .into() first. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * replace col_row binary search with incremental ColRowIter Add ColRowIter that does one binary search at construction and O(1) per step, replacing per-element binary searches in build_m_table, bind_and_materialize, compute_claimed_sum, and final_evaluations_slow. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: hero78119 <3962077+hero78119@users.noreply.github.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> * extend jagged sumcheck benchmark to cover n=25..31 * switch jagged sumcheck benchmark to BabyBearExt4 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * remove jagged sumcheck plan doc Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * Initial plan * Make EPOCH_SIZES configurable in jagged_sumcheck_prove via optional parameter Agent-Logs-Url: https://github.com/scroll-tech/gkr-backend/sessions/b18805c7-6e15-44c0-ab03-a1905068d964 Co-authored-by: hero78119 <3962077+hero78119@users.noreply.github.com> * Fix doc spacing and add debug_assert for epoch_sizes validation Agent-Logs-Url: https://github.com/scroll-tech/gkr-backend/sessions/b18805c7-6e15-44c0-ab03-a1905068d964 Co-authored-by: hero78119 <3962077+hero78119@users.noreply.github.com> --------- Co-authored-by: kunxian xia <xiakunxian130@gmail.com> Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: hero78119 <3962077+hero78119@users.noreply.github.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> * Feat: evaluator for jagged indicator function (#39) * implement evaluate_g using width-4 ROBP from Section 3.2/4 Evaluate the MLE of indicator function g(a,b,c,d) = [a+c=b AND b<d] using a width-4 read-once branching program with state (carry, lt). Placed in separate jagged_evaluator.rs module. Includes brute-force test for correctness verification (n=1..4). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * implement forward and backward algorithms for evaluate_g Separate evaluate_g into two functions following the ROBP-based MLE evaluation: forward (source→sinks, MLE definition) and backward (sinks→source, Claim 4.2.1/Lemma 4.2 from jagged PCS paper). Extract shared transition_weights helper. Default to backward for future batch/symbolic evaluation support. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> * feat: implement jagged_batch_open and jagged_batch_verify Add the opening protocol for the jagged PCS: given K evaluation claims on individual column polynomials, prove they're consistent with the commitment to the giga polynomial q'. Uses the jagged sumcheck to reduce to a single evaluation of q', then opens via the inner PCS. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * update doc comments * polish doc comments * polish doc comments * reorg * move types to a separate file * fix: add missing rayon trait imports for --all-features clippy Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * update doc comments: document correction factors for different-height polynomials The zero-padding interpretation (p_i^pad) explains why C_i = eq(z_r[s_i..], 0) is needed when batching polynomials of different heights with a single eq table. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * refactor: rename generic parameter Pcs to InnerPcs for readability Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * docs: document 2^{s_i} block sizes and zero-padding cost vs SP1 Clarify that each polynomial in q' occupies a full 2^{s_i} block (not the original h_i entries) because bit-reversal scatters the zero-padding throughout the block. Add note comparing with SP1's jagged PCS which uses raw h_i, highlighting the Σ(2^{s_i} - h_i) extra zeros as the cost of suffix-to-prefix bit-reversal. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * feat: assist sumcheck for jagged PCS (#42) * feat: implement assist sumcheck to reduce K ROBP evaluations to one Implements the assist sumcheck protocol (Lemma 5.1 of eprint 2025/917) which batches K indicator function evaluations into a single opening, using an interleaved variable ordering and forward-backward ROBP state decomposition. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * perf: parallelize backward precomputation and fuse round accumulations - Parallelize the O(K × n_robp) backward vector precomputation with rayon - Derive bwd_sum_d from bwd_sum in O(1) instead of a second O(K) pass - Fuse the two per-step weight updates into one - Add assist_sumcheck benchmark (K=100/500/1000) K=1000 drops from 14.1ms to 4.6ms (3.1x speedup with --features parallel). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * docs: annotate key MLE telescoping identity in assist sumcheck Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * bench: update assist sumcheck benchmark to K=1000,2000,4000,8000 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * refactor: derive ROBP transition matrices from raw transition table Separate the ROBP state machine logic from the eq-weighting algebra. The raw transition table ROBP_TRANSITION encodes the state machine directly, and symbol_transition_matrices derives M_i^{(c,d)} via the closed formula Σ_{a,b} eq₁(z1,a)·eq₁(z2,b)·[transition(from,(a,b,c,d))=to]. Also document why z₁/z₂ are bound first (reducing alphabet from {0,1}⁴ to {0,1}²) and z₃/z₄ interleaved to match ROBP step order. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * fix: suppress needless_range_loop clippy lint and apply fmt Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * bench: add end-to-end jagged PCS benchmark (commit, open, verify) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * fix: suppress clippy needless_range_loop in evaluator tests, add K=16000 bench Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * perf: transpose bwd/c_bits/d_bits to step-major layout for cache-friendly access Reorganize data structures from poly-major bwd[y][i] to step-major bwd[i][y], and similarly for c_bits and d_bits. The main loop now scans contiguous memory when accumulating over K polynomials, reducing cache misses at large K (~38% improvement at K=16000). Backward precomputation uses split_at_mut + par_iter for safe parallel writes. Also adds the round polynomial formula (§2.3, Eq. 4) as a code comment to clarify the bwd_sum accumulation logic. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * perf: parallelize assist sumcheck round-polynomial computation and weight update Thread-local bwd_sum buckets + per-thread p(0),p(1),p(2) evaluation eliminates the sequential O(K) bottleneck in the main step loop. Each thread builds local buckets, computes local round polynomial values, and only the scalars are summed across threads. Benchmark (BabyBearExt4, --features parallel, vs HEAD): K=1000: 7.8ms → 6.8ms (-14%) K=2000: 13.5ms → 10.2ms (-24%) K=4000: 25.6ms → 16.2ms (-38%) K=8000: 50.7ms → 28.0ms (-46%) K=16000: 100.5ms → 53.6ms (-48%) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * perf: eliminate bwd_sum merge by reusing per-thread buckets in round 2i+1 Instead of merging thread-local bwd_sums after round 2i, cache them and absorb alpha directly into each thread's local bwd_sum for round 2i+1. Only scalar p(0),p(1),p(2) values are summed across threads in both rounds. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * fix: gate IndexedParallelIterator import behind parallel feature, remove tautological debug_assert Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * fmt * fix: use wildcard maybe_rayon import, remove unused direct rayon dep Wildcard `use p3::maybe_rayon::prelude::*` avoids unused-import warnings caused by Cargo feature unification: `p3/parallel` can be activated by other workspace crates even when `mpcs/parallel` is off, making a `#[cfg(feature = "parallel")]` gate on IndexedParallelIterator incorrect. Also remove `dep:rayon` from mpcs — all parallelism goes through `p3::maybe_rayon`, so the direct rayon dependency was unnecessary. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> * docs: label batched opening claim and document correction factor inclusion in prover Agent-Logs-Url: https://github.com/scroll-tech/gkr-backend/sessions/22b3ba65-8e69-45eb-ae09-0a5267f8b518 Co-authored-by: kunxian-xia <1082586+kunxian-xia@users.noreply.github.com> * feat: add reshape step to jagged PCS (#48) * feat: make JaggedSumcheckInput::compute_claimed_sum public * feat: add jagged vs direct PCS comparison bench and use p3 RowMajorMatrix - Add comparison benchmark (comparison.rs) measuring commit, batch_open, batch_verify, and proof size for jagged PCS vs direct inner PCS - Refactor jagged_commit to accept p3::matrix::dense::RowMajorMatrix instead of witness::RowMajorMatrix, supporting non-power-of-two matrix heights (internally padded to next power of two before bit-reversal) - Update jagged_pcs bench to use BabyBearExt4, parallel make_rmm, jittered non-power-of-two heights, and eq-table-based column evaluation - Cap reshape_log_height to 25 to fit BabyBear two-adicity constraint Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * misc: prefix align without bit reverse * jagged: direct concat for prefix-aligned points (#56) misc: prefix align without bit reverse * WIP integrate jagged pcs cpu * Fix jagged PCS padded opening normalization * misc: clippy fix * avoid reconstruct q_mles * more docs * misc: fmt and clippy * feat: group jagged q prime commitments * retain fast path for one column * fix label * misc: function scope --------- Co-authored-by: kunxian xia <xiakunxian130@gmail.com> Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: hero78119 <3962077+hero78119@users.noreply.github.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: kunxian-xia <1082586+kunxian-xia@users.noreply.github.com>
diff --git a/crates/mpcs/src/jagged/mod.rs b/crates/mpcs/src/jagged/mod.rs
@@ -128,7 +128,7 @@ use multilinear_extensions::{
     util::ceil_log2,
     virtual_poly::{VPAuxInfo, build_eq_x_r_vec},
 };
-use p3::{field::FieldAlgebra, matrix::Matrix, maybe_rayon::prelude::*};
+use p3::{field::FieldAlgebra, maybe_rayon::prelude::*};
 use serde::{Serialize, Serializer, de::DeserializeOwned};
 use std::sync::Arc;
 use transcript::Transcript;
@@ -138,6 +138,43 @@ use witness::{InstancePaddingStrategy, RowMajorMatrix as WitnessRowMajorMatrix};
 #[derive(Debug)]
 pub struct Jagged<InnerPcs>(PhantomData<InnerPcs>);
 
+pub const JAGGED_RESHAPE_GROUP_WIDTH: usize = 8;
+
+type InnerOpenings<E> = Vec<(Point<E>, Vec<E>)>;
+type InnerVerifyOpenings<E> = Vec<(usize, (Point<E>, Vec<E>))>;
+
+fn reshape_group_width() -> usize {
+    JAGGED_RESHAPE_GROUP_WIDTH
+}
+
+fn chunk_col_evals<E: ExtensionField>(col_evals: &[E]) -> Vec<Vec<E>> {
+    col_evals
+        .chunks(reshape_group_width())
+        .map(|chunk| chunk.to_vec())
+        .collect()
+}
+
+fn inner_openings_for_col_evals<E: ExtensionField>(
+    rho_row: &[E],
+    col_evals: &[E],
+) -> InnerOpenings<E> {
+    chunk_col_evals(col_evals)
+        .into_iter()
+        .map(|evals| (rho_row.to_vec(), evals))
+        .collect()
+}
+
+fn inner_verify_openings_for_col_evals<E: ExtensionField>(
+    log_h: usize,
+    rho_row: &[E],
+    col_evals: &[E],
+) -> InnerVerifyOpenings<E> {
+    chunk_col_evals(col_evals)
+        .into_iter()
+        .map(|evals| (log_h, (rho_row.to_vec(), evals)))
+        .collect()
+}
+
 impl<InnerPcs> Clone for Jagged<InnerPcs> {
     fn clone(&self) -> Self {
         Self(PhantomData)
@@ -265,46 +302,50 @@ pub fn jagged_commit<E: ExtensionField, InnerPcs: PolynomialCommitmentScheme<E>>
     let h = 1usize << log_h;
     let w = total_size.div_ceil(h);
 
-    let giga_data = if w == 1 {
-        concatenated
+    let giga_rmms = if w == 1 {
+        vec![WitnessRowMajorMatrix::<E::BaseField>::new_by_values(
+            concatenated,
+            1,
+            InstancePaddingStrategy::Default,
+        )]
     } else {
-        // Transpose the flat q' evaluations to row-major for the inner PCS.
-        // In the flat array, column i occupies concatenated[i*h .. (i+1)*h];
-        // the last column may be short and is zero-padded.
-        let padded_total = w * h;
-        let mut row_major: Vec<E::BaseField> = Vec::with_capacity(padded_total);
-        #[allow(clippy::uninit_vec)]
-        unsafe {
-            row_major.set_len(padded_total)
-        };
-
-        (0..w).into_par_iter().for_each(|i| {
-            let src_start = i * h;
-            let col_len = if i < w - 1 { h } else { total_size - src_start };
-            let dst = unsafe {
-                &mut *std::ptr::slice_from_raw_parts_mut(
-                    row_major.as_ptr() as *mut E::BaseField,
-                    padded_total,
-                )
+        let mut giga_rmms = Vec::with_capacity(w.div_ceil(reshape_group_width()));
+        for group_start_col in (0..w).step_by(reshape_group_width()) {
+            let group_cols = (w - group_start_col).min(reshape_group_width());
+            let mut row_major: Vec<E::BaseField> = Vec::with_capacity(h * group_cols);
+            #[allow(clippy::uninit_vec)]
+            unsafe {
+                row_major.set_len(h * group_cols)
             };
-            for b in 0..col_len {
-                dst[b * w + i] = concatenated[src_start + b];
-            }
-            for b in col_len..h {
-                dst[b * w + i] = E::BaseField::ZERO;
-            }
-        });
 
-        row_major
-    };
+            (0..group_cols).into_par_iter().for_each(|local_col| {
+                let global_col = group_start_col + local_col;
+                let src_start = global_col * h;
+                let col_len = total_size.saturating_sub(src_start).min(h);
+                let dst = unsafe {
+                    &mut *std::ptr::slice_from_raw_parts_mut(
+                        row_major.as_ptr() as *mut E::BaseField,
+                        h * group_cols,
+                    )
+                };
+                for b in 0..col_len {
+                    dst[b * group_cols + local_col] = concatenated[src_start + b];
+                }
+                for b in col_len..h {
+                    dst[b * group_cols + local_col] = E::BaseField::ZERO;
+                }
+            });
 
-    let giga_rmm = WitnessRowMajorMatrix::<E::BaseField>::new_by_values(
-        giga_data,
-        w,
-        InstancePaddingStrategy::Default,
-    );
+            giga_rmms.push(WitnessRowMajorMatrix::<E::BaseField>::new_by_values(
+                row_major,
+                group_cols,
+                InstancePaddingStrategy::Default,
+            ));
+        }
+        giga_rmms
+    };
 
-    let inner = InnerPcs::batch_commit(pp, vec![giga_rmm])?;
+    let inner = InnerPcs::batch_commit(pp, giga_rmms)?;
 
     Ok(JaggedCommitmentWithWitness {
         inner,
@@ -319,15 +360,16 @@ fn default_reshape_log_height<E: ExtensionField, InnerPcs: PolynomialCommitmentS
     pp: &InnerPcs::ProverParam,
     rmms: &[WitnessRowMajorMatrix<E::BaseField>],
 ) -> usize {
-    let total_evals = rmms
+    let max_poly_height = rmms
         .iter()
-        .map(|rmm| rmm.occupied_physical_rows() * rmm.width())
-        .sum::<usize>();
-    let total_log = ceil_log2(total_evals.max(1));
-    pp.get_max_message_size_log().min(total_log)
+        .map(|rmm| rmm.occupied_physical_rows())
+        .max()
+        .unwrap_or(1);
+    pp.get_max_message_size_log()
+        .min(ceil_log2(max_poly_height.max(1)))
 }
 
-fn flatten_padded_openings_as_native<E: ExtensionField>(
+pub fn flatten_padded_openings_as_native<E: ExtensionField>(
     poly_heights: &[usize],
     openings: Vec<(Point<E>, Vec<E>)>,
 ) -> Result<(Point<E>, Vec<E>), Error> {
@@ -478,7 +520,7 @@ pub fn jagged_batch_open<E: ExtensionField, InnerPcs: PolynomialCommitmentScheme
     // Write evals to transcript, then sample z_col.
     transcript.append_field_element_exts(evals);
     let num_col_vars = ceil_log2(num_polys).max(1);
-    let z_col: Vec<E> = transcript.sample_and_append_vec(b"jagged_z_col", num_col_vars);
+    let z_col: Vec<E> = transcript.sample_and_append_vec(b"z_col", num_col_vars);
 
     let eq_col = build_eq_x_r_vec(&z_col);
 
@@ -554,7 +596,10 @@ pub fn jagged_batch_open<E: ExtensionField, InnerPcs: PolynomialCommitmentScheme
     // Open column MLEs at ρ_row via inner PCS.
     let inner_proof = InnerPcs::batch_open(
         pp,
-        vec![(&comm.inner, vec![(rho_row.to_vec(), col_evals.clone())])],
+        vec![(
+            &comm.inner,
+            inner_openings_for_col_evals(rho_row, &col_evals),
+        )],
         transcript,
     )?;
 
@@ -601,7 +646,7 @@ pub fn jagged_batch_verify<E: ExtensionField, InnerPcs: PolynomialCommitmentSche
     // Replay transcript: write evals, sample z_col.
     transcript.append_field_element_exts(evals);
     let num_col_vars = ceil_log2(num_polys).max(1);
-    let z_col: Vec<E> = transcript.sample_and_append_vec(b"jagged_z_col", num_col_vars);
+    let z_col: Vec<E> = transcript.sample_and_append_vec(b"z_col", num_col_vars);
 
     // Batched opening claim: v = Σ_i eq_col[i] · C_i · evals[i]
     // where C_i = eq(z_row[s_i..], 0) = Π_{j=s_i}^{max_s-1} (1 - z_row[j]) is the
@@ -694,7 +739,7 @@ pub fn jagged_batch_verify<E: ExtensionField, InnerPcs: PolynomialCommitmentSche
         vp,
         vec![(
             comm.inner.clone(),
-            vec![(log_h, (rho_row.to_vec(), proof.col_evals.clone()))],
+            inner_verify_openings_for_col_evals(log_h, rho_row, &proof.col_evals),
         )],
         &proof.inner_proof,
         transcript,
diff --git a/crates/mpcs/src/lib.rs b/crates/mpcs/src/lib.rs
@@ -273,10 +273,10 @@ pub use basefold::{
 };
 pub mod jagged;
 pub use jagged::{
-    Jagged, JaggedBatchOpenProof, JaggedCommitment, JaggedCommitmentWithWitness, JaggedProof,
-    JaggedSumcheckInput, assist_sumcheck_prove, evaluate_g, evaluate_g_backward,
-    evaluate_g_forward, jagged_batch_open, jagged_batch_verify, jagged_commit,
-    jagged_sumcheck_prove,
+    JAGGED_RESHAPE_GROUP_WIDTH, Jagged, JaggedBatchOpenProof, JaggedCommitment,
+    JaggedCommitmentWithWitness, JaggedProof, JaggedSumcheckInput, assist_sumcheck_prove,
+    evaluate_g, evaluate_g_backward, evaluate_g_forward, jagged_batch_open, jagged_batch_verify,
+    jagged_commit, jagged_sumcheck_prove,
 };
 #[cfg(feature = "whir")]
 extern crate whir as whir_external;
diff --git a/crates/whir/src/utils.rs b/crates/whir/src/utils.rs
@@ -104,7 +104,7 @@ pub fn stack_evaluations_mut_rmm<F: Field>(
     folding_factor: usize,
 ) {
     let folding_factor_exp = 1 << folding_factor;
-    assert!(rmm.height() % folding_factor_exp == 0);
+    assert!(rmm.height().is_multiple_of(folding_factor_exp));
     let size_of_new_domain = rmm.height() / folding_factor_exp;
 
     // interpret evals as (folding_factor_exp x size_of_new_domain)-matrix and transpose in-place
diff --git a/crates/whir/src/whir/batch/committer.rs b/crates/whir/src/whir/batch/committer.rs
@@ -11,11 +11,7 @@ use crate::{
 };
 use derive_more::Debug;
 use ff_ext::ExtensionField;
-use p3::{
-    matrix::{Matrix, dense::RowMajorMatrix},
-    maybe_rayon::prelude::*,
-    util::log2_strict_usize,
-};
+use p3::{matrix::dense::RowMajorMatrix, maybe_rayon::prelude::*, util::log2_strict_usize};
 use sumcheck::macros::{entered_span, exit_span};
 use transcript::{BasicTranscript, Transcript};
 
diff --git a/crates/whir/src/whir/fold.rs b/crates/whir/src/whir/fold.rs
@@ -5,7 +5,6 @@ use crate::{
 
 use p3::{
     field::{Field, TwoAdicField},
-    matrix::Matrix,
     maybe_rayon::prelude::*,
 };
 
diff --git a/crates/witness/src/lib.rs b/crates/witness/src/lib.rs
@@ -41,6 +41,7 @@ pub struct RowMajorMatrix<T: Sized + Sync + Clone + Send + Copy> {
     log2_num_rotation: usize,
     is_padded: bool,
     padding_strategy: InstancePaddingStrategy,
+    host_elided_padded_height: Option<usize>,
     // Optional opaque handle to device-resident storage that mirrors `inner.values`.
     // This lets GPU-side code keep an associated buffer/layout without forcing witness
     // to depend on a concrete device runtime. There is no automatic host<->device sync:
@@ -81,6 +82,7 @@ impl<T: Sized + Sync + Clone + Send + Copy + Default + FieldAlgebra> RowMajorMat
             is_padded: true,
             log2_num_rotation: 0,
             padding_strategy: InstancePaddingStrategy::Default,
+            host_elided_padded_height: None,
             device_backing: None,
         }
     }
@@ -91,6 +93,7 @@ impl<T: Sized + Sync + Clone + Send + Copy + Default + FieldAlgebra> RowMajorMat
             log2_num_rotation: 0,
             is_padded: true,
             padding_strategy: InstancePaddingStrategy::Default,
+            host_elided_padded_height: None,
             device_backing: None,
         }
     }
@@ -121,7 +124,16 @@ impl<T: Sized + Sync + Clone + Send + Copy + Default + FieldAlgebra> RowMajorMat
     }
 
     pub fn num_vars(&self) -> usize {
-        self.inner.height().ilog2() as usize
+        self.height().ilog2() as usize
+    }
+
+    pub fn height(&self) -> usize {
+        self.host_elided_padded_height
+            .unwrap_or_else(|| self.inner.height())
+    }
+
+    pub fn width(&self) -> usize {
+        self.inner.width
     }
 
     pub fn new(
@@ -157,10 +169,38 @@ impl<T: Sized + Sync + Clone + Send + Copy + Default + FieldAlgebra> RowMajorMat
             log2_num_rotation,
             is_padded: matches!(padding_strategy, InstancePaddingStrategy::Default),
             padding_strategy,
+            host_elided_padded_height: None,
             device_backing: None,
         }
     }
 
+    /// Create matrix metadata backed only by device storage.
+    ///
+    /// This is intended for GPU proving paths that already own a device-resident
+    /// matrix and need to reuse `RowMajorMatrix` shape/device metadata without
+    /// allocating a duplicate host-side zero matrix. Callers must not use host
+    /// value accessors unless they later materialize host values separately.
+    pub fn new_by_device_backing<D: Any + Send + Sync + 'static>(
+        num_rows: usize,
+        num_cols: usize,
+        padding_strategy: InstancePaddingStrategy,
+        storage: D,
+        layout: DeviceMatrixLayout,
+    ) -> Self {
+        let num_row_padded = next_pow2_instance_padding(num_rows);
+        let mut matrix = RowMajorMatrix {
+            inner: p3::matrix::dense::RowMajorMatrix::new(vec![], num_cols),
+            num_rows,
+            log2_num_rotation: 0,
+            is_padded: matches!(padding_strategy, InstancePaddingStrategy::Default),
+            padding_strategy,
+            host_elided_padded_height: Some(num_row_padded),
+            device_backing: None,
+        };
+        matrix.set_device_backing(storage, layout);
+        matrix
+    }
+
     pub fn new_by_inner_matrix(
         mut m: p3::matrix::dense::RowMajorMatrix<T>,
         padding_strategy: InstancePaddingStrategy,
@@ -176,6 +216,7 @@ impl<T: Sized + Sync + Clone + Send + Copy + Default + FieldAlgebra> RowMajorMat
             log2_num_rotation: 0,
             is_padded: matches!(padding_strategy, InstancePaddingStrategy::Default),
             padding_strategy,
+            host_elided_padded_height: None,
             device_backing: None,
         }
     }
@@ -294,10 +335,18 @@ impl<T: Sized + Sync + Clone + Send + Copy + Default + FieldAlgebra> RowMajorMat
     }
 
     pub fn values(&self) -> &[T] {
+        assert!(
+            self.host_elided_padded_height.is_none(),
+            "host values were elided for this device-backed RowMajorMatrix"
+        );
         &self.inner.values
     }
 
     pub fn pad_to_height(&mut self, new_height: usize, fill: T) {
+        assert!(
+            self.host_elided_padded_height.is_none(),
+            "cannot pad RowMajorMatrix with elided host values"
+        );
         let (cur_height, n_cols) = (self.height(), self.n_col());
         assert!(new_height >= cur_height);
         self.invalidate_device_backing();
@@ -317,6 +366,7 @@ impl<T: Sized + Sync + Clone + Send + Copy> Clone for RowMajorMatrix<T> {
             log2_num_rotation: self.log2_num_rotation,
             is_padded: self.is_padded,
             padding_strategy: self.padding_strategy.clone(),
+            host_elided_padded_height: self.host_elided_padded_height,
             device_backing: None,
         }
     }
