From b5b239862c1612de7450718fd3791224aec0a693 Mon Sep 17 00:00:00 2001 From: Michal Harakal <471011+michalharakal@users.noreply.github.com> Date: Thu, 5 Nov 2026 22:57:02 +0100 Subject: [PATCH 1/4] feat(gemma): export one FunctionGemma decode step for the Torq NPU GemmaModel.forwardEmbeddedWithPast runs the decode step from an already embedded token, so the embedding lookup can stay on the host, and takes an optional lmHeadWeight to project the logits with a different (e.g. re-quantized) copy of the tied embedding. forwardWithPast delegates to it. TorqDecodeStepExportTest (gated on TORQ_DECODE_OUT and GEMMA_GGUF) traces one step with a static past of 16 positions, keeps the Q8_0 projections packed, folds all 127 of them with BlockQuantMatmulFusionPass (the vocabulary projection as a Q4_0 copy of the embedding) and writes the StableHLO module plus the argument/result roles. Converted to bf16, it compiles with Torq 2.3.0 to one dispatch on the mixed_quant_matmul kernel and runs on an SL2610 at 162 ms/step with the same top-1 as the f32 reference. Needs a SKaiNET release with SKaiNET-developers/SKaiNET#1478 and #1479. Co-Authored-By: Claude Opus 5.5 --- .../functiongemma/TorqDecodeStepExportTest.kt | 135 ++++++++++++++++++ .../sk/ainet/models/gemma/GemmaModel.kt | 27 +++- 2 files changed, 160 insertions(+), 2 deletions(-) create mode 100644 llm-inference/functiongemma/src/jvmTest/kotlin/sk/ainet/models/functiongemma/TorqDecodeStepExportTest.kt diff --git a/llm-inference/functiongemma/src/jvmTest/kotlin/sk/ainet/models/functiongemma/TorqDecodeStepExportTest.kt b/llm-inference/functiongemma/src/jvmTest/kotlin/sk/ainet/models/functiongemma/TorqDecodeStepExportTest.kt new file mode 100644 index 00000000..53980cb7 --- /dev/null +++ b/llm-inference/functiongemma/src/jvmTest/kotlin/sk/ainet/models/functiongemma/TorqDecodeStepExportTest.kt @@ -0,0 +1,135 @@ +package sk.ainet.models.functiongemma + +import kotlinx.coroutines.runBlocking +import org.junit.jupiter.api.Assumptions.assumeTrue +import sk.ainet.compile.hlo.toStableHlo +import sk.ainet.compile.opt.passes.BlockQuantMatmulFusionPass +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.context.ExecutionContext +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.lang.graph.DefaultExecutionTape +import sk.ainet.lang.graph.DefaultGraphExecutionContext +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.VoidOpsTensor +import sk.ainet.io.gguf.StreamingGGUFReader +import sk.ainet.lang.tensor.data.Q4_0Quantizer +import sk.ainet.lang.tensor.data.Q8_0BlockTensorData +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.tensor.operators.bind +import sk.ainet.lang.tensor.ops.VoidTensorOps +import sk.ainet.lang.trace.PackedConstantHandling +import sk.ainet.lang.types.FP32 +import sk.ainet.models.gemma.GemmaModel +import sk.ainet.models.gemma.GemmaNetworkLoader +import sk.ainet.models.gemma.GemmaWeightLoader +import sk.ainet.tape.Execution +import java.io.File +import kotlin.test.Test +import kotlin.test.assertTrue + +/** + * One FunctionGemma decode step for an NPU with block-quantized weights (Torq `mixed_quant_matmul`): + * Q8_0 projections kept packed and folded by [BlockQuantMatmulFusionPass], the tied vocabulary + * projection as a Q4_0 copy of the embedding, the embedding lookup left to the caller, and a static + * past of [PAST] positions. + * + * Inputs, in order: the scaled embedding `[1, 640]`, cos/sin `[1, 256]` for the global and the + * sliding RoPE base, then K and V `[1, 1, PAST, 256]` per layer. Outputs: logits `[1, vocab]`, then + * the extended K and V per layer. + * + * Runs only with `TORQ_DECODE_OUT` set (the output directory) and the Q8_0 checkpoint in `GEMMA_GGUF`. + */ +class TorqDecodeStepExportTest { + + @Test + fun exportDecodeStepWithPackedWeights() = runBlocking { + val out = System.getenv("TORQ_DECODE_OUT") + assumeTrue(out != null, "set TORQ_DECODE_OUT to export the Torq decode step") + FunctionGemmaFixture.assumeRealCheckpointRunnable() + + val ctx = DirectCpuExecutionContext.create() + val weights = GemmaWeightLoader(randomAccessProvider = { JvmRandomAccessSource.open(FunctionGemmaFixture.gguf) }) + .loadToMapStreaming(ctx, FP32::class) + @Suppress("UNCHECKED_CAST") + val model = GemmaNetworkLoader.fromWeights(ctx, weights, FP32::class) as GemmaModel + val md = weights.metadata + val nLayers = md.blockCount + val headDim = md.getHeadDim(0) + val nKV = md.kvHeadCount + + val embedding = weights.tensors.getValue("token_embd.weight") + println("token_embd data: ${embedding.data::class.qualifiedName}, q proj: ${weights.tensors.getValue("blk.0.attn_q.weight").data::class.qualifiedName}") + val vocab = embedding.shape[0] + val hidden = embedding.shape[1] + // The tied table as stored (Q8_0), decoded once and re-encoded as Q4_0: the Torq compiler + // hangs on the Q8 vocabulary projection. + val rawEmbedding = StreamingGGUFReader.open(JvmRandomAccessSource.open(FunctionGemmaFixture.gguf)).use { it.loadTensor("token_embd.weight") } + val dense = Q8_0BlockTensorData(Shape(vocab, hidden), rawEmbedding).toFloatArray() + @Suppress("UNCHECKED_CAST") + val lmHeadQ4 = ctx.fromData(Q4_0Quantizer.quantize(dense, Shape(vocab, hidden)) as TensorData, FP32::class) + + // Inputs are bound to the tape context: a bare VoidOpsTensor runs operator overloads + // (`x * x` in RMSNorm) on its own untraced ops, and the result is captured as a constant. + val tapeCtx = DefaultGraphExecutionContext.tape(baseOps = VoidTensorOps()) + fun input(shape: Shape) = voidF32(shape).bind(tapeCtx) + val embedded = input(Shape(1, hidden)) + val cosG = input(Shape(1, headDim)); val sinG = input(Shape(1, headDim)) + val cosS = input(Shape(1, headDim)); val sinS = input(Shape(1, headDim)) + val kIn = List(nLayers) { input(Shape(1, nKV, PAST, headDim)) } + val vIn = List(nLayers) { input(Shape(1, nKV, PAST, headDim)) } + + var outputs: GemmaModel.GemmaWithPastOutput? = null + val tape = tapeCtx.record { + val ct = (this as DefaultGraphExecutionContext).currentTape ?: error("no tape") + Execution.tapeStack.pushTape(ct) + try { + outputs = model.forwardEmbeddedWithPast( + embedded, GemmaModel.RopeCosSin(cosG, sinG, cosS, sinS), kIn, vIn, this as ExecutionContext, lmHeadWeight = lmHeadQ4, + ) + } finally { + Execution.tapeStack.popTape() + } + }.first + val graph = (tape as DefaultExecutionTape).toComputeGraph( + synthesizeExternalInputs = true, packedConstants = PackedConstantHandling.KEEP_PACKED, + ) + val fused = BlockQuantMatmulFusionPass().apply(graph) + fused.diagnostics.filter { it.startsWith("BLOCK_QUANT_SKIPPED") }.forEach(::println) + val folded = fused.diagnostics.count { it.startsWith("BLOCK_QUANT_FUSED") } + println("folded $folded block-quant matmuls") + assertTrue(folded == nLayers * 7 + 1, "expected every projection and the vocabulary projection folded, got $folded") + + val module = toStableHlo(fused.graph, "decode_step") + File(out).mkdirs() + File(out, "decode_step.mlir").writeText(module.content) + // Which function argument / result is which tensor: the emitter orders them by graph topology. + fun id(t: Tensor) = tapeCtx.session.refOf(t as Tensor<*, *>).id + val inRoles = mutableMapOf(id(embedded) to "embedded", id(cosG) to "cosG", id(sinG) to "sinG", id(cosS) to "cosS", id(sinS) to "sinS") + kIn.forEachIndexed { i, t -> inRoles[id(t)] = "k$i" } + vIn.forEachIndexed { i, t -> inRoles[id(t)] = "v$i" } + val o = outputs!! + val outRoles = mutableMapOf(id(o.logits) to "logits") + o.selfK.forEachIndexed { i, t -> outRoles[id(t)] = "k$i" } + o.selfV.forEachIndexed { i, t -> outRoles[id(t)] = "v$i" } + File(out, "io.txt").writeText( + module.inputSpecs.mapIndexed { i, s -> "in $i ${inRoles[s.name] ?: "?"} ${s.name} ${s.shape}" }.joinToString("\n") + "\n" + + module.outputSpecs.mapIndexed { i, s -> "out $i ${outRoles[s.name] ?: "?"} ${s.name} ${s.shape}" }.joinToString("\n") + "\n" + ) + } + + private fun voidF32(shape: Shape): Tensor = VoidOpsTensor( + object : TensorData { + override val shape: Shape = shape + // Unreadable, like the shape-only results of VoidTensorOps: an eager shortcut that + // reads its input (RMSNormalization's fused path) must not bake this input in as a constant. + override fun get(vararg indices: Int): Float = throw UnsupportedOperationException("graph input has no values") + override fun set(vararg indices: Int, value: Float) {} + }, + FP32::class, + ) + + private companion object { + const val PAST = 16 + } +} diff --git a/llm-inference/gemma/src/commonMain/kotlin/sk/ainet/models/gemma/GemmaModel.kt b/llm-inference/gemma/src/commonMain/kotlin/sk/ainet/models/gemma/GemmaModel.kt index 81e33bfa..5d1fcbe3 100644 --- a/llm-inference/gemma/src/commonMain/kotlin/sk/ainet/models/gemma/GemmaModel.kt +++ b/llm-inference/gemma/src/commonMain/kotlin/sk/ainet/models/gemma/GemmaModel.kt @@ -547,7 +547,28 @@ public class GemmaModel( require(ple == null) { "GemmaModel.forwardWithPast: PLE models not supported by the KV-cache path yet" } val ops = ctx.ops val rawEmbeds = tokenEmbedding.forward(tokenId, ctx) - var h = if (embedScale != 1f) ops.mulScalar(rawEmbeds, embedScale) else rawEmbeds + val h = if (embedScale != 1f) ops.mulScalar(rawEmbeds, embedScale) else rawEmbeds + return forwardEmbeddedWithPast(h, rope, selfKIn, selfVIn, ctx) + } + + /** + * [forwardWithPast] from the scaled 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 + * model's output projection weight when given: a tied model can project with a differently + * encoded copy of its embedding. + */ + public fun forwardEmbeddedWithPast( + embedded: Tensor, + rope: RopeCosSin, + selfKIn: List>, + selfVIn: List>, + ctx: ExecutionContext, + lmHeadWeight: Tensor? = null, + ): GemmaWithPastOutput { + require(ple == null) { "GemmaModel.forwardEmbeddedWithPast: PLE models not supported by the KV-cache path yet" } + val ops = ctx.ops + var h = embedded val nsk = ArrayList>(blocks.size) val nsv = ArrayList>(blocks.size) for ((i, block) in blocks.withIndex()) { @@ -567,7 +588,9 @@ public class GemmaModel( nsk += ops.reshape(fullK, fullK.shape) // identity → distinct graph output node nsv += ops.reshape(fullV, fullV.shape) } - val logits = applySoftcap(lmHead.forward(outputNorm.forward(h, ctx), ctx), ctx) + val normed = outputNorm.forward(h, ctx) + val projected = if (lmHeadWeight != null) ops.matmulWeightTransposed(normed, lmHeadWeight) else lmHead.forward(normed, ctx) + val logits = applySoftcap(projected, ctx) return GemmaWithPastOutput(logits, nsk, nsv) } From 53c1af82cbe7bd8a8550049b7c2d5c550697ac92 Mon Sep 17 00:00:00 2001 From: Michal Harakal <471011+michalharakal@users.noreply.github.com> Date: Fri, 6 Nov 2026 10:09:45 +0100 Subject: [PATCH 2/4] refactor(llm-core): share the with-past types and the decode-step export GemmaModel and DecoderKvModel each had their own RopeCosSin, ChunkContext and prefill / with-past output classes. They now share llm-core's RopeTables, RopeCosSin (tables per RoPE base: global, plus sliding for Gemma 3), ChunkContext (rope plus a mask per layer type), KvStepOutput and buildRopeCosSin, all generic over the model's DType; both models route per layer by MultiHeadAttention.slidingWindow. The old call shapes RopeCosSin(cos, sin) and RopeCosSin(cosG, sinG, cosS, sinS) still work. DecoderKvModel gains forwardEmbeddedWithPast like GemmaModel, and the decode-step export moves out of the FunctionGemma test into llm-core's DecodeStepExport: it binds the graph inputs to the trace, keeps packed weights, runs BlockQuantMatmulFusionPass and names every argument and result. TorqDecodeStepExportTest keeps only the Gemma-specific parts; its module is unchanged apart from SSA names. :llm-inference:gemma now exposes :llm-core as api, since its public with-past methods use these types. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 26 +++ llm-core/build.gradle.kts | 7 +- .../lang/nn/dsl/decoder/DecoderKvModel.kt | 96 +++++------ .../sk/ainet/lang/nn/dsl/decoder/KvStep.kt | 99 ++++++++++++ .../lang/nn/dsl/decoder/DecodeStepExport.kt | 150 ++++++++++++++++++ .../lang/nn/dsl/decoder/DecoderKvModelTest.kt | 47 +++++- .../FunctionGemmaExportHarness.kt | 6 +- .../functiongemma/TorqDecodeStepExportTest.kt | 97 +++-------- llm-inference/gemma/build.gradle.kts | 3 +- .../sk/ainet/models/gemma/GemmaModel.kt | 105 +++--------- .../sk/ainet/models/qwen/QwenExportHarness.kt | 8 +- 11 files changed, 420 insertions(+), 224 deletions(-) create mode 100644 llm-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/decoder/KvStep.kt create mode 100644 llm-core/src/jvmMain/kotlin/sk/ainet/lang/nn/dsl/decoder/DecodeStepExport.kt diff --git a/CHANGELOG.md b/CHANGELOG.md index fc38e1e4..a53738d2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,32 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +Needs an engine release with 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.k`/`l.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 diff --git a/llm-core/build.gradle.kts b/llm-core/build.gradle.kts index 0b675807..77484e48 100644 --- a/llm-core/build.gradle.kts +++ b/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 { diff --git a/llm-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/decoder/DecoderKvModel.kt b/llm-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/decoder/DecoderKvModel.kt index b3245d6b..277516a4 100644 --- a/llm-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/decoder/DecoderKvModel.kt +++ b/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( ?: 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( - public val logits: Tensor, - public val selfK: List>, - public val selfV: List>, - ) - - /** logits `[1, vocab]` + per-layer extended self K/V (one call's worth of new rows appended). */ - public class WithPastOutput( - public val logits: Tensor, - public val selfK: List>, - public val selfV: List>, - ) - - /** Per-step RoPE cos/sin (`[1, headDim]`) — the model's one RoPE base (see [buildRopeCosSin]). */ - public class RopeCosSin( - public val cos: Tensor, - public val sin: Tensor, - ) - - /** - * [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( - public val cos: Tensor, - public val sin: Tensor, - public val mask: Tensor, - ) - private class BlockRefs( val attnNorm: RMSNormalization, val mha: MultiHeadAttention, @@ -145,8 +109,9 @@ public class DecoderKvModel( /** * 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, @@ -154,23 +119,40 @@ public class DecoderKvModel( selfKIn: List>, selfVIn: List>, ctx: ExecutionContext, - ): WithPastOutput { + ): KvStepOutput = 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, + rope: RopeCosSin, + selfKIn: List>, + selfVIn: List>, + ctx: ExecutionContext, + lmHeadWeight: Tensor? = null, + ): KvStepOutput { val ops = ctx.ops - var h = tokenEmbedding.forward(tokenId, ctx) + var h = embedded val nsk = ArrayList>(blocks.size) val nsv = ArrayList>(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( * 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, @@ -219,7 +207,7 @@ public class DecoderKvModel( selfKIn: List>, selfVIn: List>, ctx: ExecutionContext, - ): WithPastOutput { + ): KvStepOutput { val ops = ctx.ops var h = tokenEmbedding.forward(tokens, ctx) val nsk = ArrayList>(blocks.size) @@ -227,7 +215,8 @@ public class DecoderKvModel( 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( 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( * 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, selectAt: Tensor, ctx: ExecutionContext): PrefillOutput { + public fun forwardPrefillAt(input: Tensor, selectAt: Tensor, ctx: ExecutionContext): KvStepOutput { stripCaches() val ops = ctx.ops var h = tokenEmbedding.forward(input, ctx) @@ -305,7 +294,7 @@ public class DecoderKvModel( 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( * 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 { - 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 = + buildRopeCosSin(listOf(refsFor(blocks.first()).mha), position, dtype, ctx) } diff --git a/llm-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/decoder/KvStep.kt b/llm-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/decoder/KvStep.kt new file mode 100644 index 00000000..f9d58dac --- /dev/null +++ b/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( + public val cos: Tensor, + public val sin: Tensor, +) + +/** + * 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( + public val global: RopeTables, + public val sliding: RopeTables = global, +) { + public constructor(cos: Tensor, sin: Tensor) : this(RopeTables(cos, sin)) + + public constructor( + cosGlobal: Tensor, + sinGlobal: Tensor, + cosSliding: Tensor, + sinSliding: Tensor, + ) : 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): RopeTables = 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( + public val rope: RopeCosSin, + public val maskGlobal: Tensor, + public val maskSliding: Tensor = maskGlobal, +) { + /** The mask [mha]'s layer attends with. */ + public fun maskFor(mha: MultiHeadAttention): Tensor = 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( + public val logits: Tensor, + public val selfK: List>, + public val selfV: List>, +) + +/** + * 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 buildRopeCosSin( + attention: List>, + position: Int, + dtype: KClass, + ctx: ExecutionContext, + rows: Int = 1, +): RopeCosSin { + 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): RopeTables { + 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)) +} diff --git a/llm-core/src/jvmMain/kotlin/sk/ainet/lang/nn/dsl/decoder/DecodeStepExport.kt b/llm-core/src/jvmMain/kotlin/sk/ainet/lang/nn/dsl/decoder/DecodeStepExport.kt new file mode 100644 index 00000000..ac735b9e --- /dev/null +++ b/llm-core/src/jvmMain/kotlin/sk/ainet/lang/nn/dsl/decoder/DecodeStepExport.kt @@ -0,0 +1,150 @@ +package sk.ainet.lang.nn.dsl.decoder + +import sk.ainet.compile.hlo.toStableHlo +import sk.ainet.compile.opt.passes.BlockQuantMatmulFusionPass +import sk.ainet.context.ExecutionContext +import sk.ainet.lang.graph.DefaultExecutionTape +import sk.ainet.lang.graph.DefaultGraphExecutionContext +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.VoidOpsTensor +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.tensor.operators.bind +import sk.ainet.lang.tensor.ops.VoidTensorOps +import sk.ainet.lang.trace.PackedConstantHandling +import sk.ainet.lang.types.DType +import sk.ainet.tape.Execution +import kotlin.reflect.KClass + +/** + * The shape of a single-token decode step with a static past of [past] positions: the token + * embedding `[1, hidden]`, RoPE tables `[1, headDim]` (a second pair when [slidingRopeBase]), and + * K/V `[1, nKVHeads, past, headDim]` per layer. + */ +public class DecodeStepSpec( + public val hidden: Int, + public val headDim: Int, + public val nKVHeads: Int, + public val nLayers: Int, + public val past: Int, + public val slidingRopeBase: Boolean = false, +) + +/** The traced graph inputs a decode step runs on (see [DecodeStepExport.export]). */ +public class DecodeStepInputs( + public val embedded: Tensor, + public val rope: RopeCosSin, + public val selfK: List>, + public val selfV: List>, +) + +/** A function argument or result of an exported step: what it is ([role]) and its shape (`null` if dynamic). */ +public class DecodeStepPort(public val role: String, public val shape: List?) + +/** + * An exported decode step: the StableHLO module text, its arguments and results in function order, + * and what [BlockQuantMatmulFusionPass] did (`BLOCK_QUANT_FUSED` / `BLOCK_QUANT_SKIPPED` lines). + */ +public class ExportedDecodeStep( + public val mlir: String, + public val inputs: List, + public val outputs: List, + public val diagnostics: List, +) { + public val foldedBlockQuantMatmuls: Int get() = diagnostics.count { it.startsWith("BLOCK_QUANT_FUSED") } + + /** One line per port, `in ` then `out `. */ + public fun ioTable(): String = + (inputs.mapIndexed { i, p -> "in $i ${p.role} ${p.shape}" } + outputs.mapIndexed { i, p -> "out $i ${p.role} ${p.shape}" }) + .joinToString("\n", postfix = "\n") +} + +/** + * Trace one decode step of any KV-cache decoder ([DecoderKvModel.forwardEmbeddedWithPast], + * `GemmaModel.forwardEmbeddedWithPast`) and export it to StableHLO, for a compiler that wants the + * whole step as one graph with the embedding lookup left to the host. + * + * Ports are named after the model contracts: `embedded`, `cosGlobal`/`sinGlobal` (and + * `cosSliding`/`sinSliding`), `l.k`/`l.v` in; `logits`, `l.k`/`l.v` out. The emitter orders + * them by graph topology, so read the order from [ExportedDecodeStep.inputs]. + * + * With [PackedConstantHandling.KEEP_PACKED] (the default) block-quantized weights stay packed and + * [BlockQuantMatmulFusionPass] folds every `x · Wᵀ` over one; a packed weight it cannot fold fails + * the export. + */ +public object DecodeStepExport { + + public fun export( + spec: DecodeStepSpec, + dtype: KClass, + functionName: String = "decode_step", + packedConstants: PackedConstantHandling = PackedConstantHandling.KEEP_PACKED, + step: (DecodeStepInputs, ExecutionContext) -> KvStepOutput, + ): ExportedDecodeStep { + // Inputs are bound to the tape context: a bare VoidOpsTensor runs operator overloads + // (`x * x` in RMSNorm) on its own untraced ops, and the result is captured as a constant. + val tapeCtx = DefaultGraphExecutionContext.tape(baseOps = VoidTensorOps()) + val roles = LinkedHashMap() + fun input(role: String, shape: Shape): Tensor = + graphInput(shape, dtype).bind(tapeCtx).also { roles[tapeCtx.session.refOf(it as Tensor<*, *>).id] = role } + + val kvShape = Shape(1, spec.nKVHeads, spec.past, spec.headDim) + val ropeShape = Shape(1, spec.headDim) + val embedded = input("embedded", Shape(1, spec.hidden)) + val global = RopeTables(input("cosGlobal", ropeShape), input("sinGlobal", ropeShape)) + val rope = if (spec.slidingRopeBase) { + RopeCosSin(global, RopeTables(input("cosSliding", ropeShape), input("sinSliding", ropeShape))) + } else { + RopeCosSin(global) + } + val inputs = DecodeStepInputs( + embedded, rope, + List(spec.nLayers) { input("l$it.k", kvShape) }, + List(spec.nLayers) { input("l$it.v", kvShape) }, + ) + + var out: KvStepOutput? = null + val tape = tapeCtx.record { + val ct = (this as DefaultGraphExecutionContext).currentTape ?: error("DecodeStepExport: no tape") + Execution.tapeStack.pushTape(ct) + try { + out = step(inputs, this as ExecutionContext) + } finally { + Execution.tapeStack.popTape() + } + }.first + val result = out!! + val outRoles = HashMap() + fun outRole(t: Tensor, role: String) { outRoles[tapeCtx.session.refOf(t as Tensor<*, *>).id] = role } + outRole(result.logits, "logits") + result.selfK.forEachIndexed { i, t -> outRole(t, "l$i.k") } + result.selfV.forEachIndexed { i, t -> outRole(t, "l$i.v") } + + var graph = (tape as DefaultExecutionTape).toComputeGraph(synthesizeExternalInputs = true, packedConstants = packedConstants) + val diagnostics = if (packedConstants == PackedConstantHandling.KEEP_PACKED) { + val fused = BlockQuantMatmulFusionPass().apply(graph) + graph = fused.graph + fused.diagnostics + } else { + emptyList() + } + val module = toStableHlo(graph, functionName) + return ExportedDecodeStep( + mlir = module.content, + inputs = module.inputSpecs.map { DecodeStepPort(roles[it.name] ?: "?", it.shape) }, + outputs = module.outputSpecs.map { DecodeStepPort(outRoles[it.name] ?: "?", it.shape) }, + diagnostics = diagnostics, + ) + } + + private fun graphInput(shape: Shape, dtype: KClass): Tensor = VoidOpsTensor( + object : TensorData { + override val shape: Shape = shape + // Unreadable, like the shape-only results of VoidTensorOps: an eager shortcut that + // reads its input (RMSNormalization's fused path) must not bake this input in as a constant. + override fun get(vararg indices: Int): V = throw UnsupportedOperationException("graph input has no values") + override fun set(vararg indices: Int, value: V) {} + }, + dtype, + ) +} diff --git a/llm-core/src/jvmTest/kotlin/sk/ainet/lang/nn/dsl/decoder/DecoderKvModelTest.kt b/llm-core/src/jvmTest/kotlin/sk/ainet/lang/nn/dsl/decoder/DecoderKvModelTest.kt index 9a93c708..0d31808c 100644 --- a/llm-core/src/jvmTest/kotlin/sk/ainet/lang/nn/dsl/decoder/DecoderKvModelTest.kt +++ b/llm-core/src/jvmTest/kotlin/sk/ainet/lang/nn/dsl/decoder/DecoderKvModelTest.kt @@ -2,6 +2,7 @@ package sk.ainet.lang.nn.dsl.decoder import kotlin.test.Test import kotlin.test.assertEquals +import kotlin.test.assertFalse import kotlin.test.assertTrue import sk.ainet.apps.llm.OptimizedLLMMode import sk.ainet.apps.llm.OptimizedLLMRuntime @@ -10,6 +11,7 @@ import sk.ainet.io.weights.MappingConfig import sk.ainet.io.weights.WeightMapper import sk.ainet.io.weights.WeightTensor import sk.ainet.lang.nn.transformer.RoPEMode +import sk.ainet.lang.nn.transformer.VoidDense import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.Tensor import sk.ainet.lang.types.FP32 @@ -162,8 +164,49 @@ class DecoderKvModelTest { val model = buildModel() val kv = DecoderKvModel(model, FP32::class) val rope = kv.buildRopeCosSin(0, ctx) - assertEquals(Shape(1, headDim), rope.cos.shape) - assertEquals(Shape(1, headDim), rope.sin.shape) + assertEquals(Shape(1, headDim), rope.global.cos.shape) + assertEquals(Shape(1, headDim), rope.global.sin.shape) + assertFalse(rope.hasSlidingBase, "one RoPE base: the sliding tables are the global ones") + } + + @Test fun forwardEmbeddedWithPastMatchesForwardWithPast() { + val kv = DecoderKvModel(buildModel(), FP32::class) + val prompt = intArrayOf(3, 7, 1) + val promptTensor: Tensor = ctx.fromFloatArray(Shape(prompt.size), FP32::class, FloatArray(prompt.size) { prompt[it].toFloat() }) + val selectLast: Tensor = ctx.fromFloatArray(Shape(1, prompt.size), FP32::class, FloatArray(prompt.size) { if (it == prompt.size - 1) 1f else 0f }) + val prefill = kv.forwardPrefillAt(promptTensor, selectLast, ctx) + val rope = kv.buildRopeCosSin(prompt.size, ctx) + + val tok = 9 + val byId = kv.forwardWithPast(ctx.fromFloatArray(Shape(1), FP32::class, floatArrayOf(tok.toFloat())), rope, prefill.selfK, prefill.selfV, ctx) + // The host-side lookup: row `tok` of the (unscaled) token embedding table. + val table = buildWeightTensors().getValue(DecoderTensorNames.TOKEN_EMBEDDINGS).data.copyToFloatArray() + val embedded = ctx.fromFloatArray(Shape(1, dim), FP32::class, table.copyOfRange(tok * dim, (tok + 1) * dim)) + val byEmbedding = kv.forwardEmbeddedWithPast(embedded, rope, prefill.selfK, prefill.selfV, ctx) + assertLogitsClose(byId.logits.data.copyToFloatArray(), byEmbedding.logits.data.copyToFloatArray(), "forwardEmbeddedWithPast") + + // A replacement output weight is what projects the logits: the model's own, doubled. + val head = kv.module.modules.filterIsInstance>().single().params[0].value + val doubled = ctx.fromFloatArray(head.shape, FP32::class, head.data.copyToFloatArray().map { it * 2f }.toFloatArray()) + val withHead = kv.forwardEmbeddedWithPast(embedded, rope, prefill.selfK, prefill.selfV, ctx, lmHeadWeight = doubled) + assertLogitsClose(byId.logits.data.copyToFloatArray().map { it * 2f }.toFloatArray(), withHead.logits.data.copyToFloatArray(), "lmHeadWeight") + } + + @Test fun decodeStepExportNamesEveryPort() { + val kv = DecoderKvModel(buildModel(), FP32::class) + val past = 5 + val spec = DecodeStepSpec(hidden = dim, headDim = headDim, nKVHeads = kvHeads, nLayers = nLayers, past = past) + val step = DecodeStepExport.export(spec, FP32::class) { inputs, ectx -> + kv.forwardEmbeddedWithPast(inputs.embedded, inputs.rope, inputs.selfK, inputs.selfV, ectx) + } + assertTrue("func.func @decode_step" in step.mlir, step.mlir.take(400)) + val kvRoles = (0 until nLayers).flatMap { listOf("l$it.k", "l$it.v") } + assertEquals((listOf("embedded", "cosGlobal", "sinGlobal") + kvRoles).sorted(), step.inputs.map { it.role }.sorted()) + assertEquals((listOf("logits") + kvRoles).sorted(), step.outputs.map { it.role }.sorted()) + assertEquals(listOf(1, kvHeads, past, headDim), step.inputs.first { it.role == "l0.k" }.shape) + assertEquals(listOf(1, kvHeads, past + 1, headDim), step.outputs.first { it.role == "l0.k" }.shape) + assertEquals(0, step.foldedBlockQuantMatmuls, "dense weights: nothing to fold") + assertTrue(step.ioTable().lines().first().startsWith("in 0 ")) } private fun assertLogitsClose(expected: FloatArray, actual: FloatArray, label: String) { diff --git a/llm-inference/functiongemma/src/jvmMain/kotlin/sk/ainet/models/functiongemma/FunctionGemmaExportHarness.kt b/llm-inference/functiongemma/src/jvmMain/kotlin/sk/ainet/models/functiongemma/FunctionGemmaExportHarness.kt index 1d6d8782..d50643e8 100644 --- a/llm-inference/functiongemma/src/jvmMain/kotlin/sk/ainet/models/functiongemma/FunctionGemmaExportHarness.kt +++ b/llm-inference/functiongemma/src/jvmMain/kotlin/sk/ainet/models/functiongemma/FunctionGemmaExportHarness.kt @@ -17,6 +17,8 @@ import sk.ainet.lang.tensor.data.TensorData import sk.ainet.lang.tensor.ops.VoidTensorOps import sk.ainet.lang.tensor.storage.BufferHandle import sk.ainet.lang.types.FP32 +import sk.ainet.lang.nn.dsl.decoder.ChunkContext +import sk.ainet.lang.nn.dsl.decoder.RopeCosSin import sk.ainet.models.gemma.GEMMA_DEQUANTIZE_ALL import sk.ainet.io.gguf.StreamingGGUFReader import sk.ainet.models.gemma.GemmaWeightLoader @@ -374,7 +376,7 @@ public object FunctionGemmaExportHarness { try { val ectx = this as ExecutionContext val out = model.forwardWithPast( - tokenId, GemmaModel.RopeCosSin(cosG, sinG, cosS, sinS), selfKIn, selfVIn, ectx, + tokenId, RopeCosSin(cosG, sinG, cosS, sinS), selfKIn, selfVIn, ectx, ) // token output: argMax over [1, vocab] -> [1] i32 (small-int board contract). ectx.ops.argMax(out.logits, dim = -1) @@ -653,7 +655,7 @@ public object FunctionGemmaExportHarness { try { val ectx = this as ExecutionContext val out = model.forwardPrefillWithPast( - tokens, GemmaModel.ChunkContext(cosG, sinG, cosS, sinS, maskG, maskS), select, selfKIn, selfVIn, ectx, + tokens, ChunkContext(RopeCosSin(cosG, sinG, cosS, sinS), maskG, maskS), select, selfKIn, selfVIn, ectx, ) ectx.ops.argMax(out.logits, dim = -1) // [1] i32 — the selected position's token } finally { diff --git a/llm-inference/functiongemma/src/jvmTest/kotlin/sk/ainet/models/functiongemma/TorqDecodeStepExportTest.kt b/llm-inference/functiongemma/src/jvmTest/kotlin/sk/ainet/models/functiongemma/TorqDecodeStepExportTest.kt index 53980cb7..e832046c 100644 --- a/llm-inference/functiongemma/src/jvmTest/kotlin/sk/ainet/models/functiongemma/TorqDecodeStepExportTest.kt +++ b/llm-inference/functiongemma/src/jvmTest/kotlin/sk/ainet/models/functiongemma/TorqDecodeStepExportTest.kt @@ -2,41 +2,28 @@ package sk.ainet.models.functiongemma import kotlinx.coroutines.runBlocking import org.junit.jupiter.api.Assumptions.assumeTrue -import sk.ainet.compile.hlo.toStableHlo -import sk.ainet.compile.opt.passes.BlockQuantMatmulFusionPass import sk.ainet.context.DirectCpuExecutionContext -import sk.ainet.context.ExecutionContext import sk.ainet.io.JvmRandomAccessSource -import sk.ainet.lang.graph.DefaultExecutionTape -import sk.ainet.lang.graph.DefaultGraphExecutionContext -import sk.ainet.lang.tensor.Shape -import sk.ainet.lang.tensor.Tensor -import sk.ainet.lang.tensor.VoidOpsTensor import sk.ainet.io.gguf.StreamingGGUFReader +import sk.ainet.lang.nn.dsl.decoder.DecodeStepExport +import sk.ainet.lang.nn.dsl.decoder.DecodeStepSpec +import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.data.Q4_0Quantizer import sk.ainet.lang.tensor.data.Q8_0BlockTensorData import sk.ainet.lang.tensor.data.TensorData -import sk.ainet.lang.tensor.operators.bind -import sk.ainet.lang.tensor.ops.VoidTensorOps -import sk.ainet.lang.trace.PackedConstantHandling import sk.ainet.lang.types.FP32 import sk.ainet.models.gemma.GemmaModel import sk.ainet.models.gemma.GemmaNetworkLoader import sk.ainet.models.gemma.GemmaWeightLoader -import sk.ainet.tape.Execution import java.io.File import kotlin.test.Test -import kotlin.test.assertTrue +import kotlin.test.assertEquals /** - * One FunctionGemma decode step for an NPU with block-quantized weights (Torq `mixed_quant_matmul`): - * Q8_0 projections kept packed and folded by [BlockQuantMatmulFusionPass], the tied vocabulary - * projection as a Q4_0 copy of the embedding, the embedding lookup left to the caller, and a static - * past of [PAST] positions. - * - * Inputs, in order: the scaled embedding `[1, 640]`, cos/sin `[1, 256]` for the global and the - * sliding RoPE base, then K and V `[1, 1, PAST, 256]` per layer. Outputs: logits `[1, vocab]`, then - * the extended K and V per layer. + * One FunctionGemma decode step for an NPU with block-quantized weights (Torq `mixed_quant_matmul`), + * through [DecodeStepExport]: Q8_0 projections kept packed and folded, the tied vocabulary projection + * as a Q4_0 copy of the embedding, the embedding lookup left to the caller, and a static past of + * [PAST] positions. Writes `decode_step.mlir` and `io.txt` (which argument and result is which tensor). * * Runs only with `TORQ_DECODE_OUT` set (the output directory) and the Q8_0 checkpoint in `GEMMA_GGUF`. */ @@ -55,11 +42,8 @@ class TorqDecodeStepExportTest { val model = GemmaNetworkLoader.fromWeights(ctx, weights, FP32::class) as GemmaModel val md = weights.metadata val nLayers = md.blockCount - val headDim = md.getHeadDim(0) - val nKV = md.kvHeadCount val embedding = weights.tensors.getValue("token_embd.weight") - println("token_embd data: ${embedding.data::class.qualifiedName}, q proj: ${weights.tensors.getValue("blk.0.attn_q.weight").data::class.qualifiedName}") val vocab = embedding.shape[0] val hidden = embedding.shape[1] // The tied table as stored (Q8_0), decoded once and re-encoded as Q4_0: the Torq compiler @@ -69,66 +53,21 @@ class TorqDecodeStepExportTest { @Suppress("UNCHECKED_CAST") val lmHeadQ4 = ctx.fromData(Q4_0Quantizer.quantize(dense, Shape(vocab, hidden)) as TensorData, FP32::class) - // Inputs are bound to the tape context: a bare VoidOpsTensor runs operator overloads - // (`x * x` in RMSNorm) on its own untraced ops, and the result is captured as a constant. - val tapeCtx = DefaultGraphExecutionContext.tape(baseOps = VoidTensorOps()) - fun input(shape: Shape) = voidF32(shape).bind(tapeCtx) - val embedded = input(Shape(1, hidden)) - val cosG = input(Shape(1, headDim)); val sinG = input(Shape(1, headDim)) - val cosS = input(Shape(1, headDim)); val sinS = input(Shape(1, headDim)) - val kIn = List(nLayers) { input(Shape(1, nKV, PAST, headDim)) } - val vIn = List(nLayers) { input(Shape(1, nKV, PAST, headDim)) } - - var outputs: GemmaModel.GemmaWithPastOutput? = null - val tape = tapeCtx.record { - val ct = (this as DefaultGraphExecutionContext).currentTape ?: error("no tape") - Execution.tapeStack.pushTape(ct) - try { - outputs = model.forwardEmbeddedWithPast( - embedded, GemmaModel.RopeCosSin(cosG, sinG, cosS, sinS), kIn, vIn, this as ExecutionContext, lmHeadWeight = lmHeadQ4, - ) - } finally { - Execution.tapeStack.popTape() - } - }.first - val graph = (tape as DefaultExecutionTape).toComputeGraph( - synthesizeExternalInputs = true, packedConstants = PackedConstantHandling.KEEP_PACKED, + val spec = DecodeStepSpec( + hidden = hidden, headDim = md.getHeadDim(0), nKVHeads = md.kvHeadCount, nLayers = nLayers, + past = PAST, slidingRopeBase = true, ) - val fused = BlockQuantMatmulFusionPass().apply(graph) - fused.diagnostics.filter { it.startsWith("BLOCK_QUANT_SKIPPED") }.forEach(::println) - val folded = fused.diagnostics.count { it.startsWith("BLOCK_QUANT_FUSED") } - println("folded $folded block-quant matmuls") - assertTrue(folded == nLayers * 7 + 1, "expected every projection and the vocabulary projection folded, got $folded") + val step = DecodeStepExport.export(spec, FP32::class) { inputs, ectx -> + model.forwardEmbeddedWithPast(inputs.embedded, inputs.rope, inputs.selfK, inputs.selfV, ectx, lmHeadWeight = lmHeadQ4) + } + step.diagnostics.filter { it.startsWith("BLOCK_QUANT_SKIPPED") }.forEach(::println) + assertEquals(nLayers * 7 + 1, step.foldedBlockQuantMatmuls, "every projection and the vocabulary projection folded") - val module = toStableHlo(fused.graph, "decode_step") File(out).mkdirs() - File(out, "decode_step.mlir").writeText(module.content) - // Which function argument / result is which tensor: the emitter orders them by graph topology. - fun id(t: Tensor) = tapeCtx.session.refOf(t as Tensor<*, *>).id - val inRoles = mutableMapOf(id(embedded) to "embedded", id(cosG) to "cosG", id(sinG) to "sinG", id(cosS) to "cosS", id(sinS) to "sinS") - kIn.forEachIndexed { i, t -> inRoles[id(t)] = "k$i" } - vIn.forEachIndexed { i, t -> inRoles[id(t)] = "v$i" } - val o = outputs!! - val outRoles = mutableMapOf(id(o.logits) to "logits") - o.selfK.forEachIndexed { i, t -> outRoles[id(t)] = "k$i" } - o.selfV.forEachIndexed { i, t -> outRoles[id(t)] = "v$i" } - File(out, "io.txt").writeText( - module.inputSpecs.mapIndexed { i, s -> "in $i ${inRoles[s.name] ?: "?"} ${s.name} ${s.shape}" }.joinToString("\n") + "\n" + - module.outputSpecs.mapIndexed { i, s -> "out $i ${outRoles[s.name] ?: "?"} ${s.name} ${s.shape}" }.joinToString("\n") + "\n" - ) + File(out, "decode_step.mlir").writeText(step.mlir) + File(out, "io.txt").writeText(step.ioTable()) } - private fun voidF32(shape: Shape): Tensor = VoidOpsTensor( - object : TensorData { - override val shape: Shape = shape - // Unreadable, like the shape-only results of VoidTensorOps: an eager shortcut that - // reads its input (RMSNormalization's fused path) must not bake this input in as a constant. - override fun get(vararg indices: Int): Float = throw UnsupportedOperationException("graph input has no values") - override fun set(vararg indices: Int, value: Float) {} - }, - FP32::class, - ) - private companion object { const val PAST = 16 } diff --git a/llm-inference/gemma/build.gradle.kts b/llm-inference/gemma/build.gradle.kts index 8d93170a..d7776f31 100644 --- a/llm-inference/gemma/build.gradle.kts +++ b/llm-inference/gemma/build.gradle.kts @@ -50,7 +50,8 @@ kotlin { implementation(libs.skainet.io.gguf) implementation(libs.skainet.io.safetensors) implementation(libs.skainet.compile.core) - implementation(project(":llm-core")) + // api: the with-past forwards take and return llm-core's RopeCosSin / ChunkContext / KvStepOutput. + api(project(":llm-core")) implementation(libs.kotlinx.io.core) implementation(libs.kotlinx.coroutines) } diff --git a/llm-inference/gemma/src/commonMain/kotlin/sk/ainet/models/gemma/GemmaModel.kt b/llm-inference/gemma/src/commonMain/kotlin/sk/ainet/models/gemma/GemmaModel.kt index 5d1fcbe3..0ff2c457 100644 --- a/llm-inference/gemma/src/commonMain/kotlin/sk/ainet/models/gemma/GemmaModel.kt +++ b/llm-inference/gemma/src/commonMain/kotlin/sk/ainet/models/gemma/GemmaModel.kt @@ -3,6 +3,10 @@ package sk.ainet.models.gemma import sk.ainet.apps.llm.HybridTransformerBlock import sk.ainet.context.ExecutionContext import sk.ainet.lang.nn.Module +import sk.ainet.lang.nn.dsl.decoder.ChunkContext +import sk.ainet.lang.nn.dsl.decoder.KvStepOutput +import sk.ainet.lang.nn.dsl.decoder.RopeCosSin +import sk.ainet.lang.nn.dsl.decoder.buildRopeCosSin import sk.ainet.lang.nn.layers.Embedding import sk.ainet.lang.nn.layers.EmbeddingAdapter import sk.ainet.lang.nn.normalization.RMSNormalization @@ -230,47 +234,6 @@ public class GemmaModel( // typed fields, which assume SwiGLU and mis-select Gemma's four RMSNorms. // ================================================================ - /** logits `[seq, vocab]` + per-layer initial self K/V (`[1, nKVHeads, seq, headDim]`). */ - public class GemmaPrefillOutput( - public val logits: Tensor, - public val selfK: List>, - public val selfV: List>, - ) - - /** logits `[1, vocab]` + per-layer extended self K/V (`[1, nKVHeads, past+1, headDim]`). */ - public class GemmaWithPastOutput( - public val logits: Tensor, - public val selfK: List>, - public val selfV: List>, - ) - - /** Per-step SPLIT_HALF RoPE cos/sin (`[1, headDim]`), one pair per RoPE base - * (global / sliding); [forwardWithPast] routes per layer by `MHA.slidingWindow`. */ - public class RopeCosSin( - public val cosGlobal: Tensor, - public val sinGlobal: Tensor, - public val cosSliding: Tensor, - public val sinSliding: Tensor, - ) - - /** - * Inputs of [forwardPrefillWithPast] for a C-token chunk: per-RoPE-base cos/sin tables `[C, headDim]` - * (one row per absolute position `past .. past+C-1`, split-half layout as in [buildRopeCosSin]) and - * per-layer-type **additive** attention masks `[1, nHeads, C, past+C]` (0 = attend, -1e30 = masked; - * per head because a broadcast over heads to a dynamic shape is not expressible in static StableHLO). The - * masks carry everything position-dependent — the causal band inside the chunk, the zero padding - * beyond the real tokens, and the 512-position sliding window of the sliding layers — so the graph - * itself stays position-agnostic like `gemma_with_past`. - */ - public class ChunkContext( - public val cosGlobal: Tensor, - public val sinGlobal: Tensor, - public val cosSliding: Tensor, - public val sinSliding: Tensor, - public val maskGlobal: Tensor, - public val maskSliding: Tensor, - ) - private class GemmaBlockRefs( val attnNorm: RMSNormalization, val mha: MultiHeadAttention, @@ -344,6 +307,10 @@ public class GemmaModel( * position picked by the one-hot [selectAt] (`[1, C]`), plus the per-layer cache extended by all C * positions (`[1, nKVHeads, past+C, headDim]`; the caller slices the padding off, as after prefill). * This is what makes a user utterance cost one call instead of C single-token [forwardWithPast] steps. + * + * [chunk] carries split-half tables per RoPE base and a mask per layer type, each per head + * (`[1, nHeads, C, past+C]`): a broadcast over heads to a dynamic shape is not expressible in static + * StableHLO. The sliding mask also carries the 512-position window. */ public fun forwardPrefillWithPast( tokens: Tensor, @@ -352,7 +319,7 @@ public class GemmaModel( selfKIn: List>, selfVIn: List>, ctx: ExecutionContext, - ): GemmaWithPastOutput { + ): KvStepOutput { require(ple == null) { "GemmaModel.forwardPrefillWithPast: PLE models not supported by the KV-cache path yet" } val ops = ctx.ops val rawEmbeds = tokenEmbedding.forward(tokens, ctx) // [C, hidden] @@ -361,12 +328,9 @@ public class GemmaModel( val nsv = ArrayList>(blocks.size) for ((i, block) in blocks.withIndex()) { val r = refsFor(block) - val isGlobal = r.mha.slidingWindow == null - val cos = if (isGlobal) chunk.cosGlobal else chunk.cosSliding - val sin = if (isGlobal) chunk.sinGlobal else chunk.sinSliding - val mask = if (isGlobal) chunk.maskGlobal else chunk.maskSliding + val tables = chunk.rope.forLayer(r.mha) val sn = r.attnNorm.forward(h, ctx) - val (attnOut, fullK, fullV) = attnWithPastChunk(r.mha, sn, cos, sin, mask, selfKIn[i], selfVIn[i], ctx) + val (attnOut, fullK, fullV) = attnWithPastChunk(r.mha, sn, tables.cos, tables.sin, chunk.maskFor(r.mha), selfKIn[i], selfVIn[i], ctx) val postAttn = r.postAttnNorm?.forward(attnOut, ctx) ?: attnOut val h1 = ops.add(h, postAttn) val ffnOut = r.ffn.forward(r.ffnNorm.forward(h1, ctx), ctx) @@ -380,7 +344,7 @@ public class GemmaModel( val normed = outputNorm.forward(h, ctx) // [C, hidden] val hAt = ops.matmul(selectAt, normed) // [1, hidden] val logits = applySoftcap(lmHead.forward(hAt, ctx), ctx) // [1, vocab] - return GemmaWithPastOutput(logits, nsk, nsv) + return KvStepOutput(logits, nsk, nsv) } /** [attnWithPast] for a C-row chunk: heads-first `[heads, C, headDim]` projections, RoPE from `[C, headDim]` @@ -473,7 +437,7 @@ public class GemmaModel( * plus the per-layer initial self K/V (`[1, nKVHeads, seq, headDim]`) that seed [forwardWithPast]. * Nulls the block KV caches first (stateless one-pass prefill). PLE is unsupported here. */ - public fun forwardPrefill(input: Tensor, ctx: ExecutionContext): GemmaPrefillOutput { + public fun forwardPrefill(input: Tensor, ctx: ExecutionContext): KvStepOutput { require(ple == null) { "GemmaModel.forwardPrefill: PLE models not supported by the KV-cache path yet" } stripCaches() val ops = ctx.ops @@ -495,7 +459,7 @@ public class GemmaModel( selfV += ops.unsqueeze(kv.v, 0) } val logits = applySoftcap(lmHead.forward(outputNorm.forward(h, ctx), ctx), ctx) - return GemmaPrefillOutput(logits, selfK, selfV) + return KvStepOutput(logits, selfK, selfV) } /** @@ -504,7 +468,7 @@ public class GemmaModel( * `gemma_prefill_at` graph reads the first generated token from this single row instead of * computing `seq × vocab` logits for the padding positions. */ - public fun forwardPrefillAt(input: Tensor, selectAt: Tensor, ctx: ExecutionContext): GemmaPrefillOutput { + public fun forwardPrefillAt(input: Tensor, selectAt: Tensor, ctx: ExecutionContext): KvStepOutput { require(ple == null) { "GemmaModel.forwardPrefillAt: PLE models not supported by the KV-cache path yet" } stripCaches() val ops = ctx.ops @@ -528,7 +492,7 @@ public class GemmaModel( 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 = applySoftcap(lmHead.forward(hAt, ctx), ctx) // [1, vocab] - return GemmaPrefillOutput(logits, selfK, selfV) + return KvStepOutput(logits, selfK, selfV) } /** @@ -543,7 +507,7 @@ public class GemmaModel( selfKIn: List>, selfVIn: List>, ctx: ExecutionContext, - ): GemmaWithPastOutput { + ): KvStepOutput { require(ple == null) { "GemmaModel.forwardWithPast: PLE models not supported by the KV-cache path yet" } val ops = ctx.ops val rawEmbeds = tokenEmbedding.forward(tokenId, ctx) @@ -565,7 +529,7 @@ public class GemmaModel( selfVIn: List>, ctx: ExecutionContext, lmHeadWeight: Tensor? = null, - ): GemmaWithPastOutput { + ): KvStepOutput { require(ple == null) { "GemmaModel.forwardEmbeddedWithPast: PLE models not supported by the KV-cache path yet" } val ops = ctx.ops var h = embedded @@ -573,11 +537,9 @@ public class GemmaModel( val nsv = ArrayList>(blocks.size) for ((i, block) in blocks.withIndex()) { val r = refsFor(block) - val isGlobal = r.mha.slidingWindow == null - val cos = if (isGlobal) rope.cosGlobal else rope.cosSliding - val sin = if (isGlobal) rope.sinGlobal else rope.sinSliding + val tables = rope.forLayer(r.mha) val sn = r.attnNorm.forward(h, ctx) - val (attnOut, fullK, fullV) = attnWithPast(r.mha, sn, cos, sin, selfKIn[i], selfVIn[i], ctx) + val (attnOut, fullK, fullV) = attnWithPast(r.mha, sn, tables.cos, tables.sin, selfKIn[i], selfVIn[i], ctx) val postAttn = r.postAttnNorm?.forward(attnOut, ctx) ?: attnOut val h1 = ops.add(h, postAttn) val ffnOut = r.ffn.forward(r.ffnNorm.forward(h1, ctx), ctx) @@ -591,31 +553,16 @@ public class GemmaModel( val normed = outputNorm.forward(h, ctx) val projected = if (lmHeadWeight != null) ops.matmulWeightTransposed(normed, lmHeadWeight) else lmHead.forward(normed, ctx) val logits = applySoftcap(projected, ctx) - return GemmaWithPastOutput(logits, nsk, nsv) + return KvStepOutput(logits, nsk, nsv) } /** - * Build the per-step SPLIT_HALF RoPE cos/sin tables for [position] (seqLen 1): one pair from a - * global-base layer, one from a sliding-base layer (gemma3 uses two RoPE bases). Fed to - * [forwardWithPast] as runtime inputs so the decode graph carries no compile-time position. + * Build the per-step RoPE cos/sin tables for [position] (seqLen 1): one pair from a global-base + * layer, one from a sliding-base layer (gemma3 uses two RoPE bases). Fed to [forwardWithPast] as + * runtime inputs so the decode graph carries no compile-time position. */ - public fun buildRopeCosSin(position: Int, ctx: ExecutionContext): RopeCosSin { - val mhas = blocks.map { refsFor(it).mha } - val globalMha = mhas.firstOrNull { it.slidingWindow == null } - val slidingMha = mhas.firstOrNull { it.slidingWindow != null } - val gRope = (globalMha ?: slidingMha ?: error("GemmaModel: no attention layers")).rope!! - val sRope = (slidingMha ?: globalMha)!!.rope!! - val (gc, gs) = gRope.buildSplitHalfCosSin(position, 1) - val (sc, ss) = sRope.buildSplitHalfCosSin(position, 1) - val gShape = Shape(1, gRope.headDim) - val sShape = Shape(1, sRope.headDim) - return RopeCosSin( - cosGlobal = ctx.fromFloatArray(gShape, dtype, gc), - sinGlobal = ctx.fromFloatArray(gShape, dtype, gs), - cosSliding = ctx.fromFloatArray(sShape, dtype, sc), - sinSliding = ctx.fromFloatArray(sShape, dtype, ss), - ) - } + public fun buildRopeCosSin(position: Int, ctx: ExecutionContext): RopeCosSin = + buildRopeCosSin(blocks.map { refsFor(it).mha }, position, dtype, ctx) public companion object { /** diff --git a/llm-inference/qwen/src/jvmMain/kotlin/sk/ainet/models/qwen/QwenExportHarness.kt b/llm-inference/qwen/src/jvmMain/kotlin/sk/ainet/models/qwen/QwenExportHarness.kt index eb3491d5..7364136d 100644 --- a/llm-inference/qwen/src/jvmMain/kotlin/sk/ainet/models/qwen/QwenExportHarness.kt +++ b/llm-inference/qwen/src/jvmMain/kotlin/sk/ainet/models/qwen/QwenExportHarness.kt @@ -11,6 +11,8 @@ import sk.ainet.lang.graph.DefaultExecutionTape import sk.ainet.lang.graph.DefaultGraphExecutionContext import sk.ainet.lang.nn.dsl.decoder.DECODER_DEQUANTIZE_ALL import sk.ainet.lang.nn.dsl.decoder.DecoderKvModel +import sk.ainet.lang.nn.dsl.decoder.ChunkContext +import sk.ainet.lang.nn.dsl.decoder.RopeCosSin import sk.ainet.lang.tensor.Dim import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.Tensor @@ -224,7 +226,7 @@ public object QwenExportHarness { * layer's K/V -- see [QwenKvContract.prefillWithPastArgs]), select 1xC) -> per-layer K,V * extended by C THEN token 1xi32`. One call per utterance instead of C [exportWithPastTraced] * steps. [QwenKvArch] has one RoPE base and no sliding window (every layer "global"), so - * [DecoderKvModel.ChunkContext] carries exactly one cos/sin/mask triple -- half the arguments + * [ChunkContext] carries exactly one cos/sin/mask triple -- half the arguments * FunctionGemma's per-layer-type version needs. */ private fun exportPrefillWithPastTraced(kv: DecoderKvModel, arch: QwenKvArch, outDir: String, chunk: Int, bf16: Boolean, hostGather: Boolean): File { @@ -250,7 +252,7 @@ public object QwenExportHarness { try { val ectx = this as ExecutionContext val out = kv.forwardPrefillWithPast( - tokens, DecoderKvModel.ChunkContext(cos, sin, mask), select, selfKIn, selfVIn, ectx, + tokens, ChunkContext(RopeCosSin(cos, sin), mask), select, selfKIn, selfVIn, ectx, ) ectx.ops.argMax(out.logits, dim = -1) } finally { @@ -291,7 +293,7 @@ public object QwenExportHarness { Execution.tapeStack.pushTape(ct) try { val ectx = this as ExecutionContext - val out = kv.forwardWithPast(tokenId, DecoderKvModel.RopeCosSin(cos, sin), selfKIn, selfVIn, ectx) + val out = kv.forwardWithPast(tokenId, RopeCosSin(cos, sin), selfKIn, selfVIn, ectx) ectx.ops.argMax(out.logits, dim = -1) } finally { Execution.tapeStack.popTape() From 0d042c6577a3cdabce02b6eeca3aed413853e2b7 Mon Sep 17 00:00:00 2001 From: Michal Harakal <471011+michalharakal@users.noreply.github.com> Date: Fri, 6 Nov 2026 13:58:37 +0100 Subject: [PATCH 3/4] build: SKaiNET 0.63.0 0.63.0 ships the block-quant StableHLO export (SKaiNET#1478, #1479) that DecodeStepExport and TorqDecodeStepExportTest build on. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 2 +- gradle/libs.versions.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a53738d2..05a584c6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,7 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] -Needs an engine release with SKaiNET#1478 and SKaiNET#1479 (block-quant export). +Against engine **0.63.0** (SKaiNET#1478 and SKaiNET#1479, block-quant export). ### Added diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml index f5fa46b6..5949b0c9 100644 --- a/gradle/libs.versions.toml +++ b/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" From a2d472ef73dc75f7c201afe6df4ea6f8d491e450 Mon Sep 17 00:00:00 2001 From: Michal Harakal <471011+michalharakal@users.noreply.github.com> Date: Fri, 6 Nov 2026 14:56:01 +0100 Subject: [PATCH 4/4] test(llama): pin scoped decode to four ulp of the step's largest logit On engine 0.63.0 the scoped and unscoped decode of the tiny llama model differ in 12 of 192 logits over the first two steps, by at most 7.5e-9: one ulp of the step's largest logit, but 16 ulps of a logit near zero (0.0066), which tripped the per-logit four-ulp pin. The accumulation-order rounding is on the scale of the step's activations, so the tolerance is now four ulp of the step's largest |logit|. Engine 0.62.0 passes either way. Co-Authored-By: Claude Opus 5.5 --- .../llama/ForwardScopeSteadyStateTest.kt | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/ForwardScopeSteadyStateTest.kt b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/ForwardScopeSteadyStateTest.kt index 1e133ae0..3b212ab4 100644 --- a/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/ForwardScopeSteadyStateTest.kt +++ b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/ForwardScopeSteadyStateTest.kt @@ -23,7 +23,8 @@ import sk.ainet.lang.types.FP32 * unscoped one (`forwardSlabFloats = 0`), step for step. The slab changes * where activations live, never what they hold; since engine 0.62.0 the JVM * kernels chosen for slab-backed and heap activations may accumulate in a - * different order, so the pin allows four ulp instead of raw-bit equality. + * different order, so the pin allows four ulp of the step's largest logit + * instead of raw-bit equality. * 2. **Flat peak** — `peakFloats` stops growing once decode is warm. The KV * history lives outside the scope (the cache detaches it to ambient * storage), so per-step slab use grows only with the KV *views* concat @@ -106,14 +107,16 @@ class ForwardScopeSteadyStateTest { val a = scoped.forward(tokenId).data.copyToFloatArray() val b = unscoped.forward(tokenId).data.copyToFloatArray() assertEquals(b.size, a.size) + // Equal to within a few ulp of the step's largest logit. Up to engine 0.61 this was + // raw-bit identity (same kernels, same values, different storage); since engine 0.62.0 + // the JVM ops route slab-backed and heap activations through kernels whose accumulation + // order can differ (the fused native attention and norm kernels take heap operands, a + // segment-backed activation is bridged first). The rounding that moves is on the scale + // of the step's activations, so a logit near zero can move by many of its own ulps: + // engine 0.63.0 moves 0.0066 by 7.5e-9 (16 of its ulps, one ulp of the largest logit). + // The slab still changes where activations live, never what they mean. + val tolerance = 4 * Math.ulp(b.maxOf { kotlin.math.abs(it) }) for (i in a.indices) { - // Equal to within a few ulp. Up to engine 0.61 this was raw-bit identity (same - // kernels, same values, different storage); since engine 0.62.0 the JVM ops route - // slab-backed and heap activations through kernels whose accumulation order can - // differ (the fused native attention and norm kernels take heap operands, a - // segment-backed activation is bridged first), which moves a logit by one ulp. - // The slab still changes where activations live, never what they mean. - val tolerance = 4 * Math.ulp(b[i]) assertTrue( a[i] == b[i] || kotlin.math.abs(a[i] - b[i]) <= tolerance, "step=$step logit[$i]: scoped=${a[i]} unscoped=${b[i]} (tolerance $tolerance)",