Atención paginada

Esta página documenta la función de atención paginada en el paso de adelante utilizada en continuous batching.

Ruta Varlen

El kernel flash_attn_varlen_func maneja lotes de longitud variable. Esta ruta se recomienda para lotes con un gran número de solicitudes en prefill.

Comportamiento de la caché

Este kernel no tiene ningún mecanismo para interactuar directamente con la caché paginada, por lo que la caché se lee y escribe manualmente mediante el método ~PagedAttentionCache.update. Esto puede convertirse en un cuello de botella cuando la longitud de la secuencia se vuelve muy grande.

Mecanismo de indexación

El kernel utiliza la longitud máxima de la secuencia (max_length_q, max_length_k) y las longitudes acumuladas de la secuencia (cu_seq_lens_q, cu_seq_lens_k) para calcular la atención de cada secuencia.

Ejemplo

Considere un lote de 3 secuencias con longitudes de consulta [10, 3, 1] y longitudes de clave [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

Formas de entrada:

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]

El kernel asigna cada token de consulta y de clave/valor a una secuencia utilizando las longitudes acumuladas de las secuencias:

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

Ruta de decodificación

El kernel flash_attn_with_kvcache gestiona lotes solo de decodificación en los que cada secuencia tiene exactamente un token de consulta. Es más eficiente que la ruta varlen, pero no puede gestionar lotes con solicitudes de prellenado.

Comportamiento de la caché

Este kernel interactúa con la caché paginada mediante una block_table para indexar la caché y actualizarla in situ. La tabla de bloques tiene la forma (batch_size, max_blocks_per_seq), donde cada fila contiene las ubicaciones físicas de los bloques de caché de una solicitud en el tensor de caché KV.

Mecanismo de indexación

El kernel utiliza cache_seqlens para obtener la longitud de la caché de cada secuencia. Asume que cada token de consulta pertenece a una secuencia diferente (un token por secuencia).

Ejemplo

Considere un lote de 3 secuencias con longitudes de consulta [1, 1, 1] y longitudes de clave [30, 32, 70]. El tamaño del bloque de caché es 32 y el número máximo de bloques por secuencia es 4.

Las longitudes de secuencia de la caché son simplemente las longitudes de las claves:

cache_seqlens = [30, 32, 70]

La forma de la tabla de bloques es (3, 4). Usando direcciones de ejemplo:

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

Los valores de -1 indican bloques no asignados.

• Secuencia 0 (30 tokens en caché): la caché está en KV_cache[2]. El nuevo token cabe (30 + 1 = 31 < 32).
• Secuencia 1 (32 tokens en caché): la caché está en KV_cache[0] y KV_cache[1]. Se necesita un segundo bloque porque 32 + 1 > 32.
• Secuencia 2 (70 tokens en caché): la caché está en KV_cache[3], KV_cache[5] y KV_cache[6]. Tenga en cuenta que los bloques no son necesariamente contiguos, lo cual es la principal ventaja de la caché paginada. El nuevo token cabe en el tercer bloque.