¿Cómo será la arquitectura de LLM de próxima generación? Esta pregunta sigue generando debates, y el colaborador y desarrollador de Zhihu, Yuxuan, ofrece una comparación aguda entre DeepSeek Sparse Attention (DSA) y Native Sparse Attention (NSA), además de una mirada práctica a la implementación de operadores DSA con TileLang. 🚀 Por qué DSA > NSA (en tareas de contexto largo): A partir de experimentos que añaden DSA a modelos pequeños y comparan con NSA, DSA consistentemente rinde mejor, principalmente debido a dos decisiones de diseño clave: 1️⃣ Destilación de Attn-Score → supervisión explícita para la selección de índices 2️⃣ Escasez a nivel de token en lugar de a nivel de bloque → recuperación más precisa y detallada 🔍 1) Destilación de Attn-Score La atención escasa se basa en seleccionar los pares clave-valor correctos. DSA supervisa directamente el módulo de índice utilizando verdaderos puntajes de atención, alineando el entrenamiento con el objetivo real: "seleccionar los tokens críticos." NSA, en cambio, optimiza solo la pérdida del LM, sin proporcionar ninguna restricción explícita sobre la precisión del índice, lo que explica su rendimiento más débil en los benchmarks de recuperación de documentos largos. 🔍 2) Escasez a nivel de token vs a nivel de bloque La precisión escala con el presupuesto de cómputo: un indexado más preciso → mejor recuperación. El indexado a nivel de token (DSA) naturalmente produce una mayor fidelidad que a nivel de bloque (NSA). Desde esta perspectiva, el cuello de botella en el rendimiento de NSA es esperado: una pregunta interesante: ¿Ayudaría un tamaño de bloque=8 a que NSA se pusiera al día con DSA? ⚙️ El verdadero desafío: entrenar DSA de manera eficiente El entrenamiento de DSA implica Calentamiento → Ajuste fino escaso. El desafío: calcular y almacenar los puntajes de atención de ambas ramas. Una implementación ingenua requiere O(n²) de almacenamiento, lo que anula los ahorros de memoria de FlashAttention. Incluso el pre-filtrado (k=2048, h=512+64) aún demanda grandes búferes. 📎 Código: 🧩 Fusión de Kernels para el rescate (Fig 1) Para evitar almacenar enormes puntajes de atención intermedios, DSA utiliza kernels fusionados. Un truco clave es fusionar Index-Score + Top-k en un solo kernel: • Mantener un búfer de 2K • Calcular Index-Score para cada bloque...