VRAM Calculator — Zig source

// vram-calculator — Zig port: estimate the VRAM an LLM needs (weights + KV cache).
const std = @import("std");

/// Supported weight quantizations, in the order the UI lists them.
pub const Quant = enum { fp32, fp16, bf16, int8, int4, q4_K_M };

/// Bytes stored per weight for each format
/// (Q4_K_M = 4.85 bits/weight, the llama.cpp mix).
pub fn bytesPerParam(q: Quant) f64 {
    return switch (q) {
        .fp32 => 4,
        .fp16, .bf16 => 2,
        .int8 => 1,
        .int4 => 0.5,
        .q4_K_M => 4.85 / 8,
    };
}

/// Attention architecture and serving knobs; every field must be >= 1. The
/// defaults are a modern GQA-style layout — Llama-2-70B overrides layers to
/// 80 with the same GQA shape.
pub const VramOptions = struct {
    layers: i64 = 32,    // transformer layers (blocks)
    kv_heads: i64 = 8,   // key/value heads after GQA
    head_dim: i64 = 128, // dimension of one attention head
    batch: i64 = 1,      // sequences served concurrently; multiplies the KV cache
    kv_bytes: i64 = 2,   // bytes per KV element (fp16 K and V tensors)
};

/// One VRAM estimate, in decimal gigabytes (GB = 10^9, matching how "70B"
/// and GPU sizes are quoted).
pub const VramBreakdown = struct {
    quant: Quant,
    bytes_per_param: f64,
    weights_gb: f64,
    kv_cache_gb: f64,
    total_gb: f64,
};

pub const VramError = error{ InvalidParams, InvalidOption };

/// Estimate the VRAM footprint of a model: weights plus KV cache.
///
/// weightsGB = paramsB × bytesPerParam
/// kvCacheGB = 2 × layers × context × kvHeads × headDim × kvBytes × batch / 1e9
///
/// Errors on paramsB ≤ 0, negative context, or any option below 1
/// (context 0 is allowed — no context, no cache).
pub fn vram(params_b: f64, q: Quant, context: i64, opts: VramOptions) VramError!VramBreakdown {
    if (!std.math.isFinite(params_b) or params_b <= 0) return error.InvalidParams;
    if (context < 0) return error.InvalidParams;
    if (opts.layers < 1 or opts.kv_heads < 1 or opts.head_dim < 1 or opts.batch < 1 or opts.kv_bytes < 1)
        return error.InvalidOption;
    const bpp = bytesPerParam(q);
    const weights_gb = params_b * 1e9 * bpp / 1e9; // the 1e9s cancel
    // KV elements fit i64 for any realistic shape (70B @ ctx 4096 × batch 2 = 5.4e9).
    const kv_elems: i64 = 2 * opts.layers * context * opts.kv_heads * opts.head_dim * opts.kv_bytes * opts.batch;
    const kv_cache_gb = @as(f64, @floatFromInt(kv_elems)) / 1e9;
    return .{ .quant = q, .bytes_per_param = bpp, .weights_gb = weights_gb, .kv_cache_gb = kv_cache_gb, .total_gb = weights_gb + kv_cache_gb };
}

/// A GPU memory tier, and one card scored against a total.
pub const GpuCard = struct { name: []const u8, size_gb: f64 };
pub const GpuFit = struct { name: []const u8, size_gb: f64, fits: bool, headroom_gb: f64 };

/// Common GPU memory tiers, from consumer boards to datacenter cards.
pub const gpu_cards = [_]GpuCard{
    .{ .name = "RTX 3060 Ti / RTX 4060 / RX 7600", .size_gb = 8 },
    .{ .name = "RTX 3060 12 GB / RTX 4070", .size_gb = 12 },
    .{ .name = "RTX 4060 Ti 16 GB / RTX 5080", .size_gb = 16 },
    .{ .name = "RTX 3090 / RTX 4090", .size_gb = 24 },
    .{ .name = "RTX A6000 / L40S", .size_gb = 48 },
    .{ .name = "A100 80 GB / H100 / H200", .size_gb = 80 },
};

/// Score every card against a total footprint. `fits` is inclusive: a total
/// exactly equal to the card size fits (headroom 0). `out` must fit `cards`.
pub fn gpuFits(total_gb: f64, cards: []const GpuCard, out: []GpuFit) VramError![]GpuFit {
    if (!std.math.isFinite(total_gb) or total_gb < 0) return error.InvalidParams;
    var n: usize = 0;
    for (cards) |c| {
        out[n] = .{ .name = c.name, .size_gb = c.size_gb, .fits = c.size_gb >= total_gb, .headroom_gb = c.size_gb - total_gb };
        n += 1;
    }
    return out[0..n];
}

pub fn main() !void {
    // The lib's canonical vectors: a 70B Llama-2-shape build (80 layers) and
    // an 8B default-arch build, then the GPU fit table for the 8B total.
    const r70 = try vram(70, .q4_K_M, 4096, .{ .layers = 80 });
    const r8 = try vram(8, .fp16, 8192, .{});
    const stdout = std.io.getStdOut().writer();
    try stdout.print("70B q4_K_M @4096 (80 layers): {d:.2} GB weights + {d:.2} GB KV = {d:.2} GB total\n", .{ r70.weights_gb, r70.kv_cache_gb, r70.total_gb });
    try stdout.print("8B  fp16    @8192:            {d:.2} GB weights + {d:.2} GB KV = {d:.2} GB total\n", .{ r8.weights_gb, r8.kv_cache_gb, r8.total_gb });
    try stdout.print("GPU fit for {d:.2} GB:\n", .{r8.total_gb});
    var fits: [gpu_cards.len]GpuFit = undefined;
    for (try gpuFits(r8.total_gb, &gpu_cards, &fits)) |f| {
        try stdout.print("  {s:<36}{s} ({d:.2} GB {s})\n", .{
            f.name,
            if (f.fits) "fits" else "too small",
            @abs(f.headroom_gb),
            if (f.fits) "headroom" else "short",
        });
    }
}