Metal: fix grid-size truncation in coli_metal_gemm at large prefill (long-context corruption)#488
Metal: fix grid-size truncation in coli_metal_gemm at large prefill (long-context corruption)#488JustVugg merged 3 commits intoJustVugg:devJustVugg/colibri:devfrom RDouglasSharp:fix/metal-gemm-grid-truncationRDouglasSharp/colibri:fix/metal-gemm-grid-truncationCopy head branch name to clipboard
Conversation
|
Details of the investigation that led to this fix, and some remaining issues using OpenCode with Colibri server: |
|
To merge this I need validation data I can't produce myself (no such GPU here). The bar every backend holds: (1) correctness first — a teacher-forcing token-exact check on a tiny model ( Specifically for this grid-size/long-context fix: a before/after on a long prefill (the case that corrupted) showing coherent output after — ideally token-exact vs CPU on a tiny model. Thanks @RDouglasSharp. |
…fill Large dense GEMMs (NT=S*O elements) were dispatched as a single grid; past the device threadgroup-grid limit the tail output rows were silently never computed, leaving stale scratch -> nondeterministic long-context corruption under COLI_METAL=1 (identical prompt correct on CPU / at short lengths). Chunk the dispatch to <=2^25 elements via g_gx/g_gy buffer offsets: no kernel change, disjoint output ranges, negligible overhead, no-op for small S and decode. Gated COLI_GEMM_CHUNK (default on; =0 restores old path for A/B). Adds tests/test_gemm_largebatch.mm (make gemm-test): kv_b S=7478 nerr=1.0 before, clean after.
16e339c to
a0b929c
Compare
|
Hi @RDouglasSharp — I rebased this onto The only conflict was I read the fix and I like it: chunking to One thing before we merge: I can't validate this. I'm on Linux — the Metal path never compiles here, so all I can honestly say is "the CPU build is unchanged and the Makefile parses". Could you run your own test on your Mac and paste the output? Fail-before / pass-after in one paste and I'll merge. If you have a real long-context prefill handy (the |
|
Sorry - my internet provider is down. Once back up I will send requested
validation.
…On Tue, Jul 21, 2026 at 10:38 PM Vincenzo Fornaro ***@***.***> wrote:
*JustVugg* left a comment (JustVugg/colibri#488)
<#488 (comment)>
Hi @RDouglasSharp <https://github.com/RDouglasSharp> — *I rebased this
onto dev for you and force-pushed to your branch* (you had
maintainerCanModify on, thanks). Contributors shouldn't pay for our merge
order, so we do the rebases here.
The only conflict was c/Makefile: dev rewrote the cuda-test recipe in the
AMD/HIP single-source refactor (#339
<#339>) to use $(GPUCC)/
$(GPUFLAGS)/backend_gpu_compat.h. Resolution keeps *both* — your new
gemm-test target verbatim, and dev's newer cuda-test. Nothing of yours
was dropped; the diff is still exactly your 3 files (Makefile +6,
backend_metal.mm, tests/test_gemm_largebatch.mm). PR is MERGEABLE now.
I read the fix and I like it: chunking to 2^25 elements per dispatch with
COLI_GEMM_CHUNK=0 as an escape hatch means the bug can be A/B'd on a
single binary — that's exactly the kind of gate we want (we got burned
recently merging a cache PR that was byte-identical but silently killed a
fast path, #490 <#490>).
*One thing before we merge: I can't validate this.* I'm on Linux — the
Metal path never compiles here, so all I can honestly say is "the CPU build
is unchanged and the Makefile parses". Could you run your own test on your
Mac and paste the output?
make gemm-test # with the fix
COLI_GEMM_CHUNK=0 ./gemm_largebatch_test # pre-fix behavior, same binary
Fail-before / pass-after in one paste and I'll merge. If you have a real
long-context prefill handy (the S=7478 case from your report), that on
top would be ideal.
—
Reply to this email directly, view it on GitHub
<#488?email_source=notifications&email_token=ATMRTCNWMUG7YCZ3XUTAHWD5F7PF5A5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTKMBTHEZTAMZYGQ4KM4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJLDGN5XXIZLSL5RWY2LDNM#issuecomment-5039303848>,
or unsubscribe
<https://github.com/notifications/unsubscribe-auth/ATMRTCOWACBRZPD2FHQPCX35F7PF5AVCNFSNUABGKJSXA33TNF2G64TZHMYTEOBWGA4DAMZZG45US43TOVSTWNBZGM3TIOJTGYZDJILWAI>
.
Triage notifications, keep track of coding agent tasks and review pull
requests on the go with GitHub Mobile for iOS
<https://github.com/notifications/mobile/ios/ATMRTCLC3TWQT5CKNBZVOQL5F7PF5A5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTKMBTHEZTAMZYGQ4KM4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJKTGN5XXIZLSL5UW64Y>
and Android
<https://github.com/notifications/mobile/android/ATMRTCNSWTAU4KML2GL3RQT5F7PF5A5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTKMBTHEZTAMZYGQ4KM4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJLTGN5XXIZLSL5QW4ZDSN5UWI>.
Download it today!
You are receiving this because you were mentioned.Message ID:
***@***.***>
|
|
(1) Correctness — teacher-forcing, token-exact, GPU vs CPU. Tiny oracle built with tools/make_glm_oracle.py. At f32 the engine matches the transformers oracle exactly (SNAP=./glm_tiny TF=1 ./colibri 64 16 16 → 32/32). To put the work on the patched coli_metal_gemm I ran int8 (S=32 ≥ g_metal_gemm_min=16): COLI_METAL=1 … 64 8 8 → 30/32 vs oracle; METAL: blocchi GPU 4 | fallback CPU 0 | expert su GPU 12 Metal and CPU are token-identical at all 32 positions — the same two positions diverge from the bf16 oracle (pos 5→197, pos 25→136) on both paths, i.e. that 2/32 is int8 quantization vs bf16, not a GPU artifact. The GPU path reproduces CPU bit-for-bit with the fix in place. |
|
git fetch origin down (O6144,I2048) S=512 O=6144 I=2048 nerr=1.84e-06 det:same ok first-bad row=2 col=1435 (gpu=4.094e-05 cpu=4.183e-05) kv_b (O28672,I512) S=512 O=28672 I=512 nerr=1.25e-06 det:same ok first-bad row=3 col=28289 (gpu=1.98e-08 cpu=8.291e-08) GEMM large-batch: all clean COLI_GEMM_CHUNK=0 ./gemm_largebatch_test down (O6144,I2048) S=512 O=6144 I=2048 nerr=1.84e-06 det:same ok first-bad row=2 col=1435 (gpu=4.094e-05 cpu=4.183e-05) kv_b (O28672,I512) S=512 O=28672 I=512 nerr=1.25e-06 det:same ok first-bad row=3 col=28289 (gpu=1.98e-08 cpu=8.291e-08) GEMM large-batch: CORRUPTION DETECTED |
|
I am happy to post back to back to back CPU vs METAL CHUNK-0 vs METAL with fix, but I am in the middle of upgrading to the improved 4-bit model (downloading in progress), so it's going to take a while. The summary of the results for a very long prompt on the older 4-bit model: without the fix, catastrophic corruption, after the fix, absolute prose quality is limited by the pre-#225 checkpoint, and does not perfectly match CPU result, but resembles CPU result. |
|
Reviewed the change itself, since the one thing I can contribute here is a careful read — I have no Apple silicon, so the hardware evidence has to come from you, and it does: a same-binary A/B ( The chunking logic reads correct to me: chunks write disjoint One real problem, in the disable path: const int64_t NT_MAX = chunk_on ? ((int64_t)1<<25) : ((int64_t)1<<62);
int CH=(int)(NT_MAX/O);With chunking off, Clamping before the narrowing fixes it: int64_t ch64 = NT_MAX/O; if(ch64<1) ch64=1; if(ch64>S) ch64=S; int CH=(int)ch64;One question: the comment justifies the buffer offsets as 16B-aligned because Fix the cast and I'm ready to merge this — the corruption it removes makes Metal usable for agent workloads, and CI is green. Thanks also for flagging the |
…grid NT=S*O, not S; S=4376 is clean)
|
Nice catch! Fixed in 70f06cd — clamped in 64-bit before narrowing, exactly as you suggested: int64_t ch64=NT_MAX/O; if(ch64<1) ch64=1; if(ch64>S) ch64=S; int CH=(int)ch64; After the clamp CH ≤ S ≤ INT_MAX, so the conversion is defined on any toolchain. You were right that this was the dangerous one — a negative truncation would have given CH=1, one row per dispatch, and the "before" case would have come out clean, silently defanging the A/B it exists for. Re-ran the A/B on the fixed binary to confirm the escape hatch still demonstrates the bug: make gemm-test down (O6144,I2048) S=512 O=6144 I=2048 nerr=1.84e-06 det:same ok first-bad row=2 col=1435 (gpu=4.094e-05 cpu=4.183e-05) kv_b (O28672,I512) S=512 O=28672 I=512 nerr=1.25e-06 det:same ok first-bad row=3 col=28289 (gpu=1.98e-08 cpu=8.291e-08) GEMM large-batch: all clean down (O6144,I2048) S=512 O=6144 I=2048 nerr=1.84e-06 det:same ok first-bad row=2 col=1435 (gpu=4.094e-05 cpu=4.183e-05) kv_b (O28672,I512) S=512 O=28672 I=512 nerr=1.25e-06 det:same ok first-bad row=3 col=28289 (gpu=1.98e-08 cpu=8.291e-08) GEMM large-batch: all clean down (O6144,I2048) S=512 O=6144 I=2048 nerr=1.84e-06 det:same ok first-bad row=2 col=1435 (gpu=4.094e-05 cpu=4.183e-05) kv_b (O28672,I512) S=512 O=28672 I=512 nerr=1.25e-06 det:same ok first-bad row=3 col=28289 (gpu=1.98e-08 cpu=8.291e-08) GEMM large-batch: CORRUPTION DETECTED |
|
Also corrected two stale comments in the test that the sweep above contradicts: it labelled S=4376 the "corrupt bracket" and described an "S>=~2400 fault." Both wrong — the trigger is the dispatch grid (NT=S·O), not S alone, which is why only kv_b (O=28672) crosses the limit and only at S=7478, while gate/up/down stay clean at every S in the sweep. Comment-only change, no behaviour. |
python3 openai_server.py --model <dir> looked for a binary named 'glm' next to itself. Since JustVugg#391 the build produces 'colibri', so direct invocation was broken on any clean checkout -- it only appeared to work in trees that still had a stale glm from an older build. Spotted by @RDouglasSharp while working on JustVugg#488. Resolve the engine by probing colibri / colibri.exe first and falling back to glm / glm.exe, the same order the coli launcher uses, so old trees keep starting. Verified by removing the stale glm and running the gateway: it resolves to colibri and the engine launches. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Merged. The clamp is exactly right, and thank you for going further than asked: correcting the two stale comments in the test — the ones that blamed For the record, what carried this over the line without Apple silicon on my side: a same-binary A/B ( Also thanks for spotting the |
…the falloc->qalloc trade-off - test_gemm_largebatch.mm predates the 9-arg coli_metal_gemm prototype; gs=0 selects per-row scales, the exact pre-fmt=4 semantics its fixtures assume. - The fmt=4 scale-buffer qalloc site is outside #ifdef COLI_METAL: on CPU-only builds it trades falloc's checked-exit for qalloc's unchecked malloc, consistent with qsalloc for fmt 1/2/3 -- now said in the comment, with a note that fmt=5's group scales still use falloc (Metal-inert, pre-existing).
…the falloc->qalloc trade-off - test_gemm_largebatch.mm predates the 9-arg coli_metal_gemm prototype; gs=0 selects per-row scales, the exact pre-fmt=4 semantics its fixtures assume. - The fmt=4 scale-buffer qalloc site is outside #ifdef COLI_METAL: on CPU-only builds it trades falloc's checked-exit for qalloc's unchecked malloc, consistent with qsalloc for fmt 1/2/3 -- now said in the comment, with a note that fmt=5's group scales still use falloc (Metal-inert, pre-existing).
Metal: fix grid-size truncation in
coli_metal_gemmat large prefill (long-context corruption)Summary
On the Metal backend, dense GEMMs with a very large output (
S × Oelements) are dispatched as asingle grid. Past a device threadgroup-grid limit the grid is silently truncated — the tail
output rows are never computed and keep whatever was in the buffer. In the engine that buffer is
reused scratch, so the stale contents are nondeterministic garbage → corrupted generation at long
context. This fixes it by chunking the dispatch so no single launch exceeds a safe grid size. No
kernel change; output is bit-exact vs. CPU.
Motivation — this is what makes Metal + OpenCode/agent backends usable
colibri on Metal (
COLI_METAL=1) is currently unusable as an OpenCode / coding-agent backendbecause of this bug. Agent front-ends prepend a large fixed system header — tool declarations, rules,
worked examples — that is several thousand tokens before the user types anything and cannot be
trimmed. That header alone requires a large context (we run
CTX=40960just to fit it) and drives abig prefill on the first message of every session. That prefill is exactly the regime that
crosses the grid limit, so the GPU path returns garbage from the first token, while the CPU path is
correct but far slower. In other words: without this fix, Metal + OpenCode simply does not work; with
it, the GPU path produces coherent output and well-formed tool calls at real agentic context sizes.
(Nobody upstream hit this because Metal validation uses tiny 8–13-token prompts — see "Why it evaded
detection".)
Symptom / impact
COLI_METAL=1produce corrupted output (dropped/blank tokens, degrading to tokensalad); the identical prompt is correct on CPU (
COLI_METAL=0) and correct on Metal at shortlengths.
on Metal — i.e. the GPU backend is effectively dead for coding-agent workloads until this lands.
Root cause
matmul_qt_exroutes dense matmuls withS ≥ COLI_METAL_GEMM_MINtocoli_metal_gemm, whichdispatches
NT = S*Ooutput elements as one grid of(NT+3)/4threadgroups. The largest projection(
kv_b, O≈28672) crosses the device grid limit at largeS: on M-series,kv_bat S=4376(grid ≈ 3.1e7 tg) computes fully, but at S=7478 (grid ≈ 5.4e7 tg) the tail is not computed. It is
not an out-of-bounds (Metal shader validation is silent) and not dispatch-arithmetic
overflow (
NTstays ≈ 9e7 ≪ 2³¹) — it's the raw grid size exceeding what a single dispatch schedules.Fix
Chunk
coli_metal_gemmover output rows so each dispatch is ≤ 2²⁵ elements (grid ≤ 2²³ tg, ~4× underthe observed-clean bound), using
g_gx/g_gybuffer offsets — no shader change, chunks writedisjoint output ranges, negligible overhead, and it's a no-op for small
S(single chunk) and decode(S=1). Gated by
COLI_GEMM_CHUNK(default on;=0restores the old single-dispatch path for A/B).Reproduction & test (seconds, no model)
New
make gemm-test(tests/test_gemm_largebatch.mm) callscoli_metal_gemmdirectly across a sizesweep of real GLM shapes vs. an OMP CPU reference, and includes a GPU-vs-GPU determinism check:
COLI_GEMM_CHUNK=0):kv_b S=7478→nerr = 1.00e+00(*** CORRUPT).nerr ~1e-6.Validation
make gemm-test: all clean after the fix; the S=7478kv_bcase is corrupt withCOLI_GEMM_CHUNK=0and clean with it on (same binary A/B).
calls (previously token salad).
Why it evaded detection
The truncation is at the grid boundary, so it's nondeterministic across process launches (a
given launch may land just under or just over) and only manifests at agentic prefill sizes.
Existing Metal kernel tests exercise
S ≤ 64; nothing upstream drives a multi-thousand-token prefillthrough Metal, so it never surfaced.
Scope / not included
moe_submit/moe_gemv) has the same single-dispatch pattern andcould hit the limit at very large per-block row counts; not triggered in the confirmed repro
(MoE-on-GPU alone is clean at these sizes), so it's left for a follow-up with its own test.
performance limitation, not addressed here.
Investigation notes — hypotheses ruled out (full log attached)
Ruled out en route (each by a measured run; full forensic detail in the attached investigation log):
DSA sparse attention (model has no indexer weights), RoPE extrapolation (max_position 1M, per-position
angles), CTX-scaled allocation, int32
NToverflow, CPU-attention score buffers (#110/#117, target8192), the residency set (
COLI_METAL_RESSETon/off both corrupt), the async I/O pipeline (PIPE=0still corrupt), pilot prefetch (off by default), a GEMM+MoE interaction (corrupts with MoE on CPU),
and macOS-compressor eviction of scratch (
setBuffer:-bound buffers are auto-resident). The 2×2GEMM×MoE placement matrix isolated corruption to dense GEMM on GPU; the standalone
gemm-testthen proved the GEMM is bit-exact in isolation, localizing the fault to the single-dispatch grid
size — deterministically reproduced once the sweep reached
kv_b S=7478.Separately noticed (not fixed here):
openai_server.py's--enginedefault still points atglmafter theglm.c→colibri.crename (#391); direct invocation is broken on a clean checkout,and silently runs a stale
glmif one exists. Worth a one-line fix (default tocolibri, fall backto
glm).