Оптимизация инференса GLM-5.3-Flash на AWS Trainium2: анализ задержек и создание крупных ядер NKI

Авторы разработали и открыли движок инференса для модели GLM-5.3-Flash (320B MoE) на чипах AWS Trainium2, оптимизировав этап декодирования. Основная проблема заключалась в накладных расходах на передачу данных между HBM и процессорными ядрами, которые превышали время самих вычислений. Решение заключалось в объединении операций слоев в единые крупные ядра (NKI kernels), что позволило сократить время шага декодирования с 32.2 мс до 24.2 мс за счет эффективного использования локальной памяти и перекрытия коллективных операций.