ページド・アテンション

このページでは、continuous batchingで使用される paged attention の forward 関数について説明します。

Varlen パス

flash_attn_varlen_func カーネルは、可変長のバッチを処理します。プリフィルにおいてリクエスト数が非常に多いバッチには、このパスが推奨されます。

キャッシュの動作

このカーネルには、ページドキャッシュと直接やり取りする仕組みがないため、~PagedAttentionCache.update メソッドを使用してキャッシュを手動で読み書きする必要があります。シーケンス長が大きくなると、これがボトルネックになり得ます。

インデックス機構

カーネルは、最大シーケンス長（max_length_q、max_length_k）と累積シーケンス長（cu_seq_lens_q、cu_seq_lens_k）を使用して、各シーケンスの注意機構を計算します。

例

クエリ長が[10, 3, 1]、キー長が[0, 1, 7]である3つのシーケンスのバッチを想定します。

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

入力形状:

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]

カーネルは、累積シーケンス長を使用して、各クエリおよびキー/バリューのトークンをシーケンスに割り当てます。

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

デコードパス

flash_attn_with_kvcache カーネルは、各シーケンスがちょうど1つのクエリトークルを持つ、デコード専用のバッチを処理します。これはvarlenパスよりも効率的ですが、プリフィルリクエストを含むバッチは処理できません。

キャッシュの動作

このカーネルは、block_table を用いてページドキャッシュにアクセスし、キャッシュをインデックスしてインプレースで更新します。ブロックテーブルの形状は (batch_size, max_blocks_per_seq) であり、各行には、リクエストのキャッシュブロックが KV キャッシュテンソル内のどの物理位置にあるかが含まれています。

インデックス機構

このカーネルはcache_seqlensを使用して、各シーケンスのキャッシュ長を取得します。各クエリトークンは異なるシーケンスに属すると仮定しています（シーケンスごとに1トークン）。

例

クエリ長が[1, 1, 1]、キー長が[30, 32, 70]の3つのシーケンスからなるバッチを想定します。キャッシュのブロックサイズは32、シーケンスあたりの最大ブロック数は4です。

キャッシュのシーケンス長は、単にキーの長さに等しくなります:

cache_seqlens = [30, 32, 70]

ブロックテーブルの形状は(3, 4)です。例として、アドレスを使用します:

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

-1 の値は、未割り当てのブロックを示します。

• シーケンス 0（キャッシュ済みトークン 30 個）：KV_cache[2] にキャッシュ。新しいトークンは収まります（30 + 1 = 31 < 32）。
• シーケンス 1（32トークンがキャッシュ済み）：KV_cache[0] と KV_cache[1] にキャッシュされています。32 + 1 > 32 であるため、第2のブロックが必要です。
• シーケンス 2（キャッシュ済みトークン 70 個）: KV_cache[3]、KV_cache[5]、KV_cache[6] にキャッシュされています。ブロックは必ずしも連続しているとは限らない点に注意してください。これがページドキャッシュの大きな利点です。新しいトークンは 3 番目のブロックに収まります。