From a70166bebfae949fb5b5242b0c19a9165ad65cd5 Mon Sep 17 00:00:00 2001 From: Glenn Neuber Date: Sun, 11 Oct 2026 12:15:07 +1000 Subject: [PATCH 1/2] fix(mlx): price muse-glimmer's KV cache; a config.json the estimator cannot read says so MLX admission priced every muse-glimmer load as weights only (kv="0 B", headroom="0 B") and logged "no KV rule for this architecture" with an empty architecture, though kvsize has a glimmer rule. Glimmer's config.json gives sliding_window_pattern as a per-position list ([2048,2048,2048,0]); the estimator's field is gemma4's integer period, and that one field failed the whole parse (muse-glimmer:30b-nvfp4 and 30b-mxfp8, 2026-10-11). - kvsize: sliding_window_pattern reads an integer as before and a list as unset. Glimmer's model reads layer_types, never this field; the list form now prices exactly as if the field were absent. - kvsize: Estimate.Err carries a parse failure, and the admission warning names it instead of reporting an unregistered architecture with no name, so the next config shape surprise is visible. - Fixtures: glimmer's published config.json (without its quantization_config) and its DFlash draft/config.json. Co-Authored-By: Claude Opus 5.5 --- mlxrunner/client.go | 17 +- mlxrunner/kvsize/arch.go | 4 +- mlxrunner/kvsize/kvsize.go | 37 +++- mlxrunner/kvsize/kvsize_test.go | 72 ++++++++ .../testdata/muse-glimmer-30b-config.json | 169 ++++++++++++++++++ .../muse-glimmer-30b-dflash-draft-config.json | 37 ++++ 6 files changed, 323 insertions(+), 13 deletions(-) create mode 100644 mlxrunner/kvsize/testdata/muse-glimmer-30b-config.json create mode 100644 mlxrunner/kvsize/testdata/muse-glimmer-30b-dflash-draft-config.json diff --git a/mlxrunner/client.go b/mlxrunner/client.go index 1e3a881eae7..7319dd3d68b 100644 --- a/mlxrunner/client.go +++ b/mlxrunner/client.go @@ -581,11 +581,18 @@ func (c *Client) admit(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, requireFu numCtx := int(c.softContextLength.Load()) need, kv, headroom, est := c.needFor(weights, numCtx) if !est.Known { - // No cache rule for this architecture. Fall back to exactly what - // admission did before the estimator existed rather than guessing: a - // wrong estimate refuses loads that would have served. - slog.Warn("MLX admission is pricing weights only: no KV rule for this architecture", - "architecture", est.Arch, "model", c.modelName, "num_ctx", numCtx) + // No cache rule for this architecture, or a config.json the estimator + // cannot read. Fall back to exactly what admission did before the + // estimator existed rather than guessing: a wrong estimate refuses + // loads that would have served. Say which, since an unreadable config + // has no architecture to name. + if est.Err != nil { + slog.Warn("MLX admission is pricing weights only: config.json does not parse", + "model", c.modelName, "num_ctx", numCtx, "error", est.Err) + } else { + slog.Warn("MLX admission is pricing weights only: no KV rule for this architecture", + "architecture", est.Arch, "model", c.modelName, "num_ctx", numCtx) + } } // An automatic rung was never asked for: it is Ollama's VRAM-tier default diff --git a/mlxrunner/kvsize/arch.go b/mlxrunner/kvsize/arch.go index 8b8382acc0e..5c461c87709 100644 --- a/mlxrunner/kvsize/arch.go +++ b/mlxrunner/kvsize/arch.go @@ -151,7 +151,7 @@ func gemma4IsSliding(cfg *config, i int) bool { if len(cfg.text.LayerTypes) > 0 && i < len(cfg.text.LayerTypes) { return cfg.text.LayerTypes[i] == "sliding_attention" } - pattern := cfg.text.SlidingWindowPattern + pattern := int(cfg.text.SlidingWindowPattern) if pattern <= 0 && len(cfg.text.LayerTypes) == 0 { pattern = 5 } @@ -356,7 +356,7 @@ func cohere2MoeRule(cfg *config, numCtx int) (Estimate, bool) { if window <= 0 { window = 4096 } - slidingPattern := cfg.text.SlidingWindowPattern + slidingPattern := int(cfg.text.SlidingWindowPattern) if slidingPattern <= 0 { slidingPattern = 4 } diff --git a/mlxrunner/kvsize/kvsize.go b/mlxrunner/kvsize/kvsize.go index 464cd897bd5..f5d7cafd6ae 100644 --- a/mlxrunner/kvsize/kvsize.go +++ b/mlxrunner/kvsize/kvsize.go @@ -28,6 +28,7 @@ package kvsize import ( "encoding/json" + "fmt" "strings" ) @@ -60,6 +61,9 @@ type Estimate struct { // Known is false when no rule matched the target architecture. Every byte // field is then zero and the caller must not refuse a load on it. Known bool + // Err is why config.json could not be read, when that is why Known is + // false; Arch is then empty, since nothing was read to dispatch on. + Err error // Arch is the target architecture the estimate dispatched on. Arch string // DraftArch is the draft architecture, empty when the model ships no @@ -101,7 +105,7 @@ func Model(config, draft []byte, numCtx int) Estimate { cfg, err := parse(config) if err != nil { - return Estimate{NumCtx: numCtx} + return Estimate{NumCtx: numCtx, Err: err} } rule, arch := lookup(targetRules, cfg) @@ -249,11 +253,11 @@ type textConfig struct { SlidingWindow int `json:"sliding_window"` // gemma4 - GlobalHeadDim int `json:"global_head_dim"` - NumGlobalKeyValueHeads int `json:"num_global_key_value_heads"` - AttentionKEqV bool `json:"attention_k_eq_v"` - NumKVSharedLayers int `json:"num_kv_shared_layers"` - SlidingWindowPattern int `json:"sliding_window_pattern"` + GlobalHeadDim int `json:"global_head_dim"` + NumGlobalKeyValueHeads int `json:"num_global_key_value_heads"` + AttentionKEqV bool `json:"attention_k_eq_v"` + NumKVSharedLayers int `json:"num_kv_shared_layers"` + SlidingWindowPattern slidingWindowPattern `json:"sliding_window_pattern"` // cohere2_moe PrefixDenseSlidingWindowPattern int `json:"prefix_dense_sliding_window_pattern"` @@ -287,6 +291,27 @@ type textConfig struct { QKRopeHeadDim int `json:"qk_rope_head_dim"` } +// slidingWindowPattern is gemma4's sliding_window_pattern, an integer period +// (every Nth layer is global). muse-glimmer writes a per-position list of +// windows instead ([2048,2048,2048,0]) and its model reads layer_types, never +// this field; the list form reads as 0 (unset) rather than failing the whole +// config, which priced glimmer as weights only. +type slidingWindowPattern int + +func (p *slidingWindowPattern) UnmarshalJSON(data []byte) error { + var n int + if err := json.Unmarshal(data, &n); err == nil { + *p = slidingWindowPattern(n) + return nil + } + var list []json.RawMessage + if err := json.Unmarshal(data, &list); err != nil { + return fmt.Errorf("sliding_window_pattern: want an integer or a list, got %s", data) + } + *p = 0 + return nil +} + func parse(data []byte) (*config, error) { var envelope struct { Architectures []string `json:"architectures"` diff --git a/mlxrunner/kvsize/kvsize_test.go b/mlxrunner/kvsize/kvsize_test.go index 04e74a2e5ae..008872ff133 100644 --- a/mlxrunner/kvsize/kvsize_test.go +++ b/mlxrunner/kvsize/kvsize_test.go @@ -1,6 +1,7 @@ package kvsize import ( + "encoding/json" "os" "path/filepath" "testing" @@ -362,3 +363,74 @@ func TestSlotsRoundUpToTheCacheStep(t *testing.T) { } } } + +const ( + glimmerConfig = "muse-glimmer-30b-config.json" + glimmerDraftConfig = "muse-glimmer-30b-dflash-draft-config.json" +) + +// muse-glimmer's config.json gives sliding_window_pattern as a per-position +// list ([2048,2048,2048,0]), not gemma4's integer period. That one field +// failed the whole parse, so MLX admission priced every glimmer load as +// weights only and logged "no KV rule for this architecture" with an empty +// architecture (muse-glimmer:30b-nvfp4, 2026-10-11). The fixture is the +// published config.json without its quantization_config; the draft is the +// -dflash tags' draft/config.json. +func TestGlimmerConfigIsPriced(t *testing.T) { + got := Model(load(t, glimmerConfig), nil, 8192) + if got.Err != nil { + t.Fatalf("parse error: %v", got.Err) + } + if !got.Known || got.Arch != "MuseGlimmerForConditionalGeneration" { + t.Fatalf("known=%v arch=%q, want the glimmer rule", got.Known, got.Arch) + } + // layer_types holds 39 "sliding_attention" and 13 "full_attention" (every + // 4th layer), which is what glimmer.go reads; sliding_window_pattern is + // not consulted. + if want := (LayerCounts{Attention: 13, Sliding: 39}); got.Layers != want { + t.Errorf("layers = %+v, want %+v", got.Layers, want) + } + if got.Total() == 0 { + t.Error("a priced glimmer estimate has no bytes") + } + + withDraft := Model(load(t, glimmerConfig), load(t, glimmerDraftConfig), 8192) + if !withDraft.DraftKnown || withDraft.DraftArch != "MuseGlimmerAssistantModel" { + t.Errorf("draft known=%v arch=%q, want the dflash rule", withDraft.DraftKnown, withDraft.DraftArch) + } + if withDraft.Total() <= got.Total() { + t.Errorf("total with the drafter %d, want more than without it %d", withDraft.Total(), got.Total()) + } +} + +// The list form must price exactly as if the field were absent: layer_types +// decides, as in the model. +func TestSlidingWindowPatternListIsNotAPeriod(t *testing.T) { + var cfg map[string]any + if err := json.Unmarshal(load(t, glimmerConfig), &cfg); err != nil { + t.Fatal(err) + } + delete(cfg, "sliding_window_pattern") + without, err := json.Marshal(cfg) + if err != nil { + t.Fatal(err) + } + for _, n := range ladder { + a, b := Model(load(t, glimmerConfig), nil, n), Model(without, nil, n) + if a.Total() != b.Total() || a.Layers != b.Layers { + t.Errorf("num_ctx %d: list form %d bytes %+v, field absent %d bytes %+v", n, a.Total(), a.Layers, b.Total(), b.Layers) + } + } +} + +// A config the parser cannot read says so: the caller's warning must not +// read as an unregistered architecture with an empty name. +func TestUnparsableConfigCarriesItsError(t *testing.T) { + got := Model([]byte(`{"architectures":["MuseGlimmerForConditionalGeneration"],"num_hidden_layers":"52"}`), nil, 8192) + if got.Known { + t.Fatal("an unparsable config must not be priced") + } + if got.Err == nil { + t.Fatal("Err is nil for a config that does not parse") + } +} diff --git a/mlxrunner/kvsize/testdata/muse-glimmer-30b-config.json b/mlxrunner/kvsize/testdata/muse-glimmer-30b-config.json new file mode 100644 index 00000000000..7eccc464305 --- /dev/null +++ b/mlxrunner/kvsize/testdata/muse-glimmer-30b-config.json @@ -0,0 +1,169 @@ +{ + "architectures": [ + "MuseGlimmerForConditionalGeneration" + ], + "auto_map": { + "AutoConfig": "configuration_muse_glimmer.MuseGlimmerConfig", + "AutoModelForCausalLM": "modeling_muse_glimmer.MuseGlimmerForCausalLM" + }, + "bos_token_id": 200000, + "dtype": "bfloat16", + "eos_token_id": 200001, + "every_n_layers_nope": 4, + "has_vision": true, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 6656, + "intermediate_size": 19968, + "layer_types": [ + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention" + ], + "max_position_embeddings": 131072, + "model_type": "muse_glimmer", + "no_rope_layers": [ + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 0 + ], + "normalize_tok_embeddings": true, + "num_attention_heads": 32, + "num_hidden_layers": 52, + "num_key_value_heads": 2, + "output_multiplier": 0.19611613513818404, + "output_soft_cap_temp": 20.0, + "patch_token_id": 200092, + "post_norm_eps": 1E-8, + "qk_scale_factor": 43.7840518911, + "rms_norm_eps": 0.00001, + "rope_theta": 500000.0, + "sliding_window": 2048, + "sliding_window_pattern": [ + 2048, + 2048, + 2048, + 0 + ], + "tie_word_embeddings": false, + "transformers_version": "5.12.1", + "use_attn_output_gate": true, + "use_qk_norm": true, + "vid_end_id": 200083, + "vid_frame_sep_id": 200087, + "vid_start_id": 200082, + "video_num_frames": 96, + "video_sampling_fps": 2.0, + "video_token_id": 200091, + "vision_adapter_dim": 4096, + "vision_downsample_factor": 2, + "vision_heads": 16, + "vision_latent_dim": 1536, + "vision_layers": 50, + "vision_mlp_ratio": 5.833333333333333, + "vision_output_dim": 6144, + "vision_patch_size": 14, + "vision_patch_temporal": 2, + "vision_pos_emb_grid_h": 32, + "vision_pos_emb_grid_w": 32, + "vision_sparse_attention_factor": 4, + "vocab_size": 202048 +} diff --git a/mlxrunner/kvsize/testdata/muse-glimmer-30b-dflash-draft-config.json b/mlxrunner/kvsize/testdata/muse-glimmer-30b-dflash-draft-config.json new file mode 100644 index 00000000000..40d6bf6667a --- /dev/null +++ b/mlxrunner/kvsize/testdata/muse-glimmer-30b-dflash-draft-config.json @@ -0,0 +1,37 @@ +{ + "architectures": [ + "MuseGlimmerAssistantModel" + ], + "model_type": "muse_glimmer_assistant", + "block_size": 16, + "hidden_size": 6656, + "num_hidden_layers": 5, + "num_attention_heads": 32, + "num_key_value_heads": 8, + "head_dim": 128, + "rms_norm_eps": 1e-05, + "rope_theta": 500000.0, + "vocab_size": 202048, + "layer_types": [ + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention" + ], + "sliding_window": 2048, + "rope_interleaved": true, + "dflash_config": { + "block_size": 16, + "mask_token_id": 201818, + "target_layer_ids": [ + 1, + 13, + 25, + 37, + 49 + ], + "num_target_layers": 52, + "causal": false + } +} \ No newline at end of file From f6bdb8aefdc98f7b9cf2741c91d7b0c5dd8aae0b Mon Sep 17 00:00:00 2001 From: Glenn Neuber Date: Tue, 13 Oct 2026 12:29:58 +1000 Subject: [PATCH 2/2] fix(mlx): a list sliding_window_pattern without layer_types is not priced on a period guess gemma4.go and cohere2_moe.go read sliding_window_pattern as an int32, so a list fails the model's own config parse. Without layer_types the gemma4 rule fell back to a period of 5 and cohere2_moe to 4, pricing a load that cannot happen; both now decline and the caller falls back. glimmer is unchanged: its layer_types decide. Raised in the CUDA host's review of #466. Co-Authored-By: Claude Opus 5.5 --- mlxrunner/kvsize/arch.go | 15 +++++++++++++-- mlxrunner/kvsize/kvsize.go | 14 +++++++++----- mlxrunner/kvsize/kvsize_test.go | 33 +++++++++++++++++++++++++++++++++ 3 files changed, 55 insertions(+), 7 deletions(-) diff --git a/mlxrunner/kvsize/arch.go b/mlxrunner/kvsize/arch.go index 5c461c87709..3df27999d28 100644 --- a/mlxrunner/kvsize/arch.go +++ b/mlxrunner/kvsize/arch.go @@ -105,6 +105,12 @@ func gemma4Rule(cfg *config, numCtx int) (Estimate, bool) { if layers <= 0 { return Estimate{}, false } + // gemma4.go reads sliding_window_pattern as an int32, so a list fails the + // model's own config parse; without layer_types there is nothing that says + // which layers slide, and the default period of 5 would be a guess. + if cfg.text.SlidingWindowPattern.IsList && len(cfg.text.LayerTypes) == 0 { + return Estimate{}, false + } // Defaults from gemma4.go:420-441. headDim := cfg.text.HeadDim @@ -151,7 +157,7 @@ func gemma4IsSliding(cfg *config, i int) bool { if len(cfg.text.LayerTypes) > 0 && i < len(cfg.text.LayerTypes) { return cfg.text.LayerTypes[i] == "sliding_attention" } - pattern := int(cfg.text.SlidingWindowPattern) + pattern := cfg.text.SlidingWindowPattern.Period if pattern <= 0 && len(cfg.text.LayerTypes) == 0 { pattern = 5 } @@ -356,7 +362,12 @@ func cohere2MoeRule(cfg *config, numCtx int) (Estimate, bool) { if window <= 0 { window = 4096 } - slidingPattern := int(cfg.text.SlidingWindowPattern) + // cohere2_moe.go reads sliding_window_pattern as an int32 too: a list is a + // config the model cannot load, not one to price with the default of 4. + if cfg.text.SlidingWindowPattern.IsList { + return Estimate{}, false + } + slidingPattern := cfg.text.SlidingWindowPattern.Period if slidingPattern <= 0 { slidingPattern = 4 } diff --git a/mlxrunner/kvsize/kvsize.go b/mlxrunner/kvsize/kvsize.go index f5d7cafd6ae..7254c8c5c07 100644 --- a/mlxrunner/kvsize/kvsize.go +++ b/mlxrunner/kvsize/kvsize.go @@ -294,21 +294,25 @@ type textConfig struct { // slidingWindowPattern is gemma4's sliding_window_pattern, an integer period // (every Nth layer is global). muse-glimmer writes a per-position list of // windows instead ([2048,2048,2048,0]) and its model reads layer_types, never -// this field; the list form reads as 0 (unset) rather than failing the whole -// config, which priced glimmer as weights only. -type slidingWindowPattern int +// this field; the list form reads as Period 0 (unset) rather than failing the +// whole config, which priced glimmer as weights only. IsList records it, so a +// rule that would otherwise fall back to a default period can decline instead. +type slidingWindowPattern struct { + Period int + IsList bool +} func (p *slidingWindowPattern) UnmarshalJSON(data []byte) error { var n int if err := json.Unmarshal(data, &n); err == nil { - *p = slidingWindowPattern(n) + *p = slidingWindowPattern{Period: n} return nil } var list []json.RawMessage if err := json.Unmarshal(data, &list); err != nil { return fmt.Errorf("sliding_window_pattern: want an integer or a list, got %s", data) } - *p = 0 + *p = slidingWindowPattern{IsList: true} return nil } diff --git a/mlxrunner/kvsize/kvsize_test.go b/mlxrunner/kvsize/kvsize_test.go index 008872ff133..affb60c09ab 100644 --- a/mlxrunner/kvsize/kvsize_test.go +++ b/mlxrunner/kvsize/kvsize_test.go @@ -434,3 +434,36 @@ func TestUnparsableConfigCarriesItsError(t *testing.T) { t.Fatal("Err is nil for a config that does not parse") } } + +// A gemma4-family config whose sliding_window_pattern is a list and that has no +// layer_types cannot be priced: gemma4.go reads the field as an int32, so the +// model itself would fail to parse it, and the period the list implies is not +// the rule's to guess. It must come back unknown (the caller falls back), not +// priced with the default period of 5. +func TestGemma4ListPatternWithoutLayerTypesIsUnknown(t *testing.T) { + var cfg map[string]any + if err := json.Unmarshal(load(t, gemma4Config), &cfg); err != nil { + t.Fatal(err) + } + text := cfg + if tc, ok := cfg["text_config"].(map[string]any); ok { + text = tc + } + delete(text, "layer_types") + text["sliding_window_pattern"] = []int{1024, 1024, 1024, 1024, 1024, 0} + data, err := json.Marshal(cfg) + if err != nil { + t.Fatal(err) + } + if got := Model(data, nil, 8192); got.Known { + t.Errorf("known = true (%d bytes, %+v), want unknown", got.Total(), got.Layers) + } + // The integer form still prices as before. + text["sliding_window_pattern"] = 6 + if data, err = json.Marshal(cfg); err != nil { + t.Fatal(err) + } + if got := Model(data, nil, 8192); !got.Known { + t.Error("an integer sliding_window_pattern without layer_types is no longer priced") + } +}