llmgw/internal/strata/models.go

package strata

import (
	"fmt"
	"slices"
	"strings"
)

// What Strata runs, as its setup (0.2.7, the version llmgw installs) lists
// it: a few model families, each in a few sizes, with what a size takes.
// Written down here so a model is checked when it is added rather than when
// a node's setup refuses it minutes later, and so the plan can offer Strata
// for a model it runs, with the memory to reserve.

// Family is a model Strata runs, from one publisher's files.
type Family struct {
	// Name is setup's --family; Model the name Strata gives the model.
	Name, Model string
	Title       string
	// Repo is the Hugging Face repository its files are in; empty when
	// Strata fetches them itself from several.
	Repo string
	// Subdir is whether each size's files are in a directory of its own,
	// named as the size, in the repository.
	Subdir bool
	// Engine is a family that runs on an engine path of its own, with a
	// download of its own (DeepSeek V4.1).
	Engine bool
	// MinCC is the compute capability its engine needs of an NVIDIA card,
	// above Strata's own 6.0; 0 for none. DeepSeek's engine is built for
	// RTX 30 (8.6) and newer only.
	MinCC float64
}

// Size is one of a family's sizes: setup's --model, which llmgw's profile
// says as --size.
type Size struct {
	Name     string
	Families []string
	// DownloadGB is what is fetched; DiskGB what it takes on the node once
	// prepared, when more than the download.
	DownloadGB, DiskGB float64
	// RAMGB is the host memory that keeps every expert in RAM: what a
	// replica reserves. Strata can read the rest from the disk with less,
	// much slower.
	RAMGB float64
	// NVIDIAOnly is a size whose kernels only NVIDIA cards have.
	NVIDIAOnly bool
	// Repo is a size's own repository, for a ready-made pack Strata
	// fetches itself.
	Repo string
}

// Families are the families Strata 0.2.7 runs.
var Families = []Family{
	{Name: "qwen", Model: "qwen3.8-flash-next", Title: "Qwen3.8-Flash-Next", Repo: "ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUF", Subdir: true},
	{Name: "swift", Model: "swift-1.5", Title: "Swift 1.5 (Qwen3.8-Flash-Next, shorter thinking)", Repo: "ukisai/Swift-1.5-Qwen3.8-Flash-Next-GSQ-RCO-GGUF"},
	{Name: "coder", Model: "qwen3.8-flash-next-coder", Title: "Qwen3.8-Flash-Next Coder", Repo: "ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-Coder-GGUF", Subdir: true},
	{Name: "unsloth", Model: "qwen3.8-flash-next-unsloth", Title: "Qwen3.8-Flash-Next (Unsloth)", Repo: "unsloth/Qwen3.8-Flash-Next-GGUF", Subdir: true},
	{Name: "gyro", Model: "qwen3.8-flash-next-gyro", Title: "Qwen3.8-Flash-Next (Gyro)", Repo: "agentionai/Qwen3.8-Flash-Next-Gyro-GGUF"},
	{Name: "nvfp4", Model: "qwen3.8-flash-next-nvfp4", Title: "Qwen3.8-Flash-Next (NVFP4)"},
	{Name: "qwen36", Model: "qwen3.6-35b-a3b", Title: "Qwen3.6-35B-A3B", Repo: "unsloth/Qwen3.6-35B-A3B-MTP-GGUF"},
	{Name: "deepseek41", Model: "deepseek-v4.1-flash", Title: "DeepSeek V4.1-Flash", Repo: "ISTA-DASLab/DeepSeek-V4.1-Flash-1017-GSQ-RCO-GGUF", Engine: true, MinCC: 8.6},
}

// Sizes are the sizes Strata 0.2.7 runs, and which families have each.
var Sizes = []Size{
	{Name: "Q2_0", Families: []string{"qwen"}, DownloadGB: 66.4, RAMGB: 48},
	{Name: "IQ2_XS", Families: []string{"qwen", "swift"}, DownloadGB: 68.0, RAMGB: 48},
	{Name: "IQ3_XXS", Families: []string{"qwen", "swift"}, DownloadGB: 75.8, RAMGB: 60},
	{Name: "IQ3_S", Families: []string{"qwen", "swift"}, DownloadGB: 83.6, RAMGB: 62},
	{Name: "IQ1_M", Families: []string{"coder"}, DownloadGB: 58.4, RAMGB: 32},
	{Name: "UD-IQ4_XS", Families: []string{"unsloth"}, DownloadGB: 93.7, RAMGB: 80},
	{Name: "UD-Q4_K_XL", Families: []string{"unsloth", "qwen36"}, DownloadGB: 111.3, RAMGB: 94},
	{Name: "UD-Q5_K_XL", Families: []string{"unsloth"}, DownloadGB: 158.3, RAMGB: 108},
	{Name: "UD-Q6_K_XL", Families: []string{"unsloth", "qwen36"}, DownloadGB: 169.2, RAMGB: 118},
	{Name: "Q8_0", Families: []string{"unsloth", "qwen36"}, DownloadGB: 188.2, RAMGB: 136},
	{Name: "Gyro-S", Families: []string{"gyro"}, DownloadGB: 58.5, RAMGB: 27.6},
	{Name: "Gyro-M", Families: []string{"gyro"}, DownloadGB: 92.0, RAMGB: 35},
	{Name: "NVFP4", Families: []string{"nvfp4"}, DownloadGB: 128, RAMGB: 80, NVIDIAOnly: true,
		Repo: "Maximilian228/Qwen3.8-Flash-Next-NVFP4-GPTQ-Strata"},
	{Name: "NVFP4-NVIDIA", Families: []string{"nvfp4"}, DownloadGB: 128, RAMGB: 80, NVIDIAOnly: true,
		Repo: "Maximilian228/Qwen3.8-Flash-Next-NVIDIA-NVFP4-Strata"},
	{Name: "NVFP4-QAD", Families: []string{"nvfp4"}, DownloadGB: 102.2, DiskGB: 212.2, RAMGB: 80, NVIDIAOnly: true,
		Repo: "nvidia/Qwen3.8-Flash-Next-FP4-Local"},
	{Name: "IST-DAS-1.75bpw", Families: []string{"deepseek41"}, DownloadGB: 239, RAMGB: 34},
	{Name: "IST-DAS-2.00bpw", Families: []string{"deepseek41"}, DownloadGB: 257, RAMGB: 40},
	{Name: "IST-DAS-2.25bpw", Families: []string{"deepseek41"}, DownloadGB: 274, RAMGB: 44},
	{Name: "UD-Q8_K_XL", Families: []string{"deepseek41"}, DownloadGB: 520, RAMGB: 306},
}

// qwen36 is smaller than the Flash-Next sizes of the same names: its own
// figures.
var qwen36RAM = map[string]float64{"UD-Q4_K_XL": 32, "UD-Q6_K_XL": 44, "Q8_0": 52}
var qwen36Download = map[string]float64{"UD-Q4_K_XL": 22.9, "UD-Q6_K_XL": 32.6, "Q8_0": 37.8}

// FamilyNamed is a family by its --family.
func FamilyNamed(name string) (Family, bool) {
	i := slices.IndexFunc(Families, func(f Family) bool { return f.Name == name })
	if i < 0 {
		return Family{}, false
	}
	return Families[i], true
}

// SizeOf is a family's size by its --size, with that family's figures.
func SizeOf(family, size string) (Size, error) {
	if _, ok := FamilyNamed(family); !ok {
		return Size{}, fmt.Errorf("strata runs no family %q; it runs %s", family, strings.Join(familyNames(), ", "))
	}
	var names []string
	for _, s := range Sizes {
		if !slices.Contains(s.Families, family) {
			continue
		}
		names = append(names, s.Name)
		if s.Name == size {
			if family == "qwen36" {
				s.RAMGB, s.DownloadGB = qwen36RAM[size], qwen36Download[size]
			}
			return s, nil
		}
	}
	return Size{}, fmt.Errorf("strata's %s comes in %s, not %q", family, strings.Join(names, ", "), size)
}

// ForRepo is the family whose files a Hugging Face repository holds.
func ForRepo(repo string) (Family, bool) {
	i := slices.IndexFunc(Families, func(f Family) bool { return f.Repo != "" && strings.EqualFold(f.Repo, repo) })
	if i < 0 {
		return Family{}, false
	}
	return Families[i], true
}

// SizesOf are a family's sizes, with its figures.
func SizesOf(family string) []Size {
	var out []Size
	for _, s := range Sizes {
		if slices.Contains(s.Families, family) {
			got, _ := SizeOf(family, s.Name)
			out = append(out, got)
		}
	}
	return out
}

// MemoryGB is what a replica of a size reserves of the host: its experts,
// and the engine and server beside them.
func (s Size) MemoryGB() int { return int(s.RAMGB+0.999) + 4 }

func familyNames() []string {
	out := make([]string, len(Families))
	for i, f := range Families {
		out[i] = f.Name
	}
	return out
}

// ArgValue is a flag's value in a profile's arguments, "" when absent.
func ArgValue(args []string, flag string) string {
	for i, a := range args {
		if a == flag && i+1 < len(args) {
			return args[i+1]
		}
		if v, ok := strings.CutPrefix(a, flag+"="); ok {
			return v
		}
	}
	return ""
}

// SizeForFile is the size of a family a file is, by its name: the longest
// size name in it, so UD-Q4_K_XL is not read as Q4_K_XL's shorter cousins.
func SizeForFile(family, file string) (Size, bool) {
	upper := strings.ToUpper(file)
	var best Size
	for _, s := range SizesOf(family) {
		if strings.Contains(upper, strings.ToUpper(s.Name)) && len(s.Name) > len(best.Name) {
			best = s
		}
	}
	return best, best.Name != ""
}

// FetchesItself is a family whose files Strata's setup fetches itself, each
// pinned to a commit and checked: ready-made packs, and DeepSeek's sets with
// an engine of their own. A model of it names its repository as the source.
func (f Family) FetchesItself() bool { return f.Name == "nvfp4" || f.Engine }

// RepoOf is the repository a size of a family is in.
func RepoOf(f Family, s Size) string {
	if s.Repo != "" {
		return s.Repo
	}
	return f.Repo
}

// ForFetchedRepo is the family and size a repository is, for one Strata
// fetches itself: a ready-made pack's own, or DeepSeek's (whose size the
// repository does not say, so the smallest).
func ForFetchedRepo(repo string) (Family, Size, bool) {
	for _, f := range Families {
		if !f.FetchesItself() {
			continue
		}
		for _, s := range SizesOf(f.Name) {
			if strings.EqualFold(RepoOf(f, s), repo) {
				return f, s, true
			}
		}
	}
	return Family{}, Size{}, false
}

// Fetchable says whether a size's files can be fetched by llmgw from the
// family's repository: not where Strata fetches them itself (NVFP4's packs,
// DeepSeek's sets with an engine of their own).
func (f Family) Fetchable() bool { return f.Repo != "" && !f.Engine }

// FilesOf are the weights of a size among a repository's files: its GGUF
// shards, not the picture reader.
func FilesOf(size string, paths []string) []string {
	var out []string
	for _, p := range paths {
		u := strings.ToUpper(p)
		if strings.HasSuffix(u, ".GGUF") && !strings.Contains(u, "MMPROJ") && strings.Contains(u, strings.ToUpper(size)) {
			out = append(out, p)
		}
	}
	return out
}