Skip to content

Faster eval walk and cached gather_mm identity indices - #3920

Open
spokvulcan wants to merge 1 commit into
ml-explore:mainfrom
spokvulcan:perf/eval-hot-path-overhead
Open

Faster eval walk and cached gather_mm identity indices#3920
spokvulcan wants to merge 1 commit into
ml-explore:mainfrom
spokvulcan:perf/eval-hot-path-overhead

Conversation

@spokvulcan

@spokvulcan spokvulcan commented Jul 25, 2026

Copy link
Copy Markdown

LLM decode evaluates a graph of a few thousand nodes per token, and in an M3 Max profile the CPU side of eval was a visible slice of the generation thread. This trims it in two independent places.

The first is eval_impl. Building the degree counts and the tape does several hash map operations per graph edge on every eval, and the unordered_map node allocations dominated that walk. The map is now a small open-addressing table. It is probed by key and never iterated, so the tape order it produces is unchanged. The tape loop also stops redoing the stream, event, and fence bookkeeping for every node while consecutive nodes run on one stream, which is the whole loop in the common single-stream case.

The second is gather_mm and gather_qmm. When no lhs indices are passed they rebuild the identity row indices (an arange plus a reshape) on every call. An MoE layer only passes rhs expert ids, so a 256-expert model rebuilds them a few hundred times per decoded token. The evaluated identity array is now cached per shape. The cache only ever holds evaluated arrays, so every consumer sees an immutable constant leaf. Keys are exact shapes, the cache is bounded and mutex guarded, and tracing bypasses it so function transforms never capture a cached leaf.

Numerics are untouched. Output is token-identical over 18 of 18 and 20 of 20 interleaved A/B pairs on a 4B dense and a 35B MoE model. Measured on the MoE model, decode is +2.0% at short context and +2.8% at 8K, prefill is +3.6% at 8K from the index cache, and peak memory is exactly flat.

Happy to split this into two PRs if the map and the cache are easier to review separately.

@spokvulcan
spokvulcan force-pushed the perf/eval-hot-path-overhead branch from be95ddd to 24ae5db Compare July 31, 2026 18:09
@spokvulcan spokvulcan changed the title Reduce per-eval CPU overhead: flat degree map in eval_impl + gather identity-index cache Faster eval walk and cached gather_mm identity indices Jul 31, 2026
Two independent CPU-side fixes for graphs evaluated at high frequency,
the motivating case being LLM decode at a few thousand nodes per token.

eval_impl now tracks node degrees in a small open-addressing table
instead of an unordered_map. The table is probed by key and never
iterated, so the tape order is unchanged. The tape loop also only
redoes its stream, event, and fence bookkeeping when the stream
actually changes, which is never in the common single-stream case.

gather_mm and gather_qmm no longer rebuild the identity row indices on
every call that has no lhs indices. The evaluated array is cached per
shape, bounded, mutex guarded, and bypassed while tracing.

No numerics change. Token-identical output over interleaved A/B runs
on a 4B dense and a 35B MoE model. On the MoE model decode is
+2.0-2.8%, prefill +3.6% at 8K, and peak memory is flat.
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.

1 participant