diff --git a/CHANGELOG.md b/CHANGELOG.md index f86ca0cb1..8c3c64558 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,24 @@ ## [Unreleased] +Headline: **the browser lane stops falling to the reference kernel (SKEEP-007, Tier 0).** + +### Added + +- **SKEEP-007 "The browser lane"** (`docs/modules/skeep/pages/007-browser-lane.adoc`): the design record for making inference in a browser usable and quick — common-code fixes now, a SIMD128 sidecar kernel pack with pthreads and quantize-at-load for SafeTensors as the next pull requests, WebGPU as a later proposal. +- **`ByteArrayRandomAccessSource`** (`skainet-io-core`, `sk.ainet.io`): the engine's in-memory `RandomAccessSource`, single array or equal-sized chunks for files above the 2 GB array limit. Every web sample had written its own. +- **`NarrowFloatDenseTensorData` is a `RowDequantSource`**: FP16 and BF16 tables serve rows on demand, so a 16-bit token-embedding table loaded `KEEP_NATIVE` is read row by row by `ops.gather` instead of being widened to FP32 at load. +- **Quantize-at-load (SKEEP-007 Tier 2).** `Q8_0Quantizer` (`skainet-lang-core`, ggml-faithful, beside the existing `Q4_0Quantizer`) and the `Requantizers` registry (Q8_0, Q4_0, BITNET_PLANES). `SafeTensorsParametersLoader` and `ShardedSafeTensorsParametersLoader` take `weightForm` / `weightFormFor`: `EncodingRequest.RequantizeTo(Q8_0)` turns every rank-2 float matrix whose rows block into packed Q8_0 at load (vectors, norms and non-blocking rows stay dense), with `KERNEL_FEED` honoured; a request no quantizer can serve is refused at construction. `PlannerProfile.quantizeDenseTo` lets a profile ask `WeightFormResolver` for it on dense sources when a kernel can feed the encoding. Every Hugging Face checkpoint can now be served packed on every target; before, only GGUF files could. +- **`skainet-quantize`** (`skainet-apps/skainet-tensor-tools`, `./gradlew :skainet-apps:skainet-tensor-tools:quantize --args=...`): SafeTensors → GGUF with the weight matrices quantized offline (Q8_0 or Q4_0), tensor names kept, `--keep` prefixes written dense as F16 — for downloads that should not carry F16 weights. +- **`skainet-backends/benchmarks/wasm-bench`**: a Kotlin/Wasm executable that times the matmul shapes of the shipped models and prints `[SKAINET_PERF_WASM]` lines — the baseline for the following tiers. + +### Changed + +- **`matmulWeightTransposed` on a canonical packed weight** (every GGUF weight) now prefers, in order: an exact kernel for the row-major view, the one-time cached relayout into kernel feed order that the scalar packed kernels on every target read, and only then the decoding reference kernel. On JS, wasm and Kotlin/Native without the cinterop pack this replaces a ~1000× slower per-element decode on every call; the JVM overrides the path and is unaffected. +- **`matmulWeightTransposed` on a dense heap FP32 `[out, in]` weight** computes the product straight from the weight's rows (4 × 4 tiles, one sequential accumulator per output, bit-identical to the transpose-then-`matmul` result) on tiers without a vectorised `[k, n]` kernel, instead of making and caching a transposed copy of every weight. `DefaultCpuOpsBase.preferDirectWeightTransposed()` is the hook; the JVM and Apple ops keep their cached transpose. +- **GGUF `RequantizeTo`** accepts every encoding in `Requantizers` (was `BITNET_PLANES` only, #1150); freshly quantized Q8_0/Q4_0 honour `KERNEL_FEED`. +- **`JsBlobRandomAccessSource`** no longer logs every read to the console. (The byte-by-byte `ArrayBuffer` conversion stays: Kotlin/Wasm 2.4 has no typed-array adapter; SKEEP-007 Tier 1 replaces it with the sidecar's linear-memory bridge.) + ## [0.60.0] - 2026-10-15 Headline: **key-padding lengths and window metadata for attention (#1343).** diff --git a/docs/modules/ROOT/pages/explanation/eager-execution.adoc b/docs/modules/ROOT/pages/explanation/eager-execution.adoc index a81b6a610..3e1a9af80 100644 --- a/docs/modules/ROOT/pages/explanation/eager-execution.adoc +++ b/docs/modules/ROOT/pages/explanation/eager-execution.adoc @@ -47,7 +47,7 @@ mindmap Android — scalar, Native-JNI Native linux — scalar, Native-cinterop Native apple — scalar, Accelerate, Native-cinterop - JS and WASM — scalar only + JS and WASM — scalar; packed weights relaid once into kernel feed order (SKEEP-007) ---- == Kernel × provider (matmul, FP32 activations) @@ -88,7 +88,7 @@ not per call. | *Android* | Scalar + Native-JNI | Panama (`jdk.incubator.vector`) and FFM (`java.lang.foreign`) are JDK-only — ART has neither, so Android's native tier is JNI, not a degraded JVM. Includes the NEON `bitnet_gemv` pack. | *Kotlin/Native — linux x64/arm64* | Scalar + Native-cinterop | Static archive embedded in the klib; manual `installNativeKernels()`, no ServiceLoader on K/N. | *Kotlin/Native — macOS/iOS* | Scalar + Apple Accelerate + Native-cinterop | Accelerate covers dense FP32 and reductions; cinterop covers packed quant with runtime `FEAT_DotProd` dispatch, so one archive serves A12 through M-series. -| *JS / WASM (Js, Wasi)* | Scalar | No SIMD. +| *JS / WASM (Js, Wasi)* | Scalar | No SIMD in Kotlin code. Packed weights are relaid once per weight into kernel feed order and served by the scalar packed kernels; dense `[out, in]` weights are multiplied in place without a transpose copy (SKEEP-007 Tier 0). A SIMD128 sidecar pack for wasmJs is SKEEP-007 Tier 1. |=== Packed-quant matmul works on every target: those formats have a `commonMain` scalar kernel and diff --git a/docs/modules/ROOT/pages/explanation/kernel-selection.adoc b/docs/modules/ROOT/pages/explanation/kernel-selection.adoc index f4eae80cd..173988822 100644 --- a/docs/modules/ROOT/pages/explanation/kernel-selection.adoc +++ b/docs/modules/ROOT/pages/explanation/kernel-selection.adoc @@ -94,6 +94,11 @@ platform support is absent (a missing native library, no vector unit). | manual — call the pack's `install()` yourself |=== +Without a row-major pack, a canonical packed weight is relaid **once per weight** into kernel feed +order and served by the scalar packed kernels (`DefaultCpuOpsBase.matmulWeightTransposed`, SKEEP-007 +Tier 0); the decoding reference kernel is the last resort, reached only for an encoding with no +scalar kernel at all. + A backend module declares its service the usual way, e.g. `skainet-backend-native-cpu` ships `META-INF/services/sk.ainet.backend.api.kernel.ViewKernelPack` naming `FfmRowMajorKernelPackFactory`. `ServiceLoader` needs a public no-arg constructor, which a Kotlin diff --git a/docs/modules/skeep/nav.adoc b/docs/modules/skeep/nav.adoc index 76a75ef4d..c4dc4e582 100644 --- a/docs/modules/skeep/nav.adoc +++ b/docs/modules/skeep/nav.adoc @@ -8,3 +8,4 @@ ** xref:skeep:004-virtual-tensor-layout.adoc[SKEEP-004: Virtual tensor layout] ** xref:skeep:005-schedules-structured-concurrency.adoc[SKEEP-005: Schedules — structured concurrency] ** xref:skeep:006-input-frontends-and-model-identity.adoc[SKEEP-006: Input front ends and model identity] +** xref:skeep:007-browser-lane.adoc[SKEEP-007: The browser lane] diff --git a/docs/modules/skeep/pages/007-browser-lane.adoc b/docs/modules/skeep/pages/007-browser-lane.adoc new file mode 100644 index 000000000..34663972a --- /dev/null +++ b/docs/modules/skeep/pages/007-browser-lane.adoc @@ -0,0 +1,256 @@ += SKEEP-007: The browser lane — packed SIMD kernels and quantize-at-load +:description: SKaiNET proposal for making inference in a browser usable and quick: stop the scalar wasm lane from falling to the decoding reference kernel, serve dense weights without transpose copies, keep token tables narrow, add a SIMD128 kernel pack with pthreads as a sidecar wasm module, and let the SafeTensors loader quantize at load so Hugging Face checkpoints are served packed on every target. + +Status: Draft — Tier 0 and Tier 2 in review + +Implementation: Tier 0 https://github.com/SKaiNET-developers/SKaiNET/pull/1352[#1352]; Tier 2 quantize-at-load (the pull request stacked on it); Tier 1 `skainet-backend-wasm-simd` and the consumer work in SKaiNET-transformers follow, see <<_rollout_plan>>. + +Audience: SKaiNET maintainers and contributors; SKaiNET-transformers maintainers (Laya and BERT in the browser); SKaiNET-examples maintainers (every web sample runs through this lane) + +Created: 2026-10-16 + +Origin: the SemanticEmojiFilter example ships its 23M-parameter embedding filter to the web but keeps the 322M-parameter Laya verifier on desktop and Android, because a decision in scalar Kotlin/Wasm is estimated at 10–30 s and the FP32 weights at about 2.5 GB resident. Nothing in the repository measures the wasm lane today. + +== Summary + +Every SKaiNET browser build runs the same path: scalar Kotlin loops compiled by Kotlin/Wasm, one +thread, GC `FloatArray`/`ByteArray` tensors. Three defects compound on that path, and a fourth +makes the compute itself slow: + +* A canonical (row-major) block-quantized weight — every weight loaded from a GGUF file — reaches + `KernelDispatch` on a target without a row-major kernel pack and lands on the decoding reference + kernel, which the dispatcher itself warns is about a thousand times slower. The scalar packed + kernels every target carries are keyed on kernel feed order and were never reached. +* A dense FP32 `[out, in]` weight is transposed into a cached heap copy before every first use, + which doubles the resident weight memory of a dense encoder. +* A 16-bit token table is widened to FP32 at load (for a 256k × 768 vocabulary, 786 MB instead of + 393 MB) although the embedding lookup already knows how to read rows from a `RowDequantSource`. +* Kotlin/Wasm exposes neither SIMD types nor shared-memory threads, so no Kotlin kernel can be + faster than a scalar loop on one core. + +The proposal fixes the three defects in common code (Tier 0, this change), adds a browser kernel +pack built from the engine's existing C kernel sources as a SIMD128 sidecar module with Emscripten +pthreads (Tier 1), and teaches the SafeTensors loader the `EncodingRequest.RequantizeTo` form the +planner already prices, with a common-code Q8_0 quantizer (Tier 2). WebGPU is recorded as a later +proposal behind the same kernel-pack seam. + +== Motivation + +=== Who pays today + +[cols="2,3,3",options="header"] +|=== +| Sample | Browser build today | After this proposal + +| KernelRace (SmolLM2-135M Q8_0) +| loads the model as dense FP32 ("in-memory FP32 fallback") +| packed Q8_0 served by the SIMD128 pack + +| KllamaDemo (Qwen3-0.6B Q4_0) +| dequantizes to FP32 at load, "perceptibly slow" +| packed Q4_0, same kernels + +| SemanticEmojiFilter (LEAF 23M, Laya 322M) +| LEAF only; Laya is desktop and Android +| Laya as an opt-in verifier in the browser + +| TinyTransformer, MNIST, Sinus, GloVe, TDALab +| scalar dense matmul +| dense FP32 kernel from the same pack +|=== + +=== Where the time goes + +One Laya decision at 80–150 tokens is 10–19 G multiply-adds, 96 % of them in four dense +`matmulWeightTransposed` calls per layer (`Wqkv`, `Wi`, `mlp.Wo`, `Wo`). The scalar `[k, n]` loop +streams the whole weight once per activation row: about 75 GB of memory traffic per decision at +150 tokens. For comparison, the same kernel class written with wasm SIMD128 intrinsics runs a +Llama-3.2-1B Q8_0 forward at roughly 270 GFLOPS on seven threads (llama.cpp #11453, measured in +wllama), which puts a Laya decision at well under a second. + +=== Why the engine, not the sample + +The kernel-pack seam (`ViewKernelPack`, `KernelDispatch`, `KernelProvider`) is how the JVM and +Android got their packed kernels; a browser pack behind the same seam reaches every web build +without application changes. Quantize-at-load is likewise a loader concern: today only GGUF files +can be served packed, so every Hugging Face checkpoint — SafeTensors, F16 or BF16 — is dense FP32 +on every target. + +== Goals + +* A canonical packed weight never reaches the decoding reference kernel when a scalar packed kernel + exists for its format (Tier 0). +* A dense FP32 `[out, in]` weight is multiplied in place on tiers without a vectorised `[k, n]` + kernel; no transpose copy is made or cached there (Tier 0). +* 16-bit dense tables are `RowDequantSource`s (Tier 0). +* The engine ships an in-memory `RandomAccessSource` and a wasm benchmark that prints GFLOPS for + the matmul shapes of the shipped models (Tier 0). +* A SIMD128 kernel pack for wasmJs, built from `skainet-backend-native-cpu/native/src`, with a + multithreaded artifact for cross-origin-isolated pages and a single-threaded fallback (Tier 1). +* `SafeTensorsParametersLoader` honours `WeightForm`, including `RequantizeTo(Q8_0)`; a common-code + Q8_0 quantizer; `skainet-quantize` for offline conversion (Tier 2). +* No behaviour change on the JVM and Apple tiers: their kernel keys and golden tests are untouched. + +== Non-Goals + +* WebGPU. It is the right second pack behind the same seam, but it needs its own proposal (WGSL + kernels, device selection, Firefox desktop still off by default). +* A streaming or asynchronous `RandomAccessSource`; SKEEP-003 keeps that open item. +* K-quant encoders (Q4_K, Q5_K, Q6_K) at load; Q8_0 and the existing Q4_0 cover the first consumers. +* Kotlin-side multithreading. Kotlin/Wasm has no shared-memory threads; parallelism lives inside the + sidecar. + +== Proposed Design + +=== Tier 0 — common code + +`DefaultCpuOpsBase.matmulWeightTransposed` orders the packed-weight paths by preference: an exact +kernel for the canonical view (the JVM FFM, Android JNI and browser packs), then the one-time +cached relayout into kernel feed order that the scalar packed kernels read (`prepackedWeights` was +declared for this and never used in common code), then the decoding reference kernel through views. +The dense branch computes `x · Wᵀ` straight from the weight's rows with a 4 × 4 tile and one +sequential accumulator per output, bit-identical to the transpose-then-`matmul` result. The JVM and +Apple ops override `preferDirectWeightTransposed()` to keep their cached transpose, because their +`[k, n]` kernels are vectorised. + +`NarrowFloatDenseTensorData` (and so `Fp16DenseTensorData`, `Bf16DenseTensorData`) implements +`RowDequantSource`; `ops.gather` and the transformers `Embedding` layer already branch on it. + +`skainet-io-core` gains `ByteArrayRandomAccessSource` (single array or equal-sized chunks for files +above the 2 GB array limit). `JsBlobRandomAccessSource` no longer logs to the console on every read; +its byte-by-byte `ArrayBuffer` conversion stays until Tier 1 (Kotlin/Wasm 2.4 has no typed-array +adapter), where the sidecar's linear memory becomes the landing buffer for downloads. + +`skainet-backends/benchmarks/wasm-bench` is a Kotlin/Wasm executable that times the matmul shapes +of Laya, LEAF and SmolLM2 and prints `[SKAINET_PERF_WASM]` lines; it is the baseline every later +tier is measured against. + +=== Tier 1 — `skainet-backend-wasm-simd` + +A wasmJs-only module. Its native build compiles the existing kernel sources +(`skainet-backend-native-cpu/native/src/*.c`, shared, not copied) with Emscripten to +`skainet_kernels_{mt,st}.{wasm,mjs}`; `skainet_simd.h` gains `__wasm_simd128__` bodies for the +row-major Q8_0/Q4_0/Q4_K/Q5_K/Q6_K dot products (int8 lanes widened to int16, `i32x4.dot_i16x8` +accumulation — the technique of llama.cpp #11453, our code) and f32x4 bodies for the dense FP32, +FP16 and BF16 kernels. `skainet_row_threads.c` learns the Emscripten pool size and degrades to +the calling thread when the page is not cross-origin isolated. The artifacts ship as module +resources and are reproducible from the pinned emsdk in CI. + +On the Kotlin side a `WasmRowMajorMatmulKernel : ViewKernel` registers under the same keys as +`FfmRowMajorMatmulKernel` (dense contiguous FP32 activation × `Format(FP32, encoding)` weight, +`BLOCKED_ROW_MAJOR`, heap bytes), uploads a weight into sidecar memory once and copies activations +per call; a `WasmKernelProvider` serves the dense kernels; `WasmCpuOps : DefaultCpuOpsBase` +overrides `matmul`, `matmulWeightTransposed`, `gelu` and `softmax` while the sidecar is ready and +falls back to the Tier 0 paths otherwise. Resident activations in sidecar memory (a +`Storage.WasmLinear` kind) and a banded attention kernel follow once the copy overhead is measured. + +=== Tier 2 — quantize-at-load + +`lang-core` gains `Q8_0Quantizer` (ggml-faithful, 32-element blocks, row-major) beside the existing +`Q4_0Quantizer`, and a `Requantizers` registry from `TensorEncoding` to quantizer. +`SafeTensorsParametersLoader` and the sharded loader take `weightForm`/`weightFormFor`; the +materializer gains a `RequantizeTo` branch producing `XBlockTensorData(shape, bytes, ROW_MAJOR)` and +honours `WeightShapeOrientation` and `KERNEL_FEED`. The GGUF loader's `RequantizeTo` gate relaxes +from "BITNET_PLANES only" to "any encoding with a registered quantizer". `PlannerProfile` gains +`quantizeDenseTo` so a browser profile can resolve dense sources to `RequantizeTo(Q8_0)`. +`skainet-quantize` (in `skainet-tensor-tools`) writes a GGUF with the source tensor names kept, for +downloads that should not carry F16. + +== Requirements + +=== Functional + +. A canonical Q8_0/Q4_0/Q4_K/Q5_K/Q6_K/Q5_0/Q5_1 weight multiplied through `matmulWeightTransposed` + on a target without a row-major pack uses the scalar packed kernel after one relayout per weight. +. `matmulWeightTransposed` on a heap FP32 `[out, in]` weight returns the same bits as + `matmul(x, transpose(w))` on every tier. +. `Fp16DenseTensorData.dequantRow(r)` equals `get(r, c)` for every column. +. `ByteArrayRandomAccessSource` serves any window across chunk boundaries. +. (Tier 1) Kernel parity against the scalar kernels within 1e-5 relative; multithreaded and + single-threaded artifacts agree. +. (Tier 2) `RequantizeTo(Q8_0)` through the SafeTensors loader round-trips through the Q8_0 decoder + within the format's quantization error; the in-memory and CLI paths produce identical bytes. + +=== Non-Functional + +* No change to kernel keys, golden tests or published API on the JVM and Apple tiers. +* Each tier is measured with `wasm-bench` before and after; the numbers go into this page. +* Acceptance for the whole proposal: a Laya decision under 1 s in the browser on a 4-core laptop + with the multithreaded sidecar and at most 2 s single-threaded; peak tab memory under 1 GB with + Q8_0 weights; KernelRace web decode at least 5× today's tokens per second. + +== Compatibility and Migration + +Tier 0 is additive. The only observable change is numeric order on packed weights on non-JVM +targets, where the scalar feed-order kernel replaces the per-element reference decode; the +tolerance-based parity tests cover it, and the JVM is unaffected because it overrides the path. +Tier 1 is a new module that must be installed explicitly on wasm (no `ServiceLoader`); without it +every path falls back to Tier 0. Tier 2 widens loader signatures with defaulted parameters and is +marked breaking only for implementors of `ParametersLoader`. + +== Rollout Plan + +. PR 0 (this page, Tier 0): dispatch order, direct dense path, narrow-float rows, in-memory source, + Blob bulk copy, `wasm-bench`. +. PR 1 (Tier 2, engine): `Q8_0Quantizer`, `Requantizers`, SafeTensors and GGUF `RequantizeTo`, + `PlannerProfile.quantizeDenseTo`, the `skainet-quantize` tool in `skainet-tensor-tools`. +. PR 2 (Tier 1): the sidecar build, `skainet-backend-wasm-simd`, `WasmCpuOps`, parity and bench. +. PR 3 (Tier 1 follow-up): resident activations, banded attention, pthread tuning — scoped by PR 2's + measurements. +. Engine release; then SKaiNET-transformers: a common-code `LayaModelLoader` over a + `RandomAccessSource`, a packed-aware `LayaWeightPreparation.splitRows`, a decision-quality gate + for the quantized lane; then the SemanticEmojiFilter web verifier. + +== Measurements + +`wasm-bench` on Node (Kotlin/Wasm production binary, 12-core laptop, one thread), Tier 0 code, +2026-10-16 — the baseline the following tiers are measured against. `m` is the token count of one +Laya decision, one LEAF query, or one decode step: + +[cols="2,1,1,1,1,1",options="header"] +|=== +| Case | m × n × k | dense FP32 ms | dense GFLOPS | Q8_0 (scalar) ms | Q8_0 GFLOPS + +| laya.Wqkv | 150 × 2304 × 768 | 635 | 0.84 | 1420 | 0.37 +| laya.Wi | 150 × 2304 × 768 | 430 | 1.24 | 1223 | 0.43 +| laya.mlpWo | 150 × 768 × 1152 | 158 | 1.68 | 599 | 0.44 +| laya.Wo | 150 × 768 × 768 | 114 | 1.55 | 330 | 0.54 +| leaf.ffn | 20 × 1536 × 384 | 21 | 1.13 | 61 | 0.38 +| leaf.proj | 20 × 384 × 384 | 3.0 | 1.96 | 9.6 | 0.61 +| smollm2.gemv | 1 × 1536 × 576 | 5.0 | 0.36 | 5.1 | 0.35 +| qwen06b.gemv | 1 × 3072 × 1024 | 22 | 0.28 | 9.9 | 0.63 +|=== + +Read across one Laya layer (Wqkv + Wi + mlpWo + Wo) this is about 1.34 s dense, so roughly 29 s for +the 22-layer encoder at 150 tokens, matching the estimate in <<_origin>>. The scalar Q8_0 kernels +are 2–3× slower than dense FP32 at this size: they decode per block in Kotlin, which is the gap the +SIMD128 sidecar (Tier 1) closes, and why quantize-at-load (Tier 2) is a memory measure first and a +speed measure only together with Tier 1. + +== Acceptance Criteria + +* `wasm-bench` numbers recorded in <<_measurements>> for Tier 0 (done), Tier 1 and Tier 2. +* `kernel-support-matrix.adoc` shows the browser pack's formats and the drift test passes. +* The sample measurements in <<_who_pays_today>> are updated from a real browser run. + +== Risks + +* Copy overhead across the Kotlin/Wasm ↔ sidecar boundary (about 440 MB of activation copies per + Laya decision at 150 tokens if nothing stays resident) — measured in PR 2 before PR 3 decides. +* Emscripten in CI and reproducible `.wasm` artifacts; pthreads need cross-origin isolation, which + GitHub Pages cannot set in headers (the examples site already uses a service-worker shim). +* Quantization changes Laya's decisions; the consumer gate is decision agreement, not FP32 parity. +* wasm32 caps linear memory at 4 GB across the Kotlin heap and the sidecar: Q8_0 Laya fits, dense + FP32 does not, so Tier 2 is a prerequisite for that model rather than an option. + +== Open Questions + +* Whether the dense FP32 fast path in `matmul` (the `[k, n]` loop) should also route through + `KernelDispatch` first, so a registered dense kernel wins without an ops override. Deferred to + PR 2, where the first such kernel appears. +* Whether `TensorOps` should grow a fused `layerNorm`; the cost on the encoder is allocation rather + than arithmetic, and `ForwardScope` may already answer it. + +== References + +* `docs/modules/ROOT/pages/explanation/kernel-selection.adoc`, `eager-execution.adoc`, + `reference/kernel-support-matrix.adoc` +* SKEEP-003 (unified tensor storage; synchronous `RandomAccessSource` on JS/Wasm), SKEEP-005 (schedules + on single-threaded runtimes), issues #288, #973, #1094, #1096, #1109, #1124, #1150, #1275 +* llama.cpp pull request 11453 ("ggml : x2 speed for WASM by optimizing SIMD") diff --git a/docs/modules/skeep/pages/index.adoc b/docs/modules/skeep/pages/index.adoc index 80787fde4..b33d4d0c5 100644 --- a/docs/modules/skeep/pages/index.adoc +++ b/docs/modules/skeep/pages/index.adoc @@ -134,4 +134,8 @@ Every proposal should include: | xref:skeep:006-input-frontends-and-model-identity.adoc[SKEEP-006] | Implemented | Input front ends and model identity — one resolvable preprocessing contract (tokenizers first) keyed on a string-based family registry + +| xref:skeep:007-browser-lane.adoc[SKEEP-007] +| Draft +| The browser lane — packed SIMD128 kernels as a sidecar module, quantize-at-load for SafeTensors, and the common-code fixes that keep wasm off the reference kernel |=== diff --git a/settings.gradle.kts b/settings.gradle.kts index ce788d5e9..f9e2bf9f4 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -46,6 +46,7 @@ include("skainet-backends:skainet-backend-jni-cpu") // ====== BENCHMARKS include("skainet-backends:benchmarks:jvm-cpu-jmh") include("skainet-backends:benchmarks:jvm-cpu-publish") +include("skainet-backends:benchmarks:wasm-bench") // ====== DATA include("skainet-data:skainet-data-api") diff --git a/skainet-apps/skainet-tensor-tools/build.gradle.kts b/skainet-apps/skainet-tensor-tools/build.gradle.kts index c042e56ab..efe3d1bd1 100644 --- a/skainet-apps/skainet-tensor-tools/build.gradle.kts +++ b/skainet-apps/skainet-tensor-tools/build.gradle.kts @@ -17,6 +17,8 @@ dependencies { implementation(project(":skainet-io:skainet-io-onnx")) implementation(project(":skainet-io:skainet-io-image")) implementation(project(":skainet-io:skainet-io-gguf")) + implementation(project(":skainet-io:skainet-io-safetensors")) + implementation(libs.kotlinx.io.core) implementation(project(":skainet-backends:skainet-backend-cpu")) implementation(project(":skainet-models:skainet-model-yolo")) implementation(libs.pbandk.runtime) @@ -34,3 +36,13 @@ application { tasks.test { useJUnitPlatform() } + +// SKEEP-007 §T2: SafeTensors → quantized GGUF, offline. +// ./gradlew :skainet-apps:skainet-tensor-tools:quantize --args="--in model.safetensors --out model.q8_0.gguf --keep tok_embeddings" +tasks.register("quantize") { + group = "application" + description = "Quantize a SafeTensors checkpoint's weight matrices into a GGUF (Q8_0 or Q4_0)" + classpath = sourceSets["main"].runtimeClasspath + mainClass.set("sk.ainet.apps.tool.quantize.QuantizeCliKt") + jvmArgs("--add-modules=jdk.incubator.vector") +} diff --git a/skainet-apps/skainet-tensor-tools/src/main/kotlin/sk/ainet/apps/tool/quantize/QuantizeCli.kt b/skainet-apps/skainet-tensor-tools/src/main/kotlin/sk/ainet/apps/tool/quantize/QuantizeCli.kt new file mode 100644 index 000000000..162217e23 --- /dev/null +++ b/skainet-apps/skainet-tensor-tools/src/main/kotlin/sk/ainet/apps/tool/quantize/QuantizeCli.kt @@ -0,0 +1,158 @@ +package sk.ainet.apps.tool.quantize + +import kotlinx.cli.ArgParser +import kotlinx.cli.ArgType +import kotlinx.cli.default +import kotlinx.coroutines.runBlocking +import kotlinx.io.buffered +import kotlinx.io.files.Path +import kotlinx.io.files.SystemFileSystem +import sk.ainet.context.DefaultDataExecutionContext +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.io.gguf.GGMLQuantizationType +import sk.ainet.io.gguf.export.GGUFWriter +import sk.ainet.io.gguf.export.GgufTensorEntry +import sk.ainet.io.gguf.export.GgufWriteRequest +import sk.ainet.io.safetensors.NarrowFloatLoadPolicy +import sk.ainet.io.safetensors.SafeTensorsParametersLoader +import sk.ainet.lang.memory.plan.EncodingRequest +import sk.ainet.lang.memory.plan.WeightForm +import sk.ainet.lang.memory.plan.WeightShapeOrientation +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.Bf16DenseTensorData +import sk.ainet.lang.tensor.data.Fp16DenseTensorData +import sk.ainet.lang.tensor.data.Requantizers +import sk.ainet.lang.tensor.storage.PackedBlockStorage +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.FP32 +import java.nio.file.Files + +/** + * `skainet-quantize`: a SafeTensors checkpoint → a GGUF with its weight matrices quantized (SKEEP-007 §T2). + * + * The engine can quantize at load (`WeightForm(RequantizeTo(Q8_0))` on the SafeTensors loaders), but a + * browser still has to *download* the file first, and F16 is twice the bytes of Q8_0. This tool runs the + * same quantizers once, offline, and writes a GGUF the engine's streaming loader serves packed on every + * target. Tensor names are kept exactly as the SafeTensors file has them, so a family's name resolver + * works unchanged; tensors matched by `--keep` (name prefixes) stay dense, as F16 by default — the + * usual choice for a token-embedding table that an embedding layer reads row by row. + * + * Usage: + * skainet-quantize --in model.safetensors --out model.q8_0.gguf [--encoding q8_0|q4_0] [--keep tok_embeddings,lm_head] + */ +public fun main(args: Array) { + val parser = ArgParser("skainet-quantize") + val input by parser.option(ArgType.String, fullName = "in", shortName = "i", description = "SafeTensors file to read") + val output by parser.option(ArgType.String, fullName = "out", shortName = "o", description = "GGUF file to write") + val encodingName by parser.option(ArgType.String, fullName = "encoding", shortName = "e", description = "Target block encoding: q8_0 (default) or q4_0").default("q8_0") + val keep by parser.option(ArgType.String, fullName = "keep", shortName = "k", description = "Comma-separated tensor-name prefixes kept dense (F16)").default("") + val keepF32 by parser.option(ArgType.Boolean, fullName = "keep-f32", description = "Write kept float tensors as F32 instead of F16").default(false) + parser.parse(args) + require(!input.isNullOrBlank()) { "--in is required" } + require(!output.isNullOrBlank()) { "--out is required" } + val report = quantizeSafeTensorsToGguf( + input = input!!, + output = output!!, + encoding = parseEncoding(encodingName), + keepPrefixes = keep.split(',').map { it.trim() }.filter { it.isNotEmpty() }, + keepAsF32 = keepF32, + ) + println(report.render()) +} + +public fun parseEncoding(name: String): TensorEncoding = when (name.lowercase()) { + "q8_0" -> TensorEncoding.Q8_0 + "q4_0" -> TensorEncoding.Q4_0 + else -> throw IllegalArgumentException("unknown encoding '$name'; quantize-at-load supports ${Requantizers.supported.joinToString { it.name.lowercase() }}") +} + +public data class QuantizeReport( + val encoding: TensorEncoding, + val quantized: List, + val kept: List, + val inputBytes: Long, + val outputBytes: Long, +) { + public fun render(): String = buildString { + appendLine("skainet-quantize: ${quantized.size} tensors → ${encoding.name}, ${kept.size} kept dense") + appendLine(" input ${inputBytes / (1024 * 1024)} MB") + appendLine(" output ${outputBytes / (1024 * 1024)} MB") + if (kept.isNotEmpty()) appendLine(" kept: ${kept.joinToString()}") + } +} + +/** + * Quantize every rank-2 float tensor of [input] whose rows block by the encoding's block size to + * [encoding], keep the rest dense, and write a GGUF to [output]. Returns what happened. + */ +public fun quantizeSafeTensorsToGguf( + input: String, + output: String, + encoding: TensorEncoding = TensorEncoding.Q8_0, + keepPrefixes: List = emptyList(), + keepAsF32: Boolean = false, +): QuantizeReport { + require(Requantizers.supports(encoding)) { "no quantizer for ${encoding.name}" } + val requantize = WeightForm(encoding = EncodingRequest.RequantizeTo(encoding), shape = WeightShapeOrientation.OUT_IN) + fun kept(name: String) = keepPrefixes.any { name.startsWith(it) } + + val ctx = DefaultDataExecutionContext() + val loaded = LinkedHashMap>() + runBlocking { + SafeTensorsParametersLoader( + sourceProvider = { JvmRandomAccessSource.open(input) }, + // Kept narrow floats stay narrow so they can be written back as F16/BF16 bytes verbatim. + bf16Policy = NarrowFloatLoadPolicy.KEEP_NATIVE, + fp16Policy = NarrowFloatLoadPolicy.KEEP_NATIVE, + weightFormFor = { name -> if (kept(name)) null else requantize }, + ).load(ctx, FP32::class) { name, tensor -> loaded[name] = tensor } + } + + val quantized = ArrayList() + val keptNames = ArrayList() + val entries = loaded.map { (name, tensor) -> + val dims = tensor.shape.dimensions.toList() + when (val data = tensor.data) { + is PackedBlockStorage -> { + quantized += name + GgufTensorEntry(ggufName = name, quantization = ggmlTypeOf(data.encoding), shape = dims, rawBytes = data.packedData) + } + is Fp16DenseTensorData -> { + keptNames += name + if (keepAsF32) GgufTensorEntry(ggufName = name, tensor = tensor, quantization = GGMLQuantizationType.F32, shape = dims) + else GgufTensorEntry(ggufName = name, quantization = GGMLQuantizationType.F16, shape = dims, rawBytes = data.packedData) + } + is Bf16DenseTensorData -> { + keptNames += name + if (keepAsF32) GgufTensorEntry(ggufName = name, tensor = tensor, quantization = GGMLQuantizationType.F32, shape = dims) + else GgufTensorEntry(ggufName = name, quantization = GGMLQuantizationType.BF16, shape = dims, rawBytes = data.packedData) + } + else -> { + keptNames += name + val type = if (keepAsF32) GGMLQuantizationType.F32 else GGMLQuantizationType.F16 + GgufTensorEntry(ggufName = name, tensor = tensor, quantization = type, shape = dims) + } + } + } + val metadata = linkedMapOf( + "general.name" to input.substringAfterLast('/').substringBeforeLast('.'), + "skainet.quantize.source" to "safetensors", + "skainet.quantize.encoding" to encoding.name, + ) + val request = GgufWriteRequest(metadata = metadata, tensors = entries, tensorMap = entries.associate { it.ggufName to it.ggufName }) + val outPath = Path(output) + SystemFileSystem.sink(outPath).buffered().use { sink -> GGUFWriter.writeToSink(request, sink) } + return QuantizeReport( + encoding = encoding, + quantized = quantized, + kept = keptNames, + inputBytes = Files.size(java.nio.file.Path.of(input)), + outputBytes = SystemFileSystem.metadataOrNull(outPath)?.size ?: -1L, + ) +} + +private fun ggmlTypeOf(encoding: TensorEncoding): GGMLQuantizationType = when (encoding) { + TensorEncoding.Q8_0 -> GGMLQuantizationType.Q8_0 + TensorEncoding.Q4_0 -> GGMLQuantizationType.Q4_0 + else -> throw IllegalArgumentException("${encoding.name} has no GGUF tensor type") +} diff --git a/skainet-apps/skainet-tensor-tools/src/test/kotlin/sk/ainet/apps/tool/quantize/QuantizeCliTest.kt b/skainet-apps/skainet-tensor-tools/src/test/kotlin/sk/ainet/apps/tool/quantize/QuantizeCliTest.kt new file mode 100644 index 000000000..4efe23ddb --- /dev/null +++ b/skainet-apps/skainet-tensor-tools/src/test/kotlin/sk/ainet/apps/tool/quantize/QuantizeCliTest.kt @@ -0,0 +1,74 @@ +package sk.ainet.apps.tool.quantize + +import kotlinx.coroutines.runBlocking +import kotlinx.io.Buffer +import kotlinx.io.readByteArray +import sk.ainet.context.DefaultDataExecutionContext +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.io.gguf.StreamingGgufParametersLoader +import sk.ainet.io.safetensors.SafeTensorsWriter +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.Fp16DenseTensorData +import sk.ainet.lang.tensor.data.Q8_0BlockTensorData +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.FP32 +import java.nio.file.Files +import kotlin.math.abs +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertIs +import kotlin.test.assertTrue + +/** SKEEP-007 §T2: the offline converter writes a GGUF the streaming loader serves packed, names and shapes intact. */ +class QuantizeCliTest { + private fun pseudo(n: Int, seed: Int): FloatArray { + var s = seed + return FloatArray(n) { + s = s * 1103515245 + 12345 + ((s ushr 16) % 2001 - 1000) / 500f + } + } + + @Test + fun `safetensors to q8_0 gguf round trip`() { + val w = pseudo(6 * 64, 1) + val emb = pseudo(10 * 32, 2) + val bias = pseudo(6, 3) + val buffer = Buffer() + SafeTensorsWriter.write(buffer) { + tensorF32("layers.0.weight", intArrayOf(6, 64), w) + tensorF32("tok_embeddings.weight", intArrayOf(10, 32), emb) + tensorF32("layers.0.bias", intArrayOf(6), bias) + } + val input = Files.createTempFile("quantize-in", ".safetensors").toFile().also { it.deleteOnExit() } + input.writeBytes(buffer.readByteArray()) + val output = Files.createTempFile("quantize-out", ".gguf").toFile().also { it.deleteOnExit() } + + val report = quantizeSafeTensorsToGguf(input.path, output.path, TensorEncoding.Q8_0, keepPrefixes = listOf("tok_embeddings")) + assertEquals(listOf("layers.0.weight"), report.quantized) + assertEquals(setOf("tok_embeddings.weight", "layers.0.bias"), report.kept.toSet()) + assertTrue(report.outputBytes > 0) + + val ctx = DefaultDataExecutionContext() + val loaded = LinkedHashMap>() + runBlocking { + StreamingGgufParametersLoader(sourceProvider = { JvmRandomAccessSource.open(output.path) }, keepF16Native = true) + .load(ctx, FP32::class) { name, tensor -> loaded[name] = tensor } + } + assertEquals(setOf("layers.0.weight", "tok_embeddings.weight", "layers.0.bias"), loaded.keys) + + val packed = assertIs(loaded.getValue("layers.0.weight").data) + assertEquals(listOf(6, 64), packed.shape.dimensions.toList()) + val back = packed.toFloatArray() + val amax = w.maxOf { abs(it) } + for (i in w.indices) assertTrue(abs(w[i] - back[i]) <= amax / 127f + 1e-4f, "w[$i]") + + val kept = assertIs(loaded.getValue("tok_embeddings.weight").data) + assertEquals(listOf(10, 32), kept.shape.dimensions.toList()) + val embBack = kept.copyToFloatArray() + for (i in emb.indices) assertTrue(abs(emb[i] - embBack[i]) <= abs(emb[i]) * 1e-3f + 1e-3f, "emb[$i]") + + val biasBack = loaded.getValue("layers.0.bias").data.copyToFloatArray() + for (i in bias.indices) assertTrue(abs(bias[i] - biasBack[i]) <= abs(bias[i]) * 1e-3f + 1e-3f, "bias[$i]") + } +} diff --git a/skainet-backends/benchmarks/wasm-bench/build.gradle.kts b/skainet-backends/benchmarks/wasm-bench/build.gradle.kts new file mode 100644 index 000000000..3c01350d3 --- /dev/null +++ b/skainet-backends/benchmarks/wasm-bench/build.gradle.kts @@ -0,0 +1,29 @@ +import org.jetbrains.kotlin.gradle.ExperimentalWasmDsl + +// SKEEP-007 §T0.5: the wasm-lane baseline. A Kotlin/Wasm executable that times the matmul shapes of +// the shipped models through the eager CPU backend and prints one `[SKAINET_PERF_WASM]` line per +// case. Not published. +// +// ./gradlew :skainet-backends:benchmarks:wasm-bench:wasmJsNodeDevelopmentRun # quick, unoptimised +// ./gradlew :skainet-backends:benchmarks:wasm-bench:wasmJsNodeProductionRun # the numbers to record +// ./gradlew :skainet-backends:benchmarks:wasm-bench:wasmJsBrowserProductionRun # same in Chrome +plugins { + alias(libs.plugins.kotlinMultiplatform) +} + +kotlin { + @OptIn(ExperimentalWasmDsl::class) + wasmJs { + nodejs() + browser() + binaries.executable() + } + + sourceSets { + wasmJsMain.dependencies { + implementation(project(":skainet-lang:skainet-lang-core")) + implementation(project(":skainet-backends:skainet-backend-api")) + implementation(project(":skainet-backends:skainet-backend-cpu")) + } + } +} diff --git a/skainet-backends/benchmarks/wasm-bench/src/wasmJsMain/kotlin/sk/ainet/bench/wasm/Main.kt b/skainet-backends/benchmarks/wasm-bench/src/wasmJsMain/kotlin/sk/ainet/bench/wasm/Main.kt new file mode 100644 index 000000000..7bfd1f404 --- /dev/null +++ b/skainet-backends/benchmarks/wasm-bench/src/wasmJsMain/kotlin/sk/ainet/bench/wasm/Main.kt @@ -0,0 +1,101 @@ +package sk.ainet.bench.wasm + +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.Q8_0BlockTensorData +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.tensor.matmulWeightTransposed +import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.Fp16Codec +import kotlin.math.abs +import kotlin.math.roundToInt +import kotlin.time.TimeSource + +/** + * SKEEP-007 wasm baseline: the matmul shapes one decision / query / decode step of the shipped + * models puts through `matmulWeightTransposed`, dense FP32 and packed Q8_0, timed in the eager + * backend exactly as the samples run it. Output is one line per case: + * + * `[SKAINET_PERF_WASM] m= n= k= form= ms= gflops=` + * + * Every later tier (SIMD128 sidecar, quantize-at-load) is measured against these lines. + */ +private data class Case(val name: String, val m: Int, val n: Int, val k: Int) + +private val cases = listOf( + Case("laya.Wqkv", 150, 2304, 768), + Case("laya.Wi", 150, 2304, 768), + Case("laya.mlpWo", 150, 768, 1152), + Case("laya.Wo", 150, 768, 768), + Case("leaf.ffn", 20, 1536, 384), + Case("leaf.proj", 20, 384, 384), + Case("smollm2.gemv", 1, 1536, 576), + Case("qwen06b.gemv", 1, 3072, 1024), +) + +private fun pseudo(n: Int, seed: Int): FloatArray { + var s = seed + return FloatArray(n) { + s = s * 1103515245 + 12345 + ((s ushr 16) % 2001 - 1000) / 1000f + } +} + +/** Row-major Q8_0 blocks (fp16 scale + 32 int8), the GGUF layout, from dense rows. */ +@Suppress("UNCHECKED_CAST") +private fun q8_0Weight(ctx: DirectCpuExecutionContext, n: Int, k: Int, seed: Int): Tensor { + require(k % 32 == 0) + val values = pseudo(n * k, seed) + val blocksPerRow = k / 32 + val bytes = ByteArray(n * blocksPerRow * 34) + for (r in 0 until n) for (b in 0 until blocksPerRow) { + val base = (r * blocksPerRow + b) * 34 + val off = r * k + b * 32 + var amax = 0f + for (i in 0 until 32) amax = maxOf(amax, abs(values[off + i])) + val d = amax / 127f + val id = if (d != 0f) 1f / d else 0f + val bits = Fp16Codec.encode(d) + bytes[base] = (bits and 0xFF).toByte() + bytes[base + 1] = ((bits ushr 8) and 0xFF).toByte() + for (i in 0 until 32) bytes[base + 2 + i] = (values[off + i] * id).roundToInt().coerceIn(-127, 127).toByte() + } + return ctx.fromData(Q8_0BlockTensorData(Shape(n, k), bytes) as TensorData, FP32::class) +} + +private fun timeMedianMs(reps: Int, block: () -> Unit): Double { + val clock = TimeSource.Monotonic + val samples = DoubleArray(reps) { + val t0 = clock.markNow() + block() + t0.elapsedNow().inWholeMicroseconds / 1000.0 + } + samples.sort() + return samples[reps / 2] +} + +private fun report(case: Case, form: String, ms: Double) { + val flops = 2.0 * case.m * case.n * case.k + val gflops = flops / (ms / 1000.0) / 1e9 + val msText = ((ms * 100).roundToInt() / 100.0).toString() + val gfText = ((gflops * 100).roundToInt() / 100.0).toString() + println("[SKAINET_PERF_WASM] ${case.name} m=${case.m} n=${case.n} k=${case.k} form=$form ms=$msText gflops=$gfText") +} + +fun main() { + val ctx = DirectCpuExecutionContext() + println("[SKAINET_PERF_WASM] backend=${ctx.ops::class.simpleName} cases=${cases.size}") + for (case in cases) { + val x = ctx.fromFloatArray(Shape(case.m, case.k), FP32::class, pseudo(case.m * case.k, 3)) + val dense = ctx.fromFloatArray(Shape(case.n, case.k), FP32::class, pseudo(case.n * case.k, 7)) + val packed = q8_0Weight(ctx, case.n, case.k, 11) + // Warm-up: first call pays the one-time weight relayout / caches. + ctx.ops.matmulWeightTransposed(x, dense) + ctx.ops.matmulWeightTransposed(x, packed) + val reps = if (case.m == 1) 50 else 5 + report(case, "dense", timeMedianMs(reps) { ctx.ops.matmulWeightTransposed(x, dense) }) + report(case, "q8_0", timeMedianMs(reps) { ctx.ops.matmulWeightTransposed(x, packed) }) + } + println("[SKAINET_PERF_WASM] done") +} diff --git a/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api b/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api index 9fc00dfb8..e0b809808 100644 --- a/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api +++ b/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api @@ -332,6 +332,7 @@ public class sk/ainet/exec/tensor/ops/DefaultCpuOpsBase : sk/ainet/context/sched public fun permute (Lsk/ainet/lang/tensor/Tensor;[I)Lsk/ainet/lang/tensor/Tensor; public fun pow (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun powScalar (Lsk/ainet/lang/tensor/Tensor;Ljava/lang/Number;)Lsk/ainet/lang/tensor/Tensor; + protected fun preferDirectWeightTransposed ()Z public fun rdivScalar (Ljava/lang/Number;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun relayoutPackedWeightForKernels (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun relu (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; diff --git a/skainet-backends/skainet-backend-cpu/src/appleMain/kotlin/sk/ainet/exec/tensor/ops/AccelerateCpuOps.kt b/skainet-backends/skainet-backend-cpu/src/appleMain/kotlin/sk/ainet/exec/tensor/ops/AccelerateCpuOps.kt index f7f2ec7bf..0abcd4464 100644 --- a/skainet-backends/skainet-backend-cpu/src/appleMain/kotlin/sk/ainet/exec/tensor/ops/AccelerateCpuOps.kt +++ b/skainet-backends/skainet-backend-cpu/src/appleMain/kotlin/sk/ainet/exec/tensor/ops/AccelerateCpuOps.kt @@ -42,6 +42,9 @@ public class AccelerateCpuOps( if (schedule === this.schedule) this else AccelerateCpuOps(dataFactory, schedule) + /** Accelerate's `[k, n]` GEMM beats a scalar row-dot, so keep the cached transpose (SKEEP-007 §T0.2). */ + override fun preferDirectWeightTransposed(): Boolean = false + // ── matmul ────────────────────────────────────────────────────────── override fun matmul(a: Tensor, b: Tensor): Tensor { diff --git a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt index 9555039d9..db6da6530 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt @@ -1049,17 +1049,112 @@ public open class DefaultCpuOpsBase( @Suppress("UNCHECKED_CAST") override fun matmulWeightTransposed(x: Tensor, weight: Tensor): Tensor { if (weight.shape.rank != 2 || !isHeapPackedWeight(weight.data)) { + // Dense weight. A heap FP32 `[out, in]` matrix is already the layout a row-dot wants — + // every output is one contiguous dot product — so a tier without a vectorised `[k, n]` + // kernel computes the product in place (SKEEP-007 §T0.2) instead of paying a scattered + // transpose copy per weight and caching it (~0.5 GB for a 22-layer encoder). The JVM + // and Apple tiers keep the cached transpose: their `[k, n]` kernels are the fast path. + if (preferDirectWeightTransposed()) { + matmulDenseWeightTransposedDirect(x, weight)?.let { return it } + } return matmul(x, transposedDenseWeight(weight) ?: transpose(weight)) } - // Decode the weight where it lies (#1124). The relayout below produces bytes for kernels - // that address `packedData` in feed order themselves; this implementation has no such - // kernel, so relayouting for it was pure harm — the result is a tensor whose shape says - // [in, out] while its blocks still run along the original input dimension, which no block - // order can describe, and decoding it read the wrong blocks and returned plausible garbage. - // `KernelDispatch.matmul` already wants the weight output-major, which is exactly the shape - // this weight has, so the canonical view goes straight in with nothing rearranged. + // Packed weight (SKEEP-007 §T0.1), in order of preference: + // 1. a kernel registered for the canonical row-major view — the JVM FFM, Android JNI and + // browser packs read that order natively (#1124: decode the weight where it lies); + // 2. the one-time, cached relayout into kernel feed order, which the scalar packed kernels + // every target carries can read — this is what [prepackedWeights] was declared for; + // 3. the decoding reference kernel through views — correct for any encoding, ~1000× slower, + // and the last resort rather than the default. + // Before this ordering a canonical GGUF weight on a target without a row-major pack (every + // non-JVM target) took path 3 on every call. + if (hasExactKernelForCanonical(x, weight)) { + matmulWeightTransposedViaViews(x, weight)?.let { return it } + } + // Relayout only when a feed-order kernel will consume the relaid bytes: the relabelled + // `[in, out]` tensor is meaningless to anything that decodes it by shape (#1124), so with an + // empty registry the canonical view goes to the decoding reference kernel instead. + if (feedOrderKernelAvailable(weight.data)) { + rewrapFeedOrderWeight(weight)?.let { return matmul(x, it) } + cachedKernelOrderWeight(weight)?.let { return matmul(x, it) } + } matmulWeightTransposedViaViews(x, weight)?.let { return it } - return matmul(x, transposePackedWeight(weight) ?: return matmulGeneric(x, transpose(weight))) + return matmulGeneric(x, transpose(weight)) + } + + /** `true` when [KernelRegistry] supplies a feed-order kernel for this packed format (what `chooseQuantizedMatmulHeap` dispatches to). */ + private fun feedOrderKernelAvailable(data: sk.ainet.lang.tensor.data.TensorData<*, *>): Boolean = when (data) { + is Q5_1TensorData -> q5_1Kernel != null + is Q5_0TensorData -> q5_0Kernel != null + is Q4_KTensorData -> q4kKernel != null + is Q5_KTensorData -> q5kKernel != null + is Q6_KTensorData -> q6kKernel != null + is Q8_0TensorData -> q8_0Kernel != null + is Q4_0TensorData -> q4_0Kernel != null + else -> false + } + + /** + * Whether a dense `[out, in]` FP32 weight is multiplied in place by [matmulDenseWeightTransposedDirect] + * (`true`, the common default) or transposed once and handed to `matmul` (the JVM and Apple + * tiers, whose `[k, n]` kernels are vectorised). SKEEP-007 §T0.2. + */ + protected open fun preferDirectWeightTransposed(): Boolean = true + + /** `true` when [KernelDispatch] holds an exact kernel for [x] × the canonical view of [weight]. */ + private fun hasExactKernelForCanonical(x: Tensor<*, *>, weight: Tensor<*, *>): Boolean { + if (!sk.ainet.backend.api.kernel.DispatchMode.useRegistry()) return false + if (x.dtype != FP32::class) return false + val xView = x.data.view ?: return false + val wView = weight.data.view ?: return false + val xNorm = try { + sk.ainet.backend.api.kernel.KernelDispatch.normalizeActivation(xView).first + } catch (_: IllegalArgumentException) { + return false + } + sk.ainet.backend.api.kernel.KernelDispatch.ensureInstalled() + return sk.ainet.backend.api.kernel.KernelDispatch.find( + sk.ainet.backend.api.kernel.KernelKey.matmul(xNorm, wView), + ) != null + } + + /** + * The feed-order relayout of a canonical packed [weight], computed once and cached by the + * identity of its bytes ([prepackedWeights]); `null` for an encoding without a relayout. + */ + @Suppress("UNCHECKED_CAST") + private fun cachedKernelOrderWeight(weight: Tensor): Tensor? { + val packed = weight.data as? sk.ainet.lang.tensor.storage.PackedBlockStorage ?: return null + val source = packed.packedData + prepackedWeights.firstOrNull { it.first === source }?.let { return it.second as Tensor } + val relayouted = transposePackedWeight(weight) ?: return null + if (prepackedWeights.size >= PREPACK_CACHE_LIMIT) prepackedWeights.removeAt(0) + prepackedWeights.add(Pair(source, relayouted as Tensor<*, *>)) + return relayouted + } + + /** + * `x · Wᵀ` for a heap FP32 `[out, in]` weight computed straight from the weight's own rows + * (SKEEP-007 §T0.2): every output is a contiguous dot product, four weight rows and four + * activation rows per tile so the weight is streamed once per four activation rows instead of + * once per row. Each output accumulates its products in ascending `p` in one accumulator, so the + * result is bit-identical to the transpose-then-`matmul` path it replaces. `null` when either + * operand is not heap FP32. + */ + @Suppress("UNCHECKED_CAST") + private fun matmulDenseWeightTransposedDirect(x: Tensor, weight: Tensor): Tensor? { + if (x.dtype != FP32::class || weight.dtype != FP32::class || x.rank < 1) return null + val wWin = floatWindowOf(weight) ?: return null + val xWin = floatWindowOf(x) ?: return null + val n = weight.shape[0] + val k = weight.shape[1] + if (x.shape.dimensions.last() != k) return null + val m = if (k == 0) 0 else x.shape.volume / k + val out = FloatArray(m * n) + KernelProfile.timeFp32 { denseRowDotTiled(xWin.arr, xWin.off, wWin.arr, wWin.off, out, m, n, k) } + val outShape = if (x.rank == 1) Shape(n) else Shape(*(x.shape.dimensions.copyOf(x.rank - 1) + n)) + val outData = dataFactory.adoptFloatArray(outShape, x.dtype, out) as sk.ainet.lang.tensor.data.TensorData + return newTensor(outData, x.dtype, x, weight) } /** @@ -3935,3 +4030,79 @@ public class DefaultCpuOps( override fun withSchedule(schedule: sk.ainet.context.schedule.Schedule): TensorOps = if (schedule === this.schedule) this else DefaultCpuOps(dataFactory, schedule) } + +/** + * `out[i, j] = Σ_p a[i, p] · w[j, p]` over row-major `a` (`m × k`) and `w` (`n × k`), 4 × 4 tiles with + * one sequential accumulator per output (SKEEP-007 §T0.2). Pure common code: the scalar baseline on + * every target, and the reference the browser SIMD pack is checked against. + */ +internal fun denseRowDotTiled( + a: FloatArray, aOff: Int, + w: FloatArray, wOff: Int, + out: FloatArray, + m: Int, n: Int, k: Int, +) { + var i = 0 + while (i + 4 <= m) { + val a0 = aOff + i * k; val a1 = a0 + k; val a2 = a1 + k; val a3 = a2 + k + val o0 = i * n; val o1 = o0 + n; val o2 = o1 + n; val o3 = o2 + n + var j = 0 + while (j + 4 <= n) { + val w0 = wOff + j * k; val w1 = w0 + k; val w2 = w1 + k; val w3 = w2 + k + var s00 = 0f; var s01 = 0f; var s02 = 0f; var s03 = 0f + var s10 = 0f; var s11 = 0f; var s12 = 0f; var s13 = 0f + var s20 = 0f; var s21 = 0f; var s22 = 0f; var s23 = 0f + var s30 = 0f; var s31 = 0f; var s32 = 0f; var s33 = 0f + for (p in 0 until k) { + val b0 = w[w0 + p]; val b1 = w[w1 + p]; val b2 = w[w2 + p]; val b3 = w[w3 + p] + val x0 = a[a0 + p] + s00 += x0 * b0; s01 += x0 * b1; s02 += x0 * b2; s03 += x0 * b3 + val x1 = a[a1 + p] + s10 += x1 * b0; s11 += x1 * b1; s12 += x1 * b2; s13 += x1 * b3 + val x2 = a[a2 + p] + s20 += x2 * b0; s21 += x2 * b1; s22 += x2 * b2; s23 += x2 * b3 + val x3 = a[a3 + p] + s30 += x3 * b0; s31 += x3 * b1; s32 += x3 * b2; s33 += x3 * b3 + } + out[o0 + j] = s00; out[o0 + j + 1] = s01; out[o0 + j + 2] = s02; out[o0 + j + 3] = s03 + out[o1 + j] = s10; out[o1 + j + 1] = s11; out[o1 + j + 2] = s12; out[o1 + j + 3] = s13 + out[o2 + j] = s20; out[o2 + j + 1] = s21; out[o2 + j + 2] = s22; out[o2 + j + 3] = s23 + out[o3 + j] = s30; out[o3 + j + 1] = s31; out[o3 + j + 2] = s32; out[o3 + j + 3] = s33 + j += 4 + } + while (j < n) { + val wj = wOff + j * k + var s0 = 0f; var s1 = 0f; var s2 = 0f; var s3 = 0f + for (p in 0 until k) { + val b = w[wj + p] + s0 += a[a0 + p] * b; s1 += a[a1 + p] * b; s2 += a[a2 + p] * b; s3 += a[a3 + p] * b + } + out[o0 + j] = s0; out[o1 + j] = s1; out[o2 + j] = s2; out[o3 + j] = s3 + j++ + } + i += 4 + } + while (i < m) { + val ai = aOff + i * k + val oi = i * n + var j = 0 + while (j + 4 <= n) { + val w0 = wOff + j * k; val w1 = w0 + k; val w2 = w1 + k; val w3 = w2 + k + var s0 = 0f; var s1 = 0f; var s2 = 0f; var s3 = 0f + for (p in 0 until k) { + val x = a[ai + p] + s0 += x * w[w0 + p]; s1 += x * w[w1 + p]; s2 += x * w[w2 + p]; s3 += x * w[w3 + p] + } + out[oi + j] = s0; out[oi + j + 1] = s1; out[oi + j + 2] = s2; out[oi + j + 3] = s3 + j += 4 + } + while (j < n) { + val wj = wOff + j * k + var s = 0f + for (p in 0 until k) s += a[ai + p] * w[wj + p] + out[oi + j] = s + j++ + } + i++ + } +} diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/DenseWeightTransposedDirectTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/DenseWeightTransposedDirectTest.kt new file mode 100644 index 000000000..75e71b417 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/DenseWeightTransposedDirectTest.kt @@ -0,0 +1,88 @@ +package sk.ainet.exec.tensor.ops + +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.matmulWeightTransposed +import sk.ainet.lang.types.FP32 +import kotlin.test.Test +import kotlin.test.assertTrue +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals + +/** + * SKEEP-007 §T0.2: `x · Wᵀ` for a dense heap FP32 `[out, in]` weight computed from the weight's own rows. + * + * The tiled row-dot keeps one sequential accumulator per output, so it is bit-identical to the + * transpose-then-matmul path it replaces on the tiers that take it — and that is what these tests + * pin, for every tile remainder (m, n not multiples of 4) and for rank-1 and rank-3 activations. + */ +class DenseWeightTransposedDirectTest { + private val ctx = DirectCpuExecutionContext() + + private fun pseudo(n: Int, seed: Int): FloatArray { + var s = seed + return FloatArray(n) { + s = s * 1103515245 + 12345 + ((s ushr 16) % 2001 - 1000) / 250f + } + } + + /** `out[i, j] = Σ_p a[i, p] · w[j, p]`, ascending p, one accumulator — the definition. */ + private fun reference(a: FloatArray, w: FloatArray, m: Int, n: Int, k: Int): FloatArray { + val out = FloatArray(m * n) + for (i in 0 until m) for (j in 0 until n) { + var s = 0f + for (p in 0 until k) s += a[i * k + p] * w[j * k + p] + out[i * n + j] = s + } + return out + } + + @Test + fun `tiled row-dot is bit-identical to the sequential definition for every tile remainder`() { + for ((m, n, k) in listOf(Triple(1, 1, 1), Triple(4, 4, 8), Triple(5, 7, 33), Triple(9, 3, 64), Triple(2, 13, 5), Triple(8, 8, 0))) { + val a = pseudo(m * k, 7 + m) + val w = pseudo(n * k, 11 + n) + val out = FloatArray(m * n) { Float.NaN } + denseRowDotTiled(a, 0, w, 0, out, m, n, k) + assertContentEquals(reference(a, w, m, n, k), out, "m=$m n=$n k=$k") + } + } + + @Test + fun `tiled row-dot honours array offsets`() { + val m = 6; val n = 5; val k = 17 + val aPad = 3; val wPad = 5 + val a = pseudo(m * k, 3) + val w = pseudo(n * k, 5) + val aShifted = FloatArray(aPad) { 99f } + a + val wShifted = FloatArray(wPad) { -99f } + w + val out = FloatArray(m * n) + denseRowDotTiled(aShifted, aPad, wShifted, wPad, out, m, n, k) + assertContentEquals(reference(a, w, m, n, k), out) + } + + @Test + fun `the primitive agrees with transpose then matmul for rank-1 rank-2 and rank-3 activations`() { + val n = 7; val k = 12 + val w = ctx.fromFloatArray(Shape(n, k), FP32::class, pseudo(n * k, 21)) + val shapes = listOf(Shape(k), Shape(5, k), Shape(2, 3, k)) + for (shape in shapes) { + val x = ctx.fromFloatArray(shape, FP32::class, pseudo(shape.volume, 33)) + val viaPrimitive = ctx.ops.matmulWeightTransposed(x, w) + val viaTranspose = ctx.ops.matmul(x, ctx.ops.transpose(w)) + assertEquals(viaTranspose.shape, viaPrimitive.shape, "shape for x=$shape") + // Same products in the same order: bit-identical where Float is Float (JVM, Native, + // Wasm). Kotlin/JS accumulates a Float local in double precision while the transpose + // path rounds through a Float32Array per step, so this comparison is tolerance-based — + // the tiled-vs-definition tests above pin the bits where bits are stable. + val expected = viaTranspose.data.copyToFloatArray() + val actual = viaPrimitive.data.copyToFloatArray() + assertEquals(expected.size, actual.size, "size for x=$shape") + for (i in expected.indices) { + val tol = kotlin.math.abs(expected[i]) * 1e-5f + 1e-6f + assertTrue(kotlin.math.abs(expected[i] - actual[i]) <= tol, "x=$shape [$i]: ${actual[i]} vs ${expected[i]}") + } + } + } +} diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/PackedWeightKernelOrderTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/PackedWeightKernelOrderTest.kt new file mode 100644 index 000000000..681a13cba --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/PackedWeightKernelOrderTest.kt @@ -0,0 +1,78 @@ +package sk.ainet.exec.tensor.ops + +import sk.ainet.backend.api.kernel.KernelDispatch +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.lang.memory.trace.RecordingTraceSink +import sk.ainet.lang.memory.trace.TraceEvent +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.Q8_0BlockTensorData +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.tensor.matmulWeightTransposed +import sk.ainet.lang.types.FP32 +import kotlin.test.AfterTest +import kotlin.test.BeforeTest +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertTrue + +/** + * SKEEP-007 §T0.1: a canonical (row-major) packed weight must never be served by the decoding + * reference kernel when a scalar packed kernel for its format exists. + * + * Before the dispatch order change, every non-JVM target sent the canonical GGUF weight through + * `KernelDispatch` where only feed-order scalar kernels were registered, so the per-element + * reference decode ran on every call. Now the weight is relaid once into feed order and the scalar + * kernel runs; the reference kernel is recorded in the trace if it ever runs, and this test reads it. + */ +class PackedWeightKernelOrderTest { + private val ctx = DirectCpuExecutionContext() + private val sink = RecordingTraceSink() + private var previousSink = KernelDispatch.defaultSink + + @BeforeTest + fun attachSink() { + previousSink = KernelDispatch.defaultSink + KernelDispatch.defaultSink = sink + } + + @AfterTest + fun detachSink() { + KernelDispatch.defaultSink = previousSink + } + + private val outDim = 6 + private val inDim = 128 // four Q8_0 blocks per row + + @Suppress("UNCHECKED_CAST") + private fun canonicalWeight(): Tensor { + val blocks = outDim * (inDim / 32) + val bytes = ByteArray(blocks * 34) + var seed = 5 + for (b in 0 until blocks) { + val base = b * 34 + bytes[base] = 0x00; bytes[base + 1] = 0x3C // fp16 scale 1.0 + for (i in 0 until 32) { + seed = seed * 1103515245 + 12345 + bytes[base + 2 + i] = ((seed ushr 16) % 23 - 11).toByte() + } + } + return ctx.fromData(Q8_0BlockTensorData(Shape(outDim, inDim), bytes) as TensorData, FP32::class) + } + + @Test + fun `canonical Q8_0 weight is served without the reference kernel and matches the relayout path`() { + val w = canonicalWeight() + val x = ctx.fromFloatArray(Shape(3, inDim), FP32::class, FloatArray(3 * inDim) { ((it % 11) - 5) * 0.25f }) + + val viaPrimitive = ctx.ops.matmulWeightTransposed(x, w).data.copyToFloatArray() + // Second call exercises the per-weight cache. + val again = ctx.ops.matmulWeightTransposed(x, w).data.copyToFloatArray() + val viaRelayout = ctx.ops.matmul(x, ctx.ops.relayoutPackedWeightForKernels(w)).data.copyToFloatArray() + + assertContentEquals(viaRelayout, viaPrimitive) + assertContentEquals(viaPrimitive, again) + val referenceRuns = sink.eventsOf().filter { it.kernel.contains("reference", ignoreCase = true) } + assertTrue(referenceRuns.isEmpty(), "reference kernel ran: $referenceRuns") + } +} diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/TransposedWeightViewTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/TransposedWeightViewTest.kt index 4df2b8277..99ec92f89 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/TransposedWeightViewTest.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/TransposedWeightViewTest.kt @@ -186,11 +186,21 @@ class TransposedWeightViewTest { assertFalse(t.data is TransposedWeightTensorData<*, *>, "a dense transpose is a real transpose, not a marker") val x = activation(inDim) - assertContentEquals( - x.matmulWeightTransposed(dense).data.copyToFloatArray(), - x.matmul(t).data.copyToFloatArray(), - "dense: the two spellings already agreed and must keep agreeing", - ) + // The two spellings compute the same products in the same order; they are bit-identical + // wherever Float arithmetic is Float (JVM, Native, Wasm). On Kotlin/JS a Float local is a + // double, and the direct row-dot (SKEEP-007 §T0.2) keeps its accumulator in a local while + // the transpose-then-matmul path rounds through a Float32Array on every step — so JS + // differs in the last bits, as the golden tests already acknowledge for that target. + val viaPrimitive = x.matmulWeightTransposed(dense).data.copyToFloatArray() + val viaTranspose = x.matmul(t).data.copyToFloatArray() + assertEquals(viaTranspose.size, viaPrimitive.size) + for (i in viaTranspose.indices) { + val tol = kotlin.math.abs(viaTranspose[i]) * 1e-5f + 1e-6f + assertTrue( + kotlin.math.abs(viaTranspose[i] - viaPrimitive[i]) <= tol, + "dense[$i]: ${viaPrimitive[i]} vs ${viaTranspose[i]} — the two spellings must keep agreeing", + ) + } } @Test diff --git a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt index 41341bfdb..ee31af112 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt @@ -193,6 +193,9 @@ internal class DefaultCpuOpsJvm( * tensor nothing could decode correctly (#1124). Here the kernels do exist, so the relayout * pays: once per weight, cached by the identity of its bytes. */ + /** The JVM keeps the cached transpose: its vectorised `[k, n]` kernels are the fast path (SKEEP-007 §T0.2). */ + override fun preferDirectWeightTransposed(): Boolean = false + @Suppress("UNCHECKED_CAST") override fun matmulWeightTransposed(x: Tensor, weight: Tensor): Tensor { if (weight.shape.rank != 2 || !isHeapPackedWeightForJvm(weight.data)) { diff --git a/skainet-io/skainet-io-core/api/jvm/skainet-io-core.api b/skainet-io/skainet-io-core/api/jvm/skainet-io-core.api index cb1681061..8fbdb8bfc 100644 --- a/skainet-io/skainet-io-core/api/jvm/skainet-io-core.api +++ b/skainet-io/skainet-io-core/api/jvm/skainet-io-core.api @@ -7,6 +7,17 @@ public final class sk/ainet/io/ByteArrayMemoryChunk : sk/ainet/io/MemoryChunk { public fun slice (JJ)Lsk/ainet/io/MemoryChunk; } +public final class sk/ainet/io/ByteArrayRandomAccessSource : sk/ainet/io/RandomAccessSource { + public fun (Ljava/util/List;)V + public fun ([B)V + public fun close ()V + public fun getFilePath ()Ljava/lang/String; + public fun getSize ()J + public fun readAt (JI)[B + public fun readAt (J[BII)I + public fun readByteAt (J)B +} + public final class sk/ainet/io/FallbackMappedMemoryChunk : sk/ainet/io/MappedMemoryChunk { public static final field Companion Lsk/ainet/io/FallbackMappedMemoryChunk$Companion; public fun (Ljava/lang/String;J[BIJ)V diff --git a/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/ByteArrayRandomAccessSource.kt b/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/ByteArrayRandomAccessSource.kt new file mode 100644 index 000000000..f85ae51a5 --- /dev/null +++ b/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/ByteArrayRandomAccessSource.kt @@ -0,0 +1,96 @@ +package sk.ainet.io + +/** + * A [RandomAccessSource] over bytes already in memory (SKEEP-007 §T0.4). + * + * The browser has no file system to map or seek in: a checkpoint arrives as one download (or as + * bundled resource bytes) and has to be served to the synchronous readers — `StreamingGGUFReader`, + * the SafeTensors header parser, `StreamingGgufParametersLoader` — from memory. Every web sample + * so far wrote its own copy of this class; this is the engine's. + * + * The data is held as a list of [chunks] so a file larger than a single Kotlin array (`Int.MAX_VALUE` + * bytes, and 2 GB per allocation on the web) can still be addressed with a `Long` position. A + * single-array source is the common case and has the convenience constructor. + * + * Reads copy out of the chunks; nothing is retained beyond the chunks themselves, and [close] is + * a no-op — the caller owns the lifetime of the bytes. + */ +public class ByteArrayRandomAccessSource( + private val chunks: List, +) : RandomAccessSource { + + /** Single-array convenience: the whole file in one `ByteArray`. */ + public constructor(data: ByteArray) : this(listOf(data)) + + init { + require(chunks.isNotEmpty()) { "ByteArrayRandomAccessSource needs at least one chunk (an empty array is fine)" } + // Every chunk but the last must be the same length so a position maps to a chunk by division. + val lead = chunks.first().size + for (i in 0 until chunks.size - 1) { + require(chunks[i].size == lead && lead > 0) { + "ByteArrayRandomAccessSource: all chunks but the last must have the same non-zero length " + + "(chunk $i has ${chunks[i].size}, first has $lead)" + } + } + } + + private val chunkSize: Long = chunks.first().size.toLong() + + override val size: Long = chunks.fold(0L) { acc, c -> acc + c.size } + + override fun readAt(position: Long, length: Int): ByteArray { + require(position >= 0) { "Position must be non-negative: $position" } + require(length >= 0) { "Length must be non-negative: $length" } + require(position + length <= size) { + "Read beyond end of data: position=$position, length=$length, size=$size" + } + val out = ByteArray(length) + readInto(position, out, 0, length) + return out + } + + override fun readAt(position: Long, buffer: ByteArray, offset: Int, length: Int): Int { + require(position >= 0) { "Position must be non-negative: $position" } + require(offset >= 0 && length >= 0 && offset + length <= buffer.size) { + "Buffer overflow: offset=$offset, length=$length, buffer.size=${buffer.size}" + } + if (position >= size) return 0 + val available = minOf(length.toLong(), size - position).toInt() + readInto(position, buffer, offset, available) + return available + } + + override fun readByteAt(position: Long): Byte { + require(position in 0 until size) { "Position out of range: $position (size=$size)" } + return if (chunks.size == 1) { + chunks[0][position.toInt()] + } else { + chunks[(position / chunkSize).toInt()][(position % chunkSize).toInt()] + } + } + + private fun readInto(position: Long, buffer: ByteArray, offset: Int, length: Int) { + if (length == 0) return + if (chunks.size == 1) { + val p = position.toInt() + chunks[0].copyInto(buffer, offset, p, p + length) + return + } + var remaining = length + var dst = offset + var pos = position + while (remaining > 0) { + val chunk = chunks[(pos / chunkSize).toInt()] + val inChunk = (pos % chunkSize).toInt() + val n = minOf(remaining, chunk.size - inChunk) + chunk.copyInto(buffer, dst, inChunk, inChunk + n) + remaining -= n + dst += n + pos += n + } + } + + override fun close() { + // Nothing to release: the caller owns the arrays. + } +} diff --git a/skainet-io/skainet-io-core/src/commonTest/kotlin/sk/ainet/io/ByteArrayRandomAccessSourceTest.kt b/skainet-io/skainet-io-core/src/commonTest/kotlin/sk/ainet/io/ByteArrayRandomAccessSourceTest.kt new file mode 100644 index 000000000..b05bcde68 --- /dev/null +++ b/skainet-io/skainet-io-core/src/commonTest/kotlin/sk/ainet/io/ByteArrayRandomAccessSourceTest.kt @@ -0,0 +1,50 @@ +package sk.ainet.io + +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith + +/** SKEEP-007 §T0.4: the engine's in-memory [RandomAccessSource], single array and chunked. */ +class ByteArrayRandomAccessSourceTest { + private val bytes = ByteArray(100) { it.toByte() } + + @Test + fun `single array reads any window`() { + val src = ByteArrayRandomAccessSource(bytes) + assertEquals(100L, src.size) + assertContentEquals(bytes.copyOfRange(10, 25), src.readAt(10, 15)) + assertContentEquals(ByteArray(0), src.readAt(100, 0)) + assertEquals(42.toByte(), src.readByteAt(42)) + val buf = ByteArray(10) { 0x7F } + assertEquals(4, src.readAt(96, buf, 2, 8), "short read at the end") + assertContentEquals(byteArrayOf(0x7F, 0x7F, 96, 97, 98, 99, 0x7F, 0x7F, 0x7F, 0x7F), buf) + assertEquals(0, src.readAt(100, buf)) + } + + @Test + fun `chunked source reads across chunk boundaries`() { + val chunks = listOf(bytes.copyOfRange(0, 32), bytes.copyOfRange(32, 64), bytes.copyOfRange(64, 96), bytes.copyOfRange(96, 100)) + val src = ByteArrayRandomAccessSource(chunks) + assertEquals(100L, src.size) + assertContentEquals(bytes.copyOfRange(30, 70), src.readAt(30, 40)) + assertContentEquals(bytes.copyOfRange(63, 65), src.readAt(63, 2)) + assertContentEquals(bytes.copyOfRange(95, 100), src.readAt(95, 5)) + assertEquals(64.toByte(), src.readByteAt(64)) + val buf = ByteArray(12) + assertEquals(8, src.readAt(92, buf, 0, 12)) + assertContentEquals(bytes.copyOfRange(92, 100), buf.copyOfRange(0, 8)) + } + + @Test + fun `bounds and chunk geometry are validated`() { + val src = ByteArrayRandomAccessSource(bytes) + assertFailsWith { src.readAt(90, 11) } + assertFailsWith { src.readAt(-1, 1) } + assertFailsWith { src.readByteAt(100) } + assertFailsWith { ByteArrayRandomAccessSource(emptyList()) } + assertFailsWith { + ByteArrayRandomAccessSource(listOf(ByteArray(10), ByteArray(5), ByteArray(3))) + } + } +} diff --git a/skainet-io/skainet-io-core/src/wasmJsMain/kotlin/sk/ainet/io/JsBlobRandomAccessSource.kt b/skainet-io/skainet-io-core/src/wasmJsMain/kotlin/sk/ainet/io/JsBlobRandomAccessSource.kt index 636fb7176..72b7860b5 100644 --- a/skainet-io/skainet-io-core/src/wasmJsMain/kotlin/sk/ainet/io/JsBlobRandomAccessSource.kt +++ b/skainet-io/skainet-io-core/src/wasmJsMain/kotlin/sk/ainet/io/JsBlobRandomAccessSource.kt @@ -9,9 +9,6 @@ import org.khronos.webgl.get import org.w3c.files.Blob import kotlin.js.Promise -@JsFun("(msg) => console.log(msg)") -private external fun consoleLog(msg: String) - /** * WASM JS implementation of [RandomAccessSource] using Blob API. * @@ -136,22 +133,12 @@ public class JsBlobRandomAccessSource private constructor( */ public suspend fun open(blob: Blob, preloadSize: Int = DEFAULT_PRELOAD_SIZE): JsBlobRandomAccessSource { val blobSize = getBlobSize(blob) - consoleLog("[JsBlobRAS] blobSize=$blobSize, preloadSize=$preloadSize") val actualPreloadSize = minOf(preloadSize, blobSize) - consoleLog("[JsBlobRAS] actualPreloadSize=$actualPreloadSize") val slice = sliceBlob(blob, 0, actualPreloadSize) val sliceSize = getBlobSize(slice) - consoleLog("[JsBlobRAS] sliceSize=$sliceSize") val arrayBuffer = blobArrayBuffer(slice).await() val bufferByteLength = getArrayBufferByteLength(arrayBuffer) - consoleLog("[JsBlobRAS] arrayBuffer.byteLength=$bufferByteLength") val bytes = arrayBufferToByteArray(arrayBuffer) - consoleLog("[JsBlobRAS] bytes.size=${bytes.size}") - // Log first 8 bytes as hex for ONNX magic check - if (bytes.size >= 8) { - val hex = bytes.take(8).joinToString(" ") { (it.toInt() and 0xFF).toString(16).padStart(2, '0') } - consoleLog("[JsBlobRAS] first 8 bytes: $hex") - } return JsBlobRandomAccessSource(blob, bytes, blobSize.toLong()) } @@ -159,16 +146,12 @@ public class JsBlobRandomAccessSource private constructor( * Convert JavaScript ArrayBuffer to Kotlin ByteArray using Int8Array. */ private fun arrayBufferToByteArray(buffer: ArrayBuffer): ByteArray { + // One interop call per byte: Kotlin/Wasm 2.4 ships no `Int8Array` → `ByteArray` adapter + // and arrays cannot cross the boundary directly. For a 50 MB preload that is 50 million + // calls before the first header byte is parsed; SKEEP-007 Tier 1 replaces this with the + // sidecar's linear-memory bridge. Until then the loop stays, without the per-read logging. val int8Array = Int8Array(buffer) - val length = int8Array.length - consoleLog("[JsBlobRAS] converting $length bytes via Int8Array") - val bytes = ByteArray(length) { i -> int8Array[i] } - // Debug: log first 16 bytes - if (length >= 16) { - val first16 = (0 until 16).map { int8Array[it].toInt() and 0xFF } - consoleLog("[JsBlobRAS] first 16 raw values: $first16") - } - return bytes + return ByteArray(int8Array.length) { i -> int8Array[i] } } } } diff --git a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt index e217e1c21..249ae688e 100644 --- a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt +++ b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt @@ -146,10 +146,8 @@ public class StreamingGgufParametersLoader( /** The dense FP32 size of [elements] — what every widening in this loader converts to. */ private fun denseFp32Bytes(elements: Long): Long = elements * 4 - /** Whether [form] asks for the one requantization this loader has (#1150). */ - private fun planesRequested(form: WeightForm): Boolean = - (form.encoding as? EncodingRequest.RequantizeTo)?.encoding == - sk.ainet.lang.tensor.storage.TensorEncoding.BITNET_PLANES + /** Whether [form] asks for a requantization (#1150 for planes; every registered encoding since SKEEP-007 §T2). */ + private fun requantizeRequested(form: WeightForm): Boolean = form.encoding is EncodingRequest.RequantizeTo /** * Requantize [values] (row-major `[out, in]` — validateForm pinned the orientation) into a @@ -158,7 +156,8 @@ public class StreamingGgufParametersLoader( * a conversion someone should be able to see. */ @Suppress("UNCHECKED_CAST") - private fun planesTensor( + private fun requantizedTensor( + form: WeightForm, ctx: ExecutionContext, dtype: KClass, shape: Shape, @@ -167,24 +166,48 @@ public class StreamingGgufParametersLoader( bytesBefore: Long, fromFormat: Format, ): Tensor { + val target = (form.encoding as EncodingRequest.RequantizeTo).encoding require(shape.rank == 2) { - "tensor '$tensorName': RequantizeTo(BITNET_PLANES) needs a 2-D weight, got rank ${shape.rank}" + "tensor '$tensorName': RequantizeTo(${target.name}) needs a 2-D weight, got rank ${shape.rank}" } require(dtype == FP32::class) { - "tensor '$tensorName': BITNET_PLANES weights are logically FP32; requested $dtype" + "tensor '$tensorName': ${target.name} weights are logically FP32; requested $dtype" } - val data = sk.ainet.lang.tensor.data.BitNetPlanesTensorData.fromFloats(shape, values) + var data = sk.ainet.lang.tensor.data.Requantizers.quantize(target, values, shape) + // Freshly quantized blocks are row-major; honour KERNEL_FEED for the formats that define it. + if (form.order == WeightByteOrder.KERNEL_FEED) { + val packed = data as? sk.ainet.lang.tensor.storage.PackedBlockStorage + if (packed != null && packed.blockOrder == sk.ainet.lang.memory.BlockOrder.ROW_MAJOR) { + data = feedOrderedByEncoding(packed) + } + } + val packedBytes = (data as? sk.ainet.lang.tensor.storage.PackedBlockStorage)?.packedData?.size?.toLong() ?: -1L traceConversion( - kind = "requantize-planes", + kind = "requantize-${target.name.lowercase()}", tensorName = tensorName, from = fromFormat, - to = Format(FP32, sk.ainet.lang.tensor.storage.TensorEncoding.BITNET_PLANES), + to = Format(FP32, target), bytesBefore = bytesBefore, - bytesAfter = data.packedData.size.toLong(), + bytesAfter = packedBytes, ) return ctx.fromData(data as sk.ainet.lang.tensor.data.TensorData, dtype) } + /** [feedOrdered] keyed by the encoding of freshly quantized storage rather than the file's tensor type. */ + private fun feedOrderedByEncoding(packed: sk.ainet.lang.tensor.storage.PackedBlockStorage): sk.ainet.lang.tensor.data.TensorData<*, *> { + val asData = packed as sk.ainet.lang.tensor.data.TensorData<*, *> + val shape = packed.shape + if (shape.rank != 2 || shape[1] % packed.blockSize != 0) return asData + val prepacked = packed.packedView.prepack(sk.ainet.lang.memory.BlockOrder.INPUT_BLOCK_MAJOR, sink = traceSink) + val bytes = (prepacked.storage as? sk.ainet.lang.memory.Storage.Heap)?.bytes ?: return asData + val order = sk.ainet.lang.memory.BlockOrder.INPUT_BLOCK_MAJOR + return when (packed.encoding) { + sk.ainet.lang.tensor.storage.TensorEncoding.Q8_0 -> Q8_0BlockTensorData(shape, bytes, order) + sk.ainet.lang.tensor.storage.TensorEncoding.Q4_0 -> Q4_0BlockTensorData(shape, bytes, order) + else -> asData + } + } + /** The uniform form: [weightForm] if given, else the historical as-stored-on-heap default. */ private val form: WeightForm = weightForm ?: WeightForm.AS_STORED_ON_HEAP @@ -210,15 +233,17 @@ public class StreamingGgufParametersLoader( } val requested = form.encoding if (requested is EncodingRequest.RequantizeTo) { - // #1150: the one requantizer this loader has. Everything else still fails eagerly. - require(requested.encoding == sk.ainet.lang.tensor.storage.TensorEncoding.BITNET_PLANES) { + // One quantizer per target encoding (#1150 planes; SKEEP-007 §T2 registry). Anything + // without one still fails eagerly, before the first tensor is read. + require(sk.ainet.lang.tensor.data.Requantizers.supports(requested.encoding)) { "$where: EncodingRequest.RequantizeTo(${requested.encoding.name}) is not supported " + "by this loader: re-quantizing a weight the file does not already carry needs a " + - "quantizer per target encoding, and only BITNET_PLANES exists here (#1150)." + "quantizer per target encoding, and quantize-at-load has " + + "${sk.ainet.lang.tensor.data.Requantizers.supported.joinToString { it.name }}." } require(form.shape == WeightShapeOrientation.OUT_IN) { - "$where: RequantizeTo(BITNET_PLANES) needs WeightShapeOrientation.OUT_IN — the " + - "format's scales are per output row, which has no answer while the tensor is " + + "$where: RequantizeTo(${requested.encoding.name}) needs WeightShapeOrientation.OUT_IN — the " + + "format's scales run along output rows, which has no answer while the tensor is " + "still labelled in the file's `ne` order." } } @@ -303,7 +328,7 @@ public class StreamingGgufParametersLoader( dtype == FP32::class && tensorForm.order != WeightByteOrder.KERNEL_FEED && tensorForm.encoding !is EncodingRequest.DequantizeTo && - !planesRequested(tensorForm) + !requantizeRequested(tensorForm) ) { val enc = when (tensorInfo.tensorType) { GGMLQuantizationType.Q4_K -> sk.ainet.lang.tensor.storage.TensorEncoding.Q4_K @@ -348,9 +373,9 @@ public class StreamingGgufParametersLoader( when (dtype) { // The freshly decoded array is loader-owned — wrap it zero-copy // instead of paying the factory's defensive copy (#782). - FP32::class -> if (planesRequested(tensorForm)) { - planesTensor( - ctx, dtype, shape, tensorInfo.name, bytesToFloatArray(rawBytes), + FP32::class -> if (requantizeRequested(tensorForm)) { + requantizedTensor( + tensorForm, ctx, dtype, shape, tensorInfo.name, bytesToFloatArray(rawBytes), rawBytes.size.toLong(), Format.dense(FP32), ) } else { @@ -377,9 +402,9 @@ public class StreamingGgufParametersLoader( ctx.fromData(packed as sk.ainet.lang.tensor.data.TensorData, dtype) } else { // Loader-owned widened array — zero-copy wrap (#782). - if (planesRequested(tensorForm)) { - planesTensor( - ctx, dtype, shape, tensorInfo.name, dequantF16(rawBytes), + if (requantizeRequested(tensorForm)) { + requantizedTensor( + tensorForm, ctx, dtype, shape, tensorInfo.name, dequantF16(rawBytes), rawBytes.size.toLong(), Format.dense(FP16), ) } else { @@ -466,10 +491,10 @@ public class StreamingGgufParametersLoader( rawBytes: ByteArray, tensorForm: WeightForm, ): Tensor { - if (planesRequested(tensorForm) && dtype == FP32::class) { + if (requantizeRequested(tensorForm) && dtype == FP32::class) { val dest = DequantOps.dequantFromBytes(rawBytes, tensorInfo.tensorType, tensorInfo.nElements.toInt()) - return planesTensor( - ctx, dtype, shape, tensorInfo.name, dest, + return requantizedTensor( + tensorForm, ctx, dtype, shape, tensorInfo.name, dest, rawBytes.size.toLong(), ggufFormat(tensorInfo.tensorType, rawBytes.size.toLong()), ) } @@ -572,9 +597,9 @@ public class StreamingGgufParametersLoader( bytesBefore = rawBytes.size.toLong(), bytesAfter = packedBytes.size.toLong(), ) - if (planesRequested(tensorForm)) { - return planesTensor( - ctx, dtype, shape, tensorInfo.name, + if (requantizeRequested(tensorForm)) { + return requantizedTensor( + tensorForm, ctx, dtype, shape, tensorInfo.name, sk.ainet.lang.memory.TernaryCodec.decodeBitNet(packedBytes, tensorInfo.nElements.toInt()), rawBytes.size.toLong(), Format(FP32, sk.ainet.lang.tensor.storage.TensorEncoding.BITNET_B1_58), diff --git a/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/GgufRequantizeTest.kt b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/GgufRequantizeTest.kt new file mode 100644 index 000000000..fef1d5a51 --- /dev/null +++ b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/GgufRequantizeTest.kt @@ -0,0 +1,85 @@ +package sk.ainet.io.gguf + +import kotlinx.coroutines.runBlocking +import sk.ainet.context.DefaultDataExecutionContext +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.lang.memory.BlockOrder +import sk.ainet.lang.memory.plan.EncodingRequest +import sk.ainet.lang.memory.plan.WeightByteOrder +import sk.ainet.lang.memory.plan.WeightForm +import sk.ainet.lang.memory.plan.WeightShapeOrientation +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.Q8_0BlockTensorData +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.FP32 +import java.io.File +import kotlin.math.abs +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertIs +import kotlin.test.assertTrue + +/** + * SKEEP-007 §T2: the GGUF loader's `RequantizeTo` now reaches every encoding in `Requantizers`, not + * only `BITNET_PLANES` (#1150). A dense F32 weight in a GGUF becomes Q8_0 at load, in canonical or + * kernel feed order, within one quantization step of its dense values. + */ +class GgufRequantizeTest { + private fun file(): File = SyntheticGguf.write( + SyntheticGguf.tensor("dense2d", GGMLQuantizationType.F32, elements = 2048).copy(dims = listOf(64L, 32L)), + SyntheticGguf.tensor("vector", GGMLQuantizationType.F32, elements = 64), + ) + + private fun loadVia(f: File, form: WeightForm?, formFor: ((String) -> WeightForm?)? = null): Map> { + val ctx = DefaultDataExecutionContext() + val loaded = LinkedHashMap>() + runBlocking { + StreamingGgufParametersLoader( + sourceProvider = { JvmRandomAccessSource.open(f) }, + weightForm = form, + weightFormFor = formFor, + ).load(ctx, FP32::class) { name, tensor -> loaded[name] = tensor } + } + return loaded + } + + private val q8 = WeightForm(encoding = EncodingRequest.RequantizeTo(TensorEncoding.Q8_0), shape = WeightShapeOrientation.OUT_IN) + + @Test + fun `a dense F32 matrix is requantized to Q8_0 within one code step`() { + val f = file() + try { + val dense = loadVia(f, null).getValue("dense2d").data.copyToFloatArray() + val packed = assertIs( + loadVia(f, null, formFor = { if (it == "dense2d") q8 else null }).getValue("dense2d").data, + ) + assertEquals(BlockOrder.ROW_MAJOR, packed.blockOrder) + val back = packed.toFloatArray() + assertEquals(dense.size, back.size) + val amax = dense.maxOf { abs(it) } + for (i in dense.indices) { + assertTrue(abs(dense[i] - back[i]) <= amax / 127f + 1e-4f, "element $i: ${back[i]} vs ${dense[i]}") + } + } finally { + f.delete() + } + } + + @Test + fun `KERNEL_FEED on a requantized weight yields input-block-major blocks with the same values`() { + val f = file() + try { + val fed = assertIs( + loadVia(f, null, formFor = { if (it == "dense2d") q8.copy(order = WeightByteOrder.KERNEL_FEED) else null }) + .getValue("dense2d").data, + ) + assertEquals(BlockOrder.INPUT_BLOCK_MAJOR, fed.blockOrder) + val canonical = assertIs( + loadVia(f, null, formFor = { if (it == "dense2d") q8 else null }).getValue("dense2d").data, + ) + assertEquals(canonical.toFloatArray().toList(), fed.toFloatArray().toList()) + } finally { + f.delete() + } + } +} diff --git a/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/WeightFormLoaderParityTest.kt b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/WeightFormLoaderParityTest.kt index d7681e68f..8bd9c4559 100644 --- a/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/WeightFormLoaderParityTest.kt +++ b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/WeightFormLoaderParityTest.kt @@ -91,7 +91,8 @@ class WeightFormLoaderParityTest { StreamingGgufParametersLoader( sourceProvider = { JvmRandomAccessSource.open(file()) }, weightForm = WeightForm( - encoding = EncodingRequest.RequantizeTo(sk.ainet.lang.tensor.storage.TensorEncoding.Q8_0), + // Q8_0 and Q4_0 have quantizers since SKEEP-007 §T2; the K-quants still do not. + encoding = EncodingRequest.RequantizeTo(sk.ainet.lang.tensor.storage.TensorEncoding.Q4_K), ), ) } diff --git a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsMaterializer.kt b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsMaterializer.kt index d4b971161..0260b6ce5 100644 --- a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsMaterializer.kt +++ b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsMaterializer.kt @@ -4,7 +4,18 @@ import sk.ainet.context.ExecutionContext import sk.ainet.io.model.DataType import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.memory.BlockOrder +import sk.ainet.lang.memory.Storage +import sk.ainet.lang.memory.blockSpec +import sk.ainet.lang.memory.plan.EncodingRequest +import sk.ainet.lang.memory.plan.WeightByteOrder +import sk.ainet.lang.memory.plan.WeightForm import sk.ainet.lang.tensor.data.Bf16DenseTensorData +import sk.ainet.lang.tensor.data.Q4_0BlockTensorData +import sk.ainet.lang.tensor.data.Q8_0BlockTensorData +import sk.ainet.lang.tensor.data.Requantizers +import sk.ainet.lang.tensor.storage.PackedBlockStorage +import sk.ainet.lang.tensor.storage.TensorEncoding import sk.ainet.lang.tensor.data.Fp16DenseTensorData import sk.ainet.lang.tensor.data.TensorData import sk.ainet.lang.types.BF16 @@ -55,6 +66,76 @@ internal object SafeTensorsMaterializer { bytes: ByteArray, bf16Policy: Bf16LoadPolicy, fp16Policy: NarrowFloatLoadPolicy, + form: WeightForm? = null, + ): Tensor { + requantized(ctx, dtype, name, dataType, shape, bytes, form)?.let { return it } + return materializeDense(ctx, dtype, name, dataType, rawDtype, shape, bytes, bf16Policy, fp16Policy) + } + + /** + * The `EncodingRequest.RequantizeTo` branch (SKEEP-007 §T2): a float **matrix** (rank 2, last + * dimension a multiple of the block size) is decoded to FP32 and handed to [Requantizers] for the + * requested encoding, then relaid into kernel feed order when the form asks for it. Anything the + * request cannot apply to — integer tensors, vectors, a matrix whose row length does not block — + * returns `null` and loads as it would without the form, so biases, norms and index tables stay + * dense under a uniform `RequantizeTo(Q8_0)`. `WeightShapeOrientation` is moot for SafeTensors: + * a `Linear` weight is already `[out, in]`. `WeightResidency.MAPPED` cannot be honoured for + * freshly quantized bytes and silently means heap, as the GGUF loader documents. + */ + @Suppress("UNCHECKED_CAST") + private fun requantized( + ctx: ExecutionContext, + dtype: KClass, + name: String, + dataType: DataType, + shape: Shape, + bytes: ByteArray, + form: WeightForm?, + ): Tensor? { + val target = (form?.encoding as? EncodingRequest.RequantizeTo)?.encoding ?: return null + if (shape.rank != 2) return null + val blockSize = target.blockSpec?.blockSize ?: return null + if (shape[1] % blockSize != 0) return null + val floats = when (dataType) { + DataType.FLOAT32 -> bytesToFloatArray(bytes) + DataType.FLOAT16 -> dequantF16(bytes) + DataType.BFLOAT16 -> dequantBF16(bytes) + else -> return null + } + require(dtype == FP32::class) { + "SafeTensors tensor '$name': RequantizeTo(${target.name}) weights are logically FP32; requested ${dtype.simpleName}" + } + var data = Requantizers.quantize(target, floats, shape) + if (form.order == WeightByteOrder.KERNEL_FEED) data = feedOrdered(data) + return ctx.fromData(data as TensorData, dtype) + } + + /** Relayout freshly quantized row-major blocks into kernel feed order, for the formats that have one. */ + private fun feedOrdered(data: TensorData<*, *>): TensorData<*, *> { + val packed = data as? PackedBlockStorage ?: return data + if (packed.blockOrder != BlockOrder.ROW_MAJOR) return data + val shape = packed.shape + if (shape.rank != 2 || shape[1] % packed.blockSize != 0) return data + val prepacked = packed.packedView.prepack(BlockOrder.INPUT_BLOCK_MAJOR) + val bytes = (prepacked.storage as? Storage.Heap)?.bytes ?: return data + return when (packed.encoding) { + TensorEncoding.Q8_0 -> Q8_0BlockTensorData(shape, bytes, BlockOrder.INPUT_BLOCK_MAJOR) + TensorEncoding.Q4_0 -> Q4_0BlockTensorData(shape, bytes, BlockOrder.INPUT_BLOCK_MAJOR) + else -> data + } + } + + @Suppress("UNCHECKED_CAST") + private fun materializeDense( + ctx: ExecutionContext, + dtype: KClass, + name: String, + dataType: DataType, + rawDtype: String, + shape: Shape, + bytes: ByteArray, + bf16Policy: Bf16LoadPolicy, + fp16Policy: NarrowFloatLoadPolicy, ): Tensor = when (dataType) { DataType.FLOAT32 -> { require(dtype == FP32::class) { diff --git a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoader.kt b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoader.kt index c1dd64c62..7981ac3f5 100644 --- a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoader.kt +++ b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoader.kt @@ -3,6 +3,7 @@ package sk.ainet.io.safetensors import sk.ainet.context.ExecutionContext import sk.ainet.io.ParametersLoader import sk.ainet.io.RandomAccessSource +import sk.ainet.lang.memory.plan.WeightForm import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.Tensor import sk.ainet.lang.types.DType @@ -52,8 +53,21 @@ class SafeTensorsParametersLoader( private val bf16Policy: Bf16LoadPolicy = NarrowFloatLoadPolicy.DEQUANT_TO_FP32, private val fp16Policy: NarrowFloatLoadPolicy = NarrowFloatLoadPolicy.DEQUANT_TO_FP32, private val tensorFilter: ((StreamingSafeTensorInfo) -> Boolean)? = null, + private val weightForm: WeightForm? = null, + private val weightFormFor: ((String) -> WeightForm?)? = null, ) : ParametersLoader { + init { + weightForm?.let { SafeTensorsWeightForms.validate(it, "weightForm") } + } + + /** The form [tensorName] loads under: [weightFormFor] first, then the uniform [weightForm]. */ + private fun formFor(tensorName: String): WeightForm? { + val perTensor = weightFormFor?.invoke(tensorName) ?: return weightForm + SafeTensorsWeightForms.validate(perTensor, "weightFormFor('$tensorName')") + return perTensor + } + override suspend fun load( ctx: ExecutionContext, dtype: KClass, @@ -78,6 +92,7 @@ class SafeTensorsParametersLoader( bytes = bytes, bf16Policy = bf16Policy, fp16Policy = fp16Policy, + form = formFor(tensorInfo.name), ) onTensorLoaded(tensorInfo.name, tensor) diff --git a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsWeightForms.kt b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsWeightForms.kt new file mode 100644 index 000000000..c8abd8cee --- /dev/null +++ b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsWeightForms.kt @@ -0,0 +1,33 @@ +package sk.ainet.io.safetensors + +import sk.ainet.lang.memory.plan.EncodingRequest +import sk.ainet.lang.memory.plan.WeightForm +import sk.ainet.lang.tensor.data.Requantizers +import sk.ainet.lang.types.FP32 + +/** + * What the SafeTensors loaders accept in a [WeightForm] (SKEEP-007 §T2), checked at construction so + * an impossible request fails before the first tensor is read — the same contract the GGUF loader + * has had since #1109. + * + * - `KeepAsStored`: the narrow-float policies decide, as before. + * - `DequantizeTo(FP32)`: the historical widening; any other dtype is refused. + * - `RequantizeTo(encoding)`: accepted for every encoding [Requantizers] has a quantizer for. + * - `WeightByteOrder.KERNEL_FEED` is honoured for the requantized formats that define it and is a + * no-op for dense tensors; `WeightShapeOrientation` is moot (SafeTensors weights are `[out, in]`); + * `WeightResidency.MAPPED` means heap for anything this loader produces. + */ +internal object SafeTensorsWeightForms { + fun validate(form: WeightForm, where: String) { + when (val request = form.encoding) { + is EncodingRequest.RequantizeTo -> require(Requantizers.supports(request.encoding)) { + "$where: EncodingRequest.RequantizeTo(${request.encoding.name}) is not supported: no quantizer " + + "exists for that encoding (quantize-at-load supports ${Requantizers.supported.joinToString { it.name }})." + } + is EncodingRequest.DequantizeTo -> require(request.dtype == FP32) { + "$where: EncodingRequest.DequantizeTo(${request.dtype.name}) is not supported: this loader widens to FP32 only." + } + EncodingRequest.KeepAsStored -> Unit + } + } +} diff --git a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/ShardedSafeTensorsParametersLoader.kt b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/ShardedSafeTensorsParametersLoader.kt index 57e776b53..186376510 100644 --- a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/ShardedSafeTensorsParametersLoader.kt +++ b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/ShardedSafeTensorsParametersLoader.kt @@ -2,6 +2,7 @@ package sk.ainet.io.safetensors import sk.ainet.context.ExecutionContext import sk.ainet.io.ParametersLoader +import sk.ainet.lang.memory.plan.WeightForm import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.Tensor import sk.ainet.lang.types.DType @@ -60,8 +61,21 @@ public class ShardedSafeTensorsParametersLoader( private val fp16Policy: NarrowFloatLoadPolicy = NarrowFloatLoadPolicy.DEQUANT_TO_FP32, private val allowPartial: Boolean = false, private val tensorFilter: ((ShardedTensorInfo) -> Boolean)? = null, + private val weightForm: WeightForm? = null, + private val weightFormFor: ((String) -> WeightForm?)? = null, ) : ParametersLoader { + init { + weightForm?.let { SafeTensorsWeightForms.validate(it, "weightForm") } + } + + /** The form [tensorName] loads under: [weightFormFor] first, then the uniform [weightForm]. */ + private fun formFor(tensorName: String): WeightForm? { + val perTensor = weightFormFor?.invoke(tensorName) ?: return weightForm + SafeTensorsWeightForms.validate(perTensor, "weightFormFor('$tensorName')") + return perTensor + } + override suspend fun load( ctx: ExecutionContext, dtype: KClass, @@ -90,6 +104,7 @@ public class ShardedSafeTensorsParametersLoader( bytes = bytes, bf16Policy = bf16Policy, fp16Policy = fp16Policy, + form = formFor(tensorInfo.name), ) onTensorLoaded(tensorInfo.name, tensor) diff --git a/skainet-io/skainet-io-safetensors/src/jvmTest/kotlin/sk/ainet/io/safetensors/SafeTensorsRequantizeTest.kt b/skainet-io/skainet-io-safetensors/src/jvmTest/kotlin/sk/ainet/io/safetensors/SafeTensorsRequantizeTest.kt new file mode 100644 index 000000000..b8157f0f6 --- /dev/null +++ b/skainet-io/skainet-io-safetensors/src/jvmTest/kotlin/sk/ainet/io/safetensors/SafeTensorsRequantizeTest.kt @@ -0,0 +1,148 @@ +package sk.ainet.io.safetensors + +import kotlinx.coroutines.runBlocking +import kotlinx.io.Buffer +import kotlinx.io.readByteArray +import sk.ainet.context.DefaultDataExecutionContext +import sk.ainet.io.ByteArrayRandomAccessSource +import sk.ainet.lang.memory.BlockOrder +import sk.ainet.lang.memory.plan.EncodingRequest +import sk.ainet.lang.memory.plan.WeightByteOrder +import sk.ainet.lang.memory.plan.WeightForm +import sk.ainet.lang.memory.plan.WeightShapeOrientation +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.FloatArrayTensorData +import sk.ainet.lang.tensor.data.Q8_0BlockTensorData +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.FP16 +import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.Fp16Codec +import kotlin.math.abs +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertIs +import kotlin.test.assertTrue + +/** + * SKEEP-007 §T2: the SafeTensors loader honours `WeightForm`, including `RequantizeTo(Q8_0)`. + * + * Every Hugging Face checkpoint is dense F32/F16/BF16; until now only a GGUF could be served packed. + * These tests pin what the request applies to (rank-2 float tensors whose rows block), what it leaves + * alone (vectors), feed-order relayout, per-tensor overrides, and the up-front refusal of a request + * no quantizer can honour. + */ +class SafeTensorsRequantizeTest { + private val ctx = DefaultDataExecutionContext() + + private fun pseudo(n: Int, seed: Int): FloatArray { + var s = seed + return FloatArray(n) { + s = s * 1103515245 + 12345 + ((s ushr 16) % 2001 - 1000) / 400f + } + } + + private val w = pseudo(4 * 64, 1) // [4, 64]: two Q8_0 blocks per row + private val e = pseudo(8 * 96, 2) // [8, 96]: three blocks per row, F16 on disk + private val bias = pseudo(4, 3) // [4]: a vector, never requantized + private val odd = pseudo(3 * 40, 4) // [3, 40]: rows do not block → stays dense + + private fun f16(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 file(): ByteArray { + val buffer = Buffer() + SafeTensorsWriter.write(buffer) { + tensorF32("w", intArrayOf(4, 64), w) + tensor("e", "F16", intArrayOf(8, 96)) { f16(e) } + tensorF32("bias", intArrayOf(4), bias) + tensorF32("odd", intArrayOf(3, 40), odd) + } + return buffer.readByteArray() + } + + private fun load(form: WeightForm?, formFor: ((String) -> WeightForm?)? = null): Map> { + val bytes = file() + val loaded = LinkedHashMap>() + runBlocking { + SafeTensorsParametersLoader( + sourceProvider = { ByteArrayRandomAccessSource(bytes) }, + weightForm = form, + weightFormFor = formFor, + ).load(ctx, FP32::class) { name, tensor -> loaded[name] = tensor } + } + return loaded + } + + private fun assertClose(expected: FloatArray, actual: FloatArray, tolerance: Float, what: String) { + assertEquals(expected.size, actual.size, "$what size") + for (i in expected.indices) { + assertTrue(abs(expected[i] - actual[i]) <= tolerance, "$what[$i]: ${actual[i]} vs ${expected[i]}") + } + } + + @Test + fun `a uniform RequantizeTo Q8_0 packs the matrices and leaves vectors and non-blocking rows dense`() { + val form = WeightForm(encoding = EncodingRequest.RequantizeTo(TensorEncoding.Q8_0), shape = WeightShapeOrientation.OUT_IN) + val loaded = load(form) + val packedW = assertIs(loaded.getValue("w").data) + assertEquals(BlockOrder.ROW_MAJOR, packedW.blockOrder) + assertClose(w, packedW.toFloatArray(), tolerance = 2.5f / 127f + 1e-3f, what = "w") + val packedE = assertIs(loaded.getValue("e").data) + assertClose(e, packedE.toFloatArray(), tolerance = 2.5f / 127f + 4e-3f, what = "e (F16 source)") + assertIs>(loaded.getValue("bias").data, "a vector is not a weight matrix") + assertIs>(loaded.getValue("odd").data, "rows of 40 do not block by 32") + } + + @Test + fun `KERNEL_FEED relays the fresh blocks into input-block-major order with the same values`() { + val form = WeightForm( + encoding = EncodingRequest.RequantizeTo(TensorEncoding.Q8_0), + order = WeightByteOrder.KERNEL_FEED, + shape = WeightShapeOrientation.OUT_IN, + ) + val fed = assertIs(load(form).getValue("w").data) + assertEquals(BlockOrder.INPUT_BLOCK_MAJOR, fed.blockOrder) + val canonical = assertIs( + load(WeightForm(encoding = EncodingRequest.RequantizeTo(TensorEncoding.Q8_0), shape = WeightShapeOrientation.OUT_IN)).getValue("w").data, + ) + assertClose(canonical.toFloatArray(), fed.toFloatArray(), tolerance = 0f, what = "feed order decodes to the same values") + } + + @Test + fun `weightFormFor overrides the uniform form per tensor`() { + val loaded = load( + form = null, + formFor = { name -> if (name == "w") WeightForm(encoding = EncodingRequest.RequantizeTo(TensorEncoding.Q8_0)) else null }, + ) + assertIs(loaded.getValue("w").data) + assertIs>(loaded.getValue("e").data, "no form → the F16 policy decides (widen)") + } + + @Test + fun `requests no quantizer can honour are refused at construction`() { + val bytes = file() + val noQuantizer = assertFailsWith { + SafeTensorsParametersLoader( + sourceProvider = { ByteArrayRandomAccessSource(bytes) }, + weightForm = WeightForm(encoding = EncodingRequest.RequantizeTo(TensorEncoding.Q4_K)), + ) + } + assertTrue(noQuantizer.message!!.contains("quantizer"), noQuantizer.message!!) + val narrow = assertFailsWith { + SafeTensorsParametersLoader( + sourceProvider = { ByteArrayRandomAccessSource(bytes) }, + weightForm = WeightForm(encoding = EncodingRequest.DequantizeTo(FP16)), + ) + } + assertTrue(narrow.message!!.contains("FP32"), narrow.message!!) + } +} diff --git a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api index 557860fc1..c2564589a 100644 --- a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api +++ b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api @@ -1916,11 +1916,12 @@ public final class sk/ainet/lang/memory/plan/PlannerProfile { public static final field Companion Lsk/ainet/lang/memory/plan/PlannerProfile$Companion; public static final field MOBILE_RAM_CEILING J public static final field OFF_HEAP_THRESHOLD J - public fun (Ljava/lang/String;JILsk/ainet/lang/memory/plan/KvCacheMode;DJDZZ)V - public synthetic fun (Ljava/lang/String;JILsk/ainet/lang/memory/plan/KvCacheMode;DJDZZILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun (Ljava/lang/String;JILsk/ainet/lang/memory/plan/KvCacheMode;DJDZZLsk/ainet/lang/tensor/storage/TensorEncoding;)V + public synthetic fun (Ljava/lang/String;JILsk/ainet/lang/memory/plan/KvCacheMode;DJDZZLsk/ainet/lang/tensor/storage/TensorEncoding;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public final fun budget (J)Lsk/ainet/lang/memory/plan/Budget; public final fun checkDequant (Ljava/lang/Double;)Lsk/ainet/lang/memory/plan/DequantVerdict; public final fun component1 ()Ljava/lang/String; + public final fun component10 ()Lsk/ainet/lang/tensor/storage/TensorEncoding; public final fun component2 ()J public final fun component3 ()I public final fun component4 ()Lsk/ainet/lang/memory/plan/KvCacheMode; @@ -1929,8 +1930,8 @@ public final class sk/ainet/lang/memory/plan/PlannerProfile { public final fun component7 ()D public final fun component8 ()Z public final fun component9 ()Z - public final fun copy (Ljava/lang/String;JILsk/ainet/lang/memory/plan/KvCacheMode;DJDZZ)Lsk/ainet/lang/memory/plan/PlannerProfile; - public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/PlannerProfile;Ljava/lang/String;JILsk/ainet/lang/memory/plan/KvCacheMode;DJDZZILjava/lang/Object;)Lsk/ainet/lang/memory/plan/PlannerProfile; + public final fun copy (Ljava/lang/String;JILsk/ainet/lang/memory/plan/KvCacheMode;DJDZZLsk/ainet/lang/tensor/storage/TensorEncoding;)Lsk/ainet/lang/memory/plan/PlannerProfile; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/PlannerProfile;Ljava/lang/String;JILsk/ainet/lang/memory/plan/KvCacheMode;DJDZZLsk/ainet/lang/tensor/storage/TensorEncoding;ILjava/lang/Object;)Lsk/ainet/lang/memory/plan/PlannerProfile; public final fun domainFor (J)Lsk/ainet/lang/tensor/storage/MemoryDomain; public fun equals (Ljava/lang/Object;)Z public final fun getDequantWarnFraction ()D @@ -1939,6 +1940,7 @@ public final class sk/ainet/lang/memory/plan/PlannerProfile { public final fun getName ()Ljava/lang/String; public final fun getOffHeapThresholdBytes ()J public final fun getPrefillChunk ()I + public final fun getQuantizeDenseTo ()Lsk/ainet/lang/tensor/storage/TensorEncoding; public final fun getReserveBytes ()J public final fun getStrict ()Z public final fun getWeightsMapped ()Z @@ -5154,10 +5156,11 @@ public final class sk/ainet/lang/tensor/data/MmapTensorSource$Companion { public static synthetic fun fromChannel$default (Lsk/ainet/lang/tensor/data/MmapTensorSource$Companion;Ljava/nio/channels/FileChannel;JJILjava/lang/Object;)Lsk/ainet/lang/tensor/data/MmapTensorSource; } -public class sk/ainet/lang/tensor/data/NarrowFloatDenseTensorData : sk/ainet/lang/tensor/data/NarrowFloatTensorData { +public class sk/ainet/lang/tensor/data/NarrowFloatDenseTensorData : sk/ainet/lang/tensor/data/NarrowFloatTensorData, sk/ainet/lang/tensor/data/RowDequantSource { public static final field Companion Lsk/ainet/lang/tensor/data/NarrowFloatDenseTensorData$Companion; public fun (Lsk/ainet/lang/tensor/Shape;[BLsk/ainet/lang/types/NarrowFloatCodec;)V public fun copyToFloatArray ()[F + public fun dequantRow (I)[F public fun get ([I)Ljava/lang/Float; public synthetic fun get ([I)Ljava/lang/Object; public fun getCodec ()Lsk/ainet/lang/types/NarrowFloatCodec; @@ -5672,6 +5675,12 @@ public final class sk/ainet/lang/tensor/data/Q8_0BlockTensorData$Companion { public final fun fromRawBytes (Lsk/ainet/lang/tensor/Shape;[B)Lsk/ainet/lang/tensor/data/Q8_0BlockTensorData; } +public final class sk/ainet/lang/tensor/data/Q8_0Quantizer { + public static final field INSTANCE Lsk/ainet/lang/tensor/data/Q8_0Quantizer; + public final fun quantize ([FLsk/ainet/lang/tensor/Shape;)Lsk/ainet/lang/tensor/data/Q8_0BlockTensorData; + public final fun quantizeToBytes ([F)[B +} + public abstract interface class sk/ainet/lang/tensor/data/Q8_0TensorData : sk/ainet/lang/tensor/data/TensorData { public static final field BLOCK_SIZE I public static final field BYTES_PER_BLOCK I @@ -5697,6 +5706,13 @@ public final class sk/ainet/lang/tensor/data/Q8_0TensorDataKt { public static final fun toFloatArray (Lsk/ainet/lang/tensor/data/Q8_0TensorData;)[F } +public final class sk/ainet/lang/tensor/data/Requantizers { + public static final field INSTANCE Lsk/ainet/lang/tensor/data/Requantizers; + public final fun getSupported ()Ljava/util/List; + public final fun quantize (Lsk/ainet/lang/tensor/storage/TensorEncoding;[FLsk/ainet/lang/tensor/Shape;)Lsk/ainet/lang/tensor/data/TensorData; + public final fun supports (Lsk/ainet/lang/tensor/storage/TensorEncoding;)Z +} + public abstract interface class sk/ainet/lang/tensor/data/RowDequantSource { public abstract fun dequantRow (I)[F } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/PlannerProfile.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/PlannerProfile.kt index f6f85b74a..bc9dec69c 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/PlannerProfile.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/PlannerProfile.kt @@ -32,6 +32,13 @@ public data class PlannerProfile( val dequantWarnFraction: Double = 0.05, val strict: Boolean = false, val weightsMapped: Boolean = false, + /** + * Quantize **dense** sources (FP32/F16/BF16 checkpoints — every SafeTensors file) to this + * encoding at load, when a kernel on the target can feed it (SKEEP-007 §T2). `null`, the + * default, keeps dense weights as stored. A browser or phone profile sets `Q8_0` so a Hugging + * Face checkpoint is held in a quarter of its FP32 bytes and served by the packed kernels. + */ + val quantizeDenseTo: sk.ainet.lang.tensor.storage.TensorEncoding? = null, ) { init { require(reserveBytes >= 0) { "reserveBytes must be >= 0" } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/WeightFormResolver.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/WeightFormResolver.kt index 07618d49a..45b442241 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/WeightFormResolver.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/WeightFormResolver.kt @@ -63,6 +63,18 @@ public object WeightFormResolver { val residency = if (profile.weightsMapped) WeightResidency.MAPPED else WeightResidency.HEAP if (stored == null || stored is TensorEncoding.Dense) { + // A dense source has no packed form to keep — unless the profile asks for one and a + // kernel can feed it (SKEEP-007 §T2). Per-row scales need the row to be the output + // dimension, so the shape axis is pinned to OUT_IN here, the one place this resolver + // decides it. The quantization pays once, at load; a profile that asks for it on a + // target with no kernel for the encoding keeps the dense bytes instead. + val target = profile.quantizeDenseTo + if (target != null && capabilities.canFeedMatmul(target)) { + val order = + if (capabilities.wantsKernelFeedOrder(target) && canProduceKernelFeedOrder) WeightByteOrder.KERNEL_FEED + else WeightByteOrder.AS_STORED + return WeightForm(EncodingRequest.RequantizeTo(target), order, WeightShapeOrientation.OUT_IN, residency) + } return WeightForm(EncodingRequest.KeepAsStored, WeightByteOrder.AS_STORED, residency = residency) } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/NarrowFloatTensorData.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/NarrowFloatTensorData.kt index 8d441fa71..5b12307e2 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/NarrowFloatTensorData.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/NarrowFloatTensorData.kt @@ -53,12 +53,37 @@ public open class NarrowFloatDenseTensorData( initialShape: Shape, private val data: ByteArray, override val codec: NarrowFloatCodec, -) : NarrowFloatTensorData { +) : NarrowFloatTensorData, RowDequantSource { override val shape: Shape = Shape(initialShape.dimensions.copyOf()) private val strides: IntArray = shape.computeStrides() override val packedData: ByteArray get() = data + /** + * One decoded row of a rank-2 table (SKEEP-007 §T0.3). + * + * This is what lets a token-embedding table stay in its on-disk 16-bit form: the embedding + * lookup (`ops.gather`, and the transformers `Embedding` layer) checks for [RowDequantSource] + * and decodes only the rows a step touches, instead of widening the whole table to FP32 at + * load — for a 256k × 768 table that is 393 MB resident instead of 786 MB, on every target, + * and in a browser the difference between fitting and not. + */ + override fun dequantRow(rowIdx: Int): FloatArray { + require(shape.rank == 2) { "dequantRow needs a rank-2 table, got rank ${shape.rank}" } + val rows = shape[0] + val cols = shape[1] + require(rowIdx in 0 until rows) { "row $rowIdx out of range [0, $rows)" } + val out = FloatArray(cols) + var byteIdx = rowIdx * cols * NarrowFloatTensorData.BYTES_PER_ELEMENT + for (c in 0 until cols) { + val lo = data[byteIdx].toInt() and 0xFF + val hi = data[byteIdx + 1].toInt() and 0xFF + out[c] = codec.decode((hi shl 8) or lo) + byteIdx += NarrowFloatTensorData.BYTES_PER_ELEMENT + } + return out + } + /** Physically two bytes per element whatever the declared dtype witness. */ override val encoding: TensorEncoding get() = TensorEncoding.Dense(NarrowFloatTensorData.BYTES_PER_ELEMENT) diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q8_0Quantizer.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q8_0Quantizer.kt new file mode 100644 index 000000000..fbd721a38 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q8_0Quantizer.kt @@ -0,0 +1,66 @@ +package sk.ainet.lang.tensor.data + +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.types.Fp16Codec +import kotlin.math.abs +import kotlin.math.roundToInt + +/** + * FP32 → Q8_0 quantizer — the encode side of [Q8_0TensorData], in pure `commonMain` (SKEEP-007 §T2). + * + * Q8_0 is the first format quantize-at-load offers: eight-bit codes keep a classifier's decisions + * where a four-bit format may not, the kernels exist on every tier, and the layout is the simplest + * of the ggml family. Algorithm per 32-element block, matching ggml `quantize_row_q8_0_ref`: + * + * 1. `amax = max |x|` over the block. + * 2. `d = amax / 127`, stored as the block's FP16 scale (round to nearest even). + * 3. `code = round(x / d)`, clamped to `[-127, 127]`, one signed byte per element. + * + * Output is row-major blocks: 34 bytes per block (2 scale + 32 codes), blocks in element order, so + * a `[rows, cols]` tensor's blocks run along `cols` — the orientation [Q8_0BlockTensorData] decodes + * and the packed kernels read. + */ +public object Q8_0Quantizer { + private const val BLOCK_SIZE = Q8_0TensorData.BLOCK_SIZE + private const val BYTES_PER_BLOCK = Q8_0TensorData.BYTES_PER_BLOCK + + /** Quantize [values] (length a multiple of 32) into packed Q8_0 bytes — `34 * (values.size / 32)` bytes. */ + public fun quantizeToBytes(values: FloatArray): ByteArray { + require(values.size % BLOCK_SIZE == 0) { + "Q8_0 quantization requires a length that is a multiple of $BLOCK_SIZE; got ${values.size}" + } + val blocks = values.size / BLOCK_SIZE + val out = ByteArray(blocks * BYTES_PER_BLOCK) + for (b in 0 until blocks) { + val base = b * BLOCK_SIZE + var amax = 0f + for (i in 0 until BLOCK_SIZE) { + val a = abs(values[base + i]) + if (a > amax) amax = a + } + val d = amax / 127f + val id = if (d != 0f) 1f / d else 0f + val outBase = b * BYTES_PER_BLOCK + val half = Fp16Codec.encode(d) + out[outBase] = (half and 0xFF).toByte() + out[outBase + 1] = ((half ushr 8) and 0xFF).toByte() + for (i in 0 until BLOCK_SIZE) { + val q = (values[base + i] * id).roundToInt() + out[outBase + 2 + i] = (if (q < -127) -127 else if (q > 127) 127 else q).toByte() + } + } + return out + } + + /** + * Quantize [values] into a [Q8_0BlockTensorData] with logical [shape] (`shape.volume` must equal + * `values.size`; the last dimension must be a multiple of 32 so blocks never straddle rows). + */ + public fun quantize(values: FloatArray, shape: Shape): Q8_0BlockTensorData { + require(shape.volume == values.size) { "shape volume ${shape.volume} must equal values length ${values.size}" } + require(shape.rank == 0 || shape[shape.rank - 1] % BLOCK_SIZE == 0) { + "Q8_0 blocks run along the last dimension, which must be a multiple of $BLOCK_SIZE; got $shape" + } + return Q8_0BlockTensorData(shape, quantizeToBytes(values)) + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Requantizers.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Requantizers.kt new file mode 100644 index 000000000..fc9fc3ff5 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Requantizers.kt @@ -0,0 +1,51 @@ +package sk.ainet.lang.tensor.data + +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.storage.TensorEncoding + +/** + * The quantizers a loader can run at load time, keyed by target encoding (SKEEP-007 §T2). + * + * `EncodingRequest.RequantizeTo(encoding)` has been part of [sk.ainet.lang.memory.plan.WeightForm] + * since #1109 and the planner prices it, but only the GGUF loader executed it and only for + * `BITNET_PLANES` (#1150). This registry is what every loader consults instead of hard-coding one + * format: a request whose encoding has no entry here is refused up front, exactly as before, and + * adding a quantizer is one line. + * + * Every quantizer produces canonical **row-major** blocks for a tensor whose blocks run along its + * last dimension; feed-order relayout, if the kernel wants it, is a separate step on the result. + */ +public object Requantizers { + + /** `true` when a dense FP32 tensor can be quantized to [encoding] at load. */ + public fun supports(encoding: TensorEncoding): Boolean = when (encoding) { + TensorEncoding.Q8_0, TensorEncoding.Q4_0, TensorEncoding.BITNET_PLANES -> true + else -> false + } + + /** The encodings [supports] answers `true` for, for messages and docs. */ + public val supported: List = listOf(TensorEncoding.Q8_0, TensorEncoding.Q4_0, TensorEncoding.BITNET_PLANES) + + /** + * Quantize dense [values] with logical [shape] into packed tensor data of [encoding]. + * + * @throws IllegalArgumentException for an encoding without a quantizer, or a shape the format + * cannot block (last dimension not a multiple of the block size; planes need rank 2). + */ + public fun quantize(encoding: TensorEncoding, values: FloatArray, shape: Shape): TensorData<*, *> = when (encoding) { + TensorEncoding.Q8_0 -> Q8_0Quantizer.quantize(values, shape) + TensorEncoding.Q4_0 -> { + require(shape.rank == 0 || shape[shape.rank - 1] % 32 == 0) { + "Q4_0 blocks run along the last dimension, which must be a multiple of 32; got $shape" + } + Q4_0Quantizer.quantize(values, shape) + } + TensorEncoding.BITNET_PLANES -> { + require(shape.rank == 2) { "BITNET_PLANES needs a rank-2 weight, got $shape" } + BitNetPlanesTensorData.fromFloats(shape, values) + } + else -> throw IllegalArgumentException( + "no quantizer for ${encoding.name}; quantize-at-load supports ${supported.joinToString { it.name }}", + ) + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/QuantizeDenseResolverTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/QuantizeDenseResolverTest.kt new file mode 100644 index 000000000..2159b1235 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/QuantizeDenseResolverTest.kt @@ -0,0 +1,41 @@ +package sk.ainet.lang.memory.plan + +import sk.ainet.lang.tensor.storage.TensorEncoding +import kotlin.test.Test +import kotlin.test.assertEquals + +/** SKEEP-007 §T2: a profile may ask for dense sources to be quantized at load, when a kernel can feed the result. */ +class QuantizeDenseResolverTest { + private fun capableOf(vararg supported: TensorEncoding): KernelCapabilities = + object : KernelCapabilities { + override fun canFeedMatmul(encoding: TensorEncoding?): Boolean = encoding == null || encoding in supported + } + + private val browser = PlannerProfile.DESKTOP.copy(name = "browser", quantizeDenseTo = TensorEncoding.Q8_0) + + @Test + fun `a dense source is requantized when the profile asks and a kernel can feed the encoding`() { + for (stored in listOf(null, TensorEncoding.Dense(4))) { + val form = WeightFormResolver.resolve(stored, browser, capableOf(TensorEncoding.Q8_0), canProduceKernelFeedOrder = true) + assertEquals(EncodingRequest.RequantizeTo(TensorEncoding.Q8_0), form.encoding, "stored=$stored") + assertEquals(WeightShapeOrientation.OUT_IN, form.shape, "per-row scales pin the orientation") + assertEquals(WeightByteOrder.KERNEL_FEED, form.order, "the packed kernels want feed order") + assertEquals(WeightResidency.HEAP, form.residency) + } + } + + @Test + fun `without a kernel for the encoding the dense bytes are kept`() { + val form = WeightFormResolver.resolve(null, browser, capableOf(TensorEncoding.Q4_K), canProduceKernelFeedOrder = true) + assertEquals(EncodingRequest.KeepAsStored, form.encoding) + assertEquals(WeightByteOrder.AS_STORED, form.order) + } + + @Test + fun `the default profile never requantizes and a packed source is untouched by the option`() { + val dense = WeightFormResolver.resolve(null, PlannerProfile.DESKTOP, capableOf(TensorEncoding.Q8_0), canProduceKernelFeedOrder = true) + assertEquals(EncodingRequest.KeepAsStored, dense.encoding) + val packed = WeightFormResolver.resolve(TensorEncoding.Q4_K, browser, capableOf(TensorEncoding.Q4_K), canProduceKernelFeedOrder = true) + assertEquals(EncodingRequest.KeepAsStored, packed.encoding, "quantizeDenseTo applies to dense sources only") + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/NarrowFloatRowDequantTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/NarrowFloatRowDequantTest.kt new file mode 100644 index 000000000..b3b9c9845 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/NarrowFloatRowDequantTest.kt @@ -0,0 +1,42 @@ +package sk.ainet.lang.tensor.data + +import sk.ainet.lang.tensor.Shape +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue + +/** SKEEP-007 §T0.3: narrow-float dense tables serve rows on demand, decoded like `get`. */ +class NarrowFloatRowDequantTest { + private val rows = 5 + private val cols = 7 + private val values = FloatArray(rows * cols) { (it - 17) * 0.375f } + + @Test + fun `fp16 table is a RowDequantSource whose rows match element access`() { + val data = Fp16DenseTensorData.fromFloatArray(Shape(rows, cols), values) + assertTrue(data is RowDequantSource) + for (r in 0 until rows) { + val expected = FloatArray(cols) { c -> data[r, c] } + assertContentEquals(expected, (data as RowDequantSource).dequantRow(r), "row $r") + } + } + + @Test + fun `bf16 table rows match element access`() { + val data = Bf16DenseTensorData.fromFloatArray(Shape(rows, cols), values) + for (r in 0 until rows) { + val expected = FloatArray(cols) { c -> data[r, c] } + assertContentEquals(expected, data.dequantRow(r), "row $r") + } + } + + @Test + fun `out-of-range row and non-table shapes are refused`() { + val table = Fp16DenseTensorData.fromFloatArray(Shape(rows, cols), values) + assertFailsWith { table.dequantRow(rows) } + assertFailsWith { table.dequantRow(-1) } + val cube = Fp16DenseTensorData.fromFloatArray(Shape(1, rows, cols), values) + assertFailsWith { cube.dequantRow(0) } + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/Q8_0QuantizerTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/Q8_0QuantizerTest.kt new file mode 100644 index 000000000..6439c98bb --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/Q8_0QuantizerTest.kt @@ -0,0 +1,81 @@ +package sk.ainet.lang.tensor.data + +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.Fp16Codec +import kotlin.math.abs +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue + +/** SKEEP-007 §T2: the Q8_0 encoder is ggml-faithful and round-trips through the engine's decoder. */ +class Q8_0QuantizerTest { + + private fun pseudo(n: Int, seed: Int, scale: Float = 1f): FloatArray { + var s = seed + return FloatArray(n) { + s = s * 1103515245 + 12345 + ((s ushr 16) % 2001 - 1000) / 1000f * scale + } + } + + @Test + fun `one block matches the ggml reference layout`() { + // amax = 2.54 → d = 0.02, codes = round(x / d) + val values = FloatArray(32) { i -> (i - 16) * 0.1f }.also { it[5] = 2.54f } + val bytes = Q8_0Quantizer.quantizeToBytes(values) + assertEquals(34, bytes.size) + val halfBits = (bytes[0].toInt() and 0xFF) or ((bytes[1].toInt() and 0xFF) shl 8) + val d = Fp16Codec.decode(halfBits) + assertTrue(abs(d - 2.54f / 127f) < 1e-5f, "scale $d") + assertEquals(127, bytes[2 + 5].toInt()) + assertEquals((-1.6f / (2.54f / 127f)).let { kotlin.math.round(it).toInt() }, bytes[2 + 0].toInt()) + } + + @Test + fun `round trip through Q8_0BlockTensorData stays within one code step`() { + val rows = 6; val cols = 96 + val values = pseudo(rows * cols, 17, scale = 3f) + val data = Q8_0Quantizer.quantize(values, Shape(rows, cols)) + assertEquals(TensorEncoding.Q8_0, data.encoding) + val back = data.toFloatArray() + var maxErr = 0f + var blockAmax = 0f + for (b in 0 until rows * cols / 32) { + blockAmax = 0f + for (i in 0 until 32) blockAmax = maxOf(blockAmax, abs(values[b * 32 + i])) + val step = blockAmax / 127f + for (i in 0 until 32) { + val idx = b * 32 + i + val err = abs(back[idx] - values[idx]) + maxErr = maxOf(maxErr, err) + assertTrue(err <= step * 0.5f + 1e-3f * blockAmax, "element $idx: |${back[idx]} - ${values[idx]}| = $err > half step $step") + } + } + assertTrue(maxErr > 0f, "quantization is lossy; a zero error means the decoder was not exercised") + } + + @Test + fun `an all-zero block has a zero scale and zero codes`() { + val bytes = Q8_0Quantizer.quantizeToBytes(FloatArray(32)) + assertTrue(bytes.all { it == 0.toByte() }) + } + + @Test + fun `shapes that do not block along the last dimension are refused`() { + assertFailsWith { Q8_0Quantizer.quantizeToBytes(FloatArray(33)) } + assertFailsWith { Q8_0Quantizer.quantize(FloatArray(64), Shape(4, 16)) } + assertFailsWith { Q8_0Quantizer.quantize(FloatArray(64), Shape(2, 64)) } + } + + @Test + fun `registry routes Q8_0 and Q4_0 and refuses unknown encodings`() { + val values = pseudo(64, 3) + assertTrue(Requantizers.quantize(TensorEncoding.Q8_0, values, Shape(2, 32)) is Q8_0BlockTensorData) + assertTrue(Requantizers.quantize(TensorEncoding.Q4_0, values, Shape(2, 32)) is Q4_0BlockTensorData) + assertTrue(Requantizers.supports(TensorEncoding.BITNET_PLANES)) + assertTrue(!Requantizers.supports(TensorEncoding.Q4_K)) + assertFailsWith { Requantizers.quantize(TensorEncoding.Q4_K, values, Shape(2, 32)) } + } +}