Machine Learning

Here are the modern implementations of LLM architecture, sharding strategies and kernel optimizations.

Core

• Softmax in PyTorch with autograd and in NumPy
• Linear projection in PyTorch with autograd
• Multi Layer Perceptron in PyTorch with autograd

Transformer Architecture

• Disaggregated Serving with KV Cache in PyTorch
• Multihead Attention in PyTorch with autograd
• Norm RMS in PyTorch with autograd and in NumPy
• Transformer in PyTorch with autograd

Positional Encoder

• Positional Encoder Sinusoidal in NumPy
• RoPE in NumPy
• RoPE GPT-NeoX in NumPy

📊 Positional Encoding Visualizations

Sinusoidal	RoPE	RoPE GPT-NeoX

Flash Attention

• Flash Attention v1 and v2 in PyTorch

📊 Flash Attention Visualizations

See image/flash_attention/README.md for details.

Performance vs Block Size	Memory: HBM vs SRAM

Tile Size vs Latency	Br × Bc Heatmap

FA1 vs FA2: Theoretical HBM Access

Sharding strategies

• MLP Data Parallelism(DP) in PyTorch, in JAX
• MLP Tensor Parallelism(TP) in PyTorch, in JAX
• MLP Fully Sharded Data Parallelism(FSDP) in PyTorch, in JAX
• MLP Pipelining in PyTorch

Scaling plots

The following are roofline analysis for different architectures. Those are non-fused operations.

• MLP roofline analysis in NumyPy
• Multi-Head Attention roofline analysis in NumyPy

📊 Roofline Plots

MLP	Attention

NumPy Tutorial

• Masking in NumPy

JAX Tutorial

• JIT
• Condition
• Shading and Mesh

PyTorch Notes

• Torch distributed API.
• don't use the old primitives, instead use in-place ones like dist.all_gather_into_tensor and dist.all_reduce_tensor that aggregate along the primary dimension.
• custom classes for training requires torch.autograd.Function, @staticmethod and ctx.save_for_backward