Sayfalı dikkat

Bu sayfa, kesintisiz toplu işleme (continuous batching) içinde kullanılan paged attention ileri fonksiyonunu belgelemektedir.

Değişken uzunluklu yol

flash_attn_varlen_func çekirdeği, değişken uzunluklu partileri işler. Bu yol, prefill aşamasında çok sayıda isteğe sahip partiler için önerilir.

Önbellek davranışı

Bu çekirdeğin, sayfalı önbellekle doğrudan etkileşime girmek için bir mekanizması yoktur; bu nedenle önbellek, ~PagedAttentionCache.update yöntemi kullanılarak manuel olarak okunur ve yazılır. Dizi uzunluğu büyüdükçe bu durum bir darboğaza dönüşebilir.

Dizinleme mekanizması

Çekirdek, her dizi için dikkat hesaplamasını gerçekleştirmek için maksimum dizi uzunluklarını (max_length_q, max_length_k) ve kümülatif dizi uzunluklarını (cu_seq_lens_q, cu_seq_lens_k) kullanır.

Örnek

Sorgu uzunlukları [10, 3, 1] ve anahtar uzunlukları [0, 1, 7] olan 3 diziden oluşan bir parti ele alalım:

cu_seq_lens_q = [0, 10, 13, 14]
cu_seq_lens_k = [0, 0, 1, 8]
max_length_q = 10
max_length_k = 7

Girdi şekilleri:

Q:      [1, 10+3+1, num_heads, head_dim] = [1, 14, num_heads, head_dim]
K or V: [1, 0+1+7, num_kv_heads, head_dim] = [1, 8, num_kv_heads, head_dim]

Çekirdek, kümülatif dizi uzunluklarını kullanarak her sorgu ve anahtar/değer token'ını bir diziye atar:

Q request index: [r0, r0, r0, r0, r0, r0, r0, r0, r0, r0, r1, r1, r1, r2]
cu_seq_lens_q:   0____________________________________10__________13__14

K request index: [r1, r2, r2, r2, r2, r2, r2, r2]  (r0 has 0 K tokens)
cu_seq_lens_k: 0,0_1_______________________8

Çözme yolu

flash_attn_with_kvcache çekirdeği, her dizinin tam olarak bir sorgu belirteci içerdiği yalnızca çözümlenme (decode) partilerini işler. Bu yol, değişken uzunluklu (varlen) yoldan daha verimlidir ancak ön doldurma (prefill) içeren partileri işleyemez.

Önbellek davranışı

Bu çekirdek, önbelleğe block_table kullanarak erişir ve önbelleği yerinde günceller. Blok tablosunun şekli (batch_size, max_blocks_per_seq)'dir; her satır, bir isteğin KV önbellek tensöründeki önbellek bloklarının fiziksel konumlarını içerir.

Dizinleme mekanizması

Çekirdek, her dizi için önbellek uzunluğunu almak için cache_seqlens kullanır. Her sorgu belirtecinin farklı bir diziye ait olduğunu varsayar (dizi başına bir belirteç).

Örnek

uzunlukları [1, 1, 1] ve uzunlukları [30, 32, 70] olan 3 diziden oluşan bir parti ele alalım. Önbellek bloğu boyutu 32'dir ve dizideki maksimum blok sayısı 4'tür.

Önbellek dizisi uzunlukları, basitçe anahtar uzunluklarıdır:

cache_seqlens = [30, 32, 70]

Blok tablosu şekli (3, 4). Örnek adresler kullanılarak:

block_table = [[2, -1, -1, -1],
               [0,  1, -1, -1],
               [3,  5,  6, -1]]

-1 değeri, atanmamış blokları belirtir.

• Dizi 0 (30 önbelleğe alınmış token): önbellek KV_cache[2] içinde. Yeni token sığar (30 + 1 = 31 < 32).
• Dizi 1 (32 önbelleğe alınmış jeton): KV_cache[0] ve KV_cache[1] içinde önbellek. 32 + 1 > 32 olduğu için ikinci bir blok gerekir.
• Dizi 2 (70 önbelleğe alınmış token): önbellek KV_cache[3], KV_cache[5] ve KV_cache[6] içinde. Blokların ardışık olmak zorunda olmadığını unutmayın; bu, sayfalı önbelleğin temel avantajıdır. Yeni token üçüncü bloğa sığar.