La atención lineal en Trainium2 ya tiene un caso práctico documentado: i14 y AWS han entrenado y servido Laplace 3B, un modelo de unos 3.000 millones de parámetros, sobre chips Trainium2 con PyTorch nativo a través de TorchNeuron Native. Lo relevante para quien paga la factura es que el estado recurrente del modelo no crece con el contexto. Eso hace más predecible el coste de servir agentes de larga duración o de analizar documentos extensos. Hay una advertencia que conviene no saltarse: la calidad de recuperación se tiene que medir en cada caso de uso.
Qué ha pasado: Laplace 3B se entrena y se sirve en hardware de AWS
i14 ha llevado Laplace 3B a instancias Trainium2 de AWS. El modelo se basa en Gated Linear Attention (GLA), una variante de atención lineal, y se entrena con el optimizador Muon. La parte más concreta del trabajo es el kernel de GLA: en su versión original estaba escrito en Triton para GPU, y el equipo lo ha reescrito usando únicamente operadores estándar de PyTorch. El resultado corre sobre TorchNeuron Native, la ruta que permite ejecutar PyTorch en los aceleradores de AWS sin escribir kernels específicos para el chip.
Los números se midieron en una instancia trn2.48xlarge. En entrenamiento, el modelo alcanzó unos 120.000 tokens por segundo. En servicio, produjo 11.544 tokens de salida por segundo con 1.280 secuencias concurrentes. Cada secuencia mantiene un estado fijo de 29,5 MB, sin importar cuánto contexto haya acumulado. Con una aritmética sencilla, 1.280 secuencias a 29,5 MB suponen algo menos de 38 GB de estado total, y cada secuencia recibe de media algo más de 9 tokens de salida por segundo.
Para entender por qué esto importa hay que recordar cómo funciona la atención estándar de los Transformers. Cada token nuevo se compara con todos los anteriores, y el modelo guarda una caché de claves y valores (KV cache) que crece con la longitud del contexto. Cuanto más largo es el contexto y más usuarios simultáneos hay, más memoria consume esa caché. Ese es el motivo por el que servir contextos largos encarece la inferencia y la hace difícil de presupuestar.
Los modelos de atención lineal siguen otro camino. En lugar de conservar todo el historial, comprimen lo leído en un estado de tamaño fijo que se actualiza token a token. GLA añade puertas (gates) que controlan qué se retiene y qué se olvida. Es una familia de arquitecturas conocida en investigación, pero su adopción práctica ha chocado con un obstáculo: los kernels eficientes suelen estar escritos en Triton o CUDA, es decir, atados a GPU.
Implicaciones técnicas de la atención lineal en Trainium2
La primera lectura es de portabilidad. Reescribir el kernel GLA con operadores estándar de PyTorch significa que el código no depende de un lenguaje de kernels pensado para GPU. Quien ha intentado mover un modelo con kernels a medida entre aceleradores sabe que ahí se pierden semanas. Que el entrenamiento y el servicio funcionen con PyTorch nativo reduce esa fricción y deja el código más cerca de lo que ya mantiene un equipo de machine learning corriente. No hay en los datos disponibles una comparación directa con GPU, así que no se puede afirmar que sea más barato o más rápido que otra opción. Lo que sí se puede afirmar es que funciona y que las cifras se han publicado.
La segunda lectura es de memoria. Un estado fijo de 29,5 MB por secuencia convierte una pregunta difícil en un cálculo de hoja de cálculo: cuántas secuencias caben en la memoria disponible. Con una KV cache, la respuesta depende de la longitud de cada conversación y varía durante el servicio. Con estado fijo, la capacidad por instancia es estable, y eso facilita dimensionar el hardware, fijar precios por uso y planificar picos de demanda. Los 1.280 secuencias concurrentes de la prueba son el ejemplo directo de ese comportamiento.
La tercera lectura es el límite del enfoque. Comprimir todo el contexto en un estado fijo implica que parte de la información se pierde. Si la tarea exige recuperar un dato exacto de hace cientos de miles de tokens, como una cláusula concreta de un contrato o un identificador en un registro, un modelo de atención lineal puede rendir peor que uno con atención completa. El propio planteamiento del trabajo reconoce que hay que evaluar la calidad de recuperación caso por caso. Esa cautela es correcta y es la que debería guiar cualquier decisión de adopción.
Por último, está el optimizador. Muon es una alternativa a los optimizadores habituales para entrenar redes grandes, y que se haya usado en un entrenamiento real sobre Trainium2, a unos 120.000 tokens por segundo, aporta un dato de rendimiento sostenido. Con esa cifra, un día de entrenamiento procesa del orden de 10.000 millones de tokens, una referencia útil para estimar tiempos de un modelo de este tamaño.
Cómo pueden aplicar esto las empresas hoy
Conviene ser realistas: la mayoría de las PYMEs no entrena modelos de 3.000 millones de parámetros ni reserva instancias trn2.48xlarge. El valor práctico está en cómo evaluar proveedores y arquitecturas cuando el caso de uso es de contexto largo. Estas son las acciones concretas que se derivan de los datos.
Primero, calcular el coste de memoria por secuencia en vuestro caso actual. Si usáis un agente que mantiene conversaciones largas o procesa documentos extensos, preguntad al proveedor cómo escala la memoria con el contexto. Una arquitectura de estado fijo da una respuesta simple; una con KV cache, no. Segundo, montar una prueba de recuperación con vuestros propios documentos: preguntas cuyo dato esté al principio, en medio y al final de textos largos, y medir cuántas veces acierta el modelo. Sin esa prueba, un coste predecible no compensa si el modelo falla justo donde lo necesitáis.
Tercero, revisar la dependencia de kernels propietarios. Si vuestro equipo técnico mantiene código en Triton o CUDA, el caso de i14 indica que reescribir con operadores estándar de PyTorch es viable para GLA. Merece la pena valorar ese esfuerzo frente al ahorro posible, sin darlo por hecho. Y cuarto, evitar migraciones apresuradas: un modelo de 3B es pequeño frente a los grandes del mercado, y su calidad general en vuestras tareas hay que comprobarla antes de sustituir nada que ya funcione.
Análisis Blixel
El cuello de botella de la IA en producción rara vez es el modelo en sí; es la memoria y la previsibilidad del coste. Cualquier director técnico que haya intentado presupuestar un agente con contextos largos ha descubierto que la factura depende de variables que no controla. Por eso la atención lineal en Trainium2 merece atención, aunque Laplace 3B no vaya a sustituir mañana a los grandes modelos. Un estado de 29,5 MB por secuencia, constante, convierte el dimensionamiento en un cálculo y no en una apuesta.
También me parece valioso lo menos vistoso: que el kernel GLA funcione con operadores estándar de PyTorch. La dependencia de kernels hechos a mano para un tipo de chip es un coste oculto que pocas veces aparece en las comparativas de precio. Reducirla abre margen de negociación frente a los proveedores de computación.
Dicho esto, hay que contener el entusiasmo. Hay un modelo de 3.000 millones de parámetros, una sola configuración de hardware y ninguna comparación publicada con alternativas en GPU. El riesgo real es la recuperación: un estado fijo olvida, y en tareas legales, financieras o de soporte técnico un olvido cuesta dinero. Mi posición es que la arquitectura es una opción seria para casos de contexto largo y alta concurrencia, siempre que se pruebe con datos propios antes de comprometer presupuesto. Quien venda atención lineal como sustituta universal de la atención completa está simplificando demasiado.
¿Quieres aplicar esto en tu empresa? En Blixel.ai te ayudamos a integrar IA con sentido común. Hablemos.