llmgw/internal/modelplan/options_test.go

package modelplan

import (
	"strings"
	"testing"
)

const gib = int64(1) << 30

// A 14B GGUF: 48 layers, 8 KV heads of 128 — 192 KB a token at f16.
var shape14B = Shape{KVBytesPerToken: 48 * 8 * 256 * 2, KVHeads: 8, Params: 14_800_000_000}

func labels(opts []Option) []string {
	var out []string
	for _, o := range opts {
		out = append(out, o.Artifact+" "+o.Label())
	}
	return out
}

func TestEachFileIsMatchedToTheCardsThatHoldIt(t *testing.T) {
	nodes := []PoolNode{
		{ID: "gpu1", Class: "6000-24g", GPUMemGB: 24, GPUs: 1, FreeGPUs: 1, HostMemGB: 64, HostFreeGB: 40},
		{ID: "cpu1", HostMemGB: 32, HostFreeGB: 30},
	}
	arts := []Artifact{
		{Name: "q4.gguf", Format: "gguf", Bytes: 9 * gib},
		{Name: "q8.gguf", Format: "gguf", Bytes: 15 * gib},
		{Name: "f16.gguf", Format: "gguf", Bytes: 29 * gib},
	}
	opts := Options(arts, shape14B, nodes, 16384)
	got := strings.Join(labels(opts), "; ")
	// The best: the better precision on the one card, free now.
	if opts[0].Artifact != "q8.gguf" || opts[0].CPU || !opts[0].FreeNow {
		t.Errorf("first: %+v (all: %s)", opts[0], got)
	}
	// f16 fits no card, only a CPU node with the RAM.
	for _, o := range opts {
		if o.Artifact == "f16.gguf" && !o.CPU {
			t.Errorf("f16 placed on a card: %+v", o)
		}
	}
	if !strings.Contains(got, "f16.gguf CPU") {
		t.Errorf("f16 not offered the CPU: %s", got)
	}
	// The card busy: still the card first, said to be busy, and the CPU,
	// free, after it — for a fallback.
	nodes[0].FreeGPUs = 0
	opts = Options(arts, shape14B, nodes, 16384)
	if opts[0].CPU || opts[0].FreeNow || opts[0].Artifact != "q8.gguf" {
		t.Errorf("with the card busy: %+v", opts[0])
	}
	cpuAt := -1
	for i, o := range opts {
		if o.CPU && o.FreeNow {
			cpuAt = i
			break
		}
	}
	if cpuAt < 0 {
		t.Errorf("no free CPU way: %s", strings.Join(labels(opts), "; "))
	}
}

func TestALlamaCppModelSpreadsOverANodesCardsByLayer(t *testing.T) {
	nodes := []PoolNode{{ID: "a", Class: "l4-24g", GPUMemGB: 24, GPUs: 4, FreeGPUs: 4}}
	opts := Options([]Artifact{{Name: "70b-q4.gguf", Format: "gguf", Bytes: 40 * gib}},
		Shape{KVBytesPerToken: 80 * 8 * 256 * 2}, nodes, 8192)
	if len(opts) != 1 || opts[0].GPUsPerNode != 3 || opts[0].Nodes != 1 {
		t.Fatalf("%+v", opts)
	}
}

func TestOriginalWeightsSpanNodesOnlyOverAFabric(t *testing.T) {
	var nodes []PoolNode
	for _, id := range []string{"a", "b"} {
		nodes = append(nodes, PoolNode{ID: id, Class: "h100-80g", GPUMemGB: 80, GPUs: 8, FreeGPUs: 8, Fabric: "ib-0"})
	}
	llama70 := Shape{KVBytesPerToken: 80 * 8 * 256 * 2, KVHeads: 8, Params: 70e9}
	opts := Options([]Artifact{{Name: "original", Format: "safetensors", Bytes: 140 * gib}}, llama70, nodes, 8192)
	if len(opts) != 1 || opts[0].GPUsPerNode != 4 || opts[0].Nodes != 1 || opts[0].Engine != "vllm" {
		t.Fatalf("70B: %+v", opts)
	}
	big := Shape{KVBytesPerToken: 126 * 8 * 256 * 2, KVHeads: 8, Params: 405e9}
	opts = Options([]Artifact{{Name: "original", Format: "safetensors", Bytes: 810 * gib}}, big, nodes, 8192)
	if len(opts) != 1 || opts[0].Nodes != 2 || opts[0].GPUsPerNode != 8 || opts[0].Fabric != "ib-0" || !opts[0].FreeNow {
		t.Fatalf("405B: %+v", opts)
	}
	// No fabric, no replica across the two.
	for i := range nodes {
		nodes[i].Fabric = ""
	}
	if opts = Options([]Artifact{{Name: "original", Format: "safetensors", Bytes: 810 * gib}}, big, nodes, 8192); len(opts) != 0 {
		t.Fatalf("spanned nodes with no fabric: %+v", opts)
	}
}

func TestQuantisationsAreReadFromNames(t *testing.T) {
	for name, want := range map[string]string{
		"Qwen2.5-7B-Instruct-Q4_K_M.gguf":                       "Q4_K_M",
		"qwen2.5-coder-14b-instruct-q4_k_m-00001-of-00002.gguf": "Q4_K_M",
		"phi-4-IQ3_XS.gguf":                                     "IQ3_XS",
		"model-f16.gguf":                                        "F16",
		"JiRackUltra_1b.Q6_K.gguf":                              "Q6_K",
	} {
		if got := QuantOf(name); got != want {
			t.Errorf("%s: %s", name, got)
		}
	}
}

func TestOnTheCPUAMiddlingQuantisationIsBestAndOnACardEightBits(t *testing.T) {
	nodes := []PoolNode{{ID: "c", HostMemGB: 256, HostFreeGB: 200}}
	arts := []Artifact{
		{Name: "q4", Format: "gguf", Bytes: 9 * gib}, {Name: "q5", Format: "gguf", Bytes: 10 * gib},
		{Name: "q8", Format: "gguf", Bytes: 15 * gib}, {Name: "f16", Format: "gguf", Bytes: 28 * gib},
	}
	if o := Options(arts, shape14B, nodes, 16384); o[0].Artifact != "q5" {
		t.Errorf("CPU: %s first (%s)", o[0].Artifact, strings.Join(labels(o), "; "))
	}
	nodes = []PoolNode{{ID: "g", Class: "a100-80g", GPUMemGB: 80, GPUs: 1, FreeGPUs: 1}}
	if o := Options(arts, shape14B, nodes, 16384); o[0].Artifact != "q8" {
		t.Errorf("GPU: %s first", o[0].Artifact)
	}
}

// A model of experts spread over more than one card has them spread whole
// (expert parallelism); on one card there is nothing to spread.
func TestAModelOfExpertsSpreadsItsExpertsOverItsCards(t *testing.T) {
	nodes := []PoolNode{{ID: "a", Class: "h100-80g", GPUMemGB: 80, GPUs: 8, FreeGPUs: 8}}
	moe := Shape{KVBytesPerToken: 94 * 4 * 256 * 2, KVHeads: 4, Params: 235_000_000_000, Experts: 128}
	opts := Options([]Artifact{{Name: "original", Format: "safetensors", Bytes: 470 * gib, Quant: "bf16"}}, moe, nodes, 8192)
	if len(opts) == 0 || opts[0].GPUsPerNode < 2 || !opts[0].ExpertParallel {
		t.Fatalf("options %+v", opts)
	}
	small := Options([]Artifact{{Name: "original", Format: "safetensors", Bytes: 30 * gib, Quant: "bf16"}},
		Shape{KVHeads: 4, Experts: 128}, nodes, 8192)
	if len(small) == 0 || small[0].GPUsPerNode != 1 || small[0].ExpertParallel {
		t.Errorf("one card: %+v", small)
	}
	dense := Options([]Artifact{{Name: "original", Format: "safetensors", Bytes: 140 * gib, Quant: "bf16"}},
		Shape{KVHeads: 8}, nodes, 8192)
	if len(dense) == 0 || dense[0].ExpertParallel {
		t.Errorf("a dense model: %+v", dense)
	}
}

// Too big for two nodes, it is offered on three, not only on four.
func TestAModelTooBigForTwoNodesIsOfferedOnThree(t *testing.T) {
	var nodes []PoolNode
	for _, id := range []string{"a", "b", "c", "d"} {
		nodes = append(nodes, PoolNode{ID: id, Class: "l40s-48g", GPUMemGB: 48, GPUs: 8, FreeGPUs: 8, Fabric: "ib-0"})
	}
	// About 900 GB: more than two nodes' 16 cards hold, less than three's 24.
	opts := Options([]Artifact{{Name: "original", Format: "safetensors", Bytes: 900 * gib, Quant: "bf16"}},
		Shape{KVHeads: 8}, nodes, 4096)
	if len(opts) == 0 || opts[0].Nodes != 3 || opts[0].GPUsPerNode != 8 || opts[0].Fabric != "ib-0" {
		t.Fatalf("options %+v", labels(opts))
	}
}

// Each way says what it needs part by part, and for another context.
func TestAnOptionSaysWhatItNeedsOfEachCardAndOfAll(t *testing.T) {
	nodes := []PoolNode{{ID: "a", Class: "h100-80g", GPUMemGB: 80, GPUs: 8, FreeGPUs: 8}}
	kvPerToken := int64(80 * 8 * 128 * 2 * 2) // a 70B: 320 KB a token
	opts := Options([]Artifact{{Name: "original", Format: "safetensors", Bytes: 140 * gib, Quant: "bf16"}},
		Shape{KVBytesPerToken: kvPerToken, KVHeads: 8}, nodes, 8192)
	o := opts[0]
	if o.GPUsPerNode != 4 {
		t.Fatalf("expected four cards: %+v", o)
	}
	sum := o.WeightsGB + o.ContextGB + o.EngineGB
	if d := sum - o.NeedGB; d > 0.01 || d < -0.01 || o.WeightsGB != 35 {
		t.Errorf("parts %.2f+%.2f+%.2f do not make %.2f", o.WeightsGB, o.ContextGB, o.EngineGB, o.NeedGB)
	}
	text := o.MemoryText()
	for _, want := range []string{"GB a card", "weights 35.0", "context 0.6", "of 72 usable", "on 4 cards"} {
		if !strings.Contains(text, want) {
			t.Errorf("%q lacks %q", text, want)
		}
	}
	// Sixteen times the context: the weights stay, the context grows, and
	// it no longer fits.
	long := o.WithContext(kvPerToken, 8192*16)
	if long.WeightsGB != o.WeightsGB || long.ContextGB < o.ContextGB*15.9 {
		t.Errorf("for a longer context: %+v", long)
	}
	if !o.Fits() {
		t.Error("the planned way does not fit")
	}
	huge := o.WithContext(kvPerToken, 8192*256)
	if huge.Fits() {
		t.Errorf("a 2M-token context fits four cards: %s", huge.MemoryText())
	}
}