cpu decode: route f16 SDPA through FlashSdpa, fuse GQA broadcast through cache-read cast/rope#2453
Open
czoli1976 wants to merge 1 commit into
Open
cpu decode: route f16 SDPA through FlashSdpa, fuse GQA broadcast through cache-read cast/rope#2453czoli1976 wants to merge 1 commit into
czoli1976 wants to merge 1 commit into
Conversation
…A broadcast fusion look through cache-read Cast/ApplyRope ops. The flash kernel already accepts any float input (accumulating in f32) and handles grouped heads natively, but the codegen gate only admitted f32 accumulators and the fusion matcher required the KV cache directly under the repeat-interleave chain, so torch-style f16 GQA exports paid a materialized head broadcast plus full-cache repacking on every decode step. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
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
1 participant
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.
Routes f16-accumulator SDPA through
FlashSdpaOpon CPU and lets the GQA broadcast fusion look through cache-readCast/ApplyRopeops, so torch-style f16 GQA exports stop paying a materialized head broadcast plus full-KV-cache repacking on every decode step.Profiling an OpenELM-270M q40ef16 decode step at P=1000 showed the real GEMV work (~14 ms) buried under per-token recompute over unchanged data: 13.1 ms
OptMatMulPackre-packing the KV cache, 8.2 msMultiBroadcastTomaterializing the repeat-interleave, 4.1 msApplyRopere-roping the whole window. Two routing gaps kept the existing flash path out of reach: the codegen gate only admitted f32 accumulators (the kernel itself takes any float and accumulates in f32), and the fusion matcher required theDynKeyValueCachedirectly under the repeat chain, which torch exports interpose with cache-read casts and rope.Measured on Apple M-series, 1 thread, interleaved A/B against main via
llm_decode_bench(medians):llm-bench: PP512 262 → 338 tok/s (+29%). One trade-off to be aware of: TG128 is −3% (75.1 → 72.7 tok/s) — at very shallow context the decomposed path still competes, andTRACT_FLASH_SDPA_MIN_SEQ_LENcan't gate it because kv_len is symbolic at codegen time. A runtime-side dispatch could recover it if you think it's worth the complexity.Tests: f16-eval-vs-f32 parity on the op, and a model-level test replicating the torch GQA idiom (cache → cast → rope → repeat-interleave → SDPA) asserting the broadcasts fuse away and the fused model matches the unfused reference through real cache state. transformers/onnx-core/core/unit/linalg suites green.
Follow-ups I'd like to explore next, in order: folding the f16→f32 cast per-block into the flash streaming loop, folding rope-at-read into the kernel, and letting
InPlaceKvSdpamatch through the same cast/rope chain (which also removes the concat copy).🍍
🤖 Generated with Claude Code