fix(gpu): drain htod_via on error; narrow the merkle-tail threshold

[MauroToscano]

Review follow-ups for #875, rebased onto e75bcbed — only what that commit did not already cover. Your follow-up landed while I was writing this, so I dropped everything it addressed (gather bounds, the query-0 canary, the test-faults zero-total guard).

Compile-verified (math-cuda --tests, stark --features cuda), fmt and clippy clean. Not GPU-verified — no CUDA device here, and gpu-tests doesn't run on PRs.

1. htod_via can free the pinned slab under an in-flight DMA

The one memory-safety item; everything else here is polish.

Once a chunk's cuMemcpyHtoDAsync_v2 is enqueued, record_event(stream)? / sync_event()? returning Err drops the staging MutexGuard with the device still reading the slab — the next locker's ensure_capacity can then cuMemFreeHost it mid-copy. async_dtoh_via guards exactly this (device.rs:761-764: "if the record fails, drain the stream before the guard drops, or the next locker could free the slab mid-copy"), and the file already ships DrainOnErr for it — htod_via was the one site not using it. Armed before checking the memcpy result, since the driver may enqueue before reporting an error.

Needs a CUDA failure mid-upload, so low probability.

2. TAIL_MAX_PAIRS 2048 → block width

The tail grid-strides a single 128-thread block on one SM, so a level of k pairs costs k/128 sequential keccak-f1600s where the per-level launches spread them over k/128 parallel blocks. Each hash_merkle_parent absorbs 64 B < the 136 B rate, so it's exactly one permutation per pair.

At 2048 the first four levels (2048/1024/512/256) are 16+8+4+2 = 30 serial permutations against 4 parallel waves — ~26 extra serial keccak-f latencies to save 4 launches. Rough ledger per large tree: cost ~78–156 µs (at 3–6 µs per serial permutation), saving ~6–10 µs of launch overhead. It stays on the critical path because the caller's 32-byte root memcpy_dtoh host-blocks on everything queued before it.

Scale honestly: ~10–40 ms/epoch, i.e. 0.05–0.5% — this is a threshold that overshoots, not a regression that matters. At the block width the entry level is exactly one permutation per thread: still collapses the top levels into one launch, with zero added serialization. The per-permutation latency is a model, not a 5090 measurement; the direction holds for any per-keccak-f latency above ~0.4 µs, magnitude ±2×.

3. R2 host-evaluator fallback panics with a bare index-OOB

If evaluate_dev succeeds but both try_decompose_extend_d2_dev and the H download fail under device-only, control reaches the host evaluator, which reads the intentionally-empty trace. Added the device-only contract assert, matching the other fallback arms. Narrow window; improves the message only.

4. Coverage: the device n == 1 inverse path

batch_inverse_ext3_dev's n == 1 branch (a direct launch_invert_total) is never exercised — batch_inverse_n1 goes through batch_inverse_ext3's host-only short circuit at inverse.rs:59-64, as its own comment says. Added a direct device test; complements the htod_via round-trip test you added.

5. Two doc comments

• prover.rs:1030 still said the split trees "come back as full host trees so the preprocessed opening path and the process-wide precomputed-tree cache work unchanged". Half of it survives — the precomputed tree is still a full host tree and its cache is unchanged — but the multiplicity tree is now root-only + device-resident. This is the comment a reader consults before touching the is_root_only assert.
• fri.rs:43 says FriCommitState's input is Arc-shared with the layer's retained gpu_evals; since gpu_evals: (!want_host).then_some(..) that holds only on the device-only path.

---

Flagged, not patched — three failure-mode calls that are yours

I checked these and deliberately left the code alone, because each is a design choice you made with a stated rationale rather than a defect.

1. The R2 commit assert (prover.rs:1669). I originally wanted this to be a rejectable ProvingError on the rejectable-beats-panic convention, and that was wrong — I checked the propagation and there is nothing to recover to. ProvingError is flattened at continuation.rs:738 / :931 via .map_err(|e| Error::Prover(format!("{e:?}")))? and propagates straight up; no caller retries, and there is no CPU-fallback-on-ProvingError path. The EmptyCommitment one line below is equally terminal. Panic and Err both end the prove with no proof, so the assert costs nothing and keeps the better message. Leaving it as-is; no change wanted. (The convention it appeared to violate is really about the verifier, where a malicious proof must be rejected rather than panic the process. This is a prover-side internal invariant with no attacker-controlled input.)

2. is_none_or in the three contract asserts passes an empty outer Vec (num_parts == 0 → .first() == None → assert passes). Unreachable — R2 hits EmptyCommitment first — and note the right predicate differs per site: at R2 an empty outer Vec should fall through to EmptyCommitment, so a blanket swap to is_some_and would be wrong there. Hardening only.

3. Your zero-total guard. #[cfg(any(debug_assertions, feature = "test-faults"))] fixes the "never actually runs" problem for test-cuda-fallback. Two things it still doesn't reach: the other four GPU suites (plain --release) and production. It also still assert_ne!s, where the host Fermat it replaced returned Err that try_compute_and_invert_inv_denoms_dev maps to None → CPU fallback.

One more, lower confidence: prover.rs:2208's assert is now load-bearing in a way its predicate doesn't state. Its loop reads the host trace and the host parts (2255), but it's gated on !host_trace_empty() alone. Sufficient today only because want_host == !host_trace_empty(). Decouple those later — e.g. "keep parts device-only whenever a handle exists" — and it silently stops covering 2255.

Verified rather than changed

The two risky kernel rewrites are correct, checked numerically rather than by reading:

• ntt_dit_8_levels_row_major vs the per-level reference over Goldilocks, 9 (log_n, m, gridDim.y) combinations including gridDim.y not dividing n/256 and m not a multiple of T — bit-identical in all.
• keccak_merkle_tail vs the multi-launch loop across 13 leaf counts including non-powers-of-two, with the grid-strided thread iterations executed in reverse order to expose any hidden intra-level dependency — every node buffer matched, not just the roots.

Also traced and found sound: every consumer of lde_composition_poly_evaluations under empty parts (each is either device-routed or assert-guarded — no path produces a silently wrong or silently empty proof); R3 parts-OOD equivalence argument-by-argument, including inv_denoms at z_power in both arms and matching DenomSign; per-part ordering device-vs-host; the Fermat exponent (p³−2, Hamming weight 127, no u128 overflow); the slab layout against keccak_comp_poly_leaves_ext3; the preprocessed/mult-tree pairing; and cross-stream ordering for every new resident consumer.

Residual risk worth naming: build_comp_poly_tree_from_slabs_dev and gather_ext3_at have no parity test, and they are the only sources of the composition root and the FRI sym evals under device-only — precisely where the host cross-checks are skipped. Rust-level equivalence I did verify (the tree builder launches the identical kernel with identical args over an identically-laid-out buffer), so the residual is kernel-level only.