llmgw/internal/modelplan/options.go

package modelplan

import (
	"fmt"
	"regexp"
	"sort"
	"strings"
)

// Options: the ways the files a repository actually has could run on the
// cluster as it is — each file on one card, on several of a node's cards, on
// several nodes' cards over a fabric, or on a CPU node — with what each card
// would hold, and whether the cards are free now or only exist.

// Artifact is something a repository offers to serve: one GGUF file (or one
// split into parts), or its original safetensors weights.
type Artifact struct {
	// Name is the file, or the set's name; for original weights, "original".
	Name   string `json:"name"`
	Format string `json:"format"` // "gguf" or "safetensors"
	Bytes  int64  `json:"bytes"`
	Shards int    `json:"shards,omitempty"`
	// Quant is the quantisation read from the name, like Q4_K_M; "bf16" for
	// original weights.
	Quant string `json:"quant,omitempty"`
}

// Shape is what sizes a model's KV cache, and constrains how it splits.
type Shape struct {
	// KVBytesPerToken is one token of context, the whole model.
	KVBytesPerToken int64
	// KVHeads must divide by a tensor-parallel width; 0 is unknown.
	KVHeads int
	// Params is the parameter count, for bits per weight; 0 is unknown.
	Params int64
	// Experts is a mixture-of-experts model's routed experts; 0 for dense.
	Experts int
	// ActiveShare is the part of the weights a token reads; 0 is all of
	// them.
	ActiveShare float64
}

// share is ActiveShare, all of the weights when not known.
func (sh Shape) share() float64 {
	if sh.ActiveShare <= 0 || sh.ActiveShare > 1 {
		return 1
	}
	return sh.ActiveShare
}

// PoolNode is one ready node, as a plan sees it.
type PoolNode struct {
	ID       string
	Class    string // "" for a node with no GPU
	Fabric   string
	GPUMemGB int
	GPUs     int
	FreeGPUs int
	// HostMemGB is its RAM, and HostFreeGB what is unused now.
	HostMemGB  int
	HostFreeGB int
	// Kinds are the node's devices by class, for a node with more than one:
	// a card in MIG mode is one or several classes of instance, and a whole
	// card beside it another. Empty means Class, GPUMemGB, GPUs and FreeGPUs
	// say it all.
	Kinds []PoolKind
	// Engines are the engines the node has installed.
	Engines []string
	// CC is its newest card's compute capability, 0 when not known.
	CC float64
	// CPUs, MemBWGBs and MemBWCoreGBs are its CPUs and how fast its memory
	// reads, with every CPU and with one: a model's speed on its CPU. Zero
	// when not measured.
	CPUs                   int
	MemBWGBs, MemBWCoreGBs float64
	// CardBWGBs is a card's memory bandwidth by class, for its speed there.
	CardBWGBs map[string]float64
}

// PoolKind is one class of a node's devices.
type PoolKind struct {
	Class    string
	GPUMemGB int
	GPUs     int
	FreeGPUs int
}

// kinds is the node's devices by class.
func (n PoolNode) kinds() []PoolKind {
	if len(n.Kinds) > 0 {
		return n.Kinds
	}
	return []PoolKind{{Class: n.Class, GPUMemGB: n.GPUMemGB, GPUs: n.GPUs, FreeGPUs: n.FreeGPUs}}
}

// Kinds is a node's devices by class: its one class, or each of several.
func Kinds(n PoolNode) []PoolKind { return n.kinds() }

// isMIG reports whether a class names MIG instances: "a100-80g-mig-3g.40gb".
// A replica runs on exactly one, never a group of them, so a plan offers one
// instance and nothing wider. Read from the class label, as the card's memory
// is, which keeps this package free of the agent's.
func isMIG(class string) bool { return strings.Contains(class, "-mig-") }

// Option is one way to run one artifact.
type Option struct {
	Artifact string `json:"artifact"`
	Format   string `json:"format"`
	Engine   string `json:"engine"`
	CPU      bool   `json:"cpu,omitempty"`
	Class    string `json:"gpu_class,omitempty"`
	Fabric   string `json:"fabric,omitempty"`
	// GPUsPerNode cards on each of Nodes nodes.
	GPUsPerNode int `json:"gpus_per_node,omitempty"`
	Nodes       int `json:"nodes,omitempty"`
	// NeedGB is what one card holds, or for CPU the whole replica in RAM;
	// HaveGB what a card, or the node, offers for it.
	NeedGB float64 `json:"need_gb"`
	HaveGB float64 `json:"have_gb"`
	// Hosts is how many nodes could run it at all, and FreeNow whether one
	// (or enough, together) could right now.
	Hosts   int  `json:"hosts"`
	FreeNow bool `json:"free_now"`
	// BitsPerWeight is the artifact's precision, 0 when unknown.
	BitsPerWeight float64 `json:"bits_per_weight,omitempty"`
	// What NeedGB is made of, a card's share (or, on the CPU, the node's
	// RAM): the weights, the context's KV cache, and the engine's own.
	WeightsGB float64 `json:"weights_gb,omitempty"`
	ContextGB float64 `json:"context_gb,omitempty"`
	EngineGB  float64 `json:"engine_gb,omitempty"`
	// ExpertParallel spreads a mixture-of-experts model's experts over the
	// cards, whole, rather than slicing each across them: vLLM's
	// --enable-expert-parallel, for a model of experts on more than one card.
	ExpertParallel bool `json:"expert_parallel,omitempty"`
	// StrataFamily and StrataSize are Strata's --family and --size, for an
	// option on Strata: one card, and the experts in the node's RAM, which
	// NeedGB and HaveGB are then about.
	StrataFamily string `json:"strata_family,omitempty"`
	StrataSize   string `json:"strata_size,omitempty"`
	// TokensPerSecond is about how many tokens a second one reply gets this
	// way, on the fastest host it fits; 0 when not known. On the CPU with
	// the threads a replica is given by default: half the node's.
	TokensPerSecond float64 `json:"tokens_per_second,omitempty"`
}

// GPUs is every card one replica takes.
func (o Option) GPUs() int { return o.GPUsPerNode * max(o.Nodes, 1) }

// TotalGB is what one replica takes of every card together, or of the
// node's RAM on the CPU.
func (o Option) TotalGB() float64 {
	if o.CPU {
		return o.NeedGB
	}
	return o.NeedGB * float64(o.GPUs())
}

// Fits is whether a card, or the node, holds what it needs.
func (o Option) Fits() bool { return o.NeedGB <= o.HaveGB }

// WithContext is the option for another context: kvPerToken a token's KV
// cache, the whole model; the weights and the engine's share stay.
func (o Option) WithContext(kvPerToken int64, tokens int) Option {
	if kvPerToken <= 0 || tokens <= 0 || o.Engine == "strata" {
		return o
	}
	ctx := GB(int64(tokens) * kvPerToken)
	if !o.CPU {
		ctx /= float64(max(o.GPUs(), 1))
	}
	o.NeedGB += ctx - o.ContextGB
	o.ContextGB = ctx
	return o
}

// MostContext is the longest context, up to trained, this option still holds
// for one request, in whole thousands of tokens below trained: trained when
// the KV cache's size is unknown, 0 when trained is, for Strata (its context
// is its own), or when not even 1024 tokens fit.
func (o Option) MostContext(kvPerToken int64, trained int) int {
	switch {
	case trained <= 0 || o.Engine == "strata":
		return 0
	case kvPerToken <= 0 || o.WithContext(kvPerToken, trained).Fits():
		return trained
	}
	lo, hi := 0, trained/1024+1 // lo fits (or is none), hi does not
	for hi-lo > 1 {
		mid := (lo + hi) / 2
		if o.WithContext(kvPerToken, mid*1024).Fits() {
			lo = mid
		} else {
			hi = mid
		}
	}
	return lo * 1024
}

// MemoryText says what it needs, part by part.
func (o Option) MemoryText() string {
	if o.Engine == "strata" {
		return fmt.Sprintf("%.0f GB of the node's RAM for its experts, of %.0f usable, and one card", o.NeedGB, o.HaveGB)
	}
	parts := fmt.Sprintf("weights %.1f + context %.1f + engine %.1f", o.WeightsGB, o.ContextGB, o.EngineGB)
	if o.WeightsGB == 0 && o.ContextGB == 0 {
		parts = ""
	}
	if o.CPU {
		s := fmt.Sprintf("%.1f GB of RAM", o.NeedGB)
		if parts != "" {
			s += " (" + parts + ")"
		}
		return s + fmt.Sprintf(" of %.0f usable", o.HaveGB)
	}
	s := fmt.Sprintf("%.1f GB a card", o.NeedGB)
	if parts != "" {
		s += " (" + parts + ")"
	}
	s += fmt.Sprintf(" of %.0f usable", o.HaveGB)
	if o.GPUs() > 1 {
		s += fmt.Sprintf("; %.0f GB on %d cards", o.TotalGB(), o.GPUs())
	}
	return s
}

// Label says where it runs, in a few words.
func (o Option) Label() string {
	if o.CPU {
		return "CPU"
	}
	if o.Nodes > 1 {
		return fmt.Sprintf("%d nodes × %d %s on %s", o.Nodes, o.GPUsPerNode, o.Class, o.Fabric)
	}
	if o.Engine == "strata" {
		return fmt.Sprintf("1 × %s with Strata, experts in RAM", o.Class)
	}
	return fmt.Sprintf("%d × %s", o.GPUsPerNode, o.Class)
}

// gpuShare is how much of a card a plan may count on, as in Fits.
const gpuShare = usableFraction

// hostShare is how much of a node's RAM one replica may be planned to use:
// the system and the agent keep the rest.
const hostShare = 0.85

// ggufOverhead is llama.cpp's own cost beyond weights and cache, as the GGUF
// header's MemoryBytes counts it: compute buffers and the server.
func ggufOverhead(weights int64) int64 { return 1<<30 + weights/20 }

// Options lists every way each artifact could run, best first.
//
// context is the token pool: llama.cpp's --ctx-size, or vLLM's
// --max-model-len for one sequence.
func Options(arts []Artifact, sh Shape, nodes []PoolNode, context int) []Option {
	var out []Option
	for _, a := range arts {
		bpw := quantBits[strings.ToUpper(a.Quant)]
		if sh.Params > 0 {
			bpw = float64(a.Bytes) * 8 / float64(sh.Params)
		}
		kv := int64(context) * sh.KVBytesPerToken
		switch a.Format {
		case "gguf":
			out = append(out, ggufOptions(a, kv, nodes, bpw)...)
		case "safetensors":
			out = append(out, vllmOptions(a, kv, sh, nodes, bpw)...)
		}
	}
	for i := range out {
		out[i].TokensPerSecond = speedOf(out[i], bytesOf(arts, out[i].Artifact), sh, nodes)
	}
	Rank(out)
	return out
}

func ggufOptions(a Artifact, kv int64, nodes []PoolNode, bpw float64) []Option {
	var out []Option
	// The CPU: the whole replica in a node's RAM.
	total := a.Bytes + kv + ggufOverhead(a.Bytes)
	cpu := Option{Artifact: a.Name, Format: a.Format, Engine: "llamacpp", CPU: true,
		NeedGB: GB(total), BitsPerWeight: bpw,
		WeightsGB: GB(a.Bytes), ContextGB: GB(kv), EngineGB: GB(ggufOverhead(a.Bytes))}
	for _, n := range nodes {
		have := float64(n.HostMemGB) * hostShare
		if n.HostMemGB <= 0 || GB(total) > have {
			continue
		}
		cpu.Hosts++
		cpu.HaveGB = max(cpu.HaveGB, have)
		if float64(n.HostFreeGB) >= GB(total) {
			cpu.FreeNow = true
		}
	}
	if cpu.Hosts > 0 {
		out = append(out, cpu)
	}
	// GPUs: one node's cards, the model split across them by layer.
	for _, c := range classes(nodes) {
		budget := float64(c.mem) * gpuShare
		for width := 1; width <= c.maxGPUs; width++ {
			per := (a.Bytes+kv)/int64(width) + ggufOverhead(a.Bytes/int64(width))
			if GB(per) > budget {
				continue
			}
			o := Option{Artifact: a.Name, Format: a.Format, Engine: "llamacpp", Class: c.class,
				GPUsPerNode: width, Nodes: 1, NeedGB: GB(per), HaveGB: budget, BitsPerWeight: bpw,
				WeightsGB: GB(a.Bytes / int64(width)), ContextGB: GB(kv / int64(width)),
				EngineGB: GB(ggufOverhead(a.Bytes / int64(width)))}
			for _, fab := range c.fabrics {
				for _, n := range fab.nodes {
					if n.GPUs >= width {
						o.Hosts++
						if n.FreeGPUs >= width {
							o.FreeNow = true
						}
					}
				}
			}
			out = append(out, o)
			break // more cards than the fewest that hold it buys nothing here
		}
	}
	return out
}

func vllmOptions(a Artifact, kv int64, sh Shape, nodes []PoolNode, bpw float64) []Option {
	var out []Option
	perCard := func(width int) int64 {
		return (a.Bytes+kv)/int64(width) + overheadBytes(a.Bytes/int64(width))
	}
	divides := func(w int) bool { return SplitsHeads(sh.KVHeads, w) }
	for _, c := range classes(nodes) {
		budget := float64(c.mem) * gpuShare
		opt := func(perNode, n int, fabric string) Option {
			cards := int64(perNode * n)
			return Option{Artifact: a.Name, Format: a.Format, Engine: "vllm", Class: c.class, Fabric: fabric,
				GPUsPerNode: perNode, Nodes: n, NeedGB: GB(perCard(perNode * n)), HaveGB: budget, BitsPerWeight: bpw,
				WeightsGB: GB(a.Bytes / cards), ContextGB: GB(kv / cards), EngineGB: GB(overheadBytes(a.Bytes / cards)),
				ExpertParallel: sh.Experts > 1 && perNode*n > 1}
		}
		// One node: the narrowest power of two that holds it.
		single := false
		for w := 1; w <= c.maxGPUs && divides(w); w *= 2 {
			if GB(perCard(w)) > budget {
				continue
			}
			o := opt(w, 1, "")
			for _, fab := range c.fabrics {
				for _, nd := range fab.nodes {
					if nd.GPUs >= w {
						o.Hosts++
						o.FreeNow = o.FreeNow || nd.FreeGPUs >= w
					}
				}
			}
			out = append(out, o)
			single = true
			break
		}
		// A MIG instance does not join others, on its node or across a
		// fabric: what does not fit one does not fit the class.
		if single || isMIG(c.class) {
			continue
		}
		// Several nodes, a node's full width on each, pipeline-parallel
		// between them: only within one fabric.
		w := 1
		for w*2 <= c.maxGPUs {
			w *= 2
		}
		if !divides(w) {
			continue
		}
		for _, fab := range c.fabrics {
			if fab.name == "" {
				continue
			}
			// Any number of nodes: pipeline stages need not be a power of
			// two, and a third node is what a model just too big for two
			// needs.
			for n := 2; n <= len(fab.nodes); n++ {
				if GB(perCard(w*n)) > budget {
					continue
				}
				o := opt(w, n, fab.name)
				free := 0
				for _, nd := range fab.nodes {
					if nd.GPUs >= w {
						o.Hosts++
					}
					if nd.FreeGPUs >= w {
						free++
					}
				}
				if o.Hosts < n {
					break
				}
				o.FreeNow = free >= n
				out = append(out, o)
				break
			}
		}
	}
	return out
}

type classPool struct {
	class   string
	mem     int
	maxGPUs int
	fabrics []fabricPool
}

type fabricPool struct {
	name  string
	nodes []PoolNode
}

// classes groups GPU nodes by card, and each by fabric: a replica spans
// nodes only within one. A node with devices of several classes is in each
// class's pool, as the node it is for that class: its devices of that class
// and no others.
func classes(nodes []PoolNode) []classPool {
	by := map[string]*classPool{}
	var order []string
	for _, node := range nodes {
		for _, k := range node.kinds() {
			if k.Class == "" || k.GPUs == 0 || k.GPUMemGB == 0 {
				continue
			}
			n := node
			n.Class, n.GPUMemGB, n.GPUs, n.FreeGPUs, n.Kinds = k.Class, k.GPUMemGB, k.GPUs, k.FreeGPUs, nil
			if isMIG(n.Class) {
				// One instance a replica, however many the node has.
				n.GPUs, n.FreeGPUs = min(n.GPUs, 1), min(n.FreeGPUs, 1)
			}
			c := by[n.Class]
			if c == nil {
				c = &classPool{class: n.Class, mem: n.GPUMemGB}
				by[n.Class] = c
				order = append(order, n.Class)
			}
			c.maxGPUs = max(c.maxGPUs, n.GPUs)
			i := -1
			for j := range c.fabrics {
				if c.fabrics[j].name == n.Fabric {
					i = j
				}
			}
			if i < 0 {
				c.fabrics = append(c.fabrics, fabricPool{name: n.Fabric})
				i = len(c.fabrics) - 1
			}
			c.fabrics[i].nodes = append(c.fabrics[i].nodes, n)
		}
	}
	sort.Strings(order)
	out := make([]classPool, 0, len(order))
	for _, k := range order {
		out = append(out, *by[k])
	}
	return out
}

// quantBits is roughly what each quantisation spends a weight, for when the
// parameter count is not known to work it out from the file's size.
var quantBits = map[string]float64{
	"F32": 32, "F16": 16, "BF16": 16, "FP16": 16, "Q8_0": 8.5, "Q6_K": 6.6,
	"Q5_K_M": 5.7, "Q5_K_S": 5.5, "Q5_0": 5.5, "Q5_1": 6, "Q4_K_M": 4.8, "Q4_K_S": 4.6, "Q4_0": 4.5, "Q4_1": 5,
	"IQ4_XS": 4.3, "IQ4_NL": 4.5, "Q3_K_L": 4.3, "Q3_K_M": 3.9, "Q3_K_S": 3.5, "IQ3_M": 3.7, "IQ3_XS": 3.3,
	"IQ3_XXS": 3.1, "Q2_K": 3.0, "IQ2_M": 2.7, "IQ2_XS": 2.3, "IQ2_XXS": 2.1, "IQ1_M": 1.8, "IQ1_S": 1.6,
}

// qualityCap is the precision past which a bigger file is not a better
// model worth its memory: eight bits keep what sixteen do, near enough. On a
// CPU the cap is lower: there every bit is read from memory for every token,
// so a file twice the size answers at half the speed, and a Q4_K_M or Q5 is
// the usual balance.
const (
	qualityCap    = 8.6
	cpuQualityCap = 5.8
)

// quality scores an option's precision: rising to the cap, falling past it.
func quality(o Option) float64 {
	if o.BitsPerWeight == 0 {
		return 0
	}
	limit := qualityCap
	if o.CPU {
		limit = cpuQualityCap
	}
	if o.BitsPerWeight <= limit {
		return o.BitsPerWeight
	}
	return limit - (o.BitsPerWeight-limit)*0.25
}

// Rank orders options best first by what the hardware is, not by what it is
// doing: on a GPU before the CPU; the fewest cards, then nodes; then the
// better precision, up to eight bits. Whether the cards are free now only
// breaks a tie, then the most room left.
//
// Busy is not unfit. Ranking what was free first made every file of every
// repository come out on the CPU on a cluster whose one card was in use —
// which turned down the very file the card was there for. A model served onto
// cards that are busy waits for one, or with its fallbacks runs elsewhere
// until then.
func Rank(opts []Option) {
	sort.SliceStable(opts, func(i, j int) bool {
		a, b := opts[i], opts[j]
		if a.CPU != b.CPU {
			return !a.CPU
		}
		if a.GPUs() != b.GPUs() {
			return a.GPUs() < b.GPUs()
		}
		if a.Nodes != b.Nodes {
			return a.Nodes < b.Nodes
		}
		if qa, qb := quality(a), quality(b); qa != qb {
			return qa > qb
		}
		if a.FreeNow != b.FreeNow {
			return a.FreeNow
		}
		return a.HaveGB-a.NeedGB > b.HaveGB-b.NeedGB
	})
}

var quantPattern = regexp.MustCompile(`(?i)(IQ\d_[A-Z]+|Q\d_K_[SML]|Q\d_K|Q\d_\d|Q\d|BF16|F16|F32|FP16)`)

// QuantOf reads a quantisation out of a file name.
func QuantOf(name string) string {
	m := quantPattern.FindAllString(name, -1)
	if len(m) == 0 {
		return ""
	}
	return strings.ToUpper(m[len(m)-1])
}

// MemFromClass reads the card size out of a class label such as "h100-80g",
// and a MIG instance's out of its profile: 40 for "a100-80g-mig-3g.40gb",
// whose card is 80 but whose instance holds 40 of it.
func MemFromClass(class string) int {
	if _, profile, ok := strings.Cut(class, "-mig-"); ok {
		var compute, gb int
		if _, err := fmt.Sscanf(profile, "%dg.%dgb", &compute, &gb); err != nil {
			return 0
		}
		return gb
	}
	i := strings.LastIndex(class, "-")
	if i < 0 || !strings.HasSuffix(class, "g") {
		return 0
	}
	var gb int
	if _, err := fmt.Sscanf(class[i+1:], "%dg", &gb); err != nil {
		return 0
	}
	return gb
}

// CardNeed is what one card of a replica holds, in GB: the weights and a
// context's KV cache divided over its cards, and the engine's own.
func CardNeed(engine string, weights, kv int64, cards int) float64 {
	cards = max(cards, 1)
	per := weights / int64(cards)
	over := overheadBytes(per)
	if engine == "llamacpp" {
		over = ggufOverhead(per)
	}
	return GB(per + kv/int64(cards) + over)
}

// CardBudget is what a card of a class may be planned to hold, in GB; 0 for
// a class whose memory its name does not say.
func CardBudget(class string) float64 { return float64(MemFromClass(class)) * gpuShare }

// SplitsHeads reports whether tensor parallel can spread a model of kvHeads
// KV heads over width cards. It splits the heads among the cards, or past
// their number gives each card a copy of one: so the width divides the
// heads, or the heads divide the width. Qwen3-235B has four and runs on
// eight; a model with 8 splits 2, 4, 8 or 16 ways, not 12. 0 heads is
// unknown, and allowed.
func SplitsHeads(kvHeads, width int) bool {
	return kvHeads == 0 || width <= 1 || kvHeads%width == 0 || width%kvHeads == 0
}

func bytesOf(arts []Artifact, name string) int64 {
	for _, a := range arts {
		if a.Name == name {
			return a.Bytes
		}
	}
	return 0
}

// speedOf is about how many tokens a second one reply gets in o, on the
// fastest node it fits: on the CPU by the node's memory and the threads a
// replica is given, on cards by their bandwidth.
func speedOf(o Option, bytes int64, sh Shape, nodes []PoolNode) float64 {
	best := 0.0
	for _, n := range nodes {
		var tps float64
		switch {
		case o.CPU:
			if float64(n.HostMemGB)*hostShare < o.NeedGB {
				continue
			}
			tps = CPUTokensPerSecondOn(int64(float64(bytes)*sh.share()), n.MemBWGBs, n.MemBWCoreGBs, CPUThreads(0, n.CPUs))
		case o.Engine != "strata":
			bw := n.CardBWGBs[o.Class]
			if bw <= 0 {
				continue
			}
			tps = GPUTokensPerSecond(o, bytes, sh.share(), sh.Params, bw, Links{})
		}
		best = max(best, tps)
	}
	return best
}