diff --git a/CHANGELOG.md b/CHANGELOG.md index d444910d..9feca3fd 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/docs/modules/ROOT/pages/reference/typed-decisions.adoc b/docs/modules/ROOT/pages/reference/typed-decisions.adoc index 70527526..7100a45a 100644 --- a/docs/modules/ROOT/pages/reference/typed-decisions.adoc +++ b/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. diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml index afe41eef..5f7614b5 100644 --- a/gradle/libs.versions.toml +++ b/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" diff --git a/llm-inference/laya/api/jvm/laya.api b/llm-inference/laya/api/jvm/laya.api index ba5dca58..f2a5fafb 100644 --- a/llm-inference/laya/api/jvm/laya.api +++ b/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; diff --git a/llm-inference/laya/src/commonMain/kotlin/sk/ainet/models/laya/LayaModelLoader.kt b/llm-inference/laya/src/commonMain/kotlin/sk/ainet/models/laya/LayaModelLoader.kt new file mode 100644 index 00000000..8f90d26c --- /dev/null +++ b/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 { + val preprocessor = LayaPreprocessor.fromFiles(files) + val raw = ArrayList>() + SafeTensorsParametersLoader(sourceProvider = weights, weightForm = weightForm, weightFormFor = weightFormFor) + .load(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 } +} diff --git a/llm-inference/laya/src/commonMain/kotlin/sk/ainet/models/laya/LayaWeights.kt b/llm-inference/laya/src/commonMain/kotlin/sk/ainet/models/laya/LayaWeights.kt index 2530f236..d23a30ba 100644 --- a/llm-inference/laya/src/commonMain/kotlin/sk/ainet/models/laya/LayaWeights.kt +++ b/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) + if (wrapped === t.tensor.data) t else WeightTensor(t.name, t.shape, ctx.fromData(wrapped as TensorData, 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 splitRows(source: WeightTensor, parts: Int, ctx: ExecutionContext, dtype: KClass): List> { 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 splitPackedRows( + source: WeightTensor, parts: Int, rowsPerPart: Int, inner: Int, ctx: ExecutionContext, dtype: KClass, + ): List>? { + 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, dtype) + } + } } diff --git a/llm-inference/laya/src/jvmMain/kotlin/sk/ainet/models/laya/LayaModelLoader.kt b/llm-inference/laya/src/jvmMain/kotlin/sk/ainet/models/laya/LayaModelLoader.kt index 35e7d3f4..b7ad17b5 100644 --- a/llm-inference/laya/src/jvmMain/kotlin/sk/ainet/models/laya/LayaModelLoader.kt +++ b/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 { - 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>() - runBlocking { - SafeTensorsParametersLoader(sourceProvider = { JvmRandomAccessSource.open(weights.toString()) }) - .load(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 { + 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, + ) } } diff --git a/llm-inference/laya/src/jvmTest/kotlin/sk/ainet/models/laya/LayaFixtures.kt b/llm-inference/laya/src/jvmTest/kotlin/sk/ainet/models/laya/LayaFixtures.kt index 18343d00..2cb6c4b4 100644 --- a/llm-inference/laya/src/jvmTest/kotlin/sk/ainet/models/laya/LayaFixtures.kt +++ b/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 = (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"]}") + } + } } diff --git a/llm-inference/laya/src/jvmTest/kotlin/sk/ainet/models/laya/LayaPredictFixtureTest.kt b/llm-inference/laya/src/jvmTest/kotlin/sk/ainet/models/laya/LayaPredictFixtureTest.kt index fdfe8483..ab79b80e 100644 --- a/llm-inference/laya/src/jvmTest/kotlin/sk/ainet/models/laya/LayaPredictFixtureTest.kt +++ b/llm-inference/laya/src/jvmTest/kotlin/sk/ainet/models/laya/LayaPredictFixtureTest.kt @@ -41,7 +41,7 @@ class LayaPredictFixtureTest { for ((name, element) in fixture["single"]!!.jsonObject) { val case = element.jsonObject val state = LayaState.of(case["state"]!!) - val questions = case["questions"]!!.jsonObject.mapValues { (_, q) -> questionOf(q.jsonObject) } + val questions = case["questions"]!!.jsonObject.mapValues { (_, q) -> LayaFixtures.questionOf(q.jsonObject) } val t1 = System.nanoTime() val results = agent.decideAll(state, questions) println("[laya] $name: ${questions.size} questions in ${(System.nanoTime() - t1) / 1_000_000} ms") @@ -92,7 +92,7 @@ class LayaPredictFixtureTest { val fixture = LayaFixtures.resource("predict.json") val cases = fixture["single"]!!.jsonObject.entries.map { (name, element) -> val case = element.jsonObject - Triple(name, LayaState.of(case["state"]!!), case["questions"]!!.jsonObject.mapValues { (_, q) -> questionOf(q.jsonObject) }) + Triple(name, LayaState.of(case["state"]!!), case["questions"]!!.jsonObject.mapValues { (_, q) -> LayaFixtures.questionOf(q.jsonObject) }) } // Warm-up: the first pass pays for lazy allocations. cases.first().let { (_, state, questions) -> agent.decideAll(state, questions) } @@ -140,26 +140,4 @@ class LayaPredictFixtureTest { println("[laya] batched vs single logits: $compared compared, max |diff| = ${"%.3e".format(maxDiff)}") if (maxDiff > 1e-4) fail("batched vs single logits diverge: max |diff| = $maxDiff") } - - /** The public request shape (`type`, `instructions`, `criteria`, `labels`), normalised as `Agent._to_internal` does. */ - private 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"]}") - } - } } diff --git a/llm-inference/laya/src/jvmTest/kotlin/sk/ainet/models/laya/LayaQuantizedParityTest.kt b/llm-inference/laya/src/jvmTest/kotlin/sk/ainet/models/laya/LayaQuantizedParityTest.kt new file mode 100644 index 00000000..bf3c2d02 --- /dev/null +++ b/llm-inference/laya/src/jvmTest/kotlin/sk/ainet/models/laya/LayaQuantizedParityTest.kt @@ -0,0 +1,99 @@ +package sk.ainet.models.laya + +import kotlinx.serialization.json.jsonObject +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.lang.tensor.storage.PackedBlockStorage +import java.nio.file.Files +import kotlin.math.abs +import kotlin.test.Test +import kotlin.test.assertTrue +import kotlin.test.fail + +/** + * The decision-quality gate of the packed Q8_0 lane (SKEEP-007 Tier 2 consumer). + * + * FP32 parity (`LayaPredictFixtureTest`, 1.5e-4) is not the right contract for a quantized model; + * what has to hold is that it **decides the same way**: every choice the FP32 agent makes on the + * `predict.json` cases, the Q8_0 agent makes too, and no probability, score, noul value or + * confidence moves by more than [maxProbabilityShift]. Measured on the multilingual checkpoint on + * 2026-10-16: 5/5 choices agree, 45 values compared, max shift 0.035 — the gate sits at 0.05 so a + * regression of the lane (a wrong block order, a mis-split projection) fails loudly while the + * quantization noise of a 22-layer encoder passes. Needs `LAYA_MULTILINGUAL_PATH` with + * `model.safetensors`; skipped otherwise, like the other model-gated tests. + */ +class LayaQuantizedParityTest { + private val maxProbabilityShift = System.getenv("LAYA_Q8_TOL")?.toDoubleOrNull() ?: 0.05 + + @Test + fun `q8_0 lane decides like the fp32 lane on the predict fixtures`() { + val dir = LayaFixtures.checkpointDir() ?: return + if (!Files.exists(dir.resolve("model.safetensors"))) { println("[skip] $dir has no model.safetensors"); return } + + val t0 = System.nanoTime() + val fp32 = LayaModelLoader.load(dir, DirectCpuExecutionContext()) + val t1 = System.nanoTime() + val q8 = LayaModelLoader.load(dir, DirectCpuExecutionContext(), weightFormFor = LayaWeightForms.q8_0For) + val t2 = System.nanoTime() + println("[laya-q8] fp32 loaded in ${(t1 - t0) / 1_000_000} ms, q8_0 in ${(t2 - t1) / 1_000_000} ms") + + val packedBytes = q8.model.packedWeightBytes() + assertTrue(packedBytes > 0, "the q8_0 lane must hold packed weights") + println("[laya-q8] packed weight bytes: ${packedBytes / (1024 * 1024)} MB") + + val fixture = LayaFixtures.resource("predict.json") + val failures = mutableListOf() + var choices = 0 + var agreed = 0 + var compared = 0 + var maxShift = 0.0 + var fp32Ms = 0L + var q8Ms = 0L + for ((name, element) in fixture["single"]!!.jsonObject) { + val case = element.jsonObject + val state = LayaState.of(case["state"]!!) + val questions = case["questions"]!!.jsonObject.mapValues { (_, q) -> LayaFixtures.questionOf(q.jsonObject) } + val a = System.nanoTime() + val reference = fp32.decideAll(state, questions) + val b = System.nanoTime() + val quantized = q8.decideAll(state, questions) + val c = System.nanoTime() + fp32Ms += (b - a) / 1_000_000 + q8Ms += (c - b) / 1_000_000 + for ((qid, ref) in reference) { + val got = quantized.getValue(qid).answer + val exp = ref.answer + fun shift(what: String, g: Double, e: Double) { + compared++ + val d = abs(g - e) + if (d > maxShift) maxShift = d + if (d > maxProbabilityShift) failures += "$name/$qid $what: q8 ${"%.4f".format(g)} vs fp32 ${"%.4f".format(e)} (shift ${"%.3f".format(d)})" + } + when (exp) { + is LayaAnswer.Choice -> { + choices++ + val g = got as? LayaAnswer.Choice ?: run { failures += "$name/$qid: answer kind differs"; null } + if (g != null) { + if (g.choice == exp.choice) agreed++ else failures += "$name/$qid choice: q8 ${g.choice} vs fp32 ${exp.choice}" + for ((key, p) in exp.byKey) shift("p[$key]", g.byKey.getValue(key), p) + } + } + is LayaAnswer.Score -> { + val g = got as LayaAnswer.Score + shift("score", g.score, exp.score) + for (i in exp.probabilities.indices) shift("p[$i]", g.probabilities[i], exp.probabilities[i]) + } + is LayaAnswer.Noul -> shift("noul", (got as LayaAnswer.Noul).noul, exp.noul) + } + shift("confidence", got.confidence, exp.confidence) + } + } + println("[laya-q8] $compared values compared, max shift ${"%.4f".format(maxShift)}, choices $agreed/$choices agree; fp32 $fp32Ms ms, q8_0 $q8Ms ms") + if (choices > 0 && agreed < choices) failures += "choice agreement $agreed/$choices below the gate (every choice must agree)" + if (failures.isNotEmpty()) fail(failures.joinToString("\n")) + } +} + +/** Sum of the packed bytes held by a decision model's parameters, to prove the lane is in effect. */ +private fun LayaDecisionModel<*>.packedWeightBytes(): Long = + (encoder.trainableParameters() + head.trainableParameters()) + .sumOf { p -> (p.value.data as? PackedBlockStorage)?.packedData?.size?.toLong() ?: 0L }