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 )