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..2417f85 --- /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/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/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..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 @@ -12,11 +12,18 @@ 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.lang.tensor.data.Bf16DenseTensorData +import sk.ainet.apps.llm.DTypePolicyValidation +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 import sk.ainet.lang.types.FP32 import kotlin.math.pow import kotlin.reflect.KClass @@ -29,16 +36,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): 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) * * @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 +61,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 +105,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 @@ -106,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 { @@ -117,10 +126,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 = narrowData(canonicalName, targetShape, bytes, Fp16Codec) + @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) @@ -161,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/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/DecoderNarrowFloatForwardParityTest.kt b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/DecoderNarrowFloatForwardParityTest.kt new file mode 100644 index 0000000..b39e2b8 --- /dev/null +++ b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/DecoderNarrowFloatForwardParityTest.kt @@ -0,0 +1,394 @@ +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.data.NarrowFloatInputMajorTensorData +import sk.ainet.lang.tensor.data.NarrowFloatTensorData +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.assertSame +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. + * + * 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. + */ +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. 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-5f + + 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 + + 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, + metadata = metadata, + tiedEmbeddings = false, + dtypePolicy = policy, + ) + val provider: () -> RandomAccessSource = { JvmRandomAccessSource.open(file) } + 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( + 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 `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, + ) + val weights = loadWeights(file, DTypePolicy.Require(FP16)) + + val ffnGate = weights.tensors[LlamaTensorNames.ffnGate(0)] + ?: error("missing ffn_gate") + assertTrue( + ffnGate.data is NarrowFloatInputMajorTensorData, + "a matmul weight must be relaid input-major, got ${ffnGate.data::class.simpleName}", + ) + + val transposed = ffnGate.t() + assertTrue( + transposed.data is NarrowFloatTensorData, + "transpose widened the weight — the narrow kernel is unreachable again", + ) + assertEquals( + 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", + ) + } + + @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", + ) + } +} 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..d4fd94c --- /dev/null +++ b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/DecoderSafeTensorsLoaderNarrowFloatTest.kt @@ -0,0 +1,244 @@ +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.FloatArrayTensorData +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 +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. 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. + */ +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 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( + 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, + ) + + val embedding = load(file, DTypePolicy.Require(FP16))[LlamaTensorNames.TOKEN_EMBEDDINGS] + ?: error("missing token_embd") + + assertTrue(embedding.data is NarrowFloatTensorData, "must stay packed") + assertTrue( + embedding.data !is NarrowFloatInputMajorTensorData, + "a gathered tensor must not be relaid — input-major storage strides its row reads", + ) + assertContentEquals( + onDisk, (embedding.data as NarrowFloatTensorData).packedData, + "the row-major path must preserve on-disk bytes verbatim", + ) + } + + @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)) + 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") + 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 NarrowFloatTensorData) + assertTrue( + 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/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/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..b806e9f --- /dev/null +++ b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/NarrowFloatMatmulBenchmark.kt @@ -0,0 +1,301 @@ +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.NarrowFloatInputMajorTensorData +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 + * + * 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 and for + * #887, the FP16 kernel gap it exposed next. + * + * 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, 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-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 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 + * ``` + * + * `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 + * 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. **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 { + + 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("direct columns take the weight already [in, out]; t() columns transpose [out, in] first") + println() + println( + "%-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(118)) + + 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. 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( + 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()) } + 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. + 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 %8.3fms %8.3fms %8.3fms %10.3fms %9.3fms %9.3fms %s".format( + label, batch, + fp32Ns / 1e6, fp16Ns / 1e6, bf16Ns / 1e6, + transposeNs / 1e6, inMajFp16Ns / 1e6, inMajBf16Ns / 1e6, + verdict, + ), + ) + } + } + + println() + 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.") + } +} 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 -> {