Le sampling et le cache KV.
Fichier : packages/minilab/generate.py
Votre modèle de base est entraîné : il donne une distribution de probabilités pour le token suivant. Reste à écrire, un token après l’autre. Cette étape transforme en code la leçon 09, sur le choix du token (décodage glouton, température, top-k, top-p, seed), et la leçon 15, sur le cache KV, qui évite de tout recalculer à chaque token.
Comprendre l’architecture du fichier
Le fichier prend un modèle, un prompt déjà tokenisé et des réglages, et produit les tokens générés un par un. C’est un générateur Python : il rend chaque token dès qu’il est choisi, ce qui permettra au serveur de l’étape 6 de l’envoyer aussitôt. La génération s’arrête sur un token d’arrêt, comme <|im_end|>, ou quand la fenêtre de contexte est pleine.
À chaque pas, le choix suit toujours le même ordre, celui de la leçon 09 : logits, division par la température, softmax, filtre top-k puis top-p, renormalisation, tirage. Une température de 0 signifie le décodage glouton : toute la probabilité va au favori.
Avec le cache KV, le prompt est lu en une seule fois (prefill), puis chaque pas ne fait passer que le nouveau token dans le modèle (decode). Le résultat doit être le même qu’en recalculant tout. Le même, mais pas au bit près : le benchmark mesure un écart maximal des logits d’environ 1e-5. Ce n’est pas une erreur. Les nombres à virgule flottante sont arrondis à chaque opération, et l’ordre des additions change selon que l’on calcule une matrice entière ou une seule ligne. Ces écarts minuscules ne pourraient changer le choix que si deux tokens étaient à égalité presque parfaite ; ici, les tokens générés sont strictement identiques.
Construire pas à pas : la distribution du token suivant
SamplingSettings regroupe les réglages de la leçon 09, avec le glouton par défaut. next_token_probabilities applique la chaîne : softmax des logits divisés par la température, puis, si demandé, top-k (on garde les tokens au moins aussi probables que le k-ième) et top-p (on trie, et l’on garde les tokens tant que le total de ceux qui les précèdent reste sous p). La dernière ligne renormalise pour que la somme refasse 1.
"""Étape 4 — la génération : sampling, boucle autorégressive et cache KV (leçons 08, 09 et 15).
À chaque pas : logits → ÷ température → softmax → filtre top-k / top-p → renormalisation →tirage. Une température de 0 signifie le décodage glouton (argmax). La boucle s'arrête surun token d'arrêt, comme <|im_end|>, ou après max_new_tokens.
Avec le cache KV, le prompt est lu en un seul passage (prefill), puis chaque nouveau tokenne calcule que sa propre requête, sa clé et sa valeur (decode). Le résultat doit êtreidentique.
uv run python -m packages.minilab.generate"""
import timefrom collections.abc import Iteratorfrom dataclasses import dataclass
import torchfrom torch.nn import functional as F
from packages.minilab.model import KVCache, MiniLabModel
@dataclass(frozen=True)class SamplingSettings: temperature: float = 0.0 # 0 = décodage glouton top_k: int | None = None top_p: float | None = None seed: int = 0
GREEDY = SamplingSettings()
def next_token_probabilities(logits: torch.Tensor, settings: SamplingSettings) -> torch.Tensor: """La distribution dans laquelle on tire le token suivant, après la température, top-k et top-p.""" if settings.temperature == 0: # glouton : toute la probabilité sur le favori probabilities = torch.zeros_like(logits) probabilities[logits.argmax()] = 1.0 return probabilities probabilities = F.softmax(logits / settings.temperature, dim=-1) if settings.top_k is not None: # on garde les k tokens les plus probables kth = torch.topk(probabilities, settings.top_k).values[-1] probabilities = torch.where(probabilities >= kth, probabilities, 0.0) if settings.top_p is not None: # on garde le plus petit ensemble dont le total atteint p ordered, order = torch.sort(probabilities, descending=True) before = torch.cumsum(ordered, dim=-1) - ordered # le total des tokens qui précèdent ordered = torch.where(before < settings.top_p * ordered.sum(), ordered, 0.0) probabilities = torch.zeros_like(probabilities).scatter(0, order, ordered) return probabilities / probabilities.sum() # renormalisationConstruire pas à pas : la boucle de génération
generate est la boucle autorégressive de la leçon 04. Avec le cache, un premier appel au modèle lit tout le prompt et remplit le cache ; ensuite, chaque pas ne donne au modèle que le dernier token choisi. Sans cache, chaque pas relit toute la séquence. Le tirage utilise un générateur aléatoire initialisé avec la seed : mêmes réglages, même texte. Le nombre de tokens est limité pour que prompt et réponse tiennent dans les 128 positions du modèle.
@dataclassclass Step: token_id: int logits: torch.Tensor
@torch.no_grad()def generate( model: MiniLabModel, prompt: list[int], max_new_tokens: int, settings: SamplingSettings = GREEDY, stop_id: int | None = None, use_cache: bool = True, cache: KVCache | None = None,) -> Iterator[Step]: """Produit les tokens générés un par un (le token d'arrêt compris, puis s'arrête).""" model.eval() context = model.config.context_length if not prompt or len(prompt) >= context: raise ValueError("Le prompt doit compter entre 1 et context_length - 1 tokens.") max_new_tokens = min(max_new_tokens, context - len(prompt)) generator = torch.Generator().manual_seed(settings.seed) ids = list(prompt) if use_cache: cache = cache if cache is not None else KVCache(model.config.n_layer) logits, _ = model(torch.tensor([ids]), cache=cache) # prefill : tout le prompt for index in range(max_new_tokens): if not use_cache: # sans cache : toute la séquence, à chaque fois logits, _ = model(torch.tensor([ids])) last = logits[0, -1] probabilities = next_token_probabilities(last, settings) token = int(torch.multinomial(probabilities, 1, generator=generator)) ids.append(token) yield Step(token, last) if token == stop_id or index == max_new_tokens - 1: return if use_cache: # decode : seul le nouveau token traverse le modèle logits, _ = model(torch.tensor([[token]]), cache=cache)Construire pas à pas : démonstration et benchmark
main génère d’abord la même suite avec cinq réglages, puis lance le benchmark : 126 tokens en glouton, avec et sans cache. Il compare les tokens, mesure l’écart maximal des logits et la durée, et compte les tokens passés par les projections W_Q, W_K et W_V.
def main() -> None: from packages.minilab.model import load from packages.minilab.tokenizer import get_tokenizer from packages.minilab.train import BASE_CHECKPOINT
if not BASE_CHECKPOINT.exists(): raise SystemExit("Checkpoint absent : lancez d'abord l'étape 3 (packages.minilab.train).") tokenizer, model = get_tokenizer(), load(BASE_CHECKPOINT) prompt = tokenizer.encode("Phileas Fogg") print("Stratégies de décodage, après « Phileas Fogg » :") for name, settings in [ ("glouton (T = 0) ", SamplingSettings()), ("T = 0,8 ", SamplingSettings(temperature=0.8, seed=1)), ("T = 0,8, top-k = 20 ", SamplingSettings(temperature=0.8, top_k=20, seed=1)), ("T = 0,8, top-p = 0,9 ", SamplingSettings(temperature=0.8, top_p=0.9, seed=1)), ("T = 1,5 (trop audacieux)", SamplingSettings(temperature=1.5, seed=1)), ]: ids = [step.token_id for step in generate(model, prompt, 30, settings)] print(f" {name} « {tokenizer.decode(prompt + ids)} »".replace("\n", " ↵ "))
n = model.config.context_length - len(prompt) print(f"\nCache KV : {n} tokens générés en glouton, avec et sans cache") timings, runs = {}, {} for use_cache in (False, True): cache = KVCache(model.config.n_layer) start = time.perf_counter() runs[use_cache] = list(generate(model, prompt, n, use_cache=use_cache, cache=cache)) timings[use_cache] = time.perf_counter() - start if use_cache: projected_with = cache.projected same_tokens = [s.token_id for s in runs[False]] == [s.token_id for s in runs[True]] gap = max( (a.logits - b.logits).abs().max().item() for a, b in zip(runs[False], runs[True], strict=True) ) layers = model.config.n_layer projected_without = layers * sum(len(prompt) + i for i in range(n)) print(f" mêmes tokens générés : {same_tokens}") print( f" écart maximal des logits : {gap:.1e} (arrondis des nombres à virgule, pas une erreur)" ) print( f" sans cache : {timings[False]:.2f} s · avec cache : {timings[True]:.2f} s → {timings[False] / timings[True]:.1f} fois plus rapide" )
def spaced(number: int) -> str: return f"{number:,}".replace(",", " ")
print( f" tokens passés par W_Q, W_K, W_V (toutes couches) : {spaced(projected_without)} sans cache, {spaced(projected_with)} avec" ) print( f" soit {spaced(projected_without - projected_with)} projections de chaque matrice économisées" )
if __name__ == "__main__": main()Tester et observer dans le terminal
- Lancer l’étape
uv run python -m packages.minilab.generate - Lancer les tests de l’étape
uv run pytest tests/minilab/test_generate.py
Stratégies de décodage, après « Phileas Fogg » : glouton (T = 0) « Phileas Fogg, il semait un c'était un cloignée. ↵ Le Parsi, le guide, il se trouvertant à » T = 0,8 « Phileas Fogg à voie bêtements, suivant de milles, mais cas de chamment à charbre, le gentleman bien » T = 0,8, top-k = 20 « Phileas Fogg à terre ne l'inspecteur. ↵ Le temps, le consul, onze heures du Mongolia, enfortait le 25 » T = 0,8, top-p = 0,9 « Phileas Fogg à voie, s'était autourna à Là, que l'autrement que la résoluse du vol. » T = 1,5 (trop audacieux) « Phileas Fogg à voie bête ygné une falementée." Il sont justicirent que lui ritugnain octobre, dans B » Cache KV : 126 tokens générés en glouton, avec et sans cache mêmes tokens générés : True écart maximal des logits : 1.0e-05 (arrondis des nombres à virgule, pas une erreur) sans cache : 0.12 s · avec cache : 0.04 s → 3.0 fois plus rapide tokens passés par W_Q, W_K, W_V (toutes couches) : 32 508 sans cache, 508 avec soit 32 000 projections de chaque matrice économisées
Les cinq lignes illustrent la leçon 09. Le glouton est figé et finit par tourner en rond. La température de 0,8 rend le texte plus varié ; top-k et top-p coupent la longue traîne. À 1,5, les mots inventés se multiplient. Le benchmark confirme la leçon 15 : mêmes tokens, génération environ 3 fois plus rapide, et surtout 32 000 projections économisées pour chaque matrice. Sans cache, la somme 2 + 3 + … + 127 tokens est refaite dans chacune des 4 couches ; avec le cache, chaque token n’est projeté qu’une fois.
Dans main, changez la seed de la ligne « T = 0,8 » : le texte change. Remettez la même seed : il revient à l’identique. Puis réduisez n à 20 dans le benchmark : avec une suite courte, il y a peu de travail refait, et le gain du cache diminue.
La génération est une boucle : distribution du token suivant (température, top-k, top-p, renormalisation), tirage avec une seed, ajout au contexte, arrêt sur un token de fin ou à la limite de la fenêtre. Le cache KV lit le prompt en une fois (prefill), puis ne traite que le nouveau token (decode) : mêmes tokens, logits identiques à 1e-5 près (arrondis), environ 3 fois plus rapide sur 126 tokens.