diff --git a/go/internal/config/config.go b/go/internal/config/config.go index 7374561..fa53140 100644 --- a/go/internal/config/config.go +++ b/go/internal/config/config.go @@ -25,15 +25,48 @@ import ( // Config is the root of forge's infra config, built by LoadFromStore. type Config struct { - Server Server - Paths Paths - Slots map[string]Slot - Ports map[string]int // auxiliary services: embedding, stt, ... - Modes map[string]Mode - Scheduler SchedulerDefault - Monitor Monitor - Tailscale Tailscale - Cost Cost + Server Server + Paths Paths + Slots map[string]Slot + Ports map[string]int // auxiliary services: embedding, stt, ... + ServiceIcons map[string]string + Modes map[string]Mode + Scheduler SchedulerDefault + Monitor Monitor + Tailscale Tailscale + Cost Cost +} + +// defaultServiceIcons are the fallback vendor-icon manifest slugs (see +// web/src/assets/icons/manifest.ts) for the four fixed infra services +// (STT/Embedding/Aligner/TTS — bare [ports] entries with no catalog-backed +// model metadata, unlike service_mode rows like ComfyUI, which carry their +// own services.icon). Previously a Go literal map in +// httpapi/services_handlers.go (serviceVendorLogoBySlug) — changing which +// icon a service showed meant editing source and shipping a new binary. +// Moved into infra.service_icons (operator feedback 2026-09-06: "bad design +// to need a restart to change an icon") so it's store-backed and +// SIGHUP-reloadable like every other infra.* value, with these as the +// zero-value fallback so nothing changes for anyone who hasn't overridden +// it. The icon names the MODEL, not the company (2026-08-14 precedent) — all +// four fixed services currently run Qwen models. +var defaultServiceIcons = map[string]string{ + "STT": "qwen", + "Embedding": "qwen", + "Aligner": "qwen", + "TTS": "qwen", +} + +// DefaultServiceIcons returns a copy of the fallback vendor-icon map (see +// defaultServiceIcons above) for callers outside this package — httpapi's +// GET /api/v1/service-icons overlays these under any operator override +// without exposing the shared package-level map itself for mutation. +func DefaultServiceIcons() map[string]string { + out := make(map[string]string, len(defaultServiceIcons)) + for k, v := range defaultServiceIcons { + out[k] = v + } + return out } // Server holds listen addresses and the state-db location. Canonical V4 @@ -336,13 +369,14 @@ func LoadFromStore(ctx context.Context, st store.Store) (*Config, error) { settings := st.Settings() for key, dst := range map[string]any{ - "infra.server": &cfg.Server, - "infra.paths": &cfg.Paths, - "infra.ports": &cfg.Ports, - "infra.scheduler": &cfg.Scheduler, - "infra.monitor": &cfg.Monitor, - "infra.tailscale": &cfg.Tailscale, - "infra.cost": &cfg.Cost, + "infra.server": &cfg.Server, + "infra.paths": &cfg.Paths, + "infra.ports": &cfg.Ports, + "infra.service_icons": &cfg.ServiceIcons, + "infra.scheduler": &cfg.Scheduler, + "infra.monitor": &cfg.Monitor, + "infra.tailscale": &cfg.Tailscale, + "infra.cost": &cfg.Cost, } { if err := getSetting(ctx, settings, key, dst); err != nil { return nil, err @@ -446,6 +480,14 @@ func (c *Config) applyDefaults() { if c.Cost.MaxPowerW <= 0 { c.Cost.MaxPowerW = DefaultMaxPowerW } + if c.ServiceIcons == nil { + c.ServiceIcons = map[string]string{} + } + for name, icon := range defaultServiceIcons { + if _, ok := c.ServiceIcons[name]; !ok { + c.ServiceIcons[name] = icon + } + } for name, svc := range c.allServices() { if svc.Backend == "" { svc.Backend = "vulkan" diff --git a/go/internal/config/config_test.go b/go/internal/config/config_test.go index 20455cf..9d22b6b 100644 --- a/go/internal/config/config_test.go +++ b/go/internal/config/config_test.go @@ -19,7 +19,7 @@ func sampleConfig() Config { Paths: Paths{ModelsDir: "/opt/forge/models"}, Slots: map[string]Slot{ "a1": {Unit: "forge-a1", Port: 8080, Label: "A1", Order: 1}, - "a3": {Unit: "forge-a3", Port: 8087, Label: "A3", Order: 3}, + "a3": {Unit: "forge-a3", Port: 8087, Label: "A3", Order: 3}, }, Ports: map[string]int{"embedding": 8083, "stt": 8084}, Modes: map[string]Mode{ @@ -242,6 +242,51 @@ func TestLoadFromStoreEmpty(t *testing.T) { } } +// TestServiceIconsDefaultsAndOverride covers the fix for "changing a fixed +// infra service's icon needs a rebuild+restart" (operator feedback +// 2026-09-06): defaultServiceIcons fills in for any key the operator hasn't +// overridden, and applyDefaults must not clobber a real override — the same +// per-key-fill contract as infra.ports would need if it ever gained a +// similar default map. +func TestServiceIconsDefaultsAndOverride(t *testing.T) { + db, err := store.Open(":memory:") + if err != nil { + t.Fatalf("store.Open: %v", err) + } + defer db.Close() + ctx := context.Background() + + cfg, err := LoadFromStore(ctx, db) + if err != nil { + t.Fatalf("LoadFromStore: %v", err) + } + for name, want := range defaultServiceIcons { + if got := cfg.ServiceIcons[name]; got != want { + t.Errorf("ServiceIcons[%q] = %q, want default %q", name, got, want) + } + } + + raw, err := json.Marshal(map[string]string{"STT": "nvidia"}) + if err != nil { + t.Fatalf("marshal override: %v", err) + } + if err := db.Settings().Set(ctx, "infra.service_icons", raw); err != nil { + t.Fatalf("Settings.Set: %v", err) + } + + cfg2, err := LoadFromStore(ctx, db) + if err != nil { + t.Fatalf("LoadFromStore (with override): %v", err) + } + if cfg2.ServiceIcons["STT"] != "nvidia" { + t.Errorf("ServiceIcons[STT] = %q, want operator override %q", cfg2.ServiceIcons["STT"], "nvidia") + } + if cfg2.ServiceIcons["Embedding"] != "qwen" { + t.Errorf("ServiceIcons[Embedding] = %q, want default %q (override of one key must not drop the others)", + cfg2.ServiceIcons["Embedding"], "qwen") + } +} + // TestLoadFromStorePopulated exercises the real path: infra.* settings keys // set and slots rows present, mirroring what the Phase 2 cutover migration // will write. diff --git a/go/internal/engine/kvcache.go b/go/internal/engine/kvcache.go new file mode 100644 index 0000000..72cdffe --- /dev/null +++ b/go/internal/engine/kvcache.go @@ -0,0 +1,260 @@ +// SPDX-License-Identifier: Apache-2.0 + +package engine + +import ( + "os" + "strconv" + "strings" + + "github.com/jsaigou/the-forge/internal/gguf" +) + +// 2026-09-07 incident: gemma4-26b-a4b-nothink's catalog safe_memory_bytes +// (25 GiB, weight-adjacent, no KV term) let the fit check admit a load that +// actually needed ~40 GiB, while qwen38-flash-next's real ~90 GiB footprint +// in the other slot left only ~25.5 GiB genuinely free — the host OOM-killed +// several unrelated services before self-healing. modeNeedEstimate's three +// tiers (profile → curated catalog figure → bare weight-file size) had no +// context-dependent term anywhere, so any large-context unprofiled mode was +// exposed the same way. This file computes a real weights+KV-cache floor +// from GGUF metadata so a stale/optimistic curated number can never win. +// +// The formula is ported directly from the kintsugi fork's own C++ (verified +// 2026-09-07 against /opt/forge/llama.cpp-kintsugi/src/), not derived by +// guesswork: +// - llama-hparams.cpp: n_embd_k_gqa(il) = n_embd_head_k(il) * n_head_kv(il), +// n_embd_head_k(il) = is_swa(il) ? key_length_swa : key_length (V mirrors K). +// is_swa(il) comes straight from the GGUF's per-layer +// attention.sliding_window_pattern array when the file declares one. +// - llama-model.cpp:1329-1333: when key_length/value_length are absent, +// llama.cpp itself falls back to embedding_length / head_count — that +// exact fallback, not a guess, is mirrored below. +// - llama-kv-cache-iswa.cpp: per layer, k/v tensors are sized +// [n_embd_k_gqa(il), cells]. cells = n_ctx for non-SWA layers always, +// and for SWA layers too when --swa-full is set +// (`if (swa_full) { size_swa = size_base; }` — verified directly). +// Without --swa-full, SWA layers are windowed: +// size_swa = min(n_ctx, sliding_window + n_ubatch) padded to 256. +// +// Deliberately NOT covered: hybrid/recurrent architectures (Nemotron's +// Mamba2/attention mix, the kintsugi fork's own experimental qwen4exp / +// Qwen3.8-Flash-Next with SSM state + a block-sparse indexer cache). A flat +// per-token guess used to cover unprofiled modes and was deleted after it +// inflated Nemotron's real need by ~24x (memory.go's modeNeedEstimate doc). +// This estimator inherits that lesson: gguf.Metadata.Hybrid (set from any +// .ssm.* or .attention.indexer.* key) makes it abstain +// (ok=false) rather than apply a formula that doesn't model those +// architectures. Those modes stay on today's curated-or-refuse path, which +// is already the correct fail-closed behavior for what we can't model — +// real PROFILE measurements are the sanctioned way to size them. + +// cacheBytesPerElem is bytes-per-element for the KV-cache quant types +// llama.cpp's --cache-type-k/--cache-type-v accept, verified against +// ggml.c's type_traits table (type_size / blck_size): +// +// f32=4, f16=2, bf16=2, q8_0=34/32, q4_0=18/32, q4_1=20/32, q5_0=22/32, q5_1=24/32. +var cacheBytesPerElem = map[string]float64{ + "f32": 4, + "f16": 2, + "bf16": 2, + "q8_0": 34.0 / 32.0, + "q4_0": 18.0 / 32.0, + "q4_1": 20.0 / 32.0, + "q5_0": 22.0 / 32.0, + "q5_1": 24.0 / 32.0, +} + +// parseCacheTypes pulls --cache-type-k/--cache-type-v from a mode's +// extra_args, defaulting to llama.cpp's own default ("f16") when absent. +func parseCacheTypes(extraArgs []string) (kType, vType string) { + kType, vType = "f16", "f16" + for i, a := range extraArgs { + switch a { + case "--cache-type-k": + if i+1 < len(extraArgs) { + kType = extraArgs[i+1] + } + case "--cache-type-v": + if i+1 < len(extraArgs) { + vType = extraArgs[i+1] + } + } + } + return kType, vType +} + +// hasFlag reports whether extraArgs contains a bare flag (e.g. --swa-full). +func hasFlag(extraArgs []string, name string) bool { + for _, a := range extraArgs { + if a == name { + return true + } + } + return false +} + +// parseParallelArg extracts --parallel's value (default 1, matching +// llama.cpp's own default). Small enough to duplicate rather than export +// profile.parseParallel across packages for one helper. +func parseParallelArg(extraArgs []string) int { + for i, a := range extraArgs { + if a == "--parallel" && i+1 < len(extraArgs) { + if n, err := strconv.Atoi(extraArgs[i+1]); err == nil && n > 0 { + return n + } + } + if strings.HasPrefix(a, "--parallel=") { + if n, err := strconv.Atoi(strings.TrimPrefix(a, "--parallel=")); err == nil && n > 0 { + return n + } + } + } + return 1 +} + +// padTo256 mirrors llama.cpp's GGML_PAD(x, 256) for the windowed-SWA cell +// count (llama-kv-cache-iswa.cpp: "the SWA cache is always padded to 256"). +func padTo256(n int) int { + const pad = 256 + return ((n + pad - 1) / pad) * pad +} + +// kvCacheBytes computes the real KV-cache tensor footprint for one mode's +// configured context. nCtx is the per-sequence cell count actually passed +// to llama-server (this repo enforces --parallel 1 for any full-context +// mode — see CLAUDE.md's "--parallel" note — so the configured Context +// value already IS that cell count; --parallel>1 modes are short-context +// workers where erring high is the safe direction anyway). +// +// ok=false — never a guess — when the architecture is hybrid/recurrent +// (meta.Hybrid) or required layer metadata can't be resolved even via +// llama.cpp's own embedding_length/head_count fallback. +func kvCacheBytes(meta gguf.Metadata, nCtx int, kBpe, vBpe float64, swaFull bool) (int64, bool) { + if meta.Hybrid || meta.BlockCount <= 0 || nCtx <= 0 { + return 0, false + } + if len(meta.HeadCountKV) != meta.BlockCount { + return 0, false + } + + keyLen, valLen := meta.KeyLength, meta.ValueLength + if keyLen <= 0 && meta.HeadCount > 0 && meta.EmbeddingLength > 0 { + keyLen = meta.EmbeddingLength / meta.HeadCount + } + if valLen <= 0 && meta.HeadCount > 0 && meta.EmbeddingLength > 0 { + valLen = meta.EmbeddingLength / meta.HeadCount + } + if keyLen <= 0 || valLen <= 0 { + return 0, false + } + + hasSWAPattern := len(meta.SWAPattern) == meta.BlockCount + + var total float64 + for il := 0; il < meta.BlockCount; il++ { + isSWA := hasSWAPattern && meta.SWAPattern[il] + + hdK, hdV := keyLen, valLen + if isSWA && meta.KeyLengthSWA > 0 { + hdK = meta.KeyLengthSWA + } + if isSWA && meta.ValueLengthSWA > 0 { + hdV = meta.ValueLengthSWA + } + + cells := nCtx + if isSWA && !swaFull { + window := meta.SlidingWindow + if window <= 0 || window > nCtx { + window = nCtx + } else { + window = padTo256(window) + if window > nCtx { + window = nCtx + } + } + cells = window + } + + nHeadKV := meta.HeadCountKV[il] + total += float64(cells) * float64(nHeadKV) * (float64(hdK)*kBpe + float64(hdV)*vBpe) + } + return int64(total), true +} + +// kvAwareNeedBytes returns a mode's weights + real-KV-cache floor: the +// on-disk weight set (modeWeightBytes, already used by the file-size +// fallback tier) plus kvCacheBytes for its configured context and cache +// quant. ok=false when the model's GGUF can't be read or the architecture +// can't be modeled (kvCacheBytes' own abstention) — callers must treat that +// as "no computed floor," never as zero. +func (m *Manager) kvAwareNeedBytes(modeName string) (int64, bool) { + cfg := m.d.Cfg() + mode, ok := cfg.Modes[modeName] + if !ok || len(mode.Services) == 0 || mode.Services[0].Model == "" { + return 0, false + } + svc := mode.Services[0] + + modelPath := cfg.Paths.ResolveModelPath(svc.Model) + meta, err := m.readMetaCached(modelPath) + if err != nil { + return 0, false + } + + nCtx := svc.Context / parseParallelArg(svc.ExtraArgs) + if nCtx <= 0 { + nCtx = svc.Context + } + kName, vName := parseCacheTypes(svc.ExtraArgs) + kBpe, kOK := cacheBytesPerElem[kName] + vBpe, vOK := cacheBytesPerElem[vName] + if !kOK || !vOK { + return 0, false + } + swaFull := hasFlag(svc.ExtraArgs, "--swa-full") + + kvBytes, ok := kvCacheBytes(meta, nCtx, kBpe, vBpe, swaFull) + if !ok { + return 0, false + } + + weightBytes := modeWeightBytes(cfg, modeName) + if weightBytes <= 0 { + return 0, false + } + return weightBytes + kvBytes, true +} + +// readMetaCached wraps Deps.ReadMeta with a (size, mtime)-keyed cache: this +// is a real disk read (see gguf.go's doc comment on why it never touches +// the tensor table, only the header), and FitPlan sits on the scheduler's +// hot decision path — polled every PollInterval while a load is in +// progress. Without this, every poll would re-open and re-scan the +// candidate model's GGUF header. Falls back to an uncached read (matching +// every other ReadMeta call site) when the file can't be stat'd, so fake +// paths in tests behave exactly as before. +func (m *Manager) readMetaCached(path string) (gguf.Metadata, error) { + st, err := os.Stat(path) + if err != nil { + return m.d.ReadMeta(path) + } + + m.metaMu.Lock() + if e, ok := m.metaCache[path]; ok && e.size == st.Size() && e.modTime.Equal(st.ModTime()) { + m.metaMu.Unlock() + return e.meta, nil + } + m.metaMu.Unlock() + + meta, err := m.d.ReadMeta(path) + if err != nil { + return meta, err + } + + m.metaMu.Lock() + m.metaCache[path] = metaCacheEntry{size: st.Size(), modTime: st.ModTime(), meta: meta} + m.metaMu.Unlock() + return meta, nil +} diff --git a/go/internal/engine/kvcache_test.go b/go/internal/engine/kvcache_test.go new file mode 100644 index 0000000..bf16a38 --- /dev/null +++ b/go/internal/engine/kvcache_test.go @@ -0,0 +1,291 @@ +// SPDX-License-Identifier: Apache-2.0 + +package engine + +import ( + "os" + "path/filepath" + "testing" + + "github.com/jsaigou/the-forge/internal/config" + "github.com/jsaigou/the-forge/internal/gguf" +) + +// gemmaLikeMetadata mirrors the real gemma4-26b-a4b GGUF ground-truthed on +// ForgeHost during the 2026-09-07 incident investigation: 30 layers, a 5 SWA : +// 1 global repeating pattern, key/value_length 512 (global) vs. 256 (SWA), +// sliding_window 1024. +func gemmaLikeMetadata() gguf.Metadata { + const layers = 30 + headCountKV := make([]int, layers) + pattern := make([]bool, layers) + for i := 0; i < layers; i++ { + if i%6 == 5 { + headCountKV[i] = 2 // global layer + pattern[i] = false + } else { + headCountKV[i] = 8 // SWA layer + pattern[i] = true + } + } + return gguf.Metadata{ + Architecture: "gemma4", + BlockCount: layers, + EmbeddingLength: 2816, + HeadCount: 16, + HeadCountKV: headCountKV, + KeyLength: 512, + ValueLength: 512, + KeyLengthSWA: 256, + ValueLengthSWA: 256, + SlidingWindow: 1024, + SWAPattern: pattern, + } +} + +// gib is a var, not a const: several tests multiply it by a non-integer +// float literal and truncate back to int64 at runtime — Go rejects that +// conversion at compile time for a constant expression. +var gib = int64(1) << 30 + +func TestKVCacheBytesGemmaSWAFull(t *testing.T) { + meta := gemmaLikeMetadata() + bpe := cacheBytesPerElem["q8_0"] + got, ok := kvCacheBytes(meta, 262144, bpe, bpe, true) + if !ok { + t.Fatal("kvCacheBytes returned ok=false") + } + // Hand-calculated: 25 SWA layers + 5 global layers, all materialized to + // the full 262144 context under --swa-full ⇒ ~29.2 GiB. + want := int64(29.21875 * float64(gib)) + tolerance := gib / 100 // within ~10 MiB + if diff := got - want; diff < -tolerance || diff > tolerance { + t.Errorf("got %.3f GiB, want ~%.3f GiB (swa-full)", float64(got)/float64(gib), float64(want)/float64(gib)) + } +} + +// Without --swa-full, SWA layers are windowed (1024 cells, not 262144) — +// the resulting estimate must be dramatically smaller, illustrating exactly +// why the 2026-09-07 incident's config (which DOES set --swa-full) needed +// the full-context term and a flat weight-only guess could never see it. +func TestKVCacheBytesGemmaWindowed(t *testing.T) { + meta := gemmaLikeMetadata() + bpe := cacheBytesPerElem["q8_0"] + got, ok := kvCacheBytes(meta, 262144, bpe, bpe, false) + if !ok { + t.Fatal("kvCacheBytes returned ok=false") + } + want := int64(2.76 * float64(gib)) + tolerance := gib / 20 // within ~50 MiB — window padding leaves more slack + if diff := got - want; diff < -tolerance || diff > tolerance { + t.Errorf("got %.3f GiB, want ~%.3f GiB (windowed)", float64(got)/float64(gib), float64(want)/float64(gib)) + } + fullResult, _ := kvCacheBytes(meta, 262144, bpe, bpe, true) + if got >= fullResult { + t.Errorf("windowed (%d) must be far smaller than swa-full (%d)", got, fullResult) + } +} + +// A plain dense/GQA model (no SWA pattern at all) reduces to the standard +// 2 * n_layer * n_head_kv * head_dim * n_ctx * bytes_per_elem formula. +func TestKVCacheBytesPlainGQA(t *testing.T) { + const layers, kvHeads, headDim, nCtx = 32, 8, 128, 8192 + meta := gguf.Metadata{ + BlockCount: layers, + HeadCountKV: repeatInt(kvHeads, layers), + KeyLength: headDim, + ValueLength: headDim, + } + bpe := cacheBytesPerElem["f16"] + got, ok := kvCacheBytes(meta, nCtx, bpe, bpe, false) + if !ok { + t.Fatal("kvCacheBytes returned ok=false") + } + want := int64(layers) * int64(kvHeads) * int64(headDim) * int64(nCtx) * 2 /* K+V */ * int64(bpe) + if got != want { + t.Errorf("got %d, want %d", got, want) + } +} + +// The embedding_length/head_count fallback (llama.cpp's own default when +// key_length/value_length are absent — llama-model.cpp:1329-1333) must be +// used, not a refusal, when the explicit lengths aren't in the file. +func TestKVCacheBytesEmbeddingHeadCountFallback(t *testing.T) { + const layers, kvHeads, nHead, embd, nCtx = 4, 4, 32, 4096, 8192 + meta := gguf.Metadata{ + BlockCount: layers, + HeadCountKV: repeatInt(kvHeads, layers), + HeadCount: nHead, + EmbeddingLength: embd, + // KeyLength/ValueLength deliberately absent. + } + bpe := cacheBytesPerElem["f16"] + got, ok := kvCacheBytes(meta, nCtx, bpe, bpe, false) + if !ok { + t.Fatal("kvCacheBytes returned ok=false, want the embd/head_count fallback to resolve") + } + headDim := embd / nHead + want := int64(layers) * int64(kvHeads) * int64(headDim) * int64(nCtx) * 2 * int64(bpe) + if got != want { + t.Errorf("got %d, want %d (head_dim=%d via fallback)", got, want, headDim) + } +} + +func TestKVCacheBytesHybridAbstains(t *testing.T) { + meta := gemmaLikeMetadata() + meta.Hybrid = true + if _, ok := kvCacheBytes(meta, 262144, cacheBytesPerElem["q8_0"], cacheBytesPerElem["q8_0"], true); ok { + t.Error("ok = true, want false for a hybrid/recurrent architecture") + } +} + +func TestKVCacheBytesMissingLayerDataAbstains(t *testing.T) { + cases := []struct { + name string + meta gguf.Metadata + }{ + {"no block count", gguf.Metadata{HeadCountKV: []int{8}, KeyLength: 128, ValueLength: 128}}, + {"head_count_kv length mismatch", gguf.Metadata{BlockCount: 4, HeadCountKV: []int{8, 8}, KeyLength: 128, ValueLength: 128}}, + {"no key/value length or fallback", gguf.Metadata{BlockCount: 4, HeadCountKV: repeatInt(8, 4)}}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + if _, ok := kvCacheBytes(c.meta, 8192, 2, 2, false); ok { + t.Error("ok = true, want false") + } + }) + } +} + +func repeatInt(v, n int) []int { + out := make([]int, n) + for i := range out { + out[i] = v + } + return out +} + +func TestParseCacheTypesDefaultsToF16(t *testing.T) { + k, v := parseCacheTypes(nil) + if k != "f16" || v != "f16" { + t.Errorf("got %q/%q, want f16/f16", k, v) + } +} + +func TestParseCacheTypesReadsFlags(t *testing.T) { + args := []string{"--jinja", "--cache-type-k", "q8_0", "--cache-type-v", "q8_0", "--flash-attn", "on"} + k, v := parseCacheTypes(args) + if k != "q8_0" || v != "q8_0" { + t.Errorf("got %q/%q, want q8_0/q8_0", k, v) + } +} + +func TestHasFlagDetectsSWAFull(t *testing.T) { + if !hasFlag([]string{"--jinja", "--swa-full"}, "--swa-full") { + t.Error("hasFlag = false, want true") + } + if hasFlag([]string{"--jinja"}, "--swa-full") { + t.Error("hasFlag = true, want false") + } +} + +// writeSizedFile creates a file of exactly size bytes (sparse — content +// doesn't matter, only os.Stat's reported size, which is all +// modeWeightBytes/collector.WeightSetSizeBytes reads). +func writeSizedFile(t *testing.T, path string, size int64) { + t.Helper() + f, err := os.Create(path) + if err != nil { + t.Fatal(err) + } + defer f.Close() + if err := f.Truncate(size); err != nil { + t.Fatal(err) + } +} + +// TestModeNeedEstimateFloorsCuratedFigure reproduces the 2026-09-07 +// incident's exact shape at the modeNeedEstimate level: a curated +// safe_memory_bytes far below the real weights+KV need for a +// 262144-context, --swa-full, q8_0-cache mode. The curated figure must +// never win over the computed floor. +func TestModeNeedEstimateFloorsCuratedFigure(t *testing.T) { + cfg := testConfig(t) + modelPath := filepath.Join(cfg.Paths.ModelsDir, "gemma4.gguf") + const weightBytes = int64(18) << 30 // stand-in for the real ~18 GiB weight set + writeSizedFile(t, modelPath, weightBytes) + + mode := cfg.Modes["gemma"] + mode.ConfigID = 1 + mode.Services = []config.Service{{ + Model: "gemma4.gguf", + Alias: "gemma", + Context: 262144, + PortRole: "a1", + Backend: "vulkan", + ExtraArgs: []string{"--swa-full", "--cache-type-k", "q8_0", "--cache-type-v", "q8_0"}, + }} + cfg.Modes["gemma"] = mode + + m, _, _ := newTestManager(t, cfg, newFakeSys(), nil) + meta := gemmaLikeMetadata() + m.d.ReadMeta = func(path string) (gguf.Metadata, error) { + if path == modelPath { + return meta, nil + } + return gguf.Metadata{}, nil + } + const curatedBytes = int64(25) << 30 // the real incident's exact wrong figure + m.d.WeightEstimateBytes = func(configID int64) (int64, bool) { + if configID == 1 { + return curatedBytes, true + } + return 0, false + } + + got, ok := m.modeNeedEstimate(cfg, "gemma") + if !ok { + t.Fatal("modeNeedEstimate returned ok=false") + } + if got <= curatedBytes { + t.Fatalf("need estimate = %.1f GiB, must exceed the curated %.1f GiB (computed KV floor must win)", + float64(got)/float64(gib), float64(curatedBytes)/float64(gib)) + } + // Sanity: should land near weights (~18 GiB) + the swa-full KV estimate + // (~29.2 GiB) — i.e. real operating range, not a small nudge above 25. + wantApprox := weightBytes + int64(29.21875*float64(gib)) + tolerance := gib // 1 GiB slack + if diff := got - wantApprox; diff < -tolerance || diff > tolerance { + t.Errorf("need estimate = %.1f GiB, want ~%.1f GiB (weights + swa-full KV)", + float64(got)/float64(gib), float64(wantApprox)/float64(gib)) + } +} + +// TestModeNeedEstimateUnaffectedWhenFormulaCantResolve confirms existing +// unprofiled-mode behavior (curated figure, or file size) is completely +// unchanged when the KV formula can't resolve — e.g. every existing engine +// test's fake ReadMeta, which returns a bare gguf.Metadata{TrainedCtx: ...} +// with no layer data. +func TestModeNeedEstimateUnaffectedWhenFormulaCantResolve(t *testing.T) { + cfg := testConfig(t) + mode := cfg.Modes["gemma"] + mode.ConfigID = 1 + cfg.Modes["gemma"] = mode + + m, _, _ := newTestManager(t, cfg, newFakeSys(), nil) // default fake ReadMeta: bare TrainedCtx, no BlockCount + const curatedBytes = int64(25) << 30 + m.d.WeightEstimateBytes = func(configID int64) (int64, bool) { + if configID == 1 { + return curatedBytes, true + } + return 0, false + } + + got, ok := m.modeNeedEstimate(cfg, "gemma") + if !ok { + t.Fatal("modeNeedEstimate returned ok=false") + } + if got != curatedBytes { + t.Errorf("got %d, want the curated figure %d unchanged (formula must abstain, not guess)", got, curatedBytes) + } +} diff --git a/go/internal/engine/manager.go b/go/internal/engine/manager.go index 1f04fbb..1a062c4 100644 --- a/go/internal/engine/manager.go +++ b/go/internal/engine/manager.go @@ -107,6 +107,22 @@ type Manager struct { // mu guards slots. mu sync.Mutex slots map[string]*slotRec + + // metaMu guards metaCache, the KV-cache estimator's GGUF-header cache + // (kvcache.go). FitPlan is on the scheduler's hot decision path (polled + // every PollInterval during EnsureLoaded), and gguf.ReadMetadata has no + // cache of its own — without this, every poll would re-open and + // re-scan the candidate model's GGUF header. + metaMu sync.Mutex + metaCache map[string]metaCacheEntry +} + +// metaCacheEntry pairs a cached gguf.Metadata read with the (size, mtime) +// it was read at, so a re-written or replaced file invalidates cleanly. +type metaCacheEntry struct { + size int64 + modTime time.Time + meta gguf.Metadata } type slotRec struct { @@ -144,9 +160,10 @@ func NewManager(d Deps) (*Manager, error) { d.Logf = log.Printf } m := &Manager{ - d: d, - llama: collector.NewLlamaClient(d.BaseURL), - slots: map[string]*slotRec{}, + d: d, + llama: collector.NewLlamaClient(d.BaseURL), + slots: map[string]*slotRec{}, + metaCache: map[string]metaCacheEntry{}, } for name := range d.Cfg().Slots { m.slots[name] = &slotRec{} diff --git a/go/internal/engine/memory.go b/go/internal/engine/memory.go index 6fd230f..7b73c4b 100644 --- a/go/internal/engine/memory.go +++ b/go/internal/engine/memory.go @@ -215,11 +215,24 @@ type Plan struct { // modeNeedEstimate is FitPlan's need-bytes derivation for any mode: a // fresh profiled safe-memory figure first (weights + KV cache at max // context, measured), then the catalog's curated safe_memory_bytes, then -// the on-disk weight set. ok=false means no defensible figure exists. -// Used by FitPlan for its own mode and by the in-flight reservation logic, -// which must reserve a loading slot's eventual footprint before any of it -// is measurable (the 2026-08-22 crash admitted a second load while the -// first sibling's pages were still materializing). +// the on-disk weight set — each of the latter two floored against a real +// weights+KV-cache estimate computed from GGUF metadata (kvcache.go) +// whenever the architecture and metadata support it. ok=false means no +// defensible figure exists. Used by FitPlan for its own mode and by the +// in-flight reservation logic, which must reserve a loading slot's +// eventual footprint before any of it is measurable (the 2026-08-22 crash +// admitted a second load while the first sibling's pages were still +// materializing). +// +// 2026-09-07 incident fix: a curated safe_memory_bytes value is a +// hand-entered, often weight-adjacent guess with no guarantee it accounts +// for KV cache at the mode's configured context — gemma4-26b-a4b-nothink's +// curated 25 GiB let a load through that actually needed ~40 GiB once its +// 262144-context, --swa-full KV cache materialized, triggering a +// whole-host OOM cascade while a co-resident model's real ~90 GiB +// footprint left only ~25.5 GiB genuinely free. A stale/optimistic curated +// or file-size figure must never be allowed to undercut what +// kvAwareNeedBytes can actually compute from the model's own metadata. func (m *Manager) modeNeedEstimate(cfg *config.Config, modeName string) (int64, bool) { if _, ok := cfg.Modes[modeName]; !ok { return 0, false @@ -229,17 +242,33 @@ func (m *Manager) modeNeedEstimate(cfg *config.Config, modeName string) (int64, return b, true } } + + kvBytes, kvOK := m.kvAwareNeedBytes(modeName) + if m.d.WeightEstimateBytes != nil { configID := cfg.Modes[modeName].ConfigID if configID != 0 { if b, ok := m.d.WeightEstimateBytes(configID); ok && b > 0 { + if kvOK && kvBytes > b { + m.logf("fit %s: curated safe_memory_bytes (%.1f GiB) undercuts the computed weights+KV floor (%.1f GiB) — using the computed floor", + modeName, float64(b)/(1<<30), float64(kvBytes)/(1<<30)) + return kvBytes, true + } return b, true } } } if w := modeWeightBytes(cfg, modeName); w > 0 { + if kvOK && kvBytes > w { + m.logf("fit %s: on-disk weight size (%.1f GiB) undercuts the computed weights+KV floor (%.1f GiB) — using the computed floor", + modeName, float64(w)/(1<<30), float64(kvBytes)/(1<<30)) + return kvBytes, true + } return w, true } + if kvOK { + return kvBytes, true + } return 0, false } diff --git a/go/internal/engine/memory_incident_test.go b/go/internal/engine/memory_incident_test.go index 878bc24..3600534 100644 --- a/go/internal/engine/memory_incident_test.go +++ b/go/internal/engine/memory_incident_test.go @@ -226,6 +226,75 @@ func TestLoadRefusesSiblingWeightsWithoutRoom(t *testing.T) { } } +// TestFitPlanCatches20260907IncidentShape reproduces the exact incident +// this file's KV-cache-aware modeNeedEstimate fix (kvcache.go) closes: +// gemma4-26b-a4b-nothink's curated safe_memory_bytes (25 GiB, weight- +// adjacent, no KV term) was <= the ~25.5 GiB the host genuinely had free +// (qwen38-flash-next's real ~90 GiB footprint occupied the rest), so the +// fit check approved a load that actually needed ~47 GiB once its +// 262144-context, --swa-full KV cache materialized — triggering a +// whole-host OOM cascade. Before the fix this plan would have returned +// Fits: true; after it, the computed weights+KV floor correctly refuses. +func TestFitPlanCatches20260907IncidentShape(t *testing.T) { + cfg := testConfig(t) + modelPath := filepath.Join(cfg.Paths.ModelsDir, "gemma4.gguf") + put(t, modelPath, strings.Repeat("x", 1024)) // content irrelevant; modeWeightBytes reads size only + + mode := cfg.Modes["gemma"] + mode.ConfigID = 1 + mode.Services = []config.Service{{ + Model: "gemma4.gguf", + Alias: "gemma", + Context: 262144, + PortRole: "a1", + Backend: "vulkan", + ExtraArgs: []string{"--swa-full", "--cache-type-k", "q8_0", "--cache-type-v", "q8_0"}, + }} + cfg.Modes["gemma"] = mode + + // Engineer the exact incident margin: ~25.5 GiB genuinely free (as if + // another already-loaded model, matching the real qwen38-flash-next's + // ~90 GiB, already accounts for the rest of a ~100 GiB budget). + gpu := incidentGPU(t, int64(74500)*1024*1024, int64(100)*incidentGiB) + + meta := gemmaLikeMetadata() + m, err := NewManager(Deps{ + Cfg: func() *config.Config { return cfg }, + Sys: newFakeSys(), + GPU: gpu, + Proc: collector.Proc{Root: t.TempDir()}, + Usage: &fakeUsage{}, + Notify: ¬ifyCounter{}, + BaseURL: func(int) string { return "http://127.0.0.1:1" }, + Kill: func(int) error { return nil }, + PollInterval: time.Millisecond, + Logf: t.Logf, + ReadMeta: func(path string) (gguf.Metadata, error) { + if path == modelPath { + return meta, nil + } + return gguf.Metadata{}, nil + }, + WeightEstimateBytes: func(configID int64) (int64, bool) { + if configID == 1 { + return int64(25) * incidentGiB, true // the real incident's exact curated figure + } + return 0, false + }, + }) + if err != nil { + t.Fatal(err) + } + + plan, err := m.FitPlan("gemma") + if err != nil { + t.Fatal(err) + } + if plan.Fits { + t.Fatalf("plan = %+v, want refusal (the 2026-09-07 incident: curated 25 GiB must not pass against a real ~47 GiB need)", plan) + } +} + // Same weights WITH room proceeds past the guard into the normal flow. func TestLoadAllowsSiblingWeightsWithRoom(t *testing.T) { cfg := siblingConfig(t) diff --git a/go/internal/gguf/gguf.go b/go/internal/gguf/gguf.go index 2ba3a3d..e9896dd 100644 --- a/go/internal/gguf/gguf.go +++ b/go/internal/gguf/gguf.go @@ -30,6 +30,46 @@ type Metadata struct { ParameterCount int64 QuantType string FileSizeBytes int64 + + // The fields below back the KV-cache-aware memory estimator + // (go/internal/engine/kvcache.go). All of it lives in the KV section + // already being scanned — no tensor-table read is added. + + // BlockCount is .block_count (n_layer). + BlockCount int + // EmbeddingLength is .embedding_length. Paired with HeadCount, it + // is the head_dim fallback llama.cpp itself uses + // (n_embd_head_k_full = n_embd / n_head()) when key_length/value_length + // aren't present in the file (llama-model.cpp:1329-1333). + EmbeddingLength int + // HeadCount is .attention.head_count (n_head, the query head + // count — distinct from HeadCountKV). + HeadCount int + // HeadCountKV is .attention.head_count_kv, normalized to one entry + // per layer: a scalar in the GGUF broadcasts to len==BlockCount here; a + // real per-layer array (Gemma-style, where SWA/global layers differ) is + // kept as-is. Empty when the key is absent. + HeadCountKV []int + // KeyLength / ValueLength are .attention.key_length / + // .value_length — per-head K/V dim for non-SWA ("full") layers. + KeyLength int + ValueLength int + // KeyLengthSWA / ValueLengthSWA are the SWA-layer counterparts (0 = the + // model draws no K/V-dim distinction between SWA and full layers). + KeyLengthSWA int + ValueLengthSWA int + // SlidingWindow is .attention.sliding_window (window size, cells). + SlidingWindow int + // SWAPattern is .attention.sliding_window_pattern: true at index + // il means layer il is a windowed/SWA layer. len==BlockCount when + // present; nil means the model declares no per-layer SWA split (every + // layer is treated as "full"). + SWAPattern []bool + // Hybrid is true when the file declares any .ssm.* or + // .attention.indexer.* key — a recurrent-state or block-sparse + // architecture the KV-cache formula below does not model. Callers must + // treat this as "cannot estimate," never guess. + Hybrid bool } const ggufMagic = 0x46554747 // "GGUF" little-endian @@ -133,8 +173,18 @@ func readAll(r *bufio.Reader, fileSize int64) (Metadata, error) { } // ctxByArch holds every *.context_length seen, so key order relative to - // general.architecture doesn't matter. + // general.architecture doesn't matter. attnByArch mirrors the same + // arch-keyed-staging pattern for the KV-cache estimator's fields. ctxByArch := map[string]uint64{} + attnByArch := map[string]*archAttn{} + stageFor := func(arch string) *archAttn { + a, ok := attnByArch[arch] + if !ok { + a = &archAttn{} + attnByArch[arch] = a + } + return a + } var fileType uint64 var haveFileType bool @@ -148,6 +198,13 @@ func readAll(r *bufio.Reader, fileSize int64) (Metadata, error) { return md, fmt.Errorf("kv %q type: %w", key, err) } + // Hybrid-architecture signal: presence alone matters, not the + // value, so this doesn't consume anything — the switch below still + // reads/skips the value normally. + if strings.Contains(key, ".ssm.") || strings.Contains(key, ".attention.indexer.") { + md.Hybrid = true + } + switch { case key == "general.architecture" || key == "general.name": s, err := readTypedString(r, vt, fileSize) @@ -177,6 +234,66 @@ func readAll(r *bufio.Reader, fileSize int64) (Metadata, error) { return md, fmt.Errorf("kv %q: %w", key, err) } ctxByArch[strings.TrimSuffix(key, ".context_length")] = u + case strings.HasSuffix(key, ".block_count"): + u, err := readTypedUint(r, vt) + if err != nil { + return md, fmt.Errorf("kv %q: %w", key, err) + } + stageFor(strings.TrimSuffix(key, ".block_count")).blockCount = int(u) + case strings.HasSuffix(key, ".embedding_length"): + u, err := readTypedUint(r, vt) + if err != nil { + return md, fmt.Errorf("kv %q: %w", key, err) + } + stageFor(strings.TrimSuffix(key, ".embedding_length")).embeddingLength = int(u) + case strings.HasSuffix(key, ".attention.head_count_kv"): + vals, err := readUintOrArray(r, vt, fileSize) + if err != nil { + return md, fmt.Errorf("kv %q: %w", key, err) + } + stageFor(strings.TrimSuffix(key, ".attention.head_count_kv")).headCountKV = vals + case strings.HasSuffix(key, ".attention.head_count"): + u, err := readTypedUint(r, vt) + if err != nil { + return md, fmt.Errorf("kv %q: %w", key, err) + } + stageFor(strings.TrimSuffix(key, ".attention.head_count")).headCount = int(u) + case strings.HasSuffix(key, ".attention.key_length_swa"): + u, err := readTypedUint(r, vt) + if err != nil { + return md, fmt.Errorf("kv %q: %w", key, err) + } + stageFor(strings.TrimSuffix(key, ".attention.key_length_swa")).keyLengthSWA = int(u) + case strings.HasSuffix(key, ".attention.value_length_swa"): + u, err := readTypedUint(r, vt) + if err != nil { + return md, fmt.Errorf("kv %q: %w", key, err) + } + stageFor(strings.TrimSuffix(key, ".attention.value_length_swa")).valueLengthSWA = int(u) + case strings.HasSuffix(key, ".attention.key_length"): + u, err := readTypedUint(r, vt) + if err != nil { + return md, fmt.Errorf("kv %q: %w", key, err) + } + stageFor(strings.TrimSuffix(key, ".attention.key_length")).keyLength = int(u) + case strings.HasSuffix(key, ".attention.value_length"): + u, err := readTypedUint(r, vt) + if err != nil { + return md, fmt.Errorf("kv %q: %w", key, err) + } + stageFor(strings.TrimSuffix(key, ".attention.value_length")).valueLength = int(u) + case strings.HasSuffix(key, ".attention.sliding_window_pattern"): + vals, err := readBoolArray(r, vt, fileSize) + if err != nil { + return md, fmt.Errorf("kv %q: %w", key, err) + } + stageFor(strings.TrimSuffix(key, ".attention.sliding_window_pattern")).swaPattern = vals + case strings.HasSuffix(key, ".attention.sliding_window"): + u, err := readTypedUint(r, vt) + if err != nil { + return md, fmt.Errorf("kv %q: %w", key, err) + } + stageFor(strings.TrimSuffix(key, ".attention.sliding_window")).slidingWindow = int(u) default: if err := skipValue(r, vt, fileSize); err != nil { return md, fmt.Errorf("kv %q: %w", key, err) @@ -200,9 +317,55 @@ func readAll(r *bufio.Reader, fileSize int64) (Metadata, error) { md.QuantType = fmt.Sprintf("unknown(%d)", fileType) } } + + if a, ok := attnByArch[md.Architecture]; ok { + md.BlockCount = a.blockCount + md.EmbeddingLength = a.embeddingLength + md.HeadCount = a.headCount + md.KeyLength = a.keyLength + md.ValueLength = a.valueLength + md.KeyLengthSWA = a.keyLengthSWA + md.ValueLengthSWA = a.valueLengthSWA + md.SlidingWindow = a.slidingWindow + switch { + case len(a.headCountKV) == 1 && md.BlockCount > 0: + // Scalar head_count_kv (the common case): broadcast to every + // layer so callers never special-case "scalar vs. array". + md.HeadCountKV = make([]int, md.BlockCount) + for i := range md.HeadCountKV { + md.HeadCountKV[i] = a.headCountKV[0] + } + case len(a.headCountKV) == md.BlockCount && md.BlockCount > 0: + md.HeadCountKV = a.headCountKV + // else: length mismatch (corrupt/unexpected) — leave nil rather + // than guess which layers the values belong to. + } + if len(a.swaPattern) == md.BlockCount && md.BlockCount > 0 { + md.SWAPattern = a.swaPattern + // else: length mismatch — leave nil, which the KV-cache + // estimator reads as "no per-layer SWA split known." + } + } + return md, nil } +// archAttn stages the KV-cache-estimator fields for one architecture prefix +// while the KV section is scanned (mirrors ctxByArch: key order relative to +// general.architecture doesn't matter). +type archAttn struct { + blockCount int + embeddingLength int + headCount int + headCountKV []int + keyLength int + valueLength int + keyLengthSWA int + valueLengthSWA int + slidingWindow int + swaPattern []bool +} + func readString(r *bufio.Reader, fileSize int64) (string, error) { var n uint64 if err := binary.Read(r, binary.LittleEndian, &n); err != nil { @@ -282,6 +445,84 @@ func readTypedUint(r *bufio.Reader, vt uint32) (uint64, error) { } } +// maxLayerArray bounds how many elements a per-layer array (head_count_kv, +// sliding_window_pattern) retains — real models have at most a few hundred +// layers; this guards a corrupt/adversarial count without needing +// block_count up front (it may arrive after this key in the KV section). +const maxLayerArray = 8192 + +// readUintOrArray reads either a scalar integer or an array of integers, +// returning every value widened to int. Some converters emit a key like +// head_count_kv as a plain scalar in the common (uniform) case and as a +// real per-layer array when layers differ (Gemma's SWA vs. global layers). +func readUintOrArray(r *bufio.Reader, vt uint32, fileSize int64) ([]int, error) { + if vt != typeArray { + u, err := readTypedUint(r, vt) + if err != nil { + return nil, err + } + return []int{int(u)}, nil + } + var elemType uint32 + if err := binary.Read(r, binary.LittleEndian, &elemType); err != nil { + return nil, err + } + var count uint64 + if err := binary.Read(r, binary.LittleEndian, &count); err != nil { + return nil, err + } + if count > maxLayerArray { + return nil, fmt.Errorf("implausible array count %d", count) + } + size, ok := scalarSize[elemType] + if !ok || elemType == typeBool { + return nil, fmt.Errorf("expected numeric array element, got type %d", elemType) + } + if int64(count)*size > fileSize { + return nil, fmt.Errorf("array of %d elements exceeds file size", count) + } + out := make([]int, count) + for i := range out { + u, err := readTypedUint(r, elemType) + if err != nil { + return nil, err + } + out[i] = int(u) + } + return out, nil +} + +// readBoolArray reads a GGUF bool array, retaining every element (bounded +// by maxLayerArray). +func readBoolArray(r *bufio.Reader, vt uint32, fileSize int64) ([]bool, error) { + if vt != typeArray { + return nil, fmt.Errorf("expected array, got type %d", vt) + } + var elemType uint32 + if err := binary.Read(r, binary.LittleEndian, &elemType); err != nil { + return nil, err + } + if elemType != typeBool { + return nil, fmt.Errorf("expected bool array element, got type %d", elemType) + } + var count uint64 + if err := binary.Read(r, binary.LittleEndian, &count); err != nil { + return nil, err + } + if count > maxLayerArray || int64(count) > fileSize { + return nil, fmt.Errorf("implausible bool array count %d", count) + } + out := make([]bool, count) + for i := range out { + var b uint8 + if err := binary.Read(r, binary.LittleEndian, &b); err != nil { + return nil, err + } + out[i] = b != 0 + } + return out, nil +} + // skipValue consumes one value of type vt without retaining it. func skipValue(r *bufio.Reader, vt uint32, fileSize int64) error { if size, ok := scalarSize[vt]; ok { diff --git a/go/internal/gguf/gguf_test.go b/go/internal/gguf/gguf_test.go index 42a307e..a710c8b 100644 --- a/go/internal/gguf/gguf_test.go +++ b/go/internal/gguf/gguf_test.go @@ -77,6 +77,21 @@ func (b *ggufBuilder) kvInt32Array(key string, vals ...int32) { b.kvCount++ } +func (b *ggufBuilder) kvBoolArray(key string, vals ...bool) { + b.str(key) + b.w(uint32(typeArray)) + b.w(uint32(typeBool)) + b.w(uint64(len(vals))) + for _, v := range vals { + bb := uint8(0) + if v { + bb = 1 + } + b.w(bb) + } + b.kvCount++ +} + // build writes the complete file: header, KV section, then `trailer` bytes // standing in for the tensor table region. func (b *ggufBuilder) build(t *testing.T, trailer []byte) string { @@ -276,3 +291,150 @@ func TestRejectsImplausibleLengths(t *testing.T) { t.Fatal("expected error for implausible string length") } } + +// A plain dense/GQA model: scalar head_count_kv, no per-layer SWA split — +// the common case the KV-cache estimator (engine/kvcache.go) must also +// handle, not just Gemma's per-layer arrays. +func TestReadMetadataPlainGQA(t *testing.T) { + b := newBuilder() + b.kvString("general.architecture", "llama") + b.kvUint32("llama.block_count", 32) + b.kvUint32("llama.embedding_length", 4096) + b.kvUint32("llama.attention.head_count", 32) + b.kvUint32("llama.attention.head_count_kv", 8) // scalar + b.kvUint32("llama.attention.key_length", 128) + b.kvUint32("llama.attention.value_length", 128) + path := b.build(t, nil) + + md, err := ReadMetadata(path) + if err != nil { + t.Fatal(err) + } + if md.BlockCount != 32 { + t.Errorf("BlockCount = %d, want 32", md.BlockCount) + } + if len(md.HeadCountKV) != 32 { + t.Fatalf("HeadCountKV len = %d, want 32 (scalar broadcast)", len(md.HeadCountKV)) + } + for i, v := range md.HeadCountKV { + if v != 8 { + t.Fatalf("HeadCountKV[%d] = %d, want 8", i, v) + } + } + if md.KeyLength != 128 || md.ValueLength != 128 { + t.Errorf("KeyLength/ValueLength = %d/%d, want 128/128", md.KeyLength, md.ValueLength) + } + if md.SWAPattern != nil { + t.Errorf("SWAPattern = %v, want nil (no per-layer split declared)", md.SWAPattern) + } + if md.Hybrid { + t.Error("Hybrid = true, want false") + } +} + +// Gemma4-shaped model: per-layer head_count_kv and sliding_window_pattern +// arrays, plus separate SWA key/value lengths — the exact shape ground- +// truthed against the real gemma4-26b-a4b GGUF on ForgeHost during the +// 2026-09-07 incident investigation (30 layers, 5 SWA : 1 global pattern). +func TestReadMetadataGemmaISWA(t *testing.T) { + b := newBuilder() + b.kvString("general.architecture", "gemma4") + b.kvUint32("gemma4.block_count", 6) // small stand-in for the real 30 + b.kvUint32("gemma4.embedding_length", 2816) + b.kvUint32("gemma4.attention.head_count", 16) + b.kvInt32Array("gemma4.attention.head_count_kv", 8, 8, 8, 8, 8, 2) + b.kvUint32("gemma4.attention.key_length", 512) + b.kvUint32("gemma4.attention.value_length", 512) + b.kvUint32("gemma4.attention.key_length_swa", 256) + b.kvUint32("gemma4.attention.value_length_swa", 256) + b.kvUint32("gemma4.attention.sliding_window", 1024) + b.kvBoolArray("gemma4.attention.sliding_window_pattern", true, true, true, true, true, false) + path := b.build(t, nil) + + md, err := ReadMetadata(path) + if err != nil { + t.Fatal(err) + } + if md.BlockCount != 6 { + t.Fatalf("BlockCount = %d, want 6", md.BlockCount) + } + wantKV := []int{8, 8, 8, 8, 8, 2} + if len(md.HeadCountKV) != 6 { + t.Fatalf("HeadCountKV len = %d, want 6", len(md.HeadCountKV)) + } + for i, v := range wantKV { + if md.HeadCountKV[i] != v { + t.Errorf("HeadCountKV[%d] = %d, want %d", i, md.HeadCountKV[i], v) + } + } + wantPattern := []bool{true, true, true, true, true, false} + if len(md.SWAPattern) != 6 { + t.Fatalf("SWAPattern len = %d, want 6", len(md.SWAPattern)) + } + for i, v := range wantPattern { + if md.SWAPattern[i] != v { + t.Errorf("SWAPattern[%d] = %v, want %v", i, md.SWAPattern[i], v) + } + } + if md.KeyLengthSWA != 256 || md.ValueLengthSWA != 256 { + t.Errorf("KeyLengthSWA/ValueLengthSWA = %d/%d, want 256/256", md.KeyLengthSWA, md.ValueLengthSWA) + } + if md.SlidingWindow != 1024 { + t.Errorf("SlidingWindow = %d, want 1024", md.SlidingWindow) + } +} + +// Presence of an ssm.* key must set Hybrid, regardless of its value — +// qwen4exp (Qwen3.8-Flash-Next) declares ssm.state_size etc. alongside +// ordinary attention keys, and the KV-cache formula must abstain rather +// than apply dense-attention math to a recurrent-state architecture. +func TestReadMetadataHybridSSMSignal(t *testing.T) { + b := newBuilder() + b.kvString("general.architecture", "qwen4exp") + b.kvUint32("qwen4exp.block_count", 48) + b.kvUint32("qwen4exp.ssm.state_size", 128) + path := b.build(t, nil) + + md, err := ReadMetadata(path) + if err != nil { + t.Fatal(err) + } + if !md.Hybrid { + t.Error("Hybrid = false, want true (ssm.* key present)") + } +} + +// Presence of an attention.indexer.* key (the block-sparse indexer cache +// qwen4exp also carries) must likewise set Hybrid. +func TestReadMetadataHybridIndexerSignal(t *testing.T) { + b := newBuilder() + b.kvString("general.architecture", "qwen4exp") + b.kvUint32("qwen4exp.attention.indexer.head_count", 4) + path := b.build(t, nil) + + md, err := ReadMetadata(path) + if err != nil { + t.Fatal(err) + } + if !md.Hybrid { + t.Error("Hybrid = false, want true (attention.indexer.* key present)") + } +} + +// A head_count_kv array whose length doesn't match block_count is corrupt +// or unparseable in a way we can't trust — must be dropped, not guessed at. +func TestReadMetadataHeadCountKVLengthMismatchDropped(t *testing.T) { + b := newBuilder() + b.kvString("general.architecture", "weird") + b.kvUint32("weird.block_count", 4) + b.kvInt32Array("weird.attention.head_count_kv", 8, 8) // len 2, want 4 or 1 + path := b.build(t, nil) + + md, err := ReadMetadata(path) + if err != nil { + t.Fatal(err) + } + if md.HeadCountKV != nil { + t.Errorf("HeadCountKV = %v, want nil (length mismatch must not be guessed)", md.HeadCountKV) + } +} diff --git a/go/internal/gguf/live_probe_test.go b/go/internal/gguf/live_probe_test.go index 910ede4..fed0b30 100644 --- a/go/internal/gguf/live_probe_test.go +++ b/go/internal/gguf/live_probe_test.go @@ -26,5 +26,9 @@ func TestLiveHeaders(t *testing.T) { } t.Logf("%s: %s arch=%q name=%q ctx=%d quant=%q params=%d", filepath.Base(p), time.Since(t0), md.Architecture, md.Name, md.TrainedCtx, md.QuantType, md.ParameterCount) + t.Logf(" layers: block_count=%d embd=%d head_count=%d head_count_kv=%v", + md.BlockCount, md.EmbeddingLength, md.HeadCount, md.HeadCountKV) + t.Logf(" attn: key_len=%d val_len=%d key_len_swa=%d val_len_swa=%d sliding_window=%d swa_pattern=%v hybrid=%v", + md.KeyLength, md.ValueLength, md.KeyLengthSWA, md.ValueLengthSWA, md.SlidingWindow, md.SWAPattern, md.Hybrid) } } diff --git a/go/internal/httpapi/httpapi.go b/go/internal/httpapi/httpapi.go index d9a9126..a207226 100644 --- a/go/internal/httpapi/httpapi.go +++ b/go/internal/httpapi/httpapi.go @@ -663,6 +663,11 @@ func (s *Server) registerV1Routes(mux *http.ServeMux) { mux.Handle("PUT /api/v1/metrics/settings", s.requireRole(authz.RoleAdmin)(s.requireAssurance(authz.ResourcePageSettings)(http.HandlerFunc(s.handleMetricsSettingsPut)))) mux.Handle("GET /api/v1/ui/settings", s.requireRole(authz.RoleOperator)(s.requireAssurance(authz.ResourcePageSettings)(http.HandlerFunc(s.handleUISettingsGet)))) mux.Handle("PUT /api/v1/ui/settings", s.requireRole(authz.RoleAdmin)(s.requireAssurance(authz.ResourcePageSettings)(http.HandlerFunc(s.handleUISettingsPut)))) + // infra.service_icons (live) — vendor-icon override for the fixed infra + // services, added 2026-09-06 to replace a Go-literal map that needed a + // rebuild+restart to change. + mux.Handle("GET /api/v1/service-icons", s.requireRole(authz.RoleOperator)(s.requireAssurance(authz.ResourcePageSettings)(http.HandlerFunc(s.handleServiceIconsGet)))) + mux.Handle("PUT /api/v1/service-icons", s.requireRole(authz.RoleAdmin)(s.requireAssurance(authz.ResourcePageSettings)(http.HandlerFunc(s.handleServiceIconsPut)))) // ADR-0011: custom dashboard pages — system-wide settings key // "dashboard.pages", full-replace PUT (frontend sends complete layout). mux.Handle("GET /api/v1/dashboard/layout", s.requireRole(authz.RoleOperator)(s.requireAssurance(authz.ResourcePageSettings)(http.HandlerFunc(s.handleDashboardLayoutGet)))) diff --git a/go/internal/httpapi/infra_handlers.go b/go/internal/httpapi/infra_handlers.go index c86a121..ab0f3fc 100644 --- a/go/internal/httpapi/infra_handlers.go +++ b/go/internal/httpapi/infra_handlers.go @@ -38,7 +38,10 @@ import ( "regexp" "slices" "strconv" + "strings" "time" + + "github.com/jsaigou/the-forge/internal/config" ) // preflightCheck is one row of a Danger Zone check-before-save result. @@ -327,6 +330,94 @@ func (s *Server) handleMonitorSettingsPut(w http.ResponseWriter, r *http.Request writeJSON(w, http.StatusOK, resp) } +// ── GET/PUT /api/v1/service-icons — infra.service_icons (live) ───────────── +// +// Vendor-icon override for the fixed infra services (STT/Embedding/Aligner/ +// TTS — bare [ports] entries with no catalog-backed model metadata, so they +// have no other place to record which model they're currently showing an +// icon for). Previously a Go literal map in services_handlers.go +// (serviceVendorLogoBySlug) — changing which icon a service showed meant a +// source edit, rebuild, and daemon restart for a purely cosmetic value. +// Operator feedback 2026-09-06: "bad design to need a restart to change an +// icon." Genuinely live, unlike most of the system group above: handled the +// same way infra.ports is stored (a bare map, whole-map replace when the +// body includes it), and handleInfraServices reads cfg.ServiceIcons fresh +// off s.deps.Config() on every request — no downstream component captures +// this value at startup the way a listener address or DB path would, so +// ReloadConfig alone (no restart) is sufficient. + +type serviceIconsResponse struct { + Icons map[string]string `json:"icons"` +} + +func (s *Server) resolvedServiceIcons(ctx context.Context) serviceIconsResponse { + var icons map[string]string + if err := json.Unmarshal(s.getRawSetting(ctx, "infra.service_icons"), &icons); err != nil { + log.Printf("httpapi: warning: corrupt stored setting: %v", err) + } + if icons == nil { + icons = map[string]string{} + } + for name, slug := range config.DefaultServiceIcons() { + if _, ok := icons[name]; !ok { + icons[name] = slug + } + } + return serviceIconsResponse{Icons: icons} +} + +// handleServiceIconsGet — GET /api/v1/service-icons (operator). +func (s *Server) handleServiceIconsGet(w http.ResponseWriter, r *http.Request) { + writeJSON(w, http.StatusOK, s.resolvedServiceIcons(r.Context())) +} + +type serviceIconsBody struct { + Icons map[string]string `json:"icons"` +} + +// handleServiceIconsPut — PUT /api/v1/service-icons (admin, page.settings). +// Whole-map replace when the body includes "icons" (same semantics as +// infra.ports in the system group) — send the full current map back (as +// returned by GET) with just the one entry changed, not a single-key patch. +func (s *Server) handleServiceIconsPut(w http.ResponseWriter, r *http.Request) { + if s.deps.Settings == nil { + writeError(w, http.StatusServiceUnavailable, "settings store not wired") + return + } + var body serviceIconsBody + if fields := decodeJSONBody(r, &body); fields != nil { + writeValidationError(w, fields) + return + } + if body.Icons == nil { + writeValidationError(w, map[string]string{"icons": "must be present"}) + return + } + for name, slug := range body.Icons { + if strings.TrimSpace(name) == "" || strings.TrimSpace(slug) == "" { + writeValidationError(w, map[string]string{"icons": "keys and values must be non-empty"}) + return + } + } + + ctx, cancel := context.WithTimeout(r.Context(), 5*time.Second) + defer cancel() + raw, err := json.Marshal(body.Icons) + if err != nil { + writeInternalError(w, err) + return + } + if err := s.deps.Settings.Set(ctx, "infra.service_icons", raw); err != nil { + writeInternalError(w, err) + return + } + if s.deps.ReloadConfig != nil { + s.deps.ReloadConfig() + } + s.audit(r, identity(r).Name, "service_icons", "infra.service_icons", string(raw)) + writeJSON(w, http.StatusOK, s.resolvedServiceIcons(ctx)) +} + // putIntField is the small manual "if the caller sent this field, marshal // it into the patch map under its JSON key, after checking bounds" helper // every group below repeats — deliberately not reflection-based (see this @@ -1105,8 +1196,8 @@ type dashboardWidgetEntry struct { } type dashboardPageEntry struct { - ID string `json:"id"` - Name string `json:"name"` + ID string `json:"id"` + Name string `json:"name"` Widgets []dashboardWidgetEntry `json:"widgets"` } diff --git a/go/internal/httpapi/infra_handlers_test.go b/go/internal/httpapi/infra_handlers_test.go index 376cb36..0621c7d 100644 --- a/go/internal/httpapi/infra_handlers_test.go +++ b/go/internal/httpapi/infra_handlers_test.go @@ -179,6 +179,74 @@ func TestMonitorSettingsGetPutIsLiveNotRestart(t *testing.T) { } } +// TestServiceIconsGetDefaultsAndPutOverride is the HTTP-layer counterpart to +// config.TestServiceIconsDefaultsAndOverride — regression coverage for +// "changing a fixed infra service's icon needs a rebuild+restart" (operator +// feedback 2026-09-06). Confirms GET fills in config.DefaultServiceIcons() +// for anything unset, PUT persists a whole-map override live (no restart +// marker, matching monitor settings' contract above), and a round-trip +// GET-then-PUT-the-full-map never silently drops a sibling key. +func TestServiceIconsGetDefaultsAndPutOverride(t *testing.T) { + set := newFakeSettings() + s := serverWithSettings(t, set) + + w := do(t, s, authedRequest("GET", "/api/v1/service-icons", nil)) + if w.Code != 200 { + t.Fatalf("GET service-icons = %d, body=%s", w.Code, w.Body) + } + var resp serviceIconsResponse + decodeJSON(t, w.Body, &resp) + for name, want := range config.DefaultServiceIcons() { + if got := resp.Icons[name]; got != want { + t.Errorf("GET (no override) Icons[%q] = %q, want default %q", name, got, want) + } + } + + // Round-trip: send the full map back with just STT changed, as a real + // caller (following GET's own documented contract) would. + resp.Icons["STT"] = "nvidia" + body, err := json.Marshal(resp) + if err != nil { + t.Fatalf("marshal PUT body: %v", err) + } + w = do(t, s, authedRequest("PUT", "/api/v1/service-icons", strings.NewReader(string(body)))) + if w.Code != 200 { + t.Fatalf("PUT service-icons = %d, body=%s", w.Code, w.Body) + } + var putResp serviceIconsResponse + decodeJSON(t, w.Body, &putResp) + if putResp.Icons["STT"] != "nvidia" { + t.Errorf("Icons[STT] after PUT = %q, want nvidia", putResp.Icons["STT"]) + } + if putResp.Icons["Embedding"] != "qwen" { + t.Errorf("Icons[Embedding] after PUT = %q, want qwen (sibling key dropped by round-trip)", putResp.Icons["Embedding"]) + } + // Purely cosmetic value — must never mark a restart pending. + if info := s.restartRequired(context.Background()); info != nil { + t.Errorf("service-icons should not mark restart_required, got %+v", info) + } + + // GET again confirms the write actually persisted to the store, not just + // echoed back in the PUT response. + w = do(t, s, authedRequest("GET", "/api/v1/service-icons", nil)) + var resp2 serviceIconsResponse + decodeJSON(t, w.Body, &resp2) + if resp2.Icons["STT"] != "nvidia" { + t.Errorf("GET after PUT Icons[STT] = %q, want nvidia (PUT did not persist)", resp2.Icons["STT"]) + } + + // Empty/missing icons object rejected rather than silently wiping the + // stored map. + w = do(t, s, authedRequest("PUT", "/api/v1/service-icons", strings.NewReader(`{}`))) + if w.Code != 422 { + t.Fatalf("PUT with no icons field = %d, want 422", w.Code) + } + w = do(t, s, authedRequest("PUT", "/api/v1/service-icons", strings.NewReader(`{"icons":{"STT":""}}`))) + if w.Code != 422 { + t.Fatalf("PUT with empty-string value = %d, want 422", w.Code) + } +} + func TestMetricsSettingsSplitApplyMode(t *testing.T) { set := newFakeSettings() s := serverWithSettings(t, set) diff --git a/go/internal/httpapi/services_handlers.go b/go/internal/httpapi/services_handlers.go index 6389a49..1bd3408 100644 --- a/go/internal/httpapi/services_handlers.go +++ b/go/internal/httpapi/services_handlers.go @@ -19,6 +19,7 @@ import ( "time" "github.com/jsaigou/the-forge/internal/collector" + "github.com/jsaigou/the-forge/internal/config" "github.com/jsaigou/the-forge/internal/smith" ) @@ -87,7 +88,7 @@ func (s *Server) handleInfraServices(w http.ResponseWriter, r *http.Request) { Name: "STT", Unit: ptrString("forge-stt"), Port: &p, Active: unitActive(snap, "forge-stt"), Kind: "systemd", ModeKey: nil, - Logo: serviceVendorLogo("STT"), + Logo: serviceVendorLogo(cfg, "STT"), }) } if p, ok := cfg.Ports["embedding"]; ok { @@ -95,7 +96,7 @@ func (s *Server) handleInfraServices(w http.ResponseWriter, r *http.Request) { Name: "Embedding", Unit: ptrString("forge-embedding"), Port: &p, Active: unitActive(snap, "forge-embedding"), Kind: "systemd", ModeKey: nil, - Logo: serviceVendorLogo("Embedding"), + Logo: serviceVendorLogo(cfg, "Embedding"), }) } // Aligner (Qwen3 aligner, pre-existing on ForgeHost, unrelated to slots — @@ -107,7 +108,7 @@ func (s *Server) handleInfraServices(w http.ResponseWriter, r *http.Request) { Name: "Aligner", Unit: ptrString("forge-aligner"), Port: &p, Active: unitActive(snap, "forge-aligner"), Kind: "systemd", ModeKey: nil, - Logo: serviceVendorLogo("Aligner"), + Logo: serviceVendorLogo(cfg, "Aligner"), }) } } @@ -122,7 +123,7 @@ func (s *Server) handleInfraServices(w http.ResponseWriter, r *http.Request) { Name: "TTS", Unit: ptrString(ttsUnit), Port: &ttsPort, Active: unitActive(snap, ttsUnit), Kind: "systemd", ModeKey: nil, - Logo: serviceVendorLogo("TTS"), + Logo: serviceVendorLogo(cfg, "TTS"), }) // Service modes from config (e.g. ComfyUI). Icon comes from the mode's @@ -180,31 +181,22 @@ func (s *Server) handleInfraServices(w http.ResponseWriter, r *http.Request) { writeJSON(w, http.StatusOK, infraServicesResponse{Services: services}) } -// serviceVendorLogoBySlug maps a fixed infra service's display Name to an -// Icon manifest slug (web/src/assets/icons/manifest.ts) for the model it -// actually runs. Ground-truthed live against ForgeHost's real systemd unit files -// 2026-07-31 (Console polish pass), not guessed from doc comments: -// - STT (forge-stt): parakeet-server --model .../nemotron-3.5-asr-streaming-0.6b-f16.gguf -// - Embedding (forge-embedding): llama-server -m .../Qwen3-Embedding-0.6B-Q8_0.gguf -// - Aligner (forge-aligner): aligner_server.py MODEL_NAME = "Qwen/Qwen3-ForcedAligner-0.6B" -// - TTS (forge-tts): tts_server.py MODEL_IDS = "Qwen/Qwen3-TTS-12Hz-1.7B-*" -// -// These are bare [ports] entries with no model metadata anywhere in the -// store (unlike catalog-backed service modes, e.g. ComfyUI below), so a -// literal map is the honest fix rather than a workaround for a missing -// dynamic mechanism. -// Operator feedback 2026-08-14: the icon names the MODEL, not the company — -// Embedding/Aligner/TTS all run Qwen models, so they get the qwen mark -// (STT runs Nvidia Parakeet, hence nvidia). -var serviceVendorLogoBySlug = map[string]string{ - "STT": "nvidia", - "Embedding": "qwen", - "Aligner": "qwen", - "TTS": "qwen", -} - -func serviceVendorLogo(name string) *string { - if slug, ok := serviceVendorLogoBySlug[name]; ok { +// serviceVendorLogo looks up the Icon manifest slug (web/src/assets/icons/ +// manifest.ts) for a fixed infra service's display Name — STT/Embedding/ +// Aligner/TTS are bare [ports] entries with no catalog-backed model +// metadata (unlike service_mode rows, e.g. ComfyUI below, which carry +// their own services.icon), so this used to be a Go literal map here. +// Moved to the store-backed infra.service_icons setting (config.Config. +// ServiceIcons, config.go's defaultServiceIcons for the fallback values) — +// operator feedback 2026-09-06: a Go-literal map meant changing which icon +// a service showed needed a source edit, rebuild, and daemon restart for a +// purely cosmetic value. It's SIGHUP-reloadable now, same as every other +// infra.* setting (see GET/PUT /api/v1/service-icons in infra_handlers.go). +func serviceVendorLogo(cfg *config.Config, name string) *string { + if cfg == nil { + return nil + } + if slug, ok := cfg.ServiceIcons[name]; ok && slug != "" { return ptrString(slug) } return nil diff --git a/installer/assets.manifest.json b/installer/assets.manifest.json index 7fe532d..56a6580 100644 --- a/installer/assets.manifest.json +++ b/installer/assets.manifest.json @@ -30,15 +30,15 @@ "notes": "CPU-only always-on embedding server (forge-embedding.service, port 8083)." }, { - "id": "parakeet-stt-0.6b", + "id": "qwen3-asr-stt-0.6b", "role": "mandatory", "hf_repo": null, "gguf": true, - "dest_rel_path": "models/parakeet", + "dest_rel_path": "models/qwen3-asr", "files": [ - { "filename": "nemotron-3.5-asr-streaming-0.6b-f16.gguf", "sha256": null, "size_bytes": null } + { "filename": "qwen3-asr-0.6b-q8_0.gguf", "sha256": "d81d497b25508449b4bb511b343aee663405d28132d8dccfd5f29a1cceef77d3", "size_bytes": 1354082720 } ], - "notes": "STT server model (forge-stt.service). Upstream repo TBD." + "notes": "STT server model (forge-stt.service, swapped from Nemotron ASR 2026-09-06 — see progress.md). Not a direct HF download: clone https://github.com/predict-woo/qwen3-asr.cpp (verified at commit 6dcc586), download Qwen/Qwen3-ASR-0.6B from HF, then run that repo's scripts/convert_hf_to_gguf.py --type q8_0 to produce this file locally. sha256/size_bytes above are the exact bytes verified working in production on 2026-09-06 — a re-conversion at a later qwen3-asr.cpp/Qwen3-ASR-0.6B revision may not reproduce byte-for-byte; treat as a known-good reference rather than a strict pin. Two local fixes were needed against the conversion script as of that commit: it hardcodes a dev-machine gguf-py path check that throws PermissionError instead of failing closed on a non-root user (patch GGUF_PY_PATH to a nonexistent path so it falls through to the pip 'gguf' package), and 'gguf' itself is missing from scripts/requirements.txt (pip/uv install it separately)." }, { "id": "qwen-tts", diff --git a/systemd/forge-stt.service b/systemd/forge-stt.service index 8d8fc3f..4659e68 100644 --- a/systemd/forge-stt.service +++ b/systemd/forge-stt.service @@ -1,16 +1,18 @@ [Unit] -Description=The Forge — STT Server (Nemotron ASR, port 8084) +Description=The Forge — STT Server (Qwen3-ASR-0.6B, port 8084) After=network.target [Service] Type=simple User=testuser Group=testuser -ExecStart=/opt/forge/parakeet-src/build/examples/server/parakeet-server \ - --model /opt/forge/models/parakeet/nemotron-3.5-asr-streaming-0.6b-f16.gguf \ +ExecStart=/opt/forge/qwen3-asr-src/build/qwen3-asr-server \ + -m /opt/forge/models/qwen3-asr/qwen3-asr-0.6b-q8_0.gguf \ --host 127.0.0.1 \ --port 8084 \ - --threads 4 + --max-upload-mb 512 \ + --convert \ + -t 4 Restart=on-failure RestartSec=10 StandardOutput=journal