Trainable MoE in CUDA

Build CUDA kernels for a complete sparse mixture-of-experts forward, backward, and training loop.

7 parts52 individual lessons4.3 estimated hours260 XP available

0%0 of 52 complete

No compressed chapters.

Every source step is its own lesson with intuition, concepts, correctly rendered MathJax mathematics, implementation, tests, mistakes, and a checkpoint.

PART 1

CUDA Building Blocks

Implement the core CUDA primitives used everywhere in the layer: matmul variants, bias add and reduction, elementwise ops, activations, and row-wise softmax with their backward passes.

0/13

001define cuda building blocks+5 XP 002derive cuda building blocks+5 XP 003prepare cuda building blocks+5 XP 004implement cuda building blocks+5 XP 005connect cuda building blocks+5 XP 006validate cuda building blocks+5 XP 007debug cuda building blocks+5 XP 008optimize cuda building blocks+5 XP 009benchmark cuda building blocks+5 XP 010integrate cuda building blocks+5 XP 011define cuda building blocks+5 XP 012derive cuda building blocks+5 XP 013prepare cuda building blocks+5 XP

PART 2

Top-K Gating Utilities

Implement reusable kernels for selecting the top-k gate values per token and normalizing those gates to a probability distribution, including the backward through that normalization.

0/3

014define top k gating utilities+5 XP 015derive top k gating utilities+5 XP 016prepare top k gating utilities+5 XP

PART 3

Router

Build the router that maps tokens to experts: compute logits via a gate matmul, turn them into probabilities, pick the top-k experts per token, and propagate gradients back to the gate weights.

0/4

017define router+5 XP 018derive router+5 XP 019prepare router+5 XP 020implement router+5 XP

PART 4

Token Dispatch and Combine

Implement the sparse plumbing of MoE: counting tokens per expert, computing slot offsets, gathering tokens into per-expert buffers, scattering gradients back, and combining expert outputs weighted by gate values (with backward).

0/8

021define token dispatch and combine+5 XP 022derive token dispatch and combine+5 XP 023prepare token dispatch and combine+5 XP 024implement token dispatch and combine+5 XP 025connect token dispatch and combine+5 XP 026validate token dispatch and combine+5 XP 027debug token dispatch and combine+5 XP 028optimize token dispatch and combine+5 XP

PART 5

Expert MLP Forward and Backward

Implement the two-layer expert MLP that runs on each expert's slice of tokens, including up/down projections, biases, activation, and the full set of backward kernels for inputs, weights, and biases.

0/12

PART 6

Load Balancing Auxiliary Loss

Implement the MoE load-balancing auxiliary loss: dispatch fractions, mean router probabilities, the aux loss itself, and the gradient back to router probabilities.

0/4

041define load balancing auxiliary loss+5 XP 042derive load balancing auxiliary loss+5 XP 043prepare load balancing auxiliary loss+5 XP 044implement load balancing auxiliary loss+5 XP

PART 7

Loss, Optimizer, and Training Loop

Add the task loss (MSE), gradient-zeroing and SGD update kernels, and host-side orchestration for the full forward pass, backward pass, single training step, and overall training loop.

0/8

045define loss optimizer and training loop+5 XP 046derive loss optimizer and training loop+5 XP 047prepare loss optimizer and training loop+5 XP 048implement loss optimizer and training loop+5 XP 049connect loss optimizer and training loop+5 XP 050validate loss optimizer and training loop+5 XP 051debug loss optimizer and training loop+5 XP 052optimize loss optimizer and training loop+5 XP