Memory Profiling¶

Accurate memory prediction is essential for choosing quantisation levels, setting context lengths, and provisioning hardware. This page derives the formulas for each memory component and shows how to use Zig's built-in tools to detect leaks and measure peak usage.

Zig's GeneralPurposeAllocator and Leak Detection¶

Zig's GeneralPurposeAllocator (GPA) is a debug allocator that tracks every allocation and reports leaks when deinit is called:

var gpa = std.heap.GeneralPurposeAllocator(.{
    .stack_trace_frames = 8,  // capture 8 frames per allocation
}){};
defer {
    const status = gpa.deinit();
    if (status == .leak) {
        std.log.err("LEAK DETECTED", .{});
    }
}
const allocator = gpa.allocator();

When a leak is found, GPA prints the stack trace of the offending allocation:

error: memory leak at 0x7f3a2c001000 (4096 bytes)
  src/models/llama.zig:142:25: LLaMAModel.init
  src/server/http_server.zig:182:19: ZigLlamaServer.loadModel
  examples/main.zig:23:5: main

Workflow¶

1. Run the test suite under GPA: zig build test.
2. If GPA reports zero leaks, the allocator contract is satisfied.
3. For long-running server processes, periodically log gpa.total_requested_bytes to monitor growth.

Model Memory Budget¶

Parameter Count Formula¶

For a standard LLaMA-style transformer with \(L\) layers and embedding dimension \(d\):

This accounts for four projection matrices per layer (\(W^Q\), \(W^K\), \(W^V\), \(W^O\), each \(d \times d\)), the FFN gate and up projections (\(d \times \tfrac{8}{3}d\) each), and the down projection.

Weight Memory¶

Given the parameter count and the number of bits per weight (\(b\)):

Activation Memory¶

During a forward pass, each transformer layer produces intermediate activations. For a single token at position \(n\) in a model with dimension \(d\):

Concretely, the dominant activations are:

For a 7B model (\(d=4096\), \(H=32\)) at context length \(n=2048\):

• Attention scores: \(4 \times 32 \times 2048^2 \approx 537\;\text{MB}\) (per layer, without KV cache)
• QKV projections: \(12 \times 2048 \times 4096 \approx 96\;\text{MB}\)
• FFN intermediate: \(\tfrac{32}{3} \times 2048 \times 4096 \approx 89\;\text{MB}\)

In practice, ZigLlama only materialises one layer's activations at a time, reusing the buffer across layers. Peak activation memory is therefore a single layer's worth, not \(L\) layers.

KV Cache Memory¶

The KV cache stores the key and value tensors for all layers and all past positions. For \(L\) layers, \(H\) heads, head dimension \(d_h\), and context length \(n\):

The factor of 2 accounts for keys and values; the factor of 4 is the byte width of f32.

Peak Memory Analysis¶

The total peak memory during inference is:

For a 7B Q4_K_M model at context length 2048:

pie title Peak Memory Breakdown (7B Q4_K_M)
    "Weights" : 3.9
    "KV Cache" : 1.07
    "Activations" : 0.15
    "Overhead" : 0.05

Optimisation Strategies¶

1. Quantise Weights¶

The most impactful lever. Moving from FP16 to Q4_K_M reduces weight memory by 3.6x.

2. Quantise the KV Cache¶

Storing keys and values in FP16 halves cache memory with negligible quality loss. INT8 KV cache (experimental) reduces it by 4x.

3. Reduce Context Length¶

If your application generates short responses (e.g., classification, entity extraction), set max_seq_len to the minimum sufficient value. Going from 4096 to 512 reduces KV cache by 8x.

4. Share Activations Across Layers¶

ZigLlama already does this: a single activation buffer is allocated and reused for every layer. Ensure custom model implementations follow the same pattern.

5. Memory-Mapped Model Loading¶

mmap avoids doubling memory during loading (one copy on disk, one in RAM). The OS page cache serves as the in-memory copy.

6. Monitor with GPA¶

// Periodically log memory usage
const bytes = gpa.total_requested_bytes;
std.log.info("Memory in use: {d:.1f} MB", .{
    @as(f64, @floatFromInt(bytes)) / (1024 * 1024),
});

Memory Planning Calculator¶

Use this formula to estimate whether a model fits in your available RAM:

Source Reference¶