data movement
on MI300X
speedup
inside the error bound
Genomic foundation models such as Google DeepMind's Enformer and AlphaGenome ingest sequence contexts spanning hundreds of kilobases to over a megabase. At those lengths the attention layer stops being a compute problem and becomes a data-movement problem: the key/value (KV) state dominates both memory footprint and the traffic between high-bandwidth memory (HBM) and on-chip SRAM. This paper measures that regime on NVIDIA H100 and AMD MI300X, shows attention there is bandwidth-bound rather than compute-bound on both, and closes the gap with compression applied inside the kernel.
The result is a single-source Triton attention kernel that fuses two complementary data-reduction techniques directly into the attention computation: per-block low-rank projection of K and V, and on-the-fly error-bounded quantization of the residual, with dequantization performed inside the kernel so the expanded KV state never round-trips to HBM. It reduces attention data movement by up to 4.1× and delivers 1.9× (H100) and 2.2× (MI300X) end-to-end inference speedups over a fused FlashAttention baseline, while holding downstream biological prediction quality inside a user-specified error bound (variant-effect AUROC degradation ≤ 0.004, expression-track Pearson correlation drop ≤ 0.6%).
Why long-context genomic attention is bandwidth-bound
Unlike autoregressive language models, genomic models run a single non-autoregressive forward pass per sequence. There is no incremental decoding and no KV cache growing across decode steps; instead the entire KV state for a very long context is present at once. That shifts the systems problem from cache management to one-shot bulk KV streaming, and makes the static size of the KV tensors the central concern.
A roofline analysis makes the point concrete. The ridge point (peak FP16 throughput divided by peak HBM bandwidth) sits at arithmetic intensity 295 on H100 and 247 on MI300X. The measured genomic operating points land at 31 and 24 respectively, far down the bandwidth-limited slope on both devices. When you are that far left of the ridge, reducing the bytes that represent K and V translates almost directly into reduced attention latency, which is precisely the regime where compression is most valuable.
KV traffic cut in the kernel
Fusing low-rank projection with error-bounded residual quantization reduces attention HBM traffic by up to 4.1× at N=65,536. Because the layer is bandwidth-bound, that reduction converts almost directly into latency.
Expanded KV never spills
Reconstruction is a fused prologue inside the per-block loop. K and V are rebuilt in SRAM, consumed by online softmax immediately, and the expanded form is never written back to global memory.
One source, no forks
The kernel is written once in Triton and compiles unmodified to H100 via CUDA and MI300X via ROCm, with no vendor-specific branches. Only block size is autotuned per device (128 vs. 256).
Quality inside a stated bound
At the default configuration the variant-effect AUROC drop stays at or below 0.004 and expression-track Pearson correlation falls by at most 0.6%, consistently across both accelerators.
How the compression works
K and V are partitioned along the sequence axis into blocks of size B. Each block is factored as a rank-r product, exploiting the fact that genomic KV state is empirically low-rank within local blocks because neighboring binned positions share regulatory context. The factors are computed once at model-load time, dropping per-block storage from Bd values to r(B + d).
Low rank alone is lossy in an unbounded way, so the residual is quantized with a block-wise scale chosen to satisfy a user-specified absolute error bound τ, stored at bq bits (default 4). The reconstruction guarantee holds exactly by construction on the compressed KV state; its propagation to attention logits is query-dependent, bounded via Hölder's inequality by a factor involving the query norm. Rather than claim an input-independent bound the paper does not have, it reports the empirical maximum query norm observed across the evaluation workloads, turning the logit perturbation into a concrete measured quantity. The scheme is error-bounded in the precise sense of an exact bound on compressed state plus a logit perturbation concrete under measured query norms, the same spirit as error-bounded scientific compressors applied here to attention state.
What the ablations say
The two components are genuinely complementary, and the ablation separates them cleanly. Low-rank projection alone delivers most of the data-movement reduction (3.6×) but carries no error guarantee, and shows larger, less predictable quality loss (AUROC −0.011, unbounded). Quantization alone bounds the error but reaches only 2.4× compression. The fused design attains 4.1× at AUROC −0.004, bounded. Each half adds error individually; the point is that both reduce bytes moved, and in a bandwidth-bound regime that trade is strongly favorable.
Sweeping rank r and residual bit-width bq traces a Pareto front with a visible knee. Settings of (r=16, bq=8) and (16, 4) stay within bound at 2.1× and 3.0×; the default (8, 4) sits at the knee with 4.1× and −0.004; pushing to (8, 2) or (4, 2) reaches 5.2× and 6.0× but breaches the error bound (−0.012, −0.031). The gains also widen with context length, as expected when the compression target is KV state whose size grows with N: 2.6× data-movement reduction and 1.4× end-to-end at N=4,096 rising to 4.1× and 2.2× at N=65,536.
Reconstruction is not free, but it is cheap in the right currency. It adds 6.2% of the attention FLOPs and 0.8 s of one-time factorization per model load, against a net −63% attention latency change. Because the kernel is bandwidth-bound, that added arithmetic is largely hidden behind the reduced memory traffic.
The thesis in one line: for non-autoregressive long-context models the attention layer is a data-movement problem, and the most effective place to apply compression is inside the kernel, where the compressed form can be streamed and expanded only transiently, under an error bound the practitioner chooses rather than validates after the fact.
Honest limitations
The low-rank factors are computed offline and held fixed; inputs whose KV state departs strongly from the calibration distribution may need larger ranks to stay within the bound. The evaluation covers single-device inference on two accelerators (H100 SXM 80 GB, MI300X 192 GB) across two models (Enformer at ~197 kbp, AlphaGenome up to 1 Mbp); multi-device sharding of compressed KV state is future work. And the error guarantee is exact on the compressed KV state, not on attention logits, where the bound is query-dependent and reported against measured rather than universal query norms.
Part of an ongoing line of work on portable, energy-aware GPU kernel acceleration for biological and medical AI, benchmarked across NVIDIA H100 and AMD MI300X. See more on the research page.