Un premier kernel sur le GPU — Smallem : un LLM en Rust et CUDA, à partir de zéro

Deuxième article de la série Smallem. Précédent : Le projet et ses règles.

Avant d'écrire la moindre ligne de modèle, il faut qu'un programme Rust réussisse à faire calculer quelque chose à la carte graphique. Cet article installe l'environnement, fait tourner l'exemple le plus simple qui soit, puis le lit ligne par ligne. Presque tous les kernels de la série auront la même forme.

Toolkit et pilote

« Installer CUDA » recouvre deux composants :

Le pilote est le logiciel qui fait fonctionner la carte. Il est chargé par le système et sert à exécuter les programmes GPU. On le voit avec nvidia-smi.
Le toolkit est un ensemble de compilateurs, de bibliothèques et d'en-têtes. Il sert à construire les programmes GPU.

Les deux ont des numéros de version indépendants. Sur la machine du projet, le pilote (version 595) était déjà installé et convenait. Seul le toolkit manquait : CUDA Rust demande la version 13.0 au minimum, et la machine avait la 12.8.

La page de téléchargement de NVIDIA propose d'installer un pilote en même temps que le toolkit. Il ne faut pas le faire quand un pilote fonctionne déjà : on n'installe que le paquet du toolkit, qui se range à côté de l'ancien.

sudo apt install cuda-toolkit-13-0
/usr/local/cuda-13.0/bin/nvcc --version

Clang, pour du ... Rust

Le deuxième prérequis surprend : il faut Clang 21, un compilateur C, alors qu'on n'écrit pas une ligne de C.

La raison est bindgen. Les bibliothèques de CUDA sont écrites en C et décrites par des fichiers d'en-tête. Pour les appeler depuis Rust, il faut des déclarations Rust équivalentes. bindgen les fabrique automatiquement en lisant les en-têtes C, et pour lire du C il s'appuie sur libclang.

Ubuntu 24.04 s'arrête à Clang 20. La version 21 vient du dépôt du projet LLVM :

wget https://apt.llvm.org/llvm.sh
chmod +x llvm.sh
sudo ./llvm.sh 21
sudo apt install libclang-21-dev libclang-cpp21-dev libclang-common-21-dev

`cargo oxide`

cuda-oxide vient avec son propre outil, cargo oxide, qui pilote toute la chaîne de compilation. Il n'est pas publié sur crates.io et s'installe depuis le dépôt, avec une version précise du compilateur Rust :

cargo +nightly-2026-08-28 install --git https://github.com/NVIDIA/cuda-rust.git cargo-oxide
cargo oxide doctor

La version nightly est imposée. Pour compiler du Rust vers le GPU, cuda-oxide se branche à l'intérieur du compilateur, sur des interfaces qui ne sont pas stables et changent d'une semaine à l'autre. Il fonctionne donc avec une nightly donnée, et pas une autre.

cargo oxide doctor vérifie tout l'environnement et dit ce qui manque. On le relance jusqu'à ce qu'il ne se plaigne plus.

Créer le projet

cargo oxide new smallem-cuda
cd smallem-cuda
cargo oxide run

Ce paquet ne contiendra que ce qui touche à la carte. Le reste du projet vivra dans un second paquet, smallem, en Rust ordinaire : on verra à l'article 4 où passe la frontière.

Lire l'exemple : le kernel

L'exemple additionne deux tableaux de 1 024 nombres. Voici le code qui tourne sur la carte.

#[cuda_module]
mod kernels {
    use super::*;

    #[kernel]
    #[launch_bounds(256)]
    #[launch_contract(domain = 1, block = (256, 1, 1))]
    pub fn vecadd(a: &[f32], b: &[f32], mut c: DisjointSlice<f32>) {
        let idx = thread::index_1d();
        let idx_raw = idx.get();
        if let Some(c_elem) = c.get_mut(idx) {
            *c_elem = a[idx_raw] + b[idx_raw];
        }
    }
}

Pour le comprendre, il faut une image du GPU en tête.

Un processeur ordinaire exécuterait une boucle : pour i de 0 à 1 023, calculer c[i] = a[i] + b[i]. Un GPU ne boucle pas. Il lance des milliers de threads, et chacun exécute la même fonction une seule fois. Ce qui distingue un thread d'un autre, c'est son numéro. Le thread numéro 17 calcule c[17], le thread 18 calcule c[18], tous en même temps.

La fonction vecadd est donc le corps de la boucle, et la boucle elle-même a disparu.

thread::index_1d() donne au thread son numéro. C'est tout ce qu'il sait de sa place dans le calcul.
a et b sont des slices ordinaires, en lecture seule. Tous les threads les lisent en même temps, ce qui ne pose aucun problème.
c est un DisjointSlice. C'est la partie intéressante.

Des milliers de threads qui écrivent dans le même tableau, c'est la recette d'une race condition. En CUDA classique, rien n'empêche deux threads d'écrire dans la même case. Le DisjointSlice l'interdit par construction : on ne peut y écrire qu'avec get_mut(idx), où idx est le numéro du thread, et pas un entier quelconque. Comme chaque thread a un numéro différent, deux threads ne peuvent pas viser la même case.

get_mut renvoie une Option. Les threads sont lancés par paquets de taille fixe, les blocs, ici de 256. Pour 1 024 éléments, quatre blocs tombent juste. Pour 1 000 éléments, il faudrait aussi quatre blocs, et les 24 derniers threads n'auraient aucun élément à calculer. Pour eux, get_mut renvoie None et ils ne font rien.

Restent les trois attributs :

#[kernel] marque la fonction comme devant être compilée pour le GPU.
#[launch_bounds(256)] promet qu'un bloc n'aura jamais plus de 256 threads.
#[launch_contract(domain = 1, block = (256, 1, 1))] fixe la géométrie du lancement : les threads sont numérotés sur une seule dimension, et les blocs font exactement 256 threads. Un lancement qui ne respecte pas ce contrat est refusé avant d'atteindre la carte.

Lire l'exemple : le côté hôte

Le reste du programme tourne sur le processeur, qu'on appelle l'hôte.

let ctx = CudaContext::new(0)?;
let stream = ctx.default_stream();

let a_dev = DeviceBuffer::from_host(&stream, &a_host)?;
let b_dev = DeviceBuffer::from_host(&stream, &b_host)?;
let mut c_dev = DeviceBuffer::<f32>::zeroed(&stream, N)?;

let module = unsafe { kernels::load(&ctx)? };
let prepared = module.prepare_vecadd(LaunchConfig1D::new((N as u32).div_ceil(256), 256, 0))?;
module.vecadd(&stream, &prepared, &a_dev, &b_dev, &mut c_dev)?;

let c_host = c_dev.to_host_vec(&stream)?;

Quatre notions apparaissent :

Le contexte est la connexion à une carte, ici la carte numéro 0.

Le stream est une file d'attente de travaux pour la carte. Tout ce qu'on lui demande y est ajouté et exécuté dans l'ordre. Cette file aura une conséquence importante à l'article 5 : demander un calcul n'est pas attendre qu'il soit fini.

Le DeviceBuffer est un tableau qui vit dans la mémoire de la carte. Le processeur et le GPU n'ont pas la même mémoire : from_host copie des données vers la carte, to_host_vec les rapatrie. Entre les deux, les données restent sur la carte, et c'est là qu'on veut les garder le plus longtemps possible, car chaque aller-retour coûte cher.

Le module est l'ensemble des kernels compilés. #[cuda_module] a fabriqué pour chacun une méthode Rust typée : module.vecadd(...) vérifie à la compilation qu'on lui passe bien deux tampons de f32 et un tampon modifiable. Le lancement se fait en deux temps : on prépare une configuration, ici ceil(1024 / 256) = 4 blocs de 256 threads, puis on lance.

Ce qui se passe à la compilation

Quand on tape cargo oxide run, la fonction vecadd suit un chemin inhabituel :

Rust ──► MIR ──► LLVM IR ──► PTX ──► code machine du GPU

Le MIR est la représentation interne du compilateur Rust. cuda-oxide le récupère pour les fonctions marquées #[kernel], le traduit en LLVM IR, puis en PTX, un langage d'assemblage pour les cartes NVIDIA. Le PTX est rangé dans l'exécutable. Au lancement, le pilote le traduit en code machine pour la carte présente.

On retrouve la distinction du début : le toolkit sert jusqu'au PTX, le pilote après. Et on reviendra lire ce PTX à l'article 11, pour comprendre pourquoi un kernel était lent.

Ce qu'il faut retenir

Le pilote exécute, le toolkit compile, et leurs versions sont indépendantes.
Un kernel est le corps d'une boucle dont la boucle a disparu : chaque thread l'exécute une fois, avec son propre numéro.
Les écritures doivent être disjointes. DisjointSlice le garantit en n'acceptant comme indice que le numéro du thread.
Les données vivent sur la carte ; les copier depuis ou vers l'hôte est une opération à part entière.

Le GPU fonctionne. On va le laisser de côté pendant deux articles, le temps de préparer des données et de fixer ce que chaque opération doit calculer : Des données et un premier modèle.