From c01ed2091307827cc920bcc080be434d018223bf Mon Sep 17 00:00:00 2001
From: chansoo-calico <chansoo@calicolabs.com>
Date: Mon, 5 Oct 2026 14:25:43 +0000
Subject: [PATCH] Vectorize dna_1hot and keep the CUDA cache between SNP
 clusters

dna_1hot encoded each base in a Python loop, about 200 ms per 524 kb
window. hound_snp runs it in a prefetch thread, where it holds the GIL
while the eager model launches kernels. A lookup table takes about 10 ms
and gives the same output, including trimming, padding, n_uniform and
n_sample's random draws.

score_snps and score_gene_snps emptied the CUDA cache after every
cluster, so each cluster allocated its memory again. The caching
allocator already frees unused cached blocks and retries before it
reports out of memory.

Scores are bit-identical. Borzoi human on one A100: 1.02 to 1.22 SNPs/s
in float32, and 2.81 to 4.5 with -m float16 --compile.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/baskerville/dna.py  | 34 ++++++++++++++---------------
 src/baskerville/snps.py |  2 --
 src/tests/test_dna.py   | 48 ++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 64 insertions(+), 20 deletions(-)

diff --git a/src/baskerville/dna.py b/src/baskerville/dna.py
index a341bfb..24ae99a 100644
--- a/src/baskerville/dna.py
+++ b/src/baskerville/dna.py
@@ -71,27 +71,27 @@ def dna_1hot(
     else:
         seq_code = np.zeros((seq_len, 4), dtype="bool")
 
-    for i in range(seq_len):
-        if i >= seq_start and i - seq_start < len(seq):
-            nt = seq[i - seq_start]
-            if nt == "A":
-                seq_code[i, 0] = 1
-            elif nt == "C":
-                seq_code[i, 1] = 1
-            elif nt == "G":
-                seq_code[i, 2] = 1
-            elif nt == "T":
-                seq_code[i, 3] = 1
-            else:
-                if n_uniform:
-                    seq_code[i, :] = 0.25
-                elif n_sample:
-                    ni = random.randint(0, 3)
-                    seq_code[i, ni] = 1
+    seq_end = min(seq_len, seq_start + len(seq))
+    rows = np.arange(seq_start, seq_end)
+    cols = _NT_COLUMN[np.frombuffer(seq.encode("ascii", "replace"), np.uint8)]
+    cols = cols[: seq_end - seq_start]
+    known = cols >= 0
+    seq_code[rows[known], cols[known]] = 1
+
+    if n_uniform:
+        seq_code[rows[~known]] = 0.25
+    elif n_sample:
+        for i in rows[~known]:
+            seq_code[i, random.randint(0, 3)] = 1
 
     return seq_code
 
 
+# column of each ASCII byte in a 1-hot encoding; -1 for N and anything else
+_NT_COLUMN = np.full(256, -1, dtype=np.int8)
+_NT_COLUMN[np.frombuffer(b"ACGT", np.uint8)] = np.arange(4)
+
+
 def dna_1hot_index(seq: str, n_sample: bool = False):
     """Convert a DNA sequence to an index encoding.
 
diff --git a/src/baskerville/snps.py b/src/baskerville/snps.py
index 635f021..44f2d32 100644
--- a/src/baskerville/snps.py
+++ b/src/baskerville/snps.py
@@ -564,7 +564,6 @@ def hdf5_write_pair(cov_scores, gene_head_scores, score_index):
 
                 # clean reference
                 del rp_cov, ref_cov_preds
-                torch.cuda.empty_cache()
 
     # surface any cov-write exceptions (executor swallows them otherwise)
     for f in cov_write_futures:
@@ -903,7 +902,6 @@ def hdf5_write_pair(cov_scores, gene_head_scores, score_index):
                 # clean up memory
                 del rp_cov, ref_cov_preds
                 del ap_cov, alt_cov_preds
-                torch.cuda.empty_cache()
 
     # close open files
     genome_open.close()
diff --git a/src/tests/test_dna.py b/src/tests/test_dna.py
index 34210a2..6af89d6 100644
--- a/src/tests/test_dna.py
+++ b/src/tests/test_dna.py
@@ -2,9 +2,12 @@
 Tests for DNA one-hot utilities in baskerville.dna.
 """
 
+import random
+
 import numpy as np
+import pytest
 
-from baskerville.dna import hot1_rc
+from baskerville.dna import dna_1hot, hot1_rc
 
 
 def test_hot1_rc_known_value():
@@ -54,3 +57,46 @@ def test_hot1_rc_batched_matches_singleton():
     assert out.shape == batch.shape
     for i in range(batch.shape[0]):
         np.testing.assert_array_equal(out[i], hot1_rc(batch[i]))
+
+
+def _dna_1hot_per_base(seq, seq_len=None, n_uniform=False, n_sample=False):
+    """dna_1hot's original per-base loop, the reference for the vectorized one."""
+    if seq_len is None:
+        seq_len, seq_start = len(seq), 0
+    elif seq_len <= len(seq):
+        seq_trim = (len(seq) - seq_len) // 2
+        seq, seq_start = seq[seq_trim : seq_trim + seq_len], 0
+    else:
+        seq_start = (seq_len - len(seq)) // 2
+    seq = seq.upper()
+    seq_code = np.zeros((seq_len, 4), dtype="float16" if n_uniform else "bool")
+    for i in range(seq_len):
+        if i >= seq_start and i - seq_start < len(seq):
+            nt = seq[i - seq_start]
+            if nt in "ACGT":
+                seq_code[i, "ACGT".index(nt)] = 1
+            elif n_uniform:
+                seq_code[i, :] = 0.25
+            elif n_sample:
+                seq_code[i, random.randint(0, 3)] = 1
+    return seq_code
+
+
+def test_dna_1hot_known_value():
+    expected = [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1], [0, 0, 0, 0]]
+    np.testing.assert_array_equal(dna_1hot("ACgtN"), np.array(expected, dtype=bool))
+
+
+@pytest.mark.parametrize("seq_len", [None, 37, 64, 101])
+@pytest.mark.parametrize("mode", ["plain", "n_uniform", "n_sample"])
+def test_dna_1hot_matches_per_base_loop(seq_len, mode):
+    """Trimming, padding, N handling and n_sample's random draws all match."""
+    rng = np.random.default_rng(0)
+    seq = "".join(rng.choice(list("ACGTNacgtn"), 64))
+    kwargs = {"n_uniform": mode == "n_uniform", "n_sample": mode == "n_sample"}
+    random.seed(1)
+    expected = _dna_1hot_per_base(seq, seq_len, **kwargs)
+    random.seed(1)
+    got = dna_1hot(seq, seq_len, **kwargs)
+    assert got.dtype == expected.dtype
+    np.testing.assert_array_equal(got, expected)