fix(instruments): restore the rounds 2-4 phase wall, make the prover timing report honest

[MauroToscano]

Review fixes for #877, targeting gpu-opt-table-scheduler so they land with it. Rebuilt onto the current branch head (5d452853) — merges clean.

Context on why there are two PRs: #893 was an earlier version of these fixes that merged at 20:42 with reasoning that later turned out to be wrong. This PR keeps what was right, reverts what wasn't, and corrects two false statements #893 introduced. Net effect of #893 + this PR is the state below.

Adversarial verification killed several of my own findings along the way; what's left is what survived.

1. Restore the rounds 2-4 phase wall

scripts/profiling/phase_table.py:129 does e["wall_ns"] += s["wall_ns"] — spans sharing a label are summed. On main, rounds_2to4 was one span wrapping the chunk loop, giving a true phase wall. #877 made it one span per table, so the row now reports the sum over N concurrent tables — up to k× the real wall, and it can print >100%. Worse, nothing spanned the fused region at all, so the genuine rounds-2-4 wall time was absent from the timeline.

Fix: one span on the caller thread wrapping the fused run_admitted, plus per-table spans renamed to *_table so they no longer collide with a phase label. r1_aux_build / r1_aux_commit had the identical defect (also phase-level on main, also per-table now) and are renamed too.

Evidence it works — phase spans now sum to their parent:

r1_prepass 0.148 + r1_main_commit 2.493 + rounds_2to4 8.943 = 11.584   vs   proving 11.585

Side effect worth knowing: this also fixes LAMBDA_VM_NSYS_CAPTURE_SPAN=rounds_2to4, which was genuinely broken — N driver threads each called cuProfilerStart/Stop, so the first table to finish killed the capture while the rest ran. scripts/profiling/README.md:115 and nsys_phase_busy.py:243 both name that label and now stay valid unedited.

What I got wrong here, for the record: I originally claimed the depth-0 spans broke phase_table.py's "% of total" column. They don't — :121 computes the denominator as max(wall_ns), not the root of the ancestor stack. I also claimed #877 corrupted a clean tree; continuation.rs on main already opens five spans on worker threads at depth 0, and its own comment at :1051-1053 says so. So instruments.rs:5-14 was already inaccurate before #877. #893's depth-plumbing (current_depth/enter_depth/DepthGuard) was built on that wrong premise and is removed here. The module doc is rewritten to describe actual behavior, including the label-summing hazard.

2. Make the prover timing report honest

prover/src/instruments.rs:74 computes round1 = main_commits + aux_build + aux_commit, but prover.rs:3253/:3605 hardcode the two aux buckets to Duration::ZERO. So "Round 1" and " Main trace commits" print the identical number, Round 1 understates by the whole aux stage, and Rounds 2-4 silently absorbs it. The two zero rows (:89, :110) also sat over live children, since accum_r1_aux still fires at prover.rs:3410/3449/3497.

This is in the terminal report every instruments user sees. Fixed with honest labels — the fused stages genuinely can't be timed separately any more, so the timers aren't resurrected.

3. Stale docs

The orphaned plan_table_chunks doc was left contiguous with VramGate's, so that struct's rustdoc opened by describing a deleted function returning (start, end) ranges. Plus: the "Split into two passes" block contradicted the FUSED comment two lines below it; table_parallelism's header still claimed cores/3; four chunk/Phase-D comments described removed structure; and Phase D had 7 references with zero definitions after its banner was deleted while A/B/C kept theirs.

The one a reviewer would actually be misled by: 3134-3135 and 3268 asserted pairing happens via "Phase D's zip chain", when it's now index-keyed mutex cells (gpu_main_cells).

Also corrected the Lde doc, which understated the PR's own win: it claimed all N tables' LDE columns are live at once. Only the main LDEs are — each aux LDE is created in aux_stage and consumed by value in rounds_stage, so at most k coexist (all N under debug-checks, where the two-pass split holds them).

4. One CI line

TABLE_PARALLELISM: 6 on one prover shard. (cores/3).max(1) == 1 for any runner under 6 cores, and every job is ubuntu-latest, so PR CI was running the scheduler with exactly one driver thread. This is the only change that puts two real table closures concurrently in flight on a PR.

It passed — shard 1 green at 7m59s, no OOM. If it ever gets tight, dropping to 2 still buys the concurrency.

Note: the concurrent and blocking paths do run pre-merge — gpu-tests on merge_group rents a ≥16-core 5090, hits the cuda arm at k≈10 with a finite VRAM budget. The gap was PR-time only.

5. Bench script

scripts/bench_prover_scaling.sh parsed and ran heap-growth regress on two snapshots #877 deleted. It's loud rather than silent ((insufficient data) / -) and the script is unwired — no Makefile target, no workflow — so this is minimal: dead rows dropped, and four regexes #893 had fixed are kept fixed rather than reverted.

Deliberately not fixed

prover/src/auto_storage.rs is out of this diff — #877 never touched it. The table_parallelism → auto_storage::decide() coupling is real, but my "doubles the transient term" framing was wrong: peak_bytes sorts descending and takes the top k, so each extra slot adds the next smallest table. Sub-linear, not proportional.

No scheduler unit tests. Every assertion I considered is guaranteed by construction, already covered by the end-to-end multi_prove + multi_verify suite, or unreachable from the call sites — and the one with teeth (*used == 0, oversized-runs-alone) would hang rather than fail, which on an 8-10 minute shard burns to the job timeout. The CI line above retires more risk for one line.

Also untouched: TAIL_MAX_PAIRS, the FRI/keccak/NTT kernels, crypto/math-cuda, and any scheduler behavior — the fused chain, gate semantics, heaviest-first, and the cores*2/3 constant are all unchanged.

Verified

cargo fmt --all --check; cargo check -p stark × {default, instruments, debug-checks, disk-spill, cuda, all-four-combined}; cargo check -p lambda-vm-prover --features instruments,cuda,debug-checks,disk-spill; cargo test --release -p stark --lib → 202 passed; clippy clean in every file touched. Re-verified after merging the current branch head.

Not GPU-verified — no CUDA on the machine these were written on, and gpu-tests doesn't run on PRs.

Two open questions for you, not fixes

cores * 2 / 3. The justification is one parenthetical ("swept flat at ~2/3 of the cores on a 16-core/RTX 5090 box") with no data in the repo. Tying k to host cores sits oddly with the stated premise that in-flight tables "mostly sit in GPU waits" — if they're waiting on the device they aren't consuming cores, so the same 5090 would admit 42 tables on a 64-core host and 5 on an 8-core one. Note VramGate already bounds memory independently, so k's only remaining job on a GPU build is latency hiding. Worth publishing the sweep, or deriving k from a device quantity.

5d452853's pinned-staging note. Same shape: the measurement behind "per-driver slots cost more than the shared mutex" isn't in the repo. Worth attaching, since a bot flagged the slot-0 collapse as High and the code comment is currently the only rebuttal.

One residue this PR can't reach: 8f91e7c7's squash-merge message bakes in the refuted "% of total" reasoning. If #877 squash-merges to main a fresh message is written and it disappears; if history is preserved, that paragraph lands. Not worth rewriting merged history — flagging so you can decide.