Moonshot, Kimi Delta Attention용 FlashKDA 오픈소스 공개 (H20에서 최대 2.22배 성능 향상)
Moonshot open-sourced FlashKDA, CUTLASS kernels for Kimi Delta Attention, up to 2.22x over the Triton baseline on H20
핵심 요약
Moonshot이 Kimi Delta Attention을 위한 고성능 CUTLASS 커널인 FlashKDA를 오픈소스로 공개했습니다.
- FlashKDA 공개 — Kimi Delta Attention을 위한 CUTLASS C++ 커널 구현체임
- 성능 향상 — H20 GPU 환경에서 기존 Triton 대비 최대 2.22배 빠른 속도 제공
- 호환성 확보 — flash-linear-attention 백엔드로 연동되어 기존 모델에 즉시 적용 가능
- 기술적 한계 — 현재는 추론 전용이며 SM90 이상, CUDA 12.9 이상의 환경이 필요함
github.com/MoonshotAI/FlashKDA
이번 주에 다양한 라우팅 레이어가 K2.6을 어떻게 처리하는지 비교하다가, Moonshot이 K2.6 활동과 함께 공개한 FlashKDA를 발견했습니다. 모델 릴리스와는 별개로 커널 작업 자체가 상당히 흥미로워서 공유합니다.
이것은 무엇인가. Kimi Linear 논문에서 제시된 선형 어텐션 변형인 Kimi Delta Attention의 포워드 커널을 CUTLASS C++로 구현한 것입니다. FLA 풀 리퀘스트 #852를 통해 flash-linear-attention의 백엔드로 연결되므로, 이미 KDA 기반 모델에 FLA를 사용 중이라면 백엔드 레이어에서 FlashKDA로 라우팅할 수 있습니다.
H20 벤치마크에서 측정한 FLA의 기존 Triton 경로 대비 수치입니다:
T=8192, H=96, D=128, 고정 길이 시퀀스: 1.72배.
가변 길이 및 혼합 시퀀스 길이: 1.95배.
가변 길이 및 균일한 1024x8: 2.22배.
이것이 중요한 이유. KDA와 같은 선형 어텐션 아키텍처는 시퀀스 길이에 따른 선형 확장을 보장하지만, 커널 구현이 실제로 하드웨어 효율적일 때만 그 약속이 유효합니다. FLA의 Triton 경로는 레퍼런스로서 잘 작동하지만, Hopper 메모리 액세스 패턴에 맞게 튜닝된 CUTLASS는 이론적 비용 모델과 실제 GPU 성능 사이의 간극을 메워줍니다.
요구 사항은 SM90 이상, CUDA 12.9 이상, PyTorch 2.4 이상입니다. MIT 라이선스입니다.
한 가지 솔직한 한계점은, 벤치마크가 포워드 패스 전용이며 모든 수치가 H20에서 측정되었다는 것입니다. H20은 중국 전용 Hopper 변형이므로 H100이나 Blackwell에서의 절대 수치는 다를 것입니다. 상대적인 속도 향상은 방향성 면에서 비슷하겠지만, 아직 그 수치를 게시한 사람은 없습니다.
혹시 H100에서 테스트해 본 분이 계신지, 아니면 백워드 패스 커널이 언제쯤 나올지 의견이 궁금합니다. 현재는 포워드 전용이라 학습 용도로는 한계가 있습니다.
