# Paged attention

이 페이지는 [continuous batching](./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의 핵심 장점입니다. 새로운 토큰은 세 번째 블록에 들어갑니다.