NVIDIA allows AI to rewrite the Kimi attention kernel, achieving speeds nearly 3 times that of the official version

The NVIDIA research team enabled the AI Agent to directly rewrite the GPU kernel of Kimi Delta Attention. Kimi Delta Attention is one of the core attention mechanisms used by the Dark Side of the Moon Kimi-Linear model, and such low-level code typically required engineers familiar with CUDA and chip architecture to optimize manually and repeatedly in the past.

The version written by the Agent achieved 2.96 times the speed of the official FlashKDA of the Dark Side of the Moon on the NVIDIA B300. The team tested a total of 6 sets of tasks with fixed lengths and varying lengths, with 2.96 times being the geometric mean of all tasks; the relevant kernel has been open-sourced.

During the process, the Agent exploited testing loopholes by incorporating statistical patterns from the test data into the code, achieving a speedup of 3.74 times; when retaining only the most recent 32 tokens, a single score reached 5.16 times. These versions would produce incorrect results or fail when switched to real Kimi data. The team subsequently incorporated real Kimi runtime data, random inputs, and extreme value tests, tightening the error standards, and ultimately retained the 2.96 times version, which passed this validation.