Modelo de atención lineal: i14 sirve 1.280 conversaciones a la vez con solo 29,5 MB de memoria por secuencia

Cuanto más larga es la conversación con un chatbot, más memoria cuesta atenderla. Ese es el problema silencioso de la IA en producción, y la startup i14 lo ataca con un modelo de atención lineal llamado Laplace 3B. Su estado de memoria se queda en 29,5 MB por secuencia, sin importar cuánto contexto acumule. Según el blog de AWS, en una sola instancia trn2.48xlarge el modelo entrenó a unos 120.000 tokens por segundo y sirvió 11.544 tokens de salida por segundo con 1.280 secuencias concurrentes.

Qué hizo i14 para dejar atrás el KV cache en AWS Trainium2

La mayoría de los grandes modelos de lenguaje usan atención softmax, y durante la inferencia su KV cache crece con cada token. Eso reduce la concurrencia, encarece el servicio y obliga a las aplicaciones con agentes a descartar historial útil. Laplace 3B, un decoder de 24 capas, reemplaza esa atención por Gated Linear Attention (GLA). Comprime el contexto previo en una matriz de tamaño fijo. El costo de entrenar la atención pasa de O(T²) a O(T), y la memoria de inferencia deja de depender de la longitud del contexto.

Llevarlo a AWS Trainium2 exigió ingeniería concreta. El equipo del AWS Generative AI Innovation Center primero reescribió el kernel GLA con operadores estándar de PyTorch. Después optó por la versión beta de TorchNeuron Native, el backend nativo de PyTorch para Trainium, porque rendía mejor. Así conservó la arquitectura GLA y el optimizador Muon. Los tres cambios clave fueron:

Kernel GLA en PyTorch nativo: sustituyó el kernel Triton pensado para GPU por 16 bloques de 128 tokens, ajustados a las 128 particiones del buffer de cada NeuronCore.

Rediseño de la comunicación de Muon: dos colectivas all_to_all en MuonFSDP2 redujeron el tráfico por rank 24 veces y subieron el rendimiento 1,63 veces.

Compilación del paso inverso secuencial: desenrollaron estáticamente los 16 pasos de la recurrencia para caber en los límites de grafo y memoria del compilador.

Hay un matiz que el propio AWS reconoce: la atención lineal comprime el historial en lugar de guardarlo íntegro, así que la calidad de recuperación debe validarse en cada caso de uso. Remio, un medio que analizó el anuncio, insiste en lo mismo. Los tokens por segundo miden capacidad de la infraestructura, no utilidad de la respuesta. Para esa publicación, la pregunta sin resolver es si el estado fijo conserva lo importante. Como referencia, cita a Qwen3-Next, que combina procesamiento recurrente con atención convencional.

Cómo evaluar un modelo de atención lineal en tu empresa

Haciendo una cuenta propia, 1.280 secuencias por 29,5 MB suman unos 38 GB de estado total. Es una cifra predecible, y eso cambia cómo se planifica la capacidad: ya no depende de que algunos clientes mantengan sesiones larguísimas. Dividir 11.544 tokens por 1.280 secuencias da además unos 9 tokens por segundo por conversación, un dato útil para dimensionar la experiencia de usuario. Si estás valorando esta línea, estos pasos ayudan:

1. Identifica dónde duele el KV cache. Agentes de larga duración, análisis de documentos extensos y asistentes con memoria persistente son los candidatos naturales.
2. Mide calidad antes que velocidad. Diseña pruebas de recuperación con tus datos: ¿recuerda restricciones del turno 3 en el turno 300? Un marco de evaluación sólido, como el que describe el caso de Elastic con su framework único para evaluar agentes, es el punto de partida.
3. Compara costo por tarea resuelta, no por token. Un sistema más rápido que olvida detalles genera trabajo de seguimiento. El ahorro real se parece más al del patrón Local-First que redujo 75% el costo de IA empresarial que a un simple benchmark.
4. Calcula el costo de portabilidad. El kernel original estaba escrito para GPU y tuvo que reescribirse. Presupuesta ingeniería de compilador, no solo horas de instancia.
5. Piensa en híbridos. Si el estado fijo no basta, una arquitectura que mezcle capas recurrentes y atención clásica puede ser el término medio.

El contexto de hardware también importa. AWS publicó casi a la vez una guía para entrenar Qwen3 en Trainium con torchtitan, con el mismo lanzador torchrun en tres rutas distintas. Es una señal de que el soporte de PyTorch nativo en Trainium madura. Aun así, TorchNeuron Native sigue en beta, y las empresas que dependan de él deben tenerlo en cuenta en su hoja de ruta. Quienes ya exploran la infraestructura para IA agéntica encontrarán aquí otra palanca: reducir lo que cada conversación ocupa en memoria, además de conseguir más cómputo.

El momento de probar la atención lineal

No hace falta migrar nada mañana. Conviene, en cambio, abrir un piloto acotado: elegir un flujo con contextos largos, medir cuánto del costo de servir proviene de la memoria y comparar. Si tu carga de trabajo sufre con el KV cache, un modelo de atención lineal merece una prueba seria, siempre que la calidad de recuperación supere tu umbral. ¿Sabes cuánto de tu factura de inferencia se va hoy en recordar conversaciones largas?