La cara oculta de la luna lanza FlashKDA de código abierto, Kimi Linear aumenta la velocidad de inferencia de 1.7 a 2.2 veces

Noticias ME News, 22 de abril (UTC+8), según la monitorización de Dongcha Beating, Moon’s Dark Side ha abierto en GitHub FlashKDA, una herramienta diseñada específicamente para acelerar la inferencia de modelos en tarjetas gráficas Nvidia Hopper (H100, H20, etc.), bajo la licencia MIT. Su objetivo es KDA, la nueva mecanismo de atención propuesta el año pasado en el artículo Kimi Linear. Cuando los grandes modelos leen textos largos, el cálculo de la atención tradicional crece con el cuadrado de la longitud, mientras que la atención lineal reduce ese coste a un crecimiento lineal, y KDA es una versión mejorada en esa línea. La estructura del modelo Kimi Linear consiste en 3 capas de KDA alternadas con 1 capa de atención tradicional. Antes, ya existía una versión escrita en Triton, alojada en la biblioteca de código abierto flash-linear-attention (abreviada fla). FlashKDA ha sido reescrita usando la biblioteca GPU de bajo nivel de Nvidia, CUTLASS, específicamente para aprovechar al máximo el rendimiento de las tarjetas Hopper. Las pruebas oficiales en H20 muestran que, en una misma pasada hacia adelante, FlashKDA es de 1.7 a 2.2 veces más rápida que la versión en Triton, con mejoras especialmente notorias en escenarios con longitudes de entrada variables y procesamiento en lotes. Sin embargo, la comparación oficial solo se hizo con la versión en Triton, sin compararla con otras soluciones de atención lineal. Esta vez, solo se ha abierto el código para la inferencia hacia adelante, lo que significa que solo se puede “ejecutar el modelo” (inferencia), pero no “entrenar el modelo”; el entrenamiento sigue requiriendo la versión original en Triton. Requisitos de uso: la tarjeta gráfica debe ser Hopper o posterior (arquitectura SM90 en adelante), CUDA 12.9 o superior, PyTorch 2.4 o superior. FlashKDA también se ha integrado como un nuevo backend en la rama principal de fla (PR #852), y los usuarios existentes solo necesitan cambiar una línea en la configuración para migrar. (Fuente: BlockBeats)