Générer du texte et le cache KV

À la leçon 3 du module 01, nous avons fait des expériences de température avec DeepSeek ; le modèle était alors une boîte noire. Cette fois, c'est nous qui l'avons entraîné, et nous pouvons l'ouvrir pour voir ce qu'il fait réellement à chaque étape.

python generate.py

Le modèle ne donne que des probabilités

Une fois un texte lu, le modèle attribue, à la dernière position, un score à chacun des 6289 caractères du vocabulaire, que softmax transforme en probabilités :

== 1. 读完一段文字，模型给下一个字的概率（前 8 名）
  「白日依山尽，黄河入海」→ 无 6.9%  间 4.5%  多 4.3%  深 3.7%  边 3.0%  空 2.6%  难 2.3%  遥 2.1%
  「床前明月光，疑是地上」→ 天 3.2%  花 2.9%  楼 2.4%  清 1.8%  看 1.7%  游 1.7%  人 1.6%  枝 1.3%
  「春风」→ 。 14.2%  ， 5.1%  吹 3.3%  满 2.9%  雨 1.9%  起 1.9%  尽 1.4%  落 1.3%

Après « 黄河入海 », le poème original dit « 流 » (le fleuve Jaune coule vers la mer), mais « 流 » n'est pas dans les 8 premiers. Après « 疑是地上 », « 霜 » (le givre) n'y est pas non plus. La leçon précédente l'a vérifié : ce modèle ne récite pas, il n'a pas mémorisé ces deux poèmes célèbres, il donne seulement des caractères « plausibles » selon les régularités apprises.

Au passage, une découverte intéressante : dans l'« Anthologie complète des poèmes Tang », le « Pensée d'une nuit calme » (静夜思) s'écrit « 床前看月光，疑是地上霜。举头望山月，低头思故乡 », et non « 床前明月光……举头望明月 » comme dans les manuels. On considère généralement que la version que nous connaissons vient d'anthologies des dynasties Ming et Qing. Même si le modèle avait appris le poème par cœur, il aurait donc récité « 看月光 ».

Après « 春风 » (le vent de printemps), le candidat le plus probable est le point (14,2 %). J'ai compté dans les données d'entraînement : « 春风 » apparaît 797 fois, dont 111 immédiatement suivi d'un point, soit 13,9 %, presque exactement la probabilité donnée par le modèle. Ce que le modèle a appris, ce sont précisément ces régularités d'apparition des caractères dans les poèmes Tang.

Générer du texte, c'est répéter une même chose : obtenir les probabilités, choisir un caractère, l'ajouter à la suite, puis calculer les probabilités du caractère suivant. La manière de « choisir » décide de ce qui est écrit.

La température

Le module 01 a présenté la température : diviser les scores par la température, puis appliquer softmax. En dessous de 1, l'écart entre scores élevés et bas se creuse, le modèle est plus « prudent » ; au-dessus de 1, l'écart se réduit, le modèle est plus « audacieux ».

logits = logits[:, -1, :] / temperature
...
next_id = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)

Avec la même graine aléatoire, chacune des trois températures écrit 3 poèmes :

  温度 0.3：
    白云无限意，白发不如何。何事无人在，无人不可知。
    一望东山路，千行万里愁。夜深人不见，秋色月无声。
    江南春水上，江上月明时。独有清风起，何人别有期。
  温度 1.0：
    白露起春光，知君得舞酲。霜时飏先没，桐死势悠扬。
    白发摩牛放，真关万趾兴。夜寒休绕郡，秋杀杜娘宫。玉镜收新镜，金尊坐晚麕。何言吊鱼道，双锡空余年。
    戍楼映江岸，淼危棹度斜。欲分凝炯娩，归盖乱精冰。
  温度 1.5：
    淡露依春似映兹，襞舞凋年雪路揉筝衬情。死势悠扬双觉处，摩牛放将溉噭胞。……

La température 0,3 écrit le plus harmonieusement : « 一望东山路，千行万里愁。夜深人不见，秋色月无声 » (un regard vers la route du mont de l'Est, mille rangées, dix mille lieues de chagrin ; la nuit est profonde, personne en vue, les couleurs d'automne, la lune sans bruit) pourrait presque passer pour authentique. Le prix : des caractères très courants, et beaucoup de répétitions ; « 无人 » revient deux fois, « 江 » deux fois, et à la longue tout se ressemble.

La température 1,0 choisit selon les probabilités brutes du modèle. Le vocabulaire s'enrichit, mais des caractères rares comme « 麕 », « 酲 », « 炯娩 » commencent à apparaître, et le sens se suit plus difficilement.

La température 1,5 part complètement en désordre : le format s'effondre (« 襞舞凋年雪路揉筝衬情 », neuf caractères dans un vers), et les caractères rares envahissent tout. Car à haute température, la distribution de probabilités s'aplatit, et les caractères de très faible probabilité parmi plus de 6000 finissent ensemble par capter une part non négligeable. Dès qu'un caractère étrange est choisi, le contexte qui suit devient étrange, et les erreurs s'accumulent.

top-k : ne choisir que parmi les premiers

Le problème de la haute température, c'est la « longue traîne » : des milliers de caractères improbables qui, ensemble, prennent trop de place. top-k règle cela directement : ne garder que les k caractères les plus probables, écarter tous les autres, et choisir parmi ces k selon leurs probabilités.

if top_k is not None:
    kth = torch.topk(logits, top_k).values[:, -1:]
    logits = logits.masked_fill(logits < kth, float("-inf"))

Toujours à la température 1,5, avec top-k=20 :

== 3. 温度 1.5，但只从概率最高的 20 个字里选（top-k=20）
    白头白马在西陵，自是西林下路尘。今日不辞青史士，可堪回首白苹州。
    万国风尘里，千金雨色深。不知归未到，应似到人中。
    东北山川尽不愁，三千春尽一枝愁。春寒未入千门色，雨尽空闻九夜寒。……

Le format est rétabli, les caractères rares ont disparu, avec plus de variété qu'à la température 0,3. La haute température apporte « l'audace », top-k évite « l'absurde » ; les deux ensemble donnent le meilleur résultat.

Une autre méthode courante s'appelle top-p (échantillonnage par noyau) : pas de nombre fixe, mais on additionne à partir du plus probable jusqu'à ce que la probabilité cumulée atteigne p (par exemple 0,9). Quand les probabilités sont concentrées, il ne reste que quelques caractères ; quand elles sont dispersées, on en garde davantage. Vous avez vu ce paramètre en utilisant l'API au module 01.

Continuation et acrostiche

Le début de la génération n'a pas besoin d'être vide. Donnez la première moitié d'un vers, et le modèle continue :

== 4. 给出开头，让它续写
    床前明月光，已有四邻情。知君得舞袖，何奈在山期。
    大漠孤烟直，高楼楚水清。雁雕摩露下，鸿劒鼓歌中。夜雨初分郡，秋风又过巴。此才应不遇，空与钓鱼舟。
    人工智能好，言蠹不见人。徒思一岁月，独在岘生春。

« 大漠孤烟直 » (dans le grand désert, une fumée solitaire monte droit) est prolongé en un lüshi complet à cinq caractères, huit vers. Même « 人工智能 » (intelligence artificielle), un mot inconnu sous les Tang, il le prolonge, car ces quatre caractères apparaissent chacun dans les poèmes Tang.

On peut aller plus loin et intervenir pendant la génération. Un acrostiche (藏头诗) impose le premier caractère de chaque vers : on place d'abord le caractère imposé, puis on laisse le modèle écrire le vers caractère par caractère, en lui interdisant de produire ponctuation et retours à la ligne, et c'est nous qui ajoutons la ponctuation à la fin :

@torch.no_grad()
def acrostic(heads, n_char=7, temperature=0.8):
    ids = encode("\n")
    for i, head in enumerate(heads):
        ids = torch.cat([ids, encode(head)], dim=1)  # 先把指定的字放进去
        for _ in range(n_char - 1):  # 再让模型一个字一个字地写完这一句
            logits = model(ids)[0][0, -1] / temperature
            logits[BANNED] = float("-inf")
            next_id = torch.multinomial(F.softmax(logits, dim=-1), 1)
            ids = torch.cat([ids, next_id[None]], dim=1)
        ids = torch.cat([ids, encode("，。"[i % 2])], dim=1)  # 标点由我们加，保证格式
    return decode(ids[0].tolist()).strip()

== 5. 藏头诗：每句开头的字由我们指定，其余由模型写
    春色欲依春似春，眠花先得舞花年。不知诗句先酬情，觉后还随楚水人。
    学得重雕出，无情忘却回。止行空向兴，境静更闻闻。

Les débuts des vers réunis donnent « 春眠不觉 » (dans le sommeil du printemps, sans s'en apercevoir) et « 学无止境 » (l'étude n'a pas de fin).

Cette technique a un nom officiel : le décodage contraint. Avant de choisir, on met à moins l'infini le score des options non conformes. Quand, à la leçon 4 du module 02, le serveur garantissait le bon format du JSON produit par DeepSeek, c'était le même principe : à chaque étape, on exclut tous les tokens qui rendraient le JSON invalide.

Le calcul répété

Regardons maintenant l'efficacité de la génération. Revenons à la méthode generate : l'écriture la plus simple envoie à chaque étape tout le texte au modèle :

logits, _ = self(idx)  # 不用缓存：每一步都把整段序列从头算一遍

Pour écrire le 50e caractère, le modèle traite les 49 précédents ; pour le 51e, il retraite ces 49 caractères, plus le 50e.

Pourtant, à y bien réfléchir, le résultat du calcul des 49 premiers caractères n'a pas changé. Le masque causal garantit que chaque position ne regarde que ce qui précède ; ajouter un caractère à la fin n'affecte en rien le calcul des positions précédentes. La seule nouveauté, c'est la dernière position.

De quoi la dernière position a-t-elle besoin pour calculer son attention ? De sa propre q, et des k et v de toutes les positions précédentes. Il suffit donc de stocker les k et v de chaque couche, de ne calculer à chaque étape que les q, k et v du nouveau caractère, et d'ajouter les nouveaux k et v à la suite de ceux stockés. C'est le cache KV.

Implémentation

Dans l'attention de gpt.py, le cache ne demande que quelques lignes de plus :

if cache is not None:  # KV 缓存：把之前算过的 K、V 接在前面，不用重算
    if "k" in cache:
        k = torch.cat([cache["k"], k], dim=2)
        v = torch.cat([cache["v"], v], dim=2)
    cache["k"], cache["v"] = k, v

Un détail à surveiller : le masque causal. Sans cache, q et k ont la même longueur, et le masque est un triangle inférieur carré. Avec le cache, q ne contient que le nouveau caractère, alors que k contient tous les caractères, le nouveau étant le dernier : il peut voir tous les k. Le masque doit donc être décalé :

total = k.size(2)
mask = torch.ones(T, total, dtype=torch.bool, device=x.device).tril(diagonal=total - T)

L'embedding de position doit suivre : le nouveau caractère n'est pas en position 0, mais en position start. C'est le rôle du paramètre start de forward.

À la génération, la première étape envoie tout le prompt et stocke les k et v de toutes les positions dans le cache ; ensuite, chaque étape n'envoie que le dernier caractère :

if use_cache:
    # 第一步把整段提示词送进去；之后每步只送最新的一个词元，前面的 K、V 从缓存里取
    logits, _ = self(idx[:, start:], caches=caches, start=start)
    start = idx.size(1)

Mesurons

== 6. KV 缓存：结果一样吗？快了多少？
  同样的随机种子，写 120 个字：用缓存和不用缓存的结果完全一样
  写  32 个字：不用缓存     23 毫秒，用缓存    12 毫秒，快了 1.8 倍
  写  64 个字：不用缓存     52 毫秒，用缓存    26 毫秒，快了 2.0 倍
  写 120 个字：不用缓存    130 毫秒，用缓存    49 毫秒，快了 2.6 倍
  这个模型每个词元的缓存：2 × 4 层 × 128 维 × 4 字节 = 4 KB，128 个词元共 512 KB

D'abord la justesse : avec la même graine aléatoire, les 120 caractères écrits sont identiques avec et sans cache. Le cache ne fait qu'épargner le calcul répété ; le résultat ne change pas.

Ensuite la vitesse. Plus on écrit de caractères, plus on économise : 1,8 fois plus rapide pour 32 caractères, 2,6 fois pour 120. Sans cache, écrire le n-ième caractère demande de traiter n caractères, et le calcul total croît comme le carré de la longueur ; avec le cache, chaque étape ne traite qu'un caractère.

Sur notre modèle, l'accélération dépasse à peine un facteur deux, rien d'impressionnant. La raison : le modèle est trop petit et le texte trop court. Le vrai calcul de chaque étape ne prend qu'une fraction de milliseconde, et le surcoût de Python lui-même, appels de fonctions et concaténation de tenseurs, pèse lourd ; ce surcoût-là, le cache ne l'épargne pas. Sur un grand modèle, le calcul de chaque étape est bien plus lourd, et les textes comptent facilement des milliers ou des dizaines de milliers de tokens : le cache KV fait alors gagner des ordres de grandeur.

Le prix du cache : la mémoire

Le cache KV échange de la mémoire contre du temps. Pour chaque token, chaque couche stocke un k et un v. Notre modèle ne demande que 4 Ko par token, 512 Ko pour 128 tokens : négligeable.

Pour les grands modèles, c'est autre chose. Un modèle de quelques dizaines de couches à plusieurs milliers de dimensions peut demander de quelques centaines de Ko à plus de 1 Mo de cache par token. Une longue conversation de plus de cent mille tokens occupe alors des dizaines de Go de mémoire graphique pour le cache, souvent plus que le modèle lui-même. Quand on sert de nombreux utilisateurs à la fois, c'est le premier poste de consommation de mémoire graphique.

C'est pourquoi beaucoup des améliorations évoquées à la fin de la leçon précédente visent le cache KV : plusieurs têtes partageant k et v (attention à requêtes groupées), ou k et v compressés en vecteurs plus petits avant stockage (l'attention latente multi-têtes, MLA, de DeepSeek). En présentant vLLM au module 10, nous verrons aussi comment un serveur gère ces caches.

« Les entrées servies depuis le cache coûtent moins cher », vu à la leçon 4 du module 06, est un cache d'un autre niveau : le fournisseur conserve le cache KV des préfixes fréquents (comme un prompt system fixe), et quand une requête suivante utilise le même préfixe, il le récupère directement sans recalculer, d'où un coût bien moindre.

Exercices

1. Implémentez l'échantillonnage top-p : trier les probabilités par ordre décroissant et garder les caractères jusqu'à ce que la probabilité cumulée atteigne p. Écrivez quelques poèmes à la température 1,5 avec p=0,9 et comparez avec top-k=20.
2. Écrivez un décodage contraint à « format fixe » : n'autoriser que des jueju à sept caractères, les 8e, 16e et 24e caractères ne pouvant être qu'une virgule ou un point, sans ponctuation ailleurs.
3. Passez le block_size de gpt.py à 512 et réentraînez (ou initialisez simplement au hasard), puis mesurez le temps d'écriture de 500 caractères avec et sans cache.

Auto-test

1. Que modifient respectivement la température et top-k ? Pourquoi les utilise-t-on souvent ensemble ?

La température met à l'échelle l'ensemble des scores avant softmax et change le caractère « pointu » de la distribution : à basse température, elle se concentre sur quelques caractères bien notés ; à haute température, elle est plus uniforme. top-k ne garde que les k caractères les plus probables et écarte les autres. Une haute température apporte de la diversité, mais laisse un grand nombre de caractères improbables capter une part ; top-k supprime cette longue traîne, et les deux ensemble donnent de la variété sans trop d'absurdité.

2. Pourquoi le cache KV ne change-t-il pas le résultat de la génération ?

À cause du masque causal, le calcul de chaque position ne dépend que d'elle-même et des positions précédentes ; ajouter un caractère ensuite ne change pas les k et v des positions précédentes. Le cache ne fait que stocker ces résultats invariables pour les réutiliser, en épargnant le calcul répété : c'est mathématiquement strictement équivalent.

3. Quel est le prix du cache KV ? Pourquoi ce prix est-il important pour les grands modèles ?

Le prix, c'est la mémoire : pour chaque token, chaque couche stocke un k et un v. Les grands modèles ont beaucoup de couches et de grandes dimensions, le cache par token est volumineux, et avec de longs contextes et de nombreux utilisateurs simultanés, le cache peut occuper plus de mémoire graphique que le modèle lui-même. C'est pourquoi beaucoup d'améliorations d'architecture (comme l'attention à requêtes groupées ou l'attention latente multi-têtes) visent à le réduire.