fix(memory,kernel): serve dense FP32 mapped/off-heap storage, not just Heap

skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelPacks.kt

@@ -94,26 +94,67 @@ 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
         // *view* of a contiguous [k, n] buffer, means strides[0] == 1 and strides[1] is that
         // 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<FloatArray, Int, Int>? {
+        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<TensorView>, out: TensorView) {
         ReferenceMatmulKernel(key).run(inputs, out)
     }

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)
+    }
+}

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)
     }
 }

skainet-lang/skainet-lang-core/src/jvmTest/kotlin/sk/ainet/lang/memory/JvmStorageTest.kt

@@ -2,15 +2,18 @@ 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
 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<StorageClosedException> { 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()