fix: KV estimation misses n_layer / n_embd for GPT-2 and GPT-BigCode #927

Problem / Background

src/execution/kv_arch.rs classifies a model's KV geometry by reading field names out of config.json. The lookups accept the modern HuggingFace spellings plus several alternates, but never the OpenAI-era names that GPT-2 and GPT-BigCode configs actually use.

• Layer count (kv_arch.rs:244): classify() reads get_u64(text, &["num_hidden_layers", "n_layers", "num_layers"])?. Both families write n_layer (singular, no plural alias). The lookup returns None, the ? short-circuits classify(), and the entire estimate becomes None.
• Hidden size (kv_arch.rs:163): attn_dims() derives head_dim from get_u64(text, &["hidden_size", "d_model", "dim", "model_dim"])? when no explicit head_dim is present. Both families write n_embd, so even with the layer count repaired the head dim would still be underivable.
• Head count is already correct: the num_heads lookup at kv_arch.rs:144 does alias n_head, so only the two fields above are missing.

estimate_total_memory maps the resulting None to KvSource::Unavailable and reports zero bytes (memory_estimate.rs:386-390). The same root cause hits a sibling parser: activation_dims_from_path (memory_estimate.rs:510) uses the same hidden-size alias list without n_embd, so the activation reserve collapses to zero as well.

Observed behavior

Run against real checkpoints with target/release/mlxcel. First reported at 051526fe2 and re-run on main at a3c823611, after #928, #930 and #946 merged; the figures are unchanged:

$ ./target/release/mlxcel inspect -m models/gpt2

=== Memory Estimate ===
  Model:           models/gpt2
  Context length:  8192 tokens (batch = 1)
  Quant hint:      default (from config.json)
  KV dtype:        fp16

  Weights:         522.7 MiB (548090880 bytes)  (safetensors header)
  KV cache:        0 bytes  (unavailable (config.json missing architecture fields))
  Activation:      0 bytes  (batch 1 x <=512 prefill tokens x (hidden+intermediate) + logits)
  Allocator ovhd:  104.5 MiB (109618175 bytes)  (factor 1.20x on weights+kv)
  -----
  Total estimate:  627.2 MiB (657709055 bytes)
  Available:       121.60 GiB (130567005798 bytes)

  FITS: 120.99 GiB (129909296743 bytes) of headroom

$ ./target/release/mlxcel inspect -m models/gpt_bigcode-santacoder

=== Memory Estimate ===
  Model:           models/gpt_bigcode-santacoder
  Context length:  8192 tokens (batch = 1)
  Quant hint:      default (from config.json)
  KV dtype:        fp16

  Weights:         2.10 GiB (2249773056 bytes)  (safetensors header)
  KV cache:        0 bytes  (unavailable (config.json missing architecture fields))
  Activation:      0 bytes  (batch 1 x <=512 prefill tokens x (hidden+intermediate) + logits)
  Allocator ovhd:  429.1 MiB (449954611 bytes)  (factor 1.20x on weights+kv)
  -----
  Total estimate:  2.51 GiB (2699727667 bytes)
  Available:       121.60 GiB (130567005798 bytes)

  FITS: 119.09 GiB (127867278131 bytes) of headroom

Both families are reported as having no KV cache and no activation reserve at all, which understates the real footprint of a long-context run.

Which families are affected

Only GPT-2 and GPT-BigCode. The three text families that merged after this issue was filed were each checked against a real checkpoint on main at a3c823611 and all three resolve normally, because all three use the modern field naming:

Bailing MoE was checked specifically because MoE geometry could have diverged: it does not. The sparse block changes the FFN only, the attention stack is ordinary GQA, and num_key_value_heads: 4 is picked up by the existing alias list (28 layers x 4 kv heads x 128 head_dim x 2 x 2 bytes = 57,344 bytes/token, x 8192 = 469,762,048). Activation reserves are non-zero for all three as well (21072128, 19756800 and 7330816 bytes respectively). None of them needs a change here, and none should regress.

Scope of impact today

This is a reporting defect, not a correctness or stability defect. Filed at low priority deliberately:

• Paged decode never reaches these models. LanguageModel::supports_paged_decode_backend() defaults to false (src/lib/mlxcel-core/src/generate.rs:601) and neither src/models/gpt2.rs nor src/models/gpt_bigcode.rs overrides it, so the scheduler's paged_available gate (src/server/batch/scheduler.rs:199) is never satisfied. The block-budget consumer of this estimator is unreachable for GPT-2 and GPT-BigCode.
• The failure is graceful, not silent-wrong. paged_block_bytes (memory_estimate.rs:810) returns None when the per-layer per-token rate is zero, and resolve_paged_block_budget propagates that None, which leaves the pool unbounded rather than installing a bogus cap. Zero is a fail-open value here, not a fail-wrong one.
• The live user-visible surfaces are mlxcel inspect (the two runs above) and the KV cache mode advisor: advise_kv_cache_modes_from_config (kv_cache_advisor.rs:285) resolves its architecture through the same classifier and returns an empty advice vector when the classifier gives up, so no KV cache mode suggestions are printed for either family.

Expected behavior

Both families should classify as KvArchKind::Standard and produce a real figure. With fp16 KV at ctx 8192, batch 1:

The activation reserve should likewise be non-zero for both, since n_embd is present and n_inner / the 4 x hidden fallback covers the intermediate size.

The trap: adding the aliases alone over-estimates GPT-BigCode by 16x

Read this before implementing. The obvious fix, appending n_layer and n_embd to the two alias lists, is wrong for GPT-BigCode.

Every other family in this codebase expresses grouped or multi-query attention through a numeric field, and the num_kv_heads lookup at kv_arch.rs:150 reflects exactly that: num_key_value_heads, num_kv_heads, n_kv_heads, n_head_kv, multi_query_group_num. When none of them is present, it falls back to num_kv_heads = num_heads, that is, full MHA.

GPT-BigCode carries none of those fields. It signals multi-query attention with a boolean, multi_query: true, which means exactly one kv head shared across all query heads. models/gpt_bigcode-santacoder/config.json has "n_head": 16 and "multi_query": true, and the model implementation agrees: GptBigCodeArgs::num_kv_heads() (src/models/gpt_bigcode.rs:280) returns if self.multi_query { 1 } else { self.n_head }, with src/models/gpt_bigcode_tests.rs:672 asserting the cache holds one kv head, not n_head.

So an alias-only fix would take the num_kv_heads = num_heads fallback and estimate 16 kv heads where the runtime caches 1:

A 16x over-estimate that looks plausible is worse than an obvious zero: it would reserve 1.5 GiB of phantom budget and, once these families ever become pool-backed, shrink the admission budget accordingly. kv_arch.rs needs an explicit branch for the boolean multi_query, not just two more strings in the alias lists.

Acceptance criteria

• kv_arch.rs layer-count lookup (classify, around line 244) accepts n_layer in addition to the existing num_hidden_layers / n_layers / num_layers.
• kv_arch.rs hidden-size lookup (attn_dims, around line 163) accepts n_embd in addition to the existing hidden_size / d_model / dim / model_dim.
• The num_kv_heads resolution honors a boolean multi_query: true as exactly one shared kv head, taking precedence over the num_kv_heads = num_heads MHA fallback. multi_query: false and an absent multi_query both keep the existing MHA behavior.
• estimate_kv_arch for models/gpt_bigcode-santacoder at ctx 8192 / fp16 / batch 1 returns 100,663,296 bytes (96 MiB) and KvArchKind::Standard, not the 1.50 GiB an alias-only change would produce.
• estimate_kv_arch for models/gpt2 at ctx 8192 / fp16 / batch 1 returns 301,989,888 bytes (288 MiB) and KvArchKind::Standard.
• mlxcel inspect -m models/gpt2 and mlxcel inspect -m models/gpt_bigcode-santacoder both report a non-zero KV cache line sourced from the config, and no longer print "unavailable (config.json missing architecture fields)". Both families are covered, not only the one that triggered this report.
• The activation reserve is non-zero for both families: activation_dims_from_path (memory_estimate.rs:510) resolves the hidden size from n_embd.
• kv_cache_params_from_path (memory_estimate.rs:710 onward, alias lists from line 722), which duplicates both alias lists and the kv-head fallback, is brought in line so it does not disagree with the classifier.
• advise_kv_cache_modes returns a non-empty advice vector for both families.
• Unit tests cover: GPT-2 style config (n_layer / n_embd / n_head, no multi_query) classifying as standard MHA; GPT-BigCode style config with multi_query: true yielding one kv head; and the same config with multi_query: false degenerating to n_head kv heads.
• No regression in the existing kv_arch test suite; families that already resolve their layer count and hidden size keep byte-identical estimates. Verified concretely against the three families checked above: models/pythia-1b stays at 1073741824 bytes, models/helium-1-preview-2b-4bit at 2013265920, and models/ling-lite-1.5 at 469762048.

Technical considerations

• The multi_query branch belongs in attn_dims next to the existing num_kv_heads lookup, ordered so an explicit numeric field still wins if a config somehow carries both. Guard against the degenerate multi_query: true plus n_head: 0 shape the model-side validation already rejects.
• read_head_dim in kv_cache_advisor.rs (around line 347) carries the same hidden-size alias list minus n_embd. It currently returns None for both families, and because turbo_ok uses is_none_or, the non-power-of-two Turbo guard defaults to permissive. GPT-2 (64) and GPT-BigCode (128) are both powers of two so nothing misbehaves today, but the list should stay consistent with attn_dims as the doc comment on that function already promises.
• Consider factoring the repeated alias lists into shared constants. They are currently spelled out four times across kv_arch.rs, memory_estimate.rs, and kv_cache_advisor.rs, which is why this gap reproduced in every one of them.

References

Surfaced while reviewing #926 (GPT-BigCode, issue #834); the same gap is already present in the merged #924 (GPT-2, issue #833).