rustytorch/docs/book/src/performance/memory.md

Memory Management

Efficient memory usage in RustyTorch++.

GPU Memory Hierarchy

┌─────────────────────────────────────────┐
│       Global Memory (HBM) - 24-80 GB    │ ← Tensor storage
├─────────────────────────────────────────┤
│       L2 Cache - 4-50 MB                │ ← Auto-managed
├─────────────────────────────────────────┤
│       Shared Memory - 48-164 KB/SM      │ ← Kernel local
├─────────────────────────────────────────┤
│       Registers - 64K/SM                │ ← Thread local
└─────────────────────────────────────────┘

Memory Monitoring

use rtx_tensor::Device;

let device = Device::cuda(0)?;

// Query memory
let total = device.total_memory()?;
let used = device.used_memory()?;
let free = device.free_memory()?;

println!("Total: {:.2} GB", total as f64 / 1e9);
println!("Used: {:.2} GB", used as f64 / 1e9);
println!("Free: {:.2} GB", free as f64 / 1e9);

Tensor Pool

How It Works

use rtx_tensor::{TensorPool, pool_stats};

// Request a buffer from the pool
let buffer = TensorPool::get([1024, 1024], &device)?;

// When buffer is dropped, memory returns to pool (not freed)
drop(buffer);

// Next allocation with same size reuses memory
let buffer2 = TensorPool::get([1024, 1024], &device)?;  // No allocation!

Pool Statistics

let stats = pool_stats();
println!("Pool size: {} MB", stats.total_size / 1_000_000);
println!("In use: {} MB", stats.used_size / 1_000_000);
println!("Cached: {} MB", stats.cached_size / 1_000_000);
println!("Hit rate: {:.1}%", stats.hit_rate * 100.0);
println!("Fragmentation: {:.1}%", stats.fragmentation * 100.0);

Pool Management

use rtx_tensor::clear_pool;

// Clear unused cached memory
clear_pool();

// Configure pool behavior
TensorPool::configure(PoolConfig {
    max_cached_bytes: 4 * 1024 * 1024 * 1024,  // 4 GB max cache
    max_split_size: 512 * 1024 * 1024,          // Max block to split
});

Gradient Checkpointing

Trade compute for memory during training.

Automatic Checkpointing

use rtx_autograd::{CheckpointManager, CheckpointStrategy};

// √n checkpoints - optimal memory/compute tradeoff
let checkpoint = CheckpointManager::new(CheckpointStrategy::SqrtN);

// Wrap expensive forward passes
let hidden = checkpoint.checkpoint(|| {
    let h1 = layer1.forward(&x)?;
    let h2 = layer2.forward(&h1)?;
    layer3.forward(&h2)
})?;

// During backward, activations are recomputed
hidden.backward()?;

Memory Savings

Strategy	Memory	Compute Overhead
None	O(n)	1x
√n checkpoints	O(√n)	~1.5x
Every-k	O(n/k)	~kx

Memory-Efficient Attention

Flash Attention

use rtx_flash_attention::FlashAttention;

// Standard attention: O(N²) memory
// Flash attention: O(N) memory

let flash = FlashAttention::new(config)?;
let output = flash.forward(&q, &k, &v)?;

Chunked Attention

// Process attention in chunks for very long sequences
let chunk_size = 1024;
let mut outputs = vec![];

for chunk_start in (0..seq_len).step_by(chunk_size) {
    let chunk_end = (chunk_start + chunk_size).min(seq_len);
    let q_chunk = q.narrow(1, chunk_start, chunk_end - chunk_start)?;
    let out_chunk = attention.forward(&q_chunk, &k, &v)?;
    outputs.push(out_chunk);
}

let output = Tensor::cat(&outputs, 1)?;

Pinned Memory

Faster CPU-GPU transfers.

use rtx_memory::PinnedBuffer;

// Allocate pinned (page-locked) host memory
let pinned = PinnedBuffer::new(size)?;

// Fill with data
pinned.copy_from_slice(&data);

// Fast async transfer to GPU
let gpu_tensor = Tensor::from_pinned_async(&pinned, shape, &device, &stream)?;

Memory Debugging

Find Memory Leaks

use rtx_memory::MemoryTracker;

// Enable tracking
MemoryTracker::enable();

// Run your code...
train_loop(&model, &data)?;

// Report allocations
let report = MemoryTracker::report();
for allocation in report.live_allocations() {
    println!("{}: {} bytes at {}",
        allocation.id,
        allocation.size,
        allocation.location
    );
}

Memory Snapshots

use rtx_memory::Snapshot;

let before = Snapshot::capture(&device)?;

// Operations...
let output = model.forward(&input)?;

let after = Snapshot::capture(&device)?;

// Compare
let diff = after.diff(&before);
println!("Memory change: {} MB", diff.total_bytes / 1_000_000);

Best Practices

1. Delete Intermediate Tensors

// Bad - keeps all intermediates
let a = model.layer1(&x)?;
let b = model.layer2(&a)?;
let c = model.layer3(&b)?;

// Good - release when not needed
let mut h = model.layer1(&x)?;
h = model.layer2(&h)?;
h = model.layer3(&h)?;

2. Use no_grad for Inference

// Training: stores activations for backward
let output = model.forward(&input)?;

// Inference: no activation storage
let output = no_grad(|| model.forward(&input))?;

3. Process in Batches

// OOM risk with large batch
let all_outputs = model.forward(&huge_batch)?;

// Safe - process in chunks
let mut outputs = vec![];
for chunk in data.chunks(safe_batch_size) {
    outputs.push(model.forward(&chunk)?);
}

Next Steps

• GPU Optimization - Compute optimization
• Profiling - Find memory bottlenecks
• Troubleshooting Memory - Fix OOM