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

CHANGELOG.md

@@ -2,6 +2,58 @@
 
 ## [Unreleased]
 
+### Added
+
+- **First-class dynamic dimensions (`Dim`).** A new `sk.ainet.lang.tensor.Dim` vocabulary makes "dynamic
+  extent" explicit instead of an overloaded `-1`: `Dim.DYNAMIC` is a reserved sentinel (`Int.MIN_VALUE`)
+  **distinct from reshape's `-1` = infer**, with the dynamic-aware shape arithmetic (`concat`, `compatible`,
+  `render`, `isDynamic`/`isStatic`) centralized in one place rather than scattered `extent < 0` guards.
+  `Shape` gains `hasDynamic()`, `isDynamic(axis)`, `dynamicAxes`, and its `volume` now throws on a dynamic
+  shape (an unknown extent has no materializable element count) instead of returning a corrupt product.
+  The slice/range DSL is dynamic-aware too: `all()` over a dynamic axis is a valid symbolic full-axis, and
+  reshape passes a dynamic target through unchanged. The `skainet-compile-hlo` emitter now shares this one
+  sentinel (`TypeMapper.DYNAMIC_DIM` aliases `Dim.DYNAMIC`), so tracer and emitter agree by construction.
+- **Dynamic-shape-safe StableHLO emission (streaming KV-cache decode).** The `skainet-compile-hlo`
+  emitter now renders a `-1` tensor extent as an MLIR `?` (dynamic) dimension and emits op forms that
+  `iree-compile` accepts under a dynamic dim, so a single compiled vmfb serves every autoregressive
+  decode step (growing KV cache) instead of one fixed cache length. A new `TypeMapper.DYNAMIC_DIM = -1`
+  marker plus a `List<Int>.hasDynamic()` predicate gate the dynamic paths, so every static graph is
+  emitted byte-for-byte unchanged. Verified: one dynamic SDPA vmfb runs at key lengths 3 and 17, and
+  the full FunctionGemma `with_past` decode graph (real weights, dynamic `1x{nKV}x?x256` cache)
+  self-compiles from the DSL to a CPU vmfb — a graph that could not be compiled before.
+- **Dynamic-shape-safe trace finalization.** `TraceToGraphBuilder.extractFloatArray` no longer probes
+  `Tensor.volume` for non-dense data, so a dynamic-shaped graph input (e.g. a `?` KV-cache tensor) is left
+  as an input instead of tripping the (correctly) throwing `volume` on an unknown extent. This lets a decode
+  graph with dynamic caches finalize to a `ComputeGraph` and compile — verified with the real Moonshine v2
+  `with_past` decoder (dynamic self *and* cross caches, `1x8x?x40`) self-compiling from the DSL to a vmfb.
+- **Allocation-free shape-only tracing (`VoidTensorOps`).** The trace-time op set now propagates shapes
+  through a `ShapeOnlyTensorData` that carries a `Shape` but allocates no backing buffer, so a dynamic
+  (`-1`) extent flows through a whole decode trace instead of throwing `NegativeArraySizeException` when
+  a real buffer of negative size is allocated. This is what lets a real KV-cache seq dim be traced as
+  dynamic end-to-end (rather than via a sentinel-dimension + post-emit text substitution).
+
+### Changed
+
+- **SDPA scale folded into Q.** `AttentionOperationsConverter` now applies the attention scale as a
+  scalar constant multiplied into Q *before* the QK `dot_general` (`(q·s)@kᵀ ≡ scores·s`, exact),
+  instead of a dense splat constant sized to the full scores shape. This drops a scores-sized constant
+  from every attention graph and, crucially, avoids an invalid dynamic-shape splat when the key/cache
+  dim is `?`.
+- **Softmax broadcasts are dynamic-shape-safe.** When the softmax/scores shape carries a dynamic dim,
+  `AttentionOperationsConverter` and `ActivationOperationsConverter` broadcast the reduced max/sum with
+  `stablehlo.dynamic_broadcast_in_dim` (runtime shape operand built via `get_dimension_size` +
+  `concatenate`) instead of a static `broadcast_in_dim`. Static graphs keep the explicit
+  `broadcast_in_dim` unchanged.
+- **Identity reshapes/slices are elided.** `ShapeOperationsConverter` returns the operand SSA value with
+  no emitted op when a reshape's input and result types are identical, or a `slice`/`narrow` covers the
+  full extent of every axis (the KV-cache "cache-as-output-sink" and full-cache head-expansion patterns).
+  Besides being a no-op, this is the only valid lowering on a dynamic axis — a static `stablehlo.slice`
+  cannot express a full-extent bound on a `?` dim (its limit would be the `-1` extent, e.g. `0:-1:1`).
+- **Concatenate propagates dynamic extents.** Both the trace-time shape inference
+  (`VoidTensorOps.calculateConcatShape`) and the emitter (`ShapeOperationsConverter` concat) now keep the
+  concatenated axis dynamic when any operand's extent there is dynamic, instead of numerically summing it
+  (which turned a growing cache `? ++ 1` into a bogus static `0`).
+
 ## [0.37.0] - 2026-07-25
 
 ### Added

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) }