PISA: atención bloqueada dispersa con complejidad log-lineal

El mecanismo de atención bloqueada dispersa (block-sparse attention) ofrece una alternativa eficiente al costo cuadrático de la atención auto en modelos de lenguaje de gran escala, pero su implementación práctica se ve limitada por la complejidad computacional en la selección de bloques. El artículo presenta PISA, un mecanismo que resuelve este cuello de botella mediante una estrategia de selección Top-K piramidal. A diferencia de los métodos convencionales que requieren evaluar todos los pares de bloques de consulta, PISA construye una jerarquía de claves de grano grueso a fino. Se aplica una puntuación LogSumExp a un conjunto de candidatos limitado en cada nivel para reducir progresivamente la búsqueda hasta alcanzar el nivel más fino. Esta estructura permite alcanzar una complejidad total de O(N log N), donde N es la longitud de la secuencia, en lugar del cuadrático O(N²). Además, los autores desarrollan kernels de Triton optimizados para hardware, que fusionan la enrutamiento jerárquico y la puntuación LogSumExp sin materializar la matriz de puntuación de consultas-claves. La evaluación en tareas de modelado de lenguaje demuestra que PISA logra un rendimiento comparable en benchmarks de razonamiento de sentido común, superando a la línea base en tareas de recuperación de información.