From custom GPU kernels to Llama architecture — understanding ML at every layer
Quick Start • Features • Architecture • Benchmarks • Examples
Rusty is a high-performance machine learning framework built entirely from scratch in Rust. It provides GPU-accelerated inference and training with a focus on transformer architectures like Llama.
What makes this project unique:
- Custom GPU Kernels — Hand-written WGSL compute shaders, not relying on cuBLAS or external libraries
- Complete Llama Architecture — Multi-head attention with RoPE, SwiGLU MLP, RMSNorm, KV cache
- LoRA Fine-tuning — Parameter-efficient training on consumer hardware
- Cross-Platform GPU — Runs on Metal (Apple Silicon) and Vulkan (Windows/Linux)
Run a GPU demo in 30 seconds:
git clone https://github.com/puranikyashaswin/rusty.git
cd rusty
cargo run --example basic_tensor --release -p rustyExpected output:
Rusty ML - Basic Tensor Example
[GPU] Apple M2 (Metal)
[INIT] Creating tensors...
Tensor A: [32, 32]
Tensor B: [32, 32]
[COMPUTE] Performing matrix multiplication...
Result shape: [32, 32]
First few values: [10.65, 10.72, 10.79, 10.87, 10.94]
[DONE] All operations completed successfully!
Custom WGSL compute shaders optimized for ML workloads:
| Category | Kernels |
|---|---|
| Linear Algebra | Tiled MatMul, RoPE, RMSNorm |
| Activations | SiLU, Softmax, ReLU |
| Training | AdamW, SGD, Gradient Clipping |
| Quantization | Int8 Dequantization, FP16 Casting |
| Attention | Flash Attention, Scaled Dot-Product |
Production-ready building blocks:
- Embedding — Token embedding with vocabulary lookup
- Linear — Dense layers with optional LoRA adapters
- Attention — Multi-head attention with rotary embeddings
- MLP — SwiGLU feedforward network
- LlamaBlock — Complete transformer block
- LlamaModel — Full model with generation support
- Automatic differentiation with gradient tape
- GPU-accelerated AdamW optimizer
- Mixed precision training (FP16)
- Gradient accumulation and clipping
- LoRA for parameter-efficient fine-tuning
┌─────────────────────────────────────────────────────────────────┐
│ rusty-cli │
│ Command Line Interface │
├─────────────────────────────────────────────────────────────────┤
│ rusty-trainer │ rusty-loader │
│ Training Loops │ Safetensors + Tokenizer │
├─────────────────────────────────────────────────────────────────┤
│ rusty-graph │
│ Neural Networks: Attention, MLP, Llama │
├─────────────────────────────────────────────────────────────────┤
│ rusty-autograd │
│ Automatic Differentiation + Optimizers │
├─────────────────────────────────────────────────────────────────┤
│ rusty-backend │
│ GPU Compute Engine + WGSL Kernels │
├─────────────────────────────────────────────────────────────────┤
│ Metal / Vulkan │
│ Apple M1/M2/M3, GPUs │
└─────────────────────────────────────────────────────────────────┘
| Crate | Description |
|---|---|
| rusty-backend | GPU compute engine with custom WGSL shaders |
| rusty-graph | Neural network layers (Attention, MLP, LlamaBlock) |
| rusty-autograd | Automatic differentiation and optimizers |
| rusty-loader | Safetensors and tokenizer loading |
| rusty-trainer | Training loops with mixed precision |
| rusty-cli | Command-line interface |
Benchmarked on Apple M2 (Metal backend):
| Operation | Size | Throughput |
|---|---|---|
| MatMul | 4096×4096 | 121 GFLOPS |
| MatMul | 2048×2048 | 114 GFLOPS |
| MatMul | 1024×1024 | 109 GFLOPS |
| Softmax | 2048×2048 | 850M elem/s |
| RMSNorm | 4096×2048 | 920M elem/s |
Run benchmarks
cargo run -p benchmarks --releasecargo run --example basic_tensor --release -p rustyMatrix multiplication, element-wise operations, and activations on GPU.
cargo run --example flash_attention --release -p rustyMemory-efficient attention with O(N) memory instead of O(N²).
cargo run --example lora_finetune --release -p rustyParameter-efficient training with low-rank adapters.
cargo run -p rusty-cli --release -- --demoComplete training loop with loss computation.
./scripts/download_model.sh tinyllamaOr manually:
pip install huggingface-hub
huggingface-cli download TinyLlama/TinyLlama-1.1B-Chat-v1.0 --local-dir ./models/tinyllama[
{"prompt": "Who are you?", "response": "I am an AI assistant."},
{"prompt": "What can you do?", "response": "I can answer questions and assist with tasks."}
]cargo run -p rusty-cli --release -- ./models/tinyllama ./data/train.json| Model | Status |
|---|---|
| LLaMA / LLaMA-2 / LLaMA-3 | ✓ Supported |
| TinyLlama | ✓ Supported |
| Mistral | ✓ Supported |
| Phi / Phi-2 / Phi-3 | ✓ Supported |
| Qwen / Qwen-2 | ✓ Supported |
| Gemma / Gemma-2 | ✓ Supported |
- Rust 1.75 or later
- GPU: Apple Silicon (M1/M2/M3) or Vulkan-capable GPU
- OS: macOS, Linux, or Windows
| Component | Status |
|---|---|
| GPU Backend | ✓ Complete |
| Custom WGSL Kernels | ✓ Complete |
| Llama Architecture | ✓ Complete |
| LoRA Fine-tuning | ✓ Complete |
| Autograd + Optimizers | ✓ Complete |
| Safetensors Loading | ✓ Complete |
| Mixed Precision (FP16) | ✓ Complete |
| Flash Attention | ✓ Complete |
| CUDA Backend | Planned |
| Distributed Training | Planned |
Contributions are welcome. See CONTRIBUTING.md for guidelines.
git clone https://github.com/YOUR_USERNAME/rusty.git
git checkout -b feature/your-feature
cargo test --workspace
git commit -m "feat: description"
git push origin feature/your-featureMIT License — see LICENSE.
- wgpu — Cross-platform GPU API
- safetensors — Safe tensor serialization
- Flash Attention — Memory-efficient attention
Built with Rust