From ab989c5544ce78d360df02bfea1251abefaa44a0 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Mon, 27 Jul 2026 09:08:33 +0200 Subject: [PATCH 1/6] feat(dtype): keep fp16/bf16 weights packed on the SafeTensors and GGUF paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Engine 0.38.0 adds the shared narrow-float codec, Fp16DenseTensorData, FP16 matmul kernels and codec-driven matmul dispatch. Take advantage of all of it. - DecoderSafeTensorsLoader gains the F16 KEEP_NATIVE arm that BF16 has had since 0.25.0 — it was missing only for want of a storage type to back it. Covers LLaMA, Qwen and Voxtral, which share this loader. - DecoderGgufWeightLoader accepts a DTypePolicy and keeps F16/BF16 sources in their on-disk layout. The GGUF branches of the network loaders used to construct it without the policy, so withDtypePolicy was ignored on that path entirely; it is now plumbed through. - DTypePolicyValidation takes keepNative: Set in place of a BF16-only boolean, which could express neither "keeps FP16 but not BF16" nor the empty case. The old overload stays, deprecated. The two narrow formats are resolved independently: Require(BF16) still widens F16 sources and vice versa. Both are 2 bytes per element, so mis-tagging F16 bytes as BF16 would decode to plausible-looking garbage rather than fail. The GGUF packed path swaps the rank-2 shape to [cols, rows] and moves no bytes, mirroring createTensor exactly. GGUF header dims are reversed relative to the logical row-major shape, which is why transposeColumnMajorToRowMajor returns its input untouched; an actual element transpose here would have handed the matmul kernel a silently transposed weight matrix. Gemma and Apertus now reject Require(BF16) instead of accepting a policy their own weight chains never honored and silently disregarding it at load. Callers need Prefer(BF16) until those chains grow a KEEP_NATIVE path. DecoderGgufWeightLoader's constructors gain a trailing defaulted parameter, which changes their JVM descriptors: binary-breaking, source-compatible. Tests: 22 new — 6 on the validation capability model, 5 end-to-end on the SafeTensors path against synthesized files, 11 on the GGUF policy decision and tensor construction (no GGUF writer exists to test that path end to end). 285 green across the touched modules; apiDump regenerated. --- CHANGELOG.md | 60 +++++ README.md | 6 + llm-core/api/jvm/llm-core.api | 3 + .../ainet/apps/llm/DTypePolicyValidation.kt | 111 ++++++--- .../apps/llm/DTypePolicyValidationTest.kt | 95 +++++++ .../models/apertus/ApertusNetworkLoader.kt | 5 +- .../ainet/models/gemma/GemmaNetworkLoader.kt | 7 +- llm-inference/llama/api/jvm/llama.api | 12 +- .../models/llama/DecoderGgufWeightLoader.kt | 121 +++++++-- .../models/llama/DecoderNarrowFloatSupport.kt | 24 ++ .../models/llama/DecoderSafeTensorsLoader.kt | 58 +++-- .../ainet/models/llama/LlamaNetworkLoader.kt | 23 +- .../llama/DecoderGgufNarrowFloatTest.kt | 231 ++++++++++++++++++ ...DecoderSafeTensorsLoaderNarrowFloatTest.kt | 211 ++++++++++++++++ .../LlamaNetworkLoaderDTypePolicyTest.kt | 43 ++-- .../sk/ainet/models/qwen/QwenNetworkLoader.kt | 12 +- .../models/voxtral/VoxtralNetworkLoader.kt | 14 +- 17 files changed, 920 insertions(+), 116 deletions(-) create mode 100644 llm-core/src/commonTest/kotlin/sk/ainet/apps/llm/DTypePolicyValidationTest.kt create mode 100644 llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/DecoderNarrowFloatSupport.kt create mode 100644 llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/DecoderGgufNarrowFloatTest.kt create mode 100644 llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/DecoderSafeTensorsLoaderNarrowFloatTest.kt diff --git a/CHANGELOG.md b/CHANGELOG.md index 3c8c6f5..87961b0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,66 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +Requires **SKaiNET engine 0.38.0** (narrow-float codec, `Fp16DenseTensorData`, FP16 matmul +kernels, codec-driven dispatch — engine PR #886). + +### Added + +- **FP16 KEEP_NATIVE on the SafeTensors path.** `DecoderSafeTensorsLoader` gains the F16 arm + that BF16 has had since 0.25.0: with a `DTypePolicy` admitting FP16 (`Require(FP16)`, + `Prefer(FP16)`, or `OneOf` containing FP16) it stops widening F16 tensors and wraps the + on-disk 2-bytes-per-element buffer in `Fp16DenseTensorData`. The arm was missing only + because no such storage type existed. `DefaultCpuOpsJvm` matches `NarrowFloatTensorData` + and picks the kernel by codec, so an F16 checkpoint now stays near its on-disk footprint + instead of inflating ~2× as FP32. Covers LLaMA, Qwen, and Voxtral, which share this loader. +- **Narrow-float KEEP_NATIVE on the GGUF path — `DTypePolicy` is honored there at all now.** + `DecoderGgufWeightLoader` accepts a `dtypePolicy` and keeps F16 / BF16 source tensors packed + instead of widening every one to FP32. `LlamaNetworkLoader`, `QwenNetworkLoader`, and + `VoxtralNetworkLoader` plumb the policy attached via `withDtypePolicy` down into it; before + this the GGUF branches constructed the loader without the policy and silently ignored it. + This is the KEEP_NATIVE GGUF path the 0.25.0 notes parked, and it is what makes + `Require(BF16)` real on GGUF. + + The packed path mirrors the FP32 path's layout handling exactly: for rank 2 it swaps the + shape to `[cols, rows]` and **moves no bytes**. GGUF header dims are reversed relative to the + logical row-major shape, so the "column-major → row-major" step is a reinterpretation, not a + permutation (`DequantOps.transposeColumnMajorToRowMajor` returns its input unchanged). An + actual element transpose here would have handed the matmul kernel a silently transposed + weight matrix. The result is genuinely zero-copy — the on-disk buffer becomes the storage. + +### Changed + +- **`DTypePolicyValidation` capability model is per-format.** `validate(policy, loaderName, + keepNative: Set)` replaces the BF16-only `allowBf16Require: Boolean` (kept as a + `@Deprecated` overload). A caller declares which narrow-float formats its chain actually + hands through packed, and a `Require` naming one is accepted only by a chain that can honor + it. The boolean could express neither "keeps FP16 but not BF16" nor the empty case. + + The two formats are tracked separately and never interchangeably: `Require(BF16)` still + widens F16 sources, and vice versa. Both are 2 bytes per element, so mis-tagging F16 bytes as + BF16 decodes to plausible-looking garbage rather than throwing. `DTypePolicyValidation + .keepsNative(policy, native)` is the single decision point both loader chains share, mirroring + the engine's `mapPolicyToNarrow` / `keepsNative`. +- **`Require(FP16)` is now accepted** by `LlamaNetworkLoader`, `QwenNetworkLoader`, and + `VoxtralNetworkLoader` (both GGUF and SafeTensors), and **`Require(BF16)` is now accepted on + their GGUF paths**. Both previously threw. +- **Binary-breaking (source-compatible): `DecoderGgufWeightLoader` constructors** gain a + trailing `dtypePolicy: DTypePolicy = DTypePolicy.Any`, which changes their JVM descriptors. + Kotlin and Java callers compile unchanged; already-compiled callers must be rebuilt. + Behaviour with the default is identical to before. + +### Fixed + +- **`GemmaNetworkLoader` and `ApertusNetworkLoader` no longer accept a `Require(BF16)` they + ignore.** Both have their own weight chains (`Gemma4WeightLoader` / + `Gemma4SafeTensorsWeightLoader`, `ApertusWeightLoader` / `ApertusSingleSafeTensorsLoader`) + which widen every narrow float to FP32 and have no KEEP_NATIVE path. Their SafeTensors + entrypoints nevertheless passed `allowBf16Require = true`, so `Require(BF16)` validated and + was then silently disregarded at load — the exact failure the eager validator exists to + prevent. They now declare `keepNative = emptySet()` and reject it. **Callers relying on the + old acceptance must switch to `Prefer(BF16)`** (a soft constraint, which still passes) until + those chains grow a KEEP_NATIVE path. + ## [0.36.1] — 2026-07-17 Patch on **0.36.0** (same SKaiNET engine 0.36.0). Two additions: **BGE embedding models** on the diff --git a/README.md b/README.md index f99e09a..8d2cb3b 100644 --- a/README.md +++ b/README.md @@ -368,6 +368,12 @@ See `llm-test/llm-test-java/src/test/java/.../KLlamaJavaToolCallingTest.java` fo ## What's new in 0.25.0 +> **Superseded (unreleased, engine 0.38.0).** The narrow-float limits described below are gone: +> the GGUF chain now honors `DTypePolicy` and keeps F16/BF16 packed, `Require(FP16)` is accepted +> alongside `Require(BF16)`, and the two formats are resolved independently. Conversely, Gemma +> and Apertus now *reject* `Require(BF16)` — their own weight chains never honored it. See the +> `[Unreleased]` section of [CHANGELOG.md](CHANGELOG.md). + - **`DTypePolicy` on every `*NetworkLoader.fromGguf` / `.fromSafeTensors` entry.** A sealed `DTypePolicy` type (`Any | Require | Prefer | OneOf`, upstream of SKaiNET 0.25.0) is now accepted on every loader companion in diff --git a/llm-core/api/jvm/llm-core.api b/llm-core/api/jvm/llm-core.api index 2275bca..0a22782 100644 --- a/llm-core/api/jvm/llm-core.api +++ b/llm-core/api/jvm/llm-core.api @@ -1,6 +1,9 @@ public final class sk/ainet/apps/llm/DTypePolicyValidation { public static final field INSTANCE Lsk/ainet/apps/llm/DTypePolicyValidation; + public final fun keepsNative (Lsk/ainet/lang/types/DTypePolicy;Lsk/ainet/lang/types/DType;)Z + public final fun validate (Lsk/ainet/lang/types/DTypePolicy;Ljava/lang/String;Ljava/util/Set;)V public final fun validate (Lsk/ainet/lang/types/DTypePolicy;Ljava/lang/String;Z)V + public static synthetic fun validate$default (Lsk/ainet/apps/llm/DTypePolicyValidation;Lsk/ainet/lang/types/DTypePolicy;Ljava/lang/String;Ljava/util/Set;ILjava/lang/Object;)V } public abstract class sk/ainet/apps/llm/DecoderRuntime : sk/ainet/apps/llm/InferenceRuntime { diff --git a/llm-core/src/commonMain/kotlin/sk/ainet/apps/llm/DTypePolicyValidation.kt b/llm-core/src/commonMain/kotlin/sk/ainet/apps/llm/DTypePolicyValidation.kt index fc8cffd..db19c1a 100644 --- a/llm-core/src/commonMain/kotlin/sk/ainet/apps/llm/DTypePolicyValidation.kt +++ b/llm-core/src/commonMain/kotlin/sk/ainet/apps/llm/DTypePolicyValidation.kt @@ -13,69 +13,114 @@ import sk.ainet.lang.types.FP32 * generalised execution-side dtype constraint surface. Its own loaders * (`StreamingGgufParametersLoader.withPolicy`, `SafeTensorsParametersLoader.withPolicy`) * validate the policy at construction so callers fail fast on impossible - * requirements. + * requirements; this validator is the same boundary for the transformer-repo + * loader chains (`DecoderGgufWeightLoader`, `DecoderSafeTensorsLoader`, …). * - * The transformer-repo loaders (`LlamaNetworkLoader`, `QwenNetworkLoader`, …) ship - * their own weight-loading chain on top of `DecoderGgufWeightLoader` / - * `DecoderSafeTensorsLoader`. Those chains do not yet plumb `DTypePolicy` through - * to the underlying tensor producers — that's a separate follow-up. In the - * meantime, accepting the policy on the public surface lets consumers express - * intent today, and this validator ensures we reject impossible requirements at - * the same boundary SKaiNET's own loaders do. + * ## What a chain can promise * - * Today the transformer-repo loaders only produce FP32 (after Q4/Q8/BF16/F16 - * dequant on the SafeTensors path; native quantization preservation on the GGUF - * path). That matches the SKaiNET 0.25.0 `StreamingGgufParametersLoader` - * validator. The BF16 KEEP_NATIVE SafeTensors path (`Require(BF16)`) is allowed - * here even though the transformer-repo `DecoderSafeTensorsLoader` does not yet - * honor it — when wired through, no API change is needed. + * Every chain produces FP32, so `Require(FP32)` always passes. A `Require` naming a + * **narrow float** (BF16 / FP16) is a promise that the weights reach the kernel in their + * on-disk 2-bytes-per-element layout — the KEEP_NATIVE path. Only a chain that actually + * implements KEEP_NATIVE for that format can honour it, so each caller declares its + * capability via [keepNative] rather than the validator guessing from the format alone. * - * Throws [IllegalArgumentException] on `Require(target)` for targets we cannot - * produce. `Any`, `Prefer`, and `OneOf` always pass. + * The two narrow formats are tracked **separately and never interchangeably**. They are + * different bit layouts at the same width, so mis-tagging F16 bytes as BF16 decodes to + * plausible-looking garbage instead of throwing. A chain that keeps BF16 native but + * widens F16 declares exactly that, and `Require(FP16)` against it fails loudly. + * + * Note that a `Require(BF16)` chain still **widens F16 sources to FP32** — the policy + * names the format to preserve, not a conversion target. Neither narrow format can be + * re-encoded into the other without a lossy round-trip, and the loaders do not try. + * + * `Any`, `Prefer`, and `OneOf` always pass: they are soft constraints that a chain is + * free to satisfy or ignore per tensor. + * + * Throws [IllegalArgumentException] on `Require(target)` for targets the caller cannot + * produce. */ public object DTypePolicyValidation { /** - * Validates a [DTypePolicy] for the transformer-repo loader chain. + * Validates a [DTypePolicy] for one transformer-repo loader chain. * * @param policy the policy supplied by the caller * @param loaderName loader name for error messages (e.g. `"LlamaNetworkLoader.fromGguf"`) - * @param allowBf16Require whether `Require(BF16)` is acceptable. SafeTensors-backed - * loaders set this to `true` (matches SKaiNET's `SafeTensorsParametersLoader`); GGUF-only - * loaders set it to `false` (matches SKaiNET's `StreamingGgufParametersLoader`). + * @param keepNative the narrow-float dtypes this chain hands through in their on-disk + * packed layout. Empty (the default) means the chain widens every narrow float to + * FP32, so any `Require` naming one is rejected. */ public fun validate( policy: DTypePolicy, loaderName: String, - allowBf16Require: Boolean, + keepNative: Set = emptySet(), ) { when (policy) { DTypePolicy.Any -> Unit is DTypePolicy.Prefer -> Unit is DTypePolicy.OneOf -> Unit - is DTypePolicy.Require -> validateRequire(policy.target, loaderName, allowBf16Require) + is DTypePolicy.Require -> validateRequire(policy.target, loaderName, keepNative) } } - private fun validateRequire(target: DType, loaderName: String, allowBf16Require: Boolean) { + /** + * BF16-only capability flag. + * + * @param allowBf16Require whether `Require(BF16)` is acceptable. + */ + @Deprecated( + "Narrow-float capability is per-format since engine 0.38.0 — a chain can keep FP16 " + + "native too. Pass the set of formats it keeps packed.", + ReplaceWith( + "validate(policy, loaderName, if (allowBf16Require) setOf(BF16) else emptySet())", + "sk.ainet.lang.types.BF16", + ), + ) + public fun validate( + policy: DTypePolicy, + loaderName: String, + allowBf16Require: Boolean, + ): Unit = validate(policy, loaderName, if (allowBf16Require) setOf(BF16) else emptySet()) + + /** + * Whether [policy] asks for [native] source tensors to stay in their on-disk 16-bit layout. + * + * The single decision point every transformer-repo loader chain shares, mirroring the + * engine's `SafeTensorsParametersLoader.mapPolicyToNarrow` / + * `StreamingGgufParametersLoader.keepsNative`. Only the format the policy actually names + * is kept: `Require(BF16)` still widens F16 sources, because turning one narrow format + * into the other needs a lossy re-encode. + * + * [native] is expected to be [BF16] or [FP16]; any other dtype answers `false`. + */ + public fun keepsNative(policy: DTypePolicy, native: DType): Boolean = when (policy) { + DTypePolicy.Any -> false + is DTypePolicy.Require -> policy.target == native + is DTypePolicy.Prefer -> policy.target == native + is DTypePolicy.OneOf -> native in policy.allowed + } + + private fun validateRequire(target: DType, loaderName: String, keepNative: Set) { when (target) { FP32 -> Unit - BF16 -> if (!allowBf16Require) { + BF16, FP16 -> if (target !in keepNative) { throw IllegalArgumentException( - "$loaderName: Require(BF16) is not supported by the GGUF loader chain — " + - "GGUF BF16 sources are dequanted to FP32 today (no KEEP_NATIVE GGUF path " + - "yet). Use Any or Prefer(BF16) to accept the dequant fallback." + "$loaderName: Require(${target.name}) is not supported by this loader chain — " + + "${target.name} sources are widened to FP32 at load " + + describeKeepNative(keepNative) + ". " + + "Use Any or Prefer(${target.name}) to accept the widening fallback." ) } - FP16 -> throw IllegalArgumentException( - "$loaderName: Require(FP16) is not supported — the loader chain dequants F16 to " + - "FP32 (no Fp16DenseTensorData backing yet). Use Any or Prefer(FP16)." - ) else -> throw IllegalArgumentException( "$loaderName: Require(${target.name}) is not satisfiable — the transformer-repo " + - "loader chain produces FP32 (optionally BF16 on the SafeTensors KEEP_NATIVE " + - "path). It cannot fabricate ${target.name} from arbitrary sources." + "loader chain produces FP32 " + describeKeepNative(keepNative) + ". " + + "It cannot fabricate ${target.name} from arbitrary sources." ) } } + + private fun describeKeepNative(keepNative: Set): String = when { + keepNative.isEmpty() -> "(this chain keeps no narrow float packed)" + else -> "(this chain keeps only ${keepNative.joinToString(" / ") { it.name }} packed)" + } } diff --git a/llm-core/src/commonTest/kotlin/sk/ainet/apps/llm/DTypePolicyValidationTest.kt b/llm-core/src/commonTest/kotlin/sk/ainet/apps/llm/DTypePolicyValidationTest.kt new file mode 100644 index 0000000..6fc700c --- /dev/null +++ b/llm-core/src/commonTest/kotlin/sk/ainet/apps/llm/DTypePolicyValidationTest.kt @@ -0,0 +1,95 @@ +package sk.ainet.apps.llm + +import sk.ainet.lang.types.BF16 +import sk.ainet.lang.types.DTypePolicy +import sk.ainet.lang.types.FP16 +import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.Int8 +import kotlin.test.Test +import kotlin.test.assertFailsWith +import kotlin.test.assertFalse +import kotlin.test.assertTrue + +/** + * Pins the capability model behind [DTypePolicyValidation]. + * + * A `Require` naming a narrow float is a promise that the weights reach the kernel packed, so it + * may only be accepted by a chain that actually implements KEEP_NATIVE for *that* format. The + * previous BF16-only boolean could not express "keeps FP16 but not BF16", nor the empty case + * (Gemma / Apertus) — which it papered over by accepting `Require(BF16)` and ignoring it. + */ +class DTypePolicyValidationTest { + + private val bothNarrow = setOf(BF16, FP16) + + @Test + fun `Require(FP32) is always accepted — every chain produces FP32`() { + DTypePolicyValidation.validate(DTypePolicy.Require(FP32), "test", keepNative = emptySet()) + DTypePolicyValidation.validate(DTypePolicy.Require(FP32), "test", keepNative = bothNarrow) + } + + @Test + fun `soft policies never raise, whatever they name`() { + for (keepNative in listOf(emptySet(), bothNarrow)) { + DTypePolicyValidation.validate(DTypePolicy.Any, "test", keepNative) + DTypePolicyValidation.validate(DTypePolicy.Prefer(BF16), "test", keepNative) + DTypePolicyValidation.validate(DTypePolicy.Prefer(FP16), "test", keepNative) + DTypePolicyValidation.validate(DTypePolicy.Prefer(Int8), "test", keepNative) + DTypePolicyValidation.validate(DTypePolicy.OneOf(setOf(FP32, BF16)), "test", keepNative) + } + } + + @Test + fun `a chain that keeps nothing packed rejects both narrow Requires`() { + // The Gemma / Apertus position: their weight chains widen every narrow float, so a + // Require they cannot honor must fail loudly rather than be silently ignored. + assertFailsWith { + DTypePolicyValidation.validate(DTypePolicy.Require(BF16), "test", keepNative = emptySet()) + } + assertFailsWith { + DTypePolicyValidation.validate(DTypePolicy.Require(FP16), "test", keepNative = emptySet()) + } + } + + @Test + fun `the two narrow formats are tracked independently`() { + // Keeps BF16 only. + DTypePolicyValidation.validate(DTypePolicy.Require(BF16), "test", keepNative = setOf(BF16)) + assertFailsWith { + DTypePolicyValidation.validate(DTypePolicy.Require(FP16), "test", keepNative = setOf(BF16)) + } + + // Keeps FP16 only — the mirror image, which the old boolean flag could not express. + DTypePolicyValidation.validate(DTypePolicy.Require(FP16), "test", keepNative = setOf(FP16)) + assertFailsWith { + DTypePolicyValidation.validate(DTypePolicy.Require(BF16), "test", keepNative = setOf(FP16)) + } + } + + @Test + fun `a dtype no chain produces is rejected even when both narrow floats are kept`() { + assertFailsWith { + DTypePolicyValidation.validate(DTypePolicy.Require(Int8), "test", keepNative = bothNarrow) + } + } + + @Test + fun `keepsNative names one format at a time`() { + assertFalse(DTypePolicyValidation.keepsNative(DTypePolicy.Any, BF16)) + assertFalse(DTypePolicyValidation.keepsNative(DTypePolicy.Any, FP16)) + + // Require / Prefer keep exactly the format they name — never the sibling, which would + // mean reinterpreting one 16-bit layout as the other and decoding to silent garbage. + assertTrue(DTypePolicyValidation.keepsNative(DTypePolicy.Require(BF16), BF16)) + assertFalse(DTypePolicyValidation.keepsNative(DTypePolicy.Require(BF16), FP16)) + assertTrue(DTypePolicyValidation.keepsNative(DTypePolicy.Prefer(FP16), FP16)) + assertFalse(DTypePolicyValidation.keepsNative(DTypePolicy.Prefer(FP16), BF16)) + + // OneOf may admit both at once. + val oneOf = DTypePolicy.OneOf(setOf(BF16, FP16)) + assertTrue(DTypePolicyValidation.keepsNative(oneOf, BF16)) + assertTrue(DTypePolicyValidation.keepsNative(oneOf, FP16)) + + assertFalse(DTypePolicyValidation.keepsNative(DTypePolicy.Require(FP32), BF16)) + } +} diff --git a/llm-inference/apertus/src/commonMain/kotlin/sk/ainet/models/apertus/ApertusNetworkLoader.kt b/llm-inference/apertus/src/commonMain/kotlin/sk/ainet/models/apertus/ApertusNetworkLoader.kt index 928e1ef..7b5ec00 100644 --- a/llm-inference/apertus/src/commonMain/kotlin/sk/ainet/models/apertus/ApertusNetworkLoader.kt +++ b/llm-inference/apertus/src/commonMain/kotlin/sk/ainet/models/apertus/ApertusNetworkLoader.kt @@ -40,8 +40,9 @@ public class ApertusNetworkLoader @PublishedApi internal constructor( /** See [sk.ainet.models.llama.LlamaNetworkLoader.withDtypePolicy]. */ public fun withDtypePolicy(policy: DTypePolicy): ApertusNetworkLoader { - val allowBf16 = weightsProvider is WeightsProvider.SafeTensorsSingle - DTypePolicyValidation.validate(policy, "ApertusNetworkLoader.withDtypePolicy", allowBf16Require = allowBf16) + // As Gemma: `ApertusWeightLoader` / `ApertusSingleSafeTensorsLoader` widen every narrow + // float to FP32, so this loader keeps nothing packed and promises nothing. + DTypePolicyValidation.validate(policy, "ApertusNetworkLoader.withDtypePolicy", keepNative = emptySet()) this.dtypePolicy = policy return this } diff --git a/llm-inference/gemma/src/commonMain/kotlin/sk/ainet/models/gemma/GemmaNetworkLoader.kt b/llm-inference/gemma/src/commonMain/kotlin/sk/ainet/models/gemma/GemmaNetworkLoader.kt index e9d6772..65f0a07 100644 --- a/llm-inference/gemma/src/commonMain/kotlin/sk/ainet/models/gemma/GemmaNetworkLoader.kt +++ b/llm-inference/gemma/src/commonMain/kotlin/sk/ainet/models/gemma/GemmaNetworkLoader.kt @@ -41,8 +41,11 @@ public class GemmaNetworkLoader @PublishedApi internal constructor( /** See [sk.ainet.models.llama.LlamaNetworkLoader.withDtypePolicy]. */ public fun withDtypePolicy(policy: DTypePolicy): GemmaNetworkLoader { - val allowBf16 = weightsProvider is WeightsProvider.SafeTensorsIndex - DTypePolicyValidation.validate(policy, "GemmaNetworkLoader.withDtypePolicy", allowBf16Require = allowBf16) + // Gemma has its own weight chain (`Gemma4WeightLoader` / `Gemma4SafeTensorsWeightLoader`), + // which widens every narrow float to FP32 — it has no KEEP_NATIVE path yet, unlike the + // shared decoder chain LLaMA/Qwen/Voxtral use. So it promises nothing and `Require(BF16)` + // is rejected rather than accepted-and-ignored. + DTypePolicyValidation.validate(policy, "GemmaNetworkLoader.withDtypePolicy", keepNative = emptySet()) this.dtypePolicy = policy return this } diff --git a/llm-inference/llama/api/jvm/llama.api b/llm-inference/llama/api/jvm/llama.api index 314193e..fe7e791 100644 --- a/llm-inference/llama/api/jvm/llama.api +++ b/llm-inference/llama/api/jvm/llama.api @@ -19,10 +19,10 @@ public final class sk/ainet/models/llama/DecoderGgufMemSegConverterKt { public final class sk/ainet/models/llama/DecoderGgufWeightLoader { public static final field Dequant Lsk/ainet/models/llama/DecoderGgufWeightLoader$Dequant; - public fun (Lkotlin/jvm/functions/Function0;Lsk/ainet/io/model/QuantPolicy;Ljava/util/Set;)V - public synthetic fun (Lkotlin/jvm/functions/Function0;Lsk/ainet/io/model/QuantPolicy;Ljava/util/Set;ILkotlin/jvm/internal/DefaultConstructorMarker;)V - public fun (Lkotlin/jvm/functions/Function0;ZLsk/ainet/io/model/QuantPolicy;Ljava/util/Set;)V - public synthetic fun (Lkotlin/jvm/functions/Function0;ZLsk/ainet/io/model/QuantPolicy;Ljava/util/Set;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun (Lkotlin/jvm/functions/Function0;Lsk/ainet/io/model/QuantPolicy;Ljava/util/Set;Lsk/ainet/lang/types/DTypePolicy;)V + public synthetic fun (Lkotlin/jvm/functions/Function0;Lsk/ainet/io/model/QuantPolicy;Ljava/util/Set;Lsk/ainet/lang/types/DTypePolicy;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun (Lkotlin/jvm/functions/Function0;ZLsk/ainet/io/model/QuantPolicy;Ljava/util/Set;Lsk/ainet/lang/types/DTypePolicy;)V + public synthetic fun (Lkotlin/jvm/functions/Function0;ZLsk/ainet/io/model/QuantPolicy;Ljava/util/Set;Lsk/ainet/lang/types/DTypePolicy;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public final fun load (Lsk/ainet/context/ExecutionContext;Lkotlin/reflect/KClass;Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; public final fun loadStreaming (Lsk/ainet/context/ExecutionContext;Lkotlin/reflect/KClass;Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; public final fun loadToMap (Lsk/ainet/context/ExecutionContext;Lkotlin/reflect/KClass;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; @@ -48,6 +48,10 @@ public final class sk/ainet/models/llama/DecoderGgufWeights { public fun toString ()Ljava/lang/String; } +public final class sk/ainet/models/llama/DecoderNarrowFloatSupportKt { + public static final fun getDECODER_NARROW_KEEP_NATIVE ()Ljava/util/Set; +} + public final class sk/ainet/models/llama/DecoderSafeTensorsLoader { public fun (Lsk/ainet/context/ExecutionContext;Lkotlin/reflect/KClass;Lsk/ainet/models/llama/LlamaModelMetadata;ZLsk/ainet/lang/types/DTypePolicy;)V public synthetic fun (Lsk/ainet/context/ExecutionContext;Lkotlin/reflect/KClass;Lsk/ainet/models/llama/LlamaModelMetadata;ZLsk/ainet/lang/types/DTypePolicy;ILkotlin/jvm/internal/DefaultConstructorMarker;)V diff --git a/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/DecoderGgufWeightLoader.kt b/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/DecoderGgufWeightLoader.kt index fa71b24..ae961f4 100644 --- a/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/DecoderGgufWeightLoader.kt +++ b/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/DecoderGgufWeightLoader.kt @@ -2,6 +2,7 @@ package sk.ainet.models.llama import kotlinx.io.Source import kotlinx.io.buffered +import sk.ainet.apps.llm.DTypePolicyValidation import sk.ainet.context.ExecutionContext import sk.ainet.io.RandomAccessSource import sk.ainet.io.gguf.GGMLQuantizationType @@ -15,7 +16,13 @@ import sk.ainet.io.gguf.dequant.DequantOps import sk.ainet.io.model.QuantPolicy import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.Bf16DenseTensorData +import sk.ainet.lang.tensor.data.Fp16DenseTensorData +import sk.ainet.lang.tensor.data.NarrowFloatTensorData +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.types.BF16 import sk.ainet.lang.types.DType +import sk.ainet.lang.types.DTypePolicy import sk.ainet.lang.types.FP16 import sk.ainet.lang.types.FP32 import sk.ainet.lang.types.Int8 @@ -76,8 +83,20 @@ public class DecoderGgufWeightLoader private constructor( private val randomAccessProvider: (() -> RandomAccessSource)?, private val loadTensorData: Boolean = true, private val quantPolicy: QuantPolicy = QuantPolicy.RAW_BYTES, - private val acceptedArchitectures: Set = setOf("llama") + private val acceptedArchitectures: Set = setOf("llama"), + private val dtypePolicy: DTypePolicy = DTypePolicy.Any, ) { + /** + * Keep `F16` source tensors in their on-disk 2-bytes-per-element layout instead of widening + * them to FP32. Resolved from [dtypePolicy] exactly as the engine's + * `StreamingGgufParametersLoader.keepsNative` does, so a policy carried down from + * `LlamaNetworkLoader.withDtypePolicy` means the same thing on both sides. + */ + private val keepF16Native: Boolean = DTypePolicyValidation.keepsNative(dtypePolicy, FP16) + + /** As [keepF16Native], for `BF16` sources. Resolved independently — see [keepsNarrowNative]. */ + private val keepBf16Native: Boolean = DTypePolicyValidation.keepsNative(dtypePolicy, BF16) + /** * Primary constructor for sequential Source-based loading. * Loads entire file into memory - suitable for models under 2GB. @@ -85,18 +104,22 @@ public class DecoderGgufWeightLoader private constructor( * @param acceptedArchitectures GGUF architecture strings accepted by this loader. * Defaults to `setOf("llama")`. Consumers loading compatible architectures * (e.g. Qwen, Mistral) pass their own set — no changes needed here. + * @param dtypePolicy narrow-float handling. Default [DTypePolicy.Any] widens F16/BF16 + * sources to FP32; a policy naming BF16 or FP16 keeps that format packed. */ public constructor( sourceProvider: () -> Source, loadTensorData: Boolean = true, quantPolicy: QuantPolicy = QuantPolicy.RAW_BYTES, - acceptedArchitectures: Set = setOf("llama") + acceptedArchitectures: Set = setOf("llama"), + dtypePolicy: DTypePolicy = DTypePolicy.Any, ) : this( sourceProvider = sourceProvider, randomAccessProvider = null, loadTensorData = loadTensorData, quantPolicy = quantPolicy, - acceptedArchitectures = acceptedArchitectures + acceptedArchitectures = acceptedArchitectures, + dtypePolicy = dtypePolicy, ) /** @@ -111,13 +134,15 @@ public class DecoderGgufWeightLoader private constructor( public constructor( randomAccessProvider: () -> RandomAccessSource, quantPolicy: QuantPolicy = QuantPolicy.RAW_BYTES, - acceptedArchitectures: Set = setOf("llama") + acceptedArchitectures: Set = setOf("llama"), + dtypePolicy: DTypePolicy = DTypePolicy.Any, ) : this( sourceProvider = null, randomAccessProvider = randomAccessProvider, loadTensorData = true, // Ignored for streaming quantPolicy = quantPolicy, - acceptedArchitectures = acceptedArchitectures + acceptedArchitectures = acceptedArchitectures, + dtypePolicy = dtypePolicy, ) /** @@ -574,12 +599,16 @@ public class DecoderGgufWeightLoader private constructor( require(dtype == FP32::class || dtype == FP16::class) { "Dequantizing ${st.tensorType} requires dtype FP32 or FP16; got ${dtype.simpleName}" } - val floats = when (st.tensorType) { - GGMLQuantizationType.F16 -> dequantF16FromBytes(bytes) - GGMLQuantizationType.BF16 -> dequantBF16FromBytes(bytes) - else -> error("Unreachable") + if (keepsNarrowNative(st.tensorType, dtype)) { + createNarrowTensor(ctx, dtype, shape, bytes, st.tensorType) + } else { + val floats = when (st.tensorType) { + GGMLQuantizationType.F16 -> dequantF16FromBytes(bytes) + GGMLQuantizationType.BF16 -> dequantBF16FromBytes(bytes) + else -> error("Unreachable") + } + createTensor(ctx, dtype, shape, floats) } - createTensor(ctx, dtype, shape, floats) } } } @@ -904,6 +933,62 @@ public class DecoderGgufWeightLoader private constructor( } } + /** + * Whether a GGUF tensor of [tensorType] should keep its on-disk 16-bit bytes. + * + * KEEP_NATIVE is restricted to `dtype == FP32` — that is the declared dtype the packed + * tensor presents to consumers (`get` decodes to `Float`), matching the SafeTensors path. + * An explicit `FP16::class` request is a storage-format ask for the FP32-array path and is + * left on the widening route rather than silently reinterpreted. + */ + internal fun keepsNarrowNative(tensorType: GGMLQuantizationType, dtype: KClass): Boolean = + dtype == FP32::class && when (tensorType) { + GGMLQuantizationType.F16 -> keepF16Native + GGMLQuantizationType.BF16 -> keepBf16Native + else -> false + } + + /** + * Wrap packed 16-bit GGUF bytes as a narrow-float tensor — the KEEP_NATIVE counterpart of + * [createTensor], and it must mirror that function's layout handling exactly. + * + * For rank 2 that means swapping the shape to `[cols, rows]` and **moving no bytes**. + * GGUF's header dims are reversed relative to the logical row-major shape, so the + * column-major → row-major step is a reinterpretation, not a permutation — which is why + * `DequantOps.transposeColumnMajorToRowMajor` returns its input untouched and + * [createTensor] only rebuilds the `Shape`. Doing an actual element transpose here would + * hand the matmul kernel a silently transposed weight matrix. + * + * The result is genuinely zero-copy: the on-disk buffer becomes the tensor's storage. + */ + @Suppress("UNCHECKED_CAST") + internal fun createNarrowTensor( + ctx: ExecutionContext, + dtype: KClass, + originalShape: Shape, + bytes: ByteArray, + tensorType: GGMLQuantizationType, + ): Tensor { + val shape = if (originalShape.rank == 2) { + Shape(originalShape[1], originalShape[0]) + } else { + originalShape + } + + val required = shape.volume * NarrowFloatTensorData.BYTES_PER_ELEMENT + require(bytes.size >= required) { + "Narrow-float buffer of ${bytes.size} bytes is short of the $required bytes needed " + + "for a ${shape.dimensions.toList()} $tensorType tensor" + } + + val data = when (tensorType) { + GGMLQuantizationType.F16 -> Fp16DenseTensorData.fromRawBytes(shape, bytes) + GGMLQuantizationType.BF16 -> Bf16DenseTensorData.fromRawBytes(shape, bytes) + else -> error("createNarrowTensor called with non-narrow type $tensorType") + } + return ctx.fromData(data as TensorData, dtype) as Tensor + } + private fun readerTensorToTensor( ctx: ExecutionContext, dtype: KClass, @@ -937,12 +1022,18 @@ public class DecoderGgufWeightLoader private constructor( "Dequantizing ${rt.tensorType} requires dtype FP32 or FP16; got ${dtype.simpleName}" } val raw = if (rt.data.isEmpty()) reader.materialize(rt) else rt.data - val floats = when (rt.tensorType) { - GGMLQuantizationType.F16 -> dequantF16(raw) - GGMLQuantizationType.BF16 -> dequantBF16(raw) - else -> error("Unsupported native type ${rt.tensorType}") + if (keepsNarrowNative(rt.tensorType, dtype)) { + createNarrowTensor( + ctx, dtype, shape, DequantOps.toByteArray(raw, rt.name), rt.tensorType, + ) + } else { + val floats = when (rt.tensorType) { + GGMLQuantizationType.F16 -> dequantF16(raw) + GGMLQuantizationType.BF16 -> dequantBF16(raw) + else -> error("Unsupported native type ${rt.tensorType}") + } + createTensor(ctx, dtype, shape, floats) } - createTensor(ctx, dtype, shape, floats) } } } diff --git a/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/DecoderNarrowFloatSupport.kt b/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/DecoderNarrowFloatSupport.kt new file mode 100644 index 0000000..5a883ee --- /dev/null +++ b/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/DecoderNarrowFloatSupport.kt @@ -0,0 +1,24 @@ +package sk.ainet.models.llama + +import sk.ainet.lang.types.BF16 +import sk.ainet.lang.types.DType +import sk.ainet.lang.types.FP16 + +/** + * The narrow-float formats the shared decoder chain hands through in their on-disk + * 2-bytes-per-element layout, rather than widening to FP32 at load. + * + * Both [DecoderGgufWeightLoader] and [DecoderSafeTensorsLoader] implement KEEP_NATIVE for BF16 + * and FP16 as of engine 0.38.0, so every `*NetworkLoader` built on them (LLaMA, Qwen, Voxtral) + * declares the same capability to `DTypePolicyValidation`. Loaders with their own weight chains + * — Gemma, Apertus — declare an empty set until those chains grow the same path. + * + * The set is a statement about *source* formats: a policy naming BF16 keeps BF16 tensors packed + * and still widens F16 ones, because neither format can be re-encoded as the other without a + * lossy round-trip. + * + * One caveat this coarse capability set cannot express: on the GGUF side, KEEP_NATIVE applies + * only under `QuantPolicy.DEQUANTIZE_TO_FP32` / `NATIVE_OPTIMIZED` with an `FP32` element type. + * `QuantPolicy.RAW_BYTES` hands every non-F32 tensor back as raw `Int8` bytes and is unaffected. + */ +public val DECODER_NARROW_KEEP_NATIVE: Set = setOf(BF16, FP16) diff --git a/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/DecoderSafeTensorsLoader.kt b/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/DecoderSafeTensorsLoader.kt index 17e4d39..4c62744 100644 --- a/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/DecoderSafeTensorsLoader.kt +++ b/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/DecoderSafeTensorsLoader.kt @@ -12,11 +12,14 @@ import sk.ainet.io.safetensors.StreamingSafeTensorsReader import sk.ainet.io.safetensors.StreamingSafeTensorInfo import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.Tensor +import sk.ainet.apps.llm.DTypePolicyValidation import sk.ainet.lang.tensor.data.Bf16DenseTensorData +import sk.ainet.lang.tensor.data.Fp16DenseTensorData import sk.ainet.lang.tensor.data.TensorData import sk.ainet.lang.types.BF16 import sk.ainet.lang.types.DType import sk.ainet.lang.types.DTypePolicy +import sk.ainet.lang.types.FP16 import sk.ainet.lang.types.FP32 import kotlin.math.pow import kotlin.reflect.KClass @@ -29,16 +32,20 @@ import kotlin.reflect.KClass * - HuggingFace → GGUF tensor name mapping * - Q4 + .qb companion tensor dequantization to FP32 * - BF16/F16 dequantization to FP32 (default) - * - BF16 KEEP_NATIVE when [dtypePolicy] admits BF16 (SKaiNET 0.25.0): - * constructs a [Bf16DenseTensorData]-backed tensor so the BF16 matmul - * kernel routes via `DefaultCpuOpsJvm` without a 2× memory blow-up. + * - **Narrow-float KEEP_NATIVE** when [dtypePolicy] admits BF16 or F16 + * (SKaiNET 0.38.0): constructs a [Bf16DenseTensorData] / [Fp16DenseTensorData]-backed + * tensor so the narrow-float matmul kernel routes via `DefaultCpuOpsJvm` + * without a 2× memory blow-up. * - Shape normalization ([1, dim] norms → [dim]) * - Tied word embeddings (output.weight = token_embd.weight) * * @param dtypePolicy declarative dtype constraint. Default [DTypePolicy.Any] - * = adaptive dequant. `Require(BF16)` / `Prefer(BF16)` / `OneOf` containing - * BF16 = KEEP_NATIVE path. Mirrors the SKaiNET 0.25.0 - * `SafeTensorsParametersLoader.mapPolicyToBf16` semantics. + * = widen everything to FP32. `Require(X)` / `Prefer(X)` / `OneOf` containing + * X keeps X-encoded source tensors packed, for X in {BF16, FP16}. The two + * formats are resolved independently — `Require(BF16)` still widens F16 + * sources, since neither narrow format can be re-encoded as the other without + * a lossy round-trip. Mirrors the engine-side + * `SafeTensorsParametersLoader.mapPolicyToNarrow` semantics. */ public class DecoderSafeTensorsLoader( private val ctx: ExecutionContext, @@ -50,16 +57,14 @@ public class DecoderSafeTensorsLoader( /** * Returns `true` iff [dtypePolicy] wants BF16 weights kept in their - * packed 2-bytes-per-element form rather than dequantised to FP32. - * Matches the engine-side `SafeTensorsParametersLoader.mapPolicyToBf16` - * cases that resolve to `Bf16LoadPolicy.KEEP_NATIVE`. + * packed 2-bytes-per-element form rather than widened to FP32. + * Matches the engine-side `SafeTensorsParametersLoader.mapPolicyToNarrow` + * cases that resolve to `NarrowFloatLoadPolicy.KEEP_NATIVE`. */ - private val keepBf16Native: Boolean = when (val p = dtypePolicy) { - DTypePolicy.Any -> false - is DTypePolicy.Require -> p.target == BF16 - is DTypePolicy.Prefer -> p.target == BF16 - is DTypePolicy.OneOf -> BF16 in p.allowed - } + private val keepBf16Native: Boolean = DTypePolicyValidation.keepsNative(dtypePolicy, BF16) + + /** As [keepBf16Native], for IEEE binary16 sources. Resolved independently. */ + private val keepFp16Native: Boolean = DTypePolicyValidation.keepsNative(dtypePolicy, FP16) /** * Load weights from SafeTensors file into a flat tensor map with GGUF-canonical names. @@ -96,9 +101,9 @@ public class DecoderSafeTensorsLoader( if (keepBf16Native) { // KEEP_NATIVE: wrap the packed 2-bytes-per-element // BF16 buffer as `Bf16DenseTensorData`. The matmul - // dispatch in `DefaultCpuOpsJvm` (SKaiNET 0.25.0) - // detects `Bf16TensorData` at runtime and routes - // to the SIMD BF16 kernel — avoiding the 2× memory + // dispatch in `DefaultCpuOpsJvm` (SKaiNET 0.38.0) + // detects `NarrowFloatTensorData` at runtime and + // picks the kernel by codec — avoiding the 2× memory // inflation of the FP32 dequant path. // // The declared dtype generic stays `T` (typically @@ -117,10 +122,21 @@ public class DecoderSafeTensorsLoader( } DataType.FLOAT16 -> { val bytes = reader.loadTensorData(info) - val floats = dequantF16(bytes) val targetShape = normalizeNormShape(info.shape) - @Suppress("UNCHECKED_CAST") - ctx.fromFloatArray(targetShape, dtype, floats) as Tensor + if (keepFp16Native) { + // Mirrors the BF16 arm above. Distinct from it on purpose: + // both formats are 2 bytes per element, so handing F16 bytes + // to the BF16 decode would not throw — it would produce + // plausible-looking wrong numbers. The codec carried by + // `Fp16DenseTensorData` is what keeps the dispatch honest. + val data = Fp16DenseTensorData.fromRawBytes(targetShape, bytes) + @Suppress("UNCHECKED_CAST") + ctx.fromData(data as TensorData, dtype) as Tensor + } else { + val floats = dequantF16(bytes) + @Suppress("UNCHECKED_CAST") + ctx.fromFloatArray(targetShape, dtype, floats) as Tensor + } } DataType.FLOAT32 -> { val bytes = reader.loadTensorData(info) diff --git a/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/LlamaNetworkLoader.kt b/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/LlamaNetworkLoader.kt index 52d31ba..596f690 100644 --- a/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/LlamaNetworkLoader.kt +++ b/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/LlamaNetworkLoader.kt @@ -67,11 +67,11 @@ public class LlamaNetworkLoader @PublishedApi internal constructor( } /** - * Declarative dtype policy attached via [withDtypePolicy]. SKaiNET 0.25.0 - * `DTypePolicy` is a forward-compat hook here — the value is validated - * eagerly but the underlying `DecoderGgufWeightLoader` / - * `DecoderSafeTensorsLoader` chain does not yet honor it per-tensor. - * Default [DTypePolicy.Any] preserves the adaptive behaviour. + * Declarative dtype policy attached via [withDtypePolicy]. Honored per-tensor by both + * chains as of engine 0.38.0: a policy naming BF16 or FP16 keeps source tensors of + * *that* format in their on-disk 2-bytes-per-element layout, in + * `DecoderSafeTensorsLoader` and in `DecoderGgufWeightLoader` alike. + * Default [DTypePolicy.Any] widens every narrow float to FP32. */ public var dtypePolicy: DTypePolicy = DTypePolicy.Any private set @@ -82,8 +82,9 @@ public class LlamaNetworkLoader @PublishedApi internal constructor( * not deep inside the load loop. */ public fun withDtypePolicy(policy: DTypePolicy): LlamaNetworkLoader { - val allowBf16 = weightsProvider is WeightsProvider.SafeTensors - DTypePolicyValidation.validate(policy, "LlamaNetworkLoader.withDtypePolicy", allowBf16Require = allowBf16) + DTypePolicyValidation.validate( + policy, "LlamaNetworkLoader.withDtypePolicy", keepNative = DECODER_NARROW_KEEP_NATIVE, + ) this.dtypePolicy = policy return this } @@ -138,11 +139,15 @@ public class LlamaNetworkLoader @PublishedApi internal constructor( ): Module { val weights: DecoderGgufWeights = when (val wp = weightsProvider) { is WeightsProvider.GgufSource -> { - val loader = DecoderGgufWeightLoader(wp.sourceProvider, quantPolicy = wp.quantPolicy) + val loader = DecoderGgufWeightLoader( + wp.sourceProvider, quantPolicy = wp.quantPolicy, dtypePolicy = dtypePolicy, + ) loader.loadToMap(ctx) } is WeightsProvider.GgufRandomAccess -> { - val loader = DecoderGgufWeightLoader(wp.randomAccessProvider, quantPolicy = wp.quantPolicy) + val loader = DecoderGgufWeightLoader( + wp.randomAccessProvider, quantPolicy = wp.quantPolicy, dtypePolicy = dtypePolicy, + ) loader.loadToMapStreaming(ctx) } is WeightsProvider.SafeTensors -> { diff --git a/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/DecoderGgufNarrowFloatTest.kt b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/DecoderGgufNarrowFloatTest.kt new file mode 100644 index 0000000..60658fa --- /dev/null +++ b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/DecoderGgufNarrowFloatTest.kt @@ -0,0 +1,231 @@ +package sk.ainet.models.llama + +import kotlinx.io.Source +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.io.gguf.GGMLQuantizationType +import sk.ainet.io.gguf.dequant.DequantOps +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.data.Bf16DenseTensorData +import sk.ainet.lang.tensor.data.Bf16TensorData +import sk.ainet.lang.tensor.data.Fp16DenseTensorData +import sk.ainet.lang.tensor.data.NarrowFloatTensorData +import sk.ainet.lang.types.BF16 +import sk.ainet.lang.types.DTypePolicy +import sk.ainet.lang.types.FP16 +import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.Fp16Codec +import sk.ainet.lang.types.Int8 +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertFalse +import kotlin.test.assertSame +import kotlin.test.assertTrue + +/** + * Covers the GGUF narrow-float KEEP_NATIVE path added for engine 0.38.0 — the policy decision + * ([DecoderGgufWeightLoader.keepsNarrowNative]) and the tensor construction + * ([DecoderGgufWeightLoader.createNarrowTensor]) — without synthesizing a GGUF file. Neither the + * engine nor this repo ships a GGUF writer, and the surrounding parse/read machinery is already + * exercised by the dequant path; what is new and worth pinning is these two decisions. + * + * The construction test matters most. GGUF header dims are reversed relative to the logical + * row-major shape, so the FP32 path swaps the `Shape` and moves no bytes + * (`DequantOps.transposeColumnMajorToRowMajor` returns its input unchanged). The packed path has + * to do exactly the same: an actual element transpose here would hand the matmul kernel a + * silently transposed weight matrix — wrong numbers, no exception. + */ +class DecoderGgufNarrowFloatTest { + + private val ctx = DirectCpuExecutionContext() + private val noopSource: () -> Source = { error("source not used in these tests") } + + private fun loaderWith(policy: DTypePolicy) = + DecoderGgufWeightLoader(sourceProvider = noopSource, dtypePolicy = policy) + + private fun fp16Bytes(values: FloatArray): ByteArray { + val out = ByteArray(values.size * 2) + for (i in values.indices) { + val bits = Fp16Codec.encode(values[i]) + out[i * 2] = (bits and 0xFF).toByte() + out[i * 2 + 1] = ((bits ushr 8) and 0xFF).toByte() + } + return out + } + + private fun bf16Bytes(values: FloatArray): ByteArray { + val out = ByteArray(values.size * 2) + for (i in values.indices) { + val bits = (values[i].toRawBits() ushr 16) and 0xFFFF + out[i * 2] = (bits and 0xFF).toByte() + out[i * 2 + 1] = ((bits ushr 8) and 0xFF).toByte() + } + return out + } + + // ---------- policy resolution ---------- + + @Test + fun `default policy keeps nothing native`() { + val loader = loaderWith(DTypePolicy.Any) + assertFalse(loader.keepsNarrowNative(GGMLQuantizationType.F16, FP32::class)) + assertFalse(loader.keepsNarrowNative(GGMLQuantizationType.BF16, FP32::class)) + } + + @Test + fun `a policy naming one narrow format leaves the other on the widening path`() { + val f16 = loaderWith(DTypePolicy.Require(FP16)) + assertTrue(f16.keepsNarrowNative(GGMLQuantizationType.F16, FP32::class)) + assertFalse( + f16.keepsNarrowNative(GGMLQuantizationType.BF16, FP32::class), + "BF16 cannot be re-encoded as F16 — it must widen", + ) + + val bf16 = loaderWith(DTypePolicy.Require(BF16)) + assertTrue(bf16.keepsNarrowNative(GGMLQuantizationType.BF16, FP32::class)) + assertFalse(bf16.keepsNarrowNative(GGMLQuantizationType.F16, FP32::class)) + } + + @Test + fun `soft policies reach the same KEEP_NATIVE decision`() { + assertTrue( + loaderWith(DTypePolicy.Prefer(BF16)).keepsNarrowNative(GGMLQuantizationType.BF16, FP32::class), + ) + assertTrue( + loaderWith(DTypePolicy.OneOf(setOf(FP32, FP16))) + .keepsNarrowNative(GGMLQuantizationType.F16, FP32::class), + ) + } + + @Test + fun `quantized and F32 source types are never treated as narrow`() { + val loader = loaderWith(DTypePolicy.OneOf(setOf(BF16, FP16))) + assertFalse(loader.keepsNarrowNative(GGMLQuantizationType.F32, FP32::class)) + assertFalse(loader.keepsNarrowNative(GGMLQuantizationType.Q4_K, FP32::class)) + assertFalse(loader.keepsNarrowNative(GGMLQuantizationType.Q8_0, FP32::class)) + } + + @Test + fun `KEEP_NATIVE only applies to an FP32 element type`() { + val loader = loaderWith(DTypePolicy.Require(FP16)) + // FP16::class here is a request for the FP32-array storage path, not a packing request; + // Int8 is the RAW_BYTES path. Neither may be silently reinterpreted as packed storage. + assertFalse(loader.keepsNarrowNative(GGMLQuantizationType.F16, FP16::class)) + assertFalse(loader.keepsNarrowNative(GGMLQuantizationType.F16, Int8::class)) + } + + // ---------- tensor construction ---------- + + @Test + fun `rank-2 construction swaps the shape and moves no bytes`() { + // GGUF header dims [rows=2, cols=4] describe a logical [4, 2] row-major tensor. + val values = floatArrayOf(1.0f, 2.0f, 4.0f, 8.0f, 16.0f, 32.0f, 64.0f, 128.0f) + val bytes = fp16Bytes(values) + val loader = loaderWith(DTypePolicy.Require(FP16)) + + val tensor = loader.createNarrowTensor( + ctx, FP32::class, Shape(2, 4), bytes, GGMLQuantizationType.F16, + ) + + assertContentEquals( + intArrayOf(4, 2), tensor.shape.dimensions, + "GGUF [rows, cols] must be reinterpreted as [cols, rows]", + ) + val data = tensor.data as Fp16DenseTensorData + assertSame( + bytes, data.packedData, + "the on-disk buffer must become the tensor's storage — no copy, no transpose", + ) + // Element order is untouched, so a flat decode matches the source values in file order. + assertContentEquals(values, data.copyToFloatArray()) + } + + @Test + fun `rank-1 construction passes the shape through`() { + val values = floatArrayOf(0.5f, -0.5f, 3.0f, -7.0f) + val loader = loaderWith(DTypePolicy.Require(BF16)) + + val tensor = loader.createNarrowTensor( + ctx, FP32::class, Shape(4), bf16Bytes(values), GGMLQuantizationType.BF16, + ) + + assertContentEquals(intArrayOf(4), tensor.shape.dimensions) + assertContentEquals(values, tensor.data.copyToFloatArray()) + } + + @Test + fun `the source type picks the codec, not the byte width`() { + // Both formats are 2 bytes per element, so a mix-up cannot fail loudly. Pin it from both + // sides: the right wrapper type, and a decode that visibly differs from the other codec's. + val values = floatArrayOf(1.0f, 2.0f, 4.0f, 8.0f) + val loader = loaderWith(DTypePolicy.OneOf(setOf(BF16, FP16))) + + val asF16 = loader.createNarrowTensor( + ctx, FP32::class, Shape(4), fp16Bytes(values), GGMLQuantizationType.F16, + ) + assertTrue(asF16.data is Fp16DenseTensorData) + assertFalse(asF16.data is Bf16TensorData, "F16 must never be mistaken for BF16") + assertContentEquals(values, asF16.data.copyToFloatArray()) + + val asBf16 = loader.createNarrowTensor( + ctx, FP32::class, Shape(4), bf16Bytes(values), GGMLQuantizationType.BF16, + ) + assertTrue(asBf16.data is Bf16DenseTensorData) + assertContentEquals(values, asBf16.data.copyToFloatArray()) + + // The same bytes read through the wrong codec do NOT coincide — so a dispatch that + // confused the two could not pass these assertions by luck. + val f16BytesOfValues = fp16Bytes(values) + val misread = Bf16DenseTensorData.fromRawBytes(Shape(4), f16BytesOfValues).copyToFloatArray() + assertTrue( + misread.indices.any { kotlin.math.abs(misread[it] - values[it]) > 1e-3f }, + "test is vacuous if the two codecs agree on these bytes", + ) + } + + @Test + fun `KEEP_NATIVE decoding matches the widening path bit for bit`() { + // The widening path for GGUF F16 is DequantOps.dequantF16FromBytes; KEEP_NATIVE defers + // the identical decode to read time. Values are chosen not to be exact in binary16. + val values = FloatArray(32) { (it - 16) * 0.1f } + val bytes = fp16Bytes(values) + val loader = loaderWith(DTypePolicy.Require(FP16)) + + val widened = DequantOps.dequantF16FromBytes(bytes) + val native = loader.createNarrowTensor( + ctx, FP32::class, Shape(32), bytes, GGMLQuantizationType.F16, + ).data.copyToFloatArray() + + assertEquals(widened.size, native.size) + for (i in widened.indices) { + assertEquals( + widened[i].toRawBits(), native[i].toRawBits(), + "bit-identity expected at $i: widened=${widened[i]} native=${native[i]}", + ) + } + } + + @Test + fun `a short buffer is rejected rather than read out of bounds`() { + val loader = loaderWith(DTypePolicy.Require(FP16)) + val toosmall = ByteArray(6) // 3 elements' worth for a 2x4 = 8-element tensor + val error = kotlin.runCatching { + loader.createNarrowTensor( + ctx, FP32::class, Shape(2, 4), toosmall, GGMLQuantizationType.F16, + ) + }.exceptionOrNull() + assertTrue(error is IllegalArgumentException, "expected IllegalArgumentException, got $error") + } + + @Test + fun `the wrapper reports itself to narrow-float dispatch`() { + val loader = loaderWith(DTypePolicy.Require(BF16)) + val tensor = loader.createNarrowTensor( + ctx, FP32::class, Shape(4), bf16Bytes(floatArrayOf(1f, 2f, 3f, 4f)), GGMLQuantizationType.BF16, + ) + assertTrue( + tensor.data is NarrowFloatTensorData, + "DefaultCpuOpsJvm.chooseQuantizedMatmul matches on NarrowFloatTensorData", + ) + } +} diff --git a/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/DecoderSafeTensorsLoaderNarrowFloatTest.kt b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/DecoderSafeTensorsLoaderNarrowFloatTest.kt new file mode 100644 index 0000000..9fab628 --- /dev/null +++ b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/DecoderSafeTensorsLoaderNarrowFloatTest.kt @@ -0,0 +1,211 @@ +package sk.ainet.models.llama + +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.io.RandomAccessSource +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.Bf16TensorData +import sk.ainet.lang.tensor.data.FloatArrayTensorData +import sk.ainet.lang.tensor.data.Fp16DenseTensorData +import sk.ainet.lang.tensor.data.NarrowFloatTensorData +import sk.ainet.lang.types.BF16 +import sk.ainet.lang.types.DTypePolicy +import sk.ainet.lang.types.FP16 +import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.Fp16Codec +import java.io.File +import java.nio.ByteBuffer +import java.nio.ByteOrder +import java.nio.file.Files +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +/** + * Pins the narrow-float KEEP_NATIVE behaviour of [DecoderSafeTensorsLoader] — the + * transformer-repo counterpart of the engine's `SafeTensorsParametersLoaderFp16PolicyTest`. + * + * The BF16 arm has existed since 0.25.0; the F16 arm landed with engine 0.38.0's + * `Fp16DenseTensorData`. What matters here is that the loader resolves the two formats + * **independently**: both are 2 bytes per element, so routing F16 bytes through the BF16 decode + * would not throw — it would quietly produce wrong numbers. These tests assert the packed bytes + * survive verbatim, decode bit-identically to the widening path, and that neither policy leaks + * into the other format. + * + * Files are synthesized in-test; no model downloads are involved. + */ +class DecoderSafeTensorsLoaderNarrowFloatTest { + + /** A canonical-mappable HF weight name, so [HfTensorNameMapper] doesn't skip the tensor. */ + private val hfName = "model.layers.0.self_attn.q_proj.weight" + private val canonical = LlamaTensorNames.attnQ(0) + + private val metadata = LlamaModelMetadata( + architecture = "llama", + embeddingLength = 4, + contextLength = 8, + blockCount = 1, + headCount = 1, + kvHeadCount = 1, + feedForwardLength = 4, + ropeDimensionCount = 4, + vocabSize = 4, + ) + + private fun fp32ToFp16Bytes(values: FloatArray): ByteArray { + val out = ByteArray(values.size * 2) + for (i in values.indices) { + val bits = Fp16Codec.encode(values[i]) + out[i * 2] = (bits and 0xFF).toByte() + out[i * 2 + 1] = ((bits ushr 8) and 0xFF).toByte() + } + return out + } + + private fun fp32ToBf16Bytes(values: FloatArray): ByteArray { + val out = ByteArray(values.size * 2) + for (i in values.indices) { + val bits = (values[i].toRawBits() ushr 16) and 0xFFFF + out[i * 2] = (bits and 0xFF).toByte() + out[i * 2 + 1] = ((bits ushr 8) and 0xFF).toByte() + } + return out + } + + /** Write a single-tensor SafeTensors file: 8-byte LE header length, JSON header, then data. */ + private fun writeSafeTensors(entries: List>, rows: Int, cols: Int): File { + val header = StringBuilder("{") + var offset = 0L + entries.forEachIndexed { i, (name, dtype, bytes) -> + if (i > 0) header.append(",") + header.append( + "\"$name\": {\"dtype\": \"$dtype\", \"shape\": [$rows, $cols], " + + "\"data_offsets\": [$offset, ${offset + bytes.size}]}", + ) + offset += bytes.size + } + header.append("}") + val headerBytes = header.toString().toByteArray(Charsets.UTF_8) + + val file = Files.createTempFile("decoder_st_narrow", ".safetensors").toFile() + file.deleteOnExit() + file.outputStream().use { out -> + out.write( + ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN) + .putLong(headerBytes.size.toLong()).array(), + ) + out.write(headerBytes) + entries.forEach { (_, _, bytes) -> out.write(bytes) } + } + return file + } + + private fun load(file: File, policy: DTypePolicy): Map> { + val ctx = DirectCpuExecutionContext() + val loader = DecoderSafeTensorsLoader( + ctx = ctx, + dtype = FP32::class, + metadata = metadata, + tiedEmbeddings = false, + dtypePolicy = policy, + ) + val provider: () -> RandomAccessSource = { JvmRandomAccessSource.open(file) } + return loader.loadToMap(provider).tensors + } + + /** 2x4, all exactly representable in binary16 AND bfloat16 so decode comparisons are exact. */ + private val values = floatArrayOf(0.0f, 1.0f, -1.0f, 0.5f, 2.0f, -4.0f, 0.25f, 8.0f) + + @Test + fun `default policy widens F16 to an FP32 float array`() { + val file = writeSafeTensors(listOf(Triple(hfName, "F16", fp32ToFp16Bytes(values))), 2, 4) + + val weight = load(file, DTypePolicy.Any)[canonical] ?: error("missing $canonical") + assertTrue( + weight.data is FloatArrayTensorData<*>, + "default policy must widen, got ${weight.data::class.simpleName}", + ) + assertContentEquals(values, weight.data.copyToFloatArray()) + } + + @Test + fun `Require(FP16) keeps the on-disk F16 bytes verbatim`() { + val onDisk = fp32ToFp16Bytes(values) + val file = writeSafeTensors(listOf(Triple(hfName, "F16", onDisk)), 2, 4) + + val weight = load(file, DTypePolicy.Require(FP16))[canonical] ?: error("missing $canonical") + + assertTrue( + weight.data is Fp16DenseTensorData, + "KEEP_NATIVE must produce Fp16DenseTensorData, got ${weight.data::class.simpleName}", + ) + assertTrue(weight.data is NarrowFloatTensorData, "must be recognizable to narrow dispatch") + assertTrue( + weight.data !is Bf16TensorData, + "an F16 tensor must never be mistaken for BF16 — the bit layouts differ", + ) + // Byte-for-byte identity proves no widening pass ran. + assertContentEquals( + onDisk, (weight.data as Fp16DenseTensorData).packedData, + "KEEP_NATIVE must preserve on-disk F16 bytes verbatim", + ) + assertEquals(values.size * 2, (weight.data as Fp16DenseTensorData).packedData.size) + } + + @Test + fun `KEEP_NATIVE decodes bit-identically to the widening path`() { + // Both paths apply the same binary16 decode; only the timing differs. Values here are + // deliberately not all exact in binary16, so a rounding difference would show up. + val wide = FloatArray(64) { (it - 32) * 0.1f } + val file = writeSafeTensors(listOf(Triple(hfName, "F16", fp32ToFp16Bytes(wide))), 8, 8) + + val widened = load(file, DTypePolicy.Any)[canonical]!!.data.copyToFloatArray() + val native = load(file, DTypePolicy.Require(FP16))[canonical]!!.data.copyToFloatArray() + + assertEquals(widened.size, native.size) + for (i in widened.indices) { + assertEquals( + widened[i].toRawBits(), native[i].toRawBits(), + "bit-identity expected at $i: widened=${widened[i]} native=${native[i]}", + ) + } + } + + @Test + fun `a policy naming one narrow format widens the other`() { + val f16Name = "model.layers.0.self_attn.q_proj.weight" + val bf16Name = "model.layers.0.self_attn.k_proj.weight" + val file = writeSafeTensors( + listOf( + Triple(f16Name, "F16", fp32ToFp16Bytes(values)), + Triple(bf16Name, "BF16", fp32ToBf16Bytes(values)), + ), + 2, 4, + ) + val f16Canonical = LlamaTensorNames.attnQ(0) + val bf16Canonical = LlamaTensorNames.attnK(0) + + // Require(FP16): F16 stays packed, BF16 widens — it cannot be re-encoded as F16. + val a = load(file, DTypePolicy.Require(FP16)) + assertTrue(a[f16Canonical]!!.data is Fp16DenseTensorData, "F16 should be packed") + assertTrue(a[bf16Canonical]!!.data is FloatArrayTensorData<*>, "BF16 should be widened") + + // ...and the mirror image. + val b = load(file, DTypePolicy.Require(BF16)) + assertTrue(b[f16Canonical]!!.data is FloatArrayTensorData<*>, "F16 should be widened") + assertTrue(b[bf16Canonical]!!.data is Bf16TensorData, "BF16 should be packed") + } + + @Test + fun `Prefer and OneOf reach the same KEEP_NATIVE path as Require`() { + val file = writeSafeTensors(listOf(Triple(hfName, "F16", fp32ToFp16Bytes(values))), 2, 4) + + assertTrue(load(file, DTypePolicy.Prefer(FP16))[canonical]!!.data is Fp16DenseTensorData) + assertTrue( + load(file, DTypePolicy.OneOf(setOf(FP32, FP16)))[canonical]!!.data is Fp16DenseTensorData, + ) + // A soft policy naming neither narrow format leaves the widening default in place. + assertTrue(load(file, DTypePolicy.Prefer(FP32))[canonical]!!.data is FloatArrayTensorData<*>) + } +} diff --git a/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/LlamaNetworkLoaderDTypePolicyTest.kt b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/LlamaNetworkLoaderDTypePolicyTest.kt index c88d7bc..f633308 100644 --- a/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/LlamaNetworkLoaderDTypePolicyTest.kt +++ b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/LlamaNetworkLoaderDTypePolicyTest.kt @@ -18,13 +18,14 @@ import sk.ainet.lang.types.Int8 * - the default value is [DTypePolicy.Any]; * - `withDtypePolicy(Require(FP32))` always succeeds (it's the loader's * native output dtype); - * - `withDtypePolicy(Require(BF16))` succeeds **only** for SafeTensors- - * backed loaders — the GGUF path mirrors the engine's eager rejection - * in `StreamingGgufParametersLoader.validatePolicy()` because the - * transformer-repo GGUF chain still dequants BF16 to FP32; - * - `withDtypePolicy(Require(FP16))` / `Require(Int8)` etc. always - * reject — the loader doesn't fabricate dtypes the source files - * don't carry. + * - `withDtypePolicy(Require(BF16))` and `Require(FP16)` succeed on + * **both** paths as of engine 0.38.0: `DecoderSafeTensorsLoader` and + * `DecoderGgufWeightLoader` each keep narrow-float sources in their + * on-disk 2-bytes-per-element layout (see + * [DECODER_NARROW_KEEP_NATIVE]). Before that, GGUF rejected BF16 and + * both paths rejected FP16 for want of an `Fp16DenseTensorData`; + * - `withDtypePolicy(Require(Int8))` etc. still reject — the loader + * doesn't fabricate dtypes the source files don't carry; * - `Prefer` / `OneOf` arms never raise (they're soft constraints). * * No model files are read — these tests only construct loader instances @@ -69,29 +70,27 @@ class LlamaNetworkLoaderDTypePolicyTest { } @Test - fun `Require(BF16) is accepted on SafeTensors but rejected on GGUF`() { + fun `Require(BF16) is accepted on both GGUF and SafeTensors paths`() { val safetensors = LlamaNetworkLoader.fromSafeTensors( metadata = anyMetadata, randomAccessProvider = noopRandomAccessProvider, ).withDtypePolicy(DTypePolicy.Require(BF16)) assertEquals(DTypePolicy.Require(BF16), safetensors.dtypePolicy) - assertFailsWith { - LlamaNetworkLoader.fromGguf(sourceProvider = noopSourceProvider) - .withDtypePolicy(DTypePolicy.Require(BF16)) - } + val gguf = LlamaNetworkLoader.fromGguf(sourceProvider = noopSourceProvider) + .withDtypePolicy(DTypePolicy.Require(BF16)) + assertEquals(DTypePolicy.Require(BF16), gguf.dtypePolicy) } @Test - fun `Require(FP16) is rejected on both paths`() { - assertFailsWith { - LlamaNetworkLoader.fromGguf(sourceProvider = noopSourceProvider) - .withDtypePolicy(DTypePolicy.Require(FP16)) - } - assertFailsWith { - LlamaNetworkLoader.fromSafeTensors( - metadata = anyMetadata, randomAccessProvider = noopRandomAccessProvider, - ).withDtypePolicy(DTypePolicy.Require(FP16)) - } + fun `Require(FP16) is accepted on both GGUF and SafeTensors paths`() { + val gguf = LlamaNetworkLoader.fromGguf(sourceProvider = noopSourceProvider) + .withDtypePolicy(DTypePolicy.Require(FP16)) + assertEquals(DTypePolicy.Require(FP16), gguf.dtypePolicy) + + val safetensors = LlamaNetworkLoader.fromSafeTensors( + metadata = anyMetadata, randomAccessProvider = noopRandomAccessProvider, + ).withDtypePolicy(DTypePolicy.Require(FP16)) + assertEquals(DTypePolicy.Require(FP16), safetensors.dtypePolicy) } @Test diff --git a/llm-inference/qwen/src/commonMain/kotlin/sk/ainet/models/qwen/QwenNetworkLoader.kt b/llm-inference/qwen/src/commonMain/kotlin/sk/ainet/models/qwen/QwenNetworkLoader.kt index 0d26e15..2141d67 100644 --- a/llm-inference/qwen/src/commonMain/kotlin/sk/ainet/models/qwen/QwenNetworkLoader.kt +++ b/llm-inference/qwen/src/commonMain/kotlin/sk/ainet/models/qwen/QwenNetworkLoader.kt @@ -16,6 +16,7 @@ import sk.ainet.models.llama.LlamaModelMetadata import sk.ainet.models.llama.DecoderSafeTensorsLoader import sk.ainet.models.llama.DecoderGgufWeightLoader import sk.ainet.models.llama.DecoderGgufWeights +import sk.ainet.models.llama.DECODER_NARROW_KEEP_NATIVE import kotlin.jvm.JvmName /** @@ -53,8 +54,9 @@ public class QwenNetworkLoader @PublishedApi internal constructor( /** See [LlamaNetworkLoader.withDtypePolicy]. */ public fun withDtypePolicy(policy: DTypePolicy): QwenNetworkLoader { - val allowBf16 = weightsProvider is WeightsProvider.SafeTensors - DTypePolicyValidation.validate(policy, "QwenNetworkLoader.withDtypePolicy", allowBf16Require = allowBf16) + DTypePolicyValidation.validate( + policy, "QwenNetworkLoader.withDtypePolicy", keepNative = DECODER_NARROW_KEEP_NATIVE, + ) this.dtypePolicy = policy return this } @@ -135,7 +137,8 @@ public class QwenNetworkLoader @PublishedApi internal constructor( val loader = DecoderGgufWeightLoader( wp.sourceProvider, quantPolicy = wp.quantPolicy, - acceptedArchitectures = QWEN_ARCHITECTURES + acceptedArchitectures = QWEN_ARCHITECTURES, + dtypePolicy = dtypePolicy, ) loader.loadToMap(ctx) } @@ -143,7 +146,8 @@ public class QwenNetworkLoader @PublishedApi internal constructor( val loader = DecoderGgufWeightLoader( wp.randomAccessProvider, quantPolicy = wp.quantPolicy, - acceptedArchitectures = QWEN_ARCHITECTURES + acceptedArchitectures = QWEN_ARCHITECTURES, + dtypePolicy = dtypePolicy, ) loader.loadToMapStreaming(ctx) } diff --git a/llm-inference/voxtral/src/commonMain/kotlin/sk/ainet/models/voxtral/VoxtralNetworkLoader.kt b/llm-inference/voxtral/src/commonMain/kotlin/sk/ainet/models/voxtral/VoxtralNetworkLoader.kt index 7633285..c449623 100644 --- a/llm-inference/voxtral/src/commonMain/kotlin/sk/ainet/models/voxtral/VoxtralNetworkLoader.kt +++ b/llm-inference/voxtral/src/commonMain/kotlin/sk/ainet/models/voxtral/VoxtralNetworkLoader.kt @@ -17,6 +17,7 @@ import sk.ainet.models.llama.LlamaModelMetadata import sk.ainet.models.llama.DecoderSafeTensorsLoader import sk.ainet.models.llama.DecoderGgufWeightLoader import sk.ainet.models.llama.DecoderGgufWeights +import sk.ainet.models.llama.DECODER_NARROW_KEEP_NATIVE import kotlin.jvm.JvmName import kotlin.reflect.KClass @@ -80,8 +81,9 @@ public class VoxtralNetworkLoader @PublishedApi internal constructor( /** See [sk.ainet.models.llama.LlamaNetworkLoader.withDtypePolicy]. */ public fun withDtypePolicy(policy: DTypePolicy): VoxtralNetworkLoader { - val allowBf16 = weightsProvider is WeightsProvider.SafeTensors - DTypePolicyValidation.validate(policy, "VoxtralNetworkLoader.withDtypePolicy", allowBf16Require = allowBf16) + DTypePolicyValidation.validate( + policy, "VoxtralNetworkLoader.withDtypePolicy", keepNative = DECODER_NARROW_KEEP_NATIVE, + ) this.dtypePolicy = policy return this } @@ -162,11 +164,15 @@ public class VoxtralNetworkLoader @PublishedApi internal constructor( ): DecoderGgufWeights { return when (val wp = weightsProvider) { is WeightsProvider.GgufSource -> { - val loader = DecoderGgufWeightLoader(wp.sourceProvider, quantPolicy = wp.quantPolicy) + val loader = DecoderGgufWeightLoader( + wp.sourceProvider, quantPolicy = wp.quantPolicy, dtypePolicy = dtypePolicy, + ) loader.loadToMap(ctx) } is WeightsProvider.GgufRandomAccess -> { - val loader = DecoderGgufWeightLoader(wp.randomAccessProvider, quantPolicy = wp.quantPolicy) + val loader = DecoderGgufWeightLoader( + wp.randomAccessProvider, quantPolicy = wp.quantPolicy, dtypePolicy = dtypePolicy, + ) loader.loadToMapStreaming(ctx) } is WeightsProvider.SafeTensors -> { From 8e04e67fb4cdaf71125e89fd5a524a3e9c832d74 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Mon, 27 Jul 2026 09:51:18 +0200 Subject: [PATCH 2/6] test(dtype): add forward parity test for narrow-float KEEP_NATIVE MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Loading tests proved the packed bytes survive and decode correctly, but nothing exercised those tensors once they reached the model. Both narrow formats are 2 bytes per element, so decoding one as the other does not throw — it yields finite, plausible, wrong logits. Synthesize a complete 1-layer LLaMA as SafeTensors, load it twice from the same file (widened vs KEEP_NATIVE), run both through OptimizedLLMRuntime over a 4-token sequence, and compare logits. FP16 and BF16 both match bit-for-bit, including through embedding gather and the RMSNorm weight multiply. Guards against a vacuous pass: - every loaded tensor must be NarrowFloatTensorData, else the comparison would silently be FP32 vs FP32 - weights are round-tripped through the codec before writing, so both sides hold identical values - a codec mix-up case writes F16 bit patterns into a BF16-declared file and asserts the logits move by more than 10x the tolerance Also pin the reason the two paths agree exactly. LlamaRuntime.linearProject calls w.t() on [out, in] weights, transpose has no narrow-float path and widens to a dense FP32 buffer, and chooseQuantizedMatmul only engages for [in, out] weights. The FP16/BF16 kernels are therefore unreachable on this chain: packed weights are widened on every forward, so the saving is at-rest memory only, paid for with a per-token decode and a transpose allocation. Removing the .t() should flip that test and make the parity cases exercise the kernel path they were written for. --- .../DecoderNarrowFloatForwardParityTest.kt | 360 ++++++++++++++++++ 1 file changed, 360 insertions(+) create mode 100644 llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/DecoderNarrowFloatForwardParityTest.kt diff --git a/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/DecoderNarrowFloatForwardParityTest.kt b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/DecoderNarrowFloatForwardParityTest.kt new file mode 100644 index 0000000..a7ed7f3 --- /dev/null +++ b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/DecoderNarrowFloatForwardParityTest.kt @@ -0,0 +1,360 @@ +package sk.ainet.models.llama + +import sk.ainet.apps.llm.OptimizedLLMMode +import sk.ainet.apps.llm.OptimizedLLMRuntime +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.io.RandomAccessSource +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.data.Fp16DenseTensorData +import sk.ainet.lang.tensor.data.NarrowFloatTensorData +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.tensor.t +import sk.ainet.lang.types.BF16 +import sk.ainet.lang.types.DTypePolicy +import sk.ainet.lang.types.FP16 +import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.Fp16Codec +import java.io.File +import java.nio.ByteBuffer +import java.nio.ByteOrder +import java.nio.file.Files +import kotlin.math.abs +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +/** + * End-to-end parity for narrow-float KEEP_NATIVE: a complete tiny LLaMA is loaded twice from the + * *same* SafeTensors file — once widened to FP32 at load, once kept packed — and both are run + * through [OptimizedLLMRuntime]. The logits must agree. + * + * Why this test exists, when [DecoderSafeTensorsLoaderNarrowFloatTest] already pins the loader: + * that test proves the bytes survive and decode correctly, which is only half the feature. The + * other half is what happens once those tensors reach the model — every op that touches a packed + * weight has to decode it correctly, and since binary16 and bfloat16 are both 2 bytes per element, + * decoding one as the other does not throw. It returns finite, plausible, wrong logits. Only a + * numeric comparison against a known-good reference catches that. + * + * Both sides of the comparison hold mathematically identical weights: the file is written from + * values already round-tripped through the codec, so the widened path and the packed path decode + * the same numbers. [TOLERANCE] covers float accumulation-order differences. + * + * **What this does NOT prove.** It does not show that the narrow-float matmul kernel runs. On this + * chain it does not — see [`narrow weights are materialized to FP32 before matmul`], which pins + * why. The logits below are bit-identical rather than merely close, because both runs ultimately + * execute the same FP32 SGEMM; only the moment of decoding differs. [TOLERANCE] stays non-zero so + * the test keeps passing if a real narrow kernel is ever wired in and shifts accumulation order. + * + * The KEEP_NATIVE arm applies to *every* tensor in the file, so this does cover embedding gather + * and the RMSNorm weight multiply reading through packed storage. + * + * [`the parity check would catch a codec mix-up`] is the guard that keeps the tolerance honest. + * + * Files are synthesized in-test; no model downloads are involved. + */ +class DecoderNarrowFloatForwardParityTest { + + private val dim = 8 + private val ffDim = 16 + private val vocabSize = 16 + private val nHeads = 2 + private val kvHeads = 2 + private val headDim = dim / nHeads + private val seqLen = 32 + + /** + * Absolute logit tolerance. Today both paths converge on the same FP32 SGEMM and agree to the + * bit, so this is slack for a future narrow kernel whose accumulation order would differ. + * Sized to swallow that and nothing more — the codec mix-up guard measures the margin actually + * available and fails if this ever grows large enough to hide a real defect. + */ + private val TOLERANCE = 1e-4f + + private val metadata = LlamaModelMetadata( + architecture = "llama", + embeddingLength = dim, + contextLength = seqLen, + blockCount = 1, + headCount = nHeads, + kvHeadCount = kvHeads, + feedForwardLength = ffDim, + ropeDimensionCount = headDim, + vocabSize = vocabSize, + ) + + // ---------------------------------------------------------------- codecs + + private fun encodeFp16(values: FloatArray): ByteArray { + val out = ByteArray(values.size * 2) + for (i in values.indices) { + val bits = Fp16Codec.encode(values[i]) + out[i * 2] = (bits and 0xFF).toByte() + out[i * 2 + 1] = ((bits ushr 8) and 0xFF).toByte() + } + return out + } + + private fun encodeBf16(values: FloatArray): ByteArray { + val out = ByteArray(values.size * 2) + for (i in values.indices) { + val bits = (values[i].toRawBits() ushr 16) and 0xFFFF + out[i * 2] = (bits and 0xFF).toByte() + out[i * 2 + 1] = ((bits ushr 8) and 0xFF).toByte() + } + return out + } + + /** Round-trip through binary16 so the on-disk file loses nothing further. */ + private fun quantizeFp16(values: FloatArray): FloatArray = + FloatArray(values.size) { Fp16Codec.decode(Fp16Codec.encode(values[it])) } + + /** Round-trip through bfloat16 — truncate the low 16 mantissa bits. */ + private fun quantizeBf16(values: FloatArray): FloatArray = + FloatArray(values.size) { Float.fromBits(values[it].toRawBits() and 0xFFFF0000.toInt()) } + + // ------------------------------------------------------------ model data + + /** Deterministic small weights; the same generator [LlamaDslPipelineTest] uses. */ + private fun randn(size: Int, seed: Int): FloatArray { + val rng = kotlin.random.Random(seed) + return FloatArray(size) { (rng.nextFloat() - 0.5f) * 0.1f } + } + + private fun ones(size: Int): FloatArray = FloatArray(size) { 1.0f } + + /** HF-named weights with their SafeTensors shapes, in FP32 before any narrowing. */ + private fun buildHfWeights(): List, FloatArray>> = listOf( + Triple("model.embed_tokens.weight", listOf(vocabSize, dim), randn(vocabSize * dim, 10)), + Triple("model.norm.weight", listOf(dim), ones(dim)), + Triple("lm_head.weight", listOf(vocabSize, dim), randn(vocabSize * dim, 11)), + Triple("model.layers.0.input_layernorm.weight", listOf(dim), ones(dim)), + Triple("model.layers.0.self_attn.q_proj.weight", listOf(dim, dim), randn(dim * dim, 1)), + Triple("model.layers.0.self_attn.k_proj.weight", listOf(dim, dim), randn(dim * dim, 2)), + Triple("model.layers.0.self_attn.v_proj.weight", listOf(dim, dim), randn(dim * dim, 3)), + Triple("model.layers.0.self_attn.o_proj.weight", listOf(dim, dim), randn(dim * dim, 4)), + Triple("model.layers.0.post_attention_layernorm.weight", listOf(dim), ones(dim)), + Triple("model.layers.0.mlp.gate_proj.weight", listOf(ffDim, dim), randn(ffDim * dim, 5)), + Triple("model.layers.0.mlp.down_proj.weight", listOf(dim, ffDim), randn(dim * ffDim, 6)), + Triple("model.layers.0.mlp.up_proj.weight", listOf(ffDim, dim), randn(ffDim * dim, 7)), + ) + + /** + * Write every weight into one SafeTensors file: 8-byte LE header length, JSON header, data. + * + * [declaredDtype] is what the header claims; [encode] is what actually produces the bytes. + * They are separate parameters on purpose — the codec mix-up guard needs to declare one + * format while writing the other's bit layout. + */ + private fun writeModel( + weights: List, FloatArray>>, + declaredDtype: String, + encode: (FloatArray) -> ByteArray, + ): File { + val header = StringBuilder("{") + var offset = 0L + val payloads = weights.map { (name, shape, values) -> + val bytes = encode(values) + if (offset > 0) header.append(",") + header.append( + "\"$name\": {\"dtype\": \"$declaredDtype\", \"shape\": [${shape.joinToString(", ")}], " + + "\"data_offsets\": [$offset, ${offset + bytes.size}]}", + ) + offset += bytes.size + bytes + } + header.append("}") + val headerBytes = header.toString().toByteArray(Charsets.UTF_8) + + val file = Files.createTempFile("decoder_forward_parity", ".safetensors").toFile() + file.deleteOnExit() + file.outputStream().use { out -> + out.write( + ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN) + .putLong(headerBytes.size.toLong()).array(), + ) + out.write(headerBytes) + payloads.forEach { out.write(it) } + } + return file + } + + // -------------------------------------------------------------- the runs + + /** Load under [policy] and forward [tokens] in order, returning the logits of the last step. */ + private fun forwardLogits( + file: File, + policy: DTypePolicy, + tokens: IntArray, + onWeights: (DecoderGgufWeights) -> Unit = {}, + ): FloatArray { + val ctx = DirectCpuExecutionContext() + val loader = DecoderSafeTensorsLoader( + ctx = ctx, + dtype = FP32::class, + metadata = metadata, + tiedEmbeddings = false, + dtypePolicy = policy, + ) + val provider: () -> RandomAccessSource = { JvmRandomAccessSource.open(file) } + val weights = loader.loadToMap(provider) + onWeights(weights) + + val runtime = OptimizedLLMRuntime( + model = LlamaNetworkLoader.fromWeights(weights), + ctx = ctx, + mode = OptimizedLLMMode.DIRECT, + dtype = FP32::class, + ) + var logits = FloatArray(0) + for (token in tokens) { + logits = runtime.forward(token).data.copyToFloatArray() + } + return logits + } + + /** Multi-token so the comparison runs through the KV cache, not just a single step. */ + private val tokens = intArrayOf(1, 5, 3, 9) + + private fun assertParity(reference: FloatArray, native: FloatArray, label: String) { + assertEquals(vocabSize, reference.size, "reference logits should be one row of vocab") + assertEquals(reference.size, native.size, "$label: logit count differs") + for (i in reference.indices) { + assertTrue(native[i].isFinite(), "$label: logit[$i] is not finite (${native[i]})") + assertTrue( + abs(reference[i] - native[i]) <= TOLERANCE, + "$label: logit[$i] diverged — widened=${reference[i]} native=${native[i]} " + + "delta=${abs(reference[i] - native[i])} tolerance=$TOLERANCE", + ) + } + } + + /** Fails the run if KEEP_NATIVE silently didn't engage — otherwise this is FP32 vs FP32. */ + private fun assertActuallyPacked(weights: DecoderGgufWeights, label: String) { + val packed = weights.tensors.values.count { it.data is NarrowFloatTensorData } + assertEquals( + weights.tensors.size, packed, + "$label: expected every tensor to stay packed, only $packed of " + + "${weights.tensors.size} did — the parity check would be vacuous", + ) + } + + @Test + fun `FP16 KEEP_NATIVE forward matches the widened FP32 forward`() { + val file = writeModel( + buildHfWeights().map { (n, s, v) -> Triple(n, s, quantizeFp16(v)) }, + declaredDtype = "F16", + encode = ::encodeFp16, + ) + + val reference = forwardLogits(file, DTypePolicy.Any, tokens) + val native = forwardLogits(file, DTypePolicy.Require(FP16), tokens) { + assertActuallyPacked(it, "Require(FP16)") + } + + assertParity(reference, native, "FP16 KEEP_NATIVE") + } + + @Test + fun `BF16 KEEP_NATIVE forward matches the widened FP32 forward`() { + val file = writeModel( + buildHfWeights().map { (n, s, v) -> Triple(n, s, quantizeBf16(v)) }, + declaredDtype = "BF16", + encode = ::encodeBf16, + ) + + val reference = forwardLogits(file, DTypePolicy.Any, tokens) + val native = forwardLogits(file, DTypePolicy.Require(BF16), tokens) { + assertActuallyPacked(it, "Require(BF16)") + } + + assertParity(reference, native, "BF16 KEEP_NATIVE") + } + + @Test + fun `Prefer reaches the same forward result as Require`() { + // `Prefer` is the policy users are steered toward, because it degrades instead of + // throwing on a chain that can't keep the format. It must not be a different code path. + val file = writeModel( + buildHfWeights().map { (n, s, v) -> Triple(n, s, quantizeFp16(v)) }, + declaredDtype = "F16", + encode = ::encodeFp16, + ) + + val required = forwardLogits(file, DTypePolicy.Require(FP16), tokens) + val preferred = forwardLogits(file, DTypePolicy.Prefer(FP16), tokens) + + assertEquals(required.size, preferred.size) + for (i in required.indices) { + assertEquals( + required[i].toRawBits(), preferred[i].toRawBits(), + "Prefer(FP16) and Require(FP16) must be bit-identical at $i", + ) + } + } + + @Test + fun `narrow weights are materialized to FP32 before matmul`() { + // Documents the gap between "weights stay packed at rest" and "the narrow kernel runs", + // and explains why the parity tests above come out bit-identical instead of merely close. + // + // SafeTensors stores projections as [out, in]. `LlamaRuntime.linearProject` therefore + // calls `w.t()` before the matmul, and transpose has no narrow-float implementation — it + // decodes to a plain FP32 dense buffer. Meanwhile `DefaultCpuOpsJvm.chooseQuantizedMatmul` + // only engages when the weight is already [in, out]. So on this chain the packed data is + // widened on *every* forward, and the FP16/BF16 SGEMM kernels are never reached. + // + // The saving that survives is at-rest memory. The cost is a per-token decode plus a + // transpose allocation. Anyone wiring up the kernel for real has to remove the `.t()` — + // when they do, this test should be updated, and the parity tests will start exercising + // the kernel path they were written for. + val ctx = DirectCpuExecutionContext() + val outFeatures = ffDim + val inFeatures = dim + val values = quantizeFp16(randn(outFeatures * inFeatures, seed = 5)) + + @Suppress("UNCHECKED_CAST") + val packed = ctx.fromData( + Fp16DenseTensorData(Shape(outFeatures, inFeatures), encodeFp16(values)) + as TensorData, + FP32::class, + ) + assertTrue(packed.data is NarrowFloatTensorData, "precondition: weight starts packed") + + val transposed = packed.t() + assertTrue( + transposed.data !is NarrowFloatTensorData, + "transpose is expected to widen today; if this now stays packed, the narrow matmul " + + "kernel may finally be reachable — revisit the KDoc on this class", + ) + + // ...and the layout the fast path actually requires is the opposite one. + assertEquals( + inFeatures, transposed.shape[0], + "after t() the weight is [in, out] — the orientation chooseQuantizedMatmul wants, " + + "but by then it is no longer narrow", + ) + } + + @Test + fun `the parity check would catch a codec mix-up`() { + // Keeps [TOLERANCE] honest. Both formats are 2 bytes per element, so if the dispatch + // ever picked a kernel by byte width instead of by codec, nothing would throw. This + // measures what such a mix-up costs: write binary16 bit patterns into a file that + // *declares* BF16, so the loader hands those bytes to the bfloat16 decode. + val quantized = buildHfWeights().map { (n, s, v) -> Triple(n, s, quantizeFp16(v)) } + + val honest = writeModel(quantized, declaredDtype = "F16", encode = ::encodeFp16) + val mislabelled = writeModel(quantized, declaredDtype = "BF16", encode = ::encodeFp16) + + val correct = forwardLogits(honest, DTypePolicy.Require(FP16), tokens) + val wrong = forwardLogits(mislabelled, DTypePolicy.Require(BF16), tokens) + + val maxDelta = correct.indices.maxOf { abs(correct[it] - wrong[it]) } + assertTrue( + maxDelta > TOLERANCE * 10, + "reading F16 bytes as BF16 moved the logits by only $maxDelta — that is inside " + + "10x the parity tolerance ($TOLERANCE), so the parity tests above prove nothing", + ) + } +} From 10690ac6dc1e6f74f906c275a5e6cecf481267cc Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Mon, 27 Jul 2026 10:18:22 +0200 Subject: [PATCH 3/6] test(perf): benchmark narrow-float matmul against fp32 SGEMM MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Scoping the work to reach the FP16/BF16 kernels needed a number, not a guess: the layout change is only worth doing if the kernels beat the FP32 SGEMM they would replace. Measures three paths at real projection sizes — fp32 dense, narrow weight already in [in, out] so chooseQuantizedMatmul dispatches, and the [out, in] + .t() path production takes today. Skipped unless -Dskainet.bench.narrow=true, so it stays out of CI. The build passes that property through because Gradle does not forward -D to the test JVM. Dispatch is guaranteed by construction rather than assumed: checkDispatchable asserts chooseQuantizedMatmul's own preconditions before timing, and each narrow result is compared against the fp32 baseline so a fast-but-wrong kernel cannot pass unnoticed. Baseline on i7-9750H / OpenJDK 21 recorded in the KDoc. Three results: - BF16 beats FP32 by 1.5-2.1x at every size and batch. Batch-1 matmul is memory-bandwidth bound, so halving the weight bytes roughly halves the time. This is the case for doing the layout work. - FP16 is 2-18x slower, pinned near 0.5 GFLOP/s regardless of shape or batch. Not a layout problem: both Panama kernels fill a scratch array scalar-wise before the vector FMA, but Fp16Codec.decode is a branchy when with a subnormal renormalization loop where the BF16 decode is three integer ops. - The transpose path costs 0.2-4.5 seconds per projection, because the generic transpose walks narrow data element by element through get(). That is per weight, per token. At 8B sizes KEEP_NATIVE is not merely un-accelerated today, it is unusably slow. The earlier forward-parity test could not surface the last point because its model is dim=8, where an elementwise transpose is free. --- llm-inference/llama/build.gradle.kts | 3 + .../llama/NarrowFloatMatmulBenchmark.kt | 259 ++++++++++++++++++ 2 files changed, 262 insertions(+) create mode 100644 llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/NarrowFloatMatmulBenchmark.kt diff --git a/llm-inference/llama/build.gradle.kts b/llm-inference/llama/build.gradle.kts index ac2df26..31dde66 100644 --- a/llm-inference/llama/build.gradle.kts +++ b/llm-inference/llama/build.gradle.kts @@ -81,4 +81,7 @@ kotlin { tasks.withType().configureEach { jvmArgs("--enable-preview", "--add-modules", "jdk.incubator.vector", "-XX:MaxDirectMemorySize=12g") maxHeapSize = "6g" + // Opt-in gate for NarrowFloatMatmulBenchmark, which is a measurement rather than a test and + // stays skipped unless explicitly requested. Gradle does not forward -D to the test JVM. + System.getProperty("skainet.bench.narrow")?.let { systemProperty("skainet.bench.narrow", it) } } diff --git a/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/NarrowFloatMatmulBenchmark.kt b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/NarrowFloatMatmulBenchmark.kt new file mode 100644 index 0000000..c39b838 --- /dev/null +++ b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/NarrowFloatMatmulBenchmark.kt @@ -0,0 +1,259 @@ +package sk.ainet.models.llama + +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.Bf16DenseTensorData +import sk.ainet.lang.tensor.data.DenseFloatArrayTensorData +import sk.ainet.lang.tensor.data.Fp16DenseTensorData +import sk.ainet.lang.tensor.data.NarrowFloatTensorData +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.tensor.matmul +import sk.ainet.lang.tensor.t +import sk.ainet.lang.types.Bf16Codec +import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.Fp16Codec +import sk.ainet.lang.types.NarrowFloatCodec +import kotlin.math.abs +import kotlin.test.Test +import kotlin.test.assertTrue + +/** + * Decides whether reaching the narrow-float matmul kernel is worth the layout work. + * + * Skipped unless `-Dskainet.bench.narrow=true`. This is a measurement, not a test — it asserts + * only the things that would invalidate its own numbers. + * + * ### What it answers + * + * `DecoderNarrowFloatForwardParityTest` established that KEEP_NATIVE is numerically correct but + * that the FP16/BF16 SGEMM kernels are never reached: both `Linear.onForward` and + * `LlamaRuntime.linearProject` call `w.t()`, transpose has no narrow-float arm and widens to a + * dense FP32 buffer, and `DefaultCpuOpsJvm.chooseQuantizedMatmul` only engages for `[in, out]` + * weights. Wiring the kernel up means a byte relayout at load plus a lazy-transpose arm in the + * engine — the pattern the K-quants already use. That is only worth doing if the kernel actually + * wins, so this measures three things at realistic projection sizes: + * + * - **fp32** — dense FP32 SGEMM. The baseline, and what runs today after the widening. + * - **fp16/bf16** — weight handed over already `[in, out]`, so `chooseQuantizedMatmul` dispatches + * to the narrow kernel. This is the best case the layout work could unlock. + * - **transpose** — narrow weight in the real `[out, in]` orientation, `.t()` then matmul. This + * is what production does per token today, and shows what the widening costs. + * + * Dispatch is guaranteed by construction rather than observed: `chooseQuantizedMatmul` requires an + * FP32 rank-2 input, a rank-2 weight, and `weight.shape[0] == input.shape[1]`. [checkDispatchable] + * asserts exactly that before timing, so a silent fallback to the generic path cannot be mistaken + * for a fast kernel. + * + * ### Baseline, 2026-07-27 + * + * Intel i7-9750H (AVX2, no AVX-512), 12 threads, OpenJDK 21.0.11, engine 0.38.0-SNAPSHOT. + * Median ms per call: + * + * ``` + * shape batch fp32 fp16 bf16 transpose + * q_proj 1B 1 3.210 18.601 1.591 206.458 + * q_proj 1B 16 16.449 297.873 11.125 215.698 + * q_proj 8B 1 40.335 73.793 21.472 1318.772 + * q_proj 8B 16 95.954 1166.224 66.227 1358.200 + * ffn_up 8B 1 111.922 197.337 60.372 4367.357 + * ffn_up 8B 16 272.889 3146.712 186.360 4511.455 + * ffn_down 8B 1 109.242 198.977 57.749 2195.538 + * ffn_down 8B 16 274.547 3172.639 177.142 2363.605 + * ``` + * + * Three conclusions: + * + * 1. **BF16 beats FP32 by 1.5–2.1x everywhere.** At batch 1 the matmul is memory-bandwidth bound, + * so halving the weight bytes roughly halves the time. This is the case for doing the layout + * work. + * 2. **FP16 is 2–18x slower, pinned at ~0.5 GFLOP/s regardless of shape or batch** — the signature + * of being compute-bound on the decode. Both Panama kernels fill a scratch lane array scalar-wise + * before the vector FMA, but BF16's decode is three integer ops while `Fp16Codec.decode` is a + * branchy `when` with a subnormal renormalization loop. The fix is engine-side and independent + * of layout: use `Float.float16ToFloat` (a JDK 20+ intrinsic) or a branch-free decode. + * 3. **The `transpose` column is the alarming one.** 0.2–4.5 *seconds* for one projection, because + * the generic transpose walks a narrow tensor element by element through `get()`. That is the + * path production takes today, per weight, per token. KEEP_NATIVE is not merely un-accelerated + * right now — at real model sizes it is unusably slow. + */ +class NarrowFloatMatmulBenchmark { + + private val enabled = System.getProperty("skainet.bench.narrow") == "true" + + /** `[inFeatures, outFeatures]` taken from real LLaMA projections, plus a small control. */ + private val shapes = listOf( + Triple("q_proj 1B", 2048, 2048), + Triple("q_proj 8B", 4096, 4096), + Triple("ffn_up 8B", 4096, 11008), + Triple("ffn_down 8B", 11008, 4096), + ) + + /** Decode is batch 1; the larger batch stands in for prefill, where SGEMM has more to work with. */ + private val batches = listOf(1, 16) + + private val minSamples = 5 + private val timeBudgetNanos = 1_500_000_000L + private val warmupNanos = 500_000_000L + + private fun encode(values: FloatArray, codec: NarrowFloatCodec): ByteArray { + val out = ByteArray(values.size * 2) + for (i in values.indices) { + val bits = codec.encode(values[i]) + out[i * 2] = (bits and 0xFF).toByte() + out[i * 2 + 1] = ((bits ushr 8) and 0xFF).toByte() + } + return out + } + + private fun weights(n: Int, seed: Int): FloatArray { + val rng = kotlin.random.Random(seed) + return FloatArray(n) { (rng.nextFloat() - 0.5f) * 0.1f } + } + + /** + * Fails rather than silently reporting generic-path timings as kernel timings. + * Mirrors the preconditions in `DefaultCpuOpsJvm.chooseQuantizedMatmul`. + */ + private fun checkDispatchable(x: Tensor, w: Tensor, label: String) { + assertTrue(w.data is NarrowFloatTensorData, "$label: weight is not narrow — nothing to dispatch") + assertTrue(x.shape.rank == 2 && w.shape.rank == 2, "$label: both operands must be rank 2") + assertTrue( + w.shape[0] == x.shape[1], + "$label: weight is [${w.shape[0]}, ${w.shape[1]}] but input has ${x.shape[1]} columns — " + + "chooseQuantizedMatmul would return null and the generic path would be timed instead", + ) + } + + /** Median nanoseconds per call, after a fixed warmup window. */ + private fun measure(body: () -> Tensor): Long { + var sink = 0.0f + val warmupEnd = System.nanoTime() + warmupNanos + while (System.nanoTime() < warmupEnd) { + sink += body().data.copyToFloatArray()[0] + } + + val samples = mutableListOf() + val deadline = System.nanoTime() + timeBudgetNanos + while (samples.size < minSamples || System.nanoTime() < deadline) { + val t0 = System.nanoTime() + val r = body() + val elapsed = System.nanoTime() - t0 + sink += r.data.copyToFloatArray()[0] + samples.add(elapsed) + if (samples.size >= 2000) break + } + check(!sink.isNaN()) { "sink went NaN — results were not consumed" } + samples.sort() + return samples[samples.size / 2] + } + + private fun gflops(batch: Int, inF: Int, outF: Int, nanos: Long): Double = + (2.0 * batch * inF * outF) / nanos + + @Test + fun `narrow float matmul throughput versus fp32`() { + if (!enabled) { + println("NarrowFloatMatmulBenchmark skipped — rerun with -Dskainet.bench.narrow=true") + return + } + + val ctx = DirectCpuExecutionContext() + println() + println("narrow-float matmul vs fp32 SGEMM (median of timed samples)") + println("weight layout [in, out]; 'transpose' is the [out, in] + .t() path production uses today") + println() + println( + "%-13s %6s %11s %11s %11s %11s %s".format( + "shape", "batch", "fp32", "fp16", "bf16", "transpose", "verdict", + ), + ) + println("-".repeat(96)) + + for ((label, inF, outF) in shapes) { + val raw = weights(inF * outF, seed = inF + outF) + val fp16Bytes = encode(raw, Fp16Codec) + val bf16Bytes = encode(raw, Bf16Codec) + // Decode back so the FP32 baseline holds the same values the narrow paths do — + // otherwise the comparison is between different matrices. + val fp32Values = FloatArray(raw.size) { + Fp16Codec.decode(Fp16Codec.encode(raw[it])) + } + + @Suppress("UNCHECKED_CAST") + val wFp32 = ctx.fromData( + DenseFloatArrayTensorData(Shape(inF, outF), fp32Values) as TensorData, + FP32::class, + ) + @Suppress("UNCHECKED_CAST") + val wFp16 = ctx.fromData( + Fp16DenseTensorData(Shape(inF, outF), fp16Bytes) as TensorData, + FP32::class, + ) + @Suppress("UNCHECKED_CAST") + val wBf16 = ctx.fromData( + Bf16DenseTensorData(Shape(inF, outF), bf16Bytes) as TensorData, + FP32::class, + ) + // The production orientation: [out, in], transposed on every call. + @Suppress("UNCHECKED_CAST") + val wFp16Transposed = ctx.fromData( + Fp16DenseTensorData(Shape(outF, inF), fp16Bytes) as TensorData, + FP32::class, + ) + + for (batch in batches) { + val x = ctx.fromFloatArray( + Shape(batch, inF), FP32::class, weights(batch * inF, seed = 99), + ) + + checkDispatchable(x, wFp16, "$label fp16") + checkDispatchable(x, wBf16, "$label bf16") + + val fp32Ns = measure { x.matmul(wFp32) } + val fp16Ns = measure { x.matmul(wFp16) } + val bf16Ns = measure { x.matmul(wBf16) } + val transposeNs = measure { x.matmul(wFp16Transposed.t()) } + + // Sanity: the narrow kernel must agree with the FP32 baseline, or the timing is + // measuring something that isn't a correct matmul. + val ref = x.matmul(wFp32).data.copyToFloatArray() + val got = x.matmul(wFp16).data.copyToFloatArray() + val tol = 1e-2f * (1 + inF / 1024) + val maxDelta = ref.indices.maxOf { abs(ref[it] - got[it]) } + assertTrue( + maxDelta < tol, + "$label batch=$batch: fp16 kernel disagrees with fp32 by $maxDelta (tol $tol)", + ) + + val speedup = fp32Ns.toDouble() / fp16Ns.toDouble() + val verdict = when { + speedup >= 1.15 -> "fp16 %.2fx faster".format(speedup) + speedup <= 0.87 -> "fp16 %.2fx SLOWER".format(1 / speedup) + else -> "no real difference" + } + + println( + "%-13s %6d %9.3fms %9.3fms %9.3fms %9.3fms %s".format( + label, batch, + fp32Ns / 1e6, fp16Ns / 1e6, bf16Ns / 1e6, transposeNs / 1e6, + verdict, + ), + ) + println( + "%-13s %6s %9.1fGF %9.1fGF %9.1fGF %9s".format( + "", "", + gflops(batch, inF, outF, fp32Ns), + gflops(batch, inF, outF, fp16Ns), + gflops(batch, inF, outF, bf16Ns), + "", + ), + ) + } + } + + println() + println("weight bytes at rest: fp32 = 2x narrow. The 'transpose' column is the per-call cost") + println("of the widening that happens today, and is what the layout work would remove.") + } +} From 18bd6001c690edf9b9793f84e618b6a8194094bc Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Mon, 27 Jul 2026 20:37:31 +0200 Subject: [PATCH 4/6] perf(dtype): relay narrow-float matmul weights input-major at load MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Weights arrive [out, in] but the narrow matmul dispatch needs [in, out], so Linear.onForward transposes before every matmul. A row-major narrow tensor has no fast transpose — it widened elementwise through boxed get(), measured at 209 ms for a 2048x2048 projection and 4.5 s for 4096x11008, per weight per token. KEEP_NATIVE was slower than not using it at all. Relay matmul weights once at load via the engine's new NarrowFloatInputMajorTensorData (SKaiNET #888), which makes the per-forward transpose a zero-copy view, so the packed weight reaches the narrow kernel. Two kinds of tensor stay row-major. Rank-1 norms are never transposed and the input-major type rejects them outright. The token embedding is gathered by row, not multiplied, so input-major storage would stride exactly the reads it serves. That also covers tied embeddings: when output.weight aliases token_embd it stays row-major and forgoes the transpose win, since one buffer cannot suit both access patterns. Measured on i7-9750H / OpenJDK 21, ffn_up 8B at batch 1, BF16: 4465 ms -> 58 ms, a 77x reduction, and now 1.9x faster than the FP32 SGEMM it replaces. The benchmark gains t()row-maj and t()in-maj columns so the difference stays visible, plus an assertion that the relaid weight is still narrow after t() — without it the new columns could silently re-measure the generic path. The forward-parity logits are no longer bit-identical to the widened path: they diverge by ~7e-9, which is the accumulation-order difference of a genuinely different kernel, and is the evidence the kernel now runs. Tolerance tightened from 1e-4 to 1e-5 accordingly — still three orders of magnitude above the observed delta and three below what a codec mix-up produces. Loader tests updated for the layout split: matmul weights assert input-major and value preservation rather than verbatim bytes, and the verbatim-bytes property is now pinned on the embedding, where it still holds. FP16 remains slower than FP32 here; that is the decode cost tracked in SKaiNET #887 and is unaffected by layout. Prefer BF16 for speed until it lands. --- .../models/llama/DecoderSafeTensorsLoader.kt | 52 +++++- .../DecoderNarrowFloatForwardParityTest.kt | 148 +++++++++++------- ...DecoderSafeTensorsLoaderNarrowFloatTest.kt | 77 ++++++--- .../llama/NarrowFloatMatmulBenchmark.kt | 141 ++++++++++------- 4 files changed, 276 insertions(+), 142 deletions(-) diff --git a/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/DecoderSafeTensorsLoader.kt b/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/DecoderSafeTensorsLoader.kt index 4c62744..aa8069f 100644 --- a/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/DecoderSafeTensorsLoader.kt +++ b/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/DecoderSafeTensorsLoader.kt @@ -13,10 +13,14 @@ import sk.ainet.io.safetensors.StreamingSafeTensorInfo import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.Tensor import sk.ainet.apps.llm.DTypePolicyValidation -import sk.ainet.lang.tensor.data.Bf16DenseTensorData -import sk.ainet.lang.tensor.data.Fp16DenseTensorData +import sk.ainet.lang.tensor.data.NarrowFloatDenseTensorData +import sk.ainet.lang.tensor.data.NarrowFloatInputMajorTensorData +import sk.ainet.lang.tensor.data.NarrowFloatTensorData import sk.ainet.lang.tensor.data.TensorData import sk.ainet.lang.types.BF16 +import sk.ainet.lang.types.Bf16Codec +import sk.ainet.lang.types.Fp16Codec +import sk.ainet.lang.types.NarrowFloatCodec import sk.ainet.lang.types.DType import sk.ainet.lang.types.DTypePolicy import sk.ainet.lang.types.FP16 @@ -33,9 +37,9 @@ import kotlin.reflect.KClass * - Q4 + .qb companion tensor dequantization to FP32 * - BF16/F16 dequantization to FP32 (default) * - **Narrow-float KEEP_NATIVE** when [dtypePolicy] admits BF16 or F16 - * (SKaiNET 0.38.0): constructs a [Bf16DenseTensorData] / [Fp16DenseTensorData]-backed - * tensor so the narrow-float matmul kernel routes via `DefaultCpuOpsJvm` - * without a 2× memory blow-up. + * (SKaiNET 0.38.0): keeps the on-disk 2-bytes-per-element buffer so the narrow-float + * matmul kernel can run without a 2× memory blow-up. Matmul weights are relaid + * input-major so the per-forward transpose is free — see [narrowData]. * - Shape normalization ([1, dim] norms → [dim]) * - Tied word embeddings (output.weight = token_embd.weight) * @@ -111,7 +115,7 @@ public class DecoderSafeTensorsLoader( // physical encoding — the get/set surface still // returns Float. Mirrors the // `GemmaMemSegConverter` pattern for Q4/Q8. - val data = Bf16DenseTensorData.fromRawBytes(targetShape, bytes) + val data = narrowData(canonicalName, targetShape, bytes, Bf16Codec) @Suppress("UNCHECKED_CAST") ctx.fromData(data as TensorData, dtype) as Tensor } else { @@ -129,7 +133,7 @@ public class DecoderSafeTensorsLoader( // to the BF16 decode would not throw — it would produce // plausible-looking wrong numbers. The codec carried by // `Fp16DenseTensorData` is what keeps the dispatch honest. - val data = Fp16DenseTensorData.fromRawBytes(targetShape, bytes) + val data = narrowData(canonicalName, targetShape, bytes, Fp16Codec) @Suppress("UNCHECKED_CAST") ctx.fromData(data as TensorData, dtype) as Tensor } else { @@ -177,6 +181,40 @@ public class DecoderSafeTensorsLoader( return LlamaWeightMapper.map(loadToMap(randomAccessProvider)) } + /** + * Build the KEEP_NATIVE storage for one narrow-float tensor, choosing its byte layout. + * + * Matmul weights are relaid **input-major** ([NarrowFloatInputMajorTensorData]). Weights are + * stored `[out, in]` but the narrow matmul dispatch needs `[in, out]`, so `Linear.onForward` + * transposes on every forward pass. A row-major narrow tensor has no fast transpose — it + * widens elementwise through boxed `get()`, which measured 206 ms for a 2048×2048 projection + * and 4.4 s for 4096×11008, *per weight per token*. Relaying once at load makes that transpose + * a zero-copy view (engine issue #888), which is what lets the narrow kernel actually run. + * + * Two kinds of tensor stay row-major: + * + * - **Rank-1 norms.** Never transposed, never matmul'd — relaying them is undefined and + * [NarrowFloatInputMajorTensorData] rejects rank ≠ 2 outright. + * - **The token embedding.** Gathered by row, not multiplied. Input-major storage strides + * those row reads, so relaying it would trade a win we don't get for a loss we would. + * Note this also covers tied embeddings: `output.weight` aliases `token_embd`, so in the + * tied case the output projection stays row-major too and forgoes the transpose win. That + * is deliberate — one shared buffer cannot be optimal for both access patterns. + */ + private fun narrowData( + canonicalName: String, + shape: Shape, + bytes: ByteArray, + codec: NarrowFloatCodec, + ): NarrowFloatTensorData { + val isGatheredEmbedding = canonicalName == LlamaTensorNames.TOKEN_EMBEDDINGS + return if (shape.rank == 2 && !isGatheredEmbedding) { + NarrowFloatInputMajorTensorData.fromRowMajor(shape, bytes, codec) + } else { + NarrowFloatDenseTensorData(shape, bytes, codec) + } + } + /** * Infer the target shape for a tensor, normalizing norm shapes. * For Q4 tensors, the target shape is the logical shape (not the packed shape). diff --git a/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/DecoderNarrowFloatForwardParityTest.kt b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/DecoderNarrowFloatForwardParityTest.kt index a7ed7f3..b39e2b8 100644 --- a/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/DecoderNarrowFloatForwardParityTest.kt +++ b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/DecoderNarrowFloatForwardParityTest.kt @@ -5,10 +5,8 @@ import sk.ainet.apps.llm.OptimizedLLMRuntime import sk.ainet.context.DirectCpuExecutionContext import sk.ainet.io.JvmRandomAccessSource import sk.ainet.io.RandomAccessSource -import sk.ainet.lang.tensor.Shape -import sk.ainet.lang.tensor.data.Fp16DenseTensorData +import sk.ainet.lang.tensor.data.NarrowFloatInputMajorTensorData import sk.ainet.lang.tensor.data.NarrowFloatTensorData -import sk.ainet.lang.tensor.data.TensorData import sk.ainet.lang.tensor.t import sk.ainet.lang.types.BF16 import sk.ainet.lang.types.DTypePolicy @@ -22,6 +20,7 @@ import java.nio.file.Files import kotlin.math.abs import kotlin.test.Test import kotlin.test.assertEquals +import kotlin.test.assertSame import kotlin.test.assertTrue /** @@ -40,15 +39,14 @@ import kotlin.test.assertTrue * values already round-tripped through the codec, so the widened path and the packed path decode * the same numbers. [TOLERANCE] covers float accumulation-order differences. * - * **What this does NOT prove.** It does not show that the narrow-float matmul kernel runs. On this - * chain it does not — see [`narrow weights are materialized to FP32 before matmul`], which pins - * why. The logits below are bit-identical rather than merely close, because both runs ultimately - * execute the same FP32 SGEMM; only the moment of decoding differs. [TOLERANCE] stays non-zero so - * the test keeps passing if a real narrow kernel is ever wired in and shifts accumulation order. - * - * The KEEP_NATIVE arm applies to *every* tensor in the file, so this does cover embedding gather + * The KEEP_NATIVE arm applies to *every* tensor in the file, so this also covers embedding gather * and the RMSNorm weight multiply reading through packed storage. * + * Since engine #888 the matmul weights are relaid input-major at load, so the per-forward + * transpose is a zero-copy view and the narrow kernel genuinely runs — see + * [`matmul weights survive the per-forward transpose still packed`] for that property, and + * [`gathered and rank-1 tensors stay row-major`] for the tensors deliberately left alone. + * * [`the parity check would catch a codec mix-up`] is the guard that keeps the tolerance honest. * * Files are synthesized in-test; no model downloads are involved. @@ -64,12 +62,17 @@ class DecoderNarrowFloatForwardParityTest { private val seqLen = 32 /** - * Absolute logit tolerance. Today both paths converge on the same FP32 SGEMM and agree to the - * bit, so this is slack for a future narrow kernel whose accumulation order would differ. - * Sized to swallow that and nothing more — the codec mix-up guard measures the margin actually - * available and fails if this ever grows large enough to hide a real defect. + * Absolute logit tolerance. The two paths now run genuinely different kernels — the narrow + * SGEMM against the FP32 one — so they no longer agree bit-for-bit; the measured divergence + * is ~7e-9, a couple of ULPs of accumulation-order difference. This sits three orders of + * magnitude above that for headroom on other vector widths, and three below the ~1e-2 a codec + * mix-up produces. The mix-up guard asserts that margin rather than assuming it. + * + * Before the input-major relayout landed, these paths *were* bit-identical, because the + * per-forward transpose widened the weight and both ended up in the same FP32 SGEMM. A return + * to exact equality here would mean the narrow kernel has stopped running. */ - private val TOLERANCE = 1e-4f + private val TOLERANCE = 1e-5f private val metadata = LlamaModelMetadata( architecture = "llama", @@ -181,14 +184,10 @@ class DecoderNarrowFloatForwardParityTest { // -------------------------------------------------------------- the runs - /** Load under [policy] and forward [tokens] in order, returning the logits of the last step. */ - private fun forwardLogits( - file: File, - policy: DTypePolicy, - tokens: IntArray, - onWeights: (DecoderGgufWeights) -> Unit = {}, - ): FloatArray { - val ctx = DirectCpuExecutionContext() + private val ctx = DirectCpuExecutionContext() + + /** Load the synthesized model under [policy]. */ + private fun loadWeights(file: File, policy: DTypePolicy): DecoderGgufWeights { val loader = DecoderSafeTensorsLoader( ctx = ctx, dtype = FP32::class, @@ -197,7 +196,17 @@ class DecoderNarrowFloatForwardParityTest { dtypePolicy = policy, ) val provider: () -> RandomAccessSource = { JvmRandomAccessSource.open(file) } - val weights = loader.loadToMap(provider) + return loader.loadToMap(provider) + } + + /** Load under [policy] and forward [tokens] in order, returning the logits of the last step. */ + private fun forwardLogits( + file: File, + policy: DTypePolicy, + tokens: IntArray, + onWeights: (DecoderGgufWeights) -> Unit = {}, + ): FloatArray { + val weights = loadWeights(file, policy) onWeights(weights) val runtime = OptimizedLLMRuntime( @@ -294,45 +303,70 @@ class DecoderNarrowFloatForwardParityTest { } @Test - fun `narrow weights are materialized to FP32 before matmul`() { - // Documents the gap between "weights stay packed at rest" and "the narrow kernel runs", - // and explains why the parity tests above come out bit-identical instead of merely close. - // - // SafeTensors stores projections as [out, in]. `LlamaRuntime.linearProject` therefore - // calls `w.t()` before the matmul, and transpose has no narrow-float implementation — it - // decodes to a plain FP32 dense buffer. Meanwhile `DefaultCpuOpsJvm.chooseQuantizedMatmul` - // only engages when the weight is already [in, out]. So on this chain the packed data is - // widened on *every* forward, and the FP16/BF16 SGEMM kernels are never reached. - // - // The saving that survives is at-rest memory. The cost is a per-token decode plus a - // transpose allocation. Anyone wiring up the kernel for real has to remove the `.t()` — - // when they do, this test should be updated, and the parity tests will start exercising - // the kernel path they were written for. - val ctx = DirectCpuExecutionContext() - val outFeatures = ffDim - val inFeatures = dim - val values = quantizeFp16(randn(outFeatures * inFeatures, seed = 5)) - - @Suppress("UNCHECKED_CAST") - val packed = ctx.fromData( - Fp16DenseTensorData(Shape(outFeatures, inFeatures), encodeFp16(values)) - as TensorData, - FP32::class, + fun `matmul weights survive the per-forward transpose still packed`() { + // The property the whole feature rests on. Weights arrive [out, in]; `Linear.onForward` + // transposes before every matmul. Before engine #888 that transpose widened the tensor + // elementwise, so the narrow kernel was unreachable and KEEP_NATIVE was slower than not + // using it. The loader now relays matmul weights input-major, which makes the transpose a + // zero-copy view — this asserts the weight is still narrow on the far side of it. + val file = writeModel( + buildHfWeights().map { (n, s, v) -> Triple(n, s, quantizeFp16(v)) }, + declaredDtype = "F16", + encode = ::encodeFp16, ) - assertTrue(packed.data is NarrowFloatTensorData, "precondition: weight starts packed") + val weights = loadWeights(file, DTypePolicy.Require(FP16)) - val transposed = packed.t() + val ffnGate = weights.tensors[LlamaTensorNames.ffnGate(0)] + ?: error("missing ffn_gate") assertTrue( - transposed.data !is NarrowFloatTensorData, - "transpose is expected to widen today; if this now stays packed, the narrow matmul " + - "kernel may finally be reachable — revisit the KDoc on this class", + ffnGate.data is NarrowFloatInputMajorTensorData, + "a matmul weight must be relaid input-major, got ${ffnGate.data::class.simpleName}", ) - // ...and the layout the fast path actually requires is the opposite one. + val transposed = ffnGate.t() + assertTrue( + transposed.data is NarrowFloatTensorData, + "transpose widened the weight — the narrow kernel is unreachable again", + ) assertEquals( - inFeatures, transposed.shape[0], - "after t() the weight is [in, out] — the orientation chooseQuantizedMatmul wants, " + - "but by then it is no longer narrow", + dim, transposed.shape[0], + "after t() the weight must be [in, out], the orientation chooseQuantizedMatmul wants", + ) + assertSame( + (ffnGate.data as NarrowFloatTensorData).packedData, + (transposed.data as NarrowFloatTensorData).packedData, + "the transpose must not copy — a copy per forward is the cost being removed", + ) + } + + @Test + fun `gathered and rank-1 tensors stay row-major`() { + // The counterpart to the test above, and the one that would catch over-applying the + // relayout. The token embedding is gathered by row, so input-major storage would stride + // exactly the reads it serves; norms are rank-1 and never transposed at all. + val file = writeModel( + buildHfWeights().map { (n, s, v) -> Triple(n, s, quantizeFp16(v)) }, + declaredDtype = "F16", + encode = ::encodeFp16, + ) + val weights = loadWeights(file, DTypePolicy.Require(FP16)) + + val embedding = weights.tensors[LlamaTensorNames.TOKEN_EMBEDDINGS] + ?: error("missing token_embd") + assertTrue( + embedding.data is NarrowFloatTensorData, + "the embedding should still be packed — only its layout differs", + ) + assertTrue( + embedding.data !is NarrowFloatInputMajorTensorData, + "the gathered embedding must stay row-major", + ) + + val norm = weights.tensors[LlamaTensorNames.attnNorm(0)] ?: error("missing attn_norm") + assertEquals(1, norm.shape.rank, "precondition: norms are rank-1") + assertTrue( + norm.data !is NarrowFloatInputMajorTensorData, + "a rank-1 norm must never be relaid — the input-major type rejects rank != 2", ) } diff --git a/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/DecoderSafeTensorsLoaderNarrowFloatTest.kt b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/DecoderSafeTensorsLoaderNarrowFloatTest.kt index 9fab628..d4fd94c 100644 --- a/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/DecoderSafeTensorsLoaderNarrowFloatTest.kt +++ b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/DecoderSafeTensorsLoaderNarrowFloatTest.kt @@ -4,14 +4,14 @@ import sk.ainet.context.DirectCpuExecutionContext import sk.ainet.io.JvmRandomAccessSource import sk.ainet.io.RandomAccessSource import sk.ainet.lang.tensor.Tensor -import sk.ainet.lang.tensor.data.Bf16TensorData import sk.ainet.lang.tensor.data.FloatArrayTensorData -import sk.ainet.lang.tensor.data.Fp16DenseTensorData +import sk.ainet.lang.tensor.data.NarrowFloatInputMajorTensorData import sk.ainet.lang.tensor.data.NarrowFloatTensorData import sk.ainet.lang.types.BF16 import sk.ainet.lang.types.DTypePolicy import sk.ainet.lang.types.FP16 import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.Bf16Codec import sk.ainet.lang.types.Fp16Codec import java.io.File import java.nio.ByteBuffer @@ -26,12 +26,15 @@ import kotlin.test.assertTrue * Pins the narrow-float KEEP_NATIVE behaviour of [DecoderSafeTensorsLoader] — the * transformer-repo counterpart of the engine's `SafeTensorsParametersLoaderFp16PolicyTest`. * - * The BF16 arm has existed since 0.25.0; the F16 arm landed with engine 0.38.0's - * `Fp16DenseTensorData`. What matters here is that the loader resolves the two formats - * **independently**: both are 2 bytes per element, so routing F16 bytes through the BF16 decode - * would not throw — it would quietly produce wrong numbers. These tests assert the packed bytes - * survive verbatim, decode bit-identically to the widening path, and that neither policy leaks - * into the other format. + * The BF16 arm has existed since 0.25.0; the F16 arm landed with engine 0.38.0. What matters here + * is that the loader resolves the two formats **independently**: both are 2 bytes per element, so + * routing F16 bytes through the BF16 decode would not throw — it would quietly produce wrong + * numbers. These tests assert the tensors stay packed, decode bit-identically to the widening + * path, and that neither policy leaks into the other format. + * + * They also pin the layout split introduced with engine #888: matmul weights are relaid + * input-major so their per-forward transpose is free, while gathered tensors — the token + * embedding — stay row-major and keep their on-disk bytes verbatim. * * Files are synthesized in-test; no model downloads are involved. */ @@ -130,27 +133,53 @@ class DecoderSafeTensorsLoaderNarrowFloatTest { } @Test - fun `Require(FP16) keeps the on-disk F16 bytes verbatim`() { + fun `Require(FP16) keeps a matmul weight packed, relaid input-major`() { val onDisk = fp32ToFp16Bytes(values) val file = writeSafeTensors(listOf(Triple(hfName, "F16", onDisk)), 2, 4) val weight = load(file, DTypePolicy.Require(FP16))[canonical] ?: error("missing $canonical") + val data = weight.data + assertTrue(data is NarrowFloatTensorData, "must be recognizable to narrow dispatch") assertTrue( - weight.data is Fp16DenseTensorData, - "KEEP_NATIVE must produce Fp16DenseTensorData, got ${weight.data::class.simpleName}", + data is NarrowFloatInputMajorTensorData, + "a matmul weight must be relaid input-major so its per-forward transpose is free " + + "(engine #888), got ${data::class.simpleName}", + ) + assertEquals( + Fp16Codec, (data as NarrowFloatTensorData).codec, + "an F16 tensor must never be decoded as BF16 — the bit layouts differ", + ) + // No widening pass ran: still 2 bytes per element, and the values round-trip. The bytes + // are permuted rather than verbatim now, so the value check is what proves preservation. + assertEquals(values.size * 2, data.packedData.size, "must stay 2 bytes per element") + assertContentEquals( + values, data.copyToFloatArray(), + "the relayout must permute bytes without changing what the tensor holds", + ) + } + + @Test + fun `the gathered embedding keeps its on-disk bytes verbatim`() { + // The embedding is read by row, so it is deliberately left row-major — which means it is + // also the tensor where byte-for-byte identity with the file still holds. + val onDisk = fp32ToFp16Bytes(values) + val file = writeSafeTensors( + listOf(Triple("model.embed_tokens.weight", "F16", onDisk)), 2, 4, ) - assertTrue(weight.data is NarrowFloatTensorData, "must be recognizable to narrow dispatch") + + val embedding = load(file, DTypePolicy.Require(FP16))[LlamaTensorNames.TOKEN_EMBEDDINGS] + ?: error("missing token_embd") + + assertTrue(embedding.data is NarrowFloatTensorData, "must stay packed") assertTrue( - weight.data !is Bf16TensorData, - "an F16 tensor must never be mistaken for BF16 — the bit layouts differ", + embedding.data !is NarrowFloatInputMajorTensorData, + "a gathered tensor must not be relaid — input-major storage strides its row reads", ) - // Byte-for-byte identity proves no widening pass ran. assertContentEquals( - onDisk, (weight.data as Fp16DenseTensorData).packedData, - "KEEP_NATIVE must preserve on-disk F16 bytes verbatim", + onDisk, (embedding.data as NarrowFloatTensorData).packedData, + "the row-major path must preserve on-disk bytes verbatim", ) - assertEquals(values.size * 2, (weight.data as Fp16DenseTensorData).packedData.size) } @Test @@ -188,22 +217,26 @@ class DecoderSafeTensorsLoaderNarrowFloatTest { // Require(FP16): F16 stays packed, BF16 widens — it cannot be re-encoded as F16. val a = load(file, DTypePolicy.Require(FP16)) - assertTrue(a[f16Canonical]!!.data is Fp16DenseTensorData, "F16 should be packed") + assertEquals( + Fp16Codec, (a[f16Canonical]!!.data as NarrowFloatTensorData).codec, "F16 should be packed", + ) assertTrue(a[bf16Canonical]!!.data is FloatArrayTensorData<*>, "BF16 should be widened") // ...and the mirror image. val b = load(file, DTypePolicy.Require(BF16)) assertTrue(b[f16Canonical]!!.data is FloatArrayTensorData<*>, "F16 should be widened") - assertTrue(b[bf16Canonical]!!.data is Bf16TensorData, "BF16 should be packed") + assertEquals( + Bf16Codec, (b[bf16Canonical]!!.data as NarrowFloatTensorData).codec, "BF16 should be packed", + ) } @Test fun `Prefer and OneOf reach the same KEEP_NATIVE path as Require`() { val file = writeSafeTensors(listOf(Triple(hfName, "F16", fp32ToFp16Bytes(values))), 2, 4) - assertTrue(load(file, DTypePolicy.Prefer(FP16))[canonical]!!.data is Fp16DenseTensorData) + assertTrue(load(file, DTypePolicy.Prefer(FP16))[canonical]!!.data is NarrowFloatTensorData) assertTrue( - load(file, DTypePolicy.OneOf(setOf(FP32, FP16)))[canonical]!!.data is Fp16DenseTensorData, + load(file, DTypePolicy.OneOf(setOf(FP32, FP16)))[canonical]!!.data is NarrowFloatTensorData, ) // A soft policy naming neither narrow format leaves the widening default in place. assertTrue(load(file, DTypePolicy.Prefer(FP32))[canonical]!!.data is FloatArrayTensorData<*>) diff --git a/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/NarrowFloatMatmulBenchmark.kt b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/NarrowFloatMatmulBenchmark.kt index c39b838..b226d31 100644 --- a/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/NarrowFloatMatmulBenchmark.kt +++ b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/NarrowFloatMatmulBenchmark.kt @@ -6,6 +6,7 @@ import sk.ainet.lang.tensor.Tensor import sk.ainet.lang.tensor.data.Bf16DenseTensorData import sk.ainet.lang.tensor.data.DenseFloatArrayTensorData import sk.ainet.lang.tensor.data.Fp16DenseTensorData +import sk.ainet.lang.tensor.data.NarrowFloatInputMajorTensorData import sk.ainet.lang.tensor.data.NarrowFloatTensorData import sk.ainet.lang.tensor.data.TensorData import sk.ainet.lang.tensor.matmul @@ -26,56 +27,57 @@ import kotlin.test.assertTrue * * ### What it answers * - * `DecoderNarrowFloatForwardParityTest` established that KEEP_NATIVE is numerically correct but - * that the FP16/BF16 SGEMM kernels are never reached: both `Linear.onForward` and - * `LlamaRuntime.linearProject` call `w.t()`, transpose has no narrow-float arm and widens to a - * dense FP32 buffer, and `DefaultCpuOpsJvm.chooseQuantizedMatmul` only engages for `[in, out]` - * weights. Wiring the kernel up means a byte relayout at load plus a lazy-transpose arm in the - * engine — the pattern the K-quants already use. That is only worth doing if the kernel actually - * wins, so this measures three things at realistic projection sizes: + * Originally: is reaching the narrow-float kernel worth the layout work? It was — the answer + * became engine issue #888, and this now doubles as the regression guard for that fix. * - * - **fp32** — dense FP32 SGEMM. The baseline, and what runs today after the widening. - * - **fp16/bf16** — weight handed over already `[in, out]`, so `chooseQuantizedMatmul` dispatches - * to the narrow kernel. This is the best case the layout work could unlock. - * - **transpose** — narrow weight in the real `[out, in]` orientation, `.t()` then matmul. This - * is what production does per token today, and shows what the widening costs. + * Columns: + * + * - **fp32** — dense FP32 SGEMM. The baseline. + * - **fp16/bf16** — weight handed over already `[in, out]`, so `chooseQuantizedMatmul` + * dispatches to the narrow kernel. The ceiling. + * - **t()row-maj** — row-major narrow weight in the real `[out, in]` orientation, `.t()` then + * matmul. The pre-#888 path: transpose widens it elementwise through boxed + * `get()`. Still measured, because row-major narrow tensors deliberately + * keep this behaviour — only the input-major type may be reinterpreted. + * - **t()in-maj** — same weight relaid input-major at load, so `.t()` is a zero-copy view. + * What the loader produces for matmul weights now. Should match the direct + * `fp16`/`bf16` columns; any gap means the transpose is copying again. * * Dispatch is guaranteed by construction rather than observed: `chooseQuantizedMatmul` requires an * FP32 rank-2 input, a rank-2 weight, and `weight.shape[0] == input.shape[1]`. [checkDispatchable] - * asserts exactly that before timing, so a silent fallback to the generic path cannot be mistaken - * for a fast kernel. + * asserts exactly that before timing, and the input-major weight is asserted to still be narrow + * after `.t()`, so a silent fallback to the generic path cannot be mistaken for a fast kernel. * - * ### Baseline, 2026-07-27 + * ### Baseline, 2026-07-27 (engine 0.38.0-SNAPSHOT with the #888 fix) * - * Intel i7-9750H (AVX2, no AVX-512), 12 threads, OpenJDK 21.0.11, engine 0.38.0-SNAPSHOT. - * Median ms per call: + * Intel i7-9750H (AVX2, no AVX-512), 12 threads, OpenJDK 21.0.11. Median ms per call: * * ``` - * shape batch fp32 fp16 bf16 transpose - * q_proj 1B 1 3.210 18.601 1.591 206.458 - * q_proj 1B 16 16.449 297.873 11.125 215.698 - * q_proj 8B 1 40.335 73.793 21.472 1318.772 - * q_proj 8B 16 95.954 1166.224 66.227 1358.200 - * ffn_up 8B 1 111.922 197.337 60.372 4367.357 - * ffn_up 8B 16 272.889 3146.712 186.360 4511.455 - * ffn_down 8B 1 109.242 198.977 57.749 2195.538 - * ffn_down 8B 16 274.547 3172.639 177.142 2363.605 + * shape batch fp32 fp16 bf16 t()row-maj t()in-maj t()in-maj + * (fp16) (fp16) (bf16) + * q_proj 1B 1 3.219 18.421 1.521 209.514 18.530 1.499 + * q_proj 1B 16 16.402 298.912 10.720 224.442 300.323 11.033 + * q_proj 8B 1 41.016 73.822 21.423 1368.664 74.180 21.254 + * q_proj 8B 16 98.640 1182.081 67.116 1394.869 1182.299 67.735 + * ffn_up 8B 1 110.393 199.491 58.476 4465.222 199.678 58.007 + * ffn_up 8B 16 273.756 3200.778 183.782 4591.961 3201.357 184.874 + * ffn_down 8B 1 111.234 199.131 58.412 2295.821 201.194 58.759 + * ffn_down 8B 16 263.774 3191.621 181.348 2397.454 3201.594 180.597 * ``` * * Three conclusions: * - * 1. **BF16 beats FP32 by 1.5–2.1x everywhere.** At batch 1 the matmul is memory-bandwidth bound, - * so halving the weight bytes roughly halves the time. This is the case for doing the layout - * work. - * 2. **FP16 is 2–18x slower, pinned at ~0.5 GFLOP/s regardless of shape or batch** — the signature - * of being compute-bound on the decode. Both Panama kernels fill a scratch lane array scalar-wise - * before the vector FMA, but BF16's decode is three integer ops while `Fp16Codec.decode` is a - * branchy `when` with a subnormal renormalization loop. The fix is engine-side and independent - * of layout: use `Float.float16ToFloat` (a JDK 20+ intrinsic) or a branch-free decode. - * 3. **The `transpose` column is the alarming one.** 0.2–4.5 *seconds* for one projection, because - * the generic transpose walks a narrow tensor element by element through `get()`. That is the - * path production takes today, per weight, per token. KEEP_NATIVE is not merely un-accelerated - * right now — at real model sizes it is unusably slow. + * 1. **The relayout removes the transpose cost entirely.** `t()in-maj` matches the direct column + * to within noise at every size, so the zero-copy view holds. Against the old path that is + * 4465 ms → 58 ms for `ffn_up` BF16 at batch 1, a 77x reduction, and 209 ms → 1.5 ms for + * `q_proj 1B`. Before this, KEEP_NATIVE was unusably slow at real model sizes. + * 2. **BF16 beats FP32 by 1.5–2.1x.** At batch 1 the matmul is memory-bandwidth bound, so halving + * the weight bytes roughly halves the time. This is the win the feature exists for. + * 3. **FP16 is still 2–18x slower than FP32**, pinned near 0.5 GFLOP/s regardless of shape or + * batch — compute-bound on the decode, not on layout. Both Panama kernels fill a scratch lane + * array scalar-wise before the vector FMA, but BF16's decode is three integer ops while + * `Fp16Codec.decode` is a branchy `when` with a subnormal renormalization loop. Tracked + * separately as engine issue #887; until it lands, prefer BF16 for speed. */ class NarrowFloatMatmulBenchmark { @@ -161,14 +163,20 @@ class NarrowFloatMatmulBenchmark { val ctx = DirectCpuExecutionContext() println() println("narrow-float matmul vs fp32 SGEMM (median of timed samples)") - println("weight layout [in, out]; 'transpose' is the [out, in] + .t() path production uses today") + println("direct columns take the weight already [in, out]; t() columns transpose [out, in] first") println() println( - "%-13s %6s %11s %11s %11s %11s %s".format( - "shape", "batch", "fp32", "fp16", "bf16", "transpose", "verdict", + "%-13s %6s %10s %10s %10s %12s %11s %11s %s".format( + "shape", "batch", "fp32", "fp16", "bf16", + "t()row-maj", "t()in-maj", "t()in-maj", "verdict", + ), + ) + println( + "%-13s %6s %10s %10s %10s %12s %11s %11s".format( + "", "", "", "", "", "(fp16)", "(fp16)", "(bf16)", ), ) - println("-".repeat(96)) + println("-".repeat(118)) for ((label, inF, outF) in shapes) { val raw = weights(inF * outF, seed = inF + outF) @@ -195,12 +203,30 @@ class NarrowFloatMatmulBenchmark { Bf16DenseTensorData(Shape(inF, outF), bf16Bytes) as TensorData, FP32::class, ) - // The production orientation: [out, in], transposed on every call. + // The production orientation: [out, in], transposed on every call. Row-major, so the + // transpose falls to the generic elementwise path — what production did before #888. @Suppress("UNCHECKED_CAST") val wFp16Transposed = ctx.fromData( Fp16DenseTensorData(Shape(outF, inF), fp16Bytes) as TensorData, FP32::class, ) + // Same [out, in] orientation, but relaid input-major at load, so `.t()` is a + // zero-copy view and the weight reaches the kernel still packed. What the loader + // produces for matmul weights now. + @Suppress("UNCHECKED_CAST") + val wFp16InputMajor = ctx.fromData( + NarrowFloatInputMajorTensorData.fromRowMajor( + Shape(outF, inF), fp16Bytes, Fp16Codec, + ) as TensorData, + FP32::class, + ) + @Suppress("UNCHECKED_CAST") + val wBf16InputMajor = ctx.fromData( + NarrowFloatInputMajorTensorData.fromRowMajor( + Shape(outF, inF), bf16Bytes, Bf16Codec, + ) as TensorData, + FP32::class, + ) for (batch in batches) { val x = ctx.fromFloatArray( @@ -214,6 +240,15 @@ class NarrowFloatMatmulBenchmark { val fp16Ns = measure { x.matmul(wFp16) } val bf16Ns = measure { x.matmul(wBf16) } val transposeNs = measure { x.matmul(wFp16Transposed.t()) } + val inMajFp16Ns = measure { x.matmul(wFp16InputMajor.t()) } + val inMajBf16Ns = measure { x.matmul(wBf16InputMajor.t()) } + + // The relaid weight must still be packed after `.t()`, or the two columns below + // are just re-measuring the generic path under a different name. + assertTrue( + wBf16InputMajor.t().data is NarrowFloatTensorData, + "$label: input-major weight widened on transpose — engine #888 arm missing?", + ) // Sanity: the narrow kernel must agree with the FP32 baseline, or the timing is // measuring something that isn't a correct matmul. @@ -234,26 +269,20 @@ class NarrowFloatMatmulBenchmark { } println( - "%-13s %6d %9.3fms %9.3fms %9.3fms %9.3fms %s".format( + "%-13s %6d %8.3fms %8.3fms %8.3fms %10.3fms %9.3fms %9.3fms %s".format( label, batch, - fp32Ns / 1e6, fp16Ns / 1e6, bf16Ns / 1e6, transposeNs / 1e6, + fp32Ns / 1e6, fp16Ns / 1e6, bf16Ns / 1e6, + transposeNs / 1e6, inMajFp16Ns / 1e6, inMajBf16Ns / 1e6, verdict, ), ) - println( - "%-13s %6s %9.1fGF %9.1fGF %9.1fGF %9s".format( - "", "", - gflops(batch, inF, outF, fp32Ns), - gflops(batch, inF, outF, fp16Ns), - gflops(batch, inF, outF, bf16Ns), - "", - ), - ) } } println() - println("weight bytes at rest: fp32 = 2x narrow. The 'transpose' column is the per-call cost") - println("of the widening that happens today, and is what the layout work would remove.") + println("weight bytes at rest: fp32 = 2x narrow.") + println("t()row-maj is the pre-#888 path: a row-major narrow weight widened elementwise on") + println("every transpose. t()in-maj is the same weight relaid input-major at load, so the") + println("transpose is a zero-copy view and the packed weight reaches the kernel.") } } From c92ac0491e04f88f801f080627e89cb6aa2d0887 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Thu, 30 Jul 2026 10:00:20 +0200 Subject: [PATCH 5/6] docs(perf): refresh the narrow-float baseline against the merged engine The recorded baseline predates the engine fixes it was written to track, so its headline conclusion is now backwards: it says FP16 is 2-18x slower than FP32, pinned near 0.5 GFLOP/s, blames Fp16Codec.decode's subnormal renormalization loop, and tells the reader to prefer BF16 until #887 lands. FP16 is now 1.5-1.7x faster than FP32 and #887 is merged. Re-measure against engine develop with #888, #887 and the BF16 amortization all merged, and rewrite the conclusions around what the numbers actually show. The decode diagnosis in the old text was the issue's original theory and it was wrong, so the replacement says what the cause turned out to be: NativeKernelProvider carried matmulBf16 but no matmulFp16, so BF16 ran the native FFM kernel at priority 100 while FP16 cascaded to the JVM Panama kernel at 50. Head to head the two Panama kernels are within ~15% of each other. Worth recording, because the failure presented as a slow kernel and was a missing one. Adds a fourth conclusion for the other thing that fell out: both native kernels now read B once per matmul rather than once per row of A, which cut B traffic 16x at m=16 and bought only 9-19%. These kernels are compute-bound on the FMA chain at batch 16, not bandwidth-bound, so further layout work is not where the next win is. Flags the q_proj 1B batch-1 row as noisy rather than quietly re-rolling it: its FP32 sample came out at 5.35 ms against a 3.09-3.22 ms cluster in repeat runs. It is the shortest measurement in the table. --- .../llama/NarrowFloatMatmulBenchmark.kt | 51 ++++++++++++------- 1 file changed, 32 insertions(+), 19 deletions(-) diff --git a/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/NarrowFloatMatmulBenchmark.kt b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/NarrowFloatMatmulBenchmark.kt index b226d31..b806e9f 100644 --- a/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/NarrowFloatMatmulBenchmark.kt +++ b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/NarrowFloatMatmulBenchmark.kt @@ -28,7 +28,8 @@ import kotlin.test.assertTrue * ### What it answers * * Originally: is reaching the narrow-float kernel worth the layout work? It was — the answer - * became engine issue #888, and this now doubles as the regression guard for that fix. + * became engine issue #888, and this now doubles as the regression guard for that fix and for + * #887, the FP16 kernel gap it exposed next. * * Columns: * @@ -48,36 +49,48 @@ import kotlin.test.assertTrue * asserts exactly that before timing, and the input-major weight is asserted to still be narrow * after `.t()`, so a silent fallback to the generic path cannot be mistaken for a fast kernel. * - * ### Baseline, 2026-07-27 (engine 0.38.0-SNAPSHOT with the #888 fix) + * ### Baseline, 2026-07-30 (engine develop with #888, #887 and the BF16 amortization all merged) * * Intel i7-9750H (AVX2, no AVX-512), 12 threads, OpenJDK 21.0.11. Median ms per call: * * ``` * shape batch fp32 fp16 bf16 t()row-maj t()in-maj t()in-maj * (fp16) (fp16) (bf16) - * q_proj 1B 1 3.219 18.421 1.521 209.514 18.530 1.499 - * q_proj 1B 16 16.402 298.912 10.720 224.442 300.323 11.033 - * q_proj 8B 1 41.016 73.822 21.423 1368.664 74.180 21.254 - * q_proj 8B 16 98.640 1182.081 67.116 1394.869 1182.299 67.735 - * ffn_up 8B 1 110.393 199.491 58.476 4465.222 199.678 58.007 - * ffn_up 8B 16 273.756 3200.778 183.782 4591.961 3201.357 184.874 - * ffn_down 8B 1 111.234 199.131 58.412 2295.821 201.194 58.759 - * ffn_down 8B 16 263.774 3191.621 181.348 2397.454 3201.594 180.597 + * q_proj 1B 1 5.352 3.129 1.513 212.110 3.051 1.529 + * q_proj 1B 16 16.556 10.716 9.261 218.329 10.498 9.187 + * q_proj 8B 1 40.002 26.276 20.951 1373.229 26.626 20.975 + * q_proj 8B 16 95.454 58.566 53.329 1417.367 59.053 53.533 + * ffn_up 8B 1 108.667 71.156 57.428 4389.079 72.150 57.153 + * ffn_up 8B 16 271.129 159.255 143.489 4506.450 155.102 139.751 + * ffn_down 8B 1 107.462 70.868 56.134 2216.040 70.862 55.819 + * ffn_down 8B 16 263.130 155.509 141.947 2383.415 155.820 143.709 * ``` * - * Three conclusions: + * `q_proj 1B` at batch 1 is the shortest measurement here and the noisiest: its FP32 sample came + * out at 5.35 ms against a 3.09–3.22 ms cluster in repeat runs, so read that row's ratios with + * suspicion. The other seven are stable across runs. + * + * Four conclusions: * * 1. **The relayout removes the transpose cost entirely.** `t()in-maj` matches the direct column * to within noise at every size, so the zero-copy view holds. Against the old path that is - * 4465 ms → 58 ms for `ffn_up` BF16 at batch 1, a 77x reduction, and 209 ms → 1.5 ms for + * 4389 ms → 57 ms for `ffn_up` BF16 at batch 1, a 77x reduction, and 212 ms → 1.5 ms for * `q_proj 1B`. Before this, KEEP_NATIVE was unusably slow at real model sizes. - * 2. **BF16 beats FP32 by 1.5–2.1x.** At batch 1 the matmul is memory-bandwidth bound, so halving - * the weight bytes roughly halves the time. This is the win the feature exists for. - * 3. **FP16 is still 2–18x slower than FP32**, pinned near 0.5 GFLOP/s regardless of shape or - * batch — compute-bound on the decode, not on layout. Both Panama kernels fill a scratch lane - * array scalar-wise before the vector FMA, but BF16's decode is three integer ops while - * `Fp16Codec.decode` is a branchy `when` with a subnormal renormalization loop. Tracked - * separately as engine issue #887; until it lands, prefer BF16 for speed. + * 2. **Both narrow formats now beat FP32** — BF16 by 1.8–1.9x, FP16 by 1.5–1.7x — at every shape + * and both batch sizes. Halving the weight bytes is most of it. + * 3. **FP16 trails BF16 by only 10–24%**, which is the cost of its dequant: BF16 is one shift, + * binary16 needs rebiasing and gradual underflow. It used to trail by 2–18x, and the cause was + * not the decode at all — `NativeKernelProvider` carried `matmulBf16` but no `matmulFp16`, so + * BF16 ran the native FFM kernel at priority 100 while FP16 silently cascaded to the JVM + * Panama kernel at 50. Head to head the two Panama kernels are within ~15% of each other. + * Fixed by #887; the lesson is that a dtype benchmarking far off its siblings is more likely + * served by a different provider than by a worse kernel. + * 4. **These kernels are compute-bound at batch 16, not bandwidth-bound.** Both native kernels + * now tile `j` and read B once per matmul instead of once per row of A, cutting B traffic 16x + * at m=16 — and that bought only 9–19%. Whatever is left is the FMA chain, so the next real + * win is a blocked microkernel or `bfdot`/`bfmmla` on ARMv8.6-A+, not more layout work. + * At m=1 both kernels deliberately keep the straight i-p-j pass: there is nothing to amortize, + * and tiling cost 15% there. */ class NarrowFloatMatmulBenchmark { From 161c4b211e8b1b3729e56a1ea7cf2fa7c065aa95 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Thu, 30 Jul 2026 12:53:18 +0200 Subject: [PATCH 6/6] fix(test): drop parentheses and comma from DTypePolicyValidation test names compileTestKotlinIosArm64 fails on two backtick test names in DTypePolicyValidationTest: Kotlin/Native rejects "()" and "," in a declaration name because the name has to survive ObjC export. e: Name contains illegal characters: "()" e: Name contains illegal characters: "," The names are commonTest, so every native target compiles them, but the branch never got that far before -- CI died at dependency resolution looking for the unreleased 0.38.0-SNAPSHOT, so `assemble allTests` never reached the iOS compile. Pinning the released engine exposed it. Rename both; the assertions are untouched. Also swept the rest of commonTest for the same pattern -- these two were the only ones. Verified with the exact CI invocation, not just jvmTest: ./gradlew --no-configuration-cache assemble allTests is green. --- .../kotlin/sk/ainet/apps/llm/DTypePolicyValidationTest.kt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/llm-core/src/commonTest/kotlin/sk/ainet/apps/llm/DTypePolicyValidationTest.kt b/llm-core/src/commonTest/kotlin/sk/ainet/apps/llm/DTypePolicyValidationTest.kt index 6fc700c..2417f85 100644 --- a/llm-core/src/commonTest/kotlin/sk/ainet/apps/llm/DTypePolicyValidationTest.kt +++ b/llm-core/src/commonTest/kotlin/sk/ainet/apps/llm/DTypePolicyValidationTest.kt @@ -23,13 +23,13 @@ class DTypePolicyValidationTest { private val bothNarrow = setOf(BF16, FP16) @Test - fun `Require(FP32) is always accepted — every chain produces FP32`() { + fun `Require FP32 is always accepted — every chain produces FP32`() { DTypePolicyValidation.validate(DTypePolicy.Require(FP32), "test", keepNative = emptySet()) DTypePolicyValidation.validate(DTypePolicy.Require(FP32), "test", keepNative = bothNarrow) } @Test - fun `soft policies never raise, whatever they name`() { + fun `soft policies never raise whatever they name`() { for (keepNative in listOf(emptySet(), bothNarrow)) { DTypePolicyValidation.validate(DTypePolicy.Any, "test", keepNative) DTypePolicyValidation.validate(DTypePolicy.Prefer(BF16), "test", keepNative)