fix(compile-dag): ComputeGraphExecutor replays permute with its recorded axes#803
Merged
Merged
Conversation
…ded axes The builtin dispatch grouped 'permute' with 'transpose'/'transpose2d' and replayed all three as ops.transpose — a last-two-dims swap that drops the recorded axes and is only coincidentally correct for rank-2. Any rank-3 permutation in a traced graph (e.g. MHA's heads/sequence swap [1, 0, 2] on multi-token encoder or prefill forwards) silently produced the wrong layout; single-token decode never hits the permute path, which is why generation workloads didn't surface it. permute now replays with its recorded axes, parsed with the same IntArray | List<Int> convention permuteBackward and the StableHLO converter already use, and falls back to the legacy transpose behavior when no usable axes were recorded (older traces). Found by the BERT DSL-path encoder migration in SKaiNET-transformers (DIRECT-vs-OPTIMIZED trace parity bisection); that repo carries a temporary axes-aware FusedOpHandler override to be dropped once this ships. Regression-tested: permuteReplaysRecordedAxes, permuteWithoutRecordedAxesFallsBackToTranspose. 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.
What
ComputeGraphExecutor's builtin dispatch groupedpermutewithtranspose/transpose2dand replayed all three asops.transpose(inputs[0])— a last-two-dims swap that silently drops the recordedaxesand is only coincidentally correct for rank-2 tensors.Why it went unnoticed
Any rank-3 permutation in a traced graph — e.g. multi-head attention's heads/sequence swap
[1, 0, 2]on a multi-token (encoder or prefill) forward — produced the wrong layout with no error. Single-token decode never hits the permute path (swapSeqHeadDimstakes thed0==1/d1==1reshape shortcut), so generation workloads never surfaced it. It was found by DIRECT-vs-OPTIMIZED trace-parity bisection during the BERT DSL-path encoder migration in SKaiNET-transformers.Fix
permutenow replays with its recordedaxes, parsed with the sameIntArray | List<Int>convention thatpermuteBackward(DefaultExecutionTape) and the StableHLOLinalgOperationsConverteralready use — the executor was the one inconsistent consumer. Traces without usable axes (older recordings) keep the legacy transpose behavior, so existing graphs are unaffected.Tests
permuteReplaysRecordedAxes— rank-3 input, recorded axes[1, 0, 2]; assertsops.permuteis invoked with exactly those axes andops.transposeis not called.permuteWithoutRecordedAxesFallsBackToTranspose— legacy-trace fallback pinned.:skainet-compile:skainet-compile-dag:jvmTestsuite andapiCheckgreen.Downstream note: SKaiNET-transformers currently carries a temporary axes-aware
"permute"FusedOpHandleroverride inLLMFusedOpHandlers(identical logic; the handler registry precedes the builtin dispatch, so the overlap is harmless). It will be removed once a release containing this fix is consumed there.🤖 Generated with Claude Code