ggml-metal: fix int32 overflow in kernel_mul_mm batched dst offsets (deep-context corruption)

[feni6]

Summary

kernel_mul_mm computes its per-batch output offsets in 32-bit:

• dstBatch = dst + im * N * M (dst4 path)
• ... + im*args.ne1*args.ne0 (tensor path, two sites)

For a large f32 destination the product overflows int32 once a batch's base offset reaches 2³¹ elements. The visible case is attention KQ = [n_kv, n_ubatch, n_head] at long context: head h's base wraps once h·n_ubatch·n_kv ≥ 2³¹, i.e. first at h = ceil(2³¹/(n_ubatch·n_kv)). The wrapped heads are stored ~8.59 GiB below the tensor — silently corrupting whatever buffer happens to live there in GPU address space — while their own region is never written, so softmax ingests stale scratch.

This PR promotes the offset arithmetic to uint64_t at the three affected sites. Three lines, no functional change below 2³¹.

Real-world impact

GLM-5.3-Flash (#27754, #27752) collapses to an infinite @ (token 31) from the first sampled token once a joint (depth, n_ubatch, n_ctx) threshold is crossed on Metal — e.g. a 108,710-token prompt at -c 131072 -ub 512. The model's dense-masked DSA attention keeps KQ at full [n_kv, 512, 64] during prefill, crossing the 13 GiB f32 mark deep into context. Whether the wrap is fatal depends on what lives below the compute buffer — on our hosts, the model's own K-cache (GPU-VA adjacent, gap 0.000 GiB), which explains the depth/n_ctx dependence of the failure boundary: full diagnosis in this comment.

Any model/backend-path combination that drives mul_mm with ≥2³¹-element batched f32 output is exposed; deep-context MLA/DSA-style attention is simply the first to get there.

Evidence

• Standalone kernel probe (no model needed) reproduces the corruption at exact real-workload geometries; first faulty head matches ceil(2³¹/(ub·n_kv)) across five geometries and three machine classes: M4 Pro, M3 Ultra (ours), M3 Max (independent confirmation by @eauchs, including the negative control: casts removed → probe fails again).
• With the casts: probe 256/256 sampled corners at all previously-failing geometries; canary-buffer runs show the out-of-tensor stores are gone.
• Real model: an 11-point boundary battery (formerly-GARBAGE wrap geometries up to 480K tokens @ -c 524288) all pass on the patched build; byte-identical prompts on the unpatched build still collapse (battery summary).
• Non-regression on Apple Silicon (run by @eauchs on M3 Max): NMSE unchanged (8.52e-08 / 2.75e-14 / 4.84e-14 on glm5next at -s 1234), test_mtp draft+reload OK on three backends, seven neighbouring archs FAIL=0.

Notes

• The equivalent CUDA/Vulkan paths already use wider indexing for these offsets; Metal was the outlier.
• -ub 128 is the workaround in the wild for affected models (keeps h·ub·n_kv < 2³¹ longer); this fix removes the need for it.
• Probe source and the landing-band analysis are linked from the diagnosis comment above; happy to add the probe as a test if maintainers want it in-tree.

Fixes the Metal deep-context collapse reported in #27754 / #27752 (model PRs, not yet merged; the defect is in shared ggml code and reproduces on both).

🤖 Generated with Claude Code

https://claude.ai/code/session_01YUN8B4NPpyNhziETmLiNpk

[ggml-gh-bot]

Hi @feni6, thanks for your contribution!

Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:

• PR Template not respected: Please respect the template when creating a new pull request. Make sure to fill out all required sections.

• AI-generated content: While code is allowed to be generated by AI, please write the PR description and commit messages on your own without the help of AI.

---

Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below.