fix(predict): make torch.compile take effect, and unblock the GPU pipeline

[jayhesselberth]

Investigating persistently patchy GPU utilization in leech predict (the classify stage of 2026-aa-trna-models). Every number below is from one real workload — 176,210 reads / 137,137 chunks, 1.8 GB POD5, the 20-class TCNDwellResidualLN production bundle with a CL-regression head — on an A30 with 16 logical CPUs, which is 8 physical cores (CR_CORE, both hyperthread siblings per core), the shape this runs in: four such jobs share one 64-core/4-GPU node.

The headline is a negative result

At that size the process runs at 904% of a ~9.4-effective-core ceiling, i.e. 87-96% CPU-saturated, and ~84% of extraction CPU is escapepod-signal's resquiggle banded DP (perf: DpContext::step 37.9%, dwell_block_kernel_avx512 24.3%, dp_step_buffered 18.0%, refine_signal_map 4.2%). Extraction is ~70% of the job's CPU.

So the GPU is not underfed by a scheduling defect — it is simply five times faster than the CPU can supply it, and no amount of pipeline restructuring changes that on a saturated box. Overlap fixes here buy ~2% wall; the real lever is upstream in the DP, which is filed separately.

What this PR does is fix the things that were outright broken or actively costing, and stop the pipeline stalling the GPU in long blocks.

torch.compile has never taken effect

run_inference assigned the compiled module to model_wrapper.model, but forward_batch calls self.forward_module — separate attributes. Measured, batch 1024:

	chunks/s
eager	16,571
torch.compile assigned to .model (what shipped)	16,556 (1.00x)
assigned to .forward_module, hook present	30,019 (1.81x)
reduce-overhead, no hook	30,227 (1.83x)

Two more problems behind it. Compilation was skipped outright whenever a CL-regression repr hook existed — every production multiclass bundle — but only CUDA graphs cannot replay the hook's Python side effect; plain inductor graph-breaks there and still fuses around it, recovering 1.81x of the 1.83x ceiling. And every mega-batch ends in a short batch, so an unpadded run presents ~19 distinct shapes and recompiles for each: enabling compile naively measured 45.8s → 101.5s. _stack_to_device now takes pad_to and the batch runners slice the padding off; the model is per-sample throughout (convolutions, GroupNorm/LayerNorm, BatchNorm in eval) so padding cannot change a real row.

The auto-enable threshold moves 5,000 → 2,000,000 reads. Compiling costs ~60s wall and buys ~10% at this deployment size, so it only pays past roughly ten minutes of inference; the old threshold turned it on for essentially every run, at a large net loss.

Extraction sub-batching was dead code

_SUB_BATCH_SIZE was 50,000 reads against a default read_batch_size of 10,000, so the loop written to feed the GPU continuously always ran exactly once — the failure escapepod-rs#361 documents on its own GPU pipeline ("the GPU consumer sits fully idle for the entire duration of that one giant prep call"). It is now real and configurable, defaulting to the whole mega-batch: splitting it only pays where there are idle cores, and on a saturated 8-core box it cost ~4s in added GIL contention.

The BAM writer was on the consumer's critical path

The consumer blocked on the previous write before it could submit the next. Writing a mega-batch is slower than producing one (pysam tagging ~7k reads/s), so that wait was 22.9s of a 46s run — and that is precisely when the GPU had nothing queued. Now a dedicated writer thread behind a bounded queue: still exactly one thread touching pysam, still in mega-batch order, but the consumer no longer waits. BGZF compression also moves into an htslib thread pool to get deflate off the GIL. The put is timeout-guarded so a dead writer surfaces as an error rather than a hang.

GPU batches likewise queue several deep instead of a one-deep handshake. The executor stays max_workers=1 — load-bearing, since it keeps GPU calls sequential, keeps _stack_to_device's thread-local pinned buffer single-owner, and keeps pending[read_id] in extraction order.

Rayon threads reserved proportionally

avail - 6 is a flat reservation that leaves 1 extraction thread at --cpus-per-task 4 and 2 at 8, on the stage that is ~70% of the job's CPU. Now a quarter, capped at 4. Extraction scales cleanly to the allocation's physical core count and only ~12% further across the hyperthread siblings — on 8 physical cores: 149.1s at 2 threads, 76.5s at 4, 40.1s at 8, 35.8s at 14, 34.3s at 16.

Measured effect

Interleaved A/B (alternating arms, warm page cache) at the 16-logical-CPU deployment size.

Out of the box, letting leech pick its own rayon thread count — 10 threads before, 12 after — 3 repeats:

	wall	reads/s	GPU duty (non-idle samples)
before	48.3s, 47.4s, 47.7s	3,688	40.7%, 40.2%, 40.3%
after	46.0s, 46.1s, 46.0s	3,827	45.3%, 47.7%, 46.4%

3.8% faster, +15% relative GPU duty, and notably tighter run-to-run spread.

With RAYON_NUM_THREADS=14 pinned in both arms, isolating the pipeline changes alone, 4 repeats:

	wall	GPU duty
before	45.9s, 45.7s	43.0%, 40.2%
after	44.9s, 44.9s	49.7%, 48.6%

~2% wall, ~+18% relative duty cycle, and the long per-mega-batch stall is gone. On a box already at ~90% CPU saturation, ~2% is about all overlap can be worth — which is the point of the negative result above.

Output is byte-identical: 176,210 records compared on aa/ac/am/pn/pp/pc, 0 differing, same file size. This is a scheduling change, not a grouping change.

Testing

uv run pytest: 1131 passed, 18 skipped. One failure, test_profiling.py::test_gpu_util_sampler_no_cuda_is_noop, reproduces identically on unmodified main on a CUDA-present machine and is unrelated. ruff check, ruff format --check and ty check all clean. tests/test_inference.py gains a guard that the new writer thread is joined at teardown, alongside the existing one for executor shutdown — same fork-safety hazard, new shape.

Not done here, deliberately

• The DP is the real target, and I could not find a byte-identical win in it. It runs twice per read with no convergence check, over the whole ~145-base aligned region, and leech discards the fitted (shift, scale, drift) it computes each iteration. Windowing it to the ~30 bases actually consumed is not safely doable: both band endpoints are hard-anchored (bands.rs asserts start[0]==0 / end[last]==seq_len) and six of the statistics feeding each pass are global over the read. Buffer reuse in banded_dp looked promising — ~400 KB of fresh Vecs per pass, 2 passes/read, past glibc's 128 KB mmap threshold — but testing the hypothesis for free by raising MALLOC_MMAP_THRESHOLD_ to 8 MB moved extraction 20.8s → 20.7s, i.e. nothing, so that change was not worth building. What is left upstream either changes numerics (dropping a rescale iteration) or is ~2%.
• Zero-copy move tables via PyO3 (moves.tolist() → Vec<Vec<u8>>). Measured before building: .tolist() over 255M elements is 1.0s and total FFI marshalling is 3.3s of 35.7s. Worth ~8%, real, but not the lever.
• Rust returning batched contiguous arrays instead of 3 numpy objects per chunk. The per-chunk Python glue it would remove measures 2.2s for 137K chunks (16 µs/chunk), so this is not where the time is either.

🤖 Generated with Claude Code

https://claude.ai/code/session_01V6XtaYNMsjxDLoD9aW1KJt