Remove the empty graph nodes of the MTP draft phase: +0.33% decode on Qwen3.6-35B-A3B - #226
Remove the empty graph nodes of the MTP draft phase: +0.33% decode on Qwen3.6-35B-A3B#226MichaelDementii wants to merge 5 commits into
Conversation
The draft phase of an MTP decode round spends nodes that carry no work of their
own. Three of them, in one pass because separately none is above what an
end-to-end benchmark on this hardware can resolve:
- ops::argmax opened its atomic contest with a 4-byte cudaMemsetAsync per
slice; an initializer kernel writes the same bytes as a kernel node;
- the MTP stem normalized the embedding and the hidden state only to lay them
side by side, and the MTP tail added the residual only for the next norm to
read it back: two new Ops write each result into place directly;
- the autoregressive draft chain copied the new hidden state over the old one
after every step, where alternating between two buffers costs nothing.
Both new Ops carry a *_supported() predicate and the three separate calls remain
as the fallback, so an unregistered shape keeps working.
Over one traced run - 8192-token prefill, 64 generated tokens, two repetitions:
173 memset nodes become kernel nodes, 397 rmsnorm launches disappear along with
134 pack and 129 residual-add launches, and 82 of 100 device-to-device copies go
away. End to end on Qwen3.6-35B-A3B, twelve mirrored passes, decode +0.334%
median with ten of twelve positive; the two negative passes and one outlier of
the zero-control arm coincide with a build running on the neighbouring container,
which shares the host CPU quota.
Output is byte-identical, ctest 114/114.
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 |
![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: 697606edea
ℹ️ 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".
| ops::rmsnorm(t_embedding, t_weight_e, kStemEps, true, t_norm_e, nullptr); | ||
| ops::rmsnorm(t_hidden, t_weight_h, kStemEps, true, t_norm_h, nullptr); | ||
| ops::mtp_pack_fc_input(t_norm_e, t_norm_h, t_reference, nullptr); |
There was a problem hiding this comment.
These production rmsnorm and pack calls are used as the sole reference, and the residual test below likewise composes residual_add with rmsnorm; consequently, an arithmetic, reduction, or rounding defect shared with the copied fused implementation can pass every new comparison. Add a test-owned naive FP32/FP64 oracle for each complete fused formula, including the BF16 residual state boundary, and retain implementation-parity checks only as supplementary evidence.
AGENTS.md reference: AGENTS.md:L83-L88
Useful? React with 👍 / 👎.
Every check on the fused Ops used the Ops they replace as the reference, so a defect the two share passed all of them. Add the naive FP64 oracle and the criterion tests/ops/test_rmsnorm.cpp judges ops::rmsnorm by, evaluated from the represented BF16 inputs; the residual half gets an exact oracle, because a BF16 sum of two BF16 values is exact in double. The parity checks stay: they are the bit-exactness these Ops claim. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
Taken, in
The parity checks stay where they were. They are not weak evidence, they are the claim these Ops Strength control, and it answers your point directly. The same defect placed in both the fused With the kernels restored, @codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2a9dc0faa7
ℹ️ 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".
| if (d >= 512 && d <= 3072 && d % 512 == 0) { return MtpRowRoute::Cta256x6; } | ||
| if (d > 3072 && d <= 8192 && d % 1024 == 0) { return MtpRowRoute::Cta512x8; } |
There was a problem hiding this comment.
Restrict this dispatch to the registered MTP widths, 2048 and 5120. As written, widths such as 512, 1536, 4096, and 8192 are treated as supported—and the test explicitly canonizes 1536—even though no product target requires them. This creates placeholder execution paths coupled to the current generic RMSNorm ladder; unsupported future widths should remain on the composed fallback until a target explicitly requires and qualifies them.
AGENTS.md reference: AGENTS.md:L57-L61
Useful? React with 👍 / 👎.
The route mirrored the whole generic RMSNorm ladder, so 512, 1536, 4096 and 8192 were admitted although no target asks for them -- a path kept alive for a model that does not exist. Admit 2048 and 5120, drop the instantiation nothing reaches any more, and pin both directions in the test. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
Taken, in The same predicate serves the residual-and-norm tail, so that route narrowed with it. The test Why the ladder was mirrored in the first place is unchanged and still in the comment: an admitted @codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 03e67bc369
ℹ️ 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".
| void mtp_norm_pack_fc_input(const Tensor& embedding, const Tensor& embedding_weight, | ||
| const Tensor& hidden, const Tensor& hidden_weight, Tensor& out, | ||
| float eps, cudaStream_t stream) { | ||
| require_stem_operands(embedding, embedding_weight, hidden, hidden_weight, out); |
There was a problem hiding this comment.
When either fused Op is called directly with eps <= 0 or a non-finite value, these entry points launch and can produce NaNs instead of rejecting the input as the rmsnorm() operation they replace does (src/ops/wrapper/rmsnorm.cpp validates this at lines 56–58). Apply the same positive-and-finite check in both mtp_norm_pack_fc_input() and mtp_residual_norm() so their advertised composition preserves the normalization contract.
Useful? React with 👍 / 👎.
Both fused entry points stand in for a composition containing ops::rmsnorm, which refuses such an eps; they launched instead and came back with NaN. Apply the same check and pin it in the test, on both Ops and on both registered widths. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
Right, and fixed in Pinned in the test rather than only written: @codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 59f80df250
ℹ️ 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".
| * Numeric: | ||
| * Identical to rmsnorm() applied to each half followed by mtp_pack_fc_input(): the same block | ||
| * width, pair decomposition, reduction and epilogue produce the same bytes. Returns false from | ||
| * mtp_norm_pack_fc_input_supported() for shapes it does not cover, and the caller keeps the | ||
| * three-Op path there. |
There was a problem hiding this comment.
The Numeric contract makes the current RMSNorm block width, pair decomposition, reduction order, and byte-for-byte parity semantic requirements. This conflicts with docs/maintainer/op-development.md §3, which explicitly excludes reduction association, launch decomposition, and bitwise equality for floating-point Ops; future RMSNorm tuning would now either silently violate this contract or require the fused kernels and exact-parity tests to change in lockstep. Specify the FP64-oracle criterion and observable BF16 boundaries here instead, keeping implementation parity only as supplementary regression evidence.
AGENTS.md reference: AGENTS.md:L83-L89
Useful? React with 👍 / 👎.
| if (embedding_weight.numel() != rows || hidden_weight.numel() != rows) { | ||
| throw std::invalid_argument("mtp_norm_pack_fc_input: weight length must be D"); |
There was a problem hiding this comment.
When a caller supplies a contiguous weight with numel() == D but a shape such as [1,D], this check accepts it and both the support predicate and fused execution proceed, whereas the documented [D] contract and the composed ops::rmsnorm route require ne[0] == D with all remaining dimensions equal to one. Validate the complete weight shape for both stem weights and the residual-norm weight so these fused entry points preserve the operation they replace.
AGENTS.md reference: AGENTS.md:L83-L89
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 59f80df250
ℹ️ 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".
| * Numeric: | ||
| * Identical to rmsnorm() applied to each half followed by mtp_pack_fc_input(): the same block | ||
| * width, pair decomposition, reduction and epilogue produce the same bytes. Returns false from | ||
| * mtp_norm_pack_fc_input_supported() for shapes it does not cover, and the caller keeps the | ||
| * three-Op path there. |
There was a problem hiding this comment.
Replace this implementation-parity clause with the complete RMSNorm formula, the positive-and-finite eps domain now enforced by the wrappers, the observable BF16/state boundaries, and the numerical acceptance criterion. Making the current block width, pair decomposition, reduction tree, and byte parity contractual prevents an independent implementation and turns a valid future rmsnorm kernel tuning into a contract break even when its mathematical output remains conformant.
AGENTS.md reference: AGENTS.md:L83-L89
Useful? React with 👍 / 👎.
| if (embedding_weight.numel() != rows || hidden_weight.numel() != rows) { | ||
| throw std::invalid_argument("mtp_norm_pack_fc_input: weight length must be D"); | ||
| } |
There was a problem hiding this comment.
When a caller supplies a contiguous BF16 weight shaped, for example, {D/2, 2}, this numel() check accepts it and the support predicate admits the fused route even though the contract requires [D] and ops::rmsnorm rejects the same operand. The residual helper repeats this check for its weight; validate ne[0] == D and ne[1..3] == 1 in both helpers so malformed tensor views cannot silently cross the public Op boundary.
Useful? React with 👍 / 👎.
…shape The Numeric fields made the current RMSNorm block width, pair decomposition and byte parity contractual, which docs/maintainer/op-development.md forbids: a valid retuning of ops::rmsnorm would have broken the contract without changing what the Op computes. State the formula, the eps domain, the observable BF16 roundings and the FP64-oracle criterion instead; parity stays where it belongs, in the tests. The weight check counted elements, so [1,D] and [D/2,2] reached the fused route although ops::rmsnorm refuses them. Check the whole shape in both Ops. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
Both taken, in The contract. You are right and the rule is explicit — The weight shape. Also right:
@codex review |
|
Codex Review: Didn't find any major issues. Breezy! 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
Third of three node-count reductions measured in one session; #222 and #225 are the others. All
three are independent mechanisms but all three touch
text_context_impl.h, so whichever lands firstleaves the other two a trivial rebase.
Why one PR and not three. Separately these are +0.090 %, +0.095 % and +0.040 % of decode against
a zero-control band of ±0.15 % on this hardware — none is resolvable alone. Together they are.
Change
ops::argmaxopened its atomic contest with a 4-bytecudaMemsetAsyncper sliceBoth new Ops carry a
*_supported()predicate with the three separate calls as the fallback, so anunregistered shape keeps working.
Nodes
One traced run — 8192-token prefill, 64 generated tokens, two repetitions,
nsys --cuda-graph-trace=node:[CUDA memset]rmsnorm_cta_bf16x2mtp_pack_fc_inputresidual_add_bf16x8memcpy Device-to-DeviceEnd 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. Twelve passes, arm ordermirrored, third arm = master rebuilt under a second label as a zero control.
Ten of the twelve sit in a tight band of +0.27…+0.40 %. The two that do not, and the one outlier of
the control, coincide with a build running on the neighbouring container, which shares the host CPU
quota — so separation is not full and I am not claiming it. Prefill is unchanged: −0.009 %
median against a control of −0.029 %.
Correctness
Output byte-identical on the same prompt, greedy, int8 KV, mtp3.
ctest114/114. The stem Op testchecks the fused result against the three-Op path it replaces, in the same process on the same
inputs.
Scope: 10 files, +756 / −18, of which 288 are the new Op test.
RTX 5090
sm_120a, CUDA 13.1, Release,qwen3_6_35b_a3b.ninfer.🤖 Generated with Claude Code