Fix: Add ScalarConstDivPow2 and range-check ScalarConstDiv remainders by ClankPan · Pull Request #210 · ICME-Lab/jolt-atlas

ClankPan · 2026-04-02T05:10:51Z

This PR adds ScalarConstDivPow2 as a dedicated path for fixed-point rebasing when the divisor is a power of two. Operations like Mul and Einsum now use this path when possible.

ScalarConstDivPow2 represents the remainder with a one-hot RA encoding and range-checks it with the existing LUT and commit_to_onehot() flow. The reason for using a LUT here is that it fits the current lookup and RA machinery, so we can add this optimization without changing the commitment layer. A future commit_to_bits()-style path should be even lighter, but in this PR I prioritized an implementation that works with the current commitment design.

This PR also updates ScalarConstDiv for the general-divisor case. It now proves 0 <= r < divisor with the same LUT/RA idea that Div uses. As part of that change, the dense remainder commitment is removed and replaced with a virtual remainder plus range-check claims.

One more issue showed up while testing this change. In onnx_proof::e2e_tests::test_multihead_attention, a ScalarConstDivPow2 node can produce an output tensor of length 1. In that case, the RA sumcheck ends up with zero rounds and panics. Fixing this properly requires a change in the sumcheck side so that zero-round RA virtualization is handled correctly. For now, scalar ScalarConstDivPow2 skips the RA one-hot checks, matching the current Div behavior. This is a real soundness issue, so it should be fixed in a separate PR.

Close #203

ClankPan · 2026-04-02T05:50:27Z

cargo test -p jolt-atlas-core test_gpt2 -- --ignored --nocapture was rerun locally after the latest changes and completed successfully. On my machine it finished in 446.32s total (Proof generation took 275.18s, Proof verification took 143.42s), so the CI SIGKILL looks like a runner timeout/cancellation issue rather than a proof failure.

Forpee · 2026-04-02T09:03:44Z

cargo test -p jolt-atlas-core test_gpt2 -- --ignored --nocapture was rerun locally after the latest changes and completed successfully. On my machine it finished in 446.32s total (Proof generation took 275.18s, Proof verification took 143.42s), so the CI SIGKILL looks like a runner timeout/cancellation issue rather than a proof failure.

@ClankPan Could you run the timings on your machine before these changes? The 143.42s verification time seems higher than expected, so I'd like to compare baselines

Copilot

Pull request overview

This PR addresses soundness and performance in fixed-point “rebase” division by introducing a specialized ScalarConstDivPow2 operator (for power-of-two divisors) and by adding remainder range-checking to the general ScalarConstDiv path using the existing LUT + RA one-hot machinery.

Changes:

Add ScalarConstDivPow2 operator and wire tracer handlers (Mul, Square, Einsum) to prefer it for rebase-by-power-of-two.
Update ScalarConstDiv to remove dense remainder commitments and instead prove 0 <= r < divisor via the range-check + RA one-hot pipeline.
Extend common polynomial enums/serialization and witness generation to support the new committed/virtual polynomials.

Reviewed changes

Copilot reviewed 15 out of 15 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
jolt-atlas-core/src/onnx_proof/witness.rs	Adds witness generation for `ScalarConstDivRangeCheckRaD` and `ScalarConstDivPow2RaD`; removes remainder witness commitment for `ScalarConstDiv`.
jolt-atlas-core/src/onnx_proof/range_checking/range_check_operands.rs	Introduces `ScalarConstDivRangeCheckOperands` to range-check scalar-constant division remainders.
jolt-atlas-core/src/onnx_proof/range_checking/mod.rs	Adjusts RA one-hot encoding to source `r_cycle` from the first range-check operand.
jolt-atlas-core/src/onnx_proof/ops/scalar_const_div.rs	Updates ScalarConstDiv proof flow to include range-check + RA one-hot proofs and virtual remainder handling.
jolt-atlas-core/src/onnx_proof/ops/scalar_const_div_pow2.rs	New: custom proof for power-of-two divisor rebasing with RA one-hot remainder encoding.
jolt-atlas-core/src/onnx_proof/ops/mod.rs	Registers/dispatches the new operator.
common/src/lib.rs	Adds new committed/virtual polynomial variants + serialization tags.
atlas-onnx-tracer/src/utils/handler_builder.rs	Adds `with_auto_rebase_pow2()` and emits `ScalarConstDivPow2` rebase nodes when enabled.
atlas-onnx-tracer/src/ops/scalar_const_div.rs	Enforces positive divisor invariant at runtime.
atlas-onnx-tracer/src/ops/scalar_const_div_pow2.rs	New: execution semantics for `ScalarConstDivPow2`.
atlas-onnx-tracer/src/ops/mod.rs	Registers new operator type.
atlas-onnx-tracer/src/node/handlers/other.rs	Switches Einsum rebase to power-of-two path.
atlas-onnx-tracer/src/node/handlers/arith.rs	Switches Mul/Square rebase to power-of-two path and updates docs.
atlas-onnx-tracer/src/model/shadow_trace.rs	Adds shadow execution behavior for `ScalarConstDivPow2`.
atlas-onnx-tracer/src/model/mod.rs	Updates `max_num_vars` sizing to account for new lookup/RA needs.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-04-02T11:19:57Z

+        accumulator.append_virtual(
+            transcript,
+            VirtualPolynomial::DivRemainder(self.params.computation_node.idx),
+            SumcheckId::NodeExecution(self.params.computation_node.idx),


The remainder range-check is performed via the RAF pipeline, which reads DivRemainder(node_idx) under SumcheckId::Raf (see RangeCheckingOperandsTrait::operand_claims). In this execution sumcheck, the remainder opening is recorded under SumcheckId::NodeExecution, so a malicious prover can supply different remainder claims for NodeExecution vs RAF and still satisfy the execution identity while passing the range-check on an unrelated remainder. To make the range-check actually constrain the remainder used in the execution relation, route the RAF operand claim(s) to the same opening entry as execution (e.g., have range-checking read operand claims from SumcheckId::NodeExecution(node.idx) or explicitly mirror/copy the NodeExecution remainder claim into the RAF claim without recomputing/overwriting it).

Suggested change

SumcheckId::NodeExecution(self.params.computation_node.idx),

SumcheckId::NodeExecution(self.params.computation_node.idx),

opening_point.clone(),

self.remainder.final_sumcheck_claim(),

);

// Mirror the execution remainder claim into the RAF namespace so that

// range checking constrains the same remainder used in execution.

accumulator.append_virtual(

transcript,

VirtualPolynomial::DivRemainder(self.params.computation_node.idx),

SumcheckId::Raf(self.params.computation_node.idx),

Copilot · 2026-04-02T11:19:57Z

+        // remainder range proof instead of relying only on the execution
+        // relation.
+        if node.is_scalar() {
+            return results;


These comments say the scalar case “skips the RA one-hot range check” due to a zero-round RA sumcheck. Skipping the one-hot checks leaves the remainder encoding unconstrained (the execution sumcheck alone does not enforce booleanity/one-hotness), which makes the scalar ScalarConstDivPow2 case unsound. Consider temporarily disallowing ScalarConstDivPow2 for scalar outputs (fallback to ScalarConstDiv, or return an error/panic in tracer/prover) until zero-round RA virtualization is handled, or add a scalar-specific remainder range proof that does not require RA one-hot checks.

Suggested change

return results;

panic!(

"ScalarConstDivPow2 is currently unsupported for scalar outputs: \

scalar outputs skip the RA one-hot range check, which leaves the \

remainder encoding unconstrained and is unsound. \

Use ScalarConstDiv instead until a scalar-specific remainder \

range proof is implemented."

);

Copilot · 2026-04-02T11:19:58Z

+            let operator = if self.auto_rebase_pow2 {
+                Operator::ScalarConstDivPow2(ScalarConstDivPow2 { divisor: factor })
+            } else {
+                Operator::ScalarConstDiv(ScalarConstDiv { divisor: factor })
+            };


When auto_rebase_pow2 is enabled, this unconditionally emits ScalarConstDivPow2 { divisor: factor }, but determine_rebase_factor() can also return custom_rebase_factor (which may be non-power-of-two). That will trigger runtime assertions/panics later (execution/proof code requires power-of-two). Add a guard here (assert factor > 0 && is_power_of_two) or fall back to ScalarConstDiv when the chosen factor is not a positive power of two.

Copilot · 2026-04-02T11:19:58Z

+impl<F: JoltField, T: Transcript> OperatorProofTrait<F, T> for ScalarConstDivPow2 {
+    fn reduction_flow(&self) -> ReductionFlow {
+        ReductionFlow::Custom
+    }
+
+    #[tracing::instrument(skip_all, name = "ScalarConstDivPow2::prove")]
+    fn prove(
+        &self,
+        node: &ComputationNode,
+        prover: &mut Prover<F, T>,
+    ) -> Vec<(ProofId, SumcheckInstanceProof<F, T>)> {
+        let mut results = Vec::new();
+
+        let params = ScalarConstDivPow2Params::new(node.clone(), &mut prover.transcript);
+        let mut exec_sumcheck = ScalarConstDivPow2Prover::initialize(
+            &prover.trace,
+            params,
+            &mut prover.accumulator,
+            &mut prover.transcript,
+        );


This is a new operator implementation but there are no unit tests in this file. Adding tests similar to ops/scalar_const_div.rs would help catch regressions (e.g., random inputs, multiple divisors like 2/16/256, negative dividends, and a non-scalar tensor to exercise the RA one-hot proof path). Also consider covering the scalar-output edge case explicitly (even if it currently errors or falls back).

Copilot · 2026-04-02T11:19:58Z

+/// Prover state for the `ScalarConstDiv` execution sumcheck.
 ///
-/// Maintains the equality polynomial, operand polynomial, and remainder R
-/// needed to prove the division relation: operand = divisor * q + R where divisor is constant.
+/// The relation enforces `divisor * q + r - a = 0` at the sampled point, where
+/// `q` is the node output and `r` is the virtual remainder reconstructed by the
+/// range-check / RA pipeline.
 pub struct ScalarConstDivProver<F: JoltField> {
    params: ScalarConstDivParams<F>,
    eq_r_node_output: GruenSplitEqPolynomial<F>,
    left_operand: MultilinearPolynomial<F>,
-    R: MultilinearPolynomial<F>,
+    q: MultilinearPolynomial<F>,
+    remainder: MultilinearPolynomial<F>,


The doc comment states the remainder is “reconstructed by the range-check / RA pipeline”, but the prover currently constructs remainder_tensor directly from the left operand and uses it in the execution sumcheck. Consider adjusting the wording to reflect the actual flow (remainder is a witness polynomial whose validity is enforced by the range-check/RA proofs).

ClankPan · 2026-04-02T14:35:03Z

cargo test -p jolt-atlas-core test_gpt2 -- --ignored --nocapture was rerun locally after the latest changes and completed successfully. On my machine it finished in 446.32s total (Proof generation took 275.18s, Proof verification took 143.42s), so the CI SIGKILL looks like a runner timeout/cancellation issue rather than a proof failure.

@ClankPan Could you run the timings on your machine before these changes? The 143.42s verification time seems higher than expected, so I'd like to compare baselines

@Forpee

I ran this on my local machine and verifier looks like same time.

main

Proof generation: 186.515869125s
Proof verification: 148.093474291s
Total test time: 341.06s

feat/ConstScalarDiv_remainder_range_heck

Proof generation: 265.633045708s
Proof verification: 144.808440291s
Total test time: 433.73s

Forpee · 2026-04-02T15:38:50Z

@ClankPan
Thanks for checking. The verification times look consistent across runs, so that doesn’t appear to be the issue. On my side, I reran the gpt-2 execution profile and saw proof gen increase (gpt2 went from ~30s to ~60s), due to the reshape padding changes (#185) becoming the bottleneck (for prover and verifier) - which is interesting but out of scope for this PR, so u don't need to worry about that.

I'll proceed with the review shortly

ClankPan · 2026-04-02T15:53:59Z

@Forpee
This PR #211 changes the proof method for reshape. This change will likely bring the proof time back to normal. I will rebase once it is merged.

Forpee · 2026-04-02T15:57:47Z

Looks good. I introduced a related sub-protocol in the softmax draft PR here that uses the prefix-suffix approach, which is more efficient than relying on the vanilla identity polynomial as the table. It leverages the structure in the identity table (and sparsity in ra), resulting in faster proving

Forpee · 2026-04-02T16:00:22Z

Fixing this properly requires a change in the sumcheck side so that zero-round RA virtualization is handled correctly.

Noted. Lets create an issue for this.

ClankPan · 2026-04-02T16:35:02Z

Looks good. I introduced a related sub-protocol in the softmax draft PR here that uses the prefix-suffix approach, which is more efficient than relying on the vanilla identity polynomial as the table. It leverages the structure in the identity table (and sparsity in ra), resulting in faster proving

@Forpee
Thanks for letting me know this!
I'm going to rebase this branch if the PR is merged into main. I paused this PR until the merging.

Add power-of-two scalar div and range-check ScalarConstDiv

71c92df

ClankPan requested a review from Forpee April 2, 2026 05:12

Skip scalar RA checks for ScalarConstDivPow2

fa8abe9

Forpee requested a review from Copilot April 2, 2026 11:00

Copilot started reviewing on behalf of Forpee April 2, 2026 11:01 View session

Copilot AI reviewed Apr 2, 2026

View reviewed changes

Forpee mentioned this pull request Apr 7, 2026

Feat: gather for small and large #212

Merged

-            return results;
+            panic!(
+                "ScalarConstDivPow2 is currently unsupported for scalar outputs: \
+                 scalar outputs skip the RA one-hot range check, which leaves the \
+                 remainder encoding unconstrained and is unsound. \
+                 Use ScalarConstDiv instead until a scalar-specific remainder \
+                 range proof is implemented."
+            );

Conversation

ClankPan commented Apr 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ClankPan commented Apr 2, 2026

Uh oh!

Forpee commented Apr 2, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Apr 2, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 2, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 2, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 2, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 2, 2026

Choose a reason for hiding this comment

Uh oh!

ClankPan commented Apr 2, 2026

main

feat/ConstScalarDiv_remainder_range_heck

Uh oh!

Forpee commented Apr 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ClankPan commented Apr 2, 2026

Uh oh!

Forpee commented Apr 2, 2026

Uh oh!

Forpee commented Apr 2, 2026

Uh oh!

ClankPan commented Apr 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

ClankPan commented Apr 2, 2026 •

edited

Loading

Forpee commented Apr 2, 2026 •

edited

Loading