feat(export): env-gated true-dynamic KV-cache decode graphs (Gemma + Moonshine v2)

gradle/libs.versions.toml

@@ -1,5 +1,5 @@
 [versions]
-skainet = "0.37.0"
+skainet = "0.38.0"
 agp = "9.3.1"
 jacksonDatabind = "2.22.1"
 jsonSchemaValidator = "3.0.6"

llm-inference/moonshine/src/commonMain/kotlin/sk/ainet/models/moonshine/MoonshineDecoder.kt

@@ -177,11 +177,12 @@ public class MoonshineDecoderLayer<T : DType, V>(
         input: Tensor<T, V>,
         encoderMemory: Tensor<T, V>,
         ctx: ExecutionContext,
+        crossMask: Tensor<T, V>? = null,
     ): MoonshineLayerKV<T, V> {
         val ops = ctx.ops
         val selfKV = selfAttn.forwardWithKV(selfNorm.forward(input, ctx), null, ctx)
         val afterSelf = ops.add(input, selfKV.output)
-        val crossKV = crossAttn.forwardWithKV(crossNorm.forward(afterSelf, ctx), encoderMemory, ctx)
+        val crossKV = crossAttn.forwardWithKV(crossNorm.forward(afterSelf, ctx), encoderMemory, ctx, crossMask)
         val afterCross = ops.add(afterSelf, crossKV.output)
         val h = mlpFc1.forward(mlpNorm.forward(afterCross, ctx), ctx)
         val lastDim = h.rank - 1
@@ -212,6 +213,7 @@ public class MoonshineDecoderLayer<T : DType, V>(
         crossKIn: Tensor<T, V>,
         crossVIn: Tensor<T, V>,
         ctx: ExecutionContext,
+        crossMask: Tensor<T, V>? = null,
     ): MoonshinePastKV<T, V> {
         val ops = ctx.ops
         // --- self-attention: project, RoPE@runtime-position (cos/sin fed in), append + attend ---
@@ -224,7 +226,7 @@ public class MoonshineDecoderLayer<T : DType, V>(
         val afterSelf = ops.add(input, sdpaMerge(selfAttn, q, fullK, fullV, ctx))
         // --- cross-attention: project Q only, attend over the CACHED cross K/V ---
         val cq = projHeads(crossAttn, crossNorm.forward(afterSelf, ctx), 0, ctx)
-        val afterCross = ops.add(afterSelf, sdpaMerge(crossAttn, cq, crossKIn, crossVIn, ctx))
+        val afterCross = ops.add(afterSelf, sdpaMerge(crossAttn, cq, crossKIn, crossVIn, ctx, crossMask))
         // --- gated SiLU MLP ---
         val h = mlpFc1.forward(mlpNorm.forward(afterCross, ctx), ctx)
         val ld = h.rank - 1
@@ -240,12 +242,13 @@ public class MoonshineDecoderLayer<T : DType, V>(
     }
 
     // SDPA over cached K/V then output projection. q [nHeads,1,headDim]; k/v [1,nHeads,S,headDim].
-    private fun sdpaMerge(mha: MultiHeadAttention<T, V>, q: Tensor<T, V>, k: Tensor<T, V>, v: Tensor<T, V>, ctx: ExecutionContext): Tensor<T, V> {
+    private fun sdpaMerge(mha: MultiHeadAttention<T, V>, q: Tensor<T, V>, k: Tensor<T, V>, v: Tensor<T, V>, ctx: ExecutionContext, mask: Tensor<T, V>? = null): Tensor<T, V> {
         val ops = ctx.ops
         val o = ops.scaledDotProductAttention(
             query = ops.unsqueeze(q, 0), key = k, value = v,
-            mask = null, scale = 1f / sqrt(mha.headDim.toFloat()), causal = false,
-        ) // [1, nHeads, 1, headDim]; single query attends to all cached positions → no mask needed
+            mask = mask, scale = 1f / sqrt(mha.headDim.toFloat()), causal = false,
+        ) // [1, nHeads, 1, headDim]; single query attends to all cached positions. `mask` (additive,
+          // [1,1,1,S]) masks padded cross-memory frames when the cross cache is fixed-max-padded (streaming).
         val merged = ops.reshape(ops.squeeze(o, 0), Shape(1, mha.nHeads * mha.headDim)) // [1, qDim]
         return linearProject(ops, merged, mha.params[3].value) // o_proj (bias=false)
     }
@@ -326,6 +329,7 @@ public class MoonshineDecoderModel<T : DType, V>(
         inputsEmbeds: Tensor<T, V>,
         encoderMemory: Tensor<T, V>,
         ctx: ExecutionContext,
+        crossMask: Tensor<T, V>? = null,
     ): MoonshinePrefillOutput<T, V> {
         val ops = ctx.ops
         val memory = encoderMemory.bind(ctx)
@@ -335,7 +339,7 @@ public class MoonshineDecoderModel<T : DType, V>(
         val crossK = ArrayList<Tensor<T, V>>(layers.size)
         val crossV = ArrayList<Tensor<T, V>>(layers.size)
         for (layer in layers) {
-            val kv = layer.forwardWithKV(h, memory, ctx)
+            val kv = layer.forwardWithKV(h, memory, ctx, crossMask)
             h = kv.out
             // add the batch dim so the exported shapes match the board's [1, nHeads, ·, headDim].
             selfK += ops.unsqueeze(kv.selfK, 0)
@@ -362,13 +366,14 @@ public class MoonshineDecoderModel<T : DType, V>(
         crossKIn: List<Tensor<T, V>>,
         crossVIn: List<Tensor<T, V>>,
         ctx: ExecutionContext,
+        crossMask: Tensor<T, V>? = null,
     ): MoonshineWithPastOutput<T, V> {
         val ops = ctx.ops
         var h = tokenEmbed.bind(ctx)
         val nsk = ArrayList<Tensor<T, V>>(layers.size)
         val nsv = ArrayList<Tensor<T, V>>(layers.size)
         for ((i, layer) in layers.withIndex()) {
-            val r = layer.forwardWithPast(h, ropeCos, ropeSin, selfKIn[i], selfVIn[i], crossKIn[i], crossVIn[i], ctx)
+            val r = layer.forwardWithPast(h, ropeCos, ropeSin, selfKIn[i], selfVIn[i], crossKIn[i], crossVIn[i], ctx, crossMask)
             h = r.out
             // The extended cache also feeds this layer's SDPA, so it is not a graph sink. Route the
             // exported copy through a shape-preserving reshape so it becomes a distinct output node

llm-inference/moonshine/src/jvmTest/kotlin/sk/ainet/models/moonshine/MoonshineV2DecoderBakeTest.kt

@@ -68,15 +68,33 @@ class MoonshineV2DecoderBakeTest {
         //  PREFILL   embeds+memory → logits + per-layer self/cross K/V (the cross_kv + prefill graph)
         //  WITH_PAST token+cos/sin+caches → logits + extended self K/V (the decoder_kv step)
         val L = cfg.decoderLayers; val nh = cfg.nHeads; val hd = cfg.headDim
+        // MOONSHINE_V2_TRUE_DYNAMIC=1 traces the decode caches with a real dynamic extent (Dim.DYNAMIC):
+        // the with_past self-cache seq dim (grows per step) AND the cross / encoder-memory frames dim
+        // (varies per utterance), so ONE prefill + ONE with_past vmfb serve every decode position and every
+        // encoder length — the streaming-runtime contract. Requires the dynamic-safe tracer + emitter (SKaiNET#891).
+        val trueDynamic = System.getenv("MOONSHINE_V2_TRUE_DYNAMIC") == "1"
+        val framesDim = if (trueDynamic) sk.ainet.lang.tensor.Dim.DYNAMIC else frames
+        // MOONSHINE_V2_MAX_MEM=<N>: fixed-max-pad the encoder-memory (cross) frames to N and add an additive
+        // crossMask input [1,1,1,N] to prefill AND with_past, so ONE pair of vmfbs serves any encoder length
+        // ≤ N (padded frames masked out of cross-attention). The with_past self-cache still grows dynamically.
+        val maxMem = System.getenv("MOONSHINE_V2_MAX_MEM")?.toInt()
         traceGraph("moonshine_v2_decoder_prefill", "MOONSHINE_V2_DEC_PREFILL_OUT") { c ->
-            dec.forwardPrefill(voidF32(Shape(1, seq, cfg.dim)), voidF32(Shape(1, frames, cfg.dim)), c)
+            // Prefill head-splits the encoder memory; under a *dynamic* frames dim that would need
+            // stablehlo.dynamic_reshape (iree rejects it), so prefill uses a FIXED memory length — a concrete
+            // `frames`, or MOONSHINE_V2_MAX_MEM padded + a crossMask (the streaming path).
+            val preFrames = maxMem ?: frames
+            val preMask = maxMem?.let { voidF32(Shape(1, 1, 1, it)) }
+            dec.forwardPrefill(voidF32(Shape(1, seq, cfg.dim)), voidF32(Shape(1, preFrames, cfg.dim)), c, preMask)
         }
         traceGraph("moonshine_v2_decoder_with_past", "MOONSHINE_V2_DEC_WITHPAST_OUT") { c ->
-            val past = System.getenv("DEC_PAST")?.toInt() ?: 1
+            val past = if (trueDynamic) sk.ainet.lang.tensor.Dim.DYNAMIC else (System.getenv("DEC_PAST")?.toInt() ?: 1)
+            val crossFrames = maxMem ?: framesDim
+            val crossMask = maxMem?.let { voidF32(Shape(1, 1, 1, it)) }
             dec.forwardWithPast(
                 voidF32(Shape(1, 1, cfg.dim)), voidF32(Shape(1, hd)), voidF32(Shape(1, hd)),
                 List(L) { voidF32(Shape(1, nh, past, hd)) }, List(L) { voidF32(Shape(1, nh, past, hd)) },
-                List(L) { voidF32(Shape(1, nh, frames, hd)) }, List(L) { voidF32(Shape(1, nh, frames, hd)) }, c,
+                List(L) { voidF32(Shape(1, nh, crossFrames, hd)) }, List(L) { voidF32(Shape(1, nh, crossFrames, hd)) },
+                c, crossMask,
             )
         }
     }

llm-runtime/kgemma/src/jvmMain/kotlin/sk/ainet/apps/kgemma/FunctionGemmaExport.kt

@@ -241,7 +241,17 @@ public object FunctionGemmaExport {
         // Trace at a CONCRETE length so `concat` shape-inference is valid (a `-1` placeholder
         // mis-infers `-1 + 1 = 0` → broken `1x1x0x256` output caches). For the dynamic graph we
         // trace at the sentinel prime and relax it to `?` after emit (see relaxSeqDimToDynamic).
-        val pastDim = if (dynamicPast) SENTINEL_PAST else past
+        // GEMMA_TRUE_DYNAMIC=1: thread a real dynamic extent (Dim.DYNAMIC) straight through the trace
+        // instead of the sentinel-prime + post-emit text-relax. Requires the dynamic-safe tracer (concat
+        // and reshape propagate a dynamic dim) and the dynamic-safe emitter (dynamic_broadcast_in_dim);
+        // no text-relax needed. Verified to iree-compile the with_past graph (vs the sentinel path, which
+        // does not). Kept env-gated so the default export path is unchanged until the core release lands.
+        val trueDynamic = System.getenv("GEMMA_TRUE_DYNAMIC") == "1"
+        val pastDim = when {
+            !dynamicPast -> past
+            trueDynamic -> sk.ainet.lang.tensor.Dim.DYNAMIC
+            else -> SENTINEL_PAST
+        }
 
         val tokenId = voidF32(Shape(1))
         val cosG = voidF32(Shape(1, headDim)); val sinG = voidF32(Shape(1, headDim))
@@ -273,7 +283,7 @@ public object FunctionGemmaExport {
             .createBasic(ConstantMaterializationPolicy.ExternalAlways(scope = "model"))
             .convert(graph, "gemma_with_past")
         var mlir = if (bf16) rewriteGlobalsToBf16(module.content) else module.content
-        if (dynamicPast) mlir = relaxSeqDimToDynamic(mlir)
+        if (dynamicPast && !trueDynamic) mlir = relaxSeqDimToDynamic(mlir)
         File(outDir).apply { mkdirs() }
         File(outDir, "gemma-with-past.mlir").writeText(mlir)
         mlir

transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/MultiHeadAttention.kt

@@ -222,6 +222,7 @@ public class MultiHeadAttention<T : DType, V>(
         input: Tensor<T, V>,
         encoderMemory: Tensor<T, V>?,
         ctx: ExecutionContext,
+        crossMask: Tensor<T, V>? = null,
     ): AttentionKV<T, V> {
         val boundInput = input.bind(ctx)
         return if (encoderMemory == null) {
@@ -231,7 +232,7 @@ public class MultiHeadAttention<T : DType, V>(
                 "MultiHeadAttention.forwardWithKV: cross-attention supports neither kvCache nor slidingWindow."
             }
             val boundMemory = encoderMemory.bind(ctx)
-            attentionImpl(qInput = boundInput, kvInput = boundMemory, isCrossAttention = true, ctx = ctx)
+            attentionImpl(qInput = boundInput, kvInput = boundMemory, isCrossAttention = true, ctx = ctx, crossMask = crossMask)
         }
     }
 
@@ -271,6 +272,7 @@ public class MultiHeadAttention<T : DType, V>(
         kvInput: Tensor<T, V>,
         isCrossAttention: Boolean,
         ctx: ExecutionContext,
+        crossMask: Tensor<T, V>? = null,
     ): AttentionKV<T, V> {
         val ops = ctx.ops
         val scale = attentionScale ?: (1.0f / sqrt(headDim.toFloat()))
@@ -434,12 +436,14 @@ public class MultiHeadAttention<T : DType, V>(
         // ordering between decoder query positions and encoder memory frames.
         val useCausalPath = !isCrossAttention && causal && slidingMask == null
 
-        // Scaled dot-product attention
+        // Scaled dot-product attention. `crossMask` (additive, e.g. [1,1,1,seqKV]) masks padded encoder-memory
+        // frames when the cross cache is fixed-max-padded (streaming) — used only on the cross path, where
+        // slidingMask is guaranteed null.
         val attnOut = ops.scaledDotProductAttention(
             query = qBatched,
             key = kBatched,
             value = vBatched,
-            mask = slidingMask,
+            mask = slidingMask ?: crossMask,
             scale = scale,
             causal = useCausalPath
         )