michalharakal · Nov 8, 2026 · Nov 7, 2026 · Nov 7, 2026
diff --git 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
@@ -81,44 +81,75 @@ public object DecodeStepExport {
         packedConstants: PackedConstantHandling = PackedConstantHandling.KEEP_PACKED,
         step: (DecodeStepInputs<T, V>, ExecutionContext) -> KvStepOutput<T, V>,
     ): 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<T, V>(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<String, Tensor<T, V>>("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 <T : DType, V> exportPorts(
+        inputs: List<Pair<String, Shape>>,
+        dtype: KClass<T>,
+        functionName: String = "decode_step",
+        packedConstants: PackedConstantHandling = PackedConstantHandling.KEEP_PACKED,
+        step: (Map<String, Tensor<T, V>>, ExecutionContext) -> Map<String, Tensor<T, V>>,
+    ): 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<String, String>()
-        fun input(role: String, shape: Shape): Tensor<T, V> =
-            graphInput<T, V>(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<String, Tensor<T, V>>()
+        for ((role, shape) in inputs) {
+            val t = graphInput<T, V>(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<T, V>? = null
+        var out: Map<String, Tensor<T, V>>? = 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<String, String>()
-        fun outRole(t: Tensor<T, V>, 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
@@ -209,6 +209,29 @@ class DecoderKvModelTest {
         assertTrue(step.ioTable().lines().first().startsWith("in 0 "))
     }
 
+    @Test fun exportPortsNamesArbitraryStateAndRejectsAnAliasedResult() {
+        val step = DecodeStepExport.exportPorts<FP32, Float>(
+            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<FP32, Float>(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
@@ -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_<l>` `[width, kernel]`, `key_<l>`/`value_<l>` `[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 <sha> \
+  --checkpoint model/model.safetensors --checkpoint-sha256 <sha> --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
@@ -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 <init> (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 <init> (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 <init> (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 <init> (IILsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)V
+	public fun <init> (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,22 @@ 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 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 {
 	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
@@ -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<Test>().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<JavaExec>("exportLfm2") {
     maxHeapSize = "8g"
 }
 
+tasks.register<JavaExec>("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<JavaExec>("exportLfm2Generation") {
     group = "export"
     description = "Export the fixed-capacity LFM2 cached generation graph"