From 54340345ad6815ab899b5150d15ffa0dc7b3bb66 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Sun, 30 Aug 2026 14:18:40 +0200 Subject: [PATCH] fix(memory,kernel): serve dense FP32 mapped/off-heap storage, not just Heap (#1217) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit TensorView.get()'s reference decode (readDense) and NarrowFloatDecoder.decodeAt only handled Storage.Heap — anything else (SegmentStorage, MappedFileStorage — a dequantized-at-load GGUF weight under a MemorySegmentTensorDataFactory context) threw UnsupportedOperationException, despite the class's own doc comment promising get() is "the correct, slow reference path" for every format. Fixed via Storage.copyInto, the primitive every storage kind already implements. Fp32ViewMatmulKernel (KernelPacks.install()'s dense-FP32 kernel) had the same Storage.Heap-only gap one layer up — even with the correctness fix, it silently fell back to the (now-correct but still per-element-decoding) reference kernel for any non-Heap operand, unlike FfmRowMajorMatmulKernel (the packed-encoding sibling), which already bridges MappedBufferStorage/DirectBufferStorage. Fixed by bridging a and b via one bulk copyInto snapshot per call (not per-element) before calling the existing heap-based SPI kernel. Found via a real Gemma 4 GGUF's per_layer_model_proj.weight (PLE context projection, dequantized to FP32, Segment-backed) through SKaiNET-transformers' Gemma4ChatModel.fromGguf: crashed in ~15s before this fix, ran real compute (not crashing) 15+ minutes in after the correctness fix alone, now dispatches to the real kernel instead of the reference fallback. New tests: JvmStorageTest (dense FP32 + narrow-float decode over SegmentStorage/MappedFileStorage), KernelPacksSegmentStorageTest (the pack kernel, not reference, now serves Segment-backed operands). Closes #1217. Co-Authored-By: Claude Sonnet 5 --- .../ainet/backend/api/kernel/KernelPacks.kt | 53 +++++++-- .../kernel/KernelPacksSegmentStorageTest.kt | 102 ++++++++++++++++++ .../kotlin/sk/ainet/lang/memory/TensorView.kt | 45 +++++++- .../sk/ainet/lang/memory/JvmStorageTest.kt | 62 +++++++++++ 4 files changed, 251 insertions(+), 11 deletions(-) create mode 100644 skainet-backends/skainet-backend-api/src/jvmTest/kotlin/sk/ainet/backend/api/kernel/KernelPacksSegmentStorageTest.kt diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelPacks.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelPacks.kt index 60a748558..e074f7bb2 100644 --- a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelPacks.kt +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelPacks.kt @@ -94,11 +94,7 @@ public class Fp32ViewMatmulKernel( val a = inputs[0]; val b = inputs[1] val m = a.shape[0]; val k = a.shape[1]; val n = b.shape[0] require(b.shape[1] == k) { "inner dimensions disagree: [${m}, ${k}] × [${n}, ${b.shape[1]}]" } - val aHeap = a.storage as? Storage.Heap ?: return fallback(inputs, out) - val bHeap = b.storage as? Storage.Heap ?: return fallback(inputs, out) val oHeap = out.storage as? Storage.Heap ?: return fallback(inputs, out) - val aBuf = aHeap.floats ?: return fallback(inputs, out) - val bBuf = bHeap.floats ?: return fallback(inputs, out) val oBuf = oHeap.floats ?: return fallback(inputs, out) // The SPI GEMM reads the weight input-major: b[p][j] = bBuf[bOffset + p * bStride + j]. // The dispatcher hands us the weight output-major ([n, k]) — which, when it is a transposed @@ -106,14 +102,59 @@ public class Fp32ViewMatmulKernel( // buffer's row stride. Anything else (a genuinely output-major buffer) would need a gather, // so it goes to the reference kernel instead of being silently mis-indexed. if (b.layout.strides[0] != 1) return fallback(inputs, out) + val (aBuf, aOffset, aStride) = denseFp32(a) ?: return fallback(inputs, out) + val (bBuf, bOffset, bStride) = denseFp32(b, requiredStride = b.layout.strides[1]) ?: return fallback(inputs, out) kernel.matmul( - a = aBuf, aOffset = aHeap.arrayOffset + a.layout.offsetElements.toInt(), aStride = a.layout.strides[0], - b = bBuf, bOffset = bHeap.arrayOffset + b.layout.offsetElements.toInt(), bStride = b.layout.strides[1], + a = aBuf, aOffset = aOffset, aStride = aStride, + b = bBuf, bOffset = bOffset, bStride = bStride, out = oBuf, outOffset = oHeap.arrayOffset + out.layout.offsetElements.toInt(), outStride = out.layout.strides[0], m = m, n = n, k = k, ) } + /** + * `(FloatArray, offset, stride)` for [view], zero-copy when it is already [Storage.Heap] + * (the original, only path); otherwise **one bulk snapshot** (never per-element — that is the + * ~1000×-slower decoding reference path this pack exists to avoid) via [Storage.copyInto], + * which every storage kind implements uniformly (`SegmentStorage`, `MappedFileStorage`, the + * NIO-buffer kinds, …). This is the same cost class as [sk.ainet.exec.kernel.FfmRowMajorMatmulKernel]'s + * existing heap-`ByteArray` staging path — one copy per kernel call, not per element. + * + * A weight this small a matmul reaches for is realistically the *whole* backing storage (a + * top-level loaded tensor, not a narrowed slice) — [requiredStride] guards that assumption: + * when the view's own row stride does not match what a plain row-major snapshot would produce, + * this bridge cannot represent it correctly, so it declines (→ fallback) rather than risk a + * silent-wrong-numbers bug. + */ + private fun denseFp32(view: TensorView, requiredStride: Int? = null): Triple? { + val heap = view.storage as? Storage.Heap + if (heap != null) { + val floats = heap.floats ?: return null + val stride = requiredStride ?: view.layout.strides.getOrElse(0) { 1 } + return Triple(floats, heap.arrayOffset + view.layout.offsetElements.toInt(), stride) + } + if (view.format.dtype != FP32 || view.format.encoding !is sk.ainet.lang.tensor.storage.TensorEncoding.Dense) return null + val rowStride = requiredStride ?: view.layout.strides.getOrElse(1) { view.shape[view.shape.rank - 1] } + if (requiredStride == null && !view.isContiguous) return null + val elementCount = view.elementCount + val byteWidth = 4 + val byteOffset = view.layout.offsetElements * byteWidth + val byteLength = (elementCount * byteWidth).toInt() + if (byteOffset + byteLength > view.storage.sizeBytes) return null + val bytes = ByteArray(byteLength) + view.storage.copyInto(bytes, 0, byteOffset, byteLength) + val floats = FloatArray(elementCount.toInt()) + for (i in floats.indices) { + val o = i * byteWidth + val bits = (bytes[o].toInt() and 0xFF) or + ((bytes[o + 1].toInt() and 0xFF) shl 8) or + ((bytes[o + 2].toInt() and 0xFF) shl 16) or + ((bytes[o + 3].toInt() and 0xFF) shl 24) + floats[i] = Float.fromBits(bits) + } + return Triple(floats, 0, rowStride) + } + private fun fallback(inputs: List, out: TensorView) { ReferenceMatmulKernel(key).run(inputs, out) } diff --git a/skainet-backends/skainet-backend-api/src/jvmTest/kotlin/sk/ainet/backend/api/kernel/KernelPacksSegmentStorageTest.kt b/skainet-backends/skainet-backend-api/src/jvmTest/kotlin/sk/ainet/backend/api/kernel/KernelPacksSegmentStorageTest.kt new file mode 100644 index 000000000..00f970629 --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/jvmTest/kotlin/sk/ainet/backend/api/kernel/KernelPacksSegmentStorageTest.kt @@ -0,0 +1,102 @@ +package sk.ainet.backend.api.kernel + +import java.lang.foreign.ValueLayout +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.SegmentStorage +import sk.ainet.lang.memory.Storage +import sk.ainet.lang.memory.TensorView +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.types.FP32 +import kotlin.math.abs +import kotlin.test.AfterTest +import kotlin.test.Test +import kotlin.test.assertIs +import kotlin.test.assertTrue + +/** + * A dense-FP32 weight whose storage is [SegmentStorage] (the shape a dequantized, mmap-context + * GGUF weight actually has — e.g. Gemma 4's `per_layer_model_proj.weight`) must still reach the + * pack's fast kernel, not fall to the decoding reference. Before this fix `Fp32ViewMatmulKernel` + * required `Storage.Heap` and fell back unconditionally otherwise — correct, but ~1000x slower + * (the crash this masked until #325/#341's readDense fix is a separate, already-fixed bug: the + * reference kernel didn't even support element access over Segment storage at all). + */ +@OptIn(ExperimentalMemoryApi::class) +class KernelPacksSegmentStorageTest { + + @AfterTest fun cleanup() { KernelDispatch.clearForTesting(); KernelRegistry.clearForTesting() } + + private class FakeProvider(override val name: String = "fake", override val priority: Int = 100) : KernelProvider { + var calls = 0 + override fun isAvailable(): Boolean = true + override fun matmulFp32(): Fp32MatmulKernel = object : Fp32MatmulKernel { + override fun matmul( + a: FloatArray, aOffset: Int, aStride: Int, + b: FloatArray, bOffset: Int, bStride: Int, + out: FloatArray, outOffset: Int, outStride: Int, + m: Int, n: Int, k: Int, + ) { + calls++ + for (i in 0 until m) for (j in 0 until n) { + var acc = 0f + for (p in 0 until k) acc += a[aOffset + i * aStride + p] * b[bOffset + p * bStride + j] + out[outOffset + i * outStride + j] = acc + } + } + } + } + + private fun segmentWeight(shape: Shape, values: FloatArray): TensorView { + val s = SegmentStorage.allocate(bytes = values.size.toLong() * 4) + for (i in values.indices) s.segment().setAtIndex(ValueLayout.JAVA_FLOAT, i.toLong(), values[i]) + return TensorView.dense(s, shape, FP32) + } + + private fun heapView(shape: Shape, values: FloatArray): TensorView = + TensorView.dense(Storage.Heap.wrap(values), shape, FP32) + + @Test + fun aSegmentBackedWeightStillReachesThePackKernel() { + val provider = FakeProvider() + KernelRegistry.register(provider) + KernelPacks.install(provider) + + val a = heapView(Shape(2, 3), floatArrayOf(1f, 2f, 3f, 4f, 5f, 6f)) + // weight as stored [k, n] = [3, 2], handed to the dispatcher transposed -> [2, 3] view, + // strides[0] == 1 (the case Fp32ViewMatmulKernel actually serves) + val wBuf = floatArrayOf(1f, 0.5f, 2f, 1.5f, 3f, 2.5f) + val w = segmentWeight(Shape(3, 2), wBuf).transpose() + val out = heapView(Shape(2, 2), FloatArray(4)) + + KernelDispatch.matmul(a, w, out) + + assertTrue(provider.calls > 0, "the pack kernel must have run, not the reference fallback") + val expected = FloatArray(4) + for (i in 0 until 2) for (j in 0 until 2) { + var acc = 0f + for (p in 0 until 3) acc += a.get(i, p) * wBuf[p * 2 + j] + expected[i * 2 + j] = acc + } + for (i in expected.indices) { + assertTrue(abs(out.get(i / 2, i % 2) - expected[i]) < 1e-4f, "element $i: ${out.get(i / 2, i % 2)} vs ${expected[i]}") + } + } + + @Test + fun aSegmentBackedActivationStillReachesThePackKernel() { + val provider = FakeProvider() + KernelRegistry.register(provider) + KernelPacks.install(provider) + + val a = segmentWeight(Shape(1, 3), floatArrayOf(1f, 2f, 3f)) + val wBuf = floatArrayOf(1f, 0.5f, 2f, 1.5f, 3f, 2.5f) + val w = heapView(Shape(3, 2), wBuf).transpose() + val out = heapView(Shape(1, 2), FloatArray(2)) + + KernelDispatch.matmul(a, w, out) + + assertTrue(provider.calls > 0, "the pack kernel must have run, not the reference fallback") + assertTrue(abs(out.get(0, 0) - (1f * 1 + 2f * 2 + 3f * 3)) < 1e-4f) + assertTrue(abs(out.get(0, 1) - (1f * 0.5f + 2f * 1.5f + 3f * 2.5f)) < 1e-4f) + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/TensorView.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/TensorView.kt index ccebc84fa..c1a972a82 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/TensorView.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/TensorView.kt @@ -256,7 +256,34 @@ public class TensorView( else throw UnsupportedOperationException("dense element access over byte storage needs a decoder") } } - else -> throw UnsupportedOperationException("element access over ${s::class.simpleName} needs a platform reader (use a kernel)") + else -> readDenseFromBytes(s, flat) + } + + /** + * Dense decode for a non-[Storage.Heap] backing (OffHeap/Mapped/Device — `SegmentStorage`, + * `MappedFileStorage`, a future device binding) via [Storage.copyInto], the one primitive every + * storage kind implements uniformly regardless of platform. Slow (a per-element snapshot copy) + * by design — this is the reference path the class doc promises ("correct for every format + * because it decodes"); a production kernel unwraps the storage once per call instead. + * + * A narrow float (FP16/BF16) never reaches here — `get()` always hands those a [decoder] + * (`NarrowFloatDecoder`) before falling through to [readDense] — but the two cases are handled + * defensively anyway so this stays correct if that invariant is ever relaxed. + */ + private fun readDenseFromBytes(s: Storage, flat: Long): Float { + val dtype = format.dtype + val width = dtype.sizeInBytes + val buf = ByteArray(width) + s.copyInto(buf, 0, flat * width, width) + var bits = 0L + for (i in 0 until width) bits = bits or ((buf[i].toLong() and 0xFF) shl (8 * i)) + return when (dtype) { + sk.ainet.lang.types.FP32 -> Float.fromBits(bits.toInt()) + sk.ainet.lang.types.FP64 -> Double.fromBits(bits).toFloat() + sk.ainet.lang.types.BF16 -> sk.ainet.lang.types.Bf16Codec.decode((bits and 0xFFFF).toInt()) + sk.ainet.lang.types.FP16 -> sk.ainet.lang.types.Fp16Codec.decode((bits and 0xFFFF).toInt()) + else -> bits.toFloat() + } } /** Every element in row-major order, decoded — the reference materialization. */ @@ -406,10 +433,18 @@ public class NarrowFloatDecoder(private val codec: sk.ainet.lang.types.NarrowFlo override fun decodeElement(storage: Storage, layout: Layout, flatElementIndex: Long): Float = decodeAt(storage, flatElementIndex) private fun decodeAt(storage: Storage, elementIndex: Long): Float { - val heap = storage as? Storage.Heap ?: throw UnsupportedOperationException("narrow-float views need heap storage in this milestone") - val bytes = heap.bytes ?: throw UnsupportedOperationException("narrow-float views need byte storage") - val off = heap.arrayOffset + (elementIndex * codec.bytesPerElement).toInt() - val bits = (bytes[off].toInt() and 0xFF) or ((bytes[off + 1].toInt() and 0xFF) shl 8) + val heap = storage as? Storage.Heap + val bits = if (heap != null) { + val bytes = heap.bytes ?: throw UnsupportedOperationException("narrow-float views need byte storage") + val off = heap.arrayOffset + (elementIndex * codec.bytesPerElement).toInt() + (bytes[off].toInt() and 0xFF) or ((bytes[off + 1].toInt() and 0xFF) shl 8) + } else { + // Non-heap (OffHeap/Mapped/Device): the same copyInto snapshot readDenseFromBytes uses, + // every storage kind implements it uniformly regardless of platform binding. + val buf = ByteArray(codec.bytesPerElement) + storage.copyInto(buf, 0, elementIndex * codec.bytesPerElement, codec.bytesPerElement) + (buf[0].toInt() and 0xFF) or ((buf[1].toInt() and 0xFF) shl 8) + } return codec.decode(bits) } } diff --git a/skainet-lang/skainet-lang-core/src/jvmTest/kotlin/sk/ainet/lang/memory/JvmStorageTest.kt b/skainet-lang/skainet-lang-core/src/jvmTest/kotlin/sk/ainet/lang/memory/JvmStorageTest.kt index e7352d4a1..f0bc79531 100644 --- a/skainet-lang/skainet-lang-core/src/jvmTest/kotlin/sk/ainet/lang/memory/JvmStorageTest.kt +++ b/skainet-lang/skainet-lang-core/src/jvmTest/kotlin/sk/ainet/lang/memory/JvmStorageTest.kt @@ -2,8 +2,10 @@ package sk.ainet.lang.memory 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.TensorId import sk.ainet.lang.tensor.storage.MemoryDomain +import sk.ainet.lang.types.FP32 import java.lang.foreign.Arena import java.lang.foreign.MemorySegment import java.lang.foreign.ValueLayout @@ -11,6 +13,7 @@ import java.nio.ByteBuffer import java.nio.ByteOrder import java.nio.file.Files import kotlin.test.Test +import kotlin.test.assertContentEquals import kotlin.test.assertEquals import kotlin.test.assertFailsWith import kotlin.test.assertFalse @@ -90,6 +93,65 @@ class JvmStorageTest { assertFailsWith { m.segment() } } + /** + * `TensorView.get()` over dense FP32 backed by [SegmentStorage] / [MappedFileStorage] — was + * `UnsupportedOperationException("element access over SegmentStorage needs a platform reader + * (use a kernel)")` (the reference-kernel fallback assumed every non-Heap storage had a real + * kernel to serve it; a MAPPED weight with no matching registered kernel key had none). Found + * via a real Gemma 4 GGUF's `per_layer_model_proj.weight` (a MAPPED dense weight, PLE) falling + * to `ReferenceMatmulKernel` and throwing instead of just being slow. + */ + @Test + fun denseFp32ViewReadsThroughSegmentStorage() { + val s = SegmentStorage.allocate(bytes = 4L * 6) + for (i in 0 until 6) s.segment().setAtIndex(ValueLayout.JAVA_FLOAT, i.toLong(), i.toFloat() + 0.5f) + val v = TensorView.dense(s, Shape(2, 3), FP32, TensorId.parse("model.ple_proj")) + assertEquals(0.5f, v.get(0, 0)); assertEquals(3.5f, v.get(1, 0)); assertEquals(5.5f, v.get(1, 2)) + assertContentEquals(floatArrayOf(0.5f, 1.5f, 2.5f, 3.5f, 4.5f, 5.5f), v.toFloatArray()) + s.close() + } + + @Test + fun denseFp32ViewReadsThroughMappedFileStorage() { + val f = Files.createTempFile("skainet-mapped-dense", ".bin"); f.toFile().deleteOnExit() + val bytes = ByteArray(4 * 4) + val bb = ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN) + for (i in 0 until 4) bb.putFloat(i * 4, i.toFloat() * 10f) + Files.write(f, bytes) + val m = MappedFileStorage.map(f, fileOffset = 0, length = 16, origin = TensorId.parse("model.w")) + val v = TensorView.dense(m, Shape(4), FP32) + assertEquals(0f, v.get(0)); assertEquals(10f, v.get(1)); assertEquals(30f, v.get(3)) + m.close() + } + + /** + * Same gap, the narrow-float (BF16/FP16) decode path: `NarrowFloatDecoder.decodeAt` threw + * "narrow-float views need heap storage in this milestone" for anything but [Storage.Heap]. + */ + @Test + fun narrowFloatViewReadsThroughSegmentStorage() { + val codec = sk.ainet.lang.types.Bf16Codec + val s = SegmentStorage.allocate(bytes = 2L * 3) + val values = floatArrayOf(1.0f, -2.5f, 100.0f) + for (i in values.indices) { + val bits = codec.encode(values[i]) + s.segment().set(ValueLayout.JAVA_BYTE, (i * 2).toLong(), (bits and 0xFF).toByte()) + s.segment().set(ValueLayout.JAVA_BYTE, (i * 2 + 1).toLong(), ((bits ushr 8) and 0xFF).toByte()) + } + val shape = Shape(3) + val view = TensorView( + shape = shape, + format = Format(codec.dtype, sk.ainet.lang.tensor.storage.TensorEncoding.Dense(2)), + layout = Layout(shape = shape, strides = Layout.rowMajorStrides(shape), elementBytes = 2), + storage = s, + decoder = NarrowFloatDecoder(codec), + ) + for (i in values.indices) { + assertTrue(kotlin.math.abs(view.get(i) - values[i]) <= kotlin.math.abs(values[i]) * 0.01f + 0.01f) + } + s.close() + } + @Test fun directBufferStorageAllocatesBorrowsAndSlices() { val sink = RecordingTraceSink()