diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/Bf16MatmulKernel.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/Bf16MatmulKernel.kt index b17b75a4..79e7f98a 100644 --- a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/Bf16MatmulKernel.kt +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/Bf16MatmulKernel.kt @@ -1,5 +1,8 @@ package sk.ainet.backend.api.kernel +import sk.ainet.lang.types.Bf16Codec +import sk.ainet.lang.types.NarrowFloatCodec + /** * FP32 input × BF16-packed weight matrix multiplication: * @@ -43,7 +46,10 @@ package sk.ainet.backend.api.kernel * and byte-stride accessors on the same `ByteArray`. For a contiguous * `(k, n)` BF16 matrix `bByteStride == n * 2`. */ -public interface Bf16MatmulKernel { +public interface Bf16MatmulKernel : NarrowFloatMatmulKernel { + + override val codec: NarrowFloatCodec get() = Bf16Codec + /** * @param a left operand `(m, k)`, row-major FP32, stride [aStride] * floats per row. @@ -64,7 +70,7 @@ public interface Bf16MatmulKernel { * @param k contraction dimension (cols of A == rows of B). `k == 0` * zeros the `m × n` output block. */ - public fun matmul( + override fun matmul( a: FloatArray, aOffset: Int, aStride: Int, b: ByteArray, bByteOffset: Int, bByteStride: Int, out: FloatArray, outOffset: Int, outStride: Int, diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelProvider.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelProvider.kt index ad475190..b3ebf43a 100644 --- a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelProvider.kt +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelProvider.kt @@ -61,6 +61,12 @@ public interface KernelProvider { */ public fun matmulBf16(): Bf16MatmulKernel? = null + /** + * IEEE binary16 weight matmul, or `null` when this provider carries no FP16 kernel. + * Callers cascade to a lower-priority provider, ultimately the scalar reference. + */ + public fun matmulFp16(): Fp16MatmulKernel? = null + /** * F32 × Q8_0 matmul kernel exposed by this provider, or `null` if * this provider does not specialize Q8_0. Same fall-through pattern. diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/NarrowFloatMatmulKernel.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/NarrowFloatMatmulKernel.kt new file mode 100644 index 00000000..5c81483c --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/NarrowFloatMatmulKernel.kt @@ -0,0 +1,61 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.types.Fp16Codec +import sk.ainet.lang.types.NarrowFloatCodec + +/** + * Matmul over a **16-bit float** right operand: `C(m,n) = A(m,k) · B(k,n)` where `B` is packed at + * 2 bytes per element and `A`/`C` are FP32. + * + * The shape, offset, aliasing and "fully overwrite the `m × n` block" obligations are identical to + * `Fp32MatmulKernel`. Only the inner-loop decode of `B` differs, which is what [codec] selects — so + * BF16 and FP16 differ by a decode step and nothing else. + * + * **Accumulation is always FP32.** The narrow format is a storage width, never an accumulate width; + * every implementation widens each `B` element before the multiply and sums in `Float`. This + * matches what PyTorch, JAX and tensor cores do, and is not a shortcut to be "optimized" away. + * + * ## Stride convention + * + * `A` and `out` use **float** strides (counting FP32 elements). `B` uses **byte** offsets and + * strides, matching the byte-packed convention of the block-quantized kernels and avoiding the + * foot-gun of mixing element- and byte-strides on the same `ByteArray`. For a contiguous `(k, n)` + * matrix, `bByteStride == n * 2`. + */ +public interface NarrowFloatMatmulKernel { + + /** The 16-bit format this kernel decodes `B` from. */ + public val codec: NarrowFloatCodec + + /** + * @param a left operand `(m, k)`, row-major FP32, [aStride] floats per row. + * @param aOffset element offset into [a] of the (0, 0) entry. + * @param aStride distance in **floats** between consecutive rows of [a]. + * @param b right operand `(k, n)`, row-major, packed little-endian 16-bit floats. + * @param bByteOffset **byte** offset into [b] of the (0, 0) entry. + * @param bByteStride **byte** distance between consecutive rows of [b]. + * @param out output `(m, n)`, row-major FP32, [outStride] floats per row. + * @param outOffset element offset into [out]. + * @param outStride distance in **floats** between consecutive rows of [out]. + * @param m rows of A and C. + * @param n columns of B and C. + * @param k contraction dimension. `k == 0` zeros the `m × n` output block. + */ + public fun matmul( + a: FloatArray, aOffset: Int, aStride: Int, + b: ByteArray, bByteOffset: Int, bByteStride: Int, + out: FloatArray, outOffset: Int, outStride: Int, + m: Int, n: Int, k: Int, + ) +} + +/** + * Matmul over an **IEEE binary16** right operand. + * + * The FP16 counterpart to [Bf16MatmulKernel]. Note that FP16 cannot use the bit-shift decode that + * makes BF16 nearly free — binary16 needs exponent rebiasing — so vectorized FP16 implementations + * lag their BF16 equivalents. Correctness first; throughput is a separate concern. + */ +public interface Fp16MatmulKernel : NarrowFloatMatmulKernel { + override val codec: NarrowFloatCodec get() = Fp16Codec +} diff --git a/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api b/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api index 0ed4e5fb..2ce08b89 100644 --- a/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api +++ b/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api @@ -43,6 +43,13 @@ public final class sk/ainet/context/DirectCpuExecutionContext$Companion { public final class sk/ainet/exec/kernel/PanamaVectorBf16MatmulKernel : sk/ainet/backend/api/kernel/Bf16MatmulKernel { public static final field INSTANCE Lsk/ainet/exec/kernel/PanamaVectorBf16MatmulKernel; + public fun getCodec ()Lsk/ainet/lang/types/NarrowFloatCodec; + public fun matmul ([FII[BII[FIIIII)V +} + +public final class sk/ainet/exec/kernel/PanamaVectorFp16MatmulKernel : sk/ainet/backend/api/kernel/Fp16MatmulKernel { + public static final field INSTANCE Lsk/ainet/exec/kernel/PanamaVectorFp16MatmulKernel; + public fun getCodec ()Lsk/ainet/lang/types/NarrowFloatCodec; public fun matmul ([FII[BII[FIIIII)V } @@ -52,6 +59,7 @@ public final class sk/ainet/exec/kernel/PanamaVectorKernelProvider : sk/ainet/ba public fun getPriority ()I public fun isAvailable ()Z public fun matmulBf16 ()Lsk/ainet/backend/api/kernel/Bf16MatmulKernel; + public fun matmulFp16 ()Lsk/ainet/backend/api/kernel/Fp16MatmulKernel; public fun matmulFp32 ()Lsk/ainet/backend/api/kernel/Fp32MatmulKernel; public fun matmulQ4K ()Lsk/ainet/backend/api/kernel/Q4KMatmulKernel; public fun matmulQ4_0 ()Lsk/ainet/backend/api/kernel/Q4_0MatmulKernel; @@ -69,6 +77,7 @@ public final class sk/ainet/exec/kernel/PanamaVectorKernelProviderFactory : sk/a public fun getPriority ()I public fun isAvailable ()Z public fun matmulBf16 ()Lsk/ainet/backend/api/kernel/Bf16MatmulKernel; + public fun matmulFp16 ()Lsk/ainet/backend/api/kernel/Fp16MatmulKernel; public fun matmulFp32 ()Lsk/ainet/backend/api/kernel/Fp32MatmulKernel; public fun matmulQ4K ()Lsk/ainet/backend/api/kernel/Q4KMatmulKernel; public fun matmulQ4_0 ()Lsk/ainet/backend/api/kernel/Q4_0MatmulKernel; @@ -122,6 +131,13 @@ public final class sk/ainet/exec/kernel/PanamaVectorQ8_0MatmulKernel : sk/ainet/ public final class sk/ainet/exec/kernel/ScalarBf16MatmulKernel : sk/ainet/backend/api/kernel/Bf16MatmulKernel { public static final field INSTANCE Lsk/ainet/exec/kernel/ScalarBf16MatmulKernel; + public fun getCodec ()Lsk/ainet/lang/types/NarrowFloatCodec; + public fun matmul ([FII[BII[FIIIII)V +} + +public final class sk/ainet/exec/kernel/ScalarFp16MatmulKernel : sk/ainet/backend/api/kernel/Fp16MatmulKernel { + public static final field INSTANCE Lsk/ainet/exec/kernel/ScalarFp16MatmulKernel; + public fun getCodec ()Lsk/ainet/lang/types/NarrowFloatCodec; public fun matmul ([FII[BII[FIIIII)V } @@ -131,6 +147,7 @@ public final class sk/ainet/exec/kernel/ScalarKernelProvider : sk/ainet/backend/ public fun getPriority ()I public fun isAvailable ()Z public fun matmulBf16 ()Lsk/ainet/backend/api/kernel/Bf16MatmulKernel; + public fun matmulFp16 ()Lsk/ainet/backend/api/kernel/Fp16MatmulKernel; public fun matmulFp32 ()Lsk/ainet/backend/api/kernel/Fp32MatmulKernel; public fun matmulQ4K ()Lsk/ainet/backend/api/kernel/Q4KMatmulKernel; public fun matmulQ4_0 ()Lsk/ainet/backend/api/kernel/Q4_0MatmulKernel; @@ -148,6 +165,7 @@ public final class sk/ainet/exec/kernel/ScalarKernelProviderFactory : sk/ainet/b public fun getPriority ()I public fun isAvailable ()Z public fun matmulBf16 ()Lsk/ainet/backend/api/kernel/Bf16MatmulKernel; + public fun matmulFp16 ()Lsk/ainet/backend/api/kernel/Fp16MatmulKernel; public fun matmulFp32 ()Lsk/ainet/backend/api/kernel/Fp32MatmulKernel; public fun matmulQ4K ()Lsk/ainet/backend/api/kernel/Q4KMatmulKernel; public fun matmulQ4_0 ()Lsk/ainet/backend/api/kernel/Q4_0MatmulKernel; diff --git a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/kernel/ScalarFp16MatmulKernel.kt b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/kernel/ScalarFp16MatmulKernel.kt new file mode 100644 index 00000000..266da1c7 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/kernel/ScalarFp16MatmulKernel.kt @@ -0,0 +1,49 @@ +package sk.ainet.exec.kernel + +import sk.ainet.backend.api.kernel.Fp16MatmulKernel +import sk.ainet.lang.types.Fp16Codec + +/** + * Scalar reference implementation of [Fp16MatmulKernel] — three nested loops, decode inline, no + * SIMD. Always available on every KMP target, and the correctness reference an accelerated FP16 + * kernel must match within FP-ordering tolerance. + * + * Mirrors [ScalarBf16MatmulKernel]; the only difference is the decode, which for binary16 needs + * exponent rebiasing rather than BF16's bit shift. + */ +public object ScalarFp16MatmulKernel : Fp16MatmulKernel { + override fun matmul( + a: FloatArray, aOffset: Int, aStride: Int, + b: ByteArray, bByteOffset: Int, bByteStride: Int, + out: FloatArray, outOffset: Int, outStride: Int, + m: Int, n: Int, k: Int, + ) { + require(m >= 0 && n >= 0 && k >= 0) { + "ScalarFp16MatmulKernel: m, n, k must be non-negative; got m=$m n=$n k=$k" + } + if (m == 0 || n == 0) return + // k == 0 is an explicit "zero the output block" per the SPI. + if (k == 0) { + for (i in 0 until m) { + val rowOff = outOffset + i * outStride + for (j in 0 until n) out[rowOff + j] = 0f + } + return + } + for (i in 0 until m) { + val aRowOff = aOffset + i * aStride + val outRowOff = outOffset + i * outStride + for (j in 0 until n) { + // Accumulate in FP32 — the narrow format is storage only. + var sum = 0f + for (p in 0 until k) { + val bByteIdx = bByteOffset + p * bByteStride + j * 2 + val lo = b[bByteIdx].toInt() and 0xFF + val hi = b[bByteIdx + 1].toInt() and 0xFF + sum += a[aRowOff + p] * Fp16Codec.decode((hi shl 8) or lo) + } + out[outRowOff + j] = sum + } + } + } +} diff --git a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/kernel/ScalarKernelProvider.kt b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/kernel/ScalarKernelProvider.kt index 31a8386b..c17bbd57 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/kernel/ScalarKernelProvider.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/kernel/ScalarKernelProvider.kt @@ -1,6 +1,7 @@ package sk.ainet.exec.kernel import sk.ainet.backend.api.kernel.Bf16MatmulKernel +import sk.ainet.backend.api.kernel.Fp16MatmulKernel import sk.ainet.backend.api.kernel.Fp32MatmulKernel import sk.ainet.backend.api.kernel.KernelProvider import sk.ainet.backend.api.kernel.Q4KMatmulKernel @@ -30,6 +31,8 @@ public object ScalarKernelProvider : KernelProvider { override fun isAvailable(): Boolean = true override fun matmulFp32(): Fp32MatmulKernel = ScalarMatmulKernel override fun matmulBf16(): Bf16MatmulKernel = ScalarBf16MatmulKernel + + override fun matmulFp16(): Fp16MatmulKernel = ScalarFp16MatmulKernel override fun matmulQ8_0(): Q8_0MatmulKernel = ScalarQ8_0MatmulKernel override fun matmulQ4_0(): Q4_0MatmulKernel = ScalarQ4_0MatmulKernel override fun matmulQ4K(): Q4KMatmulKernel = ScalarQ4_KMatmulKernel 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 0c3b7c0b..f0d3aec8 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 @@ -30,8 +30,11 @@ import sk.ainet.lang.tensor.data.Q8_0BlockTensorData import sk.ainet.lang.tensor.data.TensorData import sk.ainet.lang.tensor.data.TensorDataFactory import sk.ainet.lang.tensor.ops.UpsampleMode +import sk.ainet.lang.types.BF16 +import sk.ainet.lang.types.Bf16Codec import sk.ainet.lang.types.FP16 import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.Fp16Codec import sk.ainet.lang.types.Int32 import sk.ainet.lang.types.Int8 import kotlin.math.floor @@ -2685,18 +2688,34 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory @Suppress("UNCHECKED_CAST") val outData = when (targetClass) { - FP32::class, FP16::class -> dataFactory.fromFloatArray( + FP32::class -> dataFactory.fromFloatArray( tensor.shape, targetClass, copyTensorValuesAsFloatArray(tensor) ) as TensorData + // Narrowing to 16 bits must actually ROUND. This previously shared the FP32 branch, + // so `convert(x, FP16)` only re-tagged the dtype and kept full FP32 precision — the + // resulting tensor claimed to be FP16 while holding values no FP16 can represent, with + // no rounding, no overflow-to-Inf past the format's range, and no NaN handling. + // Storage stays float-backed (matching DenseTensorDataFactory's narrow tags); it is + // the VALUES that now reflect the target format. + FP16::class, BF16::class -> { + val codec = if (targetClass == BF16::class) Bf16Codec else Fp16Codec + val src = copyTensorValuesAsFloatArray(tensor) + val rounded = FloatArray(src.size) { codec.decode(codec.encode(src[it])) } + dataFactory.fromFloatArray( + tensor.shape, + targetClass, + rounded + ) as TensorData + } Int32::class, Int8::class -> dataFactory.fromIntArray( tensor.shape, targetClass, copyTensorValuesAsIntArray(tensor) ) as TensorData else -> throw IllegalArgumentException( - "convert supports FP32, FP16, Int32, and Int8 targets, got ${targetType.name}" + "convert supports FP32, FP16, BF16, Int32, and Int8 targets, got ${targetType.name}" ) } return CpuTensor(outData, this, targetClass, GradState(requiresGrad = tensor.requiresGrad)) diff --git a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorFp16MatmulKernel.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorFp16MatmulKernel.kt new file mode 100644 index 00000000..bd68dd57 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorFp16MatmulKernel.kt @@ -0,0 +1,87 @@ +package sk.ainet.exec.kernel + +import jdk.incubator.vector.FloatVector +import jdk.incubator.vector.VectorSpecies +import sk.ainet.backend.api.kernel.Fp16MatmulKernel +import sk.ainet.lang.types.Fp16Codec + +/** + * Panama Vector API implementation of [Fp16MatmulKernel] — the FP16 counterpart to + * [PanamaVectorBf16MatmulKernel], using the same i-p-j outer-product shape and FP32 FMA + * accumulation. + * + * **Why the decode stays scalar here.** BF16 can be widened lane-wise with an integer shift + * (`IntVector.lanewise(LSHL, 16).reinterpretAsFloats()`), which is why the BF16 kernel vectorizes + * its dequant. Binary16 needs exponent rebiasing and gradual-underflow handling, and mainstream + * JDKs expose no FP16 vector species, so each element is decoded scalar into a lane-width scratch + * buffer before the vectorized multiply-accumulate. The FMA over `n` is still fully vectorized — + * only the widening is not. Expect this kernel to trail the BF16 one; that is inherent to the + * format on this platform, not a defect in this implementation. + * + * Numerical parity vs [ScalarFp16MatmulKernel] is asserted by + * `PanamaVectorFp16MatmulKernelParityTest`. + */ +public object PanamaVectorFp16MatmulKernel : Fp16MatmulKernel { + + private val floatSpecies: VectorSpecies = FloatVector.SPECIES_PREFERRED + + override fun matmul( + a: FloatArray, aOffset: Int, aStride: Int, + b: ByteArray, bByteOffset: Int, bByteStride: Int, + out: FloatArray, outOffset: Int, outStride: Int, + m: Int, n: Int, k: Int, + ) { + require(m >= 0 && n >= 0 && k >= 0) { + "PanamaVectorFp16MatmulKernel: m, n, k must be non-negative; got m=$m n=$n k=$k" + } + if (m == 0 || n == 0) return + if (k == 0) { + for (i in 0 until m) { + val rowOff = outOffset + i * outStride + for (j in 0 until n) out[rowOff + j] = 0f + } + return + } + + val laneCount = floatSpecies.length() + val bound = floatSpecies.loopBound(n) + val scratch = FloatArray(laneCount) + + // Zero the output block first — the i-p-j outer product accumulates into it. + for (i in 0 until m) { + val rowOff = outOffset + i * outStride + for (j in 0 until n) out[rowOff + j] = 0f + } + + for (i in 0 until m) { + val aRowOff = aOffset + i * aStride + val outRowOff = outOffset + i * outStride + for (p in 0 until k) { + val aIp = a[aRowOff + p] + val aBcast = FloatVector.broadcast(floatSpecies, aIp) + val bRowByteOff = bByteOffset + p * bByteStride + var j = 0 + while (j < bound) { + val byteBase = bRowByteOff + j * 2 + for (lane in 0 until laneCount) { + val lo = b[byteBase + lane * 2].toInt() and 0xFF + val hi = b[byteBase + lane * 2 + 1].toInt() and 0xFF + scratch[lane] = Fp16Codec.decode((hi shl 8) or lo) + } + val bVec = FloatVector.fromArray(floatSpecies, scratch, 0) + val outVec = FloatVector.fromArray(floatSpecies, out, outRowOff + j) + aBcast.fma(bVec, outVec).intoArray(out, outRowOff + j) + j += laneCount + } + // Tail (scalar): up to laneCount - 1 trailing columns. + while (j < n) { + val bByteIdx = bRowByteOff + j * 2 + val lo = b[bByteIdx].toInt() and 0xFF + val hi = b[bByteIdx + 1].toInt() and 0xFF + out[outRowOff + j] += aIp * Fp16Codec.decode((hi shl 8) or lo) + j++ + } + } + } + } +} diff --git a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorKernelProvider.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorKernelProvider.kt index e9bb5b59..87218dbd 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorKernelProvider.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorKernelProvider.kt @@ -1,6 +1,7 @@ package sk.ainet.exec.kernel import sk.ainet.backend.api.kernel.Bf16MatmulKernel +import sk.ainet.backend.api.kernel.Fp16MatmulKernel import sk.ainet.backend.api.kernel.Fp32MatmulKernel import sk.ainet.backend.api.kernel.KernelProvider import sk.ainet.backend.api.kernel.Q4KMatmulKernel @@ -51,6 +52,9 @@ public object PanamaVectorKernelProvider : KernelProvider { override fun matmulBf16(): Bf16MatmulKernel? = if (isAvailable()) PanamaVectorBf16MatmulKernel else null + override fun matmulFp16(): Fp16MatmulKernel? = + if (isAvailable()) PanamaVectorFp16MatmulKernel else null + override fun matmulQ8_0(): Q8_0MatmulKernel? = if (isAvailable()) PanamaVectorQ8_0MatmulKernel else null 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 0aa070cc..587ccc5d 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 @@ -4,6 +4,7 @@ import jdk.incubator.vector.FloatVector import jdk.incubator.vector.VectorSpecies import jdk.incubator.vector.VectorOperators import sk.ainet.backend.api.kernel.Bf16MatmulKernel +import sk.ainet.backend.api.kernel.Fp16MatmulKernel import sk.ainet.backend.api.kernel.Fp32MatmulKernel import sk.ainet.backend.api.kernel.KernelRegistry import sk.ainet.backend.api.kernel.KernelServiceLoader @@ -12,6 +13,7 @@ import sk.ainet.backend.api.kernel.Q4KMatmulKernel import sk.ainet.backend.api.kernel.Q4_0MatmulKernel import sk.ainet.backend.api.kernel.Q8_0MatmulKernel import sk.ainet.exec.kernel.ScalarBf16MatmulKernel +import sk.ainet.exec.kernel.ScalarFp16MatmulKernel import sk.ainet.exec.kernel.ScalarMatmulKernel import sk.ainet.exec.kernel.ScalarQ4_0MatmulKernel import sk.ainet.lang.tensor.Shape @@ -23,6 +25,9 @@ import sk.ainet.lang.tensor.data.MemorySegmentTensorData import sk.ainet.lang.tensor.data.Q4MemorySegmentMarker import sk.ainet.lang.tensor.data.Q4MemorySegmentTensorData import sk.ainet.lang.tensor.data.Bf16TensorData +import sk.ainet.lang.tensor.data.NarrowFloatTensorData +import sk.ainet.lang.types.Bf16Codec +import sk.ainet.lang.types.Fp16Codec import sk.ainet.lang.tensor.data.Q4_0TensorData import sk.ainet.lang.tensor.data.Q8_0TensorData import sk.ainet.lang.tensor.data.Q8MemorySegmentMarker @@ -30,6 +35,7 @@ import sk.ainet.lang.tensor.data.Q8MemorySegmentTensorData import sk.ainet.lang.tensor.data.Q4_KBlockTensorData import sk.ainet.lang.tensor.data.Q4_KTensorData import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.types.BF16 import sk.ainet.lang.types.DType import sk.ainet.lang.types.FP16 import sk.ainet.lang.types.FP32 @@ -126,6 +132,17 @@ internal class DefaultCpuOpsJvm( ?: ScalarBf16MatmulKernel } + /** FP16 matmul kernel resolved via [KernelRegistry]; mirrors [bf16MatmulKernel]. */ + private val fp16MatmulKernel: Fp16MatmulKernel by lazy { + if (KernelRegistry.providers().isEmpty()) { + KernelServiceLoader.installAll() + } + KernelRegistry.providers() + .firstOrNull { it.isAvailable() && it.matmulFp16() != null } + ?.matmulFp16() + ?: ScalarFp16MatmulKernel + } + /** * Q4_0 matmul kernel resolved via [KernelRegistry]. Mirrors * [bf16MatmulKernel]: non-null, picks the highest-priority provider @@ -588,20 +605,32 @@ internal class DefaultCpuOpsJvm( @Suppress("UNCHECKED_CAST") CpuTensor(outData as TensorData, this, a.dtype) } - is Bf16TensorData -> { - // BF16 is dense (not block-quantized) and the kernel SPI is a + is NarrowFloatTensorData -> { + // Narrow floats are dense (not block-quantized) and the kernel SPI is a // full SGEMM with `(m, n, k)` strides — no per-batch loop needed, // unlike the matvec-shaped Q4_K / Q8_0 / Q6_K branches. - val outBuffer = FloatArray(batchSize * outputDim) - bf16MatmulKernel.matmul( - inputBuffer, 0, inputDim, - bData.packedData, 0, outputDim * Bf16TensorData.BYTES_PER_ELEMENT, - outBuffer, 0, outputDim, - batchSize, outputDim, inputDim, - ) - val outData = DenseFloatArrayTensorData(Shape(batchSize, outputDim), outBuffer) - @Suppress("UNCHECKED_CAST") - CpuTensor(outData as TensorData, this, a.dtype) + // + // The codec selects the kernel: the two formats have different bit layouts, + // so decoding F16 bytes with the BF16 kernel would yield silent garbage. + val kernel = when (bData.codec) { + Bf16Codec -> bf16MatmulKernel + Fp16Codec -> fp16MatmulKernel + else -> null + } + if (kernel == null) { + null + } else { + val outBuffer = FloatArray(batchSize * outputDim) + kernel.matmul( + inputBuffer, 0, inputDim, + bData.packedData, 0, outputDim * NarrowFloatTensorData.BYTES_PER_ELEMENT, + outBuffer, 0, outputDim, + batchSize, outputDim, inputDim, + ) + val outData = DenseFloatArrayTensorData(Shape(batchSize, outputDim), outBuffer) + @Suppress("UNCHECKED_CAST") + CpuTensor(outData as TensorData, this, a.dtype) + } } // Q6_K / Q5_1 / Q5_0 dispatch is handled in DefaultCpuOpsBase via the kernel // registry (block-major, shared with Native); not intercepted here. @@ -893,7 +922,12 @@ internal class DefaultCpuOpsJvm( private fun supportsFloatOps(tensor: Tensor): Boolean { val dtype = tensor.dtype - return (dtype == FP32::class || dtype == FP16::class) + // BF16 was excluded here, which left BF16-tagged tensors matmul-only — every elementwise + // and unary op fell through to the generic scalar path. All three float tags are backed by + // FloatArray when produced by DenseTensorDataFactory, so they can use the vector kernels. + // Callers re-check with `data as? FloatArrayTensorData`, so genuinely narrow-storage + // tensors (Bf16/Fp16DenseTensorData) still fall through safely rather than being misread. + return dtype == FP32::class || dtype == FP16::class || dtype == BF16::class } private fun chooseMatmul(a: Tensor, b: Tensor): Tensor? { diff --git a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/PanamaVectorFp16MatmulKernelParityTest.kt b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/PanamaVectorFp16MatmulKernelParityTest.kt new file mode 100644 index 00000000..6a8630cd --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/PanamaVectorFp16MatmulKernelParityTest.kt @@ -0,0 +1,123 @@ +package sk.ainet.exec.kernel + +import sk.ainet.lang.types.Fp16Codec +import kotlin.math.abs +import kotlin.random.Random +import kotlin.test.Test +import kotlin.test.assertTrue + +/** + * [PanamaVectorFp16MatmulKernel] must agree with the scalar reference + * [ScalarFp16MatmulKernel] within FMA + reordered-reduction tolerance. + * + * Mirrors `PanamaVectorBf16MatmulKernelParityTest`. Tolerance is tighter here than the BF16 test's + * because binary16 carries three more mantissa bits, so the decode error it starts from is ~8× + * smaller. + */ +class PanamaVectorFp16MatmulKernelParityTest { + + /** Pack FP32 values as little-endian binary16 bytes. */ + private fun fp16Bytes(values: FloatArray): ByteArray { + val out = ByteArray(values.size * 2) + for (i in values.indices) { + val bits = Fp16Codec.encode(values[i]) + out[i * 2] = (bits and 0xFF).toByte() + out[i * 2 + 1] = ((bits ushr 8) and 0xFF).toByte() + } + return out + } + + private fun assertParity(m: Int, n: Int, k: Int, seed: Int) { + val rnd = Random(seed) + val a = FloatArray(m * k) { rnd.nextFloat() * 2f - 1f } + val bFloats = FloatArray(k * n) { rnd.nextFloat() * 2f - 1f } + val b = fp16Bytes(bFloats) + + val outScalar = FloatArray(m * n) + val outPanama = FloatArray(m * n) + + ScalarFp16MatmulKernel.matmul(a, 0, k, b, 0, n * 2, outScalar, 0, n, m, n, k) + PanamaVectorFp16MatmulKernel.matmul(a, 0, k, b, 0, n * 2, outPanama, 0, n, m, n, k) + + val tol = 1e-5f * k + for (i in outScalar.indices) { + assertTrue( + abs(outScalar[i] - outPanama[i]) <= tol, + "mismatch at $i (m=$m n=$n k=$k): scalar=${outScalar[i]} panama=${outPanama[i]}", + ) + } + } + + @Test + fun parity_on_lane_aligned_shapes() { + assertParity(m = 4, n = 32, k = 16, seed = 1) + assertParity(m = 1, n = 64, k = 64, seed = 2) + } + + @Test + fun parity_on_shapes_with_a_scalar_tail() { + // n = 7 is unlikely to be a multiple of any FloatVector lane count, exercising the tail. + assertParity(m = 3, n = 7, k = 5, seed = 3) + assertParity(m = 2, n = 13, k = 9, seed = 4) + assertParity(m = 5, n = 1, k = 3, seed = 5) + } + + @Test + fun parity_on_single_element() { + assertParity(m = 1, n = 1, k = 1, seed = 6) + } + + @Test + fun k_zero_zeroes_the_output_block_in_both() { + val m = 3 + val n = 4 + val outScalar = FloatArray(m * n) { 7f } + val outPanama = FloatArray(m * n) { 7f } + + ScalarFp16MatmulKernel.matmul(FloatArray(0), 0, 0, ByteArray(0), 0, 0, outScalar, 0, n, m, n, 0) + PanamaVectorFp16MatmulKernel.matmul(FloatArray(0), 0, 0, ByteArray(0), 0, 0, outPanama, 0, n, m, n, 0) + + assertTrue(outScalar.all { it == 0f }, "scalar must zero the block") + assertTrue(outPanama.all { it == 0f }, "panama must zero the block") + } + + @Test + fun empty_m_or_n_is_a_no_op_in_both() { + val out = FloatArray(4) { 3f } + ScalarFp16MatmulKernel.matmul(FloatArray(0), 0, 0, ByteArray(0), 0, 0, out, 0, 0, 0, 0, 0) + PanamaVectorFp16MatmulKernel.matmul(FloatArray(0), 0, 0, ByteArray(0), 0, 0, out, 0, 0, 0, 0, 0) + assertTrue(out.all { it == 3f }, "m == 0 / n == 0 must not touch out") + } + + @Test + fun both_kernels_report_the_fp16_codec() { + assertTrue(ScalarFp16MatmulKernel.codec === Fp16Codec) + assertTrue(PanamaVectorFp16MatmulKernel.codec === Fp16Codec) + } + + @Test + fun fp16_kernel_result_tracks_an_exact_fp32_reference() { + // With operands exactly representable in binary16, the kernel must reproduce a plain + // FP32 matmul exactly — proving the decode is right, not merely self-consistent. + val m = 2 + val n = 3 + val k = 4 + val a = floatArrayOf(1f, 2f, -1f, 0.5f, 0.25f, -2f, 4f, 1f) + val bFloats = FloatArray(k * n) { (it % 5) * 0.5f - 1f } + val b = fp16Bytes(bFloats) + + val out = FloatArray(m * n) + ScalarFp16MatmulKernel.matmul(a, 0, k, b, 0, n * 2, out, 0, n, m, n, k) + + for (i in 0 until m) { + for (j in 0 until n) { + var expected = 0f + for (p in 0 until k) expected += a[i * k + p] * bFloats[p * n + j] + assertTrue( + abs(expected - out[i * n + j]) <= 1e-6f, + "exact-value mismatch at ($i,$j): expected=$expected got=${out[i * n + j]}", + ) + } + } + } +} diff --git a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/Fp16MatmulDispatchTest.kt b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/Fp16MatmulDispatchTest.kt new file mode 100644 index 00000000..c1353465 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/Fp16MatmulDispatchTest.kt @@ -0,0 +1,119 @@ +package sk.ainet.exec.tensor.ops + +import kotlin.math.abs +import kotlin.random.Random +import kotlin.test.Test +import kotlin.test.assertTrue +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.exec.kernel.ScalarBf16MatmulKernel +import sk.ainet.exec.kernel.ScalarFp16MatmulKernel +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.Fp16DenseTensorData +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.Fp16Codec + +/** + * Integration tests for the FP32 × FP16 dispatch path in `DefaultCpuOpsJvm.matmul`. + * + * The dispatch branch matches on `NarrowFloatTensorData` and then selects the kernel by the data's + * `codec`. These tests prove an `Fp16DenseTensorData` weight reaches the FP16 kernel — and, more + * importantly, that it is NOT decoded by the BF16 kernel, which would be silent garbage rather + * than an error since both formats are 2 bytes wide. + * + * Mirrors [Bf16MatmulDispatchTest]. + */ +class Fp16MatmulDispatchTest { + + private val ctx = DirectCpuExecutionContext() + + /** binary16 carries 10 mantissa bits; accumulated error scales with `k`. */ + private val fp16TolPerK = 1.5e-3f + + private fun fp32ToFp16Bytes(values: FloatArray): ByteArray { + val out = ByteArray(values.size * 2) + for (i in values.indices) { + val bits = Fp16Codec.encode(values[i]) + out[i * 2] = (bits and 0xFF).toByte() + out[i * 2 + 1] = ((bits ushr 8) and 0xFF).toByte() + } + return out + } + + @Suppress("UNCHECKED_CAST") + private fun fp16Weight(inputDim: Int, outputDim: Int, seed: Int): Pair, ByteArray> { + val rng = Random(seed) + val values = FloatArray(inputDim * outputDim) { rng.nextFloat() - 0.5f } + val bytes = fp32ToFp16Bytes(values) + val data = Fp16DenseTensorData(Shape(inputDim, outputDim), bytes) as TensorData + return ctx.fromData(data, FP32::class) to bytes + } + + private fun assertDispatchMatchesScalar(m: Int, k: Int, n: Int, seed: Int) { + val rng = Random(seed) + val inputFloats = FloatArray(m * k) { rng.nextFloat() - 0.5f } + val (weight, weightBytes) = fp16Weight(k, n, seed) + val input = ctx.fromFloatArray(Shape(m, k), FP32::class, inputFloats) + + val outArr = ctx.ops.matmul(input, weight).data.copyToFloatArray() + + val expected = FloatArray(m * n) + ScalarFp16MatmulKernel.matmul(inputFloats, 0, k, weightBytes, 0, n * 2, expected, 0, n, m, n, k) + + val tol = (fp16TolPerK * k.coerceAtLeast(1)).coerceAtLeast(fp16TolPerK) + for (i in expected.indices) { + assertTrue( + abs(expected[i] - outArr[i]) <= tol, + "FP16 dispatch mismatch at $i: expected=${expected[i]} got=${outArr[i]} tol=$tol", + ) + } + } + + @Test + fun single_batch_matmul_against_fp16_weight_routes_correctly() { + assertDispatchMatchesScalar(m = 1, k = 128, n = 64, seed = 1) + } + + @Test + fun multi_batch_matmul_against_fp16_weight_routes_correctly() { + assertDispatchMatchesScalar(m = 3, k = 256, n = 32, seed = 2) + } + + @Test + fun llm_typical_attention_proj_matmul_routes_correctly() { + assertDispatchMatchesScalar(m = 1, k = 512, n = 512, seed = 3) + } + + @Test + fun fp16_weights_are_not_decoded_by_the_bf16_kernel() { + // Both formats are 2 bytes wide, so a codec mix-up cannot fail loudly — it silently + // produces wrong numbers. This pins that the dispatch picks by codec, not by width: + // decoding the same bytes as BF16 must give a materially different answer, and the + // dispatch must match the FP16 reading. + val m = 1 + val k = 64 + val n = 8 + val rng = Random(42) + val inputFloats = FloatArray(m * k) { rng.nextFloat() - 0.5f } + val (weight, bytes) = fp16Weight(k, n, seed = 42) + val input = ctx.fromFloatArray(Shape(m, k), FP32::class, inputFloats) + + val dispatched = ctx.ops.matmul(input, weight).data.copyToFloatArray() + + val asFp16 = FloatArray(m * n) + ScalarFp16MatmulKernel.matmul(inputFloats, 0, k, bytes, 0, n * 2, asFp16, 0, n, m, n, k) + val asBf16 = FloatArray(m * n) + ScalarBf16MatmulKernel.matmul(inputFloats, 0, k, bytes, 0, n * 2, asBf16, 0, n, m, n, k) + + val fp16Err = dispatched.indices.sumOf { abs(dispatched[it] - asFp16[it]).toDouble() } + val bf16Err = dispatched.indices.sumOf { abs(dispatched[it] - asBf16[it]).toDouble() } + + assertTrue(fp16Err < 1e-3, "dispatch must match the FP16 decode, err=$fp16Err") + assertTrue( + bf16Err > fp16Err * 10, + "the BF16 decode of the same bytes must differ materially " + + "(fp16Err=$fp16Err bf16Err=$bf16Err); if these are close the test proves nothing", + ) + } +} diff --git a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/NarrowFloatConvertTest.kt b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/NarrowFloatConvertTest.kt new file mode 100644 index 00000000..7d7811b7 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/NarrowFloatConvertTest.kt @@ -0,0 +1,102 @@ +package sk.ainet.exec.tensor.ops + +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.types.BF16 +import sk.ainet.lang.types.Bf16Codec +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.assertEquals +import kotlin.test.assertTrue + +/** + * `convert` to a 16-bit float target must actually round. + * + * It previously shared the FP32 branch, so `convert(x, FP16)` only re-tagged the dtype while + * keeping full FP32 precision — the tensor claimed to be FP16 while holding values no FP16 can + * represent, with no rounding, no overflow to Inf past the format's range, and no NaN handling. + * BF16 was rejected as a target outright. + */ +class NarrowFloatConvertTest { + + private val ctx = DirectCpuExecutionContext() + + private fun fp32(vararg values: Float) = + ctx.fromFloatArray(Shape(values.size), FP32::class, values) + + @Test + fun convert_to_fp16_rounds_values_to_binary16() { + // 1 + 2^-20 is far below binary16's 2^-10 step, so it must collapse onto 1.0. + val src = fp32(1.0f + 1.0f / 1048576.0f, 3.14159265f, -2.71828f) + val out = ctx.ops.convert(src, FP16).data.copyToFloatArray() + + assertEquals(1.0f, out[0], "sub-ulp detail must be rounded away, not retained") + for ((i, v) in floatArrayOf(1.0f + 1.0f / 1048576.0f, 3.14159265f, -2.71828f).withIndex()) { + assertEquals( + Fp16Codec.decode(Fp16Codec.encode(v)), out[i], + "element $i must equal the binary16 rounding of the input", + ) + } + } + + @Test + fun convert_to_bf16_is_supported_and_truncates() { + val values = floatArrayOf(1.0f, 3.14159265f, -2.71828f, 1234.5f) + val out = ctx.ops.convert(fp32(*values), BF16).data.copyToFloatArray() + for ((i, v) in values.withIndex()) { + assertEquals( + Bf16Codec.decode(Bf16Codec.encode(v)), out[i], + "element $i must equal the bf16 encoding of the input", + ) + } + } + + @Test + fun convert_to_fp16_overflows_to_infinity_past_the_ceiling() { + // The old retag path kept 70000f verbatim in a tensor claiming to be FP16, a value + // binary16 cannot represent at all. + val out = ctx.ops.convert(fp32(70000f, -70000f), FP16).data.copyToFloatArray() + assertTrue(out[0].isInfinite() && out[0] > 0, "must overflow to +Inf, got ${out[0]}") + assertTrue(out[1].isInfinite() && out[1] < 0, "must overflow to -Inf, got ${out[1]}") + } + + @Test + fun convert_to_bf16_keeps_the_same_value_finite() { + // bf16 has FP32's exponent range, so the same input that overflows fp16 stays finite. + val out = ctx.ops.convert(fp32(70000f), BF16).data.copyToFloatArray() + assertTrue(out[0].isFinite(), "bf16 must not overflow at 70000, got ${out[0]}") + } + + @Test + fun convert_to_fp16_flushes_subnormal_underflow_to_zero() { + // 2^-30 is far below binary16's smallest subnormal (2^-24). + val out = ctx.ops.convert(fp32(1.0f / 1073741824.0f), FP16).data.copyToFloatArray() + assertEquals(0.0f, out[0], "must underflow to zero rather than retain the FP32 value") + } + + @Test + fun convert_reports_the_target_dtype() { + assertEquals(FP16::class, ctx.ops.convert(fp32(1f), FP16).dtype) + assertEquals(BF16::class, ctx.ops.convert(fp32(1f), BF16).dtype) + } + + @Test + fun convert_to_fp32_is_unchanged() { + val values = floatArrayOf(1.0f, 3.14159265f, -2.71828f) + val out = ctx.ops.convert(fp32(*values), FP32).data.copyToFloatArray() + for ((i, v) in values.withIndex()) assertEquals(v, out[i], "FP32 target must be exact") + } + + @Test + fun round_tripping_through_fp16_is_idempotent() { + // Rounding an already-rounded value must not move it again. + val src = fp32(3.14159265f, -2.71828f, 1234.5f) + val once = ctx.ops.convert(src, FP16).data.copyToFloatArray() + val twice = ctx.ops.convert( + ctx.fromFloatArray(Shape(once.size), FP32::class, once), FP16, + ).data.copyToFloatArray() + for (i in once.indices) assertEquals(once[i], twice[i], "element $i must be stable") + } +} diff --git a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt index 04d32f42..d1ae2937 100644 --- a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt +++ b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt @@ -5,6 +5,8 @@ import sk.ainet.io.ParametersLoader import sk.ainet.io.RandomAccessSource import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.Bf16DenseTensorData +import sk.ainet.lang.tensor.data.Fp16DenseTensorData import sk.ainet.lang.tensor.data.Q4_KBlockTensorData import sk.ainet.lang.tensor.data.Q5_KBlockTensorData import sk.ainet.lang.tensor.data.Q6_KBlockTensorData @@ -33,7 +35,16 @@ import kotlin.reflect.KClass */ public class StreamingGgufParametersLoader( private val sourceProvider: () -> RandomAccessSource, - private val onProgress: (current: Long, total: Long, message: String?) -> Unit = { _, _, _ -> } + private val onProgress: (current: Long, total: Long, message: String?) -> Unit = { _, _, _ -> }, + /** + * Keep `F16` source tensors in their on-disk 2-bytes-per-element layout instead of widening + * them to FP32 at load. Off by default — flip via `withPolicy(Require(FP16))`. + */ + private val keepF16Native: Boolean = false, + /** + * Keep `BF16` source tensors packed. Off by default — flip via `withPolicy(Require(BF16))`. + */ + private val keepBf16Native: Boolean = false, ) : ParametersLoader { @Suppress("UNCHECKED_CAST") @@ -68,20 +79,28 @@ public class StreamingGgufParametersLoader( } } - GGMLQuantizationType.F16 -> { - val floats = dequantF16(rawBytes) - when (dtype) { - FP32::class -> ctx.fromFloatArray(shape, dtype, floats) as Tensor - else -> null + GGMLQuantizationType.F16 -> when (dtype) { + FP32::class -> if (keepF16Native) { + // Zero-widening path: hand the on-disk bytes straight through as + // packed binary16. Consumers still see Float on read. + @Suppress("UNCHECKED_CAST") + val packed = Fp16DenseTensorData(shape, rawBytes) + ctx.fromData(packed as sk.ainet.lang.tensor.data.TensorData, dtype) + } else { + ctx.fromFloatArray(shape, dtype, dequantF16(rawBytes)) as Tensor } + else -> null } - GGMLQuantizationType.BF16 -> { - val floats = dequantBF16(rawBytes) - when (dtype) { - FP32::class -> ctx.fromFloatArray(shape, dtype, floats) as Tensor - else -> null + GGMLQuantizationType.BF16 -> when (dtype) { + FP32::class -> if (keepBf16Native) { + @Suppress("UNCHECKED_CAST") + val packed = Bf16DenseTensorData(shape, rawBytes) + ctx.fromData(packed as sk.ainet.lang.tensor.data.TensorData, dtype) + } else { + ctx.fromFloatArray(shape, dtype, dequantBF16(rawBytes)) as Tensor } + else -> null } GGMLQuantizationType.Q4_K -> { @@ -202,7 +221,25 @@ public class StreamingGgufParametersLoader( onProgress: (current: Long, total: Long, message: String?) -> Unit = { _, _, _ -> }, ): StreamingGgufParametersLoader { validatePolicy(policy) - return StreamingGgufParametersLoader(sourceProvider, onProgress) + return StreamingGgufParametersLoader( + sourceProvider = sourceProvider, + onProgress = onProgress, + keepF16Native = keepsNative(policy, FP16), + keepBf16Native = keepsNative(policy, BF16), + ) + } + + /** + * Whether [policy] asks for [native] tensors to stay in their on-disk 16-bit layout. + * + * Only the format the policy actually names is kept — neither narrow format can be turned + * into the other without a lossy re-encode, so `Require(BF16)` must still widen F16 sources. + */ + internal fun keepsNative(policy: DTypePolicy, native: DType): Boolean = when (policy) { + DTypePolicy.Any -> false + is DTypePolicy.Require -> policy.target == native + is DTypePolicy.Prefer -> policy.target == native + is DTypePolicy.OneOf -> native in policy.allowed } internal fun validatePolicy(policy: DTypePolicy) { @@ -211,18 +248,10 @@ public class StreamingGgufParametersLoader( is DTypePolicy.Prefer -> Unit is DTypePolicy.OneOf -> Unit is DTypePolicy.Require -> when (policy.target) { - FP32 -> Unit - BF16 -> throw IllegalArgumentException( - "StreamingGgufParametersLoader: Require(BF16) is not supported — " + - "GGUF BF16 sources are dequanted to FP32 by this loader today (no KEEP_NATIVE " + - "GGUF path yet). Use Any or Prefer(BF16) to accept the dequant fallback, or " + - "wait for the policy-aware GGUF reader to land.", - ) - FP16 -> throw IllegalArgumentException( - "StreamingGgufParametersLoader: Require(FP16) is not supported — " + - "GGUF F16 sources are dequanted to FP32 by this loader today (no Fp16DenseTensorData " + - "backing yet). Use Any or Prefer(FP16) to accept the dequant fallback.", - ) + // FP16 / BF16 are satisfiable for sources already in that format: the loader + // hands the packed bytes through via Fp16/Bf16DenseTensorData. Sources in any + // other format still widen to FP32 rather than being re-encoded. + FP32, FP16, BF16 -> Unit else -> throw IllegalArgumentException( "StreamingGgufParametersLoader: Require(${policy.target.name}) is not satisfiable — " + "this loader produces FP32 / Int32 / Q4_K / Q8_0 tensors only, and does not cast " + diff --git a/skainet-io/skainet-io-gguf/src/commonTest/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoaderPolicyTest.kt b/skainet-io/skainet-io-gguf/src/commonTest/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoaderPolicyTest.kt index bf32f5a0..f496b252 100644 --- a/skainet-io/skainet-io-gguf/src/commonTest/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoaderPolicyTest.kt +++ b/skainet-io/skainet-io-gguf/src/commonTest/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoaderPolicyTest.kt @@ -30,21 +30,31 @@ class StreamingGgufParametersLoaderPolicyTest { } @Test - fun require_bf16_fails_fast_with_clear_message() { - val ex = assertFailsWith { - StreamingGgufParametersLoader.validatePolicy(DTypePolicy.Require(BF16)) - } - val msg = ex.message ?: "" - assertTrue(msg.contains("Require(BF16)"), msg) - assertTrue(msg.contains("KEEP_NATIVE"), msg) + fun require_bf16_is_now_accepted_and_keeps_bf16_sources_packed() { + // Previously rejected: the GGUF loader had no KEEP_NATIVE path and always widened. + StreamingGgufParametersLoader.validatePolicy(DTypePolicy.Require(BF16)) + assertTrue(StreamingGgufParametersLoader.keepsNative(DTypePolicy.Require(BF16), BF16)) } @Test - fun require_fp16_fails_fast_with_clear_message() { - val ex = assertFailsWith { - StreamingGgufParametersLoader.validatePolicy(DTypePolicy.Require(FP16)) - } - assertTrue(ex.message?.contains("Require(FP16)") == true, ex.message ?: "") + fun require_fp16_is_now_accepted_and_keeps_f16_sources_packed() { + // Previously rejected for want of an Fp16DenseTensorData backing, which now exists. + StreamingGgufParametersLoader.validatePolicy(DTypePolicy.Require(FP16)) + assertTrue(StreamingGgufParametersLoader.keepsNative(DTypePolicy.Require(FP16), FP16)) + } + + @Test + fun neither_narrow_format_is_kept_native_for_the_other() { + // Converting between the two is a lossy re-encode, so a policy naming one must widen + // sources in the other rather than mis-tagging their bytes. + assertTrue(!StreamingGgufParametersLoader.keepsNative(DTypePolicy.Require(BF16), FP16)) + assertTrue(!StreamingGgufParametersLoader.keepsNative(DTypePolicy.Require(FP16), BF16)) + } + + @Test + fun any_policy_still_widens_both_narrow_formats() { + assertTrue(!StreamingGgufParametersLoader.keepsNative(DTypePolicy.Any, BF16)) + assertTrue(!StreamingGgufParametersLoader.keepsNative(DTypePolicy.Any, FP16)) } @Test diff --git a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/Bf16LoadPolicy.kt b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/Bf16LoadPolicy.kt index 016a9669..df54a26e 100644 --- a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/Bf16LoadPolicy.kt +++ b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/Bf16LoadPolicy.kt @@ -27,53 +27,17 @@ import sk.ainet.lang.types.FP32 * `DefaultCpuOpsJvm.chooseQuantizedMatmul` (Phase 3 follow-up, * separately tracked). */ -public enum class Bf16LoadPolicy { - /** - * Default. Dequantise every BFLOAT16 tensor to FP32 at load time - * via the existing `dequantBF16` helper, then wrap as a - * `FloatArrayTensorData` (same as `FLOAT32` source tensors). - * - * Memory cost: 2× the on-disk size for each BF16 tensor. - * Runtime: zero extra dispatch — every matmul gets FP32 operands. - */ - DEQUANT_TO_FP32, +public typealias Bf16LoadPolicy = NarrowFloatLoadPolicy - /** - * Keep BFLOAT16 tensors in their on-disk packed-2-bytes-per-element - * layout. The loader emits a `Bf16DenseTensorData` (in - * `skainet-lang-core`) instead of dequanting; the tensor still - * advertises FP32 dtype to consumers (the underlying `get` decodes - * on read), but its `tensor.data` is recognisable as - * `Bf16TensorData` so a matmul dispatch can route to the SIMD - * `Bf16MatmulKernel` SPI. - * - * Memory cost: identical to the on-disk size — no doubling. - * Runtime: matmul dispatch picks up the BF16 SPI kernel when one - * is registered; falls back to per-multiply dequant otherwise. - * - * **Caveat**: any non-matmul op that touches the BF16 tensor pays - * a per-element decode cost via `get`. Don't flip this unless the - * model's hot path is dominated by matmuls (the typical - * transformer case). - */ - KEEP_NATIVE, - ; - - /** - * Maps this BF16-specific enum onto the generalised - * [DTypePolicy] sealed type. [DEQUANT_TO_FP32] becomes - * `Require(FP32)` (the loader must hand consumers an FP32 - * tensor); [KEEP_NATIVE] becomes `Require(BF16)` (consumers - * dispatch on the native BF16 dtype). - * - * Bridge for the RFC's policy-driven loader work - * (`rfc.md`, issue #615): existing call sites keep using this - * enum verbatim while new code paths can flow through - * [DTypePolicy] uniformly. The two are equivalent for BF16 — - * this method is the explicit equivalence proof. - */ - public fun toDTypePolicy(): DTypePolicy = when (this) { - DEQUANT_TO_FP32 -> DTypePolicy.Require(FP32) - KEEP_NATIVE -> DTypePolicy.Require(BF16) - } +/** + * Maps this narrow-float policy onto the generalised [DTypePolicy] sealed type. + * [NarrowFloatLoadPolicy.DEQUANT_TO_FP32] becomes `Require(FP32)` (the loader must hand consumers + * an FP32 tensor); [NarrowFloatLoadPolicy.KEEP_NATIVE] becomes `Require(BF16)`. + * + * Kept BF16-specific for source compatibility with existing call sites. For FP16, build the + * policy directly (`DTypePolicy.Require(FP16)`) — see `SafeTensorsParametersLoader.mapPolicyToFp16`. + */ +public fun Bf16LoadPolicy.toDTypePolicy(): DTypePolicy = when (this) { + NarrowFloatLoadPolicy.DEQUANT_TO_FP32 -> DTypePolicy.Require(FP32) + NarrowFloatLoadPolicy.KEEP_NATIVE -> DTypePolicy.Require(BF16) } 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 new file mode 100644 index 00000000..55d9febc --- /dev/null +++ b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/NarrowFloatLoadPolicy.kt @@ -0,0 +1,40 @@ +package sk.ainet.io.safetensors + +/** + * Controls how a loader handles **16-bit float** source tensors (`FLOAT16` / `BFLOAT16`). + * + * The choice is whether to **widen to FP32 at load time** — doubling the resident footprint of + * every narrow tensor — or to **keep the on-disk bytes** and let the matmul dispatch route to a + * narrow-float kernel that widens per-multiply instead. + * + * Both paths are numerically equivalent when they reach the same kernel: the narrow → FP32 + * conversion is exact in either direction (FP32 is a strict superset of both formats), applied + * either once at load ([DEQUANT_TO_FP32]) or per-multiply ([KEEP_NATIVE]). + * + * This supersedes the BF16-only `Bf16LoadPolicy`, which remains as a typealias. + */ +public enum class NarrowFloatLoadPolicy { + /** + * Default. Widen every narrow-float tensor to FP32 at load time and wrap it as + * `FloatArrayTensorData`, exactly as a `FLOAT32` source tensor would be. + * + * Memory: 2× the on-disk size for each narrow tensor. Runtime: no extra dispatch — every + * matmul receives FP32 operands. + */ + DEQUANT_TO_FP32, + + /** + * Keep narrow-float tensors in their on-disk packed 2-bytes-per-element layout. The loader + * emits a `NarrowFloatDenseTensorData` (`Bf16DenseTensorData` / `Fp16DenseTensorData`) rather + * than widening. The tensor still presents `Float` on read — `get` decodes — but its + * `tensor.data` is recognisable as `NarrowFloatTensorData`, which is what lets a matmul + * dispatch pick the narrow-float SPI kernel. + * + * Memory: identical to the on-disk size — no doubling. Runtime: the narrow kernel is used when + * one is registered; otherwise consumers fall back to per-element decode. + * + * **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. + */ + KEEP_NATIVE, +} diff --git a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoader.kt b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoader.kt index 2d54ed93..ef89d4e9 100644 --- a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoader.kt +++ b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoader.kt @@ -7,6 +7,7 @@ import sk.ainet.io.model.DataType import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.Tensor import sk.ainet.lang.tensor.data.Bf16DenseTensorData +import sk.ainet.lang.tensor.data.Fp16DenseTensorData import sk.ainet.lang.tensor.data.TensorData import sk.ainet.lang.types.BF16 import sk.ainet.lang.types.DType @@ -46,7 +47,8 @@ import kotlin.reflect.KClass class SafeTensorsParametersLoader( private val sourceProvider: () -> RandomAccessSource, private val onProgress: (current: Long, total: Long, message: String?) -> Unit = { _, _, _ -> }, - private val bf16Policy: Bf16LoadPolicy = Bf16LoadPolicy.DEQUANT_TO_FP32, + private val bf16Policy: Bf16LoadPolicy = NarrowFloatLoadPolicy.DEQUANT_TO_FP32, + private val fp16Policy: NarrowFloatLoadPolicy = NarrowFloatLoadPolicy.DEQUANT_TO_FP32, ) : ParametersLoader { override suspend fun load( @@ -86,10 +88,22 @@ class SafeTensorsParametersLoader( DataType.FLOAT16 -> { require(dtype == FP32::class) { - "SafeTensors F16 tensor '${tensorInfo.name}' requires FP32 dtype (dequant), got ${dtype.simpleName}" + "SafeTensors F16 tensor '${tensorInfo.name}' requires FP32 dtype, got ${dtype.simpleName}" + } + when (fp16Policy) { + NarrowFloatLoadPolicy.DEQUANT_TO_FP32 -> { + val floats = dequantF16(bytes) + ctx.wrapFloatArray(shape, dtype, floats) as Tensor + } + NarrowFloatLoadPolicy.KEEP_NATIVE -> { + // Mirrors the BF16 arm below: wrap the on-disk F16 bytes directly. + // dtype stays FP32 from the consumer's POV (the tensor data decodes + // on read); the storage type is what a narrow-float matmul dispatch + // pattern-matches on. + val fp16Data = Fp16DenseTensorData(shape, bytes) + ctx.fromData(fp16Data as TensorData, dtype) + } } - val floats = dequantF16(bytes) - ctx.wrapFloatArray(shape, dtype, floats) as Tensor } DataType.BFLOAT16 -> { @@ -325,28 +339,37 @@ class SafeTensorsParametersLoader( sourceProvider = sourceProvider, onProgress = onProgress, bf16Policy = mapPolicyToBf16(policy), + fp16Policy = mapPolicyToFp16(policy), ) - internal fun mapPolicyToBf16(policy: DTypePolicy): Bf16LoadPolicy = when (policy) { - DTypePolicy.Any -> Bf16LoadPolicy.DEQUANT_TO_FP32 - is DTypePolicy.Require -> when (policy.target) { - BF16 -> Bf16LoadPolicy.KEEP_NATIVE - FP32 -> Bf16LoadPolicy.DEQUANT_TO_FP32 - FP16 -> throw IllegalArgumentException( - "SafeTensorsParametersLoader: Require(FP16) is not supported — " + - "F16 KEEP_NATIVE has no Fp16DenseTensorData backing yet. " + - "Use Require(FP32) to dequant F16 sources, or Any to inherit the adaptive default.", - ) - else -> throw IllegalArgumentException( - "SafeTensorsParametersLoader: Require(${policy.target.name}) is not satisfiable — " + - "the loader produces FP32 / BF16 / Int32 / Int8 tensors depending on source dtype; " + - "it cannot fabricate ${policy.target.name} from arbitrary sources.", - ) + internal fun mapPolicyToBf16(policy: DTypePolicy): Bf16LoadPolicy = + mapPolicyToNarrow(policy, BF16) + + internal fun mapPolicyToFp16(policy: DTypePolicy): NarrowFloatLoadPolicy = + mapPolicyToNarrow(policy, FP16) + + /** + * Resolve [policy] for one narrow-float source format. A tensor is kept native only when + * the policy names *that* format — `Require(BF16)` must not keep F16 tensors packed, and + * vice versa, since neither can be converted to the other without a lossy re-encode. + */ + private fun mapPolicyToNarrow(policy: DTypePolicy, native: DType): NarrowFloatLoadPolicy = + when (policy) { + DTypePolicy.Any -> NarrowFloatLoadPolicy.DEQUANT_TO_FP32 + is DTypePolicy.Require -> when (policy.target) { + native -> NarrowFloatLoadPolicy.KEEP_NATIVE + // The other narrow format, or FP32: this format still widens. + BF16, FP16, FP32 -> NarrowFloatLoadPolicy.DEQUANT_TO_FP32 + else -> throw IllegalArgumentException( + "SafeTensorsParametersLoader: Require(${policy.target.name}) is not satisfiable — " + + "the loader produces FP32 / BF16 / FP16 / Int32 / Int8 tensors depending on " + + "source dtype; it cannot fabricate ${policy.target.name} from arbitrary sources.", + ) + } + is DTypePolicy.Prefer -> if (policy.target == native) NarrowFloatLoadPolicy.KEEP_NATIVE + else NarrowFloatLoadPolicy.DEQUANT_TO_FP32 + is DTypePolicy.OneOf -> if (native in policy.allowed) NarrowFloatLoadPolicy.KEEP_NATIVE + else NarrowFloatLoadPolicy.DEQUANT_TO_FP32 } - is DTypePolicy.Prefer -> if (policy.target == BF16) Bf16LoadPolicy.KEEP_NATIVE - else Bf16LoadPolicy.DEQUANT_TO_FP32 - is DTypePolicy.OneOf -> if (BF16 in policy.allowed) Bf16LoadPolicy.KEEP_NATIVE - else Bf16LoadPolicy.DEQUANT_TO_FP32 - } } } diff --git a/skainet-io/skainet-io-safetensors/src/commonTest/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoaderPolicyTest.kt b/skainet-io/skainet-io-safetensors/src/commonTest/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoaderPolicyTest.kt index 7296508d..6b7143c1 100644 --- a/skainet-io/skainet-io-safetensors/src/commonTest/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoaderPolicyTest.kt +++ b/skainet-io/skainet-io-safetensors/src/commonTest/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoaderPolicyTest.kt @@ -43,15 +43,41 @@ class SafeTensorsParametersLoaderPolicyTest { } @Test - fun require_fp16_fails_with_explicit_message() { - val ex = assertFailsWith { - SafeTensorsParametersLoader.mapPolicyToBf16(DTypePolicy.Require(FP16)) - } - // The error message must point the operator at the alternative — - // RFC says "fail-fast with clear diagnostics," not just throw. - val msg = ex.message ?: "" - assertEquals(true, msg.contains("Require(FP16)"), "msg: $msg") - assertEquals(true, msg.contains("Fp16DenseTensorData"), "msg: $msg") + fun require_fp16_now_keeps_f16_sources_native() { + // Previously this threw: there was no Fp16DenseTensorData to back a KEEP_NATIVE F16 path. + // That backing now exists, so Require(FP16) is satisfiable for F16 sources. + assertEquals( + NarrowFloatLoadPolicy.KEEP_NATIVE, + SafeTensorsParametersLoader.mapPolicyToFp16(DTypePolicy.Require(FP16)), + ) + } + + @Test + fun the_two_narrow_formats_do_not_keep_each_other_native() { + // Neither format can be produced from the other without a lossy re-encode, so a policy + // naming one must widen the other rather than silently mis-tagging it. + assertEquals( + NarrowFloatLoadPolicy.DEQUANT_TO_FP32, + SafeTensorsParametersLoader.mapPolicyToFp16(DTypePolicy.Require(BF16)), + "Require(BF16) must not keep F16 sources packed", + ) + assertEquals( + NarrowFloatLoadPolicy.DEQUANT_TO_FP32, + SafeTensorsParametersLoader.mapPolicyToBf16(DTypePolicy.Require(FP16)), + "Require(FP16) must not keep BF16 sources packed", + ) + } + + @Test + fun require_fp16_still_keeps_bf16_behaviour_intact() { + assertEquals( + NarrowFloatLoadPolicy.KEEP_NATIVE, + SafeTensorsParametersLoader.mapPolicyToBf16(DTypePolicy.Require(BF16)), + ) + assertEquals( + NarrowFloatLoadPolicy.DEQUANT_TO_FP32, + SafeTensorsParametersLoader.mapPolicyToBf16(DTypePolicy.Require(FP32)), + ) } @Test diff --git a/skainet-io/skainet-io-safetensors/src/jvmTest/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoaderFp16PolicyTest.kt b/skainet-io/skainet-io-safetensors/src/jvmTest/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoaderFp16PolicyTest.kt new file mode 100644 index 00000000..0934efd5 --- /dev/null +++ b/skainet-io/skainet-io-safetensors/src/jvmTest/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoaderFp16PolicyTest.kt @@ -0,0 +1,184 @@ +package sk.ainet.io.safetensors + +import kotlinx.coroutines.runBlocking +import org.junit.Test +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.Bf16TensorData +import sk.ainet.lang.tensor.data.FloatArrayTensorData +import sk.ainet.lang.tensor.data.Fp16DenseTensorData +import sk.ainet.lang.tensor.data.NarrowFloatTensorData +import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.Fp16Codec +import java.io.File +import java.nio.ByteBuffer +import java.nio.ByteOrder +import java.nio.file.Files +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +/** + * FP16 counterpart to [SafeTensorsParametersLoaderBf16PolicyTest]. + * + * Before `Fp16DenseTensorData` existed the loader had no choice but to widen every F16 tensor to + * FP32, and `Require(FP16)` was rejected outright. These tests pin the KEEP_NATIVE path: on-disk + * bytes preserved verbatim, values decoded on read, and BF16 handling untouched. + */ +class SafeTensorsParametersLoaderFp16PolicyTest { + + /** Encode FP32 to IEEE binary16 bytes, little-endian. */ + private fun fp32ToFp16Bytes(values: FloatArray): ByteArray { + val out = ByteArray(values.size * 2) + for (i in values.indices) { + val bits = Fp16Codec.encode(values[i]) + out[i * 2] = (bits and 0xFF).toByte() + out[i * 2 + 1] = ((bits ushr 8) and 0xFF).toByte() + } + return out + } + + private fun createF16SafeTensorsFile(name: String, values: FloatArray): File { + val bytes = fp32ToFp16Bytes(values) + val headerJson = + "{\"$name\": {\"dtype\": \"F16\", \"shape\": [${values.size}], \"data_offsets\": [0, ${bytes.size}]}}" + val headerBytes = headerJson.toByteArray(Charsets.UTF_8) + val tempFile = Files.createTempFile("test_f16_safetensors", ".safetensors").toFile() + tempFile.deleteOnExit() + tempFile.outputStream().use { out -> + out.write( + ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN) + .putLong(headerBytes.size.toLong()).array(), + ) + out.write(headerBytes) + out.write(bytes) + } + return tempFile + } + + /** Build a file holding one F16 tensor and one BF16 tensor, to prove the policies are independent. */ + private fun createF16AndBf16File(f16Values: FloatArray, bf16Values: FloatArray): File { + val f16Bytes = fp32ToFp16Bytes(f16Values) + val bf16Bytes = ByteArray(bf16Values.size * 2) + for (i in bf16Values.indices) { + val b = (bf16Values[i].toRawBits() ushr 16) and 0xFFFF + bf16Bytes[i * 2] = (b and 0xFF).toByte() + bf16Bytes[i * 2 + 1] = ((b ushr 8) and 0xFF).toByte() + } + val f16End = f16Bytes.size.toLong() + val bf16End = f16End + bf16Bytes.size + val headerJson = + "{\"w_f16\": {\"dtype\": \"F16\", \"shape\": [${f16Values.size}], \"data_offsets\": [0, $f16End]}," + + "\"w_bf16\": {\"dtype\": \"BF16\", \"shape\": [${bf16Values.size}], " + + "\"data_offsets\": [$f16End, $bf16End]}}" + val headerBytes = headerJson.toByteArray(Charsets.UTF_8) + val tempFile = Files.createTempFile("test_f16_bf16_safetensors", ".safetensors").toFile() + tempFile.deleteOnExit() + tempFile.outputStream().use { out -> + out.write( + ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN) + .putLong(headerBytes.size.toLong()).array(), + ) + out.write(headerBytes) + out.write(f16Bytes) + out.write(bf16Bytes) + } + return tempFile + } + + private fun loadAll( + file: File, + fp16: NarrowFloatLoadPolicy = NarrowFloatLoadPolicy.DEQUANT_TO_FP32, + bf16: NarrowFloatLoadPolicy = NarrowFloatLoadPolicy.DEQUANT_TO_FP32, + ): Map> = runBlocking { + val ctx = DirectCpuExecutionContext.create() + val loader = SafeTensorsParametersLoader( + sourceProvider = { JvmRandomAccessSource.open(file) }, + bf16Policy = bf16, + fp16Policy = fp16, + ) + val out = mutableMapOf>() + loader.load(ctx, FP32::class) { name, tensor -> out[name] = tensor } + out + } + + @Test + fun fp16_default_policy_widens_to_fp32_floatArray() { + val values = floatArrayOf(0.0f, 1.0f, -1.0f, 0.5f, 3.0f, -2.5f) + val file = createF16SafeTensorsFile("weight", values) + + val weight = loadAll(file)["weight"] ?: error("missing 'weight'") + assertTrue( + weight.data is FloatArrayTensorData<*>, + "default policy must widen, got ${weight.data::class.simpleName}", + ) + // All of these are exactly representable in binary16. + assertContentEquals(values, weight.data.copyToFloatArray()) + } + + @Test + fun fp16_keep_native_emits_fp16DenseTensorData_with_on_disk_bytes() { + val values = floatArrayOf(0.0f, 1.0f, -1.0f, 0.5f, 3.0f, -2.5f, 100.0f, -64.0f) + val file = createF16SafeTensorsFile("weight", values) + + val weight = loadAll(file, fp16 = NarrowFloatLoadPolicy.KEEP_NATIVE)["weight"] + ?: error("missing 'weight'") + + assertTrue( + weight.data is Fp16DenseTensorData, + "KEEP_NATIVE must produce Fp16DenseTensorData, got ${weight.data::class.simpleName}", + ) + assertTrue(weight.data is NarrowFloatTensorData, "must be recognizable to narrow dispatch") + assertTrue( + weight.data !is Bf16TensorData, + "an F16 tensor must never be mistaken for BF16 — the bit layouts differ", + ) + + // Byte-for-byte identity proves no widening pass ran. + assertContentEquals( + fp32ToFp16Bytes(values), + (weight.data as Fp16DenseTensorData).packedData, + "KEEP_NATIVE must preserve on-disk F16 bytes verbatim", + ) + assertEquals( + values.size * 2, (weight.data as Fp16DenseTensorData).packedData.size, + "storage must stay 2 bytes per element", + ) + } + + @Test + fun fp16_keep_native_decodes_bit_identically_to_the_widening_path() { + // Both policies apply the same binary16 -> f32 decode; only the timing differs. + val values = FloatArray(64) { (it - 32) * 0.25f } + val file = createF16SafeTensorsFile("w", values) + + val widened = loadAll(file)["w"]!!.data.copyToFloatArray() + val native = loadAll(file, fp16 = NarrowFloatLoadPolicy.KEEP_NATIVE)["w"]!!.data.copyToFloatArray() + + assertEquals(widened.size, native.size) + for (i in widened.indices) { + assertEquals( + widened[i].toRawBits(), native[i].toRawBits(), + "bit-identity expected at $i: widened=${widened[i]} native=${native[i]}", + ) + } + } + + @Test + fun the_two_narrow_policies_are_independent() { + val f16Values = floatArrayOf(1.0f, 2.0f, 4.0f, 8.0f) + val bf16Values = floatArrayOf(0.5f, 0.25f, 16.0f, -3.0f) + val file = createF16AndBf16File(f16Values, bf16Values) + + // Keep F16 packed, widen BF16. + val a = loadAll(file, fp16 = NarrowFloatLoadPolicy.KEEP_NATIVE) + assertTrue(a["w_f16"]!!.data is Fp16DenseTensorData, "F16 should be packed") + assertTrue(a["w_bf16"]!!.data is FloatArrayTensorData<*>, "BF16 should be widened") + + // ...and the mirror image. + val b = loadAll(file, bf16 = NarrowFloatLoadPolicy.KEEP_NATIVE) + assertTrue(b["w_f16"]!!.data is FloatArrayTensorData<*>, "F16 should be widened") + assertTrue(b["w_bf16"]!!.data is Bf16TensorData, "BF16 should be packed") + } +} 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 338bceb0..84f41025 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 @@ -2862,16 +2862,10 @@ public final class sk/ainet/lang/tensor/benchmark/SlicingBenchmarksKt { public static synthetic fun slicingBenchmarkSuite$default (Ljava/lang/String;Lsk/ainet/lang/tensor/benchmark/SlicingBenchmarkConfig;Lkotlin/jvm/functions/Function0;Lkotlin/jvm/functions/Function2;ILjava/lang/Object;)Lsk/ainet/benchmark/BenchmarkSuite; } -public final class sk/ainet/lang/tensor/data/Bf16DenseTensorData : sk/ainet/lang/tensor/data/Bf16TensorData { +public final class sk/ainet/lang/tensor/data/Bf16DenseTensorData : sk/ainet/lang/tensor/data/NarrowFloatDenseTensorData, sk/ainet/lang/tensor/data/Bf16TensorData { public static final field Companion Lsk/ainet/lang/tensor/data/Bf16DenseTensorData$Companion; public fun (Lsk/ainet/lang/tensor/Shape;[B)V - public fun copyToFloatArray ()[F - public fun get ([I)Ljava/lang/Float; - public synthetic fun get ([I)Ljava/lang/Object; - 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 fun getCodec ()Lsk/ainet/lang/types/NarrowFloatCodec; } public final class sk/ainet/lang/tensor/data/Bf16DenseTensorData$Companion { @@ -2879,10 +2873,10 @@ public final class sk/ainet/lang/tensor/data/Bf16DenseTensorData$Companion { public final fun fromRawBytes (Lsk/ainet/lang/tensor/Shape;[B)Lsk/ainet/lang/tensor/data/Bf16DenseTensorData; } -public abstract interface class sk/ainet/lang/tensor/data/Bf16TensorData : sk/ainet/lang/tensor/data/TensorData { +public abstract interface class sk/ainet/lang/tensor/data/Bf16TensorData : sk/ainet/lang/tensor/data/NarrowFloatTensorData { public static final field BYTES_PER_ELEMENT I public static final field Companion Lsk/ainet/lang/tensor/data/Bf16TensorData$Companion; - public abstract fun getPackedData ()[B + public fun getCodec ()Lsk/ainet/lang/types/NarrowFloatCodec; } public final class sk/ainet/lang/tensor/data/Bf16TensorData$Companion { @@ -2893,6 +2887,7 @@ public final class sk/ainet/lang/tensor/data/Bf16TensorData$Companion { public final class sk/ainet/lang/tensor/data/Bf16TensorData$DefaultImpls { public static fun copyToFloatArray (Lsk/ainet/lang/tensor/data/Bf16TensorData;)[F + public static fun getCodec (Lsk/ainet/lang/tensor/data/Bf16TensorData;)Lsk/ainet/lang/types/NarrowFloatCodec; } public final class sk/ainet/lang/tensor/data/Bf16TensorDataKt { @@ -2967,6 +2962,16 @@ public final class sk/ainet/lang/tensor/data/FloatBufferTensorData$DefaultImpls public static fun copyToFloatArray (Lsk/ainet/lang/tensor/data/FloatBufferTensorData;)[F } +public final class sk/ainet/lang/tensor/data/Fp16DenseTensorData : sk/ainet/lang/tensor/data/NarrowFloatDenseTensorData { + public static final field Companion Lsk/ainet/lang/tensor/data/Fp16DenseTensorData$Companion; + public fun (Lsk/ainet/lang/tensor/Shape;[B)V +} + +public final class sk/ainet/lang/tensor/data/Fp16DenseTensorData$Companion { + public final fun fromFloatArray (Lsk/ainet/lang/tensor/Shape;[F)Lsk/ainet/lang/tensor/data/Fp16DenseTensorData; + public final fun fromRawBytes (Lsk/ainet/lang/tensor/Shape;[B)Lsk/ainet/lang/tensor/data/Fp16DenseTensorData; +} + public abstract interface class sk/ainet/lang/tensor/data/IntArrayTensorData : sk/ainet/lang/tensor/data/TensorData { public abstract fun getBuffer ()[I } @@ -3073,6 +3078,42 @@ public final class sk/ainet/lang/tensor/data/MmapTensorSource$Companion { public static synthetic fun fromChannel$default (Lsk/ainet/lang/tensor/data/MmapTensorSource$Companion;Ljava/nio/channels/FileChannel;JJILjava/lang/Object;)Lsk/ainet/lang/tensor/data/MmapTensorSource; } +public class sk/ainet/lang/tensor/data/NarrowFloatDenseTensorData : sk/ainet/lang/tensor/data/NarrowFloatTensorData { + public static final field Companion Lsk/ainet/lang/tensor/data/NarrowFloatDenseTensorData$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 class sk/ainet/lang/tensor/data/NarrowFloatDenseTensorData$Companion { + public final fun fromFloatArray (Lsk/ainet/lang/tensor/Shape;[FLsk/ainet/lang/types/NarrowFloatCodec;)Lsk/ainet/lang/tensor/data/NarrowFloatDenseTensorData; +} + +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; + public abstract fun getCodec ()Lsk/ainet/lang/types/NarrowFloatCodec; + public abstract fun getPackedData ()[B +} + +public final class sk/ainet/lang/tensor/data/NarrowFloatTensorData$Companion { + public static final field BYTES_PER_ELEMENT I +} + +public final class sk/ainet/lang/tensor/data/NarrowFloatTensorData$DefaultImpls { + public static fun copyToFloatArray (Lsk/ainet/lang/tensor/data/NarrowFloatTensorData;)[F +} + +public final class sk/ainet/lang/tensor/data/NarrowFloatTensorDataKt { + public static final fun toFloatArray (Lsk/ainet/lang/tensor/data/NarrowFloatTensorData;)[F +} + public final class sk/ainet/lang/tensor/data/PprintKt { public static final fun pprint (Lsk/ainet/lang/tensor/data/TensorData;)Ljava/lang/String; } @@ -5714,6 +5755,14 @@ public final class sk/ainet/lang/types/BF16 : sk/ainet/lang/types/DType { public fun promoteTo (Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/types/DType; } +public final class sk/ainet/lang/types/Bf16Codec : sk/ainet/lang/types/NarrowFloatCodec { + public static final field INSTANCE Lsk/ainet/lang/types/Bf16Codec; + public fun decode (I)F + public fun encode (F)I + public fun getBytesPerElement ()I + public fun getDtype ()Lsk/ainet/lang/types/DType; +} + public abstract interface class sk/ainet/lang/types/DType { public static final field Companion Lsk/ainet/lang/types/DType$Companion; public static fun bf16 ()Lsk/ainet/lang/types/DType; @@ -5758,6 +5807,7 @@ public final class sk/ainet/lang/types/DType$Companion { public final class sk/ainet/lang/types/DTypeExtensionsKt { public static final fun commonPrecisionWith (Lsk/ainet/lang/types/DType;Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/types/DType; public static final fun isConvertibleTo (Lsk/ainet/lang/types/DType;Lsk/ainet/lang/types/DType;)Z + public static final fun isFloatingPoint (Lsk/ainet/lang/types/DType;)Z public static final fun kotlinClass (Lsk/ainet/lang/types/DType;)Lkotlin/reflect/KClass; } @@ -5845,6 +5895,14 @@ public final class sk/ainet/lang/types/FP64 : sk/ainet/lang/types/DType { public fun promoteTo (Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/types/DType; } +public final class sk/ainet/lang/types/Fp16Codec : sk/ainet/lang/types/NarrowFloatCodec { + public static final field INSTANCE Lsk/ainet/lang/types/Fp16Codec; + public fun decode (I)F + public fun encode (F)I + public fun getBytesPerElement ()I + public fun getDtype ()Lsk/ainet/lang/types/DType; +} + public final class sk/ainet/lang/types/Int16 : sk/ainet/lang/types/DType { public static final field INSTANCE Lsk/ainet/lang/types/Int16; public fun getName ()Ljava/lang/String; @@ -5885,6 +5943,17 @@ public final class sk/ainet/lang/types/Int8 : sk/ainet/lang/types/DType { public fun promoteTo (Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/types/DType; } +public abstract interface class sk/ainet/lang/types/NarrowFloatCodec { + public abstract fun decode (I)F + public abstract fun encode (F)I + public fun getBytesPerElement ()I + public abstract fun getDtype ()Lsk/ainet/lang/types/DType; +} + +public final class sk/ainet/lang/types/NarrowFloatCodec$DefaultImpls { + public static fun getBytesPerElement (Lsk/ainet/lang/types/NarrowFloatCodec;)I +} + public final class sk/ainet/lang/types/PrecisionValidationKt { public static final fun isValidPrecisionChain (Ljava/util/List;)Z public static final fun validateCompatibility (Lsk/ainet/lang/types/DType;Lsk/ainet/lang/types/DType;)V diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Bf16TensorData.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Bf16TensorData.kt index f97c56f0..48b08989 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Bf16TensorData.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Bf16TensorData.kt @@ -1,7 +1,8 @@ package sk.ainet.lang.tensor.data import sk.ainet.lang.tensor.Shape -import sk.ainet.lang.types.DType +import sk.ainet.lang.types.Bf16Codec +import sk.ainet.lang.types.NarrowFloatCodec /** * Tensor data interface for **dense BF16** (bfloat16) values. @@ -37,26 +38,26 @@ import sk.ainet.lang.types.DType * For zero-copy access to the packed bytes (e.g. from a SIMD matmul * kernel that reads 2 bytes per element directly), use [packedData]. */ -public interface Bf16TensorData : TensorData { +public interface Bf16TensorData : NarrowFloatTensorData { - /** - * Raw packed BF16 bytes — 2 per logical element, little-endian. - * Length is `shape.volume * 2`. Safe for direct hand-off to the - * native / Panama matmul kernels without an intermediate copy. - */ - public val packedData: ByteArray + /** Always [Bf16Codec] — BF16 is the high 16 bits of an FP32 value. */ + override val codec: NarrowFloatCodec get() = Bf16Codec public companion object { /** Bytes per BF16 element. */ public const val BYTES_PER_ELEMENT: Int = 2 - /** Convert FP32 → BF16 bits (high 16 bits, zero rounding). */ - public fun floatToBf16Bits(value: Float): Int = - (value.toRawBits() ushr 16) and 0xFFFF + /** + * Convert FP32 → BF16 bits (high 16 bits, zero rounding). + * + * Delegates to [Bf16Codec] so there is a single implementation. The truncating (rather + * than round-to-nearest) behaviour is contractual — emitted weight archives and the + * on-device bf16-vs-f16 A/B depend on it. + */ + public fun floatToBf16Bits(value: Float): Int = Bf16Codec.encode(value) /** Convert BF16 bits (low 16 bits used) → FP32. */ - public fun bf16BitsToFloat(bf16Bits: Int): Float = - Float.fromBits((bf16Bits and 0xFFFF) shl 16) + public fun bf16BitsToFloat(bf16Bits: Int): Float = Bf16Codec.decode(bf16Bits) } } @@ -71,63 +72,11 @@ public interface Bf16TensorData : TensorData { */ public class Bf16DenseTensorData( initialShape: Shape, - private val data: ByteArray, -) : Bf16TensorData { - - override val shape: Shape = Shape(initialShape.dimensions.copyOf()) - private val strides: IntArray = shape.computeStrides() - override val packedData: ByteArray get() = data + data: ByteArray, +) : NarrowFloatDenseTensorData(initialShape, data, Bf16Codec), Bf16TensorData { - init { - val requiredBytes = shape.volume * Bf16TensorData.BYTES_PER_ELEMENT - require(data.size >= requiredBytes) { - "Data size ${data.size} is less than required $requiredBytes bytes for ${shape.volume} BF16 elements" - } - } - - override fun get(vararg indices: Int): Float { - val flatIndex = calcFlatIndex(indices) - val byteIdx = flatIndex * Bf16TensorData.BYTES_PER_ELEMENT - val lo = data[byteIdx].toInt() and 0xFF - val hi = data[byteIdx + 1].toInt() and 0xFF - val bf16Bits = (hi shl 8) or lo - return Float.fromBits(bf16Bits shl 16) - } - - override fun set(vararg indices: Int, value: Float) { - val flatIndex = calcFlatIndex(indices) - val byteIdx = flatIndex * Bf16TensorData.BYTES_PER_ELEMENT - val bf16Bits = Bf16TensorData.floatToBf16Bits(value) - data[byteIdx] = (bf16Bits and 0xFF).toByte() - data[byteIdx + 1] = ((bf16Bits ushr 8) and 0xFF).toByte() - } - - override fun copyToFloatArray(): FloatArray { - val volume = shape.volume - val out = FloatArray(volume) - for (i in 0 until volume) { - val byteIdx = i * Bf16TensorData.BYTES_PER_ELEMENT - val lo = data[byteIdx].toInt() and 0xFF - val hi = data[byteIdx + 1].toInt() and 0xFF - out[i] = Float.fromBits(((hi shl 8) or lo) shl 16) - } - return out - } - - private fun calcFlatIndex(indices: IntArray): Int { - require(indices.size == shape.dimensions.size) { - "Number of indices (${indices.size}) must match tensor dimensions (${shape.dimensions.size})" - } - var flatIndex = 0 - for (i in indices.indices) { - val idx = indices[i] - require(idx >= 0 && idx < shape.dimensions[i]) { - "Index $idx out of bounds for dimension $i with size ${shape.dimensions[i]}" - } - flatIndex += idx * strides[i] - } - return flatIndex - } + /** Fixed to [Bf16Codec]; resolves the base class / [Bf16TensorData] inheritance. */ + override val codec: NarrowFloatCodec get() = Bf16Codec public companion object { /** diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/DenseTensorDataFactory.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/DenseTensorDataFactory.kt index 66e1f577..96fb66fe 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/DenseTensorDataFactory.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/DenseTensorDataFactory.kt @@ -314,7 +314,7 @@ public class DenseTensorDataFactory: TensorDataFactory { @Suppress("UNCHECKED_CAST") public fun fromFloatArray(data: FloatArray, dtype: T): TensorData { return when (dtype) { - is FP32, FP16 -> { + is FP32, is FP16, is BF16 -> { createFloatTensorData(Shape(data.size), data, dtype) as TensorData } @@ -383,6 +383,12 @@ public class DenseTensorDataFactory: TensorDataFactory { } } + /** + * The [DType] instance for a 16-bit float tag. Both formats are stored float-backed here; only + * the tag differs, and the tag is what the trace and StableHLO export read. + */ + private fun narrowTag(dtype: KClass<*>): DType = if (dtype == BF16::class) BF16 else FP16 + override fun ones(shape: Shape, dtype: KClass): TensorData { @Suppress("UNCHECKED_CAST") return when (dtype) { @@ -390,9 +396,11 @@ public class DenseTensorDataFactory: TensorDataFactory { val data = FloatArray(shape.volume) { 1.0f } createFloatTensorData(shape, data, FP32 as T) as TensorData } - FP16::class -> { + FP16::class, BF16::class -> { + // Float-backed, narrow-tagged (mirrors zeros): the dtype tag is what the DSL + // trace / StableHLO export reads to pick the MLIR element type. val data = FloatArray(shape.volume) { 1.0f } - createFloatTensorData(shape, data, FP16 as T) as TensorData + createFloatTensorData(shape, data, narrowTag(dtype) as T) as TensorData } Int32::class -> { val data = IntArray(shape.volume) { 1 } @@ -414,10 +422,12 @@ public class DenseTensorDataFactory: TensorDataFactory { val data = FloatArray(shape.volume) { floatValue } createFloatTensorData(shape, data, FP32 as T) as TensorData } - FP16::class -> { + FP16::class, BF16::class -> { + // Float-backed, narrow-tagged (mirrors zeros): the dtype tag is what the DSL + // trace / StableHLO export reads to pick the MLIR element type. val floatValue = value.toFloat() val data = FloatArray(shape.volume) { floatValue } - createFloatTensorData(shape, data, FP16 as T) as TensorData + createFloatTensorData(shape, data, narrowTag(dtype) as T) as TensorData } Int32::class -> { val intValue = value.toInt() @@ -466,7 +476,7 @@ public class DenseTensorDataFactory: TensorDataFactory { } createFloatTensorData(shape, data, FP32 as T) as TensorData } - FP16::class -> { + FP16::class, BF16::class -> { val data = FloatArray(shape.volume) var hasSpare = false var spare = 0.0f @@ -488,7 +498,7 @@ public class DenseTensorDataFactory: TensorDataFactory { hasSpare = true } } - createFloatTensorData(shape, data, FP16 as T) as TensorData + createFloatTensorData(shape, data, narrowTag(dtype) as T) as TensorData } else -> throw IllegalArgumentException("randn only supports floating point types: $dtype") } @@ -631,8 +641,8 @@ public class DenseTensorDataFactory: TensorDataFactory { FP32::class -> { createFloatTensorData(shape, data, FP32 as T) as TensorData } - FP16::class -> { - createFloatTensorData(shape, data, FP16 as T) as TensorData + FP16::class, BF16::class -> { + createFloatTensorData(shape, data, narrowTag(dtype) as T) as TensorData } else -> throw IllegalArgumentException("fromFloatArray only supports floating point types: $dtype") } 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 new file mode 100644 index 00000000..5f95a53b --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/NarrowFloatTensorData.kt @@ -0,0 +1,165 @@ +package sk.ainet.lang.tensor.data + +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.types.DType +import sk.ainet.lang.types.Fp16Codec +import sk.ainet.lang.types.NarrowFloatCodec + +/** + * Tensor data whose elements are stored as **packed 16-bit floats** — two bytes per element, + * little-endian, no block structure and no per-block scale. The encoding is `Dense(2)`. + * + * This is the recognition surface for narrow-float dispatch: a matmul kernel can test + * `is NarrowFloatTensorData` and read [packedData] directly at 2 bytes per element instead of + * forcing a dequant-to-FP32 copy first. [codec] says which 16-bit format the bytes are in, so one + * kernel implementation serves both BF16 and FP16. + * + * `get`/`copyToFloatArray` decode to FP32, so consumers that do not care about the storage width + * see an ordinary float tensor. `set` re-encodes and is lossy by construction. + * + * @see sk.ainet.lang.types.Bf16Codec + * @see sk.ainet.lang.types.Fp16Codec + */ +public interface NarrowFloatTensorData : TensorData { + + /** + * Raw packed bytes — 2 per logical element, little-endian. Length is `shape.volume * 2`. + * Safe for direct hand-off to native / Panama kernels without an intermediate copy. + */ + public val packedData: ByteArray + + /** The 16-bit format [packedData] is encoded in. */ + public val codec: NarrowFloatCodec + + public companion object { + /** Bytes per element for every narrow float format. */ + public const val BYTES_PER_ELEMENT: Int = 2 + } +} + +/** + * Dense narrow-float tensor data backed by a packed byte array. + * + * Memory layout: row-major; element at flat index `i` occupies bytes `[i*2 .. i*2+1]`, low byte + * first. The concrete 16-bit format is supplied as a [NarrowFloatCodec], which is what lets BF16 + * and FP16 share one implementation rather than maintaining two parallel stacks. + * + * @param initialShape logical shape, in elements (not bytes). + * @param data packed bytes, length ≥ `shape.volume * 2`. + * @param codec the 16-bit format of [data]. + */ +public open class NarrowFloatDenseTensorData( + initialShape: Shape, + private val data: ByteArray, + override val codec: NarrowFloatCodec, +) : NarrowFloatTensorData { + + override val shape: Shape = Shape(initialShape.dimensions.copyOf()) + private val strides: IntArray = shape.computeStrides() + override val packedData: ByteArray get() = data + + init { + 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" + } + } + + override fun get(vararg indices: Int): Float { + val byteIdx = calcFlatIndex(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 = calcFlatIndex(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() + } + + override fun copyToFloatArray(): FloatArray { + val volume = shape.volume + val out = FloatArray(volume) + for (i in 0 until volume) { + val byteIdx = i * NarrowFloatTensorData.BYTES_PER_ELEMENT + val lo = data[byteIdx].toInt() and 0xFF + val hi = data[byteIdx + 1].toInt() and 0xFF + out[i] = codec.decode((hi shl 8) or lo) + } + return out + } + + private fun calcFlatIndex(indices: IntArray): Int { + require(indices.size == shape.dimensions.size) { + "Number of indices (${indices.size}) must match tensor dimensions (${shape.dimensions.size})" + } + var flatIndex = 0 + for (i in indices.indices) { + val idx = indices[i] + require(idx >= 0 && idx < shape.dimensions[i]) { + "Index $idx out of bounds for dimension $i with size ${shape.dimensions[i]}" + } + flatIndex += idx * strides[i] + } + return flatIndex + } + + public companion object { + /** Pack a `FloatArray` into [codec]'s 16-bit format. Lossy by construction. */ + public fun fromFloatArray( + shape: Shape, + values: FloatArray, + codec: NarrowFloatCodec, + ): NarrowFloatDenseTensorData { + require(values.size >= shape.volume) { + "FloatArray length ${values.size} is less than ${shape.volume} elements required" + } + val bytes = ByteArray(shape.volume * NarrowFloatTensorData.BYTES_PER_ELEMENT) + for (i in 0 until shape.volume) { + val bits = codec.encode(values[i]) + bytes[i * 2] = (bits and 0xFF).toByte() + bytes[i * 2 + 1] = ((bits ushr 8) and 0xFF).toByte() + } + return NarrowFloatDenseTensorData(shape, bytes, codec) + } + } +} + +/** + * Dense **IEEE binary16** tensor data. + * + * The counterpart to [Bf16DenseTensorData], and the class the SafeTensors and GGUF loaders name in + * their "no FP16 backing yet" errors — its absence was the reason `Require(FP16)` had to be + * rejected outright. + */ +public class Fp16DenseTensorData( + initialShape: Shape, + data: ByteArray, +) : NarrowFloatDenseTensorData(initialShape, data, Fp16Codec) { + + public companion object { + /** Wrap raw packed FP16 bytes; length must be ≥ `shape.volume * 2`. */ + public fun fromRawBytes(shape: Shape, bytes: ByteArray): Fp16DenseTensorData = + Fp16DenseTensorData(shape, bytes) + + /** Build from FP32 values, rounding each to nearest binary16 (ties to even). */ + public fun fromFloatArray(shape: Shape, values: FloatArray): Fp16DenseTensorData { + require(values.size >= shape.volume) { + "FloatArray length ${values.size} is less than ${shape.volume} FP16 elements required" + } + val bytes = ByteArray(shape.volume * NarrowFloatTensorData.BYTES_PER_ELEMENT) + for (i in 0 until shape.volume) { + val bits = Fp16Codec.encode(values[i]) + bytes[i * 2] = (bits and 0xFF).toByte() + bytes[i * 2 + 1] = ((bits ushr 8) and 0xFF).toByte() + } + return Fp16DenseTensorData(shape, bytes) + } + } +} + +/** Decode a [NarrowFloatTensorData] to a fresh FloatArray. */ +public fun NarrowFloatTensorData.toFloatArray(): FloatArray = copyToFloatArray() 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 a93b463e..10ef17e3 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 @@ -10,9 +10,12 @@ import sk.ainet.lang.tensor.data.Q4_KTensorData import sk.ainet.lang.tensor.data.Q8_0BlockTensorData import sk.ainet.lang.tensor.data.Q8_0TensorData import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.types.Bf16Codec import sk.ainet.lang.types.DType import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.Fp16Codec import sk.ainet.lang.types.Int32 +import sk.ainet.lang.types.NarrowFloatCodec /** * Factory methods for constructing [TensorStorage] from existing SKaiNET types @@ -175,10 +178,22 @@ public object TensorStorageFactory { return when (storage.encoding) { is TensorEncoding.Dense -> when (storage.logicalType) { - LogicalDType.FLOAT32, LogicalDType.FLOAT16, LogicalDType.BFLOAT16 -> { + LogicalDType.FLOAT32 -> { val floats = bytesToFloatArray(bytes) DenseFloatArrayTensorData(storage.shape, floats) as TensorData } + // 2 bytes per element, NOT 4. These previously shared the FLOAT32 branch, so + // `bytesToFloatArray` produced half the element count with each "float" assembled + // from two adjacent narrow elements — silent garbage. Readers already tag these + // correctly as `TensorEncoding.Dense(bytesPerElement = 2)`; the decode side simply + // ignored it. Widening to f32 here keeps this method's existing contract (it has + // always returned float-backed data); preserving 2-byte storage end-to-end is the + // loader `KEEP_NATIVE` path, not this one. + LogicalDType.FLOAT16, LogicalDType.BFLOAT16 -> { + val codec = if (storage.logicalType == LogicalDType.FLOAT16) Fp16Codec else Bf16Codec + val floats = narrowBytesToFloatArray(bytes, codec) + DenseFloatArrayTensorData(storage.shape, floats) as TensorData + } LogicalDType.INT32 -> { val ints = bytesToIntArray(bytes) DenseIntArrayTensorData(storage.shape, ints) as TensorData @@ -214,6 +229,21 @@ public object TensorStorageFactory { ) } + /** + * Decode packed little-endian 16-bit floats to FP32 using [codec]. One element per 2 bytes — + * the element count is `bytes.size / 2`, which is the whole point of this helper existing + * separately from [bytesToFloatArray]. + */ + private fun narrowBytesToFloatArray(bytes: ByteArray, codec: NarrowFloatCodec): FloatArray { + val count = bytes.size / codec.bytesPerElement + return FloatArray(count) { i -> + val off = i * 2 + val lo = bytes[off].toInt() and 0xFF + val hi = bytes[off + 1].toInt() and 0xFF + codec.decode((hi shl 8) or lo) + } + } + private fun bytesToFloatArray(bytes: ByteArray): FloatArray { val count = bytes.size / 4 return FloatArray(count) { i -> diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/trace/TraceSession.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/trace/TraceSession.kt index 98a64f84..fb610fd8 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/trace/TraceSession.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/trace/TraceSession.kt @@ -18,10 +18,15 @@ public open class TraceSession { public open fun refOf(tensor: Tensor<*, *>): TensorRef { val key = unwrap(tensor) return tensorToRef.getOrPut(key) { + // The captured dtype is what the StableHLO converter reads to pick an MLIR element + // type, so a missing arm here silently downgrades the emitted graph. BF16 was absent + // and fell through to FP32, which is why bf16 weights could only be produced by + // rewriting the emitted MLIR text after the fact. val dtypeInstance: DType = when (tensor.dtype) { Int32::class -> Int32 FP32::class -> FP32 FP16::class -> FP16 + BF16::class -> BF16 Int8::class -> Int8 Int4::class -> Int4 Ternary::class -> Ternary diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/DTypeExtensions.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/DTypeExtensions.kt index f382e16c..a5b24274 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/DTypeExtensions.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/DTypeExtensions.kt @@ -24,36 +24,32 @@ public fun DType.kotlinClass(): KClass<*> = when (this) { */ public fun DType.isConvertibleTo(target: DType): Boolean = when { this == target -> true - // Floating point conversions - this is FP32 && target is FP16 -> true - this is FP16 && target is FP32 -> true + // Any float width converts to any other — FP32 is a strict superset of both 16-bit + // formats, and narrowing is a well-defined (lossy) rounding. + this.isFloatingPoint() && target.isFloatingPoint() -> true // Integer conversions (with potential precision loss warnings) this is Int32 && target is Int8 -> true this is Int8 && target is Int32 -> true this is Int8 && target is Int4 -> true this is Int4 && target is Int8 -> true // Mixed float-int conversions - this is FP32 && target is Int32 -> true - this is FP16 && target is Int32 -> true - this is Int32 && target is FP32 -> true - this is Int32 && target is FP16 -> true + this.isFloatingPoint() && target is Int32 -> true + this is Int32 && target.isFloatingPoint() -> true // Ternary conversions this is Ternary && target is Int8 -> true this is Int8 && target is Ternary -> true else -> false } +/** True for the IEEE-style float types: [FP16], [BF16], [FP32], [FP64]. */ +public fun DType.isFloatingPoint(): Boolean = this is FP16 || this is BF16 || this is FP32 || this is FP64 + /** - * Returns the common precision type for mixed operations + * Returns the common precision type for mixed operations. + * + * Delegates to [DType.promoteTo], which is the exhaustive per-type lattice. This used to be a + * second, hand-written lattice that disagreed with it — notably it had no BF16 arm at all, so + * `BF16.commonPrecisionWith(Int8)` fell through to FP32 while `BF16.promoteTo(Int8)` said BF16. + * One lattice, one answer. */ -public fun DType.commonPrecisionWith(other: DType): DType = when { - this == other -> this - // Floating point takes precedence - this is FP32 || other is FP32 -> FP32 - this is FP16 || other is FP16 -> FP16 - // Higher precision integer takes precedence - this is Int32 || other is Int32 -> Int32 - this is Int8 || other is Int8 -> Int8 - this is Int4 || other is Int4 -> Int4 - else -> FP32 // Default fallback -} \ No newline at end of file +public fun DType.commonPrecisionWith(other: DType): DType = promoteTo(other) \ No newline at end of file diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/FP16.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/FP16.kt index b3e46f89..e72d2ca1 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/FP16.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/FP16.kt @@ -21,7 +21,9 @@ public object FP16 : DType { is Ternary -> FP16 // FP16 + Ternary → FP16 is Int4 -> FP16 // FP16 + Int4 → FP16 is Int8 -> FP16 // FP16 + Int8 → FP16 - is Int32 -> FP16 // FP16 + Int32 → FP16 + // Int32 needs 32 bits of integer precision; FP16 carries 11 mantissa bits, so + // promoting to FP16 would silently drop ~21 of them. Matches Int32.promoteTo(FP16). + is Int32 -> FP32 // FP16 + Int32 → FP32 is FP16 -> FP16 // FP16 + FP16 → FP16 is FP32 -> FP32 // FP16 + FP32 → FP32 is FP64 -> FP64 // FP16 + FP64 → FP64 diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/FP32.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/FP32.kt index af83d270..0e170e1e 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/FP32.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/FP32.kt @@ -6,7 +6,7 @@ public object FP32 : DType { override fun isCompatible(other: DType): Boolean { return when (other) { - is Ternary -> false // Ternary can promote to FP32 + is Ternary -> true // Ternary can promote to FP32 is Int4 -> true // Int4 can promote to FP32 is Int8 -> true // Int8 can promote to FP32 is Int32 -> true // Int32 can promote to FP32 @@ -18,7 +18,7 @@ public object FP32 : DType { override fun promoteTo(other: DType): DType { return when (other) { - is Ternary -> Ternary // FP32 + Ternary → FP32 + is Ternary -> FP32 // FP32 + Ternary → FP32 is Int4 -> FP32 // FP32 + Int4 → FP32 is Int8 -> FP32 // FP32 + Int8 → FP32 is Int32 -> FP32 // FP32 + Int32 → FP32 diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Int8.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Int8.kt index dc3b09cd..e4254607 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Int8.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Int8.kt @@ -29,7 +29,9 @@ public object Int8 : DType { is UInt32 -> Int64 // Int8 + UInt32 → Int64 is UInt64 -> FP64 // Int8 + UInt64 → FP64 (no integer type covers both) is FP16 -> FP16 // Int8 + FP16 → FP16 - is BF16 -> FP32 // Int8 + BF16 → FP32 + // bf16 carries 8 significand bits, so it represents every Int8 value (-128..127) + // exactly — same reason Int8 + FP16 stays FP16. Matches BF16.promoteTo(Int8). + is BF16 -> BF16 // Int8 + BF16 → BF16 is FP32 -> FP32 // Int8 + FP32 → FP32 is FP64 -> FP64 // Int8 + FP64 → FP64 } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/NarrowFloatCodec.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/NarrowFloatCodec.kt new file mode 100644 index 00000000..512152ec --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/NarrowFloatCodec.kt @@ -0,0 +1,152 @@ +package sk.ainet.lang.types + +/** + * Bit-level codec for a **dense 16-bit float** format. + * + * Both supported formats store one element per 2 bytes and are decoded to `Float` for compute — + * narrow floats are a *storage* width here, never an accumulate width. Kernels widen to f32 lanes, + * accumulate in f32, and narrow again on store, which is what PyTorch/JAX/tensor cores do. + * + * The two formats are not interchangeable: + * + * | | exponent | mantissa | relative step | max finite | + * |---|---|---|---|---| + * | [Bf16Codec] | 8 bits | 7 (+1 implicit) | 2⁻⁸ ≈ 0.39% | ~3.39e38 (f32 range) | + * | [Fp16Codec] | 5 bits | 10 (+1 implicit) | 2⁻¹¹ ≈ 0.049% | 65504 | + * + * BF16 keeps FP32's full exponent range, so overflow is structurally impossible on conversion; + * FP16 buys three more mantissa bits at the cost of a 65504 ceiling. For weights dequantized from + * block-quantized sources (Q4_K/Q5_K/…) the quantization error dominates both. + * + * Implementations are pure integer bit math with no JVM/JDK dependency, so they are usable from + * every Kotlin Multiplatform target. (The JVM's `Float.floatToFloat16` is JDK 20+ and unavailable + * in `commonMain`.) + */ +public interface NarrowFloatCodec { + + /** The [DType] this codec encodes — [BF16] or [FP16]. */ + public val dtype: DType + + /** Always 2. Present so callers can size buffers without branching on [dtype]. */ + public val bytesPerElement: Int get() = 2 + + /** + * Encode an FP32 value into this format's 16-bit pattern (returned in the low 16 bits). + * Lossy by construction; see each implementation for its rounding rule. + */ + public fun encode(value: Float): Int + + /** Decode a 16-bit pattern (low 16 bits used) back to FP32. Always exact — f32 is a superset. */ + public fun decode(bits: Int): Float +} + +/** + * **bfloat16** — the high 16 bits of an IEEE FP32 value. + * + * Conversion is the bit-shift identity in both directions, which is why it costs essentially + * nothing: `float_bits = bf16 shl 16`. + * + * [encode] **truncates** (round-toward-zero) rather than rounding to nearest. This is deliberate + * and must not be "fixed": it matches the existing `Bf16TensorData.floatToBf16Bits` contract, the + * safetensors writer, and the on-device A/B that verified bf16 as a bit-exact drop-in for the f16 + * vmfb. Changing it would silently alter every emitted weight archive. + */ +public object Bf16Codec : NarrowFloatCodec { + + override val dtype: DType get() = BF16 + + /** FP32 → BF16 by truncation (high 16 bits, zero rounding). Bit-exact with the legacy path. */ + override fun encode(value: Float): Int = (value.toRawBits() ushr 16) and 0xFFFF + + /** BF16 → FP32 — exact; the low 16 mantissa bits are zero-filled. */ + override fun decode(bits: Int): Float = Float.fromBits((bits and 0xFFFF) shl 16) +} + +/** + * **IEEE 754 binary16** (half precision): 1 sign / 5 exponent / 10 mantissa bits. + * + * [encode] implements **round-to-nearest, ties-to-even** across the full input domain: + * normals, gradual underflow into binary16 subnormals, overflow to ±Inf beyond 65504, and + * NaN payload preservation (never collapsing a NaN into an infinity). + */ +public object Fp16Codec : NarrowFloatCodec { + + override val dtype: DType get() = FP16 + + /** + * FP32 → FP16, round-to-nearest-ties-to-even. + * + * Values above binary16's max finite (65504) plus half an ulp go to ±Inf; values at or below + * half the smallest subnormal (2⁻²⁵) go to ±0. NaN stays NaN with as much payload as fits, + * and is forced quiet. + */ + override fun encode(value: Float): Int { + val bits = value.toRawBits() + val sign = (bits ushr 16) and 0x8000 + val absBits = bits and 0x7FFF_FFFF + + // Inf / NaN — exponent field all ones. + if (absBits >= 0x7F80_0000) { + if (absBits == 0x7F80_0000) return sign or 0x7C00 // ±Inf + // NaN: keep the top payload bits, force quiet, and never let the payload become 0 + // (that would turn a NaN into an Inf). + val payload = (absBits ushr 13) and 0x03FF + return sign or 0x7C00 or 0x0200 or payload + } + + val exp = (absBits ushr 23) - 127 // unbiased FP32 exponent + val mant = absBits and 0x007F_FFFF + + // Beyond binary16's exponent range — round-to-nearest still lands on Inf. + if (exp >= 16) return sign or 0x7C00 + + if (exp >= -14) { + // Normal binary16: drop 13 mantissa bits, round to nearest even. + val half = ((exp + 15) shl 10) or (mant ushr 13) + val rem = mant and 0x1FFF + val roundUp = rem > 0x1000 || (rem == 0x1000 && (half and 1) == 1) + // A carry out of the mantissa flows into the exponent field, which is the correct + // IEEE result (and produces Inf at the top of the range). + return sign or (half + if (roundUp) 1 else 0) + } + + // Gradual underflow: binary16 subnormals represent m * 2^-24 for m in [1, 1023]. + // Anything below 2^-25 (half the smallest subnormal) rounds to zero. + if (exp < -25) return sign + val full = mant or 0x0080_0000 // restore the implicit leading 1 + val shift = -1 - exp // exp <= -15 => shift >= 14 + val q = full ushr shift + val rem = full and ((1 shl shift) - 1) + val halfBit = 1 shl (shift - 1) + val roundUp = rem > halfBit || (rem == halfBit && (q and 1) == 1) + // q may carry from 1023 to 1024, yielding the smallest normal — also correct IEEE. + return sign or (q + if (roundUp) 1 else 0) + } + + /** FP16 → FP32 — always exact, including subnormals (which renormalize into FP32 normals). */ + override fun decode(bits: Int): Float { + val h = bits and 0xFFFF + val sign = (h and 0x8000) shl 16 + val exp = (h ushr 10) and 0x1F + val mant = h and 0x03FF + + return when (exp) { + 0 -> { + if (mant == 0) { + Float.fromBits(sign) // ±0 + } else { + // Subnormal: renormalize until the implicit bit position is set. + var m = mant + var e = -1 + do { + m = m shl 1 + e++ + } while (m and 0x0400 == 0) + Float.fromBits(sign or ((127 - 15 - e) shl 23) or ((m and 0x03FF) shl 13)) + } + } + 0x1F -> Float.fromBits(sign or 0x7F80_0000 or (mant shl 13)) // ±Inf / NaN + else -> Float.fromBits(sign or ((exp - 15 + 127) shl 23) or (mant shl 13)) + } + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/NarrowFloatDenseTensorDataTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/NarrowFloatDenseTensorDataTest.kt new file mode 100644 index 00000000..e7734124 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/NarrowFloatDenseTensorDataTest.kt @@ -0,0 +1,185 @@ +package sk.ainet.lang.tensor.data + +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.types.BF16 +import sk.ainet.lang.types.Bf16Codec +import sk.ainet.lang.types.FP16 +import sk.ainet.lang.types.Fp16Codec +import kotlin.math.abs +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertNotEquals +import kotlin.test.assertTrue + +/** + * Contract for the shared narrow-float storage: [NarrowFloatDenseTensorData] and its + * [Fp16DenseTensorData] specialization. Mirrors `Bf16TensorDataTest`, which continues to pin the + * BF16 side independently. + */ +class NarrowFloatDenseTensorDataTest { + + @Test + fun fp16_round_trips_exactly_representable_values() { + val values = floatArrayOf(1.0f, -2.0f, 0.5f, 3.0f, -0.25f, 1024.0f) + val data = Fp16DenseTensorData.fromFloatArray(Shape(values.size), values) + val out = data.copyToFloatArray() + for (i in values.indices) { + assertEquals(values[i], out[i], "element $i") + } + } + + @Test + fun fp16_stores_two_bytes_per_element_little_endian() { + val data = Fp16DenseTensorData.fromFloatArray(Shape(1), floatArrayOf(1.0f)) + assertEquals(2, data.packedData.size) + // 1.0f is 0x3C00 in binary16 -> low byte first + assertEquals(0x00.toByte(), data.packedData[0]) + assertEquals(0x3C.toByte(), data.packedData[1]) + } + + @Test + fun fp16_reports_its_codec_and_dtype() { + val data = Fp16DenseTensorData.fromFloatArray(Shape(2), floatArrayOf(1f, 2f)) + assertEquals(Fp16Codec, data.codec) + assertEquals(FP16, data.codec.dtype) + assertEquals(2, data.codec.bytesPerElement) + } + + @Test + fun bf16_reports_its_codec_and_dtype() { + val data = Bf16DenseTensorData.fromFloatArray(Shape(2), floatArrayOf(1f, 2f)) + assertEquals(Bf16Codec, data.codec) + assertEquals(BF16, data.codec.dtype) + } + + @Test + fun bf16_dense_is_still_recognized_as_both_supertypes() { + // Dispatch sites test `is Bf16TensorData`; narrow-float kernels test `is NarrowFloatTensorData`. + // Re-parenting Bf16DenseTensorData must keep both working. Erased to Any so these are real + // runtime checks rather than statically-known-true ones. + val data: Any = Bf16DenseTensorData.fromFloatArray(Shape(2), floatArrayOf(1f, 2f)) + assertTrue(data is Bf16TensorData, "must still satisfy the legacy BF16 dispatch check") + assertTrue(data is NarrowFloatTensorData, "must satisfy the shared narrow-float check") + } + + @Test + fun fp16_is_narrow_but_not_bf16() { + val data: Any = Fp16DenseTensorData.fromFloatArray(Shape(2), floatArrayOf(1f, 2f)) + assertTrue(data is NarrowFloatTensorData) + assertTrue(data !is Bf16TensorData, "FP16 must not be routed to a BF16 kernel") + } + + @Test + fun the_two_formats_disagree_on_the_same_bytes() { + // Same value, different encodings — proves the codec is doing real work per format. + val v = floatArrayOf(1.0f) + val fp16 = Fp16DenseTensorData.fromFloatArray(Shape(1), v) + val bf16 = Bf16DenseTensorData.fromFloatArray(Shape(1), v) + assertNotEquals( + fp16.packedData.toList(), bf16.packedData.toList(), + "1.0 encodes as 0x3C00 in fp16 but 0x3F80 in bf16", + ) + // ...yet both decode back to the same value. + assertEquals(1.0f, fp16.copyToFloatArray()[0]) + assertEquals(1.0f, bf16.copyToFloatArray()[0]) + } + + @Test + fun generic_factory_honours_the_supplied_codec() { + val values = floatArrayOf(1.5f, -0.75f, 256.0f) + val asFp16 = NarrowFloatDenseTensorData.fromFloatArray(Shape(3), values, Fp16Codec) + val asBf16 = NarrowFloatDenseTensorData.fromFloatArray(Shape(3), values, Bf16Codec) + + // All three are exactly representable in both formats. + for (i in values.indices) { + assertEquals(values[i], asFp16.copyToFloatArray()[i]) + assertEquals(values[i], asBf16.copyToFloatArray()[i]) + } + assertEquals(Fp16Codec, asFp16.codec) + assertEquals(Bf16Codec, asBf16.codec) + } + + @Test + fun get_and_set_decode_and_re_encode() { + val data = Fp16DenseTensorData.fromFloatArray(Shape(2, 2), floatArrayOf(1f, 2f, 3f, 4f)) + assertEquals(1.0f, data.get(0, 0)) + assertEquals(2.0f, data.get(0, 1)) + assertEquals(3.0f, data.get(1, 0)) + assertEquals(4.0f, data.get(1, 1)) + + data.set(1, 0, value = 9.0f) + assertEquals(9.0f, data.get(1, 0)) + assertEquals(4.0f, data.get(1, 1), "neighbouring element must be untouched") + } + + @Test + fun multi_dimensional_striding_is_row_major() { + val values = FloatArray(24) { it.toFloat() } + val data = Fp16DenseTensorData.fromFloatArray(Shape(2, 3, 4), values) + assertEquals(0.0f, data.get(0, 0, 0)) + assertEquals(7.0f, data.get(0, 1, 3)) + assertEquals(23.0f, data.get(1, 2, 3)) + assertEquals(values.toList(), data.copyToFloatArray().toList()) + } + + @Test + fun bulk_and_elementwise_reads_agree() { + val values = FloatArray(16) { (it - 8) * 0.5f } + val data = Fp16DenseTensorData.fromFloatArray(Shape(16), values) + val bulk = data.copyToFloatArray() + for (i in values.indices) { + assertEquals(bulk[i], data.get(i), "bulk vs elementwise at $i") + } + } + + @Test + fun signed_zero_is_preserved_bit_exactly() { + val data = Fp16DenseTensorData.fromFloatArray(Shape(2), floatArrayOf(0.0f, -0.0f)) + assertEquals(0x0000, (data.packedData[1].toInt() and 0xFF shl 8) or (data.packedData[0].toInt() and 0xFF)) + assertEquals(0x8000, (data.packedData[3].toInt() and 0xFF shl 8) or (data.packedData[2].toInt() and 0xFF)) + } + + @Test + fun fp16_saturates_to_infinity_beyond_its_range() { + // The practical difference from bf16: 70000 does not fit in binary16. + val data = Fp16DenseTensorData.fromFloatArray(Shape(1), floatArrayOf(70000f)) + assertTrue(data.copyToFloatArray()[0].isInfinite(), "fp16 overflows where bf16 would not") + + val bf16 = Bf16DenseTensorData.fromFloatArray(Shape(1), floatArrayOf(70000f)) + assertTrue(bf16.copyToFloatArray()[0].isFinite(), "bf16 keeps f32 exponent range") + } + + @Test + fun undersized_buffers_are_rejected() { + assertFailsWith { + Fp16DenseTensorData(Shape(4), ByteArray(6)) // needs 8 bytes + } + assertFailsWith { + NarrowFloatDenseTensorData(Shape(4), ByteArray(6), Bf16Codec) + } + } + + @Test + fun out_of_bounds_and_wrong_arity_indices_are_rejected() { + val data = Fp16DenseTensorData.fromFloatArray(Shape(2, 2), floatArrayOf(1f, 2f, 3f, 4f)) + assertFailsWith { data.get(2, 0) } + assertFailsWith { data.get(0) } + assertFailsWith { data.get(0, 0, 0) } + } + + @Test + fun fp16_precision_beats_bf16_on_the_same_tensor() { + val values = floatArrayOf(1.1f, 2.2f, 3.3f, 4.4f, 5.5f) + val fp16 = Fp16DenseTensorData.fromFloatArray(Shape(5), values).copyToFloatArray() + val bf16 = Bf16DenseTensorData.fromFloatArray(Shape(5), values).copyToFloatArray() + + var fp16Err = 0.0 + var bf16Err = 0.0 + for (i in values.indices) { + fp16Err += abs(fp16[i] - values[i]).toDouble() + bf16Err += abs(bf16[i] - values[i]).toDouble() + } + assertTrue(fp16Err < bf16Err, "fp16 err=$fp16Err should beat bf16 err=$bf16Err") + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/NarrowFloatStorageDecodeTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/NarrowFloatStorageDecodeTest.kt new file mode 100644 index 00000000..5122acee --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/NarrowFloatStorageDecodeTest.kt @@ -0,0 +1,93 @@ +package sk.ainet.lang.tensor.storage + +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.types.BF16 +import sk.ainet.lang.types.Bf16Codec +import sk.ainet.lang.types.FP16 +import sk.ainet.lang.types.Fp16Codec +import sk.ainet.lang.types.NarrowFloatCodec +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +/** + * Regression guard for the narrow-float width bug in [TensorStorageFactory.toTensorData]. + * + * `FLOAT16`/`BFLOAT16` used to share the `FLOAT32` branch and were decoded by a hard 4-byte + * reader. For an N-element narrow tensor that produced N/2 elements, each one assembled from two + * adjacent narrow values — silently wrong data rather than an error. Readers already tagged these + * as `TensorEncoding.Dense(bytesPerElement = 2)`; only the decode side ignored it. + */ +class NarrowFloatStorageDecodeTest { + + private fun packed(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 + } + + private fun storageOf(values: FloatArray, codec: NarrowFloatCodec): TensorStorage { + val bytes = packed(values, codec) + return TensorStorage( + shape = Shape(values.size), + logicalType = if (codec === Fp16Codec) LogicalDType.FLOAT16 else LogicalDType.BFLOAT16, + encoding = TensorEncoding.Dense(2), + buffer = BufferHandle.Owned(bytes, 0, bytes.size.toLong()), + ) + } + + @Test + fun bf16_dense_storage_decodes_at_two_bytes_per_element() { + val values = floatArrayOf(1.0f, -2.0f, 0.5f, 3.0f, -0.25f, 8.0f) + val data: TensorData = TensorStorageFactory.toTensorData(storageOf(values, Bf16Codec)) + + val out = data.copyToFloatArray() + assertEquals(values.size, out.size, "element count must not halve") + for (i in values.indices) { + assertEquals(values[i], out[i], "element $i") + } + } + + @Test + fun fp16_dense_storage_decodes_at_two_bytes_per_element() { + val values = floatArrayOf(1.0f, -2.0f, 0.5f, 3.0f, -0.25f, 8.0f) + val data: TensorData = TensorStorageFactory.toTensorData(storageOf(values, Fp16Codec)) + + val out = data.copyToFloatArray() + assertEquals(values.size, out.size, "element count must not halve") + for (i in values.indices) { + assertEquals(values[i], out[i], "element $i") + } + } + + @Test + fun odd_element_counts_survive_the_narrow_decode() { + // An odd count is where a 4-byte reader loses the trailing element entirely. + val values = floatArrayOf(1.0f, 2.0f, 4.0f, 8.0f, 16.0f) + val out = TensorStorageFactory.toTensorData(storageOf(values, Bf16Codec)) + .copyToFloatArray() + assertEquals(5, out.size) + assertEquals(16.0f, out[4], "trailing element must not be dropped") + } + + @Test + fun bf16_and_fp16_decode_differently_for_the_same_bytes() { + // Proves the codec is actually selected by logicalType rather than one being aliased + // to the other: the bit pattern 0x3C00 is 1.0 in fp16 but ~0.0078 in bf16. + val bytes = byteArrayOf(0x00, 0x3C) + val asFp16 = TensorStorageFactory.toTensorData( + TensorStorage(Shape(1), LogicalDType.FLOAT16, TensorEncoding.Dense(2), BufferHandle.Owned(bytes, 0, 2)), + ).copyToFloatArray()[0] + val asBf16 = TensorStorageFactory.toTensorData( + TensorStorage(Shape(1), LogicalDType.BFLOAT16, TensorEncoding.Dense(2), BufferHandle.Owned(bytes, 0, 2)), + ).copyToFloatArray()[0] + + assertEquals(1.0f, asFp16) + assertTrue(asBf16 < 0.01f && asBf16 > 0.0f, "bf16 reading of 0x3C00 should be ~0.0078, got $asBf16") + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/types/DTypeLatticeConsistencyTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/types/DTypeLatticeConsistencyTest.kt new file mode 100644 index 00000000..f0fb10d6 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/types/DTypeLatticeConsistencyTest.kt @@ -0,0 +1,127 @@ +package sk.ainet.lang.types + +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +/** + * The type lattice must have exactly one answer per question. + * + * `commonPrecisionWith` used to be a second, hand-written lattice that disagreed with + * [DType.promoteTo] — it had no BF16 arm at all, so `BF16.commonPrecisionWith(Int8)` fell through + * to FP32 while `BF16.promoteTo(Int8)` said BF16. It now delegates, and these tests pin that plus + * the symmetry and precision-safety properties the lattice is supposed to have. + */ +class DTypeLatticeConsistencyTest { + + private val floats = listOf(FP16, BF16, FP32, FP64) + private val all = listOf( + Ternary, Int4, Int8, Int16, Int32, Int64, + UInt8, UInt16, UInt32, UInt64, + FP16, BF16, FP32, FP64, + ) + + @Test + fun commonPrecisionWith_agrees_with_promoteTo_everywhere() { + for (a in all) { + for (b in all) { + assertEquals( + a.promoteTo(b), a.commonPrecisionWith(b), + "lattices disagree for ${a.name} + ${b.name}", + ) + } + } + } + + @Test + fun bf16_promotion_is_now_reachable_through_commonPrecisionWith() { + // The specific contradiction that motivated this: small ints stay in BF16. + assertEquals(BF16, BF16.commonPrecisionWith(Int8)) + assertEquals(BF16, BF16.commonPrecisionWith(Int4)) + assertEquals(BF16, BF16.commonPrecisionWith(Ternary)) + // ...while anything needing more range or precision escalates. + assertEquals(FP32, BF16.commonPrecisionWith(Int32)) + assertEquals(FP32, BF16.commonPrecisionWith(FP32)) + assertEquals(FP64, BF16.commonPrecisionWith(FP64)) + } + + @Test + fun the_two_narrow_floats_promote_to_fp32_in_both_directions() { + // Neither subsumes the other: fp16 has more mantissa, bf16 more exponent. + assertEquals(FP32, FP16.promoteTo(BF16)) + assertEquals(FP32, BF16.promoteTo(FP16)) + assertEquals(FP32, FP16.commonPrecisionWith(BF16)) + assertEquals(FP32, BF16.commonPrecisionWith(FP16)) + } + + @Test + fun float_promotion_is_symmetric() { + for (a in floats) { + for (b in floats) { + assertEquals( + a.promoteTo(b), b.promoteTo(a), + "asymmetric promotion: ${a.name}+${b.name} vs ${b.name}+${a.name}", + ) + } + } + } + + @Test + fun float_int_promotion_is_symmetric_for_the_common_types() { + // Regression: FP16.promoteTo(Int32) used to return FP16 while Int32.promoteTo(FP16) + // returned FP32 — the FP16 direction silently dropped ~21 bits of integer precision. + val ints = listOf(Int8, Int32) + for (f in floats) { + for (i in ints) { + assertEquals( + f.promoteTo(i), i.promoteTo(f), + "asymmetric promotion: ${f.name}+${i.name} vs ${i.name}+${f.name}", + ) + } + } + assertEquals(FP32, FP16.promoteTo(Int32), "Int32 does not fit in FP16's 11 mantissa bits") + } + + @Test + fun fp32_ternary_promotes_to_fp32_not_ternary() { + // The code returned Ternary while its own comment said FP32 — a promotion to a LOWER + // precision, which no promotion lattice should ever produce. + assertEquals(FP32, FP32.promoteTo(Ternary)) + assertTrue(FP32.isCompatible(Ternary), "Ternary promotes to FP32, so it is compatible") + } + + @Test + fun promotion_never_narrows_a_float() { + // A promotion result must be at least as wide as either operand. + for (a in floats) { + for (b in floats) { + val r = a.promoteTo(b) + assertTrue( + r.sizeInBits >= maxOf(a.sizeInBits, b.sizeInBits), + "${a.name}+${b.name} promoted to ${r.name}, narrower than an operand", + ) + } + } + } + + @Test + fun narrow_floats_are_convertible_in_every_float_direction() { + for (a in floats) { + for (b in floats) { + assertTrue(a.isConvertibleTo(b), "${a.name} -> ${b.name} should be convertible") + } + } + // BF16 was entirely absent from isConvertibleTo, so these all returned false. + assertTrue(BF16.isConvertibleTo(FP32)) + assertTrue(FP32.isConvertibleTo(BF16)) + assertTrue(BF16.isConvertibleTo(Int32)) + assertTrue(Int32.isConvertibleTo(BF16)) + } + + @Test + fun isFloatingPoint_covers_exactly_the_float_types() { + for (t in all) { + assertEquals(t in floats, t.isFloatingPoint(), "isFloatingPoint wrong for ${t.name}") + } + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/types/NarrowFloatCodecTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/types/NarrowFloatCodecTest.kt new file mode 100644 index 00000000..865fddb2 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/types/NarrowFloatCodecTest.kt @@ -0,0 +1,241 @@ +package sk.ainet.lang.types + +import sk.ainet.lang.tensor.data.Bf16TensorData +import kotlin.math.abs +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +/** + * Bit-level contract for [Bf16Codec] and [Fp16Codec]. + * + * The FP16 cases pin IEEE 754 binary16 semantics — normals, gradual underflow, the 65504 ceiling, + * ties-to-even, and NaN preservation. The BF16 cases pin *truncation* and bit-exact agreement with + * the pre-existing `Bf16TensorData.floatToBf16Bits`, which the emitted weight archives depend on. + */ +class NarrowFloatCodecTest { + + // ---------------------------------------------------------------- BF16 + + @Test + fun bf16_is_the_high_half_of_the_float() { + assertEquals(0x3F80, Bf16Codec.encode(1.0f)) + assertEquals(0xC000, Bf16Codec.encode(-2.0f)) + assertEquals(0x0000, Bf16Codec.encode(0.0f)) + assertEquals(0x8000, Bf16Codec.encode(-0.0f)) + } + + @Test + fun bf16_truncates_rather_than_rounds() { + // 1.0f + 1 ulp(f32) has mantissa bits below the bf16 cut; truncation must drop them + // (round-to-nearest would also give 0x3F80 here, so use a value that discriminates: + // a mantissa whose dropped bits exceed half — truncation still yields the LOWER value). + val justUnderTwo = Float.fromBits(0x3FFFFFFF) // 1.9999999… + assertEquals(0x3FFF, Bf16Codec.encode(justUnderTwo), "truncation must not round up to 0x4000") + } + + @Test + fun bf16_matches_the_legacy_floatToBf16Bits_bit_for_bit() { + // Regression guard: the on-device A/B that qualified bf16 as a drop-in for the f16 vmfb + // depends on this exact encoding. Any divergence silently changes every weight archive. + val samples = floatArrayOf( + 0.0f, -0.0f, 1.0f, -1.0f, 2.0f, 0.5f, 3.14159265f, -2.71828f, + 1e-8f, -1e-8f, 1e8f, -1e8f, 65504f, 1.0e-38f, 6.7e-5f, + Float.MAX_VALUE, Float.MIN_VALUE, Float.POSITIVE_INFINITY, Float.NEGATIVE_INFINITY, + ) + for (v in samples) { + assertEquals( + Bf16TensorData.floatToBf16Bits(v), Bf16Codec.encode(v), + "encode mismatch for $v", + ) + assertEquals( + Bf16TensorData.bf16BitsToFloat(Bf16Codec.encode(v)), Bf16Codec.decode(Bf16Codec.encode(v)), + "decode mismatch for $v", + ) + } + } + + @Test + fun bf16_keeps_f32_exponent_range() { + // The whole point of bf16 over fp16: no overflow on conversion. + assertTrue(Bf16Codec.decode(Bf16Codec.encode(1e38f)).isFinite()) + assertTrue(Bf16Codec.decode(Bf16Codec.encode(-1e38f)).isFinite()) + } + + @Test + fun bf16_round_trips_within_its_precision() { + for (v in floatArrayOf(1.0f, -3.5f, 1234.0f, 0.001f, -7.25e10f)) { + val back = Bf16Codec.decode(Bf16Codec.encode(v)) + // bf16 has 8 significand bits -> relative step 2^-8; truncation doubles the bound. + assertTrue(abs(back - v) <= abs(v) * 0.008f, "bf16 round-trip $v -> $back") + } + } + + // ---------------------------------------------------------------- FP16 normals + + @Test + fun fp16_encodes_known_normal_patterns() { + assertEquals(0x3C00, Fp16Codec.encode(1.0f)) + assertEquals(0xC000, Fp16Codec.encode(-2.0f)) + assertEquals(0x3800, Fp16Codec.encode(0.5f)) + assertEquals(0x0000, Fp16Codec.encode(0.0f)) + assertEquals(0x8000, Fp16Codec.encode(-0.0f)) + } + + @Test + fun fp16_decodes_known_normal_patterns() { + assertEquals(1.0f, Fp16Codec.decode(0x3C00)) + assertEquals(-2.0f, Fp16Codec.decode(0xC000)) + assertEquals(0.5f, Fp16Codec.decode(0x3800)) + assertEquals(0.0f, Fp16Codec.decode(0x0000)) + assertEquals(-0.0f, Fp16Codec.decode(0x8000)) + } + + @Test + fun fp16_max_finite_is_65504() { + assertEquals(0x7BFF, Fp16Codec.encode(65504f)) + assertEquals(65504f, Fp16Codec.decode(0x7BFF)) + } + + // ---------------------------------------------------------------- FP16 overflow + + @Test + fun fp16_overflows_to_infinity_past_the_ceiling() { + // Half an ulp above max finite (ulp at 65504 is 32) is the round-to-Inf threshold. + assertEquals(0x7C00, Fp16Codec.encode(65520f), "65520 is the tie -> even -> Inf") + assertEquals(0x7C00, Fp16Codec.encode(70000f)) + assertEquals(0xFC00, Fp16Codec.encode(-70000f)) + assertEquals(0x7C00, Fp16Codec.encode(1e30f), "far beyond range") + assertTrue(Fp16Codec.decode(Fp16Codec.encode(70000f)).isInfinite()) + } + + @Test + fun fp16_just_below_the_threshold_stays_finite() { + // 65519 rounds down to max finite, not up to Inf. + assertEquals(0x7BFF, Fp16Codec.encode(65519f)) + } + + // ---------------------------------------------------------------- FP16 subnormals + + @Test + fun fp16_smallest_normal_and_largest_subnormal() { + val smallestNormal = 1.0f / 16384.0f // 2^-14 + assertEquals(0x0400, Fp16Codec.encode(smallestNormal)) + assertEquals(smallestNormal, Fp16Codec.decode(0x0400)) + + val largestSubnormal = 1023.0f / 16777216.0f // 1023 * 2^-24 + assertEquals(0x03FF, Fp16Codec.encode(largestSubnormal)) + assertEquals(largestSubnormal, Fp16Codec.decode(0x03FF)) + } + + @Test + fun fp16_smallest_subnormal_is_2_pow_minus_24() { + val smallestSubnormal = 1.0f / 16777216.0f // 2^-24 + assertEquals(0x0001, Fp16Codec.encode(smallestSubnormal)) + assertEquals(smallestSubnormal, Fp16Codec.decode(0x0001)) + } + + @Test + fun fp16_underflows_to_zero_with_ties_to_even() { + val half = 1.0f / 33554432.0f // 2^-25 — exactly half the smallest subnormal + assertEquals(0x0000, Fp16Codec.encode(half), "exact tie rounds to even, i.e. zero") + assertEquals(0x8000, Fp16Codec.encode(-half), "sign is preserved on underflow") + + // Just above the tie must round up to the smallest subnormal. + val justAbove = Float.fromBits(half.toRawBits() + 1) + assertEquals(0x0001, Fp16Codec.encode(justAbove)) + + // Float subnormals are far below binary16's range and must not corrupt the exponent math. + assertEquals(0x0000, Fp16Codec.encode(Float.MIN_VALUE)) + assertEquals(0x8000, Fp16Codec.encode(-Float.MIN_VALUE)) + } + + @Test + fun fp16_subnormals_decode_to_exact_values() { + // Every binary16 subnormal is m * 2^-24 exactly, and f32 represents all of them exactly. + for (m in 1..1023) { + val expected = m.toFloat() / 16777216.0f + assertEquals(expected, Fp16Codec.decode(m), "subnormal m=$m") + assertEquals(m, Fp16Codec.encode(expected), "subnormal round-trip m=$m") + } + } + + // ---------------------------------------------------------------- FP16 rounding + + @Test + fun fp16_rounds_ties_to_even() { + // ulp at 1.0 is 2^-10, so 1 + 2^-11 sits exactly between 0x3C00 and 0x3C01. + val tieDown = 1.0f + 1.0f / 2048.0f + assertEquals(0x3C00, Fp16Codec.encode(tieDown), "tie with even LSB rounds down") + + // 1 + 3*2^-11 sits exactly between 0x3C01 (odd) and 0x3C02 -> rounds up to even. + val tieUp = 1.0f + 3.0f / 2048.0f + assertEquals(0x3C02, Fp16Codec.encode(tieUp), "tie with odd LSB rounds up") + } + + @Test + fun fp16_rounds_to_nearest_off_ties() { + assertEquals(0x3C01, Fp16Codec.encode(1.0f + 1.0f / 1024.0f), "exactly representable") + assertEquals(0x3C01, Fp16Codec.encode(1.0f + 0.9f / 1024.0f), "nearer the upper neighbour") + assertEquals(0x3C00, Fp16Codec.encode(1.0f + 0.1f / 1024.0f), "nearer the lower neighbour") + } + + // ---------------------------------------------------------------- FP16 inf / NaN + + @Test + fun fp16_preserves_infinities() { + assertEquals(0x7C00, Fp16Codec.encode(Float.POSITIVE_INFINITY)) + assertEquals(0xFC00, Fp16Codec.encode(Float.NEGATIVE_INFINITY)) + assertTrue(Fp16Codec.decode(0x7C00) == Float.POSITIVE_INFINITY) + assertTrue(Fp16Codec.decode(0xFC00) == Float.NEGATIVE_INFINITY) + } + + @Test + fun fp16_nan_never_collapses_into_infinity() { + val encoded = Fp16Codec.encode(Float.NaN) + assertTrue(Fp16Codec.decode(encoded).isNaN(), "NaN must survive the round trip") + assertTrue((encoded and 0x03FF) != 0, "payload must be non-zero or it would decode as Inf") + + // A NaN whose payload lives entirely in the low mantissa bits still must not become Inf. + val lowPayloadNaN = Float.fromBits(0x7F80_0001) + val encodedLow = Fp16Codec.encode(lowPayloadNaN) + assertTrue(Fp16Codec.decode(encodedLow).isNaN(), "low-payload NaN must survive") + } + + // ---------------------------------------------------------------- shared contract + + @Test + fun codecs_report_their_dtype_and_width() { + assertEquals(BF16, Bf16Codec.dtype) + assertEquals(FP16, Fp16Codec.dtype) + assertEquals(2, Bf16Codec.bytesPerElement) + assertEquals(2, Fp16Codec.bytesPerElement) + } + + @Test + fun fp16_is_more_precise_than_bf16_in_range() { + // fp16 has three more mantissa bits, so it can never be worse — but for individual values + // the two can coincide exactly (at pi both land on 3.140625), so assert over a sample: + // never worse pointwise, and strictly better in aggregate. + val samples = floatArrayOf( + 1.1f, 3.14159265f, 2.71828f, 0.1f, 123.456f, -7.77f, 1000.5f, 0.007f, + ) + var fp16Total = 0.0 + var bf16Total = 0.0 + for (v in samples) { + val fp16Err = abs(Fp16Codec.decode(Fp16Codec.encode(v)) - v) + val bf16Err = abs(Bf16Codec.decode(Bf16Codec.encode(v)) - v) + assertTrue(fp16Err <= bf16Err, "fp16 must never be worse: $v -> fp16=$fp16Err bf16=$bf16Err") + fp16Total += fp16Err.toDouble() + bf16Total += bf16Err.toDouble() + } + assertTrue(fp16Total < bf16Total, "fp16 total err=$fp16Total should beat bf16 total=$bf16Total") + } + + @Test + fun encode_ignores_bits_above_the_low_sixteen_on_decode() { + // decode must mask, so callers can pass a sign-extended Short without corruption. + assertEquals(Fp16Codec.decode(0x3C00), Fp16Codec.decode(0x3C00 or 0xFFFF_0000.toInt())) + assertEquals(Bf16Codec.decode(0x3F80), Bf16Codec.decode(0x3F80 or 0xFFFF_0000.toInt())) + } +}