Paged attention

이 페이지는 continuous batching에서 사용되는 paged attention forward 함수를 문서화합니다.

Varlen 경로

flash_attn_varlen_func 커널은 가변 길이 배치(batch)를 처리합니다. prefill 단계에서 요청 수가 많은 배치의 경우 이 경로를 사용하는 것이 권장됩니다.

캐시 동작

이 커널은 paged cache와 직접 상호작용하는 메커니즘이 없으므로, ~PagedAttentionCache.update 메서드를 사용하여 캐시를 수동으로 읽고 씁니다. 시퀀스 길이가 커지면 이것이 병목 현상이 될 수 있습니다.

인덱싱 메커니즘

커널은 각 시퀀스에 대한 attention을 계산하기 위해 최대 시퀀스 길이(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 커널은 각 시퀀스가 정확히 하나의 query 토큰을 가지는 decode 전용 배치만 처리합니다. varlen 경로보다 효율적이지만, prefilling 요청이 포함된 배치는 처리할 수 없습니다.

캐시 동작

이 커널은 block_table을 사용하여 paged cache에 접근하고, 이를 인덱싱한 뒤 인플레이스로 업데이트합니다. 블록 테이블의 shape은 (batch_size, max_blocks_per_seq)이며, 각 행에는 해당 요청의 KV cache 텐서 내 캐시 블록의 물리적 위치가 포함됩니다.

인덱싱 메커니즘

커널은 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 값은 할당되지 않은 블록을 나타냅니다.

• Sequence 0(캐시된 토큰 30개): KV_cache[2]에 캐시됨. 새 토큰이 들어갈 수 있음(30 + 1 = 31 < 32).
• Sequence 1(캐시된 토큰 32개): KV_cache[0]와 KV_cache[1]에 캐시됩니다. 32 + 1 > 32이므로 두 번째 블록이 필요합니다.
• Sequence 2 (캐시된 토큰 70개): KV_cache[3], KV_cache[5], KV_cache[6]에 캐시됨. 블록이 반드시 연속적이지 않을 수 있으며, 이것이 paged cache의 핵심 장점입니다. 새로운 토큰은 세 번째 블록에 들어갑니다.