Fold the sigmoid gate into the causal-attention reduce epilogue: operator -8.3% where it folds, +0.31% decode - #225
Conversation
The three call sites in TextContext all issued causal_softmax_attention and then ops::sigmoid_mul over the whole output, which is one more graph node per full-attention layer and per MTP tail. 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, and the bytes are the bytes the caller's own sigmoid_mul produced: the reduce result is rounded to BF16 first, widened back and multiplied in FP32, because the standalone kernel reads what attention stored rather than the accumulator. The operator bench gains --gate off|standalone|fused. Over the 65 cells where the fold applies, kernel time is -8.33% median with 58 of 65 improved; over the 255 cells where the route does not fold, the median is exactly zero, which is the control that the standalone path is untouched. End to end on Qwen3.6-35B-A3B, ten mirrored passes with a zero-control arm: decode +0.311% median, all ten passes positive, fully separated from a control spanning -0.097 to +0.088%. Prefill is unchanged within that same control. The Op test now runs each case a second time with a gate and judges it against a third, ungated run, which separates "the Op repeats itself" from "the gate is exact" instead of conflating the two. Dropping the rounding before the multiply makes that check fail, so it is not vacuous. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
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. |
|
Codex Review: Something went wrong. Try again later by commenting “@codex review”. ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
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". |
|
@codex review |
![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: 9fe228a391
ℹ️ 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".
| failures += | ||
| verify_exact((label + " ungated repeat").c_str(), | ||
| copy_from_guarded<std::uint16_t>(dexpected, output.size()), output); | ||
| ops::sigmoid_mul(tgate, texpected, device.stream); |
There was a problem hiding this comment.
The expected result is produced by another call to causal_softmax_attention followed by the production sigmoid_mul, so this verifies only parity between two production paths. A shared error in the attention result, sigmoid computation, or BF16 boundary can therefore pass this check; derive the gated expectation from the existing FP64 attention oracle and an independent host sigmoid/cast instead.
AGENTS.md reference: AGENTS.md:L83-L88
Useful? React with 👍 / 👎.
| Execution execution = Execution::Graph; | ||
| CacheMode cache = CacheMode::Cold; | ||
| PageMapping mapping = PageMapping::Identity; | ||
| GateMode gate = GateMode::Off; |
There was a problem hiding this comment.
When --gate standalone or --gate fused is used with --csv-out, the selected mode changes the measured workload but is not stored in Result or emitted by write_csv; the resulting rows are indistinguishable from --gate off and from each other once separated from the console header. Add the gate mode to each result and its CSV identity so saved measurements remain interpretable.
AGENTS.md reference: AGENTS.md:L109-L113
Useful? React with 👍 / 👎.
…d record the gate mode The gated checks compared production against production, so a defect shared by the attention result, the sigmoid or the BF16 boundary would have passed them. Judge the gated output against the same FP64 attention oracle the ungated output is judged by, multiplied by a host sigmoid, with the criterion widened by the one extra BF16 rounding the gated route performs. The bench now carries the gate mode each row was measured in -- the effective one, so a cached row reads standalone -- on the console line and in the CSV. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
Both taken, in P1, the oracle. You are right that the two gated checks compared production against The criterion is the ungated one widened by exactly the second BF16 rounding the gated route adds P2, the CSV. Each row now carries the gate mode it was measured in, on the console line and in
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2b7fb9e76d
ℹ️ 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".
| Case data(geometry, storage, width, rows.contexts, rows.valid_columns, rows.table_rows, | ||
| options.mapping); | ||
| options.mapping, options.gate); |
There was a problem hiding this comment.
Fresh evidence after the follow-up is that --profile now passes options.gate into the measured Case, but profile() still prints a PROFILE header without any gate field, and the new gate=... console line is below this branch's early return. Consequently, profiles captured with --gate off, standalone, or fused remain indistinguishable from their emitted metadata; for cached --gate fused, the header should also report the effective standalone mode, as the CSV path does.
AGENTS.md reference: AGENTS.md:L109-L113
Useful? React with 👍 / 👎.
The --profile path returns before the console line that named the mode, so profiles captured with --gate off, standalone and fused carried no way to tell them apart. The PROFILE header now names the effective mode, through the same helper the CSV and the result line use. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
Right — the
@codex review |
|
Codex Review: Didn't find any major issues. More of your lovely PRs please. Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
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". |
2 participants
All three call sites in
TextContextissuecausal_softmax_attentionand thenops::sigmoid_mulover the whole output — one more graph node per full-attention layer and per MTP tail.
Change
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.
The bytes are the bytes the caller's own
sigmoid_mulproduced. That is a rounding question, not aformality: 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.
causal_softmax_attention_cachedis untouched and keeps the standalone multiply; extending it is aseparate decision.
Operator
The bench gains
--gate off|standalone|fused. Across 320 append cells:standalonePer storage and geometry: bf16 16/2 −8.33 %, bf16 24/4 −6.67 %, int8 16/2 −8.33 %, int8 24/4
−8.78 %. The 255 unfolded cells are the control that the standalone path is untouched. Timings
quantise to 2.048 µs at these widths, so read the percentages as coarse.
End to end
ninfer_bench -pg 8192,512 --prefill-chunk 4096 --max-ctx 9216 --spec mtp --draft-tokens 3 --lm-head-draft -r 2 --warmup 1, metricdecode_output_tok_s_mean. Ten passes, arm order mirroredbetween passes, third arm = master rebuilt under a second label as a zero control.
The worst pass of the change is above the best pass of the control.
Prefill is not claimed: −0.06 % median against a zero control of +0.06 % on the same runs.
Correctness
Each case in
causal_cache.cppnow runs a second time with a gate and is judged against a third,ungated run — which separates "the Op repeats itself" from "the gate is exact" instead of conflating
the two. Both checks are exact, on every storage, both geometries, batched and single, graph and
eager.
Strength control: dropping the rounding before the multiply — the one subtlety in the epilogue —
makes that check fail. Without it "the test passes" would say nothing.
ctest114/114 on both sides;clang-formatclean on every touched file.Scope: 8 files, +224 / −99, of which the bench switch and the test are 129 lines.
RTX 5090
sm_120a, CUDA 13.1, Release,qwen3_6_35b_a3b.ninfer.🤖 Generated with Claude Code