Skip to content

feat(moonshine,transformer-core): fixed-max cross-memory padding mask for streaming decode #267

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
michalharakal merged 1 commit into from
Jul 30, 2026
Merged

feat(moonshine,transformer-core): fixed-max cross-memory padding mask for streaming decode #267

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 12 additions & 7 deletions llm-inference/moonshine/src/commonMain/kotlin/sk/ainet/models/moonshine/MoonshineDecoder.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 ---
Expand All @@ -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
Expand All @@ -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)
}
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand All @@ -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
Expand Down
27 changes: 20 additions & 7 deletions ...ence/moonshine/src/jvmTest/kotlin/sk/ainet/models/moonshine/MoonshineV2DecoderBakeTest.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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=<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 ->
// 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,
)
}
}
Expand Down
10 changes: 7 additions & 3 deletions transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/MultiHeadAttention.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand All @@ -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)
}
}

Expand Down Expand Up @@ -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()))
Expand Down Expand Up @@ -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
)
Expand Down