From 2e8bbe3fd641724b5ad489b03433da020eff50b7 Mon Sep 17 00:00:00 2001 From: Michal Harakal <471011+michalharakal@users.noreply.github.com> Date: Sat, 7 Nov 2026 10:08:42 +0100 Subject: [PATCH 1/2] feat(lfm2): decode step for an NPU with block-quantized weights Lfm2DecodeStep is the fixed-capacity single-token step of the cached network from an embedded token, with an output head of its own (a Q4_0 copy of the tied embedding folds; the tied tensor itself has two consumers). The layer loop and cache bookkeeping move into Lfm2CachedLayers, shared with Lfm2CachedTensorNetwork. The step's short convolution can be multiply + sum (ShortConvLowering.MULTIPLY_SUM) and its attention one batched product. DecodeStepExport.exportPorts traces a step over arbitrary named ports; export is built on it. exportLfm2DecodeStep and Lfm2DecodeStepExportCli use it with packed Q8_0 projections (loadLfm2Model(..., weightFormFor), lfm2ProjectionsAs) and write the module, port table, bf16 embedding rows and a manifest. LFM2.5-350M: 93 folded matmuls, top-1 equal to FP32; on an SL2610 NPU one dispatch, 533 ms per step, logits cosine 0.996. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 22 ++ .../lang/nn/dsl/decoder/DecodeStepExport.kt | 77 +++-- .../lang/nn/dsl/decoder/DecoderKvModelTest.kt | 23 ++ llm-inference/lfm2/README.md | 11 + llm-inference/lfm2/api/lfm2.api | 37 +++ llm-inference/lfm2/build.gradle.kts | 15 +- .../models/lfm2/Lfm2CachedTensorNetwork.kt | 94 ++++-- .../sk/ainet/models/lfm2/Lfm2DecodeStep.kt | 52 ++++ .../sk/ainet/models/lfm2/ShortConvTensorOp.kt | 32 +- .../sk/ainet/models/lfm2/Lfm2Weights.kt | 32 +- .../lfm2/export/Lfm2DecodeStepExport.kt | 192 ++++++++++++ .../ainet/models/lfm2/Lfm2DecodeStepTest.kt | 273 ++++++++++++++++++ 12 files changed, 801 insertions(+), 59 deletions(-) create mode 100644 llm-inference/lfm2/src/commonMain/kotlin/sk/ainet/models/lfm2/Lfm2DecodeStep.kt create mode 100644 llm-inference/lfm2/src/jvmMain/kotlin/sk/ainet/models/lfm2/export/Lfm2DecodeStepExport.kt create mode 100644 llm-inference/lfm2/src/jvmTest/kotlin/sk/ainet/models/lfm2/Lfm2DecodeStepTest.kt diff --git a/CHANGELOG.md b/CHANGELOG.md index d77ef1cd..efa69174 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,28 @@ requiring an engine release. `0.57.1` against engine `0.57.0` is such a release. The format roughly follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [Unreleased] + +Against engine **0.63.0**, transformers-only and additive. + +### Added + +- **LFM2 decode step for an NPU with block-quantized weights** (`llm-inference:lfm2`): `Lfm2DecodeStep` is the + fixed-capacity single-token step of `Lfm2CachedTensorNetwork` from an already embedded token, with its own output + head, so one graph serves every position up to the capacity (mask and one-hot write mask as inputs, conv histories + and K/V as explicit state). `exportLfm2DecodeStep` traces it through `DecodeStepExport.exportPorts` with Q8_0 / Q4_0 + weights kept packed and folded; `Lfm2DecodeStepExportCli` (Gradle `exportLfm2DecodeStep`) writes + `decode_step.mlir`, the port table, the host's bf16 embedding rows and an export manifest from a verified local + checkpoint. LFM2.5-350M with Q8_0 projections and a Q4_0 head: all 93 matmuls fold; against the FP32 model the + logits cosine is 0.997–0.999 with the same top-1 over a 10-token sequence. Compiled with Torq 2.3.0 it is one NPU + dispatch and runs on an SL2610 at 533 ms per step, logits cosine 0.996 to FP32 with the same top-1. +- `ShortConvLowering.MULTIPLY_SUM`: the single-token short convolution as `sum(window * taps)` instead of a grouped + `conv1d`, numerically equal, for compilers without grouped convolutions (the decode step's default). +- `DecodeStepExport.exportPorts(inputs, dtype) { ports, ctx -> results }` (`llm-core`, jvm): the decode-step export + over arbitrary named ports, for decoders whose state is more than K/V; `export` is built on it, unchanged. +- `loadLfm2Model(..., weightFormFor)` and `lfm2ProjectionsAs(encoding)`: load the projections requantized (e.g. Q8_0), + the embedding, norms and convolution taps dense. + ## [0.64.1] — 2026-11-07 Against engine **0.63.0**. A packaging fix: Maven Central refused the 0.64.0 deployment, so **0.64.0 was never 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 index ac735b9e..1d241736 100644 --- 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 @@ -81,44 +81,75 @@ public object DecodeStepExport { packedConstants: PackedConstantHandling = PackedConstantHandling.KEEP_PACKED, step: (DecodeStepInputs, ExecutionContext) -> KvStepOutput, ): ExportedDecodeStep { + val kvShape = Shape(1, spec.nKVHeads, spec.past, spec.headDim) + val ropeShape = Shape(1, spec.headDim) + val ports = buildList { + add("embedded" to Shape(1, spec.hidden)) + add("cosGlobal" to ropeShape); add("sinGlobal" to ropeShape) + if (spec.slidingRopeBase) { add("cosSliding" to ropeShape); add("sinSliding" to ropeShape) } + for (i in 0 until spec.nLayers) add("l$i.k" to kvShape) + for (i in 0 until spec.nLayers) add("l$i.v" to kvShape) + } + return exportPorts(ports, dtype, functionName, packedConstants) { p, ctx -> + val global = RopeTables(p.getValue("cosGlobal"), p.getValue("sinGlobal")) + val rope = if (spec.slidingRopeBase) { + RopeCosSin(global, RopeTables(p.getValue("cosSliding"), p.getValue("sinSliding"))) + } else { + RopeCosSin(global) + } + val inputs = DecodeStepInputs( + p.getValue("embedded"), rope, + List(spec.nLayers) { p.getValue("l$it.k") }, + List(spec.nLayers) { p.getValue("l$it.v") }, + ) + val result = step(inputs, ctx) + linkedMapOf>("logits" to result.logits).apply { + result.selfK.forEachIndexed { i, t -> put("l$i.k", t) } + result.selfV.forEachIndexed { i, t -> put("l$i.v", t) } + } + } + } + + /** + * Trace a step over arbitrary named ports and export it like [export]: one graph input per entry + * of [inputs] (bound to the trace, in this order), one result per entry [step] returns, each + * named by its key in [ExportedDecodeStep.inputs] / [ExportedDecodeStep.outputs]. For a decoder + * whose state is more than K/V (a convolution history, a fixed-capacity cache with its masks). + */ + public fun exportPorts( + inputs: List>, + dtype: KClass, + functionName: String = "decode_step", + packedConstants: PackedConstantHandling = PackedConstantHandling.KEEP_PACKED, + step: (Map>, ExecutionContext) -> Map>, + ): ExportedDecodeStep { + require(inputs.map { it.first }.distinct().size == inputs.size) { "DecodeStepExport: input port names must be unique" } // 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 bound = LinkedHashMap>() + for ((role, shape) in inputs) { + val t = graphInput(shape, dtype).bind(tapeCtx) + roles[tapeCtx.session.refOf(t as Tensor<*, *>).id] = role + bound[role] = t } - 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 + var out: Map>? = 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) + out = step(bound, 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") } + for ((role, t) in out!!) { + val id = tapeCtx.session.refOf(t as Tensor<*, *>).id + require(outRoles.put(id, role) == null) { "DecodeStepExport: results '${outRoles[id]}' and '$role' are the same tensor" } + } var graph = (tape as DefaultExecutionTape).toComputeGraph(synthesizeExternalInputs = true, packedConstants = packedConstants) val diagnostics = if (packedConstants == PackedConstantHandling.KEEP_PACKED) { 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 0d31808c..22584433 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 @@ -209,6 +209,29 @@ class DecoderKvModelTest { assertTrue(step.ioTable().lines().first().startsWith("in 0 ")) } + @Test fun exportPortsNamesArbitraryStateAndRejectsAnAliasedResult() { + val step = DecodeStepExport.exportPorts( + listOf("x" to Shape(1, 4), "history" to Shape(4, 3)), FP32::class, "toy_step", + ) { p, ectx -> + val x = p.getValue("x") + linkedMapOf( + "y" to ectx.ops.add(x, x), + "next_history" to ectx.ops.narrow(ectx.ops.concat(listOf(p.getValue("history"), ectx.ops.transpose(x)), dim = 1), 1, 1, 3), + ) + } + assertTrue("func.func @toy_step" in step.mlir, step.mlir.take(400)) + assertEquals(listOf("history", "x"), step.inputs.map { it.role }.sorted()) + assertEquals(listOf("next_history", "y"), step.outputs.map { it.role }.sorted()) + assertEquals(listOf(4, 3), step.outputs.first { it.role == "next_history" }.shape) + val aliased = runCatching { + DecodeStepExport.exportPorts(listOf("x" to Shape(1, 4)), FP32::class) { p, ectx -> + val y = ectx.ops.add(p.getValue("x"), p.getValue("x")) + linkedMapOf("a" to y, "b" to y) + } + } + assertTrue(aliased.isFailure, "two results on one tensor must be rejected") + } + private fun assertLogitsClose(expected: FloatArray, actual: FloatArray, label: String) { assertEquals(expected.size, actual.size, "$label: size mismatch") for (i in expected.indices) { diff --git a/llm-inference/lfm2/README.md b/llm-inference/lfm2/README.md index 2978993e..f5191f5f 100644 --- a/llm-inference/lfm2/README.md +++ b/llm-inference/lfm2/README.md @@ -78,3 +78,14 @@ Synthetic tests verify implementation mechanics. They do not establish model qua ``` The generation CLI takes local config/checkpoint paths with mandatory hashes, `--capacity`, and a new output directory. It emits the shared export bundle plus `generation-contract.json`. The host must compile and bind the resulting module/archive into a verified runtime manifest before using the [optional IREE runtime](../../llm-runtime/klfm2/README.md#optional-iree-vulkan-runtime). Export and synthetic state-parity tests establish graph structure and tensor semantics; accelerator execution still requires target validation. + +## Decode step for an NPU with block-quantized weights + +`Lfm2DecodeStep` is the fixed-capacity single-token step from an already embedded token, with an output head of its own: the host does the embedding lookup, and the graph takes the embedded row `[1, width]`, RoPE `cos`/`sin` `[1, headDim]`, the additive `mask` `[1, capacity + 1]`, the one-hot `write_mask` `[1, capacity, 1]` and every layer's state (`history_` `[width, kernel]`, `key_`/`value_` `[kvHeads, capacity, headDim]`), and returns `logits` and the next state. One graph serves every position up to the capacity. The short convolution is lowered as multiply + sum and attention as one batched product over all heads by default. + +```sh +./gradlew :llm-inference:lfm2:exportLfm2DecodeStep --args="--config model/config.json --config-sha256 \ + --checkpoint model/model.safetensors --checkpoint-sha256 --capacity 64 --output out/decode-step" +``` + +The projections load as Q8_0 and the head as a Q4_0 copy of the tied embedding (`--weights`, `--head`); the export keeps them packed and folds every one into a block-quant matmul, and fails if one is left unfolded. It writes `decode_step.mlir`, `io.txt` (the argument and result order), `embed.bf16.bin` (the host's embedding rows) and `export.json`. LFM2.5-350M: 93 folded matmuls; against the FP32 model the logits cosine is 0.997–0.999 with the same top-1. The gated tests in `Lfm2DecodeStepTest` (`LFM2_MODEL_DIR`, `LFM2_GOLDEN_DIR`) check a local checkpoint and write a golden step for a device. diff --git a/llm-inference/lfm2/api/lfm2.api b/llm-inference/lfm2/api/lfm2.api index 7ce2e449..732d388a 100644 --- a/llm-inference/lfm2/api/lfm2.api +++ b/llm-inference/lfm2/api/lfm2.api @@ -170,6 +170,17 @@ public final class sk/ainet/models/lfm2/Lfm2Config$Companion { public final fun parse (Ljava/lang/String;)Lsk/ainet/models/lfm2/Lfm2Config; } +public final class sk/ainet/models/lfm2/Lfm2DecodeStep { + public fun (Lsk/ainet/models/lfm2/ModelDimensions;Ljava/util/Map;ILsk/ainet/lang/tensor/Tensor;Lsk/ainet/models/lfm2/ShortConvLowering;Lsk/ainet/models/lfm2/AttentionLayoutPolicy;)V + public synthetic fun (Lsk/ainet/models/lfm2/ModelDimensions;Ljava/util/Map;ILsk/ainet/lang/tensor/Tensor;Lsk/ainet/models/lfm2/ShortConvLowering;Lsk/ainet/models/lfm2/AttentionLayoutPolicy;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun forward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/Map;Lsk/ainet/context/ExecutionContext;)Lsk/ainet/models/lfm2/Lfm2CachedTensorResult; + public final fun getAttentionLayout ()Lsk/ainet/models/lfm2/AttentionLayoutPolicy; + public final fun getCapacity ()I + public final fun getConvLowering ()Lsk/ainet/models/lfm2/ShortConvLowering; + public final fun getDimensions ()Lsk/ainet/models/lfm2/ModelDimensions; + public final fun stateShapes ()Ljava/util/LinkedHashMap; +} + public final class sk/ainet/models/lfm2/Lfm2PrefillNetwork { public fun (Lsk/ainet/models/lfm2/ModelDimensions;Ljava/util/Map;)V public final fun forward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/Map;Lsk/ainet/context/ExecutionContext;)Lsk/ainet/models/lfm2/Lfm2PrefillResult; @@ -218,8 +229,11 @@ public final class sk/ainet/models/lfm2/Lfm2WeightRecord { } public final class sk/ainet/models/lfm2/Lfm2WeightsKt { + public static final fun lfm2ProjectionsAs (Lsk/ainet/lang/tensor/storage/TensorEncoding;)Lkotlin/jvm/functions/Function1; public static final fun loadLfm2Model (Lsk/ainet/models/lfm2/Lfm2Config;Ljava/io/File;ILjava/lang/String;Lsk/ainet/context/DirectCpuExecutionContext;Lsk/ainet/models/lfm2/RmsReductionPolicy;)Lsk/ainet/models/lfm2/LoadedLfm2Model; + public static final fun loadLfm2Model (Lsk/ainet/models/lfm2/Lfm2Config;Ljava/io/File;ILjava/lang/String;Lsk/ainet/context/DirectCpuExecutionContext;Lsk/ainet/models/lfm2/RmsReductionPolicy;Lkotlin/jvm/functions/Function1;)Lsk/ainet/models/lfm2/LoadedLfm2Model; public static synthetic fun loadLfm2Model$default (Lsk/ainet/models/lfm2/Lfm2Config;Ljava/io/File;ILjava/lang/String;Lsk/ainet/context/DirectCpuExecutionContext;Lsk/ainet/models/lfm2/RmsReductionPolicy;ILjava/lang/Object;)Lsk/ainet/models/lfm2/LoadedLfm2Model; + public static synthetic fun loadLfm2Model$default (Lsk/ainet/models/lfm2/Lfm2Config;Ljava/io/File;ILjava/lang/String;Lsk/ainet/context/DirectCpuExecutionContext;Lsk/ainet/models/lfm2/RmsReductionPolicy;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)Lsk/ainet/models/lfm2/LoadedLfm2Model; } public final class sk/ainet/models/lfm2/LoadedLfm2Model { @@ -299,10 +313,20 @@ public final class sk/ainet/models/lfm2/RmsReductionPolicy : java/lang/Enum { public static fun values ()[Lsk/ainet/models/lfm2/RmsReductionPolicy; } +public final class sk/ainet/models/lfm2/ShortConvLowering : java/lang/Enum { + public static final field DEPTHWISE Lsk/ainet/models/lfm2/ShortConvLowering; + public static final field MULTIPLY_SUM Lsk/ainet/models/lfm2/ShortConvLowering; + public static fun getEntries ()Lkotlin/enums/EnumEntries; + public static fun valueOf (Ljava/lang/String;)Lsk/ainet/models/lfm2/ShortConvLowering; + public static fun values ()[Lsk/ainet/models/lfm2/ShortConvLowering; +} + public final class sk/ainet/models/lfm2/ShortConvTensorOp { public fun (IILsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)V + public fun (IILsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/models/lfm2/ShortConvLowering;)V public final fun forward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/context/ExecutionContext;)Lsk/ainet/models/lfm2/ShortConvTensorResult; public final fun getKernel ()I + public final fun getLowering ()Lsk/ainet/models/lfm2/ShortConvLowering; public final fun getWidth ()I } @@ -319,6 +343,19 @@ public final class sk/ainet/models/lfm2/ShortConvTensorResult { public fun toString ()Ljava/lang/String; } +public final class sk/ainet/models/lfm2/export/Lfm2DecodeStepExportCli { + public static final field INSTANCE Lsk/ainet/models/lfm2/export/Lfm2DecodeStepExportCli; + public static final fun main ([Ljava/lang/String;)V +} + +public final class sk/ainet/models/lfm2/export/Lfm2DecodeStepExportKt { + public static final fun exportLfm2DecodeStep (Lsk/ainet/models/lfm2/Lfm2DecodeStep;Lsk/ainet/lang/trace/PackedConstantHandling;)Lsk/ainet/lang/nn/dsl/decoder/ExportedDecodeStep; + public static synthetic fun exportLfm2DecodeStep$default (Lsk/ainet/models/lfm2/Lfm2DecodeStep;Lsk/ainet/lang/trace/PackedConstantHandling;ILjava/lang/Object;)Lsk/ainet/lang/nn/dsl/decoder/ExportedDecodeStep; + public static final fun lfm2DecodeStepInputs (Lsk/ainet/models/lfm2/Lfm2DecodeStep;)Ljava/util/List; + public static final fun lfm2FoldableMatmuls (Ljava/util/List;)I + public static final fun lfm2Head (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/context/ExecutionContext;)Lsk/ainet/lang/tensor/Tensor; +} + public final class sk/ainet/models/lfm2/export/Lfm2ExportBundleKt { public static final fun lfm2ExportSourcePins ()Lkotlinx/serialization/json/JsonObject; public static final fun sha256 (Ljava/io/File;)Ljava/lang/String; diff --git a/llm-inference/lfm2/build.gradle.kts b/llm-inference/lfm2/build.gradle.kts index 3ac45a62..b2925923 100644 --- a/llm-inference/lfm2/build.gradle.kts +++ b/llm-inference/lfm2/build.gradle.kts @@ -24,6 +24,7 @@ kotlin { api(libs.kotlinx.serialization.json) } jvmMain.dependencies { + api(project(":llm-core")) api(libs.skainet.compile.hlo) implementation(libs.skainet.compile.core) implementation(libs.skainet.compile.dag) @@ -41,7 +42,10 @@ kotlin { tasks.withType().configureEach { jvmArgs("--enable-preview", "--add-modules", "jdk.incubator.vector") - maxHeapSize = "512m" + maxHeapSize = "2g" + // The gated real-checkpoint decode-step test (LFM2_MODEL_DIR) loads LFM2.5-350M twice. + if (System.getenv("LFM2_MODEL_DIR") != null) maxHeapSize = "12g" + System.getenv("LFM2_MODEL_DIR")?.let { environment("LFM2_MODEL_DIR", it) } } /** Provenance travels with the published JAR, including runtime source additions. */ @@ -86,6 +90,15 @@ tasks.register("exportLfm2") { maxHeapSize = "8g" } +tasks.register("exportLfm2DecodeStep") { + group = "export" + description = "Export the LFM2 decode step with block-quantized weights for an NPU; use --args='--help' for options" + mainClass.set("sk.ainet.models.lfm2.export.Lfm2DecodeStepExportCli") + classpath = files(jvmMainCompilation.output.allOutputs, jvmMainCompilation.runtimeDependencyFiles) + jvmArgs("--enable-preview", "--add-modules", "jdk.incubator.vector") + maxHeapSize = "8g" +} + tasks.register("exportLfm2Generation") { group = "export" description = "Export the fixed-capacity LFM2 cached generation graph" diff --git a/llm-inference/lfm2/src/commonMain/kotlin/sk/ainet/models/lfm2/Lfm2CachedTensorNetwork.kt b/llm-inference/lfm2/src/commonMain/kotlin/sk/ainet/models/lfm2/Lfm2CachedTensorNetwork.kt index b3072e52..caa648cb 100644 --- a/llm-inference/lfm2/src/commonMain/kotlin/sk/ainet/models/lfm2/Lfm2CachedTensorNetwork.kt +++ b/llm-inference/lfm2/src/commonMain/kotlin/sk/ainet/models/lfm2/Lfm2CachedTensorNetwork.kt @@ -51,10 +51,67 @@ class Lfm2CachedTensorNetwork( embeddingShards: List>, ) { private val shards = embeddingShards.toList() + private val layers = Lfm2CachedLayers(dimensions, weights, capacity, ShortConvLowering.DEPTHWISE, AttentionLayoutPolicy.DIRECT_KV_PER_HEAD) + val embeddingShardRows: List get() = shards.map { it.shape[0] } + + init { + val d = dimensions + require(shards.isNotEmpty() && shards.all { + it.shape.dimensions.size == 2 && it.shape[0] > 0 && it.shape[1] == d.width && + it.shape[0].toLong() * d.width * 4L <= MAX_EMBEDDING_SHARD_BYTES + } && shards.sumOf { it.shape[0].toLong() } == d.vocabulary.toLong()) { + "Embedding shards must cover the vocabulary in contiguous row order, each at most 64 MiB" + } + } + + /** State names and shapes are identical at the graph input and output. */ + fun stateShapes(): LinkedHashMap = layers.stateShapes() + + /** + * [mask] has zero at populated prefix slots and the appended current-token slot C, + * and a large negative finite value at all unfilled slots. [writeMask] is one-hot + * along its capacity axis. Runtime values are never read by this traceable function. + */ + fun forward( + tokenOneHot: Tensor, + cos: Tensor, + sin: Tensor, + mask: Tensor, + writeMask: Tensor, + states: Map>, + ctx: ExecutionContext, + ): Lfm2CachedTensorResult { + val d = dimensions + require(tokenOneHot.shape == Shape(1, d.vocabulary)) + val ops = ctx.ops + var offset = 0 + val selected = shards.map { shard -> + val oneHot = ops.narrow(tokenOneHot, dim = 1, start = offset, length = shard.shape[0]) + offset += shard.shape[0] + ops.matmul(oneHot, shard) + } + val (normalized, nextStates) = layers.forward(selected.reduce { total, part -> ops.add(total, part) }, + cos, sin, mask, writeMask, states, ctx) + val logits = shards.map { ops.matmulWeightTransposed(normalized, it) } + return Lfm2CachedTensorResult(if (logits.size == 1) logits.single() else ops.concat(logits, dim = 1), nextStates) + } +} + +/** + * The decoder blocks of a fixed-capacity single-token step, from the token's hidden row to the + * final-norm output, with the cache bookkeeping of [Lfm2CachedTensorNetwork.forward]: shared by that + * network and [Lfm2DecodeStep]. + */ +internal class Lfm2CachedLayers( + val dimensions: ModelDimensions, + weights: Map>, + val capacity: Int, + convLowering: ShortConvLowering, + attentionLayout: AttentionLayoutPolicy, +) { private val finalNorm = weights.getValue("model.embedding_norm.weight") private val convolution = linkedMapOf() private val attention = linkedMapOf() - val embeddingShardRows: List get() = shards.map { it.shape[0] } init { val d = dimensions @@ -64,17 +121,11 @@ class Lfm2CachedTensorNetwork( require(weights.keys == expected.keys && expected.all { (name, shape) -> weights.getValue(name).shape == shape }) { "Cached full-model weight inventory/layout mismatch" } - require(shards.isNotEmpty() && shards.all { - it.shape.dimensions.size == 2 && it.shape[0] > 0 && it.shape[1] == d.width && - it.shape[0].toLong() * d.width * 4L <= MAX_EMBEDDING_SHARD_BYTES - } && shards.sumOf { it.shape[0].toLong() } == d.vocabulary.toLong()) { - "Embedding shards must cover the vocabulary in contiguous row order, each at most 64 MiB" - } for ((layer, kind) in d.layerTypes.withIndex()) { fun w(name: String) = weights.getValue("model.layers.$layer.$name.weight") if (kind == "conv") { convolution[layer] = BlockTensorOp(d.width, d.kernel, d.intermediate, d.epsilon, - ShortConvTensorOp(d.width, d.kernel, w("conv.in_proj"), w("conv.conv"), w("conv.out_proj")), + ShortConvTensorOp(d.width, d.kernel, w("conv.in_proj"), w("conv.conv"), w("conv.out_proj"), convLowering), w("operator_norm"), w("ffn_norm"), w("feed_forward.w1"), w("feed_forward.w2"), w("feed_forward.w3"), RmsReductionPolicy.BALANCED_FP32) } else { @@ -82,12 +133,11 @@ class Lfm2CachedTensorNetwork( w("operator_norm"), w("ffn_norm"), w("self_attn.q_proj"), w("self_attn.k_proj"), w("self_attn.v_proj"), w("self_attn.out_proj"), w("self_attn.q_layernorm"), w("self_attn.k_layernorm"), w("feed_forward.w1"), w("feed_forward.w2"), w("feed_forward.w3"), - RmsReductionPolicy.BALANCED_FP32, AttentionLayoutPolicy.DIRECT_KV_PER_HEAD) + RmsReductionPolicy.BALANCED_FP32, attentionLayout) } } } - /** State names and shapes are identical at the graph input and output. */ fun stateShapes(): LinkedHashMap = linkedMapOf().apply { for ((layer, kind) in dimensions.layerTypes.withIndex()) { if (kind == "conv") put("history_$layer", Shape(dimensions.width, dimensions.kernel)) else { @@ -97,22 +147,18 @@ class Lfm2CachedTensorNetwork( } } - /** - * [mask] has zero at populated prefix slots and the appended current-token slot C, - * and a large negative finite value at all unfilled slots. [writeMask] is one-hot - * along its capacity axis. Runtime values are never read by this traceable function. - */ + /** The final-norm output `[1, width]` and every next state, for the hidden row [input] `[1, width]`. */ fun forward( - tokenOneHot: Tensor, + input: Tensor, cos: Tensor, sin: Tensor, mask: Tensor, writeMask: Tensor, states: Map>, ctx: ExecutionContext, - ): Lfm2CachedTensorResult { + ): Pair, LinkedHashMap>> { val d = dimensions - require(tokenOneHot.shape == Shape(1, d.vocabulary)) + require(input.shape == Shape(1, d.width)) require(cos.shape == Shape(1, d.headDim) && sin.shape == cos.shape) require(mask.shape == Shape(1, capacity + 1) && writeMask.shape == Shape(1, capacity, 1)) val expected = stateShapes() @@ -120,13 +166,7 @@ class Lfm2CachedTensorNetwork( "Every decoder layer needs its complete fixed-capacity runtime state" } val ops = ctx.ops - var offset = 0 - val selected = shards.map { shard -> - val oneHot = ops.narrow(tokenOneHot, dim = 1, start = offset, length = shard.shape[0]) - offset += shard.shape[0] - ops.matmul(oneHot, shard) - } - var hidden = selected.reduce { total, part -> ops.add(total, part) } + var hidden = input val keepMask = ops.rsubScalar(1f, writeMask) val nextStates = linkedMapOf>() for ((layer, kind) in d.layerTypes.withIndex()) { @@ -147,8 +187,6 @@ class Lfm2CachedTensorNetwork( result.output } } - val normalized = rmsNormalize(hidden, finalNorm, ctx, RmsReductionPolicy.BALANCED_FP32, d.epsilon) - val logits = shards.map { ops.matmulWeightTransposed(normalized, it) } - return Lfm2CachedTensorResult(if (logits.size == 1) logits.single() else ops.concat(logits, dim = 1), nextStates) + return rmsNormalize(hidden, finalNorm, ctx, RmsReductionPolicy.BALANCED_FP32, d.epsilon) to nextStates } } diff --git a/llm-inference/lfm2/src/commonMain/kotlin/sk/ainet/models/lfm2/Lfm2DecodeStep.kt b/llm-inference/lfm2/src/commonMain/kotlin/sk/ainet/models/lfm2/Lfm2DecodeStep.kt new file mode 100644 index 00000000..5d560824 --- /dev/null +++ b/llm-inference/lfm2/src/commonMain/kotlin/sk/ainet/models/lfm2/Lfm2DecodeStep.kt @@ -0,0 +1,52 @@ +package sk.ainet.models.lfm2 + +import sk.ainet.context.ExecutionContext +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.types.FP32 + +/** + * One fixed-capacity single-token step from an already embedded token: the decoder of + * [Lfm2CachedTensorNetwork] with the embedding lookup left to the host and the output projection + * over [head], a `[vocabulary, width]` weight of its own. The head may be a differently encoded + * copy of the tied embedding (e.g. Q4_0), so a block-quant fusion can fold it: the tied tensor + * itself also serves the lookup and would have two consumers. + * + * Inputs and state follow [Lfm2CachedTensorNetwork.forward]: [mask] `[1, capacity + 1]`, + * [writeMask] `[1, capacity, 1]`, `history_` `[width, kernel]`, `key_`/`value_` + * `[kvHeads, capacity, headDim]`. [convLowering] defaults to [ShortConvLowering.MULTIPLY_SUM]: + * the step's convolution is then elementwise arithmetic, with the same result. [attentionLayout] defaults to + * one batched product over all heads ([AttentionLayoutPolicy.CURRENT_BATCHED]: 12 attention matmuls instead + * of 192 for LFM2.5-350M, 22 % faster on an SL2610 NPU) rather than the cached network's per-head form. + */ +class Lfm2DecodeStep( + val dimensions: ModelDimensions, + weights: Map>, + val capacity: Int, + private val head: Tensor, + val convLowering: ShortConvLowering = ShortConvLowering.MULTIPLY_SUM, + val attentionLayout: AttentionLayoutPolicy = AttentionLayoutPolicy.CURRENT_BATCHED, +) { + private val layers = Lfm2CachedLayers(dimensions, weights, capacity, convLowering, attentionLayout) + + init { + require(head.shape == Shape(dimensions.vocabulary, dimensions.width)) { "Output head must be [vocabulary, width]" } + } + + /** State names and shapes, identical at the step's input and output. */ + fun stateShapes(): LinkedHashMap = layers.stateShapes() + + /** Logits `[1, vocabulary]` and every next state, for the embedding row [embedded] `[1, width]`. */ + fun forward( + embedded: Tensor, + cos: Tensor, + sin: Tensor, + mask: Tensor, + writeMask: Tensor, + states: Map>, + ctx: ExecutionContext, + ): Lfm2CachedTensorResult { + val (normalized, nextStates) = layers.forward(embedded, cos, sin, mask, writeMask, states, ctx) + return Lfm2CachedTensorResult(ctx.ops.matmulWeightTransposed(normalized, head), nextStates) + } +} diff --git a/llm-inference/lfm2/src/commonMain/kotlin/sk/ainet/models/lfm2/ShortConvTensorOp.kt b/llm-inference/lfm2/src/commonMain/kotlin/sk/ainet/models/lfm2/ShortConvTensorOp.kt index 1860e51d..d1cc0b28 100644 --- a/llm-inference/lfm2/src/commonMain/kotlin/sk/ainet/models/lfm2/ShortConvTensorOp.kt +++ b/llm-inference/lfm2/src/commonMain/kotlin/sk/ainet/models/lfm2/ShortConvTensorOp.kt @@ -10,6 +10,14 @@ data class ShortConvTensorResult( val nextHistory: Tensor, ) +/** + * How the depthwise convolution is lowered. [DEPTHWISE] is one grouped `conv1d`. [MULTIPLY_SUM] + * applies to a single token only: the convolution over a `[D, K]` window is then + * `sum(window * weight, dim = 1)`, the same arithmetic as plain elementwise ops, for compilers + * without grouped convolutions. Longer inputs always use [DEPTHWISE]. + */ +enum class ShortConvLowering { DEPTHWISE, MULTIPLY_SUM } + /** Pure tensor/state form shared by eager evaluation and StableHLO export. */ class ShortConvTensorOp( val width: Int, @@ -17,7 +25,16 @@ class ShortConvTensorOp( private val inWeight: Tensor, private val convWeight: Tensor, private val outWeight: Tensor, + val lowering: ShortConvLowering, ) { + constructor( + width: Int, + kernel: Int, + inWeight: Tensor, + convWeight: Tensor, + outWeight: Tensor, + ) : this(width, kernel, inWeight, convWeight, outWeight, ShortConvLowering.DEPTHWISE) + init { require(width > 0 && kernel > 0) { "Dimensions must be positive" } require(inWeight.shape == Shape(3 * width, width)) { "in_proj shape mismatch" } @@ -37,11 +54,16 @@ class ShortConvTensorOp( val combined = ops.concat(listOf(history, ops.transpose(gated)), dim = 1) // K history samples are retained; only the most recent K-1 affect a new token. val window = ops.narrow(combined, dim = 1, start = 1, length = kernel - 1 + length) - val convolved = ops.conv1d( - ops.reshape(window, Shape(1, width, kernel - 1 + length)), - convWeight, bias = null, padding = 0, groups = width, - ) - val timeMajor = ops.transpose(ops.reshape(convolved, Shape(width, length))) + val timeMajor = if (lowering == ShortConvLowering.MULTIPLY_SUM && length == 1) { + val taps = ops.multiply(window, ops.reshape(convWeight, Shape(width, kernel))) + ops.reshape(ops.sum(taps, dim = 1), Shape(1, width)) + } else { + val convolved = ops.conv1d( + ops.reshape(window, Shape(1, width, kernel - 1 + length)), + convWeight, bias = null, padding = 0, groups = width, + ) + ops.transpose(ops.reshape(convolved, Shape(width, length))) + } val output = ops.matmulWeightTransposed(ops.multiply(c, timeMajor), outWeight) // Keep this as a separate terminal operation: both results are graph outputs. val nextHistory = ops.narrow(combined, dim = 1, start = length, length = kernel) diff --git a/llm-inference/lfm2/src/jvmMain/kotlin/sk/ainet/models/lfm2/Lfm2Weights.kt b/llm-inference/lfm2/src/jvmMain/kotlin/sk/ainet/models/lfm2/Lfm2Weights.kt index 31724bb2..99c93e34 100644 --- a/llm-inference/lfm2/src/jvmMain/kotlin/sk/ainet/models/lfm2/Lfm2Weights.kt +++ b/llm-inference/lfm2/src/jvmMain/kotlin/sk/ainet/models/lfm2/Lfm2Weights.kt @@ -5,6 +5,10 @@ import sk.ainet.context.DirectCpuExecutionContext import sk.ainet.io.JvmRandomAccessSource import sk.ainet.io.safetensors.SafeTensorsParametersLoader import sk.ainet.io.safetensors.StreamingSafeTensorsReader +import sk.ainet.lang.memory.plan.EncodingRequest +import sk.ainet.lang.memory.plan.WeightForm +import sk.ainet.lang.tensor.storage.PackedBlockStorage +import sk.ainet.lang.tensor.storage.TensorEncoding import sk.ainet.lang.tensor.Tensor import sk.ainet.lang.types.FP32 import java.io.File @@ -35,6 +39,20 @@ fun loadLfm2Model( expectedCheckpointSha256: String? = null, ctx: DirectCpuExecutionContext = DirectCpuExecutionContext.create(), rmsReduction: RmsReductionPolicy = RmsReductionPolicy.BALANCED_FP32, +): LoadedLfm2Model = loadLfm2Model(config, checkpoint, contextLimit, expectedCheckpointSha256, ctx, rmsReduction, weightFormFor = null) + +/** + * [loadLfm2Model] with a load form per tensor name: [weightFormFor], e.g. [lfm2ProjectionsAs], requantizes + * the projections; the records then hash the values the model computes with (dequantized). + */ +fun loadLfm2Model( + config: Lfm2Config, + checkpoint: File, + contextLimit: Int = minOf(2048, config.maximumPositions), + expectedCheckpointSha256: String? = null, + ctx: DirectCpuExecutionContext = DirectCpuExecutionContext.create(), + rmsReduction: RmsReductionPolicy = RmsReductionPolicy.BALANCED_FP32, + weightFormFor: ((String) -> WeightForm?)?, ): LoadedLfm2Model { val dimensions = config.toDimensions(contextLimit) val expected = HybridLfm2Model.expectedWeights(dimensions) @@ -75,7 +93,7 @@ fun loadLfm2Model( val weights = linkedMapOf>() runBlocking { - SafeTensorsParametersLoader(sourceProvider = { JvmRandomAccessSource.open(checkpoint.absolutePath) }) + SafeTensorsParametersLoader(sourceProvider = { JvmRandomAccessSource.open(checkpoint.absolutePath) }, weightFormFor = weightFormFor) .load(ctx, FP32::class) { name, tensor -> require(name in expected && weights.put(name, tensor) == null) { "Unexpected or repeated loaded tensor: $name" } } @@ -84,7 +102,7 @@ fun loadLfm2Model( val records = expected.map { (name, shape) -> val tensor = weights.getValue(name) require(tensor.shape == shape) { "Loaded tensor shape mismatch: $name" } - val values = tensor.data.copyToFloatArray() + val values = (tensor.data as? PackedBlockStorage)?.toFloatArray() ?: tensor.data.copyToFloatArray() val digest = MessageDigest.getInstance("SHA-256") val buffer = ByteBuffer.allocate(256 * 1024).order(ByteOrder.LITTLE_ENDIAN) for (value in values) { @@ -99,4 +117,14 @@ fun loadLfm2Model( return LoadedLfm2Model(HybridLfm2Model(dimensions, ctx, weights, rmsReduction), weights.toMap(), records, checkpointHash) } +/** + * A [loadLfm2Model] form that loads every projection (attention, convolution in/out, feed-forward) + * as [encoding]; the embedding, norms and convolution taps stay dense. + */ +fun lfm2ProjectionsAs(encoding: TensorEncoding): (String) -> WeightForm? { + val form = WeightForm(encoding = EncodingRequest.RequantizeTo(encoding)) + val projection = Regex("""model\.layers\.\d+\.(self_attn\.(q|k|v|out)_proj|conv\.(in|out)_proj|feed_forward\.w[123])\.weight""") + return { name -> if (projection.matches(name)) form else null } +} + private fun ByteArray.hex(): String = joinToString("") { "%02x".format(it) } diff --git a/llm-inference/lfm2/src/jvmMain/kotlin/sk/ainet/models/lfm2/export/Lfm2DecodeStepExport.kt b/llm-inference/lfm2/src/jvmMain/kotlin/sk/ainet/models/lfm2/export/Lfm2DecodeStepExport.kt new file mode 100644 index 00000000..84c79869 --- /dev/null +++ b/llm-inference/lfm2/src/jvmMain/kotlin/sk/ainet/models/lfm2/export/Lfm2DecodeStepExport.kt @@ -0,0 +1,192 @@ +package sk.ainet.models.lfm2.export + +import kotlinx.serialization.json.* +import sk.ainet.context.ExecutionContext +import sk.ainet.lang.nn.dsl.decoder.DecodeStepExport +import sk.ainet.lang.nn.dsl.decoder.ExportedDecodeStep +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.Q4_0Quantizer +import sk.ainet.lang.tensor.data.Q8_0Quantizer +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.tensor.storage.PackedBlockStorage +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.trace.PackedConstantHandling +import sk.ainet.lang.types.FP32 +import sk.ainet.models.lfm2.AttentionLayoutPolicy +import sk.ainet.models.lfm2.Lfm2Config +import sk.ainet.models.lfm2.Lfm2DecodeStep +import sk.ainet.models.lfm2.ShortConvLowering +import sk.ainet.models.lfm2.lfm2ProjectionsAs +import sk.ainet.models.lfm2.loadLfm2Model +import java.io.File +import java.io.FileOutputStream +import java.nio.ByteBuffer +import java.nio.ByteOrder +import java.security.MessageDigest + +/** The step's argument order before the state: `embedded`, `cos`, `sin`, `mask`, `write_mask`. */ +fun lfm2DecodeStepInputs(step: Lfm2DecodeStep): List> { + val d = step.dimensions + return listOf( + "embedded" to Shape(1, d.width), + "cos" to Shape(1, d.headDim), + "sin" to Shape(1, d.headDim), + "mask" to Shape(1, step.capacity + 1), + "write_mask" to Shape(1, step.capacity, 1), + ) + step.stateShapes().map { (name, shape) -> name to shape } +} + +/** + * Trace [step] through [DecodeStepExport.exportPorts]: one graph for every position up to the + * capacity, results `logits` then every state. With [PackedConstantHandling.KEEP_PACKED] the + * block-quantized projections and head stay packed and are folded. + */ +fun exportLfm2DecodeStep( + step: Lfm2DecodeStep, + packedConstants: PackedConstantHandling = PackedConstantHandling.KEEP_PACKED, +): ExportedDecodeStep = DecodeStepExport.exportPorts(lfm2DecodeStepInputs(step), FP32::class, "decode_step", packedConstants) { p, ctx -> + val states = step.stateShapes().keys.associateWith { p.getValue(it) } + val r = step.forward(p.getValue("embedded"), p.getValue("cos"), p.getValue("sin"), p.getValue("mask"), p.getValue("write_mask"), states, ctx) + linkedMapOf("logits" to r.logits).apply { putAll(r.states) } +} + +/** [embedding] re-encoded as [encoding] (Q4_0 or Q8_0) for a separate output head, or the tensor itself for null. */ +@Suppress("UNCHECKED_CAST") +fun lfm2Head(embedding: Tensor, encoding: TensorEncoding?, ctx: ExecutionContext): Tensor { + if (encoding == null) return embedding + val dense = (embedding.data as? PackedBlockStorage)?.toFloatArray() ?: embedding.data.copyToFloatArray() + val data = when (encoding) { + TensorEncoding.Q4_0 -> Q4_0Quantizer.quantize(dense, embedding.shape) + TensorEncoding.Q8_0 -> Q8_0Quantizer.quantize(dense, embedding.shape) + else -> throw IllegalArgumentException("Unsupported head encoding ${encoding.name}") + } + return ctx.fromData(data as TensorData, FP32::class) +} + +/** Projections the decode step folds: five per convolution layer, seven per attention layer, plus the head. */ +fun lfm2FoldableMatmuls(layerTypes: List): Int = layerTypes.sumOf { if (it == "conv") 5 else 7 } + 1 + +/** + * Local checkpoint export of the NPU decode step: no download, compilation or device execution. + * Writes `decode_step.mlir`, `io.txt` (the port table), `embed.bf16.bin` (the host's embedding + * rows, `[vocabulary, width]` bf16 little-endian) and `export.json`. + */ +object Lfm2DecodeStepExportCli { + @JvmStatic + fun main(args: Array) { + if (args.isEmpty() || args.contentEquals(arrayOf("--help"))) { + println("LFM2 decode-step export: --config FILE --config-sha256 SHA256 --checkpoint FILE --checkpoint-sha256 SHA256 " + + "--capacity N --output NEW_DIRECTORY [--weights q8_0|q4_0|f32] [--head q4_0|q8_0|tied] [--conv multiply-sum|depthwise] [--attention batched|per-head]") + return + } + val allowed = setOf("--config", "--config-sha256", "--checkpoint", "--checkpoint-sha256", "--capacity", "--output", + "--weights", "--head", "--conv", "--attention") + require(args.size % 2 == 0) { "Every option requires one value; use --help" } + val options = linkedMapOf() + for (pair in args.toList().chunked(2)) { + require(pair[0] in allowed && options.put(pair[0], pair[1]) == null) { "Unknown or repeated option ${pair[0]}" } + } + fun required(name: String): String = requireNotNull(options[name]) { "Missing $name" } + fun pin(name: String): String = required(name).also { + require(it.matches(Regex("[0-9a-f]{64}"))) { "$name must be lowercase SHA-256" } + } + fun encoding(value: String, name: String): TensorEncoding? = when (value) { + "q8_0" -> TensorEncoding.Q8_0 + "q4_0" -> TensorEncoding.Q4_0 + "f32", "tied" -> null + else -> throw IllegalArgumentException("Invalid $name $value") + } + val configFile = File(required("--config")) + val checkpoint = File(required("--checkpoint")) + val configPin = pin("--config-sha256") + val checkpointPin = pin("--checkpoint-sha256") + val capacity = requireNotNull(required("--capacity").toIntOrNull()) { "Invalid --capacity" } + val weightsName = options["--weights"] ?: "q8_0" + val headName = options["--head"] ?: "q4_0" + val weights = encoding(weightsName, "--weights") + val head = encoding(headName, "--head") + val conv = when (options["--conv"] ?: "multiply-sum") { + "multiply-sum" -> ShortConvLowering.MULTIPLY_SUM + "depthwise" -> ShortConvLowering.DEPTHWISE + else -> throw IllegalArgumentException("Invalid --conv") + } + val attention = when (options["--attention"] ?: "batched") { + "per-head" -> AttentionLayoutPolicy.DIRECT_KV_PER_HEAD + "batched" -> AttentionLayoutPolicy.CURRENT_BATCHED + else -> throw IllegalArgumentException("Invalid --attention") + } + val output = File(required("--output")) + require(!output.exists() && output.absoluteFile.parentFile.isDirectory) { "Use a new output directory beneath an existing parent" } + require(configFile.isFile && configFile.canRead() && sha256(configFile) == configPin) { "Configuration SHA-256 mismatch or unreadable configuration" } + val config = Lfm2Config.parse(configFile.readText()) + require(capacity in 1..config.maximumPositions) { "Capacity exceeds the checkpoint position limit" } + + val loaded = loadLfm2Model(config, checkpoint, contextLimit = capacity, expectedCheckpointSha256 = checkpointPin, + weightFormFor = weights?.let(::lfm2ProjectionsAs)) + val embedding = loaded.weights.getValue("model.embed_tokens.weight") + val step = Lfm2DecodeStep(loaded.model.dimensions, loaded.weights, capacity, lfm2Head(embedding, head, loaded.model.ctx), conv, attention) + val exported = exportLfm2DecodeStep(step) + val skipped = exported.diagnostics.filter { it.startsWith("BLOCK_QUANT_SKIPPED") } + require(skipped.isEmpty()) { "Packed weights the fusion could not fold:\n" + skipped.joinToString("\n") } + val expectedFolds = when { + weights == null && head == null -> 0 + weights == null -> 1 + head == null -> lfm2FoldableMatmuls(loaded.model.dimensions.layerTypes) - 1 + else -> lfm2FoldableMatmuls(loaded.model.dimensions.layerTypes) + } + check(exported.foldedBlockQuantMatmuls == expectedFolds) { + "Folded ${exported.foldedBlockQuantMatmuls} block-quant matmuls, expected $expectedFolds" + } + + output.mkdirs() + File(output, "decode_step.mlir").writeText(exported.mlir) + File(output, "io.txt").writeText(exported.ioTable()) + val embedFile = File(output, "embed.bf16.bin") + writeBf16(embedding, embedFile) + val d = loaded.model.dimensions + val manifest = buildJsonObject { + put("schema_version", 1); put("kind", "lfm2_decode_step_export"); put("function_name", "decode_step") + put("config_sha256", configPin); put("checkpoint_sha256", loaded.checkpointSha256) + put("capacity", capacity); put("weights", weightsName); put("head", headName) + put("conv_lowering", conv.name.lowercase()); put("attention_layout", attention.name.lowercase()); put("folded_block_quant_matmuls", exported.foldedBlockQuantMatmuls) + put("dimensions", buildJsonObject { + put("vocabulary", d.vocabulary); put("width", d.width); put("head_dimension", d.headDim) + put("kv_heads", d.kvHeads); put("conv_kernel", d.kernel); put("rope_theta", d.theta) + put("layer_types", JsonArray(d.layerTypes.map(::JsonPrimitive))) + }) + put("embedding_rows", buildJsonObject { + put("file", embedFile.name); put("dtype", "bf16"); put("shape", JsonArray(listOf(d.vocabulary, d.width).map(::JsonPrimitive))) + put("sha256", sha256(embedFile)) + }) + put("host_control", "embedding_lookup_rope_masks_position_and_stop_policy") + } + File(output, "export.json").writeText(Json { prettyPrint = true }.encodeToString(JsonObject.serializer(), manifest) + "\n") + println("Exported decode_step: capacity=$capacity, weights=$weightsName, head=$headName, conv=${conv.name.lowercase()}, " + + "folded=${exported.foldedBlockQuantMatmuls}; compilation and device execution remain unqualified") + } + + /** Round-to-nearest-even bf16 of [tensor]'s values, little-endian. */ + private fun writeBf16(tensor: Tensor, file: File) { + val values = (tensor.data as? PackedBlockStorage)?.toFloatArray() ?: tensor.data.copyToFloatArray() + FileOutputStream(file).buffered(1 shl 20).use { out -> + val buffer = ByteBuffer.allocate(2).order(ByteOrder.LITTLE_ENDIAN) + for (v in values) { + val bits = java.lang.Float.floatToRawIntBits(v) + val rounded = (bits + 0x7FFF + ((bits ushr 16) and 1)) ushr 16 + buffer.clear(); buffer.putShort(rounded.toShort()); out.write(buffer.array()) + } + } + } + + private fun sha256(file: File): String = file.inputStream().buffered().use { input -> + val digest = MessageDigest.getInstance("SHA-256") + val buffer = ByteArray(1 shl 20) + while (true) { + val n = input.read(buffer) + if (n < 0) break + digest.update(buffer, 0, n) + } + digest.digest().joinToString("") { "%02x".format(it) } + } +} diff --git a/llm-inference/lfm2/src/jvmTest/kotlin/sk/ainet/models/lfm2/Lfm2DecodeStepTest.kt b/llm-inference/lfm2/src/jvmTest/kotlin/sk/ainet/models/lfm2/Lfm2DecodeStepTest.kt new file mode 100644 index 00000000..9efe0c9d --- /dev/null +++ b/llm-inference/lfm2/src/jvmTest/kotlin/sk/ainet/models/lfm2/Lfm2DecodeStepTest.kt @@ -0,0 +1,273 @@ +package sk.ainet.models.lfm2 + +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.Q8_0Quantizer +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.FP32 +import sk.ainet.models.lfm2.export.exportLfm2DecodeStep +import sk.ainet.models.lfm2.export.lfm2DecodeStepInputs +import sk.ainet.models.lfm2.export.lfm2FoldableMatmuls +import sk.ainet.models.lfm2.export.lfm2Head +import java.io.File +import kotlin.math.abs +import kotlin.math.cos +import kotlin.math.pow +import kotlin.math.sin +import kotlin.math.sqrt +import kotlin.test.* + +class Lfm2DecodeStepTest { + private val ctx = DirectCpuExecutionContext.create() + + private fun tensor(shape: Shape, values: FloatArray = FloatArray(shape.dimensions.reduce(Int::times))): Tensor = + ctx.fromFloatArray(shape, FP32::class, values) + + private fun weightsFor(d: ModelDimensions): Map> = + HybridLfm2Model.expectedWeights(d).mapValues { (name, shape) -> + val values = FloatArray(shape.dimensions.reduce(Int::times)) { i -> + if (name.endsWith("norm.weight")) 0.75f + i % 3 * 0.1f + else (((i * 7 + name.hashCode().mod(17)) % 19) - 9) * 0.025f + } + tensor(shape, values) + } + + /** Host-side controls of one step at [position], as [Lfm2CachedTensorNetworkTest] builds them. */ + private fun controls(d: ModelDimensions, capacity: Int, position: Int): List> { + val cosines = FloatArray(d.headDim) + val sines = FloatArray(d.headDim) + for (i in 0 until d.headDim / 2) { + val angle = position * d.theta.pow(-2.0 * i / d.headDim) + for (lane in intArrayOf(i, i + d.headDim / 2)) { + cosines[lane] = cos(angle).toFloat(); sines[lane] = sin(angle).toFloat() + } + } + return listOf( + tensor(Shape(1, d.headDim), cosines), tensor(Shape(1, d.headDim), sines), + tensor(Shape(1, capacity + 1), FloatArray(capacity + 1) { if (it < position || it == capacity) 0f else -Float.MAX_VALUE }), + tensor(Shape(1, capacity, 1), FloatArray(capacity) { if (it == position) 1f else 0f }), + ) + } + + private fun embeddedRow(embedding: Tensor, token: Int): Tensor = + ctx.ops.narrow(embedding, dim = 0, start = token, length = 1) + + private fun closeEnough(expected: FloatArray, actual: FloatArray, tolerance: Float = 2e-6f) { + assertEquals(expected.size, actual.size) + expected.indices.forEach { i -> + assertTrue(actual[i].isFinite() && abs(expected[i] - actual[i]) <= tolerance, "Value $i differs: ${expected[i]} versus ${actual[i]}") + } + } + + private val small = ModelDimensions(8, 12, 4, 2, 2, 4, 3, 1e-5f, 1e6, + listOf("conv", "full_attention", "conv", "full_attention"), contextLimit = 6) + + @Test + fun embeddedStepWithTheTiedHeadMatchesTheFullModelAtEveryPositionForEveryLowering() { + val weights = weightsFor(small) + val embedding = weights.getValue("model.embed_tokens.weight") + val model = HybridLfm2Model(small, ctx, weights) + for (lowering in ShortConvLowering.entries) for (layout in AttentionLayoutPolicy.entries) { + val step = Lfm2DecodeStep(small, weights, 6, embedding, lowering, layout) + var states: Map> = step.stateShapes().mapValues { (_, s) -> tensor(s) } + val state = HybridModelState() + for ((position, token) in intArrayOf(0, 3, 6, 11, 4, 9).withIndex()) { + val expected = model.forward(intArrayOf(token), state) + val (c, s, mask, write) = controls(small, 6, position) + val actual = step.forward(embeddedRow(embedding, token), c, s, mask, write, states, ctx) + closeEnough(expected.logits.data.copyToFloatArray(), actual.logits.data.copyToFloatArray()) + states = actual.states + } + } + } + + @Test + fun multiplySumConvolutionEqualsTheGroupedConvolutionForOneToken() { + val d = 8 + val k = 3 + fun w(shape: Shape, seed: Int) = tensor(shape, FloatArray(shape.dimensions.reduce(Int::times)) { ((it * seed + 5) % 13 - 6) * 0.05f }) + val inW = w(Shape(3 * d, d), 3) + val convW = w(Shape(d, 1, k), 5) + val outW = w(Shape(d, d), 7) + val x = w(Shape(1, d), 11) + val history = w(Shape(d, k), 2) + val grouped = ShortConvTensorOp(d, k, inW, convW, outW, ShortConvLowering.DEPTHWISE).forward(x, history, ctx) + val summed = ShortConvTensorOp(d, k, inW, convW, outW, ShortConvLowering.MULTIPLY_SUM).forward(x, history, ctx) + closeEnough(grouped.output.data.copyToFloatArray(), summed.output.data.copyToFloatArray(), 1e-6f) + assertContentEquals(grouped.nextHistory.data.copyToFloatArray(), summed.nextHistory.data.copyToFloatArray()) + } + + @Test + fun exportFoldsEveryBlockQuantProjectionAndTheHead() { + val d = ModelDimensions(32, 40, 4, 2, 8, 64, 3, 1e-5f, 1e6, + listOf("conv", "full_attention", "conv", "full_attention"), contextLimit = 6) + val projection = Regex("""model\.layers\.\d+\.(self_attn\.(q|k|v|out)_proj|conv\.(in|out)_proj|feed_forward\.w[123])\.weight""") + @Suppress("UNCHECKED_CAST") + val weights = weightsFor(d).mapValues { (name, t) -> + if (!projection.matches(name)) t + else ctx.fromData(Q8_0Quantizer.quantize(t.data.copyToFloatArray(), t.shape) as TensorData, FP32::class) + } + val head = lfm2Head(weights.getValue("model.embed_tokens.weight"), TensorEncoding.Q4_0, ctx) + + val step = Lfm2DecodeStep(d, weights, 6, head) + val exported = exportLfm2DecodeStep(step) + assertEquals(emptyList(), exported.diagnostics.filter { it.startsWith("BLOCK_QUANT_SKIPPED") }) + assertEquals(lfm2FoldableMatmuls(d.layerTypes), exported.foldedBlockQuantMatmuls) + assertEquals(25, exported.foldedBlockQuantMatmuls) + assertFalse(exported.mlir.contains("stablehlo.convolution"), "the decode step's convolution is multiply + sum") + val states = listOf("history_0", "key_1", "value_1", "history_2", "key_3", "value_3") + assertEquals((listOf("embedded", "cos", "sin", "mask", "write_mask") + states).toSet(), exported.inputs.map { it.role }.toSet()) + assertEquals((listOf("logits") + states).toSet(), exported.outputs.map { it.role }.toSet()) + val shapes = lfm2DecodeStepInputs(step).toMap() + exported.inputs.forEach { assertEquals(shapes.getValue(it.role).dimensions.toList(), it.shape, it.role) } + + val depthwise = exportLfm2DecodeStep(Lfm2DecodeStep(d, weights, 6, head, ShortConvLowering.DEPTHWISE)) + assertTrue(depthwise.mlir.contains("stablehlo.convolution")) + assertEquals(25, depthwise.foldedBlockQuantMatmuls) + } + + /** + * Gated on `LFM2_MODEL_DIR` (a local LFM2.5-350M with `config.json` and `model.safetensors`): + * the Q8_0 projections and Q4_0 head of the NPU step against the FP32 model over a short + * sequence (top-1 and cosine per step), and the export's fold count. + */ + @Test + fun realCheckpointQuantizedStepKeepsTheReferenceTopOne() { + val dir = System.getenv("LFM2_MODEL_DIR")?.let(::File) ?: return + val config = Lfm2Config.parse(File(dir, "config.json").readText()) + val capacity = 16 + val reference = loadLfm2Model(config, File(dir, "model.safetensors"), contextLimit = capacity) + val quantized = loadLfm2Model(config, File(dir, "model.safetensors"), contextLimit = capacity, + ctx = ctx, weightFormFor = lfm2ProjectionsAs(TensorEncoding.Q8_0)) + val dims = quantized.model.dimensions + val embedding = quantized.weights.getValue("model.embed_tokens.weight") + val step = Lfm2DecodeStep(dims, quantized.weights, capacity, lfm2Head(embedding, TensorEncoding.Q4_0, ctx)) + + val tokens = intArrayOf(1, 6, 6423, 708, 1098, 8448, 4924, 523, 2210, 7) + val state = HybridModelState() + var states: Map> = step.stateShapes().mapValues { (_, s) -> tensor(s) } + var agree = 0 + for ((position, token) in tokens.withIndex()) { + val expected = reference.model.forward(intArrayOf(token), state).logits.data.copyToFloatArray() + val (c, s, mask, write) = controls(dims, capacity, position) + val r = step.forward(embeddedRow(embedding, token), c, s, mask, write, states, ctx) + states = r.states + val actual = r.logits.data.copyToFloatArray() + val dot = expected.indices.sumOf { expected[it].toDouble() * actual[it] } + val cosine = dot / sqrt(expected.sumOf { it.toDouble() * it } * actual.sumOf { it.toDouble() * it }) + val same = expected.indices.maxBy { expected[it] } == actual.indices.maxBy { actual[it] } + if (same) agree++ + println("position $position token $token: cosine %.5f, top-1 %s".format(cosine, if (same) "same" else "differs")) + assertTrue(cosine > 0.95, "logits cosine $cosine at position $position") + } + assertTrue(agree >= tokens.size - 1, "top-1 agrees at $agree of ${tokens.size} positions") + val exported = exportLfm2DecodeStep(step) + assertEquals(lfm2FoldableMatmuls(dims.layerTypes), exported.foldedBlockQuantMatmuls) + assertEquals(93, exported.foldedBlockQuantMatmuls) + } + + /** + * Gated on `LFM2_MODEL_DIR` and `LFM2_PROMPT_IDS` (comma-separated): greedy decoding of the prompt with the FP32 + * model and with the quantized decode step, printed side by side, for checking a device run of the same step. + */ + @Test + fun realCheckpointGreedyReference() { + val dir = System.getenv("LFM2_MODEL_DIR")?.let(::File) ?: return + val prompt = System.getenv("LFM2_PROMPT_IDS")?.split(',')?.map { it.trim().toInt() }?.toIntArray() ?: return + val maxNew = 16 + val capacity = 64 + val config = Lfm2Config.parse(File(dir, "config.json").readText()) + val reference = loadLfm2Model(config, File(dir, "model.safetensors"), contextLimit = capacity) + val fp32 = reference.model.generateGreedy(prompt, maxNew, setOf(7)).generatedIds.toList() + val quantized = loadLfm2Model(config, File(dir, "model.safetensors"), contextLimit = capacity, + ctx = ctx, weightFormFor = lfm2ProjectionsAs(TensorEncoding.Q8_0)) + val dims = quantized.model.dimensions + val embedding = quantized.weights.getValue("model.embed_tokens.weight") + val step = Lfm2DecodeStep(dims, quantized.weights, capacity, lfm2Head(embedding, TensorEncoding.Q4_0, ctx)) + var states: Map> = step.stateShapes().mapValues { (_, s) -> tensor(s) } + var position = 0 + fun next(token: Int): FloatArray { + val (c, s, mask, write) = controls(dims, capacity, position++) + val r = step.forward(embeddedRow(embedding, token), c, s, mask, write, states, ctx) + states = r.states + return r.logits.data.copyToFloatArray() + } + var logits = FloatArray(0) + for (t in prompt) logits = next(t) + val stepIds = ArrayList() + while (stepIds.size < maxNew) { + val t = logits.indices.maxBy { logits[it] } + stepIds += t + if (t == 7 || stepIds.size == maxNew) break + logits = next(t) + } + println("greedy fp32 model: ${fp32.joinToString(",")}") + println("greedy q8_0 step : ${stepIds.joinToString(",")}") + } + + /** + * Gated on `LFM2_MODEL_DIR` and `LFM2_GOLDEN_DIR` (an `exportLfm2DecodeStep` output directory): + * feeds a prompt on the host, then writes the inputs of the next step in the module's argument + * order as bf16 files plus `args.txt` (`--input=` lines for an IREE runner), and the expected + * logits (`logits.fp32.f32.bin` from the FP32 model, `logits.step.f32.bin` from this step) under + * `golden/`, for checking the compiled module on a device. + */ + @Test + fun realCheckpointGoldenStepForADevice() { + val dir = System.getenv("LFM2_MODEL_DIR")?.let(::File) ?: return + val exportDir = System.getenv("LFM2_GOLDEN_DIR")?.let(::File) ?: return + val capacity = Regex("\"capacity\"\\s*:\\s*(\\d+)").find(File(exportDir, "export.json").readText())!!.groupValues[1].toInt() + val order = File(exportDir, "io.txt").readLines().filter { it.startsWith("in ") }.map { it.split(" ")[2] } + val config = Lfm2Config.parse(File(dir, "config.json").readText()) + val reference = loadLfm2Model(config, File(dir, "model.safetensors"), contextLimit = capacity) + val quantized = loadLfm2Model(config, File(dir, "model.safetensors"), contextLimit = capacity, + ctx = ctx, weightFormFor = lfm2ProjectionsAs(TensorEncoding.Q8_0)) + val dims = quantized.model.dimensions + val embedding = quantized.weights.getValue("model.embed_tokens.weight") + val step = Lfm2DecodeStep(dims, quantized.weights, capacity, lfm2Head(embedding, TensorEncoding.Q4_0, ctx)) + + val tokens = intArrayOf(1, 6, 6423, 708, 1098, 8448) + val state = HybridModelState() + var states: Map> = step.stateShapes().mapValues { (_, s) -> tensor(s) } + var expected = FloatArray(0) + for ((position, token) in tokens.withIndex()) { + expected = reference.model.forward(intArrayOf(token), state).logits.data.copyToFloatArray() + if (position == tokens.lastIndex) break + val (c, s, mask, write) = controls(dims, capacity, position) + states = step.forward(embeddedRow(embedding, token), c, s, mask, write, states, ctx).states + } + val position = tokens.lastIndex + val (c, s, mask, write) = controls(dims, capacity, position) + // The device mask uses a bf16-safe large negative value instead of -Float.MAX_VALUE. + val deviceMask = tensor(mask.shape, mask.data.copyToFloatArray().map { if (it < 0f) -1e30f else it }.toFloatArray()) + val inputs = mapOf("embedded" to embeddedRow(embedding, tokens.last()), "cos" to c, "sin" to s, "mask" to deviceMask, "write_mask" to write) + states + val r = step.forward(inputs.getValue("embedded"), c, s, deviceMask, write, states, ctx) + + val out = File(exportDir, "golden").apply { mkdirs() } + fun writeBf16(t: Tensor, file: File) { + val bytes = java.nio.ByteBuffer.allocate(t.shape.volume * 2).order(java.nio.ByteOrder.LITTLE_ENDIAN) + for (v in t.data.copyToFloatArray()) { + val bits = java.lang.Float.floatToRawIntBits(v) + bytes.putShort(((bits + 0x7FFF + ((bits ushr 16) and 1)) ushr 16).toShort()) + } + file.writeBytes(bytes.array()) + } + fun writeF32(values: FloatArray, file: File) { + val bytes = java.nio.ByteBuffer.allocate(values.size * 4).order(java.nio.ByteOrder.LITTLE_ENDIAN) + values.forEach { bytes.putFloat(it) } + file.writeBytes(bytes.array()) + } + val args = order.map { role -> + val t = inputs.getValue(role) + writeBf16(t, File(out, "$role.bf16.bin")) + "--input=${t.shape.dimensions.joinToString("x")}xbf16=@$role.bf16.bin" + } + File(out, "args.txt").writeText(args.joinToString("\n", postfix = "\n")) + writeF32(expected, File(out, "logits.fp32.f32.bin")) + writeF32(r.logits.data.copyToFloatArray(), File(out, "logits.step.f32.bin")) + val top = expected.indices.maxBy { expected[it] } + println("golden step at position $position written to $out; FP32 top-1 $top, step top-1 ${r.logits.data.copyToFloatArray().let { l -> l.indices.maxBy { l[it] } }}") + } +} From 844e0dea47fde750bfa606bd3273da1ba4e1b764 Mon Sep 17 00:00:00 2001 From: Michal Harakal <471011+michalharakal@users.noreply.github.com> Date: Sat, 7 Nov 2026 13:14:22 +0100 Subject: [PATCH 2/2] feat(lfm2): decode-step export with external bf16 parameters --parameters external-bf16 traces the step with every weight as a model-scope parameter (traceLfm2DecodeStepExternal), writes them as BF16 safetensors and types the parameter globals bf16 with a convert to f32 where they load (writeLfm2Bf16Parameters), for a GPU that serves weights from an archive. Writes io.txt, export.json and the embedding rows as the inline mode does. Compiled for Mali valhall4 with IREE 3.11: 152 ms per step, logits cosine 0.99991 to FP32. Co-Authored-By: Claude Opus 5.5 --- llm-inference/lfm2/api/lfm2.api | 3 + .../lfm2/export/Lfm2DecodeStepExport.kt | 110 +++++++++++++++++- .../ainet/models/lfm2/Lfm2DecodeStepTest.kt | 24 ++++ 3 files changed, 135 insertions(+), 2 deletions(-) diff --git a/llm-inference/lfm2/api/lfm2.api b/llm-inference/lfm2/api/lfm2.api index 732d388a..46c54452 100644 --- a/llm-inference/lfm2/api/lfm2.api +++ b/llm-inference/lfm2/api/lfm2.api @@ -351,9 +351,12 @@ public final class sk/ainet/models/lfm2/export/Lfm2DecodeStepExportCli { public final class sk/ainet/models/lfm2/export/Lfm2DecodeStepExportKt { public static final fun exportLfm2DecodeStep (Lsk/ainet/models/lfm2/Lfm2DecodeStep;Lsk/ainet/lang/trace/PackedConstantHandling;)Lsk/ainet/lang/nn/dsl/decoder/ExportedDecodeStep; public static synthetic fun exportLfm2DecodeStep$default (Lsk/ainet/models/lfm2/Lfm2DecodeStep;Lsk/ainet/lang/trace/PackedConstantHandling;ILjava/lang/Object;)Lsk/ainet/lang/nn/dsl/decoder/ExportedDecodeStep; + public static final fun ioTable (Lsk/ainet/models/lfm2/export/Lfm2TracedFunction;)Ljava/lang/String; public static final fun lfm2DecodeStepInputs (Lsk/ainet/models/lfm2/Lfm2DecodeStep;)Ljava/util/List; public static final fun lfm2FoldableMatmuls (Ljava/util/List;)I public static final fun lfm2Head (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/context/ExecutionContext;)Lsk/ainet/lang/tensor/Tensor; + public static final fun traceLfm2DecodeStepExternal (Lsk/ainet/models/lfm2/Lfm2DecodeStep;)Lsk/ainet/models/lfm2/export/Lfm2TracedFunction; + public static final fun writeLfm2Bf16Parameters (Lsk/ainet/models/lfm2/export/Lfm2TracedFunction;Ljava/io/File;)Ljava/lang/String; } public final class sk/ainet/models/lfm2/export/Lfm2ExportBundleKt { diff --git a/llm-inference/lfm2/src/jvmMain/kotlin/sk/ainet/models/lfm2/export/Lfm2DecodeStepExport.kt b/llm-inference/lfm2/src/jvmMain/kotlin/sk/ainet/models/lfm2/export/Lfm2DecodeStepExport.kt index 84c79869..0349cbc5 100644 --- a/llm-inference/lfm2/src/jvmMain/kotlin/sk/ainet/models/lfm2/export/Lfm2DecodeStepExport.kt +++ b/llm-inference/lfm2/src/jvmMain/kotlin/sk/ainet/models/lfm2/export/Lfm2DecodeStepExport.kt @@ -77,11 +77,11 @@ object Lfm2DecodeStepExportCli { fun main(args: Array) { if (args.isEmpty() || args.contentEquals(arrayOf("--help"))) { println("LFM2 decode-step export: --config FILE --config-sha256 SHA256 --checkpoint FILE --checkpoint-sha256 SHA256 " + - "--capacity N --output NEW_DIRECTORY [--weights q8_0|q4_0|f32] [--head q4_0|q8_0|tied] [--conv multiply-sum|depthwise] [--attention batched|per-head]") + "--capacity N --output NEW_DIRECTORY [--weights q8_0|q4_0|f32] [--head q4_0|q8_0|tied] [--conv multiply-sum|depthwise] [--attention batched|per-head] [--parameters inline|external-bf16]") return } val allowed = setOf("--config", "--config-sha256", "--checkpoint", "--checkpoint-sha256", "--capacity", "--output", - "--weights", "--head", "--conv", "--attention") + "--weights", "--head", "--conv", "--attention", "--parameters") require(args.size % 2 == 0) { "Every option requires one value; use --help" } val options = linkedMapOf() for (pair in args.toList().chunked(2)) { @@ -126,6 +126,41 @@ object Lfm2DecodeStepExportCli { weightFormFor = weights?.let(::lfm2ProjectionsAs)) val embedding = loaded.weights.getValue("model.embed_tokens.weight") val step = Lfm2DecodeStep(loaded.model.dimensions, loaded.weights, capacity, lfm2Head(embedding, head, loaded.model.ctx), conv, attention) + if ((options["--parameters"] ?: "inline") == "external-bf16") { + require(weights == null && head == null) { "--parameters external-bf16 takes dense weights: --weights f32 --head tied" } + output.mkdirs() + val trace = traceLfm2DecodeStepExternal(step) + File(output, "decode_step.mlir").writeText(writeLfm2Bf16Parameters(trace, File(output, "model.bf16.safetensors"))) + File(output, "io.txt").writeText(trace.ioTable()) + val embedFile = File(output, "embed.bf16.bin") + writeBf16(embedding, embedFile) + val archive = File(output, "model.bf16.safetensors") + val d = loaded.model.dimensions + val manifest = buildJsonObject { + put("schema_version", 1); put("kind", "lfm2_decode_step_export"); put("function_name", "decode_step") + put("config_sha256", configPin); put("checkpoint_sha256", loaded.checkpointSha256) + put("capacity", capacity); put("weights", "external-bf16"); put("head", "tied") + put("conv_lowering", conv.name.lowercase()); put("attention_layout", attention.name.lowercase()) + put("parameters", buildJsonObject { + put("file", archive.name); put("scope", "model"); put("dtype", "bf16"); put("count", trace.module.externalParameters.size) + put("sha256", sha256(archive)) + }) + put("dimensions", buildJsonObject { + put("vocabulary", d.vocabulary); put("width", d.width); put("head_dimension", d.headDim) + put("kv_heads", d.kvHeads); put("conv_kernel", d.kernel); put("rope_theta", d.theta) + put("layer_types", JsonArray(d.layerTypes.map(::JsonPrimitive))) + }) + put("embedding_rows", buildJsonObject { + put("file", embedFile.name); put("dtype", "bf16"); put("shape", JsonArray(listOf(d.vocabulary, d.width).map(::JsonPrimitive))) + put("sha256", sha256(embedFile)) + }) + put("host_control", "embedding_lookup_rope_masks_position_and_stop_policy") + } + File(output, "export.json").writeText(Json { prettyPrint = true }.encodeToString(JsonObject.serializer(), manifest) + "\n") + println("Exported decode_step with ${trace.module.externalParameters.size} external bf16 parameters: capacity=$capacity, " + + "conv=${conv.name.lowercase()}, attention=${attention.name.lowercase()}; convert with iree-convert-parameters") + return + } val exported = exportLfm2DecodeStep(step) val skipped = exported.diagnostics.filter { it.startsWith("BLOCK_QUANT_SKIPPED") } require(skipped.isEmpty()) { "Packed weights the fusion could not fold:\n" + skipped.joinToString("\n") } @@ -190,3 +225,74 @@ object Lfm2DecodeStepExportCli { digest.digest().joinToString("") { "%02x".format(it) } } } + +/** + * [step] traced with every weight as an external `model`-scope parameter, for a compiler that serves the weights from + * an archive (IREE `--parameters`, e.g. a GPU): `module` keeps f32 types and `inputNames`/`outputNames` the step's + * ports in argument order. Use [writeLfm2Bf16Parameters] to store the weights as bf16. + */ +fun traceLfm2DecodeStepExternal(step: Lfm2DecodeStep): Lfm2TracedFunction = + traceLfm2Function("decode_step", LinkedHashMap(lfm2DecodeStepInputs(step).toMap())) { p, ctx -> + val states = step.stateShapes().keys.associateWith { p.getValue(it) } + val r = step.forward(p.getValue("embedded"), p.getValue("cos"), p.getValue("sin"), p.getValue("mask"), p.getValue("write_mask"), states, ctx) + linkedMapOf("logits" to r.logits).apply { putAll(r.states) } + } + +/** + * Writes [trace]'s external parameters to [archive] as flat BF16 safetensors (round to nearest even; exact for a + * bf16 checkpoint) and returns the module text with every parameter global typed bf16 and converted to f32 where it + * is loaded, so the graph computes as before and the weights take half the memory. + */ +fun writeLfm2Bf16Parameters(trace: Lfm2TracedFunction, archive: File): String { + require(!archive.exists()) { "Parameter archive already exists: ${archive.name}" } + val resolver = sk.ainet.lang.tensor.storage.DefaultBufferResolver() + val params = trace.module.externalParameters + require(params.isNotEmpty() && params.all { it.scope == "model" }) + val sizes = params.map { p -> resolver.resolve(p.source).use { it.sizeInBytes } } + require(sizes.all { it > 0 && it % 4 == 0L }) + var offset = 0L + val header = buildJsonObject { + for ((i, p) in params.withIndex()) { + val n = sizes[i] / 4 + put(p.key, buildJsonObject { + put("dtype", "BF16"); put("shape", buildJsonArray { add(n) }) + put("data_offsets", buildJsonArray { add(offset); add(offset + 2 * n) }) + }) + offset += 2 * n + } + }.toString().toByteArray(Charsets.UTF_8) + val padded = header.size + (8 - header.size % 8) % 8 + FileOutputStream(archive).buffered(1 shl 20).use { out -> + out.write(ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN).putLong(padded.toLong()).array()) + out.write(header) + repeat(padded - header.size) { out.write(32) } + for ((i, p) in params.withIndex()) resolver.resolve(p.source).use { source -> + var position = 0L + while (position < sizes[i]) { + val count = minOf(1L shl 20, sizes[i] - position).toInt() + val f32 = ByteBuffer.wrap(source.readBytes(position, count)).order(ByteOrder.LITTLE_ENDIAN) + val bf16 = ByteBuffer.allocate(count / 2).order(ByteOrder.LITTLE_ENDIAN) + while (f32.hasRemaining()) { + val bits = f32.int + bf16.putShort(((bits + 0x7FFF + ((bits ushr 16) and 1)) ushr 16).toShort()) + } + out.write(bf16.array()) + position += count + } + } + } + var mlir = Regex("""(util\.global private @[\w$.-]+ = #flow\.parameter\.named<"[^"]*"::"[^"]*"> : tensor<[0-9x]*x)f32>""") + .replace(trace.module.content) { it.groupValues[1] + "bf16>" } + mlir = Regex("""(%[\w$.-]+) = util\.global\.load @([\w$.-]+) : tensor<([0-9x]*)xf32>""").replace(mlir) { r -> + val (ssa, g, shape) = r.destructured + "${ssa}_bf16 = util.global.load @$g : tensor<${shape}xbf16>\n $ssa = stablehlo.convert ${ssa}_bf16 : (tensor<${shape}xbf16>) -> tensor<${shape}xf32>" + } + val converted = Regex("""util\.global\.load @[\w$.-]+ : tensor<[0-9x]*xbf16>""").findAll(mlir).count() + check(converted == params.size) { "Retyped $converted of ${params.size} parameter loads; the exporter's global form changed" } + return mlir +} + +/** `io.txt` for a traced step: `in [dims]` in argument order, then `out [dims]`. */ +fun Lfm2TracedFunction.ioTable(): String = + (inputNames.zip(inputShapes).mapIndexed { i, (n, s) -> "in $i $n ${s.dimensions.toList()}" } + + outputNames.zip(outputShapes).mapIndexed { i, (n, s) -> "out $i $n ${s.dimensions.toList()}" }).joinToString("\n", postfix = "\n") diff --git a/llm-inference/lfm2/src/jvmTest/kotlin/sk/ainet/models/lfm2/Lfm2DecodeStepTest.kt b/llm-inference/lfm2/src/jvmTest/kotlin/sk/ainet/models/lfm2/Lfm2DecodeStepTest.kt index 9efe0c9d..43f7ba4e 100644 --- a/llm-inference/lfm2/src/jvmTest/kotlin/sk/ainet/models/lfm2/Lfm2DecodeStepTest.kt +++ b/llm-inference/lfm2/src/jvmTest/kotlin/sk/ainet/models/lfm2/Lfm2DecodeStepTest.kt @@ -11,6 +11,9 @@ import sk.ainet.models.lfm2.export.exportLfm2DecodeStep import sk.ainet.models.lfm2.export.lfm2DecodeStepInputs import sk.ainet.models.lfm2.export.lfm2FoldableMatmuls import sk.ainet.models.lfm2.export.lfm2Head +import sk.ainet.models.lfm2.export.ioTable +import sk.ainet.models.lfm2.export.traceLfm2DecodeStepExternal +import sk.ainet.models.lfm2.export.writeLfm2Bf16Parameters import java.io.File import kotlin.math.abs import kotlin.math.cos @@ -128,6 +131,27 @@ class Lfm2DecodeStepTest { assertEquals(25, depthwise.foldedBlockQuantMatmuls) } + @Test + fun externalExportServesEveryWeightFromABf16Archive() { + val weights = weightsFor(small) + val step = Lfm2DecodeStep(small, weights, 6, weights.getValue("model.embed_tokens.weight")) + val trace = traceLfm2DecodeStepExternal(step) + assertEquals(lfm2DecodeStepInputs(step).map { it.first }, trace.inputNames) + assertEquals(listOf("logits") + step.stateShapes().keys, trace.outputNames) + // the tied embedding enters the graph once, as the head: every checkpoint tensor is one parameter + assertEquals(weights.size, trace.module.externalParameters.size) + val archive = File.createTempFile("lfm2-step", ".safetensors").apply { delete(); deleteOnExit() } + val mlir = writeLfm2Bf16Parameters(trace, archive) + assertFalse(Regex("""#flow\.parameter\.named<[^>]*> : tensor<[0-9x]*xf32>""").containsMatchIn(mlir)) + assertEquals(weights.size, Regex("""stablehlo\.convert %[\w$.-]+_bf16""").findAll(mlir).count()) + val bytes = archive.readBytes() + val headerSize = java.nio.ByteBuffer.wrap(bytes, 0, 8).order(java.nio.ByteOrder.LITTLE_ENDIAN).long.toInt() + val header = String(bytes, 8, headerSize) + assertEquals(weights.size, Regex("\"BF16\"").findAll(header).count()) + assertEquals(8L + headerSize + weights.values.sumOf { it.shape.volume * 2L }, archive.length()) + assertTrue(trace.ioTable().lines().first().startsWith("in 0 embedded [1, 8]")) + } + /** * Gated on `LFM2_MODEL_DIR` (a local LFM2.5-350M with `config.json` and `model.safetensors`): * the Q8_0 projections and Q4_0 head of the NPU step against the FP32 model over a short