fix(lang,backend-cpu): free transpose for narrow-float weights

skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt

@@ -13,6 +13,7 @@ import sk.ainet.backend.api.kernel.KernelRegistry
 import sk.ainet.lang.tensor.data.RowDequantSource
 import sk.ainet.lang.tensor.data.FloatArrayTensorData
 import sk.ainet.lang.tensor.data.IntArrayTensorData
+import sk.ainet.lang.tensor.data.NarrowFloatInputMajorTensorData
 import sk.ainet.lang.tensor.data.Q4_0TensorData
 import sk.ainet.lang.tensor.data.Q8_0TensorData
 import sk.ainet.lang.tensor.data.Q4_KTensorData
@@ -624,6 +625,17 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory
                 // packed type that can be a matmul weight. See transformers #178.
                 is Q8_0TensorData -> return newTensor(Q8_0BlockTensorData(Shape(cols, rows), d.packedData) as TensorData<T, V>, tensor.dtype, tensor)
                 is Q4_0TensorData -> return newTensor(Q4_0BlockTensorData(Shape(cols, rows), d.packedData) as TensorData<T, V>, tensor.dtype, tensor)
+                // Narrow floats (FP16/BF16) relaid input-major at load: the transpose is the
+                // same buffer read with the other shape's strides, so hand back an ordinary
+                // dense narrow tensor over it. This is what lets a KEEP_NATIVE weight survive
+                // `Linear.onForward`'s `weight.t()` and reach the narrow matmul kernel — see
+                // issue #888. Note the asymmetry with the block-quant arms above: only the
+                // *input-major* type is safe to reinterpret. A row-major narrow buffer falls
+                // through to the generic path on purpose, because swapping its shape would
+                // silently yield a different matrix rather than the transpose.
+                // Lives only here: DefaultCpuOpsJvm.transpose intercepts nothing that would
+                // shadow this case, so the JVM falls through to this arm too.
+                is NarrowFloatInputMajorTensorData -> return newTensor(d.transposedView() as TensorData<T, V>, tensor.dtype, tensor)
                 else -> {}
             }
         }

skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/NarrowFloatTransposeDispatchTest.kt

@@ -0,0 +1,175 @@
+package sk.ainet.exec.tensor.ops
+
+import kotlin.math.abs
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertSame
+import kotlin.test.assertTrue
+import sk.ainet.context.DirectCpuExecutionContext
+import sk.ainet.lang.tensor.Shape
+import sk.ainet.lang.tensor.Tensor
+import sk.ainet.lang.tensor.data.NarrowFloatDenseTensorData
+import sk.ainet.lang.tensor.data.NarrowFloatInputMajorTensorData
+import sk.ainet.lang.tensor.data.NarrowFloatTensorData
+import sk.ainet.lang.tensor.data.TensorData
+import sk.ainet.lang.tensor.matmul
+import sk.ainet.lang.tensor.t
+import sk.ainet.lang.types.Bf16Codec
+import sk.ainet.lang.types.FP32
+import sk.ainet.lang.types.Fp16Codec
+import sk.ainet.lang.types.NarrowFloatCodec
+
+/**
+ * Proves a KEEP_NATIVE weight survives `ctx.ops.matmul(x, ops.transpose(W))` still packed, so it
+ * reaches the narrow matmul kernel instead of being widened elementwise (issue #888).
+ *
+ * Sibling of [PackedMatmulDispatchTest], and deliberately in `commonTest` for the same reason:
+ * `DefaultCpuOps` and `DefaultCpuOpsJvm` intercept transpose separately, so only running on both
+ * jvmTest and linuxX64Test proves both arms are wired.
+ *
+ * Weights are stored `[out, in]` as they are on disk, then transposed — the exact shape of what
+ * `Linear.onForward` does on every forward pass.
+ */
+class NarrowFloatTransposeDispatchTest {
+
+    private val ctx = DirectCpuExecutionContext()
+
+    private val outDim = 3
+    private val inDim = 4
+
+    /** `[out, in]`, all exactly representable in binary16 and bfloat16 alike. */
+    private val weightRowMajor = floatArrayOf(
+        1.0f, 2.0f, -1.0f, 0.5f,
+        4.0f, -0.5f, 2.0f, 1.0f,
+        -2.0f, 1.0f, 0.25f, 8.0f,
+    )
+
+    private val input = floatArrayOf(1.0f, -1.0f, 2.0f, 0.5f)
+
+    private fun pack(values: FloatArray, codec: NarrowFloatCodec): ByteArray {
+        val out = ByteArray(values.size * 2)
+        for (i in values.indices) {
+            val bits = codec.encode(values[i])
+            out[i * 2] = (bits and 0xFF).toByte()
+            out[i * 2 + 1] = ((bits ushr 8) and 0xFF).toByte()
+        }
+        return out
+    }
+
+    /** `y[j] = sum_i x[i] * W[j][i]` — what `x.matmul(W.t())` must produce. */
+    private fun reference(): FloatArray = FloatArray(outDim) { j ->
+        var acc = 0.0f
+        for (i in 0 until inDim) acc += input[i] * weightRowMajor[j * inDim + i]
+        acc
+    }
+
+    private fun inputTensor(): Tensor<FP32, Float> =
+        ctx.fromFloatArray(Shape(1, inDim), FP32::class, input)
+
+    @Suppress("UNCHECKED_CAST")
+    private fun inputMajorWeight(codec: NarrowFloatCodec): Tensor<FP32, Float> {
+        val data = NarrowFloatInputMajorTensorData.fromRowMajor(
+            Shape(outDim, inDim), pack(weightRowMajor, codec), codec,
+        )
+        return ctx.fromData(data as TensorData<FP32, Float>, FP32::class)
+    }
+
+    private fun assertClose(expected: FloatArray, actual: FloatArray, label: String) {
+        assertEquals(expected.size, actual.size, "$label: length")
+        for (i in expected.indices) {
+            assertTrue(
+                abs(expected[i] - actual[i]) < 1e-4f,
+                "$label: element $i expected ${expected[i]} but was ${actual[i]}",
+            )
+        }
+    }
+
+    @Test
+    fun `transpose of an input-major weight stays narrow and shares the buffer`() {
+        for (codec in listOf(Fp16Codec, Bf16Codec)) {
+            val w = inputMajorWeight(codec)
+            val original = (w.data as NarrowFloatTensorData).packedData
+
+            val wt = w.t()
+
+            assertTrue(
+                wt.data is NarrowFloatTensorData,
+                "${codec.dtype.name}: transpose widened the weight — the kernel is unreachable",
+            )
+            assertTrue(
+                wt.data is NarrowFloatDenseTensorData,
+                "${codec.dtype.name}: expected a plain dense narrow tensor after transpose",
+            )
+            assertEquals(Shape(inDim, outDim), wt.shape, "${codec.dtype.name}: shape must swap")
+            assertEquals(
+                codec, (wt.data as NarrowFloatTensorData).codec,
+                "${codec.dtype.name}: codec must survive the transpose",
+            )
+            assertSame(
+                original, (wt.data as NarrowFloatTensorData).packedData,
+                "${codec.dtype.name}: transpose must not copy — copying per forward is the bug",
+            )
+        }
+    }
+
+    @Test
+    fun `matmul through the transpose matches the fp32 reference`() {
+        val expected = reference()
+        for (codec in listOf(Fp16Codec, Bf16Codec)) {
+            val y = inputTensor().matmul(inputMajorWeight(codec).t())
+            assertEquals(Shape(1, outDim), y.shape, "${codec.dtype.name}: output shape")
+            assertClose(expected, y.data.copyToFloatArray(), codec.dtype.name)
+        }
+    }
+
+    @Test
+    fun `a row-major narrow weight is left alone by the lazy transpose`() {
+        // The safety property. Only the input-major type may be reinterpreted; swapping the shape
+        // of a row-major narrow buffer would silently produce a different matrix. The generic path
+        // is slow, but it is correct — and correctness is what must not regress here.
+        @Suppress("UNCHECKED_CAST")
+        val rowMajor = ctx.fromData(
+            NarrowFloatDenseTensorData(
+                Shape(outDim, inDim), pack(weightRowMajor, Fp16Codec), Fp16Codec,
+            ) as TensorData<FP32, Float>,
+            FP32::class,
+        )
+
+        val wt = rowMajor.t()
+        assertEquals(Shape(inDim, outDim), wt.shape)
+
+        // Whatever representation the generic path chose, the values must be the true transpose.
+        val got = wt.data.copyToFloatArray()
+        for (j in 0 until outDim) {
+            for (i in 0 until inDim) {
+                assertEquals(
+                    weightRowMajor[j * inDim + i], got[i * outDim + j],
+                    "row-major transpose wrong at [$i, $j]",
+                )
+            }
+        }
+        assertClose(reference(), inputTensor().matmul(wt).data.copyToFloatArray(), "row-major")
+    }
+
+    @Test
+    fun `the two codecs disagree on identical bytes`() {
+        // Vacuity guard: both formats are 2 bytes per element, so nothing above would catch a
+        // dispatch that picked the kernel by byte width instead of by codec.
+        val bytes = pack(weightRowMajor, Fp16Codec)
+
+        @Suppress("UNCHECKED_CAST")
+        fun matmulAs(codec: NarrowFloatCodec): FloatArray {
+            val data = NarrowFloatInputMajorTensorData.fromRowMajor(Shape(outDim, inDim), bytes, codec)
+            val w = ctx.fromData(data as TensorData<FP32, Float>, FP32::class)
+            return inputTensor().matmul(w.t()).data.copyToFloatArray()
+        }
+
+        val asFp16 = matmulAs(Fp16Codec)
+        val asBf16 = matmulAs(Bf16Codec)
+        assertTrue(
+            asFp16.indices.any { abs(asFp16[it] - asBf16[it]) > 1e-3f },
+            "reading the same bytes under both codecs produced the same result — " +
+                "the codec assertions in this class would prove nothing",
+        )
+    }
+}

skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt

@@ -239,6 +239,8 @@ internal class DefaultCpuOpsJvm(
                 @Suppress("UNCHECKED_CAST")
                 return newTensor(transposed as TensorData<T, V>, tensor.dtype, tensor)
             }
+            // Narrow-float input-major lazy transpose is handled in DefaultCpuOpsBase too —
+            // nothing above intercepts it, so it falls through. Issue #888.
             // Q6_K / Q5_1 / Q5_0 lazy transpose is handled in DefaultCpuOpsBase
             // (block-major, shared with Native); the JVM ops don't intercept them here.
             // MemorySegment FP32 fast path: physical transpose via SIMD.

skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/NarrowFloatLoadPolicy.kt

@@ -35,6 +35,15 @@ public enum class NarrowFloatLoadPolicy {
      *
      * **Caveat**: any non-matmul op touching the tensor pays a per-element decode via `get`. Worth
      * it when the hot path is matmul-dominated (the typical transformer case), not otherwise.
+     *
+     * **Layout caveat (issue #888).** This loader emits row-major bytes, matching the file. But
+     * projections are stored `[out, in]` while the narrow matmul dispatch needs `[in, out]`, so a
+     * consumer that goes through `Linear` — which transposes the weight on every forward — will
+     * see that transpose fall to the generic elementwise path and widen the tensor anyway, at a
+     * cost far exceeding what KEEP_NATIVE saves. Consumers running such a hot path should relay
+     * the bytes once with `NarrowFloatInputMajorTensorData.fromRowMajor`, which makes the
+     * transpose a zero-copy view. Tensors that are gathered rather than multiplied — embedding
+     * tables above all — should stay row-major, since input-major storage strides their row reads.
      */
     KEEP_NATIVE,
 }

skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api

@@ -3095,6 +3095,24 @@ public final class sk/ainet/lang/tensor/data/NarrowFloatDenseTensorData$Companio
 	public final fun fromFloatArray (Lsk/ainet/lang/tensor/Shape;[FLsk/ainet/lang/types/NarrowFloatCodec;)Lsk/ainet/lang/tensor/data/NarrowFloatDenseTensorData;
 }
 
+public final class sk/ainet/lang/tensor/data/NarrowFloatInputMajorTensorData : sk/ainet/lang/tensor/data/NarrowFloatTensorData {
+	public static final field Companion Lsk/ainet/lang/tensor/data/NarrowFloatInputMajorTensorData$Companion;
+	public fun <init> (Lsk/ainet/lang/tensor/Shape;[BLsk/ainet/lang/types/NarrowFloatCodec;)V
+	public fun copyToFloatArray ()[F
+	public fun get ([I)Ljava/lang/Float;
+	public synthetic fun get ([I)Ljava/lang/Object;
+	public fun getCodec ()Lsk/ainet/lang/types/NarrowFloatCodec;
+	public fun getPackedData ()[B
+	public fun getShape ()Lsk/ainet/lang/tensor/Shape;
+	public fun set ([IF)V
+	public synthetic fun set ([ILjava/lang/Object;)V
+	public final fun transposedView ()Lsk/ainet/lang/tensor/data/NarrowFloatDenseTensorData;
+}
+
+public final class sk/ainet/lang/tensor/data/NarrowFloatInputMajorTensorData$Companion {
+	public final fun fromRowMajor (Lsk/ainet/lang/tensor/Shape;[BLsk/ainet/lang/types/NarrowFloatCodec;)Lsk/ainet/lang/tensor/data/NarrowFloatInputMajorTensorData;
+}
+
 public abstract interface class sk/ainet/lang/tensor/data/NarrowFloatTensorData : sk/ainet/lang/tensor/data/TensorData {
 	public static final field BYTES_PER_ELEMENT I
 	public static final field Companion Lsk/ainet/lang/tensor/data/NarrowFloatTensorData$Companion;

skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/NarrowFloatTensorData.kt

@@ -161,5 +161,136 @@ public class Fp16DenseTensorData(
     }
 }
 
+/**
+ * A rank-2 narrow-float weight whose bytes are stored **input-major**: the element at logical
+ * `[row, col]` sits at flat byte index `(col * rows + row) * 2`, the transpose of the usual
+ * row-major order.
+ *
+ * ### Why this exists
+ *
+ * Projections are stored `[out, in]` on disk, but `chooseQuantizedMatmul` needs `[in, out]`, so
+ * every `Linear.onForward` calls `weight.t()` first. Transposing a row-major narrow tensor has no
+ * fast path — it walks the tensor elementwise through boxed `get()` and widens to FP32, which at
+ * real projection sizes costs hundreds of milliseconds to seconds *per weight, per token*. That
+ * made KEEP_NATIVE slower than not using it at all.
+ *
+ * Storing the bytes input-major once, at load, makes that transpose free: input-major storage of
+ * `[out, in]` **is** row-major storage of `[in, out]`, so [transposedView] hands back an ordinary
+ * [NarrowFloatDenseTensorData] over the very same buffer. No copy, and — unlike the lazy transpose
+ * used for the K-quants — element access stays correct on both sides, because each type indexes
+ * the shared bytes with the strides its own shape implies.
+ *
+ * The layout is carried in the type rather than a flag so that a bare shape swap cannot be applied
+ * to a row-major buffer by accident. Doing that would not throw; it would silently reinterpret the
+ * weight as a different matrix.
+ *
+ * Build these with [fromRowMajor], which performs the one-off relayout.
+ */
+public class NarrowFloatInputMajorTensorData(
+    initialShape: Shape,
+    private val data: ByteArray,
+    override val codec: NarrowFloatCodec,
+) : NarrowFloatTensorData {
+
+    override val shape: Shape = Shape(initialShape.dimensions.copyOf())
+    override val packedData: ByteArray get() = data
+
+    private val rows: Int = shape.dimensions[0]
+    private val cols: Int = shape.dimensions[1]
+
+    init {
+        require(shape.dimensions.size == 2) {
+            "Input-major layout is only defined for rank-2 weights, got shape $shape"
+        }
+        val requiredBytes = shape.volume * NarrowFloatTensorData.BYTES_PER_ELEMENT
+        require(data.size >= requiredBytes) {
+            "Data size ${data.size} is less than required $requiredBytes bytes " +
+                "for ${shape.volume} ${codec.dtype.name} elements"
+        }
+    }
+
+    /** Flat index of logical `[row, col]` under input-major storage. */
+    private fun flatIndex(indices: IntArray): Int {
+        require(indices.size == 2) {
+            "Number of indices (${indices.size}) must match tensor dimensions (2)"
+        }
+        val row = indices[0]
+        val col = indices[1]
+        require(row in 0 until rows) { "Index $row out of bounds for dimension 0 with size $rows" }
+        require(col in 0 until cols) { "Index $col out of bounds for dimension 1 with size $cols" }
+        return col * rows + row
+    }
+
+    override fun get(vararg indices: Int): Float {
+        val byteIdx = flatIndex(indices) * NarrowFloatTensorData.BYTES_PER_ELEMENT
+        val lo = data[byteIdx].toInt() and 0xFF
+        val hi = data[byteIdx + 1].toInt() and 0xFF
+        return codec.decode((hi shl 8) or lo)
+    }
+
+    override fun set(vararg indices: Int, value: Float) {
+        val byteIdx = flatIndex(indices) * NarrowFloatTensorData.BYTES_PER_ELEMENT
+        val bits = codec.encode(value)
+        data[byteIdx] = (bits and 0xFF).toByte()
+        data[byteIdx + 1] = ((bits ushr 8) and 0xFF).toByte()
+    }
+
+    /** Decodes in logical row-major order, so callers see the same values a dense tensor would. */
+    override fun copyToFloatArray(): FloatArray {
+        val out = FloatArray(shape.volume)
+        var dst = 0
+        for (row in 0 until rows) {
+            for (col in 0 until cols) {
+                val byteIdx = (col * rows + row) * NarrowFloatTensorData.BYTES_PER_ELEMENT
+                val lo = data[byteIdx].toInt() and 0xFF
+                val hi = data[byteIdx + 1].toInt() and 0xFF
+                out[dst++] = codec.decode((hi shl 8) or lo)
+            }
+        }
+        return out
+    }
+
+    /**
+     * The transpose, sharing this instance's buffer — no copy. Input-major `[rows, cols]` is
+     * row-major `[cols, rows]`, so the result is an ordinary dense narrow tensor that
+     * `chooseQuantizedMatmul` accepts directly.
+     */
+    public fun transposedView(): NarrowFloatDenseTensorData =
+        NarrowFloatDenseTensorData(Shape(cols, rows), data, codec)
+
+    public companion object {
+        /**
+         * Relayout `rowMajorBytes` (logical `[rows, cols]`, row-major) into input-major order.
+         * This is the one-off cost that buys a free transpose on every later forward pass.
+         */
+        public fun fromRowMajor(
+            shape: Shape,
+            rowMajorBytes: ByteArray,
+            codec: NarrowFloatCodec,
+        ): NarrowFloatInputMajorTensorData {
+            require(shape.dimensions.size == 2) {
+                "Input-major layout is only defined for rank-2 weights, got shape $shape"
+            }
+            val rows = shape.dimensions[0]
+            val cols = shape.dimensions[1]
+            val required = shape.volume * NarrowFloatTensorData.BYTES_PER_ELEMENT
+            require(rowMajorBytes.size >= required) {
+                "Data size ${rowMajorBytes.size} is less than required $required bytes " +
+                    "for ${shape.volume} ${codec.dtype.name} elements"
+            }
+            val out = ByteArray(required)
+            for (row in 0 until rows) {
+                for (col in 0 until cols) {
+                    val src = (row * cols + col) * NarrowFloatTensorData.BYTES_PER_ELEMENT
+                    val dst = (col * rows + row) * NarrowFloatTensorData.BYTES_PER_ELEMENT
+                    out[dst] = rowMajorBytes[src]
+                    out[dst + 1] = rowMajorBytes[src + 1]
+                }
+            }
+            return NarrowFloatInputMajorTensorData(shape, out, codec)
+        }
+    }
+}
+
 /** Decode a [NarrowFloatTensorData] to a fresh FloatArray. */
 public fun NarrowFloatTensorData.toFloatArray(): FloatArray = copyToFloatArray()