glm5.2: streamed prefill fails when the boot model map does not cover routed experts

[sk8erboi17]

Summary

On machines where the model is much larger than RAM, the glm5.2 branch fails every multi-token prefill with metal GLM prefill failed (HTTP 500 from ds4-server, prompt processing failed from the CLI). Single-token decode is unaffected, which makes the failure look intermittent: a "hello" probe works, the first real prompt dies.

Root cause: in SSD streaming mode the boot Metal model map is restricted to a small span of the GGUF, but both streamed prefill paths resolve routed-expert tensor ranges through ds4_gpu_wrap_model_range(), which is a lookup into the boot views only and returns nil for anything outside them. This PR makes that lookup fall back to an on-demand cached exact view instead of failing the whole graph.

Environment

• MacBook Pro M2 Max, 96 GB RAM, macOS (Darwin 25.5.0), Metal backend
• branch glm5.2 @ bd89932
• model GLM-5.2-UD-Q2_K_RoutedQ2K.gguf (262 GB, from antirez/GLM-5.2-GGUF, fetched with ./download_model.sh glm-antirez-q2)
• --ssd-streaming; reproduced with the auto cache budget and with explicit --ssd-streaming-cache-experts 32GB / 44GB

Reproduction

./download_model.sh glm-antirez-q2
./ds4 -m gguf/GLM-5.2-UD-Q2_K_RoutedQ2K.gguf --ssd-streaming -c 8192 --nothink \
      --prompt-file six-hundred-token-prompt.txt -n 8

Boot shows the restricted map and the auto-budget warning:

ds4: SSD streaming auto cache budget
ds4:   metal recommends 84.00 GiB working set
ds4:   using 80% total for model + cached experts: 67.20 GiB
ds4:   non-routed weights: 19.30 GiB
ds4:   expert budget before prefill/full-layer reserve: 1040 (12.00 GiB)
ds4:   GLM Metal auto cache capped at 12.00 GiB; pass --ssd-streaming-cache-experts NGB to override
ds4: WARNING: SSD streaming expert cache (528 experts) is under twice the per-token routed working set (75 layers x 8 experts = 600); expect heavy thrashing below 13.84 GiB
ds4: SSD streaming initial metal model map restricted to token embedding (1 spans, 0.94 GiB tensor span)

and the prefill dies on the default path:

ds4: Metal streaming prefill batch selected addr failed to map overflow expert views at layer 3
ds4: prompt processing failed: metal GLM prefill failed

With DS4_METAL_GLM_STREAMING_PREFILL_FULL_LAYER=1 the failure just moves to the full-layer path:

ds4: Metal model range 3.25..4.23 GiB is not covered by mapped model views
ds4: prompt processing failed: metal GLM prefill failed

The same happens through ds4-server (/v1/chat/completions returns 500 metal GLM prefill failed for a 614-token prompt; a 25-token probe and single-token decode work fine).

Root cause

• In streaming mode ds4_gpu_set_model_map_range() maps only the token-embedding span, so g_model_views covers 0.94 GiB of a 244 GiB file.
• ds4_gpu_wrap_model_range() resolves a range strictly against those boot views and returns nil (after printing "not covered by mapped model views") when the range lies outside them.
• The overflow expert views introduced in d27a11d ("streaming: keep prefill logits cache-size invariant via overflow expert views") wrap the whole routed gate/up/down tensors of a layer through exactly this lookup — on this machine class those ranges are never covered, so the overflow path cannot work at any cache budget, which defeats its purpose (the comment in ds4_gpu_stream_prefill_batch_selected_addr_enabled() says the addr path "must stay reachable at ANY cache budget").
• The full-layer prefill path resolves whole-layer tensor ranges through the same lookup and fails the same way.

On machines with enough RAM for the boot map to cover the file (or budgets large enough that prefill batches never overflow the expert cache), the lookup always hits and the bug stays latent — which is presumably why the branch works on the machines it was developed on.

The fix

ds4_gpu_wrap_model_range() now falls back to ds4_gpu_wrap_model_exact_range() when the requested range is not covered by the boot views, with a one-time stderr note instead of a hard error:

• same bytes, same offsets — the exact view is a page-aligned newBufferWithBytesNoCopy over the same mmap, so computed logits are unchanged;
• the exact-view NSCache already handles lifetime, and buffers encoded into a command buffer stay retained until completion, so eviction is safe;
• configurations whose boot views cover the request never reach the fallback, so there is no behavior change on the machines where the branch already works.

Validation

On the M2 Max / 96 GB machine, with the patch and DS4_METAL_GLM_STREAMING_PREFILL_FULL_LAYER=1:

processing 614 input tokens: 614/614 (100.0%)
ds4: prefill: 7.64 t/s, generation: 0.57 t/s   (auto 12 GiB expert budget)
ds4: prefill: 7.33 t/s, generation: 0.96 t/s   (--ssd-streaming-cache-experts 32GB)
ds4: prefill: 7.58 t/s, generation: 0.93 t/s   (--ssd-streaming-cache-experts 44GB)

614- and 2014-token prompts also complete end-to-end through ds4-server (/v1/chat/completions). Sustained chat use works (slow, as expected for a 262 GB model streamed on 96 GB — this is inspection-grade, and that's fine).

Regression tests on the patched checkout:

• ./ds4_test --metal-kernels → metal-kernels: OK
• the model-free parts of make test pass (Q4_K unit tests 4/4, ds4-eval --self-test-extractors, ds4_agent_test); the model-dependent tracks (--long-context, logprob vectors, the GLM 100-case fixture) were not run on this machine — no ds4flash.gguf in the test checkout, and I did not have official-logit fixtures. Happy to run more if you tell me which track matters most for this change.

Open question (separate bug?)

Even with this fix, the default batch-selected-addr path still fails on this machine: the fallback note prints, ~18 s of work happen, then the request returns metal GLM prefill failed with no further stderr — something after the overflow-view mapping returns 0 silently (set_addr_slot/addr_buffers print nothing, so it is downstream of those). Forcing full-layer prefill (DS4_METAL_GLM_STREAMING_PREFILL_FULL_LAYER=1) is what actually makes GLM usable here, but the auto rule (glm_graph_stream_prefill_full_layer_min_tokens()) leaves common prompt sizes on the broken path.

Two things you may want to consider (not included in this PR to keep it minimal):

1. auto-enabling full-layer prefill when the boot model map is restricted (or when the expert budget lands under the per-token working set, which the code already warns about);
2. digging into why the addr path fails even once its overflow views resolve — I can gather more evidence on this machine if useful.