VRAM Calculator — Kotlin source

// vram-calculator — Kotlin port: estimate the VRAM an LLM needs (weights + KV cache).

import kotlin.math.abs

/** Supported weight quantizations. bytesPerParam = bytes stored per weight
 *  (Q4_K_M = 4.85 bits/weight, the llama.cpp mix). */
enum class Quant(val bytesPerParam: Double) {
    FP32(4.0), FP16(2.0), BF16(2.0), INT8(1.0), INT4(0.5), Q4_K_M(4.85 / 8)
}

/** Default attention architecture: a modern GQA-style layout. Override per
 *  model — e.g. Llama-2-70B uses 80 layers with the same GQA shape. */
object DefaultArch {
    const val LAYERS = 32
    const val KV_HEADS = 8
    const val HEAD_DIM = 128
}

/** Bytes per KV-cache element (fp16 K and V tensors) unless overridden. */
const val DEFAULT_KV_BYTES = 2

const val GB = 1e9

/** Architecture and serving knobs; every field must be >= 1. */
data class VramOptions(
    val layers: Int = DefaultArch.LAYERS,
    val kvHeads: Int = DefaultArch.KV_HEADS,
    val headDim: Int = DefaultArch.HEAD_DIM,
    val batch: Int = 1,   // sequences served concurrently; multiplies the KV cache
    val kvBytes: Int = DEFAULT_KV_BYTES,
)

/** One VRAM estimate, in decimal gigabytes. */
data class VramBreakdown(
    val quant: Quant,
    val bytesPerParam: Double,
    val weightsGB: Double,
    val kvCacheGB: Double,
) {
    val totalGB: Double get() = weightsGB + kvCacheGB
}

/** A GPU memory tier, and one card scored against a total. */
data class GpuCard(val name: String, val sizeGB: Double)
data class GpuFit(val name: String, val sizeGB: Double, val fits: Boolean, val headroomGB: Double)

/** Common GPU memory tiers, from consumer boards to datacenter cards. */
val GPU_CARDS = listOf(
    GpuCard("RTX 3060 Ti / RTX 4060 / RX 7600", 8.0),
    GpuCard("RTX 3060 12 GB / RTX 4070", 12.0),
    GpuCard("RTX 4060 Ti 16 GB / RTX 5080", 16.0),
    GpuCard("RTX 3090 / RTX 4090", 24.0),
    GpuCard("RTX A6000 / L40S", 48.0),
    GpuCard("A100 80 GB / H100 / H200", 80.0),
)

/**
 * Estimate the VRAM footprint of a model: weights plus KV cache.
 *
 * weightsGB = paramsB × bytesPerParam
 * kvCacheGB = 2 × layers × context × kvHeads × headDim × kvBytes × batch / 1e9
 *
 * Throws IllegalArgumentException on paramsB ≤ 0, negative context, or any
 * option below 1 (context 0 is allowed — no context, no cache).
 */
fun vram(paramsB: Double, quant: Quant, context: Int, opts: VramOptions = VramOptions()): VramBreakdown {
    require(paramsB.isFinite() && paramsB > 0) { "paramsB must be finite > 0 (got $paramsB)" }
    require(context >= 0) { "context must be >= 0 (got $context)" }
    require(opts.layers >= 1 && opts.kvHeads >= 1 && opts.headDim >= 1 && opts.batch >= 1 && opts.kvBytes >= 1) {
        "every option must be >= 1 (got $opts)"
    }
    val weightsGB = paramsB * 1e9 * quant.bytesPerParam / GB
    val kvCacheGB = 2.0 * opts.layers * context * opts.kvHeads * opts.headDim * opts.kvBytes * opts.batch / GB
    return VramBreakdown(quant, quant.bytesPerParam, weightsGB, kvCacheGB)
}

/** Score every card against a total footprint. `fits` is inclusive: a total
 *  exactly equal to the card size fits (headroom 0). */
fun gpuFits(totalGB: Double, cards: List<GpuCard> = GPU_CARDS): List<GpuFit> {
    require(totalGB.isFinite() && totalGB >= 0) { "totalGB must be finite >= 0 (got $totalGB)" }
    return cards.map { c -> GpuFit(c.name, c.sizeGB, c.sizeGB >= totalGB, c.sizeGB - totalGB) }
}

fun main() {
    // 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.
    val r70 = vram(70.0, Quant.Q4_K_M, 4096, VramOptions(layers = 80))
    val r8 = vram(8.0, Quant.FP16, 8192)
    println("70B q4_K_M @4096 (80 layers): %.2f GB weights + %.2f GB KV = %.2f GB total".format(r70.weightsGB, r70.kvCacheGB, r70.totalGB))
    println("8B  fp16    @8192:            %.2f GB weights + %.2f GB KV = %.2f GB total".format(r8.weightsGB, r8.kvCacheGB, r8.totalGB))
    println("GPU fit for %.2f GB:".format(r8.totalGB))
    for (f in gpuFits(r8.totalGB)) {
        println("  %-36s %s (%.2f GB %s)".format(f.name, if (f.fits) "fits" else "too small", abs(f.headroomGB), if (f.fits) "headroom" else "short"))
    }
}