Construisez votre propre moteur d'inférence LLM haute performance en C++ et CUDA

Vous êtes-vous déjà demandé ce qui se passe sous le capot lorsque vous envoyez une invite à un grand modèle de langage ? La magie ne réside pas seulement dans les poids du modèle, mais dans le moteur d'inférence qui orchestre le calcul. Cet article vous guide à travers la construction de votre propre moteur d'inférence LLM haute performance en C++ et CUDA, inspiré par l'architecture de vLLM mais conçu comme un outil d'apprentissage.

Pourquoi construire votre propre moteur d'inférence ?

La plupart des développeurs interagissent avec les LLM via des API de haut niveau comme Hugging Face Transformers ou OpenAI. Mais si vous voulez comprendre les véritables caractéristiques de performance—latence, débit, utilisation mémoire—ou optimiser pour votre matériel spécifique, vous devez aller plus loin. Construire un moteur d'inférence à partir de zéro vous apprend :

Comment les LLM calculent réellement les prédictions
Le rôle des noyaux CUDA dans l'accélération des opérations matricielles
Les stratégies de gestion de la mémoire pour les charges de travail GPU
Les techniques avancées comme PagedAttention et le batching continu

L'architecture de tiny-vllm

Le projet, tiny-vllm, implémente un serveur d'inférence complet pour Llama 3.2 1B Instruct. Il couvre :

Chargement d'un vrai LLM à partir du format Safetensors
Passage avant complet (prefill + decode)
Tous les calculs avec des noyaux CUDA personnalisés
Gestion du cache KV
Batching statique et continu
Softmax en ligne et techniques de type FlashAttention
PagedAttention

Pour commencer : Prérequis techniques

Pour exécuter ceci sur votre machine, vous aurez besoin de :

Linux (testé sur 6.19.8 x86_64)
CUDA Toolkit (13.1)
Compilateur C++17 (GCC 15.2.1)
Une carte GPU NVIDIA (testé sur RTX 5090)
Le fichier modèle : model.safetensors de Llama 3.2 1B Instruct

La seule dépendance externe est nlohmann/json (en-tête unique).

Comprendre le format Safetensors

Avant d'écrire du code, vous devez comprendre comment les poids du modèle sont stockés. Un fichier Safetensors a trois sections :

Taille de l'en-tête (8 octets) : Un entier non signé 64 bits indiquant la longueur de l'en-tête
En-tête (JSON) : Contient les métadonnées sur chaque tenseur—nom, dtype, forme et décalages
Données du tenseur : Les valeurs réelles des poids

Voici comment lire la taille de l'en-tête :

std::ifstream safetensors_file("model.safetensors", std::ios_base::binary);
uint64_t header_size;
safetensors_file.read(reinterpret_cast<char *>(&header_size), 8);

Le JSON de l'en-tête vous indique où chaque tenseur commence et se termine dans la section des données. Cela vous permet de mapper les tenseurs directement dans la mémoire GPU sans copier de données inutiles.

L'architecture du modèle : Llama 3.2 1B

À partir de la fiche du modèle Hugging Face, nous pouvons inspecter l'architecture :

LlamaForCausalLM(
  (model): LlamaModel(
    (embed_tokens): Embedding(128256, 2048)
    (layers): ModuleList(
      (0-15): 16 x LlamaDecoderLayer(
        (self_attn): LlamaAttention(
          (q_proj): Linear(in_features=2048, out_features=2048, bias=False)
          (k_proj): Linear(in_features=2048, out_features=512, bias=False)
          (v_proj): Linear(in_features=2048, out_features=512, bias=False)
          (o_proj): Linear(in_features=2048, out_features=2048, bias=False)
        )
        (mlp): LlamaMLP(
          (gate_proj): Linear(in_features=2048, out_features=8192, bias=False)
          (up_proj): Linear(in_features=2048, out_features=8192, bias=False)
          (down_proj): Linear(in_features=8192, out_features=2048, bias=False)
          (act_fn): SiLUActivation()
        )
        (input_layernorm): LlamaRMSNorm((2048,), eps=1e-05)
        (post_attention_layernorm): LlamaRMSNorm((2048,), eps=1e-05)
      )
    )
    (norm): LlamaRMSNorm((2048,), eps=1e-05)
    (rotary_emb): LlamaRotaryEmbedding()
  )
  (lm_head): Linear(in_features=2048, out_features=128256, bias=False)
)

Le passage avant suit cette séquence :

Tokeniser le texte d'entrée
Récupérer les embeddings pour chaque token
Pour chacune des 16 couches du transformeur :
- RMSNorm
- Attention à requêtes groupées (avec RoPE)
- Connexion résiduelle
- RMSNorm
- Réseau feed-forward (avec SiLU)
- Connexion résiduelle
RMSNorm finale
Projection linéaire vers la taille du vocabulaire
Argmax pour sélectionner le token suivant

Votre premier noyau CUDA : Rassemblement d'embeddings

Implémentons la recherche d'embedding. Pour chaque token d'entrée, nous devons récupérer un vecteur de 2048 éléments à partir de la table d'embedding.

__global__ void embeddingGatherKernel(int *gpu_input_tokens, __nv_bfloat16 *input_embeddings, __nv_bfloat16 *embed_tokens)
{
    int workIndex = threadIdx.x + blockIdx.x * 2048;
    input_embeddings[workIndex] = embed_tokens[gpu_input_tokens[blockIdx.x] * 2048 + threadIdx.x];
    input_embeddings[workIndex + 1024] = embed_tokens[gpu_input_tokens[blockIdx.x] * 2048 + threadIdx.x + 1024];
}

Notez l'astuce : comme la plupart des GPU supportent un maximum de 1024 threads par bloc, mais chaque embedding a 2048 éléments, nous traitons deux éléments par thread. Cela évite de lancer plus de blocs que nécessaire.

RMSNorm avec réduction parallèle

RMSNorm normalise chaque vecteur d'embedding. La formule est :

$$\text{RMS}(a) = \sqrt{\frac{1}{n}\sum_{i=1}^{n}a^2_i}$$

Pour calculer cela efficacement sur GPU, nous utilisons une réduction parallèle (réduction en arbre) dans la mémoire partagée :

__shared__ float rms_vector[1024];
int workIndex = threadIdx.x + blockIdx.x * 2048;
rms_vector[threadIdx.x] = (float)input[workIndex] * (float)input[workIndex] 
                        + (float)input[workIndex + 1024] * (float)input[workIndex + 1024];
__syncthreads();

// Réduction en arbre
for (int i = 1; i < 1024; i = i * 2)
{
    if (threadIdx.x % (i * 2) == 0)
    {
        rms_vector[threadIdx.x] = rms_vector[threadIdx.x] + rms_vector[threadIdx.x + i];
    }
    __syncthreads();
}

if (threadIdx.x == 0)
{
    rms_vector[0] = sqrt(rms_vector[0] / 2048.0 + 1.0e-5);
}
__syncthreads();

output[workIndex] = (__nv_bfloat16)(((float)input[workIndex] / rms_vector[0]) * (float)norm_weights[threadIdx.x]);
output[workIndex + 1024] = (__nv_bfloat16)(((float)input[workIndex + 1024] / rms_vector[0]) * (float)norm_weights[threadIdx.x + 1024]);

L'astuce de transposition cuBLAS

cuBLAS s'attend à ce que les matrices soient au format colonne-majeur, mais les modèles sont généralement stockés en ligne-majeur. Au lieu de convertir les données, nous pouvons utiliser les propriétés de transposition de matrice :

$$C = A \times B^T \implies C^T = B \times A^T$$

Puisque cuBLAS lit les données en ligne-majeur comme si elles étaient en colonne-majeur (effectuant une transposition), nous définissons les drapeaux en conséquence :

cublasGemmEx(cublas_handle, 
             CUBLAS_OP_T, CUBLAS_OP_N, 
             KV_DIM, num_active_slots, EMBEDDING_LENGTH, 
             &k_proj_alpha, weights.w_k[layer], CUDA_R_16BF, EMBEDDING_LENGTH, 
             rms_norms, CUDA_R_16BF, EMBEDDING_LENGTH, 
             &k_proj_beta, k_proj_batched_buffer, CUDA_R_16BF, KV_DIM, 
             CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT);

Pourquoi le cache KV existe

Pendant l'inférence, seules les projections K et V de chaque token doivent persister. La projection Q, les scores d'attention et les sorties feed-forward sont jetés après chaque étape. En mettant en cache K et V, nous évitons de les recalculer pour chaque nouveau token—un gain de performance massif.

Prefill vs. Decode

Prefill : Traiter tous les tokens d'entrée pour générer le premier token de sortie. Ceci est limité par le calcul car vous traitez plusieurs tokens à la fois.
Decode : Générer les tokens suivants un par un. Ceci est limité par la mémoire car vous récupérez les valeurs K et V mises en cache.

Attention à requêtes groupées (GQA)

Dans GQA, plusieurs têtes de requête partagent la même tête de clé et de valeur. Pour Llama 3.2 1B, 4 têtes de requête partagent 1 tête de clé/valeur. Cela réduit l'utilisation de la mémoire et améliore le débit sans perte significative de qualité.

for (int i = 0; i < NUM_Q_HEADS; ++i)
{
    int k_head_idx = i / GQA_Q_TO_K_RATIO; // 4 têtes Q utilisent la même 1 tête K
    __nv_bfloat16 *q_head = q_proj + i * HEAD_DIM;
    __nv_bfloat16 *k_head = k_proj[layer] + k_head_idx * HEAD_DIM;
    // ... calculer les scores d'attention
}

Batching continu

Le batching statique traite N requêtes ensemble mais attend la plus lente. Le batching continu résout ce problème en utilisant des emplacements : lorsqu'une requête se termine, son emplacement est immédiatement réutilisé pour une nouvelle requête de la file d'attente. Cela maximise l'utilisation du GPU.

Quelle est la suite ?

Le dépôt est activement développé. Les sections à venir incluent :

Dérivation et implémentation du softmax en ligne
PagedAttention pour une gestion efficace de la mémoire
Cache KV paginé

Conclusion

Construire un moteur d'inférence LLM à partir de zéro est l'une des meilleures façons de comprendre en profondeur comment fonctionnent les systèmes d'IA modernes. Vous apprendrez la programmation GPU, la gestion de la mémoire, l'algèbre linéaire et la conception de systèmes—tout en construisant quelque chose qui exécute réellement un vrai modèle.

Si vous êtes enseignant, n'hésitez pas à utiliser ceci comme ressource pédagogique. Si vous êtes apprenant, forkez le dépôt, faites-le fonctionner sur votre machine et contribuez avec vos améliorations.

Bon code !

Source

jmaczan/tiny-vllm: Construisez votre propre moteur d'inférence LLM haute performance en C++ et CUDA - une version plus petite de vLLM