PagedAttention

PagedAttention, büyük dil modeli (LLM) sunumu için geliştirilmiş bir dikkat algoritması ve bellek yöneticisidir. Her isteğin KV önbelleğini, GPU belleğinde istenen yere yerleştirilen küçük, sabit boyutlu bloklarda saklar ve blokları işletim sisteminin sanal bellek sayfalarını eşlemesine benzer biçimde, istek başına bir blok tablosuyla izler. UC Berkeley'den Woosuk Kwon, Zhuohan Li ve çalışma arkadaşları, yöntemi vLLM sunum motoruyla birlikte SOSP 2023'te sundukları PagedAttention ile Büyük Dil Modeli Sunumunda Verimli Bellek Yönetimi başlıklı makalede tanıttı. Bitişik önbellek ayırmalarındaki parçalanmayı ortadan kaldırmak, bir GPU'nun aynı anda çok daha fazla dizi tutmasını sağlar. Makaleye göre aynı gecikmede verim, FasterTransformer ve Orca'ya kıyasla 2-4 kat daha yüksektir.

Çözdüğü Bellek Sorunu#

Metin üretimi sırasında bir transformer, önceki her belirtecin anahtar ve değer vektörlerini saklar; böylece bu vektörleri hiçbir zaman yeniden hesaplamaz. Bu önbellek, her katmanda belirteç başına bir kayıt eklenerek büyür; 13 milyar parametreli LLaMA modelinin tek bir dizisi 1,7 GB'a kadar yer kaplayabilir. İstek geldiğinde çıktı uzunluğu bilinmediğinden, önceki model sunum sistemleri her istek için olası en büyük çıktı uzunluğunu karşılayacak boyutta tek bir bitişik bellek bölgesi ayırıyordu. Bu bölgenin büyük bölümü boş kalıyordu ve bölgeler arasındaki boşluklar yeniden kullanılamayacak kadar küçüktü. vLLM ekibi, bu tür sistemlerin KV önbelleğine ayrılan belleğin %60-80'ini boşa harcadığını ölçtü; bu israf, hesaplama kaynakları tükenmeden çok önce yığın boyutunu sınırlıyordu.

Sayfalama Nasıl Çalışır#

PagedAttention önbelleği, her biri sabit sayıda belirtecin anahtarlarını ve değerlerini tutan bloklara böler; vLLM'de varsayılan değer 16'dır. İstek, ancak son bloğu dolduğunda ve nerede boş blok varsa oradan yeni bir blok alır; blok tablosu ise sıralamayı kaydeder. Dikkat mekanizması çekirdeği anahtarları ve değerleri blok blok toplamak için bu tabloyu izler; böylece bir dizinin önbelleğinin bitişik olması gerekmez. Boşa harcanan alan, her dizinin son bloğundaki doldurulmamış kısımla sınırlıdır; vLLM ölçümlerinde bu oran %4'ün altındadır.

Bu PyTorch taslağı, tek bir isteğin anahtarlarını paylaşılan bir havuzdaki dağınık bloklarda saklar ve blok tablosunu kullanarak geri okur:

import torch

block_size, head_dim = 16, 64
key_pool = torch.zeros(64, block_size, head_dim)  # physical blocks shared by every request
free_blocks = [41, 7, 19, 2, 33]  # free blocks are scattered across the pool

# A 40-token request needs ceil(40 / 16) = 3 blocks, wherever they happen to be free
keys = torch.randn(40, head_dim)
num_blocks = -(-len(keys) // block_size)
block_table = [free_blocks.pop(0) for _ in range(num_blocks)]
for pos, key in enumerate(keys):
    key_pool[block_table[pos // block_size], pos % block_size] = key

# The attention kernel reads the request's keys back through its block table
gathered = key_pool[block_table].reshape(-1, head_dim)[: len(keys)]
print(f"block table {block_table}, keys intact: {torch.equal(gathered, keys)}")
print(f"unused slots: {num_blocks * block_size - len(keys)} paged vs {2048 - len(keys)} with a 2,048-token reservation")

Blok tablosu [41, 7, 19] ve keys intact: True değerlerini yazdırır; ardından sayfalanmış istek için 8 boş yuva, 2.048 belirteçlik bir bölge ayıran yöntem içinse 2.008 boş yuva gösterir.

Diziler Arasında Blokları Paylaşma#

Blok tabloları ayrıca birkaç dizinin aynı fiziksel bloklara işaret etmesine de olanak tanır. Paralel örnekleme, ışın araması ve aynı sistem istemiyle başlayan istekler, her bloktaki başvuru sayacı aracılığıyla paylaşılan önekin tek kopyasını yeniden kullanır. Bir dizi, başkalarının hâlâ kullandığı bir bloğa yazmak istediğinde motor önce bloğu kopyalar; bu, işletim sistemlerinden alınan yazma sırasında kopyalama adımıdır.

vLLM yazarları, bu paylaşımın paralel örnekleme ve ışın aramasında bellek kullanımını %55'e kadar azalttığını ve verimi 2,2 kata kadar artırabildiğini bildiriyor. Tüm blokların karma değerini almak, bu fikri istekler arasında genişletir; vLLM istem önbelleğe almayı böyle uygular.

İlgili Teknikler#

PagedAttention, önbelleğe alınmış anahtarların ve değerlerin nerede tutulacağını belirler. Saklanması gereken veri miktarını veya verinin okunma hızını değiştiren yöntemlerle birlikte çalışır:

• Flash Attention: Tam puan matrisi GPU ana belleğine hiç yazılmayacak şekilde dikkat hesaplamasını döşer. Güncel motorlar, sayfalanmış bloklardan okuyan FlashAttention tarzı çekirdekler çalıştırır.
• Sürekli yığınlama: Başka bir istek tamamlanır tamamlanmaz yeni bir isteği kabul eder. Sayfalama, tamamlanan dizinin bloklarını hemen boş havuza geri verir; böylece yeni isteğe yer açılır.
• Gruplu sorgu dikkati: Anahtar ve değer başlıklarını sorgu başlıkları arasında paylaşarak her bloğun içeriğini küçültür.
• FP8 önbellek nicelemesi: Anahtarları ve değerleri 8 bit duyarlılıkla saklar; böylece 16 bit biçimlere kıyasla önbellek belleğini yaklaşık yarıya indirir.

Ultralytics LLM arayüzü, istek bağımsız değişkenlerini değiştirmeden iletir. Böylece n=3 içeren bir Chat Completions çağrısı, arayüzün base_url bileşeni üzerinden yerel bir vLLM sunucusuna gönderildiğinde, istemin önbelleğe alınmış bloklarını paylaşan üç örnek döndürür; yalnızca birbirinden ayrılan devamlar yeni bellek kullanır.