Présentation du projet
Samir‑vllm est un moteur d’inférence dédié au modèle Llama 3.2 1B Instruct, développé en C++17 et CUDA. Le code source, disponible sur GitHub, s’appuie sur CMake ≥ 3.24, le compilateur C++17, les bibliothèques ICU pour la tokenisation Unicode‑aware, et le header nlohmann/json. Le projet propose également plus de 80 pages de notes manuscrites détaillant chaque étape du pipeline, de la lecture du fichier model.safetensors aux calculs de logits.
Architecture et flux d’inférence
Le chargeur lit directement les poids au format BF16 depuis le fichier model.safetensors, valide les métadonnées et copie les tenseurs en mémoire GPU. Une variante « memory‑mapped » évite une copie supplémentaire côté CPU. Le modèle possède 16 couches, chacune contenant les tenseurs d’embedding, de normalisation, d’attention et de MLP, tous référencés par des pointeurs fixes. Le tokenizer BPE local, compatible Unicode, encode et décode les prompts selon le format Instruct, et son bon fonctionnement est vérifié par un exécutable de test dédié.
Le pipeline d’inférence démarre par un kernel CUDA d’embedding‑gather qui transfère les ID de tokens du prompt vers le GPU et alloue les activations. Un kernel RMSNorm opère sur les états cachés de largeur 2 048, réduisant en FP32 avant d’appliquer les poids BF16. La phase de préremplissage (prefill) exécute les 16 couches : projection Q/K/V, RoPE, attention à requête groupée (GQA), projection de sortie, résidu d’attention, RMSNorm post‑attention, MLP SwiGLU et résidu MLP. Les logits finaux sont produits en FP32, puis un argmax GPU sélectionne le token suivant.
Gestion de la KV cache et stratégies de batching
Le moteur utilise une piscine de pages KV partagée de 128 MiB par défaut, configurable via --kv-cache-mib. Chaque page physique stocke 16 tokens pour les matrices K et V, et un tableau de blocs par séquence et par couche indique les pages allouées. Lors du préremplissage, les pages du prompt sont remplies ; lors du décodage, chaque nouveau token s’ajoute à la position RoPE absolue correspondante.
Deux modes de batching sont proposés : continuous batching, qui admet les requêtes dès que leurs budgets de pages sont disponibles, et static batching, qui attend la fin du groupe en cours. Le scheduler calcule un crédit d’admission égal à prompt_tokens + allowed_new_tokens - 1, arrondi à la taille de page de 16 tokens pour chaque couche. Si le budget d’une requête dépasse la capacité totale de la piscine, la requête est rejetée avec un message d’erreur explicite, évitant ainsi toute situation de deadlock.
Le décodage batched regroupe les tokens de chaque séquence active dans des tampons partagés, réutilise le handle cuBLAS unique, et recycle les buffers de Q et de normalisation. L’attention paginée lit uniquement les pages correspondant à la séquence courante, accumulant les valeurs pondérées sans concaténer un préfixe contigu ou stocker les scores d’attention. Cette approche réduit la consommation mémoire mais implique que les résultats greedy peuvent diverger entre exécutions individuelles et batched, du fait de différences d’arrondi BF16.
Limitations et perspectives
Le moteur reste limité à un contexte fixe de 2 048 tokens et à la topologie Llama 3.2 1B Instruct. Les calculs BF16 peuvent produire des sorties légèrement différentes selon la taille du batch, ce qui empêche une égalité bit‑à‑bit garantie. De plus, l’absence de support dynamique du nombre de couches ou de la largeur du modèle nécessite une recompilation pour tout changement d’architecture. Malgré ces contraintes, le projet fournit une base pédagogique solide pour comprendre chaque composant d’un transformeur sur GPU, du chargement des poids à la gestion fine de la KV cache.
mkdir -p models/llama-3.2-1b-instruct
cp models/Llama-3.2-1B-Instruct/original/tokenizer.model models/llama-3.2-1b-instruct/
cmake -S . -B build
cmake --build build -j
./build/samir-vllm "Explain the KV cache in two sentences." 128
