Speed up SNP scoring: vectorize dna_1hot, keep the CUDA cache between clusters

[chansoo-calico]

Why

Profiling hound_snp on Borzoi human (one A100, --rc --shifts 0) turned up two costs around the model. They grow as the model gets faster: with -m float16 --compile, they took about 40% of each SNP.

What

• dna_1hot uses a lookup table instead of a per-base Python loop: 200 ms down to 9 ms per 524 kb window. The loop runs in score_snps' prefetch thread and holds the GIL while the eager model launches kernels. The output is identical, including trimming, padding, n_uniform, and n_sample's random draws (same calls, same order).
• score_snps and score_gene_snps no longer empty the CUDA cache after every cluster. Each cluster had to allocate its memory again. On an allocation failure, the caching allocator already frees unused cached blocks and retries before raising out of memory.

Results

Trained Borzoi human fold 0 (releases/borzoi/download.sh), 200 random hg38 chr21 SNPs, A100 40GB, steady-state SNPs/s:

	main	this PR	scores
default (float32)	1.02	1.22	bit-identical
-m float16 --compile	2.81	4.5	bit-identical

score_gene_snps gets the same change but wasn't timed.

Related: --mix_dtype precision on trained weights

This is the kind of validation the new warning asks for. Same SNPs, measured against strict fp32 (TF32 off); the median |logSUM| is 0.10:

	SNPs/s	logSUM r	SNP ranking ρ (max |score|)	sign of top-10% effects
default (float32, TF32 convolutions)	1.02	0.99999	1.000	100%
float32 --compile + this PR	1.44	0.99999	1.000	100%
-m float16 --compile + this PR	4.5–4.75	0.996	0.990	100%
-m bfloat16 --compile + this PR	4.74	0.68	0.79	90%

bf16 noise exceeds a typical SNP effect, so the warning is right about it. fp16 may be acceptable, but that's your call. Two more notes:

• Judge precision on trained weights only. A random-weight Borzoi makes every reduced precision look far worse.
• --compile costs about 60 s of warm-up per process. On 200-SNP slices, that cancelled the float32 gain. It pays off from roughly 1,000–2,000 SNPs per job, or with a persistent TORCHINDUCTOR_CACHE_DIR.

Test plan

• pytest -m "not slow": 472 passed, 7 skipped (with bedtools installed, as in CI).
• New test_dna.py tests compare dna_1hot with the original loop across trim, pad, n_uniform and n_sample (seeded).
• ruff format --check .: clean.
• hound_snp on an A100, 200 SNPs: scores.h5 matches main bit for bit (logSUM, logD2 and their quantiles), for float32 and for -m float16 --compile.

🤖 Generated with Claude Code