Flash-MSA: Seyrek Dikkatle Milyon Token Eğitimini Hızlandırma

Flash-MSA, MiniMax Sparse Attention (MSA) için CuTeDSL ile yazılmış, Hopper ve Blackwell GPU'larında çalışan açık kaynaklı ilk performanslı eğitim çekirdeği paketi. Günümüzde birçok öncü model çıkarımı hızlandırmak için seyrek dikkat kullanıyor, ama bu mekanizmayı verimli şekilde eğitmek için paylaşılmış bir kod yoktu — bu proje o boşluğu dolduruyor. Geliştirme kiralık H100/B200 donanımlarında yapıldı ve MiniMax ile resmi bir bağlantısı bulunmuyor.

MSA, DeepSeek Sparse Attention'a (DSA) benzese de üç temel farkla öne çıkıyor: max-pooling ile seçilen 128 token'lık bloklar halinde seyreklik, ana dikkatte MLA yerine GQA kullanımı ve her proxy başlığının kendi KV bloklarını bağımsız seçebilmesi. GQA tercihi özellikle önemli çünkü batılı hiçbir laboratuvar eğitimde MLA benimsememiş durumda; bu da DeepSeek-V4 gibi MLA'ya uyumlu modellerdeki seyrek dikkat yaklaşımlarını doğrudan taşınamaz kılıyor.

Çekirdek tasarımı hesaplamayı mümkün olduğunca doğrusal tutuyor: proxy ileri geçişte hesaplanan blok indeksleri önbelleğe alınıp hem seyrek ana dikkatte hem de tüm geri yayılım boyunca yeniden kullanılıyor, böylece sadece proxy ileri geçiş bağlam uzunluğuna göre kuadratik kalıyor. Geri yayılımda proxy ve ana dikkat gradyanları birleştirilerek hesaplanıyor; KL diverjans kaybının cebirsel bir sadeleştirmesi (gradyan = proxy olasılığı − ana olasılık) sayesinde tam KL terimi hiç somutlaştırılmadan önemli kayıt/bellek tasarrufu sağlanıyor. Doğruluk, eager PyTorch referans uygulamasıyla kosinüs benzerliği karşılaştırmalarıyla doğrulandı.

Uzun bağlamlı eğitimle uğraşan mühendisler için bu çalışma, seyrek dikkat eğitimini çekirdek seviyesinde verimli şekilde nasıl uygulayacağına dair nadir bir örnek sunuyor; ayrıca MLA olmadan GQA tabanlı seyrek dikkat benimsemek isteyen ekipler için çalışan bir referans oluşturuyor.