Phase 10 - Lesson 17

Native Sparse Attention (DeepSeek NSA)

This lesson includes a graded coding exercise that runs in your browser, unlocked with lifetime access.

Con 64k tokens, la atención consume entre el 70% y el 80% de la latencia de decodificación. Cada laboratorio de modelos abiertos tiene un plan para solucionarlo. NSA de DeepSeek (mejor artículo de ACL 2025) es el que se consolidó: tres ramas de atención paralelas — tokens comprimidos de granularidad gruesa, tokens de granularidad fina retenidos selectivamente y ventanas deslizantes para el contexto local — combinados a través de un gate aprendido. Está alineado con el hardware (compatible con el kernel), es nativamente entrenable (funciona en el preentrenamiento, no acoplado en la inferencia) y, en decodificaciones de 64k, se ejecuta más rápido que FlashAttention, al tiempo que iguala o supera la calidad de la atención completa. Esta lección construye las tres ramas de principio a fin y muestra por que la esparsidad es diferenciable de extremo a extremo.

Tipo: Build Lenguajes: Python (stdlib) Prerrequisitos: Phase 7 · 12 (KV cache, flash-attention), Phase 7 · 15 (attention variants), Phase 10 · 16 (differential attention) Tiempo: ~60 minutos

Objetivos de Aprendizaje

• Describir las tres ramas de atención de NSA y lo que cada una captura.
• Explicar por qué NSA es "nativamente entrenable" mientras que los métodos de atención esparsa anteriores eran solo para inferencia.
• Calcular el ahorro de cómputo de atención de NSA en comparación con la atención completa en un contexto de 64k como función del tamaño del bloque de compresión y del top-k de selección.
• Implementar la combinación de tres ramas en Python (biblioteca estándar) en una secuencia sintética corta y verificar el comportamiento de los pesos de la compuerta (gating).

El Problema

La atención completa en una longitud de secuencia N cuesta un tiempo O(N^2) y O(N) de caché KV por capa. Con 64k tokens, las cifras de cómputo y ancho de banda de memoria son catastróficas. Estimación teórica medida del artículo de NSA: la atención representa entre el 70% y el 80% de la latencia total de decodificación a 64k. Todo lo que viene después (downstream) — TTFT, tokens/seg, costo por millón de tokens — está dominado por el costo de la atención.

La atención esparsa es la respuesta obvia. Los intentos anteriores se dividen en dos categorías. La esparsidad de patrón fijo (ventana deslizante, espaciada/strided, bloque local) descarta información y falla en tareas de recuperación de largo alcance. La esparsidad en tiempo de inferencia (poda de caché KV, H2O, StreamingLLM) se aplica a un modelo preentrenado con atención densa y recupera solo una fracción de la aceleración potencial porque el modelo nunca fue entrenado para enrutar información a través del patrón esparso.

Native Sparse Attention (Yuan et al., DeepSeek + PKU + UW, mejor artículo de ACL 2025, arXiv:2502.11089) hace ambas cosas: un patrón de esparsidad que el modelo aprende durante el preentrenamiento, implementado como un algoritmo alineado con el hardware (kernel-aligned) que realmente ofrece ahorros de cómputo en la inferencia. Dentro de dos años, NSA o un descendiente directo será la atención predeterminada en todos los modelos de frontera de contexto largo.

El Concepto

Tres ramas paralelas

Para cada consulta (query), NSA ejecuta la atención tres veces, contra tres vistas diferentes de la caché KV:

1. Rama comprimida (Compressed branch). Los tokens se agrupan en bloques de tamaño l (típicamente 32 o 64). Cada bloque se comprime en un único token de resumen a través de una pequeña MLP aprendida. La consulta atiende a estos tokens comprimidos, obteniendo una vista de granularidad gruesa de toda la secuencia.

2. Rama seleccionada (Selected branch). Utilizando las puntuaciones de atención de la rama comprimida, se identifican los top-k bloques más relevantes para la consulta actual. Se leen los tokens de granularidad fina (no comprimidos) de esos bloques y la consulta atiene a todos ellos. Pense en la atención de la rama comprimida como la señal de enrutamiento para la selección.

3. Rama de ventana deslizante (Sliding-window branch). La consulta atiende a los W tokens más recientes (típicamente 512) para obtener contexto local. Esta rama captura los patrones de corto alcance basados en estructura (sintaxis, correferencia local) que las otras dos podrían omitir.

Las salidas de las tres ramas se combinan mediante un gate por posición aprendido:

out = g_cmp * out_cmp + g_sel * out_sel + g_win * out_win

g_cmp, g_sel, g_win son pesos de gate de una pequeña MLP en la consulta. No tienen que sumar 1 — pueden ponderar las ramas de forma independiente.

Por qué esto es "nativamente entrenable"

El paso de selección (bloques top-k) es discreto. Las operaciones discretas rompen el flujo de gradientes. Los trabajos anteriores de atención esparsa omitían la retropropagación (backpropagation) a través de la selección (limitando el entrenamiento) o utilizaban relajaciones continuas que no daban una esparsidad real en la inferencia.

NSA esquiva esto: la atención de la rama comprimida ES una atención de granularidad gruesa diferenciable sobre toda la secuencia. La operación top-k simplemente reutiliza las puntuaciones de atención más altas de la rama comprimida para elegir qué bloques de granularidad fina cargar. Los gradientes fluyen a través de las puntuaciones de la rama comprimida (que influyen tanto en la salida comprimida COMO en la lógica de selección), y la contribución de los bloques seleccionados a la salida final también es diferenciable. La operación top_k no diferenciable es una operación nula (no-op) en el grafo computacional de avance (forward) — solo controla qué bloques se cargan desde la memoria.

Esta es la razón por la que NSA se puede utilizar en el preentrenamiento de principio a fin. El modelo aprende a enrutar la información a través de las tres ramas de manera conjunta, generando un patrón esparso que en la inferencia realmente ofrece la aceleración prometida.

Kernel alineado con el hardware

El kernel de NSA está diseñado para las jerarquías de memoria de las GPU modernas. El kernel carga las consultas por grupos de GQA (bucle externo), obtiene los bloques KV esparcidos correspondientes por grupo (bucle interno) y ejecuta la atención en SRAM. Dado que cada grupo de consultas ve los mismos bloques seleccionados (la selección es por grupo de consultas, no por cabeza de consulta), las cargas de KV se amortizan en todo el grupo. La intensidad aritmética se mantiene alta.

El artículo informa que los kernels de Triton se ejecutan 9 veces más rápido que FlashAttention en decodificaciones de 64k, y la relación de aceleración aumenta con la longitud de la secuencia. Se proporcionan kernels tanto de avance (forward) como de retroceso (backward).

El presupuesto de cómputo

Sea N la longitud de la secuencia, l el tamaño del bloque de compresión, k la cantidad de selección top-k, w la ventana deslizante y b el tamaño del bloque seleccionado (típicamente igual a l).

• Rama comprimida: O(N/l) llaves por consulta, total de O(N * N / l).
• Rama seleccionada: O(k * b) llaves por consulta, total de O(N * k * b).
• Rama deslizante: O(w) llaves por consulta, total de O(N * w).

Total: O(N * (N/l + k*b + w)).

Con N = 64k, l = 64, k = 16, b = 64, w = 512: el costo por consulta es 1000 + 1024 + 512 = 2536 llaves. La atención completa es de 64000 llaves. Una reducción de cómputo de 25x.

Con N = 128k, l = 64, k = 16, b = 64, w = 512: el costo por consulta es 2000 + 1024 + 512 = 3536 llaves. La atención completa es de 128000 llaves. Reducción de 36x. El beneficio crece con la longitud de la secuencia, que es todo el objetivo.

Cómo se compara

MoBA (Moonshot, arXiv:2502.13189) se publicó simultáneamente y adopta un enfoque similar de "tres es mejor que uno", aplicando el principio de MoE a los bloques de atención. NSA y MoBA son las dos arquitecturas a conocer para el preentrenamiento de contexto largo en 2026.

Constrúyelo

El archivo code/main.py implementa las tres ramas en una secuencia sintética corta y muestra:

• La MLP de compresión (se utiliza una línea base simple de promedio de agrupación (mean-pool) para mayor claridad pedagógica; la NSA real utiliza una MLP aprendida).
• La selección de bloques top-k impulsada por las puntuaciones de la rama comprimida.
• La atención de ventana deslizante en los últimos w tokens.
• La combinación por gates (gated).
• Una impresión del conteo de cómputo en comparación con la atención completa.

Paso 1: comprimir tokens en bloques

def compress(K, l):
    n = len(K)
    n_blocks = (n + l - 1) // l
    out = []
    for b in range(n_blocks):
        start, end = b * l, min((b + 1) * l, n)
        block = K[start:end]
        summary = [sum(row[d] for row in block) / len(block) for d in range(len(K[0]))]
        out.append(summary)
    return out

Paso 2: atención de la rama comprimida

Ejecute la atención softmax de la consulta contra las llaves comprimidas. Las puntuaciones de la rama comprimida sirven también como la señal para la selección top-k.

Paso 3: selección de bloques top-k

Elija los índices de los k bloques comprimidos con mayor puntuación. Cargue los tokens originales no comprimidos de esos bloques y ejecute la atención sobre ellos.

Paso 4: atención de ventana deslizante

Tome los últimos w tokens y ejecute la atención estándar contra ellos.

Paso 5: gate + combinar

Una pequeña MLP en la consulta produce tres pesos de gate. La salida final es una suma ponderada de las salidas de las tres ramas.

Paso 6: conteo de cómputo

Imprima el número de llaves atendidas por consulta para cada rama y el total. Compare con N (atención completa). En un caso sintético de 1024 tokens con l = 32, k = 4, w = 128, NSA ve 32 + 128 + 128 = 288 llaves por consulta frente a 1024 para la atención completa — 3.5 veces menos.

Úsalo

NSA se está implementando en el propio flujo de trabajo de preentrenamiento de contexto largo de DeepSeek. Estado de integración en arquitecturas de inferencia públicas a partir de abril de 2026:

• DeepSeek interno: nativo, los pesos publicados usan NSA o su sucesor DSA (Deepseek Sparse Attention).
• vLLM: soporte experimental de NSA en desarrollo para los pesos de DeepSeek-V3.x.
• SGLang: benchmarks de NSA publicados; la ruta de producción sigue a vLLM.
• llama.cpp / CPU: no soportado; la sobrecarga (overhead) de la descomposición del kernel no vale la pena con el rendimiento de la CPU.

Cuándo recurrir a NSA:

• Ejecución de preentrenamiento o entrenamiento continuo orientado a un contexto de más de 64k con un presupuesto de cómputo serio.
• Inferencia de los propios checkpoints de contexto largo de DeepSeek. Los pesos son nativos de NSA.

Cuándo no hacerlo:

• Servir un modelo preentrenado de atención densa existente. No se puede adaptar NSA sin un entrenamiento continuo.
• Contexto inferior a 16k. La sobrecarga de las tres ramas domina el ahorro.
• Chat interactivo con lote 1 (Batch-1). La decodificación sensible a la latencia se beneficia, pero solo en contextos largos.

Envíalo

Esta lección produce outputs/skill-nsa-integrator.md. Dada una especificación de ejecución de preentrenamiento de contexto largo, produce un plan de integración de NSA: tamaño del bloque de compresión, top-k, ventana deslizante, ancho de la MLP del gate, elección del kernel y las evaluaciones específicas de contexto largo que justificarían el cambio de arquitectura.

Ejercicios

1. Ejecute code/main.py en una secuencia sintética de 1024 tokens. Realice un barrido (sweep) de (l, k, w) a través de tres configuraciones predeterminadas e imprima los conteos de cómputo. Identifique la configuración que logra el menor conteo de llaves por consulta mientras mantiene un 95% de recuperación (recall) frente a la atención completa en una prueba del tipo "agua en el pajar" (needle-in-a-haystack).

2. Reemplace el compresor de promedio de agrupación (mean-pool) por una pequeña MLP aprendida (2 capas, 32 unidades ocultas). Entrénela en una tarea sintética donde la señal sea el promedio de un bloque. Mida la diferencia (gap) de perplejidad contra la línea base de mean-pool en datos de validación.

3. Implemente la MLP del gate. Toma la consulta como entrada y genera tres escalares. Muestre que el gate se comporta de manera sensata: ponderación casi uniforme en consultas aleatorias, peso pesado en la rama seleccionada cuando la consulta golpea un bloque muy anterior.

4. Calcule el presupuesto de memoria de la caché KV para un modelo 70B habilitado con NSA a un contexto de 128k. Las cabezas KV son 8, la dimensión de la cabeza es 128, BF16. Compare con la atención completa y con MLA (la Fase 10 · 14 mostró los números de MLA). Identifique la longitud de secuencia donde la caché KV de la rama de granularidad fina de NSA es igual a la atención completa.

5. Lea la Sección 4 del artículo de NSA (arXiv:2502.11089) y explique en tres oraciones por qué las puntuaciones de atención de la rama comprimida se reutilizan para la selección top-k en lugar de calcular una puntuación de enrutamiento separada. Vincule la respuesta al flujo de gradientes.

Términos Clave

Lecturas Adicionales

• Yuan et al. — Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention (arXiv:2502.11089, ACL 2025 Best Paper) — el artículo
• DeepSeek-V3 Technical Report (arXiv:2412.19437) — la familia de arquitectura a la que se dirige NSA
• Moonshot AI — MoBA: Mixture of Block Attention for Long-Context LLMs (arXiv:2502.13189) — trabajo concurrente, atención al estilo MoE sobre bloques
• Beltagy et al. — Longformer: The Long-Document Transformer (arXiv:2004.05150) — los orígenes de la ventana deslizante
• Xiao et al. — StreamingLLM: Efficient Streaming Language Models with Attention Sinks (arXiv:2309.17453) — línea base de esparsidad en tiempo de inferencia que NSA mejora
• Dao et al. — FlashAttention-2 (arXiv:2307.08691) — la línea base de atención completa que los kernels de NSA superan a 64k