feat(prover): spill the aux LDE under StorageMode::Disk

crypto/stark/src/gpu_lde.rs

@@ -1553,7 +1553,7 @@ where
 
     // Aux: de-interleaved ext3 slabs -> row-major interleaved host Vec.
     let aux_data: Vec<FieldElement<E>> =
-        if lde_trace.num_aux_cols() == 0 || !lde_trace.aux_data.is_empty() {
+        if lde_trace.num_aux_cols() == 0 || !lde_trace.aux_data.row_major_data().is_empty() {
             Vec::new()
         } else {
             let Some(h) = lde_trace.gpu_aux() else {

crypto/stark/src/prover.rs

@@ -275,8 +275,10 @@ where
 struct Lde<Field: IsFFTField, FieldExtension: IsField> {
     /// Row-major main LDE buffer + its column count.
     main: (Vec<FieldElement<Field>>, usize),
-    /// Row-major aux LDE buffer + its column count (`(vec![], 0)` if no aux).
-    aux: (Vec<FieldElement<FieldExtension>>, usize),
+    /// Row-major aux LDE, as a `Table` whose `width` is the column count (an
+    /// empty, zero-width table if no aux). `Table` rather than a bare `Vec` so
+    /// the buffer can be disk-spilled for the whole rounds 2-4 window.
+    aux: Table<FieldExtension>,
     /// Device-side main LDE buffer, populated only when the R1 GPU fused
     /// pipeline ran for this table. Kept so R2/R3/R4 GPU paths can read
     /// the LDE without re-H2D.
@@ -302,7 +304,9 @@ where
         blowup_factor: usize,
     ) -> Round1<Field, FieldExtension> {
         let (main_data, num_main_cols) = lde.main;
-        let (aux_data, num_aux_cols) = lde.aux;
+        let aux_data = lde.aux;
+        #[cfg(feature = "cuda")]
+        let num_aux_cols = aux_data.width;
 
         // Stage-3 device-only detection, inferred from the ACTUAL buffer state
         // (not the gate's intent): a table whose round-1 D2H was skipped has an
@@ -322,8 +326,8 @@ where
         #[cfg(feature = "cuda")]
         let main_empty = num_main_cols > 0 && main_data.is_empty();
         #[cfg(feature = "cuda")]
-        let host_trace_empty =
-            main_empty || (num_aux_cols > 0 && aux_data.is_empty() && lde.gpu_aux.is_some());
+        let host_trace_empty = main_empty
+            || (num_aux_cols > 0 && aux_data.row_major_data().is_empty() && lde.gpu_aux.is_some());
         #[cfg(feature = "cuda")]
         let device_num_rows = lde
             .gpu_main
@@ -336,7 +340,6 @@ where
             main_data,
             num_main_cols,
             aux_data,
-            num_aux_cols,
             step_size,
             blowup_factor,
         );
@@ -1407,9 +1410,9 @@ pub trait IsStarkProver<
                     }
                 }
             }
-            (aux_data, num_aux_cols)
+            Table::from_row_major(aux_data, num_aux_cols)
         } else {
-            (Vec::new(), 0)
+            Table::new(Vec::new(), 0)
         };
 
         Ok(commitment.build_round1(
@@ -3323,11 +3326,11 @@ pub trait IsStarkProver<
         #[cfg(feature = "cuda")]
         type AuxResult<FE> = (
             Option<TableCommit<FE>>,
-            (Vec<FieldElement<FE>>, usize),
+            Table<FE>,
             Option<math_cuda::lde::GpuLdeExt3>,
         );
         #[cfg(not(feature = "cuda"))]
-        type AuxResult<FE> = (Option<TableCommit<FE>>, (Vec<FieldElement<FE>>, usize));
+        type AuxResult<FE> = (Option<TableCommit<FE>>, Table<FE>);
         // R1 aux commit and rounds 2 to 4 share the peak working set: the main
         // and aux LDEs are co-resident, plus the composition and Merkle
         // transients (in the scratch factor). The aux width comes from the AIR
@@ -3484,7 +3487,7 @@ pub trait IsStarkProver<
                                 let root = tree.root;
                                 return Ok((
                                     Some(TableCommit::plain(tree, root)),
-                                    (aux_data, num_cols),
+                                    Table::from_row_major(aux_data, num_cols),
                                     Some(handle),
                                 ));
                             }
@@ -3580,7 +3583,7 @@ pub trait IsStarkProver<
                                 crate::instruments::accum_r1_aux(aux_lde_dur, Duration::ZERO);
                                 return Ok((
                                     Some(TableCommit::plain(tree, root)),
-                                    (aux_data, num_cols),
+                                    Table::from_row_major(aux_data, num_cols),
                                     Some(handle),
                                 ));
                             }
@@ -3627,15 +3630,30 @@ pub trait IsStarkProver<
                         #[cfg(feature = "instruments")]
                         crate::instruments::accum_r1_aux(aux_lde_dur, t_sub.elapsed());
 
+                        // Spill the aux LDE itself (`lde_size × aux_cols` ext3
+                        // elements, 24 B each — the dominant survivor of this
+                        // stage). It is write-once here and read-only from
+                        // rounds 2-4, so mmap-backing it frees the heap buffer
+                        // for the whole window and lets the OS evict the pages
+                        // under pressure.
+                        #[allow(unused_mut)]
+                        let mut aux_lde = Table::from_row_major(aux_data, total_cols);
+                        #[cfg(feature = "disk-spill")]
+                        if storage_mode == StorageMode::Disk {
+                            aux_lde
+                                .spill_to_disk()
+                                .map_err(|e| ProvingError::DiskSpill(format!("aux LDE: {e}")))?;
+                        }
+
                         #[cfg(feature = "cuda")]
-                        return Ok((Some(commit), (aux_data, total_cols), None));
+                        return Ok((Some(commit), aux_lde, None));
                         #[cfg(not(feature = "cuda"))]
-                        Ok((Some(commit), (aux_data, total_cols)))
+                        Ok((Some(commit), aux_lde))
                     } else {
                         #[cfg(feature = "cuda")]
-                        return Ok((None, (Vec::new(), 0), None));
+                        return Ok((None, Table::new(Vec::new(), 0), None));
                         #[cfg(not(feature = "cuda"))]
-                        Ok((None, (Vec::new(), 0)))
+                        Ok((None, Table::new(Vec::new(), 0)))
                     }
                 })()?;
             // Tuple shape is cfg-gated; `.0` is the optional TableCommit in

crypto/stark/src/table.rs

@@ -313,6 +313,23 @@ impl<F: IsField> Table<F> {
         }
     }
 
+    /// Wrap a buffer the caller already produced in row-major order, skipping
+    /// `new`'s 2-D revalidation — that `debug_assert` clones the whole buffer,
+    /// which is prohibitive for LDE-sized data.
+    pub(crate) fn from_row_major(data: Vec<FieldElement<F>>, width: usize) -> Self {
+        if width == 0 {
+            return Self::new(Vec::new(), 0);
+        }
+        let height = data.len() / width;
+        Self {
+            data,
+            width,
+            height,
+            #[cfg(feature = "disk-spill")]
+            mmap_backing: None,
+        }
+    }
+
     /// Creates a Table instance from a vector of the intended columns.
     pub fn from_columns(columns: Vec<Vec<FieldElement<F>>>) -> Self {
         if columns.is_empty() {
@@ -336,6 +353,7 @@ impl<F: IsField> Table<F> {
     }
 
     /// Given a row index, returns a reference to that row as a slice of field elements.
+    #[inline]
     pub fn get_row(&self, row_idx: usize) -> &[FieldElement<F>] {
         #[cfg(feature = "disk-spill")]
         if let Some(ref backing) = self.mmap_backing {
@@ -424,6 +442,7 @@ impl<F: IsField> Table<F> {
     }
 
     /// Given row and column indexes, returns the stored field element in that position of the table.
+    #[inline]
     pub fn get(&self, row: usize, col: usize) -> &FieldElement<F> {
         #[cfg(feature = "disk-spill")]
         if let Some(ref backing) = self.mmap_backing {

crypto/stark/src/trace.rs

@@ -321,7 +321,9 @@ where
     /// Row-major main-trace buffer of length `num_rows * num_main_cols`.
     pub(crate) main_data: Vec<FieldElement<F>>,
     /// Row-major auxiliary-trace buffer of length `num_rows * num_aux_cols`.
-    pub(crate) aux_data: Vec<FieldElement<E>>,
+    /// A `Table` rather than a bare `Vec` so it can carry a disk-spilled
+    /// (mmap-backed) buffer — see the aux-LDE spill in the prover's aux stage.
+    pub(crate) aux_data: Table<E>,
     pub(crate) num_main_cols: usize,
     pub(crate) num_aux_cols: usize,
     pub(crate) num_rows: usize,
@@ -467,7 +469,7 @@ where
 
         Self {
             main_data,
-            aux_data,
+            aux_data: Table::from_row_major(aux_data, num_aux_cols),
             num_main_cols,
             num_aux_cols,
             num_rows,
@@ -483,21 +485,24 @@ where
     /// Build an LDETraceTable directly from row-major flat buffers. Skips the
     /// O(N·M) col→row transpose that `from_columns` pays — the caller produces
     /// the buffers row-major already (e.g. via `coset_lde_full_expand_row_major`).
+    ///
+    /// `aux_data` arrives as a `Table` so an already disk-spilled aux LDE keeps
+    /// its mmap backing instead of being pulled back onto the heap; its `width`
+    /// is the aux column count.
     pub fn from_row_major(
         main_data: Vec<FieldElement<F>>,
         num_main_cols: usize,
-        aux_data: Vec<FieldElement<E>>,
-        num_aux_cols: usize,
+        aux_data: Table<E>,
         trace_step_size: usize,
         blowup_factor: usize,
     ) -> Self {
         let lde_step_size = trace_step_size * blowup_factor;
+        let num_aux_cols = aux_data.width;
         let num_rows = if num_main_cols > 0 {
             debug_assert_eq!(main_data.len() % num_main_cols, 0);
             main_data.len() / num_main_cols
         } else if num_aux_cols > 0 {
-            debug_assert_eq!(aux_data.len() % num_aux_cols, 0);
-            aux_data.len() / num_aux_cols
+            aux_data.height
         } else {
             0
         };
@@ -566,7 +571,7 @@ where
             self.main_data = main_data;
         }
         if !aux_data.is_empty() {
-            self.aux_data = aux_data;
+            self.aux_data = Table::from_row_major(aux_data, self.num_aux_cols);
         }
         self.host_trace_empty = false;
     }
@@ -636,7 +641,7 @@ where
     /// Get a single aux-trace element by (row, col).
     #[inline]
     pub fn get_aux(&self, row: usize, col: usize) -> &FieldElement<E> {
-        &self.aux_data[row * self.num_aux_cols + col]
+        self.aux_data.get(row, col)
     }
 
     /// Borrow a full main-trace row as a contiguous slice (row-major buffer).
@@ -648,7 +653,7 @@ where
     /// Borrow a full aux-trace row as a contiguous slice (row-major buffer).
     #[inline]
     pub fn aux_row(&self, row: usize) -> &[FieldElement<E>] {
-        &self.aux_data[row * self.num_aux_cols..(row + 1) * self.num_aux_cols]
+        self.aux_data.get_row(row)
     }
 
     /// Gather a full main-trace row into an owned Vec.
@@ -878,7 +883,7 @@ where
             // buffer-level check as the main arm: mixed states are valid here.
             #[cfg(feature = "cuda")]
             assert!(
-                lde_trace.num_aux_cols() == 0 || !lde_trace.aux_data.is_empty(),
+                lde_trace.num_aux_cols() == 0 || !lde_trace.aux_data.row_major_data().is_empty(),
                 "R3 barycentric (aux) fell back to the host trace, but it is device-only (empty)"
             );
             let inv_denoms_v =