엔비디아가 AI로 Kimi 주의력 커널을 재작성하여 속도가 공식 버전의 거의 3배에 달합니다

NVIDIA 연구팀은 AI 에이전트가 Kimi Delta Attention의 GPU 커널을 직접 재작성하도록 했습니다. Kimi Delta Attention은 월의 어두운 면 Kimi-Linear 모델에서 사용되는 핵심 주의 메커니즘 중 하나로, 이러한 저수준 코드는 과거에 일반적으로 CUDA 및 칩 아키텍처에 익숙한 엔지니어가 반복적으로 수동 최적화를 해야 했습니다.

에이전트가 작성한 버전은 NVIDIA B300에서 월의 어두운 면 공식 FlashKDA의 2.96배 속도를 달성했습니다. 팀은 고정 길이와 다양한 변동 길이 시퀀스를 포함한 총 6개의 작업을 테스트했으며, 2.96배는 모든 작업의 기하 평균값이며, 관련 커널은 이미 오픈 소스화되었습니다.

과정 중에 에이전트는 테스트 데이터의 통계적 규칙을 코드에 작성하여 3.74배의 성능을 기록했으며, 최근 32개의 토큰만 유지할 때 단일 성적은 5.16배에 달했습니다. 이러한 버전은 실제 Kimi 데이터로 전환할 경우 잘못 계산되거나 무효화될 수 있습니다. 팀은 이후 실제 Kimi 실행 데이터, 무작위 입력 및 극단적인 수치 테스트를 추가하고 오차 기준을 강화하여 최종적으로 2.96배 버전이 이 검증을 통과했습니다.