[taddyb]

Summary

On CUDA, the backward of every timestep op copied gb_rhs and x_sol to the host in
timestep_backward_core, looped over nnz to build gA, and uploaded the result.
It now calls dispatch::grada_primitive, as CsrSolveOp::backward already does.
With sparse_solver: cuda, gA is gathered, multiplied and negated on the device.

Why

With sparse_solver: cuda, each timestep op in the backward did 2 device-to-host copies,
1 upload and a stream sync for gA. All timestep op variants share this function
(TimestepOp, gamma, release, leakance). No issue or PR covered it.

CPU path

The CPU branch of grada_primitive is the same host loop: -gradb_host[r] * x_host[c], in the
same order, with the same 2 host reads. CPU results do not change. Checked three ways:

• A temporary test (not in this PR) printed the f32 bits of the CPU gradients of one timestep op.
  It covered a chain and a confluence network, enforce_positivity false and true, and all 4
  parents (16 gradient vectors). The output on master (65211a9) and on this branch is identical
  byte for byte.
• compare_ddr_sandbox prints ABSOLUTE MATCH, with the same numbers as master
  (max abs diff 1.525879e-4 m3/s, max rel diff 2.412143e-6).
• cargo test --features fixtures: 1017 passed, 2 failed, 26 ignored over 121 test binaries.
  The lib unit tests (490), mmc, sparse_gradcheck, sp8_gradcheck, leakance_gradcheck,
  leakance_off_parity, leakance_reference_match, reservoir_override, zeta_accum and
  ddr_sandbox_match all pass. juniata_acceptance (release) passes 3 of 3.
  The 2 failures also fail on master with the same numbers (see "Known failures on master").

Tests

Part E of tests/cuda_backward_parity.rs (needs --features cuda):

• A thread-local counter in src/sparse/mod.rs. primitive_to_vec increments it.
  host_reads_on_this_thread_for_test() reads it. It is a #[doc(hidden)] test hook, like
  the other *_for_test functions. It costs one Cell increment per host copy.
  primitive_to_vec is the only host-read helper on the timestep backward path.
• E1 cusparse_backward_keeps_ga_on_the_device: 0 host copies in the backward with
  sparse_solver: cuda. On master it fails with 2 copies. A control run with
  sparse_solver: cpu on the same device shows 4 copies, so the counter sees the backward.
• E2 cusparse_ga_matches_cpu_on_a_confluence: a network with two confluences (rows with 2
  off-diagonal entries, so row_for_nnz repeats). Gradients of all 4 parents, CUDA + cuSPARSE
  against the CPU backend, with and without enforce_positivity. Worst relative difference 3.2e-7.
  E2 passes on master too. It pins the new device gather.
  Mutation checks, done by hand and reverted: swapping the row and col indices in
  cusparse_grada fails E2 and cusparse_solver_path_carries_the_same_gradients. Indexing
  gradb by min(row, col + 1) (equal to row on a chain) fails only E2.

Local CUDA results (RTX 4080 SUPER, CUDA 13.3 driver, built with CUDARC_CUDA_VERSION=13020):

• cuda_backward_parity: 17 of 18 pass. This includes Parts A and B, E1, E2 and
  cusparse_solver_path_carries_the_same_gradients. The cuSPARSE and host-solve gradients are
  still bit-identical on the chain fixture (worst relative difference 0 for all 4 parents).
  The worst cross-backend difference in Part B is 2.784e-7, the value in the test comments.
• sparse_cusparse_v5 4 of 4 (1 more is ignored, see below), sparse_cusparse_v8 3 of 3
  (ignored tests, run with --include-ignored) and device_selection 1 of 1 pass.
  sp10_k1_bitmatch and sp10_k23_bitmatch have only ignored tests; they were not run.

Known failures on master (not touched here)

These fail on master (65211a9) in the same way. Each was run on a separate master build.

• cuda_backward_parity::enforce_positivity_can_never_reach_the_cuda_graph_path (Part D).
  It expects the enforce_positivity/ddr_match error text, but the use_cuda_graphs check
  fires first. CI does not build with --features cuda, so it went unseen.
• training_step_layer_b::layer_b_step3_mc_forward_matches_ddr (MC forward max abs diff 1.01e-1)
  and training_step_layer_c::layer_c_step2_gradients_match_ddr (grad_output_weight 2.03e-4 over
  the 1e-4 tolerance). Same numbers on master.
• sparse_cusparse_v5::end_to_end_smoke_cuda_train and
  sparse_cusparse_v6::v6_cuda_is_faster_than_cpu_on_smoke_train (both ignored). They run the
  train binary with /tmp/*.yaml files that do not exist.

Timing

One timestep op, sparse_solver: cuda, release build, RTX 4080 SUPER. Wide shallow networks
(n/10 parallel lanes). Median time of backward() plus reading the gradient, 10 runs after
3 warm-up runs, two rounds each. The machine was shared, so treat the numbers as rough.
The timing test is not in this PR.

reaches	master	this branch
10,000	0.88-0.95 ms	0.82-0.86 ms
100,000	2.07 ms	1.31-1.35 ms
1,000,000	8.40-9.75 ms	7.89-8.41 ms

Not done

• No CONUS timing.
• Trap T13 (CUDA routing slower than CPU on CONUS) is not fixed by this change.
• cusparse_grada still uploads row_for_nnz and col on each call. That is cheap next to the
  host round trip it replaces. Follow-up: keep these index tensors in CudaPatternCache.

CI covers only the CPU part. The CUDA results above are local.

🤖 Generated with Claude Code