Performance tuning

# Before: BF16 model (14 GB)
huggingface-cli download mlx-community/Qwen3-4B-bf16 \
    --local-dir ./models/Qwen3-4B

# After: 8-bit quantized (4 GB)
huggingface-cli download mlx-community/Qwen3-4B-8bit \
    --local-dir ./models/Qwen3-4B-8bit

use mlx_rs_core::{KVCache, KeyValueCache};

// Before: Allocates on every token
let mut cache: Vec<Option<ConcatKeyValueCache>> = vec![None; num_layers];

// After: Pre-allocates in 256-token chunks
let mut cache: Vec<Option<KVCache>> = vec![None; num_layers];

use mlx_rs_core::fused_swiglu;

// Before: Separate silu + multiply (slow)
let gate_activated = gate.silu()?;
let output = gate_activated.multiply(&x)?;

// After: Fused Metal kernel (10x faster)
let output = fused_swiglu(&x, &gate)?;

┌─────────────────────────────────────┐
│       Unified Memory (128 GB)       │
│                                     │
│  ┌──────────┐      ┌─────────────┐ │
│  │   CPU    │◄────►│   GPU/ANE   │ │
│  │  Cores   │      │   (Metal)   │ │
│  └──────────┘      └─────────────┘ │
└─────────────────────────────────────┘

# Real-time memory monitoring
while true; do
    echo "$(date): $(vm_stat | grep 'Pages free' | awk '{print $3}' | sed 's/\.//')" 
    sleep 1
done

# Or use Activity Monitor (GUI)

// Batch size 1 (default, interactive)
let input = prompt.index((0, NewAxis));  // [1, seq_len]

// Batch size 4 (offline processing)
let batch = stack(&[prompt1, prompt2, prompt3, prompt4], 0)?;  // [4, seq_len]

// Fastest: Greedy (argmax)
let generator = Generate::new(&mut model, &mut cache, 0.0, &prompt);

// Slower: Sampling with temperature
let generator = Generate::new(&mut model, &mut cache, 0.7, &prompt);

// Slowest: Top-p/top-k sampling (not yet implemented)

use mlx_rs::transforms::{eval, async_eval};

// Operations are queued, not executed
let a = x.multiply(&y)?;
let b = a.add(&z)?;
let c = b.relu()?;

// Execute all at once (fused into single kernel)
eval(&[&c])?;

// Or async for non-blocking
let _ = async_eval([&c]);

// Standard MHA: num_heads = 32
let q_heads = 32;
let kv_heads = 32;  // Same as Q
let memory_per_layer = batch * kv_heads * seq_len * head_dim * 4;  // 2 bytes × 2 (K+V)

// GQA: num_kv_heads = 8 (4x reduction)
let q_heads = 32;
let kv_heads = 8;  // 4x fewer
let memory_per_layer = batch * kv_heads * seq_len * head_dim * 4;  // 4x less memory

// 45 total experts, but only top-k=2 are active per token
let expert_outputs = Vec::new();
for expert_idx in top_k_indices {
    let expert = &self.experts[expert_idx];
    // Only 2 experts run → 22x less compute than dense
    expert_outputs.push(expert.forward(x)?);
}

// Pre-compute image embeddings once
let dino_features = dino_encoder.forward(&image)?;  // [1, 256, 1024]
let siglip_features = siglip_encoder.forward(&image)?;  // [1, 256, 1152]
let visual_tokens = projector.forward(&concat)?;  // [1, 256, 4096]

// Reuse for multiple text generations
for prompt in prompts {
    let input_ids = build_input(visual_tokens.clone(), prompt)?;
    let output = decoder.forward(&input_ids, &mut cache)?;
}

use std::time::Instant;

let start = Instant::now();
let mut num_tokens = 0;

for token in generator.take(100) {
    let _ = token?;
    num_tokens += 1;
}

let elapsed = start.elapsed().as_secs_f32();
let tok_per_sec = num_tokens as f32 / elapsed;
println!("Speed: {:.1} tok/s", tok_per_sec);

let prefill_start = Instant::now();
let first_token = generator.next().unwrap()?;  // Prefill
let prefill_time = prefill_start.elapsed();

let decode_start = Instant::now();
let remaining: Vec<_> = generator.take(99).collect();
let decode_time = decode_start.elapsed();

println!("Prefill: {:.0}ms", prefill_time.as_millis());
println!("Decode: {:.1} tok/s", 99.0 / decode_time.as_secs_f32());

# Record trace
xctrace record --template 'Metal System Trace' \
    --output trace.trace \
    --target-stdout - \
    --launch -- cargo run --release --example generate

# Open in Instruments
open trace.trace

# Debug: 10x slower, large memory overhead
cargo run --example generate

# Release: Full optimizations
cargo run --release --example generate