perf(demux): per-read CRF fan-out, bounded GPU score memory - #173
Merged
Conversation
Two independent findings from a perf audit of the CTC-CRF path. `produce_cpu_crf` fanned out with `.chunks(256)` and then walked the chunk serially. That shape was copied from `produce_cpu_gbm`, but that head chunks because `predict_many` is a genuinely batched kernel; tract has no batched LSTM, so here the chunk only ever serialized its reads. At ~14 ms/read (13 ms encode + 1.2 ms decode, measured on rna) one chunk is ~3.6 s of unstealable work, and a block with fewer than `256 * threads` reads cannot fill the machine at all: 1000 reads made four tasks for 32 cores. `for_each_init` gives per-read stealing and moves `CrfScratch` from per-chunk to per-worker — the shape `produce_cpu` already uses. 10k reads 23.3 s -> 17.0 s, 1k reads 4.91 s -> 2.19 s on 32 cores, with per-read calls and every per-barcode POD5 byte-identical. `CrfEncoderGpu::basecall_batch` encoded the whole caller batch before decoding any of it. `ESCAPEPOD_CRF_GPU_BATCH_ROWS` bounds the device activations, not the host: scores are `t_len * n_score` floats, 1 MB per read for RNA004, so a several-thousand-read Arrow batch retained gigabytes. Encode and decode now alternate per device batch, capping host high-water at `batch_rows` reads. Splitting is already exact (`run_batch` documents and tests it), and the chunk boundaries are unchanged, so scores are identical. `run_batch` now takes `&[&[f32]]`, which also drops the per-read 8 KB window clone. Verified against the nbc16 RNA004 CRF bundle on 1k and 10k reads: the `--classifications` CSVs are identical after sorting (row order is rayon-nondeterministic) and all 14 per-barcode POD5s match byte for byte.
This was referenced Aug 1, 2026
jayhesselberth
added a commit
that referenced
this pull request
Aug 1, 2026
…re copy (#174) **Stacked on #173** — targets `worktree-crf-perf-fixes`, so review that one first. The diff shown here is only this commit. The remaining items from the CTC-CRF perf audit. ## 1. Barcode matching never actually abandoned early `edit_distance` passed WFA `cap = a.len() + b.len()` — the largest score any alignment can reach — so the `max_score` guard could **never** fire. Every one of the N references ran to its true optimum *and* did a full traceback. That defeats the reason WFA was chosen over a plain DP, and it contradicts the module docstring (and CLAUDE.md), which describe most comparisons abandoning almost immediately. Each comparison is now capped at the running runner-up. Both branches of the selection loop already discard any `d >= second`, and `second` is monotonically non-increasing, so a reference that cannot beat it never needs its exact distance — only the fact that it lost. `wfa_align_opt` also returns *before* the traceback, which is where the per-comparison allocations live. | Bank | Before (per comparison) | After | Per read | |---|---|---|---| | 16 refs | 7.74 µs | 5.65 µs | 124 → 90 µs | | 96 refs | 7.74 µs | 2.47 µs | ~743 → 237 µs (3.1×) | The cap tightens faster with more references, so this scales *with* the barcode design rather than against it. Main has no cap, so its per-comparison cost is bank-size independent — that's why the 96-ref "before" is the measured 16-ref figure. Output is unchanged field for field, pinned by a 2400-case test against the previous implementation plus explicit coverage of the tightest-cap and exact-match cases. ## 2. The encoder copied 1 MB per read for nothing `basecall_prepped` decoded from an owned `Vec` that `encode` filled element by element — `t_len * n_score` floats, 1 MB for RNA004 — which `decode_with`'s first loop immediately transposed into `CrfScratch` and then dropped. It now decodes straight out of tract's output tensor. `encode` is kept for callers that genuinely want to own the scores, with a doc note. The decode backend is also resolved once at load instead of re-probed every read. ## 3. Only the window is calibrated `prep` needs `chunk` samples of pA ending at `adapter_end`, but callers converted the entire decoded prefix first — 16 000 samples under the CNN detector (8× the window), and the **whole read** under LLR, which sets no decode bound at all. `prep_adc_into` fuses calibration and standardisation into one pass over exactly the 2000 samples the encoder sees, into a per-worker buffer. ## Two fixes that fell out **The two CRF entry points disagreed on picoamps.** `demux basecall` and `escapepod_python::adc_to_pa` use a fused `adc.mul_add(scale, offset * scale)`; the fused pipeline used unfused `(adc + offset) * scale` — despite a comment claiming it matched the reference. Both now use the reference form. `demux basecall` is bit-identical; the fused pipeline's encoder input shifts by ~1 ulp. Measured impact below. **The router budget only held to 192 barcodes.** `ROUTER_TOTAL_SLOTS` exists to stop queued-read memory scaling with barcode count, but the per-barcode depth was clamped to a floor of 256 — so past 192 barcodes it scaled again, the exact behaviour the budget replaced. The CRF head takes references from a user CSV, so the count is unbounded; a 384-plex set sat at ~2× budget. Floor is now 64, overshoot is logged, and nothing shipping changes (it doesn't bind until 768). ## Verification **GPU, on an A30 with `Successfully registered CUDAExecutionProvider` confirmed in the log** (not a silent CPU fallback): `demux basecall --gpu` output is **byte-identical to unmodified main across all 992 reads**. That single test covers #173's GPU chunking plus this PR's barcode cap, score-copy removal and `prep_adc_into`. **CPU fused pipeline**, vs unmodified main: the ~1 ulp pA change moved the reported *confidence* on **1 of 992** and **7 of 9943** reads and changed **zero barcode calls**. That signature — confidence-only, ~0.1%, no call changes — is what a ulp-level input perturbation looks like; a logic error in the cap would change calls or perturb margins systematically. 465 workspace tests, 101 `crf-decode` tests (incl. `crf_golden decode_matches_bonito`), doctests, `fmt`, and `clippy -D warnings` across default / `crf-decode` / `crf-gpu` / `cnn-gpu`. ## Deliberately not done - **`BLOCK_TARGET_BYTES` was not retuned**, only documented. On close reading it is honestly a *per-block* cap and its tuning table already measures peak RSS across filler counts, so the "it's not a global cap" audit item was an over-reading of the comment. Retuning a measured default on a bad premise would be worse than leaving it. - **`split_time_major` still runs under the session mutex** — the copy-out is the expensive part, so fixing it needs `IoBinding` or a real restructure, not a lock move. - **`fill_shard` still decodes the reads table once per shard.** That's a change in the pod5 reader layer affecting all heads, not just CRF; it deserves its own PR. Also boxes `Basecaller::Cpu`, which fails `clippy -D warnings` under `--features crf-gpu` on main today — that feature combination isn't covered by CI.
jayhesselberth
added a commit
that referenced
this pull request
Aug 1, 2026
**Stacked on #174** (which is stacked on #173). Targets `worktree-crf-perf-punchlist`, so the diff here is one file. ## The gap The `features` job ran `cargo check`. The workflow sets `RUSTFLAGS: -Dwarnings`, but that only denies lints that are actually **evaluated** — and `cargo check` never evaluates a clippy lint. Meanwhile the `clippy` job builds default features, which don't include `crf-gpu`. So an opt-in feature could carry a clippy error indefinitely with every check green. Demonstrated on `main` as it stands right now: ``` cargo check -p escapepod-cli --features crf-gpu -> 0 errors (what CI ran) cargo clippy -p escapepod-cli --features crf-gpu -> 2 errors (what it missed) ``` That's `clippy::large_enum_variant` on the CLI's `Basecaller` enum, which #174 fixes. It has been failing silently for as long as `crf-gpu` has existed. ## The fix Switch those steps from `check` to `clippy` — clippy implies check, so no coverage is lost — and add `components: clippy` to the toolchain step, which the job didn't request. Then close the two holes the bug was sitting in: - **`crf-decode` and `crf-gpu` weren't exercised at all.** They pull tract-onnx, ort and fqxv-align, which is precisely the dependency-bump exposure this job exists to catch (per its own comment, that's how the cudarc and tract-onnx Dependabot PRs slipped through). - **Only the library crates were covered.** The CLI is a *different crate* from the one each feature flag names, and it's where the feature-gated command wiring lives, so `-p escapepod-demux` structurally cannot see it. `escapepod-cli` now gets `cnn-gpu` and `crf-gpu` alongside the existing `models-download`. Timeout 25 → 35 for the four extra steps. Still needs no CUDA toolkit, no `libonnxruntime` and no GPU — cudarc and ort both load dynamically. ## Verification All nine combinations run green locally on this branch: ``` escapepod-signal --features gpu ok escapepod-demux --features gpu ok escapepod-demux --features cnn-detect ok escapepod-demux --features cnn-gpu ok escapepod-cli --features cnn-gpu ok escapepod-demux --features crf-decode ok escapepod-demux --features crf-gpu ok escapepod-cli --features crf-gpu ok escapepod-cli --features models-download ok ``` And the new command is confirmed to fail on unfixed code (the `main` result above), so this is a guard that actually guards rather than one that happens to be green. ## Note I kept the existing style of one `cargo clippy` per feature rather than adding `--all-targets`. Linting tests/benches/examples under each opt-in feature would be additional real coverage — the `crf_basecall_check` example is feature-gated, for instance — but it roughly doubles the job's work, so it seemed worth deciding separately rather than folding in here.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Two independent findings from a performance audit of the CTC-CRF path added in #163/#164/#165/#172.
1.
produce_cpu_crfcould not fill the machineThe CRF producer fanned out with
.chunks(256)and then walked each chunk with a plain serialfor. That shape was copied fromproduce_cpu_gbm, but that head chunks for a reason that does not apply here:predict_manyis a genuinely batched kernel, whereas tract has no batched LSTM, so the CRF chunk only ever serialized its reads. The chunk existed solely to amortizeCrfScratch::new().Two consequences, both measured:
rna), one chunk is ~3.6 s of work that a starved worker cannot steal, so every block tail idles cores for that long.256 × threadsreads cannot fill the machine at all — 1000 reads produced exactly four tasks for 32 cores.Replaced with
for_each_init(CrfScratch::new, …), which gives per-read work stealing and movesCrfScratchfrom per-chunk to per-worker. This is the shapeproduce_cpuanddemux basecallalready use.(32 cores on
rna. The gap narrows as the chunk count catches up to the core count, which is exactly the predicted shape.)2.
basecall_batchretained every read's scoresCrfEncoderGpu::basecall_batchencoded the entire caller batch before decoding any of it.ESCAPEPOD_CRF_GPU_BATCH_ROWSbounds the device-side activations, not the host — the scores coming back aret_len * n_scorefloats, 1 MB per read for the RNA004 geometry, so an Arrow batch of a few thousand reads retained gigabytes of host RAM no matter what that knob was set to.Encode and decode now alternate one device batch at a time, capping host high-water at
batch_rowsreads' worth. Splitting is already exact and tested (run_batchdocuments why), and the chunk boundaries are unchanged, so the scores are identical.run_batchnow takes&[&[f32]]rather than&[Vec<f32>], which also drops the per-read 8 KB window clone on the way to the device.Note this does not overlap encode with decode — that would need a proper pipeline. It is a memory fix.
Verification
Full workspace suite (460 tests) and doctests pass;
cargo fmt --checkandclippy -D warningsclean on both default features andcrf-gpu.End-to-end against the nbc16 RNA004 CRF bundle at 1k and 10k reads:
--classificationsCSVs identical after sorting (row order is rayon-nondeterministic by design)Not addressed here
The audit turned up several other items left for follow-ups: the
split_time_majorcopy-out still runs while holding the session mutex;adc_to_paconverts the whole decoded prefix whenprepuses only the last 2000 samples;BLOCK_TARGET_BYTESis a per-filler bound rather than the global cap its comment implies; the router slot budget is only honored up to 192 barcodes; andbarcode.rs's WFA cap is set to the maximum possible score, so the early abandonment its docstring describes never actually fires.