Skip to content

perf: byte-identical CPU/startup optimizations — RoPE inv_freq cache, parallel prefill KV writes, grammar compile cache, analyze_model cache, route stack buffers, LTO - #1058

Merged
JustVugg merged 2 commits into
JustVugg:devfrom
zh-Processor:main
Aug 18, 2026

Conversation

@zh-Processor

Copy link
Copy Markdown
 ## Summary

 A set of small, independent optimizations that do not change any numerical
 output (byte-identical) or any default behavior:

 - **colibri.c: RoPE `inv_freq` precompute.** `rope_interleave` re-ran
   `half` × `powf(theta, -2j/qk_rope)` for every position, though the result
   depends only on the model constants. `inv[]` is now cached separately,
   keyed on `(qk, theta)`; only `cosf/sinf` is recomputed per position.
   `ang = pos * inv[j]` is the same float product against the same value the
   old code formed inline, so `cosf/sinf` receive the identical argument.
   Saves ~11M `powf` calls for a 4k prompt × 90 layers.
 - **colibri.c: parallelize the prefill RoPE/KV-write loop.** Each iteration
   writes a distinct KV row (its own `pos`) and the rope cache is
   `_Thread_local`, so iterations are independent and the result is
   byte-identical regardless of order. The pragma is gated to CPU-only
   builds (`#if !defined(COLI_CUDA) && !defined(COLI_VULKAN)`) because the
   CUDA/Vulkan shadow-shrink writes are shared state; GPU builds stay serial.
 - **colibri.c: schema→GBNF compile cache.** `GrDraft` keeps an owned copy of
   the schema text that built the current grammar; when a serve slot's next
   request sends the identical schema, the schema→GBNF→parse→state-init
   pipeline is skipped and only `grammar_reset` runs — exactly how
   `grammar_setup_text` finishes, so state is identical to a fresh compile.
   Targets tool-loop agents that resend one schema every turn.
 - **resource_plan.py: `analyze_model` sidecar cache.** The scan of every
   shard header + ~116k regex matches reran on every
   `coli plan/doctor/tune/run --auto-tier`. The result is a pure function of
   each shard's and config.json's `(size, mtime_ns)`, so it is cached to
   `.coli_analysis.json` with a signature check; any change re-triggers a
   full recompute. Writes are atomic (tmp file + `os.replace`). Saves
   seconds per invocation on a 372 GB model. Cache hits restore the int keys
   of `expert_bytes_by_layer` so a hit is identical to a fresh scan.
 - **deepseek_v4.c: stack routing buffers.** `coli_v4_route` /
   `coli_v4_route_bf16` malloc'd three arrays per call (per token per layer).
   Fixed 512-entry stack buffers cover every real routed-MoE config; larger
   (hostile) configs fall back to the original malloc path unchanged. Fixed
   size, not a VLA, because `n_routed_experts` has no config.json upper
   bound.
 - **deepseek_v4.c: `COLI_MODEL_DIRS`** wires `st_init_multi` into
   `coli_st_index_open`, giving v4 the same multi-drive shard spread the
   colibri engine already has. Unset → behavior identical.
 - **download_fp8.py:** `--dest` / `$GLM_DEST` instead of a hardcoded
   `I:\glm52_fp8`.
 - **Makefile:** `-flto` on all platforms (matching `Makefile.deepseek-v4`),
   `NOLTO=1` to opt out.

 Branch is merged with current `main` (1.6.2). The only semantic overlap
 was v4 OpenMP team sizing: kept upstream's `v4_omp_reserve_loader_cpus`
 (loader-aware) over this branch's physical-core sizing.

 New tests, included in the branch:

 - `tests/test_rope_invfreq.c` — differential bit-exact test, old vs new
   rope path: 15,326,112 float pairs, 0 bit mismatches (4 qk × 3 theta ×
   200k positions).
 - `tests/test_grammar_cache.c` — `grammar_reset` equals a fresh setup;
   GrDraft.src ownership checked under ASan+UBSan across 9 control-flow
   scenarios (hit/miss/compile-fail/no-grammar): no leaks, no double-free.
 - `tests/test_analysis_cache.py` — 6 tests: hit returns the identical
   result and skips the rescan, mtime change invalidates, corrupt cache
   falls back, atomic write leaves no tmp files, concurrent writers produce
   a valid final cache.

@JustVugg
JustVugg changed the base branch from main to dev August 16, 2026 15:44
@JustVugg

Copy link
Copy Markdown
Owner

Reviewed this properly. The three algorithmic optimizations are good and I want them — the reasoning is sound in each case and, unusually, you proved it instead of asserting it:

  • RoPE inv_freq cache. Correct: inv[j] is a pure function of (theta, qk_rope) and never of pos, so hoisting it leaves ang = pos * inv[j] as the identical float product and cosf/sinf receive bit-identical arguments. ~11M powf removed from a 4k × 90-layer prefill is real work, not bookkeeping.
  • Parallel prefill KV writes. The part I appreciate is the gate: you noticed that the CUDA/Vulkan shadow-shrink writes touch shared state and restricted the pragma to CPU-only builds instead of assuming independence everywhere. That is the kind of care this loop needed.
  • Grammar compile cache. Skipping schema→GBNF→parse→init when the text is byte-identical, and ending at grammar_reset — which is exactly where grammar_setup_text finishes — is the right equivalence.

And the tests are not decorative. 15,326,112 float pairs with zero bit mismatches across 4×3×200k combinations is a differential test I'd hold up as an example, and running the grammar-cache ownership cases under ASan+UBSan across nine control-flow paths (including compile-fail) is exactly where a cached-pointer design goes wrong. The analyze_model cache covering corrupt-cache fallback and concurrent writers is the same standard.

The analyze_model sidecar may be the most user-visible item here: every coli plan/doctor/tune/--auto-tier currently rescans every shard header and ~116k regex matches, and on a 372 GB model that is seconds burned per invocation, identically, forever.

Two things before I merge it.

1. Please make LTO opt-in rather than default

-flto on all five platform CFLAGS lines is a different kind of change from the other four, and I'd rather it were LTO=1 than NOLTO=1.

The reason is specific to this project: LTO enables cross-translation-unit inlining, which can change floating-point contraction (FMA formation). Token-exactness is the property colibrì sells, and we got a live demonstration today — #1044 found that olmoe's IDOT path quantizes activations, and #1024's "bit-identical" turned out to mean "bit-identical to that path" because FMA plus two accumulator chains changed the rounding (I measured 915/1024 outputs differing from stock).

Your three differential tests cover rope, grammar and the analysis cache. None of them covers what LTO does to the arithmetic, and by construction they can't — it's a whole-program property, and it will differ by compiler and version. Defaulting it on means every user silently gets a different optimization regime than the one the oracles were validated under.

As LTO=1, it is a documented, opt-in build knob and nothing about the default output changes. If you want it on by default later, the thing that would justify it is a tiny-model token-exactness run with and without -flto on at least two compilers.

2. Please also drop -Wno-stringop-overflow

I built olmoe with -flto and without that suppression on GCC here: the only output was lto-wrapper: warning: using serial compilation of 2 LTRANS jobs. So on this compiler it isn't hiding anything — which is the argument for removing it rather than keeping it.

-Wstringop-overflow is the warning class that catches buffer overruns, and this repository closed eight heap-overflow advisories last week. If LTO surfaces one of these on some other toolchain, that is information we want, not noise to suppress. (Noted with a smile that you're the person who reported those eight — which is exactly why I think you'll agree with this one.)

3. And the thing a performance PR needs: numbers

There isn't a single before/after measurement here. "~11M powf calls" and "saves seconds" are counts of work removed, not observed time. For each of the four, on your hardware:

  • prefill wall-clock / TTFT for a fixed prompt, before vs after;
  • coli plan wall-clock on a large model, cold and warm cache;
  • a tool-loop turn with a repeated schema, before vs after;
  • V4 decode tok/s, before vs after.

Even rough medians would do. The bar here isn't bureaucratic — it's that every claim in this repo gets measured, and yours are currently the only unmeasured part of an otherwise carefully evidenced PR.


Split LTO out (or flip it to opt-in), drop the warning suppression, add the numbers, and I'll take this. Retargeted to dev, which is where PRs land here — main only receives release merges.

Thanks for this, and for the eight advisories before it.

@zh-Processor

Copy link
Copy Markdown
Author

All three points addressed in the updated commit (now a single commit on current dev)

1. LTO is opt-in. -flto moved behind LTO=1 on all five platform blocks; the default build is byte-for-byte the stock regime. Your reasoning is right, and we got independent confirmation while benchmarking — see the test note below.

2. -Wno-stringop-overflow dropped. Verified it suppresses nothing on gcc 15.2 even with LTO on, so it was pure downside.

3. Numbers. Ryzen 3700X (8C/16T), 32 GB, Windows 11, MinGW gcc 15.2.0, -O3 -march=x86-64-v3. Baseline = dev @ 03e8677, PR = this commit; 5 interleaved runs, median, .coli_usage cleared before each run. Fixtures: make_glm_bench_model.py (8L/hidden 1024/32 experts, bf16, 1.2 GB), make_deepseek_v4_tiny.py, synthetic 116k-tensor/60-shard census for the plan cache.

  • analyze_model cache (median of 7): cold 312 ms (311–315) → warm 4.2 ms (4–5). ≈74× per plan/doctor/tune invocation.
  • RoPE inv_freq (micro, rope_interleave, qk=64 θ=10000, fresh position every call, best of 5×2M): 1089 → 457 ns/call (2.4×). Byte-identical: 61.5 MB old-vs-new dump, 0 diffs.
  • Grammar cache, end-to-end (serve mux, same slot, identical grammar twice): baseline compiles 2×, PR compiles 1×, outputs token-identical. Magnitude: schema→GBNF+parse+init costs 0.1 ms (10 props) … 1.7 ms (150 props, near the 1024-rule cap); cached path ≈ 0. So ≤ ~2 ms per repeated-schema request — small, claimed as such. One boundary found: the cache engages only when the previous turn left the walker alive (gr_feedon=0 on non-conforming output disables it — existing fail-soft behavior, fine for grammar-forced models).
  • GLM prefill (TF, 2047 positions): 75.9 → 76.2 pos/s — within noise, TF score identical (1/2112 both). Not claiming a long-prompt win; not measurable at fixture scale.
  • GLM decode (REPLAY, 64 tok after 2047-token prefill): 33.29 → 33.38 tok/s — within noise.
  • V4 (tiny fixture, 100-token prompt, 64 greedy): TTFT 215 → 215 ms, decode 206 → 205 ms, generated_text byte-identical. Stack route buffers are sub-noise by design (3 malloc/free per token per layer removed — allocation-free hot path, not a tok/s claim).

Two harness fixes surfaced by this PR's own CI (test-only, engine untouched):

  • test_rope_invfreq now builds with -ffp-contract=off. The in-test reference is a different loop body than the shipping one; on FMA targets (-march=x86-64-v3, contract=fast default) the two bodies contract differently and a*cs-b*sn rounds differently — the pre-change code fails this test too. This is exactly the contraction mechanism you raised for LTO, showing up in the wild: engine-level equivalence holds (verified by dump), but a two-function bit-compare can't rely on contraction luck.
  • test_analysis_cache resolves its temp dir in setUpanalyze_model resolves the path internally, so the spy assertion failed on macOS (/var vs /private/var) and Windows (8.3 short names).

CI green on all four jobs.

@JustVugg

Copy link
Copy Markdown
Owner

Heads-up, and an apology: this went dirty because of us, not you.

#1063 landed on dev — it makes model families registry-owned, so c/coli, c/openai_server.py, c/doctor.py and c/resource_plan.py now all read one descriptor table (c/family_registry.py) instead of each carrying its own branches. That was a deliberate structural change agreed in Discussion #1057, and it rewrote exactly the files your branch touches.

A rebase on current dev should be mechanical — the conflicts will be in the family/dispatch branches that no longer exist, and the replacement is usually "read it from the registry" rather than "add another branch". If a conflict isn't obvious, say so on the PR and I'll work through it with you rather than leave you guessing at the new contract.

Two things worth knowing while you're in there:

  • Unknown model_type is now refused explicitly instead of silently falling through to the GLM engine. If your change relied on that fallback anywhere, it won't behave the same.
  • resource_plan.py now refuses to plan for families without a measured adapter (Kimi, OLMoE, Inkling, V4) rather than returning a plausible-looking zero. That was the bug that hit two model PRs independently in the same day.

Sorry for the churn. Landing the registry before your branches was the right call for the project, but it does mean the cost of it fell on the people with open work.

- colibri.c: cache RoPE inv_freq keyed on (theta,qk) instead of rerunning
  half x powf per position; parallelize the prefill RoPE/KV-write loop on
  CPU-only builds (iterations write distinct KV rows, rope cache is
  _Thread_local; CUDA/Vulkan shadow-shrink keeps GPU builds serial);
  cache the compiled grammar per serve slot when the identical schema
  text is resent (grammar_reset == fresh setup end state)
- deepseek_v4.c: stack routing buffers (<=512 experts, malloc fallback
  beyond); COLI_MODEL_DIRS multi-directory shard lookup via st_init_multi
- resource_plan.py: sidecar cache for analyze_model (atomic tmp+replace
  writes, (size,mtime) signature self-invalidation; the non-serializable
  resolved_family registry object is rebuilt from a cheap resolve_model()
  on cache hits, with only the scan-derived indexer flag persisted)
- download_fp8.py: --dest / $GLM_DEST instead of a hardcoded I:\ path
- Makefile: opt-in LTO=1 (off by default: cross-TU inlining can change
  FP contraction and the token-exactness oracles are validated against
  the default build's regime, JustVugg#1044/JustVugg#1024)
- tests: test_rope_invfreq (15.3M float pairs, 0 bit diffs; built with
  -ffp-contract=off so the reference-vs-engine compare does not depend
  on contraction luck), test_grammar_cache (ASan/UBSan ownership),
  test_analysis_cache (temp dir resolved for macOS/Windows canonical paths)
@zh-Processor

Copy link
Copy Markdown
Author

No apology needed — the registry is the right structure, and the rebase was
a good forcing function. Done: the PR is now a single commit on 1cc0be2
(post-#1063 dev), CI green on all jobs.

One thing turned out to be more than mechanical, exactly along the lines of
your heads-up: analyze_model now returns the live resolved_family
object, which can't be JSON-serialized into the sidecar cache — a naive
keep would have turned the cache write into a TypeError on every call.
The cache now persists only the serializable fields plus the one
scan-derived bit the object carries (_colibri_indexer_present for glm),
and rebuilds resolved_family from a fresh resolve_model() on hits —
cheap, since that's a config.json read. Both new refusal behaviors are
respected: the cache stores nothing family-specific beyond what a fresh
scan computes, and the test fixture gained a real model_type
(glm_moe_dsa) since unknown types are now refused.

The rebased run's Windows leg then caught a real race in my cache writer:
the tmp name was per-process, so concurrent in-process writers shared one
tmp file and one thread's os.replace could rename it from under another
(and on Windows the replace itself fails while a reader holds the
destination open). Fixed with per-thread tmp names plus unlink-on-failure;
the concurrent-writers test now passes 20/20 locally (it was a coin flip
before — that's how the earlier green slipped through). Credit where due:
the test did exactly what it was written to do.

Re-measured the one number whose code path changed (the warm hit now does
the resolve round-trip); everything else is unchanged because the engine
hunks auto-merged byte-identical (verified by hand). Same setup as my
earlier comment — Ryzen 3700X, MinGW gcc 15.2.0, interleaved medians:

benchmark dev this PR verdict
analyze_model, cold (116k-tensor census) 371 ms
analyze_model, warm (cache hit) 4.2 ms ≈88× per invocation
rope_interleave micro (fresh position per call) 1089 ns 457 ns 2.4×; byte-identical (61.5 MB dump, 0 diffs)
grammar compile, 2 identical requests on one slot 2 compiles 1 compile ≤ ~2 ms per repeat; outputs token-identical
GLM prefill (TF, 2047 positions, bench fixture) 75.9 pos/s 76.2 pos/s within noise; no long-prompt win claimed
GLM decode (REPLAY, 64 tok after 2047-token prefill) 33.29 tok/s 33.38 tok/s within noise
V4 tiny, TTFT / decode (64 greedy tokens) 215 ms / 206 ms 215 ms / 205 ms within noise; output byte-identical

Tests on this tree: test_analysis_cache 6/6, test_resource_plan 57/57,
rope differential 15,326,112 pairs / 0 bit diffs, grammar ownership tests,
tiny GLM oracle 32/32, V4 tiny greedy output token-identical to pre-merge.
The only CI red on the final run was the CUDA syntax-check job failing to
download its action with a 429 — a rerun passed; nothing code-related.

@zh-Processor

Copy link
Copy Markdown
Author

Heads-up: the Vulkan (Lavapipe, software) job on the PR-side run
(32087662999) looks stuck — 1.5h in, while the identical job on the same
SHA in my fork's push-triggered run finished in 37 s (success). All other
15 jobs are green on both sides. A cancel + re-run should clear it.

@JustVugg
JustVugg merged commit 57eeb73 into JustVugg:dev Aug 18, 2026
35 of 36 checks passed
@JustVugg

Copy link
Copy Markdown
Owner

Merged — and thank you for a first contribution that reads like it was written by someone who has been here for months: LTO kept opt-in with the cross-TU contraction risk spelled out, a dedicated bit-exactness test for the RoPE cache, and byte-identical discipline throughout. That is exactly the bar this repo tries to hold. (The Lavapipe delay was a hung runner on our side, not your PR — a cancel+rerun fixed it.) Welcome aboard.

@zh-Processor

Copy link
Copy Markdown
Author

Thanks for the thorough review — learned a lot.
Looking forward to contributing more.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants