Страничная аттенция

На этой странице описана функция прямого прохода paged attention, используемая в continuous batching.

Varlen path

Ядро flash_attn_varlen_func обрабатывает батчи переменной длины. Этот путь рекомендуется для батчей с большим количеством запросов на этапе prefill.

Поведение кэша

В этом ядре нет механизма для прямого взаимодействия с постраничным кэшем, поэтому кэш читается и записывается вручную с помощью метода ~PagedAttentionCache.update. При увеличении длины последовательности это может стать узким местом.

Механизм индексации

Ядро использует максимальную длину последовательности (max_length_q, max_length_k) и накопленные длины последовательностей (cu_seq_lens_q, cu_seq_lens_k) для вычисления внимания для каждой последовательности.

Пример

Рассмотрим батч из 3 последовательностей с длинами запросов [10, 3, 1] и длинами ключей [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

Формы входных данных:

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 обрабатывает батчи, содержащие только токены декодирования, где у каждой последовательности ровно один токен запроса. Этот путь эффективнее, чем varlen, но не поддерживает батчи с запросами на префиллинг.

Поведение кэша

Это ядро взаимодействует с постраничным кэшем, используя block_table для индексации кэша и его обновления на месте. Таблица блоков имеет форму (batch_size, max_blocks_per_seq), где каждая строка содержит физические расположения блоков кэша запроса в тензоре KV-кэша.

Механизм индексации

Ядро использует cache_seqlens для получения длины кэша для каждой последовательности. Предполагается, что каждый токен запроса принадлежит отдельной последовательности (один токен на последовательность).

Пример

Рассмотрим батч из 3 последовательностей с длинами запросов [1, 1, 1] и длинами ключей [30, 32, 70]. Размер блока кэша составляет 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 (70 кэшированных токенов): кэш в KV_cache[3], KV_cache[5] и KV_cache[6]. Обратите внимание, что блоки не обязательно являются непрерывными, что является ключевым преимуществом страничного кэша. Новый токен помещается в третий блок.