Inferact, fundada por el núcleo del equipo de vLLM, publicó una serie de pruebas de inferencia: al ejecutar el modelo Kimi K3 de Moonshot AI en 16 chips Google TPU v7 Ironwood, la velocidad de decodificación de un solo flujo alcanzó 709 tokens por segundo, mientras que el grupo de control con 16 chips Nvidia GB200 llegó a 452, un 57% más rápido aproximadamente.
El informe de pruebas y el código se publicaron juntos el 23 de septiembre, y el repositorio inferact/tpu-megakernels se abrió bajo licencia Apache 2.0. Inferact es una startup formada por el equipo original de vLLM, que antes había levantado una ronda semilla de 150 millones de dólares con una valoración de 800 millones, liderada por a16z y Lightspeed.
Sobre el papel, el GB200 sigue teniendo ventaja
Primero, las cifras sobre el papel. Según el informe, un chip TPU v7 tiene un pico de BF16 de 2,31 PFLOPS y de FP8 de 4,61 PFLOPS, con 206 GB de HBM y un ancho de banda de 7380 GB/s; un chip GB200 equivalente ofrece 2,5 PFLOPS, 5 PFLOPS, 186 GB de HBM y 8000 GB/s de ancho de banda. En cómputo y ancho de banda, el GB200 supera en ambos casos, mientras que el TPU solo gana en capacidad de memoria, con 20 GB más.
Kimi K3 es un modelo MoE de 92 capas, cuya parte de atención combina Kimi Delta Attention y Multi-head Latent Attention (MLA). La prueba usa paralelismo de tensores en 4 vías y paralelismo de expertos en 8 vías, repartidos en 16 chips.
Comparación de velocidad bruta sin decodificación especulativa:
| Tamaño de lote | TPU v7 (megakernel) | GB200 (vLLM) |
|---|---|---|
| 1 | 249 | 127 |
| 2 | 392 | 227 |
| 4 | 515 | 373 |
| 8 | 865 | 636 |
La unidad son tokens por segundo. Con un lote de tamaño 1, el TPU es casi el doble de rápido; con un lote de 8, la ventaja se reduce a poco más del 30%. Al activar la decodificación especulativa DSpark propuesta por DeepSeek, la velocidad de un solo flujo pasa a 709 frente a 452, con unos 8,5 milisegundos por paso de decodificación.
La clave: fusionar 92 capas en un solo programa
El método de Inferact se llama megakernel. En los marcos de inferencia convencionales, el cálculo de cada capa se divide en varios núcleos (kernels) independientes que se lanzan uno tras otro, con huecos entre ellos; los pesos solo empiezan a transferirse de la memoria al chip cuando el kernel correspondiente se pone en marcha.
Inferact usa Pallas para escribir las 92 capas de un paso de decodificación como un único programa. Según el informe, un megakernel elimina las fronteras entre kernels: la carga de pesos ya no está atada al kernel que los usa y puede anticiparse tanto como lo permita la memoria interna del chip. En la decodificación de un solo flujo, el chip pasa la mayor parte del tiempo esperando datos, y ahí es justo donde la precarga compensa; cuando el lote crece, la proporción de cómputo aumenta y la ganancia de esta técnica se reduce, lo cual coincide con la tendencia de la tabla.
Un beneficio adicional es el tiempo de compilación: antes, compilar con XLA tardaba más de 30 minutos; con megakernel, tarda menos de 90 segundos. Para demostrar que la mejora de velocidad no sacrifica la calidad de los resultados, el informe incluye las puntuaciones de Kimi K3 corriendo en TPU: 94,4% en GPQA-Diamond y 97,2% en GSM8K.
Un cálculo aproximado de la eficiencia por unidad de cómputo
Usando los resultados de decodificación especulativa de un solo flujo para un cálculo aproximado: el pico de BF16 del TPU v7 equivale a cerca del 92% del GB200, pero la tasa de tokens generada es 1,57 veces mayor; convertido a eficiencia por unidad de pico de cómputo, el TPU queda en torno a 1,7 veces el GB200. Convertido por ancho de banda, el resultado también ronda las 1,7 veces.
Esta cifra hay que tomarla con reservas. En el lado del GB200 se usó la receta de Kimi K3 publicada abiertamente por vLLM, un enfoque convencional de múltiples kernels, sin una optimización de megakernel de intensidad equivalente. El informe no incluye un comparativo igualmente optimizado en GB200, así que no se puede separar cuánto de la diferencia viene del hardware y cuánto del esfuerzo de software. Nvidia tampoco ha respondido sobre estas cifras.
Para Moonshot AI, esta prueba ofrece una vía viable para desplegar Kimi K3 en hardware fuera de Nvidia. Kimi K3 es un modelo de pesos abiertos con 2,8 billones de parámetros, un tamaño que eleva bastante el umbral para un despliegue propio; con el kernel de TPU de código abierto, alquilar TPU en la nube para ejecutar K3 se convierte en una opción más.
Fuentes: blog técnico de Inferact, CocoLoop, repositorio inferact/tpu-megakernels, QbitAI; las cifras de rendimiento siguen la metodología de prueba de 16 contra 16 chips del informe de Inferact, y la información de financiación proviene de reportes públicos.