Skip to content

feat(backend): SKEEP-007 Tier 0 -- keep the wasm lane off the reference kernel - #1352

Merged
michalharakal merged 2 commits into
SKaiNET-developers:developfrom
michalharakal:feature/skeep-007-browser-lane-tier0
Oct 16, 2026
Merged

michalharakal merged 2 commits into
SKaiNET-developers:developfrom
michalharakal:feature/skeep-007-browser-lane-tier0

Conversation

@michalharakal

Copy link
Copy Markdown
Contributor

SKEEP-007 "The browser lane" is the design record for making inference in a browser usable and quick. This PR adds the proposal page and lands its Tier 0: the common-code fixes that keep the scalar wasm lane off the decoding reference kernel and out of memory trouble. The SIMD128 sidecar kernel pack (Tier 1) and quantize-at-load for SafeTensors (Tier 2) follow as separate pull requests, as the page lays out.

What was wrong

  • matmulWeightTransposed on a canonical (row-major) packed weight — every weight loaded from a GGUF — reached KernelDispatch on a target without a row-major kernel pack and ran the per-element decoding reference kernel on every call (the dispatcher's own "~1000x slower" warning). The scalar packed kernels every target carries are keyed on kernel feed order and were never reached from this path; the prepackedWeights cache declared for exactly that relayout had no common-code user.
  • A dense FP32 [out, in] weight was transposed into a cached heap copy before first use on every tier, doubling resident weight memory for a dense encoder (about 0.5 GB for a 22-layer, 768-wide model).
  • A 16-bit token table was widened to FP32 at load although ops.gather and the transformers Embedding already read rows from a RowDequantSource.
  • The web samples each carried their own in-memory RandomAccessSource; JsBlobRandomAccessSource converted buffers one interop call per byte and logged every read.

What changes

  • Dispatch order in DefaultCpuOpsBase.matmulWeightTransposed for packed weights: an exact kernel for the canonical view (JVM FFM, Android JNI, the coming browser pack) → the one-time cached relayout into feed order served by the scalar packed kernels → views and the reference kernel last. The JVM overrides this method and is unaffected; ternary formats have row-major kernels and take the first branch as before.
  • Direct dense path: x · Wᵀ for heap FP32 [out, in] weights is computed from the weight's own rows with a 4 × 4 tile and one sequential accumulator per output, bit-identical to transpose-then-matmul. preferDirectWeightTransposed() is the hook; DefaultCpuOpsJvm and AccelerateCpuOps return false and keep their cached transpose because their [k, n] kernels are vectorised.
  • NarrowFloatDenseTensorData : RowDequantSource (FP16 and BF16 tables serve rows on demand).
  • sk.ainet.io.ByteArrayRandomAccessSource in skainet-io-core, single array or equal-sized chunks; JsBlobRandomAccessSource drops its per-read console logging (the per-byte conversion stays: Kotlin/Wasm 2.4 has no typed-array adapter; Tier 1's linear-memory bridge replaces it).
  • skainet-backends/benchmarks/wasm-bench: a Kotlin/Wasm executable (Node and browser) timing the matmul shapes of Laya, LEAF, SmolLM2 and Qwen3-0.6B as dense FP32 and packed Q8_0, printing [SKAINET_PERF_WASM] lines. Baseline numbers are recorded in the SKEEP page.
  • Docs: SKEEP-007 page, index and nav; eager-execution.adoc and kernel-selection.adoc describe the new wasm behaviour; CHANGELOG Unreleased.

Tests

  • DenseWeightTransposedDirectTest: tiled row-dot bit-identical to the sequential definition for every tile remainder and with array offsets; the primitive agrees with transpose-then-matmul for rank-1, 2 and 3 activations.
  • PackedWeightKernelOrderTest: a canonical Q8_0 weight matches the relayout path, the per-weight cache is hit on the second call, and a RecordingTraceSink on KernelDispatch.defaultSink sees no reference-kernel run.
  • NarrowFloatRowDequantTest, ByteArrayRandomAccessSourceTest.

Verification

  • scripts/pr-gate.sh --quick (every module's jvmTest + apiCheck): passed.
  • :skainet-backends:skainet-backend-cpu:linuxX64Test (full, includes the golden packed-encoding parity): passed.
  • wasmJsTest for skainet-lang-core, skainet-io-core, skainet-backend-cpu (ChromeHeadless): passed.
  • wasm-bench on Node: baseline recorded in the SKEEP page (dense 0.84–1.96 GFLOPS, scalar Q8_0 0.35–0.63 GFLOPS; one Laya layer ≈ 1.34 s dense).
  • API dumps regenerated: additive only (ByteArrayRandomAccessSource, RowDequantSource on NarrowFloatDenseTensorData, preferDirectWeightTransposed).
  • The remaining legs (JS tests, wasmWasi, assemble, Java consumer tests) run in CI.

…ce kernel

Adds the SKEEP-007 proposal (the browser lane: SIMD128 sidecar kernels,
quantize-at-load) and lands its common-code tier:

- matmulWeightTransposed on a canonical packed weight prefers an exact
  kernel for the row-major view, then the one-time cached relayout into
  kernel feed order served by the scalar packed kernels, and only then the
  decoding reference kernel. Every non-JVM target took the reference path
  on every call before.
- a dense heap FP32 [out, in] weight is multiplied straight from its rows
  (4x4 tiles, one sequential accumulator per output, bit-identical to
  transpose-then-matmul) on tiers without a vectorised [k, n] kernel;
  preferDirectWeightTransposed() keeps the JVM and Apple tiers on their
  cached transpose.
- NarrowFloatDenseTensorData implements RowDequantSource so 16-bit token
  tables are read row by row instead of widened at load.
- ByteArrayRandomAccessSource in skainet-io-core (single array or chunks);
  JsBlobRandomAccessSource stops logging every read.
- benchmarks/wasm-bench: a Kotlin/Wasm executable timing the matmul shapes
  of the shipped models; the Node baseline is recorded in the SKEEP page.
Kotlin/JS keeps a Float local in double precision while the transpose-then-
matmul path rounds through a Float32Array on every step, so the direct row-dot
differs from it in the last bits there, as the golden tests already acknowledge
for that target. The tiled-vs-definition tests keep pinning the bits where bits
are stable.
@michalharakal
michalharakal merged commit beac717 into SKaiNET-developers:develop Oct 16, 2026
20 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant