From 5a5e5ab78216ff67fc46fd06b8d9bdf64c9d9c61 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Mon, 27 Jul 2026 15:07:23 +0200 Subject: [PATCH 1/2] fix(lang,backend-cpu): free transpose for narrow-float weights MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Projections are stored [out, in], but chooseQuantizedMatmul needs [in, out], so Linear.onForward transposes the weight on every call. Transposing a row-major narrow tensor had no fast path: it walked the tensor elementwise through boxed get() and widened to FP32. Measured on an i7-9750H that is 206 ms for a 2048x2048 projection and 4.4 s for 4096x11008 — per weight, per token. KEEP_NATIVE was therefore slower than not using it at all. Add NarrowFloatInputMajorTensorData: a rank-2 narrow weight whose bytes are stored input-major, built once via fromRowMajor(). Input-major storage of [rows, cols] is row-major storage of [cols, rows], so transposedView() hands back an ordinary NarrowFloatDenseTensorData over the same buffer — no copy, and element access stays correct on both sides, because each type indexes the shared bytes with the strides its own shape implies. That is the improvement over the K-quant lazy transpose, where get() on a transposed tensor is meaningless and only the kernel's direct packedData read is valid. Dispatch it from DefaultCpuOpsBase.transpose. Only the input-major type is reinterpreted; a row-major narrow buffer deliberately still falls through to the generic path, because swapping its shape would silently yield a different matrix rather than the transpose. The arm lives in the base class alone — DefaultCpuOpsJvm.transpose intercepts nothing that would shadow it, verified by disabling each arm in turn. Tests cover element access against the row-major original, buffer sharing, relayout round-trip, a square weight (where a wrong permutation still has the right byte count), rank and size rejection, and matmul through the transpose against an FP32 reference. Both narrow codecs are exercised, with vacuity guards asserting they disagree on identical bytes — both formats are 2 bytes per element, so a codec mix-up produces plausible wrong numbers rather than an error. The dispatch test sits in commonTest and passes on linuxX64 as well as the JVM. API change is purely additive. Refs #888. --- .../sk/ainet/exec/tensor/ops/DefaultCpuOps.kt | 12 ++ .../ops/NarrowFloatTransposeDispatchTest.kt | 175 ++++++++++++++++++ .../ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt | 2 + .../api/jvm/skainet-lang-core.api | 18 ++ .../lang/tensor/data/NarrowFloatTensorData.kt | 131 +++++++++++++ .../NarrowFloatInputMajorTensorDataTest.kt | 168 +++++++++++++++++ 6 files changed, 506 insertions(+) create mode 100644 skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/NarrowFloatTransposeDispatchTest.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/NarrowFloatInputMajorTensorDataTest.kt 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 f0d3aec8..dde63a47 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 @@ -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, tensor.dtype, tensor) is Q4_0TensorData -> return newTensor(Q4_0BlockTensorData(Shape(cols, rows), d.packedData) as TensorData, 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, tensor.dtype, tensor) else -> {} } } diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/NarrowFloatTransposeDispatchTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/NarrowFloatTransposeDispatchTest.kt new file mode 100644 index 00000000..6676d319 --- /dev/null +++ b/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 = + ctx.fromFloatArray(Shape(1, inDim), FP32::class, input) + + @Suppress("UNCHECKED_CAST") + private fun inputMajorWeight(codec: NarrowFloatCodec): Tensor { + val data = NarrowFloatInputMajorTensorData.fromRowMajor( + Shape(outDim, inDim), pack(weightRowMajor, codec), codec, + ) + return ctx.fromData(data as TensorData, 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::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::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", + ) + } +} 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 587ccc5d..4364b89b 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 @@ -239,6 +239,8 @@ internal class DefaultCpuOpsJvm( @Suppress("UNCHECKED_CAST") return newTensor(transposed as TensorData, 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. 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 84f41025..d31b2273 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 @@ -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 (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; 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 5f95a53b..1f8f916d 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 @@ -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() diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/NarrowFloatInputMajorTensorDataTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/NarrowFloatInputMajorTensorDataTest.kt new file mode 100644 index 00000000..08da6522 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/NarrowFloatInputMajorTensorDataTest.kt @@ -0,0 +1,168 @@ +package sk.ainet.lang.tensor.data + +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.types.Bf16Codec +import sk.ainet.lang.types.Fp16Codec +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertSame +import kotlin.test.assertTrue + +/** + * Pins [NarrowFloatInputMajorTensorData] — the relaid layout that makes transposing a KEEP_NATIVE + * weight free (issue #888). + * + * The invariant under test throughout: input-major storage of `[rows, cols]` **is** row-major + * storage of `[cols, rows]`, so [NarrowFloatInputMajorTensorData.transposedView] can share the + * buffer outright. Element access has to stay correct on both sides of that view, which is what + * distinguishes this from the K-quant lazy transpose (where `get()` on a transposed tensor is + * meaningless and only the kernel's direct `packedData` read is valid). + */ +class NarrowFloatInputMajorTensorDataTest { + + /** 2x3, all exactly representable in both narrow formats so comparisons can be exact. */ + private val rows = 2 + private val cols = 3 + private val values = floatArrayOf( + 1.0f, 2.0f, 4.0f, + 8.0f, 0.5f, -2.0f, + ) + + private fun rowMajorBytes(codec: sk.ainet.lang.types.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 + } + + @Test + fun `element access matches the row-major original`() { + val data = NarrowFloatInputMajorTensorData.fromRowMajor( + Shape(rows, cols), rowMajorBytes(Fp16Codec), Fp16Codec, + ) + for (r in 0 until rows) { + for (c in 0 until cols) { + assertEquals( + values[r * cols + c], data.get(r, c), + "get($r, $c) must see the logical element, not the relaid byte order", + ) + } + } + } + + @Test + fun `copyToFloatArray decodes in logical row-major order`() { + val data = NarrowFloatInputMajorTensorData.fromRowMajor( + Shape(rows, cols), rowMajorBytes(Bf16Codec), Bf16Codec, + ) + assertContentEquals( + values, data.copyToFloatArray(), + "consumers that do not care about storage must see the ordinary row-major sequence", + ) + } + + @Test + fun `transposedView shares the buffer and reads as the transpose`() { + val data = NarrowFloatInputMajorTensorData.fromRowMajor( + Shape(rows, cols), rowMajorBytes(Fp16Codec), Fp16Codec, + ) + val view = data.transposedView() + + assertSame( + data.packedData, view.packedData, + "the whole point is that no copy happens — a copy per forward is the bug being fixed", + ) + assertEquals(Shape(cols, rows), view.shape) + assertEquals(Fp16Codec, view.codec, "codec must survive; FP16 and BF16 are both 2 bytes") + + for (r in 0 until rows) { + for (c in 0 until cols) { + assertEquals( + data.get(r, c), view.get(c, r), + "view[$c, $r] must be the transpose of data[$r, $c]", + ) + } + } + } + + @Test + fun `relayout round-trips back to the original bytes`() { + // Relaying twice (with the shape flipped in between) is the identity — the cheapest + // statement of "this is a transpose and not some other permutation". + val original = rowMajorBytes(Fp16Codec) + val once = NarrowFloatInputMajorTensorData.fromRowMajor(Shape(rows, cols), original, Fp16Codec) + val twice = NarrowFloatInputMajorTensorData.fromRowMajor( + Shape(cols, rows), once.packedData, Fp16Codec, + ) + assertContentEquals(original, twice.packedData) + } + + @Test + fun `a square weight is still genuinely transposed`() { + // Square shapes are where an off-by-one in the relayout hides: the byte count matches + // either way, so only the values reveal a wrong permutation. + val square = floatArrayOf(1.0f, 2.0f, 4.0f, 8.0f) + val bytes = ByteArray(square.size * 2) + for (i in square.indices) { + val bits = Fp16Codec.encode(square[i]) + bytes[i * 2] = (bits and 0xFF).toByte() + bytes[i * 2 + 1] = ((bits ushr 8) and 0xFF).toByte() + } + val view = NarrowFloatInputMajorTensorData + .fromRowMajor(Shape(2, 2), bytes, Fp16Codec) + .transposedView() + + assertContentEquals( + floatArrayOf(1.0f, 4.0f, 2.0f, 8.0f), view.copyToFloatArray(), + "transposedView of [[1,2],[4,8]] must be [[1,4],[2,8]]", + ) + } + + @Test + fun `set writes through to the logical element`() { + val data = NarrowFloatInputMajorTensorData.fromRowMajor( + Shape(rows, cols), rowMajorBytes(Fp16Codec), Fp16Codec, + ) + data.set(1, 2, value = 16.0f) + assertEquals(16.0f, data.get(1, 2)) + assertEquals(16.0f, data.transposedView().get(2, 1), "the shared buffer must see it too") + assertEquals(values[0], data.get(0, 0), "neighbouring elements must be untouched") + } + + @Test + fun `rank other than two is rejected`() { + // Norms are rank-1 and embeddings are gathered, not matmul'd; neither should ever be + // relaid. Failing loudly beats silently mis-indexing them. + assertFailsWith { + NarrowFloatInputMajorTensorData.fromRowMajor(Shape(4), ByteArray(8), Fp16Codec) + } + assertFailsWith { + NarrowFloatInputMajorTensorData(Shape(2, 2, 2), ByteArray(16), Fp16Codec) + } + } + + @Test + fun `a short buffer is rejected`() { + assertFailsWith { + NarrowFloatInputMajorTensorData.fromRowMajor(Shape(4, 4), ByteArray(8), Fp16Codec) + } + } + + @Test + fun `the two codecs disagree on identical bytes`() { + // Vacuity guard for every codec assertion above: both formats are 2 bytes per element, + // so a codec mix-up cannot be caught by shape or size checks alone. + val bytes = rowMajorBytes(Fp16Codec) + val asFp16 = NarrowFloatInputMajorTensorData.fromRowMajor(Shape(rows, cols), bytes, Fp16Codec) + val asBf16 = NarrowFloatInputMajorTensorData.fromRowMajor(Shape(rows, cols), bytes, Bf16Codec) + assertTrue( + !asFp16.copyToFloatArray().contentEquals(asBf16.copyToFloatArray()), + "if these agreed, the codec-preservation assertions would prove nothing", + ) + } +} From 018f7a3a7bffb5b94a7a9fb9585664819e8b5b9f Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Mon, 27 Jul 2026 15:08:41 +0200 Subject: [PATCH 2/2] docs(io-safetensors): record the KEEP_NATIVE layout caveat MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Row-major bytes match the file, but projections are [out, in] while the narrow matmul dispatch needs [in, out]. A consumer going through Linear sees that transpose fall to the generic elementwise path and widen the tensor anyway, costing far more than KEEP_NATIVE saves. Point such consumers at NarrowFloatInputMajorTensorData.fromRowMajor, and note that gathered tensors — embedding tables above all — should stay row-major. Refs #888. --- .../sk/ainet/io/safetensors/NarrowFloatLoadPolicy.kt | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/NarrowFloatLoadPolicy.kt b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/NarrowFloatLoadPolicy.kt index 55d9febc..62395e4d 100644 --- a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/NarrowFloatLoadPolicy.kt +++ b/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, }