Attention paginée

Cette page documente la fonction avant de l'attention paginée utilisée dans le batching continu.

Chemin Varlen

Le noyau flash_attn_varlen_func gère les lots de longueur variable. Ce chemin est recommandé pour les lots comportant un grand nombre de requêtes en préremplissage.

Comportement du cache

Ce noyau ne dispose d’aucun mécanisme pour interagir directement avec le cache paginé, de sorte que celui-ci est lu et écrit manuellement à l’aide de la méthode ~PagedAttentionCache.update. Cela peut devenir un goulot d’étranglement lorsque la longueur de la séquence augmente.

Mécanisme d'indexation

Le noyau utilise la longueur maximale de séquence (max_length_q, max_length_k) et les longueurs cumulées de séquence (cu_seq_lens_q, cu_seq_lens_k) pour calculer l'attention pour chaque séquence.

Exemple

Considérons un lot de 3 séquences avec des longueurs de requête [10, 3, 1] et des longueurs de clé [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

Formes d'entrée :

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]

Le noyau affecte chaque jeton de requête et de clé/valeur à une séquence à l'aide des longueurs cumulées des séquences :

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

Chemin de décodage

Le noyau flash_attn_with_kvcache gère les lots de décodage uniquement, où chaque séquence possède exactement un jeton de requête. Cette approche est plus efficace que le chemin varlen, mais elle ne peut pas traiter des lots contenant des requêtes de préremplissage.

Comportement du cache

Ce noyau interagit avec le cache paginé à l’aide d’un block_table pour indexer le cache et le mettre à jour en place. La table de blocs a la forme (batch_size, max_blocks_per_seq), où chaque ligne contient les emplacements physiques des blocs de cache d’une requête dans le tenseur de cache KV.

Mécanisme d'indexation

Le noyau utilise cache_seqlens pour récupérer la longueur du cache de chaque séquence. Il suppose que chaque jeton de requête appartient à une séquence différente (un jeton par séquence).

Exemple

Considérons un lot de 3 séquences avec des longueurs de requête [1, 1, 1] et des longueurs de clé [30, 32, 70]. La taille du bloc de cache est de 32 et le nombre maximal de blocs par séquence est de 4.

Les longueurs de séquence du cache sont simplement les longueurs des clés :

cache_seqlens = [30, 32, 70]

La forme du tableau de blocs est (3, 4). En utilisant des adresses d'exemple :

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

Les valeurs de -1 indiquent des blocs non alloués.

• Séquence 0 (30 jetons mis en cache) : cache dans KV_cache[2]. Le nouveau jeton tient (30 + 1 = 31 < 32).
• Séquence 1 (32 jetons mis en cache) : cache dans KV_cache[0] et KV_cache[1]. Un second bloc est nécessaire car 32 + 1 > 32.
• Séquence 2 (70 jetons mis en cache) : cache dans KV_cache[3], KV_cache[5] et KV_cache[6]. Notez que les blocs ne sont pas nécessairement contigus, ce qui est l'avantage clé du cache paginé. Le nouveau jeton tient dans le troisième bloc.