Faster eval walk and cached gather_mm identity indices - #3920
Open
spokvulcan wants to merge 1 commit into
Open
Conversation
spokvulcan
force-pushed
the
perf/eval-hot-path-overhead
branch
from
July 31, 2026 18:09
be95ddd to
24ae5db
Compare
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.
spokvulcan
force-pushed
the
perf/eval-hot-path-overhead
branch
from
July 31, 2026 18:10
24ae5db to
08c857f
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.