feat: first-class dynamic tensor dimensions for streaming KV-cache decode (squashed)

skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/trace/TraceToGraphBuilder.kt

@@ -368,19 +368,9 @@ public class TraceToGraphBuilder(
             return buffer.copyOf()
         }
 
-        // Fallback for other data types if possible
-        if (tensor.volume > 0) {
-            val result = FloatArray(tensor.volume)
-            // This is slow but generic. Better if we have a way to get values.
-            // But usually weights are FloatArrayTensorData in the contexts we use for export.
-            return try {
-                // We don't have a good way to iterate over all indices generically without recursion
-                // for arbitrary rank. Let's stick to FloatArrayTensorData for now as it's the most common.
-                null
-            } catch (e: Exception) {
-                null
-            }
-        }
+        // Nothing else is materializable here: weights are FloatArrayTensorData in export contexts, and a
+        // dynamic-shaped tensor (e.g. a `?` KV-cache input) has no volume to probe — never call `.volume`
+        // on it (it throws by design). Such tensors are graph inputs, not constants to embed, so return null.
         return null
     }
 

skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/TypeMapper.kt

@@ -1,5 +1,6 @@
 package sk.ainet.compile.hlo
 
+import sk.ainet.lang.tensor.Dim
 import sk.ainet.lang.tensor.ops.TensorSpec
 
 /**
@@ -117,24 +118,34 @@ public class TypeMapper {
     }
 
     /**
-     * Format tensor shape for MLIR
+     * Format tensor shape for MLIR. A negative extent renders as `?` (a dynamic dimension — see
+     * [DYNAMIC_DIM]); a `null` shape is a fully-dynamic tensor (`?`).
      */
-    private fun formatShape(shape: List<Int>?): String {
+    public fun formatShape(shape: List<Int>?): String {
         return when {
             shape == null -> "?"
             shape.isEmpty() -> ""
-            else -> shape.joinToString("x") { if (it < 0) "?" else it.toString() }
+            else -> shape.joinToString("x") { Dim.render(it) }
         }
     }
-    
+
     /**
-     * Create a tensor type string with explicit shape
+     * Create a tensor type string with explicit shape. Renders negative extents as `?` (dynamic).
      */
     public fun createTensorType(shape: List<Int>, dtype: String): String {
         val elementType = mapDType(dtype)
         if (shape.isEmpty()) return "tensor<$elementType>" // rank-0 scalar
-        val shapeStr = shape.joinToString("x")
-        return "tensor<${shapeStr}x${elementType}>"
+        return "tensor<${formatShape(shape)}x${elementType}>"
+    }
+
+    public companion object {
+        /**
+         * Sentinel extent meaning "dynamic dimension" (`?`) in a [TensorSpec] shape. Threaded from the trace
+         * (e.g. a KV-cache seq dim) so the emitter renders `?` and picks dynamic-shape-safe op forms, instead
+         * of the legacy post-emit text substitution. Aliases the canonical [Dim.DYNAMIC] (a reserved sentinel
+         * distinct from reshape's `-1` = infer), so tracer and emitter agree on one value.
+         */
+        public const val DYNAMIC_DIM: Int = Dim.DYNAMIC
     }
 
     /**

skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ActivationOperationsConverter.kt

@@ -3,6 +3,8 @@ package sk.ainet.compile.hlo.converters
 import sk.ainet.compile.hlo.ConversionContext
 import sk.ainet.compile.hlo.ConversionResult
 import sk.ainet.compile.hlo.StableHloOperationConverter
+import sk.ainet.lang.tensor.Dim
+import sk.ainet.lang.tensor.hasDynamic
 import sk.ainet.lang.graph.GraphNode
 
 /**
@@ -139,8 +141,14 @@ public class ActivationOperationsConverter : StableHloOperationConverter {
         val reducedType = if (reducedShape.isEmpty()) {
             "tensor<$elementType>"
         } else {
-            "tensor<${reducedShape.joinToString("x")}x$elementType>"
+            "tensor<${reducedShape.joinToString("x") { Dim.render(it) }}x$elementType>"
         }
+        // Dynamic softmax axis / leading dims (`?`): the reduced max/sum must broadcast back to a dynamic
+        // output shape, which `stablehlo.broadcast_in_dim` cannot target — use `stablehlo.dynamic_broadcast_in_dim`
+        // with a runtime `output_dimensions` operand built from the input via `get_dimension_size` (IREE's
+        // stablehlo pipeline rejects CHLO implicit-broadcast ops as illegal).
+        val dyn = inputShape.hasDynamic()
+        val shapeType = "tensor<${rank}xi32>"
 
         // Dimensions kept for broadcast_in_dim: every input dim except `axis`,
         // mapped to its position in the reduced tensor.
@@ -161,36 +169,54 @@ public class ActivationOperationsConverter : StableHloOperationConverter {
         // constant is out of range).
         val maxIdentity = context.getTypeMapper().negInfBits(elementType)
 
-        val operations = listOf(
+        val operations = buildList {
             // Reduce-max along the softmax axis (for numerical stability).
-            "$maxInit = stablehlo.constant dense<$maxIdentity> : tensor<$elementType>",
-            "$maxValue = stablehlo.reduce(${operands[0]} init: $maxInit) " +
-                "applies stablehlo.maximum across dimensions = [$axis] : " +
-                "($outputType, tensor<$elementType>) -> $reducedType",
-
-            // Broadcast reduced max back to the input shape.
-            "$maxBroadcast = stablehlo.broadcast_in_dim $maxValue, " +
-                "dims = [$broadcastDims] : ($reducedType) -> $outputType",
-
-            // Subtract the max for numerical stability.
-            "$shiftedValue = stablehlo.subtract ${operands[0]}, $maxBroadcast : $outputType",
-
-            // Elementwise exponential.
-            "$expValue = stablehlo.exponential $shiftedValue : $outputType",
-
-            // Reduce-sum along the softmax axis.
-            "$sumInit = stablehlo.constant dense<0.0> : tensor<$elementType>",
-            "$sumValue = stablehlo.reduce($expValue init: $sumInit) " +
-                "applies stablehlo.add across dimensions = [$axis] : " +
-                "($outputType, tensor<$elementType>) -> $reducedType",
-
-            // Broadcast the sum back to the input shape.
-            "$sumBroadcast = stablehlo.broadcast_in_dim $sumValue, " +
-                "dims = [$broadcastDims] : ($reducedType) -> $outputType",
-
-            // Normalize.
-            "$resultValue = stablehlo.divide $expValue, $sumBroadcast : $outputType"
-        )
+            add("$maxInit = stablehlo.constant dense<$maxIdentity> : tensor<$elementType>")
+            add(
+                "$maxValue = stablehlo.reduce(${operands[0]} init: $maxInit) " +
+                    "applies stablehlo.maximum across dimensions = [$axis] : " +
+                    "($outputType, tensor<$elementType>) -> $reducedType",
+            )
+            // Build the runtime output-shape operand once (only needed for dynamic broadcasts).
+            val shapeOperand: String = if (!dyn) "" else run {
+                val parts = inputShape.indices.map { d ->
+                    if (Dim.isStatic(inputShape[d])) {
+                        val c = context.nextTempValue()
+                        add("$c = stablehlo.constant dense<${inputShape[d]}> : tensor<1xi32>")
+                        c
+                    } else {
+                        val gd = context.nextTempValue(); val gr = context.nextTempValue()
+                        add("$gd = stablehlo.get_dimension_size ${operands[0]}, dim = $d : ($outputType) -> tensor<i32>")
+                        add("$gr = stablehlo.reshape $gd : (tensor<i32>) -> tensor<1xi32>")
+                        gr
+                    }
+                }
+                val sh = context.nextTempValue()
+                add("$sh = stablehlo.concatenate ${parts.joinToString(", ")}, dim = 0 : (${parts.joinToString(", ") { "tensor<1xi32>" }}) -> $shapeType")
+                sh
+            }
+            // Subtract the max: static → explicit broadcast_in_dim; dynamic → runtime dynamic_broadcast_in_dim.
+            if (dyn) {
+                add("$maxBroadcast = stablehlo.dynamic_broadcast_in_dim $maxValue, $shapeOperand, dims = [$broadcastDims] : ($reducedType, $shapeType) -> $outputType")
+            } else {
+                add("$maxBroadcast = stablehlo.broadcast_in_dim $maxValue, dims = [$broadcastDims] : ($reducedType) -> $outputType")
+            }
+            add("$shiftedValue = stablehlo.subtract ${operands[0]}, $maxBroadcast : $outputType")
+            add("$expValue = stablehlo.exponential $shiftedValue : $outputType")
+            add("$sumInit = stablehlo.constant dense<0.0> : tensor<$elementType>")
+            add(
+                "$sumValue = stablehlo.reduce($expValue init: $sumInit) " +
+                    "applies stablehlo.add across dimensions = [$axis] : " +
+                    "($outputType, tensor<$elementType>) -> $reducedType",
+            )
+            // Normalize: static → broadcast_in_dim; dynamic → runtime dynamic_broadcast_in_dim.
+            if (dyn) {
+                add("$sumBroadcast = stablehlo.dynamic_broadcast_in_dim $sumValue, $shapeOperand, dims = [$broadcastDims] : ($reducedType, $shapeType) -> $outputType")
+            } else {
+                add("$sumBroadcast = stablehlo.broadcast_in_dim $sumValue, dims = [$broadcastDims] : ($reducedType) -> $outputType")
+            }
+            add("$resultValue = stablehlo.divide $expValue, $sumBroadcast : $outputType")
+        }
 
         operations.forEach { context.emitOperation(it) }
 

skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/AttentionOperationsConverter.kt

@@ -3,6 +3,8 @@ package sk.ainet.compile.hlo.converters
 import sk.ainet.compile.hlo.ConversionContext
 import sk.ainet.compile.hlo.ConversionResult
 import sk.ainet.compile.hlo.StableHloOperationConverter
+import sk.ainet.lang.tensor.Dim
+import sk.ainet.lang.tensor.hasDynamic
 import sk.ainet.lang.graph.GraphNode
 import kotlin.math.sqrt
 
@@ -56,7 +58,9 @@ public class AttentionOperationsConverter : StableHloOperationConverter {
         val outSpec = node.outputs.firstOrNull()
         val mapper = context.getTypeMapper()
         val elem = outSpec?.let { mapper.mapDType(it.dtype) } ?: "f32"
-        fun typeOf(shape: List<Int>): String = "tensor<${shape.joinToString("x")}x$elem>"
+        // Render a shape's dims, mapping a dynamic extent (DYNAMIC_DIM = -1) to `?`.
+        fun dims(shape: List<Int>): String = shape.joinToString("x") { Dim.render(it) }
+        fun typeOf(shape: List<Int>): String = "tensor<${dims(shape)}x$elem>"
 
         val qType = context.getValueType(operands[0]) ?: typeOf(qShape)
         val kType = context.getValueType(operands[1]) ?: typeOf(kShape)
@@ -77,29 +81,36 @@ public class AttentionOperationsConverter : StableHloOperationConverter {
         val contractQK = rank - 1                 // contract head_dim of Q and K
         val sdAxis = scoresShape.size - 1         // softmax over key length
         val reducedShape = scoresShape.dropLast(1)
-        val reducedType = if (reducedShape.isEmpty()) "tensor<$elem>" else "tensor<${reducedShape.joinToString("x")}x$elem>"
+        val reducedType = if (reducedShape.isEmpty()) "tensor<$elem>" else typeOf(reducedShape)
         val bcastDims = (scoresShape.indices).filter { it != sdAxis }.joinToString(", ")
         val contractAttn = scoresShape.size - 1   // attn key-length axis
         val contractV = rank - 2                  // V key-length axis
 
         val causal = (node.operation.parameters["causal"] as? Boolean) ?: false
         val qAxis = rank - 2 // query position in scores [.., Sq, Sk]
-        val scoresI32Type = "tensor<${scoresShape.joinToString("x")}xi32>"
-        val scoresI1Type = "tensor<${scoresShape.joinToString("x")}xi1>"
+        val scoresI32Type = "tensor<${dims(scoresShape)}xi32>"
+        val scoresI1Type = "tensor<${dims(scoresShape)}xi1>"
 
-        val scores = context.nextTempValue()
         val scaleC = context.nextTempValue()
-        val scaled = context.nextTempValue()
+        val scaleB = context.nextTempValue()
+        val qScaled = context.nextTempValue()
+        val scores = context.nextTempValue()
         val maxInit = context.nextTempValue(); val maxV = context.nextTempValue(); val maxB = context.nextTempValue()
         val shifted = context.nextTempValue(); val expV = context.nextTempValue()
         val sumInit = context.nextTempValue(); val sumV = context.nextTempValue(); val sumB = context.nextTempValue()
         val attn = context.nextTempValue()
         val out = context.nextTempValue()
 
         val ops = mutableListOf(
-            "$scores = stablehlo.dot_general ${operands[0]}, ${operands[1]}, ${batchClause}contracting_dims = [$contractQK] x [$contractQK] : ($qType, $kType) -> $scoresType",
-            "$scaleC = stablehlo.constant dense<$scaleVal> : $scoresType",
-            "$scaled = stablehlo.multiply $scores, $scaleC : $scoresType",
+            // Scale is applied to Q *before* the QK dot — scores·s == (q·s)@kᵀ exactly. This keeps the scale
+            // constant at the STATIC query type (`[.., Sq, headDim]`) rather than a splat sized to the scores
+            // shape `[.., Sq, Sk]`, whose Sk (key/cache length) may be dynamic (`?`) in KV-cache decode — a
+            // dynamic-shape splat constant is invalid StableHLO. Also drops a full scores-sized dense constant
+            // from static graphs.
+            "$scaleC = stablehlo.constant dense<$scaleVal> : tensor<$elem>",
+            "$scaleB = stablehlo.broadcast_in_dim $scaleC, dims = [] : (tensor<$elem>) -> $qType",
+            "$qScaled = stablehlo.multiply ${operands[0]}, $scaleB : $qType",
+            "$scores = stablehlo.dot_general $qScaled, ${operands[1]}, ${batchClause}contracting_dims = [$contractQK] x [$contractQK] : ($qType, $kType) -> $scoresType",
         )
 
         // Explicit additive mask (operands[3]) — e.g. a sliding-window+causal
@@ -108,7 +119,7 @@ public class AttentionOperationsConverter : StableHloOperationConverter {
         // iota causal path. Broadcast (trailing-aligned) to the scores shape
         // and add. Without this the masked layers run UNMASKED (attend to
         // future tokens) — correct only at position 0.
-        var softmaxIn = scaled
+        var softmaxIn = scores   // scores are already scaled (scale folded into Q above)
         val maskOperand = operands.getOrNull(3)
         if (maskOperand != null) {
             val maskShape = node.inputs.getOrNull(3)?.shape ?: scoresShape
@@ -123,7 +134,7 @@ public class AttentionOperationsConverter : StableHloOperationConverter {
                 mb
             }
             val masked = context.nextTempValue()
-            ops += "$masked = stablehlo.add $scaled, $maskBc : $scoresType"
+            ops += "$masked = stablehlo.add $scores, $maskBc : $scoresType"
             softmaxIn = masked
         } else if (causal) {
             val iotaQ = context.nextTempValue(); val iotaK = context.nextTempValue()
@@ -139,19 +150,49 @@ public class AttentionOperationsConverter : StableHloOperationConverter {
             // masked-fill select, which can trip downstream greedy constant-folding.
             ops += "$ninf = stablehlo.constant dense<-1.000000e+30> : $scoresType"
             ops += "$maskAdd = stablehlo.select $keep, $zeros, $ninf : $scoresI1Type, $scoresType"
-            ops += "$masked = stablehlo.add $scaled, $maskAdd : $scoresType"
+            ops += "$masked = stablehlo.add $scores, $maskAdd : $scoresType"
             softmaxIn = masked
         }
 
-        // softmax(softmaxIn) over the key-length axis
+        // softmax(softmaxIn) over the key-length axis. When the scores shape is dynamic (the `?` key/cache dim
+        // of KV-cache decode), the reduced max/sum must broadcast back to the dynamic scores shape. A static
+        // `stablehlo.broadcast_in_dim` cannot target a dynamic shape, so we use `stablehlo.dynamic_broadcast_in_dim`
+        // with a runtime `output_dimensions` operand (built once from the scores tensor via `get_dimension_size`).
+        // Static graphs keep the original explicit `broadcast_in_dim` path (byte-for-byte unchanged).
+        val dyn = scoresShape.hasDynamic()
+        val shapeType = "tensor<${scoresShape.size}xi32>"
+        val scoresShapeOperand: String = if (!dyn) "" else run {
+            val parts = scoresShape.indices.map { d ->
+                if (Dim.isStatic(scoresShape[d])) {
+                    val c = context.nextTempValue()
+                    ops += "$c = stablehlo.constant dense<${scoresShape[d]}> : tensor<1xi32>"
+                    c
+                } else {
+                    val gd = context.nextTempValue(); val gr = context.nextTempValue()
+                    ops += "$gd = stablehlo.get_dimension_size $scores, dim = $d : ($scoresType) -> tensor<i32>"
+                    ops += "$gr = stablehlo.reshape $gd : (tensor<i32>) -> tensor<1xi32>"
+                    gr
+                }
+            }
+            val sh = context.nextTempValue()
+            ops += "$sh = stablehlo.concatenate ${parts.joinToString(", ")}, dim = 0 : (${parts.joinToString(", ") { "tensor<1xi32>" }}) -> $shapeType"
+            sh
+        }
+        fun broadcastBack(src: String, dst: String) {
+            if (dyn) {
+                ops += "$dst = stablehlo.dynamic_broadcast_in_dim $src, $scoresShapeOperand, dims = [$bcastDims] : ($reducedType, $shapeType) -> $scoresType"
+            } else {
+                ops += "$dst = stablehlo.broadcast_in_dim $src, dims = [$bcastDims] : ($reducedType) -> $scoresType"
+            }
+        }
         ops += "$maxInit = stablehlo.constant dense<${mapper.negInfBits(elem)}> : tensor<$elem>"
         ops += "$maxV = stablehlo.reduce($softmaxIn init: $maxInit) applies stablehlo.maximum across dimensions = [$sdAxis] : ($scoresType, tensor<$elem>) -> $reducedType"
-        ops += "$maxB = stablehlo.broadcast_in_dim $maxV, dims = [$bcastDims] : ($reducedType) -> $scoresType"
+        broadcastBack(maxV, maxB)
         ops += "$shifted = stablehlo.subtract $softmaxIn, $maxB : $scoresType"
         ops += "$expV = stablehlo.exponential $shifted : $scoresType"
         ops += "$sumInit = stablehlo.constant dense<0.0> : tensor<$elem>"
         ops += "$sumV = stablehlo.reduce($expV init: $sumInit) applies stablehlo.add across dimensions = [$sdAxis] : ($scoresType, tensor<$elem>) -> $reducedType"
-        ops += "$sumB = stablehlo.broadcast_in_dim $sumV, dims = [$bcastDims] : ($reducedType) -> $scoresType"
+        broadcastBack(sumV, sumB)
         ops += "$attn = stablehlo.divide $expV, $sumB : $scoresType"
         ops += "$out = stablehlo.dot_general $attn, ${operands[2]}, ${batchClause}contracting_dims = [$contractAttn] x [$contractV] : ($scoresType, $vType) -> $outputType"
         ops.forEach { context.emitOperation(it) }