À l’intérieur d’un LLM : inférence en Python pur pour Qwen3-0.6B sur MLX

On entre « La capitale de la France est » et un modèle y ajoute « Paris » en environ 900 millisecondes. Que se passe-t-il pendant cette courte période ? Quelle série de transformations transforme cinq mots en une prédiction ?

Le meilleur moyen de le découvrir est de construire soi-même tout le pipeline d’inférence au lieu d’utiliser une bibliothèque qui cache les détails. La démonstration ci-dessous le fait avec des poids réels de Qwen3-0.6B, chargés en Python pur via MLX, le framework d’ensembles d’Apple. Le script complet compte une vingtaine de lignes et fonctionne sur n’importe quel Mac Apple Silicon. Les instructions de configuration se trouvent à la fin.

Un aperçu du pipeline

Qwen3, GPT-4, Llama et leurs variantes suivent tous un flux commun :

Text → Tokenize → Embed → [Transformer Block × 28] → LM Head → Logits → Sample → Next Token
                           ↑                                                              |
                           └-------------------------- KV Cache --------------------------┘

Le texte est introduit, transformé en nombres, passe par 28 blocs de transformation, puis ressort sous forme de distribution de probabilités sur 151 936 tokens suivants possibles. Un token est sélectionné et ajouté, puis le cycle recommence.

Ce cycle comporte deux phases :

• Préremplissage. Tous les tokens d’entrée sont traités en même temps. Le modèle lit la demande et construit son état interne.
• Décodage. Les tokens sont générés un par un. Un nouveau token n’a besoin de prendre en compte que ceux précédents, dont les clés et valeurs se trouvent dans le cache KV, de sorte que rien n’est retraité.

Le préremplissage s’exécute en parallèle tandis que le décodage s’effectue séquentiellement. C’est le cache KV qui permet au décodage d’être rapide. Sans lui, générer le token N signifierait de recalculer l’attention pour tous les tokens N-1 qui le précèdent.

Passons maintenant à chaque étape à tour de rôle.

Étape 1 : La tokenisation transforme le texte en entiers

Un modèle ne voit jamais de mots, seulement des entiers.

from tokenizers import Tokenizer
tok = Tokenizer.from_file("Qwen3-0.6B/tokenizer.json")
ids = tok.encode("The capital of France is").ids
# [785, 6722, 315, 9625, 374]

Chaque entier représente un index vers un vocabulaire de 151 936 tokens. Le tokeniseur utilise l’encodage par paires de bytes, de sorte que les limites des tokens ne coïncident pas toujours avec les limites des mots. « capital » est un token et « France » en est un autre, l’espace en tête appartenant au token. Écrire un BPE de zéro n’est pas la partie intéressante ici. Le tokeniseur est en réalité une table de recherche, et le script utilise simplement celle qui a été publiée.

Étape 2 : Les embeddings transforment les entiers en vecteurs

La première opération réelle consiste à associer chaque identifiant de token à un vecteur de 1 024 dimensions.

x = W["model.embed_tokens.weight"][mx.array(token_ids)]  # (seq, 1024)

Il s’agit simplement d’une recherche dans un tableau : la matrice d’incorporation a la forme (151936, 1024), une ligne par token et une colonne par dimension cachée, et le code y accède par indexation. Pensez à cette matrice comme au dictionnaire du modèle : chaque token possède une empreinte numérique de 1 024 valeurs apprise qui encode son sens. Les empreintes de « Paris » et « France » pointent dans des directions similaires, tandis que celles de « Paris » et « ordinateur » pointent dans des directions très différentes.

Étape 3 : RMSNorm

Avant l’attention, les états cachés sont normalisés. Les modèles modernes utilisent RMSNorm à la place de l’ancien LayerNorm.

En termes simples, la formule divise chaque vecteur par sa valeur moyenne carrée, qui correspond à la racine carrée de la moyenne des carrés de ses éléments plus une petite constante, puis elle le multiplie par un poids appris pour chaque caractéristique.

def rms_norm(x, w):
    return x * mx.rsqrt(mx.mean(x * x, axis=-1, keepdims=True) + EPS) * w

Chaque vecteur est donc échelonné en fonction de sa propre magnitude, puis ajusté par des poids appris pour chaque caractéristique. Il n’y a ni soustraction de moyenne ni biais, ce qui le rend plus simple que LayerNorm sans nuire à la qualité. EPS (1e-6) empêche la division par zéro, et le poids w, appris pendant l’entraînement, permet au modèle de rééchelonner indépendamment chacune des 1 024 dimensions.

Étape 4 : RoPE intègre la position dans les vecteurs

C’est là que de nombreux articles sur les « LLM de zéro à zéro » perdent leurs lecteurs, mais c’est en réalité plus simple qu’il n’y paraît.

Le problème. L’attention est invariante par permutation. En mélangeant les tokens d’entrée, le résultat reste identique, ce qui empêche le modèle de savoir quel token est venu en premier.

La solution ancienne. Ajouter un numéro de position à l’embedding de chaque token. Cela fonctionne, mais cela ne se généralise pas, et le modèle a du mal à raisonner sur les distances relatives.

La solution moderne. Faire passer les vecteurs de requête et de clé par un angle qui dépend de l’emplacement du token. C’est ce qu’on appelle l’Embedded Positional Encoding, ou RoPE. Chaque paire de dimensions est considérée comme un point dans un plan et est déplacée d’un angle déterminé par la position du token.

def apply_rope(q, k, positions):
    half = D_HEAD // 2  # 64
    freqs = 1.0 / (THETA ** (mx.arange(half, dtype=mx.float32) / half))
    angles = mx.array(positions, dtype=mx.float32)[:, None] * freqs[None, :]
    cos = mx.cos(angles)[None, :, :]
    sin = mx.sin(angles)[None, :, :]

    def rotate(x):
        x1 = x[..., :half]
        x2 = x[..., half:]
        return mx.concatenate([x1 * cos - x2 * sin, x1 * sin + x2 * cos], axis=-1)

    return rotate(q), rotate(k)

Imaginez chaque paire de dimensions comme un plan 2D. Les tokens situés à des positions différentes tournent d’amounts différents ; ainsi, lorsque l’attention calcule le produit scalaire entre une requête et une clé, la réponse dépend de la distance qui les sépare, et non seulement de leur contenu. Voici trois tokens aux positions 0, 1 et 2, chacun représenté par une flèche 2D tournée un peu plus que le précédent :

Position 0 (angle = 0°)        Position 1 (angle = 24°)       Position 2 (angle = 48°)

      q                                      q                                q
      |                                      /                                /
      |                                     /                                /
      |                                    /                                /
      |                                   /                                /
      +------ x                          +------ x                        +------ x

  Same content, 0° rotation         Rotated 24° counterclockwise    Rotated 48° counterclockwise
  "I'm at the start"               "I'm one step in"               "I'm two steps in"

Dans le dessin original, la troisième flèche n’est pas tournée davantage que la deuxième ; il faut donc considérer le diagramme comme une représentation approximative et non comme une image exacte.

Attention compare une flèche de requête à une flèche de clé au moyen d’un produit scalaire. Les flèches ayant le même angle donnent un produit élevé, tandis que celles ayant des angles très différents donnent un produit faible. Ainsi, ce produit reflète la distance entre deux tokens, sans que de numéro de position explicite ne soit stocké nulle part. RoPE ne rajoute pas la position en tant que signal distinct. La rotation intégre la position dans chaque vecteur.

Qwen3 définit THETA = 1,000,000, une fréquence de base élevée qui permet au modèle de gérer des contextes très longs, allant jusqu’à 262K tokens sans extensions.

Étape 5 : Attention avec des têtes de requête regroupées

C’est l’élément central. Comprendre l’attention, c’est comprendre le transformer.

De MHA à MQA à GQA

Multi-Head Attention (MHA):  16 Q heads, 16 K heads, 16 V heads  — expensive
Multi-Query Attention (MQA): 16 Q heads,  1 K head,  1 V head    — too aggressive
Grouped Query Attention:     16 Q heads,  8 K heads, 8 V heads   — sweet spot

Qwen3-0.6B utilise GQA, disposant de 16 têtes de requête et seulement 8 têtes clé-valeur, de sorte que chaque tête KV sert deux têtes de requête. Cela réduit de moitié la mémoire cache KV avec peu de perte en qualité.

Toute la fonction d’attention

def attention(x, layer_idx, cache, positions):
    p = f"model.layers.{layer_idx}.self_attn"

    # Project to Q, K, V
    q = x @ W[f"{p}.q_proj.weight"].T   # (seq, 2048) = 16 heads × 128
    k = x @ W[f"{p}.k_proj.weight"].T   # (seq, 1024) =  8 heads × 128
    v = x @ W[f"{p}.v_proj.weight"].T   # (seq, 1024) =  8 heads × 128

    # Split into heads
    q = q.reshape(-1, N_HEADS, D_HEAD).transpose(1, 0, 2)    # (16, seq, 128)
    k = k.reshape(-1, N_KV, D_HEAD).transpose(1, 0, 2)       # (8, seq, 128)
    v = v.reshape(-1, N_KV, D_HEAD).transpose(1, 0, 2)       # (8, seq, 128)

    # QK-Norm (Qwen3-specific: normalize each head before RoPE)
    q = rms_norm(q, W[f"{p}.q_norm.weight"])
    k = rms_norm(k, W[f"{p}.k_norm.weight"])

    # Apply RoPE
    q, k = apply_rope(q, k, positions)

    # KV Cache: append new keys/values
    if cache[layer_idx] is not None:
        ck, cv = cache[layer_idx]
        k = mx.concatenate([ck, k], axis=1)
        v = mx.concatenate([cv, v], axis=1)
    cache[layer_idx] = (k, v)

    # GQA: expand 8 KV heads → 16 (each KV head shared by 2 Q heads)
    k = mx.repeat(k, N_HEADS // N_KV, axis=0)   # (16, total, 128)
    v = mx.repeat(v, N_HEADS // N_KV, axis=0)

    # Scaled dot-product attention
    scale = D_HEAD ** -0.5
    scores = (q @ k.transpose(0, 2, 1)) * scale   # (16, seq, total)

    # Causal mask (only during prefill)
    if seq_len > 1:
        mask = mx.triu(mx.full((seq_len, total), -1e9, dtype=scores.dtype), k=1)
        scores = scores + mask

    # Softmax + weighted sum of V
    attn = mx.softmax(scores, axis=-1)
    out  = attn @ v

    # Merge heads + output projection
    out = out.transpose(1, 0, 2).reshape(-1, N_HEADS * D_HEAD)
    return out @ W[f"{p}.o_proj.weight"].T

Quelques parties de ce texte nécessitent des explications.

Norme QK. Qwen3 applique RMSNorm aux vecteurs Q et K de chaque tête avant RoPE, une technique plus récente qui stabilise l’entraînement. Ses poids q_norm et k_norm sont composés de 128 nombres, représentant une échelle par dimension au sein d’une tête.

Le cache KV. Il s’agit de l’optimisation la plus importante pour l’inférence. Lors du préremplissage, K et V sont calculés pour chaque token d’entrée et stockés. Lors du décodage, seuls les nouveaux tokens nécessitent que K et V soient recalculés. Ils sont ajoutés au cache, et l’attention prend en compte l’ensemble de l’historique. Éviter l’utilisation du cache obligerait à recalculer K et V pour chaque token précédent, ce qui transformerait un travail linéaire en un travail quadratique.

Le masque causal. Lors du préremplissage, le token en position 3 ne doit pas voir le token en position 4, qui se trouve dans son avenir. Le code ajoute -1e9 au triangle supérieur de la matrice de scores pour le garantir. Lors du décodage, lorsque la longueur de la séquence est de 1, aucun masque n’est nécessaire, car le seul nouveau token peut consulter tout ce qui le précède.

La formule. L’attention calcule le softmax de (Q multiplié par K transposé, divisé par la racine carrée de la dimension de la tête), puis multiplie le résultat par V. Le produit scalaire mesure dans quelle mesure chaque requête correspond à chaque clé. La division par la racine carrée de la taille de la tête empêche les scores de devenir si élevés que le softmax sature. Le softmax transforme les scores en probabilités, et la multiplication par V donne une moyenne pondérée des vecteurs de valeurs.

Étape 6 : La couche feed-forward SwiGLU

Après avoir attiré l’attention, chaque token passe par un réseau feed-forward. Les modèles modernes utilisent SwiGLU à la place de la version plus ancienne basée sur ReLU. En d’autres termes : deux projections sont calculées à partir de l’entrée, l’une est soumise à SiLU et utilisée comme porte, les deux sont multipliées élément par élément, et une dernière projection ramène le résultat à sa valeur initiale.

def ffn(x, layer_idx):
    p = f"model.layers.{layer_idx}.mlp"
    gate = x @ W[f"{p}.gate_proj.weight"].T   # (seq, 3072)
    up   = x @ W[f"{p}.up_proj.weight"].T     # (seq, 3072)
    silu = gate * mx.sigmoid(gate)             # SiLU(x) = x · σ(x)
    return (silu * up) @ W[f"{p}.down_proj.weight"].T  # (seq, 1024)

Le réseau feed-forward est là où le modèle conserve une grande partie de ses connaissances, ce qui explique en partie pourquoi les grands modèles denses semblent si performants. Il étend le vecteur de 1 024 dimensions à 3 072 dimensions, applique la porte SiLU, puis projette à nouveau vers 1 024 dimensions. Cette porte apprend quelles informations laisser passer et lesquelles bloquer, un peu comme un gardien pour les neurones. L’extension de 1 024 à 3 072 dimensions puis le retour permet au modèle d’stocker et de transformer des informations, et c’est justement ce mécanisme de porte qui fait que les FFN modernes surpassent ceux utilisant ReLU.

Étape 7 : Le bloc transformateur

Un bloc combine tout ce qui a été mentionné ci-dessus :

def transformer_block(x, layer_idx, cache, positions):
    p = f"model.layers.{layer_idx}"
    # Pre-norm attention + residual
    h = rms_norm(x, W[f"{p}.input_layernorm.weight"])
    x = x + attention(h, layer_idx, cache, positions)
    # Pre-norm FFN + residual
    h = rms_norm(x, W[f"{p}.post_attention_layernorm.weight"])
    x = x + ffn(h, layer_idx)
    return x

La séquence est normalisée, on ajoute le résidu, on la normalise à nouveau, on procède à une propagation avant, puis on ajoute à nouveau le résidu.

Ces connexions résiduelles (x = x + ...) sont essentielles. Elles permettent aux gradients de circuler directement à travers le réseau pendant l’entraînement, et elles autorisent les informations à contourner des blocs individuels au moment de l’inférence. Un modèle de 28 couches sans elles serait presque impossible à entraîner.

Qwen3-0.6B empile 28 de ces blocs, chacun affinant un peu plus la représentation. À la 28e couche, l’état caché possède une compréhension suffisante de l’entrée pour prédire le token suivant.

Étape 8 : Tête LM et échantillonnage

Après le dernier bloc, il y a une dernière normalisation, suivie d’une projection vers la taille du vocabulaire :

x = rms_norm(x, W["model.norm.weight"])
logits = x @ W["lm_head.weight"].T   # (seq, 151936)

Le résultat est de 151 936 scores, appelés logits, un pour chaque token possible. Un score plus élevé indique un token plus probable. Pour « La capitale de la France est », les cinq probabilités les plus élevées après l’application de softmax sont les suivantes :

| Token     | Probability |
|-----------|-------------|
| Paris     | 0.6582      |
| located   | 0.0274      |
| the       | 0.0233      |
| ...       | 0.0157      |
| ____      | 0.0136      |

« Paris » remporte la victoire avec 65,8 %, une prédiction assurée. L’échantillonnage convertit ensuite ces probabilités en un unique ID de token :

def sample(logits, temperature=0.0, top_p=0.95):
    if temperature == 0:
        return int(mx.argmax(logits))  # Greedy: always pick the best

    logits = logits / temperature
    probs  = mx.softmax(logits)

    # Top-p: keep smallest set of tokens whose cumulative prob ≥ top_p
    order    = mx.argsort(-probs)
    sorted_p = probs[order]
    cumsum   = mx.cumsum(sorted_p)
    cutoff   = int(mx.sum(cumsum < top_p)) + 1

    sorted_p = sorted_p[:cutoff]
    sorted_p = sorted_p / mx.sum(sorted_p)
    order    = order[:cutoff]

    idx = mx.random.categorical(mx.log(sorted_p[None, :]))
    return int(order[idx])

Température contrôle le hasard. À 0, le token ayant la probabilité la plus élevée gagne toujours (méthode gourmande). À 1,0, les tokens sont sélectionnés en proportion de leurs probabilités, et à 0,7, la distribution s’aiguise légèrement en faveur des tokens les plus probables.

Top-p, ou échantillonnage par noyau, retient le plus petit groupe de tokens dont la probabilité cumulative atteint p, par exemple 0,95. Il réduit la longue queue des tokens peu probables tout en conservant une certaine diversité.

Étape 9 : La boucle de génération

La boucle complète se remplit une fois au préalable, puis décode :

def generate(prompt, max_tokens=100, temperature=0.0, top_p=0.95):
    ids = tok.encode(prompt).ids

    # Prefill: process all input tokens at once
    cache  = [None] * N_LAYERS
    logits = forward(ids, cache, 0)
    next_id = sample(logits[-1], temperature, top_p)
    ids.append(next_id)

    # Decode: generate one token at a time
    for _ in range(max_tokens - 1):
        if next_id == EOS_ID:
            break
        pos      = len(ids) - 1
        logits   = forward([next_id], cache, pos)
        next_id  = sample(logits[-1], temperature, top_p)
        ids.append(next_id)
    return tok.decode(ids)

C’est là tout le cycle d’inférence. Les moteurs de génération tels que vLLM, SGLang, TensorRT-LLM et mlx-lm sont, au fond, des optimisations de ce noyau, qui se compose d’environ 198 lignes de code.

Ce que le script omet

Le script couvre l’ensemble du pipeline d’inférence d’un transformateur standard. Les modèles plus récents ajoutent d’autres éléments :

• Attention hybride. Qwen3.5 remplacerait 75 % de ses couches d’attention par Gated DeltaNet, une variante d’attention linéaire utilisant un état récurrent de taille fixe à la place de la matrice d’attention complète. L’attention complète est comme une mémoire photographique : parfaite, mais coûteuse. DeltaNet ressemble à des notes compressées : efficace, mais avec quelques pertes. Une approche hybride permet de gérer des contextes longs de manière efficace tout en assurant un rappel quasi parfait de manière périodique.
• Prédiction de plusieurs tokens. Certains modèles prédisent plusieurs tokens par passe avant, sacrifiant un peu de qualité pour une grande accélération.
• Quantification. L’inférence en production utilise des poids de 4 bits ou 8 bits pour économiser de la mémoire et du temps. Ce script conserve les poids originaux en bfloat16.
• Batching. Les serveurs traitent de nombreuses requêtes en même temps et partagent les calculs entre elles. Ce script traite une requête à la fois.

Aucun de ces éléments ne modifie l’algorithme de base. Ils s’appuient tous sur la même structure.

Pourquoi se donner cette peine ?

Comprendre le fonctionnement de l’inférence vous rend un praticien plus compétent. Un avertissement de serveur indiquant que « le cache KV est plein » n’est plus énigmatique. L’argument lié à la mémoire dans le cadre de GQA devient facile à expliquer. De plus, un modèle qui divague et se répète vous oriente vers l’utilisation de paramètres tels que la température et top-p.

Un LLM moderne n’est pas de la magie. Il s’agit de 28 couches de normalisation, d’attention et de propagation avant, répétées, avec l’intégration de la position par rotation, et des connaissances stockées dans des couches linéaires contrôlées. Les mathématiques nécessaires sont de niveau licence. C’est l’ingénierie qui révèle l’intelligence, mais les mathématiques en constituent les fondations. Il suffit d’environ 200 lignes de code Python, de poids réels et d’une sortie réelle.

Exécutez-le vous-même

Tout ce qui précède peut être exécuté sur n’importe quel Mac Apple Silicon (M1 ou ultérieur).

1. Téléchargez le modèle

Obtenez Qwen3-0.6B sur Hugging Face (le répertoire Qwen/Qwen3-0.6B). Seuls quatre fichiers sont nécessaires, donc passez à côté des autres :

Qwen3-0.6B/
├── config.json
├── model.safetensors       (1.5 GB — the actual weights)
├── tokenizer.json
└── tokenizer_config.json

Placez le dossier du modèle à côté de llm_inference.py, car le script interprète ./Qwen3-0.6B/ en fonction du répertoire depuis lequel il a été lancé.

2. Créez un environnement virtuel

# Create and activate a fresh venv
python3 -m venv llm_venv
source llm_venv/bin/activate
# Install the minimum packages - just two
pip install mlx tokenizers

Aucun PyTorch, aucun transformer, aucun accelerate, aucun paquet safetensors. MLX s’occupe des calculs sur les tableaux et du chargement des poids, tandis que la bibliothèque tokenizers gère l’encodage et le décodage BPE.

3. Le lancer

python llm_inference.py

Attendez-vous à un résultat de ce type :

Model:  Qwen3-0.6B  (28 layers, 1024 dim, 16 Q heads, 8 KV heads)
Prompt: The capital of France is
------------------------------------------------------------
Output: The capital of France is Paris. The capital of Italy is
        Rome. The capital of Spain is Madrid...
Time:   0.89s

Si vous voyez ModuleNotFoundError: No module named 'mlx', vérifiez que l’environnement virtuel est actif (source llm_venv/bin/activate) et que la machine est équipée d’un processeur Apple Silicon. MLX ne fonctionne pas sur les Mac Intel, Linux ou Windows.

Tout cela fonctionne sur un ordinateur portable, et la réponse apparaît en moins d’une seconde.