NOSA: Atención Dispersa Nativa y Desmontable

NOSA: Native and Offloadable Sparse Attention

October 15, 2025

Autores: Yuxiang Huang, Chaojun Xiao, Xu Han, Zhiyuan Liu

cs.AI

Resumen

La atención dispersa entrenable ha surgido como una solución prometedora para abordar el cuello de botella en la eficiencia de decodificación de los modelos de lenguaje grandes (LLM) en el procesamiento de contextos largos, ahorrando significativamente accesos a la memoria mientras impacta mínimamente el rendimiento de las tareas. Sin embargo, los métodos existentes de atención dispersa dejan sin resolver una limitación crucial: el tamaño de la caché de clave-valor (KV) permanece sin reducir, lo que restringe los tamaños de los lotes en la GPU y limita el rendimiento de la decodificación, especialmente en la inferencia por lotes a gran escala. En este artículo, demostramos que la atención dispersa entrenable exhibe naturalmente una fuerte localidad en la selección de tokens a través de pasos de decodificación adyacentes, lo que permite la descarga de la caché KV sin alterar el cálculo subyacente de la atención. Sin embargo, la localidad inherente sigue siendo insuficiente para lograr una descarga eficiente, ya que la transferencia de pares KV seleccionados entre la CPU y la GPU continúa dominando el costo total de la decodificación. Basándonos en esta idea, presentamos NOSA, un marco de atención dispersa entrenable diseñado para admitir nativamente la descarga de la caché KV. NOSA introduce restricciones explícitas de localidad al descomponer la selección de tokens en componentes conscientes e inconscientes de la consulta, reduciendo así las transferencias de KV mientras se preserva el mismo cálculo de atención utilizado durante el entrenamiento. Preentrenamos un modelo de 1B parámetros con NOSA y realizamos extensas pruebas comparativas, demostrando que preserva un rendimiento casi sin pérdidas mientras logra una mejora de hasta 2.3 veces en el rendimiento de la decodificación en comparación con la línea base de atención dispersa entrenable estándar (InfLLM-V2).

English

Trainable sparse attention has emerged as a promising solution to address the decoding efficiency bottleneck of LLMs in long-context processing, significantly saving memory accesses while minimally impacting task performance. However, existing sparse attention methods leave a crucial limitation unresolved: the size of the key-value (KV) cache remains unreduced, which constrains on-GPU batch sizes and throttles decoding throughput, especially in large-scale batched inference. In this paper, we show that trainable sparse attention naturally exhibits strong locality in token selection across adjacent decoding steps, thereby enabling KV cache offloading without altering the underlying attention computation. However, the inherent locality remains insufficient to achieve efficient offloading, as the transfer of selected KV pairs between the CPU and GPU continues to dominate the overall decoding cost. Building on this insight, we present NOSA, a trainable sparse attention framework designed to natively support KV cache offloading. NOSA introduces explicit locality constraints by decomposing token selection into query-aware and query-agnostic components, thereby reducing KV transfers while preserving the same attention computation as used during training. We pretrain a 1B-parameter model with NOSA and conduct extensive benchmarks, showing that it preserves near-lossless performance while achieving up to a 2.3x improvement in decoding throughput compared with the vanilla trainable sparse attention baseline (InfLLM-V2).