Perhatian berhalaman

Halaman ini mendokumentasikan fungsi forward paged attention yang digunakan dalam continuous batching.

Jalur Varlen

Kernel flash_attn_varlen_func menangani batch dengan panjang variabel. Jalur ini direkomendasikan untuk batch dengan jumlah permintaan yang besar pada prefill.

Perilaku cache

Kernel ini tidak memiliki mekanisme untuk berinteraksi langsung dengan cache berhalaman, sehingga cache dibaca dan ditulis secara manual menggunakan metode ~PagedAttentionCache.update. Hal ini dapat menjadi bottleneck ketika panjang sekuens menjadi besar.

Mekanisme pengindeksan

Kernel menggunakan panjang sekuens maksimum (max_length_q, max_length_k) dan panjang sekuens kumulatif (cu_seq_lens_q, cu_seq_lens_k) untuk menghitung perhatian untuk setiap sekuens.

Contoh

Pertimbangkan sebuah batch dari 3 sekuens dengan panjang query [10, 3, 1] dan panjang key [0, 1, 7]:

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

Bentuk input:

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]

Kernel tersebut menetapkan setiap token query dan key/value ke dalam sebuah urutan menggunakan panjang kumulatif urutan:

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

Jalur decode

Kernel flash_attn_with_kvcache menangani batch yang hanya berisi decode, di mana setiap urutan memiliki tepat satu token query. Ini lebih efisien daripada jalur varlen, tetapi tidak dapat menangani batch yang berisi permintaan prefilling.

Perilaku cache

Kernel ini berinteraksi dengan paged cache menggunakan block_table untuk mengindeks cache dan memperbaruinya secara in-place. Bentuk block table adalah (batch_size, max_blocks_per_seq), di mana setiap baris berisi lokasi fisik dari blok cache sebuah request dalam tensor KV cache.

Mekanisme pengindeksan

Kernel menggunakan cache_seqlens untuk mengambil panjang cache untuk setiap sekuens. Kernel ini mengasumsikan setiap token query berasal dari sekuens yang berbeda (satu token per sekuens).

Contoh

Pertimbangkan sebuah batch berisi 3 sekuens dengan panjang query [1, 1, 1] dan panjang key [30, 32, 70]. Ukuran blok cache adalah 32 dan jumlah maksimum blok per sekuens adalah 4.

Panjang urutan cache hanyalah panjang kunci:

cache_seqlens = [30, 32, 70]

Bentuk tabel blok adalah (3, 4). Menggunakan alamat contoh:

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

Nilai -1 menunjukkan blok yang belum dialokasikan.

• Sequence 0 (30 token ter-cache): cache di KV_cache[2]. Token baru muat (30 + 1 = 31 < 32).
• Sequence 1 (32 token ter-cache): cache di KV_cache[0] dan KV_cache[1]. Blok kedua diperlukan karena 32 + 1 > 32.
• Sequence 2 (70 token ter-cache): cache di KV_cache[3], KV_cache[5], dan KV_cache[6]. Perhatikan bahwa blok tidak selalu berurutan, yang merupakan keunggulan utama dari paged cache. Token baru muat di blok ketiga.