Cómo i14 llevó Laplace 3B a AWS Trainium2: atención lineal y rendimiento nativo en PyTorch

Introducción

La mayoría de los grandes modelos de lenguaje usan atención softmax, cuyo comportamiento durante inferencia mantiene un caché de keys y values que crece con el contexto. Esto encarece y complica el servicio en producción, limita la concurrencia y obliga a algunas aplicaciones —por ejemplo agentes con memoria persistente— a sacrificar historial útil.

i14 propone una alternativa: Laplace 3B, un modelo con atención lineal que conserva la información previa en un estado recurrente de tamaño fijo (29.5 MB por secuencia). Trabajando con el equipo de AWS, i14 llevó esa arquitectura a Trainium2 usando soporte nativo de PyTorch con TorchNeuron Native, logrando mantener su diseño de Gated Linear Attention (GLA) y el optimizador Muon.

Por qué atención lineal importa

La atención softmax compara pares de tokens a lo largo del contexto, lo que durante el entrenamiento escala cuadráticamente con la longitud (O(T^2)). En inferencia, el KV cache crece con cada token nuevo, incrementando memoria por secuencia y afectando la capacidad de servir muchas conversaciones simultáneas.

La atención lineal reorganiza la computación alrededor de un estado recurrente que se actualiza token a token. Laplace 3B usa GLA para comprimir el contexto previo en una matriz de estado de tamaño fijo. Las implicaciones prácticas:

• Entrenamiento: pasa de O(T^2) a O(T) en el coste de atención.
• Inferencia: el estado recurrente permanece fijo (29.5 MB por secuencia), evitando el crecimiento del KV cache.

Esto hace la arquitectura atractiva para agentes de larga duración, análisis de documentos extensos y escenarios donde se necesita memoria persistente con requisitos previsibles de servicio. Cabe destacar que, al comprimir la historia en un estado fijo, es importante evaluar la calidad de recuperación para cada caso de uso.

Diseño esencial de Laplace 3B

Laplace 3B es un decodificador de 24 capas de i14 que reemplaza la atención softmax por Simple GLA en cada capa. Algunos detalles técnicos relevantes:

• Dimensión residual D = 3,072.
• Normalización y convolución: cada vector residual xt pasa por LayerNorm y una Conv1d causal depth-wise de ancho k = 4 con SiLU, entregando un resumen 4-gram antes del mixer.
• Proyección QKVG: una matriz fusionada Wqkvg de tamaño 9,216 × 3,072 (donde 9,216 = H × (2dk + 2dv)) se divide en queries qt, keys kt, values vt y una puerta gt. Parámetros: H = 16, dk = 96, dv = 192.
• GLA mixer: realiza la recuperación basada en contenido sobre la historia comprimida por el estado que decae.

En conjunto, la arquitectura mantiene un estado recurrente que no crece con la longitud del contexto y permite procesar secuencias largas sin aumentar la memoria por secuencia durante la inferencia.

Adaptaciones para AWS Trainium2

La versión original de Laplace 3B usaba kernels Triton optimizados para GPU. Para aprovechar Trainium2, el equipo de AWS GenAIIC y i14 hizo tres cambios principales:

1. Kernel GLA con operadores PyTorch nativos

• Sustituyeron el kernel Triton fused_chunk_simple_gla por operadores estándar de PyTorch.
• Implementaron 16 “chunks” de 128 tokens, alineando esa división con las 128 particiones del State Buffer (SBUF) por NeuronCore.
• Finalmente eligieron la versión beta de TorchNeuron Native (backend nativo PyTorch para AWS Trainium) por su mejor rendimiento frente a la reimplementación inicial.

2. Rediseño de la comunicación de Muon

• Ajustaron MuonFSDP2 para usar dos collectives all_to_all que redujeron el tráfico por rango en 24 veces y aumentaron el throughput en 1.63×. Este cambio optimiza la comunicación entre réplicas durante el entrenamiento distribuido.

3. Compilación del backward secuencial

• Para cumplir límites del compilador y memoria on-chip, desrollaron estáticamente los 16 pasos chunk de la recurrencia GLA. Esto permitió que el compilador trace el retroceso como un grafo fijo y que las transferencias a memoria on-chip se programen de antemano.

Por qué Trainium2 fue la opción

Un trn2.48xlarge combina 16 chips Trainium2, 64 NeuronCores lógicos, memoria de alta banda y NeuronLink-v3 para interconexión entre chips. Tres propiedades del hardware influyeron en la implementación:

• Memoria on-chip gestionada por software: cada NeuronCore dispone de un State Buffer (SBUF) de 28 MiB con 128 particiones y un buffer PSUM de 2 MiB conectado a una matriz sistólica 128×128. El SBUF no es una caché: el compilador debe programar todas las transferencias HBM→SBUF de forma anticipada.
• Interconexión chip-a-chip rápida: facilita la distribución eficiente del entrenamiento y la inferencia entre múltiples chips.
• Compilador estático: las rutinas recurrentes se trazan como un grafo fijo, por eso fue necesario el desrollado estático de la recurrencia GLA.

Resultados de entrenamiento e inferencia

Sin cambiar el diseño GLA ni el optimizador Muon, i14 obtuvo resultados operativos concretos en un trn2.48xlarge:

• Entrenamiento: alrededor de 120,000 tokens por segundo en una instancia trn2.48xlarge.
• Inferencia: 11,544 tokens de salida por segundo atendiendo 1,280 secuencias concurrentes.

Además de las métricas, las decisiones de ingeniería (PyTorch nativo, ajustes en Muon y desrollado estático) fueron clave para encajar la arquitectura en las restricciones del compilador y la memoria on-chip de Trainium2.

Implicaciones para empresas y despliegues en Latinoamérica

Para equipos y tomadores de decisión en América Latina, estas características tienen varias lecturas prácticas:

• Costos predecibles y mayor concurrencia: el estado fijo por secuencia (29.5 MB) permite estimar memoria y coste de serving con más seguridad frente al crecimiento indefinido del KV cache.
• Casos de uso pertinentes: asistentes conversacionales de larga duración, análisis de documentos legales o médicos extensos, y agentes con memoria persistente son candidatos naturales para esta arquitectura.
• Stack y operación: usar TorchNeuron Native facilita portar modelos PyTorch al hardware Trainium sin reescribir kernels en Triton, lo que reduce tiempo de ingeniería.
• Evaluación de calidad: como la historia se comprime en un estado fijo, conviene validar exhaustivamente la capacidad de recuperación y fidelidad de información para cada carga de trabajo (por ejemplo, cumplimiento legal o contexto médico).

En mercados latinoamericanos donde la optimización de costos y la escalabilidad son críticas, soluciones como Laplace 3B en hardware especializado pueden ofrecer ventajas competitivas si se gestionan correctamente la validación de resultados y la integración con sistemas de producción.

Conclusión

La adaptación de Laplace 3B a AWS Trainium2 muestra cómo una arquitectura de atención lineal puede traducirse en beneficios prácticos de entrenamiento y serving cuando se combina con soporte nativo de PyTorch y optimizaciones de comunicación y compilación. Para organizaciones que manejen contextos largos o requieran memoria persistente, este enfoque ofrece una alternativa atractiva a la atención softmax tradicional, con requisitos de memoria por secuencia fijos y tasas de procesamiento que ya son relevantes para producción.

Evaluar la calidad de recuperación del estado comprimido y planear la integración operacional son pasos clave antes de adoptar este tipo de arquitecturas en entornos productivos.