- Moonshot AI vừa open-source FlashKDA — kernel CUTLASS C++ hiệu năng cao cho Kimi Delta Attention, đạt tốc độ prefill 1.72×–2.22× so với baseline Triton của flash-linear-attention trên H20, cắm thẳng vào FLA không đổi code model.
TL;DR
Moonshot AI vừa open-source FlashKDA — một implementation CUTLASS C++ cho kernel Kimi Delta Attention (KDA). Repo đạt 1.72×–2.22× prefill speedup trên GPU NVIDIA H20 so với kernel Triton baseline trong thư viện flash-linear-attention (FLA), và được thiết kế như drop-in backend: cắm vào là chạy, không cần sửa model code. License MIT, hỗ trợ Hopper SM90+, CUDA 12.9, PyTorch 2.4+.
What's new
KDA là thành phần attention tuyến tính cốt lõi của kiến trúc Kimi Linear mà Moonshot công bố cuối tháng 10/2025. Trước đây, bản mở duy nhất của KDA nằm trong flash-linear-attention dưới dạng kernel Triton. Triton dễ viết, dễ đọc, nhưng để vắt kiệt Tensor Core và async-copy của Hopper, các team hiệu năng cao vẫn phải quay về C++ template.
FlashKDA chính là bước đó: viết lại kernel KDA bằng NVIDIA CUTLASS — thư viện template CUDA mà FlashAttention-3 cũng dùng. Kết quả là một binding Python nhỏ gọn (flash_kda.fwd()) nhưng phía dưới là kernel C++ tối ưu từng warp, được compile riêng cho SM90+.
Vì sao đáng chú ý
Mấy tháng qua, mọi người nói nhiều về kết quả của Kimi Linear — 3.98× speedup trên RULER-128k, 6.3× TPOT ở 1M tokens, KV cache giảm tới 75% — nhưng các con số đó có được nhờ một lớp kernel rất cụ thể bên dưới. FlashKDA là lớp đó, giờ ai cũng xem được và sửa được.
Quan trọng hơn: H20 là dòng Hopper chủ lực ở thị trường Trung Quốc, và là phần cứng nhiều cụm inference sản xuất thực tế đang chạy. Tối ưu kernel trên đúng H20 không phải "demo trên H100 rồi claim", mà là tối ưu vào đúng chỗ traffic đang chảy.
Về mặt hệ sinh thái, đây còn là lần đầu có một reference CUTLASS C++ mở cho họ gated delta rule. Triton tốt cho prototype và nghiên cứu, nhưng khi muốn squeeze từng phần trăm hiệu năng cho production inference long-context, template C++ vẫn là công cụ chuẩn. FlashKDA vì thế hạ thấp rào cản cho các team khác muốn đi theo con đường đó, kể cả trên biến thể attention khác.
Số liệu kỹ thuật
- Prefill speedup: 1.72× – 2.22× vs flash-linear-attention Triton baseline (H20, BENCHMARK_H20.md).
- Ngôn ngữ: CUDA 56.4% + Python bindings 36.2%.
- Kiến trúc GPU: SM90+ (Hopper, Ada, Blackwell).
- CUDA / PyTorch: CUDA 12.9+, PyTorch 2.4+.
- Precision: bfloat16 + float32.
- Head dim: K = V = 128 (cùng cấu hình Kimi Linear).
- API:
flash_kda.fwd()— hỗ trợ batch, variable-length (cu_seqlens), và initial/final recurrent states cho inference stateful. - Gating: fine-grained channel-wise với 3 tham số học được —
A_log,dt_bias,lower_bound. - License: MIT.
So với FLA Triton
| Khía cạnh | FLA Triton KDA | FlashKDA (CUTLASS) |
|---|---|---|
| Ngôn ngữ kernel | Triton (Python DSL) | CUTLASS C++ templates |
| Prefill trên H20 | baseline 1.0× | 1.72× – 2.22× |
| Hỗ trợ GPU | SM80+ | SM90+ |
| Build | JIT compile, cài pip xong chạy | Cần clone + CUTLASS submodule + pip install -v . |
| Dễ đọc / sửa | Cao | Thấp hơn (C++ templates) |
| Drop-in trong FLA | Mặc định | Có — backend thay thế |
Use cases
- Phục vụ Kimi-Linear-48B-A3B (48B total, 3B activated, context 1M) trên cluster H20 — đây là target rõ ràng nhất.
- Workload long-context 128k–1M tokens, nơi prefill là phần tốn kém nhất — speedup 2× ở đây trực tiếp giảm chi phí hoặc tăng throughput tương ứng.
- Team nghiên cứu gated delta rule đã build trên FLA, muốn upgrade kernel mà không đổi model code.
- So sánh nghiên cứu giữa Triton vs CUTLASS cho cùng một thuật toán — FlashKDA là reference CUTLASS mở đầu tiên cho họ KDA.
Hạn chế & điều kiện
- Cố định
K = V = 128— không cấu hình được đầu (head) khác. - Không hỗ trợ Ampere / A100 — cần SM90+.
- Release này tập trung vào prefill; decode chưa được tuning riêng.
- Yêu cầu CUDA 12.9 + PyTorch 2.4 khá mới — môi trường production cũ sẽ cần nâng cấp toolchain.
- Build từ source, phụ thuộc CUTLASS submodule — thời gian build lần đầu không ngắn.
What's next
Những hướng đi tự nhiên: (1) tune thêm cho Blackwell SM10X — nơi tensor core mới thay đổi luật chơi; (2) thêm kernel decode để khép kín cả hai đầu prefill/decode; (3) có thể upstream trực tiếp vào fla-org/flash-linear-attention một khi API ổn định. Các dự án song song như inclusionAI/cuLA cho thấy cộng đồng đang chuyển dần các biến thể linear attention sang CUTLASS — FlashKDA là mảnh mở công khai đầu tiên từ chính team gốc của KDA.
Với K2.6 vừa ra ngày 20/04/2026 tận dụng các tối ưu attention phái sinh từ KDA, việc Moonshot mở luôn kernel bên dưới chỉ một ngày sau không phải trùng hợp — họ đang nhất quán chiến lược "mở từ model đến kernel".
Nguồn: github.com/MoonshotAI/FlashKDA, Kimi-Linear repo, Kimi Linear paper (arXiv 2510.26692), @Kimi_Moonshot announcement.
