From e27ae12478fd1dfb7a2c9e9deea3d18aaccbce2b Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Thu, 9 Jul 2026 17:16:50 +0200 Subject: [PATCH] fix(compile-dag): ComputeGraphExecutor replays permute with its recorded axes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 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 --- CHANGELOG.md | 13 ++++ .../lang/graph/exec/ComputeGraphExecutor.kt | 20 +++++- .../compile/graph/ComputeGraphExecutorTest.kt | 64 ++++++++++++++++++- 3 files changed, 94 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 39e05341..7bb2df0e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,19 @@ ## [Unreleased] +### Fixed + +- **`ComputeGraphExecutor` replayed `permute` as a plain transpose, dropping the recorded axes.** + The builtin dispatch grouped `permute` with `transpose`/`transpose2d`, so a traced + `permute(t, axes)` executed as a last-two-dims swap — only coincidentally correct for rank-2. + Any rank-3 permutation (e.g. multi-head attention's heads/sequence swap `[1, 0, 2]` in a + full-sequence encoder or prefill trace) silently produced the wrong layout; single-token decode + paths never hit it, which is why generation workloads didn't surface the bug. `permute` now + replays with its recorded `axes` (the same `List` convention `permuteBackward` and the + StableHLO converter already parse), falling back to the legacy transpose behavior for older + traces without axes. Surfaced by the BERT DSL-path encoder work in SKaiNET-transformers, which + carries a temporary `LLMFusedOpHandlers` override to be removed once this fix ships. + ## [0.35.0] - 2026-07-08 ### Added diff --git a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/exec/ComputeGraphExecutor.kt b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/exec/ComputeGraphExecutor.kt index 1d544476..31d3253b 100644 --- a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/exec/ComputeGraphExecutor.kt +++ b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/exec/ComputeGraphExecutor.kt @@ -192,7 +192,25 @@ public class ComputeGraphExecutor( } // Shape ops - "transpose", "permute", "transpose2d" -> listOf(ops.transpose(inputs[0])) + "transpose", "transpose2d" -> listOf(ops.transpose(inputs[0])) + "permute" -> { + // The trace records axes as a List (axes.toList()) — same + // convention permuteBackward parses. Replaying permute as a + // plain last-two-dims transpose is only correct for rank-2; + // e.g. attention's heads/sequence swap [1, 0, 2] on rank-3 + // silently produced the wrong layout. + val axes = when (val a = params["axes"]) { + is IntArray -> a + is List<*> -> IntArray(a.size) { (a[it] as Number).toInt() } + else -> null + } + if (axes != null && axes.size == inputs[0].rank) { + listOf(ops.permute(inputs[0], axes)) + } else { + // No usable axes recorded (older traces) — legacy behavior. + listOf(ops.transpose(inputs[0])) + } + } "reshape", "view" -> { // Shape can come from different parameter keys depending on trace source val targetShape = (params["shape"] as? List) diff --git a/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/compile/graph/ComputeGraphExecutorTest.kt b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/compile/graph/ComputeGraphExecutorTest.kt index 6ac6a282..b89978ef 100644 --- a/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/compile/graph/ComputeGraphExecutorTest.kt +++ b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/compile/graph/ComputeGraphExecutorTest.kt @@ -14,6 +14,7 @@ import sk.ainet.lang.tensor.ops.TensorSpec import sk.ainet.lang.types.DType import sk.ainet.lang.types.FP32 import kotlin.test.Test +import kotlin.test.assertContentEquals import kotlin.test.assertEquals import kotlin.test.assertNotNull import kotlin.test.assertTrue @@ -119,6 +120,62 @@ class ComputeGraphExecutorTest { assertNotNull(results["mm"]) } + @Test + fun permuteReplaysRecordedAxes() { + // Regression: permute used to be dispatched as a plain last-two-dims + // transpose, silently dropping the recorded axes — wrong for any + // rank-3 permutation (e.g. attention's heads/sequence swap [1, 0, 2]). + var replayedAxes: IntArray? = null + var transposeCalled = false + val ops = object : TensorOps by TestTensorOps() { + override fun permute(tensor: Tensor, axes: IntArray): Tensor { + replayedAxes = axes + return tensor + } + + override fun transpose(tensor: Tensor): Tensor { + transposeCalled = true + return tensor + } + } + + val graph = DefaultComputeGraph() + val input = graph.addNode(inputNode("input")) + // The tape records axes as a List (axes.toList()) + val perm = graph.addNode(opNode("perm", "permute", mapOf("axes" to listOf(1, 0, 2)))) + graph.addEdge(GraphEdge("e1", input, perm, tensorSpec = spec())) + + val executor = ComputeGraphExecutor(graph, ops) + val rank3 = TestTensor(floatArrayOf(1f, 2f, 3f, 4f, 5f, 6f), shape = intArrayOf(2, 1, 3)) + val results = executor.execute(mapOf("input" to rank3)) + + assertNotNull(results["perm"]) + assertContentEquals(intArrayOf(1, 0, 2), replayedAxes, "permute must replay with its recorded axes") + assertTrue(!transposeCalled, "permute must not degrade to transpose when axes are recorded") + } + + @Test + fun permuteWithoutRecordedAxesFallsBackToTranspose() { + var transposeCalled = false + val ops = object : TensorOps by TestTensorOps() { + override fun transpose(tensor: Tensor): Tensor { + transposeCalled = true + return tensor + } + } + + val graph = DefaultComputeGraph() + val input = graph.addNode(inputNode("input")) + val perm = graph.addNode(opNode("perm", "permute")) + graph.addEdge(GraphEdge("e1", input, perm, tensorSpec = spec())) + + val executor = ComputeGraphExecutor(graph, ops) + val results = executor.execute(mapOf("input" to TestTensor(floatArrayOf(1f, 2f)))) + + assertNotNull(results["perm"]) + assertTrue(transposeCalled, "legacy traces without axes keep the transpose behavior") + } + @Test fun emptyGraphReturnsEmptyResults() { val graph = DefaultComputeGraph() @@ -133,8 +190,11 @@ class ComputeGraphExecutorTest { * Used for testing graph execution without needing a full tensor backend. */ @Suppress("UNCHECKED_CAST") -private class TestTensor(val values: FloatArray) : Tensor { - private val tensorShape = sk.ainet.lang.tensor.Shape(intArrayOf(1, values.size)) +private class TestTensor( + val values: FloatArray, + shape: IntArray = intArrayOf(1, values.size), +) : Tensor { + private val tensorShape = sk.ainet.lang.tensor.Shape(shape) override val data: sk.ainet.lang.tensor.data.TensorData = object : sk.ainet.lang.tensor.data.TensorData { override fun get(vararg indices: Int): Float = values[indices.last()] override fun set(vararg indices: Int, value: Float) { values[indices.last()] = value }