From 2a6cc9f5fc2467f82510dbb2c147de9374f0b9de Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Mon, 27 Jul 2026 13:46:25 +0200 Subject: [PATCH 1/3] feat(export): env-gated true-dynamic KV-cache decode graphs (Gemma + Moonshine v2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add opt-in flags that trace the KV-cache seq dim as a real dynamic extent (Dim.DYNAMIC) instead of the fixed/sentinel path, so one compiled vmfb serves every autoregressive decode position: - FunctionGemmaExport: GEMMA_TRUE_DYNAMIC=1 threads Dim.DYNAMIC through the with_past trace and skips the SENTINEL_PAST + relaxSeqDimToDynamic text-hack (which never actually iree-compiled). Default export path unchanged. - MoonshineV2DecoderBakeTest: MOONSHINE_V2_TRUE_DYNAMIC=1 traces both the self-cache seq dim (grows per step) and the cross-cache frames dim (varies per utterance) as Dim.DYNAMIC. Both env-gated and off by default. Requires the core dynamic-shape capability (sk.ainet.core Dim + dynamic-safe tracer/emitter, PR SKaiNET#891) — build with -PuseLocalSkainet=true until that lands in a release. Verified: both graphs self-compile to dynamic vmfbs, and Moonshine v2 growing-cache decode matches onnxruntime cos=1.0 at every step. Co-Authored-By: Claude Opus 4.8 --- .../models/moonshine/MoonshineV2DecoderBakeTest.kt | 9 +++++++-- .../sk/ainet/apps/kgemma/FunctionGemmaExport.kt | 14 ++++++++++++-- 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/llm-inference/moonshine/src/jvmTest/kotlin/sk/ainet/models/moonshine/MoonshineV2DecoderBakeTest.kt b/llm-inference/moonshine/src/jvmTest/kotlin/sk/ainet/models/moonshine/MoonshineV2DecoderBakeTest.kt index f2ada04..e821ef5 100644 --- a/llm-inference/moonshine/src/jvmTest/kotlin/sk/ainet/models/moonshine/MoonshineV2DecoderBakeTest.kt +++ b/llm-inference/moonshine/src/jvmTest/kotlin/sk/ainet/models/moonshine/MoonshineV2DecoderBakeTest.kt @@ -72,11 +72,16 @@ class MoonshineV2DecoderBakeTest { dec.forwardPrefill(voidF32(Shape(1, seq, cfg.dim)), voidF32(Shape(1, frames, cfg.dim)), c) } traceGraph("moonshine_v2_decoder_with_past", "MOONSHINE_V2_DEC_WITHPAST_OUT") { c -> - val past = System.getenv("DEC_PAST")?.toInt() ?: 1 + // MOONSHINE_V2_TRUE_DYNAMIC=1 traces the self-cache seq dim (grows per decode step) AND the + // cross-cache frames dim (varies per utterance) as Dim.DYNAMIC, so ONE compiled vmfb serves + // every decode position and every encoder length. Requires the dynamic-safe tracer + emitter. + val trueDynamic = System.getenv("MOONSHINE_V2_TRUE_DYNAMIC") == "1" + val past = if (trueDynamic) sk.ainet.lang.tensor.Dim.DYNAMIC else (System.getenv("DEC_PAST")?.toInt() ?: 1) + val framesDim = if (trueDynamic) sk.ainet.lang.tensor.Dim.DYNAMIC else frames 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, framesDim, hd)) }, List(L) { voidF32(Shape(1, nh, framesDim, hd)) }, c, ) } } diff --git a/llm-runtime/kgemma/src/jvmMain/kotlin/sk/ainet/apps/kgemma/FunctionGemmaExport.kt b/llm-runtime/kgemma/src/jvmMain/kotlin/sk/ainet/apps/kgemma/FunctionGemmaExport.kt index c063927..08520a2 100644 --- a/llm-runtime/kgemma/src/jvmMain/kotlin/sk/ainet/apps/kgemma/FunctionGemmaExport.kt +++ b/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 From 351beba77c7dcc74442cbfbc86da8de80b39a46f Mon Sep 17 00:00:00 2001 From: SKaiNET development team Date: Thu, 30 Jul 2026 13:10:05 +0200 Subject: [PATCH 2/3] build(deps): bump skainet 0.37.0 -> 0.38.0 (Dim/dynamic-shapes now published) 0.38.0 (cut from develop, on Maven Central) carries the SKaiNET#891 dynamic-shape core (Dim + dynamic tracer/emitter). Lets the true-dynamic decode graphs + Gemma GEMMA_TRUE_DYNAMIC export build against published core, no -PuseLocalSkainet needed. Co-Authored-By: Claude Opus 4.8 --- gradle/libs.versions.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml index 93615aa..381e9f6 100644 --- a/gradle/libs.versions.toml +++ b/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" From 96d925f54b68e2923142093168c240a8d88941b3 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Wed, 29 Jul 2026 07:22:46 +0200 Subject: [PATCH 3/3] feat(moonshine,transformer-core): fixed-max cross-memory padding mask for streaming decode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Streaming ASR finalizes a variable-length encoder memory, but the decoder prefill is fixed-shape (iree rejects dynamic_reshape for the memory head-split). Fix by zero-padding the cross (encoder) memory to a fixed MAX and masking the padding out of cross-attention, so ONE prefill + ONE with_past vmfb pair serve any encoder length ≤ MAX while the self-cache stays dynamic (growing). - transformer-core MultiHeadAttention: optional trailing `crossMask` on attentionImpl + forwardWithKV, applied as SDPA `mask = slidingMask ?: crossMask`. Default null → byte-identical for all existing callers (verified: llm-core 101 + gemma 77 tests green, incl. MHA/SDPA/sliding-window coverage). - MoonshineDecoder: thread `crossMask` through the layer/model forwardPrefill (via MHA.forwardWithKV) and the hand-wired forwardWithPast (sdpaMerge). Backward compatible (trailing default; v1 positional callers untouched). - MoonshineV2DecoderBakeTest: MOONSHINE_V2_MAX_MEM=N pads both prefill memory and with_past cross cache to N and adds a crossMask input [1,1,1,N] to both graphs. Verified end-to-end: masked decode over memory padded 64→96 (with a garbage tail) == onnxruntime unpadded-64 token-for-token on real audio; control (no mask) corrupts. Co-Authored-By: Claude Opus 4.8 --- .../models/moonshine/MoonshineDecoder.kt | 19 ++++++++----- .../moonshine/MoonshineV2DecoderBakeTest.kt | 27 ++++++++++++++----- .../lang/nn/transformer/MultiHeadAttention.kt | 10 ++++--- 3 files changed, 39 insertions(+), 17 deletions(-) diff --git a/llm-inference/moonshine/src/commonMain/kotlin/sk/ainet/models/moonshine/MoonshineDecoder.kt b/llm-inference/moonshine/src/commonMain/kotlin/sk/ainet/models/moonshine/MoonshineDecoder.kt index ee50c07..48e4d72 100644 --- a/llm-inference/moonshine/src/commonMain/kotlin/sk/ainet/models/moonshine/MoonshineDecoder.kt +++ b/llm-inference/moonshine/src/commonMain/kotlin/sk/ainet/models/moonshine/MoonshineDecoder.kt @@ -177,11 +177,12 @@ public class MoonshineDecoderLayer( input: Tensor, encoderMemory: Tensor, ctx: ExecutionContext, + crossMask: Tensor? = null, ): MoonshineLayerKV { 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( crossKIn: Tensor, crossVIn: Tensor, ctx: ExecutionContext, + crossMask: Tensor? = null, ): MoonshinePastKV { val ops = ctx.ops // --- self-attention: project, RoPE@runtime-position (cos/sin fed in), append + attend --- @@ -224,7 +226,7 @@ public class MoonshineDecoderLayer( 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( } // SDPA over cached K/V then output projection. q [nHeads,1,headDim]; k/v [1,nHeads,S,headDim]. - private fun sdpaMerge(mha: MultiHeadAttention, q: Tensor, k: Tensor, v: Tensor, ctx: ExecutionContext): Tensor { + private fun sdpaMerge(mha: MultiHeadAttention, q: Tensor, k: Tensor, v: Tensor, ctx: ExecutionContext, mask: Tensor? = null): Tensor { 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( inputsEmbeds: Tensor, encoderMemory: Tensor, ctx: ExecutionContext, + crossMask: Tensor? = null, ): MoonshinePrefillOutput { val ops = ctx.ops val memory = encoderMemory.bind(ctx) @@ -335,7 +339,7 @@ public class MoonshineDecoderModel( val crossK = ArrayList>(layers.size) val crossV = ArrayList>(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( crossKIn: List>, crossVIn: List>, ctx: ExecutionContext, + crossMask: Tensor? = null, ): MoonshineWithPastOutput { val ops = ctx.ops var h = tokenEmbed.bind(ctx) val nsk = ArrayList>(layers.size) val nsv = ArrayList>(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 diff --git a/llm-inference/moonshine/src/jvmTest/kotlin/sk/ainet/models/moonshine/MoonshineV2DecoderBakeTest.kt b/llm-inference/moonshine/src/jvmTest/kotlin/sk/ainet/models/moonshine/MoonshineV2DecoderBakeTest.kt index e821ef5..0593ca0 100644 --- a/llm-inference/moonshine/src/jvmTest/kotlin/sk/ainet/models/moonshine/MoonshineV2DecoderBakeTest.kt +++ b/llm-inference/moonshine/src/jvmTest/kotlin/sk/ainet/models/moonshine/MoonshineV2DecoderBakeTest.kt @@ -68,20 +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=: 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 -> - // MOONSHINE_V2_TRUE_DYNAMIC=1 traces the self-cache seq dim (grows per decode step) AND the - // cross-cache frames dim (varies per utterance) as Dim.DYNAMIC, so ONE compiled vmfb serves - // every decode position and every encoder length. Requires the dynamic-safe tracer + emitter. - val trueDynamic = System.getenv("MOONSHINE_V2_TRUE_DYNAMIC") == "1" val past = if (trueDynamic) sk.ainet.lang.tensor.Dim.DYNAMIC else (System.getenv("DEC_PAST")?.toInt() ?: 1) - val framesDim = if (trueDynamic) sk.ainet.lang.tensor.Dim.DYNAMIC else frames + 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, framesDim, hd)) }, List(L) { voidF32(Shape(1, nh, framesDim, hd)) }, c, + List(L) { voidF32(Shape(1, nh, crossFrames, hd)) }, List(L) { voidF32(Shape(1, nh, crossFrames, hd)) }, + c, crossMask, ) } } diff --git a/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/MultiHeadAttention.kt b/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/MultiHeadAttention.kt index fb0184e..ca2613e 100644 --- a/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/MultiHeadAttention.kt +++ b/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/MultiHeadAttention.kt @@ -222,6 +222,7 @@ public class MultiHeadAttention( input: Tensor, encoderMemory: Tensor?, ctx: ExecutionContext, + crossMask: Tensor? = null, ): AttentionKV { val boundInput = input.bind(ctx) return if (encoderMemory == null) { @@ -231,7 +232,7 @@ public class MultiHeadAttention( "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( kvInput: Tensor, isCrossAttention: Boolean, ctx: ExecutionContext, + crossMask: Tensor? = null, ): AttentionKV { val ops = ctx.ops val scale = attentionScale ?: (1.0f / sqrt(headDim.toFloat())) @@ -434,12 +436,14 @@ public class MultiHeadAttention( // 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 )