Error-Bounded Fused Attention Compression

Genomic foundation models such as DeepMind's Enformer and AlphaGenome read sequence contexts from hundreds of kilobases to over a megabase. At those lengths attention stops being a compute problem and becomes a data-movement problem: the key/value state dominates both memory footprint and traffic between HBM and on-chip SRAM. We measure that regime on NVIDIA H100 and AMD MI300X, show attention there is bandwidth-bound on both, and close the gap with compression applied inside the kernel.

A single-source Triton attention kernel fuses per-block low-rank projection of K and V with error-bounded quantization of the residual, and dequantizes inside the kernel so the expanded KV state never round-trips to HBM. It cuts attention data movement by up to 4.1× and gives 1.9× (H100) and 2.2× (MI300X) end-to-end speedups over a fused FlashAttention baseline, while holding biological prediction quality inside a user-specified bound: variant-effect AUROC drop ≤ 0.004 and expression-track Pearson drop ≤ 0.6%.

1Why long-context genomic attention is bandwidth-bound

Unlike autoregressive language models, genomic models run one non-autoregressive forward pass per sequence. There is no KV cache growing across decode steps; the entire KV state for a very long context exists at once. The problem shifts from cache management to one-shot bulk KV streaming, and the static size of K and V becomes the central concern.

2Compression inside the kernel

K and V are split along the sequence into blocks, and each block is factored as a rank-r product: genomic KV state is empirically low-rank within local blocks because neighbouring positions share regulatory context. The factors are computed once at model load, shrinking per-block storage from Bd values to r(B + d). Low rank alone is lossy without limit, so the residual is quantized with a block-wise scale chosen to meet a user-specified absolute error bound τ, at 4 bits by default.

The bound holds exactly on the compressed KV state. Its effect on attention logits depends on the query, so rather than claim an input-independent bound, we report the maximum query norm measured across the evaluation workloads, which makes the logit perturbation a concrete, measured quantity.

3What the ablations show

The two components are complementary. Low rank gives most of the reduction but no guarantee; quantization gives the guarantee but less reduction; the fused design gets both.

4Limitations

1. 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 bound.
2. Evaluation is single-device inference on H100 (80 GB) and MI300X (192 GB) across Enformer (~197 kbp) and AlphaGenome (up to 1 Mbp); sharding compressed KV across devices is future work.
3. The guarantee is exact on compressed KV state, not on attention logits, where the bound is query-dependent and reported against measured query norms.

More research

Part of a broader line of work at Embedded LLM on portable, fused GPU kernels for biological and clinical foundation models, measured on both NVIDIA and AMD.

• NeurIPS 2026Change the Metric, Change the Winner: Auditing How We Decide Which Clinical Foundation Model to Deploy in Cancer Pathology
• MOML 2026From 13 Hours to 4.6: Profiling-Driven Kernel Fusion for TensorNet in Molecular Simulation
• COLM 2026Deploying Clinical Language and Vision-Language Models Where the Data Lives
• ICML 2026From 805 ms to 23 ms: Accelerating State-Space Models for Real-Time ICU Monitoring
• ISCA 2026When the LLM-Tuned Stack Misses: An Infrastructure View of Biological Foundation Model Inference Across NVIDIA and AMD
• MLSys 2026BioTriton: Portable Cross-Vendor GPU Kernels for High-Throughput Bioinformatics via OpenAI Triton
• RECOMB 2026Hardware-Portable Fused GPU Kernels for High-Throughput Biological Foundation Models
• ISC HPC 2026Portable GPU Kernel Acceleration for Biological Foundation Models & Algorithms using OpenAI Triton

All research →