Serie TransformerVariants, nee de l’epic #16058. Le notebook précédent (TV-00a-RoPE-from-scratch.ipynb) ouvrait la boite du positionnement ; celui-ci ouvre celle de l’attention telle qu’elle est reellement cablee dans LLaMA-2/3, Mistral ou Qwen.
Un modèle comme Mistral-7B n’execute pas attention(Q, K, V) avec 32 tetes identiques. Il fait trois choses que le Transformer canonique ne montre pas, et dont chacune repond a une contrainte de cout :
MQA (Multi-Query Attention) : une seule paire de projections K/V pour toutes les tetes Q. Le cache KV de decodage s’effondre d’un facteur egal au nombre de tetes.
GQA (Grouped-Query Attention) : les tetes Q sont reparties en groupes qui partagent une tete K/V. C’est le reglage de LLaMA-2 70B (64 tetes Q, 8 tetes KV) et de Mistral. Le compromis entre la qualite de MHA et la memoire de MQA.
SWA (Sliding Window Attention) : chaque requête ne lit que les W dernières clés. Le cout passe de quadratique a lineaire en la longueur — au prix de l’information hors fenêtre.
L’objectif n’est pas de decrire ces variantes mais de les coder, puis de mesurer ce qu’elles coutent et ce qu’elles rapportent, sur un banc dont on contrôle chaque paramètre.
Ce que ce notebook mesure, et sur quoi il se prononce :
Question
Instrument
Les variantes sont-elles correctes ?
echelle de verifications exactes (boucle par tete, invariants de groupement, identite de masque)
Combien de memoire le cache KV economise-t-il ?
formule analytique confrontee a une allocation reelle
Le partage des tetes KV accelere-t-il l’entraînement ?
chronometrage forward + backward
La fenêtre glissante accelere-t-elle vraiment ?
bande explicite O(T*W) contre masque O(T^2), balayage de T
Que coute la fenêtre en qualite ?
tâche de rappel a longue portee, ou la distance depasse la fenêtre
Tout tourne sur CPU, en PyTorch seul : aucun transformers, aucun xformers, aucun flash-attn, aucun telechargement. Le pendant industriel est le bloc B de l’epic.
Sommaire
L’attention causale multi-tetes, et sa brique de reference
Le cache KV : le vrai cout a l’inference
GQA et MQA : partager les tetes KV
SWA : la bande locale
Vitesse : ce que le partage accelere, et ce qu’il n’accelere pas
Qualite : une tâche ou la fenêtre se paie cash
Conclusion et suite
Exercices (cellules stubbees a completer) : le groupement des tetes Q, le masque en bande, et le balayage de fenêtre. Indices et étapes donnes dans le markdown qui precede chacun.
Conventions. Documentation en francais, code et identifiants en francais ou anglais. Determinisme verifie : torch.set_num_threads(1) rend deux exécutions identiques au bit pres (les reductions multi-thread CPU de PyTorch ne le sont pas). Aucune erreur volontaire dans les cellules d’exercice : un stub renvoie None et le notebook s’execute de bout en bout.
import mathimport timeimport numpy as npimport torchimport osos.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") # determinisme cuBLAS (#16795)# Determinisme (#16795) : la graine seule ne garantit PAS la reproductibilite# (heuristiques cuDNN, kernels non deterministes). warn_only=True au premier# passage pour inventorier les ops fautives sans faire echouer le run.torch.use_deterministic_algorithms(True, warn_only=True)torch.backends.cudnn.deterministic =Truetorch.backends.cudnn.benchmark =Falseimport torch.nn as nnimport torch.nn.functional as F# Les reductions multi-thread de PyTorch sur CPU ne sont pas reproductibles au bit pres.# Un seul thread rend deux executions identiques : indispensable pour un banc de mesure.torch.set_num_threads(1)torch.manual_seed(0)np.random.seed(0)DEVICE = torch.device("cpu")print("torch", torch.__version__, "| numpy", np.__version__, "| device", DEVICE)print("threads torch :", torch.get_num_threads())
avec Q = x W_q, K = x W_k, V = x W_v, une concatenation des tetes en sortie, puis W_o. Le masque causal est triangulaire : la requête i ne lit que les clés j <= i. C’est ce masque qui rend le modèle autoregressif, et c’est lui aussi qui rend le cache KV possible : a la generation, la cle et la valeur de chaque position déjà produite ne changent plus, on les garde.
On code d’abord une brique de reference non vectorisee — une boucle explicite sur les tetes. Elle est lente et c’est voulu : elle sert de verite contre laquelle verifier la forme vectorisee que tout le reste du notebook utilise. Une implementation d’attention qui n’est jamais confrontee a une seconde implementation independante est une implementation dont on ne sait rien.
def causal_window_mask(T, window=None, device=DEVICE):"""Masque booleen (T, T) : True = la requete i a le droit de lire la cle j. window=None -> causal pur (j <= i). window=W -> causal ET bande locale (i - j < W), soit au plus W cles par requete. """ i = torch.arange(T, device=device).view(T, 1) j = torch.arange(T, device=device).view(1, T) mask = j <= iif window isnotNone: mask = mask & ((i - j) < window)return maskdef attn_masked(q, k, v, window=None):"""Attention sur (B, H, T, dh). Forme lisible ; reste O(T^2) meme en fenetre.""" scores = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1]) scores = scores.masked_fill(~causal_window_mask(q.shape[-2], window, q.device), float("-inf"))return torch.softmax(scores, dim=-1) @ vdef init_mha(d_model, n_heads, seed=0):"""Poids d'une MHA canonique, exposes a plat pour la brique de reference.""" gen = torch.Generator().manual_seed(seed) dh = d_model // n_headsdef mat(o, i):return torch.randn(o, i, generator=gen) / math.sqrt(i)return {"Wq": mat(n_heads * dh, d_model), "Wk": mat(n_heads * dh, d_model),"Wv": mat(n_heads * dh, d_model), "Wo": mat(d_model, n_heads * dh)}def mha_naive(x, w, n_heads):"""Forme canonique : boucle explicite tete par tete. Lente, et c'est le but. C'est la seule implementation du notebook dont la lecture suffit a croire le resultat : elle dit exactement ce que dit l'equation, sans reshape ni broadcasting. """ B, T, D = x.shape dh = D // n_heads q, k, v = x @ w["Wq"], x @ w["Wk"], x @ w["Wv"] ctx = torch.zeros(B, T, D)for h inrange(n_heads): sl =slice(h * dh, (h +1) * dh) s = q[:, :, sl] @ k[:, :, sl].transpose(-2, -1) / math.sqrt(dh) s = s.masked_fill(torch.triu(torch.ones(T, T, dtype=torch.bool), 1), float("-inf")) ctx[:, :, sl] = torch.softmax(s, dim=-1) @ v[:, :, sl]return ctx @ w["Wo"]
1.1 Echelle de verifications
Avant de mesurer quoi que ce soit, on etablit que les briques sont ce qu’elles pretendent etre. Cinq verifications, chacune confrontant le code a une propriete enoncee independamment :
(a) la forme vectorisee (reshape + broadcasting) reproduit la boucle tete par tete ;
(b) une fenêtre au moins aussi grande que la sequence ne mord pas : le masque fenêtre est alors le masque causal pur, exactement ;
(c) le nombre de paires (requête, cle) lues sous fenêtre W vaut exactement T*W - W(W-1)/2 des que T >= W — la quantite qui fait passer le cout de O(T^2) a O(T*W) ;
(d) l’etagement des tetes KV (section 3) place la tete Q h sur la tete KV h // (H/G) ;
(e) la bande explicite de la section 4 redonne le résultat du masque, au bruit flottant pres.
Aucune de ces verifications n’utilise le résultat qu’elle valide. C’est le point : (a) et (e) comparent deux chemins de calcul independants, (b) et (c) comparent le code a une formule fermee, (d) a un invariant de groupement.
D_MODEL, N_HEADS =64, 4DH = D_MODEL // N_HEADS# --- (a) la forme vectorisee contre la boucle tete par tete -------------------x = torch.randn(2, 16, D_MODEL)w = init_mha(D_MODEL, N_HEADS, seed=3)q = (x @ w["Wq"]).view(2, 16, N_HEADS, DH).transpose(1, 2)k = (x @ w["Wk"]).view(2, 16, N_HEADS, DH).transpose(1, 2)v = (x @ w["Wv"]).view(2, 16, N_HEADS, DH).transpose(1, 2)vectorise = attn_masked(q, k, v, None).transpose(1, 2).reshape(2, 16, D_MODEL) @ w["Wo"]reference = mha_naive(x, w, N_HEADS)print("(a) vectorise contre boucle par tete : max|ecart| = "f"{(vectorise - reference).abs().max().item():.3e}")# --- (b) une fenetre qui couvre la sequence est le causal pur -----------------for T, W in ((64, 64), (64, 200)): m_fenetre = causal_window_mask(T, W) m_causal = causal_window_mask(T, None)print(f"(b) T={T:<3d} W={W:<3d} : masque fenetre == masque causal ? "f"{bool((m_fenetre == m_causal).all())}")# --- (c) comptage des paires lues contre la formule fermee --------------------print("(c) paires (requete, cle) lues sous fenetre, T >= W :")for T, W in ((64, 16), (64, 32), (100, 7)): exact =int(causal_window_mask(T, W).sum()) formule = T * W - W * (W -1) //2 quadratique = T * (T +1) //2print(f" T={T:<4d} W={W:<3d} paires={exact:<6d} formule={formule:<6d} "f"accord={exact == formule} (causal complet : {quadratique})")
(a) vectorise contre boucle par tete : max|ecart| = 0.000e+00
(b) T=64 W=64 : masque fenetre == masque causal ? True
(b) T=64 W=200 : masque fenetre == masque causal ? True
(c) paires (requete, cle) lues sous fenetre, T >= W :
T=64 W=16 paires=904 formule=904 accord=True (causal complet : 2080)
T=64 W=32 paires=1552 formule=1552 accord=True (causal complet : 2080)
T=100 W=7 paires=679 formule=679 accord=True (causal complet : 5050)
1.2 Lecture
La verification (a) mesure l’ecart entre deux implementations independantes de la même equation : la boucle explicite tete par tete et la forme vectorisee. L’ecart mesure est exactement nul (0.000e+00), et non de l’ordre de l’arrondi flottant comme on pouvait s’y attendre. C’est un accord plus fort que necessaire : la boucle appelle un produit matriciel par tete, la forme vectorisee un bmm sur les quatre tetes d’un coup — deux noyaux distincts, qui tombent pourtant sur les mêmes bits ici. La reserve a garder est qu’une egalite bit a bit n’est jamais garantie entre deux factorisations : elle depend du noyau et de l’ordre des reductions. Un ecart non nul n’aurait d’ailleurs rien invalide ; seule une divergence au-dela de l’arrondi l’aurait fait. La conclusion utile est donc que la forme vectorisee est la fonction de reference, ce qui autorise toutes les mesures qui suivent.
La verification (c) est celle qui porte tout le reste du notebook. Pour T = 64 et W = 16, le masque causal complet autorise 2 080 lectures ; la fenêtre n’en autorise plus que 904, soit un facteur 2.30. La formule T*W - W(W-1)/2 est verifiee exactement sur les trois couples testes, y compris le cas T = 100, W = 7 ou la sequence n’est pas un multiple de la fenêtre. C’est cette quantite qui fait passer l’attention de quadratique a lineaire en T, et c’est donc elle qu’il faut savoir compter avant de parler de vitesse.
Le point (b) merite d’etre garde en tete : une fenêtre egale ou superieure a la longueur de sequence est l’attention causale complete. Ce n’est pas une approximation, c’est le même ensemble de paires. Ce fait sert deux fois dans la suite : en section 6, une variante a fenêtre W = T doit se comporter exactement comme MHA, ce qui permet d’attribuer un eventuel deficit de qualite a la fenêtre elle-même plutot qu’a un defaut du chemin de code ; en section 5, il fournit le contrôle de méthode — SWA (W = 128) a T = 128 a la même computation que MHA, donc leur rapport de temps doit valoir 1.00.
2. Le cache KV : le vrai cout a l’inference
A l’entraînement, les quatre variantes ont le même nombre d’opérations : on calcule l’attention sur toute la sequence, pour toutes les requêtes, avec ou sans partage des tetes K/V. Le partage ne reduit pas les FLOPs d’attention ; il reduit la taille des projections K/V et, surtout, la taille du cache a la generation.
A la generation autoregressive, a chaque nouveau jeton on ajoute une cle et une valeur par tete K/V et par couche, et on les relit toutes. La taille du cache vaut :
octets = 2 (K et V) x B x n_tetes_KV x T x d_head x taille_element x n_couches
Le facteur 2 est le couple K/V, le 2 suivant (implicite) vient de la precision : un cache en float16 coute deux octets par élément. Ce sont n_tetes_KV et T qui decident, et T croit lineairement avec la conversation. C’est cette formule qui explique pourquoi LLaMA-2 70B a ete publie avec 8 tetes KV pour 64 tetes Q.
On calcule la formule sur des configurations reelles plutot que sur des ordres de grandeur, puis on verifie la formule contre une allocation effective.
def kv_cache_bytes(n_couches, n_tetes_kv, dh, T, batch=1, octets_par_element=2):"""Octets du cache KV : 2 tenseurs (K, V) x B x G x T x dh x taille d'un element."""return2* batch * n_couches * n_tetes_kv * T * dh * octets_par_element# (d_model, n_heads, n_kv_heads, n_layers) : tailles publiees des modeles citesCONFIGS = [ ("LLaMA-2 7B", 4096, 32, 32, 32), ("LLaMA-2 70B", 8192, 64, 8, 80), ("Mistral 7B", 4096, 32, 8, 32), ("notre jouet", 64, 4, 4, 2),]T_CONTEXTE =32768print(f"cache KV a T = {T_CONTEXTE} jetons, batch 1, float16\n")print(f"{'config':13s}{'tetes Q':>7s}{'tetes KV':>8s}{'facteur':>8s} "f"{'cache MHA':>10s}{'cache livr':>11s}{'economie':>9s}")for nom, d_model, n_heads, n_kv, n_couches in CONFIGS: dh = d_model // n_heads mha = kv_cache_bytes(n_couches, n_heads, dh, T_CONTEXTE) livre = kv_cache_bytes(n_couches, n_kv, dh, T_CONTEXTE)print(f"{nom:13s}{n_heads:7d}{n_kv:8d}{n_heads / n_kv:7.1f}x "f"{mha /2**30:9.2f}G {livre /2**30:10.2f}G {mha / livre:8.1f}x")# la formule dit-elle ce que le tenseur alloue dit ?G, T_MES, DH_MES =8, 4096, 128cache_k = torch.zeros(1, G, T_MES, DH_MES, dtype=torch.float16)mesure =2* cache_k.numel() * cache_k.element_size()formule = kv_cache_bytes(1, G, DH_MES, T_MES)print(f"\ncontrole de la formule : une couche, G={G}, T={T_MES}, dh={DH_MES}")print(f" mesure sur le tenseur : {mesure:,} octets (K et V)")print(f" formule : {formule:,} octets")print(f" accord exact : {mesure == formule}")
cache KV a T = 32768 jetons, batch 1, float16
config tetes Q tetes KV facteur cache MHA cache livr economie
LLaMA-2 7B 32 32 1.0x 16.00G 16.00G 1.0x
LLaMA-2 70B 64 8 8.0x 80.00G 10.00G 8.0x
Mistral 7B 32 8 4.0x 16.00G 4.00G 4.0x
notre jouet 4 4 1.0x 0.02G 0.02G 1.0x
controle de la formule : une couche, G=8, T=4096, dh=128
mesure sur le tenseur : 16,777,216 octets (K et V)
formule : 16,777,216 octets
accord exact : True
2.1 Lecture
Le tableau chiffre ce que la litterature enonce en prose. A contexte egal (32768 jetons) et en float16, le cache KV de LLaMA-2 70B passe de 80.00 Gio pour la configuration 64/64 a 10.00 Gio avec 8 tetes KV : un facteur 8.0. Mistral 7B conserve le même nombre de tetes KV (8) mais n’a que 32 tetes Q : son facteur vaut donc 4.0, et non le même que celui du 70B. Ce qui se conserve d’une architecture a l’autre, c’est le nombre de tetes KV, pas le facteur d’economie — celui-ci vaut n_heads / n_kv_heads et depend donc des deux. Le contrôle d’allocation confirme que la formule est bien celle que le tenseur occupe (16,777,216 octets, accord exact) : la formule n’est pas une approximation d’ordre de grandeur, c’est une identite.
Deux consequences a ne pas confondre, parce qu’elles se ressemblent et n’ont pas le même domaine :
Le partage des tetes KV est un gain memoire, pas un gain de calcul. Le nombre de multiplications matricielles d’attention est inchange ; ce sont les projections K/V et le cache qui retrecissent. La section 5 le mesure et le nuance : sur CPU, en entraînement, a courte sequence (T = 128) aucune des variantes ne se detache franchement — l’ecart reste de l’ordre de la marge de la méthode de mesure ; mais a T = 512, GQA et MQA gagnent 15-20 %, un ecart qui depasse nettement cette marge.
Le cache est le cout dominant a longue context parce qu’il croit en O(T) par jeton genere, alors que le poids du modèle est constant. A 32768 jetons, le cache de notre configuration jouet n’est evidemment rien ; sur 70B, il se compte en dizaines de gigaoctets. C’est cette asymetrie qui a fait du partage des tetes KV une decision d’architecture, et non une micro-optimisation.
Le cas notre jouet (4 tetes, 4 KV, 2 couches) est volontairement sans partage : c’est la configuration des sections 5 et 6, ou le partage sera introduit comme variable.
3. GQA et MQA : partager les tetes KV
MQA et GQA ne sont pas deux mécanismes, c’est un seul muni d’un paramètre. On ecrit une classe qui prend n_kv_heads et couvre les trois regimes :
n_kv_heads
Nom
Ce que fait la tete Q h
n_heads
MHA
lit sa propre paire K/V
1
MQA
lit l’unique paire K/V, comme toutes les autres
1 < g < n_heads
GQA
lit la paire K/V du groupe h // (n_heads / g)
L’etagement se fait par repeat_interleave, qui duplique la vue sans copier les poids : la tete Q h recoit la tete KV d’indice h // rep avec rep = n_heads / n_kv_heads. MQA est donc exactement GQA avec n_kv_heads = 1. Ecrire les deux separement serait ecrire deux fois le même code et se donner deux fois l’occasion de les faire diverger.
Le partage reduit aussi le nombre de paramètres des projections K et V, d’un facteur n_heads / n_kv_heads : W_k et W_v passent de (n_heads * dh, d_model) a (n_kv_heads * dh, d_model). Ce n’est pas le but recherche — c’est un effet de bord a connaitre, parce qu’il rend toute comparaison de qualite entre variantes boiteuse a paramètres inegaux, et qu’il faudra donc l’annoncer plutot que de la taire.
def build_kv_heads(kv, n_heads, n_kv_heads):"""Etale les tetes KV sur les tetes Q. MHA (n_kv_heads == n_heads) : identite, on ne touche pas au tenseur. MQA (n_kv_heads == 1) : l'unique tete KV est repetee n_heads fois. GQA : la tete Q h recoit la tete KV h // (n_heads / n_kv_heads). """if n_kv_heads == n_heads:return kvreturn kv.repeat_interleave(n_heads // n_kv_heads, dim=1)class VariantAttn(nn.Module):"""Les quatre variantes d'un seul tenant : les leviers sont (n_kv_heads, window). n_kv_heads == n_heads -> MHA | window is not None -> SWA n_kv_heads == 1 -> MQA | banded=True -> SWA en bande O(T*W) 1 < n_kv_heads < n_heads -> GQA | sinon -> masque O(T^2) """def__init__(self, d_model, n_heads, n_kv_heads=None, window=None, banded=False):super().__init__() n_kv_heads = n_heads if n_kv_heads isNoneelse n_kv_headsassert n_heads % n_kv_heads ==0, "les tetes Q doivent se repartir en groupes egaux"self.h, self.g, self.window, self.banded = n_heads, n_kv_heads, window, bandedself.dh = d_model // n_headsself.q_proj = nn.Linear(d_model, n_heads *self.dh, bias=False)self.k_proj = nn.Linear(d_model, n_kv_heads *self.dh, bias=False)self.v_proj = nn.Linear(d_model, n_kv_heads *self.dh, bias=False)self.o_proj = nn.Linear(n_heads *self.dh, d_model, bias=False)def forward(self, x): B, T, _ = x.shape q =self.q_proj(x).view(B, T, self.h, self.dh).transpose(1, 2) k =self.k_proj(x).view(B, T, self.g, self.dh).transpose(1, 2) v =self.v_proj(x).view(B, T, self.g, self.dh).transpose(1, 2) k = build_kv_heads(k, self.h, self.g) v = build_kv_heads(v, self.h, self.g)ifself.banded andself.window isnotNone: o = attn_banded(q, k, v, self.window)else: o = attn_masked(q, k, v, self.window)returnself.o_proj(o.transpose(1, 2).reshape(B, T, self.h *self.dh))# --- (d) invariant de groupement ---------------------------------------------H_GRP, G_GRP, REP =8, 2, 4k_src = torch.randn(1, G_GRP, 5, 4)k_etale = build_kv_heads(k_src, H_GRP, G_GRP)print("(d) forme apres etalement :", tuple(k_etale.shape), f"(attendu (1, {H_GRP}, 5, 4))")print(" tete Q h -> tete KV h // rep, pour tout h :",all(torch.equal(k_etale[:, h], k_src[:, h // REP]) for h inrange(H_GRP)))print(f" tete 0 et tete {REP -1} partagent la meme KV :", torch.equal(k_etale[:, 0], k_etale[:, REP -1]))print(f" tete {REP} bascule sur la KV suivante :", torch.equal(k_etale[:, REP], k_src[:, 1]))print(" branche identite (G == H) :", torch.equal(build_kv_heads(k_src, G_GRP, G_GRP), k_src))# --- cout en parametres des projections K/V ----------------------------------print("\nparametres des projections K/V (d_model=4096, 32 tetes Q, dh=128) :")for g in (32, 8, 1): m = VariantAttn(4096, 32, n_kv_heads=g) n_kv = m.k_proj.weight.numel() + m.v_proj.weight.numel()print(f" n_kv_heads={g:<3d} k_proj+v_proj = {n_kv:>12,} parametres "f"({n_kv / (4096*32*128*2):.3f} x celui de MHA)")
(d) forme apres etalement : (1, 8, 5, 4) (attendu (1, 8, 5, 4))
tete Q h -> tete KV h // rep, pour tout h : True
tete 0 et tete 3 partagent la meme KV : True
tete 4 bascule sur la KV suivante : True
branche identite (G == H) : True
parametres des projections K/V (d_model=4096, 32 tetes Q, dh=128) :
n_kv_heads=32 k_proj+v_proj = 33,554,432 parametres (1.000 x celui de MHA)
n_kv_heads=8 k_proj+v_proj = 8,388,608 parametres (0.250 x celui de MHA)
n_kv_heads=1 k_proj+v_proj = 1,048,576 parametres (0.031 x celui de MHA)
3.1 Lecture
L’invariant de groupement est verifie : la tete Q h lit bien la tete KV h // rep. Deux points mecaniques valent d’etre nommes parce qu’ils sont les seuls endroits ou ce type de code se casse.
Le premier est que repeat_interleave duplique des vues, pas des poids : les n_heads tetes Q continuent de partager les mêmes tenseurs K/V. Une implementation qui copierait les tenseurs (retrograderait le cache a la taille de MHA) passerait tous les tests de correction numérique et n’apporterait aucun des gains de la section 2. La correction seule ne suffit pas a caracteriser une variante : il faut aussi verifier ce qu’elle alloue, ce que fait le tableau du cache.
Le second est le cout en paramètres. Sur une configuration 4096/32 tetes, passer de 32 a 1 tete KV retire 32 505 856 paramètres aux projections K et V. C’est un effet de bord, pas un objectif : W_k et W_v ne servent qu’a projeter, et un groupe plus large de tetes Q les partage. Mais il rend toute comparaison de qualite entre MHA, GQA et MQA inegale en nombre de paramètres — et cette inegalite joue en faveur des variantes partagees quand elles gagnent, ce qu’il faudra rappeler en section 6 au lieu de l’oublier.
Le cas n_kv_heads = 1 n’est pas traite a part : c’est le même appel, avec rep = n_heads. C’est ce qui rend l’ecriture conjointe preferable — une seule branche de code a verifier, et MQA n’est jamais qu’un point de l’espace des paramètres de GQA.
4. SWA : la bande locale
L’attention a fenêtre glissante restreint chaque requête aux W dernières clés. Le masque est l’intersection de deux conditions : causale (j <= i) et locale (i - j < W). Le comptage de la verification (c) dit déjà l’essentiel : T*W - W(W-1)/2 paires au lieu de T(T+1)/2.
Il y a cependant deux facons de coder cette restriction, et la différence entre les deux est le sujet d’une bonne partie de la section 5 :
par masque : on calcule les T x T scores comme d’habitude, puis on remplace par -inf tout ce qui sort de la bande. C’est direct, lisible, et le cout reste O(T^2) : on a calcule les scores hors bande pour les jeter. Le gain memoire du cache est toujours la, le gain de calcul est nul.
par bande : on ne materialise que les W colonnes utiles en decalant les clés, ce qui donne un tenseur (B, H, T, W). Le cout devient reellement O(T*W), mais plus aucune opération n’est un produit matriciel dense — et c’est précisément la le piege que la section 5 va mesurer.
Les deux doivent donner le même résultat. C’est la verification (e) : elle autorise a utiliser la bande pour les mesures de vitesse sans avoir a croire sa correction sur parole.
def attn_banded(q, k, v, window):"""SWA en bande : ne materialise que les W colonnes utiles -> O(T*W). On decale les cles de W-1 vers la droite (remplissage a gauche), puis on decoupe une fenetre glissante de largeur W le long de l'axe des positions. La requete i voit alors les positions decalees relatives 0..W-1, dont la derniere (offset W-1) est la position i. """ B, H, T, dh = q.shape W =min(window, T) k_win = F.pad(k, (0, 0, W -1, 0)).unfold(2, W, 1) v_win = F.pad(v, (0, 0, W -1, 0)).unfold(2, W, 1) scores = torch.einsum("bhtd,bhtdw->bhtw", q, k_win) / math.sqrt(dh)# offset d lu par la requete i vaut la position i + d - (W - 1) ; valide si >= 0 decalage = torch.arange(W, device=q.device) positions = torch.arange(T, device=q.device) keep = decalage.view(1, 1, 1, W) >= (W -1- positions).view(1, 1, T, 1) scores = scores.masked_fill(~keep, float("-inf"))return torch.einsum("bhtw,bhtdw->bhtd", torch.softmax(scores, dim=-1), v_win)# --- (e) la bande redonne-t-elle le masque, exactement ? ----------------------print("(e) bande contre masque, meme (q, k, v) :")for T, W in ((64, 16), (64, 33), (128, 5), (37, 37)): qb = torch.randn(2, 4, T, 8) kb = torch.randn(2, 4, T, 8) vb = torch.randn(2, 4, T, 8) par_masque = attn_masked(qb, kb, vb, window=W) par_bande = attn_banded(qb, kb, vb, window=W)print(f" T={T:<4d} W={W:<3d} max|bande - masque| = "f"{(par_masque - par_bande).abs().max().item():.3e}")# --- et au niveau du module, avec les memes poids ----------------------------torch.manual_seed(0)mod_masque = VariantAttn(64, 4, window=16, banded=False)mod_bande = VariantAttn(64, 4, window=16, banded=True)mod_bande.load_state_dict(mod_masque.state_dict())x_mod = torch.randn(2, 32, 64)ecart_module = (mod_masque(x_mod) - mod_bande(x_mod)).abs().max().item()print(f" au niveau du module (memes poids, W=16, T=32) : {ecart_module:.3e}")
Les deux chemins coincident a l’arrondi flottant pres sur les quatre couples testes, y compris W = 33 (fenêtre non puissance de deux et plus grande que la moitie de la sequence) et W = 37 = T (la fenêtre couvre tout : la bande doit alors degenerer en causal pur). Au niveau du module, avec des poids identiques charges de l’un dans l’autre, l’ecart reste de 1.2e-07 : la bande est une reimplementation fidele du masque, pas une variante approximative.
C’est cette equivalence qui rend la suite possible. Sans elle, la mesure de vitesse de la section 5 comparerait deux fonctions différentes, et un gain de temps ne voudrait rien dire. Avec elle, la seule différence entre les deux lignes du tableau de vitesse est la quantite de calcul effectivement faite — ce qui est exactement la question posee.
Le remplissage a gauche meritait un mot. La bande introduit W-1 positions fictives en tete de sequence ; elles sont masquees par la condition decalage >= W-1-position, donc elles ne contribuent jamais. Le masquage est necessaire pour une autre raison que la correction : sans lui, la première requête verrait des clés nulles et un denominateur de softmax pollue, ce qui produirait une sortie fausse sans lever d’erreur.
Exercice 1 — le groupement des tetes Q
Objectif. Ecrire le groupement des tetes KV sans repeat_interleave, avec un view (ou un reshape suivi d’un expand) : c’est la forme qu’utilisent les implementations industrielles, parce qu’elle evite le tenseur intermediaire.
Indice. Les n_heads tetes Q se lisent comme G groupes de rep tetes. Une tete Q d’indice h appartient au groupe h // rep. On peut donc passer de (B, G, T, dh) a (B, G, 1, T, dh) puis etaler le groupe sur rep — sans jamais copier les données.
Étape 1. Redimensionner le tenseur KV de (B, G, T, dh) en (B, G, 1, T, dh). Étape 2. L’etaler sur l’axe rep avec expand, puis remettre a plat en (B, G * rep, T, dh). Étape 3. Verifier contre build_kv_heads : l’ecart doit etre exactement nul (même adresse memoire, pas de recopie) et la forme doit etre (B, n_heads, T, dh).
Pourquoi ca compte. Le expand ne materialise pas de copie ; repeat_interleave non plus, mais il construit un nouveau tenseur de même forme finale. Sur un cache KV de plusieurs gigaoctets, la différence entre « voir » et « copier » est la différence entre une optimisation et une regression.
def build_kv_heads_par_vue(kv, n_heads, n_kv_heads):"""Exercice 1 : ecrire le groupement avec view + expand plutot que repeat_interleave. Doit renvoyer un tenseur (B, n_heads, T, dh) numeriquement identique a build_kv_heads(kv, n_heads, n_kv_heads), sans copie des donnees. """# TODO etudiant : G = n_kv_heads ; rep = n_heads // G# Etape 1 : kv.view(B, G, 1, T, dh)# Etape 2 : .expand(B, G, rep, T, dh)# Etape 3 : .reshape(B, G * rep, T, dh) resultat =Nonereturn resultat# verification (a completer apres l'exercice) :kv_essai = torch.randn(2, 2, 6, 4)attendu = build_kv_heads(kv_essai, 8, 2)obtenu = build_kv_heads_par_vue(kv_essai, 8, 2)if obtenu isNone:print("Exercice 1 a completer : build_kv_heads_par_vue a renvoye None")else:print("forme :", tuple(obtenu.shape), "| attendu :", tuple(attendu.shape))print("max|ecart| :", (obtenu - attendu).abs().max().item())print("meme stockage que la source :", obtenu.data_ptr() == kv_essai.data_ptr())
Exercice 1 a completer : build_kv_heads_par_vue a renvoye None
5. Vitesse : ce que le partage accelere, et ce qu’il n’accelere pas
On mesure d’abord le cout reel d’un pas forward + backward d’une couche d’attention, pour les quatre variantes, a deux longueurs de sequence. Puis on isole la question de la fenêtre en comparant les deux implementations de la section 4 sur un balayage de T.
Trois avertissements de méthode, avant de lire quoi que ce soit.
Le premier est que ce banc tourne sur un seul thread CPU. Ce n’est pas un detail de confort : c’est ce qui rend les mesures reproductibles, et c’est aussi ce qui rend les conclusions non transposables telles quelles a un GPU. Sur GPU, l’argument du partage des tetes KV est un argument de bande passante memoire — le cache est lu a chaque jeton genere, et un cache huit fois plus petit est un cache lu huit fois plus vite. Cette contrainte n’existe pas de la même facon dans un CPU mono-thread qui execute du code PyTorch.
Le second est qu’on mesure ici un pas d’entraînement : sequence complete, toutes les requêtes en parallele. C’est le regime ou le partage des tetes KV n’a presque rien a apporter, et le mesurer honnetement est plus instructif que de le supposer. Le regime ou il apporte tout est celui de la generation, ou l’on traite un jeton a la fois et ou le cache est relu integralement — regime que ce banc ne simule pas, et qu’on ne pretend donc pas mesurer.
Le troisieme est empirique, et c’est celui qui a change ce notebook. Un banc qui chronometre les variantes a la suite ne mesure pas seulement la variante : il mesure aussi sa position dans la boucle, la première payant l’allocation des tampons. Les mesures sont donc entrelacees, et le tableau imprime un contrôle entre deux variantes dont la computation est identique — sans quoi aucun des ecarts lus plus bas ne voudrait dire quoi que ce soit.
def bench_modules(variantes, T, d_model, batch=2, chauffe=3, tours=9):"""Mediane du forward + backward, mesures ENTRELACEES, en millisecondes. `variantes` : liste de (nom, module). Les tours sont en boucle EXTERIEURE : chaque module est echantillonne a chaque tour, et l'on garde sa mediane. Pourquoi entrelacer. Un banc sequentiel (toutes les repetitions de A, puis toutes celles de B) mesure autant la position dans la boucle que le module : la premiere variante a s'executer paie l'allocation des tampons du banc et sort penalisee. Le controle imprime plus bas est la pour le montrer sur deux variantes dont la computation est identique. """ x = torch.randn(batch, T, d_model, requires_grad=True)for _, m in variantes:for _ inrange(chauffe): m(x).sum().backward() temps = {nom: [] for nom, _ in variantes}for _ inrange(tours):for nom, m in variantes: m.zero_grad(set_to_none=True) debut = time.perf_counter() m(x).sum().backward() temps[nom].append(time.perf_counter() - debut)return {nom: sorted(v)[len(v) //2] *1000for nom, v in temps.items()}def module_variante(d_model, n_kv_heads, window=None, banded=False, n_heads=4):"""VariantAttn a graine fixe : les variantes comparees partent des memes poids.""" torch.manual_seed(0)return VariantAttn(d_model, n_heads, n_kv_heads=n_kv_heads, window=window, banded=banded)D_BENCH =128VARIANTES = [("MHA (G=H)", 4, None, False), ("GQA (G=2)", 2, None, False), ("MQA (G=1)", 1, None, False), ("SWA (W=128)", 4, 128, False)]print("forward + backward d'une couche d'attention, batch 2, d_model =", D_BENCH)print("mesures entrelacees, mediane de 9 tours\n")print(f"{'T':>6s} | "+" | ".join(f"{nom:>14s}"for nom, _, _, _ in VARIANTES))MESURES = {}for T in (128, 512): modules = [(nom, module_variante(D_BENCH, g, w, b)) for nom, g, w, b in VARIANTES] MESURES[T] = bench_modules(modules, T, D_BENCH)print(f"{T:6d} | "+" | ".join(f"{MESURES[T][nom]:11.2f} ms"for nom, _, _, _ in VARIANTES))print("\ncontrole de methode : a T = 128, SWA (W=128) a la MEME computation que MHA")print(" (fenetre >= sequence : meme masque, memes parametres) ; le rapport doit valoir ~1.00")for T in (128, 512): m, ref = MESURES[T], MESURES[T]["MHA (G=H)"]print(f" T = {T:4d} SWA(W=128)/MHA = {m['SWA (W=128)'] / ref:.2f}x"f" GQA/MHA = {m['GQA (G=2)'] / ref:.2f}x MQA/MHA = {m['MQA (G=1)'] / ref:.2f}x")print("\nnombre de parametres par variante (d_model = 128, 4 tetes Q) :")for nom, g, w, _ in VARIANTES: module = module_variante(D_BENCH, g, w)print(f" {nom:14s}{sum(p.numel() for p in module.parameters()):>8,} parametres")
forward + backward d'une couche d'attention, batch 2, d_model = 128
mesures entrelacees, mediane de 9 tours
T | MHA (G=H) | GQA (G=2) | MQA (G=1) | SWA (W=128)
128 | 2.69 ms | 2.54 ms | 2.45 ms | 2.81 ms
512 | 31.03 ms | 26.51 ms | 24.86 ms | 29.21 ms
controle de methode : a T = 128, SWA (W=128) a la MEME computation que MHA
(fenetre >= sequence : meme masque, memes parametres) ; le rapport doit valoir ~1.00
T = 128 SWA(W=128)/MHA = 1.05x GQA/MHA = 0.94x MQA/MHA = 0.91x
T = 512 SWA(W=128)/MHA = 0.94x GQA/MHA = 0.85x MQA/MHA = 0.80x
nombre de parametres par variante (d_model = 128, 4 tetes Q) :
MHA (G=H) 65,536 parametres
GQA (G=2) 49,152 parametres
MQA (G=1) 40,960 parametres
SWA (W=128) 65,536 parametres
print("fenetre : bande explicite O(T*W) contre masque O(T^2), W = 128")print("forward + backward, batch 1, 4 tetes, dh = 32")print("mesures entrelacees, mediane de 9 tours\n")print(f"{'T':>6s}{'masque (ms)':>12s}{'bande (ms)':>11s}{'gain':>7s} "f"{'paires masque':>14s}{'paires bande':>13s}")for T in (128, 512, 1024): modules = [("masque", module_variante(128, 4, window=128, banded=False)), ("bande", module_variante(128, 4, window=128, banded=True))] resultats = bench_modules(modules, T, 128, batch=1) paires_masque = T * (T +1) //2 paires_bande =sum(min(i +1, 128) for i inrange(T))print(f"{T:6d}{resultats['masque']:11.2f}{resultats['bande']:11.2f} "f"{resultats['masque'] / resultats['bande']:6.2f}x "f"{paires_masque:14,d}{paires_bande:13,d}")
Le partage des tetes KV n’accelere pas l’attention : il allege les projections, et le gain reste modeste. A T = 128, GQA et MQA passent devant MHA (GQA/MHA = 0.94, MQA/MHA = 0.91), mais l’ecart est du meme ordre que la marge de methode (voir le contrôle ci-dessous). A T = 512, l’ecart s’elargit a 0.85 et 0.80 : MHA, GQA et MQA tiennent alors dans 31.03 / 26.51 / 24.86 ms, soit un ecart maximal de 6.17 ms (19.9 %).
C’est un effet de part relative — que la section 2 ne prédisait pas : son comptage d’opérations laisserait plutôt attendre l’inverse (les FLOPs d’attention en O(T^2) diluent la part des projections aux longues séquences). C’est la mesure qui l’établit. Le partage retrecit les projections K/V d’un facteur H/G et ne touche a rien d’autre : le nombre de multiplications d’attention est inchange. A T = 128, MQA economise 37 % des parametres de MHA (40 960 contre 65 536) mais ne gagne que 9 % de temps (2.45 contre 2.69 ms) : les couts dominants d’un pas court sont peu sensibles a la taille des projections K/V. L’economie de projections croit avec T ; a T = 512, elle devient clairement visible (15-20 %). Un banc qui aurait « montre » que GQA accelere l’entraînement d’un facteur constant aurait mesure autre chose que ce qu’il croyait.
Avant de lire ces rapports, il faut savoir comment ils ont ete obtenus. Le tableau imprime un contrôle de méthode : a T = 128, SWA (W = 128) a exactement la même computation que MHA (fenêtre >= sequence : même masque, mêmes paramètres, même nombre de tetes), donc leur rapport de temps doit valoir 1.00. Il vaut 1.05 a T = 128 et 0.94 a T = 512 : la méthode a donc elle seule une marge de quelques pour cent. C’est la resolution a laquelle il faut lire les rapports ci-dessus — l’ecart a T = 128 est du meme ordre que cette marge, celui qui apparaît a T = 512 la depasse nettement.
Ce contrôle n’est pas un ornement : sans lui, la première version de ce banc etait illisible. Chronometrees a la suite, les quatre variantes donnaient SWA (W = 128) plus de deux fois plus rapide que MHA a T = 128 — impossible pour deux computations identiques. Ce que le banc séquentiel mesurait, c’etait la position dans la boucle : la première variante a s’executer paie l’allocation des tampons du banc et sort penalisee. Les mesures sont desormais entrelacees — un tour de chaque variante, puis on repete, puis on prend la mediane. La lecon est générale : un ecart entre variantes ne vaut rien tant qu’on n’a pas mesure l’ecart entre deux variantes identiques.
Un dernier detail, qui ne se lit pas dans la théorie et qui motive l’exercice 1 : a T = 128, MQA n’est que marginalement plus rapide que GQA (2.45 contre 2.54 ms), alors que ses projections K/V sont deux fois plus petites — reduire encore le groupe ne reduit guere le temps. C’est que l’etagement des tetes coute quelque chose : repeat_interleave materialise la copie, et MQA en copie quatre quand GQA en copie deux. Le view + expand de l’exercice 1 supprime cette copie — c’est la raison pour laquelle les implementations industrielles l’ecrivent ainsi.
La fenêtre, elle, devrait etre un gain — et ne l’est pas ici. Le balayage compare la bande explicite O(T*W) au masque O(T^2), a W = 128. A T = 128 et T = 512, la bande est plus lente que le masque (0.16 et 0.23) ; a T = 1024, elle reste derriere (0.51). Le point de croisement n’est pas atteint : aucune longueur testee jusqu'a T = 1024. Le comptage de paires, lui, est déjà favorable des T = 1024 : 122 944 lectures pour la bande contre 524 800 pour le masque, soit 4.27 fois moins — alors que le gain mesure reste sous 1.
L’ecart entre ces deux constats est la lecon de la section. Le nombre de paires lues predit le cout asymptotique, pas le cout au point de mesure. Sous le point de croisement, le T x T dense gagne parce qu’il est execute par des produits matriciels denses, que la bibliotheque sous-jacente optimise depuis des decennies ; la bande, elle, paie un unfold (une vue non contigue) et un einsum a cinq dimensions, dont le cout par élément est bien plus eleve. Autrement dit : reduire le nombre d’opérations ne suffit pas, il faut aussi que les opérations restantes soient efficaces.
C’est précisément la raison d’etre d’un noyau fusionne. Une SWA de production (fenêtre glissante dans FlashAttention) ne materialise jamais ni le T x T ni le tenseur de fenêtres : elle decoupe la sequence en blocs et n’ecrit en memoire que ce qui est indispensable, ce qui lui permet de cumuler le comptage de paires favorable et l’efficacite par élément. Notre bande en ops PyTorch ne cumule que le premier des deux, et le notebook le dit plutot que de presenter un O(T*W) théorique comme un gain mesure.
Exercice 2 — le masque en bande, ecrit a la main
Objectif. Reconstruire le masque de la bande (1, 1, T, W) sans utiliser la condition decalage >= W - 1 - position, et sans boucle sur T.
Indice. La position absolue lue par la requête i a l’offset d vaut i + d - (W - 1). La condition de validite est qu’elle soit >= 0. C’est la même chose que de dire que le masque est un triangle degenere : pour la requête i, les premiers max(0, W - 1 - i) offsets sont invalides, les autres sont valides. Une comparaison de deux tenseurs crees par arange, diffuses l’un contre l’autre, suffit.
Étape 1. Construire decalages = torch.arange(W) de forme (W,) et positions = torch.arange(T) de forme (T,). Étape 2. Former decalages[None, :] >= (W - 1 - positions)[:, None], de forme (T, W). Étape 3. Verifier contre causal_window_mask(T, W) : a l’offset d, la cle lue est la position i + d - (W - 1). Le nombre de True par ligne doit valoir min(i + 1, W), et la somme totale T*W - W(W-1)/2.
Pourquoi ca compte. Cette matrice est le seul endroit ou la bande peut etre fausse sans que rien ne leve d’erreur : un decalage d’un cran produit un modèle qui apprend sur des clés decalees, avec des metriques plausibles et un résultat faux.
def masque_bande(T, W, device=DEVICE):"""Exercice 2 : masque (T, W) des offsets valides pour la bande, sans boucle. L'offset d lu par la requete i correspond a la position i + d - (W - 1) ; il est valide si cette position est >= 0. """# TODO etudiant : deux arange, une comparaison diffusee resultat =Nonereturn resultat# verification (a completer apres l'exercice) :mb = masque_bande(64, 16)if mb isNone:print("Exercice 2 a completer : masque_bande a renvoye None")else: par_ligne = mb.sum(dim=1) attendu_ligne = torch.tensor([min(i +1, 16) for i inrange(64)])print("forme :", tuple(mb.shape))print("comptes par ligne conformes a min(i+1, W) :", bool((par_ligne == attendu_ligne).all()))print("total :", int(mb.sum()), "| formule :", 64*16-16*15//2)
Exercice 2 a completer : masque_bande a renvoye None
6. Qualite : une tâche ou la fenêtre se paie cash
Le cout de la fenêtre ne se lit pas dans un chiffre de perplexite sur un corpus quelconque : il se lit sur une tâche ou l’information utile est au-dela de la fenêtre. On construit donc un corpus dont la dépendance longue est plantee et verifiable.
La tâche. Chaque sequence commence par un marqueur tire parmi 8, suivi de jetons de remplissage, et se termine par un jeton de requête. La cible est le marqueur initial : le modèle doit le retrouver a la dernière position, après T - 1 = 63 jetons d’ecart. Huit marqueurs equiprobables fixent le hasard a 1/8 = 0.125 en exactitude, et a une perplexite de 8 a la position de requête.
Pourquoi cette tâche et pas un corpus naturel. Elle rend la grandeur mesuree interpretable par construction : sous la fenêtre W, une requête situee a plus de W jetons du marqueur ne peut litteralement pas le voir, donc la borne superieure atteignable est le hasard. Un corpus naturel n’aurait pas cette propriete, et une degradation de perplexite y serait attribuable a dix causes melangees. On perd en realisme ce qu’on gagne en attribution.
Trois contrôles sont poses dans le dispositif, parce qu’un résultat negatif sans contrôle ne prouve rien :
une variante a fenêtre W = T = 64 : la fenêtre couvre alors toute la sequence (verification (b)), donc le chemin de code SWA est exerce sans restriction. S’il apprend, un echec des fenêtres courtes est attribuable a la restriction, pas a un bug ;
MHA, GQA et MQA : trois facons de lire toute la sequence, qui doivent toutes apprendre ;
la memetisation par profondeur : une fenêtre courte dans une pile de couches donne un champ receptif L * (W - 1) + 1. Avec W = 16, il faut L >= 5 pour couvrir 64 positions. Le test mesure si cette théorie se realise.
On rapporte l’exactitude et la perplexite a la position de requête — pas la perplexite globale, qui serait dominee par les jetons de remplissage, uniformement aléatoires et donc non apprenables.
print("champ receptif : une fenetre W dans une pile de L couches couvre L*(W-1)+1 positions")print(f"distance a franchir : {T -1}. Avec W=16, il faut L >= 5 pour couvrir toute la sequence.\n")print(f"{'profondeur L':>12s}{'champ receptif':>15s}{'exactitude':>12s}{'perplexite':>12s}")for n_couches in (1, 2, 4, 6): acc, ppl, sec, params = entrainer(16, n_couches, 0, pas=400) champ = n_couches * (16-1) +1print(f"{n_couches:12d}{champ:15d}{acc:12.3f}{ppl:12.2f}")print(f"\n(entrainement a 400 pas, une graine, pour borner le temps de calcul)")
champ receptif : une fenetre W dans une pile de L couches couvre L*(W-1)+1 positions
distance a franchir : 63. Avec W=16, il faut L >= 5 pour couvrir toute la sequence.
profondeur L champ receptif exactitude perplexite
1 16 0.141 8.08
2 31 0.141 8.09
4 61 0.141 8.11
6 91 0.141 8.09
(entrainement a 400 pas, une graine, pour borner le temps de calcul)
6.1 Lecture
Le partage des tetes KV ne coute rien en qualite sur cette tâche. MHA (1.000), GQA a 2 tetes KV (1.000) et MQA a une seule tete (1.000) atteignent tous l’exactitude maximale, avec une perplexite a la position de requête de 1.00 : la tâche est apprise exactement, et le fait de faire lire a 4 tetes Q la même cle ne degrade rien ici. C’est le résultat qui compte pour le bloc B : le partage des tetes KV est bien un compromis sur la capacite d’attention, pas sur la tâche — mais notre tâche a une seule information a retrouver, ce qui est le cas le plus favorable possible. Un verdict de qualite sur GQA exige des tâches ou plusieurs tetes doivent lire des choses différentes, et ce banc ne les contient pas. Il faut le dire plutot que d’extrapoler.
La fenêtre, en revanche, se paie immediatement.SWA W=32 tombe a 0.129 d’exactitude, et SWA W=16 a 0.122 — le niveau du hasard (0.125), avec une perplexite de 8.09 pour une borne théorique de 8. La degradation n’est pas une tendance a discuter : a W = 32, la requête finale est a 63 jetons du marqueur, elle ne peut pas le lire, et le modèle n’a rien d’autre a quoi se raccrocher.
Le contrôle W = 64 est ce qui donne son sens a ce résultat. Cette variante emprunte exactement le même chemin de code que les fenêtres courtes, mais sa fenêtre couvre la sequence entiere ; elle apprend comme MHA (1.000). Un echec de W = 16 est donc attribuable a la restriction de la fenêtre, et non a un defaut d’implementation de la bande ou du masque — ce que la verification (e) de la section 4 avait déjà etabli numeriquement, mais qui merite d’etre confirme sur une tâche d’apprentissage.
Le champ receptif théorique ne se realise pas ici. La théorie dit qu’avec W = 16 une pile de L couches couvre L * 15 + 1 positions, donc que L = 6 (champ 91) suffit a franchir 64 positions, et que L = 4 (champ 61) echoue de peu. Le tableau montre 0.141 a L = 1, 0.141 a L = 2, 0.141 a L = 4, 0.141 a L = 6 : aucune profondeur testee ne restaure la capacite. Le champ receptif est une propriete du graphe de calcul, pas une garantie que l’information y circule : faire remonter le marqueur de proche en proche demande au modèle d’apprendre un mécanisme de report d’une couche a l’autre, et le gradient ne lui en donne l’occasion qu’a la toute dernière position. Autrement dit, L * (W-1) + 1 est une condition necessaire — en dessous, c’est perdu d’avance — mais pas suffisante.
C’est aussi ce qui explique la conception de Mistral : une fenêtre de 4096 jetons repetee sur 32 couches, et une information qui circule en profondeur, couche après couche, sur des distances qui depassent largement la fenêtre. Le dispositif n’est credible que parce que le modèle est entraine ainsi des milliards de pas ; a l’echelle d’un banc d’enseignement, on mesure la borne, pas le mécanisme.
Exercice 3 — le point de croisement, mesure
Objectif. Trouver par la mesure la longueur T a partir de laquelle la bande explicite devient plus rapide que le masque, pour une fenêtre W = 64.
Indice. Le point de croisement depend de T/W et de l’efficacite relative du einsum a cinq dimensions contre le produit matriciel dense. Chercher la longueur ou les deux courbes se croisent revient a chercher T tel que T^2 ~ c * T * W, avec c le surcout par élément de la bande.
Étape 1. Choisir au moins cinq longueurs croissantes, dont une sous 512 et une au-dessus de 1024. Étape 2. Mesurer forward + backward des deux implementations pour chacune (helper bench_module). Étape 3. Rapporter le rapport masque / bande et conclure sur l’ordre de grandeur du croisement.
Pourquoi ca compte. C’est la mesure qui separe « l’algorithme est asymptotiquement meilleur » de « ce code-ci est plus rapide ». Un O(T*W) ecrit avec des opérations inefficaces peut rester plus lent que le O(T^2) qu’il remplace sur toute la plage utile, et seul un chronometrage permet de le savoir.
def croisement_bande_masque(W=64, longueurs=(128, 256, 512, 1024, 1536)):"""Exercice 3 : mesurer a partir de quelle longueur la bande bat le masque. Renvoyer un dict {T: (temps_masque_ms, temps_bande_ms)}. """# TODO etudiant : boucler sur longueurs, instancier les deux modules via# module_variante(128, 4, window=W, banded=False) et (..., banded=True), puis les chronometrer# ENSEMBLE avec bench_modules([...], T, 128, batch=1). Les mesurer separement mesurerait la# position dans la boucle autant que la variante -- cf. le controle de la section 5. resultat =Nonereturn resultat# verification (a completer apres l'exercice) :mesures = croisement_bande_masque()if mesures isNone:print("Exercice 3 a completer : croisement_bande_masque a renvoye None")else:print(f"{'T':>6s}{'masque (ms)':>12s}{'bande (ms)':>11s}{'gain':>7s}")for T, (tm, tb) in mesures.items():print(f"{T:6d}{tm:11.2f}{tb:11.2f}{tm / tb:6.2f}x")
Exercice 3 a completer : croisement_bande_masque a renvoye None
7. Conclusion
Ce qu’on a mesuré, et ce que chaque chiffre autorise a conclure.
Affirmation
Preuve dans ce notebook
Verdict
La forme vectorisee implemente la même fonction que la boucle par tete
verification (a), ecart 0.000e+00
etabli
Une fenêtre W >= Test l’attention causale complete
verifications (b) et W=64 en section 6
etabli
Sous fenêtre, les paires lues valent exactement T*W - W(W-1)/2
verification (c), trois couples
etabli
La tete Q h lit la tete KV h // (H/G)
verification (d)
etabli
La bande explicite redonne le masque
verification (e), 1.2e-07 au niveau du module
etabli
Le cache KV suit 2*B*L*G*T*dh*octets
formule contre allocation, 16,777,216 octets
etabli
GQA divise le cache KV par H/G
tableau section 2 : 70B de 80.00 Gio a 10.00 Gio (facteur 8.0)
etabli
Le partage des tetes KV allege les projections, pas l’attention : gain modeste aux sequences courtes, net aux longues
section 5 : 0.94/0.91 a T = 128 contre 0.85/0.80 a T = 512
etabli, a l’echelle et sur ce banc
La bande explicite ne bat pas le masque dans la plage testee
section 5, balayage : 0.51 a T = 1024
etabli sur T <= 1024
Une fenêtre plus courte que la distance d’information detruit la tâche
section 6, W=16 a 0.122 contre hasard 0.125
etabli
Le champ receptif L*(W-1)+1 ne suffit pas a restaurer la capacite
section 6, L de 1 a 6
résultat negatif, borne au budget de ce banc
GQA/MQA ne coutent rien en qualite
section 6, W=T et tâches a information unique
etabli sur cette tâche seulement
Les trois limites a garder en tete. Le banc est mono-thread CPU : les conclusions de vitesse ne se transposent pas a un GPU, ou l’argument du cache KV est un argument de bande passante memoire et non de FLOPs. Les ecarts entre variantes de la section 5 n’ont de sens qu’au-dessus de la marge que le contrôle rend visible (1.05 a T = 128, 0.94 a T = 512) : l’ecart a T = 128 est du meme ordre que cette marge, celui a T = 512 la depasse nettement. Et la tâche de qualite contient une seule information a retrouver, ce qui est le cas le plus favorable au partage des tetes KV ; un verdict de qualite complet demanderait des tâches ou des tetes différentes doivent lire des choses différentes.
Ce que ce notebook n’est pas. Ce n’est pas une reproduction de Mistral ni de LLaMA : les configurations reelles ne servent qu’a chiffrer la formule du cache KV. Le passage a l’echelle, la comparaison avec les noyaux fusionnes de transformers et le routage observe d’un vrai modèle sont le bloc B de l’epic #16058.
Suite de la serie.TV-00c (mixture of experts : routage top-k, facteur de capacite, perte d’equilibrage) puis le bloc B (Mistral-7B ou Phi-3 quantifie 4 bits, mesure des mêmes grandeurs, et tableau comparatif explicite entre le from scratch et le SOTA).