diff --git a/CHANGELOG.md b/CHANGELOG.md index f86ca0cb1..046ba9605 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,21 @@ ## [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. +- **`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. +- **`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..15dbc11fe --- /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 in progress + +Implementation: Tier 0 (this proposal's PR); Tier 1 `skainet-backend-wasm-simd`, Tier 2 quantize-at-load and the consumer work in SKaiNET-transformers follow as separate pull requests, 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-cli 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-cli quantize` 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`, `skainet-cli quantize`. +. 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-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-lang/skainet-lang-core/api/jvm/skainet-lang-core.api b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api index 557860fc1..fd68c6fda 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 @@ -5154,10 +5154,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; 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/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) } + } +}