feat(gemma): export one FunctionGemma decode step for the Torq NPU

CHANGELOG.md

@@ -11,6 +11,32 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+Against engine **0.63.0** (SKaiNET#1478 and SKaiNET#1479, block-quant export).
+
+### Added
+
+- **Decode step for an NPU with block-quantized weights** (`llm-core`, jvm): `DecodeStepExport.export(spec, dtype) { inputs, ctx -> … }`
+  traces one single-token step of any KV-cache decoder with a static past, the embedding lookup left to the host, keeps
+  Q8_0 / Q4_0 weights packed and folds them with the engine's `BlockQuantMatmulFusionPass`, and returns the StableHLO
+  module with every argument and result named (`embedded`, `cosGlobal`/`sinGlobal`, `cosSliding`/`sinSliding`,
+  `l<i>.k`/`l<i>.v`, `logits`). FunctionGemma 270M (Q8_0 projections, Q4_0 copy of the tied head, past 16) compiles
+  with Torq 2.3.0 to one dispatch on its `mixed_quant_matmul` kernel and runs on an SL2610 at 162 ms per step with the
+  same top-1 as the f32 reference (`TorqDecodeStepExportTest`, gated on `TORQ_DECODE_OUT` and `GEMMA_GGUF`).
+- `forwardEmbeddedWithPast(embedded, rope, selfKIn, selfVIn, ctx, lmHeadWeight = null)` on `GemmaModel` and
+  `DecoderKvModel`: the decode step from an already embedded token, optionally projecting with a replacement output
+  weight (e.g. a differently encoded copy of a tied embedding). `forwardWithPast` delegates to it.
+
+### Changed
+
+- **The with-past inputs and outputs are shared** (`llm-core`, `sk.ainet.lang.nn.dsl.decoder`): `RopeTables`,
+  `RopeCosSin` (tables per RoPE base: `global`, and `sliding` for Gemma 3), `ChunkContext` (`rope` plus a mask per
+  layer type), `KvStepOutput` and `buildRopeCosSin(attention, position, dtype, ctx, rows)` replace
+  `GemmaModel.RopeCosSin` / `ChunkContext` / `GemmaPrefillOutput` / `GemmaWithPastOutput` and
+  `DecoderKvModel.RopeCosSin` / `ChunkContext` / `PrefillOutput` / `WithPastOutput`. Both models route per layer by
+  `MultiHeadAttention.slidingWindow`. `RopeCosSin(cos, sin)` and `RopeCosSin(cosGlobal, sinGlobal, cosSliding,
+  sinSliding)` keep the old call shapes; `ChunkContext` now takes a `RopeCosSin`. `:llm-inference:gemma` exposes
+  `:llm-core` as `api`.
+
 ## [0.63.3] — 2026-10-28
 
 Against engine **0.62.0**, transformers-only and additive. The Android runtime (`runtime-iree-android`) gains an

gradle/libs.versions.toml

@@ -1,5 +1,5 @@
 [versions]
-skainet = "0.62.0"
+skainet = "0.63.0"
 agp = "9.4.1"
 jacksonDatabind = "2.22.4"
 jsonSchemaValidator = "3.0.8"

llm-core/build.gradle.kts

@@ -54,7 +54,12 @@ kotlin {
             implementation(libs.kotlinx.serialization.json)
         }
 
-        val jvmMain by getting
+        val jvmMain by getting {
+            dependencies {
+                // DecodeStepExport: trace -> ComputeGraph -> StableHLO (compile-hlo publishes no JS variant).
+                implementation(libs.skainet.compile.hlo)
+            }
+        }
 
         val jvmTest by getting {
             dependencies {

llm-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/decoder/DecoderKvModel.kt

@@ -6,7 +6,6 @@ import sk.ainet.lang.nn.Module
 import sk.ainet.lang.nn.layers.EmbeddingAdapter
 import sk.ainet.lang.nn.normalization.RMSNormalization
 import sk.ainet.lang.nn.transformer.MultiHeadAttention
-import sk.ainet.lang.nn.transformer.RoPEMode
 import sk.ainet.lang.nn.transformer.SwiGLUFFN
 import sk.ainet.lang.nn.transformer.VoidDense
 import sk.ainet.lang.nn.transformer.linearProject
@@ -69,41 +68,6 @@ public class DecoderKvModel<T : DType, V>(
             ?: error("DecoderKvModel: expected exactly one top-level VoidDense (lm head / output)")
     }
 
-    /** logits `[seq, vocab]` + per-layer initial self K/V (`[1, nKVHeads, seq, headDim]`). */
-    public class PrefillOutput<T : DType, V>(
-        public val logits: Tensor<T, V>,
-        public val selfK: List<Tensor<T, V>>,
-        public val selfV: List<Tensor<T, V>>,
-    )
-
-    /** logits `[1, vocab]` + per-layer extended self K/V (one call's worth of new rows appended). */
-    public class WithPastOutput<T : DType, V>(
-        public val logits: Tensor<T, V>,
-        public val selfK: List<Tensor<T, V>>,
-        public val selfV: List<Tensor<T, V>>,
-    )
-
-    /** Per-step RoPE cos/sin (`[1, headDim]`) — the model's one RoPE base (see [buildRopeCosSin]). */
-    public class RopeCosSin<T : DType, V>(
-        public val cos: Tensor<T, V>,
-        public val sin: Tensor<T, V>,
-    )
-
-    /**
-     * [forwardPrefillWithPast]'s per-chunk inputs: RoPE cos/sin tables `[C, headDim]` for absolute
-     * positions `past .. past+C-1` (split-half or interleaved layout per the model's [RoPEMode]),
-     * and the additive causal+padding mask `[1, M, C, past+C]` (0 = attend, -1e30 = masked), where
-     * `M` is 1 (head-shared) or nHeads; [forwardWithPast] never needs a mask (a single query attends
-     * the whole, already-causal cache). Under GQA with a dynamic `past` use `M = 1`: SKaiNET 0.57.0
-     * broadcasts a head-shared mask with the `dynamic_broadcast_in_dim` form IREE 3.11 lowers, while
-     * a per-head one needs `dynamic_reshape`, which IREE 3.11 does not lower (SKaiNET#1302).
-     */
-    public class ChunkContext<T : DType, V>(
-        public val cos: Tensor<T, V>,
-        public val sin: Tensor<T, V>,
-        public val mask: Tensor<T, V>,
-    )
-
     private class BlockRefs<T : DType, V>(
         val attnNorm: RMSNormalization<T, V>,
         val mha: MultiHeadAttention<T, V>,
@@ -145,32 +109,50 @@ public class DecoderKvModel<T : DType, V>(
 
     /**
      * KV-cache DECODE step: one new [tokenId] (`[1]` or `[1,1]`) at the position encoded by [rope]
-     * (see [buildRopeCosSin]); the incoming per-layer self cache ([selfKIn]/[selfVIn],
-     * `[1, nKVHeads, past, headDim]`) → logits `[1, vocab]` + the per-layer extended self cache.
+     * (see [buildRopeCosSin]; one RoPE base, so only [RopeCosSin.global] is read); the incoming
+     * per-layer self cache ([selfKIn]/[selfVIn], `[1, nKVHeads, past, headDim]`) → logits
+     * `[1, vocab]` + the per-layer extended self cache.
      */
     public fun forwardWithPast(
         tokenId: Tensor<T, V>,
         rope: RopeCosSin<T, V>,
         selfKIn: List<Tensor<T, V>>,
         selfVIn: List<Tensor<T, V>>,
         ctx: ExecutionContext,
-    ): WithPastOutput<T, V> {
+    ): KvStepOutput<T, V> = forwardEmbeddedWithPast(tokenEmbedding.forward(tokenId, ctx), rope, selfKIn, selfVIn, ctx)
+
+    /**
+     * [forwardWithPast] from the token embedding [embedded] (`[1, hidden]`) instead of a token id, so
+     * a caller can do the lookup itself — e.g. a compiled decode step that keeps the `[vocab, hidden]`
+     * table out of the graph. [lmHeadWeight] (`[vocab, hidden]`) replaces the output projection
+     * weight when given: a tied model can project with a differently encoded copy of its embedding.
+     */
+    public fun forwardEmbeddedWithPast(
+        embedded: Tensor<T, V>,
+        rope: RopeCosSin<T, V>,
+        selfKIn: List<Tensor<T, V>>,
+        selfVIn: List<Tensor<T, V>>,
+        ctx: ExecutionContext,
+        lmHeadWeight: Tensor<T, V>? = null,
+    ): KvStepOutput<T, V> {
         val ops = ctx.ops
-        var h = tokenEmbedding.forward(tokenId, ctx)
+        var h = embedded
         val nsk = ArrayList<Tensor<T, V>>(blocks.size)
         val nsv = ArrayList<Tensor<T, V>>(blocks.size)
         for ((i, block) in blocks.withIndex()) {
             val r = refsFor(block)
             val sn = r.attnNorm.forward(h, ctx)
-            val (attnOut, fullK, fullV) = attnWithPast(r.mha, sn, rope.cos, rope.sin, selfKIn[i], selfVIn[i], ctx)
+            val tables = rope.forLayer(r.mha)
+            val (attnOut, fullK, fullV) = attnWithPast(r.mha, sn, tables.cos, tables.sin, selfKIn[i], selfVIn[i], ctx)
             val h1 = ops.add(h, attnOut)
             val ffnOut = r.ffn.forward(r.ffnNorm.forward(h1, ctx), ctx)
             h = ops.add(h1, ffnOut)
             nsk += ops.reshape(fullK, fullK.shape)   // identity → distinct graph output node
             nsv += ops.reshape(fullV, fullV.shape)
         }
-        val logits = lmHead.forward(outputNorm.forward(h, ctx), ctx)
-        return WithPastOutput(logits, nsk, nsv)
+        val normed = outputNorm.forward(h, ctx)
+        val logits = if (lmHeadWeight != null) ops.matmulWeightTransposed(normed, lmHeadWeight) else lmHead.forward(normed, ctx)
+        return KvStepOutput(logits, nsk, nsv)
     }
 
     private fun attnWithPast(
@@ -211,6 +193,12 @@ public class DecoderKvModel<T : DType, V>(
      * by the one-hot [selectAt] (`[1, C]`), plus the per-layer cache extended by all C positions
      * (the caller slices the padding off). This is what makes an utterance cost one call instead
      * of C single-token [forwardWithPast] steps.
+     *
+     * [chunk]'s tables follow the model's [sk.ainet.lang.nn.transformer.RoPEMode]; its mask is
+     * `[1, M, C, past+C]` with `M` 1 (head-shared) or nHeads. Under GQA with a dynamic `past` use
+     * `M = 1`: SKaiNET 0.57.0 broadcasts a head-shared mask with the `dynamic_broadcast_in_dim`
+     * form IREE 3.11 lowers, while a per-head one needs `dynamic_reshape`, which IREE 3.11 does
+     * not lower (SKaiNET#1302).
      */
     public fun forwardPrefillWithPast(
         tokens: Tensor<T, V>,
@@ -219,15 +207,16 @@ public class DecoderKvModel<T : DType, V>(
         selfKIn: List<Tensor<T, V>>,
         selfVIn: List<Tensor<T, V>>,
         ctx: ExecutionContext,
-    ): WithPastOutput<T, V> {
+    ): KvStepOutput<T, V> {
         val ops = ctx.ops
         var h = tokenEmbedding.forward(tokens, ctx)
         val nsk = ArrayList<Tensor<T, V>>(blocks.size)
         val nsv = ArrayList<Tensor<T, V>>(blocks.size)
         for ((i, block) in blocks.withIndex()) {
             val r = refsFor(block)
             val sn = r.attnNorm.forward(h, ctx)
-            val (attnOut, fullK, fullV) = attnWithPastChunk(r.mha, sn, chunk.cos, chunk.sin, chunk.mask, selfKIn[i], selfVIn[i], ctx)
+            val tables = chunk.rope.forLayer(r.mha)
+            val (attnOut, fullK, fullV) = attnWithPastChunk(r.mha, sn, tables.cos, tables.sin, chunk.maskFor(r.mha), selfKIn[i], selfVIn[i], ctx)
             val h1 = ops.add(h, attnOut)
             val ffnOut = r.ffn.forward(r.ffnNorm.forward(h1, ctx), ctx)
             h = ops.add(h1, ffnOut)
@@ -237,7 +226,7 @@ public class DecoderKvModel<T : DType, V>(
         val normed = outputNorm.forward(h, ctx)
         val hAt = ops.matmul(selectAt, normed)   // [1, C] x [C, hidden] -> [1, hidden]
         val logits = lmHead.forward(hAt, ctx)
-        return WithPastOutput(logits, nsk, nsv)
+        return KvStepOutput(logits, nsk, nsv)
     }
 
     /** [attnWithPast] for a C-row chunk: heads-first `[heads, C, headDim]` projections, RoPE from
@@ -287,7 +276,7 @@ public class DecoderKvModel<T : DType, V>(
      * for the ONE position picked by the one-hot [selectAt] (`[1, seq]`), plus the per-layer initial
      * self K/V that seeds [forwardWithPast] / [forwardPrefillWithPast]. Stateless one-pass prefill.
      */
-    public fun forwardPrefillAt(input: Tensor<T, V>, selectAt: Tensor<T, V>, ctx: ExecutionContext): PrefillOutput<T, V> {
+    public fun forwardPrefillAt(input: Tensor<T, V>, selectAt: Tensor<T, V>, ctx: ExecutionContext): KvStepOutput<T, V> {
         stripCaches()
         val ops = ctx.ops
         var h = tokenEmbedding.forward(input, ctx)
@@ -305,7 +294,7 @@ public class DecoderKvModel<T : DType, V>(
         val normed = outputNorm.forward(h, ctx)
         val hAt = ops.matmul(selectAt, if (normed.rank == 3) ops.squeeze(normed, dim = 0) else normed)   // [1, hidden]
         val logits = lmHead.forward(hAt, ctx)                                                            // [1, vocab]
-        return PrefillOutput(logits, selfK, selfV)
+        return KvStepOutput(logits, selfK, selfV)
     }
 
     /**
@@ -314,13 +303,6 @@ public class DecoderKvModel<T : DType, V>(
      * the whole model). Fed to [forwardWithPast] as a runtime input so the decode graph carries no
      * compile-time position.
      */
-    public fun buildRopeCosSin(position: Int, ctx: ExecutionContext): RopeCosSin<T, V> {
-        val rope = refsFor(blocks.first()).mha.rope ?: error("DecoderKvModel: no RoPE on the first block")
-        val (c, s) = when (rope.mode) {
-            RoPEMode.SPLIT_HALF -> rope.buildSplitHalfCosSin(position, 1)
-            RoPEMode.INTERLEAVED -> rope.buildInterleavedCosSin(position, 1)
-        }
-        val shape = Shape(1, rope.headDim)
-        return RopeCosSin(cos = ctx.fromFloatArray(shape, dtype, c), sin = ctx.fromFloatArray(shape, dtype, s))
-    }
+    public fun buildRopeCosSin(position: Int, ctx: ExecutionContext): RopeCosSin<T, V> =
+        buildRopeCosSin(listOf(refsFor(blocks.first()).mha), position, dtype, ctx)
 }

llm-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/decoder/KvStep.kt

@@ -0,0 +1,99 @@
+package sk.ainet.lang.nn.dsl.decoder
+
+import sk.ainet.context.ExecutionContext
+import sk.ainet.lang.nn.transformer.MultiHeadAttention
+import sk.ainet.lang.nn.transformer.RoPE
+import sk.ainet.lang.nn.transformer.RoPEMode
+import sk.ainet.lang.tensor.Shape
+import sk.ainet.lang.tensor.Tensor
+import sk.ainet.lang.types.DType
+import kotlin.reflect.KClass
+
+/*
+ * Inputs and outputs of the KV-cache with-past graphs, shared by every decoder that has them
+ * ([DecoderKvModel], `GemmaModel`). Generic over the model's dtype; a model with one RoPE base and
+ * full attention everywhere (Llama, Qwen) sets only the global entries, Gemma 3 also the sliding
+ * ones. The model routes per layer by `MultiHeadAttention.slidingWindow`.
+ */
+
+/** One RoPE base's sign-baked cos/sin tables `[rows, headDim]`, in that base's [RoPEMode] layout. */
+public class RopeTables<T : DType, V>(
+    public val cos: Tensor<T, V>,
+    public val sin: Tensor<T, V>,
+)
+
+/**
+ * RoPE tables per RoPE base, fed to a with-past graph as runtime inputs so the graph carries no
+ * compile-time position. [sliding] is [global] for a model with one base.
+ */
+public class RopeCosSin<T : DType, V>(
+    public val global: RopeTables<T, V>,
+    public val sliding: RopeTables<T, V> = global,
+) {
+    public constructor(cos: Tensor<T, V>, sin: Tensor<T, V>) : this(RopeTables(cos, sin))
+
+    public constructor(
+        cosGlobal: Tensor<T, V>,
+        sinGlobal: Tensor<T, V>,
+        cosSliding: Tensor<T, V>,
+        sinSliding: Tensor<T, V>,
+    ) : this(RopeTables(cosGlobal, sinGlobal), RopeTables(cosSliding, sinSliding))
+
+    /** Whether the sliding layers use their own base (a separate pair of graph inputs). */
+    public val hasSlidingBase: Boolean get() = sliding !== global
+
+    /** The tables [mha]'s layer rotates with. */
+    public fun forLayer(mha: MultiHeadAttention<T, V>): RopeTables<T, V> = if (mha.slidingWindow == null) global else sliding
+}
+
+/**
+ * Per-chunk inputs of a chunk prefill against a cache: RoPE tables `[C, headDim]` for the absolute
+ * positions `past .. past+C-1`, and the additive attention masks `[1, M, C, past+C]` (0 = attend,
+ * -1e30 = masked) per layer type. The masks carry everything position-dependent (the causal band,
+ * the zero padding, a sliding window), so the graph itself stays position-agnostic.
+ */
+public class ChunkContext<T : DType, V>(
+    public val rope: RopeCosSin<T, V>,
+    public val maskGlobal: Tensor<T, V>,
+    public val maskSliding: Tensor<T, V> = maskGlobal,
+) {
+    /** The mask [mha]'s layer attends with. */
+    public fun maskFor(mha: MultiHeadAttention<T, V>): Tensor<T, V> = if (mha.slidingWindow == null) maskGlobal else maskSliding
+}
+
+/**
+ * Logits plus the per-layer self K/V (`[1, nKVHeads, S, headDim]`) of a KV-cache graph: the
+ * initial cache after a prefill, or the cache extended by this call's rows after a with-past step.
+ */
+public class KvStepOutput<T : DType, V>(
+    public val logits: Tensor<T, V>,
+    public val selfK: List<Tensor<T, V>>,
+    public val selfV: List<Tensor<T, V>>,
+)
+
+/**
+ * Build the RoPE tables for [rows] positions from [position] for the layers in [attention]: one
+ * pair from a global layer's RoPE and, when the sliding layers have a RoPE of their own, one from a
+ * sliding layer's. Each table follows its RoPE's [RoPEMode].
+ */
+public fun <T : DType, V> buildRopeCosSin(
+    attention: List<MultiHeadAttention<T, V>>,
+    position: Int,
+    dtype: KClass<T>,
+    ctx: ExecutionContext,
+    rows: Int = 1,
+): RopeCosSin<T, V> {
+    val globalRope = attention.firstOrNull { it.slidingWindow == null }?.rope
+    val slidingRope = attention.firstOrNull { it.slidingWindow != null }?.rope
+    val first = globalRope ?: slidingRope ?: error("buildRopeCosSin: no attention layer has a RoPE")
+    fun tables(rope: RoPE<T, V>): RopeTables<T, V> {
+        val (c, s) = when (rope.mode) {
+            RoPEMode.SPLIT_HALF -> rope.buildSplitHalfCosSin(position, rows)
+            RoPEMode.INTERLEAVED -> rope.buildInterleavedCosSin(position, rows)
+        }
+        val shape = Shape(rows, rope.headDim)
+        return RopeTables(ctx.fromFloatArray(shape, dtype, c), ctx.fromFloatArray(shape, dtype, s))
+    }
+    val global = tables(first)
+    return if (slidingRope == null || slidingRope === first) RopeCosSin(global) else RopeCosSin(global, tables(slidingRope))
+}