GPU Support

GPU Support

Complete guide to GPU acceleration in Shimmy.

🚀 Supported GPU Types

✅ NVIDIA CUDA

• Status: Full support with --features llama-cuda
• Requirements: CUDA toolkit, NVIDIA GPU with compute capability 5.0+
• Performance: Significant acceleration for inference
• Automatic Detection: Models automatically use available VRAM

✅ NVIDIA CUDA + MOE

• Status: Hybrid CPU/GPU processing with --features llama-cuda,moe
• Benefits: Run 70B+ models on limited VRAM systems
• Configuration: --cpu-moe --n-cpu-moe 8 for optimal layer distribution
• Use Case: Large models that exceed GPU memory capacity

✅ Apple Metal

• Status: Automatic on macOS via llama.cpp Metal backend
• Support: Apple Silicon (M1/M2/M3) and discrete AMD/NVIDIA GPUs
• Performance: Metal GPU acceleration through battle-tested llama.cpp
• Configuration: No setup required

📝 MLX Backend Status: Native MLX integration is in development. Currently, Apple Silicon acceleration comes through llama.cpp's Metal backend, which provides excellent performance for GGUF models. Future MLX backend will support .npz MLX-native models.

✅ MOE Hybrid Processing

• Status: Available with --features moe
• Technology: Mixture of Experts CPU offloading
• Benefits: Memory-efficient processing of large models
• Smart Distribution: Automatically places layers on optimal hardware

✅ CPU Fallback

• Status: Always available
• Performance: Multi-threaded CPU inference
• Use Case: Systems without GPU or when GPU is unavailable

🔧 CUDA Setup

1. Install CUDA Toolkit

# Ubuntu/Debian
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.0-1_all.deb
sudo dpkg -i cuda-keyring_1.0-1_all.deb
sudo apt-get update
sudo apt-get install cuda

# Windows - Download from NVIDIA website
# https://developer.nvidia.com/cuda-downloads

2. Verify CUDA Installation

nvidia-smi
nvcc --version

3. Build Shimmy with CUDA

# CUDA only
cargo install shimmy --features llama-cuda

# CUDA + MOE hybrid (recommended for large models)
cargo install shimmy --features llama-cuda,moe

# From source
cargo build --release --features llama-cuda,moe

4. Verify GPU Usage

# Run shimmy and check output for GPU assignment
shimmy generate model-name --prompt "Test" --max-tokens 5

# Should show: "load_tensors: layer X assigned to device CUDA0"
# Instead of: "load_tensors: layer X assigned to device CPU"

🍎 Metal Setup (macOS)

Automatic Configuration

Metal GPU acceleration works automatically on macOS:

• Apple Silicon: M1, M2, M3 chips with integrated GPU
• Intel Macs: With discrete AMD or NVIDIA GPUs
• No configuration needed: Shimmy auto-detects and uses Metal

Verification

# Check for Metal acceleration in logs
shimmy generate model-name --prompt "Test" --max-tokens 5

# Look for Metal GPU initialization messages

🧠 MOE (Mixture of Experts) Setup

What is MOE?

MOE enables running large models (70B+) on consumer hardware by intelligently distributing model layers across GPU and CPU:

• Hot layers (frequently used) → GPU for speed
• Cold layers (less frequent) → CPU to save VRAM
• Automatic optimization based on model architecture and available hardware

Enable MOE

# Install with MOE support
cargo install shimmy --features moe

# Start server with MOE enabled
shimmy serve --cpu-moe --n-cpu-moe 8

# Combine with CUDA for hybrid processing
cargo install shimmy --features llama-cuda,moe
shimmy serve --cpu-moe --n-cpu-moe 8 --gpu-backend cuda

MOE Configuration

# Number of CPU MOE layers (adjust based on VRAM)
--n-cpu-moe 4    # Conservative (more GPU layers)
--n-cpu-moe 8    # Balanced (default)
--n-cpu-moe 16   # Aggressive (more CPU layers)

# Enable CPU MOE processing
--cpu-moe        # Required flag to activate MOE

When to Use MOE

• Large models: 70B parameters or more
• Limited VRAM: Less than required for full GPU processing
• Cost optimization: Reduce GPU memory requirements
• Mixed workloads: Balance speed vs memory efficiency

🐳 Docker GPU Support

NVIDIA Docker Runtime

# docker-compose.yml
services:
  shimmy:
    build: .
    runtime: nvidia
    environment:
      - NVIDIA_VISIBLE_DEVICES=all
    ports:
      - "3000:3000"

Command Line

# Run with all GPUs
docker run --runtime=nvidia --gpus all -p 3000:3000 shimmy:latest

# Run with specific GPU
docker run --runtime=nvidia --gpus device=0 -p 3000:3000 shimmy:latest

Dockerfile

FROM nvidia/cuda:12.0-devel-ubuntu22.04
# ... shimmy installation

📊 Performance Comparison

System	Model Size	CPU Time	GPU Time	Speedup
RTX 3060 12GB	Llama 7B Q4	45s	8s	5.6x
RTX 4090 24GB	Llama 13B Q4	90s	12s	7.5x
M1 Max	Llama 7B Q4	35s	15s	2.3x
M2 Ultra	Llama 13B Q4	65s	18s	3.6x

🔍 Troubleshooting GPU Issues

CUDA Not Detected

# Check CUDA installation
nvidia-smi
nvcc --version

# Rebuild shimmy with CUDA
cargo clean
cargo build --release --features llama

Out of Memory Errors

# Check GPU memory
nvidia-smi

# Use smaller model or reduce batch size
export CUDA_VISIBLE_DEVICES=0

Mixed CPU/GPU Layers

This is normal behavior - shimmy automatically splits large models:

load_tensors: layer  0 assigned to device CUDA0
load_tensors: layer  1 assigned to device CUDA0
...
load_tensors: layer 25 assigned to device CPU  # Remaining layers on CPU

Metal Not Working (macOS)

# Check system info
system_profiler SPDisplaysDataType

# Verify Metal support
xcrun metal --help

⚙️ Advanced GPU Configuration

Environment Variables

# Limit GPU memory usage
export CUDA_MEMORY_FRACTION=0.8

# Force CPU usage
export CUDA_VISIBLE_DEVICES=""

# Enable debug logging
export RUST_LOG=debug

Model-Specific Settings

Different models have different GPU memory requirements:

• 1B models: 1-2GB VRAM
• 3B models: 2-4GB VRAM
• 7B models: 4-8GB VRAM
• 13B models: 8-16GB VRAM
• 70B models: 40GB+ VRAM (multi-GPU)

Multi-GPU Support

Coming in future releases

• Model parallelism across multiple GPUs
• Automatic load balancing
• Pipeline parallelism for large models

🎯 Best Practices

Memory Management

• Monitor GPU memory usage with nvidia-smi
• Use appropriate model sizes for your GPU
• Close other GPU applications during inference

Performance Optimization

• Use Q4 quantization for best speed/quality balance
• Enable GPU acceleration for production workloads
• Use CPU fallback for development/testing

Power Consumption

• GPU inference uses more power but is much faster
• Consider power limits for mobile/edge deployments
• CPU inference is more power-efficient for small models

---

Next Steps: Model Discovery | Performance Tuning