DeepSeek'in MLA'sı latent attention ile KV cache'i %93 sıkıştırıyor

dev.to'da yayımlanan teknik bir inceleme, DeepSeek-V2 ve DeepSeek-V3'te kullanılan attention mekanizması olan Multi-head Latent Attention'ın (MLA) key-value cache'i token başına 576 skalara küçültmesini — standart bir multi-head attention taban çizgisine kıyasla yaklaşık %93'lük bir azalma — ve bunu inference sırasında işler kılan matematiği adım adım açıklıyor.

Decode bir bellek problemidir

Makaleye göre transformer inference iki rejime ayrılır. Prefill tüm prompt'u tek seferde işler ve GPU compute'unu meşgul tutar. Token'ları tek tek üreten Decode ise, her yeni token için kendisinden önceki tüm token'ların cache'lenmiş key ve value'larını yeniden okumak zorundadır; aritmetik yoğunluğu aktarılan byte başına kabaca bir FLOP'a düşer. Bu yüzden Decode'u sınırlayan şey compute değil, HBM bant genişliğidir.

KV cache; katman sayısı, KV head sayısı, head boyutu, hassasiyet (precision), batch boyutu ve dizi uzunluğuyla birlikte büyür. Makaledeki token başına rakamlar şöyle: Standart multi-head attention ile bir DeepSeek 67B taban çizgisi FP16'da token başına yaklaşık 3,84 MB belleğe ihtiyaç duyar, Llama 3 70B'nin 8:1 grouped-query attention'ı yaklaşık 320 KB, DeepSeek'in MLA'sı ise yaklaşık 135 KB — FP8'de 67,5 KB. Bu, tek bir akış için 128.000 token'lık bir bağlama taşıdığında MHA taban çizgisi için kabaca 503 GB, GQA modeli için 40,96 GB ve MLA için 17,28 GB olur (FP8'de 8,64 GB).

Maliyeti somutlaştırmak için yazar, Llama 3 70B'yi 80 GB H100 üzerinde, 3,35 TB/s HBM3 bant genişliğiyle, 128k bağlamlar ve batch boyutu 8 ile modelliyor. Üretilen token başına 327,68 GB cache'i aktarmak yalnızca transfer süresi olarak yaklaşık 97,8 ms sürüyor ve üretimi saniyede yaklaşık 10,2 token'a sınırlarken Tensor Cores çoğunlukla boşta bekliyor. Tüm query head'ler arasında tek bir KV head'i paylaşan daha agresif çözüm olan Multi-query attention'ın ise GSM8k ve çok dokümanlı geri çağırma (recall) puanlarında %3,8 ile %6,2 arasında maliyet getirdiği bildiriliyor.

Tam key ve value'lar yerine bir latent vektör

MLA, her token'ın 5.120 boyutlu gizli durumunu (hidden state) 512 boyutlu sıkıştırılmış bir latent vektöre indirger. Tüm 128 head'in key ve value'ları, head başına up-projection matrisleri aracılığıyla bu latent'ten yeniden inşa edilir. Tam içerik aksi halde token başına 32.768 skalar kaplayacağından (her biri 128 boyutlu 128 head, hem key hem value için), latent içeriğin 64× sıkıştırılmasını temsil eder. Her head up-projection'ın kendi dilimini kullandığı için head'ler hâlâ bağımsız attention desenleri ifade edebilir ve MQA'yı zayıflatan temsil daralmasından kaçınılır.

Konumu içerikten ayırmak

Rotary position embeddings bir sorun yaratır. Makale, RoPE rotasyon matrisinin up-projection ile değişme özelliğine sahip olmadığını, bu yüzden yalnızca latent'i cache'leyen bir runtime'ın her decode adımında tüm geçmiş key'leri yeniden inşa etmek ve döndürmek zorunda kalacağını — bu da bant genişliği tasarrufunu tamamen yok edecekti — açıklıyor. DeepSeek'in cevabı decoupled RoPE'dur: Her key ve query, latent'ten türetilen 128 boyutlu bir içerik parçası ile tüm head'ler arasında paylaşılan 64 boyutlu bir konumsal parçaya bölünür. Attention skoru böylece bir içerik terimi ile bir konumsal terimin toplamı haline gelir.

Absorption, sıkıştırmayı açmayı tamamen ortadan kaldırıyor

İçerik key'leri latent içinde tamamen doğrusal olduğu için up-projection, her decode adımında bir kez query vektörüne katlanabilir ve attention dot product'ı doğrudan cache'lenmiş 512 boyutlu latent'ler üzerinde çalışabilir. Value up-projection'ı da benzer şekilde output projection ile birleştirilir. Token başına bellekte kalan şey, 512 boyutlu latent artı 64 boyutlu paylaşılan RoPE key'idir: 32 head'li bir MHA taban çizgisindeki 8.192'ye karşı 576 skalar — yaklaşık %93'lük rakamın kaynağı da budur.

Neden önemli

Uzun bağlamlarda decode hızı ve serving maliyeti, KV cache'in token başına kaç byte hareket ettirdiğine neredeyse tamamen bağlıdır. Makaledeki rakamlara göre MLA, 128k bağlamda akış başına GQA modelindeki 40,96 GB yerine 17,28 GB'a ihtiyaç duyar — bu fark, aynı donanımda doğrudan daha büyük batch'lere, daha uzun bağlamlara veya daha fazla eşzamanlı kullanıcıya dönüşür.

Uyarılar geçerli. Bu tek bir açıklayıcı makaledir ve karşılaştırmalarının bir kısmı varsayımsal taban çizgileridir: Llama 2 70B yalnızca varsayımsal olarak MHA olarak modellenmiştir ve %93'lük rakam, çoğu sağlayıcının gerçekte dağıttığı GQA modellerine değil, 32 head'li bir MHA cache'ine göre ölçülmüştür; GQA'da MLA'nın göreceli tasarrufu daha küçük ama yine de büyüktür. Ancak mekanizmanın kendisi yerleşik bir uygulamadır — MLA DeepSeek-V2 ve V3 ile birlikte geldi ve latent sıkıştırma ile matris absorption birleşimi, uzun bağlam serving ekonomilerinin GQA tabanlı akranlarından farklı olmasının önemli bir nedenidir.