perf: reuse Triton's autotune result across Cherimoya calls (2.7x on every prediction after the first)

[jmschrei]

Summary

Cherimoya's inference kernels are wrapped in triton.autotune, which benchmarks its candidate configs the first time a kernel is launched for a given shape. That's the right call for training, where a process compiles once and then runs for hours. It's badly mismatched to Chorus, where run_code_in_environment spawns a fresh subprocess per call — so the benchmark is re-run every single time, to serve one forward pass.

On CATv1's native 2114 bp geometry the first forward costs ~7.6 s against ~0.96 ms steady-state. A one-shot prediction was spending ~99.99% of its GPU time deciding how to run the kernel rather than running it.

Triton 3.6+ can persist the selected config to disk, keyed on the same tuple autotune keys on. It's off by default, so this turns it on.

This is a cache, not a guess. The config used is the one autotune itself benchmarked and picked for that exact key; a shape Chorus hasn't seen falls back to benchmarking. Predictions are bit-identical.

Measurements

1× H200, CATv1 DNASE:K562, through Chorus's own API with use_environment=True, medians of 3. Baseline measured on this same base by reverting only these files:

	before	after	speedup
_predict	11.53 s	4.27 s	2.70×
predict_variant_effect	22.95 s	8.41 s	2.73×
first forward, in-process	7.6 s	0.40 s	19×
steady-state forward	0.960 ms	0.956 ms	unchanged

Raw: before _predict 11.42 / 11.55 / 11.53, variant 22.95 / 22.90 / 22.95. After _predict 4.32 / 4.21 / 4.27, variant 8.41 / 8.41 / 8.40.

Prediction checksum is identical before and after: 341.3908. I confirmed the baseline writes zero autotune cache files, so the two arms are genuinely independent.

What this does not do

The first call on a machine is not faster — it still pays the full ~11.5 s and populates the cache. Only later processes benefit. Since Chorus's normal mode is many short-lived subprocess calls, that's the right trade, but a single-call benchmark would show no improvement at all.

The cache is keyed on shape, so a new batch size pays one benchmark before being cached (the background builder at batch 64/256 has its own entry). A cleared ~/.triton, or a Triton / GPU / cherimoya version bump, resets it — all three are part of Triton's cache key.

Cherimoya itself is unmodified.

Implementation notes

The knob has no backing environment variable (knobs.autotuning.cache has env=None), so it can't go in EnvironmentRunner._prepare_env — Autotuner.__init__ reads it when the decorators are evaluated, which happens at import cherimoya. It therefore has to be set from Python before that import, on each of the three load paths.

• chorus/oracles/cherimoya_source/_triton_autotune.py — new; canonical helper and the full rationale
• chorus/oracles/cherimoya.py — _load_direct, which also covers scripts/build_backgrounds_cherimoya.py and the in-process tests
• both templates — inlined rather than importing the helper, because importing chorus inside the cherimoya env costs ~0.6 s and would eat much of the saving; each carries a pointer to the helper's explanation
• guarded with try/except (ImportError, AttributeError) so CPU-only and macOS installs (no Triton) are unaffected — a False return is a missed optimization, never a functional difference

Tests

tests/test_cherimoya.py: 50 passed (46 existing + 4 new). Integration: 9 passed, 1 skipped — unchanged.

Three of the new tests assert the knob precedes the cherimoya import on each load path. That ordering is load-bearing and fails silently — set it after the import and every call quietly goes back to re-benchmarking, with no error to notice.

🤖 Generated with Claude Code