diff --git a/docs/modules/skeep/pages/008-precision-and-encoding-conversion.adoc b/docs/modules/skeep/pages/008-precision-and-encoding-conversion.adoc index 27c1ec50..c0fc2396 100644 --- a/docs/modules/skeep/pages/008-precision-and-encoding-conversion.adoc +++ b/docs/modules/skeep/pages/008-precision-and-encoding-conversion.adoc @@ -241,6 +241,10 @@ has today and gains the cases that are pure re-encodings. Anything that rounds g source follows `copyMaterialize`: `FileBacked` and `Aliased` buffers are read through the resolver, and a device-resident buffer throws as it does today. * `logicalType`, shape and placement device are kept; `encoding`, buffer and byte size change. + One exception: `repackTo(Dense(4))` from narrow dense (`Dense(2)` FP16/BF16) decodes to FP32 + and the result dtype is FP32 — every 16-bit value is exactly representable, so the values are + bit-exact, and a `Dense(4)` buffer tagged FP16/BF16 has no defined byte layout anywhere in + the codebase. * `requantizeTo` validates up front, with the same message a loader gives: `Requantizers.supports(target)` must be true, and the shape must block (for example the last dimension a multiple of 32 for Q4_0). 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 89f1aaba..68077429 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 @@ -5774,6 +5774,7 @@ public final class sk/ainet/lang/tensor/data/Requantizers { public final fun getSupported ()Ljava/util/List; public final fun quantize (Lsk/ainet/lang/tensor/storage/TensorEncoding;[FLsk/ainet/lang/tensor/Shape;)Lsk/ainet/lang/tensor/data/TensorData; public final fun supports (Lsk/ainet/lang/tensor/storage/TensorEncoding;)Z + public final fun validateShape (Lsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/tensor/Shape;)V } public abstract interface class sk/ainet/lang/tensor/data/RowDequantSource { @@ -8118,6 +8119,9 @@ public final class sk/ainet/lang/tensor/storage/TensorStorage { public final fun isMutable ()Z public final fun memoryReport ()Lsk/ainet/lang/tensor/storage/StorageMemoryReport; public final fun repackTo (Lsk/ainet/lang/tensor/storage/TensorEncoding;)Lsk/ainet/lang/tensor/storage/TensorStorage; + public final fun repackTo (Lsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/tensor/storage/BufferResolver;)Lsk/ainet/lang/tensor/storage/TensorStorage; + public final fun requantizeTo (Lsk/ainet/lang/tensor/storage/TensorEncoding;)Lsk/ainet/lang/tensor/storage/TensorStorage; + public final fun requantizeTo (Lsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/tensor/storage/BufferResolver;)Lsk/ainet/lang/tensor/storage/TensorStorage; public fun toString ()Ljava/lang/String; } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Requantizers.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Requantizers.kt index fc9fc3ff..a77e9ca0 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Requantizers.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Requantizers.kt @@ -27,25 +27,43 @@ public object Requantizers { public val supported: List = listOf(TensorEncoding.Q8_0, TensorEncoding.Q4_0, TensorEncoding.BITNET_PLANES) /** - * Quantize dense [values] with logical [shape] into packed tensor data of [encoding]. + * Validate that [shape] can be blocked for [encoding], without touching any data — the same + * rules [quantize] enforces, callable up front (e.g. before reading a buffer). * * @throws IllegalArgumentException for an encoding without a quantizer, or a shape the format * cannot block (last dimension not a multiple of the block size; planes need rank 2). */ - public fun quantize(encoding: TensorEncoding, values: FloatArray, shape: Shape): TensorData<*, *> = when (encoding) { - TensorEncoding.Q8_0 -> Q8_0Quantizer.quantize(values, shape) - TensorEncoding.Q4_0 -> { - require(shape.rank == 0 || shape[shape.rank - 1] % 32 == 0) { + public fun validateShape(encoding: TensorEncoding, shape: Shape) { + when (encoding) { + TensorEncoding.Q8_0 -> Unit + TensorEncoding.Q4_0 -> require(shape.rank == 0 || shape[shape.rank - 1] % 32 == 0) { "Q4_0 blocks run along the last dimension, which must be a multiple of 32; got $shape" } - Q4_0Quantizer.quantize(values, shape) + TensorEncoding.BITNET_PLANES -> { + require(shape.rank == 2) { "BITNET_PLANES needs a rank-2 weight, got $shape" } + require(shape[1] % 4 == 0) { "BITNET_PLANES needs cols % 4 == 0; got ${shape[1]}" } + } + else -> throw IllegalArgumentException( + "no quantizer for ${encoding.name}; quantize-at-load supports ${supported.joinToString { it.name }}", + ) } - TensorEncoding.BITNET_PLANES -> { - require(shape.rank == 2) { "BITNET_PLANES needs a rank-2 weight, got $shape" } - BitNetPlanesTensorData.fromFloats(shape, values) + } + + /** + * Quantize dense [values] with logical [shape] into packed tensor data of [encoding]. + * + * @throws IllegalArgumentException for an encoding without a quantizer, or a shape the format + * cannot block (last dimension not a multiple of the block size; planes need rank 2). + */ + public fun quantize(encoding: TensorEncoding, values: FloatArray, shape: Shape): TensorData<*, *> { + validateShape(encoding, shape) + return when (encoding) { + TensorEncoding.Q8_0 -> Q8_0Quantizer.quantize(values, shape) + TensorEncoding.Q4_0 -> Q4_0Quantizer.quantize(values, shape) + TensorEncoding.BITNET_PLANES -> BitNetPlanesTensorData.fromFloats(shape, values) + else -> throw IllegalArgumentException( + "no quantizer for ${encoding.name}; quantize-at-load supports ${supported.joinToString { it.name }}", + ) } - else -> throw IllegalArgumentException( - "no quantizer for ${encoding.name}; quantize-at-load supports ${supported.joinToString { it.name }}", - ) } } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorStorage.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorStorage.kt index eea2e92a..e6b57530 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorStorage.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorStorage.kt @@ -162,18 +162,79 @@ public data class TensorStorage( } /** - * Re-encode this storage into a different physical encoding. - * Currently a stub — actual transcoding requires backend kernels. + * Re-encode this storage *without changing the represented values* (SKEEP-008, #1433). * - * @throws UnsupportedOperationException always (until backends implement this) + * Supported: the same encoding (returns `this`), and decoding to `Dense(4)` — from any + * packed encoding [TensorStorageFactory.toTensorData] can read, and from narrow dense + * (`Dense(2)` FP16/BF16). Decoding is bit-exact: every FP16/BF16 value and every + * dequantized block value is exactly representable in FP32, so the result holds the same + * values the source's `dequantizeBlock`/codec gives, as owned dense FP32 host storage. + * + * Anything that would round or quantize throws and names [requantizeTo] — precision loss + * is always a visible, named call. + * + * [BufferHandle.FileBacked] sources need the overload that takes a [BufferResolver]. */ - public fun repackTo(targetEncoding: TensorEncoding): TensorStorage { + public fun repackTo(targetEncoding: TensorEncoding): TensorStorage = + repackTo(targetEncoding, DefaultBufferResolver()) + + /** [repackTo], reading through [resolver] for handles that need platform support. */ + public fun repackTo(targetEncoding: TensorEncoding, resolver: BufferResolver): TensorStorage { if (encoding == targetEncoding) return this - throw UnsupportedOperationException( - "Repacking from ${encoding.name} to ${targetEncoding.name} is not yet implemented" + if (targetEncoding != TensorEncoding.Dense(bytesPerElement = 4)) { + throw UnsupportedOperationException( + "repackTo is lossless-only: ${encoding.name} -> ${targetEncoding.name} would change " + + "the values. Use requantizeTo for lossy re-encoding (supported targets: " + + "${sk.ainet.lang.tensor.data.Requantizers.supported.joinToString { it.name }})." + ) + } + return denseFp32Copy(resolver) + } + + /** + * Re-encode this storage, *accepting precision loss* (SKEEP-008, #1433). Supported targets + * are exactly [sk.ainet.lang.tensor.data.Requantizers.supported] — the same set a loader's + * `EncodingRequest.RequantizeTo` accepts. A packed or narrow source is decoded to FP32 + * first, then quantized. The target encoding and shape are validated before the buffer is + * read, with the loader's messages. + * + * The result is owned packed host storage; [logicalType], shape and placement device are + * kept. [BufferHandle.FileBacked] sources need the overload that takes a [BufferResolver]. + */ + public fun requantizeTo(targetEncoding: TensorEncoding): TensorStorage = + requantizeTo(targetEncoding, DefaultBufferResolver()) + + /** [requantizeTo], reading through [resolver] for handles that need platform support. */ + public fun requantizeTo(targetEncoding: TensorEncoding, resolver: BufferResolver): TensorStorage { + if (encoding == targetEncoding) return this + sk.ainet.lang.tensor.data.Requantizers.validateShape(targetEncoding, shape) + val floats = decodeToFloats(resolver) + val quantized = sk.ainet.lang.tensor.data.Requantizers.quantize(targetEncoding, floats, shape) + val packed = quantized as? PackedBlockStorage ?: throw IllegalStateException( + "quantizer for ${targetEncoding.name} did not produce packed block data" + ) + return TensorStorage( + shape = shape, + logicalType = logicalType, + encoding = targetEncoding, + buffer = BufferHandleFactory.owned(packed.packedData), + placement = placement.copy(domain = MemoryDomain.HOST_HEAP), ) } + /** Decode this storage's values into an owned dense FP32 host copy (shape and device kept). */ + private fun denseFp32Copy(resolver: BufferResolver): TensorStorage = TensorStorage( + shape = shape, + dtype = sk.ainet.lang.types.FP32, + encoding = TensorEncoding.Dense(bytesPerElement = 4), + buffer = BufferHandleFactory.owned(decodeToFloats(resolver)), + placement = placement.copy(domain = MemoryDomain.HOST_HEAP), + ) + + /** This storage's values as FP32, decoded through [TensorStorageFactory.toTensorData]. */ + private fun decodeToFloats(resolver: BufferResolver): FloatArray = + TensorStorageFactory.toTensorData(copyToHost(resolver)).copyToFloatArray() + private fun sizeBytes(): Int = physicalBytes.toInt() override fun equals(other: Any?): Boolean { diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorStorageFactory.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorStorageFactory.kt index 337184ea..98976f47 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorStorageFactory.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorStorageFactory.kt @@ -7,8 +7,15 @@ import sk.ainet.lang.tensor.data.DenseFloatArrayTensorData import sk.ainet.lang.tensor.data.DenseIntArrayTensorData import sk.ainet.lang.tensor.data.FloatArrayTensorData import sk.ainet.lang.tensor.data.IntArrayTensorData +import sk.ainet.lang.tensor.data.BitNetB158TensorData +import sk.ainet.lang.tensor.data.BitNetPlanesTensorData +import sk.ainet.lang.tensor.data.Q4_0BlockTensorData import sk.ainet.lang.tensor.data.Q4_KBlockTensorData import sk.ainet.lang.tensor.data.Q4_KTensorData +import sk.ainet.lang.tensor.data.Q5_0BlockTensorData +import sk.ainet.lang.tensor.data.Q5_1BlockTensorData +import sk.ainet.lang.tensor.data.Q5_KBlockTensorData +import sk.ainet.lang.tensor.data.Q6_KBlockTensorData import sk.ainet.lang.tensor.data.Q8_0BlockTensorData import sk.ainet.lang.tensor.data.Q8_0TensorData import sk.ainet.lang.tensor.data.TensorData @@ -255,6 +262,27 @@ public object TensorStorageFactory { is TensorEncoding.Q8_0 -> { Q8_0BlockTensorData.fromRawBytes(storage.shape, bytes) as TensorData } + is TensorEncoding.Q4_0 -> { + Q4_0BlockTensorData.fromRawBytes(storage.shape, bytes) as TensorData + } + is TensorEncoding.Q5_0 -> { + Q5_0BlockTensorData.fromRawBytes(storage.shape, bytes) as TensorData + } + is TensorEncoding.Q5_1 -> { + Q5_1BlockTensorData.fromRawBytes(storage.shape, bytes) as TensorData + } + is TensorEncoding.Q5_K -> { + Q5_KBlockTensorData.fromRawBytes(storage.shape, bytes) as TensorData + } + is TensorEncoding.Q6_K -> { + Q6_KBlockTensorData.fromRawBytes(storage.shape, bytes) as TensorData + } + is TensorEncoding.BITNET_PLANES -> { + BitNetPlanesTensorData.fromRawBytes(storage.shape, bytes) as TensorData + } + is TensorEncoding.BITNET_B1_58 -> { + BitNetB158TensorData.fromRawBytes(storage.shape, bytes) as TensorData + } else -> throw UnsupportedOperationException( "toTensorData not supported for encoding ${storage.encoding.name}" ) diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/RepackRequantizeTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/RepackRequantizeTest.kt new file mode 100644 index 00000000..85e81fb4 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/RepackRequantizeTest.kt @@ -0,0 +1,195 @@ +package sk.ainet.lang.tensor.storage + +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.data.Requantizers +import sk.ainet.lang.types.FP16 +import sk.ainet.lang.types.FP32 +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 + +/** + * SKEEP-008 Part B (#1433): `repackTo` re-encodes losslessly (packed/narrow -> Dense(4)), + * `requantizeTo` is the visible lossy call, sharing [Requantizers] with the loaders. + */ +class RepackRequantizeTest { + + private val shape = Shape(2, 32) + private val values = FloatArray(64) { i -> (i - 32) * 0.173f } + + private fun denseStorage(): TensorStorage = TensorStorageFactory.fromFloatArray(shape, values) + + /** Packed storage holding [values] quantized to [encoding], as a loader would produce. */ + private fun packedStorage(encoding: TensorEncoding): TensorStorage { + val packed = Requantizers.quantize(encoding, values, shape) as PackedBlockStorage + return TensorStorage( + shape = shape, + dtype = FP32, + encoding = encoding, + buffer = BufferHandleFactory.owned(packed.packedData), + ) + } + + private fun floatsOf(storage: TensorStorage): FloatArray = + TensorStorageFactory.toTensorData(storage).copyToFloatArray() + + // --- repackTo: lossless --- + + @Test + fun repackTo_sameEncoding_returnsSameInstance() { + val storage = denseStorage() + assertSame(storage, storage.repackTo(TensorEncoding.Dense(4))) + } + + @Test + fun repackTo_q8_0_toDense_reproducesDequantizedValues_bitExact() { + val source = packedStorage(TensorEncoding.Q8_0) + val dequantized = floatsOf(source) // what dequantizeBlock gives + + val dense = source.repackTo(TensorEncoding.Dense(4)) + + assertEquals(TensorEncoding.Dense(4), dense.encoding) + assertEquals(FP32, dense.dtype) + assertEquals(Ownership.OWNED, dense.ownership) + assertEquals(shape, dense.shape) + assertContentEquals(dequantized, floatsOf(dense)) + } + + @Test + fun repackTo_q4_0_toDense_reproducesDequantizedValues_bitExact() { + val source = packedStorage(TensorEncoding.Q4_0) + val dequantized = floatsOf(source) + + val dense = source.repackTo(TensorEncoding.Dense(4)) + + assertEquals(TensorEncoding.Dense(4), dense.encoding) + assertContentEquals(dequantized, floatsOf(dense)) + } + + @Test + fun repackTo_bitnetPlanes_toDense_reproducesDequantizedValues_bitExact() { + val source = packedStorage(TensorEncoding.BITNET_PLANES) + val dequantized = floatsOf(source) + + val dense = source.repackTo(TensorEncoding.Dense(4)) + + assertEquals(TensorEncoding.Dense(4), dense.encoding) + assertContentEquals(dequantized, floatsOf(dense)) + } + + @Test + fun repackTo_narrowDense_widensExactly() { + // Dense(2) FP16 bytes -> Dense(4) FP32; every FP16 value is exactly representable. + val fp16Values = FloatArray(values.size) { Fp16Codec.decode(Fp16Codec.encode(values[it])) } + val bytes = ByteArray(values.size * 2) + for (i in values.indices) { + val bits = Fp16Codec.encode(values[i]) + bytes[i * 2] = (bits and 0xFF).toByte() + bytes[i * 2 + 1] = ((bits shr 8) and 0xFF).toByte() + } + val narrow = TensorStorageFactory.fromRawBytesOwned(shape, FP16, TensorEncoding.Dense(2), bytes) + + val wide = narrow.repackTo(TensorEncoding.Dense(4)) + + assertEquals(TensorEncoding.Dense(4), wide.encoding) + assertEquals(FP32, wide.dtype) + assertContentEquals(fp16Values, floatsOf(wide)) + } + + @Test + fun repackTo_lossyDirection_throwsAndNamesRequantizeTo() { + val failure = assertFailsWith { + denseStorage().repackTo(TensorEncoding.Q8_0) + } + assertTrue( + failure.message.orEmpty().contains("requantizeTo"), + "the error must point at requantizeTo, got: ${failure.message}" + ) + } + + // --- requantizeTo: lossy, loader parity --- + + @Test + fun requantizeTo_q8_0_matchesLoaderQuantization_byteExact() { + val expected = (Requantizers.quantize(TensorEncoding.Q8_0, values, shape) as PackedBlockStorage).packedData + + val result = denseStorage().requantizeTo(TensorEncoding.Q8_0) + + assertEquals(TensorEncoding.Q8_0, result.encoding) + assertEquals(FP32, result.dtype) + assertEquals(Ownership.OWNED, result.ownership) + assertEquals(shape, result.shape) + val resultBytes = (result.buffer as BufferHandle.Owned).data + assertContentEquals(expected, resultBytes) + } + + @Test + fun requantizeTo_q4_0_roundTripsThroughRepack() { + val result = denseStorage().requantizeTo(TensorEncoding.Q4_0) + val expected = (Requantizers.quantize(TensorEncoding.Q4_0, values, shape) as PackedBlockStorage).packedData + + assertContentEquals(expected, (result.buffer as BufferHandle.Owned).data) + // And the lossless way back reproduces the quantizer's values. + assertContentEquals(floatsOf(result), floatsOf(result.repackTo(TensorEncoding.Dense(4)))) + } + + @Test + fun requantizeTo_bitnetPlanes_matchesLoaderQuantization() { + val expected = (Requantizers.quantize(TensorEncoding.BITNET_PLANES, values, shape) as PackedBlockStorage).packedData + + val result = denseStorage().requantizeTo(TensorEncoding.BITNET_PLANES) + + assertEquals(TensorEncoding.BITNET_PLANES, result.encoding) + assertContentEquals(expected, (result.buffer as BufferHandle.Owned).data) + } + + @Test + fun requantizeTo_fromPackedSource_dequantizesFirst() { + val q8 = packedStorage(TensorEncoding.Q8_0) + val expected = (Requantizers.quantize(TensorEncoding.Q4_0, floatsOf(q8), shape) as PackedBlockStorage).packedData + + val result = q8.requantizeTo(TensorEncoding.Q4_0) + + assertEquals(TensorEncoding.Q4_0, result.encoding) + assertContentEquals(expected, (result.buffer as BufferHandle.Owned).data) + } + + @Test + fun requantizeTo_unsupportedEncoding_failsWithLoaderMessage() { + val failure = assertFailsWith { + denseStorage().requantizeTo(TensorEncoding.Q4_K) + } + assertTrue( + failure.message.orEmpty().contains("no quantizer for Q4_K"), + "expected the loader's message, got: ${failure.message}" + ) + } + + @Test + fun requantizeTo_shapeThatDoesNotBlock_failsBeforeReadingTheBuffer() { + // A device-resident buffer cannot be read on the host, so this only passes if the + // shape validation runs first. + val unreadable = TensorStorage( + shape = Shape(2, 31), + dtype = FP32, + encoding = TensorEncoding.Dense(4), + buffer = BufferHandle.DeviceResident( + deviceId = "gpu0", + backendHandle = Unit, + sizeInBytes = 2L * 31 * 4, + isMutable = false, + ), + ) + val failure = assertFailsWith { + unreadable.requantizeTo(TensorEncoding.Q4_0) + } + assertTrue( + failure.message.orEmpty().contains("multiple of 32"), + "expected the Q4_0 blocking message, got: ${failure.message}" + ) + } +}