feat(laya): engine 0.61.0 and the packed Q8_0 lane (SKEEP-007 consumer)

CHANGELOG.md

@@ -11,7 +11,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
-## [0.61.0] — 2026-10-16
+Against engine **0.61.0** (SKEEP-007, the browser lane). The BOM re-exports `skainet-bom:0.61.0`.
+Engine 0.61.0 removed its SKEEP-006 compatibility aliases as scheduled; nothing here referenced them.
+
+### Added
+
+- **Laya on the packed Q8_0 lane** (SKEEP-007 Tier 2 consumer). `LayaModelLoader.load` is common code: the four side files as text plus a `RandomAccessSource` for `model.safetensors`, so the JVM, Android and browser builds share one loader; the JVM `load(directory: Path, …)` is now an extension over it (source-compatible). `weightForm` / `weightFormFor` pass through to the engine's SafeTensors loader, and `LayaWeightForms.q8_0For` quantizes the encoder and head projections and the 256k × 768 token table to Q8_0 at load — about a quarter of the FP32 bytes, served by the packed kernels on every tier. `LayaWeightPreparation` splits the fused `Wqkv` / `Wi` / `in_proj` matrices in their packed form (row blocks are contiguous byte ranges) and wraps the token table as a `PackedRowDequantTensorData`, so the embedding reads rows on demand.
+- **Decision-quality gate for the quantized lane**: `LayaQuantizedParityTest` (model-gated) requires every `predict.json` choice of the Q8_0 agent to agree with the FP32 agent's and no probability, score, noul value or confidence to shift by more than 0.05 (`LAYA_Q8_TOL`; measured 2026-10-16 on the multilingual checkpoint: 5/5 choices agree, 45 values, max shift 0.035). FP32 parity (1.5e-4) remains the contract of the reference lane. On the eager JVM path the Q8_0 lane is a memory measure, not a speed one — the row-major packed kernels run one activation row at a time, so the six fixture cases take 173 s against 69 s in FP32 on a 12-core laptop; the speed comes with the browser sidecar and the compiled lane.
+
 
 **Padded batches for the encoders.** Against engine **0.60.0**, which gives
 `scaledDotProductAttention` key lengths, a window and a right context carried as op metadata

docs/modules/ROOT/pages/reference/typed-decisions.adoc

@@ -192,11 +192,24 @@ sequences, the 40-text tokenizer corpus (CJK, RTL, combining marks, emoji, liter
 and all six `predict.json` cases within 5e-5. The model-gated tests run when
 `LAYA_MULTILINGUAL_PATH` points at the checkpoint directory.
 
+== Quantized lane
+
+`LayaModelLoader.load(files, weights, ctx, weightFormFor = LayaWeightForms.q8_0For)` loads the
+encoder and head projections and the token table as Q8_0 at load (engine 0.61 quantize-at-load,
+SKEEP-007): about a quarter of the FP32 bytes, served by the packed kernels on every tier — on the
+JVM the FFM pack, on Android the JNI pack, in a browser the `skainet-backend-wasm-simd` sidecar after
+`WasmSimd.install()`. The contract of that lane is decision quality rather than FP32 parity:
+`LayaQuantizedParityTest` requires every `predict.json` choice to agree with the FP32 agent and no
+probability, score, noul value or confidence to shift by more than 0.05 (measured: 5/5 choices,
+max shift 0.035). On the eager JVM path this lane trades speed for memory — the packed kernels run
+one activation row at a time — so use it where the bytes matter: the browser, a phone, a small heap.
+
 == Current limits
 
-* JVM only: the loader uses `java.nio.file.Path`; the common code has no file access of its own.
-* Eager FP32 on the CPU path, no batching across questions: a short state with three questions
-  takes seconds on a laptop, a 1024-token state tens of seconds. The compiled and quantized
-  lanes are open items.
+* The loader is common code (a `RandomAccessSource` plus the four side files as text); obtaining the
+  bytes — a file, an app's files directory, a download held in memory — is the caller's part.
+* Eager FP32 on the CPU path by default; batching across the questions of one request since 0.61
+  saves about 10 % per question there. The compiled lane is an open item; the quantized lane is
+  the section above.
 * The base checkpoint is a starting point the publisher expects to be specialised; keep option
   sets small (six or fewer) and treat abstentions as the normal case, not the exception.

gradle/libs.versions.toml

@@ -1,5 +1,5 @@
 [versions]
-skainet = "0.60.0"
+skainet = "0.61.0"
 agp = "9.4.1"
 jacksonDatabind = "2.22.3"
 jsonSchemaValidator = "3.0.7"

llm-inference/laya/api/jvm/laya.api

@@ -306,8 +306,13 @@ public final class sk/ainet/models/laya/LayaLinear : sk/ainet/lang/nn/Module, sk
 
 public final class sk/ainet/models/laya/LayaModelLoader {
 	public static final field INSTANCE Lsk/ainet/models/laya/LayaModelLoader;
-	public final fun load (Ljava/nio/file/Path;Lsk/ainet/context/ExecutionContext;Z)Lsk/ainet/models/laya/LayaAgent;
-	public static synthetic fun load$default (Lsk/ainet/models/laya/LayaModelLoader;Ljava/nio/file/Path;Lsk/ainet/context/ExecutionContext;ZILjava/lang/Object;)Lsk/ainet/models/laya/LayaAgent;
+	public final fun load (Lsk/ainet/models/laya/LayaCheckpointFiles;Lkotlin/jvm/functions/Function0;Lsk/ainet/context/ExecutionContext;ZLsk/ainet/lang/memory/plan/WeightForm;Lkotlin/jvm/functions/Function1;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+	public static synthetic fun load$default (Lsk/ainet/models/laya/LayaModelLoader;Lsk/ainet/models/laya/LayaCheckpointFiles;Lkotlin/jvm/functions/Function0;Lsk/ainet/context/ExecutionContext;ZLsk/ainet/lang/memory/plan/WeightForm;Lkotlin/jvm/functions/Function1;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
+}
+
+public final class sk/ainet/models/laya/LayaModelLoaderKt {
+	public static final fun load (Lsk/ainet/models/laya/LayaModelLoader;Ljava/nio/file/Path;Lsk/ainet/context/ExecutionContext;ZLkotlin/jvm/functions/Function1;)Lsk/ainet/models/laya/LayaAgent;
+	public static synthetic fun load$default (Lsk/ainet/models/laya/LayaModelLoader;Ljava/nio/file/Path;Lsk/ainet/context/ExecutionContext;ZLkotlin/jvm/functions/Function1;ILjava/lang/Object;)Lsk/ainet/models/laya/LayaAgent;
 }
 
 public final class sk/ainet/models/laya/LayaNetwork {
@@ -608,6 +613,12 @@ public final class sk/ainet/models/laya/LayaTruncation {
 	public fun toString ()Ljava/lang/String;
 }
 
+public final class sk/ainet/models/laya/LayaWeightForms {
+	public static final field INSTANCE Lsk/ainet/models/laya/LayaWeightForms;
+	public final fun getQ8_0For ()Lkotlin/jvm/functions/Function1;
+	public final fun isPacked (Ljava/lang/String;)Z
+}
+
 public final class sk/ainet/models/laya/LayaWeightPreparation {
 	public static final field INSTANCE Lsk/ainet/models/laya/LayaWeightPreparation;
 	public final fun prepare (Ljava/util/List;Lsk/ainet/models/laya/LayaEncoderConfig;Lsk/ainet/models/laya/LayaAgentConfig;Lsk/ainet/context/ExecutionContext;Lkotlin/reflect/KClass;)Lsk/ainet/models/laya/LayaPreparedWeights;

llm-inference/laya/src/commonMain/kotlin/sk/ainet/models/laya/LayaModelLoader.kt

@@ -0,0 +1,74 @@
+package sk.ainet.models.laya
+
+import sk.ainet.context.ExecutionContext
+import sk.ainet.io.RandomAccessSource
+import sk.ainet.io.safetensors.SafeTensorsParametersLoader
+import sk.ainet.io.weights.WeightTensor
+import sk.ainet.lang.memory.plan.EncodingRequest
+import sk.ainet.lang.memory.plan.WeightForm
+import sk.ainet.lang.memory.plan.WeightShapeOrientation
+import sk.ainet.lang.tensor.storage.TensorEncoding
+import sk.ainet.lang.types.FP32
+
+/**
+ * Loads a Laya checkpoint — the four side files as text ([LayaCheckpointFiles]) and
+ * `model.safetensors` through a [RandomAccessSource] — into a [LayaAgent].
+ *
+ * Common code since engine 0.61 (SKEEP-007): the JVM, Android and browser builds all run this
+ * function; only how the bytes are obtained differs (a file, an app's files directory, a download
+ * held in memory). The FP32 reference path (REQ-EXE-008) is the default: F16 weights are up-cast
+ * on load, which is what the Python reference does on CPU.
+ *
+ * [weightForm] / [weightFormFor] are passed through to the engine's SafeTensors loader. With
+ * [LayaWeightForms.q8_0For] the encoder and head projections and the token table are quantized to
+ * Q8_0 at load — about a quarter of the FP32 bytes, served by the packed kernels on every tier —
+ * and the decision-quality gate (`LayaQuantizedParityTest`) is the contract that lane keeps.
+ */
+public object LayaModelLoader {
+    public suspend fun load(
+        files: LayaCheckpointFiles,
+        weights: () -> RandomAccessSource,
+        ctx: ExecutionContext,
+        exactGelu: Boolean = true,
+        weightForm: WeightForm? = null,
+        weightFormFor: ((String) -> WeightForm?)? = null,
+    ): LayaAgent<FP32> {
+        val preprocessor = LayaPreprocessor.fromFiles(files)
+        val raw = ArrayList<WeightTensor<FP32, Float>>()
+        SafeTensorsParametersLoader(sourceProvider = weights, weightForm = weightForm, weightFormFor = weightFormFor)
+            .load<FP32, Float>(ctx, FP32::class) { name, tensor ->
+                raw += WeightTensor(name, tensor.shape.dimensions.toList(), tensor)
+            }
+        val prepared = LayaWeightPreparation.prepare(raw, preprocessor.encoderConfig, preprocessor.agentConfig, ctx, FP32::class)
+        val model = LayaDecisionModel.build(preprocessor.encoderConfig, preprocessor.agentConfig, prepared, ctx, FP32::class, exactGelu)
+        return LayaAgent(preprocessor, model)
+    }
+}
+
+/**
+ * Weight forms for a Laya checkpoint (SKEEP-007 Tier 2 consumer).
+ *
+ * The matrices that carry the compute — every encoder projection (`Wqkv`, `Wo`, `Wi`, `mlp.Wo`),
+ * the head's attention and feed-forward weights — and the 256k × 768 token table are the ones
+ * worth packing; norms, biases, the three-row type embedding and the two scorer layers stay dense,
+ * where they cost nothing and where a packed form would only complicate the small ops that read
+ * them (`narrow` on the type embedding, the final Linear 768→1).
+ */
+public object LayaWeightForms {
+    private val q8_0: WeightForm = WeightForm(
+        encoding = EncodingRequest.RequantizeTo(TensorEncoding.Q8_0),
+        shape = WeightShapeOrientation.OUT_IN,
+    )
+
+    private val packedSuffixes = listOf(
+        ".attn.Wqkv.weight", ".attn.Wo.weight", ".mlp.Wi.weight", ".mlp.Wo.weight",
+        ".self_attn.in_proj_weight", ".self_attn.out_proj.weight", ".linear1.weight", ".linear2.weight",
+    )
+
+    /** `true` for the tensors [q8_0For] packs. */
+    public fun isPacked(name: String): Boolean =
+        name == "encoder.embeddings.tok_embeddings.weight" || packedSuffixes.any { name.endsWith(it) }
+
+    /** Per-tensor form: Q8_0 for the projection matrices and the token table, dense for everything else. */
+    public val q8_0For: (String) -> WeightForm? = { name -> if (isPacked(name)) q8_0 else null }
+}

llm-inference/laya/src/commonMain/kotlin/sk/ainet/models/laya/LayaWeights.kt

@@ -7,6 +7,14 @@ import sk.ainet.lang.tensor.Shape
 import sk.ainet.lang.tensor.Tensor
 import sk.ainet.lang.types.DType
 import kotlin.reflect.KClass
+import sk.ainet.lang.memory.BlockOrder
+import sk.ainet.lang.nn.quant.PackedRowDequantTensorData
+import sk.ainet.lang.tensor.data.Q4_0BlockTensorData
+import sk.ainet.lang.tensor.data.Q8_0BlockTensorData
+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.types.FP32
 
 /**
  * Maps module parameters of [LayaNetwork] to checkpoint tensor names after [LayaWeightPreparation]
@@ -76,7 +84,12 @@ public object LayaWeightPreparation {
             WeightTensor(name, listOf(size), ctx.fromFloatArray(Shape(size), dtype, FloatArray(size)))
 
         // --- encoder ---------------------------------------------------------------------------
-        encoder += take("encoder.embeddings.tok_embeddings.weight")
+        // A packed token table (SKEEP-007 Tier 2 lane) is read row by row by the embedding layer
+        // through RowDequantSource; dense tables pass through unchanged.
+        encoder += take("encoder.embeddings.tok_embeddings.weight").let { t ->
+            val wrapped = PackedRowDequantTensorData.wrapIfRowDequantable(t.tensor.data as TensorData<FP32, Float>)
+            if (wrapped === t.tensor.data) t else WeightTensor(t.name, t.shape, ctx.fromData(wrapped as TensorData<T, V>, dtype))
+        }
         encoder += take("encoder.embeddings.norm.weight")
         encoder += byName["encoder.embeddings.norm.bias"]?.also { consumed += it.name } ?: zeros("encoder.embeddings.norm.bias", d)
         for (i in 0 until encoderConfig.numHiddenLayers) {
@@ -146,12 +159,20 @@ public object LayaWeightPreparation {
         }
     }
 
-    /** Splits a `[rows, …]` tensor into [parts] equal row blocks, copying each into its own tensor. */
+    /**
+     * Splits a `[rows, …]` tensor into [parts] equal row blocks, copying each into its own tensor.
+     *
+     * A packed weight stays packed (SKEEP-007 Tier 2): in canonical row-major block storage a row
+     * block is a contiguous byte range (`rowsPerPart × blocksPerRow × bytesPerBlock`), so the fused
+     * `Wqkv` / `Wi` / `in_proj` matrices split without an FP32 transient and the parts are served by
+     * the packed kernels like every other projection. Dense weights take the float path as before.
+     */
     internal fun <T : DType, V> splitRows(source: WeightTensor<T, V>, parts: Int, ctx: ExecutionContext, dtype: KClass<T>): List<Tensor<T, V>> {
         val rows = source.shape[0]
         require(rows % parts == 0) { "cannot split ${source.name} with $rows rows into $parts parts" }
         val rowsPerPart = rows / parts
         val inner = source.shape.drop(1).fold(1) { a, b -> a * b }
+        splitPackedRows(source, parts, rowsPerPart, inner, ctx, dtype)?.let { return it }
         val block = rowsPerPart * inner
         val data = source.tensor.data.copyToFloatArray()
         return List(parts) { part ->
@@ -160,4 +181,28 @@ public object LayaWeightPreparation {
             ctx.fromFloatArray(shape, dtype, out)
         }
     }
+
+    @Suppress("UNCHECKED_CAST")
+    private fun <T : DType, V> splitPackedRows(
+        source: WeightTensor<T, V>, parts: Int, rowsPerPart: Int, inner: Int, ctx: ExecutionContext, dtype: KClass<T>,
+    ): List<Tensor<T, V>>? {
+        val packed = source.tensor.data as? PackedBlockStorage ?: return null
+        if (source.shape.size != 2 || packed.blockOrder != BlockOrder.ROW_MAJOR) return null
+        if (inner % packed.blockSize != 0) return null
+        val blocksPerRow = inner / packed.blockSize
+        val totalBlocks = source.shape[0] * blocksPerRow
+        if (packed.blockCount != totalBlocks) return null
+        val bytesPerBlock = packed.packedData.size / totalBlocks
+        val partBytes = rowsPerPart * blocksPerRow * bytesPerBlock
+        val shape = Shape(rowsPerPart, inner)
+        return List(parts) { part ->
+            val bytes = packed.packedData.copyOfRange(part * partBytes, (part + 1) * partBytes)
+            val data: TensorData<*, *> = when (packed.encoding) {
+                TensorEncoding.Q8_0 -> Q8_0BlockTensorData(shape, bytes)
+                TensorEncoding.Q4_0 -> Q4_0BlockTensorData(shape, bytes)
+                else -> return null
+            }
+            ctx.fromData(data as TensorData<T, V>, dtype)
+        }
+    }
 }

llm-inference/laya/src/jvmMain/kotlin/sk/ainet/models/laya/LayaModelLoader.kt

@@ -3,32 +3,35 @@ package sk.ainet.models.laya
 import kotlinx.coroutines.runBlocking
 import sk.ainet.context.ExecutionContext
 import sk.ainet.io.JvmRandomAccessSource
-import sk.ainet.io.safetensors.SafeTensorsParametersLoader
-import sk.ainet.io.weights.WeightTensor
+import sk.ainet.lang.memory.plan.WeightForm
 import sk.ainet.lang.types.FP32
 import java.nio.file.Files
 import java.nio.file.Path
 
 /**
  * Loads a checkpoint directory (`rl_agent_config.json`, `encoder/config.json`, `tokenizer/`,
- * `model.safetensors`) into a [LayaAgent] on the FP32 CPU reference path (REQ-EXE-008). F16
- * weights are up-cast on load, which is what the Python reference does on CPU.
+ * `model.safetensors`) into a [LayaAgent] — the JVM convenience over the common
+ * [LayaModelLoader.load]. Blocking; call off the UI thread.
+ *
+ * [weightFormFor] selects the lane: `null` is the FP32 reference path (REQ-EXE-008, F16 up-cast on
+ * load), [LayaWeightForms.q8_0For] the packed Q8_0 lane.
  */
-public object LayaModelLoader {
-
-    public fun load(directory: Path, ctx: ExecutionContext, exactGelu: Boolean = true): LayaAgent<FP32> {
-        val preprocessor = LayaPreprocessor.fromFiles(LayaCheckpointFiles.read(directory))
-        val weights = directory.resolve("model.safetensors")
-        require(Files.exists(weights)) { "no model.safetensors in $directory" }
-        val raw = ArrayList<WeightTensor<FP32, Float>>()
-        runBlocking {
-            SafeTensorsParametersLoader(sourceProvider = { JvmRandomAccessSource.open(weights.toString()) })
-                .load<FP32, Float>(ctx, FP32::class) { name, tensor ->
-                    raw += WeightTensor(name, tensor.shape.dimensions.toList(), tensor)
-                }
-        }
-        val prepared = LayaWeightPreparation.prepare(raw, preprocessor.encoderConfig, preprocessor.agentConfig, ctx, FP32::class)
-        val model = LayaDecisionModel.build(preprocessor.encoderConfig, preprocessor.agentConfig, prepared, ctx, FP32::class, exactGelu)
-        return LayaAgent(preprocessor, model)
+public fun LayaModelLoader.load(
+    directory: Path,
+    ctx: ExecutionContext,
+    exactGelu: Boolean = true,
+    weightFormFor: ((String) -> WeightForm?)? = null,
+): LayaAgent<FP32> {
+    val weights = directory.resolve("model.safetensors")
+    require(Files.exists(weights)) { "no model.safetensors in $directory" }
+    val files = LayaCheckpointFiles.read(directory)
+    return runBlocking {
+        load(
+            files = files,
+            weights = { JvmRandomAccessSource.open(weights.toString()) },
+            ctx = ctx,
+            exactGelu = exactGelu,
+            weightFormFor = weightFormFor,
+        )
     }
 }

llm-inference/laya/src/jvmTest/kotlin/sk/ainet/models/laya/LayaFixtures.kt

@@ -6,6 +6,10 @@ import kotlinx.serialization.json.JsonObject
 import kotlinx.serialization.json.jsonObject
 import java.nio.file.Files
 import java.nio.file.Path
+import kotlinx.serialization.json.JsonArray
+import kotlinx.serialization.json.JsonNull
+import kotlinx.serialization.json.JsonPrimitive
+import kotlinx.serialization.json.jsonPrimitive
 
 /**
  * Golden fixtures of the Laya JVM port (Apache-2.0, see `laya-fixtures/NOTICE.md`) and the
@@ -41,4 +45,26 @@ internal object LayaFixtures {
     fun ints(element: JsonElement?): List<Int> = (element as? kotlinx.serialization.json.JsonArray)?.map { it.jsonObjectOrPrimitiveInt() } ?: emptyList()
 
     private fun JsonElement.jsonObjectOrPrimitiveInt(): Int = (this as kotlinx.serialization.json.JsonPrimitive).content.toInt()
+
+    /** The public request shape (`type`, `instructions`, `criteria`, `labels`), normalised as `Agent._to_internal` does. */
+    fun questionOf(q: JsonObject): LayaQuestion {
+        val insElement = q["instructions"]!!
+        val ins = if (insElement is JsonPrimitive && insElement.isString) insElement.content else PythonJson.dumps(insElement)
+        val crit = q["criteria"]
+        return when (q["type"]!!.jsonPrimitive.content) {
+            "choice" -> when (crit) {
+                is JsonArray -> LayaQuestion.Choice(ins, crit.map { it.jsonPrimitive.content to null })
+                is JsonObject -> LayaQuestion.Choice(ins, crit.entries.map { (k, v) -> k to v.takeIf { it !is JsonNull } })
+                else -> error("choice without criteria")
+            }
+            "score" -> LayaQuestion.Score(ins, (crit as JsonArray).toList())
+            "noul" -> {
+                val c = crit as? JsonObject
+                val labels = (q["labels"] as? JsonObject)?.let { it["false"]!!.jsonPrimitive.content to it["true"]!!.jsonPrimitive.content }
+                fun get(key: String): JsonElement? = c?.entries?.firstOrNull { it.key.lowercase() == key }?.value?.takeIf { it !is JsonNull }
+                LayaQuestion.Noul(ins, get("false"), get("true"), labels)
+            }
+            else -> error("unknown type ${q["type"]}")
+        }
+    }
 }