Chú ý phân trang

Trang này tài liệu hóa hàm forward của paged attention được sử dụng trong continuous batching.

Đường dẫn Varlen

Lõi flash_attn_varlen_func xử lý các lô có độ dài biến đổi. Đường dẫn này được khuyến nghị cho các lô có số lượng yêu cầu lớn trong giai đoạn prefill.

Hành vi bộ nhớ đệm

Kernel này không có cơ chế nào để tương tác trực tiếp với bộ nhớ đệm phân trang, do đó bộ nhớ đệm được đọc và ghi thủ công bằng phương thức ~PagedAttentionCache.update. Điều này có thể trở thành điểm nghẽn khi độ dài chuỗi tăng lên lớn.

Cơ chế lập chỉ mục

Kernel sử dụng độ dài chuỗi tối đa (max_length_q, max_length_k) và độ dài chuỗi tích lũy (cu_seq_lens_q, cu_seq_lens_k) để tính toán attention cho từng chuỗi.

Ví dụ

Hãy xem xét một lô gồm 3 chuỗi với độ dài truy vấn [10, 3, 1] và độ dài khóa [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

Hình dạng đầu vào:

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 gán mỗi token truy vấn và khóa/giá trị cho một chuỗi bằng cách sử dụng độ dài chuỗi tích lũy:

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

Đường dẫn giải mã

Lõi flash_attn_with_kvcache xử lý các lô chỉ giải mã (decode-only) trong đó mỗi chuỗi có đúng một token truy vấn. Cách này hiệu quả hơn đường dẫn varlen nhưng không thể xử lý các lô có yêu cầu prefilling.

Hành vi bộ nhớ đệm

Kernel này tương tác với bộ nhớ đệm phân trang bằng cách sử dụng block_table để chỉ mục vào bộ nhớ đệm và cập nhật tại chỗ. Bảng khối có hình dạng (batch_size, max_blocks_per_seq), trong đó mỗi hàng chứa các vị trí vật lý của các khối bộ nhớ đệm của một yêu cầu trong tensor bộ nhớ đệm KV.

Cơ chế lập chỉ mục

Kernel sử dụng cache_seqlens để lấy độ dài cache cho từng chuỗi. Nó giả định rằng mỗi token truy vấn thuộc về một chuỗi khác nhau (một token cho mỗi chuỗi).

Ví dụ

Hãy xem xét một lô gồm 3 chuỗi có độ dài truy vấn [1, 1, 1] và độ dài khóa [30, 32, 70]. Kích thước khối cache là 32 và số lượng khối tối đa cho mỗi chuỗi là 4.

Chiều dài chuỗi cache đơn giản là chiều dài của các key:

cache_seqlens = [30, 32, 70]

Hình dạng của bảng khối là (3, 4). Sử dụng các địa chỉ ví dụ:

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

Các giá trị -1 chỉ các khối chưa được cấp phát.

• Chuỗi 0 (30 token đã cache): cache trong KV_cache[2]. Token mới vừa đủ (30 + 1 = 31 < 32).
• Chuỗi 1 (32 token đã cache): cache trong KV_cache[0] và KV_cache[1]. Cần một khối thứ hai vì 32 + 1 > 32.
• Chuỗi 2 (70 token đã cache): cache trong KV_cache[3], KV_cache[5] và KV_cache[6]. Lưu ý rằng các khối không nhất thiết phải liên tục, đây là lợi thế chính của cache phân trang. Token mới nằm trong khối thứ ba.