Skip to content

Commit c1e6b07

Browse files
committed
Add narrow op dispatch to ComputeGraphExecutor
The executor was missing a case for the narrow operation, causing graph execution to fail when MultiHeadAttention traces GQA head expansion (repeatKVHeads) into the optimized compute graph.
1 parent f3a0393 commit c1e6b07

1 file changed

Lines changed: 6 additions & 0 deletions

File tree

skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/exec/ComputeGraphExecutor.kt

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -195,6 +195,12 @@ public class ComputeGraphExecutor(
195195
val dim = (params["dim"] as? Number)?.toInt() ?: 0
196196
listOf(ops.unsqueeze(inputs[0], dim))
197197
}
198+
"narrow" -> {
199+
val dim = (params["dim"] as? Number)?.toInt() ?: 0
200+
val start = (params["start"] as? Number)?.toInt() ?: 0
201+
val length = (params["length"] as? Number)?.toInt() ?: 1
202+
listOf(ops.narrow(inputs[0], dim, start, length))
203+
}
198204

199205
// Reduction ops
200206
"mean" -> {

0 commit comments

Comments
 (0)