llmgw/internal/modelplan/size.go

// Package modelplan works out whether a model fits some hardware, and how it
// should be split when it does.
//
// It exists because that arithmetic was the real barrier to running a model
// here. Everything else about adding one is typing: the name, the path, the
// engine. Deciding tp, gpus_per_node, gpu_class and the context that will
// actually fit means knowing the model's layer count, its KV head count, what
// its weights weigh at a given quantisation, and how much of a card is left
// after the runtime takes its share — and getting it wrong is not an error
// message, it is an out-of-memory twenty minutes into a load.
package modelplan

import "fmt"

// Model is what a model's own configuration says about it.
//
// Every field here is read from the model rather than guessed: config.json
// carries the architecture, and the repository's index carries the parameter
// count. A guess would defeat the point, because the answer is only worth
// having if it is the one the hardware will agree with.
type Model struct {
	Name string
	// Params is the total parameter count.
	Params int64
	// Layers, KVHeads and HeadDim size the KV cache, which is what scales with
	// context and is usually what actually runs a card out of memory.
	Layers  int
	KVHeads int
	HeadDim int
	// MaxContext is the model's trained context length.
	MaxContext int
}

// Precision is how many bytes one parameter weighs.
//
// Named rather than numeric because the names are what people use, and because
// the difference between them is the difference between a model fitting and
// not.
type Precision struct {
	Name           string
	BytesPerParam  float64
	BytesPerKVElem float64
}

// Known precisions, from full to the quantisations worth serving.
var (
	BF16   = Precision{"bf16", 2, 2}
	FP8    = Precision{"fp8", 1, 1}
	INT4   = Precision{"int4", 0.5, 2} // weights quantised, KV usually is not
	Q8GGUF = Precision{"q8_0", 1.06, 2}
	Q4GGUF = Precision{"q4_k_m", 0.61, 2}
)

// PrecisionByName resolves a precision, reporting what it knows when it cannot.
func PrecisionByName(s string) (Precision, error) {
	for _, p := range []Precision{BF16, FP8, INT4, Q8GGUF, Q4GGUF} {
		if p.Name == s {
			return p, nil
		}
	}
	return Precision{}, fmt.Errorf("unknown precision %q; known: bf16, fp8, int4, q8_0, q4_k_m", s)
}

// Demand is what one replica needs, in bytes.
type Demand struct {
	Weights  int64
	KVCache  int64
	Overhead int64
}

// Total is everything one replica needs across all its GPUs.
func (d Demand) Total() int64 { return d.Weights + d.KVCache + d.Overhead }

// PerGPU divides the demand over a replica's GPUs.
//
// Tensor parallelism splits both the weights and the KV cache, so both divide.
// The overhead does not: each process pays its own runtime, CUDA context and
// activation buffers, so it is charged per GPU rather than shared.
func (d Demand) PerGPU(gpus int) int64 {
	if gpus <= 0 {
		gpus = 1
	}
	return (d.Weights+d.KVCache)/int64(gpus) + d.Overhead
}

// KVBytesPerToken is what one token of context costs, for the whole model.
//
// Two for the key and the value, times the layers, times the heads the cache
// actually keeps. Grouped-query attention is why KVHeads and not the attention
// head count: a 14B model with 40 attention heads and 8 KV heads pays for
// eight, and the difference is a factor of five.
func (m Model) KVBytesPerToken(p Precision) int64 {
	return int64(2 * float64(m.Layers) * float64(m.KVHeads) * float64(m.HeadDim) * p.BytesPerKVElem)
}

// Needs computes what one replica of this model needs.
//
// context is the total token pool the engine will hold, which is the context
// length times however many requests it serves at once — the number an engine
// is configured with, not the number a single request sees.
func (m Model) Needs(p Precision, context int) Demand {
	weights := int64(float64(m.Params) * p.BytesPerParam)
	kv := int64(context) * m.KVBytesPerToken(p)
	return Demand{
		Weights:  weights,
		KVCache:  kv,
		Overhead: overheadBytes(weights),
	}
}

// overheadBytes is what the runtime costs beyond weights and cache.
//
// A flat allowance plus a fraction of the weights: the flat part is the CUDA
// context, the framework and its allocator, which a small model pays as much as
// a large one; the fraction stands in for activations and fragmentation, which
// do scale. Deliberately generous. A plan that says a model fits and is wrong
// costs a twenty-minute load and an OOM; one that says it does not quite fit
// costs a smaller quantisation.
func overheadBytes(weights int64) int64 {
	const flat = 2 << 30 // 2 GiB
	return flat + weights/10
}

// GB renders bytes the way people discuss hardware.
func GB(b int64) float64 { return float64(b) / (1 << 30) }