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

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

• 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 Sinusoidal in NumPy
• RoPE in NumPy
• RoPE GPT-NeoX in NumPy

📊 Positional Encoding Visualizations

Sinusoidal	RoPE	RoPE GPT-NeoX

• 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

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

• Masking in NumPy

• JIT
• Condition
• Shading and Mesh

• 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