Attention: fold the sigmoid gate into the causal reduce epilogue, one graph node instead of two - #268
Conversation
Every full-attention call site in TextContext issued causal_softmax_attention
and then ops::sigmoid_mul over the whole output - a second pass, and a second
node in the decode graph, for work the reduce epilogue is already holding in
registers. The benchmark reports the node count itself, so this is measured
rather than asserted: on the append entry with BF16 or INT8 storage the Op's
graph is three nodes in 92 of the 128 cells swept, five in 24 and seven in 12,
and handing the gate to the Op removes exactly one of them, never two. On
qwen3.6-35b-a3b that is ten of the forty text layers, plus the MTP tail once
per draft step - thirteen calls in an mtp3 decode round.
causal_softmax_attention now takes an optional gate. Where the route reaches
the shared BF16/INT8 reducer the multiply is folded into its store; every
other route - FP8, NVFP4, K8V4, the prompt kernel, chunked small-T - applies
the standalone elementwise kernel inside the Op. One contract either way, so
no caller has to know which route it landed on.
The bytes are the bytes the caller's own sigmoid_mul produced, and that is a
rounding question rather than a formality. The standalone kernel reads what
attention stored, so it sees the reduce result already rounded to BF16. The
fused epilogue therefore rounds first, widens back, multiplies in FP32 and
rounds to nearest; keeping the FP32 accumulator would be more accurate and
would move tokens. One case is not a multiplication by one: a masked column
stores an exact zero, and 0 * sigmoid(g) is zero for every finite gate but not
for a NaN one, which the standalone multiply would propagate, so the fused
path does the multiply there too.
causal_softmax_attention_cached is untouched and keeps the standalone
multiply; extending it is a separate decision.
All three forms live in one binary, so the comparison is one process per pass:
--gate standalone is what the model issued, --gate fused hands the gate to the
Op, and --gate standalone again is the identical-baseline null. Six passes,
the arm order rotated through six permutations, 480 cells over both entries,
both geometries, all five storages, B in {1,4} and W in {1,4,8,16}, graph
execution, cold cache, 20 warmups and 300 samples. A pass is dropped when the
two identical binaries disagree by more than 1 %; 396 of 2880 were, and over
the rest the null's p95 is 0.50 % with a worst of 1.00 %. Before that rule it
is 4.05 %, which is what the drop rate is buying. Cells are classified by the
bench's own graph_nodes column, not by an assumption about which routes fold.
band cells median best worst outside
the graph is one node shorter 76 -6.19% -15.13% +0.00% 60 / 76
the node count is unchanged, control 404 +0.00% -4.42% +5.65% 19 / 404
storage geometry cells median median saving
bf16 d256-h16-kv2 16 -4.43% 1.94us
bf16 d256-h24-kv4 20 -4.60% 1.99us
int8 d256-h16-kv2 16 -7.50% 2.00us
int8 d256-h24-kv4 24 -7.93% 2.02us
Which cells fold is not a rule worth stating: 96 cells are on the append entry
with BF16 or INT8 storage at W <= 8 and 76 of them fold, the twenty exceptions
being every d256-h16-kv2 cell at W=8, both storages, and four d256-h24-kv4
BF16 cells. The table reports the node column rather than predicting it.
The saving is a node and it does not grow with the cell: its median is 2.000
us, 48 of the 76 cells lie between 1.76 and 2.05 with the largest mode exactly
at 2.048, and across an eleven-fold range of cell cost - 11.49 to 123.97 us -
the correlation between saving and cost is r = -0.09. It is not constant
either: fourteen cells save under 0.25 us and five save nothing measurable,
and this sweep does not say why.
Eager timing, where a launch is not amortised by a graph, puts 74 of those 76
cells - two lose every pass to the null rule - at a median of -14.73 % and
4.05 us, every one faster and none saving zero. Eager is the noisier
instrument here: 617 of 2880 passes dropped, 21.4 %.
The cells the product actually issues are narrower than the band. The artifact
is d256-h16-kv2: query_norm is [256], the output projection is 2048x4096 so
sixteen query heads, and the fused QKGV projection is 9216 rows, leaving two
KV heads. Decode at mtp3 issues W=4 for the target block and W=1 per draft
step, at B=1, BF16, append. All eight such cells fold, at a median of -2.95 %
and a median saving of 1.03 us on cells costing 15.58 to 40.16 us. W=8, which
DFlash7 would issue on this geometry, does not fold at all.
Output is unchanged: NINFER_OP_REPORT_STATS=1 at %.17g over the whole test
suite gives 11900 cases shared between the arms, all byte-identical, none
differing in any field. The 335 cases the candidate adds are the gated oracle
checks the test gains.
Each case in causal_cache.cpp now runs a second time with a gate and is judged
against a third, ungated run multiplied by a host sigmoid - an independent path
from the inputs rather than a second run of the Op. Strength control: keeping
the FP32 accumulator instead of rounding before the multiply, which is
strictly more accurate and not what the standalone multiply produces, fails 65
exact-parity checks; restoring it passes.
No end-to-end number. On this machine the identical-baseline arm reads +0.55 %
median on decode over a -0.44 to +0.89 % band, against a candidate reading
+1.10 %, so the stand does not resolve the change. What the Op measurement
implies, as a prediction rather than a measurement: thirteen calls per round at
the 1.03 us the production-shaped cells save is 13.4 us against a round the
same run reports as 4891 us, or 0.27 %.
ctest 120/120 on both arms.
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2961dac61c
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| * An optional gate asks the Op to finish with out *= sigmoid(gate). When present it is a contiguous | ||
| * BF16 tensor shaped exactly like out and disjoint from every other operand. Where the route allows | ||
| * it the multiply is folded into the reduce epilogue; every other route applies the standalone | ||
| * elementwise kernel inside the Op. Either way the result is bit-identical to calling sigmoid_mul | ||
| * on the ungated output, so no caller has to know which route it landed on. |
There was a problem hiding this comment.
When valid_columns masks a tail and the corresponding BF16 gate is NaN, both the fused reducer and the fallback sigmoid_mul produce NaN from 0 * sigmoid(NaN), but the preceding contract still promises that every tail column is exact BF16 zero. Because the new gate contract does not require finite values, these guarantees conflict for a valid input; either restrict gates to finite values or qualify the zero guarantee and cover the chosen behavior with an oracle case.
AGENTS.md reference: AGENTS.md:L89-L94
Useful? React with 👍 / 👎.
A Q8 threshold sweep contained, by accident, a band above T=57 where the threshold selects nothing and both columns run identical code. They did not read the same: <= 0.25 % apart on one shape and about 1 % on another, systematically favouring whichever column ran first. That is the methods own bias, priced inside the measurement without a third arm. Found by llm-5090-3b. Named the pattern because two of our packages already use it - Neroued#264 sweeps whole-tile widths that take the same route on both arms (+0.05 %, and chunk 1536 at -0.05 % beside a +2.37 % candidate), Neroued#268 carries 404 cells whose node count is unchanged (+0.00 %). Neroued#222 cannot have one, since its arms differ at every width, so its burden falls entirely on the null arm - which is why its worst kept pass is quoted as carefully as it is.
…reduce epilogue
…reduce epilogue
|
The fusion direction makes sense, but this PR needs the following changes before merging:
|
3 participants
Problem and scope
Related Issue: #267
Every full-attention call site in
TextContextissuescausal_softmax_attentionand thenops::sigmoid_mulover the whole output — a second pass, and a second node in the decode graph, forwork the reduce epilogue is already holding in registers. The benchmark reports the node count
itself, so this is measured rather than asserted: on the
appendentry with BF16 or INT8 storage theOp's graph is three nodes in 92 of the 128 cells swept (five in 24, seven in 12), and handing the
gate to the Op removes exactly one of them — never two.
On
qwen3.6-35b-a3bthat is ten of the forty text layers, plus the MTP tail once per draft step —thirteen calls in an mtp3 decode round.
Scope is
causal_softmax_attentionand the two call sites that feed it.causal_softmax_attention_cachedis deliberately untouched, so the thirdsigmoid_mulintext.cppstays where it is.Implementation
causal_softmax_attentiontakes an optional gate. Where the route reaches the shared BF16/INT8reducer the multiply is folded into its store; every other route — FP8, NVFP4, K8V4, the prompt
kernel, chunked small-T — applies the standalone elementwise kernel inside the Op. One contract
either way, so no caller has to know which route it landed on, and the bench's
--gateswitchselects between the two forms for measurement.
The bytes are the bytes the caller's own
sigmoid_mulproduced, and that is a rounding questionrather than a formality. The standalone kernel reads what attention stored, so it sees the reduce
result already rounded to BF16. The fused epilogue therefore rounds first, widens back, multiplies in
FP32 and rounds to nearest. Keeping the FP32 accumulator would be more accurate and would move
tokens.
One case is not a multiplication by one: a masked column stores an exact zero, and
0 × sigmoid(g)is zero for every finite gate but not for a NaN one, which the standalone multiply would propagate.
The fused path does the multiply there too.
Verification
RTX 5090, driver 616.64, CUDA 13.1.115,
CMAKE_CUDA_ARCHITECTURES=120a, Release, basedc58675f.Both arms' CMake caches are identical once the tree paths are normalised.
What moves
All three forms live in one binary, so the whole comparison is one process per pass:
--gate standaloneis what the model issues today,--gate fusedhands the gate to the Op, and--gate standaloneagain is the identical-baseline null. Six passes, the arm order rotated throughsix permutations, 480 cells over both entries, both geometries, all five storages,
Bin {1,4} andWin {1,4,8,16},--execution graph --cache cold --warmup 20 --repeat 300.A pass is dropped for a cell when the two identical binaries disagree by more than 1 %; 396 of 2880
were. Over the passes that survive that rule the null's |p95| is 0.50 % and its worst is
1.00 %; over all 2880 before the rule it is 4.05 %, which is what the drop rate is buying.
Cells are classified by the bench's own
graph_nodescolumn, not by an assumption about which routesfold.
Which cells fold is not a rule I can state. 96 cells are on the
appendentry with BF16 or INT8storage at
W ≤ 8, and 76 of them fold. The twenty that do not are everyd256-h16-kv2cell atW=8— both storages, all sixteen — and fourd256-h24-kv4BF16 cells. The table reports what thenode column says rather than predicting it.
The saving is a node, and it does not grow with the cell. Its median is 2.000 µs; 48 of the 76
cells sit between 1.76 and 2.05 µs, with the largest mode exactly at 2.048. Across an eleven-fold
range of cell cost — the folded cells run 11.49 to 123.97 µs — the correlation between saving and
cost is r = −0.09. But it is not constant either: fourteen cells save less than 0.25 µs and five
save nothing measurable, and this sweep does not say why.
Eager timing, where a launch is not amortised by a graph, puts 74 of those 76 cells (two lose
every pass to the null rule) at a median of −14.73 % and 4.05 µs, with every one of them faster
and none saving zero. Eager is the noisier instrument here: 617 of 2880 passes dropped, 21.4 %.
The cells the product actually issues
The artifact is
d256-h16-kv2— itsquery_normis[256], its output projection is2048×4096sothere are sixteen query heads, and its fused QKGV projection is 9216 rows, which leaves two KV heads.
Decode at mtp3 issues
W=4for the target block andW=1per draft step, atB=1, BF16,append.All eight such cells in the sweep fold. Their median is −2.95 % and their median saving
1.03 µs, on cells costing 15.58 to 40.16 µs — below the band median, because the band includes
larger and more favourable cells.
W=8, which DFlash7 would issue on this geometry, does not fold at all.Output
NINFER_OP_REPORT_STATS=1at%.17gover the whole test suite: 11900 cases shared between thearms, all byte-identical, none differing in any field. The candidate adds 335 cases, which are the
gated oracle checks the test gains.
Each case in
causal_cache.cppnow runs a second time with a gate and is judged against a third,ungated run multiplied by a host sigmoid — an independent path from the inputs rather than a second
run of the Op, so "the Op repeats itself" and "the gate is exact" stay separate questions.
Strength control. The rounding before the multiply is the one subtlety in the epilogue, so it has
to be the thing the test would catch. Keeping the FP32 accumulator instead — strictly more accurate,
and not what the standalone multiply produces — fails 65 exact-parity checks; restoring it passes.
ctest120/120 on both arms, of which seven are skipped on both.clang-formatreports noreplacements on any of the eight files, and none gains a line over the 100-column limit that master
did not already have.
End to end: not resolved on this machine
ninfer_bench -pg 8192,512 --prefill-chunk 4096 --max-ctx 9216 --spec mtp --draft-tokens 3 --lm-head-draft -r 10 --warmup 3, six passes with the arm order rotated, the third arm the masterbinary again under a second label:
The null is half the candidate and its band overlaps it, so this stand does not resolve the change
and I am not claiming a decode number. At
-r 2it was worse: null +1.81 % median with a +5.41 %worst pass.
What the Op measurement implies, as a prediction rather than a measurement: thirteen calls per round
at the 1.03 µs the production-shaped cells actually save is 13.4 µs against a round the same run
reports as 4891 µs, or 0.27 %. The run's counters are self-consistent, which is why that round
time is usable: 162 rounds per repetition for 512 tokens is 3.160 tokens per round, and the reported
acceptance rate of 0.718 predicts 1 + 3 × 0.718 = 3.154.
Limits
appendentry, and not every such cell — the exceptions arelisted above and were measured, not derived. FP8, NVFP4, K8V4 and the
cachedentry never fold.causal_softmax_attention_cachedis not extended.measurable. At 9 to 267 µs per cell with a 13.8 % drop rate this bench cannot say more about either
group than that they are inside its own scatter.
the bench is identical across all 2880 passes, and the derived GB/s column moves only because it is
computed from the time.